Merge pull request #34200 from BerriAI/litellm_internal_staging
Some checks failed
CodeQL / Analyze (actions) (push) Has been cancelled
CodeQL / Analyze (javascript-typescript) (push) Has been cancelled
CodeQL / Analyze (python) (push) Has been cancelled
CodSpeed Benchmarks / benchmarks (push) Has been cancelled
Helm unit test / unit-test (push) Has been cancelled
Scorecard supply-chain security / Scorecard analysis (push) Has been cancelled
LiteLLM Rust / rustfmt, clippy, test (push) Has been cancelled
GitHub Actions Security Analysis / zizmor (push) Has been cancelled

chore(ci): promote internal staging to main
This commit is contained in:
yuneng-jiang 2026-07-21 19:03:23 -07:00 • committed by GitHub
commit f2479cc704
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
606 changed files with 44613 additions and 8558 deletions

View file

@ -30,7 +30,7 @@ body:
id: steps-to-reproduce
attributes:
label: Steps to Reproduce
description: Please provide detailed steps to reproduce this bug(A curl/python code to reproduce the bug)
description: Please provide a numbered list of the exact steps to reproduce this bug (include a curl/python snippet to reproduce it). Number each step (1., 2., 3., ...) in the order you performed them.
placeholder: |
1. config.yaml file/ .env file/ etc.
2. Run the following code...

View file

@ -1,3 +1,18 @@
## TLDR
<!-- Fill in the bullets below and keep each one short and concrete: one line per bullet, roughly 10 words max
This section must be extremely human parsable, comprehensible, and readable: its target audience is humans, not AI agents -->
Problem this solves:
- <blah>
- ...
How it solves it:
- <blah>
- ...
## Relevant issues
<!-- e.g., "Fixes #000" -->

View file

@ -1,61 +0,0 @@
name: Create Daily OSS Branch
on:
schedule:
- cron: "0 16 * * 1-5" # 9am PT during daylight saving time, weekdays.
workflow_dispatch:
inputs:
date:
description: "Branch date in YYYY_MM_DD format. Defaults to today's UTC date."
required: false
type: string
permissions:
contents: write
jobs:
create-oss-branch:
if: github.repository == 'BerriAI/litellm'
runs-on: ubuntu-latest
timeout-minutes: 10
steps:
- name: Checkout repository
uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
fetch-depth: 0
persist-credentials: false
- name: Create dated OSS branch
env:
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
REQUESTED_DATE: ${{ inputs.date }}
run: |
set -euo pipefail
if [ -n "${REQUESTED_DATE}" ]; then
if ! echo "${REQUESTED_DATE}" | grep -Eq '^[0-9]{4}_[0-9]{2}_[0-9]{2}$'; then
echo "::error::date must use YYYY_MM_DD format, got '${REQUESTED_DATE}'"
exit 1
fi
BRANCH_DATE="${REQUESTED_DATE}"
else
BRANCH_DATE="$(date -u +'%Y_%m_%d')"
fi
BRANCH_NAME="litellm_oss_daily_${BRANCH_DATE}"
echo "Creating branch: ${BRANCH_NAME}"
git config user.name "github-actions[bot]"
git config user.email "github-actions[bot]@users.noreply.github.com"
git fetch origin main "${BRANCH_NAME}" || true
if git show-ref --verify --quiet "refs/remotes/origin/${BRANCH_NAME}"; then
echo "Branch ${BRANCH_NAME} already exists. Skipping creation."
exit 0
fi
git checkout -b "${BRANCH_NAME}" origin/main
git push "https://x-access-token:${GITHUB_TOKEN}@github.com/${GITHUB_REPOSITORY}.git" "${BRANCH_NAME}"
echo "Successfully created and pushed branch: ${BRANCH_NAME}"

View file

@ -31,12 +31,12 @@ jobs:
echo "PR head repo: $HEAD_REPO"
echo "PR head branch: $HEAD_REF"
if [ "$HEAD_REPO" != "$BASE_REPO" ]; then
echo "::error::PRs to main must originate from the canonical repository ($BASE_REPO), not a fork ($HEAD_REPO). External contributors should open PRs against the current daily OSS branch (named litellm_oss_daily_YYYY_MM_DD; a fresh one is cut each weekday, so target the most recent) instead."
echo "::error::PRs to main must originate from the canonical repository ($BASE_REPO), not a fork ($HEAD_REPO). External contributors should open PRs against 'litellm_internal_staging' instead."
exit 1
fi
if [ "$HEAD_REF" = "litellm_internal_staging" ] || [[ "$HEAD_REF" == litellm_hotfix_?* ]]; then
echo "Allowed source branch."
exit 0
fi
echo "::error::PRs to main must originate from 'litellm_internal_staging' or a 'litellm_hotfix_*' branch. Got: '$HEAD_REF'. If this is a contribution, retarget the PR against the current daily OSS branch (named litellm_oss_daily_YYYY_MM_DD; a fresh one is cut each weekday, so target the most recent) instead."
echo "::error::PRs to main must originate from 'litellm_internal_staging' or a 'litellm_hotfix_*' branch. Got: '$HEAD_REF'. If this is a contribution, retarget the PR against 'litellm_internal_staging' instead."
exit 1

View file

@ -58,6 +58,8 @@ jobs:
# free OSS, run as a pinned, checksum-verified binary; no GitHub Action
# dependency and no vendor SaaS callout.
- name: Scan image for fixable HIGH/CRITICAL CVEs
env:
GRYPE_MATCH_PYTHON_USING_CPES: "true"
run: |
"$RUNNER_TEMP/grype" litellm-image-scan:${{ github.sha }} \
--only-fixed \

View file

@ -1,50 +0,0 @@
name: OSS Daily Guardrails
on:
push:
branches:
- "litellm_oss_daily_20*"
pull_request:
branches:
- "litellm_oss_daily_20*"
- litellm_internal_staging
permissions:
contents: read
concurrency:
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
cancel-in-progress: true
jobs:
oss-safe-checks:
name: Run OSS daily safe checks
if: startsWith(github.ref_name, 'litellm_oss_daily_20') || startsWith(github.head_ref, 'litellm_oss_daily_20') || startsWith(github.base_ref, 'litellm_oss_daily_20')
runs-on: ubuntu-latest
timeout-minutes: 10
steps:
- name: Checkout repository
uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
persist-credentials: false
- name: Set up Python
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
with:
python-version: "3.12"
- name: Set up uv
uses: ./.github/actions/setup-uv-with-retries
with:
version: "0.10.9"
- name: Run secret scan test
run: |
uv run --frozen --with 'pytest==9.0.2' pytest tests/litellm/test_no_hardcoded_secrets.py -v
- name: Run Ruff
run: |
uv sync --frozen
cd litellm
uv run --no-sync ruff check .

View file

@ -115,6 +115,9 @@ jobs:
- name: check_fastuuid_usage
run: uv run --no-sync python ./tests/code_coverage_tests/check_fastuuid_usage.py
- name: check_e2e_no_raw_requests
run: uv run --no-sync python ./tests/code_coverage_tests/check_e2e_no_raw_requests.py
- name: memory_test
run: uv run --no-sync python ./tests/code_coverage_tests/memory_test.py

View file

@ -0,0 +1,57 @@
name: UI Unit Tests
permissions:
contents: read
on:
pull_request:
branches:
- main
- litellm_internal_staging
- litellm_oss_staging
- "litellm_**"
push:
branches:
- litellm_internal_staging
concurrency:
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
cancel-in-progress: true
jobs:
ui-unit-tests:
runs-on: ubuntu-latest
timeout-minutes: 20
defaults:
run:
working-directory: ui/litellm-dashboard
steps:
- name: Checkout repository
uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
fetch-depth: 0
persist-credentials: false
- name: Setup Node.js
uses: actions/setup-node@a0853c24544627f65ddf259abe73b1d18a591444 # v5.0.0
with:
node-version: "20"
cache: "npm"
cache-dependency-path: ui/litellm-dashboard/package-lock.json
- name: Install dependencies
run: npm ci
- name: Run UI unit tests (Vitest)
env:
CI: "true"
BASE_SHA: ${{ github.event.pull_request.base.sha }}
run: |
if [ -n "$BASE_SHA" ]; then
echo "Pull request: running only tests related to changes since $BASE_SHA"
npm run test -- --run --changed "$BASE_SHA" --passWithNoTests \
--pool forks --poolOptions.forks.maxForks=4
else
echo "Push to $GITHUB_REF_NAME: running the full suite"
npm run test -- --run --pool forks --poolOptions.forks.maxForks=4
fi

View file

@ -61,5 +61,11 @@ jobs:
- name: Run Clippy
run: cargo clippy --workspace --all-targets --locked -- -D warnings
- name: Run Clippy with Bedrock auth
run: cargo clippy -p litellm-core --all-targets --features bedrock-auth --locked -- -D warnings
- name: Run Rust tests
run: cargo test --workspace --locked
- name: Run core tests with Bedrock auth
run: cargo test -p litellm-core --features bedrock-auth --locked

View file

@ -19,7 +19,7 @@ Same thing for bug fixes. The tests should make it so that this specific bug can
End-to-end tests belong in `tests/e2e/` and must follow the harness conventions documented in that directory's `CLAUDE.md`
When creating PRs, don't set base to `main`. `litellm_internal_staging` serves that purpose for internal contributors; external / OSS contributions target the current daily OSS branch instead, named `litellm_oss_daily_YYYY_MM_DD` (a fresh one is cut each weekday, so use the most recent)
When creating PRs, don't set base to `main`. `litellm_internal_staging` is the default base branch and serves that purpose for both internal and external / OSS contributions
When writing a PR body, treat the comments and imperative instructions inside @.github/pull_request_template.md as rules to follow, not just layout. Agent harnesses may strip HTML comments from copies of that file injected into context, so read .github/pull_request_template.md from disk before writing a PR body to make sure you see every comment rule

View file

