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

This commit is contained in:
Yucheng He 2026-09-30 14:10:49 -07:00
commit 63f6095c35
859 changed files with 72305 additions and 13512 deletions

View file

@ -3419,7 +3419,7 @@ workflows:
name: integration-<< matrix.suite >>
matrix:
parameters:
suite: [management, accounting, database, providers, mcp, sdk, cost, browser]
suite: [management, accounting, database, providers, mcp, sdk, cost, security, browser]
- integration_contracts:
name: integration-extensions
suite: extensions

View file

@ -89,6 +89,7 @@ legacy_paths() {
proxy-db-auth-checks)
echo tests/unit/proxy/auth/test_auth_checks.py
echo tests/unit/proxy/auth/test_user_api_key_auth.py
echo tests/unit/proxy/test_credential_slot_registry.py
echo tests/unit/proxy/test_deprecated_key_grace_period.py ;;
proxy-db-budgets)
echo tests/unit/proxy/auth/test_default_end_user_budget_simple.py

View file

@ -42,7 +42,7 @@ case "$subject" in
;;
esac
ALLOWED_TYPES="feat|fix|docs|style|refactor|perf|test|build|ci|chore|revert"
ALLOWED_TYPES="feat|fix|docs|style|refactor|perf|test|build|ci|chore|revert|security"
# Description must not start with an uppercase letter — kept in sync with the
# subjectPattern in .github/workflows/conventional-commits.yml so the local
# hook is the strictly tighter of the two gates. (Without this guard, a commit
@ -61,7 +61,7 @@ cat >&2 <<EOF
Expected: <type>(<scope>)!: <description>
(description must start with a lowercase letter)
Allowed types: feat, fix, docs, style, refactor, perf, test, build, ci, chore, revert
Allowed types: feat, fix, docs, style, refactor, perf, test, build, ci, chore, revert, security
Examples:
feat(router): add weighted round-robin strategy
fix(bedrock): decouple STS region from aws_region_name

8
.github/CODEOWNERS vendored
View file

@ -1,10 +1,2 @@
/ui/ @yuneng-berri @ryan-crabbe-berri
/litellm/proxy/_experimental/out/ @yuneng-berri @ryan-crabbe-berri
/ui/Dockerfile
/ui/nginx.conf
/ui/litellm-dashboard/src/lib/http/schema.d.ts
/ui/litellm-dashboard/tsconfig.tsbuildinfo
/model_prices_and_context_window.json @mateo-berri @ryan-crabbe-berri @kerry-berri
/litellm/model_prices_and_context_window_backup.json @mateo-berri @ryan-crabbe-berri @kerry-berri
/litellm-proxy-extras/litellm_proxy_extras/migrations/ @yuneng-berri @ryan-crabbe-berri
/.github/CODEOWNERS @yuneng-berri

View file

@ -41,6 +41,7 @@ jobs:
ci
chore
revert
security
requireScope: false
subjectPattern: ^(?![A-Z]).+$
subjectPatternError: |

View file

@ -15,6 +15,8 @@ jobs:
runs-on: ubuntu-latest
permissions:
contents: write
outputs:
version: ${{ steps.version.outputs.version }}
steps:
- name: Require main
env:
@ -64,3 +66,14 @@ jobs:
sha: context.sha,
});
core.info(`Created branch ${branchName} at ${context.sha}`);
linear-release:
name: Move the Linear release to rc
needs: create-rc-branch
permissions:
contents: read
uses: ./.github/workflows/linear-release.yml
with:
rc_version: ${{ needs.create-rc-branch.outputs.version }}
secrets:
LINEAR_API_KEY: ${{ secrets.LINEAR_API_KEY }}

131
.github/workflows/linear-release.yml vendored Normal file
View file

@ -0,0 +1,131 @@
name: Linear Release
on:
push:
branches:
- main
- "rc/**"
release:
types: [published]
workflow_call:
inputs:
rc_version:
description: "X.Y.0 release whose rc branch was just cut"
required: true
type: string
secrets:
LINEAR_API_KEY:
required: true
permissions: {}
jobs:
linear-release:
name: Linear Release
if: github.repository == 'BerriAI/litellm'
runs-on: ubuntu-latest
permissions:
contents: read
steps:
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
fetch-depth: 0
persist-credentials: false
- name: Plan
id: plan
env:
EVENT: ${{ github.event_name }}
REF_NAME: ${{ github.ref_name }}
BEFORE: ${{ github.event.before }}
CREATED: ${{ github.event.created }}
RC_VERSION: ${{ inputs.rc_version }}
RELEASE_TAG: ${{ github.event.release.tag_name }}
PRERELEASE: ${{ github.event.release.prerelease }}
run: |
set -euo pipefail
sync_base="${BEFORE}"
if [ "${CREATED}" = "true" ]; then
sync_base=""
fi
if [ -n "${RC_VERSION}" ]; then
echo "version=${RC_VERSION}" >> "$GITHUB_OUTPUT"
echo "stage=rc" >> "$GITHUB_OUTPUT"
elif [ "${EVENT}" = "release" ]; then
if [ "${PRERELEASE}" = "true" ] || ! echo "${RELEASE_TAG}" | grep -qE '^v[0-9]+\.[0-9]+\.0$'; then
echo "::notice::${RELEASE_TAG} is not an X.Y.0 stable release; nothing to complete"
exit 0
fi
echo "version=${RELEASE_TAG#v}" >> "$GITHUB_OUTPUT"
echo "complete=true" >> "$GITHUB_OUTPUT"
elif [ "${REF_NAME}" = "main" ]; then
version="$(python3 .github/scripts/read_rc_version.py | cut -d= -f2)"
status=0
git ls-remote --exit-code --heads origin "rc/${version}" > /dev/null || status=$?
case "${status}" in
0)
IFS=. read -r major minor _ <<< "${version}"
version="${major}.$((minor + 1)).0"
;;
2) ;;
*)
echo "::error::could not check whether rc/${version} exists (git ls-remote exit ${status})"
exit 1
;;
esac
echo "version=${version}" >> "$GITHUB_OUTPUT"
echo "sync_base=${sync_base}" >> "$GITHUB_OUTPUT"
echo "main=true" >> "$GITHUB_OUTPUT"
else
echo "version=${REF_NAME#rc/}" >> "$GITHUB_OUTPUT"
echo "sync_base=${sync_base}" >> "$GITHUB_OUTPUT"
echo "stage=rc" >> "$GITHUB_OUTPUT"
fi
- name: Sync commits into the release
if: steps.plan.outputs.sync_base != ''
uses: linear/linear-release-action@d4af10092984f9bc6d5efa075b242bdf01333463 # v0.18.0
with:
access_key: ${{ secrets.LINEAR_API_KEY }}
command: sync
name: LiteLLM ${{ steps.plan.outputs.version }}
version: ${{ steps.plan.outputs.version }}
base_ref: ${{ steps.plan.outputs.sync_base }}
cli_version: v0.18.0
- name: Keep the main stage unless the rc branch was cut during this run
id: main_stage
if: steps.plan.outputs.main == 'true'
env:
VERSION: ${{ steps.plan.outputs.version }}
run: |
set -euo pipefail
status=0
git ls-remote --exit-code --heads origin "rc/${VERSION}" > /dev/null || status=$?
case "${status}" in
0) echo "::notice::rc/${VERSION} was cut during this run; leaving the release in its rc stage" ;;
2) echo "stage=main" >> "$GITHUB_OUTPUT" ;;
*)
echo "::error::could not check whether rc/${VERSION} exists (git ls-remote exit ${status})"
exit 1
;;
esac
- name: Move the release to its stage
if: steps.plan.outputs.stage != '' || steps.main_stage.outputs.stage != ''
uses: linear/linear-release-action@d4af10092984f9bc6d5efa075b242bdf01333463 # v0.18.0
with:
access_key: ${{ secrets.LINEAR_API_KEY }}
command: update
stage: ${{ steps.plan.outputs.stage || steps.main_stage.outputs.stage }}
version: ${{ steps.plan.outputs.version }}
cli_version: v0.18.0
- name: Complete the release
if: steps.plan.outputs.complete == 'true'
uses: linear/linear-release-action@d4af10092984f9bc6d5efa075b242bdf01333463 # v0.18.0
with:
access_key: ${{ secrets.LINEAR_API_KEY }}
command: complete
version: ${{ steps.plan.outputs.version }}
cli_version: v0.18.0

View file

