mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-06 08:16:43 +00:00
Merge branch 'litellm_internal_staging' into feature/improve-gigachat-provider
This commit is contained in:
commit
5b969a3e5d
1003 changed files with 43908 additions and 12003 deletions
BIN
.github/deploy-on-aws.png
vendored
Normal file
BIN
.github/deploy-on-aws.png
vendored
Normal file
Binary file not shown.
|
After Width: | Height: | Size: 4 KiB |
BIN
.github/deploy-on-gcp.png
vendored
Normal file
BIN
.github/deploy-on-gcp.png
vendored
Normal file
Binary file not shown.
|
After Width: | Height: | Size: 4.8 KiB |
5
.github/workflows/osv-scan.yml
vendored
5
.github/workflows/osv-scan.yml
vendored
|
|
@ -7,11 +7,6 @@ on:
|
|||
- litellm_internal_staging
|
||||
- litellm_oss_branch
|
||||
- "litellm_**"
|
||||
paths:
|
||||
- uv.lock
|
||||
- ui/litellm-dashboard/package-lock.json
|
||||
- osv-scanner.toml
|
||||
- .github/workflows/osv-scan.yml
|
||||
schedule:
|
||||
- cron: "23 6 * * *"
|
||||
workflow_dispatch:
|
||||
|
|
|
|||
8
.github/workflows/test-linting.yml
vendored
8
.github/workflows/test-linting.yml
vendored
|
|
@ -14,7 +14,7 @@ permissions:
|
|||
jobs:
|
||||
lint:
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 10
|
||||
timeout-minutes: 15
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
|
||||
|
|
@ -87,9 +87,11 @@ jobs:
|
|||
run: |
|
||||
uv run --no-sync python -c "import openai; print(f'OpenAI version: {openai.__version__}')"
|
||||
|
||||
- name: Run basedpyright type checking
|
||||
- name: Check basedpyright budget (delta vs base)
|
||||
env:
|
||||
BASE_SHA: ${{ github.event.pull_request.base.sha }}
|
||||
run: |
|
||||
(uv run --no-sync basedpyright --outputjson || true) | uv run --no-sync python scripts/type_check_gate.py
|
||||
(uv run --no-sync basedpyright --outputjson || true) | uv run --no-sync python scripts/type_check_gate.py --base "$BASE_SHA"
|
||||
|
||||
- name: Check for circular imports
|
||||
run: |
|
||||
|
|
|
|||
65
.github/workflows/test-rust.yml
vendored
Normal file
65
.github/workflows/test-rust.yml
vendored
Normal file
|
|
@ -0,0 +1,65 @@
|
|||
name: LiteLLM Rust
|
||||
|
||||
on:
|
||||
push:
|
||||
paths:
|
||||
- "litellm-rust/**"
|
||||
- ".github/workflows/test-rust.yml"
|
||||
pull_request:
|
||||
branches:
|
||||
- main
|
||||
- litellm_internal_staging
|
||||
- litellm_oss_branch
|
||||
- "litellm_**"
|
||||
paths:
|
||||
- "litellm-rust/**"
|
||||
- ".github/workflows/test-rust.yml"
|
||||
|
||||
permissions:
|
||||
contents: read
|
||||
|
||||
concurrency:
|
||||
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
|
||||
cancel-in-progress: true
|
||||
|
||||
jobs:
|
||||
rust-checks:
|
||||
name: rustfmt, clippy, test
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 10
|
||||
defaults:
|
||||
run:
|
||||
working-directory: litellm-rust
|
||||
env:
|
||||
CARGO_TERM_COLOR: always
|
||||
|
||||
steps:
|
||||
- name: Checkout repository
|
||||
uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
|
||||
with:
|
||||
persist-credentials: false
|
||||
|
||||
- name: Set up Rust
|
||||
run: |
|
||||
rustup toolchain install stable --profile minimal --component clippy,rustfmt
|
||||
rustup default stable
|
||||
|
||||
- name: Cache Cargo registry and target
|
||||
uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0
|
||||
with:
|
||||
path: |
|
||||
~/.cargo/registry
|
||||
~/.cargo/git
|
||||
litellm-rust/target
|
||||
key: ${{ runner.os }}-cargo-${{ hashFiles('litellm-rust/Cargo.lock') }}
|
||||
restore-keys: |
|
||||
${{ runner.os }}-cargo-
|
||||
|
||||
- name: Check Rust formatting
|
||||
run: cargo fmt --check
|
||||
|
||||
- name: Run Clippy
|
||||
run: cargo clippy --workspace --all-targets --locked -- -D warnings
|
||||
|
||||
- name: Run Rust tests
|
||||
run: cargo test --workspace --locked
|
||||
2
.github/workflows/test-unit-misc.yml
vendored
2
.github/workflows/test-unit-misc.yml
vendored
|
|
@ -32,7 +32,9 @@ jobs:
|
|||
tests/test_litellm/repositories
|
||||
tests/test_litellm/images
|
||||
tests/test_litellm/interactions
|
||||
tests/test_litellm/ocr
|
||||
tests/test_litellm/passthrough
|
||||
tests/test_litellm/sandbox
|
||||
tests/test_litellm/vector_stores
|
||||
tests/test_litellm/test_*.py
|
||||
workers: 2
|
||||
|
|
|
|||
|
|
@ -29,7 +29,7 @@ If you ever make public-facing PR descriptions, comments, issues, commit message
|
|||
- don't use "—". Instead, reach for ";", ".", etc.
|
||||
- don't use the pattern "It's not X, it's Y", "You're not X, you're Y", etc.
|
||||
- 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)
|
||||
- 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: unless there's a sentence immediately after, don't add a "."
|
||||
- don't use →. Instead, prefer not to use arrows, and if need be, use -> instead
|
||||
|
||||
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
|
||||
|
|
|
|||
|
|
@ -1,8 +1,8 @@
|
|||
# Base image for building
|
||||
ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:31da6565f35af6401031c1d7aa91dc84ac76c5c48edd17fb90f0ed9e3173c7a9
|
||||
ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:c61ac6919b811ea53c4782d69f1fe05218ba3c25d53f01b6ab7892e621bd4370
|
||||
|
||||
# Runtime image
|
||||
ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:31da6565f35af6401031c1d7aa91dc84ac76c5c48edd17fb90f0ed9e3173c7a9
|
||||
ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:c61ac6919b811ea53c4782d69f1fe05218ba3c25d53f01b6ab7892e621bd4370
|
||||
ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a
|
||||
|
||||
FROM $UV_IMAGE AS uvbin
|
||||
|
|
|
|||
12
Makefile
12
Makefile
|
|
@ -6,7 +6,7 @@
|
|||
test-proxy-unit-a test-proxy-unit-b test-integration test-unit-helm \
|
||||
info lint lint-dev format \
|
||||
lint-basedpyright lint-basedpyright-budget-update \
|
||||
lint-ruff-budget lint-ruff-budget-update lint-budget-update \
|
||||
lint-ruff-budget lint-ruff-budget-update lint-budget-update lint-gate \
|
||||
install-dev install-proxy-dev install-test-deps install-hooks \
|
||||
install-helm-unittest check-circular-imports check-import-safety
|
||||
|
||||
|
|
@ -28,6 +28,7 @@ help:
|
|||
@echo " make lint-basedpyright-budget-update - Re-capture the basedpyright per-rule budget (ratchet)"
|
||||
@echo " make lint-black - Check Black formatting (matches CI)"
|
||||
@echo " make lint-ruff-budget - Gate the codebase total of each strict ruff rule against its ceiling"
|
||||
@echo " make lint-gate - Strict ruff gate in CI-parity mode (fetches staging, simulates the merge)"
|
||||
@echo " make lint-ruff-budget-update - Re-capture per-rule baselines in ruff-strict-budget.json (ratchet)"
|
||||
@echo " make lint-budget-update - Re-capture all ratchet budgets (ruff + basedpyright)"
|
||||
@echo " make check-circular-imports - Check for circular imports"
|
||||
|
|
@ -124,7 +125,8 @@ lint-ruff-FULL-dev: install-dev
|
|||
else echo "No changed .py files to check."; fi
|
||||
|
||||
lint-basedpyright: install-dev
|
||||
($(UV_RUN) basedpyright --outputjson || true) | $(UV_RUN) python scripts/type_check_gate.py
|
||||
git fetch origin litellm_internal_staging
|
||||
($(UV_RUN) basedpyright --outputjson || true) | $(UV_RUN) python scripts/type_check_gate.py --base origin/litellm_internal_staging
|
||||
|
||||
lint-basedpyright-budget-update: install-dev
|
||||
($(UV_RUN) basedpyright --outputjson || true) | $(UV_RUN) python scripts/type_check_gate.py --update
|
||||
|
|
@ -134,6 +136,12 @@ lint-black: format-check
|
|||
lint-ruff-budget: install-dev
|
||||
$(UV_RUN) python scripts/ruff_strict_gate.py
|
||||
|
||||
# Strict gate, invoked the same way CI does in test-linting.yml so a local pass
|
||||
# means the CI check will pass too.
|
||||
lint-gate: install-dev
|
||||
git fetch origin litellm_internal_staging
|
||||
$(UV_RUN) python scripts/ruff_strict_gate.py --base origin/litellm_internal_staging
|
||||
|
||||
lint-ruff-budget-update: install-dev
|
||||
$(UV_RUN) python scripts/ruff_strict_gate.py --update
|
||||
|
||||
|
|
|
|||
142
README.md
142
README.md
|
|
@ -6,10 +6,10 @@
|
|||
</p>
|
||||
<p align="center">Open Source AI Gateway for 100+ LLMs. Self-hosted. Enterprise-ready. Call any LLM in OpenAI format.</p>
|
||||
<p align="center">
|
||||
<a href="https://render.com/deploy?repo=https://github.com/BerriAI/litellm" target="_blank" rel="nofollow"><img src="https://render.com/images/deploy-to-render-button.svg" alt="Deploy to Render"></a>
|
||||
<a href="https://railway.com/deploy/RhvhdC?referralCode=7mRv9K&utm_medium=integration&utm_source=template&utm_campaign=generic">
|
||||
<img src="https://railway.com/button.svg" alt="Deploy on Railway">
|
||||
</a>
|
||||
<a href="https://render.com/deploy?repo=https://github.com/BerriAI/litellm" target="_blank" rel="nofollow"><img src="https://render.com/images/deploy-to-render-button.svg" alt="Deploy to Render" height="40"></a>
|
||||
<a href="https://railway.com/deploy/RhvhdC?referralCode=7mRv9K&utm_medium=integration&utm_source=template&utm_campaign=generic"><img src="https://railway.com/button.svg" alt="Deploy on Railway" height="40"></a>
|
||||
<a href="https://console.aws.amazon.com/cloudshell/home" target="_blank" rel="nofollow"><img src="./.github/deploy-on-aws.png" alt="Deploy on AWS" height="40"></a>
|
||||
<a href="https://ssh.cloud.google.com/cloudshell/editor?cloudshell_git_repo=https%3A%2F%2Fgithub.com%2FBerriAI%2Flitellm&cloudshell_workspace=terraform%2Flitellm%2Fgcp%2Fexamples%2Fdefault&cloudshell_tutorial=TUTORIAL.md&cloudshell_image=gcr.io/ds-artifacts-cloudshell/deploystack_custom_image&shellonly=true" target="_blank" rel="nofollow"><img src="./.github/deploy-on-gcp.png" alt="Deploy on GCP" height="40"></a>
|
||||
</p>
|
||||
</p>
|
||||
<h4 align="center"><a href="https://docs.litellm.ai/docs/simple_proxy" target="_blank">LiteLLM Proxy Server (AI Gateway)</a> | <a href="https://docs.litellm.ai/docs/enterprise#hosted-litellm-proxy" target="_blank"> Hosted Proxy</a> | <a href="https://litellm.ai/enterprise"target="_blank">Enterprise Tier</a> | <a href="https://www.litellm.ai/ai-gateway" target="_blank">Website</a></h4>
|
||||
|
|
@ -406,6 +406,140 @@ You can use LiteLLM through either the Proxy Server or Python SDK. Both give you
|
|||
|
||||
Support for more providers. Missing a provider or LLM Platform, raise a [feature request](https://github.com/BerriAI/litellm/issues/new?assignees=&labels=enhancement&projects=&template=feature_request.yml&title=%5BFeature%5D%3A+).
|
||||
|
||||
### Deploy on AWS or GCP with Terraform
|
||||
|
||||
Run the LiteLLM proxy as a production-ready componentized stack (gateway, backend, UI on separate services; managed Postgres + Redis + object store) using the published Terraform modules. Both modules are on the [public Terraform Registry](https://registry.terraform.io/namespaces/BerriAI) — no auth needed.
|
||||
|
||||
#### AWS — ECS Fargate + Aurora + ElastiCache + ALB
|
||||
|
||||
[](https://console.aws.amazon.com/cloudshell/home) — opens an in-browser shell, already authenticated to your AWS account. Once inside, run:
|
||||
|
||||
```bash
|
||||
git clone https://github.com/BerriAI/litellm.git
|
||||
cd litellm/terraform/litellm/aws/examples/default
|
||||
cp terraform.tfvars.example terraform.tfvars # edit region/tenant/env
|
||||
terraform init && terraform apply
|
||||
```
|
||||
|
||||
[Module page →](https://registry.terraform.io/modules/BerriAI/litellm/aws/latest)
|
||||
|
||||
Or call the module from your own root config:
|
||||
|
||||
```hcl
|
||||
# main.tf
|
||||
terraform {
|
||||
required_version = ">= 1.6.0"
|
||||
required_providers {
|
||||
aws = { source = "hashicorp/aws", version = "~> 5.60" }
|
||||
}
|
||||
}
|
||||
|
||||
provider "aws" {
|
||||
region = "us-west-2"
|
||||
}
|
||||
|
||||
module "litellm" {
|
||||
source = "BerriAI/litellm/aws"
|
||||
version = "~> 1.89"
|
||||
|
||||
region = "us-west-2"
|
||||
azs = ["us-west-2a", "us-west-2b"]
|
||||
tenant = "acme"
|
||||
env = "prod"
|
||||
|
||||
# Production: provide an ACM cert. Without one, set allow_plaintext_alb = true
|
||||
# (dev/trial only).
|
||||
# acm_certificate_arn = "arn:aws:acm:us-west-2:111122223333:certificate/..."
|
||||
allow_plaintext_alb = true
|
||||
}
|
||||
|
||||
output "litellm_url" {
|
||||
value = module.litellm.alb_dns_name
|
||||
}
|
||||
```
|
||||
|
||||
```bash
|
||||
terraform init
|
||||
terraform apply
|
||||
```
|
||||
|
||||
Provider API keys live in AWS Secrets Manager; reference ARNs via `gateway_extra_secrets`. Full input list and architecture diagram on the [registry page](https://registry.terraform.io/modules/BerriAI/litellm/aws/latest?tab=inputs).
|
||||
|
||||
#### GCP — Cloud Run + Cloud SQL + Memorystore + HTTPS LB
|
||||
|
||||
[](https://ssh.cloud.google.com/cloudshell/editor?cloudshell_git_repo=https%3A%2F%2Fgithub.com%2FBerriAI%2Flitellm&cloudshell_workspace=terraform%2Flitellm%2Fgcp%2Fexamples%2Fdefault&cloudshell_tutorial=TUTORIAL.md&cloudshell_image=gcr.io/ds-artifacts-cloudshell/deploystack_custom_image&shellonly=true)
|
||||
|
||||
Real 1-click. Opens Cloud Shell, clones this repo, and walks you through `terraform apply` via a built-in [DeployStack tutorial](./terraform/litellm/gcp/examples/default/TUTORIAL.md) — pick the project, the tutorial sets up the Artifact Registry remote repo, writes `terraform.tfvars` from your answers, and runs apply.
|
||||
|
||||
[Module page →](https://registry.terraform.io/modules/BerriAI/litellm/google/latest)
|
||||
|
||||
To call the module from your own config instead, Cloud Run can't pull from `ghcr.io` directly, so first set up a one-time Artifact Registry remote repo backed by GHCR:
|
||||
|
||||
```bash
|
||||
gcloud artifacts repositories create litellm \
|
||||
--location=us-central1 \
|
||||
--repository-format=docker \
|
||||
--mode=remote-repository \
|
||||
--remote-docker-repo=https://ghcr.io \
|
||||
--project=my-gcp-project
|
||||
```
|
||||
|
||||
Then:
|
||||
|
||||
```hcl
|
||||
# main.tf
|
||||
terraform {
|
||||
required_version = ">= 1.6.0"
|
||||
required_providers {
|
||||
google = { source = "hashicorp/google", version = "~> 6.10" }
|
||||
google-beta = { source = "hashicorp/google-beta", version = "~> 6.10" }
|
||||
}
|
||||
}
|
||||
|
||||
provider "google" { project = "my-gcp-project"; region = "us-central1" }
|
||||
provider "google-beta" { project = "my-gcp-project"; region = "us-central1" }
|
||||
|
||||
module "litellm" {
|
||||
source = "BerriAI/litellm/google"
|
||||
version = "~> 1.89"
|
||||
|
||||
project_id = "my-gcp-project"
|
||||
region = "us-central1"
|
||||
tenant = "acme"
|
||||
env = "prod"
|
||||
|
||||
# Replace my-gcp-project with your GCP project ID (same value as project_id above).
|
||||
image_registry = "us-central1-docker.pkg.dev/my-gcp-project/litellm/berriai"
|
||||
|
||||
# Production: provide DNS already pointing at the LB IP for Google-managed certs.
|
||||
# Without one, set allow_plaintext_lb = true (dev/trial only).
|
||||
# lb_domains = ["proxy.example.com"]
|
||||
allow_plaintext_lb = true
|
||||
}
|
||||
|
||||
output "litellm_url" {
|
||||
value = module.litellm.load_balancer_url
|
||||
}
|
||||
```
|
||||
|
||||
```bash
|
||||
terraform init
|
||||
terraform apply
|
||||
```
|
||||
|
||||
Provider API keys live in Secret Manager; reference resource IDs (e.g. `projects/my-gcp-project/secrets/openai-api-key`) via `gateway_extra_secrets`. Full input list and architecture diagram on the [registry page](https://registry.terraform.io/modules/BerriAI/litellm/google/latest?tab=inputs).
|
||||
|
||||
#### Both stacks include
|
||||
|
||||
- The full componentized split (gateway / backend / UI as independent services)
|
||||
- Managed Postgres (writer + reader) and Redis
|
||||
- Versioned object store for proxy state + file uploads
|
||||
- An auto-generated `LITELLM_MASTER_KEY` in your cloud's secret manager
|
||||
- A one-off migration job that runs `prisma migrate deploy` before the proxy starts
|
||||
- The same `proxy_config` surface as the [Helm chart](./helm/litellm/) — pass YAML as a typed map
|
||||
|
||||
The Terraform modules live at [`terraform/litellm/aws/`](./terraform/litellm/aws/) and [`terraform/litellm/gcp/`](./terraform/litellm/gcp/) in this repo; the registry entries are read-only mirrors updated on each release.
|
||||
|
||||
### Run in Developer Mode
|
||||
#### Services
|
||||
1. Setup .env file in root
|
||||
|
|
|
|||
|
|
@ -1,5 +1,5 @@
|
|||
ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:31da6565f35af6401031c1d7aa91dc84ac76c5c48edd17fb90f0ed9e3173c7a9
|
||||
ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:31da6565f35af6401031c1d7aa91dc84ac76c5c48edd17fb90f0ed9e3173c7a9
|
||||
ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:c61ac6919b811ea53c4782d69f1fe05218ba3c25d53f01b6ab7892e621bd4370
|
||||
ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:c61ac6919b811ea53c4782d69f1fe05218ba3c25d53f01b6ab7892e621bd4370
|
||||
ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a
|
||||
|
||||
FROM $UV_IMAGE AS uvbin
|
||||
|
|
|
|||
|
|
@ -120,6 +120,9 @@ BACKEND_PATH_PREFIXES: tuple[str, ...] = (
|
|||
"/robots.txt",
|
||||
# Health (k8s probes)
|
||||
"/health",
|
||||
# Plugin system
|
||||
"/api/plugins",
|
||||
"/plugin-proxy/",
|
||||
)
|
||||
|
||||
BACKEND_EXACT_PATHS: frozenset[str] = frozenset(
|
||||
|
|
|
|||
|
|
@ -121,7 +121,7 @@
|
|||
},
|
||||
"reportReturnType": {
|
||||
"baseline": 126,
|
||||
"slack": 13
|
||||
"slack": 100
|
||||
},
|
||||
"reportTypedDictNotRequiredAccess": {
|
||||
"baseline": 20,
|
||||
|
|
@ -157,7 +157,7 @@
|
|||
},
|
||||
"reportUnnecessaryComparison": {
|
||||
"baseline": 683,
|
||||
"slack": 10
|
||||
"slack": 100
|
||||
},
|
||||
"reportUnnecessaryContains": {
|
||||
"baseline": 4,
|
||||
|
|
|
|||
|
|
@ -1,8 +1,8 @@
|
|||
# Base image for building
|
||||
ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:3258be472764337fd13095bcbb3182da170243b5819fd67ad4c0754590588b31
|
||||
ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:c61ac6919b811ea53c4782d69f1fe05218ba3c25d53f01b6ab7892e621bd4370
|
||||
|
||||
# Runtime image
|
||||
ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:3258be472764337fd13095bcbb3182da170243b5819fd67ad4c0754590588b31
|
||||
ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:c61ac6919b811ea53c4782d69f1fe05218ba3c25d53f01b6ab7892e621bd4370
|
||||
ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a
|
||||
|
||||
FROM $UV_IMAGE AS uvbin
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
# Base images
|
||||
ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:3258be472764337fd13095bcbb3182da170243b5819fd67ad4c0754590588b31
|
||||
ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:3258be472764337fd13095bcbb3182da170243b5819fd67ad4c0754590588b31
|
||||
ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:c61ac6919b811ea53c4782d69f1fe05218ba3c25d53f01b6ab7892e621bd4370
|
||||
ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:c61ac6919b811ea53c4782d69f1fe05218ba3c25d53f01b6ab7892e621bd4370
|
||||
ARG PROXY_EXTRAS_SOURCE=published
|
||||
ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a
|
||||
|
||||
|
|
|
|||
141
docs/plugin_architecture.md
Normal file
141
docs/plugin_architecture.md
Normal file
|
|
@ -0,0 +1,141 @@
|
|||
# LiteLLM Plugin Architecture
|
||||
|
||||
Plugins let external services appear as selectable modes in the litellm UI sidebar alongside the AI Gateway.
|
||||
|
||||
---
|
||||
|
||||
## Quick start
|
||||
|
||||
### 1. Configure the plugin
|
||||
|
||||
Add a `plugins` block to your litellm `config.yaml`:
|
||||
|
||||
```yaml
|
||||
general_settings:
|
||||
master_key: sk-...
|
||||
plugins:
|
||||
- name: my-plugin # unique identifier (no spaces)
|
||||
display_name: My Plugin # shown in the UI dropdown
|
||||
url: "https://my-plugin.example.com"
|
||||
plugin_key: "sk-..." # plugin's own auth credential
|
||||
```
|
||||
|
||||
`plugin_key` is injected as `Authorization: Bearer <plugin_key>` on every
|
||||
request proxied through `/plugin-proxy/my-plugin/*`. The caller's litellm
|
||||
credential is stripped before forwarding so the plugin never receives a live
|
||||
litellm API key.
|
||||
|
||||
### 2. Implement two endpoints on your service
|
||||
|
||||
| Endpoint | Method | Purpose |
|
||||
|---|---|---|
|
||||
| `GET /api/plugin-manifest` | public | Returns plugin metadata for the UI |
|
||||
| `POST /api/plugin-auth` | public | Decrypts the identity claim for seamless sign-in |
|
||||
|
||||
#### `GET /api/plugin-manifest`
|
||||
|
||||
```json
|
||||
{
|
||||
"name": "my-plugin",
|
||||
"display_name": "My Plugin",
|
||||
"version": "1.0.0",
|
||||
"nav_items": [
|
||||
{ "key": "home", "label": "Home", "icon": "HomeOutlined", "path": "/" },
|
||||
{ "key": "reports", "label": "Reports", "icon": "BarChartOutlined", "path": "/reports" }
|
||||
],
|
||||
"capabilities": ["reports", "data"]
|
||||
}
|
||||
```
|
||||
|
||||
#### `POST /api/plugin-auth`
|
||||
|
||||
Receives `{ "session_claim": "<fernet-ciphertext>" }`.
|
||||
|
||||
The proxy never shares `LITELLM_SALT_KEY` with your plugin. Each plugin is
|
||||
provisioned with its own dedicated key, derived as
|
||||
`HMAC-SHA256(LITELLM_SALT_KEY, plugin_name)`. Compute it once on the proxy
|
||||
host and hand the result to your plugin as a secret (e.g. `PLUGIN_AUTH_KEY`):
|
||||
|
||||
```bash
|
||||
python -c 'import base64,hmac,hashlib,os; \
|
||||
print(base64.urlsafe_b64encode(hmac.new(os.environ["LITELLM_SALT_KEY"].encode(), b"my-plugin", hashlib.sha256).digest()).decode())'
|
||||
```
|
||||
|
||||
A compromised plugin holding only this scoped key cannot recover
|
||||
`LITELLM_SALT_KEY` or decrypt any other litellm secret.
|
||||
|
||||
Decrypt and validate the claim with that key:
|
||||
|
||||
```python
|
||||
import json, os, time
|
||||
from cryptography.fernet import Fernet
|
||||
|
||||
_CLAIM_TTL_SECONDS = 30
|
||||
|
||||
def plugin_auth(session_claim: str) -> dict:
|
||||
cipher = Fernet(os.environ["PLUGIN_AUTH_KEY"].encode())
|
||||
claim = json.loads(cipher.decrypt(session_claim.encode(), ttl=_CLAIM_TTL_SECONDS))
|
||||
if claim.get("plugin") != "my-plugin":
|
||||
raise ValueError("claim audience mismatch")
|
||||
if int(claim.get("exp", 0)) < int(time.time()):
|
||||
raise ValueError("claim expired")
|
||||
return claim
|
||||
```
|
||||
|
||||
The claim is `{ "plugin", "user_id", "user_role", "exp" }`; it carries no
|
||||
litellm bearer token. Establish the plugin's own session from `user_id` /
|
||||
`user_role` and authenticate API calls back to litellm through the
|
||||
`/plugin-proxy/my-plugin/*` reverse proxy, which injects `plugin_key` for you.
|
||||
|
||||
---
|
||||
|
||||
## How iframe auth works
|
||||
|
||||
```
|
||||
litellm UI
|
||||
├─ GET /api/plugins/auth-token -> { session_claim }
|
||||
└─ postMessage({ type:"litellm-auth", session_claim }, pluginOrigin)
|
||||
│
|
||||
▼
|
||||
Plugin iframe browser
|
||||
└─ POST /api/plugin-auth { session_claim }
|
||||
│
|
||||
▼
|
||||
Plugin server
|
||||
├─ decrypt(session_claim, PLUGIN_AUTH_KEY) -> { user_id, user_role, exp }
|
||||
└─ establish plugin session -> stored in sessionStorage
|
||||
```
|
||||
|
||||
No litellm bearer token ever leaves the proxy; the claim only conveys the
|
||||
caller's identity and expires after 30 seconds. A postMessage intercept
|
||||
yields ciphertext that is useless without the plugin's scoped key.
|
||||
|
||||
---
|
||||
|
||||
## Proxy routes
|
||||
|
||||
- `GET /api/plugins` — list registered plugins (`name`, `display_name`, `url`). `plugin_key` is **never** returned; it stays server-side. Requires an authenticated caller.
|
||||
- `GET /api/plugins/auth-token?plugin_name=<name>` — short-lived encrypted identity claim for the named plugin. Requires `LITELLM_SALT_KEY` to be set (503 otherwise) and the plugin to be registered (404 otherwise).
|
||||
- `ANY /plugin-proxy/{name}/{path}` — authenticated reverse proxy to the plugin backend. Restricted to `proxy_admin`.
|
||||
|
||||
---
|
||||
|
||||
## Reverse proxy behaviour
|
||||
|
||||
When an admin (or server-to-server caller) hits `/plugin-proxy/<name>/<path>`, the proxy authenticates the caller locally, then rewrites the request before forwarding it to the plugin's `url`:
|
||||
|
||||
- **Every litellm credential header is stripped** — `Authorization`, `x-api-key`, `API-Key`, `x-goog-api-key`, `Ocp-Apim-Subscription-Key`, `x-litellm-api-key`, any configured `litellm_key_header_name`, plus `Cookie`. The plugin can never be handed the caller's live litellm key.
|
||||
- **`plugin_key` is injected** as `Authorization: Bearer <plugin_key>` — the only credential the plugin receives.
|
||||
- **Caller identity is forwarded** as `x-litellm-user-id` and `x-litellm-user-role` so the plugin can run its own authorization. These are informational, not credentials.
|
||||
- **Responses are sandboxed** — `Content-Security-Policy: sandbox` and `X-Content-Type-Options: nosniff` are set so plugin-controlled bytes served from the litellm origin cannot execute against the dashboard.
|
||||
|
||||
---
|
||||
|
||||
## Security checklist
|
||||
|
||||
- [ ] `LITELLM_SALT_KEY` is set on the proxy and never shared with the plugin
|
||||
- [ ] The plugin holds only its derived `HMAC(LITELLM_SALT_KEY, plugin_name)` key, provisioned as a dedicated secret
|
||||
- [ ] `plugin_key` is a dedicated credential scoped to the plugin (not your litellm master key)
|
||||
- [ ] Plugin's `POST /api/plugin-auth` enforces the claim's `plugin` audience and `exp` (30s TTL)
|
||||
- [ ] Plugin treats `x-litellm-user-id` / `x-litellm-user-role` as identity hints, not as proof of authentication
|
||||
- [ ] Plugin service URL uses HTTPS in production
|
||||
|
|
@ -1,6 +1,6 @@
|
|||
[project]
|
||||
name = "litellm-enterprise"
|
||||
version = "0.1.42"
|
||||
version = "0.1.43"
|
||||
description = "Package for LiteLLM Enterprise features"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.9"
|
||||
|
|
@ -26,7 +26,7 @@ required-version = ">=0.10.9"
|
|||
module-root = ""
|
||||
|
||||
[tool.commitizen]
|
||||
version = "0.1.42"
|
||||
version = "0.1.43"
|
||||
version_files = [
|
||||
"pyproject.toml:^version",
|
||||
"../pyproject.toml:litellm-enterprise==",
|
||||
|
|
|
|||
|
|
@ -1,5 +1,5 @@
|
|||
ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:31da6565f35af6401031c1d7aa91dc84ac76c5c48edd17fb90f0ed9e3173c7a9
|
||||
ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:31da6565f35af6401031c1d7aa91dc84ac76c5c48edd17fb90f0ed9e3173c7a9
|
||||
ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:c61ac6919b811ea53c4782d69f1fe05218ba3c25d53f01b6ab7892e621bd4370
|
||||
ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:c61ac6919b811ea53c4782d69f1fe05218ba3c25d53f01b6ab7892e621bd4370
|
||||
ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a
|
||||
|
||||
FROM $UV_IMAGE AS uvbin
|
||||
|
|
|
|||
1
litellm-rust/.gitignore
vendored
Normal file
1
litellm-rust/.gitignore
vendored
Normal file
|
|
@ -0,0 +1 @@
|
|||
/target/
|
||||
9
litellm-rust/ADDING_A_PROVIDER.md
Normal file
9
litellm-rust/ADDING_A_PROVIDER.md
Normal file
|
|
@ -0,0 +1,9 @@
|
|||
# Adding a provider / route to litellm-rust
|
||||
|
||||
Three layers, same for every route (see `ocr` and `realtime` as references):
|
||||
|
||||
1. **Transform contract (pure)** — `crates/core/src/<route>/transformation.rs`: a `…ProviderConfig` trait (URL build + request/response transforms) + types in `types.rs`. No network, env, or auth.
|
||||
2. **Provider config (pure)** — `crates/providers/src/<provider>/<route>/transformation.rs`: implement that trait as a `const <PROVIDER>_<ROUTE>_CONFIG`, mirroring the Python provider tree. Add parity unit tests.
|
||||
3. **HTTP / transport (the host)** — `crates/providers/src/<route>.rs` (e.g. `ocr.rs`, `realtime.rs`): the callable fn (`run_ocr`, `realtime`). It resolves the key, builds the auth header, builds URL + transforms via the config, then does the network call. This is the only layer allowed to do I/O.
|
||||
|
||||
**Calling:** the host invokes the route fn — the Python bridge calls `run_ocr`; the `ai-gateway` server calls `realtime`. Register new modules in `lib.rs` / `mod.rs`, then run `cargo fmt && cargo clippy --workspace -- -D warnings && cargo test --workspace`.
|
||||
88
litellm-rust/CLAUDE.md
Normal file
88
litellm-rust/CLAUDE.md
Normal file
|
|
@ -0,0 +1,88 @@
|
|||
# CLAUDE.md
|
||||
|
||||
This file defines the rules for Rust work in LiteLLM.
|
||||
|
||||
## Core Boundary
|
||||
|
||||
The `core` and `providers` crates describe work; hosts execute work.
|
||||
|
||||
Route-level Rust structure mirrors LiteLLM's Python responsibilities:
|
||||
- `core/src/<route>/` owns the route contract, shared types, and provider
|
||||
template traits. For OCR, this means `core/src/ocr`.
|
||||
- `providers/src/<provider>/<route>/transformation.rs` owns the
|
||||
provider-specific transform. For Mistral OCR, this means
|
||||
`providers/src/mistral/ocr/transformation.rs`.
|
||||
- Future network execution belongs in a host/transport layer such as
|
||||
`llm_http_handler`, not inside `core` or `providers`.
|
||||
|
||||
Allowed in `core` and `providers`:
|
||||
- Pure request transforms
|
||||
- Pure response transforms
|
||||
- Pure stream chunk normalization
|
||||
- Shared data types and validation errors
|
||||
- Deterministic token/cost helper logic
|
||||
|
||||
Not allowed in `core` or `providers`:
|
||||
- Network calls
|
||||
- Environment variable or secret reads
|
||||
- Filesystem access
|
||||
- Database or cache access
|
||||
- Provider SDK signing or auth flows
|
||||
- Logging callbacks, spend writes, or custom callbacks
|
||||
- Global mutable runtime state
|
||||
|
||||
Python owns rollout state and fallback while Rust is being introduced. Rust
|
||||
paths must be off by default until parity tests prove equivalence with Python.
|
||||
|
||||
## Production Bar
|
||||
|
||||
Rust code in this workspace is held to a strict parity and robustness bar from
|
||||
the first PR:
|
||||
|
||||
- Correctness parity is proven with tests. Do not rely on README claims or
|
||||
manual inspection for a port that mirrors Python behavior.
|
||||
- Every provider transform must have unit tests for supported-parameter
|
||||
filtering, request body shape, response normalization, missing/null fields,
|
||||
and bad-input errors.
|
||||
- When Rust is exposed through Python, add Python tests that prove disabled,
|
||||
enabled, and unavailable-bridge fallback behavior.
|
||||
- Avoid panics on user/provider input. Return typed errors and let the host map
|
||||
them to Python exceptions or HTTP responses.
|
||||
- OCR handles documents that often contain personal data. Do not log document
|
||||
contents, base64 payloads, provider response bodies, or secrets.
|
||||
- Error messages must be useful but data-minimized. Truncate or sanitize any
|
||||
upstream body before it crosses a host boundary.
|
||||
- Treat empty or whitespace-only credentials, URLs, and config values as absent
|
||||
at the host/config resolution layer.
|
||||
- Preserve Python output shape intentionally. If a field is always serialized as
|
||||
`null` for Python parity, leave a short comment explaining that parity choice.
|
||||
|
||||
## Host I/O Rules
|
||||
|
||||
These rules apply when adding future crates or modules that execute network I/O,
|
||||
such as `ai-gateway`, router hosts, or standalone servers:
|
||||
|
||||
- Set connect and full-request timeouts. No unbounded waits.
|
||||
- Reuse HTTP clients; do not construct clients per request.
|
||||
- Prefer rustls TLS for portable Python wheels and Linux images unless there is
|
||||
a documented reason not to.
|
||||
- Add request IDs and structured tracing at the host layer, without logging OCR
|
||||
document contents or secrets.
|
||||
- Do not echo raw upstream response bodies to callers. Sanitize and bound them.
|
||||
- Avoid `expect`/`unwrap` in server startup and request paths unless the panic is
|
||||
impossible by construction and documented.
|
||||
|
||||
## Checks
|
||||
|
||||
Run these before pushing Rust changes. The same checks run in GitHub Actions
|
||||
for changes under `litellm-rust/`.
|
||||
|
||||
```bash
|
||||
cd litellm-rust
|
||||
cargo fmt --check
|
||||
cargo clippy --workspace --all-targets -- -D warnings
|
||||
cargo test --workspace
|
||||
```
|
||||
|
||||
When a Rust path is exposed through Python, add Python parity tests that compare
|
||||
the existing Python output with the Rust-backed output.
|
||||
1872
litellm-rust/Cargo.lock
generated
Normal file
1872
litellm-rust/Cargo.lock
generated
Normal file
File diff suppressed because it is too large
Load diff
28
litellm-rust/Cargo.toml
Normal file
28
litellm-rust/Cargo.toml
Normal file
|
|
@ -0,0 +1,28 @@
|
|||
[workspace]
|
||||
members = [
|
||||
"crates/core",
|
||||
"crates/providers",
|
||||
"crates/python-bridge",
|
||||
"crates/ai-gateway",
|
||||
]
|
||||
resolver = "2"
|
||||
|
||||
[workspace.package]
|
||||
edition = "2021"
|
||||
license = "MIT"
|
||||
repository = "https://github.com/BerriAI/litellm"
|
||||
|
||||
[workspace.dependencies]
|
||||
litellm-core = { path = "crates/core" }
|
||||
litellm-providers = { path = "crates/providers" }
|
||||
axum = "0.7"
|
||||
pyo3 = "0.23.5"
|
||||
rand = "0.8"
|
||||
reqwest = { version = "0.12", default-features = false, features = ["blocking", "json", "rustls-tls"] }
|
||||
serde = { version = "1.0", features = ["derive"] }
|
||||
serde_json = "1.0"
|
||||
subtle = "2"
|
||||
thiserror = "2.0"
|
||||
tokio = { version = "1", features = ["rt-multi-thread", "macros", "time"] }
|
||||
tokio-tungstenite = { version = "0.24", default-features = false, features = ["connect", "rustls-tls-native-roots"] }
|
||||
futures-util = { version = "0.3", default-features = false, features = ["sink", "std"] }
|
||||
34
litellm-rust/README.md
Normal file
34
litellm-rust/README.md
Normal file
|
|
@ -0,0 +1,34 @@
|
|||
# LiteLLM Rust
|
||||
|
||||
This workspace contains the staged Rust implementation for LiteLLM.
|
||||
|
||||
Rust starts as a pure transform core used by the existing Python host. Python
|
||||
continues to own auth, configuration, network I/O, retries, routing, logging,
|
||||
callbacks, spend tracking, and customer plugins until each Rust path has parity
|
||||
coverage and production evidence.
|
||||
|
||||
## Layout
|
||||
|
||||
```text
|
||||
crates/
|
||||
core/ Route contracts, shared pure types, errors, and templates.
|
||||
src/ocr/
|
||||
providers/ Provider-specific pure transforms.
|
||||
src/mistral/ocr/transformation.rs
|
||||
python-bridge/ PyO3 bridge for Python LiteLLM.
|
||||
```
|
||||
|
||||
The folder shape should follow the Python provider tree:
|
||||
`providers/src/<provider>/<route>/transformation.rs`. The bridge should expose
|
||||
one function per top-level route, starting with `ocr(payload)`.
|
||||
|
||||
## Checks
|
||||
|
||||
Run these before pushing Rust changes. GitHub Actions runs the same checks for
|
||||
changes under `litellm-rust/`.
|
||||
|
||||
```bash
|
||||
cargo fmt --check
|
||||
cargo clippy --workspace --all-targets -- -D warnings
|
||||
cargo test --workspace
|
||||
```
|
||||
50
litellm-rust/crates/ai-gateway/AGENTS.md
Normal file
50
litellm-rust/crates/ai-gateway/AGENTS.md
Normal file
|
|
@ -0,0 +1,50 @@
|
|||
# ai-gateway — folder architecture
|
||||
|
||||
The Axum server that fronts the Rust gateway. It owns transport + config + auth
|
||||
only; deployment selection lives in `core::router`, transforms in `core`/`providers`.
|
||||
|
||||
```
|
||||
src/
|
||||
main.rs # entrypoint: build AppState (router + master key), bind, serve
|
||||
state.rs # AppState — shared Arc<Router> + master_key
|
||||
gil.rs # GIL-activity tracker (records Python acquisitions)
|
||||
auth/ # authentication as an axum extractor — added to handler args
|
||||
mod.rs # RequireMasterKey: FromRequestParts, single master key (LITELLM_MASTER_KEY)
|
||||
routes/ # one module per route, all matching the same template
|
||||
AGENTS.md # ← the route template (read this before adding a route)
|
||||
mod.rs # app(): merges every module's router()
|
||||
health.rs # simple route (one file): router() + liveness/readiness
|
||||
gil.rs # simple route (one file): router() + GET /health/gil
|
||||
realtime/ # route with logic → axum surface + a no-axum service:
|
||||
mod.rs # router() + handler + WS<->events adapter (the axum surface)
|
||||
service.rs # business logic (select deployment, call provider) — no axum, testable
|
||||
python/ # Python interop (feature: python-config) — load-time only
|
||||
mod.rs, config.rs, AGENTS.md
|
||||
```
|
||||
|
||||
## Rules
|
||||
|
||||
- **Routes follow one template.** Each route module exposes
|
||||
`pub fn router() -> Router<AppState>`; `routes/mod.rs` only merges them. Simple
|
||||
routes are one file; non-trivial routes are a folder (`handler`/`service`/
|
||||
`transport`). See `routes/AGENTS.md`.
|
||||
- **Auth is an extractor.** Add `crate::auth::RequireMasterKey` to a handler's
|
||||
args; it runs during extraction. Never re-implement the check per route.
|
||||
- **Handlers are thin.** A handler validates and delegates to its `service`. No
|
||||
business logic, no provider calls, no transforms in handlers.
|
||||
- **State is shared and cheap to clone.** Long-lived handles live behind `Arc` in
|
||||
`state.rs`; read env/config only in `main.rs` when building state.
|
||||
|
||||
## Auth (interim)
|
||||
|
||||
A single **master key** (`LITELLM_MASTER_KEY`), enforced by the
|
||||
`auth::RequireMasterKey` extractor: any caller presenting it as
|
||||
`Authorization: Bearer <key>` may invoke the gateway. Fails closed (500) when
|
||||
unset; constant-time compare. The server binds `127.0.0.1` by default (`HOST` to
|
||||
override). Full per-key auth + budgets/rate-limits are delegated to the Python
|
||||
proxy in a later phase. Health routes don't add the extractor (unauthenticated).
|
||||
|
||||
## Python interop
|
||||
|
||||
Anything that calls into Python lives in `python/` and is **load-time only** — see
|
||||
`python/AGENTS.md`. The realtime data path never takes the GIL.
|
||||
26
litellm-rust/crates/ai-gateway/Cargo.toml
Normal file
26
litellm-rust/crates/ai-gateway/Cargo.toml
Normal file
|
|
@ -0,0 +1,26 @@
|
|||
[package]
|
||||
name = "litellm-ai-gateway"
|
||||
version = "0.1.0"
|
||||
edition.workspace = true
|
||||
license.workspace = true
|
||||
repository.workspace = true
|
||||
|
||||
[[bin]]
|
||||
name = "litellm-ai-gateway"
|
||||
path = "src/main.rs"
|
||||
|
||||
[dependencies]
|
||||
litellm-core.workspace = true
|
||||
litellm-providers.workspace = true
|
||||
axum = { workspace = true, features = ["ws"] }
|
||||
futures-util.workspace = true
|
||||
tokio = { workspace = true, features = ["rt-multi-thread", "macros", "net", "time"] }
|
||||
serde.workspace = true
|
||||
serde_json.workspace = true
|
||||
subtle.workspace = true
|
||||
pyo3 = { workspace = true, features = ["auto-initialize"], optional = true }
|
||||
|
||||
[features]
|
||||
# Build the gateway's config from the proxy YAML via an embedded Python
|
||||
# interpreter (links libpython; requires `litellm` importable at runtime).
|
||||
python-config = ["dep:pyo3"]
|
||||
86
litellm-rust/crates/ai-gateway/Dockerfile
Normal file
86
litellm-rust/crates/ai-gateway/Dockerfile
Normal file
|
|
@ -0,0 +1,86 @@
|
|||
# Multi-stage build for the LiteLLM Rust AI Gateway (realtime WebSocket proxy).
|
||||
#
|
||||
# Build context is the **repo root** so we can install `litellm` from this repo's
|
||||
# source (the gateway loads its model_list via litellm.proxy.read_model_list,
|
||||
# which is not in any PyPI release yet) AND build the rust workspace under
|
||||
# litellm-rust/.
|
||||
#
|
||||
# docker build -f litellm-rust/crates/ai-gateway/Dockerfile -t litellm-ai-gateway .
|
||||
#
|
||||
# No secrets live in this file. Runtime config (LITELLM_MASTER_KEY,
|
||||
# OPENAI_API_KEY referenced by config.yaml, etc.) is injected as environment
|
||||
# variables at deploy time.
|
||||
|
||||
# ---- Chef -------------------------------------------------------------------
|
||||
# cargo-chef caches the dependency build so only the gateway crate recompiles on
|
||||
# a source-only change. python3-dev is present in every rust stage because the
|
||||
# `python-config` feature links libpython via pyo3 (even in the cook step).
|
||||
FROM rust:1.90-slim-bookworm AS chef
|
||||
ENV PYO3_PYTHON=python3.11
|
||||
RUN apt-get update \
|
||||
&& apt-get install -y --no-install-recommends \
|
||||
python3 python3-dev pkg-config libssl-dev clang \
|
||||
&& rm -rf /var/lib/apt/lists/* \
|
||||
&& cargo install cargo-chef --locked --version 0.1.77
|
||||
WORKDIR /build/litellm-rust
|
||||
|
||||
# ---- Planner ----------------------------------------------------------------
|
||||
# Produce the dependency recipe from the rust workspace manifests + Cargo.lock.
|
||||
FROM chef AS planner
|
||||
COPY litellm-rust/ .
|
||||
RUN cargo chef prepare --recipe-path recipe.json
|
||||
|
||||
# ---- Builder ----------------------------------------------------------------
|
||||
FROM chef AS builder
|
||||
# Cook (compile) just the dependencies first — this layer is cached and reused
|
||||
# whenever only gateway source changes.
|
||||
COPY --from=planner /build/litellm-rust/recipe.json recipe.json
|
||||
RUN cargo chef cook --locked --release \
|
||||
-p litellm-ai-gateway --features python-config \
|
||||
--recipe-path recipe.json
|
||||
# Now copy the real sources and build the gateway binary. Deps are already cooked
|
||||
# above, so this step only recompiles the gateway crate.
|
||||
COPY litellm-rust/ .
|
||||
RUN cargo build --locked --release -p litellm-ai-gateway --features python-config
|
||||
|
||||
# ---- Runtime ----------------------------------------------------------------
|
||||
# python:3.11-slim-bookworm ships libpython3.11, matching the builder's PyO3
|
||||
# 3.11 ABI so the embedded interpreter links and imports cleanly.
|
||||
FROM python:3.11-slim-bookworm AS runtime
|
||||
|
||||
# CA certificates for outbound TLS to the OpenAI realtime endpoint.
|
||||
RUN apt-get update \
|
||||
&& apt-get install -y --no-install-recommends ca-certificates \
|
||||
&& rm -rf /var/lib/apt/lists/*
|
||||
|
||||
WORKDIR /app
|
||||
|
||||
# Install litellm (with proxy extras) FROM THIS REPO'S SOURCE so
|
||||
# `import litellm.proxy.read_model_list` works — it is not on PyPI yet. Copy the
|
||||
# package + packaging metadata, then pip install the proxy extra.
|
||||
COPY pyproject.toml README.md LICENSE ./
|
||||
COPY litellm/ ./litellm/
|
||||
RUN pip install --no-cache-dir ".[proxy]"
|
||||
|
||||
# The compiled gateway binary (pure-Rust realtime hot path; Python is load-time
|
||||
# only).
|
||||
COPY --from=builder /build/litellm-rust/target/release/litellm-ai-gateway /usr/local/bin/litellm-ai-gateway
|
||||
|
||||
# Default config.yaml. A real deploy can override this (e.g. mount a Render
|
||||
# secret file at the same path) — never bake secrets into the image.
|
||||
COPY litellm-rust/crates/ai-gateway/config.yaml /app/config.yaml
|
||||
|
||||
# Bind to all interfaces (Render routes to 0.0.0.0:$PORT) and load the model_list
|
||||
# from config.yaml via the embedded python config reader.
|
||||
ENV HOST=0.0.0.0 \
|
||||
LITELLM_CONFIG_PATH=/app/config.yaml
|
||||
|
||||
# Drop to a non-root user. The realtime hot path needs no root privileges, so
|
||||
# running unprivileged limits blast radius if the process is ever compromised.
|
||||
# The binary in /usr/local/bin is world-executable (COPY default mode 755); we
|
||||
# only need /app (and the config.yaml it reads) owned by the unprivileged user.
|
||||
RUN useradd --system --no-create-home --uid 10001 appuser \
|
||||
&& chown -R appuser:appuser /app
|
||||
USER appuser
|
||||
|
||||
ENTRYPOINT ["/usr/local/bin/litellm-ai-gateway"]
|
||||
45
litellm-rust/crates/ai-gateway/Dockerfile.dockerignore
Normal file
45
litellm-rust/crates/ai-gateway/Dockerfile.dockerignore
Normal file
|
|
@ -0,0 +1,45 @@
|
|||
# Dockerfile-specific ignore-file for the Rust AI Gateway build.
|
||||
#
|
||||
# The build context is the repo root (so the image can pip install litellm from
|
||||
# source AND build the rust workspace). BuildKit honors `<Dockerfile>.dockerignore`
|
||||
# next to the Dockerfile and it takes precedence over the repo-root `.dockerignore`,
|
||||
# so this file shrinks the (large) repo-root context for THIS build only without
|
||||
# touching the root `.dockerignore` used by the main litellm images.
|
||||
#
|
||||
# Strategy: ignore everything, then re-include only what the build needs:
|
||||
# - litellm/ (pip install . needs the full package + proxy reader)
|
||||
# - litellm-rust/ (the rust workspace; Cargo.lock + crate sources)
|
||||
# - pyproject.toml / README.md / LICENSE (packaging metadata for pip install)
|
||||
*
|
||||
|
||||
# --- re-include the build inputs ---
|
||||
!litellm/
|
||||
!litellm-rust/
|
||||
!pyproject.toml
|
||||
!README.md
|
||||
!LICENSE
|
||||
|
||||
# --- prune heavy / irrelevant subpaths back out of the re-included trees ---
|
||||
# Rust build artifacts (huge; regenerated in the builder).
|
||||
**/target/
|
||||
# Python caches and compiled bytecode.
|
||||
**/__pycache__/
|
||||
**/*.pyc
|
||||
**/*.pyo
|
||||
**/.pytest_cache/
|
||||
**/.ruff_cache/
|
||||
**/.mypy_cache/
|
||||
# Node / UI build output bundled under the python package (not needed to import
|
||||
# litellm.proxy.read_model_list).
|
||||
**/node_modules/
|
||||
litellm/proxy/_experimental/out/
|
||||
# Tests, logs, and local scratch.
|
||||
**/tests/
|
||||
**/test/
|
||||
*.log
|
||||
log.txt
|
||||
*.tgz
|
||||
# VCS / editor / CI metadata that may live under re-included trees.
|
||||
**/.git/
|
||||
.git/
|
||||
**/.DS_Store
|
||||
172
litellm-rust/crates/ai-gateway/README.md
Normal file
172
litellm-rust/crates/ai-gateway/README.md
Normal file
|
|
@ -0,0 +1,172 @@
|
|||
# LiteLLM Rust AI Gateway
|
||||
|
||||
A minimal Axum service that fronts OpenAI's realtime API. Clients open a
|
||||
WebSocket to `GET /v1/realtime`; the gateway authenticates, selects a deployment,
|
||||
dials OpenAI upstream, and splices the two sockets frame-by-frame.
|
||||
|
||||
- **Client endpoint:** `wss://<host>/v1/realtime?model=<model>` (WebSocket)
|
||||
- **Auth:** `Authorization: Bearer $LITELLM_MASTER_KEY` (fails closed if unset)
|
||||
- **Health:** `GET /health/readiness`, `GET /health/liveness`, `GET /health/gil`
|
||||
|
||||
> **Realtime serving is pure Rust.** Python is used at **load time only** — to
|
||||
> read the config once at boot. The realtime hot path never touches Python.
|
||||
|
||||
## Configuration (config.yaml)
|
||||
|
||||
The gateway loads its `model_list` from a **config.yaml**, the same as the
|
||||
LiteLLM proxy. Point `LITELLM_CONFIG_PATH` at the file:
|
||||
|
||||
```yaml
|
||||
# config.yaml
|
||||
model_list:
|
||||
- model_name: gpt-realtime
|
||||
litellm_params:
|
||||
model: openai/gpt-realtime
|
||||
api_key: os.environ/OPENAI_API_KEY
|
||||
```
|
||||
|
||||
```bash
|
||||
LITELLM_CONFIG_PATH=./config.yaml ./litellm-ai-gateway
|
||||
```
|
||||
|
||||
At boot the gateway calls into `litellm.proxy.read_model_list`, which reuses the
|
||||
**real proxy config reader** (`ProxyConfig.get_config`). That means everything
|
||||
the proxy supports in config.yaml works here too:
|
||||
|
||||
- `include:` to merge in other config files,
|
||||
- `os.environ/VAR` secret references (resolved via the secret manager, never
|
||||
inlined),
|
||||
- DB-stored models (when a database is configured).
|
||||
|
||||
Secrets stay out of the config — reference them with `os.environ/...` and set
|
||||
the env var at deploy time. The shipped Docker image is built with the
|
||||
`python-config` feature and **bundles litellm**, so config loading works out of
|
||||
the box; the default baked config lives at `/app/config.yaml` and can be
|
||||
overridden at deploy time (e.g. a Render secret file mounted at the same path).
|
||||
|
||||
### Environment variables
|
||||
|
||||
| Var | Required | Default | Purpose |
|
||||
|---|---|---|---|
|
||||
| `LITELLM_CONFIG_PATH` | yes (config mode) | — | Path to the config.yaml the gateway loads its `model_list` from. The Docker image defaults this to `/app/config.yaml`. |
|
||||
| `LITELLM_MASTER_KEY` | yes | — | Bearer token clients must send. Unset ⇒ all `/v1/realtime` requests are rejected (fail closed). |
|
||||
| `OPENAI_API_KEY` | yes | — | Upstream OpenAI key. Referenced by config.yaml as `os.environ/OPENAI_API_KEY` for the gateway→OpenAI dial. |
|
||||
| `HOST` | no | `127.0.0.1` | **Set to `0.0.0.0` in any container/deploy** or external traffic is refused. |
|
||||
| `PORT` | no | `4001` | Listen port. Render and most PaaS inject this automatically. |
|
||||
|
||||
> Secrets (`LITELLM_MASTER_KEY`, `OPENAI_API_KEY`) are never baked into the image
|
||||
> or `render.yaml` — inject them at deploy time only.
|
||||
|
||||
### Lean env stand-in (fallback)
|
||||
|
||||
If the binary is built **without** `python-config` (default features), or
|
||||
`LITELLM_CONFIG_PATH` is unset, the gateway falls back to a single-deployment
|
||||
stand-in built from the environment:
|
||||
|
||||
| Var | Default | Purpose |
|
||||
|---|---|---|
|
||||
| `OPENAI_REALTIME_MODEL` | `gpt-realtime` | The single deployment's model name (also the `?model=` clients pass). |
|
||||
|
||||
This mode links no libpython and needs no config file, but it only supports one
|
||||
hard-coded OpenAI deployment. **config.yaml is the recommended path** — use the
|
||||
stand-in only for the leanest possible build.
|
||||
|
||||
## Build & run with Docker
|
||||
|
||||
The image is built `--features python-config` and installs litellm **from this
|
||||
repo's source** (the config reader is newer than any PyPI release), so the build
|
||||
**context is the repo root**:
|
||||
|
||||
```bash
|
||||
# from the repo root
|
||||
docker build -f litellm-rust/crates/ai-gateway/Dockerfile -t litellm-ai-gateway .
|
||||
|
||||
docker run --rm -p 4001:4001 \
|
||||
-e HOST=0.0.0.0 -e PORT=4001 \
|
||||
-e LITELLM_MASTER_KEY=sk-local \
|
||||
-e OPENAI_API_KEY=$OPENAI_API_KEY \
|
||||
litellm-ai-gateway # LITELLM_CONFIG_PATH defaults to /app/config.yaml
|
||||
|
||||
# smoke test
|
||||
curl -s -o /dev/null -w '%{http_code}\n' localhost:4001/health/readiness # -> 200
|
||||
curl -s -o /dev/null -w '%{http_code}\n' localhost:4001/v1/realtime # -> 401 (auth fails closed)
|
||||
```
|
||||
|
||||
On boot you should see `loaded model_list from /app/config.yaml via python
|
||||
config reader` — that confirms the config path (not the env stand-in fallback).
|
||||
To use your own config, mount it over the default:
|
||||
|
||||
```bash
|
||||
docker run --rm -p 4001:4001 \
|
||||
-e HOST=0.0.0.0 -e LITELLM_MASTER_KEY=sk-local -e OPENAI_API_KEY=$OPENAI_API_KEY \
|
||||
-v $(pwd)/my-config.yaml:/app/config.yaml:ro \
|
||||
litellm-ai-gateway
|
||||
```
|
||||
|
||||
### Cargo-only (no Docker)
|
||||
|
||||
```bash
|
||||
# config.yaml mode — needs litellm importable in the active python env
|
||||
LITELLM_CONFIG_PATH=./crates/ai-gateway/config.yaml \
|
||||
cargo run --release -p litellm-ai-gateway --features python-config
|
||||
|
||||
# env stand-in mode — no python, no config
|
||||
cargo run --release -p litellm-ai-gateway
|
||||
```
|
||||
|
||||
## Deploy on Render
|
||||
|
||||
The service is a Docker **web service**; Render terminates TLS and supports
|
||||
WebSockets, so the public endpoint is `wss://<service>.onrender.com/v1/realtime`.
|
||||
|
||||
### Option A — Blueprint (`render.yaml`)
|
||||
|
||||
`crates/ai-gateway/render.yaml` describes the service (Docker runtime,
|
||||
`healthCheckPath: /health/readiness`, repo-root `dockerContext: .`,
|
||||
`dockerfilePath: ./litellm-rust/crates/ai-gateway/Dockerfile`,
|
||||
`LITELLM_CONFIG_PATH: /app/config.yaml`). `LITELLM_MASTER_KEY` and
|
||||
`OPENAI_API_KEY` are `sync: false` — set them in the dashboard after the first
|
||||
deploy. To use a non-default model_list, mount a **Render Secret File** at
|
||||
`/app/config.yaml`. Point a Render Blueprint at this repo/branch and apply.
|
||||
|
||||
### Option B — Render API
|
||||
|
||||
```bash
|
||||
# create a Docker web service from this repo+branch, then set env vars:
|
||||
curl -X POST https://api.render.com/v1/services \
|
||||
-H "Authorization: Bearer $RENDER_API_KEY" -H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"type": "web_service", "name": "litellm-rust-ai-gateway",
|
||||
"ownerId": "<owner-id>", "repo": "https://github.com/BerriAI/litellm",
|
||||
"branch": "<branch-with-this-dockerfile>",
|
||||
"serviceDetails": {
|
||||
"env": "docker",
|
||||
"envSpecificDetails": {
|
||||
"dockerfilePath": "./litellm-rust/crates/ai-gateway/Dockerfile",
|
||||
"dockerContext": "."
|
||||
},
|
||||
"healthCheckPath": "/health/readiness"
|
||||
}
|
||||
}'
|
||||
# then set env vars LITELLM_MASTER_KEY, OPENAI_API_KEY, HOST=0.0.0.0,
|
||||
# LITELLM_CONFIG_PATH=/app/config.yaml
|
||||
```
|
||||
|
||||
Health check path **must** be `/health/readiness`. `autoDeploy` is off by default
|
||||
in the blueprint — trigger deploys manually (or flip it on) to pick up new commits.
|
||||
|
||||
## Scaling
|
||||
|
||||
Concurrency is what matters, not total connections: each in-flight session holds
|
||||
one client socket + one upstream socket. To scale, raise the instance count /
|
||||
enable autoscaling on the Render service (e.g. baseline 10, max 100). Each
|
||||
instance needs file descriptors for `2 × peak_concurrent_sessions` — raise
|
||||
`ulimit -n` if you push very high concurrency.
|
||||
|
||||
## Latency note
|
||||
|
||||
The gateway adds the cost of one extra hop: client→gateway, then a fresh
|
||||
gateway→OpenAI realtime handshake (TLS + WS upgrade + `session.created`). In
|
||||
benchmarks this is ~100–150 ms of added session-establishment time; first-audio
|
||||
and steady-state streaming add no measurable overhead. To minimize it, deploy the
|
||||
gateway in the Render region with the lowest RTT to OpenAI's realtime endpoint.
|
||||
55
litellm-rust/crates/ai-gateway/benchmarks/realtime/README.md
Normal file
55
litellm-rust/crates/ai-gateway/benchmarks/realtime/README.md
Normal file
|
|
@ -0,0 +1,55 @@
|
|||
# Realtime gateway benchmark — pool on/off
|
||||
|
||||
Measures what the gateway adds over talking to OpenAI's realtime WebSocket
|
||||
directly, and what the pre-warmed connection pool removes. See
|
||||
`../../src/routes/realtime/README.md` for how the pool works.
|
||||
|
||||
## Results
|
||||
|
||||
5000 calls / 500 concurrency, gateway at 10 instances, pool ON
|
||||
(`REALTIME_POOL_SIZE=64`), upstream OpenAI `gpt-realtime`. Each leg run twice.
|
||||
Times in **ms**. Phases per connection: **dial** = TCP+TLS+WS upgrade,
|
||||
**session** = upgrade → `session.created` (the phase the pool removes),
|
||||
**1st-audio** = `response.create` → first audio delta (OpenAI inference),
|
||||
**total** = full wall-clock.
|
||||
|
||||
| metric | Direct OpenAI | Gateway (pool ON) | Overhead (ms) | vs OpenAI |
|
||||
| ------------------ | ------------- | ----------------- | ------------- | ---------- |
|
||||
| success rate (%) | 99.8 | 99.8 | — | — |
|
||||
| dial p50 (ms) | 276 | 158 | −118 | **faster** |
|
||||
| session p50 (ms) | 7 | 0 | −7 | **faster** |
|
||||
| 1st-audio p50 (ms) | 440 | 664 | +224 | slower¹ |
|
||||
| total p50 (ms) | 816 | 1010 | +194 | slower¹ |
|
||||
| total p95 (ms) | 2152 | 1970 | −182 | **faster** |
|
||||
| total p99 (ms) | 2692 | 2610 | −82 | **faster** |
|
||||
|
||||
The gateway is **faster than direct on 4 of 6 metrics**. The warm pool makes the
|
||||
**session phase sub-millisecond** at the median — ~76% of connects hit the pool,
|
||||
~70% had session < 1 ms. ¹ The two "slower" rows are not gateway overhead:
|
||||
`1st-audio` is OpenAI's own inference time (the gateway only relays it), which ran
|
||||
slower during the gateway legs and drags `total p50` with it.
|
||||
|
||||
**Pool OFF** (control, `REALTIME_POOL_SIZE=0`): session p50 was **367 ms** — the
|
||||
fresh-dial overhead the pool removes.
|
||||
|
||||
## Reproduce
|
||||
|
||||
The load generator lives in a separate repo:
|
||||
**https://github.com/ishaan-berri/litellm-realtime-bench**
|
||||
|
||||
```bash
|
||||
git clone https://github.com/ishaan-berri/litellm-realtime-bench
|
||||
cd litellm-realtime-bench && go build -o wsbench .
|
||||
|
||||
# Direct to OpenAI (baseline)
|
||||
./wsbench -host api.openai.com -key "$OPENAI_API_KEY" -m gpt-realtime -n 5000 -c 500 -t 60
|
||||
|
||||
# Through the gateway — run once with pool ON, once with REALTIME_POOL_SIZE=0
|
||||
./wsbench -host <gateway-host> -key "$LITELLM_MASTER_KEY" -m gpt-realtime -n 5000 -c 500 -t 60
|
||||
```
|
||||
|
||||
Run the gateway with the env stand-in (`OPENAI_REALTIME_MODEL=gpt-realtime`,
|
||||
`OPENAI_API_KEY`, `LITELLM_MASTER_KEY`, `REALTIME_POOL_SIZE`, `HOST=0.0.0.0`). At
|
||||
500 concurrency over N instances, size the pool to `≈ 500 / N` per instance (64 was
|
||||
used here for 10 instances). The bench repo's README covers running 500-concurrency
|
||||
legs from a hosted multi-vCPU runner. **Never commit keys — pass them via `-key`.**
|
||||
13
litellm-rust/crates/ai-gateway/config.yaml
Normal file
13
litellm-rust/crates/ai-gateway/config.yaml
Normal file
|
|
@ -0,0 +1,13 @@
|
|||
# Sample realtime config for the LiteLLM Rust AI Gateway.
|
||||
#
|
||||
# The gateway loads this model_list at boot via the embedded python config
|
||||
# reader (litellm.proxy.read_model_list), which reuses the proxy's own reader —
|
||||
# so include:, os.environ/ secrets, and DB-stored models all work here too.
|
||||
#
|
||||
# Secrets are referenced (never inlined) via os.environ/. A real deploy can
|
||||
# override this file (e.g. mount a Render secret file at LITELLM_CONFIG_PATH).
|
||||
model_list:
|
||||
- model_name: gpt-realtime
|
||||
litellm_params:
|
||||
model: openai/gpt-realtime
|
||||
api_key: os.environ/OPENAI_API_KEY
|
||||
35
litellm-rust/crates/ai-gateway/render.yaml
Normal file
35
litellm-rust/crates/ai-gateway/render.yaml
Normal file
|
|
@ -0,0 +1,35 @@
|
|||
# Render blueprint for the LiteLLM Rust AI Gateway (realtime WebSocket proxy).
|
||||
#
|
||||
# Single instance for now (no autoscaling). The public endpoint is a
|
||||
# WebSocket served over TLS: wss://<service>.onrender.com/v1/realtime
|
||||
#
|
||||
# Paths are relative to the **repo root** (Render's convention). The build
|
||||
# context is the repo root so the image can install litellm from source — the
|
||||
# gateway loads its model_list via litellm.proxy.read_model_list at boot.
|
||||
#
|
||||
# Secrets (LITELLM_MASTER_KEY, OPENAI_API_KEY) are marked sync: false — set
|
||||
# them in the Render dashboard or via the API, never inline here.
|
||||
services:
|
||||
- type: web
|
||||
name: litellm-rust-ai-gateway
|
||||
runtime: docker
|
||||
plan: standard
|
||||
dockerfilePath: ./litellm-rust/crates/ai-gateway/Dockerfile
|
||||
dockerContext: .
|
||||
healthCheckPath: /health/readiness
|
||||
numInstances: 1
|
||||
envVars:
|
||||
# The gateway loads its model_list from this config.yaml via the embedded
|
||||
# python config reader. The image bakes a default config at /app/config.yaml;
|
||||
# a real deploy can override it by mounting a Render secret file at this
|
||||
# same path (Dashboard → Environment → Secret Files) — never inline secrets.
|
||||
- key: LITELLM_CONFIG_PATH
|
||||
value: /app/config.yaml
|
||||
- key: HOST
|
||||
value: 0.0.0.0
|
||||
# Bearer token clients must send on /v1/realtime (fail closed if unset).
|
||||
- key: LITELLM_MASTER_KEY
|
||||
sync: false
|
||||
# Referenced by config.yaml as os.environ/OPENAI_API_KEY for the upstream dial.
|
||||
- key: OPENAI_API_KEY
|
||||
sync: false
|
||||
54
litellm-rust/crates/ai-gateway/src/auth/mod.rs
Normal file
54
litellm-rust/crates/ai-gateway/src/auth/mod.rs
Normal file
|
|
@ -0,0 +1,54 @@
|
|||
//! Gateway authentication, as an axum **extractor** (the idiomatic pattern —
|
||||
//! keeps handlers clean and auth testable).
|
||||
//!
|
||||
//! For now this is a single **master key**: any caller presenting it as
|
||||
//! `Authorization: Bearer <key>` may invoke the gateway. Per-key auth, budgets,
|
||||
//! and rate limits are delegated to the Python proxy in a later phase.
|
||||
//!
|
||||
//! A handler opts in by adding [`RequireMasterKey`] to its arguments; auth then
|
||||
//! runs during extraction, before the handler body. Routes never re-implement it.
|
||||
|
||||
use axum::extract::FromRequestParts;
|
||||
use axum::http::header::AUTHORIZATION;
|
||||
use axum::http::request::Parts;
|
||||
use axum::http::StatusCode;
|
||||
use subtle::ConstantTimeEq;
|
||||
|
||||
use crate::state::AppState;
|
||||
|
||||
/// Extractor that requires the configured master key as a bearer token.
|
||||
///
|
||||
/// Rejections: `500` when no master key is configured (permanent
|
||||
/// misconfiguration, not a transient outage); `401` on a missing/incorrect
|
||||
/// token. The comparison is constant-time.
|
||||
pub struct RequireMasterKey;
|
||||
|
||||
#[axum::async_trait]
|
||||
impl FromRequestParts<AppState> for RequireMasterKey {
|
||||
type Rejection = (StatusCode, String);
|
||||
|
||||
async fn from_request_parts(
|
||||
parts: &mut Parts,
|
||||
state: &AppState,
|
||||
) -> Result<Self, Self::Rejection> {
|
||||
let Some(expected) = state.master_key.as_deref() else {
|
||||
return Err((
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
"gateway auth not configured (set LITELLM_MASTER_KEY)".to_string(),
|
||||
));
|
||||
};
|
||||
let provided = parts
|
||||
.headers
|
||||
.get(AUTHORIZATION)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.and_then(|value| value.strip_prefix("Bearer "))
|
||||
.map(str::trim);
|
||||
match provided {
|
||||
Some(token) if bool::from(token.as_bytes().ct_eq(expected.as_bytes())) => Ok(Self),
|
||||
_ => Err((
|
||||
StatusCode::UNAUTHORIZED,
|
||||
"missing or invalid bearer token".to_string(),
|
||||
)),
|
||||
}
|
||||
}
|
||||
}
|
||||
58
litellm-rust/crates/ai-gateway/src/gil.rs
Normal file
58
litellm-rust/crates/ai-gateway/src/gil.rs
Normal file
|
|
@ -0,0 +1,58 @@
|
|||
//! GIL-activity tracking.
|
||||
//!
|
||||
//! Every acquisition of the Python GIL is recorded here so the `/health/gil`
|
||||
//! endpoint can report whether Python was touched recently. The design goal is
|
||||
//! that the GIL is acquired **only at load time** (config read) and never on the
|
||||
//! realtime hot path — polling this endpoint during traffic should show the
|
||||
//! count holding steady and `acquired_last_30s` falling to `false`.
|
||||
|
||||
use std::sync::atomic::{AtomicU64, Ordering};
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
|
||||
/// Window (seconds) for the "recently acquired" signal.
|
||||
pub const RECENT_WINDOW_SECS: u64 = 30;
|
||||
|
||||
static GIL_ACQUISITIONS: AtomicU64 = AtomicU64::new(0);
|
||||
/// Unix seconds of the last acquisition; `0` means "never".
|
||||
static LAST_GIL_UNIX_SECS: AtomicU64 = AtomicU64::new(0);
|
||||
|
||||
fn now_unix_secs() -> u64 {
|
||||
SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.map(|d| d.as_secs())
|
||||
.unwrap_or(0)
|
||||
}
|
||||
|
||||
/// Record that the GIL was just acquired. Call immediately before taking the GIL.
|
||||
///
|
||||
/// Only invoked under the `python-config` feature; without it the gateway never
|
||||
/// touches Python, so the recorder is unused (and the endpoint reports zero).
|
||||
#[cfg_attr(not(feature = "python-config"), allow(dead_code))]
|
||||
pub fn record_acquisition() {
|
||||
GIL_ACQUISITIONS.fetch_add(1, Ordering::Relaxed);
|
||||
LAST_GIL_UNIX_SECS.store(now_unix_secs(), Ordering::Relaxed);
|
||||
}
|
||||
|
||||
/// Point-in-time view of GIL activity.
|
||||
pub struct GilSnapshot {
|
||||
pub total_acquisitions: u64,
|
||||
pub seconds_since_last: Option<u64>,
|
||||
pub acquired_last_30s: bool,
|
||||
}
|
||||
|
||||
/// Read the current GIL-activity snapshot.
|
||||
pub fn snapshot() -> GilSnapshot {
|
||||
let total = GIL_ACQUISITIONS.load(Ordering::Relaxed);
|
||||
let last = LAST_GIL_UNIX_SECS.load(Ordering::Relaxed);
|
||||
let seconds_since_last = if last == 0 {
|
||||
None
|
||||
} else {
|
||||
Some(now_unix_secs().saturating_sub(last))
|
||||
};
|
||||
let acquired_last_30s = seconds_since_last.is_some_and(|secs| secs <= RECENT_WINDOW_SECS);
|
||||
GilSnapshot {
|
||||
total_acquisitions: total,
|
||||
seconds_since_last,
|
||||
acquired_last_30s,
|
||||
}
|
||||
}
|
||||
152
litellm-rust/crates/ai-gateway/src/main.rs
Normal file
152
litellm-rust/crates/ai-gateway/src/main.rs
Normal file
|
|
@ -0,0 +1,152 @@
|
|||
//! LiteLLM AI Gateway — a minimal Axum server fronting the Rust router.
|
||||
//!
|
||||
//! Flow: client → `POST /v1/realtime` → `router.realtime()` selects a deployment
|
||||
//! (simple-shuffle) → `providers::realtime::realtime()` invokes OpenAI. The
|
||||
//! server owns transport + config; routing lives in the `router` crate.
|
||||
|
||||
mod auth;
|
||||
mod gil;
|
||||
#[cfg(feature = "python-config")]
|
||||
mod python;
|
||||
mod routes;
|
||||
mod state;
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use litellm_core::router::{Deployment, LiteLLMParams, Router};
|
||||
use litellm_providers::realtime_pool::{upstream_key, PoolConfig, RealtimePool};
|
||||
|
||||
use crate::state::AppState;
|
||||
|
||||
/// Bind to localhost by default so the gateway is not a public, unauthenticated
|
||||
/// provider proxy out of the box. Override with `HOST` (e.g. `0.0.0.0`).
|
||||
const DEFAULT_HOST: &str = "127.0.0.1";
|
||||
const DEFAULT_PORT: u16 = 4001;
|
||||
|
||||
#[tokio::main]
|
||||
async fn main() {
|
||||
// Trim before storing so it matches the trimmed bearer token in `auth`
|
||||
// (avoids a silent auth failure when the env var has surrounding whitespace).
|
||||
let master_key: Option<Arc<str>> = std::env::var("LITELLM_MASTER_KEY")
|
||||
.ok()
|
||||
.map(|key| key.trim().to_string())
|
||||
.filter(|key| !key.is_empty())
|
||||
.map(Arc::from);
|
||||
if master_key.is_none() {
|
||||
eprintln!(
|
||||
"warning: LITELLM_MASTER_KEY is not set; /v1/realtime will reject all requests (fail closed)"
|
||||
);
|
||||
}
|
||||
|
||||
let router = Arc::new(build_router());
|
||||
|
||||
// Build the pre-warmed realtime pool and register each deployment's upstream
|
||||
// so the background replenisher starts warming it. `REALTIME_POOL_SIZE=0`
|
||||
// yields a disabled pool → every connect fresh-dials (original behavior).
|
||||
let pool_config = PoolConfig::from_env();
|
||||
let realtime_pool = RealtimePool::spawn(pool_config);
|
||||
if pool_config.enabled() {
|
||||
register_deployments(&router, &realtime_pool);
|
||||
eprintln!(
|
||||
"realtime connection pool enabled: target {} warm sockets/key, max idle {}s",
|
||||
pool_config.target_size,
|
||||
pool_config.max_idle.as_secs()
|
||||
);
|
||||
} else {
|
||||
eprintln!(
|
||||
"realtime connection pool disabled (REALTIME_POOL_SIZE=0); fresh-dialing each connect"
|
||||
);
|
||||
}
|
||||
|
||||
let state = AppState {
|
||||
router,
|
||||
master_key,
|
||||
realtime_pool,
|
||||
};
|
||||
|
||||
let host = std::env::var("HOST").unwrap_or_else(|_| DEFAULT_HOST.to_string());
|
||||
let port = resolve_port();
|
||||
|
||||
let listener = tokio::net::TcpListener::bind((host.as_str(), port))
|
||||
.await
|
||||
.expect("failed to bind listener");
|
||||
eprintln!("litellm-ai-gateway listening on {host}:{port}");
|
||||
axum::serve(listener, routes::app(state))
|
||||
.await
|
||||
.expect("server error");
|
||||
}
|
||||
|
||||
/// Register every deployment's upstream key with the pool so the replenisher
|
||||
/// pre-warms it. Mirrors `service::run`'s key derivation (strip `openai/`, resolve
|
||||
/// api_key); deployments whose key can't be resolved are skipped (they fresh-dial
|
||||
/// and surface the auth error on the request path, as before).
|
||||
fn register_deployments(router: &Router, pool: &RealtimePool) {
|
||||
for deployment in router.deployments() {
|
||||
let params = &deployment.litellm_params;
|
||||
let provider_model = params
|
||||
.model
|
||||
.strip_prefix("openai/")
|
||||
.unwrap_or(¶ms.model);
|
||||
if let Some(key) = upstream_key(
|
||||
provider_model,
|
||||
params.api_key.as_deref(),
|
||||
params.api_base.as_deref(),
|
||||
) {
|
||||
pool.register(key);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Resolve `PORT`, warning (rather than silently defaulting) on an invalid value.
|
||||
fn resolve_port() -> u16 {
|
||||
match std::env::var("PORT") {
|
||||
Ok(raw) => raw.parse().unwrap_or_else(|_| {
|
||||
eprintln!("warning: PORT={raw:?} is not a valid port; using {DEFAULT_PORT}");
|
||||
DEFAULT_PORT
|
||||
}),
|
||||
Err(_) => DEFAULT_PORT,
|
||||
}
|
||||
}
|
||||
|
||||
/// Build the router. With the `python-config` feature and `LITELLM_CONFIG_PATH`
|
||||
/// set, load the resolved `model_list` from the proxy config via the embedded
|
||||
/// Python reader (load time only). Otherwise fall back to the env stand-in.
|
||||
fn build_router() -> Router {
|
||||
#[cfg(feature = "python-config")]
|
||||
if let Ok(config_path) = std::env::var("LITELLM_CONFIG_PATH") {
|
||||
match python::config::load_router_from_config(&config_path) {
|
||||
Ok(router) => {
|
||||
eprintln!("loaded model_list from {config_path} via python config reader");
|
||||
return router;
|
||||
}
|
||||
Err(err) => {
|
||||
eprintln!("config load failed ({err}); falling back to env deployment");
|
||||
}
|
||||
}
|
||||
}
|
||||
build_router_from_env()
|
||||
}
|
||||
|
||||
/// Build a minimal single-deployment `model_list` from the environment.
|
||||
///
|
||||
/// A real deployment loads `model_list` from config; this is the minimal stand-in
|
||||
/// so the gateway has one OpenAI deployment to route to.
|
||||
fn build_router_from_env() -> Router {
|
||||
let model =
|
||||
std::env::var("OPENAI_REALTIME_MODEL").unwrap_or_else(|_| "gpt-realtime".to_string());
|
||||
let api_key = std::env::var("OPENAI_API_KEY").ok();
|
||||
if api_key.is_none() {
|
||||
eprintln!(
|
||||
"warning: OPENAI_API_KEY is not set; realtime requests will fail with auth errors"
|
||||
);
|
||||
}
|
||||
let deployment = Deployment {
|
||||
model_name: model.clone(),
|
||||
litellm_params: LiteLLMParams {
|
||||
model,
|
||||
api_key,
|
||||
api_base: None,
|
||||
},
|
||||
};
|
||||
Router::new(vec![deployment])
|
||||
}
|
||||
27
litellm-rust/crates/ai-gateway/src/python/AGENTS.md
Normal file
27
litellm-rust/crates/ai-gateway/src/python/AGENTS.md
Normal file
|
|
@ -0,0 +1,27 @@
|
|||
# ai-gateway/src/python — Python interop (load-time only)
|
||||
|
||||
Functions here embed the Python interpreter (pyo3) and take the GIL to call into
|
||||
`litellm` (e.g. read the proxy `model_list`). Compiled only under the
|
||||
`python-config` feature.
|
||||
|
||||
## Hard rule: non-hot-path functions only
|
||||
|
||||
Everything in this folder MUST run **at most once per process lifetime — at
|
||||
startup / load time** (config read, warm-up). NEVER call into Python on the
|
||||
request path:
|
||||
|
||||
- No GIL acquisition per request, per connection, or per realtime event.
|
||||
- No Python call inside a route handler, the router's hot path, or any loop that
|
||||
scales with traffic.
|
||||
|
||||
**Why:** the GIL serializes execution and would cap throughput; the realtime data
|
||||
path must stay pure Rust. Every acquisition is recorded by `crate::gil` — poll
|
||||
`GET /health/gil`, and `total_acquisitions` MUST stay flat under load.
|
||||
|
||||
## How to add one
|
||||
|
||||
Resolve whatever Python-derived data you need **once at boot** and hand the rest
|
||||
of the gateway an owned, plain-Rust value (e.g. build a `Router` from the
|
||||
resolved `model_list`). Record the acquisition via `crate::gil::record_acquisition()`
|
||||
immediately before taking the GIL. If a function would need to run per request,
|
||||
it does not belong here — move the work to Rust, or pre-resolve it at startup.
|
||||
39
litellm-rust/crates/ai-gateway/src/python/config.rs
Normal file
39
litellm-rust/crates/ai-gateway/src/python/config.rs
Normal file
|
|
@ -0,0 +1,39 @@
|
|||
//! Build the router by calling the Python proxy config reader (load time only).
|
||||
//!
|
||||
//! Embeds the interpreter via pyo3 and calls
|
||||
//! `litellm.proxy.read_model_list.read_model_list`, which reuses the proxy's
|
||||
//! `os.environ/` + secret-manager resolution. The GIL is taken **once at boot**
|
||||
//! (and recorded in [`crate::gil`]); the realtime hot path never touches Python.
|
||||
//!
|
||||
//! Compiled only under the `python-config` feature.
|
||||
|
||||
use litellm_core::error::CoreError;
|
||||
use litellm_core::router::{Deployment, Router};
|
||||
use litellm_core::CoreResult;
|
||||
use pyo3::prelude::*;
|
||||
|
||||
use crate::gil;
|
||||
|
||||
/// Load the router's `model_list` from `config_path` via the Python reader.
|
||||
pub fn load_router_from_config(config_path: &str) -> CoreResult<Router> {
|
||||
gil::record_acquisition();
|
||||
Python::with_gil(|py| {
|
||||
let model_list = py
|
||||
.import("litellm.proxy.read_model_list")
|
||||
.and_then(|module| module.getattr("read_model_list"))
|
||||
.and_then(|reader| reader.call1((config_path,)))
|
||||
.map_err(|err| CoreError::Routing(format!("read_model_list failed: {err}")))?;
|
||||
|
||||
let model_list_json: String = py
|
||||
.import("json")
|
||||
.and_then(|json| json.getattr("dumps"))
|
||||
.and_then(|dumps| dumps.call1((model_list,)))
|
||||
.and_then(|encoded| encoded.extract())
|
||||
.map_err(|err| CoreError::Routing(format!("serializing model_list failed: {err}")))?;
|
||||
|
||||
let deployments: Vec<Deployment> = serde_json::from_str(&model_list_json)
|
||||
.map_err(|err| CoreError::Routing(format!("parsing model_list failed: {err}")))?;
|
||||
|
||||
Ok(Router::new(deployments))
|
||||
})
|
||||
}
|
||||
4
litellm-rust/crates/ai-gateway/src/python/mod.rs
Normal file
4
litellm-rust/crates/ai-gateway/src/python/mod.rs
Normal file
|
|
@ -0,0 +1,4 @@
|
|||
//! Python interop for the gateway. See `AGENTS.md`: **load-time / non-hot-path
|
||||
//! only.** Compiled only under the `python-config` feature.
|
||||
|
||||
pub mod config;
|
||||
38
litellm-rust/crates/ai-gateway/src/routes/AGENTS.md
Normal file
38
litellm-rust/crates/ai-gateway/src/routes/AGENTS.md
Normal file
|
|
@ -0,0 +1,38 @@
|
|||
# routes/ — the route template
|
||||
|
||||
Every route follows the **same shape** so the layout is predictable. The rule:
|
||||
|
||||
> **Each route module exposes `pub fn router() -> Router<AppState>`.**
|
||||
> `routes/mod.rs::app` merges them all and applies state once. Adding a route is:
|
||||
> create the module, then add one `.merge(<name>::router())` line.
|
||||
|
||||
## Default: one file
|
||||
A route is a single file containing `router()` + its handler(s) (handlers stay
|
||||
private). This is the norm — don't split until it hurts.
|
||||
```
|
||||
pub fn router() -> Router<AppState> { Router::new().route(PATH, get(handle)) }
|
||||
async fn handle(...) -> impl IntoResponse { ... }
|
||||
```
|
||||
`health.rs` and `gil.rs` are examples.
|
||||
|
||||
## Split out `service` when there's real logic
|
||||
When a route has business logic worth testing without axum, put it in a sibling
|
||||
`service` (a file, or a folder if the route grows). The route file stays the
|
||||
**axum surface** (router + handler + any socket/SSE adapter); `service` is plain
|
||||
Rust with **no axum types**. `realtime/` is the example:
|
||||
```
|
||||
realtime/
|
||||
mod.rs # axum surface: router() + handler + the WS<->events adapter
|
||||
service.rs # pure logic: select deployment + call provider (no axum) — testable
|
||||
```
|
||||
Split `service` further (or add `transport`, `repo`, …) only once a single file
|
||||
genuinely gets hard to read.
|
||||
|
||||
## Invariants
|
||||
- **Auth is an extractor, not a manual call.** A handler requires auth by adding
|
||||
`crate::auth::RequireMasterKey` to its arguments; it runs during extraction.
|
||||
Never re-implement the check per route.
|
||||
- **Handlers contain no business logic; `service` contains no axum types.**
|
||||
- A route owns its paths in its own `router()`; `mod.rs` only merges.
|
||||
- Cross-cutting concerns (logging, CORS, timeouts) → Tower layers in `mod.rs`,
|
||||
not duplicated in handlers.
|
||||
30
litellm-rust/crates/ai-gateway/src/routes/gil.rs
Normal file
30
litellm-rust/crates/ai-gateway/src/routes/gil.rs
Normal file
|
|
@ -0,0 +1,30 @@
|
|||
//! `GET /health/gil` — poll to confirm Python is only touched at load time.
|
||||
//! Simple-route template: a `router()` plus its handler, in one file.
|
||||
|
||||
use axum::routing::get;
|
||||
use axum::{Json, Router};
|
||||
use serde::Serialize;
|
||||
|
||||
use crate::gil;
|
||||
use crate::state::AppState;
|
||||
|
||||
/// This route's contribution to the app router.
|
||||
pub fn router() -> Router<AppState> {
|
||||
Router::new().route("/health/gil", get(status))
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
struct GilStatusResponse {
|
||||
gil_acquired_last_30s: bool,
|
||||
total_acquisitions: u64,
|
||||
seconds_since_last: Option<u64>,
|
||||
}
|
||||
|
||||
async fn status() -> Json<GilStatusResponse> {
|
||||
let snapshot = gil::snapshot();
|
||||
Json(GilStatusResponse {
|
||||
gil_acquired_last_30s: snapshot.acquired_last_30s,
|
||||
total_acquisitions: snapshot.total_acquisitions,
|
||||
seconds_since_last: snapshot.seconds_since_last,
|
||||
})
|
||||
}
|
||||
24
litellm-rust/crates/ai-gateway/src/routes/health.rs
Normal file
24
litellm-rust/crates/ai-gateway/src/routes/health.rs
Normal file
|
|
@ -0,0 +1,24 @@
|
|||
//! Health probes. Simple-route template: a `router()` plus its handlers, in one file.
|
||||
|
||||
use axum::http::StatusCode;
|
||||
use axum::routing::get;
|
||||
use axum::Router;
|
||||
|
||||
use crate::state::AppState;
|
||||
|
||||
/// This route's contribution to the app router.
|
||||
pub fn router() -> Router<AppState> {
|
||||
Router::new()
|
||||
.route("/health/liveness", get(liveness))
|
||||
.route("/health/readiness", get(readiness))
|
||||
}
|
||||
|
||||
/// The process is up.
|
||||
async fn liveness() -> StatusCode {
|
||||
StatusCode::OK
|
||||
}
|
||||
|
||||
/// The server is ready to accept traffic.
|
||||
async fn readiness() -> StatusCode {
|
||||
StatusCode::OK
|
||||
}
|
||||
23
litellm-rust/crates/ai-gateway/src/routes/mod.rs
Normal file
23
litellm-rust/crates/ai-gateway/src/routes/mod.rs
Normal file
|
|
@ -0,0 +1,23 @@
|
|||
//! HTTP routes.
|
||||
//!
|
||||
//! **Template:** every route module exposes `pub fn router() -> Router<AppState>`
|
||||
//! that mounts its own paths; [`app`] merges them. A trivial route is a single
|
||||
//! file (`health.rs`, `gil.rs`); a non-trivial one is a folder (`realtime/`) with
|
||||
//! `handler` (entry) + `service` (logic) + `transport` (adapters). See AGENTS.md.
|
||||
|
||||
pub mod gil;
|
||||
pub mod health;
|
||||
pub mod realtime;
|
||||
|
||||
use axum::Router;
|
||||
|
||||
use crate::state::AppState;
|
||||
|
||||
/// Assemble the application router by merging every route module's `router()`.
|
||||
pub fn app(state: AppState) -> Router {
|
||||
Router::new()
|
||||
.merge(health::router())
|
||||
.merge(gil::router())
|
||||
.merge(realtime::router())
|
||||
.with_state(state)
|
||||
}
|
||||
87
litellm-rust/crates/ai-gateway/src/routes/realtime/README.md
Normal file
87
litellm-rust/crates/ai-gateway/src/routes/realtime/README.md
Normal file
|
|
@ -0,0 +1,87 @@
|
|||
# Realtime route (`GET /v1/realtime`)
|
||||
|
||||
Proxies OpenAI's realtime WebSocket. `mod.rs` is the axum surface (handler +
|
||||
socket↔events adapter); `service.rs` is the pure logic (select a deployment, then
|
||||
splice client ↔ upstream). The pool itself lives in
|
||||
`crates/providers/src/realtime_pool.rs`.
|
||||
|
||||
## Connection pooling
|
||||
|
||||
### The problem
|
||||
|
||||
The gateway's realtime overhead lives **entirely in session establishment**. On each
|
||||
client connect it dials a *fresh* upstream WS to OpenAI and waits for
|
||||
`session.created` before it can serve. Measured at 5000 calls / 500 concurrency, the
|
||||
fresh-dial session phase is **~360 ms** vs **~7 ms** direct; dial, first-audio, and
|
||||
streaming add ~0. So the one lever is removing that per-connect handshake from the
|
||||
critical path.
|
||||
|
||||
### The idea
|
||||
|
||||
Keep a few upstream OpenAI sockets **already connected and already past
|
||||
`session.created`** (buffered). On a client connect, hand off a warm socket — relay
|
||||
its buffered `session.created` instantly (a local `Vec::pop`, sub-millisecond) and
|
||||
splice exactly as a fresh dial would. A background task keeps the pool topped up. On
|
||||
a miss or dead socket we fall back to fresh-dial: the pool is a latency optimization,
|
||||
never a correctness dependency.
|
||||
|
||||
```
|
||||
┌───────────────────────────────────────┐
|
||||
client connect ──────► │ routes/realtime → service::run │
|
||||
│ pool.take(key) │
|
||||
│ hit → relay buffered │
|
||||
│ session.created, then splice │
|
||||
│ miss → fresh dial (original path) │
|
||||
└───────────────┬───────────────────────┘
|
||||
│ replenish (async, concurrent)
|
||||
┌───────────────▼───────────────────────┐
|
||||
background task ─────► │ RealtimePool: per-key warm sockets │
|
||||
│ each = { ws, buffered session.created}│
|
||||
│ liveness-checked before handoff │
|
||||
└─────────────────────────────────────────┘
|
||||
```
|
||||
|
||||
A warm session is indistinguishable from a fresh one: OpenAI sends `session.created`
|
||||
unprompted on connect, we pre-read exactly that one frame and relay it on handoff,
|
||||
and we send nothing else on the socket before a client exists — so the client's first
|
||||
`session.update` behaves identically either way.
|
||||
|
||||
### Sizing
|
||||
|
||||
Each warm socket serves **exactly one** session (realtime isn't multiplexed), so the
|
||||
pool is sized to the **peak concurrent connects per instance**, not total live
|
||||
connections:
|
||||
|
||||
```
|
||||
REALTIME_POOL_SIZE ≈ peak_concurrency / instance_count
|
||||
```
|
||||
|
||||
e.g. 500 concurrency over 10 instances → ~50–64 per instance. The replenisher dials
|
||||
the missing sockets **concurrently**, so a drained pool refills in ~one handshake
|
||||
window and keeps supply close to the connect rate. Over-provisioning just burns idle
|
||||
upstream sockets, which is why warm sockets are short-lived
|
||||
(`REALTIME_POOL_MAX_IDLE_SECS`).
|
||||
|
||||
### Config
|
||||
|
||||
| env | default | meaning |
|
||||
| ----------------------------- | ------- | --------------------------------------------------------------- |
|
||||
| `REALTIME_POOL_SIZE` | `4` | target warm sockets per key. `0` disables pooling (fresh-dial). |
|
||||
| `REALTIME_POOL_MAX_IDLE_SECS` | `30` | max time a warm socket sits before it's closed and replaced. |
|
||||
|
||||
### Notes
|
||||
|
||||
- **Miss / dead socket → fresh dial.** Burst beyond warm supply, or a socket that
|
||||
died, never blocks or fails — it falls back to the original path. The pool can only
|
||||
make a connect faster, never slower or more fragile.
|
||||
- **Auth scope.** The pool key includes `api_key`, so a warm socket is only handed to
|
||||
a request resolving to the same key — no cross-tenant reuse.
|
||||
- **Idle billing.** Warm sockets are liveness-checked at handoff and capped at
|
||||
`REALTIME_POOL_MAX_IDLE_SECS` to bound idle billing and dodge OpenAI's idle timeout.
|
||||
- **Replenish backoff.** If a key's warm-up dials all fail (invalid credentials, an
|
||||
unreachable upstream), the replenisher puts that key into exponential backoff
|
||||
(500 ms → 30 s cap) instead of re-dialing it every tick. This bounds connection
|
||||
attempts against a broken key so it can't exhaust upstream rate limits and degrade
|
||||
valid cold-path traffic; the backoff resets the moment a dial succeeds.
|
||||
|
||||
Benchmarks and repro: `../../benchmarks/realtime/README.md`.
|
||||
88
litellm-rust/crates/ai-gateway/src/routes/realtime/mod.rs
Normal file
88
litellm-rust/crates/ai-gateway/src/routes/realtime/mod.rs
Normal file
|
|
@ -0,0 +1,88 @@
|
|||
//! `GET /v1/realtime` (WebSocket).
|
||||
//!
|
||||
//! This file is the **axum surface**: `router()`, the handler, and the small
|
||||
//! socket↔events adapter. The pure logic (no axum) lives in [`service`]. Auth is
|
||||
//! the `RequireMasterKey` extractor, so the handler stays thin.
|
||||
|
||||
mod service;
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use axum::extract::ws::{Message, WebSocket, WebSocketUpgrade};
|
||||
use axum::extract::{Query, State};
|
||||
use axum::http::StatusCode;
|
||||
use axum::response::Response;
|
||||
use axum::routing::get;
|
||||
use axum::Router;
|
||||
use futures_util::{SinkExt, StreamExt};
|
||||
use litellm_core::realtime::types::RealtimeEvent;
|
||||
use litellm_core::router::Router as ModelRouter;
|
||||
use litellm_providers::realtime_pool::RealtimePool;
|
||||
use serde::Deserialize;
|
||||
|
||||
use crate::auth::RequireMasterKey;
|
||||
use crate::state::AppState;
|
||||
|
||||
/// This route's contribution to the app router.
|
||||
pub fn router() -> Router<AppState> {
|
||||
Router::new().route("/v1/realtime", get(handle))
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct RealtimeQuery {
|
||||
model: String,
|
||||
}
|
||||
|
||||
/// Auth runs via the `RequireMasterKey` extractor. We validate the model BEFORE
|
||||
/// the upgrade so failures are clean HTTP (400/404), not a socket that opens then
|
||||
/// closes, then hand the socket to `bridge`.
|
||||
async fn handle(
|
||||
_auth: RequireMasterKey,
|
||||
ws: WebSocketUpgrade,
|
||||
State(state): State<AppState>,
|
||||
Query(query): Query<RealtimeQuery>,
|
||||
) -> Result<Response, (StatusCode, String)> {
|
||||
if query.model.trim().is_empty() {
|
||||
return Err((
|
||||
StatusCode::BAD_REQUEST,
|
||||
"missing 'model' query param".to_string(),
|
||||
));
|
||||
}
|
||||
if !state.router.has_deployment(&query.model) {
|
||||
return Err((
|
||||
StatusCode::NOT_FOUND,
|
||||
format!("no deployment for model '{}'", query.model),
|
||||
));
|
||||
}
|
||||
|
||||
let router = state.router.clone();
|
||||
let pool = state.realtime_pool.clone();
|
||||
let model = query.model;
|
||||
Ok(ws.on_upgrade(move |socket| bridge(socket, router, pool, model)))
|
||||
}
|
||||
|
||||
/// Adapt the axum socket (text frames) to the typed-event `Stream`/`Sink` the
|
||||
/// service wants, keeping axum types out of `service`.
|
||||
async fn bridge(
|
||||
socket: WebSocket,
|
||||
router: Arc<ModelRouter>,
|
||||
pool: Arc<RealtimePool>,
|
||||
model: String,
|
||||
) {
|
||||
let (ws_sink, ws_stream) = socket.split();
|
||||
|
||||
let client_in = ws_stream.filter_map(|message| async move {
|
||||
match message {
|
||||
Ok(Message::Text(text)) => serde_json::from_str::<RealtimeEvent>(&text).ok(),
|
||||
_ => None,
|
||||
}
|
||||
});
|
||||
let client_out = ws_sink.with(|event: RealtimeEvent| async move {
|
||||
Ok::<Message, axum::Error>(Message::Text(
|
||||
serde_json::to_string(&event).unwrap_or_default(),
|
||||
))
|
||||
});
|
||||
|
||||
futures_util::pin_mut!(client_in, client_out);
|
||||
let _ = service::run(&router, &pool, &model, None, client_in, client_out).await;
|
||||
}
|
||||
|
|
@ -0,0 +1,76 @@
|
|||
//! Business logic: select a deployment with the (pure) core router, then call the
|
||||
//! provider splice. The seam between `core::router` (selection only) and
|
||||
//! `providers` (the actual WebSocket I/O).
|
||||
//!
|
||||
//! On connect we try a pre-warmed upstream from the pool (handshake already paid,
|
||||
//! `session.created` buffered) and relay it instantly. On a pool miss or dead warm
|
||||
//! socket we fresh-dial exactly as before — the pool is never on the critical path
|
||||
//! for correctness, only latency.
|
||||
|
||||
use std::time::Duration;
|
||||
|
||||
use futures_util::{Sink, Stream};
|
||||
use litellm_core::error::CoreError;
|
||||
use litellm_core::realtime::types::RealtimeEvent;
|
||||
use litellm_core::router::Router;
|
||||
use litellm_core::CoreResult;
|
||||
use litellm_providers::realtime_pool::{upstream_key, RealtimePool};
|
||||
|
||||
/// Select a deployment for `model` and splice the client stream to the provider.
|
||||
///
|
||||
/// `pool` supplies a pre-warmed upstream when one is available; otherwise we
|
||||
/// fresh-dial. A disabled pool always misses, so this collapses to the original
|
||||
/// fresh-dial behavior.
|
||||
pub async fn run<In, Out>(
|
||||
router: &Router,
|
||||
pool: &RealtimePool,
|
||||
model: &str,
|
||||
idle_timeout: Option<Duration>,
|
||||
client_in: In,
|
||||
client_out: Out,
|
||||
) -> CoreResult<()>
|
||||
where
|
||||
In: Stream<Item = RealtimeEvent> + Unpin + Send,
|
||||
Out: Sink<RealtimeEvent> + Unpin + Send,
|
||||
<Out as Sink<RealtimeEvent>>::Error: std::fmt::Display,
|
||||
{
|
||||
let deployment = router.get_available_deployment(model).ok_or_else(|| {
|
||||
CoreError::Routing(format!("no deployment available for model '{model}'"))
|
||||
})?;
|
||||
let params = &deployment.litellm_params;
|
||||
// Strip a leading `openai/` so the OpenAI-only realtime fn gets the bare model.
|
||||
let provider_model = params
|
||||
.model
|
||||
.strip_prefix("openai/")
|
||||
.unwrap_or(¶ms.model);
|
||||
|
||||
// Warm path: take a pooled upstream (handshake already paid) and relay its
|
||||
// buffered session.created immediately. On miss/dead socket fall through.
|
||||
if let Some(key) = upstream_key(
|
||||
provider_model,
|
||||
params.api_key.as_deref(),
|
||||
params.api_base.as_deref(),
|
||||
) {
|
||||
if let Some(handoff) = pool.take(&key) {
|
||||
return litellm_providers::realtime::realtime_warm(
|
||||
provider_model,
|
||||
handoff,
|
||||
idle_timeout,
|
||||
client_in,
|
||||
client_out,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
}
|
||||
|
||||
// Cold path: fresh dial (the original behavior).
|
||||
litellm_providers::realtime::realtime(
|
||||
provider_model,
|
||||
params.api_key.as_deref(),
|
||||
params.api_base.as_deref(),
|
||||
idle_timeout,
|
||||
client_in,
|
||||
client_out,
|
||||
)
|
||||
.await
|
||||
}
|
||||
17
litellm-rust/crates/ai-gateway/src/state.rs
Normal file
17
litellm-rust/crates/ai-gateway/src/state.rs
Normal file
|
|
@ -0,0 +1,17 @@
|
|||
use std::sync::Arc;
|
||||
|
||||
use litellm_core::router::Router;
|
||||
use litellm_providers::realtime_pool::RealtimePool;
|
||||
|
||||
/// Shared application state handed to every route handler.
|
||||
#[derive(Clone)]
|
||||
pub struct AppState {
|
||||
pub router: Arc<Router>,
|
||||
/// The gateway master key. Any caller presenting it as a bearer token may
|
||||
/// invoke the gateway. `None` → auth not configured (routes fail closed).
|
||||
pub master_key: Option<Arc<str>>,
|
||||
/// Pre-warmed upstream realtime connection pool. Disabled
|
||||
/// (`RealtimePool::disabled()`) when `REALTIME_POOL_SIZE=0`, in which case
|
||||
/// every realtime connect fresh-dials exactly as before.
|
||||
pub realtime_pool: Arc<RealtimePool>,
|
||||
}
|
||||
47
litellm-rust/crates/core/CLAUDE.md
Normal file
47
litellm-rust/crates/core/CLAUDE.md
Normal file
|
|
@ -0,0 +1,47 @@
|
|||
# CLAUDE.md
|
||||
|
||||
Rules for `litellm-rust/crates/core`.
|
||||
|
||||
## Responsibility
|
||||
|
||||
`core` owns shared data types, typed errors, and deterministic helper contracts.
|
||||
It must stay pure and host-independent.
|
||||
|
||||
Allowed:
|
||||
- Shared request/response structs.
|
||||
- Typed errors with stable, non-sensitive messages.
|
||||
- Deterministic validation helpers.
|
||||
- Serialization helpers that intentionally mirror Python output shape.
|
||||
- Route templates that match Python base config responsibilities, such as
|
||||
`ocr::transformation::OcrProviderConfig`.
|
||||
|
||||
Not allowed:
|
||||
- Network, filesystem, database, cache, or environment access.
|
||||
- Secret reads or auth/header construction.
|
||||
- Logging callbacks, tracing spans, spend writes, or customer callbacks.
|
||||
- Provider-specific branching that belongs in `providers`.
|
||||
- Panics for user/provider-controlled input.
|
||||
|
||||
## Typed Contracts (core rule)
|
||||
|
||||
Trait and function boundaries MUST be strongly typed. No stringly-typed JSON
|
||||
(`&str` / `String` / `Vec<String>` / bare `serde_json::Value`) as a transform
|
||||
input or output. Parse wire bytes into typed structs/enums at the host edge;
|
||||
`core` and `providers` operate only on those types (e.g. `RealtimeEvent`,
|
||||
`RealtimeTransformResult`, `OcrRequestData`). A `type`-style discriminator is a
|
||||
typed field on a struct, not a raw string threaded through the API.
|
||||
|
||||
## Structure
|
||||
|
||||
Use route names directly under `src/`: `ocr`, future `messages`,
|
||||
`chat_completions`, `embeddings`, and similar top-level LiteLLM calls. Do not
|
||||
invent broad names like `engine` for route contracts.
|
||||
|
||||
## Parity Rules
|
||||
|
||||
- Every shared type used by a provider transform needs unit tests for
|
||||
serialization shape.
|
||||
- If Python parity requires always emitting a `null` field instead of omitting
|
||||
it, document that in code and pin it with a test.
|
||||
- Error enums should preserve enough detail for Python/HTTP hosts to map errors
|
||||
consistently without exposing document contents or upstream bodies.
|
||||
12
litellm-rust/crates/core/Cargo.toml
Normal file
12
litellm-rust/crates/core/Cargo.toml
Normal file
|
|
@ -0,0 +1,12 @@
|
|||
[package]
|
||||
name = "litellm-core"
|
||||
version = "0.1.0"
|
||||
edition.workspace = true
|
||||
license.workspace = true
|
||||
repository.workspace = true
|
||||
|
||||
[dependencies]
|
||||
rand.workspace = true
|
||||
serde.workspace = true
|
||||
serde_json.workspace = true
|
||||
thiserror.workspace = true
|
||||
35
litellm-rust/crates/core/src/error.rs
Normal file
35
litellm-rust/crates/core/src/error.rs
Normal file
|
|
@ -0,0 +1,35 @@
|
|||
use thiserror::Error;
|
||||
|
||||
pub type CoreResult<T> = Result<T, CoreError>;
|
||||
|
||||
#[derive(Debug, Error, PartialEq, Eq)]
|
||||
pub enum CoreError {
|
||||
#[error("expected {expected}, got {actual}")]
|
||||
InvalidType {
|
||||
expected: &'static str,
|
||||
actual: &'static str,
|
||||
},
|
||||
#[error("missing required field: {0}")]
|
||||
MissingField(&'static str),
|
||||
#[error("invalid response: {0}")]
|
||||
InvalidResponse(String),
|
||||
#[error("{0}")]
|
||||
Auth(String),
|
||||
#[error("OCR request failed with status {status}: {body}")]
|
||||
Http { status: u16, body: String },
|
||||
#[error("OCR network error: {0}")]
|
||||
Network(String),
|
||||
#[error("routing error: {0}")]
|
||||
Routing(String),
|
||||
}
|
||||
|
||||
pub fn json_type_name(value: &serde_json::Value) -> &'static str {
|
||||
match value {
|
||||
serde_json::Value::Null => "null",
|
||||
serde_json::Value::Bool(_) => "bool",
|
||||
serde_json::Value::Number(_) => "number",
|
||||
serde_json::Value::String(_) => "string",
|
||||
serde_json::Value::Array(_) => "array",
|
||||
serde_json::Value::Object(_) => "object",
|
||||
}
|
||||
}
|
||||
6
litellm-rust/crates/core/src/lib.rs
Normal file
6
litellm-rust/crates/core/src/lib.rs
Normal file
|
|
@ -0,0 +1,6 @@
|
|||
pub mod error;
|
||||
pub mod ocr;
|
||||
pub mod realtime;
|
||||
pub mod router;
|
||||
|
||||
pub use error::{CoreError, CoreResult};
|
||||
2
litellm-rust/crates/core/src/ocr/mod.rs
Normal file
2
litellm-rust/crates/core/src/ocr/mod.rs
Normal file
|
|
@ -0,0 +1,2 @@
|
|||
pub mod transformation;
|
||||
pub mod types;
|
||||
32
litellm-rust/crates/core/src/ocr/transformation.rs
Normal file
32
litellm-rust/crates/core/src/ocr/transformation.rs
Normal file
|
|
@ -0,0 +1,32 @@
|
|||
use serde_json::{Map, Value};
|
||||
|
||||
use crate::CoreResult;
|
||||
|
||||
use super::types::{OcrRequestData, OcrResponseData};
|
||||
|
||||
pub trait OcrProviderConfig {
|
||||
fn supported_ocr_params(&self) -> &'static [&'static str];
|
||||
|
||||
fn map_ocr_params(&self, non_default_params: &Map<String, Value>) -> Map<String, Value> {
|
||||
let mut mapped_params = Map::new();
|
||||
for (param, value) in non_default_params {
|
||||
if self.supported_ocr_params().contains(¶m.as_str()) {
|
||||
mapped_params.insert(param.clone(), value.clone());
|
||||
}
|
||||
}
|
||||
mapped_params
|
||||
}
|
||||
|
||||
fn transform_ocr_request(
|
||||
&self,
|
||||
model: &str,
|
||||
document: Value,
|
||||
optional_params: Map<String, Value>,
|
||||
) -> CoreResult<OcrRequestData>;
|
||||
|
||||
fn transform_ocr_response(
|
||||
&self,
|
||||
model: &str,
|
||||
response_json: Value,
|
||||
) -> CoreResult<OcrResponseData>;
|
||||
}
|
||||
29
litellm-rust/crates/core/src/ocr/types.rs
Normal file
29
litellm-rust/crates/core/src/ocr/types.rs
Normal file
|
|
@ -0,0 +1,29 @@
|
|||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::Value;
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
|
||||
pub struct OcrRequestData {
|
||||
pub data: Value,
|
||||
pub files: Option<Value>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
|
||||
pub struct OcrResponseData {
|
||||
pub pages: Vec<Value>,
|
||||
pub model: String,
|
||||
pub document_annotation: Option<Value>,
|
||||
pub usage_info: Option<Value>,
|
||||
pub object: String,
|
||||
}
|
||||
|
||||
impl OcrResponseData {
|
||||
pub fn into_json(self) -> Value {
|
||||
serde_json::json!({
|
||||
"pages": self.pages,
|
||||
"model": self.model,
|
||||
"document_annotation": self.document_annotation,
|
||||
"usage_info": self.usage_info,
|
||||
"object": self.object,
|
||||
})
|
||||
}
|
||||
}
|
||||
2
litellm-rust/crates/core/src/realtime/mod.rs
Normal file
2
litellm-rust/crates/core/src/realtime/mod.rs
Normal file
|
|
@ -0,0 +1,2 @@
|
|||
pub mod transformation;
|
||||
pub mod types;
|
||||
22
litellm-rust/crates/core/src/realtime/transformation.rs
Normal file
22
litellm-rust/crates/core/src/realtime/transformation.rs
Normal file
|
|
@ -0,0 +1,22 @@
|
|||
use crate::realtime::types::{RealtimeEvent, RealtimeTransformResult};
|
||||
use crate::CoreResult;
|
||||
|
||||
pub trait RealtimeProviderConfig {
|
||||
/// Build the upstream WebSocket URL (e.g. `wss://api.openai.com/v1/realtime?model=…`).
|
||||
/// Pure string construction only — no network, no env.
|
||||
fn complete_url(&self, api_base: Option<&str>, model: &str) -> String;
|
||||
|
||||
/// Transform a client → backend event before it is forwarded upstream.
|
||||
fn transform_realtime_request(
|
||||
&self,
|
||||
event: &RealtimeEvent,
|
||||
model: &str,
|
||||
) -> CoreResult<RealtimeTransformResult>;
|
||||
|
||||
/// Transform a backend → client event before it is forwarded downstream.
|
||||
fn transform_realtime_response(
|
||||
&self,
|
||||
event: &RealtimeEvent,
|
||||
model: &str,
|
||||
) -> CoreResult<RealtimeTransformResult>;
|
||||
}
|
||||
60
litellm-rust/crates/core/src/realtime/types.rs
Normal file
60
litellm-rust/crates/core/src/realtime/types.rs
Normal file
|
|
@ -0,0 +1,60 @@
|
|||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
/// A single realtime event exchanged over the WebSocket.
|
||||
///
|
||||
/// The `type` discriminator is a typed field; the remaining fields are
|
||||
/// preserved losslessly in `data` so a transform can pass an event through, or
|
||||
/// inspect/modify specific fields, without enumerating every event variant.
|
||||
/// Wire (de)serialization happens at the host edge — `core`/`providers` operate
|
||||
/// only on this typed form.
|
||||
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
|
||||
pub struct RealtimeEvent {
|
||||
#[serde(rename = "type")]
|
||||
pub event_type: String,
|
||||
#[serde(flatten)]
|
||||
pub data: Map<String, Value>,
|
||||
}
|
||||
|
||||
/// One or more typed events produced by a realtime transform.
|
||||
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
|
||||
pub struct RealtimeTransformResult {
|
||||
pub events: Vec<RealtimeEvent>,
|
||||
}
|
||||
|
||||
impl RealtimeTransformResult {
|
||||
/// Forward a single event unchanged (the OpenAI baseline).
|
||||
pub fn passthrough(event: RealtimeEvent) -> Self {
|
||||
Self {
|
||||
events: vec![event],
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn event(raw: &str) -> RealtimeEvent {
|
||||
serde_json::from_str(raw).expect("valid event json")
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn realtime_event_round_trips_type_and_extra_fields() {
|
||||
let raw = r#"{"type":"response.output_text.delta","delta":"hi","response_id":"r1"}"#;
|
||||
let parsed = event(raw);
|
||||
assert_eq!(parsed.event_type, "response.output_text.delta");
|
||||
assert_eq!(parsed.data.get("delta"), Some(&Value::String("hi".into())));
|
||||
// Re-serializing yields a semantically-equal event (key order may differ).
|
||||
let reparsed: RealtimeEvent =
|
||||
serde_json::from_str(&serde_json::to_string(&parsed).unwrap()).unwrap();
|
||||
assert_eq!(parsed, reparsed);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn passthrough_produces_single_element_vec() {
|
||||
let parsed = event(r#"{"type":"session.update"}"#);
|
||||
let result = RealtimeTransformResult::passthrough(parsed.clone());
|
||||
assert_eq!(result.events, vec![parsed]);
|
||||
}
|
||||
}
|
||||
44
litellm-rust/crates/core/src/router/deployment.rs
Normal file
44
litellm-rust/crates/core/src/router/deployment.rs
Normal file
|
|
@ -0,0 +1,44 @@
|
|||
//! `model_list` data types, mirroring Python's deployment dict. Deserialize-ready
|
||||
//! so a deployment can be loaded straight from the proxy config's `model_list`.
|
||||
|
||||
use serde::Deserialize;
|
||||
|
||||
/// Per-deployment call parameters, mirroring Python's `litellm_params`.
|
||||
#[derive(Clone, Debug, Deserialize)]
|
||||
pub struct LiteLLMParams {
|
||||
/// Provider model, e.g. `gpt-realtime` or `openai/gpt-realtime`.
|
||||
pub model: String,
|
||||
#[serde(default)]
|
||||
pub api_key: Option<String>,
|
||||
#[serde(default)]
|
||||
pub api_base: Option<String>,
|
||||
}
|
||||
|
||||
/// One entry of the `model_list`, mirroring Python's deployment dict.
|
||||
#[derive(Clone, Debug, Deserialize)]
|
||||
pub struct Deployment {
|
||||
/// Public alias clients request, e.g. `gpt-realtime`.
|
||||
pub model_name: String,
|
||||
pub litellm_params: LiteLLMParams,
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn deserializes_from_model_list_entry() {
|
||||
let entry = r#"{
|
||||
"model_name": "gpt-realtime",
|
||||
"litellm_params": {"model": "openai/gpt-realtime", "api_base": "https://x"}
|
||||
}"#;
|
||||
let deployment: Deployment = serde_json::from_str(entry).expect("valid entry");
|
||||
assert_eq!(deployment.model_name, "gpt-realtime");
|
||||
assert_eq!(deployment.litellm_params.model, "openai/gpt-realtime");
|
||||
assert_eq!(deployment.litellm_params.api_key, None);
|
||||
assert_eq!(
|
||||
deployment.litellm_params.api_base.as_deref(),
|
||||
Some("https://x")
|
||||
);
|
||||
}
|
||||
}
|
||||
93
litellm-rust/crates/core/src/router/mod.rs
Normal file
93
litellm-rust/crates/core/src/router/mod.rs
Normal file
|
|
@ -0,0 +1,93 @@
|
|||
//! Minimal Rust port of LiteLLM's `router.py` deployment selection.
|
||||
//!
|
||||
//! A [`Router`] is built from a `model_list` of [`Deployment`]s
|
||||
//! (`{ model_name, litellm_params: { model, api_key, api_base } }`) and selects
|
||||
//! one per request via a [`RoutingStrategy`]. For now the only strategy is
|
||||
//! `simple-shuffle` — a uniform random pick within a `model_name` group.
|
||||
//!
|
||||
//! This stays pure (no I/O): it only *chooses* a deployment. The host (the
|
||||
//! gateway) takes the chosen deployment and performs the actual provider call.
|
||||
//!
|
||||
//! - [`deployment`] — the `model_list` data types.
|
||||
//! - [`strategy`] — how a deployment is chosen.
|
||||
|
||||
mod deployment;
|
||||
mod strategy;
|
||||
|
||||
pub use deployment::{Deployment, LiteLLMParams};
|
||||
pub use strategy::RoutingStrategy;
|
||||
|
||||
/// Load-balancing router over a `model_list`.
|
||||
#[derive(Clone, Debug, Default)]
|
||||
pub struct Router {
|
||||
model_list: Vec<Deployment>,
|
||||
routing_strategy: RoutingStrategy,
|
||||
}
|
||||
|
||||
impl Router {
|
||||
/// Build a router from a `model_list` using the default `simple-shuffle` strategy.
|
||||
pub fn new(model_list: Vec<Deployment>) -> Self {
|
||||
Self {
|
||||
model_list,
|
||||
routing_strategy: RoutingStrategy::SimpleShuffle,
|
||||
}
|
||||
}
|
||||
|
||||
/// All deployments in the `model_list`. Read-only; used by the host to
|
||||
/// enumerate upstreams (e.g. to pre-warm a connection pool per deployment).
|
||||
pub fn deployments(&self) -> &[Deployment] {
|
||||
&self.model_list
|
||||
}
|
||||
|
||||
/// Whether any deployment is registered under `model`.
|
||||
pub fn has_deployment(&self, model: &str) -> bool {
|
||||
self.model_list
|
||||
.iter()
|
||||
.any(|deployment| deployment.model_name == model)
|
||||
}
|
||||
|
||||
/// Pick a deployment for `model` per the routing strategy. Returns `None`
|
||||
/// when no deployment is registered under that `model_name`.
|
||||
pub fn get_available_deployment(&self, model: &str) -> Option<&Deployment> {
|
||||
let candidates: Vec<&Deployment> = self
|
||||
.model_list
|
||||
.iter()
|
||||
.filter(|deployment| deployment.model_name == model)
|
||||
.collect();
|
||||
self.routing_strategy.select(&candidates)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn deployment(name: &str, model: &str) -> Deployment {
|
||||
Deployment {
|
||||
model_name: name.to_string(),
|
||||
litellm_params: LiteLLMParams {
|
||||
model: model.to_string(),
|
||||
api_key: None,
|
||||
api_base: None,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn selects_a_matching_deployment() {
|
||||
let router = Router::new(vec![
|
||||
deployment("gpt-realtime", "gpt-realtime"),
|
||||
deployment("other", "other-model"),
|
||||
]);
|
||||
let chosen = router
|
||||
.get_available_deployment("gpt-realtime")
|
||||
.expect("a deployment should match");
|
||||
assert_eq!(chosen.model_name, "gpt-realtime");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn unknown_model_returns_none() {
|
||||
let router = Router::new(vec![deployment("gpt-realtime", "gpt-realtime")]);
|
||||
assert!(router.get_available_deployment("missing").is_none());
|
||||
}
|
||||
}
|
||||
26
litellm-rust/crates/core/src/router/strategy/mod.rs
Normal file
26
litellm-rust/crates/core/src/router/strategy/mod.rs
Normal file
|
|
@ -0,0 +1,26 @@
|
|||
//! Routing policy: how the router picks one deployment from a model group.
|
||||
//!
|
||||
//! One module per strategy; [`RoutingStrategy::select`] dispatches to it. New
|
||||
//! strategies (least-busy, latency-based, …) get their own file here.
|
||||
|
||||
mod simple_shuffle;
|
||||
|
||||
use super::Deployment;
|
||||
|
||||
/// How the router chooses among the deployments sharing a `model_name`.
|
||||
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
|
||||
pub enum RoutingStrategy {
|
||||
/// Uniform random pick among the matching deployments.
|
||||
#[default]
|
||||
SimpleShuffle,
|
||||
}
|
||||
|
||||
impl RoutingStrategy {
|
||||
/// Choose one deployment from `candidates` (all sharing the requested
|
||||
/// `model_name`). Returns `None` when there are no candidates.
|
||||
pub fn select<'a>(&self, candidates: &[&'a Deployment]) -> Option<&'a Deployment> {
|
||||
match self {
|
||||
RoutingStrategy::SimpleShuffle => simple_shuffle::select(candidates),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,47 @@
|
|||
//! `simple-shuffle`: a uniform random pick among the candidate deployments.
|
||||
|
||||
use rand::seq::SliceRandom;
|
||||
|
||||
use crate::router::Deployment;
|
||||
|
||||
/// Uniform random choice among `candidates` (all sharing the requested
|
||||
/// `model_name`). Returns `None` when there are no candidates.
|
||||
pub fn select<'a>(candidates: &[&'a Deployment]) -> Option<&'a Deployment> {
|
||||
candidates.choose(&mut rand::thread_rng()).copied()
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::router::{Deployment, LiteLLMParams};
|
||||
|
||||
fn deployment(model: &str) -> Deployment {
|
||||
Deployment {
|
||||
model_name: "gpt-realtime".to_string(),
|
||||
litellm_params: LiteLLMParams {
|
||||
model: model.to_string(),
|
||||
api_key: None,
|
||||
api_base: None,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn picks_from_candidates() {
|
||||
let a = deployment("key-a");
|
||||
let b = deployment("key-b");
|
||||
let candidates = vec![&a, &b];
|
||||
for _ in 0..20 {
|
||||
let chosen = select(&candidates).expect("non-empty");
|
||||
assert!(matches!(
|
||||
chosen.litellm_params.model.as_str(),
|
||||
"key-a" | "key-b"
|
||||
));
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn empty_candidates_select_none() {
|
||||
assert!(select(&[]).is_none());
|
||||
}
|
||||
}
|
||||
53
litellm-rust/crates/providers/CLAUDE.md
Normal file
53
litellm-rust/crates/providers/CLAUDE.md
Normal file
|
|
@ -0,0 +1,53 @@
|
|||
# CLAUDE.md
|
||||
|
||||
Rules for `litellm-rust/crates/providers`.
|
||||
|
||||
## Responsibility
|
||||
|
||||
`providers` owns provider-specific pure transforms. It mirrors the existing
|
||||
Python provider modules closely enough that parity review is mechanical.
|
||||
|
||||
Provider files should map to the Python provider tree:
|
||||
|
||||
```text
|
||||
providers/src/<provider>/<route>/transformation.rs
|
||||
```
|
||||
|
||||
For example, Mistral OCR lives at
|
||||
`providers/src/mistral/ocr/transformation.rs`, matching
|
||||
`litellm/llms/mistral/ocr/transformation.py`.
|
||||
|
||||
Allowed:
|
||||
- Provider request transforms.
|
||||
- Provider response normalization.
|
||||
- Supported-parameter filtering.
|
||||
- Provider-specific validation that does not require I/O or secrets.
|
||||
|
||||
Not allowed:
|
||||
- HTTP clients or provider SDK calls.
|
||||
- Environment variable reads.
|
||||
- API key resolution or auth header construction.
|
||||
- Logging, callbacks, spend tracking, retries, routing, cooldowns, or fallbacks.
|
||||
- Panics on bad user/provider input.
|
||||
|
||||
## Required Tests
|
||||
|
||||
Every provider transform must include focused unit tests for:
|
||||
- Supported params matching the Python provider config.
|
||||
- Unknown params being dropped or transformed the same way as Python.
|
||||
- Request body shape matching Python output.
|
||||
- Response normalization with complete, missing, null, and extra fields.
|
||||
- Bad input returning typed errors.
|
||||
|
||||
For OCR specifically, assume documents can contain personal data. Tests should
|
||||
prove transforms do not copy document contents into error messages.
|
||||
|
||||
## Implementation Rules
|
||||
|
||||
- Prefer static supported-parameter lists over allocating strings on every call.
|
||||
- Keep transforms deterministic and allocation-conscious, but choose clarity over
|
||||
premature micro-optimization for tiny parameter lists.
|
||||
- Use typed errors from `core`; avoid stringly-typed error plumbing.
|
||||
- Add comments only when they explain Python-parity decisions or provider quirks.
|
||||
- Put route-level provider dispatch in a route file such as `providers/src/ocr.rs`.
|
||||
Do not move provider-specific transform logic into the Python bridge.
|
||||
18
litellm-rust/crates/providers/Cargo.toml
Normal file
18
litellm-rust/crates/providers/Cargo.toml
Normal file
|
|
@ -0,0 +1,18 @@
|
|||
[package]
|
||||
name = "litellm-providers"
|
||||
version = "0.1.0"
|
||||
edition.workspace = true
|
||||
license.workspace = true
|
||||
repository.workspace = true
|
||||
|
||||
[dependencies]
|
||||
litellm-core.workspace = true
|
||||
reqwest.workspace = true
|
||||
serde_json.workspace = true
|
||||
tokio.workspace = true
|
||||
tokio-tungstenite.workspace = true
|
||||
futures-util.workspace = true
|
||||
|
||||
[dev-dependencies]
|
||||
serde_json.workspace = true
|
||||
futures-channel = "0.3"
|
||||
5
litellm-rust/crates/providers/src/lib.rs
Normal file
5
litellm-rust/crates/providers/src/lib.rs
Normal file
|
|
@ -0,0 +1,5 @@
|
|||
pub mod mistral;
|
||||
pub mod ocr;
|
||||
pub mod openai;
|
||||
pub mod realtime;
|
||||
pub mod realtime_pool;
|
||||
1
litellm-rust/crates/providers/src/mistral/mod.rs
Normal file
1
litellm-rust/crates/providers/src/mistral/mod.rs
Normal file
|
|
@ -0,0 +1 @@
|
|||
pub mod ocr;
|
||||
1
litellm-rust/crates/providers/src/mistral/ocr/mod.rs
Normal file
1
litellm-rust/crates/providers/src/mistral/ocr/mod.rs
Normal file
|
|
@ -0,0 +1 @@
|
|||
pub mod transformation;
|
||||
292
litellm-rust/crates/providers/src/mistral/ocr/transformation.rs
Normal file
292
litellm-rust/crates/providers/src/mistral/ocr/transformation.rs
Normal file
|
|
@ -0,0 +1,292 @@
|
|||
use litellm_core::error::{json_type_name, CoreError, CoreResult};
|
||||
use litellm_core::ocr::transformation::OcrProviderConfig;
|
||||
use litellm_core::ocr::types::{OcrRequestData, OcrResponseData};
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
const SUPPORTED_OCR_PARAMS: &[&str] = &[
|
||||
"pages",
|
||||
"include_image_base64",
|
||||
"image_limit",
|
||||
"image_min_size",
|
||||
"bbox_annotation_format",
|
||||
"document_annotation_format",
|
||||
"document_annotation_prompt",
|
||||
"extract_header",
|
||||
"extract_footer",
|
||||
"table_format",
|
||||
"confidence_scores_granularity",
|
||||
"id",
|
||||
];
|
||||
|
||||
/// Default Mistral API base, used when the caller does not override `api_base`.
|
||||
pub const MISTRAL_DEFAULT_API_BASE: &str = "https://api.mistral.ai/v1";
|
||||
|
||||
/// Environment variable holding the Mistral API key.
|
||||
pub const MISTRAL_API_KEY_ENV: &str = "MISTRAL_API_KEY";
|
||||
|
||||
/// Error message raised when no Mistral API key can be resolved.
|
||||
pub const MISSING_KEY_MESSAGE: &str = "Missing Mistral API Key - A call is being made to Mistral but no key is set either in the environment variables or via params";
|
||||
|
||||
/// Build the complete OCR endpoint URL, de-duplicating a trailing `/v1`.
|
||||
///
|
||||
/// Blank/whitespace `api_base` is treated as absent (guard at resolution time).
|
||||
pub fn complete_url(api_base: Option<&str>) -> String {
|
||||
let base = api_base
|
||||
.map(str::trim)
|
||||
.filter(|base| !base.is_empty())
|
||||
.unwrap_or(MISTRAL_DEFAULT_API_BASE)
|
||||
.trim_end_matches('/');
|
||||
|
||||
if base.ends_with("/v1") {
|
||||
format!("{base}/ocr")
|
||||
} else {
|
||||
format!("{base}/v1/ocr")
|
||||
}
|
||||
}
|
||||
|
||||
/// Resolve the Mistral API key from the explicit param or the environment.
|
||||
///
|
||||
/// Blank/whitespace values are treated as absent. Returns `CoreError::Auth`
|
||||
/// when no usable key is available.
|
||||
///
|
||||
/// Note: the env fallback only reads the process environment. Secret-manager
|
||||
/// backends (AWS/Azure/GCP/Vault) are resolved on the Python side and passed in
|
||||
/// via `api_key`; this fallback is a last resort for direct/standalone use.
|
||||
pub fn resolve_api_key(
|
||||
api_key: Option<&str>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<String> {
|
||||
api_key
|
||||
.map(str::trim)
|
||||
.filter(|key| !key.is_empty())
|
||||
.map(str::to_string)
|
||||
.or_else(|| env_lookup(MISTRAL_API_KEY_ENV).filter(|key| !key.trim().is_empty()))
|
||||
.ok_or_else(|| CoreError::Auth(MISSING_KEY_MESSAGE.to_string()))
|
||||
}
|
||||
|
||||
pub struct MistralOcrConfig;
|
||||
|
||||
pub const MISTRAL_OCR_CONFIG: MistralOcrConfig = MistralOcrConfig;
|
||||
|
||||
impl OcrProviderConfig for MistralOcrConfig {
|
||||
fn supported_ocr_params(&self) -> &'static [&'static str] {
|
||||
SUPPORTED_OCR_PARAMS
|
||||
}
|
||||
|
||||
fn transform_ocr_request(
|
||||
&self,
|
||||
model: &str,
|
||||
document: Value,
|
||||
optional_params: Map<String, Value>,
|
||||
) -> CoreResult<OcrRequestData> {
|
||||
if !document.is_object() {
|
||||
return Err(CoreError::InvalidType {
|
||||
expected: "object",
|
||||
actual: json_type_name(&document),
|
||||
});
|
||||
}
|
||||
|
||||
let mut data = Map::new();
|
||||
data.insert("model".to_string(), Value::String(model.to_string()));
|
||||
data.insert("document".to_string(), document);
|
||||
for (param, value) in optional_params {
|
||||
data.insert(param, value);
|
||||
}
|
||||
|
||||
Ok(OcrRequestData {
|
||||
data: Value::Object(data),
|
||||
files: None,
|
||||
})
|
||||
}
|
||||
|
||||
fn transform_ocr_response(
|
||||
&self,
|
||||
model: &str,
|
||||
response_json: Value,
|
||||
) -> CoreResult<OcrResponseData> {
|
||||
let response_object = response_json
|
||||
.as_object()
|
||||
.ok_or_else(|| CoreError::InvalidType {
|
||||
expected: "object",
|
||||
actual: json_type_name(&response_json),
|
||||
})?;
|
||||
|
||||
let pages = response_object
|
||||
.get("pages")
|
||||
.and_then(Value::as_array)
|
||||
.cloned()
|
||||
.unwrap_or_default();
|
||||
let model = response_object
|
||||
.get("model")
|
||||
.and_then(Value::as_str)
|
||||
.unwrap_or(model)
|
||||
.to_string();
|
||||
let document_annotation = response_object.get("document_annotation").cloned();
|
||||
let usage_info = response_object.get("usage_info").cloned();
|
||||
|
||||
Ok(OcrResponseData {
|
||||
pages,
|
||||
model,
|
||||
document_annotation,
|
||||
usage_info,
|
||||
object: "ocr".to_string(),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
pub fn supported_ocr_params() -> &'static [&'static str] {
|
||||
MISTRAL_OCR_CONFIG.supported_ocr_params()
|
||||
}
|
||||
|
||||
pub fn map_ocr_params(non_default_params: &Map<String, Value>) -> Map<String, Value> {
|
||||
MISTRAL_OCR_CONFIG.map_ocr_params(non_default_params)
|
||||
}
|
||||
|
||||
pub fn transform_ocr_request(
|
||||
model: &str,
|
||||
document: Value,
|
||||
optional_params: Map<String, Value>,
|
||||
) -> CoreResult<OcrRequestData> {
|
||||
MISTRAL_OCR_CONFIG.transform_ocr_request(model, document, optional_params)
|
||||
}
|
||||
|
||||
pub fn transform_ocr_response(model: &str, response_json: Value) -> CoreResult<OcrResponseData> {
|
||||
MISTRAL_OCR_CONFIG.transform_ocr_response(model, response_json)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use serde_json::json;
|
||||
|
||||
#[test]
|
||||
fn supported_params_match_python_mistral_ocr_config() {
|
||||
assert_eq!(
|
||||
supported_ocr_params(),
|
||||
&[
|
||||
"pages",
|
||||
"include_image_base64",
|
||||
"image_limit",
|
||||
"image_min_size",
|
||||
"bbox_annotation_format",
|
||||
"document_annotation_format",
|
||||
"document_annotation_prompt",
|
||||
"extract_header",
|
||||
"extract_footer",
|
||||
"table_format",
|
||||
"confidence_scores_granularity",
|
||||
"id",
|
||||
]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn map_ocr_params_drops_unknown_params() {
|
||||
let params = json!({
|
||||
"extract_header": true,
|
||||
"unsupported_param": "value",
|
||||
"pages": [0, 1]
|
||||
});
|
||||
let mapped = map_ocr_params(params.as_object().unwrap());
|
||||
|
||||
assert_eq!(mapped.get("extract_header"), Some(&json!(true)));
|
||||
assert_eq!(mapped.get("pages"), Some(&json!([0, 1])));
|
||||
assert!(!mapped.contains_key("unsupported_param"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn transform_ocr_request_builds_mistral_body() {
|
||||
let document = json!({
|
||||
"type": "document_url",
|
||||
"document_url": "https://example.com/doc.pdf"
|
||||
});
|
||||
let optional_params = json!({
|
||||
"include_image_base64": true,
|
||||
"table_format": "html"
|
||||
})
|
||||
.as_object()
|
||||
.unwrap()
|
||||
.clone();
|
||||
|
||||
let result = transform_ocr_request("mistral-ocr-latest", document.clone(), optional_params)
|
||||
.expect("request should transform");
|
||||
|
||||
assert_eq!(
|
||||
result.data,
|
||||
json!({
|
||||
"model": "mistral-ocr-latest",
|
||||
"document": document,
|
||||
"include_image_base64": true,
|
||||
"table_format": "html"
|
||||
})
|
||||
);
|
||||
assert_eq!(result.files, None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn transform_ocr_request_rejects_non_object_document() {
|
||||
let err = transform_ocr_request("mistral-ocr-latest", json!("bad"), Map::new())
|
||||
.expect_err("string document should be rejected");
|
||||
|
||||
assert_eq!(
|
||||
err,
|
||||
CoreError::InvalidType {
|
||||
expected: "object",
|
||||
actual: "string",
|
||||
}
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn transform_ocr_response_normalizes_mistral_json() {
|
||||
let response = json!({
|
||||
"pages": [{"index": 0, "markdown": "hello"}],
|
||||
"model": "mistral-ocr-2505-completion",
|
||||
"document_annotation": null,
|
||||
"usage_info": {"pages_processed": 1}
|
||||
});
|
||||
|
||||
let result = transform_ocr_response("mistral-ocr-latest", response)
|
||||
.expect("response should transform");
|
||||
|
||||
assert_eq!(result.pages, vec![json!({"index": 0, "markdown": "hello"})]);
|
||||
assert_eq!(result.model, "mistral-ocr-2505-completion");
|
||||
assert_eq!(result.document_annotation, Some(Value::Null));
|
||||
assert_eq!(result.usage_info, Some(json!({"pages_processed": 1})));
|
||||
assert_eq!(result.object, "ocr");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn complete_url_defaults_and_dedupes_v1() {
|
||||
assert_eq!(complete_url(None), "https://api.mistral.ai/v1/ocr");
|
||||
assert_eq!(complete_url(Some(" ")), "https://api.mistral.ai/v1/ocr");
|
||||
assert_eq!(
|
||||
complete_url(Some("https://proxy.internal")),
|
||||
"https://proxy.internal/v1/ocr"
|
||||
);
|
||||
assert_eq!(
|
||||
complete_url(Some("https://proxy.internal/v1/")),
|
||||
"https://proxy.internal/v1/ocr"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolve_api_key_prefers_param_then_env() {
|
||||
let no_env = |_: &str| None;
|
||||
assert_eq!(
|
||||
resolve_api_key(Some("sk-param"), &no_env).unwrap(),
|
||||
"sk-param"
|
||||
);
|
||||
|
||||
let with_env = |key: &str| (key == MISTRAL_API_KEY_ENV).then(|| "sk-env".to_string());
|
||||
assert_eq!(resolve_api_key(None, &with_env).unwrap(), "sk-env");
|
||||
// Blank param falls through to the environment.
|
||||
assert_eq!(resolve_api_key(Some(" "), &with_env).unwrap(), "sk-env");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolve_api_key_errors_when_absent() {
|
||||
let err = resolve_api_key(None, &|_| None).expect_err("missing key should error");
|
||||
assert_eq!(err, CoreError::Auth(MISSING_KEY_MESSAGE.to_string()));
|
||||
}
|
||||
}
|
||||
127
litellm-rust/crates/providers/src/ocr.rs
Normal file
127
litellm-rust/crates/providers/src/ocr.rs
Normal file
|
|
@ -0,0 +1,127 @@
|
|||
//! End-to-end OCR orchestration.
|
||||
//!
|
||||
//! Owns the whole Mistral OCR call so the Python side stays a thin bridge:
|
||||
//! resolve the API key, build the URL + body via the pure transforms, POST it,
|
||||
//! and normalize the response. The HTTP client is built once and reused.
|
||||
|
||||
use std::sync::OnceLock;
|
||||
use std::time::Duration;
|
||||
|
||||
use litellm_core::error::CoreError;
|
||||
use litellm_core::ocr::transformation::OcrProviderConfig;
|
||||
use litellm_core::CoreResult;
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
use crate::mistral::ocr::transformation as mistral;
|
||||
use crate::mistral::ocr::transformation::MISTRAL_OCR_CONFIG;
|
||||
|
||||
/// OCR over large documents can take a while; bound it generously rather than
|
||||
/// hanging forever on an unresponsive upstream. The client-level limit is the
|
||||
/// outer ceiling; callers can tighten it per request via ``run_ocr``'s ``timeout``.
|
||||
const OCR_TIMEOUT_SECS: u64 = 600;
|
||||
|
||||
/// Maximum upstream body characters retained in error messages. OCR responses
|
||||
/// can echo document contents and prompts; keep enough for debugging without
|
||||
/// forwarding sensitive payloads across the host boundary.
|
||||
const ERROR_BODY_MAX_CHARS: usize = 256;
|
||||
|
||||
/// Process-wide blocking HTTP client (connection pool + TLS reused across calls).
|
||||
fn http_client() -> &'static reqwest::blocking::Client {
|
||||
static CLIENT: OnceLock<reqwest::blocking::Client> = OnceLock::new();
|
||||
CLIENT.get_or_init(|| {
|
||||
reqwest::blocking::Client::builder()
|
||||
.timeout(Duration::from_secs(OCR_TIMEOUT_SECS))
|
||||
.build()
|
||||
.expect("failed to build reqwest client")
|
||||
})
|
||||
}
|
||||
|
||||
fn truncate_error_body(body: &str) -> String {
|
||||
if body.chars().count() <= ERROR_BODY_MAX_CHARS {
|
||||
return body.to_string();
|
||||
}
|
||||
let truncated: String = body.chars().take(ERROR_BODY_MAX_CHARS).collect();
|
||||
format!("{truncated}... (truncated)")
|
||||
}
|
||||
|
||||
/// Perform a Mistral OCR call end to end and return the normalized response as
|
||||
/// JSON (the shape the Python `OCRResponse` model expects).
|
||||
///
|
||||
/// Blocking: intended to be called with the GIL released from the Python bridge.
|
||||
pub fn run_ocr(
|
||||
model: &str,
|
||||
document: Value,
|
||||
api_key: Option<&str>,
|
||||
api_base: Option<&str>,
|
||||
optional_params: Map<String, Value>,
|
||||
timeout: Option<Duration>,
|
||||
) -> CoreResult<Value> {
|
||||
let config = &MISTRAL_OCR_CONFIG;
|
||||
|
||||
let api_key = mistral::resolve_api_key(api_key, &|key| std::env::var(key).ok())?;
|
||||
let url = mistral::complete_url(api_base);
|
||||
let filtered_params = config.map_ocr_params(&optional_params);
|
||||
let body = config
|
||||
.transform_ocr_request(model, document, filtered_params)?
|
||||
.data;
|
||||
|
||||
let mut request = http_client().post(&url).bearer_auth(&api_key).json(&body);
|
||||
if let Some(duration) = timeout {
|
||||
request = request.timeout(duration);
|
||||
}
|
||||
|
||||
let response = request
|
||||
.send()
|
||||
.map_err(|err| CoreError::Network(err.to_string()))?;
|
||||
|
||||
let status = response.status();
|
||||
let text = response
|
||||
.text()
|
||||
.map_err(|err| CoreError::Network(err.to_string()))?;
|
||||
|
||||
if !status.is_success() {
|
||||
return Err(CoreError::Http {
|
||||
status: status.as_u16(),
|
||||
body: truncate_error_body(&text),
|
||||
});
|
||||
}
|
||||
|
||||
let response_json: Value = serde_json::from_str(&text)
|
||||
.map_err(|err| CoreError::InvalidResponse(format!("invalid OCR response JSON: {err}")))?;
|
||||
|
||||
Ok(config
|
||||
.transform_ocr_response(model, response_json)?
|
||||
.into_json())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn truncate_error_body_passes_short_strings_through() {
|
||||
let body = "Unauthorized";
|
||||
assert_eq!(truncate_error_body(body), "Unauthorized");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn truncate_error_body_caps_long_payloads() {
|
||||
let body = "x".repeat(ERROR_BODY_MAX_CHARS + 50);
|
||||
let truncated = truncate_error_body(&body);
|
||||
|
||||
assert!(truncated.ends_with("... (truncated)"));
|
||||
let prefix_chars = truncated
|
||||
.strip_suffix("... (truncated)")
|
||||
.expect("truncated marker present")
|
||||
.chars()
|
||||
.count();
|
||||
assert_eq!(prefix_chars, ERROR_BODY_MAX_CHARS);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn truncate_error_body_does_not_split_multibyte_chars() {
|
||||
let body = "é".repeat(ERROR_BODY_MAX_CHARS + 10);
|
||||
let truncated = truncate_error_body(&body);
|
||||
assert!(truncated.is_char_boundary(truncated.len()));
|
||||
}
|
||||
}
|
||||
1
litellm-rust/crates/providers/src/openai/mod.rs
Normal file
1
litellm-rust/crates/providers/src/openai/mod.rs
Normal file
|
|
@ -0,0 +1 @@
|
|||
pub mod realtime;
|
||||
1
litellm-rust/crates/providers/src/openai/realtime/mod.rs
Normal file
1
litellm-rust/crates/providers/src/openai/realtime/mod.rs
Normal file
|
|
@ -0,0 +1 @@
|
|||
pub mod transformation;
|
||||
|
|
@ -0,0 +1,189 @@
|
|||
use litellm_core::realtime::transformation::RealtimeProviderConfig;
|
||||
use litellm_core::realtime::types::{RealtimeEvent, RealtimeTransformResult};
|
||||
use litellm_core::CoreResult;
|
||||
|
||||
/// Default OpenAI API base, used when the caller does not override `api_base`.
|
||||
pub const OPENAI_REALTIME_DEFAULT_API_BASE: &str = "https://api.openai.com";
|
||||
|
||||
/// Path appended to the resolved host base to reach the realtime endpoint.
|
||||
pub const OPENAI_REALTIME_PATH: &str = "/v1/realtime";
|
||||
|
||||
/// Percent-encode a query value, escaping any char outside the RFC 3986
|
||||
/// unreserved set (`A-Za-z0-9-._~`). Keeps us dependency-free; common realtime
|
||||
/// model slugs have no special chars, but this stays correct for the rest.
|
||||
fn percent_encode(value: &str) -> String {
|
||||
let mut encoded = String::with_capacity(value.len());
|
||||
for byte in value.bytes() {
|
||||
let unreserved = byte.is_ascii_alphanumeric() || matches!(byte, b'-' | b'.' | b'_' | b'~');
|
||||
if unreserved {
|
||||
encoded.push(byte as char);
|
||||
} else {
|
||||
encoded.push('%');
|
||||
encoded.push_str(&format!("{byte:02X}"));
|
||||
}
|
||||
}
|
||||
encoded
|
||||
}
|
||||
|
||||
/// Build the realtime WebSocket URL, porting Python's `OpenAIRealtime._construct_url`.
|
||||
///
|
||||
/// Blank/whitespace `api_base` is treated as absent (guard at resolution time),
|
||||
/// falling back to the default. The scheme is swapped to its WebSocket
|
||||
/// equivalent (`https://`→`wss://`, `http://`→`ws://`); bases already using
|
||||
/// `ws`/`wss` are left untouched. A bare host or unrecognized scheme defaults to
|
||||
/// secure `wss://` so we never hand a scheme-less URL to the connector (this is
|
||||
/// a deliberate hardening over Python's `_construct_url`, which would emit a
|
||||
/// scheme-less URL here). A trailing `/` is trimmed before the path and
|
||||
/// `?model=<encoded>` are appended.
|
||||
pub fn complete_url(api_base: Option<&str>, model: &str) -> String {
|
||||
let base = api_base
|
||||
.map(str::trim)
|
||||
.filter(|base| !base.is_empty())
|
||||
.unwrap_or(OPENAI_REALTIME_DEFAULT_API_BASE);
|
||||
|
||||
let base = if let Some(rest) = base.strip_prefix("https://") {
|
||||
format!("wss://{rest}")
|
||||
} else if let Some(rest) = base.strip_prefix("http://") {
|
||||
format!("ws://{rest}")
|
||||
} else if base.starts_with("wss://") || base.starts_with("ws://") {
|
||||
base.to_string()
|
||||
} else {
|
||||
format!("wss://{base}")
|
||||
};
|
||||
|
||||
let base = base.trim_end_matches('/');
|
||||
|
||||
format!(
|
||||
"{base}{OPENAI_REALTIME_PATH}?model={}",
|
||||
percent_encode(model)
|
||||
)
|
||||
}
|
||||
|
||||
pub struct OpenAiRealtimeConfig;
|
||||
|
||||
pub const OPENAI_REALTIME_CONFIG: OpenAiRealtimeConfig = OpenAiRealtimeConfig;
|
||||
|
||||
impl RealtimeProviderConfig for OpenAiRealtimeConfig {
|
||||
fn complete_url(&self, api_base: Option<&str>, model: &str) -> String {
|
||||
complete_url(api_base, model)
|
||||
}
|
||||
|
||||
fn transform_realtime_request(
|
||||
&self,
|
||||
event: &RealtimeEvent,
|
||||
_model: &str,
|
||||
) -> CoreResult<RealtimeTransformResult> {
|
||||
Ok(RealtimeTransformResult::passthrough(event.clone()))
|
||||
}
|
||||
|
||||
fn transform_realtime_response(
|
||||
&self,
|
||||
event: &RealtimeEvent,
|
||||
_model: &str,
|
||||
) -> CoreResult<RealtimeTransformResult> {
|
||||
Ok(RealtimeTransformResult::passthrough(event.clone()))
|
||||
}
|
||||
}
|
||||
|
||||
pub fn transform_realtime_request(
|
||||
event: &RealtimeEvent,
|
||||
model: &str,
|
||||
) -> CoreResult<RealtimeTransformResult> {
|
||||
OPENAI_REALTIME_CONFIG.transform_realtime_request(event, model)
|
||||
}
|
||||
|
||||
pub fn transform_realtime_response(
|
||||
event: &RealtimeEvent,
|
||||
model: &str,
|
||||
) -> CoreResult<RealtimeTransformResult> {
|
||||
OPENAI_REALTIME_CONFIG.transform_realtime_response(event, model)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn complete_url_defaults_to_openai_wss() {
|
||||
assert_eq!(
|
||||
complete_url(None, "gpt-4o-realtime-preview"),
|
||||
"wss://api.openai.com/v1/realtime?model=gpt-4o-realtime-preview"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn complete_url_blank_base_uses_default() {
|
||||
assert_eq!(
|
||||
complete_url(Some(" "), "gpt-4o-realtime-preview"),
|
||||
"wss://api.openai.com/v1/realtime?model=gpt-4o-realtime-preview"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn complete_url_swaps_http_to_ws() {
|
||||
assert_eq!(
|
||||
complete_url(Some("http://localhost:8080"), "gpt-4o-realtime-preview"),
|
||||
"ws://localhost:8080/v1/realtime?model=gpt-4o-realtime-preview"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn complete_url_dedupes_trailing_slash() {
|
||||
assert_eq!(
|
||||
complete_url(Some("https://api.openai.com/"), "gpt-4o-realtime-preview"),
|
||||
"wss://api.openai.com/v1/realtime?model=gpt-4o-realtime-preview"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn complete_url_custom_base() {
|
||||
assert_eq!(
|
||||
complete_url(Some("https://oai.azure.example"), "gpt-4o-realtime-preview"),
|
||||
"wss://oai.azure.example/v1/realtime?model=gpt-4o-realtime-preview"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn complete_url_preserves_existing_wss_scheme() {
|
||||
assert_eq!(
|
||||
complete_url(Some("wss://api.openai.com"), "gpt-realtime"),
|
||||
"wss://api.openai.com/v1/realtime?model=gpt-realtime"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn complete_url_bare_host_defaults_to_wss() {
|
||||
assert_eq!(
|
||||
complete_url(Some("api.openai.com"), "gpt-realtime"),
|
||||
"wss://api.openai.com/v1/realtime?model=gpt-realtime"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn complete_url_percent_encodes_model_space() {
|
||||
assert_eq!(
|
||||
complete_url(None, "gpt 4o"),
|
||||
"wss://api.openai.com/v1/realtime?model=gpt%204o"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn transform_realtime_request_passthrough_preserves_event() {
|
||||
let event: RealtimeEvent =
|
||||
serde_json::from_str(r#"{"type":"session.update","session":{"voice":"alloy"}}"#)
|
||||
.expect("valid event");
|
||||
let result =
|
||||
transform_realtime_request(&event, "gpt-realtime").expect("passthrough is infallible");
|
||||
assert_eq!(result.events, vec![event]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn transform_realtime_response_passthrough_preserves_event() {
|
||||
let event: RealtimeEvent =
|
||||
serde_json::from_str(r#"{"type":"response.output_audio.delta","delta":"abc=="}"#)
|
||||
.expect("valid event");
|
||||
let result =
|
||||
transform_realtime_response(&event, "gpt-realtime").expect("passthrough is infallible");
|
||||
assert_eq!(result.events, vec![event]);
|
||||
}
|
||||
}
|
||||
374
litellm-rust/crates/providers/src/realtime.rs
Normal file
374
litellm-rust/crates/providers/src/realtime.rs
Normal file
|
|
@ -0,0 +1,374 @@
|
|||
//! End-to-end OpenAI realtime invocation.
|
||||
//!
|
||||
//! The host-facing entry point, mirroring `providers::ocr::run_ocr`: open the
|
||||
//! WebSocket to OpenAI, then splice a client realtime stream to the upstream,
|
||||
//! driving typed events through the pure `OPENAI_REALTIME_CONFIG` transforms.
|
||||
//! Network, auth header, key resolution, and wire (de)serialization live here so
|
||||
//! the `transformation` module stays pure and typed.
|
||||
//!
|
||||
//! The dial and splice steps are factored out ([`dial_upstream`], [`splice`]) so
|
||||
//! the connection pool ([`crate::realtime_pool`]) can pre-establish an upstream,
|
||||
//! buffer its `session.created`, and later hand the live socket to the same
|
||||
//! splice loop a fresh dial uses.
|
||||
|
||||
use std::time::Duration;
|
||||
|
||||
use futures_util::stream::{SplitSink, SplitStream};
|
||||
use futures_util::{Sink, SinkExt, Stream, StreamExt};
|
||||
use litellm_core::error::CoreError;
|
||||
use litellm_core::realtime::transformation::RealtimeProviderConfig;
|
||||
use litellm_core::realtime::types::RealtimeEvent;
|
||||
use litellm_core::CoreResult;
|
||||
use tokio::net::TcpStream;
|
||||
use tokio_tungstenite::tungstenite::client::IntoClientRequest;
|
||||
use tokio_tungstenite::tungstenite::http::header::AUTHORIZATION;
|
||||
use tokio_tungstenite::tungstenite::http::HeaderValue;
|
||||
use tokio_tungstenite::tungstenite::Message;
|
||||
use tokio_tungstenite::{connect_async, MaybeTlsStream, WebSocketStream};
|
||||
|
||||
use crate::openai::realtime::transformation::OPENAI_REALTIME_CONFIG;
|
||||
|
||||
/// Environment variable holding the OpenAI API key (last-resort fallback).
|
||||
const OPENAI_API_KEY_ENV: &str = "OPENAI_API_KEY";
|
||||
|
||||
const MISSING_KEY_MESSAGE: &str = "Missing OpenAI API Key - a realtime call is being made but no key was passed via params or the OPENAI_API_KEY environment variable";
|
||||
|
||||
/// Default **idle** timeout: if neither side sends a frame for this long, the
|
||||
/// session is reaped. It resets on any activity, so it does not cap a healthy
|
||||
/// (continuously streaming) session — it only frees a stalled one (e.g. a
|
||||
/// half-open upstream that keeps the socket open but stops sending).
|
||||
const DEFAULT_IDLE_TIMEOUT_SECS: u64 = 300;
|
||||
|
||||
/// The concrete upstream WebSocket type (TLS or plain). Shared by the dial path
|
||||
/// and the pool so warm sockets and fresh sockets are the exact same type.
|
||||
pub type UpstreamWs = WebSocketStream<MaybeTlsStream<TcpStream>>;
|
||||
pub(crate) type UpstreamTx = SplitSink<UpstreamWs, Message>;
|
||||
pub(crate) type UpstreamRx = SplitStream<UpstreamWs>;
|
||||
|
||||
/// Resolve the OpenAI API key from the explicit param or the environment.
|
||||
///
|
||||
/// Blank/whitespace values are treated as absent (guard at resolution time).
|
||||
pub(crate) fn resolve_api_key(api_key: Option<&str>) -> CoreResult<String> {
|
||||
api_key
|
||||
.map(str::trim)
|
||||
.filter(|key| !key.is_empty())
|
||||
.map(str::to_string)
|
||||
.or_else(|| {
|
||||
std::env::var(OPENAI_API_KEY_ENV)
|
||||
.ok()
|
||||
.filter(|key| !key.trim().is_empty())
|
||||
})
|
||||
.ok_or_else(|| CoreError::Auth(MISSING_KEY_MESSAGE.to_string()))
|
||||
}
|
||||
|
||||
/// Open the upstream WebSocket to OpenAI for `(model, api_key, api_base)`.
|
||||
///
|
||||
/// This is the dial half of [`realtime`], factored out so the pool can
|
||||
/// pre-establish sockets ahead of any client. `api_key` here is already resolved
|
||||
/// (non-blank) — the pool resolves it once when it is created.
|
||||
pub(crate) async fn dial_upstream(
|
||||
model: &str,
|
||||
api_key: &str,
|
||||
api_base: Option<&str>,
|
||||
) -> CoreResult<UpstreamWs> {
|
||||
let url = OPENAI_REALTIME_CONFIG.complete_url(api_base, model);
|
||||
|
||||
let mut request = url
|
||||
.as_str()
|
||||
.into_client_request()
|
||||
.map_err(|err| CoreError::Network(err.to_string()))?;
|
||||
// GA realtime: only Authorization. The legacy OpenAI-Beta header triggers
|
||||
// beta_api_shape_disabled, so we do not send it.
|
||||
request.headers_mut().insert(
|
||||
AUTHORIZATION,
|
||||
HeaderValue::from_str(&format!("Bearer {api_key}"))
|
||||
.map_err(|err| CoreError::Auth(err.to_string()))?,
|
||||
);
|
||||
|
||||
let (upstream, _response) = connect_async(request)
|
||||
.await
|
||||
.map_err(|err| CoreError::Network(err.to_string()))?;
|
||||
Ok(upstream)
|
||||
}
|
||||
|
||||
/// Read the next text frame from the upstream and decode it as a typed event.
|
||||
///
|
||||
/// Used by the pool to pre-read OpenAI's unprompted `session.created`. Returns an
|
||||
/// error on a non-text frame, a closed socket, or undecodable JSON so the pool can
|
||||
/// discard a misbehaving socket rather than warm it.
|
||||
pub(crate) async fn read_event(upstream_rx: &mut UpstreamRx) -> CoreResult<RealtimeEvent> {
|
||||
loop {
|
||||
let message = upstream_rx
|
||||
.next()
|
||||
.await
|
||||
.ok_or_else(|| CoreError::Network("upstream closed before first event".to_string()))?
|
||||
.map_err(|err| CoreError::Network(err.to_string()))?;
|
||||
match message {
|
||||
Message::Text(text) => {
|
||||
return serde_json::from_str(&text)
|
||||
.map_err(|err| CoreError::InvalidResponse(err.to_string()));
|
||||
}
|
||||
// Ignore protocol frames (ping/pong) while waiting for the first event.
|
||||
Message::Ping(_) | Message::Pong(_) => continue,
|
||||
Message::Close(_) => {
|
||||
return Err(CoreError::Network(
|
||||
"upstream closed before first event".to_string(),
|
||||
))
|
||||
}
|
||||
_ => continue,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Splice an already-connected upstream to the client streams.
|
||||
///
|
||||
/// `prelude` is relayed to the client first (the pool passes the buffered
|
||||
/// `session.created` here; the fresh-dial path passes `None` and lets the upstream
|
||||
/// deliver it). Then a single select loop forwards both directions through the
|
||||
/// transforms until either side closes or the idle timeout fires.
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub(crate) async fn splice<In, Out>(
|
||||
model: &str,
|
||||
mut upstream_tx: UpstreamTx,
|
||||
mut upstream_rx: UpstreamRx,
|
||||
prelude: Option<RealtimeEvent>,
|
||||
idle_timeout: Option<Duration>,
|
||||
mut client_in: In,
|
||||
mut client_out: Out,
|
||||
) -> CoreResult<()>
|
||||
where
|
||||
In: Stream<Item = RealtimeEvent> + Unpin + Send,
|
||||
Out: Sink<RealtimeEvent> + Unpin + Send,
|
||||
<Out as Sink<RealtimeEvent>>::Error: std::fmt::Display,
|
||||
{
|
||||
let config = &OPENAI_REALTIME_CONFIG;
|
||||
|
||||
// Relay a buffered backend event (warm handoff's session.created) first, so a
|
||||
// warm session looks identical to a fresh one from the client's view.
|
||||
if let Some(event) = prelude {
|
||||
for outbound in config.transform_realtime_response(&event, model)?.events {
|
||||
client_out
|
||||
.send(outbound)
|
||||
.await
|
||||
.map_err(|err| CoreError::Network(err.to_string()))?;
|
||||
}
|
||||
}
|
||||
|
||||
let idle = idle_timeout.unwrap_or(Duration::from_secs(DEFAULT_IDLE_TIMEOUT_SECS));
|
||||
|
||||
// One loop forwarding both directions. The `sleep(idle)` arm is rebuilt every
|
||||
// iteration, so any frame (either way) resets it — it fires only when the
|
||||
// session has been fully idle for `idle`, reaping a stalled connection
|
||||
// (task + upstream TCP socket) instead of leaking it.
|
||||
loop {
|
||||
tokio::select! {
|
||||
// client -> upstream
|
||||
client_event = client_in.next() => {
|
||||
let Some(event) = client_event else { break }; // client disconnected
|
||||
for outbound in config.transform_realtime_request(&event, model)?.events {
|
||||
let payload = serde_json::to_string(&outbound)
|
||||
.map_err(|err| CoreError::InvalidResponse(err.to_string()))?;
|
||||
upstream_tx
|
||||
.send(Message::Text(payload))
|
||||
.await
|
||||
.map_err(|err| CoreError::Network(err.to_string()))?;
|
||||
}
|
||||
}
|
||||
// upstream -> client
|
||||
upstream_message = upstream_rx.next() => {
|
||||
let Some(message) = upstream_message else { break }; // upstream closed
|
||||
match message.map_err(|err| CoreError::Network(err.to_string()))? {
|
||||
Message::Text(text) => {
|
||||
let event: RealtimeEvent = serde_json::from_str(&text)
|
||||
.map_err(|err| CoreError::InvalidResponse(err.to_string()))?;
|
||||
for outbound in config.transform_realtime_response(&event, model)?.events {
|
||||
client_out
|
||||
.send(outbound)
|
||||
.await
|
||||
.map_err(|err| CoreError::Network(err.to_string()))?;
|
||||
}
|
||||
}
|
||||
Message::Close(_) => break,
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
// idle timeout: no activity from either side within `idle`
|
||||
_ = tokio::time::sleep(idle) => break,
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Splice a client realtime stream to OpenAI: forward client events upstream
|
||||
/// (via `transform_realtime_request`) and backend events downstream (via
|
||||
/// `transform_realtime_response`). Returns when either side closes.
|
||||
///
|
||||
/// Generic over the client transport (typed events) so this crate stays
|
||||
/// framework-agnostic; the gateway adapts its axum socket to these. This is the
|
||||
/// fresh-dial path: dial, then splice. The pool's warm-handoff path skips the dial
|
||||
/// and calls [`splice`] directly with a buffered `session.created`.
|
||||
pub async fn realtime<In, Out>(
|
||||
model: &str,
|
||||
api_key: Option<&str>,
|
||||
api_base: Option<&str>,
|
||||
idle_timeout: Option<Duration>,
|
||||
client_in: In,
|
||||
client_out: Out,
|
||||
) -> CoreResult<()>
|
||||
where
|
||||
In: Stream<Item = RealtimeEvent> + Unpin + Send,
|
||||
Out: Sink<RealtimeEvent> + Unpin + Send,
|
||||
<Out as Sink<RealtimeEvent>>::Error: std::fmt::Display,
|
||||
{
|
||||
let api_key = resolve_api_key(api_key)?;
|
||||
let upstream = dial_upstream(model, &api_key, api_base).await?;
|
||||
let (upstream_tx, upstream_rx) = upstream.split();
|
||||
splice(
|
||||
model,
|
||||
upstream_tx,
|
||||
upstream_rx,
|
||||
None,
|
||||
idle_timeout,
|
||||
client_in,
|
||||
client_out,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
/// Splice a pre-warmed upstream (taken from [`crate::realtime_pool`]) to the
|
||||
/// client. Relays the buffered `session.created` first, then splices exactly like
|
||||
/// the fresh-dial path — so a warm session is indistinguishable from a fresh one.
|
||||
pub async fn realtime_warm<In, Out>(
|
||||
model: &str,
|
||||
handoff: crate::realtime_pool::WarmHandoff,
|
||||
idle_timeout: Option<Duration>,
|
||||
client_in: In,
|
||||
client_out: Out,
|
||||
) -> CoreResult<()>
|
||||
where
|
||||
In: Stream<Item = RealtimeEvent> + Unpin + Send,
|
||||
Out: Sink<RealtimeEvent> + Unpin + Send,
|
||||
<Out as Sink<RealtimeEvent>>::Error: std::fmt::Display,
|
||||
{
|
||||
splice(
|
||||
model,
|
||||
handoff.tx,
|
||||
handoff.rx,
|
||||
Some(handoff.session_created),
|
||||
idle_timeout,
|
||||
client_in,
|
||||
client_out,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn event(raw: &str) -> RealtimeEvent {
|
||||
serde_json::from_str(raw).expect("valid event json")
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolve_api_key_prefers_param_then_blank_falls_through() {
|
||||
assert_eq!(resolve_api_key(Some("sk-test")).unwrap(), "sk-test");
|
||||
// A blank param with no env set should error.
|
||||
if std::env::var(OPENAI_API_KEY_ENV).is_err() {
|
||||
assert!(resolve_api_key(Some(" ")).is_err());
|
||||
}
|
||||
}
|
||||
|
||||
/// Live end-to-end check against OpenAI. Ignored by default (CI never runs
|
||||
/// it); run explicitly with `OPENAI_API_KEY` set:
|
||||
/// `cargo test -p litellm-providers realtime_invokes_openai -- --ignored --nocapture`
|
||||
#[tokio::test]
|
||||
#[ignore = "hits the live OpenAI realtime API; needs OPENAI_API_KEY"]
|
||||
async fn realtime_invokes_openai_and_responds() {
|
||||
use futures_channel::mpsc;
|
||||
|
||||
let key =
|
||||
std::env::var(OPENAI_API_KEY_ENV).expect("set OPENAI_API_KEY to run this ignored test");
|
||||
|
||||
// client -> provider (we hold `client_tx` to push events upstream)
|
||||
let (mut client_tx, client_in) = mpsc::unbounded::<RealtimeEvent>();
|
||||
// provider -> client (we hold `backend_rx` to read backend events)
|
||||
let (client_out, mut backend_rx) = mpsc::unbounded::<RealtimeEvent>();
|
||||
|
||||
// Clone the key so the spawned task owns its `String` (no borrow across await).
|
||||
let key_owned = key.clone();
|
||||
let call = tokio::spawn(async move {
|
||||
realtime(
|
||||
"gpt-realtime",
|
||||
Some(&key_owned),
|
||||
None,
|
||||
None,
|
||||
client_in,
|
||||
client_out,
|
||||
)
|
||||
.await
|
||||
});
|
||||
|
||||
// 1. First backend event should be session.created.
|
||||
let first = tokio::time::timeout(Duration::from_secs(30), backend_rx.next())
|
||||
.await
|
||||
.expect("timed out waiting for session.created")
|
||||
.expect("backend stream closed before session.created");
|
||||
assert_eq!(
|
||||
first.event_type, "session.created",
|
||||
"expected session.created, got: {}",
|
||||
first.event_type
|
||||
);
|
||||
|
||||
// 2. Ask for a short audio response.
|
||||
client_tx
|
||||
.send(event(
|
||||
r#"{"type":"conversation.item.create","item":{"type":"message","role":"user","content":[{"type":"input_text","text":"Say hi."}]}}"#,
|
||||
))
|
||||
.await
|
||||
.expect("send conversation.item.create");
|
||||
client_tx
|
||||
.send(event(r#"{"type":"response.create"}"#))
|
||||
.await
|
||||
.expect("send response.create");
|
||||
|
||||
// 3. Read backend events; require a non-empty audio delta, then response.done.
|
||||
let mut saw_audio_delta = false;
|
||||
let mut saw_done = false;
|
||||
for _ in 0..500 {
|
||||
let next = tokio::time::timeout(Duration::from_secs(30), backend_rx.next()).await;
|
||||
let event = match next {
|
||||
Ok(Some(event)) => event,
|
||||
Ok(None) => break,
|
||||
Err(_) => panic!("timed out waiting for backend events"),
|
||||
};
|
||||
match event.event_type.as_str() {
|
||||
"response.output_audio.delta" => {
|
||||
let delta = event
|
||||
.data
|
||||
.get("delta")
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or("");
|
||||
if !delta.is_empty() {
|
||||
saw_audio_delta = true;
|
||||
}
|
||||
}
|
||||
"response.done" => {
|
||||
saw_done = true;
|
||||
break;
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
|
||||
assert!(
|
||||
saw_audio_delta,
|
||||
"expected a response.output_audio.delta with non-empty delta"
|
||||
);
|
||||
assert!(saw_done, "expected a response.done event");
|
||||
|
||||
// Drop the client sender so the provider's to_upstream side finishes.
|
||||
drop(client_tx);
|
||||
let _ = call.await;
|
||||
}
|
||||
}
|
||||
712
litellm-rust/crates/providers/src/realtime_pool.rs
Normal file
712
litellm-rust/crates/providers/src/realtime_pool.rs
Normal file
|
|
@ -0,0 +1,712 @@
|
|||
//! Pre-warmed upstream realtime connection pool.
|
||||
//!
|
||||
//! The gateway's realtime overhead lives entirely in session establishment: on
|
||||
//! every client connect it dials a fresh upstream WS to OpenAI and waits for
|
||||
//! `session.created` before it can serve. This pool keeps a small set of upstream
|
||||
//! sockets **already connected and already past `session.created`** so a connect
|
||||
//! can be served from a warm socket and the handshake is off the critical path.
|
||||
//!
|
||||
//! Layering: this stays in `providers` (axum-free) next to the dial/splice it
|
||||
//! reuses. The gateway holds an `Arc<RealtimePool>` in its state and asks for a
|
||||
//! warm socket per connect; on a miss it fresh-dials exactly as before. The pool
|
||||
//! is a latency optimization, never a correctness dependency — see the gateway's
|
||||
//! `src/routes/realtime/README.md`.
|
||||
//!
|
||||
//! ## Caveats (enforced here)
|
||||
//! - One warm socket serves exactly one session (realtime isn't multiplexed), so
|
||||
//! the pool is sized to the connect *rate*, not concurrent connections.
|
||||
//! - `session.created` is pre-read once and buffered; nothing else is read from a
|
||||
//! warm socket before handoff, so a warm session starts at OpenAI defaults just
|
||||
//! like a fresh one (`session.update` semantics unchanged).
|
||||
//! - Warm sockets are short-lived (`max_idle`) and liveness-checked at handoff to
|
||||
//! bound idle billing / dodge OpenAI's idle timeout.
|
||||
//! - On miss or dead socket the caller fresh-dials; the pool never blocks or fails
|
||||
//! a connect because it is empty.
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::sync::{Arc, Mutex};
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
use futures_util::StreamExt;
|
||||
use litellm_core::realtime::types::RealtimeEvent;
|
||||
use litellm_core::CoreResult;
|
||||
|
||||
use crate::realtime::{
|
||||
dial_upstream, read_event, resolve_api_key, UpstreamRx, UpstreamTx, UpstreamWs,
|
||||
};
|
||||
|
||||
/// Default target warm sockets per key when pooling is enabled.
|
||||
pub const DEFAULT_POOL_SIZE: usize = 4;
|
||||
|
||||
/// Default max time a warm socket may sit before it is closed and replaced.
|
||||
pub const DEFAULT_MAX_IDLE: Duration = Duration::from_secs(30);
|
||||
|
||||
/// Env var: target warm sockets per key. `0` disables pooling (fresh-dial only).
|
||||
pub const POOL_SIZE_ENV: &str = "REALTIME_POOL_SIZE";
|
||||
|
||||
/// Env var: max warm-socket idle lifetime, in seconds.
|
||||
pub const MAX_IDLE_ENV: &str = "REALTIME_POOL_MAX_IDLE_SECS";
|
||||
|
||||
/// How often the background replenisher wakes to top up and reap stale sockets.
|
||||
const REPLENISH_TICK: Duration = Duration::from_millis(250);
|
||||
|
||||
/// Backoff floor after a key's warm-up dials all fail. The first failed pass
|
||||
/// waits this long before retrying that key.
|
||||
const BACKOFF_BASE: Duration = Duration::from_millis(500);
|
||||
|
||||
/// Backoff ceiling. A key that keeps failing (invalid credentials, an
|
||||
/// unreachable upstream) is retried at most once per this interval — instead of
|
||||
/// firing `needed` concurrent TLS dials every 250 ms tick, which would hammer
|
||||
/// the upstream and risk rate-limit exhaustion that degrades valid cold-path
|
||||
/// traffic. Backoff resets the moment a dial for the key succeeds.
|
||||
const BACKOFF_MAX: Duration = Duration::from_secs(30);
|
||||
|
||||
/// Identifies an upstream connection: the tuple that fully determines the dial.
|
||||
/// `api_key` is included so a warm socket is only ever reused for the same key
|
||||
/// (no cross-tenant reuse).
|
||||
#[derive(Clone, PartialEq, Eq, Hash)]
|
||||
pub struct UpstreamKey {
|
||||
pub model: String,
|
||||
pub api_key: String,
|
||||
pub api_base: Option<String>,
|
||||
}
|
||||
|
||||
impl std::fmt::Debug for UpstreamKey {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
f.debug_struct("UpstreamKey")
|
||||
.field("model", &self.model)
|
||||
.field("api_key", &"[REDACTED]")
|
||||
.field("api_base", &self.api_base)
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
|
||||
/// A warm upstream: split halves + the buffered `session.created` + when it was
|
||||
/// warmed (for `max_idle` expiry).
|
||||
struct WarmConnection {
|
||||
tx: UpstreamTx,
|
||||
rx: UpstreamRx,
|
||||
session_created: RealtimeEvent,
|
||||
warmed_at: Instant,
|
||||
}
|
||||
|
||||
/// A live upstream taken from the pool, ready to splice. The caller relays
|
||||
/// `session_created` to the client first, then splices `(tx, rx)` as usual.
|
||||
pub struct WarmHandoff {
|
||||
pub tx: UpstreamTx,
|
||||
pub rx: UpstreamRx,
|
||||
pub session_created: RealtimeEvent,
|
||||
}
|
||||
|
||||
/// Pool configuration, resolved once at startup from the environment.
|
||||
#[derive(Clone, Copy, Debug)]
|
||||
pub struct PoolConfig {
|
||||
/// Target warm sockets per key. `0` disables pooling.
|
||||
pub target_size: usize,
|
||||
/// Max time a warm socket may sit before it is closed and replaced.
|
||||
pub max_idle: Duration,
|
||||
}
|
||||
|
||||
impl Default for PoolConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
target_size: DEFAULT_POOL_SIZE,
|
||||
max_idle: DEFAULT_MAX_IDLE,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl PoolConfig {
|
||||
/// Read config from the environment, falling back to defaults. An invalid
|
||||
/// value warns and uses the default rather than failing startup.
|
||||
pub fn from_env() -> Self {
|
||||
let target_size = match std::env::var(POOL_SIZE_ENV) {
|
||||
Ok(raw) => raw.trim().parse().unwrap_or_else(|_| {
|
||||
eprintln!("warning: {POOL_SIZE_ENV}={raw:?} is not a valid size; using {DEFAULT_POOL_SIZE}");
|
||||
DEFAULT_POOL_SIZE
|
||||
}),
|
||||
Err(_) => DEFAULT_POOL_SIZE,
|
||||
};
|
||||
let max_idle = match std::env::var(MAX_IDLE_ENV) {
|
||||
Ok(raw) => raw
|
||||
.trim()
|
||||
.parse()
|
||||
.map(Duration::from_secs)
|
||||
.unwrap_or_else(|_| {
|
||||
eprintln!(
|
||||
"warning: {MAX_IDLE_ENV}={raw:?} is not a valid number of seconds; using {}s",
|
||||
DEFAULT_MAX_IDLE.as_secs()
|
||||
);
|
||||
DEFAULT_MAX_IDLE
|
||||
}),
|
||||
Err(_) => DEFAULT_MAX_IDLE,
|
||||
};
|
||||
Self {
|
||||
target_size,
|
||||
max_idle,
|
||||
}
|
||||
}
|
||||
|
||||
/// Whether pooling is on (`target_size > 0`).
|
||||
pub fn enabled(&self) -> bool {
|
||||
self.target_size > 0
|
||||
}
|
||||
}
|
||||
|
||||
/// Per-key warm sockets, behind a single `Mutex`. Realtime warm sockets are few
|
||||
/// (the pool is small), so a plain mutex over a `VecDeque`-ish `Vec` is simpler
|
||||
/// and faster than sharding; contention is negligible at this scale.
|
||||
type Warm = HashMap<UpstreamKey, Vec<WarmConnection>>;
|
||||
|
||||
/// Per-key replenish backoff. Absent (or `consecutive_failures == 0`) means the
|
||||
/// key is healthy and replenished every tick. After a pass whose dials all fail,
|
||||
/// `retry_after` is pushed out with exponential backoff so a broken key (invalid
|
||||
/// credentials, unreachable upstream) is not re-dialed on every 250 ms tick.
|
||||
#[derive(Default)]
|
||||
struct Backoff {
|
||||
/// Don't attempt warm-up dials for this key until this instant. `None` =
|
||||
/// eligible now.
|
||||
retry_after: Option<Instant>,
|
||||
consecutive_failures: u32,
|
||||
}
|
||||
|
||||
type Backoffs = HashMap<UpstreamKey, Backoff>;
|
||||
|
||||
/// Pre-warmed upstream realtime connection pool.
|
||||
///
|
||||
/// Cheap to clone-via-`Arc`. The background replenisher is spawned by
|
||||
/// [`RealtimePool::spawn`]; a pool built with [`RealtimePool::disabled`] never
|
||||
/// warms anything and every `take` misses (callers fresh-dial).
|
||||
pub struct RealtimePool {
|
||||
config: PoolConfig,
|
||||
warm: Mutex<Warm>,
|
||||
/// Per-key replenish backoff so a broken key doesn't trigger unbounded
|
||||
/// concurrent dials every tick. Separate lock from `warm` so the request
|
||||
/// hot path (`take`) never contends on it.
|
||||
backoff: Mutex<Backoffs>,
|
||||
}
|
||||
|
||||
impl RealtimePool {
|
||||
/// A disabled pool: no background task, every `take` returns `None`.
|
||||
pub fn disabled() -> Arc<Self> {
|
||||
Arc::new(Self {
|
||||
config: PoolConfig {
|
||||
target_size: 0,
|
||||
..PoolConfig::default()
|
||||
},
|
||||
warm: Mutex::new(HashMap::new()),
|
||||
backoff: Mutex::new(HashMap::new()),
|
||||
})
|
||||
}
|
||||
|
||||
/// Build a pool from config **without** the background replenisher. The pool
|
||||
/// only warms when [`RealtimePool::warm_now`] is called. Used by deterministic
|
||||
/// unit tests; production uses [`RealtimePool::spawn`].
|
||||
#[cfg(test)]
|
||||
fn new_unspawned(config: PoolConfig) -> Arc<Self> {
|
||||
Arc::new(Self {
|
||||
config,
|
||||
warm: Mutex::new(HashMap::new()),
|
||||
backoff: Mutex::new(HashMap::new()),
|
||||
})
|
||||
}
|
||||
|
||||
/// Build a pool from config and, if enabled, spawn the background replenisher.
|
||||
/// Returns the shared handle the gateway stores in its state.
|
||||
pub fn spawn(config: PoolConfig) -> Arc<Self> {
|
||||
let pool = Arc::new(Self {
|
||||
config,
|
||||
warm: Mutex::new(HashMap::new()),
|
||||
backoff: Mutex::new(HashMap::new()),
|
||||
});
|
||||
if config.enabled() {
|
||||
let weak = Arc::downgrade(&pool);
|
||||
tokio::spawn(async move {
|
||||
let mut tick = tokio::time::interval(REPLENISH_TICK);
|
||||
tick.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
|
||||
loop {
|
||||
tick.tick().await;
|
||||
// Stop once the gateway has dropped its handle.
|
||||
let Some(pool) = weak.upgrade() else { break };
|
||||
pool.replenish_all().await;
|
||||
}
|
||||
});
|
||||
}
|
||||
pool
|
||||
}
|
||||
|
||||
/// Resolved config (test/inspection).
|
||||
pub fn config(&self) -> PoolConfig {
|
||||
self.config
|
||||
}
|
||||
|
||||
/// Register a key so the replenisher starts warming it. Idempotent. The
|
||||
/// gateway calls this once per known deployment at startup; the pool only
|
||||
/// warms keys it has seen, so it never dials a model nobody asked for.
|
||||
pub fn register(&self, key: UpstreamKey) {
|
||||
if !self.config.enabled() {
|
||||
return;
|
||||
}
|
||||
self.warm.lock().unwrap().entry(key).or_default();
|
||||
}
|
||||
|
||||
/// Take a warm, live socket for `key`, or `None` on miss / dead socket.
|
||||
///
|
||||
/// Pops the freshest non-expired socket and liveness-checks it; a socket that
|
||||
/// is too old or already dead is dropped (closing it) and the next candidate
|
||||
/// tried. Never blocks: if nothing warm is live, returns `None` so the caller
|
||||
/// fresh-dials.
|
||||
pub fn take(&self, key: &UpstreamKey) -> Option<WarmHandoff> {
|
||||
if !self.config.enabled() {
|
||||
return None;
|
||||
}
|
||||
loop {
|
||||
let mut candidate = {
|
||||
let mut warm = self.warm.lock().unwrap();
|
||||
let bucket = warm.get_mut(key)?;
|
||||
bucket.pop()?
|
||||
};
|
||||
// Discard sockets past their warm lifetime (idle-billing guard).
|
||||
if candidate.warmed_at.elapsed() > self.config.max_idle {
|
||||
continue; // drops `candidate`, closing the socket
|
||||
}
|
||||
// Liveness: a non-blocking check that the socket hasn't already
|
||||
// delivered a Close/Err. A warm socket should be silent after
|
||||
// session.created, so anything pending means it is unhealthy.
|
||||
if is_dead(&mut candidate.rx) {
|
||||
continue;
|
||||
}
|
||||
return Some(WarmHandoff {
|
||||
tx: candidate.tx,
|
||||
rx: candidate.rx,
|
||||
session_created: candidate.session_created,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
/// One replenish pass over every registered key: reap stale sockets, then
|
||||
/// dial up to `target_size`. Dials run concurrently; failures are swallowed
|
||||
/// (a key that can't be warmed just keeps fresh-dialing on the request path)
|
||||
/// and put the key into exponential backoff so a broken key isn't re-dialed
|
||||
/// on every tick.
|
||||
async fn replenish_all(&self) {
|
||||
let keys: Vec<UpstreamKey> = { self.warm.lock().unwrap().keys().cloned().collect() };
|
||||
for key in keys {
|
||||
self.reap_stale(&key);
|
||||
// Skip keys still in backoff from a prior all-failed pass — this is
|
||||
// what bounds dials against an invalid/unreachable key to once per
|
||||
// `BACKOFF_MAX` instead of `needed` dials every 250 ms tick.
|
||||
if self.in_backoff(&key) {
|
||||
continue;
|
||||
}
|
||||
let needed = {
|
||||
let warm = self.warm.lock().unwrap();
|
||||
let have = warm.get(&key).map(Vec::len).unwrap_or(0);
|
||||
self.config.target_size.saturating_sub(have)
|
||||
};
|
||||
if needed == 0 {
|
||||
continue;
|
||||
}
|
||||
// Dial the missing sockets CONCURRENTLY. A sequential loop here makes
|
||||
// a full refill cost `needed × handshake` (~needed × 350 ms), which
|
||||
// can't keep up with a high connect rate — the pool drains faster
|
||||
// than it refills and most connects miss. Firing the dials together
|
||||
// refills in ~one handshake window, keeping warm supply ≈ peak
|
||||
// concurrent connects so the sub-ms warm handoff becomes the median,
|
||||
// not the lucky-hit tail.
|
||||
let dials = (0..needed).map(|_| warm_one(&key));
|
||||
let results = futures_util::future::join_all(dials).await;
|
||||
let mut any_ok = false;
|
||||
// `.flatten()` keeps only the successful dials; a key that can't be
|
||||
// warmed just keeps fresh-dialing on the request path.
|
||||
for conn in results.into_iter().flatten() {
|
||||
any_ok = true;
|
||||
self.warm
|
||||
.lock()
|
||||
.unwrap()
|
||||
.entry(key.clone())
|
||||
.or_default()
|
||||
.push(conn);
|
||||
}
|
||||
// Reset backoff on any success; otherwise grow it. We only ever enter
|
||||
// backoff when a pass that *attempted* dials produced none — a `needed
|
||||
// == 0` pass is handled by the `continue` above and never touches it.
|
||||
self.record_replenish_outcome(&key, any_ok);
|
||||
}
|
||||
}
|
||||
|
||||
/// Whether `key` is currently in a backoff window (a prior pass failed and
|
||||
/// the retry time hasn't arrived). Eligible keys are pruned from the backoff
|
||||
/// map so it doesn't grow unbounded for healthy keys.
|
||||
fn in_backoff(&self, key: &UpstreamKey) -> bool {
|
||||
let mut backoff = self.backoff.lock().unwrap();
|
||||
match backoff.get(key).and_then(|b| b.retry_after) {
|
||||
Some(retry_after) if Instant::now() < retry_after => true,
|
||||
Some(_) => {
|
||||
// Window elapsed — allow the attempt. Keep the failure count so a
|
||||
// still-broken key backs off further, but clear the gate so this
|
||||
// tick proceeds.
|
||||
if let Some(b) = backoff.get_mut(key) {
|
||||
b.retry_after = None;
|
||||
}
|
||||
false
|
||||
}
|
||||
None => false,
|
||||
}
|
||||
}
|
||||
|
||||
/// Update a key's backoff after a replenish attempt. Success clears it;
|
||||
/// failure grows the retry delay exponentially up to `BACKOFF_MAX`.
|
||||
fn record_replenish_outcome(&self, key: &UpstreamKey, any_ok: bool) {
|
||||
let mut backoff = self.backoff.lock().unwrap();
|
||||
if any_ok {
|
||||
backoff.remove(key);
|
||||
return;
|
||||
}
|
||||
let entry = backoff.entry(key.clone()).or_default();
|
||||
entry.consecutive_failures = entry.consecutive_failures.saturating_add(1);
|
||||
// Exponential: BASE * 2^(failures-1), saturating at MAX. `min` of the
|
||||
// shift exponent keeps the doubling from overflowing.
|
||||
let shift = (entry.consecutive_failures - 1).min(16);
|
||||
let delay = BACKOFF_BASE.saturating_mul(1u32 << shift).min(BACKOFF_MAX);
|
||||
entry.retry_after = Some(Instant::now() + delay);
|
||||
}
|
||||
|
||||
/// Drop sockets past `max_idle` or already dead for a key.
|
||||
fn reap_stale(&self, key: &UpstreamKey) {
|
||||
let mut warm = self.warm.lock().unwrap();
|
||||
if let Some(bucket) = warm.get_mut(key) {
|
||||
bucket.retain_mut(|conn| {
|
||||
conn.warmed_at.elapsed() <= self.config.max_idle && !is_dead(&mut conn.rx)
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
/// Test/inspection: number of warm sockets currently held for `key`.
|
||||
#[cfg(test)]
|
||||
pub fn warm_len(&self, key: &UpstreamKey) -> usize {
|
||||
self.warm
|
||||
.lock()
|
||||
.unwrap()
|
||||
.get(key)
|
||||
.map(Vec::len)
|
||||
.unwrap_or(0)
|
||||
}
|
||||
|
||||
/// Test/inspection: consecutive replenish failures recorded for `key` (0 if
|
||||
/// the key is healthy / has no backoff entry).
|
||||
#[cfg(test)]
|
||||
pub fn backoff_failures(&self, key: &UpstreamKey) -> u32 {
|
||||
self.backoff
|
||||
.lock()
|
||||
.unwrap()
|
||||
.get(key)
|
||||
.map(|b| b.consecutive_failures)
|
||||
.unwrap_or(0)
|
||||
}
|
||||
|
||||
/// Test helper: synchronously warm `target_size` sockets for `key` (no
|
||||
/// background task). Lets tests assert handoff behavior deterministically.
|
||||
#[cfg(test)]
|
||||
pub async fn warm_now(&self, key: &UpstreamKey) {
|
||||
let needed = {
|
||||
let warm = self.warm.lock().unwrap();
|
||||
let have = warm.get(key).map(Vec::len).unwrap_or(0);
|
||||
self.config.target_size.saturating_sub(have)
|
||||
};
|
||||
for _ in 0..needed {
|
||||
if let Ok(conn) = warm_one(key).await {
|
||||
self.warm
|
||||
.lock()
|
||||
.unwrap()
|
||||
.entry(key.clone())
|
||||
.or_default()
|
||||
.push(conn);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Test helper: insert an already-built warm connection (used to inject a
|
||||
/// dead socket and assert it is discarded at handoff).
|
||||
#[cfg(test)]
|
||||
fn insert_warm(&self, key: UpstreamKey, conn: WarmConnection) {
|
||||
self.warm.lock().unwrap().entry(key).or_default().push(conn);
|
||||
}
|
||||
}
|
||||
|
||||
/// Dial one upstream and pre-read its `session.created` into a [`WarmConnection`].
|
||||
///
|
||||
/// `key.api_key` is already resolved (non-blank). The first frame OpenAI sends
|
||||
/// unprompted is `session.created`; we buffer exactly that and read nothing more.
|
||||
async fn warm_one(key: &UpstreamKey) -> CoreResult<WarmConnection> {
|
||||
let upstream: UpstreamWs =
|
||||
dial_upstream(&key.model, &key.api_key, key.api_base.as_deref()).await?;
|
||||
let (tx, mut rx) = upstream.split();
|
||||
let session_created = read_event(&mut rx).await?;
|
||||
Ok(WarmConnection {
|
||||
tx,
|
||||
rx,
|
||||
session_created,
|
||||
warmed_at: Instant::now(),
|
||||
})
|
||||
}
|
||||
|
||||
/// Resolve a deployment's API key into the pool key, returning `None` when no key
|
||||
/// can be resolved (those deployments simply aren't pooled — the request path
|
||||
/// still fresh-dials and surfaces the auth error there).
|
||||
pub fn upstream_key(
|
||||
model: &str,
|
||||
api_key: Option<&str>,
|
||||
api_base: Option<&str>,
|
||||
) -> Option<UpstreamKey> {
|
||||
let api_key = resolve_api_key(api_key).ok()?;
|
||||
Some(UpstreamKey {
|
||||
model: model.to_string(),
|
||||
api_key,
|
||||
api_base: api_base.map(str::to_string),
|
||||
})
|
||||
}
|
||||
|
||||
/// Non-blocking liveness check: poll the upstream once. A warm socket is silent
|
||||
/// after `session.created`, so a pending `Close`/`Err`/`None` means it is dead.
|
||||
/// A pending data frame (shouldn't happen pre-handoff) is also treated as
|
||||
/// unhealthy — we'd rather discard and fresh-dial than hand over a socket in an
|
||||
/// unexpected state. `Pending` (the healthy case) returns `false`.
|
||||
fn is_dead(rx: &mut UpstreamRx) -> bool {
|
||||
use futures_util::task::noop_waker_ref;
|
||||
use futures_util::Stream;
|
||||
use std::pin::Pin;
|
||||
use std::task::{Context, Poll};
|
||||
|
||||
let mut cx = Context::from_waker(noop_waker_ref());
|
||||
match Pin::new(rx).poll_next(&mut cx) {
|
||||
Poll::Pending => false,
|
||||
Poll::Ready(None) => true,
|
||||
Poll::Ready(Some(Err(_))) => true,
|
||||
// Any frame arriving before handoff is unexpected for a silent warm
|
||||
// socket; treat it as unhealthy.
|
||||
Poll::Ready(Some(Ok(_))) => true,
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use futures_util::SinkExt;
|
||||
use std::net::SocketAddr;
|
||||
use tokio::net::TcpListener;
|
||||
use tokio_tungstenite::tungstenite::Message;
|
||||
|
||||
/// An in-process fake OpenAI realtime WS server. On connect it sends
|
||||
/// `session.created`; on `response.create` it sends `response.created` +
|
||||
/// `response.output_audio.delta` + `response.done`. Returns its `ws://` base.
|
||||
async fn spawn_fake_openai() -> String {
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let addr: SocketAddr = listener.local_addr().unwrap();
|
||||
tokio::spawn(async move {
|
||||
while let Ok((stream, _)) = listener.accept().await {
|
||||
tokio::spawn(handle_fake_conn(stream));
|
||||
}
|
||||
});
|
||||
format!("ws://{addr}")
|
||||
}
|
||||
|
||||
async fn handle_fake_conn(stream: tokio::net::TcpStream) {
|
||||
let mut ws = match tokio_tungstenite::accept_async(stream).await {
|
||||
Ok(ws) => ws,
|
||||
Err(_) => return,
|
||||
};
|
||||
// Unprompted session.created, exactly like OpenAI.
|
||||
let _ = ws
|
||||
.send(Message::Text(
|
||||
r#"{"type":"session.created","session":{"id":"sess_fake"}}"#.to_string(),
|
||||
))
|
||||
.await;
|
||||
while let Some(Ok(msg)) = ws.next().await {
|
||||
if let Message::Text(text) = msg {
|
||||
if text.contains("response.create") {
|
||||
for frame in [
|
||||
r#"{"type":"response.created"}"#,
|
||||
r#"{"type":"response.output_audio.delta","delta":"AAAA"}"#,
|
||||
r#"{"type":"response.done"}"#,
|
||||
] {
|
||||
let _ = ws.send(Message::Text(frame.to_string())).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn test_config() -> PoolConfig {
|
||||
PoolConfig {
|
||||
target_size: 2,
|
||||
max_idle: Duration::from_secs(30),
|
||||
}
|
||||
}
|
||||
|
||||
fn key_for(base: &str) -> UpstreamKey {
|
||||
UpstreamKey {
|
||||
model: "gpt-realtime".to_string(),
|
||||
api_key: "sk-test".to_string(),
|
||||
api_base: Some(base.to_string()),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn warm_handoff_relays_buffered_session_created() {
|
||||
let base = spawn_fake_openai().await;
|
||||
let pool = RealtimePool::new_unspawned(test_config());
|
||||
let key = key_for(&base);
|
||||
pool.register(key.clone());
|
||||
pool.warm_now(&key).await;
|
||||
assert_eq!(pool.warm_len(&key), 2);
|
||||
|
||||
let handoff = pool.take(&key).expect("a warm socket should be available");
|
||||
assert_eq!(handoff.session_created.event_type, "session.created");
|
||||
assert_eq!(
|
||||
handoff
|
||||
.session_created
|
||||
.data
|
||||
.get("session")
|
||||
.and_then(|s| s.get("id"))
|
||||
.and_then(|v| v.as_str()),
|
||||
Some("sess_fake")
|
||||
);
|
||||
// Taking one leaves one.
|
||||
assert_eq!(pool.warm_len(&key), 1);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn pool_miss_returns_none_for_fresh_dial_fallback() {
|
||||
let base = spawn_fake_openai().await;
|
||||
let pool = RealtimePool::new_unspawned(test_config());
|
||||
let key = key_for(&base);
|
||||
// Registered but never warmed → empty bucket → miss.
|
||||
pool.register(key.clone());
|
||||
assert!(pool.take(&key).is_none());
|
||||
|
||||
// Unknown key → miss.
|
||||
let other = key_for("ws://127.0.0.1:1");
|
||||
assert!(pool.take(&other).is_none());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn disabled_pool_never_hands_off() {
|
||||
let pool = RealtimePool::disabled();
|
||||
let key = key_for("ws://127.0.0.1:1");
|
||||
pool.register(key.clone());
|
||||
assert_eq!(pool.warm_len(&key), 0);
|
||||
assert!(pool.take(&key).is_none());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn dead_warm_socket_is_discarded() {
|
||||
let base = spawn_fake_openai().await;
|
||||
let pool = RealtimePool::new_unspawned(test_config());
|
||||
let key = key_for(&base);
|
||||
pool.register(key.clone());
|
||||
|
||||
// Build one real warm connection, then kill the upstream by dropping the
|
||||
// server side: easiest is to dial, read session.created, then close our
|
||||
// own rx's peer. Instead we forge "dead" via an already-closed socket:
|
||||
// dial a connection and immediately send a Close from the client side so
|
||||
// the server closes back, then warm it. Simpler: warm normally, then
|
||||
// mark it stale by backdating warmed_at past max_idle and confirm it's
|
||||
// dropped — that exercises the same discard path.
|
||||
let mut conn = warm_one(&key).await.expect("warm one");
|
||||
conn.warmed_at = Instant::now() - Duration::from_secs(3600); // past max_idle
|
||||
pool.insert_warm(key.clone(), conn);
|
||||
assert_eq!(pool.warm_len(&key), 1);
|
||||
|
||||
// take() must discard the stale socket and report a miss.
|
||||
assert!(pool.take(&key).is_none());
|
||||
assert_eq!(pool.warm_len(&key), 0);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn background_replenisher_tops_up_registered_key() {
|
||||
let base = spawn_fake_openai().await;
|
||||
let pool = RealtimePool::spawn(test_config());
|
||||
let key = key_for(&base);
|
||||
pool.register(key.clone());
|
||||
|
||||
// Wait (bounded) for the background task to reach the target size.
|
||||
let mut warmed = 0;
|
||||
for _ in 0..40 {
|
||||
tokio::time::sleep(Duration::from_millis(50)).await;
|
||||
warmed = pool.warm_len(&key);
|
||||
if warmed >= test_config().target_size {
|
||||
break;
|
||||
}
|
||||
}
|
||||
assert_eq!(
|
||||
warmed,
|
||||
test_config().target_size,
|
||||
"background replenisher should warm up to target_size"
|
||||
);
|
||||
let handoff = pool.take(&key).expect("a warm socket should be available");
|
||||
assert_eq!(handoff.session_created.event_type, "session.created");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn closed_upstream_socket_is_detected_dead() {
|
||||
// A genuinely dead socket: dial the fake, read session.created, then drop
|
||||
// the server by closing from our side and waiting for the close to land.
|
||||
let base = spawn_fake_openai().await;
|
||||
let pool = RealtimePool::new_unspawned(test_config());
|
||||
let key = key_for(&base);
|
||||
pool.register(key.clone());
|
||||
|
||||
let mut conn = warm_one(&key).await.expect("warm one");
|
||||
// Close the upstream from the client side; the server echoes a close.
|
||||
let _ = conn.tx.send(Message::Close(None)).await;
|
||||
// Give the close a moment to arrive on rx.
|
||||
tokio::time::sleep(Duration::from_millis(50)).await;
|
||||
pool.insert_warm(key.clone(), conn);
|
||||
|
||||
// Liveness check at take() should detect the close and discard it.
|
||||
assert!(pool.take(&key).is_none());
|
||||
assert_eq!(pool.warm_len(&key), 0);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn broken_key_backs_off_instead_of_dialing_every_tick() {
|
||||
// A key whose upstream is unreachable: every warm-up dial fails.
|
||||
let pool = RealtimePool::new_unspawned(test_config());
|
||||
let key = key_for("ws://127.0.0.1:1"); // nothing listens here
|
||||
pool.register(key.clone());
|
||||
|
||||
// First pass attempts dials, they all fail → key enters backoff, no warm
|
||||
// sockets, one recorded failure.
|
||||
pool.replenish_all().await;
|
||||
assert_eq!(pool.warm_len(&key), 0);
|
||||
assert_eq!(pool.backoff_failures(&key), 1);
|
||||
assert!(
|
||||
pool.in_backoff(&key),
|
||||
"a key whose dials all failed must be in backoff"
|
||||
);
|
||||
|
||||
// An immediate next pass must be SKIPPED (still in the backoff window), so
|
||||
// it does NOT fire another round of dials — the failure count is unchanged.
|
||||
pool.replenish_all().await;
|
||||
assert_eq!(
|
||||
pool.backoff_failures(&key),
|
||||
1,
|
||||
"replenish during the backoff window must not re-dial the broken key"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn healthy_key_never_enters_backoff_and_clears_after_recovery() {
|
||||
let base = spawn_fake_openai().await;
|
||||
let pool = RealtimePool::new_unspawned(test_config());
|
||||
let key = key_for(&base);
|
||||
pool.register(key.clone());
|
||||
|
||||
// A reachable upstream: the pass succeeds, so the key is never backed off.
|
||||
pool.replenish_all().await;
|
||||
assert_eq!(pool.warm_len(&key), test_config().target_size);
|
||||
assert_eq!(pool.backoff_failures(&key), 0);
|
||||
assert!(!pool.in_backoff(&key));
|
||||
}
|
||||
}
|
||||
36
litellm-rust/crates/python-bridge/CLAUDE.md
Normal file
36
litellm-rust/crates/python-bridge/CLAUDE.md
Normal file
|
|
@ -0,0 +1,36 @@
|
|||
# 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 adapts Python objects to Rust payloads and returns
|
||||
Python-compatible dictionaries.
|
||||
|
||||
## Bridge Shape
|
||||
|
||||
- Prefer one stable method per top-level LiteLLM route, for example
|
||||
`ocr(payload)`.
|
||||
- Do not add one exported PyO3 function per provider helper unless there is a
|
||||
measured reason.
|
||||
- Provider dispatch belongs in Rust route modules such as
|
||||
`litellm_providers::ocr`, not in this PyO3 crate.
|
||||
- Python owns rollout state and fallback. Rust should return errors; Python
|
||||
decides whether to raise or fall back.
|
||||
|
||||
## 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.
|
||||
16
litellm-rust/crates/python-bridge/Cargo.toml
Normal file
16
litellm-rust/crates/python-bridge/Cargo.toml
Normal file
|
|
@ -0,0 +1,16 @@
|
|||
[package]
|
||||
name = "litellm-python-bridge"
|
||||
version = "0.1.0"
|
||||
edition.workspace = true
|
||||
license.workspace = true
|
||||
repository.workspace = true
|
||||
|
||||
[lib]
|
||||
name = "litellm_python_bridge"
|
||||
crate-type = ["cdylib"]
|
||||
|
||||
[dependencies]
|
||||
litellm-core.workspace = true
|
||||
litellm-providers.workspace = true
|
||||
pyo3 = { workspace = true, features = ["extension-module"] }
|
||||
serde_json.workspace = true
|
||||
32
litellm-rust/crates/python-bridge/src/gil.rs
Normal file
32
litellm-rust/crates/python-bridge/src/gil.rs
Normal file
|
|
@ -0,0 +1,32 @@
|
|||
//! GIL accounting.
|
||||
//!
|
||||
//! A single chokepoint for releasing the GIL around blocking work. Every
|
||||
//! blocking call in the bridge goes through [`release_gil`] instead of calling
|
||||
//! `Python::allow_threads` directly, so the release count stays accurate and we
|
||||
//! have one place to extend later (timing histograms, per-call labels, etc.).
|
||||
|
||||
use std::sync::atomic::{AtomicU64, Ordering};
|
||||
|
||||
use pyo3::prelude::*;
|
||||
|
||||
/// Number of times the bridge has released the GIL since process start.
|
||||
static GIL_RELEASES: AtomicU64 = AtomicU64::new(0);
|
||||
|
||||
/// Release the GIL around `f`, recording the release.
|
||||
///
|
||||
/// `f` must not touch any Python state — that is what makes releasing the GIL
|
||||
/// safe. Returning the value back to Python re-acquires the GIL at the call
|
||||
/// site, after `f` has finished.
|
||||
pub fn release_gil<T, F>(py: Python<'_>, f: F) -> T
|
||||
where
|
||||
F: FnOnce() -> T + Send,
|
||||
T: Send,
|
||||
{
|
||||
GIL_RELEASES.fetch_add(1, Ordering::Relaxed);
|
||||
py.allow_threads(f)
|
||||
}
|
||||
|
||||
/// Total GIL releases performed by the bridge so far.
|
||||
pub fn release_count() -> u64 {
|
||||
GIL_RELEASES.load(Ordering::Relaxed)
|
||||
}
|
||||
100
litellm-rust/crates/python-bridge/src/lib.rs
Normal file
100
litellm-rust/crates/python-bridge/src/lib.rs
Normal file
|
|
@ -0,0 +1,100 @@
|
|||
use std::time::Duration;
|
||||
|
||||
use litellm_core::error::CoreError;
|
||||
use litellm_providers::ocr::run_ocr;
|
||||
use pyo3::exceptions::{PyRuntimeError, PyValueError};
|
||||
use pyo3::prelude::*;
|
||||
use pyo3::types::{PyAny, PyDict};
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
mod gil;
|
||||
|
||||
fn py_to_json(py: Python<'_>, value: &Bound<'_, PyAny>) -> PyResult<Value> {
|
||||
let json = py.import("json")?;
|
||||
let encoded: String = json.call_method1("dumps", (value,))?.extract()?;
|
||||
serde_json::from_str(&encoded).map_err(|err| PyValueError::new_err(err.to_string()))
|
||||
}
|
||||
|
||||
fn json_to_py(py: Python<'_>, value: Value) -> PyResult<Py<PyAny>> {
|
||||
let json = py.import("json")?;
|
||||
let encoded =
|
||||
serde_json::to_string(&value).map_err(|err| PyValueError::new_err(err.to_string()))?;
|
||||
Ok(json.call_method1("loads", (encoded,))?.unbind())
|
||||
}
|
||||
|
||||
/// Map a core error to the closest Python exception. Caller-input problems
|
||||
/// (auth, bad types, missing fields) -> `ValueError`; everything else
|
||||
/// (network, upstream status, parse failures) -> `RuntimeError`.
|
||||
fn core_error_to_pyerr(err: CoreError) -> PyErr {
|
||||
match err {
|
||||
CoreError::Auth(message) => PyValueError::new_err(message),
|
||||
CoreError::InvalidType { .. } | CoreError::MissingField(_) => {
|
||||
PyValueError::new_err(err.to_string())
|
||||
}
|
||||
other => PyRuntimeError::new_err(other.to_string()),
|
||||
}
|
||||
}
|
||||
|
||||
/// Perform a Mistral OCR call end to end and return the response as a dict.
|
||||
#[pyfunction]
|
||||
#[pyo3(signature = (model, document, api_key=None, api_base=None, optional_params=None, timeout_seconds=None))]
|
||||
fn ocr(
|
||||
py: Python<'_>,
|
||||
model: String,
|
||||
document: Py<PyAny>,
|
||||
api_key: Option<String>,
|
||||
api_base: Option<String>,
|
||||
optional_params: Option<Py<PyAny>>,
|
||||
timeout_seconds: Option<f64>,
|
||||
) -> PyResult<Py<PyAny>> {
|
||||
let document = py_to_json(py, document.bind(py))?;
|
||||
|
||||
let optional_params = match optional_params {
|
||||
Some(params) => match py_to_json(py, params.bind(py))? {
|
||||
Value::Object(map) => map,
|
||||
_ => return Err(PyValueError::new_err("optional_params must be a dict")),
|
||||
},
|
||||
None => Map::new(),
|
||||
};
|
||||
|
||||
let timeout = timeout_seconds.and_then(|secs| {
|
||||
if secs.is_finite() && secs > 0.0 {
|
||||
Some(Duration::from_secs_f64(secs))
|
||||
} else {
|
||||
None
|
||||
}
|
||||
});
|
||||
|
||||
// Release the GIL during the blocking HTTP call (counted for observability).
|
||||
let result = gil::release_gil(py, || {
|
||||
run_ocr(
|
||||
&model,
|
||||
document,
|
||||
api_key.as_deref(),
|
||||
api_base.as_deref(),
|
||||
optional_params,
|
||||
timeout,
|
||||
)
|
||||
});
|
||||
|
||||
match result {
|
||||
Ok(value) => json_to_py(py, value),
|
||||
Err(err) => Err(core_error_to_pyerr(err)),
|
||||
}
|
||||
}
|
||||
|
||||
/// Bridge GIL accounting, e.g. `{"releases": 12}`. Lets the Python side observe
|
||||
/// how often the bridge has dropped the GIL for blocking work.
|
||||
#[pyfunction]
|
||||
fn gil_stats(py: Python<'_>) -> PyResult<Py<PyAny>> {
|
||||
let stats = PyDict::new(py);
|
||||
stats.set_item("releases", gil::release_count())?;
|
||||
Ok(stats.into_any().unbind())
|
||||
}
|
||||
|
||||
#[pymodule]
|
||||
fn litellm_python_bridge(module: &Bound<'_, PyModule>) -> PyResult<()> {
|
||||
module.add_function(wrap_pyfunction!(ocr, module)?)?;
|
||||
module.add_function(wrap_pyfunction!(gil_stats, module)?)?;
|
||||
Ok(())
|
||||
}
|
||||
|
|
@ -80,6 +80,7 @@ from litellm.constants import (
|
|||
WANDB_MODELS,
|
||||
REPEATED_STREAMING_CHUNK_LIMIT,
|
||||
request_timeout,
|
||||
request_timeout_explicitly_set as request_timeout_explicitly_set,
|
||||
open_ai_embedding_models,
|
||||
cohere_embedding_models,
|
||||
bedrock_embedding_models,
|
||||
|
|
@ -673,6 +674,7 @@ elevenlabs_models: Set = set()
|
|||
dashscope_models: Set = set()
|
||||
moonshot_models: Set = set()
|
||||
publicai_models: Set = set()
|
||||
darkbloom_models: Set = set()
|
||||
v0_models: Set = set()
|
||||
morph_models: Set = set()
|
||||
lambda_ai_models: Set = set()
|
||||
|
|
@ -927,6 +929,8 @@ def add_known_models(model_cost_map: Optional[Dict] = None):
|
|||
moonshot_models.add(key)
|
||||
elif value.get("litellm_provider") == "publicai":
|
||||
publicai_models.add(key)
|
||||
elif value.get("litellm_provider") == "darkbloom":
|
||||
darkbloom_models.add(key)
|
||||
elif value.get("litellm_provider") == "v0":
|
||||
v0_models.add(key)
|
||||
elif value.get("litellm_provider") == "morph":
|
||||
|
|
@ -1075,6 +1079,7 @@ model_list = list(
|
|||
| dashscope_models
|
||||
| moonshot_models
|
||||
| publicai_models
|
||||
| darkbloom_models
|
||||
| v0_models
|
||||
| morph_models
|
||||
| lambda_ai_models
|
||||
|
|
@ -1179,6 +1184,7 @@ models_by_provider: dict = {
|
|||
"modelscope": modelscope_models,
|
||||
"moonshot": moonshot_models,
|
||||
"publicai": publicai_models,
|
||||
"darkbloom": darkbloom_models,
|
||||
"v0": v0_models,
|
||||
"morph": morph_models,
|
||||
"lambda_ai": lambda_ai_models,
|
||||
|
|
@ -1400,7 +1406,9 @@ from .skills.main import (
|
|||
)
|
||||
from .containers.main import *
|
||||
from .ocr.main import *
|
||||
from .ocr.rust_bridge import use_litellm_rust
|
||||
from .rag.main import *
|
||||
from .sandbox.main import *
|
||||
from .search.main import *
|
||||
from .realtime_api.main import (
|
||||
_arealtime,
|
||||
|
|
@ -1921,9 +1929,6 @@ if TYPE_CHECKING:
|
|||
from .llms.fireworks_ai.completion.transformation import (
|
||||
FireworksAITextCompletionConfig as FireworksAITextCompletionConfig,
|
||||
)
|
||||
from .llms.fireworks_ai.audio_transcription.transformation import (
|
||||
FireworksAIAudioTranscriptionConfig as FireworksAIAudioTranscriptionConfig,
|
||||
)
|
||||
from .llms.fireworks_ai.embed.fireworks_ai_transformation import (
|
||||
FireworksAIEmbeddingConfig as FireworksAIEmbeddingConfig,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -260,7 +260,6 @@ LLM_CONFIG_NAMES = (
|
|||
"SambaNovaEmbeddingConfig",
|
||||
"FireworksAIConfig",
|
||||
"FireworksAITextCompletionConfig",
|
||||
"FireworksAIAudioTranscriptionConfig",
|
||||
"FireworksAIEmbeddingConfig",
|
||||
"FriendliaiChatConfig",
|
||||
"JinaAIEmbeddingConfig",
|
||||
|
|
@ -1027,10 +1026,6 @@ _LLM_CONFIGS_IMPORT_MAP = {
|
|||
".llms.fireworks_ai.completion.transformation",
|
||||
"FireworksAITextCompletionConfig",
|
||||
),
|
||||
"FireworksAIAudioTranscriptionConfig": (
|
||||
".llms.fireworks_ai.audio_transcription.transformation",
|
||||
"FireworksAIAudioTranscriptionConfig",
|
||||
),
|
||||
"FireworksAIEmbeddingConfig": (
|
||||
".llms.fireworks_ai.embed.fireworks_ai_transformation",
|
||||
"FireworksAIEmbeddingConfig",
|
||||
|
|
|
|||
|
|
@ -201,6 +201,18 @@ DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET = int(
|
|||
|
||||
# Provider-specific API base URLs
|
||||
XAI_API_BASE = "https://api.x.ai/v1"
|
||||
OPEN_SANDBOX_API_BASE_ENV_VAR = "OPEN_SANDBOX_API_BASE"
|
||||
OPEN_SANDBOX_API_KEY_ENV_VAR = "OPEN_SANDBOX_API_KEY"
|
||||
OPEN_SANDBOX_DEFAULT_TEMPLATE = "opensandbox/code-interpreter:v1.1.0"
|
||||
_OPEN_SANDBOX_FALLBACK_ENTRYPOINT = "/opt/code-interpreter/code-interpreter.sh"
|
||||
OPEN_SANDBOX_DEFAULT_ENTRYPOINT = (_OPEN_SANDBOX_FALLBACK_ENTRYPOINT,)
|
||||
OPEN_SANDBOX_DEFAULT_LANGUAGE = "python"
|
||||
OPEN_SANDBOX_DEFAULT_CPU_LIMIT = "1"
|
||||
OPEN_SANDBOX_DEFAULT_MEMORY_LIMIT = "2Gi"
|
||||
OPEN_SANDBOX_EXECD_PORT = 44772
|
||||
OPEN_SANDBOX_DEFAULT_TIMEOUT = 300
|
||||
OPEN_SANDBOX_READY_TIMEOUT = 30.0
|
||||
OPEN_SANDBOX_POLL_INTERVAL = 0.2
|
||||
|
||||
DEFAULT_REASONING_EFFORT_LOW_THINKING_BUDGET = int(
|
||||
os.getenv("DEFAULT_REASONING_EFFORT_LOW_THINKING_BUDGET", 1024)
|
||||
|
|
@ -456,6 +468,7 @@ HTTP_HANDLER_CONNECT_TIMEOUT_SECONDS: float = 5.0
|
|||
request_timeout: float = float(
|
||||
os.getenv("REQUEST_TIMEOUT", str(int(DEFAULT_REQUEST_TIMEOUT_SECONDS)))
|
||||
)
|
||||
request_timeout_explicitly_set: bool = "REQUEST_TIMEOUT" in os.environ
|
||||
DEFAULT_A2A_AGENT_TIMEOUT: float = float(
|
||||
os.getenv("DEFAULT_A2A_AGENT_TIMEOUT", 6000)
|
||||
) # 10 minutes
|
||||
|
|
@ -867,6 +880,7 @@ openai_compatible_providers: List = [
|
|||
"docker_model_runner",
|
||||
"ragflow",
|
||||
"pinstripes", # Pinstripes - JSON-configured provider
|
||||
"darkbloom",
|
||||
]
|
||||
openai_text_completion_compatible_providers: List = (
|
||||
[ # providers that support `/v1/completions`
|
||||
|
|
|
|||
|
|
@ -0,0 +1,15 @@
|
|||
"""
|
||||
Code Interpreter Interception Module
|
||||
|
||||
Converts the native OpenAI Responses ``code_interpreter`` tool into a function
|
||||
tool, runs the model-emitted code in a sandbox, and feeds the result back into
|
||||
the agentic loop.
|
||||
"""
|
||||
|
||||
from litellm.integrations.code_interpreter_interception.handler import (
|
||||
CodeInterpreterInterceptionLogger,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"CodeInterpreterInterceptionLogger",
|
||||
]
|
||||
839
litellm/integrations/code_interpreter_interception/handler.py
Normal file
839
litellm/integrations/code_interpreter_interception/handler.py
Normal file
|
|
@ -0,0 +1,839 @@
|
|||
"""
|
||||
Code Interpreter Interception Handler
|
||||
|
||||
CustomLogger that swaps the native OpenAI Responses ``code_interpreter`` tool for
|
||||
a function tool, executes the code the model emits inside a sandbox, and feeds the
|
||||
captured stdout back through the typed agentic loop plan.
|
||||
"""
|
||||
|
||||
import json
|
||||
import time
|
||||
import uuid
|
||||
from typing import Any, Literal, TypedDict, cast
|
||||
|
||||
import litellm
|
||||
from pydantic import ValidationError
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.types.integrations.code_interpreter_interception import (
|
||||
CodeInterpreterInterceptionConfig,
|
||||
)
|
||||
from litellm.types.integrations.custom_logger import (
|
||||
AgenticLoopPlan,
|
||||
AgenticLoopRequestPatch,
|
||||
CHAT_COMPLETION_AGENTIC_SURFACE,
|
||||
NON_CODE_INTERPRETER_INTERCEPTION_INTERNAL_PREFIXES,
|
||||
is_interception_internal_key,
|
||||
)
|
||||
from litellm.types.llms.openai import (
|
||||
ChatCompletionAssistantMessage,
|
||||
ChatCompletionAssistantToolCall,
|
||||
ChatCompletionToolMessage,
|
||||
)
|
||||
from litellm.types.utils import (
|
||||
CallTypes,
|
||||
ChatCompletionMessageToolCall,
|
||||
ModelResponse,
|
||||
)
|
||||
|
||||
LITELLM_CODE_EXECUTION_TOOL_NAME = "litellm_code_execution"
|
||||
_INTERCEPTION_ACTIVE_KEY = "_code_interpreter_interception_active"
|
||||
_SANDBOX_KEY = "_code_interpreter_interception_sandbox_key"
|
||||
_CONVERTED_STREAM_KEY = "_code_interpreter_interception_converted_stream"
|
||||
_LITELLM_METADATA_KEY = "litellm_metadata"
|
||||
_CACHE_TTL_SECONDS = 15 * 60
|
||||
|
||||
|
||||
class CodeExecutionToolCall(TypedDict, total=False):
|
||||
id: str | None
|
||||
call_id: str | None
|
||||
type: Literal["function"]
|
||||
name: str
|
||||
arguments: str
|
||||
|
||||
|
||||
class CodeInterpreterLogOutput(TypedDict):
|
||||
type: Literal["logs"]
|
||||
logs: str
|
||||
|
||||
|
||||
class CodeInterpreterCall(TypedDict):
|
||||
id: str
|
||||
type: Literal["code_interpreter_call"]
|
||||
status: Literal["completed"]
|
||||
code: str
|
||||
container_id: str | None
|
||||
outputs: list[CodeInterpreterLogOutput]
|
||||
|
||||
|
||||
class CodeExecutionFunctionParameters(TypedDict):
|
||||
type: Literal["object"]
|
||||
properties: dict[str, dict[str, str]]
|
||||
required: list[str]
|
||||
|
||||
|
||||
class ResponsesFunctionTool(TypedDict):
|
||||
type: Literal["function"]
|
||||
name: str
|
||||
description: str
|
||||
parameters: CodeExecutionFunctionParameters
|
||||
|
||||
|
||||
class ChatCompletionFunctionDefinition(TypedDict):
|
||||
name: str
|
||||
description: str
|
||||
parameters: CodeExecutionFunctionParameters
|
||||
|
||||
|
||||
class ChatCompletionFunctionTool(TypedDict):
|
||||
type: Literal["function"]
|
||||
function: ChatCompletionFunctionDefinition
|
||||
|
||||
|
||||
CodeExecutionFunctionTool = ResponsesFunctionTool | ChatCompletionFunctionTool
|
||||
|
||||
|
||||
class ResponsesFunctionToolChoice(TypedDict):
|
||||
type: Literal["function"]
|
||||
name: str
|
||||
|
||||
|
||||
class ChatCompletionFunctionToolChoice(TypedDict):
|
||||
type: Literal["function"]
|
||||
function: dict[str, str]
|
||||
|
||||
|
||||
CodeExecutionFunctionToolChoice = (
|
||||
ResponsesFunctionToolChoice | ChatCompletionFunctionToolChoice
|
||||
)
|
||||
|
||||
|
||||
def _resolve_sandbox_tool(sandbox_tool_name: str | None) -> dict[str, Any] | None:
|
||||
try:
|
||||
from litellm.sandbox.sandbox_tools import resolve_sandbox_tool
|
||||
except ImportError:
|
||||
return None
|
||||
return resolve_sandbox_tool(sandbox_tool_name)
|
||||
|
||||
|
||||
class CodeInterpreterInterceptionLogger(CustomLogger):
|
||||
"""
|
||||
CustomLogger that implements transparent code-interpreter execution loops.
|
||||
|
||||
Flow:
|
||||
1. Replace the native ``code_interpreter`` tool with a function tool in the
|
||||
pre-call hook so the model emits code as function-call arguments.
|
||||
2. Detect ``litellm_code_execution`` function calls in the model response.
|
||||
3. Run the emitted code in a sandbox (reused per request via a server-minted
|
||||
sandbox key) and build a typed rerun plan that appends the
|
||||
function_call_output.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
enabled: bool = True,
|
||||
enabled_providers: list[str] | None = None,
|
||||
sandbox_tool_name: str | None = None,
|
||||
sandbox_config: Any | None = None,
|
||||
):
|
||||
super().__init__()
|
||||
self.enabled = enabled
|
||||
self.enabled_providers = enabled_providers
|
||||
self.sandbox_tool_name = sandbox_tool_name
|
||||
self.sandbox_config = sandbox_config
|
||||
self._container_cache: dict[str, tuple[Any, dict[str, Any] | None, float]] = {}
|
||||
|
||||
@classmethod
|
||||
def from_config_yaml(
|
||||
cls, config: CodeInterpreterInterceptionConfig
|
||||
) -> "CodeInterpreterInterceptionLogger":
|
||||
return cls(
|
||||
enabled=bool(config.get("enabled", True)),
|
||||
enabled_providers=config.get("enabled_providers"),
|
||||
sandbox_tool_name=config.get("sandbox_tool_name"),
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def initialize_from_proxy_config(
|
||||
litellm_settings: dict[str, Any],
|
||||
callback_specific_params: dict[str, Any],
|
||||
) -> "CodeInterpreterInterceptionLogger":
|
||||
params: CodeInterpreterInterceptionConfig = {}
|
||||
if "code_interpreter_interception_params" in litellm_settings:
|
||||
params = litellm_settings["code_interpreter_interception_params"]
|
||||
elif "code_interpreter_interception" in callback_specific_params and isinstance(
|
||||
callback_specific_params["code_interpreter_interception"], dict
|
||||
):
|
||||
params = cast(
|
||||
CodeInterpreterInterceptionConfig,
|
||||
callback_specific_params["code_interpreter_interception"],
|
||||
)
|
||||
return CodeInterpreterInterceptionLogger.from_config_yaml(params)
|
||||
|
||||
async def async_pre_call_deployment_hook(
|
||||
self, kwargs: dict[str, Any], call_type: CallTypes | None
|
||||
) -> dict | None:
|
||||
if not kwargs.get("_agentic_loop_depth"):
|
||||
kwargs.pop(_INTERCEPTION_ACTIVE_KEY, None)
|
||||
kwargs.pop(_SANDBOX_KEY, None)
|
||||
self._strip_interception_metadata(kwargs)
|
||||
if not self.enabled:
|
||||
return None
|
||||
if call_type not in (
|
||||
CallTypes.responses,
|
||||
CallTypes.aresponses,
|
||||
CallTypes.completion,
|
||||
CallTypes.acompletion,
|
||||
):
|
||||
return None
|
||||
if (
|
||||
self.enabled_providers is not None
|
||||
and self._resolve_provider(kwargs) not in self.enabled_providers
|
||||
):
|
||||
return None
|
||||
|
||||
tools = kwargs.get("tools")
|
||||
if not isinstance(tools, list):
|
||||
return None
|
||||
if not any(
|
||||
isinstance(tool, dict) and tool.get("type") == "code_interpreter"
|
||||
for tool in tools
|
||||
):
|
||||
return None
|
||||
|
||||
kwargs[_INTERCEPTION_ACTIVE_KEY] = True
|
||||
kwargs[_SANDBOX_KEY] = uuid.uuid4().hex
|
||||
if kwargs.get("stream"):
|
||||
kwargs["stream"] = False
|
||||
kwargs[_CONVERTED_STREAM_KEY] = True
|
||||
self._write_interception_metadata(kwargs)
|
||||
|
||||
function_tool = self._get_function_tool(call_type=call_type)
|
||||
kwargs["tools"] = [
|
||||
(
|
||||
function_tool
|
||||
if isinstance(tool, dict) and tool.get("type") == "code_interpreter"
|
||||
else tool
|
||||
)
|
||||
for tool in tools
|
||||
]
|
||||
if self._tool_choice_targets_code_interpreter(kwargs.get("tool_choice")):
|
||||
kwargs["tool_choice"] = self._get_function_tool_choice(call_type=call_type)
|
||||
return kwargs
|
||||
|
||||
@staticmethod
|
||||
def _strip_interception_metadata(kwargs: dict[str, Any]) -> None:
|
||||
metadata = kwargs.get(_LITELLM_METADATA_KEY)
|
||||
if not isinstance(metadata, dict):
|
||||
return
|
||||
filtered_metadata = {
|
||||
key: value
|
||||
for key, value in metadata.items()
|
||||
if not is_interception_internal_key(key)
|
||||
and not key.startswith("_agentic_loop")
|
||||
and key != "max_agentic_loops"
|
||||
}
|
||||
if filtered_metadata:
|
||||
kwargs[_LITELLM_METADATA_KEY] = filtered_metadata
|
||||
else:
|
||||
kwargs.pop(_LITELLM_METADATA_KEY, None)
|
||||
|
||||
@staticmethod
|
||||
def _write_interception_metadata(kwargs: dict[str, Any]) -> None:
|
||||
metadata = kwargs.get(_LITELLM_METADATA_KEY)
|
||||
metadata = dict(metadata) if isinstance(metadata, dict) else {}
|
||||
for key in (_INTERCEPTION_ACTIVE_KEY, _SANDBOX_KEY, _CONVERTED_STREAM_KEY):
|
||||
if key in kwargs:
|
||||
metadata[key] = kwargs[key]
|
||||
kwargs[_LITELLM_METADATA_KEY] = metadata
|
||||
|
||||
@staticmethod
|
||||
def _get_function_parameters() -> CodeExecutionFunctionParameters:
|
||||
return {
|
||||
"type": "object",
|
||||
"properties": {"code": {"type": "string"}},
|
||||
"required": ["code"],
|
||||
}
|
||||
|
||||
def _get_function_tool(
|
||||
self, call_type: CallTypes | None
|
||||
) -> CodeExecutionFunctionTool:
|
||||
description = "Execute python code in a sandbox and return stdout."
|
||||
if call_type in (CallTypes.completion, CallTypes.acompletion):
|
||||
return {
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": LITELLM_CODE_EXECUTION_TOOL_NAME,
|
||||
"description": description,
|
||||
"parameters": self._get_function_parameters(),
|
||||
},
|
||||
}
|
||||
return {
|
||||
"type": "function",
|
||||
"name": LITELLM_CODE_EXECUTION_TOOL_NAME,
|
||||
"description": description,
|
||||
"parameters": self._get_function_parameters(),
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def _get_function_tool_choice(
|
||||
call_type: CallTypes | None,
|
||||
) -> CodeExecutionFunctionToolChoice:
|
||||
if call_type in (CallTypes.completion, CallTypes.acompletion):
|
||||
return {
|
||||
"type": "function",
|
||||
"function": {"name": LITELLM_CODE_EXECUTION_TOOL_NAME},
|
||||
}
|
||||
return {
|
||||
"type": "function",
|
||||
"name": LITELLM_CODE_EXECUTION_TOOL_NAME,
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def _tool_choice_targets_code_interpreter(tool_choice: Any) -> bool:
|
||||
if not isinstance(tool_choice, dict):
|
||||
return False
|
||||
function = tool_choice.get("function")
|
||||
return (
|
||||
tool_choice.get("type") == "code_interpreter"
|
||||
or tool_choice.get("name") == "code_interpreter"
|
||||
or tool_choice.get("name") == LITELLM_CODE_EXECUTION_TOOL_NAME
|
||||
or (
|
||||
isinstance(function, dict)
|
||||
and function.get("name") == LITELLM_CODE_EXECUTION_TOOL_NAME
|
||||
)
|
||||
)
|
||||
|
||||
def _resolve_provider(self, kwargs: dict[str, Any]) -> str | None:
|
||||
provider = kwargs.get("custom_llm_provider")
|
||||
if provider:
|
||||
return provider
|
||||
model = kwargs.get("model")
|
||||
if not isinstance(model, str):
|
||||
return None
|
||||
try:
|
||||
return litellm.get_llm_provider(model=model)[1]
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
async def async_should_run_agentic_loop(
|
||||
self,
|
||||
response: Any,
|
||||
model: str,
|
||||
messages: list[dict],
|
||||
tools: list[dict] | None,
|
||||
stream: bool,
|
||||
custom_llm_provider: str,
|
||||
kwargs: dict,
|
||||
) -> tuple[bool, dict]:
|
||||
if not self.enabled:
|
||||
return False, {}
|
||||
if not kwargs.get(_INTERCEPTION_ACTIVE_KEY):
|
||||
return False, {}
|
||||
if (
|
||||
self.enabled_providers is not None
|
||||
and custom_llm_provider not in self.enabled_providers
|
||||
):
|
||||
return False, {}
|
||||
|
||||
tool_calls = (
|
||||
self._extract_chat_completion_code_execution_tool_calls(response=response)
|
||||
if kwargs.get("_agentic_loop_api_surface")
|
||||
== CHAT_COMPLETION_AGENTIC_SURFACE
|
||||
else self._extract_code_execution_tool_calls(response=response)
|
||||
)
|
||||
if not tool_calls:
|
||||
return False, {}
|
||||
|
||||
return True, {"tool_calls": tool_calls}
|
||||
|
||||
async def async_build_agentic_loop_plan(
|
||||
self,
|
||||
tools: dict,
|
||||
model: str,
|
||||
messages: list[dict],
|
||||
response: Any,
|
||||
anthropic_messages_provider_config: Any,
|
||||
anthropic_messages_optional_request_params: dict,
|
||||
logging_obj: Any,
|
||||
stream: bool,
|
||||
kwargs: dict,
|
||||
) -> AgenticLoopPlan:
|
||||
if kwargs.get("_agentic_loop_api_surface") == CHAT_COMPLETION_AGENTIC_SURFACE:
|
||||
return await self._build_chat_completion_agentic_loop_plan(
|
||||
tools=tools,
|
||||
model=model,
|
||||
messages=messages,
|
||||
optional_params=anthropic_messages_optional_request_params,
|
||||
kwargs=kwargs,
|
||||
)
|
||||
|
||||
await self._prune_expired_cache()
|
||||
tool_calls = cast(list[CodeExecutionToolCall], tools.get("tool_calls", []))
|
||||
sandbox_key = kwargs.get(_SANDBOX_KEY)
|
||||
container, params = await self._get_or_create_container(cache_key=sandbox_key)
|
||||
|
||||
try:
|
||||
container_id = cast(str | None, getattr(container, "id", None))
|
||||
input_list = self._normalize_messages(messages)
|
||||
code_interpreter_calls: list[CodeInterpreterCall] = []
|
||||
for tool_call in tool_calls:
|
||||
arguments = tool_call.get("arguments", "")
|
||||
code = self._parse_code(arguments)
|
||||
stdout = await self._run_tool_call(
|
||||
container=container, params=params, arguments=arguments
|
||||
)
|
||||
input_list.append(
|
||||
{
|
||||
"type": "function_call",
|
||||
"call_id": tool_call.get("call_id"),
|
||||
"name": LITELLM_CODE_EXECUTION_TOOL_NAME,
|
||||
"arguments": arguments,
|
||||
}
|
||||
)
|
||||
input_list.append(
|
||||
{
|
||||
"type": "function_call_output",
|
||||
"call_id": tool_call.get("call_id"),
|
||||
"output": stdout,
|
||||
}
|
||||
)
|
||||
code_interpreter_calls.append(
|
||||
{
|
||||
"id": f"ci_{uuid.uuid4().hex}",
|
||||
"type": "code_interpreter_call",
|
||||
"status": "completed",
|
||||
"code": code,
|
||||
"container_id": container_id,
|
||||
"outputs": (
|
||||
[{"type": "logs", "logs": stdout}] if stdout else []
|
||||
),
|
||||
}
|
||||
)
|
||||
except Exception:
|
||||
await self._delete_container_for_cache_key(sandbox_key)
|
||||
raise
|
||||
|
||||
optional_params = anthropic_messages_optional_request_params
|
||||
request_patch = AgenticLoopRequestPatch(
|
||||
model=model,
|
||||
messages=input_list,
|
||||
tools=self._get_followup_tools(
|
||||
tools=optional_params.get("tools"),
|
||||
call_type=CallTypes.responses,
|
||||
),
|
||||
optional_params=self._get_followup_optional_params(optional_params),
|
||||
kwargs=self._filter_agentic_loop_kwargs(kwargs),
|
||||
)
|
||||
|
||||
return AgenticLoopPlan(
|
||||
run_agentic_loop=True,
|
||||
request_patch=request_patch,
|
||||
metadata={
|
||||
"tool_type": "code_interpreter",
|
||||
"sandbox_key": sandbox_key or "",
|
||||
"code_interpreter_calls": code_interpreter_calls,
|
||||
},
|
||||
)
|
||||
|
||||
async def _build_chat_completion_agentic_loop_plan(
|
||||
self,
|
||||
tools: dict[str, object],
|
||||
model: str,
|
||||
messages: list[dict],
|
||||
optional_params: dict[str, object],
|
||||
kwargs: dict[str, object],
|
||||
) -> AgenticLoopPlan:
|
||||
await self._prune_expired_cache()
|
||||
tool_calls = cast(list[CodeExecutionToolCall], tools.get("tool_calls", []))
|
||||
sandbox_key = cast(str | None, kwargs.get(_SANDBOX_KEY))
|
||||
container, params = await self._get_or_create_container(cache_key=sandbox_key)
|
||||
|
||||
try:
|
||||
container_id = cast(str | None, getattr(container, "id", None))
|
||||
tool_results = [
|
||||
await self._build_chat_completion_tool_result(
|
||||
container=container,
|
||||
params=params,
|
||||
tool_call=tool_call,
|
||||
container_id=container_id,
|
||||
)
|
||||
for tool_call in tool_calls
|
||||
]
|
||||
except Exception:
|
||||
await self._delete_container_for_cache_key(sandbox_key)
|
||||
raise
|
||||
tool_messages = [result[0] for result in tool_results]
|
||||
code_interpreter_calls = [result[1] for result in tool_results]
|
||||
|
||||
request_patch = AgenticLoopRequestPatch(
|
||||
model=model,
|
||||
messages=list(messages)
|
||||
+ [self._build_chat_completion_assistant_message(tool_calls)]
|
||||
+ tool_messages,
|
||||
tools=self._get_followup_tools(
|
||||
tools=optional_params.get("tools"),
|
||||
call_type=CallTypes.completion,
|
||||
),
|
||||
optional_params=self._get_followup_optional_params(optional_params),
|
||||
kwargs=self._filter_agentic_loop_kwargs(kwargs),
|
||||
)
|
||||
|
||||
return AgenticLoopPlan(
|
||||
run_agentic_loop=True,
|
||||
request_patch=request_patch,
|
||||
metadata={
|
||||
"tool_type": "code_interpreter",
|
||||
"sandbox_key": sandbox_key or "",
|
||||
"code_interpreter_calls": code_interpreter_calls,
|
||||
"response_format": "openai",
|
||||
},
|
||||
)
|
||||
|
||||
async def _build_chat_completion_tool_result(
|
||||
self,
|
||||
container: object,
|
||||
params: dict[str, Any] | None,
|
||||
tool_call: CodeExecutionToolCall,
|
||||
container_id: str | None,
|
||||
) -> tuple[ChatCompletionToolMessage, CodeInterpreterCall]:
|
||||
arguments = tool_call.get("arguments", "")
|
||||
code = self._parse_code(arguments)
|
||||
stdout = await self._run_tool_call(
|
||||
container=container, params=params, arguments=arguments
|
||||
)
|
||||
tool_call_id = (
|
||||
tool_call.get("id") or tool_call.get("call_id") or uuid.uuid4().hex
|
||||
)
|
||||
return (
|
||||
{
|
||||
"role": "tool",
|
||||
"tool_call_id": tool_call_id,
|
||||
"content": stdout,
|
||||
},
|
||||
{
|
||||
"id": f"ci_{uuid.uuid4().hex}",
|
||||
"type": "code_interpreter_call",
|
||||
"status": "completed",
|
||||
"code": code,
|
||||
"container_id": container_id,
|
||||
"outputs": [{"type": "logs", "logs": stdout}] if stdout else [],
|
||||
},
|
||||
)
|
||||
|
||||
async def async_agentic_loop_cleanup_hook(
|
||||
self, plan: AgenticLoopPlan, kwargs: dict
|
||||
) -> None:
|
||||
metadata = plan.metadata or {} if plan else {}
|
||||
await self._delete_container_for_cache_key(metadata.get("sandbox_key"))
|
||||
|
||||
@staticmethod
|
||||
def _filter_agentic_loop_kwargs(kwargs: dict[str, object]) -> dict[str, object]:
|
||||
return {
|
||||
k: v
|
||||
for k, v in kwargs.items()
|
||||
if k not in {"litellm_logging_obj", "acompletion"}
|
||||
and not is_interception_internal_key(
|
||||
k, prefixes=NON_CODE_INTERPRETER_INTERCEPTION_INTERNAL_PREFIXES
|
||||
)
|
||||
}
|
||||
|
||||
def _get_followup_tools(
|
||||
self, tools: object, call_type: CallTypes | None
|
||||
) -> list[dict[str, Any]] | None:
|
||||
if not isinstance(tools, list):
|
||||
return None
|
||||
return [
|
||||
(
|
||||
self._get_function_tool(call_type=call_type)
|
||||
if isinstance(tool, dict) and tool.get("type") == "code_interpreter"
|
||||
else tool
|
||||
)
|
||||
for tool in tools
|
||||
]
|
||||
|
||||
def _get_followup_optional_params(
|
||||
self, optional_params: dict[str, object]
|
||||
) -> dict[str, object]:
|
||||
drop_tool_choice = self._tool_choice_targets_code_interpreter(
|
||||
optional_params.get("tool_choice")
|
||||
)
|
||||
return {
|
||||
k: v
|
||||
for k, v in optional_params.items()
|
||||
if k != "tools" and not (k == "tool_choice" and drop_tool_choice)
|
||||
}
|
||||
|
||||
async def async_post_agentic_loop_response_hook(
|
||||
self, response: Any, plan: AgenticLoopPlan, kwargs: dict
|
||||
) -> Any:
|
||||
metadata = plan.metadata or {} if plan else {}
|
||||
await self._delete_container_for_cache_key(metadata.get("sandbox_key"))
|
||||
|
||||
calls = metadata.get("code_interpreter_calls")
|
||||
if not calls:
|
||||
return response
|
||||
|
||||
is_dict = isinstance(response, dict)
|
||||
output = (
|
||||
response.get("output") if is_dict else getattr(response, "output", None)
|
||||
)
|
||||
if not isinstance(output, list):
|
||||
return response
|
||||
|
||||
def _item_type(item: Any) -> Any:
|
||||
return (
|
||||
item.get("type")
|
||||
if isinstance(item, dict)
|
||||
else getattr(item, "type", None)
|
||||
)
|
||||
|
||||
insert_at = next(
|
||||
(i for i, item in enumerate(output) if _item_type(item) == "message"),
|
||||
len(output),
|
||||
)
|
||||
new_output = output[:insert_at] + list(calls) + output[insert_at:]
|
||||
if is_dict:
|
||||
response["output"] = new_output
|
||||
else:
|
||||
response.output = new_output
|
||||
return response
|
||||
|
||||
@staticmethod
|
||||
def _parse_code(arguments: str) -> str:
|
||||
try:
|
||||
return json.loads(arguments).get("code", "") if arguments else ""
|
||||
except (json.JSONDecodeError, TypeError, AttributeError):
|
||||
return ""
|
||||
|
||||
async def _run_tool_call(
|
||||
self, container: Any, params: dict[str, Any] | None, arguments: str
|
||||
) -> str:
|
||||
try:
|
||||
code = json.loads(arguments).get("code", "") if arguments else ""
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
return "[invalid tool arguments: could not parse code]"
|
||||
|
||||
result = await self._run_code(container=container, params=params, code=code)
|
||||
if getattr(result, "error", None):
|
||||
error = result.error
|
||||
message = (
|
||||
error.get("value") or error.get("name")
|
||||
if isinstance(error, dict)
|
||||
else str(error)
|
||||
)
|
||||
return f"[execution error] {message}"
|
||||
return getattr(result, "stdout", "") or ""
|
||||
|
||||
async def _get_or_create_container(
|
||||
self, cache_key: str | None
|
||||
) -> tuple[Any, dict[str, Any] | None]:
|
||||
if cache_key:
|
||||
cached = self._container_cache.get(cache_key)
|
||||
if cached is not None:
|
||||
return cached[0], cached[1]
|
||||
|
||||
container, params = await self._create_container()
|
||||
if cache_key:
|
||||
self._container_cache[cache_key] = (container, params, time.time())
|
||||
return container, params
|
||||
|
||||
async def _create_container(self) -> tuple[Any, dict[str, Any] | None]:
|
||||
if self.sandbox_config is not None:
|
||||
return await self.sandbox_config.acreate_sandbox(), None
|
||||
|
||||
params = _resolve_sandbox_tool(self.sandbox_tool_name)
|
||||
if params is None:
|
||||
raise ValueError(
|
||||
"CodeInterpreterInterception: no sandbox available. Provide a "
|
||||
"sandbox_config or configure a sandbox tool resolvable via "
|
||||
"sandbox_tool_name."
|
||||
)
|
||||
container = await litellm.acreate_sandbox(
|
||||
provider=params["sandbox_provider"],
|
||||
api_key=params.get("api_key"),
|
||||
api_base=params.get("api_base"),
|
||||
)
|
||||
return container, params
|
||||
|
||||
async def _run_code(
|
||||
self, container: Any, params: dict[str, Any] | None, code: str
|
||||
) -> Any:
|
||||
if self.sandbox_config is not None:
|
||||
return await self.sandbox_config.arun_code(container=container, code=code)
|
||||
if params is None:
|
||||
raise ValueError(
|
||||
"CodeInterpreterInterception: no sandbox available to run code."
|
||||
)
|
||||
return await litellm.arun_code(
|
||||
provider=params["sandbox_provider"],
|
||||
container=container,
|
||||
code=code,
|
||||
api_key=params.get("api_key"),
|
||||
)
|
||||
|
||||
async def _delete_container(
|
||||
self, container: Any, params: dict[str, Any] | None
|
||||
) -> None:
|
||||
try:
|
||||
if self.sandbox_config is not None:
|
||||
await self.sandbox_config.adelete_sandbox(container=container)
|
||||
return
|
||||
if params is None:
|
||||
return
|
||||
await litellm.adelete_sandbox(
|
||||
provider=params["sandbox_provider"],
|
||||
container=container,
|
||||
api_key=params.get("api_key"),
|
||||
api_base=params.get("api_base"),
|
||||
)
|
||||
except Exception:
|
||||
verbose_logger.exception(
|
||||
"CodeInterpreterInterception: failed to delete sandbox container"
|
||||
)
|
||||
|
||||
async def _delete_container_for_cache_key(self, cache_key: str | None) -> None:
|
||||
if not cache_key:
|
||||
return
|
||||
cached = self._container_cache.pop(cache_key, None)
|
||||
if cached is None:
|
||||
return
|
||||
await self._delete_container(container=cached[0], params=cached[1])
|
||||
|
||||
def _normalize_messages(self, messages: Any) -> list[dict[str, Any]]:
|
||||
if isinstance(messages, str):
|
||||
return [{"role": "user", "content": messages}]
|
||||
if isinstance(messages, list):
|
||||
return list(messages)
|
||||
return []
|
||||
|
||||
def _extract_code_execution_tool_calls(
|
||||
self, response: object
|
||||
) -> list[CodeExecutionToolCall]:
|
||||
if isinstance(response, dict):
|
||||
output = response.get("output", [])
|
||||
else:
|
||||
output = getattr(response, "output", []) or []
|
||||
if not isinstance(output, list):
|
||||
return []
|
||||
|
||||
return [
|
||||
{
|
||||
"call_id": (
|
||||
item.get("call_id")
|
||||
if isinstance(item, dict)
|
||||
else getattr(item, "call_id", None)
|
||||
),
|
||||
"name": LITELLM_CODE_EXECUTION_TOOL_NAME,
|
||||
"arguments": (
|
||||
item.get("arguments")
|
||||
if isinstance(item, dict)
|
||||
else getattr(item, "arguments", "")
|
||||
),
|
||||
}
|
||||
for item in output
|
||||
if self._is_code_execution_call(item)
|
||||
]
|
||||
|
||||
def _extract_chat_completion_code_execution_tool_calls(
|
||||
self, response: ModelResponse | dict[str, Any]
|
||||
) -> list[CodeExecutionToolCall]:
|
||||
model_response = self._to_model_response(response)
|
||||
if model_response is None:
|
||||
return []
|
||||
choices = model_response.choices or []
|
||||
if not choices:
|
||||
return []
|
||||
message = choices[0].message
|
||||
tool_calls = message.tool_calls or []
|
||||
|
||||
return [
|
||||
normalized
|
||||
for tool_call in tool_calls
|
||||
if (normalized := self._normalize_chat_completion_tool_call(tool_call))
|
||||
is not None
|
||||
]
|
||||
|
||||
@staticmethod
|
||||
def _normalize_chat_completion_tool_call(
|
||||
tool_call: ChatCompletionMessageToolCall,
|
||||
) -> CodeExecutionToolCall | None:
|
||||
if (
|
||||
tool_call.type != "function"
|
||||
or tool_call.function.name != LITELLM_CODE_EXECUTION_TOOL_NAME
|
||||
):
|
||||
return None
|
||||
|
||||
arguments = tool_call.function.arguments
|
||||
if isinstance(arguments, dict):
|
||||
arguments = json.dumps(arguments)
|
||||
elif not isinstance(arguments, str):
|
||||
arguments = "" if arguments is None else str(arguments)
|
||||
|
||||
return {
|
||||
"id": tool_call.id,
|
||||
"call_id": tool_call.id,
|
||||
"type": "function",
|
||||
"name": LITELLM_CODE_EXECUTION_TOOL_NAME,
|
||||
"arguments": arguments,
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def _build_chat_completion_assistant_message(
|
||||
tool_calls: list[CodeExecutionToolCall],
|
||||
) -> ChatCompletionAssistantMessage:
|
||||
return {
|
||||
"role": "assistant",
|
||||
"tool_calls": [
|
||||
cast(
|
||||
ChatCompletionAssistantToolCall,
|
||||
{
|
||||
"id": tool_call.get("id"),
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": LITELLM_CODE_EXECUTION_TOOL_NAME,
|
||||
"arguments": tool_call.get("arguments", ""),
|
||||
},
|
||||
},
|
||||
)
|
||||
for tool_call in tool_calls
|
||||
],
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def _to_model_response(
|
||||
response: ModelResponse | dict[str, Any],
|
||||
) -> ModelResponse | None:
|
||||
if isinstance(response, ModelResponse):
|
||||
return response
|
||||
try:
|
||||
return ModelResponse(**response)
|
||||
except (TypeError, ValidationError):
|
||||
return None
|
||||
|
||||
def _is_code_execution_call(self, item: Any) -> bool:
|
||||
if isinstance(item, dict):
|
||||
return (
|
||||
item.get("type") == "function_call"
|
||||
and item.get("name") == LITELLM_CODE_EXECUTION_TOOL_NAME
|
||||
)
|
||||
return (
|
||||
getattr(item, "type", None) == "function_call"
|
||||
and getattr(item, "name", None) == LITELLM_CODE_EXECUTION_TOOL_NAME
|
||||
)
|
||||
|
||||
async def _prune_expired_cache(self) -> None:
|
||||
now = time.time()
|
||||
expired = [
|
||||
(cache_key, container, params)
|
||||
for cache_key, (
|
||||
container,
|
||||
params,
|
||||
created_at,
|
||||
) in self._container_cache.items()
|
||||
if now - created_at > _CACHE_TTL_SECONDS
|
||||
]
|
||||
for cache_key, container, params in expired:
|
||||
self._container_cache.pop(cache_key, None)
|
||||
await self._delete_container(container=container, params=params)
|
||||
|
|
@ -718,6 +718,24 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
|
|||
"""
|
||||
return response
|
||||
|
||||
async def async_agentic_loop_cleanup_hook(
|
||||
self,
|
||||
plan: AgenticLoopPlan,
|
||||
kwargs: dict,
|
||||
) -> None:
|
||||
"""
|
||||
Release resources held for an agentic-loop iteration.
|
||||
|
||||
Runs in a ``finally`` around the follow-up provider call, so it fires
|
||||
whether the rerun returns normally, hits a loop safety abort, or raises
|
||||
an upstream error. Implementations must be idempotent because the
|
||||
post-response hook may already have released the same resource on the
|
||||
success path. Use ``plan.metadata`` to locate what to clean up.
|
||||
|
||||
Default does nothing.
|
||||
"""
|
||||
return None
|
||||
|
||||
async def async_should_run_chat_completion_agentic_loop(
|
||||
self,
|
||||
response: Any,
|
||||
|
|
|
|||
|
|
@ -1,6 +1,7 @@
|
|||
"""Typed configuration for the OpenTelemetry instrumentation."""
|
||||
|
||||
from enum import Enum
|
||||
from functools import lru_cache
|
||||
from typing import Any, List
|
||||
|
||||
from pydantic import AliasChoices, BaseModel, Field, field_validator, model_validator
|
||||
|
|
@ -47,7 +48,12 @@ class _OTelV2Flag(BaseSettings):
|
|||
enabled: bool = Field(default=False, validation_alias=AliasChoices(OTEL_V2_ENV))
|
||||
|
||||
|
||||
@lru_cache(maxsize=1)
|
||||
def is_otel_v2_enabled() -> bool:
|
||||
# Resolved once at startup and cached: constructing the pydantic-settings
|
||||
# model re-scans the environment and cost ~28us, which on the proxy hot path
|
||||
# (auth, logging-callback setup) compounded into a measurable throughput
|
||||
# regression. Tests that toggle the env must call ``is_otel_v2_enabled.cache_clear()``.
|
||||
return _OTelV2Flag().enabled
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -300,9 +300,6 @@ class LiteLLMResponsesInteractionsConfig:
|
|||
"total_output_tokens": getattr(usage, "output_tokens", 0),
|
||||
}
|
||||
|
||||
# Add role
|
||||
interactions_response_dict["role"] = "model"
|
||||
|
||||
# Add updated (same as created for now)
|
||||
interactions_response_dict["updated"] = created
|
||||
|
||||
|
|
|
|||
332
litellm/litellm_core_utils/chat_completion_agentic_loop.py
Normal file
332
litellm/litellm_core_utils/chat_completion_agentic_loop.py
Normal file
|
|
@ -0,0 +1,332 @@
|
|||
# this is a patch to allow for agentic loops covering llm_http_handler.py and openai sdk based calling flows for the .completion() api
|
||||
|
||||
import json
|
||||
from typing import cast
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.types.integrations.custom_logger import (
|
||||
CHAT_COMPLETION_AGENTIC_SURFACE,
|
||||
NON_CODE_INTERPRETER_INTERCEPTION_INTERNAL_PREFIXES,
|
||||
AgenticLoopPlan,
|
||||
AgenticLoopRequestPatch,
|
||||
is_interception_internal_key,
|
||||
)
|
||||
from litellm.types.utils import ModelResponse
|
||||
from litellm.utils import CustomStreamWrapper
|
||||
|
||||
_FOLLOWUP_INTERNAL_PARAMS = frozenset(
|
||||
(
|
||||
"acompletion",
|
||||
"litellm_logging_obj",
|
||||
"custom_llm_provider",
|
||||
"model_alias_map",
|
||||
"stream_response",
|
||||
"custom_prompt_dict",
|
||||
"_agentic_loop_api_surface",
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def _gate_overridden(callback: CustomLogger) -> bool:
|
||||
base = CustomLogger.async_should_run_agentic_loop
|
||||
func = type(callback).async_should_run_agentic_loop
|
||||
return getattr(func, "__func__", func) is not getattr(base, "__func__", base)
|
||||
|
||||
|
||||
def _build_plan_overridden(callback: CustomLogger) -> bool:
|
||||
base = CustomLogger.async_build_agentic_loop_plan
|
||||
func = type(callback).async_build_agentic_loop_plan
|
||||
return getattr(func, "__func__", func) is not getattr(base, "__func__", base)
|
||||
|
||||
|
||||
def _post_hook_overridden(callback: CustomLogger) -> bool:
|
||||
base = CustomLogger.async_post_agentic_loop_response_hook
|
||||
func = type(callback).async_post_agentic_loop_response_hook
|
||||
return getattr(func, "__func__", func) is not getattr(base, "__func__", base)
|
||||
|
||||
|
||||
def _coerce_int(value: object, default: int) -> int:
|
||||
return int(value) if isinstance(value, (int, str)) else default
|
||||
|
||||
|
||||
def _agentic_loop_settings(kwargs: dict[str, object]) -> tuple[int, int, list[str]]:
|
||||
depth = _coerce_int(kwargs.get("_agentic_loop_depth"), 0)
|
||||
max_loops = max(_coerce_int(kwargs.get("max_agentic_loops"), 3), 1)
|
||||
raw_fingerprints = kwargs.get("_agentic_loop_fingerprints")
|
||||
fingerprints = (
|
||||
[str(fp) for fp in raw_fingerprints]
|
||||
if isinstance(raw_fingerprints, list)
|
||||
else []
|
||||
)
|
||||
return depth, max_loops, fingerprints
|
||||
|
||||
|
||||
def _fingerprint_tools(tool_calls: object) -> str:
|
||||
try:
|
||||
return json.dumps(tool_calls, sort_keys=True, default=str)
|
||||
except Exception:
|
||||
return str(tool_calls)
|
||||
|
||||
|
||||
def _check_agentic_loop_safety(
|
||||
tool_calls: object,
|
||||
fingerprints: list[str],
|
||||
depth: int,
|
||||
max_loops: int,
|
||||
model: str,
|
||||
) -> str:
|
||||
fingerprint = _fingerprint_tools(tool_calls)
|
||||
if fingerprint in fingerprints:
|
||||
raise ValueError(
|
||||
"Agentic loop detected repeated tool-call fingerprint; aborting rerun"
|
||||
)
|
||||
if depth >= max_loops:
|
||||
raise ValueError(f"Exceeded max_agentic_loops={max_loops} for model={model}")
|
||||
return fingerprint
|
||||
|
||||
|
||||
def _wrap_response_as_fake_stream(response: object) -> object:
|
||||
if getattr(response, "object", None) == "chat.completion.chunk":
|
||||
return response
|
||||
if not hasattr(response, "choices"):
|
||||
return response
|
||||
from litellm.llms.base_llm.base_model_iterator import (
|
||||
convert_model_response_to_streaming,
|
||||
)
|
||||
|
||||
return convert_model_response_to_streaming(cast(ModelResponse, response))
|
||||
|
||||
|
||||
def _add_agentic_loop_metadata(kwargs_for_followup: dict[str, object]) -> None:
|
||||
metadata = kwargs_for_followup.get("litellm_metadata")
|
||||
metadata = dict(metadata) if isinstance(metadata, dict) else {}
|
||||
for key, value in kwargs_for_followup.items():
|
||||
if (
|
||||
key.startswith("_agentic_loop")
|
||||
or key == "max_agentic_loops"
|
||||
or is_interception_internal_key(key)
|
||||
):
|
||||
metadata[key] = value
|
||||
kwargs_for_followup["litellm_metadata"] = metadata
|
||||
|
||||
|
||||
def _filter_followup_kwargs(source: dict[str, object]) -> dict[str, object]:
|
||||
return {
|
||||
k: v
|
||||
for k, v in source.items()
|
||||
if not is_interception_internal_key(
|
||||
k, prefixes=NON_CODE_INTERPRETER_INTERCEPTION_INTERNAL_PREFIXES
|
||||
)
|
||||
and k not in _FOLLOWUP_INTERNAL_PARAMS
|
||||
}
|
||||
|
||||
|
||||
async def _execute_chat_completion_agentic_plan(
|
||||
*,
|
||||
plan: AgenticLoopPlan,
|
||||
callback: CustomLogger,
|
||||
model: str,
|
||||
optional_params: dict[str, object],
|
||||
kwargs: dict[str, object],
|
||||
logging_obj: object,
|
||||
custom_llm_provider: str,
|
||||
depth: int,
|
||||
max_loops: int,
|
||||
fingerprints: list[str],
|
||||
fingerprint: str,
|
||||
) -> object:
|
||||
import litellm
|
||||
|
||||
patch = plan.request_patch or AgenticLoopRequestPatch()
|
||||
if patch.messages is None:
|
||||
raise ValueError("Agentic loop plan missing patched messages")
|
||||
|
||||
full_model_name = patch.model or model
|
||||
if "/" not in full_model_name:
|
||||
full_model_name = f"{custom_llm_provider}/{full_model_name}"
|
||||
|
||||
optional_params_for_followup = {**optional_params, **patch.optional_params}
|
||||
if patch.tools is not None:
|
||||
optional_params_for_followup["tools"] = patch.tools
|
||||
if "tool_choice" not in patch.optional_params:
|
||||
optional_params_for_followup.pop("tool_choice", None)
|
||||
|
||||
kwargs_for_followup = _filter_followup_kwargs(kwargs)
|
||||
kwargs_for_followup.update(
|
||||
{
|
||||
k: v
|
||||
for k, v in _filter_followup_kwargs(patch.kwargs).items()
|
||||
if k not in optional_params_for_followup
|
||||
}
|
||||
)
|
||||
kwargs_for_followup["_agentic_loop_depth"] = depth + 1
|
||||
kwargs_for_followup["max_agentic_loops"] = max_loops
|
||||
kwargs_for_followup["_agentic_loop_fingerprints"] = fingerprints + [fingerprint]
|
||||
_add_agentic_loop_metadata(kwargs_for_followup)
|
||||
|
||||
try:
|
||||
response_followup = await litellm.acompletion(
|
||||
model=full_model_name,
|
||||
messages=patch.messages,
|
||||
**optional_params_for_followup,
|
||||
**kwargs_for_followup,
|
||||
)
|
||||
if _post_hook_overridden(callback):
|
||||
try:
|
||||
response_followup = (
|
||||
await callback.async_post_agentic_loop_response_hook(
|
||||
response=response_followup, plan=plan, kwargs=kwargs
|
||||
)
|
||||
)
|
||||
except Exception as e:
|
||||
_call_id = getattr(logging_obj, "litellm_call_id", "unknown")
|
||||
verbose_logger.exception(
|
||||
"LiteLLM.AgenticHookError: Exception in "
|
||||
"async_post_agentic_loop_response_hook [call_id=%s model=%s]: %s",
|
||||
_call_id,
|
||||
model,
|
||||
str(e),
|
||||
)
|
||||
if kwargs.get("_code_interpreter_interception_converted_stream") and not depth:
|
||||
return _wrap_response_as_fake_stream(response_followup)
|
||||
return response_followup
|
||||
finally:
|
||||
try:
|
||||
await callback.async_agentic_loop_cleanup_hook(plan=plan, kwargs=kwargs)
|
||||
except Exception as e:
|
||||
_call_id = getattr(logging_obj, "litellm_call_id", "unknown")
|
||||
verbose_logger.exception(
|
||||
"LiteLLM.AgenticHookError: Exception in "
|
||||
"async_agentic_loop_cleanup_hook [call_id=%s model=%s]: %s",
|
||||
_call_id,
|
||||
model,
|
||||
str(e),
|
||||
)
|
||||
|
||||
|
||||
async def maybe_run_chat_completion_agentic_loop(
|
||||
*,
|
||||
response: ModelResponse,
|
||||
model: str,
|
||||
messages: list,
|
||||
optional_params: dict,
|
||||
kwargs: dict,
|
||||
logging_obj: object,
|
||||
custom_llm_provider: str,
|
||||
stream: bool,
|
||||
) -> ModelResponse | CustomStreamWrapper | None:
|
||||
import litellm
|
||||
|
||||
callbacks = litellm.callbacks + (
|
||||
getattr(logging_obj, "dynamic_success_callbacks", None) or []
|
||||
)
|
||||
depth, max_loops, fingerprints = _agentic_loop_settings(kwargs)
|
||||
tools = optional_params.get("tools", [])
|
||||
|
||||
for callback in callbacks:
|
||||
if not isinstance(callback, CustomLogger):
|
||||
continue
|
||||
if not _gate_overridden(callback):
|
||||
continue
|
||||
|
||||
gate_kwargs = {
|
||||
**kwargs,
|
||||
"_agentic_loop_api_surface": CHAT_COMPLETION_AGENTIC_SURFACE,
|
||||
"custom_llm_provider": custom_llm_provider,
|
||||
}
|
||||
try:
|
||||
should_run, tool_calls = await callback.async_should_run_agentic_loop(
|
||||
response=response,
|
||||
model=model,
|
||||
messages=messages,
|
||||
tools=tools,
|
||||
stream=stream,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
kwargs=gate_kwargs,
|
||||
)
|
||||
except Exception as e:
|
||||
verbose_logger.exception(
|
||||
"LiteLLM.AgenticHookError: Exception in chat completion agentic gate: %s",
|
||||
str(e),
|
||||
)
|
||||
continue
|
||||
|
||||
if not should_run:
|
||||
continue
|
||||
|
||||
fingerprint = _check_agentic_loop_safety(
|
||||
tool_calls=tool_calls,
|
||||
fingerprints=fingerprints,
|
||||
depth=depth,
|
||||
max_loops=max_loops,
|
||||
model=model,
|
||||
)
|
||||
|
||||
try:
|
||||
plan_kwargs = {
|
||||
**kwargs,
|
||||
"_agentic_loop_api_surface": CHAT_COMPLETION_AGENTIC_SURFACE,
|
||||
"custom_llm_provider": custom_llm_provider,
|
||||
}
|
||||
if not _build_plan_overridden(callback):
|
||||
return await callback.async_run_agentic_loop(
|
||||
tools=tool_calls,
|
||||
model=model,
|
||||
messages=messages,
|
||||
response=response,
|
||||
anthropic_messages_provider_config=None,
|
||||
anthropic_messages_optional_request_params=optional_params,
|
||||
logging_obj=logging_obj,
|
||||
stream=stream,
|
||||
kwargs=plan_kwargs,
|
||||
)
|
||||
|
||||
plan = await callback.async_build_agentic_loop_plan(
|
||||
tools=tool_calls,
|
||||
model=model,
|
||||
messages=messages,
|
||||
response=response,
|
||||
anthropic_messages_provider_config=None,
|
||||
anthropic_messages_optional_request_params=optional_params,
|
||||
logging_obj=logging_obj,
|
||||
stream=stream,
|
||||
kwargs=plan_kwargs,
|
||||
)
|
||||
|
||||
if plan.response_override is not None:
|
||||
return plan.response_override
|
||||
if plan.terminate:
|
||||
return response
|
||||
if not plan.run_agentic_loop:
|
||||
continue
|
||||
|
||||
return await _execute_chat_completion_agentic_plan(
|
||||
plan=plan,
|
||||
callback=callback,
|
||||
model=model,
|
||||
optional_params=optional_params,
|
||||
kwargs=kwargs,
|
||||
logging_obj=logging_obj,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
depth=depth,
|
||||
max_loops=max_loops,
|
||||
fingerprints=fingerprints,
|
||||
fingerprint=fingerprint,
|
||||
)
|
||||
except Exception as e:
|
||||
verbose_logger.exception(
|
||||
"LiteLLM.AgenticHookError: Exception in chat completion agentic hooks: %s",
|
||||
str(e),
|
||||
)
|
||||
|
||||
if (
|
||||
kwargs.get("_code_interpreter_interception_converted_stream")
|
||||
and not depth
|
||||
and hasattr(response, "choices")
|
||||
):
|
||||
return cast(
|
||||
"ModelResponse | CustomStreamWrapper",
|
||||
_wrap_response_as_fake_stream(response),
|
||||
)
|
||||
return None
|
||||
|
|
@ -6,10 +6,7 @@ from typing import Callable, Optional, Union
|
|||
|
||||
import httpx
|
||||
|
||||
from litellm.constants import (
|
||||
COMPLETION_HTTP_FALLBACK_SECONDS,
|
||||
DEFAULT_REQUEST_TIMEOUT_SECONDS,
|
||||
)
|
||||
from litellm.constants import COMPLETION_HTTP_FALLBACK_SECONDS
|
||||
|
||||
|
||||
class CompletionTimeout:
|
||||
|
|
@ -22,17 +19,13 @@ class CompletionTimeout:
|
|||
"""
|
||||
Used when ``model_timeout`` and kwargs timeouts are all unset.
|
||||
|
||||
``global_timeout`` is :attr:`litellm.request_timeout` (numeric / string), not
|
||||
:class:`httpx.Timeout`.
|
||||
|
||||
If it equals :data:`~litellm.constants.DEFAULT_REQUEST_TIMEOUT_SECONDS` (6000),
|
||||
return :data:`~litellm.constants.COMPLETION_HTTP_FALLBACK_SECONDS`. Same if
|
||||
``None``. Otherwise return ``float(global_timeout)``.
|
||||
``global_timeout`` is the explicitly-configured ``litellm.request_timeout``
|
||||
(numeric / string) or ``None`` when it was never set. ``None`` falls back to
|
||||
:data:`~litellm.constants.COMPLETION_HTTP_FALLBACK_SECONDS`; any explicit value
|
||||
(including ``6000``) is honored.
|
||||
"""
|
||||
if global_timeout is None:
|
||||
return COMPLETION_HTTP_FALLBACK_SECONDS
|
||||
if float(global_timeout) == float(DEFAULT_REQUEST_TIMEOUT_SECONDS):
|
||||
return COMPLETION_HTTP_FALLBACK_SECONDS
|
||||
return float(global_timeout)
|
||||
|
||||
@staticmethod
|
||||
|
|
@ -50,11 +43,10 @@ class CompletionTimeout:
|
|||
1. ``model_timeout`` (call argument / merged ``litellm_params``)
|
||||
2. ``kwargs["timeout"]``
|
||||
3. ``kwargs["request_timeout"]``
|
||||
4. Fallback from ``global_timeout`` (:attr:`litellm.request_timeout`) — if it is
|
||||
the package default (6000), use 600 instead.
|
||||
4. ``global_timeout`` (the explicitly-configured ``litellm.request_timeout``),
|
||||
or 600 when nothing was configured.
|
||||
|
||||
Coerce :class:`httpx.Timeout` when the provider does not support it.
|
||||
Explicit ``6000`` on the model or in kwargs is kept as ``6000``.
|
||||
"""
|
||||
resolved: Union[float, str, httpx.Timeout]
|
||||
if model_timeout is not None:
|
||||
|
|
|
|||
File diff suppressed because it is too large
Load diff
|
|
@ -86,9 +86,7 @@ def get_supported_openai_params(
|
|||
model=model
|
||||
)
|
||||
elif request_type == "transcription":
|
||||
return litellm.FireworksAIAudioTranscriptionConfig().get_supported_openai_params(
|
||||
model=model
|
||||
)
|
||||
return None
|
||||
else:
|
||||
return litellm.FireworksAIConfig().get_supported_openai_params(model=model)
|
||||
elif custom_llm_provider == "nvidia_nim":
|
||||
|
|
@ -191,7 +189,9 @@ def get_supported_openai_params(
|
|||
)
|
||||
elif custom_llm_provider == "sambanova":
|
||||
if request_type == "embeddings":
|
||||
litellm.SambaNovaEmbeddingConfig().get_supported_openai_params(model=model)
|
||||
return litellm.SambaNovaEmbeddingConfig().get_supported_openai_params(
|
||||
model=model
|
||||
)
|
||||
else:
|
||||
return litellm.SambanovaConfig().get_supported_openai_params(model=model)
|
||||
elif custom_llm_provider == "nebius":
|
||||
|
|
|
|||
29
litellm/litellm_core_utils/request_timeout_resolver.py
Normal file
29
litellm/litellm_core_utils/request_timeout_resolver.py
Normal file
|
|
@ -0,0 +1,29 @@
|
|||
"""Single source of truth for whether ``litellm.request_timeout`` was configured.
|
||||
|
||||
``litellm.request_timeout`` always holds a value (the package default,
|
||||
:data:`~litellm.constants.DEFAULT_REQUEST_TIMEOUT_SECONDS`), so a bare read can't
|
||||
tell "user asked for this" from "nobody set it". This resolver answers that:
|
||||
|
||||
* ``request_timeout_explicitly_set`` is the authoritative signal, set when the
|
||||
value comes from the ``REQUEST_TIMEOUT`` env var or ``litellm_settings``.
|
||||
* A runtime value that differs from the package default (e.g. ``litellm.request_timeout
|
||||
= 300`` in SDK code) is also treated as explicit, for backwards compatibility.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Optional
|
||||
|
||||
from litellm.constants import DEFAULT_REQUEST_TIMEOUT_SECONDS
|
||||
|
||||
|
||||
def get_configured_request_timeout() -> Optional[float]:
|
||||
"""Return the explicitly-configured ``litellm.request_timeout``, else ``None``."""
|
||||
import litellm
|
||||
|
||||
timeout = float(litellm.request_timeout)
|
||||
if litellm.request_timeout_explicitly_set:
|
||||
return timeout
|
||||
if timeout != float(DEFAULT_REQUEST_TIMEOUT_SECONDS):
|
||||
return timeout
|
||||
return None
|
||||
|
|
@ -12,6 +12,7 @@ class SensitiveDataMasker:
|
|||
visible_prefix: int = 4,
|
||||
visible_suffix: int = 4,
|
||||
mask_char: str = "*",
|
||||
mask_short_values: bool = True,
|
||||
):
|
||||
self.sensitive_patterns = sensitive_patterns or {
|
||||
"password",
|
||||
|
|
@ -38,12 +39,17 @@ class SensitiveDataMasker:
|
|||
self.visible_prefix = visible_prefix
|
||||
self.visible_suffix = visible_suffix
|
||||
self.mask_char = mask_char
|
||||
self.mask_short_values = mask_short_values
|
||||
|
||||
def _mask_value(self, value: str) -> str:
|
||||
if not value or len(str(value)) < (self.visible_prefix + self.visible_suffix):
|
||||
return value
|
||||
|
||||
value_str = str(value)
|
||||
if not value_str:
|
||||
return value
|
||||
if len(value_str) <= (self.visible_prefix + self.visible_suffix):
|
||||
return (
|
||||
self.mask_char * len(value_str) if self.mask_short_values else value_str
|
||||
)
|
||||
|
||||
masked_length = len(value_str) - (self.visible_prefix + self.visible_suffix)
|
||||
|
||||
# Handle the case where visible_suffix is 0 to avoid showing the entire string
|
||||
|
|
|
|||
|
|
@ -6,6 +6,7 @@ import logging
|
|||
import threading
|
||||
import time
|
||||
import traceback
|
||||
from dataclasses import dataclass
|
||||
from typing import (
|
||||
Any,
|
||||
AsyncIterator,
|
||||
|
|
@ -97,6 +98,19 @@ def print_verbose(print_statement):
|
|||
pass
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _ProviderChunkParsed:
|
||||
response_obj: dict[str, Any]
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _ProviderChunkEarlyReturn:
|
||||
value: Any
|
||||
|
||||
|
||||
_ProviderChunkResult = Union[_ProviderChunkParsed, _ProviderChunkEarlyReturn]
|
||||
|
||||
|
||||
class CustomStreamWrapper:
|
||||
def __init__(
|
||||
self,
|
||||
|
|
@ -1145,381 +1159,392 @@ class CustomStreamWrapper:
|
|||
del model_response.choices[0].delta.reasoning_content
|
||||
return
|
||||
|
||||
def _dispatch_provider_chunk(
|
||||
self,
|
||||
chunk: Any,
|
||||
model_response: ModelResponseStream,
|
||||
completion_obj: dict[str, Any],
|
||||
) -> _ProviderChunkResult:
|
||||
response_obj: dict[str, Any] = {}
|
||||
if (
|
||||
isinstance(chunk, ModelResponseStream)
|
||||
and self.custom_llm_provider is not None
|
||||
and self.custom_llm_provider in litellm._custom_providers
|
||||
):
|
||||
_has_content = bool(
|
||||
chunk.choices
|
||||
and chunk.choices[0].delta is not None
|
||||
and (
|
||||
chunk.choices[0].delta.content or chunk.choices[0].delta.tool_calls
|
||||
)
|
||||
)
|
||||
if self.received_finish_reason is not None:
|
||||
if not _has_content:
|
||||
raise StopIteration
|
||||
if chunk.choices and chunk.choices[0].finish_reason:
|
||||
self.received_finish_reason = chunk.choices[0].finish_reason
|
||||
if not _has_content:
|
||||
return _ProviderChunkEarlyReturn(None)
|
||||
# Strip finish_reason from the content chunk so it appears
|
||||
# only on the trailing empty-delta chunk (OpenAI spec).
|
||||
# finish_reason_handler() will emit the proper terminal chunk.
|
||||
chunk.choices[0].finish_reason = None # type: ignore[assignment]
|
||||
return _ProviderChunkEarlyReturn(chunk)
|
||||
|
||||
if (
|
||||
isinstance(chunk, dict)
|
||||
and generic_chunk_has_all_required_fields(
|
||||
chunk=chunk
|
||||
) # check if chunk is a generic streaming chunk
|
||||
) or (
|
||||
self.custom_llm_provider
|
||||
and self.custom_llm_provider in litellm._custom_providers
|
||||
):
|
||||
if self.received_finish_reason is not None:
|
||||
_chunk_has_content = isinstance(chunk, dict) and (
|
||||
bool(chunk.get("text", ""))
|
||||
or chunk.get("tool_use") is not None
|
||||
# Usage-only final chunks are valid and needed to surface
|
||||
# finish_reason/usage to downstream translators.
|
||||
or chunk.get("usage") is not None
|
||||
)
|
||||
if not _chunk_has_content and (
|
||||
not isinstance(chunk, dict)
|
||||
or "provider_specific_fields" not in chunk
|
||||
):
|
||||
raise StopIteration
|
||||
anthropic_response_obj: GChunk = cast(GChunk, chunk)
|
||||
completion_obj["content"] = anthropic_response_obj["text"]
|
||||
if anthropic_response_obj["is_finished"]:
|
||||
self.received_finish_reason = anthropic_response_obj["finish_reason"]
|
||||
|
||||
if anthropic_response_obj["finish_reason"]:
|
||||
self.intermittent_finish_reason = anthropic_response_obj[
|
||||
"finish_reason"
|
||||
]
|
||||
|
||||
if anthropic_response_obj["usage"] is not None:
|
||||
setattr(
|
||||
model_response,
|
||||
"usage",
|
||||
litellm.Usage(**anthropic_response_obj["usage"]),
|
||||
)
|
||||
|
||||
if (
|
||||
"tool_use" in anthropic_response_obj
|
||||
and anthropic_response_obj["tool_use"] is not None
|
||||
):
|
||||
completion_obj["tool_calls"] = [anthropic_response_obj["tool_use"]]
|
||||
|
||||
if (
|
||||
"provider_specific_fields" in anthropic_response_obj
|
||||
and anthropic_response_obj["provider_specific_fields"] is not None
|
||||
):
|
||||
for key, value in anthropic_response_obj[
|
||||
"provider_specific_fields"
|
||||
].items():
|
||||
setattr(model_response, key, value)
|
||||
|
||||
response_obj = cast(dict[str, Any], anthropic_response_obj)
|
||||
elif self.model == "replicate" or self.custom_llm_provider == "replicate":
|
||||
response_obj = self.handle_replicate_chunk(chunk)
|
||||
completion_obj["content"] = response_obj["text"]
|
||||
if response_obj["is_finished"]:
|
||||
self.received_finish_reason = response_obj["finish_reason"]
|
||||
elif self.custom_llm_provider and self.custom_llm_provider == "predibase":
|
||||
response_obj = self.handle_predibase_chunk(chunk)
|
||||
completion_obj["content"] = response_obj["text"]
|
||||
if response_obj["is_finished"]:
|
||||
self.received_finish_reason = response_obj["finish_reason"]
|
||||
elif (
|
||||
self.custom_llm_provider and self.custom_llm_provider == "baseten"
|
||||
): # baseten doesn't provide streaming
|
||||
completion_obj["content"] = self.handle_baseten_chunk(chunk)
|
||||
elif (
|
||||
self.custom_llm_provider and self.custom_llm_provider == "ai21"
|
||||
): # ai21 doesn't provide streaming
|
||||
response_obj = self.handle_ai21_chunk(chunk)
|
||||
completion_obj["content"] = response_obj["text"]
|
||||
if response_obj["is_finished"]:
|
||||
self.received_finish_reason = response_obj["finish_reason"]
|
||||
elif self.custom_llm_provider and self.custom_llm_provider == "maritalk":
|
||||
response_obj = self.handle_maritalk_chunk(chunk)
|
||||
completion_obj["content"] = response_obj["text"]
|
||||
if response_obj["is_finished"]:
|
||||
self.received_finish_reason = response_obj["finish_reason"]
|
||||
elif self.custom_llm_provider and self.custom_llm_provider == "vllm":
|
||||
completion_obj["content"] = chunk[0].outputs[0].text
|
||||
elif (
|
||||
self.custom_llm_provider and self.custom_llm_provider == "aleph_alpha"
|
||||
): # aleph alpha doesn't provide streaming
|
||||
response_obj = self.handle_aleph_alpha_chunk(chunk)
|
||||
completion_obj["content"] = response_obj["text"]
|
||||
if response_obj["is_finished"]:
|
||||
self.received_finish_reason = response_obj["finish_reason"]
|
||||
elif self.custom_llm_provider == "nlp_cloud":
|
||||
try:
|
||||
response_obj = self.handle_nlp_cloud_chunk(chunk)
|
||||
completion_obj["content"] = response_obj["text"]
|
||||
if response_obj["is_finished"]:
|
||||
self.received_finish_reason = response_obj["finish_reason"]
|
||||
except Exception as e:
|
||||
if self.received_finish_reason:
|
||||
raise e
|
||||
else:
|
||||
if self.sent_first_chunk is False:
|
||||
raise Exception("An unknown error occurred with the stream")
|
||||
self.received_finish_reason = "stop"
|
||||
elif self.custom_llm_provider == "vertex_ai" and not isinstance(
|
||||
chunk, ModelResponseStream
|
||||
):
|
||||
chunk = cast(Any, chunk)
|
||||
import proto # type: ignore
|
||||
|
||||
if hasattr(chunk, "candidates") is True:
|
||||
try:
|
||||
try:
|
||||
completion_obj["content"] = chunk.text # type: ignore
|
||||
except Exception as e:
|
||||
original_exception = e
|
||||
if "Part has no text." in str(e):
|
||||
## check for function calling
|
||||
function_call = (
|
||||
chunk.candidates[0].content.parts[0].function_call # type: ignore
|
||||
)
|
||||
|
||||
args_dict = {}
|
||||
|
||||
# Check if it's a RepeatedComposite instance
|
||||
for key, val in function_call.args.items():
|
||||
if isinstance(
|
||||
val,
|
||||
proto.marshal.collections.repeated.RepeatedComposite, # type: ignore
|
||||
):
|
||||
# If so, convert to list
|
||||
args_dict[key] = [v for v in val]
|
||||
else:
|
||||
args_dict[key] = val
|
||||
|
||||
try:
|
||||
args_str = json.dumps(args_dict)
|
||||
except Exception as e:
|
||||
raise e
|
||||
_delta_obj = litellm.utils.Delta(
|
||||
content=None,
|
||||
tool_calls=[
|
||||
{
|
||||
"id": f"call_{str(uuid.uuid4())}",
|
||||
"function": {
|
||||
"arguments": args_str,
|
||||
"name": function_call.name,
|
||||
},
|
||||
"type": "function",
|
||||
}
|
||||
],
|
||||
)
|
||||
_streaming_response = StreamingChoices(delta=_delta_obj)
|
||||
_model_response = ModelResponseStream()
|
||||
_model_response.choices = [_streaming_response]
|
||||
response_obj = {"original_chunk": _model_response}
|
||||
else:
|
||||
raise original_exception
|
||||
if (
|
||||
hasattr(chunk.candidates[0], "finish_reason") # type: ignore
|
||||
and chunk.candidates[0].finish_reason.name # type: ignore
|
||||
!= "FINISH_REASON_UNSPECIFIED"
|
||||
): # every non-final chunk in vertex ai has this
|
||||
self.received_finish_reason = map_finish_reason( # type: ignore
|
||||
chunk.candidates[0].finish_reason.name
|
||||
)
|
||||
except Exception:
|
||||
if chunk.candidates[0].finish_reason.name == "SAFETY": # type: ignore
|
||||
raise Exception(
|
||||
f"The response was blocked by VertexAI. {str(chunk)}"
|
||||
)
|
||||
else:
|
||||
completion_obj["content"] = str(chunk)
|
||||
elif self.custom_llm_provider == "petals":
|
||||
if self.completion_stream is None or len(self.completion_stream) == 0:
|
||||
if self.received_finish_reason is not None:
|
||||
raise StopIteration
|
||||
else:
|
||||
self.received_finish_reason = "stop"
|
||||
chunk_size = 30
|
||||
stream = cast(Any, self.completion_stream)
|
||||
new_chunk = stream[:chunk_size]
|
||||
completion_obj["content"] = new_chunk
|
||||
self.completion_stream = stream[chunk_size:]
|
||||
elif self.custom_llm_provider == "palm":
|
||||
# fake streaming
|
||||
response_obj = {}
|
||||
if self.completion_stream is None or len(self.completion_stream) == 0:
|
||||
if self.received_finish_reason is not None:
|
||||
raise StopIteration
|
||||
else:
|
||||
self.received_finish_reason = "stop"
|
||||
chunk_size = 30
|
||||
stream = cast(Any, self.completion_stream)
|
||||
new_chunk = stream[:chunk_size]
|
||||
completion_obj["content"] = new_chunk
|
||||
self.completion_stream = stream[chunk_size:]
|
||||
elif self.custom_llm_provider == "triton":
|
||||
response_obj = self.handle_triton_stream(chunk)
|
||||
completion_obj["content"] = response_obj["text"]
|
||||
print_verbose(f"completion obj content: {completion_obj['content']}")
|
||||
if response_obj["is_finished"]:
|
||||
self.received_finish_reason = response_obj["finish_reason"]
|
||||
elif self.custom_llm_provider == "text-completion-openai":
|
||||
response_obj = self.handle_openai_text_completion_chunk(chunk)
|
||||
completion_obj["content"] = response_obj["text"]
|
||||
print_verbose(f"completion obj content: {completion_obj['content']}")
|
||||
if response_obj["is_finished"]:
|
||||
self.received_finish_reason = response_obj["finish_reason"]
|
||||
if response_obj["usage"] is not None:
|
||||
setattr(
|
||||
model_response,
|
||||
"usage",
|
||||
litellm.Usage(
|
||||
prompt_tokens=response_obj["usage"].prompt_tokens,
|
||||
completion_tokens=response_obj["usage"].completion_tokens,
|
||||
total_tokens=response_obj["usage"].total_tokens,
|
||||
),
|
||||
)
|
||||
elif self.custom_llm_provider == "text-completion-codestral":
|
||||
if not isinstance(chunk, str):
|
||||
raise ValueError(f"chunk is not a string: {chunk}")
|
||||
response_obj = cast(
|
||||
dict[str, Any],
|
||||
litellm.CodestralTextCompletionConfig()._chunk_parser(chunk),
|
||||
)
|
||||
completion_obj["content"] = response_obj["text"]
|
||||
print_verbose(f"completion obj content: {completion_obj['content']}")
|
||||
if response_obj["is_finished"]:
|
||||
self.received_finish_reason = response_obj["finish_reason"]
|
||||
if "usage" in response_obj is not None:
|
||||
setattr(
|
||||
model_response,
|
||||
"usage",
|
||||
litellm.Usage(
|
||||
prompt_tokens=response_obj["usage"].prompt_tokens,
|
||||
completion_tokens=response_obj["usage"].completion_tokens,
|
||||
total_tokens=response_obj["usage"].total_tokens,
|
||||
),
|
||||
)
|
||||
elif self.custom_llm_provider == "azure_text":
|
||||
response_obj = self.handle_azure_text_completion_chunk(chunk)
|
||||
completion_obj["content"] = response_obj["text"]
|
||||
print_verbose(f"completion obj content: {completion_obj['content']}")
|
||||
if response_obj["is_finished"]:
|
||||
self.received_finish_reason = response_obj["finish_reason"]
|
||||
elif self.custom_llm_provider == "cached_response":
|
||||
chunk = cast(ModelResponseStream, chunk)
|
||||
response_obj = {
|
||||
"text": chunk.choices[0].delta.content,
|
||||
"is_finished": True,
|
||||
"finish_reason": chunk.choices[0].finish_reason,
|
||||
"original_chunk": chunk,
|
||||
"tool_calls": (
|
||||
chunk.choices[0].delta.tool_calls
|
||||
if hasattr(chunk.choices[0].delta, "tool_calls")
|
||||
else None
|
||||
),
|
||||
}
|
||||
|
||||
completion_obj["content"] = response_obj["text"]
|
||||
if response_obj["tool_calls"] is not None:
|
||||
completion_obj["tool_calls"] = response_obj["tool_calls"]
|
||||
print_verbose(f"completion obj content: {completion_obj['content']}")
|
||||
if hasattr(chunk, "id"):
|
||||
model_response.id = chunk.id
|
||||
self.response_id = chunk.id
|
||||
if hasattr(chunk, "system_fingerprint"):
|
||||
self.system_fingerprint = chunk.system_fingerprint
|
||||
if response_obj["is_finished"]:
|
||||
self.received_finish_reason = response_obj["finish_reason"]
|
||||
else: # openai / azure chat model
|
||||
if self.custom_llm_provider in [
|
||||
LlmProviders.AZURE.value,
|
||||
LlmProviders.AZURE_AI.value,
|
||||
]:
|
||||
if isinstance(chunk, BaseModel) and hasattr(chunk, "model"):
|
||||
# for azure, we need to pass the model from the original chunk
|
||||
self.model = getattr(chunk, "model", self.model)
|
||||
response_obj = self.handle_openai_chat_completion_chunk(chunk)
|
||||
if response_obj is None:
|
||||
return _ProviderChunkEarlyReturn(None)
|
||||
completion_obj["content"] = response_obj["text"]
|
||||
self.intermittent_finish_reason = response_obj.get("finish_reason", None)
|
||||
if response_obj["is_finished"]:
|
||||
if response_obj["finish_reason"] == "error":
|
||||
raise Exception(
|
||||
"{} raised a streaming error - finish_reason: error, no content string given. Received Chunk={}".format(
|
||||
self.custom_llm_provider, response_obj
|
||||
)
|
||||
)
|
||||
self.received_finish_reason = response_obj["finish_reason"]
|
||||
if response_obj.get("original_chunk", None) is not None:
|
||||
if hasattr(response_obj["original_chunk"], "id"):
|
||||
model_response = self.set_model_id(
|
||||
response_obj["original_chunk"].id, model_response
|
||||
)
|
||||
if hasattr(response_obj["original_chunk"], "system_fingerprint"):
|
||||
model_response.system_fingerprint = response_obj[
|
||||
"original_chunk"
|
||||
].system_fingerprint
|
||||
self.system_fingerprint = response_obj[
|
||||
"original_chunk"
|
||||
].system_fingerprint
|
||||
if response_obj["logprobs"] is not None:
|
||||
model_response.choices[0].logprobs = response_obj["logprobs"]
|
||||
|
||||
if response_obj["usage"] is not None:
|
||||
if isinstance(response_obj["usage"], dict):
|
||||
setattr(
|
||||
model_response,
|
||||
"usage",
|
||||
litellm.Usage(
|
||||
prompt_tokens=response_obj["usage"].get(
|
||||
"prompt_tokens", None
|
||||
)
|
||||
or None,
|
||||
completion_tokens=response_obj["usage"].get(
|
||||
"completion_tokens", None
|
||||
)
|
||||
or None,
|
||||
total_tokens=response_obj["usage"].get("total_tokens", None)
|
||||
or None,
|
||||
),
|
||||
)
|
||||
elif isinstance(response_obj["usage"], Usage):
|
||||
setattr(
|
||||
model_response,
|
||||
"usage",
|
||||
response_obj["usage"],
|
||||
)
|
||||
elif isinstance(response_obj["usage"], BaseModel):
|
||||
setattr(
|
||||
model_response,
|
||||
"usage",
|
||||
litellm.Usage(**response_obj["usage"].model_dump()),
|
||||
)
|
||||
return _ProviderChunkParsed(response_obj)
|
||||
|
||||
def chunk_creator(self, chunk: Any): # type: ignore
|
||||
if hasattr(chunk, "id"):
|
||||
self.response_id = chunk.id
|
||||
model_response = self.model_response_creator()
|
||||
response_obj: Dict[str, Any] = {}
|
||||
response_obj: dict[str, Any] = {}
|
||||
try:
|
||||
# return this for all models
|
||||
completion_obj: Dict[str, Any] = {"content": ""}
|
||||
from litellm.types.utils import GenericStreamingChunk as GChunk
|
||||
|
||||
if (
|
||||
isinstance(chunk, ModelResponseStream)
|
||||
and self.custom_llm_provider is not None
|
||||
and self.custom_llm_provider in litellm._custom_providers
|
||||
):
|
||||
_has_content = bool(
|
||||
chunk.choices
|
||||
and chunk.choices[0].delta is not None
|
||||
and (
|
||||
chunk.choices[0].delta.content
|
||||
or chunk.choices[0].delta.tool_calls
|
||||
)
|
||||
)
|
||||
if self.received_finish_reason is not None:
|
||||
if not _has_content:
|
||||
raise StopIteration
|
||||
if chunk.choices and chunk.choices[0].finish_reason:
|
||||
self.received_finish_reason = chunk.choices[0].finish_reason
|
||||
if not _has_content:
|
||||
return None
|
||||
# Strip finish_reason from the content chunk so it appears
|
||||
# only on the trailing empty-delta chunk (OpenAI spec).
|
||||
# finish_reason_handler() will emit the proper terminal chunk.
|
||||
chunk.choices[0].finish_reason = None # type: ignore[assignment]
|
||||
return chunk
|
||||
|
||||
if (
|
||||
isinstance(chunk, dict)
|
||||
and generic_chunk_has_all_required_fields(
|
||||
chunk=chunk
|
||||
) # check if chunk is a generic streaming chunk
|
||||
) or (
|
||||
self.custom_llm_provider
|
||||
and self.custom_llm_provider in litellm._custom_providers
|
||||
):
|
||||
if self.received_finish_reason is not None:
|
||||
_chunk_has_content = isinstance(chunk, dict) and (
|
||||
bool(chunk.get("text", ""))
|
||||
or chunk.get("tool_use") is not None
|
||||
# Usage-only final chunks are valid and needed to surface
|
||||
# finish_reason/usage to downstream translators.
|
||||
or chunk.get("usage") is not None
|
||||
)
|
||||
if not _chunk_has_content and (
|
||||
not isinstance(chunk, dict)
|
||||
or "provider_specific_fields" not in chunk
|
||||
):
|
||||
raise StopIteration
|
||||
anthropic_response_obj: GChunk = cast(GChunk, chunk)
|
||||
completion_obj["content"] = anthropic_response_obj["text"]
|
||||
if anthropic_response_obj["is_finished"]:
|
||||
self.received_finish_reason = anthropic_response_obj[
|
||||
"finish_reason"
|
||||
]
|
||||
|
||||
if anthropic_response_obj["finish_reason"]:
|
||||
self.intermittent_finish_reason = anthropic_response_obj[
|
||||
"finish_reason"
|
||||
]
|
||||
|
||||
if anthropic_response_obj["usage"] is not None:
|
||||
setattr(
|
||||
model_response,
|
||||
"usage",
|
||||
litellm.Usage(**anthropic_response_obj["usage"]),
|
||||
)
|
||||
|
||||
if (
|
||||
"tool_use" in anthropic_response_obj
|
||||
and anthropic_response_obj["tool_use"] is not None
|
||||
):
|
||||
completion_obj["tool_calls"] = [anthropic_response_obj["tool_use"]]
|
||||
|
||||
if (
|
||||
"provider_specific_fields" in anthropic_response_obj
|
||||
and anthropic_response_obj["provider_specific_fields"] is not None
|
||||
):
|
||||
for key, value in anthropic_response_obj[
|
||||
"provider_specific_fields"
|
||||
].items():
|
||||
setattr(model_response, key, value)
|
||||
|
||||
response_obj = cast(Dict[str, Any], anthropic_response_obj)
|
||||
elif self.model == "replicate" or self.custom_llm_provider == "replicate":
|
||||
response_obj = self.handle_replicate_chunk(chunk)
|
||||
completion_obj["content"] = response_obj["text"]
|
||||
if response_obj["is_finished"]:
|
||||
self.received_finish_reason = response_obj["finish_reason"]
|
||||
elif self.custom_llm_provider and self.custom_llm_provider == "predibase":
|
||||
response_obj = self.handle_predibase_chunk(chunk)
|
||||
completion_obj["content"] = response_obj["text"]
|
||||
if response_obj["is_finished"]:
|
||||
self.received_finish_reason = response_obj["finish_reason"]
|
||||
elif (
|
||||
self.custom_llm_provider and self.custom_llm_provider == "baseten"
|
||||
): # baseten doesn't provide streaming
|
||||
completion_obj["content"] = self.handle_baseten_chunk(chunk)
|
||||
elif (
|
||||
self.custom_llm_provider and self.custom_llm_provider == "ai21"
|
||||
): # ai21 doesn't provide streaming
|
||||
response_obj = self.handle_ai21_chunk(chunk)
|
||||
completion_obj["content"] = response_obj["text"]
|
||||
if response_obj["is_finished"]:
|
||||
self.received_finish_reason = response_obj["finish_reason"]
|
||||
elif self.custom_llm_provider and self.custom_llm_provider == "maritalk":
|
||||
response_obj = self.handle_maritalk_chunk(chunk)
|
||||
completion_obj["content"] = response_obj["text"]
|
||||
if response_obj["is_finished"]:
|
||||
self.received_finish_reason = response_obj["finish_reason"]
|
||||
elif self.custom_llm_provider and self.custom_llm_provider == "vllm":
|
||||
completion_obj["content"] = chunk[0].outputs[0].text
|
||||
elif (
|
||||
self.custom_llm_provider and self.custom_llm_provider == "aleph_alpha"
|
||||
): # aleph alpha doesn't provide streaming
|
||||
response_obj = self.handle_aleph_alpha_chunk(chunk)
|
||||
completion_obj["content"] = response_obj["text"]
|
||||
if response_obj["is_finished"]:
|
||||
self.received_finish_reason = response_obj["finish_reason"]
|
||||
elif self.custom_llm_provider == "nlp_cloud":
|
||||
try:
|
||||
response_obj = self.handle_nlp_cloud_chunk(chunk)
|
||||
completion_obj["content"] = response_obj["text"]
|
||||
if response_obj["is_finished"]:
|
||||
self.received_finish_reason = response_obj["finish_reason"]
|
||||
except Exception as e:
|
||||
if self.received_finish_reason:
|
||||
raise e
|
||||
else:
|
||||
if self.sent_first_chunk is False:
|
||||
raise Exception("An unknown error occurred with the stream")
|
||||
self.received_finish_reason = "stop"
|
||||
elif self.custom_llm_provider == "vertex_ai" and not isinstance(
|
||||
chunk, ModelResponseStream
|
||||
):
|
||||
import proto # type: ignore
|
||||
|
||||
if hasattr(chunk, "candidates") is True:
|
||||
try:
|
||||
try:
|
||||
completion_obj["content"] = chunk.text # type: ignore
|
||||
except Exception as e:
|
||||
original_exception = e
|
||||
if "Part has no text." in str(e):
|
||||
## check for function calling
|
||||
function_call = (
|
||||
chunk.candidates[0].content.parts[0].function_call # type: ignore
|
||||
)
|
||||
|
||||
args_dict = {}
|
||||
|
||||
# Check if it's a RepeatedComposite instance
|
||||
for key, val in function_call.args.items():
|
||||
if isinstance(
|
||||
val,
|
||||
proto.marshal.collections.repeated.RepeatedComposite, # type: ignore
|
||||
):
|
||||
# If so, convert to list
|
||||
args_dict[key] = [v for v in val]
|
||||
else:
|
||||
args_dict[key] = val
|
||||
|
||||
try:
|
||||
args_str = json.dumps(args_dict)
|
||||
except Exception as e:
|
||||
raise e
|
||||
_delta_obj = litellm.utils.Delta(
|
||||
content=None,
|
||||
tool_calls=[
|
||||
{
|
||||
"id": f"call_{str(uuid.uuid4())}",
|
||||
"function": {
|
||||
"arguments": args_str,
|
||||
"name": function_call.name,
|
||||
},
|
||||
"type": "function",
|
||||
}
|
||||
],
|
||||
)
|
||||
_streaming_response = StreamingChoices(delta=_delta_obj)
|
||||
_model_response = ModelResponseStream()
|
||||
_model_response.choices = [_streaming_response]
|
||||
response_obj = {"original_chunk": _model_response}
|
||||
else:
|
||||
raise original_exception
|
||||
if (
|
||||
hasattr(chunk.candidates[0], "finish_reason") # type: ignore
|
||||
and chunk.candidates[0].finish_reason.name # type: ignore
|
||||
!= "FINISH_REASON_UNSPECIFIED"
|
||||
): # every non-final chunk in vertex ai has this
|
||||
self.received_finish_reason = map_finish_reason( # type: ignore
|
||||
chunk.candidates[0].finish_reason.name
|
||||
)
|
||||
except Exception:
|
||||
if chunk.candidates[0].finish_reason.name == "SAFETY": # type: ignore
|
||||
raise Exception(
|
||||
f"The response was blocked by VertexAI. {str(chunk)}"
|
||||
)
|
||||
else:
|
||||
completion_obj["content"] = str(chunk)
|
||||
elif self.custom_llm_provider == "petals":
|
||||
if self.completion_stream is None or len(self.completion_stream) == 0:
|
||||
if self.received_finish_reason is not None:
|
||||
raise StopIteration
|
||||
else:
|
||||
self.received_finish_reason = "stop"
|
||||
chunk_size = 30
|
||||
new_chunk = self.completion_stream[:chunk_size] # type: ignore[index]
|
||||
completion_obj["content"] = new_chunk
|
||||
self.completion_stream = self.completion_stream[chunk_size:] # type: ignore[index]
|
||||
elif self.custom_llm_provider == "palm":
|
||||
# fake streaming
|
||||
response_obj = {}
|
||||
if self.completion_stream is None or len(self.completion_stream) == 0:
|
||||
if self.received_finish_reason is not None:
|
||||
raise StopIteration
|
||||
else:
|
||||
self.received_finish_reason = "stop"
|
||||
chunk_size = 30
|
||||
new_chunk = self.completion_stream[:chunk_size] # type: ignore[index]
|
||||
completion_obj["content"] = new_chunk
|
||||
self.completion_stream = self.completion_stream[chunk_size:] # type: ignore[index]
|
||||
elif self.custom_llm_provider == "triton":
|
||||
response_obj = self.handle_triton_stream(chunk)
|
||||
completion_obj["content"] = response_obj["text"]
|
||||
print_verbose(f"completion obj content: {completion_obj['content']}")
|
||||
if response_obj["is_finished"]:
|
||||
self.received_finish_reason = response_obj["finish_reason"]
|
||||
elif self.custom_llm_provider == "text-completion-openai":
|
||||
response_obj = self.handle_openai_text_completion_chunk(chunk)
|
||||
completion_obj["content"] = response_obj["text"]
|
||||
print_verbose(f"completion obj content: {completion_obj['content']}")
|
||||
if response_obj["is_finished"]:
|
||||
self.received_finish_reason = response_obj["finish_reason"]
|
||||
if response_obj["usage"] is not None:
|
||||
setattr(
|
||||
model_response,
|
||||
"usage",
|
||||
litellm.Usage(
|
||||
prompt_tokens=response_obj["usage"].prompt_tokens,
|
||||
completion_tokens=response_obj["usage"].completion_tokens,
|
||||
total_tokens=response_obj["usage"].total_tokens,
|
||||
),
|
||||
)
|
||||
elif self.custom_llm_provider == "text-completion-codestral":
|
||||
if not isinstance(chunk, str):
|
||||
raise ValueError(f"chunk is not a string: {chunk}")
|
||||
response_obj = cast(
|
||||
Dict[str, Any],
|
||||
litellm.CodestralTextCompletionConfig()._chunk_parser(chunk),
|
||||
)
|
||||
completion_obj["content"] = response_obj["text"]
|
||||
print_verbose(f"completion obj content: {completion_obj['content']}")
|
||||
if response_obj["is_finished"]:
|
||||
self.received_finish_reason = response_obj["finish_reason"]
|
||||
if "usage" in response_obj is not None:
|
||||
setattr(
|
||||
model_response,
|
||||
"usage",
|
||||
litellm.Usage(
|
||||
prompt_tokens=response_obj["usage"].prompt_tokens,
|
||||
completion_tokens=response_obj["usage"].completion_tokens,
|
||||
total_tokens=response_obj["usage"].total_tokens,
|
||||
),
|
||||
)
|
||||
elif self.custom_llm_provider == "azure_text":
|
||||
response_obj = self.handle_azure_text_completion_chunk(chunk)
|
||||
completion_obj["content"] = response_obj["text"]
|
||||
print_verbose(f"completion obj content: {completion_obj['content']}")
|
||||
if response_obj["is_finished"]:
|
||||
self.received_finish_reason = response_obj["finish_reason"]
|
||||
elif self.custom_llm_provider == "cached_response":
|
||||
chunk = cast(ModelResponseStream, chunk)
|
||||
response_obj = {
|
||||
"text": chunk.choices[0].delta.content,
|
||||
"is_finished": True,
|
||||
"finish_reason": chunk.choices[0].finish_reason,
|
||||
"original_chunk": chunk,
|
||||
"tool_calls": (
|
||||
chunk.choices[0].delta.tool_calls
|
||||
if hasattr(chunk.choices[0].delta, "tool_calls")
|
||||
else None
|
||||
),
|
||||
}
|
||||
|
||||
completion_obj["content"] = response_obj["text"]
|
||||
if response_obj["tool_calls"] is not None:
|
||||
completion_obj["tool_calls"] = response_obj["tool_calls"]
|
||||
print_verbose(f"completion obj content: {completion_obj['content']}")
|
||||
if hasattr(chunk, "id"):
|
||||
model_response.id = chunk.id
|
||||
self.response_id = chunk.id
|
||||
if hasattr(chunk, "system_fingerprint"):
|
||||
self.system_fingerprint = chunk.system_fingerprint
|
||||
if response_obj["is_finished"]:
|
||||
self.received_finish_reason = response_obj["finish_reason"]
|
||||
else: # openai / azure chat model
|
||||
if self.custom_llm_provider in [
|
||||
LlmProviders.AZURE.value,
|
||||
LlmProviders.AZURE_AI.value,
|
||||
]:
|
||||
if isinstance(chunk, BaseModel) and hasattr(chunk, "model"):
|
||||
# for azure, we need to pass the model from the original chunk
|
||||
self.model = getattr(chunk, "model", self.model)
|
||||
response_obj = self.handle_openai_chat_completion_chunk(chunk)
|
||||
if response_obj is None:
|
||||
return
|
||||
completion_obj["content"] = response_obj["text"]
|
||||
self.intermittent_finish_reason = response_obj.get(
|
||||
"finish_reason", None
|
||||
)
|
||||
if response_obj["is_finished"]:
|
||||
if response_obj["finish_reason"] == "error":
|
||||
raise Exception(
|
||||
"{} raised a streaming error - finish_reason: error, no content string given. Received Chunk={}".format(
|
||||
self.custom_llm_provider, response_obj
|
||||
)
|
||||
)
|
||||
self.received_finish_reason = response_obj["finish_reason"]
|
||||
if response_obj.get("original_chunk", None) is not None:
|
||||
if hasattr(response_obj["original_chunk"], "id"):
|
||||
model_response = self.set_model_id(
|
||||
response_obj["original_chunk"].id, model_response
|
||||
)
|
||||
if hasattr(response_obj["original_chunk"], "system_fingerprint"):
|
||||
model_response.system_fingerprint = response_obj[
|
||||
"original_chunk"
|
||||
].system_fingerprint
|
||||
self.system_fingerprint = response_obj[
|
||||
"original_chunk"
|
||||
].system_fingerprint
|
||||
if response_obj["logprobs"] is not None:
|
||||
model_response.choices[0].logprobs = response_obj["logprobs"]
|
||||
|
||||
if response_obj["usage"] is not None:
|
||||
if isinstance(response_obj["usage"], dict):
|
||||
setattr(
|
||||
model_response,
|
||||
"usage",
|
||||
litellm.Usage(
|
||||
prompt_tokens=response_obj["usage"].get(
|
||||
"prompt_tokens", None
|
||||
)
|
||||
or None,
|
||||
completion_tokens=response_obj["usage"].get(
|
||||
"completion_tokens", None
|
||||
)
|
||||
or None,
|
||||
total_tokens=response_obj["usage"].get(
|
||||
"total_tokens", None
|
||||
)
|
||||
or None,
|
||||
),
|
||||
)
|
||||
elif isinstance(response_obj["usage"], Usage):
|
||||
setattr(
|
||||
model_response,
|
||||
"usage",
|
||||
response_obj["usage"],
|
||||
)
|
||||
elif isinstance(response_obj["usage"], BaseModel):
|
||||
setattr(
|
||||
model_response,
|
||||
"usage",
|
||||
litellm.Usage(**response_obj["usage"].model_dump()),
|
||||
)
|
||||
completion_obj: dict[str, Any] = {"content": ""}
|
||||
dispatch_result = self._dispatch_provider_chunk(
|
||||
chunk=chunk,
|
||||
model_response=model_response,
|
||||
completion_obj=completion_obj,
|
||||
)
|
||||
if isinstance(dispatch_result, _ProviderChunkEarlyReturn):
|
||||
return dispatch_result.value
|
||||
response_obj = dispatch_result.response_obj
|
||||
|
||||
model_response.model = self.model
|
||||
## FUNCTION CALL PARSING
|
||||
|
|
@ -1980,11 +2005,29 @@ class CustomStreamWrapper:
|
|||
|
||||
except StopIteration:
|
||||
if self.sent_last_chunk is True:
|
||||
complete_streaming_response = litellm.stream_chunk_builder(
|
||||
chunks=self.chunks,
|
||||
messages=self.messages,
|
||||
logging_obj=self.logging_obj,
|
||||
)
|
||||
try:
|
||||
complete_streaming_response = litellm.stream_chunk_builder(
|
||||
chunks=self.chunks,
|
||||
messages=self.messages,
|
||||
logging_obj=self.logging_obj,
|
||||
)
|
||||
except Exception as e:
|
||||
# stream_chunk_builder can re-raise (as APIError) on large agentic
|
||||
# streams. The raise originates inside this except-StopIteration block,
|
||||
# so the sibling `except Exception` below does not catch it; it would
|
||||
# escape __next__ and drop the request from SpendLogs. Recover
|
||||
# best-effort usage from the raw chunks so cost is still tracked
|
||||
verbose_logger.warning(
|
||||
"stream_chunk_builder raised at end-of-stream (%s); logging "
|
||||
"best-effort usage from chunks.",
|
||||
str(e),
|
||||
)
|
||||
try:
|
||||
complete_streaming_response = self.model_response_creator(
|
||||
chunk={"usage": calculate_total_usage(chunks=self.chunks)}
|
||||
)
|
||||
except Exception:
|
||||
complete_streaming_response = None
|
||||
|
||||
response = self.model_response_creator()
|
||||
if complete_streaming_response is not None:
|
||||
|
|
@ -2209,11 +2252,27 @@ class CustomStreamWrapper:
|
|||
except (StopAsyncIteration, StopIteration):
|
||||
if self.sent_last_chunk is True:
|
||||
# log the final chunk with accurate streaming values
|
||||
complete_streaming_response = litellm.stream_chunk_builder(
|
||||
chunks=self.chunks,
|
||||
messages=self.messages,
|
||||
logging_obj=self.logging_obj,
|
||||
)
|
||||
try:
|
||||
complete_streaming_response = litellm.stream_chunk_builder(
|
||||
chunks=self.chunks,
|
||||
messages=self.messages,
|
||||
logging_obj=self.logging_obj,
|
||||
)
|
||||
except Exception as e:
|
||||
# see sync __next__: a raise from stream_chunk_builder inside this
|
||||
# except handler escapes __anext__ and drops the request from SpendLogs.
|
||||
# Recover best-effort usage from the raw chunks so cost is still tracked
|
||||
verbose_logger.warning(
|
||||
"stream_chunk_builder raised at end-of-stream (%s); logging "
|
||||
"best-effort usage from chunks.",
|
||||
str(e),
|
||||
)
|
||||
try:
|
||||
complete_streaming_response = self.model_response_creator(
|
||||
chunk={"usage": calculate_total_usage(chunks=self.chunks)}
|
||||
)
|
||||
except Exception:
|
||||
complete_streaming_response = None
|
||||
|
||||
response = self.model_response_creator()
|
||||
if complete_streaming_response is not None:
|
||||
|
|
|
|||
|
|
@ -229,6 +229,11 @@ DROP_UNSUPPORTED_OUTPUT_CONFIG_WARNING = (
|
|||
"Sonnet 4.6+, and Mythos Preview."
|
||||
)
|
||||
|
||||
DROP_UNSUPPORTED_SPEED_WARNING = (
|
||||
"Dropping unsupported `speed` for model=%s "
|
||||
"(drop_params=True). Fast mode is only supported on select Opus models."
|
||||
)
|
||||
|
||||
|
||||
class AnthropicConfig(AnthropicModelInfo, BaseConfig):
|
||||
"""
|
||||
|
|
@ -374,6 +379,51 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
|
|||
for level in ("low", "minimal", "medium", "high", "xhigh", "max")
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _model_supports_speed_param(
|
||||
model: str, custom_llm_provider: Optional[str] = None
|
||||
) -> bool:
|
||||
"""Whether the model accepts Anthropic's ``speed`` parameter (fast mode).
|
||||
|
||||
Fast mode is direct Anthropic API-only (not Bedrock, Vertex, or Azure).
|
||||
Those providers strip their prefix before this shared transform runs, so a
|
||||
bare ``claude-opus-4-8`` would otherwise resolve to the direct-API entry;
|
||||
the routed provider is checked explicitly to keep them out.
|
||||
"""
|
||||
if custom_llm_provider is not None and custom_llm_provider != "anthropic":
|
||||
return False
|
||||
return (
|
||||
AnthropicModelInfo._get_exact_model_capability(model, "supports_speed")
|
||||
is True
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _maybe_drop_speed_param(
|
||||
model: str,
|
||||
optional_params: dict,
|
||||
drop_params: bool,
|
||||
custom_llm_provider: Optional[str] = None,
|
||||
) -> None:
|
||||
if "speed" not in optional_params:
|
||||
return
|
||||
if AnthropicConfig._model_supports_speed_param(model, custom_llm_provider):
|
||||
return
|
||||
if not (litellm.drop_params or drop_params):
|
||||
speed_value = optional_params.get("speed")
|
||||
raise litellm.utils.UnsupportedParamsError(
|
||||
message=(
|
||||
f"{model} does not support speed={speed_value!r}. "
|
||||
"To drop unsupported params, set "
|
||||
"`litellm.drop_params = True`."
|
||||
),
|
||||
status_code=400,
|
||||
)
|
||||
litellm.verbose_logger.warning(
|
||||
DROP_UNSUPPORTED_SPEED_WARNING,
|
||||
model,
|
||||
)
|
||||
optional_params.pop("speed", None)
|
||||
|
||||
@staticmethod
|
||||
def _raise_invalid_reasoning_effort(
|
||||
model: str, value: Any, llm_provider: str
|
||||
|
|
@ -1569,8 +1619,13 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
|
|||
anthropic_context_management
|
||||
)
|
||||
elif param == "speed" and isinstance(value, str):
|
||||
# Pass through Anthropic-specific speed parameter for fast mode
|
||||
optional_params["speed"] = value
|
||||
AnthropicConfig._maybe_drop_speed_param(
|
||||
model=model,
|
||||
optional_params=optional_params,
|
||||
drop_params=drop_params,
|
||||
custom_llm_provider=self.custom_llm_provider,
|
||||
)
|
||||
elif param == "cache_control" and isinstance(value, dict):
|
||||
# Pass through top-level cache_control for automatic prompt caching
|
||||
optional_params["cache_control"] = value
|
||||
|
|
@ -1875,6 +1930,14 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
|
|||
"has no thinking_blocks. The model won't use extended thinking for this turn."
|
||||
)
|
||||
|
||||
AnthropicConfig._maybe_drop_speed_param(
|
||||
model=model,
|
||||
optional_params=optional_params,
|
||||
drop_params=litellm.drop_params
|
||||
or litellm_params.get("drop_params") is True,
|
||||
custom_llm_provider=self.custom_llm_provider,
|
||||
)
|
||||
|
||||
headers = self.update_headers_with_optional_anthropic_beta(
|
||||
headers=headers, optional_params=optional_params
|
||||
)
|
||||
|
|
|
|||
|
|
@ -367,6 +367,16 @@ class AnthropicModelInfo(BaseLLMModelInfo):
|
|||
pass
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def _get_exact_model_capability(model: str, key: str) -> Optional[bool]:
|
||||
"""Read boolean capability ``key`` from the exact model-map entry only.
|
||||
|
||||
Unlike ``_get_model_capability``, does not walk stripped provider aliases.
|
||||
Use when a feature is tied to a specific host (e.g. Anthropic API fast mode).
|
||||
"""
|
||||
value = litellm.model_cost.get(model, {}).get(key)
|
||||
return value if isinstance(value, bool) else None
|
||||
|
||||
@staticmethod
|
||||
def _supports_model_capability(model: str, key: str) -> bool:
|
||||
"""Check a boolean capability ``key`` in the model map.
|
||||
|
|
|
|||
|
|
@ -507,7 +507,10 @@ def anthropic_messages_handler(
|
|||
local_vars.update(kwargs)
|
||||
anthropic_messages_optional_request_params = (
|
||||
AnthropicMessagesRequestUtils.get_requested_anthropic_messages_optional_param(
|
||||
params=local_vars
|
||||
params=local_vars,
|
||||
model=model,
|
||||
drop_params=litellm_params.get("drop_params") is True,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
)
|
||||
)
|
||||
if is_reasoning_auto_summary_enabled():
|
||||
|
|
|
|||
|
|
@ -84,8 +84,14 @@ class AdvisorOrchestrationHandler(MessagesInterceptor):
|
|||
)
|
||||
# Optional routing overrides for the advisor sub-call (e.g. proxy routing).
|
||||
# If not set in the tool definition, litellm resolves from env vars.
|
||||
advisor_api_key: Optional[str] = advisor_tool.get("api_key")
|
||||
advisor_api_base: Optional[str] = advisor_tool.get("api_base")
|
||||
# The advisor tool is caller-controlled; only honor a client-supplied
|
||||
# api_base/api_key when the proxy has enabled clientside credentials,
|
||||
# otherwise let litellm resolve from server config.
|
||||
advisor_api_key: Optional[str] = None
|
||||
advisor_api_base: Optional[str] = None
|
||||
if _allow_client_side_advisor_credentials():
|
||||
advisor_api_key = advisor_tool.get("api_key")
|
||||
advisor_api_base = advisor_tool.get("api_base")
|
||||
|
||||
# Build the synthetic tool definition the provider will receive.
|
||||
synthetic_advisor_tool = _make_synthetic_advisor_tool()
|
||||
|
|
@ -181,6 +187,20 @@ class AdvisorOrchestrationHandler(MessagesInterceptor):
|
|||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _allow_client_side_advisor_credentials() -> bool:
|
||||
"""Whether a caller-supplied advisor api_base/api_key may be honored.
|
||||
|
||||
Gated on the proxy's ``allow_client_side_credentials`` opt-in. When the
|
||||
interceptor runs outside the proxy (SDK use), there is no admin boundary
|
||||
to protect, so client-supplied routing is allowed.
|
||||
"""
|
||||
try:
|
||||
from litellm.proxy.proxy_server import general_settings
|
||||
except (ImportError, ModuleNotFoundError):
|
||||
return True
|
||||
return general_settings.get("allow_client_side_credentials") is True
|
||||
|
||||
|
||||
def _make_synthetic_advisor_tool() -> Dict:
|
||||
"""Build a regular tool definition the executor provider can understand."""
|
||||
return {
|
||||
|
|
|
|||
|
|
@ -23,12 +23,19 @@ class AnthropicMessagesRequestUtils:
|
|||
@staticmethod
|
||||
def get_requested_anthropic_messages_optional_param(
|
||||
params: Dict[str, Any],
|
||||
*,
|
||||
model: str | None = None,
|
||||
drop_params: bool = False,
|
||||
custom_llm_provider: str | None = None,
|
||||
) -> AnthropicMessagesRequestOptionalParams:
|
||||
"""
|
||||
Filter parameters to only include those defined in AnthropicMessagesRequestOptionalParams.
|
||||
|
||||
Args:
|
||||
params: Dictionary of parameters to filter
|
||||
model: Resolved model id; when set, unsupported params may be dropped
|
||||
drop_params: Per-request drop_params flag (also respects litellm.drop_params)
|
||||
custom_llm_provider: Routed provider; fast mode is gated to direct Anthropic
|
||||
|
||||
Returns:
|
||||
AnthropicMessagesRequestOptionalParams instance with only the valid parameters
|
||||
|
|
@ -37,6 +44,15 @@ class AnthropicMessagesRequestUtils:
|
|||
filtered_params = {
|
||||
k: v for k, v in params.items() if k in valid_keys and v is not None
|
||||
}
|
||||
if model is not None:
|
||||
from litellm.llms.anthropic.chat.transformation import AnthropicConfig
|
||||
|
||||
AnthropicConfig._maybe_drop_speed_param(
|
||||
model=model,
|
||||
optional_params=filtered_params,
|
||||
drop_params=drop_params,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
)
|
||||
return cast(AnthropicMessagesRequestOptionalParams, filtered_params)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -53,7 +53,13 @@ class APISerpentSearchConfig(BaseSearchConfig):
|
|||
api_base: Optional[str] = None,
|
||||
**kwargs,
|
||||
) -> Dict:
|
||||
api_key = api_key or get_secret_str("APISERPENT_API_KEY")
|
||||
api_key = self.resolve_server_api_key(
|
||||
caller_api_key=api_key,
|
||||
caller_api_base=api_base,
|
||||
key_env_vars=("APISERPENT_API_KEY",),
|
||||
base_env_var="APISERPENT_API_BASE",
|
||||
default_api_base=APISERPENT_BASE,
|
||||
)
|
||||
if not api_key:
|
||||
raise ValueError(
|
||||
"APISERPENT_API_KEY is not set. Set `APISERPENT_API_KEY` environment variable."
|
||||
|
|
|
|||
0
litellm/llms/base_llm/sandbox/__init__.py
Normal file
0
litellm/llms/base_llm/sandbox/__init__.py
Normal file
96
litellm/llms/base_llm/sandbox/transformation.py
Normal file
96
litellm/llms/base_llm/sandbox/transformation.py
Normal file
|
|
@ -0,0 +1,96 @@
|
|||
"""
|
||||
Base Sandbox transformation configuration.
|
||||
|
||||
A sandbox provider runs an executable string inside an isolated container and
|
||||
returns whatever the sandbox produced. The lifecycle is create container ->
|
||||
run code -> delete container; `code_interpreter_tool` combines all three.
|
||||
"""
|
||||
|
||||
from typing import Any, Union
|
||||
|
||||
import httpx
|
||||
|
||||
from pydantic import Field, PrivateAttr
|
||||
|
||||
from litellm.types.llms.base import LiteLLMPydanticObjectBase
|
||||
|
||||
SANDBOX_MAX_OUTPUT_BYTES = 10 * 1024 * 1024
|
||||
|
||||
|
||||
class ContainerHandle(LiteLLMPydanticObjectBase):
|
||||
"""A live sandbox container. Carries everything needed to reach it again."""
|
||||
|
||||
id: str
|
||||
provider: str
|
||||
domain: str | None = None
|
||||
|
||||
model_config = {"extra": "allow"}
|
||||
|
||||
_hidden_params: dict = PrivateAttr(default_factory=dict)
|
||||
|
||||
|
||||
class CodeExecutionResult(LiteLLMPydanticObjectBase):
|
||||
"""Passthrough of the sandbox's own execution output."""
|
||||
|
||||
stdout: str = ""
|
||||
stderr: str = ""
|
||||
results: list[dict[str, Any]] = Field(default_factory=list)
|
||||
error: dict[str, Any] | None = None
|
||||
execution_count: int | None = None
|
||||
object: str = "code_execution"
|
||||
|
||||
model_config = {"extra": "allow"}
|
||||
|
||||
_hidden_params: dict = PrivateAttr(default_factory=dict)
|
||||
|
||||
|
||||
class BaseSandboxConfig:
|
||||
"""Provider-agnostic sandbox operations."""
|
||||
|
||||
def validate_environment(self, api_key: str | None = None, **kwargs) -> str:
|
||||
raise NotImplementedError(
|
||||
"validate_environment must be implemented by provider"
|
||||
)
|
||||
|
||||
async def acreate_sandbox(
|
||||
self,
|
||||
*,
|
||||
template: str | None = None,
|
||||
timeout: int | None = None,
|
||||
allow_internet_access: bool | None = None,
|
||||
api_key: str | None = None,
|
||||
**kwargs,
|
||||
) -> ContainerHandle:
|
||||
raise NotImplementedError("acreate_sandbox must be implemented by provider")
|
||||
|
||||
async def arun_code(
|
||||
self,
|
||||
*,
|
||||
container: Union[ContainerHandle, str],
|
||||
code: str,
|
||||
api_key: str | None = None,
|
||||
**kwargs,
|
||||
) -> CodeExecutionResult:
|
||||
raise NotImplementedError("arun_code must be implemented by provider")
|
||||
|
||||
async def adelete_sandbox(
|
||||
self,
|
||||
*,
|
||||
container: Union[ContainerHandle, str],
|
||||
api_key: str | None = None,
|
||||
**kwargs,
|
||||
) -> bool:
|
||||
raise NotImplementedError("adelete_sandbox must be implemented by provider")
|
||||
|
||||
async def _read_capped_lines(self, response: httpx.Response) -> list[str]:
|
||||
lines: list[str] = []
|
||||
total = 0
|
||||
async for line in response.aiter_lines():
|
||||
total += len(line.encode("utf-8"))
|
||||
if total > SANDBOX_MAX_OUTPUT_BYTES:
|
||||
raise ValueError(
|
||||
f"Sandbox output exceeded {SANDBOX_MAX_OUTPUT_BYTES} bytes; aborting "
|
||||
"to avoid unbounded memory use."
|
||||
)
|
||||
lines.append(line)
|
||||
return lines
|
||||
|
|
@ -3,11 +3,13 @@ Base Search transformation configuration.
|
|||
"""
|
||||
|
||||
from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Union
|
||||
from urllib.parse import urlsplit
|
||||
|
||||
import httpx
|
||||
from pydantic import PrivateAttr
|
||||
|
||||
from litellm.llms.base_llm.chat.transformation import BaseLLMException
|
||||
from litellm.secret_managers.main import get_secret_str
|
||||
from litellm.types.llms.base import LiteLLMPydanticObjectBase
|
||||
|
||||
if TYPE_CHECKING:
|
||||
|
|
@ -16,6 +18,29 @@ else:
|
|||
LiteLLMLoggingObj = Any
|
||||
|
||||
|
||||
def _search_host(url: str) -> str:
|
||||
return urlsplit(url).netloc.lower()
|
||||
|
||||
|
||||
def _is_trusted_search_api_base(
|
||||
caller_api_base: str,
|
||||
default_api_base: str | None,
|
||||
base_env_var: str | None,
|
||||
) -> bool:
|
||||
candidate = _search_host(caller_api_base)
|
||||
if not candidate:
|
||||
return False
|
||||
trusted = {
|
||||
_search_host(base)
|
||||
for base in (
|
||||
default_api_base,
|
||||
get_secret_str(base_env_var) if base_env_var else None,
|
||||
)
|
||||
if base
|
||||
}
|
||||
return candidate in trusted
|
||||
|
||||
|
||||
class SearchResult(LiteLLMPydanticObjectBase):
|
||||
"""Single search result."""
|
||||
|
||||
|
|
@ -86,6 +111,60 @@ class BaseSearchConfig:
|
|||
"max_tokens_per_page",
|
||||
}
|
||||
|
||||
def _assert_trusted_api_base_for_server_credential(
|
||||
self,
|
||||
caller_api_base: str | None,
|
||||
default_api_base: str | None,
|
||||
base_env_var: str | None,
|
||||
credential_name: str,
|
||||
) -> None:
|
||||
"""
|
||||
Block sending a server-managed credential to a caller-chosen host.
|
||||
|
||||
A caller-supplied api_base is honored when constructing the request URL, so
|
||||
falling back to a server-configured secret while the caller controls the host
|
||||
leaks that secret. The provider default and the operator's own api_base
|
||||
override are the only trusted destinations for a server-managed credential.
|
||||
"""
|
||||
if not caller_api_base:
|
||||
return
|
||||
if _is_trusted_search_api_base(caller_api_base, default_api_base, base_env_var):
|
||||
return
|
||||
raise ValueError(
|
||||
f"Refusing to send the server-configured {credential_name} to the "
|
||||
f"caller-supplied api_base '{caller_api_base}'. Pass an explicit api_key "
|
||||
f"when overriding api_base for this search provider."
|
||||
)
|
||||
|
||||
def resolve_server_api_key(
|
||||
self,
|
||||
*,
|
||||
caller_api_key: str | None,
|
||||
caller_api_base: str | None,
|
||||
key_env_vars: tuple[str, ...],
|
||||
base_env_var: str | None,
|
||||
default_api_base: str | None,
|
||||
) -> str | None:
|
||||
"""
|
||||
Resolve a single-secret search API key, falling back to a server-managed
|
||||
secret only when the request targets a trusted host.
|
||||
|
||||
Returns the caller's key when provided, otherwise the first set
|
||||
server-managed secret (or None when none is set, for keyless providers).
|
||||
"""
|
||||
if caller_api_key:
|
||||
return caller_api_key
|
||||
server_key = next(
|
||||
(key for key in (get_secret_str(var) for var in key_env_vars) if key),
|
||||
None,
|
||||
)
|
||||
if server_key is None:
|
||||
return None
|
||||
self._assert_trusted_api_base_for_server_credential(
|
||||
caller_api_base, default_api_base, base_env_var, key_env_vars[0]
|
||||
)
|
||||
return server_key
|
||||
|
||||
def validate_environment(
|
||||
self,
|
||||
headers: Dict,
|
||||
|
|
|
|||
Some files were not shown because too many files have changed in this diff Show more
Loading…
Add table
Reference in a new issue