diff --git a/.github/workflows/issue_fixed_comment.yml b/.github/workflows/issue_fixed_comment.yml new file mode 100644 index 00000000000..92993d319a7 --- /dev/null +++ b/.github/workflows/issue_fixed_comment.yml @@ -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' }} diff --git a/AGENTS.md b/AGENTS.md index a1e8f6f618d..cade08bdd02 100644 --- a/AGENTS.md +++ b/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] # ` 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_.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_.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 # `. `# 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: ` + +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: ` + - 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: ` +- 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 diff --git a/CLAUDE.md b/CLAUDE.md deleted file mode 100644 index b9753ab864b..00000000000 --- a/CLAUDE.md +++ /dev/null @@ -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] # ` 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_.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_.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 # `. `# 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: ` - -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: ` - - 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: ` -- 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 diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 0443f1bed75..82cad680a70 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -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** diff --git a/GEMINI.md b/GEMINI.md index 41921fdff4d..5fc00e0b5ae 100644 --- a/GEMINI.md +++ b/GEMINI.md @@ -1 +1 @@ -Read @CLAUDE.md for coding guidelines +Read @AGENTS.md for coding guidelines diff --git a/README.md b/README.md index 901cc5b0cea..3f3ea0bd60b 100644 --- a/README.md +++ b/README.md @@ -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** diff --git a/ci_cd/generate_model_prices_schema.py b/ci_cd/generate_model_prices_schema.py index ab29b70bdd4..8eec07dadda 100644 --- a/ci_cd/generate_model_prices_schema.py +++ b/ci_cd/generate_model_prices_schema.py @@ -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) diff --git a/enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py b/enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py index f40ced302ce..2114dfd9849 100644 --- a/enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py +++ b/enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py @@ -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, ) diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 5a8a613204f..d4b32659ba1 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -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", diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index de6eacc62ee..8634dce92d0 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -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" diff --git a/litellm-rust/crates/core/Cargo.toml b/litellm-rust/crates/core/Cargo.toml index 3995a235778..ab04fb8d4ae 100644 --- a/litellm-rust/crates/core/Cargo.toml +++ b/litellm-rust/crates/core/Cargo.toml @@ -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 diff --git a/litellm-rust/crates/core/src/ocr/client.rs b/litellm-rust/crates/core/src/ocr/client.rs index c05622932b1..c7b4751bd9e 100644 --- a/litellm-rust/crates/core/src/ocr/client.rs +++ b/litellm-rust/crates/core/src/ocr/client.rs @@ -14,7 +14,3 @@ pub async fn perform( ) -> Result { litellm_host::run::run(ocr_machine(client.clone()), &LocalOcrHost::new(request)).await } - -pub async fn ocr(request: LiteLLMOcrRequest) -> Result { - perform(&OcrClient::shared()?, request).await -} diff --git a/litellm-rust/crates/core/tests/ocr.rs b/litellm-rust/crates/core/tests/ocr.rs index 41a650945bc..e7a8fc0abc1 100644 --- a/litellm-rust/crates/core/tests/ocr.rs +++ b/litellm-rust/crates/core/tests/ocr.rs @@ -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 { diff --git a/litellm-rust/crates/http/AGENTS.md b/litellm-rust/crates/http/AGENTS.md new file mode 100644 index 00000000000..08fa34bd799 --- /dev/null +++ b/litellm-rust/crates/http/AGENTS.md @@ -0,0 +1 @@ +- https://github.com/BerriAI/litellm-docs/blob/main/docs/guides/security_settings.md diff --git a/litellm-rust/crates/http/Cargo.toml b/litellm-rust/crates/http/Cargo.toml new file mode 100644 index 00000000000..0ac09a9d155 --- /dev/null +++ b/litellm-rust/crates/http/Cargo.toml @@ -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 diff --git a/litellm-rust/crates/http/src/config.rs b/litellm-rust/crates/http/src/config.rs new file mode 100644 index 00000000000..10f28b44eec --- /dev/null +++ b/litellm-rust/crates/http/src/config.rs @@ -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, + pub key_exchange_group: Option, + pub tls12_cipher_suites: Option>, + pub force_ipv4: bool, + pub http2: bool, + pub user_agent: Option, + pub trust_proxy_env: bool, + pub connect_timeout: Duration, + pub tcp_keepalive: Option, + pub pool_idle_timeout: Duration, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct Resolution { + pub config: HttpClientConfig, + pub unsupported: Vec, +} + +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::) + .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 { + 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, 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, + ) { + 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 + )); + } +} diff --git a/litellm-rust/crates/http/src/error.rs b/litellm-rust/crates/http/src/error.rs new file mode 100644 index 00000000000..697d0cf59c8 --- /dev/null +++ b/litellm-rust/crates/http/src/error.rs @@ -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 for Error { + fn from(error: reqwest::Error) -> Self { + Self::Client(error.without_url().to_string()) + } +} diff --git a/litellm-rust/crates/http/src/lib.rs b/litellm-rust/crates/http/src/lib.rs new file mode 100644 index 00000000000..ddbc3b63b08 --- /dev/null +++ b/litellm-rust/crates/http/src/lib.rs @@ -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}; diff --git a/litellm-rust/crates/http/src/pool.rs b/litellm-rust/crates/http/src/pool.rs new file mode 100644 index 00000000000..330d6de29e8 --- /dev/null +++ b/litellm-rust/crates/http/src/pool.rs @@ -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, + ttl: Duration, + clients: Mutex, +} + +impl HttpClientPool { + pub fn new(media_resolver: Arc) -> Self { + Self::with_ttl(media_resolver, CLIENT_TTL) + } + + pub fn with_ttl(media_resolver: Arc, ttl: Duration) -> Self { + Self { + media_resolver, + ttl, + clients: Mutex::default(), + } + } + + pub fn client( + &self, + config: &HttpClientConfig, + variant: ClientVariant, + ) -> Result { + 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, Arc>>) { + 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() + ); + } +} diff --git a/litellm-rust/crates/http/src/proxy.rs b/litellm-rust/crates/http/src/proxy.rs new file mode 100644 index 00000000000..4dc4bf778b8 --- /dev/null +++ b/litellm-rust/crates/http/src/proxy.rs @@ -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::() + .is_ok_and(|uri| self.0.intercept(&uri).is_some()) + } +} diff --git a/litellm-rust/crates/http/src/settings.rs b/litellm-rust/crates/http/src/settings.rs new file mode 100644 index 00000000000..8ac7ef92568 --- /dev/null +++ b/litellm-rust/crates/http/src/settings.rs @@ -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, + pub ssl_cert_file: Option, + pub ssl_certificate: Option, + pub ssl_security_level: Option, + pub ssl_ecdh_curve: Option, + pub force_ipv4: Option, + pub http2: Option, + pub aiohttp_trust_env: Option, + pub disable_aiohttp_trust_env: Option, + pub disable_aiohttp_transport: Option, + pub user_agent: Option, + pub tcp_keepalive: Option, + pub pool_idle_timeout: Option, +} + +impl HttpSettingsLayer { + pub fn from_environment(env: &(dyn Fn(&str) -> Option + 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::().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, + pub ssl_cert_file: Option, + pub ssl_certificate: Option, + pub ssl_security_level: Option, + pub ssl_ecdh_curve: Option, + pub force_ipv4: bool, + pub http2: bool, + pub user_agent: Option, + pub trust_proxy_env: bool, + pub connect_timeout: Duration, + pub tcp_keepalive: Option, + 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, + ) -> 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 { + None + } + + fn env_of( + values: &'static [(&'static str, &'static str)], + ) -> impl Fn(&str) -> Option + 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, + ) { + 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); + } +} diff --git a/litellm-rust/crates/http/src/tls.rs b/litellm-rust/crates/http/src/tls.rs new file mode 100644 index 00000000000..aaae2b659e3 --- /dev/null +++ b/litellm-rust/crates/http/src/tls.rs @@ -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 { + 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 { + 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>, + pub(crate) unsupported: Vec, +} + +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 = 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 = 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 { + 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 { + 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 { + 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>, 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>, Error> { + CertificateDer::pem_slice_iter(&read(path)?) + .collect::>() + .map_err(|error| invalid_pem(path, error)) +} + +fn read(path: &Path) -> Result, 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); + +impl ServerCertVerifier for NoVerification { + fn verify_server_cert( + &self, + _end_entity: &CertificateDer<'_>, + _intermediates: &[CertificateDer<'_>], + _server_name: &ServerName<'_>, + _ocsp_response: &[u8], + _now: UnixTime, + ) -> Result { + Ok(ServerCertVerified::assertion()) + } + + fn verify_tls12_signature( + &self, + _message: &[u8], + _cert: &CertificateDer<'_>, + _dss: &DigitallySignedStruct, + ) -> Result { + Ok(HandshakeSignatureValid::assertion()) + } + + fn verify_tls13_signature( + &self, + _message: &[u8], + _cert: &CertificateDer<'_>, + _dss: &DigitallySignedStruct, + ) -> Result { + Ok(HandshakeSignatureValid::assertion()) + } + + fn supported_verify_schemes(&self) -> Vec { + 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 { + tls.crypto_provider() + .kx_groups + .iter() + .map(|group| group.name()) + .collect() + } + + fn offered_tls12_suites(tls: &ClientConfig) -> Vec { + 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 + )); + } +} diff --git a/litellm-rust/crates/llms/Cargo.toml b/litellm-rust/crates/llms/Cargo.toml index d295e4407ba..a81a4427b4d 100644 --- a/litellm-rust/crates/llms/Cargo.toml +++ b/litellm-rust/crates/llms/Cargo.toml @@ -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" diff --git a/litellm-rust/crates/llms/src/base_llm/ocr/transformation.rs b/litellm-rust/crates/llms/src/base_llm/ocr/transformation.rs index 8321dcfb4ce..f215546849d 100644 --- a/litellm-rust/crates/llms/src/base_llm/ocr/transformation.rs +++ b/litellm-rust/crates/llms/src/base_llm/ocr/transformation.rs @@ -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; diff --git a/litellm-rust/crates/llms/src/custom_httpx/llm_http_handler.rs b/litellm-rust/crates/llms/src/custom_httpx/llm_http_handler.rs index fdd568d83fd..58dc03eea2d 100644 --- a/litellm-rust/crates/llms/src/custom_httpx/llm_http_handler.rs +++ b/litellm-rust/crates/llms/src/custom_httpx/llm_http_handler.rs @@ -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 { - let document_fetcher = MediaFetcher::new().map_err(transport::Error::from)?; + pub fn new( + pool: &HttpClientPool, + config: &HttpClientConfig, + url_policy: UrlPolicy, + vertex_auth: VertexAuth, + ) -> Result { 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 { - static CLIENT: OnceLock> = 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::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( @@ -296,6 +279,8 @@ pub fn body_document(body: &Value) -> Result { #[cfg(test)] mod tests { + use std::time::Duration; + use super::*; #[tokio::test] diff --git a/litellm-rust/crates/llms/src/custom_httpx/media.rs b/litellm-rust/crates/llms/src/custom_httpx/media.rs index 0b7fa30e34b..572e7f12e54 100644 --- a/litellm-rust/crates/llms/src/custom_httpx/media.rs +++ b/litellm-rust/crates/llms/src/custom_httpx/media.rs @@ -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, +} + +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 bool + Send + Sync>; + #[derive(Clone)] pub struct MediaFetcher { - client: reqwest::Client, + pinned: reqwest::Client, + unpinned: reqwest::Client, + uses_proxy: ProxyMatch, address_resolver: Arc, + url_policy: UrlPolicy, allow_private_network: bool, } @@ -63,26 +97,39 @@ pub struct DownloadedMedia { } impl MediaFetcher { - pub fn new() -> Result { - Self::with_resolvers(Arc::new(PublicDnsResolver), Arc::new(SystemAddressResolver)) + pub fn new( + pool: &HttpClientPool, + config: &HttpClientConfig, + url_policy: UrlPolicy, + ) -> Result { + 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( - transport_resolver: Arc, + fn with_resolution( + pool: &HttpClientPool, + config: &HttpClientConfig, + url_policy: UrlPolicy, address_resolver: Arc, - ) -> Result - 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 { 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 { 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::() { - 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::() + { + 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))); + } } diff --git a/litellm-rust/crates/llms/src/custom_httpx/transport.rs b/litellm-rust/crates/llms/src/custom_httpx/transport.rs index 172dd96476a..c42cdf410f6 100644 --- a/litellm-rust/crates/llms/src/custom_httpx/transport.rs +++ b/litellm-rust/crates/llms/src/custom_httpx/transport.rs @@ -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 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 { + 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; diff --git a/litellm-rust/crates/python-bridge/AGENTS.md b/litellm-rust/crates/python-bridge/AGENTS.md index 5dccfb4aca8..a19a709e60c 100644 --- a/litellm-rust/crates/python-bridge/AGENTS.md +++ b/litellm-rust/crates/python-bridge/AGENTS.md @@ -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///`. + +## 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. diff --git a/litellm-rust/crates/python-bridge/CLAUDE.md b/litellm-rust/crates/python-bridge/CLAUDE.md deleted file mode 100644 index e55bb192cdd..00000000000 --- a/litellm-rust/crates/python-bridge/CLAUDE.md +++ /dev/null @@ -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///`. - -## 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. diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index e9b7f384406..c66701548d1 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -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 diff --git a/litellm-rust/crates/python-bridge/python_settings.json b/litellm-rust/crates/python-bridge/python_settings.json new file mode 100644 index 00000000000..a6f5ee9c6f4 --- /dev/null +++ b/litellm-rust/crates/python-bridge/python_settings.json @@ -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" + ] +} diff --git a/litellm-rust/crates/python-bridge/src/http.rs b/litellm-rust/crates/python-bridge/src/http.rs new file mode 100644 index 00000000000..d174dccaa56 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/http.rs @@ -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 = + LazyLock::new(|| HttpClientPool::new(Arc::new(PublicDnsResolver))); + +static REPORTED_UNSUPPORTED: LazyLock>> = LazyLock::new(Mutex::default); + +pub(crate) fn pool() -> &'static HttpClientPool { + &POOL +} + +pub(crate) fn call_config( + py: Python<'_>, + kwargs: &Bound<'_, PyDict>, + asynchronous: bool, +) -> PyResult { + 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>, + unsupported: Vec, +) -> Vec { + 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 { + 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> { + Ok(kwargs + .get_item("ssl_verify")? + .and_then(|value| ssl_verify(&value))) +} + +fn for_call(call_ssl_verify: Option, 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, +} + +#[derive(FromPyObject)] +struct PythonHttpSettings<'py> { + ssl_verify: Bound<'py, PyAny>, + ssl_certificate: Option, + ssl_security_level: Option, + ssl_ecdh_curve: Option, + 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 { + 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 { + if let Ok(enabled) = value.extract::() { + return Some(if enabled { + SslVerify::Enabled + } else { + SslVerify::Disabled + }); + } + value + .extract::() + .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::(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); + } +} diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index ca699e7c483..7eba0d201be 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -1,7 +1,9 @@ mod credentials; mod diagnostics; mod errors; +mod http; mod marshal; +mod python_settings; mod routes; mod token_counter; diff --git a/litellm-rust/crates/python-bridge/src/python_settings.rs b/litellm-rust/crates/python-bridge/src/python_settings.rs new file mode 100644 index 00000000000..79921d67452 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/python_settings.rs @@ -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> { + 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 = locals + .get_item("keys") + .unwrap() + .unwrap() + .extract::>() + .unwrap() + .into_iter() + .collect(); + let read: BTreeSet = PythonSettings::ALL + .map(|group| group.name().to_owned()) + .into(); + assert_eq!(read, declared); + }); + } +} diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs b/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs index 8afa1e2a906..f9d7024c824 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs @@ -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 = 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> { - 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 }, diff --git a/litellm/constants.py b/litellm/constants.py index a7d4eba0f15..d62cad74a36 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -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", diff --git a/litellm/integrations/websearch_interception/handler.py b/litellm/integrations/websearch_interception/handler.py index 587da997f94..b7238629c85 100644 --- a/litellm/integrations/websearch_interception/handler.py +++ b/litellm/integrations/websearch_interception/handler.py @@ -44,6 +44,8 @@ from litellm.types.integrations.custom_logger import ( from litellm.types.integrations.websearch_interception import ( AnthropicSearchQuery, AnthropicServerToolUseBlock, + SearchFailed, + SearchOutcome, WebSearchInterceptionConfig, ) from litellm.types.llms.anthropic import AnthropicThinkingParam @@ -332,16 +334,8 @@ class WebSearchInterceptionLogger(CustomLogger): None, ) - # Execute search — keep the structured SearchResponse so the native - # block can carry per-result url/title/page_age. - try: - if kwargs is None: - search_result_text, structured = await self._execute_search(query) - else: - search_result_text, structured = await self._execute_search(query, kwargs=kwargs) - except Exception as e: - verbose_logger.error("WebSearchInterception: Short-circuit search failed: %s", e) - search_result_text, structured = f"Search failed: {e}", None + outcome: Final = await self._short_circuit_search_outcome(query, kwargs=kwargs) + search_result_text: Final = WebSearchTransformation.search_outcome_text(outcome) content: Final[list[dict[str, object]]] = [] if native_tool is not None: @@ -356,10 +350,7 @@ class WebSearchInterceptionLogger(CustomLogger): } ) content.append( - WebSearchTransformation.build_web_search_tool_result_block( - tool_use_id=tool_use_id, - search_response=structured, - ) + WebSearchTransformation.build_web_search_outcome_block(tool_use_id=tool_use_id, outcome=outcome) ) # Keep the text block so non-native short-circuit callers (Claude Code, # github_copilot, etc.) see the same payload they always have. @@ -934,7 +925,7 @@ class WebSearchInterceptionLogger(CustomLogger): tool_calls: Final = tools["tool_calls"] thinking_blocks: Final = tools.get("thinking_blocks", []) - request_patch, structured_results = await self._build_anthropic_request_patch( + request_patch, search_outcomes = await self._build_anthropic_request_patch( model=model, messages=messages, tool_calls=tool_calls, @@ -953,17 +944,21 @@ class WebSearchInterceptionLogger(CustomLogger): # pre-build the Anthropic-native ``web_search_tool_result`` blocks now # (while we still have the structured SearchResponse list) and stash # them on plan metadata for the post-hook to inject. - if kwargs.get(WEBSEARCH_EMIT_NATIVE_BLOCKS_KEY): - metadata[WEBSEARCH_NATIVE_BLOCKS_METADATA_KEY] = self._build_native_result_blocks( - tool_calls=tool_calls, - structured_results=structured_results, - ) + if not kwargs.get(WEBSEARCH_EMIT_NATIVE_BLOCKS_KEY): + return AgenticLoopPlan(run_agentic_loop=True, request_patch=request_patch, metadata=metadata) - return AgenticLoopPlan( - run_agentic_loop=True, - request_patch=request_patch, - metadata=metadata, + metadata[WEBSEARCH_NATIVE_BLOCKS_METADATA_KEY] = self._build_native_result_blocks( + tool_calls=tool_calls, + search_outcomes=search_outcomes, ) + every_search_failed: Final = bool(search_outcomes) and all( + isinstance(outcome, SearchFailed) for outcome in search_outcomes + ) + if every_search_failed: + return AgenticLoopPlan( + run_agentic_loop=False, terminate=True, stop_reason="web_search_failed", metadata=metadata + ) + return AgenticLoopPlan(run_agentic_loop=True, request_patch=request_patch, metadata=metadata) async def async_post_agentic_loop_response_hook( self, @@ -992,7 +987,7 @@ class WebSearchInterceptionLogger(CustomLogger): @staticmethod def _build_native_result_blocks( tool_calls: list[dict], - structured_results: list[SearchResponse | None], + search_outcomes: Sequence[SearchOutcome], ) -> tuple[Mapping[str, object], ...]: """ Build a ``server_tool_use`` + ``web_search_tool_result`` pair per tool_call. @@ -1004,10 +999,10 @@ class WebSearchInterceptionLogger(CustomLogger): """ return tuple( block - for i, tool_call in enumerate(tool_calls) + for tool_call, outcome in zip(tool_calls, search_outcomes, strict=True) for block in WebSearchInterceptionLogger._native_result_pair( query=WebSearchInterceptionLogger._tool_call_query(tool_call), - search_response=structured_results[i] if i < len(structured_results) else None, + outcome=outcome, ) ) @@ -1022,15 +1017,12 @@ class WebSearchInterceptionLogger(CustomLogger): @staticmethod def _native_result_pair( query: str, - search_response: SearchResponse | None, + outcome: SearchOutcome, ) -> tuple[Mapping[str, object], Mapping[str, object]]: tool_use_id: Final = f"srvtoolu_{uuid.uuid4().hex}" return ( AnthropicServerToolUseBlock(id=tool_use_id, input=AnthropicSearchQuery(query=query)).model_dump(), - WebSearchTransformation.build_web_search_tool_result_block( - tool_use_id=tool_use_id, - search_response=search_response, - ), + WebSearchTransformation.build_web_search_outcome_block(tool_use_id=tool_use_id, outcome=outcome), ) @staticmethod @@ -1306,7 +1298,7 @@ class WebSearchInterceptionLogger(CustomLogger): kwargs: Mapping[str, object], ) -> "AnthropicMessagesResponse | AsyncIterator[object]": """Legacy path: execute search + build patch + run follow-up call.""" - request_patch, structured_results = await self._build_anthropic_request_patch( + request_patch, search_outcomes = await self._build_anthropic_request_patch( model=model, messages=messages, tool_calls=tool_calls, @@ -1344,7 +1336,7 @@ class WebSearchInterceptionLogger(CustomLogger): if kwargs.get(WEBSEARCH_EMIT_NATIVE_BLOCKS_KEY): native_blocks: Final = self._build_native_result_blocks( tool_calls=tool_calls, - structured_results=structured_results, + search_outcomes=search_outcomes, ) response = self._inject_native_blocks(response, native_blocks) @@ -1359,15 +1351,9 @@ class WebSearchInterceptionLogger(CustomLogger): anthropic_messages_optional_request_params: dict, logging_obj: "LiteLLMLoggingObj | None", kwargs: dict, - ) -> tuple[AgenticLoopRequestPatch, list[SearchResponse | None]]: + ) -> tuple[AgenticLoopRequestPatch, tuple[SearchOutcome, ...]]: """ Execute litellm.search() and build follow-up request patch. - - Returns the patch alongside the parallel list of structured - ``SearchResponse`` objects (one per tool_call, ``None`` when the - search failed or the tool_call had no query). The caller uses these - to optionally build Anthropic-native ``web_search_tool_result`` - content blocks for the final response. """ # Extract search queries from tool_use blocks @@ -1385,27 +1371,10 @@ class WebSearchInterceptionLogger(CustomLogger): # Execute searches in parallel verbose_logger.debug("WebSearchInterception: Executing %s search(es) in parallel", len(search_tasks)) search_results: Final = await asyncio.gather(*search_tasks, return_exceptions=True) - - # Split the gathered (text, structured) tuples into two parallel lists. - # The text list feeds the follow-up model call; the structured list - # is returned to the caller for native-block emission. - final_search_results: Final[list[str]] = [] - structured_results: Final[list[SearchResponse | None]] = [] - for i, result in enumerate(search_results): - if isinstance(result, Exception): - verbose_logger.error("WebSearchInterception: Search %s failed with error: %s", i, result) - final_search_results.append(f"Search failed: {result}") - structured_results.append(None) - elif isinstance(result, tuple) and len(result) == 2: - text_value, structured_value = result - final_search_results.append(cast(str, text_value) if isinstance(text_value, str) else str(text_value)) - structured_results.append(structured_value if isinstance(structured_value, SearchResponse) else None) - else: - # Defensive: legacy callers / unexpected shape — preserve text, - # drop structure. - verbose_logger.debug("WebSearchInterception: Unexpected result type %s at index %s", type(result), i) - final_search_results.append(str(result)) - structured_results.append(None) + search_outcomes: Final = tuple(WebSearchTransformation.search_outcome(result) for result in search_results) + final_search_results: Final = tuple( + WebSearchTransformation.search_outcome_text(outcome) for outcome in search_outcomes + ) # Build assistant and user messages using transformation assistant_message, user_message = WebSearchTransformation.transform_response( @@ -1449,7 +1418,18 @@ class WebSearchInterceptionLogger(CustomLogger): optional_params=optional_params_without_max_tokens, kwargs=kwargs_for_followup, ) - return patch, structured_results + return patch, search_outcomes + + async def _short_circuit_search_outcome(self, query: str, kwargs: Mapping[str, object] | None) -> SearchOutcome: + try: + result: Final = ( + await self._execute_search(query) + if kwargs is None + else await self._execute_search(query, kwargs=kwargs) + ) + except Exception as e: + return WebSearchTransformation.search_outcome(e) + return WebSearchTransformation.search_outcome(result) async def _execute_search( self, query: str, kwargs: Mapping[str, object] | None = None diff --git a/litellm/integrations/websearch_interception/transformation.py b/litellm/integrations/websearch_interception/transformation.py index fe4b6583c55..47af73570fc 100644 --- a/litellm/integrations/websearch_interception/transformation.py +++ b/litellm/integrations/websearch_interception/transformation.py @@ -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: """ diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py index f8eb15dca88..e24fa004448 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -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, diff --git a/litellm/litellm_core_utils/llm_response_utils/response_metadata.py b/litellm/litellm_core_utils/llm_response_utils/response_metadata.py index c83c266a17e..93701b3c1e7 100644 --- a/litellm/litellm_core_utils/llm_response_utils/response_metadata.py +++ b/litellm/litellm_core_utils/llm_response_utils/response_metadata.py @@ -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 diff --git a/litellm/litellm_core_utils/logging_utils.py b/litellm/litellm_core_utils/logging_utils.py index 0f14b461d3d..5be9dd7be2f 100644 --- a/litellm/litellm_core_utils/logging_utils.py +++ b/litellm/litellm_core_utils/logging_utils.py @@ -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: diff --git a/litellm/litellm_core_utils/realtime_streaming.py b/litellm/litellm_core_utils/realtime_streaming.py index fa567bdf4c9..d2fbb26bb02 100644 --- a/litellm/litellm_core_utils/realtime_streaming.py +++ b/litellm/litellm_core_utils/realtime_streaming.py @@ -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 = "", diff --git a/litellm/llms/anthropic/common_utils.py b/litellm/llms/anthropic/common_utils.py index 2de9ab41d95..3328e20e2ea 100644 --- a/litellm/llms/anthropic/common_utils.py +++ b/litellm/llms/anthropic/common_utils.py @@ -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( diff --git a/litellm/llms/base_llm/audio_transcription/transformation.py b/litellm/llms/base_llm/audio_transcription/transformation.py index b323c4812b5..2296909cfe1 100644 --- a/litellm/llms/base_llm/audio_transcription/transformation.py +++ b/litellm/llms/base_llm/audio_transcription/transformation.py @@ -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, diff --git a/litellm/llms/base_llm/realtime/transcription_protocol.py b/litellm/llms/base_llm/realtime/transcription_protocol.py new file mode 100644 index 00000000000..c3264911d46 --- /dev/null +++ b/litellm/llms/base_llm/realtime/transcription_protocol.py @@ -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 diff --git a/litellm/llms/base_llm/realtime/transformation.py b/litellm/llms/base_llm/realtime/transformation.py index e44cccc1a62..1f4ad29fa74 100644 --- a/litellm/llms/base_llm/realtime/transformation.py +++ b/litellm/llms/base_llm/realtime/transformation.py @@ -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, diff --git a/litellm/llms/bedrock/batches/handler.py b/litellm/llms/bedrock/batches/handler.py index 6239973eb7c..fd4c3dc1659 100644 --- a/litellm/llms/bedrock/batches/handler.py +++ b/litellm/llms/bedrock/batches/handler.py @@ -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:::/``""" 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: diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index 1a32fec45e3..3e412b5ad24 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/litellm/llms/bedrock/chat/converse_transformation.py @@ -108,6 +108,7 @@ BEDROCK_COMPUTER_USE_TOOLS: Final = [ "bash_", "text_editor_", ] +BEDROCK_OPENAI_COMPAT_MIN_MAX_TOKENS: Final = 16 # Beta header patterns that are not supported by Bedrock Converse API # These will be filtered out to prevent errors @@ -378,6 +379,10 @@ class AmazonConverseConfig(BaseConfig): def _is_openai_gpt_reasoning_model(model: str) -> bool: return re.search(r"openai\.gpt-\d", model) is not None + @staticmethod + def _requires_min_max_tokens(model: str) -> bool: + return re.search(r"openai\.gpt-\d|xai\.grok-", model) is not None + def _is_nova_2_model(self, model: str) -> bool: """ Check if the model is a Nova 2 model that supports reasoningConfig. @@ -1000,7 +1005,11 @@ class AmazonConverseConfig(BaseConfig): is_thinking_enabled=is_thinking_enabled, ) if param == "max_tokens" or param == "max_completion_tokens": - optional_params["maxTokens"] = value + optional_params["maxTokens"] = ( + max(value, BEDROCK_OPENAI_COMPAT_MIN_MAX_TOKENS) + if isinstance(value, int) and self._requires_min_max_tokens(model) + else value + ) if param == "stream": optional_params["stream"] = value if param == "stop": diff --git a/litellm/llms/custom_httpx/http_handler.py b/litellm/llms/custom_httpx/http_handler.py index 05dff0cb9d8..6b90394043f 100644 --- a/litellm/llms/custom_httpx/http_handler.py +++ b/litellm/llms/custom_httpx/http_handler.py @@ -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) diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 477d10a3cbd..8ceb9a0fea9 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -5906,7 +5906,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 +6287,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: diff --git a/litellm/llms/fireworks_ai/cache_pricing.py b/litellm/llms/fireworks_ai/cache_pricing.py new file mode 100644 index 00000000000..f5e49cad01a --- /dev/null +++ b/litellm/llms/fireworks_ai/cache_pricing.py @@ -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 + ), + }, + }, + ) diff --git a/litellm/llms/fireworks_ai/cost_calculator.py b/litellm/llms/fireworks_ai/cost_calculator.py index 1795a700d25..4b6ca7c9896 100644 --- a/litellm/llms/fireworks_ai/cost_calculator.py +++ b/litellm/llms/fireworks_ai/cost_calculator.py @@ -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, diff --git a/litellm/llms/meta/realtime/transformation.py b/litellm/llms/meta/realtime/transformation.py index 1b8943f0cee..442c79255af 100644 --- a/litellm/llms/meta/realtime/transformation.py +++ b/litellm/llms/meta/realtime/transformation.py @@ -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:] diff --git a/litellm/llms/vertex_ai/audio_transcription/realtime_backend.py b/litellm/llms/vertex_ai/audio_transcription/realtime_backend.py new file mode 100644 index 00000000000..859a883463c --- /dev/null +++ b/litellm/llms/vertex_ai/audio_transcription/realtime_backend.py @@ -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)) diff --git a/litellm/llms/vertex_ai/audio_transcription/realtime_transformation.py b/litellm/llms/vertex_ai/audio_transcription/realtime_transformation.py new file mode 100644 index 00000000000..ac23901accb --- /dev/null +++ b/litellm/llms/vertex_ai/audio_transcription/realtime_transformation.py @@ -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] diff --git a/litellm/llms/vertex_ai/audio_transcription/transformation.py b/litellm/llms/vertex_ai/audio_transcription/transformation.py index db3504c9a6a..b1284e15def 100644 --- a/litellm/llms/vertex_ai/audio_transcription/transformation.py +++ b/litellm/llms/vertex_ai/audio_transcription/transformation.py @@ -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: diff --git a/litellm/llms/vertex_ai/realtime/transformation.py b/litellm/llms/vertex_ai/realtime/transformation.py index fe59034c27b..9fed6d52f0e 100644 --- a/litellm/llms/vertex_ai/realtime/transformation.py +++ b/litellm/llms/vertex_ai/realtime/transformation.py @@ -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), + ) diff --git a/litellm/llms/xai/audio_transcription/__init__.py b/litellm/llms/xai/audio_transcription/__init__.py new file mode 100644 index 00000000000..c7910cf1f6b --- /dev/null +++ b/litellm/llms/xai/audio_transcription/__init__.py @@ -0,0 +1,3 @@ +from .transformation import XAIAudioTranscriptionConfig + +__all__ = ["XAIAudioTranscriptionConfig"] diff --git a/litellm/llms/xai/audio_transcription/transformation.py b/litellm/llms/xai/audio_transcription/transformation.py new file mode 100644 index 00000000000..feeabed0d9c --- /dev/null +++ b/litellm/llms/xai/audio_transcription/transformation.py @@ -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}"} diff --git a/litellm/main.py b/litellm/main.py index ac8fa507728..38184db1d10 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -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 diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 30b08e54410..f8c585873de 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -19244,7 +19244,20 @@ "supports_function_calling": true, "supports_anthropic_thinking_payload": true, "supports_prompt_caching": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_reasoning": true, + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_vision": true, + "supports_audio_input": true, + "supports_video_input": true }, "databricks/databricks-gemini-2-5-pro": { "cache_creation_input_token_cost": 1.24999e-06, @@ -19265,7 +19278,21 @@ "supports_function_calling": true, "supports_anthropic_thinking_payload": true, "supports_prompt_caching": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "deprecation_date": "2026-10-02", + "supports_reasoning": true, + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_vision": true, + "supports_audio_input": true, + "supports_video_input": true }, "databricks/databricks-gemini-3-1-flash-lite": { "cache_creation_input_token_cost": 3.1248e-07, @@ -19285,7 +19312,19 @@ "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_function_calling": true, "supports_prompt_caching": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_vision": true, + "supports_audio_input": true, + "supports_video_input": true }, "databricks/databricks-gemini-3-1-flash-image": { "litellm_provider": "databricks", @@ -19347,7 +19386,20 @@ "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_function_calling": true, "supports_prompt_caching": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_reasoning": true, + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_vision": true, + "supports_audio_input": true, + "supports_video_input": true }, "databricks/databricks-gemini-3-flash": { "cache_creation_input_token_cost": 6.2503e-07, @@ -19367,7 +19419,19 @@ "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_function_calling": true, "supports_prompt_caching": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_vision": true, + "supports_audio_input": true, + "supports_video_input": true }, "databricks/databricks-gemini-3-pro": { "cache_creation_input_token_cost": 2.49998e-06, @@ -21433,7 +21497,11 @@ "mode": "chat", "supports_tool_choice": true, "supports_function_calling": true, - "supports_image_size": false + "supports_image_size": false, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_vision": true, + "supports_audio_input": true }, "deepinfra/google/gemini-2.5-pro": { "max_tokens": 1000000, @@ -21444,7 +21512,11 @@ "litellm_provider": "deepinfra", "mode": "chat", "supports_tool_choice": true, - "supports_function_calling": true + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_vision": true, + "supports_audio_input": true }, "deepinfra/google/gemma-3-12b-it": { "max_tokens": 131072, @@ -29340,26 +29412,6 @@ "supports_parallel_function_calling": true, "supports_vision": true }, - "github_copilot/gemini-2.5-pro": { - "litellm_provider": "github_copilot", - "max_input_tokens": 128000, - "max_output_tokens": 64000, - "max_tokens": 64000, - "mode": "chat", - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_vision": true - }, - "github_copilot/gemini-3-pro-preview": { - "litellm_provider": "github_copilot", - "max_input_tokens": 128000, - "max_output_tokens": 64000, - "max_tokens": 64000, - "mode": "chat", - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_vision": true - }, "github_copilot/gpt-3.5-turbo": { "litellm_provider": "github_copilot", "max_input_tokens": 16384, @@ -30014,17 +30066,6 @@ "output_cost_per_token": 8.8e-07, "supports_function_calling": true }, - "gmi/google/gemini-3-pro-preview": { - "input_cost_per_token": 2e-06, - "litellm_provider": "gmi", - "max_input_tokens": 1048576, - "max_output_tokens": 65536, - "max_tokens": 65536, - "mode": "chat", - "output_cost_per_token": 1.2e-05, - "supports_function_calling": true, - "supports_vision": true - }, "gmi/google/gemini-3-flash-preview": { "input_cost_per_token": 5e-07, "litellm_provider": "gmi", @@ -30034,7 +30075,8 @@ "mode": "chat", "output_cost_per_token": 3e-06, "supports_function_calling": true, - "supports_vision": true + "supports_vision": true, + "supports_system_messages": true }, "gmi/moonshotai/Kimi-K2-Thinking": { "input_cost_per_token": 8e-07, @@ -40163,7 +40205,12 @@ "supports_response_schema": true, "supports_vision": true, "supports_native_streaming": true, - "supports_image_size": false + "supports_image_size": false, + "supports_reasoning": true, + "supports_system_messages": true, + "supports_pdf_input": true, + "supports_audio_input": true, + "supports_video_input": true }, "oci/google.gemini-2.5-pro": { "input_cost_per_token": 1.25e-06, @@ -40177,7 +40224,12 @@ "supports_function_calling": true, "supports_response_schema": true, "supports_vision": true, - "supports_native_streaming": true + "supports_native_streaming": true, + "supports_reasoning": true, + "supports_system_messages": true, + "supports_pdf_input": true, + "supports_audio_input": true, + "supports_video_input": true }, "oci/google.gemini-2.5-flash-lite": { "input_cost_per_token": 7.5e-08, @@ -40192,7 +40244,12 @@ "supports_response_schema": true, "supports_vision": true, "supports_native_streaming": true, - "supports_image_size": false + "supports_image_size": false, + "supports_reasoning": false, + "supports_system_messages": true, + "supports_pdf_input": true, + "supports_audio_input": true, + "supports_video_input": true }, "oci/cohere.command-a-vision": { "input_cost_per_token": 1.56e-06, @@ -41435,7 +41492,7 @@ "max_tokens": 65535, "mode": "chat", "output_cost_per_token": 2.5e-06, - "supports_audio_output": true, + "supports_audio_output": false, "supports_function_calling": true, "supports_response_schema": true, "supports_system_messages": true, @@ -41449,7 +41506,8 @@ "supports_audio_input": true, "supports_pdf_input": true, "supports_reasoning": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/google/gemini-2.5-pro": { "cache_creation_input_token_cost": 3.75e-07, @@ -41462,7 +41520,7 @@ "max_tokens": 65536, "mode": "chat", "output_cost_per_token": 1e-05, - "supports_audio_output": true, + "supports_audio_output": false, "supports_function_calling": true, "supports_response_schema": true, "supports_system_messages": true, @@ -41478,7 +41536,8 @@ "supports_audio_input": true, "supports_pdf_input": true, "supports_reasoning": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/google/gemini-3-pro-preview": { "cache_read_input_token_cost": 2e-07, @@ -41563,7 +41622,8 @@ "supports_url_context": true, "supports_vision": true, "supports_web_search": false, - "tpm": 800000 + "tpm": 800000, + "supports_video_input": true }, "openrouter/google/gemini-3.1-flash-lite-preview": { "cache_creation_input_token_cost": 8.33333333333333e-08, @@ -41690,7 +41750,8 @@ "supports_system_messages": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/gryphe/mythomax-l2-13b": { "input_cost_per_token": 8e-08, @@ -44413,12 +44474,16 @@ "output_cost_per_token": 1.2e-05, "litellm_provider": "replicate", "mode": "chat", - "supports_function_calling": true, - "supports_parallel_function_calling": true, + "supports_function_calling": false, + "supports_parallel_function_calling": false, "supports_vision": true, "supports_system_messages": true, - "supports_tool_choice": true, - "supports_response_schema": true + "supports_tool_choice": false, + "supports_response_schema": false, + "input_cost_per_token_above_200k_tokens": 4e-06, + "output_cost_per_token_above_200k_tokens": 1.8e-05, + "supports_audio_input": true, + "supports_video_input": true }, "replicate/anthropic/claude-4.5-sonnet": { "input_cost_per_token": 3e-06, @@ -44487,17 +44552,19 @@ "supports_response_schema": true }, "replicate/google/gemini-2.5-flash": { - "input_cost_per_token": 2.5e-06, + "input_cost_per_token": 3e-07, "output_cost_per_token": 2.5e-06, "litellm_provider": "replicate", "mode": "chat", - "supports_function_calling": true, - "supports_parallel_function_calling": true, + "supports_function_calling": false, + "supports_parallel_function_calling": false, "supports_vision": true, "supports_system_messages": true, - "supports_tool_choice": true, - "supports_response_schema": true, - "supports_image_size": false + "supports_tool_choice": false, + "supports_response_schema": false, + "supports_image_size": false, + "supports_reasoning": true, + "supports_video_input": true }, "replicate/openai/gpt-oss-120b": { "input_cost_per_token": 1.8e-07, @@ -48076,10 +48143,15 @@ "supports_function_calling": true, "supports_tool_choice": true, "supports_response_schema": true, - "supports_image_size": false + "supports_image_size": false, + "cache_read_input_token_cost": 3e-08, + "supports_reasoning": true, + "supports_pdf_input": true, + "supports_web_search": true, + "supports_prompt_caching": true }, "vercel_ai_gateway/google/gemini-2.5-pro": { - "input_cost_per_token": 2.5e-06, + "input_cost_per_token": 1.25e-06, "litellm_provider": "vercel_ai_gateway", "max_input_tokens": 1048576, "max_output_tokens": 65536, @@ -48089,7 +48161,15 @@ "supports_vision": true, "supports_function_calling": true, "supports_tool_choice": true, - "supports_response_schema": true + "supports_response_schema": true, + "input_cost_per_token_above_200k_tokens": 2.5e-06, + "output_cost_per_token_above_200k_tokens": 1.5e-05, + "cache_read_input_token_cost": 1.25e-07, + "cache_read_input_token_cost_above_200k_tokens": 2.5e-07, + "supports_reasoning": true, + "supports_pdf_input": true, + "supports_web_search": true, + "supports_prompt_caching": true }, "vercel_ai_gateway/google/gemini-embedding-001": { "input_cost_per_token": 1.5e-07, @@ -48825,7 +48905,8 @@ "mode": "audio_transcription", "source": "https://cloud.google.com/speech-to-text/pricing", "supported_endpoints": [ - "/v1/audio/transcriptions" + "/v1/audio/transcriptions", + "/v1/realtime" ] }, "vertex_ai/claude-3-5-haiku": { @@ -62224,7 +62305,8 @@ "supports_response_schema": true, "supports_reasoning": true, "supports_vision": true, - "source": "https://deepinfra.com/pricing" + "source": "https://deepinfra.com/pricing", + "supports_audio_input": true }, "deepinfra/XiaomiMiMo/MiMo-V2.5": { "max_tokens": 262144, @@ -62524,7 +62606,8 @@ "supports_response_schema": true, "supports_reasoning": true, "supports_vision": true, - "source": "https://deepinfra.com/pricing" + "source": "https://deepinfra.com/pricing", + "supports_audio_input": true }, "deepinfra/google/gemini-3.7-flash": { "max_tokens": 1000000, @@ -62538,7 +62621,8 @@ "supports_response_schema": true, "supports_reasoning": true, "supports_vision": true, - "source": "https://deepinfra.com/pricing" + "source": "https://deepinfra.com/pricing", + "supports_audio_input": true }, "deepinfra/inclusionAI/Ling-3.0-flash": { "max_tokens": 131072, @@ -62970,7 +63054,8 @@ "supports_response_schema": true, "supports_reasoning": true, "supports_vision": true, - "source": "https://deepinfra.com/pricing" + "source": "https://deepinfra.com/pricing", + "supports_audio_input": true }, "deepinfra/XiaomiMiMo/MiMo-V2.5-Pro": { "max_tokens": 1048576, @@ -63375,6 +63460,34 @@ "video" ] }, + "xai/grok-voice-transcribe-1.0": { + "input_cost_per_second": 2.778e-05, + "litellm_provider": "xai", + "metadata": { + "calculation": "$0.10/3600 seconds = $0.00002778 per second", + "original_pricing_per_hour": 0.1 + }, + "mode": "audio_transcription", + "output_cost_per_second": 0.0, + "source": "https://docs.x.ai/developers/pricing", + "supported_endpoints": [ + "/v1/audio/transcriptions" + ] + }, + "xai/grok-voice-transcribe-2.0": { + "input_cost_per_second": 2.778e-05, + "litellm_provider": "xai", + "metadata": { + "calculation": "$0.10/3600 seconds = $0.00002778 per second", + "original_pricing_per_hour": 0.1 + }, + "mode": "audio_transcription", + "output_cost_per_second": 0.0, + "source": "https://docs.x.ai/developers/pricing", + "supported_endpoints": [ + "/v1/audio/transcriptions" + ] + }, "low/1024-x-1024/grok-imagine-image-2.0": { "input_cost_per_image": 0.04, "litellm_provider": "xai", @@ -65365,7 +65478,8 @@ "deprecation_date": "2026-10-20", "input_cost_per_audio_token": 3e-07, "supports_prompt_caching": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/google/gemini-3.5-flash": { "cache_creation_input_token_cost": 8.33333333333333e-08, @@ -65388,7 +65502,8 @@ "cache_read_input_token_cost": 1.5e-07, "input_cost_per_audio_token": 3e-06, "supports_prompt_caching": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/google/gemini-3.5-flash-lite": { "cache_creation_input_token_cost": 8.33333333333333e-08, @@ -65411,7 +65526,8 @@ "cache_read_input_token_cost": 3e-08, "input_cost_per_audio_token": 3e-07, "supports_prompt_caching": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/google/gemini-3.6-flash": { "cache_creation_input_token_cost": 4.16666666666667e-08, @@ -65434,7 +65550,8 @@ "cache_read_input_token_cost": 7.5e-08, "input_cost_per_audio_token": 7.5e-07, "supports_prompt_caching": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/google/gemini-3.7-flash": { "cache_creation_input_token_cost": 4.16666666666667e-08, @@ -65457,7 +65574,8 @@ "cache_read_input_token_cost": 7.5e-08, "input_cost_per_audio_token": 7.5e-07, "supports_prompt_caching": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/google/gemini-3.8-flash": { "cache_creation_input_token_cost": 4.16666666666667e-08, @@ -65480,7 +65598,8 @@ "cache_read_input_token_cost": 7.5e-08, "input_cost_per_audio_token": 7.5e-07, "supports_prompt_caching": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/openai/gpt-4o-mini": { "input_cost_per_token": 1.5e-07, @@ -67108,7 +67227,8 @@ "supports_pdf_input": true, "supports_audio_input": true, "supports_prompt_caching": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/qwen/qwen3-max-thinking": { "input_cost_per_token": 7.8e-07, @@ -71974,7 +72094,8 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/google/gemini-2.5-flash:batch": { "cache_read_input_audio_token_cost": 1e-07, @@ -71997,7 +72118,8 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/google/gemini-2.5-pro:batch": { "cache_read_input_audio_token_cost": 1.25e-07, @@ -72023,7 +72145,8 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/google/gemini-3-flash-preview:batch": { "input_cost_per_audio_token": 5e-07, @@ -72043,7 +72166,8 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/google/gemini-3.1-flash-lite:batch": { "cache_read_input_audio_token_cost": 2.5e-08, @@ -72065,7 +72189,8 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/google/gemini-3.1-pro-preview:batch": { "input_cost_per_audio_token": 1e-06, @@ -72087,7 +72212,8 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/google/gemini-3.5-flash-lite:batch": { "cache_read_input_audio_token_cost": 1.5e-08, @@ -72109,7 +72235,8 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/google/gemini-3.5-flash:batch": { "cache_read_input_audio_token_cost": 1.5e-07, @@ -72131,7 +72258,8 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/google/gemini-3.6-flash:batch": { "cache_creation_input_token_cost": 4.16666666666667e-08, @@ -72154,7 +72282,8 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/google/gemini-3.7-flash:batch": { "cache_creation_input_token_cost": 4.16666666666667e-08, @@ -72177,7 +72306,8 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/google/gemini-3.8-flash:batch": { "cache_creation_input_token_cost": 4.16666666666667e-08, @@ -72200,7 +72330,8 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/ibm-granite/granite-4.0-h-micro": { "input_cost_per_token": 1.7e-08, diff --git a/litellm/proxy/_experimental/mcp_server/AGENTS.md b/litellm/proxy/_experimental/mcp_server/AGENTS.md index a35e5d1ce10..d9e0bfa3589 100644 --- a/litellm/proxy/_experimental/mcp_server/AGENTS.md +++ b/litellm/proxy/_experimental/mcp_server/AGENTS.md @@ -1,6 +1,6 @@ # Experimental MCP Server Change Guidelines -Read @../../../../CLAUDE.md and @CLAUDE.md before changing this package. +Read @../../../../AGENTS.md before changing this package. This directory owns the proxy-hosted MCP server implementation. Keep changes inside the module that owns the behavior, and only reach outside this package @@ -14,13 +14,13 @@ Respect the current package boundaries: ```text litellm/proxy/_experimental/mcp_server/ AGENTS.md - CLAUDE.md server.py # ASGI/MCP route handling, sessions, tool calls [PR7: 7-arm only — move BYOK/OAuth pre-fetch into resolver] mcp_server_manager.py # upstream server registry, clients, tool routing [PR7: _create_mcp_client swaps resolve_mcp_auth -> resolve_credentials] auth/ user_api_key_auth_mcp.py # LiteLLM admission auth and MCP request headers token_exchange.py # OAuth token exchange handling [unchanged; V1TokenExchangeAdapter delegates here] litellm_auth_handler.py # authenticated-user adapter for MCP sessions + client_allowlist.py # gateway-level client application allowlist (mcp_allowed_clients); leaf module, no litellm.proxy imports outbound_credentials/ # NEW — typed upstream-credential resolution (resolve_credentials + arms) __init__.py # public surface: resolve_credentials, the configs, CredError result.py # Ok | Error union (pure stdlib) @@ -67,8 +67,10 @@ module materially harder to understand. auth, SSE, streamable HTTP, and stdio as separate flows. Do not collapse them behind a single generic branch unless tests prove every mode still behaves correctly. -- Be especially careful with legacy `delegate_auth_to_upstream: true`. The local - `CLAUDE.md` explains its admitted replacement and public discovery contract. +- Be especially careful with legacy `delegate_auth_to_upstream: true`. `auth_type: oauth2` + with `delegate_auth_to_upstream: true` is deprecated: LiteLLM admission is required + for matching MCP routes. Use `auth_type: oauth_delegate` for client-forwarded OAuth. + OAuth discovery endpoints stay public so clients can start the RFC 9728 flow. - Keep database-backed fields in sync across migrations, typed models under `litellm/types/mcp.py` or `litellm/types/mcp_server/`, config loading, this package, and dashboard state when the field is user-visible. diff --git a/litellm/proxy/_experimental/mcp_server/CLAUDE.md b/litellm/proxy/_experimental/mcp_server/CLAUDE.md deleted file mode 100644 index 7f8d06b4570..00000000000 --- a/litellm/proxy/_experimental/mcp_server/CLAUDE.md +++ /dev/null @@ -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 diff --git a/litellm/proxy/_experimental/mcp_server/client_allowlist.py b/litellm/proxy/_experimental/mcp_server/client_allowlist.py new file mode 100644 index 00000000000..524642600b3 --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/client_allowlist.py @@ -0,0 +1,171 @@ +""" +Gateway-level allowlist of MCP client applications (``general_settings.mcp_allowed_clients``). + +Each entry pairs an admin-chosen ``alias`` (shown in the dashboard and logs) with the ``value`` that +identifies the client. Only the value is compared, exactly and case-sensitively. +A caller that authenticated with a JWT is identified by the claim named in +``litellm_jwtauth.mcp_client_id_jwt_field``, a value asserted by the identity provider. +Every other caller is identified by the header named in ``general_settings.mcp_client_id_header``, +which the client picks itself, so that source is a policy control rather than a security boundary. +While the allowlist is set, a caller with no usable identity source is rejected. +""" + +from collections.abc import Mapping +from dataclasses import dataclass +from types import MappingProxyType +from typing import Final, Literal + +from pydantic import TypeAdapter, ValidationError +from typing_extensions import ReadOnly, TypedDict + +from litellm._logging import verbose_logger +from litellm.litellm_core_utils.dot_notation_indexing import get_nested_value +from litellm.types.mcp import MCPAllowedClient + +MCP_ALLOWED_CLIENTS_SETTING: Final = "mcp_allowed_clients" +MCP_CLIENT_ID_HEADER_SETTING: Final = "mcp_client_id_header" +MCP_CLIENT_ID_JWT_FIELD_SETTING: Final = "mcp_client_id_jwt_field" +_JWT_AUTH_SETTING: Final = "litellm_jwtauth" + +_ALLOWED_CLIENTS_ADAPTER: Final[TypeAdapter[list[MCPAllowedClient]]] = TypeAdapter(list[MCPAllowedClient]) +_OPTIONAL_NAME_ADAPTER: Final[TypeAdapter[str | None]] = TypeAdapter(str | None) +_OPTIONAL_MAPPING_ADAPTER: Final[TypeAdapter[dict[str, object] | None]] = TypeAdapter(dict[str, object] | None) +_NOBODY: Final[Mapping[str, str]] = MappingProxyType({}) + + +class MCPClientForbiddenBody(TypedDict): + error: ReadOnly[Literal["Forbidden"]] + details: ReadOnly[str] + + +@dataclass(frozen=True, slots=True) +class MCPClientAllowlist: + """``aliases_by_value`` maps each admitted identity value to the alias the admin gave it.""" + + aliases_by_value: Mapping[str, str] + jwt_field: str | None + header: str | None + + +@dataclass(frozen=True, slots=True) +class MCPClientIdentity: + client_id: str + source: Literal["jwt", "header"] + source_name: str + + @property + def description(self) -> str: + return f"'{self.client_id}' (from {'JWT claim' if self.source == 'jwt' else 'header'} '{self.source_name}')" + + +@dataclass(frozen=True, slots=True) +class MCPClientRejection: + details: str + + @property + def response_body(self) -> MCPClientForbiddenBody: + body: Final[MCPClientForbiddenBody] = {"error": "Forbidden", "details": self.details} + return body + + +def _unidentified_rejection(reason: str) -> MCPClientRejection: + return MCPClientRejection( + details=f"{reason} This gateway only admits client applications listed in {MCP_ALLOWED_CLIENTS_SETTING}." + ) + + +def parse_allowed_mcp_clients(raw_setting: object) -> Mapping[str, str] | None: + """Value-to-alias mapping; None when the setting is absent (not enforced). A malformed setting admits nobody.""" + if raw_setting is None: + return None + try: + clients: Final = _ALLOWED_CLIENTS_ADAPTER.validate_python(raw_setting) + except ValidationError: + verbose_logger.warning( + "%s is not a list of {alias, value} entries (%r); rejecting every MCP client until it is fixed", + MCP_ALLOWED_CLIENTS_SETTING, + raw_setting, + ) + return _NOBODY + return MappingProxyType({client.value: client.alias for client in clients}) + + +def _parse_optional_name(setting_name: str, raw_setting: object) -> str | None: + try: + name: Final = _OPTIONAL_NAME_ADAPTER.validate_python(raw_setting) + except ValidationError: + verbose_logger.warning("%s is not a string (%r); ignoring it", setting_name, raw_setting) + return None + return name or None + + +def _jwt_field_from_general_settings(general_settings: Mapping[str, object]) -> str | None: + try: + jwt_auth: Final = _OPTIONAL_MAPPING_ADAPTER.validate_python(general_settings.get(_JWT_AUTH_SETTING)) + except ValidationError: + return None + if jwt_auth is None: + return None + return _parse_optional_name( + f"{_JWT_AUTH_SETTING}.{MCP_CLIENT_ID_JWT_FIELD_SETTING}", jwt_auth.get(MCP_CLIENT_ID_JWT_FIELD_SETTING) + ) + + +def load_mcp_client_allowlist(general_settings: Mapping[str, object]) -> MCPClientAllowlist | None: + """None when ``mcp_allowed_clients`` is unset, which admits every client.""" + allowed_clients: Final = parse_allowed_mcp_clients(general_settings.get(MCP_ALLOWED_CLIENTS_SETTING)) + if allowed_clients is None: + return None + header: Final = _parse_optional_name( + MCP_CLIENT_ID_HEADER_SETTING, general_settings.get(MCP_CLIENT_ID_HEADER_SETTING) + ) + return MCPClientAllowlist( + aliases_by_value=allowed_clients, + jwt_field=_jwt_field_from_general_settings(general_settings), + header=header.lower() if header is not None else None, + ) + + +def resolve_mcp_client_identity( + allowlist: MCPClientAllowlist, + jwt_claims: Mapping[str, object] | None, + headers: Mapping[str, str], +) -> MCPClientIdentity | MCPClientRejection: + """A JWT caller is identified by its configured claim alone, so a header can never override the IdP.""" + if jwt_claims is not None and allowlist.jwt_field is not None: + claim: Final[object] = get_nested_value(data=jwt_claims, key_path=allowlist.jwt_field) + if isinstance(claim, str) and claim: + return MCPClientIdentity(client_id=claim, source="jwt", source_name=allowlist.jwt_field) + return _unidentified_rejection( + f"The JWT presented has no '{allowlist.jwt_field}' claim naming the client application." + ) + if allowlist.header is None: + configured: Final = ( + f"litellm_jwtauth.{MCP_CLIENT_ID_JWT_FIELD_SETTING} for JWT callers or {MCP_CLIENT_ID_HEADER_SETTING}" + ) + return _unidentified_rejection( + f"No client identity source is configured for this request; set {configured} in general_settings." + ) + header_value: Final = headers.get(allowlist.header) + if header_value: + return MCPClientIdentity(client_id=header_value, source="header", source_name=allowlist.header) + return _unidentified_rejection(f"The request has no '{allowlist.header}' header naming the client application.") + + +def check_mcp_client_allowed( + allowlist: MCPClientAllowlist | None, + jwt_claims: Mapping[str, object] | None, + headers: Mapping[str, str], +) -> MCPClientRejection | None: + if allowlist is None: + return None + identity: Final = resolve_mcp_client_identity(allowlist, jwt_claims, headers) + if isinstance(identity, MCPClientRejection): + return identity + alias: Final = allowlist.aliases_by_value.get(identity.client_id) + if alias is None: + return MCPClientRejection( + details=f"MCP client {identity.description} is not listed in this gateway's {MCP_ALLOWED_CLIENTS_SETTING}." + ) + verbose_logger.debug("Admitted MCP client '%s' identified as %s", alias, identity.description) + return None diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index 6a0ab5bdec5..9d895d755dd 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -193,6 +193,7 @@ if MCP_AVAILABLE: filter_tools_by_allowed_tools, filter_tools_by_key_team_permissions, fire_mcp_tool_call_failure_logging, + reject_disallowed_mcp_client, ) ######################################################## @@ -875,6 +876,7 @@ if MCP_AVAILABLE: MCPRequestHandler, ) + reject_disallowed_mcp_client(request.headers, user_api_key_dict) try: mcp_server_name = _as_query_str(mcp_server_name) toolset_name = _as_query_str(toolset_name) @@ -1078,6 +1080,7 @@ if MCP_AVAILABLE: proxy_logging_obj, ) + reject_disallowed_mcp_client(request.headers, user_api_key_dict) try: user_api_key_dict = await acting_user_auth(user_api_key_dict) data = await request.json() diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index a7ca2775fb1..73366b701d2 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -47,6 +47,11 @@ from litellm.proxy._experimental.mcp_server.byok_credential_cache import ( cache_byok_credential, get_cached_byok_credential, ) +from litellm.proxy._experimental.mcp_server.client_allowlist import ( + MCPClientAllowlist, + check_mcp_client_allowed, + load_mcp_client_allowlist, +) from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( get_request_base_url, ) @@ -72,6 +77,7 @@ from litellm.proxy._experimental.mcp_server.oauth_utils import ( get_route_relative_request_path, well_known_root_suffix, ) +from litellm.proxy._experimental.mcp_server.ui_session_utils import is_ui_session_credential from litellm.proxy._experimental.mcp_server.utils import ( LITELLM_MCP_SERVER_DESCRIPTION, LITELLM_MCP_SERVER_NAME, @@ -3701,6 +3707,25 @@ if MCP_AVAILABLE: mcp_servers_from_path = [servers_and_path] return mcp_servers_from_path + def _load_mcp_client_allowlist() -> MCPClientAllowlist | None: + from litellm.proxy.proxy_server import general_settings + + return load_mcp_client_allowlist(general_settings) + + def reject_disallowed_mcp_client(headers: Mapping[str, str], user_api_key_auth: UserAPIKeyAuth | None) -> None: + """Gate every MCP tool surface on ``mcp_allowed_clients``; the dashboard's own session is not a client app.""" + if user_api_key_auth is not None and is_ui_session_credential(user_api_key_auth): + return + rejection: Final = check_mcp_client_allowed( + allowlist=_load_mcp_client_allowlist(), + jwt_claims=user_api_key_auth.jwt_claims if user_api_key_auth is not None else None, + headers=headers, + ) + if rejection is None: + return + verbose_logger.warning("Rejected MCP request from a disallowed client application: %s", rejection.details) + raise HTTPException(status_code=403, detail=rejection.response_body) + async def extract_mcp_auth_context(scope, path): """ Extracts mcp_servers from the path and processes the MCP request for auth context. @@ -4537,6 +4562,7 @@ if MCP_AVAILABLE: oauth2_headers, raw_headers, ) = await extract_mcp_auth_context(scope, path) + reject_disallowed_mcp_client(StarletteRequest(scope).headers, user_api_key_auth) scoped_server_endpoint: Final = len(_get_mcp_servers_in_path(path) or []) == 1 # Extract client IP for MCP access control @@ -4865,6 +4891,7 @@ if MCP_AVAILABLE: oauth2_headers, raw_headers, ) = await extract_mcp_auth_context(scope, path) + reject_disallowed_mcp_client(StarletteRequest(scope).headers, user_api_key_auth) scoped_server_endpoint: Final = len(_get_mcp_servers_in_path(path) or []) == 1 # Extract client IP for MCP access control diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 0097d4b6e92..8d8fada6af7 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -34,6 +34,7 @@ from litellm.types.llms.openai import ( ResponsesAPIResponse, ) from litellm.types.mcp import ( + MCPAllowedClient, MCPAuth, MCPAuthType, MCPCredentials, @@ -897,6 +898,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 @@ -2902,6 +2906,14 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase): None, description="Custom CIDR ranges that define internal/private networks for MCP access control. When set, only these ranges are treated as internal. Defaults to RFC 1918 private ranges (10.0.0.0/8, 172.16.0.0/12, 192.168.0.0/16, 127.0.0.0/8).", ) + mcp_allowed_clients: list[MCPAllowedClient] | None = Field( + None, + description="MCP client applications admitted by the gateway, each an {alias, value} pair where alias is the name shown in the dashboard and logs and value is the identity that must match exactly. When set, every MCP request must carry a client identity equal to one of the values: a JWT caller is identified by the claim named in litellm_jwtauth.mcp_client_id_jwt_field, any other caller by the header named in mcp_client_id_header. A request with no resolvable identity, or an unlisted one, is rejected with 403. Unset means every client is admitted.", + ) + mcp_client_id_header: str | None = Field( + None, + description="Request header whose value names the calling MCP client application (for example 'x-mcp-client') for callers that did not authenticate with a JWT, used only while mcp_allowed_clients is set. The client picks this value itself, so it is a policy control rather than a security boundary; prefer litellm_jwtauth.mcp_client_id_jwt_field where callers use JWTs.", + ) mcp_trusted_proxy_ranges: list[str] | None = Field( None, description="CIDR ranges of trusted reverse proxies. When set, X-Forwarded-For and X-Forwarded-* origin headers are only trusted from these IPs.", @@ -5118,6 +5130,15 @@ class LiteLLM_JWTAuth(LiteLLMPydanticObjectBase): "then agent_name, and the request is rejected when it matches neither." ), ) + mcp_client_id_jwt_field: str | None = Field( + default=None, + description=( + "The field in the JWT token that identifies the MCP client application (harness) making the request, " + "e.g. 'azp' or 'client_id'. Supports dot notation. Only consulted while general_settings.mcp_allowed_clients " + "is set: the claim value must be listed there or the MCP request is rejected with 403. Distinct from " + "agent_id_jwt_field, which identifies an AI agent rather than the client software." + ), + ) public_key_ttl: float = 600 public_key_stale_ttl: float = Field( default=DEFAULT_JWKS_STALE_TTL, diff --git a/litellm/proxy/hooks/batch_rate_limiter.py b/litellm/proxy/hooks/batch_rate_limiter.py index ab6e10ca76b..a5b6cabf519 100644 --- a/litellm/proxy/hooks/batch_rate_limiter.py +++ b/litellm/proxy/hooks/batch_rate_limiter.py @@ -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"] diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index cdec5922fff..a6b00be1091 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -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"] diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index b03f1e4348c..9a973755894 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -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 diff --git a/litellm/proxy/management_endpoints/team_admin_field_permissions.py b/litellm/proxy/management_endpoints/team_admin_field_permissions.py index 56d455494c6..6038775d96b 100644 --- a/litellm/proxy/management_endpoints/team_admin_field_permissions.py +++ b/litellm/proxy/management_endpoints/team_admin_field_permissions.py @@ -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) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 009b8cb728f..7791f034fba 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -17249,6 +17249,8 @@ _GENERAL_SETTINGS_CONFIG_LIST_FIELD_TYPES: Final[Mapping[str, str]] = MappingPro "maximum_spend_logs_cleanup_run_budget": "String", "maximum_spend_logs_cleanup_batch_timeout": "String", "mcp_internal_ip_ranges": "List", + "mcp_allowed_clients": "TypedDictionary", + "mcp_client_id_header": "String", "mcp_trusted_proxy_ranges": "List", "mcp_xff_num_trusted_hops": "Integer", "always_include_stream_usage": "Boolean", diff --git a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py index fd160636d46..75431383fbd 100644 --- a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py +++ b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py @@ -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)}." ) }, ) diff --git a/litellm/realtime_api/main.py b/litellm/realtime_api/main.py index d70295e9a2a..0e83edab5e1 100644 --- a/litellm/realtime_api/main.py +++ b/litellm/realtime_api/main.py @@ -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), diff --git a/litellm/rust_bridge/settings.py b/litellm/rust_bridge/settings.py new file mode 100644 index 00000000000..e170f93b198 --- /dev/null +++ b/litellm/rust_bridge/settings.py @@ -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(), + ) diff --git a/litellm/types/integrations/websearch_interception.py b/litellm/types/integrations/websearch_interception.py index 7926b9eee0a..a7c2e7f2315 100644 --- a/litellm/types/integrations/websearch_interception.py +++ b/litellm/types/integrations/websearch_interception.py @@ -2,11 +2,15 @@ Type definitions for WebSearch Interception integration. """ -from typing import Literal, TypedDict +from dataclasses import dataclass +from typing import TYPE_CHECKING, Literal, TypeAlias, TypedDict from pydantic import BaseModel from typing_extensions import ReadOnly +if TYPE_CHECKING: + from litellm.llms.base_llm.search.transformation import SearchResponse + class AnthropicSearchQuery(BaseModel): """``input`` of an Anthropic ``server_tool_use`` block for a web search.""" @@ -27,6 +31,31 @@ class AnthropicServerToolUseBlock(BaseModel): input: AnthropicSearchQuery +WebSearchToolResultErrorCode: TypeAlias = Literal[ + "invalid_tool_input", + "unavailable", + "max_uses_exceeded", + "too_many_requests", + "query_too_long", + "request_too_large", +] + + +@dataclass(frozen=True, slots=True) +class SearchSucceeded: + text: str + response: "SearchResponse | None" + + +@dataclass(frozen=True, slots=True) +class SearchFailed: + error_code: WebSearchToolResultErrorCode + message: str + + +SearchOutcome: TypeAlias = SearchSucceeded | SearchFailed + + class WebSearchInterceptionConfig(TypedDict, total=False): """ Configuration parameters for WebSearchInterceptionLogger. diff --git a/litellm/types/llms/vertex_ai_speech_to_text.py b/litellm/types/llms/vertex_ai_speech_to_text.py index 8995d98385b..d07a5bbc192 100644 --- a/litellm/types/llms/vertex_ai_speech_to_text.py +++ b/litellm/types/llms/vertex_ai_speech_to_text.py @@ -1,4 +1,6 @@ -from pydantic import BaseModel +from typing import Annotated, Literal + +from pydantic import BaseModel, ConfigDict, Field from typing_extensions import TypedDict @@ -38,3 +40,66 @@ class VertexSpeechToTextResponseMetadata(BaseModel): class VertexSpeechToTextRecognizeResponse(BaseModel): results: list[VertexSpeechToTextResult] = [] metadata: VertexSpeechToTextResponseMetadata | None = None + + +class VertexSpeechStreamingConfigure(BaseModel): + model_config = ConfigDict(frozen=True) + kind: Literal["configure"] = "configure" + model: str + language_codes: tuple[str, ...] + sample_rate_hertz: int + + +class VertexSpeechStreamingFinishTurn(BaseModel): + model_config = ConfigDict(frozen=True) + kind: Literal["finish_turn"] = "finish_turn" + + +class VertexSpeechStreamingDiscardTurn(BaseModel): + model_config = ConfigDict(frozen=True) + kind: Literal["discard_turn"] = "discard_turn" + + +VertexSpeechStreamingCommandUnion = ( + VertexSpeechStreamingConfigure | VertexSpeechStreamingFinishTurn | VertexSpeechStreamingDiscardTurn +) +VertexSpeechStreamingCommand = Annotated[VertexSpeechStreamingCommandUnion, Field(discriminator="kind")] + + +class VertexSpeechStreamingResult(BaseModel): + model_config = ConfigDict(frozen=True) + transcript: str + is_final: bool + + +class VertexSpeechStreamingResponse(BaseModel): + model_config = ConfigDict(frozen=True) + kind: Literal["response"] = "response" + speech_event: Literal["none", "begin", "end"] + results: tuple[VertexSpeechStreamingResult, ...] + billed_seconds: float + + +class VertexSpeechStreamingConfigured(BaseModel): + model_config = ConfigDict(frozen=True) + kind: Literal["configured"] = "configured" + + +class VertexSpeechStreamingTurnFinished(BaseModel): + model_config = ConfigDict(frozen=True) + kind: Literal["turn_finished"] = "turn_finished" + + +class VertexSpeechStreamingTurnDiscarded(BaseModel): + model_config = ConfigDict(frozen=True) + kind: Literal["turn_discarded"] = "turn_discarded" + billed_seconds: float + + +VertexSpeechStreamingEventUnion = ( + VertexSpeechStreamingResponse + | VertexSpeechStreamingConfigured + | VertexSpeechStreamingTurnFinished + | VertexSpeechStreamingTurnDiscarded +) +VertexSpeechStreamingEvent = Annotated[VertexSpeechStreamingEventUnion, Field(discriminator="kind")] diff --git a/litellm/types/mcp.py b/litellm/types/mcp.py index c5a26c997b7..e0fd3e9a69d 100644 --- a/litellm/types/mcp.py +++ b/litellm/types/mcp.py @@ -91,6 +91,22 @@ class MCPPublicServer(BaseModel): mcp_info: dict[str, Any] | None = None +class MCPAllowedClient(BaseModel): + """One entry of `general_settings.mcp_allowed_clients`.""" + + model_config = ConfigDict(frozen=True) + + alias: str = Field( + min_length=1, + description="Human-readable name for this client application, shown in the dashboard and in gateway logs.", + ) + value: str = Field( + min_length=1, + description="Exact value of the JWT claim named in litellm_jwtauth.mcp_client_id_jwt_field, or of the " + "mcp_client_id_header header, that identifies this client application. Matched case-sensitively.", + ) + + class MCPToolSearchSettings(BaseModel): """`litellm_settings.mcp_tool_search`: how the native `mcp_tool_search` virtual tool ranks the caller's tools.""" diff --git a/litellm/utils.py b/litellm/utils.py index f2315651a53..e30e8cde86d 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -8730,6 +8730,10 @@ class ProviderConfigManager: ) return ElevenLabsAudioTranscriptionConfig() + elif litellm.LlmProviders.XAI == provider: + from litellm.llms.xai.audio_transcription.transformation import XAIAudioTranscriptionConfig + + return XAIAudioTranscriptionConfig() elif litellm.LlmProviders.OPENAI == provider: if "gpt-4o" in model: return litellm.OpenAIGPTAudioTranscriptionConfig() diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 30b08e54410..f8c585873de 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -19244,7 +19244,20 @@ "supports_function_calling": true, "supports_anthropic_thinking_payload": true, "supports_prompt_caching": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_reasoning": true, + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_vision": true, + "supports_audio_input": true, + "supports_video_input": true }, "databricks/databricks-gemini-2-5-pro": { "cache_creation_input_token_cost": 1.24999e-06, @@ -19265,7 +19278,21 @@ "supports_function_calling": true, "supports_anthropic_thinking_payload": true, "supports_prompt_caching": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "deprecation_date": "2026-10-02", + "supports_reasoning": true, + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_vision": true, + "supports_audio_input": true, + "supports_video_input": true }, "databricks/databricks-gemini-3-1-flash-lite": { "cache_creation_input_token_cost": 3.1248e-07, @@ -19285,7 +19312,19 @@ "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_function_calling": true, "supports_prompt_caching": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_vision": true, + "supports_audio_input": true, + "supports_video_input": true }, "databricks/databricks-gemini-3-1-flash-image": { "litellm_provider": "databricks", @@ -19347,7 +19386,20 @@ "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_function_calling": true, "supports_prompt_caching": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_reasoning": true, + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_vision": true, + "supports_audio_input": true, + "supports_video_input": true }, "databricks/databricks-gemini-3-flash": { "cache_creation_input_token_cost": 6.2503e-07, @@ -19367,7 +19419,19 @@ "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_function_calling": true, "supports_prompt_caching": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text" + ], + "supports_vision": true, + "supports_audio_input": true, + "supports_video_input": true }, "databricks/databricks-gemini-3-pro": { "cache_creation_input_token_cost": 2.49998e-06, @@ -21433,7 +21497,11 @@ "mode": "chat", "supports_tool_choice": true, "supports_function_calling": true, - "supports_image_size": false + "supports_image_size": false, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_vision": true, + "supports_audio_input": true }, "deepinfra/google/gemini-2.5-pro": { "max_tokens": 1000000, @@ -21444,7 +21512,11 @@ "litellm_provider": "deepinfra", "mode": "chat", "supports_tool_choice": true, - "supports_function_calling": true + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_vision": true, + "supports_audio_input": true }, "deepinfra/google/gemma-3-12b-it": { "max_tokens": 131072, @@ -29340,26 +29412,6 @@ "supports_parallel_function_calling": true, "supports_vision": true }, - "github_copilot/gemini-2.5-pro": { - "litellm_provider": "github_copilot", - "max_input_tokens": 128000, - "max_output_tokens": 64000, - "max_tokens": 64000, - "mode": "chat", - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_vision": true - }, - "github_copilot/gemini-3-pro-preview": { - "litellm_provider": "github_copilot", - "max_input_tokens": 128000, - "max_output_tokens": 64000, - "max_tokens": 64000, - "mode": "chat", - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_vision": true - }, "github_copilot/gpt-3.5-turbo": { "litellm_provider": "github_copilot", "max_input_tokens": 16384, @@ -30014,17 +30066,6 @@ "output_cost_per_token": 8.8e-07, "supports_function_calling": true }, - "gmi/google/gemini-3-pro-preview": { - "input_cost_per_token": 2e-06, - "litellm_provider": "gmi", - "max_input_tokens": 1048576, - "max_output_tokens": 65536, - "max_tokens": 65536, - "mode": "chat", - "output_cost_per_token": 1.2e-05, - "supports_function_calling": true, - "supports_vision": true - }, "gmi/google/gemini-3-flash-preview": { "input_cost_per_token": 5e-07, "litellm_provider": "gmi", @@ -30034,7 +30075,8 @@ "mode": "chat", "output_cost_per_token": 3e-06, "supports_function_calling": true, - "supports_vision": true + "supports_vision": true, + "supports_system_messages": true }, "gmi/moonshotai/Kimi-K2-Thinking": { "input_cost_per_token": 8e-07, @@ -40163,7 +40205,12 @@ "supports_response_schema": true, "supports_vision": true, "supports_native_streaming": true, - "supports_image_size": false + "supports_image_size": false, + "supports_reasoning": true, + "supports_system_messages": true, + "supports_pdf_input": true, + "supports_audio_input": true, + "supports_video_input": true }, "oci/google.gemini-2.5-pro": { "input_cost_per_token": 1.25e-06, @@ -40177,7 +40224,12 @@ "supports_function_calling": true, "supports_response_schema": true, "supports_vision": true, - "supports_native_streaming": true + "supports_native_streaming": true, + "supports_reasoning": true, + "supports_system_messages": true, + "supports_pdf_input": true, + "supports_audio_input": true, + "supports_video_input": true }, "oci/google.gemini-2.5-flash-lite": { "input_cost_per_token": 7.5e-08, @@ -40192,7 +40244,12 @@ "supports_response_schema": true, "supports_vision": true, "supports_native_streaming": true, - "supports_image_size": false + "supports_image_size": false, + "supports_reasoning": false, + "supports_system_messages": true, + "supports_pdf_input": true, + "supports_audio_input": true, + "supports_video_input": true }, "oci/cohere.command-a-vision": { "input_cost_per_token": 1.56e-06, @@ -41435,7 +41492,7 @@ "max_tokens": 65535, "mode": "chat", "output_cost_per_token": 2.5e-06, - "supports_audio_output": true, + "supports_audio_output": false, "supports_function_calling": true, "supports_response_schema": true, "supports_system_messages": true, @@ -41449,7 +41506,8 @@ "supports_audio_input": true, "supports_pdf_input": true, "supports_reasoning": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/google/gemini-2.5-pro": { "cache_creation_input_token_cost": 3.75e-07, @@ -41462,7 +41520,7 @@ "max_tokens": 65536, "mode": "chat", "output_cost_per_token": 1e-05, - "supports_audio_output": true, + "supports_audio_output": false, "supports_function_calling": true, "supports_response_schema": true, "supports_system_messages": true, @@ -41478,7 +41536,8 @@ "supports_audio_input": true, "supports_pdf_input": true, "supports_reasoning": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/google/gemini-3-pro-preview": { "cache_read_input_token_cost": 2e-07, @@ -41563,7 +41622,8 @@ "supports_url_context": true, "supports_vision": true, "supports_web_search": false, - "tpm": 800000 + "tpm": 800000, + "supports_video_input": true }, "openrouter/google/gemini-3.1-flash-lite-preview": { "cache_creation_input_token_cost": 8.33333333333333e-08, @@ -41690,7 +41750,8 @@ "supports_system_messages": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/gryphe/mythomax-l2-13b": { "input_cost_per_token": 8e-08, @@ -44413,12 +44474,16 @@ "output_cost_per_token": 1.2e-05, "litellm_provider": "replicate", "mode": "chat", - "supports_function_calling": true, - "supports_parallel_function_calling": true, + "supports_function_calling": false, + "supports_parallel_function_calling": false, "supports_vision": true, "supports_system_messages": true, - "supports_tool_choice": true, - "supports_response_schema": true + "supports_tool_choice": false, + "supports_response_schema": false, + "input_cost_per_token_above_200k_tokens": 4e-06, + "output_cost_per_token_above_200k_tokens": 1.8e-05, + "supports_audio_input": true, + "supports_video_input": true }, "replicate/anthropic/claude-4.5-sonnet": { "input_cost_per_token": 3e-06, @@ -44487,17 +44552,19 @@ "supports_response_schema": true }, "replicate/google/gemini-2.5-flash": { - "input_cost_per_token": 2.5e-06, + "input_cost_per_token": 3e-07, "output_cost_per_token": 2.5e-06, "litellm_provider": "replicate", "mode": "chat", - "supports_function_calling": true, - "supports_parallel_function_calling": true, + "supports_function_calling": false, + "supports_parallel_function_calling": false, "supports_vision": true, "supports_system_messages": true, - "supports_tool_choice": true, - "supports_response_schema": true, - "supports_image_size": false + "supports_tool_choice": false, + "supports_response_schema": false, + "supports_image_size": false, + "supports_reasoning": true, + "supports_video_input": true }, "replicate/openai/gpt-oss-120b": { "input_cost_per_token": 1.8e-07, @@ -48076,10 +48143,15 @@ "supports_function_calling": true, "supports_tool_choice": true, "supports_response_schema": true, - "supports_image_size": false + "supports_image_size": false, + "cache_read_input_token_cost": 3e-08, + "supports_reasoning": true, + "supports_pdf_input": true, + "supports_web_search": true, + "supports_prompt_caching": true }, "vercel_ai_gateway/google/gemini-2.5-pro": { - "input_cost_per_token": 2.5e-06, + "input_cost_per_token": 1.25e-06, "litellm_provider": "vercel_ai_gateway", "max_input_tokens": 1048576, "max_output_tokens": 65536, @@ -48089,7 +48161,15 @@ "supports_vision": true, "supports_function_calling": true, "supports_tool_choice": true, - "supports_response_schema": true + "supports_response_schema": true, + "input_cost_per_token_above_200k_tokens": 2.5e-06, + "output_cost_per_token_above_200k_tokens": 1.5e-05, + "cache_read_input_token_cost": 1.25e-07, + "cache_read_input_token_cost_above_200k_tokens": 2.5e-07, + "supports_reasoning": true, + "supports_pdf_input": true, + "supports_web_search": true, + "supports_prompt_caching": true }, "vercel_ai_gateway/google/gemini-embedding-001": { "input_cost_per_token": 1.5e-07, @@ -48825,7 +48905,8 @@ "mode": "audio_transcription", "source": "https://cloud.google.com/speech-to-text/pricing", "supported_endpoints": [ - "/v1/audio/transcriptions" + "/v1/audio/transcriptions", + "/v1/realtime" ] }, "vertex_ai/claude-3-5-haiku": { @@ -62224,7 +62305,8 @@ "supports_response_schema": true, "supports_reasoning": true, "supports_vision": true, - "source": "https://deepinfra.com/pricing" + "source": "https://deepinfra.com/pricing", + "supports_audio_input": true }, "deepinfra/XiaomiMiMo/MiMo-V2.5": { "max_tokens": 262144, @@ -62524,7 +62606,8 @@ "supports_response_schema": true, "supports_reasoning": true, "supports_vision": true, - "source": "https://deepinfra.com/pricing" + "source": "https://deepinfra.com/pricing", + "supports_audio_input": true }, "deepinfra/google/gemini-3.7-flash": { "max_tokens": 1000000, @@ -62538,7 +62621,8 @@ "supports_response_schema": true, "supports_reasoning": true, "supports_vision": true, - "source": "https://deepinfra.com/pricing" + "source": "https://deepinfra.com/pricing", + "supports_audio_input": true }, "deepinfra/inclusionAI/Ling-3.0-flash": { "max_tokens": 131072, @@ -62970,7 +63054,8 @@ "supports_response_schema": true, "supports_reasoning": true, "supports_vision": true, - "source": "https://deepinfra.com/pricing" + "source": "https://deepinfra.com/pricing", + "supports_audio_input": true }, "deepinfra/XiaomiMiMo/MiMo-V2.5-Pro": { "max_tokens": 1048576, @@ -63375,6 +63460,34 @@ "video" ] }, + "xai/grok-voice-transcribe-1.0": { + "input_cost_per_second": 2.778e-05, + "litellm_provider": "xai", + "metadata": { + "calculation": "$0.10/3600 seconds = $0.00002778 per second", + "original_pricing_per_hour": 0.1 + }, + "mode": "audio_transcription", + "output_cost_per_second": 0.0, + "source": "https://docs.x.ai/developers/pricing", + "supported_endpoints": [ + "/v1/audio/transcriptions" + ] + }, + "xai/grok-voice-transcribe-2.0": { + "input_cost_per_second": 2.778e-05, + "litellm_provider": "xai", + "metadata": { + "calculation": "$0.10/3600 seconds = $0.00002778 per second", + "original_pricing_per_hour": 0.1 + }, + "mode": "audio_transcription", + "output_cost_per_second": 0.0, + "source": "https://docs.x.ai/developers/pricing", + "supported_endpoints": [ + "/v1/audio/transcriptions" + ] + }, "low/1024-x-1024/grok-imagine-image-2.0": { "input_cost_per_image": 0.04, "litellm_provider": "xai", @@ -65365,7 +65478,8 @@ "deprecation_date": "2026-10-20", "input_cost_per_audio_token": 3e-07, "supports_prompt_caching": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/google/gemini-3.5-flash": { "cache_creation_input_token_cost": 8.33333333333333e-08, @@ -65388,7 +65502,8 @@ "cache_read_input_token_cost": 1.5e-07, "input_cost_per_audio_token": 3e-06, "supports_prompt_caching": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/google/gemini-3.5-flash-lite": { "cache_creation_input_token_cost": 8.33333333333333e-08, @@ -65411,7 +65526,8 @@ "cache_read_input_token_cost": 3e-08, "input_cost_per_audio_token": 3e-07, "supports_prompt_caching": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/google/gemini-3.6-flash": { "cache_creation_input_token_cost": 4.16666666666667e-08, @@ -65434,7 +65550,8 @@ "cache_read_input_token_cost": 7.5e-08, "input_cost_per_audio_token": 7.5e-07, "supports_prompt_caching": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/google/gemini-3.7-flash": { "cache_creation_input_token_cost": 4.16666666666667e-08, @@ -65457,7 +65574,8 @@ "cache_read_input_token_cost": 7.5e-08, "input_cost_per_audio_token": 7.5e-07, "supports_prompt_caching": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/google/gemini-3.8-flash": { "cache_creation_input_token_cost": 4.16666666666667e-08, @@ -65480,7 +65598,8 @@ "cache_read_input_token_cost": 7.5e-08, "input_cost_per_audio_token": 7.5e-07, "supports_prompt_caching": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/openai/gpt-4o-mini": { "input_cost_per_token": 1.5e-07, @@ -67108,7 +67227,8 @@ "supports_pdf_input": true, "supports_audio_input": true, "supports_prompt_caching": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/qwen/qwen3-max-thinking": { "input_cost_per_token": 7.8e-07, @@ -71974,7 +72094,8 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/google/gemini-2.5-flash:batch": { "cache_read_input_audio_token_cost": 1e-07, @@ -71997,7 +72118,8 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/google/gemini-2.5-pro:batch": { "cache_read_input_audio_token_cost": 1.25e-07, @@ -72023,7 +72145,8 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/google/gemini-3-flash-preview:batch": { "input_cost_per_audio_token": 5e-07, @@ -72043,7 +72166,8 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/google/gemini-3.1-flash-lite:batch": { "cache_read_input_audio_token_cost": 2.5e-08, @@ -72065,7 +72189,8 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/google/gemini-3.1-pro-preview:batch": { "input_cost_per_audio_token": 1e-06, @@ -72087,7 +72212,8 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/google/gemini-3.5-flash-lite:batch": { "cache_read_input_audio_token_cost": 1.5e-08, @@ -72109,7 +72235,8 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/google/gemini-3.5-flash:batch": { "cache_read_input_audio_token_cost": 1.5e-07, @@ -72131,7 +72258,8 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/google/gemini-3.6-flash:batch": { "cache_creation_input_token_cost": 4.16666666666667e-08, @@ -72154,7 +72282,8 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/google/gemini-3.7-flash:batch": { "cache_creation_input_token_cost": 4.16666666666667e-08, @@ -72177,7 +72306,8 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/google/gemini-3.8-flash:batch": { "cache_creation_input_token_cost": 4.16666666666667e-08, @@ -72200,7 +72330,8 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "supports_web_search": false + "supports_web_search": false, + "supports_video_input": true }, "openrouter/ibm-granite/granite-4.0-h-micro": { "input_cost_per_token": 1.7e-08, diff --git a/pyproject.toml b/pyproject.toml index dfe84a28d52..a017e39e084 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -129,6 +129,12 @@ grpc = [ # Newest non-yanked release older than the 30-day cutoff. "grpcio==1.78.0", ] +stt-vertex-chirp = [ + # Google Cloud Speech-to-Text v2 streaming (gRPC) for Chirp models on + # /v1/realtime. Imported lazily inside the backend so litellm core stays + # usable without it. + "google-cloud-speech>=2.40.0,<3.0", +] stt-nvidia-riva = [ # NVIDIA Riva STT provider (gRPC). These are imported lazily inside the # provider handler so litellm core remains usable without them. @@ -152,6 +158,7 @@ proxy-runtime = [ # Keep these in a dedicated extra so uv-based images preserve the same # feature surface without forcing the base SDK install to grow. "google-cloud-aiplatform>=1.133.0,<2.0", + "google-cloud-speech>=2.40.0,<3.0", "google-genai>=1.37.0,<2.0", "anthropic[vertex]>=0.84.0,<1.0", "grpcio==1.78.0", @@ -270,6 +277,7 @@ ci = [ "langgraph>=1.2.4,<1.3.0", "langgraph-prebuilt>=1.1.0,<1.3.0", "claude-agent-sdk==0.1.44", + "google-cloud-speech==2.40.0", ] healthcheck = [ "httpx==0.28.1", @@ -329,9 +337,6 @@ litellm-enterprise = { workspace = true } [tool.uv.workspace] members = ["enterprise", "litellm-proxy-extras"] -[tool.isort] -profile = "black" - [tool.commitizen] version = "1.103.0" version_files = [ diff --git a/scripts/comment-fixed-issue.test.ts b/scripts/comment-fixed-issue.test.ts new file mode 100644 index 00000000000..f9cd41d96d8 --- /dev/null +++ b/scripts/comment-fixed-issue.test.ts @@ -0,0 +1,234 @@ +import { describe, expect, test } from "bun:test"; + +import type { Comment, GitHubApi } from "./auto-close-duplicates"; +import { + FIXED_MARKER, + closerOf, + commentFixedIssue, + fixedBody, + nextMinor, + parseVersion, + placement, + readConfig, + releaseCandidate, + type ClosedIssue, + type FixedConfig, +} from "./comment-fixed-issue"; + +const MERGE_COMMIT = "68c4c82ac977b48b2b81ee8d633d5771307c6162"; + +const mergedPr = { + __typename: "PullRequest" as const, + number: 41767, + merged: true, + baseRefName: "main", + mergeCommit: { oid: MERGE_COMMIT }, +}; + +type Closer = ClosedIssue["timelineItems"]["nodes"][number]["closer"]; + +const closedBy = (closer: Closer, state: ClosedIssue["state"] = "CLOSED"): ClosedIssue => ({ + state, + timelineItems: { nodes: [{ closer }] }, +}); + +const pyproject = (version: string): string => + `[project]\nname = "litellm"\nversion = "${version}"\n\n[tool.commitizen]\nversion = "${version}"\n`; + +const config: FixedConfig = { repo: "BerriAI/litellm", issueNumber: 41750, defaultBranch: "main", dryRun: false }; + +interface World { + readonly issue?: ClosedIssue | null; + readonly comments?: readonly Comment[]; + readonly version?: string; + // Which existing rc.1 tags contain the merge commit; a tag absent from the map does not exist + readonly tags?: Readonly>; +} + +function fakeApi(world: World = {}): { readonly api: GitHubApi; readonly writes: string[] } { + const writes: string[] = []; + const tags = world.tags ?? {}; + const api: GitHubApi = { + request: async (method: string, path: string, body?: object): Promise => { + if (method === "POST" && path === "/graphql") { + return { data: { repository: { issue: world.issue === undefined ? closedBy(mergedPr) : world.issue } } } as T; + } + if (method !== "GET") { + writes.push(`${method} ${path} ${JSON.stringify(body)}`); + return {} as T; + } + if (path.startsWith("/repos/BerriAI/litellm/issues/41750/comments")) { + return (world.comments ?? []) as T; + } + if (path === `/repos/BerriAI/litellm/contents/pyproject.toml?ref=${MERGE_COMMIT}`) { + return { content: btoa(pyproject(world.version ?? "1.103.0")).replace(/(.{60})/g, "$1\n") } as T; + } + const matching = /^\/repos\/BerriAI\/litellm\/git\/matching-refs\/tags\/(.+)$/.exec(path); + if (matching !== null) { + return (matching[1] in tags ? [{ ref: `refs/tags/${matching[1]}` }] : []) as T; + } + const compare = /^\/repos\/BerriAI\/litellm\/compare\/(.+)\.\.\.(.+)$/.exec(path); + if (compare !== null && compare[2] === MERGE_COMMIT) { + return { status: tags[compare[1]] ? "behind" : "ahead" } as T; + } + throw new Error(`unexpected ${method} ${path}`); + }, + }; + return { api, writes }; +} + +describe("closerOf", () => { + test("a pull request merged into the default branch is the fix", () => { + expect(closerOf(closedBy(mergedPr), "main")).toEqual({ kind: "pull_request", number: 41767, mergeCommit: MERGE_COMMIT }); + }); + + test("an issue closed by hand, by a commit, or by an unmerged pull request gets no comment", () => { + expect(closerOf(closedBy(null), "main")).toEqual({ kind: "skip", reason: "closed by hand, not by a pull request" }); + expect(closerOf(closedBy({ __typename: "Commit", oid: MERGE_COMMIT }), "main").kind).toBe("skip"); + expect(closerOf(closedBy({ ...mergedPr, merged: false }), "main").kind).toBe("skip"); + expect(closerOf(closedBy({ ...mergedPr, mergeCommit: null }), "main").kind).toBe("skip"); + }); + + test("a pull request merged into a release branch is not a fix on main", () => { + const verdict = closerOf(closedBy({ ...mergedPr, baseRefName: "release/1.102.0rc2" }), "main"); + expect(verdict).toEqual({ kind: "skip", reason: "#41767 merged into release/1.102.0rc2, not main" }); + }); + + test("an issue reopened after the close event is left alone", () => { + expect(closerOf(closedBy(mergedPr, "OPEN"), "main")).toEqual({ kind: "skip", reason: "the issue is open again" }); + }); +}); + +describe("version helpers", () => { + test("parseVersion reads the project version and ignores everything else", () => { + expect(parseVersion(pyproject("1.103.0"))).toBe("1.103.0"); + expect(parseVersion('[project]\nversion = "1.103.0rc1"\n')).toBeUndefined(); + expect(parseVersion("[project]\nname = 'litellm'\n")).toBeUndefined(); + }); + + test("the first rc of a version is the release that carries a fix merged under it", () => { + expect(releaseCandidate("1.103.0")).toBe("v1.103.0-rc.1"); + }); + + test("nextMinor bumps the minor and resets the patch", () => { + expect(nextMinor("1.103.0")).toBe("1.104.0"); + expect(nextMinor("1.99.4")).toBe("1.100.0"); + }); +}); + +describe("placement", () => { + test("no rc yet: the fix ships in the rc.1 of the version at the merge commit", async () => { + const { api } = fakeApi({ version: "1.103.0" }); + expect(await placement(api, "BerriAI/litellm", MERGE_COMMIT)).toEqual({ kind: "release", tag: "v1.103.0-rc.1", shipped: false }); + }); + + test("rc.1 already cut with the commit in it: the fix is out", async () => { + const { api } = fakeApi({ version: "1.102.0", tags: { "v1.102.0-rc.1": true } }); + expect(await placement(api, "BerriAI/litellm", MERGE_COMMIT)).toEqual({ kind: "release", tag: "v1.102.0-rc.1", shipped: true }); + }); + + test("rc.1 cut before the merge while main still said that version: the fix waits for the next minor", async () => { + const { api } = fakeApi({ version: "1.102.0", tags: { "v1.102.0-rc.1": false } }); + expect(await placement(api, "BerriAI/litellm", MERGE_COMMIT)).toEqual({ kind: "release", tag: "v1.103.0-rc.1", shipped: false }); + }); + + test("keeps walking minors while each rc.1 exists without the commit, then gives up", async () => { + const twoTaken = fakeApi({ version: "1.102.0", tags: { "v1.102.0-rc.1": false, "v1.103.0-rc.1": false } }); + expect(await placement(twoTaken.api, "BerriAI/litellm", MERGE_COMMIT)).toEqual({ kind: "release", tag: "v1.104.0-rc.1", shipped: false }); + + const allTaken = fakeApi({ + version: "1.102.0", + tags: { "v1.102.0-rc.1": false, "v1.103.0-rc.1": false, "v1.104.0-rc.1": false, "v1.105.0-rc.1": false }, + }); + expect((await placement(allTaken.api, "BerriAI/litellm", MERGE_COMMIT)).kind).toBe("skip"); + }); + + test("a pyproject without a version line is a skip, not a comment", async () => { + const { api } = fakeApi({ version: "not-a-version" }); + expect((await placement(api, "BerriAI/litellm", MERGE_COMMIT)).kind).toBe("skip"); + }); +}); + +describe("fixedBody", () => { + test("names the pull request and the first release, and carries the marker the rerun looks for", () => { + const body = fixedBody(41767, { tag: "v1.103.0-rc.1", shipped: false }); + expect(body.startsWith(FIXED_MARKER)).toBe(true); + expect(body).toContain("Fixed by #41767."); + expect(body).toContain("ships in v1.103.0-rc.1 and up"); + expect(body).toContain("dev pre-release"); + }); + + test("a release that is already out says so instead of promising one", () => { + const body = fixedBody(41767, { tag: "v1.102.0-rc.1", shipped: true }); + expect(body).toContain("is in v1.102.0-rc.1 and up"); + expect(body).not.toContain("ships in"); + }); + + test("stays within the 25-word comment rule either way", () => { + for (const shipped of [true, false]) { + const words = fixedBody(41767, { tag: "v1.103.0-rc.1", shipped }).replace(FIXED_MARKER, "").trim().split(/\s+/); + expect(words.length).toBeGreaterThanOrEqual(15); + expect(words.length).toBeLessThanOrEqual(25); + } + }); +}); + +describe("commentFixedIssue", () => { + test("a real run posts one comment naming the pull request and the release", async () => { + const { api, writes } = fakeApi(); + const verdict = await commentFixedIssue(api, config); + expect(verdict).toMatchObject({ kind: "commented", pullRequest: 41767, tag: "v1.103.0-rc.1" }); + expect(writes).toHaveLength(1); + expect(writes[0]).toContain("POST /repos/BerriAI/litellm/issues/41750/comments"); + expect(writes[0]).toContain("Fixed by #41767. This ships in v1.103.0-rc.1 and up"); + }); + + test("a dry run renders the comment and writes nothing", async () => { + const { api, writes } = fakeApi(); + const verdict = await commentFixedIssue(api, { ...config, dryRun: true }); + expect(verdict.kind).toBe("commented"); + expect(writes).toEqual([]); + }); + + test("an issue that already carries the comment is not commented twice", async () => { + const existing: Comment = { + id: 1, + body: `${FIXED_MARKER}\nFixed by #41767. This ships in v1.103.0-rc.1 and up.`, + created_at: "2026-09-18T00:00:00Z", + user: { type: "Bot", login: "github-actions[bot]" }, + }; + const { api, writes } = fakeApi({ comments: [existing] }); + expect(await commentFixedIssue(api, config)).toEqual({ kind: "skip", reason: "already carries a fixed-in comment" }); + expect(writes).toEqual([]); + }); + + test("a hand-closed issue never reaches the release lookup or the API writes", async () => { + const { api, writes } = fakeApi({ issue: closedBy(null) }); + expect((await commentFixedIssue(api, config)).kind).toBe("skip"); + expect(writes).toEqual([]); + }); + + test("a number that is not an issue in the repository is a skip", async () => { + const { api, writes } = fakeApi({ issue: null }); + expect(await commentFixedIssue(api, config)).toEqual({ kind: "skip", reason: "not an issue in this repository" }); + expect(writes).toEqual([]); + }); +}); + +describe("readConfig", () => { + const env = { GITHUB_TOKEN: "t", GITHUB_REPOSITORY: "BerriAI/litellm", ISSUE_NUMBER: "41750", DEFAULT_BRANCH: "main" }; + + test("reads the four inputs and treats anything but the literal true as a real run", () => { + expect(readConfig(env)).toEqual({ token: "t", repo: "BerriAI/litellm", issueNumber: 41750, defaultBranch: "main", dryRun: false }); + expect(readConfig({ ...env, DRY_RUN: "true" }).dryRun).toBe(true); + expect(readConfig({ ...env, DRY_RUN: "false" }).dryRun).toBe(false); + }); + + test("refuses a missing token, repo, branch or a bad issue number", () => { + expect(() => readConfig({ ...env, GITHUB_TOKEN: undefined })).toThrow("GITHUB_TOKEN"); + expect(() => readConfig({ ...env, GITHUB_REPOSITORY: "litellm" })).toThrow("owner/repo"); + expect(() => readConfig({ ...env, DEFAULT_BRANCH: "" })).toThrow("DEFAULT_BRANCH"); + expect(() => readConfig({ ...env, ISSUE_NUMBER: "0" })).toThrow("ISSUE_NUMBER"); + expect(() => readConfig({ ...env, ISSUE_NUMBER: "abc" })).toThrow("ISSUE_NUMBER"); + }); +}); diff --git a/scripts/comment-fixed-issue.ts b/scripts/comment-fixed-issue.ts new file mode 100644 index 00000000000..480b5e90249 --- /dev/null +++ b/scripts/comment-fixed-issue.ts @@ -0,0 +1,224 @@ +#!/usr/bin/env bun + +import { githubApi, listAll, type Comment, type GitHubApi } from "./auto-close-duplicates"; + +declare const process: { readonly env: Readonly> }; + +export interface FixedConfig { + readonly repo: string; + readonly issueNumber: number; + readonly defaultBranch: string; + readonly dryRun: boolean; +} + +interface PullRequestCloser { + readonly __typename: "PullRequest"; + readonly number: number; + readonly merged: boolean; + readonly baseRefName: string; + readonly mergeCommit: { readonly oid: string } | null; +} + +interface CommitCloser { + readonly __typename: "Commit"; + readonly oid: string; +} + +export interface ClosedIssue { + readonly state: "OPEN" | "CLOSED"; + readonly timelineItems: { + readonly nodes: readonly { readonly closer: PullRequestCloser | CommitCloser | null }[]; + }; +} + +interface TimelineResponse { + readonly data?: { readonly repository?: { readonly issue: ClosedIssue | null } }; +} + +interface MatchingRef { + readonly ref: string; +} + +interface Comparison { + readonly status: "ahead" | "behind" | "identical" | "diverged"; +} + +interface FileContent { + readonly content: string; +} + +export type Closer = + | { readonly kind: "pull_request"; readonly number: number; readonly mergeCommit: string } + | { readonly kind: "skip"; readonly reason: string }; + +export type Placement = + | { readonly kind: "release"; readonly tag: string; readonly shipped: boolean } + | { readonly kind: "skip"; readonly reason: string }; + +export type FixedVerdict = + | { readonly kind: "commented"; readonly pullRequest: number; readonly tag: string; readonly body: string } + | { readonly kind: "skip"; readonly reason: string }; + +export const FIXED_MARKER = ""; +const MAX_MINOR_BUMPS = 3; + +export const CLOSER_QUERY = `query($owner: String!, $name: String!, $number: Int!) { + repository(owner: $owner, name: $name) { + issue(number: $number) { + state + timelineItems(last: 1, itemTypes: [CLOSED_EVENT]) { + nodes { + ... on ClosedEvent { + closer { + __typename + ... on PullRequest { number merged baseRefName mergeCommit { oid } } + ... on Commit { oid } + } + } + } + } + } + } +}`; + +const skip = (reason: string): { readonly kind: "skip"; readonly reason: string } => ({ kind: "skip", reason }); + +export function closerOf(issue: ClosedIssue, defaultBranch: string): Closer { + if (issue.state !== "CLOSED") { + return skip("the issue is open again"); + } + const closer = issue.timelineItems.nodes[0]?.closer ?? null; + if (closer === null) { + return skip("closed by hand, not by a pull request"); + } + if (closer.__typename === "Commit") { + return skip(`closed by commit ${closer.oid.slice(0, 10)}, not by a pull request`); + } + if (!closer.merged || closer.mergeCommit === null) { + return skip(`closed by #${closer.number}, which is not merged`); + } + if (closer.baseRefName !== defaultBranch) { + return skip(`#${closer.number} merged into ${closer.baseRefName}, not ${defaultBranch}`); + } + return { kind: "pull_request", number: closer.number, mergeCommit: closer.mergeCommit.oid }; +} + +export function parseVersion(pyproject: string): string | undefined { + return /^version = "(\d+\.\d+\.\d+)"$/m.exec(pyproject)?.[1]; +} + +export function releaseCandidate(version: string): string { + return `v${version}-rc.1`; +} + +export function nextMinor(version: string): string { + const [major, minor] = version.split(".").map(Number); + return `${major}.${minor + 1}.0`; +} + +async function tagExists(api: GitHubApi, repo: string, tag: string): Promise { + const refs = await api.request("GET", `/repos/${repo}/git/matching-refs/tags/${tag}`); + return refs.some((ref) => ref.ref === `refs/tags/${tag}`); +} + +async function tagContains(api: GitHubApi, repo: string, tag: string, sha: string): Promise { + const comparison = await api.request("GET", `/repos/${repo}/compare/${tag}...${sha}`); + return comparison.status === "behind" || comparison.status === "identical"; +} + +// The first rc of a version is cut straight from main, so a fix merged while pyproject says X.Y.Z ships in +// vX.Y.Z-rc.1 unless that rc was already cut without it, in which case it waits for the next minor's rc.1 +async function firstReleaseWith( + api: GitHubApi, + repo: string, + sha: string, + version: string, + bumpsLeft: number, +): Promise { + const tag = releaseCandidate(version); + if (!(await tagExists(api, repo, tag))) { + return { kind: "release", tag, shipped: false }; + } + if (await tagContains(api, repo, tag, sha)) { + return { kind: "release", tag, shipped: true }; + } + if (bumpsLeft === 0) { + return skip(`${tag} exists without ${sha.slice(0, 10)} and the next ${MAX_MINOR_BUMPS} rc.1 tags are taken too`); + } + return firstReleaseWith(api, repo, sha, nextMinor(version), bumpsLeft - 1); +} + +export async function placement(api: GitHubApi, repo: string, mergeCommit: string): Promise { + const file = await api.request("GET", `/repos/${repo}/contents/pyproject.toml?ref=${mergeCommit}`); + const version = parseVersion(atob(file.content.replace(/\n/g, ""))); + if (version === undefined) { + return skip(`pyproject.toml at ${mergeCommit.slice(0, 10)} has no version line`); + } + return firstReleaseWith(api, repo, mergeCommit, version, MAX_MINOR_BUMPS); +} + +export function fixedBody(pullRequest: number, release: { readonly tag: string; readonly shipped: boolean }): string { + const availability = release.shipped + ? `This is in ${release.tag} and up, so upgrading to that release or any newer one picks it up.` + : `This ships in ${release.tag} and up, and the next dev pre-release cut from main will carry it too.`; + return `${FIXED_MARKER}\nFixed by #${pullRequest}. ${availability}`; +} + +export async function commentFixedIssue(api: GitHubApi, config: FixedConfig): Promise { + const [owner, name] = config.repo.split("/"); + const response = await api.request("POST", "/graphql", { + query: CLOSER_QUERY, + variables: { owner, name, number: config.issueNumber }, + }); + const issue = response.data?.repository?.issue ?? null; + if (issue === null) { + return skip("not an issue in this repository"); + } + const closer = closerOf(issue, config.defaultBranch); + if (closer.kind === "skip") { + return closer; + } + const issuePath = `/repos/${config.repo}/issues/${config.issueNumber}`; + const comments = await listAll(api, `${issuePath}/comments`); + if (comments.some((comment) => comment.body.includes(FIXED_MARKER))) { + return skip("already carries a fixed-in comment"); + } + const release = await placement(api, config.repo, closer.mergeCommit); + if (release.kind === "skip") { + return release; + } + const body = fixedBody(closer.number, release); + if (!config.dryRun) { + await api.request("POST", `${issuePath}/comments`, { body }); + } + return { kind: "commented", pullRequest: closer.number, tag: release.tag, body }; +} + +export function readConfig(env: Readonly>): FixedConfig & { readonly token: string } { + const token = env.GITHUB_TOKEN; + const repo = env.GITHUB_REPOSITORY; + const defaultBranch = env.DEFAULT_BRANCH; + if (!token || !repo || !/^[\w.-]+\/[\w.-]+$/.test(repo) || !defaultBranch) { + throw new Error("GITHUB_TOKEN, GITHUB_REPOSITORY (owner/repo) and DEFAULT_BRANCH are required"); + } + const issueNumber = Number(env.ISSUE_NUMBER); + if (!Number.isInteger(issueNumber) || issueNumber <= 0) { + throw new Error(`ISSUE_NUMBER must be a positive integer, got "${env.ISSUE_NUMBER}"`); + } + return { token, repo, issueNumber, defaultBranch, dryRun: env.DRY_RUN === "true" }; +} + +function describe(config: FixedConfig, verdict: FixedVerdict): string { + if (verdict.kind === "skip") { + return `#${config.issueNumber}: skipped, ${verdict.reason}`; + } + if (config.dryRun) { + return `#${config.issueNumber}: DRY RUN, set the ISSUE_FIXED_COMMENT_ENABLED repo variable to true to post this:\n\n${verdict.body}`; + } + return `#${config.issueNumber}: commented, fixed by #${verdict.pullRequest} in ${verdict.tag}`; +} + +if (import.meta.main) { + const { token, ...config } = readConfig(process.env); + console.log(describe(config, await commentFixedIssue(githubApi(token), config))); +} diff --git a/scripts/test_tool_allowlist_script.py b/scripts/test_tool_allowlist_script.py index 9503a21219c..f94aac60f80 100644 --- a/scripts/test_tool_allowlist_script.py +++ b/scripts/test_tool_allowlist_script.py @@ -3,10 +3,10 @@ Standalone script to test tool allowlist enforcement and tool name extraction. Run from repo root: - poetry run python scripts/test_tool_allowlist_script.py + uv run python scripts/test_tool_allowlist_script.py Or run the unit tests: - poetry run pytest tests/test_litellm/proxy/test_tools_allowlist_enforcement.py -v + uv run pytest tests/test_litellm/proxy/test_tools_allowlist_enforcement.py -v """ import asyncio @@ -148,7 +148,7 @@ def main(): asyncio.run(test_check_tools_allowlist()) print("Done. For full unit tests run:") print( - " poetry run pytest tests/test_litellm/proxy/test_tools_allowlist_enforcement.py -v" + " uv run pytest tests/test_litellm/proxy/test_tools_allowlist_enforcement.py -v" ) diff --git a/tests/code_coverage_tests/check_e2e_no_raw_requests.py b/tests/code_coverage_tests/check_e2e_no_raw_requests.py index fe6a77fc26c..3f40cc3ee1e 100644 --- a/tests/code_coverage_tests/check_e2e_no_raw_requests.py +++ b/tests/code_coverage_tests/check_e2e_no_raw_requests.py @@ -5,7 +5,7 @@ anywhere; a small allowlist grandfathers the files that legitimately make raw ca (the transport itself, the root conftest liveness probe, the claude_code version resolver's constant registry URL fetch, and the mcp OAuth client, whose httpx client is the object the official mcp SDK's streamable_http_client requires and so -cannot go through the sync requests transport). Referenced by tests/e2e/CLAUDE.md.""" +cannot go through the sync requests transport). Referenced by tests/e2e/AGENTS.md.""" from __future__ import annotations diff --git a/tests/code_coverage_tests/test_e2e_changed_gate.py b/tests/code_coverage_tests/test_e2e_changed_gate.py index 101816c7f11..5ae0863baf0 100644 --- a/tests/code_coverage_tests/test_e2e_changed_gate.py +++ b/tests/code_coverage_tests/test_e2e_changed_gate.py @@ -142,7 +142,7 @@ def select_tests(changed: tuple[str, ...]) -> tuple[str, ...]: ("tests/e2e/guardrails/test_bedrock_guardrail_e2e.py",), ("tests/e2e/guardrails/test_bedrock_guardrail_e2e.py",), ), - (("tests/e2e/logging/helpers.py", "docs/my-website/docs/index.md", "tests/e2e/CLAUDE.md"), ()), + (("tests/e2e/logging/helpers.py", "docs/my-website/docs/index.md", "tests/e2e/AGENTS.md"), ()), ( ("tests/e2e/logging/test_datadog_e2e.py", "tests/e2e/logging/test_datadog_e2e.py"), ("tests/e2e/logging/test_datadog_e2e.py",), diff --git a/tests/e2e/CLAUDE.md b/tests/e2e/AGENTS.md similarity index 99% rename from tests/e2e/CLAUDE.md rename to tests/e2e/AGENTS.md index 0cdc0fdb124..967a85f1255 100644 --- a/tests/e2e/CLAUDE.md +++ b/tests/e2e/AGENTS.md @@ -1,6 +1,6 @@ # e2e harness conventions -Code-style rules for writing tests under `tests/e2e/`. The harness already encodes the plumbing; your job is the feature-specific behavior, not reinventing it. For what a complete test must do (the lifecycle contract, asserting both recorded state and enforced behavior) and how to run a suite, see `CONTRIBUTING.md` in this directory. Repo-wide conventions live in the root `CLAUDE.md` +Code-style rules for writing tests under `tests/e2e/`. The harness already encodes the plumbing; your job is the feature-specific behavior, not reinventing it. For what a complete test must do (the lifecycle contract, asserting both recorded state and enforced behavior) and how to run a suite, see `CONTRIBUTING.md` in this directory. Repo-wide conventions live in the root `AGENTS.md` ## Suite folders diff --git a/tests/e2e/CONTRIBUTING.md b/tests/e2e/CONTRIBUTING.md index 20073e5d68f..2afcc563824 100644 --- a/tests/e2e/CONTRIBUTING.md +++ b/tests/e2e/CONTRIBUTING.md @@ -2,7 +2,7 @@ This directory holds the live end-to-end suites that prove product correctness against a real running proxy and real provider APIs. The goal of this guide is simple: when you ship a feature, you add e2e coverage that walks that feature the way production does, across every route and edge case it touches, so a later change that breaks it fails here first -Read this before adding a test and i recommend reading through CLAUDE.md +Read this before adding a test and i recommend reading through AGENTS.md When contributing to this directory, please first discuss the change you wish to make via issue or pull request. We require screenshots and proof of your tests working on a live proxy. @@ -134,7 +134,7 @@ One sharp edge: a replayed response reuses the recorded provider response id, an Another sharp edge, same root: record and replay derive every per-test token deterministically (the model name included, so a replay regenerates the exact requests the record run sent), which means an edge-wired deployment left in the database by an interrupted earlier run carries the same model name as the fresh one the current run registers. The proxy then holds two deployments under one model group and load-balances across both, and because the leftover's `api_base` points at the earlier run's edge process, which is gone, the calls that land on it fail with a connection error that reads like a transport bug rather than the stale row it is. Give each record or replay run a fresh database, or let a run finish so its own teardown deletes what it registered, and never reuse one long-lived proxy across back-to-back record/replay sessions. CI hands every job its own empty database and its own proxy, so it never sees this -Replay answers any provider call that drifted from the recording with an HTTP 599 whose body names the computed and closest recorded keys, so the test fails loudly instead of silently going live, and a bundle older than seven days fails at collection time naming its age; either way the fix is to re-record. Only tests that register edge-wired deployments participate: everything else hits its provider live in every mode, so record exactly the suite you replay. If the proxy runs in a container, set `E2E_PROVIDER_EDGE_ADVERTISE_HOST` (e.g. `host.docker.internal`) so the api_base the proxy stores can reach the edge on the pytest host, and `E2E_PROVIDER_EDGE_BIND_HOST=0.0.0.0` so the edge accepts it. The suites wired to the edge today are `quota_management/spend_tracking/test_provider_edge_spend_e2e.py`, `llm_translation/test_chat_completions_contract_e2e.py`, the OpenAI registrations in `llm_translation/test_embeddings_endpoint_e2e.py`, the Anthropic tests in `llm_translation/test_messages_e2e.py`, streamed and not, and the OpenAI batch deployment behind `batches/`. A streamed response replays as the chunk sequence the provider sent rather than one buffered body. See `CLAUDE.md` in this directory for the bundle format, the edge design, and the current limits (Bedrock). The scheduled CI record/replay lane is described above +Replay answers any provider call that drifted from the recording with an HTTP 599 whose body names the computed and closest recorded keys, so the test fails loudly instead of silently going live, and a bundle older than seven days fails at collection time naming its age; either way the fix is to re-record. Only tests that register edge-wired deployments participate: everything else hits its provider live in every mode, so record exactly the suite you replay. If the proxy runs in a container, set `E2E_PROVIDER_EDGE_ADVERTISE_HOST` (e.g. `host.docker.internal`) so the api_base the proxy stores can reach the edge on the pytest host, and `E2E_PROVIDER_EDGE_BIND_HOST=0.0.0.0` so the edge accepts it. The suites wired to the edge today are `quota_management/spend_tracking/test_provider_edge_spend_e2e.py`, `llm_translation/test_chat_completions_contract_e2e.py`, the OpenAI registrations in `llm_translation/test_embeddings_endpoint_e2e.py`, the Anthropic tests in `llm_translation/test_messages_e2e.py`, streamed and not, and the OpenAI batch deployment behind `batches/`. A streamed response replays as the chunk sequence the provider sent rather than one buffered body. See `AGENTS.md` in this directory for the bundle format, the edge design, and the current limits (Bedrock). The scheduled CI record/replay lane is described above Tests marked `@pytest.mark.e2e` hard-fail when no proxy answers `/health/liveliness`, so a run that goes red with `No live proxy` at setup means the proxy isn't up; they never skip for a missing proxy, so an absent proxy can't be mistaken for a pass diff --git a/tests/e2e/batches/COVERAGE.md b/tests/e2e/batches/COVERAGE.md index b36d8937ad0..d18bed6c088 100644 --- a/tests/e2e/batches/COVERAGE.md +++ b/tests/e2e/batches/COVERAGE.md @@ -12,7 +12,7 @@ cost write-back via a cross-run marker baton (design below). Only supported cells are tested. The capability table in `capabilities.py` holds one row per supported (provider, scenario) pair, so there are no skipped cells in the parametrized run. The batches suite never skips: missing provider creds or upstream -failures are hard test failures (see `tests/e2e/CLAUDE.md`). +failures are hard test failures (see `tests/e2e/AGENTS.md`). | Provider | create | retrieve | cancel | list | content download | file backing | |-----------|--------|----------|--------|------|------------------|--------------| diff --git a/tests/e2e/claude_code/cron_vm/run_daily.sh b/tests/e2e/claude_code/cron_vm/run_daily.sh index 00d3e66e5bc..e878007d8a3 100755 --- a/tests/e2e/claude_code/cron_vm/run_daily.sh +++ b/tests/e2e/claude_code/cron_vm/run_daily.sh @@ -288,7 +288,7 @@ else # Download the tarball and Astral's official .sha256 sidecar to disk # and verify the digest before extracting/executing anything. This # closes the supply-chain trust gap of piping a remote binary - # straight into `tar -xzO ... > file ; chmod +x` (see CLAUDE.md + # straight into `tar -xzO ... > file ; chmod +x` (see AGENTS.md # "CI Supply-Chain Safety"). curl -fsSL --output "${UV_TMPDIR}/${UV_TARBALL_NAME}" "${UV_DOWNLOAD_URL}" curl -fsSL --output "${UV_TMPDIR}/${UV_TARBALL_NAME}.sha256" "${UV_DOWNLOAD_URL}.sha256" diff --git a/tests/e2e/coverage_registry/README.md b/tests/e2e/coverage_registry/README.md index da6aee84cc4..4f9845bab87 100644 --- a/tests/e2e/coverage_registry/README.md +++ b/tests/e2e/coverage_registry/README.md @@ -3,7 +3,7 @@ This directory is the **denominator** for e2e test coverage: the set of behaviors we want covered, one row per behavior, checked into the repo so coverage is a number we can track instead of a guess. It implements the plan in the "E2E Coverage Tracking" -note; the naming grammar lives in `tests/e2e/CLAUDE.md`. +note; the naming grammar lives in `tests/e2e/AGENTS.md`. ## The model diff --git a/tests/e2e/coverage_registry/__init__.py b/tests/e2e/coverage_registry/__init__.py index 959b3327194..0eb153011a6 100644 --- a/tests/e2e/coverage_registry/__init__.py +++ b/tests/e2e/coverage_registry/__init__.py @@ -3,6 +3,6 @@ `schema.py` defines one validated row per customer-noticeable behavior (a "cell"). The `*.yaml` files hold the rows, one file per id-prefix. `registry.py` loads and validates them; `collector.py` diffs the registry against the `@pytest.mark.covers` -markers on the live tests and reports coverage per module. See tests/e2e/CLAUDE.md +markers on the live tests and reports coverage per module. See tests/e2e/AGENTS.md for the naming grammar. """ diff --git a/tests/e2e/coverage_registry/mcp.yaml b/tests/e2e/coverage_registry/mcp.yaml index 05013389d77..bf511ad6b06 100644 --- a/tests/e2e/coverage_registry/mcp.yaml +++ b/tests/e2e/coverage_registry/mcp.yaml @@ -1,4 +1,4 @@ -# MCP module. Grounded in litellm/proxy/_experimental/mcp_server/. See tests/e2e/CLAUDE.md for the grammar. +# MCP module. Grounded in litellm/proxy/_experimental/mcp_server/. See tests/e2e/AGENTS.md for the grammar. - id: mcp.list_tools.api_key.succeeds module: mcp tier: P0 diff --git a/tests/e2e/llm_translation/realtime/REALTIME_COVERAGE_MATRIX.md b/tests/e2e/llm_translation/realtime/REALTIME_COVERAGE_MATRIX.md index a6e32b88479..c85471da90d 100644 --- a/tests/e2e/llm_translation/realtime/REALTIME_COVERAGE_MATRIX.md +++ b/tests/e2e/llm_translation/realtime/REALTIME_COVERAGE_MATRIX.md @@ -49,7 +49,7 @@ kept commented out in `PROVIDERS` until they pass end-to-end here; re-enable the uncommenting their entry. Every provider is provisioned and asserted; the suite never skips a provider. Per -`tests/e2e/CLAUDE.md` there is no sanctioned skip: the whole-suite proxy-liveness +`tests/e2e/AGENTS.md` there is no sanctioned skip: the whole-suite proxy-liveness probe hard-fails when no proxy answers, and a provider whose credentials or upstream realtime model are missing on the gateway is likewise a hard failure, not a skip. Give the gateway each provider's credentials to turn its tests green. diff --git a/tests/e2e/llm_translation/realtime/conftest.py b/tests/e2e/llm_translation/realtime/conftest.py index 752737e830e..804a9b9c649 100644 --- a/tests/e2e/llm_translation/realtime/conftest.py +++ b/tests/e2e/llm_translation/realtime/conftest.py @@ -29,7 +29,7 @@ def realtime_models(client: RealtimeClient) -> Iterator[dict[str, str]]: provider-id -> model-name map the tests connect with; delete them on teardown. Every provider is provisioned (never skipped): a provider whose credentials or upstream model are missing on the gateway hard-fails its test, per the suite's - fail-on-behavior contract in tests/e2e/CLAUDE.md.""" + fail-on-behavior contract in tests/e2e/AGENTS.md.""" records = tuple((provider.id, *client.provision(provider)) for provider in PROVIDERS) try: yield {provider_id: model_name for provider_id, model_name, _ in records} diff --git a/tests/e2e/llm_translation/realtime/realtime_client.py b/tests/e2e/llm_translation/realtime/realtime_client.py index 3ffca7e8b88..7c4a9cc4af9 100644 --- a/tests/e2e/llm_translation/realtime/realtime_client.py +++ b/tests/e2e/llm_translation/realtime/realtime_client.py @@ -38,7 +38,7 @@ class RealtimeProvider: the suite registers through /model/new (the gateway resolves the os.environ/* credential refs), so the suite is self-contained and never depends on a static gateway model_list. Every provider here is provisioned and asserted: per - tests/e2e/CLAUDE.md the suite never skips a provider, so a provider whose + tests/e2e/AGENTS.md the suite never skips a provider, so a provider whose credentials or upstream realtime model are missing on the gateway is a hard failure, not a skip.""" @@ -98,7 +98,7 @@ PROVIDERS = ( def realtime_model(provider: RealtimeProvider, provisioned: Mapping[str, str]) -> str: """Return the provisioned deployment name for this provider. Every provider in PROVIDERS is provisioned at session start, so a missing entry is a harness bug, - never an environment skip - the suite hard-fails instead (see tests/e2e/CLAUDE.md).""" + never an environment skip - the suite hard-fails instead (see tests/e2e/AGENTS.md).""" model = provisioned.get(provider.id) assert model is not None, ( f"{provider.id} was not provisioned; the realtime_models fixture is broken" diff --git a/tests/test_litellm/integrations/websearch_interception/test_websearch_agentic_loop_cap.py b/tests/test_litellm/integrations/websearch_interception/test_websearch_agentic_loop_cap.py index 40fd8c4e9e6..57e3ba59456 100644 --- a/tests/test_litellm/integrations/websearch_interception/test_websearch_agentic_loop_cap.py +++ b/tests/test_litellm/integrations/websearch_interception/test_websearch_agentic_loop_cap.py @@ -12,19 +12,23 @@ config.yaml through to the settings the loop actually reads. """ import json -from unittest.mock import MagicMock +from unittest.mock import MagicMock, patch import pytest import litellm +from litellm.exceptions import AuthenticationError, RateLimitError from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.websearch_interception.handler import ( + WEBSEARCH_EMIT_NATIVE_BLOCKS_KEY, WebSearchInterceptionLogger, ) +from litellm.integrations.websearch_interception.tools import get_litellm_web_search_tool +from litellm.litellm_core_utils.agentic_loop_settings import DEFAULT_MAX_AGENTIC_LOOPS from litellm.llms.anthropic.experimental_pass_through.messages.fake_stream_iterator import ( FakeAnthropicMessagesStreamIterator, ) -from litellm.litellm_core_utils.agentic_loop_settings import DEFAULT_MAX_AGENTIC_LOOPS +from litellm.llms.base_llm.search.transformation import SearchResponse, SearchResult from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler from litellm.secret_managers.main import get_secret from litellm.types.integrations.custom_logger import ( @@ -490,6 +494,135 @@ class TestOuterFramePostHookStillRuns: assert result["stop_reason"] == "end_turn" +def _response_asking_for_searches(*queries: str) -> dict: + return { + "id": "msg_123", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4-5", + "content": [ + {"id": f"toolu_internal_{index}", "type": "tool_use", "name": INTERNAL_TOOL_NAME, "input": {"query": query}} + for index, query in enumerate(queries, start=1) + ], + "stop_reason": "tool_use", + "usage": {"input_tokens": 10, "output_tokens": 5}, + } + + +class TestFailedSearchEndsTheTurn: + """ + A search that failed used to come back to the client as an empty successful + ``web_search_tool_result`` while the model was re-asked the same query until + the loop cap tripped. When the client sent a native web search tool, the + turn now ends after the first failed search, with Anthropic's + ``web_search_tool_result_error`` object in the tool result and no follow-up + model call. An iteration where some search still succeeded keeps its + follow-up call. + """ + + def setup_method(self): + self.handler = BaseLLMHTTPHandler() + self.logger = WebSearchInterceptionLogger(enabled_providers=["anthropic"]) + self.followup_calls: list[dict] = [] + + async def _fake_acreate(self, **call_kwargs): + self.followup_calls.append(call_kwargs) + return { + "id": "msg_followup", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4-5", + "content": [{"type": "text", "text": "final answer"}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 20, "output_tokens": 5}, + } + + async def _run(self, response: dict, converted_stream: bool = False): + return await self.handler._call_agentic_completion_hooks( + response=response, + model="claude-sonnet-4-5", + messages=[{"role": "user", "content": "who won the world cup"}], + anthropic_messages_provider_config=MagicMock(), + anthropic_messages_optional_request_params={"tools": [get_litellm_web_search_tool()]}, + logging_obj=_logging_obj(self.logger, converted_stream=converted_stream), + stream=False, + custom_llm_provider="anthropic", + kwargs={"_agentic_loop_depth": 0, "max_agentic_loops": 3, WEBSEARCH_EMIT_NATIVE_BLOCKS_KEY: True}, + ) + + @pytest.mark.asyncio + async def test_all_failed_iteration_ends_the_turn_without_a_follow_up_call(self, monkeypatch): + monkeypatch.setattr("litellm.anthropic_interface.messages.acreate", self._fake_acreate) + + with patch.object( + self.logger, + "_execute_search", + side_effect=AuthenticationError("401 Unauthorized", llm_provider="tavily", model="tavily"), + ): + result = await self._run(_response_asking_for_searches("who won the world cup")) + + assert self.followup_calls == [] + assert result["stop_reason"] == "end_turn" + assert INTERNAL_TOOL_NAME not in _tool_use_names(result) + assert _block_types(result) == ["server_tool_use", "web_search_tool_result"] + server_tool_use, tool_result = result["content"] + assert server_tool_use["id"].startswith("srvtoolu_") + assert server_tool_use["input"] == {"query": "who won the world cup"} + assert tool_result["tool_use_id"] == server_tool_use["id"] + assert tool_result["content"] == {"type": "web_search_tool_result_error", "error_code": "unavailable"} + + @pytest.mark.asyncio + async def test_all_failed_iteration_streams_the_error_block(self, monkeypatch): + monkeypatch.setattr("litellm.anthropic_interface.messages.acreate", self._fake_acreate) + + with patch.object( + self.logger, + "_execute_search", + side_effect=AuthenticationError("401 Unauthorized", llm_provider="tavily", model="tavily"), + ): + result = await self._run(_response_asking_for_searches("who won the world cup"), converted_stream=True) + + assert self.followup_calls == [] + assert isinstance(result, FakeAnthropicMessagesStreamIterator) + events = _stream_events(result.response) + started = [event["content_block"] for event in events if event["type"] == "content_block_start"] + assert [block["type"] for block in started] == ["server_tool_use", "web_search_tool_result"] + assert started[1]["tool_use_id"] == started[0]["id"] + assert started[1]["content"] == {"type": "web_search_tool_result_error", "error_code": "unavailable"} + assert [event["delta"]["stop_reason"] for event in events if event["type"] == "message_delta"] == ["end_turn"] + + @pytest.mark.asyncio + async def test_mixed_iteration_keeps_the_follow_up_call(self, monkeypatch): + monkeypatch.setattr("litellm.anthropic_interface.messages.acreate", self._fake_acreate) + + async def search(query, kwargs=None): + if query == "fails": + raise RateLimitError("slow down", llm_provider="tavily", model="tavily") + found = SearchResult(title="Result", url="https://example.com", snippet="A result.", date=None) + return ("Title: Result\nURL: https://example.com", SearchResponse(results=[found])) + + with patch.object(self.logger, "_execute_search", side_effect=search): + result = await self._run(_response_asking_for_searches("fails", "works")) + + assert len(self.followup_calls) == 1 + tool_results = self.followup_calls[0]["messages"][-1]["content"] + assert [block["type"] for block in tool_results] == ["tool_result", "tool_result"] + assert tool_results[0]["content"] == "Search failed: litellm.RateLimitError: slow down" + assert tool_results[1]["content"] == "Title: Result\nURL: https://example.com" + assert result["stop_reason"] == "end_turn" + assert _block_types(result) == [ + "server_tool_use", + "web_search_tool_result", + "server_tool_use", + "web_search_tool_result", + "text", + ] + assert result["content"][0]["input"] == {"query": "fails"} + assert result["content"][1]["content"] == {"type": "web_search_tool_result_error", "error_code": "too_many_requests"} + assert result["content"][2]["input"] == {"query": "works"} + assert result["content"][3]["content"][0]["url"] == "https://example.com" + + class TestMaxAgenticLoopsConfigKnob: def test_from_config_yaml_reads_the_knob(self): logger = WebSearchInterceptionLogger.from_config_yaml( diff --git a/tests/test_litellm/integrations/websearch_interception/test_websearch_native_blocks.py b/tests/test_litellm/integrations/websearch_interception/test_websearch_native_blocks.py index c859f9b2f55..291fb5a5941 100644 --- a/tests/test_litellm/integrations/websearch_interception/test_websearch_native_blocks.py +++ b/tests/test_litellm/integrations/websearch_interception/test_websearch_native_blocks.py @@ -10,6 +10,13 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest +from litellm.exceptions import ( + APIConnectionError, + AuthenticationError, + BadRequestError, + RateLimitError, + Timeout, +) from litellm.integrations.websearch_interception.handler import ( WEBSEARCH_EMIT_NATIVE_BLOCKS_KEY, WEBSEARCH_NATIVE_BLOCKS_METADATA_KEY, @@ -27,6 +34,10 @@ from litellm.types.integrations.custom_logger import ( AgenticLoopPlan, AgenticLoopRequestPatch, ) +from litellm.types.integrations.websearch_interception import ( + SearchFailed, + SearchSucceeded, +) def _make_search_response() -> SearchResponse: @@ -48,6 +59,10 @@ def _make_search_response() -> SearchResponse: ) +def _succeeded_outcome() -> SearchSucceeded: + return SearchSucceeded(text="Title: LiteLLM Docs\nURL: https://docs.litellm.ai/", response=_make_search_response()) + + class TestIsAnthropicNativeWebSearchTool: """The detector must match native tools without catching look-alikes.""" @@ -227,12 +242,10 @@ class TestBuildPlanAttachesBlocks: messages=[{"role": "user", "content": "hi"}], max_tokens=1024, ) - structured = [_make_search_response()] - with patch.object( logger, "_build_anthropic_request_patch", - new=AsyncMock(return_value=(patch_obj, structured)), + new=AsyncMock(return_value=(patch_obj, (_succeeded_outcome(),))), ): plan = await logger.async_build_agentic_loop_plan( tools={"tool_calls": tool_calls, "thinking_blocks": []}, @@ -277,7 +290,7 @@ class TestBuildPlanAttachesBlocks: with patch.object( logger, "_build_anthropic_request_patch", - new=AsyncMock(return_value=(patch_obj, [_make_search_response()])), + new=AsyncMock(return_value=(patch_obj, (_succeeded_outcome(),))), ): plan = await logger.async_build_agentic_loop_plan( tools={"tool_calls": tool_calls, "thinking_blocks": []}, @@ -294,6 +307,145 @@ class TestBuildPlanAttachesBlocks: assert WEBSEARCH_NATIVE_BLOCKS_METADATA_KEY not in plan.metadata +class TestFailedSearchOutcome: + """A search that raises becomes a ``web_search_tool_result_error`` block, coded by exception type.""" + + @pytest.mark.parametrize( + ("error", "expected_code"), + [ + (RateLimitError("slow down", llm_provider="tavily", model="tavily"), "too_many_requests"), + (BadRequestError("bad query", model="tavily", llm_provider="tavily"), "invalid_tool_input"), + (AuthenticationError("401 Unauthorized", llm_provider="tavily", model="tavily"), "unavailable"), + (APIConnectionError("connection refused", llm_provider="tavily", model="tavily"), "unavailable"), + (Timeout("timed out", model="tavily", llm_provider="tavily"), "unavailable"), + (RuntimeError("boom"), "unavailable"), + ], + ) + def test_error_block_carries_the_mapped_error_code(self, error, expected_code): + outcome = WebSearchTransformation.search_outcome(error) + + assert outcome == SearchFailed(error_code=expected_code, message=str(error)) + assert WebSearchTransformation.build_web_search_outcome_block("srvtoolu_x", outcome) == { + "type": "web_search_tool_result", + "tool_use_id": "srvtoolu_x", + "content": {"type": "web_search_tool_result_error", "error_code": expected_code}, + } + assert WebSearchTransformation.search_outcome_text(outcome) == f"Search failed: {error}" + + def test_succeeded_outcome_still_yields_result_items(self): + outcome = WebSearchTransformation.search_outcome(("Title: x", _make_search_response())) + + assert outcome == SearchSucceeded(text="Title: x", response=_make_search_response()) + block = WebSearchTransformation.build_web_search_outcome_block("srvtoolu_x", outcome) + assert [item["type"] for item in block["content"]] == ["web_search_result", "web_search_result"] + assert block["content"][0]["url"] == "https://docs.litellm.ai/" + assert WebSearchTransformation.search_outcome_text(outcome) == "Title: x" + + @pytest.mark.asyncio + async def test_all_failed_iteration_terminates_when_native_blocks_are_emitted(self): + logger = WebSearchInterceptionLogger(enabled_providers=["bedrock"]) + tool_calls = [ + {"id": "toolu_one", "type": "tool_use", "name": "litellm_web_search", "input": {"query": "q1"}}, + {"id": "toolu_two", "type": "tool_use", "name": "litellm_web_search", "input": {"query": "q2"}}, + ] + + with patch.object( + logger, + "_execute_search", + side_effect=AuthenticationError("401 Unauthorized", llm_provider="tavily", model="tavily"), + ): + plan = await logger.async_build_agentic_loop_plan( + tools={"tool_calls": tool_calls, "thinking_blocks": []}, + model="bedrock/claude", + messages=[{"role": "user", "content": "hi"}], + response=MagicMock(), + anthropic_messages_provider_config=None, + anthropic_messages_optional_request_params={}, + logging_obj=MagicMock(model_call_details={}), + stream=False, + kwargs={WEBSEARCH_EMIT_NATIVE_BLOCKS_KEY: True}, + ) + + assert plan.run_agentic_loop is False + assert plan.terminate is True + assert plan.stop_reason == "web_search_failed" + blocks = plan.metadata[WEBSEARCH_NATIVE_BLOCKS_METADATA_KEY] + assert [b["type"] for b in blocks] == [ + "server_tool_use", + "web_search_tool_result", + "server_tool_use", + "web_search_tool_result", + ] + assert blocks[1]["tool_use_id"] == blocks[0]["id"] + assert blocks[1]["content"] == {"type": "web_search_tool_result_error", "error_code": "unavailable"} + assert blocks[3]["tool_use_id"] == blocks[2]["id"] + assert blocks[3]["content"] == {"type": "web_search_tool_result_error", "error_code": "unavailable"} + + @pytest.mark.asyncio + async def test_all_failed_iteration_keeps_the_follow_up_without_native_blocks(self): + logger = WebSearchInterceptionLogger(enabled_providers=["bedrock"]) + tool_calls = [ + {"id": "toolu_one", "type": "tool_use", "name": "litellm_web_search", "input": {"query": "q1"}}, + ] + + with patch.object( + logger, + "_execute_search", + side_effect=AuthenticationError("401 Unauthorized", llm_provider="tavily", model="tavily"), + ): + plan = await logger.async_build_agentic_loop_plan( + tools={"tool_calls": tool_calls, "thinking_blocks": []}, + model="bedrock/claude", + messages=[{"role": "user", "content": "hi"}], + response=MagicMock(), + anthropic_messages_provider_config=None, + anthropic_messages_optional_request_params={}, + logging_obj=MagicMock(model_call_details={}), + stream=False, + kwargs={}, + ) + + assert plan.run_agentic_loop is True + assert plan.terminate is False + assert plan.request_patch is not None + tool_results = plan.request_patch.messages[-1]["content"] + assert "Search failed: litellm.AuthenticationError: 401 Unauthorized" in tool_results[0]["content"] + + @pytest.mark.asyncio + async def test_mixed_iteration_keeps_the_follow_up_and_pairs_each_block(self): + logger = WebSearchInterceptionLogger(enabled_providers=["bedrock"]) + tool_calls = [ + {"id": "toolu_one", "type": "tool_use", "name": "litellm_web_search", "input": {"query": "fails"}}, + {"id": "toolu_two", "type": "tool_use", "name": "litellm_web_search", "input": {"query": "works"}}, + ] + + async def search(query, kwargs=None): + if query == "fails": + raise RateLimitError("slow down", llm_provider="tavily", model="tavily") + return ("Title: x", _make_search_response()) + + with patch.object(logger, "_execute_search", side_effect=search): + plan = await logger.async_build_agentic_loop_plan( + tools={"tool_calls": tool_calls, "thinking_blocks": []}, + model="bedrock/claude", + messages=[{"role": "user", "content": "hi"}], + response=MagicMock(), + anthropic_messages_provider_config=None, + anthropic_messages_optional_request_params={}, + logging_obj=MagicMock(model_call_details={}), + stream=False, + kwargs={WEBSEARCH_EMIT_NATIVE_BLOCKS_KEY: True}, + ) + + assert plan.run_agentic_loop is True + assert plan.terminate is False + blocks = plan.metadata[WEBSEARCH_NATIVE_BLOCKS_METADATA_KEY] + assert blocks[0]["input"] == {"query": "fails"} + assert blocks[1]["content"] == {"type": "web_search_tool_result_error", "error_code": "too_many_requests"} + assert blocks[2]["input"] == {"query": "works"} + assert blocks[3]["content"][0]["url"] == "https://docs.litellm.ai/" + + class TestPostHookInjectsBlocks: """The post-hook must prepend blocks; absent metadata is a no-op.""" @@ -437,13 +589,17 @@ class TestShortCircuitEmitsNativeBlocks: assert block_types == ["text"] @pytest.mark.asyncio - async def test_native_short_circuit_failure_still_emits_blocks(self): - """Search failure on native path: emit blocks with empty results + - the legacy text-error block, so the client gets a well-formed - response instead of a malformed half-shape.""" + async def test_native_short_circuit_failure_emits_the_error_block(self): + """Search failure on native path: the tool result carries Anthropic's + error object (rendered as "Web search error: " by the client) + next to the legacy text-error block.""" logger = WebSearchInterceptionLogger(enabled_providers=["github_copilot"]) - with patch.object(logger, "_execute_search", side_effect=RuntimeError("boom")): + with patch.object( + logger, + "_execute_search", + side_effect=RateLimitError("slow down", llm_provider="tavily", model="tavily"), + ): result = await logger.try_short_circuit_search( model="github_copilot/claude-sonnet-4", messages=[{"role": "user", "content": "search query"}], @@ -455,9 +611,10 @@ class TestShortCircuitEmitsNativeBlocks: block_types = [b["type"] for b in result["content"]] assert block_types == ["server_tool_use", "web_search_tool_result", "text"] tool_result = result["content"][1] - assert tool_result["content"] == [] + assert tool_result["tool_use_id"] == result["content"][0]["id"] + assert tool_result["content"] == {"type": "web_search_tool_result_error", "error_code": "too_many_requests"} text_block = result["content"][2] - assert "Search failed" in text_block["text"] + assert text_block["text"] == "Search failed: litellm.RateLimitError: slow down" class TestLegacyPathMatchesNewPath: @@ -489,7 +646,7 @@ class TestLegacyPathMatchesNewPath: patch.object( logger, "_build_anthropic_request_patch", - new=AsyncMock(return_value=(patch_obj, [_make_search_response()])), + new=AsyncMock(return_value=(patch_obj, (_succeeded_outcome(),))), ), patch( "litellm.integrations.websearch_interception.handler.anthropic_messages.acreate", diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py index 686a792fa0f..4fe3d410ef8 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py @@ -15,6 +15,7 @@ from litellm.litellm_core_utils.llm_cost_calc.utils import ( _is_off_peak, _is_within_off_peak_window, apply_off_peak_pricing, + apply_provider_cache_read_default, calculate_cache_writing_cost, generic_cost_per_token, get_billed_token_rates, @@ -96,6 +97,19 @@ def test_generic_cost_per_token_bills_cache_reads_at_input_rate_when_no_cache_re assert completion_cost == pytest.approx(380 * 9.7e-7) +def test_apply_provider_cache_read_default_only_derives_a_rate_for_fireworks() -> None: + openai_info: ModelInfo = {"input_cost_per_token": 2e-6} + fireworks_info: ModelInfo = {"input_cost_per_token": 2e-6} + + assert apply_provider_cache_read_default(openai_info, "openai") is openai_info + assert apply_provider_cache_read_default(openai_info, None) is openai_info + + processed_fireworks_info = apply_provider_cache_read_default(fireworks_info, "fireworks_ai") + + assert processed_fireworks_info is not fireworks_info + assert processed_fireworks_info["cache_read_input_token_cost"] == pytest.approx(2e-6 * 0.5) + + def test_generic_cost_per_token_prefers_audio_per_second_rate() -> None: model_info: ModelInfo = { "key": "gemini-embedding-2", @@ -239,9 +253,7 @@ def test_reasoning_tokens_no_price_set(_local_model_cost_map): model_cost_map["input_cost_per_token"] * usage.prompt_tokens, 10, ) - print(f"completion_cost: {completion_cost}") expected_completion_cost = model_cost_map["output_cost_per_token"] * usage.completion_tokens - print(f"expected_completion_cost: {expected_completion_cost}") assert round(completion_cost, 10) == round( expected_completion_cost, 10, diff --git a/tests/test_litellm/litellm_core_utils/llm_response_utils/test_response_metadata.py b/tests/test_litellm/litellm_core_utils/llm_response_utils/test_response_metadata.py index a06b6bbf3cc..3379879a8a6 100644 --- a/tests/test_litellm/litellm_core_utils/llm_response_utils/test_response_metadata.py +++ b/tests/test_litellm/litellm_core_utils/llm_response_utils/test_response_metadata.py @@ -9,11 +9,14 @@ import asyncio import datetime from unittest.mock import MagicMock +import pytest + import litellm.litellm_core_utils.llm_response_utils.response_metadata as response_metadata_mod import litellm.proxy.common_request_processing as common_request_processing_mod from litellm.litellm_core_utils.litellm_logging import Logging from litellm.litellm_core_utils.llm_response_utils.response_metadata import ( ResponseMetadata, + _union_duration_ms, response_timing_metrics, update_response_metadata, ) @@ -70,9 +73,7 @@ class TestCallbackDurationMs: def test_update_response_metadata_includes_callback_duration(self): """End-to-end: update_response_metadata should propagate callback_duration_ms.""" result = ModelResponse() - logging_obj = self._make_logging_obj( - callback_duration_ms=5.5, llm_api_duration_ms=800.0 - ) + logging_obj = self._make_logging_obj(callback_duration_ms=5.5, llm_api_duration_ms=800.0) logging_obj._response_cost_calculator = MagicMock(return_value=0.001) logging_obj.litellm_call_id = "test-call-id" @@ -231,11 +232,24 @@ class TestResponseTimingMetrics: START = datetime.datetime(2025, 1, 1, 0, 0, 0) END = datetime.datetime(2025, 1, 1, 0, 0, 1) - def _make_logging_obj(self, llm_api_duration_ms=None, caching_details=None): + def _make_logging_obj( + self, + llm_api_duration_ms: float | None = None, + llm_api_timing_windows: object = None, + caching_details: dict[str, object] | None = None, + received_at: datetime.datetime | str | None = None, + ) -> MagicMock: logging_obj = MagicMock() logging_obj.model_call_details = {} if llm_api_duration_ms is not None: logging_obj.model_call_details["llm_api_duration_ms"] = llm_api_duration_ms + if received_at is not None or llm_api_timing_windows is not None: + metadata = {} + if received_at is not None: + metadata["litellm_received_at"] = received_at + if llm_api_timing_windows is not None: + metadata["llm_api_timing_windows"] = llm_api_timing_windows + logging_obj.model_call_details["litellm_params"] = {"metadata": metadata} logging_obj.caching_details = caching_details return logging_obj @@ -246,6 +260,104 @@ class TestResponseTimingMetrics: "litellm_overhead_time_ms": 100.0, } + def test_window_starts_at_proxy_receive_when_stamped(self): + received_at = self.START.astimezone(datetime.timezone.utc) - datetime.timedelta(seconds=3) + logging_obj = self._make_logging_obj(llm_api_duration_ms=900.0, received_at=received_at) + + result = response_timing_metrics(self.START, self.END, logging_obj) + + assert result["_response_ms"] == pytest.approx(4000.0) + assert result["litellm_overhead_time_ms"] == pytest.approx(3100.0) + + def test_receive_anchored_window_subtracts_all_provider_attempts(self): + logging_obj = self._make_logging_obj( + llm_api_duration_ms=300.0, + llm_api_timing_windows=( + (self.START.timestamp(), self.START.timestamp() + 0.3), + (self.START.timestamp() + 0.4, self.START.timestamp() + 0.8), + ), + received_at=self.START, + ) + + result = response_timing_metrics(self.START, self.END, logging_obj) + + assert result["_response_ms"] == pytest.approx(1000.0) + assert result["litellm_overhead_time_ms"] == pytest.approx(300.0) + + def test_sdk_window_subtracts_current_provider_attempt(self): + logging_obj = self._make_logging_obj( + llm_api_duration_ms=300.0, + llm_api_timing_windows=((self.START.timestamp(), self.START.timestamp() + 0.3),), + ) + + result = response_timing_metrics(self.START, self.END, logging_obj) + + assert result["_response_ms"] == pytest.approx(1000.0) + assert result["litellm_overhead_time_ms"] == pytest.approx(700.0) + + def test_receive_anchored_window_unions_nested_and_retry_windows(self): + windows = ( + (self.START.timestamp(), self.START.timestamp() + 0.3), + (self.START.timestamp(), self.START.timestamp() + 0.3), + (self.START.timestamp() + 0.4, self.START.timestamp() + 0.7), + ) + logging_obj = self._make_logging_obj( + llm_api_duration_ms=300.0, + llm_api_timing_windows=windows, + received_at=self.START, + ) + + result = response_timing_metrics(self.START, self.END, logging_obj) + + assert result["litellm_overhead_time_ms"] == pytest.approx(400.0) + assert _union_duration_ms(windows, self.START.timestamp(), self.END.timestamp()) == pytest.approx(600.0) + + def test_receive_anchored_window_ignores_seeded_windows_outside_window(self): + windows = ( + (self.START.timestamp() - 10.0, self.START.timestamp() - 1.0), + (self.END.timestamp() + 1.0, self.END.timestamp() + 2.0), + ) + logging_obj = self._make_logging_obj( + llm_api_duration_ms=300.0, + llm_api_timing_windows=windows, + received_at=self.START, + ) + + result = response_timing_metrics(self.START, self.END, logging_obj) + + assert result["litellm_overhead_time_ms"] == pytest.approx(700.0) + + def test_receive_anchored_window_falls_back_to_current_provider_attempt(self): + logging_obj = self._make_logging_obj( + llm_api_duration_ms=300.0, + received_at=self.START, + ) + + result = response_timing_metrics(self.START, self.END, logging_obj) + + assert result["_response_ms"] == pytest.approx(1000.0) + assert result["litellm_overhead_time_ms"] == pytest.approx(700.0) + + def test_cache_hit_window_starts_at_proxy_receive_when_stamped(self): + received_at = self.START.astimezone(datetime.timezone.utc) - datetime.timedelta(seconds=3) + logging_obj = self._make_logging_obj( + caching_details={"cache_hit": True, "cache_duration_ms": 250.0}, + received_at=received_at, + ) + + result = response_timing_metrics(self.START, self.END, logging_obj) + + assert result["_response_ms"] == pytest.approx(4000.0) + assert result["litellm_overhead_time_ms"] == pytest.approx(3750.0) + + def test_non_datetime_proxy_receive_falls_back_to_start_time(self): + logging_obj = self._make_logging_obj(llm_api_duration_ms=900.0, received_at="bad") + + result = response_timing_metrics(self.START, self.END, logging_obj) + + assert result["_response_ms"] == pytest.approx(1000.0) + assert result["litellm_overhead_time_ms"] == pytest.approx(100.0) + def test_overhead_omitted_when_no_provider_or_cache_duration_recorded(self): logging_obj = self._make_logging_obj() assert response_timing_metrics(self.START, self.END, logging_obj) == {"_response_ms": 1000.0} @@ -358,6 +470,28 @@ class TestDetailedTiming: assert hidden.get("timing_pre_processing_ms") == 20.0 assert hidden.get("timing_post_processing_ms") == 10.0 # 530 - 20 - 500 + def test_detailed_timing_pre_processing_uses_receive_anchor(self, monkeypatch): + monkeypatch.setattr(response_metadata_mod, "LITELLM_DETAILED_TIMING", True) + + result = ModelResponse() + received_at = datetime.datetime.now(datetime.timezone.utc) + start = received_at + datetime.timedelta(milliseconds=200) + api_call_start = start.replace(tzinfo=None) + end = start + datetime.timedelta(milliseconds=530) + logging_obj = self._make_logging_obj( + llm_api_duration_ms=500.0, + api_call_start_time=api_call_start, + ) + logging_obj.model_call_details["litellm_params"] = {"metadata": {"litellm_received_at": received_at}} + + metadata = ResponseMetadata(result) + metadata.set_timing_metrics(start, end, logging_obj) + metadata.apply() + + hidden = result._hidden_params + assert hidden.get("timing_pre_processing_ms") == pytest.approx(200.0) + assert hidden.get("timing_post_processing_ms") == pytest.approx(30.0) + def test_detailed_timing_absent_when_disabled(self, monkeypatch): """When LITELLM_DETAILED_TIMING is false, no detailed timing keys.""" monkeypatch.setattr(response_metadata_mod, "LITELLM_DETAILED_TIMING", False) @@ -377,9 +511,7 @@ class TestDetailedTiming: def test_detailed_timing_headers_in_custom_headers(self, monkeypatch): """When LITELLM_DETAILED_TIMING is true, headers flow to get_custom_headers.""" - monkeypatch.setattr( - common_request_processing_mod, "LITELLM_DETAILED_TIMING", True - ) + monkeypatch.setattr(common_request_processing_mod, "LITELLM_DETAILED_TIMING", True) user_api_key_dict = UserAPIKeyAuth(api_key="sk-test") hidden_params = { @@ -402,9 +534,7 @@ class TestDetailedTiming: def test_detailed_timing_headers_absent_when_disabled(self, monkeypatch): """When LITELLM_DETAILED_TIMING is false, no timing headers emitted.""" - monkeypatch.setattr( - common_request_processing_mod, "LITELLM_DETAILED_TIMING", False - ) + monkeypatch.setattr(common_request_processing_mod, "LITELLM_DETAILED_TIMING", False) user_api_key_dict = UserAPIKeyAuth(api_key="sk-test") hidden_params = { diff --git a/tests/test_litellm/litellm_core_utils/test_logging_utils.py b/tests/test_litellm/litellm_core_utils/test_logging_utils.py index b446021a7dc..672595b85d6 100644 --- a/tests/test_litellm/litellm_core_utils/test_logging_utils.py +++ b/tests/test_litellm/litellm_core_utils/test_logging_utils.py @@ -2,18 +2,42 @@ Tests for litellm.litellm_core_utils.logging_utils — base64 truncation helpers. """ +import datetime import threading +from unittest.mock import MagicMock import pytest from litellm.litellm_core_utils import logging_utils from litellm.litellm_core_utils.logging_utils import ( - format_base64_size, + _set_duration_in_model_call_details, _truncate_base64_in_string, + format_base64_size, truncate_base64_in_messages, truncate_base64_in_messages_async, ) + +class TestSetDurationInModelCallDetails: + def test_records_provider_attempt_windows_in_shared_metadata(self): + metadata = {"request_id": "test"} + logging_obj = MagicMock() + logging_obj.model_call_details = {"litellm_params": {"metadata": metadata}} + first_start = datetime.datetime(2025, 1, 1, 0, 0, 0) + first_end = first_start + datetime.timedelta(milliseconds=300) + second_start = datetime.datetime(2025, 1, 1, 0, 0, 1) + second_end = second_start + datetime.timedelta(milliseconds=700) + + _set_duration_in_model_call_details(logging_obj, first_start, first_end) + _set_duration_in_model_call_details(logging_obj, second_start, second_end) + + assert metadata["llm_api_timing_windows"] == ( + (first_start.timestamp(), first_end.timestamp()), + (second_start.timestamp(), second_end.timestamp()), + ) + assert logging_obj.model_call_details["llm_api_duration_ms"] == pytest.approx(700.0) + + # --------------------------------------------------------------------------- # format_base64_size # --------------------------------------------------------------------------- @@ -157,10 +181,7 @@ class TestTruncateBase64InMessages: } ] result = truncate_base64_in_messages(messages) - assert ( - result[0]["content"][0]["image_url"]["url"] - == f"data:image/png;base64,{short}" - ) + assert result[0]["content"][0]["image_url"]["url"] == f"data:image/png;base64,{short}" # --------------------------------------------------------------------------- diff --git a/tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py b/tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py index 133d6e502f4..945033c5cac 100644 --- a/tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py +++ b/tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py @@ -1828,7 +1828,11 @@ class TestAnthropicThinkingSignatureSelfHeal: assert out[0] is msgs[0] - def test_flatten_unencrypted_web_search_results_leaves_error_blocks_alone(self): + def test_flatten_unencrypted_web_search_results_flattens_error_blocks(self): + """A failed intercepted search is replayed by the client as the error + object LiteLLM emitted. Anthropic rejects a replayed ``server_tool_use`` + it never issued, so the pair is flattened to text the same way a + successful unencrypted result is.""" from litellm.llms.anthropic.common_utils import ( flatten_unencrypted_web_search_results_in_anthropic_messages, ) @@ -1837,6 +1841,7 @@ class TestAnthropicThinkingSignatureSelfHeal: { "role": "assistant", "content": [ + {"type": "server_tool_use", "id": "srvtoolu_1", "name": "web_search", "input": {"query": "q"}}, { "type": "web_search_tool_result", "tool_use_id": "srvtoolu_1", @@ -1844,14 +1849,18 @@ class TestAnthropicThinkingSignatureSelfHeal: "type": "web_search_tool_result_error", "error_code": "max_uses_exceeded", }, - } + }, ], } ] - out = flatten_unencrypted_web_search_results_in_anthropic_messages(msgs) + once = flatten_unencrypted_web_search_results_in_anthropic_messages(msgs) + twice = flatten_unencrypted_web_search_results_in_anthropic_messages(once) - assert out[0] is msgs[0] + assert once[0]["content"] == [ + {"type": "text", "text": "Web search results for 'q':\n\nSearch failed: max_uses_exceeded"} + ] + assert json.dumps(twice) == json.dumps(once) def test_sanitize_tool_use_ids_in_anthropic_messages(self): from litellm.llms.anthropic.common_utils import ( diff --git a/tests/test_litellm/llms/base_llm/realtime/__init__.py b/tests/test_litellm/llms/base_llm/realtime/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/base_llm/realtime/test_transcription_protocol.py b/tests/test_litellm/llms/base_llm/realtime/test_transcription_protocol.py new file mode 100644 index 00000000000..d38cc480e85 --- /dev/null +++ b/tests/test_litellm/llms/base_llm/realtime/test_transcription_protocol.py @@ -0,0 +1,128 @@ +import base64 +import json + +import pytest + +from litellm.llms.base_llm.realtime.transcription_protocol import ( + RealtimeTranscriptionProtocolError, + completed_event, + decode_pcm16_append, + parse_transcription_session_update, + transcription_session, +) + + +def _session_update(session: dict[str, object]) -> str: + return json.dumps({"type": "session.update", "session": session}) + + +def test_ga_layout_parses_format_language_and_turn_detection(): + update = parse_transcription_session_update( + _session_update( + { + "type": "transcription", + "audio": { + "input": { + "format": {"type": "audio/pcm", "rate": 16_000, "channels": 1}, + "transcription": {"model": "chirp_3", "language": "pt-BR", "prompt": "names"}, + "turn_detection": {"type": "server_vad", "threshold": 0.5}, + } + }, + } + ) + ) + assert update.session_type == "transcription" + assert update.audio_format is not None + assert (update.audio_format.layout, update.audio_format.rate, update.audio_format.channels) == ("ga", 16_000, 1) + assert update.audio_format.is_pcm16 + assert (update.model, update.language) == ("chirp_3", "pt-BR") + assert update.unsupported_transcription_keys == ("prompt",) + assert update.turn_detection_type == "server_vad" + assert not update.turn_detection_disabled + + +def test_beta_layout_parses_flat_fields(): + update = parse_transcription_session_update( + json.dumps( + { + "type": "transcription_session.update", + "session": { + "input_audio_format": "pcm16", + "input_audio_transcription": {"model": "whisper-1"}, + "turn_detection": None, + }, + } + ) + ) + assert update.audio_format is not None + assert (update.audio_format.layout, update.audio_format.encoding) == ("beta", "pcm16") + assert update.audio_format.is_pcm16 + assert update.model == "whisper-1" + assert update.turn_detection_disabled + + +def test_absent_turn_detection_is_not_disabled(): + update = parse_transcription_session_update(_session_update({"audio": {"input": {"transcription": {}}}})) + assert update.turn_detection is None + assert not update.turn_detection_disabled + + +@pytest.mark.parametrize( + ("payload", "message"), + [ + ("not json", "invalid JSON object"), + ("[]", "must be a JSON object"), + (json.dumps({"type": "response.create"}), "expected session.update"), + (_session_update({}), "requires a session object"), + (_session_update({"input_audio_format": "pcm16", "audio": {"input": {"format": "pcm16"}}}), "either beta or GA"), + (_session_update({"input_audio_transcription": {}, "audio": {"input": {"transcription": {}}}}), "either beta or GA"), + (_session_update({"audio": {"input": {"format": {"rate": "fast"}}}}), "must be an integer"), + (_session_update({"audio": {"input": {"format": {"rate": True}}}}), "must be an integer"), + (_session_update({"audio": {"input": {"transcription": {"language": 7}}}}), "must be a string"), + (_session_update({"audio": {"input": {"transcription": []}}}), "must be an object"), + ], +) +def test_malformed_session_updates_are_rejected(payload: str, message: str): + with pytest.raises(RealtimeTranscriptionProtocolError, match=message): + parse_transcription_session_update(payload) + + +def test_decode_pcm16_append_returns_the_raw_samples(): + assert decode_pcm16_append(base64.b64encode(b"\x01\x02\x03\x04").decode()) == b"\x01\x02\x03\x04" + + +@pytest.mark.parametrize( + ("audio", "message"), + [ + (None, "must be a base64 string"), + ("@@@", "must be valid base64"), + (base64.b64encode(b"\x01\x02\x03").decode(), "complete samples"), + ], +) +def test_decode_pcm16_append_rejects_bad_audio(audio: object, message: str): + with pytest.raises(RealtimeTranscriptionProtocolError, match=message): + decode_pcm16_append(audio) + + +def test_decode_pcm16_append_enforces_the_backlog_limit(): + with pytest.raises(RealtimeTranscriptionProtocolError, match="backlog limit"): + decode_pcm16_append(base64.b64encode(b"\x00" * 8).decode(), max_encoded_bytes=4) + + +def test_transcription_session_reflects_negotiated_settings(): + manual = transcription_session(session_id="sess_1", model="chirp_3", sample_rate=16_000, language=None, server_vad=False) + assert manual["id"] == "sess_1" + assert manual["audio"]["input"] == { + "format": {"type": "audio/pcm", "rate": 16_000}, + "transcription": {"model": "chirp_3"}, + "turn_detection": None, + } + vad = transcription_session(session_id="sess_1", model="chirp_3", sample_rate=24_000, language="en-US", server_vad=True) + assert vad["audio"]["input"]["transcription"] == {"model": "chirp_3", "language": "en-US"} + assert vad["audio"]["input"]["turn_detection"] == {"type": "server_vad"} + + +def test_completed_event_carries_usage_only_when_billed(): + assert "usage" not in completed_event("item_1", "hello", None) + billed = completed_event("item_1", "hello", {"type": "duration", "seconds": 2.5}) + assert (billed["item_id"], billed["transcript"], billed["usage"]) == ("item_1", "hello", {"type": "duration", "seconds": 2.5}) diff --git a/tests/test_litellm/llms/bedrock/batches/test_handler.py b/tests/test_litellm/llms/bedrock/batches/test_handler.py index 03daafcad72..e69098a460d 100644 --- a/tests/test_litellm/llms/bedrock/batches/test_handler.py +++ b/tests/test_litellm/llms/bedrock/batches/test_handler.py @@ -8,10 +8,14 @@ the tests don't hit AWS. from __future__ import annotations +import json +from collections.abc import Iterator, Mapping from datetime import datetime, timedelta, timezone +from typing import Final from unittest.mock import MagicMock, patch import pytest +from botocore.awsrequest import AWSPreparedRequest, AWSResponse from litellm.llms.bedrock.batches.handler import ( # noqa: E402 @@ -570,3 +574,58 @@ def test_cancel_batch_stops_and_polls_the_job_with_the_tagged_session(monkeypatc fake_bedrock.stop_model_invocation_job.assert_called_once_with(jobIdentifier=JOB_ARN) assert batch.status == "cancelled" assert [kwargs["aws_access_key_id"] for kwargs in bedrock_client_kwargs] == ["ASIABATCHCANCELTAGGED"] * 2 + + +class _JsonBody: + def __init__(self, payload: bytes) -> None: + self._payload: Final = payload + + def stream(self) -> Iterator[bytes]: + return iter((self._payload,)) + + +class _AuthorizationRecorder: + def __init__(self, body: Mapping[str, object]) -> None: + self._payload: Final = json.dumps(body, default=str).encode() + self.authorization_headers: tuple[str, ...] = () + + def send(self, request: AWSPreparedRequest) -> AWSResponse: + raw_authorization: Final = request.headers["Authorization"] + authorization: Final = ( + raw_authorization.decode() if isinstance(raw_authorization, bytes) else str(raw_authorization) + ) + self.authorization_headers = (*self.authorization_headers, authorization) + return AWSResponse(request.url, 200, {"content-type": "application/json"}, _JsonBody(self._payload)) + + +def test_retrieve_signs_with_deployment_credentials_when_env_bearer_token_is_set(monkeypatch): + """A proxy-wide AWS_BEARER_TOKEN_BEDROCK must not override the deployment's own SigV4 credentials.""" + monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", "env-bearer-token") + recorder: Final = _AuthorizationRecorder(_fake_boto3_response()) + + with patch("botocore.httpsession.URLLib3Session.send", recorder.send): + batch = BedrockBatchesHandler._handle_model_invocation_job_status( + batch_id=JOB_ARN, + aws_access_key_id="AKIADEPLOYMENTKEY", + aws_secret_access_key="deployment-secret", + ) + + assert batch.status == "completed" + assert len(recorder.authorization_headers) == 1 + assert recorder.authorization_headers[0].startswith("AWS4-HMAC-SHA256 Credential=AKIADEPLOYMENTKEY/") + + +def test_cancel_signs_with_deployment_credentials_when_env_bearer_token_is_set(monkeypatch): + monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", "env-bearer-token") + recorder: Final = _AuthorizationRecorder(_fake_boto3_response(status="Stopped")) + + with patch("botocore.httpsession.URLLib3Session.send", recorder.send): + batch = BedrockBatchesHandler.cancel_batch( + batch_id=JOB_ARN, + aws_access_key_id="AKIADEPLOYMENTKEY", + aws_secret_access_key="deployment-secret", + ) + + assert batch.status == "cancelled" + assert len(recorder.authorization_headers) == 2 + assert all(h.startswith("AWS4-HMAC-SHA256 Credential=AKIADEPLOYMENTKEY/") for h in recorder.authorization_headers) diff --git a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py index cd33e8d34d8..db5da28c024 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py @@ -458,6 +458,34 @@ def test_reasoning_with_forced_tool_choice_switches_to_auto(): assert optional_params["tool_choice"] == {"auto": {}} +@pytest.mark.parametrize( + "model, param, value, expected_max_tokens", + [ + ("us.openai.gpt-6-astra", "max_tokens", 1, 16), + ("us.openai.gpt-6-astra", "max_completion_tokens", 1, 16), + ("us.openai.gpt-6-astra", "max_tokens", 64, 64), + ("us.xai.grok-4.6", "max_tokens", 1, 16), + ("global.xai.grok-4.6", "max_completion_tokens", 1, 16), + ("us.xai.grok-4.6", "max_tokens", 32, 32), + ("anthropic.claude-sonnet-4-5-20250929-v1:0", "max_tokens", 1, 1), + ("arn:aws:bedrock:us-east-1:123456789012:inference-profile/us.openai.gpt-6-astra", "max_tokens", 1, 16), + ("arn:aws:bedrock:us-east-1:123456789012:inference-profile/global.xai.grok-4.6", "max_tokens", 1, 16), + ("arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/abc123xyz", "max_tokens", 1, 1), + ], +) +def test_map_openai_params_enforces_minimum_max_tokens_for_openai_compat_models( + model: str, param: str, value: int, expected_max_tokens: int +): + optional_params = AmazonConverseConfig().map_openai_params( + non_default_params={param: value}, + optional_params={}, + model=model, + drop_params=False, + ) + + assert optional_params["maxTokens"] == expected_max_tokens + + @pytest.mark.parametrize( "model", [ diff --git a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py index 95dceccb2f5..82220d53375 100644 --- a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py +++ b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py @@ -1,4 +1,5 @@ import asyncio +import base64 import json import logging import threading @@ -2064,6 +2065,7 @@ async def _run_async_realtime_with_backend_failure(client_ws): provider_config = Mock() provider_config.get_complete_url.return_value = "wss://backend.example/live" provider_config.validate_environment.return_value = {} + provider_config.open_backend = AsyncMock(return_value=None) with patch.object( handler, @@ -3707,3 +3709,149 @@ def test_image_edit_handler_keeps_the_sync_transform(): assert config.transform_calls == ["sync"] assert captured["body"] == {"transformed_by": "sync"} assert response.data[0].b64_json == "sync" + + +class _ScriptedClientWebSocket(_FakeClientWebSocket): + def __init__(self, messages: list[str], last_event_type: str) -> None: + super().__init__() + self._messages: Final = list(messages) + self._last_event_type: Final = last_event_type + self._backend_done: Final = asyncio.Event() + + async def receive_text(self) -> str: + if self._messages: + return self._messages.pop(0) + await asyncio.wait_for(self._backend_done.wait(), timeout=5) + raise RuntimeError("client went away") + + async def send_text(self, payload: str) -> None: + await super().send_text(payload) + if json.loads(payload).get("type") == self._last_event_type: + self._backend_done.set() + + def sent_events(self) -> list[dict[str, object]]: + return [json.loads(payload) for name, payload in self.events if name == "send_text"] + + +@pytest.mark.asyncio +async def test_async_realtime_bridges_a_transcription_session_through_the_provider_backend(): + import websockets.exceptions # noqa: F401 # binds the submodule so async_realtime's except clause resolves, as in the proxy process + + from datetime import timedelta + + from google.cloud.speech_v2.types import ( + RecognitionResponseMetadata, + SpeechRecognitionAlternative, + StreamingRecognitionResult, + StreamingRecognizeResponse, + ) + + from litellm.llms.vertex_ai.audio_transcription.realtime_backend import SpeechStreamingBackend + from litellm.llms.vertex_ai.audio_transcription.realtime_transformation import VertexChirpRealtimeConfig + + def google_response(transcript: str, is_final: bool, billed: float) -> StreamingRecognizeResponse: + return StreamingRecognizeResponse( + results=[ + StreamingRecognitionResult( + alternatives=[SpeechRecognitionAlternative(transcript=transcript)], is_final=is_final + ) + ], + metadata=RecognitionResponseMetadata(total_billed_duration=timedelta(seconds=billed)), + ) + + class FakeTransport: + async def close(self) -> None: + return None + + class FakeSpeechClient: + transport = FakeTransport() + + def __init__(self) -> None: + self.requests: Final[list[object]] = [] + + async def streaming_recognize(self, requests=None): + return self._respond(requests) + + async def _respond(self, requests): + script = [google_response("four score", False, 0.0), google_response("Four score and seven", True, 2.0)] + async for request in requests: + self.requests.append(request) + if request.audio and script: + yield script.pop(0) + + speech_client = FakeSpeechClient() + + async def resolve_access_token() -> str: + return "token" + + provider_config = VertexChirpRealtimeConfig( + resolve_access_token=resolve_access_token, + project="proj-1", + location="us", + backend_factory=lambda target: SpeechStreamingBackend( + target, client_factory=lambda target, access_token: speech_client + ), + ) + audio = base64.b64encode(b"\x00\x01" * 800).decode() + client_ws = _ScriptedClientWebSocket( + [ + json.dumps( + { + "type": "session.update", + "session": { + "type": "transcription", + "audio": { + "input": { + "format": {"type": "audio/pcm", "rate": 16000}, + "transcription": {"model": "chirp_3", "language": "en"}, + "turn_detection": {"type": "server_vad"}, + } + }, + }, + } + ), + json.dumps({"type": "input_audio_buffer.append", "audio": audio}), + json.dumps({"type": "input_audio_buffer.append", "audio": audio}), + json.dumps({"type": "input_audio_buffer.commit"}), + ], + last_event_type="conversation.item.input_audio_transcription.completed", + ) + logging_obj = Mock() + logging_obj.litellm_trace_id = "trace_1" + logging_obj.model_call_details = {} + logging_obj.dispatch_success_handlers = AsyncMock() + logging_obj.dispatch_failure_handlers = AsyncMock() + handler = BaseLLMHTTPHandler() + + with patch.object(handler, "_open_realtime_backend_ws", AsyncMock(side_effect=AssertionError("dialed a websocket"))) as dial: + await handler.async_realtime( + model="chirp_3", + websocket=client_ws, + logging_obj=logging_obj, + provider_config=provider_config, + headers={}, + query_params={"model": "chirp_3", "intent": "transcription"}, + ) + + dial.assert_not_awaited() + events = client_ws.sent_events() + assert [event["type"] for event in events] == [ + "session.created", + "session.updated", + "input_audio_buffer.speech_started", + "conversation.item.input_audio_transcription.delta", + "conversation.item.input_audio_transcription.delta", + "input_audio_buffer.speech_stopped", + "conversation.item.input_audio_transcription.completed", + ] + assert events[0]["session"]["audio"]["input"]["transcription"] == {"model": "chirp_3"} + assert events[1]["session"]["audio"]["input"] == { + "format": {"type": "audio/pcm", "rate": 16000}, + "transcription": {"model": "chirp_3", "language": "en-US"}, + "turn_detection": {"type": "server_vad"}, + } + assert [event["delta"] for event in events[3:5]] == ["four score", " and seven"] + assert events[6]["transcript"] == "Four score and seven" + assert events[6]["usage"] == {"type": "duration", "seconds": 2.0} + assert speech_client.requests[0].streaming_config.config.model == "chirp_3" + assert [bytes(request.audio) for request in speech_client.requests[1:]] == [b"\x00\x01" * 800, b"\x00\x01" * 800] diff --git a/tests/test_litellm/llms/fireworks_ai/test_fireworks_ai_cache_pricing.py b/tests/test_litellm/llms/fireworks_ai/test_fireworks_ai_cache_pricing.py new file mode 100644 index 00000000000..c21943cfe75 --- /dev/null +++ b/tests/test_litellm/llms/fireworks_ai/test_fireworks_ai_cache_pricing.py @@ -0,0 +1,65 @@ +from copy import deepcopy + +import pytest + +from litellm.constants import FIREWORKS_AI_DEFAULT_CACHE_READ_RATE_RATIO +from litellm.llms.fireworks_ai.cache_pricing import with_default_cache_read_rate +from litellm.types.utils import ModelInfo + + +def test_explicit_cache_read_rate_and_missing_input_rate_keep_the_entry_untouched() -> None: + explicit_info: ModelInfo = {"input_cost_per_token": 2e-6, "cache_read_input_token_cost": 1e-6} + no_input_rate_info: ModelInfo = {"output_cost_per_token": 3e-6} + + assert with_default_cache_read_rate(explicit_info) is explicit_info + assert with_default_cache_read_rate(no_input_rate_info) is no_input_rate_info + + +def test_missing_cache_read_rate_is_derived_for_standard_and_off_peak_without_mutating_the_entry() -> None: + model_info: ModelInfo = { + "input_cost_per_token": 2e-6, + "off_peak_pricing": { + "hours_utc": "14:00-00:00", + "input_cost_per_token": 1e-6, + "output_cost_per_token": 3e-6, + }, + } + original: ModelInfo = deepcopy(model_info) + + derived = with_default_cache_read_rate(model_info) + + assert model_info == original + assert derived is not model_info + assert derived["cache_read_input_token_cost"] == pytest.approx(2e-6 * FIREWORKS_AI_DEFAULT_CACHE_READ_RATE_RATIO) + assert derived["off_peak_pricing"] == { + "hours_utc": "14:00-00:00", + "input_cost_per_token": 1e-6, + "output_cost_per_token": 3e-6, + "cache_read_input_token_cost": 1e-6 * FIREWORKS_AI_DEFAULT_CACHE_READ_RATE_RATIO, + } + + +def test_off_peak_window_without_its_own_input_rate_reuses_the_standard_derived_rate() -> None: + model_info: ModelInfo = { + "input_cost_per_token": 2e-6, + "off_peak_pricing": {"hours_utc": "14:00-00:00", "output_cost_per_token": 3e-6}, + } + + derived = with_default_cache_read_rate(model_info) + + assert derived["off_peak_pricing"]["cache_read_input_token_cost"] == derived["cache_read_input_token_cost"] + + +def test_string_rates_from_config_are_coerced_before_the_discount_is_applied() -> None: + model_info: ModelInfo = { + "input_cost_per_token": "2e-6", + "off_peak_pricing": { + "hours_utc": "14:00-00:00", + "input_cost_per_token": "1e-6", + }, + } + + derived = with_default_cache_read_rate(model_info) + + assert derived["cache_read_input_token_cost"] == pytest.approx(1e-6) + assert derived["off_peak_pricing"]["cache_read_input_token_cost"] == pytest.approx(5e-7) diff --git a/tests/test_litellm/llms/fireworks_ai/test_fireworks_ai_cost_calculator.py b/tests/test_litellm/llms/fireworks_ai/test_fireworks_ai_cost_calculator.py index 1bee310d9d3..52222f22a51 100644 --- a/tests/test_litellm/llms/fireworks_ai/test_fireworks_ai_cost_calculator.py +++ b/tests/test_litellm/llms/fireworks_ai/test_fireworks_ai_cost_calculator.py @@ -5,6 +5,11 @@ from typing import Final import pytest import litellm +from litellm.litellm_core_utils.llm_cost_calc.utils import ( + calculate_prompt_caching_savings, + generic_cost_per_token, + get_token_type_cost_breakdown, +) from litellm.llms.fireworks_ai.cost_calculator import cost_per_token from litellm.types.utils import ( CompletionTokensDetailsWrapper, @@ -48,11 +53,13 @@ STANDARD_CACHE_READ_COST = 1.5e-08 def _register_off_peak_model( - off_peak_pricing: OffPeakPricing, cache_read_cost: float | None = STANDARD_CACHE_READ_COST + off_peak_pricing: OffPeakPricing, + cache_read_cost: float | None = STANDARD_CACHE_READ_COST, + model: str = OFF_PEAK_MODEL, ) -> None: - litellm.model_cost = { # test-quality-ok: the save/restore conftest returns litellm.model_cost to the original object after each test, so replacing the map for this entry leaks nothing + litellm.model_cost = { # test-quality-ok: conftest restores litellm.model_cost after each test **litellm.model_cost, - f"fireworks_ai/{OFF_PEAK_MODEL}": { + f"fireworks_ai/{model}": { "litellm_provider": "fireworks_ai", "mode": "chat", "input_cost_per_token": STANDARD_INPUT_COST, @@ -103,9 +110,8 @@ def test_off_peak_rates_left_unset_keep_the_standard_rates(): assert math.isclose(completion_cost, 200 * STANDARD_OUTPUT_COST, rel_tol=1e-10) -def test_off_peak_window_bills_cached_tokens_at_the_off_peak_input_rate_without_a_cache_read_rate(): - """Most fireworks_ai price-map entries carry no cache_read_input_token_cost, so cached tokens - fall back to the input rate, and inside the window that has to be the off-peak one.""" +def test_off_peak_window_bills_cached_tokens_at_the_discounted_off_peak_input_rate_without_a_cache_read_rate(): + """Entries without a cache-read rate use Fireworks' documented 50% cached-token discount.""" _register_off_peak_model( {"hours_utc": OFF_PEAK_WINDOW, "input_cost_per_token": 1e-08, "output_cost_per_token": 2e-08}, cache_read_cost=None, @@ -114,12 +120,116 @@ def test_off_peak_window_bills_cached_tokens_at_the_off_peak_input_rate_without_ prompt_cost, completion_cost = cost_per_token(model=OFF_PEAK_MODEL, usage=usage, current_time=INSIDE_WINDOW) - assert math.isclose(prompt_cost, 1000 * 1e-08, rel_tol=1e-10) + assert math.isclose(prompt_cost, (700 * 1e-08) + (300 * 1e-08 * 0.5), rel_tol=1e-10) assert math.isclose(completion_cost, 200 * 2e-08, rel_tol=1e-10) peak_prompt_cost, _ = cost_per_token(model=OFF_PEAK_MODEL, usage=usage, current_time=OUTSIDE_WINDOW) - assert math.isclose(peak_prompt_cost, 1000 * STANDARD_INPUT_COST, rel_tol=1e-10) + assert math.isclose( + peak_prompt_cost, + (700 * STANDARD_INPUT_COST) + (300 * STANDARD_INPUT_COST * 0.5), + rel_tol=1e-10, + ) + + no_input_rate_model = "accounts/fireworks/models/off-peak-no-input-rate-test" + _register_off_peak_model( + {"hours_utc": OFF_PEAK_WINDOW, "output_cost_per_token": 2e-08}, + cache_read_cost=None, + model=no_input_rate_model, + ) + + standard_cache_prompt_cost, _ = cost_per_token(model=no_input_rate_model, usage=usage, current_time=INSIDE_WINDOW) + + assert math.isclose( + standard_cache_prompt_cost, + (700 * STANDARD_INPUT_COST) + (300 * STANDARD_INPUT_COST * 0.5), + rel_tol=1e-10, + ) + + +def test_an_entry_without_a_cache_read_rate_bills_cached_tokens_at_the_documented_default_discount(): + """Fireworks documents a default 50% cached-token discount for serverless models: + https://docs.fireworks.ai/guides/prompt-caching, accessed 2026-09-19.""" + model = "accounts/fireworks/models/default-cache-read-test" + litellm.model_cost = { # test-quality-ok: conftest restores litellm.model_cost after each test + **litellm.model_cost, + f"fireworks_ai/{model}": { + "litellm_provider": "fireworks_ai", + "mode": "chat", + "input_cost_per_token": INPUT_COST, + "output_cost_per_token": OUTPUT_COST, + }, + } + usage = _usage(prompt_tokens=1000, cached_tokens=300, completion_tokens=200) + + prompt_cost, completion_cost = cost_per_token(model=model, usage=usage) + + assert math.isclose(prompt_cost, (700 * INPUT_COST) + (300 * INPUT_COST * 0.5), rel_tol=1e-10) + assert prompt_cost < 1000 * INPUT_COST + assert math.isclose(completion_cost, 200 * OUTPUT_COST, rel_tol=1e-10) + + +def test_fireworks_cache_read_rates_match_breakdown_and_caching_savings(): + model = "accounts/fireworks/models/breakdown-cache-read-test" + litellm.model_cost = { # test-quality-ok: the save/restore conftest returns litellm.model_cost to the original object after each test, so replacing the map for this entry leaks nothing + **litellm.model_cost, + f"fireworks_ai/{model}": { + "litellm_provider": "fireworks_ai", + "mode": "chat", + "input_cost_per_token": INPUT_COST, + "output_cost_per_token": OUTPUT_COST, + }, + } + usage = _usage(prompt_tokens=1000, cached_tokens=300, completion_tokens=200) + + breakdown = get_token_type_cost_breakdown( + model=model, + custom_llm_provider="fireworks_ai", + usage=usage, + ) + prompt_cost, _ = cost_per_token(model=model, usage=usage) + savings = calculate_prompt_caching_savings( + model_info=litellm.get_model_info(model=model, custom_llm_provider="fireworks_ai"), + usage=usage, + custom_llm_provider="fireworks_ai", + ) + + assert math.isclose(breakdown.cache_read_cost, 300 * INPUT_COST * 0.5, rel_tol=1e-10) + assert math.isclose(breakdown.rates.cache_read_input_token_cost, INPUT_COST * 0.5, rel_tol=1e-10) + assert math.isclose( + (700 * breakdown.rates.input_cost_per_token) + breakdown.cache_read_cost, prompt_cost, rel_tol=1e-10 + ) + assert math.isclose(savings, 300 * INPUT_COST * 0.5, rel_tol=1e-10) + + +def test_generic_cost_per_token_applies_fireworks_cache_read_default_with_or_without_model_info(): + model = "accounts/fireworks/models/generic-cache-read-test" + litellm.model_cost = { # test-quality-ok: conftest restores litellm.model_cost after each test + **litellm.model_cost, + f"fireworks_ai/{model}": { + "litellm_provider": "fireworks_ai", + "mode": "chat", + "input_cost_per_token": INPUT_COST, + "output_cost_per_token": OUTPUT_COST, + }, + } + usage = _usage(prompt_tokens=1000, cached_tokens=300, completion_tokens=200) + expected_prompt_cost = (700 * INPUT_COST) + (300 * INPUT_COST * 0.5) + + implicit_model_info_cost, _ = generic_cost_per_token( + model=model, + usage=usage, + custom_llm_provider="fireworks_ai", + ) + explicit_model_info_cost, _ = generic_cost_per_token( + model=model, + usage=usage, + custom_llm_provider="fireworks_ai", + model_info=litellm.get_model_info(model=model, custom_llm_provider="fireworks_ai"), + ) + + assert math.isclose(implicit_model_info_cost, expected_prompt_cost, rel_tol=1e-10) + assert math.isclose(explicit_model_info_cost, expected_prompt_cost, rel_tol=1e-10) def test_off_peak_defaults_to_the_current_time(): diff --git a/tests/test_litellm/llms/vertex_ai/audio_transcription/test_vertex_ai_realtime_backend.py b/tests/test_litellm/llms/vertex_ai/audio_transcription/test_vertex_ai_realtime_backend.py new file mode 100644 index 00000000000..d5e88706e23 --- /dev/null +++ b/tests/test_litellm/llms/vertex_ai/audio_transcription/test_vertex_ai_realtime_backend.py @@ -0,0 +1,491 @@ +import asyncio +import json +from collections.abc import AsyncIterator, Callable, Sequence +from dataclasses import replace +from datetime import timedelta +from typing import Final + +import pytest +from google.cloud.speech_v2.types import ( + RecognitionResponseMetadata, + SpeechRecognitionAlternative, + StreamingRecognitionResult, + StreamingRecognizeRequest, + StreamingRecognizeResponse, +) +from websockets.exceptions import ConnectionClosedError, ConnectionClosedOK + +from litellm.llms.vertex_ai.audio_transcription.realtime_backend import REQUEST_QUEUE_SIZE, SpeechStreamingBackend +from litellm.llms.vertex_ai.audio_transcription.realtime_transformation import SpeechStreamingTarget + + +async def _static_token() -> str: + return "token" + + +TARGET: Final = SpeechStreamingTarget( + api_endpoint="us-speech.googleapis.com", + recognizer="projects/proj-1/locations/us/recognizers/_", + resolve_access_token=_static_token, +) +CONFIGURE: Final = json.dumps( + {"kind": "configure", "model": "chirp_3", "language_codes": ["en-US"], "sample_rate_hertz": 16_000} +) +FINISH_TURN: Final = json.dumps({"kind": "finish_turn"}) +DISCARD_TURN: Final = json.dumps({"kind": "discard_turn"}) +ScriptItem = StreamingRecognizeResponse | Exception | asyncio.Event + + +def _response( + transcript: str | None, + *, + is_final: bool = False, + billed: float = 0.0, + event: str = "SPEECH_EVENT_TYPE_UNSPECIFIED", +) -> StreamingRecognizeResponse: + results = ( + [] + if transcript is None + else [ + StreamingRecognitionResult( + alternatives=[SpeechRecognitionAlternative(transcript=transcript)], is_final=is_final + ) + ] + ) + return StreamingRecognizeResponse( + results=results, + speech_event_type=event, + metadata=RecognitionResponseMetadata(total_billed_duration=timedelta(seconds=billed)), + ) + + +class _FakeTransport: + def __init__(self) -> None: + self.closed = False + + async def close(self) -> None: + self.closed = True + + +class _FakeSpeechClient: + def __init__(self, *scripts: Sequence[ScriptItem]) -> None: + self.transport: Final = _FakeTransport() + self.streams: Final[list[list[StreamingRecognizeRequest]]] = [] + self._scripts: Final = [list(script) for script in scripts] + + async def streaming_recognize( + self, requests: AsyncIterator[StreamingRecognizeRequest] | None = None + ) -> AsyncIterator[StreamingRecognizeResponse]: + assert requests is not None + script: Final = self._scripts.pop(0) if self._scripts else [] + received: Final[list[StreamingRecognizeRequest]] = [] + self.streams.append(received) + return self._respond(requests, script, received) + + async def _respond( + self, + requests: AsyncIterator[StreamingRecognizeRequest], + script: list[ScriptItem], + received: list[StreamingRecognizeRequest], + ) -> AsyncIterator[StreamingRecognizeResponse]: + async for request in requests: + received.append(request) + if request.audio and script: + yield await self._next(script) + while script: + yield await self._next(script) + + @staticmethod + async def _next(script: list[ScriptItem]) -> StreamingRecognizeResponse: + item: Final = script.pop(0) + if isinstance(item, asyncio.Event): + await item.wait() + return await _FakeSpeechClient._next(script) + if isinstance(item, Exception): + raise item + return item + + +def _backend(client: _FakeSpeechClient, **kwargs: object) -> SpeechStreamingBackend: + return SpeechStreamingBackend(TARGET, client_factory=lambda target, access_token: client, **kwargs) + + +async def _recv(backend: SpeechStreamingBackend) -> dict[str, object]: + message: Final = await asyncio.wait_for(backend.recv(), timeout=2) + assert isinstance(message, str) + return json.loads(message) + + +async def _transcript(backend: SpeechStreamingBackend) -> str: + event: Final = await _recv(backend) + assert event["kind"] == "response", event + (result,) = event["results"] + return result["transcript"] + + +async def _configure(backend: SpeechStreamingBackend) -> None: + await backend.send(CONFIGURE) + assert await _recv(backend) == {"kind": "configured"} + + +async def _until(condition: Callable[[], bool]) -> None: + async def poll() -> None: + while not condition(): + await asyncio.sleep(0) + + await asyncio.wait_for(poll(), timeout=2) + + +def _audio(stream: list[StreamingRecognizeRequest]) -> list[bytes]: + return [bytes(request.audio) for request in stream[1:]] + + +@pytest.mark.asyncio +async def test_audio_streams_through_one_recognize_call_with_the_config_first(): + client = _FakeSpeechClient([_response("hello"), _response("hello world", is_final=True, billed=2.0)]) + async with _backend(client) as backend: + await _configure(backend) + await backend.send(b"\x01\x02") + await backend.send(b"\x03\x04") + await backend.send(FINISH_TURN) + first, second, finished = [await _recv(backend) for _ in range(3)] + assert first == { + "kind": "response", + "speech_event": "none", + "results": [{"transcript": "hello", "is_final": False}], + "billed_seconds": 0.0, + } + assert second["results"] == [{"transcript": "hello world", "is_final": True}] + assert second["billed_seconds"] == 2.0 + assert finished == {"kind": "turn_finished"} + (requests,) = client.streams + assert requests[0].recognizer == TARGET.recognizer + config = requests[0].streaming_config + assert config.config.model == "chirp_3" + assert list(config.config.language_codes) == ["en-US"] + assert config.config.explicit_decoding_config.sample_rate_hertz == 16_000 + assert config.config.explicit_decoding_config.audio_channel_count == 1 + assert config.config.explicit_decoding_config.encoding.name == "LINEAR16" + assert config.streaming_features.interim_results + assert config.streaming_features.enable_voice_activity_events + assert _audio(requests) == [b"\x01\x02", b"\x03\x04"] + assert client.transport.closed + + +@pytest.mark.asyncio +async def test_voice_activity_events_are_relayed(): + client = _FakeSpeechClient( + [_response(None, event="SPEECH_ACTIVITY_BEGIN"), _response(None, event="SPEECH_ACTIVITY_END")] + ) + async with _backend(client) as backend: + await _configure(backend) + await backend.send(b"\x00\x00") + await backend.send(b"\x00\x00") + begin, end = [await _recv(backend) for _ in range(2)] + assert (begin["speech_event"], begin["results"]) == ("begin", []) + assert end["speech_event"] == "end" + + +@pytest.mark.asyncio +async def test_audio_before_configure_is_rejected(): + backend = _backend(_FakeSpeechClient()) + with pytest.raises(RuntimeError, match="before the Speech-to-Text stream was configured"): + await backend.send(b"\x00\x00") + + +@pytest.mark.asyncio +async def test_stream_failure_closes_the_session_with_1011_and_the_reason(): + client = _FakeSpeechClient([PermissionError("IAM_PERMISSION_DENIED: speech.recognizers.recognize")]) + async with _backend(client) as backend: + await _configure(backend) + await backend.send(b"\x00\x00") + with pytest.raises(ConnectionClosedError) as excinfo: + await backend.recv() + assert excinfo.value.rcvd is not None + assert excinfo.value.rcvd.code == 1011 + assert "IAM_PERMISSION_DENIED" in excinfo.value.rcvd.reason + assert client.transport.closed + + +@pytest.mark.asyncio +async def test_close_reports_a_normal_closure_to_both_directions(): + client = _FakeSpeechClient([_response("hi")]) + backend = _backend(client) + await _configure(backend) + await backend.send(b"\x00\x00") + assert await _transcript(backend) == "hi" + await backend.close() + with pytest.raises(ConnectionClosedOK): + await backend.recv() + with pytest.raises(ConnectionClosedOK): + await backend.send(b"\x00\x00") + assert client.transport.closed + + +@pytest.mark.asyncio +async def test_turn_commands_without_audio_answer_immediately(): + backend = _backend(_FakeSpeechClient()) + await _configure(backend) + await backend.send(FINISH_TURN) + assert await _recv(backend) == {"kind": "turn_finished"} + await backend.send(DISCARD_TURN) + assert await _recv(backend) == {"kind": "turn_discarded", "billed_seconds": 0.0} + + +@pytest.mark.asyncio +async def test_discard_turn_cancels_the_open_stream_and_the_next_turn_starts_fresh(): + client = _FakeSpeechClient([_response("draft")], [_response("again", is_final=True)]) + async with _backend(client) as backend: + await _configure(backend) + await backend.send(b"\x01\x01") + assert await _transcript(backend) == "draft" + await backend.send(DISCARD_TURN) + assert await _recv(backend) == {"kind": "turn_discarded", "billed_seconds": 0.0} + await backend.send(b"\x02\x02") + assert await _transcript(backend) == "again" + assert [_audio(stream) for stream in client.streams] == [[b"\x01\x01"], [b"\x02\x02"]] + + +@pytest.mark.asyncio +async def test_discard_turn_drops_its_queued_results_and_keeps_google_billed_seconds(): + client = _FakeSpeechClient( + [_response("draft"), _response("leftover", is_final=True, billed=2.0)], + [_response("fresh", is_final=True, billed=1.0)], + ) + async with _backend(client) as backend: + await _configure(backend) + await backend.send(b"\x01\x01") + assert await _transcript(backend) == "draft" + await backend.send(b"\x02\x02") + await _until(lambda: len(client.streams[0]) == 3) + await backend.send(DISCARD_TURN) + assert await _recv(backend) == {"kind": "turn_discarded", "billed_seconds": 2.0} + assert backend._discarded_turns == frozenset() + await backend.send(b"\x03\x03") + fresh = await _recv(backend) + assert fresh["results"] == [{"transcript": "fresh", "is_final": True}] + assert fresh["billed_seconds"] == 3.0 + + +@pytest.mark.asyncio +async def test_discard_turn_keeps_the_queued_results_of_the_turn_finished_before_it(): + client = _FakeSpeechClient([_response("one", is_final=True, billed=2.0)], [_response("two")]) + async with _backend(client) as backend: + await _configure(backend) + await backend.send(b"\x01\x01") + await backend.send(FINISH_TURN) + await backend.send(b"\x02\x02") + await _until(lambda: len(client.streams) == 2 and len(client.streams[1]) == 2) + await backend.send(DISCARD_TURN) + assert await _transcript(backend) == "one" + assert await _recv(backend) == {"kind": "turn_finished"} + assert await _recv(backend) == {"kind": "turn_discarded", "billed_seconds": 2.0} + + +@pytest.mark.asyncio +async def test_billed_seconds_accumulate_across_turns(): + client = _FakeSpeechClient( + [_response("one", is_final=True, billed=2.0)], [_response("two", is_final=True, billed=3.0)] + ) + async with _backend(client) as backend: + await _configure(backend) + await backend.send(b"\x00\x00") + await backend.send(FINISH_TURN) + first = await _recv(backend) + assert await _recv(backend) == {"kind": "turn_finished"} + await backend.send(b"\x00\x00") + await backend.send(FINISH_TURN) + second = await _recv(backend) + assert await _recv(backend) == {"kind": "turn_finished"} + assert (first["billed_seconds"], second["billed_seconds"]) == (2.0, 5.0) + assert len(client.streams) == 2 + + +@pytest.mark.asyncio +async def test_streams_rotate_before_the_five_minute_limit_without_ending_the_turn(): + now = [0.0] + client = _FakeSpeechClient( + [_response("first"), _response("first half", is_final=True, billed=239.0)], + [_response("second", billed=1.0)], + ) + async with _backend(client, clock=lambda: now[0], rotation_seconds=240.0) as backend: + await _configure(backend) + await backend.send(b"\x01\x01") + assert await _transcript(backend) == "first" + now[0] = 239.0 + await backend.send(b"\x02\x02") + assert await _transcript(backend) == "first half" + now[0] = 240.0 + await backend.send(b"\x03\x03") + second = await _recv(backend) + assert second["results"] == [{"transcript": "second", "is_final": False}] + assert second["billed_seconds"] == 240.0 + await backend.send(FINISH_TURN) + assert await _recv(backend) == {"kind": "turn_finished"} + assert [_audio(stream) for stream in client.streams] == [[b"\x01\x01", b"\x02\x02"], [b"\x03\x03"]] + assert client.streams[1][0].streaming_config.config.model == "chirp_3" + + +@pytest.mark.asyncio +async def test_turn_finished_follows_results_that_arrive_after_a_rotation(): + now = [0.0] + client = _FakeSpeechClient( + [_response("one"), _response("one two", is_final=True)], + [_response("three")], + ) + async with _backend(client, clock=lambda: now[0], rotation_seconds=240.0) as backend: + await _configure(backend) + await backend.send(b"\x01\x01") + assert await _transcript(backend) == "one" + now[0] = 240.0 + await backend.send(b"\x02\x02") + await backend.send(FINISH_TURN) + assert await _transcript(backend) == "one two" + assert await _transcript(backend) == "three" + assert await _recv(backend) == {"kind": "turn_finished"} + + +@pytest.mark.asyncio +async def test_rotation_waits_for_a_pause_in_speech(): + now = [0.0] + client = _FakeSpeechClient( + [ + _response(None, event="SPEECH_ACTIVITY_BEGIN"), + _response("still talking"), + _response("still talking", is_final=True, event="SPEECH_ACTIVITY_END"), + ], + [_response("next")], + ) + async with _backend( + client, clock=lambda: now[0], rotation_seconds=240.0, rotation_deadline_seconds=280.0 + ) as backend: + await _configure(backend) + await backend.send(b"\x01\x01") + assert (await _recv(backend))["speech_event"] == "begin" + now[0] = 250.0 + await backend.send(b"\x02\x02") + assert await _transcript(backend) == "still talking" + now[0] = 260.0 + await backend.send(b"\x03\x03") + assert (await _recv(backend))["speech_event"] == "end" + now[0] = 261.0 + await backend.send(b"\x04\x04") + assert await _transcript(backend) == "next" + assert [_audio(stream) for stream in client.streams] == [[b"\x01\x01", b"\x02\x02", b"\x03\x03"], [b"\x04\x04"]] + + +@pytest.mark.asyncio +async def test_rotation_is_forced_at_the_deadline_during_continuous_speech(): + now = [0.0] + client = _FakeSpeechClient( + [_response(None, event="SPEECH_ACTIVITY_BEGIN"), _response("still talking")], + [_response("cut off")], + ) + async with _backend( + client, clock=lambda: now[0], rotation_seconds=240.0, rotation_deadline_seconds=280.0 + ) as backend: + await _configure(backend) + await backend.send(b"\x01\x01") + assert (await _recv(backend))["speech_event"] == "begin" + now[0] = 279.0 + await backend.send(b"\x02\x02") + assert await _transcript(backend) == "still talking" + now[0] = 280.0 + await backend.send(b"\x03\x03") + assert await _transcript(backend) == "cut off" + assert [_audio(stream) for stream in client.streams] == [[b"\x01\x01", b"\x02\x02"], [b"\x03\x03"]] + + +@pytest.mark.asyncio +async def test_every_stream_opens_its_own_client_with_a_freshly_resolved_token(): + now = [0.0] + tokens = iter(("token-1", "token-2")) + seen_tokens: list[str] = [] + clients = [_FakeSpeechClient([_response("first")]), _FakeSpeechClient([_response("second")])] + unopened = iter(clients) + + async def resolve_access_token() -> str: + return next(tokens) + + def open_client(target: SpeechStreamingTarget, access_token: str) -> _FakeSpeechClient: + seen_tokens.append(access_token) + return next(unopened) + + backend = SpeechStreamingBackend( + replace(TARGET, resolve_access_token=resolve_access_token), + client_factory=open_client, + clock=lambda: now[0], + rotation_seconds=240.0, + ) + async with backend: + await _configure(backend) + await backend.send(b"\x01\x01") + assert await _transcript(backend) == "first" + now[0] = 240.0 + await backend.send(b"\x02\x02") + assert await _transcript(backend) == "second" + assert clients[0].transport.closed + assert not clients[1].transport.closed + assert seen_tokens == ["token-1", "token-2"] + assert [len(client.streams) for client in clients] == [1, 1] + assert clients[1].transport.closed + + +@pytest.mark.asyncio +async def test_close_releases_a_rotated_stream_that_never_started_relaying(): + now = [0.0] + hold = asyncio.Event() + clients = [_FakeSpeechClient([_response("first"), hold]), _FakeSpeechClient([_response("never")])] + unopened = iter(clients) + backend = SpeechStreamingBackend( + TARGET, + client_factory=lambda target, access_token: next(unopened), + clock=lambda: now[0], + rotation_seconds=240.0, + ) + await _configure(backend) + await backend.send(b"\x01\x01") + assert await _transcript(backend) == "first" + now[0] = 240.0 + await backend.send(b"\x02\x02") + await asyncio.sleep(0) + assert clients[1].streams == [] + await backend.close() + assert [client.transport.closed for client in clients] == [True, True] + + +@pytest.mark.asyncio +async def test_discard_turn_cancels_every_stream_of_the_turn(): + now = [0.0] + hold = asyncio.Event() + client = _FakeSpeechClient( + [_response("draft"), hold, _response("never delivered")], + [_response("fresh", is_final=True)], + ) + async with _backend(client, clock=lambda: now[0], rotation_seconds=240.0) as backend: + await _configure(backend) + await backend.send(b"\x01\x01") + assert await _transcript(backend) == "draft" + now[0] = 240.0 + await backend.send(b"\x02\x02") + await backend.send(DISCARD_TURN) + assert await _recv(backend) == {"kind": "turn_discarded", "billed_seconds": 0.0} + await backend.send(b"\x03\x03") + assert await _transcript(backend) == "fresh" + assert [_audio(stream) for stream in client.streams] == [[b"\x01\x01"], [b"\x03\x03"]] + + +@pytest.mark.asyncio +async def test_audio_sends_block_once_the_request_queue_is_full(): + hold = asyncio.Event() + client = _FakeSpeechClient([hold, _response("late", is_final=True)]) + async with _backend(client) as backend: + await _configure(backend) + for _ in range(REQUEST_QUEUE_SIZE + 1): + await backend.send(b"\x00\x00") + blocked = asyncio.create_task(backend.send(b"\x00\x00")) + await asyncio.sleep(0) + assert not blocked.done() + hold.set() + await asyncio.wait_for(blocked, timeout=2) + assert await _transcript(backend) == "late" diff --git a/tests/test_litellm/llms/vertex_ai/audio_transcription/test_vertex_ai_realtime_transformation.py b/tests/test_litellm/llms/vertex_ai/audio_transcription/test_vertex_ai_realtime_transformation.py new file mode 100644 index 00000000000..84c3a4e244a --- /dev/null +++ b/tests/test_litellm/llms/vertex_ai/audio_transcription/test_vertex_ai_realtime_transformation.py @@ -0,0 +1,424 @@ +import base64 +import json +from typing import Final +from unittest.mock import MagicMock + +import pytest + +from litellm.llms.base_llm.realtime.transformation import RealtimeBackend +from litellm.llms.vertex_ai.audio_transcription.realtime_transformation import ( + MAX_AUDIO_MESSAGE_BYTES, + ChirpProtocolError, + ChirpSessionConfig, + SpeechStreamingTarget, + VertexChirpRealtimeConfig, + is_vertex_speech_to_text_model, + new_words, + parse_chirp_session_update, +) +from litellm.llms.vertex_ai.common_utils import VertexAIError +from litellm.types.llms.vertex_ai_speech_to_text import ( + VertexSpeechStreamingConfigured, + VertexSpeechStreamingResponse, + VertexSpeechStreamingResult, + VertexSpeechStreamingTurnDiscarded, + VertexSpeechStreamingTurnFinished, +) +from litellm.types.realtime import RealtimeResponseTransformInput + +MODEL: Final = "chirp_3" +EMPTY_TRANSFORM_INPUT: Final[RealtimeResponseTransformInput] = { + "session_configuration_request": None, + "current_output_item_id": None, + "current_response_id": None, + "current_delta_chunks": None, + "current_item_chunks": None, + "current_conversation_id": None, + "current_delta_type": None, +} +DELTA: Final = "conversation.item.input_audio_transcription.delta" +COMPLETED: Final = "conversation.item.input_audio_transcription.completed" + + +def _event(event_type: str, **fields: object) -> str: + return json.dumps({"type": event_type, **fields}) + + +def _ga_session_update( + rate: int = 24_000, turn_detection: str | None = "server_vad", language: str | None = "en", model: str = MODEL +) -> str: + transcription = {"model": model} if language is None else {"model": model, "language": language} + return _event( + "session.update", + session={ + "type": "transcription", + "audio": { + "input": { + "format": {"type": "audio/pcm", "rate": rate}, + "turn_detection": None if turn_detection is None else {"type": turn_detection}, + "transcription": transcription, + } + }, + }, + ) + + +async def _token() -> str: + return "token" + + +def _config(location: str | None = "us") -> VertexChirpRealtimeConfig: + return VertexChirpRealtimeConfig(resolve_access_token=_token, project="proj-1", location=location) + + +def _configured( + rate: int = 24_000, turn_detection: str | None = "server_vad", language: str | None = "en" +) -> VertexChirpRealtimeConfig: + config = _config() + config.transform_session_created_event(MODEL, "sess_1") + config.transform_realtime_request(_ga_session_update(rate, turn_detection, language), MODEL) + return config + + +def _backend_events(config: VertexChirpRealtimeConfig, frame: object) -> list[dict[str, object]]: + assert hasattr(frame, "model_dump_json") + response = config.transform_realtime_response(frame.model_dump_json(), MODEL, MagicMock(), EMPTY_TRANSFORM_INPUT)[ + "response" + ] + assert isinstance(response, list) + return response + + +def _response( + *results: tuple[str, bool], speech_event: str = "none", billed_seconds: float = 0.0 +) -> VertexSpeechStreamingResponse: + return VertexSpeechStreamingResponse( + speech_event=speech_event, + results=tuple(VertexSpeechStreamingResult(transcript=text, is_final=final) for text, final in results), + billed_seconds=billed_seconds, + ) + + +def _types(events: list[dict[str, object]]) -> list[object]: + return [event["type"] for event in events] + + +def _commands(config: VertexChirpRealtimeConfig, payload: str) -> list[object]: + return [ + json.loads(command) if isinstance(command, str) else command + for command in config.transform_realtime_request(payload, MODEL) + ] + + +@pytest.mark.parametrize( + ("model", "expected"), + [ + ("vertex_ai/chirp_3", True), + ("chirp_3", True), + ("chirp_2", False), + ("gemini-live-2.5-flash", False), + ("vertex_ai/gemini-2.0-flash-live-preview-04-09", False), + ("vertex_ai/gemini-3.5-transcribe-live-preview", False), + ("gemini-3.5-transcribe-preview", False), + ], +) +def test_is_vertex_speech_to_text_model(model: str, expected: bool): + assert is_vertex_speech_to_text_model(model) is expected + + +def test_ga_session_update_maps_to_a_speech_config(): + config = parse_chirp_session_update(_ga_session_update(16_000, "server_vad", "pt"), "vertex_ai/chirp_3") + assert config == ChirpSessionConfig(model=MODEL, language="pt-BR", sample_rate=16_000, server_vad=True) + assert json.loads(config.configure_command()) == { + "kind": "configure", + "model": MODEL, + "language_codes": ["pt-BR"], + "sample_rate_hertz": 16_000, + } + + +def test_beta_session_update_defaults_the_rate_and_auto_detects_the_language(): + config = parse_chirp_session_update( + _event( + "transcription_session.update", + session={ + "input_audio_format": "pcm16", + "input_audio_transcription": {"model": MODEL}, + "turn_detection": None, + }, + ), + MODEL, + ) + assert config == ChirpSessionConfig(model=MODEL, language=None, sample_rate=24_000, server_vad=False) + assert json.loads(config.configure_command())["language_codes"] == ["auto"] + + +@pytest.mark.parametrize( + ("payload", "message"), + [ + (_event("session.update", session={"type": "realtime_voice"}), "transcription sessions only"), + (_ga_session_update(model="gemini-live-2.5-flash"), "cannot be changed"), + (_event("session.update", session={"audio": {"input": {"format": {"type": "audio/pcmu"}}}}), "pcm16"), + ( + _event("session.update", session={"audio": {"input": {"format": {"type": "audio/pcm", "channels": 2}}}}), + "mono", + ), + (_ga_session_update(rate=4_000), "sample rates"), + (_ga_session_update(rate=96_000), "sample rates"), + (_ga_session_update(turn_detection="semantic_vad"), "server_vad"), + ], +) +def test_unsupported_session_settings_are_rejected(payload: str, message: str): + with pytest.raises(ChirpProtocolError, match=message): + parse_chirp_session_update(payload, MODEL) + + +def test_session_update_configures_once_and_later_updates_are_ignored(): + config = _config() + config.transform_session_created_event(MODEL, "sess_1") + first = _commands(config, _ga_session_update(16_000)) + assert first == [{"kind": "configure", "model": MODEL, "language_codes": ["en-US"], "sample_rate_hertz": 16_000}] + assert config.is_setup_message(first[0]) + assert _commands(config, _ga_session_update(8_000)) == [] + + +def test_audio_and_commits_before_session_update_are_rejected(): + config = _config() + with pytest.raises(ChirpProtocolError, match=r"session\.update must configure"): + config.transform_realtime_request( + _event("input_audio_buffer.append", audio=base64.b64encode(b"\x00\x00").decode()), MODEL + ) + with pytest.raises(ChirpProtocolError, match=r"session\.update must configure"): + config.transform_realtime_request(_event("input_audio_buffer.commit"), MODEL) + + +def test_append_is_split_into_google_sized_chunks(): + config = _configured() + audio = bytes(range(256)) * 250 + chunks = config.transform_realtime_request( + _event("input_audio_buffer.append", audio=base64.b64encode(audio).decode()), MODEL + ) + assert [len(chunk) for chunk in chunks] == [ + MAX_AUDIO_MESSAGE_BYTES, + MAX_AUDIO_MESSAGE_BYTES, + 64_000 - 2 * MAX_AUDIO_MESSAGE_BYTES, + ] + assert b"".join(chunk for chunk in chunks if isinstance(chunk, bytes)) == audio + + +def test_commit_end_and_clear_map_to_turn_commands(): + config = _configured() + assert _commands(config, _event("input_audio_buffer.commit")) == [{"kind": "finish_turn"}] + assert _commands(config, _event("input_audio_buffer.end")) == [{"kind": "finish_turn"}] + assert _commands(config, _event("input_audio_buffer.clear")) == [{"kind": "discard_turn"}] + + +def test_unsupported_client_events_are_dropped(): + assert _commands(_configured(), _event("response.create")) == [] + + +def test_connect_announces_a_session_with_chirp_defaults(): + event = _config().transform_session_created_event(MODEL, "sess_1") + assert event["type"] == "session.created" + assert event["session"]["id"] == "sess_1" + assert event["session"]["audio"]["input"] == { + "format": {"type": "audio/pcm", "rate": 24_000}, + "transcription": {"model": MODEL}, + "turn_detection": {"type": "server_vad"}, + } + + +def test_configured_backend_reports_the_negotiated_session(): + config = _configured(rate=16_000, turn_detection=None, language="pt-BR") + events = _backend_events(config, VertexSpeechStreamingConfigured()) + assert _types(events) == ["session.created"] + session = events[0]["session"] + assert isinstance(session, dict) + assert session["id"] == "sess_1" + assert session["audio"]["input"] == { + "format": {"type": "audio/pcm", "rate": 16_000}, + "transcription": {"model": MODEL, "language": "pt-BR"}, + "turn_detection": None, + } + + +def test_backend_frames_before_session_update_are_an_error(): + config = _config() + config.transform_session_created_event(MODEL, "sess_1") + with pytest.raises(ChirpProtocolError, match=r"session\.update must configure"): + _backend_events(config, VertexSpeechStreamingConfigured()) + + +def test_server_vad_turn_streams_new_words_then_completes_with_usage(): + config = _configured() + assert _types(_backend_events(config, _response(speech_event="begin"))) == ["input_audio_buffer.speech_started"] + first = _backend_events(config, _response(("four score", False))) + assert [(event["type"], event["delta"]) for event in first] == [(DELTA, "four score")] + second = _backend_events(config, _response(("four score and seven", False))) + assert [event["delta"] for event in second] == [" and seven"] + final = _backend_events(config, _response(("Four score and seven years ago.", True), billed_seconds=3.5)) + assert _types(final) == [DELTA, "input_audio_buffer.speech_stopped", COMPLETED] + assert final[0]["delta"] == " years ago." + assert final[2]["transcript"] == "Four score and seven years ago." + assert final[2]["usage"] == {"type": "duration", "seconds": 3.5} + assert {event["item_id"] for event in (*first, *second, *final)} == {first[0]["item_id"]} + assert _backend_events(config, _response(speech_event="end")) == [] + + +def test_server_vad_final_result_completes_before_the_interim_that_follows_it(): + config = _configured() + _backend_events(config, _response(speech_event="begin")) + events = _backend_events(config, _response(("four score", True), ("and seven", False))) + assert _types(events) == [ + DELTA, + "input_audio_buffer.speech_stopped", + COMPLETED, + "input_audio_buffer.speech_started", + DELTA, + ] + assert events[2]["transcript"] == "four score" + assert events[4]["delta"] == "and seven" + assert events[4]["item_id"] != events[2]["item_id"] + assert events[4]["item_id"] == events[3]["item_id"] + finished = _backend_events(config, _response(("and seven years", True))) + assert [(event["type"], event.get("delta", event.get("transcript"))) for event in finished] == [ + (DELTA, " years"), + ("input_audio_buffer.speech_stopped", None), + (COMPLETED, "and seven years"), + ] + assert {event["item_id"] for event in finished} == {events[4]["item_id"]} + + +def test_manual_turn_keeps_the_interim_that_follows_a_final_in_the_same_frame(): + config = _configured(turn_detection=None) + first = _backend_events(config, _response(("four score", True), ("and seven", False))) + assert [(event["type"], event["delta"]) for event in first] == [(DELTA, "four score"), (DELTA, " and seven")] + second = _backend_events(config, _response(("and seven years", True))) + assert [event["delta"] for event in second] == [" years"] + completed = _backend_events(config, VertexSpeechStreamingTurnFinished()) + assert [(event["type"], event["transcript"]) for event in completed] == [(COMPLETED, "four score and seven years")] + assert {event["item_id"] for event in (*first, *second, *completed)} == {first[0]["item_id"]} + + +def test_manual_turns_complete_on_commit_without_speech_events(): + config = _configured(turn_detection=None) + assert _backend_events(config, _response(speech_event="begin")) == [] + first = _backend_events(config, _response(("hello there", True), billed_seconds=1.25)) + assert [(event["type"], event["delta"]) for event in first] == [(DELTA, "hello there")] + second = _backend_events(config, _response(("world", True))) + assert [event["delta"] for event in second] == [" world"] + completed = _backend_events(config, VertexSpeechStreamingTurnFinished()) + assert _types(completed) == [COMPLETED] + assert completed[0]["transcript"] == "hello there world" + assert completed[0]["usage"] == {"type": "duration", "seconds": 1.25} + assert _backend_events(config, VertexSpeechStreamingTurnFinished()) == [] + + +def test_clear_discards_the_open_turn(): + config = _configured(turn_detection=None) + draft = _backend_events(config, _response(("draft", False))) + assert _backend_events(config, VertexSpeechStreamingTurnDiscarded(billed_seconds=0.0)) == [] + assert _backend_events(config, VertexSpeechStreamingTurnFinished()) == [] + fresh = _backend_events(config, _response(("again", False))) + assert fresh[0]["delta"] == "again" + assert fresh[0]["item_id"] != draft[0]["item_id"] + + +def test_cleared_audio_keeps_google_billed_seconds_for_the_close_flush(): + config = _configured(turn_detection=None) + assert _backend_events(config, _response(("draft", False), billed_seconds=1.0)) != [] + assert _backend_events(config, VertexSpeechStreamingTurnDiscarded(billed_seconds=2.5)) == [] + assert config.unbilled_usage_on_session_close(MODEL) == {"type": "duration", "seconds": 2.5} + + +def test_usage_is_billed_once_across_turns_and_flushed_on_close(): + config = _configured() + first = _backend_events(config, _response(("one", True), billed_seconds=2.0)) + second = _backend_events(config, _response(("two", True), billed_seconds=5.0)) + assert first[-1]["usage"] == {"type": "duration", "seconds": 2.0} + assert second[-1]["usage"] == {"type": "duration", "seconds": 3.0} + assert config.unbilled_usage_on_session_close(MODEL) is None + assert _backend_events(config, _response(billed_seconds=6.5)) == [] + assert config.unbilled_usage_on_session_close(MODEL) == {"type": "duration", "seconds": 1.5} + assert config.unbilled_usage_on_session_close(MODEL) is None + + +class _NullBackend: + async def __aenter__(self) -> "_NullBackend": + return self + + async def __aexit__(self, exc_type: object, exc_value: object, traceback: object) -> None: + return None + + async def send(self, message: str | bytes) -> None: + return None + + async def recv(self, decode: bool | None = None) -> str | bytes: + return "" + + async def close(self) -> None: + return None + + +@pytest.mark.asyncio +async def test_open_backend_targets_the_regional_speech_endpoint(): + targets: list[SpeechStreamingTarget] = [] + + def factory(target: SpeechStreamingTarget) -> RealtimeBackend: + targets.append(target) + return _NullBackend() + + config = VertexChirpRealtimeConfig( + resolve_access_token=_token, project="proj-1", location=None, backend_factory=factory + ) + url = config.get_complete_url(None, "vertex_ai/chirp_3") + assert url == "us-speech.googleapis.com" + assert config.validate_environment({}, MODEL, "https://" + url) == {} + backend = await config.open_backend(url, {}) + assert isinstance(backend, _NullBackend) + assert targets == [ + SpeechStreamingTarget( + api_endpoint="us-speech.googleapis.com", + recognizer="projects/proj-1/locations/us/recognizers/_", + resolve_access_token=_token, + ) + ] + assert await targets[0].resolve_access_token() == "token" + + +@pytest.mark.parametrize( + ("location", "api_base", "endpoint"), + [ + ("global", None, "speech.googleapis.com"), + ("europe-west4", None, "europe-west4-speech.googleapis.com"), + ("us", "https://speech-proxy.internal:8443/v2", "speech-proxy.internal:8443"), + ], +) +def test_get_complete_url_honors_location_and_api_base(location: str, api_base: str | None, endpoint: str): + assert _config(location).get_complete_url(api_base, MODEL) == endpoint + + +def test_get_complete_url_rejects_non_speech_models(): + with pytest.raises(ValueError, match="Unsupported Speech-to-Text streaming model"): + _config().get_complete_url(None, "gemini-live-2.5-flash") + + +@pytest.mark.parametrize("location", ["bad loc", "../us"]) +def test_invalid_locations_are_rejected_up_front(location: str): + with pytest.raises(VertexAIError): + _config(location) + + +@pytest.mark.parametrize( + ("previous", "current", "delta"), + [ + ("", "hello", "hello"), + ("hello", "hello world", " world"), + ("hello", "Hello, world", " world"), + ("hello world", "hello world", ""), + ("hello there", "hello world", " world"), + ("hello world", "hello", ""), + ], +) +def test_new_words(previous: str, current: str, delta: str): + assert new_words(previous, current) == delta diff --git a/tests/test_litellm/llms/xai/test_xai_audio_transcription_transformation.py b/tests/test_litellm/llms/xai/test_xai_audio_transcription_transformation.py new file mode 100644 index 00000000000..e2e3fc3d4dc --- /dev/null +++ b/tests/test_litellm/llms/xai/test_xai_audio_transcription_transformation.py @@ -0,0 +1,188 @@ +import httpx +import pytest + +import litellm +from litellm.llms.base_llm.audio_transcription.transformation import ( + AudioTranscriptionRequestData, +) +from litellm.llms.custom_httpx.http_handler import HTTPHandler +from litellm.llms.xai.audio_transcription.transformation import ( + XAIAudioTranscriptionConfig, + XAIAudioTranscriptionError, +) +from litellm.types.utils import LlmProviders +from litellm.utils import ProviderConfigManager + +CONFIG = XAIAudioTranscriptionConfig() + +WAV_BYTES = b"RIFF" + b"\x00" * 64 + + +def test_transform_request_serializes_provider_params(): + result = CONFIG.transform_audio_transcription_request( + model="grok-voice-transcribe-2.0", + audio_file=WAV_BYTES, + optional_params={ + "language": "en", + "diarize": True, + "keyterm": ["LiteLLM", "Grok"], + }, + litellm_params={}, + ) + + assert isinstance(result, AudioTranscriptionRequestData) + data = result.data + assert data["model"] == "grok-voice-transcribe-2.0" + assert data["language"] == "en" + assert data["diarize"] == "true" + assert data["keyterm"] == ["LiteLLM", "Grok"] + filename, content, content_type = result.files["file"] + assert content == WAV_BYTES + assert isinstance(filename, str) + assert isinstance(content_type, str) + + +def test_transform_request_flattens_extra_body(): + result = CONFIG.transform_audio_transcription_request( + model="grok-voice-transcribe-1.0", + audio_file=WAV_BYTES, + optional_params={ + "language": "en", + "extra_body": {"diarize": False, "channels": 2}, + }, + litellm_params={}, + ) + assert result.data["diarize"] == "false" + assert result.data["channels"] == "2" + assert "extra_body" not in result.data + + +@pytest.mark.parametrize( + "api_base,expected", + [ + (None, "https://api.x.ai/v1/stt"), + ("https://api.x.ai/v1", "https://api.x.ai/v1/stt"), + ("https://api.x.ai/v1/", "https://api.x.ai/v1/stt"), + ("https://proxy.example/", "https://proxy.example/v1/stt"), + ], +) +def test_get_complete_url(api_base, expected): + url = CONFIG.get_complete_url( + api_base=api_base, + api_key=None, + model="grok-voice-transcribe-2.0", + optional_params={}, + litellm_params={}, + ) + assert url == expected + + +def test_validate_environment_sets_bearer_header(): + headers = CONFIG.validate_environment( + headers={}, + model="grok-voice-transcribe-2.0", + messages=[], + optional_params={}, + litellm_params={}, + api_key="sk-test", + ) + assert headers["Authorization"] == "Bearer sk-test" + assert "Content-Type" not in headers + + +def test_validate_environment_requires_key(monkeypatch): + monkeypatch.delenv("XAI_API_KEY", raising=False) + monkeypatch.setattr(litellm, "xai_key", None) + with pytest.raises(ValueError, match="xAI API key is required"): + CONFIG.validate_environment( + headers={}, + model="grok-voice-transcribe-2.0", + messages=[], + optional_params={}, + litellm_params={}, + api_key=None, + ) + + +def test_transform_response_maps_xai_shape(): + raw = httpx.Response( + 200, + json={ + "text": "hello world", + "language": "en", + "duration": 3.2, + "words": [ + {"text": "hello", "start": 0.0, "end": 0.5, "speaker": 1}, + {"text": "world", "start": 0.5, "end": 1.0}, + ], + }, + request=httpx.Request("POST", "https://api.x.ai/v1/stt"), + ) + response = CONFIG.transform_audio_transcription_response(raw_response=raw) + + assert response.text == "hello world" + assert response["language"] == "en" + assert response["duration"] == 3.2 + assert response["task"] == "transcribe" + assert response["words"] == [ + {"word": "hello", "start": 0.0, "end": 0.5, "speaker": 1}, + {"word": "world", "start": 0.5, "end": 1.0}, + ] + assert response._hidden_params["audio_transcription_duration"] == 3.2 + + +def test_transform_response_raises_on_error_status(): + raw = httpx.Response( + 400, + json={ + "code": "Client specified an invalid argument", + "error": "Incorrect API key provided", + }, + request=httpx.Request("POST", "https://api.x.ai/v1/stt"), + ) + with pytest.raises(XAIAudioTranscriptionError) as exc: + CONFIG.transform_audio_transcription_response(raw_response=raw) + assert exc.value.status_code == 400 + assert "Incorrect API key provided" in exc.value.message + + +def test_transcription_routes_to_xai_stt(monkeypatch): + monkeypatch.delenv("XAI_API_KEY", raising=False) + monkeypatch.setattr(litellm, "xai_key", None) + captured: dict = {} + + def handler(request: httpx.Request) -> httpx.Response: + captured["request"] = request + return httpx.Response( + 200, + json={"text": "transcribed text", "language": "en", "duration": 1.5}, + request=request, + ) + + http_handler = HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(handler))) + response = litellm.transcription( + model="xai/grok-voice-transcribe-2.0", + file=("sample.wav", WAV_BYTES, "audio/wav"), + api_key="sk-test", + diarize=True, + keyterm=["LiteLLM"], + client=http_handler, + ) + + request = captured["request"] + assert str(request.url) == "https://api.x.ai/v1/stt" + assert request.headers["Authorization"] == "Bearer sk-test" + body = request.content.decode("utf-8", errors="replace") + assert 'name="model"' in body and "grok-voice-transcribe-2.0" in body + assert 'name="diarize"' in body and "true" in body + assert 'name="keyterm"' in body and "LiteLLM" in body + assert 'name="file"' in body + assert response.text == "transcribed text" + + +def test_provider_config_manager_returns_xai_config(): + config = ProviderConfigManager.get_provider_audio_transcription_config( + model="grok-voice-transcribe-2.0", + provider=LlmProviders.XAI, + ) + assert isinstance(config, XAIAudioTranscriptionConfig) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_client_allowlist.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_client_allowlist.py new file mode 100644 index 00000000000..d1f852e504e --- /dev/null +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_client_allowlist.py @@ -0,0 +1,217 @@ +from collections.abc import Mapping +from typing import Final + +import pytest + +from litellm.proxy._experimental.mcp_server.client_allowlist import ( + MCP_ALLOWED_CLIENTS_SETTING, + MCP_CLIENT_ID_HEADER_SETTING, + MCP_CLIENT_ID_JWT_FIELD_SETTING, + MCPClientAllowlist, + MCPClientIdentity, + MCPClientRejection, + check_mcp_client_allowed, + load_mcp_client_allowlist, + parse_allowed_mcp_clients, + resolve_mcp_client_identity, +) + +ANTIGRAVITY: Final = {"alias": "Antigravity CLI", "value": "antigravity-cli"} +CODEX: Final = {"alias": "Codex", "value": "codex-mcp-client"} +ANTIGRAVITY_ONLY: Final[Mapping[str, str]] = {"antigravity-cli": "Antigravity CLI"} +JWT_ONLY: Final = MCPClientAllowlist(aliases_by_value=ANTIGRAVITY_ONLY, jwt_field="azp", header=None) +HEADER_ONLY: Final = MCPClientAllowlist(aliases_by_value=ANTIGRAVITY_ONLY, jwt_field=None, header="x-mcp-client") +JWT_AND_HEADER: Final = MCPClientAllowlist(aliases_by_value=ANTIGRAVITY_ONLY, jwt_field="azp", header="x-mcp-client") +NO_SOURCE: Final = MCPClientAllowlist(aliases_by_value=ANTIGRAVITY_ONLY, jwt_field=None, header=None) +NO_HEADERS: Final[Mapping[str, str]] = {} + + +_ALLOWLIST_SETTING_CASES: Final[tuple[tuple[object, Mapping[str, str] | None], ...]] = ( + (None, None), + ([], {}), + ([ANTIGRAVITY], ANTIGRAVITY_ONLY), + ([ANTIGRAVITY, CODEX], {"antigravity-cli": "Antigravity CLI", "codex-mcp-client": "Codex"}), + ( + [ANTIGRAVITY, {"alias": "Antigravity (prod)", "value": "antigravity-cli"}], + {"antigravity-cli": "Antigravity (prod)"}, + ), + ( + [ANTIGRAVITY, {"alias": "Antigravity CLI", "value": "antigravity-prod"}], + {**ANTIGRAVITY_ONLY, "antigravity-prod": "Antigravity CLI"}, + ), + (["antigravity-cli"], {}), + ("antigravity-cli", {}), + ([ANTIGRAVITY, 1], {}), + ([{"alias": "Antigravity CLI"}], {}), + ([{"value": "antigravity-cli"}], {}), + ([{"alias": "", "value": "antigravity-cli"}], {}), + ([{"alias": "Antigravity CLI", "value": ""}], {}), + ([{"alias": "Antigravity CLI", "value": ["antigravity-cli"]}], {}), + (ANTIGRAVITY, {}), +) + + +@pytest.mark.parametrize(("raw_setting", "expected"), _ALLOWLIST_SETTING_CASES) +def test_parse_allowed_mcp_clients(raw_setting: object, expected: Mapping[str, str] | None) -> None: + assert parse_allowed_mcp_clients(raw_setting) == expected + + +def test_load_returns_none_when_the_allowlist_setting_is_absent_even_if_identity_sources_are_set() -> None: + settings: Final = {"litellm_jwtauth": {"mcp_client_id_jwt_field": "azp"}, "mcp_client_id_header": "x-mcp-client"} + assert load_mcp_client_allowlist(settings) is None + + +def test_load_reads_the_jwt_field_from_litellm_jwtauth_and_lowercases_the_header_name() -> None: + settings: Final = { + "mcp_allowed_clients": [ANTIGRAVITY, CODEX], + "litellm_jwtauth": {"user_id_jwt_field": "sub", "mcp_client_id_jwt_field": "resource_access.mcp.client"}, + "mcp_client_id_header": "X-MCP-Client", + } + assert load_mcp_client_allowlist(settings) == MCPClientAllowlist( + aliases_by_value={"antigravity-cli": "Antigravity CLI", "codex-mcp-client": "Codex"}, + jwt_field="resource_access.mcp.client", + header="x-mcp-client", + ) + + +@pytest.mark.parametrize( + "settings", + ( + {"mcp_allowed_clients": [ANTIGRAVITY]}, + {"mcp_allowed_clients": [ANTIGRAVITY], "litellm_jwtauth": {}, "mcp_client_id_header": ""}, + {"mcp_allowed_clients": [ANTIGRAVITY], "litellm_jwtauth": {"mcp_client_id_jwt_field": ""}}, + {"mcp_allowed_clients": [ANTIGRAVITY], "litellm_jwtauth": "azp", "mcp_client_id_header": ["x"]}, + ), +) +def test_load_without_a_usable_identity_source_keeps_the_allowlist_but_no_source( + settings: Mapping[str, object], +) -> None: + assert load_mcp_client_allowlist(settings) == NO_SOURCE + + +@pytest.mark.parametrize("raw_setting", ("antigravity-cli", ["antigravity-cli"], [{"alias": "Antigravity CLI"}])) +def test_load_malformed_allowlist_admits_nobody(raw_setting: object) -> None: + loaded: Final = load_mcp_client_allowlist({"mcp_allowed_clients": raw_setting}) + assert loaded is not None + assert loaded.aliases_by_value == {} + assert check_mcp_client_allowed(loaded, {"azp": "antigravity-cli"}, {"x-mcp-client": "antigravity-cli"}) is not None + + +def test_only_the_value_identifies_a_client_never_its_alias() -> None: + assert check_mcp_client_allowed(JWT_ONLY, {"azp": "Antigravity CLI"}, NO_HEADERS) is not None + assert check_mcp_client_allowed(HEADER_ONLY, None, {"x-mcp-client": "Antigravity CLI"}) is not None + + +def test_two_clients_may_share_an_alias_and_both_are_admitted() -> None: + settings: Final = { + "mcp_allowed_clients": [ + {"alias": "Coding CLI", "value": "cli-dev"}, + {"alias": "Coding CLI", "value": "cli-prod"}, + ], + "litellm_jwtauth": {"mcp_client_id_jwt_field": "azp"}, + } + loaded: Final = load_mcp_client_allowlist(settings) + assert check_mcp_client_allowed(loaded, {"azp": "cli-dev"}, NO_HEADERS) is None + assert check_mcp_client_allowed(loaded, {"azp": "cli-prod"}, NO_HEADERS) is None + assert check_mcp_client_allowed(loaded, {"azp": "Coding CLI"}, NO_HEADERS) is not None + + +def test_unconfigured_allowlist_admits_callers_with_no_identity_at_all() -> None: + assert check_mcp_client_allowed(None, None, NO_HEADERS) is None + assert check_mcp_client_allowed(None, {"azp": "claude-code"}, {"x-mcp-client": "claude-code"}) is None + + +def test_jwt_claim_identifies_the_client() -> None: + assert resolve_mcp_client_identity(JWT_ONLY, {"azp": "antigravity-cli"}, NO_HEADERS) == MCPClientIdentity( + client_id="antigravity-cli", source="jwt", source_name="azp" + ) + assert check_mcp_client_allowed(JWT_ONLY, {"azp": "antigravity-cli"}, NO_HEADERS) is None + + +def test_nested_jwt_claim_path_is_resolved_with_dot_notation() -> None: + nested: Final = MCPClientAllowlist( + aliases_by_value=ANTIGRAVITY_ONLY, jwt_field="resource_access.mcp.client", header=None + ) + claims: Final = {"resource_access": {"mcp": {"client": "antigravity-cli"}}} + assert check_mcp_client_allowed(nested, claims, NO_HEADERS) is None + + +def test_unlisted_jwt_client_is_rejected_and_the_rejection_names_it() -> None: + rejection: Final = check_mcp_client_allowed(JWT_ONLY, {"azp": "claude-code"}, NO_HEADERS) + assert isinstance(rejection, MCPClientRejection) + assert "'claude-code'" in rejection.details + assert "azp" in rejection.details + assert MCP_ALLOWED_CLIENTS_SETTING in rejection.details + assert rejection.response_body == {"error": "Forbidden", "details": rejection.details} + + +@pytest.mark.parametrize("claims", ({"sub": "user-1"}, {"azp": ""}, {"azp": 42}, {"azp": ["antigravity-cli"]})) +def test_jwt_without_a_usable_client_claim_is_rejected(claims: Mapping[str, object]) -> None: + rejection: Final = check_mcp_client_allowed(JWT_ONLY, claims, NO_HEADERS) + assert isinstance(rejection, MCPClientRejection) + assert "azp" in rejection.details + + +def test_matching_is_exact_not_prefix_or_case_insensitive() -> None: + for spoof in ("Antigravity-Cli", "antigravity-cli-sdk", " antigravity-cli"): + assert check_mcp_client_allowed(JWT_ONLY, {"azp": spoof}, NO_HEADERS) is not None + assert check_mcp_client_allowed(HEADER_ONLY, None, {"x-mcp-client": spoof}) is not None + + +def test_configured_header_identifies_callers_without_a_jwt() -> None: + headers: Final = {"x-mcp-client": "antigravity-cli"} + assert resolve_mcp_client_identity(HEADER_ONLY, None, headers) == MCPClientIdentity( + client_id="antigravity-cli", source="header", source_name="x-mcp-client" + ) + assert check_mcp_client_allowed(HEADER_ONLY, None, headers) is None + assert check_mcp_client_allowed(HEADER_ONLY, {}, headers) is None + + +def test_unlisted_or_missing_header_is_rejected() -> None: + unlisted: Final = check_mcp_client_allowed(HEADER_ONLY, None, {"x-mcp-client": "claude-code"}) + assert isinstance(unlisted, MCPClientRejection) + assert "'claude-code'" in unlisted.details + for headers in (NO_HEADERS, {"x-mcp-client": ""}, {"x-other": "antigravity-cli"}): + missing = check_mcp_client_allowed(HEADER_ONLY, None, headers) + assert isinstance(missing, MCPClientRejection) + assert "x-mcp-client" in missing.details + + +def test_header_is_not_consulted_when_it_is_not_configured() -> None: + rejection: Final = check_mcp_client_allowed(JWT_ONLY, None, {"x-mcp-client": "antigravity-cli"}) + assert isinstance(rejection, MCPClientRejection) + assert MCP_CLIENT_ID_HEADER_SETTING in rejection.details + assert MCP_CLIENT_ID_JWT_FIELD_SETTING in rejection.details + + +def test_jwt_caller_is_judged_by_its_claim_even_when_the_header_would_pass() -> None: + spoofed_header: Final = {"x-mcp-client": "antigravity-cli"} + assert check_mcp_client_allowed(JWT_AND_HEADER, {"azp": "claude-code"}, spoofed_header) is not None + assert check_mcp_client_allowed(JWT_AND_HEADER, {"sub": "user-1"}, spoofed_header) is not None + assert check_mcp_client_allowed(JWT_AND_HEADER, {"azp": "antigravity-cli"}, {"x-mcp-client": "claude-code"}) is None + + +def test_jwt_caller_with_an_empty_claim_set_cannot_fall_back_to_the_header() -> None: + rejection: Final = check_mcp_client_allowed(JWT_AND_HEADER, {}, {"x-mcp-client": "antigravity-cli"}) + assert isinstance(rejection, MCPClientRejection) + assert "azp" in rejection.details + + +def test_non_jwt_caller_falls_back_to_the_header_when_both_sources_are_configured() -> None: + assert check_mcp_client_allowed(JWT_AND_HEADER, None, {"x-mcp-client": "antigravity-cli"}) is None + assert check_mcp_client_allowed(JWT_AND_HEADER, None, {"x-mcp-client": "claude-code"}) is not None + + +def test_allowlist_with_no_identity_source_rejects_everyone_and_says_what_to_configure() -> None: + rejection: Final = check_mcp_client_allowed( + NO_SOURCE, {"azp": "antigravity-cli"}, {"x-mcp-client": "antigravity-cli"} + ) + assert isinstance(rejection, MCPClientRejection) + assert MCP_CLIENT_ID_JWT_FIELD_SETTING in rejection.details + assert MCP_CLIENT_ID_HEADER_SETTING in rejection.details + + +def test_empty_allowlist_rejects_an_identified_client() -> None: + empty: Final = MCPClientAllowlist(aliases_by_value={}, jwt_field="azp", header="x-mcp-client") + assert check_mcp_client_allowed(empty, {"azp": "antigravity-cli"}, NO_HEADERS) is not None + assert check_mcp_client_allowed(empty, None, {"x-mcp-client": "antigravity-cli"}) is not None diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py index 0e2ceb98987..33e14736357 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py @@ -1,4 +1,5 @@ import asyncio +import contextlib import contextvars import os from datetime import datetime, timedelta @@ -17,6 +18,8 @@ from mcp.types import ( TextContent, TextResourceContents, ) +from pydantic import TypeAdapter +from starlette.types import Receive, Scope, Send from litellm.proxy._types import ( LiteLLM_MCPServerTable, @@ -2035,6 +2038,220 @@ async def test_mcp_routing_initialize_to_stateful_no_session_to_stateless( assert not any(name.startswith(b"x-mcp-debug") for name in headers) +_FORBIDDEN_BODY_ADAPTER: Final[TypeAdapter[dict[str, str]]] = TypeAdapter(dict[str, str]) +_BODY_CHUNK_ADAPTER: Final[TypeAdapter[bytes]] = TypeAdapter(bytes) +_INITIALIZE: Final = ( + b'{"jsonrpc":"2.0","id":0,"method":"initialize","params":{"protocolVersion":"2025-06-18",' + b'"capabilities":{},"clientInfo":{"name":"antigravity-cli","version":"1.0.0"}}}' +) +_TOOLS_LIST: Final = b'{"jsonrpc":"2.0","id":2,"method":"tools/list","params":{}}' +_ALLOWLIST_SETTINGS: Final[dict[str, object]] = { + "mcp_allowed_clients": [{"alias": "Antigravity CLI", "value": "antigravity-cli"}], + "litellm_jwtauth": {"mcp_client_id_jwt_field": "azp"}, + "mcp_client_id_header": "x-mcp-client", +} +_LISTED_JWT: Final[dict[str, object]] = {"azp": "antigravity-cli", "sub": "user-1"} +_UNLISTED_JWT: Final[dict[str, object]] = {"azp": "claude-code", "sub": "user-1"} +_LISTED_HEADER: Final[list[tuple[bytes, bytes]]] = [(b"x-mcp-client", b"antigravity-cli")] +_UNLISTED_HEADER: Final[list[tuple[bytes, bytes]]] = [(b"x-mcp-client", b"claude-code")] + + +async def _drain_body(receive: Receive) -> bytes: + first: Final = await receive() + body: Final = _BODY_CHUNK_ADAPTER.validate_python(first.get("body", b"")) + if not first.get("more_body", False): + return body + return body + await _drain_body(receive) + + +def _client_allowlist_patches( + settings: dict[str, object], jwt_claims: dict[str, object] | None +) -> contextlib.ExitStack: + stack: Final = contextlib.ExitStack() + stack.enter_context( + patch( # test-quality-ok: the ASGI handler resolves auth through a module-level function; no injection seam + "litellm.proxy._experimental.mcp_server.server.extract_mcp_auth_context", + new_callable=AsyncMock, + return_value=(UserAPIKeyAuth(user_id="allowlist-user", jwt_claims=jwt_claims), None, None, None, None, {}), + ) + ) + stack.enter_context( + patch( # test-quality-ok: module flag guarding lazy session-manager startup; no injection seam + "litellm.proxy._experimental.mcp_server.server._SESSION_MANAGERS_INITIALIZED", True + ) + ) + stack.enter_context( + patch( # test-quality-ok: the allowlist is read off this module global; no injection seam + "litellm.proxy.proxy_server.general_settings", settings + ) + ) + return stack + + +def _forbidden_body(denied: HTTPException) -> dict[str, str]: + return _FORBIDDEN_BODY_ADAPTER.validate_python(denied.detail) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("jwt_claims", "headers", "request_body", "expected_fragment"), + ( + (_UNLISTED_JWT, [], _INITIALIZE, "MCP client 'claude-code' (from JWT claim 'azp')"), + (_UNLISTED_JWT, _LISTED_HEADER, _INITIALIZE, "MCP client 'claude-code' (from JWT claim 'azp')"), + ({"sub": "user-1"}, _LISTED_HEADER, _INITIALIZE, "no 'azp' claim"), + (None, _UNLISTED_HEADER, _INITIALIZE, "MCP client 'claude-code' (from header 'x-mcp-client')"), + (None, [], _INITIALIZE, "no 'x-mcp-client' header"), + (_UNLISTED_JWT, [(b"mcp-session-id", b"session-1")], _TOOLS_LIST, "MCP client 'claude-code'"), + ), +) +async def test_streamable_http_rejects_unlisted_client_before_any_session_work( + jwt_claims: dict[str, object] | None, + headers: list[tuple[bytes, bytes]], + request_body: bytes, + expected_fragment: str, +) -> None: + from litellm.proxy._experimental.mcp_server import server as mcp_module + + scope: Final[Scope] = {"type": "http", "method": "POST", "path": "/mcp", "headers": headers} + receive: Final = AsyncMock(return_value={"type": "http.request", "body": request_body, "more_body": False}) + send: Final = AsyncMock() + stateful_handle: Final = AsyncMock() + stateless_handle: Final = AsyncMock() + session_cap: Final = AsyncMock(return_value=True) + + with ( + _client_allowlist_patches(_ALLOWLIST_SETTINGS, jwt_claims), + patch( # test-quality-ok: session managers are module singletons; the downstream call is the observable + "litellm.proxy._experimental.mcp_server.server.session_manager_stateful", + SimpleNamespace(handle_request=stateful_handle), + ), + patch( # test-quality-ok: session managers are module singletons; the downstream call is the observable + "litellm.proxy._experimental.mcp_server.server.session_manager_stateless", + SimpleNamespace(handle_request=stateless_handle), + ), + patch( # test-quality-ok: module-level cap check; asserting it is never reached is the point + "litellm.proxy._experimental.mcp_server.server._enforce_stateful_session_cap_for_owner", session_cap + ), + pytest.raises(HTTPException) as denied, + ): + await mcp_module.handle_streamable_http_mcp(scope, receive, send) + + assert denied.value.status_code == 403 + body: Final = _forbidden_body(denied.value) + assert body["error"] == "Forbidden" + assert expected_fragment in body["details"] + assert "mcp_allowed_clients" in body["details"] + receive.assert_not_awaited() + send.assert_not_awaited() + stateful_handle.assert_not_awaited() + stateless_handle.assert_not_awaited() + session_cap.assert_not_awaited() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("settings", "jwt_claims", "headers"), + ( + (_ALLOWLIST_SETTINGS, _LISTED_JWT, []), + (_ALLOWLIST_SETTINGS, _LISTED_JWT, _UNLISTED_HEADER), + (_ALLOWLIST_SETTINGS, None, _LISTED_HEADER), + ({}, _UNLISTED_JWT, _UNLISTED_HEADER), + ({}, None, []), + ), +) +async def test_streamable_http_admits_listed_or_unrestricted_clients_and_hands_the_body_downstream( + settings: dict[str, object], jwt_claims: dict[str, object] | None, headers: list[tuple[bytes, bytes]] +) -> None: + from litellm.proxy._experimental.mcp_server import server as mcp_module + + scope: Final[Scope] = {"type": "http", "method": "POST", "path": "/mcp", "headers": headers} + receive: Final = AsyncMock( + side_effect=[ + {"type": "http.request", "body": _INITIALIZE[:20], "more_body": True}, + {"type": "http.request", "body": _INITIALIZE[20:], "more_body": False}, + ] + ) + send: Final = AsyncMock() + downstream_bodies: Final[list[bytes]] = [] + + async def handle_request(_: Scope, downstream_receive: Receive, __: Send) -> None: + downstream_bodies.append(await _drain_body(downstream_receive)) + + stateful_handle: Final = AsyncMock(side_effect=handle_request) + stateless_handle: Final = AsyncMock() + + with ( + _client_allowlist_patches(settings, jwt_claims), + patch( # test-quality-ok: session managers are module singletons; the downstream call is the observable + "litellm.proxy._experimental.mcp_server.server.session_manager_stateful", + SimpleNamespace(handle_request=stateful_handle), + ), + patch( # test-quality-ok: session managers are module singletons; the downstream call is the observable + "litellm.proxy._experimental.mcp_server.server.session_manager_stateless", + SimpleNamespace(handle_request=stateless_handle), + ), + ): + await mcp_module.handle_streamable_http_mcp(scope, receive, send) + + assert downstream_bodies == [_INITIALIZE] + stateless_handle.assert_not_awaited() + send.assert_not_awaited() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("jwt_claims", "headers", "admitted"), + ( + (_LISTED_JWT, [], True), + (None, _LISTED_HEADER, True), + (_UNLISTED_JWT, _LISTED_HEADER, False), + (None, _UNLISTED_HEADER, False), + (None, [], False), + ), +) +async def test_sse_endpoint_applies_the_same_client_allowlist( + jwt_claims: dict[str, object] | None, headers: list[tuple[bytes, bytes]], admitted: bool +) -> None: + from litellm.proxy._experimental.mcp_server import server as mcp_module + + scope: Final[Scope] = {"type": "http", "method": "POST", "path": "/mcp/sse", "headers": headers} + receive: Final = AsyncMock(return_value={"type": "http.request", "body": _INITIALIZE, "more_body": False}) + send: Final = AsyncMock() + downstream_bodies: Final[list[bytes]] = [] + + async def handle_request(_: Scope, downstream_receive: Receive, __: Send) -> None: + downstream_bodies.append(await _drain_body(downstream_receive)) + + with ( + _client_allowlist_patches(_ALLOWLIST_SETTINGS, jwt_claims), + patch( # test-quality-ok: module-level pre-auth probe unrelated to the allowlist under test; no injection seam + "litellm.proxy._experimental.mcp_server.server._raise_preemptive_401_for_unauthenticated_servers", + new_callable=AsyncMock, + ), + patch( # test-quality-ok: module-level upstream auth probe unrelated to the allowlist under test; no injection seam + "litellm.proxy._experimental.mcp_server.server._check_passthrough_upstream_auth", + new_callable=AsyncMock, + ), + patch.object( # test-quality-ok: SSE manager is a module singleton; the downstream call is the observable + mcp_module.sse_session_manager, "handle_request", side_effect=handle_request + ), + ): + if admitted: + await mcp_module.handle_sse_mcp(scope, receive, send) + assert downstream_bodies == [_INITIALIZE] + send.assert_not_awaited() + return + with pytest.raises(HTTPException) as denied: + await mcp_module.handle_sse_mcp(scope, receive, send) + + assert denied.value.status_code == 403 + body: Final = _forbidden_body(denied.value) + assert body["error"] == "Forbidden" + assert "mcp_allowed_clients" in body["details"] + assert downstream_bodies == [] + send.assert_not_awaited() + + @pytest.mark.asyncio async def test_mcp_routing_chunked_initialize_to_stateful(): """ @@ -5422,11 +5639,12 @@ async def test_get_tools_from_mcp_servers_logs_list_tools_to_spendlogs_when_enab Ensure list-tools logging path calls `async_success_handler` when enabled. """ try: + from mcp.types import Tool as MCPTool + from litellm.proxy._experimental.mcp_server.server import ( _get_tools_from_mcp_servers, ) from litellm.proxy._types import UserAPIKeyAuth - from mcp.types import Tool as MCPTool except ImportError: pytest.skip("MCP server not available") @@ -8376,10 +8594,10 @@ async def test_fire_mcp_tool_call_logging_iserror_logs_failure(): """Regression test: a CallToolResult with isError=True must go down the failure logging path (async_failure_handler + post_call_failure_hook), never async_success_handler.""" + from litellm.proxy._experimental.mcp_server.exceptions import MCPToolResultError from litellm.proxy._experimental.mcp_server.server import ( _fire_mcp_tool_call_logging, ) - from litellm.proxy._experimental.mcp_server.exceptions import MCPToolResultError logging_obj = _mock_mcp_logging_obj() proxy_logging_mock = _mock_mcp_proxy_logging() @@ -8729,11 +8947,11 @@ async def test_call_mcp_tool_skips_failure_hook_for_upstream_auth_error(): caller-must-reauth signal, not a failed call, so call_mcp_tool must re-raise it WITHOUT firing post_call_failure_hook (which records a failure and can trip LLM exception alerts). The streamable handler downgrades it to an informational isError result afterward.""" + from litellm.proxy._experimental.mcp_server.exceptions import MCPUpstreamAuthError from litellm.proxy._experimental.mcp_server.server import ( call_mcp_tool, global_mcp_server_manager, ) - from litellm.proxy._experimental.mcp_server.exceptions import MCPUpstreamAuthError from litellm.proxy._types import MCPTransport, UserAPIKeyAuth from litellm.types.mcp_server.mcp_server_manager import MCPServer diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py index 4ec4ae31ca6..ac3ad9ed89e 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py @@ -4313,3 +4313,131 @@ class TestV1ResolvedOauth2Gate: assert rest_endpoints._v1_resolved_oauth2_server_ids(["oauth2-srv"]) == set() assert rest_endpoints._v1_resolved_oauth2_server_ids(["oauth2-srv", "delegate-srv"]) == {"delegate-srv"} + + +_CLIENT_ALLOWLIST_SETTINGS: Final[dict[str, object]] = { + "mcp_allowed_clients": [{"alias": "Antigravity CLI", "value": "antigravity-cli"}], + "litellm_jwtauth": {"mcp_client_id_jwt_field": "azp"}, + "mcp_client_id_header": "x-mcp-client", +} + + +class TestClientAllowlistOnRestRoutes: + """``mcp_allowed_clients`` must gate the REST tool facade exactly like the /mcp transports, + otherwise an unlisted harness can list and call tools by switching to /mcp-rest.""" + + pytestmark = pytest.mark.asyncio + + @staticmethod + def _stub_listing(monkeypatch: pytest.MonkeyPatch) -> list[UserAPIKeyAuth]: + listed_for: list[UserAPIKeyAuth] = [] + + async def fake_contexts(user_api_key_auth: UserAPIKeyAuth) -> list[UserAPIKeyAuth]: + listed_for.append(user_api_key_auth) + return [user_api_key_auth] + + async def fake_get_allowed_mcp_servers( + user_api_key_auth: UserAPIKeyAuth | None = None, + *, + keyless_source: bool = False, + ) -> list[str]: + return [] + + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", _CLIENT_ALLOWLIST_SETTINGS, raising=False) + monkeypatch.setattr(rest_endpoints, "build_effective_auth_contexts", fake_contexts, raising=False) + monkeypatch.setattr( + rest_endpoints.global_mcp_server_manager, + "get_allowed_mcp_servers", + fake_get_allowed_mcp_servers, + raising=False, + ) + return listed_for + + @pytest.mark.parametrize( + ("caller", "headers", "expected_fragment"), + ( + (UserAPIKeyAuth(jwt_claims={"azp": "claude-code"}), {"x-mcp-client": "antigravity-cli"}, "'claude-code'"), + (UserAPIKeyAuth(jwt_claims={}), {"x-mcp-client": "antigravity-cli"}, "no 'azp' claim"), + (UserAPIKeyAuth(), {"x-mcp-client": "claude-code"}, "'claude-code'"), + (UserAPIKeyAuth(), {}, "no 'x-mcp-client' header"), + ), + ) + async def test_tools_list_rejects_unlisted_clients_before_resolving_servers( + self, + monkeypatch: pytest.MonkeyPatch, + caller: UserAPIKeyAuth, + headers: dict[str, str], + expected_fragment: str, + ) -> None: + listed_for: Final = self._stub_listing(monkeypatch) + request: Final = _build_request(headers, path="/mcp-rest/tools/list", method="GET") + + with pytest.raises(HTTPException) as denied: + await rest_endpoints.list_tool_rest_api( + request, server_id=None, mcp_server_name=None, toolset_name=None, user_api_key_dict=caller + ) + + assert denied.value.status_code == 403 + assert denied.value.detail["error"] == "Forbidden" + assert expected_fragment in denied.value.detail["details"] + assert "mcp_allowed_clients" in denied.value.detail["details"] + assert listed_for == [] + + @pytest.mark.parametrize( + ("caller", "headers"), + ( + (UserAPIKeyAuth(jwt_claims={"azp": "antigravity-cli"}), {"x-mcp-client": "claude-code"}), + (UserAPIKeyAuth(), {"x-mcp-client": "antigravity-cli"}), + ), + ) + async def test_tools_list_admits_listed_clients( + self, monkeypatch: pytest.MonkeyPatch, caller: UserAPIKeyAuth, headers: dict[str, str] + ) -> None: + listed_for: Final = self._stub_listing(monkeypatch) + request: Final = _build_request(headers, path="/mcp-rest/tools/list", method="GET") + + result: Final = await rest_endpoints.list_tool_rest_api( + request, server_id=None, mcp_server_name=None, toolset_name=None, user_api_key_dict=caller + ) + + assert result["tools"] == [] + assert listed_for == [caller] + + async def test_dashboard_session_is_not_treated_as_a_client_application( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + from litellm.constants import UI_SESSION_TOKEN_TEAM_ID + + listed_for: Final = self._stub_listing(monkeypatch) + session: Final = UserAPIKeyAuth(team_id=UI_SESSION_TOKEN_TEAM_ID, user_id="admin-user", user_role="proxy_admin") + request: Final = _build_request(path="/mcp-rest/tools/list", method="GET") + + result: Final = await rest_endpoints.list_tool_rest_api( + request, server_id=None, mcp_server_name=None, toolset_name=None, user_api_key_dict=session + ) + + assert result["tools"] == [] + assert listed_for == [session] + + async def test_tools_call_rejects_unlisted_clients_before_reading_the_body( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", _CLIENT_ALLOWLIST_SETTINGS, raising=False) + acting: Final = AsyncMock() + monkeypatch.setattr(rest_endpoints, "acting_user_auth", acting, raising=False) + request: Final = _build_request( + {"x-mcp-client": "antigravity-cli"}, + path="/mcp-rest/tools/call", + method="POST", + json_body={"server_id": "server-1", "name": "demo-tool", "arguments": {}}, + ) + + with pytest.raises(HTTPException) as denied: + await rest_endpoints.call_tool_rest_api( + request, user_api_key_dict=UserAPIKeyAuth(jwt_claims={"azp": "claude-code"}) + ) + + assert denied.value.status_code == 403 + assert denied.value.detail["error"] == "Forbidden" + assert "'claude-code'" in denied.value.detail["details"] + acting.assert_not_awaited() diff --git a/tests/test_litellm/proxy/auth/test_route_checks.py b/tests/test_litellm/proxy/auth/test_route_checks.py index 72c59223549..2e0d2c710a4 100644 --- a/tests/test_litellm/proxy/auth/test_route_checks.py +++ b/tests/test_litellm/proxy/auth/test_route_checks.py @@ -4038,3 +4038,37 @@ def test_team_key_without_service_account_marker_still_rejected(): valid_token=valid_token, request_data={}, ) + + +@pytest.mark.parametrize("route", ["/project/new", "/project/update"]) +def test_project_write_routes_reach_endpoint_for_internal_user(route): + """The route gate lets a non-admin through so /project/new and /project/update can apply the + team_admin_editable_team_fields projects permission themselves, instead of a blanket 401.""" + valid_token = UserAPIKeyAuth(user_id="team_admin", user_role=LitellmUserRoles.INTERNAL_USER.value) + request = MagicMock(spec=Request) + request.query_params = {} + + RouteChecks.non_proxy_admin_allowed_routes_check( + user_obj=None, + _user_role=LitellmUserRoles.INTERNAL_USER.value, + route=route, + request=request, + valid_token=valid_token, + request_data={}, + ) + + +def test_project_delete_route_stays_proxy_admin_only(): + valid_token = UserAPIKeyAuth(user_id="team_admin", user_role=LitellmUserRoles.INTERNAL_USER.value) + request = MagicMock(spec=Request) + request.query_params = {} + + with pytest.raises(Exception, match="Only proxy admin can be used to generate, delete, update"): + RouteChecks.non_proxy_admin_allowed_routes_check( + user_obj=None, + _user_role=LitellmUserRoles.INTERNAL_USER.value, + route="/project/delete", + request=request, + valid_token=valid_token, + request_data={}, + ) diff --git a/tests/test_litellm/proxy/hooks/test_batch_rate_limiter.py b/tests/test_litellm/proxy/hooks/test_batch_rate_limiter.py index 919e9c79828..930f62fcd10 100644 --- a/tests/test_litellm/proxy/hooks/test_batch_rate_limiter.py +++ b/tests/test_litellm/proxy/hooks/test_batch_rate_limiter.py @@ -6,7 +6,10 @@ batch under a per-minute RPM/TPM budget. Scopes that configure `tpd_limit` are charged against a 24h token window instead of their minute counters. """ -from datetime import datetime +import time +from collections.abc import Iterator +from datetime import datetime, timezone +from typing import Final import pytest from fastapi import HTTPException @@ -257,3 +260,39 @@ def test_online_descriptors_ignore_tpd_limit(): model_has_failures=False, ) assert [(d["key"], d["rate_limit"]["window_size"]) for d in descriptors] == [("api_key", rate_limiter.window_size)] + + +@pytest.fixture(params=["Europe/Paris", "Asia/Kolkata", "America/Los_Angeles"]) +def process_timezone(request: pytest.FixtureRequest, monkeypatch: pytest.MonkeyPatch) -> Iterator[str]: + monkeypatch.setenv("TZ", request.param) + time.tzset() + yield request.param + monkeypatch.undo() + time.tzset() + + +@pytest.mark.skipif(not hasattr(time, "tzset"), reason="switching the process timezone needs time.tzset()") +@pytest.mark.asyncio +async def test_batch_rate_limit_error_reports_reset_time_in_utc_on_a_non_utc_proxy(process_timezone: str) -> None: + window_start: Final = datetime(2026, 9, 13, 8, 0, 0, tzinfo=timezone.utc) + clock: Final = _Clock(window_start) + _internal_usage_cache, _rate_limiter, batch_limiter = _make_limiters(clock) + user_api_key_dict: Final = UserAPIKeyAuth(api_key=hash_token("tpd-key-utc"), rpm_limit=1, tpd_limit=1000) + + await batch_limiter._check_and_increment_batch_counters( + user_api_key_dict=user_api_key_dict, + data={}, + batch_usage=BatchFileUsage(total_tokens=600, request_count=6), + ) + clock.now = datetime(2026, 9, 13, 11, 0, 0, tzinfo=timezone.utc) + with pytest.raises(HTTPException) as exc: + await batch_limiter._check_and_increment_batch_counters( + user_api_key_dict=user_api_key_dict, + data={}, + batch_usage=BatchFileUsage(total_tokens=600, request_count=6), + ) + + assert exc.value.status_code == 429 + assert exc.value.headers["retry-after"] == str(BATCH_TPD_WINDOW_SECONDS - 3 * 3600) + assert exc.value.headers["reset_at"] == "2026-09-14 08:00:00 UTC" + assert str(exc.value.detail).endswith("Limit resets at: 2026-09-14 08:00:00 UTC") diff --git a/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py index 5889b1b513f..4907b4ea054 100644 --- a/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py +++ b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py @@ -7,10 +7,10 @@ import logging import os import sys import time -from collections.abc import Sequence +from collections.abc import Iterator, Sequence from contextlib import contextmanager -from datetime import datetime, timedelta -from typing import Any, Dict, List, Optional +from datetime import datetime, timedelta, timezone +from typing import Any, Dict, Final, List, Optional import pytest from fastapi import HTTPException @@ -21,10 +21,12 @@ from litellm.caching.caching import DualCache from litellm.caching.in_memory_cache import InMemoryCache from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError from litellm.proxy.hooks.parallel_request_limiter_v3 import ( PARALLEL_REQUEST_SLOT_TTL_SECONDS, ParallelSlotAcquisition, RateLimitDescriptor, + RateLimitResponse, RequestRateLimiterStash, _request_stash, get_or_create_request_stash, @@ -6911,3 +6913,51 @@ def test_success_tpm_accounting_skips_team_model_pool_when_key_owns_model_tpm_li assert handler.create_rate_limit_keys("model_per_key", f"{hash_token('sk-pool')}:test-model", "tokens") in charged_keys team_pool_key = handler.create_rate_limit_keys("model_per_team", "t:test-model", "tokens") assert (team_pool_key in charged_keys) is charges_team_model_pool + + +@pytest.fixture(params=["Europe/Paris", "Asia/Kolkata", "America/Los_Angeles"]) +def process_timezone(request: pytest.FixtureRequest, monkeypatch: pytest.MonkeyPatch) -> Iterator[str]: + monkeypatch.setenv("TZ", request.param) + time.tzset() + yield request.param + monkeypatch.undo() + time.tzset() + + +@pytest.mark.skipif(not hasattr(time, "tzset"), reason="switching the process timezone needs time.tzset()") +def test_rate_limit_error_reports_reset_time_in_utc_on_a_non_utc_proxy(process_timezone: str) -> None: + now: Final = datetime(2026, 9, 4, 21, 53, 21, tzinfo=timezone.utc) + handler: Final = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(DualCache()), time_provider=lambda: now + ) + expected_reset: Final = (now + timedelta(seconds=handler.window_size)).strftime("%Y-%m-%d %H:%M:%S UTC") + over_limit: Final[RateLimitResponse] = { + "overall_code": "OVER_LIMIT", + "statuses": [ + { + "code": "OVER_LIMIT", + "descriptor_key": "api_key", + "limit_remaining": 0, + "rate_limit_type": "requests", + "current_limit": 2, + } + ], + } + + with pytest.raises(ProxyRateLimitError) as exc_info: + handler._handle_rate_limit_error( + response=over_limit, + descriptors=[{"key": "api_key", "value": "sk-test", "rate_limit": None}], + requested_model="gpt-4o-mini", + ) + + assert exc_info.value.status_code == 429 + assert exc_info.value.headers == { + "retry-after": str(handler.window_size), + "rate_limit_type": "requests", + "reset_at": expected_reset, + } + assert exc_info.value.detail == ( + "Rate limit exceeded for api_key: sk-test. Limit type: requests. " + f"Current limit: 2, Remaining: 0. Limit resets at: {expected_reset}" + ) diff --git a/tests/test_litellm/proxy/management_endpoints/test_project_org_authz.py b/tests/test_litellm/proxy/management_endpoints/test_project_org_authz.py index a06d79306ab..ce08030739d 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_project_org_authz.py +++ b/tests/test_litellm/proxy/management_endpoints/test_project_org_authz.py @@ -7,12 +7,18 @@ Unit tests for the VERIA-55 fixes: member of. """ +from types import MappingProxyType +from typing import Final from unittest.mock import AsyncMock, MagicMock import pytest from fastapi import HTTPException -from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth +from litellm.models.team import LiteLLM_TeamTable +from litellm.proxy._types import LitellmUserRoles, Member, UserAPIKeyAuth + +_PROJECTS_ENABLED: Final = MappingProxyType({"team_admin_editable_team_fields": ["projects"]}) +_PROJECTS_DISABLED: Final = MappingProxyType({"team_admin_editable_team_fields": ["max_budget"]}) # --------------------------------------------------------------------------- @@ -20,11 +26,9 @@ from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth # --------------------------------------------------------------------------- -def _make_prisma_with_team(team_id: str, admins: list): +def _make_prisma_with_team(team_id: str, admins: list, members_with_roles: tuple[Member, ...] = ()): prisma = MagicMock() - team_row = MagicMock() - team_row.team_id = team_id - team_row.admins = admins + team_row = LiteLLM_TeamTable(team_id=team_id, admins=admins, members_with_roles=list(members_with_roles)) prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=team_row) return prisma @@ -49,6 +53,7 @@ async def test_project_perm_check_uses_current_team_not_caller_supplied(): user_api_key_dict=caller, team_id="team-A", prisma_client=prisma, + general_settings=_PROJECTS_ENABLED, ) assert has_perm is False prisma.db.litellm_teamtable.find_unique.assert_awaited_once() @@ -70,10 +75,105 @@ async def test_project_perm_check_allows_team_admin_of_existing_team(): user_api_key_dict=alice, team_id="team-A", prisma_client=prisma, + general_settings=_PROJECTS_ENABLED, ) assert has_perm is True +@pytest.mark.asyncio +async def test_project_perm_check_allows_members_with_roles_admin(): + """Team admins added through /team/member_add live in members_with_roles, not the legacy admins list.""" + from litellm_enterprise.proxy.management_endpoints.project_endpoints import ( + _check_user_permission_for_project, + ) + + prisma = _make_prisma_with_team( + team_id="team-A", + admins=[], + members_with_roles=(Member(user_id="carol", role="admin"), Member(user_id="dave", role="user")), + ) + carol = UserAPIKeyAuth(user_id="carol", user_role=LitellmUserRoles.INTERNAL_USER.value) + dave = UserAPIKeyAuth(user_id="dave", user_role=LitellmUserRoles.INTERNAL_USER.value) + + assert ( + await _check_user_permission_for_project( + user_api_key_dict=carol, team_id="team-A", prisma_client=prisma, general_settings=_PROJECTS_ENABLED + ) + is True + ) + assert ( + await _check_user_permission_for_project( + user_api_key_dict=dave, team_id="team-A", prisma_client=prisma, general_settings=_PROJECTS_ENABLED + ) + is False + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("general_settings", [MappingProxyType({}), _PROJECTS_DISABLED]) +async def test_project_perm_check_denies_team_admin_unless_projects_permission_configured(general_settings): + from litellm_enterprise.proxy.management_endpoints.project_endpoints import ( + _check_user_permission_for_project, + ) + + prisma = _make_prisma_with_team( + team_id="team-A", admins=["alice"], members_with_roles=(Member(user_id="carol", role="admin"),) + ) + + for user_id in ("alice", "carol"): + caller = UserAPIKeyAuth(user_id=user_id, user_role=LitellmUserRoles.INTERNAL_USER.value) + has_perm = await _check_user_permission_for_project( + user_api_key_dict=caller, + team_id="team-A", + prisma_client=prisma, + general_settings=general_settings, + ) + assert has_perm is False + prisma.db.litellm_teamtable.find_unique.assert_not_called() + + +@pytest.mark.asyncio +async def test_project_perm_check_require_admin_denies_team_admin_even_when_configured(): + """/project/delete passes require_admin=True, so the projects permission must not open it up.""" + from litellm_enterprise.proxy.management_endpoints.project_endpoints import ( + _check_user_permission_for_project, + ) + + prisma = _make_prisma_with_team(team_id="team-A", admins=["alice"]) + alice = UserAPIKeyAuth(user_id="alice", user_role=LitellmUserRoles.INTERNAL_USER.value) + + has_perm = await _check_user_permission_for_project( + user_api_key_dict=alice, + team_id=None, + prisma_client=prisma, + general_settings=_PROJECTS_ENABLED, + require_admin=True, + ) + assert has_perm is False + prisma.db.litellm_teamtable.find_unique.assert_not_called() + + +@pytest.mark.asyncio +async def test_project_perm_check_uses_injected_team_object_for_reassignment_target(): + from litellm_enterprise.proxy.management_endpoints.project_endpoints import ( + _check_user_permission_for_project, + ) + + prisma = _make_prisma_with_team(team_id="team-A", admins=["alice"]) + alice = UserAPIKeyAuth(user_id="alice", user_role=LitellmUserRoles.INTERNAL_USER.value) + target_team = LiteLLM_TeamTable(team_id="team-B", members_with_roles=[Member(user_id="erin", role="admin")]) + + has_perm = await _check_user_permission_for_project( + user_api_key_dict=alice, + team_id="team-B", + prisma_client=prisma, + general_settings=_PROJECTS_ENABLED, + team_object=target_team, + ) + assert has_perm is False + prisma.db.litellm_teamtable.find_unique.assert_not_called() + + @pytest.mark.asyncio async def test_project_perm_check_proxy_admin_always_allowed(): from litellm_enterprise.proxy.management_endpoints.project_endpoints import ( @@ -90,6 +190,7 @@ async def test_project_perm_check_proxy_admin_always_allowed(): user_api_key_dict=admin, team_id="team-A", prisma_client=prisma, + general_settings=MappingProxyType({}), ) assert has_perm is True # Admin shortcut should not even hit the DB. diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_admin_field_permissions.py b/tests/test_litellm/proxy/management_endpoints/test_team_admin_field_permissions.py index 5b31089f91e..1a72d1de393 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_admin_field_permissions.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_admin_field_permissions.py @@ -9,6 +9,7 @@ from litellm.proxy.management_endpoints.team_admin_field_permissions import ( changed_team_fields, resolve_team_admin_editable_fields, team_admin_edit_verdict, + team_admin_may_manage_projects, team_admin_request_or_raise, ) @@ -31,6 +32,25 @@ class TestResolveTeamAdminEditableFields: def test_malformed_setting_fails_closed(self, raw): assert resolve_team_admin_editable_fields({"team_admin_editable_team_fields": raw}, _SUPPORTED) == frozenset() + def test_projects_permission_is_not_a_team_field(self): + configured = {"team_admin_editable_team_fields": ["projects", "tpm_limit"]} + assert resolve_team_admin_editable_fields(configured, _SUPPORTED) == frozenset({"tpm_limit"}) + + +class TestTeamAdminMayManageProjects: + def test_missing_setting_denies(self): + assert team_admin_may_manage_projects({}) is False + + def test_team_fields_alone_do_not_grant_projects(self): + assert team_admin_may_manage_projects({"team_admin_editable_team_fields": ["tpm_limit", "max_budget"]}) is False + + def test_projects_entry_grants(self): + assert team_admin_may_manage_projects({"team_admin_editable_team_fields": ["max_budget", "projects"]}) is True + + @pytest.mark.parametrize("raw", ["projects", 7, [1, 2]]) + def test_malformed_setting_denies(self, raw): + assert team_admin_may_manage_projects({"team_admin_editable_team_fields": raw}) is False + class TestChangedTeamFields: def test_team_id_alone_changes_nothing(self): diff --git a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py index f4490519554..88d38d74f49 100644 --- a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py +++ b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py @@ -333,6 +333,36 @@ async def test_arrival_time_prefers_litellm_received_at_over_time_time(): assert updated_data["proxy_server_request"]["arrival_time"] == received_at.timestamp() +@pytest.mark.asyncio +async def test_proxy_clears_client_supplied_timing_windows(): + request_mock = MagicMock(spec=Request) + request_mock.url = MagicMock() + request_mock.url.path = "/v1/chat/completions" + request_mock.url.__str__.return_value = "http://localhost/v1/chat/completions" + request_mock.method = "POST" + request_mock.query_params = {} + request_mock.headers = {"Content-Type": "application/json"} + request_mock.client = MagicMock() + request_mock.client.host = "127.0.0.1" + request_mock.state = SimpleNamespace(litellm_received_at=datetime.now(timezone.utc)) + + user_api_key_dict = UserAPIKeyAuth(api_key="hashed-key", metadata={}, team_metadata={}) + + updated_data = await add_litellm_data_to_request( + data={ + "model": "gpt-3.5-turbo", + "metadata": {"llm_api_timing_windows": ((0.0, 1.0),)}, + }, + request=request_mock, + user_api_key_dict=user_api_key_dict, + proxy_config=MagicMock(), + general_settings={}, + version="test-version", + ) + + assert updated_data["metadata"]["llm_api_timing_windows"] == () + + @pytest.mark.asyncio async def test_arrival_time_falls_back_to_time_time_without_litellm_received_at(): """Callers that never went through user_api_key_auth (no stamp on request.state) diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 2200372b567..8ec24e25326 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -14314,6 +14314,29 @@ async def test_update_general_settings_keeps_yaml_openai_websocket_passthrough() assert ps.general_settings["enable_openai_websocket_passthrough"] is False +def test_settings_store_exposes_dashboard_saved_mcp_client_allowlist_to_the_mcp_gateway() -> None: + from litellm.proxy._experimental.mcp_server.client_allowlist import MCPClientAllowlist, load_mcp_client_allowlist + from litellm.proxy.proxy_server import ProxyConfig + + settings: Final = ProxyConfig().settings + settings.load_yaml({"litellm_jwtauth": {"mcp_client_id_jwt_field": "azp"}}) + assert load_mcp_client_allowlist(settings) is None + + settings.apply_db_row( + "general_settings", + { + "mcp_allowed_clients": [{"alias": "Antigravity CLI", "value": "antigravity-cli"}], + "mcp_client_id_header": "X-MCP-Client", + }, + ) + assert load_mcp_client_allowlist(settings) == MCPClientAllowlist( + aliases_by_value={"antigravity-cli": "Antigravity CLI"}, jwt_field="azp", header="x-mcp-client" + ) + + settings.apply_db_row("general_settings", {"mcp_client_id_header": "X-MCP-Client"}) + assert load_mcp_client_allowlist(settings) is None + + async def test_token_counter_keeps_the_event_loop_free_during_a_huggingface_count(monkeypatch): from tests.large_text import text from tests.test_litellm.litellm_core_utils.event_loop_lag import ( diff --git a/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py b/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py index fc5733b1a82..9860d1bf94a 100644 --- a/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py +++ b/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py @@ -3291,7 +3291,7 @@ class TestTeamAdminEditableTeamFieldsSetting: def test_patch_rejects_field_names_the_proxy_does_not_support(self, monkeypatch): mock_prisma = self._as_proxy_admin(monkeypatch) monkeypatch.setattr( - "litellm.proxy.ui_crud_endpoints.proxy_setting_endpoints.SUPPORTED_TEAM_ADMIN_EDITABLE_TEAM_FIELDS", + "litellm.proxy.ui_crud_endpoints.proxy_setting_endpoints.SUPPORTED_TEAM_ADMIN_PERMISSIONS", frozenset({"tpm_limit"}), ) @@ -3336,6 +3336,26 @@ class TestTeamAdminEditableTeamFieldsSetting: assert stored["team_admin_editable_team_fields"] == enabled assert general_settings["team_admin_editable_team_fields"] == enabled + def test_patch_accepts_the_projects_permission_and_project_endpoints_see_it(self, monkeypatch): + from litellm.proxy.management_endpoints.team_admin_field_permissions import ( + team_admin_may_manage_projects, + ) + + mock_prisma = self._as_proxy_admin(monkeypatch) + general_settings: dict = {} + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", general_settings) + assert team_admin_may_manage_projects(general_settings) is False + + try: + response = client.patch("/update/ui_settings", json={"team_admin_editable_team_fields": ["projects"]}) + finally: + app.dependency_overrides.clear() + + assert response.status_code == 200 + stored = json.loads(mock_prisma.db.litellm_uisettings.upsert.call_args.kwargs["data"]["create"]["ui_settings"]) + assert stored["team_admin_editable_team_fields"] == ["projects"] + assert team_admin_may_manage_projects(general_settings) is True + def test_patch_with_an_empty_list_turns_team_admin_editing_off_again(self, monkeypatch): mock_prisma = self._as_proxy_admin(monkeypatch) general_settings: dict = {"team_admin_editable_team_fields": ["tpm_limit"]} @@ -3372,6 +3392,7 @@ class TestTeamAdminEditableTeamFieldsSetting: assert field_schema["type"] == "array" assert field_schema["items"]["type"] == "string" assert "tpm_limit" in field_schema["items"]["enum"] + assert "projects" in field_schema["items"]["enum"] class TestSyncUiSettingsToGeneralSettings: diff --git a/tests/test_litellm/realtime_api/test_main.py b/tests/test_litellm/realtime_api/test_main.py index 643e65af47c..86b25b2f9c8 100644 --- a/tests/test_litellm/realtime_api/test_main.py +++ b/tests/test_litellm/realtime_api/test_main.py @@ -499,3 +499,69 @@ async def test_arealtime_azure_env_beta_protocol_wins_over_a_ga_client(monkeypat assert await _azure_backend_url_dialed_for(_GA_CLIENT) == ( "wss://my-endpoint.openai.azure.com/openai/realtime?api-version=2024-10-01-preview&deployment=gpt-realtime" ) + + +async def _vertex_provider_config_for(monkeypatch, model: str, vertex_location: str | None): + from litellm.llms.vertex_ai.realtime.transformation import VertexAIRealtimeConfig + from litellm.llms.vertex_ai.audio_transcription.realtime_transformation import VertexChirpRealtimeConfig + + captured: dict[str, object] = {} + + def mock_get_llm_provider(model, api_base, api_key): + return model.removeprefix("vertex_ai/"), "vertex_ai", None, api_base + + async def mock_token_resolver(**kwargs): + return "access-token", kwargs["project_id"] + + async def mock_async_realtime(**kwargs): + captured.update(kwargs) + + monkeypatch.setattr(realtime_main, "get_llm_provider", mock_get_llm_provider) + monkeypatch.setattr(realtime_main, "vertex_access_token_resolver", mock_token_resolver) + monkeypatch.setattr(realtime_main.base_llm_http_handler, "async_realtime", mock_async_realtime) + monkeypatch.setattr(litellm, "vertex_location", None) + monkeypatch.delenv("VERTEXAI_LOCATION", raising=False) + await realtime_main._arealtime.__wrapped__( + model=model, + websocket=MagicMock(), + litellm_logging_obj=FakeLogging(), + query_params={"model": model, "intent": "transcription"}, + vertex_credentials="fake-credentials", + vertex_project="proj-1", + vertex_location=vertex_location, + ) + provider_config = captured["provider_config"] + assert isinstance(provider_config, (VertexAIRealtimeConfig, VertexChirpRealtimeConfig)) + return provider_config, captured["model"] + + +@pytest.mark.asyncio +async def test_arealtime_routes_chirp_models_to_the_speech_to_text_backend(monkeypatch): + from litellm.llms.vertex_ai.audio_transcription.realtime_transformation import VertexChirpRealtimeConfig + + provider_config, model = await _vertex_provider_config_for(monkeypatch, "vertex_ai/chirp_3", None) + assert isinstance(provider_config, VertexChirpRealtimeConfig) + assert model == "chirp_3" + assert provider_config.get_complete_url(None, model) == "us-speech.googleapis.com" + assert provider_config.validate_environment({}, model, "https://us-speech.googleapis.com") == {} + + +@pytest.mark.asyncio +async def test_arealtime_routes_chirp_models_to_the_configured_speech_region(monkeypatch): + provider_config, model = await _vertex_provider_config_for(monkeypatch, "vertex_ai/chirp_3", "europe-west4") + assert provider_config.get_complete_url(None, model) == "europe-west4-speech.googleapis.com" + + +@pytest.mark.asyncio +async def test_arealtime_keeps_gemini_live_on_the_vertex_realtime_websocket(monkeypatch): + from litellm.llms.vertex_ai.realtime.transformation import VertexAIRealtimeConfig + + provider_config, model = await _vertex_provider_config_for(monkeypatch, "vertex_ai/gemini-live-2.5-flash", None) + assert isinstance(provider_config, VertexAIRealtimeConfig) + assert provider_config.get_complete_url(None, model).startswith("wss://us-central1-aiplatform.googleapis.com/") + + +@pytest.mark.asyncio +async def test_realtime_health_check_names_the_batch_mode_for_chirp_models(): + with pytest.raises(ValueError, match="mode audio_transcription"): + await realtime_main._realtime_health_check(model="chirp_3", custom_llm_provider="vertex_ai", api_key=None) diff --git a/tests/test_litellm/rust_bridge/test_settings.py b/tests/test_litellm/rust_bridge/test_settings.py new file mode 100644 index 00000000000..f75145c2b2c --- /dev/null +++ b/tests/test_litellm/rust_bridge/test_settings.py @@ -0,0 +1,75 @@ +import dataclasses +import logging +from pathlib import Path +from typing import Final + +import pytest +from pydantic import TypeAdapter + +import litellm +from litellm.llms.custom_httpx.http_handler import default_user_agent +from litellm.rust_bridge import settings + +CONTRACT_PATH: Final = Path(__file__).parents[3] / "litellm-rust/crates/python-bridge/python_settings.json" + + +def test_the_rust_contract_matches_the_returned_fields() -> None: + contract: Final = TypeAdapter(dict[str, list[str]]).validate_json(CONTRACT_PATH.read_text()) + + assert contract == { + "http_settings": [field.name for field in dataclasses.fields(settings.http_settings())], + "url_policy": [field.name for field in dataclasses.fields(settings.url_policy())], + } + + +def test_url_policy_reads_the_litellm_globals(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "user_url_validation", False) + monkeypatch.setattr(litellm, "user_url_allowed_hosts", ["docs.internal:8443"]) + + assert settings.url_policy() == settings.UrlPolicy( + user_url_validation=False, + user_url_allowed_hosts=["docs.internal:8443"], + ) + + +def test_http_settings_reads_the_litellm_globals(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "ssl_verify", "/etc/ssl/corp.pem") + monkeypatch.setattr(litellm, "ssl_certificate", "/etc/ssl/client.pem") + monkeypatch.setattr(litellm, "ssl_security_level", "DEFAULT@SECLEVEL=1") + monkeypatch.setattr(litellm, "ssl_ecdh_curve", "X25519") + monkeypatch.setattr(litellm, "force_ipv4", True) + monkeypatch.setattr(litellm, "http2", True) + monkeypatch.setattr(litellm, "aiohttp_trust_env", True) + monkeypatch.setattr(litellm, "disable_aiohttp_trust_env", True) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + + assert settings.http_settings() == settings.HttpSettings( + ssl_verify="/etc/ssl/corp.pem", + ssl_certificate="/etc/ssl/client.pem", + ssl_security_level="DEFAULT@SECLEVEL=1", + ssl_ecdh_curve="X25519", + force_ipv4=True, + http2=True, + aiohttp_trust_env=True, + disable_aiohttp_trust_env=True, + disable_aiohttp_transport=True, + user_agent=default_user_agent(), + ) + + +def test_http_settings_ignores_environment_overrides(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("LITELLM_USER_AGENT", "operator/1") + monkeypatch.setenv("SSL_VERIFY", "false") + monkeypatch.setattr(litellm, "ssl_verify", True) + + result: Final = settings.http_settings() + + assert result.user_agent == default_user_agent() + assert result.ssl_verify is True + + +def test_warn_reaches_the_litellm_logger(caplog: pytest.LogCaptureFixture) -> None: + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + settings.warn("ssl_ecdh_curve 'secp521r1' is not supported") + + assert [record.getMessage() for record in caplog.records] == ["ssl_ecdh_curve 'secp521r1' is not supported"] diff --git a/tests/test_litellm/test_model_prices_schema.py b/tests/test_litellm/test_model_prices_schema.py index 2f9b11a16b7..a9015397b32 100644 --- a/tests/test_litellm/test_model_prices_schema.py +++ b/tests/test_litellm/test_model_prices_schema.py @@ -125,6 +125,55 @@ def test_schema_accepts_cache_creation_cost_inside_a_pricing_tier(committed_sche assert validator.is_valid({"some-model": entry}) +OFF_PEAK_ENTRY: Final = MappingProxyType( + { + "litellm_provider": "openrouter", + "mode": "chat", + "input_cost_per_token": 2e-6, + "output_cost_per_token": 8e-6, + "off_peak_pricing": { + "hours_utc": "16:30-00:30", + "windows": [{"hours_utc": ["00:30-02:00"], "weekdays": [6, "Sunday", "mon", "THURS"]}], + "weekday_timezone": "Asia/Shanghai", + "input_cost_per_token": 1e-6, + "output_cost_per_token": 4e-6, + "cache_read_input_token_cost": 1e-7, + }, + } +) + + +def test_generator_classifies_off_peak_pricing_as_a_windowed_rate_block(): + generator = load_generator() + schema = json.loads(generator.render(generator.build_schema({"some-model": dict(OFF_PEAK_ENTRY)}))) + validator = build_validator(schema) + assert validator.is_valid({"some-model": dict(OFF_PEAK_ENTRY)}) + + +@pytest.mark.parametrize( + "block", + [ + {"hours_utc": "16:30-00:30", "input_cost_per_token": "1e-6"}, + {"hours_utc": "16:30-00:30", "input_cost_per_token": -1e-6}, + {"hours_utc": 1630, "input_cost_per_token": 1e-6}, + {"hours_utc": "16:30-00:30", "discount": 0.5}, + {"windows": [{"weekdays": [6]}], "input_cost_per_token": 1e-6}, + {"windows": [{"hours_utc": "00:30-02:00", "weekdays": [0]}], "input_cost_per_token": 1e-6}, + {"windows": [], "input_cost_per_token": 1e-6}, + {"input_cost_per_token": 1e-6}, + {"hours_utc": "16:30", "input_cost_per_token": 1e-6}, + {"hours_utc": "25:00-01:00", "input_cost_per_token": 1e-6}, + {"hours_utc": ["16:30-00:30", "4pm-midnight"], "input_cost_per_token": 1e-6}, + {"windows": [{"hours_utc": "00:30-02:00", "weekdays": ["Funday"]}], "input_cost_per_token": 1e-6}, + ], +) +def test_generated_off_peak_schema_rejects_malformed_blocks(block: dict): + generator = load_generator() + schema = json.loads(generator.render(generator.build_schema({"some-model": dict(OFF_PEAK_ENTRY)}))) + validator = build_validator(schema) + assert not validator.is_valid({"some-model": {**OFF_PEAK_ENTRY, "off_peak_pricing": block}}) + + def find_duplicate_keys(path: Path) -> list[str]: duplicates: list[str] = [] diff --git a/tests/test_litellm/test_router_retry_non_retryable_errors.py b/tests/test_litellm/test_router_retry_non_retryable_errors.py index 0728947eafe..98a7db7a079 100644 --- a/tests/test_litellm/test_router_retry_non_retryable_errors.py +++ b/tests/test_litellm/test_router_retry_non_retryable_errors.py @@ -10,12 +10,24 @@ Verifies that: Regression tests for https://github.com/BerriAI/litellm/issues/21343 """ +import asyncio +import datetime +from collections.abc import Awaitable, Callable +from typing import Final from unittest.mock import AsyncMock, patch import pytest import litellm from litellm import Router +from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.litellm_core_utils.llm_response_utils.response_metadata import ( + _union_duration_ms, + response_timing_metrics, +) +from litellm.litellm_core_utils.logging_utils import track_llm_api_timing +from litellm.litellm_core_utils.rules import Rules +from litellm.utils import function_setup def _make_rate_limit_error(message="Rate limited"): @@ -274,3 +286,74 @@ async def test_not_found_error_in_retry_loop_raises_immediately(): # Only 2 calls: initial + first retry that hits non-retryable assert call_count == 2 + + +@pytest.mark.asyncio +async def test_retry_attempts_accumulate_timing_in_shared_request_metadata(): + received_at: Final = datetime.datetime.now() + metadata: dict[str, object] = { + "model_group": "test-model", + "litellm_received_at": received_at, + } + logging_obj_raw, _ = function_setup( + "acompletion", + Rules(), + datetime.datetime.now(), + model="test-model", + messages=[{"role": "user", "content": "test"}], + metadata=metadata, + litellm_call_id="retry-timing-test", + is_async_call=True, + ) + assert isinstance(logging_obj_raw, Logging) + logging_obj: Final[Logging] = logging_obj_raw + attempt_numbers: list[int] = [] + metadata_ids: list[int] = [] + + @track_llm_api_timing() + async def timed_attempt(*, logging_obj: Logging, **kwargs: object) -> str: + del kwargs + attempt_numbers.append(len(attempt_numbers) + 1) + metadata_ids.append(id(logging_obj.model_call_details["litellm_params"]["metadata"])) + await asyncio.sleep(0.01) + if len(attempt_numbers) == 1: + raise _make_rate_limit_error() + return "success" + + async def invoke(original_function: Callable[..., Awaitable[str]], *args: object, **kwargs: object) -> str: + return await original_function(*args, **kwargs) + + router = _create_router(num_retries=1) + with ( + patch.object(router, "make_call", new=AsyncMock(side_effect=invoke)), + patch.object( + router, + "_async_get_healthy_deployments", + new=AsyncMock(return_value=(["d1"], ["d1"])), + ), + patch.object(router, "_time_to_sleep_before_retry", return_value=0), + ): + result = await router.async_function_with_retries( + original_function=timed_attempt, + model="test-model", + messages=[{"role": "user", "content": "test"}], + metadata=metadata, + logging_obj=logging_obj, + num_retries=1, + ) + + request_metadata: Final = logging_obj.model_call_details["litellm_params"]["metadata"] + windows: Final = request_metadata["llm_api_timing_windows"] + end_time: Final = datetime.datetime.fromtimestamp(max(window[1] for window in windows)) + timing_metrics: Final = response_timing_metrics(received_at, end_time, logging_obj) + assert result == "success" + assert attempt_numbers == [1, 2] + assert request_metadata is metadata + assert metadata_ids == [id(metadata), id(metadata)] + assert len(windows) == 2 + union_duration_ms: Final = _union_duration_ms(windows, received_at.timestamp(), end_time.timestamp()) + assert union_duration_ms is not None + total_response_time_ms: Final = (end_time.timestamp() - received_at.timestamp()) * 1000 + assert timing_metrics["litellm_overhead_time_ms"] == pytest.approx( + round(total_response_time_ms - union_duration_ms, 4) + ) diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index 5adfb2aea4c..2fda5dfc490 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -1060,7 +1060,7 @@ def test_max_tokens_consistency(): if len(inconsistencies) > 10: error_msg += f"\n ... and {len(inconsistencies) - 10} more\n" - error_msg += "\nTo fix these inconsistencies, run: poetry run python fix_max_tokens_inconsistencies.py" + error_msg += "\nTo fix these inconsistencies, run: uv run python fix_max_tokens_inconsistencies.py" raise AssertionError(error_msg) diff --git a/tests/test_litellm_rust/ocr/test_lifecycle.py b/tests/test_litellm_rust/ocr/test_lifecycle.py index 5fca927bea3..264a666c685 100644 --- a/tests/test_litellm_rust/ocr/test_lifecycle.py +++ b/tests/test_litellm_rust/ocr/test_lifecycle.py @@ -39,6 +39,7 @@ async def test_proxy_metadata_remains_python_owned(ocr_server: RecordingServer) ) events: Final = await recorder.wait_for_async("async_log_success_event") assert response.pages[0].markdown == "native OCR response" + assert response._hidden_params["additional_headers"]["x-litellm-rust"] == "true" assert events[0].kwargs["litellm_params"]["metadata"]["user_api_key_auth"].user_id == "ocr-user" assert "metadata" not in ocr_server.requests[0].body diff --git a/ui/litellm-dashboard/CLAUDE.md b/ui/litellm-dashboard/AGENTS.md similarity index 100% rename from ui/litellm-dashboard/CLAUDE.md rename to ui/litellm-dashboard/AGENTS.md diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPNetworkSettings.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPNetworkSettings.test.tsx index 9526f5de074..4a486bfc648 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPNetworkSettings.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPNetworkSettings.test.tsx @@ -1,7 +1,8 @@ -import { render, screen, waitFor } from "@testing-library/react"; +import { fireEvent, render, screen, waitFor, within } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import { describe, it, expect, vi, beforeEach } from "vitest"; import MCPNetworkSettings from "./MCPNetworkSettings"; +import { toast } from "@/lib/toast"; import { getGeneralSettingsCall, updateConfigFieldSetting, @@ -16,8 +17,31 @@ vi.mock("@/components/networking", () => ({ fetchMCPClientIp: vi.fn(), })); +vi.mock("@/lib/toast", () => ({ + toast: { success: vi.fn(), fromError: vi.fn() }, +})); + const renderSettings = () => render(); +const ANTIGRAVITY = { alias: "Antigravity CLI", value: "antigravity-cli" }; +const CODEX = { alias: "Codex", value: "codex-mcp-client" }; + +const clientCard = (alias: string) => screen.getByRole("button", { name: new RegExp(`^${alias}`) }); + +const fillClientDialog = async (alias: string, value: string) => { + const dialog = await screen.findByRole("dialog"); + fireEvent.change(within(dialog).getByRole("textbox", { name: "Alias" }), { target: { value: alias } }); + fireEvent.change(within(dialog).getByRole("textbox", { name: "Value" }), { target: { value } }); + return dialog; +}; + +const addClient = async (alias: string, value: string) => { + await userEvent.click(screen.getByRole("button", { name: "Add client" })); + const dialog = await fillClientDialog(alias, value); + await userEvent.click(within(dialog).getByRole("button", { name: "Add" })); + await waitFor(() => expect(screen.queryByRole("dialog")).not.toBeInTheDocument()); +}; + describe("MCPNetworkSettings", () => { beforeEach(() => { vi.clearAllMocks(); @@ -82,25 +106,399 @@ describe("MCPNetworkSettings", () => { expect(screen.getByText("203.0.113.0/24")).toBeInTheDocument(); }); - it("saves the configured ranges", async () => { + it("saves the configured ranges once they change", async () => { + vi.mocked(fetchMCPClientIp).mockResolvedValue("203.0.113.45"); vi.mocked(getGeneralSettingsCall).mockResolvedValue([ { field_name: "mcp_internal_ip_ranges", field_value: ["10.0.0.0/8"] }, ]); renderSettings(); + await userEvent.click(await screen.findByText("203.0.113.0/24")); await userEvent.click(await screen.findByRole("button", { name: /Save/ })); await waitFor(() => - expect(updateConfigFieldSetting).toHaveBeenCalledWith("tok", "mcp_internal_ip_ranges", ["10.0.0.0/8"]), + expect(updateConfigFieldSetting).toHaveBeenCalledWith("tok", "mcp_internal_ip_ranges", [ + "10.0.0.0/8", + "203.0.113.0/24", + ]), ); - expect(deleteConfigFieldSetting).not.toHaveBeenCalled(); + expect(deleteConfigFieldSetting).not.toHaveBeenCalledWith("tok", "mcp_internal_ip_ranges"); }); - it("clears the setting instead of saving an empty list", async () => { + it("clears a stored range setting instead of saving an empty list", async () => { + vi.mocked(getGeneralSettingsCall).mockResolvedValue([ + { field_name: "mcp_internal_ip_ranges", field_value: ["10.0.0.0/8"] }, + ]); + renderSettings(); - await userEvent.click(await screen.findByRole("button", { name: /Save/ })); + await userEvent.click(await screen.findByRole("button", { name: "Remove 10.0.0.0/8" })); + await userEvent.click(screen.getByRole("button", { name: /Save/ })); await waitFor(() => expect(deleteConfigFieldSetting).toHaveBeenCalledWith("tok", "mcp_internal_ip_ranges")); expect(updateConfigFieldSetting).not.toHaveBeenCalled(); }); + + it("does not write settings that were never stored and are still empty", async () => { + renderSettings(); + await userEvent.click(await screen.findByRole("button", { name: /Save/ })); + + await waitFor(() => expect(toast.success).toHaveBeenCalledWith("MCP network settings saved")); + expect(deleteConfigFieldSetting).not.toHaveBeenCalled(); + expect(updateConfigFieldSetting).not.toHaveBeenCalled(); + }); + + it("labels the section Allowed Clients and renders each stored client as a card showing alias and value", async () => { + vi.mocked(getGeneralSettingsCall).mockResolvedValue([ + { field_name: "mcp_allowed_clients", field_value: [ANTIGRAVITY, CODEX] }, + ]); + + renderSettings(); + + expect(await screen.findByText("Allowed Clients")).toBeVisible(); + expect(screen.queryByText(/Allowed Client IDs/)).not.toBeInTheDocument(); + expect(screen.queryByText(/Allowed Client Applications/)).not.toBeInTheDocument(); + expect(clientCard("Antigravity CLI")).toHaveTextContent("antigravity-cli"); + expect(clientCard("Codex")).toHaveTextContent("codex-mcp-client"); + expect(screen.queryByRole("dialog")).not.toBeInTheDocument(); + }); + + it("opens an edit dialog when a client card is clicked, prefilled with that client's alias and value", async () => { + vi.mocked(getGeneralSettingsCall).mockResolvedValue([ + { field_name: "mcp_allowed_clients", field_value: [ANTIGRAVITY, CODEX] }, + ]); + + renderSettings(); + await screen.findByText("Allowed Clients"); + await userEvent.click(clientCard("Codex")); + + const dialog = await screen.findByRole("dialog", { name: "Edit client" }); + expect(within(dialog).getByRole("textbox", { name: "Alias" })).toHaveValue("Codex"); + expect(within(dialog).getByRole("textbox", { name: "Value" })).toHaveValue("codex-mcp-client"); + }); + + it("warns that a stored allowlist in the old plain-string shape denies every client and lets Save remove it", async () => { + vi.mocked(getGeneralSettingsCall).mockResolvedValue([ + { field_name: "mcp_allowed_clients", field_value: ["antigravity-cli"] }, + ]); + + renderSettings(); + + expect(await screen.findByText(/stored allowlist is not a list of alias and value pairs/)).toBeVisible(); + expect(screen.queryByText("antigravity-cli")).not.toBeInTheDocument(); + + await userEvent.click(screen.getByRole("button", { name: /Save/ })); + + await waitFor(() => expect(deleteConfigFieldSetting).toHaveBeenCalledWith("tok", "mcp_allowed_clients")); + expect(updateConfigFieldSetting).not.toHaveBeenCalled(); + await waitFor(() => expect(screen.queryByText(/stored allowlist is not a list/)).not.toBeInTheDocument()); + }); + + it("replaces a stored allowlist in the old plain-string shape with the clients the admin adds", async () => { + vi.mocked(getGeneralSettingsCall).mockResolvedValue([ + { field_name: "mcp_allowed_clients", field_value: ["antigravity-cli"] }, + ]); + + renderSettings(); + + await screen.findByText(/stored allowlist is not a list of alias and value pairs/); + await addClient(ANTIGRAVITY.alias, ANTIGRAVITY.value); + await userEvent.click(screen.getByRole("button", { name: /Save/ })); + + await waitFor(() => + expect(updateConfigFieldSetting).toHaveBeenCalledWith("tok", "mcp_allowed_clients", [ANTIGRAVITY]), + ); + expect(deleteConfigFieldSetting).not.toHaveBeenCalledWith("tok", "mcp_allowed_clients"); + await waitFor(() => expect(screen.queryByText(/stored allowlist is not a list/)).not.toBeInTheDocument()); + }); + + it("treats a stored entry with an empty alias or value as denying every client, like the gateway does", async () => { + vi.mocked(getGeneralSettingsCall).mockResolvedValue([ + { field_name: "mcp_allowed_clients", field_value: [ANTIGRAVITY, { alias: "", value: "claude-code" }] }, + ]); + + renderSettings(); + + expect(await screen.findByText(/stored allowlist is not a list of alias and value pairs/)).toBeVisible(); + expect(screen.queryByRole("button", { name: /^Antigravity CLI/ })).not.toBeInTheDocument(); + }); + + it("adds clients as alias and value pairs and saves them under mcp_allowed_clients", async () => { + renderSettings(); + await screen.findByText("Allowed Clients"); + + await addClient(" Antigravity CLI ", " antigravity-cli "); + await addClient("Codex", "codex-mcp-client"); + await userEvent.click(screen.getByRole("button", { name: /Save/ })); + + await waitFor(() => + expect(updateConfigFieldSetting).toHaveBeenCalledWith("tok", "mcp_allowed_clients", [ANTIGRAVITY, CODEX]), + ); + expect(deleteConfigFieldSetting).not.toHaveBeenCalledWith("tok", "mcp_allowed_clients"); + }); + + it("edits a stored client's value through its dialog and saves the new value", async () => { + vi.mocked(getGeneralSettingsCall).mockResolvedValue([ + { field_name: "mcp_allowed_clients", field_value: [ANTIGRAVITY] }, + ]); + + renderSettings(); + await screen.findByText("Allowed Clients"); + await userEvent.click(clientCard("Antigravity CLI")); + const dialog = await screen.findByRole("dialog"); + fireEvent.change(within(dialog).getByRole("textbox", { name: "Value" }), { + target: { value: " 0oa1b2c3d4e5f6g7h8i9 " }, + }); + await userEvent.click(within(dialog).getByRole("button", { name: "Done" })); + + await waitFor(() => expect(screen.queryByRole("dialog")).not.toBeInTheDocument()); + expect(clientCard("Antigravity CLI")).toHaveTextContent("0oa1b2c3d4e5f6g7h8i9"); + await userEvent.click(screen.getByRole("button", { name: /Save/ })); + + await waitFor(() => + expect(updateConfigFieldSetting).toHaveBeenCalledWith("tok", "mcp_allowed_clients", [ + { alias: "Antigravity CLI", value: "0oa1b2c3d4e5f6g7h8i9" }, + ]), + ); + }); + + it("keeps a stored client untouched when its dialog is cancelled", async () => { + vi.mocked(getGeneralSettingsCall).mockResolvedValue([ + { field_name: "mcp_allowed_clients", field_value: [ANTIGRAVITY] }, + ]); + + renderSettings(); + await screen.findByText("Allowed Clients"); + await userEvent.click(clientCard("Antigravity CLI")); + const dialog = await screen.findByRole("dialog"); + fireEvent.change(within(dialog).getByRole("textbox", { name: "Value" }), { target: { value: "changed" } }); + await userEvent.click(within(dialog).getByRole("button", { name: "Cancel" })); + + await waitFor(() => expect(screen.queryByRole("dialog")).not.toBeInTheDocument()); + expect(clientCard("Antigravity CLI")).toHaveTextContent("antigravity-cli"); + await userEvent.click(screen.getByRole("button", { name: /Save/ })); + + await waitFor(() => expect(toast.success).toHaveBeenCalledWith("MCP network settings saved")); + expect(updateConfigFieldSetting).not.toHaveBeenCalled(); + }); + + it("will not add a client that has an alias but no value", async () => { + renderSettings(); + await screen.findByText("Allowed Clients"); + + await userEvent.click(screen.getByRole("button", { name: "Add client" })); + const dialog = await fillClientDialog("Antigravity CLI", " "); + + expect(within(dialog).getByRole("button", { name: "Add" })).toBeDisabled(); + }); + + it("adds nothing when the add dialog is cancelled", async () => { + vi.mocked(getGeneralSettingsCall).mockResolvedValue([ + { field_name: "mcp_allowed_clients", field_value: [ANTIGRAVITY] }, + ]); + + renderSettings(); + await screen.findByText("Allowed Clients"); + await userEvent.click(screen.getByRole("button", { name: "Add client" })); + const dialog = await fillClientDialog("Codex", "codex-mcp-client"); + await userEvent.click(within(dialog).getByRole("button", { name: "Cancel" })); + + await waitFor(() => expect(screen.queryByRole("dialog")).not.toBeInTheDocument()); + expect(screen.queryByText("Codex")).not.toBeInTheDocument(); + await userEvent.click(screen.getByRole("button", { name: /Save/ })); + + await waitFor(() => expect(toast.success).toHaveBeenCalledWith("MCP network settings saved")); + expect(updateConfigFieldSetting).not.toHaveBeenCalled(); + expect(deleteConfigFieldSetting).not.toHaveBeenCalled(); + }); + + it("removes the right client from the middle of the list", async () => { + vi.mocked(getGeneralSettingsCall).mockResolvedValue([ + { + field_name: "mcp_allowed_clients", + field_value: [ANTIGRAVITY, { alias: "Claude Code", value: "claude-code" }, CODEX], + }, + ]); + + renderSettings(); + await screen.findByText("Allowed Clients"); + await userEvent.click(clientCard("Claude Code")); + await userEvent.click(within(await screen.findByRole("dialog")).getByRole("button", { name: "Remove client" })); + + await waitFor(() => expect(screen.queryByRole("dialog")).not.toBeInTheDocument()); + expect(screen.queryByText("claude-code")).not.toBeInTheDocument(); + await userEvent.click(screen.getByRole("button", { name: /Save/ })); + + await waitFor(() => + expect(updateConfigFieldSetting).toHaveBeenCalledWith("tok", "mcp_allowed_clients", [ANTIGRAVITY, CODEX]), + ); + }); + + it("removes a client and clears the setting when the list becomes empty", async () => { + vi.mocked(getGeneralSettingsCall).mockResolvedValue([ + { field_name: "mcp_allowed_clients", field_value: [{ alias: "Claude Code", value: "claude-code" }] }, + ]); + + renderSettings(); + await screen.findByText("Allowed Clients"); + await userEvent.click(clientCard("Claude Code")); + await userEvent.click(within(await screen.findByRole("dialog")).getByRole("button", { name: "Remove client" })); + + await waitFor(() => expect(screen.queryByRole("dialog")).not.toBeInTheDocument()); + expect(screen.queryByText("claude-code")).not.toBeInTheDocument(); + + await userEvent.click(screen.getByRole("button", { name: /Save/ })); + + await waitFor(() => expect(deleteConfigFieldSetting).toHaveBeenCalledWith("tok", "mcp_allowed_clients")); + expect(updateConfigFieldSetting).not.toHaveBeenCalledWith("tok", "mcp_allowed_clients", expect.anything()); + }); + + it("warns that a stored empty allowlist denies every client and lets Save remove it", async () => { + vi.mocked(getGeneralSettingsCall).mockResolvedValue([{ field_name: "mcp_allowed_clients", field_value: [] }]); + + renderSettings(); + + expect(await screen.findByText(/An empty allowlist is currently stored, so every client is denied/)).toBeVisible(); + + await userEvent.click(screen.getByRole("button", { name: /Save/ })); + + await waitFor(() => expect(deleteConfigFieldSetting).toHaveBeenCalledWith("tok", "mcp_allowed_clients")); + expect(updateConfigFieldSetting).not.toHaveBeenCalled(); + await waitFor(() => expect(screen.queryByText(/An empty allowlist is currently stored/)).not.toBeInTheDocument()); + }); + + it("does not show the deny-all warning when no allowlist is stored", async () => { + vi.mocked(getGeneralSettingsCall).mockResolvedValue([{ field_name: "mcp_allowed_clients", field_value: null }]); + + renderSettings(); + + await screen.findByText("Allowed Clients"); + expect(screen.queryByText(/every client is denied/)).not.toBeInTheDocument(); + }); + + it("explains that JWT callers are identified by the configured claim and others by the opt-in header", async () => { + renderSettings(); + + expect(await screen.findByText(/litellm_jwtauth\.mcp_client_id_jwt_field/)).toBeVisible(); + expect(screen.getByText(/Clients pick this value themselves, so it is a policy control/)).toBeVisible(); + expect(screen.queryByText(/clientInfo/)).not.toBeInTheDocument(); + }); + + it("renders the stored client identity header once settings load", async () => { + vi.mocked(getGeneralSettingsCall).mockResolvedValue([ + { field_name: "mcp_client_id_header", field_value: "x-mcp-client" }, + ]); + + renderSettings(); + + expect(await screen.findByRole("textbox", { name: "Client identity header" })).toHaveValue("x-mcp-client"); + }); + + it("saves a newly typed client identity header under mcp_client_id_header", async () => { + renderSettings(); + fireEvent.change(await screen.findByRole("textbox", { name: "Client identity header" }), { + target: { value: " x-mcp-client " }, + }); + + await userEvent.click(screen.getByRole("button", { name: /Save/ })); + + await waitFor(() => + expect(updateConfigFieldSetting).toHaveBeenCalledWith("tok", "mcp_client_id_header", "x-mcp-client"), + ); + expect(updateConfigFieldSetting).not.toHaveBeenCalledWith("tok", "mcp_allowed_clients", expect.anything()); + expect(deleteConfigFieldSetting).not.toHaveBeenCalled(); + }); + + it("clears a stored client identity header when the field is emptied, so only JWT identity is trusted", async () => { + vi.mocked(getGeneralSettingsCall).mockResolvedValue([ + { field_name: "mcp_client_id_header", field_value: "x-mcp-client" }, + ]); + + renderSettings(); + fireEvent.change(await screen.findByRole("textbox", { name: "Client identity header" }), { + target: { value: "" }, + }); + + await userEvent.click(screen.getByRole("button", { name: /Save/ })); + + await waitFor(() => expect(deleteConfigFieldSetting).toHaveBeenCalledWith("tok", "mcp_client_id_header")); + expect(updateConfigFieldSetting).not.toHaveBeenCalled(); + }); + + it("does not rewrite an unchanged client identity header on save", async () => { + vi.mocked(getGeneralSettingsCall).mockResolvedValue([ + { field_name: "mcp_client_id_header", field_value: "x-mcp-client" }, + ]); + + renderSettings(); + await screen.findByRole("textbox", { name: "Client identity header" }); + await userEvent.click(screen.getByRole("button", { name: /Save/ })); + + await waitFor(() => expect(toast.success).toHaveBeenCalled()); + expect(updateConfigFieldSetting).not.toHaveBeenCalled(); + expect(deleteConfigFieldSetting).not.toHaveBeenCalled(); + }); + + it("keeps the private ranges and the allowed clients as independent settings on save", async () => { + vi.mocked(getGeneralSettingsCall).mockResolvedValue([ + { field_name: "mcp_internal_ip_ranges", field_value: ["10.0.0.0/8"] }, + { field_name: "mcp_allowed_clients", field_value: [ANTIGRAVITY] }, + ]); + + renderSettings(); + await screen.findByText("Allowed Clients"); + await addClient("Codex", "codex-mcp-client"); + await userEvent.click(screen.getByRole("button", { name: /Save/ })); + + await waitFor(() => + expect(updateConfigFieldSetting).toHaveBeenCalledWith("tok", "mcp_allowed_clients", [ANTIGRAVITY, CODEX]), + ); + expect(updateConfigFieldSetting).not.toHaveBeenCalledWith("tok", "mcp_internal_ip_ranges", expect.anything()); + expect(deleteConfigFieldSetting).not.toHaveBeenCalled(); + }); + + it("still saves the allowed clients when the private range write fails, and reports the failure", async () => { + vi.mocked(getGeneralSettingsCall).mockResolvedValue([ + { field_name: "mcp_internal_ip_ranges", field_value: ["10.0.0.0/8"] }, + ]); + const rangeFailure = new Error("Field name=mcp_internal_ip_ranges not in config"); + vi.mocked(deleteConfigFieldSetting).mockRejectedValue(rangeFailure); + + renderSettings(); + await userEvent.click(await screen.findByRole("button", { name: "Remove 10.0.0.0/8" })); + await addClient("Codex", "codex-mcp-client"); + await userEvent.click(screen.getByRole("button", { name: /Save/ })); + + await waitFor(() => expect(updateConfigFieldSetting).toHaveBeenCalledWith("tok", "mcp_allowed_clients", [CODEX])); + await waitFor(() => expect(toast.fromError).toHaveBeenCalledWith(rangeFailure)); + expect(toast.success).not.toHaveBeenCalled(); + }); + + it("writes the private ranges and the allowed clients one after the other, never concurrently", async () => { + vi.mocked(fetchMCPClientIp).mockResolvedValue("203.0.113.45"); + let finishRangeWrite: (() => void) | undefined; + vi.mocked(updateConfigFieldSetting).mockImplementation( + (_token, fieldName) => + new Promise((resolve) => { + if (fieldName === "mcp_internal_ip_ranges") { + finishRangeWrite = resolve; + } else { + resolve(); + } + }), + ); + + renderSettings(); + await userEvent.click(await screen.findByText("203.0.113.0/24")); + await addClient("Codex", "codex-mcp-client"); + await userEvent.click(screen.getByRole("button", { name: /Save/ })); + + await waitFor(() => + expect(updateConfigFieldSetting).toHaveBeenCalledWith("tok", "mcp_internal_ip_ranges", ["203.0.113.0/24"]), + ); + expect(updateConfigFieldSetting).not.toHaveBeenCalledWith("tok", "mcp_allowed_clients", expect.anything()); + + finishRangeWrite?.(); + await waitFor(() => expect(updateConfigFieldSetting).toHaveBeenCalledWith("tok", "mcp_allowed_clients", [CODEX])); + await waitFor(() => expect(toast.success).toHaveBeenCalledWith("MCP network settings saved")); + }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPNetworkSettings.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPNetworkSettings.tsx index 8b4d2a58652..2fd62c7f1ef 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPNetworkSettings.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPNetworkSettings.tsx @@ -1,11 +1,21 @@ -import React, { useState, useEffect } from "react"; +import React, { useState, useEffect, useId } from "react"; import { Save, Plus, X } from "lucide-react"; import { Badge } from "@/components/ui/badge"; import { Button } from "@/components/ui/button"; import { Card } from "@/components/ui/card"; +import { + Dialog, + DialogContent, + DialogDescription, + DialogFooter, + DialogHeader, + DialogTitle, +} from "@/components/ui/dialog"; import { Input } from "@/components/ui/input"; +import { Label } from "@/components/ui/label"; import { UiLoadingSpinner } from "@/components/ui/ui-loading-spinner"; import { DeprecationBanner } from "@/components/DeprecationBanner"; +import { toast } from "@/lib/toast"; import { getGeneralSettingsCall, updateConfigFieldSetting, @@ -26,12 +36,142 @@ function ipToSlash24(ip: string): string { return `${parts[0]}.${parts[1]}.${parts[2]}.0/24`; } +export interface AllowedClient { + readonly alias: string; + readonly value: string; +} + +interface AllowedClientRow extends AllowedClient { + readonly key: string; +} + +interface ClientDraft extends AllowedClient { + readonly key: string | null; +} + +const isAllowedClient = (entry: unknown): entry is AllowedClient => { + if (typeof entry !== "object" || entry === null) return false; + const { alias, value } = entry as Partial>; + return typeof alias === "string" && typeof value === "string" && !isIncomplete({ alias, value }); +}; + +type StoredAllowlist = + | { readonly kind: "absent" } + | { readonly kind: "clients"; readonly clients: AllowedClient[] } + | { readonly kind: "malformed" }; + +const ABSENT: StoredAllowlist = { kind: "absent" }; + +const parseStoredClients = (fieldValue: unknown): StoredAllowlist => { + if (fieldValue === null || fieldValue === undefined) return ABSENT; + if (Array.isArray(fieldValue) && fieldValue.every(isAllowedClient)) { + return { kind: "clients", clients: fieldValue.map(({ alias, value }) => ({ alias, value })) }; + } + return { kind: "malformed" }; +}; + +let nextRowKey = 0; +const newRow = (client: AllowedClient): AllowedClientRow => ({ ...client, key: `client-${nextRowKey++}` }); + +const trimClient = ({ alias, value }: AllowedClient): AllowedClient => ({ alias: alias.trim(), value: value.trim() }); + +const isIncomplete = ({ alias, value }: AllowedClient) => alias === "" || value === ""; + +const sameList = (a: string[], b: string[]) => a.length === b.length && a.every((value, i) => value === b[i]); + +const sameClients = (a: AllowedClient[], b: AllowedClient[]) => + a.length === b.length && a.every((client, i) => client.alias === b[i].alias && client.value === b[i].value); + +const unchangedSinceLoad = (value: string[], stored: string[] | null) => + stored === null ? value.length === 0 : value.length > 0 && sameList(value, stored); + +const clientsUnchangedSinceLoad = (value: AllowedClient[], stored: StoredAllowlist) => { + switch (stored.kind) { + case "absent": + return value.length === 0; + case "clients": + return value.length > 0 && sameClients(value, stored.clients); + case "malformed": + return false; + } +}; + +const headerUnchangedSinceLoad = (value: string, stored: string | null) => + stored === null ? value === "" : value !== "" && value === stored; + +interface AllowedClientDialogProps { + readonly draft: ClientDraft | null; + readonly onChange: (draft: ClientDraft) => void; + readonly onCommit: () => void; + readonly onRemove: () => void; + readonly onClose: () => void; +} + +const AllowedClientDialog: React.FC = ({ draft, onChange, onCommit, onRemove, onClose }) => { + const aliasId = useId(); + const valueId = useId(); + if (draft === null) return null; + return ( + !open && onClose()}> + + + {draft.key === null ? "Add client" : "Edit client"} + + The alias is the name shown in the dashboard and gateway logs. The value is the exact JWT claim or header + value that identifies the client, such as the OAuth client ID your identity provider issues. + + +
+
+ + onChange({ ...draft, alias: e.target.value })} + /> +
+
+ + onChange({ ...draft, value: e.target.value })} + /> +
+
+ + {draft.key !== null && ( + + )} + + + + + + ); +}; + const MCPNetworkSettings: React.FC = ({ accessToken }) => { const [loading, setLoading] = useState(true); const [saving, setSaving] = useState(false); const [privateRanges, setPrivateRanges] = useState([]); + const [allowedClients, setAllowedClients] = useState([]); + const [clientIdHeader, setClientIdHeader] = useState(""); + const [storedRanges, setStoredRanges] = useState(null); + const [storedClients, setStoredClients] = useState(ABSENT); + const [storedClientIdHeader, setStoredClientIdHeader] = useState(null); const [currentIp, setCurrentIp] = useState(null); const [rangeDraft, setRangeDraft] = useState(""); + const [clientDraft, setClientDraft] = useState(null); useEffect(() => { loadSettings(); @@ -44,8 +184,18 @@ const MCPNetworkSettings: React.FC = ({ accessToken }) try { const settings = await getGeneralSettingsCall(accessToken); for (const field of settings) { - if (field.field_name === "mcp_internal_ip_ranges" && field.field_value) { + if (field.field_name === "mcp_internal_ip_ranges" && Array.isArray(field.field_value)) { setPrivateRanges(field.field_value); + setStoredRanges(field.field_value); + } + if (field.field_name === "mcp_allowed_clients") { + const stored = parseStoredClients(field.field_value); + setAllowedClients(stored.kind === "clients" ? stored.clients.map(newRow) : []); + setStoredClients(stored); + } + if (field.field_name === "mcp_client_id_header" && typeof field.field_value === "string") { + setClientIdHeader(field.field_value); + setStoredClientIdHeader(field.field_value); } } } catch (error) { @@ -63,20 +213,56 @@ const MCPNetworkSettings: React.FC = ({ accessToken }) } }; + const persistRanges = async (token: string) => { + if (unchangedSinceLoad(privateRanges, storedRanges)) return; + if (privateRanges.length > 0) { + await updateConfigFieldSetting(token, "mcp_internal_ip_ranges", privateRanges); + setStoredRanges(privateRanges); + return; + } + await deleteConfigFieldSetting(token, "mcp_internal_ip_ranges"); + setStoredRanges(null); + }; + + const persistAllowedClients = async (token: string) => { + const clients = allowedClients.map(({ alias, value }) => ({ alias, value })); + if (clientsUnchangedSinceLoad(clients, storedClients)) return; + if (clients.length > 0) { + await updateConfigFieldSetting(token, "mcp_allowed_clients", clients); + setStoredClients({ kind: "clients", clients }); + return; + } + await deleteConfigFieldSetting(token, "mcp_allowed_clients"); + setStoredClients(ABSENT); + }; + + const persistClientIdHeader = async (token: string) => { + const value = clientIdHeader.trim(); + if (headerUnchangedSinceLoad(value, storedClientIdHeader)) return; + if (value !== "") { + await updateConfigFieldSetting(token, "mcp_client_id_header", value); + setStoredClientIdHeader(value); + return; + } + await deleteConfigFieldSetting(token, "mcp_client_id_header"); + setStoredClientIdHeader(null); + }; + const handleSave = async () => { if (!accessToken) return; setSaving(true); - try { - if (privateRanges.length > 0) { - await updateConfigFieldSetting(accessToken, "mcp_internal_ip_ranges", privateRanges); - } else { - await deleteConfigFieldSetting(accessToken, "mcp_internal_ip_ranges"); - } - } catch (error) { - console.error("Failed to save MCP network settings:", error); - } finally { - setSaving(false); + const [rangeResult] = await Promise.allSettled([persistRanges(accessToken)]); + const [clientResult] = await Promise.allSettled([persistAllowedClients(accessToken)]); + const [headerResult] = await Promise.allSettled([persistClientIdHeader(accessToken)]); + setSaving(false); + const failures = [rangeResult, clientResult, headerResult].filter( + (result): result is PromiseRejectedResult => result.status === "rejected", + ); + if (failures.length === 0) { + toast.success("MCP network settings saved"); + return; } + failures.forEach((failure) => toast.fromError(failure.reason)); }; const addSuggestedRange = (range: string) => { @@ -86,17 +272,37 @@ const MCPNetworkSettings: React.FC = ({ accessToken }) }; // Commas separate entries, matching the old tokenised input. - const commitDraft = () => { - const added = rangeDraft + const splitDraft = (draft: string, existing: string[]) => + draft .split(",") .map((r) => r.trim()) - .filter((r) => r !== "" && !privateRanges.includes(r)); + .filter((r) => r !== "" && !existing.includes(r)); + + const commitDraft = () => { + const added = splitDraft(rangeDraft, privateRanges); if (added.length > 0) { setPrivateRanges([...privateRanges, ...added]); } setRangeDraft(""); }; + const commitClientDraft = () => { + if (clientDraft === null) return; + const client = trimClient(clientDraft); + setAllowedClients( + clientDraft.key === null + ? [...allowedClients, newRow(client)] + : allowedClients.map((row) => (row.key === clientDraft.key ? { ...row, ...client } : row)), + ); + setClientDraft(null); + }; + + const removeDraftedClient = () => { + if (clientDraft === null) return; + setAllowedClients(allowedClients.filter((row) => row.key !== clientDraft.key)); + setClientDraft(null); + }; + if (loading) { return (
@@ -106,6 +312,8 @@ const MCPNetworkSettings: React.FC = ({ accessToken }) } const suggestedRange = currentIp ? ipToSlash24(currentIp) : null; + const storedAllowlistIsMalformed = storedClients.kind === "malformed"; + const storedAllowlistIsEmpty = storedClients.kind === "clients" && storedClients.clients.length === 0; return (
@@ -178,12 +386,88 @@ const MCPNetworkSettings: React.FC = ({ accessToken })

+
+

Allowed Clients

+

+ Only the MCP client applications listed here can use the gateway. Leave empty to allow every client. A client + that authenticates with a JWT is identified by the claim named in litellm_jwtauth.mcp_client_id_jwt_field in + your proxy config (for example azp or client_id), which your identity provider asserts and the client cannot + change. Any other client is identified by the request header configured below, if you enable one. +

+
+ + + {storedAllowlistIsMalformed && ( +

+ The stored allowlist is not a list of alias and value pairs, so every client is denied. Add the clients you + want and save to replace it, or save with the list empty to remove it and allow every client again. +

+ )} + {storedAllowlistIsEmpty && ( +

+ An empty allowlist is currently stored, so every client is denied. Save with the list empty to remove it and + allow every client again. +

+ )} + {allowedClients.length > 0 && ( +
+ {allowedClients.map((row) => ( + + ))} +
+ )} + +

+ Click a client to edit or remove it. Leave the list empty to allow every client. Every MCP request from an + unlisted client, or from one with no resolvable identity, gets a 403. +

+ +
+

Client Identity Header (less secure)

+
+ setClientIdHeader(e.target.value)} + /> +

+ Optional header whose value names the client for callers without a JWT identity. Clients pick this value + themselves, so it is a policy control rather than a security boundary. Without it, callers that do not carry + the JWT claim are rejected while the allowlist is set. +

+ +
+ + setClientDraft(null)} + />
); }; diff --git a/ui/litellm-dashboard/src/components/team/teamAdminEditAccess.test.ts b/ui/litellm-dashboard/src/components/team/teamAdminEditAccess.test.ts index da3f9bf8289..1ef3816b5c4 100644 --- a/ui/litellm-dashboard/src/components/team/teamAdminEditAccess.test.ts +++ b/ui/litellm-dashboard/src/components/team/teamAdminEditAccess.test.ts @@ -13,6 +13,7 @@ describe("teamAdminFieldLabel", () => { ["tpm_limit", "Tokens per minute Limit (TPM)"], ["rpm_limit", "Requests per minute Limit (RPM)"], ["max_budget", "Max Budget (USD)"], + ["projects", "Create and update projects"], ])("names %s the way the team settings form does", (field, label) => { expect(teamAdminFieldLabel(field)).toBe(label); }); diff --git a/ui/litellm-dashboard/src/components/team/teamAdminEditAccess.ts b/ui/litellm-dashboard/src/components/team/teamAdminEditAccess.ts index b878af03df6..5706eeafe82 100644 --- a/ui/litellm-dashboard/src/components/team/teamAdminEditAccess.ts +++ b/ui/litellm-dashboard/src/components/team/teamAdminEditAccess.ts @@ -47,6 +47,7 @@ const TEAM_ADMIN_FIELD_LABELS: ReadonlyMap = new Map([ ["tpm_limit", "Tokens per minute Limit (TPM)"], ["rpm_limit", "Requests per minute Limit (RPM)"], ["max_budget", "Max Budget (USD)"], + ["projects", "Create and update projects"], ]); export const teamAdminFieldLabel = (field: string): string => TEAM_ADMIN_FIELD_LABELS.get(field) ?? field; diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 136b8f26784..9f71880c610 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -26927,6 +26927,16 @@ export interface components { * @description Maximum retention period for spend logs (e.g., '7d' for 7 days). Logs older than this will be deleted. */ maximum_spend_logs_retention_period?: string | null; + /** + * Mcp Allowed Clients + * @description MCP client applications admitted by the gateway, each an {alias, value} pair where alias is the name shown in the dashboard and logs and value is the identity that must match exactly. When set, every MCP request must carry a client identity equal to one of the values: a JWT caller is identified by the claim named in litellm_jwtauth.mcp_client_id_jwt_field, any other caller by the header named in mcp_client_id_header. A request with no resolvable identity, or an unlisted one, is rejected with 403. Unset means every client is admitted. + */ + mcp_allowed_clients?: components["schemas"]["MCPAllowedClient"][] | null; + /** + * Mcp Client Id Header + * @description Request header whose value names the calling MCP client application (for example 'x-mcp-client') for callers that did not authenticate with a JWT, used only while mcp_allowed_clients is set. The client picks this value itself, so it is a policy control rather than a security boundary; prefer litellm_jwtauth.mcp_client_id_jwt_field where callers use JWTs. + */ + mcp_client_id_header?: string | null; /** * Mcp Internal Ip Ranges * @description Custom CIDR ranges that define internal/private networks for MCP access control. When set, only these ranges are treated as internal. Defaults to RFC 1918 private ranges (10.0.0.0/8, 172.16.0.0/12, 192.168.0.0/16, 127.0.0.0/8). @@ -32563,6 +32573,22 @@ export interface components { */ status?: "healthy" | "unhealthy"; }; + /** + * MCPAllowedClient + * @description One entry of `general_settings.mcp_allowed_clients`. + */ + MCPAllowedClient: { + /** + * Alias + * @description Human-readable name for this client application, shown in the dashboard and in gateway logs. + */ + alias: string; + /** + * Value + * @description Exact value of the JWT claim named in litellm_jwtauth.mcp_client_id_jwt_field, or of the mcp_client_id_header header, that identifies this client application. Matched case-sensitively. + */ + value: string; + }; /** MCPConnectorEntry */ MCPConnectorEntry: { /** Args */ diff --git a/uv.lock b/uv.lock index f04fa5a17c1..ab9582ed134 100644 --- a/uv.lock +++ b/uv.lock @@ -2615,6 +2615,23 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/3d/f7/661d7a9023e877a226b5683429c3662f75a29ef45cb1464cf39adb689218/google_cloud_resource_manager-1.17.0-py3-none-any.whl", hash = "sha256:e479baf4b014a57f298e01b8279e3290b032e3476d69c8e5e1427af8f82739a5", size = 404403, upload-time = "2026-03-26T22:15:26.57Z" }, ] +[[package]] +name = "google-cloud-speech" +version = "2.40.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "google-api-core", version = "2.25.2", source = { registry = "https://pypi.org/simple" }, extra = ["grpc"], marker = "python_full_version >= '3.14'" }, + { name = "google-api-core", version = "2.30.3", source = { registry = "https://pypi.org/simple" }, extra = ["grpc"], marker = "python_full_version < '3.14'" }, + { name = "google-auth" }, + { name = "grpcio" }, + { name = "proto-plus" }, + { name = "protobuf" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/5a/c1/5dc9795314f4aefea0b01b02e9f5486a198341ecc15fe47f89a61c68df63/google_cloud_speech-2.40.0.tar.gz", hash = "sha256:e89e688e4ce0b926754038bf992d0d0f065c5f1c3503bb20e6c46d08b63658fc", size = 404366, upload-time = "2026-06-03T16:13:59.506Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/cc/78/afeca8d597fab54bdd823f857aad15d6f9c4628ff3cb72aa237d01700721/google_cloud_speech-2.40.0-py3-none-any.whl", hash = "sha256:7cc0302b3b9ca33d2eae9669da94a44316601a240942895362ac70e765b9f39c", size = 345427, upload-time = "2026-06-03T16:12:40.909Z" }, +] + [[package]] name = "google-cloud-storage" version = "3.4.1" @@ -4566,6 +4583,7 @@ proxy-runtime = [ { name = "ddtrace" }, { name = "detect-secrets" }, { name = "google-cloud-aiplatform" }, + { name = "google-cloud-speech" }, { name = "google-genai" }, { name = "grpcio" }, { name = "langfuse" }, @@ -4593,6 +4611,9 @@ stt-nvidia-riva = [ { name = "nvidia-riva-client" }, { name = "soundfile" }, ] +stt-vertex-chirp = [ + { name = "google-cloud-speech" }, +] utils = [ { name = "numpydoc" }, ] @@ -4607,6 +4628,7 @@ ci = [ { name = "blockbuster" }, { name = "claude-agent-sdk" }, { name = "detect-secrets" }, + { name = "google-cloud-speech" }, { name = "google-generativeai" }, { name = "jsonlines" }, { name = "langchain" }, @@ -4721,6 +4743,8 @@ requires-dist = [ { name = "google-cloud-aiplatform", marker = "extra == 'proxy-runtime'", specifier = ">=1.133.0,<2.0" }, { name = "google-cloud-iam", marker = "extra == 'extra-proxy'", specifier = ">=2.19.1,<3.0" }, { name = "google-cloud-kms", marker = "extra == 'extra-proxy'", specifier = ">=2.24.2,<3.0" }, + { name = "google-cloud-speech", marker = "extra == 'proxy-runtime'", specifier = ">=2.40.0,<3.0" }, + { name = "google-cloud-speech", marker = "extra == 'stt-vertex-chirp'", specifier = ">=2.40.0,<3.0" }, { name = "google-genai", marker = "extra == 'proxy-runtime'", specifier = ">=1.37.0,<2.0" }, { name = "granian", marker = "extra == 'proxy'", specifier = ">=2.7.4,<3.0" }, { name = "grpcio", marker = "extra == 'grpc'", specifier = "==1.78.0" }, @@ -4787,7 +4811,7 @@ requires-dist = [ { name = "uvloop", marker = "sys_platform != 'win32' and extra == 'proxy'", specifier = ">=0.22.1,<1.0" }, { name = "websockets", marker = "extra == 'proxy'", specifier = ">=15.0.1,<16.0" }, ] -provides-extras = ["proxy", "cli", "extra-proxy", "utils", "caching", "mcp", "saml", "semantic-router", "mlflow", "grpc", "stt-nvidia-riva", "google", "bedrock-realtime", "proxy-runtime"] +provides-extras = ["proxy", "cli", "extra-proxy", "utils", "caching", "mcp", "saml", "semantic-router", "mlflow", "grpc", "stt-vertex-chirp", "stt-nvidia-riva", "google", "bedrock-realtime", "proxy-runtime"] [package.metadata.requires-dev] ci = [ @@ -4799,6 +4823,7 @@ ci = [ { name = "blockbuster", specifier = "==1.5.26" }, { name = "claude-agent-sdk", specifier = "==0.1.44" }, { name = "detect-secrets", specifier = "==1.5.0" }, + { name = "google-cloud-speech", specifier = "==2.40.0" }, { name = "google-generativeai", specifier = "==0.8.6" }, { name = "jsonlines", specifier = "==4.0.0" }, { name = "langchain", specifier = "==1.3.9" },