@ -99,7 +99,7 @@
"limit": 0
},
"reportUnknownArgumentType": {
"limit": 44358
"limit": 44802
},
"reportUnknownLambdaType": {
"limit": 109

View file

@ -1,37 +0,0 @@
# Publish MCP servers in the AI Hub
Set `litellm_settings.public_mcp_servers` to the concrete IDs of the servers you want listed in the public AI Hub. Pin `server_id` in each configuration entry so the publication list stays stable across deployments
```yaml
mcp_servers:
documentation:
server_id: documentation-mcp
url: https://mcp.example.com/mcp
transport: http
available_on_public_internet: true
litellm_settings:
public_mcp_hub_strict_whitelist: true
public_mcp_servers:
- documentation-mcp
```
Use `documentation-mcp`, the `server_id`, in the publication list. The configuration key `documentation`, display names, and aliases are not publication IDs. Database-created servers use the ID returned by `/v1/mcp/server`
The dashboard's **AI Hub > MCP Hub > Manage MCP Hub Visibility** dialog edits this same list. Its YAML example includes the selected server IDs. With database-backed configuration (`store_model_in_db: true`), a value declared in YAML is owned by that file: edit the file and reload, or remove that key from YAML to let the dashboard manage it in the database. File-backed deployments can save the list directly to their configuration file
To remove all explicit entries, save an empty selection in the dialog or configure:
```yaml
litellm_settings:
public_mcp_hub_strict_whitelist: true
public_mcp_servers: []
```
## Hub listing and network access
The **Hub listing** column in AI Hub identifies servers that appear in `/public/mcp_hub`. The dashboard derives this status from the current registry and publication settings. Setting `mcp_info.is_public` on a server does not publish it; that response field is derived metadata. `mcp_info.is_public_explicit` identifies registered servers included in the explicit publication list
Gateway cards and server details show **All Networks** when `available_on_public_internet` is enabled or the server is explicitly published in `public_mcp_servers`. They show **Internal Only** when both are false. The per-server flag defaults to `true`; explicit publication overrides a disabled flag for compatibility. Older proxies that omit the metadata needed to determine access show **Unknown**. These labels describe allowed client IPs; authentication and tool permissions still apply
The default `public_mcp_hub_strict_whitelist: true` lists only registered servers in `public_mcp_servers`. Legacy mode (`false`) additionally lists registered servers with `available_on_public_internet: true`. In legacy mode, clearing the explicit publication list leaves these automatically listed servers visible. Enable strict mode when the publication list should fully determine hub visibility

View file

@ -24,7 +24,7 @@ model_list:
- model_name: sagemaker-completion-model
litellm_params:
model: sagemaker/berri-benchmarking-Llama-2-70b-chat-hf-4
input_cost_per_second: 0.000420
cost_per_second: 0.000420
- model_name: text-embedding-ada-002
litellm_params:
model: azure/azure-embedding-model

View file

@ -1,6 +1,6 @@
[project]
name = "litellm-enterprise"
version = "0.1.71"
version = "0.1.72"
description = "Package for LiteLLM Enterprise features"
readme = "README.md"
requires-python = ">=3.9"
@ -26,7 +26,7 @@ required-version = ">=0.10.9"
module-root = ""
[tool.commitizen]
version = "0.1.71"
version = "0.1.72"
version_files = [
"pyproject.toml:^version",
"../pyproject.toml:litellm-enterprise==",

View file

@ -0,0 +1,97 @@
-- AlterTable
ALTER TABLE "LiteLLM_AgentsTable" ADD COLUMN IF NOT EXISTS "enabled" BOOLEAN NOT NULL DEFAULT true,
ADD COLUMN IF NOT EXISTS "execution_mode" TEXT NOT NULL DEFAULT 'autonomous',
ADD COLUMN IF NOT EXISTS "identity_managed" BOOLEAN NOT NULL DEFAULT false;
-- AlterTable
ALTER TABLE "LiteLLM_SpendLogs" ADD COLUMN IF NOT EXISTS "billing_agent_id" TEXT;
-- CreateTable
CREATE TABLE IF NOT EXISTS "LiteLLM_AgentIdentity" (
"agent_id" TEXT NOT NULL,
"active" BOOLEAN NOT NULL DEFAULT true,
"provider" TEXT NOT NULL,
"issuer" TEXT NOT NULL,
"tenant_id" TEXT NOT NULL,
"client_id" TEXT NOT NULL,
"service_principal_id" TEXT,
"required_roles" TEXT[] DEFAULT ARRAY[]::TEXT[],
"required_scopes" TEXT[] DEFAULT ARRAY['user_impersonation']::TEXT[],
"revision" TEXT NOT NULL,
"last_authenticated_at" TIMESTAMP(3),
CONSTRAINT "LiteLLM_AgentIdentity_pkey" PRIMARY KEY ("agent_id")
);
-- CreateTable
CREATE TABLE IF NOT EXISTS "LiteLLM_RetiredAgentIdentity" (
"binding_id" TEXT NOT NULL,
"agent_id" TEXT,
"provider" TEXT NOT NULL,
"issuer" TEXT NOT NULL,
"tenant_id" TEXT NOT NULL,
"client_id" TEXT NOT NULL,
CONSTRAINT "LiteLLM_RetiredAgentIdentity_pkey" PRIMARY KEY ("binding_id")
);
-- CreateTable
CREATE TABLE IF NOT EXISTS "LiteLLM_RetiredAgent" (
"original_agent_id" TEXT NOT NULL,
"retired_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
CONSTRAINT "LiteLLM_RetiredAgent_pkey" PRIMARY KEY ("original_agent_id")
);
-- CreateTable
CREATE TABLE IF NOT EXISTS "LiteLLM_VerifiedSubject" (
"subject_id" TEXT NOT NULL,
"issuer" TEXT NOT NULL,
"tenant_id" TEXT NOT NULL,
"oid" TEXT NOT NULL,
"kind" TEXT NOT NULL DEFAULT 'human',
"user_id" TEXT,
"verified_via" TEXT NOT NULL DEFAULT 'sso_interactive',
"verified_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
CONSTRAINT "LiteLLM_VerifiedSubject_pkey" PRIMARY KEY ("subject_id")
);
-- CreateIndex
CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_AgentIdentity_provider_tenant_id_client_id_key" ON "LiteLLM_AgentIdentity"("provider", "tenant_id", "client_id");
-- CreateIndex
CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_AgentIdentity_issuer_service_principal_id_key" ON "LiteLLM_AgentIdentity"("issuer", "service_principal_id");
-- CreateIndex
CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_RetiredAgentIdentity_provider_tenant_id_client_id_key" ON "LiteLLM_RetiredAgentIdentity"("provider", "tenant_id", "client_id");
-- CreateIndex
CREATE INDEX IF NOT EXISTS "LiteLLM_VerifiedSubject_user_id_idx" ON "LiteLLM_VerifiedSubject"("user_id");
-- CreateIndex
CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_VerifiedSubject_issuer_tenant_id_oid_key" ON "LiteLLM_VerifiedSubject"("issuer", "tenant_id", "oid");
-- AddForeignKey
DO $$
BEGIN
IF NOT EXISTS (SELECT 1 FROM pg_constraint WHERE conname = 'LiteLLM_AgentIdentity_agent_id_fkey') THEN
ALTER TABLE "LiteLLM_AgentIdentity" ADD CONSTRAINT "LiteLLM_AgentIdentity_agent_id_fkey" FOREIGN KEY ("agent_id") REFERENCES "LiteLLM_AgentsTable"("agent_id") ON DELETE CASCADE ON UPDATE CASCADE;
END IF;
END $$;
-- AddForeignKey
DO $$
BEGIN
IF NOT EXISTS (SELECT 1 FROM pg_constraint WHERE conname = 'LiteLLM_RetiredAgentIdentity_agent_id_fkey') THEN
ALTER TABLE "LiteLLM_RetiredAgentIdentity" ADD CONSTRAINT "LiteLLM_RetiredAgentIdentity_agent_id_fkey" FOREIGN KEY ("agent_id") REFERENCES "LiteLLM_AgentsTable"("agent_id") ON DELETE SET NULL ON UPDATE CASCADE;
END IF;
END $$;
-- AddForeignKey
DO $$
BEGIN
IF NOT EXISTS (SELECT 1 FROM pg_constraint WHERE conname = 'LiteLLM_VerifiedSubject_user_id_fkey') THEN
ALTER TABLE "LiteLLM_VerifiedSubject" ADD CONSTRAINT "LiteLLM_VerifiedSubject_user_id_fkey" FOREIGN KEY ("user_id") REFERENCES "LiteLLM_UserTable"("user_id") ON DELETE CASCADE ON UPDATE CASCADE;
END IF;
END $$;

View file

@ -0,0 +1,2 @@
-- AlterTable
ALTER TABLE "LiteLLM_MCPServerTable" ADD COLUMN IF NOT EXISTS "pinned_tools" JSONB DEFAULT '{}';

View file

@ -0,0 +1,19 @@
CREATE TABLE IF NOT EXISTS "LiteLLM_DailyModelUsage" (
"date" TEXT NOT NULL,
"model_group" TEXT NOT NULL,
"model" TEXT NOT NULL,
"custom_llm_provider" TEXT NOT NULL,
"task_type" TEXT NOT NULL,
"spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0,
"prompt_tokens" BIGINT NOT NULL DEFAULT 0,
"completion_tokens" BIGINT NOT NULL DEFAULT 0,
"request_count" BIGINT NOT NULL DEFAULT 0,
"successful_requests" BIGINT NOT NULL DEFAULT 0,
"failed_requests" BIGINT NOT NULL DEFAULT 0,
"created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
"updated_at" TIMESTAMP(3) NOT NULL,
CONSTRAINT "LiteLLM_DailyModelUsage_pkey" PRIMARY KEY ("date", "model_group", "model", "custom_llm_provider", "task_type")
);
CREATE INDEX IF NOT EXISTS "LiteLLM_DailyModelUsage_date_idx" ON "LiteLLM_DailyModelUsage"("date");
CREATE INDEX IF NOT EXISTS "LiteLLM_DailyModelUsage_model_group_idx" ON "LiteLLM_DailyModelUsage"("model_group");

View file

@ -78,6 +78,11 @@ model LiteLLM_AgentsTable {
object_permission_id String?
object_permission LiteLLM_ObjectPermissionTable? @relation(fields: [object_permission_id], references: [object_permission_id])
spend Float @default(0.0)
identity_managed Boolean @default(false)
enabled Boolean @default(true)
execution_mode String @default("autonomous")
identity LiteLLM_AgentIdentity?
retired_identities LiteLLM_RetiredAgentIdentity[]
tpm_limit Int?
rpm_limit Int?
session_tpm_limit Int?
@ -88,6 +93,56 @@ model LiteLLM_AgentsTable {
updated_by String
}
model LiteLLM_AgentIdentity {
agent_id String @id
active Boolean @default(true)
agent LiteLLM_AgentsTable @relation(fields: [agent_id], references: [agent_id], onDelete: Cascade)
provider String
issuer String
tenant_id String
client_id String
service_principal_id String?
required_roles String[] @default([])
required_scopes String[] @default(["user_impersonation"])
revision String @default(uuid())
last_authenticated_at DateTime?
@@unique([provider, tenant_id, client_id])
@@unique([issuer, service_principal_id])
}
model LiteLLM_RetiredAgentIdentity {
binding_id String @id @default(uuid())
agent_id String?
agent LiteLLM_AgentsTable? @relation(fields: [agent_id], references: [agent_id], onDelete: SetNull)
provider String
issuer String
tenant_id String
client_id String
@@unique([provider, tenant_id, client_id])
}
model LiteLLM_RetiredAgent {
original_agent_id String @id
retired_at DateTime @default(now())
}
model LiteLLM_VerifiedSubject {
subject_id String @id @default(uuid())
issuer String
tenant_id String
oid String
kind String @default("human")
user_id String?
user LiteLLM_UserTable? @relation(fields: [user_id], references: [user_id], onDelete: Cascade)
verified_via String @default("sso_interactive")
verified_at DateTime @default(now())
@@unique([issuer, tenant_id, oid])
@@index([user_id])
}
model LiteLLM_OrganizationTable {
organization_id String @id @default(uuid())
organization_alias String
@ -241,6 +296,7 @@ model LiteLLM_DeletedTeamTable {
// Track spend, rate limit, budget Users
model LiteLLM_UserTable {
verified_subjects LiteLLM_VerifiedSubject[]
user_id String @id
user_alias String?
team_id String?
@ -322,6 +378,7 @@ model LiteLLM_MCPServerTable {
allowed_tools String[] @default([])
tool_name_to_display_name Json? @default("{}")
tool_name_to_description Json? @default("{}")
pinned_tools Json? @default("{}")
extra_headers String[] @default([])
static_headers Json? @default("{}")
// Admin-configured environment variables interpolated into static_headers
@ -674,6 +731,7 @@ model LiteLLM_SpendLogs {
session_id String?
status String?
mcp_namespaced_tool_name String?
billing_agent_id String?
agent_id String?
proxy_server_request Json? @default("{}")
litellm_call_id String?
@ -1259,6 +1317,26 @@ model LiteLLM_DailyToolSpend {
@@id([date, tool_name])
}
model LiteLLM_DailyModelUsage {
date String
model_group String
model String
custom_llm_provider String
task_type String
spend Float @default(0.0)
prompt_tokens BigInt @default(0)
completion_tokens BigInt @default(0)
request_count BigInt @default(0)
successful_requests BigInt @default(0)
failed_requests BigInt @default(0)
created_at DateTime @default(now())
updated_at DateTime @updatedAt
@@id([date, model_group, model, custom_llm_provider, task_type])
@@index([date])
@@index([model_group])
}
// Gateway request counts recorded at the ASGI edge by
// BillableRequestMetricsMiddleware. This is the source of truth for SGR
// (successful gateway requests): it counts what the proxy actually answered,

View file

@ -1,6 +1,6 @@
[project]
name = "litellm-proxy-extras"
version = "0.4.102"
version = "0.4.103"
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.102"
version = "0.4.103"
version_files = [
"pyproject.toml:^version",
"../pyproject.toml:litellm-proxy-extras==",

View file

@ -3552,8 +3552,10 @@ name = "litellm-cache-response"
version = "0.1.0"
dependencies = [
"litellm-cache",
"litellm-cache-gcs",
"litellm-cache-memory",
"litellm-cache-redis",
"litellm-http",
"py_literal",
"redis",
"redis-test",
@ -3562,6 +3564,7 @@ dependencies = [
"serde_json",
"sha2 0.10.9",
"tokio",
"wiremock",
]
[[package]]
@ -3645,13 +3648,18 @@ dependencies = [
"litellm-auth",
"litellm-auth-aws",
"litellm-auth-gcp",
"litellm-cache",
"litellm-cache-memory",
"litellm-cache-response",
"litellm-core-utils",
"litellm-framing",
"litellm-host",
"litellm-host-native",
"litellm-http",
"litellm-llms",
"litellm-llms-types",
"litellm-secrets",
"litellm-tracing",
"litellm-types",
"mime_guess",
"moka",
"rand 0.8.7",
@ -3667,6 +3675,7 @@ dependencies = [
"time",
"tokio",
"tokio-tungstenite",
"tokio-util",
"tracing",
"url",
"veil",
@ -3678,13 +3687,12 @@ name = "litellm-core-utils"
version = "0.1.0"
dependencies = [
"fancy-regex 0.19.2",
"litellm-llms-types",
"litellm-tracing",
"litellm-types",
"rstest",
"serde",
"serde_json",
"serde_path_to_error",
"serde_with",
"strum",
"thiserror 2.0.19",
"url",
@ -3806,15 +3814,19 @@ dependencies = [
"bytes",
"futures-util",
"litellm-auth",
"litellm-cache-memory",
"litellm-cache-response",
"litellm-core",
"litellm-gateway-auth",
"litellm-host",
"litellm-host-http",
"litellm-http",
"litellm-llms",
"litellm-llms-types",
"litellm-router",
"litellm-secrets",
"litellm-types",
"rstest",
"serde",
"serde_json",
"thiserror 2.0.19",
"tokio",
@ -3907,12 +3919,24 @@ dependencies = [
"futures-util",
"http 1.4.2",
"litellm-host",
"litellm-host-native",
"rstest",
"serde_json",
"thiserror 2.0.19",
"tokio",
]
[[package]]
name = "litellm-host-native"
version = "0.1.0"
dependencies = [
"futures-util",
"litellm-host",
"rstest",
"serde_json",
"tokio",
]
[[package]]
name = "litellm-host-python"
version = "0.1.0"
@ -3973,9 +3997,9 @@ dependencies = [
"litellm-framing",
"litellm-host",
"litellm-http",
"litellm-llms-types",
"litellm-python-compat",
"litellm-secrets",
"litellm-types",
"reqwest 0.12.28",
"rstest",
"serde",
@ -3989,13 +4013,26 @@ dependencies = [
"url",
]
[[package]]
name = "litellm-llms-types"
version = "0.1.0"
dependencies = [
"macro_rules_attribute",
"rstest",
"schemars 1.2.2",
"serde",
"serde_json",
"serde_with",
"strum",
]
[[package]]
name = "litellm-model-catalog"
version = "0.1.0"
dependencies = [
"indexmap 2.14.0",
"jsonschema",
"litellm-types",
"litellm-llms-types",
"rstest",
"schemars 1.2.2",
"serde",
@ -4033,12 +4070,13 @@ dependencies = [
"litellm-host-python",
"litellm-http",
"litellm-llms",
"litellm-llms-types",
"litellm-secrets",
"litellm-secrets-aws",
"litellm-secrets-types",
"litellm-token-counter",
"litellm-traces",
"litellm-tracing",
"litellm-types",
"pyo3",
"pyo3-async-runtimes",
"qdrant-client",
@ -4314,6 +4352,21 @@ dependencies = [
"tiktoken-rs",
]
[[package]]
name = "litellm-traces"
version = "0.1.0"
dependencies = [
"litellm-http",
"rstest",
"serde",
"serde_json",
"testcontainers-modules",
"thiserror 2.0.19",
"time",
"tokio",
"url",
]
[[package]]
name = "litellm-tracing"
version = "0.1.0"
@ -4328,17 +4381,6 @@ dependencies = [
"tracing-subscriber",
]
[[package]]
name = "litellm-types"
version = "0.1.0"
dependencies = [
"rstest",
"schemars 1.2.2",
"serde",
"serde_json",
"strum",
]
[[package]]
name = "litemap"
version = "0.8.2"
@ -5678,6 +5720,7 @@ checksum = "16a1cfa75cc186dd73d5818e510e042e40927bccc9c236b061cea97e1eb08029"
dependencies = [
"base64 0.23.1",
"bytes",
"encoding_rs",
"futures-core",
"futures-util",
"h2 0.4.15",
@ -5689,6 +5732,7 @@ dependencies = [
"hyper-util",
"js-sys",
"log",
"mime",
"percent-encoding",
"pin-project-lite",
"quinn",
@ -6919,6 +6963,7 @@ dependencies = [
"memchr",
"parse-display",
"pin-project-lite",
"reqwest 0.13.5",
"serde",
"serde_json",
"serde_with",

View file

@ -12,6 +12,7 @@ repository = "https://github.com/BerriAI/litellm"
litellm-config = { path = "crates/config" }
litellm-router = { path = "crates/router" }
litellm-tracing = { path = "crates/tracing" }
litellm-traces = { path = "crates/traces" }
litellm-core = { path = "crates/core" }
litellm-gateway-mcp = { path = "crates/gateway-mcp" }
litellm-gateway = { path = "crates/gateway" }
@ -22,6 +23,7 @@ litellm-gateway-ui = { path = "crates/gateway-ui" }
litellm-coroutine = { path = "crates/coroutine" }
litellm-host = { path = "crates/host" }
litellm-host-http = { path = "crates/host-http" }
litellm-host-native = { path = "crates/host-native" }
litellm-callbacks-legacy-python = { path = "crates/callbacks-legacy-python" }
litellm-framing = { path = "crates/framer" }
litellm-auth = { path = "crates/auth" }
@ -38,7 +40,7 @@ litellm-secrets-azure = { path = "crates/secrets-azure" }
litellm-secrets-cyberark = { path = "crates/secrets-cyberark" }
litellm-http = { path = "crates/http" }
litellm-llms = { path = "crates/llms" }
litellm-types = { path = "crates/types" }
litellm-llms-types = { path = "crates/llms-types" }
litellm-core-utils = { path = "crates/core-utils" }
litellm-db = { path = "crates/db" }
litellm-db-testing = { path = "crates/db-testing" }
@ -73,6 +75,7 @@ proptest = "1.7.0"
pyo3 = "0.29.2"
pyo3-async-runtimes = { version = "0.29.0", features = ["tokio-runtime"] }
rand = "0.8"
macro_rules_attribute = "0.2.3"
schemars = "1"
reqwest = { version = "0.12", default-features = false, features = ["json", "multipart", "rustls-tls", "http2", "stream"] }
qdrant-client = { version = "1.19.0", default-features = false }

View file

@ -39,7 +39,33 @@ impl TokenProviderHandle {
Self(caller)
}
pub fn from_callback<F, Fut>(acquire: F) -> Self
where
F: Fn() -> Fut + Send + Sync + 'static,
Fut: Future<Output = Result<ResolvedCredential, Error>> + Send + 'static,
{
Self::new(Arc::new(CallbackTokenProvider(acquire)))
}
pub async fn acquire(&self) -> Result<ResolvedCredential, Error> {
self.0.acquire().await
}
}
struct CallbackTokenProvider<F>(F);
impl<F> std::fmt::Debug for CallbackTokenProvider<F> {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
formatter.write_str("CallbackTokenProvider")
}
}
impl<F, Fut> TokenProvider for CallbackTokenProvider<F>
where
F: Fn() -> Fut + Send + Sync,
Fut: Future<Output = Result<ResolvedCredential, Error>> + Send + 'static,
{
fn acquire(&self) -> TokenFuture<'_> {
Box::pin((self.0)())
}
}

View file

@ -0,0 +1,104 @@
use std::{
error::Error as StdError,
future::{Future, poll_fn},
sync::{
Arc,
atomic::{AtomicBool, AtomicUsize, Ordering},
},
task::Poll,
time::{Duration, SystemTime},
};
use litellm_auth_types::{
Error, ErrorDetail, ResolvedCredential, SecretValue, TokenProviderHandle,
};
use rstest::rstest;
fn credential(index: usize, access_token: bool) -> ResolvedCredential {
let token = SecretValue::new(format!("credential-{index}"));
if access_token {
return ResolvedCredential::AccessToken {
token,
expires_on: Some(SystemTime::UNIX_EPOCH + Duration::from_secs(index as u64)),
};
}
ResolvedCredential::Static(token)
}
#[rstest]
#[case::static_secret(false)]
#[case::access_token(true)]
#[tokio::test]
async fn callbacks_acquire_fresh_credentials_on_demand(#[case] access_token: bool) {
let calls = Arc::new(AtomicUsize::new(0));
let callback_calls = calls.clone();
let provider = TokenProviderHandle::from_callback(move || {
let index = callback_calls.fetch_add(1, Ordering::SeqCst);
async move {
tokio::task::yield_now().await;
Ok(credential(index, access_token))
}
});
let cloned = provider.clone();
assert_eq!(calls.load(Ordering::SeqCst), 0);
assert_eq!(
provider.acquire().await.unwrap(),
credential(0, access_token)
);
assert_eq!(calls.load(Ordering::SeqCst), 1);
assert_eq!(cloned.acquire().await.unwrap(), credential(1, access_token));
assert_eq!(calls.load(Ordering::SeqCst), 2);
}
#[rstest]
#[tokio::test]
async fn callback_errors_preserve_the_original_source() {
let provider = TokenProviderHandle::from_callback(|| async {
Err(Error::CredentialAcquisition(ErrorDetail::failed(
"caller credential",
std::io::Error::from(std::io::ErrorKind::PermissionDenied),
)))
});
let error = provider.acquire().await.unwrap_err();
assert!(matches!(error, Error::CredentialAcquisition(_)));
let source = std::iter::successors(Some(&error as &(dyn StdError + 'static)), |error| {
(*error).source()
})
.find_map(|error| error.downcast_ref::<std::io::Error>())
.unwrap();
assert_eq!(source.kind(), std::io::ErrorKind::PermissionDenied);
}
struct Release(Arc<AtomicBool>);
impl Drop for Release {
fn drop(&mut self) {
self.0.store(true, Ordering::SeqCst);
}
}
#[rstest]
#[tokio::test]
async fn cancelling_acquisition_drops_the_callback_future() {
let released = Arc::new(AtomicBool::new(false));
let callback_released = released.clone();
let provider = TokenProviderHandle::from_callback(move || {
let released = callback_released.clone();
async move {
let _release = Release(released);
std::future::pending().await
}
});
let mut acquisition = Box::pin(provider.acquire());
poll_fn(|context| {
assert!(acquisition.as_mut().poll(context).is_pending());
assert!(!released.load(Ordering::SeqCst));
Poll::Ready(())
})
.await;
drop(acquisition);
assert!(released.load(Ordering::SeqCst));
}

View file

@ -56,6 +56,7 @@ async fn set_writes_encoded_object_and_headers(#[future(awt)] server: MockServer
)]
#[case::missing("missing", ResponseTemplate::new(404), Ok(None))]
#[case::server_error("server-error", ResponseTemplate::new(500), Err(Error::Unavailable))]
#[case::unauthorized("unauthorized", ResponseTemplate::new(401), Err(Error::Unavailable))]
#[case::invalid(
"invalid",
ResponseTemplate::new(200).set_body_string("not json"),

View file

@ -0,0 +1,29 @@
# Response caching
Design this crate for shared Rust execution used by the Python SDK and the Rust gateway. The Python SDK will remain, with more core execution moving to Rust and Python callbacks staying in Python. The Rust gateway is still evolving and is intended to replace the Python proxy. Keep response-cache policy independent of Python, HTTP serving, and either proxy's configuration format
Separate what is cached, how a hit is matched, and where entries are stored. Chat Completions, Messages, Responses, and embeddings are API workloads. Exact and semantic matching are lookup behaviors. Memory, Redis, disk, and object stores are storage choices. Embeddings are inference too, so do not use an inference-cache name to imply a category that excludes embeddings. Consult the existing Python cache and caching handler for behavior and compatibility contracts without copying their class structure
Storage traits, codecs, and backend capabilities belong in `litellm-cache` and the storage crates. Keep storage reusable for value types beyond LLM responses. This crate owns response entries, matching and freshness semantics, the Python-compatible response codec, and deferred-write policy. Core owns route-specific request identity, response encoding and reconstruction, embedding partial-hit orchestration, and stream capture and replay. Boundaries own configuration translation, resource construction, and caller identity
Construct and inject the response-cache service at the Python bridge or gateway boundary, as with the HTTP client. Reuse it across calls. Core and provider code must not discover cache configuration through Python globals, process configuration, or backend-specific factories
Keep `ResponseCache<B>` generic over its storage backend. Preserve typed backend contexts and capability bounds internally. Inject an object-safe service into core for runtime backend selection, so storage types do not spread through route and host types. Keep API request and response types statically typed. Add a generic parameter only where it preserves a useful type relationship or capability
Keep the core service contract narrow. Lookup and store must not require connection testing, ping, flush, deletion, counters, queues, or scripts. Require batch operations where a consumer needs partial hits, and keep management capabilities on their own interfaces. An exact-only adapter must remain explicit about its matching restriction. Supporting semantic matching requires a defined lookup-context and embedding execution contract, not just a renamed trait
Separate reusable resources from per-call policy. Backend configuration, namespace, default expiry, and entry limits belong to the configured service or backend. Read/write controls, expiry and freshness overrides, and authenticated caller scope belong to the call. Passing call options must not replace or mutate the route's configured service
Keep cache misses and storage failures distinguishable in return values. Core owns the decision to continue with provider execution after a cache failure. A read can reject an entry for freshness while the backend still retains it. Preserve the timestamp at which a response was produced when writing it later
Define lookup placement explicitly relative to authorization, deployment and credential resolution, and request-transforming callbacks. Cache identity must account for every input that affects reuse, including API surface and caller scope, while preserving intentional Python caching groups. Preserve existing keys and response formats unless changing them is an explicit migration decision
Cache normalized provider results before caller-specific response transformations. Hits must still run the applicable response processing, success callbacks, and cache-hit accounting. Keep callback execution in the host. Python cache implementations and semantic embedders that require the caller's task must use the existing host-operation mechanism rather than Python calls from a Rust worker. Preserve legacy fallback until that contract is supported
Keep unary caching independent of stream-only methods. Store streams only after successful exhaustion and protocol completion. Errors, incomplete streams, cancellation, and oversized entries must not populate the cache. Embedding batches need ordered partial results and reconstruction around the uncached inputs
Test each contract in its owner: storage capabilities in backend tests, envelopes and freshness here, reuse and replay in core, Python callback and fallback behavior at the bridge, and HTTP behavior at the gateway. Run backend contract checks and Python response-codec fixtures before exposing a new backend
`ScopedCache` requires an explicit shared or isolated scope at construction. Per-call `CachePolicy` controls reads, writes, expiry, and freshness without replacing the attached scope or service. `CacheOptions` binds that policy to an explicit scope for storage requests and has no default sharing policy. Versioned native envelopes reject incompatible API surfaces and versions as misses; this envelope is distinct from the legacy Python response codec
Response storage is not the source of budget or rate-limit coordination dependencies. Keep counters, reservations, and atomic admission operations out of `ResponseCacheService`, including when both services happen to use Redis

View file

@ -13,9 +13,12 @@ serde_json.workspace = true
sha2.workspace = true
[dev-dependencies]
litellm-cache-gcs.workspace = true
litellm-http = { workspace = true, features = ["test-support"] }
litellm-cache-memory.workspace = true
litellm-cache-redis.workspace = true
redis = "1.7.0"
redis-test = "1.0.4"
rstest.workspace = true
tokio.workspace = true
wiremock = "0.6.5"

View file

@ -1,51 +0,0 @@
# Response cache
`ResponseCache<B>` adds request keys, independent read/write controls, response envelopes, and freshness checks to any `B: BaseCache<Value = CacheEntry>`
## Ownership
`litellm-cache` defines typed storage, codec, and capability traits. `BaseCache` is only get, set, TTL, and pipeline writes. Everything else is an optional capability a backend implements only where its Python class defines the method: `DisconnectCache`, `ConnectionCache` (`test_connection`), `PingCache`, `BatchCache`, `DeleteCache`, `FlushCache`, counters, queues, TTL, scan, and scripts. Memory, Redis, disk, S3, GCS, and Azure Blob implement those traits without depending on response policy, so other consumers can store their own value types in the same backends
Semantic backends (Redis, Valkey, Qdrant) are generic over their embedder and codec, and share one prompt and embedding contract from `litellm_cache::semantic`. They take a `SemanticCacheContext`, so `ResponseCache` drives them the same way it drives exact backends
`litellm-cache-response` owns response keys, controls, entries, the Python-compatible response codec, and `WriteBuffer`, the backend-neutral deferred-write policy. It has no runtime dependency on a specific cache backend or Python
`ExactResponseCache` is the object-safe view of a `ResponseCache` over an exact backend. `ConnectionProbe` is the object-safe `test_connection`, implemented only when the backend implements `ConnectionCache`, so a host holds one next to its `ExactResponseCache` and reports the operation as unsupported otherwise, as Python's `BaseCache` does. Lookup, store, batch, and flush never require it
## Native Rust use
```rust
use std::{sync::Arc, time::Duration};
use litellm_cache_memory::InMemoryCache;
use litellm_cache_response::{CacheKeyInput, ResponseCache, ResponseCacheRequest};
use serde_json::json;
let cache = ResponseCache::new(Arc::new(InMemoryCache::default()));
let request = ResponseCacheRequest::new(CacheKeyInput {
preset: Some("example:key".into()),
..Default::default()
});
let now = Duration::from_secs(100);
cache.store(&request, json!({"answer": 7}), now)?;
assert_eq!(cache.async_lookup(&request, now).await?, Some(json!({"answer": 7})));
```
For Redis, inject `RedisCache::new(url, ttl, ResponseCacheCodec)` instead. Namespaces are optional and existing namespace prefixes are preserved
Callers supply Unix time for response freshness. Backend TTL uses its own clock. A read can reject an entry through `max_age` even while the backend still retains it
## Python integration
The bridge activates backends through the Rust catalog in `litellm/rust_bridge/catalog.py`. Every cache rule ships as `PYTHON_ONLY`, so SDK, Router, and proxy calls stay on Python and construct no native cache resources until a rule is changed
When a rule selects a backend, the Python `Cache` facade builds the native runtime from its own configuration and routes its storage calls (sync and async lookup and store, and pipelined batch store) to it. Stream replay, embedding partial-hit merging, response reconstruction, and callbacks stay in Python on top of that native store. The Python backend object remains for its direct API
Object responses are written as they are, and every other response shape is written as a serialized string, which is the pair of shapes Python reads. A string on the wire is therefore always a serialized response, so string-valued responses round trip. Typed backends such as memory never pass through the codec
Native cache handles must be recreated after fork. Native errors propagate to the host, which owns the existing fail-open and logging policy
## Adding another backend
Implement `BaseCache` for the backend with its associated value type and the capability traits its Python class supports, and accept a `CacheCodec` when wire serialization is needed. `ResponseCache<B>` then works without another response implementation
Run the `litellm-cache-testing` contract checks the backend's capabilities allow, and run response fixtures with `ResponseCacheCodec`, including both Python envelope encodings, before adding a catalog rule

View file

@ -4,6 +4,7 @@ mod codec;
mod embedding;
mod exact;
mod response;
mod service;
pub use buffer::WriteBuffer;
pub use caching::{
@ -14,3 +15,8 @@ pub use codec::ResponseCacheCodec;
pub use embedding::PartialHits;
pub use exact::{ConnectionProbe, ExactResponseCache};
pub use response::{ResponseCache, ResponseCacheRequest};
pub use service::{
CacheOptions, CachePolicy, CacheScope, ResponseCacheConfig, ResponseCacheService,
ResponseEnvelope, ScopedCache,
};

View file

@ -7,7 +7,9 @@ use litellm_cache::{
};
use serde_json::Value;
use crate::{CacheControls, CacheEntry, CacheKeyInput, PartialHits, cache_key};
use crate::{
CacheControls, CacheEntry, CacheKeyInput, PartialHits, ResponseCacheConfig, cache_key,
};
#[derive(Clone)]
pub struct ResponseCacheRequest<C: CacheContext = litellm_cache::ExactCacheContext> {
@ -50,6 +52,7 @@ where
B::Context: Default + PartialEq,
{
backend: Arc<B>,
config: ResponseCacheConfig,
}
impl<B> ResponseCache<B>
@ -58,7 +61,18 @@ where
B::Context: Default + PartialEq,
{
pub fn new(backend: Arc<B>) -> Self {
Self { backend }
Self {
backend,
config: ResponseCacheConfig::default(),
}
}
pub fn with_config(self, config: ResponseCacheConfig) -> Self {
Self { config, ..self }
}
pub fn config(&self) -> &ResponseCacheConfig {
&self.config
}
pub fn backend(&self) -> &B {
@ -221,7 +235,7 @@ where
response: Value,
now: Duration,
) -> Result<(), Error> {
if !request.controls.writes() {
if !request.controls.writes() || !self.fits(&response) {
return Ok(());
}
self.backend.set_cache(
@ -240,7 +254,7 @@ where
response: Value,
now: Duration,
) -> Result<(), Error> {
if !request.controls.writes() {
if !request.controls.writes() || !self.fits(&response) {
return Ok(());
}
self.backend
@ -277,7 +291,7 @@ where
) -> Result<(), Error> {
let writable = entries
.into_iter()
.filter(|(request, _, _)| request.controls.writes())
.filter(|(request, response, _)| request.controls.writes() && self.fits(response))
.map(|(request, response, now)| {
(
cache_key(&request.key),
@ -312,6 +326,11 @@ where
Ok(())
}
fn fits(&self, response: &Value) -> bool {
self.config.max_entry_bytes == usize::MAX
|| response.to_string().len() <= self.config.max_entry_bytes
}
fn partial_hits(
requests: &[ResponseCacheRequest<B::Context>],
readable: Vec<(usize, &ResponseCacheRequest<B::Context>)>,

View file

@ -0,0 +1,185 @@
use std::{future::Future, pin::Pin, time::Duration};
use litellm_cache::{BaseCache, Error, ExactCacheContext};
use serde_json::Value;
use crate::{
CacheControls, CacheEntry, CacheKeyField, CacheKeyInput, ResponseCache, ResponseCacheRequest,
};
type CacheFuture<'a, T> = Pin<Box<dyn Future<Output = Result<T, Error>> + Send + 'a>>;
#[derive(Clone)]
pub struct ResponseCacheConfig {
pub namespace: String,
pub max_entry_bytes: usize,
}
impl Default for ResponseCacheConfig {
fn default() -> Self {
Self {
namespace: String::new(),
max_entry_bytes: usize::MAX,
}
}
}
pub trait ResponseCacheService: Send + Sync {
fn config(&self) -> &ResponseCacheConfig;
fn lookup<'a>(
&'a self,
request: &'a ResponseCacheRequest,
now: Duration,
) -> CacheFuture<'a, Option<Value>>;
fn store<'a>(
&'a self,
request: &'a ResponseCacheRequest,
response: Value,
now: Duration,
) -> CacheFuture<'a, ()>;
}
impl<B> ResponseCacheService for ResponseCache<B>
where
B: BaseCache<Value = CacheEntry, Context = ExactCacheContext>,
{
fn config(&self) -> &ResponseCacheConfig {
self.config()
}
fn lookup<'a>(
&'a self,
request: &'a ResponseCacheRequest,
now: Duration,
) -> CacheFuture<'a, Option<Value>> {
Box::pin(self.async_lookup(request, now))
}
fn store<'a>(
&'a self,
request: &'a ResponseCacheRequest,
response: Value,
now: Duration,
) -> CacheFuture<'a, ()> {
Box::pin(self.async_store(request, response, now))
}
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum CacheScope {
Shared,
Isolated(String),
}
#[derive(Clone, Copy, Default)]
pub struct CachePolicy {
pub caching: Option<bool>,
pub no_cache: bool,
pub no_store: bool,
pub ttl: Option<Duration>,
pub max_age: Option<Duration>,
}
impl CachePolicy {
pub fn enabled(&self) -> bool {
self.caching != Some(false) && !(self.no_cache && self.no_store)
}
}
#[derive(Clone)]
pub struct CacheOptions {
pub policy: CachePolicy,
pub scope: CacheScope,
}
impl CacheOptions {
pub fn new(scope: CacheScope) -> Self {
Self {
policy: CachePolicy::default(),
scope,
}
}
pub fn request(self, namespace: &str, surface: &str, mut input: Value) -> ResponseCacheRequest {
input.sort_all_objects();
let scope = match self.scope {
CacheScope::Shared => String::new(),
CacheScope::Isolated(scope) => serde_json::json!(["isolated", scope]).to_string(),
};
ResponseCacheRequest {
key: CacheKeyInput {
namespace: Some(format!("{namespace}:inference-v2")),
fields: [
("surface", surface.to_owned()),
("scope", scope),
("request", input.to_string()),
]
.into_iter()
.map(|(name, value)| CacheKeyField {
name: name.into(),
value: Some(value),
api_parameter: true,
internal_parameter: false,
})
.collect(),
..Default::default()
},
controls: CacheControls {
configured: true,
supported_call_type: true,
native_backend: true,
default_on: true,
caching: self.policy.caching,
no_cache: self.policy.no_cache,
no_store: self.policy.no_store,
..Default::default()
},
context: ExactCacheContext {
ttl: self.policy.ttl,
},
max_age: self.policy.max_age,
}
}
}
#[derive(serde::Serialize, serde::Deserialize)]
pub struct ResponseEnvelope<T> {
version: u32,
surface: String,
output: T,
}
impl<T> ResponseEnvelope<T> {
pub fn new(surface: &str, output: T) -> Self {
Self {
version: 1,
surface: surface.into(),
output,
}
}
pub fn decode(self, surface: &str) -> Option<T> {
(self.version == 1 && self.surface == surface).then_some(self.output)
}
}
#[derive(Clone)]
pub struct ScopedCache {
pub service: std::sync::Arc<dyn ResponseCacheService>,
pub scope: CacheScope,
}
impl ScopedCache {
pub fn new(service: std::sync::Arc<dyn ResponseCacheService>, scope: CacheScope) -> Self {
Self { service, scope }
}
pub fn options(&self, policy: Option<CachePolicy>) -> CacheOptions {
CacheOptions {
policy: policy.unwrap_or_default(),
scope: self.scope.clone(),
}
}
}

View file

@ -18,7 +18,7 @@ use litellm_cache_response::{
WriteBuffer, cache_key,
};
use redis_test::MockCmd;
use rstest::rstest;
use rstest::{fixture, rstest};
use serde_json::{Value, json};
use support::{keyed, memory, redis, request};
@ -648,3 +648,129 @@ async fn write_buffer_clear_drops_pending_entries(memory: Memory, request: Respo
assert_eq!(memory.lookup(&request, now).unwrap(), None);
assert_eq!(memory.lookup(&other, now).unwrap(), None);
}
#[rstest]
#[case::python_sync("{'timestamp': 100.0, 'response': '{\"answer\": 7}'}")]
#[case::python_async(r#"{"timestamp":100.0,"response":{"answer":7}}"#)]
#[case::bare_response(r#"{"answer":7}"#)]
#[tokio::test]
async fn gcs_reads_python_entries_and_writes_python_compatible_envelopes(
#[case] encoded: &str,
#[values(false, true)] asynchronous: bool,
#[future(awt)] gcs: (wiremock::MockServer, Gcs),
) {
use wiremock::{
Mock, ResponseTemplate,
matchers::{body_json, header, method, path, query_param},
};
let (server, cache) = gcs;
let response = json!({"answer": 7});
Mock::given(method("GET"))
.and(path("/storage/v1/b/bucket/o/cache%2Fpython"))
.and(query_param("alt", "media"))
.and(header("authorization", "Bearer token"))
.respond_with(ResponseTemplate::new(200).set_body_string(encoded))
.expect(1)
.mount(&server)
.await;
Mock::given(method("POST"))
.and(path("/upload/storage/v1/b/bucket/o"))
.and(query_param("uploadType", "media"))
.and(query_param("name", "cache/native"))
.and(header("authorization", "Bearer token"))
.and(header("content-type", "application/json"))
.and(body_json(json!({"timestamp": 102.0, "response": response})))
.respond_with(ResponseTemplate::new(200))
.expect(1)
.mount(&server)
.await;
let lookup = if asynchronous {
cache
.async_lookup(&keyed("python"), Duration::from_secs(102))
.await
} else {
cache.lookup(&keyed("python"), Duration::from_secs(102))
};
assert_eq!(lookup.unwrap(), Some(response.clone()));
let request = ResponseCacheRequest {
context: litellm_cache::ExactCacheContext {
ttl: Some(Duration::from_secs(12)),
},
..keyed("native")
};
let stored = if asynchronous {
cache
.async_store(&request, response, Duration::from_secs(102))
.await
} else {
cache.store(&request, response, Duration::from_secs(102))
};
assert_eq!(stored, Ok(()));
let requests = server.received_requests().await.unwrap();
let upload = requests
.iter()
.find(|request| request.method.as_str() == "POST")
.unwrap();
assert_eq!(
upload.url.query(),
Some("uploadType=media&name=cache%2Fnative")
);
}
#[rstest]
#[tokio::test]
async fn gcs_batch_reads_preserve_order_and_treat_invalid_entries_as_misses(
#[future(awt)] gcs: (wiremock::MockServer, Gcs),
) {
use wiremock::{
Mock, ResponseTemplate,
matchers::{method, path},
};
let (server, cache) = gcs;
Mock::given(method("GET"))
.and(path("/storage/v1/b/bucket/o/cache%2Fhit"))
.respond_with(
ResponseTemplate::new(200)
.set_body_json(json!({"timestamp": 100.0, "response": {"answer":7}})),
)
.mount(&server)
.await;
Mock::given(method("GET"))
.and(path("/storage/v1/b/bucket/o/cache%2Finvalid"))
.respond_with(ResponseTemplate::new(200).set_body_string("not an entry"))
.mount(&server)
.await;
Mock::given(method("GET"))
.and(path("/storage/v1/b/bucket/o/cache%2Fmissing"))
.respond_with(ResponseTemplate::new(404))
.mount(&server)
.await;
let requests = [keyed("hit"), keyed("missing"), keyed("invalid")];
let partial = cache
.async_lookup_batch(&requests, Duration::from_secs(102))
.await
.unwrap();
assert_eq!(partial.values, vec![Some(json!({"answer":7})), None, None]);
assert_eq!(partial.missing_indices, vec![1, 2]);
}
type Gcs = ResponseCache<litellm_cache_gcs::GcsCache<litellm_cache_response::ResponseCacheCodec>>;
#[fixture]
async fn gcs() -> (wiremock::MockServer, Gcs) {
let server = wiremock::MockServer::start().await;
let cache = ResponseCache::new(Arc::new(litellm_cache_gcs::GcsCache::with_token_source(
litellm_cache_gcs::GcsConfig {
bucket_name: "bucket".into(),
gcs_path: Some("cache".into()),
path_service_account: None,
endpoint: server.uri(),
},
litellm_http::Client::plain_for_test(),
litellm_cache_response::ResponseCacheCodec,
Arc::new(litellm_cache_gcs::StaticTokenSource("token".into())),
)));
(server, cache)
}

View file

@ -0,0 +1,185 @@
use std::{
sync::{
Arc,
atomic::{AtomicU64, Ordering},
},
time::Duration,
};
use litellm_cache::ExactCacheContext;
use litellm_cache_memory::InMemoryCache;
use litellm_cache_response::{
CacheEntry, CacheKeyInput, ResponseCache, ResponseCacheConfig, ResponseCacheRequest,
ResponseCacheService,
};
use rstest::rstest;
use serde_json::json;
#[rstest]
#[tokio::test]
async fn service_honors_per_call_expiry_and_freshness() {
let clock = Arc::new(AtomicU64::new(0));
let cache_clock = clock.clone();
let cache: Arc<dyn ResponseCacheService> = Arc::new(ResponseCache::new(Arc::new(
InMemoryCache::with_clock(Some(100), Some(Duration::from_secs(60)), move || {
Duration::from_secs(cache_clock.load(Ordering::SeqCst))
}),
)));
let request = ResponseCacheRequest {
context: ExactCacheContext {
ttl: Some(Duration::from_secs(5)),
},
..ResponseCacheRequest::new(CacheKeyInput {
preset: Some("entry".into()),
..Default::default()
})
};
cache
.store(&request, json!({"answer":7}), Duration::ZERO)
.await
.unwrap();
assert_eq!(
cache.lookup(&request, Duration::ZERO).await.unwrap(),
Some(json!({"answer":7}))
);
let stale_request = ResponseCacheRequest {
max_age: Some(Duration::from_secs(1)),
..request.clone()
};
clock.store(2, Ordering::SeqCst);
assert_eq!(
cache
.lookup(&stale_request, Duration::from_secs(2))
.await
.unwrap(),
None
);
assert!(
cache
.lookup(&request, Duration::from_secs(2))
.await
.unwrap()
.is_some()
);
clock.store(6, Ordering::SeqCst);
assert_eq!(
cache
.lookup(&request, Duration::from_secs(6))
.await
.unwrap(),
None
);
}
#[rstest]
#[tokio::test]
async fn entry_limit_applies_to_sync_async_and_batch_writes() {
let storage = Arc::new(InMemoryCache::<CacheEntry>::default());
let cache = ResponseCache::new(storage.clone()).with_config(ResponseCacheConfig {
namespace: "service-test".into(),
max_entry_bytes: json!({"answer":7}).to_string().len(),
});
let small = json!({"answer":7});
let large = json!({"answer":"too large"});
let request = |key: &str| {
ResponseCacheRequest::new(CacheKeyInput {
preset: Some(key.into()),
..Default::default()
})
};
cache
.store(&request("sync"), large.clone(), Duration::ZERO)
.unwrap();
cache
.async_store(&request("async"), large.clone(), Duration::ZERO)
.await
.unwrap();
cache
.async_store_batch(
vec![
(request("batch-large"), large),
(request("batch-small"), small.clone()),
],
Duration::ZERO,
)
.await
.unwrap();
let service: Arc<dyn ResponseCacheService> = Arc::new(cache);
service
.store(&request("service"), small.clone(), Duration::ZERO)
.await
.unwrap();
for key in ["sync", "async", "batch-large"] {
assert!(storage.get_cache(key).unwrap().is_none());
}
for key in ["batch-small", "service"] {
assert_eq!(
service.lookup(&request(key), Duration::ZERO).await.unwrap(),
Some(small.clone())
);
}
}
#[rstest]
#[case::same_scope("tenant-a", "tenant-a", true)]
#[case::different_scope("tenant-a", "tenant-b", false)]
#[case::empty_isolated_scope("", "", true)]
#[tokio::test]
async fn isolated_policy_controls_actual_entry_reuse(
#[case] first: &str,
#[case] second: &str,
#[case] hit: bool,
#[values(false, true)] override_policy: bool,
) {
use litellm_cache_response::{CachePolicy, CacheScope, ScopedCache};
let service = Arc::new(ResponseCache::new(Arc::new(
InMemoryCache::<CacheEntry>::default(),
)));
let request = |scope| {
ScopedCache::new(service.clone(), scope)
.options(override_policy.then_some(CachePolicy {
ttl: Some(Duration::from_secs(30)),
..CachePolicy::default()
}))
.request("test", "messages", json!({"prompt":"hello"}))
};
service
.async_store(
&request(CacheScope::Isolated(first.into())),
json!({"answer":7}),
Duration::ZERO,
)
.await
.unwrap();
assert_eq!(
service
.async_lookup(
&request(CacheScope::Isolated(second.into())),
Duration::ZERO
)
.await
.unwrap(),
hit.then(|| json!({"answer":7}))
);
assert_eq!(
service
.async_lookup(&request(CacheScope::Shared), Duration::ZERO)
.await
.unwrap(),
None
);
}
#[rstest]
#[case::valid(1, "messages", Some(7))]
#[case::unknown_version(2, "messages", None)]
#[case::another_surface(1, "responses", None)]
fn envelopes_require_a_matching_surface_and_version(
#[case] version: u32,
#[case] surface: &str,
#[case] expected: Option<u32>,
) {
let envelope: litellm_cache_response::ResponseEnvelope<u32> =
serde_json::from_value(json!({"version":version,"surface":surface,"output":7})).unwrap();
assert_eq!(envelope.decode("messages"), expected);
}

View file

@ -1,12 +1,13 @@
- Target invariants, not completion claims
- This crate is the legacy `@client` wrapper as the native call sees it, and nothing else: the `Logging` contract (`function_setup`, the deployment hooks, `pre_call`/`post_call`, the sync and async success and failure fan-out, the deferred proxy release, the argument sharing those callbacks rely on)
- This crate owns compatibility for all existing Python callbacks and loggers, including `CustomLogger`. `mapping.rs` owns the executable call bindings and the inventory of Python-owned hooks. A Python-owned entry records an existing path, never permission to invoke it a second time. The native call adapter preserves the `Logging` contract (`function_setup`, the deployment hooks, `pre_call`/`post_call`, the sync and async success and failure fan-out, the deferred proxy release, the argument sharing those callbacks rely on)
- Smell test: if a future callback host (`callbacks-v1-python`, WASM, in-process Rust) could share a piece of this crate, it does not belong here
- SDK request policy (credential inheritance, the budget and retry-count limits) is the driver's preflight, supplied by `python-bridge`; this crate only adopts the keyword view it produces
- The driver in `litellm-host-python`, the routes and core see one `PythonCallHooks`; they never learn which Python objects consume a call
- SDK request policy (credential inheritance, the budget and retry-count limits) is a separate hook supplied by `python-bridge`; compose it after this adapter so logging adopts the final keyword view before policy mutates or rejects it
- The driver in `litellm-host-python`, the routes and core see one `PythonCallHooks` using the shared `CallEvent`; they never learn which Python objects consume a call
- Every litellm Python internal Rust still borrows is a variant of `LegacyPython`, grouped by subsystem, with its signature pinned in `python_contract.json`
- The enum only shrinks: when Rust owns a subsystem, delete its group rather than adding a Rust path beside it
- Calling a user's own callback directly is permanent Python surface and gets its own type outside `LegacyPython`
- `PublicCall` is the caller's call as `Logging` sees it: the positional arguments, the keyword view as the call rewrites it (setup, deployment hook, preflight) and the bound request object backing omitted keywords; routes hand it over through `run_legacy_call` and keep no copy
- `PublicCall` is the caller's call as `Logging` sees it: the positional arguments, the keyword view as the call rewrites it (setup, deployment hook, preflight) and the bound request object backing omitted keywords; shared bridge composition hands it to `LegacyLogging`; routes use the neutral call boundary
- `LoggingOperation` selects legacy logging entrypoints and response handling. It belongs here rather than in shared inference data contracts
- `setup` reuses a `Logging` passed as `litellm_logging_obj` (the proxy and Router) and otherwise builds one through `function_setup`; which callbacks run is `Logging`'s decision, never this crate's
- Callbacks receive the caller's own objects and may mutate them; this crate alone carries that obligation
- Retain complete boundary arguments, opaque values, aliases, omitted/default distinctions and deliberate copies; preserve the deployment-hook kwargs view

View file

@ -2,12 +2,14 @@
//! raises is answered with the same `Logging` calls, in the same order, as the Python
//! `@client` path makes them.
use crate::LoggingOperation;
use litellm_host_python::PythonOwned;
use litellm_host::event::{
FailureOrigin, MachineEvent, RequestContext, Timing, WireRequest, epoch_seconds,
use litellm_host::{
interceptors::{RawResponse, RequestContext, WireRequest},
lifecycle::{FailureOrigin, Timing, epoch_seconds},
};
use litellm_host_python::{HookEvent, HookStep, PythonCallHooks, from_py, missing_state, to_py};
use litellm_host_python::{HookStep, from_py, missing_state, to_py};
use pyo3::{
exceptions::{PyBaseException, PyException},
gc::{PyTraverseError, PyVisit},
@ -24,22 +26,10 @@ use crate::{
setup,
};
/// What the legacy contract needs to know about the route it is logging.
#[derive(Clone, Copy, Debug)]
pub struct LegacySurface {
pub call_type: &'static str,
/// What `Logging.pre_call` is told the input was.
pub input_description: &'static str,
/// How a streamed response is billed; `None` for a route that never streams.
pub stream: Option<PassThroughStream>,
}
/// The pass-through billing a streamed response goes through once its chunks are in.
#[derive(Clone, Copy, Debug)]
pub struct PassThroughStream {
pub url_route: &'static str,
/// A value of Python's `EndpointType`.
pub endpoint_type: &'static str,
struct PassThroughStream {
url_route: &'static str,
endpoint_type: &'static str,
}
/// What the Messages stream iterator keeps for its end-of-stream billing.
@ -55,7 +45,7 @@ struct LoggedRequest {
}
pub struct LegacyLogging {
surface: LegacySurface,
operation: LoggingOperation,
call: PublicCall,
logger: Option<PythonLogger>,
start: Py<PyAny>,
@ -66,6 +56,7 @@ pub struct LegacyLogging {
stream: Option<DeliveredStream>,
asynchronous: bool,
internal: bool,
cache_key: Option<String>,
}
fn datetime(py: Python<'_>, epoch_seconds: f64) -> PyResult<Py<PyAny>> {
@ -79,12 +70,12 @@ fn is_cancellation(py: Python<'_>, error: &PyErr) -> bool {
impl LegacyLogging {
pub fn new(
py: Python<'_>,
surface: LegacySurface,
operation: LoggingOperation,
call: PublicCall,
asynchronous: bool,
) -> Self {
Self {
surface,
operation,
call,
logger: None,
start: py.None(),
@ -95,9 +86,47 @@ impl LegacyLogging {
stream: None,
asynchronous,
internal: false,
cache_key: None,
}
}
fn call_type(&self) -> &'static str {
match (self.operation, self.asynchronous) {
(LoggingOperation::Completion, false) => "completion",
(LoggingOperation::Completion, true) => "acompletion",
(LoggingOperation::Responses, false) => "responses",
(LoggingOperation::Responses, true) => "aresponses",
(LoggingOperation::Messages, _) => "anthropic_messages",
(LoggingOperation::Ocr, false) => "ocr",
(LoggingOperation::Ocr, true) => "aocr",
}
}
fn input_description(&self) -> &'static str {
match self.operation {
LoggingOperation::Completion => "Chat completions",
LoggingOperation::Responses => "Responses",
LoggingOperation::Messages => "Messages",
LoggingOperation::Ocr => "OCR document processing",
}
}
fn stream_billing(&self) -> Option<PassThroughStream> {
match self.operation {
LoggingOperation::Messages => Some(PassThroughStream {
url_route: "/v1/messages",
endpoint_type: "anthropic",
}),
LoggingOperation::Completion | LoggingOperation::Responses | LoggingOperation::Ocr => {
None
}
}
}
pub(crate) fn adopt_arguments(&mut self, py: Python<'_>, arguments: &Py<PyDict>) {
self.call.set_kwargs(arguments.clone_ref(py));
}
/// Deployment hooks are awaited, and Python's synchronous `@client` wrapper never
/// runs them.
fn runs_deployment_hooks(&self) -> bool {
@ -112,7 +141,6 @@ impl LegacyLogging {
/// The keyword view the rest of the call reads: a copy, so the deployment hook's own
/// dict is left as the hook returned it, carrying the logger as `@client` injects it.
/// The driver's preflight rewrites this same dict before the host projects from it.
fn prepare(&mut self, py: Python<'_>) -> PyResult<HookStep<Self, Py<PyDict>>> {
let prepared = self.call.kwargs().bind(py).copy()?;
prepared.set_item("litellm_logging_obj", self.logger()?.object(py))?;
@ -181,17 +209,17 @@ impl LegacyLogging {
fn stream_success(&self, py: Python<'_>, stream: &DeliveredStream) -> PyResult<()> {
let logger = self.logger()?;
let billing = self.surface.stream.ok_or_else(missing_state)?;
let billing = self.stream_billing().ok_or_else(missing_state)?;
let billed = Streaming::Success.call(
py,
(
logger.object(py),
billing.url_route,
billing.endpoint_type,
&self
.request
.as_ref()
.map(|request| request.body.clone_ref(py)),
&self.request.as_ref().map_or_else(
|| self.call.kwargs().clone_ref(py),
|request| request.body.clone_ref(py),
),
&stream.chunks,
&self.start,
&self.end,
@ -211,9 +239,12 @@ impl LegacyLogging {
/// partial usage. The sync path has no loop to schedule that on, so it falls back to
/// the plain failure handler.
fn stream_failure(&mut self, py: Python<'_>) -> PyResult<HookStep<Self, ()>> {
let (Some(logger), Some(error), Some(stream), Some(billing)) =
(&self.logger, &self.error, &self.stream, self.surface.stream)
else {
let (Some(logger), Some(error), Some(stream), Some(billing)) = (
&self.logger,
&self.error,
&self.stream,
self.stream_billing(),
) else {
return Ok(HookStep::Ready(()));
};
if !self.asynchronous {
@ -224,10 +255,10 @@ impl LegacyLogging {
(
logger.object(py),
billing.endpoint_type,
&self
.request
.as_ref()
.map(|request| request.body.clone_ref(py)),
&self.request.as_ref().map_or_else(
|| self.call.kwargs().clone_ref(py),
|request| request.body.clone_ref(py),
),
&stream.chunks,
error,
),
@ -309,8 +340,8 @@ impl LegacyLogging {
}
}
impl PythonCallHooks for LegacyLogging {
fn prepare_arguments(
impl LegacyLogging {
pub(crate) fn prepare_call(
&mut self,
py: Python<'_>,
arguments: Py<PyDict>,
@ -321,7 +352,7 @@ impl PythonCallHooks for LegacyLogging {
self.internal = is_internal_call(py)?;
let result = setup(
py,
self.surface.call_type,
self.call_type(),
self.call.args(),
self.call.kwargs(),
&self.start,
@ -331,14 +362,14 @@ impl PythonCallHooks for LegacyLogging {
self.call.set_kwargs(result.kwargs()?);
if self.runs_deployment_hooks() {
return Ok(HookStep::Await(
DeploymentHooks::before_call(py, self.call.kwargs(), self.surface.call_type)?,
DeploymentHooks::before_call(py, self.call.kwargs(), self.call_type())?,
Self::resume_begin,
));
}
self.prepare(py)
}
fn before_provider_request(
pub(crate) fn pre_call(
&mut self,
py: Python<'_>,
wire: Box<WireRequest>,
@ -367,7 +398,7 @@ impl PythonCallHooks for LegacyLogging {
});
self.logger()?.pre_call(
py,
self.surface.input_description,
self.input_description(),
context.api_key.as_ref().map(|api_key| api_key.expose()),
&body,
&headers,
@ -384,7 +415,7 @@ impl PythonCallHooks for LegacyLogging {
})))
}
fn transform_response(
pub(crate) fn transform_public_response(
&mut self,
py: Python<'_>,
response: Py<PyAny>,
@ -398,7 +429,7 @@ impl PythonCallHooks for LegacyLogging {
py,
self.call.kwargs(),
&self.response,
self.surface.call_type,
self.call_type(),
)?,
Self::resume_after_success,
));
@ -406,67 +437,105 @@ impl PythonCallHooks for LegacyLogging {
self.finalize(py)
}
fn on_event(&mut self, py: Python<'_>, event: HookEvent<'_>) -> PyResult<HookStep<Self, ()>> {
match event {
HookEvent::Started { .. } => Ok(HookStep::Ready(())),
HookEvent::Machine(MachineEvent::ResponseReceived { raw }) => {
let api_key = self
.request
.as_ref()
.and_then(|request| request.context.api_key.as_ref())
.map(|api_key| api_key.expose());
self.logger()?.post_call(
py,
&raw.body,
api_key,
self.request.as_ref().map(|request| &request.body),
self.request.as_ref().map(|request| &request.headers),
)?;
Ok(HookStep::Ready(()))
}
HookEvent::Succeeded { timing, response } => {
self.end = Some(datetime(py, timing.end_time)?);
self.response = Some(response.clone_ref(py));
match &self.stream {
Some(stream) => self.stream_success(py, stream)?,
None => self.dispatch_success(py)?,
}
Ok(HookStep::Ready(()))
}
HookEvent::Failed {
timing,
origin,
error,
} => {
self.end = Some(datetime(py, timing.end_time)?);
self.error = Some(error.clone_ref(py).into_value(py));
if self.stream.is_some() {
return self.stream_failure(py);
}
if origin == FailureOrigin::Call
&& self.logger.is_some()
&& self.runs_deployment_hooks()
{
let error = self.error.as_ref().ok_or_else(missing_state)?;
return Ok(HookStep::Await(
DeploymentHooks::after_failure(
py,
self.call.kwargs(),
error,
self.surface.call_type,
)?,
Self::resume_deployment_failure,
));
}
self.dispatch_failure(py)
}
}
pub(crate) fn result_ready(
&mut self,
py: Python<'_>,
facts: &litellm_host::interceptors::ExecutionFacts,
) -> PyResult<HookStep<Self, ()>> {
use litellm_host::interceptors::ResultSource;
let logger = self.logger()?.object(py);
let params = logger
.getattr("litellm_params")?
.cast_into::<PyDict>()?
.copy()?;
params.set_item("custom_llm_provider", &facts.provider.provider)?;
crate::python::Logging::Update.call(
py,
(
&logger,
self.call.kwargs(),
&facts.provider.model,
logger.getattr("optional_params")?,
params,
&facts.provider.provider,
),
)?;
let details = logger.getattr("model_call_details")?;
self.cache_key = match &facts.source {
ResultSource::Provider => None,
ResultSource::Cache { key } => Some(key.clone()),
};
details.set_item("cache_hit", self.cache_key.is_some())?;
details.set_item("cache_key", self.cache_key.as_deref())?;
Ok(HookStep::Ready(()))
}
fn on_stream_open(&mut self, py: Python<'_>) -> PyResult<()> {
if self.surface.stream.is_none() {
pub(crate) fn post_call(
&mut self,
py: Python<'_>,
raw: &RawResponse,
) -> PyResult<HookStep<Self, ()>> {
let api_key = self
.request
.as_ref()
.and_then(|request| request.context.api_key.as_ref())
.map(|api_key| api_key.expose());
self.logger()?.post_call(
py,
&raw.body,
api_key,
self.request.as_ref().map(|request| &request.body),
self.request.as_ref().map(|request| &request.headers),
)?;
Ok(HookStep::Ready(()))
}
pub(crate) fn succeeded(
&mut self,
py: Python<'_>,
timing: Timing,
response: &Py<PyAny>,
) -> PyResult<HookStep<Self, ()>> {
self.end = Some(datetime(py, timing.end_time)?);
self.response = Some(response.clone_ref(py));
match &self.stream {
Some(stream) => self.stream_success(py, stream)?,
None => self.dispatch_success(py)?,
}
Ok(HookStep::Ready(()))
}
pub(crate) fn failed(
&mut self,
py: Python<'_>,
timing: Timing,
origin: FailureOrigin,
error: &PyErr,
) -> PyResult<HookStep<Self, ()>> {
self.end = Some(datetime(py, timing.end_time)?);
self.error = Some(error.clone_ref(py).into_value(py));
if self.stream.is_some() {
return self.stream_failure(py);
}
if origin == FailureOrigin::Call && self.logger.is_some() && self.runs_deployment_hooks() {
let error = self.error.as_ref().ok_or_else(missing_state)?;
return Ok(HookStep::Await(
DeploymentHooks::after_failure(py, self.call.kwargs(), error, self.call_type())?,
Self::resume_deployment_failure,
));
}
self.dispatch_failure(py)
}
pub(crate) fn stream_opened(&mut self, py: Python<'_>, head: &Py<PyAny>) -> PyResult<()> {
if self.stream_billing().is_none() {
return Err(missing_state());
}
if let Some(key) = &self.cache_key {
head.bind(py).set_item("cache_key", key)?;
head.bind(py).set_item("cache_hit", true)?;
}
Streaming::Opened.call(py, (self.logger()?.object(py),))?;
self.stream = Some(DeliveredStream {
chunks: PyList::empty(py).unbind(),
@ -475,7 +544,7 @@ impl PythonCallHooks for LegacyLogging {
Ok(())
}
fn on_stream_chunk(&mut self, py: Python<'_>, chunk: &Py<PyAny>) -> PyResult<()> {
pub(crate) fn stream_chunk(&mut self, py: Python<'_>, chunk: &Py<PyAny>) -> PyResult<()> {
let stream = self.stream.as_mut().ok_or_else(missing_state)?;
if stream.first_chunk.is_none() {
stream.first_chunk = Some(datetime(py, epoch_seconds())?);
@ -519,8 +588,9 @@ impl PythonOwned for LegacyLogging {
mod deployment_hooks_tests {
use std::ffi::CStr;
use litellm_host::event::{FailureOrigin, Timing};
use litellm_host_python::{HookEvent, HookStep, PythonCallHooks};
use litellm_host::hooks::CallHooks;
use litellm_host::lifecycle::{FailureOrigin, Timing};
use litellm_host_python::{HookStep, PythonCallEvent};
use pyo3::exceptions::asyncio::CancelledError;
use pyo3::prelude::*;
use pyo3::types::PyDict;
@ -579,6 +649,47 @@ kwargs = {'logger': logger, 'document': document}
matches!(step, HookStep::Await(_, _))
}
#[rstest]
#[case::sync_completion(crate::LoggingOperation::Completion, false, "completion")]
#[case::async_completion(crate::LoggingOperation::Completion, true, "acompletion")]
#[case::sync_responses(crate::LoggingOperation::Responses, false, "responses")]
#[case::async_responses(crate::LoggingOperation::Responses, true, "aresponses")]
#[case::sync_messages(crate::LoggingOperation::Messages, false, "anthropic_messages")]
#[case::async_messages(crate::LoggingOperation::Messages, true, "anthropic_messages")]
#[case::sync_ocr(crate::LoggingOperation::Ocr, false, "ocr")]
#[case::async_ocr(crate::LoggingOperation::Ocr, true, "aocr")]
fn operation_selects_the_legacy_setup_and_deployment_hook_contract(
#[case] operation: crate::LoggingOperation,
#[case] asynchronous: bool,
#[case] expected: &str,
) {
Python::initialize();
Python::attach(|py| {
let locals = namespace(py, CALL);
let mut logging = LegacyLogging {
operation,
..legacy_call(py, &locals, asynchronous)
};
let kwargs = local(&locals, "kwargs")
.cast_into::<PyDict>()
.unwrap()
.unbind();
let step = logging.prepare_arguments(py, kwargs, 0.0).unwrap();
assert_eq!(awaits_deployment_hook(&step), asynchronous);
locals.set_item("expected", expected).unwrap();
locals.set_item("asynchronous", asynchronous).unwrap();
run(
py,
&locals,
c"
assert logger.setup_call_type == expected
if asynchronous:
assert logger.calls == [('pre_hook', expected)]
",
);
});
}
#[rstest]
#[case::synchronous(false)]
#[case::asynchronous(true)]
@ -775,7 +886,7 @@ assert finalized is replacement
)
.unwrap();
let failure = PyErr::from_value(local(&locals, "failure"));
let failed = HookEvent::Failed {
let failed = PythonCallEvent::Failed {
timing: TIMING,
origin: FailureOrigin::Call,
error: &failure,
@ -808,8 +919,10 @@ mod payload_tests {
use std::ffi::CStr;
use litellm_auth::SecretValue;
use litellm_host::event::{MachineEvent, RawResponse, RequestContext, WireRequest};
use litellm_host_python::{HookEvent, HookStep, PythonCallHooks, PythonOwned, to_py};
use litellm_host::hooks::CallHooks;
use litellm_host::interceptors::{RawResponse, RequestContext, WireRequest};
use litellm_host::lifecycle::ExecutionEvent;
use litellm_host_python::{HookStep, PythonCallEvent, PythonOwned, to_py};
use proptest::prelude::*;
use pyo3::gc::{PyTraverseError, PyVisit};
use pyo3::prelude::*;
@ -833,6 +946,7 @@ class PayloadLogger(StubLogger):
def pre_call(self, input, api_key, additional_args):
self.record('pre_call', None)
self.pre = additional_args
self.pre_input = input
self.pre_api_key = api_key
on_pre_call(additional_args)
@ -924,13 +1038,18 @@ check = lambda: None
let step = logging
.before_provider_request(py, Box::new(wire), context)
.unwrap();
let raw = MachineEvent::ResponseReceived {
raw: RawResponse {
body: "raw response".into(),
},
let raw = RawResponse {
body: "raw response".into(),
};
assert!(matches!(
logging.on_event(py, HookEvent::Machine(&raw)).unwrap(),
logging
.on_event(
py,
PythonCallEvent::Execution(ExecutionEvent::ProviderResponseReceived {
raw: &raw
})
)
.unwrap(),
HookStep::Ready(())
));
(logging, step)
@ -975,6 +1094,57 @@ check = lambda: None
}
}
#[rstest]
#[case::completion(crate::LoggingOperation::Completion, "Chat completions")]
#[case::responses(crate::LoggingOperation::Responses, "Responses")]
#[case::messages(crate::LoggingOperation::Messages, "Messages")]
#[case::ocr(crate::LoggingOperation::Ocr, "OCR document processing")]
fn prepared_arguments_replace_the_legacy_view_without_losing_callback_aliases(
#[case] operation: crate::LoggingOperation,
#[case] description: &str,
) {
Python::initialize();
Python::attach(|py| {
let locals = namespace(py, PAYLOAD_LOGGER);
run(
py,
&locals,
c"
original = [0]
replacement = [1]
kwargs['pages'] = original
prepared = {'pages': replacement}
",
);
let mut logging = LegacyLogging {
operation,
logger: Some(PythonLogger::new(local(&locals, "logger").unbind())),
..legacy_call(py, &locals, false)
};
let prepared = local(&locals, "prepared")
.cast_into::<pyo3::types::PyDict>()
.unwrap()
.unbind();
logging.arguments_prepared(py, &prepared).unwrap();
let wire = WireRequest {
body: json!({"pages": [1]}),
..route_wire()
};
let (_, step) = send_and_receive(py, &mut logging, wire, &route_context());
assert!(matches!(step, HookStep::Ready(_)));
locals.set_item("description", description).unwrap();
run(
py,
&locals,
c"
assert logger.pre['complete_input_dict']['pages'] is replacement
assert logger.pre_input == description
assert original == [0]
",
);
});
}
#[rstest::rstest]
fn a_cycle_through_the_retained_headers_is_collected() {
Python::initialize();
@ -1459,8 +1629,9 @@ def check():
mod terminal_tests {
use std::ffi::CStr;
use litellm_host::event::{FailureOrigin, Timing};
use litellm_host_python::{HookEvent, HookStep, PythonCallHooks, PythonOwned};
use litellm_host::hooks::CallHooks;
use litellm_host::lifecycle::{FailureOrigin, Timing};
use litellm_host_python::{HookStep, PythonCallEvent, PythonOwned};
use pyo3::exceptions::PyRuntimeError;
use pyo3::exceptions::asyncio::CancelledError;
use pyo3::prelude::*;
@ -1492,7 +1663,7 @@ mod terminal_tests {
logging
.on_event(
py,
HookEvent::Succeeded {
PythonCallEvent::Succeeded {
timing: TIMING,
response: &response,
},
@ -1509,7 +1680,7 @@ mod terminal_tests {
logging
.on_event(
py,
HookEvent::Failed {
PythonCallEvent::Failed {
timing: TIMING,
origin: FailureOrigin::Host,
error: &failure,
@ -1558,6 +1729,77 @@ assert hasattr(logger, '_native_pending_logging') == getattr(logger, '_defer_asy
});
}
#[rstest]
fn dropped_observations_preserve_deferred_success_and_response_identity() {
Python::initialize();
Python::attach(|py| {
let locals = namespace(
py,
c"response = object()\nlogger._defer_async_logging = True",
);
let mut logging = logged(py, &locals, true);
let response = local(&locals, "response").unbind();
let event = PythonCallEvent::Succeeded {
timing: TIMING,
response: &response,
};
let (sender, receiver) = litellm_host::observation::observation_channel(
std::num::NonZeroUsize::new(1).unwrap(),
);
drop(receiver);
sender.emit(event.snapshot());
assert!(matches!(
logging.on_event(py, event).unwrap(),
HookStep::Ready(())
));
assert_eq!(sender.dropped_events(), 1);
run(py, &locals, c"
assert logger.names() == ['sync_success_for_async_call'], logger.calls
logger._native_pending_logging.release(True)
logger._native_pending_logging.release(True)
assert logger.names() == ['sync_success_for_async_call', 'async_success_handler', 'enqueued'], logger.calls
assert logger.calls[0][1] is response
assert logger.calls[1][1] is response
");
});
}
#[rstest]
fn stream_bindings_deliver_collected_chunks_in_order_without_success_fan_out() {
Python::initialize();
Python::attach(|py| {
let locals = namespace(py, c"first = b'first'\nlast = b'last'\nresponse = None");
let mut logging = LegacyLogging {
operation: crate::LoggingOperation::Messages,
..logged(py, &locals, true)
};
logging
.on_stream_open(py, &pyo3::types::PyDict::new(py).into_any().unbind())
.unwrap();
logging
.on_stream_chunk(py, &local(&locals, "first").unbind())
.unwrap();
logging
.on_stream_chunk(py, &local(&locals, "last").unbind())
.unwrap();
assert!(matches!(
succeed(py, &locals, &mut logging),
HookStep::Ready(())
));
run(
py,
&locals,
c"
assert logger.names() == ['stream_opened', 'stream_success'], logger.calls
chunks = logger.calls[1][1]
assert len(chunks) == 2
assert chunks[0] is first
assert chunks[1] is last
",
);
});
}
#[rstest]
#[case::synchronous(false, &["failure_handler"])]
#[case::asynchronous(true, &[])]

View file

@ -3,16 +3,13 @@
//! lifetime. No other callback host has that obligation, which is why nothing outside
//! this crate holds them.
use litellm_host::{call::HostedCompletion, machine::Machine, protocol::Protocol};
use litellm_host_python::{Preflight, PythonBinding, PythonHostCalls, lookup, run_call};
use litellm_host_python::lookup;
use pyo3::{
gc::{PyTraverseError, PyVisit},
prelude::*,
types::{PyDict, PyTuple},
};
use crate::{LegacyLogging, LegacySurface};
pub struct PublicCall {
args: Py<PyTuple>,
kwargs: Py<PyDict>,
@ -34,6 +31,10 @@ impl PublicCall {
})
}
pub fn arguments(&self, py: Python<'_>) -> Py<PyDict> {
self.kwargs.clone_ref(py)
}
pub(crate) fn args(&self) -> &Py<PyTuple> {
&self.args
}
@ -64,35 +65,6 @@ impl PublicCall {
}
}
/// Runs one native call under the legacy `Logging` contract: the protocol host projects from
/// the keyword view the contract prepares and `preflight` rewrites, and the contract
/// observes the call.
pub fn run_legacy_call<H, M>(
py: Python<'_>,
surface: LegacySurface,
call: PublicCall,
start: impl FnOnce(<H::Protocol as Protocol>::Request) -> M + Send + Sync + 'static,
host: H,
preflight: Preflight,
asynchronous: bool,
) -> PyResult<Py<PyAny>>
where
H: PythonBinding + PythonHostCalls<H::Protocol> + 'static,
M: Machine<Protocol = H::Protocol> + 'static,
M::Complete: Into<HostedCompletion<<H::Protocol as Protocol>::Response>>,
{
let arguments = call.kwargs.clone_ref(py);
run_call(
py,
start,
host,
LegacyLogging::new(py, surface, call, asynchronous),
preflight,
arguments,
asynchronous,
)
}
#[cfg(test)]
mod tests {
use super::*;

View file

@ -2,7 +2,7 @@
//! the deferred and worker-submitted success paths, and the sync-callbacks-for-async-calls
//! duplication. All of it expires with the legacy callback contract.
use litellm_host::event::{RequestContext, WireRequest};
use litellm_host::interceptors::{RequestContext, WireRequest};
use litellm_host_python::to_py;
use pyo3::{exceptions::PyBaseException, prelude::*, types::PyDict};

View file

@ -2,25 +2,31 @@
//! sync and async callback registries it fans out to, the deployment hooks and the deferred
//! proxy release. All of it sits behind one
//! [`PythonCallHooks`](litellm_host_python::PythonCallHooks), so the driver, the routes and
//! core never learn which Python object is on the other end. The SDK's own request policy
//! (credential inheritance, the budget and retry limits) is the driver's preflight, not this
//! crate's.
//! core never learn which Python object is on the other end.
//!
//! Legacy callbacks receive the caller's own objects and may mutate them. [`PublicCall`]
//! is where those objects live, and [`run_legacy_call`] is how a route hands them over
//! without keeping a copy.
//! is where those objects live.
mod adapter;
mod call;
mod callbacks;
mod deferred;
mod logger;
mod mapping;
mod python;
pub(crate) use adapter::LegacyLogging;
pub use adapter::{LegacySurface, PassThroughStream};
pub use call::{PublicCall, run_legacy_call};
pub use adapter::LegacyLogging;
pub use call::PublicCall;
pub(crate) use callbacks::{LegacyCallbacks, is_internal_call};
pub(crate) use logger::{DeploymentHooks, PythonLogger, finalize, setup};
pub use mapping::{CallBoundary, CallbackMapping, Dispatch, callback_mappings};
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum LoggingOperation {
Completion,
Responses,
Messages,
Ocr,
}
#[cfg(test)]
mod test_support;

View file

@ -0,0 +1,288 @@
use litellm_host::{
hooks::CallHooks,
interceptors::{RawResponse, RequestContext, WireRequest},
lifecycle::{ExecutionEvent, FailureOrigin, Timing},
};
use litellm_host_python::{HookStep, PythonCallEvent, PythonRuntime};
use pyo3::{prelude::*, types::PyDict};
use crate::LegacyLogging;
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum CallBoundary {
PrepareArguments,
BeforeProviderRequest,
AfterProviderResponse,
TransformResponse,
Succeeded,
Failed,
StreamOpened,
StreamChunk,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Dispatch {
Call(CallBoundary),
Python(&'static str),
DeclarationOnly,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct CallbackMapping {
pub callback: &'static str,
pub dispatch: Dispatch,
}
struct Binding<H> {
boundary: CallBoundary,
invoke: H,
callbacks: &'static [&'static str],
}
impl<H> Binding<H> {
fn mappings(&self) -> impl Iterator<Item = CallbackMapping> {
self.callbacks.iter().map(|callback| CallbackMapping {
callback,
dispatch: Dispatch::Call(self.boundary),
})
}
}
type Step<T> = PyResult<HookStep<LegacyLogging, T>>;
type Prepare = fn(&mut LegacyLogging, Python<'_>, Py<PyDict>, f64) -> Step<Py<PyDict>>;
type Before =
fn(&mut LegacyLogging, Python<'_>, Box<WireRequest>, &RequestContext) -> Step<Box<WireRequest>>;
type After = fn(&mut LegacyLogging, Python<'_>, &RawResponse) -> Step<()>;
type Transform = fn(&mut LegacyLogging, Python<'_>, Py<PyAny>, Timing) -> Step<Py<PyAny>>;
type Success = fn(&mut LegacyLogging, Python<'_>, Timing, &Py<PyAny>) -> Step<()>;
type Failure = fn(&mut LegacyLogging, Python<'_>, Timing, FailureOrigin, &PyErr) -> Step<()>;
type Open = fn(&mut LegacyLogging, Python<'_>, &Py<PyAny>) -> PyResult<()>;
type Chunk = fn(&mut LegacyLogging, Python<'_>, &Py<PyAny>) -> PyResult<()>;
const PREPARE: Binding<Prepare> = Binding {
boundary: CallBoundary::PrepareArguments,
invoke: LegacyLogging::prepare_call,
callbacks: &["async_pre_call_deployment_hook"],
};
const BEFORE: Binding<Before> = Binding {
boundary: CallBoundary::BeforeProviderRequest,
invoke: LegacyLogging::pre_call,
callbacks: &["log_pre_api_call", "log_input_event"],
};
const AFTER: Binding<After> = Binding {
boundary: CallBoundary::AfterProviderResponse,
invoke: LegacyLogging::post_call,
callbacks: &["log_post_api_call"],
};
const TRANSFORM: Binding<Transform> = Binding {
boundary: CallBoundary::TransformResponse,
invoke: LegacyLogging::transform_public_response,
callbacks: &["async_post_call_success_deployment_hook"],
};
const SUCCESS: Binding<Success> = Binding {
boundary: CallBoundary::Succeeded,
invoke: LegacyLogging::succeeded,
callbacks: &[
"log_success_event",
"async_log_success_event",
"logging_hook",
"async_logging_hook",
"redact_standard_logging_payload_from_model_call_details",
"log_event",
"async_log_event",
],
};
const FAILURE: Binding<Failure> = Binding {
boundary: CallBoundary::Failed,
invoke: LegacyLogging::failed,
callbacks: &[
"async_post_call_failure_deployment_hook",
"log_failure_event",
"async_log_failure_event",
"log_model_group_rate_limit_error",
"log_event",
"async_log_event",
],
};
const OPEN: Binding<Open> = Binding {
boundary: CallBoundary::StreamOpened,
invoke: LegacyLogging::stream_opened,
callbacks: &[],
};
const CHUNK: Binding<Chunk> = Binding {
boundary: CallBoundary::StreamChunk,
invoke: LegacyLogging::stream_chunk,
callbacks: &[],
};
pub fn callback_mappings() -> impl Iterator<Item = CallbackMapping> {
PREPARE
.mappings()
.chain(BEFORE.mappings())
.chain(AFTER.mappings())
.chain(TRANSFORM.mappings())
.chain(SUCCESS.mappings())
.chain(FAILURE.mappings())
.chain(OPEN.mappings())
.chain(CHUNK.mappings())
.chain(PYTHON_CALLBACKS.iter().copied())
}
impl CallHooks<PythonRuntime> for LegacyLogging {
fn prepare_arguments(
&mut self,
py: Python<'_>,
arguments: Py<PyDict>,
started_at: f64,
) -> Step<Py<PyDict>> {
(PREPARE.invoke)(self, py, arguments, started_at)
}
fn arguments_prepared(&mut self, py: Python<'_>, arguments: &Py<PyDict>) -> PyResult<()> {
self.adopt_arguments(py, arguments);
Ok(())
}
fn before_provider_request(
&mut self,
py: Python<'_>,
wire: Box<WireRequest>,
context: &RequestContext,
) -> Step<Box<WireRequest>> {
(BEFORE.invoke)(self, py, wire, context)
}
fn transform_response(
&mut self,
py: Python<'_>,
response: Py<PyAny>,
timing: Timing,
) -> Step<Py<PyAny>> {
(TRANSFORM.invoke)(self, py, response, timing)
}
fn on_event(&mut self, py: Python<'_>, event: PythonCallEvent<'_>) -> Step<()> {
match event {
PythonCallEvent::Started { .. } | PythonCallEvent::Cancelled { .. } => {
Ok(HookStep::Ready(()))
}
PythonCallEvent::Execution(ExecutionEvent::ResultReady { facts }) => {
self.result_ready(py, &facts)
}
PythonCallEvent::Execution(ExecutionEvent::ProviderResponseReceived { raw }) => {
(AFTER.invoke)(self, py, raw)
}
PythonCallEvent::Succeeded { timing, response } => {
(SUCCESS.invoke)(self, py, timing, response)
}
PythonCallEvent::Failed {
timing,
origin,
error,
} => (FAILURE.invoke)(self, py, timing, origin, error),
}
}
fn on_stream_open(&mut self, py: Python<'_>, head: &Py<PyAny>) -> PyResult<()> {
(OPEN.invoke)(self, py, head)
}
fn on_stream_chunk(&mut self, py: Python<'_>, chunk: &Py<PyAny>) -> PyResult<()> {
(CHUNK.invoke)(self, py, chunk)
}
}
macro_rules! python_callbacks {
($($dispatch:expr => [$($callback:literal),* $(,)?]),* $(,)?) => {
const PYTHON_CALLBACKS: &[CallbackMapping] = &[
$($(CallbackMapping { callback: $callback, dispatch: $dispatch },)*)*
];
};
}
python_callbacks! {
Dispatch::Python("litellm.router") => [
"async_pre_routing_hook",
"async_filter_deployments",
"pre_call_check",
"async_pre_call_check",
],
Dispatch::Python("litellm.router_utils.fallback_event_handlers") => [
"log_success_fallback_event",
"log_failure_fallback_event",
],
Dispatch::Python("litellm.proxy.utils") => [
"async_pre_call_hook",
"async_post_call_response_headers_hook",
"async_post_call_failure_hook",
"async_post_call_success_hook",
"async_moderation_hook",
"async_post_call_streaming_hook",
"async_post_call_streaming_iterator_hook",
"async_filter_listed_models",
],
Dispatch::Python("litellm.litellm_core_utils.litellm_logging") => [
"async_get_chat_completion_prompt",
"get_chat_completion_prompt",
"log_stream_event",
"async_log_stream_event",
"async_post_mcp_tool_call_hook",
],
Dispatch::Python("litellm.llms.anthropic.pass_through.messages.handler") => [
"async_pre_request_hook",
],
Dispatch::Python("litellm.litellm_core_utils.streaming_handler") => [
"async_post_call_streaming_deployment_hook",
],
Dispatch::Python("litellm.responses.streaming_iterator") => [
"async_post_call_streaming_deployment_hook",
],
Dispatch::Python("litellm.main") => [
"translate_completion_input_params",
"translate_completion_output_params",
"translate_completion_output_params_streaming",
],
Dispatch::Python("litellm.integrations.argilla") => ["async_dataset_hook"],
Dispatch::Python("litellm.proxy.management_helpers.audit_logs") => ["async_log_audit_log_event"],
Dispatch::Python("litellm.llms.custom_httpx.llm_http_handler") => [
"async_should_run_agentic_loop",
"async_run_agentic_loop",
"async_build_agentic_loop_plan",
"async_post_agentic_loop_response_hook",
"async_agentic_loop_cleanup_hook",
"async_should_run_chat_completion_agentic_loop",
"async_run_chat_completion_agentic_loop",
"async_build_chat_completion_agentic_loop_plan",
],
Dispatch::Python("litellm.litellm_core_utils.chat_completion_agentic_loop") => [
"async_should_run_agentic_loop",
"async_run_agentic_loop",
"async_build_agentic_loop_plan",
"async_post_agentic_loop_response_hook",
"async_agentic_loop_cleanup_hook",
],
Dispatch::Python("litellm.llms.openai.openai") => [
"async_should_run_chat_completion_agentic_loop",
"async_run_chat_completion_agentic_loop",
],
Dispatch::Python("litellm.proxy.spend_tracking.cold_storage_handler") => [
"get_proxy_server_request_from_cold_storage_with_object_key",
],
Dispatch::Python("litellm.integrations.custom_logger") => [
"truncate_standard_logging_payload_content",
"redacts_messages_itself",
"handle_callback_failure",
"get_callback_env_vars",
],
Dispatch::DeclarationOnly => [
"async_log_pre_api_call",
"async_log_input_event",
],
}

View file

@ -3,7 +3,7 @@ use std::ffi::CStr;
use pyo3::prelude::*;
use pyo3::types::{PyDict, PyTuple};
use crate::{LegacyLogging, LegacySurface, PublicCall};
use crate::{LegacyLogging, PublicCall};
/// The parameters of every `callbacks_legacy_python` function, as the real module declares them.
/// `tests/unit/rust_bridge/test_callbacks_legacy_python.py` pins this file to the Python
@ -45,11 +45,14 @@ def contracted(name, fake):
if not hasattr(legacy, 'is_internal'):
legacy.is_internal = contextvars.ContextVar('is_internal_call', default=False)
def setup(call_type, args, kwargs, start, asynchronous):
logger = kwargs['logger_factory'](kwargs) if 'logger_factory' in kwargs else kwargs['logger']
logger.setup_call_type = call_type
return types.SimpleNamespace(logger=logger, kwargs=kwargs)
FAKES = {
'setup': lambda call_type, args, kwargs, start, asynchronous: types.SimpleNamespace(
logger=kwargs['logger_factory'](kwargs) if 'logger_factory' in kwargs else kwargs['logger'],
kwargs=kwargs,
),
'setup': setup,
'finalize': lambda response, logger, kwargs, start, end: logger.record('finalize', response),
'update_logging': lambda logger, kwargs, model, optional_params, litellm_params, provider: logger.update_from_kwargs(
kwargs=kwargs,
@ -82,10 +85,10 @@ FAKES = {
),
'after_deployment_failure': lambda kwargs, error, call_type: kwargs['logger'].hook('failure', error, call_type),
'stream_opened': lambda logger: logger.record('stream_opened', None),
'stream_success': lambda logger, request_body, chunks, start, end, first_chunk: logger.record(
'stream_success': lambda logger, url_route, endpoint_type, request_body, chunks, start, end, first_chunk: logger.record(
'stream_success', list(chunks)
),
'stream_failure': lambda logger, request_body, chunks, error: logger.record('stream_failure', error),
'stream_failure': lambda logger, endpoint_type, request_body, chunks, error: logger.record('stream_failure', error),
}
assert FAKES.keys() == CONTRACT.keys(), sorted(FAKES.keys() ^ CONTRACT.keys())
for name, fake in FAKES.items():
@ -186,14 +189,5 @@ pub(crate) fn legacy_call(
.map(|kwargs| kwargs.cast_into::<PyDict>().unwrap())
.unwrap_or_else(|| PyDict::new(py));
let call = PublicCall::capture(&request, &PyTuple::empty(py), &kwargs).unwrap();
LegacyLogging::new(
py,
LegacySurface {
call_type: "test",
input_description: "test input",
stream: None,
},
call,
asynchronous,
)
LegacyLogging::new(py, crate::LoggingOperation::Ocr, call, asynchronous)
}

View file

@ -8,11 +8,10 @@ repository.workspace = true
[dependencies]
fancy-regex.workspace = true
litellm-tracing.workspace = true
litellm-types.workspace = true
litellm-llms-types.workspace = true
serde.workspace = true
serde_json.workspace = true
serde_path_to_error = "0.1"
serde_with.workspace = true
strum.workspace = true
thiserror.workspace = true
url.workspace = true

View file

@ -2,7 +2,7 @@
use std::time::{SystemTime, UNIX_EPOCH};
use litellm_types::utils::{ChatCompletionsUsage, PromptTokensDetails};
use litellm_llms_types::formats::chat_completions::{ChatCompletionsUsage, PromptTokensDetails};
/// OpenAI finish reasons, mirroring Python's `_FINISH_REASON_MAP` for the
/// reasons the providers on this route can emit. Python warns and falls back to

View file

@ -1,4 +1,4 @@
use litellm_types::utils::{ProviderSpecificHeader, ProviderSpecificHeaders};
use litellm_llms_types::headers::{ProviderSpecificHeader, ProviderSpecificHeaders};
use serde_json::{Map, Value};
pub fn get_provider_specific_headers(

View file

@ -10,7 +10,7 @@
//! `_bedrock_converse_messages_pt` for the text-only surface this route
//! accepts; anything richer is declined upstream by the capability gate.
use litellm_types::llms::openai::{ChatMessage, ChatMessageContent};
use litellm_llms_types::formats::chat_completions::{ChatMessage, ChatMessageContent};
use strum::IntoStaticStr;
pub const EMPTY_TEXT_PLACEHOLDER: &str =

View file

@ -1,12 +1,3 @@
use serde::{
Deserializer,
de::{Error, Visitor},
};
use serde_with::DeserializeAs;
pub struct LaxI64;
pub struct FiniteF64;
pub fn parse_str_bool(value: &str) -> Option<bool> {
let token = value.trim_matches(|character: char| {
character.is_whitespace() || matches!(character, '\u{1c}'..='\u{1f}')
@ -22,129 +13,12 @@ pub fn parse_redis_bool(value: &str) -> bool {
value == "1" || value.eq_ignore_ascii_case("true") || value.eq_ignore_ascii_case("yes")
}
impl<'de> DeserializeAs<'de, i64> for LaxI64 {
fn deserialize_as<D: Deserializer<'de>>(deserializer: D) -> Result<i64, D::Error> {
deserializer.deserialize_any(Self)
}
}
impl<'de> Visitor<'de> for LaxI64 {
type Value = i64;
fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
formatter.write_str("an integer in the i64 range")
}
fn visit_i64<E: Error>(self, value: i64) -> Result<i64, E> {
Ok(value)
}
fn visit_u64<E: Error>(self, value: u64) -> Result<i64, E> {
i64::try_from(value).map_err(E::custom)
}
fn visit_f64<E: Error>(self, value: f64) -> Result<i64, E> {
integral_float(value).ok_or_else(|| E::custom("expected an integer in the i64 range"))
}
fn visit_str<E: Error>(self, value: &str) -> Result<i64, E> {
integer_string(value.trim())
.ok_or_else(|| E::custom("expected an integer in the i64 range"))
}
fn visit_bool<E: Error>(self, value: bool) -> Result<i64, E> {
Ok(i64::from(value))
}
}
impl<'de> DeserializeAs<'de, f64> for FiniteF64 {
fn deserialize_as<D: Deserializer<'de>>(deserializer: D) -> Result<f64, D::Error> {
deserializer.deserialize_any(Self)
}
}
impl<'de> Visitor<'de> for FiniteF64 {
type Value = f64;
fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
formatter.write_str("a finite number")
}
fn visit_i64<E: Error>(self, value: i64) -> Result<f64, E> {
Ok(value as f64)
}
fn visit_u64<E: Error>(self, value: u64) -> Result<f64, E> {
Ok(value as f64)
}
fn visit_f64<E: Error>(self, value: f64) -> Result<f64, E> {
value
.is_finite()
.then_some(value)
.ok_or_else(|| E::custom("expected a finite number"))
}
fn visit_str<E: Error>(self, value: &str) -> Result<f64, E> {
self.visit_f64(value.trim().parse::<f64>().map_err(E::custom)?)
}
fn visit_bool<E: Error>(self, value: bool) -> Result<f64, E> {
Ok(f64::from(value))
}
}
fn integer_string(value: &str) -> Option<i64> {
let integer = match value.split_once('.') {
Some((integer, fraction)) => {
if fraction.is_empty() || !fraction.bytes().all(|byte| byte == b'0') {
return None;
}
integer
}
None => value,
};
if integer.starts_with('_') || integer.ends_with('_') || integer.contains("__") {
return None;
}
let digits = integer.strip_prefix(['+', '-']).unwrap_or(integer);
if digits.is_empty()
|| digits.starts_with('_')
|| !digits
.bytes()
.all(|byte| byte.is_ascii_digit() || byte == b'_')
{
return None;
}
integer.replace('_', "").parse().ok()
}
fn integral_float(value: f64) -> Option<i64> {
(value.is_finite()
&& value.fract() == 0.0
&& value >= i64::MIN as f64
&& value < -(i64::MIN as f64))
.then_some(value as i64)
}
#[cfg(test)]
mod tests {
use rstest::rstest;
use serde::{Deserialize, Serialize};
use serde_json::json;
use serde_with::serde_as;
use super::*;
#[serde_as]
#[derive(Debug, Deserialize, Serialize, PartialEq)]
struct Numbers {
#[serde_as(deserialize_as = "Option<Vec<LaxI64>>")]
integers: Option<Vec<i64>>,
#[serde_as(deserialize_as = "Option<FiniteF64>")]
float: Option<f64>,
}
#[rstest]
#[case::trimmed_true(" True ", Some(true))]
#[case::control_whitespace_true("\u{1c}TRUE\u{1f}", Some(true))]
@ -160,73 +34,4 @@ mod tests {
) {
assert_eq!(parse_str_bool(input), expected, "{input:?}");
}
#[test]
fn adapters_compose_and_serialize_as_numbers() {
let numbers: Numbers = serde_json::from_value(json!({
"integers": ["9007199254740993.0", "1_000", " +2.000 ", 3.0, true],
"float": " 1.5 "
}))
.unwrap();
assert_eq!(
serde_json::to_value(numbers).unwrap(),
json!({
"integers": [9_007_199_254_740_993_i64, 1000, 2, 3, 1], "float": 1.5
})
);
for input in [json!({}), json!({"integers": null, "float": null})] {
assert_eq!(
serde_json::from_value::<Numbers>(input).unwrap(),
Numbers {
integers: None,
float: None,
}
);
}
}
#[test]
fn integer_bounds_and_invalid_values_are_checked() {
for input in [
json!(i64::MIN),
json!(i64::MAX),
json!(i64::MAX.to_string()),
] {
assert!(serde_json::from_value::<Numbers>(json!({"integers": [input]})).is_ok());
}
for input in [
json!(u64::MAX),
json!(9_223_372_036_854_775_808_u64),
json!(9_223_372_036_854_775_808.0),
json!("-9223372036854775809"),
json!("1.0000000000000001"),
json!("1e3"),
json!("2."),
json!(".0"),
json!("_2"),
json!("2__0"),
json!(2.5),
json!(null),
json!({}),
] {
assert!(serde_json::from_value::<Numbers>(json!({"integers": [input]})).is_err());
}
}
#[test]
fn floats_reject_nonfinite_and_invalid_values() {
for input in [
json!("NaN"),
json!("inf"),
json!("-inf"),
json!("1e999"),
json!([]),
] {
assert!(serde_json::from_value::<Numbers>(json!({"float": input})).is_err());
}
for (input, expected) in [(json!(2), 2.0), (json!(2.5), 2.5), (json!(true), 1.0)] {
let numbers: Numbers = serde_json::from_value(json!({"float": input})).unwrap();
assert_eq!(numbers.float, Some(expected));
}
}
}

View file

@ -2,7 +2,7 @@ litellm-core owns route orchestration. Messages and HTTP Responses return `litel
Hosts assemble route objects from shared `CoreResources`, HTTP settings, and secret sources. Each route owns its provider client and authentication dependencies. Gateway routes live for the gateway lifetime; Python assembles routes per call from its settings snapshot
Chat Completions, Messages, and OCR execute through their route objects. Calls pass `RouteHooks` directly; use `&()` when no hooks are needed. Construction does no work; preparation and lifecycle observation begin when the future is polled. Handlers accept `RouteHooks`, never a concrete `ChannelHooks`. Native observers receive start and terminal events through the shared call runner; a stream retains its lifecycle until exhaustion, error, or drop. A host channel has no native observer because its driver owns terminal dispatch
Chat Completions, Messages, Responses, and OCR execute through their route objects. Calls pass `Interceptors` and an optional `ObservationSender` separately; use `&()` for no hooks and `None` for no observer. Construction does no work; preparation and lifecycle observation begin when the future is polled. Handlers accept `Interceptors`, never a concrete `ChannelInterceptors`. Native observers receive start and terminal events through the shared call runner; a stream retains its lifecycle until exhaustion, error, or drop. Hosted routes leave terminal observation to their driver
`route.rs` declares the concrete `Protocol` and implements a route method that accepts a typed request and constructs a `litellm_host::call::HostedMachine` with `hosted_call`. The shared call plumbing owns stream opening, delivery, backpressure, and detachment. Request decoding belongs to the boundary before the machine starts. Route closures only supply execution dependencies and route-specific host capabilities such as an OCR token provider. Use `run_hosted` for a native host so detachment is reported as cancellation. Python uses its own shared driver and preserves caller-task callback execution
@ -10,17 +10,17 @@ Responses WebSocket sessions remain separate from the HTTP call driver because a
## Crate layering
For Messages, Responses, Chat Completions, OCR, and other API formats, `core/src/<format>/` owns orchestration. Shared API data contracts belong in `litellm-types`, adapter contracts and shared transformation machinery in `llms/src/base_llm/<format>/`, and provider policy in `llms/src/<provider>/<format>/`. A repeated format directory name does not imply interchangeable responsibilities. Select concrete adapters here, then invoke their contracts instead of applying one provider's policy to every call. Route types describe call envelopes and execution state, not duplicate public payload schemas
For Messages, Responses, Chat Completions, OCR, and other API formats, `core/src/<format>/` owns orchestration. Shared API data contracts belong in `litellm-llms-types`, adapter contracts and shared transformation machinery in `llms/src/base_llm/<format>/`, and provider policy in `llms/src/<provider>/<format>/`. A repeated format directory name does not imply interchangeable responsibilities. Select concrete adapters here, then invoke their contracts instead of applying one provider's policy to every call. Route types describe call envelopes and execution state, not duplicate public payload schemas
Each crate mirrors one top-level Python package, so a Rust path reads as its Python path with the crate name in place of the package directory. Dependencies only point down:
Crates separate API data, transformations, transport, and orchestration. Python package names identify counterparts, not ownership. Dependencies only point down:
- `litellm-types` mirrors `litellm/types/`: pure serde data, no I/O
- `litellm-llms-types` owns shared inference API contracts, grouped by format: pure serde data and shape validation, no I/O
- `litellm-core-utils` mirrors `litellm/litellm_core_utils/`: pure helpers (provider resolution, prompt factory, call arguments, settings lookup and layer merge), no network I/O
- `litellm-http` is Rust-only and route-neutral: settings resolution, the pooled `reqwest` clients, TLS, proxies, the SSRF-safe media fetcher, request and header helpers, and transport errors. Python's `litellm/llms/custom_httpx/` is split by responsibility instead of mirrored: its transport half lives here, its OCR handler in `litellm-llms`
- `litellm-llms` mirrors `litellm/llms/`: `base_llm/<api>/transformation.rs`, `<provider>/<api>/transformation.rs`, and `base_llm/ocr/handler.rs` (the OCR request handler)
- `litellm-core` mirrors the route packages (`litellm/ocr/`, `litellm/messages/`, ...): entrypoints, route request types, provider dispatch, the route machine, and hooks
A route module owns the call entrypoint, route request types (`*Request<'a>`), credential fallback, provider dispatch, and the handler glue that runs a provider config. Provider code never imports from core; when it needs the caller's hooks mid-call it goes through `litellm_llms::base_llm::ocr::handler::CallHooks`, the provider-level hooks OCR implements over its host until it folds into `litellm_host::hooks::RouteHooks`. Import every item from its canonical path. Never re-export another crate's items or give an item a second public path; the only re-export allowed is a private submodule surfacing its item at its module root (`mod error; pub use error::Error;`). Handlers belong in core or llms, never in a host crate
A route module owns the call entrypoint, route request types (`*Request<'a>`), credential fallback, provider dispatch, and the handler glue that runs a provider config. Provider code never imports from core; when it needs the caller's hooks mid-call it goes through `litellm_llms::base_llm::ocr::handler::CallHooks`, the provider-level hooks OCR implements over its host until it folds into `litellm_host::interceptors::Interceptors`. Import every item from its canonical path. Never re-export another crate's items or give an item a second public path; the only re-export allowed is a private submodule surfacing its item at its module root (`mod error; pub use error::Error;`). Handlers belong in core or llms, never in a host crate
## Error placement
@ -33,3 +33,17 @@ Scope follows the concept, not the first caller. An error type under `litellm-ll
`litellm_llms::Error` (`crates/llms/src/error.rs`) is the one transformation error for every provider and API. `base_llm/ocr/error.rs` is the recorded exception until OCR folds into it
Not here: serving HTTP (axum routes, extractors), config file reading, rollout state, databases, or callback execution of any kind. Core runs each route as a machine that yields host operations and call events; which integrations consume those events is the host's business.
## Response caching and accounting boundary
Attach a `litellm_cache_response::ScopedCache` with `route.with_cache(cache)`. Cached and uncached routes use the same `execute` and `machine` methods. `CallOptions` carries a scope-free `CachePolicy` and observation; per-call policy never replaces the attached scope or service
Messages groups per-call dependencies in `CallContext` and explicitly sequences cache lookup, provider execution, result acceptance, and cache storage. Provider transport does not own cache orchestration. Stream capture remains in the shared cache implementation
Core owns request identity, typed response reconstruction and stream capture/replay. `cache-response` owns cache policy, namespacing, scope encoding, versioned envelopes and freshness. The SDK explicitly chooses shared scope. The gateway derives isolated scope from authenticated identity before attaching its service
Core delivers `ExecutionFacts` through the awaited `ResultReady` host operation for both provider and cached results, before public response processing or stream opening. Facts carry resolved model/provider and result source, including the hit key. Usage remains in the typed response or delivered stream, where completion and cancellation determine what was actually reported. Passive observation is not an accounting delivery mechanism
Core does not calculate prices, charge budgets, or update rate-limit counters. The legacy Python callback adapter translates execution facts into the existing Python logging contract; Python remains the accounting owner on that path. Native gateway accounting belongs to gateway dependencies, independently of `host-python`. Response-cache services expose no coordination counters or reservation APIs. A shared Redis deployment does not make response storage and accounting coordination the same dependency
Cache lookup follows provider preparation, credential resolution and the request interceptor. Keys describe the effective provider URL, authenticated headers and rewritten body. Signed requests bypass caching until the signing identity has a stable cache representation

View file

@ -6,8 +6,12 @@ license.workspace = true
repository.workspace = true
[dependencies]
litellm-cache.workspace = true
litellm-cache-response.workspace = true
litellm-framing.workspace = true
tokio-util = { version = "0.7", features = ["codec"] }
litellm-secrets.workspace = true
litellm-types.workspace = true
litellm-llms-types.workspace = true
litellm-core-utils.workspace = true
litellm-host.workspace = true
bytes.workspace = true
@ -36,8 +40,10 @@ url.workspace = true
veil.workspace = true
[dev-dependencies]
litellm-cache-memory.workspace = true
litellm-http = { workspace = true, features = ["test-support"] }
litellm-auth-gcp.workspace = true
litellm-host-native.workspace = true
litellm-llms = { workspace = true, features = ["test-support"] }
rstest.workspace = true
rstest_reuse.workspace = true

View file

@ -0,0 +1,376 @@
use std::{
future::Future,
marker::PhantomData,
sync::Arc,
time::{Duration, SystemTime, UNIX_EPOCH},
};
use bytes::{Bytes, BytesMut};
use futures_util::{StreamExt, TryStreamExt, stream};
use litellm_cache_response::{
CacheOptions, CachePolicy, ResponseCacheRequest, ResponseCacheService, ResponseEnvelope,
ScopedCache, cache_key,
};
use litellm_host::{
call::{CallOutput, OutputOf},
interceptors::{ExecutionFacts, Interceptors, ProviderIdentity, ResultSource, WireRequest},
lifecycle::{CallEvent, ExecutionEvent},
observation::ObservationSender,
protocol::Protocol,
};
use serde::{Deserialize, Serialize, de::DeserializeOwned};
use serde_json::Value;
use tokio_util::codec::Decoder;
use crate::RouteError;
pub trait Cachable: Protocol<Error = RouteError> {
const SURFACE: &'static str;
fn reusable(_response: &Self::Response) -> bool {
true
}
}
pub struct CacheRequest {
pub identity: ProviderIdentity,
pub input: Value,
}
impl CacheRequest {
pub fn from_wire(identity: ProviderIdentity, wire: Option<&WireRequest>) -> Self {
Self {
input: wire.map_or(Value::Null, |wire| {
serde_json::json!({
"provider": identity.provider,
"model": identity.model,
"url": wire.url,
"headers": wire.headers,
"body": wire.body,
})
}),
identity,
}
}
}
pub trait StreamCachable: Cachable {
const TERMINAL_EVENT: &'static str;
fn replay(data: Bytes) -> Option<OutputOf<Self>>;
fn bytes(chunk: &Self::Chunk) -> &[u8];
}
#[derive(Serialize, Deserialize)]
#[serde(tag = "kind", content = "value")]
pub enum CachedOutput<R> {
Response(R),
Stream(String),
}
struct CacheSession {
service: Arc<dyn ResponseCacheService>,
request: ResponseCacheRequest,
}
impl CacheSession {
fn prepare<P: Cachable>(
service: Option<Arc<dyn ResponseCacheService>>,
options: Option<CacheOptions>,
request: &CacheRequest,
) -> Option<Self> {
let options = options.filter(|options| options.policy.enabled())?;
let service = service?;
let input = request.input.clone();
let request = options.request(&service.config().namespace, P::SURFACE, input);
Some(Self { service, request })
}
async fn lookup<P: Cachable>(&self) -> Option<CachedOutput<P::Response>>
where
P::Response: DeserializeOwned,
{
if !self.request.controls.reads() {
return None;
}
match self.service.lookup(&self.request, now()).await {
Ok(Some(value)) => {
serde_json::from_value::<ResponseEnvelope<CachedOutput<P::Response>>>(value)
.ok()
.and_then(|entry| entry.decode(P::SURFACE))
}
Ok(None) => None,
Err(_) => {
tracing::warn!("response cache lookup failed");
None
}
}
}
async fn store(&self, entry: Value) {
if !self.request.controls.writes() {
return;
}
if self
.service
.store(&self.request, entry, now())
.await
.is_err()
{
tracing::warn!("response cache write failed");
}
}
async fn store_response<P: Cachable>(&self, response: &P::Response)
where
P::Response: Serialize,
{
if !self.request.controls.writes() || !P::reusable(response) {
return;
}
if let Ok(value) = serde_json::to_value(response)
&& let Ok(entry) = serde_json::to_value(ResponseEnvelope::new(
P::SURFACE,
CachedOutput::Response(value),
))
{
self.store(entry).await;
}
}
}
pub async fn execute_unary<P, F, Fut>(
request: CacheRequest,
cache: Option<Arc<dyn ResponseCacheService>>,
options: Option<CacheOptions>,
interceptors: &impl Interceptors<RouteError>,
observers: Option<&ObservationSender>,
provider: F,
) -> Result<P::Response, RouteError>
where
P: Cachable,
P::Response: Serialize + DeserializeOwned,
F: FnOnce() -> Fut,
Fut: Future<Output = Result<P::Response, RouteError>>,
{
let identity = request.identity.clone();
crate::diagnostic::provider(&identity.model, &identity.provider);
let session = CacheSession::prepare::<P>(cache, options, &request);
let hit = match &session {
Some(session) => session.lookup::<P>().await.and_then(|entry| match entry {
CachedOutput::Response(response) => Some((response, cache_key(&session.request.key))),
CachedOutput::Stream(_) => None,
}),
None => None,
};
let (response, source) = match hit {
Some((response, key)) => (response, ResultSource::Cache { key }),
None => (provider().await?, ResultSource::Provider),
};
let from_provider = source == ResultSource::Provider;
publish(
ExecutionFacts {
provider: identity,
source,
},
interceptors,
observers,
)
.await?;
if from_provider && let Some(session) = session {
session.store_response::<P>(&response).await;
}
Ok(response)
}
pub async fn execute_streaming<P, F, Fut>(
request: CacheRequest,
cache: Option<Arc<dyn ResponseCacheService>>,
options: Option<CacheOptions>,
interceptors: &impl Interceptors<RouteError>,
observers: Option<&ObservationSender>,
provider: F,
) -> Result<OutputOf<P>, RouteError>
where
P: StreamCachable,
P::Response: Serialize + DeserializeOwned,
F: FnOnce() -> Fut,
Fut: Future<Output = Result<OutputOf<P>, RouteError>>,
{
let identity = request.identity.clone();
crate::diagnostic::provider(&identity.model, &identity.provider);
let session = CacheSession::prepare::<P>(cache, options, &request);
let cache = CallCache::<P> {
session,
protocol: PhantomData,
};
let hit = cache.lookup().await;
let (output, source) = match hit {
Some(hit) => hit,
None => (provider().await?, ResultSource::Provider),
};
publish(
ExecutionFacts {
provider: identity,
source: source.clone(),
},
interceptors,
observers,
)
.await?;
Ok(cache.finish(output, &source).await)
}
pub(crate) struct CallCache<P> {
session: Option<CacheSession>,
protocol: PhantomData<P>,
}
impl<P: StreamCachable> CallCache<P> {
pub(crate) fn from_wire(
cache: Option<&ScopedCache>,
policy: CachePolicy,
identity: &ProviderIdentity,
wire: &WireRequest,
) -> Self {
let session = cache.and_then(|cache| {
if !policy.enabled() {
return None;
}
let options = cache.options(Some(policy));
let request = CacheRequest::from_wire(identity.clone(), Some(wire));
Some(CacheSession {
request: options.request(
&cache.service.config().namespace,
P::SURFACE,
request.input,
),
service: cache.service.clone(),
})
});
Self {
session,
protocol: PhantomData,
}
}
pub(crate) async fn lookup(&self) -> Option<(OutputOf<P>, ResultSource)>
where
P::Response: DeserializeOwned,
{
let session = self.session.as_ref()?;
let output = match session.lookup::<P>().await? {
CachedOutput::Response(response) => CallOutput::Complete(response),
CachedOutput::Stream(data) => P::replay(Bytes::from(data))?,
};
Some((
output,
ResultSource::Cache {
key: cache_key(&session.request.key),
},
))
}
pub(crate) async fn finish(self, output: OutputOf<P>, source: &ResultSource) -> OutputOf<P>
where
P::Response: Serialize,
{
let Some(session) = self.session.filter(|session| {
*source == ResultSource::Provider && session.request.controls.writes()
}) else {
return output;
};
match output {
CallOutput::Complete(response) => {
session.store_response::<P>(&response).await;
CallOutput::Complete(response)
}
CallOutput::Stream { head, chunks } => CallOutput::Stream {
head,
chunks: capture_stream::<P>(chunks, session),
},
}
}
}
fn capture_stream<P: StreamCachable>(
chunks: futures_util::stream::BoxStream<'static, Result<P::Chunk, RouteError>>,
session: CacheSession,
) -> futures_util::stream::BoxStream<'static, Result<P::Chunk, RouteError>> {
stream::try_unfold(
(chunks, Some(Vec::<u8>::new()), session),
|(mut chunks, captured, session)| async move {
match chunks.try_next().await? {
Some(chunk) => {
let captured = captured.and_then(|mut data| {
let bytes = P::bytes(&chunk);
if data.len().saturating_add(bytes.len())
> session.service.config().max_entry_bytes
{
return None;
}
data.extend_from_slice(bytes);
Some(data)
});
Ok(Some((chunk, (chunks, captured, session))))
}
None => {
if let Some(data) = captured
&& let Ok(text) = String::from_utf8(data)
&& successful_stream(&text, P::TERMINAL_EVENT)
&& let Ok(entry) = serde_json::to_value(ResponseEnvelope::new(
P::SURFACE,
CachedOutput::<Value>::Stream(text),
))
{
session.store(entry).await;
}
Ok::<_, RouteError>(None)
}
}
},
)
.boxed()
}
fn now() -> Duration {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap_or_default()
}
fn successful_stream(text: &str, terminal: &str) -> bool {
let mut pending = BytesMut::from(text.as_bytes());
let mut codec = litellm_framing::sse::SseCodec::default();
let mut complete = false;
loop {
let event = match codec.decode(&mut pending) {
Ok(Some(event)) => event,
Ok(None) => return complete && pending.is_empty(),
Err(_) => return false,
};
let Ok(value) = serde_json::from_str::<Value>(&event.data) else {
return false;
};
let Some(kind) = value.get("type").and_then(Value::as_str) else {
return false;
};
if matches!(kind, "error" | "response.failed" | "response.incomplete") {
return false;
}
complete |= kind == terminal;
}
}
async fn publish(
facts: ExecutionFacts,
interceptors: &impl Interceptors<RouteError>,
observers: Option<&ObservationSender>,
) -> Result<(), RouteError> {
if let Some(observers) = observers {
observers.emit(CallEvent::Execution(ExecutionEvent::ResultReady {
facts: facts.clone(),
}));
}
interceptors.result_ready(facts).await
}

View file

@ -1,16 +1,14 @@
use litellm_host::{lifecycle::ExecutionEvent, observation::ObservationSender};
use std::time::Duration;
use litellm_auth::AuthServices;
use litellm_host::{
event::{MachineEvent, RawResponse, RequestContext, WireRequest},
hooks::RouteHooks,
};
use litellm_host::interceptors::{Interceptors, RawResponse, RequestContext, WireRequest};
use litellm_http::{Client, outbound::OutboundRequest, request::truncate_error_body};
use litellm_llms::base_llm::{
auth::{Authenticated, resolve_auth},
chat::transformation::ProviderChatResponseData,
};
use litellm_types::utils::ChatCompletionsResponse;
use litellm_llms_types::formats::chat_completions::ChatCompletionsResponse;
use serde_json::Value;
use super::Error;
@ -23,7 +21,10 @@ pub(super) async fn execute(
http: &Client,
auth: &AuthServices,
request: ProviderChatCompletionsRequest,
hooks: &impl RouteHooks<Error>,
cache: Option<litellm_cache_response::ScopedCache>,
cache_options: Option<litellm_cache_response::CachePolicy>,
interceptors: &impl Interceptors<Error>,
observers: Option<&ObservationSender>,
) -> Result<ChatCompletionsResponse, Error> {
let ProviderChatCompletionsRequest {
model,
@ -45,7 +46,11 @@ pub(super) async fn execute(
api_key,
};
let authenticated = resolve_auth(auth, environment, &|key| secrets.get(key)).await?;
let wire = hooks
let identity = litellm_host::interceptors::ProviderIdentity {
model: context.model.clone(),
provider: context.custom_llm_provider.clone(),
};
let wire = interceptors
.before_provider_request(
WireRequest {
url,
@ -55,55 +60,72 @@ pub(super) async fn execute(
context,
)
.await?;
let outbound = outbound_request(
Authenticated {
headers: wire.headers,
signer: authenticated.signer,
let cache = cache.filter(|_| authenticated.signer.is_none());
let cache_request =
crate::caching::CacheRequest::from_wire(identity, cache.as_ref().map(|_| &wire));
crate::caching::execute_unary::<super::route::ChatCompletions, _, _>(
cache_request,
cache.as_ref().map(|cache| cache.service.clone()),
cache.as_ref().map(|cache| cache.options(cache_options)),
interceptors,
observers,
|| async move {
let outbound = outbound_request(
Authenticated {
headers: wire.headers,
signer: authenticated.signer,
},
wire.url,
&wire.body,
timeout,
)?;
let response = crate::outbound::send(outbound, http).await.map_err(|err| {
// Failing to establish the connection means the request never went out,
// so the host can still serve it. Everything else here, a timeout
// above all, may have reached the provider and been answered.
if err.is_connect() || err.is_builder() {
Error::Transport(litellm_http::transport::Error::Connect(err.to_string()))
} else {
Error::Transport(litellm_http::transport::Error::Network(err.to_string()))
}
})?;
let status = response.status();
let text = response.text().await.map_err(|err| {
Error::Transport(litellm_http::transport::Error::Network(err.to_string()))
})?;
if !status.is_success() {
return Err(Error::Transport(litellm_http::transport::Error::Http {
status: status.as_u16(),
body: truncate_error_body(&text),
}));
}
let raw = RawResponse { body: text.clone() };
if let Some(observers) = observers {
observers.emit(litellm_host::lifecycle::CallEvent::Execution(
ExecutionEvent::ProviderResponseReceived { raw: raw.clone() },
));
}
interceptors
.after_provider_response(raw)
.await
.map_err(Error::post_call)?;
let body: Value = serde_json::from_str(&text).map_err(|err| {
Error::InvalidResponse(litellm_llms::ErrorDetail::invalid(
"chat completions response JSON",
err,
))
})?;
config
.transform_response(&model, ProviderChatResponseData { body })
.map_err(Error::from)
.map_err(as_response_error)
},
wire.url,
&wire.body,
timeout,
)?;
let response = crate::outbound::send(outbound, http).await.map_err(|err| {
// Failing to establish the connection means the request never went out,
// so the host can still serve it. Everything else here, a timeout
// above all, may have reached the provider and been answered.
if err.is_connect() || err.is_builder() {
Error::Transport(litellm_http::transport::Error::Connect(err.to_string()))
} else {
Error::Transport(litellm_http::transport::Error::Network(err.to_string()))
}
})?;
let status = response.status();
let text = response.text().await.map_err(|err| {
Error::Transport(litellm_http::transport::Error::Network(err.to_string()))
})?;
if !status.is_success() {
return Err(Error::Transport(litellm_http::transport::Error::Http {
status: status.as_u16(),
body: truncate_error_body(&text),
}));
}
hooks
.on_event(MachineEvent::ResponseReceived {
raw: RawResponse { body: text.clone() },
})
.await
.map_err(Error::post_call)?;
let body: Value = serde_json::from_str(&text).map_err(|err| {
Error::InvalidResponse(litellm_llms::ErrorDetail::invalid(
"chat completions response JSON",
err,
))
})?;
config
.transform_response(&model, ProviderChatResponseData { body })
.map_err(Error::from)
.map_err(as_response_error)
)
.await
}
/// Re-tag an error raised while normalizing a response the provider already
@ -168,7 +190,7 @@ mod tests {
raw: Mutex<Vec<String>>,
}
impl RouteHooks<Error> for RecordingHooks {
impl Interceptors<Error> for RecordingHooks {
async fn before_provider_request(
&self,
wire: WireRequest,
@ -188,8 +210,7 @@ mod tests {
})
}
async fn on_event(&self, event: MachineEvent) -> Result<(), Error> {
let MachineEvent::ResponseReceived { raw } = event;
async fn after_provider_response(&self, raw: RawResponse) -> Result<(), Error> {
self.raw.lock().unwrap().push(raw.body);
Ok(())
}
@ -223,13 +244,16 @@ mod tests {
)
.mount(&upstream)
.await;
let hooks = RecordingHooks::default();
let interceptors = RecordingHooks::default();
execute(
&Client::plain_for_test(),
&AuthServices::default(),
prepared(&upstream.uri()),
&hooks,
None,
None,
&interceptors,
None,
)
.await
.expect("chat completions call succeeds");
@ -240,14 +264,17 @@ mod tests {
assert_eq!(sent["system"], "added by the host");
assert_eq!(request.headers["x-host"], "seen");
assert_eq!(request.headers["x-api-key"], "sk-test");
let [context] = <[RequestContext; 1]>::try_from(hooks.contexts.into_inner().unwrap())
.unwrap_or_else(|seen| panic!("before_provider_request runs once, saw {}", seen.len()));
let [context] =
<[RequestContext; 1]>::try_from(interceptors.contexts.into_inner().unwrap())
.unwrap_or_else(|seen| {
panic!("before_provider_request runs once, saw {}", seen.len())
});
assert_eq!(
(context.model.as_str(), context.custom_llm_provider.as_str()),
("claude-sonnet-4-5", "anthropic")
);
assert_eq!(context.optional_params, json!({"max_tokens": 16}));
assert_eq!(hooks.raw.into_inner().unwrap(), [ANTHROPIC_MESSAGE]);
assert_eq!(interceptors.raw.into_inner().unwrap(), [ANTHROPIC_MESSAGE]);
}
#[rstest]
@ -258,13 +285,16 @@ mod tests {
.respond_with(ResponseTemplate::new(500).set_body_string("boom"))
.mount(&upstream)
.await;
let hooks = RecordingHooks::default();
let interceptors = RecordingHooks::default();
let error = execute(
&Client::plain_for_test(),
&AuthServices::default(),
prepared(&upstream.uri()),
&hooks,
None,
None,
&interceptors,
None,
)
.await
.expect_err("the upstream failure fails the call");
@ -273,7 +303,7 @@ mod tests {
error,
Error::Transport(litellm_http::transport::Error::Http { status: 500, .. })
));
assert!(hooks.raw.into_inner().unwrap().is_empty());
assert!(interceptors.raw.into_inner().unwrap().is_empty());
}
#[rstest::rstest]

View file

@ -1,10 +1,11 @@
use litellm_host::observation::ObservationSender;
pub mod route;
pub mod types;
pub use crate::error::RouteError as Error;
mod common_utils;
pub(crate) mod handler;
mod prepare;
use litellm_types::utils::ChatCompletionsResponse;
use litellm_llms_types::formats::chat_completions::ChatCompletionsResponse;
use prepare::{prepare_provider_request, resolve_request};
use crate::chat_completions::types::ChatCompletionsRequest;
@ -17,6 +18,7 @@ pub struct ChatCompletionsRoute {
http: litellm_http::Client,
auth: Arc<AuthServices>,
secrets: Arc<dyn SecretSource>,
cache: Option<litellm_cache_response::ScopedCache>,
}
impl ChatCompletionsRoute {
@ -29,44 +31,63 @@ impl ChatCompletionsRoute {
http,
auth,
secrets,
cache: None,
}
}
pub fn with_cache(self, cache: litellm_cache_response::ScopedCache) -> Self {
Self {
cache: Some(cache),
..self
}
}
pub async fn execute(
&self,
request: ChatCompletionsRequest<'_>,
hooks: &impl litellm_host::hooks::RouteHooks<Error>,
interceptors: &impl litellm_host::interceptors::Interceptors<Error>,
options: impl Into<crate::CallOptions>,
) -> Result<ChatCompletionsResponse, Error> {
litellm_host::lifecycle::observe_unary(hooks.observer(), self.run(request, hooks)).await
let crate::CallOptions {
cache: cache_options,
observers,
} = options.into();
litellm_host::lifecycle::observe_unary(
observers.clone(),
self.run_call(
request.into(),
cache_options,
interceptors,
observers.as_ref(),
),
)
.await
}
#[tracing::instrument(name = "litellm.route", skip_all, fields(
route = "chat_completions",
model = %request.model,
provider,
resolved_model,
stream = false,
outcome
))]
async fn run(
&self,
request: ChatCompletionsRequest<'_>,
hooks: &impl litellm_host::hooks::RouteHooks<Error>,
cache_options: Option<litellm_cache_response::CachePolicy>,
interceptors: &impl litellm_host::interceptors::Interceptors<Error>,
observers: Option<&ObservationSender>,
) -> Result<ChatCompletionsResponse, Error> {
crate::diagnostic::unary(async {
let resolved = resolve_request(request)?;
let snapshot = self
.secrets
.resolve(&resolved.config.secret_names())
.await?;
let prepared = prepare_provider_request(resolved, snapshot)?;
crate::diagnostic::provider(&prepared.model, &prepared.custom_llm_provider);
let execute: futures_util::future::BoxFuture<
'_,
Result<ChatCompletionsResponse, Error>,
> = Box::pin(handler::execute(&self.http, &self.auth, prepared, hooks));
execute.await
})
.await
let resolved = resolve_request(request)?;
let snapshot = self
.secrets
.resolve(&resolved.config.secret_names())
.await?;
let prepared = prepare_provider_request(resolved, snapshot)?;
crate::diagnostic::provider(&prepared.model, &prepared.custom_llm_provider);
let execute: futures_util::future::BoxFuture<'_, Result<ChatCompletionsResponse, Error>> =
Box::pin(handler::execute(
&self.http,
&self.auth,
prepared,
self.cache.clone(),
cache_options,
interceptors,
observers,
));
execute.await
}
}

View file

@ -2,8 +2,8 @@ use litellm_auth::SecretValue;
use litellm_core_utils::settings::Lookup;
use litellm_http::request::with_default_headers;
use litellm_llms::base_llm::{auth::ValidatedEnvironment, chat::transformation::BaseConfig};
use litellm_llms_types::formats::chat_completions::ChatMessage;
use litellm_secrets::source::Secrets;
use litellm_types::llms::openai::ChatMessage;
use serde_json::Value;
use super::{

View file

@ -1,10 +1,11 @@
use litellm_host::observation::ObservationSender;
use std::convert::Infallible;
use litellm_host::{
call::{CallOutput, HostedMachine, hosted_call},
protocol::Protocol,
};
use litellm_types::utils::ChatCompletionsResponse;
use litellm_llms_types::formats::chat_completions::ChatCompletionsResponse;
use super::{
ChatCompletionsRoute, Error,
@ -23,22 +24,59 @@ impl Protocol for ChatCompletions {
}
impl ChatCompletionsRoute {
pub fn machine(self, call: ChatCompletionsCall) -> HostedMachine<ChatCompletions> {
pub fn machine(
self,
call: ChatCompletionsCall,
options: impl Into<crate::CallOptions>,
) -> HostedMachine<ChatCompletions> {
let crate::CallOptions {
cache: cache_options,
observers,
} = options.into();
hosted_call(
call,
move |call: ChatCompletionsCall, _, hooks| async move {
let request = ChatCompletionsRequest {
model: &call.model,
messages: call.messages,
optional_params: call.optional_params,
api_key: call.api_key.as_deref(),
api_base: call.api_base.as_deref(),
custom_llm_provider: call.custom_llm_provider.as_deref(),
extra_headers: call.extra_headers,
timeout: call.timeout,
};
self.run(request, &hooks).await.map(CallOutput::Complete)
observers,
move |call, _, interceptors, observers| async move {
self.run_call(call, cache_options, &interceptors, observers.as_ref())
.await
.map(CallOutput::Complete)
},
)
}
#[tracing::instrument(name = "litellm.route", skip_all, fields(
route = "chat_completions",
model = %call.model,
provider,
resolved_model,
stream = false,
outcome
))]
pub(super) async fn run_call(
&self,
call: ChatCompletionsCall,
cache_options: Option<litellm_cache_response::CachePolicy>,
interceptors: &impl litellm_host::interceptors::Interceptors<Error>,
observers: Option<&ObservationSender>,
) -> Result<ChatCompletionsResponse, Error> {
crate::diagnostic::unary(async {
let request = ChatCompletionsRequest {
model: &call.model,
messages: call.messages,
optional_params: call.optional_params,
api_key: call.api_key.as_deref(),
api_base: call.api_base.as_deref(),
custom_llm_provider: call.custom_llm_provider.as_deref(),
extra_headers: call.extra_headers,
timeout: call.timeout,
};
self.run(request, cache_options, interceptors, observers)
.await
})
.await
}
}
impl crate::caching::Cachable for ChatCompletions {
const SURFACE: &'static str = "chat_completions";
}

View file

@ -3,7 +3,7 @@ use std::time::Duration;
use litellm_auth::SecretValue;
use litellm_llms::base_llm::{auth::ValidatedEnvironment, chat::transformation::BaseConfig};
use litellm_types::llms::openai::ChatMessage;
use litellm_llms_types::formats::chat_completions::ChatMessage;
use serde_json::{Map, Value};
/// A `/chat/completions` call as it crosses into the core.

View file

@ -0,0 +1,48 @@
use litellm_cache_response::CachePolicy;
use litellm_host::{
interceptors::{ExecutionFacts, Interceptors, RawResponse},
lifecycle::{CallEvent, ExecutionEvent},
observation::ObservationSender,
};
use crate::{CallOptions, RouteError};
pub(crate) struct CallContext<'a, I> {
pub interceptors: &'a I,
pub observers: Option<ObservationSender>,
pub cache: CachePolicy,
}
impl<'a, I: Interceptors<RouteError>> CallContext<'a, I> {
pub fn new(interceptors: &'a I, options: CallOptions) -> Self {
Self {
interceptors,
observers: options.observers,
cache: options.cache.unwrap_or_default(),
}
}
pub async fn result_ready(&self, facts: ExecutionFacts) -> Result<(), RouteError> {
if let Some(observers) = &self.observers {
observers.emit(CallEvent::Execution(ExecutionEvent::ResultReady {
facts: facts.clone(),
}));
}
self.interceptors.result_ready(facts).await
}
pub async fn response_received(&self, body: &str) -> Result<(), RouteError> {
let raw = RawResponse {
body: body.to_owned(),
};
if let Some(observers) = &self.observers {
observers.emit(CallEvent::Execution(
ExecutionEvent::ProviderResponseReceived { raw: raw.clone() },
));
}
self.interceptors
.after_provider_response(raw)
.await
.map_err(RouteError::post_call)
}
}

View file

@ -1,6 +1,8 @@
mod context;
mod diagnostic;
pub mod audio_transcription;
pub mod caching;
pub mod chat_completions;
pub mod constants;
pub mod error;
@ -12,3 +14,27 @@ pub mod resources;
pub mod responses;
pub use error::RouteError;
#[derive(Clone, Default)]
pub struct CallOptions {
pub cache: Option<litellm_cache_response::CachePolicy>,
pub observers: Option<litellm_host::observation::ObservationSender>,
}
impl From<Option<litellm_host::observation::ObservationSender>> for CallOptions {
fn from(observers: Option<litellm_host::observation::ObservationSender>) -> Self {
Self {
cache: None,
observers,
}
}
}
impl From<litellm_cache_response::CachePolicy> for CallOptions {
fn from(cache: litellm_cache_response::CachePolicy) -> Self {
Self {
cache: Some(cache),
observers: None,
}
}
}

View file

@ -1,4 +1,4 @@
This directory owns provider-independent Messages call orchestration: the entrypoint, call envelopes, provider selection, credential resolution, transport coordination, hooks, and stream lifecycle. Shared API data contracts belong in `litellm-types::messages`, adapter contracts and execution inputs in `llms/src/base_llm/messages`, and provider implementations in `llms/src/<provider>/messages`
This directory owns provider-independent Messages call orchestration: the entrypoint, call envelopes, provider selection, credential resolution, transport coordination, hooks, and stream lifecycle. Shared API data contracts belong in `litellm-llms-types::formats::messages`, adapter contracts and execution inputs in `llms/src/base_llm/messages`, and provider implementations in `llms/src/<provider>/messages`
Select concrete provider adapters and invoke their contracts. Delegate authentication policy, beta selection, payload rewriting, and response interpretation to those adapters. Keep provider policy out of request preparation and transport handlers. Calling a concrete provider helper for every provider is still a policy dependency

View file

@ -3,7 +3,7 @@ pub(super) use litellm_http::request::truncate_error_body;
use litellm_llms::{
anthropic::messages::transformation::ANTHROPIC_MESSAGES_CONFIG,
azure_ai::messages::transformation::AZURE_ANTHROPIC_MESSAGES_CONFIG,
base_llm::messages::transformation::BaseAnthropicMessagesConfig,
base_llm::messages::transformation::BaseMessagesConfig,
bedrock::messages::invoke_transformations::anthropic_claude3_transformation::BEDROCK_ANTHROPIC_MESSAGES_CONFIG,
};
use serde_json::{Map, Value};
@ -30,7 +30,7 @@ impl MessagesProvider {
.into()
}
pub(crate) fn config(self) -> &'static dyn BaseAnthropicMessagesConfig {
pub(crate) fn config(self) -> &'static dyn BaseMessagesConfig {
match self {
Self::Anthropic => &ANTHROPIC_MESSAGES_CONFIG,
Self::AzureAi => &AZURE_ANTHROPIC_MESSAGES_CONFIG,

View file

@ -2,95 +2,144 @@ use std::time::Duration;
use bytes::Bytes;
use futures_util::{StreamExt, TryStreamExt, stream::BoxStream};
use litellm_auth::AuthServices;
use litellm_host::{
event::{MachineEvent, RawResponse, RequestContext, WireRequest},
hooks::RouteHooks,
};
use litellm_host::interceptors::{Interceptors, ProviderIdentity, RequestContext, WireRequest};
use litellm_http::transport::Error as TransportError;
use litellm_llms::base_llm::{
auth::{Authenticated, resolve_auth},
messages::{
streaming::{ByteStream, StreamDecoder, encode_anthropic_sse},
transformation::BaseAnthropicMessagesConfig,
transformation::BaseMessagesConfig,
},
};
use litellm_llms_types::formats::messages::MessagesResponse;
use litellm_tracing::ByteChunk;
use litellm_types::llms::anthropic_messages::anthropic_response::AnthropicMessagesResponse;
use serde_json::Value;
use super::{
Error, MessagesResponse, common_utils::truncate_error_body, prepare::ProviderMessagesRequest,
Error, MessagesCallResponse, MessagesRoute, common_utils::truncate_error_body,
prepare::ProviderMessagesRequest,
};
use crate::{constants::MESSAGES_TIMEOUT_SECS, outbound::outbound_request};
use crate::{constants::MESSAGES_TIMEOUT_SECS, context::CallContext, outbound::outbound_request};
pub(super) async fn execute(
http: &litellm_http::Client,
auth: &AuthServices,
request: ProviderMessagesRequest,
hooks: &impl RouteHooks<Error>,
) -> Result<MessagesResponse, Error> {
let ProviderMessagesRequest {
provider,
url,
body,
environment,
timeout,
api_key,
} = request;
let stream = body.params.stream == Some(true);
let context = RequestContext {
model: body.model.clone(),
custom_llm_provider: provider.as_str().to_string(),
optional_params: serde_json::to_value(&body.params).map_err(serialize_failure)?,
secret_fields: Vec::new(),
api_key,
};
let authenticated = resolve_auth(auth, environment, &|key| std::env::var(key).ok()).await?;
let wire = hooks
.before_provider_request(
WireRequest {
url,
headers: authenticated.headers,
body: serde_json::to_value(&body).map_err(serialize_failure)?,
pub(super) struct ProviderCall {
pub identity: ProviderIdentity,
pub wire: WireRequest,
provider: super::common_utils::MessagesProvider,
signer: Option<litellm_auth_aws::SigV4Signer>,
timeout: Option<Duration>,
stream: bool,
}
impl ProviderCall {
pub fn cacheable(&self) -> bool {
self.signer.is_none()
}
}
impl MessagesRoute {
pub(super) async fn prepare_outbound(
&self,
request: ProviderMessagesRequest,
context: &CallContext<'_, impl Interceptors<Error>>,
) -> Result<ProviderCall, Error> {
let ProviderMessagesRequest {
provider,
url,
body,
environment,
timeout,
api_key,
} = request;
let request_context = RequestContext {
model: body.model.clone(),
custom_llm_provider: provider.as_str().to_string(),
optional_params: serde_json::to_value(&body.params).map_err(serialize_failure)?,
secret_fields: Vec::new(),
api_key,
};
let authenticated =
resolve_auth(&self.auth, environment, &|key| std::env::var(key).ok()).await?;
let identity = ProviderIdentity {
model: request_context.model.clone(),
provider: request_context.custom_llm_provider.clone(),
};
let wire = context
.interceptors
.before_provider_request(
WireRequest {
url,
headers: authenticated.headers,
body: serde_json::to_value(&body).map_err(serialize_failure)?,
},
request_context,
)
.await?;
let stream = match wire.body.get("stream") {
None | Some(Value::Null) => false,
Some(Value::Bool(stream)) => *stream,
Some(value) => {
return Err(Error::InvalidRequest(
litellm_llms::ErrorDetail::InvalidValue {
field: "stream",
expected: "a boolean",
actual: value.clone(),
},
));
}
};
Ok(ProviderCall {
identity,
wire,
provider,
signer: authenticated.signer,
timeout,
stream,
})
}
pub(super) async fn call_provider(
&self,
request: ProviderCall,
context: &CallContext<'_, impl Interceptors<Error>>,
) -> Result<MessagesCallResponse, Error> {
let ProviderCall {
identity,
wire,
provider,
signer,
timeout,
stream,
} = request;
let provider_name = provider.as_str();
log_request_body(provider_name, stream, &wire.body);
let response = send(
&self.http,
Authenticated {
headers: wire.headers,
signer,
},
context,
&wire.url,
&wire.body,
timeout,
)
.await?;
let provider_name = provider.as_str();
log_request_body(provider_name, stream, &wire.body);
let response = send(
http,
Authenticated {
headers: wire.headers,
signer: authenticated.signer,
},
&wire.url,
&wire.body,
timeout,
)
.await?;
if !response.status().is_success() {
return Err(provider_error(response).await);
if !response.status().is_success() {
return Err(provider_error(response).await);
}
let config = provider.config();
if stream {
return Ok(streaming_response(
response,
config.stream_decoder(),
provider_name,
));
}
let text = response.text().await.map_err(network)?;
log_response_body(&text);
context.response_received(&text).await?;
decode_response(config, &identity.model, &text)
.map(|message| MessagesCallResponse::Complete(Box::new(message)))
}
let config = provider.config();
if stream {
return Ok(streaming_response(
response,
config.stream_decoder(),
provider_name,
));
}
let text = response.text().await.map_err(network)?;
log_response_body(&text);
hooks
.on_event(MachineEvent::ResponseReceived {
raw: RawResponse { body: text.clone() },
})
.await
.map_err(Error::post_call)?;
decode_response(config, &body.model, &text)
.map(|message| MessagesResponse::Complete(Box::new(message)))
}
fn serialize_failure(err: serde_json::Error) -> Error {
@ -135,10 +184,10 @@ async fn provider_error(response: reqwest::Response) -> Error {
}
fn decode_response(
config: &dyn BaseAnthropicMessagesConfig,
config: &dyn BaseMessagesConfig,
model: &str,
text: &str,
) -> Result<AnthropicMessagesResponse, Error> {
) -> Result<MessagesResponse, Error> {
let response = serde_json::from_str(text).map_err(|err| {
Error::InvalidResponse(litellm_llms::ErrorDetail::invalid(
"messages response JSON",
@ -154,7 +203,7 @@ fn streaming_response(
response: reqwest::Response,
decoder: Option<StreamDecoder>,
provider: &'static str,
) -> MessagesResponse {
) -> MessagesCallResponse {
let headers = response
.headers()
.iter()
@ -171,7 +220,7 @@ fn streaming_response(
.boxed(),
Some(decode) => decoded_chunks(response, decode, provider),
};
MessagesResponse::Stream {
MessagesCallResponse::Stream {
head: super::route::MessagesStreamHead { headers },
chunks,
}
@ -245,7 +294,7 @@ mod tests {
.send()
.await
.unwrap();
let MessagesResponse::Stream { mut chunks, .. } =
let MessagesCallResponse::Stream { mut chunks, .. } =
streaming_response(response, Some(anthropic_sse_event_stream), "test")
else {
panic!("a streaming response returns chunks");

View file

@ -4,18 +4,23 @@ mod prepare;
pub mod route;
mod types;
use futures_util::FutureExt;
use litellm_auth::AuthServices;
use litellm_host::interceptors::{ExecutionFacts, Interceptors, ResultSource};
use crate::{caching::CallCache, context::CallContext};
use litellm_secrets::source::SecretSource;
use std::sync::Arc;
pub use crate::error::RouteError as Error;
pub use types::{MessagesCall, MessagesResponse, MessagesShaping, messages_body};
pub use types::{MessagesCall, MessagesCallResponse, MessagesShaping, messages_body};
#[derive(Clone)]
pub struct MessagesRoute {
http: litellm_http::Client,
auth: Arc<AuthServices>,
secrets: Arc<dyn SecretSource>,
cache: Option<litellm_cache_response::ScopedCache>,
}
impl MessagesRoute {
@ -28,15 +33,27 @@ impl MessagesRoute {
http,
auth,
secrets,
cache: None,
}
}
#[must_use]
pub fn with_cache(self, cache: litellm_cache_response::ScopedCache) -> Self {
Self {
cache: Some(cache),
..self
}
}
pub async fn execute(
&self,
call: MessagesCall,
hooks: &impl litellm_host::hooks::RouteHooks<Error>,
) -> Result<MessagesResponse, Error> {
litellm_host::lifecycle::observe_call(hooks.observer(), self.run(call, hooks)).await
interceptors: &impl litellm_host::interceptors::Interceptors<Error>,
options: impl Into<crate::CallOptions>,
) -> Result<MessagesCallResponse, Error> {
let context = CallContext::new(interceptors, options.into());
litellm_host::lifecycle::observe_call(context.observers.clone(), self.run(call, context))
.await
}
#[tracing::instrument(name = "litellm.route", skip_all, fields(
@ -50,14 +67,33 @@ impl MessagesRoute {
async fn run(
&self,
call: MessagesCall,
hooks: &impl litellm_host::hooks::RouteHooks<Error>,
) -> Result<MessagesResponse, Error> {
context: CallContext<'_, impl Interceptors<Error>>,
) -> Result<MessagesCallResponse, Error> {
crate::diagnostic::call(async {
let request = prepare::prepare(call, self.secrets.as_ref()).await?;
crate::diagnostic::provider(&request.body.model, request.provider.as_str());
let execute: futures_util::future::BoxFuture<'_, Result<MessagesResponse, Error>> =
Box::pin(handler::execute(&self.http, &self.auth, request, hooks));
execute.await
let prepared = prepare::prepare(call, self.secrets.as_ref()).await?;
crate::diagnostic::provider(&prepared.body.model, prepared.provider.as_str());
let request = self.prepare_outbound(prepared, &context).boxed().await?;
let cache = CallCache::<route::Messages>::from_wire(
self.cache.as_ref().filter(|_| request.cacheable()),
context.cache,
&request.identity,
&request.wire,
);
let identity = request.identity.clone();
let (output, source) = match cache.lookup().await {
Some(hit) => hit,
None => (
self.call_provider(request, &context).await?,
ResultSource::Provider,
),
};
context
.result_ready(ExecutionFacts {
provider: identity,
source: source.clone(),
})
.await?;
Ok(cache.finish(output, &source).await)
})
.await
}

View file

@ -9,8 +9,8 @@ use litellm_http::request::with_default_headers;
use litellm_llms::base_llm::{
auth::ValidatedEnvironment, messages::context::MessagesTransformContext,
};
use litellm_llms_types::formats::messages::MessagesRequest;
use litellm_secrets::source::SecretSource;
use litellm_types::llms::anthropic_messages::anthropic_request::AnthropicMessagesRequest;
use super::{
Error, MessagesCall,
@ -27,7 +27,7 @@ struct ResolvedProvider {
pub(super) struct ProviderMessagesRequest {
pub(super) provider: MessagesProvider,
pub(super) url: String,
pub(super) body: AnthropicMessagesRequest,
pub(super) body: MessagesRequest,
pub(super) environment: ValidatedEnvironment,
pub(super) timeout: Option<Duration>,
/// The caller's own credential, reported to the host beside the wire request.
@ -79,7 +79,7 @@ fn prepare_provider_request(
let env_lookup = |key: &str| secrets.get(key);
let sanitized = config.shape_request(
AnthropicMessagesRequest { model, ..body },
MessagesRequest { model, ..body },
shaping.reasoning_auto_summary,
)?;
let trimmed = without_additional_drop_params(sanitized, &shaping.additional_drop_params)?;
@ -124,9 +124,9 @@ fn prepare_provider_request(
}
fn without_additional_drop_params(
request: AnthropicMessagesRequest,
request: MessagesRequest,
paths: &[String],
) -> Result<AnthropicMessagesRequest, Error> {
) -> Result<MessagesRequest, Error> {
if paths.is_empty() {
return Ok(request);
}
@ -134,7 +134,7 @@ fn without_additional_drop_params(
let trimmed = paths
.iter()
.fold(params, |params, path| delete_nested_value(params, path));
Ok(AnthropicMessagesRequest {
Ok(MessagesRequest {
params: serde_json::from_value(trimmed).map_err(invalid_request)?,
..request
})
@ -143,7 +143,7 @@ fn without_additional_drop_params(
#[cfg(test)]
mod tests {
use litellm_llms::base_llm::auth::resolve_auth;
use litellm_types::utils::ProviderSpecificHeaders;
use litellm_llms_types::headers::ProviderSpecificHeaders;
use rstest::{fixture, rstest};
use serde_json::{Map, Value, json};
@ -155,7 +155,7 @@ mod tests {
MessagesShaping::default()
}
fn body(value: Value) -> AnthropicMessagesRequest {
fn body(value: Value) -> MessagesRequest {
serde_json::from_value(value).unwrap()
}

View file

@ -5,11 +5,11 @@ use litellm_host::{
call::{HostedCompletion, HostedMachine, hosted_call},
protocol::Protocol,
};
use litellm_types::llms::anthropic_messages::anthropic_response::AnthropicMessagesResponse;
use litellm_llms_types::formats::messages::MessagesResponse;
use super::{Error, MessagesCall};
pub type MessagesOutput = HostedCompletion<Box<AnthropicMessagesResponse>>;
pub type MessagesOutput = HostedCompletion<Box<MessagesResponse>>;
/// The upstream response as the caller sees it at stream hand-off, before any chunk.
pub struct MessagesStreamHead {
@ -19,7 +19,7 @@ pub struct MessagesStreamHead {
pub struct Messages;
impl Protocol for Messages {
type Response = Box<AnthropicMessagesResponse>;
type Response = Box<MessagesResponse>;
type Error = Error;
type Request = MessagesCall;
type HostCall = Infallible;
@ -30,9 +30,49 @@ impl Protocol for Messages {
pub type MessagesMachine = HostedMachine<Messages>;
impl super::MessagesRoute {
pub fn machine(self, request: super::MessagesCall) -> MessagesMachine {
hosted_call(request, move |call, _, hooks| async move {
self.run(call, &hooks).await
})
pub fn machine(
self,
request: super::MessagesCall,
options: impl Into<crate::CallOptions>,
) -> MessagesMachine {
let crate::CallOptions {
cache: cache_options,
observers,
} = options.into();
hosted_call(
request,
observers,
move |call, _, interceptors, observers| async move {
let context = crate::context::CallContext::new(
&interceptors,
crate::CallOptions {
cache: cache_options,
observers,
},
);
self.run(call, context).await
},
)
}
}
impl crate::caching::Cachable for Messages {
const SURFACE: &'static str = "messages";
}
impl crate::caching::StreamCachable for Messages {
const TERMINAL_EVENT: &'static str = "message_stop";
fn replay(data: bytes::Bytes) -> Option<litellm_host::call::OutputOf<Self>> {
Some(litellm_host::call::CallOutput::Stream {
head: MessagesStreamHead {
headers: Vec::new(),
},
chunks: Box::pin(futures_util::stream::iter([Ok(data)])),
})
}
fn bytes(chunk: &Self::Chunk) -> &[u8] {
chunk.as_ref()
}
}

View file

@ -2,12 +2,10 @@ use std::time::Duration;
use bytes::Bytes;
use litellm_host::call::CallOutput;
use litellm_llms::base_llm::messages::context::MessagesModelCapabilities as AnthropicModelCapabilities;
use litellm_types::{
llms::anthropic_messages::{
anthropic_request::AnthropicMessagesRequest, anthropic_response::AnthropicMessagesResponse,
},
utils::ProviderSpecificHeaders,
use litellm_llms::base_llm::messages::context::MessagesModelCapabilities;
use litellm_llms_types::{
formats::messages::{MessagesRequest, MessagesResponse},
headers::ProviderSpecificHeaders,
};
use serde::{Deserialize, Serialize};
use serde_json::{Map, Value};
@ -15,7 +13,7 @@ use serde_json::{Map, Value};
use super::Error;
pub struct MessagesCall {
pub body: AnthropicMessagesRequest,
pub body: MessagesRequest,
pub api_key: Option<String>,
pub api_base: Option<String>,
pub custom_llm_provider: Option<String>,
@ -25,7 +23,7 @@ pub struct MessagesCall {
pub shaping: MessagesShaping,
}
pub fn messages_body(body: Map<String, Value>) -> Result<AnthropicMessagesRequest, Error> {
pub fn messages_body(body: Map<String, Value>) -> Result<MessagesRequest, Error> {
serde_json::from_value(Value::Object(body)).map_err(invalid_request)
}
@ -33,13 +31,13 @@ pub(super) fn invalid_request(err: serde_json::Error) -> Error {
Error::InvalidRequest(format!("invalid Anthropic messages request: {err}").into())
}
pub type MessagesResponse =
CallOutput<Box<AnthropicMessagesResponse>, super::route::MessagesStreamHead, Bytes, Error>;
pub type MessagesCallResponse =
CallOutput<Box<MessagesResponse>, super::route::MessagesStreamHead, Bytes, Error>;
#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
pub struct MessagesShaping {
#[serde(default)]
pub capabilities: AnthropicModelCapabilities,
pub capabilities: MessagesModelCapabilities,
#[serde(default)]
pub drop_params: bool,
#[serde(default)]
@ -76,9 +74,9 @@ mod tests {
#[case::partial_capabilities(
json!({"capabilities": {"supports_reasoning": true}}),
MessagesShaping {
capabilities: AnthropicModelCapabilities {
capabilities: MessagesModelCapabilities {
supports_reasoning: true,
..AnthropicModelCapabilities::default()
..MessagesModelCapabilities::default()
},
..MessagesShaping::default()
},
@ -100,7 +98,7 @@ mod tests {
"additional_drop_params": ["metadata.user_id", "thinking"]
}),
MessagesShaping {
capabilities: AnthropicModelCapabilities {
capabilities: MessagesModelCapabilities {
supports_reasoning: true,
supports_adaptive_thinking: true,
thinking_always_on: false,

View file

@ -1,9 +1,9 @@
use litellm_host::observation::ObservationSender;
use std::sync::Arc;
use litellm_host::hooks::RouteHooks;
use litellm_llms::base_llm::ocr::{
error::Error, handler::OcrClient, transformation::LiteLLMOcrResponse,
};
use litellm_host::interceptors::Interceptors;
use litellm_llms::base_llm::ocr::{error::Error, handler::OcrClient};
use litellm_llms_types::formats::ocr::LiteLLMOcrResponse;
use super::{
handler::perform_ocr_request,
@ -23,9 +23,14 @@ impl OcrRoute {
pub async fn execute(
&self,
request: LiteLLMOcrRequest,
hooks: &impl RouteHooks<Error>,
interceptors: &impl Interceptors<Error>,
observers: Option<ObservationSender>,
) -> Result<LiteLLMOcrResponse, Error> {
litellm_host::lifecycle::observe_unary(hooks.observer(), self.run(request, hooks)).await
litellm_host::lifecycle::observe_unary(
observers.clone(),
self.run(request, interceptors, observers.as_ref()),
)
.await
}
#[tracing::instrument(name = "litellm.route", skip_all, fields(
@ -39,7 +44,8 @@ impl OcrRoute {
pub(super) async fn run(
&self,
request: LiteLLMOcrRequest,
hooks: &impl RouteHooks<Error>,
interceptors: &impl Interceptors<Error>,
observers: Option<&ObservationSender>,
) -> Result<LiteLLMOcrResponse, Error> {
crate::diagnostic::unary(async {
let caller_document = matches!(&request.document, OcrDocumentInput::Document(_));
@ -48,8 +54,9 @@ impl OcrRoute {
Box::pin(perform_ocr_request(
&self.client,
prepared,
hooks,
interceptors,
caller_document,
observers,
));
execute.await
})

View file

@ -1,10 +1,8 @@
use std::{collections::BTreeMap as Map, io::Read, path::Path};
use base64::{Engine, engine::general_purpose::STANDARD};
use litellm_llms::base_llm::ocr::{
error::Error,
transformation::{OCR_INLINE_MAX_BYTES, OcrDocument},
};
use litellm_llms::base_llm::ocr::{error::Error, transformation::OCR_INLINE_MAX_BYTES};
use litellm_llms_types::formats::ocr::OcrDocument;
use crate::ocr::types::OcrDocumentInput;

View file

@ -1,13 +1,12 @@
use futures_util::future::BoxFuture;
use litellm_host::{
event::{MachineEvent, RawResponse, RequestContext, WireRequest},
hooks::RouteHooks,
};
use litellm_host::interceptors::{Interceptors, RawResponse, RequestContext, WireRequest};
use litellm_host::{lifecycle::ExecutionEvent, observation::ObservationSender};
use litellm_llms::base_llm::ocr::{
error::Error,
handler::{CallHooks, OcrClient},
transformation::{LiteLLMOcrResponse, PreparedOcrRequest},
transformation::PreparedOcrRequest,
};
use litellm_llms_types::formats::ocr::LiteLLMOcrResponse;
use serde_json::Value;
use super::{arguments::is_secret_param, prepare::prepare_request, provider_config::OcrConfigKind};
@ -16,8 +15,9 @@ use crate::ocr::types::ResolvedOcrRequest;
pub(crate) async fn perform_ocr_request(
client: &OcrClient,
request: ResolvedOcrRequest,
host: &impl RouteHooks<Error>,
host: &impl Interceptors<Error>,
caller_document: bool,
observers: Option<&ObservationSender>,
) -> Result<LiteLLMOcrResponse, Error> {
request.response_format()?;
let config = request.config;
@ -27,19 +27,26 @@ pub(crate) async fn perform_ocr_request(
.await
.map_err(|error| Error::Secret(std::sync::Arc::new(error)))?;
let request = prepare_request(request, caller_document, client, secrets);
let hooks = OcrCallHooks::new(host, &request, config);
config.ocr(client, &request, &hooks).await
let interceptors = OcrCallHooks::new(host, &request, config, observers);
config.ocr(client, &request, &interceptors).await
}
struct OcrCallHooks<'a, H> {
hooks: &'a H,
interceptors: &'a H,
context: RequestContext,
observers: Option<&'a ObservationSender>,
}
impl<'a, H> OcrCallHooks<'a, H> {
fn new(hooks: &'a H, request: &PreparedOcrRequest, config: OcrConfigKind) -> Self {
fn new(
interceptors: &'a H,
request: &PreparedOcrRequest,
config: OcrConfigKind,
observers: Option<&'a ObservationSender>,
) -> Self {
Self {
hooks,
interceptors,
observers,
context: RequestContext {
model: request.model.clone(),
custom_llm_provider: <&str>::from(config.provider()).to_owned(),
@ -56,22 +63,26 @@ impl<'a, H> OcrCallHooks<'a, H> {
}
}
impl<H: RouteHooks<Error>> CallHooks<Error> for OcrCallHooks<'_, H> {
impl<H: Interceptors<Error>> CallHooks<Error> for OcrCallHooks<'_, H> {
fn before_provider_request(
&self,
wire: WireRequest,
) -> BoxFuture<'_, Result<WireRequest, Error>> {
Box::pin(
self.hooks
self.interceptors
.before_provider_request(wire, self.context.clone()),
)
}
fn response_received<'a>(&'a self, body: &'a [u8]) -> BoxFuture<'a, Result<(), Error>> {
Box::pin(self.hooks.on_event(MachineEvent::ResponseReceived {
raw: RawResponse {
body: String::from_utf8_lossy(body).into_owned(),
},
}))
let raw = RawResponse {
body: String::from_utf8_lossy(body).into_owned(),
};
if let Some(observers) = self.observers {
observers.emit(litellm_host::lifecycle::CallEvent::Execution(
ExecutionEvent::ProviderResponseReceived { raw: raw.clone() },
));
}
Box::pin(self.interceptors.after_provider_response(raw))
}
}

View file

@ -80,17 +80,18 @@ mod tests {
use futures_util::future::BoxFuture;
use litellm_core_utils::call_arguments::{CallArguments, compose_body, parse_options};
use litellm_host::event::WireRequest;
use litellm_host::interceptors::WireRequest;
use litellm_llms::{
base_llm::ocr::{
error::Error,
handler::{CallHooks, OcrClient},
transformation::{BaseOcrConfig, OcrResponseFormat},
transformation::BaseOcrConfig,
},
cohere::ocr::transformation::CohereParseConfig,
mistral::ocr::transformation::MistralOcrConfig,
vertex_ai::ocr::transformation::VertexAiOcrConfig,
};
use litellm_llms_types::formats::ocr::OcrResponseFormat;
use serde_json::{Value, json};
use super::*;
@ -100,7 +101,7 @@ mod tests {
wire::{OcrWireRequest, decode_request},
};
/// Stands in for a host with no hooks registered.
/// Stands in for a host with no interceptors registered.
struct NoHooks;
impl CallHooks<Error> for NoHooks {

View file

@ -14,8 +14,7 @@ use litellm_llms::{
error::Error,
handler::{self, CallHooks, OcrClient},
transformation::{
BaseOcrConfig, LiteLLMOcrResponse, OcrCredentialInputs, OcrDocument, OcrResponseFormat,
PreparedOcrRequest, ResolvedOcrCredentials,
BaseOcrConfig, OcrCredentialInputs, PreparedOcrRequest, ResolvedOcrCredentials,
},
},
cohere::ocr::transformation::CohereParseConfig,
@ -25,6 +24,7 @@ use litellm_llms::{
deepseek_transformation::VertexAIDeepSeekOCRConfig, transformation::VertexAiOcrConfig,
},
};
use litellm_llms_types::formats::ocr::{LiteLLMOcrResponse, OcrDocument, OcrResponseFormat};
macro_rules! with_config {
($kind:expr, $config:ident => $body:expr) => {
@ -133,9 +133,9 @@ impl OcrConfigKind {
self,
client: &OcrClient,
request: &PreparedOcrRequest,
hooks: &dyn CallHooks<Error>,
interceptors: &dyn CallHooks<Error>,
) -> Result<LiteLLMOcrResponse, Error> {
with_config!(self, config => handler::ocr(&config, client, request, hooks).await)
with_config!(self, config => handler::ocr(&config, client, request, interceptors).await)
}
}

View file

@ -1,11 +1,13 @@
use litellm_auth::ResolvedCredential;
use litellm_auth::{ResolvedCredential, TokenProviderHandle};
use litellm_host::observation::ObservationSender;
use litellm_host::{
call::{CallOutput, HostedMachine, hosted_call},
machine::{HostTokenProvider, TokenProtocol},
machine::HostServices,
protocol::Protocol,
protocol::Reply,
};
use litellm_llms::base_llm::ocr::{error::Error, transformation::LiteLLMOcrResponse};
use litellm_llms::base_llm::ocr::error::Error;
use litellm_llms_types::formats::ocr::LiteLLMOcrResponse;
use crate::ocr::types::{LiteLLMOcrRequest, OcrDocumentInput};
@ -31,27 +33,38 @@ impl Protocol for Ocr {
type StreamHead = std::convert::Infallible;
}
impl TokenProtocol for Ocr {
fn acquire_token_op(reply: Reply<ResolvedCredential>) -> OcrOp {
OcrOp::AcquireAzureAdToken(reply)
}
}
pub type OcrMachine = HostedMachine<Ocr>;
fn caller_token_provider(services: HostServices<Ocr>) -> TokenProviderHandle {
TokenProviderHandle::from_callback(move || {
let host_services = services.clone();
async move {
host_services
.call(OcrOp::AcquireAzureAdToken)
.await
.map_err(|error| {
litellm_auth::Error::CredentialAcquisition(error.to_string().into())
})
}
})
}
impl crate::ocr::OcrRoute {
pub fn machine(self, request: OcrCall) -> OcrMachine {
pub fn machine(self, request: OcrCall, observers: Option<ObservationSender>) -> OcrMachine {
hosted_call(
request,
move |projection: OcrCall, services, hooks| async move {
observers,
move |projection: OcrCall, services, interceptors, observers| async move {
let request = LiteLLMOcrRequest {
azure_ad_token_provider: projection
.caller_token
.then(|| HostTokenProvider::handle(services))
.then(|| caller_token_provider(services))
.or(projection.request.azure_ad_token_provider),
..projection.request
};
self.run(request, &hooks).await.map(CallOutput::Complete)
self.run(request, &interceptors, observers.as_ref())
.await
.map(CallOutput::Complete)
},
)
}

View file

@ -5,10 +5,9 @@ use litellm_auth::{InputSource, SecretValue, TokenProviderHandle};
use litellm_core_utils::call_arguments::CallArguments;
use litellm_llms::base_llm::ocr::{
error::Error,
transformation::{
OcrCredentialInputs, OcrDocument, OcrResponseFormat, OcrTransportConfig, response_format,
},
transformation::{OcrCredentialInputs, OcrTransportConfig, response_format},
};
use litellm_llms_types::formats::ocr::{OcrDocument, OcrResponseFormat};
use serde_json::{Map, Value};
use super::provider_config::{OcrConfigKind, resolve_provider_config};
@ -222,7 +221,7 @@ mod tests {
use super::*;
fn document() -> OcrDocument {
OcrDocument::try_from(
serde_json::from_value(
json!({"type":"document_url","document_url":"data:application/pdf;base64,YWJj"}),
)
.unwrap()

View file

@ -1,10 +1,8 @@
use std::{collections::BTreeMap, time::Duration};
use litellm_auth::{InputSource, SecretValue};
use litellm_llms::base_llm::ocr::{
error::Error,
transformation::{OcrDocument, decode_request_value},
};
use litellm_llms::base_llm::ocr::{error::Error, transformation::decode_request_value};
use litellm_llms_types::formats::ocr::OcrDocument;
use serde::Deserialize;
use serde_json::{Map, Value};

View file

@ -1,10 +1,9 @@
use litellm_host::lifecycle::ExecutionEvent;
use litellm_host::observation::ObservationSender;
use std::time::Duration;
use futures_util::StreamExt;
use litellm_host::{
event::{MachineEvent, RawResponse, WireRequest},
hooks::RouteHooks,
};
use litellm_host::interceptors::{Interceptors, RawResponse, WireRequest};
use litellm_llms::base_llm::auth::{Authenticated, resolve_auth};
use super::{
@ -16,10 +15,17 @@ pub(super) async fn execute(
http: &litellm_http::Client,
auth: &litellm_auth::AuthServices,
request: ProviderResponsesRequest,
hooks: &impl RouteHooks<Error>,
cache: Option<litellm_cache_response::ScopedCache>,
cache_options: Option<litellm_cache_response::CachePolicy>,
interceptors: &impl Interceptors<Error>,
observers: Option<&ObservationSender>,
) -> Result<ResponsesOutput, Error> {
let authenticated = resolve_auth(auth, request.environment, &|_| None).await?;
let wire = hooks
let identity = litellm_host::interceptors::ProviderIdentity {
model: request.context.model.clone(),
provider: request.context.custom_llm_provider.clone(),
};
let wire = interceptors
.before_provider_request(
WireRequest {
url: request.url,
@ -29,61 +35,80 @@ pub(super) async fn execute(
request.context,
)
.await?;
let stream = match wire.body.get("stream") {
None => false,
Some(serde_json::Value::Bool(value)) => *value,
Some(_) => return Err(Error::InvalidRequest("stream must be a boolean".into())),
};
let outbound = crate::outbound::outbound_request(
Authenticated {
headers: wire.headers,
signer: authenticated.signer,
let cache = cache.filter(|_| authenticated.signer.is_none());
let cache_request =
crate::caching::CacheRequest::from_wire(identity, cache.as_ref().map(|_| &wire));
crate::caching::execute_streaming::<super::route::Responses, _, _>(
cache_request,
cache.as_ref().map(|cache| cache.service.clone()),
cache.as_ref().map(|cache| cache.options(cache_options)),
interceptors,
observers,
|| async move {
let stream = match wire.body.get("stream") {
None => false,
Some(serde_json::Value::Bool(value)) => *value,
Some(_) => return Err(Error::InvalidRequest("stream must be a boolean".into())),
};
let outbound = crate::outbound::outbound_request(
Authenticated {
headers: wire.headers,
signer: authenticated.signer,
},
wire.url,
&wire.body,
Some(request.timeout.unwrap_or(Duration::from_secs(600))),
)?;
let response = crate::outbound::send(outbound, http)
.await
.map_err(network)?;
let status = response.status().as_u16();
if !response.status().is_success() {
let body = response.text().await.map_err(network)?;
return Err(litellm_http::transport::Error::Http {
status,
body: litellm_http::request::truncate_error_body(&body),
}
.into());
}
if stream {
let headers = response
.headers()
.iter()
.filter_map(|(name, value)| {
Some((name.to_string(), value.to_str().ok()?.to_owned()))
})
.collect();
let chunks = response
.bytes_stream()
.map(|chunk| chunk.map_err(network))
.boxed();
return Ok(ResponsesOutput::Stream {
head: ResponsesStreamHead { headers },
chunks,
});
}
let body = response.text().await.map_err(network)?;
let raw = RawResponse { body: body.clone() };
if let Some(observers) = observers {
observers.emit(litellm_host::lifecycle::CallEvent::Execution(
ExecutionEvent::ProviderResponseReceived { raw: raw.clone() },
));
}
interceptors
.after_provider_response(raw)
.await
.map_err(Error::post_call)?;
let value = serde_json::from_str(&body)
.map_err(|error| Error::InvalidResponse(error.to_string().into()))?;
request
.config
.transform_response_api_response(value)
.map(ResponsesOutput::Complete)
.map_err(Error::from)
},
wire.url,
&wire.body,
Some(request.timeout.unwrap_or(Duration::from_secs(600))),
)?;
let response = crate::outbound::send(outbound, http)
.await
.map_err(network)?;
let status = response.status().as_u16();
if !response.status().is_success() {
let body = response.text().await.map_err(network)?;
return Err(litellm_http::transport::Error::Http {
status,
body: litellm_http::request::truncate_error_body(&body),
}
.into());
}
if stream {
let headers = response
.headers()
.iter()
.filter_map(|(name, value)| Some((name.to_string(), value.to_str().ok()?.to_owned())))
.collect();
let chunks = response
.bytes_stream()
.map(|chunk| chunk.map_err(network))
.boxed();
return Ok(ResponsesOutput::Stream {
head: ResponsesStreamHead { headers },
chunks,
});
}
let body = response.text().await.map_err(network)?;
hooks
.on_event(MachineEvent::ResponseReceived {
raw: RawResponse { body: body.clone() },
})
.await
.map_err(Error::post_call)?;
let value = serde_json::from_str(&body)
.map_err(|error| Error::InvalidResponse(error.to_string().into()))?;
request
.config
.transform_response_api_response(value)
.map(ResponsesOutput::Complete)
.map_err(Error::from)
)
.await
}
fn network(error: reqwest::Error) -> Error {

View file

@ -1,4 +1,5 @@
pub use crate::error::RouteError as Error;
use litellm_host::observation::ObservationSender;
pub mod websocket;
mod handler;
@ -9,7 +10,7 @@ pub mod types;
use std::sync::Arc;
use litellm_auth::AuthServices;
use litellm_host::hooks::RouteHooks;
use litellm_host::interceptors::Interceptors;
use litellm_secrets::source::SecretSource;
use types::{ResponsesCall, ResponsesOutput};
@ -18,6 +19,7 @@ pub struct ResponsesRoute {
http: litellm_http::Client,
auth: Arc<AuthServices>,
secrets: Arc<dyn SecretSource>,
cache: Option<litellm_cache_response::ScopedCache>,
}
impl ResponsesRoute {
@ -30,15 +32,32 @@ impl ResponsesRoute {
http,
auth,
secrets,
cache: None,
}
}
pub fn with_cache(self, cache: litellm_cache_response::ScopedCache) -> Self {
Self {
cache: Some(cache),
..self
}
}
pub async fn execute(
&self,
call: ResponsesCall,
hooks: &impl RouteHooks<Error>,
interceptors: &impl Interceptors<Error>,
options: impl Into<crate::CallOptions>,
) -> Result<ResponsesOutput, Error> {
litellm_host::lifecycle::observe_call(hooks.observer(), self.run(call, hooks)).await
let crate::CallOptions {
cache: cache_options,
observers,
} = options.into();
litellm_host::lifecycle::observe_call(
observers.clone(),
self.run(call, cache_options, interceptors, observers.as_ref()),
)
.await
}
#[tracing::instrument(name = "litellm.route", skip_all, fields(
@ -52,18 +71,36 @@ impl ResponsesRoute {
async fn run(
&self,
call: ResponsesCall,
hooks: &impl RouteHooks<Error>,
cache_options: Option<litellm_cache_response::CachePolicy>,
interceptors: &impl litellm_host::interceptors::Interceptors<Error>,
observers: Option<&ObservationSender>,
) -> Result<ResponsesOutput, Error> {
crate::diagnostic::call(async {
let request = prepare::prepare(call, self.secrets.as_ref()).await?;
crate::diagnostic::provider(
&request.context.model,
&request.context.custom_llm_provider,
);
let execute: futures_util::future::BoxFuture<'_, Result<ResponsesOutput, Error>> =
Box::pin(handler::execute(&self.http, &self.auth, request, hooks));
execute.await
self.run_provider(call, cache_options, interceptors, observers)
.await
})
.await
}
async fn run_provider(
&self,
call: ResponsesCall,
cache_options: Option<litellm_cache_response::CachePolicy>,
interceptors: &impl Interceptors<Error>,
observers: Option<&ObservationSender>,
) -> Result<ResponsesOutput, Error> {
let request = prepare::prepare(call, self.secrets.as_ref()).await?;
crate::diagnostic::provider(&request.context.model, &request.context.custom_llm_provider);
let execute: futures_util::future::BoxFuture<'_, Result<ResponsesOutput, Error>> =
Box::pin(handler::execute(
&self.http,
&self.auth,
request,
self.cache.clone(),
cache_options,
interceptors,
observers,
));
execute.await
}
}

View file

@ -1,4 +1,4 @@
use litellm_host::event::RequestContext;
use litellm_host::interceptors::RequestContext;
use litellm_llms::{
base_llm::responses::transformation::BaseResponsesApiConfig,
openai::responses::transformation::OpenAiResponsesApiConfig,
@ -16,14 +16,9 @@ pub(super) async fn prepare(
call: ResponsesCall,
secrets: &dyn SecretSource,
) -> Result<ProviderResponsesRequest, Error> {
let provider = call.custom_llm_provider.as_deref().unwrap_or("openai");
if provider != "openai" {
return Err(Error::Unsupported("native HTTP responses provider"));
}
let model = call.model.strip_prefix("openai/").unwrap_or(&call.model);
if model.is_empty() || model.contains('/') {
return Err(Error::InvalidProvider(call.model));
}
let identity = resolve_provider(&call.model, call.custom_llm_provider.as_deref())?;
let provider = identity.provider.as_str();
let model = identity.model.as_str();
let config: &'static dyn BaseResponsesApiConfig = &OpenAiResponsesApiConfig;
let snapshot = secrets
.resolve(config.secret_names(call.api_key.as_deref(), call.api_base.as_deref()))
@ -56,3 +51,21 @@ pub(super) async fn prepare(
timeout: call.timeout,
})
}
pub(super) fn resolve_provider(
model: &str,
custom_llm_provider: Option<&str>,
) -> Result<litellm_host::interceptors::ProviderIdentity, Error> {
let provider = custom_llm_provider.unwrap_or("openai");
if provider != "openai" {
return Err(Error::Unsupported("native HTTP responses provider"));
}
let resolved = model.strip_prefix("openai/").unwrap_or(model);
if resolved.is_empty() || resolved.contains('/') {
return Err(Error::InvalidProvider(model.into()));
}
Ok(litellm_host::interceptors::ProviderIdentity {
model: resolved.into(),
provider: provider.into(),
})
}

View file

@ -5,7 +5,7 @@ use litellm_host::{
call::{HostedMachine, hosted_call},
protocol::Protocol,
};
use litellm_types::responses::main::ResponsesApiResponse;
use litellm_llms_types::formats::responses::ResponsesApiResponse;
use super::{
Error, ResponsesRoute,
@ -24,9 +24,51 @@ impl Protocol for Responses {
}
impl ResponsesRoute {
pub fn machine(self, call: ResponsesCall) -> HostedMachine<Responses> {
hosted_call(call, move |call, _, hooks| async move {
self.run(call, &hooks).await
})
pub fn machine(
self,
call: ResponsesCall,
options: impl Into<crate::CallOptions>,
) -> HostedMachine<Responses> {
let crate::CallOptions {
cache: cache_options,
observers,
} = options.into();
hosted_call(
call,
observers,
move |call, _, interceptors, observers| async move {
self.run(call, cache_options, &interceptors, observers.as_ref())
.await
},
)
}
}
impl crate::caching::Cachable for Responses {
const SURFACE: &'static str = "responses";
fn reusable(response: &Self::Response) -> bool {
response
.extra
.get("status")
.and_then(serde_json::Value::as_str)
== Some("completed")
}
}
impl crate::caching::StreamCachable for Responses {
const TERMINAL_EVENT: &'static str = "response.completed";
fn replay(data: bytes::Bytes) -> Option<litellm_host::call::OutputOf<Self>> {
Some(litellm_host::call::CallOutput::Stream {
head: ResponsesStreamHead {
headers: Vec::new(),
},
chunks: Box::pin(futures_util::stream::iter([Ok(data)])),
})
}
fn bytes(chunk: &Self::Chunk) -> &[u8] {
chunk.as_ref()
}
}

View file

@ -5,7 +5,7 @@ use litellm_host::call::CallOutput;
use litellm_llms::base_llm::{
auth::ValidatedEnvironment, responses::transformation::BaseResponsesApiConfig,
};
use litellm_types::responses::main::ResponsesApiResponse;
use litellm_llms_types::formats::responses::ResponsesApiResponse;
use serde_json::{Map, Value};
use super::Error;
@ -32,6 +32,6 @@ pub(super) struct ProviderResponsesRequest {
pub environment: ValidatedEnvironment,
pub url: String,
pub body: Value,
pub context: litellm_host::event::RequestContext,
pub context: litellm_host::interceptors::RequestContext,
pub timeout: Option<Duration>,
}

View file

@ -2,7 +2,7 @@ use std::{collections::HashMap, sync::Arc, time::Duration};
use futures_util::{SinkExt, StreamExt};
use litellm_http::websocket::{UpstreamWebSocket, connect_upstream};
use litellm_types::responses::streaming_websocket::ResponsesWsEventType;
use litellm_llms_types::formats::responses::streaming_websocket::ResponsesWsEventType;
use tokio::sync::Mutex;
use tokio_tungstenite::tungstenite::{
Message,

File diff suppressed because it is too large Load diff

View file

@ -1,8 +1,13 @@
use litellm_host::interceptors::RawResponse;
use litellm_host::{
interceptors::{ExecutionFacts, ResultSource},
lifecycle::ExecutionEvent,
};
use std::time::Duration;
use litellm_core::chat_completions::{Error, types::ChatCompletionsRequest};
use litellm_http::transport::Error as TransportError;
use litellm_types::utils::ChatCompletionsResponse;
use litellm_llms_types::formats::chat_completions::ChatCompletionsResponse;
use rstest::{fixture, rstest};
use serde_json::{Map, Value, json};
use wiremock::ResponseTemplate;
@ -13,7 +18,7 @@ use support::*;
const ANTHROPIC_MESSAGE: &str = r#"{"id":"msg_1","type":"message","role":"assistant","model":"claude-sonnet-4-5-20260101","content":[{"type":"text","text":"hello"}],"stop_reason":"end_turn","stop_sequence":null,"usage":{"input_tokens":11,"output_tokens":4}}"#;
async fn complete(request: ChatCompletionsRequest<'_>) -> Result<ChatCompletionsResponse, Error> {
chat_completions_route().execute(request, &()).await
chat_completions_route().execute(request, &(), None).await
}
fn object(value: Value) -> Map<String, Value> {
@ -253,7 +258,7 @@ async fn direct_and_hosted_calls_share_hooks_and_lifecycle(
#[case] hosted: bool,
) {
use litellm_core::chat_completions::route::ChatCompletions;
use litellm_host::{call::HostedCompletion, event::CallEvent};
use litellm_host::{call::HostedCompletion, lifecycle::CallEvent};
let upstream = upstream([anthropic_response(ANTHROPIC_MESSAGE)]).await;
let base = upstream.uri();
@ -265,8 +270,9 @@ async fn direct_and_hosted_calls_share_hooks_and_lifecycle(
.into(),
);
let response = if hosted {
let result = litellm_host::in_process::run_hosted(
chat_completions_route().machine(host.request().unwrap()),
let result = litellm_host_native::in_process::run_hosted(
chat_completions_route()
.machine(host.request().unwrap(), Some(host.events.0.sender.clone())),
host.runtime(),
)
.await
@ -290,6 +296,7 @@ async fn direct_and_hosted_calls_share_hooks_and_lifecycle(
timeout: call.timeout,
},
&host,
Some(host.events.0.sender.clone()),
)
.await
.unwrap()
@ -307,7 +314,13 @@ async fn direct_and_hosted_calls_share_hooks_and_lifecycle(
&events[..],
[
CallEvent::Started { .. },
CallEvent::Machine(_),
CallEvent::Execution(ExecutionEvent::ProviderResponseReceived { .. }),
CallEvent::Execution(ExecutionEvent::ResultReady {
facts: ExecutionFacts {
source: ResultSource::Provider,
..
}
}),
CallEvent::Succeeded { .. }
]
));
@ -318,12 +331,9 @@ async fn direct_and_hosted_calls_share_hooks_and_lifecycle(
async fn a_post_call_hook_failure_never_looks_safe_to_retry(
request: ChatCompletionsRequest<'static>,
) {
use litellm_host::{
event::{MachineEvent, RequestContext, WireRequest},
hooks::RouteHooks,
};
use litellm_host::interceptors::{Interceptors, RequestContext, WireRequest};
struct FailingHook;
impl RouteHooks<Error> for FailingHook {
impl Interceptors<Error> for FailingHook {
async fn before_provider_request(
&self,
wire: WireRequest,
@ -331,7 +341,7 @@ async fn a_post_call_hook_failure_never_looks_safe_to_retry(
) -> Result<WireRequest, Error> {
Ok(wire)
}
async fn on_event(&self, _: MachineEvent) -> Result<(), Error> {
async fn after_provider_response(&self, _: RawResponse) -> Result<(), Error> {
Err(Error::InvalidRequest("callback rejected".into()))
}
}
@ -344,6 +354,7 @@ async fn a_post_call_hook_failure_never_looks_safe_to_retry(
..request
},
&FailingHook,
None,
)
.await
.unwrap_err();

View file

@ -1,7 +1,11 @@
use litellm_host::lifecycle::ExecutionEvent;
use std::sync::Mutex;
use litellm_core::messages::route::Messages;
use litellm_host::event::{CallEvent, MachineEvent, RequestContext, WireRequest};
use litellm_core::messages::{MessagesCallResponse, route::Messages};
use litellm_host::{
interceptors::{ExecutionFacts, RequestContext, ResultSource, WireRequest},
lifecycle::CallEvent,
};
use litellm_llms::base_llm::messages::context::MessagesModelCapabilities as AnthropicModelCapabilities;
use rstest::rstest;
@ -14,8 +18,10 @@ type Rewrite = Box<dyn Fn(WireRequest) -> Result<WireRequest, Error> + Send + Sy
struct RecordingHost {
call: LocalMessagesHost,
rewrite: Rewrite,
events: Mutex<Vec<CallEvent>>,
events: super::support::Observations,
optional_params: Mutex<Vec<Value>>,
facts: Mutex<Vec<ExecutionFacts>>,
reject_result: bool,
}
impl RecordingHost {
@ -23,8 +29,10 @@ impl RecordingHost {
Self {
call: LocalMessagesHost::new(call),
rewrite,
events: Mutex::new(Vec::new()),
events: super::support::Observations::default(),
optional_params: Mutex::new(Vec::new()),
facts: Mutex::new(Vec::new()),
reject_result: false,
}
}
@ -38,7 +46,7 @@ impl RecordingHost {
.unwrap()
.iter()
.filter_map(|event| match event {
CallEvent::Machine(MachineEvent::ResponseReceived { raw }) => {
CallEvent::Execution(ExecutionEvent::ProviderResponseReceived { raw }) => {
Some(raw.body.clone())
}
_ => None,
@ -51,24 +59,32 @@ impl RecordingHost {
pub fn request(&self) -> Result<MessagesCall, Error> {
self.call.request()
}
pub fn runtime(&self) -> litellm_host::in_process::Host<'_, (), Self, ()> {
litellm_host::in_process::Host {
pub fn runtime(&self) -> litellm_host_native::in_process::Host<'_, (), Self, ()> {
litellm_host_native::in_process::Host {
services: &(),
hooks: self,
interceptors: self,
stream: &(),
observer: Some(self),
observers: Some(&self.events.sender),
}
}
}
impl litellm_host::lifecycle::CallObserver for RecordingHost {
fn observe(&self, event: litellm_host::event::CallEvent) {
self.events.lock().unwrap().push(event.clone());
fn observe(&self, event: litellm_host::lifecycle::CallEvent) {
self.events.sender.emit(event);
}
}
impl litellm_host::hooks::RouteHooks<<Messages as litellm_host::protocol::Protocol>::Error>
impl litellm_host::interceptors::Interceptors<<Messages as litellm_host::protocol::Protocol>::Error>
for RecordingHost
{
async fn result_ready(&self, facts: ExecutionFacts) -> Result<(), Error> {
self.facts.lock().unwrap().push(facts);
if self.reject_result {
return Err(Error::Unsupported("result rejected"));
}
Ok(())
}
async fn before_provider_request(
&self,
wire: WireRequest,
@ -80,20 +96,107 @@ impl litellm_host::hooks::RouteHooks<<Messages as litellm_host::protocol::Protoc
.push(context.optional_params.clone());
(self.rewrite)(wire)
}
async fn on_event(
async fn after_provider_response(
&self,
event: litellm_host::event::MachineEvent,
raw: litellm_host::interceptors::RawResponse,
) -> Result<(), <Messages as litellm_host::protocol::Protocol>::Error> {
litellm_host::lifecycle::CallObserver::observe(
self,
litellm_host::event::CallEvent::Machine(event),
litellm_host::lifecycle::CallEvent::Execution(
litellm_host::lifecycle::ExecutionEvent::ProviderResponseReceived { raw },
),
);
Ok(())
}
}
#[rstest]
#[case::native_unary(false, false)]
#[case::native_stream(true, false)]
#[case::hosted_unary(false, true)]
#[case::hosted_stream(true, true)]
#[tokio::test]
async fn rejected_results_are_not_delivered_or_cached(
call: MessagesCall,
#[case] streaming: bool,
#[case] hosted: bool,
) {
use futures_util::TryStreamExt;
use litellm_cache_memory::InMemoryCache;
use litellm_cache_response::{CacheScope, ResponseCache, ScopedCache};
let response = if streaming {
ResponseTemplate::new(200).set_body_raw(
"event: message_stop\ndata: {\"type\":\"message_stop\"}\n\n",
"text/event-stream",
)
} else {
message_response()
};
let upstream = upstream([response.clone(), response]).await;
let route = messages_route(no_secrets()).with_cache(ScopedCache::new(
Arc::new(ResponseCache::new(Arc::new(InMemoryCache::new(
Some(100),
Some(Duration::from_secs(60)),
)))),
CacheScope::Shared,
));
for (reject, expected_requests, cached) in [
(true, 1, false),
(false, 2, false),
(true, 2, true),
(false, 2, true),
] {
let request = authenticated(
with_fields(
MessagesCall {
body: call.body.clone(),
..super::call()
},
json!({"stream": streaming}),
),
upstream.uri(),
);
let host = RecordingHost {
reject_result: reject,
..RecordingHost::passthrough(request)
};
let result = if hosted {
litellm_host_native::in_process::run_hosted(
route.clone().machine(host.request().unwrap(), None),
host.runtime(),
)
.await
.map(|_| ())
} else {
match route.execute(host.request().unwrap(), &host, None).await {
Ok(MessagesCallResponse::Complete(_)) => Ok(()),
Ok(MessagesCallResponse::Stream { chunks, .. }) => {
chunks.try_collect::<Vec<_>>().await.map(|_| ())
}
Err(error) => Err(error),
}
};
assert_eq!(
result,
if reject {
Err(Error::Unsupported("result rejected"))
} else {
Ok(())
}
);
assert_eq!(received(&upstream).await.len(), expected_requests);
let facts = host.facts.lock().unwrap();
assert_eq!(facts.len(), 1);
assert_eq!(
matches!(facts[0].source, ResultSource::Cache { .. }),
cached
);
}
}
async fn run_through(host: &RecordingHost) -> Result<MessagesOutput, Error> {
litellm_host::in_process::run_hosted(
litellm_host_native::in_process::run_hosted(
machine(Arc::new(RecordingSecrets::empty()))(host.request()?),
host.runtime(),
)
@ -137,6 +240,81 @@ async fn what_before_send_returns_is_what_the_provider_receives(call: MessagesCa
assert_eq!(request.header("x-api-key"), Some("sk-ant"));
}
#[rstest]
#[case::enable(false, json!(true), Some(true))]
#[case::disable(true, json!(false), Some(false))]
#[case::null(true, Value::Null, Some(false))]
#[case::invalid(false, json!("true"), None)]
#[tokio::test]
async fn response_mode_follows_the_intercepted_request(
call: MessagesCall,
traces: TraceCapture,
#[case] original_stream: bool,
#[case] rewritten_stream: Value,
#[case] expected_stream: Option<bool>,
) {
use futures_util::TryStreamExt;
let sse = "event: message_stop\ndata: {\"type\":\"message_stop\"}\n\n";
let response = if expected_stream == Some(true) {
ResponseTemplate::new(200).set_body_raw(sse, "text/event-stream")
} else {
message_response()
};
let upstream = upstream([response]).await;
let rewrite = rewritten_stream.clone();
let host = RecordingHost::new(
authenticated(
with_fields(call, json!({"stream": original_stream})),
upstream.uri(),
),
Box::new(move |wire| {
let mut body = wire.body;
body["stream"] = rewrite.clone();
Ok(WireRequest { body, ..wire })
}),
);
let result = traces
.logger()
.instrument(async {
let output = messages_route(no_secrets())
.execute(host.request()?, &host, None)
.await?;
match output {
MessagesCallResponse::Stream { chunks, .. } => {
assert_eq!(expected_stream, Some(true));
assert_eq!(
chunks.try_collect::<Vec<_>>().await?.concat(),
sse.as_bytes()
);
}
MessagesCallResponse::Complete(message) => {
assert_eq!(expected_stream, Some(false));
assert_eq!(*message, serde_json::from_value(message_body()).unwrap());
}
}
Ok::<_, Error>(())
})
.await;
let summaries = traces.summaries("litellm.route");
assert_eq!(summaries.len(), 1);
let Some(expected_stream) = expected_stream else {
assert!(matches!(result, Err(Error::InvalidRequest(_))));
assert!(received(&upstream).await.is_empty());
assert_eq!(summaries[0]["outcome"], "failure");
return;
};
result.unwrap();
assert_eq!(
only_request(&upstream).await.json()["stream"],
rewritten_stream
);
assert_eq!(host.raw_responses().len(), usize::from(!expected_stream));
assert_eq!(summaries[0]["stream"], expected_stream);
assert_eq!(summaries[0]["outcome"], "success");
}
#[rstest]
#[tokio::test]
async fn a_before_send_failure_never_sends(call: MessagesCall) {

View file

@ -8,10 +8,8 @@ use litellm_core::messages::{
route::{Messages, MessagesMachine, MessagesOutput},
};
use litellm_http::{HttpSettings, Resolution};
use litellm_llms_types::formats::messages::{MessagesRequest, MessagesResponse};
use litellm_secrets::source::SecretSource;
use litellm_types::llms::anthropic_messages::{
anthropic_request::AnthropicMessagesRequest, anthropic_response::AnthropicMessagesResponse,
};
use rstest::fixture;
use serde_json::{Map, Value, json};
use wiremock::ResponseTemplate;
@ -35,7 +33,7 @@ fn object(value: Value) -> Map<String, Value> {
map
}
fn body(value: Value) -> AnthropicMessagesRequest {
fn body(value: Value) -> MessagesRequest {
serde_json::from_value(value).unwrap()
}
@ -99,7 +97,7 @@ fn headers<'a>(pairs: impl IntoIterator<Item = (&'a str, &'a str)>) -> Option<Ma
}
fn machine(secrets: Arc<dyn SecretSource>) -> impl FnOnce(MessagesCall) -> MessagesMachine {
move |request| messages_route(secrets).machine(request)
move |request| messages_route(secrets).machine(request, None)
}
async fn run_with(
@ -107,7 +105,8 @@ async fn run_with(
call: MessagesCall,
) -> Result<MessagesOutput, Error> {
let host = LocalMessagesHost::new(call);
litellm_host::in_process::run_hosted(machine(secrets)(host.request()?), host.runtime()).await
litellm_host_native::in_process::run_hosted(machine(secrets)(host.request()?), host.runtime())
.await
}
/// Runs the route with a secret source that knows nothing, so no environment leaks in.
@ -115,7 +114,7 @@ async fn run(call: MessagesCall) -> Result<MessagesOutput, Error> {
run_with(Arc::new(RecordingSecrets::empty()), call).await
}
async fn run_message(call: MessagesCall) -> AnthropicMessagesResponse {
async fn run_message(call: MessagesCall) -> MessagesResponse {
match run(call).await.expect("messages call succeeds") {
MessagesOutput::Complete(message) => *message,
MessagesOutput::StreamEnded | MessagesOutput::Detached => {
@ -144,39 +143,41 @@ impl LocalMessagesHost {
.take()
.ok_or_else(|| Error::InvalidRequest("messages request was already projected".into()))
}
pub fn runtime(&self) -> litellm_host::in_process::Host<'_, (), Self, ()> {
litellm_host::in_process::Host {
pub fn runtime(&self) -> litellm_host_native::in_process::Host<'_, (), Self, ()> {
litellm_host_native::in_process::Host {
services: &(),
hooks: self,
interceptors: self,
stream: &(),
observer: Some(self),
observers: None,
}
}
}
impl litellm_host::lifecycle::CallObserver for LocalMessagesHost {
fn observe(&self, _: litellm_host::event::CallEvent) {}
fn observe(&self, _: litellm_host::lifecycle::CallEvent) {}
}
impl litellm_host::hooks::RouteHooks<<Messages as litellm_host::protocol::Protocol>::Error>
impl litellm_host::interceptors::Interceptors<<Messages as litellm_host::protocol::Protocol>::Error>
for LocalMessagesHost
{
async fn before_provider_request(
&self,
wire: litellm_host::event::WireRequest,
_: litellm_host::event::RequestContext,
wire: litellm_host::interceptors::WireRequest,
_: litellm_host::interceptors::RequestContext,
) -> Result<
litellm_host::event::WireRequest,
litellm_host::interceptors::WireRequest,
<Messages as litellm_host::protocol::Protocol>::Error,
> {
Ok(wire)
}
async fn on_event(
async fn after_provider_response(
&self,
event: litellm_host::event::MachineEvent,
raw: litellm_host::interceptors::RawResponse,
) -> Result<(), <Messages as litellm_host::protocol::Protocol>::Error> {
litellm_host::lifecycle::CallObserver::observe(
self,
litellm_host::event::CallEvent::Machine(event),
litellm_host::lifecycle::CallEvent::Execution(
litellm_host::lifecycle::ExecutionEvent::ProviderResponseReceived { raw },
),
);
Ok(())
}

View file

@ -1,6 +1,8 @@
use litellm_llms::base_llm::messages::context::{MessagesModelCapabilities, SupportedEffortTiers};
use litellm_types::llms::anthropic::{AnthropicBeta, BetaSet};
use litellm_types::utils::{ProviderSpecificHeader, ProviderSpecificHeaders};
use litellm_llms_types::{
headers::{ProviderSpecificHeader, ProviderSpecificHeaders},
providers::anthropic::{AnthropicBeta, BetaSet},
};
use rstest::rstest;
use super::*;

View file

@ -1,17 +1,27 @@
use litellm_core::messages::{MessagesResponse, messages_body};
use litellm_core::messages::{MessagesCallResponse, messages_body};
use litellm_host::{
interceptors::{ExecutionFacts, ResultSource},
lifecycle::ExecutionEvent,
};
use litellm_http::transport::Error as TransportError;
use rstest::rstest;
use super::*;
#[rstest]
#[case::without_hooks(false)]
#[case::with_hooks(true)]
#[case::neither(false, false)]
#[case::hooks_only(true, false)]
#[case::observer_only(false, true)]
#[case::both(true, true)]
#[tokio::test]
async fn calls_defer_execution_until_polled(call: MessagesCall, #[case] with_hooks: bool) {
async fn calls_defer_execution_until_polled(
call: MessagesCall,
#[case] with_hooks: bool,
#[case] with_observer: bool,
) {
use futures_util::future::BoxFuture;
use litellm_host::event::CallEvent;
use litellm_host::lifecycle::CallEvent;
let upstream = upstream([message_response()]).await;
let secrets = Arc::new(RecordingSecrets::new([("ANTHROPIC_API_KEY", "test-key")]));
@ -21,17 +31,19 @@ async fn calls_defer_execution_until_polled(call: MessagesCall, #[case] with_hoo
..call
});
let request = host.request().unwrap();
let future: BoxFuture<'_, Result<MessagesResponse, Error>> = if with_hooks {
Box::pin(route.execute(request, &host))
let observer: Option<litellm_host::observation::ObservationSender> =
with_observer.then(|| host.events.0.sender.clone());
let future: BoxFuture<'_, Result<MessagesCallResponse, Error>> = if with_hooks {
Box::pin(route.execute(request, &host, observer))
} else {
Box::pin(route.execute(request, &()))
Box::pin(route.execute(request, &(), observer))
};
assert!(secrets.requested().is_empty());
assert!(host.events.0.lock().unwrap().is_empty());
assert!(received(&upstream).await.is_empty());
let MessagesResponse::Complete(response) = future.await.unwrap() else {
let MessagesCallResponse::Complete(response) = future.await.unwrap() else {
panic!("expected a completed message");
};
assert_eq!(
@ -43,18 +55,41 @@ async fn calls_defer_execution_until_polled(call: MessagesCall, #[case] with_hoo
assert_eq!(sent.header("x-api-key"), Some("test-key"));
assert_eq!(sent.header("x-hook"), with_hooks.then_some("called"));
let events = host.events.0.lock().unwrap();
if with_hooks {
assert!(matches!(
&events[..],
[
CallEvent::Started { .. },
CallEvent::Machine(_),
CallEvent::Succeeded { .. }
]
));
} else {
assert!(events.is_empty());
}
assert!(matches!(
(with_hooks, with_observer, events.as_slice()),
(false, false, [])
| (true, false, [])
| (
false,
true,
[
CallEvent::Started { .. },
CallEvent::Execution(ExecutionEvent::ProviderResponseReceived { .. }),
CallEvent::Execution(ExecutionEvent::ResultReady {
facts: ExecutionFacts {
source: ResultSource::Provider,
..
}
}),
CallEvent::Succeeded { .. }
]
)
| (
true,
true,
[
CallEvent::Started { .. },
CallEvent::Execution(ExecutionEvent::ProviderResponseReceived { .. }),
CallEvent::Execution(ExecutionEvent::ResultReady {
facts: ExecutionFacts {
source: ResultSource::Provider,
..
}
}),
CallEvent::Succeeded { .. }
]
)
));
}
#[rstest]
@ -246,11 +281,12 @@ async fn the_facade_sends_through_the_injected_http_pool_configuration(call: Mes
..call
},
&(),
None,
)
.await
.expect("messages request succeeds");
let MessagesResponse::Complete(message) = response else {
let MessagesCallResponse::Complete(message) = response else {
panic!("a non-streaming request returns a message");
};
assert_eq!(message.id, "msg_1");
@ -301,3 +337,120 @@ async fn message_route_summary_excludes_payload_diagnostics(
assert!(summaries[0].get("body").is_none());
assert!(!format!("{:?}", traces.records()).contains("private-key-sentinel"));
}
#[rstest]
#[case::uncached(false, 2)]
#[case::cached(true, 1)]
#[tokio::test]
async fn route_uses_injected_dependencies_and_optional_cache(
#[case] caching: bool,
#[case] expected_requests: usize,
) {
use litellm_cache_memory::InMemoryCache;
use litellm_cache_response::{CacheScope, ResponseCache, ScopedCache};
use litellm_core::messages::MessagesRoute;
let upstream = upstream([message_response(), message_response()]).await;
let resources = resources();
let route = MessagesRoute::new(
provider_http(&resources, &http_config()),
resources.auth.clone(),
Arc::new(RecordingSecrets::new([("ANTHROPIC_API_KEY", "route-key")])),
);
let route = if caching {
route.with_cache(ScopedCache::new(
Arc::new(ResponseCache::new(Arc::new(InMemoryCache::new(
Some(100),
Some(Duration::from_secs(60)),
)))),
CacheScope::Shared,
))
} else {
route
};
for _ in 0..2 {
let request = MessagesCall {
api_base: Some(upstream.uri()),
..super::call()
};
let MessagesCallResponse::Complete(response) =
route.execute(request, &(), None).await.unwrap()
else {
panic!("expected a completed message");
};
assert_eq!(
response.content,
message_body()["content"].as_array().unwrap().as_slice()
);
}
let requests = received(&upstream).await;
assert_eq!(requests.len(), expected_requests);
assert_eq!(requests[0].header("x-api-key"), Some("route-key"));
}
#[rstest]
#[tokio::test]
async fn cache_overrides_preserve_the_routes_isolated_scope(call: MessagesCall) {
use litellm_cache_memory::InMemoryCache;
use litellm_cache_response::{CachePolicy, CacheScope, ResponseCache, ScopedCache};
let first_body = message_body();
let second_body = Value::Object(
first_body
.as_object()
.unwrap()
.iter()
.map(|(key, value)| {
(
key.clone(),
if key == "id" {
json!("msg_second")
} else {
value.clone()
},
)
})
.collect(),
);
let upstream = upstream([
json_response(first_body.clone()),
json_response(second_body.clone()),
])
.await;
let service = Arc::new(ResponseCache::new(Arc::new(InMemoryCache::new(
Some(100),
Some(Duration::from_secs(60)),
))));
let first = messages_route(no_secrets()).with_cache(ScopedCache::new(
service.clone(),
CacheScope::Isolated("first".into()),
));
let second = messages_route(no_secrets()).with_cache(ScopedCache::new(
service,
CacheScope::Isolated("second".into()),
));
for (route, expected) in [
(&first, &first_body),
(&second, &second_body),
(&first, &first_body),
(&second, &second_body),
] {
let request = MessagesCall {
body: call.body.clone(),
api_key: Some("same-key".into()),
api_base: Some(upstream.uri()),
..super::call()
};
let override_options = CachePolicy {
ttl: Some(Duration::from_secs(30)),
..CachePolicy::default()
};
let MessagesCallResponse::Complete(response) =
route.execute(request, &(), override_options).await.unwrap()
else {
panic!("expected a completed message");
};
assert_eq!(response.id, expected["id"].as_str().unwrap());
}
assert_eq!(received(&upstream).await.len(), 2);
}

View file

@ -1,12 +1,14 @@
use std::sync::{Mutex, mpsc};
use std::{
ops::ControlFlow,
sync::{Mutex, mpsc},
};
use bytes::Bytes;
use futures_util::{StreamExt, TryStreamExt};
use litellm_core::messages::{
MessagesResponse,
MessagesCallResponse,
route::{Messages, MessagesStreamHead},
};
use litellm_host::protocol::Demand;
use litellm_tracing::{Logger, Metadata, Record, Sink};
use rstest::rstest;
use tokio::{
@ -60,12 +62,12 @@ impl RecordingStreamHost {
}
}
fn record(&self, op: Seen) -> Demand {
fn record(&self, op: Seen) -> ControlFlow<()> {
let mut seen = self.seen.lock().unwrap();
seen.push(op);
match seen.len() < self.detach_after {
true => Demand::More,
false => Demand::Detached,
true => ControlFlow::Continue(()),
false => ControlFlow::Break(()),
}
}
}
@ -74,47 +76,49 @@ impl RecordingStreamHost {
pub fn request(&self) -> Result<MessagesCall, Error> {
self.call.request()
}
pub fn runtime(&self) -> litellm_host::in_process::Host<'_, (), Self, Self> {
litellm_host::in_process::Host {
pub fn runtime(&self) -> litellm_host_native::in_process::Host<'_, (), Self, Self> {
litellm_host_native::in_process::Host {
services: &(),
hooks: self,
interceptors: self,
stream: self,
observer: Some(self),
observers: None,
}
}
}
impl litellm_host::in_process::StreamConsumer<Messages> for RecordingStreamHost {
async fn open_stream(&self, head: MessagesStreamHead) -> Result<Demand, Error> {
impl litellm_host_native::in_process::StreamConsumer<Messages> for RecordingStreamHost {
async fn open_stream(&self, head: MessagesStreamHead) -> Result<ControlFlow<()>, Error> {
Ok(self.record(Seen::Open(head.headers)))
}
async fn send_chunk(&self, chunk: Bytes) -> Result<Demand, Error> {
async fn send_chunk(&self, chunk: Bytes) -> Result<ControlFlow<()>, Error> {
Ok(self.record(Seen::Deliver(chunk)))
}
}
impl litellm_host::lifecycle::CallObserver for RecordingStreamHost {
fn observe(&self, _: litellm_host::event::CallEvent) {}
fn observe(&self, _: litellm_host::lifecycle::CallEvent) {}
}
impl litellm_host::hooks::RouteHooks<<Messages as litellm_host::protocol::Protocol>::Error>
impl litellm_host::interceptors::Interceptors<<Messages as litellm_host::protocol::Protocol>::Error>
for RecordingStreamHost
{
async fn before_provider_request(
&self,
wire: litellm_host::event::WireRequest,
_: litellm_host::event::RequestContext,
wire: litellm_host::interceptors::WireRequest,
_: litellm_host::interceptors::RequestContext,
) -> Result<
litellm_host::event::WireRequest,
litellm_host::interceptors::WireRequest,
<Messages as litellm_host::protocol::Protocol>::Error,
> {
Ok(wire)
}
async fn on_event(
async fn after_provider_response(
&self,
event: litellm_host::event::MachineEvent,
raw: litellm_host::interceptors::RawResponse,
) -> Result<(), <Messages as litellm_host::protocol::Protocol>::Error> {
litellm_host::lifecycle::CallObserver::observe(
self,
litellm_host::event::CallEvent::Machine(event),
litellm_host::lifecycle::CallEvent::Execution(
litellm_host::lifecycle::ExecutionEvent::ProviderResponseReceived { raw },
),
);
Ok(())
}
@ -136,7 +140,7 @@ fn sse_response() -> ResponseTemplate {
}
async fn stream_through(host: &RecordingStreamHost) -> Result<MessagesOutput, Error> {
litellm_host::in_process::run_hosted(
litellm_host_native::in_process::run_hosted(
machine(Arc::new(RecordingSecrets::empty()))(host.request()?),
host.runtime(),
)
@ -344,11 +348,12 @@ async fn the_sdk_returns_stream_headers_and_every_sse_byte(
..streaming(call, upstream.uri())
},
&(),
None,
)
.await
.unwrap();
let MessagesResponse::Stream { head, chunks } = response else {
let MessagesCallResponse::Stream { head, chunks } = response else {
panic!("a streaming request returns a stream");
};
for (name, value) in UPSTREAM_HEADERS {
@ -364,7 +369,7 @@ async fn the_sdk_returns_stream_headers_and_every_sse_byte(
async fn the_sdk_returns_http_errors_before_opening_a_stream(call: MessagesCall) {
let upstream = upstream([ResponseTemplate::new(429).set_body_string("slow down")]).await;
let error = messages_route(no_secrets())
.execute(streaming(call, upstream.uri()), &())
.execute(streaming(call, upstream.uri()), &(), None)
.await
.err()
.expect("upstream failure is returned by messages()");
@ -395,13 +400,14 @@ async fn dropping_the_sdk_stream_closes_the_unfinished_upstream(
..streaming(call, base)
},
&(),
None,
),
)
.await
.expect("messages() returns before the upstream finishes")
.unwrap();
let MessagesResponse::Stream { mut chunks, .. } = response else {
let MessagesCallResponse::Stream { mut chunks, .. } = response else {
panic!("a streaming request returns a stream");
};
if read_chunk {
@ -431,11 +437,12 @@ async fn the_sdk_yields_a_body_error_once_after_delivered_chunks(call: MessagesC
..streaming(call, base)
},
&(),
None,
)
.await
.unwrap();
let MessagesResponse::Stream { mut chunks, .. } = response else {
let MessagesCallResponse::Stream { mut chunks, .. } = response else {
panic!("a streaming request returns a stream");
};
assert_eq!(

View file

@ -3,8 +3,10 @@ use std::{
time::Duration,
};
use litellm_host::event::{CallEvent, MachineEvent};
use litellm_host::lifecycle::CallEvent;
use litellm_host::lifecycle::ExecutionEvent;
use litellm_llms::base_llm::ocr::settings::OcrSettings;
use rstest::rstest;
use super::*;
@ -190,7 +192,7 @@ async fn client_settings_choose_the_api_version_and_the_inch_to_pixel_dpi() {
});
let result = route
.execute(read_request(&upstream.uri(), json!({})), &())
.execute(read_request(&upstream.uri(), json!({})), &(), None)
.await
.unwrap();
@ -281,7 +283,7 @@ async fn response_received_fires_for_the_submission_and_the_completed_poll() {
let recorder = observed.clone();
let host =
LocalOcrHost::new(read_request(&upstream.uri(), json!({}))).with_observer(move |event| {
if let CallEvent::Machine(MachineEvent::ResponseReceived { raw }) = event {
if let CallEvent::Execution(ExecutionEvent::ProviderResponseReceived { raw }) = event {
recorder.lock().unwrap().push(raw.body.clone());
}
});
@ -354,7 +356,7 @@ async fn the_polling_deadline_bounds_the_retry_delay() {
let error = tokio::time::timeout(
Duration::from_secs(1),
route.execute(read_request(&upstream.uri(), json!({})), &()),
route.execute(read_request(&upstream.uri(), json!({})), &(), None),
)
.await
.expect("the deadline cuts the retry delay short")

View file

@ -1,6 +1,6 @@
use base64::Engine;
use litellm_core::ocr::types::OcrDocumentInput;
use litellm_host::event::WireRequest;
use litellm_host::interceptors::WireRequest;
use rstest::rstest;
use wiremock::{Mock, matchers::any};
@ -208,8 +208,8 @@ async fn configured_client_preserves_document_url_policy(#[case] allowed: bool)
json!({"type": "document_url", "document_url": document_url}),
json!({}),
));
let result = litellm_host::in_process::run_hosted(
route.machine(host.request().unwrap()),
let result = litellm_host_native::in_process::run_hosted(
route.machine(host.request().unwrap(), None),
host.runtime(),
)
.await;

View file

@ -1,18 +1,27 @@
use std::sync::{Arc, Mutex};
use litellm_host::interceptors::RawResponse;
use litellm_host::lifecycle::ExecutionEvent;
use std::sync::{
Arc, Mutex,
atomic::{AtomicUsize, Ordering},
};
use litellm_core::ocr::{
route::{Ocr, OcrCall, OcrOp},
types::OcrDocumentInput,
};
use litellm_host::event::{CallEvent, MachineEvent, RequestContext, WireRequest};
use litellm_host::{
interceptors::{RequestContext, WireRequest},
lifecycle::CallEvent,
};
use rstest::rstest;
use super::*;
pub(crate) fn event_name(event: &CallEvent) -> &'static str {
match event {
CallEvent::Execution(ExecutionEvent::ResultReady { .. }) => "result_ready",
CallEvent::Started { .. } => "started",
CallEvent::Machine(MachineEvent::ResponseReceived { .. }) => "response",
CallEvent::Execution(ExecutionEvent::ProviderResponseReceived { .. }) => "response",
CallEvent::Succeeded { .. } => "success",
CallEvent::Failed { .. } => "failure",
CallEvent::Cancelled { .. } => "cancelled",
@ -22,15 +31,12 @@ pub(crate) fn event_name(event: &CallEvent) -> &'static str {
fn recording_host(
request: LiteLLMOcrRequest,
events: Arc<Mutex<Vec<&'static str>>>,
interceptions: Arc<AtomicUsize>,
block: bool,
) -> LocalOcrHost {
let before_send_events = events.clone();
LocalOcrHost::new(request)
.with_before_send(move |wire, _| {
before_send_events
.lock()
.unwrap()
.push("before_provider_request");
interceptions.fetch_add(1, Ordering::SeqCst);
match block {
true => Err(Error::InvalidRequest("blocked".into())),
false => Ok(wire),
@ -41,22 +47,22 @@ fn recording_host(
#[rstest::rstest]
#[tokio::test]
async fn hooks_run_in_order_and_one_success_is_emitted() {
async fn interception_runs_once_and_observers_receive_ordered_success_events() {
let upstream = upstream([pages_response()]).await;
let events = Arc::new(Mutex::new(Vec::new()));
let interceptions = Arc::new(AtomicUsize::new(0));
perform_with(recording_host(
ocr_request("mistral/model", &upstream.uri(), json!({})),
events.clone(),
interceptions.clone(),
false,
))
.await
.unwrap();
assert_eq!(
*events.lock().unwrap(),
["started", "before_provider_request", "response", "success"]
);
assert_eq!(interceptions.load(Ordering::SeqCst), 1);
assert_eq!(*events.lock().unwrap(), ["started", "response", "success"]);
assert_eq!(received(&upstream).await.len(), 1);
}
@ -65,10 +71,12 @@ async fn hooks_run_in_order_and_one_success_is_emitted() {
async fn a_blocking_before_send_prevents_the_call_and_emits_one_failure() {
let upstream = upstream([pages_response()]).await;
let events = Arc::new(Mutex::new(Vec::new()));
let interceptions = Arc::new(AtomicUsize::new(0));
let error = perform_with(recording_host(
ocr_request("mistral/model", &upstream.uri(), json!({})),
events.clone(),
interceptions.clone(),
true,
))
.await
@ -78,10 +86,8 @@ async fn a_blocking_before_send_prevents_the_call_and_emits_one_failure() {
matches!(&error, Error::InvalidRequest(message) if message == "blocked"),
"{error:?}"
);
assert_eq!(
*events.lock().unwrap(),
["started", "before_provider_request", "failure"]
);
assert_eq!(interceptions.load(Ordering::SeqCst), 1);
assert_eq!(*events.lock().unwrap(), ["started", "failure"]);
assert!(received(&upstream).await.is_empty());
}
@ -90,19 +96,19 @@ async fn a_blocking_before_send_prevents_the_call_and_emits_one_failure() {
async fn an_upstream_failure_emits_one_terminal_failure() {
let upstream = upstream([status_response(500, json!({"error": "failed"}))]).await;
let events = Arc::new(Mutex::new(Vec::new()));
let interceptions = Arc::new(AtomicUsize::new(0));
let result = perform_with(recording_host(
ocr_request("mistral/model", &upstream.uri(), json!({})),
events.clone(),
interceptions.clone(),
false,
))
.await;
assert!(result.is_err());
assert_eq!(
*events.lock().unwrap(),
["started", "before_provider_request", "failure"]
);
assert_eq!(interceptions.load(Ordering::SeqCst), 1);
assert_eq!(*events.lock().unwrap(), ["started", "failure"]);
assert_eq!(received(&upstream).await.len(), 1);
}
@ -114,7 +120,7 @@ async fn an_invalid_provider_response_is_observed_before_normalization_fails() {
let recorder = observed.clone();
let host = LocalOcrHost::new(ocr_request("mistral/model", &upstream.uri(), json!({})))
.with_observer(move |event| {
if let CallEvent::Machine(MachineEvent::ResponseReceived { raw }) = event {
if let CallEvent::Execution(ExecutionEvent::ProviderResponseReceived { raw }) = event {
recorder.lock().unwrap().push(raw.body.clone());
}
});
@ -208,16 +214,16 @@ impl CallerTokenHost {
caller_token: true,
})
}
pub fn runtime(&self) -> litellm_host::in_process::Host<'_, Self, Self, ()> {
litellm_host::in_process::Host {
pub fn runtime(&self) -> litellm_host_native::in_process::Host<'_, Self, Self, ()> {
litellm_host_native::in_process::Host {
services: self,
hooks: self,
interceptors: self,
stream: &(),
observer: Some(self),
observers: None,
}
}
}
impl litellm_host::services::HostCallHandler<Ocr> for CallerTokenHost {
impl litellm_host_native::services::HostCallHandler<Ocr> for CallerTokenHost {
async fn handle_host_call(&self, op: OcrOp) -> Result<(), Error> {
match op {
OcrOp::AcquireAzureAdToken(reply) => {
@ -232,9 +238,9 @@ impl litellm_host::services::HostCallHandler<Ocr> for CallerTokenHost {
}
impl litellm_host::lifecycle::CallObserver for CallerTokenHost {
fn observe(&self, _: litellm_host::event::CallEvent) {}
fn observe(&self, _: litellm_host::lifecycle::CallEvent) {}
}
impl litellm_host::hooks::RouteHooks<<Ocr as litellm_host::protocol::Protocol>::Error>
impl litellm_host::interceptors::Interceptors<<Ocr as litellm_host::protocol::Protocol>::Error>
for CallerTokenHost
{
async fn before_provider_request(
@ -263,13 +269,15 @@ impl litellm_host::hooks::RouteHooks<<Ocr as litellm_host::protocol::Protocol>::
.collect();
Ok(WireRequest { headers, ..wire })
}
async fn on_event(
async fn after_provider_response(
&self,
event: litellm_host::event::MachineEvent,
raw: litellm_host::interceptors::RawResponse,
) -> Result<(), <Ocr as litellm_host::protocol::Protocol>::Error> {
litellm_host::lifecycle::CallObserver::observe(
self,
litellm_host::event::CallEvent::Machine(event),
litellm_host::lifecycle::CallEvent::Execution(
litellm_host::lifecycle::ExecutionEvent::ProviderResponseReceived { raw },
),
);
Ok(())
}
@ -288,8 +296,8 @@ async fn the_callers_azure_token_is_acquired_before_before_send_which_can_still_
trace: Mutex::new(Vec::new()),
};
litellm_host::in_process::run_hosted(
ocr_route().machine(host.request().unwrap()),
litellm_host_native::in_process::run_hosted(
ocr_route().machine(host.request().unwrap(), None),
host.runtime(),
)
.await
@ -312,15 +320,11 @@ async fn the_callers_azure_token_is_acquired_before_before_send_which_can_still_
#[rstest]
#[tokio::test]
async fn direct_execution_uses_hooks_without_a_machine() {
use litellm_host::{hooks::RouteHooks, lifecycle::CallObserver};
use litellm_host::interceptors::Interceptors;
struct Hooks(Arc<super::support::CallEvents>);
impl RouteHooks<Error> for Hooks {
fn observer(&self) -> Option<Arc<dyn CallObserver>> {
Some(self.0.clone())
}
struct Hooks;
impl Interceptors<Error> for Hooks {
async fn before_provider_request(
&self,
wire: WireRequest,
@ -336,8 +340,7 @@ async fn direct_execution_uses_hooks_without_a_machine() {
})
}
async fn on_event(&self, event: MachineEvent) -> Result<(), Error> {
self.0.observe(CallEvent::Machine(event));
async fn after_provider_response(&self, _: RawResponse) -> Result<(), Error> {
Ok(())
}
}
@ -348,10 +351,11 @@ async fn direct_execution_uses_hooks_without_a_machine() {
.await;
let events = Arc::new(super::support::CallEvents::default());
let route = ocr_route();
let hooks = Hooks(events.clone());
let interceptors = Hooks;
let builder = route.execute(
ocr_request("mistral/model", &upstream.uri(), json!({})),
&hooks,
&interceptors,
Some(events.0.sender.clone()),
);
assert!(events.0.lock().unwrap().is_empty());
assert!(received(&upstream).await.is_empty());
@ -365,7 +369,7 @@ async fn direct_execution_uses_hooks_without_a_machine() {
&events.0.lock().unwrap()[..],
[
CallEvent::Started { .. },
CallEvent::Machine(_),
CallEvent::Execution(_),
CallEvent::Succeeded { .. }
]
));

View file

@ -1,4 +1,3 @@
use litellm_host::protocol::HookRequest;
use std::{
sync::{
Arc,
@ -12,17 +11,16 @@ use litellm_core::ocr::{
types::OcrDocumentInput,
};
use litellm_host::{
event::{CallEvent, WireRequest},
hooks::RouteHooks,
interceptors::{Interceptors, WireRequest},
machine::{HostFailure, Machine, MachineStep},
protocol::Suspension,
services::HostCallHandler,
protocol::{HostRequest, InterceptRequest},
};
use litellm_host_native::services::HostCallHandler;
use litellm_llms::base_llm::ocr::transformation::OcrTransportConfig;
use rstest::rstest;
use tokio::{io::AsyncReadExt, net::TcpListener, sync::Notify};
use super::{lifecycle::event_name, *};
use super::*;
/// Drives the machine by hand, answering every op through `host` except `before_provider_request`,
/// which `intercept` answers so a test can fail or cancel exactly there.
@ -34,7 +32,7 @@ async fn drive_until(
Vec<&'static str>,
OcrMachine,
) {
let mut machine = ocr_route().machine(host.request().unwrap());
let mut machine = ocr_route().machine(host.request().unwrap(), None);
let mut ops = Vec::new();
let outcome = loop {
let op = match machine.resume().await {
@ -43,23 +41,30 @@ async fn drive_until(
Err(error) => break Err(error),
};
let answer = match op {
Suspension::Stream(stream) => match stream {
HostRequest::Intercept(InterceptRequest::ResultReady { facts, reply }) => host
.result_ready(facts)
.await
.map(|()| reply.send(()))
.map_err(HostFailure::Error),
HostRequest::Stream(stream) => match stream {
litellm_host::protocol::StreamDelivery::Open(head, _) => match head {},
litellm_host::protocol::StreamDelivery::Chunk(chunk, _) => match chunk {},
},
Suspension::HostCall(op) => {
HostRequest::HostCall(op) => {
ops.push(match op {
OcrOp::AcquireAzureAdToken(_) => "AcquireAzureAdToken",
});
host.handle_host_call(op).await.map_err(HostFailure::Error)
}
Suspension::Hook(HookRequest::BeforeProviderRequest { wire, reply, .. }) => {
HostRequest::Intercept(InterceptRequest::BeforeProviderRequest {
wire, reply, ..
}) => {
ops.push("BeforeSend");
intercept(*wire).map(|wire| reply.send(wire))
}
Suspension::Hook(HookRequest::Event(event, reply)) => {
ops.push(event_name(&CallEvent::Machine(event.clone())));
host.on_event(event)
HostRequest::Intercept(InterceptRequest::AfterProviderResponse { raw, reply }) => {
ops.push("response");
host.after_provider_response(raw)
.await
.map(|()| reply.send(()))
.map_err(HostFailure::Error)
@ -80,10 +85,11 @@ async fn drive_until_notified(machine: &mut OcrMachine, host: &LocalOcrHost, sto
_ = stop.notified() => break,
step = machine.resume() => {
match step.unwrap() {
MachineStep::Suspended(Suspension::HostCall(op)) => host.handle_host_call(op).await.unwrap(),
MachineStep::Suspended(Suspension::Hook(HookRequest::BeforeProviderRequest { wire, reply, .. })) => reply.send(*wire),
MachineStep::Suspended(Suspension::Hook(HookRequest::Event(_, reply))) => reply.send(()),
MachineStep::Suspended(Suspension::Stream(stream)) => match stream {
MachineStep::Suspended(HostRequest::Intercept(InterceptRequest::ResultReady { reply, .. })) => reply.send(()),
MachineStep::Suspended(HostRequest::HostCall(op)) => host.handle_host_call(op).await.unwrap(),
MachineStep::Suspended(HostRequest::Intercept(InterceptRequest::BeforeProviderRequest { wire, reply, .. })) => reply.send(*wire),
MachineStep::Suspended(HostRequest::Intercept(InterceptRequest::AfterProviderResponse { reply, .. })) => reply.send(()),
MachineStep::Suspended(HostRequest::Stream(stream)) => match stream {
litellm_host::protocol::StreamDelivery::Open(head, _) => match head {},
litellm_host::protocol::StreamDelivery::Chunk(chunk, _) => match chunk {},
},
@ -181,15 +187,16 @@ async fn a_before_send_failure_ends_the_call_without_reaching_transport(
async fn resuming_before_answering_keeps_the_pending_operation() {
let upstream = upstream([pages_response()]).await;
let request = ocr_request("mistral/model", &upstream.uri(), json!({}));
let mut machine = ocr_route().machine(OcrCall {
request,
caller_token: false,
});
let Ok(MachineStep::Suspended(Suspension::Hook(HookRequest::BeforeProviderRequest {
wire,
reply,
..
}))) = machine.resume().await
let mut machine = ocr_route().machine(
OcrCall {
request,
caller_token: false,
},
None,
);
let Ok(MachineStep::Suspended(HostRequest::Intercept(
InterceptRequest::BeforeProviderRequest { wire, reply, .. },
))) = machine.resume().await
else {
panic!("expected the provider request hook");
};
@ -197,8 +204,8 @@ async fn resuming_before_answering_keeps_the_pending_operation() {
reply.send(*wire);
assert!(matches!(
machine.resume().await,
Ok(MachineStep::Suspended(Suspension::Hook(
HookRequest::Event(_, _)
Ok(MachineStep::Suspended(HostRequest::Intercept(
InterceptRequest::AfterProviderResponse { .. }
)))
));
}
@ -244,7 +251,7 @@ async fn interrupt_drops_provider_captures_before_returning() {
},
)));
let host = LocalOcrHost::new(request);
let mut machine = ocr_route().machine(host.request().unwrap());
let mut machine = ocr_route().machine(host.request().unwrap(), None);
drive_until_notified(&mut machine, &host, &entered).await;
assert!(!dropped.load(Ordering::SeqCst));
@ -280,7 +287,7 @@ async fn interrupting_an_in_flight_provider_request_closes_its_connection() {
while socket.read(&mut buffer).await.unwrap() != 0 {}
});
let host = LocalOcrHost::new(ocr_request("mistral/model", &base, json!({})));
let mut machine = ocr_route().machine(host.request().unwrap());
let mut machine = ocr_route().machine(host.request().unwrap(), None);
drive_until_notified(&mut machine, &host, &received).await;
let cancelled = Error::InvalidRequest("cancelled".into());

View file

@ -5,12 +5,12 @@ use litellm_core::ocr::{
types::{LiteLLMOcrRequest, OcrDocumentInput},
wire::{OcrWireRequest, decode_request},
};
use litellm_host::event::{CallEvent, RequestContext, WireRequest};
use litellm_llms::base_llm::ocr::{
error::Error,
settings::OcrSettings,
transformation::{LiteLLMOcrResponse, OcrDocument},
use litellm_host::{
interceptors::{RequestContext, WireRequest},
lifecycle::CallEvent,
};
use litellm_llms::base_llm::ocr::{error::Error, settings::OcrSettings};
use litellm_llms_types::formats::ocr::{LiteLLMOcrResponse, OcrDocument};
use serde_json::{Map, Value, json};
use std::sync::Mutex;
use wiremock::{MockServer, ResponseTemplate};
@ -57,13 +57,22 @@ fn ocr_route_with(settings: OcrSettings) -> OcrRoute {
}
async fn perform(request: LiteLLMOcrRequest) -> Result<LiteLLMOcrResponse, Error> {
ocr_route().execute(request, &()).await
ocr_route().execute(request, &(), None).await
}
async fn perform_with(host: LocalOcrHost) -> Result<LiteLLMOcrResponse, Error> {
litellm_host::in_process::run_hosted(ocr_route().machine(host.request()?), host.runtime())
.await
.map(completed)
let result = litellm_host_native::in_process::run_hosted(
ocr_route().machine(host.request()?, None),
host.runtime(),
)
.await
.map(completed);
if let Some(observer) = &host.observer {
for event in host.events.0.lock().unwrap().iter() {
observer(event);
}
}
result
}
fn wire(model: &str, base: &str, document: Value, options: Value) -> OcrWireRequest {
@ -155,6 +164,7 @@ struct LocalOcrHost {
request: Mutex<Option<LiteLLMOcrRequest<OcrDocumentInput>>>,
before_provider_request: Option<BeforeSend>,
observer: Option<Observer>,
events: support::CallEvents,
}
impl LocalOcrHost {
@ -163,6 +173,7 @@ impl LocalOcrHost {
request: Mutex::new(Some(request)),
before_provider_request: None,
observer: None,
events: support::CallEvents::default(),
}
}
@ -199,16 +210,16 @@ impl LocalOcrHost {
})
.ok_or_else(|| Error::InvalidRequest("OCR request was already projected".into()))
}
pub fn runtime(&self) -> litellm_host::in_process::Host<'_, Self, Self, ()> {
litellm_host::in_process::Host {
pub fn runtime(&self) -> litellm_host_native::in_process::Host<'_, Self, Self, ()> {
litellm_host_native::in_process::Host {
services: self,
hooks: self,
interceptors: self,
stream: &(),
observer: Some(self),
observers: Some(&self.events.0.sender),
}
}
}
impl litellm_host::services::HostCallHandler<Ocr> for LocalOcrHost {
impl litellm_host_native::services::HostCallHandler<Ocr> for LocalOcrHost {
async fn handle_host_call(&self, op: OcrOp) -> Result<(), Error> {
match op {
OcrOp::AcquireAzureAdToken(_) => {
@ -221,13 +232,11 @@ impl litellm_host::services::HostCallHandler<Ocr> for LocalOcrHost {
}
impl litellm_host::lifecycle::CallObserver for LocalOcrHost {
fn observe(&self, event: litellm_host::event::CallEvent) {
if let Some(observer) = &self.observer {
observer(&event);
}
fn observe(&self, event: litellm_host::lifecycle::CallEvent) {
self.events.0.sender.emit(event);
}
}
impl litellm_host::hooks::RouteHooks<<Ocr as litellm_host::protocol::Protocol>::Error>
impl litellm_host::interceptors::Interceptors<<Ocr as litellm_host::protocol::Protocol>::Error>
for LocalOcrHost
{
async fn before_provider_request(
@ -240,13 +249,15 @@ impl litellm_host::hooks::RouteHooks<<Ocr as litellm_host::protocol::Protocol>::
None => Ok(wire),
}
}
async fn on_event(
async fn after_provider_response(
&self,
event: litellm_host::event::MachineEvent,
raw: litellm_host::interceptors::RawResponse,
) -> Result<(), <Ocr as litellm_host::protocol::Protocol>::Error> {
litellm_host::lifecycle::CallObserver::observe(
self,
litellm_host::event::CallEvent::Machine(event),
litellm_host::lifecycle::CallEvent::Execution(
litellm_host::lifecycle::ExecutionEvent::ProviderResponseReceived { raw },
),
);
Ok(())
}

View file

@ -172,7 +172,7 @@ async fn missing_credentials_come_from_the_injected_secret_source(
})
.unwrap();
route.execute(request, &()).await.unwrap();
route.execute(request, &(), None).await.unwrap();
assert_eq!(source.requested(), MistralOcrConfig.secret_names());
assert_eq!(
@ -201,6 +201,7 @@ async fn the_client_uses_the_injected_http_pool_configuration() {
.execute(
ocr_request("mistral/model", &upstream.uri(), json!({})),
&(),
None,
)
.await
.unwrap();

View file

@ -1,6 +1,7 @@
use litellm_host::lifecycle::ExecutionEvent;
use std::sync::{Arc, Mutex};
use litellm_host::event::{CallEvent, MachineEvent, WireRequest};
use litellm_host::{interceptors::WireRequest, lifecycle::CallEvent};
use rstest::rstest;
use super::*;
@ -147,7 +148,7 @@ async fn response_received_fires_once_for_the_parse_response() {
let recorder = observed.clone();
let host = LocalOcrHost::new(ocr_request("reducto/parse-v3", &upstream.uri(), json!({})))
.with_observer(move |event| {
if let CallEvent::Machine(MachineEvent::ResponseReceived { raw }) = event {
if let CallEvent::Execution(ExecutionEvent::ProviderResponseReceived { raw }) = event {
recorder.lock().unwrap().push(raw.body.clone());
}
});

View file

@ -55,6 +55,7 @@ async fn configured_project_and_location_apply_when_the_call_sets_neither() {
.execute(
ocr_request("vertex_ai/mistral-ocr-maas", &upstream.uri(), json!({})),
&(),
None,
)
.await
.unwrap();

View file

@ -123,7 +123,7 @@ async fn auth_survives_per_call_clients_without_freezing_settings_or_secrets(
input_sources: Default::default(),
timeout_seconds: Some(5.0),
}).unwrap();
let result = route.execute(request, &()).await.unwrap();
let result = route.execute(request, &(), None).await.unwrap();
assert!(!result.pages.is_empty());
}
let requests = upstream.received_requests().await.unwrap();

View file

@ -1,3 +1,7 @@
use litellm_host::{
interceptors::{ExecutionFacts, ResultSource},
lifecycle::ExecutionEvent,
};
use std::sync::Arc;
use futures_util::TryStreamExt;
@ -5,7 +9,7 @@ use litellm_core::responses::{
route::Responses,
types::{ResponsesCall, ResponsesOutput},
};
use litellm_host::{call::HostedCompletion, event::CallEvent};
use litellm_host::{call::HostedCompletion, lifecycle::CallEvent};
use rstest::{fixture, rstest};
use serde_json::json;
use wiremock::ResponseTemplate;
@ -39,8 +43,9 @@ async fn http_responses_share_execution_and_hooks(call: ResponsesCall, #[case] h
..call
});
let response = if hosted {
let HostedCompletion::Complete(response) = litellm_host::in_process::run_hosted(
responses_route(no_secrets()).machine(host.request().unwrap()),
let HostedCompletion::Complete(response) = litellm_host_native::in_process::run_hosted(
responses_route(no_secrets())
.machine(host.request().unwrap(), Some(host.events.0.sender.clone())),
host.runtime(),
)
.await
@ -51,7 +56,7 @@ async fn http_responses_share_execution_and_hooks(call: ResponsesCall, #[case] h
} else {
let call = host.request.lock().unwrap().take().unwrap();
let ResponsesOutput::Complete(response) = responses_route(no_secrets())
.execute(call, &host)
.execute(call, &host, Some(host.events.0.sender.clone()))
.await
.unwrap()
else {
@ -69,7 +74,13 @@ async fn http_responses_share_execution_and_hooks(call: ResponsesCall, #[case] h
&host.events.0.lock().unwrap()[..],
[
CallEvent::Started { .. },
CallEvent::Machine(_),
CallEvent::Execution(ExecutionEvent::ProviderResponseReceived { .. }),
CallEvent::Execution(ExecutionEvent::ResultReady {
facts: ExecutionFacts {
source: ResultSource::Provider,
..
}
}),
CallEvent::Succeeded { .. }
]
));
@ -95,8 +106,9 @@ async fn streaming_keeps_headers_and_bytes_and_finishes_after_consumption(
});
let (headers, bytes) = if hosted {
assert_eq!(
litellm_host::in_process::run_hosted(
responses_route(no_secrets()).machine(host.request().unwrap()),
litellm_host_native::in_process::run_hosted(
responses_route(no_secrets())
.machine(host.request().unwrap(), Some(host.events.0.sender.clone())),
host.runtime(),
)
.await
@ -110,13 +122,24 @@ async fn streaming_keeps_headers_and_bytes_and_finishes_after_consumption(
} else {
let call = host.request.lock().unwrap().take().unwrap();
let ResponsesOutput::Stream { head, chunks } = responses_route(no_secrets())
.execute(call, &host)
.execute(call, &host, Some(host.events.0.sender.clone()))
.await
.unwrap()
else {
panic!()
};
assert_eq!(host.events.0.lock().unwrap().len(), 1);
assert!(matches!(
&host.events.0.lock().unwrap()[..],
[
CallEvent::Started { .. },
CallEvent::Execution(ExecutionEvent::ResultReady {
facts: ExecutionFacts {
source: ResultSource::Provider,
..
}
}),
]
));
(
head.headers,
chunks.try_collect::<Vec<_>>().await.unwrap().concat(),
@ -126,7 +149,16 @@ async fn streaming_keeps_headers_and_bytes_and_finishes_after_consumption(
assert_eq!(bytes, body.as_bytes());
assert!(matches!(
&host.events.0.lock().unwrap()[..],
[CallEvent::Started { .. }, CallEvent::Succeeded { .. }]
[
CallEvent::Started { .. },
CallEvent::Execution(ExecutionEvent::ResultReady {
facts: ExecutionFacts {
source: ResultSource::Provider,
..
}
}),
CallEvent::Succeeded { .. }
]
));
}
@ -147,7 +179,7 @@ async fn provider_failures_emit_failure_once(
let call = host.request.lock().unwrap().take().unwrap();
assert!(
responses_route(no_secrets())
.execute(call, &host)
.execute(call, &host, Some(host.events.0.sender.clone()))
.await
.is_err()
);
@ -190,7 +222,7 @@ async fn credentials_and_endpoint_are_resolved_only_when_needed(
..call
};
responses_route(secrets.clone())
.execute(call, &())
.execute(call, &(), None)
.await
.unwrap();
assert_eq!(
@ -224,7 +256,7 @@ async fn unsupported_providers_fail_before_secrets_or_transport(
};
assert!(
responses_route(secrets.clone())
.execute(call, &())
.execute(call, &(), None)
.await
.is_err()
);
@ -257,15 +289,17 @@ async fn route_tracing_covers_native_and_hosted_outcomes(
.logger()
.instrument(async {
if hosted {
litellm_host::in_process::run_hosted(
route.clone().machine(host.request().unwrap()),
litellm_host_native::in_process::run_hosted(
route
.clone()
.machine(host.request().unwrap(), Some(host.events.0.sender.clone())),
host.runtime(),
)
.await
.map(|_| ())
} else {
route
.execute(host.request().unwrap(), &())
.execute(host.request().unwrap(), &(), None)
.await
.map(|_| ())
}
@ -311,6 +345,7 @@ async fn stream_trace_survives_handoff_and_closes_before_the_stream_object_is_dr
..call
},
&(),
None,
)
.await
})
@ -350,6 +385,7 @@ async fn preparation_failure_is_traced_but_unpolled_builders_are_not(
..call
},
&(),
None,
))
});
assert!(traces.records().is_empty());
@ -363,6 +399,7 @@ async fn preparation_failure_is_traced_but_unpolled_builders_are_not(
..self::call()
},
&(),
None,
)
.await
})

View file

@ -3,7 +3,10 @@
#![allow(dead_code)] // each test binary compiles this module on its own and uses a different subset
use std::sync::{Arc, Mutex};
use std::{
ops::ControlFlow,
sync::{Arc, Mutex},
};
use futures_util::future::BoxFuture;
use litellm_http::{
@ -248,11 +251,43 @@ pub struct RecordingCall<P: litellm_host::protocol::Protocol> {
}
#[derive(Default)]
pub struct CallEvents(pub Mutex<Vec<litellm_host::event::CallEvent>>);
pub struct CallEvents(pub Observations);
pub struct Observations {
pub sender: litellm_host::observation::ObservationSender,
receiver: Mutex<tokio::sync::mpsc::Receiver<litellm_host::lifecycle::CallEvent>>,
recorded: Mutex<Vec<litellm_host::lifecycle::CallEvent>>,
}
impl Default for Observations {
fn default() -> Self {
let (sender, receiver) = litellm_host::observation::observation_channel(
std::num::NonZeroUsize::new(128).unwrap(),
);
Self {
sender,
receiver: Mutex::new(receiver),
recorded: Mutex::new(Vec::new()),
}
}
}
impl Observations {
pub fn lock(
&self,
) -> std::sync::LockResult<std::sync::MutexGuard<'_, Vec<litellm_host::lifecycle::CallEvent>>>
{
let mut events = self.recorded.lock()?;
let mut receiver = self.receiver.lock().unwrap();
while let Ok(event) = receiver.try_recv() {
events.push(event);
}
Ok(events)
}
}
impl litellm_host::lifecycle::CallObserver for CallEvents {
fn observe(&self, event: litellm_host::event::CallEvent) {
self.0.lock().unwrap().push(event);
fn observe(&self, event: litellm_host::lifecycle::CallEvent) {
self.0.sender.emit(event);
}
}
@ -267,19 +302,15 @@ impl<P: litellm_host::protocol::Protocol> RecordingCall<P> {
}
}
impl<P: litellm_host::protocol::Protocol> litellm_host::hooks::RouteHooks<P::Error>
impl<P: litellm_host::protocol::Protocol> litellm_host::interceptors::Interceptors<P::Error>
for RecordingCall<P>
{
fn observer(&self) -> Option<Arc<dyn litellm_host::lifecycle::CallObserver>> {
Some(self.events.clone())
}
async fn before_provider_request(
&self,
wire: litellm_host::event::WireRequest,
_: litellm_host::event::RequestContext,
) -> Result<litellm_host::event::WireRequest, P::Error> {
Ok(litellm_host::event::WireRequest {
wire: litellm_host::interceptors::WireRequest,
_: litellm_host::interceptors::RequestContext,
) -> Result<litellm_host::interceptors::WireRequest, P::Error> {
Ok(litellm_host::interceptors::WireRequest {
headers: wire
.headers
.into_iter()
@ -289,12 +320,10 @@ impl<P: litellm_host::protocol::Protocol> litellm_host::hooks::RouteHooks<P::Err
})
}
async fn on_event(&self, event: litellm_host::event::MachineEvent) -> Result<(), P::Error> {
self.events
.0
.lock()
.unwrap()
.push(litellm_host::event::CallEvent::Machine(event));
async fn after_provider_response(
&self,
_: litellm_host::interceptors::RawResponse,
) -> Result<(), P::Error> {
Ok(())
}
}
@ -311,34 +340,28 @@ where
.take()
.ok_or_else(|| litellm_host::machine::MachineFault::Abandoned.into())
}
pub fn runtime(&self) -> litellm_host::in_process::Host<'_, (), Self, Self> {
litellm_host::in_process::Host {
pub fn runtime(&self) -> litellm_host_native::in_process::Host<'_, (), Self, Self> {
litellm_host_native::in_process::Host {
services: &(),
hooks: self,
interceptors: self,
stream: self,
observer: Some(self),
observers: Some(&self.events.0.sender),
}
}
}
impl<P> litellm_host::in_process::StreamConsumer<P> for RecordingCall<P>
impl<P> litellm_host_native::in_process::StreamConsumer<P> for RecordingCall<P>
where
P: litellm_host::protocol::Protocol<HostCall = std::convert::Infallible>,
P::Error: From<litellm_host::machine::MachineFault>,
{
async fn open_stream(
&self,
head: P::StreamHead,
) -> Result<litellm_host::protocol::Demand, P::Error> {
async fn open_stream(&self, head: P::StreamHead) -> Result<ControlFlow<()>, P::Error> {
*self.head.lock().unwrap() = Some(head);
Ok(litellm_host::protocol::Demand::More)
Ok(ControlFlow::Continue(()))
}
async fn send_chunk(
&self,
chunk: P::Chunk,
) -> Result<litellm_host::protocol::Demand, P::Error> {
async fn send_chunk(&self, chunk: P::Chunk) -> Result<ControlFlow<()>, P::Error> {
self.chunks.lock().unwrap().push(chunk);
Ok(litellm_host::protocol::Demand::More)
Ok(ControlFlow::Continue(()))
}
}
impl<P> litellm_host::lifecycle::CallObserver for RecordingCall<P>
@ -346,8 +369,8 @@ where
P: litellm_host::protocol::Protocol<HostCall = std::convert::Infallible>,
P::Error: From<litellm_host::machine::MachineFault>,
{
fn observe(&self, event: litellm_host::event::CallEvent) {
self.events.0.lock().unwrap().push(event.clone());
fn observe(&self, event: litellm_host::lifecycle::CallEvent) {
self.events.0.sender.emit(event);
}
}

View file

@ -203,6 +203,85 @@ fn threshold_tiers_and_boundaries() {
assert_eq!(calculate(&specification, &flex).unwrap().input(), 600.0);
}
#[rstest]
#[case::ultrafast_above_threshold(ServiceTier::Ultrafast, 300_000, 9_301_000.0, 37_000.0)]
#[case::ultrafast_at_threshold(ServiceTier::Ultrafast, 272_000, 544_500.0, 5_000.0)]
#[case::standard_above_threshold(ServiceTier::Standard, 300_000, 3_300_600.0, 13_000.0)]
#[case::priority_above_threshold(ServiceTier::Priority, 300_000, 5_701_000.0, 23_000.0)]
fn tiered_long_context_rates_are_selected_by_service_tier(
#[case] service_tier: ServiceTier,
#[case] prompt_tokens: u64,
#[case] expected_input: f64,
#[case] expected_output: f64,
) {
let standard = Rates {
cache_read: Rate::Value(3.0),
..rates(Rate::Value(1.0), Rate::Value(2.0))
};
let tiers = [
TierRates {
tier: ServiceTier::Priority,
rates: Rates {
cache_read: Rate::Value(5.0),
..rates(Rate::Value(3.0), Rate::Value(4.0))
},
},
TierRates {
tier: ServiceTier::Ultrafast,
rates: Rates {
cache_read: Rate::Value(7.0),
..rates(Rate::Value(2.0), Rate::Value(5.0))
},
},
];
let threshold_tiers = [
TierRates {
tier: ServiceTier::Priority,
rates: Rates {
cache_read: Rate::Value(29.0),
..rates(Rate::Value(19.0), Rate::Value(23.0))
},
},
TierRates {
tier: ServiceTier::Ultrafast,
rates: Rates {
cache_read: Rate::Value(41.0),
..rates(Rate::Value(31.0), Rate::Value(37.0))
},
},
];
let thresholds = [ThresholdRates {
above_prompt_tokens: 272_000,
standard: Rates {
cache_read: Rate::Value(17.0),
..rates(Rate::Value(11.0), Rate::Value(13.0))
},
tiers: &threshold_tiers,
}];
let pricing = Pricing {
standard,
tiers: &tiers,
thresholds: &thresholds,
off_peak: None,
};
let base = request();
let long_context_request = Request {
usage: Usage {
prompt_tokens,
completion_tokens: 1_000,
cache_read_tokens: 100,
cache_write_tokens: 0,
..base.usage
},
service_tier,
..base
};
let cost = calculate(&pricing, &long_context_request).unwrap();
assert_eq!(cost.input(), expected_input);
assert_eq!(cost.output(), expected_output);
}
#[test]
fn compile_rejects_ambiguous_rates() {
let duplicate = ThresholdRates {

View file

@ -6,6 +6,7 @@ license.workspace = true
repository.workspace = true
[dependencies]
litellm-cache-response.workspace = true
axum = { workspace = true, features = ["json", "multipart", "original-uri"] }
base64.workspace = true
bytes.workspace = true
@ -13,15 +14,18 @@ litellm-auth.workspace = true
litellm-gateway-auth.workspace = true
litellm-core.workspace = true
litellm-host-http.workspace = true
litellm-host.workspace = true
litellm-http.workspace = true
litellm-llms.workspace = true
litellm-router.workspace = true
litellm-secrets.workspace = true
litellm-types.workspace = true
litellm-llms-types.workspace = true
serde.workspace = true
serde_json.workspace = true
thiserror.workspace = true
[dev-dependencies]
litellm-cache-memory.workspace = true
futures-util.workspace = true
tokio = { workspace = true, features = ["io-util"] }
rstest.workspace = true

View file

@ -0,0 +1,111 @@
use std::time::Duration;
use litellm_cache_response::{CacheOptions, CachePolicy, CacheScope};
use litellm_gateway_auth::AuthenticatedRequest;
use serde::Deserialize;
use serde_json::{Map, Value};
use crate::Error;
#[derive(Default, Deserialize)]
#[serde(default, deny_unknown_fields)]
struct Controls {
#[serde(rename = "no-cache")]
no_cache: bool,
#[serde(rename = "no-store")]
no_store: bool,
ttl: Option<f64>,
#[serde(rename = "s-maxage", alias = "s-max-age")]
max_age: Option<f64>,
}
type Prepared = (Map<String, Value>, CacheOptions);
pub(crate) fn prepare(
identity: &AuthenticatedRequest,
body: Map<String, Value>,
) -> Result<Prepared, Error> {
let controls: Controls = match body.get("cache").filter(|value| !value.is_null()) {
Some(value) => serde_json::from_value(value.clone())
.map_err(|error| Error::InvalidBody(error.to_string()))?,
None => Controls::default(),
};
let caching: Option<bool> = body
.get("caching")
.filter(|value| !value.is_null())
.map(|value| serde_json::from_value(value.clone()))
.transpose()
.map_err(|error| Error::InvalidBody(error.to_string()))?;
let caller = identity.caller();
let options = CacheOptions {
policy: CachePolicy {
caching,
no_cache: controls.no_cache,
no_store: controls.no_store,
ttl: controls.ttl.map(duration).transpose()?,
max_age: controls.max_age.map(duration).transpose()?,
},
scope: CacheScope::Isolated(
serde_json::json!([
caller.principal().authority(),
caller.principal().subject(),
caller.authentication().credential_id
])
.to_string(),
),
};
Ok((
body.into_iter()
.filter(|(name, _)| !matches!(name.as_str(), "cache" | "caching"))
.collect(),
options,
))
}
fn duration(seconds: f64) -> Result<Duration, Error> {
Duration::try_from_secs_f64(seconds)
.ok()
.filter(|duration| !duration.is_zero())
.ok_or_else(|| Error::InvalidBody("cache durations must be finite and positive".into()))
}
#[derive(Clone, Default)]
pub(crate) struct CacheHeaders(std::sync::Arc<std::sync::OnceLock<String>>);
impl litellm_host::interceptors::Interceptors<litellm_core::RouteError> for CacheHeaders {
async fn before_provider_request(
&self,
wire: litellm_host::interceptors::WireRequest,
_: litellm_host::interceptors::RequestContext,
) -> Result<litellm_host::interceptors::WireRequest, litellm_core::RouteError> {
Ok(wire)
}
async fn after_provider_response(
&self,
_: litellm_host::interceptors::RawResponse,
) -> Result<(), litellm_core::RouteError> {
Ok(())
}
async fn result_ready(
&self,
facts: litellm_host::interceptors::ExecutionFacts,
) -> Result<(), litellm_core::RouteError> {
if let litellm_host::interceptors::ResultSource::Cache { key } = facts.source {
let _ = self.0.set(key);
}
Ok(())
}
}
impl CacheHeaders {
pub(crate) fn apply(&self, mut response: axum::response::Response) -> axum::response::Response {
if let Some(key) = self.0.get()
&& let Ok(value) = axum::http::HeaderValue::from_str(key)
{
response.headers_mut().insert("x-litellm-cache-key", value);
}
response
}
}

View file

@ -42,12 +42,21 @@ async fn handle(
) -> Result<Response, Error> {
let deployment = request::resolve_deployment(gateway, &body)?;
request::authorize_model(identity, deployment, &body).await?;
let (body, cache_options) = crate::caching::prepare(identity, body)?;
let route = gateway.chat_completions.clone();
let route = match &gateway.cache {
Some(cache) => route.with_cache(litellm_cache_response::ScopedCache::new(
cache.clone(),
cache_options.scope.clone(),
)),
None => route,
};
let messages = body.get("messages").cloned().unwrap_or_default();
let headers = crate::caching::CacheHeaders::default();
let response = litellm_host_http::serve_unary(
gateway
.chat_completions
.clone()
.machine(ChatCompletionsCall {
route.machine(
ChatCompletionsCall {
model: deployment.model.clone(),
messages,
optional_params: body
@ -59,11 +68,14 @@ async fn handle(
custom_llm_provider: deployment.custom_llm_provider.clone(),
extra_headers: None,
timeout: deployment.timeout,
}),
(),
},
cache_options.policy,
),
(),
headers.clone(),
litellm_host_http::Unary::new(Json),
None,
)
.await?;
Ok(response)
Ok(headers.apply(response))
}

View file

@ -4,6 +4,7 @@
//! maps a public model name to its deployment and runs the core route.
mod audio_transcription;
mod caching;
mod chat_completions;
mod error;
pub mod messages;
@ -27,6 +28,7 @@ pub use litellm_router::{Deployment, Router as ModelRouter};
pub use request::{JsonObject, RequestId};
pub struct Gateway {
cache: Option<Arc<dyn litellm_cache_response::ResponseCacheService>>,
pub audio_transcription: AudioTranscriptionRoute,
pub chat_completions: ChatCompletionsRoute,
pub messages: MessagesRoute,
@ -39,6 +41,13 @@ pub struct Gateway {
}
impl Gateway {
pub fn with_cache(self, cache: Arc<dyn litellm_cache_response::ResponseCacheService>) -> Self {
Self {
cache: Some(cache),
..self
}
}
pub fn new(
resources: CoreResources,
http: HttpClientConfig,
@ -48,6 +57,7 @@ impl Gateway {
let provider = resources.pool.client(&http, ClientVariant::Provider)?;
let auth = resources.auth.clone();
Ok(Self {
cache: None,
audio_transcription: AudioTranscriptionRoute::new(
provider.clone(),
auth.clone(),

View file

@ -12,7 +12,7 @@ use axum::{
};
use litellm_core::messages::{MessagesCall, messages_body, route::Messages};
use litellm_host_http::Sse;
use litellm_types::utils::{ProviderSpecificHeader, ProviderSpecificHeaders};
use litellm_llms_types::headers::{ProviderSpecificHeader, ProviderSpecificHeaders};
use serde_json::{Map, Value};
use crate::{Deployment, Error, Gateway, JsonObject, RequestId, request};
@ -43,11 +43,23 @@ async fn handle(
) -> Result<Response, Error> {
let deployment = request::resolve_deployment(gateway, &body)?;
request::authorize_model(identity, deployment, &body).await?;
let (body, cache_options) = crate::caching::prepare(identity, body)?;
let route = gateway.messages.clone();
let route = match &gateway.cache {
Some(cache) => route.with_cache(litellm_cache_response::ScopedCache::new(
cache.clone(),
cache_options.scope.clone(),
)),
None => route,
};
let call = project(deployment, body, headers)?;
let machine = gateway.messages.clone().machine(call);
let machine = route.machine(call, cache_options.policy);
let stream =
Sse::<Messages, _, _>::new(Json, |error| Bytes::from(Error::from(error).sse_frame()));
Ok(litellm_host_http::serve(machine, (), (), stream).await?)
let headers = crate::caching::CacheHeaders::default();
let response = litellm_host_http::serve(machine, (), headers.clone(), stream, None).await?;
Ok(headers.apply(response))
}
fn project(

View file

@ -4,7 +4,8 @@ use std::sync::Arc;
use axum::{Json, extract::State, http::HeaderMap, response::IntoResponse};
use litellm_auth::SecretValue;
use litellm_core::ocr::types::{LiteLLMOcrRequest, OcrConnectionInputs, OcrDocumentInput};
use litellm_llms::base_llm::ocr::transformation::OcrDocument;
use litellm_llms::base_llm::ocr::transformation::decode_request_value;
use litellm_llms_types::formats::ocr::OcrDocument;
use serde_json::Value;
use crate::{
@ -42,7 +43,11 @@ async fn handle(
file_name: upload.file_name,
mime_type: upload.mime_type,
},
None => OcrDocument::try_from(body.get("document").cloned().unwrap_or_default())?.into(),
None => decode_request_value::<OcrDocument>(
body.get("document").cloned().unwrap_or_default(),
"document",
)?
.into(),
};
let format = body
.get("req_format")
@ -70,7 +75,7 @@ async fn handle(
..Default::default()
},
)?;
let response = gateway.ocr.execute(call, &()).await?;
let response = gateway.ocr.execute(call, &(), None).await?;
match response.provider_native_response {
Some(native) => Ok(Value::Object(native)),
None => Ok(response.into_json()),

View file

@ -15,6 +15,16 @@ pub(crate) async fn create(
) -> Result<Response, Error> {
let deployment = request::resolve_deployment(&gateway, &body)?;
request::authorize_model(&identity, deployment, &body).await?;
let (body, cache_options) = crate::caching::prepare(&identity, body)?;
let route = gateway.responses.clone();
let route = match &gateway.cache {
Some(cache) => route.with_cache(litellm_cache_response::ScopedCache::new(
cache.clone(),
cache_options.scope.clone(),
)),
None => route,
};
let call = ResponsesCall {
model: deployment.model.clone(),
input: body.get("input").cloned().unwrap_or_default(),
@ -28,7 +38,7 @@ pub(crate) async fn create(
extra_headers: None,
timeout: deployment.timeout,
};
let machine = gateway.responses.clone().machine(call);
let machine = route.machine(call, cache_options.policy);
let stream = Sse::<Responses, _, _>::new(Json, |error| {
let error = Error::from(error);
Bytes::from(format!(
@ -36,5 +46,7 @@ pub(crate) async fn create(
json!({"type": "error", "code": error.status().as_u16().to_string(), "message": error.to_string(), "param": null})
))
});
Ok(litellm_host_http::serve(machine, (), (), stream).await?)
let headers = crate::caching::CacheHeaders::default();
let response = litellm_host_http::serve(machine, (), headers.clone(), stream, None).await?;
Ok(headers.apply(response))
}

View file

@ -0,0 +1,185 @@
mod support;
use std::{sync::Arc, time::Duration};
use axum::body::to_bytes;
use litellm_cache_memory::InMemoryCache;
use litellm_cache_response::{
CacheKeyInput, ResponseCache, ResponseCacheConfig, ResponseCacheRequest, ResponseCacheService,
};
use rstest::rstest;
use serde_json::{Value, json};
use wiremock::{Mock, MockServer, ResponseTemplate, matchers::method};
#[rstest]
#[case::chat("/v1/chat/completions", "anthropic/test-model", false)]
#[case::messages("/v1/messages", "anthropic/test-model", false)]
#[case::responses("/v1/responses", "openai/test-model", false)]
#[case::messages_stream("/v1/messages", "anthropic/test-model", true)]
#[case::responses_stream("/v1/responses", "openai/test-model", true)]
#[tokio::test]
async fn all_inference_endpoints_share_native_cache(
#[case] path: &str,
#[case] model: &str,
#[case] stream: bool,
#[values("s-maxage", "s-max-age")] max_age: &str,
) {
let upstream = MockServer::start().await;
let is_responses = path.ends_with("responses");
let provider_body = if is_responses {
json!({"id":"response-1", "model":"test-model", "status":"completed", "output":[]})
} else {
json!({"id":"message-1", "model":"test-model", "type":"message", "role":"assistant", "content":[{"type":"text","text":"hello"}], "stop_reason":"end_turn", "usage":{"input_tokens":1,"output_tokens":1}})
};
let terminal = if is_responses {
"response.completed"
} else {
"message_stop"
};
let events = format!("event: {terminal}\ndata: {{\"type\":\"{terminal}\"}}\n\n");
let template = if stream {
ResponseTemplate::new(200).set_body_raw(events.clone(), "text/event-stream")
} else {
ResponseTemplate::new(200).set_body_json(provider_body)
};
Mock::given(method("POST"))
.respond_with(template)
.expect(2)
.mount(&upstream)
.await;
let cache: Arc<dyn ResponseCacheService> = Arc::new(
ResponseCache::new(Arc::new(InMemoryCache::new(
Some(100),
Some(Duration::from_secs(60)),
)))
.with_config(ResponseCacheConfig {
namespace: "gateway-test".into(),
max_entry_bytes: 4096,
}),
);
let app = support::app_with_cache(model, &upstream.uri(), cache.clone());
let request = if is_responses {
json!({"model":"public/model", "input":"hello", "stream":stream, "cache":{(max_age):600}})
} else {
json!({"model":"public/model", "messages":[{"role":"user","content":"hello"}], "max_tokens":16, "stream":stream, "cache":{(max_age):600}})
};
let first = support::post(app.clone(), path, request.clone()).await;
assert_eq!(first.status(), 200);
assert!(!first.headers().contains_key("x-litellm-cache-key"));
let first = to_bytes(first.into_body(), 4096).await.unwrap();
let second = support::post(app.clone(), path, request.clone()).await;
assert_eq!(second.status(), 200);
let cache_key = second.headers().get("x-litellm-cache-key").unwrap().clone();
assert!(!cache_key.as_bytes().is_empty());
let stored = cache
.lookup(
&ResponseCacheRequest::new(CacheKeyInput {
preset: Some(cache_key.to_str().unwrap().into()),
..Default::default()
}),
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap(),
)
.await
.unwrap();
assert!(
stored.is_some(),
"the header must identify the stored entry"
);
let second = to_bytes(second.into_body(), 4096).await.unwrap();
if stream {
assert_eq!(first, events);
assert_eq!(second, first);
} else {
assert_eq!(
serde_json::from_slice::<Value>(&first).unwrap(),
serde_json::from_slice::<Value>(&second).unwrap()
);
}
let bypass_request = Value::Object(
request
.as_object()
.unwrap()
.iter()
.map(|(name, value)| {
(
name.clone(),
if name == "cache" {
json!({"no-cache": true, "no-store": true})
} else {
value.clone()
},
)
})
.collect(),
);
let bypassed = support::post(app.clone(), path, bypass_request).await;
assert_eq!(bypassed.status(), 200);
assert!(!bypassed.headers().contains_key("x-litellm-cache-key"));
to_bytes(bypassed.into_body(), 4096).await.unwrap();
let restored = support::post(app, path, request).await;
assert_eq!(restored.status(), 200);
assert_eq!(
restored.headers().get("x-litellm-cache-key"),
Some(&cache_key)
);
assert_eq!(to_bytes(restored.into_body(), 4096).await.unwrap(), second);
}
#[rstest]
#[case::different_subject("issuer", "tenant-b")]
#[case::different_authority("other-issuer", "tenant-a")]
#[tokio::test]
async fn authenticated_callers_do_not_share_cached_responses(
#[case] authority: &str,
#[case] subject: &str,
) {
use litellm_gateway_auth::{Principal, PrincipalKind};
let upstream = MockServer::start().await;
Mock::given(method("POST"))
.respond_with(ResponseTemplate::new(200).set_body_json(json!({
"id":"response-1", "model":"test-model", "status":"completed", "output":[]
})))
.expect(2)
.mount(&upstream)
.await;
let cache: Arc<dyn ResponseCacheService> = Arc::new(ResponseCache::new(Arc::new(
InMemoryCache::new(Some(100), Some(Duration::from_secs(60))),
)));
let first_caller = support::app_with_cache_for_principal(
"openai/test-model",
&upstream.uri(),
cache.clone(),
Principal::new("issuer".into(), "tenant-a".into(), PrincipalKind::Service),
);
let second_caller = support::app_with_cache_for_principal(
"openai/test-model",
&upstream.uri(),
cache,
Principal::new(authority.into(), subject.into(), PrincipalKind::Service),
);
let body = json!({"model":"public/model","input":"same prompt"});
let first = support::post(first_caller.clone(), "/v1/responses", body.clone()).await;
assert_eq!(first.status(), 200);
assert!(!first.headers().contains_key("x-litellm-cache-key"));
let first_hit = support::post(first_caller.clone(), "/v1/responses", body.clone()).await;
assert_eq!(first_hit.status(), 200);
let first_key = first_hit.headers().get("x-litellm-cache-key").unwrap();
let second = support::post(second_caller.clone(), "/v1/responses", body.clone()).await;
assert_eq!(second.status(), 200);
assert!(!second.headers().contains_key("x-litellm-cache-key"));
let second_hit = support::post(second_caller, "/v1/responses", body.clone()).await;
assert_eq!(second_hit.status(), 200);
assert_ne!(
second_hit.headers().get("x-litellm-cache-key").unwrap(),
first_key
);
let first_again = support::post(first_caller, "/v1/responses", body).await;
assert_eq!(first_again.status(), 200);
assert_eq!(
first_again.headers().get("x-litellm-cache-key"),
Some(first_key)
);
}

View file

@ -2,7 +2,8 @@ mod support;
use axum::{body::Body, http::Request};
use litellm_gateway_inference::Error;
use litellm_llms::base_llm::ocr::{error::Error as OcrError, transformation::OcrDocument};
use litellm_llms::base_llm::ocr::{error::Error as OcrError, transformation::decode_request_value};
use litellm_llms_types::formats::ocr::OcrDocument;
use rstest::rstest;
use serde_json::{Value, json};
use tower::ServiceExt;
@ -146,7 +147,7 @@ async fn malformed_multipart_uses_an_openai_error_envelope(
#[rstest]
#[case::missing_document(
"/v1/ocr", "mistral/test-ocr", "",
Error::Ocr(OcrDocument::try_from(Value::Null).unwrap_err()),
Error::Ocr(decode_request_value::<OcrDocument>(Value::Null, "document").unwrap_err()),
)]
#[case::empty_document(
"/v1/ocr",

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