scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-29 06:33:28 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 200

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

Browse source 
# Conflicts:
#	litellm/integrations/otel/mappers/langfuse.py
#	litellm/types/integrations/langfuse_otel.py
This commit is contained in:
yucheng 2026-10-20 16:28:00 -07:00
commit b474b2eabc
1566 changed files with 109902 additions and 18574 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

90
.githooks/pre-push
View file

@ -1,90 +0,0 @@
#!/usr/bin/env bash
#
# pre-push — enforce Conventional Branches
# https://conventional-branch.github.io/
#
# Branch format: <type>/<description>
# <type> must be one of: feature, bugfix, hotfix, release, chore
#
# Protected branches (always allowed):
# - main
# - dependabot/*
# - gh-readonly-queue/*
#
# Tag pushes and branch deletions are skipped.
# Bypass: git push --no-verify.
set -eu
ZERO_OID="0000000000000000000000000000000000000000"
ZERO_OID_SHA256="0000000000000000000000000000000000000000000000000000000000000000"
ALLOWED_TYPES="feature|bugfix|hotfix|release|chore"
BRANCH_PATTERN="^(${ALLOWED_TYPES})/.+"
PROTECTED_NAMES="main"
PROTECTED_PREFIXES="dependabot/ gh-readonly-queue/"
is_protected() {
branch="$1"
for name in $PROTECTED_NAMES; do
if [ "$branch" = "$name" ]; then
return 0
fi
done
for prefix in $PROTECTED_PREFIXES; do
case "$branch" in "$prefix"*) return 0 ;; esac
done
return 1
}
invalid=""
while read -r local_ref local_oid remote_ref remote_oid; do
# Branch deletion (no local commit being pushed).
if [ "$local_oid" = "$ZERO_OID" ] || [ "$local_oid" = "$ZERO_OID_SHA256" ]; then
continue
fi
# Only validate branch pushes; ignore tags and other ref namespaces.
case "$remote_ref" in
refs/heads/*) ;;
*) continue ;;
esac
branch="${remote_ref#refs/heads/}"
if is_protected "$branch"; then
continue
fi
if ! printf '%s' "$branch" | grep -Eq "$BRANCH_PATTERN"; then
invalid="$invalid $branch"
fi
done
if [ -n "$invalid" ]; then
cat >&2 <<EOF
✗ Branch name does not follow Conventional Branches.
Invalid:$invalid
Expected: <type>/<description>
Allowed types: feature, bugfix, hotfix, release, chore
Examples:
feature/weighted-round-robin
bugfix/streaming-empty-chunks
chore/bump-deps
hotfix/auth-bypass
Protected (always allowed): main, dependabot/*, gh-readonly-queue/*.
See https://conventional-branch.github.io/
Rename with: git branch -m <new-name>
To bypass (use sparingly): git push --no-verify
EOF
exit 1
fi
exit 0

2
.github/CODEOWNERS vendored
View file

@ -1,2 +1,4 @@
/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
/helm/ @yassin-berriai
/terraform/ @yassin-berriai

9
.github/PULL_REQUEST_TEMPLATE/general.md vendored
View file

@ -4,6 +4,13 @@
Drop every section you have nothing to put in, heading included: a bare "## Relevant issues" or
"## Affected release" with nothing under it must not appear in the final description -->
<!-- Two or three plain sentences a teammate can read in 10 seconds, written the way you'd explain it in Slack
Say what the PR does for users and why, then anything that keeps it working over time (a sync job, a config someone owns)
No headings, bullets, tables, or code names here. Everything else goes inside the Details block below -->
<details>
<summary>Details</summary>
## TLDR
<!-- Fill in the bullets below and keep each one short and concrete: one line per bullet, roughly 10 words max
@ -164,3 +171,5 @@ Example checklists:
## Final Attestation
- [ ] The tests check the right things, including the edge cases, and regressions in the respective real-world customer use-cases are not possible after this PR
</details>

2
.github/ci-coverage-allowlist.yml vendored
View file

@ -5,7 +5,7 @@ description: >-
test_paths:
- reason: >-
litellm.agent() end-to-end suite. It drives the real claude, codex and opencode CLIs and
litellm.agent() end-to-end suite. It drives the real claude, codex, opencode and pi CLIs and
deepagents against a live LiteLLM AI Gateway, so it needs those binaries on PATH plus
LITELLM_PROXY_API_BASE / LITELLM_PROXY_API_KEY, and skips without them. Run manually
before changing litellm/harness; the mocked coverage runs in tests/unit/harness and

11
.github/workflows/test-unit.yml vendored
View file

@ -139,6 +139,7 @@ jobs:
test-path: |-
tests/unit/batch_completion
tests/unit/decisions
tests/unit/extraction
tests/unit/litellm_core_utils
tests/unit/systemone
python-version: "3.12"
@ -266,6 +267,7 @@ jobs:
tests/unit/llms/featherless_ai
tests/unit/llms/firecrawl
tests/unit/llms/fireworks_ai
tests/unit/llms/friendliai
tests/unit/llms/gdc
tests/unit/llms/gemini
tests/unit/llms/gigachat
@ -321,6 +323,7 @@ jobs:
tests/unit/llms/pass_through
tests/unit/llms/perplexity
tests/unit/llms/pg_vector
tests/unit/llms/pi
tests/unit/llms/publicai
tests/unit/llms/ragflow
tests/unit/llms/recraft
@ -410,6 +413,7 @@ jobs:
tests/unit/embeddings
tests/unit/endpoints
tests/unit/files
tests/unit/fine_tuning
tests/unit/images
tests/unit/interactions
tests/unit/messages
@ -430,6 +434,13 @@ jobs:
- shard: rust-bridge-harness
test-path: >-
tests/test_litellm_rust/test_dispatch.py
tests/test_litellm_rust/test_public_entries.py
tests/test_litellm_rust/chat_completions/test_dispatch.py
tests/test_litellm_rust/embeddings/test_dispatch.py
tests/test_litellm_rust/messages/test_dispatch.py
tests/test_litellm_rust/ocr/test_dispatch.py
tests/test_litellm_rust/responses/test_dispatch.py
tests/unit/compression
tests/unit/harness
tests/unit/rust_bridge

10
CONTRIBUTING.md
View file

@ -39,12 +39,12 @@ git clone https://github.com/YOUR_USERNAME/litellm.git
cd litellm
# Create a new branch for your feature (see "Commit and Branch Conventions" below)
git checkout -b feature/your-feature
git checkout -b litellm_your_feature
# Install development dependencies
make install-dev
# Install git hooks that enforce commit + branch conventions (one-time, opt-in)
# Install git hooks that enforce commit conventions (one-time, opt-in)
make install-hooks
# Verify your setup works
@ -55,7 +55,7 @@ That's it! Your local development environment is ready.
## Commit and Branch Conventions
Commits follow [Conventional Commits](https://www.conventionalcommits.org/en/v1.0.0/) and branches follow [Conventional Branches](https://conventional-branch.github.io/). Run `make install-hooks` once per clone to enable the local git hooks that enforce these — see the [contributor docs](https://docs.litellm.ai/docs/extras/contributing_code#commit-and-branch-conventions) for the full type list, examples, the protected-branch bypass list, and how to opt out.
Commits follow [Conventional Commits](https://www.conventionalcommits.org/en/v1.0.0/). Internal contributors name branches with the `litellm_` prefix and no `/`, and fork branches can use any name. Run `make install-hooks` once per clone to enable the local hook that enforces the commit convention. See the [contributor docs](https://docs.litellm.ai/docs/extras/contributing_code#commit-and-branch-conventions) for details
### 2. Development Workflow
@ -78,8 +78,8 @@ uv run pytest tests/unit/<your_test_file>.py -v
git add .
git commit -m "feat(scope): your descriptive commit message"
# Push and create a PR (branch must follow Conventional Branches — see above)
git push origin feature/your-feature
# Push and create a PR
git push origin litellm_your_feature
```
## Adding Testing

4
Makefile
View file

@ -22,7 +22,7 @@ help:
@echo " make install-proxy-dev-ci - Install proxy dev dependencies (CI-compatible)"
@echo " make install-test-deps - Install the full local test environment"
@echo " make install-helm-unittest - Install helm unittest plugin"
@echo " make install-hooks - Install git hooks (Conventional Commits + Branches)"
@echo " make install-hooks - Install git hooks (Conventional Commits)"
@echo " make check - Run CI-equivalent lint on staged files, or on the diff vs the base branch when nothing is staged"
@echo " make pre-commit - Legacy alias for make check"
@echo " make format - Apply ruff format code formatting"
@ -120,7 +120,7 @@ install-helm-unittest:
helm plugin install https://github.com/helm-unittest/helm-unittest --version v0.8.2; \
}
# Install git hooks that enforce Conventional Commits and Conventional Branches.
# Install git hooks that enforce Conventional Commits.
# Opt-in: not chained into install-dev.
install-hooks:
./scripts/install_git_hooks.sh

6
README.md
View file

@ -287,7 +287,7 @@ For MCP OAuth, an upstream may advertise dynamic client registration but refuse
</details>
<details>
<summary><b>Agents</b> - Run Claude Code, Codex, OpenCode or Deep Agents on any model (Python SDK)</summary>
<summary><b>Agents</b> - Run Claude Code, Codex, OpenCode, pi or Deep Agents on any model (Python SDK)</summary>
### Python SDK - Agents
@ -296,7 +296,7 @@ import litellm
from litellm import Harness, sandbox
result = litellm.agent(
Harness.CLAUDE_CODE, # or Harness.CODEX, Harness.OPENCODE, Harness.DEEPAGENTS
Harness.CLAUDE_CODE, # or Harness.CODEX, Harness.OPENCODE, Harness.PI, Harness.DEEPAGENTS
"Find why tests/test_router.py is flaky and fix it.",
sandbox=sandbox.local("./repo"),
model="litellm_proxy/claude-sonnet-4-5", # a model group on your AI Gateway
@ -305,7 +305,7 @@ result = litellm.agent(
print(result.text, result.cost, [f.path for f in result.files])
```
Set `LITELLM_PROXY_API_BASE` and `LITELLM_PROXY_API_KEY` and every model call the agent makes goes through your AI Gateway, tagged `harness,claude_code`. Drop the `litellm_proxy/` prefix to call a provider directly. Install `starlette uvicorn` plus the agent's CLI (`claude`, `codex` or `opencode`), or `deepagents langchain-litellm` for Deep Agents.
Set `LITELLM_PROXY_API_BASE` and `LITELLM_PROXY_API_KEY` and every model call the agent makes goes through your AI Gateway, tagged `harness,claude_code`. Drop the `litellm_proxy/` prefix to call a provider directly. Install `starlette uvicorn` plus the agent's CLI (`claude`, `codex`, `opencode` or `pi`), or `deepagents langchain-litellm` for Deep Agents.
[**Docs: Agent Harnesses**](https://docs.litellm.ai/docs/harness)

41
cookbook/litellm_proxy_server/batch_api/bedrock/README.md Normal file
View file

@ -0,0 +1,41 @@
# Bedrock batch endpoints
Set `aws_bedrock_control_plane_endpoint` on a deployment when batch create, retrieve, and cancel requests must use a custom Bedrock control-plane endpoint. The proxy signs these requests for the configured URL. Without this setting, batch requests use the standard regional Bedrock endpoint
```yaml
model_list:
- model_name: bedrock-batch
litellm_params:
model: bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0
aws_region_name: us-east-1
aws_access_key_id: os.environ/AWS_ACCESS_KEY_ID
aws_secret_access_key: os.environ/AWS_SECRET_ACCESS_KEY
aws_batch_role_arn: os.environ/AWS_BATCH_ROLE_ARN
s3_output_bucket_name: os.environ/AWS_S3_OUTPUT_BUCKET_NAME
aws_bedrock_control_plane_endpoint: os.environ/BEDROCK_CONTROL_PLANE_ENDPOINT
```
Set `BEDROCK_CONTROL_PLANE_ENDPOINT` to your Bedrock control-plane VPC endpoint or compatible gateway base URL, without `/model-invocation-job`. The proxy adds the operation path. The endpoint must be reachable from the proxy, and the configured AWS region must match the endpoint's signing region
`api_base` and `aws_bedrock_runtime_endpoint` configure runtime inference separately. They do not override the batch control-plane endpoint. `s3_endpoint_url` configures S3 file access separately
The control-plane setting belongs in deployment configuration. Sending `aws_bedrock_control_plane_endpoint` in an API request body is rejected. There is no automatic `AWS_BEDROCK_CONTROL_PLANE_ENDPOINT` environment fallback, so reference the desired environment variable explicitly as shown above
Start the proxy with this configuration and set `LITELLM_API_BASE` to its URL. Set `BEDROCK_BATCH_INPUT_URI` to a valid input JSONL file in S3 that the batch service role can read
```bash
curl -sS "$LITELLM_API_BASE/v1/batches" \
-H "Authorization: Bearer $LITELLM_MASTER_KEY" \
-H 'Content-Type: application/json' \
-d "{\"model\":\"bedrock-batch\",\"input_file_id\":\"$BEDROCK_BATCH_INPUT_URI\",\"endpoint\":\"/v1/chat/completions\",\"completion_window\":\"24h\"}"
```
Use the returned `id` as `BATCH_ID` to retrieve or cancel the same job through the deployment's configured endpoint
```bash
curl -sS "$LITELLM_API_BASE/v1/batches/$BATCH_ID?model=bedrock-batch" \
-H "Authorization: Bearer $LITELLM_MASTER_KEY"
curl -sS -X POST "$LITELLM_API_BASE/v1/batches/$BATCH_ID/cancel?model=bedrock-batch" \
-H "Authorization: Bearer $LITELLM_MASTER_KEY"
```

16
enterprise/litellm_enterprise/integrations/custom_guardrail.py
View file

@ -20,9 +20,15 @@ class EnterpriseCustomGuardrailHelper:
the tag's value (the mode for that tag).
- If no request tag matches, fall back to default mode(s).
"""
from litellm.litellm_core_utils.core_helpers import (
get_metadata_variable_name_from_kwargs,
)
from litellm.litellm_core_utils.litellm_logging import (
StandardLoggingPayloadSetup,
)
from litellm.router_strategy.tag_based_routing import (
without_baggage_mapped_tags,
)
from litellm.proxy._types import CommonProxyErrors
from litellm.proxy.proxy_server import premium_user
@ -36,9 +42,13 @@ class EnterpriseCustomGuardrailHelper:
proxy_server_request = data.get("proxy_server_request", {})
request_tags = StandardLoggingPayloadSetup.get_request_tags(
litellm_params=data,
proxy_server_request=proxy_server_request,
request_tags = without_baggage_mapped_tags(
StandardLoggingPayloadSetup.get_request_tags(
litellm_params=data,
proxy_server_request=proxy_server_request,
),
data.get(get_metadata_variable_name_from_kwargs(data)),
keep_serving_deployment_tags=True,
)
# Check if any request tag matches a Mode tag key

197
enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py
View file

@ -4,6 +4,7 @@ Polls LiteLLM_ManagedObjectTable to check if the batch job is complete, and if t
import asyncio
import json
import re
from collections.abc import Mapping, Sequence
from dataclasses import replace as dataclasses_replace
from datetime import datetime, timedelta, timezone
@ -13,6 +14,7 @@ from typing import TYPE_CHECKING, Final, List, Literal, Optional, Protocol, Tupl
from litellm_enterprise.proxy.common_utils.managed_object_attribution import (
ManagedObjectCreatorAttribution,
ManagedObjectRow,
row_metadata,
)
import litellm
@ -39,7 +41,12 @@ if TYPE_CHECKING:
from litellm.types.utils import LiteLLMBatch
def _error_names_resource(error: Exception, resource_id: str) -> bool:
return resource_id in re.findall(r"""[^\s"'`]+""", str(error))
CHECK_BATCH_COST_USER_AGENT = "LiteLLM Proxy/CheckBatchCost"
RECENT_BATCH_POLL_WINDOW: Final = timedelta(hours=48)
PROVIDER_TERMINAL_BATCH_STATUSES: Final[Tuple[str, ...]] = (
"completed",
@ -84,11 +91,6 @@ def _team_table(prisma_client: "PrismaClient") -> "TableActions[prisma_models.Li
return TeamRepository(prisma_client).table
def _row_metadata(row: object) -> dict[str, object]:
metadata: Final = getattr(row, "metadata", None)
return cast(dict[str, object], metadata) if isinstance(metadata, dict) else {} # cast-ok: prisma Json object column
class CheckBatchCost:
def __init__(
self,
@ -109,6 +111,11 @@ class CheckBatchCost:
# the guaranteed-failing primary query on every subsequent cycle.
self._has_batch_processed_column: bool = True
self.batch_processed_support_confirmed: bool = False
self._poll_cursor: tuple[datetime, str] | None = None
self._poll_before: datetime = datetime.now(timezone.utc)
self._arrival_cursor: tuple[datetime, str] | None = None
self._arrival_before: datetime = self._poll_before
self._poll_arrivals: bool = False
self._line_item_stores: Final[set[asyncio.Task[None]]] = set()
self._line_item_store_turn: Final = asyncio.Lock()
self._poll_idle: Final = asyncio.Event()
@ -161,29 +168,9 @@ class CheckBatchCost:
f"(older than {MANAGED_OBJECT_STALENESS_CUTOFF_DAYS} days) as stale_expired"
)
if not self._has_batch_processed_column:
return
# A row already in a terminal status is never rewritten by the sweep above, so
# without this it keeps a poll-page slot forever and starves newer batches.
retired: Final = await _managed_object_table(self.prisma_client).update_many(
where={
"file_purpose": "batch",
"batch_processed": False,
"status": {"in": ["complete", "completed"]},
"created_at": {"lt": cutoff},
},
data={"batch_processed": True},
)
if retired > 0:
verbose_proxy_logger.warning(
f"CheckBatchCost: gave up on {retired} completed managed objects older than "
f"{MANAGED_OBJECT_STALENESS_CUTOFF_DAYS} days that were never costed"
)
async def _fallback_find_jobs(self) -> "Sequence[ManagedObjectRow]":
"""Query batch jobs without the batch_processed filter (for older schemas)."""
return await _managed_object_table(self.prisma_client).find_many(
return await self._find_jobs(
where={
"file_purpose": "batch",
"status": {
@ -197,10 +184,59 @@ class CheckBatchCost:
]
},
},
take=MAX_OBJECTS_PER_POLL_CYCLE,
order={"created_at": "asc"},
)
async def _find_jobs(self, where: Mapping[str, object]) -> Sequence[ManagedObjectRow]:
arrivals: Final = self._poll_arrivals
page: Final = await self._find_job_page(where, arrivals=arrivals)
self._poll_arrivals = not arrivals if page else arrivals
if page:
return page
return await self._find_job_page(where, arrivals=not arrivals)
async def _find_job_page(self, where: Mapping[str, object], *, arrivals: bool) -> Sequence[ManagedObjectRow]:
cursor: Final = self._arrival_cursor if arrivals else self._poll_cursor
if cursor is None:
if arrivals:
self._arrival_before = datetime.now(timezone.utc)
else:
self._poll_before = datetime.now(timezone.utc) - RECENT_BATCH_POLL_WINDOW
before: Final = self._arrival_before if arrivals else self._poll_before
page: Final = await _managed_object_table(self.prisma_client).find_many(
where={
**where,
"created_at": {
"lte": before,
**({"gt": before - RECENT_BATCH_POLL_WINDOW} if arrivals else {}),
},
**(
{
"OR": [
{"created_at": {"gt": cursor[0]}},
{"created_at": cursor[0], "id": {"gt": cursor[1]}},
]
}
if cursor is not None
else {}
),
},
take=MAX_OBJECTS_PER_POLL_CYCLE,
order=[{"created_at": "asc"}, {"id": "asc"}],
)
next_cursor: Final = (page[-1].created_at, page[-1].id) if len(page) == MAX_OBJECTS_PER_POLL_CYCLE else None
if arrivals:
self._arrival_cursor = next_cursor
else:
self._poll_cursor = next_cursor
return page
def _unprocessed_snapshot(self, job: ManagedObjectRow) -> Mapping[str, object]:
return {
"id": job.id,
"updated_at": job.updated_at,
**({"batch_processed": False} if self._has_batch_processed_column else {}),
}
async def _retire_job(self, job: "ManagedObjectRow", reason: str) -> None:
"""
Take a row that can never be costed out of the poll page. Leaving it selectable
@ -208,24 +244,19 @@ class CheckBatchCost:
once enough such rows accumulate no newer batch is ever reached. Older schemas
without batch_processed can only be excluded through the status filter.
"""
data: Final = (
{"batch_processed": True}
if self._has_batch_processed_column
else {"status": "stale_expired"}
)
data: Final = {"batch_processed": True} if self._has_batch_processed_column else {"status": "stale_expired"}
try:
await _managed_object_table(self.prisma_client).update(
where={"id": job.id},
retired: Final = await _managed_object_table(self.prisma_client).update_many(
where=self._unprocessed_snapshot(job),
data=data,
)
except Exception as db_err:
verbose_proxy_logger.error(
f"CheckBatchCost: failed to retire uncostable job {job.id} ({reason}): {db_err}"
)
verbose_proxy_logger.error(f"CheckBatchCost: failed to retire uncostable job {job.id} ({reason}): {db_err}")
return
if retired == 0:
return
verbose_proxy_logger.warning(
f"CheckBatchCost: job {job.id} can never be costed ({reason}), "
"so it will no longer be polled"
f"CheckBatchCost: job {job.id} can never be costed ({reason}), so it will no longer be polled"
)
async def _claim_job_for_costing(self, job: "ManagedObjectRow", response: "LiteLLMBatch") -> bool:
@ -312,13 +343,13 @@ class CheckBatchCost:
from litellm.exceptions import NotFoundError
return isinstance(error, (NotFoundError, openai.NotFoundError)) and batch_id in str(error)
return isinstance(error, (NotFoundError, openai.NotFoundError)) and _error_names_resource(error, batch_id)
def _batch_deployment_exists(self, model_id: str) -> bool:
"""A 404 only proves the batch is gone when it came from the batch's own
deployment. Once that deployment leaves the router, default fallbacks can
silently send the retrieve to a provider that never saw the batch, so its
404 must not retire the row; the staleness sweep bounds it instead."""
404 must not retire a potentially recoverable row."""
return self.llm_router.get_deployment(model_id=model_id) is not None
@staticmethod
@ -333,7 +364,7 @@ class CheckBatchCost:
if not output_file_id:
return False
return isinstance(error, (NotFoundError, openai.NotFoundError)) and output_file_id in str(error)
return isinstance(error, (NotFoundError, openai.NotFoundError)) and _error_names_resource(error, output_file_id)
async def _finalize_unbilled_terminal_job(
self, job: "ManagedObjectRow", response: "LiteLLMBatch"
@ -355,18 +386,18 @@ class CheckBatchCost:
db_batch_object=job,
unified_batch_id=_is_base64_encoded_unified_file_id(job.unified_object_id),
)
update_data: Final[dict] = {
update_data: Final[dict[str, object]] = {
"status": response.status,
"file_object": response.model_dump_json(),
**({"batch_processed": True} if self._has_batch_processed_column else {}),
}
await _managed_object_table(self.prisma_client).update(
where={"id": job.id},
finalized: Final = await _managed_object_table(self.prisma_client).update_many(
where=self._unprocessed_snapshot(job),
data=update_data,
)
verbose_proxy_logger.info(
f"CheckBatchCost: marked job {job.id} as {response.status} in DB"
)
if finalized == 0:
return False
verbose_proxy_logger.info(f"CheckBatchCost: marked job {job.id} as {response.status} in DB")
return True
except Exception as db_err:
verbose_proxy_logger.error(
@ -619,7 +650,11 @@ class CheckBatchCost:
from litellm.files.types import FileContentCallOptions
from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
from litellm.litellm_core_utils.litellm_logging import deployment_pricing_model_info, mask_api_base_credentials
from litellm.litellm_core_utils.litellm_logging import (
deployment_pricing_model_info,
mask_api_base_credentials,
resolve_vertex_location_for_cost,
)
from litellm.proxy.openai_files_endpoints.common_utils import (
_is_base64_encoded_unified_file_id,
)
@ -748,6 +783,12 @@ class CheckBatchCost:
custom_llm_provider=batch_file_provider,
model_name=model_name,
model_info=deployment_model_info,
vertex_location=resolve_vertex_location_for_cost(
custom_llm_provider=llm_provider,
litellm_params=deployment_info.litellm_params.model_dump(exclude_none=True),
optional_params=None,
model=model_name,
),
)
error_file_failed_requests: Final = await count_error_file_failed_requests(
response,
@ -807,20 +848,26 @@ class CheckBatchCost:
)
return None
try:
await logging_obj.async_success_handler(
result=response,
batch_cost=batch_result.cost,
batch_usage=batch_result.usage,
batch_models=batch_result.models,
batch_successful_requests=batch_result.successful_requests,
batch_failed_requests=batch_result.failed_requests,
batch_prompt_cost=batch_result.prompt_cost,
batch_completion_cost=batch_result.completion_cost,
)
except Exception:
await self._release_job_claim(job, original_status, original_file_object)
raise
from litellm.proxy.hooks.proxy_track_cost_callback import track_batch_accounting
with track_batch_accounting(logging_obj.litellm_call_id) as accounting:
try:
await logging_obj.async_success_handler(
result=response,
batch_cost=batch_result.cost,
batch_usage=batch_result.usage,
batch_models=batch_result.models,
batch_successful_requests=batch_result.successful_requests,
batch_failed_requests=batch_result.failed_requests,
batch_prompt_cost=batch_result.prompt_cost,
batch_completion_cost=batch_result.completion_cost,
)
finally:
if accounting.retryable_failure:
await self._release_job_claim(job, original_status, original_file_object)
if accounting.retryable_failure:
return None
# Record batch duration (completed_at - created_at)
if prom_logger and response.completed_at and response.created_at:
@ -889,7 +936,10 @@ class CheckBatchCost:
)
from litellm.batches.batch_utils import fetch_batch_managed_file_content
from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider
from litellm.litellm_core_utils.litellm_logging import deployment_pricing_model_info
from litellm.litellm_core_utils.litellm_logging import (
deployment_pricing_model_info,
resolve_vertex_location_for_cost,
)
try:
deployment_info: Final = self.llm_router.get_deployment(model_id=model_id)
@ -933,6 +983,12 @@ class CheckBatchCost:
cost_provider=cost_provider,
model_info=deployment_pricing_model_info(model_id=model_id, deployment_model=litellm_model_name),
metadata=line_item_metadata(await self._creator_attribution.batch_metadata(job, batch_id)),
vertex_location=resolve_vertex_location_for_cost(
custom_llm_provider=llm_provider,
litellm_params=deployment_info.litellm_params.model_dump(exclude_none=True),
optional_params=None,
model=model_name,
),
)
input_content: Final = (
b""
@ -1003,7 +1059,9 @@ class CheckBatchCost:
if (has_key_row and key_row is None) or (team_id and team_row is None):
return True
settings: Final = get_dynamic_logging_metadata(
UserAPIKeyAuth(metadata=_row_metadata(key_row), team_metadata=_row_metadata(team_row), team_id=team_id),
UserAPIKeyAuth(
metadata=dict(row_metadata(key_row)), team_metadata=dict(row_metadata(team_row)), team_id=team_id
),
proxy_config,
)
callback_vars: Final = (settings.callback_vars if settings is not None else None) or {}
@ -1028,6 +1086,7 @@ class CheckBatchCost:
"""
try:
from litellm.integrations.prometheus import PrometheusLogger
prom_logger = PrometheusLogger.get_instance()
except Exception as e:
verbose_proxy_logger.error(f"CheckBatchCost: could not get Prometheus logger: {e}")
@ -1052,7 +1111,7 @@ class CheckBatchCost:
# endpoint may transition a batch to "complete" before
# CheckBatchCost runs. The batch_processed=False filter
# already prevents reprocessing finished batches.
jobs = await _managed_object_table(self.prisma_client).find_many(
jobs = await self._find_jobs(
where={
"file_purpose": "batch",
"batch_processed": False,
@ -1065,8 +1124,6 @@ class CheckBatchCost:
]
},
},
take=MAX_OBJECTS_PER_POLL_CYCLE,
order={"created_at": "asc"},
)
self.batch_processed_support_confirmed = True
except Exception as query_err:
@ -1133,7 +1190,7 @@ class CheckBatchCost:
continue
verbose_proxy_logger.error(
f"CheckBatchCost: failed to track cost for batch {batch_id} "
f"(job {job.id}); leaving it unprocessed so the next poll retries: {tracking_err}"
f"(job {job.id}): {tracking_err}"
)
self._record_error(prom_logger, "cost_tracking_error")
continue

117
enterprise/litellm_enterprise/proxy/common_utils/check_responses_cost.py
View file

@ -7,6 +7,12 @@ attribution, priced against the deployment that served it. The poll's status che
with BACKGROUND_RESPONSE_STATUS_POLL_CALL_ORIGIN and client reads are marked managed by the
managed files hook, so both stay free. The spend event carries the id the client holds, the row's
unified_object_id, so its spend log is found under that id with the _background_cost suffix.
The spend event reaches the creating key's or team's own logging callbacks the way that creator's
requests do. Its message content is redacted when the creator's key or team logging turns message
logging off, else as the proxy is set, the way the batch poll decides it. The row is claimed before
the creator is read, so reads that lose the race skip the lookups. A creator whose lookup fails, or
whose logging settings the logger rejects, is still billed, through the proxy's own callbacks only
and with its message content redacted.
A read that finds no row, because the create could not store one, stores it under the reading key
and bills it the same way, so a response the create left untracked is still billed once. A streamed
background create is free on its own; its stream stores the row under the creating key and bills
@ -29,6 +35,7 @@ from typing import TYPE_CHECKING, Final, cast
from litellm_enterprise.proxy.common_utils.managed_object_attribution import (
ManagedObjectCreatorAttribution,
ManagedObjectRow,
ResponseCreator,
)
from pydantic import BaseModel, ConfigDict, TypeAdapter
@ -54,12 +61,29 @@ from litellm.types.utils import (
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
from litellm.proxy.proxy_server import ProxyConfig
from litellm.proxy.utils import PrismaClient, ProxyLogging
from litellm.repositories.prisma_protocols import TableActions
from litellm.router import Router
_STORED_RESPONSE_ADAPTER: Final[TypeAdapter[dict[str, object]]] = TypeAdapter(dict[str, object])
_CALLBACK_LIST_ADAPTER: Final[TypeAdapter[tuple[object, ...]]] = TypeAdapter(tuple[object, ...])
class _RoutedModelInfo(BaseModel):
model_config = ConfigDict(frozen=True)
id: str | None = None
class _RoutedDeployment(BaseModel):
model_config = ConfigDict(frozen=True)
model_info: _RoutedModelInfo = _RoutedModelInfo()
_ROUTED_DEPLOYMENTS: Final = TypeAdapter(tuple[_RoutedDeployment, ...])
_RETRIEVE_RESPONSE_CALL_TYPE: Final = "aget_responses"
_CALLBACK_LIST_FIELDS: Final = frozenset({"success_callback", "failure_callback", "litellm_disabled_callbacks"})
_REJECTED_CREATOR_LOGGING: Final[Mapping[str, object]] = MappingProxyType({"turn_off_message_logging": True})
@dataclass(frozen=True, slots=True)
@ -76,6 +100,7 @@ class BilledBackgroundResponse:
custom_llm_provider: str | None
litellm_params: Mapping[str, object]
metadata: Mapping[str, object]
creator_logging: Mapping[str, object]
@dataclass(frozen=True, slots=True)
@ -94,11 +119,51 @@ def _without_read_price(response: ResponsesAPIResponse) -> ResponsesAPIResponse:
return unpriced
async def log_billed_background_response(billed: BilledBackgroundResponse) -> BilledCharge:
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
def _string_callbacks(callbacks: object) -> tuple[str, ...]:
if not isinstance(callbacks, list):
return ()
return tuple(
callback for callback in _CALLBACK_LIST_ADAPTER.validate_python(callbacks) if isinstance(callback, str)
)
logging_obj: Final = LiteLLMLogging(
def _running_proxy_config() -> "ProxyConfig":
from litellm.proxy.proxy_server import proxy_config
return proxy_config
def creator_logging_fields(creator: ResponseCreator, proxy_config: "ProxyConfig | None") -> Mapping[str, object]:
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.litellm_pre_call_utils import key_and_team_logging_fields
if not creator.settings_read:
return _REJECTED_CREATOR_LOGGING
try:
return key_and_team_logging_fields(
user_api_key_dict=UserAPIKeyAuth(
metadata=dict(creator.key_metadata),
team_id=creator.team_id,
team_metadata=dict(creator.team_metadata),
),
proxy_config=proxy_config if proxy_config is not None else _running_proxy_config(),
)
except Exception as e:
verbose_proxy_logger.error(
f"CheckResponsesCost: could not read the creator's logging settings, so only the proxy's "
f"callbacks get the billed event: {e}"
)
return _REJECTED_CREATOR_LOGGING
def _billed_logging_obj(billed: BilledBackgroundResponse, creator_logging: Mapping[str, object]) -> "LiteLLMLogging":
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
from litellm.utils import split_dynamic_success_callbacks
sync_callbacks, async_callbacks = split_dynamic_success_callbacks(
_string_callbacks(creator_logging.get("success_callback"))
)
return LiteLLMLogging(
model=billed.model,
messages=[{"role": "user", "content": "<retrieve_response>"}],
stream=False,
@ -106,7 +171,27 @@ async def log_billed_background_response(billed: BilledBackgroundResponse) -> Bi
start_time=datetime.now(),
litellm_call_id=str(uuid.uuid4()),
function_id=str(uuid.uuid4()),
dynamic_success_callbacks=list(sync_callbacks),
dynamic_async_success_callbacks=list(async_callbacks) or None,
kwargs={key: value for key, value in creator_logging.items() if key not in _CALLBACK_LIST_FIELDS},
)
def _creator_or_proxy_logging_obj(billed: BilledBackgroundResponse) -> "LiteLLMLogging":
try:
return _billed_logging_obj(billed, billed.creator_logging)
except Exception as e:
verbose_proxy_logger.error(
f"CheckResponsesCost: the creator's logging settings were rejected, so only the proxy's "
f"callbacks get the billed event: {e}"
)
return _billed_logging_obj(billed, _REJECTED_CREATOR_LOGGING)
async def log_billed_background_response(billed: BilledBackgroundResponse) -> BilledCharge:
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
logging_obj: Final = _creator_or_proxy_logging_obj(billed)
logging_obj.update_environment_variables(
litellm_params={**billed.litellm_params, "metadata": dict(billed.metadata)},
optional_params={},
@ -174,6 +259,7 @@ class CheckResponsesCost:
prisma_client: "PrismaClient",
llm_router: "Router",
log_billed_response: BilledResponseLogger = log_billed_background_response,
proxy_config: "ProxyConfig | None" = None,
):
from litellm.proxy.utils import PrismaClient, ProxyLogging
from litellm.router import Router
@ -183,6 +269,7 @@ class CheckResponsesCost:
self.llm_router: Router = llm_router
self._creator_attribution: Final = ManagedObjectCreatorAttribution(prisma_client, "CheckResponsesCost")
self._log_billed_response: Final = log_billed_response
self._proxy_config: Final = proxy_config
def _resolve_deployment(self, response_id: str) -> Deployment | None:
model_id: str | None = ResponsesAPIRequestUtils.get_model_id_from_response_id(response_id)
@ -294,11 +381,12 @@ class CheckResponsesCost:
data={"status": job.status or PENDING_RESPONSE_STATUSES[0]},
)
async def _billed_response(
def _billed_response(
self,
job: ManagedObjectRow,
target: _PollTarget,
response: ResponsesAPIResponse,
creator: ResponseCreator,
) -> BilledBackgroundResponse:
deployment: Final = target.deployment
return BilledBackgroundResponse(
@ -310,12 +398,26 @@ class CheckResponsesCost:
litellm_params=_deployment_litellm_params(deployment) if deployment is not None else MappingProxyType({}),
metadata=MappingProxyType(
{
**(await self._creator_attribution.response_metadata(job, job.unified_object_id)),
**creator.metadata,
INTERNAL_CALL_ORIGIN_METADATA_KEY: BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN,
**target.model_metadata,
**(_deployment_metadata(deployment) if deployment is not None else {}),
**self._requested_model_group(target, job.team_id),
}
),
creator_logging=creator_logging_fields(creator, self._proxy_config),
)
def _requested_model_group(self, target: _PollTarget, team_id: str | None) -> Mapping[str, object]:
requested: Final = target.model_metadata.get("model_group")
if target.deployment is None or not isinstance(requested, str):
return MappingProxyType({})
routed: Final = _ROUTED_DEPLOYMENTS.validate_python(
self.llm_router.get_model_list(model_name=requested, team_id=team_id) or ()
)
serving_id: Final = target.deployment.model_info.id
return MappingProxyType(
{"model_group": requested} if any(row.model_info.id == serving_id for row in routed) else {}
)
async def _bill_once(
@ -328,7 +430,8 @@ class CheckResponsesCost:
verbose_proxy_logger.info(f"Response {job.unified_object_id} was already billed elsewhere, skipping")
return None
try:
charge: Final = await self._log_billed_response(await self._billed_response(job, target, response))
creator: Final = await self._creator_attribution.response_creator(job, job.unified_object_id)
charge: Final = await self._log_billed_response(self._billed_response(job, target, response, creator))
except BaseException:
await self._release_billing_claim(job)
raise

207
enterprise/litellm_enterprise/proxy/common_utils/managed_object_attribution.py
View file

@ -9,12 +9,17 @@ from __future__ import annotations
from collections.abc import Mapping
from dataclasses import dataclass
from typing import TYPE_CHECKING, Final, Protocol, cast
from datetime import datetime
from typing import TYPE_CHECKING, Final, Protocol, TypeGuard, cast
from pydantic import TypeAdapter
from litellm._logging import verbose_proxy_logger
from litellm.constants import CLI_SESSION_KEY_PREFIX
from litellm.repositories.table_repositories import (
DeletedVerificationTokenRepository,
DeprecatedVerificationTokenRepository,
)
from litellm.repositories.team_repository import TeamRepository
from litellm.repositories.user_repository import UserRepository
from litellm.repositories.verification_token_repository import VerificationTokenRepository
@ -27,6 +32,12 @@ if TYPE_CHECKING:
class ManagedObjectRow(Protocol):
@property
def created_at(self) -> datetime: ...
@property
def updated_at(self) -> datetime: ...
@property
def id(self) -> str: ...
@ -63,6 +74,7 @@ class ManagedObjectRow(Protocol):
_TAG_LIST_ADAPTER: Final[TypeAdapter[list[object]]] = TypeAdapter(list[object])
_KEY_OWNER_SCOPE_PREFIX: Final = "key:"
_MAX_ROTATION_HOPS: Final = 16
def _string_tags(request_tags: object) -> list[str] | None:
@ -89,6 +101,20 @@ def _team_table(prisma_client: PrismaClient) -> TableActions[prisma_models.LiteL
return TeamRepository(prisma_client).table
def _rotated_token_table(
prisma_client: PrismaClient,
) -> TableActions[prisma_models.LiteLLM_DeprecatedVerificationToken]:
return DeprecatedVerificationTokenRepository(prisma_client).table
def _deleted_token_table(prisma_client: PrismaClient) -> TableActions[prisma_models.LiteLLM_DeletedVerificationToken]:
return DeletedVerificationTokenRepository(prisma_client).table
def _deleted_team_table(prisma_client: PrismaClient) -> TableActions[prisma_models.LiteLLM_DeletedTeamTable]:
return TeamRepository(prisma_client).deleted_table
@dataclass(frozen=True, slots=True)
class _CreatorRows:
key: prisma_models.LiteLLM_VerificationToken | None
@ -96,6 +122,15 @@ class _CreatorRows:
user: prisma_models.LiteLLM_UserTable | None
@dataclass(frozen=True, slots=True)
class ResponseCreator:
metadata: Mapping[str, object]
key_metadata: Mapping[str, object]
team_id: str | None
team_metadata: Mapping[str, object]
settings_read: bool = True
class ManagedObjectCreatorAttribution:
"""Rebuilds the spend-tracking metadata of whoever created a managed object, so the poll
job's spend row is attributed like the creating request: key hash, team, organization,
@ -107,8 +142,15 @@ class ManagedObjectCreatorAttribution:
user's alias, because a resolvable name is more useful on the spend row than a null.
user_api_key_org_id is resolved too: the spend update writer reads it off this metadata
to increment organization spend, so leaving it out drops the cost from org accounting.
A background response also carries its end user and the key, team, and user budgets, the
way the read that billed it before did, and never names a key-scoped creator as its user.
A background response also carries its end user, the creating key's project and its alias,
and the key, team, and user budgets and per-model budgets, the way the read that billed it
before did, and never names a key-scoped creator as its user. A key that sets no budget of
its own takes its budget tier's, the same rule auth applies. Its creator's key and team
metadata come along too, so the billed event reaches the key's and team's logging callbacks
like that read did. A rotated key's metadata is read from the live key its rotations lead
to, and a deleted key's or team's from its deleted row. A lookup that fails leaves the
creator's settings unread and the attribution to the rows that still load, so the
response still bills and still counts toward their budgets.
"""
def __init__(self, prisma_client: PrismaClient, poller_name: str) -> None:
@ -119,17 +161,70 @@ class ManagedObjectCreatorAttribution:
rows: Final = await self._rows(job, object_id)
return {"user_api_key_user_id": job.created_by, **_shared_metadata(job, rows)}
async def response_metadata(self, job: ManagedObjectRow, object_id: str) -> Mapping[str, object]:
rows: Final = await self._rows(job, object_id)
end_user_id: Final = getattr(job, "end_user_id", None)
return {
"user_api_key_user_id": _creator_user_id(job.created_by),
**_shared_metadata(job, rows),
**({"user_api_key_end_user_id": end_user_id} if end_user_id else {}),
**_budget_fields(rows.key, "user_api_key_spend", "user_api_key_max_budget"),
**_budget_fields(rows.team, "user_api_key_team_spend", "user_api_key_team_max_budget"),
**_budget_fields(rows.user, "user_api_key_user_spend", "user_api_key_user_max_budget"),
}
async def response_creator(self, job: ManagedObjectRow, object_id: str) -> ResponseCreator:
try:
return await self._read_response_creator(job, object_id)
except Exception as e:
verbose_proxy_logger.error(
f"{self._poller_name}: could not read the creator of {object_id}, so it bills without their "
f"settings: {e}"
)
return ResponseCreator(
metadata=_response_creator_metadata(job, await self._rows(job, object_id)),
key_metadata={},
team_id=getattr(job, "team_id", None),
team_metadata={},
settings_read=False,
)
async def _read_response_creator(self, job: ManagedObjectRow, object_id: str) -> ResponseCreator:
api_key: Final = getattr(job, "api_key", None)
team_id: Final = getattr(job, "team_id", None)
rows: Final = _CreatorRows(
key=await self._live_key_row(api_key, job.created_by),
team=await self._live_team_row(team_id),
user=await self._user_row(object_id, _creator_user_id(job.created_by)),
)
return ResponseCreator(
metadata=_response_creator_metadata(job, rows),
key_metadata=row_metadata(rows.key) if rows.key is not None else await self._retired_key_metadata(api_key),
team_id=team_id,
team_metadata=(
row_metadata(rows.team) if rows.team is not None else await self._deleted_team_metadata(team_id)
),
)
async def _retired_key_metadata(self, api_key: str | None) -> Mapping[str, object]:
if not _names_a_key_row(api_key):
return {}
current_token: Final = await self._token_rotated_to(api_key, _MAX_ROTATION_HOPS)
current_key: Final = (
await _token_table(self._prisma_client).find_unique(where={"token": current_token})
if current_token != api_key
else None
)
if current_key is not None:
return row_metadata(current_key)
return row_metadata(
await _deleted_token_table(self._prisma_client).find_first(
where={"token": current_token}, order={"deleted_at": "desc"}
)
)
async def _token_rotated_to(self, token: str, hops_left: int) -> str:
rotation: Final = await _rotated_token_table(self._prisma_client).find_unique(where={"token": token})
if rotation is None or hops_left == 0:
return token
return await self._token_rotated_to(rotation.active_token_id, hops_left - 1)
async def _deleted_team_metadata(self, team_id: str | None) -> Mapping[str, object]:
if not team_id:
return {}
return row_metadata(
await _deleted_team_table(self._prisma_client).find_first(
where={"team_id": team_id}, order={"deleted_at": "desc"}
)
)
async def _rows(self, job: ManagedObjectRow, object_id: str) -> _CreatorRows:
return _CreatorRows(
@ -152,23 +247,57 @@ class ManagedObjectCreatorAttribution:
async def _key_row(
self, object_id: str, api_key: str | None, created_by: str | None
) -> prisma_models.LiteLLM_VerificationToken | None:
if not api_key or _is_cli_session_key(api_key, created_by):
return None
try:
return await _token_table(self._prisma_client).find_unique(where={"token": api_key})
return await self._live_key_row(api_key, created_by)
except Exception as e:
verbose_proxy_logger.error(f"{self._poller_name}: could not look up the creating key for {object_id}: {e}")
return None
async def _team_row(self, object_id: str, team_id: str | None) -> prisma_models.LiteLLM_TeamTable | None:
if not team_id:
async def _live_key_row(
self, api_key: str | None, created_by: str | None
) -> prisma_models.LiteLLM_VerificationToken | None:
if not api_key or _is_cli_session_key(api_key, created_by):
return None
return await _token_table(self._prisma_client).find_unique(
where={"token": api_key}, include={"litellm_budget_table": True, "litellm_project_table": True}
)
async def _team_row(self, object_id: str, team_id: str | None) -> prisma_models.LiteLLM_TeamTable | None:
try:
return await _team_table(self._prisma_client).find_unique(where={"team_id": team_id})
return await self._live_team_row(team_id)
except Exception as e:
verbose_proxy_logger.error(f"{self._poller_name}: could not look up team {team_id} for {object_id}: {e}")
return None
async def _live_team_row(self, team_id: str | None) -> prisma_models.LiteLLM_TeamTable | None:
if not team_id:
return None
return await _team_table(self._prisma_client).find_unique(where={"team_id": team_id})
def _response_creator_metadata(job: ManagedObjectRow, rows: _CreatorRows) -> Mapping[str, object]:
end_user_id: Final = getattr(job, "end_user_id", None)
return {
"user_api_key_user_id": _creator_user_id(job.created_by),
**_shared_metadata(job, rows),
**({"user_api_key_end_user_id": end_user_id} if end_user_id else {}),
"user_api_key_project_id": getattr(rows.key, "project_id", None),
"user_api_key_project_alias": getattr(getattr(rows.key, "litellm_project_table", None), "project_alias", None),
**_key_budget_fields(rows.key),
**_budget_fields(
rows.team,
"user_api_key_team_spend",
"user_api_key_team_max_budget",
"user_api_key_team_model_max_budget",
),
**_budget_fields(
rows.user,
"user_api_key_user_spend",
"user_api_key_user_max_budget",
"user_api_key_user_model_max_budget",
),
}
def _shared_metadata(job: ManagedObjectRow, rows: _CreatorRows) -> Mapping[str, object]:
api_key: Final = getattr(job, "api_key", None)
@ -196,6 +325,19 @@ def _shared_metadata(job: ManagedObjectRow, rows: _CreatorRows) -> Mapping[str,
}
def row_metadata(row: object) -> Mapping[str, object]:
metadata: Final = getattr(row, "metadata", None)
if not isinstance(metadata, dict):
return {}
return cast(Mapping[str, object], metadata) # cast-ok: prisma Json object column
def _names_a_key_row(api_key: object) -> TypeGuard[str]:
from litellm.litellm_core_utils.litellm_logging import is_valid_sha256_hash
return isinstance(api_key, str) and is_valid_sha256_hash(api_key)
def _is_cli_session_key(api_key: str, created_by: str | None) -> bool:
return bool(created_by) and api_key == f"{CLI_SESSION_KEY_PREFIX}-{created_by}"
@ -212,7 +354,28 @@ def _organization_id(row: object) -> str | None:
return cast("str | None", getattr(row, "organization_id", None)) if row is not None else None
def _budget_fields(row: object, spend_key: str, max_budget_key: str) -> Mapping[str, object]:
def _budget_fields(row: object, spend_key: str, max_budget_key: str, model_max_budget_key: str) -> Mapping[str, object]:
if row is None:
return {}
return {spend_key: getattr(row, "spend", None), max_budget_key: getattr(row, "max_budget", None)}
return {
spend_key: getattr(row, "spend", None),
max_budget_key: getattr(row, "max_budget", None),
model_max_budget_key: getattr(row, "model_max_budget", None),
}
def _key_budget_fields(key: object) -> Mapping[str, object]:
if key is None:
return {}
tier: Final = getattr(key, "litellm_budget_table", None)
return {
"user_api_key_spend": getattr(key, "spend", None),
"user_api_key_max_budget": _own_or_tier_budget(key, tier, "max_budget"),
"user_api_key_model_max_budget": _own_or_tier_budget(key, tier, "model_max_budget"),
}
def _own_or_tier_budget(key: object, tier: object, field: str) -> object:
own: Final = getattr(key, field, None)
tier_value: Final = getattr(tier, field, None)
return own if tier_value is None or own not in (None, {}) else tier_value

63
enterprise/litellm_enterprise/proxy/hooks/managed_files.py
View file

@ -6,6 +6,7 @@ import base64
import json
import time
from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence
from itertools import chain
from types import MappingProxyType
from typing import (
TYPE_CHECKING,
@ -479,6 +480,13 @@ def _with_managed_background_marker(request: Mapping[str, object]) -> dict[str,
_MANAGED_FILES_TARGET: Final = "managed_files"
_PROVIDER_FILE_ID_KWARGS: Final = MappingProxyType(
{
CallTypes.acreate_batch: ("input_file_id",),
CallTypes.acreate_fine_tuning_job: ("training_file", "validation_file"),
}
)
_PROVIDER_FILE_ID_KEYS: Final = tuple(dict.fromkeys(chain.from_iterable(_PROVIDER_FILE_ID_KWARGS.values())))
_PROVIDER_FILE_RETRIEVE_RETRY_DELAYS_SECONDS: Final = (0.5, 1.0, 2.0)
@ -1562,6 +1570,10 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
if isinstance(file_id, str) and file_id
)
await self.check_file_ids_access(fine_tuning_file_ids, user_api_key_dict)
if fine_tuning_file_ids:
data["model_file_id_mapping"] = await self.get_model_file_id_mapping(
fine_tuning_file_ids, user_api_key_dict.parent_otel_span
)
return data
@ -1576,24 +1588,27 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
if request_kwargs is None:
return healthy_deployments
input_file_id = cast(Optional[str], request_kwargs.get("input_file_id"))
model_file_id_mapping = cast(
Optional[Dict[str, Dict[str, str]]],
request_kwargs.get("model_file_id_mapping"),
request: Final = cast(Mapping[str, object], request_kwargs)
model_file_id_mapping: Final = cast(Mapping[str, Mapping[str, str]], request.get("model_file_id_mapping") or {})
requested_file_ids: Final = (request.get(key) for key in _PROVIDER_FILE_ID_KEYS)
model_ids_per_file: Final = tuple(
frozenset(model_file_id_mapping[file_id])
for file_id in requested_file_ids
if isinstance(file_id, str) and file_id in model_file_id_mapping
)
allowed_model_ids: Final[frozenset[str]] = (
model_ids_per_file[0].intersection(*model_ids_per_file[1:]) if model_ids_per_file else frozenset()
)
allowed_model_ids = []
if input_file_id and model_file_id_mapping:
model_id_dict = model_file_id_mapping.get(input_file_id, {})
allowed_model_ids = list(model_id_dict.keys())
if len(allowed_model_ids) == 0:
return healthy_deployments
return [
holders: Final[list[dict[str, object]]] = [
deployment
for deployment in healthy_deployments
if deployment.get("model_info", {}).get("id") in allowed_model_ids
]
return holders if holders or "input_file_id" in request else healthy_deployments
async def async_pre_call_deployment_hook(
self, kwargs: Dict[str, Any], call_type: Optional[CallTypes]
@ -1601,25 +1616,21 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
"""
Allow modifying the request just before it's sent to the deployment.
"""
accessor_key: Optional[str] = None
if call_type and call_type == CallTypes.acreate_batch:
accessor_key = "input_file_id"
elif call_type and call_type == CallTypes.acreate_fine_tuning_job:
accessor_key = "training_file"
else:
accessor_keys: Final = _PROVIDER_FILE_ID_KWARGS.get(call_type) if call_type else None
if accessor_keys is None:
return kwargs
if accessor_key:
model_file_id_mapping = cast(Optional[Dict[str, Dict[str, str]]], kwargs.get("model_file_id_mapping"))
# model_info may be at top-level or nested under litellm_metadata
# (batch/file operations use litellm_metadata)
model_id = cast(Optional[str], kwargs.get("model_info", {}).get("id", None))
if model_id is None:
model_id = cast(
Optional[str],
kwargs.get("litellm_metadata", {}).get("model_info", {}).get("id", None),
)
for accessor_key in accessor_keys:
input_file_id = cast(Optional[str], kwargs.get(accessor_key))
model_file_id_mapping = cast(Optional[Dict[str, Dict[str, str]]], kwargs.get("model_file_id_mapping"))
# model_info may be at top-level or nested under litellm_metadata
# (batch/file operations use litellm_metadata)
model_id = cast(Optional[str], kwargs.get("model_info", {}).get("id", None))
if model_id is None:
model_id = cast(
Optional[str],
kwargs.get("litellm_metadata", {}).get("model_info", {}).get("id", None),
)
mapped_file_id: Optional[str] = None
if input_file_id and model_file_id_mapping and model_id:
mapped_file_id = model_file_id_mapping.get(input_file_id, {}).get(model_id, None)
@ -1813,7 +1824,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
),
)
async def get_model_file_id_mapping(self, file_ids: List[str], litellm_parent_otel_span: Span) -> dict:
async def get_model_file_id_mapping(self, file_ids: Sequence[str], litellm_parent_otel_span: Span) -> dict:
"""
Get model-specific file IDs for a list of proxy file IDs.
Returns a dictionary mapping litellm_proxy/ file_id -> model_id -> model_file_id

4
enterprise/pyproject.toml
View file

@ -1,6 +1,6 @@
[project]
name = "litellm-enterprise"
version = "0.1.78"
version = "0.1.79"
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.78"
version = "0.1.79"
version_files = [
"pyproject.toml:^version",
"../pyproject.toml:litellm-enterprise==",

4
gateway/routes/allowlist.py
View file

@ -1,7 +1,7 @@
"""Path allowlist for the gateway component.
The gateway exposes the LLM data-plane surface: chat/completions, embeddings,
audio, batches, files, fine-tuning, rerank, decisions, ocr, rag, video, search, image,
audio, batches, files, fine-tuning, rerank, decisions, extraction, ocr, rag, video, search, image,
responses, vector stores, passthrough providers, realtime websockets, MCP
tool-call endpoints, and operational endpoints (/health, /metrics, and the
/debug/memory/summary read of the serving worker's RSS).
@ -64,6 +64,8 @@ GATEWAY_PATH_PREFIXES: tuple[str, ...] = (
"/decisions",
"/v1/systemone",
"/systemone",
"/v1/extraction",
"/extraction",
"/v1/ocr",
"/ocr",
"/v1/rag/",

13
helm/litellm-helm/README.md
View file

@ -70,11 +70,18 @@ proxyConfigMap:
proxy_config:
general_settings:
master_key: os.environ/PROXY_MASTER_KEY
model_list: []
```
The chart ships no models, so add them under `proxy_config.model_list` or from the Admin UI with `store_model_in_db: true`
```
proxy_config:
model_list:
- model_name: gpt-3.5-turbo
- model_name: gpt-5.6
litellm_params:
model: gpt-3.5-turbo
api_key: eXaMpLeOnLy
model: openai/gpt-5.6
api_key: os.environ/OPENAI_API_KEY
```
#### Example using existing `proxyConfigMap` instead of creating it:

26
helm/litellm-helm/tests/deployment_tests.yaml
View file

@ -270,6 +270,32 @@ tests:
- matchRegex:
path: data["config.yaml"]
pattern: '(?m)^\s*health_check_details:\s*false$'
- it: should not render any placeholder models when the user sets proxy_config without model_list
template: configmap-litellm.yaml
set:
proxy_config.general_settings.store_model_in_db: true
asserts:
- notMatchRegex:
path: data["config.yaml"]
pattern: 'model_name:'
- matchRegex:
path: data["config.yaml"]
pattern: 'store_model_in_db: true'
- matchRegex:
path: data["config.yaml"]
pattern: 'master_key: os.environ/PROXY_MASTER_KEY'
- it: should render a user-supplied model_list as the only models
template: configmap-litellm.yaml
set:
proxy_config.model_list:
- model_name: my-model
litellm_params:
model: openai/my-model
api_key: os.environ/OPENAI_API_KEY
asserts:
- matchRegex:
path: data["config.yaml"]
pattern: 'model_list:\n- litellm_params:\n api_key: os.environ/OPENAI_API_KEY\n model: openai/my-model\n model_name: my-model\n$'
- it: should allow overriding liveness, readiness, and startup probes
template: deployment.yaml
set:

12
helm/litellm-helm/values.yaml
View file

@ -180,17 +180,7 @@ proxyConfigMap:
# Examples: https://github.com/BerriAI/litellm/tree/main/litellm/proxy/example_config_yaml
# Reference: https://docs.litellm.ai/docs/proxy/configs
proxy_config:
model_list:
# At least one model must exist for the proxy to start.
- model_name: gpt-3.5-turbo
litellm_params:
model: gpt-3.5-turbo
api_key: eXaMpLeOnLy
- model_name: fake-openai-endpoint
litellm_params:
model: openai/fake
api_key: fake-key
api_base: https://exampleopenaiendpoint-production.up.railway.app/
model_list: []
general_settings:
master_key: os.environ/PROXY_MASTER_KEY

4
litellm-proxy-extras/pyproject.toml
View file

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

57
litellm-rust/crates/inference-chat/src/common_utils.rs
View file

@ -6,9 +6,11 @@ use litellm_llms::{
bedrock::chat::converse_transformation::BEDROCK_CHAT_COMPLETIONS_CONFIG,
openai_like::chat::transformation::OPENAI_LIKE_CHAT_COMPLETIONS_CONFIG,
};
use litellm_llms_types::formats::chat_completions::ChatMessage;
use serde_json::{Map, Value};
use super::Error;
use crate::types::{ChatCompletionsCall, ResolvedChatCompletionsRequest};
const HEADER_CONTEXT: &str = "chat completions";
@ -42,3 +44,58 @@ pub(super) fn string_headers(
) -> Result<Vec<(String, String)>, Error> {
shared_string_headers(HEADER_CONTEXT, extra_headers).map_err(Error::from)
}
pub(super) fn resolve_request(
call: ChatCompletionsCall,
config: &'static dyn BaseConfig,
) -> Result<ResolvedChatCompletionsRequest, Error> {
let messages = serde_json::from_value::<Vec<ChatMessage>>(call.messages).map_err(|err| {
Error::InvalidRequest(litellm_llms::ErrorDetail::invalid(
"chat completions messages",
err,
))
})?;
if messages.is_empty() {
return Err(Error::InvalidRequest(
"chat completions requires at least one message".into(),
));
}
if let Some(reason) = config.unsupported_reason(&messages, &call.optional_params) {
return Err(Error::Unsupported(reason.0));
}
Ok(ResolvedChatCompletionsRequest {
messages,
optional_params: call.optional_params,
api_key: call.api_key,
api_base: call.api_base,
extra_headers: call.extra_headers,
timeout: call.timeout,
})
}
#[cfg(test)]
mod tests {
use rstest::rstest;
use super::chat_completions_provider;
use litellm_core_utils::get_llm_provider_logic::LlmProviders;
#[rstest]
#[case::anthropic(LlmProviders::Anthropic, ("anthropic-version", "2023-06-01"))]
#[case::bedrock(LlmProviders::Bedrock, ("Content-Type", "application/json"))]
fn resolves_supported_provider_configs(
#[case] provider: LlmProviders,
#[case] expected_header: (&str, &str),
) {
let config = chat_completions_provider(provider)
.expect("provider has a chat completions config")
.config();
assert!(config.default_headers().contains(&expected_header));
}
#[rstest]
#[case::openai(LlmProviders::Openai)]
fn rejects_providers_without_a_chat_config(#[case] provider: LlmProviders) {
assert!(chat_completions_provider(provider).is_none());
}
}

116
litellm-rust/crates/inference-chat/src/handler.rs
View file

@ -3,44 +3,66 @@ use std::time::Duration;
use litellm_auth::AuthServices;
use litellm_host::interceptors::{Interceptors, RawResponse, RequestContext, WireRequest};
use litellm_http::{Client, outbound::OutboundRequest, request::truncate_error_body};
use litellm_http::{
Client,
outbound::OutboundRequest,
request::{truncate_error_body, with_default_headers},
};
use litellm_llms::base_llm::{
auth::{Authenticated, resolve_auth},
auth::{Authenticated, ValidatedEnvironment, resolve_auth},
chat::transformation::BaseConfig,
chat::transformation::ProviderChatResponseData,
};
use litellm_llms_types::formats::chat_completions::ChatCompletionsResponse;
use litellm_secrets::source::Secrets;
use serde_json::Value;
use super::Error;
use crate::{constants::CHAT_COMPLETIONS_TIMEOUT_SECS, types::ProviderChatCompletionsRequest};
use crate::{
common_utils::string_headers, constants::CHAT_COMPLETIONS_TIMEOUT_SECS,
types::ResolvedChatCompletionsRequest,
};
#[expect(clippy::too_many_arguments)] // mirrors BaseLLMHTTPHandler.completion, which takes each stage input separately
pub(super) async fn execute(
http: &Client,
auth: &AuthServices,
request: ProviderChatCompletionsRequest,
config: &'static dyn BaseConfig,
provider: litellm_inference::provider::ResolvedProvider<'_>,
call: ResolvedChatCompletionsRequest,
secrets: Secrets,
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,
custom_llm_provider,
config,
url,
body,
optional_params,
environment,
secrets,
timeout,
api_key,
} = request;
let validated = config.validate_environment(
string_headers(call.extra_headers)?,
call.api_key.as_deref(),
provider.model,
&call.optional_params,
&|key| secrets.get(key),
)?;
let url = config.get_complete_url(
call.api_base.as_deref(),
provider.model,
&call.optional_params,
&|key| secrets.get(key),
)?;
let body = config
.transform_request(provider.model, call.messages, call.optional_params.clone())?
.body;
let environment = ValidatedEnvironment {
headers: with_default_headers(validated.headers, config.default_headers()),
auth: validated.auth,
};
let provider_name = <&'static str>::from(provider.provider);
let context = RequestContext {
model: model.clone(),
custom_llm_provider,
optional_params: Value::Object(optional_params),
model: provider.model.to_owned(),
custom_llm_provider: provider_name.to_owned(),
optional_params: Value::Object(call.optional_params),
secret_fields: Vec::new(),
api_key,
api_key: call.api_key.map(litellm_auth::SecretValue::new),
};
let authenticated = resolve_auth(auth, environment, &|key| secrets.get(key)).await?;
let identity = litellm_host::interceptors::ProviderIdentity {
@ -76,7 +98,7 @@ pub(super) async fn execute(
},
wire.url,
&wire.body,
timeout,
call.timeout,
)?;
let response = litellm_inference::outbound::send(outbound, http)
@ -121,7 +143,7 @@ pub(super) async fn execute(
))
})?;
config
.transform_response(&model, ProviderChatResponseData { body })
.transform_response(provider.model, ProviderChatResponseData { body })
.map_err(Error::from)
.map_err(as_response_error)
},
@ -134,7 +156,7 @@ pub(super) async fn execute(
///
/// A config reports the same variants on either side of the call: a missing
/// field or an unsupported block can mean "this request cannot be translated"
/// during prepare and "this response cannot be normalized" here. Only the
/// before the request is sent and "this response cannot be normalized" here. Only the
/// second kind has already been billed, and a host that keeps a reference
/// implementation must not retry those, so collapse them to one variant that
/// can only mean the provider was already called.
@ -178,9 +200,10 @@ mod tests {
use super::*;
use crate::{
prepare::{prepare_provider_request, resolve_request},
types::ChatCompletionsRequest,
common_utils::{chat_completions_provider, resolve_request},
types::{ChatCompletionsCall, ResolvedChatCompletionsRequest},
};
use litellm_inference::provider::resolve_llm_provider;
const ANTHROPIC_MESSAGE: &str = r#"{"id":"msg_1","type":"message","role":"assistant","model":"claude-sonnet-4-5","content":[{"type":"text","text":"hello"}],"stop_reason":"end_turn","stop_sequence":null,"usage":{"input_tokens":11,"output_tokens":4}}"#;
@ -217,22 +240,33 @@ mod tests {
}
}
fn prepared(api_base: &str) -> ProviderChatCompletionsRequest {
prepare_provider_request(
resolve_request(ChatCompletionsRequest {
model: "anthropic/claude-sonnet-4-5",
fn prepared(
api_base: &str,
) -> (
litellm_inference::provider::ResolvedProvider<'static>,
&'static dyn BaseConfig,
ResolvedChatCompletionsRequest,
) {
let provider =
resolve_llm_provider("anthropic/claude-sonnet-4-5", None, "chat completions").unwrap();
let config = chat_completions_provider(provider.provider)
.expect("provider has a chat completions config")
.config();
let request = resolve_request(
ChatCompletionsCall {
model: "anthropic/claude-sonnet-4-5".into(),
messages: json!([{"role": "user", "content": "hi"}]),
optional_params: json!({"max_tokens": 16}).as_object().unwrap().clone(),
api_key: Some("sk-test"),
api_base: Some(api_base),
api_key: Some("sk-test".into()),
api_base: Some(api_base.into()),
custom_llm_provider: None,
extra_headers: None,
timeout: None,
})
.unwrap(),
std::sync::Arc::new(|_: &str| None),
},
config,
)
.unwrap()
.unwrap();
(provider, config, request)
}
#[rstest]
@ -247,10 +281,14 @@ mod tests {
.await;
let interceptors = RecordingHooks::default();
let (provider, config, request) = prepared(&upstream.uri());
execute(
&Client::plain_for_test(),
&AuthServices::default(),
prepared(&upstream.uri()),
config,
provider,
request,
std::sync::Arc::new(|_: &str| None),
None,
None,
&interceptors,
@ -288,10 +326,14 @@ mod tests {
.await;
let interceptors = RecordingHooks::default();
let (provider, config, request) = prepared(&upstream.uri());
let error = execute(
&Client::plain_for_test(),
&AuthServices::default(),
prepared(&upstream.uri()),
config,
provider,
request,
std::sync::Arc::new(|_: &str| None),
None,
None,
&interceptors,

53
litellm-rust/crates/inference-chat/src/lib.rs
View file

@ -5,14 +5,13 @@ pub use litellm_inference::RouteError as Error;
mod common_utils;
pub mod constants;
pub(crate) mod handler;
mod prepare;
use litellm_llms_types::formats::chat_completions::ChatCompletionsResponse;
use prepare::{prepare_provider_request, resolve_request};
use litellm_auth::AuthServices;
use litellm_inference::provider::resolve_llm_provider;
use litellm_secrets::source::SecretSource;
use std::sync::Arc;
use types::ChatCompletionsRequest;
use types::{ChatCompletionsCall, ChatCompletionsRequest};
#[derive(Clone)]
pub struct ChatCompletionsRoute {
@ -55,7 +54,7 @@ impl ChatCompletionsRoute {
} = options.into();
litellm_host::lifecycle::observe_unary(
observers.clone(),
self.run_call(
self.run(
request.into(),
cache_options,
interceptors,
@ -65,30 +64,52 @@ impl ChatCompletionsRoute {
.await
}
#[tracing::instrument(name = "litellm.route", skip_all, fields(
route = "chat_completions",
model = %call.model,
provider,
resolved_model,
stream = false,
outcome
))]
async fn run(
&self,
request: ChatCompletionsRequest<'_>,
call: ChatCompletionsCall,
cache_options: Option<litellm_cache_response::CachePolicy>,
interceptors: &impl litellm_host::interceptors::Interceptors<Error>,
observers: Option<&ObservationSender>,
) -> Result<ChatCompletionsResponse, Error> {
let resolved = resolve_request(request)?;
let snapshot = self
.secrets
.resolve(&resolved.config.secret_names())
.await?;
let prepared = prepare_provider_request(resolved, snapshot)?;
litellm_inference::diagnostic::provider(&prepared.model, &prepared.custom_llm_provider);
let execute: futures_util::future::BoxFuture<'_, Result<ChatCompletionsResponse, Error>> =
Box::pin(handler::execute(
litellm_inference::diagnostic::unary(async {
let model = call.model.clone();
let custom_llm_provider = call.custom_llm_provider.clone();
let provider =
resolve_llm_provider(&model, custom_llm_provider.as_deref(), "chat completions")?;
let config = common_utils::chat_completions_provider(provider.provider)
.ok_or_else(|| Error::InvalidProvider(<&str>::from(provider.provider).to_string()))?
.config();
let call = common_utils::resolve_request(call, config)?;
litellm_inference::diagnostic::provider(
provider.model,
<&'static str>::from(provider.provider),
);
let secrets = self.secrets.resolve(&config.secret_names()).await?;
let execute: futures_util::future::BoxFuture<
'_,
Result<ChatCompletionsResponse, Error>,
> = Box::pin(handler::execute(
&self.http,
&self.auth,
prepared,
config,
provider,
call,
secrets,
self.cache.clone(),
cache_options,
interceptors,
observers,
));
execute.await
execute.await
})
.await
}
}

638
litellm-rust/crates/inference-chat/src/prepare.rs
View file

@ -1,638 +0,0 @@
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 serde_json::Value;
use super::{
Error,
common_utils::{chat_completions_provider, string_headers},
};
use crate::types::{
ChatCompletionsRequest, ProviderChatCompletionsRequest, ResolvedChatCompletionsRequest,
};
use litellm_inference::provider::resolve_llm_provider;
pub(super) struct ResolvedProvider {
pub(super) model: String,
pub(super) custom_llm_provider: String,
pub(super) config: &'static dyn BaseConfig,
}
pub(super) fn resolve_provider_config<'a>(
model: &'a str,
custom_llm_provider: Option<&'a str>,
) -> Result<ResolvedProvider, Error> {
let provider_info = resolve_llm_provider(model, custom_llm_provider, "chat completions")?;
let config = chat_completions_provider(provider_info.provider)
.ok_or_else(|| Error::InvalidProvider(<&str>::from(provider_info.provider).to_string()))?
.config();
Ok(ResolvedProvider {
model: provider_info.model.to_string(),
custom_llm_provider: <&str>::from(provider_info.provider).to_string(),
config,
})
}
pub(super) fn parse_messages(messages: Value) -> Result<Vec<ChatMessage>, Error> {
serde_json::from_value(messages).map_err(|err| {
Error::InvalidRequest(litellm_llms::ErrorDetail::invalid(
"chat completions messages",
err,
))
})
}
pub(super) fn resolve_request(
request: ChatCompletionsRequest<'_>,
) -> Result<ResolvedChatCompletionsRequest<'_>, Error> {
let ResolvedProvider {
model,
custom_llm_provider,
config,
} = resolve_provider_config(request.model, request.custom_llm_provider)?;
let messages = parse_messages(request.messages)?;
if messages.is_empty() {
return Err(Error::InvalidRequest(
"chat completions requires at least one message".into(),
));
}
if let Some(reason) = config.unsupported_reason(&messages, &request.optional_params) {
return Err(Error::Unsupported(reason.0));
}
Ok(ResolvedChatCompletionsRequest {
model,
custom_llm_provider,
config,
messages,
optional_params: request.optional_params,
api_key: request.api_key,
api_base: request.api_base,
extra_headers: request.extra_headers,
timeout: request.timeout,
})
}
fn validate_environment(
request: &ResolvedChatCompletionsRequest<'_>,
model: &str,
config: &dyn BaseConfig,
secrets: &dyn Lookup,
) -> Result<ValidatedEnvironment, Error> {
let env_lookup = |key: &str| secrets.get(key);
let forwarded = string_headers(request.extra_headers.clone())?;
let validated = config.validate_environment(
forwarded,
request.api_key,
model,
&request.optional_params,
&env_lookup,
)?;
Ok(ValidatedEnvironment {
headers: with_default_headers(validated.headers, config.default_headers()),
auth: validated.auth,
})
}
#[tracing::instrument(name = "litellm.prepare", level = "debug", skip_all)]
pub(super) fn prepare_provider_request(
request: ResolvedChatCompletionsRequest<'_>,
secrets: Secrets,
) -> Result<ProviderChatCompletionsRequest, Error> {
let environment =
validate_environment(&request, &request.model, request.config, secrets.as_ref())?;
let model = request.model;
let config = request.config;
let env_lookup = |key: &str| secrets.get(key);
let url = config.get_complete_url(
request.api_base,
&model,
&request.optional_params,
&env_lookup,
)?;
let transformed =
config.transform_request(&model, request.messages, request.optional_params.clone())?;
Ok(ProviderChatCompletionsRequest {
model,
custom_llm_provider: request.custom_llm_provider,
config,
url,
body: transformed.body,
optional_params: request.optional_params,
environment,
secrets,
timeout: request.timeout,
api_key: request.api_key.map(|key| SecretValue::new(key.to_string())),
})
}
#[cfg(test)]
mod tests {
use litellm_auth::CredentialPlacement;
use litellm_llms::base_llm::auth::{AuthScheme, resolve_auth};
use serde_json::{Map, Value, json};
use super::{prepare_provider_request, resolve_request};
use crate::{
Error,
types::{ChatCompletionsRequest, ProviderChatCompletionsRequest},
};
fn prepare_chat_completions_call(
request: ChatCompletionsRequest<'_>,
) -> Result<ProviderChatCompletionsRequest, Error> {
prepare_provider_request(
resolve_request(request)?,
std::sync::Arc::new(|_: &str| None),
)
}
/// The headers as they go on the wire, credential applied.
fn wire_headers(prepared: &ProviderChatCompletionsRequest) -> Vec<(String, String)> {
tokio::runtime::Builder::new_current_thread()
.build()
.unwrap()
.block_on(resolve_auth(
&litellm_auth::AuthServices::default(),
prepared.environment.clone(),
&|_| None,
))
.unwrap()
.headers
}
fn request<'a>(
model: &'a str,
provider: Option<&'a str>,
messages: Value,
optional_params: Value,
) -> ChatCompletionsRequest<'a> {
ChatCompletionsRequest {
model,
messages,
optional_params: match optional_params {
Value::Object(map) => map,
other => panic!("params must be an object, got {other}"),
},
api_key: Some("sk-test"),
api_base: None,
custom_llm_provider: provider,
extra_headers: None,
timeout: None,
}
}
fn preparation_error(request: ChatCompletionsRequest<'_>) -> Error {
prepare_chat_completions_call(request)
.err()
.expect("request preparation should fail")
}
#[test]
fn resolves_the_provider_from_the_model_prefix() {
let prepared = prepare_chat_completions_call(request(
"anthropic/claude-sonnet-4-5",
None,
json!([{"role": "user", "content": "hi"}]),
json!({"max_tokens": 16}),
))
.expect("prepares");
assert_eq!(prepared.model, "claude-sonnet-4-5");
assert_eq!(prepared.url, "https://api.anthropic.com/v1/messages");
assert_eq!(prepared.body["model"], json!("claude-sonnet-4-5"));
}
#[test]
fn strips_an_explicit_provider_prefix_from_the_model() {
let prepared = prepare_chat_completions_call(request(
"anthropic/claude-sonnet-4-5",
Some("anthropic"),
json!([{"role": "user", "content": "hi"}]),
json!({}),
))
.expect("prepares");
assert_eq!(prepared.model, "claude-sonnet-4-5");
}
#[test]
fn adds_the_auth_and_default_headers() {
let prepared = prepare_chat_completions_call(request(
"claude-sonnet-4-5",
Some("anthropic"),
json!([{"role": "user", "content": "hi"}]),
json!({}),
))
.expect("prepares");
assert!(
wire_headers(&prepared).contains(&("x-api-key".to_string(), "sk-test".to_string()))
);
assert!(
wire_headers(&prepared)
.contains(&("anthropic-version".to_string(), "2023-06-01".to_string()))
);
assert!(matches!(
prepared.environment.auth,
AuthScheme::Credential {
placement: CredentialPlacement::Header("x-api-key"),
..
}
));
}
#[test]
fn the_deployment_credential_replaces_a_caller_supplied_auth_header() {
// Python builds `{**headers, **anthropic_headers}`, so the deployment's key
// overwrites a forwarded one. Honouring the caller's would let whoever sends
// the request choose the Anthropic principal it bills to.
let mut call = request(
"claude-sonnet-4-5",
Some("anthropic"),
json!([{"role": "user", "content": "hi"}]),
json!({}),
);
call.extra_headers = Some(Map::from_iter([(
"X-Api-Key".to_string(),
json!("sk-caller"),
)]));
let prepared = prepare_chat_completions_call(call).expect("prepares");
let headers = wire_headers(&prepared);
let keys: Vec<_> = headers
.iter()
.filter(|(name, _)| name.eq_ignore_ascii_case("x-api-key"))
.collect();
assert_eq!(keys.len(), 1, "got {:?}", headers);
assert_eq!(keys[0].1, "sk-test");
}
#[test]
fn a_forwarded_authorization_header_suppresses_the_resolved_api_key_header() {
// Anthropic's `validate_environment` pops `x-api-key` and sets `authorization`
// for an OAuth token, so re-adding the key here would put the credential into
// a header the host removed on purpose.
let mut call = request(
"claude-sonnet-4-5",
Some("anthropic"),
json!([{"role": "user", "content": "hi"}]),
json!({}),
);
call.extra_headers = Some(Map::from_iter([
(
"Authorization".to_string(),
json!("Bearer sk-ant-oat01-token"),
),
("X-Api-Key".to_string(), json!("sk-caller")),
]));
let prepared = prepare_chat_completions_call(call).expect("prepares");
assert!(
!wire_headers(&prepared)
.iter()
.any(|(name, value)| name.eq_ignore_ascii_case("x-api-key") && value == "sk-test"),
"the resolved key must not be applied over an OAuth bearer, got {:?}",
wire_headers(&prepared)
);
assert!(
wire_headers(&prepared)
.iter()
.any(|(name, value)| name.eq_ignore_ascii_case("authorization")
&& value == "Bearer sk-ant-oat01-token")
);
}
#[test]
fn an_unrelated_forwarded_authorization_does_not_defer_the_resolved_key() {
// Only an OAuth bearer replaces the credential. Python sends the deployment's
// `x-api-key` alongside any other forwarded `authorization`, so deferring on
// the mere presence of that header would drop the deployment's auth.
let mut call = request(
"claude-sonnet-4-5",
Some("anthropic"),
json!([{"role": "user", "content": "hi"}]),
json!({}),
);
call.extra_headers = Some(Map::from_iter([
("Authorization".to_string(), json!("Bearer unrelated")),
("X-Api-Key".to_string(), json!("sk-caller")),
]));
let prepared = prepare_chat_completions_call(call).expect("prepares");
let headers = wire_headers(&prepared);
let keys: Vec<_> = headers
.iter()
.filter(|(name, _)| name.eq_ignore_ascii_case("x-api-key"))
.collect();
assert_eq!(keys.len(), 1, "got {:?}", headers);
assert_eq!(keys[0].1, "sk-test");
assert!(
wire_headers(&prepared)
.iter()
.any(|(name, value)| name.eq_ignore_ascii_case("authorization")
&& value == "Bearer unrelated"),
"the unrelated authorization must survive, got {:?}",
wire_headers(&prepared)
);
}
#[rstest::rstest]
fn rejects_empty_messages_before_resolving_credentials() {
let call = ChatCompletionsRequest {
api_key: None,
..request("claude-sonnet-4-5", Some("anthropic"), json!([]), json!({}))
};
assert!(matches!(preparation_error(call), Error::InvalidRequest(_)));
}
#[test]
fn rejects_an_unknown_provider() {
assert_eq!(
preparation_error(request(
"openai/gpt-4o",
None,
json!([{"role": "user", "content": "hi"}]),
json!({}),
)),
Error::InvalidProvider("openai".to_string())
);
}
#[test]
fn rejects_a_model_with_no_resolvable_provider() {
assert!(matches!(
preparation_error(request(
"claude-sonnet-4-5",
None,
json!([{"role": "user", "content": "hi"}]),
json!({}),
)),
Error::InvalidProvider(_)
));
}
#[test]
fn rejects_an_empty_or_malformed_message_list() {
assert_eq!(
preparation_error(request(
"anthropic/claude-sonnet-4-5",
None,
json!([]),
json!({}),
)),
Error::InvalidRequest(
"chat completions requires at least one message"
.to_string()
.into()
)
);
assert!(matches!(
preparation_error(request(
"anthropic/claude-sonnet-4-5",
None,
json!("not a list"),
json!({}),
)),
Error::InvalidRequest(_)
));
}
#[test]
fn rejects_non_string_extra_headers() {
let mut call = request(
"anthropic/claude-sonnet-4-5",
None,
json!([{"role": "user", "content": "hi"}]),
json!({}),
);
call.extra_headers = Some(Map::from_iter([("x-trace".to_string(), json!(7))]));
assert_eq!(
preparation_error(call),
Error::Headers(litellm_http::request::HeaderError {
context: "chat completions",
name: "x-trace".to_string(),
actual: "number",
})
);
}
#[test]
fn prepares_a_bedrock_call_without_resolving_credentials() {
let mut call = request(
"bedrock/us-east-1/anthropic.claude-v2",
None,
json!([{"role": "user", "content": "hi"}]),
json!({"maxTokens": 16}),
);
call.api_key = None;
let prepared = prepare_chat_completions_call(call).expect("prepares");
assert_eq!(
prepared.url,
"https://bedrock-runtime.us-east-1.amazonaws.com/model/anthropic.claude-v2/converse"
);
assert!(matches!(
&prepared.environment.auth,
AuthScheme::AwsSigV4 { region, service: "bedrock", .. } if region == "us-east-1"
));
// SigV4 signs the serialized body, so prepare must not have added an
// Authorization header; the signer does it over the bytes sent.
assert!(
!prepared
.environment
.headers
.iter()
.any(|(name, _)| name.eq_ignore_ascii_case("authorization"))
);
assert_eq!(prepared.body["inferenceConfig"], json!({"maxTokens": 16}));
}
#[tokio::test]
async fn a_forwarded_client_header_does_not_enter_the_bedrock_signature() {
// Python signs only the AWS header set and reattaches the rest, so a header
// the caller forwarded rides along without joining the canonical request.
// Signing it makes Converse 403 on a deployment that works on Python.
let mut call = request(
"bedrock/us-east-1/anthropic.claude-v2",
None,
json!([{"role": "user", "content": "hi"}]),
json!({
"maxTokens": 16,
"aws_access_key_id": "AKIDEXAMPLE",
"aws_secret_access_key": "wJalrXUtnFEMI/K7MDENG+bPxRfiCYEXAMPLEKEY"
}),
);
// A key would resolve to a bearer token and never reach the signer.
call.api_key = None;
call.extra_headers = Some(Map::from_iter([(
"x-request-id".to_string(),
json!("abc-123"),
)]));
let prepared = prepare_chat_completions_call(call).expect("prepares");
let authenticated = resolve_auth(
&litellm_auth::AuthServices::default(),
prepared.environment,
&|_| None,
)
.await
.expect("resolves");
let signed = crate::handler::outbound_request(
authenticated,
prepared.url,
&prepared.body,
prepared.timeout,
)
.expect("signs");
let authorization = signed
.header("authorization")
.expect("carries an authorization header")
.to_string();
assert!(
authorization.starts_with("AWS4-HMAC-SHA256"),
"expected a SigV4 signature, got {authorization}"
);
assert!(
!authorization.contains("x-request-id"),
"forwarded header reached SignedHeaders: {authorization}"
);
// It still goes on the wire, it is just not part of the signature.
assert!(
signed
.headers()
.iter()
.any(|(name, value)| name == "x-request-id" && value == "abc-123"),
"forwarded header was dropped instead of reattached"
);
}
#[rstest::rstest]
#[case::authorization("Authorization")]
#[case::amz_date("x-amz-date")]
#[case::security_token("x-amz-security-token")]
#[case::date("Date")]
#[tokio::test]
async fn rejects_a_forwarded_header_the_signer_computes(#[case] forwarded: &str) {
let call = ChatCompletionsRequest {
api_key: None,
extra_headers: Some(Map::from_iter([(forwarded.to_string(), json!("forged"))])),
..request(
"bedrock/us-east-1/anthropic.claude-v2",
None,
json!([{"role": "user", "content": "hi"}]),
json!({
"maxTokens": 16,
"aws_access_key_id": "AKIDEXAMPLE",
"aws_secret_access_key": "wJalrXUtnFEMI/K7MDENG+bPxRfiCYEXAMPLEKEY"
}),
)
};
let prepared = prepare_chat_completions_call(call).expect("prepares");
let authenticated = resolve_auth(
&litellm_auth::AuthServices::default(),
prepared.environment,
&|_| None,
)
.await
.expect("resolves");
let error = crate::handler::outbound_request(
authenticated,
prepared.url,
&prepared.body,
prepared.timeout,
)
.expect_err("conflicting signing headers must fail");
assert!(
matches!(error, Error::Unsupported(_)),
"{forwarded} returned {error:?}"
);
}
#[test]
fn a_bedrock_deployment_bearer_outranks_a_forwarded_authorization() {
// `get_request_headers` assigns `headers["Authorization"]` unconditionally
// once a bearer token resolves, so the deployment's identity wins on
// Python. Keeping the caller's would authorize and bill the call as a
// different principal, and only when the deployment carries `rust: true`.
let mut call = request(
"bedrock/us-east-1/anthropic.claude-v2",
None,
json!([{"role": "user", "content": "hi"}]),
json!({"maxTokens": 16}),
);
call.extra_headers = Some(Map::from_iter([(
"Authorization".to_string(),
json!("Bearer caller-supplied"),
)]));
let prepared = prepare_chat_completions_call(call).expect("prepares");
let headers = wire_headers(&prepared);
let authorizations: Vec<_> = headers
.iter()
.filter(|(name, _)| name.eq_ignore_ascii_case("authorization"))
.map(|(_, value)| value.as_str())
.collect();
assert_eq!(
authorizations,
vec!["Bearer sk-test"],
"the deployment token must be the only authorization on the wire"
);
}
#[test]
fn an_anthropic_forwarded_oauth_bearer_still_outranks_the_resolved_key() {
// The opposite precedence, and deliberate: Anthropic's own transform
// honours a forwarded OAuth bearer, so the Bedrock fix above must not be
// generalized into a rule that the configured key always wins.
//
// An OAuth bearer is the whole of that exception. This forwarded a plain
// `x-api-key` until round 17, which read as the same claim and was not:
// Python overwrites a forwarded `x-api-key` with the deployment's.
let mut call = request(
"claude-sonnet-4-5",
Some("anthropic"),
json!([{"role": "user", "content": "hi"}]),
json!({}),
);
call.extra_headers = Some(Map::from_iter([(
"authorization".to_string(),
json!("Bearer sk-ant-oat01-forwarded"),
)]));
let prepared = prepare_chat_completions_call(call).expect("prepares");
let headers = wire_headers(&prepared);
let keys: Vec<_> = headers
.iter()
.filter(|(name, _)| name.eq_ignore_ascii_case("x-api-key"))
.map(|(_, value)| value.as_str())
.collect();
assert!(keys.is_empty(), "got {:?}", headers);
assert!(
wire_headers(&prepared)
.iter()
.any(|(name, value)| name.eq_ignore_ascii_case("authorization")
&& value == "Bearer sk-ant-oat01-forwarded")
);
}
#[test]
fn a_bedrock_api_key_is_sent_as_a_bearer_token_instead_of_being_signed() {
// The configured bearer identity has its own account and quota boundary,
// so a request carrying one must not be signed as whatever principal the
// host's AWS credentials resolve to.
let prepared = prepare_chat_completions_call(request(
"bedrock/us-east-1/anthropic.claude-v2",
None,
json!([{"role": "user", "content": "hi"}]),
json!({"maxTokens": 16}),
))
.expect("prepares");
assert!(matches!(
&prepared.environment.auth,
AuthScheme::Credential { placement: CredentialPlacement::Bearer, secret }
if secret.expose() == "sk-test"
));
assert!(
wire_headers(&prepared)
.iter()
.any(|(name, value)| name.eq_ignore_ascii_case("authorization")
&& value == "Bearer sk-test"),
"prepare did not carry the bearer token"
);
}
}

40
litellm-rust/crates/inference-chat/src/route.rs
View file

@ -1,4 +1,3 @@
use litellm_host::observation::ObservationSender;
use std::convert::Infallible;
use litellm_host::{
@ -7,10 +6,7 @@ use litellm_host::{
};
use litellm_llms_types::formats::chat_completions::ChatCompletionsResponse;
use super::{
ChatCompletionsRoute, Error,
types::{ChatCompletionsCall, ChatCompletionsRequest},
};
use super::{ChatCompletionsRoute, Error, types::ChatCompletionsCall};
pub struct ChatCompletions;
@ -37,44 +33,12 @@ impl ChatCompletionsRoute {
call,
observers,
move |call, _, interceptors, observers| async move {
self.run_call(call, cache_options, &interceptors, observers.as_ref())
self.run(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> {
litellm_inference::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 litellm_inference::caching::Cachable for ChatCompletions {

29
litellm-rust/crates/inference-chat/src/types.rs
View file

@ -1,8 +1,5 @@
use litellm_secrets::source::Secrets;
use std::time::Duration;
use litellm_auth::SecretValue;
use litellm_llms::base_llm::{auth::ValidatedEnvironment, chat::transformation::BaseConfig};
use litellm_llms_types::formats::chat_completions::ChatMessage;
use serde_json::{Map, Value};
@ -49,31 +46,11 @@ impl From<ChatCompletionsRequest<'_>> for ChatCompletionsCall {
}
}
pub struct ResolvedChatCompletionsRequest<'a> {
pub model: String,
pub custom_llm_provider: String,
pub config: &'static dyn BaseConfig,
pub struct ResolvedChatCompletionsRequest {
pub messages: Vec<ChatMessage>,
pub optional_params: Map<String, Value>,
pub api_key: Option<&'a str>,
pub api_base: Option<&'a str>,
pub api_key: Option<String>,
pub api_base: Option<String>,
pub extra_headers: Option<Map<String, Value>>,
pub timeout: Option<Duration>,
}
pub struct ProviderChatCompletionsRequest {
pub model: String,
pub custom_llm_provider: String,
pub config: &'static dyn BaseConfig,
pub url: String,
pub body: Value,
/// The route's parameters before the provider transformation, reported to the host
/// beside the wire request.
pub optional_params: Map<String, Value>,
/// The forwarded and default headers plus how the call authenticates; the credential
/// itself is applied when the request is sent.
pub environment: ValidatedEnvironment,
pub secrets: Secrets,
pub timeout: Option<Duration>,
pub api_key: Option<SecretValue>,
}

439
litellm-rust/crates/inference-chat/tests/request.rs Normal file
View file

@ -0,0 +1,439 @@
use std::time::Duration;
use litellm_inference_chat::{Error, types::ChatCompletionsRequest};
use rstest::rstest;
use serde_json::{Map, Value, json};
use wiremock::ResponseTemplate;
mod support;
use support::*;
const ANTHROPIC_MESSAGE: &str = r#"{"id":"msg_1","type":"message","role":"assistant","model":"claude-sonnet-4-5","content":[{"type":"text","text":"hello"}],"stop_reason":"end_turn","stop_sequence":null,"usage":{"input_tokens":11,"output_tokens":4}}"#;
fn request<'a>(
model: &'a str,
provider: Option<&'a str>,
messages: Value,
optional_params: Value,
) -> ChatCompletionsRequest<'a> {
let Value::Object(optional_params) = optional_params else {
panic!("params must be an object");
};
ChatCompletionsRequest {
model,
messages,
optional_params,
api_key: Some("sk-test"),
api_base: None,
custom_llm_provider: provider,
extra_headers: None,
timeout: Some(Duration::from_secs(10)),
}
}
async fn complete(
request: ChatCompletionsRequest<'_>,
) -> Result<litellm_llms_types::formats::chat_completions::ChatCompletionsResponse, Error> {
chat_completions_route().execute(request, &(), None).await
}
fn object(value: Value) -> Map<String, Value> {
let Value::Object(map) = value else {
panic!("expected an object");
};
map
}
fn anthropic_response() -> ResponseTemplate {
ResponseTemplate::new(200).set_body_raw(ANTHROPIC_MESSAGE, "application/json")
}
fn bedrock_response() -> ResponseTemplate {
json_response(json!({
"output": {"message": {"role": "assistant", "content": [{"text": "hello"}]}},
"stopReason": "end_turn",
"usage": {"inputTokens": 11, "outputTokens": 4, "totalTokens": 15}
}))
}
#[rstest]
#[tokio::test]
async fn resolves_the_provider_from_the_model_prefix() {
let upstream = upstream([anthropic_response()]).await;
let base = upstream.uri();
complete(ChatCompletionsRequest {
api_base: Some(&base),
..request(
"anthropic/claude-sonnet-4-5",
None,
json!([{"role": "user", "content": "hi"}]),
json!({"max_tokens": 16}),
)
})
.await
.expect("call succeeds");
let sent = only_request(&upstream).await;
assert_eq!(sent.url.path(), "/v1/messages");
assert_eq!(sent.json()["model"], "claude-sonnet-4-5");
}
#[rstest]
#[tokio::test]
async fn strips_an_explicit_provider_prefix_from_the_model() {
let upstream = upstream([anthropic_response()]).await;
let base = upstream.uri();
complete(ChatCompletionsRequest {
api_base: Some(&base),
..request(
"anthropic/claude-sonnet-4-5",
Some("anthropic"),
json!([{"role": "user", "content": "hi"}]),
json!({}),
)
})
.await
.expect("call succeeds");
assert_eq!(
only_request(&upstream).await.json()["model"],
"claude-sonnet-4-5"
);
}
#[rstest]
#[tokio::test]
async fn adds_the_auth_and_default_headers() {
let upstream = upstream([anthropic_response()]).await;
let base = upstream.uri();
complete(ChatCompletionsRequest {
api_base: Some(&base),
..request(
"claude-sonnet-4-5",
Some("anthropic"),
json!([{"role": "user", "content": "hi"}]),
json!({}),
)
})
.await
.expect("call succeeds");
let sent = only_request(&upstream).await;
assert_eq!(sent.header_values("x-api-key"), ["sk-test"]);
assert_eq!(sent.header("anthropic-version"), Some("2023-06-01"));
}
#[rstest]
#[tokio::test]
async fn the_deployment_credential_replaces_a_caller_supplied_auth_header() {
let upstream = upstream([anthropic_response()]).await;
let base = upstream.uri();
complete(ChatCompletionsRequest {
api_base: Some(&base),
extra_headers: Some(object(json!({"X-Api-Key": "sk-caller"}))),
..request(
"claude-sonnet-4-5",
Some("anthropic"),
json!([{"role": "user", "content": "hi"}]),
json!({}),
)
})
.await
.expect("call succeeds");
assert_eq!(
only_request(&upstream).await.header_values("x-api-key"),
["sk-test"]
);
}
#[rstest]
#[tokio::test]
async fn a_forwarded_authorization_header_suppresses_the_resolved_api_key_header() {
let upstream = upstream([anthropic_response()]).await;
let base = upstream.uri();
complete(ChatCompletionsRequest {
api_base: Some(&base),
extra_headers: Some(object(json!({
"Authorization": "Bearer sk-ant-oat01-token",
"X-Api-Key": "sk-caller"
}))),
..request(
"claude-sonnet-4-5",
Some("anthropic"),
json!([{"role": "user", "content": "hi"}]),
json!({}),
)
})
.await
.expect("call succeeds");
let sent = only_request(&upstream).await;
assert_eq!(sent.header_values("x-api-key"), ["sk-caller"]);
assert_eq!(
sent.header_values("authorization"),
["Bearer sk-ant-oat01-token"]
);
}
#[rstest]
#[tokio::test]
async fn an_unrelated_forwarded_authorization_does_not_defer_the_resolved_key() {
let upstream = upstream([anthropic_response()]).await;
let base = upstream.uri();
complete(ChatCompletionsRequest {
api_base: Some(&base),
extra_headers: Some(object(json!({
"Authorization": "Bearer unrelated",
"X-Api-Key": "sk-caller"
}))),
..request(
"claude-sonnet-4-5",
Some("anthropic"),
json!([{"role": "user", "content": "hi"}]),
json!({}),
)
})
.await
.expect("call succeeds");
let sent = only_request(&upstream).await;
assert_eq!(sent.header_values("x-api-key"), ["sk-test"]);
assert_eq!(sent.header_values("authorization"), ["Bearer unrelated"]);
}
#[rstest]
#[tokio::test]
async fn rejects_empty_messages_before_resolving_credentials() {
let error = complete(ChatCompletionsRequest {
api_key: None,
..request("claude-sonnet-4-5", Some("anthropic"), json!([]), json!({}))
})
.await
.expect_err("empty messages are rejected");
assert!(matches!(error, Error::InvalidRequest(_)));
}
#[rstest]
#[tokio::test]
async fn rejects_an_unknown_provider() {
assert_eq!(
complete(request(
"openai/gpt-4o",
None,
json!([{"role": "user", "content": "hi"}]),
json!({}),
))
.await
.unwrap_err(),
Error::InvalidProvider("openai".into())
);
}
#[rstest]
#[tokio::test]
async fn rejects_a_model_with_no_resolvable_provider() {
assert!(matches!(
complete(request(
"claude-sonnet-4-5",
None,
json!([{"role": "user", "content": "hi"}]),
json!({}),
))
.await,
Err(Error::InvalidProvider(_))
));
}
#[rstest]
#[case::empty(json!([]), Some("chat completions requires at least one message"))]
#[case::malformed(json!("not a list"), None)]
#[tokio::test]
async fn rejects_an_empty_or_malformed_message_list(
#[case] messages: Value,
#[case] expected: Option<&str>,
) {
let error = complete(request(
"anthropic/claude-sonnet-4-5",
None,
messages,
json!({}),
))
.await
.unwrap_err();
if let Some(expected) = expected {
assert_eq!(error, Error::InvalidRequest(expected.to_string().into()));
} else {
assert!(matches!(error, Error::InvalidRequest(_)));
}
}
#[rstest]
#[tokio::test]
async fn rejects_non_string_extra_headers() {
assert_eq!(
complete(ChatCompletionsRequest {
extra_headers: Some(Map::from_iter([("x-trace".to_string(), json!(7))])),
..request(
"anthropic/claude-sonnet-4-5",
None,
json!([{"role": "user", "content": "hi"}]),
json!({}),
)
})
.await
.unwrap_err(),
Error::Headers(litellm_http::request::HeaderError {
context: "chat completions",
name: "x-trace".to_string(),
actual: "number",
})
);
}
#[rstest]
#[tokio::test]
async fn a_forwarded_client_header_does_not_enter_the_bedrock_signature() {
let upstream = upstream([bedrock_response()]).await;
let base = upstream.uri();
let mut call = request(
"bedrock/us-east-1/anthropic.claude-v2",
None,
json!([{"role": "user", "content": "hi"}]),
json!({
"maxTokens": 16,
"aws_access_key_id": "AKIDEXAMPLE",
"aws_secret_access_key": "wJalrXUtnFEMI/K7MDENG+bPxRfiCYEXAMPLEKEY"
}),
);
call.api_key = None;
call.api_base = Some(&base);
call.extra_headers = Some(Map::from_iter([(
"x-request-id".to_string(),
json!("abc-123"),
)]));
complete(call).await.expect("call succeeds");
let sent = only_request(&upstream).await;
let authorization = sent.header("authorization").expect("signed request");
assert!(authorization.starts_with("AWS4-HMAC-SHA256"));
assert!(!authorization.contains("x-request-id"));
assert_eq!(sent.header("x-request-id"), Some("abc-123"));
assert_eq!(sent.url.path(), "/model/anthropic.claude-v2/converse");
let body: Value = serde_json::from_slice(&sent.body).unwrap();
assert_eq!(body["inferenceConfig"], json!({"maxTokens": 16}));
assert_eq!(
body["messages"],
json!([{"role": "user", "content": [{"text": "hi"}]}])
);
}
#[rstest]
#[case::authorization("Authorization")]
#[case::amz_date("x-amz-date")]
#[case::security_token("x-amz-security-token")]
#[case::date("Date")]
#[tokio::test]
async fn rejects_a_forwarded_header_the_signer_computes(#[case] forwarded: &str) {
let upstream = wiremock::MockServer::start().await;
let base = upstream.uri();
let call = ChatCompletionsRequest {
api_key: None,
api_base: Some(&base),
extra_headers: Some(Map::from_iter([(forwarded.to_string(), json!("forged"))])),
..request(
"bedrock/us-east-1/anthropic.claude-v2",
None,
json!([{"role": "user", "content": "hi"}]),
json!({
"maxTokens": 16,
"aws_access_key_id": "AKIDEXAMPLE",
"aws_secret_access_key": "wJalrXUtnFEMI/K7MDENG+bPxRfiCYEXAMPLEKEY"
}),
)
};
let error = complete(call)
.await
.expect_err("conflicting signing headers fail");
assert!(matches!(error, Error::Unsupported(_)), "{error:?}");
assert!(received(&upstream).await.is_empty());
}
#[rstest]
#[tokio::test]
async fn a_bedrock_deployment_bearer_outranks_a_forwarded_authorization() {
let upstream = upstream([bedrock_response()]).await;
let base = upstream.uri();
complete(ChatCompletionsRequest {
api_base: Some(&base),
extra_headers: Some(Map::from_iter([(
"Authorization".to_string(),
json!("Bearer caller-supplied"),
)])),
..request(
"bedrock/us-east-1/anthropic.claude-v2",
None,
json!([{"role": "user", "content": "hi"}]),
json!({"maxTokens": 16}),
)
})
.await
.expect("call succeeds");
assert_eq!(
only_request(&upstream).await.header_values("authorization"),
["Bearer sk-test"]
);
}
#[rstest]
#[tokio::test]
async fn an_anthropic_forwarded_oauth_bearer_still_outranks_the_resolved_key() {
let upstream = upstream([anthropic_response()]).await;
let base = upstream.uri();
complete(ChatCompletionsRequest {
api_base: Some(&base),
extra_headers: Some(Map::from_iter([(
"authorization".to_string(),
json!("Bearer sk-ant-oat01-forwarded"),
)])),
..request(
"claude-sonnet-4-5",
Some("anthropic"),
json!([{"role": "user", "content": "hi"}]),
json!({}),
)
})
.await
.expect("call succeeds");
let sent = only_request(&upstream).await;
assert!(sent.header_values("x-api-key").is_empty());
assert_eq!(
sent.header_values("authorization"),
["Bearer sk-ant-oat01-forwarded"]
);
}
#[rstest]
#[tokio::test]
async fn a_bedrock_api_key_is_sent_as_a_bearer_token_instead_of_being_signed() {
let upstream = upstream([bedrock_response()]).await;
let base = upstream.uri();
complete(ChatCompletionsRequest {
api_base: Some(&base),
..request(
"bedrock/us-east-1/anthropic.claude-v2",
None,
json!([{"role": "user", "content": "hi"}]),
json!({"maxTokens": 16}),
)
})
.await
.expect("call succeeds");
assert_eq!(
only_request(&upstream).await.header_values("authorization"),
["Bearer sk-test"]
);
}

4
litellm-rust/crates/inference-messages/AGENTS.md
View file

@ -7,8 +7,8 @@
- 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
- Route types such as `MessagesCall`, prepared requests, and response wrappers containing live streams describe execution. Reuse the shared Messages payload types inside them instead of defining another request or response schema here
- Route types such as `MessagesCall`, provider calls, and response wrappers containing live streams describe execution. Reuse the shared Messages payload types inside them instead of defining another request or response schema here
- The route returns `litellm_host::call::CallOutput`: a completed response, or a stream head and chunks
- Per-call dependencies are grouped in `litellm_inference::context::CallContext`. `src/lib.rs` explicitly sequences cache lookup, provider execution, result acceptance (`CallContext::result_ready`), and cache storage
- Per-call dependencies are grouped in `litellm_inference::context::CallContext`. `handler::execute` explicitly sequences the provider call build in Python's handler order, cache lookup, provider execution, result acceptance (`CallContext::result_ready`), and cache storage
- Preserve the order of validation, normalization, caller-requested parameter removal, and provider transformation when that order affects observable behavior
- Test provider dispatch, auth precedence, header handling, transformations, and responses through behavior, not source structure

22
litellm-rust/crates/inference-messages/src/common_utils.rs
View file

@ -82,17 +82,17 @@ mod tests {
use litellm_core_utils::get_llm_provider_logic::LlmProviders;
#[rstest]
#[case::tencent("tencent")]
#[case::anthropic("anthropic")]
#[case::azure_ai("azure_ai")]
#[case::bedrock("bedrock")]
#[case::deepseek("deepseek")]
#[case::edenai("edenai")]
#[case::minimax("minimax")]
#[case::vertex_ai("vertex_ai")]
fn provider_keeps_its_python_name(#[case] name: &str) {
let provider =
messages_provider(name.parse::<LlmProviders>().unwrap(), "claude-sonnet-4-5").unwrap();
#[case::tencent(LlmProviders::Tencent, "tencent")]
#[case::anthropic(LlmProviders::Anthropic, "anthropic")]
#[case::azure_ai(LlmProviders::AzureAi, "azure_ai")]
#[case::bedrock(LlmProviders::Bedrock, "bedrock")]
#[case::deepseek(LlmProviders::Deepseek, "deepseek")]
#[case::edenai(LlmProviders::Edenai, "edenai")]
#[case::minimax(LlmProviders::Minimax, "minimax")]
#[case::vertex_ai(LlmProviders::VertexAi, "vertex_ai")]
#[case::openrouter(LlmProviders::Openrouter, "openrouter")]
fn provider_keeps_its_python_name(#[case] provider: LlmProviders, #[case] name: &str) {
let provider = messages_provider(provider, "claude-sonnet-4-5").unwrap();
assert_eq!(provider.as_str(), name);
}

821
litellm-rust/crates/inference-messages/src/handler.rs
View file

@ -1,148 +1,271 @@
use std::time::Duration;
use bytes::Bytes;
use futures_util::{StreamExt, TryStreamExt, stream::BoxStream};
use litellm_host::interceptors::{Interceptors, ProviderIdentity, RequestContext, WireRequest};
use litellm_http::transport::Error as TransportError;
use futures_util::{FutureExt, StreamExt, TryStreamExt, stream::BoxStream};
use litellm_auth::{AuthServices, SecretValue};
use litellm_core_utils::{
dot_notation_indexing::delete_nested_value,
get_provider_specific_headers::get_provider_specific_headers, settings::Lookup,
};
use litellm_host::interceptors::{
Cost, ExecutionFacts, Interceptors, ProviderIdentity, RequestContext, ResultSource, WireRequest,
};
use litellm_http::{request::with_default_headers, transport::Error as TransportError};
use litellm_llms::base_llm::{
auth::{Authenticated, resolve_auth},
auth::{Authenticated, ValidatedEnvironment, resolve_auth},
messages::{
context::MessagesTransformContext,
streaming::{self, ByteStream, StreamDecoder},
transformation::BaseMessagesConfig,
},
};
use litellm_llms_types::formats::messages::MessagesResponse;
use litellm_llms_types::formats::messages::{MessagesRequest, MessagesResponse};
use litellm_secrets::source::Secrets;
use litellm_tracing::ByteChunk;
use serde_json::Value;
use super::{
Error, MessagesCallResponse, MessagesRoute, common_utils::truncate_error_body,
prepare::ProviderMessagesRequest,
Error, MessagesCall, MessagesCallResponse, MessagesRoute,
common_utils::{MessagesProvider, string_headers, truncate_error_body},
constants::MESSAGES_TIMEOUT_SECS,
route::Messages,
types::invalid_request,
};
use crate::constants::MESSAGES_TIMEOUT_SECS;
use litellm_inference::{context::CallContext, outbound::outbound_request};
use litellm_inference::{caching::CallCache, context::CallContext, outbound::outbound_request};
pub(super) struct ProviderCall {
pub identity: ProviderIdentity,
pub wire: WireRequest,
pub(super) provider: super::common_utils::MessagesProvider,
struct ProviderCall {
identity: ProviderIdentity,
wire: WireRequest,
provider: MessagesProvider,
signer: Option<litellm_auth_aws::SigV4Signer>,
timeout: Option<Duration>,
stream: bool,
}
impl ProviderCall {
pub fn cacheable(&self) -> bool {
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: provider
.config()
.wire_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 execute(
route: &MessagesRoute,
provider: MessagesProvider,
model: String,
call: MessagesCall,
secrets: Secrets,
context: &CallContext<'_, impl Interceptors<Error>>,
) -> Result<MessagesCallResponse, Error> {
let request = provider_call(
&route.auth,
provider,
model,
call,
secrets.as_ref(),
context,
)
.boxed()
.await?;
let cache = CallCache::<Messages>::from_wire(
route.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 => (
call_provider(&route.http, request, context).await?,
ResultSource::Provider,
),
};
let reported = match &output {
MessagesCallResponse::Complete(response) => provider.config().reported_cost(response),
MessagesCallResponse::Stream { .. } => None,
};
context
.result_ready(ExecutionFacts {
provider: identity,
cost: Cost::settle(&source, reported),
source: source.clone(),
})
}
.await?;
Ok(cache.finish(output, &source).await)
}
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,
async fn provider_call(
auth: &AuthServices,
provider: MessagesProvider,
model: String,
call: MessagesCall,
secrets: &(dyn Lookup + Send + Sync),
context: &CallContext<'_, impl Interceptors<Error>>,
) -> Result<ProviderCall, Error> {
let MessagesCall {
body,
api_key,
api_base,
litellm_params,
extra_headers,
provider_specific_header,
timeout,
shaping,
..
} = call;
let connection_params = litellm_router_types::LitellmParams {
api_base: api_base.clone().or(litellm_params.api_base.clone()),
..litellm_params
};
let config = provider.config();
let env_lookup = |key: &str| secrets.get(key);
let sanitized = config.shape_request(
MessagesRequest { model, ..body },
shaping.settings.reasoning_auto_summary,
)?;
let trimmed =
without_additional_drop_params(sanitized, &shaping.settings.additional_drop_params)?;
let scoped =
get_provider_specific_headers(provider_specific_header.as_ref(), provider.as_str());
let forwarded = string_headers(Some(
extra_headers.into_iter().flatten().chain(scoped).collect(),
))?;
let validated = config.validate_environment(
forwarded,
api_key.as_deref(),
&trimmed.model,
&trimmed.params,
&connection_params,
&env_lookup,
)?;
let body = config.transform_anthropic_messages_request(
trimmed,
&MessagesTransformContext::new(shaping.capabilities, shaping.settings.drop_params),
)?;
let environment = ValidatedEnvironment {
headers: config.request_headers(
with_default_headers(validated.headers, config.default_headers()),
&body,
),
auth: validated.auth,
};
let url = config.get_complete_url(
api_base.as_deref(),
&body.model,
&connection_params,
body.params.stream == Some(true),
&env_lookup,
)?;
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: api_key.map(SecretValue::new),
};
let authenticated = resolve_auth(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: config.wire_body(serde_json::to_value(&body).map_err(serialize_failure)?),
},
&wire.url,
&wire.body,
timeout,
request_context,
)
.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 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(),
},
));
}
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)))
};
Ok(ProviderCall {
identity,
wire,
provider,
signer: authenticated.signer,
timeout: timeout.or_else(|| shaping.settings.request_timeout()),
stream,
})
}
fn without_additional_drop_params(
request: MessagesRequest,
paths: &[String],
) -> Result<MessagesRequest, Error> {
if paths.is_empty() {
return Ok(request);
}
let params = serde_json::to_value(request.params).map_err(invalid_request)?;
let trimmed = paths
.iter()
.fold(params, |params, path| delete_nested_value(params, path));
Ok(MessagesRequest {
params: serde_json::from_value(trimmed).map_err(invalid_request)?,
..request
})
}
async fn call_provider(
http: &litellm_http::Client,
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(
http,
Authenticated {
headers: wire.headers,
signer,
},
&wire.url,
&wire.body,
timeout,
)
.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)))
}
fn serialize_failure(err: serde_json::Error) -> Error {
@ -271,10 +394,482 @@ fn log_chunk(provider: &str, stage: &str, data: &bytes::Bytes) {
#[cfg(test)]
mod tests {
use rstest::rstest;
use litellm_llms_types::headers::ProviderSpecificHeaders;
use rstest::{fixture, rstest};
use serde_json::{Map, Value, json};
use wiremock::{Mock, MockServer, ResponseTemplate, matchers::any};
use super::*;
use crate::{MessagesSettings, MessagesShaping, common_utils::messages_provider};
use litellm_inference::provider::resolve_llm_provider;
#[fixture]
fn shaping() -> MessagesShaping {
MessagesShaping::default()
}
fn body(value: Value) -> MessagesRequest {
serde_json::from_value(value).unwrap()
}
async fn provider_call_for(call: MessagesCall) -> Result<ProviderCall, Error> {
provider_call_with_secrets(call, &|_: &str| None).await
}
async fn provider_call_with_secrets(
call: MessagesCall,
secrets: &(dyn Lookup + Send + Sync),
) -> Result<ProviderCall, Error> {
let provider = resolve_llm_provider(
&call.body.model,
call.custom_llm_provider.as_deref(),
"messages",
)?;
let config = messages_provider(provider.provider, provider.model)
.ok_or_else(|| Error::InvalidProvider(<&str>::from(provider.provider).to_string()))?;
let model = provider.model.to_owned();
let interceptors = ();
let context = CallContext::new(&interceptors, litellm_inference::CallOptions::default());
provider_call(
&AuthServices::default(),
config,
model,
call,
secrets,
&context,
)
.await
}
async fn wire_body(fields: Value, shaping: MessagesShaping) -> Result<Value, Error> {
provider_call_for(MessagesCall {
body: body(fields),
api_key: Some("sk-test".into()),
api_base: Some("https://anthropic.test".into()),
custom_llm_provider: Some("anthropic".into()),
litellm_params: Default::default(),
extra_headers: None,
provider_specific_header: None,
timeout: None,
shaping,
})
.await
.map(|call| call.wire.body)
}
#[rstest]
#[case::api_key(
&[("ANTHROPIC_API_KEY", "sk-secret")],
&[("x-api-key", "sk-secret")],
"https://api.anthropic.com/v1/messages"
)]
#[case::auth_token(
&[("ANTHROPIC_AUTH_TOKEN", "token")],
&[("authorization", "Bearer token")],
"https://api.anthropic.com/v1/messages"
)]
#[case::api_base(
&[("ANTHROPIC_API_KEY", "sk-secret"), ("ANTHROPIC_API_BASE", "https://gateway.test")],
&[("x-api-key", "sk-secret")],
"https://gateway.test/v1/messages"
)]
#[case::sdk_base_url(
&[("ANTHROPIC_API_KEY", "sk-secret"), ("ANTHROPIC_BASE_URL", "https://sdk.test")],
&[("x-api-key", "sk-secret")],
"https://sdk.test/v1/messages"
)]
#[tokio::test]
async fn credentials_and_base_come_from_the_resolved_secrets(
shaping: MessagesShaping,
#[case] secrets: &[(&str, &str)],
#[case] expected_auth: &[(&str, &str)],
#[case] expected_url: &str,
) {
let lookup = |name: &str| {
secrets
.iter()
.find(|(key, _)| *key == name)
.map(|(_, value)| value.to_string())
};
let call = provider_call_with_secrets(
MessagesCall {
body: body(
json!({"model": "claude-test", "messages": [{"role": "user", "content": "hi"}], "max_tokens": 16}),
),
api_key: None,
api_base: None,
custom_llm_provider: Some("anthropic".into()),
litellm_params: Default::default(),
extra_headers: None,
provider_specific_header: None,
timeout: None,
shaping,
},
&lookup,
)
.await
.unwrap();
let auth: Vec<(&str, &str)> = call
.wire
.headers
.iter()
.filter(|(name, _)| matches!(name.as_str(), "x-api-key" | "authorization"))
.map(|(name, value)| (name.as_str(), value.as_str()))
.collect();
assert_eq!(
(auth.as_slice(), call.wire.url.as_str()),
(expected_auth, expected_url)
);
}
#[rstest]
#[tokio::test]
async fn missing_credentials_are_reported_before_the_request_is_transformed(
shaping: MessagesShaping,
) {
let result = provider_call_for(MessagesCall {
body: body(
json!({"model": "claude-test", "messages": [{"role": "user", "content": "hi"}]}),
),
api_key: None,
api_base: None,
custom_llm_provider: Some("anthropic".into()),
litellm_params: Default::default(),
extra_headers: None,
provider_specific_header: None,
timeout: None,
shaping,
})
.await;
assert!(
matches!(
result,
Err(Error::Auth(litellm_auth::Error::MissingApiKey { .. }))
),
"{:?}",
result.err()
);
}
#[rstest]
#[case::nothing_configured(None, None, None)]
#[case::configured_setting(None, Some(1800.0), Some(Duration::from_secs(1800)))]
#[case::call_timeout_wins(
Some(Duration::from_secs(30)),
Some(1800.0),
Some(Duration::from_secs(30))
)]
#[case::call_timeout_only(Some(Duration::from_secs(30)), None, Some(Duration::from_secs(30)))]
#[case::unusable_setting_ignored(None, Some(-1.0), None)]
#[case::zero_setting_ignored(None, Some(0.0), None)]
#[tokio::test]
async fn timeout_falls_back_to_the_configured_request_timeout(
shaping: MessagesShaping,
#[case] call_timeout: Option<Duration>,
#[case] request_timeout: Option<f64>,
#[case] expected: Option<Duration>,
) {
let call = provider_call_for(MessagesCall {
body: body(json!({"model": "claude-test", "messages": [{"role": "user", "content": "hi"}], "max_tokens": 16})),
api_key: Some("sk-test".into()),
api_base: Some("https://anthropic.test".into()),
custom_llm_provider: Some("anthropic".into()),
litellm_params: Default::default(),
extra_headers: None,
provider_specific_header: None,
timeout: call_timeout,
shaping: MessagesShaping {
settings: MessagesSettings { request_timeout, ..shaping.settings },
..shaping
},
})
.await
.unwrap();
assert_eq!(call.timeout, expected);
}
#[rstest]
#[case::top_level_and_nested_paths(
json!({
"max_tokens": 1024,
"thinking": {"type": "enabled", "budget_tokens": 2048},
"context_management": {"edits": [{"type": "clear_thinking_20251015"}]},
"metadata": {"user_id": "u1"},
"tools": [{"name": "lookup", "input_schema": {"type": "object"}, "input_examples": [{"q": "x"}]}]
}),
&["thinking", "context_management", "tools[*].input_examples"],
json!({
"max_tokens": 1024,
"metadata": {"user_id": "u1"},
"tools": [{"name": "lookup", "input_schema": {"type": "object"}}]
}),
)]
#[case::no_paths(
json!({"max_tokens": 16, "safeguards": [{"type": "dangerous_tool_use"}]}),
&[],
json!({"max_tokens": 16, "safeguards": [{"type": "dangerous_tool_use"}]}),
)]
#[case::model_and_messages_are_never_dropped(
json!({"max_tokens": 16}),
&["model", "messages", "messages[0].content"],
json!({"max_tokens": 16}),
)]
#[tokio::test]
async fn wire_body_drops_configured_paths(
shaping: MessagesShaping,
#[case] fields: Value,
#[case] additional_drop_params: &[&str],
#[case] expected_fields: Value,
) {
let with_messages = |fields: Value| -> Value {
let Value::Object(fields) = fields else {
unreachable!()
};
Value::Object(
[
("model".to_string(), json!("claude-test")),
(
"messages".to_string(),
json!([{"role": "user", "content": "hi"}]),
),
]
.into_iter()
.chain(fields)
.collect(),
)
};
let shaping = MessagesShaping {
settings: MessagesSettings {
additional_drop_params: additional_drop_params
.iter()
.map(ToString::to_string)
.collect(),
..shaping.settings
},
..shaping
};
assert_eq!(
wire_body(with_messages(fields), shaping).await,
Ok(with_messages(expected_fields))
);
}
#[rstest]
#[case::model_prefix_picks_the_provider(
"azure_ai/claude-test",
None,
&[("x-priority", "extra"), ("x-scoped", "azure_ai")]
)]
#[case::explicit_provider(
"claude-test",
Some("anthropic"),
&[("x-priority", "scoped"), ("x-scoped", "anthropic")]
)]
#[case::provider_prefix_on_an_anthropic_model(
"anthropic/claude-test",
None,
&[("x-priority", "scoped"), ("x-scoped", "anthropic")]
)]
#[tokio::test]
async fn provider_specific_headers_follow_the_resolved_provider(
shaping: MessagesShaping,
#[case] model: &str,
#[case] custom_llm_provider: Option<&str>,
#[case] expected: &[(&str, &str)],
) {
let configured: ProviderSpecificHeaders = serde_json::from_value(json!([
{"custom_llm_provider": "azure_ai", "extra_headers": {"x-scoped": "azure_ai"}},
{"custom_llm_provider": "anthropic", "extra_headers": {"x-scoped": "anthropic", "x-priority": "scoped"}}
]))
.unwrap();
let call = provider_call_for(MessagesCall {
body: body(
json!({"model": model, "messages": [{"role": "user", "content": "hi"}], "max_tokens": 16}),
),
api_key: Some("sk-test".into()),
api_base: Some("https://resource.services.ai.azure.com".into()),
custom_llm_provider: custom_llm_provider.map(Into::into),
litellm_params: Default::default(),
extra_headers: Some(Map::from_iter([("x-priority".into(), json!("extra"))])),
provider_specific_header: Some(configured),
timeout: None,
shaping,
})
.await
.unwrap();
let caller_headers: Vec<(&str, &str)> = call
.wire
.headers
.iter()
.filter(|(name, _)| matches!(name.as_str(), "x-priority" | "x-scoped"))
.map(|(name, value)| (name.as_str(), value.as_str()))
.collect();
assert_eq!(caller_headers, expected);
}
#[rstest]
#[tokio::test]
async fn wire_body_carries_the_provider_stripped_model(shaping: MessagesShaping) {
assert_eq!(
wire_body(
json!({
"model": "anthropic/claude-test",
"messages": [{"role": "user", "content": "hi"}],
"max_tokens": 16
}),
shaping,
)
.await,
Ok(json!({
"model": "claude-test",
"messages": [{"role": "user", "content": "hi"}],
"max_tokens": 16
}))
);
}
#[rstest]
#[tokio::test]
async fn dropped_thinking_display_is_not_restored_by_auto_summary(shaping: MessagesShaping) {
let shaping = MessagesShaping {
settings: MessagesSettings {
reasoning_auto_summary: true,
additional_drop_params: vec!["thinking.display".to_string()],
..shaping.settings
},
..shaping
};
assert_eq!(
wire_body(
json!({
"model": "claude-test",
"messages": [{"role": "user", "content": "hi"}],
"max_tokens": 4096,
"thinking": {"type": "enabled", "budget_tokens": 2048}
}),
shaping,
)
.await,
Ok(json!({
"model": "claude-test",
"messages": [{"role": "user", "content": "hi"}],
"max_tokens": 4096,
"thinking": {"type": "enabled", "budget_tokens": 2048}
}))
);
}
#[rstest]
#[tokio::test]
async fn dropping_an_invalid_metadata_user_id_does_not_skip_its_validation(
shaping: MessagesShaping,
) {
let shaping = MessagesShaping {
settings: MessagesSettings {
additional_drop_params: vec!["metadata.user_id".to_string()],
..shaping.settings
},
..shaping
};
assert!(matches!(
wire_body(
json!({
"model": "claude-test",
"messages": [{"role": "user", "content": "hi"}],
"max_tokens": 16,
"metadata": {"user_id": 123}
}),
shaping,
)
.await,
Err(Error::InvalidRequest(_))
));
}
#[rstest]
#[tokio::test]
async fn wire_body_rejects_invalid_metadata_before_the_call(shaping: MessagesShaping) {
assert_eq!(
wire_body(
json!({
"model": "claude-test",
"messages": [{"role": "user", "content": "hi"}],
"max_tokens": 16,
"metadata": {"user_id": 123}
}),
shaping,
)
.await,
Err(Error::InvalidRequest(
litellm_llms::ErrorDetail::InvalidValue {
field: "metadata.user_id",
expected: "a string",
actual: json!(123),
}
))
);
}
#[rstest]
#[case::mantle(
"bedrock_mantle/claude-test",
"https://bedrock-mantle.us-west-2.api.aws",
"/anthropic/v1/messages",
"bedrock"
)]
#[case::platform(
"bedrock/claude_platform/claude-test",
"https://aws-external-anthropic.us-west-2.api.aws",
"/v1/messages",
"aws-external-anthropic"
)]
#[tokio::test]
async fn explicit_endpoint_region_is_used_by_the_signer(
shaping: MessagesShaping,
#[case] model: &str,
#[case] base: &str,
#[case] path: &str,
#[case] service: &str,
) {
let call = provider_call_with_secrets(MessagesCall {
body: body(json!({"model": model, "max_tokens": 16, "messages": [{"role": "user", "content": "hi"}]})),
api_key: None,
api_base: Some(base.into()),
custom_llm_provider: None,
litellm_params: crate::litellm_params(Map::from_iter([
("model".into(), json!(model)),
("aws_access_key_id".into(), json!("AKIDEXAMPLE")),
("aws_secret_access_key".into(), json!("secret")),
]))
.unwrap(),
extra_headers: None,
provider_specific_header: None,
timeout: None,
shaping,
}, &|name: &str| match name {
"AWS_REGION_NAME" => Some("us-east-1".into()),
"ANTHROPIC_AWS_WORKSPACE_ID" => Some("workspace".into()),
_ => None,
})
.await
.unwrap();
assert_eq!(call.wire.url, format!("{base}{path}"));
let signed = outbound_request(
Authenticated {
headers: call.wire.headers,
signer: call.signer,
},
call.wire.url,
&call.wire.body,
None,
)
.unwrap();
let authorization = signed.header("authorization").unwrap_or_default();
assert!(
authorization.contains(&format!("/us-west-2/{service}/aws4_request")),
"{authorization}"
);
}
#[rstest]
#[case::event(

55
litellm-rust/crates/inference-messages/src/lib.rs
View file

@ -1,18 +1,19 @@
mod common_utils;
mod constants;
mod handler;
mod prepare;
pub mod route;
mod types;
use futures_util::FutureExt;
use litellm_auth::AuthServices;
use litellm_host::interceptors::{Cost, ExecutionFacts, Interceptors, ResultSource};
use litellm_inference::{caching::CallCache, context::CallContext};
use litellm_host::interceptors::Interceptors;
use litellm_inference::context::CallContext;
use litellm_inference::provider::resolve_llm_provider;
use litellm_secrets::source::SecretSource;
use std::sync::Arc;
use crate::common_utils::messages_provider;
pub use litellm_inference::RouteError as Error;
pub use types::{
MessagesCall, MessagesCallResponse, MessagesSettings, MessagesShaping, litellm_params,
@ -74,39 +75,23 @@ impl MessagesRoute {
context: CallContext<'_, impl Interceptors<Error>>,
) -> Result<MessagesCallResponse, Error> {
litellm_inference::diagnostic::call(async {
let prepared = prepare::prepare(call, self.secrets.as_ref()).await?;
let provider = resolve_llm_provider(
&call.body.model,
call.custom_llm_provider.as_deref(),
"messages",
)?;
let config = messages_provider(provider.provider, provider.model).ok_or_else(|| {
Error::InvalidProvider(<&str>::from(provider.provider).to_string())
})?;
litellm_inference::diagnostic::provider(
&prepared.body.model,
prepared.provider.as_str(),
provider.model,
<&'static str>::from(provider.provider),
);
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 config = request.provider.config();
let (output, source) = match cache.lookup().await {
Some(hit) => hit,
None => (
self.call_provider(request, &context).await?,
ResultSource::Provider,
),
};
let reported = match &output {
MessagesCallResponse::Complete(response) => config.reported_cost(response),
MessagesCallResponse::Stream { .. } => None,
};
context
.result_ready(ExecutionFacts {
provider: identity,
cost: Cost::settle(&source, reported),
source: source.clone(),
})
.await?;
Ok(cache.finish(output, &source).await)
let model = provider.model.to_owned();
let secrets = self.secrets.resolve(config.config().secret_names()).await?;
handler::execute(self, config, model, call, secrets, &context)
.boxed()
.await
})
.await
}

554
litellm-rust/crates/inference-messages/src/prepare.rs
View file

@ -1,554 +0,0 @@
use std::time::Duration;
use litellm_auth::SecretValue;
use litellm_core_utils::{
dot_notation_indexing::delete_nested_value,
get_provider_specific_headers::get_provider_specific_headers, settings::Lookup,
};
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 super::{
Error, MessagesCall,
common_utils::{MessagesProvider, messages_provider, string_headers},
types::invalid_request,
};
use litellm_inference::provider::resolve_llm_provider;
struct ResolvedProvider {
model: String,
provider: MessagesProvider,
}
pub(super) struct ProviderMessagesRequest {
pub(super) provider: MessagesProvider,
pub(super) url: String,
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.
pub(super) api_key: Option<SecretValue>,
}
#[tracing::instrument(name = "litellm.prepare", level = "debug", skip_all)]
pub(super) async fn prepare(
call: MessagesCall,
secrets: &dyn SecretSource,
) -> Result<ProviderMessagesRequest, Error> {
let resolved = resolve_provider(&call.body.model, call.custom_llm_provider.as_deref())?;
let secrets = secrets
.resolve(resolved.provider.config().secret_names())
.await?;
prepare_provider_request(call, resolved, secrets.as_ref())
}
fn resolve_provider(
model: &str,
custom_llm_provider: Option<&str>,
) -> Result<ResolvedProvider, Error> {
let resolved = resolve_llm_provider(model, custom_llm_provider, "messages")?;
let provider = messages_provider(resolved.provider, resolved.model)
.ok_or_else(|| Error::InvalidProvider(<&str>::from(resolved.provider).to_string()))?;
Ok(ResolvedProvider {
model: resolved.model.to_string(),
provider,
})
}
fn prepare_provider_request(
call: MessagesCall,
resolved: ResolvedProvider,
secrets: &dyn Lookup,
) -> Result<ProviderMessagesRequest, Error> {
let ResolvedProvider { model, provider } = resolved;
let MessagesCall {
body,
api_key,
api_base,
litellm_params,
extra_headers,
provider_specific_header,
timeout,
shaping,
..
} = call;
let connection_params = litellm_router_types::LitellmParams {
api_base: api_base.clone().or(litellm_params.api_base.clone()),
..litellm_params
};
let config = provider.config();
let env_lookup = |key: &str| secrets.get(key);
let sanitized = config.shape_request(
MessagesRequest { model, ..body },
shaping.settings.reasoning_auto_summary,
)?;
let trimmed =
without_additional_drop_params(sanitized, &shaping.settings.additional_drop_params)?;
let transformed = config.transform_anthropic_messages_request(
trimmed,
&MessagesTransformContext::new(shaping.capabilities, shaping.settings.drop_params),
)?;
let scoped =
get_provider_specific_headers(provider_specific_header.as_ref(), provider.as_str());
let forwarded = string_headers(Some(
extra_headers.into_iter().flatten().chain(scoped).collect(),
))?;
let forwarded = config.request_headers(
with_default_headers(forwarded, config.default_headers()),
&transformed,
);
let environment = config.validate_environment(
forwarded,
api_key.as_deref(),
&transformed.model,
&connection_params,
&env_lookup,
)?;
let url = config.get_complete_url(
api_base.as_deref(),
&transformed.model,
&connection_params,
transformed.params.stream == Some(true),
&env_lookup,
)?;
Ok(ProviderMessagesRequest {
provider,
url,
body: transformed,
environment,
timeout: timeout.or_else(|| shaping.settings.request_timeout()),
api_key: api_key.map(SecretValue::new),
})
}
fn without_additional_drop_params(
request: MessagesRequest,
paths: &[String],
) -> Result<MessagesRequest, Error> {
if paths.is_empty() {
return Ok(request);
}
let params = serde_json::to_value(request.params).map_err(invalid_request)?;
let trimmed = paths
.iter()
.fold(params, |params, path| delete_nested_value(params, path));
Ok(MessagesRequest {
params: serde_json::from_value(trimmed).map_err(invalid_request)?,
..request
})
}
#[cfg(test)]
mod tests {
use litellm_llms::base_llm::auth::resolve_auth;
use litellm_llms_types::headers::ProviderSpecificHeaders;
use rstest::{fixture, rstest};
use serde_json::{Map, Value, json};
use super::*;
use crate::{MessagesSettings, MessagesShaping};
#[fixture]
fn shaping() -> MessagesShaping {
MessagesShaping::default()
}
fn body(value: Value) -> MessagesRequest {
serde_json::from_value(value).unwrap()
}
fn prepare(call: MessagesCall) -> Result<ProviderMessagesRequest, Error> {
prepare_with_secrets(call, &|_: &str| None)
}
fn prepare_with_secrets(
call: MessagesCall,
secrets: &dyn Lookup,
) -> Result<ProviderMessagesRequest, Error> {
let resolved = resolve_provider(&call.body.model, call.custom_llm_provider.as_deref())?;
prepare_provider_request(call, resolved, secrets)
}
/// The headers as they go on the wire, credential applied.
fn wire_headers(prepared: &ProviderMessagesRequest) -> Vec<(String, String)> {
tokio::runtime::Builder::new_current_thread()
.build()
.unwrap()
.block_on(resolve_auth(
&litellm_auth::AuthServices::default(),
prepared.environment.clone(),
&|_| None,
))
.unwrap()
.headers
}
#[rstest]
#[case::api_key(
&[("ANTHROPIC_API_KEY", "sk-secret")],
&[("x-api-key", "sk-secret")],
"https://api.anthropic.com/v1/messages"
)]
#[case::auth_token(
&[("ANTHROPIC_AUTH_TOKEN", "token")],
&[("authorization", "Bearer token")],
"https://api.anthropic.com/v1/messages"
)]
#[case::api_base(
&[("ANTHROPIC_API_KEY", "sk-secret"), ("ANTHROPIC_API_BASE", "https://gateway.test")],
&[("x-api-key", "sk-secret")],
"https://gateway.test/v1/messages"
)]
#[case::sdk_base_url(
&[("ANTHROPIC_API_KEY", "sk-secret"), ("ANTHROPIC_BASE_URL", "https://sdk.test")],
&[("x-api-key", "sk-secret")],
"https://sdk.test/v1/messages"
)]
fn credentials_and_base_come_from_the_resolved_secrets(
shaping: MessagesShaping,
#[case] secrets: &[(&str, &str)],
#[case] expected_auth: &[(&str, &str)],
#[case] expected_url: &str,
) {
let lookup = |name: &str| {
secrets
.iter()
.find(|(key, _)| *key == name)
.map(|(_, value)| value.to_string())
};
let prepared = prepare_with_secrets(
MessagesCall {
body: body(
json!({"model": "claude-test", "messages": [{"role": "user", "content": "hi"}], "max_tokens": 16}),
),
api_key: None,
api_base: None,
custom_llm_provider: Some("anthropic".into()),
litellm_params: Default::default(),
extra_headers: None,
provider_specific_header: None,
timeout: None,
shaping,
},
&lookup,
)
.unwrap();
let headers = wire_headers(&prepared);
let auth: Vec<(&str, &str)> = headers
.iter()
.filter(|(name, _)| matches!(name.as_str(), "x-api-key" | "authorization"))
.map(|(name, value)| (name.as_str(), value.as_str()))
.collect();
assert_eq!(
(auth.as_slice(), prepared.url.as_str()),
(expected_auth, expected_url)
);
}
#[rstest]
#[case::nothing_configured(None, None, None)]
#[case::configured_setting(None, Some(1800.0), Some(Duration::from_secs(1800)))]
#[case::call_timeout_wins(
Some(Duration::from_secs(30)),
Some(1800.0),
Some(Duration::from_secs(30))
)]
#[case::call_timeout_only(Some(Duration::from_secs(30)), None, Some(Duration::from_secs(30)))]
#[case::unusable_setting_ignored(None, Some(-1.0), None)]
#[case::zero_setting_ignored(None, Some(0.0), None)]
fn timeout_falls_back_to_the_configured_request_timeout(
shaping: MessagesShaping,
#[case] call_timeout: Option<Duration>,
#[case] request_timeout: Option<f64>,
#[case] expected: Option<Duration>,
) {
let prepared = prepare(MessagesCall {
body: body(json!({"model": "claude-test", "messages": [{"role": "user", "content": "hi"}], "max_tokens": 16})),
api_key: Some("sk-test".into()),
api_base: Some("https://anthropic.test".into()),
custom_llm_provider: Some("anthropic".into()),
litellm_params: Default::default(),
extra_headers: None,
provider_specific_header: None,
timeout: call_timeout,
shaping: MessagesShaping {
settings: MessagesSettings { request_timeout, ..shaping.settings },
..shaping
},
})
.unwrap();
assert_eq!(prepared.timeout, expected);
}
fn prepared_body(fields: Value, shaping: MessagesShaping) -> Result<Value, Error> {
prepare(MessagesCall {
body: body(fields),
api_key: Some("sk-test".into()),
api_base: Some("https://anthropic.test".into()),
custom_llm_provider: Some("anthropic".into()),
litellm_params: Default::default(),
extra_headers: None,
provider_specific_header: None,
timeout: None,
shaping,
})
.map(|prepared| serde_json::to_value(prepared.body).unwrap())
}
#[rstest]
#[case::top_level_and_nested_paths(
json!({
"max_tokens": 1024,
"thinking": {"type": "enabled", "budget_tokens": 2048},
"context_management": {"edits": [{"type": "clear_thinking_20251015"}]},
"metadata": {"user_id": "u1"},
"tools": [{"name": "lookup", "input_schema": {"type": "object"}, "input_examples": [{"q": "x"}]}]
}),
&["thinking", "context_management", "tools[*].input_examples"],
json!({
"max_tokens": 1024,
"metadata": {"user_id": "u1"},
"tools": [{"name": "lookup", "input_schema": {"type": "object"}}]
}),
)]
#[case::no_paths(
json!({"max_tokens": 16, "safeguards": [{"type": "dangerous_tool_use"}]}),
&[],
json!({"max_tokens": 16, "safeguards": [{"type": "dangerous_tool_use"}]}),
)]
#[case::model_and_messages_are_never_dropped(
json!({"max_tokens": 16}),
&["model", "messages", "messages[0].content"],
json!({"max_tokens": 16}),
)]
fn prepared_body_drops_configured_paths(
shaping: MessagesShaping,
#[case] fields: Value,
#[case] additional_drop_params: &[&str],
#[case] expected_fields: Value,
) {
let with_messages = |fields: Value| -> Value {
let Value::Object(fields) = fields else {
unreachable!()
};
Value::Object(
[
("model".to_string(), json!("claude-test")),
(
"messages".to_string(),
json!([{"role": "user", "content": "hi"}]),
),
]
.into_iter()
.chain(fields)
.collect(),
)
};
let shaping = MessagesShaping {
settings: MessagesSettings {
additional_drop_params: additional_drop_params
.iter()
.map(ToString::to_string)
.collect(),
..shaping.settings
},
..shaping
};
assert_eq!(
prepared_body(with_messages(fields), shaping),
Ok(with_messages(expected_fields))
);
}
#[rstest]
#[case::model_prefix_picks_the_provider(
"azure_ai/claude-test",
None,
&[("x-priority", "extra"), ("x-scoped", "azure_ai")]
)]
#[case::explicit_provider(
"claude-test",
Some("anthropic"),
&[("x-priority", "scoped"), ("x-scoped", "anthropic")]
)]
#[case::provider_prefix_on_an_anthropic_model(
"anthropic/claude-test",
None,
&[("x-priority", "scoped"), ("x-scoped", "anthropic")]
)]
fn provider_specific_headers_follow_the_resolved_provider(
shaping: MessagesShaping,
#[case] model: &str,
#[case] custom_llm_provider: Option<&str>,
#[case] expected: &[(&str, &str)],
) {
let configured: ProviderSpecificHeaders = serde_json::from_value(json!([
{"custom_llm_provider": "azure_ai", "extra_headers": {"x-scoped": "azure_ai"}},
{"custom_llm_provider": "anthropic", "extra_headers": {"x-scoped": "anthropic", "x-priority": "scoped"}}
]))
.unwrap();
let prepared = prepare(MessagesCall {
body: body(
json!({"model": model, "messages": [{"role": "user", "content": "hi"}], "max_tokens": 16}),
),
api_key: Some("sk-test".into()),
api_base: Some("https://resource.services.ai.azure.com".into()),
custom_llm_provider: custom_llm_provider.map(Into::into),
litellm_params: Default::default(),
extra_headers: Some(Map::from_iter([("x-priority".into(), json!("extra"))])),
provider_specific_header: Some(configured),
timeout: None,
shaping,
})
.unwrap();
let caller_headers: Vec<(&str, &str)> = prepared
.environment
.headers
.iter()
.filter(|(name, _)| matches!(name.as_str(), "x-priority" | "x-scoped"))
.map(|(name, value)| (name.as_str(), value.as_str()))
.collect();
assert_eq!(caller_headers, expected);
}
#[rstest]
fn prepared_body_carries_the_provider_stripped_model(shaping: MessagesShaping) {
assert_eq!(
prepared_body(
json!({
"model": "anthropic/claude-test",
"messages": [{"role": "user", "content": "hi"}],
"max_tokens": 16
}),
shaping,
),
Ok(json!({
"model": "claude-test",
"messages": [{"role": "user", "content": "hi"}],
"max_tokens": 16
}))
);
}
#[rstest]
fn dropped_thinking_display_is_not_restored_by_auto_summary(shaping: MessagesShaping) {
let shaping = MessagesShaping {
settings: MessagesSettings {
reasoning_auto_summary: true,
additional_drop_params: vec!["thinking.display".to_string()],
..shaping.settings
},
..shaping
};
assert_eq!(
prepared_body(
json!({
"model": "claude-test",
"messages": [{"role": "user", "content": "hi"}],
"max_tokens": 4096,
"thinking": {"type": "enabled", "budget_tokens": 2048}
}),
shaping,
),
Ok(json!({
"model": "claude-test",
"messages": [{"role": "user", "content": "hi"}],
"max_tokens": 4096,
"thinking": {"type": "enabled", "budget_tokens": 2048}
}))
);
}
#[rstest]
fn dropping_an_invalid_metadata_user_id_does_not_skip_its_validation(shaping: MessagesShaping) {
let shaping = MessagesShaping {
settings: MessagesSettings {
additional_drop_params: vec!["metadata.user_id".to_string()],
..shaping.settings
},
..shaping
};
assert!(matches!(
prepared_body(
json!({
"model": "claude-test",
"messages": [{"role": "user", "content": "hi"}],
"max_tokens": 16,
"metadata": {"user_id": 123}
}),
shaping,
),
Err(Error::InvalidRequest(_))
));
}
#[rstest]
fn prepared_body_rejects_invalid_metadata_before_the_call(shaping: MessagesShaping) {
assert_eq!(
prepared_body(
json!({
"model": "claude-test",
"messages": [{"role": "user", "content": "hi"}],
"max_tokens": 16,
"metadata": {"user_id": 123}
}),
shaping,
),
Err(Error::InvalidRequest(
litellm_llms::ErrorDetail::InvalidValue {
field: "metadata.user_id",
expected: "a string",
actual: json!(123),
}
))
);
}
#[rstest]
#[case::mantle(
"bedrock_mantle/claude-test",
"https://bedrock-mantle.us-west-2.api.aws",
"/anthropic/v1/messages",
"bedrock"
)]
#[case::platform(
"bedrock/claude_platform/claude-test",
"https://aws-external-anthropic.us-west-2.api.aws",
"/v1/messages",
"aws-external-anthropic"
)]
fn explicit_endpoint_region_is_used_by_the_signer(
shaping: MessagesShaping,
#[case] model: &str,
#[case] base: &str,
#[case] path: &str,
#[case] service: &str,
) {
let prepared = prepare_with_secrets(MessagesCall {
body: body(json!({"model": model, "max_tokens": 16, "messages": [{"role": "user", "content": "hi"}]})),
api_key: None,
api_base: Some(base.into()),
custom_llm_provider: None,
litellm_params: Default::default(),
extra_headers: None,
provider_specific_header: None,
timeout: None,
shaping,
}, &|name: &str| match name {
"AWS_REGION_NAME" => Some("us-east-1".into()),
"ANTHROPIC_AWS_WORKSPACE_ID" => Some("workspace".into()),
_ => None,
}).unwrap();
assert_eq!(prepared.url, format!("{base}{path}"));
assert!(
matches!(prepared.environment.auth, litellm_llms::base_llm::auth::AuthScheme::AwsSigV4 { region, service: actual, .. } if region == "us-west-2" && actual == service)
);
}
}

4
litellm-rust/crates/inference-ocr/src/arguments.rs
View file

@ -44,8 +44,8 @@ pub fn consumed_optional_param_names(
model: &str,
custom_llm_provider: Option<&str>,
) -> Result<Vec<&'static str>, Error> {
let (model, config) = resolve_provider_config(model, custom_llm_provider)?;
let provider_fields = config.get_supported_ocr_params(&model);
let (provider, config) = resolve_provider_config(model, custom_llm_provider)?;
let provider_fields = config.get_supported_ocr_params(provider.model);
let auth_fields: &[&str] = match config {
OcrConfigKind::AwsTextract | OcrConfigKind::AwsTextractAnalyze => AWS_AUTH_OPTION_FIELDS,
OcrConfigKind::AzureAi

81
litellm-rust/crates/inference-ocr/src/client.rs
View file

@ -1,81 +0,0 @@
use litellm_host::observation::ObservationSender;
use std::sync::Arc;
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,
types::{LiteLLMOcrRequest, OcrDocumentInput, ResolvedOcrRequest},
};
#[derive(Clone)]
pub struct OcrRoute {
client: OcrClient,
}
impl OcrRoute {
pub fn new(client: OcrClient) -> Self {
Self { client }
}
pub async fn execute(
&self,
request: LiteLLMOcrRequest,
interceptors: &impl Interceptors<Error>,
observers: Option<ObservationSender>,
) -> Result<LiteLLMOcrResponse, Error> {
litellm_host::lifecycle::observe_unary(
observers.clone(),
self.run(request, interceptors, observers.as_ref()),
)
.await
}
#[tracing::instrument(name = "litellm.route", skip_all, fields(
route = "ocr",
model = %request.model,
resolved_model = %request.model,
provider = <&str>::from(request.config.provider()),
stream = false,
outcome
))]
pub(super) async fn run(
&self,
request: LiteLLMOcrRequest,
interceptors: &impl Interceptors<Error>,
observers: Option<&ObservationSender>,
) -> Result<LiteLLMOcrResponse, Error> {
litellm_inference::diagnostic::unary(async {
let caller_document = matches!(&request.document, OcrDocumentInput::Document(_));
let prepared = prepare_request_document(request).await?;
let execute: futures_util::future::BoxFuture<'_, Result<LiteLLMOcrResponse, Error>> =
Box::pin(perform_ocr_request(
&self.client,
prepared,
interceptors,
caller_document,
observers,
));
execute.await
})
.await
}
}
#[tracing::instrument(name = "litellm.prepare", level = "debug", skip_all)]
async fn prepare_request_document(
request: LiteLLMOcrRequest<OcrDocumentInput>,
) -> Result<ResolvedOcrRequest, Error> {
if let OcrDocumentInput::Document(_) = &request.document {
return request.map_document(super::document::prepare_document);
}
let logger = litellm_tracing::Logger::current();
let span = tracing::Span::current();
tokio::task::spawn_blocking(move || {
logger.scope(|| span.in_scope(|| request.map_document(super::document::prepare_document)))
})
.await
.map_err(|error| Error::DocumentTask(Arc::new(error)))?
}

19
litellm-rust/crates/inference-ocr/src/document.rs
View file

@ -1,10 +1,10 @@
use std::{collections::BTreeMap as Map, io::Read, path::Path};
use std::{collections::BTreeMap as Map, io::Read, path::Path, sync::Arc};
use base64::{Engine, engine::general_purpose::STANDARD};
use litellm_llms::base_llm::ocr::{error::Error, transformation::OCR_INLINE_MAX_BYTES};
use litellm_llms_types::formats::ocr::OcrDocument;
use crate::types::OcrDocumentInput;
use crate::types::{LiteLLMOcrRequest, OcrDocumentInput, ResolvedOcrRequest};
pub fn prepare_document(input: OcrDocumentInput) -> Result<OcrDocument, Error> {
match input {
@ -24,6 +24,21 @@ pub fn prepare_document(input: OcrDocumentInput) -> Result<OcrDocument, Error> {
}
}
pub(crate) async fn resolve_document(
request: LiteLLMOcrRequest<OcrDocumentInput>,
) -> Result<ResolvedOcrRequest, Error> {
if let OcrDocumentInput::Document(_) = &request.document {
return request.map_document(prepare_document);
}
let logger = litellm_tracing::Logger::current();
let span = tracing::Span::current();
tokio::task::spawn_blocking(move || {
logger.scope(|| span.in_scope(|| request.map_document(prepare_document)))
})
.await
.map_err(|error| Error::DocumentTask(Arc::new(error)))?
}
pub fn read_path_document(path: &Path, mime_type: Option<&str>) -> Result<OcrDocument, Error> {
let mut bytes = Vec::new();
std::fs::File::open(path)

296
litellm-rust/crates/inference-ocr/src/handler.rs
View file

@ -4,29 +4,45 @@ use litellm_host::{lifecycle::ExecutionEvent, observation::ObservationSender};
use litellm_llms::base_llm::ocr::{
error::Error,
handler::{CallHooks, OcrClient},
transformation::PreparedOcrRequest,
transformation::{OcrConnection, PreparedOcrRequest},
};
use litellm_llms_types::formats::ocr::LiteLLMOcrResponse;
use litellm_secrets::source::Secrets;
use serde_json::Value;
use super::{arguments::is_secret_param, prepare::prepare_request, provider_config::OcrConfigKind};
use crate::types::ResolvedOcrRequest;
use super::{arguments::is_secret_param, provider_config::OcrConfigKind};
use crate::types::{LiteLLMOcrRequest, ResolvedOcrRequest};
pub(crate) async fn perform_ocr_request(
pub(crate) async fn execute(
client: &OcrClient,
request: ResolvedOcrRequest,
config: OcrConfigKind,
call: ResolvedOcrRequest,
secrets: Secrets,
host: &impl Interceptors<Error>,
caller_document: bool,
observers: Option<&ObservationSender>,
) -> Result<LiteLLMOcrResponse, Error> {
request.response_format()?;
let config = request.config;
let secrets = client
.secret_source()
.resolve(&config.secret_names())
.await
.map_err(|error| Error::Secret(std::sync::Arc::new(error)))?;
let request = prepare_request(request, caller_document, client, secrets);
call.response_format()?;
let LiteLLMOcrRequest {
model,
document,
credentials,
transport,
optional_params,
input_sources,
azure_ad_token_provider,
..
} = call;
let resolved = config.resolve_credentials(credentials, &secrets);
let request = PreparedOcrRequest {
model,
document,
connection: OcrConnection::new(resolved, transport, client.settings().clone(), secrets),
caller_document,
optional_params,
input_sources,
azure_ad_token_provider,
};
let interceptors = OcrCallHooks::new(host, &request, config, observers);
config.ocr(client, &request, &interceptors).await
}
@ -86,3 +102,257 @@ impl<H: Interceptors<Error>> CallHooks<Error> for OcrCallHooks<'_, H> {
Box::pin(self.interceptors.after_provider_response(raw))
}
}
#[cfg(test)]
mod tests {
use std::time::Duration;
use litellm_core_utils::call_arguments::{CallArguments, compose_body, parse_options};
use litellm_llms_types::formats::ocr::OcrResponseFormat;
use rstest::rstest;
use serde_json::{Value, json};
use wiremock::{Mock, MockServer, ResponseTemplate, matchers::any};
use super::*;
use crate::{
OcrRoute,
types::LiteLLMOcrRequest,
wire::{OcrWireRequest, decode_request},
};
fn client() -> OcrClient {
OcrClient::for_test(
litellm_http::Client::plain_for_test(),
litellm_http::Client::no_redirect_for_test(),
)
.with_secrets(litellm_inference_testing::no_secrets())
}
fn request(model: &str, base: &str, document: Value, options: Value) -> LiteLLMOcrRequest {
decode_request(OcrWireRequest {
model: model.into(),
document,
api_key: Some(litellm_auth::SecretValue::new("test-key")),
api_base: Some(base.into()),
custom_llm_provider: None,
extra_headers: None,
optional_params: options.as_object().unwrap().clone(),
input_sources: Default::default(),
timeout_seconds: Some(2.0),
})
.unwrap()
}
async fn run_request(request: LiteLLMOcrRequest) -> Result<LiteLLMOcrResponse, Error> {
OcrRoute::new(client()).execute(request, &(), None).await
}
async fn upstream(request_count: u64, response: Value) -> MockServer {
let server = MockServer::start().await;
Mock::given(any())
.respond_with(ResponseTemplate::new(200).set_body_json(response))
.up_to_n_times(request_count)
.mount(&server)
.await;
server
}
fn image(url: &str) -> Value {
json!({"type": "image_url", "image_url": url})
}
fn body(request: &wiremock::Request) -> Value {
serde_json::from_slice(&request.body).unwrap()
}
fn assert_mistral_request(request: &wiremock::Request) {
assert_eq!(
request
.headers
.get("authorization")
.unwrap()
.to_str()
.unwrap(),
"Bearer test-key"
);
assert_eq!(
request
.headers
.get("content-type")
.unwrap()
.to_str()
.unwrap(),
"application/json"
);
assert_eq!(
body(request),
json!({
"model": "mistral-ocr-maas",
"document": {"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"},
"pages": [0, 2],
"include_image_base64": true,
"unknown": "preserved"
})
);
}
#[rstest]
#[tokio::test]
async fn cohere_body_keeps_native_document_fields_and_untyped_overrides() {
let upstream = upstream(
1,
json!({"pages": [{"index": 0, "markdown": {"content": "hello"}}]}),
)
.await;
let request = request(
"cohere/parse",
&upstream.uri(),
image("https://example.com/original.png"),
json!({
"output_format": "markdown", "timeout": 30,
"extra_body": {
"output_format": {"future": true},
"document": {"type": "image_url", "image_url": "https://example.com/a.png",
"provider_options": {"nested": [false, 0, null]}}
}
}),
);
run_request(request).await.unwrap();
let [request] = upstream
.received_requests()
.await
.unwrap()
.try_into()
.unwrap();
assert_eq!(
body(&request),
json!({
"model": "parse", "output_format": {"future": true},
"document": {"type": "image_url", "image_url": "https://example.com/a.png",
"provider_options": {"nested": [false, 0, null]}}
})
);
}
#[rstest]
#[tokio::test]
async fn explicit_null_options_use_defaults_before_http() {
let upstream = upstream(
1,
json!({"pages": [{"index": 0, "markdown": {"content": "hello"}}]}),
)
.await;
let request = request(
"cohere/parse",
&upstream.uri(),
image("https://example.com/a.png"),
json!({"output_format": null, "req_format": null}),
);
assert_eq!(
request.response_format().unwrap(),
OcrResponseFormat::Litellm
);
run_request(request).await.unwrap();
let [request] = upstream
.received_requests()
.await
.unwrap()
.try_into()
.unwrap();
let body = body(&request);
assert_eq!(body["output_format"], "markdown");
assert!(body.get("req_format").is_none());
}
#[rstest]
#[tokio::test]
async fn direct_and_vertex_mistral_build_the_same_request_and_share_normalization() {
let upstream = upstream(
2,
json!({
"pages": [{"index": 0, "markdown": "hello"}],
"extra": "preserved"
}),
)
.await;
let options = json!({
"pages": [0, 2],
"include_image_base64": true,
"vertex_project": "project-1",
"vertex_location": "us-central1",
"unknown": "preserved"
});
let document =
json!({"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"});
let direct_call = request(
"mistral/mistral-ocr-maas",
&upstream.uri(),
document.clone(),
options.clone(),
);
let vertex_call = request(
"vertex_ai/mistral-ocr-maas",
&upstream.uri(),
document,
options,
);
assert_eq!(direct_call.transport.timeout, Some(Duration::from_secs(2)));
assert_eq!(vertex_call.transport.timeout, Some(Duration::from_secs(2)));
let direct_response = run_request(direct_call).await.unwrap().into_json();
let vertex_response = run_request(vertex_call).await.unwrap().into_json();
let [direct_http, vertex_http] = upstream
.received_requests()
.await
.unwrap()
.try_into()
.unwrap();
assert_eq!(direct_http.url.path(), "/v1/ocr");
assert_eq!(
vertex_http.url.path(),
"/v1/projects/project-1/locations/us-central1/publishers/mistralai/models/mistral-ocr-maas:rawPredict"
);
assert_mistral_request(&direct_http);
assert_mistral_request(&vertex_http);
assert_eq!(direct_response, vertex_response);
assert_eq!(direct_response["model"], "mistral-ocr-maas");
assert_eq!(direct_response["object"], "ocr");
assert_eq!(direct_response["extra"], "preserved");
}
#[derive(serde::Deserialize)]
struct KnownParams {
pages: Option<Vec<i64>>,
}
#[rstest]
fn parsed_provider_params_separates_known_and_extra_params() {
let arguments: CallArguments = serde_json::from_value(json!({
"pages": [0, 2],
"future_ocr_option": true,
"extra_body": {"provider_option": "value"}
}))
.unwrap();
let known: KnownParams = parse_options(&arguments).unwrap();
assert_eq!(known.pages, Some(vec![0, 2]));
assert_eq!(arguments["future_ocr_option"], true);
assert_eq!(arguments["extra_body"], json!({"provider_option": "value"}));
assert_eq!(
arguments
.iter()
.filter(|(name, _)| name.as_str() != "pages")
.count(),
2
);
assert_eq!(
compose_body(&arguments, &json!({"pages": known.pages}), &["pages"]).unwrap(),
json!({
"pages": [0, 2], "future_ocr_option": true, "provider_option": "value"
})
);
}
}

78
litellm-rust/crates/inference-ocr/src/lib.rs
View file

@ -1,10 +1,82 @@
pub mod arguments;
mod client;
pub use client::OcrRoute;
pub mod document;
pub(crate) mod handler;
pub(crate) mod prepare;
pub mod provider_config;
pub mod route;
pub mod types;
pub mod wire;
use std::sync::Arc;
use litellm_host::interceptors::Interceptors;
use litellm_host::observation::ObservationSender;
use litellm_llms::base_llm::ocr::{error::Error, handler::OcrClient};
use litellm_llms_types::formats::ocr::LiteLLMOcrResponse;
use crate::{
handler::execute,
types::{LiteLLMOcrRequest, OcrDocumentInput},
};
#[derive(Clone)]
pub struct OcrRoute {
client: OcrClient,
}
impl OcrRoute {
pub fn new(client: OcrClient) -> Self {
Self { client }
}
pub async fn execute(
&self,
request: LiteLLMOcrRequest,
interceptors: &impl Interceptors<Error>,
observers: Option<ObservationSender>,
) -> Result<LiteLLMOcrResponse, Error> {
litellm_host::lifecycle::observe_unary(
observers.clone(),
self.run(request, interceptors, observers.as_ref()),
)
.await
}
#[tracing::instrument(name = "litellm.route", skip_all, fields(
route = "ocr",
model = %call.model,
resolved_model = %call.model,
provider = <&str>::from(call.config.provider()),
stream = false,
outcome
))]
async fn run(
&self,
call: LiteLLMOcrRequest,
interceptors: &impl Interceptors<Error>,
observers: Option<&ObservationSender>,
) -> Result<LiteLLMOcrResponse, Error> {
litellm_inference::diagnostic::unary(async {
let caller_document = matches!(&call.document, OcrDocumentInput::Document(_));
let call = document::resolve_document(call).await?;
let config = call.config;
let secrets = self
.client
.secret_source()
.resolve(&config.secret_names())
.await
.map_err(|error| Error::Secret(Arc::new(error)))?;
let execute: futures_util::future::BoxFuture<'_, Result<LiteLLMOcrResponse, Error>> =
Box::pin(execute(
&self.client,
config,
call,
secrets,
interceptors,
caller_document,
observers,
));
execute.await
})
.await
}
}

309
litellm-rust/crates/inference-ocr/src/prepare.rs
View file

@ -1,309 +0,0 @@
use litellm_auth::{InputSource, SecretValue, Sourced};
use litellm_core_utils::get_llm_provider_logic::LlmProviders;
use litellm_llms::base_llm::ocr::{
handler::OcrClient,
transformation::{OcrConnection, OcrCredentialInputs, PreparedOcrRequest},
};
use litellm_secrets::source::Secrets;
use crate::types::{LiteLLMOcrRequest, ResolvedOcrRequest};
pub(crate) fn prepare_request(
request: ResolvedOcrRequest,
caller_document: bool,
client: &OcrClient,
secrets: Secrets,
) -> PreparedOcrRequest {
let credentials = request.credentials.clone();
let (preferred_api_key_env, api_base_env) = match request.config.provider() {
LlmProviders::Mistral => (
Some("MISTRAL_AZURE_API_KEY"),
Some("MISTRAL_AZURE_API_BASE"),
),
LlmProviders::AzureAi => (None, Some("AZURE_AI_API_BASE")),
_ => (None, None),
};
let secret = |name: &str| secrets.truthy(name);
let dynamic_api_key = credentials.dynamic_api_key.or_else(|| {
credentials.api_key.clone().or_else(|| {
preferred_api_key_env
.into_iter()
.chain(request.config.get_api_key_env_var())
.find_map(secret)
.map(|value| Sourced::new(SecretValue::new(value), InputSource::Environment))
})
});
let dynamic_api_base = credentials.dynamic_api_base.or_else(|| {
credentials.api_base.clone().or_else(|| {
api_base_env
.and_then(secret)
.map(|value| Sourced::new(value, InputSource::Environment))
})
});
let resolved = request
.config
.resolve_connection_params(OcrCredentialInputs {
dynamic_api_key,
dynamic_api_base,
..credentials
});
let LiteLLMOcrRequest {
model,
document,
transport,
optional_params,
input_sources,
azure_ad_token_provider,
..
} = request;
PreparedOcrRequest {
model,
document,
connection: OcrConnection::new(resolved, transport, client.settings().clone(), secrets),
caller_document,
optional_params,
input_sources,
azure_ad_token_provider,
}
}
#[cfg(test)]
mod tests {
use std::time::Duration;
use futures_util::future::BoxFuture;
use litellm_core_utils::call_arguments::{CallArguments, compose_body, parse_options};
use litellm_host::interceptors::WireRequest;
use litellm_llms::{
base_llm::ocr::{
error::Error,
handler::{CallHooks, OcrClient},
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::*;
use crate::{
document::prepare_document,
types::LiteLLMOcrRequest,
wire::{OcrWireRequest, decode_request},
};
/// Stands in for a host with no interceptors registered.
struct NoHooks;
impl CallHooks<Error> for NoHooks {
fn before_provider_request(
&self,
wire: WireRequest,
) -> BoxFuture<'_, Result<WireRequest, Error>> {
Box::pin(async move { Ok(wire) })
}
fn response_received<'a>(&'a self, _body: &'a [u8]) -> BoxFuture<'a, Result<(), Error>> {
Box::pin(async { Ok(()) })
}
}
fn client() -> OcrClient {
OcrClient::for_test(
litellm_http::Client::plain_for_test(),
litellm_http::Client::no_redirect_for_test(),
)
}
fn request(model: &str, base: &str, document: Value, options: Value) -> LiteLLMOcrRequest {
decode_request(OcrWireRequest {
model: model.into(),
document,
api_key: Some(litellm_auth::SecretValue::new("test-key")),
api_base: Some(base.into()),
custom_llm_provider: None,
extra_headers: None,
optional_params: options.as_object().unwrap().clone(),
input_sources: Default::default(),
timeout_seconds: Some(2.0),
})
.unwrap()
}
fn prepared(request: LiteLLMOcrRequest) -> PreparedOcrRequest {
prepare_request(
request.map_document(prepare_document).unwrap(),
true,
&client(),
std::sync::Arc::new(litellm_core_utils::settings::ProcessEnvironment),
)
}
fn image(url: &str) -> Value {
json!({"type": "image_url", "image_url": url})
}
#[rstest::rstest]
#[tokio::test]
async fn cohere_body_keeps_native_document_fields_and_untyped_overrides() {
let request = request(
"cohere/parse",
"https://example.com",
image("https://example.com/original.png"),
json!({
"output_format": "markdown", "timeout": 30,
"extra_body": {
"output_format": {"future": true},
"document": {"type": "image_url", "image_url": "https://example.com/a.png",
"provider_options": {"nested": [false, 0, null]}}
}
}),
);
let http = CohereParseConfig
.prepare_request(&prepared(request), &client(), &NoHooks)
.await
.unwrap();
let body: Value = serde_json::from_slice(http.body()).unwrap();
assert_eq!(
body,
json!({
"model": "parse", "output_format": {"future": true},
"document": {"type": "image_url", "image_url": "https://example.com/a.png",
"provider_options": {"nested": [false, 0, null]}}
})
);
}
#[rstest::rstest]
#[tokio::test]
async fn explicit_null_options_use_defaults_before_http() {
let request = request(
"cohere/parse",
"https://example.com",
image("https://example.com/a.png"),
json!({"output_format": null, "req_format": null}),
);
assert_eq!(
request.response_format().unwrap(),
OcrResponseFormat::Litellm
);
let http = CohereParseConfig
.prepare_request(&prepared(request), &client(), &NoHooks)
.await
.unwrap();
let body: Value = serde_json::from_slice(http.body()).unwrap();
assert_eq!(body["output_format"], "markdown");
assert!(body.get("req_format").is_none());
}
#[rstest::rstest]
#[tokio::test]
async fn direct_and_vertex_mistral_build_the_same_request_and_share_normalization() {
let options = json!({
"pages": [0, 2],
"include_image_base64": true,
"vertex_project": "project-1",
"vertex_location": "us-central1",
"unknown": "preserved"
});
let document =
json!({"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"});
let direct = prepared(request(
"mistral/mistral-ocr-maas",
"https://mistral.test",
document.clone(),
options.clone(),
));
let vertex = prepared(request(
"vertex_ai/mistral-ocr-maas",
"https://vertex.test",
document,
options,
));
let direct_http = MistralOcrConfig
.prepare_request(&direct, &client(), &NoHooks)
.await
.unwrap();
let vertex_http = VertexAiOcrConfig
.prepare_request(&vertex, &client(), &NoHooks)
.await
.unwrap();
assert_eq!(direct_http.url(), "https://mistral.test/v1/ocr");
assert_eq!(
vertex_http.url(),
"https://vertex.test/v1/projects/project-1/locations/us-central1/publishers/mistralai/models/mistral-ocr-maas:rawPredict"
);
for http in [&direct_http, &vertex_http] {
assert_eq!(http.header("authorization").unwrap(), "Bearer test-key");
assert_eq!(http.header("content-type").unwrap(), "application/json");
assert_eq!(http.timeout(), Some(Duration::from_secs(2)));
let body: Value = serde_json::from_slice(http.body()).unwrap();
assert_eq!(
body,
json!({
"model": "mistral-ocr-maas",
"document": {"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"},
"pages": [0, 2],
"include_image_base64": true,
"unknown": "preserved"
})
);
}
let payload = serde_json::to_vec(
&json!({"pages": [{"index": 0, "markdown": "hello"}], "extra": "preserved"}),
)
.unwrap();
let direct_response = MistralOcrConfig
.transform_ocr_response(&direct.model, &payload, OcrResponseFormat::Litellm)
.unwrap()
.into_json();
let vertex_response = VertexAiOcrConfig
.transform_ocr_response(&vertex.model, &payload, OcrResponseFormat::Litellm)
.unwrap()
.into_json();
assert_eq!(direct_response, vertex_response);
assert_eq!(direct_response["model"], "mistral-ocr-maas");
assert_eq!(direct_response["object"], "ocr");
assert_eq!(direct_response["extra"], "preserved");
}
#[derive(serde::Deserialize)]
struct KnownParams {
pages: Option<Vec<i64>>,
}
#[rstest::rstest]
fn parsed_provider_params_separates_known_and_extra_params() {
let arguments: CallArguments = serde_json::from_value(json!({
"pages": [0, 2],
"future_ocr_option": true,
"extra_body": {"provider_option": "value"}
}))
.unwrap();
let known: KnownParams = parse_options(&arguments).unwrap();
assert_eq!(known.pages, Some(vec![0, 2]));
assert_eq!(arguments["future_ocr_option"], true);
assert_eq!(arguments["extra_body"], json!({"provider_option": "value"}));
assert_eq!(
arguments
.iter()
.filter(|(name, _)| name.as_str() != "pages")
.count(),
2
);
assert_eq!(
compose_body(&arguments, &json!({"pages": known.pages}), &["pages"]).unwrap(),
json!({
"pages": [0, 2], "future_ocr_option": true, "provider_option": "value"
})
);
}
}

106
litellm-rust/crates/inference-ocr/src/provider_config.rs
View file

@ -1,5 +1,6 @@
use litellm_core_utils::get_llm_provider_logic::LlmProviders;
use litellm_core_utils::get_llm_provider_logic::{CustomLlmProvider, get_custom_llm_provider};
use litellm_auth::{InputSource, SecretValue, Sourced};
use litellm_core_utils::get_llm_provider_logic::{LlmProviders, get_custom_llm_provider};
use litellm_inference::provider::{ResolvedProvider, resolve_llm_provider};
use litellm_llms::{
aws_textract::ocr::{
analyze_transformation::TextractAnalyzeDocumentConfig, common_utils::TextractOperation,
@ -25,6 +26,7 @@ use litellm_llms::{
},
};
use litellm_llms_types::formats::ocr::{LiteLLMOcrResponse, OcrDocument, OcrResponseFormat};
use litellm_secrets::source::Secrets;
macro_rules! with_config {
($kind:expr, $config:ident => $body:expr) => {
@ -115,7 +117,18 @@ impl OcrConfigKind {
}
pub(crate) fn secret_names(self) -> Vec<&'static str> {
with_config!(self, config => config.secret_names())
let names = with_config!(self, config => config.secret_names());
let fallback_names: &[&str] = match self.provider() {
LlmProviders::Mistral => &["MISTRAL_AZURE_API_KEY", "MISTRAL_AZURE_API_BASE"],
LlmProviders::AzureAi => &["AZURE_AI_API_BASE"],
_ => &[],
};
let additional_names = fallback_names
.iter()
.copied()
.filter(|name| !names.contains(name))
.collect::<Vec<_>>();
names.into_iter().chain(additional_names).collect()
}
pub(crate) fn get_health_check_document(self) -> OcrDocument {
@ -129,6 +142,43 @@ impl OcrConfigKind {
with_config!(self, config => config.resolve_connection_params(inputs))
}
pub(crate) fn resolve_credentials(
self,
credentials: OcrCredentialInputs,
secrets: &Secrets,
) -> ResolvedOcrCredentials {
let (preferred_api_key_env, api_base_env) = match self.provider() {
LlmProviders::Mistral => (
Some("MISTRAL_AZURE_API_KEY"),
Some("MISTRAL_AZURE_API_BASE"),
),
LlmProviders::AzureAi => (None, Some("AZURE_AI_API_BASE")),
_ => (None, None),
};
let secret = |name: &str| secrets.truthy(name);
let dynamic_api_key = credentials.dynamic_api_key.or_else(|| {
credentials.api_key.clone().or_else(|| {
preferred_api_key_env
.into_iter()
.chain(self.get_api_key_env_var())
.find_map(secret)
.map(|value| Sourced::new(SecretValue::new(value), InputSource::Environment))
})
});
let dynamic_api_base = credentials.dynamic_api_base.or_else(|| {
credentials.api_base.clone().or_else(|| {
api_base_env
.and_then(secret)
.map(|value| Sourced::new(value, InputSource::Environment))
})
});
self.resolve_connection_params(OcrCredentialInputs {
dynamic_api_key,
dynamic_api_base,
..credentials
})
}
pub(crate) async fn ocr(
self,
client: &OcrClient,
@ -164,7 +214,7 @@ pub fn passthrough_response(
endpoint: &str,
body: &[u8],
) -> Result<Option<LiteLLMOcrResponse>, Error> {
let (model, config) = resolve_provider_config(model, Some("azure_ai"))?;
let (provider, config) = resolve_provider_config(model, Some("azure_ai"))?;
let segments: Vec<&str> = endpoint
.split('/')
.filter(|segment| !segment.is_empty())
@ -173,7 +223,7 @@ pub fn passthrough_response(
OcrConfigKind::AzureAi => segments == AZURE_AI_OCR_PATH,
OcrConfigKind::AzureCohere => segments == AZURE_COHERE_PARSE_PATH,
OcrConfigKind::AzureDocumentIntelligence => {
segments == AzureDocumentIntelligenceOcrConfig::analyze_path(&model)?
segments == AzureDocumentIntelligenceOcrConfig::analyze_path(provider.model)?
}
other => {
let provider: &'static str = other.provider().into();
@ -183,24 +233,28 @@ pub fn passthrough_response(
if !is_ocr_endpoint {
return Ok(None);
}
with_config!(config, config => config.transform_ocr_response(&model, body, OcrResponseFormat::Litellm))
with_config!(config, config => config.transform_ocr_response(provider.model, body, OcrResponseFormat::Litellm))
.map(Some)
}
pub(crate) fn resolve_provider_config(
model: &str,
custom_llm_provider: Option<&str>,
) -> Result<(String, OcrConfigKind), Error> {
let provider =
get_custom_llm_provider(model, custom_llm_provider).unwrap_or(CustomLlmProvider {
model,
custom_llm_provider: LlmProviders::Mistral.into(),
});
let llm_provider = provider
.custom_llm_provider
.parse::<LlmProviders>()
.map_err(|_| Error::InvalidProvider(provider.custom_llm_provider.to_string()))?;
let config = match llm_provider {
pub(crate) fn resolve_provider_config<'a>(
model: &'a str,
custom_llm_provider: Option<&'a str>,
) -> Result<(ResolvedProvider<'a>, OcrConfigKind), Error> {
let provider = match resolve_llm_provider(model, custom_llm_provider, "ocr") {
Ok(provider) => provider,
Err(_) if get_custom_llm_provider(model, custom_llm_provider).is_none() => {
ResolvedProvider {
model,
provider: LlmProviders::Mistral,
}
}
Err(litellm_inference::RouteError::InvalidProvider(provider)) => {
return Err(Error::InvalidProvider(provider));
}
Err(error) => return Err(Error::InvalidProvider(error.to_string())),
};
let config = match provider.provider {
LlmProviders::AwsTextract => match TextractOperation::from_model(provider.model)? {
TextractOperation::DetectDocumentText => OcrConfigKind::AwsTextract,
TextractOperation::AnalyzeDocument => OcrConfigKind::AwsTextractAnalyze,
@ -227,11 +281,11 @@ pub(crate) fn resolve_provider_config(
LlmProviders::VertexAi => OcrConfigKind::VertexAi,
_ => {
return Err(Error::InvalidProvider(
provider.custom_llm_provider.to_string(),
<&'static str>::from(provider.provider).to_string(),
));
}
};
Ok((provider.model.to_string(), config))
Ok((provider, config))
}
fn is_document_intelligence_model(model: &str) -> bool {
@ -243,7 +297,7 @@ fn is_document_intelligence_model(model: &str) -> bool {
mod tests {
use std::collections::HashSet;
use litellm_auth::{InputSource, Sourced};
use litellm_auth::Sourced;
use litellm_llms::{
base_llm::ocr::document::InlineDocument, cohere::ocr::transformation::validate_document,
};
@ -484,8 +538,8 @@ mod tests {
#[case] expected_config: OcrConfigKind,
) {
let expected_model = qualified_model.split_once('/').unwrap().1;
let (model, config) = resolve_provider_config(qualified_model, None).unwrap();
assert_eq!(model, expected_model);
let (provider, config) = resolve_provider_config(qualified_model, None).unwrap();
assert_eq!(provider.model, expected_model);
assert_eq!(config, expected_config);
}
@ -528,7 +582,7 @@ mod tests {
expected_config
);
assert_eq!(
resolve_provider_config(model, None).unwrap().0,
resolve_provider_config(model, None).unwrap().0.model,
model.split_once('/').unwrap().1
);
}

13
litellm-rust/crates/inference-ocr/src/types.rs
View file

@ -75,6 +75,8 @@ impl OcrConnectionInputs {
pub struct LiteLLMOcrRequest<D = OcrDocumentInput> {
pub model: String,
pub(crate) requested_model: String,
pub(crate) custom_llm_provider: Option<String>,
pub document: D,
pub credentials: OcrCredentialInputs,
pub transport: OcrTransportConfig,
@ -91,7 +93,8 @@ impl LiteLLMOcrRequest {
custom_llm_provider: Option<&str>,
optional_params: CallArguments,
) -> Result<Self, Error> {
let (model, config) = resolve_provider_config(&model, custom_llm_provider)?;
let (provider, config) = resolve_provider_config(&model, custom_llm_provider)?;
let resolved_model = provider.model.to_owned();
let default_transport = OcrTransportConfig::default();
let max_response_bytes = optional_params
.get("max_response_bytes")
@ -116,7 +119,9 @@ impl LiteLLMOcrRequest {
.collect();
Ok(Self {
model,
model: resolved_model,
requested_model: model,
custom_llm_provider: custom_llm_provider.map(str::to_owned),
document: document.into(),
credentials: OcrCredentialInputs::default(),
transport,
@ -135,6 +140,8 @@ impl<D> LiteLLMOcrRequest<D> {
) -> Result<LiteLLMOcrRequest<T>, E> {
Ok(LiteLLMOcrRequest {
model: self.model,
requested_model: self.requested_model,
custom_llm_provider: self.custom_llm_provider,
document: map(self.document)?,
credentials: self.credentials,
transport: self.transport,
@ -148,6 +155,8 @@ impl<D> LiteLLMOcrRequest<D> {
pub fn with_document<T>(self, document: T) -> LiteLLMOcrRequest<T> {
LiteLLMOcrRequest {
model: self.model,
requested_model: self.requested_model,
custom_llm_provider: self.custom_llm_provider,
document,
credentials: self.credentials,
transport: self.transport,

13
litellm-rust/crates/inference-ocr/tests/ocr/lifecycle.rs
View file

@ -66,6 +66,19 @@ async fn interception_runs_once_and_observers_receive_ordered_success_events() {
assert_eq!(received(&upstream).await.len(), 1);
}
#[rstest::rstest]
#[case::mutated_model("model-b")]
#[tokio::test]
async fn provider_request_uses_the_request_model_after_it_is_mutated(#[case] model: &str) {
let upstream = upstream([pages_response()]).await;
let mut request = ocr_request("mistral/model-a", &upstream.uri(), json!({}));
request.model = model.into();
perform(request).await.unwrap();
assert_eq!(only_request(&upstream).await.json()["model"], json!(model));
}
#[rstest::rstest]
#[tokio::test]
async fn a_blocking_before_send_prevents_the_call_and_emits_one_failure() {

61
litellm-rust/crates/inference-responses/src/handler.rs
View file

@ -3,39 +3,61 @@ use litellm_host::observation::ObservationSender;
use std::time::Duration;
use futures_util::StreamExt;
use litellm_host::interceptors::{Interceptors, RawResponse, WireRequest};
use litellm_llms::base_llm::auth::{Authenticated, resolve_auth};
use litellm_host::interceptors::{Interceptors, RawResponse, RequestContext, WireRequest};
use litellm_llms::base_llm::auth::{AuthScheme, Authenticated, resolve_auth};
use litellm_secrets::source::Secrets;
use super::{
Error,
types::{ProviderResponsesRequest, ResponsesOutput, ResponsesStreamHead},
Error, ResponsesRoute,
types::{ResponsesCall, ResponsesOutput, ResponsesProvider, ResponsesStreamHead},
};
pub(super) async fn execute(
http: &litellm_http::Client,
auth: &litellm_auth::AuthServices,
request: ProviderResponsesRequest,
cache: Option<litellm_cache_response::ScopedCache>,
route: &ResponsesRoute,
ResponsesProvider { identity, config }: ResponsesProvider,
call: ResponsesCall,
secrets: Secrets,
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 identity = litellm_host::interceptors::ProviderIdentity {
model: request.context.model.clone(),
provider: request.context.custom_llm_provider.clone(),
let lookup = |name: &str| secrets.get(name);
let environment = config.validate_environment(
litellm_http::request::string_headers("responses", call.extra_headers)?,
call.api_key.as_deref(),
&lookup,
)?;
let url = config.get_complete_url(call.api_base.as_deref(), &lookup);
let context = RequestContext {
model: identity.model.clone(),
custom_llm_provider: identity.provider.clone(),
optional_params: serde_json::Value::Object(call.optional_params.clone()),
secret_fields: Vec::new(),
api_key: match &environment.auth {
AuthScheme::Credential { secret, .. } => Some(secret.clone()),
_ => None,
},
};
let body = config.transform_responses_api_request(
&identity.model,
call.input,
call.optional_params,
)?;
let authenticated = resolve_auth(&route.auth, environment, &|_| None).await?;
let wire = interceptors
.before_provider_request(
WireRequest {
url: request.url,
url,
headers: authenticated.headers,
body: request.body,
body,
},
request.context,
context,
)
.await?;
let cache = cache.filter(|_| authenticated.signer.is_none());
let cache = route
.cache
.clone()
.filter(|_| authenticated.signer.is_none());
let cache_request = litellm_inference::caching::CacheRequest::from_wire(
identity,
cache.as_ref().map(|_| &wire),
@ -59,9 +81,9 @@ pub(super) async fn execute(
},
wire.url,
&wire.body,
Some(request.timeout.unwrap_or(Duration::from_secs(600))),
Some(call.timeout.unwrap_or(Duration::from_secs(600))),
)?;
let response = litellm_inference::outbound::send(outbound, http)
let response = litellm_inference::outbound::send(outbound, &route.http)
.await
.map_err(network)?;
let status = response.status().as_u16();
@ -103,8 +125,7 @@ pub(super) async fn execute(
.map_err(Error::post_call)?;
let value = serde_json::from_str(&body)
.map_err(|error| Error::InvalidResponse(error.to_string().into()))?;
request
.config
config
.transform_response_api_response(value)
.map(ResponsesOutput::Complete)
.map_err(Error::from)

108
litellm-rust/crates/inference-responses/src/lib.rs
View file

@ -7,14 +7,30 @@ pub mod types;
pub mod websocket;
mod handler;
mod prepare;
use std::sync::Arc;
use litellm_auth::AuthServices;
use litellm_core_utils::get_llm_provider_logic::{
CustomLlmProvider, LlmProviders, get_custom_llm_provider,
};
use litellm_host::interceptors::Interceptors;
use litellm_host::interceptors::ProviderIdentity;
use litellm_llms::{
base_llm::responses::transformation::BaseResponsesApiConfig,
openai::responses::transformation::OpenAiResponsesApiConfig,
openrouter::responses::transformation::OPENROUTER_RESPONSES_CONFIG,
};
use litellm_secrets::source::SecretSource;
use types::{ResponsesCall, ResponsesOutput};
use types::{ResponsesCall, ResponsesOutput, ResponsesProvider};
fn responses_provider(provider: LlmProviders) -> Option<&'static dyn BaseResponsesApiConfig> {
match provider {
LlmProviders::Openai => Some(&OpenAiResponsesApiConfig),
LlmProviders::Openrouter => Some(&OPENROUTER_RESPONSES_CONFIG),
_ => None,
}
}
#[derive(Clone)]
pub struct ResponsesRoute {
@ -78,34 +94,68 @@ impl ResponsesRoute {
observers: Option<&ObservationSender>,
) -> Result<ResponsesOutput, Error> {
litellm_inference::diagnostic::call(async {
self.run_provider(call, cache_options, interceptors, observers)
.await
let explicit = call
.custom_llm_provider
.as_deref()
.filter(|provider| !provider.is_empty());
let (provider_name, resolved_model) =
match get_custom_llm_provider(&call.model, explicit) {
Some(CustomLlmProvider {
model,
custom_llm_provider,
}) => (custom_llm_provider, model),
None => ("openai", call.model.as_str()),
};
if resolved_model.is_empty() {
return Err(Error::InvalidProvider(call.model));
}
let provider = provider_name
.parse::<LlmProviders>()
.map_err(|_| Error::Unsupported("native HTTP responses provider"))?;
let config = responses_provider(provider)
.ok_or(Error::Unsupported("native HTTP responses provider"))?;
let wire_model = match provider {
LlmProviders::Openai if resolved_model.contains('/') => {
return Err(Error::InvalidProvider(call.model));
}
LlmProviders::Openai => resolved_model.into(),
LlmProviders::Openrouter if explicit.is_some() && !resolved_model.contains('/') => {
format!("openrouter/{resolved_model}")
}
LlmProviders::Openrouter => resolved_model.into(),
_ => return Err(Error::Unsupported("native HTTP responses provider")),
};
let provider = ResponsesProvider {
identity: ProviderIdentity {
model: wire_model,
provider: provider_name.into(),
},
config,
};
litellm_inference::diagnostic::provider(
&provider.identity.model,
&provider.identity.provider,
);
let secrets = self
.secrets
.resolve(
provider
.config
.secret_names(call.api_key.as_deref(), call.api_base.as_deref()),
)
.await?;
let execute: futures_util::future::BoxFuture<'_, Result<ResponsesOutput, Error>> =
Box::pin(handler::execute(
self,
provider,
call,
secrets,
cache_options,
interceptors,
observers,
));
execute.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?;
litellm_inference::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
}
}

162
litellm-rust/crates/inference-responses/src/prepare.rs
View file

@ -1,162 +0,0 @@
use litellm_core_utils::get_llm_provider_logic::{
CustomLlmProvider, LlmProviders, get_custom_llm_provider,
};
use litellm_host::interceptors::{ProviderIdentity, RequestContext};
use litellm_llms::{
base_llm::responses::transformation::BaseResponsesApiConfig,
openai::responses::transformation::OpenAiResponsesApiConfig,
openrouter::responses::transformation::OPENROUTER_RESPONSES_CONFIG,
};
use litellm_secrets::source::SecretSource;
use serde_json::Value;
use super::{
Error,
types::{ProviderResponsesRequest, ResponsesCall},
};
#[tracing::instrument(name = "litellm.prepare", level = "debug", skip_all)]
pub(super) async fn prepare(
call: ResponsesCall,
secrets: &dyn SecretSource,
) -> Result<ProviderResponsesRequest, Error> {
let ResponsesProvider { identity, config } =
resolve_provider(&call.model, call.custom_llm_provider.as_deref())?;
let provider = identity.provider.as_str();
let model = identity.model.as_str();
let snapshot = secrets
.resolve(config.secret_names(call.api_key.as_deref(), call.api_base.as_deref()))
.await?;
let lookup = |name: &str| snapshot.get(name);
let environment = config.validate_environment(
litellm_http::request::string_headers("responses", call.extra_headers)?,
call.api_key.as_deref(),
&lookup,
)?;
let context = RequestContext {
model: model.into(),
custom_llm_provider: provider.into(),
optional_params: Value::Object(call.optional_params.clone()),
secret_fields: Vec::new(),
api_key: match &environment.auth {
litellm_llms::base_llm::auth::AuthScheme::Credential { secret, .. } => {
Some(secret.clone())
}
_ => None,
},
};
let body = config.transform_responses_api_request(model, call.input, call.optional_params)?;
Ok(ProviderResponsesRequest {
url: config.get_complete_url(call.api_base.as_deref(), &lookup),
config,
environment,
body,
context,
timeout: call.timeout,
})
}
pub(super) struct ResponsesProvider {
pub identity: ProviderIdentity,
pub config: &'static dyn BaseResponsesApiConfig,
}
/// Python's `get_llm_provider`: a bare model is OpenAI's, and with an explicit OpenRouter
/// provider a model without a vendor segment keeps the `openrouter/` prefix as its wire id.
pub(super) fn resolve_provider(
model: &str,
custom_llm_provider: Option<&str>,
) -> Result<ResponsesProvider, Error> {
let explicit = custom_llm_provider.filter(|provider| !provider.is_empty());
let (provider, resolved) = match get_custom_llm_provider(model, explicit) {
Some(CustomLlmProvider {
model,
custom_llm_provider,
}) => (custom_llm_provider, model),
None => ("openai", model),
};
if resolved.is_empty() {
return Err(Error::InvalidProvider(model.into()));
}
let (config, wire_model): (&'static dyn BaseResponsesApiConfig, String) =
match provider.parse::<LlmProviders>() {
Ok(LlmProviders::Openai) if resolved.contains('/') => {
return Err(Error::InvalidProvider(model.into()));
}
Ok(LlmProviders::Openai) => (&OpenAiResponsesApiConfig, resolved.into()),
Ok(LlmProviders::Openrouter) if explicit.is_some() && !resolved.contains('/') => (
&OPENROUTER_RESPONSES_CONFIG,
format!("openrouter/{resolved}"),
),
Ok(LlmProviders::Openrouter) => (&OPENROUTER_RESPONSES_CONFIG, resolved.into()),
_ => return Err(Error::Unsupported("native HTTP responses provider")),
};
Ok(ResponsesProvider {
identity: ProviderIdentity {
model: wire_model,
provider: provider.into(),
},
config,
})
}
#[cfg(test)]
mod tests {
use rstest::rstest;
use super::*;
/// Pinned against Python's `get_llm_provider` on 2026-10-11.
#[rstest]
#[case::bare("gpt-5", None, "openai", "gpt-5")]
#[case::openai_prefix("openai/gpt-5", None, "openai", "gpt-5")]
#[case::openai_explicit("gpt-5", Some("openai"), "openai", "gpt-5")]
#[case::openrouter_prefix(
"openrouter/anthropic/claude",
None,
"openrouter",
"anthropic/claude"
)]
#[case::openrouter_explicit(
"openrouter/anthropic/claude",
Some("openrouter"),
"openrouter",
"anthropic/claude"
)]
#[case::openrouter_explicit_unprefixed(
"anthropic/claude",
Some("openrouter"),
"openrouter",
"anthropic/claude"
)]
#[case::openrouter_native_id(
"openrouter/auto",
Some("openrouter"),
"openrouter",
"openrouter/auto"
)]
#[case::openrouter_native_id_implicit("openrouter/auto", None, "openrouter", "auto")]
fn providers_and_wire_models_match_python(
#[case] model: &str,
#[case] custom_llm_provider: Option<&str>,
#[case] provider: &str,
#[case] wire_model: &str,
) {
let resolved = resolve_provider(model, custom_llm_provider).unwrap();
assert_eq!(resolved.identity.provider, provider);
assert_eq!(resolved.identity.model, wire_model);
}
#[rstest]
#[case::other_prefix("other/model", None)]
#[case::other_explicit("model", Some("other"))]
#[case::openai_with_vendor_segment("openai/other/model", None)]
#[case::empty_model("openai/", None)]
#[case::conflicting_prefix("openrouter/anthropic/claude", Some("openai"))]
fn everything_else_is_rejected(#[case] model: &str, #[case] custom_llm_provider: Option<&str>) {
assert!(matches!(
resolve_provider(model, custom_llm_provider),
Err(Error::InvalidProvider(_) | Error::Unsupported(_))
));
}
}

19
litellm-rust/crates/inference-responses/src/types.rs
View file

@ -2,9 +2,8 @@ use std::time::Duration;
use bytes::Bytes;
use litellm_host::call::CallOutput;
use litellm_llms::base_llm::{
auth::ValidatedEnvironment, responses::transformation::BaseResponsesApiConfig,
};
use litellm_host::interceptors::ProviderIdentity;
use litellm_llms::base_llm::responses::transformation::BaseResponsesApiConfig;
use litellm_llms_types::formats::responses::ResponsesApiResponse;
use serde_json::{Map, Value};
@ -21,17 +20,13 @@ pub struct ResponsesCall {
pub timeout: Option<Duration>,
}
pub(crate) struct ResponsesProvider {
pub identity: ProviderIdentity,
pub config: &'static dyn BaseResponsesApiConfig,
}
pub struct ResponsesStreamHead {
pub headers: Vec<(String, String)>,
}
pub type ResponsesOutput = CallOutput<ResponsesApiResponse, ResponsesStreamHead, Bytes, Error>;
pub(super) struct ProviderResponsesRequest {
pub config: &'static dyn BaseResponsesApiConfig,
pub environment: ValidatedEnvironment,
pub url: String,
pub body: Value,
pub context: litellm_host::interceptors::RequestContext,
pub timeout: Option<Duration>,
}

105
litellm-rust/crates/inference-responses/tests/responses.rs
View file

@ -7,6 +7,7 @@ use std::sync::Arc;
use futures_util::TryStreamExt;
use litellm_host::{call::HostedCompletion, lifecycle::CallEvent};
use litellm_inference_responses::{
Error,
route::Responses,
types::{ResponsesCall, ResponsesOutput},
};
@ -32,6 +33,110 @@ fn call() -> ResponsesCall {
}
}
#[rstest]
#[case::bare_openai("gpt-5", None, "gpt-5")]
#[case::openai_prefix("openai/gpt-5", None, "gpt-5")]
#[case::empty_explicit_provider("openai/gpt-5", Some(""), "gpt-5")]
#[case::explicit_openai("gpt-5", Some("openai"), "gpt-5")]
#[case::openrouter_prefix("openrouter/anthropic/claude", None, "anthropic/claude")]
#[case::explicit_openrouter_prefix(
"openrouter/anthropic/claude",
Some("openrouter"),
"anthropic/claude"
)]
#[case::explicit_openrouter_unprefixed("anthropic/claude", Some("openrouter"), "anthropic/claude")]
#[case::explicit_openrouter_native_id("auto", Some("openrouter"), "openrouter/auto")]
#[case::implicit_openrouter_native_id("openrouter/auto", None, "auto")]
#[tokio::test]
async fn provider_resolution_preserves_the_wire_model(
#[case] model: &str,
#[case] custom_llm_provider: Option<&str>,
#[case] expected_wire_model: &str,
) {
let response = json!({
"id": "response-1",
"model": expected_wire_model,
"output": [{"type": "message", "content": []}],
"usage": {"total_tokens": 7}
});
let upstream = upstream([json_response(response)]).await;
responses_route(no_secrets())
.execute(
ResponsesCall {
model: model.into(),
api_key: Some("test-key".into()),
api_base: Some(upstream.uri()),
custom_llm_provider: custom_llm_provider.map(str::to_owned),
..call()
},
&(),
None,
)
.await
.unwrap();
assert_eq!(
only_request(&upstream).await.json()["model"],
expected_wire_model
);
}
#[rstest]
#[case::explicit_unknown_provider("test-model", Some("unknown-provider"))]
#[case::unsupported_responses_provider("test-model", Some("anthropic"))]
#[tokio::test]
async fn unsupported_providers_fail_before_transport(
#[case] model: &str,
#[case] custom_llm_provider: Option<&str>,
) {
let upstream = upstream([]).await;
let result = responses_route(no_secrets())
.execute(
ResponsesCall {
model: model.into(),
api_base: Some(upstream.uri()),
custom_llm_provider: custom_llm_provider.map(str::to_owned),
..call()
},
&(),
None,
)
.await;
assert!(matches!(
result,
Err(Error::Unsupported("native HTTP responses provider"))
));
assert!(received(&upstream).await.is_empty());
}
#[rstest]
#[case::openai_with_vendor_segment("openai/other/model", None)]
#[case::empty_model("openai/", None)]
#[case::conflicting_prefix("openrouter/anthropic/claude", Some("openai"))]
#[tokio::test]
async fn invalid_provider_models_fail_before_transport(
#[case] model: &str,
#[case] custom_llm_provider: Option<&str>,
) {
let upstream = upstream([]).await;
let result = responses_route(no_secrets())
.execute(
ResponsesCall {
model: model.into(),
api_base: Some(upstream.uri()),
custom_llm_provider: custom_llm_provider.map(str::to_owned),
..call()
},
&(),
None,
)
.await;
assert!(matches!(result, Err(Error::InvalidProvider(_))));
assert!(received(&upstream).await.is_empty());
}
#[rstest]
#[case::direct(false)]
#[case::hosted(true)]

60
litellm-rust/crates/inference-transcription/src/handler.rs
View file

@ -1,28 +1,57 @@
use std::time::Duration;
use litellm_http::{Client, request::truncate_error_body};
use litellm_llms::base_llm::auth::resolve_auth;
use litellm_http::{
Client,
request::{string_headers, truncate_error_body, with_default_headers},
};
use litellm_llms::base_llm::{
audio_transcription::transformation::BaseAudioTranscriptionConfig,
auth::{ValidatedEnvironment, resolve_auth},
};
use litellm_secrets::source::Secrets;
use serde_json::Value;
use super::Error;
use crate::{
constants::AUDIO_TRANSCRIPTION_TIMEOUT_SECS, types::ProviderAudioTranscriptionRequest,
};
use crate::{constants::AUDIO_TRANSCRIPTION_TIMEOUT_SECS, types::AudioTranscriptionRequest};
use litellm_inference::provider::ResolvedProvider;
pub async fn execute_audio_transcription_provider_call(
pub(super) async fn execute(
http: &Client,
auth: &litellm_auth::AuthServices,
request: ProviderAudioTranscriptionRequest,
config: &'static dyn BaseAudioTranscriptionConfig,
provider: ResolvedProvider<'_>,
call: AudioTranscriptionRequest<'_>,
secrets: Secrets,
) -> Result<Value, Error> {
let env_lookup = |key: &str| request.secrets.get(key);
let authenticated = resolve_auth(auth, request.environment.clone(), &env_lookup).await?;
let env_lookup = |key: &str| secrets.get(key);
let forwarded = string_headers("audio transcription", call.extra_headers)?;
let validated = config.validate_environment(
forwarded,
provider.model,
&call.optional_params,
&env_lookup,
)?;
let environment = ValidatedEnvironment {
headers: with_default_headers(validated.headers, config.default_headers()),
auth: validated.auth,
};
let url = config.get_complete_url(
call.api_base,
provider.model,
&call.optional_params,
&env_lookup,
)?;
let filtered_params = config.map_transcription_params(&call.optional_params);
let body = config
.transform_audio_transcription_request(provider.model, call.audio, filtered_params)?
.body;
let authenticated = resolve_auth(auth, environment, &env_lookup).await?;
let outbound = litellm_inference::outbound::outbound_request(
authenticated,
request.url.clone(),
&request.body,
url,
&body,
Some(
request
.timeout
call.timeout
.unwrap_or(Duration::from_secs(AUDIO_TRANSCRIPTION_TIMEOUT_SECS)),
),
)?;
@ -47,8 +76,7 @@ pub async fn execute_audio_transcription_provider_call(
error,
))
})?;
Ok(request
.config
.transform_audio_transcription_response(&request.model, response_json)?
Ok(config
.transform_audio_transcription_response(provider.model, response_json)?
.into_json())
}

43
litellm-rust/crates/inference-transcription/src/lib.rs
View file

@ -2,16 +2,29 @@ pub mod types;
pub use litellm_inference::RouteError as Error;
mod constants;
mod handler;
mod prepare;
pub use handler::execute_audio_transcription_provider_call;
use litellm_auth::AuthServices;
use litellm_core_utils::get_llm_provider_logic::LlmProviders;
use litellm_inference::provider::resolve_llm_provider;
use litellm_llms::{
base_llm::audio_transcription::transformation::BaseAudioTranscriptionConfig,
bedrock::audio_transcription::BEDROCK_AUDIO_TRANSCRIPTION_CONFIG,
};
use litellm_secrets::source::SecretSource;
pub use prepare::prepare_audio_transcription_provider_call;
use serde_json::Value;
use std::sync::Arc;
use crate::types::AudioTranscriptionRequest;
fn audio_transcription_provider(
provider: LlmProviders,
) -> Option<&'static dyn BaseAudioTranscriptionConfig> {
match provider {
LlmProviders::Bedrock => Some(&BEDROCK_AUDIO_TRANSCRIPTION_CONFIG),
_ => None,
}
}
#[derive(Clone)]
pub struct AudioTranscriptionRoute {
http: litellm_http::Client,
@ -32,6 +45,10 @@ impl AudioTranscriptionRoute {
}
}
pub async fn execute(&self, request: AudioTranscriptionRequest<'_>) -> Result<Value, Error> {
self.run(request).await
}
#[tracing::instrument(name = "litellm.route", skip_all, fields(
route = "audio_transcription",
model = request.model,
@ -40,13 +57,23 @@ impl AudioTranscriptionRoute {
stream = false,
outcome
))]
pub async fn execute(&self, request: AudioTranscriptionRequest<'_>) -> Result<Value, Error> {
async fn run(&self, request: AudioTranscriptionRequest<'_>) -> Result<Value, Error> {
litellm_inference::diagnostic::unary(async {
let request =
prepare_audio_transcription_provider_call(request, self.secrets.as_ref()).await?;
litellm_inference::diagnostic::provider(&request.model, &request.custom_llm_provider);
let provider = resolve_llm_provider(
request.model,
request.custom_llm_provider,
"audio transcription",
)?;
let config = audio_transcription_provider(provider.provider).ok_or_else(|| {
Error::InvalidProvider(<&str>::from(provider.provider).to_string())
})?;
litellm_inference::diagnostic::provider(
provider.model,
<&'static str>::from(provider.provider),
);
let secrets = self.secrets.resolve(&config.secret_names()).await?;
let execute: futures_util::future::BoxFuture<'_, Result<Value, Error>> = Box::pin(
execute_audio_transcription_provider_call(&self.http, &self.auth, request),
handler::execute(&self.http, &self.auth, config, provider, request, secrets),
);
execute.await
})

65
litellm-rust/crates/inference-transcription/src/prepare.rs
View file

@ -1,65 +0,0 @@
use litellm_core_utils::get_llm_provider_logic::LlmProviders;
use litellm_http::request::string_headers;
use litellm_http::request::with_default_headers;
use litellm_llms::{
base_llm::{
audio_transcription::transformation::BaseAudioTranscriptionConfig,
auth::ValidatedEnvironment,
},
bedrock::audio_transcription::BEDROCK_AUDIO_TRANSCRIPTION_CONFIG,
};
use litellm_secrets::source::SecretSource;
use super::Error;
use crate::types::{AudioTranscriptionRequest, ProviderAudioTranscriptionRequest};
use litellm_inference::provider::resolve_llm_provider;
fn provider_config(provider: LlmProviders) -> Option<&'static dyn BaseAudioTranscriptionConfig> {
match provider {
LlmProviders::Bedrock => Some(&BEDROCK_AUDIO_TRANSCRIPTION_CONFIG),
_ => None,
}
}
#[tracing::instrument(name = "litellm.prepare", level = "debug", skip_all)]
pub async fn prepare_audio_transcription_provider_call(
request: AudioTranscriptionRequest<'_>,
secrets: &dyn SecretSource,
) -> Result<ProviderAudioTranscriptionRequest, Error> {
let provider_info = resolve_llm_provider(
request.model,
request.custom_llm_provider,
"audio transcription",
)?;
let model = provider_info.model.to_string();
let config = provider_config(provider_info.provider)
.ok_or_else(|| Error::InvalidProvider(<&str>::from(provider_info.provider).to_string()))?;
let snapshot = secrets.resolve(&config.secret_names()).await?;
let env_lookup = |key: &str| snapshot.get(key);
let forwarded = string_headers("audio transcription", request.extra_headers)?;
let validated =
config.validate_environment(forwarded, &model, &request.optional_params, &env_lookup)?;
let environment = ValidatedEnvironment {
headers: with_default_headers(validated.headers, config.default_headers()),
auth: validated.auth,
};
let url = config.get_complete_url(
request.api_base,
&model,
&request.optional_params,
&env_lookup,
)?;
let filtered_params = config.map_transcription_params(&request.optional_params);
let transformed =
config.transform_audio_transcription_request(&model, request.audio, filtered_params)?;
Ok(ProviderAudioTranscriptionRequest {
model,
custom_llm_provider: <&str>::from(provider_info.provider).to_string(),
config,
url,
body: transformed.body,
environment,
secrets: snapshot,
timeout: request.timeout,
})
}

38
litellm-rust/crates/inference-transcription/src/types.rs
View file

@ -1,9 +1,5 @@
use litellm_secrets::source::Secrets;
use std::time::Duration;
use litellm_llms::base_llm::{
audio_transcription::transformation::BaseAudioTranscriptionConfig, auth::ValidatedEnvironment,
};
use serde_json::{Map, Value};
pub struct AudioTranscriptionRequest<'a> {
@ -16,37 +12,3 @@ pub struct AudioTranscriptionRequest<'a> {
pub optional_params: Map<String, Value>,
pub timeout: Option<Duration>,
}
#[derive(Clone)]
pub struct ProviderAudioTranscriptionRequest {
pub model: String,
pub custom_llm_provider: String,
pub config: &'static dyn BaseAudioTranscriptionConfig,
pub url: String,
pub body: Value,
pub environment: ValidatedEnvironment,
pub secrets: Secrets,
pub timeout: Option<Duration>,
}
impl ProviderAudioTranscriptionRequest {
pub fn model(&self) -> &str {
&self.model
}
pub fn custom_llm_provider(&self) -> &str {
&self.custom_llm_provider
}
pub fn url(&self) -> &str {
&self.url
}
pub fn body(&self) -> &Value {
&self.body
}
pub fn with_body(self, body: Value) -> Self {
Self { body, ..self }
}
}

3
litellm-rust/crates/llms/src/anthropic/messages/transformation.rs
View file

@ -74,6 +74,7 @@ impl BaseMessagesConfig for AnthropicMessagesConfig {
headers: Headers,
api_key: Option<&str>,
_model: &str,
_optional_params: &MessagesOptionalParams,
_litellm_params: &LitellmParams,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Result<ValidatedEnvironment, Error> {
@ -887,6 +888,7 @@ mod tests {
headers(forwarded),
api_key,
"claude",
&MessagesOptionalParams::default(),
&LitellmParams::default(),
&env(vars),
)
@ -1206,6 +1208,7 @@ mod tests {
Vec::new(),
None,
"claude",
&MessagesOptionalParams::default(),
&LitellmParams::default(),
&record,
);

5
litellm-rust/crates/llms/src/azure_ai/messages/transformation.rs
View file

@ -1,7 +1,7 @@
use litellm_auth::SecretValue;
use litellm_http::request::{has_bearer_auth, has_header};
use litellm_llms_types::{
formats::messages::MessagesRequest,
formats::messages::{MessagesOptionalParams, MessagesRequest},
providers::anthropic::{BetaProvider, DEFAULT_HEADERS, MESSAGES_PATH},
};
use litellm_router_types::LitellmParams;
@ -83,6 +83,7 @@ impl BaseMessagesConfig for AzureAnthropicMessagesConfig {
headers: Headers,
api_key: Option<&str>,
_model: &str,
_optional_params: &MessagesOptionalParams,
_litellm_params: &LitellmParams,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Result<ValidatedEnvironment, Error> {
@ -237,6 +238,7 @@ mod tests {
.collect(),
api_key,
"claude",
&MessagesOptionalParams::default(),
&LitellmParams::default(),
&|_| None,
)
@ -622,6 +624,7 @@ mod tests {
Vec::new(),
None,
"claude",
&MessagesOptionalParams::default(),
&LitellmParams::default(),
&record,
);

6
litellm-rust/crates/llms/src/base_llm/messages/transformation.rs
View file

@ -1,4 +1,6 @@
use litellm_llms_types::formats::messages::{MessagesRequest, MessagesResponse};
use litellm_llms_types::formats::messages::{
MessagesOptionalParams, MessagesRequest, MessagesResponse,
};
use litellm_router_types::LitellmParams;
use serde_json::Value;
@ -55,6 +57,7 @@ pub trait BaseMessagesConfig: Sync {
headers: Headers,
api_key: Option<&str>,
model: &str,
optional_params: &MessagesOptionalParams,
litellm_params: &LitellmParams,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Result<ValidatedEnvironment, Error>;
@ -110,6 +113,7 @@ mod tests {
headers: Headers,
_api_key: Option<&str>,
_model: &str,
_optional_params: &MessagesOptionalParams,
_litellm_params: &LitellmParams,
_env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Result<ValidatedEnvironment, Error> {

25
litellm-rust/crates/llms/src/bedrock/claude_platform/messages/transformation.rs
View file

@ -5,7 +5,7 @@ use litellm_auth_aws::{AwsCredentialSource, resolve_aws_region};
use litellm_core_utils::settings::resolve_non_empty;
use litellm_http::request::{header_value, with_header, without_headers};
use litellm_llms_types::{
formats::messages::MessagesRequest,
formats::messages::{MessagesOptionalParams, MessagesRequest},
providers::anthropic::{BetaProvider, DEFAULT_HEADERS, MESSAGES_PATH},
};
use litellm_router_types::LitellmParams;
@ -60,6 +60,7 @@ impl BaseMessagesConfig for BedrockClaudePlatformMessagesConfig {
headers: Headers,
api_key: Option<&str>,
_model: &str,
optional_params: &MessagesOptionalParams,
params: &LitellmParams,
env: &dyn Fn(&str) -> Option<String>,
) -> Result<ValidatedEnvironment, Error> {
@ -67,11 +68,8 @@ impl BaseMessagesConfig for BedrockClaudePlatformMessagesConfig {
.filter(|value| !value.trim().is_empty())
.or_else(|| {
WORKSPACE_FIELDS.iter().find_map(|field| {
params
.extra
.get(*field)?
.as_str()
.filter(|value| !value.trim().is_empty())
workspace_value(optional_params.extra.get(*field))
.or_else(|| workspace_value(params.extra.get(*field)))
})
});
let workspace = resolve_non_empty(workspace, env, WORKSPACE_ENVS)
@ -182,6 +180,10 @@ impl BaseMessagesConfig for BedrockClaudePlatformMessagesConfig {
}
}
fn workspace_value(value: Option<&Value>) -> Option<&str> {
value?.as_str().filter(|value| !value.trim().is_empty())
}
fn platform_region(
base: Option<&str>,
params: &LitellmParams,
@ -260,6 +262,7 @@ mod tests {
vec![],
Some("key"),
"claude-test",
&MessagesOptionalParams::default(),
&LitellmParams::default(),
&|_| None
)
@ -280,6 +283,7 @@ mod tests {
],
Some("explicit"),
"claude-test",
&MessagesOptionalParams::default(),
&LitellmParams::default(),
&|name| (name == "ANTHROPIC_AWS_WORKSPACE_ID").then(|| "workspace".into()),
)
@ -345,7 +349,14 @@ mod tests {
.get_complete_url(None, "claude-test", &params, false, &env)
.unwrap();
let environment = BEDROCK_CLAUDE_PLATFORM_MESSAGES_CONFIG
.validate_environment(vec![], None, "claude-test", &params, &env)
.validate_environment(
vec![],
None,
"claude-test",
&MessagesOptionalParams::default(),
&params,
&env,
)
.unwrap();
assert_eq!(
url,

5
litellm-rust/crates/llms/src/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.rs
View file

@ -16,7 +16,7 @@ use litellm_auth_aws::{
resolve_bedrock_region,
};
use litellm_llms_types::formats::messages::{
MessagesRequest,
MessagesOptionalParams, MessagesRequest,
streaming::{MessagesStreamEvent, MessagesStreamUsage},
};
use litellm_llms_types::providers::bedrock::{
@ -146,6 +146,7 @@ impl BaseMessagesConfig for AmazonAnthropicClaudeMessagesConfig {
headers: Headers,
api_key: Option<&str>,
model: &str,
_optional_params: &MessagesOptionalParams,
litellm_params: &LitellmParams,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Result<ValidatedEnvironment, Error> {
@ -618,6 +619,7 @@ mod tests {
Vec::new(),
None,
"anthropic.claude-3",
&MessagesOptionalParams::default(),
&litellm_params,
&|_| None,
)
@ -658,6 +660,7 @@ mod tests {
vec![("authorization".into(), "Bearer forwarded".into())],
api_key,
"anthropic.claude-3",
&MessagesOptionalParams::default(),
&LitellmParams::default(),
&env,
)

22
litellm-rust/crates/llms/src/bedrock_mantle/messages/transformation.rs
View file

@ -6,7 +6,7 @@ use litellm_auth_aws::bedrock_model_id_and_region;
use litellm_core_utils::settings::resolve_non_empty;
use litellm_http::request::{with_header, without_headers};
use litellm_llms_types::{
formats::messages::MessagesRequest,
formats::messages::{MessagesOptionalParams, MessagesRequest},
providers::anthropic::{AnthropicBeta, BetaProvider, BetaSet, DEFAULT_HEADERS},
};
use litellm_router_types::LitellmParams;
@ -59,6 +59,7 @@ impl BaseMessagesConfig for BedrockMantleAnthropicMessagesConfig {
headers: Headers,
api_key: Option<&str>,
model: &str,
_optional_params: &MessagesOptionalParams,
params: &LitellmParams,
env: &dyn Fn(&str) -> Option<String>,
) -> Result<ValidatedEnvironment, Error> {
@ -317,7 +318,14 @@ mod tests {
"https://private.test/anthropic/v1/messages"
);
let environment = BEDROCK_MANTLE_MESSAGES_CONFIG
.validate_environment(vec![], None, model, &params, &|_| None)
.validate_environment(
vec![],
None,
model,
&MessagesOptionalParams::default(),
&params,
&|_| None,
)
.unwrap();
assert!(
matches!(environment.auth, AuthScheme::AwsSigV4 { region: actual, service: "bedrock", .. } if actual == region)
@ -347,7 +355,14 @@ mod tests {
.get_complete_url(None, "claude-test", &params, false, &|_| None)
.unwrap();
let environment = BEDROCK_MANTLE_MESSAGES_CONFIG
.validate_environment(vec![], None, "claude-test", &params, &|_| None)
.validate_environment(
vec![],
None,
"claude-test",
&MessagesOptionalParams::default(),
&params,
&|_| None,
)
.unwrap();
assert_eq!(
url,
@ -401,6 +416,7 @@ mod tests {
],
key,
"claude-test",
&MessagesOptionalParams::default(),
&LitellmParams::default(),
&|name| match name {
"BEDROCK_MANTLE_API_KEY" => Some("mantle".into()),

2
litellm-rust/crates/llms/src/deepseek/messages/transformation.rs
View file

@ -104,6 +104,7 @@ impl BaseMessagesConfig for DeepSeekAnthropicMessagesConfig {
headers: Headers,
api_key: Option<&str>,
_model: &str,
_optional_params: &MessagesOptionalParams,
_litellm_params: &LitellmParams,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Result<ValidatedEnvironment, Error> {
@ -298,6 +299,7 @@ mod tests {
.collect(),
api_key,
"deepseek-v4-pro",
&MessagesOptionalParams::default(),
&LitellmParams::default(),
env,
)

1
litellm-rust/crates/llms/src/edenai/messages/transformation.rs
View file

@ -36,6 +36,7 @@ impl BaseMessagesConfig for EdenAIAnthropicMessagesConfig {
headers: Headers,
api_key: Option<&str>,
_model: &str,
_optional_params: &MessagesOptionalParams,
_params: &LitellmParams,
env: &dyn Fn(&str) -> Option<String>,
) -> Result<ValidatedEnvironment, Error> {

10
litellm-rust/crates/llms/src/minimax/messages/transformation.rs
View file

@ -42,6 +42,7 @@ impl BaseMessagesConfig for MinimaxMessagesConfig {
headers: Headers,
api_key: Option<&str>,
_model: &str,
_optional_params: &MessagesOptionalParams,
_params: &LitellmParams,
env: &dyn Fn(&str) -> Option<String>,
) -> Result<ValidatedEnvironment, Error> {
@ -183,7 +184,14 @@ mod tests {
.flatten()
};
let environment = MINIMAX_MESSAGES_CONFIG
.validate_environment(vec![], key, "test-model", &Default::default(), &env)
.validate_environment(
vec![],
key,
"test-model",
&MessagesOptionalParams::default(),
&Default::default(),
&env,
)
.unwrap();
assert!(matches!(
environment.auth,

1
litellm-rust/crates/llms/src/openrouter/messages/transformation.rs
View file

@ -34,6 +34,7 @@ impl BaseMessagesConfig for OpenRouterAnthropicMessagesConfig {
headers: Headers,
api_key: Option<&str>,
_model: &str,
_optional_params: &MessagesOptionalParams,
_params: &LitellmParams,
env: &dyn Fn(&str) -> Option<String>,
) -> Result<ValidatedEnvironment, Error> {

10
litellm-rust/crates/llms/src/tencent/messages/transformation.rs
View file

@ -43,6 +43,7 @@ impl BaseMessagesConfig for TencentAnthropicMessagesConfig {
headers: Headers,
api_key: Option<&str>,
_model: &str,
_optional_params: &MessagesOptionalParams,
_params: &LitellmParams,
env: &dyn Fn(&str) -> Option<String>,
) -> Result<ValidatedEnvironment, Error> {
@ -229,7 +230,14 @@ mod tests {
.flatten()
};
let environment = TENCENT_MESSAGES_CONFIG
.validate_environment(vec![], key, "test-model", &Default::default(), &env)
.validate_environment(
vec![],
key,
"test-model",
&MessagesOptionalParams::default(),
&Default::default(),
&env,
)
.unwrap();
assert!(matches!(
environment.auth,

2
litellm-rust/crates/llms/src/vertex_ai/messages/transformation.rs
View file

@ -114,6 +114,7 @@ impl BaseMessagesConfig for VertexAiPartnerModelsAnthropicMessagesConfig {
headers: Headers,
_api_key: Option<&str>,
_model: &str,
_optional_params: &MessagesOptionalParams,
litellm_params: &LitellmParams,
_env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Result<ValidatedEnvironment, Error> {
@ -439,6 +440,7 @@ mod tests {
.collect(),
Some("ignored-api-key"),
"claude-sonnet-4-5",
&MessagesOptionalParams::default(),
&params(Some("proj"), None),
&no_env,
)

2
litellm-rust/crates/python-bridge/src/cache/mod.rs vendored
View file

@ -4,7 +4,7 @@ mod selection;
pub(crate) use native::NativeCacheHandle;
pub(crate) use python::{CacheCall, PythonCache};
pub(crate) use selection::{Cached, Selection, admit_native, configure, configured_native};
pub(crate) use selection::{Cached, Selection, configure, configured_native};
use litellm_cache::Error;
use pyo3::{

2
litellm-rust/crates/python-bridge/src/cache/native/v2.rs vendored
View file

@ -279,7 +279,7 @@ pub(in crate::cache) fn configured(
)> {
let handle = native_handle(configured)?.ok_or_else(|| {
pyo3::exceptions::PyRuntimeError::new_err(
"the configured cache changed to a Python cache after native admission",
"native inference needs a v2 cache, the catalog keeps Python caches on Python",
)
})?;
let cache = handle.extract::<PyRef<'_, NativeCacheHandle>>()?;

15
litellm-rust/crates/python-bridge/src/cache/selection.rs vendored
View file

@ -76,21 +76,6 @@ fn selected_cache<'py>(
Ok(Some(configured))
}
pub(crate) fn admit_native(
py: Python<'_>,
kwargs: &Bound<'_, PyDict>,
call_type: &str,
) -> PyResult<()> {
if let Some(configured) = selected_cache(py, kwargs, call_type)?
&& native::v2::native_handle(&configured)?.is_none()
{
return Err(crate::errors::RustBridgeDeclined::new_err(
"the configured cache requires Python inference",
));
}
Ok(())
}
pub(crate) fn configured_native(
py: Python<'_>,
kwargs: &Bound<'_, PyDict>,

1
litellm-rust/crates/python-bridge/src/routes/inference.rs
View file

@ -177,7 +177,6 @@ where
} else {
R::SYNC_CALL_TYPE
};
crate::cache::admit_native(py, &call.kwargs, call_type)?;
let (arguments, hooks) = super::call_hooks(py, R::OPERATION, &call, asynchronous)?;
super::run_public_call(
py,

40
litellm-rust/crates/python-bridge/src/routes/responses/mod.rs
View file

@ -8,59 +8,19 @@ use litellm_auth::AuthServices;
use litellm_cache_response::{CachePolicy, ScopedCache};
use litellm_callbacks_legacy_python::LoggingOperation;
use litellm_host::{call::HostedMachine, protocol::Protocol};
use litellm_host_python::present;
use litellm_inference_responses::{ResponsesRoute, route::Responses};
use litellm_secrets::source::SecretSource;
use pyo3::prelude::*;
use serde_json::Value;
pub(crate) use websocket::ResponsesWebSocketConnection;
use super::{
NativeCall,
inference::{InferenceHost, InferenceRoute, run_inference},
};
use crate::errors::RustBridgeDeclined;
const ROUTE_HOST_MODULE: &str = "litellm.rust_bridge.host.routes.responses";
fn run_responses(py: Python<'_>, call: NativeCall<'_>, asynchronous: bool) -> PyResult<Py<PyAny>> {
let resolved = call.resolved()?;
if let Some(reason) = py
.import(ROUTE_HOST_MODULE)?
.getattr("decline_reason")?
.call1((&resolved,))?
.extract::<Option<String>>()?
{
return Err(RustBridgeDeclined::new_err(reason));
}
let argument = |name: &str| present(&call.kwargs, &call.base, name);
let model = argument("model")?
.ok_or_else(|| pyo3::exceptions::PyValueError::new_err("model is required"))?
.extract::<String>()?;
let provider = argument("custom_llm_provider")?
.map(|value| value.extract::<String>())
.transpose()?;
if provider
.as_deref()
.is_some_and(|provider| provider != "openai")
|| model
.strip_prefix("openai/")
.unwrap_or(&model)
.contains('/')
{
return Err(RustBridgeDeclined::new_err(
"native HTTP responses provider",
));
}
if argument("stream")?
.map(|value| litellm_host_python::from_py::<Value>(&value))
.transpose()?
.is_some_and(|value| value == Value::Bool(true))
{
return Err(RustBridgeDeclined::new_err(
"native Python responses streaming",
));
}
let host = InferenceHost::new(call.view()?, ROUTE_HOST_MODULE);
run_inference::<ResponsesRoute, _>(py, call, asynchronous, ResponsesPythonHost::new(host))
}

25
litellm/__init__.py
View file

@ -7,6 +7,7 @@ _apply_process_defaults()
### INIT VARIABLES #########################
import threading
import os
import datetime
from collections.abc import Mapping, Sequence
from typing import (
@ -62,7 +63,6 @@ from litellm.constants import (
modelscope_models,
empower_models,
together_ai_models,
baseten_models,
WANDB_MODELS,
REPEATED_STREAMING_CHUNK_LIMIT,
request_timeout,
@ -354,6 +354,7 @@ llm_guard_mode: Literal["all", "key-specific", "request-specific"] = "all"
guardrail_name_config_map: Dict[str, GuardrailItem] = {}
include_cost_in_streaming_usage: bool = False
reasoning_auto_summary: bool = False
reasoning_effort_rounding: Literal["up", "down"] = "up"
### PROMPTS ####
from litellm.types.prompts.init_prompts import PromptSpec
@ -550,6 +551,7 @@ from litellm.litellm_core_utils.get_model_cost_map import get_model_cost_map, ma
model_cost = get_model_cost_map(url=model_cost_map_url)
cost_discount_config: Dict[str, float] = {} # Provider-specific cost discounts {"vertex_ai": 0.05} = 5% discount
apply_search_free_tier: bool = False
cost_margin_config: Dict[
str, Union[float, Dict[str, float]]
] = {} # Provider-specific or global cost margins. Examples:
@ -624,6 +626,7 @@ bedrock_models: Set = set()
bedrock_converse_models: Set = set(BEDROCK_CONVERSE_MODELS)
fal_ai_models: Set = set()
fireworks_ai_models: Set = set()
baseten_models: Set[str] = set()
fireworks_ai_embedding_models: Set = set()
deepinfra_models: Set = set()
perplexity_models: Set = set()
@ -742,6 +745,16 @@ def is_openai_finetune_model(key: str) -> bool:
return key.startswith("ft:") and not key.count(":") > 1
def _is_past_deprecation_date(value: Mapping[str, object]) -> bool:
deprecation_date: Final = value.get("deprecation_date")
if not isinstance(deprecation_date, str):
return False
try:
return datetime.date.fromisoformat(deprecation_date) <= datetime.date.today()
except ValueError:
return False
def _populate_provider_model_sets(model_cost_map: Mapping[str, Mapping[str, object]]) -> None:
for key, value in tuple(model_cost_map.items()):
if value.get("litellm_provider") == "openai" and not is_openai_finetune_model(key):
@ -842,6 +855,11 @@ def _populate_provider_model_sets(model_cost_map: Mapping[str, Mapping[str, obje
# ignore the 'up-to', '-to-' model names -> not real models. just for cost tracking based on model params.
if "-to-" not in key and "fireworks-ai-default" not in key:
fireworks_ai_models.add(key)
elif value.get("litellm_provider") == "baseten":
if _is_past_deprecation_date(value):
baseten_models.discard(key)
else:
baseten_models.add(key)
elif value.get("litellm_provider") == "fireworks_ai-embedding-models":
# ignore the 'up-to', '-to-' model names -> not real models. just for cost tracking based on model params.
if "-to-" not in key:
@ -1487,6 +1505,7 @@ from .rag.main import *
from .sandbox.main import *
from .decisions.main import *
from .systemone.main import *
from .extraction.main import *
from .tool_loop import ToolLoopMaxRoundsExceeded, arun_tool_loop, run_tool_loop
from .search.main import *
from .realtime_api.main import (
@ -1693,6 +1712,9 @@ if TYPE_CHECKING:
from .llms.cloudflare.decisions.transformation import (
CloudflareDecisionsConfig as CloudflareDecisionsConfig,
)
from .llms.cloudflare.extraction.transformation import (
CloudflareExtractionConfig as CloudflareExtractionConfig,
)
from .llms.strands_decider.decisions.transformation import (
StrandsDeciderDecisionsConfig as StrandsDeciderDecisionsConfig,
)
@ -2368,6 +2390,7 @@ _AGENT_EXPORTS: Final = frozenset(
"ClaudeCodeOptions",
"CodexOptions",
"OpenCodeOptions",
"PiOptions",
"DeepAgentsOptions",
"ToolLoopOptions",
}

2
litellm/_lazy_imports_registry.py
View file

@ -166,6 +166,7 @@ LLM_CONFIG_NAMES: Final = (
"OpenRouterDecisionsConfig",
"FireworksAIDecisionsConfig",
"CloudflareDecisionsConfig",
"CloudflareExtractionConfig",
"StrandsDeciderDecisionsConfig",
"DatabricksDecisionsConfig",
"HostedVLLMDecisionsConfig",
@ -735,6 +736,7 @@ _LLM_CONFIGS_IMPORT_MAP: Final = {
"OpenRouterDecisionsConfig": (".llms.openrouter.decisions.transformation", "OpenRouterDecisionsConfig"),
"FireworksAIDecisionsConfig": (".llms.fireworks_ai.decisions.transformation", "FireworksAIDecisionsConfig"),
"CloudflareDecisionsConfig": (".llms.cloudflare.decisions.transformation", "CloudflareDecisionsConfig"),
"CloudflareExtractionConfig": (".llms.cloudflare.extraction.transformation", "CloudflareExtractionConfig"),
"StrandsDeciderDecisionsConfig": (
".llms.strands_decider.decisions.transformation",
"StrandsDeciderDecisionsConfig",

2
litellm/a2a_protocol/litellm_completion_bridge/handler.py
View file

@ -80,7 +80,7 @@ class A2ACompletionBridgeHandler:
@staticmethod
def _build_completion_params(
params: dict[str, Any],
litellm_params: Mapping[str, Any],
litellm_params: Mapping[str, object],
api_base: str | None,
agent_extra_headers: Mapping[str, str] | None,
*,

4
litellm/a2a_protocol/main.py
View file

@ -602,9 +602,9 @@ def _build_streaming_logging_obj(
litellm_params: dict[str, object] | None,
metadata: dict[str, object] | None,
proxy_server_request: dict[str, object] | None,
start_time: datetime.datetime,
) -> Logging:
"""Build logging object for streaming A2A requests."""
start_time: Final = datetime.datetime.now()
model: Final = f"a2a_agent/{agent_name}"
logging_obj: Final = Logging(
@ -718,6 +718,7 @@ async def asend_message_streaming(
if request is None:
raise ValueError("request is required")
called_at: Final = datetime.datetime.now()
_raw_logging_obj: Final = kwargs.get("litellm_logging_obj")
logging_obj: Logging | None = _raw_logging_obj if isinstance(_raw_logging_obj, Logging) else None
@ -751,6 +752,7 @@ async def asend_message_streaming(
litellm_params=litellm_params,
metadata=metadata,
proxy_server_request=proxy_server_request,
start_time=called_at,
)
verbose_logger.info("A2A send_message_streaming request_id=%s, agent=%s", request.id, agent_name)

4
litellm/a2a_protocol/streaming_iterator.py
View file

@ -41,7 +41,7 @@ class A2AStreamingIterator:
self.request = request
self.logging_obj = logging_obj
self.agent_name = agent_name
self.start_time = datetime.now()
self.start_time = logging_obj.start_time
# Collect chunks for token counting
self.chunks: list[SendStreamingMessageResponse] = []
@ -54,6 +54,8 @@ class A2AStreamingIterator:
async def __anext__(self) -> "SendStreamingMessageResponse":
try:
chunk: Final = await self.stream.__anext__()
if self.logging_obj.completion_start_time is None:
self.logging_obj.update_completion_start_time(completion_start_time=datetime.now())
# Store chunk
self.chunks.append(chunk)

7
litellm/batches/batch_line_items.py
View file

@ -38,6 +38,7 @@ class BatchLineItemSource:
cost_provider: BatchFileProvider
model_info: ModelInfo | None
metadata: StandardBatchLineItemMetadata
vertex_location: str | None = None
def line_item_metadata(attribution: Mapping[str, object]) -> StandardBatchLineItemMetadata:
@ -140,7 +141,11 @@ def _line_item(
source: BatchLineItemSource, custom_id: str | None, request: Mapping[str, object] | None, entry: dict[str, object]
) -> StandardBatchLineItem:
priced: Final = batch_cost_and_usage(
(entry,), source.cost_provider, model_name=source.model, model_info=source.model_info
(entry,),
source.cost_provider,
model_name=source.model,
model_info=source.model_info,
vertex_location=source.vertex_location,
)
succeeded: Final = priced.successful_requests > 0
return StandardBatchLineItem(

37
litellm/batches/batch_utils.py
View file

@ -77,9 +77,10 @@ async def calculate_batch_cost_and_usage(
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "mistral"],
model_name: str | None = None,
model_info: ModelInfo | None = None,
vertex_location: str | None = None,
) -> BatchCostUsageResult:
"""Calculate the cost and usage of a batch, see batch_cost_and_usage."""
return batch_cost_and_usage(file_content_dictionary, custom_llm_provider, model_name, model_info)
return batch_cost_and_usage(file_content_dictionary, custom_llm_provider, model_name, model_info, vertex_location)
def batch_cost_and_usage(
@ -87,6 +88,7 @@ def batch_cost_and_usage(
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "mistral"],
model_name: str | None = None,
model_info: ModelInfo | None = None,
vertex_location: str | None = None,
) -> BatchCostUsageResult:
"""
Calculate the cost and usage of a batch.
@ -96,16 +98,21 @@ def batch_cost_and_usage(
pricing. Threaded through to batch_cost_calculator so that
deployment-specific pricing (e.g. input_cost_per_token_batches)
is used instead of the global cost map.
vertex_location: The Vertex AI location the batch ran on; non-global
locations apply the model's regional-endpoint uplift multiplier.
"""
first_row: Final = file_content_dictionary[0] if file_content_dictionary else None
if _uses_native_vertex_output(custom_llm_provider, model_name, first_row):
return calculate_vertex_ai_batch_cost_and_usage(file_content_dictionary, model_name, model_info=model_info)
return calculate_vertex_ai_batch_cost_and_usage(
file_content_dictionary, model_name, model_info=model_info, vertex_location=vertex_location
)
return _aggregate_batch_cost_usage_models(
entries=file_content_dictionary,
custom_llm_provider=custom_llm_provider,
model_name=model_name,
model_info=model_info,
vertex_location=vertex_location,
)
@ -115,6 +122,7 @@ async def handle_completed_batch(
model_name: str | None = None,
litellm_params: dict[str, object] | None = None,
model_info: ModelInfo | None = None,
vertex_location: str | None = None,
) -> BatchCostUsageResult:
"""Fetch a completed batch's output file and aggregate its cost, usage, and
models in a single pass over the JSONL lines, so the parsed file content is
@ -128,6 +136,8 @@ async def handle_completed_batch(
model_info: Optional deployment-level model info with custom pricing,
threaded through so a deployment's configured rates win over the
global cost map.
vertex_location: The Vertex AI location the batch ran on; non-global
locations apply the model's regional-endpoint uplift multiplier.
"""
# A completed batch whose request lines all failed has no output file - the
# results are written to a separate error_file_id and output_file_id is None.
@ -155,7 +165,10 @@ async def handle_completed_batch(
output_file_result: Final = (
calculate_vertex_ai_batch_cost_and_usage(
_iter_batch_output_entries(file_content), model_name, model_info=model_info
_iter_batch_output_entries(file_content),
model_name,
model_info=model_info,
vertex_location=vertex_location,
)
if _uses_native_vertex_output(
custom_llm_provider, model_name, next(_iter_batch_output_entries(file_content), None)
@ -165,6 +178,7 @@ async def handle_completed_batch(
custom_llm_provider=custom_llm_provider,
model_name=model_name,
model_info=model_info,
vertex_location=vertex_location,
)
)
@ -203,6 +217,7 @@ def _classify_output_line_stats(
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock", "mistral"],
model_name: str | None,
model_info: ModelInfo | None,
vertex_location: str | None = None,
) -> Iterator[_BatchOutputLineStats | _LineOutcome]:
"""Classify every output line in a single pass, so counting failures never needs
a second read of a potentially huge output file. A line the provider reported as
@ -213,7 +228,7 @@ def _classify_output_line_stats(
if not _batch_response_was_successful(entry, custom_llm_provider):
yield _LineOutcome.PROVIDER_FAILED
continue
stats = _safe_output_line_stats(entry, custom_llm_provider, model_name, model_info)
stats = _safe_output_line_stats(entry, custom_llm_provider, model_name, model_info, vertex_location)
yield stats if stats is not None else _LineOutcome.UNCOSTABLE
@ -222,13 +237,14 @@ def _safe_output_line_stats(
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock", "mistral"],
model_name: str | None,
model_info: ModelInfo | None,
vertex_location: str | None = None,
) -> _BatchOutputLineStats | None:
"""Return the stats for one provider-successful batch output line, or None when
it cannot be costed, so a single bad line never aborts the whole batch's cost
accounting."""
custom_id: Final = entry.get("custom_id") if isinstance(entry, dict) else None
try:
return _compute_output_line_stats(entry, custom_llm_provider, model_name, model_info)
return _compute_output_line_stats(entry, custom_llm_provider, model_name, model_info, vertex_location)
except Exception as e: # noqa: BLE001 # any single line's costing failure must not abort the whole batch
verbose_logger.warning(
"batch output line could not be costed, so it is billed at $0 and the rest of the batch "
@ -244,6 +260,7 @@ def _compute_output_line_stats(
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock", "mistral"],
model_name: str | None,
model_info: ModelInfo | None,
vertex_location: str | None = None,
) -> _BatchOutputLineStats:
response_body: Final = _get_response_from_batch_job_output_file(entry, custom_llm_provider)
usage: Final = _get_batch_job_usage_from_response_body(response_body, custom_llm_provider)
@ -258,6 +275,7 @@ def _compute_output_line_stats(
model_name=model_name,
response_model=response_model,
model_info=model_info,
vertex_location=vertex_location,
)
return _BatchOutputLineStats(
prompt_cost=line_prompt_cost,
@ -287,6 +305,7 @@ def _output_line_cost(
model_name: str | None,
response_model: str | None,
model_info: ModelInfo | None,
vertex_location: str | None = None,
) -> tuple[float, float]:
"""(prompt_cost, completion_cost) for one output line, priced at batch rates."""
from litellm.cost_calculator import batch_cost_calculator, ocr_batch_cost
@ -307,6 +326,7 @@ def _output_line_cost(
model=cost_model,
custom_llm_provider=custom_llm_provider,
model_info=model_info,
vertex_location=vertex_location,
)
@ -315,11 +335,14 @@ def _aggregate_batch_cost_usage_models(
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock", "mistral"],
model_name: str | None = None,
model_info: ModelInfo | None = None,
vertex_location: str | None = None,
) -> BatchCostUsageResult:
"""Aggregate cost, usage, models, and pass/fail counts from batch output
entries in a single pass, holding one small stats record per line instead
of the parsed file."""
all_results: Final = tuple(_classify_output_line_stats(entries, custom_llm_provider, model_name, model_info))
all_results: Final = tuple(
_classify_output_line_stats(entries, custom_llm_provider, model_name, model_info, vertex_location)
)
line_stats: Final = tuple(result for result in all_results if isinstance(result, _BatchOutputLineStats))
failed_requests: Final = sum(1 for result in all_results if result is _LineOutcome.PROVIDER_FAILED)
successful_requests: Final = len(all_results) - failed_requests
@ -366,6 +389,7 @@ def calculate_vertex_ai_batch_cost_and_usage(
vertex_ai_batch_responses: Iterable[dict[str, object]],
model_name: str | None = None,
model_info: ModelInfo | None = None,
vertex_location: str | None = None,
) -> BatchCostUsageResult:
"""
Cost and usage of a native Vertex predictions.jsonl, one
@ -384,6 +408,7 @@ def calculate_vertex_ai_batch_cost_and_usage(
model_info=model_info,
calculate_usage=VertexGeminiConfig.calculate_usage,
cost_calculator=batch_cost_calculator,
vertex_location=vertex_location,
)
for row in vertex_ai_batch_responses
)

26
litellm/caching/dual_cache.py
View file

@ -59,16 +59,6 @@ class PendingBatchRead:
previous_access_times: dict[str, float | None]
@dataclass(frozen=True, slots=True)
class DeclaredBatchRead:
"""A ``async_batch_get_cache`` split in two: the memory half done, the Redis half declared on a ``RedisBatch``
so it rides that batch's next round trip, resolved later with ``async_resolve_batch_get``."""
keys: tuple[str, ...]
pending: PendingBatchRead
result: BatchResult[Mapping[str, object]] | None
def _log_deferred_increment_failure(future: asyncio.Future[float]) -> None:
if future.cancelled():
return
@ -382,22 +372,6 @@ class DualCache(BaseCache):
await self.in_memory_cache.async_set_cache(key, value, **self._backfill_kwargs(kwargs))
return merged
async def declare_batch_get(self, keys: Sequence[str], batch: RedisBatch) -> DeclaredBatchRead:
pending: Final = await self._prepare_batch_get(
list(keys),
local_only=False,
throttle_redis=False,
)
return DeclaredBatchRead(
keys=tuple(keys),
pending=pending,
result=batch.mget(pending.redis_keys) if pending.redis_keys else None,
)
async def async_resolve_batch_get(self, declared: DeclaredBatchRead) -> list[object | None]:
redis_result: Final = None if declared.result is None else await declared.result
return await self._apply_batch_get(declared.pending, redis_result)
async def async_batch_get_cache(
self,
keys: list,

47
litellm/caching/redis_cache.py
View file

@ -134,7 +134,12 @@ _INCREMENT_WITH_FLOOR_LUA: Final = (
"return count"
)
_INCREMENT_IF_PRESENT_LUA: Final = (
"if redis.call('EXISTS', KEYS[1]) == 0 then return false end return redis.call('INCRBYFLOAT', KEYS[1], ARGV[1])"
)
_LUA_COUNT: Final = TypeAdapter(int)
_LUA_OPTIONAL_TOTAL: Final[TypeAdapter[float | None]] = TypeAdapter(float | None)
_OPTIONAL_COUNTS: Final = TypeAdapter(tuple[int | None, ...])
@ -675,7 +680,7 @@ class RedisCache(BaseCache):
# redis namespaces
self.namespace = namespace
# for high traffic, we store the redis results in memory and then batch write to redis
self.redis_batch_writing_buffer: list = []
self.redis_batch_writing_buffer: list[tuple[str, object]] = []
if redis_flush_size is None:
self.redis_flush_size: int = 100
else:
@ -1568,6 +1573,46 @@ class RedisCache(BaseCache):
count: Final = await _redis_client.eval(_INCREMENT_WITH_FLOOR_LUA, 1, namespaced_key, value, ttl)
return _LUA_COUNT.validate_python(count)
@_redis_circuit_breaker_guard
async def async_increment_if_present(
self, key: str, value: float, parent_otel_span: Span | None = None
) -> float | None:
_redis_client: Final = self._async_commands()
namespaced_key: Final = self.check_and_fix_namespace(key=key)
start_time: Final = time.time()
try:
total: Final = _LUA_OPTIONAL_TOTAL.validate_python(
await _redis_client.eval(_INCREMENT_IF_PRESENT_LUA, 1, namespaced_key, value)
)
except Exception as e:
failed_at: Final = time.time()
asyncio.create_task(
self.service_logger_obj.async_service_failure_hook(
service=ServiceTypes.REDIS,
duration=failed_at - start_time,
error=e,
call_type="async_increment_if_present",
caller=_get_call_stack_info(),
start_time=start_time,
end_time=failed_at,
parent_otel_span=parent_otel_span,
)
)
raise
end_time: Final = time.time()
asyncio.create_task(
self.service_logger_obj.async_service_success_hook(
service=ServiceTypes.REDIS,
duration=end_time - start_time,
call_type="async_increment_if_present",
caller=_get_call_stack_info(),
start_time=start_time,
end_time=end_time,
parent_otel_span=parent_otel_span,
)
)
return total
async def flush_cache_buffer(self):
print_verbose(f"flushing to redis....reached size of buffer {len(self.redis_batch_writing_buffer)}")
await self.async_set_cache_pipeline(self.redis_batch_writing_buffer)

35
litellm/chat_completions/dispatch.py
View file

@ -1,15 +1,19 @@
import inspect
from collections.abc import Awaitable, Callable, Coroutine, Mapping
from types import MappingProxyType
from typing import Final, TypeAlias, cast # noqa: TID251 # native binding selects a sync result or an async awaitable
from litellm import main
from litellm.rust_bridge.catalog import Route, RouteContext
from litellm.rust_bridge import detectors
from litellm.rust_bridge.catalog import Feature, Route, RouteContext
from litellm.rust_bridge.chat_completions import detectors as route_detectors
from litellm.rust_bridge.chat_completions.entrypoints import (
NATIVE_ACOMPLETION,
NATIVE_COMPLETION,
)
from litellm.rust_bridge.dispatch import PublicDispatch
from litellm.rust_bridge.public_call import (
AsyncFlag,
NativeCall,
bind,
native_call,
@ -44,6 +48,7 @@ def _python_acompletion() -> PythonAcompletion:
_PYTHON_COMPLETION: Final = _python_completion()
_COMPLETION: Final = signature(_PYTHON_COMPLETION)
_FLAG: Final = AsyncFlag("acompletion", sync=_COMPLETION)
_PYTHON_ACOMPLETION: Final = _python_acompletion()
_ACOMPLETION: Final = signature(_PYTHON_ACOMPLETION)
@ -61,11 +66,32 @@ def _public_request(
return native_call(legacy, args, kwargs)
_FEATURES: Final[detectors.Detectors] = MappingProxyType(
{
**detectors.COMMON,
Feature.MOCK_RESPONSE: route_detectors.chat_mock_response,
Feature.AGENTIC_LOOP_HOOK: route_detectors.chat_agentic_loop_hook,
Feature.RESPONSE_CACHE: detectors.response_cache,
Feature.STREAMING: detectors.streaming,
Feature.DROP_PARAMS: detectors.drop_params,
Feature.MODIFY_PARAMS: detectors.modify_params,
Feature.EXTRA_BODY: detectors.extra_body,
Feature.RETRIES_FALLBACKS: route_detectors.chat_retries_fallbacks,
Feature.PROMPT_MANAGEMENT: detectors.prompt_management,
Feature.CACHE_CONTROL_INJECTION: detectors.cache_control_injection,
Feature.MODEL_PARAM_GATING: route_detectors.chat_model_param_gating,
Feature.MODEL_LIST: route_detectors.model_list,
Feature.MODEL_ALIAS_MAP: route_detectors.model_alias_map,
}
)
def _context(request: NativeCall) -> RouteContext:
return RouteContext(
Route.CHAT_COMPLETIONS,
provider=optional_str(request.resolved.get("custom_llm_provider")),
model=str(request.resolved["model"]),
features=detectors.detected(_FEATURES, request),
)
@ -73,7 +99,6 @@ _DISPATCH: Final = PublicDispatch(
route=Route.CHAT_COMPLETIONS,
request=lambda args, kwargs: _public_request(_COMPLETION, args, kwargs),
context=_context,
bypass=lambda request: request.kwargs.get("acompletion") is True,
)
_ADISPATCH: Final = PublicDispatch(
@ -88,9 +113,11 @@ def completion(
**kwargs: object, # kwargs-ok: preserve the public chat completions call shape
) -> ChatResult | Coroutine[object, object, ChatResult]:
python: Final = _PYTHON_COMPLETION
if _FLAG.requested(kwargs):
return acompletion(**_FLAG.keywords(args, kwargs))
return _DISPATCH.run(
args,
kwargs,
_FLAG.dropped(kwargs),
python=python,
binding=NATIVE_COMPLETION,
native=native_call_hook,
@ -101,7 +128,7 @@ async def acompletion(*args: object, **kwargs: object) -> ChatResult: # kwargs-
python: Final = _PYTHON_ACOMPLETION
return await _ADISPATCH.arun(
args,
kwargs,
_FLAG.dropped(kwargs),
python=python,
binding=NATIVE_ACOMPLETION,
native=native_call_hook,

42
litellm/constants.py
View file

@ -25,6 +25,8 @@ CLAUDE_CODE_CLIENT: Final = "claude-code"
MICROSOFT_365_COPILOT_DEFAULT_TIME_ZONE: Final = "UTC"
GITHUB_COPILOT_PER_USER_AUTH_TYPE: Final = "per_user_oauth"
GITHUB_COPILOT_AUTH_TYPE_KEY: Final = "github_copilot_auth_type"
GITHUB_COPILOT_HOST_KEY: Final = "github_copilot_host"
GITHUB_COPILOT_DEFAULT_HOST: Final = "github.com"
GITHUB_COPILOT_USER_TOKEN_SAFETY_MARGIN_SECONDS: Final = 60
GITHUB_COPILOT_USER_CREDENTIAL_CACHE_TTL_SECONDS: Final = 60
GITHUB_COPILOT_DEVICE_FLOW_CACHE_PREFIX: Final = "github_copilot_device_flow"
@ -40,6 +42,14 @@ class ServerStreamingClassification(str, Enum):
SERVER_STREAMING_CLASSIFICATION_MARKER: Final = ServerStreamingClassification.MARKER
DEFER_PYDANTIC_BUILD: Final = os.getenv("DEFER_PYDANTIC_BUILD", "true") in ("true", "1", "on")
DEFAULT_HEALTH_CHECK_PROMPT: Final = str(os.getenv("DEFAULT_HEALTH_CHECK_PROMPT", "test from litellm"))
KUBERNETES_POD_DISCOVERY_REFRESH_INTERVAL_SECONDS: Final = get_env_float(
"KUBERNETES_POD_DISCOVERY_REFRESH_INTERVAL_SECONDS", 5.0
)
KUBERNETES_POD_DISCOVERY_IDLE_EVICTION_SECONDS: Final = get_env_float(
"KUBERNETES_POD_DISCOVERY_IDLE_EVICTION_SECONDS", 300.0
)
KUBERNETES_POD_DISCOVERY_DNS_TIMEOUT_SECONDS: Final = 2.0
KUBERNETES_POD_ROUTING_KEY: Final = "kubernetes_pod_routing"
AZURE_DEFAULT_RESPONSES_API_VERSION: Final = str(os.getenv("AZURE_DEFAULT_RESPONSES_API_VERSION", "preview"))
AZURE_OPENAI_AUDIO_PROVIDERS: Final = frozenset({"azure", "azure_ai"})
ROUTER_MAX_FALLBACKS: Final = get_env_int("ROUTER_MAX_FALLBACKS", 5)
@ -183,6 +193,11 @@ DEFAULT_IMAGE_HEIGHT: Final = get_env_int("DEFAULT_IMAGE_HEIGHT", 300)
# Maps to OpenAI's 50 MB payload limit - requests with images exceeding this size will be rejected
# Set MAX_IMAGE_URL_DOWNLOAD_SIZE_MB=0 to disable image URL handling entirely
MAX_IMAGE_URL_DOWNLOAD_SIZE_MB: Final = get_env_float("MAX_IMAGE_URL_DOWNLOAD_SIZE_MB", 50.0)
MAX_DECOMPRESSED_REQUEST_BODY_SIZE_MB: Final = get_env_float("MAX_DECOMPRESSED_REQUEST_BODY_SIZE_MB", 50.0)
MAX_REQUEST_BODY_DECOMPRESSION_RATIO: Final = get_env_int("MAX_REQUEST_BODY_DECOMPRESSION_RATIO", 200)
MAX_DECOMPRESSED_REQUEST_BODIES_IN_FLIGHT_MB: Final = get_env_float(
"MAX_DECOMPRESSED_REQUEST_BODIES_IN_FLIGHT_MB", 200.0
)
MAX_SIZE_PER_ITEM_IN_MEMORY_CACHE_IN_KB: Final = get_env_int(
"MAX_SIZE_PER_ITEM_IN_MEMORY_CACHE_IN_KB", 1024
) # 1MB = 1024KB
@ -243,7 +258,17 @@ MCP_PER_USER_TOKEN_EXPIRY_BUFFER_SECONDS: Final = get_env_int("MCP_PER_USER_TOKE
# MCP timeout defaults (seconds). Override via env vars for slow/custom MCP servers.
MCP_CLIENT_TIMEOUT: Final = get_env_float("LITELLM_MCP_CLIENT_TIMEOUT", 60.0)
MCP_MAX_SSE_EVENT_SIZE_BYTES: Final = get_env_int_in_range(
"LITELLM_MCP_MAX_SSE_EVENT_SIZE_BYTES", 64 * 1024 * 1024, 1, sys.maxsize
)
MCP_TOOL_LISTING_TIMEOUT: Final = get_env_float("LITELLM_MCP_TOOL_LISTING_TIMEOUT", 30.0)
MCP_UPSTREAM_RATE_LIMIT_MAX_RETRIES: Final = get_env_int("LITELLM_MCP_UPSTREAM_RATE_LIMIT_MAX_RETRIES", 3)
MCP_UPSTREAM_RATE_LIMIT_BASE_DELAY_SECONDS: Final = get_env_float(
"LITELLM_MCP_UPSTREAM_RATE_LIMIT_BASE_DELAY_SECONDS", 0.5
)
MCP_UPSTREAM_RATE_LIMIT_MAX_DELAY_SECONDS: Final = get_env_float(
"LITELLM_MCP_UPSTREAM_RATE_LIMIT_MAX_DELAY_SECONDS", 5.0
)
MCP_METADATA_TIMEOUT: Final = get_env_float("LITELLM_MCP_METADATA_TIMEOUT", 10.0)
MCP_HEALTH_CHECK_TIMEOUT: Final = get_env_float("LITELLM_MCP_HEALTH_CHECK_TIMEOUT", 10.0)
AGENT_HEALTH_CHECK_TIMEOUT_SECONDS: Final = get_env_float("LITELLM_AGENT_HEALTH_CHECK_TIMEOUT", 5.0)
@ -1034,7 +1059,7 @@ openai_compatible_endpoints: Final[list[str]] = [
]
openai_compatible_providers: Final[list] = [
openai_compatible_providers: Final[list[str]] = [
"anyscale",
"groq",
"nvidia_nim",
@ -1256,14 +1281,6 @@ together_ai_models: Final[set] = set(
# supports all together ai models, just pass in the model id e.g. completion(model="together_computer/replit_code_3b",...)
baseten_models: Final[set] = set(
[
"qvv0xeq",
"q841o8w",
"31dxrj3",
]
) # FALCON 7B # WizardLM # Mosaic ML
featherless_ai_models: Final[set] = set(
[
"featherless-ai/Qwerky-72B",
@ -1633,6 +1650,7 @@ OUTPUT_TOKEN_CEILING_PARAMS: Final = frozenset({"max_tokens", "max_completion_to
CLIENT_OUTPUT_CEILING_METADATA_KEY: Final = "_client_output_ceiling"
CONSUMED_REQUEST_TAGS_METADATA_KEY: Final = "_consumed_request_tags"
ROUTING_REQUEST_TAGS_METADATA_KEY: Final = "_routing_request_tags"
BAGGAGE_MAPPED_TAGS_METADATA_KEY: Final = "_baggage_mapped_tags"
VECTOR_STORE_CALLER_METADATA_KWARG: Final = "litellm_vector_store_caller_metadata"
ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY: Final = "_litellm_router_usage_counted_tokens"
INTERNAL_CALL_ORIGIN_METADATA_KEY: Final = "internal_call_origin"
@ -1799,6 +1817,7 @@ LITELLM_EXECUTED_BATCH_CONCURRENCY: Final = max(1, get_env_int("LITELLM_EXECUTED
########################### CLI SSO AUTHENTICATION CONSTANTS ###########################
LITELLM_CLI_SOURCE_IDENTIFIER: Final = "litellm-cli"
LITELLM_CLI_SESSION_TOKEN_PREFIX: Final = "litellm-session-token"
CP_RETURN_TO_COOKIE_PREFIX: Final = "return_to.v1."
CLI_SSO_SESSION_CACHE_KEY_PREFIX: Final = "cli_sso_session"
CLI_SSO_SESSION_TTL_SECONDS: Final = 600
CLI_SESSION_KEY_PREFIX: Final = "cli-session"
@ -1842,6 +1861,9 @@ SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS: Final = get_env_float(
SPEND_LOG_CLEANUP_RUN_BUDGET_SECONDS: Final = get_env_float("SPEND_LOG_CLEANUP_RUN_BUDGET_SECONDS", 300.0)
SPEND_LOG_CLEANUP_BATCH_TIMEOUT_SECONDS: Final = get_env_float("SPEND_LOG_CLEANUP_BATCH_TIMEOUT_SECONDS", 30.0)
SPEND_LOG_CLEANUP_REMAINING_COUNT_CAP: Final = get_env_int("SPEND_LOG_CLEANUP_REMAINING_COUNT_CAP", 100000)
SPEND_LOG_SIZE_CHECK_JOB_NAME: Final = "spend_log_size_check"
SPEND_LOG_SIZE_CHECK_INTERVAL: Final = os.getenv("SPEND_LOG_SIZE_CHECK_INTERVAL", "1d")
SPEND_LOG_SIZE_CLEANUP_RECHECK_BATCHES: Final = get_env_int("SPEND_LOG_SIZE_CLEANUP_RECHECK_BATCHES", 10)
SCHEDULED_JOB_SHUTDOWN_FINISH_TIMEOUT_SECONDS: Final = get_env_float(
"SCHEDULED_JOB_SHUTDOWN_FINISH_TIMEOUT_SECONDS", 5.0
)
@ -1963,6 +1985,7 @@ LITELLM_SETTINGS_SAFE_DB_OVERRIDES: Final = [
"public_model_groups",
"public_model_groups_links",
"cost_discount_config",
"apply_search_free_tier",
"cost_margin_config",
"block_requests_for_models_without_pricing",
"budget_exceeded_throttle_percentage",
@ -2303,6 +2326,7 @@ HARNESS_STDERR_TAIL_LINES: Final = 40
HARNESS_STREAM_READ_CHUNK_BYTES: Final = 64 * 1024
HARNESS_EVENT_QUEUE_MAX_SIZE: Final = 1024
HARNESS_PROCESS_KILL_GRACE_SECONDS: Final = 5.0
HARNESS_RESUME_COMMAND_TIMEOUT_SECONDS: Final = 30.0
HARNESS_SNAPSHOT_SKIP_DIRS: Final = frozenset(
{
".git",

153
litellm/cost_calculator.py
View file

@ -45,6 +45,7 @@ from litellm.litellm_core_utils.llm_cost_calc.utils import (
get_regional_uplift_multiplier,
get_service_tier_cost_key,
get_token_type_cost_breakdown,
get_vertex_regional_endpoint_uplift,
parse_prompt_tokens_details,
select_cost_metric_for_model,
)
@ -111,6 +112,7 @@ from litellm.llms.xai.cost_calculator import cost_per_token as xai_cost_per_toke
from litellm.responses.utils import ResponseAPILoggingUtils
from litellm.types.agents import LiteLLMSendMessageResponse
from litellm.types.decisions import DecisionsResponse, DecisionsUsage, OpenAIDecisionResponse, OpenAIDecisionUsage
from litellm.types.extraction import ExtractionResponse, ExtractionUsage
from litellm.types.llms.base import CachedTokensDetails, LiteLLMBaseModel
from litellm.types.llms.openai import (
HttpxBinaryResponseContent,
@ -137,6 +139,7 @@ from litellm.types.utils import (
TranscriptionUsageTokensObject,
Usage,
VectorStoreSearchResponse,
is_server_derived_pricing_key,
)
from litellm.utils import (
CallTypes,
@ -149,6 +152,8 @@ from litellm.utils import (
TextCompletionResponse,
TranscriptionResponse,
cached_get_model_info_helper,
model_info_with_pricing_overrides,
runtime_model_cost_registrations,
token_counter,
unmapped_model_generalization,
)
@ -437,6 +442,7 @@ def cost_per_token(
### REQUEST MODEL ###
request_model: str | None = None, # original request model for router detection
custom_model_info: OCRPricing | None = None,
bedrock_model_info: ModelInfo | None = None,
) -> tuple[float, float]:
"""
Calculates the cost per token for a given model, prompt tokens, and completion tokens.
@ -676,6 +682,7 @@ def cost_per_token(
model=model,
custom_llm_provider=custom_llm_provider,
data_residency=data_residency,
vertex_location=vertex_location,
)
elif call_type == "atranscription" or call_type == "transcription":
if _transcription_usage_has_token_details(usage_block):
@ -738,6 +745,15 @@ def cost_per_token(
elif custom_llm_provider == "anthropic":
return anthropic_cost_per_token(model=model, usage=usage_block, service_tier=service_tier)
elif custom_llm_provider == "bedrock":
resolved_bedrock_info: Final = bedrock_model_info or _bedrock_pricing_info(model, None, "bedrock", False)
if resolved_bedrock_info is not None:
return generic_cost_per_token(
model=model,
usage=usage_block,
custom_llm_provider="bedrock",
service_tier=service_tier,
model_info=resolved_bedrock_info,
)
return bedrock_cost_per_token(model=model, usage=usage_block, service_tier=service_tier)
elif custom_llm_provider == "openai":
return openai_cost_per_token(
@ -860,11 +876,46 @@ def _cost_map_entry_prices_anything(entry: Mapping[str, object]) -> bool:
)
def _declares_rates(model_info: ModelInfo | None) -> bool:
entry: Final = _raw_cost_map_entry(model_info["key"]) if model_info is not None else None
def _bedrock_registration_has_prices(entry: Mapping[str, object]) -> bool:
return _cost_map_entry_prices_anything(entry) or any(
value is not None and field.startswith("cache_") and is_server_derived_pricing_key(field)
for field, value in entry.items()
)
def _bedrock_pricing_info(
model: str | None,
invoked_model_id: str | None,
custom_llm_provider: str | None,
explicit_pricing: bool,
) -> ModelInfo | None:
if explicit_pricing or custom_llm_provider != "bedrock" or not model:
return None
registration: Final = next(
(
(key, entry)
for key, entry in runtime_model_cost_registrations(model, "bedrock")
if _bedrock_registration_has_prices(entry)
),
None,
)
profile: Final = _cost_map_model_info(invoked_model_id, "bedrock") if invoked_model_id else None
if profile is None or not _declares_rates(profile):
return _cost_map_model_info(registration[0], None) if registration is not None else None
if registration is None or registration[0] == profile["key"]:
return profile
return model_info_with_pricing_overrides(profile, registration[1])
def cost_map_key_prices_anything(key: str) -> bool:
entry: Final = _raw_cost_map_entry(key)
return entry is not None and _cost_map_entry_prices_anything(entry)
def _declares_rates(model_info: ModelInfo | None) -> bool:
return model_info is not None and cost_map_key_prices_anything(model_info["key"])
def _base_model_pricing_provider(base_model: str, custom_llm_provider: str | None) -> str | None:
"""The provider to price an explicit base_model under, when the deployment's own provider has
no rates for it (e.g. chatgpt/ subscription rows) but the base model's own row does"""
@ -886,6 +937,16 @@ def _pricing_provider_for_cost_candidate(
return current_provider if isinstance(current_provider, str) else None
def _invoked_model_id_for_cost(
optional_params: object,
logging_obj: LitellmLoggingObject | None,
) -> str | None:
recorded: Final[object] = getattr(logging_obj, "model_call_details", None)
return _get_hidden_str_for_cost_calc(optional_params, "model_id") or _get_hidden_str_for_cost_calc(
recorded, "model_id"
)
def select_model_name_for_cost_calc(
model: str | None,
completion_response: object | None,
@ -894,6 +955,7 @@ def select_model_name_for_cost_calc(
custom_llm_provider: str | None = None,
router_model_id: str | None = None,
region_name: str | None = None,
invoked_model_id: str | None = None,
) -> str | None:
"""
1. If custom pricing is true, return received model name
@ -914,6 +976,13 @@ def select_model_name_for_cost_calc(
hidden_params: Final[dict | None] = getattr(completion_response, "_hidden_params", None)
provider_response_model: Final = _get_hidden_str_for_cost_calc(hidden_params, "provider_response_model")
explicit_pricing: Final = custom_pricing is True or base_model is not None
configured_model: Final = model or (
completion_response_model if isinstance(completion_response_model, str) else None
)
bedrock_model_info: Final = _bedrock_pricing_info(
configured_model, invoked_model_id, custom_llm_provider, explicit_pricing
)
bedrock_pricing_model: Final = bedrock_model_info["key"] if bedrock_model_info is not None else None
priced_from_response: Final = provider_response_model is not None or completion_response_model is not None
priced_region: Final = (
_get_hidden_str_for_cost_calc(hidden_params, "region_name") or region_name
@ -931,8 +1000,8 @@ def select_model_name_for_cost_calc(
else:
return_model = model
elif base_model is not None or provider_response_model is not None:
return_model = base_model if base_model is not None else provider_response_model
elif base_model is not None or bedrock_pricing_model is not None or provider_response_model is not None:
return_model = base_model if base_model is not None else bedrock_pricing_model or provider_response_model
elif completion_response_model is None and hidden_params is not None:
if hidden_params.get("model", None) is not None and len(hidden_params["model"]) > 0:
@ -1103,6 +1172,10 @@ def get_usage_object(
if isinstance(completion_response, (DecisionsResponse, OpenAIDecisionResponse)):
return None if completion_response.usage is None else _decisions_usage(completion_response.usage)
if isinstance(completion_response, ExtractionResponse):
return _extraction_usage(completion_response.usage)
if isinstance(completion_response, HttpxBinaryResponseContent):
return completion_response.usage
if usage_obj is None:
return None
if isinstance(usage_obj, Usage):
@ -1157,6 +1230,14 @@ def _decisions_usage(usage: DecisionsUsage | OpenAIDecisionUsage) -> Usage:
)
def _extraction_usage(usage: ExtractionUsage) -> Usage:
return Usage(
prompt_tokens=usage.input_tokens,
completion_tokens=usage.output_tokens,
total_tokens=usage.input_tokens + usage.output_tokens,
)
def _is_known_usage_objects(usage_obj):
"""Returns True if the usage obj is a known Usage type"""
return (
@ -1164,6 +1245,7 @@ def _is_known_usage_objects(usage_obj):
or isinstance(usage_obj, ResponseAPIUsage)
or isinstance(usage_obj, DecisionsUsage)
or isinstance(usage_obj, OpenAIDecisionUsage)
or isinstance(usage_obj, ExtractionUsage)
or TranscriptionUsageObjectTransformation.is_transcription_usage_object(usage_obj)
)
@ -1193,6 +1275,8 @@ def infer_call_type(call_type: CallTypesLiteral | None, completion_response: obj
return "send_message"
elif isinstance(completion_response, (DecisionsResponse, OpenAIDecisionResponse)):
return "adecisions"
elif isinstance(completion_response, ExtractionResponse):
return "aextraction"
return call_type
@ -1441,6 +1525,7 @@ def completion_cost(
litellm_model_name: str | None = None,
router_model_id: str | None = None,
litellm_logging_obj: LitellmLoggingObject | None = None,
pricing_provider: str | None = None,
### SERVICE TIER ###
service_tier: str | None = None, # for OpenAI service tier pricing
### DATA RESIDENCY ###
@ -1461,6 +1546,8 @@ def completion_cost(
"""
try:
call_type = infer_call_type(call_type, completion_response) or "completion"
if pricing_provider is not None:
custom_llm_provider = pricing_provider # rebind-ok: model_info.pricing_provider overrides the routing provider for cost-map lookups
if call_type == CallTypes.aresponses_websocket.value and isinstance(
completion_response, LiteLLMRealtimeStreamLoggingObject
@ -1486,6 +1573,7 @@ def completion_cost(
optional_params=optional_params,
custom_pricing=custom_pricing,
base_model=base_model,
pricing_provider=pricing_provider,
standard_built_in_tools_params=standard_built_in_tools_params,
litellm_model_name=litellm_model_name,
router_model_id=router_model_id,
@ -1544,6 +1632,10 @@ def completion_cost(
if base_model is not None and custom_pricing is not True
else None
)
invoked_model_id: Final = _invoked_model_id_for_cost(
_LOGGED_LITELLM_PARAMS.validate_python(optional_params) if optional_params is not None else None,
litellm_logging_obj,
)
selected_model: Final = select_model_name_for_cost_calc(
model=model,
completion_response=completion_response,
@ -1552,8 +1644,15 @@ def completion_cost(
base_model=base_model,
router_model_id=router_model_id,
region_name=region_name,
invoked_model_id=invoked_model_id,
)
bedrock_model_info: Final = _bedrock_pricing_info(
model or _get_response_model(completion_response),
invoked_model_id,
_get_provider_for_cost_calc(model, custom_llm_provider),
explicit_pricing,
)
requested_model: Final = model
request_pricing_provider: Final = custom_llm_provider if isinstance(custom_llm_provider, str) else None
raw_candidates: Final = (
@ -1565,10 +1664,10 @@ def completion_cost(
*(((model, request_pricing_provider),) if model is not None else ()),
)
potential_model_names: Final = (
raw_candidates if explicit_pricing else _priced_candidates_before_unpriced_rules(raw_candidates)
raw_candidates if custom_pricing is True else _priced_candidates_before_unpriced_rules(raw_candidates)
)
for idx, (model, pricing_provider) in enumerate(potential_model_names):
for idx, (model, candidate_pricing_provider) in enumerate(potential_model_names):
try:
if verbose_logger.isEnabledFor(logging.DEBUG):
verbose_logger.debug("selected model name for cost calculation: %s", model)
@ -1589,6 +1688,8 @@ def completion_cost(
)
if isinstance(usage_obj, (DecisionsUsage, OpenAIDecisionUsage)):
_usage = _decisions_usage(usage_obj).model_dump()
elif isinstance(usage_obj, ExtractionUsage):
_usage = _extraction_usage(usage_obj).model_dump()
elif usage_obj is None:
_usage = {}
elif isinstance(usage_obj, BaseModel):
@ -1635,7 +1736,9 @@ def completion_cost(
hidden_params = getattr(completion_response, "_hidden_params", None)
if hidden_params is not None:
custom_llm_provider = hidden_params.get("custom_llm_provider", custom_llm_provider or None)
custom_llm_provider = pricing_provider or hidden_params.get(
"custom_llm_provider", custom_llm_provider or None
)
region_name = hidden_params.get("region_name", region_name)
else:
@ -1921,7 +2024,7 @@ def completion_cost(
prompt_tokens=prompt_tokens or 0,
completion_tokens=completion_tokens or 0,
custom_llm_provider=_pricing_provider_for_cost_candidate(
candidate_provider=pricing_provider,
candidate_provider=candidate_pricing_provider,
request_provider=request_pricing_provider,
current_provider=custom_llm_provider,
),
@ -1942,6 +2045,7 @@ def completion_cost(
vertex_location=vertex_location,
response=completion_response,
custom_model_info=_ocr_model_info(litellm_logging_obj, custom_pricing, router_model_id),
bedrock_model_info=bedrock_model_info,
)
# Get additional costs from provider (e.g., routing fees, infrastructure costs)
@ -2024,7 +2128,7 @@ def completion_cost(
_token_type_breakdown = get_token_type_cost_breakdown(
model=model,
custom_llm_provider=_pricing_provider_for_cost_candidate(
candidate_provider=pricing_provider,
candidate_provider=candidate_pricing_provider,
request_provider=request_pricing_provider,
current_provider=custom_llm_provider,
),
@ -2033,6 +2137,7 @@ def completion_cost(
data_residency=data_residency,
vertex_location=vertex_location,
custom_cost_per_token=custom_cost_per_token,
model_info=bedrock_model_info,
)
_reasoning_cost = _token_type_breakdown.reasoning_cost
_cache_read_cost = _token_type_breakdown.cache_read_cost
@ -2110,7 +2215,8 @@ def response_cost_calculator(
| Response
| SearchResponse
| DecisionsResponse
| OpenAIDecisionResponse,
| OpenAIDecisionResponse
| ExtractionResponse,
model: str,
custom_llm_provider: str | None,
call_type: Literal[
@ -2136,6 +2242,8 @@ def response_cost_calculator(
"adecisions",
"systemone",
"asystemone",
"extraction",
"aextraction",
],
optional_params: dict,
cache_hit: bool | None = None,
@ -2146,6 +2254,7 @@ def response_cost_calculator(
litellm_model_name: str | None = None,
router_model_id: str | None = None,
litellm_logging_obj: LitellmLoggingObject | None = None,
pricing_provider: str | None = None,
### SERVICE TIER ###
service_tier: str | None = None, # for OpenAI service tier pricing
### DATA RESIDENCY ###
@ -2186,6 +2295,7 @@ def response_cost_calculator(
litellm_model_name=litellm_model_name,
router_model_id=router_model_id,
litellm_logging_obj=litellm_logging_obj,
pricing_provider=pricing_provider,
service_tier=service_tier,
data_residency=data_residency,
vertex_location=vertex_location,
@ -2335,7 +2445,16 @@ def pricing_entry_for_cost_calc(
custom_llm_provider=custom_llm_provider,
router_model_id=router_model_id,
region_name=region_name,
invoked_model_id=_invoked_model_id_for_cost(None, litellm_logging_obj),
)
bedrock_model_info: Final = _bedrock_pricing_info(
model or _get_response_model(completion_response),
_invoked_model_id_for_cost(None, litellm_logging_obj),
_get_provider_for_cost_calc(model, custom_llm_provider),
custom_pricing is True or base_model is not None,
)
if bedrock_model_info is not None:
return bedrock_model_info["key"], bedrock_model_info
candidates: Final = (selected_model, _get_response_model(completion_response), model)
resolved: Final = next(
(info for info in (_cost_map_model_info(name, custom_llm_provider) for name in candidates if name) if info),
@ -2783,6 +2902,8 @@ def batch_cost_calculator(
custom_llm_provider: str | None = None,
model_info: ModelInfo | None = None,
data_residency: str | None = None,
*,
vertex_location: str | None = None,
) -> tuple[float, float]:
"""
Calculate the cost of a batch job.
@ -2872,7 +2993,8 @@ def batch_cost_calculator(
total_prompt_cost *= uplift
total_completion_cost *= uplift
return total_prompt_cost, total_completion_cost
vertex_uplift: Final = get_vertex_regional_endpoint_uplift(model_info, vertex_location)
return total_prompt_cost * vertex_uplift, total_completion_cost * vertex_uplift
def _batch_prompt_cost(
@ -3113,14 +3235,19 @@ class ResponsesWebSocketTokenUsageProcessor(BaseTokenUsageProcessor):
return response_ids[-1] if response_ids else None
@staticmethod
def create_logging_object(results: Sequence[Mapping[str, object]]) -> LiteLLMRealtimeStreamLoggingObject:
def create_logging_object(
results: Sequence[Mapping[str, object]], litellm_model_name: str | None = None
) -> LiteLLMRealtimeStreamLoggingObject:
tier_partition: Final = ResponsesWebSocketTokenUsageProcessor.partition_results_by_service_tier(results)
return LiteLLMRealtimeStreamLoggingObject(
logging_object: Final = LiteLLMRealtimeStreamLoggingObject(
id=ResponsesWebSocketTokenUsageProcessor.response_id_from_responses_ws_results(results),
usage=ResponsesWebSocketTokenUsageProcessor.collect_and_combine_usage_from_responses_ws_results(results),
results=cast(OpenAIRealtimeStreamList, list(results)), # cast-ok: raw WS frames, unvalidated field
service_tier=next(iter(tier_partition)) if len(tier_partition) == 1 else None,
)
if litellm_model_name is not None:
logging_object.hidden_params = {"litellm_model_name": litellm_model_name}
return logging_object
_TRANSCRIPTION_COMPLETED_EVENT_TYPE: Final = "conversation.item.input_audio_transcription.completed"

68
litellm/decisions/call.py
View file

@ -1,3 +1,4 @@
import json
from collections.abc import Mapping, Sequence
from dataclasses import dataclass, field, replace
from functools import partial
@ -6,9 +7,11 @@ from typing import Final, Literal, TypeAlias
import httpx
from pydantic import BaseModel, ConfigDict, StrictBool, TypeAdapter, ValidationError
from pydantic_core import to_jsonable_python
from typing_extensions import assert_never
import litellm
from litellm.cost_calculator import cost_map_key_prices_anything
from litellm.litellm_core_utils.core_helpers import RESPONSE_COST_HEADER, normalize_drop_params
from litellm.litellm_core_utils.get_provider_specific_headers import ProviderSpecificHeaderUtils
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
@ -19,6 +22,7 @@ from litellm.litellm_core_utils.prompt_templates.image_handling import (
from litellm.llms.base_llm.chat.transformation import BaseLLMException
from litellm.llms.base_llm.decisions.transformation import (
BaseDecisionsConfig,
DecisionsResponseValidationError,
image_media_type,
image_param,
input_images,
@ -56,6 +60,7 @@ from litellm.types.decisions import (
OpenAIDecisionResponse,
UnsupportedDecisionsRequest,
)
from litellm.types.llms.openai import ChatCompletionUserMessage
from litellm.types.utils import LlmProviders, ProviderSpecificHeader
from litellm.utils import ProviderConfigManager, get_model_cost_key
@ -415,6 +420,30 @@ def _outbound_headers(
return MappingProxyType({**(extra_headers or {}), **forwarded})
_LOGGED_FIELDS_ADAPTER: Final[TypeAdapter[Mapping[str, object]]] = TypeAdapter(Mapping[str, object])
def _jsonable_fallback(value: object) -> object:
return dict(_LOGGED_FIELDS_ADAPTER.validate_python(value)) if isinstance(value, Mapping) else str(value)
def _logged_request_message(caller_request: object) -> ChatCompletionUserMessage:
caller_fields: Final = _LOGGED_FIELDS_ADAPTER.validate_python(
to_jsonable_python(caller_request, fallback=_jsonable_fallback, bytes_mode="base64")
)
sent_fields: Final = json.dumps(
dict((name, value) for name, value in caller_fields.items() if value not in (None, [], {}))
)
return ChatCompletionUserMessage(role="user", content=sent_fields)
def log_caller_request(kwargs: Mapping[str, object], caller_request: object) -> None:
"""Record every request field the caller sent as the logged `messages`, before any validation."""
logging_obj: Final = kwargs.get("litellm_logging_obj")
if isinstance(logging_obj, LiteLLMLoggingObj):
logging_obj.update_messages([_logged_request_message(caller_request)])
def prepare_call(
*,
model: str,
@ -660,6 +689,13 @@ def _request_body(call: DecisionsCall) -> Mapping[str, object]:
return body
def _pricing_model(response: DecisionsIRResponse, call: DecisionsCall) -> str:
returned: Final = get_model_cost_key(f"{call.custom_llm_provider}/{response.model}")
if returned is not None and cost_map_key_prices_anything(returned):
return returned
return f"{call.custom_llm_provider}/{call.model}"
def _format_response(response: DecisionsIRResponse, call: DecisionsCall) -> DecisionsResponse | OpenAIDecisionResponse:
formatted: Final = _formatted_response(response, call)
provider_cost: Final = call.provider_config.provider_reported_cost(response)
@ -667,10 +703,7 @@ def _format_response(response: DecisionsIRResponse, call: DecisionsCall) -> Deci
{
"model": f"{call.custom_llm_provider}/{call.model}",
"custom_llm_provider": call.custom_llm_provider,
"provider_response_model": (
get_model_cost_key(f"{call.custom_llm_provider}/{response.model}")
or f"{call.custom_llm_provider}/{call.model}"
),
"provider_response_model": _pricing_model(response, call),
**({"additional_headers": {RESPONSE_COST_HEADER: provider_cost}} if provider_cost is not None else {}),
}
)
@ -700,18 +733,21 @@ def _formatted_response(
def _map_upstream_exception(error: Exception, call: DecisionsCall) -> Exception:
if isinstance(error, BaseLLMException) and error.status_code_is_synthesized:
provider_label: Final = f"{call.custom_llm_provider[0].upper()}{call.custom_llm_provider[1:]}Exception"
return litellm.APIConnectionError(
message=f"{provider_label} - {error.message}",
llm_provider=call.custom_llm_provider,
model=f"{call.custom_llm_provider}/{call.model}",
)
return litellm.exception_type(
model=f"{call.custom_llm_provider}/{call.model}",
custom_llm_provider=call.custom_llm_provider,
original_exception=error,
)
provider_label: Final = f"{call.custom_llm_provider[0].upper()}{call.custom_llm_provider[1:]}Exception"
model: Final = f"{call.custom_llm_provider}/{call.model}"
match error:
case DecisionsResponseValidationError():
return litellm.InternalServerError(
message=f"{provider_label} - {error.message}", llm_provider=call.custom_llm_provider, model=model
)
case BaseLLMException(status_code_is_synthesized=True):
return litellm.APIConnectionError(
message=f"{provider_label} - {error.message}", llm_provider=call.custom_llm_provider, model=model
)
case _:
return litellm.exception_type(
model=model, custom_llm_provider=call.custom_llm_provider, original_exception=error
)
async def asend(call: DecisionsCall) -> DecisionsResponse | OpenAIDecisionResponse:

11
litellm/decisions/main.py
View file

@ -17,6 +17,7 @@ from litellm.decisions.call import (
SystemOneQuestions,
SystemOneRequestFields,
asend,
log_caller_request,
prepare_call,
reject_images_next_to_input,
send,
@ -78,6 +79,16 @@ def _prepare_decisions_call(
extra_headers: Mapping[str, str] | None,
kwargs: Mapping[str, object],
) -> DecisionsCall:
log_caller_request(
kwargs,
{
"input": decision_input,
"questions": questions,
"audio": audio,
"videos": videos,
**{key: kwargs[key] for key in _SYSTEMONE_FIELDS if key in kwargs},
},
)
return prepare_call(
model=model,
fields=_request_fields(model, decision_input, questions, audio, videos, kwargs),

18
litellm/embeddings/dispatch.py
View file

@ -9,7 +9,15 @@ from litellm.rust_bridge.embeddings.entrypoints import (
NATIVE_AEMBEDDING,
NATIVE_EMBEDDING,
)
from litellm.rust_bridge.public_call import NativeCall, bind, native_call, native_call_hook, optional_str, signature
from litellm.rust_bridge.public_call import (
AsyncFlag,
NativeCall,
bind,
native_call,
native_call_hook,
optional_str,
signature,
)
from litellm.types.utils import EmbeddingResponse
__all__ = ("aembedding", "embedding")
@ -24,6 +32,7 @@ _PYTHON_AEMBEDDING: Final = cast( # cast-ok: [LIT006] preserve the legacy publi
PythonAembedding, main.aembedding
)
_EMBEDDING_SIGNATURE: Final = signature(_PYTHON_EMBEDDING)
_FLAG: Final = AsyncFlag("aembedding", sync=_EMBEDDING_SIGNATURE)
def _public_request(
@ -50,7 +59,6 @@ _DISPATCH: Final = PublicDispatch(
route=Route.EMBEDDINGS,
request=lambda args, kwargs: _public_request(_EMBEDDING_SIGNATURE, args, kwargs),
context=_context,
bypass=lambda request: request.kwargs.get("aembedding") is True,
)
_ADISPATCH: Final = PublicDispatch(
@ -64,9 +72,11 @@ def embedding(
*args: object,
**kwargs: object, # kwargs-ok: preserve the public embedding call shape
) -> EmbeddingResponse | Coroutine[object, object, EmbeddingResponse]:
if _FLAG.requested(kwargs):
return aembedding(**_FLAG.keywords(args, kwargs))
return _DISPATCH.run(
args,
kwargs,
_FLAG.dropped(kwargs),
python=_PYTHON_EMBEDDING,
binding=NATIVE_EMBEDDING,
native=native_call_hook,
@ -76,7 +86,7 @@ def embedding(
async def aembedding(*args: object, **kwargs: object) -> EmbeddingResponse: # kwargs-ok: preserve the public call shape
return await _ADISPATCH.arun(
args,
kwargs,
_FLAG.dropped(kwargs),
python=_PYTHON_AEMBEDDING,
binding=NATIVE_AEMBEDDING,
native=native_call_hook,

9
litellm/endpoints/speech/speech_to_completion_bridge/handler.py
View file

@ -6,6 +6,8 @@ from typing import TYPE_CHECKING, Final
from typing_extensions import TypedDict
from litellm._internal_context import internal_sub_call
if TYPE_CHECKING:
from litellm import LiteLLMLoggingObj
from litellm.types.llms.openai import HttpxBinaryResponseContent
@ -111,9 +113,10 @@ class SpeechToCompletionBridgeHandler:
voice=voice,
)
result: Final = completion(
**request_data,
)
with internal_sub_call():
result: Final = completion(
**request_data,
)
requested_response_format: Final = optional_params.get("response_format")
if isinstance(result, ModelResponse):

15
litellm/endpoints/speech/speech_to_completion_bridge/transformation.py
View file

@ -5,11 +5,12 @@ from typing import TYPE_CHECKING, Final, cast
from typing_extensions import NotRequired, ReadOnly, TypedDict
from litellm.constants import OPENAI_CHAT_COMPLETION_PARAMS
from litellm.litellm_core_utils.hidden_params import get_hidden_params, served_from_cache, set_hidden_param
if TYPE_CHECKING:
from litellm import Logging as LiteLLMLoggingObj
from litellm.types.llms.openai import ChatCompletionUserMessage, HttpxBinaryResponseContent
from litellm.types.utils import ModelResponse
from litellm.types.utils import ModelResponse, Usage
def _completion_response_cost(model_response: "ModelResponse") -> float | None:
@ -20,6 +21,13 @@ def _completion_response_cost(model_response: "ModelResponse") -> float | None:
return response_cost if isinstance(response_cost, float) else None
def _completion_usage(model_response: "ModelResponse") -> "Usage | None":
from litellm.types.utils import Usage
usage: Final = getattr(model_response, "usage", None)
return usage if isinstance(usage, Usage) else None
GEMINI_TTS_CHAT_AUDIO_FORMAT: Final = "pcm16"
GEMINI_TTS_RAW_RESPONSE_FORMAT: Final = "pcm"
GEMINI_TTS_SUPPORTED_RESPONSE_FORMATS: Final = frozenset({"wav", GEMINI_TTS_RAW_RESPONSE_FORMAT})
@ -179,4 +187,9 @@ class SpeechToCompletionBridgeTransformationHandler:
)
binary_response: Final = HttpxBinaryResponseContent(response)
binary_response.set_response_cost(_completion_response_cost(model_response))
binary_response.usage = _completion_usage(model_response)
cached_hidden_params: Final = get_hidden_params(model_response) if served_from_cache(model_response) else None
if cached_hidden_params is not None:
set_hidden_param(binary_response, "cache_hit", True)
set_hidden_param(binary_response, "cache_key", cached_hidden_params.get("cache_key"))
return binary_response

13
litellm/experimental_mcp_client/client.py
View file

@ -95,6 +95,7 @@ from pydantic import AnyUrl, TypeAdapter
from litellm._logging import verbose_logger
from litellm.constants import (
MCP_CLIENT_TIMEOUT,
MCP_MAX_SSE_EVENT_SIZE_BYTES,
MCP_NPM_CACHE_DIR,
MCP_TOOL_LISTING_MAX_PAGES,
MCP_TOOL_LISTING_TIMEOUT,
@ -292,6 +293,7 @@ class _MCPResponseStream(httpx2.AsyncByteStream):
class _MCPHTTPClient(httpx2.AsyncClient):
cleanup_scope: anyio.CancelScope | None = None
cleanup_errors: tuple[Exception, ...] = ()
post_failure: asyncio.Future[Exception] | None = None
def _record_cleanup_error(self, error: Exception) -> None:
if self.cleanup_scope is not None and self.cleanup_scope.shield:
@ -319,7 +321,12 @@ class _MCPHTTPClient(httpx2.AsyncClient):
response: Final = await super().send(request, stream=stream, auth=auth, follow_redirects=follow_redirects)
if request.method == "POST" and response.is_error and response.status_code != 404:
await response.aclose()
response.raise_for_status()
try:
response.raise_for_status()
except httpx2.HTTPStatusError as error:
if self.post_failure is not None and not self.post_failure.done():
self.post_failure.set_result(error)
raise
if stream:
response.stream = _MCPResponseStream(
cast(httpx2.AsyncByteStream, response.stream), self._record_cleanup_error
@ -641,6 +648,7 @@ class MCPClient:
transport_ctx: Final = streamable_http_client(
url=self.server_url,
http_client=http_client,
max_sse_event_size=MCP_MAX_SSE_EVENT_SIZE_BYTES,
)
return transport_ctx, http_client
@ -757,6 +765,8 @@ class MCPClient:
read_stream: Final = transport[0]
write_stream: Final = transport[1]
stream_error: Final[asyncio.Future[Exception]] = asyncio.get_running_loop().create_future()
if isinstance(http_client, _MCPHTTPClient):
http_client.post_failure = stream_error # rebind-ok: operation-scoped stream error
_notify_stream_error(on_stream_error, stream_error)
session_kwargs: Final = (
@ -823,6 +833,7 @@ class MCPClient:
if isinstance(http_client, _MCPHTTPClient):
http_client.cleanup_errors = ()
http_client.cleanup_scope = None
http_client.post_failure = None # rebind-ok: clear operation-scoped stream error
await anyio.lowlevel.checkpoint_if_cancelled()
if cleanup_scope.cancel_called:
raise (

3
litellm/extraction/__init__.py Normal file
View file

@ -0,0 +1,3 @@
from litellm.extraction.main import aextraction, extraction
__all__ = ["aextraction", "extraction"]

351
litellm/extraction/main.py Normal file
View file

@ -0,0 +1,351 @@
import asyncio
from collections.abc import Mapping, Sequence
from dataclasses import dataclass
from functools import partial
from types import MappingProxyType
from typing import Final, TypeAlias
import httpx
from pydantic import TypeAdapter, ValidationError
import litellm
from litellm.decisions.call import log_caller_request
from litellm.litellm_core_utils.core_helpers import get_metadata_variable_name_from_kwargs
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.litellm_core_utils.tool_fields import DropPolicy, drop_policy
from litellm.llms.base_llm.chat.transformation import BaseLLMException
from litellm.llms.base_llm.extraction.transformation import BaseExtractionConfig
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler
from litellm.types.extraction import (
EXTRACTION_STATE_REWRITTEN_METADATA_KEY,
ExtractionAnswer,
ExtractionBase64Image,
ExtractionField,
ExtractionOptions,
ExtractionRequestBody,
ExtractionResponse,
)
from litellm.types.router import GenericLiteLLMParams
from litellm.types.utils import LlmProviders, is_litellm_owned_kwarg
from litellm.utils import (
ProviderConfigManager,
client, # pyright: ignore[reportUnknownVariableType] # client is an untyped decorator
get_model_cost_key,
)
ExtractionFields: TypeAlias = Mapping[str, ExtractionField | Mapping[str, object]]
ExtractionImages: TypeAlias = Sequence[str | ExtractionBase64Image | Mapping[str, object]]
ExtractionOptionsParam: TypeAlias = ExtractionOptions | Mapping[str, object]
_REQUEST_ADAPTER: Final[TypeAdapter[ExtractionRequestBody]] = TypeAdapter(ExtractionRequestBody)
_METADATA_ADAPTER: Final[TypeAdapter[Mapping[str, object]]] = TypeAdapter(Mapping[str, object])
_HANDLER: Final = BaseLLMHTTPHandler()
_LOCAL_KWARGS: Final = frozenset(
(
*GenericLiteLLMParams.model_fields,
"extra_headers",
"base_url",
"default_headers",
"deployment_id",
"additional_drop_params",
)
)
@dataclass(frozen=True, slots=True, repr=False)
class _ExtractionCall:
model: str
requested_model: str
custom_llm_provider: str
provider_config: BaseExtractionConfig
body: Mapping[str, object]
api_base: str
api_key: str | None
logging_obj: LiteLLMLoggingObj | None
headers: Mapping[str, str]
timeout: float | httpx.Timeout | None
litellm_params: Mapping[str, object]
client: HTTPHandler | AsyncHTTPHandler | None
def supported_extraction_providers() -> tuple[str, ...]:
return tuple(
provider.value
for provider in LlmProviders
if ProviderConfigManager.get_provider_extraction_config(model="", provider=provider) is not None
)
def _provider_config(model: str, custom_llm_provider: str) -> BaseExtractionConfig:
provider: Final = next((member for member in LlmProviders if member.value == custom_llm_provider), None)
provider_config: Final = (
None if provider is None else ProviderConfigManager.get_provider_extraction_config(model, provider)
)
if provider_config is None:
raise litellm.BadRequestError(
message=(
f"Unknown Extraction provider '{custom_llm_provider}'. "
f"Supported providers: {', '.join(supported_extraction_providers())}"
),
model=model,
llm_provider=custom_llm_provider,
)
return provider_config
def _request_fields(
*,
state: str | None,
fields: ExtractionFields | None,
images: ExtractionImages | None,
reference_date: str | None,
options: ExtractionOptionsParam | None,
) -> Mapping[str, object]:
sent: Final[Mapping[str, object | None]] = MappingProxyType(
{"state": state, "fields": fields, "images": images, "reference_date": reference_date, "options": options}
)
return MappingProxyType({name: value for name, value in sent.items() if value is not None})
def _is_unsupported(name: str, value: object, policy: DropPolicy) -> bool:
return (
value is not None
and value is not False
and name not in _LOCAL_KWARGS
and not is_litellm_owned_kwarg(name)
and not policy.drops(name)
)
def _unsupported_params(kwargs: Mapping[str, object]) -> tuple[str, ...]:
policy: Final = drop_policy(kwargs)
return tuple(sorted(name for name, value in kwargs.items() if _is_unsupported(name, value, policy)))
def _prepare_call(
*,
model: str,
request_fields: Mapping[str, object],
api_key: str | None,
api_base: str | None,
timeout: float | httpx.Timeout | None,
custom_llm_provider: str | None,
extra_headers: Mapping[str, str] | None,
kwargs: Mapping[str, object],
) -> _ExtractionCall:
log_caller_request(kwargs, request_fields)
upstream_model, provider, dynamic_api_key, dynamic_api_base = litellm.get_llm_provider(
model=model,
custom_llm_provider=custom_llm_provider,
api_base=api_base,
api_key=api_key,
)
provider_config: Final = _provider_config(model, provider)
unsupported: Final = _unsupported_params(kwargs)
if unsupported:
raise litellm.UnsupportedParamsError(
message=(
f"{provider} does not support parameters: {list(unsupported)}, for model={model}. "
"To drop these, set `litellm.drop_params=True` or for proxy:\n\n"
"`litellm_settings:\n drop_params: true`\n"
),
model=model,
llm_provider=provider,
)
try:
request: Final = _REQUEST_ADAPTER.validate_python(request_fields)
except ValidationError as error:
raise litellm.BadRequestError(
message=f"Invalid Extraction request: {error}", model=model, llm_provider=provider
) from error
canonical_model: Final = provider_config.canonical_model(upstream_model)
resolved_api_base: Final = provider_config.resolve_api_base(dynamic_api_base or api_base)
if resolved_api_base is None:
raise litellm.BadRequestError(
message=provider_config.missing_api_base_message(provider), model=model, llm_provider=provider
)
resolved_api_key: Final = provider_config.resolve_api_key(dynamic_api_key or api_key)
if resolved_api_key is None:
raise litellm.BadRequestError(
message=f"Missing API key for Extraction provider '{provider}'", model=model, llm_provider=provider
)
logging_obj: Final = kwargs.get("litellm_logging_obj")
http_client: Final = kwargs.get("client")
if isinstance(logging_obj, LiteLLMLoggingObj):
logging_obj.update_from_kwargs( # pyright: ignore[reportUnknownMemberType] # Logging takes untyped dicts
kwargs=dict(kwargs),
model=canonical_model,
litellm_params={
"litellm_call_id": kwargs.get("litellm_call_id"),
"api_base": provider_config.get_complete_url(resolved_api_base, canonical_model),
},
custom_llm_provider=provider,
)
return _ExtractionCall(
model=canonical_model,
requested_model=model,
custom_llm_provider=provider,
provider_config=provider_config,
body=provider_config.transform_extraction_request(canonical_model, request),
api_base=resolved_api_base,
api_key=resolved_api_key,
logging_obj=logging_obj if isinstance(logging_obj, LiteLLMLoggingObj) else None,
headers=MappingProxyType(dict(extra_headers or {})),
timeout=timeout,
litellm_params=kwargs,
client=http_client if isinstance(http_client, HTTPHandler | AsyncHTTPHandler) else None,
)
def _state_rewritten_by_guardrail(litellm_params: Mapping[str, object]) -> bool:
metadata: Final = litellm_params.get(get_metadata_variable_name_from_kwargs(litellm_params))
return (
isinstance(metadata, Mapping)
and _METADATA_ADAPTER.validate_python(metadata).get(EXTRACTION_STATE_REWRITTEN_METADATA_KEY) is True
)
def _without_spans(answers: Mapping[str, ExtractionAnswer]) -> Mapping[str, ExtractionAnswer]:
return {name: answer.model_copy(update={"span": None, "spans": None}) for name, answer in answers.items()}
def _formatted_response(response: ExtractionResponse, call: _ExtractionCall) -> ExtractionResponse:
spans_index_hidden_text: Final = response.state_text is None and _state_rewritten_by_guardrail(call.litellm_params)
formatted: Final = response.model_copy(
update={
"model": call.requested_model,
**({"answers": _without_spans(response.answers)} if spans_index_hidden_text else {}),
}
)
formatted.set_hidden_params(
{
"model": f"{call.custom_llm_provider}/{call.model}",
"custom_llm_provider": call.custom_llm_provider,
"provider_response_model": (
get_model_cost_key(f"{call.custom_llm_provider}/{response.model}")
or f"{call.custom_llm_provider}/{call.model}"
),
}
)
return formatted
def _upstream_exception(error: Exception, call: _ExtractionCall) -> Exception:
if isinstance(error, BaseLLMException) and error.status_code_is_synthesized:
return litellm.APIConnectionError(
message=f"{call.custom_llm_provider.capitalize()}Exception - {error.message}",
llm_provider=call.custom_llm_provider,
model=f"{call.custom_llm_provider}/{call.model}",
)
return litellm.exception_type( # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # untyped
model=f"{call.custom_llm_provider}/{call.model}",
custom_llm_provider=call.custom_llm_provider,
original_exception=error,
)
async def _asend(call: _ExtractionCall) -> ExtractionResponse:
try:
response: Final = await _HANDLER.aextraction(
model=call.model,
custom_llm_provider=call.custom_llm_provider,
logging_obj=call.logging_obj,
provider_config=call.provider_config,
body=call.body,
api_base=call.api_base,
api_key=call.api_key,
headers=call.headers,
timeout=call.timeout,
litellm_params=call.litellm_params,
client=call.client if isinstance(call.client, AsyncHTTPHandler) else None,
)
except Exception as error:
raise _upstream_exception(error, call) from error
return _formatted_response(response, call)
def _send(call: _ExtractionCall) -> ExtractionResponse:
try:
response: Final = _HANDLER.extraction(
model=call.model,
custom_llm_provider=call.custom_llm_provider,
logging_obj=call.logging_obj,
provider_config=call.provider_config,
body=call.body,
api_base=call.api_base,
api_key=call.api_key,
headers=call.headers,
timeout=call.timeout,
litellm_params=call.litellm_params,
client=call.client if isinstance(call.client, HTTPHandler) else None,
)
except Exception as error:
raise _upstream_exception(error, call) from error
return _formatted_response(response, call)
@client
async def aextraction(
model: str,
fields: ExtractionFields,
state: str | None = None,
images: ExtractionImages | None = None,
reference_date: str | None = None,
options: ExtractionOptionsParam | None = None,
api_key: str | None = None,
api_base: str | None = None,
timeout: float | httpx.Timeout | None = None,
custom_llm_provider: str | None = None,
extra_headers: Mapping[str, str] | None = None,
**kwargs: object, # kwargs-ok: litellm_params forwarded by @client
) -> ExtractionResponse:
call: Final = await asyncio.to_thread(
partial(
_prepare_call,
model=model,
request_fields=_request_fields(
state=state, fields=fields, images=images, reference_date=reference_date, options=options
),
api_key=api_key,
api_base=api_base,
timeout=timeout,
custom_llm_provider=custom_llm_provider,
extra_headers=extra_headers,
kwargs=kwargs,
)
)
return await _asend(call)
@client
def extraction(
model: str,
fields: ExtractionFields,
state: str | None = None,
images: ExtractionImages | None = None,
reference_date: str | None = None,
options: ExtractionOptionsParam | None = None,
api_key: str | None = None,
api_base: str | None = None,
timeout: float | httpx.Timeout | None = None,
custom_llm_provider: str | None = None,
extra_headers: Mapping[str, str] | None = None,
**kwargs: object, # kwargs-ok: litellm_params forwarded by @client
) -> ExtractionResponse:
call: Final = _prepare_call(
model=model,
request_fields=_request_fields(
state=state, fields=fields, images=images, reference_date=reference_date, options=options
),
api_key=api_key,
api_base=api_base,
timeout=timeout,
custom_llm_provider=custom_llm_provider,
extra_headers=extra_headers,
kwargs=kwargs,
)
return _send(call)
__all__ = ["aextraction", "extraction"]

7
litellm/fine_tuning/main.py
View file

@ -173,6 +173,7 @@ def create_fine_tuning_job(
try:
_is_async: Final = kwargs.pop("acreate_fine_tuning_job", False) is True
optional_params: Final = GenericLiteLLMParams(**kwargs)
provider_model: Final = model.removeprefix(f"{custom_llm_provider}/")
# handle hyperparameters
hyperparameters = hyperparameters or {} # original hyperparameters
@ -218,7 +219,7 @@ def create_fine_tuning_job(
)
create_fine_tuning_job_data_dict = _build_fine_tuning_job_data(
model,
provider_model,
training_file,
_oai_hyperparameters,
suffix,
@ -263,7 +264,7 @@ def create_fine_tuning_job(
extra_body = _prepare_azure_extra_body(extra_body, kwargs, azure_specific_hyperparams)
create_fine_tuning_job_data_dict = _build_fine_tuning_job_data(
model,
provider_model,
training_file,
_oai_hyperparameters,
suffix,
@ -298,7 +299,7 @@ def create_fine_tuning_job(
response = vertex_fine_tuning_apis_instance.create_fine_tuning_job(
_is_async=_is_async,
create_fine_tuning_job_data=_build_fine_tuning_job_data(
model,
provider_model,
training_file,
_oai_hyperparameters,
suffix,

3
litellm/google_genai/main.py
View file

@ -16,6 +16,7 @@ from litellm.llms.base_llm.google_genai.transformation import (
BaseGoogleGenAIGenerateContentConfig,
)
from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler
from litellm.llms.vertex_ai.vertex_llm_base import VertexBase
from litellm.types.llms.base import LiteLLMBaseModel
from litellm.types.router import GenericLiteLLMParams
from litellm.types.utils import CallTypes
@ -194,12 +195,14 @@ class GenerateContentHelper:
if litellm_logging_obj is None:
raise ValueError("litellm_logging_obj is required, but got None")
vertex_location: Final = VertexBase.explicit_vertex_ai_location(dict(litellm_params))
litellm_logging_obj.update_from_kwargs(
kwargs=kwargs,
model=model,
optional_params=dict(generate_content_config_dict),
litellm_params={
"litellm_call_id": litellm_call_id,
**({"vertex_location": vertex_location} if vertex_location else {}),
},
custom_llm_provider=custom_llm_provider,
)

4
litellm/google_genai/streaming_iterator.py
View file

@ -71,7 +71,7 @@ class BaseGoogleGenAIGenerateContentStreamingIterator:
):
self.litellm_logging_obj = litellm_logging_obj
self.request_body = request_body
self.start_time = datetime.now()
self.start_time = litellm_logging_obj.start_time
self.collected_chunks: list[bytes] = []
self.model = model
self.custom_llm_provider = custom_llm_provider
@ -192,6 +192,8 @@ class AsyncGoogleGenAIGenerateContentStreamingIterator(BaseGoogleGenAIGenerateCo
async def __anext__(self):
try:
chunk: Final = await _anext_google_genai_sse_chunk(self.stream_iterator)
if self.litellm_logging_obj.completion_start_time is None:
self.litellm_logging_obj.update_completion_start_time(completion_start_time=datetime.now())
self.collected_chunks.append(chunk)
return chunk
except StopAsyncIteration:

4
litellm/harness/__init__.py
View file

@ -1,4 +1,4 @@
"""Agent harnesses: run Claude Code, Codex, OpenCode, Deep Agents or Tool Loop on any LiteLLM model.
"""Agent harnesses: run Claude Code, Codex, OpenCode, pi, Deep Agents or Tool Loop on any LiteLLM model.
The entrypoints live on the top-level package:
@ -30,6 +30,7 @@ from litellm.harness.options import (
CodexOptions,
DeepAgentsOptions,
OpenCodeOptions,
PiOptions,
ToolLoopOptions,
)
from litellm.harness.runtime import (
@ -78,6 +79,7 @@ __all__ = (
"OpenCodeOptions",
"OptionsMismatch",
"OutputInvalid",
"PiOptions",
"Reasoning",
"Result",
"SandboxError",

119
litellm/harness/endpoint.py
View file

@ -22,7 +22,7 @@ from typing import TYPE_CHECKING, Any, Final, Protocol
import httpx
import openai
from pydantic import ConfigDict, TypeAdapter
from pydantic import ConfigDict, TypeAdapter, ValidationError
import litellm
from litellm.constants import (
@ -172,16 +172,57 @@ def usage_from_mapping(usage: object) -> tuple[int, int]:
return _as_int(input_tokens), _as_int(output_tokens)
def usage_from_body(body: object) -> tuple[int, int]:
"""Usage from a non-streaming JSON response body."""
if not isinstance(body, Mapping):
return 0, 0
if isinstance(body.get("usage"), Mapping):
return usage_from_mapping(body["usage"])
response = body.get("response")
if isinstance(response, Mapping):
return usage_from_mapping(response.get("usage"))
return 0, 0
@dataclass(frozen=True, slots=True)
class PromptCache:
"""Prompt tokens read from or written to the provider's cache. Anthropic's usage counts them outside input_tokens,
while OpenAI's prompt_tokens and the Responses API's input_tokens already include them."""
read_tokens: int = 0
creation_tokens: int = 0
outside_input: bool = False
def prompt_tokens(self, input_tokens: int) -> int:
if self.outside_input:
return input_tokens + self.read_tokens + self.creation_tokens
return input_tokens
NO_PROMPT_CACHE: Final = PromptCache()
def json_object(value: object) -> Mapping[str, object]:
try:
return _JSON_OBJECT_ADAPTER.validate_python(value)
except ValidationError:
return MappingProxyType({})
def prompt_cache_from_mapping(usage: object) -> PromptCache:
"""Cache token counts from a usage dict using OpenAI, Anthropic, or Responses field names."""
fields: Final = json_object(usage)
details: Final = json_object(fields.get("prompt_tokens_details", fields.get("input_tokens_details")))
return PromptCache(
read_tokens=_as_int(fields.get("cache_read_input_tokens") or details.get("cached_tokens")),
creation_tokens=_as_int(
fields.get("cache_creation_input_tokens")
or details.get("cache_write_tokens")
or details.get("cache_creation_tokens")
),
outside_input="prompt_tokens" not in fields and "input_tokens_details" not in fields,
)
def usage_mapping_of_body(body: object) -> Mapping[str, object]:
fields: Final = json_object(body)
usage: Final = json_object(fields.get("usage"))
return usage if usage else json_object(json_object(fields.get("response")).get("usage"))
def usage_mapping_of_json(raw: bytes) -> Mapping[str, object]:
try:
return usage_mapping_of_body(_JSON_VALUE_ADAPTER.validate_python(json.loads(raw)))
except ValueError:
return MappingProxyType({})
class SSEUsageParser:
@ -190,6 +231,7 @@ class SSEUsageParser:
def __init__(self) -> None:
self.input_tokens = 0
self.output_tokens = 0
self.prompt_cache = NO_PROMPT_CACHE
self._buffer = b""
def feed(self, chunk: bytes) -> None:
@ -226,21 +268,23 @@ class SSEUsageParser:
elif event_type == "response.completed":
self._absorb_response_completed(event)
elif isinstance(event.get("usage"), Mapping):
self._set(*usage_from_mapping(event["usage"]))
self._absorb_usage(event["usage"])
def _absorb_message_start(self, event: Mapping[str, object]) -> None:
message = event.get("message")
if isinstance(message, Mapping):
self._set(*usage_from_mapping(message.get("usage")))
self._absorb_usage(json_object(event.get("message")).get("usage"))
def _absorb_message_delta(self, event: Mapping[str, object]) -> None:
# message_delta output_tokens is cumulative for the whole message.
self._set(*usage_from_mapping(event.get("usage")))
self._absorb_usage(event.get("usage"))
def _absorb_response_completed(self, event: Mapping[str, object]) -> None:
response = event.get("response")
if isinstance(response, Mapping):
self._set(*usage_from_mapping(response.get("usage")))
self._absorb_usage(json_object(event.get("response")).get("usage"))
def _absorb_usage(self, usage: object) -> None:
self._set(*usage_from_mapping(usage))
prompt_cache: Final = prompt_cache_from_mapping(usage)
if prompt_cache.read_tokens or prompt_cache.creation_tokens:
self.prompt_cache = prompt_cache
def _set(self, input_tokens: int, output_tokens: int) -> None:
if input_tokens:
@ -249,13 +293,19 @@ class SSEUsageParser:
self.output_tokens = output_tokens
def compute_cost(model: str | None, input_tokens: int, output_tokens: int) -> float:
"""Cost from LiteLLM's price map. Never raises; unknown models cost 0.0."""
def compute_cost(
model: str | None, input_tokens: int, output_tokens: int, prompt_cache: PromptCache = NO_PROMPT_CACHE
) -> float:
"""Cost from LiteLLM's price map, cached prompt tokens at their cache prices. Never raises; unknown models cost 0.0."""
if not model or not (input_tokens or output_tokens):
return 0.0
try:
prompt_cost, completion_cost = litellm.cost_per_token(
model=model, prompt_tokens=input_tokens, completion_tokens=output_tokens
model=model,
prompt_tokens=prompt_cache.prompt_tokens(input_tokens),
completion_tokens=output_tokens,
cache_read_input_tokens=prompt_cache.read_tokens,
cache_creation_input_tokens=prompt_cache.creation_tokens,
)
return float(prompt_cost) + float(completion_cost)
except Exception: # accounting must never break a call; the price-map lookup raises bare Exception
@ -565,9 +615,10 @@ class ModelEndpoint:
input_tokens: int,
output_tokens: int,
cost: float | None,
prompt_cache: PromptCache,
) -> None:
if cost is None:
cost = compute_cost(model, input_tokens, output_tokens)
cost = compute_cost(model, input_tokens, output_tokens, prompt_cache)
self.usage.add(input_tokens, output_tokens, cost)
async def _forward(self, request: Request, route: str, body: Mapping[str, object]) -> Response:
@ -617,14 +668,12 @@ class ModelEndpoint:
) -> None:
if is_sse:
parser.close()
tokens = (parser.input_tokens, parser.output_tokens)
else:
try:
body: Final[object] = _JSON_VALUE_ADAPTER.validate_python(json.loads(collected))
tokens = usage_from_body(body)
except ValueError:
tokens = (0, 0)
self._record(model, tokens[0], tokens[1], header_cost(upstream.headers))
self._record(
model, parser.input_tokens, parser.output_tokens, header_cost(upstream.headers), parser.prompt_cache
)
return
usage: Final = usage_mapping_of_json(collected)
self._record(model, *usage_from_mapping(usage), header_cost(upstream.headers), prompt_cache_from_mapping(usage))
def _sdk_kwargs(
self, body: Mapping[str, object]
@ -667,9 +716,9 @@ class ModelEndpoint:
return self._responses.StreamingResponse(
self._sdk_stream(route, response, model), media_type=SSE_MEDIA_TYPE
)
data = to_jsonable(response)
input_tokens, output_tokens = usage_from_body(data)
self._record(model, input_tokens, output_tokens, hidden_cost(response))
data: Final = to_jsonable(response)
usage: Final = usage_mapping_of_body(data)
self._record(model, *usage_from_mapping(usage), hidden_cost(response), prompt_cache_from_mapping(usage))
return self._json(data)
async def _sdk_stream(self, route: str, iterator: AsyncIterable[object], model: str | None) -> AsyncIterator[bytes]:
@ -688,4 +737,4 @@ class ModelEndpoint:
yield f"event: error\ndata: {json.dumps(error_body(e, message))}\n\n".encode()
finally:
parser.close()
self._record(model, parser.input_tokens, parser.output_tokens, None)
self._record(model, parser.input_tokens, parser.output_tokens, None, parser.prompt_cache)

86
litellm/harness/handlers/cli_handler.py
View file

@ -12,17 +12,27 @@ from __future__ import annotations
import asyncio
import os
from collections import deque
from collections.abc import AsyncIterator, Sequence
from collections.abc import AsyncIterator, Mapping, Sequence
from dataclasses import dataclass
from typing import Final
from litellm._logging import verbose_logger
from litellm.constants import HARNESS_STDERR_TAIL_LINES, HARNESS_STREAM_READ_CHUNK_BYTES
from litellm.constants import (
HARNESS_RESUME_COMMAND_TIMEOUT_SECONDS,
HARNESS_STDERR_TAIL_LINES,
HARNESS_STREAM_READ_CHUNK_BYTES,
)
from litellm.harness.context import SessionContext
from litellm.harness.errors import HarnessInstallFailed, SandboxError
from litellm.harness.errors import HarnessError, HarnessInstallFailed, SandboxError
from litellm.harness.handlers.base import BaseHarnessHandler
from litellm.harness.sandbox.base import Process, Sandbox
from litellm.harness.sandbox.base import CompletedRun, Process, Sandbox
from litellm.harness.types import Event
from litellm.llms.base_llm.harness.transformation import BaseCLIHarnessConfig, HarnessSessionSetup
from litellm.llms.base_llm.harness.transformation import (
BaseCLIHarnessConfig,
HarnessSessionCopy,
HarnessSessionSetup,
HarnessTurnError,
)
from litellm.llms.base_llm.harness.utils import decode_json_line, read_skill_files
# Link <private_dir>/<dir> to a LiteLLM-owned cache dir so a later session can resume.
@ -81,6 +91,52 @@ async def persist_dir(sandbox: Sandbox, link_path: str, cache_subpath: str) -> N
)
@dataclass(frozen=True, slots=True)
class ResumeFailed:
reason: str
async def bounded_resume_run(
sandbox: Sandbox, cmd: Sequence[str], env: Mapping[str, str]
) -> CompletedRun | ResumeFailed:
try:
run: Final = await sandbox.run(cmd, env=env, timeout=HARNESS_RESUME_COMMAND_TIMEOUT_SECONDS)
except SandboxError as exc:
return ResumeFailed(str(exc))
return run if run.exit_code == 0 else ResumeFailed(run.stderr.strip() or f"exit code {run.exit_code}")
async def import_session_copy(sandbox: Sandbox, copy: HarnessSessionCopy, env: Mapping[str, str]) -> str | ResumeFailed:
try:
await sandbox.write(copy.path, copy.data)
except SandboxError as exc:
return ResumeFailed(str(exc))
imported: Final = await bounded_resume_run(sandbox, copy.import_command, env)
return imported if isinstance(imported, ResumeFailed) else copy.session_id
async def resumable_session_id(
config: BaseCLIHarnessConfig[object, object],
sandbox: Sandbox,
setup_env: Mapping[str, str],
private_dir: str,
native_session_id: str,
) -> str | ResumeFailed:
export_cmd: Final = config.transform_resume_export_command(native_session_id)
if export_cmd is None:
return native_session_id
exported: Final = await bounded_resume_run(sandbox, export_cmd, setup_env)
if isinstance(exported, ResumeFailed):
return exported
try:
copy: Final = config.transform_resume_copy(exported.stdout, private_dir, sandbox.workdir)
except HarnessError as exc:
return ResumeFailed(str(exc))
if copy is None:
return native_session_id
return await import_session_copy(sandbox, copy, setup_env)
async def copy_skills(sandbox: Sandbox, skills: Sequence[str], skills_root: str) -> None:
for skill in skills:
name = os.path.basename(os.path.realpath(os.fspath(skill)))
@ -89,13 +145,14 @@ async def copy_skills(sandbox: Sandbox, skills: Sequence[str], skills_root: str)
class CLIHarnessHandler(BaseHarnessHandler):
config: BaseCLIHarnessConfig
config: BaseCLIHarnessConfig[object, object]
def __init__(self, config: BaseCLIHarnessConfig) -> None:
super().__init__(config)
self._private_dir: str | None = None
self._setup: HarnessSessionSetup | None = None
self._native_id: str | None = None
self._pending_resume: str | None = None
self._proc: Process | None = None
async def start(self, ctx: SessionContext) -> None:
@ -119,6 +176,9 @@ class CLIHarnessHandler(BaseHarnessHandler):
async def turn(self, ctx: SessionContext, prompt: str) -> AsyncIterator[Event]:
if self._setup is None or self._private_dir is None:
raise RuntimeError("CLIHarnessHandler.turn() called before start()")
if self._pending_resume is not None:
self._native_id = await self._prepared_session_id(ctx, self._setup, self._private_dir, self._pending_resume)
self._pending_resume = None
request = self.config.transform_turn_request(ctx, self._setup, self._private_dir, prompt, self._native_id)
argv: Final = list(request.argv) # mutable-ok: Sandbox.exec takes list[str]
proc = await ctx.sandbox.exec(argv, env=request.env, cwd=request.cwd)
@ -158,4 +218,16 @@ class CLIHarnessHandler(BaseHarnessHandler):
return self._native_id
async def resume(self, ctx: SessionContext, native_session_id: str) -> None:
self._native_id = native_session_id
if native_session_id != self._native_id:
self._native_id = native_session_id
self._pending_resume = native_session_id
async def _prepared_session_id(
self, ctx: SessionContext, setup: HarnessSessionSetup, private_dir: str, saved_id: str
) -> str:
target: Final = await resumable_session_id(self.config, ctx.sandbox, setup.env, private_dir, saved_id)
match target:
case ResumeFailed(reason=reason):
raise HarnessTurnError(f"could not prepare session {saved_id} for resume: {reason}")
case session_id:
return session_id

2
litellm/harness/handlers/deepagents_handler.py
View file

@ -183,7 +183,7 @@ class DeepAgentsHandler(BaseHarnessHandler):
break
resume: dict[str, object] = {} # mutable-ok: Command(resume=) payload, filled per answered approval
for interrupt in state.interrupts:
decisions: list[dict[str, Any]] = [] # mutable-ok: HITL decisions collected across awaited approvals
decisions: list[dict[str, str]] = [] # mutable-ok: HITL decisions collected across awaited approvals
for request in approval_requests(getattr(interrupt, "value", None)):
approval = Approval(
tool=normalized_tool_name(str(request.get("name") or "")),

29
litellm/harness/options.py
View file

@ -4,28 +4,37 @@ from __future__ import annotations
from collections.abc import Mapping, Sequence
from dataclasses import dataclass, field
from typing import Literal
from typing import Literal, TypeAlias
PiThinkingLevel: TypeAlias = Literal["off", "minimal", "low", "medium", "high", "xhigh", "max"]
@dataclass(frozen=True)
class ClaudeCodeOptions:
config: Mapping[str, object] = field(default_factory=dict)
env: Mapping[str, str] = field(default_factory=dict)
config: Mapping[str, object] = field(default_factory=dict[str, object])
env: Mapping[str, str] = field(default_factory=dict[str, str])
@dataclass(frozen=True)
class CodexOptions:
reasoning_effort: Literal["low", "medium", "high", "xhigh"] | None = None
web_search: bool = False
config: Mapping[str, object] = field(default_factory=dict)
env: Mapping[str, str] = field(default_factory=dict)
config: Mapping[str, object] = field(default_factory=dict[str, object])
env: Mapping[str, str] = field(default_factory=dict[str, str])
@dataclass(frozen=True)
class OpenCodeOptions:
agent: str = "build"
config: Mapping[str, object] = field(default_factory=dict)
env: Mapping[str, str] = field(default_factory=dict)
config: Mapping[str, object] = field(default_factory=dict[str, object])
env: Mapping[str, str] = field(default_factory=dict[str, str])
@dataclass(frozen=True)
class PiOptions:
thinking: PiThinkingLevel | None = None
config: Mapping[str, object] = field(default_factory=dict[str, object])
env: Mapping[str, str] = field(default_factory=dict[str, str])
@dataclass(frozen=True)
@ -36,7 +45,9 @@ class DeepAgentsOptions:
@dataclass(frozen=True)
class ToolLoopOptions:
completion_kwargs: Mapping[str, object] = field(default_factory=dict)
completion_kwargs: Mapping[str, object] = field(default_factory=dict[str, object])
HarnessOptions = ClaudeCodeOptions | CodexOptions | OpenCodeOptions | DeepAgentsOptions | ToolLoopOptions
HarnessOptions: TypeAlias = (
ClaudeCodeOptions | CodexOptions | OpenCodeOptions | PiOptions | DeepAgentsOptions | ToolLoopOptions
)

16
litellm/harness/sandbox/local.py
View file

@ -14,7 +14,7 @@ from typing import Final
from litellm.constants import HARNESS_PROCESS_KILL_GRACE_SECONDS
from litellm.harness.errors import SandboxError
from litellm.harness.sandbox.base import CompletedRun
from litellm.harness.sandbox.base import CompletedRun, Process
from litellm.harness.sandbox.snapshot import snapshot_local
_SECRET_PREFIXES: Final = (
@ -103,7 +103,7 @@ class SubprocessHandle:
await self._proc.wait()
async def _read_all(handle: SubprocessHandle) -> tuple[bytes, bytes, int]:
async def _read_all(handle: Process) -> tuple[bytes, bytes, int]:
if handle.stdin is not None:
handle.stdin.close()
stdout, stderr = await asyncio.gather(handle.stdout.read(), handle.stderr.read())
@ -111,13 +111,16 @@ async def _read_all(handle: SubprocessHandle) -> tuple[bytes, bytes, int]:
return stdout, stderr, exit_code
async def collect_output(handle: SubprocessHandle, cmd: Sequence[str], timeout: float | None) -> CompletedRun:
"""Close stdin, read stdout/stderr to EOF; kill and raise SandboxError on timeout."""
async def collect_output(handle: Process, cmd: Sequence[str], timeout: float | None) -> CompletedRun:
"""Close stdin, read stdout/stderr to EOF; kill on cancel, and kill and raise SandboxError on timeout."""
try:
stdout, stderr, code = await asyncio.wait_for(_read_all(handle), timeout)
except asyncio.TimeoutError:
await handle.kill()
raise SandboxError(f"command timed out after {timeout}s: {cmd[0]}")
except asyncio.CancelledError:
await handle.kill()
raise
return CompletedRun(
stdout=stdout.decode("utf-8", errors="replace"),
stderr=stderr.decode("utf-8", errors="replace"),
@ -180,14 +183,15 @@ class LocalSandbox:
self._check_open()
if not cmd:
raise SandboxError("exec() needs a non-empty command")
work_dir: Final = self._resolve_cwd(cwd)
try:
proc = await asyncio.create_subprocess_exec(
*cmd,
stdin=asyncio.subprocess.PIPE,
stdout=asyncio.subprocess.PIPE,
stderr=asyncio.subprocess.PIPE,
cwd=self._resolve_cwd(cwd),
env=self.child_env(env),
cwd=work_dir,
env=self.child_env(MappingProxyType({**(env or {}), "PWD": work_dir})),
start_new_session=True,
)
except (FileNotFoundError, PermissionError) as exc:

2
litellm/harness/sync.py
View file

@ -411,7 +411,7 @@ def agent(
options: HarnessOptions | None = None,
install: bool = False,
) -> Result | EventStream:
"""Run an agent harness (Claude Code, Codex, OpenCode, Deep Agents, Tool Loop) on one prompt.
"""Run an agent harness (Claude Code, Codex, OpenCode, pi, Deep Agents, Tool Loop) on one prompt.
Returns a Result. With stream=True it returns an iterator of events instead.
Prefix the model with `litellm_proxy/` to route every model call through your

1
litellm/harness/types.py
View file

@ -26,6 +26,7 @@ class Harness(Enum):
CLAUDE_CODE = "claude_code"
CODEX = "codex"
OPENCODE = "opencode"
PI = "pi"
DEEPAGENTS = "deepagents"
TOOL_LOOP = "tool_loop"

21
litellm/integrations/anthropic_cache_control_hook.py
View file

@ -342,6 +342,20 @@ EXTERNAL_BREAKPOINTS_STAMP: Final = "_litellm_external_breakpoints"
CUSTOM_LLM_PROVIDER_STAMP: Final = "_litellm_custom_llm_provider"
def _accepts_cache_control(part: object) -> bool:
match part:
case {"type": "thinking" | "redacted_thinking"}:
return False
case dict():
return True
case _:
return False
def _last_index_accepting_cache_control(parts: Sequence[object]) -> int | None:
return next((index for index in reversed(range(len(parts))) if _accepts_cache_control(parts[index])), None)
class AnthropicCacheControlHook(CustomPromptManagement):
@staticmethod
def _request_value(request_kwargs: object, key: str) -> object:
@ -740,10 +754,11 @@ class AnthropicCacheControlHook(CustomPromptManagement):
# 1. if string, insert cache control in the message
if isinstance(message_content, str):
message["cache_control"] = control
# 2. list of objects - only apply to last item per Anthropic spec
# 2. list of objects - only apply to the last item Claude accepts a breakpoint on
elif isinstance(message_content, list):
if len(message_content) > 0 and isinstance(message_content[-1], dict):
message_content[-1]["cache_control"] = control # pyright: ignore[reportGeneralTypeIssues] # loose runtime dict
markable_index: Final = _last_index_accepting_cache_control(message_content)
if markable_index is not None:
message_content[markable_index]["cache_control"] = control # pyright: ignore[reportGeneralTypeIssues] # loose runtime dict
return message
@staticmethod

13
litellm/integrations/arize/_utils.py
View file

@ -884,7 +884,7 @@ def _set_session_and_user_attrs(span: "Span", kwargs: dict, standard_logging_pay
safe_set_attribute(span, "litellm.key_alias", str(key_alias))
def _set_response_cost_attr(span: "Span", standard_logging_payload) -> None:
def _set_response_cost_attr(span: "Span", standard_logging_payload: Mapping[str, object] | None) -> None:
"""Emit cost attributes from the StandardLoggingPayload when present.
Uses the OpenInference `llm.cost.total` key so Arize / Phoenix can
@ -892,15 +892,22 @@ def _set_response_cost_attr(span: "Span", standard_logging_payload) -> None:
single total in `StandardLoggingPayload.response_cost`, so we cannot
split it into prompt/completion. We also keep the legacy
`llm.response.cost` key for back-compat with any consumer querying it.
When the price lookup failed, the call's total is unknown, so no cost is
emitted and the backend does not record the call as free. A guardrail fee
in `response_cost` stays on the guardrail span, as in OTel v2.
"""
if not isinstance(standard_logging_payload, dict):
return
cost: Final = standard_logging_payload.get("response_cost")
if cost is None:
cost_lookup_failed: Final = (
standard_logging_payload.get("response_cost_failure_debug_info") is not None
and standard_logging_payload.get("cache_hit") is not True
)
if not isinstance(cost, (int, float, str)) or cost_lookup_failed:
return
try:
cost_value: Final = float(cost)
except (TypeError, ValueError):
except ValueError:
return
safe_set_attribute(span, "llm.cost.total", cost_value)
safe_set_attribute(span, "llm.response.cost", cost_value)

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