@ -322,7 +322,7 @@ npm run build
## Submitting Your PR
1. **Push your branch**: `git push origin your-feature-branch`
2. **Create a PR**: Go to GitHub and open a pull request against the current daily OSS branch, named `litellm_oss_daily_YYYY_MM_DD`. A fresh one is cut each weekday, so pick the most recent from the [branch list](https://github.com/BerriAI/litellm/branches/all?query=litellm_oss_daily). Do not target `main`.
2. **Create a PR**: Go to GitHub and open a pull request against [`litellm_internal_staging`](https://github.com/BerriAI/litellm/tree/litellm_internal_staging), which is the default base branch. Do not target `main`.
3. **Fill out the PR template**: Provide clear description of changes
4. **Wait for review**: Maintainers will review and provide feedback
5. **Address feedback**: Make requested changes and push updates

View file

@ -0,0 +1,17 @@
-- AlterTable
ALTER TABLE "LiteLLM_DailyUserSpend" ADD COLUMN IF NOT EXISTS "compression_saved_tokens" BIGINT NOT NULL DEFAULT 0;
-- AlterTable
ALTER TABLE "LiteLLM_DailyOrganizationSpend" ADD COLUMN IF NOT EXISTS "compression_saved_tokens" BIGINT NOT NULL DEFAULT 0;
-- AlterTable
ALTER TABLE "LiteLLM_DailyEndUserSpend" ADD COLUMN IF NOT EXISTS "compression_saved_tokens" BIGINT NOT NULL DEFAULT 0;
-- AlterTable
ALTER TABLE "LiteLLM_DailyAgentSpend" ADD COLUMN IF NOT EXISTS "compression_saved_tokens" BIGINT NOT NULL DEFAULT 0;
-- AlterTable
ALTER TABLE "LiteLLM_DailyTeamSpend" ADD COLUMN IF NOT EXISTS "compression_saved_tokens" BIGINT NOT NULL DEFAULT 0;
-- AlterTable
ALTER TABLE "LiteLLM_DailyTagSpend" ADD COLUMN IF NOT EXISTS "compression_saved_tokens" BIGINT NOT NULL DEFAULT 0;

View file

@ -0,0 +1,23 @@
-- AlterTable
ALTER TABLE "LiteLLM_DailyUserSpend" ADD COLUMN IF NOT EXISTS "compression_savings_spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0;
ALTER TABLE "LiteLLM_DailyUserSpend" ADD COLUMN IF NOT EXISTS "prompt_caching_savings_spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0;
-- AlterTable
ALTER TABLE "LiteLLM_DailyOrganizationSpend" ADD COLUMN IF NOT EXISTS "compression_savings_spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0;
ALTER TABLE "LiteLLM_DailyOrganizationSpend" ADD COLUMN IF NOT EXISTS "prompt_caching_savings_spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0;
-- AlterTable
ALTER TABLE "LiteLLM_DailyEndUserSpend" ADD COLUMN IF NOT EXISTS "compression_savings_spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0;
ALTER TABLE "LiteLLM_DailyEndUserSpend" ADD COLUMN IF NOT EXISTS "prompt_caching_savings_spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0;
-- AlterTable
ALTER TABLE "LiteLLM_DailyAgentSpend" ADD COLUMN IF NOT EXISTS "compression_savings_spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0;
ALTER TABLE "LiteLLM_DailyAgentSpend" ADD COLUMN IF NOT EXISTS "prompt_caching_savings_spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0;
-- AlterTable
ALTER TABLE "LiteLLM_DailyTeamSpend" ADD COLUMN IF NOT EXISTS "compression_savings_spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0;
ALTER TABLE "LiteLLM_DailyTeamSpend" ADD COLUMN IF NOT EXISTS "prompt_caching_savings_spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0;
-- AlterTable
ALTER TABLE "LiteLLM_DailyTagSpend" ADD COLUMN IF NOT EXISTS "compression_savings_spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0;
ALTER TABLE "LiteLLM_DailyTagSpend" ADD COLUMN IF NOT EXISTS "prompt_caching_savings_spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0;

View file

@ -0,0 +1,9 @@
-- CreateTable
CREATE TABLE IF NOT EXISTS "LiteLLM_SSOIdentityAssertion" (
"user_id" TEXT NOT NULL,
"assertion_b64" TEXT NOT NULL,
"created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
"updated_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
CONSTRAINT "LiteLLM_SSOIdentityAssertion_pkey" PRIMARY KEY ("user_id")
);

View file

@ -403,6 +403,15 @@ model LiteLLM_MCPServerOAuthClient {
updated_at DateTime @default(now()) @updatedAt @map("updated_at")
}
// The enterprise IdP identity assertion captured at SSO login, one row per user.
// assertion_b64 is an encrypted JSON payload: {id_token, refresh_token?, issuer?, expires_at?}.
model LiteLLM_SSOIdentityAssertion {
user_id String @id
assertion_b64 String
created_at DateTime @default(now()) @map("created_at")
updated_at DateTime @default(now()) @updatedAt @map("updated_at")
}
// Generate Tokens for Proxy
model LiteLLM_VerificationToken {
token String @id
@ -736,6 +745,9 @@ model LiteLLM_DailyUserSpend {
completion_tokens BigInt @default(0)
cache_read_input_tokens BigInt @default(0)
cache_creation_input_tokens BigInt @default(0)
compression_saved_tokens BigInt @default(0)
compression_savings_spend Float @default(0.0)
prompt_caching_savings_spend Float @default(0.0)
spend Float @default(0.0)
api_requests BigInt @default(0)
successful_requests BigInt @default(0)
@ -767,6 +779,9 @@ model LiteLLM_DailyOrganizationSpend {
completion_tokens BigInt @default(0)
cache_read_input_tokens BigInt @default(0)
cache_creation_input_tokens BigInt @default(0)
compression_saved_tokens BigInt @default(0)
compression_savings_spend Float @default(0.0)
prompt_caching_savings_spend Float @default(0.0)
spend Float @default(0.0)
api_requests BigInt @default(0)
successful_requests BigInt @default(0)
@ -798,6 +813,9 @@ model LiteLLM_DailyEndUserSpend {
completion_tokens BigInt @default(0)
cache_read_input_tokens BigInt @default(0)
cache_creation_input_tokens BigInt @default(0)
compression_saved_tokens BigInt @default(0)
compression_savings_spend Float @default(0.0)
prompt_caching_savings_spend Float @default(0.0)
spend Float @default(0.0)
api_requests BigInt @default(0)
successful_requests BigInt @default(0)
@ -828,6 +846,9 @@ model LiteLLM_DailyAgentSpend {
completion_tokens BigInt @default(0)
cache_read_input_tokens BigInt @default(0)
cache_creation_input_tokens BigInt @default(0)
compression_saved_tokens BigInt @default(0)
compression_savings_spend Float @default(0.0)
prompt_caching_savings_spend Float @default(0.0)
spend Float @default(0.0)
api_requests BigInt @default(0)
successful_requests BigInt @default(0)
@ -858,6 +879,9 @@ model LiteLLM_DailyTeamSpend {
completion_tokens BigInt @default(0)
cache_read_input_tokens BigInt @default(0)
cache_creation_input_tokens BigInt @default(0)
compression_saved_tokens BigInt @default(0)
compression_savings_spend Float @default(0.0)
prompt_caching_savings_spend Float @default(0.0)
spend Float @default(0.0)
api_requests BigInt @default(0)
successful_requests BigInt @default(0)
@ -890,6 +914,9 @@ model LiteLLM_DailyTagSpend {
completion_tokens BigInt @default(0)
cache_read_input_tokens BigInt @default(0)
cache_creation_input_tokens BigInt @default(0)
compression_saved_tokens BigInt @default(0)
compression_savings_spend Float @default(0.0)
prompt_caching_savings_spend Float @default(0.0)
spend Float @default(0.0)
api_requests BigInt @default(0)
successful_requests BigInt @default(0)

View file

@ -1,6 +1,6 @@
[project]
name = "litellm-proxy-extras"
version = "0.4.79"
version = "0.4.80"
description = "Additional files for the LiteLLM Proxy. Reduces the size of the main litellm package."
readme = "README.md"
requires-python = ">=3.9"
@ -26,7 +26,7 @@ required-version = ">=0.10.9"
module-root = ""
[tool.commitizen]
version = "0.4.79"
version = "0.4.80"
version_files = [
"pyproject.toml:^version",
"../pyproject.toml:litellm-proxy-extras==",

View file

@ -39,6 +39,11 @@ Route-level Rust structure mirrors LiteLLM's Python responsibilities:
- Network execution lives in the host crate `ai-gateway` (`ai-gateway/src/io/`),
never inside `core`.
Call-hook and lifecycle instrumentation, including phase timing, usage
accumulation, and callback payload construction, always lives in `core`.
Hosts feed observed events into core and dispatch the completed payloads through
their I/O logger; hosts must not own callback orchestration.
Allowed in `core`:
- Pure request transforms
- Pure response transforms
@ -57,6 +62,13 @@ Not allowed in `core`:
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.
A new provider/route may instead be implemented rust-only with no Python
reference; then the Python interface is a thin dispatch that calls Rust with no
fallback, and you state the rust-only choice explicitly in the PR. Either way
the Python side stays minimal (it only marshals inputs and calls the Rust
interface), never add a per-route feature flag, and never push provider
dispatch into `litellm/main.py`; put it in a thin dispatch class under
`litellm/llms/<provider>/<route>/`.
## Production Bar

1299
litellm-rust/Cargo.lock generated

File diff suppressed because it is too large Load diff

View file

@ -7,7 +7,8 @@ members = [
resolver = "2"
[workspace.package]
edition = "2021"
edition = "2024"
rust-version = "1.88"
license = "MIT"
repository = "https://github.com/BerriAI/litellm"

View file

@ -39,11 +39,17 @@ Rules for adding or changing an LLM provider/route in `litellm-rust`. OCR (`MIST
19. Every provider transform ships tests for: supported-param filtering, request body shape, response normalization, missing/null fields, bad input, and `*_match_python` fixture parity.
20. Lifecycle/hook tests cover hook order, success + failure callback payloads, pre-call guardrail blocking before any provider I/O, during-call body mutation, and provider-error mapping.
21. Rust paths stay off by default and behind Python parity tests (disabled / enabled-equals-Python / bridge-unavailable fallback) until parity is proven.
21. When a route has a Python reference implementation, the Rust path stays off by default and behind Python parity tests (disabled / enabled-equals-Python / bridge-unavailable fallback) until parity is proven. A new provider/route may instead be implemented rust-only with no Python reference; then the Python interface is a thin dispatch to Rust with no fallback, and tests cover the rust-backed path plus the unavailable-bridge error. State the rust-only choice explicitly in the PR.
## Python bridge (SDK side)
22. A Python -> Rust bridge keeps the Python side minimal: the Python interface only marshals inputs and calls the Rust interface, with no transform, handler, or business logic. Aim for well under 100 lines of interface code per route; if the Python grows past that, the logic belongs in Rust.
23. Do not bloat `litellm/main.py`. A route's provider dispatch lives in a thin dispatch class under `litellm/llms/<provider>/<route>/` that calls the Rust bridge; `main.py` only instantiates it and calls its sync/async method.
24. Do not add new feature flags unless explicitly requested. Reuse the existing litellm rust rollout mechanism (`use_litellm_rust`); never introduce a per-route env flag such as `LITELLM_USE_RUST_<ROUTE>`.
## Checks before push
22. Run, and keep green:
25. Run, and keep green:
```bash
cd litellm-rust
cargo fmt --check

View file

@ -14,7 +14,7 @@ path = "src/main.rs"
required-features = ["server"]
[dependencies]
litellm-core.workspace = true
litellm-core = { workspace = true, features = ["bedrock-auth"] }
# reqwest (rustls + json) is used by io/ocr and ships realtime logs to the
# Python proxy callbacks API.
reqwest.workspace = true
@ -41,3 +41,4 @@ python-config = ["dep:pyo3"]
[dev-dependencies]
futures-channel = "0.3"
tower = { version = "0.5.3", features = ["util"] }

View file

@ -0,0 +1,48 @@
use std::collections::BTreeMap;
use litellm_core::CoreResult;
use litellm_core::audio_transcription::transformation::AudioTranscriptionProviderConfig;
use litellm_core::error::CoreError;
use litellm_core::providers::bedrock::audio_transcription::BEDROCK_AUDIO_TRANSCRIPTION_CONFIG;
use serde_json::{Map, Value};
pub(super) fn audio_transcription_provider_config(
provider: &str,
) -> Option<&'static dyn AudioTranscriptionProviderConfig> {
match provider {
"bedrock" => Some(&BEDROCK_AUDIO_TRANSCRIPTION_CONFIG),
_ => None,
}
}
pub(super) fn string_headers(
headers: Option<Map<String, Value>>,
) -> CoreResult<BTreeMap<String, String>> {
headers
.unwrap_or_default()
.into_iter()
.map(|(key, value)| {
value
.as_str()
.map(|value| (key.clone(), value.to_string()))
.ok_or_else(|| {
CoreError::InvalidRequest(format!(
"audio transcription extra_headers.{key} must be a string"
))
})
})
.collect()
}
pub(super) fn has_header(headers: &BTreeMap<String, String>, name: &str) -> bool {
headers.keys().any(|key| key.eq_ignore_ascii_case(name))
}
pub(super) fn truncate_error_body(body: &str) -> String {
let truncated: String = body.chars().take(256).collect();
if truncated.chars().count() == body.chars().count() {
truncated
} else {
format!("{truncated}... (truncated)")
}
}

View file

@ -0,0 +1,89 @@
use std::time::SystemTime;
use litellm_core::CoreResult;
use litellm_core::audio_transcription::transformation::AudioTranscriptionAuth;
use litellm_core::error::CoreError;
use litellm_core::providers::bedrock::audio_transcription::aws_auth_config;
use litellm_core::providers::bedrock::aws_base::{resolve_credentials, sign_bedrock_post};
use serde_json::Value;
use super::common_utils::truncate_error_body;
use super::types::ProviderAudioTranscriptionRequest;
use crate::client::http_client;
pub(crate) async fn execute_audio_transcription_provider_call(
request: ProviderAudioTranscriptionRequest,
) -> CoreResult<Value> {
let body = serde_json::to_vec(&request.body).map_err(|error| {
CoreError::InvalidRequest(format!("invalid audio request body: {error}"))
})?;
let mut request_builder = http_client().post(&request.url).body(body.clone());
for (key, value) in &request.upstream_headers {
request_builder = request_builder.header(key, value);
}
if let Some(duration) = request.timeout {
request_builder = request_builder.timeout(duration);
}
let response = request_builder
.send()
.await
.map_err(|error| CoreError::Network(error.to_string()))?;
let status = response.status();
let text = response
.text()
.await
.map_err(|error| CoreError::Network(error.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(|error| {
CoreError::InvalidResponse(format!("invalid audio response JSON: {error}"))
})?;
Ok(request
.config
.transform_transcription_response(&request.model, response_json)?
.into_json())
}
pub(crate) async fn sign_request(
request: &ProviderAudioTranscriptionRequest,
optional_params: &serde_json::Map<String, Value>,
) -> CoreResult<ProviderAudioTranscriptionRequest> {
let env_lookup = environment_lookup;
let auth = request
.config
.auth_strategy(&request.model, optional_params, &env_lookup)?;
let body = serde_json::to_vec(&request.body).map_err(|error| {
CoreError::InvalidRequest(format!("invalid audio request body: {error}"))
})?;
let mut headers = super::common_utils::string_headers(None)?;
headers.insert("Content-Type".to_string(), "application/json".to_string());
headers.extend(request.upstream_headers.iter().cloned());
match auth {
AudioTranscriptionAuth::Bearer => {}
AudioTranscriptionAuth::AwsSigV4 { region, .. } => {
let credentials =
resolve_credentials(aws_auth_config(optional_params, &env_lookup), &env_lookup)
.await?;
headers.extend(sign_bedrock_post(
&request.url,
&body,
&headers,
&region,
&credentials,
SystemTime::now(),
)?);
}
}
Ok(ProviderAudioTranscriptionRequest {
upstream_headers: headers.into_iter().collect(),
..request.clone()
})
}
pub(super) fn environment_lookup(key: &str) -> Option<String> {
std::env::var(key).ok()
}

View file

@ -0,0 +1,300 @@
use std::future::Future;
use std::pin::Pin;
use litellm_core::CoreResult;
use litellm_core::audio_transcription::transformation::AudioTranscriptionAuth;
use litellm_core::call_lifecycle::{CallLifecycleContext, CallLifecycleHooks, CallLifecycleTiming};
use litellm_core::error::CoreError;
use serde_json::{Map, Value, json};
use super::common_utils::{audio_transcription_provider_config, has_header, string_headers};
use super::handler::sign_request;
use super::types::{PreparedAudioTranscriptionRequest, ProviderAudioTranscriptionRequest};
use crate::integrations::custom_guardrail::{
CustomGuardrailRunner, GuardrailContext, GuardrailError, GuardrailRequest,
};
use crate::integrations::custom_logger::{
CallType, CallbackTiming, CallbackValue, CustomLoggerRunner, LoggingError, ModelCallDetails,
};
use crate::integrations::types::{
RequestMetadata, StandardLoggingMetadata, StandardLoggingPayload,
};
pub(crate) struct AudioTranscriptionLifecycleHooks {
logger_runner: CustomLoggerRunner,
guardrail_runner: CustomGuardrailRunner,
request_metadata: RequestMetadata,
}
type AudioFuture<'a, T> = Pin<Box<dyn Future<Output = CoreResult<T>> + Send + 'a>>;
type AudioLogFuture<'a> = Pin<Box<dyn Future<Output = ()> + Send + 'a>>;
impl AudioTranscriptionLifecycleHooks {
pub(crate) fn new(
logger_runner: CustomLoggerRunner,
guardrail_runner: CustomGuardrailRunner,
request_metadata: RequestMetadata,
) -> Self {
Self {
logger_runner,
guardrail_runner,
request_metadata,
}
}
async fn run_pre_call_guardrails(
&self,
request: PreparedAudioTranscriptionRequest,
) -> CoreResult<PreparedAudioTranscriptionRequest> {
if self.guardrail_runner.is_empty() {
return Ok(request);
}
let (guardrail_request, _) = self
.guardrail_runner
.run_pre_call(
&guardrail_context(&self.request_metadata),
GuardrailRequest::new(json!({
"model": request.model,
"custom_llm_provider": request.custom_llm_provider,
"audio": request.audio,
"optional_params": request.optional_params,
})),
)
.await
.map_err(guardrail_error_to_core_error)?;
let Value::Object(mut data) = guardrail_request.data else {
return Err(CoreError::InvalidRequest(
"audio transcription pre_call guardrail must return an object".to_string(),
));
};
let audio = data.remove("audio").ok_or_else(|| {
CoreError::InvalidRequest("audio transcription guardrail removed audio".to_string())
})?;
let optional_params = match data.remove("optional_params") {
Some(Value::Object(value)) => value,
Some(_) => {
return Err(CoreError::InvalidRequest(
"audio transcription optional_params must be an object".to_string(),
));
}
None => Map::new(),
};
Ok(PreparedAudioTranscriptionRequest {
audio,
optional_params,
..request
})
}
async fn prepare_provider_request(
&self,
request: PreparedAudioTranscriptionRequest,
) -> CoreResult<ProviderAudioTranscriptionRequest> {
let config = audio_transcription_provider_config(&request.custom_llm_provider)
.ok_or_else(|| CoreError::InvalidProvider(request.custom_llm_provider.clone()))?;
let env_lookup = super::handler::environment_lookup;
let headers = string_headers(request.extra_headers)?;
let url = config.complete_url(
request.api_base.as_deref(),
&request.model,
&request.optional_params,
&env_lookup,
)?;
let filtered_params = config.map_transcription_params(&request.optional_params);
let body = config.transform_transcription_request(
&request.model,
request.audio,
filtered_params,
)?;
let auth = config.auth_strategy(&request.model, &request.optional_params, &env_lookup)?;
let mut upstream_headers = headers.into_iter().collect::<Vec<_>>();
if matches!(auth, AudioTranscriptionAuth::Bearer)
&& !has_header(
&upstream_headers
.iter()
.cloned()
.collect::<std::collections::BTreeMap<_, _>>(),
"authorization",
)
&& let Some(api_key) = request.api_key.as_deref()
{
upstream_headers.push(("Authorization".to_string(), format!("Bearer {api_key}")));
}
let provider_request = ProviderAudioTranscriptionRequest {
model: request.model,
config,
url,
body: body.body,
upstream_headers,
timeout: request.timeout,
};
let provider_request = self.run_during_call_guardrails(provider_request).await?;
sign_request(&provider_request, &request.optional_params).await
}
async fn run_during_call_guardrails(
&self,
request: ProviderAudioTranscriptionRequest,
) -> CoreResult<ProviderAudioTranscriptionRequest> {
if self.guardrail_runner.is_empty() {
return Ok(request);
}
let (guardrail_request, _) = self
.guardrail_runner
.run_during_call(
&guardrail_context(&self.request_metadata),
GuardrailRequest::new(json!({
"model": request.model,
"custom_llm_provider": "bedrock",
"url": request.url,
"body": request.body,
})),
)
.await
.map_err(guardrail_error_to_core_error)?;
let Value::Object(mut data) = guardrail_request.data else {
return Err(CoreError::InvalidRequest(
"audio transcription during_call guardrail must return an object".to_string(),
));
};
let body = data.remove("body").ok_or_else(|| {
CoreError::InvalidRequest("audio transcription guardrail removed body".to_string())
})?;
Ok(ProviderAudioTranscriptionRequest { body, ..request })
}
fn logging_payload(
&self,
context: &CallLifecycleContext,
timing: &CallLifecycleTiming,
) -> StandardLoggingPayload {
StandardLoggingPayload {
id: context.litellm_call_id.clone(),
litellm_call_id: context.litellm_call_id.clone(),
call_type: context.call_type.clone(),
model: context.model.clone(),
custom_llm_provider: context.custom_llm_provider.clone(),
response_cost: 0.0,
prompt_tokens: 0,
completion_tokens: 0,
total_tokens: 0,
start_time: timing.start_time,
end_time: timing.end_time,
stream: false,
metadata: StandardLoggingMetadata {
user_api_key_hash: self.request_metadata.user_api_key_hash.clone(),
user_api_key_user_id: self.request_metadata.user_api_key_user_id.clone(),
user_api_key_team_id: self.request_metadata.user_api_key_team_id.clone(),
..Default::default()
},
messages: None,
}
}
}
impl CallLifecycleHooks<PreparedAudioTranscriptionRequest, ProviderAudioTranscriptionRequest, Value>
for AudioTranscriptionLifecycleHooks
{
type PreCallFuture<'a> = AudioFuture<'a, PreparedAudioTranscriptionRequest>;
type DuringCallFuture<'a> = AudioFuture<'a, ProviderAudioTranscriptionRequest>;
type SuccessFuture<'a> = AudioLogFuture<'a>;
type FailureFuture<'a> = AudioLogFuture<'a>;
fn async_pre_call_hook<'a>(
&'a self,
_context: &'a CallLifecycleContext,
request: PreparedAudioTranscriptionRequest,
) -> Self::PreCallFuture<'a> {
Box::pin(async move { self.run_pre_call_guardrails(request).await })
}
fn async_during_call_hook<'a>(
&'a self,
_context: &'a CallLifecycleContext,
request: PreparedAudioTranscriptionRequest,
) -> Self::DuringCallFuture<'a> {
Box::pin(async move { self.prepare_provider_request(request).await })
}
fn async_log_success_event<'a>(
&'a self,
context: &'a CallLifecycleContext,
response: &'a Value,
timing: &'a CallLifecycleTiming,
) -> Self::SuccessFuture<'a> {
Box::pin(async move {
if self.logger_runner.is_empty() {
return;
}
self.logger_runner
.async_log_success_event(
&ModelCallDetails::from_standard_logging_payload(
self.logging_payload(context, timing),
),
&CallbackValue::new("audio_transcription", response.clone()),
CallbackTiming::new(timing.start_time, timing.end_time),
)
.await;
})
}
fn async_log_failure_event<'a>(
&'a self,
context: &'a CallLifecycleContext,
error: &'a CoreError,
timing: &'a CallLifecycleTiming,
) -> Self::FailureFuture<'a> {
Box::pin(async move {
if self.logger_runner.is_empty() {
return;
}
let logging_error = LoggingError {
message: error.to_string(),
kind: core_error_kind(error).to_string(),
};
self.logger_runner
.async_log_failure_event(
&ModelCallDetails::from_standard_logging_payload(
self.logging_payload(context, timing),
)
.with_failure_error(logging_error.clone()),
Some(&CallbackValue::new(
"error",
json!({"message": logging_error.message, "kind": logging_error.kind}),
)),
CallbackTiming::new(timing.start_time, timing.end_time),
)
.await;
})
}
}
fn guardrail_context(metadata: &RequestMetadata) -> GuardrailContext {
GuardrailContext {
call_type: CallType::Other("audio_transcription".to_string()),
selected_guardrails: Vec::new(),
metadata: std::collections::HashMap::new(),
user_api_key_hash: metadata.user_api_key_hash.clone(),
user_api_key_user_id: metadata.user_api_key_user_id.clone(),
user_api_key_team_id: metadata.user_api_key_team_id.clone(),
trace_parent: None,
}
}
fn guardrail_error_to_core_error(error: GuardrailError) -> CoreError {
CoreError::InvalidRequest(format!("{}: {}", error.kind, error.message))
}
fn core_error_kind(error: &CoreError) -> &'static str {
match error {
CoreError::Auth(_) => "AuthError",
CoreError::InvalidProvider(_) => "InvalidProvider",
CoreError::InvalidRequest(_) => "InvalidRequest",
CoreError::InvalidType { .. } => "InvalidType",
CoreError::MissingField(_) => "MissingField",
CoreError::Http { .. } => "HttpError",
CoreError::InvalidResponse(_) => "InvalidResponse",
CoreError::Network(_) => "NetworkError",
CoreError::Routing(_) => "RoutingError",
}
}

View file

@ -0,0 +1,25 @@
use litellm_core::CoreResult;
use litellm_core::call_lifecycle::CallLifecycle;
use serde_json::Value;
mod common_utils;
mod handler;
mod hooks;
mod prepare;
mod types;
pub use types::AudioTranscriptionRequest;
use handler::execute_audio_transcription_provider_call;
use prepare::{PreparedAudioTranscriptionCall, prepare_audio_transcription_call};
pub async fn audio_transcription(request: AudioTranscriptionRequest<'_>) -> CoreResult<Value> {
let PreparedAudioTranscriptionCall { request, hooks } =
prepare_audio_transcription_call(request);
CallLifecycle::default()
.run_request(request, &hooks, execute_audio_transcription_provider_call)
.await
}
#[cfg(test)]
mod tests;

View file

@ -0,0 +1,55 @@
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::{SystemTime, UNIX_EPOCH};
use litellm_core::routing_utils::provider::{CustomLlmProvider, get_custom_llm_provider};
use super::hooks::AudioTranscriptionLifecycleHooks;
use super::types::{AudioTranscriptionRequest, PreparedAudioTranscriptionRequest};
use crate::integrations::custom_guardrail::CustomGuardrailRunner;
use crate::integrations::custom_logger::CustomLoggerRunner;
pub(crate) struct PreparedAudioTranscriptionCall {
pub(crate) request: PreparedAudioTranscriptionRequest,
pub(crate) hooks: AudioTranscriptionLifecycleHooks,
}
pub(crate) fn prepare_audio_transcription_call(
request: AudioTranscriptionRequest<'_>,
) -> PreparedAudioTranscriptionCall {
let call_id = request
.litellm_call_id
.map(str::to_string)
.unwrap_or_else(new_audio_transcription_call_id);
let provider_info = get_custom_llm_provider(request.model, request.custom_llm_provider)
.unwrap_or(CustomLlmProvider {
model: request.model,
custom_llm_provider: "bedrock",
});
PreparedAudioTranscriptionCall {
request: PreparedAudioTranscriptionRequest {
model: provider_info.model.to_string(),
custom_llm_provider: provider_info.custom_llm_provider.to_string(),
litellm_call_id: call_id,
audio: request.audio,
api_key: request.api_key.map(str::to_string),
api_base: request.api_base.map(str::to_string),
extra_headers: request.extra_headers,
optional_params: request.optional_params,
timeout: request.timeout,
},
hooks: AudioTranscriptionLifecycleHooks::new(
CustomLoggerRunner::new(request.callbacks),
CustomGuardrailRunner::new(request.guardrails),
request.request_metadata,
),
}
}
fn new_audio_transcription_call_id() -> String {
static COUNTER: AtomicU64 = AtomicU64::new(1);
let sequence = COUNTER.fetch_add(1, Ordering::Relaxed);
let timestamp = SystemTime::now()
.duration_since(UNIX_EPOCH)
.map_or(0, |duration| duration.as_nanos());
format!("audio-transcription-{timestamp}-{sequence}")
}

View file

@ -0,0 +1,53 @@
use std::io::{Read, Write};
use std::net::TcpListener;
use std::thread;
use serde_json::{Map, json};
use super::{AudioTranscriptionRequest, audio_transcription};
#[tokio::test]
async fn bedrock_request_is_signed_and_contains_audio() {
let listener = TcpListener::bind("127.0.0.1:0").expect("listener");
let address = listener.local_addr().expect("address");
let server = thread::spawn(move || {
let (mut stream, _) = listener.accept().expect("connection");
let mut request = Vec::new();
let mut buffer = [0_u8; 16_384];
let count = stream.read(&mut buffer).expect("request");
request.extend_from_slice(&buffer[..count]);
let request = String::from_utf8_lossy(&request);
assert!(request.contains("POST /model/mistral.voxtral-mini-3b-2507/converse"));
assert!(request.contains("authorization: AWS4-HMAC-SHA256"));
assert!(request.contains("x-amz-date:"));
assert!(request.contains("\"bytes\":\"AQI=\""));
assert!(request.contains("Transcribe the audio. Respond with only the transcript."));
let response = b"HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: 53\r\nConnection: close\r\n\r\n{\"output\":{\"message\":{\"content\":[{\"text\":\"hello\"}]}}}";
stream.write_all(response).expect("response");
});
let optional_params = Map::from_iter([
("aws_access_key_id".to_string(), json!("access-key")),
("aws_secret_access_key".to_string(), json!("secret-key")),
("aws_region_name".to_string(), json!("us-east-1")),
]);
let api_base = format!("http://{address}");
let response = audio_transcription(AudioTranscriptionRequest {
model: "mistral.voxtral-mini-3b-2507",
audio: json!({"data": "AQI=", "format": "wav", "filename": "audio.wav"}),
api_key: None,
api_base: Some(&api_base),
custom_llm_provider: Some("bedrock"),
extra_headers: None,
optional_params,
timeout: None,
callbacks: Vec::new(),
guardrails: Vec::new(),
request_metadata: Default::default(),
litellm_call_id: None,
})
.await
.expect("transcription");
assert_eq!(response, json!({"text": "hello"}));
server.join().expect("server");
}

View file

@ -0,0 +1,58 @@
use std::sync::Arc;
use std::time::Duration;
use litellm_core::audio_transcription::transformation::AudioTranscriptionProviderConfig;
use litellm_core::call_lifecycle::{CallLifecycleContext, CallLifecycleRequest};
use serde_json::{Map, Value};
use crate::integrations::custom_guardrail::CustomGuardrail;
use crate::integrations::custom_logger::CustomLogger;
use crate::integrations::types::RequestMetadata;
pub struct AudioTranscriptionRequest<'a> {
pub model: &'a str,
pub audio: Value,
pub api_key: Option<&'a str>,
pub api_base: Option<&'a str>,
pub custom_llm_provider: Option<&'a str>,
pub extra_headers: Option<Map<String, Value>>,
pub optional_params: Map<String, Value>,
pub timeout: Option<Duration>,
pub callbacks: Vec<Arc<dyn CustomLogger>>,
pub guardrails: Vec<Arc<dyn CustomGuardrail>>,
pub request_metadata: RequestMetadata,
pub litellm_call_id: Option<&'a str>,
}
pub(crate) struct PreparedAudioTranscriptionRequest {
pub(crate) model: String,
pub(crate) custom_llm_provider: String,
pub(crate) litellm_call_id: String,
pub(crate) audio: Value,
pub(crate) api_key: Option<String>,
pub(crate) api_base: Option<String>,
pub(crate) extra_headers: Option<Map<String, Value>>,
pub(crate) optional_params: Map<String, Value>,
pub(crate) timeout: Option<Duration>,
}
impl CallLifecycleRequest for PreparedAudioTranscriptionRequest {
fn lifecycle_context(&self) -> CallLifecycleContext {
CallLifecycleContext::new(
"audio_transcription",
self.model.clone(),
self.custom_llm_provider.clone(),
self.litellm_call_id.clone(),
)
}
}
#[derive(Clone)]
pub(crate) struct ProviderAudioTranscriptionRequest {
pub(crate) model: String,
pub(crate) config: &'static dyn AudioTranscriptionProviderConfig,
pub(crate) url: String,
pub(crate) body: Value,
pub(crate) upstream_headers: Vec<(String, String)>,
pub(crate) timeout: Option<Duration>,
}

View file

@ -9,9 +9,9 @@
//! runs during extraction, before the handler body. Routes never re-implement it.
use axum::extract::FromRequestParts;
use axum::http::StatusCode;
use axum::http::header::AUTHORIZATION;
use axum::http::request::Parts;
use axum::http::StatusCode;
use sha2::{Digest, Sha256};
use subtle::ConstantTimeEq;

View file

@ -1,13 +1,13 @@
use std::sync::OnceLock;
use std::time::Duration;
const OCR_TIMEOUT_SECS: u64 = 600;
const HTTP_CLIENT_TIMEOUT_SECS: u64 = 600;
pub(super) fn http_client() -> &'static reqwest::Client {
pub(crate) fn http_client() -> &'static reqwest::Client {
static CLIENT: OnceLock<reqwest::Client> = OnceLock::new();
CLIENT.get_or_init(|| {
reqwest::Client::builder()
.timeout(Duration::from_secs(OCR_TIMEOUT_SECS))
.timeout(Duration::from_secs(HTTP_CLIENT_TIMEOUT_SECS))
.build()
.expect("failed to build reqwest client")
})

View file

@ -40,3 +40,19 @@ pub(crate) const MESSAGES_CONNECT_TIMEOUT_SECS: u64 = 10;
/// Max characters of an upstream error body echoed across the host boundary
/// before truncation, so provider bodies are bounded and data-minimized.
pub(crate) const MESSAGES_ERROR_BODY_MAX_CHARS: usize = 256;
pub(crate) const DEFAULT_RESPONSES_WS_CONNECT_TIMEOUT_SECS: u64 = 10;
pub(crate) const DEFAULT_RESPONSES_WS_IDLE_TIMEOUT_SECS: u64 = 300;
/// HTTP path for the non-streaming Anthropic Messages route.
#[cfg(feature = "server")]
pub(crate) const MESSAGES_ROUTE_PATH: &str = "/v1/messages";
/// Provider name used by the Anthropic Messages route when a deployment's
/// provider model does not carry an explicit provider prefix.
pub(crate) const ANTHROPIC_MESSAGES_PROVIDER: &str = "anthropic";
/// Request headers owned by the gateway and never forwarded upstream.
#[cfg(feature = "server")]
pub(crate) const MESSAGES_HEADERS_NOT_FORWARDED: &[&str] =
&["authorization", "connection", "content-length", "host"];

View file

@ -0,0 +1 @@
pub use crate::audio_transcription::{AudioTranscriptionRequest, audio_transcription};

View file

@ -1 +1 @@
pub use crate::messages::{messages, MessagesRequest};
pub use crate::messages::{MessagesRequest, messages};

View file

@ -1,4 +1,6 @@
pub mod audio_transcription;
pub mod messages;
pub mod ocr;
pub mod realtime;
pub mod realtime_pool;
pub mod responses_ws;

View file

@ -1 +1 @@
pub use crate::ocr::{ocr, OcrRequest};
pub use crate::ocr::{OcrRequest, ocr};

View file

@ -15,16 +15,16 @@ use std::time::Duration;
use futures_util::stream::{SplitSink, SplitStream};
use futures_util::{Sink, SinkExt, Stream, StreamExt};
use litellm_core::CoreResult;
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 tokio_tungstenite::tungstenite::client::IntoClientRequest;
use tokio_tungstenite::tungstenite::http::HeaderValue;
use tokio_tungstenite::tungstenite::http::header::AUTHORIZATION;
use tokio_tungstenite::{MaybeTlsStream, WebSocketStream, connect_async};
use litellm_core::providers::openai::realtime::transformation::OPENAI_REALTIME_CONFIG;
@ -113,7 +113,7 @@ pub(crate) async fn read_event(upstream_rx: &mut UpstreamRx) -> CoreResult<Realt
Message::Close(_) => {
return Err(CoreError::Network(
"upstream closed before first event".to_string(),
))
));
}
_ => continue,
}

View file

@ -28,11 +28,11 @@ 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 litellm_core::realtime::types::RealtimeEvent;
use crate::io::realtime::{
dial_upstream, read_event, resolve_api_key, UpstreamRx, UpstreamTx, UpstreamWs,
UpstreamRx, UpstreamTx, UpstreamWs, dial_upstream, read_event, resolve_api_key,
};
/// Default target warm sockets per key when pooling is enabled.
@ -473,8 +473,8 @@ pub fn upstream_key(
/// 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 futures_util::task::noop_waker_ref;
use std::pin::Pin;
use std::task::{Context, Poll};
@ -523,15 +523,15 @@ mod tests {
))
.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;
}
if let Message::Text(text) = msg
&& 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;
}
}
}

View file

