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

This commit is contained in:
mateo-berri 2026-10-29 17:21:28 -07:00
commit 17822fd186
504 changed files with 31504 additions and 5645 deletions

View file

@ -1,14 +1,14 @@
#!/usr/bin/env bash
set -euo pipefail
if [ "${GITHUB_ACTIONS:-}" = true ]; then
echo "Integration contracts are owned by CircleCI" >&2
exit 1
fi
suite="${1:?integration suite required}"
mode="${2:-standard}"
side="${3:-}"
requested_workers="${INTEGRATION_WORKERS:-}"
license_env=()
if [ -n "${LITELLM_LICENSE:-}" ]; then
license_env=("LITELLM_LICENSE=$LITELLM_LICENSE")
fi
if [ "$mode" = replica ]; then
results="test-results/integration-${suite}-replica"
elif [ "$mode" = parity ]; then
@ -84,7 +84,7 @@ if [ "$suite" = browser ]; then
export LITELLM_UI_PATH="$PWD/ui/litellm-dashboard/out"
test -f "$LITELLM_UI_PATH/index.html"
fi
export INTEGRATION_SEED="$(.venv/bin/python -c 'import hashlib,os; print(int(hashlib.sha256((os.environ.get("CIRCLE_SHA1", "local") + os.environ.get("CIRCLE_WORKFLOW_ID", "local")).encode()).hexdigest()[:8],16))')"
export INTEGRATION_SEED="$(.venv/bin/python -c 'import hashlib,os; print(int(hashlib.sha256(((os.environ.get("CIRCLE_SHA1") or os.environ.get("GITHUB_SHA") or "local") + (os.environ.get("CIRCLE_WORKFLOW_ID") or os.environ.get("GITHUB_RUN_ID") or "local")).encode()).hexdigest()[:8],16))')"
export INTEGRATION_ORDER_SEED="$INTEGRATION_SEED"
uv run --no-sync prisma generate --schema litellm/proxy/schema.prisma > "$results/prisma-generate.log" 2>&1
@ -107,7 +107,7 @@ sudo iptables -N integration_only
guard_created=true
sudo iptables -A integration_only -o lo -j ACCEPT
sudo iptables -A integration_only -m conntrack --ctstate ESTABLISHED,RELATED -j ACCEPT
for service in postgres-db redis-cache; do
for service in "${INTEGRATION_POSTGRES_CONTAINER:-postgres-db}" "${INTEGRATION_REDIS_CONTAINER:-redis-cache}"; do
address="$(docker inspect --format '{{range .NetworkSettings.Networks}}{{.IPAddress}}{{end}}' "$service")"
sudo iptables -A integration_only -d "$address" -j ACCEPT
done
@ -145,6 +145,9 @@ fi
if [ "$suite" = mcp ]; then
export INTEGRATION_WORKERS=2 INTEGRATION_COVERAGE=1
fi
if [ -n "$requested_workers" ]; then
export INTEGRATION_WORKERS="$requested_workers"
fi
isolated_workers=false
if [ "${INTEGRATION_WORKERS:-1}" -gt 1 ]; then
isolated_workers=true
@ -185,7 +188,7 @@ start_proxy() {
"${database_env[@]}" REDIS_HOST="$REDIS_HOST" REDIS_PORT="$REDIS_PORT" \
INTEGRATION_UPSTREAM_URL="$INTEGRATION_UPSTREAM_URL" \
LITELLM_MASTER_KEY="$LITELLM_MASTER_KEY" LITELLM_SALT_KEY="$LITELLM_SALT_KEY" LITELLM_UI_PATH="$LITELLM_UI_PATH" PROXY_BASE_URL="http://127.0.0.1:$port" \
LITELLM_LICENSE="${LITELLM_LICENSE:-}" \
"${license_env[@]}" \
LITELLM_MODE=PRODUCTION STORE_MODEL_IN_DB=True LITELLM_ENABLE_MCP_STDIO=true "${cost_map_env[@]}" \
AWS_EC2_METADATA_DISABLED=true DO_NOT_TRACK=1 COVERAGE_FILE="$coverage_data" \
"${proxy_command[@]}" --config tests/integration/proxy_config.yaml \
@ -257,7 +260,7 @@ env -i PATH="$PATH" HOME="$HOME" PYTHONPATH="$PYTHONPATH" \
"${shared_services[@]}" "${worker_coverage[@]}" \
INTEGRATION_WORKERS="${INTEGRATION_WORKERS:-1}" \
LITELLM_MODE=PRODUCTION \
LITELLM_LICENSE="${LITELLM_LICENSE:-}" \
"${license_env[@]}" \
INTEGRATION_SEED="$INTEGRATION_SEED" \
INTEGRATION_ORDER_SEED="$INTEGRATION_ORDER_SEED" \
INTEGRATION_CASE_SHARDS="${CIRCLE_NODE_TOTAL:-1}" INTEGRATION_CASE_SHARD_INDEX="${CIRCLE_NODE_INDEX:-0}" \

119
.github/scripts/assert_secret_absent.py vendored Normal file
View file

@ -0,0 +1,119 @@
import base64
import io
import os
import sys
import urllib.parse
import zipfile
from collections.abc import Callable, Iterator, Mapping
from dataclasses import dataclass
from functools import reduce
from itertools import chain
from pathlib import Path
from types import MappingProxyType
from typing import Final, TextIO
SECRET_VARIABLE: Final = "SCANNED_SECRET"
MIN_ENCODED_LENGTH: Final = 12
MAX_ARCHIVE_DEPTH: Final = 3
MASK: Final = "***"
Forms = tuple[tuple[str, str], ...]
@dataclass(frozen=True, slots=True)
class Hit:
location: str
form: str
def _aligned_encodings(raw: bytes, name: str, encode: Callable[[bytes], bytes]) -> Iterator[tuple[str, str]]:
for offset in range(3):
chunk: Final = raw[offset:]
yield f"{name} offset {offset}", encode(chunk[: len(chunk) - len(chunk) % 3]).decode()
def secret_forms(secret: str) -> Forms:
encoded: Final = (
("url-encoded", urllib.parse.quote(secret, safe="")),
*_aligned_encodings(secret.encode(), "base64", base64.b64encode),
*_aligned_encodings(secret.encode(), "base64url", base64.urlsafe_b64encode),
)
by_value: Final = MappingProxyType(
{value: name for name, value in encoded if value != secret and len(value) >= MIN_ENCODED_LENGTH}
)
return (("raw", secret), *((name, value) for value, name in by_value.items()))
def mask(text: str, forms: Forms) -> str:
return reduce(lambda masked, form: masked.replace(form[1], MASK), forms, text)
def _name_hits(location: str, name: str, forms: Forms) -> Iterator[Hit]:
for form, value in forms:
if value in name:
yield Hit(mask(location, forms), f"{form} in name")
def _content_hits(location: str, data: bytes, forms: Forms, depth: int) -> Iterator[Hit]:
for form, value in forms:
if value.encode() in data:
yield Hit(mask(location, forms), form)
if depth >= MAX_ARCHIVE_DEPTH or not zipfile.is_zipfile(io.BytesIO(data)):
return
with zipfile.ZipFile(io.BytesIO(data)) as archive:
for member in archive.infolist():
if member.is_dir():
continue
yield from _content_hits(f"{location}!{member.filename}", archive.read(member), forms, depth + 1)
def _walk(root: Path) -> Iterator[tuple[Path, bool]]:
for directory, subdirectories, files in os.walk(root):
yield from ((Path(directory, name), True) for name in subdirectories if Path(directory, name).is_symlink())
yield from ((Path(directory, name), False) for name in files)
def _path_hits(root: Path, path: Path, linked_directory: bool, forms: Forms) -> Iterator[Hit]:
location: Final = str(path.relative_to(root))
if linked_directory:
yield Hit(mask(location, forms), "directory symlink, not scanned")
return
yield from _name_hits(location, location, forms)
yield from _content_hits(location, path.read_bytes(), forms, 0)
def scan(root: Path, secret: str) -> tuple[int, tuple[Hit, ...]]:
forms: Final = secret_forms(secret)
entries: Final = tuple(_walk(root))
hits: Final = tuple(chain.from_iterable(_path_hits(root, path, linked, forms) for path, linked in entries))
return len(entries), hits
def main(argv: tuple[str, ...], environ: Mapping[str, str], out: TextIO) -> int:
if len(argv) != 1:
print(f"usage: assert_secret_absent.py <directory>, with the secret in ${SECRET_VARIABLE}", file=out)
return 2
secret: Final = environ.get(SECRET_VARIABLE, "")
root: Final = Path(argv[0])
if not secret:
print(f"{SECRET_VARIABLE} is empty, nothing to scan for", file=out)
return 0
if not root.is_dir():
print(f"{root} is not a directory, nothing to scan", file=out)
return 0
try:
scanned, hits = scan(root, secret)
except Exception as error:
print(f"scan failed with {type(error).__name__}, treating the results as unsafe", file=out)
return 1
for hit in hits:
print(f"secret found: {hit.location} ({hit.form})", file=out)
if hits:
print(f"{len(hits)} copies of the secret under {root}, refusing to upload", file=out)
return 1
print(f"no copy of the secret in {scanned} files under {root}", file=out)
return 0
if __name__ == "__main__":
sys.exit(main(tuple(sys.argv[1:]), os.environ, sys.stdout))

View file

@ -7,6 +7,8 @@ import zipfile
from pathlib import Path
from typing import Final
COST_MAP: Final = Path(__file__).resolve().parents[2] / "model_prices_and_context_window.json"
CHILD_SCRIPT: Final = """
from importlib.util import module_from_spec, spec_from_file_location
from pathlib import Path
@ -34,6 +36,9 @@ else:
after = module.gil_stats()
if not isinstance(after.get("releases"), int):
raise AssertionError(f"native module unusable after panic: {after!r}")
if module.bundled_model_cost_map() != Path(sys.argv[2]).read_bytes():
raise AssertionError("the wheel embeds a different model_prices_and_context_window.json than this checkout")
"""
@ -55,7 +60,9 @@ def main() -> int:
native_path: Final = Path(temporary_directory) / Path(native_members[0].filename).name
native_path.write_bytes(archive.read(native_members[0]))
result: Final = subprocess.run((sys.executable, "-c", CHILD_SCRIPT, str(native_path)), check=False)
result: Final = subprocess.run(
(sys.executable, "-c", CHILD_SCRIPT, str(native_path), str(COST_MAP)), check=False
)
if result.returncode != 0:
sys.stderr.write(f"native wheel smoke test exited with status {result.returncode}\n")

View file

@ -46,3 +46,24 @@ jobs:
HEAD_REF: ${{ github.event.pull_request.head.ref }}
run: |
uv run --no-project python ci_cd/cost_map_guard.py --base "$MERGE_BASE" --head "$HEAD_SHA" --head-ref "$HEAD_REF"
- name: Check whether the pull request changes the cost map
id: changed
env:
MERGE_BASE: ${{ steps.revisions.outputs.merge_base }}
HEAD_SHA: ${{ github.event.pull_request.head.sha }}
run: |
if git diff --quiet "$MERGE_BASE" "$HEAD_SHA" -- model_prices_and_context_window.json; then
echo "cost_map=false" >> "$GITHUB_OUTPUT"
else
echo "cost_map=true" >> "$GITHUB_OUTPUT"
fi
- if: steps.changed.outputs.cost_map == 'true'
run: rustup toolchain install --no-self-update
- name: Load the cost map through model-catalog
if: steps.changed.outputs.cost_map == 'true'
env:
HEAD_SHA: ${{ github.event.pull_request.head.sha }}
run: |
git show "${HEAD_SHA}:model_prices_and_context_window.json" > "$RUNNER_TEMP/model_prices_and_context_window.json"
cargo run --locked --quiet --manifest-path litellm-rust/Cargo.toml -p litellm-model-catalog --bin check-registry \
-- "$RUNNER_TEMP/model_prices_and_context_window.json"

View file

@ -13,8 +13,10 @@ on:
- migrations/Dockerfile
- migrations/run.py
- gateway/Dockerfile
- gateway/launch.py
- gateway/main.py
- gateway/routes/allowlist.py
- litellm_cli/**
- backend/Dockerfile
- backend/main.py
- litellm-rust/**

199
.github/workflows/test-integration.yml vendored Normal file
View file

@ -0,0 +1,199 @@
name: integration
on:
pull_request:
types: [labeled]
schedule:
- cron: "5 8-20/3 * * 1-6"
timezone: America/Los_Angeles
workflow_dispatch:
inputs:
suite:
description: Integration group to run
type: choice
options:
- all
- providers
- extensions
- accounting
- management
- management-replica
- database
- database-replica
- mcp
- security
- cost
- sdk
- browser
default: all
workers:
description: Workers per group, blank keeps each group's default
type: string
default: ""
permissions:
contents: read
concurrency:
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}-${{ github.event.label.name }}-${{ inputs.suite }}-${{ inputs.workers }}
cancel-in-progress: ${{ github.event_name == 'pull_request' }}
jobs:
integration:
name: integration ${{ matrix.suite }}
if: github.event_name != 'pull_request' || github.event.label.name == 'run-ci'
environment: integration-license
runs-on: ${{ fromJSON('{"providers":"ubuntu-latest-64-cores","extensions":"ubuntu-latest-64-cores","accounting":"ubuntu-latest-32-cores","management":"ubuntu-latest-32-cores","management-replica":"ubuntu-latest-32-cores"}')[matrix.suite] || 'ubuntu-latest-16-cores' }}
timeout-minutes: 60
strategy:
fail-fast: false
matrix:
suite: ${{ fromJSON((inputs.suite || 'all') == 'all' && '["providers","extensions","accounting","management","management-replica","database","database-replica","mcp","security","cost","sdk","browser"]' || format('["{0}"]', inputs.suite)) }}
env:
UV_PYTHON: "3.12"
INTEGRATION_WORKERS: ${{ inputs.workers || fromJSON('{"providers":"32","extensions":"48","accounting":"16","management":"16","management-replica":"16","cost":"16","mcp":"8","security":"8","database":"4","database-replica":"4"}')[matrix.suite] || '' }}
services:
postgres:
image: postgres:16@sha256:e17e86066e5ef83e0952a9347f5c792b7ece00972e2aa787a6986f471b3dd3d5
command: -c max_connections=1000 -c shared_preload_libraries=pg_stat_statements -c pg_stat_statements.track=all -c pg_stat_statements.max=20000
env:
POSTGRES_USER: postgres
POSTGRES_PASSWORD: postgres
POSTGRES_DB: circle_test
ports:
- 5432:5432
options: >-
--shm-size 1g
--health-cmd "pg_isready -U postgres"
--health-interval 5s
--health-timeout 5s
--health-retries 20
redis:
image: redis:7-alpine@sha256:7aec734b2bb298a1d769fd8729f13b8514a41bf90fcdd1f38ec52267fbaa8ee6
ports:
- 6379:6379
options: >-
--health-cmd "redis-cli ping"
--health-interval 5s
--health-timeout 5s
--health-retries 20
steps:
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
persist-credentials: false
- name: Set up Python
timeout-minutes: 3
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
with:
python-version: ${{ env.UV_PYTHON }}
- name: Set up uv
timeout-minutes: 3
uses: ./.github/actions/setup-uv-with-retries
with:
version: "0.10.9"
- name: Cache the Rust build
timeout-minutes: 5
uses: ./.github/actions/cache-cargo-build
- name: Cache Prisma binaries
timeout-minutes: 5
uses: ./.github/actions/cache-prisma-binaries
- name: Install dependencies
timeout-minutes: 20
run: .github/scripts/uv_sync_with_retries.sh --frozen --all-groups --all-extras
- name: Install redis-server 7.2.16
timeout-minutes: 3
run: |
mkdir -p "$RUNNER_TEMP/redis-bin"
cid="$(docker create redis:7.2.16@sha256:0637954999d01b7c9ce9167db2da50656e2590d3b884f1c600c5f63bb6e6773c)"
docker cp "${cid}:/usr/local/bin/redis-server" "$RUNNER_TEMP/redis-bin/redis-server"
docker rm "$cid"
"$RUNNER_TEMP/redis-bin/redis-server" --version | grep -q 'v=7.2.16'
echo "$RUNNER_TEMP/redis-bin" >> "$GITHUB_PATH"
- name: Install MCP SDK1 peer
if: matrix.suite == 'mcp'
timeout-minutes: 5
run: |
uv venv --python 3.12 .venv-mcp-peer
uv pip install --no-config --exclude-newer "3 days" --python .venv-mcp-peer 'mcp==1.28.1' 'langchain-mcp-adapters==0.2.1'
echo "MCP_TEST_PEER_PYTHON=$PWD/.venv-mcp-peer/bin/python" >> "$GITHUB_ENV"
- name: Set up Node.js
if: matrix.suite == 'browser'
timeout-minutes: 5
uses: actions/setup-node@a0853c24544627f65ddf259abe73b1d18a591444 # v5.0.0
with:
node-version-file: ui/litellm-dashboard/.nvmrc
cache: "npm"
cache-dependency-path: |
ui/litellm-dashboard/package-lock.json
tests/e2e/ui/package-lock.json
- name: Install browser dependencies and build the dashboard
if: matrix.suite == 'browser'
timeout-minutes: 20
run: |
(cd ui/litellm-dashboard && npm ci)
(cd tests/e2e/ui && npm ci)
sudo env PATH="$PATH" DEBIAN_FRONTEND=noninteractive \
node tests/e2e/ui/node_modules/@playwright/test/cli.js install-deps chromium
node tests/e2e/ui/node_modules/@playwright/test/cli.js install chromium
(cd ui/litellm-dashboard && NEXT_TELEMETRY_DISABLED=1 npm run build)
- name: Run integration contracts
timeout-minutes: 45
env:
SUITE: ${{ matrix.suite }}
INTEGRATION_POSTGRES_CONTAINER: ${{ job.services.postgres.id }}
INTEGRATION_REDIS_CONTAINER: ${{ job.services.redis.id }}
LITELLM_LICENSE: ${{ secrets.LITELLM_LICENSE }}
run: |
case "$SUITE" in
*-replica) bash .circleci/scripts/run_integration.sh "${SUITE%-replica}" replica ;;
*) bash .circleci/scripts/run_integration.sh "$SUITE" ;;
esac
- name: Verify every case ran exactly once
if: ${{ !cancelled() && matrix.suite != 'browser' }}
timeout-minutes: 2
run: |
record="$(find test-results -mindepth 2 -maxdepth 2 -name execution.json -print -quit)"
test -n "$record"
mkdir -p "$RUNNER_TEMP/shards/0"
cp "$record" "$RUNNER_TEMP/shards/0/execution.json"
.venv/bin/python -m tests.integration._support.sharding "$RUNNER_TEMP/shards" --shards 1
- name: Report MCP coverage
if: ${{ !cancelled() && matrix.suite == 'mcp' }}
timeout-minutes: 5
run: |
coverage_dir=test-results/integration-mcp/coverage
test -f "$coverage_dir/coverage.txt" && exit 0
export COVERAGE_FILE="$PWD/$coverage_dir/data"
.venv/bin/python -m coverage combine --rcfile=tests/integration/mcp_coverage.toml
.venv/bin/python -m coverage report --rcfile=tests/integration/mcp_coverage.toml > "$coverage_dir/coverage.txt"
.venv/bin/python -m coverage html --rcfile=tests/integration/mcp_coverage.toml -d "$coverage_dir/html"
tail -n 1 "$coverage_dir/coverage.txt"
- name: Check results for the license
id: secret-scan
if: ${{ !cancelled() }}
timeout-minutes: 5
env:
SCANNED_SECRET: ${{ secrets.LITELLM_LICENSE }}
run: python .github/scripts/assert_secret_absent.py test-results
- name: Upload results
if: ${{ !cancelled() && steps.secret-scan.outcome == 'success' }}
timeout-minutes: 10
uses: actions/upload-artifact@4cec3d8aa04e39d1a68397de0c4cd6fb9dce8ec1 # v4.6.1
with:
name: integration-${{ matrix.suite }}-${{ github.run_id }}-${{ github.run_attempt }}
path: test-results/
retention-days: 14

View file

@ -322,6 +322,7 @@ jobs:
tests/unit/llms/openai_like
tests/unit/llms/opencode
tests/unit/llms/opensandbox
tests/unit/llms/opensearch
tests/unit/llms/openrouter
tests/unit/llms/ovhcloud
tests/unit/llms/parallel_ai

View file

@ -141,12 +141,56 @@ response = client.chat.completions.create(
)
```
[**Docs: LLM Providers**](https://docs.litellm.ai/docs/providers)
</details>
<details>
<summary><b>Auto Routing</b> - Choose a model for each prompt (Python SDK + AI Gateway)</summary>
### Python SDK
[**Docs: Autorouter SDK**](https://docs.litellm.ai/docs/autorouter_sdk)
Install with `uv add litellm` and set `OPENAI_API_KEY` for the models below.
```python
import asyncio
from litellm import Router
router = Router(model_list=[
{"model_name": "fast", "litellm_params": {"model": "openai/gpt-6-luna"}},
{"model_name": "capable", "litellm_params": {"model": "openai/gpt-6.1-sol"}},
{"model_name": "auto", "litellm_params": {
"model": "auto_router/complexity_router",
"complexity_router_config": {
"classifier_type": "heuristic",
"tiers": {"SIMPLE": "fast", "MEDIUM": "fast",
"COMPLEX": "capable", "REASONING": "capable"},
"return_raw_model_name": True,
},
"complexity_router_default_model": "capable",
}},
])
response = asyncio.run(router.acompletion(
model="auto", messages=[{"role": "user", "content": "What is 2+2?"}],
))
print(response.model, response.choices[0].message.content)
```
### AI Gateway (Proxy Server)
[**Create an Auto Router**](https://docs.litellm.ai/docs/auto_router/setup), then call its model name. Replace `auto` and the URL below with yours, and set `LITELLM_API_KEY` to your virtual key.
```shell
curl http://localhost:4000/v1/chat/completions \
-H "Authorization: Bearer $LITELLM_API_KEY" \
-H "Content-Type: application/json" \
-d '{"model":"auto","messages":[{"role":"user","content":"What is 2+2?"}]}'
```
</details>
<details>
<summary><b>Agents</b> - Invoke A2A Agents (Python SDK + AI Gateway)</summary>
<summary><b>A2A</b> - Invoke A2A Agents (Python SDK + AI Gateway)</summary>
[**Supported Providers**](https://docs.litellm.ai/docs/a2a#add-a2a-agents) - LangGraph, Vertex AI Agent Engine, Azure AI Foundry, Bedrock AgentCore, Pydantic AI
@ -442,7 +486,7 @@ You can use LiteLLM through either the Proxy Server or Python SDK. Both give you
<tr>
<th style={{width: '14%'}}></th>
<th style={{width: '43%'}}><strong><a href="https://docs.litellm.ai/docs/simple_proxy">LiteLLM AI Gateway</a></strong></th>
<th style={{width: '43%'}}><strong><a href="https://docs.litellm.ai/docs/">LiteLLM Python SDK</a></strong></th>
<th style={{width: '43%'}}><strong><a href="https://docs.litellm.ai/docs/python_sdk">LiteLLM Python SDK</a></strong></th>
</tr>
</thead>
<tbody>

View file

@ -24,6 +24,7 @@ ENV UV_PYTHON_DOWNLOADS=0
COPY pyproject.toml README.md rust-toolchain.toml ./
COPY .cargo ./.cargo
COPY litellm-rust ./litellm-rust
COPY model_prices_and_context_window.json ./
RUN --mount=type=cache,target=/root/.cache/uv \
mkdir -p litellm/rust_bridge /wheels && \

View file

@ -15,9 +15,10 @@ ignore:
flag_management:
default_rules:
carryforward: true
# Dead flags no CI job uploads anymore: their carried-forward sessions were
# measured against old revisions, and the stale line maps mark comment lines
# of since-edited files as missed, sinking patch coverage on unrelated PRs.
# Dead flags no CI job uploads anymore, and flags whose job is path-gated on
# PRs: their carried-forward sessions were measured against old revisions, and
# the stale line maps mark comment lines of since-edited files as missed,
# sinking patch coverage on unrelated PRs.
individual_flags:
- name: proxy-mgmt-behavior
carryforward: false
@ -99,6 +100,10 @@ flag_management:
carryforward: false
- name: proxy-db-endpoints-and-responses
carryforward: false
- name: roi-postgres
carryforward: false
- name: redis-compat
carryforward: false
component_management:
individual_components:

View file

@ -398,7 +398,7 @@ class CheckBatchCost:
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 a potentially recoverable row."""
return self.llm_router.get_deployment(model_id=model_id) is not None
return self.llm_router.registry().deployment(model_id) is not None
@staticmethod
def _is_output_file_gone_at_provider(error: Exception, output_file_id: Optional[str]) -> bool:
@ -564,7 +564,7 @@ class CheckBatchCost:
def _get_deployment_id_for_bare_model(
self, bare_model_name: str, llm_provider: str
) -> Optional[str]:
model_group = self.llm_router.resolve_model_name_from_model_id(bare_model_name)
model_group = self.llm_router.registry().resolve(bare_model_name, kind="from_id")
deployment_id = (
self._get_deployment_id_for_provider(model_group, llm_provider) if model_group else None
)
@ -580,7 +580,7 @@ class CheckBatchCost:
) -> Optional[str]:
from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider
for deployment in self.llm_router.get_model_list(model_name=None) or []:
for deployment in self.llm_router.registry().deployments() or []:
litellm_params = deployment.get("litellm_params") or {}
actual_model = litellm_params.get("model")
if not isinstance(actual_model, str):
@ -623,8 +623,9 @@ class CheckBatchCost:
"""
from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider
for deployment_id in self.llm_router.get_model_ids(model_name=model_group):
deployment_info = self.llm_router.get_deployment(model_id=deployment_id)
registry = self.llm_router.registry()
for deployment_id in registry.names(model_group, kind="ids"):
deployment_info = registry.deployment(deployment_id)
if deployment_info is None:
continue
try:
@ -743,7 +744,7 @@ class CheckBatchCost:
except (IndexError, AttributeError):
pass
credentials: Final = self.llm_router.get_deployment_credentials_with_provider(model_id) or {}
credentials: Final = self.llm_router.registry().credentials(model_id, form="with_provider") or {}
_file_content = await afile_content(
file_id=cast(str, raw_output_file_id),
_litellm_internal_model_credentials=MappingProxyType(dict(credentials)),
@ -766,7 +767,7 @@ class CheckBatchCost:
except Exception:
pass
deployment_info = self.llm_router.get_deployment(model_id=model_id)
deployment_info = self.llm_router.registry().deployment(model_id)
if deployment_info is None:
verbose_proxy_logger.info(
f"Skipping job {job.unified_object_id} because it is not a valid deployment info"
@ -996,7 +997,7 @@ class CheckBatchCost:
)
try:
deployment_info: Final = self.llm_router.get_deployment(model_id=model_id)
deployment_info: Final = self.llm_router.registry().deployment(model_id)
if deployment_info is None:
verbose_proxy_logger.error(
f"CheckBatchCost: deployment {model_id} is no longer in the router, so the line items of batch "
@ -1009,7 +1010,7 @@ class CheckBatchCost:
custom_llm_provider=deployment_info.litellm_params.custom_llm_provider,
)
cost_provider: Final = cast(BatchFileProvider, llm_provider)
credentials: Final = self.llm_router.get_deployment_credentials_with_provider(model_id) or {}
credentials: Final = self.llm_router.registry().credentials(model_id, form="with_provider") or {}
file_access: Final = {
**credentials,
"_litellm_internal_model_credentials": MappingProxyType(dict(credentials)),

View file

@ -273,7 +273,7 @@ class CheckResponsesCost:
def _resolve_deployment(self, response_id: str) -> Deployment | None:
model_id: str | None = ResponsesAPIRequestUtils.get_model_id_from_response_id(response_id)
return self.llm_router.get_deployment(model_id=model_id) if model_id is not None else None
return self.llm_router.registry().deployment(model_id) if model_id is not None else None
async def _get_response(
self,
@ -415,7 +415,7 @@ class CheckResponsesCost:
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 ()
self.llm_router.registry().deployments(requested, team_id=team_id) or ()
)
serving_id: Final = target.deployment.model_info.id
return MappingProxyType(

View file

@ -245,7 +245,7 @@ def _provider_file_retrieve_credentials(
if llm_router is None or model_id is None:
return None
try:
credentials: Final = llm_router.get_deployment_credentials_with_provider(model_id)
credentials: Final = llm_router.registry().credentials(model_id, form="with_provider")
except Exception as error:
verbose_logger.warning(
"Failed to retrieve credentials for provider file "
@ -2650,7 +2650,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
k: v for k, v in data.items() if k not in ("model", "file_id", "_litellm_internal_model_credentials")
}
for model_id, model_file_id in specific_model_file_id_mapping.items():
credentials = llm_router.get_deployment_credentials_with_provider(model_id=model_id)
credentials = llm_router.registry().credentials(model_id, form="with_provider")
delete_data = {
**filtered_data,
**(
@ -2694,7 +2694,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
# against the deployment's configured bucket, which they only
# trust from this immutable server-side snapshot, never from
# request params.
credentials = llm_router.get_deployment_credentials_with_provider(model_id=model_id)
credentials = llm_router.registry().credentials(model_id, form="with_provider")
if credentials is not None:
data["_litellm_internal_model_credentials"] = cast(Dict, MappingProxyType(dict(credentials)))
else:

View file

@ -290,7 +290,7 @@ def _check_team_project_limits(
def _router_access_group_names(llm_router: "Router | None") -> frozenset[str]:
return frozenset(llm_router.get_model_access_groups()) if llm_router is not None else frozenset()
return frozenset(llm_router.registry().access_groups()) if llm_router is not None else frozenset()
def _project_models_expanding_at_request_time(

View file

@ -40,6 +40,7 @@ ENV UV_PYTHON_DOWNLOADS=0
COPY pyproject.toml README.md rust-toolchain.toml ./
COPY .cargo ./.cargo
COPY litellm-rust ./litellm-rust
COPY model_prices_and_context_window.json ./
RUN --mount=type=cache,target=/root/.cache/uv \
mkdir -p litellm/rust_bridge /wheels && \
@ -144,6 +145,7 @@ COPY --from=builder --chown=nonroot:nonroot /app/.venv /app/.venv
COPY --from=builder --chown=nonroot:nonroot /app/litellm /app/litellm
COPY --from=builder --chown=nonroot:nonroot /app/enterprise /app/enterprise
COPY --from=builder --chown=nonroot:nonroot /app/gateway /app/gateway
COPY --from=builder --chown=nonroot:nonroot /app/litellm_cli /app/litellm_cli
COPY --from=builder --chown=nonroot:nonroot /app/docker/component_entrypoint.sh /app/docker/component_entrypoint.sh
COPY --from=builder --chown=nonroot:nonroot /app/schema.prisma /app/schema.prisma
COPY --from=builder /opt/prisma /opt/prisma

View file

@ -8,6 +8,11 @@ handed the loopback URL it listens on. A pre-existing ``DATABASE_URL`` wins in
wins under token auth too, so exporting it here is enough for every worker to
pick the pooled URL up unchanged.
LiteLLM imports wait for ``main``: every uvicorn worker imports this module again
before it can answer its first 5 s health ping, so a slow host would kill each one.
Workers then load the gateway app through ``worker_app`` without garbage collection
pauses, which hold the GIL long enough to miss a ping.
Run with:
python -m gateway.launch --workers 4 --host 0.0.0.0 --port 4000
"""
@ -15,19 +20,19 @@ Run with:
import os
import sys
from collections.abc import Callable, Mapping, Sequence
from typing import Final
from typing import TYPE_CHECKING, Final
from uvicorn.main import main as uvicorn_main
from litellm.proxy.db.db_url_settings import DatabaseURLSettings
from litellm.proxy.db.pgbouncer import (
PgBouncerError,
PgBouncerSettings,
export_pooled_database_url,
start_in_container_pgbouncer,
)
from litellm_cli.main import without_gc_pauses
GATEWAY_APP: Final = "gateway.main:app"
if TYPE_CHECKING:
from fastapi import FastAPI
from litellm.proxy.db.db_url_settings import DatabaseURLSettings
from litellm.proxy.db.pgbouncer import PgBouncerError, PgBouncerSettings
GATEWAY_APP: Final = "gateway.launch:worker_app"
KEEPALIVE_FLAG: Final = "--timeout-keep-alive"
@ -35,21 +40,23 @@ def uvicorn_argv(argv: Sequence[str], environ: Mapping[str, str]) -> tuple[str,
"""Honor ``KEEPALIVE_TIMEOUT`` like ``proxy_cli.py`` does, unless the flag was passed explicitly."""
keepalive: Final = environ.get("KEEPALIVE_TIMEOUT")
if keepalive is None or any(arg == KEEPALIVE_FLAG or arg.startswith(f"{KEEPALIVE_FLAG}=") for arg in argv):
return (GATEWAY_APP, *argv)
return (GATEWAY_APP, *argv, KEEPALIVE_FLAG, keepalive)
return (GATEWAY_APP, "--factory", *argv)
return (GATEWAY_APP, "--factory", *argv, KEEPALIVE_FLAG, keepalive)
def pool_database_url(
settings: DatabaseURLSettings,
pgbouncer: PgBouncerSettings,
settings: "DatabaseURLSettings",
pgbouncer: "PgBouncerSettings",
environ: Mapping[str, str],
) -> str | PgBouncerError | None:
) -> "str | PgBouncerError | None":
"""Start the in-container PgBouncer and return its loopback URL, or None when ``pgbouncer.enabled`` is off.
The upstream URL is whatever ``apply_to_env`` assembled from the discrete
``DATABASE_*`` vars (or an operator-pinned ``DATABASE_URL``). Under token
auth the pooler mints and renews the upstream token itself.
"""
from litellm.proxy.db.pgbouncer import PgBouncerError, start_in_container_pgbouncer
if not pgbouncer.enabled:
return None
upstream_url: Final = environ.get("DATABASE_URL")
@ -58,11 +65,20 @@ def pool_database_url(
return start_in_container_pgbouncer(pgbouncer, upstream_url, token_auth=settings.token_auth())
def worker_app() -> "FastAPI":
with without_gc_pauses():
from gateway.main import app
return app
def _serve(argv: Sequence[str]) -> None:
uvicorn_main(tuple(argv), prog_name="uvicorn")
def main(argv: Sequence[str], serve: Callable[[Sequence[str]], None] = _serve) -> None:
from litellm.proxy.db.db_url_settings import DatabaseURLSettings
from litellm.proxy.db.pgbouncer import PgBouncerError, PgBouncerSettings, export_pooled_database_url
settings: Final = DatabaseURLSettings.from_env()
settings.apply_to_env()
pooled_url: Final = pool_database_url(settings, PgBouncerSettings(), os.environ)

View file

@ -259,3 +259,9 @@ itself.
`lensWorker.mode` selects `disabled`, `bundled` or `external`. Bundled mode generates separate service and identity-signing credentials. External mode requires `lensWorker.gateway.secretName` and `lensWorker.serviceTokenSecret.name`; their keys default to `gateway-secret` and `service-token`. Provision distinct values matching the external Lens deployment
Follow the [Lens Helm connection guide](https://github.com/BerriAI/lens/blob/main/helm/lens/README.md#connect-a-gateway) for complete values, verification, GitOps credential requirements and existing-data migration. Source-chart installation requires a built Lens image until the first signed Lens release is published
## Shutdown accounting
`GRACEFUL_SHUTDOWN_TIMEOUT` defaults to 30 seconds. Pending success callbacks share the time remaining since SIGTERM or the first `/health/drain` call. Callback flushing reserves the existing 10-second teardown headroom for cancellation, tracked accounting tasks and persistence. Exhausting this budget can discard callbacks, which is logged
Set `terminationGracePeriodSeconds` to at least `GRACEFUL_SHUTDOWN_TIMEOUT + 10`, allowing additional time for preStop hooks and database writes. This is not a hard shutdown bound or a guarantee of persistence. Database retries retain their own limits, and cancellation-resistant callbacks can remain alive. The existing preStop hook can also be followed by Uvicorn's separate connection drain, as described in `values.yaml`

View file

@ -14,11 +14,9 @@ either, so the migration keeps today's behavior for it.
"""
import os
import re
from collections.abc import Callable, Iterator
from collections.abc import Callable
from dataclasses import dataclass
from functools import reduce
from itertools import accumulate
from pathlib import Path
from types import MappingProxyType
from typing import TYPE_CHECKING, Final, Literal
@ -43,7 +41,7 @@ MIGRATION_DATA_CHECK_TIMEOUT_ENV_VAR: Final = "LITELLM_MIGRATION_DATA_CHECK_TIME
DEFAULT_MIGRATION_DATA_CHECK_TIMEOUT_SECONDS: Final = 600.0
REFUSAL: Final = "Existing rows would make a pending schema change fail:"
Kind = Literal["too_long", "duplicates", "nulls", "any_row", "orphans", "check_violations"]
Kind = Literal["too_long"]
@dataclass(frozen=True, slots=True)
@ -54,9 +52,6 @@ class Check:
statement: str
source: str
index: str | None = None
referenced_table: str | None = None
referenced_columns: tuple[str, ...] = ()
expression: str | None = None
@dataclass(frozen=True, slots=True)
@ -100,33 +95,6 @@ class ColumnType:
ColumnTypeLookup = Callable[[str, str], ColumnType | None]
_IDENT: Final = r'"?(\w+)"?'
_QUALIFIED: Final = r'(?:"?\w+"?\s*\.\s*)?"?(\w+)"?'
_CREATE_INDEX: Final = re.compile(
rf"^\s*CREATE\s+(UNIQUE\s+)?INDEX\s+(?:CONCURRENTLY\s+)?(?:IF\s+NOT\s+EXISTS\s+)?{_IDENT}\s+ON\s+(?:ONLY\s+)?"
rf"{_QUALIFIED}\s*(?:USING\s+(\w+)\s*)?\((.*?)\)",
re.IGNORECASE | re.DOTALL,
)
_ALTER_TABLE: Final = re.compile(
rf"^\s*ALTER\s+TABLE\s+(?:IF\s+EXISTS\s+)?(?:ONLY\s+)?{_QUALIFIED}\s+(.*)$", re.IGNORECASE | re.DOTALL
)
_SET_NOT_NULL: Final = re.compile(rf"^ALTER\s+(?:COLUMN\s+)?{_IDENT}\s+SET\s+NOT\s+NULL$", re.IGNORECASE)
_ADD_COLUMN: Final = re.compile(
rf"^ADD\s+(?:COLUMN\s+)?(?:IF\s+NOT\s+EXISTS\s+)?(?!(?:CONSTRAINT|FOREIGN|UNIQUE|CHECK|PRIMARY|EXCLUDE)\b){_IDENT}\s+(.*)$",
re.IGNORECASE | re.DOTALL,
)
_CONSTRAINT: Final = rf"^ADD\s+(?:CONSTRAINT\s+{_IDENT}\s+)?"
_ADD_FOREIGN_KEY: Final = re.compile(
rf"{_CONSTRAINT}FOREIGN\s+KEY\s*\(([^)]*)\)\s*REFERENCES\s+{_QUALIFIED}\s*\(([^)]*)\)(.*)$",
re.IGNORECASE | re.DOTALL,
)
_ADD_UNIQUE: Final = re.compile(rf"{_CONSTRAINT}UNIQUE\s*\(([^)]*)\)", re.IGNORECASE | re.DOTALL)
_ADD_CHECK: Final = re.compile(rf"{_CONSTRAINT}CHECK\s*\((.*)\)\s*(NOT\s+VALID)?\s*$", re.IGNORECASE | re.DOTALL)
_PLAIN_COLUMN: Final = re.compile(rf"^{_IDENT}(?:\s+(?:ASC|DESC|NULLS\s+(?:FIRST|LAST)))*$", re.IGNORECASE)
_COMMENT: Final = re.compile(r"--[^\n]*")
_NOT_NULL: Final = re.compile(r"\bNOT\s+NULL\b", re.IGNORECASE)
_HAS_DEFAULT: Final = re.compile(r"\bDEFAULT\b|\bGENERATED\b", re.IGNORECASE)
def migration_data_check_timeout_seconds() -> float:
"""Seconds one count query may run before the preflight gives up on verifying it."""
@ -145,117 +113,6 @@ def migration_data_check_timeout_seconds() -> float:
return DEFAULT_MIGRATION_DATA_CHECK_TIMEOUT_SECONDS
def statements(migration_sql: str) -> tuple[str, ...]:
"""The statements of one migration file, comments stripped, empty ones dropped."""
parts: Final = _COMMENT.sub("", migration_sql).split(";")
return tuple(part.strip() for part in parts if part.strip())
def _split_top_level(text: str) -> tuple[str, ...]:
"""Split on the commas outside parentheses."""
depth: Final = accumulate((1 if char == "(" else -1 if char == ")" else 0) for char in text)
cuts: Final = tuple(
position for position, (char, level) in enumerate(zip(text, depth, strict=True)) if char == "," and level == 0
)
starts: Final = (0, *(cut + 1 for cut in cuts))
stops: Final = (*cuts, len(text))
pieces: Final = tuple(text[start:stop].strip() for start, stop in zip(starts, stops, strict=True))
return tuple(piece for piece in pieces if piece)
def _plain_columns(column_list: str) -> tuple[str, ...] | None:
"""The column names of an index or constraint column list, None when one entry is an
expression the derivation cannot turn into a count."""
matches: Final = tuple(_PLAIN_COLUMN.match(entry) for entry in _split_top_level(column_list))
if any(match is None for match in matches):
return None
return tuple(match.group(1) for match in matches if match is not None)
def _index_checks(match: "re.Match[str]", statement_text: str, source: str) -> Iterator[Check]:
unique, name, table, method, column_list = match.groups()
columns: Final = _plain_columns(column_list)
if columns is None:
return
label: Final = f'CREATE {"UNIQUE " if unique else ""}INDEX "{name}"'
if method is None or method.lower() == "btree":
yield Check("too_long", table, columns, label, source, index=name)
if unique:
yield Check("duplicates", table, columns, label, source, index=name)
def _alter_table_checks(table: str, action: str, source: str) -> Iterator[Check]:
if (set_not_null := _SET_NOT_NULL.match(action)) is not None:
yield Check("nulls", table, (set_not_null.group(1),), f'ALTER TABLE "{table}" {action}', source)
return
if (add_column := _ADD_COLUMN.match(action)) is not None:
column, body = add_column.groups()
if _NOT_NULL.search(body) is not None and _HAS_DEFAULT.search(body) is None:
yield Check(
"any_row", table, (column,), f'ALTER TABLE "{table}" ADD COLUMN "{column}" ... NOT NULL', source
)
return
if (foreign_key := _ADD_FOREIGN_KEY.match(action)) is not None:
constraint, columns, referenced_table, referenced_columns, tail = foreign_key.groups()
own: Final = _plain_columns(columns)
theirs: Final = _plain_columns(referenced_columns)
if own is None or theirs is None or len(own) != len(theirs) or "NOT VALID" in tail.upper():
return
yield Check(
"orphans",
table,
own,
_constraint_label(table, constraint, "FOREIGN KEY"),
source,
referenced_table=referenced_table,
referenced_columns=theirs,
)
return
if (unique := _ADD_UNIQUE.match(action)) is not None:
constraint, columns = unique.groups()
unique_columns: Final = _plain_columns(columns)
if unique_columns is not None:
yield Check("duplicates", table, unique_columns, _constraint_label(table, constraint, "UNIQUE"), source)
return
if (check := _ADD_CHECK.match(action)) is not None:
constraint, expression, not_valid = check.groups()
if not_valid is None:
yield Check(
"check_violations",
table,
(),
_constraint_label(table, constraint, "CHECK"),
source,
expression=expression,
)
def _constraint_label(table: str, constraint: str | None, kind: str) -> str:
named: Final = f'CONSTRAINT "{constraint}" ' if constraint else ""
return f'ALTER TABLE "{table}" ADD {named}{kind}'
def checks_for_statement(statement_text: str, source: str) -> tuple[Check, ...]:
"""The count queries one SQL statement needs before it can be applied, empty for a
statement that cannot fail on existing rows or that the derivation cannot read."""
if (index := _CREATE_INDEX.match(statement_text)) is not None:
return tuple(_index_checks(index, statement_text, source))
if (alter := _ALTER_TABLE.match(statement_text)) is not None:
table, actions = alter.groups()
return tuple(
check for action in _split_top_level(actions) for check in _alter_table_checks(table, action, source)
) # comprehension-ok: flattening the per-action checks of one statement
return ()
def checks_for_migration(name: str, migration_sql: str) -> tuple[Check, ...]:
return tuple(
check
for statement_text in statements(migration_sql)
for check in checks_for_statement(statement_text, f"migration {name}")
) # comprehension-ok: flattening the per-statement checks of one migration
def checks_for_index(table: str, name: str, columns: tuple[str, ...]) -> tuple[Check, ...]:
"""The entry-size check a request-log index needs before `CREATE INDEX CONCURRENTLY`."""
return (Check("too_long", table, columns, f'CREATE INDEX "{name}"', "request-log index", index=name),)
@ -424,42 +281,7 @@ def _entry_size_query(check: Check, schema: str, types: tuple[ColumnType, ...],
def _count_query(check: Check, schema: str, types: tuple[ColumnType, ...], limit_bytes: int) -> "sql.Composed":
from psycopg import sql
table: Final = sql.Identifier(schema, check.table)
columns: Final = sql.SQL(", ").join(sql.Identifier(column) for column in check.columns)
match check.kind:
case "too_long":
return _entry_size_query(check, schema, types, limit_bytes)
case "duplicates":
not_null: Final = sql.SQL(" AND ").join(
sql.SQL("{} IS NOT NULL").format(sql.Identifier(column)) for column in check.columns
)
return sql.SQL(
"SELECT count(*) FROM (SELECT 1 FROM {t} WHERE {nn} GROUP BY {cols} HAVING count(*) > 1) AS d"
).format(t=table, nn=not_null, cols=columns)
case "nulls":
return sql.SQL("SELECT count(*) FROM {t} WHERE {col} IS NULL").format(
t=table, col=sql.Identifier(check.columns[0])
)
case "any_row":
return sql.SQL("SELECT count(*) FROM {t}").format(t=table)
case "orphans":
pairs: Final = sql.SQL(" AND ").join(
sql.SQL("parent.{} = child.{}").format(sql.Identifier(theirs), sql.Identifier(own))
for own, theirs in zip(check.columns, check.referenced_columns, strict=True)
)
present: Final = sql.SQL(" AND ").join(
sql.SQL("child.{} IS NOT NULL").format(sql.Identifier(column)) for column in check.columns
)
return sql.SQL(
"SELECT count(*) FROM {t} AS child WHERE {present} AND NOT EXISTS (SELECT 1 FROM {ref} AS parent WHERE {pairs})"
).format(t=table, present=present, ref=sql.Identifier(schema, check.referenced_table or ""), pairs=pairs)
case "check_violations":
return sql.SQL("SELECT count(*) FROM {t} WHERE NOT ({expr})").format(
t=table,
expr=sql.SQL(check.expression or "true"), # pyright: ignore[reportArgumentType] # copied from our own migration file
)
return _entry_size_query(check, schema, types, limit_bytes)
def _applies(
@ -474,15 +296,7 @@ def _applies(
return False
if check.index is not None and _index_is_valid(connection, schema, check.index):
return False
if check.kind == "any_row":
return types is None
if types is None:
return False
if check.kind == "too_long":
return any(column.can_overflow_btree for column in types)
if check.kind == "orphans":
return check.referenced_table is not None and _table_exists(connection, schema, check.referenced_table)
return True
return types is not None and any(column.can_overflow_btree for column in types)
def _too_long_detail(check: Check, types: tuple[ColumnType, ...], values: tuple[object, ...], limit_bytes: int) -> str:
@ -515,21 +329,7 @@ def _run_check(
rows: Final = _int(values[0], 0) if values else 0
if rows == 0:
return None
match check.kind:
case "too_long":
return BlockedChange(check, rows, _too_long_detail(check, types, values, limit_bytes))
case "duplicates":
return BlockedChange(check, rows, "duplicated value groups")
case "nulls":
return BlockedChange(check, rows, "NULL")
case "any_row":
return BlockedChange(
check, rows, "rows, and a NOT NULL column with no DEFAULT cannot be added to a table with rows"
)
case "orphans":
return BlockedChange(check, rows, f'with no matching row in "{check.referenced_table}"')
case "check_violations":
return BlockedChange(check, rows, f"failing CHECK ({check.expression})")
return BlockedChange(check, rows, _too_long_detail(check, types, values, limit_bytes))
def find_blockers(
@ -571,19 +371,7 @@ def find_blockers(
def _remedy(check: Check) -> str:
joined: Final = ", ".join(check.columns)
match check.kind:
case "too_long":
return f'Shorten or hash the long "{check.table}" ({joined}) values so every row\'s index entry fits.'
case "duplicates":
return f'Remove or merge the duplicate "{check.table}" ({joined}) values.'
case "nulls":
return f'Fill the NULL "{check.table}"."{joined}" values, or give the column a default first.'
case "any_row":
return "Give the new column a DEFAULT in the migration, or add it nullable and fill it afterwards."
case "orphans":
return f'Delete or repoint the "{check.table}" ({joined}) values with no "{check.referenced_table}" row.'
case "check_violations":
return f'Fix the "{check.table}" rows that fail the expression before adding the constraint.'
return f'Shorten or hash the long "{check.table}" ({joined}) values so every row\'s index entry fits.'
def _describe(finding: Finding) -> str:
@ -607,30 +395,4 @@ def _describe(finding: Finding) -> str:
def report(findings: tuple[Finding, ...]) -> str:
"""The operator-facing refusal: one block per finding, then what to do next."""
lines: Final = "\n".join(_describe(finding) for finding in findings)
return (
f"{REFUSAL}\n{lines}\n"
"No schema change was made. Fix the rows and start again, or run `litellm --check_migrations_against_existing_data` with the new image to check again."
)
def pending_change_checks(
connection: "psycopg.Connection[tuple[object, ...]]",
schema: str,
migrations_dir: Path | None,
indexes: tuple[tuple[str, str, tuple[str, ...]], ...],
) -> tuple[Check, ...]:
"""Every check the pending migrations and the request-log indexes need. `migrations_dir`
is None when the schema is pushed rather than migrated, so only the indexes are checked."""
migration_checks: Final = (
tuple(
check
for name, migration_sql in pending_migrations(connection, schema, migrations_dir)
for check in checks_for_migration(name, migration_sql)
) # comprehension-ok: flattening the per-migration checks
if migrations_dir is not None
else ()
)
index_checks: Final = tuple(
check for table, name, columns in indexes for check in checks_for_index(table, name, columns)
) # comprehension-ok: flattening the per-index checks
return migration_checks + index_checks
return f"{REFUSAL}\n{lines}\nNo index was built. Fix the rows and the next index build picks it up."

View file

@ -0,0 +1,184 @@
"""`litellm --check_migrations_against_existing_data`: run the pending migrations for real, each
inside its own transaction that is always rolled back, so Postgres itself says whether this
database's rows and schema take them. Each migration holds the locks the real upgrade will,
and releases them before the next one runs."""
from __future__ import annotations
import re
import time
from collections.abc import Callable
from dataclasses import dataclass
from functools import reduce
from pathlib import Path
from typing import TYPE_CHECKING, Final
from litellm_proxy_extras.migration_data_checks import pending_migrations
if TYPE_CHECKING:
import psycopg
_CONCURRENTLY: Final = re.compile(
r"\b((?:CREATE\s+(?:UNIQUE\s+)?INDEX|DROP\s+INDEX|REINDEX\s+\w+)\s+)CONCURRENTLY\s+", re.IGNORECASE
)
UNDEFINED_SQLSTATES: Final = frozenset({"42703", "42P01", "42704", "42883"})
Migration = tuple[str, str]
@dataclass(frozen=True, slots=True)
class Applied:
name: str
seconds: float
replayed: tuple[str, ...] = ()
@dataclass(frozen=True, slots=True)
class Rejected:
name: str
error: str
seconds: float
@dataclass(frozen=True, slots=True)
class Baseline:
"""A database with tables but no `_prisma_migrations`: the boot baselines it instead of replaying
history, so the dry run asks the same question the boot does. `refusal` is None when the boot
will baseline, otherwise its reason (a custom migration folder, or the schema diff)."""
refusal: str | None
@property
def matches(self) -> bool:
return self.refusal is None
@dataclass(frozen=True, slots=True)
class DryRun:
applied: tuple[Applied, ...]
rejected: Rejected | None
baseline: Baseline | None = None
@property
def ok(self) -> bool:
return self.rejected is None and (self.baseline is None or self.baseline.matches)
def transactional(migration_sql: str) -> str:
"""The migration as it can run inside a transaction: `CONCURRENTLY` index builds become plain ones."""
return _CONCURRENTLY.sub(r"\1", migration_sql)
def _attempt(
connection: psycopg.Connection[tuple[object, ...]],
schema: str,
chain: tuple[Migration, ...],
clock: Callable[[], float],
) -> tuple[float, psycopg.Error | None]:
import psycopg
from psycopg import sql
started: Final = clock()
try:
with connection.transaction(force_rollback=True):
connection.execute(sql.SQL("SET LOCAL search_path TO {}").format(sql.Identifier(schema)))
for _, migration_sql in chain:
connection.execute(transactional(migration_sql))
except psycopg.Error as exc:
return clock() - started, exc
return clock() - started, None
def _run(
connection: psycopg.Connection[tuple[object, ...]],
schema: str,
pending: tuple[Migration, ...],
position: int,
done: DryRun,
clock: Callable[[], float],
) -> DryRun:
"""One migration in its own rolled-back transaction. A migration that fails because something
an earlier pending migration creates does not exist yet is rerun with those replayed in front."""
if done.rejected is not None:
return done
name: Final = pending[position][0]
seconds, error = _attempt(connection, schema, (pending[position],), clock)
if error is None:
return DryRun(done.applied + (Applied(name, seconds),), None, done.baseline)
if error.sqlstate not in UNDEFINED_SQLSTATES or position == 0:
return DryRun(done.applied, Rejected(name, str(error).strip(), seconds), done.baseline)
earlier: Final = pending[:position]
chained_seconds, chained_error = _attempt(connection, schema, earlier + (pending[position],), clock)
if chained_error is None:
replayed: Final = tuple(earlier_name for earlier_name, _ in earlier)
return DryRun(done.applied + (Applied(name, chained_seconds, replayed),), None, done.baseline)
return DryRun(done.applied, Rejected(name, str(chained_error).strip(), chained_seconds), done.baseline)
def _has_ledger(connection: psycopg.Connection[tuple[object, ...]], schema: str) -> bool:
from psycopg import sql
row: Final = connection.execute(
"SELECT to_regclass(%s)", (sql.Identifier(schema, "_prisma_migrations").as_string(connection),)
).fetchone()
return row is not None and row[0] is not None
def _has_litellm_tables(connection: psycopg.Connection[tuple[object, ...]], schema: str) -> bool:
row: Final = connection.execute(
"SELECT 1 FROM information_schema.tables WHERE table_schema = %s AND table_name LIKE 'LiteLLM\\_%%' LIMIT 1",
(schema,),
).fetchone()
return row is not None
def dry_run(
connection: psycopg.Connection[tuple[object, ...]],
schema: str,
migrations_dir: Path,
baseline_check: Callable[[], str | None] | None = None,
clock: Callable[[], float] = time.monotonic,
) -> DryRun:
"""Apply every pending migration in order, each in its own transaction that is rolled back
before the next one starts. The first migration Postgres rejects ends the run with its error.
`baseline_check` (None when the boot would baseline, otherwise why it refuses) stands in for
the replay on a database that has tables but no migration history, which the boot baselines."""
if baseline_check is not None and not _has_ledger(connection, schema) and _has_litellm_tables(connection, schema):
return DryRun((), None, Baseline(baseline_check()))
pending: Final = pending_migrations(connection, schema, migrations_dir)
return reduce(
lambda done, position: _run(connection, schema, pending, position, done, clock),
range(len(pending)),
DryRun((), None),
)
def _timing(applied: Applied) -> str:
replay: Final = f" (depends on {', '.join(applied.replayed)}, replayed in front of it)" if applied.replayed else ""
return f"\n {applied.name}: {applied.seconds:.1f} s{replay}"
def report(result: DryRun) -> str:
"""What the operator reads: the per-migration time the real upgrade will hold its locks,
or the migration Postgres rejected with Postgres's own words. Nothing was applied either way."""
if result.baseline is not None:
if result.baseline.matches:
return (
"No migration history and the schema matches this build, "
"the proxy baselines it at boot, nothing to dry-run."
)
return f"No migration history and the proxy cannot baseline it at boot: {result.baseline.refusal}"
timings: Final = "".join(_timing(applied) for applied in result.applied)
if result.rejected is None:
count: Final = len(result.applied)
if count == 0:
return "No pending migrations, this database is on this version's schema."
return (
f"{count} pending migration(s) applied one at a time in a dry run, each rolled back before the next "
f"so its locks are released, nothing changed:{timings}"
)
return (
f"Migration {result.rejected.name} fails on this database after {result.rejected.seconds:.1f} s, "
f"nothing was applied.\n{result.rejected.error}"
+ (f"\nApplied before it in the dry run:{timings}" if timings else "")
)

View file

@ -86,6 +86,18 @@ def migration_files(directory: Path) -> tuple[tuple[str, str], ...]:
)
def has_packaged_history(migrations_dir: Path) -> bool:
"""Whether `migrations_dir` carries exactly the migration files and schema.prisma this package ships,
the only history the boot baselines an existing database onto."""
packaged_dir: Final = Path(__file__).parent
migrations: Final = migration_files(migrations_dir)
return (
bool(migrations)
and migrations == migration_files(packaged_dir)
and (migrations_dir / "schema.prisma").read_bytes() == (packaged_dir / "schema.prisma").read_bytes()
)
def baseline_current_schema(
coordinator: MigrationCoordinator,
schema: str,
@ -95,14 +107,9 @@ def baseline_current_schema(
) -> None:
from psycopg import sql
packaged_dir: Final = Path(__file__).parent
migrations: Final = migration_files(migrations_dir)
if (
not migrations
or migrations != migration_files(packaged_dir)
or (migrations_dir / "schema.prisma").read_bytes() != (packaged_dir / "schema.prisma").read_bytes()
):
if not has_packaged_history(migrations_dir):
raise RuntimeError("Cannot automatically baseline an existing database with custom migration history.")
migrations: Final = migration_files(migrations_dir)
coordinator.acquire_prisma_lock()
existing: Final = coordinator.connection.execute(

View file

@ -11,18 +11,13 @@ from collections.abc import Callable
from contextlib import AbstractContextManager, nullcontext
from dataclasses import dataclass, replace
from pathlib import Path
from typing import TYPE_CHECKING, Final, Optional, Protocol, Union
from typing import TYPE_CHECKING, Final, Protocol
from litellm_proxy_extras import prisma_toolchain
from litellm_proxy_extras._logging import logger
from litellm_proxy_extras._redaction import _redact_credentials
from litellm_proxy_extras.migration_data_checks import (
Finding,
find_blockers,
migration_data_check_timeout_seconds,
pending_change_checks,
report,
)
from litellm_proxy_extras.migration_dry_run import dry_run
from litellm_proxy_extras.migration_dry_run import report as dry_run_report
from litellm_proxy_extras.migration_lock import held_migration_lock
from litellm_proxy_extras.prisma_toolchain import (
PRISMA_COMMAND_TIMEOUT_ENV_VAR,
@ -36,7 +31,6 @@ from litellm_proxy_extras.replica_identity import (
apply_replica_identity_full,
)
from litellm_proxy_extras.request_log_indexes import (
REQUEST_LOG_INDEXES,
ensure_request_log_indexes,
filter_request_log_index_diff,
)
@ -52,7 +46,7 @@ class LensCheckConnect(Protocol):
) -> "psycopg.Connection[tuple[object, ...]]": ...
def str_to_bool(value: Optional[str]) -> bool:
def str_to_bool(value: str | None) -> bool:
if value is None:
return False
return value.lower() in ("true", "1", "t", "y", "yes")
@ -169,7 +163,7 @@ def _without_sql_comments(statement: str) -> str:
).strip()
def _without_spend_logs_pk_clauses(statement: str) -> Optional[str]:
def _without_spend_logs_pk_clauses(statement: str) -> str | None:
prefix_match = _SPEND_LOGS_ALTER_RE.match(statement)
if not prefix_match:
return statement
@ -217,7 +211,7 @@ def _max_migration_timestamp(names) -> int:
return max(_migration_timestamp(n) for n in names)
def _redacted_command(command: object) -> Union[str, tuple[str, ...], list[str]]:
def _redacted_command(command: object) -> str | tuple[str, ...] | list[str]:
if isinstance(command, tuple):
return tuple(_redact_credentials(str(argument)) for argument in command)
if isinstance(command, list):
@ -404,7 +398,7 @@ class ProxyExtrasDBManager:
return tuple(row) if row is not None else ()
@staticmethod
def _failed_migration_logs(migration_name: str, started_at: str) -> Optional[str]:
def _failed_migration_logs(migration_name: str, started_at: str) -> str | None:
row: Final = ProxyExtrasDBManager._read_migration_ledger(
"SELECT logs FROM {} WHERE migration_name = %s AND started_at = %s::timestamptz "
"AND finished_at IS NULL AND rolled_back_at IS NULL",
@ -721,10 +715,10 @@ class ProxyExtrasDBManager:
return row is not None
@staticmethod
def _prisma_schema_param(url: str) -> Optional[str]:
def _prisma_schema_param(url: str) -> str | None:
"""The `schema` query param Prisma uses to pick its target schema,
or None when the URL does not set one."""
from urllib.parse import urlparse, parse_qsl
from urllib.parse import parse_qsl, urlparse
return next(
(v for k, v in parse_qsl(urlparse(url).query) if k == "schema"),
@ -1424,7 +1418,6 @@ class ProxyExtrasDBManager:
bool: True if setup was successful, False otherwise
"""
ensure_prisma_toolchain(prisma_command=_get_prisma_command(), prisma_env=_get_prisma_env())
ProxyExtrasDBManager.raise_if_rows_block_pending_changes(use_migrate=use_migrate)
migrated = ProxyExtrasDBManager._run_migrations(use_migrate=use_migrate, use_v2_resolver=use_v2_resolver)
if not migrated:
return False
@ -1433,46 +1426,69 @@ class ProxyExtrasDBManager:
return True
@staticmethod
def migration_data_check_findings(use_migrate: bool = True) -> tuple[Finding, ...]:
"""What in the existing rows would make a pending migration or request-log index
fail, read with count queries only: no write, no table lock. Empty when nothing
blocks, when no database is configured, or when it cannot be reached, so the
migration step that follows reports the connection problem the way it does today."""
def _baseline_refusal(prisma_dir: Path) -> str | None:
"""Why `baseline_current_schema` would refuse this database at boot, None when it would baseline:
the migrations folder must be the packaged one, then `prisma migrate diff` must find the live
database and this build's schema.prisma identical."""
from litellm_proxy_extras.migration_lock import migration_environment
from litellm_proxy_extras.migration_recovery import has_packaged_history
if not has_packaged_history(prisma_dir):
return (
f"{prisma_dir} is a custom migration history (LITELLM_MIGRATION_DIR), it differs from the "
"migrations this build ships, and the proxy only baselines onto the packaged history."
)
try:
prisma_toolchain.run_prisma(
(
_get_prisma_command(),
"migrate",
"diff",
"--from-schema-datasource",
str(prisma_dir / "schema.prisma"),
"--to-schema-datamodel",
str(prisma_dir / "schema.prisma"),
"--exit-code",
),
timeout=prisma_toolchain.prisma_command_timeout(),
env=migration_environment(_get_prisma_env()),
)
except subprocess.CalledProcessError as exc:
diff: Final = str(exc.stdout or exc.stderr or exc).strip()
return f"the schema differs from this build. prisma migrate diff:\n{diff}"
except subprocess.TimeoutExpired as exc:
return f"prisma migrate diff did not finish within {exc.timeout} s"
return None
@staticmethod
def migration_dry_run_summary(use_migrate: bool = True) -> tuple[bool, str]:
"""`--check_migrations_against_existing_data`: the pending migrations applied for real on
DIRECT_URL or DATABASE_URL inside one transaction that is rolled back, so Postgres reports
what the upgrade would hit. (True, the per-migration timings) when they all apply,
(False, Postgres's own error) when one does not."""
import psycopg
if not use_migrate:
return (True, "The schema is pushed rather than migrated, there is nothing to dry-run.")
database_url: Final = os.environ.get("DATABASE_URL")
if not database_url:
return ()
return (False, "DATABASE_URL is not set.")
direct_url: Final = ProxyExtrasDBManager._strip_prisma_query_params(
os.environ.get("DIRECT_URL") or database_url
)
schema: Final = ProxyExtrasDBManager._prisma_schema_param(database_url) or "public"
migrations_dir: Final = Path(ProxyExtrasDBManager._get_prisma_dir()) / "migrations" if use_migrate else None
indexes: Final = tuple((index.table, index.name, index.columns) for index in REQUEST_LOG_INDEXES)
prisma_dir: Final = Path(ProxyExtrasDBManager._get_prisma_dir())
try:
with psycopg.connect(direct_url, connect_timeout=10, autocommit=True) as connection:
checks: Final = pending_change_checks(connection, schema, migrations_dir, indexes)
return find_blockers(connection, schema, checks, migration_data_check_timeout_seconds())
result: Final = dry_run(
connection,
schema,
prisma_dir / "migrations",
baseline_check=lambda: ProxyExtrasDBManager._baseline_refusal(prisma_dir),
)
except psycopg.OperationalError as exc:
logger.warning(
"Could not reach the database for the pre-migration row check: %s", _redact_credentials(str(exc))
)
return ()
@staticmethod
def migration_data_check_report(use_migrate: bool = True) -> Optional[str]:
"""The operator-facing refusal for `migration_data_check_findings`, None when nothing blocks."""
findings: Final = ProxyExtrasDBManager.migration_data_check_findings(use_migrate=use_migrate)
return report(findings) if findings else None
@staticmethod
def raise_if_rows_block_pending_changes(use_migrate: bool) -> None:
"""Refuse to run migrations that the existing rows would make fail, before any is
applied. The RuntimeError carries the report and reaches the operator through the
proxy's "Database migration cannot proceed" exit."""
refusal: Final = ProxyExtrasDBManager.migration_data_check_report(use_migrate=use_migrate)
if refusal is not None:
raise RuntimeError(refusal)
return (False, f"Could not reach the database: {_redact_credentials(str(exc))}")
return (result.ok, dry_run_report(result))
@staticmethod
def build_request_log_indexes(build: Callable[[str, str], bool] = ensure_request_log_indexes) -> bool:
@ -1724,5 +1740,4 @@ class ProxyExtrasDBManager:
time.sleep(random.randrange(5, 15))
finally:
os.chdir(original_dir)
pass
return False

View file

@ -4345,6 +4345,7 @@ dependencies = [
"litellm-inference-transcription",
"litellm-llms",
"litellm-llms-types",
"litellm-model-catalog",
"litellm-router-types",
"litellm-secrets",
"litellm-secrets-aws",

View file

@ -53,6 +53,7 @@ pub(super) const RULES: &[Rule] = &[
|mapping| {
mapping.error_str.contains("invalid_request_error")
&& !mapping.error_str.contains("Incorrect API key provided")
&& !matches!(mapping.status, Some(401 | 404))
},
PublicError::BadRequest,
),
@ -132,6 +133,23 @@ mod tests {
assert_eq!(classified(Some(400), text), Some(expected));
}
#[rstest::rstest]
#[case::unauthorized(Some(401), None)]
#[case::not_found(Some(404), None)]
#[case::forbidden(Some(403), Some(PublicError::BadRequest))]
#[case::rate_limited(Some(429), Some(PublicError::BadRequest))]
#[case::server_error(Some(500), Some(PublicError::BadRequest))]
#[case::no_status(None, Some(PublicError::BadRequest))]
fn invalid_request_errors_leave_401_and_404_to_the_status_mapping(
#[case] status: Option<u16>,
#[case] expected: Option<PublicError>,
) {
assert_eq!(
classified(status, "invalid_request_error bad field"),
expected
);
}
#[rstest::rstest]
#[case::rate_limit_before_context_window(
"rate limit and This model's maximum context length is 10",

View file

@ -25,3 +25,5 @@ A new `N`, image size or resolution in an existing key shape parses with no Rust
## Tests
`BUNDLED_REGISTRY` embeds the root `model_prices_and_context_window.json`, so the tests over it fail a JSON change with a bad value or a key no type reads
`Catalog::problems` is the one definition of a mergeable cost map. The bundled test asserts it is empty, and the cost map guard workflow runs it on a pull request's JSON through `cargo run --bin check-registry <path>`

View file

@ -10,6 +10,7 @@ schema = ["dep:schemars", "litellm-cost/schema"]
[dependencies]
litellm-cost = { workspace = true, features = ["serde"] }
litellm-error.workspace = true
indexmap = { version = "2.14.0", features = ["serde"] }
jiff = { version = "0.2.37", default-features = false, features = ["std", "tzdb-bundle-always"] }
@ -24,5 +25,4 @@ time.workspace = true
[dev-dependencies]
jsonschema = { version = "0.55.1", default-features = false }
litellm-error.workspace = true
rstest.workspace = true

View file

@ -0,0 +1,26 @@
use std::{env, fs, path::PathBuf, process::ExitCode};
use litellm_model_catalog::Catalog;
fn main() -> ExitCode {
let Some(path) = env::args_os().nth(1).map(PathBuf::from) else {
eprintln!("usage: check-registry <model_prices_and_context_window.json>");
return ExitCode::from(2);
};
let body = match fs::read(&path) {
Ok(body) => body,
Err(error) => {
eprintln!("{}: {error}", path.display());
return ExitCode::FAILURE;
}
};
let problems = Catalog::problems(&body);
for problem in &problems {
eprintln!("{}", litellm_error::Report::new(problem));
}
if problems.is_empty() {
ExitCode::SUCCESS
} else {
ExitCode::FAILURE
}
}

View file

@ -5,7 +5,7 @@ use indexmap::IndexMap;
use serde::Deserialize;
use serde_json::Value;
use crate::error::{Error, Loaded, Rejects, Repr};
use crate::error::{Error, Loaded, Rejects, Repr, pointer_child, pointer_for_key};
use crate::fallback::{FallbackGeneralizations, FallbackRule};
use crate::model::Model;
use crate::store::{Rows, Store};
@ -211,6 +211,36 @@ impl Catalog {
})
}
pub fn problems(body: &[u8]) -> Vec<Error> {
let Loaded {
value: catalog,
errors,
} = match Self::parse(body) {
Ok(loaded) => loaded,
Err(fatal) => return vec![fatal],
};
let unread = catalog.rows().entries.iter().flat_map(|(model, entry)| {
EntryWire::read(entry.fields())
.unknown
.into_iter()
.map(move |(key, _)| {
Error::new(
pointer_child(&pointer_for_key(model), &key),
Repr::UnreadField,
)
})
});
let no_fallback_rules = catalog
.fallback_rules()
.is_none_or(<[_]>::is_empty)
.then(|| Error::from_key(FALLBACK_GENERALIZATIONS, Repr::NoFallbackRules));
errors
.into_iter()
.chain(unread)
.chain(no_fallback_rules)
.collect()
}
pub fn row_errors(model: &str, fields: &Fields) -> Vec<Error> {
let (_, _, errors) = read_row(fields);
errors
@ -404,3 +434,85 @@ fn unknown_timezone(wire: &EntryWire) -> Option<Error> {
)
})
}
#[cfg(test)]
mod tests {
use rstest::rstest;
use super::*;
const RULES: &str = r#""fallback_generalizations":{"rules":[{"name":"family","pattern":"^new-","model_info":{"mode":"chat"}}]}"#;
#[rstest]
#[case::clean(format!(r#"{{{RULES},"m":{{"litellm_provider":"test"}}}}"#), vec![])]
#[case::not_json("{".to_owned(), vec![("", "json")])]
#[case::row_error(format!(r#"{{{RULES},"m":{{}}}}"#), vec![("/m/litellm_provider", "missing_provider")])]
#[case::unread_key(
format!(r#"{{{RULES},"m":{{"litellm_provider":"test","typo_cost":1}}}}"#),
vec![("/m/typo_cost", "unread_field")]
)]
#[case::empty_model_key(
format!(r#"{{{RULES},"":{{"litellm_provider":"test","future_field":true}}}}"#),
vec![("//future_field", "unread_field")]
)]
#[case::unread_key_on_invalid_row(
format!(r#"{{{RULES},"m":{{"typo_cost":1}}}}"#),
vec![("/m/litellm_provider", "missing_provider"), ("/m/typo_cost", "unread_field")]
)]
#[case::escaped_unread_key_on_invalid_row(
format!(r#"{{{RULES},"open/router~v1":{{"typo/cost~":1}}}}"#),
vec![("/open~1router~0v1/litellm_provider", "missing_provider"), ("/open~1router~0v1/typo~1cost~0", "unread_field")]
)]
#[case::no_rules(
r#"{"m":{"litellm_provider":"test"}}"#.to_owned(),
vec![("/fallback_generalizations", "no_fallback_rules")]
)]
#[case::empty_rules(
r#"{"fallback_generalizations":{"rules":[]},"m":{"litellm_provider":"test"}}"#.to_owned(),
vec![("/fallback_generalizations", "no_fallback_rules")]
)]
fn problems_list_what_a_cost_map_change_must_fix(
#[case] body: String,
#[case] expected: Vec<(&str, &str)>,
) {
let problems: Vec<_> = Catalog::problems(body.as_bytes())
.iter()
.map(|problem| (problem.path().to_owned(), <&str>::from(problem.kind())))
.collect();
let expected: Vec<_> = expected
.into_iter()
.map(|(path, kind)| (path.to_owned(), kind))
.collect();
assert_eq!(problems, expected);
}
#[rstest]
fn runtime_parsing_preserves_custom_fields_and_recovers_valid_rows() {
let body = br#"{
"good":{"litellm_provider":"test","custom_metadata":true},
"bad":{"litellm_provider":"test","input_cost_per_token":"wrong"}
}"#;
let loaded = Catalog::parse(body).unwrap();
assert_eq!(loaded.errors.len(), 1);
assert_eq!(loaded.errors[0].path(), "/bad/input_cost_per_token");
let (catalog, rejected) = loaded.value.without_rejected(&loaded.errors);
assert_eq!(rejected.len(), 1);
assert!(catalog.lookup("bad").is_none());
let good = catalog.lookup("good").unwrap();
assert!(good.entry.model().is_some());
assert_eq!(
good.entry.field("custom_metadata"),
Some(&Value::Bool(true))
);
assert!(
Catalog::problems(body)
.iter()
.any(|problem| problem.kind() == crate::ErrorKind::UnreadField)
);
assert!(
Catalog::problems(body)
.iter()
.any(|problem| problem.kind() == crate::ErrorKind::NoFallbackRules)
);
}
}

View file

@ -109,6 +109,12 @@ pub(crate) enum Repr {
#[source]
cause: Option<DateCause>,
},
#[error("not a field litellm reads")]
#[strum_discriminants(strum(serialize = "unread_field"))]
UnreadField,
#[error("has no rules")]
#[strum_discriminants(strum(serialize = "no_fallback_rules"))]
NoFallbackRules,
}
/// What a problem takes out of the served map
@ -148,7 +154,7 @@ impl Error {
| Repr::NoOffPeakHours
| Repr::EmptyList
| Repr::InvalidDate { .. } => Rejects::Row,
Repr::MissingProvider => Rejects::Nothing,
Repr::MissingProvider | Repr::UnreadField | Repr::NoFallbackRules => Rejects::Nothing,
}
}
}
@ -334,6 +340,8 @@ mod tests {
#[case::no_off_peak_hours(Repr::NoOffPeakHours, Rejects::Row)]
#[case::empty_list(Repr::EmptyList, Rejects::Row)]
#[case::invalid_date(Repr::InvalidDate { cause: None }, Rejects::Row)]
#[case::unread_field(Repr::UnreadField, Rejects::Nothing)]
#[case::no_fallback_rules(Repr::NoFallbackRules, Rejects::Nothing)]
fn policy_classifies_each_error(#[case] error: Repr, #[case] expected: Rejects) {
assert_eq!(Error::from(error).rejects(), expected);
}

View file

@ -366,18 +366,14 @@ fn written_is_the_document_in_order_including_sections_and_bad_entries() {
}
#[rstest]
fn the_bundled_registry_parses_with_no_problems_and_every_key_understood() {
let loaded = Catalog::parse(BUNDLED_REGISTRY).unwrap();
let catalog = loaded.value;
assert!(loaded.errors.is_empty(), "{:#?}", loaded.errors);
let unrecognized: Vec<_> = catalog.unrecognized().collect();
assert!(unrecognized.is_empty(), "{unrecognized:#?}");
assert!(
catalog
.fallback_rules()
.is_some_and(|rules| !rules.is_empty())
);
fn the_bundled_registry_has_no_problems() {
let problems: Vec<_> = Catalog::problems(BUNDLED_REGISTRY)
.iter()
.map(report)
.collect();
assert!(problems.is_empty(), "{problems:#?}");
const REFERENCE_MODEL_COUNT: usize = 4303;
let catalog = Catalog::parse(BUNDLED_REGISTRY).unwrap().value;
assert!(catalog.model_count() as f64 >= REFERENCE_MODEL_COUNT as f64 * 0.9);
}

View file

@ -106,6 +106,18 @@ GIL handling to `litellm-host-python`.
## Data Handling
- Reuse large immutable bundled data and parsed catalogs across calls, including in
replacement APIs. Share Rust storage and cache repeated Python projections instead
of reparsing, serializing or copying the full payload on every access
- Cache owned Python objects with `PyOnceLock<Py<T>>`, then bind and clone references
while attached. Do not use blocking `OnceLock::get_or_init` while attached
([PyO3 initialization](https://pyo3.rs/v0.29.2/faq.html),
[reference ownership](https://pyo3.rs/v0.29.2/types.html))
- Never reuse Python objects across interpreters or interpreter reinitialization.
Once-only caches are for immutable data. Reloadable/custom data needs versioned
entries or invalidation, and caller-mutable Python containers must not be shared
- Test unchanged contents and allocation reuse, not just equal values. Keep this
coverage when replacing a bundled-data accessor with a catalog API
- OCR payloads can contain personal data and large base64 images. Do not log
payloads or provider responses.
- Avoid copying large payloads more than needed. The current JSON round-trip is

View file

@ -49,6 +49,7 @@ litellm-secrets-types.workspace = true
litellm-llms-types.workspace = true
litellm-router-types.workspace = true
litellm-host-python.workspace = true
litellm-model-catalog.workspace = true
litellm-token-counter = { path = "../token-counter", default-features = false }
pyo3.workspace = true
prost.workspace = true

View file

@ -8,6 +8,7 @@ mod execution;
mod http;
mod logger;
mod marshal;
mod model_registry;
mod preflight;
mod python_settings;
mod routes;
@ -30,6 +31,8 @@ mod _native {
#[pymodule_export]
use crate::logger::NativeDiagnosticProcessor;
#[pymodule_export]
use crate::model_registry::bundled_model_cost_map;
#[pymodule_export]
use crate::routes::audio_transcription::{atranscription, transcription};
#[pymodule_export]
use crate::routes::chat_completions::{acompletion, completion};
@ -110,6 +113,7 @@ mod tests {
"NativeClickHouseSpendConfig",
"NativeClickHouseSpendStorage",
"trace_encode_error",
"bundled_model_cost_map",
"TokenCounter",
"Tokenizer",
"gil_stats",

View file

@ -0,0 +1,30 @@
use pyo3::prelude::*;
use pyo3::sync::PyOnceLock;
use pyo3::types::PyBytes;
#[pyfunction]
pub fn bundled_model_cost_map(py: Python<'_>) -> Bound<'_, PyBytes> {
static BYTES: PyOnceLock<Py<PyBytes>> = PyOnceLock::new();
BYTES
.get_or_init(py, || {
PyBytes::new(py, litellm_model_catalog::BUNDLED_REGISTRY).unbind()
})
.bind(py)
.clone()
}
#[cfg(test)]
mod tests {
use super::*;
#[rstest::rstest]
fn repeated_calls_reuse_the_bundled_registry_bytes() {
Python::initialize();
let first = Python::attach(|py| bundled_model_cost_map(py).unbind());
Python::attach(|py| {
let second = bundled_model_cost_map(py);
assert!(first.bind(py).is(&second));
assert_eq!(second.as_bytes(), litellm_model_catalog::BUNDLED_REGISTRY);
});
}
}

View file

@ -520,7 +520,7 @@ aiohttp_trust_env: bool = False # set to true to use HTTP_ Proxy settings
disable_aiohttp_transport: bool = False # Set this to true to use httpx instead
disable_aiohttp_trust_env: bool = False # When False, aiohttp will respect HTTP(S)_PROXY env vars
force_ipv4: bool = False # when True, litellm will force ipv4 for all LLM requests. Some users have seen httpx ConnectionError when using ipv6.
http2: bool = False
http2: bool | None = None
network_mock: bool = False # When True, use mock transport — no real network calls
####### STOP SEQUENCE LIMIT #######

View file

@ -57,7 +57,7 @@ def resolve_embedding_max_input_tokens(
return configured_max_input_tokens
if router is None:
return None
deployment_max_input_tokens, _ = router.get_configured_token_limits(embedding_model)
deployment_max_input_tokens, _ = router.registry().group(embedding_model, field="token_limits")
return deployment_max_input_tokens

View file

@ -830,6 +830,13 @@ class LLMCachingHandler:
if is_internal_call.get():
return
if logging_obj.defer_async_logging and not logging_obj.stream:
from litellm.utils import dispatch_success_logging
logging_obj.model_call_details["cache_hit"] = cache_hit # rebind-ok: capture cache status before deferring
dispatch_success_logging(logging_obj, cached_result, start_time, end_time, False, False)
return
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(

View file

@ -29,6 +29,9 @@ it does not belong to. Queued clients are therefore bucketed by what it takes to
close them, and each bucket is ordered by deadline, so a reap walks the entries
that are due rather than the whole queue.
A bucket whose loop has been collected is dropped, and its clients are left to the
collector, as they are when the queue is full.
The queue holds its clients weakly, so waiting out a grace window never keeps
alive anything the collector would have reclaimed first.
"""
@ -36,6 +39,7 @@ alive anything the collector would have reclaimed first.
import asyncio
import contextlib
import inspect
import itertools
import threading
import time
import weakref
@ -44,6 +48,8 @@ from collections.abc import Awaitable, Callable, Iterator
from dataclasses import dataclass, replace
from typing import Final
import httpx
from litellm.constants import (
EVICTED_LLM_CLIENT_CLOSE_GRACE_SECONDS,
EVICTED_LLM_CLIENT_CLOSE_MAX_PENDING,
@ -53,6 +59,9 @@ _CLOSABLE_ANYWHERE: Final = "closable-anywhere"
_CLOSABLE_ON_ANY_LOOP: Final = "closable-on-any-loop"
_BucketKey = str | int
_LOOP_TOKENS: Final[weakref.WeakKeyDictionary[asyncio.AbstractEventLoop, int]] = weakref.WeakKeyDictionary()
_LOOP_TOKEN_COUNTER: Final = itertools.count()
_LOOP_TOKEN_LOCK: Final = threading.Lock()
@dataclass(frozen=True, slots=True)
@ -63,12 +72,12 @@ class _PendingClose:
collector would otherwise have reclaimed first.
``needs_loop`` is set for a client whose close is a coroutine; those can only
be closed from the event loop they were evicted on, recorded in ``loop_id``.
be closed from the event loop they were evicted on, recorded in ``loop_token``.
A client that closes synchronously carries neither constraint.
"""
client_ref: "weakref.ref[object]"
loop_id: int | None
loop_token: int | None
needs_loop: bool
close_after: float
@ -77,16 +86,24 @@ def _bucket_key(pending: _PendingClose) -> _BucketKey:
"""Which reaps can close this entry: any at all, any running a loop, or one loop's."""
if not pending.needs_loop:
return _CLOSABLE_ANYWHERE
if pending.loop_id is None:
if pending.loop_token is None:
return _CLOSABLE_ON_ANY_LOOP
return pending.loop_id
return pending.loop_token
def _running_loop_id() -> int | None:
def running_loop_token() -> int | None:
try:
return id(asyncio.get_running_loop())
loop: Final = asyncio.get_running_loop()
except RuntimeError:
return None
try:
token: Final[int | None] = _LOOP_TOKENS.get(loop)
if token is not None:
return token
with _LOOP_TOKEN_LOCK:
return _LOOP_TOKENS.setdefault(loop, next(_LOOP_TOKEN_COUNTER))
except TypeError:
return id(loop)
def _close_function(client: object) -> Callable[[], object] | None:
@ -94,13 +111,24 @@ def _close_function(client: object) -> Callable[[], object] | None:
return close_fn
def _transport_of(client: object) -> object:
"""The httpx transport behind an SDK wrapper, a litellm handler, or a bare client."""
def _mounted_transports(holder: object) -> tuple[object, ...]:
mounts: Final[object] = getattr(holder, "_mounts", None)
if not isinstance(mounts, dict):
return ()
return tuple(
mounted
for mounted in mounts.values() # pyright: ignore[reportUnknownVariableType] # untyped httpx mount map
if isinstance(mounted, (httpx.AsyncBaseTransport, httpx.BaseTransport))
)
def _transports_of(client: object) -> tuple[object, ...]:
"""The httpx transports behind an SDK wrapper, a litellm handler, or a bare client, proxy mounts included."""
for holder in (getattr(client, "_client", None), getattr(client, "client", None), client):
transport: object = getattr(holder, "_transport", None)
if transport is not None:
return transport
return None
return (transport, *_mounted_transports(holder))
return ()
def _connection_is_idle(connection: object) -> bool:
@ -123,6 +151,17 @@ def _pool_has_busy_connection(transport: object) -> bool | None:
)
def _transport_in_flight(transport: object) -> bool:
spread_in_flight: Final[object] = getattr(transport, "has_request_in_flight", None)
if callable(spread_in_flight):
return bool(spread_in_flight())
pooled_busy: Final = _pool_has_busy_connection(transport)
if pooled_busy is not None:
return pooled_busy
session: Final[object] = getattr(transport, "client", None)
return bool(getattr(getattr(session, "connector", None), "_acquired", None))
def _has_connection_in_flight(client: object) -> bool:
"""Whether the client is servicing a request right now.
@ -138,12 +177,7 @@ def _has_connection_in_flight(client: object) -> bool:
try:
if getattr(getattr(client, "connection_pool", None), "_in_use_connections", None):
return True
transport: Final = _transport_of(client)
pooled_busy: Final = _pool_has_busy_connection(transport)
if pooled_busy is not None:
return pooled_busy
session: Final[object] = getattr(transport, "client", None)
return bool(getattr(getattr(session, "connector", None), "_acquired", None))
return any(_transport_in_flight(transport) for transport in _transports_of(client))
except Exception: # noqa: BLE001 - a client that cannot report its state is treated as idle
return False
@ -167,6 +201,8 @@ class EvictedClientCloser:
self._clock = clock
self._owned: weakref.WeakSet[object] = weakref.WeakSet()
self._buckets: dict[_BucketKey, deque[_PendingClose]] = {} # mutable-ok: deadline-ordered queues
self._dead_loop_tokens: deque[int] = deque() # mutable-ok: loop finalizers append without taking the queue lock
self._watched_loop_tokens: set[int] = set() # mutable-ok: tracks tokens until their dead buckets are discarded
self._pending_count = 0
self._queue_lock = threading.Lock() # the cache is reachable from every worker thread's loop
self._close_tasks: set[asyncio.Task[None]] = set() # mutable-ok: strong refs to running closes
@ -197,13 +233,19 @@ class EvictedClientCloser:
close_fn: Final = _close_function(client)
if close_fn is None:
return
loop_token: Final = running_loop_token()
needs_loop: Final = inspect.iscoroutinefunction(close_fn)
if self._pending_count >= self._max_pending:
return
self._discard_dead_loops()
if self._pending_count >= self._max_pending:
return
if needs_loop and loop_token is not None:
self._watch_loop(loop_token)
self._enqueue(
_PendingClose(
client_ref=weakref.ref(client),
loop_id=_running_loop_id(),
needs_loop=inspect.iscoroutinefunction(close_fn),
loop_token=loop_token,
needs_loop=needs_loop,
close_after=self._clock() + self._grace_seconds,
)
)
@ -217,7 +259,7 @@ class EvictedClientCloser:
if not self._pending_count:
return
now: Final = self._clock()
for pending in self._take_due(_running_loop_id(), now):
for pending in self._take_due(running_loop_token(), now):
client = pending.client_ref()
if client is None:
continue
@ -230,6 +272,29 @@ class EvictedClientCloser:
def pending_count(self) -> int:
return self._pending_count
def _watch_loop(self, loop_token: int) -> None:
if loop_token in self._watched_loop_tokens:
return
try:
weakref.finalize(asyncio.get_running_loop(), self._dead_loop_tokens.append, loop_token)
except TypeError:
return
self._watched_loop_tokens.add(loop_token)
def _discard_dead_loops(self) -> None:
with self._queue_lock:
self._discard_dead_loops_locked()
def _discard_dead_loops_locked(self) -> None:
while self._dead_loop_tokens:
self._discard_dead_loop_token_locked(self._dead_loop_tokens.popleft())
def _discard_dead_loop_token_locked(self, loop_token: int) -> None:
self._watched_loop_tokens.discard(loop_token)
bucket: Final = self._buckets.pop(loop_token, None)
if bucket is not None:
self._pending_count -= len(bucket)
def _enqueue(self, pending: _PendingClose) -> None:
"""Append to the entry's bucket, dropping any dead entries it queues behind.
@ -238,6 +303,7 @@ class EvictedClientCloser:
the front rather than having to be searched for.
"""
with self._queue_lock:
self._discard_dead_loops_locked()
bucket: Final = self._buckets.setdefault(_bucket_key(pending), deque())
while bucket and bucket[0].client_ref() is None:
bucket.popleft()
@ -245,9 +311,12 @@ class EvictedClientCloser:
bucket.append(pending)
self._pending_count += 1
def _take_due(self, loop_id: int | None, now: float) -> tuple[_PendingClose, ...]:
buckets = (_CLOSABLE_ANYWHERE,) if loop_id is None else (_CLOSABLE_ANYWHERE, _CLOSABLE_ON_ANY_LOOP, loop_id)
def _take_due(self, loop_token: int | None, now: float) -> tuple[_PendingClose, ...]:
buckets = (
(_CLOSABLE_ANYWHERE,) if loop_token is None else (_CLOSABLE_ANYWHERE, _CLOSABLE_ON_ANY_LOOP, loop_token)
)
with self._queue_lock:
self._discard_dead_loops_locked()
return tuple(pending for key in buckets for pending in self._drain_locked(key, now))
def _drain_locked(self, key: _BucketKey, now: float) -> Iterator[_PendingClose]:

View file

@ -2,10 +2,9 @@
Add the event loop to the cache key, to prevent event loop closed errors.
"""
import asyncio
from typing import Final
from .evicted_client_closer import EvictedClientCloser, default_evicted_client_closer
from .evicted_client_closer import EvictedClientCloser, default_evicted_client_closer, running_loop_token
from .in_memory_cache import InMemoryCache
@ -43,24 +42,22 @@ class LLMClientCache(InMemoryCache):
self.evicted_client_closer.schedule(evicted)
self.evicted_client_closer.reap()
def update_cache_key_with_event_loop(self, key):
def update_cache_key_with_event_loop(self, key: str) -> str:
"""
Add the event loop to the cache key, to prevent event loop closed errors.
If none, use the key as is.
"""
try:
event_loop: Final = asyncio.get_running_loop()
stringified_event_loop: Final = str(id(event_loop))
return f"{key}-{stringified_event_loop}"
except RuntimeError: # handle no current running event loop
return key
token: Final = running_loop_token()
return f"{key}-{token}" if token is not None else key
def set_cache(self, key: str, value: object, litellm_owned_client: bool = False, **kwargs):
def _scoped_key(self, key: str, loop_bound: bool) -> str:
return self.update_cache_key_with_event_loop(key) if loop_bound else key
def set_cache(self, key: str, value: object, litellm_owned_client: bool = False, loop_bound: bool = True, **kwargs):
"""``litellm_owned_client`` marks a client litellm built, so it may be closed once evicted."""
if litellm_owned_client:
self.evicted_client_closer.mark_owned(value)
key = self.update_cache_key_with_event_loop(key)
return super().set_cache(key, value, **kwargs)
return super().set_cache(self._scoped_key(key, loop_bound), value, **kwargs)
async def async_set_cache(self, key: str, value: object, litellm_owned_client: bool = False, **kwargs):
if litellm_owned_client:
@ -68,13 +65,13 @@ class LLMClientCache(InMemoryCache):
key = self.update_cache_key_with_event_loop(key)
return await super().async_set_cache(key, value, **kwargs)
def get_cache(self, key, **kwargs):
key = self.update_cache_key_with_event_loop(key)
def get_cache(self, key: str, loop_bound: bool = True, **kwargs):
scoped_key: Final = self._scoped_key(key, loop_bound)
self.evicted_client_closer.reap()
return super().get_cache(key, **kwargs)
return super().get_cache(scoped_key, **kwargs)
async def async_get_cache(self, key, **kwargs):
async def async_get_cache(self, key: str, **kwargs):
key = self.update_cache_key_with_event_loop(key)
return await super().async_get_cache(key, **kwargs)

View file

@ -139,6 +139,7 @@ class AuthenticationError(openai.AuthenticationError):
litellm_debug_info: str | None = None,
max_retries: int | None = None,
num_retries: int | None = None,
body: Mapping[str, object] | None = None,
):
self.status_code = 401
self.message = f"litellm.AuthenticationError: {message}"
@ -152,7 +153,7 @@ class AuthenticationError(openai.AuthenticationError):
request=httpx.Request(method="GET", url="https://litellm.ai"), # mock request object
)
super().__init__(
self.message, response=self.response, body=None
self.message, response=self.response, body=body
) # Call the base class constructor with the parameters it needs
def __str__(self):
@ -189,6 +190,7 @@ class NotFoundError(openai.NotFoundError):
litellm_debug_info: str | None = None,
max_retries: int | None = None,
num_retries: int | None = None,
body: Mapping[str, object] | None = None,
):
self.status_code = 404
self.message = f"litellm.NotFoundError: {message}"
@ -202,7 +204,7 @@ class NotFoundError(openai.NotFoundError):
request=httpx.Request(method="GET", url="https://litellm.ai"), # mock request object
)
super().__init__(
self.message, response=self.response, body=None
self.message, response=self.response, body=body
) # Call the base class constructor with the parameters it needs
def __str__(self):
@ -321,6 +323,17 @@ class UnsupportedEndpointError(BadRequestError):
pass
class EndpointNotSupportedByModelError(BadRequestError):
def __init__(self, message: str, model: str, llm_provider: str = "") -> None:
super().__init__( # pyright: ignore[reportUnknownMemberType] # inherited public exception constructor is untyped
message=message, model=model, llm_provider=llm_provider
)
self.message = message
self.type = "invalid_request_error"
self.param = "model"
self.code = "400"
class UnprocessableEntityError(openai.UnprocessableEntityError):
def __init__(
self,

View file

@ -68,6 +68,11 @@ from litellm.utils import (
)
_PROVIDER_ROUTING_HINTS: Final = TypeAdapter(ProviderRoutingHints)
_PROVIDERS_RESOLVING_THEIR_OWN_IMAGE_BASE: Final = (
LlmProviders.DASHSCOPE,
LlmProviders.QWENCLOUD,
LlmProviders.QWEN_AI_PLATFORM,
)
# Cache for ImageEditRequestUtils to avoid repeated __getattr__ calls
_ImageEditRequestUtils_cache: type["ImageEditRequestUtils"] | None = None
@ -271,6 +276,7 @@ def image_generation(
"""
try:
args: Final = locals()
requested_api_base: Final = api_base
aimg_generation: Final = kwargs.get("aimg_generation", False)
litellm_call_id: Final = kwargs.get("litellm_call_id", None)
logger_fn: Final = kwargs.get("logger_fn", None)
@ -468,7 +474,11 @@ def image_generation(
raise ValueError(f"image generation config is not supported for {custom_llm_provider}")
# Resolve api_base from litellm.api_base if not explicitly provided
_api_base: Final = api_base or litellm.api_base
_api_base: Final = (
requested_api_base
if custom_llm_provider in _PROVIDERS_RESOLVING_THEIR_OWN_IMAGE_BASE
else api_base or litellm.api_base
)
litellm_params_dict["api_base"] = _api_base
return llm_http_handler.image_generation_handler(

View file

@ -349,7 +349,7 @@ class SlackAlerting(CustomBatchLogger):
except Exception:
return 0
async def send_daily_reports(self, router) -> bool:
async def send_daily_reports(self, router: Router) -> bool:
"""
Send a daily report on:
- Top 5 deployments with most failed requests
@ -365,7 +365,7 @@ class SlackAlerting(CustomBatchLogger):
False -> if not sent
"""
ids: Final = router.get_model_ids()
ids: Final = router.registry().names(kind="ids")
# get keys
failed_request_keys: Final = [f"{id}:{SlackAlertingCacheKeys.failed_requests_key.value}" for id in ids]
@ -430,7 +430,7 @@ class SlackAlerting(CustomBatchLogger):
message += "\tNone\n"
for i in range(len(top_5_failed)):
key = failed_request_keys[top_5_failed[i]].split(":")[0]
_deployment = router.get_model_info(key)
_deployment = router.registry().deployment_dict(key)
if isinstance(_deployment, dict):
deployment_name = _deployment["litellm_params"].get("model", "")
else:
@ -452,7 +452,7 @@ class SlackAlerting(CustomBatchLogger):
message += "\tNone\n"
for i in range(len(top_5_slowest)):
key = latency_keys[top_5_slowest[i]].split(":")[0]
_deployment = router.get_model_info(key)
_deployment = router.registry().deployment_dict(key)
if _deployment is not None:
deployment_name = _deployment["litellm_params"].get("model", "")
else:
@ -802,7 +802,7 @@ class SlackAlerting(CustomBatchLogger):
if self.llm_router is None:
return
deployment: Final = self.llm_router.get_deployment(model_id=deployment_id)
deployment: Final = self.llm_router.registry().deployment(deployment_id)
if deployment is None:
return
@ -953,7 +953,7 @@ class SlackAlerting(CustomBatchLogger):
return
### EXTRACT MODEL DETAILS ###
deployment: Final = self.llm_router.get_deployment(model_id=deployment_id)
deployment: Final = self.llm_router.registry().deployment(deployment_id)
if deployment is None:
return

View file

@ -102,18 +102,14 @@ class AgentOps(OpenTelemetry):
"Connection": "keep-alive",
}
client: Final = get_httpx_client()
try:
response: Final = client.post(
url=auth_endpoint,
headers=headers,
json={"api_key": api_key},
timeout=10,
)
response: Final = get_httpx_client().post(
url=auth_endpoint,
headers=headers,
json={"api_key": api_key},
timeout=10,
)
if response.status_code != 200:
raise Exception(f"Failed to fetch auth token: {response.text}")
if response.status_code != 200:
raise Exception(f"Failed to fetch auth token: {response.text}")
return _DECODED_JSON.validate_python(response.json())
finally:
client.close()
return _DECODED_JSON.validate_python(response.json())

View file

@ -35,6 +35,8 @@ from litellm.types.utils import StandardCallbackDynamicParams
if TYPE_CHECKING:
from opentelemetry.trace import Span as _Span
from litellm.proxy._types import UserAPIKeyAuth
Span = _Span
else:
Span = Any
@ -126,6 +128,26 @@ def _starts_its_own_trace(span: object, kwargs: Mapping[str, object]) -> bool:
)
def _a_destination_records_the_call(kwargs: Mapping[str, object]) -> bool:
"""Whether a key's or team's ``langfuse_otel`` destination gets this call from the OTel v2 fan-out, which records
every call but one whose logging payload shows it never reached a provider."""
from litellm.integrations.otel.model.metadata import LLMCallEvent
if not _a_destination_carries_the_trace():
return False
call: Final = LLMCallEvent.from_dict(kwargs)
return call.payload is None or call.reached_a_provider
def _a_destination_carries_the_trace() -> bool:
"""Whether this request resolved a key's or team's ``langfuse_otel`` destination, whose fan-out sends it the whole
trace."""
from litellm.integrations.otel.model.config import ExporterOwner
from litellm.integrations.otel.plumbing.context import destination_backends
return ExporterOwner.LANGFUSE_OTEL in destination_backends()
class LangfuseOtelLogger(OpenTelemetry):
def __init__(self, config=None, *args, **kwargs):
# Prevent LangfuseOtelLogger from modifying global environment variables by constructing config manually
@ -541,6 +563,27 @@ class LangfuseOtelLogger(OpenTelemetry):
"""
return None
def _emit_once(
self,
kwargs: dict[str, object], # mutable-ok: the v1 dedupe marker is written into the request's kwargs
*scope: object,
) -> bool:
"""Whether to emit this span: once per request, and never for a call a Langfuse destination records."""
return not _a_destination_records_the_call(kwargs) and super()._emit_once(kwargs, *scope) # pyright: ignore[reportUnknownMemberType] # the v1 hook takes a bare dict
async def async_post_call_failure_hook(
self,
request_data: dict[str, object], # mutable-ok: the v1 hook contract
original_exception: Exception,
user_api_key_dict: "UserAPIKeyAuth",
traceback_str: str | None = None,
) -> None:
"""Record the proxy failure of a request no Langfuse destination carries."""
if not _a_destination_carries_the_trace():
await super().async_post_call_failure_hook( # pyright: ignore[reportUnknownMemberType] # the v1 hook takes a bare dict
request_data, original_exception, user_api_key_dict, traceback_str
)
async def async_service_success_hook(self, *args, **kwargs):
"""
Langfuse should not receive service success logs.

View file

@ -13,9 +13,9 @@ from opentelemetry.trace.status import Status, StatusCode
from litellm.integrations.otel.mappers import resolve_mappers
from litellm.integrations.otel.mappers.base import AttributeMapper, AttrValue, SpanData
from litellm.integrations.otel.mappers.langfuse import LANGFUSE_REQUEST_METADATA_ATTRIBUTES
from litellm.integrations.otel.mappers.langfuse import LANGFUSE_REQUEST_METADATA_ATTRIBUTES, LangfuseMapper
from litellm.integrations.otel.mappers.openinference import fit_indexed_messages
from litellm.integrations.otel.model.config import OpenTelemetryV2Config
from litellm.integrations.otel.model.config import ExporterOwner, OpenTelemetryV2Config
from litellm.integrations.otel.model.payloads import (
GuardrailSpanData,
LLMCallSpanData,
@ -34,6 +34,7 @@ from litellm.integrations.otel.model.spans import (
mcp_tool_call_span_name,
service_span_name,
)
from litellm.integrations.otel.plumbing.context import destination_backends
from litellm.integrations.otel.plumbing.events import GenAIEventRecorder
from litellm.integrations.otel.plumbing.providers import to_otel_span_kind
@ -168,6 +169,9 @@ class SpanEmitter:
self._mappers: list[AttributeMapper] = (
list(mappers) if mappers is not None else resolve_mappers(config.mapper_names)
)
self._langfuse: Final = (
None if any(isinstance(mapper, LangfuseMapper) for mapper in self._mappers) else LangfuseMapper()
)
# Bounded LRU (ordered by insertion / most-recent touch). Storing keys
# only — the value is unused — so it behaves like a capped set.
self._emitted: OrderedDict[tuple[str, SpanRole], None] = OrderedDict()
@ -295,17 +299,27 @@ class SpanEmitter:
)
else None
)
mapped: Final = MappingProxyType(
{key: value for mapper in self._mappers for key, value in mapper.map(data).items()}
configured: Final = {key: value for mapper in self._mappers for key, value in mapper.map(data).items()}
for_langfuse: Final = (
dict(
sorted(
self._langfuse.map(data).items(),
key=lambda item: item[0] not in LANGFUSE_REQUEST_METADATA_ATTRIBUTES,
)
)
if self._langfuse is not None and ExporterOwner.LANGFUSE_OTEL in destination_backends()
else _NO_ATTRIBUTES
)
mapped: Final = MappingProxyType({**configured, **for_langfuse})
stamped_later: Final = error_attributes(error) if error else _NO_ATTRIBUTES
reserved: Final = len(stamped_later.keys() - mapped.keys())
overwrites: Final = carried_keys(span) & mapped.keys()
# The Langfuse request metadata sheds first, so a span at its attribute limit loses nothing else to it.
# The Langfuse request metadata, and the Langfuse vocabulary only a destination asked for, shed first,
# so a span at its attribute limit loses nothing else to them.
fitted: Final = fit_indexed_messages(
MappingProxyType({key: value for key, value in mapped.items() if key not in overwrites}),
attribute_budget(span, reserved),
LANGFUSE_REQUEST_METADATA_ATTRIBUTES,
LANGFUSE_REQUEST_METADATA_ATTRIBUTES.union(for_langfuse),
)
for key, value in mapped.items():
if key in overwrites or key in fitted:

View file

@ -198,6 +198,7 @@ class OpenTelemetryV2(CustomLogger):
tracer_provider: TracerProvider | None = None,
logger_provider: LoggerProvider | None = None,
meter_provider: "MeterProvider | None" = None,
console_default: bool = False,
**kwargs: Unpack[_CustomLoggerOptions],
) -> None:
super().__init__(**kwargs)
@ -206,7 +207,7 @@ class OpenTelemetryV2(CustomLogger):
self._tracer_provider: TracerProvider = (
tracer_provider
if tracer_provider is not None
else build_tracer_provider(self.config, tenant_overrides=True)
else build_tracer_provider(self.config, tenant_overrides=True, console_default=console_default)
)
self.tracer: Tracer = get_tracer(self._tracer_provider, LITELLM_TRACER_NAME)
self._metrics_recorder = self._init_metrics(meter_provider)
@ -217,7 +218,9 @@ class OpenTelemetryV2(CustomLogger):
mappers=resolve_mappers(self.config.mapper_names),
event_recorder=self._init_events(logger_provider),
)
self._tenant_tracers = TenantTracerCache(self.config, callback_name, LITELLM_TRACER_NAME)
self._tenant_tracers = TenantTracerCache(
self.config, callback_name, LITELLM_TRACER_NAME, console_default=console_default
)
self._open_llm_calls: OrderedDict[str, _LLMCallSpan] = OrderedDict()
self._init_otel_logger_on_litellm_proxy()
@ -541,9 +544,7 @@ class OpenTelemetryV2(CustomLogger):
# gate rejection carries ``is_no_upstream_call`` and gets no span. A cache
# hit never hands off to a provider, so ``cache_hit`` substitutes for the
# upstream stamp, since the payload proves the call completed.
if carrier is None and (
call.is_no_upstream_call or call.payload is None or (not call.upstream_started and not call.cache_hit)
):
if carrier is None and not call.reached_a_provider:
return None
try:
return self._finish_carrier(carrier, call, start_time, end_time)
@ -1091,6 +1092,7 @@ def build_otel_v2_logger(
logger_provider: LoggerProvider | None = None,
meter_provider: "MeterProvider | None" = None,
settings: Mapping[str, object] = MappingProxyType({}),
console_default: bool = False,
) -> OpenTelemetryV2:
return _logger_class(config)(
config=config,
@ -1098,6 +1100,7 @@ def build_otel_v2_logger(
tracer_provider=tracer_provider,
logger_provider=logger_provider,
meter_provider=meter_provider,
console_default=console_default,
**settings,
)

View file

@ -103,6 +103,18 @@ class LangfuseMapper:
),
}
#: Every attribute this mapper can write on a model call.
LLM_CALL_ATTRIBUTES: Final[frozenset[str]] = frozenset(
{
*_LLM_CALL_ATTRS,
*_BLOB_ATTRS,
LANGFUSE_TRACE_NAME,
LANGFUSE_TRACE_USER_ID,
LANGFUSE_TRACE_SESSION_ID,
LANGFUSE_TRACE_TAGS,
}
)
def map(self, data: SpanData) -> AttributeMap:
match data:
case LLMCallSpanData():

View file

@ -121,6 +121,17 @@ class ExporterSpec(LiteLLMBaseModel):
)
def is_unconfigured_placeholder(spec: ExporterSpec) -> bool:
"""Whether ``spec`` is the one ``_normalize`` folds in when nothing was configured.
No field set is what says the operator asked for nothing: an exporter they did
configure survives, even ``OTEL_EXPORTER=console`` whose value matches the default,
and so does the gated spec ``credential_gated_exporters`` appends, which would
otherwise eat itself when one preset layers onto another.
"""
return not spec.model_fields_set
class _EnvWithoutBareExcludedServices(PydanticBaseSettingsSource):
def __init__(self, settings_cls: type[BaseSettings], env_settings: PydanticBaseSettingsSource) -> None:
super().__init__(settings_cls)

View file

@ -272,6 +272,11 @@ class LLMCallEvent:
trace: TraceControls
session_id: str | None
@property
def reached_a_provider(self) -> bool:
"""Whether the call was handed to a provider or answered from the response cache."""
return not self.is_no_upstream_call and self.payload is not None and (self.upstream_started or self.cache_hit)
@classmethod
def from_dict(cls, kwargs: Mapping[str, object]) -> LLMCallEvent:
raw_payload: Final = kwargs.get("standard_logging_object")

View file

@ -47,6 +47,7 @@ from litellm.integrations.otel.mappers.langfuse import (
LANGFUSE_OBSERVATION_OUTPUT,
LANGFUSE_REQUEST_METADATA_ATTRIBUTES,
LANGFUSE_TRACE_NAME,
LangfuseMapper,
)
from litellm.integrations.otel.model.config import (
CaptureMessageContent,
@ -55,6 +56,7 @@ from litellm.integrations.otel.model.config import (
OpenTelemetryV2Config,
capture_message_content_from_env,
captures_span_content,
is_unconfigured_placeholder,
)
from litellm.integrations.otel.model.semconv import (
DB,
@ -490,19 +492,25 @@ def _is_litellm_span(span: ReadableSpan) -> bool:
return scope is not None and scope.name == LITELLM_TRACER_NAME
def without_langfuse_request_metadata(span: ReadableSpan) -> ReadableSpan:
"""``span`` without the request metadata the Langfuse mapper writes for a Langfuse destination.
def without_destination_langfuse_attributes(span: ReadableSpan, langfuse_mapped: bool) -> ReadableSpan:
"""``span`` without the Langfuse attributes only a Langfuse destination asked for.
Only LiteLLM's own spans lose it: any other span keeps every key its own code set.
That is the request metadata the Langfuse mapper writes and, on a model call written by
a logger whose own mappers leave the Langfuse vocabulary out (``langfuse_mapped`` off),
every Langfuse attribute, since that logger wrote them for a key's or team's Langfuse
destination alone. Only LiteLLM's own spans lose them: any other span keeps every key
its own code set.
"""
attributes: Final = span.attributes or _NO_ATTRIBUTES
if not (_is_litellm_span(span) and any(key in LANGFUSE_REQUEST_METADATA_ATTRIBUTES for key in attributes)):
dropped: Final = (
LANGFUSE_REQUEST_METADATA_ATTRIBUTES
if langfuse_mapped or not is_llm_call_span(span)
else LangfuseMapper.LLM_CALL_ATTRIBUTES
)
if not (_is_litellm_span(span) and any(key in dropped for key in attributes)):
return span
return _CountedSpanView(
span,
MappingProxyType(
{key: value for key, value in attributes.items() if key not in LANGFUSE_REQUEST_METADATA_ATTRIBUTES}
),
span, MappingProxyType({key: value for key, value in attributes.items() if key not in dropped})
)
@ -658,6 +666,7 @@ def _for_destination(
span: ReadableSpan,
destination: "OtelDestination",
default_capture: str,
langfuse_mapped: bool,
) -> ReadableSpan:
"""The view of ``span`` a tenant destination receives.
@ -672,13 +681,16 @@ def _for_destination(
and the query string of the request URL, which can hold the same key. A
destination whose effective setting disables span content loses the prompt and
response content too, and a destination for any backend but Langfuse loses the
request metadata the Langfuse mapper writes on LiteLLM's own spans. The span itself
stays, so the tenant still gets the whole trace tree.
Langfuse attributes only a Langfuse destination asked for (see
:func:`without_destination_langfuse_attributes`). The span itself stays, so the
tenant still gets the whole trace tree.
"""
extra: Final = destination.resource_attributes
defaults: Final = destination.resource_defaults
source: Final = (
span if destination.callback_name == ExporterOwner.LANGFUSE_OTEL else without_langfuse_request_metadata(span)
span
if destination.callback_name == ExporterOwner.LANGFUSE_OTEL
else without_destination_langfuse_attributes(span, langfuse_mapped)
)
attributes: Final = source.attributes or _NO_ATTRIBUTES
database: Final = _is_database_span(attributes)
@ -728,8 +740,10 @@ class TenantFanOutSpanProcessor(SpanProcessor):
default_capture_by_backend: Mapping[str, str] = MappingProxyType({}),
pending_drains: int = _MAX_PENDING_DRAINS,
drain_pool: _DrainPool | None = None,
langfuse_mapped: bool = True,
) -> None:
self._operator_sinks: Final = operator_sinks
self._langfuse_mapped: Final = langfuse_mapped
self._excluded_db_systems: Final = excluded_db_systems
self._default_capture: Final = default_capture
self._default_capture_by_backend: Final = default_capture_by_backend
@ -822,7 +836,8 @@ class TenantFanOutSpanProcessor(SpanProcessor):
try:
processor.on_end(
_scoped(
_for_destination(span, destination, self._default_capture_for(destination)), destination.span_scope
_for_destination(span, destination, self._default_capture_for(destination), self._langfuse_mapped),
destination.span_scope,
)
)
except Exception as exc: # noqa: BLE001 # one destination's failure must not cost the others their span
@ -1198,17 +1213,19 @@ class _OverriddenBackendFilter(SpanProcessor):
return self._inner.force_flush(timeout_millis)
class _LangfuseRequestMetadataFilter(SpanProcessor):
"""Strip the Langfuse request metadata from LiteLLM's own spans before ``inner`` exports them."""
class _DestinationLangfuseAttributeFilter(SpanProcessor):
"""Strip the Langfuse attributes only a Langfuse destination asked for from LiteLLM's own spans before ``inner``
exports them."""
def __init__(self, inner: SpanProcessor) -> None:
def __init__(self, inner: SpanProcessor, langfuse_mapped: bool) -> None:
self._inner: Final = inner
self._langfuse_mapped: Final = langfuse_mapped
def on_start(self, span: SDKSpan, parent_context: Context | None = None) -> None:
self._inner.on_start(span, parent_context)
def on_end(self, span: ReadableSpan) -> None:
self._inner.on_end(without_langfuse_request_metadata(span))
self._inner.on_end(without_destination_langfuse_attributes(span, self._langfuse_mapped))
def shutdown(self) -> None:
self._inner.shutdown()
@ -1491,13 +1508,18 @@ def _spec_processor(
spec: ExporterSpec,
use_simple_processor: bool | None,
strip_message_content: bool,
strip_langfuse_request_metadata: bool,
strip_destination_langfuse_attributes: bool,
langfuse_mapped: bool,
) -> SpanProcessor:
exporting: Final = _processor_for(
_exporter_from_spec(spec),
(spec.use_simple_processor if spec.use_simple_processor is not None else use_simple_processor),
)
scoped: Final = _LangfuseRequestMetadataFilter(exporting) if strip_langfuse_request_metadata else exporting
scoped: Final = (
_DestinationLangfuseAttributeFilter(exporting, langfuse_mapped)
if strip_destination_langfuse_attributes
else exporting
)
return _MessageContentFilter(scoped) if strip_message_content else scoped
@ -1530,6 +1552,7 @@ def build_tracer_provider(
use_simple_processor: bool | None = None,
tenant_overrides: bool = False,
strip_message_content: bool = False,
console_default: bool = False,
) -> TracerProvider:
"""Build the shared :class:`TracerProvider`.
@ -1548,11 +1571,16 @@ def build_tracer_provider(
``config.langfuse_span_scope`` narrows the exporter owned by ``langfuse_otel``
alone; a collector or any other backend in the same config keeps the full tree.
In a config that writes the Langfuse vocabulary, every exporter but ``langfuse_otel``'s
goes without the request metadata the Langfuse mapper writes on LiteLLM's own spans.
Every exporter but ``langfuse_otel``'s goes without the Langfuse attributes only a
Langfuse destination asked for (see :func:`without_destination_langfuse_attributes`),
which any logger writes for a request with a Langfuse destination.
``strip_message_content`` removes prompt and response content before every
exporter, for a routed provider whose spans can carry content another
destination asked for.
The console spec ``config._normalize`` folds in when nothing was configured
prints only under ``console_default``, which the ``otel`` callback passes: a
vendor preset or the proxy's background global logger was never asked to print.
"""
provider: Final = TracerProvider(resource=build_resource(config))
provider.sampler = _RequestExclusionSampler(provider.sampler)
@ -1564,19 +1592,20 @@ def build_tracer_provider(
provider.add_span_processor(_processor_for(exporter, use_simple_processor))
return provider
langfuse_mapped: Final = _maps_langfuse(config)
# ``config._normalize`` guarantees at least one spec (it folds the top-level
# ``exporter``/``endpoint``/``headers`` fields in when ``exporters`` is empty).
writes_langfuse: Final = "langfuse" in config.mapper_names or any(
spec.owner is ExporterOwner.LANGFUSE_OTEL for spec in config.exporters
)
for spec in config.exporters:
if spec.requires_headers and not spec.headers:
continue
if is_unconfigured_placeholder(spec) and not console_default:
continue
processor = _spec_processor(
spec,
use_simple_processor,
strip_message_content,
strip_langfuse_request_metadata=writes_langfuse and spec.owner is not ExporterOwner.LANGFUSE_OTEL,
strip_destination_langfuse_attributes=spec.owner is not ExporterOwner.LANGFUSE_OTEL,
langfuse_mapped=langfuse_mapped,
)
owner = spec.owner.value if tenant_overrides and spec.owner is not None else None
scope = _operator_scope(config, spec)
@ -1613,8 +1642,9 @@ def attach_tenant_fan_out(
so exactly one fan-out lands. ``configs`` name the operator's own exporters, one
config per v2 logger since each keeps its own provider and still writes its
account, so an additive destination pointing at any of them is delivered once
rather than twice. ``excluded_db_systems`` only filters what the fan-out
delivers, never the operator's own exporters.
rather than twice. The first is the config of the logger that owns ``provider``,
whose model calls the fan-out carries. ``excluded_db_systems`` only filters what
the fan-out delivers, never the operator's own exporters.
"""
with _FAN_OUT_ATTACH_LOCK:
if any(isinstance(processor, TenantFanOutSpanProcessor) for processor in _attached_processors(provider)):
@ -1625,10 +1655,16 @@ def attach_tenant_fan_out(
excluded_db_systems=excluded_db_systems,
default_capture=_operator_capture_mode(configs),
default_capture_by_backend=MappingProxyType(dict(_owned_capture_modes(configs))),
langfuse_mapped=any(_maps_langfuse(config) for config in configs[:1]),
)
)
def _maps_langfuse(config: OpenTelemetryV2Config) -> bool:
"""Whether a logger built from ``config`` writes the Langfuse vocabulary on every model call."""
return "langfuse" in config.mapper_names
def _operator_capture_mode(configs: Sequence[OpenTelemetryV2Config]) -> str:
"""The mode of the operator's generic ``otel`` callback, so ``callback_settings.otel`` counts as the global default"""
return next(

View file

@ -163,10 +163,12 @@ class TenantTracerCache:
config: OpenTelemetryV2Config,
callback_name: str | None,
tracer_name: str,
console_default: bool = False,
) -> None:
self._config = config
self._callback_name = callback_name
self._tracer_name = tracer_name
self._console_default = console_default
# Guards the three mutable structures below: ``pre_call`` can run on
# thread-pool workers concurrently with the event loop, so cache
# updates, span counts, and retirement must be atomic.
@ -284,6 +286,7 @@ class TenantTracerCache:
built: Final = build_tracer_provider(
self._routed_config(credential_headers, project_headers, endpoint, service_name),
strip_message_content=not self._config.capture_span_content,
console_default=self._console_default,
)
self._providers[cache_key] = built
return built

View file

@ -111,7 +111,7 @@ _REQUIRED_HEADERS_BY_CALLBACK: Final[Mapping[str, frozenset[str]]] = MappingProx
_NO_ATTRS: Final[Mapping[str, str]] = MappingProxyType({})
def _arize_project() -> Mapping[str, str]:
def _arize_project(_params: StandardCallbackDynamicParams) -> Mapping[str, str]:
from litellm.integrations.arize.arize import ArizeLogger
project: Final = ArizeLogger.get_arize_config().project_name
@ -124,10 +124,27 @@ def _arize_default_project() -> Mapping[str, str]:
return MappingProxyType({"model_id": ARIZE_DEFAULT_PROJECT})
#: Resource attributes the operator chose for a backend, applied to every export the
#: team receives over the span's own resource: ``ARIZE_PROJECT_NAME`` for Arize.
_RESOURCE_ATTRIBUTES_BY_CALLBACK: Final[Mapping[str, Callable[[], Mapping[str, str]]]] = MappingProxyType(
{"arize": _arize_project}
def _stands_in_for_the_legacy_langfuse_logger() -> bool:
"""Whether a Langfuse destination replaces the older Langfuse logger's copy, which a proxy without Langfuse keys
of its own sends for a key's or team's ``langfuse_otel``."""
return not (os.environ.get("LANGFUSE_PUBLIC_KEY") and os.environ.get("LANGFUSE_SECRET_KEY"))
def _langfuse_environment(params: StandardCallbackDynamicParams) -> Mapping[str, str]:
"""The key's or team's ``langfuse_environment``, else ``LANGFUSE_TRACING_ENVIRONMENT``, where the older Langfuse
logger wrote it."""
from litellm.types.integrations.langfuse_otel import LangfuseSpanAttributes
if not _stands_in_for_the_legacy_langfuse_logger():
return _NO_ATTRS
environment: Final = params.get("langfuse_environment") or os.environ.get("LANGFUSE_TRACING_ENVIRONMENT")
return MappingProxyType({LangfuseSpanAttributes.LANGFUSE_ENVIRONMENT.value: environment} if environment else {})
#: Resource attributes a backend gets on every export the team receives, over the
#: span's own resource: ``ARIZE_PROJECT_NAME`` for Arize, the Langfuse environment.
_RESOURCE_ATTRIBUTES_BY_CALLBACK: Final[Mapping[str, Callable[[StandardCallbackDynamicParams], Mapping[str, str]]]] = (
MappingProxyType({"arize": _arize_project, "langfuse_otel": _langfuse_environment})
)
#: Resource attributes a backend needs on every export to accept it, filled only where
@ -137,12 +154,14 @@ _RESOURCE_DEFAULTS_BY_CALLBACK: Final[Mapping[str, Callable[[], Mapping[str, str
)
def _resource_attributes(callback_name: str, service_name: str | None) -> Mapping[str, str]:
def _resource_attributes(
callback_name: str, service_name: str | None, params: StandardCallbackDynamicParams
) -> Mapping[str, str]:
backend: Final = _RESOURCE_ATTRIBUTES_BY_CALLBACK.get(callback_name)
return MappingProxyType(
{
**({"service.name": service_name} if service_name else {}),
**(backend() if backend is not None else {}),
**(backend(params) if backend is not None else {}),
}
)
@ -214,7 +233,7 @@ def destination_for(
return OtelDestination(
endpoint=endpoint,
headers=MappingProxyType(dict(headers)),
resource_attributes=_resource_attributes(callback_name, service_name),
resource_attributes=_resource_attributes(callback_name, service_name, params),
resource_defaults=_resource_defaults(callback_name),
callback_name=callback_name,
protocol=protocol,

View file

@ -3,7 +3,7 @@
from collections.abc import Iterable
from typing import Final
from litellm.integrations.otel.model.config import ExporterOwner, ExporterSpec
from litellm.integrations.otel.model.config import ExporterOwner, ExporterSpec, is_unconfigured_placeholder
def ensure_mappers(mapper_names: Iterable[str], *names: str) -> list[str]:
@ -35,14 +35,3 @@ def credential_gated_exporters(
*(spec for spec in exporters if not is_unconfigured_placeholder(spec)),
ExporterSpec(owner=owner, requires_headers=True),
)
def is_unconfigured_placeholder(spec: "ExporterSpec") -> bool:
"""Whether ``spec`` is the one ``_normalize`` folds in when nothing was configured.
No field set is what says the operator asked for nothing: an exporter they did
configure survives, even ``OTEL_EXPORTER=console`` whose value matches the default,
and so does the gated spec this module appends, which would otherwise eat itself
when one preset layers onto another.
"""
return not spec.model_fields_set

View file

@ -468,15 +468,16 @@ def _bounded_requested_model_label(requested_model: object, router_originated: b
llm_router: Final = _get_proxy_llm_router()
if llm_router is None:
return requested_model if router_originated else UNRECOGNIZED_REQUESTED_MODEL_LABEL
if llm_router.is_recognized_model(requested_model):
registry: Final = llm_router.registry()
if registry.resolve(requested_model, kind="recognized"):
return requested_model
if requested_model in llm_router.team_public_model_names:
if requested_model in registry.names(kind="team_public"):
return requested_model
if llm_router.pattern_router.route(requested_model) is not None:
if registry.patterns().route(requested_model) is not None:
return requested_model
if any(
team_pattern_router.route(requested_model) is not None
for team_pattern_router in llm_router.team_pattern_routers.values()
for team_pattern_router in registry.patterns(team=True).values()
):
return requested_model
return UNRECOGNIZED_REQUESTED_MODEL_LABEL
@ -3410,7 +3411,7 @@ class PrometheusLogger(CustomLogger):
OpenAI / Anthropic / Azure return remaining tokens/requests in response
headers, but Bedrock and Vertex AI do not. This fallback computes
``configured_limit - current_usage`` via
``Router.get_remaining_model_group_usage`` so the same gauges are
``Router.health().usage(model_group, kind="remaining")`` so the same gauges are
emitted for every provider when tpm/rpm is configured on the
deployment.
"""
@ -3437,7 +3438,7 @@ class PrometheusLogger(CustomLogger):
return
try:
remaining_usage: Final = await llm_router.get_remaining_model_group_usage(model_group)
remaining_usage: Final = await llm_router.health().usage(model_group, kind="remaining")
except Exception as e:
verbose_logger.exception(
"Prometheus: get_remaining_model_group_usage failed for model_group=%s: %s",

View file

@ -40,7 +40,6 @@ from litellm.litellm_core_utils.llm_judge import (
)
from litellm.litellm_core_utils.redact_messages import should_redact_message_logging
from litellm.llms.base_llm.base_utils import type_to_response_format_param
from litellm.router_utils.common_utils import resolve_model_group_alias
from litellm.types.llms.base import LiteLLMBaseModel
from litellm.types.management_endpoints.auto_router_endpoints import ShadowEvalDirection
from litellm.types.utils import SHADOW_EVAL_JUDGE_CALL_ORIGIN, SHADOW_EVAL_ROUTER_CALL_ORIGIN
@ -790,9 +789,7 @@ class ActiveShadowEvalJob(LiteLLMBaseModel):
def _canonical_group(router: "Router | None", model_group: str) -> str:
"""A model group in the one spelling both a job's scope and a request's model compare
under: an alias resolves to its target so the two never fail to match on spelling."""
return (
resolve_model_group_alias(router.model_group_alias, model_group) if router is not None else None
) or model_group
return (router.registry().resolve(model_group, kind="alias") if router is not None else None) or model_group
def _scope_admits(router: "Router | None", job: "ActiveShadowEvalJob", model_group: str) -> bool:

View file

@ -79,6 +79,7 @@ if TYPE_CHECKING:
from litellm.llms.base_llm.anthropic_messages.transformation import (
BaseAnthropicMessagesConfig,
)
from litellm.router import Router
from litellm.types.llms.anthropic_messages.anthropic_response import (
AnthropicMessagesResponse,
)
@ -132,7 +133,7 @@ class _SearchToolLitellmParams(TypedDict, total=False):
class _SearchToolConfig(TypedDict, total=False):
search_tool_name: str
search_tool_name: ReadOnly[str]
litellm_params: ReadOnly[_SearchToolLitellmParams | None]
@ -1768,10 +1769,13 @@ class WebSearchInterceptionLogger(CustomLogger):
return None
def _select_search_tool_from_router(self, llm_router: object) -> "_SearchToolConfig | None":
if llm_router is None or not hasattr(llm_router, "search_tools"):
def _select_search_tool_from_router(self, llm_router: "Router | None") -> "_SearchToolConfig | None":
if llm_router is None:
return None
search_tools: Final = tuple(getattr(llm_router, "search_tools", None) or ())
settings: Final = llm_router.settings()
if not settings.is_set("search_tools"):
return None
search_tools: Final = tuple(settings.search_tools or ())
return self._select_search_tool_from_list(search_tools=search_tools, source="router")
def _select_search_tool_from_list(

View file

@ -413,6 +413,25 @@ def _map_openai_exception(
body=getattr(original_exception, "body", None),
)
elif "invalid_request_error" in error_str and "Incorrect API key provided" not in error_str:
invalid_request_status: Final = getattr(original_exception, "status_code", None)
if invalid_request_status == 401:
raise AuthenticationError(
message=f"{exception_provider} - {message}",
llm_provider=custom_llm_provider,
model=model,
response=response,
litellm_debug_info=extra_information,
body=getattr(original_exception, "body", None),
)
if invalid_request_status == 404:
raise NotFoundError(
message=f"{exception_provider} - {message}",
model=model,
llm_provider=custom_llm_provider,
response=response,
litellm_debug_info=extra_information,
body=getattr(original_exception, "body", None),
)
raise BadRequestError(
message=f"{exception_provider} - {message}",
llm_provider=custom_llm_provider,

View file

@ -31,7 +31,13 @@ EXTRACTION_HEALTH_CHECK_FIELDS: Final[Mapping[str, Mapping[str, str]]] = Mapping
)
def native_health_check_mode(model: str, custom_llm_provider: str | None) -> Literal["anthropic_messages"] | None:
def native_health_check_mode(
model: str, custom_llm_provider: str | None
) -> Literal["anthropic_messages", "realtime"] | None:
if custom_llm_provider == "soniox":
from litellm.llms.soniox.common_utils import soniox_health_check_mode
return soniox_health_check_mode(model)
if custom_llm_provider != "bedrock_mantle":
return None
from litellm.llms.bedrock_mantle.common_utils import mantle_health_check_mode

View file

@ -1,6 +1,7 @@
# What is this?
## Common Utility file for Logging handler
# Logging function -> log the exact model details + what's being sent | Non-Blocking
import asyncio
import copy
import datetime
import functools
@ -96,6 +97,7 @@ from litellm.litellm_core_utils.internal_call_metadata import (
MODEL_ACCESS_GROUP_METADATA_KEY,
is_unbilled_non_inference_call,
)
from litellm.litellm_core_utils.llm_cost_calc.finalization import FinalizedCost, PreparedCost
from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import (
bedrock_guardrail_trace,
cost_breakdown_with_guardrail,
@ -813,6 +815,11 @@ class Logging(LiteLLMLoggingBaseClass):
# Initialize cost breakdown field
self.cost_breakdown: CostBreakdown | None = None
self.finalized_cost: FinalizedCost | None = None
self.prepared_cost: PreparedCost | None = None
self.cost_allowance_discount: float | None = None
self.cost_finalization_lock = asyncio.Lock()
self.deferred_success_result: tuple[object, datetime.datetime, datetime.datetime] | None = None
self.billed_token_rates: BilledTokenRates | None = None
self._provider_reported_cost: float | None = None
@ -874,6 +881,11 @@ class Logging(LiteLLMLoggingBaseClass):
"current_attempt_api_call_start_time": None,
"litellm_params": litellm_params,
}
new_call.finalized_cost = None
new_call.prepared_cost = None
new_call.cost_allowance_discount = None
new_call.cost_finalization_lock = asyncio.Lock()
new_call.deferred_success_result = None
new_call.streaming_chunks = []
new_call.sync_streaming_chunks = []
return new_call
@ -2462,6 +2474,55 @@ class Logging(LiteLLMLoggingBaseClass):
**kwargs,
)
def prepare_response_cost(self, response: object, now: datetime.datetime) -> PreparedCost | None:
if self.prepared_cost is not None:
return self.prepared_cost
recorded_cache_hit: Final[object] = self.model_call_details.get("cache_hit")
deferred: Final = self.deferred_success_result
prepared: Final = self.prepare_success_logging(
result=deferred[0] if deferred is not None else response,
start_time=deferred[1] if deferred is not None else self.start_time,
end_time=deferred[2] if deferred is not None else now.replace(tzinfo=None),
cache_hit=recorded_cache_hit if isinstance(recorded_cache_hit, bool) else None,
standard_logging_object=None,
build_logging_payload=True,
publish_logging_payload=False,
)
payload: Final = prepared[3]
raw_cost: Final[object] = self.model_call_details.get("response_cost")
if payload is None or not isinstance(raw_cost, (int, float)):
return None
recorded_usage: Final[object] = self.model_call_details.get("combined_usage_object")
response_usage: Final[object] = getattr(prepared[2], "usage", None)
normalized_usage: Final = (
recorded_usage
if isinstance(recorded_usage, Usage)
else response_usage
if isinstance(response_usage, Usage)
else None
)
usage: Final = StandardLoggingPayloadSetup.get_usage_from_response_obj(
_extract_response_obj_and_hidden_params(prepared[2], None)[0],
combined_usage_object=normalized_usage,
).model_copy(deep=True)
self.prepared_cost = PreparedCost(
raw_cost,
{
**payload,
"response": None,
"messages": None,
"model_map_information": copy.deepcopy(payload["model_map_information"]),
"cost_breakdown": copy.deepcopy(payload["cost_breakdown"]),
},
usage,
)
self.record_partial_usage_for_failure(usage, raw_cost)
return self.prepared_cost
def record_finalized_cost(self, finalized: FinalizedCost) -> None:
self.finalized_cost = finalized
self.model_call_details["response_cost"] = finalized.amount
def record_partial_usage_for_failure(self, usage: Usage, response_cost: float) -> None:
"""Stash what an interrupted stream already consumed so the failure log bills it instead of zero."""
self.model_call_details["combined_usage_object"] = usage
@ -2812,12 +2873,16 @@ class Logging(LiteLLMLoggingBaseClass):
logging_result, "_hidden_params", {}
)
if self.model_call_details.get("cache_hit") is True:
if self.finalized_cost is not None:
self.model_call_details["response_cost"] = self.finalized_cost.amount
elif self.model_call_details.get("cache_hit") is True:
self.model_call_details["response_cost"] = 0.0
elif "response_cost" in hidden_params:
self.model_call_details["response_cost"] = hidden_params["response_cost"]
self._record_zero_cost_diagnostic(logging_result, hidden_params["response_cost"])
elif (existing_cost := self.model_call_details.get("response_cost")) is not None and existing_cost != 0:
elif (existing_cost := self.model_call_details.get("response_cost")) is not None and (
existing_cost != 0 or self.call_type == CallTypes.pass_through.value # pyright: ignore[reportUnknownMemberType] # Logging.call_type is untyped
):
# Preserve response_cost if already calculated (e.g., by pass-through
# handlers like Gemini/Vertex which call completion_cost directly).
# Do not preserve 0 from failure_handler on intermediate router retries.
@ -2936,7 +3001,7 @@ class Logging(LiteLLMLoggingBaseClass):
standard_logging_object: StandardLoggingPayload | None = None,
build_logging_payload: bool = True,
) -> tuple[datetime.datetime, datetime.datetime, object]:
prepared: Final = self._prepare_success_logging(
prepared: Final = self.prepare_success_logging(
result=result,
start_time=start_time,
end_time=end_time,
@ -2947,7 +3012,7 @@ class Logging(LiteLLMLoggingBaseClass):
)
return prepared[0], prepared[1], prepared[2]
def _prepare_success_logging(
def prepare_success_logging(
self,
result: object,
start_time: datetime.datetime | None,
@ -2968,7 +3033,11 @@ class Logging(LiteLLMLoggingBaseClass):
self.model_call_details["log_event_type"] = "successful_api_call"
self.model_call_details["end_time"] = end_time
self.model_call_details["cache_hit"] = cache_hit
self.model_call_details["cache_hit"] = (
self.prepared_cost.payload["cache_hit"]
if cache_hit is None and self.prepared_cost is not None
else cache_hit
)
if self.call_type == CallTypes.anthropic_messages.value:
result = self._anthropic_messages_logged_response(result=result)
@ -3007,6 +3076,7 @@ class Logging(LiteLLMLoggingBaseClass):
if (
litellm.max_budget
and publish_logging_payload
and self.stream is False
and result is not None
and isinstance(result, dict)
@ -3037,10 +3107,17 @@ class Logging(LiteLLMLoggingBaseClass):
publish_logging_payload: bool,
) -> StandardLoggingPayload | None:
if standard_logging_object is not None:
effective_payload: Final = (
self.finalized_cost.logging_payload(standard_logging_object)
if self.finalized_cost is not None
else standard_logging_object
)
if self.finalized_cost is not None:
self.model_call_details["response_cost"] = self.finalized_cost.amount
if not publish_logging_payload:
return standard_logging_object.copy()
self.model_call_details["standard_logging_object"] = standard_logging_object
return standard_logging_object
return effective_payload.copy()
self.model_call_details["standard_logging_object"] = effective_payload
return effective_payload
if result is None or self.stream is True:
# Streaming reaches here before its cost is known, so the cost
# is seeded to None, but only when nothing has already
@ -3281,7 +3358,7 @@ class Logging(LiteLLMLoggingBaseClass):
verbose_logger.debug("Logging Details LiteLLM-Success Call: Cache_hit=%s", cache_hit)
if not self.should_run_logging(event_type="sync_success"): # prevent double logging
return
start_time, end_time, result, request_logging_payload = self._prepare_success_logging(
start_time, end_time, result, request_logging_payload = self.prepare_success_logging(
start_time=start_time,
end_time=end_time,
result=result,
@ -4204,7 +4281,9 @@ class Logging(LiteLLMLoggingBaseClass):
# chunks already delivered; the router stashes that recovered usage as
# ``combined_usage_object`` and pre-computes its cost, so preserve it
# here instead of zeroing the spend on an otherwise-failed request.
if self.model_call_details.get("combined_usage_object") is None:
if self.finalized_cost is not None:
self.model_call_details["response_cost"] = self.finalized_cost.amount
elif self.model_call_details.get("combined_usage_object") is None:
self.model_call_details["response_cost"] = 0
if hasattr(exception, "headers") and isinstance(exception.headers, dict):
@ -5544,7 +5623,7 @@ def build_custom_logger_compatible_class(
return callback
otel_settings: Final = _get_custom_logger_settings_from_proxy_server(callback_name=logging_integration)
otel_logger_v2: Final = build_otel_v2_logger(
config=OpenTelemetryV2Config(**otel_settings), settings=otel_settings
config=OpenTelemetryV2Config(**otel_settings), settings=otel_settings, console_default=True
)
_in_memory_loggers.append(otel_logger_v2)
_maybe_auto_initialize_arize_phoenix(_in_memory_loggers)
@ -7367,6 +7446,12 @@ def _cache_hit_cost_breakdown(served: CostBreakdown | None) -> CostBreakdown:
)
def _finalized_cost_of(logging_obj: object) -> FinalizedCost | None:
"""Duck-typed logging objects (and test doubles) carry no finalized cost."""
finalized: Final = getattr(logging_obj, "finalized_cost", None)
return finalized if isinstance(finalized, FinalizedCost) else None
def get_standard_logging_object_payload(
kwargs: dict | None,
init_response_obj: object,
@ -7470,7 +7555,10 @@ def get_standard_logging_object_payload(
raw_response_cost: Final = kwargs.get("response_cost")
llm_response_cost: Final[float] = raw_response_cost or 0.0
guardrail_cost: Final = guardrail_information_cost(metadata.get("standard_logging_guardrail_information"))
response_cost: Final[float] = llm_response_cost + guardrail_cost
finalized_cost: Final = _finalized_cost_of(logging_obj)
response_cost: Final[float] = (
finalized_cost.amount if finalized_cost is not None else llm_response_cost + guardrail_cost
)
# clean up litellm hidden params
clean_hidden_params: Final = StandardLoggingPayloadSetup.get_hidden_params(hidden_params)
@ -7537,13 +7625,17 @@ def get_standard_logging_object_payload(
):
model_name = response_model_name
request_cost_breakdown: Final = cost_breakdown_with_guardrail(
(
_cache_hit_cost_breakdown(uncached_price.cost_breakdown)
if uncached_price is not None
else logging_obj.cost_breakdown
),
guardrail_cost,
request_cost_breakdown: Final = (
finalized_cost.logging_breakdown()
if finalized_cost is not None
else cost_breakdown_with_guardrail(
(
_cache_hit_cost_breakdown(uncached_price.cost_breakdown)
if uncached_price is not None
else logging_obj.cost_breakdown
),
guardrail_cost,
)
)
captured_baseline: Final = logging_obj.baseline_observation
autorouter_savings: Final = (
@ -7600,6 +7692,7 @@ def get_standard_logging_object_payload(
model=model_name,
metadata=clean_metadata,
cache_key=clean_hidden_params["cache_key"],
cost_finalized=finalized_cost is not None,
response_cost=response_cost,
cost_breakdown=request_cost_breakdown,
autorouter_savings=autorouter_savings,

View file

@ -0,0 +1,61 @@
from collections.abc import Mapping
from dataclasses import dataclass
from types import MappingProxyType
from typing import Final, TypeAlias
from pydantic import TypeAdapter
from litellm.types.utils import CostBreakdown, StandardLoggingPayload, Usage
BreakdownValue: TypeAlias = float | str | None | Mapping[str, float]
_BREAKDOWN_FIELDS: Final = TypeAdapter(Mapping[str, BreakdownValue])
_COST_BREAKDOWN: Final = TypeAdapter(CostBreakdown)
@dataclass(frozen=True, slots=True)
class PreparedCost:
amount: float
payload: StandardLoggingPayload
usage: Usage
@dataclass(frozen=True, slots=True)
class FinalizedCost:
amount: float
breakdown: Mapping[str, BreakdownValue] | None
allowance_discount: float
def logging_payload(self, payload: StandardLoggingPayload) -> StandardLoggingPayload:
return {
**payload,
"response_cost": self.amount,
"cost_breakdown": self.logging_breakdown(),
"cost_finalized": True,
}
def logging_breakdown(self) -> CostBreakdown | None:
return _COST_BREAKDOWN.validate_python(self.breakdown) if self.breakdown is not None else None
def finalize_cost(amount: float, breakdown: CostBreakdown | None, allowance_discount: float = 0.0) -> FinalizedCost:
discount: Final = min(max(amount, 0.0), max(allowance_discount, 0.0))
total: Final = amount - discount
existing: Final = _BREAKDOWN_FIELDS.validate_python(breakdown) if breakdown is not None else None
adjusted: Final = (
{
**(existing or {}),
"total_cost": total,
"original_cost": (breakdown or {}).get("original_cost", amount),
"discount_amount": (breakdown or {}).get("discount_amount", 0.0) + discount,
}
if discount > 0
else existing
)
frozen: Final = (
MappingProxyType(
{key: MappingProxyType(value) if isinstance(value, Mapping) else value for key, value in adjusted.items()}
)
if adjusted is not None
else None
)
return FinalizedCost(amount=total, breakdown=frozen, allowance_discount=discount)

View file

@ -108,7 +108,7 @@ def judge_target(router: Router | None, model: str, team_id: str | None = None)
only for its own team and a team's own deployment resolves for nobody else, so asking
without it answers for a caller who does not exist.
"""
served: Final = router.resolved_litellm_models(model, team_id=team_id) if router is not None else ()
served: Final = router.registry().names(model, kind="litellm_models", team_id=team_id) if router is not None else ()
if served:
return JudgeTarget("router", frozenset(_provider_qualified(m) or m for m in served))
qualified: Final = _provider_qualified(model)
@ -148,7 +148,7 @@ def attached_auth(request_data: object) -> UserAPIKeyAuth | None:
def _named_route(router: Router, deployment_id: str) -> bool:
deployment: Final = router.get_deployment(model_id=deployment_id)
deployment: Final = router.registry().deployment(deployment_id)
if deployment is None:
return False
public_names: Final = (deployment.model_name, deployment.model_info.team_public_model_name or "")
@ -167,7 +167,7 @@ def caller_team_for(router: Router | None, model: str, auth: UserAPIKeyAuth | No
team_id: Final = auth.team_id if auth is not None else None
if team_id is None or router is None or judge_target(router, model).via != "nothing":
return None
team_route: Final = router.get_candidate_model_ids_for_route(model, team_id)
team_route: Final = router.registry().names(model, kind="route_candidates", team_id=team_id)
return team_id if team_route and all(_named_route(router, deployment_id) for deployment_id in team_route) else None

View file

@ -3,7 +3,7 @@ import json
import re
import time
import traceback
from collections.abc import Mapping, Sequence
from collections.abc import Generator, Iterator, Mapping, Sequence
from typing import Final, Literal, cast
from pydantic import ConfigDict, TypeAdapter, ValidationError, with_config
@ -19,7 +19,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import (
from litellm.llms.ollama.common_utils import OllamaResponseIdNormalizer
from litellm.types.llms.databricks import DatabricksTool
from litellm.types.llms.openai import (
ChatCompletionReasoningItem,
ChatCompletionAnnotation,
ChatCompletionThinkingBlock,
ImageURLListItem,
OpenAIModerationResponse,
@ -59,6 +59,31 @@ _CHOICES_FIELDS: Final[frozenset] = frozenset(Choices.model_fields.keys())
_MODEL_RESPONSE_FIELDS: Final[frozenset] = frozenset(ModelResponse.model_fields.keys()) | {"usage"}
_ANNOTATION_ENTRIES: Final = TypeAdapter(tuple[object, ...])
def _is_provider_specific_annotation(annotation: object) -> bool:
match annotation:
case {"type": "url_citation"}:
return False
case {"type": _}:
return True
case _:
return False
def _annotation_entries(annotations: object) -> tuple[object, ...]:
return _ANNOTATION_ENTRIES.validate_python(annotations) if isinstance(annotations, list) else ()
def partition_annotations(annotations: object) -> tuple[object, tuple[object, ...]]:
entries: Final = _annotation_entries(annotations)
others: Final = tuple(annotation for annotation in entries if _is_provider_specific_annotation(annotation))
if not others:
return annotations, ()
return [annotation for annotation in entries if not _is_provider_specific_annotation(annotation)] or None, others
def _normalize_images_for_message(
images: list[dict] | None,
) -> list[ImageURLListItem] | None:
@ -145,31 +170,89 @@ def _split_assembled_content_for_replay(content: str | None) -> list[str]:
return _REPLAY_CONTENT_SLICE_RE.findall(content)
def _clear_later_replay_slice_metadata(choice: StreamingChoices) -> None:
# Rebuild the delta as content-only so every accumulate-able field (role,
# tool_calls, reasoning_content, thinking_blocks, audio, images,
# annotations, ...) is dropped on later slices instead of an enumerated
# subset; repeating any of them makes downstream handlers that accumulate
# streamed deltas collect it once per slice, and a field added to Delta
# later can't silently re-introduce the duplication.
choice.delta = Delta(content=choice.delta.content)
choice.logprobs = None
if hasattr(choice, "enhancements"):
del choice.enhancements
_REPLAY_REASONING_FIELDS: Final = frozenset({"reasoning_content", "thinking_blocks", "reasoning_items"})
_REPLAY_LEAD_FIELDS: Final = _REPLAY_REASONING_FIELDS | {"role", "provider_specific_fields"}
_REPLAY_TOOL_FIELDS: Final = frozenset({"tool_calls", "function_call"})
_REPLAY_TEXT_FIELDS: Final = frozenset({"content", "annotations"})
def _first_replay_delta(
message: Mapping[str, object], reasoning_items: Sequence[ChatCompletionReasoningItem] | None
) -> Delta:
return Delta(
content=message.get("content"),
role=message["role"],
function_call=message.get("function_call"),
tool_calls=message.get("tool_calls"),
reasoning_items=None if reasoning_items is None else list(reasoning_items),
provider_specific_fields=_cached_provider_specific_fields(
_CACHED_MESSAGE_PROVIDER_FIELDS.validate_python(message)
),
def _fields_in(fields: Mapping[str, object], keys: frozenset[str]) -> dict[str, object]:
return {key: value for key, value in fields.items() if key in keys}
def _replay_deltas(delta: Delta) -> tuple[Delta, ...]:
fields: Final[dict[str, object]] = delta.model_dump(exclude_none=True)
lead: Final = _fields_in(fields, _REPLAY_LEAD_FIELDS) if _REPLAY_REASONING_FIELDS.intersection(fields) else {}
pieces: Final = _split_assembled_content_for_replay(delta.content) or ([delta.content] if delta.content else [])
if not pieces and not lead:
return (delta,)
if not pieces:
held_back: Final = lead.keys() | {"content"}
rest: Final = {key: value for key, value in fields.items() if key not in held_back}
return (Delta.model_validate(lead),) + ((Delta.model_validate(rest),) if rest else ())
tools: Final = _fields_in(fields, _REPLAY_TOOL_FIELDS)
annotations: Final = _fields_in(fields, frozenset({"annotations"}))
moved: Final = lead.keys() | _REPLAY_TOOL_FIELDS | _REPLAY_TEXT_FIELDS
first_slice: Final = {key: value for key, value in fields.items() if key not in moved}
last_index: Final = len(pieces) - 1
return (
((Delta.model_validate(lead),) if lead else ())
+ tuple(
Delta.model_validate(
{
**(first_slice if index == 0 else {}),
"content": piece,
**(annotations if index == last_index else {}),
}
)
for index, piece in enumerate(pieces)
)
+ ((Delta.model_validate(tools),) if tools else ())
)
def _replay_frame(
response: ModelResponseStream, delta: Delta, finish_reason: str | None, keeps_choice_metadata: bool
) -> ModelResponseStream:
frame: Final = response.model_copy(deep=True)
choice: Final = frame.choices[0]
choice.delta = delta
choice.finish_reason = finish_reason
if not keeps_choice_metadata:
choice.logprobs = None
if hasattr(choice, "enhancements"):
del choice.enhancements
return frame
def _replay_frames(response: ModelResponseStream) -> Iterator[ModelResponseStream]:
if len(response.choices) != 1:
yield response
return
usage: Final = getattr(response, "usage", None)
if usage is not None:
delattr(response, "usage")
choice: Final = response.choices[0]
deltas: Final = _replay_deltas(choice.delta)
first_text: Final = next((index for index, delta in enumerate(deltas) if delta.content), 0)
for index, delta in enumerate(deltas):
yield _replay_frame(response, delta, finish_reason=None, keeps_choice_metadata=index == first_text)
finish_frame: Final = _replay_frame(
response, Delta(), finish_reason=choice.finish_reason, keeps_choice_metadata=False
)
if usage is not None:
setattr(finish_frame, "usage", usage)
yield finish_frame
def _first_replay_delta(message: Mapping[str, object]) -> Delta:
return Delta.model_validate(
{
**message,
"provider_specific_fields": _cached_provider_specific_fields(
_CACHED_MESSAGE_PROVIDER_FIELDS.validate_python(message)
),
}
)
@ -234,7 +317,7 @@ async def convert_to_streaming_response_async(
t["index"] = index
pydantic_tool_calls.append(ChatCompletionDeltaToolCall(**t))
choice["message"]["tool_calls"] = pydantic_tool_calls
delta = _first_replay_delta(choice["message"], choice["message"].get("reasoning_items"))
delta = _first_replay_delta(choice["message"])
finish_reason = choice.get("finish_reason", None)
if finish_reason is None:
@ -266,41 +349,14 @@ async def convert_to_streaming_response_async(
if "model" in response_object:
model_response_object.model = response_object["model"]
# Replay cached content with per-word cadence so stream=true cache hits
# don't arrive as a single SSE frame. Multi-choice (n>1) responses and
# unsplittable content (None/empty/whitespace-free) keep the original
# single-yield behavior.
slices: list[str] = []
if len(model_response_object.choices) == 1:
slices = _split_assembled_content_for_replay(model_response_object.choices[0].delta.content)
if len(slices) <= 1:
yield model_response_object
await asyncio.sleep(0)
return
# Detach usage from the base object so we can re-attach it only to the
# final slice chunk. A non-None usage always lives in __pydantic_extra__
# here (set via setattr above), so delattr cannot fail.
original_usage: Final = getattr(model_response_object, "usage", None)
if original_usage is not None:
delattr(model_response_object, "usage")
original_finish_reason: Final = model_response_object.choices[0].finish_reason
last_idx: Final = len(slices) - 1
for i, piece in enumerate(slices):
slice_chunk = model_response_object.model_copy(deep=True)
slice_chunk.choices[0].delta.content = piece
if i > 0:
_clear_later_replay_slice_metadata(slice_chunk.choices[0])
slice_chunk.choices[0].finish_reason = original_finish_reason if i == last_idx else None
if i == last_idx and original_usage is not None:
setattr(slice_chunk, "usage", original_usage)
yield slice_chunk
for frame in _replay_frames(model_response_object):
yield frame
await asyncio.sleep(0)
def convert_to_streaming_response(
response_object: dict | None = None,
):
) -> Generator[ModelResponseStream, None, None]:
# used for yielding Cache hits when stream == True
if response_object is None:
raise Exception("Error in response object format")
@ -356,30 +412,7 @@ def convert_to_streaming_response(
if "model" in response_object:
model_response_object.model = response_object["model"]
# Replay cached content with per-word cadence on sync cache-hit paths
# (S3Cache, sync completion()). See convert_to_streaming_response_async
# for the full rationale — this mirrors its tail.
slices: list[str] = []
if len(model_response_object.choices) == 1:
slices = _split_assembled_content_for_replay(model_response_object.choices[0].delta.content)
if len(slices) <= 1:
yield model_response_object
return
original_usage: Final = getattr(model_response_object, "usage", None)
if original_usage is not None:
delattr(model_response_object, "usage")
original_finish_reason: Final = model_response_object.choices[0].finish_reason
last_idx: Final = len(slices) - 1
for i, piece in enumerate(slices):
slice_chunk = model_response_object.model_copy(deep=True)
slice_chunk.choices[0].delta.content = piece
if i > 0:
_clear_later_replay_slice_metadata(slice_chunk.choices[0])
slice_chunk.choices[0].finish_reason = original_finish_reason if i == last_idx else None
if i == last_idx and original_usage is not None:
setattr(slice_chunk, "usage", original_usage)
yield slice_chunk
yield from _replay_frames(model_response_object)
from collections import defaultdict
@ -787,6 +820,9 @@ def convert_to_model_response_object(
provider_specific_fields = dict(choice["message"].get("provider_specific_fields", None) or {})
for f in choice["message"].keys() - _MESSAGE_FIELDS:
provider_specific_fields[f] = choice["message"][f]
annotations, other_annotations = partition_annotations(choice["message"].get("annotations", None))
if other_annotations:
provider_specific_fields["annotations"] = list(other_annotations)
# Handle reasoning models that display `reasoning_content` within `content`
reasoning_content, content = extract_reasoning_content(choice["message"])
@ -809,7 +845,9 @@ def convert_to_model_response_object(
reasoning_content=reasoning_content,
thinking_blocks=thinking_blocks,
reasoning_items=choice["message"].get("reasoning_items", None),
annotations=choice["message"].get("annotations", None),
annotations=cast( # cast-ok: provider annotation dicts pass through whole, keeping keys such as url_citation.content
list[ChatCompletionAnnotation] | None, annotations
),
images=_normalize_images_for_message(choice["message"].get("images", None)),
)
finish_reason = choice.get("finish_reason", None)

View file

@ -6,7 +6,7 @@ import atexit
import contextvars
import inspect
import logging
from collections.abc import Coroutine, Iterator
from collections.abc import Awaitable, Callable, Coroutine, Iterator
from typing import Final
from typing_extensions import TypedDict
@ -66,14 +66,19 @@ class LoggingWorker:
max_queue_size: int = LOGGING_WORKER_MAX_QUEUE_SIZE,
concurrency: int = LOGGING_WORKER_CONCURRENCY,
timeout_summary_window: float = LOGGING_WORKER_TIMEOUT_SUMMARY_WINDOW_SECONDS,
*,
retry_sleep: Callable[[float], Awaitable[None]] = asyncio.sleep,
):
self.timeout = timeout
self.max_queue_size = max_queue_size
self.concurrency = concurrency
self.timeout_summary_window = timeout_summary_window
self._queue: asyncio.Queue[LoggingTask] | None = None
self._worker_task: asyncio.Task | None = None
self._running_tasks: set[asyncio.Task] = set()
self._worker_task: asyncio.Task[None] | None = None
self._running_tasks: set[asyncio.Task[None]] = set()
self._overflow_tasks: set[asyncio.Task[None]] = set() # mutable-ok: done callbacks discard refs
self._pending_enqueues: dict[int, LoggingTask] = {} # mutable-ok: retries and shutdown share callback ownership
self._retry_sleep = retry_sleep
self._dequeued_tasks: dict[int, LoggingTask] = {} # mutable-ok: refs so flush can rescue never-started tasks
self._sem: asyncio.Semaphore | None = None
self._bound_loop: asyncio.AbstractEventLoop | None = None
@ -168,6 +173,7 @@ class LoggingWorker:
self._sem = None
self._worker_task = None
self._running_tasks.clear()
self._overflow_tasks.clear()
self._timeout_summary_task = None
self._timeout_burst_count = 0
self._queue = new_queue
@ -185,6 +191,10 @@ class LoggingWorker:
self._sem = asyncio.Semaphore(self.concurrency)
if self._worker_task is None or self._worker_task.done():
self._worker_task = asyncio.create_task(self._worker_loop())
if not self._overflow_tasks:
self._aggressive_clear_in_progress = False
for pending in tuple(self._pending_enqueues.values()):
self._schedule_delayed_enqueue_retry(pending)
async def _process_log_task(
self, task: LoggingTask, sem: asyncio.Semaphore, queue: "asyncio.Queue[LoggingTask]"
@ -278,12 +288,24 @@ class LoggingWorker:
# Capture the current context when enqueueing
task: Final = LoggingTask(coroutine=coroutine, context=contextvars.copy_context())
self._enqueue_task(task)
def _enqueue_task(self, task: LoggingTask) -> None:
if self._queue is None:
return
try:
self._queue.put_nowait(task)
except asyncio.QueueFull:
# Queue is full - handle it appropriately
verbose_logger.exception("LoggingWorker queue is full")
self._pending_enqueues[id(task)] = task
self._handle_queue_full(task)
else:
self._pending_enqueues.pop(id(task), None)
def _start_overflow_task(self, coroutine: Coroutine[object, object, None]) -> None:
task: Final = asyncio.create_task(coroutine)
self._overflow_tasks.add(task)
task.add_done_callback(self._overflow_tasks.discard)
def _should_start_aggressive_clear(self) -> bool:
"""
@ -326,7 +348,7 @@ class LoggingWorker:
if self._should_start_aggressive_clear():
self._mark_aggressive_clear_started()
# Schedule clearing as async task so enqueue returns immediately (non-blocking)
asyncio.create_task(self._aggressively_clear_queue_async(task))
self._start_overflow_task(self._aggressively_clear_queue_async(task))
else:
# Cooldown active or clear in progress, schedule a delayed retry
self._schedule_delayed_enqueue_retry(task)
@ -363,7 +385,7 @@ class LoggingWorker:
delay: Final = self._calculate_retry_delay()
# Schedule the retry as a background task
asyncio.create_task(self._retry_enqueue_task(task, delay))
self._start_overflow_task(self._retry_enqueue_task(task, delay))
except RuntimeError:
# No event loop, drop the task as we can't schedule a retry
pass
@ -373,17 +395,8 @@ class LoggingWorker:
Retry enqueueing the task after delay, preserving original context.
This is called as a background task from _schedule_delayed_enqueue_retry.
"""
await asyncio.sleep(delay)
# Try to enqueue the task directly, preserving its original context
if self._queue is None:
return
try:
self._queue.put_nowait(task)
except asyncio.QueueFull:
# Still full - handle it appropriately (clear or retry again)
self._handle_queue_full(task)
await self._retry_sleep(delay)
self._enqueue_task(task)
def _extract_tasks_from_queue(self) -> list[LoggingTask]:
"""
@ -423,15 +436,10 @@ class LoggingWorker:
return
extracted_tasks: Final = self._extract_tasks_from_queue()
# Add new task to extracted tasks to process directly
if new_task is not None:
self._track_dequeued(new_task)
extracted_tasks.append(new_task)
# Process extracted tasks directly
self._enqueue_task(new_task)
if extracted_tasks:
await self._process_extracted_tasks(extracted_tasks)
await asyncio.gather(*(self._process_single_task(task) for task in extracted_tasks))
except Exception as e:
verbose_logger.exception("LoggingWorker error during aggressive clear: %s", e)
finally:
@ -439,7 +447,7 @@ class LoggingWorker:
self._aggressive_clear_in_progress = False
async def _process_single_task(self, task: LoggingTask) -> None:
"""Process a single task and mark it done."""
"""Process a task extracted from the queue and mark it done."""
if self._queue is None:
return
@ -455,17 +463,6 @@ class LoggingWorker:
self._untrack_dequeued(task)
self._queue.task_done()
async def _process_extracted_tasks(self, tasks: list[LoggingTask]) -> None:
"""
Process tasks that were extracted from the queue to make room.
Processes them concurrently without semaphore limits for maximum speed.
"""
if not tasks or self._queue is None:
return
# Process all tasks concurrently for maximum speed
await asyncio.gather(*[self._process_single_task(task) for task in tasks])
def ensure_initialized_and_enqueue(self, async_coroutine: Coroutine[object, None, None]) -> None:
"""
Ensure the logging worker is initialized and enqueue the coroutine.
@ -473,32 +470,63 @@ class LoggingWorker:
self.start()
self.enqueue(async_coroutine)
async def stop(self) -> None:
@property
def pending_count(self) -> int:
return (
(self._queue.qsize() if self._queue is not None else 0)
+ len(self._dequeued_tasks)
+ len(self._pending_enqueues)
)
async def stop(self, *, discard_pending: bool = False, timeout: float | None = None) -> None:
"""Stop the logging worker and clean up resources."""
if discard_pending and self._queue is not None:
discarded: Final = self._drain_pending(self._queue)
for task in discarded:
task["coroutine"].close()
self._queue.task_done()
if discarded:
verbose_logger.warning("LoggingWorker shutdown: %d queued callbacks discarded", len(discarded))
if self._timeout_summary_task is not None:
self._timeout_summary_task.cancel()
self._timeout_summary_task = None
self._emit_timeout_summary()
if self._worker_task is None and not self._running_tasks:
# No worker launched and no in-flight tasks to drain.
return
if discard_pending:
rejected: Final = tuple(self._pending_enqueues.values())
self._pending_enqueues.clear()
for task in rejected:
task["coroutine"].close()
if rejected:
verbose_logger.warning("LoggingWorker shutdown: %d overflow callbacks discarded", len(rejected))
tasks_to_cancel: Final[list[asyncio.Task]] = list(self._running_tasks)
if self._worker_task:
# Include the main worker loop so it stops fetching work.
tasks_to_cancel.append(self._worker_task)
tasks_to_cancel: Final = (
*self._running_tasks,
*self._overflow_tasks,
*((self._worker_task,) if self._worker_task is not None else ()),
)
if not tasks_to_cancel:
return
for task in tasks_to_cancel:
# Propagate cancellation to every pending task.
task.cancel()
# Wait for cancellation to settle; ignore errors raised during shutdown.
await asyncio.gather(*tasks_to_cancel, return_exceptions=True)
done, pending = await asyncio.wait(tasks_to_cancel, timeout=timeout)
await asyncio.gather(*done, return_exceptions=True)
if pending:
verbose_logger.warning("LoggingWorker shutdown: %d tasks still alive after cancellation", len(pending))
self._worker_task = None
if discard_pending and self._queue is not None:
for task in self._unstarted_dequeued_tasks():
task["coroutine"].close()
self._untrack_dequeued(task)
self._queue.task_done()
if self._worker_task is not None and self._worker_task.done():
self._worker_task = None
# Drop references to completed tasks so we can restart cleanly.
self._running_tasks.clear()
self._running_tasks.difference_update(done)
async def flush(self) -> None:
"""Flush the logging queue.
@ -516,6 +544,9 @@ class LoggingWorker:
return
self.start()
await self._queue.join()
while self._overflow_tasks:
await asyncio.wait(tuple(self._overflow_tasks))
await self._queue.join()
async def clear_queue(self):
"""

View file

@ -424,6 +424,109 @@ def anthropic_image_source_to_openai_url(image_source: Mapping[str, object]) ->
return None
@dataclass(frozen=True, slots=True)
class InlineDocument:
filename: str
file_data: str
def _document_media_type(source: object) -> str:
match source:
case {"media_type": str() as media_type} if media_type:
return media_type
case _:
return "application/pdf"
def anthropic_inline_document(document: Mapping[str, object]) -> InlineDocument | None:
source: Final = document.get("source")
match source:
case {"type": "base64", "data": str() as data} if data:
title: Final = document.get("title")
return InlineDocument(
filename=title if isinstance(title, str) and title else "document.pdf",
file_data=f"data:{_document_media_type(source)};base64,{data}",
)
case _:
return None
def is_pdf_file_part(part: object) -> bool:
match part:
case {"type": "file", "file": {"file_data": str() as file_data}} if file_data.startswith("data:"):
return file_data.startswith("data:application/pdf")
case {"type": "file", "file": {"format": str() as file_format}}:
return file_format == "application/pdf"
case {"type": "file", "file": {"filename": str() as filename}}:
return filename.lower().endswith(".pdf")
case _:
return False
_CONTENT_PARTS: Final = TypeAdapter(tuple[object, ...])
def content_parts(message: Mapping[str, object]) -> Sequence[object]:
content: Final = message.get("content")
return _CONTENT_PARTS.validate_python(content) if isinstance(content, list) else ()
def _has_pdf_file_part(message: Mapping[str, object]) -> bool:
return any(is_pdf_file_part(part) for part in content_parts(message))
def messages_have_pdf_file_part(messages: Sequence[Mapping[str, object]]) -> bool:
return any(_has_pdf_file_part(message) for message in messages)
def _message_without_pdf_file_parts(message: AllMessageValues) -> AllMessageValues:
if not _has_pdf_file_part(message):
return message
kept: Final = [part for part in content_parts(message) if not is_pdf_file_part(part)]
return cast(AllMessageValues, {**message, "content": kept}) # cast-ok: same message with fewer content parts
def _is_pdf_only_user_turn(message: AllMessageValues) -> bool:
parts: Final = content_parts(message)
return message["role"] == "user" and bool(parts) and all(is_pdf_file_part(part) for part in parts)
def messages_without_pdf_file_parts(messages: Sequence[AllMessageValues]) -> tuple[AllMessageValues, ...]:
return tuple(
_message_without_pdf_file_parts(message) for message in messages if not _is_pdf_only_user_turn(message)
)
_CONTENT_PART: Final = TypeAdapter(Mapping[str, object])
def _file_part_as_image_url(part: object) -> object:
match part:
case {"type": "file", "file": {"file_data": str() as file_data}} if file_data:
fields: Final = _CONTENT_PART.validate_python(part)
return {key: value for key, value in fields.items() if key != "file"} | {
"type": "image_url",
"image_url": {"url": file_data},
}
case _:
return part
def message_with_file_parts_as_image_urls(message: Mapping[str, object]) -> Mapping[str, object]:
parts: Final = content_parts(message)
if not parts:
return message
return {**message, "content": [_file_part_as_image_url(part) for part in parts]}
def _typed_message_with_file_parts_as_image_urls(message: AllMessageValues) -> AllMessageValues:
return cast(AllMessageValues, message_with_file_parts_as_image_urls(message)) # cast-ok: same message, image parts
def messages_with_file_parts_as_image_urls(messages: Sequence[AllMessageValues]) -> tuple[AllMessageValues, ...]:
return tuple(_typed_message_with_file_parts_as_image_urls(message) for message in messages)
def _image_part_url(part: Mapping[str, object]) -> str | None:
"""The image URL carried by one content part, whichever of the three dialects wrote it."""
part_type: Final = part.get("type")

View file

@ -20,6 +20,7 @@ from litellm.constants import (
REALTIME_SESSION_SUCCESS_LOGGED_KEY,
)
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
from litellm.llms.base_llm.realtime.transcription_protocol import RealtimeTranscriptionProtocolError
from litellm.llms.base_llm.realtime.transformation import BaseRealtimeConfig, RealtimeBackend
from litellm.types.llms.openai import (
OpenAIRealtimeEvents,
@ -1281,24 +1282,29 @@ class RealTimeStreaming:
verbose_logger.exception("Connection closed in backend to client send messages - %s", e)
close: Final = backend_close_from(e)
self._flush_unbilled_transcription_usage()
if self._backend_refused_session(close):
await self.log_backend_refusal(e)
else:
await self.log_messages()
await self._log_session_end(e if self._backend_refused_session(close) else None)
return close
except asyncio.CancelledError:
self._flush_unbilled_transcription_usage()
await self.log_messages()
await self._log_session_end(None)
raise
except Exception as e:
verbose_logger.exception("Error in backend to client send messages: %s", e)
self._flush_unbilled_transcription_usage()
await self.log_messages()
await self._log_session_end(None)
return BackendClose(code=1011, reason="proxy failed while relaying the upstream websocket")
def _backend_refused_session(self, close: BackendClose) -> bool:
return close.code != 1000 and not self._backend_sent_frames
async def _log_session_end(self, transport_refusal: Exception | None) -> None:
provider_refusal: Final = None if self.provider_config is None else self.provider_config.session_refusal()
refusal: Final = provider_refusal or transport_refusal
if refusal is None:
await self.log_messages()
return
await self.log_backend_refusal(refusal)
async def log_backend_refusal(self, error: Exception) -> None:
if not self.logging_obj:
return
@ -1691,7 +1697,13 @@ class RealTimeStreaming:
# ``_maybe_send_guardrail_turn_detection_update`` would retry.
if msg_type == "session.update" and self.provider_config is None:
await self._send_session_gate_once()
sent = await self._send_to_backend(message)
try:
sent = await self._send_to_backend(message)
except RealtimeTranscriptionProtocolError as e:
await self.websocket.send_text(realtime_error_event(str(e), error_type="invalid_request_error"))
if self._uses_deferred_backend_setup():
raise
continue
if guardrail_turn_detection_injected and sent:
self._guardrail_turn_detection_update_sent = True

View file

@ -4,12 +4,13 @@ post-call guardrail rewrote rather than the provider response the proxy assemble
from __future__ import annotations
from collections.abc import Sequence
from itertools import chain
from typing import Final, Literal
from pydantic import BaseModel, Field, TypeAdapter, ValidationError
from pydantic import BaseModel, ConfigDict, Field, TypeAdapter, ValidationError
from litellm.types.llms.base import LiteLLMBaseModel
from litellm.types.utils import ModelResponse, ModelResponseStream
from litellm.types.utils import ModelResponse, ModelResponseStream, TextCompletionResponse
SERVED_OUTPUT_TEXTS_KEY: Final = "served_output_texts"
@ -44,6 +45,19 @@ class _ChatChoices(LiteLLMBaseModel):
choices: list[object]
class _TextCompletionChoice(LiteLLMBaseModel):
model_config = ConfigDict(frozen=True)
index: int | None = None
text: str | None = None
class _TextCompletionChunk(LiteLLMBaseModel):
model_config = ConfigDict(frozen=True)
choices: tuple[_TextCompletionChoice, ...]
def _as_json_object(response: object) -> dict[str, object] | None:
candidate: Final = response.model_dump() if isinstance(response, BaseModel) else response
try:
@ -108,8 +122,13 @@ def served_output_texts(response: object) -> ServedTexts | None:
def served_stream_output_texts(chunks: Sequence[object]) -> ServedTexts | None:
if chunks and all(isinstance(chunk, ModelResponseStream) for chunk in chunks):
return _chat_stream_texts(tuple(chunk for chunk in chunks if isinstance(chunk, ModelResponseStream)))
chat_chunks: Final = tuple(chunk for chunk in chunks if isinstance(chunk, ModelResponseStream))
if chunks and len(chat_chunks) == len(chunks):
return _indexed_texts_by_choice(tuple(chain.from_iterable(_chat_chunk_texts(chunk) for chunk in chat_chunks)))
if chunks and all(isinstance(chunk, TextCompletionResponse) for chunk in chunks):
return _indexed_texts_by_choice(
tuple(chain.from_iterable(_text_completion_chunk_texts(chunk) for chunk in chunks))
)
from litellm.llms.anthropic.messages.stream_assembly import (
assemble_anthropic_sse_stream,
is_anthropic_sse_stream,
@ -121,19 +140,26 @@ def served_stream_output_texts(chunks: Sequence[object]) -> ServedTexts | None:
return _chat_texts(assembled) if assembled is not None else None
def _chat_stream_choice_text(chunks: Sequence[ModelResponseStream], index: int) -> str | None:
contents: Final = tuple(
content
for chunk in chunks
def _chat_chunk_texts(chunk: ModelResponseStream) -> tuple[tuple[int, str | None], ...]:
return tuple(
(choice.index or 0, content if isinstance(content := choice.delta.content, str) else None)
for choice in chunk.choices
if (choice.index or 0) == index and isinstance(content := choice.delta.content, str)
)
return "".join(contents) if contents else None
def _chat_stream_texts(chunks: Sequence[ModelResponseStream]) -> ServedTexts | None:
choice_count: Final = max(((choice.index or 0) + 1 for chunk in chunks for choice in chunk.choices), default=0)
texts: Final = tuple(_chat_stream_choice_text(chunks, index) for index in range(choice_count))
def _text_completion_chunk_texts(chunk: object) -> tuple[tuple[int, str | None], ...]:
parsed: Final = _TextCompletionChunk.model_validate(_as_json_object(chunk))
return tuple((choice.index or 0, choice.text) for choice in parsed.choices)
def _joined_choice_text(indexed_texts: Sequence[tuple[int, str | None]], index: int) -> str | None:
texts: Final = tuple(text for text_index, text in indexed_texts if text_index == index and text is not None)
return "".join(texts) if texts else None
def _indexed_texts_by_choice(indexed_texts: Sequence[tuple[int, str | None]]) -> ServedTexts | None:
choice_count: Final = max((index + 1 for index, _ in indexed_texts), default=0)
texts: Final = tuple(_joined_choice_text(indexed_texts, index) for index in range(choice_count))
return texts if any(text is not None for text in texts) else None
@ -170,6 +196,8 @@ def _choice_with_text(choice: object, text: str | None) -> object:
if choice_obj is None or text is None:
return choice
message: Final = _as_json_object(choice_obj.get("message"))
if message is None or message.get("content") == text:
return choice
return {**choice_obj, "message": {**message, "content": text}}
if message is not None:
return choice if message.get("content") == text else {**choice_obj, "message": {**message, "content": text}}
if isinstance(choice_obj.get("text"), str):
return {**choice_obj, "text": text}
return choice

View file

@ -0,0 +1,27 @@
from collections.abc import Iterable
from typing import Final
from litellm.types.utils import ModelResponseStream, StreamingChoices
def is_empty_streaming_choice(choice: StreamingChoices) -> bool:
if choice.finish_reason is not None:
return False
if getattr(choice, "logprobs", None) is not None:
return False
delta: Final = choice.delta if hasattr(choice, "delta") else None
if delta is None:
return True
return _all_values_none(delta.model_dump().values())
def _all_values_none(values: Iterable[object]) -> bool:
return all(value is None for value in values)
def is_injected_stream_usage_artifact(chunk: object) -> bool:
if not isinstance(chunk, ModelResponseStream):
return False
if chunk.provider_specific_fields is not None:
return False
return all(is_empty_streaming_choice(choice) for choice in chunk.choices or [])

View file

@ -19,6 +19,7 @@ from litellm.litellm_core_utils.core_helpers import (
)
from litellm.litellm_core_utils.hidden_params import HIDDEN_PARAMS_ATTR, set_hidden_params
from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import BEDROCK_TRACE_FIELDS, object_items, object_mapping
from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import partition_annotations
from litellm.litellm_core_utils.prompt_templates.common_utils import (
get_content_from_model_response, # pyright: ignore[reportUnknownVariableType] # response param accepts an untyped dict
)
@ -1700,10 +1701,17 @@ def stream_chunk_builder(
# Merge annotations from ALL chunks — providers may spread
# them across multiple streaming chunks or send them only in
# the final chunk.
all_annotations: Final[list] = []
all_annotations: Final[list[object]] = []
for ac in annotation_chunks:
all_annotations.extend(ac["choices"][0]["delta"]["annotations"])
response["choices"][0]["message"]["annotations"] = all_annotations
message_annotations, provider_annotations = partition_annotations(all_annotations)
if message_annotations is not None:
response["choices"][0]["message"]["annotations"] = message_annotations
if provider_annotations:
response.choices[0].message.provider_specific_fields = {
**(response.choices[0].message.provider_specific_fields or {}),
"annotations": list(provider_annotations),
}
audio_chunks: Final = [
chunk
@ -1766,7 +1774,10 @@ def stream_chunk_builder(
if combined_provider_fields:
_choice = response.choices[0]
_choice.message.provider_specific_fields = combined_provider_fields
_choice.message.provider_specific_fields = {
**(_choice.message.provider_specific_fields or {}),
**combined_provider_fields,
}
completion_output = get_content_from_model_response(response)

View file

@ -6,6 +6,7 @@ import copy
import json
import re
from collections.abc import Container, Mapping, MutableMapping, Sequence
from dataclasses import dataclass
from datetime import datetime, timezone
from types import MappingProxyType
from typing import Annotated, Any, ClassVar, Final, Literal, TypeVar
@ -1774,6 +1775,48 @@ def anthropic_request_errors_on_thinking_binding_mismatch(request_body: object)
return binding is not None and binding.get("prefix_mismatch_behavior") == "error"
_CLAUDE_STRICT_TOOLS_REJECTION: Final = re.compile(
r"Too many strict tools \(\d+\)"
r"|Schemas contains too many (?:optional parameters|parameters with union types)"
r"|[Cc]ompiled grammar (?:is too large|size)"
r"|tools\.\d+\.custom(?:\.strict)?: "
r"|Grammar compilation timed out"
)
_CLAUDE_TOOLS_ADAPTER: Final = TypeAdapter(tuple[Mapping[str, object], ...])
def is_claude_strict_tools_rejection(status_code: int, error_text: str) -> bool:
return status_code == 400 and _CLAUDE_STRICT_TOOLS_REJECTION.search(error_text) is not None
@dataclass(frozen=True, slots=True)
class ToolsWithoutStrict:
tools: tuple[Mapping[str, object], ...]
dropped_count: int
def tools_without_strict(tools: object) -> ToolsWithoutStrict | None:
try:
validated: Final = _CLAUDE_TOOLS_ADAPTER.validate_python(tools)
except ValidationError:
return None
dropped_count: Final = sum(tool.get("strict") is True for tool in validated)
if dropped_count == 0:
return None
return ToolsWithoutStrict(
tools=tuple({key: value for key, value in tool.items() if key != "strict"} for tool in validated),
dropped_count=dropped_count,
)
def warn_claude_strict_tools_dropped(dropped_count: int) -> None:
litellm.verbose_logger.warning(
"Claude rejected the request's strict schemas, so litellm removed `strict` from every tool (%d had "
"strict true) and sent the request again because drop_params is set.",
dropped_count,
)
def strip_thinking_blocks_from_anthropic_messages(messages: Sequence[object]) -> list[object]:
"""
Return a new message list with thinking / redacted_thinking content blocks removed

View file

@ -1,6 +1,7 @@
# What is this?
## Translates OpenAI call to Anthropic `/v1/messages` format
import copy
import io
import itertools
import json
import traceback
@ -181,6 +182,44 @@ class _ToolCallPart:
)
@dataclass(frozen=True, slots=True)
class _ThinkingPiece:
text: str
signature: str
index: int | None
@staticmethod
def of(chunk: "ModelResponseStream") -> "_ThinkingPiece | None":
choices: Final = _optional_attr_sequence(chunk, "choices")
if len(choices) != 1 or _optional_attr(choices[0], "finish_reason") is not None:
return None
delta: Final = _optional_attr(choices[0], "delta")
if _optional_attr(delta, "content") or _optional_attr(delta, "tool_calls"):
return None
blocks: Final = _optional_attr_sequence(delta, "thinking_blocks")
if not blocks:
reasoning: Final = _optional_str_attr(delta, "reasoning_content")
return _ThinkingPiece(text=reasoning, signature="", index=None) if reasoning else None
block: Final = as_str_mapping(blocks[0]) if len(blocks) == 1 else None
if block is None or block.get("type") != "thinking":
return None
text: Final = block.get("thinking")
signature: Final = block.get("signature")
index: Final = block.get("index")
return _ThinkingPiece(
text=text if isinstance(text, str) else "",
signature=signature if isinstance(signature, str) else "",
index=index if isinstance(index, int) else None,
)
def block(self, *, text: str, signature: str) -> ChatCompletionThinkingBlock:
if self.index is None:
return ChatCompletionThinkingBlock(type="thinking", thinking=text, signature=signature or None)
return ChatCompletionThinkingBlock(
type="thinking", thinking=text, signature=signature or None, index=self.index
)
class _CombinedChunkSplitter:
"""
Splits a streaming chunk that carries BOTH response content and a
@ -209,6 +248,10 @@ class _CombinedChunkSplitter:
self._buffer: deque[ModelResponseStream] = deque()
self._held_calls: Mapping[int, tuple[ModelResponseStream, _ToolCallPart]] = MappingProxyType({})
self._live_call_index: int | None = None
self._held_signature: ModelResponseStream | None = None
self._unsigned_thinking = io.StringIO()
self._unsigned_thinking_length = 0
self._thinking_index: int | None = None
self._closed = False
async def aclose(self) -> None:
@ -458,7 +501,45 @@ class _CombinedChunkSplitter:
for split_chunk in self._split_by_payload_kind(combined_chunk) # comprehension-ok: flattens nested splits
)
block_chunks: Final = tuple(itertools.chain.from_iterable(map(self._split_thinking_blocks, split_chunks)))
return tuple(itertools.chain.from_iterable(map(self._route_tool_calls, block_chunks)))
signed_chunks: Final = tuple(itertools.chain.from_iterable(map(self._defer_signature, block_chunks)))
return tuple(itertools.chain.from_iterable(map(self._route_tool_calls, signed_chunks)))
def _end_thinking_run(self) -> "tuple[ModelResponseStream, ...]":
held: Final = self._held_signature
self._held_signature = None
self._unsigned_thinking = io.StringIO()
self._unsigned_thinking_length = 0
self._thinking_index = None
return (held,) if held is not None else ()
def _defer_signature(self, chunk: "ModelResponseStream") -> "tuple[ModelResponseStream, ...]":
"""Stream a signed thinking chunk's new text as thinking and send its signature once, when the block ends.
Anthropic-shaped producers stream unsigned text and then sign it, either with no text or
restating the unsigned text, which passes through untouched. Mistral signs every chunk while
each one still carries new text, which would otherwise collapse into ``signature_delta``
events and lose the text.
"""
piece: Final = _ThinkingPiece.of(chunk)
released: Final = self._end_thinking_run() if piece is None or piece.index != self._thinking_index else ()
if piece is None:
return (*released, chunk)
self._thinking_index = piece.index
if not piece.signature:
self._unsigned_thinking_length += self._unsigned_thinking.write(piece.text)
return (*released, chunk)
self._held_signature = None
if not piece.text or self._restates_unsigned_thinking(piece.text):
return (*released, chunk)
signature_chunk: Final = self._with_thinking_block(
chunk, piece.block(text="", signature=piece.signature), keep_role=False
)
self._clear_usage(signature_chunk)
self._held_signature = signature_chunk
return (*released, self._with_thinking_block(chunk, piece.block(text=piece.text, signature=""), keep_role=True))
def _restates_unsigned_thinking(self, text: str) -> bool:
return len(text) == self._unsigned_thinking_length and text == self._unsigned_thinking.getvalue()
def __iter__(self) -> "Iterator[ModelResponseStream]":
return self
@ -493,7 +574,7 @@ class _CombinedChunkSplitter:
try:
chunk = next(self._sync_iter)
except StopIteration:
self._buffer.extend(self._release_held_calls())
self._buffer.extend((*self._end_thinking_run(), *self._release_held_calls()))
if not self._buffer:
raise
break
@ -512,7 +593,7 @@ class _CombinedChunkSplitter:
try:
chunk = await self._async_iter.__anext__()
except StopAsyncIteration:
self._buffer.extend(self._release_held_calls())
self._buffer.extend((*self._end_thinking_run(), *self._release_held_calls()))
if not self._buffer:
raise
break

View file

@ -123,6 +123,7 @@ from openai.types.chat.chat_completion_chunk import Choice as OpenAIStreamingCho
from litellm.litellm_core_utils.prompt_templates.common_utils import (
anthropic_image_source_to_openai_url,
anthropic_inline_document,
parse_tool_call_arguments,
reasoning_content_from_thinking_blocks,
with_prompt_cache_breakpoint,
@ -185,6 +186,8 @@ from litellm.types.llms.openai import (
AllMessageValues,
ChatCompletionAssistantMessage,
ChatCompletionAssistantToolCall,
ChatCompletionFileObject,
ChatCompletionFileObjectFile,
ChatCompletionImageObject,
ChatCompletionImageUrlObject,
ChatCompletionRedactedThinkingBlock,
@ -498,7 +501,9 @@ class LiteLLMAnthropicMessagesAdapter:
for m in ordered_messages:
user_message: ChatCompletionUserMessage | None = None
tool_message_list: list[ChatCompletionToolMessage] = []
new_user_content_list: list[ChatCompletionTextObject | ChatCompletionImageObject] = []
new_user_content_list: list[
ChatCompletionTextObject | ChatCompletionImageObject | ChatCompletionFileObject
] = []
if m["role"] == "system":
system_message = self._translate_midturn_system_message_to_openai(m, model)
if system_message is not None:
@ -531,13 +536,8 @@ class LiteLLMAnthropicMessagesAdapter:
self._add_prompt_cache_breakpoint_if_present(content, image_obj)
)
elif content.get("type") == "document":
# Convert Anthropic document format (PDF, etc.) to OpenAI format
source = content.get("source", {})
openai_image_url = self._translate_anthropic_image_to_openai(cast(dict, source))
if openai_image_url:
image_url_obj = ChatCompletionImageUrlObject(url=openai_image_url)
doc_obj = ChatCompletionImageObject(type="image_url", image_url=image_url_obj)
doc_obj = self._translate_anthropic_document_to_openai(content)
if doc_obj is not None:
self._add_cache_control_if_applicable(content, doc_obj, model)
new_user_content_list.append(doc_obj)
elif content.get("type") == "tool_result":
@ -1357,7 +1357,7 @@ class LiteLLMAnthropicMessagesAdapter:
text: Final = tool_result_text_part(block)
return None if text is None else ChatCompletionTextObject(type="text", text=text)
case "image" | "document":
return self._tool_result_image_part(block.get("source"))
return self._source_image_url_part(block.get("source"))
case "tool_reference":
return ChatCompletionToolReferenceObject(
type="tool_reference", tool_name=str(block.get("tool_name") or "")
@ -1365,7 +1365,18 @@ class LiteLLMAnthropicMessagesAdapter:
case _:
return None
def _tool_result_image_part(self, image_source: object) -> ChatCompletionImageObject | None:
def _translate_anthropic_document_to_openai(
self, document: Mapping[str, object]
) -> ChatCompletionFileObject | ChatCompletionImageObject | None:
inline_document: Final = anthropic_inline_document(document)
if inline_document is None:
return self._source_image_url_part(document.get("source"))
return ChatCompletionFileObject(
type="file",
file=ChatCompletionFileObjectFile(filename=inline_document.filename, file_data=inline_document.file_data),
)
def _source_image_url_part(self, image_source: object) -> ChatCompletionImageObject | None:
if not isinstance(image_source, dict):
return None
openai_image_url = self._translate_anthropic_image_to_openai(image_source)

View file

@ -396,7 +396,8 @@ def _resolve_advisor_router(advisor_model: str) -> "Router | None":
return None
if llm_router is None:
return None
if llm_router.is_recognized_model(advisor_model) or llm_router.pattern_router.route(advisor_model):
registry: Final = llm_router.registry()
if registry.resolve(advisor_model, kind="recognized") or registry.patterns().route(advisor_model):
return llm_router
return None

View file

@ -13,6 +13,7 @@ from typing import Any, Final, cast
from litellm.litellm_core_utils.prompt_templates.common_utils import (
TOOL_RESULT_IMAGE_BOUNDARY,
TOOL_RESULT_IMAGE_PLACEHOLDER,
anthropic_inline_document,
encrypted_reasoning_signature,
responses_reasoning_items_from_thinking_blocks,
with_prompt_cache_breakpoint,
@ -106,27 +107,18 @@ class LiteLLMAnthropicToResponsesAPIAdapter:
block: Mapping[str, object],
) -> dict[str, str] | None: # mutable-ok: API message payload
"""Convert an Anthropic document block to a Responses input_file part."""
inline_document: Final = anthropic_inline_document(block)
if inline_document is not None:
return {
"type": "input_file",
"filename": inline_document.filename,
"file_data": inline_document.file_data,
}
raw_source: Final = block.get("source")
if not isinstance(raw_source, Mapping):
return None
source: Final = cast(Mapping[str, object], raw_source) # cast-ok: untrusted client payload
source_type: Final = source.get("type")
if source_type == "base64":
data: Final = source.get("data")
if not isinstance(data, str) or not data:
return None
raw_media_type: Final = source.get("media_type")
media_type: Final = (
raw_media_type if isinstance(raw_media_type, str) and raw_media_type else "application/pdf"
)
raw_title: Final = block.get("title")
filename: Final = raw_title if isinstance(raw_title, str) and raw_title else "document.pdf"
return {
"type": "input_file",
"filename": filename,
"file_data": f"data:{media_type};base64,{data}",
}
if source_type == "url":
if source.get("type") == "url":
url: Final = source.get("url")
if not isinstance(url, str) or not url:
return None

View file

@ -1,7 +1,7 @@
import copy
import enum
import re
from collections.abc import AsyncIterator, Collection, Iterator
from collections.abc import AsyncIterator, Collection, Iterator, Mapping
from typing import TYPE_CHECKING, Final, cast
import httpx
@ -10,10 +10,11 @@ from httpx import Response
import litellm
from litellm._logging import verbose_logger
from litellm.litellm_core_utils.prompt_templates.common_utils import (
audio_or_image_in_message_content,
content_parts,
convert_content_list_to_str,
drop_tool_reference_parts_from_tool_messages,
filter_value_from_dict,
is_pdf_file_part,
)
from litellm.llms.azure.common_utils import BaseAzureLLM
from litellm.llms.azure_ai.common_utils import (
@ -46,6 +47,20 @@ NON_OPENAI_SPEC_MESSAGE_FIELDS: Final = (
"cache_control",
)
LIST_ONLY_CONTENT_PART_TYPES: Final = frozenset({"image_url", "input_audio"})
def _is_list_only_content_part(part: object) -> bool:
match part:
case {"type": str() as part_type} if part_type in LIST_ONLY_CONTENT_PART_TYPES:
return True
case _:
return is_pdf_file_part(part)
def _has_list_only_content_part(message: Mapping[str, object]) -> bool:
return any(_is_list_only_content_part(part) for part in content_parts(message))
class AzureAIChatCompletionResponseIterator(OpenAIChatCompletionResponseIterator):
def chunk_parser(self, chunk: dict[str, object]) -> ModelResponseStream: # mutable-ok: mirrors override contract
@ -246,7 +261,7 @@ class AzureAIStudioConfig(OpenAIConfig):
these with "Extra inputs are not permitted", which breaks multi-turn
Anthropic-format clients that echo thinking blocks back as history.
2. Transforms list content to a string.
3. If message contains an image or audio, send as is (user-intended)
3. If message contains an image, audio, or PDF file part, send as is (user-intended)
Operates on a deep copy so the caller's messages keep their thinking blocks
and provider metadata, which a fallback to another provider still needs.
@ -257,8 +272,7 @@ class AzureAIStudioConfig(OpenAIConfig):
for field in NON_OPENAI_SPEC_MESSAGE_FIELDS:
filter_value_from_dict(message_dict, field)
# Do nothing if the message contains an image or audio
if audio_or_image_in_message_content(message):
if _has_list_only_content_part(message):
continue
texts = convert_content_list_to_str(message=message)

View file

@ -4,6 +4,7 @@ from typing import TYPE_CHECKING, Any, Final
import httpx
from litellm.litellm_core_utils.core_helpers import drops_unsupported_params
from litellm.types.llms.anthropic_messages.anthropic_response import (
AnthropicMessagesResponse,
)
@ -201,11 +202,23 @@ class BaseAnthropicMessagesConfig(ABC):
anthropic_request_errors_on_thinking_binding_mismatch,
is_anthropic_invalid_thinking_block_error,
is_anthropic_thinking_binding_mismatch_error,
is_claude_strict_tools_rejection,
tools_without_strict,
validated_object_mapping,
)
if e.response.status_code != 400:
return False
error_text: Final[str] = e.response.text
body: Final = validated_object_mapping(request_body)
drops_rejected_strict_tools: Final = (
drops_unsupported_params(litellm_params.get("drop_params"))
and is_claude_strict_tools_rejection(e.response.status_code, error_text)
and body is not None
and tools_without_strict(body.get("tools")) is not None
)
if drops_rejected_strict_tools:
return True
caller_asked_for_this_error: Final = anthropic_request_errors_on_thinking_binding_mismatch(
request_body
) and is_anthropic_thinking_binding_mismatch_error(error_text)
@ -219,11 +232,22 @@ class BaseAnthropicMessagesConfig(ABC):
"""
from litellm.llms.anthropic.common_utils import (
is_anthropic_invalid_thinking_block_error,
is_claude_strict_tools_rejection,
strip_thinking_blocks_from_anthropic_messages_request_dict,
tools_without_strict,
warn_claude_strict_tools_dropped,
)
if e.response.status_code == 400 and is_anthropic_invalid_thinking_block_error(e.response.text):
strip_thinking_blocks_from_anthropic_messages_request_dict(request_data)
without_strict: Final = (
tools_without_strict(request_data.get("tools"))
if is_claude_strict_tools_rejection(e.response.status_code, e.response.text)
else None
)
if without_strict is not None:
warn_claude_strict_tools_dropped(without_strict.dropped_count)
request_data["tools"] = list(without_strict.tools) # rebind-ok: the retry loop re-signs this same dict
return request_data
def native_output_fallback_litellm_params(

View file

@ -4,7 +4,6 @@ from types import TracebackType
from typing import TYPE_CHECKING, Any, Protocol
import httpx
from typing_extensions import Self
from litellm.types.llms.openai import OpenAIRealtimeStreamSessionEvents
from litellm.types.realtime import (
@ -24,7 +23,7 @@ else:
class RealtimeBackend(Protocol):
async def __aenter__(self) -> Self: ...
async def __aenter__(self) -> object: ...
async def __aexit__(
self,
@ -97,9 +96,15 @@ class BaseRealtimeConfig(ABC):
def unbilled_usage_on_session_close(self, model: str) -> RealtimeInputAudioTranscriptionUsage | None:
return None
def session_refusal(self) -> Exception | None:
return None
async def open_backend(self, url: str, headers: Mapping[str, str]) -> RealtimeBackend | None:
return None
def wrap_backend(self, backend: RealtimeBackend) -> RealtimeBackend:
return backend
def transform_session_created_event(
self,
model: str,

View file

@ -93,9 +93,12 @@ class RouterVectorStoreEmbeddingExecutor:
def _router_serves(self, model: str) -> bool:
team_id: Final = self.metadata.get("user_api_key_team_id")
resolved: Final = self.router.resolved_litellm_models(model, team_id if isinstance(team_id, str) else None)
registry: Final = self.router.registry()
resolved: Final = registry.names(
model, kind="litellm_models", team_id=team_id if isinstance(team_id, str) else None
)
deployment_models: Final = (
deployment.get("litellm_params", {}).get("model") for deployment in self.router.get_model_list() or ()
deployment.get("litellm_params", {}).get("model") for deployment in registry.deployments() or ()
)
return bool(resolved) or model in deployment_models
@ -124,6 +127,9 @@ class BaseVectorStoreConfig:
def validate_create_vector_store(self) -> None:
return None
def get_async_httpx_client_params(self) -> Mapping[str, object]:
return MappingProxyType({})
def get_supported_openai_params(self, model: str) -> list[VECTOR_STORE_OPENAI_PARAMS]:
return []
@ -226,6 +232,9 @@ class BaseVectorStoreConfig:
``get_complete_url`` adds a query string, like Azure, override this."""
return f"{api_base}/{encoded_vector_store_id}"
def get_vector_store_list_url(self, api_base: str) -> str:
return api_base
def get_error_class(self, error_message: str, status_code: int, headers: dict | httpx.Headers) -> BaseLLMException:
from ..chat.transformation import BaseLLMException

View file

@ -2,11 +2,11 @@ from collections.abc import Callable, Coroutine
from functools import partial
from typing import Final, TypeAlias
from litellm.llms.anthropic.common_utils import is_claude_strict_tools_rejection
from litellm.llms.base_llm.chat.transformation import BaseLLMException
from litellm.llms.bedrock.common_utils import (
NativeOutputFallback,
is_bedrock_native_output_rejection,
is_bedrock_strict_tools_rejection,
warn_bedrock_native_output_rejected,
warn_bedrock_strict_tools_dropped,
)
@ -34,7 +34,7 @@ def _native_output_retry(
def _strict_tools_retry(
error: BaseLLMException, fallback: NativeOutputFallbackFactory | None, model: str
) -> NativeOutputFallback | None:
if fallback is None or not is_bedrock_strict_tools_rejection(error.status_code, error.message):
if fallback is None or not is_claude_strict_tools_rejection(error.status_code, error.message):
return None
retry: Final = fallback()
if retry is not None:

View file

@ -366,21 +366,10 @@ class NativeOutputFallback:
litellm_params: Mapping[str, object]
_BEDROCK_STRICT_TOOLS_REJECTION: Final = re.compile(
r"Too many strict tools \(\d+\)"
r"|Schemas contains too many (?:optional parameters|parameters with union types)"
r"|[Cc]ompiled grammar (?:is too large|size)"
r"|tools\.\d+\.custom(?:\.strict)?: "
r"|Grammar compilation timed out"
)
_TOOLS_ADAPTER: Final = TypeAdapter(tuple[Mapping[str, object], ...])
_TOOL_SPEC_ADAPTER: Final = TypeAdapter(Mapping[str, object])
def is_bedrock_strict_tools_rejection(status_code: int, error_text: str) -> bool:
return status_code == 400 and _BEDROCK_STRICT_TOOLS_REJECTION.search(error_text) is not None
def warn_bedrock_strict_tools_dropped(model: str) -> None:
verbose_logger.warning(
"Bedrock rejected the strict tool schemas for model=%s, so litellm sent the request again without `strict` "

View file

@ -22,6 +22,7 @@ import httpx
import litellm
from litellm._logging import verbose_logger
from litellm.litellm_core_utils.core_helpers import drops_unsupported_params
from litellm.litellm_core_utils.url_utils import join_api_base
from litellm.llms.base_llm.chat.transformation import BaseLLMException
from litellm.llms.base_llm.responses.codex_compat import drop_unsupported_tools, normalize_codex_input_items
@ -196,7 +197,9 @@ class BedrockMantleResponsesAPIConfig(BedrockMantleAuthMixin, OpenAIResponsesAPI
params: Final = cast( # cast-ok: the base signature leaves the params dict untyped
"ResponsesAPIOptionalRequestParams", response_api_optional_request_params
)
hoisted: Final = hoist_additional_tools(input, params.get("tools"))
hoisted: Final = hoist_additional_tools(
input, params.get("tools"), drop_params=drops_unsupported_params(litellm_params.drop_params)
)
normalized_input, rewritten_types = normalize_codex_input_items(hoisted.input)
if rewritten_types:
verbose_logger.warning(

View file

@ -0,0 +1,128 @@
import weakref
from collections.abc import AsyncIterator, Callable
from typing import Final
import httpx
HTTP2_STREAMS_PER_CONNECTION: Final = 100
class _Lane:
__slots__ = ("in_flight", "transport")
def __init__(self, transport: httpx.AsyncBaseTransport) -> None:
self.transport: Final = transport
self.in_flight = 0
def acquire(self) -> None:
self.in_flight += 1
def release(self) -> None:
self.in_flight -= 1
class _Lanes:
def __init__(
self,
new_transport: Callable[[httpx.Limits], httpx.AsyncBaseTransport],
limits: httpx.Limits,
streams_per_connection: int,
) -> None:
self._new_transport: Final = new_transport
self._limits: Final = limits
self._streams_per_connection: Final = streams_per_connection
self.lanes: tuple[_Lane, ...] = ()
def acquire(self) -> _Lane:
free: Final = next((lane for lane in self.lanes if lane.in_flight < self._streams_per_connection), None)
lane: Final = free if free is not None else self._open_or_share()
lane.acquire()
return lane
def _open_or_share(self) -> _Lane:
cap: Final = self._limits.max_connections
if cap is not None and len(self.lanes) * self._streams_per_connection >= cap:
return min(self.lanes, key=lambda lane: lane.in_flight)
return self._open()
def _open(self) -> _Lane:
lane: Final = _Lane(self._new_transport(self._next_lane_limits()))
self.lanes = (*self.lanes, lane)
return lane
def _next_lane_limits(self) -> httpx.Limits:
cap: Final = self._limits.max_connections
if cap is None:
return self._limits
return httpx.Limits(
max_connections=min(self._streams_per_connection, cap - len(self.lanes) * self._streams_per_connection),
max_keepalive_connections=self._limits.max_keepalive_connections,
keepalive_expiry=self._limits.keepalive_expiry,
)
class _AsyncReleasingStream(httpx.AsyncByteStream):
def __init__(self, stream: httpx.AsyncByteStream, release: Callable[[], None]) -> None:
self._stream: Final = stream
self._release: Final = weakref.finalize(self, release)
async def __aiter__(self) -> AsyncIterator[bytes]:
try:
async for chunk in self._stream:
yield chunk
except BaseException:
self._release()
raise
async def aclose(self) -> None:
self._release()
await self._stream.aclose()
def _allowed_over_http2(name: bytes, value: bytes) -> bool:
return name.lower() != b"te" or value.strip().lower() == b"trailers"
def _without_te_codings(request: httpx.Request) -> httpx.Request:
if request.url.scheme != "https":
return request
kept: Final = tuple((name, value) for name, value in request.headers.raw if _allowed_over_http2(name, value))
if len(kept) == len(request.headers.raw):
return request
return httpx.Request(
request.method, request.url, headers=kept, stream=request.stream, extensions=request.extensions
)
def _with_stream(response: httpx.Response, stream: httpx.AsyncByteStream) -> httpx.Response:
return httpx.Response(response.status_code, headers=response.headers, stream=stream, extensions=response.extensions)
class AsyncHTTP2ConnectionSpread(httpx.AsyncBaseTransport):
def __init__(
self,
new_transport: Callable[[httpx.Limits], httpx.AsyncBaseTransport],
limits: httpx.Limits,
streams_per_connection: int = HTTP2_STREAMS_PER_CONNECTION,
) -> None:
self._lanes: Final = _Lanes(new_transport, limits, streams_per_connection)
async def handle_async_request(self, request: httpx.Request) -> httpx.Response:
lane: Final = self._lanes.acquire()
try:
response: Final = await lane.transport.handle_async_request(_without_te_codings(request))
except BaseException:
lane.release()
raise
stream: Final = response.stream
if not isinstance(stream, httpx.AsyncByteStream):
lane.release()
return response
return _with_stream(response, _AsyncReleasingStream(stream, lane.release))
def has_request_in_flight(self) -> bool:
return any(lane.in_flight > 0 for lane in self._lanes.lanes)
async def aclose(self) -> None:
for lane in self._lanes.lanes:
await lane.transport.aclose()

View file

@ -18,6 +18,7 @@ import certifi
import httpx
from aiohttp import ClientSession, DummyCookieJar, TCPConnector
from httpx import USE_CLIENT_DEFAULT, AsyncHTTPTransport, HTTPTransport
from httpx._config import DEFAULT_LIMITS
from httpx._types import CertTypes, RequestFiles
from httpx._utils import get_environment_proxies
from pydantic import TypeAdapter
@ -43,6 +44,7 @@ from litellm.litellm_core_utils.logging_utils import track_llm_api_timing
from litellm.litellm_core_utils.request_timeout_resolver import (
get_configured_request_timeout,
)
from litellm.llms.custom_httpx.http2_connection_spread import AsyncHTTP2ConnectionSpread
from litellm.types.llms.custom_http import *
if TYPE_CHECKING:
@ -73,7 +75,7 @@ _RequestContent: TypeAlias = str | bytes | Iterable[bytes] | AsyncIterable[bytes
_IPV4_LOCAL_ADDRESS: Final = "0.0.0.0"
_HttpxTransportT = TypeVar("_HttpxTransportT", HTTPTransport, AsyncHTTPTransport)
_HttpxTransportT: Final = TypeVar("_HttpxTransportT", httpx.BaseTransport, httpx.AsyncBaseTransport)
HTTPEventHooks: TypeAlias = Mapping[str, list[Callable[..., object]]]
@ -84,6 +86,36 @@ def http2_enabled() -> bool:
return litellm.http2 is True or str_to_bool(os.getenv("LITELLM_HTTP2", "False")) is True
_HTTP2_OFF_VALUES: Final = frozenset({"false", "0", "no", "off"})
def _is_http2_off(value: bool | int | str | None) -> bool:
return str(value).strip().lower() in _HTTP2_OFF_VALUES
def http2_disabled() -> bool:
return _is_http2_off(litellm.http2) or _is_http2_off(os.getenv("LITELLM_HTTP2"))
def _resolve_http2(provider_default: bool, ssl_verify: VerifyTypes | None) -> bool:
if http2_enabled():
return True
return provider_default and not http2_disabled() and not isinstance(ssl_verify, ssl.SSLContext)
def _aiohttp_env_proxies_disabled() -> bool:
from litellm.secret_managers.main import str_to_bool
return litellm.disable_aiohttp_trust_env or str_to_bool(os.getenv("DISABLE_AIOHTTP_TRUST_ENV", "False")) is True
_AIOHTTP_EQUIVALENT_LIMITS: Final = httpx.Limits(
max_connections=AIOHTTP_CONNECTOR_LIMIT if AIOHTTP_CONNECTOR_LIMIT > 0 else None,
max_keepalive_connections=None,
keepalive_expiry=AIOHTTP_KEEPALIVE_TIMEOUT,
)
def _environment_proxy_mounts(
build_proxy_transport: Callable[[str], _HttpxTransportT],
) -> Mapping[str, _HttpxTransportT | None]:
@ -145,7 +177,7 @@ def build_aiohttp_keepalive_socket_factory() -> Callable[[_AddrInfo], socket.soc
_build_aiohttp_keepalive_socket_factory = build_aiohttp_keepalive_socket_factory
def get_default_headers() -> dict:
def get_default_headers() -> dict[str, str]:
"""
Get default headers for HTTP requests.
@ -280,9 +312,7 @@ def _prepare_request_data_and_content(
# Cache for SSL contexts to avoid creating duplicate contexts with the same configuration
# Key: tuple of (cafile, ssl_security_level, ssl_ecdh_curve)
# Value: ssl.SSLContext
_ssl_context_cache: Final[dict[tuple[str | None, str | None, str | None], ssl.SSLContext]] = {}
_ssl_context_cache: Final[dict[tuple[str | None, str | None, str | None, bool], ssl.SSLContext]] = {}
def _create_ssl_context(
@ -375,6 +405,7 @@ def get_ssl_verify(
def get_ssl_configuration(
ssl_verify: VerifyTypes | None = None,
http2: bool = False,
) -> bool | str | ssl.SSLContext:
"""
Unified SSL configuration function that handles ssl_context and ssl_verify logic.
@ -397,6 +428,8 @@ def get_ssl_configuration(
- False: Disable SSL verification
- True: Enable SSL verification
- str: Path to CA bundle file
http2: Whether the context is for an HTTP/2 client. httpx sets ALPN on the context it is given,
so HTTP/2 clients get their own cached context instead of sharing the HTTP/1.1 one
Returns:
Union[bool, str, ssl.SSLContext]: Appropriate SSL configuration
@ -423,7 +456,7 @@ def get_ssl_configuration(
if ssl_verify is not False:
# Create cache key from configuration parameters
cache_key: Final = (cafile, ssl_security_level, ssl_ecdh_curve)
cache_key: Final = (cafile, ssl_security_level, ssl_ecdh_curve, http2)
# Check if we have a cached SSL context for this configuration
if cache_key not in _ssl_context_cache:
@ -439,6 +472,11 @@ def get_ssl_configuration(
return ssl_verify
def _httpx_default_verify_unless_configured(verify: VerifyTypes) -> VerifyTypes:
litellm_defaults: Final = (get_ssl_configuration(True), get_ssl_configuration(True, http2=True))
return True if any(verify is default for default in litellm_defaults) else verify
_shared_realtime_ssl_context: bool | str | ssl.SSLContext | None = None
@ -662,6 +700,7 @@ class AsyncHTTPHandler:
shared_session: Optional["ClientSession"] = None,
transport: httpx.AsyncBaseTransport | None = None,
follow_redirects: bool = True,
http2: bool = False,
):
self.timeout = timeout
self.event_hooks = event_hooks
@ -669,6 +708,7 @@ class AsyncHTTPHandler:
self.shared_session = shared_session
self.transport = transport
self.follow_redirects = follow_redirects
self.http2 = http2
self._owns_client = True
self._client = self.create_client(
timeout=timeout,
@ -711,8 +751,11 @@ class AsyncHTTPHandler:
follow_redirects=self.follow_redirects,
trust_env=False,
)
use_http2: Final = _resolve_http2(self.http2, ssl_verify)
replaces_aiohttp: Final = use_http2 and AsyncHTTPHandler.should_use_aiohttp_transport()
limits: Final = _AIOHTTP_EQUIVALENT_LIMITS if replaces_aiohttp else DEFAULT_LIMITS
# Get unified SSL configuration
ssl_config: Final = get_ssl_configuration(ssl_verify)
ssl_config: Final = get_ssl_configuration(ssl_verify, http2=use_http2)
# An SSL certificate used by the requested host to authenticate the client.
# /path/to/client.pem
@ -726,14 +769,23 @@ class AsyncHTTPHandler:
ssl_context=ssl_config if isinstance(ssl_config, ssl.SSLContext) else None,
ssl_verify=ssl_config if isinstance(ssl_config, bool) else None,
shared_session=shared_session,
http2=use_http2,
cert=cert,
limits=limits,
)
# Get default headers (User-Agent, overridable via LITELLM_USER_AGENT)
default_headers: Final = get_default_headers()
trust_env: Final = not (replaces_aiohttp and _aiohttp_env_proxies_disabled())
return httpx.AsyncClient(
transport=transport,
mounts=AsyncHTTPHandler.create_httpx_proxy_mounts(transport, verify=ssl_config, cert=cert),
mounts=AsyncHTTPHandler.create_httpx_proxy_mounts(
transport, verify=ssl_config, cert=cert, http2=use_http2, limits=limits
)
if trust_env
else None,
limits=limits,
event_hooks=event_hooks,
timeout=timeout,
verify=ssl_config,
@ -741,7 +793,8 @@ class AsyncHTTPHandler:
headers=default_headers,
cookies=blocked_cookie_jar(),
follow_redirects=self.follow_redirects,
http2=http2_enabled(),
http2=use_http2,
trust_env=trust_env,
)
async def close(self):
@ -1152,6 +1205,17 @@ class AsyncHTTPHandler:
# Mirrors LiteLLMAiohttpTransport._background_close_tasks.
_finalizer_close_tasks: ClassVar[set["asyncio.Task[None]"]] = set() # mutable-ok: strong refs for pending closes
@classmethod
async def _close_client_at_finalization(cls, client: httpx.AsyncClient) -> None:
try:
await client.aclose()
except (OSError, RuntimeError) as exc:
verbose_logger.debug("Error closing client at finalization: %s", exc)
finally:
task: Final = asyncio.current_task()
if task is not None:
cls._finalizer_close_tasks.discard(task)
@classmethod
def _on_finalizer_close_done(cls, task: "asyncio.Task[None]") -> None:
cls._finalizer_close_tasks.discard(task)
@ -1219,8 +1283,8 @@ class AsyncHTTPHandler:
# here is the cross-loop path the transport refuses.
self._dispose_wrapped_aiohttp_session()
return
task: Final = loop.create_task(self._client.aclose())
cls: Final = type(self)
task: Final = loop.create_task(cls._close_client_at_finalization(self._client))
cls._finalizer_close_tasks.add(task)
task.add_done_callback(cls._on_finalizer_close_done)
except Exception:
@ -1231,7 +1295,10 @@ class AsyncHTTPHandler:
ssl_context: ssl.SSLContext | None = None,
ssl_verify: bool | None = None,
shared_session: Optional["ClientSession"] = None,
) -> LiteLLMAiohttpTransport | AsyncHTTPTransport | None:
http2: bool = False,
cert: CertTypes | None = None,
limits: httpx.Limits = DEFAULT_LIMITS,
) -> LiteLLMAiohttpTransport | AsyncHTTPTransport | AsyncHTTP2ConnectionSpread | None:
"""
- Creates a transport for httpx.AsyncClient
- if litellm.force_ipv4 is True, it will return AsyncHTTPTransport with local_address="0.0.0.0"
@ -1249,7 +1316,7 @@ class AsyncHTTPHandler:
#########################################################
# AIOHTTP TRANSPORT is off by default
#########################################################
if AsyncHTTPHandler.should_use_aiohttp_transport():
if not http2 and AsyncHTTPHandler.should_use_aiohttp_transport():
return AsyncHTTPHandler.create_aiohttp_transport(
ssl_context=ssl_context,
ssl_verify=ssl_verify,
@ -1259,7 +1326,12 @@ class AsyncHTTPHandler:
#########################################################
# HTTPX TRANSPORT is used when aiohttp is not installed
#########################################################
return AsyncHTTPHandler._create_httpx_transport()
return AsyncHTTPHandler._create_httpx_transport(
http2=http2,
verify=ssl_context if ssl_context is not None else ssl_verify is not False,
cert=cert,
limits=limits,
)
_create_async_transport = create_async_transport
@ -1406,28 +1478,57 @@ class AsyncHTTPHandler:
_create_aiohttp_transport = create_aiohttp_transport
@staticmethod
def _create_httpx_transport() -> AsyncHTTPTransport | None:
def _create_httpx_transport(
http2: bool = False,
verify: ssl.SSLContext | bool = True,
cert: CertTypes | None = None,
limits: httpx.Limits = DEFAULT_LIMITS,
) -> AsyncHTTPTransport | AsyncHTTP2ConnectionSpread | None:
"""
Creates an AsyncHTTPTransport
- If HTTP/2 is on, it spreads requests over as many HTTP/2 connections as the stream limit needs
- If force_ipv4 is True, it will create an AsyncHTTPTransport with local_address set to "0.0.0.0"
- [Default] If force_ipv4 is False, it will return None
- [Default] Otherwise it will return None
"""
local_address: Final = _IPV4_LOCAL_ADDRESS if litellm.force_ipv4 else None
transport_verify: Final = (
_httpx_default_verify_unless_configured(verify)
if litellm.force_ipv4 and not AsyncHTTPHandler.should_use_aiohttp_transport()
else verify
)
if http2 or http2_enabled():
return AsyncHTTP2ConnectionSpread(
lambda lane_limits: AsyncHTTPTransport(
local_address=local_address, verify=transport_verify, cert=cert, http2=True, limits=lane_limits
),
limits,
)
if litellm.force_ipv4:
return AsyncHTTPTransport(local_address=_IPV4_LOCAL_ADDRESS, http2=http2_enabled())
else:
return None
return AsyncHTTPTransport(local_address=local_address, verify=transport_verify, cert=cert, limits=limits)
return None
@staticmethod
def create_httpx_proxy_mounts(
transport: LiteLLMAiohttpTransport | AsyncHTTPTransport | None,
transport: LiteLLMAiohttpTransport | AsyncHTTPTransport | AsyncHTTP2ConnectionSpread | None,
verify: VerifyTypes,
cert: CertTypes | None,
) -> Mapping[str, AsyncHTTPTransport | None] | None:
if not isinstance(transport, AsyncHTTPTransport):
http2: bool = False,
limits: httpx.Limits = DEFAULT_LIMITS,
) -> Mapping[str, httpx.AsyncBaseTransport | None] | None:
if not isinstance(transport, (AsyncHTTPTransport, AsyncHTTP2ConnectionSpread)):
return None
if http2 or http2_enabled():
return _environment_proxy_mounts(
lambda proxy_url: AsyncHTTP2ConnectionSpread(
lambda lane_limits: AsyncHTTPTransport(
proxy=proxy_url, verify=verify, cert=cert, http2=True, limits=lane_limits
),
limits,
)
)
return _environment_proxy_mounts(
lambda proxy_url: AsyncHTTPTransport(proxy=proxy_url, verify=verify, cert=cert, http2=http2_enabled())
lambda proxy_url: AsyncHTTPTransport(proxy=proxy_url, verify=verify, cert=cert, limits=limits)
)
_create_httpx_proxy_mounts = create_httpx_proxy_mounts
@ -1453,8 +1554,9 @@ class HTTPHandler:
self._client = self.create_client() if client is None else client
def create_client(self) -> httpx.Client:
use_http2: Final = http2_enabled()
# Get unified SSL configuration
ssl_config: Final = get_ssl_configuration(self.ssl_verify)
ssl_config: Final = get_ssl_configuration(self.ssl_verify, http2=use_http2)
# An SSL certificate used by the requested host to authenticate the client.
# /path/to/client.pem
@ -1463,17 +1565,19 @@ class HTTPHandler:
# Get default headers (User-Agent, overridable via LITELLM_USER_AGENT)
default_headers: Final = get_default_headers() if not self.disable_default_headers else None
transport: Final = self._create_sync_transport(verify=ssl_config, cert=cert, http2=use_http2)
# Create a client with a connection pool
return httpx.Client(
transport=self._create_sync_transport(),
mounts=self._create_sync_proxy_mounts(verify=ssl_config, cert=cert),
transport=transport,
mounts=self._create_sync_proxy_mounts(transport, verify=ssl_config, cert=cert, http2=use_http2),
timeout=self.timeout if self.timeout is not None else _DEFAULT_TIMEOUT,
verify=ssl_config,
cert=cert,
headers=default_headers,
cookies=blocked_cookie_jar(),
follow_redirects=self.follow_redirects,
http2=http2_enabled(),
http2=use_http2,
)
@property
@ -1769,27 +1873,34 @@ class HTTPHandler:
except Exception:
pass
def _create_sync_transport(self) -> HTTPTransport | None:
@staticmethod
def _create_sync_transport(verify: VerifyTypes, cert: CertTypes | None, http2: bool) -> HTTPTransport | None:
"""
Create an HTTP transport with IPv4 only if litellm.force_ipv4 is True.
Otherwise, return None.
Otherwise, return litellm.sync_transport or None.
Some users have seen httpx ConnectionError when using ipv6 - forcing ipv4 resolves the issue for them
"""
if litellm.force_ipv4:
return HTTPTransport(local_address=_IPV4_LOCAL_ADDRESS, http2=http2_enabled())
else:
return getattr(litellm, "sync_transport", None)
return HTTPTransport(
local_address=_IPV4_LOCAL_ADDRESS,
verify=_httpx_default_verify_unless_configured(verify),
cert=cert,
http2=http2,
)
return getattr(litellm, "sync_transport", None)
@staticmethod
def _create_sync_proxy_mounts(
transport: httpx.BaseTransport | None,
verify: VerifyTypes,
cert: CertTypes | None,
) -> Mapping[str, HTTPTransport | None] | None:
if not litellm.force_ipv4:
http2: bool,
) -> Mapping[str, httpx.BaseTransport | None] | None:
if transport is None or transport is getattr(litellm, "sync_transport", None):
return None
return _environment_proxy_mounts(
lambda proxy_url: HTTPTransport(proxy=proxy_url, verify=verify, cert=cert, http2=http2_enabled())
lambda proxy_url: HTTPTransport(proxy=proxy_url, verify=verify, cert=cert, http2=http2)
)
@ -1873,7 +1984,7 @@ def get_httpx_client(params: dict | None = None) -> HTTPHandler:
cache = LLMClientCache()
setattr(litellm, "in_memory_llm_clients_cache", cache)
_cached_client: Final = cache.get_cache(_cache_key_name)
_cached_client: Final = cache.get_cache(_cache_key_name, loop_bound=False)
if _cached_client:
return _cached_client
@ -1889,6 +2000,7 @@ def get_httpx_client(params: dict | None = None) -> HTTPHandler:
value=_new_client,
ttl=_DEFAULT_TTL_FOR_HTTPX_CLIENTS,
litellm_owned_client=True,
loop_bound=False,
)
return _new_client

View file

@ -26,6 +26,7 @@ from typing import (
from urllib.parse import parse_qs, urlencode, urlparse, urlunparse
import httpx
from h2.events import ConnectionTerminated
from httpx import USE_CLIENT_DEFAULT
from httpx._types import FileContent
from openai.types.file_deleted import FileDeleted
@ -371,6 +372,34 @@ class _SignedBodyKwargs(TypedDict, total=False):
json: ReadOnly[dict[str, object]]
_GOAWAY_RETRY_MESSAGE: Final = "Vector store search connection was closed by an HTTP/2 GOAWAY, retrying once"
def _closed_by_http2_goaway(error: httpx.RemoteProtocolError) -> bool:
cause: Final = error.__cause__
return cause is not None and any(isinstance(arg, ConnectionTerminated) for arg in cause.args)
async def _asend_with_goaway_retry(send: Callable[[], Awaitable[httpx.Response]]) -> httpx.Response:
try:
return await send()
except httpx.RemoteProtocolError as error:
if not _closed_by_http2_goaway(error):
raise
verbose_logger.warning(_GOAWAY_RETRY_MESSAGE)
return await send()
def _send_with_goaway_retry(send: Callable[[], httpx.Response]) -> httpx.Response:
try:
return send()
except httpx.RemoteProtocolError as error:
if not _closed_by_http2_goaway(error):
raise
verbose_logger.warning(_GOAWAY_RETRY_MESSAGE)
return send()
def _is_sized_json_body(response: httpx.Response) -> bool:
headers: Final = httpx.Headers(response.headers)
media_types: Final = (value.split(";")[0].strip().lower() for value in headers.get_list("content-type"))
@ -2438,8 +2467,7 @@ class BaseLLMHTTPHandler:
)
else:
verbose_logger.debug(
"Anthropic /v1/messages: invalid thinking signature; "
"stripping thinking blocks and retrying (attempt %s/%s).",
"Anthropic /v1/messages: retrying after a recoverable HTTP error (attempt %s/%s).",
len(recovered_errors) + 2,
max_attempts,
)
@ -6923,7 +6951,7 @@ class BaseLLMHTTPHandler:
try:
ssl_context: Final = realtime_ssl_for_url(url)
provider_backend: Final = await provider_config.open_backend(url, headers)
backend_ws: Final = (
backend_ws: Final = provider_config.wrap_backend(
provider_backend
if provider_backend is not None
else await self._open_realtime_backend_ws(websockets, url, headers, ssl_context)
@ -6961,9 +6989,7 @@ class BaseLLMHTTPHandler:
await backend_ws.send(_session_config)
realtime_streaming.session_configuration_request = _session_config
# For providers that defer setup until client session.update, optionally
# send synthetic session.created to unblock clients waiting on connect.
if not provider_config.requires_session_configuration():
if not _session_config:
synthetic_session: Final = provider_config.transform_session_created_event(
model=model,
logging_session_id=logging_obj.litellm_trace_id,
@ -10629,7 +10655,10 @@ class BaseLLMHTTPHandler:
if client is None or not isinstance(client, AsyncHTTPHandler):
async_httpx_client = get_async_httpx_client(
llm_provider=litellm.LlmProviders(custom_llm_provider),
params={"ssl_verify": litellm_params.get("ssl_verify", None)},
params={
"ssl_verify": litellm_params.get("ssl_verify", None),
**vector_store_provider_config.get_async_httpx_client_params(),
},
)
else:
async_httpx_client = client
@ -10701,11 +10730,8 @@ class BaseLLMHTTPHandler:
request_data: Final = json.dumps(request_body) if signed_json_body is None else signed_json_body
try:
response: Final = await async_httpx_client.post(
url=url,
headers=headers,
data=request_data,
timeout=timeout,
response: Final = await _asend_with_goaway_retry(
partial(async_httpx_client.post, url=url, headers=headers, data=request_data, timeout=timeout)
)
except httpx.TimeoutException:
raise vector_store_provider_config.get_error_class(
@ -10838,11 +10864,8 @@ class BaseLLMHTTPHandler:
request_data: Final = json.dumps(request_body) if signed_json_body is None else signed_json_body
try:
response: Final = sync_httpx_client.post(
url=url,
headers=headers,
data=request_data,
timeout=timeout,
response: Final = _send_with_goaway_retry(
partial(sync_httpx_client.post, url=url, headers=headers, data=request_data, timeout=timeout)
)
except httpx.TimeoutException:
raise vector_store_provider_config.get_error_class(
@ -11152,7 +11175,7 @@ class BaseLLMHTTPHandler:
litellm_params=dict(litellm_params),
)
url: Final = api_base
url: Final = vector_store_provider_config.get_vector_store_list_url(api_base)
params: Final[dict[str, object]] = {}
if after is not None:
@ -11231,7 +11254,7 @@ class BaseLLMHTTPHandler:
litellm_params=dict(litellm_params),
)
url: Final = api_base
url: Final = vector_store_provider_config.get_vector_store_list_url(api_base)
params: Final[dict[str, object]] = {}
if after is not None:

View file

@ -10,6 +10,7 @@ from litellm.secret_managers.main import get_secret_str
from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolParam
from ...openai.chat.gpt_transformation import OpenAIGPTConfig
from ..common_utils import resolve_dashscope_openai_api_base
class DashScopeChatConfig(OpenAIGPTConfig):
@ -55,11 +56,9 @@ class DashScopeChatConfig(OpenAIGPTConfig):
def _get_openai_compatible_provider_info(
self, api_base: str | None, api_key: str | None
) -> tuple[str | None, str | None]:
api_base = (
api_base or get_secret_str("DASHSCOPE_API_BASE") or "https://dashscope.aliyuncs.com/compatible-mode/v1"
)
resolved_api_base: Final = resolve_dashscope_openai_api_base(api_base)
dynamic_api_key: Final = api_key or get_secret_str("DASHSCOPE_API_KEY")
return api_base, dynamic_api_key
return resolved_api_base, dynamic_api_key
def get_openai_compatible_provider_info(
self,
@ -69,7 +68,7 @@ class DashScopeChatConfig(OpenAIGPTConfig):
return self._get_openai_compatible_provider_info(api_base, api_key)
def _resolve_chat_api_base(self, api_base: str | None) -> str:
return api_base or "https://dashscope.aliyuncs.com/compatible-mode/v1"
return resolve_dashscope_openai_api_base(api_base)
def get_complete_url(
self,

View file

@ -2,6 +2,7 @@
Common utilities for the DashScope LLM provider.
"""
from dataclasses import dataclass
from typing import TYPE_CHECKING, Final
from urllib.parse import urlparse
@ -19,10 +20,123 @@ if TYPE_CHECKING:
from litellm.llms.base_llm.videos.transformation import BaseVideoConfig
DASHSCOPE_CHAT_COMPATIBLE_PATH: Final = "/compatible-mode/v1"
DASHSCOPE_RERANK_PATH: Final = "/compatible-api/v1/reranks"
DASHSCOPE_ANTHROPIC_PATH: Final = "/apps/anthropic"
DASHSCOPE_NATIVE_V1_PATH: Final = "/api/v1"
DASHSCOPE_COMPATIBLE_API_PATH: Final = "/compatible-api/v1"
DASHSCOPE_RERANK_PATH: Final = f"{DASHSCOPE_COMPATIBLE_API_PATH}/reranks"
DASHSCOPE_IMAGE_PATH: Final = "/api/v1/services/aigc/multimodal-generation/generation"
DASHSCOPE_DEFAULT_ORIGIN: Final = "https://dashscope.aliyuncs.com"
_STANDARD_DASHSCOPE_HOSTS: Final = frozenset(
{
"dashscope.aliyuncs.com",
"dashscope-intl.aliyuncs.com",
"dashscope-us.aliyuncs.com",
}
)
_WORKSPACE_HOST_SUFFIX: Final = ".maas.aliyuncs.com"
_EXCLUDED_WORKSPACE_PREFIXES: Final = ("trial.", "token-plan.")
_PROTOCOL_PATH_PREFIXES: Final = (
DASHSCOPE_CHAT_COMPATIBLE_PATH,
DASHSCOPE_ANTHROPIC_PATH,
DASHSCOPE_NATIVE_V1_PATH,
DASHSCOPE_COMPATIBLE_API_PATH,
)
def _rerank_base_for_chat_shaped_base(api_base: str | None) -> str | None:
@dataclass(frozen=True, slots=True)
class DashScopeEndpoints:
origin: str
openai: str
anthropic: str
native: str
def _is_dashscope_host(host: str) -> bool:
if host in _STANDARD_DASHSCOPE_HOSTS:
return True
return host.endswith(_WORKSPACE_HOST_SUFFIX) and not host.startswith(_EXCLUDED_WORKSPACE_PREFIXES)
def _dashscope_origin(api_base: str | None) -> str | None:
if not api_base:
return None
parsed: Final = urlparse(api_base)
if parsed.scheme not in ("http", "https") or not parsed.netloc:
return None
if parsed.query or parsed.fragment:
return None
host: Final = (parsed.hostname or "").lower()
if not _is_dashscope_host(host):
return None
path: Final = parsed.path.rstrip("/")
if path and not any(path == prefix or path.startswith(f"{prefix}/") for prefix in _PROTOCOL_PATH_PREFIXES):
return None
return f"{parsed.scheme}://{parsed.netloc}"
def resolve_dashscope_endpoints(api_base: str | None) -> DashScopeEndpoints | None:
"""
Derive DashScope's protocol-specific bases from any standard DashScope base.
Accepts the domain root, an OpenAI-compatible base, an Anthropic Messages
base, or a native `/api/v1` base for `dashscope.aliyuncs.com`,
`dashscope-intl.aliyuncs.com`, `dashscope-us.aliyuncs.com`, and
`{workspace}.{region}.maas.aliyuncs.com`. Returns None for any other host
or path so callers keep the caller-supplied base untouched.
"""
origin: Final = _dashscope_origin(api_base)
if origin is None:
return None
return DashScopeEndpoints(
origin=origin,
openai=f"{origin}{DASHSCOPE_CHAT_COMPATIBLE_PATH}",
anthropic=f"{origin}{DASHSCOPE_ANTHROPIC_PATH}",
native=f"{origin}{DASHSCOPE_NATIVE_V1_PATH}",
)
def is_dashscope_native_endpoint_url(api_base: str) -> bool:
return urlparse(api_base).path.rstrip("/").startswith(f"{DASHSCOPE_NATIVE_V1_PATH}/")
def resolve_dashscope_openai_api_base(
api_base: str | None,
env_var: str = "DASHSCOPE_API_BASE",
default: str = f"{DASHSCOPE_DEFAULT_ORIGIN}{DASHSCOPE_CHAT_COMPATIBLE_PATH}",
) -> str:
candidate: Final = api_base or get_secret_str(env_var)
if candidate and is_dashscope_native_endpoint_url(candidate):
return candidate
endpoints: Final = resolve_dashscope_endpoints(candidate)
if endpoints is not None:
return endpoints.openai
return candidate or default
def resolve_dashscope_anthropic_api_base(
api_base: str | None,
env_var: str = "DASHSCOPE_API_BASE",
default: str = f"{DASHSCOPE_DEFAULT_ORIGIN}{DASHSCOPE_ANTHROPIC_PATH}",
) -> str:
candidate: Final = api_base or get_secret_str(env_var)
endpoints: Final = resolve_dashscope_endpoints(candidate)
if endpoints is not None:
return endpoints.anthropic
return candidate or default
def dashscope_base_selects_anthropic_messages(api_base: str | None) -> bool:
candidate: Final = api_base or get_secret_str("DASHSCOPE_API_BASE")
if not candidate:
return False
path: Final = urlparse(candidate).path.rstrip("/")
if path.endswith(DASHSCOPE_ANTHROPIC_PATH) or f"{DASHSCOPE_ANTHROPIC_PATH}/" in path:
return True
return not path and resolve_dashscope_endpoints(candidate) is not None
def _chat_shaped_aliyun_origin(api_base: str | None) -> str | None:
if api_base is None:
return None
parsed: Final = urlparse(api_base)
@ -30,14 +144,24 @@ def _rerank_base_for_chat_shaped_base(api_base: str | None) -> str | None:
on_aliyun_host: Final = host == "aliyuncs.com" or host.endswith(".aliyuncs.com")
if not on_aliyun_host or parsed.path.rstrip("/") != DASHSCOPE_CHAT_COMPATIBLE_PATH:
return None
return f"{parsed.scheme}://{parsed.netloc}{DASHSCOPE_RERANK_PATH}"
return f"{parsed.scheme}://{parsed.netloc}"
def _is_compatible_api_base(api_base: str) -> bool:
path: Final = urlparse(api_base).path.rstrip("/")
return path == DASHSCOPE_COMPATIBLE_API_PATH or path.startswith(f"{DASHSCOPE_COMPATIBLE_API_PATH}/")
def resolve_dashscope_family_rerank_api_base(api_base: str | None, env_var: str, default_rerank_base: str) -> str:
remapped: Final = _rerank_base_for_chat_shaped_base(api_base)
if api_base is not None and remapped is None:
if api_base is None:
return get_secret_str(env_var) or default_rerank_base
if _is_compatible_api_base(api_base):
return api_base
return get_secret_str(env_var) or remapped or default_rerank_base
chat_origin: Final = _chat_shaped_aliyun_origin(api_base)
if chat_origin is not None:
return get_secret_str(env_var) or f"{chat_origin}{DASHSCOPE_RERANK_PATH}"
endpoints: Final = resolve_dashscope_endpoints(api_base)
return f"{endpoints.origin}{DASHSCOPE_RERANK_PATH}" if endpoints is not None else api_base
def get_dashscope_family_embedding_config(custom_llm_provider: str) -> "BaseEmbeddingConfig":

View file

@ -25,7 +25,7 @@ from litellm.secret_managers.main import get_secret_str
from litellm.types.llms.openai import AllEmbeddingInputValues, AllMessageValues
from litellm.types.utils import EmbeddingResponse, Usage
from ..common_utils import DashScopeError
from ..common_utils import DashScopeError, resolve_dashscope_openai_api_base
DEFAULT_API_BASE: Final = "https://dashscope.aliyuncs.com/compatible-mode/v1"
@ -91,7 +91,7 @@ class DashScopeEmbeddingConfig(BaseEmbeddingConfig):
return resolved_api_key
def _resolve_embedding_api_base(self, api_base: str | None) -> str:
return api_base or get_secret_str("DASHSCOPE_API_BASE") or DEFAULT_API_BASE
return resolve_dashscope_openai_api_base(api_base, default=DEFAULT_API_BASE)
def validate_environment(
self,

View file

@ -39,6 +39,13 @@ from litellm.types.llms.openai import (
)
from litellm.types.utils import ImageObject, ImageResponse
from ..common_utils import (
DASHSCOPE_CHAT_COMPATIBLE_PATH,
DASHSCOPE_IMAGE_PATH,
is_dashscope_native_endpoint_url,
resolve_dashscope_endpoints,
)
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj
from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer
@ -53,8 +60,6 @@ _DECODED_JSON: Final = TypeAdapter(object)
DEFAULT_API_BASE: Final = "https://dashscope-intl.aliyuncs.com/api/v1/services/aigc/multimodal-generation/generation"
CHAT_COMPATIBLE_MODE_PATH: Final = "/compatible-mode/v1"
# Maps OpenAI size strings (WxH) to DashScope size strings (W*H)
OPENAI_TO_DASHSCOPE_SIZE: Final[dict] = {
"256x256": "256*256",
@ -102,8 +107,14 @@ class DashScopeImageGenerationConfig(BaseImageGenerationConfig):
raise ValueError("DASHSCOPE_API_KEY is not set")
return resolved_api_key
def _resolve_image_api_base(self, image_api_base: str | None) -> str:
return image_api_base or get_secret_str("DASHSCOPE_API_BASE_IMAGE") or DEFAULT_API_BASE
def _image_api_base_env(self) -> str | None:
return get_secret_str("DASHSCOPE_API_BASE_IMAGE")
def _default_image_api_base(self) -> str:
return DEFAULT_API_BASE
def _configured_chat_api_base(self) -> str | None:
return get_secret_str("DASHSCOPE_API_BASE")
def get_complete_url(
self,
@ -114,10 +125,18 @@ class DashScopeImageGenerationConfig(BaseImageGenerationConfig):
litellm_params: dict,
stream: bool | None = None,
) -> str:
image_api_base: Final = (
api_base if api_base and not api_base.rstrip("/").endswith(CHAT_COMPATIBLE_MODE_PATH) else None
)
return self._resolve_image_api_base(image_api_base)
image_env_api_base: Final = self._image_api_base_env()
configured_api_base: Final = api_base or self._configured_chat_api_base()
if not configured_api_base or configured_api_base.rstrip("/").endswith(DASHSCOPE_CHAT_COMPATIBLE_PATH):
return image_env_api_base or self._default_image_api_base()
if is_dashscope_native_endpoint_url(configured_api_base):
return configured_api_base
endpoints: Final = resolve_dashscope_endpoints(configured_api_base)
if endpoints is None:
return configured_api_base
if not api_base and image_env_api_base:
return image_env_api_base
return f"{endpoints.origin}{DASHSCOPE_IMAGE_PATH}"
def validate_environment(
self,

View file

@ -0,0 +1,7 @@
"""
DashScope Anthropic-compatible Messages Module
"""
from .transformation import DashScopeAnthropicMessagesConfig
__all__ = ["DashScopeAnthropicMessagesConfig"]

View file

@ -0,0 +1,96 @@
"""
DashScope Anthropic-compatible Messages transformation config.
"""
from typing import ClassVar, Final
import litellm
from litellm.exceptions import AuthenticationError
from litellm.litellm_core_utils.url_utils import ensure_api_route
from litellm.llms.anthropic.messages.native.transformation import (
AnthropicMessagesConfig,
)
from litellm.secret_managers.main import get_secret_str, normalize_nonempty_secret_str
from ..common_utils import (
dashscope_base_selects_anthropic_messages,
resolve_dashscope_anthropic_api_base,
)
class DashScopeAnthropicMessagesConfig(AnthropicMessagesConfig):
"""
DashScope exposes an Anthropic-compatible Messages API at
``{origin}/apps/anthropic`` for every DashScope domain, including
``dashscope.aliyuncs.com`` and ``{workspace}.{region}.maas.aliyuncs.com``.
"""
_workload_identity_eligible: ClassVar[bool] = False
@property
def custom_llm_provider(self) -> str | None:
return "dashscope"
def should_strip_billing_metadata(self) -> bool:
return True
def handles_web_search_natively(self) -> bool:
return False
def serves_api_base(self, model: str, api_base: str | None) -> bool:
return dashscope_base_selects_anthropic_messages(api_base)
@staticmethod
def get_api_key(api_key: str | None = None) -> str | None:
return (
normalize_nonempty_secret_str(api_key)
or normalize_nonempty_secret_str(get_secret_str("DASHSCOPE_API_KEY"))
or normalize_nonempty_secret_str(litellm.api_key)
)
@staticmethod
def get_api_base(api_base: str | None = None) -> str:
return resolve_dashscope_anthropic_api_base(api_base)
def validate_anthropic_messages_environment(
self,
headers: dict[str, str], # mutable-ok: mirrors the legacy base override signature
model: str,
messages: list[object], # mutable-ok: mirrors the legacy base override signature
optional_params: dict[str, object], # mutable-ok: mirrors the legacy base override signature
litellm_params: dict[str, object], # mutable-ok: mirrors the legacy base override signature
api_key: str | None = None,
api_base: str | None = None,
) -> tuple[dict[str, str], str | None]: # mutable-ok: mirrors the base signature
resolved_api_key: Final = self.get_api_key(api_key=api_key)
if resolved_api_key is None:
present_headers: Final = {name.lower() for name in headers}
if "x-api-key" not in present_headers and "authorization" not in present_headers:
raise AuthenticationError(
message=(
"Missing DashScope API key. Set DASHSCOPE_API_KEY or pass api_key explicitly. "
"Anthropic credentials are never sent to a DashScope endpoint."
),
llm_provider="dashscope",
model=model,
)
return super().validate_anthropic_messages_environment( # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType, reportReturnType] # base signature uses an unparameterized dict
headers=headers,
model=model,
messages=messages,
optional_params=optional_params,
litellm_params=litellm_params,
api_key=resolved_api_key,
api_base=api_base,
)
def get_complete_url(
self,
api_base: str | None,
api_key: str | None,
model: str,
optional_params: dict[str, object], # mutable-ok: mirrors the base get_complete_url signature
litellm_params: dict[str, object], # mutable-ok: mirrors the base get_complete_url signature
stream: bool | None = None,
) -> str:
return ensure_api_route(self.get_api_base(api_base=api_base), "v1/messages")

View file

@ -3,7 +3,7 @@ from typing import Final
from litellm.secret_managers.main import get_secret_str
from .chat.transformation import DashScopeChatConfig
from .common_utils import resolve_dashscope_family_rerank_api_base
from .common_utils import resolve_dashscope_family_rerank_api_base, resolve_dashscope_openai_api_base
from .embed.transformation import DashScopeEmbeddingConfig
from .image_generation.transformation import DashScopeImageGenerationConfig
from .rerank.transformation import DashScopeRerankConfig
@ -38,7 +38,7 @@ class QwenAIPlatformChatConfig(DashScopeChatConfig):
return self._resolve_chat_api_base(api_base), _resolve_qwen_ai_platform_api_key(api_key)
def _resolve_chat_api_base(self, api_base: str | None) -> str:
return api_base or get_secret_str("QWEN_AI_PLATFORM_API_BASE") or QWEN_AI_PLATFORM_API_BASE
return resolve_dashscope_openai_api_base(api_base, "QWEN_AI_PLATFORM_API_BASE", QWEN_AI_PLATFORM_API_BASE)
class QwenAIPlatformEmbeddingConfig(DashScopeEmbeddingConfig):
@ -46,7 +46,7 @@ class QwenAIPlatformEmbeddingConfig(DashScopeEmbeddingConfig):
return _require_qwen_ai_platform_api_key(api_key)
def _resolve_embedding_api_base(self, api_base: str | None) -> str:
return api_base or get_secret_str("QWEN_AI_PLATFORM_API_BASE") or QWEN_AI_PLATFORM_API_BASE
return resolve_dashscope_openai_api_base(api_base, "QWEN_AI_PLATFORM_API_BASE", QWEN_AI_PLATFORM_API_BASE)
class QwenAIPlatformRerankConfig(DashScopeRerankConfig):
@ -63,8 +63,14 @@ class QwenAIPlatformImageGenerationConfig(DashScopeImageGenerationConfig):
def _resolve_api_key(self, api_key: str | None) -> str:
return _require_qwen_ai_platform_api_key(api_key)
def _resolve_image_api_base(self, image_api_base: str | None) -> str:
return image_api_base or get_secret_str("QWEN_AI_PLATFORM_API_BASE_IMAGE") or QWEN_AI_PLATFORM_IMAGE_API_BASE
def _image_api_base_env(self) -> str | None:
return get_secret_str("QWEN_AI_PLATFORM_API_BASE_IMAGE")
def _default_image_api_base(self) -> str:
return QWEN_AI_PLATFORM_IMAGE_API_BASE
def _configured_chat_api_base(self) -> str | None:
return get_secret_str("QWEN_AI_PLATFORM_API_BASE")
class QwenAIPlatformVideoConfig(DashScopeVideoConfig):

View file

@ -3,7 +3,7 @@ from typing import Final
from litellm.secret_managers.main import get_secret_str
from .chat.transformation import DashScopeChatConfig
from .common_utils import resolve_dashscope_family_rerank_api_base
from .common_utils import resolve_dashscope_family_rerank_api_base, resolve_dashscope_openai_api_base
from .embed.transformation import DashScopeEmbeddingConfig
from .image_generation.transformation import DashScopeImageGenerationConfig
from .rerank.transformation import DashScopeRerankConfig
@ -38,7 +38,7 @@ class QwenCloudChatConfig(DashScopeChatConfig):
return self._resolve_chat_api_base(api_base), _resolve_qwencloud_api_key(api_key)
def _resolve_chat_api_base(self, api_base: str | None) -> str:
return api_base or get_secret_str("QWENCLOUD_API_BASE") or QWENCLOUD_API_BASE
return resolve_dashscope_openai_api_base(api_base, "QWENCLOUD_API_BASE", QWENCLOUD_API_BASE)
class QwenCloudEmbeddingConfig(DashScopeEmbeddingConfig):
@ -46,7 +46,7 @@ class QwenCloudEmbeddingConfig(DashScopeEmbeddingConfig):
return _require_qwencloud_api_key(api_key)
def _resolve_embedding_api_base(self, api_base: str | None) -> str:
return api_base or get_secret_str("QWENCLOUD_API_BASE") or QWENCLOUD_API_BASE
return resolve_dashscope_openai_api_base(api_base, "QWENCLOUD_API_BASE", QWENCLOUD_API_BASE)
class QwenCloudRerankConfig(DashScopeRerankConfig):
@ -63,8 +63,14 @@ class QwenCloudImageGenerationConfig(DashScopeImageGenerationConfig):
def _resolve_api_key(self, api_key: str | None) -> str:
return _require_qwencloud_api_key(api_key)
def _resolve_image_api_base(self, image_api_base: str | None) -> str:
return image_api_base or get_secret_str("QWENCLOUD_API_BASE_IMAGE") or QWENCLOUD_IMAGE_API_BASE
def _image_api_base_env(self) -> str | None:
return get_secret_str("QWENCLOUD_API_BASE_IMAGE")
def _default_image_api_base(self) -> str:
return QWENCLOUD_IMAGE_API_BASE
def _configured_chat_api_base(self) -> str | None:
return get_secret_str("QWENCLOUD_API_BASE")
class QwenCloudVideoConfig(DashScopeVideoConfig):

View file

@ -38,6 +38,8 @@ from litellm.types.videos.utils import (
)
from litellm.utils import get_model_info
from ..common_utils import resolve_dashscope_endpoints
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
@ -206,6 +208,9 @@ def _parse_task_response(raw_response: httpx.Response) -> _TaskResponse:
def _normalized_api_base(api_base: str) -> str:
endpoints: Final = resolve_dashscope_endpoints(api_base)
if endpoints is not None:
return endpoints.origin
trimmed: Final = api_base.rstrip("/")
if trimmed.endswith(DASHSCOPE_COMPATIBLE_MODE_PATH):
return trimmed[: -len(DASHSCOPE_COMPATIBLE_MODE_PATH)]

View file

@ -38,6 +38,15 @@ else:
_EMPTY_FUNCTION: Final[Mapping[str, object]] = MappingProxyType({})
_JSON_VALUE: Final = TypeAdapter(object)
_CONTENT_PART: Final = TypeAdapter(Mapping[str, object])
def _file_part_data(part: Mapping[str, object]) -> str | None:
match part:
case {"type": "file", "file": {"file_data": str() as file_data}} if file_data:
return file_data
case _:
return None
def is_valid_json(value: str) -> bool:
@ -295,7 +304,7 @@ class GigaChatConfig(BaseConfig):
def _transform_list_content(self, content: Sequence) -> tuple[str, Sequence[object]]:
"""
Extract text and image attachments from a multimodal message content list.
Extract text, image and file attachments from a multimodal message content list.
Args:
content: List of content parts (OpenAI multimodal format)
@ -320,6 +329,10 @@ class GigaChatConfig(BaseConfig):
file_id = self._upload_image(upload_url)
if file_id:
attachments.append(file_id)
elif (file_data := _file_part_data(_CONTENT_PART.validate_python(part))) and (
uploaded_file_id := self._upload_image(file_data)
):
attachments.append(uploaded_file_id)
text: Final = "\n".join(texts) if texts else ""
return text, attachments

View file

@ -20,6 +20,7 @@ import litellm
from litellm._logging import verbose_logger
from litellm.constants import DEFAULT_OCI_CHAT_MAX_TOKENS
from litellm.litellm_core_utils.logging_utils import track_llm_api_timing
from litellm.litellm_core_utils.prompt_templates.common_utils import messages_with_file_parts_as_image_urls
from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
from litellm.litellm_core_utils.url_utils import join_api_base
from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException
@ -681,7 +682,9 @@ class OCIChatConfig(BaseConfig):
servingMode=serving_mode,
chatRequest=OCIChatRequestPayload(
apiFormat=vendor.value,
messages=adapt_messages_to_generic_oci_standard(messages),
messages=adapt_messages_to_generic_oci_standard(
list(messages_with_file_parts_as_image_urls(messages))
),
**self._get_optional_params(vendor, optional_params, model),
),
)

View file

@ -275,6 +275,10 @@ def build_output_token_limit_response(
return e.response.headers, (_aiter_once(chunk) if is_async else _iter_once(chunk))
def _is_loop_bound(client_initialization_params: Mapping[str, object]) -> bool:
return client_initialization_params.get("is_async") is not False
class BaseOpenAILLM:
"""
Base class for OpenAI LLMs for getting their httpx clients and SSL verification settings
@ -289,7 +293,9 @@ class BaseOpenAILLM:
client_initialization_params=client_initialization_params,
client_type=client_type,
)
_cached_client: Final = litellm.in_memory_llm_clients_cache.get_cache(_cache_key)
_cached_client: Final = litellm.in_memory_llm_clients_cache.get_cache(
_cache_key, loop_bound=_is_loop_bound(client_initialization_params)
)
return _cached_client
@staticmethod
@ -330,6 +336,7 @@ class BaseOpenAILLM:
value=openai_client,
ttl=_DEFAULT_TTL_FOR_HTTPX_CLIENTS,
litellm_owned_client=litellm_owned_client,
loop_bound=_is_loop_bound(client_initialization_params),
)
@staticmethod

View file

@ -5,18 +5,30 @@ This module provides guardrail translation support for OpenAI's text completion
The handler processes the 'prompt' parameter for guardrails.
"""
from collections.abc import Mapping
import json
from collections.abc import Mapping, Sequence
from itertools import chain
from typing import TYPE_CHECKING, Final
from litellm._logging import verbose_proxy_logger
from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation
from litellm.types.utils import GenericGuardrailAPIInputs
from litellm.llms.base_llm.guardrail_translation.base_translation import (
BaseTranslation,
StreamingScanKey,
StreamTransformSink,
)
from litellm.llms.base_llm.guardrail_translation.utils import (
blocked_chat_stream_usage,
stream_item_field,
stream_item_items,
)
from litellm.types.utils import GenericGuardrailAPIInputs, TextChoices, TextCompletionResponse, Usage
if TYPE_CHECKING:
from litellm.integrations.custom_guardrail import CustomGuardrail
from fastapi import HTTPException
from litellm.integrations.custom_guardrail import CustomGuardrail, ModifyResponseException
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth
from litellm.types.utils import TextCompletionResponse
class OpenAITextCompletionHandler(BaseTranslation):
@ -32,6 +44,9 @@ class OpenAITextCompletionHandler(BaseTranslation):
- A list of strings (for batch completions)
"""
delivers_ended_stream_rewrites = True
assembles_streamed_response = True
async def process_input_messages(
self,
data: dict,
@ -125,7 +140,7 @@ class OpenAITextCompletionHandler(BaseTranslation):
guardrail_to_apply: "CustomGuardrail",
litellm_logging_obj: "LiteLLMLoggingObj | None" = None,
user_api_key_dict: "UserAPIKeyAuth | None" = None,
request_data: dict | None = None,
request_data: dict[str, object] | None = None,
) -> "TextCompletionResponse":
"""
Process output response by applying guardrails to completion text.
@ -162,11 +177,7 @@ class OpenAITextCompletionHandler(BaseTranslation):
if "response" not in request_data:
request_data["response"] = response
# Add user API key metadata with prefixed keys
if "litellm_metadata" not in request_data:
user_metadata: Final = self.transform_user_api_key_dict_to_metadata(user_api_key_dict)
if user_metadata:
request_data["litellm_metadata"] = user_metadata
self.merge_user_api_key_metadata_into_request(request_data, user_api_key_dict)
inputs: Final = GenericGuardrailAPIInputs(texts=texts_to_check)
# Include model information from the response if available
@ -195,3 +206,168 @@ class OpenAITextCompletionHandler(BaseTranslation):
)
return response
async def process_output_streaming_response(
self,
responses_so_far: list[object], # mutable-ok: rewrites the caller's buffered chunks in place
guardrail_to_apply: "CustomGuardrail",
litellm_logging_obj: "LiteLLMLoggingObj | None" = None,
user_api_key_dict: "UserAPIKeyAuth | None" = None,
request_data: dict[str, object] | None = None, # mutable-ok: the hooks share and mutate the request payload
stream_transform_sink: StreamTransformSink | None = None,
deliver_ended_stream_rewrites: bool = False,
) -> Sequence[object]:
chunks: Final = _text_completion_chunks(responses_so_far)
if not self._check_streaming_has_ended(chunks):
return responses_so_far
assembled: Final = self._assemble_ended_stream(chunks)
pre_guardrail_texts: Final = _choice_texts(assembled)
await self.process_output_response(
response=assembled,
guardrail_to_apply=guardrail_to_apply,
litellm_logging_obj=litellm_logging_obj,
user_api_key_dict=user_api_key_dict,
request_data=request_data,
)
if not deliver_ended_stream_rewrites:
return responses_so_far
rewrites: Final = tuple(
after for before, after in zip(pre_guardrail_texts, _choice_texts(assembled)) if after != before
)
for choice_index, text in rewrites:
self._write_choice_text(chunks, choice_index, text)
return responses_so_far
def get_streaming_scan_key(self, responses_so_far: Sequence[object]) -> StreamingScanKey | None:
chunks: Final = _text_completion_chunks(responses_so_far)
if not self._check_streaming_has_ended(chunks):
return None
assembled: Final = self._assemble_ended_stream(chunks)
return StreamingScanKey(texts=tuple(text for _, text in _choice_texts(assembled)), stream_ended=True)
def released_stream_as_ended(self, responses_so_far: Sequence[object]) -> tuple[object, ...]:
chunks: Final = _text_completion_chunks(responses_so_far)
streamed_choices: Final = tuple(chain.from_iterable(chunk.choices for chunk in chunks))
finished_indices: Final = frozenset(
_choice_index(choice) for choice in streamed_choices if choice.finish_reason
)
unfinished_indices: Final = tuple(
dict.fromkeys(
_choice_index(choice) for choice in streamed_choices if _choice_index(choice) not in finished_indices
)
)
if not unfinished_indices:
return tuple(responses_so_far)
terminator: Final = TextCompletionResponse(
id=chunks[0].id,
model=chunks[0].model,
object="text_completion",
choices=[TextChoices(index=index, text="", finish_reason="length") for index in unfinished_indices],
)
return (*responses_so_far, terminator)
def build_block_sse_chunks(
self,
exc: "ModifyResponseException",
stream_started: bool = False,
responses_so_far: Sequence[object] | None = None,
) -> Sequence[bytes]:
streamed: Final = _text_completion_chunks(responses_so_far or ())
source: Final = streamed[0] if streamed else TextCompletionResponse(model=exc.model)
message_chunk: Final = TextCompletionResponse(
id=source.id,
created=source.created,
model=source.model or exc.model,
object="text_completion",
choices=[TextChoices(index=0, text=exc.message)],
).model_copy(update={"usage": None})
final_chunk: Final = TextCompletionResponse(
id=source.id,
created=source.created,
model=source.model or exc.model,
object="text_completion",
choices=[TextChoices(index=0, text="", finish_reason="content_filter")],
).model_copy(update={"usage": _blocked_stream_usage(exc.original_response)})
return tuple(
f"data: {chunk.model_dump_json(exclude_none=True)}\n\n".encode() for chunk in (message_chunk, final_chunk)
)
def build_stream_error_items(
self,
exc: "HTTPException",
responses_so_far: Sequence[object] | None = None,
) -> Sequence[bytes] | None:
from litellm.proxy.common_request_processing import sse_error_payload
_, error_obj = sse_error_payload(exc)
return (f'data: {{"error": {json.dumps(error_obj)}}}\n\n'.encode(),)
def _check_streaming_has_ended(self, responses_so_far: Sequence[object]) -> bool:
streamed_choices: Final = chain.from_iterable(stream_item_items(item, "choices") for item in responses_so_far)
return any(stream_item_field(choice, "finish_reason") is not None for choice in streamed_choices)
@staticmethod
def _assemble_ended_stream(responses_so_far: Sequence[TextCompletionResponse]) -> TextCompletionResponse:
streamed_choices: Final = tuple(chain.from_iterable(item.choices for item in responses_so_far))
choice_indices: Final = tuple(dict.fromkeys(_choice_index(choice) for choice in streamed_choices))
assembled_choices: Final = [
TextChoices(
index=index,
text="".join(_choice_text(choice) for choice in streamed_choices if _choice_index(choice) == index),
finish_reason=next(
(
choice.finish_reason
for choice in streamed_choices
if _choice_index(choice) == index and choice.finish_reason
),
None,
),
)
for index in choice_indices
]
first: Final = responses_so_far[0]
return TextCompletionResponse(
id=first.id,
created=first.created,
model=first.model,
object="text_completion",
choices=assembled_choices,
usage=next((item.usage for item in reversed(responses_so_far) if item.usage is not None), None),
)
@staticmethod
def _write_choice_text(responses_so_far: Sequence[TextCompletionResponse], choice_index: int, text: str) -> None:
streamed_choices: Final = chain.from_iterable(item.choices for item in responses_so_far)
head, *tail = (choice for choice in streamed_choices if _choice_index(choice) == choice_index)
head.text = text
for choice in tail:
if isinstance(stream_item_field(choice, "text"), str):
choice.text = ""
def _blocked_stream_usage(original_response: object) -> Usage | None:
prompt_tokens, completion_tokens = blocked_chat_stream_usage(original_response)
if isinstance(original_response, list) and prompt_tokens == completion_tokens == 0:
return None
return Usage(
prompt_tokens=prompt_tokens,
completion_tokens=completion_tokens,
total_tokens=prompt_tokens + completion_tokens,
)
def _text_completion_chunks(responses_so_far: Sequence[object]) -> tuple[TextCompletionResponse, ...]:
return tuple(item for item in responses_so_far if isinstance(item, TextCompletionResponse))
def _choice_text(choice: TextChoices) -> str:
text: Final = stream_item_field(choice, "text")
return text if isinstance(text, str) else ""
def _choice_index(choice: TextChoices) -> int:
return choice.index or 0
def _choice_texts(response: TextCompletionResponse) -> tuple[tuple[int, str], ...]:
return tuple((_choice_index(choice), _choice_text(choice)) for choice in response.choices)

View file

View file

@ -0,0 +1,415 @@
import hashlib
import json
import re
from collections.abc import Mapping, Sequence
from functools import partial
from typing import TYPE_CHECKING, Final, Literal, NoReturn
from urllib.parse import quote, urlsplit
import httpx
from pydantic import ConfigDict, Field, TypeAdapter, ValidationError
from litellm._logging import verbose_logger
from litellm.exceptions import AuthenticationError, BadRequestError
from litellm.litellm_core_utils.core_helpers import drops_unsupported_params
from litellm.litellm_core_utils.optional_imports import ensure_optional_import
from litellm.llms.base_llm.chat.transformation import BaseLLMException
from litellm.llms.base_llm.vector_store.transformation import (
BaseQueryEmbeddingVectorStoreConfig,
VectorStoreEmbeddingExecutor,
)
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, SignsRequestsWithAWS, sign_aws_json_post
from litellm.types.llms.base import LiteLLMBaseModel
from litellm.types.llms.bedrock import AwsAuthParams
from litellm.types.router import GenericLiteLLMParams
from litellm.types.utils import EmbeddingResponse
from litellm.types.vector_stores import (
BaseVectorStoreAuthCredentials,
VectorStoreCreateOptionalRequestParams,
VectorStoreIndexEndpoints,
VectorStoreResultContent,
VectorStoreSearchOptionalRequestParams,
VectorStoreSearchResponse,
VectorStoreSearchResult,
)
if TYPE_CHECKING:
from botocore.credentials import Credentials
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
DEFAULT_EMBEDDING_FIELD: Final = "embedding"
DEFAULT_TEXT_FIELD: Final = "text"
DEFAULT_MAX_NUM_RESULTS: Final = 10
MAX_MAX_NUM_RESULTS: Final = 50
_PARAM_PREFIX: Final = "opensearch_"
_AWS_HOSTS: Final = (
re.compile(r"\.(?P<region>[a-z0-9-]+)\.(?P<service>aoss|aoss-fips|es)\.amazonaws\.com(?:\.cn)?\Z"),
re.compile(r"\.(?P<service>aoss|aos)\.(?P<region>[a-z0-9-]+)\.on\.aws\Z"),
)
_ACCOUNT_ID_LABEL: Final = re.compile(r"\A\d{12}\.")
_QUERY_VECTOR: Final = TypeAdapter(tuple[float, ...])
_ATTRIBUTE_VALUE_TYPES: Final = (str, int, float, bool)
_SEARCH_ONLY_MESSAGE: Final = (
"The OpenSearch vector store is search-only. Create the k-NN index in OpenSearch, "
"then register it here with its name as opensearch_index_name or vector_store_id."
)
_PER_ACCOUNT_MESSAGE: Final = (
"api_base is an OpenSearch Serverless per-account endpoint, which needs a collection header the vector store "
"does not send. Use the collection endpoint instead, for example https://<collection-id>.<region>.aoss.amazonaws.com."
)
_SIGV4_ONLY_MESSAGE: Final = (
"The OpenSearch vector store signs every request with AWS SigV4 and does not support api_key. "
"Remove api_key and give the vector store AWS credentials instead."
)
_CREDENTIALS_HINT: Final = (
"Set aws_access_key_id and aws_secret_access_key, aws_role_name, or aws_profile_name on the vector store, "
"or run LiteLLM with an AWS identity."
)
def _bad_request(message: str) -> BadRequestError:
return BadRequestError(message=message, model=None, llm_provider="opensearch")
def _unauthenticated(reason: str) -> AuthenticationError:
return AuthenticationError(
message=f"No usable AWS credentials to sign the OpenSearch request ({reason}). {_CREDENTIALS_HINT}",
llm_provider="opensearch",
model=None,
)
def _usable_credentials(credentials: object) -> "Credentials":
from botocore.credentials import Credentials
if isinstance(credentials, Credentials):
return credentials
raise _unauthenticated("none were found")
class _OpenSearchParams(LiteLLMBaseModel):
model_config = ConfigDict(frozen=True, extra="ignore")
litellm_embedding_model: str | None = None
opensearch_embedding_field: str | None = None
opensearch_index_name: str | None = None
opensearch_text_field: str | None = None
@property
def embedding_field(self) -> str:
return self.opensearch_embedding_field or DEFAULT_EMBEDDING_FIELD
@property
def text_field(self) -> str:
return self.opensearch_text_field or DEFAULT_TEXT_FIELD
def index_name(self, vector_store_id: str) -> str:
return self.opensearch_index_name or vector_store_id
def require_embedding_model(self) -> str:
if not self.litellm_embedding_model:
raise _bad_request(
"litellm_embedding_model is required for the OpenSearch vector store. It must be the model that "
f"produced the vectors in '{self.embedding_field}', for example openai/text-embedding-3-small."
)
return self.litellm_embedding_model
class _SigningRegion(LiteLLMBaseModel):
model_config = ConfigDict(frozen=True, extra="ignore")
aws_region_name: str | None = None
class _SearchOptions(LiteLLMBaseModel):
model_config = ConfigDict(frozen=True, extra="ignore")
max_num_results: int | None = None
filters: Mapping[str, object] | None = None
ranking_options: Mapping[str, object] | None = None
rewrite_query: bool | None = None
@property
def unsupported(self) -> tuple[str, ...]:
requested: Final = (
("filters", self.filters),
("ranking_options", self.ranking_options),
("rewrite_query", self.rewrite_query),
)
return tuple(name for name, value in requested if value)
@property
def k(self) -> int:
if self.max_num_results is None:
return DEFAULT_MAX_NUM_RESULTS
if not 1 <= self.max_num_results <= MAX_MAX_NUM_RESULTS:
raise _bad_request(
f"max_num_results must be between 1 and {MAX_MAX_NUM_RESULTS}, got {self.max_num_results}"
)
return self.max_num_results
_KNOWN_PARAMS: Final = frozenset(name for name in _OpenSearchParams.model_fields if name.startswith(_PARAM_PREFIX))
class _LoggedSearch(LiteLLMBaseModel):
model_config = ConfigDict(frozen=True, extra="ignore")
query: str | tuple[str, ...] = ""
litellm_params: _OpenSearchParams = _OpenSearchParams()
@property
def search_query(self) -> str:
return BaseQueryEmbeddingVectorStoreConfig.query_text(self.query)
class _Hit(LiteLLMBaseModel):
model_config = ConfigDict(frozen=True)
id: str = Field(alias="_id")
score: float | None = Field(default=None, alias="_score")
source: Mapping[str, object] = Field(default_factory=dict, alias="_source")
class _Hits(LiteLLMBaseModel):
model_config = ConfigDict(frozen=True)
hits: Sequence[_Hit]
class _SearchResponse(LiteLLMBaseModel):
model_config = ConfigDict(frozen=True)
hits: _Hits
def aws_signing_target(url: str) -> tuple[Literal["aoss", "es"], str | None]:
hostname: Final = urlsplit(url).hostname or ""
match: Final = next(filter(None, (pattern.search(hostname) for pattern in _AWS_HOSTS)), None)
if match is None:
return "es", None
return ("aoss" if match["service"].startswith("aoss") else "es"), match["region"]
def _search_result(hit: _Hit, params: _OpenSearchParams) -> VectorStoreSearchResult:
text: Final = hit.source.get(params.text_field)
filename: Final = hit.source.get("filename")
return VectorStoreSearchResult(
score=hit.score,
content=[VectorStoreResultContent(type="text", text=text)] if isinstance(text, str) else [],
file_id=hit.id,
filename=filename if isinstance(filename, str) else None,
attributes={
key: value
for key, value in hit.source.items()
if key not in (params.text_field, params.embedding_field) and isinstance(value, _ATTRIBUTE_VALUE_TYPES)
},
)
class OpenSearchVectorStoreConfig(BaseQueryEmbeddingVectorStoreConfig, SignsRequestsWithAWS):
def __init__(self, aws: BaseAWSLLM | None = None) -> None:
self._aws: Final = aws if aws is not None else BaseAWSLLM()
def get_auth_credentials(self, litellm_params: Mapping[str, object]) -> BaseVectorStoreAuthCredentials:
return BaseVectorStoreAuthCredentials()
def get_vector_store_endpoints_by_type(self) -> VectorStoreIndexEndpoints:
return VectorStoreIndexEndpoints(read=[], write=[])
def validate_environment(
self, headers: Mapping[str, object], litellm_params: GenericLiteLLMParams | None
) -> dict[str, object]: # mutable-ok: the shared HTTP handler requires writable headers
params: Final = dict(litellm_params or {})
self._signing_credentials(params, self._signing_region(params, str(params.get("api_base") or "")))
return {**headers, "Content-Type": "application/json"}
def _signing_region(self, params: Mapping[str, object], api_base: str) -> str:
host_region: Final = aws_signing_target(api_base)[1]
return self._aws.get_aws_region_name_for_non_llm_api_calls(
host_region or _SigningRegion.model_validate(params).aws_region_name
)
def _signing_credentials(self, params: Mapping[str, object], region: str) -> "Credentials":
if params.get("api_key"):
raise _bad_request(_SIGV4_ONLY_MESSAGE)
auth: Final = AwsAuthParams.model_validate(params)
if bool(auth.aws_access_key_id) != bool(auth.aws_secret_access_key):
raise _unauthenticated("aws_access_key_id and aws_secret_access_key must be set together")
ensure_optional_import("botocore")
from botocore.exceptions import (
CredentialRetrievalError,
NoCredentialsError,
PartialCredentialsError,
ProfileNotFound,
)
try:
return _usable_credentials(self._aws.resolve_credentials(auth, region))
except CredentialRetrievalError as error:
verbose_logger.warning("OpenSearch vector store AWS credential source failed: %s", error)
raise _unauthenticated("the AWS credential source failed, see the LiteLLM logs") from None
except (NoCredentialsError, PartialCredentialsError, ProfileNotFound) as error:
raise _unauthenticated(str(error)) from None
def get_complete_url(self, api_base: str | None, litellm_params: Mapping[str, object]) -> str:
parsed: Final = urlsplit(api_base or "")
if parsed.scheme not in ("http", "https") or not parsed.hostname:
raise _bad_request(
"api_base must be the OpenSearch endpoint URL, for example "
"https://<collection-id>.<region>.aoss.amazonaws.com or https://<domain-endpoint>."
)
if _ACCOUNT_ID_LABEL.match(parsed.hostname) and aws_signing_target(api_base or "")[0] == "aoss":
raise _bad_request(_PER_ACCOUNT_MESSAGE)
return (api_base or "").rstrip("/")
def get_vector_store_url(self, api_base: str, encoded_vector_store_id: str) -> NoReturn:
raise _bad_request(_SEARCH_ONLY_MESSAGE)
def get_vector_store_list_url(self, api_base: str) -> NoReturn:
raise _bad_request(_SEARCH_ONLY_MESSAGE)
@staticmethod
def _params(
litellm_params: Mapping[str, object],
optional_params: VectorStoreSearchOptionalRequestParams,
extra_body: Mapping[str, object] | None,
) -> tuple[_OpenSearchParams, int]:
unknown: Final = sorted(
key for key in litellm_params if key.startswith(_PARAM_PREFIX) and key not in _KNOWN_PARAMS
)
if unknown:
raise _bad_request(
f"Unknown OpenSearch vector store parameter(s): {', '.join(unknown)}. "
f"Supported: {', '.join(sorted(_KNOWN_PARAMS))}."
)
if extra_body:
raise _bad_request("The OpenSearch vector store does not support extra_body.")
try:
options: Final = _SearchOptions.model_validate(optional_params)
except ValidationError as error:
raise _bad_request(f"Invalid vector store search parameters: {error}") from None
if options.unsupported and not drops_unsupported_params(litellm_params.get("drop_params")):
raise _bad_request(
f"The OpenSearch vector store does not support {', '.join(options.unsupported)} yet. "
"Remove it from the request, or set drop_params to drop it."
)
try:
return _OpenSearchParams.model_validate(litellm_params), options.k
except ValidationError:
raise _bad_request(
"opensearch_embedding_field, opensearch_index_name, opensearch_text_field, and litellm_embedding_model "
"must be strings."
) from None
@staticmethod
def _search_request(
vector_store_id: str,
embedding: EmbeddingResponse,
k: int,
params: _OpenSearchParams,
api_base: str,
) -> tuple[str, dict[str, object]]: # mutable-ok: the provider contract returns a writable JSON request body
vector: Final = _QUERY_VECTOR.validate_python(embedding.data[0]["embedding"]) # pyright: ignore[reportUnknownMemberType] # EmbeddingResponse.data is an untyped list, the vector is validated here
return (
f"{api_base}/{quote(params.index_name(vector_store_id), safe='')}/_search",
{
"size": k,
"query": {"knn": {params.embedding_field: {"vector": vector, "k": k}}},
"_source": {"excludes": [params.embedding_field]},
},
)
def transform_search_vector_store_request(
self,
vector_store_id: str,
query: str | Sequence[str],
vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams,
api_base: str,
litellm_logging_obj: "LiteLLMLoggingObj",
litellm_params: Mapping[str, object],
extra_body: Mapping[str, object] | None = None,
embedding_executor: VectorStoreEmbeddingExecutor | None = None,
) -> tuple[str, dict[str, object]]: # mutable-ok: the provider contract returns a writable JSON request body
params, k = self._params(litellm_params, vector_store_search_optional_params, extra_body)
embedding: Final = self.query_embedding_executor(embedding_executor, None).embed(
params.require_embedding_model(), self.query_text(query), self.query_embedding_configuration(litellm_params)
)
return self._search_request(vector_store_id, embedding, k, params, api_base)
async def atransform_search_vector_store_request(
self,
vector_store_id: str,
query: str | Sequence[str],
vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams,
api_base: str,
litellm_logging_obj: "LiteLLMLoggingObj",
litellm_params: Mapping[str, object],
extra_body: Mapping[str, object] | None = None,
embedding_executor: VectorStoreEmbeddingExecutor | None = None,
) -> tuple[str, dict[str, object]]: # mutable-ok: the provider contract returns a writable JSON request body
params, k = self._params(litellm_params, vector_store_search_optional_params, extra_body)
embedding: Final = await self.query_embedding_executor(embedding_executor, None).aembed(
params.require_embedding_model(), self.query_text(query), self.query_embedding_configuration(litellm_params)
)
return self._search_request(vector_store_id, embedding, k, params, api_base)
def sign_request(
self,
headers: Mapping[str, str],
optional_params: Mapping[str, object],
request_data: Mapping[str, object],
api_base: str,
api_key: str | None = None,
) -> tuple[dict[str, str], bytes]: # mutable-ok: the shared HTTP handler requires writable headers
region: Final = self._signing_region(optional_params, api_base)
body: Final = json.dumps(request_data)
signed: Final = sign_aws_json_post(
get_credentials=partial(self._signing_credentials, optional_params, region),
service_name=aws_signing_target(api_base)[0],
aws_region_name=region,
url=api_base,
body=body,
headers={
"Content-Type": "application/json",
"x-amz-content-sha256": hashlib.sha256(body.encode()).hexdigest(),
},
)
signed_headers: Final = dict(signed.headers.items())
signed_names: Final = frozenset(name.lower() for name in signed_headers)
return {
**{name: value for name, value in headers.items() if name.lower() not in signed_names},
**signed_headers,
}, body.encode()
def transform_search_vector_store_response(
self, response: httpx.Response, litellm_logging_obj: "LiteLLMLoggingObj"
) -> VectorStoreSearchResponse:
if not response.is_success:
raise BaseLLMException(status_code=response.status_code, message=response.text, headers=response.headers)
try:
hits: Final = _SearchResponse.model_validate_json(response.content).hits.hits
except ValidationError:
raise BaseLLMException(
status_code=502,
message="OpenSearch answered the search without a hits.hits list.",
headers=response.headers,
) from None
logged: Final = _LoggedSearch.model_validate(litellm_logging_obj.model_call_details)
return VectorStoreSearchResponse(
object="vector_store.search_results.page",
search_query=logged.search_query,
data=[_search_result(hit, logged.litellm_params) for hit in hits],
)
def validate_create_vector_store(self) -> NoReturn:
raise _bad_request(_SEARCH_ONLY_MESSAGE)
def transform_create_vector_store_request(
self, vector_store_create_optional_params: VectorStoreCreateOptionalRequestParams, api_base: str
) -> NoReturn:
raise _bad_request(_SEARCH_ONLY_MESSAGE)
def transform_create_vector_store_response(self, response: httpx.Response) -> NoReturn:
raise _bad_request(_SEARCH_ONLY_MESSAGE)

View file

@ -11,6 +11,7 @@ import httpx
from pydantic import BaseModel
import litellm
from litellm.litellm_core_utils.prompt_templates.common_utils import message_with_file_parts_as_image_urls
from litellm.litellm_core_utils.url_utils import join_api_base
from litellm.types.llms.openai import AllMessageValues
from litellm.types.utils import ModelResponse
@ -104,7 +105,7 @@ def _message_role_mapping(role: str) -> type[ChatMessage]:
def _messages_to_sap_template(messages: list[AllMessageValues]) -> list:
template: Final = []
for message in messages:
folded_message = _fold_message_cache_control(message)
folded_message = _fold_message_cache_control(message_with_file_parts_as_image_urls(message))
message_role_class = _message_role_mapping(str(folded_message["role"]))
validated_message = validate_dict(folded_message, message_role_class)
template.append(validated_message)

View file

@ -112,7 +112,7 @@ def _convert_image_url_to_anthropic(block: Mapping[str, object]) -> object:
pdf/text data URIs become document blocks rather than malformed image blocks.
"""
image_url: Final = block.get("image_url")
url: Final = image_url if isinstance(image_url, str) else _image_url_field(image_url, "url")
url: Final = image_url if isinstance(image_url, str) else _str_field(image_url, "url")
if not url:
return block
@ -121,18 +121,31 @@ def _convert_image_url_to_anthropic(block: Mapping[str, object]) -> object:
if select_anthropic_content_block_type_for_file(_data_uri_media_type(url)) == "document"
else create_anthropic_image_param(
image_url if isinstance(image_url, dict) else url,
format=_image_url_field(image_url, "format"),
format=_str_field(image_url, "format"),
is_bedrock_invoke=True,
)
)
return _with_cache_control(converted, block)
def _convert_file_to_anthropic(block: Mapping[str, object]) -> object:
file_data: Final = _str_field(block.get("file"), "file_data")
if not file_data:
return block
return _with_cache_control(
anthropic_process_openai_file_message({"type": "file", "file": {"file_data": file_data}}), block
)
def _with_cache_control(converted: Mapping[str, object], block: Mapping[str, object]) -> Mapping[str, object]:
cache_control: Final = block.get("cache_control")
if cache_control is None:
return converted
return {**converted, "cache_control": cache_control}
def _image_url_field(image_url: object, key: str) -> str | None:
value: Final = image_url.get(key) if isinstance(image_url, dict) else None
def _str_field(container: object, key: str) -> str | None:
value: Final = container.get(key) if isinstance(container, dict) else None
return value if isinstance(value, str) else None
@ -141,15 +154,23 @@ def _data_uri_media_type(url: str) -> str:
return match.group(1) if match else ""
def _convert_image_url_blocks_to_anthropic(content: object) -> object:
def _convert_media_block_to_anthropic(block: object) -> object:
mapping: Final = as_str_mapping(block)
if mapping is None:
return block
match mapping.get("type"):
case "image_url":
return _convert_image_url_to_anthropic(mapping)
case "file":
return _convert_file_to_anthropic(mapping)
case _:
return block
def _convert_media_blocks_to_anthropic(content: object) -> object:
if not isinstance(content, list):
return content
return [
_convert_image_url_to_anthropic(block)
if isinstance(block, Mapping) and block.get("type") == "image_url"
else block
for block in content
]
return [_convert_media_block_to_anthropic(block) for block in content]
def _convert_tool_result_to_anthropic(
@ -391,7 +412,7 @@ class SnowflakeConfig(SnowflakeBaseConfig, OpenAIGPTConfig):
conversation.append(
{
"role": role,
"content": _convert_image_url_blocks_to_anthropic(content),
"content": _convert_media_blocks_to_anthropic(content),
}
)

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