@ -0,0 +1,549 @@
use std::collections::HashMap;
use std::sync::Arc;
use std::time::Duration;
use futures_util::stream::{SplitSink, SplitStream};
use futures_util::{Sink, SinkExt, Stream, StreamExt};
use litellm_core::providers::openai::responses::transformation::OPENAI_RESPONSES_WS_CONFIG;
use litellm_core::responses::types::ResponsesWsEvent;
use litellm_core::responses::websocket::ResponsesWebSocketProviderConfig;
use litellm_core::{CoreError, CoreResult};
use tokio::net::TcpStream;
use tokio::sync::Mutex;
use tokio_tungstenite::tungstenite::Message;
use tokio_tungstenite::tungstenite::client::IntoClientRequest;
use tokio_tungstenite::tungstenite::http::HeaderValue;
use tokio_tungstenite::tungstenite::http::header::{AUTHORIZATION, HeaderName};
use tokio_tungstenite::{MaybeTlsStream, WebSocketStream, connect_async};
use crate::constants::{
DEFAULT_RESPONSES_WS_CONNECT_TIMEOUT_SECS, DEFAULT_RESPONSES_WS_IDLE_TIMEOUT_SECS,
};
const OPENAI_API_KEY_ENV: &str = "OPENAI_API_KEY";
const MISSING_KEY_MESSAGE: &str = "Missing OpenAI API Key - a Responses WebSocket call is being made but no key was passed via params or the OPENAI_API_KEY environment variable";
pub type ResponsesUpstreamWs = WebSocketStream<MaybeTlsStream<TcpStream>>;
type UpstreamTx = SplitSink<ResponsesUpstreamWs, Message>;
type UpstreamRx = SplitStream<ResponsesUpstreamWs>;
#[derive(Clone)]
pub struct ResponsesWebSocketConnection {
socket: Arc<Mutex<Option<ResponsesUpstreamWs>>>,
}
impl ResponsesWebSocketConnection {
pub async fn connect_url(
url: &str,
headers: &HashMap<String, String>,
timeout: Option<Duration>,
) -> CoreResult<Self> {
let mut request = url
.into_client_request()
.map_err(|error| CoreError::Network(error.to_string()))?;
for (name, value) in headers {
let header_name = name
.parse::<HeaderName>()
.map_err(|error| CoreError::InvalidRequest(error.to_string()))?;
let header_value = HeaderValue::from_str(value)
.map_err(|error| CoreError::InvalidRequest(error.to_string()))?;
request.headers_mut().insert(header_name, header_value);
}
let connect = connect_async(request);
let result = match timeout {
Some(timeout) => tokio::time::timeout(timeout, connect).await.map_err(|_| {
CoreError::Network("Responses WebSocket connection timed out".to_string())
})?,
None => connect.await,
};
let (socket, _) = result.map_err(|error| match error {
tokio_tungstenite::tungstenite::Error::Http(response) => CoreError::Http {
status: response.status().as_u16(),
body: String::new(),
},
other => CoreError::Network(other.to_string()),
})?;
Ok(Self {
socket: Arc::new(Mutex::new(Some(socket))),
})
}
pub async fn send_text(&self, text: String) -> CoreResult<()> {
let mut socket = self.socket.lock().await;
let Some(socket) = socket.as_mut() else {
return Err(CoreError::Network(
"Responses WebSocket is closed".to_string(),
));
};
socket
.send(Message::Text(text))
.await
.map_err(|error| CoreError::Network(error.to_string()))
}
pub async fn recv_text(&self) -> CoreResult<Option<String>> {
let mut socket_guard = self.socket.lock().await;
let Some(socket) = socket_guard.as_mut() else {
return Ok(None);
};
match socket.next().await {
Some(Ok(Message::Text(text))) => Ok(Some(text)),
Some(Ok(Message::Binary(bytes))) => String::from_utf8(bytes.to_vec())
.map(Some)
.map_err(|error| CoreError::InvalidResponse(error.to_string())),
Some(Ok(Message::Close(_))) | None => Ok(None),
Some(Ok(_)) => Ok(None),
Some(Err(error)) => Err(CoreError::Network(error.to_string())),
}
}
pub async fn close(&self) -> CoreResult<()> {
let mut socket = self.socket.lock().await;
if let Some(socket) = socket.as_mut() {
socket
.close(None)
.await
.map_err(|error| CoreError::Network(error.to_string()))?;
}
*socket = None;
Ok(())
}
}
pub(crate) fn resolve_api_key(api_key: Option<&str>) -> CoreResult<String> {
api_key
.map(str::trim)
.filter(|value| !value.is_empty())
.map(str::to_string)
.or_else(|| {
std::env::var(OPENAI_API_KEY_ENV)
.ok()
.filter(|value| !value.trim().is_empty())
})
.ok_or_else(|| CoreError::Auth(MISSING_KEY_MESSAGE.to_string()))
}
async fn dial_upstream(
model: &str,
api_key: &str,
api_base: Option<&str>,
) -> CoreResult<ResponsesUpstreamWs> {
let url = OPENAI_RESPONSES_WS_CONFIG.complete_websocket_url(api_base, model);
let mut request = url
.as_str()
.into_client_request()
.map_err(|error| CoreError::Network(error.to_string()))?;
request.headers_mut().insert(
AUTHORIZATION,
HeaderValue::from_str(&format!("Bearer {api_key}"))
.map_err(|error| CoreError::Auth(error.to_string()))?,
);
let result = tokio::time::timeout(
Duration::from_secs(DEFAULT_RESPONSES_WS_CONNECT_TIMEOUT_SECS),
connect_async(request),
)
.await
.map_err(|_| CoreError::Network("Responses WebSocket connection timed out".to_string()))?;
result
.map(|(socket, _)| socket)
.map_err(|error| match error {
tokio_tungstenite::tungstenite::Error::Http(response) => CoreError::Http {
status: response.status().as_u16(),
body: String::new(),
},
other => CoreError::Network(other.to_string()),
})
}
pub struct ResponsesWebSocketStreaming;
impl ResponsesWebSocketStreaming {
pub async fn bidirectional_forward<In, Out>(
model: &str,
upstream_tx: UpstreamTx,
upstream_rx: UpstreamRx,
idle_timeout: Option<Duration>,
observe: impl FnMut(&ResponsesWsEvent) + Send,
client_in: In,
client_out: Out,
) -> CoreResult<()>
where
In: Stream<Item = ResponsesWsEvent> + Unpin + Send,
Out: Sink<ResponsesWsEvent> + Unpin + Send,
Out::Error: std::fmt::Display,
{
splice(
model,
upstream_tx,
upstream_rx,
idle_timeout,
observe,
client_in,
client_out,
)
.await
}
}
pub(crate) async fn splice<In, Out>(
model: &str,
mut upstream_tx: UpstreamTx,
mut upstream_rx: UpstreamRx,
idle_timeout: Option<Duration>,
mut observe: impl FnMut(&ResponsesWsEvent) + Send,
mut client_in: In,
mut client_out: Out,
) -> CoreResult<()>
where
In: Stream<Item = ResponsesWsEvent> + Unpin + Send,
Out: Sink<ResponsesWsEvent> + Unpin + Send,
Out::Error: std::fmt::Display,
{
let idle =
idle_timeout.unwrap_or_else(|| Duration::from_secs(DEFAULT_RESPONSES_WS_IDLE_TIMEOUT_SECS));
loop {
tokio::select! {
event = client_in.next() => {
let Some(event) = event else { break };
for outbound in OPENAI_RESPONSES_WS_CONFIG
.transform_ws_request(&event, model)?
.events
{
let payload = serde_json::to_string(&outbound)
.map_err(|error| CoreError::InvalidResponse(error.to_string()))?;
upstream_tx.send(Message::Text(payload))
.await
.map_err(|error| CoreError::Network(error.to_string()))?;
}
}
message = upstream_rx.next() => {
let Some(message) = message else { break };
match message.map_err(|error| CoreError::Network(error.to_string()))? {
Message::Text(text) => {
let event = serde_json::from_str::<ResponsesWsEvent>(&text)
.map_err(|error| CoreError::InvalidResponse(error.to_string()))?;
observe(&event);
for outbound in OPENAI_RESPONSES_WS_CONFIG
.transform_ws_response(&event, model)?
.events
{
client_out.send(outbound)
.await
.map_err(|error| CoreError::Network(error.to_string()))?;
}
}
Message::Close(_) => break,
_ => {}
}
}
_ = tokio::time::sleep(idle) => break,
}
}
Ok(())
}
#[allow(clippy::too_many_arguments)]
pub async fn async_responses_websocket<In, Out>(
model: &str,
api_key: Option<&str>,
api_base: Option<&str>,
first_frame: Option<ResponsesWsEvent>,
idle_timeout: Option<Duration>,
mut observe: impl FnMut(&ResponsesWsEvent) + Send,
client_in: In,
client_out: Out,
) -> CoreResult<()>
where
In: Stream<Item = ResponsesWsEvent> + Unpin + Send,
Out: Sink<ResponsesWsEvent> + Unpin + Send,
Out::Error: std::fmt::Display,
{
let key = resolve_api_key(api_key)?;
let upstream = dial_upstream(model, &key, api_base).await?;
let (mut upstream_tx, upstream_rx) = upstream.split();
if let Some(first_frame) = first_frame {
for outbound in OPENAI_RESPONSES_WS_CONFIG
.transform_ws_request(&first_frame, model)?
.events
{
let payload = serde_json::to_string(&outbound)
.map_err(|error| CoreError::InvalidResponse(error.to_string()))?;
upstream_tx
.send(Message::Text(payload))
.await
.map_err(|error| CoreError::Network(error.to_string()))?;
}
}
ResponsesWebSocketStreaming::bidirectional_forward(
model,
upstream_tx,
upstream_rx,
idle_timeout,
&mut observe,
client_in,
client_out,
)
.await
}
#[allow(clippy::too_many_arguments)]
pub async fn responses_ws<In, Out>(
model: &str,
api_key: Option<&str>,
api_base: Option<&str>,
first_frame: Option<ResponsesWsEvent>,
idle_timeout: Option<Duration>,
observe: impl FnMut(&ResponsesWsEvent) + Send,
client_in: In,
client_out: Out,
) -> CoreResult<()>
where
In: Stream<Item = ResponsesWsEvent> + Unpin + Send,
Out: Sink<ResponsesWsEvent> + Unpin + Send,
Out::Error: std::fmt::Display,
{
async_responses_websocket(
model,
api_key,
api_base,
first_frame,
idle_timeout,
observe,
client_in,
client_out,
)
.await
}
#[cfg(test)]
mod tests {
use super::*;
use futures_channel::mpsc;
use futures_util::{SinkExt, StreamExt};
use litellm_core::responses::types::ResponsesWsEventType;
use serde_json::json;
use tokio::io::AsyncWriteExt;
use tokio::net::TcpListener;
use tokio_tungstenite::accept_async;
async fn websocket_base() -> (String, tokio::task::JoinHandle<()>) {
let listener = TcpListener::bind("127.0.0.1:0").await.expect("bind");
let address = listener.local_addr().expect("local address");
let task = tokio::spawn(async move {
let (stream, _) = listener.accept().await.expect("accept");
let mut socket = accept_async(stream).await.expect("websocket handshake");
while let Some(Ok(Message::Text(text))) = socket.next().await {
let request: serde_json::Value = serde_json::from_str(&text).expect("request json");
let model = request
.get("model")
.and_then(serde_json::Value::as_str)
.or_else(|| {
request
.get("response")
.and_then(serde_json::Value::as_object)
.and_then(|response| {
response.get("model").and_then(serde_json::Value::as_str)
})
})
.expect("enforced model");
socket
.send(Message::Text(
json!({
"type": "response.created",
"response": {
"id": format!("resp-{model}"),
"model": model,
"extra": "preserved"
}
})
.to_string(),
))
.await
.expect("created event");
socket
.send(Message::Text(
json!({
"type": "response.completed",
"response": {
"id": format!("resp-{model}"),
"model": model,
"usage": {
"input_tokens": 1,
"output_tokens": 2,
"total_tokens": 3
}
}
})
.to_string(),
))
.await
.expect("completed event");
}
});
(format!("http://{address}"), task)
}
fn event(value: serde_json::Value) -> ResponsesWsEvent {
serde_json::from_value(value).expect("event")
}
#[test]
fn explicit_nonblank_key_wins() {
assert_eq!(
resolve_api_key(Some(" explicit ")).expect("key"),
"explicit"
);
}
#[test]
fn blank_key_is_not_accepted_without_environment_key() {
if std::env::var(OPENAI_API_KEY_ENV).is_err() {
assert!(resolve_api_key(Some(" ")).is_err());
}
}
#[tokio::test]
async fn forwards_events_sequentially_and_enforces_model() {
let (api_base, server) = websocket_base().await;
let (client_tx, client_rx) = mpsc::unbounded();
let (output_tx, mut output_rx) = mpsc::unbounded();
let (observed_tx, observed_rx) = mpsc::unbounded();
client_tx
.unbounded_send(event(json!({
"type": "response.create",
"model": "wrong"
})))
.expect("first request");
client_tx
.unbounded_send(event(json!({
"type": "response.create",
"response": {"model": "also-wrong"}
})))
.expect("second request");
let task = tokio::spawn(async move {
responses_ws(
"authorized-model",
Some("test-key"),
Some(&api_base),
None,
Some(Duration::from_secs(1)),
move |event| {
observed_tx
.unbounded_send(event.clone())
.expect("observe event");
},
client_rx,
output_tx,
)
.await
});
let first = output_rx.next().await.expect("first output");
let second = output_rx.next().await.expect("second output");
let third = output_rx.next().await.expect("third output");
let fourth = output_rx.next().await.expect("fourth output");
drop(client_tx);
task.await.expect("splice task").expect("successful splice");
server.await.expect("server task");
assert_eq!(first.event_type, ResponsesWsEventType::ResponseCreated);
assert_eq!(first.model(), Some("authorized-model"));
assert_eq!(first.data["response"]["extra"], "preserved");
assert_eq!(second.event_type, ResponsesWsEventType::ResponseCompleted);
assert_eq!(third.event_type, ResponsesWsEventType::ResponseCreated);
assert_eq!(fourth.event_type, ResponsesWsEventType::ResponseCompleted);
let observed: Vec<_> = observed_rx.collect().await;
assert_eq!(observed.len(), 4);
assert!(
observed
.iter()
.all(|event| event.event_type != ResponsesWsEventType::ResponseCreate)
);
}
#[tokio::test]
async fn idle_timeout_ends_without_upstream_events() {
let listener = TcpListener::bind("127.0.0.1:0").await.expect("bind");
let address = listener.local_addr().expect("address");
let server = tokio::spawn(async move {
let (stream, _) = listener.accept().await.expect("accept");
let _socket = accept_async(stream).await.expect("handshake");
tokio::time::sleep(Duration::from_secs(1)).await;
});
let (_client_tx, client_rx) = mpsc::unbounded::<ResponsesWsEvent>();
let (output_tx, mut output_rx) = mpsc::unbounded();
let result = responses_ws(
"model",
Some("key"),
Some(&format!("http://{address}")),
None,
Some(Duration::from_millis(20)),
|_| {},
client_rx,
output_tx,
)
.await;
assert!(result.is_ok());
assert!(output_rx.next().await.is_none());
server.abort();
}
#[tokio::test]
async fn dial_http_status_is_preserved() {
let listener = TcpListener::bind("127.0.0.1:0").await.expect("bind");
let address = listener.local_addr().expect("address");
let server = tokio::spawn(async move {
let (mut stream, _) = listener.accept().await.expect("accept");
stream
.write_all(b"HTTP/1.1 401 Unauthorized\r\nContent-Length: 0\r\n\r\n")
.await
.expect("response");
});
let (_client_tx, client_rx) = mpsc::unbounded::<ResponsesWsEvent>();
let (output_tx, _output_rx) = mpsc::unbounded();
let error = responses_ws(
"model",
Some("key"),
Some(&format!("http://{address}")),
None,
Some(Duration::from_millis(20)),
|_| {},
client_rx,
output_tx,
)
.await
.expect_err("status error");
assert!(matches!(error, CoreError::Http { status: 401, .. }));
server.await.expect("server task");
}
#[tokio::test]
async fn dial_http_500_status_is_preserved() {
let listener = TcpListener::bind("127.0.0.1:0").await.expect("bind");
let address = listener.local_addr().expect("address");
let server = tokio::spawn(async move {
let (mut stream, _) = listener.accept().await.expect("accept");
stream
.write_all(b"HTTP/1.1 500 Internal Server Error\r\nContent-Length: 0\r\n\r\n")
.await
.expect("response");
});
let (_client_tx, client_rx) = mpsc::unbounded::<ResponsesWsEvent>();
let (output_tx, _output_rx) = mpsc::unbounded();
let error = responses_ws(
"model",
Some("key"),
Some(&format!("http://{address}")),
None,
Some(Duration::from_millis(20)),
|_| {},
client_rx,
output_tx,
)
.await
.expect_err("status error");
assert!(matches!(error, CoreError::Http { status: 500, .. }));
server.await.expect("server task");
}
}

View file

@ -11,6 +11,8 @@
//! binary turns on. The `python-config` feature additionally pulls in [`python`]
//! for the load-time config reader.
pub mod audio_transcription;
mod client;
pub mod io;
pub mod messages;
pub mod ocr;
@ -26,9 +28,6 @@ pub mod routes;
#[cfg(feature = "server")]
pub mod state;
// Realtime request logging. Only the server serves realtime, so these are
// `server`-gated; `io::realtime` exposes the generic `observe` hook while the
// collector and callback fan-out live here.
mod constants;
pub mod integrations;
#[cfg(feature = "server")]

View file

@ -11,7 +11,7 @@
use std::sync::Arc;
use litellm_ai_gateway::io::realtime_pool::{upstream_key, PoolConfig, RealtimePool};
use litellm_ai_gateway::io::realtime_pool::{PoolConfig, RealtimePool, upstream_key};
use litellm_ai_gateway::routes;
use litellm_ai_gateway::state::AppState;
use litellm_core::router::{Deployment, LiteLLMParams, Router};

View file

@ -1,7 +1,8 @@
use litellm_core::error::{json_type_name, CoreError};
use litellm_core::messages::transformation::AnthropicMessagesProviderConfig;
use litellm_core::providers::azure_ai::messages::transformation::AZURE_ANTHROPIC_MESSAGES_CONFIG;
use litellm_core::CoreResult;
use litellm_core::error::{CoreError, json_type_name};
use litellm_core::messages::transformation::AnthropicMessagesProviderConfig;
use litellm_core::providers::anthropic::messages::transformation::ANTHROPIC_MESSAGES_CONFIG;
use litellm_core::providers::azure_ai::messages::transformation::AZURE_ANTHROPIC_MESSAGES_CONFIG;
use serde_json::{Map, Value};
use crate::constants::MESSAGES_ERROR_BODY_MAX_CHARS;
@ -18,6 +19,7 @@ pub(super) fn messages_provider_config(
provider: &str,
) -> Option<&'static dyn AnthropicMessagesProviderConfig> {
match provider {
"anthropic" => Some(&ANTHROPIC_MESSAGES_CONFIG),
"azure_ai" => Some(&AZURE_ANTHROPIC_MESSAGES_CONFIG),
_ => None,
}
@ -48,3 +50,15 @@ pub(super) fn has_header(headers: &[(String, String)], name: &str) -> bool {
.iter()
.any(|(key, _)| key.eq_ignore_ascii_case(name))
}
pub(super) fn has_bearer_auth(headers: &[(String, String)]) -> bool {
headers.iter().any(|(name, value)| {
if !name.eq_ignore_ascii_case("authorization") {
return false;
}
let value = value.trim();
value.len() > 7
&& value[..7].eq_ignore_ascii_case("bearer ")
&& !value[7..].trim().is_empty()
})
}

View file

@ -1,10 +1,11 @@
use litellm_core::error::CoreError;
use litellm_core::CoreResult;
use litellm_core::error::CoreError;
use serde_json::Value;
use super::client::http_client;
use super::common_utils::truncate_error_body;
use super::types::ProviderMessagesRequest;
use crate::constants::ANTHROPIC_MESSAGES_PROVIDER;
pub(super) async fn execute_messages_provider_call(
request: ProviderMessagesRequest,
@ -45,3 +46,38 @@ pub(super) async fn execute_messages_provider_call(
CoreError::InvalidResponse(format!("failed to serialize messages response: {err}"))
})
}
pub(super) async fn execute_messages_provider_stream(
request: ProviderMessagesRequest,
) -> CoreResult<reqwest::Response> {
if request.provider != ANTHROPIC_MESSAGES_PROVIDER {
return Err(CoreError::InvalidRequest(
"streaming messages is not supported for this provider".to_string(),
));
}
let mut request_builder = http_client().post(&request.url).json(&request.body);
for (key, value) in &request.upstream_headers {
request_builder = request_builder.header(key, value);
}
if let Some(duration) = request.timeout {
request_builder = request_builder.timeout(duration);
}
let response = request_builder
.send()
.await
.map_err(|err| CoreError::Network(err.to_string()))?;
let status = response.status();
if !status.is_success() {
let text = response
.text()
.await
.map_err(|err| CoreError::Network(err.to_string()))?;
return Err(CoreError::Http {
status: status.as_u16(),
body: truncate_error_body(&text),
});
}
Ok(response)
}

View file

@ -9,12 +9,40 @@ mod types;
pub use types::MessagesRequest;
use handler::execute_messages_provider_call;
use handler::{execute_messages_provider_call, execute_messages_provider_stream};
use prepare::prepare_messages_call;
pub async fn messages(request: MessagesRequest<'_>) -> CoreResult<Value> {
match execute_messages(request, false).await? {
MessagesResponse::Json(body) => Ok(body),
MessagesResponse::Stream(response) => {
drop(response);
Err(litellm_core::CoreError::InvalidResponse(
"non-streaming messages execution returned a stream".to_string(),
))
}
}
}
pub(crate) enum MessagesResponse {
Json(Value),
Stream(reqwest::Response),
}
pub(crate) async fn execute_messages(
request: MessagesRequest<'_>,
stream: bool,
) -> CoreResult<MessagesResponse> {
let prepared = prepare_messages_call(request)?;
execute_messages_provider_call(prepared).await
if stream {
execute_messages_provider_stream(prepared)
.await
.map(MessagesResponse::Stream)
} else {
execute_messages_provider_call(prepared)
.await
.map(MessagesResponse::Json)
}
}
#[cfg(test)]

View file

@ -1,9 +1,9 @@
use litellm_core::messages::transformation::MessagesAuthStrategy;
use litellm_core::routing_utils::provider::{get_custom_llm_provider, CustomLlmProvider};
use litellm_core::CoreError;
use litellm_core::CoreResult;
use litellm_core::messages::transformation::MessagesAuthStrategy;
use litellm_core::routing_utils::provider::{CustomLlmProvider, get_custom_llm_provider};
use super::common_utils::{has_header, messages_provider_config, string_headers};
use super::common_utils::{has_bearer_auth, has_header, messages_provider_config, string_headers};
use super::types::{MessagesRequest, ProviderMessagesRequest};
pub(super) fn prepare_messages_call(
@ -33,7 +33,9 @@ pub(super) fn prepare_messages_call(
let mut headers = string_headers(request.extra_headers)?;
let auth_strategy = config.auth_strategy();
if !has_header(&headers, auth_strategy.header_name()) {
let already_authorized = has_header(&headers, auth_strategy.header_name())
|| (config.accepts_bearer_auth() && has_bearer_auth(&headers));
if !already_authorized {
let api_key = config.resolve_api_key(request.api_key, &env_lookup)?;
let auth_header = match auth_strategy {
MessagesAuthStrategy::Bearer => {
@ -62,6 +64,7 @@ pub(super) fn prepare_messages_call(
})?;
Ok(ProviderMessagesRequest {
provider: provider.to_string(),
model,
config,
url,

View file

@ -1,14 +1,14 @@
use std::time::Duration;
use litellm_core::error::CoreError;
use serde_json::{json, Map, Value};
use serde_json::{Map, Value, json};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::{TcpListener, TcpStream};
use super::common_utils::{
has_header, messages_provider_config, string_headers, truncate_error_body,
has_bearer_auth, has_header, messages_provider_config, string_headers, truncate_error_body,
};
use super::{messages, MessagesRequest};
use super::{MessagesRequest, messages};
async fn read_http_request(socket: &mut TcpStream) -> String {
let mut request = Vec::new();
@ -52,9 +52,9 @@ fn write_response(body: &str) -> String {
}
#[test]
fn provider_config_only_resolves_azure_ai() {
fn provider_config_resolves_anthropic_and_azure_ai() {
assert!(messages_provider_config("anthropic").is_some());
assert!(messages_provider_config("azure_ai").is_some());
assert!(messages_provider_config("anthropic").is_none());
assert!(messages_provider_config("openai").is_none());
}
@ -85,6 +85,34 @@ fn has_header_is_case_insensitive() {
assert!(!has_header(&headers, "authorization"));
}
#[test]
fn has_bearer_auth_requires_a_nonempty_bearer_token() {
assert!(has_bearer_auth(&[(
"Authorization".to_string(),
"Bearer tok".to_string()
)]));
assert!(has_bearer_auth(&[(
"authorization".to_string(),
"bearer tok".to_string()
)]));
assert!(!has_bearer_auth(&[(
"authorization".to_string(),
"Bearer ".to_string()
)]));
assert!(!has_bearer_auth(&[(
"authorization".to_string(),
String::new()
)]));
assert!(!has_bearer_auth(&[(
"authorization".to_string(),
"Basic abc".to_string()
)]));
assert!(!has_bearer_auth(&[(
"x-api-key".to_string(),
"sk".to_string()
)]));
}
#[tokio::test]
async fn messages_round_trip_builds_azure_request_and_passes_response_through() {
let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds");
@ -148,6 +176,52 @@ async fn messages_round_trip_builds_azure_request_and_passes_response_through()
);
}
#[tokio::test]
async fn messages_round_trip_builds_native_anthropic_request() {
let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds");
let addr = listener.local_addr().expect("addr");
let server = tokio::spawn(async move {
let (mut socket, _) = listener.accept().await.expect("accepts request");
let request = read_http_request(&mut socket).await;
let response_body = r#"{"id":"msg_1","type":"message","role":"assistant","content":[{"type":"text","text":"hi"}],"model":"claude-sonnet-4-5","stop_reason":"end_turn","usage":{"input_tokens":1,"output_tokens":2}}"#;
socket
.write_all(write_response(response_body).as_bytes())
.await
.expect("writes response");
request
});
let response = messages(MessagesRequest {
model: "claude-sonnet-4-5",
body: json!({
"model": "claude-sonnet-4-5",
"max_tokens": 1024,
"messages": [{"role": "user", "content": "hi"}]
}),
api_key: Some("sk-ant"),
api_base: Some(&format!("http://{addr}")),
custom_llm_provider: Some("anthropic"),
extra_headers: None,
timeout: Some(Duration::from_secs(5)),
})
.await
.expect("messages request succeeds");
assert_eq!(response["content"][0]["text"], "hi");
assert_eq!(response["stop_reason"], "end_turn");
let request = server.await.expect("server task completes");
let (head, _) = request.split_once("\r\n\r\n").expect("has body");
assert!(head.starts_with("POST /v1/messages "), "{head}");
let head_lower = head.to_ascii_lowercase();
assert!(head_lower.contains("x-api-key: sk-ant"), "{head}");
assert!(
head_lower.contains("anthropic-version: 2023-06-01"),
"{head}"
);
}
#[tokio::test]
async fn messages_does_not_duplicate_auth_when_x_api_key_supplied() {
let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds");
@ -206,6 +280,112 @@ async fn messages_does_not_duplicate_auth_when_x_api_key_supplied() {
assert!(!head.contains("rust-fallback-key"), "{head}");
}
#[tokio::test]
async fn messages_forwards_entra_id_bearer_without_requiring_api_key() {
let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds");
let addr = listener.local_addr().expect("addr");
let server = tokio::spawn(async move {
let (mut socket, _) = listener.accept().await.expect("accepts request");
let request = read_http_request(&mut socket).await;
let response_body =
r#"{"id":"msg_3","type":"message","role":"assistant","content":[],"model":"m"}"#;
socket
.write_all(write_response(response_body).as_bytes())
.await
.expect("writes response");
request
});
let mut headers = Map::new();
headers.insert(
"Authorization".to_string(),
Value::String("Bearer entra-token".to_string()),
);
messages(MessagesRequest {
model: "claude-sonnet-4-5",
body: json!({"model": "claude-sonnet-4-5", "max_tokens": 8, "messages": []}),
api_key: None,
api_base: Some(&format!("http://{addr}")),
custom_llm_provider: Some("azure_ai"),
extra_headers: Some(headers),
timeout: Some(Duration::from_secs(5)),
})
.await
.expect("entra id request succeeds without api key");
let request = server.await.expect("server task completes");
let head = request
.split_once("\r\n\r\n")
.expect("has body")
.0
.to_ascii_lowercase();
assert!(head.contains("authorization: bearer entra-token"), "{head}");
assert!(!head.contains("x-api-key"), "{head}");
}
#[tokio::test]
async fn messages_requires_auth_when_no_key_and_no_header() {
let err = messages(MessagesRequest {
model: "claude-sonnet-4-5",
body: json!({"model": "claude-sonnet-4-5", "max_tokens": 8, "messages": []}),
api_key: None,
api_base: Some("http://127.0.0.1:1"),
custom_llm_provider: Some("azure_ai"),
extra_headers: None,
timeout: Some(Duration::from_millis(50)),
})
.await
.expect_err("missing auth errors");
assert!(matches!(err, CoreError::Auth(_)));
}
#[tokio::test]
async fn messages_ignores_malformed_authorization_and_uses_api_key() {
let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds");
let addr = listener.local_addr().expect("addr");
let server = tokio::spawn(async move {
let (mut socket, _) = listener.accept().await.expect("accepts request");
let request = read_http_request(&mut socket).await;
let response_body =
r#"{"id":"msg_4","type":"message","role":"assistant","content":[],"model":"m"}"#;
socket
.write_all(write_response(response_body).as_bytes())
.await
.expect("writes response");
request
});
let mut headers = Map::new();
headers.insert(
"Authorization".to_string(),
Value::String("Bearer ".to_string()),
);
messages(MessagesRequest {
model: "claude-sonnet-4-5",
body: json!({"model": "claude-sonnet-4-5", "max_tokens": 8, "messages": []}),
api_key: Some("sk-azure"),
api_base: Some(&format!("http://{addr}")),
custom_llm_provider: Some("azure_ai"),
extra_headers: Some(headers),
timeout: Some(Duration::from_secs(5)),
})
.await
.expect("falls back to api key");
let request = server.await.expect("server task completes");
let head = request
.split_once("\r\n\r\n")
.expect("has body")
.0
.to_ascii_lowercase();
assert!(head.contains("x-api-key: sk-azure"), "{head}");
}
#[tokio::test]
async fn messages_maps_provider_error_status_to_http_error() {
let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds");
@ -248,12 +428,12 @@ async fn messages_rejects_unsupported_provider() {
body: json!({"model": "claude-3-5-sonnet", "max_tokens": 8, "messages": []}),
api_key: Some("sk"),
api_base: Some("http://127.0.0.1:1"),
custom_llm_provider: Some("anthropic"),
custom_llm_provider: Some("openai"),
extra_headers: None,
timeout: Some(Duration::from_millis(50)),
})
.await
.expect_err("unsupported provider errors");
assert!(matches!(err, CoreError::InvalidProvider(provider) if provider == "anthropic"));
assert!(matches!(err, CoreError::InvalidProvider(provider) if provider == "openai"));
}

View file

@ -14,6 +14,7 @@ pub struct MessagesRequest<'a> {
}
pub(crate) struct ProviderMessagesRequest {
pub(crate) provider: String,
pub(crate) model: String,
pub(crate) config: &'static dyn AnthropicMessagesProviderConfig,
pub(crate) url: String,

View file

@ -1,11 +1,11 @@
use std::net::IpAddr;
use std::time::{Duration, Instant};
use base64::engine::general_purpose::STANDARD as BASE64_STANDARD;
use base64::Engine;
use base64::engine::general_purpose::STANDARD as BASE64_STANDARD;
use litellm_core::CoreResult;
use litellm_core::error::CoreError;
use litellm_core::ocr::transformation::OcrProviderConfig;
use litellm_core::CoreResult;
use reqwest::Url;
use serde_json::{Map, Value};
@ -18,7 +18,7 @@ use litellm_core::providers::vertex_ai::ocr::transformation::{
VERTEX_AI_DEEPSEEK_OCR_CONFIG, VERTEX_AI_OCR_CONFIG,
};
use super::client::http_client;
use crate::client::http_client;
const ERROR_BODY_MAX_CHARS: usize = 256;
const AZURE_DOCUMENT_INTELLIGENCE_POLL_TIMEOUT_SECS: u64 = 120;

View file

@ -1,11 +1,11 @@
use litellm_core::CoreResult;
use litellm_core::error::CoreError;
use litellm_core::ocr::transformation::OcrResponseHandling;
use litellm_core::CoreResult;
use serde_json::Value;
use super::client::http_client;
use super::common_utils::{poll_document_intelligence, truncate_error_body};
use super::types::ProviderOcrRequest;
use crate::client::http_client;
pub(crate) async fn execute_ocr_provider_call(request: ProviderOcrRequest) -> CoreResult<Value> {
let mut request_builder = http_client().post(&request.url).json(&request.body);

View file

@ -1,11 +1,11 @@
use std::future::Future;
use std::pin::Pin;
use litellm_core::CoreResult;
use litellm_core::call_lifecycle::{CallLifecycleContext, CallLifecycleHooks, CallLifecycleTiming};
use litellm_core::error::CoreError;
use litellm_core::ocr::transformation::OcrAuthStrategy;
use litellm_core::CoreResult;
use serde_json::{json, Map, Value};
use serde_json::{Map, Value, json};
use super::common_utils::{
convert_document_url_to_data_uri, has_header, ocr_provider_config, string_headers,
@ -292,7 +292,7 @@ fn parse_ocr_pre_call_guardrail_request(
Some(_) => {
return Err(CoreError::InvalidRequest(
"OCR pre_call guardrail optional_params must be an object".to_string(),
))
));
}
None => Map::new(),
};

View file

@ -1,8 +1,7 @@
use litellm_core::call_lifecycle::CallLifecycle;
use litellm_core::CoreResult;
use litellm_core::call_lifecycle::CallLifecycle;
use serde_json::Value;
mod client;
mod common_utils;
mod handler;
mod hooks;
@ -12,7 +11,7 @@ mod types;
pub use types::OcrRequest;
use handler::execute_ocr_provider_call;
use prepare::{prepare_ocr_call, PreparedOcrCall};
use prepare::{PreparedOcrCall, prepare_ocr_call};
pub async fn ocr(request: OcrRequest<'_>) -> CoreResult<Value> {
let PreparedOcrCall { request, hooks } = prepare_ocr_call(request);

View file

@ -1,7 +1,7 @@
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::{SystemTime, UNIX_EPOCH};
use litellm_core::routing_utils::provider::{get_custom_llm_provider, CustomLlmProvider};
use litellm_core::routing_utils::provider::{CustomLlmProvider, get_custom_llm_provider};
use super::hooks::OcrLifecycleHooks;
use super::types::{OcrRequest, PreparedOcrRequest};

View file

@ -3,12 +3,12 @@ use std::time::Duration;
use litellm_core::error::CoreError;
use litellm_core::ocr::transformation::OcrResponseHandling;
use serde_json::{json, Map, Value};
use serde_json::{Map, Value, json};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::{TcpListener, TcpStream};
use super::common_utils::{has_header, ocr_provider_config, string_headers, truncate_error_body};
use super::{ocr, OcrRequest};
use super::{OcrRequest, ocr};
use crate::integrations::custom_guardrail::{
CustomGuardrail, GuardrailContext, GuardrailDecision, GuardrailError, GuardrailEventHook,
GuardrailFuture, GuardrailRequest,
@ -228,19 +228,23 @@ fn truncate_error_body_does_not_split_multibyte_chars() {
#[test]
fn ocr_dispatch_supports_migrated_providers() {
assert!(ocr_provider_config("mistral", "mistral-ocr-latest").is_some());
assert!(ocr_provider_config("azure_ai", "pixtral-12b-2409")
.expect("azure ai config resolves")
.requires_data_uri_document());
assert!(
ocr_provider_config("azure_ai", "pixtral-12b-2409")
.expect("azure ai config resolves")
.requires_data_uri_document()
);
assert_eq!(
ocr_provider_config("azure_ai", "doc-intelligence/prebuilt-read")
.expect("document intelligence config resolves")
.response_handling(),
OcrResponseHandling::AzureDocumentIntelligencePoll
);
assert!(ocr_provider_config("vertex_ai", "deepseek-ocr-maas")
.expect("vertex deepseek config resolves")
.supported_ocr_params()
.contains(&"temperature"));
assert!(
ocr_provider_config("vertex_ai", "deepseek-ocr-maas")
.expect("vertex deepseek config resolves")
.supported_ocr_params()
.contains(&"temperature")
);
assert!(ocr_provider_config("openai", "gpt-4o").is_none());
}

View file

@ -7,9 +7,9 @@
//!
//! Compiled only under the `python-config` feature.
use litellm_core::CoreResult;
use litellm_core::error::CoreError;
use litellm_core::router::{Deployment, Router};
use litellm_core::CoreResult;
use pyo3::prelude::*;
use crate::gil;

View file

@ -1,8 +1,8 @@
//! Health probes. Simple-route template: a `router()` plus its handlers, in one file.
use axum::Router;
use axum::http::StatusCode;
use axum::routing::get;
use axum::Router;
use crate::state::AppState;

View file

@ -0,0 +1,513 @@
//! `POST /v1/messages`, the Anthropic Messages HTTP surface.
mod service;
use axum::Router;
use axum::body::Body;
use axum::extract::{Json, State};
use axum::http::StatusCode;
use axum::http::header::{CACHE_CONTROL, CONTENT_TYPE, HeaderMap, HeaderValue};
use axum::response::{IntoResponse, Response};
use axum::routing::post;
use litellm_core::CoreError;
use serde_json::{Map, Value};
use crate::auth::RequireMasterKey;
use crate::constants::{MESSAGES_HEADERS_NOT_FORWARDED, MESSAGES_ROUTE_PATH};
use crate::state::AppState;
/// This route's contribution to the app router.
pub fn router() -> Router<AppState> {
Router::new().route(MESSAGES_ROUTE_PATH, post(handle))
}
async fn handle(
_auth: RequireMasterKey,
State(state): State<AppState>,
headers: HeaderMap,
Json(body): Json<Value>,
) -> Result<Response, MessagesRouteError> {
let extra_headers = forwarded_headers(&headers)?;
match service::run(&state.router, body, extra_headers)
.await
.map_err(MessagesRouteError::from)?
{
service::MessagesResponse::Json(body) => Ok(Json(body).into_response()),
service::MessagesResponse::Stream(upstream) => stream_response(upstream),
}
}
fn stream_response(upstream: reqwest::Response) -> Result<Response, MessagesRouteError> {
let content_type = upstream
.headers()
.get(CONTENT_TYPE)
.cloned()
.unwrap_or_else(|| HeaderValue::from_static("text/event-stream"));
let mut response = Response::builder()
.status(
StatusCode::from_u16(upstream.status().as_u16()).map_err(|error| {
MessagesRouteError(CoreError::InvalidResponse(format!(
"invalid upstream response status: {error}"
)))
})?,
)
.header(CONTENT_TYPE, content_type);
if let Some(value) = upstream.headers().get(CACHE_CONTROL) {
response = response.header(CACHE_CONTROL, value);
}
response
.body(Body::from_stream(upstream.bytes_stream()))
.map_err(|error| {
MessagesRouteError(CoreError::InvalidResponse(format!(
"failed to build streaming response: {error}"
)))
})
}
fn forwarded_headers(headers: &HeaderMap) -> Result<Option<Map<String, Value>>, CoreError> {
let forwarded = headers
.iter()
.filter(|(name, _)| {
!MESSAGES_HEADERS_NOT_FORWARDED
.iter()
.any(|excluded| name.as_str().eq_ignore_ascii_case(excluded))
})
.map(|(name, value)| {
let value = value.to_str().map_err(|_| {
CoreError::InvalidRequest(format!("invalid value for header {}", name.as_str()))
})?;
Ok((name.to_string(), Value::String(value.to_string())))
})
.collect::<Result<Map<_, _>, CoreError>>()?;
Ok((!forwarded.is_empty()).then_some(forwarded))
}
#[derive(Debug)]
struct MessagesRouteError(CoreError);
impl From<CoreError> for MessagesRouteError {
fn from(error: CoreError) -> Self {
Self(error)
}
}
impl IntoResponse for MessagesRouteError {
fn into_response(self) -> Response {
let (status, message) = match self.0 {
CoreError::InvalidRequest(message) => (StatusCode::BAD_REQUEST, message),
CoreError::InvalidProvider(_) | CoreError::Routing(_) => (
StatusCode::NOT_FOUND,
"no messages deployment is configured for this model".to_string(),
),
CoreError::Auth(_) => (
StatusCode::BAD_GATEWAY,
"messages provider authentication failed".to_string(),
),
CoreError::Http { .. }
| CoreError::Network(_)
| CoreError::InvalidResponse(_)
| CoreError::InvalidType { .. }
| CoreError::MissingField(_) => (
StatusCode::BAD_GATEWAY,
"messages provider request failed".to_string(),
),
};
(
status,
Json(serde_json::json!({"error": {"message": message}})),
)
.into_response()
}
}
#[cfg(test)]
mod tests {
use std::sync::Arc;
use axum::body::Body;
use axum::http::Request;
use axum::http::StatusCode;
use axum::http::header::{CACHE_CONTROL, CONTENT_TYPE};
use litellm_core::router::{Deployment, LiteLLMParams, Router as ModelRouter};
use serde_json::json;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::TcpListener;
use tower::ServiceExt;
use super::super::app;
use crate::io::realtime_pool::RealtimePool;
use crate::state::AppState;
fn state(model: &str, api_base: String, master_key: Option<&str>) -> AppState {
state_with_provider(model, model, api_base, master_key)
}
fn state_with_provider(
model_alias: &str,
provider_model: &str,
api_base: String,
master_key: Option<&str>,
) -> AppState {
AppState {
router: Arc::new(ModelRouter::new(vec![Deployment {
model_name: model_alias.to_string(),
litellm_params: LiteLLMParams {
model: format!("anthropic/{provider_model}"),
api_key: Some("upstream-key".to_string()),
api_base: Some(api_base),
},
}])),
master_key: master_key.map(Arc::from),
loggers: Arc::new(Vec::new()),
realtime_pool: RealtimePool::disabled(),
}
}
async fn upstream(listener: TcpListener) -> (String, tokio::task::JoinHandle<String>) {
let address = listener.local_addr().expect("listener has address");
let server = tokio::spawn(async move {
let (mut socket, _) = listener.accept().await.expect("accepts request");
let mut request = Vec::new();
let mut buffer = [0_u8; 4096];
loop {
let read = socket.read(&mut buffer).await.expect("reads request");
request.extend_from_slice(&buffer[..read]);
if request.windows(4).any(|window| window == b"\r\n\r\n") {
break;
}
}
let request = String::from_utf8(request).expect("request is utf8");
let content_length = request
.lines()
.find_map(|line| {
let (name, value) = line.split_once(':')?;
name.eq_ignore_ascii_case("content-length")
.then(|| value.trim().parse::<usize>().ok())
.flatten()
})
.unwrap_or(0);
let header_end = request.find("\r\n\r\n").expect("request has headers") + 4;
let mut full_request = request.into_bytes();
while full_request.len().saturating_sub(header_end) < content_length {
let read = socket.read(&mut buffer).await.expect("reads body");
full_request.extend_from_slice(&buffer[..read]);
}
let request = String::from_utf8(full_request).expect("request is utf8");
let body = r#"{"id":"msg_1","type":"message","role":"assistant","content":[],"model":"claude-test"}"#;
let response = format!(
"HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}",
body.len(),
body
);
socket
.write_all(response.as_bytes())
.await
.expect("writes response");
request
});
(format!("http://{address}"), server)
}
async fn streaming_upstream(
listener: TcpListener,
status: u16,
content_type: &'static str,
body: &'static str,
) -> (String, tokio::task::JoinHandle<String>) {
let address = listener.local_addr().expect("listener has address");
let server = tokio::spawn(async move {
let (mut socket, _) = listener.accept().await.expect("accepts request");
let mut request = Vec::new();
let mut buffer = [0_u8; 4096];
loop {
let read = socket.read(&mut buffer).await.expect("reads request");
request.extend_from_slice(&buffer[..read]);
if request.windows(4).any(|window| window == b"\r\n\r\n") {
break;
}
}
let request_text = String::from_utf8(request).expect("request is utf8");
let content_length = request_text
.lines()
.find_map(|line| {
let (name, value) = line.split_once(':')?;
name.eq_ignore_ascii_case("content-length")
.then(|| value.trim().parse::<usize>().ok())
.flatten()
})
.unwrap_or(0);
let header_end = request_text.find("\r\n\r\n").expect("request has headers") + 4;
let mut full_request = request_text.into_bytes();
while full_request.len().saturating_sub(header_end) < content_length {
let read = socket.read(&mut buffer).await.expect("reads body");
full_request.extend_from_slice(&buffer[..read]);
}
let response = format!(
"HTTP/1.1 {status} OK\r\ncontent-type: {content_type}\r\ncache-control: no-cache\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{body}",
body.len()
);
socket
.write_all(response.as_bytes())
.await
.expect("writes response");
String::from_utf8(full_request).expect("request is utf8")
});
(format!("http://{address}"), server)
}
#[tokio::test]
async fn route_constructs_anthropic_upstream_request() {
let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds");
let (api_base, server) = upstream(listener).await;
let app = app(state("claude-test", api_base, Some("master-key")));
let response = app
.oneshot(
Request::builder()
.method("POST")
.uri("/v1/messages")
.header("authorization", "Bearer master-key")
.header("x-api-key", "request-upstream-key")
.header("anthropic-beta", "beta-feature")
.header("content-type", "application/json")
.body(Body::from(
json!({
"model": "claude-test",
"max_tokens": 16,
"messages": [{"role": "user", "content": "hello"}]
})
.to_string(),
))
.expect("request builds"),
)
.await
.expect("route responds");
assert_eq!(response.status(), StatusCode::OK);
let body = axum::body::to_bytes(response.into_body(), usize::MAX)
.await
.expect("response body reads");
assert_eq!(
serde_json::from_slice::<serde_json::Value>(&body).expect("json")["id"],
"msg_1"
);
let upstream_request = server.await.expect("upstream task completes");
let (head, body) = upstream_request
.split_once("\r\n\r\n")
.expect("upstream request has body");
let head = head.to_ascii_lowercase();
assert!(head.contains("x-api-key: request-upstream-key"));
assert!(head.contains("anthropic-beta: beta-feature"));
assert!(!head.contains("authorization: bearer master-key"));
let body: serde_json::Value = serde_json::from_str(body).expect("upstream body is json");
assert_eq!(body["model"], "claude-test");
assert_eq!(body["messages"][0]["content"], "hello");
}
#[tokio::test]
async fn route_substitutes_model_alias_with_provider_model_upstream() {
let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds");
let (api_base, server) = upstream(listener).await;
let app = app(state_with_provider(
"production",
"claude-sonnet-4-5",
api_base,
Some("master-key"),
));
let response = app
.oneshot(
Request::builder()
.method("POST")
.uri("/v1/messages")
.header("authorization", "Bearer master-key")
.header("content-type", "application/json")
.body(Body::from(
json!({
"model": "production",
"max_tokens": 16,
"messages": [{"role": "user", "content": "hello"}]
})
.to_string(),
))
.expect("request builds"),
)
.await
.expect("route responds");
assert_eq!(response.status(), StatusCode::OK);
let upstream_request = server.await.expect("upstream task completes");
let (_, upstream_body) = upstream_request
.split_once("\r\n\r\n")
.expect("upstream request has body");
let upstream_body: serde_json::Value =
serde_json::from_str(upstream_body).expect("upstream body is json");
assert_eq!(upstream_body["model"], "claude-sonnet-4-5");
assert_ne!(upstream_body["model"], "production");
}
#[tokio::test]
async fn route_streams_anthropic_events_without_buffering_or_reordering() {
let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds");
let events = "event: message_start\ndata: {\"type\":\"message_start\"}\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\"}\n\nevent: message_stop\ndata: {\"type\":\"message_stop\"}\n\n";
let (api_base, server) =
streaming_upstream(listener, 200, "text/event-stream", events).await;
let app = app(state("claude-test", api_base, Some("master-key")));
let response = app
.oneshot(
Request::builder()
.method("POST")
.uri("/v1/messages")
.header("authorization", "Bearer master-key")
.header("content-type", "application/json")
.body(Body::from(
json!({
"model": "claude-test",
"max_tokens": 16,
"stream": true,
"messages": [{"role": "user", "content": "hello"}]
})
.to_string(),
))
.expect("request builds"),
)
.await
.expect("route responds");
assert_eq!(response.status(), StatusCode::OK);
assert_eq!(
response
.headers()
.get(CONTENT_TYPE)
.unwrap()
.to_str()
.unwrap(),
"text/event-stream"
);
assert_eq!(
response
.headers()
.get(CACHE_CONTROL)
.unwrap()
.to_str()
.unwrap(),
"no-cache"
);
let response_body = axum::body::to_bytes(response.into_body(), usize::MAX)
.await
.expect("response body reads");
assert_eq!(response_body, events.as_bytes());
let upstream_request = server.await.expect("upstream task completes");
let (_, upstream_body) = upstream_request
.split_once("\r\n\r\n")
.expect("upstream request has body");
assert_eq!(
serde_json::from_str::<serde_json::Value>(upstream_body)
.expect("upstream body is json")["stream"],
true
);
}
#[tokio::test]
async fn route_maps_streaming_upstream_errors_before_starting_response() {
let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds");
let (api_base, server) = streaming_upstream(
listener,
429,
"application/json",
r#"{"error":"rate limited"}"#,
)
.await;
let app = app(state("claude-test", api_base, Some("master-key")));
let response = app
.oneshot(
Request::builder()
.method("POST")
.uri("/v1/messages")
.header("authorization", "Bearer master-key")
.header("content-type", "application/json")
.body(Body::from(
json!({
"model": "claude-test",
"max_tokens": 16,
"stream": true,
"messages": [{"role": "user", "content": "hello"}]
})
.to_string(),
))
.expect("request builds"),
)
.await
.expect("route responds");
assert_eq!(response.status(), StatusCode::BAD_GATEWAY);
let response_body = axum::body::to_bytes(response.into_body(), usize::MAX)
.await
.expect("response body reads");
assert_eq!(
serde_json::from_slice::<serde_json::Value>(&response_body).expect("error is json")["error"]
["message"],
"messages provider request failed"
);
server.await.expect("upstream task completes");
}
#[tokio::test]
async fn route_rejects_missing_master_key() {
let app = app(state(
"claude-test",
"http://127.0.0.1:1".to_string(),
Some("master-key"),
));
let response = app
.oneshot(
Request::builder()
.method("POST")
.uri("/v1/messages")
.header("content-type", "application/json")
.body(Body::from("{}"))
.expect("request builds"),
)
.await
.expect("route responds");
assert_eq!(response.status(), StatusCode::UNAUTHORIZED);
}
#[tokio::test]
async fn route_rejects_invalid_master_key() {
let app = app(state(
"claude-test",
"http://127.0.0.1:1".to_string(),
Some("master-key"),
));
let response = app
.oneshot(
Request::builder()
.method("POST")
.uri("/v1/messages")
.header("authorization", "Bearer wrong-key")
.header("content-type", "application/json")
.body(Body::from("{}"))
.expect("request builds"),
)
.await
.expect("route responds");
assert_eq!(response.status(), StatusCode::UNAUTHORIZED);
}
#[tokio::test]
async fn route_rejects_malformed_json_without_panicking() {
let app = app(state(
"claude-test",
"http://127.0.0.1:1".to_string(),
Some("master-key"),
));
let response = app
.oneshot(
Request::builder()
.method("POST")
.uri("/v1/messages")
.header("authorization", "Bearer master-key")
.header("content-type", "application/json")
.body(Body::from("{not-json"))
.expect("request builds"),
)
.await
.expect("route responds");
assert_eq!(response.status(), StatusCode::BAD_REQUEST);
}
}

View file

@ -0,0 +1,64 @@
use std::sync::Arc;
use litellm_core::router::Router;
use litellm_core::{CoreError, CoreResult};
use serde_json::{Map, Value};
use crate::constants::ANTHROPIC_MESSAGES_PROVIDER;
use crate::messages::{MessagesRequest, execute_messages};
pub(crate) enum MessagesResponse {
Json(Value),
Stream(reqwest::Response),
}
pub async fn run(
router: &Arc<Router>,
body: Value,
extra_headers: Option<Map<String, Value>>,
) -> CoreResult<MessagesResponse> {
let model = body
.get("model")
.and_then(Value::as_str)
.map(str::trim)
.filter(|model| !model.is_empty())
.ok_or_else(|| CoreError::InvalidRequest("messages body requires a model".to_string()))?;
let deployment = router.get_available_deployment(model).ok_or_else(|| {
CoreError::Routing(format!("no deployment available for model '{model}'"))
})?;
let provider_model = deployment.litellm_params.model.as_str();
let upstream_model = provider_model
.split_once('/')
.map_or(provider_model, |(_, model)| model);
let custom_llm_provider = if provider_model.contains('/') {
None
} else {
Some(ANTHROPIC_MESSAGES_PROVIDER)
};
let mut body = body;
body.as_object_mut()
.ok_or_else(|| CoreError::InvalidRequest("messages body must be an object".to_string()))?
.insert(
"model".to_string(),
Value::String(upstream_model.to_string()),
);
let request = MessagesRequest {
model: provider_model,
body,
api_key: deployment.litellm_params.api_key.as_deref(),
api_base: deployment.litellm_params.api_base.as_deref(),
custom_llm_provider,
extra_headers,
timeout: None,
};
let stream = request.body.get("stream").and_then(Value::as_bool) == Some(true);
execute_messages(request, stream)
.await
.map(|response| match response {
crate::messages::MessagesResponse::Json(body) => MessagesResponse::Json(body),
crate::messages::MessagesResponse::Stream(upstream) => {
MessagesResponse::Stream(upstream)
}
})
}

View file

@ -7,7 +7,9 @@
pub mod gil;
pub mod health;
pub mod messages;
pub mod realtime;
pub mod responses;
use axum::Router;
@ -18,6 +20,8 @@ pub fn app(state: AppState) -> Router {
Router::new()
.merge(health::router())
.merge(gil::router())
.merge(messages::router())
.merge(realtime::router())
.merge(responses::router())
.with_state(state)
}

View file

@ -6,17 +6,17 @@
mod service;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::{SystemTime, UNIX_EPOCH};
use crate::io::realtime_pool::RealtimePool;
use axum::Router;
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;

View file

@ -9,12 +9,12 @@
use std::time::Duration;
use crate::io::realtime_pool::{upstream_key, RealtimePool};
use crate::io::realtime_pool::{RealtimePool, upstream_key};
use futures_util::{Sink, Stream};
use litellm_core::CoreResult;
use litellm_core::error::CoreError;
use litellm_core::realtime::types::RealtimeEvent;
use litellm_core::router::Router;
use litellm_core::CoreResult;
/// Select a deployment for `model` and splice the client stream to the provider.
///

View file

@ -0,0 +1,348 @@
mod service;
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::{SystemTime, UNIX_EPOCH};
use axum::Router;
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 futures_util::{Sink, SinkExt, StreamExt};
use litellm_core::responses::types::{ResponsesErrorFrame, ResponsesWsEvent, ResponsesWsEventType};
use litellm_core::router::Router as ModelRouter;
use serde::Deserialize;
use crate::auth::RequireMasterKey;
use crate::integrations::custom_logger::CustomLogger;
use crate::integrations::types::RequestMetadata;
use crate::state::AppState;
static CALL_SEQ: AtomicU64 = AtomicU64::new(0);
fn new_call_id() -> String {
let nanos = SystemTime::now()
.duration_since(UNIX_EPOCH)
.map(|duration| duration.as_nanos())
.unwrap_or(0);
let sequence = CALL_SEQ.fetch_add(1, Ordering::Relaxed);
format!("respws-{nanos:x}-{sequence:x}")
}
pub fn router() -> Router<AppState> {
Router::new()
.route("/v1/responses", get(handle))
.route("/responses", get(handle))
}
#[derive(Debug, Deserialize)]
struct ResponsesQuery {
model: Option<String>,
}
async fn handle(
_auth: RequireMasterKey,
ws: WebSocketUpgrade,
State(state): State<AppState>,
Query(query): Query<ResponsesQuery>,
) -> Result<Response, (StatusCode, String)> {
if let Some(model) = query.model.as_deref() {
validate_model(&state.router, model)?;
}
let router = state.router.clone();
let loggers = state.loggers.clone();
let master_key = state.master_key.clone();
Ok(ws.on_upgrade(move |socket| bridge(socket, router, loggers, master_key, query.model)))
}
fn validate_model(router: &ModelRouter, model: &str) -> Result<(), (StatusCode, String)> {
if model.trim().is_empty() {
return Err((
StatusCode::BAD_REQUEST,
"missing 'model' query param".to_string(),
));
}
let Some(deployment) = router.get_available_deployment(model) else {
return Err((
StatusCode::NOT_FOUND,
format!("no deployment for model '{model}'"),
));
};
if deployment.litellm_params.model.contains('/')
&& !deployment.litellm_params.model.starts_with("openai/")
{
return Err((
StatusCode::BAD_REQUEST,
"Responses WebSocket route supports OpenAI deployments only".to_string(),
));
}
Ok(())
}
async fn send_error_and_close<S>(sink: &mut S, message: String)
where
S: futures_util::Sink<Message> + Unpin,
S::Error: std::fmt::Display,
{
if let Ok(payload) = serde_json::to_string(&ResponsesErrorFrame::invalid_request(message)) {
let _ = sink.send(Message::Text(payload)).await;
}
let _ = sink
.send(Message::Close(Some(axum::extract::ws::CloseFrame {
code: 1008,
reason: "Pre-call error".into(),
})))
.await;
let _ = sink.close().await;
}
struct ResponseClientSink {
sink: futures_util::stream::SplitSink<WebSocket, Message>,
}
impl Sink<ResponsesWsEvent> for ResponseClientSink {
type Error = axum::Error;
fn poll_ready(
mut self: std::pin::Pin<&mut Self>,
context: &mut std::task::Context<'_>,
) -> std::task::Poll<Result<(), Self::Error>> {
std::pin::Pin::new(&mut self.sink).poll_ready(context)
}
fn start_send(
mut self: std::pin::Pin<&mut Self>,
item: ResponsesWsEvent,
) -> Result<(), Self::Error> {
let payload = serde_json::to_string(&item).map_err(axum::Error::new)?;
std::pin::Pin::new(&mut self.sink).start_send(Message::Text(payload))
}
fn poll_flush(
mut self: std::pin::Pin<&mut Self>,
context: &mut std::task::Context<'_>,
) -> std::task::Poll<Result<(), Self::Error>> {
std::pin::Pin::new(&mut self.sink).poll_flush(context)
}
fn poll_close(
mut self: std::pin::Pin<&mut Self>,
context: &mut std::task::Context<'_>,
) -> std::task::Poll<Result<(), Self::Error>> {
std::pin::Pin::new(&mut self.sink).poll_close(context)
}
}
impl ResponseClientSink {
async fn close_with_code(&mut self, code: u16, reason: &'static str) {
let _ = self
.sink
.send(Message::Close(Some(axum::extract::ws::CloseFrame {
code,
reason: reason.into(),
})))
.await;
let _ = self.sink.close().await;
}
}
async fn bridge(
socket: WebSocket,
router: Arc<ModelRouter>,
loggers: Arc<Vec<Arc<dyn CustomLogger>>>,
master_key: Option<Arc<str>>,
requested_model: Option<String>,
) {
let (mut ws_sink, ws_stream) = socket.split();
let (model, first_frame, stream) = if let Some(model) = requested_model {
(model, None, ws_stream)
} else {
let mut stream = ws_stream;
let first = match stream.next().await {
Some(Ok(Message::Text(text))) => {
match serde_json::from_str::<ResponsesWsEvent>(&text) {
Ok(event) => event,
Err(_) => {
send_error_and_close(
&mut ws_sink,
"Invalid JSON in response.create event".to_string(),
)
.await;
return;
}
}
}
_ => {
send_error_and_close(&mut ws_sink, "Missing response.create event".to_string())
.await;
return;
}
};
let Some(model) = first.model().filter(|value| !value.trim().is_empty()) else {
send_error_and_close(
&mut ws_sink,
"Missing model in response.create event".to_string(),
)
.await;
return;
};
if first.event_type != ResponsesWsEventType::ResponseCreate {
send_error_and_close(
&mut ws_sink,
"First frame must be a response.create event".to_string(),
)
.await;
return;
}
(model.to_string(), Some(first), stream)
};
if let Err((status, message)) = validate_model(&router, &model) {
let _ = status;
let _ = message;
send_error_and_close(&mut ws_sink, "Unknown model deployment".to_string()).await;
return;
}
let call_id = new_call_id();
let metadata = RequestMetadata {
user_api_key_hash: master_key.as_deref().map(crate::auth::hash_token),
..RequestMetadata::default()
};
let client_in = Box::pin(stream.filter_map(|message| async move {
match message {
Ok(Message::Text(text)) => serde_json::from_str::<ResponsesWsEvent>(&text).ok(),
_ => None,
}
}));
let mut client_out = ResponseClientSink { sink: ws_sink };
let result = service::run(
&router,
&model,
first_frame,
None,
loggers,
call_id,
metadata,
client_in,
&mut client_out,
)
.await;
if result.is_err() {
client_out
.close_with_code(1011, "Internal server error")
.await;
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::io::realtime_pool::RealtimePool;
use crate::state::AppState;
use axum::body::Body;
use axum::http::Request;
use litellm_core::router::Router as ModelRouter;
use serde_json::json;
use std::pin::Pin;
use std::sync::Arc;
use std::task::{Context, Poll};
use tower::ServiceExt;
struct RecordingSink {
messages: Vec<Message>,
}
impl Sink<Message> for RecordingSink {
type Error = std::convert::Infallible;
fn poll_ready(
self: Pin<&mut Self>,
_context: &mut Context<'_>,
) -> Poll<Result<(), Self::Error>> {
Poll::Ready(Ok(()))
}
fn start_send(mut self: Pin<&mut Self>, item: Message) -> Result<(), Self::Error> {
self.messages.push(item);
Ok(())
}
fn poll_flush(
self: Pin<&mut Self>,
_context: &mut Context<'_>,
) -> Poll<Result<(), Self::Error>> {
Poll::Ready(Ok(()))
}
fn poll_close(
self: Pin<&mut Self>,
_context: &mut Context<'_>,
) -> Poll<Result<(), Self::Error>> {
Poll::Ready(Ok(()))
}
}
#[tokio::test]
async fn pre_call_error_matches_python_frame_and_close() {
let mut sink = RecordingSink {
messages: Vec::new(),
};
send_error_and_close(&mut sink, "missing model".to_string()).await;
let Message::Text(payload) = &sink.messages[0] else {
panic!("expected error text frame");
};
assert_eq!(
serde_json::from_str::<serde_json::Value>(payload).expect("error json"),
json!({
"type": "error",
"error": {
"type": "invalid_request_error",
"message": "missing model"
}
})
);
assert_eq!(
sink.messages[1],
Message::Close(Some(axum::extract::ws::CloseFrame {
code: 1008,
reason: "Pre-call error".into(),
}))
);
}
fn state() -> AppState {
AppState {
router: Arc::new(ModelRouter::default()),
master_key: Some(Arc::from("master-key")),
loggers: Arc::new(Vec::new()),
realtime_pool: RealtimePool::disabled(),
}
}
#[tokio::test]
async fn auth_rejects_responses_upgrade_before_handler() {
let request = Request::builder()
.uri("/responses?model=known")
.body(Body::empty())
.expect("request");
let response = router()
.with_state(state())
.oneshot(request)
.await
.expect("response");
assert_eq!(response.status(), StatusCode::UNAUTHORIZED);
}
#[test]
fn unknown_query_model_is_rejected_before_upgrade() {
assert_eq!(
validate_model(&ModelRouter::default(), "unknown").expect_err("unknown model"),
(
StatusCode::NOT_FOUND,
"no deployment for model 'unknown'".to_string()
)
);
}
}

View file

@ -0,0 +1,156 @@
use std::sync::Arc;
use std::time::Duration;
use futures_util::{Sink, Stream};
use litellm_core::call_lifecycle::{CallLifecycle, CallLifecycleContext};
use litellm_core::responses::instrumentation::{
ResponsesWsCallbackPayload, ResponsesWsInstrumentation, ResponsesWsLogOutcome,
ResponsesWsMetadata,
};
use litellm_core::responses::types::ResponsesWsEvent;
use litellm_core::{CoreError, CoreResult};
use crate::integrations::custom_logger::{
CallbackTiming, CallbackValue, CustomLogger, CustomLoggerRunner, LoggingError, ModelCallDetails,
};
use crate::integrations::types::RequestMetadata;
#[allow(clippy::too_many_arguments)]
pub async fn run<In, Out>(
router: &litellm_core::router::Router,
model: &str,
first_frame: Option<ResponsesWsEvent>,
idle_timeout: Option<Duration>,
loggers: Arc<Vec<Arc<dyn CustomLogger>>>,
call_id: String,
metadata: RequestMetadata,
client_in: In,
client_out: Out,
) -> CoreResult<()>
where
In: Stream<Item = ResponsesWsEvent> + Unpin + Send,
Out: Sink<ResponsesWsEvent> + Unpin + Send,
Out::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;
let provider_model = params
.model
.strip_prefix("openai/")
.unwrap_or(&params.model);
if params.model.contains('/') && !params.model.starts_with("openai/") {
return Err(CoreError::InvalidProvider(
"Responses WebSocket route supports OpenAI deployments only".to_string(),
));
}
let instrumentation = Arc::new(ResponsesWsInstrumentation::new(
call_id.clone(),
model,
ResponsesWsMetadata {
user_api_key_hash: metadata.user_api_key_hash,
user_api_key_user_id: metadata.user_api_key_user_id,
user_api_key_team_id: metadata.user_api_key_team_id,
},
));
let observer_instrumentation = Arc::clone(&instrumentation);
let context = CallLifecycleContext::new("responses_websocket", model, "openai", call_id);
let result = CallLifecycle::default()
.run(context, (), instrumentation.as_ref(), |_| async move {
crate::io::responses_ws::async_responses_websocket(
provider_model,
params.api_key.as_deref(),
params.api_base.as_deref(),
first_frame,
idle_timeout,
move |event| {
observer_instrumentation.observe(event);
},
client_in,
client_out,
)
.await
})
.await;
let outcome = instrumentation.take_or_build_outcome(result.is_ok());
dispatch_outcome(loggers, outcome).await;
result
}
async fn dispatch_outcome(
loggers: Arc<Vec<Arc<dyn CustomLogger>>>,
outcome: ResponsesWsLogOutcome,
) {
let runner = CustomLoggerRunner::new(loggers.as_ref().clone());
match outcome {
ResponsesWsLogOutcome::Success { payload, callback } => {
let (details, response, start_time, end_time) = logging_values(payload, callback, None);
let _ = runner
.async_log_success_event(
&details,
&response,
CallbackTiming::new(start_time, end_time),
)
.await;
}
ResponsesWsLogOutcome::Failure {
payload,
callback,
error_message,
error_kind,
} => {
let error = LoggingError {
message: error_message,
kind: error_kind,
};
let (details, response, start_time, end_time) =
logging_values(payload, callback, Some(error));
let _ = runner
.async_log_failure_event(
&details,
Some(&response),
CallbackTiming::new(start_time, end_time),
)
.await;
}
}
}
fn logging_values(
payload: litellm_core::responses::instrumentation::ResponsesWsLogPayload,
callback: ResponsesWsCallbackPayload,
error: Option<LoggingError>,
) -> (ModelCallDetails, CallbackValue, f64, f64) {
let start_time = payload.start_time;
let end_time = payload.end_time;
let callback = CallbackValue::new(callback.object, callback.value);
let details = ModelCallDetails::from_standard_logging_payload(
crate::integrations::types::StandardLoggingPayload {
id: payload.id,
litellm_call_id: payload.litellm_call_id,
call_type: payload.call_type,
model: payload.model,
custom_llm_provider: payload.custom_llm_provider,
response_cost: payload.response_cost,
prompt_tokens: payload.usage.prompt_tokens,
completion_tokens: payload.usage.completion_tokens,
total_tokens: payload.usage.total_tokens,
start_time: payload.start_time,
end_time: payload.end_time,
stream: payload.stream,
metadata: crate::integrations::types::StandardLoggingMetadata {
user_api_key_hash: payload.metadata.user_api_key_hash,
user_api_key_user_id: payload.metadata.user_api_key_user_id,
user_api_key_team_id: payload.metadata.user_api_key_team_id,
..Default::default()
},
messages: None,
},
);
let details = match error {
Some(error) => details.with_failure_error(error),
None => details,
};
(details, callback, start_time, end_time)
}

View file

@ -10,6 +10,25 @@ rand.workspace = true
serde.workspace = true
serde_json.workspace = true
thiserror.workspace = true
sha2.workspace = true
aws-config = { version = "1.9.0", default-features = false, features = ["rustls", "rt-tokio"], optional = true }
aws-credential-types = { version = "1.3.0", features = ["hardcoded-credentials"], optional = true }
aws-sdk-sts = { version = "1.108.0", default-features = false, features = ["rustls", "rt-tokio"], optional = true }
aws-sigv4 = { version = "1.5.1", optional = true }
aws-types = { version = "1.4.0", optional = true }
aws-smithy-runtime-api = { version = "1.13.0", optional = true }
[features]
default = []
bedrock-auth = [
"dep:aws-config",
"dep:aws-credential-types",
"dep:aws-sdk-sts",
"dep:aws-sigv4",
"dep:aws-types",
"dep:aws-smithy-runtime-api",
]
[dev-dependencies]
reqwest.workspace = true
tokio = { workspace = true, features = ["macros", "rt-multi-thread"] }

View file

@ -0,0 +1,2 @@
pub mod transformation;
pub mod types;

View file

@ -0,0 +1,57 @@
use serde_json::{Map, Value};
use crate::CoreResult;
use super::types::{AudioTranscriptionRequestData, AudioTranscriptionResponseData};
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum AudioTranscriptionAuth {
Bearer,
AwsSigV4 {
region: String,
service: &'static str,
},
}
pub trait AudioTranscriptionProviderConfig: Sync {
fn supported_transcription_params(&self) -> &'static [&'static str];
fn map_transcription_params(&self, params: &Map<String, Value>) -> Map<String, Value> {
params
.iter()
.filter(|(key, _)| {
self.supported_transcription_params()
.contains(&key.as_str())
})
.map(|(key, value)| (key.clone(), value.clone()))
.collect()
}
fn transform_transcription_request(
&self,
model: &str,
audio: Value,
optional_params: Map<String, Value>,
) -> CoreResult<AudioTranscriptionRequestData>;
fn transform_transcription_response(
&self,
model: &str,
response_json: Value,
) -> CoreResult<AudioTranscriptionResponseData>;
fn complete_url(
&self,
api_base: Option<&str>,
model: &str,
optional_params: &Map<String, Value>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> CoreResult<String>;
fn auth_strategy(
&self,
model: &str,
optional_params: &Map<String, Value>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> CoreResult<AudioTranscriptionAuth>;
}

View file

@ -0,0 +1,20 @@
use serde::{Deserialize, Serialize};
use serde_json::Value;
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct AudioTranscriptionRequestData {
pub body: Value,
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct AudioTranscriptionResponseData {
pub text: String,
}
impl AudioTranscriptionResponseData {
pub fn into_json(self) -> Value {
serde_json::json!({
"text": self.text,
})
}
}

View file

@ -0,0 +1,258 @@
use std::cmp::Reverse;
use std::collections::{BinaryHeap, HashMap};
use std::time::{Duration, SystemTime, UNIX_EPOCH};
const DEFAULT_MAX_SIZE_IN_MEMORY: usize = 200;
const DEFAULT_TTL: Duration = Duration::from_secs(600);
pub struct InMemoryCache<V: Clone> {
pub cache_dict: HashMap<String, V>,
pub ttl_dict: HashMap<String, Duration>,
pub expiration_heap: BinaryHeap<Reverse<(Duration, String)>>,
pub max_size_in_memory: usize,
pub default_ttl: Duration,
now: Box<dyn Fn() -> Duration + Send + Sync>,
}
impl<V: Clone> Default for InMemoryCache<V> {
fn default() -> Self {
Self::new(None, None)
}
}
impl<V: Clone> InMemoryCache<V> {
pub fn new(max_size_in_memory: Option<usize>, default_ttl: Option<Duration>) -> Self {
Self::with_clock(max_size_in_memory, default_ttl, || {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap_or_default()
})
}
pub fn with_clock(
max_size_in_memory: Option<usize>,
default_ttl: Option<Duration>,
now: impl Fn() -> Duration + Send + Sync + 'static,
) -> Self {
Self {
cache_dict: HashMap::new(),
ttl_dict: HashMap::new(),
expiration_heap: BinaryHeap::new(),
max_size_in_memory: max_size_in_memory.unwrap_or(DEFAULT_MAX_SIZE_IN_MEMORY),
default_ttl: default_ttl.unwrap_or(DEFAULT_TTL),
now: Box::new(now),
}
}
pub fn evict_cache(&mut self) {
if self.max_size_in_memory == 0 {
return;
}
let current_time = (self.now)();
while let Some(Reverse((expiration_time, key))) = self.expiration_heap.peek().cloned() {
if self.ttl_dict.get(&key).copied() != Some(expiration_time) {
self.expiration_heap.pop();
} else if expiration_time <= current_time {
self.expiration_heap.pop();
self.remove_key(&key);
} else {
break;
}
}
while self.cache_dict.len() >= self.max_size_in_memory {
let Some(Reverse((expiration_time, key))) = self.expiration_heap.pop() else {
break;
};
if self.ttl_dict.get(&key).copied() == Some(expiration_time) {
self.remove_key(&key);
}
}
}
pub fn allow_ttl_override(&self, key: &str) -> bool {
match self.ttl_dict.get(key).copied() {
None => true,
Some(expiration_time) => expiration_time < (self.now)(),
}
}
pub fn set_cache(&mut self, key: impl Into<String>, value: V, ttl: Option<Duration>) {
if self.max_size_in_memory == 0 {
return;
}
self.evict_cache();
let key = key.into();
self.cache_dict.insert(key.clone(), value);
if self.allow_ttl_override(&key) {
let expiration_time = (self.now)() + ttl.unwrap_or(self.default_ttl);
self.ttl_dict.insert(key.clone(), expiration_time);
self.expiration_heap.push(Reverse((expiration_time, key)));
}
}
// Generic values intentionally omit Python's per-item size check.
pub fn get_cache(&mut self, key: &str) -> Option<V> {
if self.cache_dict.contains_key(key) {
if self.is_key_expired(key) {
self.remove_key(key);
return None;
}
return self.cache_dict.get(key).cloned();
}
None
}
pub fn get_ttl(&self, key: &str) -> Option<Duration> {
self.ttl_dict.get(key).copied()
}
pub fn delete_cache(&mut self, key: &str) {
self.remove_key(key);
}
pub fn flush_cache(&mut self) {
self.cache_dict.clear();
self.ttl_dict.clear();
self.expiration_heap.clear();
}
fn is_key_expired(&self, key: &str) -> bool {
self.ttl_dict
.get(key)
.is_some_and(|expiration_time| *expiration_time < (self.now)())
}
fn remove_key(&mut self, key: &str) {
self.cache_dict.remove(key);
self.ttl_dict.remove(key);
}
}
#[cfg(test)]
mod tests {
use std::sync::{
Arc,
atomic::{AtomicU64, Ordering},
};
use super::InMemoryCache;
use std::time::Duration;
fn cache(now: Arc<AtomicU64>, max_size: usize, default_ttl: Duration) -> InMemoryCache<String> {
InMemoryCache::with_clock(Some(max_size), Some(default_ttl), move || {
Duration::from_secs(now.load(Ordering::Relaxed))
})
}
#[test]
fn ttl_expiry_is_deterministic() {
let now = Arc::new(AtomicU64::new(100));
let mut cache = cache(now.clone(), 10, Duration::from_secs(60));
cache.set_cache("key", "value".to_string(), None);
assert_eq!(cache.get_cache("key"), Some("value".to_string()));
now.store(161, Ordering::Relaxed);
assert_eq!(cache.get_cache("key"), None);
assert_eq!(cache.get_ttl("key"), None);
}
#[test]
fn default_and_per_set_ttl_are_applied() {
let now = Arc::new(AtomicU64::new(100));
let mut cache = cache(now.clone(), 10, Duration::from_secs(60));
cache.set_cache("default", "value".to_string(), None);
cache.set_cache("custom", "value".to_string(), Some(Duration::from_secs(20)));
assert_eq!(cache.get_ttl("default"), Some(Duration::from_secs(160)));
assert_eq!(cache.get_ttl("custom"), Some(Duration::from_secs(120)));
}
#[test]
fn unexpired_entries_do_not_allow_ttl_override() {
let now = Arc::new(AtomicU64::new(100));
let mut cache = cache(now.clone(), 10, Duration::from_secs(60));
cache.set_cache("key", "first".to_string(), Some(Duration::from_secs(20)));
cache.set_cache("key", "second".to_string(), Some(Duration::from_secs(80)));
assert_eq!(cache.get_cache("key"), Some("second".to_string()));
assert_eq!(cache.get_ttl("key"), Some(Duration::from_secs(120)));
now.store(121, Ordering::Relaxed);
cache.set_cache("key", "third".to_string(), Some(Duration::from_secs(80)));
assert_eq!(cache.get_ttl("key"), Some(Duration::from_secs(201)));
}
#[test]
fn max_size_evicts_earliest_expiration() {
let now = Arc::new(AtomicU64::new(100));
let mut cache = cache(now, 2, Duration::from_secs(60));
cache.set_cache("early", "value".to_string(), Some(Duration::from_secs(10)));
cache.set_cache("late", "value".to_string(), Some(Duration::from_secs(20)));
cache.set_cache("new", "value".to_string(), Some(Duration::from_secs(30)));
assert_eq!(cache.get_cache("early"), None);
assert!(cache.get_cache("late").is_some());
assert!(cache.get_cache("new").is_some());
}
#[test]
fn expired_entries_are_evicted_before_live_entries() {
let now = Arc::new(AtomicU64::new(100));
let mut cache = cache(now.clone(), 3, Duration::from_secs(60));
cache.set_cache(
"expired-one",
"value".to_string(),
Some(Duration::from_secs(10)),
);
cache.set_cache(
"expired-two",
"value".to_string(),
Some(Duration::from_secs(20)),
);
cache.set_cache("live", "value".to_string(), Some(Duration::from_secs(100)));
now.store(121, Ordering::Relaxed);
cache.set_cache("new", "value".to_string(), Some(Duration::from_secs(100)));
assert_eq!(cache.get_cache("expired-one"), None);
assert_eq!(cache.get_cache("expired-two"), None);
assert!(cache.get_cache("live").is_some());
assert!(cache.get_cache("new").is_some());
}
#[test]
fn stale_heap_entries_are_skipped() {
let now = Arc::new(AtomicU64::new(100));
let mut cache = cache(now, 1, Duration::from_secs(60));
cache.set_cache(
"removed",
"value".to_string(),
Some(Duration::from_secs(10)),
);
cache.delete_cache("removed");
cache.set_cache("kept", "value".to_string(), Some(Duration::from_secs(20)));
cache.set_cache("new", "value".to_string(), Some(Duration::from_secs(30)));
assert_eq!(cache.get_cache("removed"), None);
assert_eq!(cache.get_cache("kept"), None);
assert!(cache.get_cache("new").is_some());
}
#[test]
fn delete_and_flush_remove_values_and_ttls() {
let now = Arc::new(AtomicU64::new(100));
let mut cache = cache(now, 10, Duration::from_secs(60));
cache.set_cache("one", "value".to_string(), None);
cache.set_cache("two", "value".to_string(), None);
cache.delete_cache("one");
assert_eq!(cache.get_cache("one"), None);
cache.flush_cache();
assert!(cache.cache_dict.is_empty());
assert!(cache.ttl_dict.is_empty());
assert!(cache.expiration_heap.is_empty());
}
#[test]
fn zero_max_size_does_not_cache() {
let now = Arc::new(AtomicU64::new(100));
let mut cache = cache(now, 0, Duration::from_secs(60));
cache.set_cache("key", "value".to_string(), None);
assert_eq!(cache.get_cache("key"), None);
assert!(cache.cache_dict.is_empty());
}
}

View file

@ -0,0 +1 @@
pub mod in_memory_cache;

View file

@ -0,0 +1,3 @@
pub const OPENAI_DEFAULT_API_BASE: &str = "https://api.openai.com";
pub const OPENAI_RESPONSES_DEFAULT_API_BASE: &str = "https://api.openai.com/v1";
pub const OPENAI_RESPONSES_PATH: &str = "/responses";

View file

@ -1,9 +1,13 @@
pub mod audio_transcription;
pub mod caching;
pub mod call_lifecycle;
pub mod constants;
pub mod error;
pub mod messages;
pub mod ocr;
pub mod providers;
pub mod realtime;
pub mod responses;
pub mod router;
pub mod routing_utils;

View file

@ -35,6 +35,10 @@ pub trait AnthropicMessagesProviderConfig: Sync {
MessagesAuthStrategy::Header("x-api-key")
}
fn accepts_bearer_auth(&self) -> bool {
false
}
fn default_headers(&self) -> &'static [(&'static str, &'static str)] {
&[
("anthropic-version", "2023-06-01"),

View file

@ -5,7 +5,7 @@ use crate::messages::types::{
MessageContent, SystemPrompt,
};
use crate::providers::anthropic::messages::transformation::{
non_empty, AnthropicMessagesConfig, ANTHROPIC_MESSAGES_CONFIG,
ANTHROPIC_MESSAGES_CONFIG, AnthropicMessagesConfig, non_empty,
};
use serde_json::{Map, Value};
@ -163,6 +163,10 @@ impl AnthropicMessagesProviderConfig for AzureAnthropicMessagesConfig {
self.anthropic.auth_strategy()
}
fn accepts_bearer_auth(&self) -> bool {
true
}
fn default_headers(&self) -> &'static [(&'static str, &'static str)] {
self.anthropic.default_headers()
}
@ -294,6 +298,11 @@ mod tests {
);
}
#[test]
fn accepts_bearer_auth_for_entra_id() {
assert!(AZURE_ANTHROPIC_MESSAGES_CONFIG.accepts_bearer_auth());
}
#[test]
fn default_headers_match_python() {
assert_eq!(

View file

@ -1,9 +1,9 @@
use std::collections::BTreeSet;
use crate::error::{json_type_name, CoreError, CoreResult};
use crate::error::{CoreError, CoreResult, json_type_name};
use crate::ocr::transformation::{OcrAuthStrategy, OcrProviderConfig, OcrResponseHandling};
use crate::ocr::types::{OcrRequestData, OcrResponseData};
use serde_json::{json, Map, Value};
use serde_json::{Map, Value, json};
use crate::providers::mistral::ocr::transformation::MISTRAL_OCR_CONFIG;
@ -206,11 +206,11 @@ pub fn complete_document_intelligence_url(
AZURE_DOCUMENT_INTELLIGENCE_API_VERSION
);
if let Some(pages) = optional_params.get("pages") {
if let Some(normalized) = normalize_pages_param(pages)? {
url.push_str("&pages=");
url.push_str(&normalized);
}
if let Some(pages) = optional_params.get("pages")
&& let Some(normalized) = normalize_pages_param(pages)?
{
url.push_str("&pages=");
url.push_str(&normalized);
}
Ok(url)
@ -231,7 +231,7 @@ fn document_url_from_mistral_document(document: &Value) -> CoreResult<&str> {
other => {
return Err(CoreError::InvalidRequest(format!(
"Invalid document type: {other}. Must be 'document_url' or 'image_url'"
)))
)));
}
};
object

View file

@ -0,0 +1,310 @@
use serde_json::{Map, Value, json};
use crate::audio_transcription::transformation::{
AudioTranscriptionAuth, AudioTranscriptionProviderConfig,
};
use crate::audio_transcription::types::{
AudioTranscriptionRequestData, AudioTranscriptionResponseData,
};
use crate::error::{CoreError, CoreResult, json_type_name};
use super::aws_base::AwsAuthConfig;
use super::constants::{
AWS_REGION, AWS_REGION_NAME, BEDROCK_RUNTIME_ENDPOINT_TEMPLATE, BEDROCK_SERVICE,
DEFAULT_BEDROCK_REGION,
};
const SUPPORTED_PARAMS: &[&str] = &["language", "prompt", "temperature", "response_format"];
pub static BEDROCK_AUDIO_TRANSCRIPTION_CONFIG: BedrockAudioTranscriptionConfig =
BedrockAudioTranscriptionConfig;
pub struct BedrockAudioTranscriptionConfig;
pub fn bedrock_model_id_and_region(model: &str) -> (String, Option<String>) {
let mut stripped = model;
for prefix in ["bedrock/converse/", "bedrock/", "converse/"] {
if let Some(value) = stripped.strip_prefix(prefix) {
stripped = value;
break;
}
}
let mut region = None;
if let Some((candidate, remainder)) = stripped.split_once('/')
&& is_bedrock_region(candidate)
{
region = Some(candidate.to_string());
stripped = remainder;
}
for prefix in ["nova-2/", "nova/"] {
if let Some(value) = stripped.strip_prefix(prefix) {
stripped = value;
break;
}
}
if region.is_none() {
region = stripped
.strip_prefix("arn:")
.and_then(|value| value.split(':').nth(3))
.filter(|value| !value.is_empty())
.map(str::to_string);
}
(stripped.to_string(), region)
}
fn is_bedrock_region(value: &str) -> bool {
value.len() > 3
&& value.contains('-')
&& value
.chars()
.all(|char| char.is_ascii_alphanumeric() || char == '-')
}
pub fn resolve_bedrock_region(
model_region: Option<&str>,
optional_params: &Map<String, Value>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> String {
if let Some(region) = optional_params
.get("aws_region_name")
.and_then(Value::as_str)
{
return region.to_string();
}
if let Some(region) = model_region {
return region.to_string();
}
env_lookup(AWS_REGION_NAME)
.or_else(|| env_lookup(AWS_REGION))
.unwrap_or_else(|| DEFAULT_BEDROCK_REGION.to_string())
}
fn audio_fields(audio: Value) -> CoreResult<(String, String)> {
let object = audio.as_object().ok_or_else(|| CoreError::InvalidType {
expected: "object",
actual: json_type_name(&audio),
})?;
let data = object
.get("data")
.and_then(Value::as_str)
.filter(|value| !value.is_empty())
.ok_or(CoreError::MissingField("audio.data"))?;
let format = object
.get("format")
.and_then(Value::as_str)
.filter(|value| matches!(*value, "wav" | "mp3" | "flac" | "ogg"))
.ok_or_else(|| {
CoreError::InvalidRequest("audio.format must be wav, mp3, flac, or ogg".to_string())
})?;
Ok((data.to_string(), format.to_string()))
}
fn optional_string<'a>(params: &'a Map<String, Value>, key: &str) -> Option<&'a str> {
params
.get(key)
.and_then(Value::as_str)
.filter(|value| !value.is_empty())
}
impl AudioTranscriptionProviderConfig for BedrockAudioTranscriptionConfig {
fn supported_transcription_params(&self) -> &'static [&'static str] {
SUPPORTED_PARAMS
}
fn transform_transcription_request(
&self,
_model: &str,
audio: Value,
optional_params: Map<String, Value>,
) -> CoreResult<AudioTranscriptionRequestData> {
let (data, format) = audio_fields(audio)?;
let mut instruction = "Transcribe the audio. Respond with only the transcript.".to_string();
if let Some(language) = optional_string(&optional_params, "language") {
instruction.push_str(&format!(" The audio language is {language}."));
}
if let Some(prompt) = optional_string(&optional_params, "prompt") {
instruction.push_str(&format!(" Additional context: {prompt}"));
}
let mut inference_config = Map::from_iter([("maxTokens".to_string(), json!(4096))]);
if let Some(temperature) = optional_params.get("temperature") {
inference_config.insert("temperature".to_string(), temperature.clone());
}
Ok(AudioTranscriptionRequestData {
body: json!({
"messages": [{
"role": "user",
"content": [
{"audio": {"format": format, "source": {"bytes": data}}},
{"text": instruction}
]
}],
"system": [{"text": "You are a transcription assistant."}],
"inferenceConfig": inference_config,
}),
})
}
fn transform_transcription_response(
&self,
_model: &str,
response_json: Value,
) -> CoreResult<AudioTranscriptionResponseData> {
let content = response_json
.get("output")
.and_then(|value| value.get("message"))
.and_then(|value| value.get("content"))
.and_then(Value::as_array)
.ok_or_else(|| {
CoreError::InvalidResponse("Bedrock response has no output content".to_string())
})?;
let mut text = String::new();
for block in content {
if let Some(value) = block.get("text").and_then(Value::as_str) {
text.push_str(value);
}
}
Ok(AudioTranscriptionResponseData { text })
}
fn complete_url(
&self,
api_base: Option<&str>,
model: &str,
optional_params: &Map<String, Value>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> CoreResult<String> {
let (model_id, model_region) = bedrock_model_id_and_region(model);
let region = resolve_bedrock_region(model_region.as_deref(), optional_params, env_lookup);
let endpoint = optional_params
.get("aws_bedrock_runtime_endpoint")
.and_then(Value::as_str)
.or(api_base)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(str::to_string)
.unwrap_or_else(|| BEDROCK_RUNTIME_ENDPOINT_TEMPLATE.replace("{region}", &region));
Ok(format!(
"{}/model/{model_id}/converse",
endpoint.trim_end_matches('/')
))
}
fn auth_strategy(
&self,
model: &str,
optional_params: &Map<String, Value>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> CoreResult<AudioTranscriptionAuth> {
let (_, model_region) = bedrock_model_id_and_region(model);
Ok(AudioTranscriptionAuth::AwsSigV4 {
region: resolve_bedrock_region(model_region.as_deref(), optional_params, env_lookup),
service: BEDROCK_SERVICE,
})
}
}
pub fn aws_auth_config(
optional_params: &Map<String, Value>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> AwsAuthConfig {
let value = |key: &str| {
optional_params
.get(key)
.and_then(Value::as_str)
.map(str::to_string)
};
let env = |key: &str| env_lookup(key);
AwsAuthConfig {
access_key_id: value("aws_access_key_id").or_else(|| env("AWS_ACCESS_KEY_ID")),
secret_access_key: value("aws_secret_access_key").or_else(|| env("AWS_SECRET_ACCESS_KEY")),
session_token: value("aws_session_token").or_else(|| env("AWS_SESSION_TOKEN")),
region_name: value("aws_region_name").or_else(|| env(AWS_REGION_NAME)),
session_name: value("aws_session_name").or_else(|| env("AWS_SESSION_NAME")),
profile_name: value("aws_profile_name").or_else(|| env("AWS_PROFILE_NAME")),
role_name: value("aws_role_name").or_else(|| env("AWS_ROLE_NAME")),
web_identity_token: value("aws_web_identity_token")
.or_else(|| env("AWS_WEB_IDENTITY_TOKEN")),
sts_endpoint: value("aws_sts_endpoint").or_else(|| env("AWS_STS_ENDPOINT")),
external_id: value("aws_external_id").or_else(|| env("AWS_EXTERNAL_ID")),
}
}
#[cfg(test)]
mod tests {
use super::*;
fn no_env(_: &str) -> Option<String> {
None
}
#[test]
fn request_matches_python_shape() {
let params = Map::from_iter([
("language".to_string(), json!("en")),
("prompt".to_string(), json!("Speaker names")),
("temperature".to_string(), json!(0)),
("timestamp_granularities".to_string(), json!(["word"])),
]);
let params = BEDROCK_AUDIO_TRANSCRIPTION_CONFIG.map_transcription_params(&params);
let result = BEDROCK_AUDIO_TRANSCRIPTION_CONFIG
.transform_transcription_request(
"mistral.voxtral-mini-3b-2507",
json!({"data": "AQI=", "format": "wav", "filename": "sample.wav"}),
params,
)
.expect("request");
assert_eq!(
result.body,
json!({
"messages": [{
"role": "user",
"content": [
{"audio": {"format": "wav", "source": {"bytes": "AQI="}}},
{"text": "Transcribe the audio. Respond with only the transcript. The audio language is en. Additional context: Speaker names"}
]
}],
"system": [{"text": "You are a transcription assistant."}],
"inferenceConfig": {"maxTokens": 4096, "temperature": 0}
})
);
}
#[test]
fn response_concatenates_content_blocks() {
let result = BEDROCK_AUDIO_TRANSCRIPTION_CONFIG
.transform_transcription_response(
"model",
json!({"output": {"message": {"content": [{"text": "hello "}, {"text": "world"}]}}}),
)
.expect("response");
assert_eq!(result.text, "hello world");
assert_eq!(result.into_json(), json!({"text": "hello world"}));
}
#[test]
fn invalid_audio_is_rejected() {
let result = BEDROCK_AUDIO_TRANSCRIPTION_CONFIG.transform_transcription_request(
"model",
json!({"data": "AQI="}),
Map::new(),
);
assert!(result.is_err());
}
#[test]
fn region_and_url_precedence_match_python() {
let params = Map::from_iter([("aws_region_name".to_string(), json!("eu-west-1"))]);
let url = BEDROCK_AUDIO_TRANSCRIPTION_CONFIG
.complete_url(
None,
"bedrock/us-east-1/mistral.voxtral-mini-3b-2507",
&params,
&no_env,
)
.expect("url");
assert_eq!(
url,
"https://bedrock-runtime.eu-west-1.amazonaws.com/model/mistral.voxtral-mini-3b-2507/converse"
);
}
}

View file

@ -0,0 +1,726 @@
use std::collections::BTreeMap;
use std::sync::{Mutex, OnceLock};
use std::time::Duration;
use std::time::{SystemTime, UNIX_EPOCH};
use crate::caching::in_memory_cache::InMemoryCache;
use crate::error::{CoreError, CoreResult};
use aws_credential_types::Credentials;
use aws_credential_types::provider::ProvideCredentials;
use aws_sigv4::http_request::{
SignableBody, SignableRequest, SigningParams, SigningSettings, sign,
};
use aws_sigv4::sign::v4;
use aws_smithy_runtime_api::client::identity::Identity;
use sha2::{Digest, Sha256};
use super::constants::{
AWS_ACCESS_KEY_ID, AWS_EXTERNAL_ID, AWS_PROFILE_NAME, AWS_REGION_NAME, AWS_ROLE_ARN,
AWS_ROLE_NAME, AWS_SECRET_ACCESS_KEY, AWS_SESSION_NAME, AWS_SESSION_TOKEN, AWS_STS_ENDPOINT,
AWS_WEB_IDENTITY_TOKEN, AWS_WEB_IDENTITY_TOKEN_FILE, BEDROCK_SERVICE,
DEFAULT_SESSION_NAME_PREFIX,
};
const STATIC_CREDENTIALS_TTL: Duration = Duration::from_secs(3600 - 60);
const AMBIENT_CREDENTIALS_TTL: Duration = Duration::from_secs(600);
static IAM_CREDENTIALS_CACHE: OnceLock<Mutex<InMemoryCache<Credentials>>> = OnceLock::new();
fn credential_cache_ttl(flow: &AwsAuthFlow) -> Option<Duration> {
match flow {
AwsAuthFlow::StaticKeys { .. } => Some(STATIC_CREDENTIALS_TTL),
AwsAuthFlow::DefaultChain => Some(AMBIENT_CREDENTIALS_TTL),
AwsAuthFlow::WebIdentity { .. }
| AwsAuthFlow::AssumeRole { .. }
| AwsAuthFlow::Profile { .. }
| AwsAuthFlow::SessionToken { .. } => None,
}
}
#[derive(Clone, Debug, Default, PartialEq, Eq)]
pub struct AwsAuthConfig {
pub access_key_id: Option<String>,
pub secret_access_key: Option<String>,
pub session_token: Option<String>,
pub region_name: Option<String>,
pub session_name: Option<String>,
pub profile_name: Option<String>,
pub role_name: Option<String>,
pub web_identity_token: Option<String>,
pub sts_endpoint: Option<String>,
pub external_id: Option<String>,
}
impl AwsAuthConfig {
fn with_environment(self, env_lookup: &(dyn Fn(&str) -> Option<String> + Sync)) -> Self {
Self {
access_key_id: self.access_key_id.or_else(|| env_lookup(AWS_ACCESS_KEY_ID)),
secret_access_key: self
.secret_access_key
.or_else(|| env_lookup(AWS_SECRET_ACCESS_KEY)),
session_token: self.session_token.or_else(|| env_lookup(AWS_SESSION_TOKEN)),
region_name: self.region_name.or_else(|| env_lookup(AWS_REGION_NAME)),
session_name: self.session_name.or_else(|| env_lookup(AWS_SESSION_NAME)),
profile_name: self.profile_name.or_else(|| env_lookup(AWS_PROFILE_NAME)),
role_name: self.role_name.or_else(|| env_lookup(AWS_ROLE_NAME)),
web_identity_token: self
.web_identity_token
.or_else(|| env_lookup(AWS_WEB_IDENTITY_TOKEN)),
sts_endpoint: self.sts_endpoint.or_else(|| env_lookup(AWS_STS_ENDPOINT)),
external_id: self.external_id.or_else(|| env_lookup(AWS_EXTERNAL_ID)),
}
}
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum AwsAuthFlow {
WebIdentity {
token: String,
role: String,
session_name: String,
},
AssumeRole {
role: String,
session_name: Option<String>,
},
Profile {
name: String,
},
SessionToken {
access_key_id: String,
secret_access_key: String,
session_token: String,
},
StaticKeys {
access_key_id: String,
secret_access_key: String,
region_name: String,
},
DefaultChain,
}
fn cache_key(config: &AwsAuthConfig, flow: &AwsAuthFlow) -> String {
let mut hasher = Sha256::new();
hasher.update(format!("{config:?}:{flow:?}"));
format!("{:x}", hasher.finalize())
}
fn get_cached_credentials(key: &str) -> Option<Credentials> {
let cache = IAM_CREDENTIALS_CACHE.get_or_init(|| Mutex::new(InMemoryCache::default()));
let mut entries = cache.lock().ok()?;
entries.get_cache(key)
}
fn set_cached_credentials(key: String, credentials: Credentials, ttl: Duration) {
let cache = IAM_CREDENTIALS_CACHE.get_or_init(|| Mutex::new(InMemoryCache::default()));
if let Ok(mut entries) = cache.lock() {
entries.set_cache(key, credentials, Some(ttl));
}
}
fn role_identity(arn: &str) -> Option<(&str, &str, &str)> {
let mut parts = arn.splitn(6, ':');
let ("arn", partition, _, _, account, resource) = (
parts.next()?,
parts.next()?,
parts.next()?,
parts.next()?,
parts.next()?,
parts.next()?,
) else {
return None;
};
let role = if let Some(role) = resource.strip_prefix("role/") {
role.rsplit('/').next()?
} else {
resource.strip_prefix("assumed-role/")?.split('/').next()?
};
Some((partition, account, role))
}
fn same_role_arns(target: &str, caller: &str) -> bool {
role_identity(target) == role_identity(caller)
}
pub fn classify_auth(
config: AwsAuthConfig,
env_lookup: &(dyn Fn(&str) -> Option<String> + Sync),
) -> AwsAuthFlow {
let config = config.with_environment(env_lookup);
if let (Some(token), Some(role), Some(session_name)) = (
config.web_identity_token.clone(),
config.role_name.clone(),
config.session_name.clone(),
) {
return AwsAuthFlow::WebIdentity {
token,
role,
session_name,
};
}
if let Some(role) = config.role_name.clone() {
return AwsAuthFlow::AssumeRole {
role,
session_name: config.session_name.clone(),
};
}
if let Some(name) = config.profile_name {
return AwsAuthFlow::Profile { name };
}
if let (Some(access_key_id), Some(secret_access_key), Some(session_token)) = (
config.access_key_id.clone(),
config.secret_access_key.clone(),
config.session_token,
) {
return AwsAuthFlow::SessionToken {
access_key_id,
secret_access_key,
session_token,
};
}
if let (Some(access_key_id), Some(secret_access_key), Some(region_name)) = (
config.access_key_id,
config.secret_access_key,
config.region_name,
) {
return AwsAuthFlow::StaticKeys {
access_key_id,
secret_access_key,
region_name,
};
}
AwsAuthFlow::DefaultChain
}
pub async fn resolve_credentials(
config: AwsAuthConfig,
env_lookup: &(dyn Fn(&str) -> Option<String> + Sync),
) -> CoreResult<Credentials> {
let resolved = config.clone().with_environment(env_lookup);
let flow = classify_auth(config, env_lookup);
match flow {
AwsAuthFlow::SessionToken {
access_key_id,
secret_access_key,
session_token,
} => Ok(Credentials::new(
access_key_id,
secret_access_key,
Some(session_token),
None,
"litellm-static-session",
)),
AwsAuthFlow::StaticKeys {
access_key_id,
secret_access_key,
region_name,
} => {
let flow = AwsAuthFlow::StaticKeys {
access_key_id: access_key_id.clone(),
secret_access_key: secret_access_key.clone(),
region_name,
};
let key = cache_key(&resolved, &flow);
if let Some(credentials) = get_cached_credentials(&key) {
return Ok(credentials);
}
let credentials = Credentials::new(
access_key_id,
secret_access_key,
None,
None,
"litellm-static",
);
set_cached_credentials(
key,
credentials.clone(),
credential_cache_ttl(&flow).unwrap_or(STATIC_CREDENTIALS_TTL),
);
Ok(credentials)
}
AwsAuthFlow::Profile { name } => {
let provider = aws_config::profile::ProfileFileCredentialsProvider::builder()
.profile_name(name)
.build();
provider.provide_credentials().await.map_err(|error| {
CoreError::Auth(format!("AWS profile credentials failed: {error}"))
})
}
AwsAuthFlow::AssumeRole { role, session_name } => {
if is_already_running_as_role(&role, &resolved).await? {
let ambient_flow = AwsAuthFlow::DefaultChain;
let key = cache_key(&resolved, &ambient_flow);
if let Some(credentials) = get_cached_credentials(&key) {
return Ok(credentials);
}
let provider =
aws_config::default_provider::credentials::DefaultCredentialsChain::builder()
.build()
.await;
let credentials = provider.provide_credentials().await.map_err(|error| {
CoreError::Auth(format!("AWS default credentials failed: {error}"))
})?;
set_cached_credentials(
key,
credentials.clone(),
credential_cache_ttl(&ambient_flow).unwrap_or(AMBIENT_CREDENTIALS_TTL),
);
return Ok(credentials);
}
let mut loader = aws_config::defaults(aws_config::BehaviorVersion::latest());
if let Some(region) = resolved.region_name.clone() {
loader = loader.region(aws_types::region::Region::new(region));
}
if let Some(endpoint) = resolved.sts_endpoint.clone() {
loader = loader.endpoint_url(endpoint);
}
if let (Some(access_key_id), Some(secret_access_key)) =
(resolved.access_key_id, resolved.secret_access_key)
{
loader = loader.credentials_provider(Credentials::new(
access_key_id,
secret_access_key,
resolved.session_token,
None,
"litellm-role-source",
));
}
let sdk_config = loader.load().await;
let builder = aws_config::sts::AssumeRoleProvider::builder(role);
let builder = match session_name {
Some(name) => builder.session_name(name),
None => builder.session_name(default_session_name()),
};
let builder = match resolved.external_id {
Some(id) => builder.external_id(id),
None => builder,
};
let provider = builder.configure(&sdk_config).build().await;
provider
.provide_credentials()
.await
.map_err(|error| CoreError::Auth(format!("AWS role credentials failed: {error}")))
}
AwsAuthFlow::WebIdentity {
token,
role,
session_name,
} => {
let mut loader = aws_config::defaults(aws_config::BehaviorVersion::latest());
if let Some(region) = resolved.region_name {
loader = loader.region(aws_types::region::Region::new(region));
}
if let Some(endpoint) = resolved.sts_endpoint {
loader = loader.endpoint_url(endpoint);
}
let sdk_config = loader.load().await;
let client = aws_sdk_sts::Client::new(&sdk_config);
let response = client
.assume_role_with_web_identity()
.role_arn(role)
.role_session_name(session_name)
.web_identity_token(token)
.send()
.await
.map_err(|error| {
CoreError::Auth(format!("AWS web identity credentials failed: {error}"))
})?;
let credentials = response.credentials().ok_or_else(|| {
CoreError::Auth("AWS web identity response had no credentials".to_string())
})?;
let expiration = SystemTime::try_from(*credentials.expiration()).map_err(|error| {
CoreError::Auth(format!("AWS web identity expiration was invalid: {error}"))
})?;
Ok(Credentials::new(
credentials.access_key_id(),
credentials.secret_access_key(),
Some(credentials.session_token().to_string()),
Some(expiration),
"litellm-web-identity",
))
}
AwsAuthFlow::DefaultChain => {
let key = cache_key(&resolved, &AwsAuthFlow::DefaultChain);
if let Some(credentials) = get_cached_credentials(&key) {
return Ok(credentials);
}
let provider =
aws_config::default_provider::credentials::DefaultCredentialsChain::builder()
.build()
.await;
let credentials = provider.provide_credentials().await.map_err(|error| {
CoreError::Auth(format!("AWS default credentials failed: {error}"))
})?;
set_cached_credentials(
key,
credentials.clone(),
credential_cache_ttl(&AwsAuthFlow::DefaultChain).unwrap_or(AMBIENT_CREDENTIALS_TTL),
);
Ok(credentials)
}
}
}
async fn is_already_running_as_role(role: &str, config: &AwsAuthConfig) -> CoreResult<bool> {
if role_identity(role).is_none() {
return Ok(false);
}
if let (Ok(current_role), Ok(token_file)) = (
std::env::var(AWS_ROLE_ARN),
std::env::var(AWS_WEB_IDENTITY_TOKEN_FILE),
) && !token_file.is_empty()
{
return Ok(same_role_arns(role, &current_role));
}
let mut loader = aws_config::defaults(aws_config::BehaviorVersion::latest());
if let Some(region) = config.region_name.clone() {
loader = loader.region(aws_types::region::Region::new(region));
}
if let Some(endpoint) = config.sts_endpoint.clone() {
loader = loader.endpoint_url(endpoint);
}
let sdk_config = loader.load().await;
let response = match aws_sdk_sts::Client::new(&sdk_config)
.get_caller_identity()
.send()
.await
{
Ok(response) => response,
Err(_) => return Ok(false),
};
Ok(response
.arn()
.is_some_and(|caller| same_role_arns(role, caller)))
}
fn default_session_name() -> String {
let seconds = SystemTime::now()
.duration_since(UNIX_EPOCH)
.map_or(0, |duration| duration.as_secs());
format!("{DEFAULT_SESSION_NAME_PREFIX}-{seconds}")
}
pub fn sign_bedrock_post(
url: &str,
body: &[u8],
headers: &BTreeMap<String, String>,
region: &str,
credentials: &Credentials,
signing_time: SystemTime,
) -> CoreResult<BTreeMap<String, String>> {
let identity: Identity = credentials.clone().into();
let params = v4::SigningParams::builder()
.identity(&identity)
.region(region)
.name(BEDROCK_SERVICE)
.time(signing_time)
.settings(SigningSettings::default())
.build()
.map(SigningParams::from)
.map_err(|error| CoreError::Auth(format!("AWS signing parameters failed: {error}")))?;
let header_refs = headers
.iter()
.map(|(name, value)| (name.as_str(), value.as_str()));
let request = SignableRequest::new("POST", url, header_refs, SignableBody::Bytes(body))
.map_err(|error| CoreError::Auth(format!("AWS signable request failed: {error}")))?;
let (instructions, _) = sign(request, &params)
.map_err(|error| CoreError::Auth(format!("AWS request signing failed: {error}")))?
.into_parts();
Ok(instructions
.headers()
.map(|(name, value)| {
let normalized_name = match name {
"authorization" => "Authorization",
"x-amz-date" => "X-Amz-Date",
"x-amz-security-token" => "X-Amz-Security-Token",
_ => name,
};
(normalized_name.to_string(), value.to_string())
})
.collect())
}
#[cfg(test)]
mod tests {
use super::*;
fn no_env(_: &str) -> Option<String> {
None
}
fn parity_inputs() -> (String, Vec<u8>, BTreeMap<String, String>) {
(
"https://bedrock-runtime.us-east-1.amazonaws.com/model/amazon.titan-text-express-v1/invoke"
.to_string(),
br#"{"input":"hello"}"#.to_vec(),
BTreeMap::from([("Content-Type".to_string(), "application/json".to_string())]),
)
}
#[test]
fn classification_preserves_python_precedence() {
let config = AwsAuthConfig {
access_key_id: Some("ak".into()),
secret_access_key: Some("sk".into()),
session_token: Some("token".into()),
region_name: Some("us-east-1".into()),
session_name: Some("session".into()),
profile_name: Some("profile".into()),
role_name: Some("role".into()),
web_identity_token: Some("oidc".into()),
..Default::default()
};
assert!(matches!(
classify_auth(config, &no_env),
AwsAuthFlow::WebIdentity { .. }
));
}
#[test]
fn classification_covers_fallthroughs() {
let env = |key: &str| match key {
AWS_PROFILE_NAME => Some("profile".into()),
_ => None,
};
assert!(matches!(
classify_auth(AwsAuthConfig::default(), &env),
AwsAuthFlow::Profile { .. }
));
assert!(matches!(
classify_auth(
AwsAuthConfig {
access_key_id: Some("ak".into()),
secret_access_key: Some("sk".into()),
session_token: Some("token".into()),
..Default::default()
},
&no_env
),
AwsAuthFlow::SessionToken { .. }
));
assert!(matches!(
classify_auth(
AwsAuthConfig {
access_key_id: Some("ak".into()),
secret_access_key: Some("sk".into()),
region_name: Some("us-east-1".into()),
..Default::default()
},
&no_env
),
AwsAuthFlow::StaticKeys { .. }
));
assert_eq!(
classify_auth(AwsAuthConfig::default(), &no_env),
AwsAuthFlow::DefaultChain
);
}
#[tokio::test]
async fn static_credentials_do_not_use_network() {
let credentials = resolve_credentials(
AwsAuthConfig {
access_key_id: Some("ak".into()),
secret_access_key: Some("sk".into()),
region_name: Some("us-east-1".into()),
..Default::default()
},
&no_env,
)
.await
.expect("static credentials");
assert_eq!(credentials.access_key_id(), "ak");
assert_eq!(credentials.session_token(), None);
}
#[test]
fn cache_policy_matches_python_flows() {
assert_eq!(
credential_cache_ttl(&AwsAuthFlow::StaticKeys {
access_key_id: "ak".into(),
secret_access_key: "sk".into(),
region_name: "us-east-1".into(),
}),
Some(STATIC_CREDENTIALS_TTL)
);
assert_eq!(
credential_cache_ttl(&AwsAuthFlow::DefaultChain),
Some(AMBIENT_CREDENTIALS_TTL)
);
assert_eq!(
credential_cache_ttl(&AwsAuthFlow::SessionToken {
access_key_id: "ak".into(),
secret_access_key: "sk".into(),
session_token: "token".into(),
}),
None
);
assert_eq!(
credential_cache_ttl(&AwsAuthFlow::Profile {
name: "profile".into()
}),
None
);
assert_eq!(
credential_cache_ttl(&AwsAuthFlow::AssumeRole {
role: "arn:aws:iam::123456789012:role/demo".into(),
session_name: None,
}),
None
);
assert_eq!(
credential_cache_ttl(&AwsAuthFlow::WebIdentity {
token: "token".into(),
role: "arn:aws:iam::123456789012:role/demo".into(),
session_name: "session".into(),
}),
None
);
}
#[test]
fn cache_round_trip_preserves_credentials() {
let key = format!("cache-test-{}", std::process::id());
let credentials = Credentials::new("cache-ak", "cache-sk", None, None, "test");
set_cached_credentials(key.clone(), credentials.clone(), STATIC_CREDENTIALS_TTL);
assert_eq!(
get_cached_credentials(&key).map(|value| value.access_key_id().to_string()),
Some("cache-ak".to_string())
);
}
#[test]
fn same_role_comparison_matches_partition_account_and_role() {
assert!(same_role_arns(
"arn:aws:iam::123456789012:role/path/demo",
"arn:aws:sts::123456789012:assumed-role/demo/session"
));
assert!(!same_role_arns(
"arn:aws:iam::123456789012:role/demo",
"arn:aws:iam::999999999999:role/demo"
));
assert!(!same_role_arns(
"arn:aws:iam::123456789012:role/demo",
"arn:aws-cn:iam::123456789012:role/demo"
));
assert!(!same_role_arns(
"arn:aws:iam::123456789012:user/demo",
"arn:aws:iam::123456789012:role/demo"
));
}
#[test]
fn signing_matches_botocore_golden_vector() {
let (url, body, headers) = parity_inputs();
let credentials = Credentials::new(
"AKIDEXAMPLE",
"wJalrXUtnFEMI/K7MDENG+bPxRfiCYEXAMPLEKEY",
Some("session-token".to_string()),
None,
"test",
);
let signed = sign_bedrock_post(
&url,
&body,
&headers,
"us-east-1",
&credentials,
UNIX_EPOCH + std::time::Duration::from_secs(1_704_164_645),
)
.expect("golden signature");
assert_eq!(
signed.get("X-Amz-Date").map(String::as_str),
Some("20240102T030405Z")
);
assert_eq!(
signed.get("X-Amz-Security-Token").map(String::as_str),
Some("session-token")
);
assert_eq!(
signed.get("Authorization").map(String::as_str),
Some(
"AWS4-HMAC-SHA256 Credential=AKIDEXAMPLE/20240102/us-east-1/bedrock/aws4_request, SignedHeaders=content-type;host;x-amz-date;x-amz-security-token, Signature=55c027ef47527d3ad63f1735f9d099efdbc99f296ff914bd94e727e24ec0e464"
)
);
}
#[test]
fn signing_without_session_token_omits_security_header() {
let (url, body, headers) = parity_inputs();
let credentials = Credentials::new(
"AKIDEXAMPLE",
"wJalrXUtnFEMI/K7MDENG+bPxRfiCYEXAMPLEKEY",
None,
None,
"test",
);
let signed = sign_bedrock_post(
&url,
&body,
&headers,
"us-east-1",
&credentials,
UNIX_EPOCH + std::time::Duration::from_secs(1_704_164_645),
)
.expect("signature");
assert!(!signed.contains_key("X-Amz-Security-Token"));
}
#[ignore]
#[tokio::test]
async fn live_bedrock_invoke_model_returns_200() -> Result<(), Box<dyn std::error::Error>> {
let access_key_id = std::env::var("AWS_BEDROCK_TEST_ACCESS_KEY_ID")?;
let secret_access_key = std::env::var("AWS_BEDROCK_TEST_SECRET_ACCESS_KEY")?;
let body = br#"{"anthropic_version":"bedrock-2023-05-31","max_tokens":1,"messages":[{"role":"user","content":[{"type":"text","text":"ping"}]}]}"#.to_vec();
let headers =
BTreeMap::from([("Content-Type".to_string(), "application/json".to_string())]);
let credentials = resolve_credentials(
AwsAuthConfig {
access_key_id: Some(access_key_id),
secret_access_key: Some(secret_access_key),
region_name: Some("us-west-2".to_string()),
..Default::default()
},
&no_env,
)
.await?;
let client = reqwest::Client::new();
let mut failures = Vec::new();
for region in ["us-west-2", "us-east-1"] {
let url = format!(
"https://bedrock-runtime.{region}.amazonaws.com/model/us.anthropic.claude-opus-4-8/invoke"
);
let signed_headers = sign_bedrock_post(
&url,
&body,
&headers,
region,
&credentials,
SystemTime::now(),
)?;
let mut request = client.post(&url).body(body.clone());
for (name, value) in &headers {
request = request.header(name, value);
}
for (name, value) in signed_headers {
request = request.header(name, value);
}
let response = request.send().await?;
let status = response.status();
let response_body = response.text().await?;
let snippet: String = response_body.chars().take(240).collect();
println!("region={region} status={status} response={snippet}");
if status == reqwest::StatusCode::OK {
return Ok(());
}
failures.push(format!("{region}: {status} {snippet}"));
}
panic!(
"no Bedrock region returned HTTP 200: {}",
failures.join("; ")
);
}
}

View file

@ -0,0 +1,18 @@
pub const AWS_ACCESS_KEY_ID: &str = "AWS_ACCESS_KEY_ID";
pub const AWS_SECRET_ACCESS_KEY: &str = "AWS_SECRET_ACCESS_KEY";
pub const AWS_SESSION_TOKEN: &str = "AWS_SESSION_TOKEN";
pub const AWS_REGION_NAME: &str = "AWS_REGION_NAME";
pub const AWS_REGION: &str = "AWS_REGION";
pub const AWS_SESSION_NAME: &str = "AWS_SESSION_NAME";
pub const AWS_PROFILE_NAME: &str = "AWS_PROFILE_NAME";
pub const AWS_ROLE_NAME: &str = "AWS_ROLE_NAME";
pub const AWS_WEB_IDENTITY_TOKEN: &str = "AWS_WEB_IDENTITY_TOKEN";
pub const AWS_ROLE_ARN: &str = "AWS_ROLE_ARN";
pub const AWS_WEB_IDENTITY_TOKEN_FILE: &str = "AWS_WEB_IDENTITY_TOKEN_FILE";
pub const AWS_STS_ENDPOINT: &str = "AWS_STS_ENDPOINT";
pub const AWS_EXTERNAL_ID: &str = "AWS_EXTERNAL_ID";
pub const BEDROCK_SERVICE: &str = "bedrock";
pub const DEFAULT_SESSION_NAME_PREFIX: &str = "litellm-session";
pub const DEFAULT_BEDROCK_REGION: &str = "us-west-2";
pub const BEDROCK_RUNTIME_ENDPOINT_TEMPLATE: &str =
"https://bedrock-runtime.{region}.amazonaws.com";

View file

@ -0,0 +1,8 @@
//! User-directed exception: this base provider owns AWS auth I/O for parity
//! with Python's `BaseAWSLLM`; the broader core purity guidance is reconciled
//! separately.
#[cfg(feature = "bedrock-auth")]
pub mod audio_transcription;
pub mod aws_base;
mod constants;

View file

@ -1,4 +1,4 @@
use crate::error::{json_type_name, CoreError, CoreResult};
use crate::error::{CoreError, CoreResult, json_type_name};
use crate::ocr::transformation::OcrProviderConfig;
use crate::ocr::types::{OcrRequestData, OcrResponseData};
use serde_json::{Map, Value};

View file

@ -1,5 +1,7 @@
pub mod anthropic;
pub mod azure_ai;
#[cfg(feature = "bedrock-auth")]
pub mod bedrock;
pub mod mistral;
pub mod openai;
pub mod vertex_ai;

View file

@ -1 +1,2 @@
pub mod realtime;
pub mod responses;

View file

@ -1,6 +1,6 @@
use crate::CoreResult;
use crate::realtime::transformation::RealtimeProviderConfig;
use crate::realtime::types::{RealtimeEvent, RealtimeTransformResult};
use crate::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";

View file

@ -0,0 +1 @@
pub mod transformation;

View file

@ -0,0 +1,48 @@
use crate::CoreResult;
use crate::responses::types::{ResponsesWsEvent, ResponsesWsTransformResult};
use crate::responses::websocket::{ResponsesWebSocketProviderConfig, enforce_model};
pub struct OpenAIResponsesWsConfig;
pub const OPENAI_RESPONSES_WS_CONFIG: OpenAIResponsesWsConfig = OpenAIResponsesWsConfig;
impl ResponsesWebSocketProviderConfig for OpenAIResponsesWsConfig {
fn supports_native_websocket(&self) -> bool {
true
}
fn transform_ws_request(
&self,
event: &ResponsesWsEvent,
model: &str,
) -> CoreResult<ResponsesWsTransformResult> {
Ok(ResponsesWsTransformResult::passthrough(enforce_model(
event, model,
)))
}
fn transform_ws_response(
&self,
event: &ResponsesWsEvent,
_model: &str,
) -> CoreResult<ResponsesWsTransformResult> {
Ok(ResponsesWsTransformResult::passthrough(event.clone()))
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn openai_config_is_native_and_enforces_model() {
let event: ResponsesWsEvent =
serde_json::from_value(serde_json::json!({"type":"response.create"}))
.expect("valid event");
let result = OPENAI_RESPONSES_WS_CONFIG
.transform_ws_request(&event, "gpt-5")
.expect("valid transform");
assert_eq!(result.events[0].model(), Some("gpt-5"));
assert!(OPENAI_RESPONSES_WS_CONFIG.supports_native_websocket());
}
}

View file

@ -1,7 +1,7 @@
use crate::error::{json_type_name, CoreError, CoreResult};
use crate::error::{CoreError, CoreResult, json_type_name};
use crate::ocr::transformation::OcrProviderConfig;
use crate::ocr::types::{OcrRequestData, OcrResponseData};
use serde_json::{json, Map, Value};
use serde_json::{Map, Value, json};
use crate::providers::mistral::ocr::transformation::MISTRAL_OCR_CONFIG;
@ -140,7 +140,7 @@ fn document_content_item(document: &Value) -> CoreResult<Value> {
other => {
return Err(CoreError::InvalidRequest(format!(
"Unsupported document type: {other}. Expected 'image_url' or 'document_url'"
)))
)));
}
};
let url = object

View file

@ -1,5 +1,5 @@
use crate::realtime::types::{RealtimeEvent, RealtimeTransformResult};
use crate::CoreResult;
use crate::realtime::types::{RealtimeEvent, RealtimeTransformResult};
pub trait RealtimeProviderConfig {
/// Build the upstream WebSocket URL (e.g. `wss://api.openai.com/v1/realtime?model=…`).

View file

@ -0,0 +1,365 @@
use std::future::Future;
use std::pin::Pin;
use std::sync::Mutex;
use std::time::{SystemTime, UNIX_EPOCH};
use serde_json::Value;
use crate::call_lifecycle::{CallLifecycleContext, CallLifecycleHooks, CallLifecycleTiming};
use crate::responses::types::{ResponsesWsEvent, ResponsesWsEventType};
use crate::{CoreError, CoreResult};
#[derive(Clone, Debug, Default, PartialEq, Eq)]
pub struct ResponsesWsUsage {
pub prompt_tokens: u64,
pub completion_tokens: u64,
pub total_tokens: u64,
}
#[derive(Clone, Debug, Default, PartialEq, Eq)]
pub struct ResponsesWsMetadata {
pub user_api_key_hash: Option<String>,
pub user_api_key_user_id: Option<String>,
pub user_api_key_team_id: Option<String>,
}
#[derive(Clone, Debug, PartialEq)]
pub struct ResponsesWsLogPayload {
pub id: String,
pub litellm_call_id: String,
pub call_type: String,
pub model: String,
pub custom_llm_provider: String,
pub response_cost: f64,
pub usage: ResponsesWsUsage,
pub start_time: f64,
pub end_time: f64,
pub stream: bool,
pub metadata: ResponsesWsMetadata,
}
#[derive(Clone, Debug, PartialEq)]
pub enum ResponsesWsLogOutcome {
Success {
payload: ResponsesWsLogPayload,
callback: ResponsesWsCallbackPayload,
},
Failure {
payload: ResponsesWsLogPayload,
callback: ResponsesWsCallbackPayload,
error_message: String,
error_kind: String,
},
}
#[derive(Clone, Debug, PartialEq)]
pub struct ResponsesWsCallbackPayload {
pub object: String,
pub value: Value,
}
struct InstrumentationState {
litellm_call_id: String,
id: String,
model: String,
usage: ResponsesWsUsage,
start_time: f64,
end_time: f64,
metadata: ResponsesWsMetadata,
outcome: Option<ResponsesWsLogOutcome>,
}
pub struct ResponsesWsInstrumentation {
state: Mutex<InstrumentationState>,
}
impl ResponsesWsInstrumentation {
pub fn new(
litellm_call_id: impl Into<String>,
model: impl Into<String>,
metadata: ResponsesWsMetadata,
) -> Self {
let litellm_call_id = litellm_call_id.into();
let now = epoch_seconds();
Self {
state: Mutex::new(InstrumentationState {
id: litellm_call_id.clone(),
litellm_call_id,
model: model.into(),
usage: ResponsesWsUsage::default(),
start_time: now,
end_time: now,
metadata,
outcome: None,
}),
}
}
pub fn observe(&self, event: &ResponsesWsEvent) {
if !matches!(
event.event_type,
ResponsesWsEventType::ResponseCreated
| ResponsesWsEventType::ResponseCompleted
| ResponsesWsEventType::ResponseFailed
| ResponsesWsEventType::ResponseIncomplete
| ResponsesWsEventType::Error
) {
return;
}
let Ok(mut state) = self.state.lock() else {
return;
};
let Some(response) = event.data.get("response").and_then(Value::as_object) else {
return;
};
if let Some(id) = response
.get("id")
.and_then(Value::as_str)
.filter(|value| !value.is_empty())
{
state.id = id.to_string();
state.litellm_call_id = id.to_string();
}
if let Some(model) = response
.get("model")
.and_then(Value::as_str)
.filter(|value| !value.is_empty())
{
state.model = model.to_string();
}
let Some(usage) = response.get("usage").and_then(Value::as_object) else {
return;
};
if let Some(input) = usage.get("input_tokens").and_then(Value::as_u64) {
state.usage.prompt_tokens += input;
}
if let Some(output) = usage.get("output_tokens").and_then(Value::as_u64) {
state.usage.completion_tokens += output;
}
state.usage.total_tokens += usage
.get("total_tokens")
.and_then(Value::as_u64)
.unwrap_or_else(|| {
usage
.get("input_tokens")
.and_then(Value::as_u64)
.unwrap_or(0)
+ usage
.get("output_tokens")
.and_then(Value::as_u64)
.unwrap_or(0)
});
}
pub fn success_outcome(&self) -> ResponsesWsLogOutcome {
let mut state = self
.state
.lock()
.unwrap_or_else(|poisoned| poisoned.into_inner());
state.end_time = epoch_seconds();
ResponsesWsLogOutcome::Success {
payload: build_payload(&state),
callback: ResponsesWsCallbackPayload {
object: "responses_websocket".to_string(),
value: Value::Null,
},
}
}
pub fn failure_outcome(&self) -> ResponsesWsLogOutcome {
let mut state = self
.state
.lock()
.unwrap_or_else(|poisoned| poisoned.into_inner());
state.end_time = epoch_seconds();
ResponsesWsLogOutcome::Failure {
payload: build_payload(&state),
callback: ResponsesWsCallbackPayload {
object: "error".to_string(),
value: serde_json::json!({
"message": "Responses WebSocket session ended in failure",
"kind": "ResponsesWebSocketError",
}),
},
error_message: "Responses WebSocket session ended in failure".to_string(),
error_kind: "ResponsesWebSocketError".to_string(),
}
}
pub fn take_outcome(&self) -> Option<ResponsesWsLogOutcome> {
self.state
.lock()
.unwrap_or_else(|poisoned| poisoned.into_inner())
.outcome
.take()
}
pub fn take_or_build_outcome(&self, success: bool) -> ResponsesWsLogOutcome {
self.take_outcome().unwrap_or_else(|| {
if success {
self.success_outcome()
} else {
self.failure_outcome()
}
})
}
}
type LifecycleFuture<'a, T> = Pin<Box<dyn Future<Output = CoreResult<T>> + Send + 'a>>;
impl CallLifecycleHooks<(), (), ()> for ResponsesWsInstrumentation {
type PreCallFuture<'a> = LifecycleFuture<'a, ()>;
type DuringCallFuture<'a> = LifecycleFuture<'a, ()>;
type SuccessFuture<'a> = Pin<Box<dyn Future<Output = ()> + Send + 'a>>;
type FailureFuture<'a> = Pin<Box<dyn Future<Output = ()> + Send + 'a>>;
fn async_pre_call_hook<'a>(
&'a self,
_context: &'a CallLifecycleContext,
request: (),
) -> Self::PreCallFuture<'a> {
Box::pin(async move { Ok(request) })
}
fn async_during_call_hook<'a>(
&'a self,
_context: &'a CallLifecycleContext,
request: (),
) -> Self::DuringCallFuture<'a> {
Box::pin(async move { Ok(request) })
}
fn async_log_success_event<'a>(
&'a self,
_context: &'a CallLifecycleContext,
_response: &'a (),
_timing: &'a CallLifecycleTiming,
) -> Self::SuccessFuture<'a> {
Box::pin(async move {
let outcome = self.success_outcome();
if let Ok(mut state) = self.state.lock() {
state.outcome = Some(outcome);
}
})
}
fn async_log_failure_event<'a>(
&'a self,
_context: &'a CallLifecycleContext,
_error: &'a CoreError,
_timing: &'a CallLifecycleTiming,
) -> Self::FailureFuture<'a> {
Box::pin(async move {
let outcome = self.failure_outcome();
if let Ok(mut state) = self.state.lock() {
state.outcome = Some(outcome);
}
})
}
}
fn build_payload(state: &InstrumentationState) -> ResponsesWsLogPayload {
ResponsesWsLogPayload {
id: state.id.clone(),
litellm_call_id: state.litellm_call_id.clone(),
call_type: "responses_websocket".to_string(),
model: state.model.clone(),
custom_llm_provider: "openai".to_string(),
response_cost: 0.0,
usage: state.usage.clone(),
start_time: state.start_time,
end_time: state.end_time,
stream: true,
metadata: state.metadata.clone(),
}
}
fn epoch_seconds() -> f64 {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.map(|duration| duration.as_secs_f64())
.unwrap_or(0.0)
}
#[cfg(test)]
mod tests {
use super::*;
fn event(value: Value) -> ResponsesWsEvent {
serde_json::from_value(value).expect("valid Responses WebSocket event")
}
#[test]
fn accumulates_upstream_usage_and_identity() {
let instrumentation =
ResponsesWsInstrumentation::new("call-1", "gpt-5", ResponsesWsMetadata::default());
instrumentation.observe(&event(serde_json::json!({
"type": "response.completed",
"response": {
"id": "resp-1",
"model": "gpt-5-mini",
"usage": {
"input_tokens": 3,
"output_tokens": 5,
"total_tokens": 8
}
}
})));
let ResponsesWsLogOutcome::Success { payload, .. } = instrumentation.success_outcome()
else {
panic!("expected success outcome");
};
assert_eq!(payload.id, "resp-1");
assert_eq!(payload.model, "gpt-5-mini");
assert_eq!(payload.usage.prompt_tokens, 3);
assert_eq!(payload.usage.completion_tokens, 5);
assert_eq!(payload.usage.total_tokens, 8);
assert!(payload.end_time >= payload.start_time);
}
#[test]
fn builds_failure_payload_without_dispatching_callbacks() {
let instrumentation =
ResponsesWsInstrumentation::new("call-1", "gpt-5", ResponsesWsMetadata::default());
assert!(matches!(
instrumentation.failure_outcome(),
ResponsesWsLogOutcome::Failure { .. }
));
}
#[tokio::test]
async fn lifecycle_records_success_outcome_for_provider_completion() {
let instrumentation =
ResponsesWsInstrumentation::new("call-1", "gpt-5", ResponsesWsMetadata::default());
let result = crate::call_lifecycle::CallLifecycle::default()
.run(
crate::call_lifecycle::CallLifecycleContext::new(
"responses_websocket",
"gpt-5",
"openai",
"call-1",
),
(),
&instrumentation,
|_| async { Ok::<(), CoreError>(()) },
)
.await;
assert!(result.is_ok());
assert!(matches!(
instrumentation.take_outcome(),
Some(ResponsesWsLogOutcome::Success { .. })
));
}
#[test]
fn builds_outcome_when_lifecycle_did_not_record_one() {
let instrumentation =
ResponsesWsInstrumentation::new("call-1", "gpt-5", ResponsesWsMetadata::default());
assert!(matches!(
instrumentation.take_or_build_outcome(true),
ResponsesWsLogOutcome::Success { .. }
));
}
}

View file

@ -0,0 +1,3 @@
pub mod instrumentation;
pub mod types;
pub mod websocket;

View file

@ -0,0 +1,166 @@
use serde::{Deserialize, Deserializer, Serialize, Serializer};
use serde_json::{Map, Value};
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum ResponsesWsEventType {
ResponseCreate,
ResponseCreated,
ResponseCompleted,
ResponseFailed,
ResponseIncomplete,
Error,
Other(String),
}
impl ResponsesWsEventType {
pub fn as_str(&self) -> &str {
match self {
Self::ResponseCreate => "response.create",
Self::ResponseCreated => "response.created",
Self::ResponseCompleted => "response.completed",
Self::ResponseFailed => "response.failed",
Self::ResponseIncomplete => "response.incomplete",
Self::Error => "error",
Self::Other(value) => value,
}
}
}
impl Serialize for ResponsesWsEventType {
fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
where
S: Serializer,
{
serializer.serialize_str(self.as_str())
}
}
impl<'de> Deserialize<'de> for ResponsesWsEventType {
fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
where
D: Deserializer<'de>,
{
let value = String::deserialize(deserializer)?;
Ok(match value.as_str() {
"response.create" => Self::ResponseCreate,
"response.created" => Self::ResponseCreated,
"response.completed" => Self::ResponseCompleted,
"response.failed" => Self::ResponseFailed,
"response.incomplete" => Self::ResponseIncomplete,
"error" => Self::Error,
_ => Self::Other(value),
})
}
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct ResponsesWsEvent {
#[serde(rename = "type")]
pub event_type: ResponsesWsEventType,
#[serde(flatten)]
pub data: Map<String, Value>,
}
impl ResponsesWsEvent {
pub fn model(&self) -> Option<&str> {
let model = self.data.get("model").and_then(Value::as_str);
if model.is_some() {
return model;
}
self.data
.get("response")
.and_then(Value::as_object)
.and_then(|response| response.get("model"))
.and_then(Value::as_str)
}
pub fn is_response_create(&self) -> bool {
self.event_type == ResponsesWsEventType::ResponseCreate
}
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct ResponsesWsTransformResult {
pub events: Vec<ResponsesWsEvent>,
}
impl ResponsesWsTransformResult {
pub fn passthrough(event: ResponsesWsEvent) -> Self {
Self {
events: vec![event],
}
}
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct ResponsesErrorFrame {
#[serde(rename = "type")]
pub frame_type: &'static str,
pub error: ResponsesErrorBody,
}
impl ResponsesErrorFrame {
pub fn invalid_request(message: impl Into<String>) -> Self {
Self {
frame_type: "error",
error: ResponsesErrorBody {
error_type: "invalid_request_error",
message: message.into(),
},
}
}
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct ResponsesErrorBody {
#[serde(rename = "type")]
pub error_type: &'static str,
pub message: String,
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn event_type_round_trips_known_and_unknown_values() {
let known: ResponsesWsEventType =
serde_json::from_str("\"response.completed\"").expect("valid event type");
assert_eq!(known, ResponsesWsEventType::ResponseCompleted);
let unknown: ResponsesWsEventType =
serde_json::from_str("\"response.output_text.delta\"").expect("valid event type");
assert_eq!(
unknown,
ResponsesWsEventType::Other("response.output_text.delta".to_string())
);
}
#[test]
fn error_frame_matches_proxy_shape() {
let frame = ResponsesErrorFrame::invalid_request("missing model");
assert_eq!(
serde_json::to_value(frame).expect("serializable"),
serde_json::json!({
"type": "error",
"error": {
"type": "invalid_request_error",
"message": "missing model"
}
})
);
}
#[test]
fn model_reads_flat_and_nested_create_shapes() {
let flat: ResponsesWsEvent =
serde_json::from_value(serde_json::json!({"type":"response.create","model":"gpt-5"}))
.expect("valid event");
let nested: ResponsesWsEvent = serde_json::from_value(serde_json::json!({
"type":"response.create",
"response":{"model":"gpt-5-mini"}
}))
.expect("valid event");
assert_eq!(flat.model(), Some("gpt-5"));
assert_eq!(nested.model(), Some("gpt-5-mini"));
}
}

View file

@ -0,0 +1,188 @@
use crate::CoreResult;
use crate::constants::{OPENAI_RESPONSES_DEFAULT_API_BASE, OPENAI_RESPONSES_PATH};
use crate::responses::types::{ResponsesWsEvent, ResponsesWsEventType, ResponsesWsTransformResult};
pub trait ResponsesWebSocketProviderConfig: Sync {
fn supports_native_websocket(&self) -> bool {
false
}
fn model_in_websocket_url(&self) -> bool {
true
}
fn complete_websocket_url(&self, api_base: Option<&str>, model: &str) -> String {
complete_websocket_url(api_base, model, self.model_in_websocket_url())
}
fn transform_ws_request(
&self,
event: &ResponsesWsEvent,
model: &str,
) -> CoreResult<ResponsesWsTransformResult>;
fn transform_ws_response(
&self,
event: &ResponsesWsEvent,
model: &str,
) -> CoreResult<ResponsesWsTransformResult>;
}
pub fn complete_websocket_url(
api_base: Option<&str>,
model: &str,
model_in_websocket_url: bool,
) -> String {
let base = api_base
.map(str::trim)
.filter(|value| !value.is_empty())
.unwrap_or(OPENAI_RESPONSES_DEFAULT_API_BASE);
let (base_without_query, query) = base
.split_once('?')
.map_or((base, None), |(value, query)| (value, Some(query)));
let response_url = format!(
"{}{}",
base_without_query.trim_end_matches('/'),
OPENAI_RESPONSES_PATH
);
let scheme_flipped = if let Some(rest) = response_url.strip_prefix("https://") {
format!("wss://{rest}")
} else if let Some(rest) = response_url.strip_prefix("http://") {
format!("ws://{rest}")
} else {
response_url
};
let url = query.map_or(scheme_flipped.clone(), |value| {
format!("{scheme_flipped}?{value}")
});
if !model_in_websocket_url
|| query.is_some_and(|value| {
value
.split('&')
.any(|part| part.split('=').next() == Some("model"))
})
{
return url;
}
format!(
"{url}{}model={}",
if query.is_some() { "&" } else { "?" },
percent_encode(model)
)
}
fn percent_encode(value: &str) -> String {
value
.bytes()
.map(|byte| {
if byte.is_ascii_alphanumeric() || matches!(byte, b'-' | b'.' | b'_' | b'~') {
format!("{}", byte as char)
} else {
format!("%{byte:02X}")
}
})
.collect()
}
pub fn enforce_model(event: &ResponsesWsEvent, model: &str) -> ResponsesWsEvent {
if !event.is_response_create() {
return event.clone();
}
let mut enforced = event.clone();
let has_flat_model = enforced.data.contains_key("model");
if let Some(response) = enforced
.data
.get_mut("response")
.and_then(serde_json::Value::as_object_mut)
{
response.insert(
"model".to_string(),
serde_json::Value::String(model.to_string()),
);
if has_flat_model {
enforced.data.insert(
"model".to_string(),
serde_json::Value::String(model.to_string()),
);
}
} else {
enforced.data.insert(
"model".to_string(),
serde_json::Value::String(model.to_string()),
);
}
enforced
}
pub fn is_terminal_event(event_type: &ResponsesWsEventType) -> bool {
matches!(
event_type,
ResponsesWsEventType::ResponseCreated
| ResponsesWsEventType::ResponseCompleted
| ResponsesWsEventType::ResponseFailed
| ResponsesWsEventType::ResponseIncomplete
| ResponsesWsEventType::Error
)
}
#[cfg(test)]
mod tests {
use super::*;
fn event(value: serde_json::Value) -> ResponsesWsEvent {
serde_json::from_value(value).expect("valid event")
}
#[test]
fn url_construction_matches_python_defaults_and_query_behavior() {
assert_eq!(
complete_websocket_url(None, "gpt-5", true),
"wss://api.openai.com/v1/responses?model=gpt-5"
);
assert_eq!(
complete_websocket_url(Some("http://localhost:8080/"), "gpt 5", true),
"ws://localhost:8080/responses?model=gpt%205"
);
assert_eq!(
complete_websocket_url(Some("https://example.test/v1?foo=bar"), "gpt-5", true),
"wss://example.test/v1/responses?foo=bar&model=gpt-5"
);
assert_eq!(
complete_websocket_url(Some("https://example.test?model=existing"), "gpt-5", true),
"wss://example.test/responses?model=existing"
);
}
#[test]
fn enforce_model_overrides_flat_and_nested_values() {
let flat = enforce_model(
&event(serde_json::json!({"type":"response.create","model":"wrong"})),
"gpt-5",
);
assert_eq!(flat.model(), Some("gpt-5"));
let nested = enforce_model(
&event(serde_json::json!({
"type":"response.create",
"model":"wrong",
"response":{"model":"also-wrong"}
})),
"gpt-5",
);
assert_eq!(nested.model(), Some("gpt-5"));
assert_eq!(
nested
.data
.get("response")
.and_then(|value| value.get("model")),
Some(&serde_json::json!("gpt-5"))
);
let nested_without_flat = enforce_model(
&event(serde_json::json!({
"type":"response.create",
"response":{"model":"also-wrong"}
})),
"gpt-5",
);
assert!(!nested_without_flat.data.contains_key("model"));
}
}

View file

@ -17,7 +17,13 @@ Python-compatible dictionaries.
- 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.
decides whether to raise or fall back. For a rust-only provider/route (no
Python reference), the Python side is a thin dispatch that calls Rust and
raises when the bridge is unavailable, with no fallback.
- Keep the Python interface minimal (well under 100 lines per route): it only
marshals inputs and calls Rust. Do not add per-route feature flags, and do
not put provider dispatch in `litellm/main.py`; it lives in a thin dispatch
class under `litellm/llms/<provider>/<route>/`.
## Data Handling

View file

@ -10,7 +10,7 @@ name = "_native"
crate-type = ["cdylib"]
[dependencies]
litellm-core.workspace = true
litellm-core = { workspace = true, features = ["bedrock-auth"] }
litellm-ai-gateway = { workspace = true, default-features = false }
pyo3 = { workspace = true, features = ["extension-module"] }
pyo3-async-runtimes.workspace = true

View file

@ -1,7 +1,12 @@
use std::collections::HashMap;
use std::time::Duration;
use litellm_ai_gateway::io::messages::{messages as run_messages, MessagesRequest};
use litellm_ai_gateway::io::ocr::{ocr as run_ocr, OcrRequest};
use litellm_ai_gateway::io::audio_transcription::{
AudioTranscriptionRequest, audio_transcription as run_audio_transcription,
};
use litellm_ai_gateway::io::messages::{MessagesRequest, messages as run_messages};
use litellm_ai_gateway::io::ocr::{OcrRequest, ocr as run_ocr};
use litellm_ai_gateway::io::responses_ws::ResponsesWebSocketConnection as RustResponsesWebSocketConnection;
use litellm_core::error::CoreError;
use pyo3::exceptions::{PyRuntimeError, PyValueError};
use pyo3::prelude::*;
@ -65,6 +70,76 @@ fn optional_timeout(timeout_seconds: Option<f64>) -> Option<Duration> {
})
}
fn marshal_headers(
py: Python<'_>,
headers: Option<Py<PyAny>>,
) -> PyResult<HashMap<String, String>> {
let value = match headers {
Some(headers) => py_to_json(py, headers.bind(py))?,
None => Value::Object(Map::new()),
};
let Value::Object(headers) = value else {
return Err(PyValueError::new_err("headers must be a dict"));
};
headers
.into_iter()
.map(|(name, value)| {
value
.as_str()
.map(|value| (name, value.to_string()))
.ok_or_else(|| PyValueError::new_err("header values must be strings"))
})
.collect()
}
#[pyclass]
struct ResponsesWebSocketConnection {
inner: RustResponsesWebSocketConnection,
}
#[pymethods]
impl ResponsesWebSocketConnection {
#[classmethod]
#[pyo3(signature = (url, headers=None, timeout_seconds=None))]
fn connect<'py>(
_cls: &Bound<'py, pyo3::types::PyType>,
py: Python<'py>,
url: String,
headers: Option<Py<PyAny>>,
timeout_seconds: Option<f64>,
) -> PyResult<Bound<'py, PyAny>> {
let headers = marshal_headers(py, headers)?;
let timeout = optional_timeout(timeout_seconds);
pyo3_async_runtimes::tokio::future_into_py(py, async move {
let inner = RustResponsesWebSocketConnection::connect_url(&url, &headers, timeout)
.await
.map_err(core_error_to_pyerr)?;
Python::attach(|py| Py::new(py, ResponsesWebSocketConnection { inner }))
})
}
fn send_text<'py>(&self, py: Python<'py>, text: String) -> PyResult<Bound<'py, PyAny>> {
let inner = self.inner.clone();
pyo3_async_runtimes::tokio::future_into_py(py, async move {
inner.send_text(text).await.map_err(core_error_to_pyerr)
})
}
fn recv_text<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
let inner = self.inner.clone();
pyo3_async_runtimes::tokio::future_into_py(py, async move {
inner.recv_text().await.map_err(core_error_to_pyerr)
})
}
fn close<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
let inner = self.inner.clone();
pyo3_async_runtimes::tokio::future_into_py(py, async move {
inner.close().await.map_err(core_error_to_pyerr)
})
}
}
fn marshal_inputs(
py: Python<'_>,
document: Py<PyAny>,
@ -172,6 +247,93 @@ fn aocr(
})
}
#[pyfunction]
#[pyo3(signature = (model, audio, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, optional_params=None, timeout_seconds=None))]
#[allow(clippy::too_many_arguments)]
fn transcription(
py: Python<'_>,
model: String,
audio: Py<PyAny>,
api_key: Option<String>,
api_base: Option<String>,
custom_llm_provider: Option<String>,
extra_headers: Option<Py<PyAny>>,
optional_params: Option<Py<PyAny>>,
timeout_seconds: Option<f64>,
) -> PyResult<Py<PyAny>> {
let audio = py_to_json(py, audio.bind(py))?;
let extra_headers = match extra_headers {
Some(headers) => Some(optional_object_to_map(py, "extra_headers", Some(headers))?),
None => None,
};
let optional_params = optional_object_to_map(py, "optional_params", optional_params)?;
let timeout = optional_timeout(timeout_seconds);
let result = gil::release_gil(py, || {
pyo3_async_runtimes::tokio::get_runtime().block_on(run_audio_transcription(
AudioTranscriptionRequest {
model: &model,
audio,
api_key: api_key.as_deref(),
api_base: api_base.as_deref(),
custom_llm_provider: custom_llm_provider.as_deref(),
extra_headers,
optional_params,
timeout,
callbacks: Vec::new(),
guardrails: Vec::new(),
request_metadata: Default::default(),
litellm_call_id: None,
},
))
});
match result {
Ok(value) => json_to_py(py, value),
Err(err) => Err(core_error_to_pyerr(err)),
}
}
#[pyfunction]
#[pyo3(signature = (model, audio, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, optional_params=None, timeout_seconds=None))]
#[allow(clippy::too_many_arguments)]
fn atranscription(
py: Python<'_>,
model: String,
audio: Py<PyAny>,
api_key: Option<String>,
api_base: Option<String>,
custom_llm_provider: Option<String>,
extra_headers: Option<Py<PyAny>>,
optional_params: Option<Py<PyAny>>,
timeout_seconds: Option<f64>,
) -> PyResult<Bound<'_, PyAny>> {
let audio = py_to_json(py, audio.bind(py))?;
let extra_headers = match extra_headers {
Some(headers) => Some(optional_object_to_map(py, "extra_headers", Some(headers))?),
None => None,
};
let optional_params = optional_object_to_map(py, "optional_params", optional_params)?;
let timeout = optional_timeout(timeout_seconds);
pyo3_async_runtimes::tokio::future_into_py(py, async move {
let value = run_audio_transcription(AudioTranscriptionRequest {
model: &model,
audio,
api_key: api_key.as_deref(),
api_base: api_base.as_deref(),
custom_llm_provider: custom_llm_provider.as_deref(),
extra_headers,
optional_params,
timeout,
callbacks: Vec::new(),
guardrails: Vec::new(),
request_metadata: Default::default(),
litellm_call_id: None,
})
.await
.map_err(core_error_to_pyerr)?;
Python::attach(|py| json_to_py(py, value))
})
}
type MarshaledMessagesInputs = (Value, Option<Map<String, Value>>, Option<Duration>);
fn marshal_messages_inputs(
@ -269,8 +431,11 @@ fn gil_stats(py: Python<'_>) -> PyResult<Py<PyAny>> {
fn _native(module: &Bound<'_, PyModule>) -> PyResult<()> {
module.add_function(wrap_pyfunction!(ocr, module)?)?;
module.add_function(wrap_pyfunction!(aocr, module)?)?;
module.add_function(wrap_pyfunction!(transcription, module)?)?;
module.add_function(wrap_pyfunction!(atranscription, module)?)?;
module.add_function(wrap_pyfunction!(messages, module)?)?;
module.add_function(wrap_pyfunction!(amessages, module)?)?;
module.add_class::<ResponsesWebSocketConnection>()?;
module.add_function(wrap_pyfunction!(gil_stats, module)?)?;
Ok(())
}

View file

@ -58,12 +58,12 @@ class DiskCache(BaseCache):
return return_val
def increment_cache(self, key, value: int, **kwargs) -> int:
# get the value
cached_value = self.get_cache(key=key)
init_value = cached_value if isinstance(cached_value, int) else 0
value = init_value + value
self.set_cache(key, value, **kwargs)
return value
with self.disk_cache.transact():
cached_value = self.get_cache(key=key)
init_value = cached_value if isinstance(cached_value, int) else 0
new_value = init_value + value
self.set_cache(key, new_value, **kwargs)
return new_value
async def async_get_cache(self, key, **kwargs):
return self.get_cache(key=key, **kwargs)
@ -76,12 +76,7 @@ class DiskCache(BaseCache):
return return_val
async def async_increment(self, key, value: int, **kwargs) -> int:
# get the value
cached_value = await self.async_get_cache(key=key)
init_value = cached_value if isinstance(cached_value, int) else 0
value = init_value + value
await self.async_set_cache(key, value, **kwargs)
return value
return self.increment_cache(key=key, value=value, **kwargs)
def flush_cache(self):
self.disk_cache.clear()

View file

@ -12,6 +12,7 @@ import json
import sys
import time
import heapq
import threading
from typing import TYPE_CHECKING, Any, List, Optional
if TYPE_CHECKING:
@ -46,6 +47,7 @@ class InMemoryCache(BaseCache):
self.cache_dict: dict = {}
self.ttl_dict: dict = {}
self.expiration_heap: list[tuple[float, str]] = []
self._increment_lock = threading.Lock()
def check_value_size(self, value: Any):
"""
@ -223,12 +225,13 @@ class InMemoryCache(BaseCache):
return_val.append(val)
return return_val
def increment_cache(self, key, value: int, **kwargs) -> int:
# get the value
init_value = self.get_cache(key=key) or 0
value = init_value + value
self.set_cache(key, value, **kwargs)
return value
def increment_cache(self, key, value: float, **kwargs) -> float:
with self._increment_lock:
# keep read-modify-write atomic
init_value = self.get_cache(key=key) or 0
value = init_value + value
self.set_cache(key, value, **kwargs)
return value
async def async_get_cache(self, key, **kwargs):
return self.get_cache(key=key, **kwargs)
@ -241,11 +244,7 @@ class InMemoryCache(BaseCache):
return return_val
async def async_increment(self, key, value: float, **kwargs) -> float:
# get the value
init_value = await self.async_get_cache(key=key) or 0
value = init_value + value
await self.async_set_cache(key, value, **kwargs)
return value
return self.increment_cache(key=key, value=value, **kwargs)
async def async_increment_pipeline(
self, increment_list: List["RedisPipelineIncrementOperation"], **kwargs

View file

@ -1292,6 +1292,7 @@ MAXIMUM_TRACEBACK_LINES_TO_LOG = int(os.getenv("MAXIMUM_TRACEBACK_LINES_TO_LOG",
X_LITELLM_DISABLE_CALLBACKS = "x-litellm-disable-callbacks"
LITELLM_METADATA_FIELD = "litellm_metadata"
OLD_LITELLM_METADATA_FIELD = "metadata"
RETURN_RAW_MODEL_NAME_METADATA_KEY = "_complexity_router_return_raw_model_name"
LITELLM_TRUNCATED_PAYLOAD_FIELD = "litellm_truncated"
LITELLM_TRUNCATION_DB_SAFEGUARD_NOTE = (
"Truncation is a DB storage safeguard. "
@ -1468,6 +1469,7 @@ _batch_polling_env = os.getenv("PROXY_BATCH_POLLING_ENABLED", "true").lower()
PROXY_BATCH_POLLING_ENABLED = _batch_polling_env == "true"
PROXY_BUDGET_RESCHEDULER_MAX_TIME = int(os.getenv("PROXY_BUDGET_RESCHEDULER_MAX_TIME", 605))
PROXY_BATCH_WRITE_AT = int(os.getenv("PROXY_BATCH_WRITE_AT", 10)) # in seconds, increased from 10
PROXY_CONFIG_RELOAD_INTERVAL_SECONDS = get_env_int("PROXY_CONFIG_RELOAD_INTERVAL_SECONDS", 30)
# APScheduler Configuration - MEMORY LEAK FIX
# These settings prevent memory leaks in APScheduler's normalize() and _apply_jitter() functions

View file

@ -9,6 +9,7 @@ from openai.types.chat import ChatCompletionToolParam
from openai.types.responses.function_tool_param import FunctionToolParam
from openai.types.shared_params.function_definition import FunctionDefinition
from litellm.types.llms.anthropic import AnthropicMessagesTool
from litellm.types.utils import ChatCompletionMessageToolCall
@ -75,6 +76,20 @@ def transform_mcp_tool_to_openai_responses_api_tool(
)
def transform_mcp_tool_to_anthropic_tool(mcp_tool: MCPTool) -> AnthropicMessagesTool:
"""Convert an MCP tool to an Anthropic Messages API tool."""
from litellm.litellm_core_utils.prompt_templates.common_utils import (
sanitize_input_schema_for_anthropic,
)
return AnthropicMessagesTool(
name=mcp_tool.name,
description=mcp_tool.description or "",
input_schema=sanitize_input_schema_for_anthropic(mcp_tool.inputSchema),
type="custom",
)
async def load_mcp_tools(
session: ClientSession, format: Literal["mcp", "openai"] = "mcp"
) -> Union[List[MCPTool], List[ChatCompletionToolParam]]:

View file

@ -91,7 +91,9 @@ class AnthropicCacheControlHook(CustomPromptManagement):
# Pass through non-message injection points for provider-specific handling
if remaining_points:
non_default_params["cache_control_injection_points"] = remaining_points
non_default_params["cache_control_injection_points"] = AnthropicCacheControlHook._stamped_as_judged(
remaining_points
)
return model, processed_messages, non_default_params
@ -310,6 +312,35 @@ class AnthropicCacheControlHook(CustomPromptManagement):
return ChatCompletionCachedContent(type="ephemeral", ttl=ttl)
return ChatCompletionCachedContent(type="ephemeral")
@staticmethod
def _stamped_as_judged(points: list[CacheControlInjectionPoint]) -> list[dict[str, object]]:
"""Mark written-back points as having passed the client cache_control judgment.
Builds copies because config-owned point dicts are shared across
requests; mutating them would leak the stamp into future requests.
"""
return [{**point, "_litellm_judged": True} for point in points]
@staticmethod
def _should_stand_down(
points: list[CacheControlInjectionPoint],
messages: list[AllMessageValues],
system: str | list | None,
tools: list | None,
) -> bool:
"""Whether configured injection points must yield to client-set cache_control.
Points that a prior pass over this request already judged and wrote
back carry the internal judged stamp; any re-entry (acompletion
re-entering completion, the async-to-sync /v1/messages dispatch,
interceptor sub-calls reusing the request kwargs) must not re-judge
them, because by then the messages carry litellm's own injected marks
and the judgment would misread those as client breakpoints.
"""
if all(point.get("_litellm_judged") for point in points):
return False
return AnthropicCacheControlHook._request_has_cache_control(messages, system, tools)
@staticmethod
def _request_has_cache_control(
messages: list[AllMessageValues],
@ -322,7 +353,9 @@ class AnthropicCacheControlHook(CustomPromptManagement):
stand down entirely rather than add more, per the auto-caching contract.
Tools count: they are a breakpoint the client can mark, they count toward
the provider's four-block limit, and caching only the tool definitions is
a common pattern, so injecting alongside them can exceed the cap.
a common pattern, so injecting alongside them can exceed the cap. Tools
carry the mark either at the top level (Anthropic shape) or nested under
``function`` (OpenAI shape); the Anthropic chat transform accepts both.
"""
if any(AnthropicCacheControlHook._count_cache_control_blocks(msg) for msg in messages):
return True
@ -330,7 +363,14 @@ class AnthropicCacheControlHook(CustomPromptManagement):
if any(isinstance(block, dict) and block.get("cache_control") is not None for block in system):
return True
if tools is not None:
return any(isinstance(tool, dict) and tool.get("cache_control") is not None for tool in tools)
return any(
isinstance(tool, dict)
and (
tool.get("cache_control") is not None
or (isinstance(tool.get("function"), dict) and tool["function"].get("cache_control") is not None)
)
for tool in tools
)
return False
@staticmethod
@ -392,13 +432,23 @@ class AnthropicCacheControlHook(CustomPromptManagement):
custom_llm_provider: str | None,
tools: list | None = None,
) -> None:
"""For /chat/completions: add default injection points to the request params.
"""For /chat/completions: resolve the injection points the request should carry.
No-op when injection points are already configured (explicit config wins).
Seeding the param lets the existing prompt-management gate and the
AnthropicCacheControlHook run unchanged.
Configured injection points win over the automatic defaults, but stand
down entirely when the client already marked its own cache_control
breakpoints (messages or tools): injecting alongside them clashes with
the client's caching strategy and can exceed the provider's four-block
limit. The judgment happens once per request; points a prior pass
wrote back carry the judged stamp and are never re-judged (see
``_should_stand_down``). Seeding the param lets the existing
prompt-management gate and the AnthropicCacheControlHook run
unchanged.
"""
if non_default_params.get("cache_control_injection_points"):
if AnthropicCacheControlHook._should_stand_down(
non_default_params["cache_control_injection_points"], messages, None, tools
):
non_default_params.pop("cache_control_injection_points")
return
points = AnthropicCacheControlHook.get_default_injection_points(
messages=messages,
@ -421,18 +471,26 @@ class AnthropicCacheControlHook(CustomPromptManagement):
) -> Tuple[List[Dict], str | list | None]:
"""Extract cache_control_injection_points from kwargs and apply if present.
When none are configured but ``litellm.enable_anthropic_prompt_caching``
is on, synthesize default breakpoints for the native /v1/messages path.
Pops the key from kwargs; if remaining (non-message) points exist they
are written back so downstream transforms can handle them.
Configured points stand down entirely when the client already marked
its own cache_control breakpoints anywhere in the request. The
judgment happens once per request; points a prior pass wrote back
carry the judged stamp and are never re-judged (see
``_should_stand_down``). When none are configured but
``litellm.enable_anthropic_prompt_caching`` is on, synthesize default
breakpoints for the native /v1/messages path. Pops the key from kwargs;
if remaining (non-message) points exist they are written back so
downstream transforms can handle them.
"""
typed_messages = cast(list[AllMessageValues], messages) # cast-ok: Anthropic-shaped dicts from v1/messages
configured = cast( # cast-ok: kwargs is untyped; this key only holds the documented injection-point list
list[CacheControlInjectionPoint] | None, kwargs.pop("cache_control_injection_points", None)
)
if configured and AnthropicCacheControlHook._should_stand_down(configured, typed_messages, system, tools):
return messages, system
injection_points: list[CacheControlInjectionPoint] = configured or []
if not injection_points and model is not None:
injection_points = AnthropicCacheControlHook.get_default_injection_points(
messages=cast(list[AllMessageValues], messages), # cast-ok: Anthropic-shaped dicts from v1/messages
messages=typed_messages,
system=system,
tools=tools,
model=model,
@ -447,7 +505,7 @@ class AnthropicCacheControlHook(CustomPromptManagement):
injection_points=injection_points,
)
if remaining:
kwargs["cache_control_injection_points"] = remaining
kwargs["cache_control_injection_points"] = AnthropicCacheControlHook._stamped_as_judged(remaining)
return messages, system
@property

View file

@ -14,6 +14,7 @@ from litellm.compression import compress
from litellm.integrations.custom_logger import CustomLogger
from litellm.types.integrations.compression_interception import (
CompressionInterceptionConfig,
CompressionSavingsMetadata,
)
from litellm.types.integrations.custom_logger import (
AgenticLoopPlan,
@ -25,6 +26,41 @@ LITELLM_CONTENT_RETRIEVE_TOOL_NAME = "litellm_content_retrieve"
_CACHE_TTL_SECONDS = 15 * 60
def _compression_savings_from_counts(
original_tokens: object, compressed_tokens: object
) -> CompressionSavingsMetadata | None:
if isinstance(original_tokens, bool) or not isinstance(original_tokens, int):
return None
if isinstance(compressed_tokens, bool) or not isinstance(compressed_tokens, int):
return None
if compressed_tokens < 0 or original_tokens < compressed_tokens:
return None
return CompressionSavingsMetadata(
tokens_before=original_tokens,
tokens_after=compressed_tokens,
tokens_saved=original_tokens - compressed_tokens,
source="compression_interception",
)
def _record_compression_savings(kwargs: dict[str, object], savings: CompressionSavingsMetadata) -> None:
"""
Attach savings to the request's litellm metadata so they land in the
SpendLog row's metadata JSON under ``compression_savings``.
``/v1/messages`` requests carry proxy metadata under ``litellm_metadata``
(the ``metadata`` key is Anthropic's own API field). The existing dict is
updated in place because the proxy and the logging object hold references
to the same object; replacing it would orphan writes made through those
references.
"""
existing = kwargs.get("litellm_metadata")
if isinstance(existing, dict):
existing["compression_savings"] = savings
return
kwargs["litellm_metadata"] = {"compression_savings": savings}
class CompressionInterceptionLogger(CustomLogger):
"""
CustomLogger that implements transparent prompt compression + retrieval loops.
@ -130,6 +166,12 @@ class CompressionInterceptionLogger(CustomLogger):
call_id = str(uuid.uuid4())
kwargs["litellm_call_id"] = call_id
self._compression_cache_by_call_id[call_id] = (cache, time.time())
savings = _compression_savings_from_counts(
original_tokens=compressed.get("original_tokens"),
compressed_tokens=compressed.get("compressed_tokens"),
)
if savings is not None:
_record_compression_savings(kwargs=kwargs, savings=savings)
verbose_logger.debug(
"CompressionInterception: compressed request [call_id=%s original=%d compressed=%d cached_keys=%d]",
call_id,

View file

@ -1,8 +1,8 @@
import base64
import json # <--- NEW
import json
import os
from datetime import datetime
from typing import TYPE_CHECKING, Any, Optional, Union
from typing import TYPE_CHECKING, Any, Dict, Optional, Union
from litellm._logging import verbose_logger
from litellm.integrations.arize import _utils
@ -25,6 +25,8 @@ else:
LANGFUSE_CLOUD_EU_ENDPOINT = "https://cloud.langfuse.com/api/public/otel"
LANGFUSE_CLOUD_US_ENDPOINT = "https://us.cloud.langfuse.com/api/public/otel"
LANGFUSE_INGESTION_VERSION_HEADER = "x-langfuse-ingestion-version"
LANGFUSE_INGESTION_VERSION = "4"
class LangfuseOtelLogger(OpenTelemetry):
@ -326,7 +328,9 @@ class LangfuseOtelLogger(OpenTelemetry):
return OpenTelemetryConfig(
exporter="otlp_http",
endpoint=endpoint,
headers=f"Authorization={auth_header}",
headers=LangfuseOtelLogger._format_otel_headers(
LangfuseOtelLogger._build_langfuse_otel_headers(auth_header)
),
)
@staticmethod
@ -338,6 +342,26 @@ class LangfuseOtelLogger(OpenTelemetry):
auth_header = base64.b64encode(auth_string.encode()).decode()
return f"Basic {auth_header}"
@staticmethod
def _build_langfuse_otel_headers(auth_header: str) -> Dict[str, str]:
"""
Build the OTLP header set Langfuse expects.
`x-langfuse-ingestion-version: 4` selects Langfuse's v4 ingestion path;
without it spans fall back to the older transformation path.
"""
return {
"Authorization": auth_header,
LANGFUSE_INGESTION_VERSION_HEADER: LANGFUSE_INGESTION_VERSION,
}
@staticmethod
def _format_otel_headers(headers: Dict[str, str]) -> str:
"""
Serialize a header mapping into the comma-separated OTLP header string
"""
return ",".join(f"{key}={value}" for key, value in headers.items())
def construct_dynamic_otel_headers(
self, standard_callback_dynamic_params: StandardCallbackDynamicParams
) -> Optional[dict]:
@ -358,7 +382,7 @@ class LangfuseOtelLogger(OpenTelemetry):
public_key=dynamic_langfuse_public_key,
secret_key=dynamic_langfuse_secret_key,
)
dynamic_headers["Authorization"] = auth_header
dynamic_headers.update(LangfuseOtelLogger._build_langfuse_otel_headers(auth_header))
return dynamic_headers

View file

@ -57,6 +57,35 @@ def safe_divide(
return numerator / denominator
def coerce_token_limit(value: object) -> int | None:
"""
Coerce a max_input_tokens / max_output_tokens value to an int, treating a
malformed value as absent.
A deployment's model_info is registered into litellm.model_cost verbatim, so a
config value like "128,000" or "" reaches the /v1/models listing uncoerced from
both the router index and the cost map. Returning None omits that one limit
instead of failing the whole listing.
Args:
value: The raw configured or cost-map value
Returns:
The value as an int, or None if it is missing or not a usable number.
Bools are rejected because True/False is never a meaningful token limit.
"""
if isinstance(value, bool):
return None
if isinstance(value, int):
return value
if isinstance(value, (str, float)):
try:
return int(value)
except (TypeError, ValueError, OverflowError):
return None
return None
_FINISH_REASON_MAP: dict[str, OpenAIChatCompletionFinishReason] = {
# Anthropic
"stop_sequence": "stop",

View file

@ -7,8 +7,8 @@ duration_in_seconds is used in diff parts of the code base, example
"""
import re
import time
from datetime import datetime, timedelta, timezone, tzinfo
import time as time_module
from datetime import datetime, time, timedelta, timezone, tzinfo
from typing import Optional, Tuple
from zoneinfo import ZoneInfo
@ -61,7 +61,7 @@ def duration_in_seconds(duration: str) -> int:
elif unit == "w":
return value * 604800
elif unit == "mo":
now = time.time()
now = time_module.time()
current_time = datetime.fromtimestamp(now)
# Calculate target month and year, handling overflow past December
@ -94,12 +94,17 @@ def duration_in_seconds(duration: str) -> int:
raise ValueError(f"Unsupported duration unit, passed duration: {duration}")
def get_next_standardized_reset_time(duration: str, current_time: datetime, timezone_str: str = "UTC") -> datetime:
def get_next_standardized_reset_time(
duration: str,
current_time: datetime,
timezone_str: str = "UTC",
reset_time_of_day: time = time(0, 0),
) -> datetime:
"""
Get the next standardized reset time based on the duration.
All durations will reset at predictable intervals, aligned from the current time:
- Nd: If N=1, reset at next midnight; if N>1, reset every N days from now
- Nd: If N=1, reset at the next `reset_time_of_day`; if N>1, reset every N days from now
- Nh: Every N hours, aligned to hour boundaries (e.g., 1:00, 2:00)
- Nm: Every N minutes, aligned to minute boundaries (e.g., 1:05, 1:10)
- Ns: Every N seconds, aligned to second boundaries
@ -108,12 +113,15 @@ def get_next_standardized_reset_time(duration: str, current_time: datetime, time
- duration: Duration string (e.g. "30s", "30m", "30h", "30d")
- current_time: Current datetime
- timezone_str: Timezone string (e.g. "UTC", "US/Eastern", "Asia/Kolkata")
- reset_time_of_day: Wall-clock time the reset lands on for day/week/month
durations (defaults to midnight). Ignored for sub-day durations, where a
time-of-day is meaningless.
Returns:
- Next reset time at a standardized interval in the specified timezone
"""
# Set up timezone and normalize current time
current_time, tz = _setup_timezone(current_time, timezone_str)
current_time, _ = _setup_timezone(current_time, timezone_str)
# Parse duration
value, unit = _parse_duration(duration)
@ -126,9 +134,9 @@ def get_next_standardized_reset_time(duration: str, current_time: datetime, time
# Handle different time units
if unit == "d":
return _handle_day_reset(current_time, base_midnight, value, tz)
return _handle_day_reset(current_time, base_midnight, value, reset_time_of_day)
elif unit == "w":
return _handle_day_reset(current_time, base_midnight, value * 7, tz)
return _handle_day_reset(current_time, base_midnight, value * 7, reset_time_of_day)
elif unit == "h":
return _handle_hour_reset(current_time, base_midnight, value)
elif unit == "m":
@ -136,7 +144,7 @@ def get_next_standardized_reset_time(duration: str, current_time: datetime, time
elif unit == "s":
return _handle_second_reset(current_time, base_midnight, value)
elif unit == "mo":
return _handle_month_reset(current_time, base_midnight, value)
return _handle_month_reset(current_time, base_midnight, value, reset_time_of_day)
else:
# Unrecognized unit, default to next midnight
return base_midnight + timedelta(days=1)
@ -175,46 +183,58 @@ def _parse_duration(duration: str) -> Tuple[Optional[int], Optional[str]]:
return int(value), unit
def _handle_day_reset(current_time: datetime, base_midnight: datetime, value: int, tz: tzinfo) -> datetime:
def _apply_time_of_day(dt: datetime, reset_time_of_day: time) -> datetime:
"""Set the wall-clock time of `dt` to `reset_time_of_day`, keeping its date and tzinfo."""
return dt.replace(
hour=reset_time_of_day.hour,
minute=reset_time_of_day.minute,
second=reset_time_of_day.second,
microsecond=reset_time_of_day.microsecond,
)
def _next_occurrence(
boundary_midnight: datetime,
reset_time_of_day: time,
current_time: datetime,
period: timedelta,
) -> datetime:
"""Place the reset at `reset_time_of_day` on the boundary day, rolling forward one
`period` if that instant has already passed (or is exactly now)."""
candidate = _apply_time_of_day(boundary_midnight, reset_time_of_day)
if candidate <= current_time:
return candidate + period
return candidate
def _first_of_next_month(first_of_month: datetime) -> datetime:
"""Given the 1st of some month, return the 1st of the following month."""
if first_of_month.month == 12:
return first_of_month.replace(year=first_of_month.year + 1, month=1)
return first_of_month.replace(month=first_of_month.month + 1)
def _handle_day_reset(
current_time: datetime,
base_midnight: datetime,
value: int,
reset_time_of_day: time,
) -> datetime:
"""Handle day-based reset times."""
# Handle zero value - immediate expiration
if value == 0:
return current_time
if value == 1: # Daily reset at midnight
return base_midnight + timedelta(days=1)
elif value == 7: # Weekly reset on Monday at midnight
if value == 1: # Daily reset at the configured time of day
return _next_occurrence(base_midnight, reset_time_of_day, current_time, timedelta(days=1))
elif value == 7: # Weekly reset on Monday at the configured time of day
days_until_monday = (7 - current_time.weekday()) % 7
if days_until_monday == 0: # If today is Monday
days_until_monday = 7
return base_midnight + timedelta(days=days_until_monday)
elif value == 30: # Monthly reset on 1st at midnight
# Get 1st of next month at midnight
if current_time.month == 12:
next_reset = datetime(
year=current_time.year + 1,
month=1,
day=1,
hour=0,
minute=0,
second=0,
microsecond=0,
tzinfo=tz,
)
else:
next_reset = datetime(
year=current_time.year,
month=current_time.month + 1,
day=1,
hour=0,
minute=0,
second=0,
microsecond=0,
tzinfo=tz,
)
return next_reset
else: # Custom day value - next interval is value days from current
return current_time.replace(hour=0, minute=0, second=0, microsecond=0) + timedelta(days=value)
upcoming_monday = base_midnight + timedelta(days=days_until_monday)
return _next_occurrence(upcoming_monday, reset_time_of_day, current_time, timedelta(days=7))
elif value == 30: # Monthly reset on 1st at the configured time of day
return _handle_month_reset(current_time, base_midnight, 1, reset_time_of_day)
else: # Custom day value - next interval is value days from the start of today
return _apply_time_of_day(base_midnight + timedelta(days=value), reset_time_of_day)
def _handle_hour_reset(current_time: datetime, base_midnight: datetime, value: int) -> datetime:
@ -316,36 +336,30 @@ def _handle_second_reset(current_time: datetime, base_midnight: datetime, value:
return current_time.replace(hour=next_hour, minute=next_minute, second=next_second, microsecond=0)
def _handle_month_reset(current_time: datetime, base_midnight: datetime, value: int) -> datetime:
def _handle_month_reset(
current_time: datetime,
base_midnight: datetime,
value: int,
reset_time_of_day: time,
) -> datetime:
"""
Handle monthly reset times. For monthly resets, we always reset at the start of the next month.
Handle monthly reset times. Resets land on the 1st at `reset_time_of_day`; if the
1st of the current month at that time has already passed, roll to the 1st of next month.
Args:
current_time: Current datetime
base_midnight: Midnight of current day
value: Number of months (currently only supports 1 month resets)
reset_time_of_day: Wall-clock time the reset lands on
Returns:
datetime: First day of next month at midnight
datetime: First day of the next reset month at `reset_time_of_day`
"""
if value != 1:
raise ValueError("Monthly resets currently only support 1 month intervals")
# Get the first day of next month
if current_time.month == 12:
next_month = 1
next_year = current_time.year + 1
else:
next_month = current_time.month + 1
next_year = current_time.year
return datetime(
year=next_year,
month=next_month,
day=1,
hour=0,
minute=0,
second=0,
microsecond=0,
tzinfo=current_time.tzinfo,
)
first_of_this_month = base_midnight.replace(day=1)
candidate = _apply_time_of_day(first_of_this_month, reset_time_of_day)
if candidate <= current_time:
return _apply_time_of_day(_first_of_next_month(first_of_this_month), reset_time_of_day)
return candidate

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