mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-11 08:38:24 +00:00
Merge remote-tracking branch 'origin/main' into litellm_model_info_wildcard_routed_price
This commit is contained in:
commit
1c54f0784f
573 changed files with 37964 additions and 6119 deletions
|
|
@ -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
119
.github/scripts/assert_secret_absent.py
vendored
Normal 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))
|
||||
9
.github/scripts/smoke_test_native_wheel.py
vendored
9
.github/scripts/smoke_test_native_wheel.py
vendored
|
|
@ -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")
|
||||
|
|
|
|||
21
.github/workflows/cost-map-guard.yml
vendored
21
.github/workflows/cost-map-guard.yml
vendored
|
|
@ -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"
|
||||
|
|
|
|||
2
.github/workflows/image-scan.yml
vendored
2
.github/workflows/image-scan.yml
vendored
|
|
@ -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
199
.github/workflows/test-integration.yml
vendored
Normal 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":"32","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
|
||||
4
.github/workflows/test-linting.yml
vendored
4
.github/workflows/test-linting.yml
vendored
|
|
@ -376,10 +376,10 @@ jobs:
|
|||
exit 0
|
||||
fi
|
||||
files="$(git diff --name-only "$base" HEAD)"
|
||||
if grep -Eq '^(litellm/(proxy|types)/|ui/litellm-dashboard/(src/lib/http/schema\.d\.ts|scripts/gen-api-types\.mjs|package(-lock)?\.json)$|\.github/workflows/test-linting\.yml$)' <<< "$files"; then
|
||||
if grep -Eq '^(litellm/|enterprise/|(pyproject\.toml|uv\.lock)$|ui/litellm-dashboard/(src/lib/http/schema\.d\.ts|scripts/gen-api-types\.mjs|package(-lock)?\.json)$|\.github/workflows/test-linting\.yml$)' <<< "$files"; then
|
||||
echo "relevant=true" >> "$GITHUB_OUTPUT"
|
||||
else
|
||||
echo "No proxy, types or generator changes in this pull request, nothing to verify."
|
||||
echo "No litellm, enterprise, dependency or generator changes in this pull request, nothing to verify."
|
||||
echo "relevant=false" >> "$GITHUB_OUTPUT"
|
||||
fi
|
||||
|
||||
|
|
|
|||
|
|
@ -12,6 +12,7 @@ on:
|
|||
paths:
|
||||
- "terraform/provider/**"
|
||||
- "litellm/proxy/**"
|
||||
- "enterprise/**"
|
||||
- ".github/workflows/test-terraform-provider.yml"
|
||||
|
||||
permissions:
|
||||
|
|
|
|||
1
.github/workflows/test-unit.yml
vendored
1
.github/workflows/test-unit.yml
vendored
|
|
@ -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
|
||||
|
|
|
|||
50
README.md
50
README.md
|
|
@ -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>
|
||||
|
|
|
|||
|
|
@ -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 && \
|
||||
|
|
|
|||
11
codecov.yaml
11
codecov.yaml
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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)),
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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`
|
||||
|
|
|
|||
|
|
@ -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."
|
||||
|
|
|
|||
184
litellm-proxy-extras/litellm_proxy_extras/migration_dry_run.py
Normal file
184
litellm-proxy-extras/litellm_proxy_extras/migration_dry_run.py
Normal 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 "")
|
||||
)
|
||||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
1
litellm-rust/Cargo.lock
generated
1
litellm-rust/Cargo.lock
generated
|
|
@ -4345,6 +4345,7 @@ dependencies = [
|
|||
"litellm-inference-transcription",
|
||||
"litellm-llms",
|
||||
"litellm-llms-types",
|
||||
"litellm-model-catalog",
|
||||
"litellm-router-types",
|
||||
"litellm-secrets",
|
||||
"litellm-secrets-aws",
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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>`
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
26
litellm-rust/crates/model-catalog/src/bin/check-registry.rs
Normal file
26
litellm-rust/crates/model-catalog/src/bin/check-registry.rs
Normal 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
|
||||
}
|
||||
}
|
||||
|
|
@ -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)
|
||||
);
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
30
litellm-rust/crates/python-bridge/src/model_registry.rs
Normal file
30
litellm-rust/crates/python-bridge/src/model_registry.rs
Normal 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);
|
||||
});
|
||||
}
|
||||
}
|
||||
|
|
@ -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 #######
|
||||
|
|
@ -685,6 +685,7 @@ moonshot_models: Set = set()
|
|||
publicai_models: Set = set()
|
||||
darkbloom_models: Set = set()
|
||||
openinfer_models: Final[Set[str]] = set()
|
||||
kyma_models: Final[set[str]] = set() # mutable-ok: populated when the shared model catalog loads
|
||||
streamlake_models: Final[Set[str]] = set()
|
||||
zerogpu_models: Final[Set[str]] = set()
|
||||
v0_models: Set = set()
|
||||
|
|
@ -981,6 +982,8 @@ def _populate_provider_model_sets(model_cost_map: Mapping[str, Mapping[str, obje
|
|||
darkbloom_models.add(key)
|
||||
elif value.get("litellm_provider") == "openinfer":
|
||||
openinfer_models.add(key)
|
||||
elif value.get("litellm_provider") == "kyma":
|
||||
kyma_models.add(key)
|
||||
elif value.get("litellm_provider") == "streamlake":
|
||||
streamlake_models.add(key)
|
||||
elif value.get("litellm_provider") == "zerogpu":
|
||||
|
|
@ -1159,6 +1162,7 @@ model_list = list(
|
|||
| publicai_models
|
||||
| darkbloom_models
|
||||
| openinfer_models
|
||||
| kyma_models
|
||||
| streamlake_models
|
||||
| zerogpu_models
|
||||
| v0_models
|
||||
|
|
@ -1290,6 +1294,7 @@ def _build_models_by_provider() -> dict:
|
|||
"publicai": publicai_models,
|
||||
"darkbloom": darkbloom_models,
|
||||
"openinfer": openinfer_models,
|
||||
"kyma": kyma_models,
|
||||
"streamlake": streamlake_models,
|
||||
"zerogpu": zerogpu_models,
|
||||
"v0": v0_models,
|
||||
|
|
|
|||
|
|
@ -36,6 +36,8 @@ REDIS_FAMILIES_METADATA_KEY: Final = "families"
|
|||
|
||||
_service_caller: Final[ContextVar[str | None]] = ContextVar("service_caller", default=None)
|
||||
|
||||
_mirrored_call: Final[ContextVar[bool]] = ContextVar("mirrored_call", default=False)
|
||||
|
||||
|
||||
@contextmanager
|
||||
def internal_sub_call() -> Generator[None]:
|
||||
|
|
@ -123,6 +125,30 @@ def current_service_caller() -> str | None:
|
|||
return _service_caller.get()
|
||||
|
||||
|
||||
@contextmanager
|
||||
def mirrored_call() -> Generator[None]:
|
||||
"""A silent_model mirror of a request: every litellm call inside, its retries and fallbacks included."""
|
||||
token: Final = _mirrored_call.set(True)
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
_restore(_mirrored_call, token)
|
||||
|
||||
|
||||
def in_mirrored_call() -> bool:
|
||||
return _mirrored_call.get()
|
||||
|
||||
|
||||
@contextmanager
|
||||
def outside_mirrored_call() -> Generator[None]:
|
||||
"""Logging a mirrored call: the calls its callbacks start are not mirrors themselves."""
|
||||
token: Final = _mirrored_call.set(False)
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
_restore(_mirrored_call, token)
|
||||
|
||||
|
||||
@contextmanager
|
||||
def pinned_billing_time(moment: datetime) -> Generator[None]:
|
||||
"""Price every rate lookup inside this block at ``moment`` rather than at each one's own clock read."""
|
||||
|
|
|
|||
325
litellm/anthropic_beta_gated_fields.py
Normal file
325
litellm/anthropic_beta_gated_fields.py
Normal file
|
|
@ -0,0 +1,325 @@
|
|||
"""Body fields that an ``anthropic-beta`` value gates.
|
||||
|
||||
``anthropic_beta_headers_config.json`` says which beta each provider leg
|
||||
carries. When a leg drops a beta, the request fields that beta unlocks
|
||||
would make the provider 400. Under ``drop_params`` the same request degrades on
|
||||
that leg by dropping those fields too; without it the request goes out as sent
|
||||
and the provider answers its own 400, like any unsupported param. Every helper
|
||||
here is copy-on-write: the caller's containers are returned untouched when
|
||||
nothing is stripped, and a shallow copy is returned otherwise, so a fallback
|
||||
leg still sends the full request.
|
||||
"""
|
||||
|
||||
from collections.abc import Mapping, Sequence, Set
|
||||
from dataclasses import dataclass
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
|
||||
from pydantic import TypeAdapter, ValidationError
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.litellm_core_utils.core_helpers import drops_unsupported_params
|
||||
from litellm.types.llms.anthropic import (
|
||||
ANTHROPIC_BETA_HEADER_VALUES,
|
||||
ANTHROPIC_MID_CONVERSATION_OUTPUT_CONFIG_BETA_HEADER,
|
||||
ANTHROPIC_MID_CONVERSATION_TOOL_CHANGES_BETA_HEADER,
|
||||
ANTHROPIC_THINKING_BINDING_CONTROLS_BETA_HEADER,
|
||||
ANTHROPIC_THINKING_DISPLAY_UPDATES_BETA_HEADER,
|
||||
)
|
||||
|
||||
ANTHROPIC_MID_CONVERSATION_SYSTEM_CLEAR_AT_BETA_HEADER: Final = "mid-conversation-system-clear-at-2026-08-21"
|
||||
ANTHROPIC_INLINE_TOOLS_BETA_HEADER: Final = "inline-tools-2026-09-15"
|
||||
|
||||
THINKING_BLOCK_BINDING: Final = "thinking.block_binding"
|
||||
THINKING_DISPLAY_UPDATES: Final = "thinking.display=updates"
|
||||
SYSTEM_CLEAR_AT: Final = "messages[].system.clear_at"
|
||||
SYSTEM_OUTPUT_CONFIG: Final = "messages[].system.output_config"
|
||||
SYSTEM_TOOL_CHANGES: Final = "messages[].system.content[].tool_addition|tool_removal"
|
||||
SYSTEM_MESSAGE_EMPTIED: Final = "messages[].system (emptied by the strip)"
|
||||
|
||||
_TOOL_CHANGE_BLOCK_TYPES: Final = frozenset({"tool_addition", "tool_removal"})
|
||||
|
||||
BETA_GATED_FIELDS: Final[Mapping[str, tuple[str, ...]]] = MappingProxyType(
|
||||
{
|
||||
ANTHROPIC_THINKING_BINDING_CONTROLS_BETA_HEADER: (THINKING_BLOCK_BINDING,),
|
||||
ANTHROPIC_THINKING_DISPLAY_UPDATES_BETA_HEADER: (THINKING_DISPLAY_UPDATES,),
|
||||
ANTHROPIC_MID_CONVERSATION_SYSTEM_CLEAR_AT_BETA_HEADER: (SYSTEM_CLEAR_AT,),
|
||||
ANTHROPIC_MID_CONVERSATION_OUTPUT_CONFIG_BETA_HEADER: (SYSTEM_OUTPUT_CONFIG,),
|
||||
ANTHROPIC_BETA_HEADER_VALUES.PER_TURN_CONTROL_2026_07_01.value: (SYSTEM_OUTPUT_CONFIG,),
|
||||
ANTHROPIC_MID_CONVERSATION_TOOL_CHANGES_BETA_HEADER: (SYSTEM_TOOL_CHANGES,),
|
||||
ANTHROPIC_INLINE_TOOLS_BETA_HEADER: (SYSTEM_TOOL_CHANGES,),
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _betas_unlocking(field: str) -> frozenset[str]:
|
||||
return frozenset(beta for beta, fields in BETA_GATED_FIELDS.items() if field in fields)
|
||||
|
||||
|
||||
_FIELD_BETAS: Final[Mapping[str, frozenset[str]]] = MappingProxyType(
|
||||
{
|
||||
field: _betas_unlocking(field) for fields in BETA_GATED_FIELDS.values() for field in fields
|
||||
} # comprehension-ok: flattening the beta-to-fields map into a field-to-betas index
|
||||
)
|
||||
|
||||
|
||||
def field_is_gated(field: str, active_betas: Set[str]) -> bool:
|
||||
"""True when no active beta unlocks ``field``, so the field must be stripped."""
|
||||
return _FIELD_BETAS[field].isdisjoint(active_betas)
|
||||
|
||||
|
||||
_MAPPING_ADAPTER: Final = TypeAdapter(Mapping[str, object])
|
||||
_SEQUENCE_ADAPTER: Final = TypeAdapter(tuple[object, ...])
|
||||
_HEADERS_ADAPTER: Final = TypeAdapter(Mapping[str, str])
|
||||
|
||||
|
||||
def _as_mapping(value: object) -> Mapping[str, object] | None:
|
||||
try:
|
||||
return _MAPPING_ADAPTER.validate_python(value, strict=True)
|
||||
except ValidationError:
|
||||
return None
|
||||
|
||||
|
||||
def _as_sequence(value: object) -> tuple[object, ...] | None:
|
||||
try:
|
||||
return _SEQUENCE_ADAPTER.validate_python(value)
|
||||
except ValidationError:
|
||||
return None
|
||||
|
||||
|
||||
def _is_tool_change_block(block: object) -> bool:
|
||||
mapping: Final = _as_mapping(block)
|
||||
block_type: Final = mapping.get("type") if mapping is not None else None
|
||||
return isinstance(block_type, str) and block_type in _TOOL_CHANGE_BLOCK_TYPES
|
||||
|
||||
|
||||
def _header_beta_values(headers: Mapping[str, str]) -> frozenset[str]:
|
||||
joined: Final = ",".join(value for key, value in headers.items() if key.lower() == "anthropic-beta")
|
||||
return frozenset(piece.strip() for piece in joined.split(",") if piece.strip())
|
||||
|
||||
|
||||
def _validated_headers(headers: object) -> Mapping[str, str]:
|
||||
try:
|
||||
return _HEADERS_ADAPTER.validate_python(headers)
|
||||
except ValidationError:
|
||||
return MappingProxyType({})
|
||||
|
||||
|
||||
def active_beta_values(headers: object, body_betas: object = ()) -> frozenset[str]:
|
||||
"""The ``anthropic-beta`` values a request will send: the comma-joined header (key matched
|
||||
case-insensitively) plus the body's ``anthropic_beta`` list on legs that carry betas in the body."""
|
||||
body_values: Final = _as_sequence(body_betas) or ()
|
||||
body_betas_as_strings: Final = frozenset(beta for beta in body_values if isinstance(beta, str))
|
||||
return _header_beta_values(_validated_headers(headers)) | body_betas_as_strings
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class StrippedThinking:
|
||||
thinking: object
|
||||
stripped: tuple[str, ...]
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class StrippedMessages:
|
||||
messages: Sequence[object]
|
||||
stripped: tuple[str, ...]
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class StrippedRequest:
|
||||
request: Mapping[str, object]
|
||||
stripped: tuple[str, ...]
|
||||
|
||||
|
||||
def strip_beta_gated_thinking(thinking: object, active_betas: Set[str], *, drop_params: object) -> StrippedThinking:
|
||||
"""Drop ``thinking`` sub-fields whose beta is absent. Returns the same object when nothing is gated."""
|
||||
mapping: Final = _as_mapping(thinking)
|
||||
if mapping is None or not drops_unsupported_params(drop_params):
|
||||
return StrippedThinking(thinking=thinking, stripped=())
|
||||
binding_gated: Final = "block_binding" in mapping and field_is_gated(THINKING_BLOCK_BINDING, active_betas)
|
||||
display_gated: Final = mapping.get("display") == "updates" and field_is_gated(
|
||||
THINKING_DISPLAY_UPDATES, active_betas
|
||||
)
|
||||
if not binding_gated and not display_gated:
|
||||
return StrippedThinking(thinking=thinking, stripped=())
|
||||
gated_keys: Final = frozenset(
|
||||
{*(("block_binding",) if binding_gated else ()), *(("display",) if display_gated else ())}
|
||||
)
|
||||
stripped: Final = (
|
||||
*((THINKING_BLOCK_BINDING,) if binding_gated else ()),
|
||||
*((THINKING_DISPLAY_UPDATES,) if display_gated else ()),
|
||||
)
|
||||
return StrippedThinking(
|
||||
thinking={key: value for key, value in mapping.items() if key not in gated_keys},
|
||||
stripped=stripped,
|
||||
)
|
||||
|
||||
|
||||
def _strip_beta_gated_system_message(message: Mapping[str, object], active_betas: Set[str]) -> StrippedMessages:
|
||||
"""One mid-conversation ``role: system`` message. ``messages`` holds the rewritten message, or nothing
|
||||
when the strip left it with neither content nor ``output_config``, since such a system turn is itself a 400."""
|
||||
clear_at_gated: Final = "clear_at" in message and field_is_gated(SYSTEM_CLEAR_AT, active_betas)
|
||||
output_config_gated: Final = "output_config" in message and field_is_gated(SYSTEM_OUTPUT_CONFIG, active_betas)
|
||||
content: Final = message.get("content")
|
||||
blocks: Final = _as_sequence(content) if not isinstance(content, str) else None
|
||||
tool_changes_gated: Final = (
|
||||
blocks is not None
|
||||
and any(_is_tool_change_block(block) for block in blocks)
|
||||
and field_is_gated(SYSTEM_TOOL_CHANGES, active_betas)
|
||||
)
|
||||
if not clear_at_gated and not output_config_gated and not tool_changes_gated:
|
||||
return StrippedMessages(messages=(message,), stripped=())
|
||||
kept_content: Final[object] = (
|
||||
[block for block in blocks if not _is_tool_change_block(block)]
|
||||
if tool_changes_gated and blocks is not None
|
||||
else content
|
||||
)
|
||||
gated_keys: Final = frozenset(
|
||||
{*(("clear_at",) if clear_at_gated else ()), *(("output_config",) if output_config_gated else ())}
|
||||
)
|
||||
rewritten: Final = {
|
||||
key: (kept_content if key == "content" else value) for key, value in message.items() if key not in gated_keys
|
||||
}
|
||||
stripped: Final = (
|
||||
*((SYSTEM_CLEAR_AT,) if clear_at_gated else ()),
|
||||
*((SYSTEM_OUTPUT_CONFIG,) if output_config_gated else ()),
|
||||
*((SYSTEM_TOOL_CHANGES,) if tool_changes_gated else ()),
|
||||
)
|
||||
if kept_content in (None, "", []) and "output_config" not in rewritten:
|
||||
return StrippedMessages(messages=(), stripped=(*stripped, SYSTEM_MESSAGE_EMPTIED))
|
||||
return StrippedMessages(messages=(rewritten,), stripped=stripped)
|
||||
|
||||
|
||||
def _strip_beta_gated_message(message: object, active_betas: Set[str]) -> StrippedMessages:
|
||||
mapping: Final = _as_mapping(message)
|
||||
if mapping is None or mapping.get("role") != "system":
|
||||
return StrippedMessages(messages=(message,), stripped=())
|
||||
return _strip_beta_gated_system_message(mapping, active_betas)
|
||||
|
||||
|
||||
def _ends_on_assistant(messages: Sequence[object]) -> bool:
|
||||
last: Final = _as_mapping(messages[-1]) if messages else None
|
||||
return last is not None and last.get("role") == "assistant"
|
||||
|
||||
|
||||
def strip_beta_gated_system_messages(
|
||||
messages: Sequence[object], active_betas: Set[str], *, drop_params: object
|
||||
) -> StrippedMessages:
|
||||
"""Drop the gated fields of every mid-conversation ``role: system`` message. Returns the same
|
||||
sequence when nothing is gated, or when dropping an emptied system turn would leave the conversation
|
||||
ending on an assistant turn, which the model would read as a prefill to continue."""
|
||||
if not drops_unsupported_params(drop_params):
|
||||
return StrippedMessages(messages=messages, stripped=())
|
||||
per_message: Final = tuple(_strip_beta_gated_message(message, active_betas) for message in messages)
|
||||
stripped: Final = tuple(
|
||||
name for result in per_message for name in result.stripped
|
||||
) # comprehension-ok: flattening per-message strip reports
|
||||
if not stripped:
|
||||
return StrippedMessages(messages=messages, stripped=())
|
||||
kept: Final = [
|
||||
message for result in per_message for message in result.messages
|
||||
] # comprehension-ok: flattening per-message results
|
||||
if _ends_on_assistant(kept) and not _ends_on_assistant(messages):
|
||||
return StrippedMessages(messages=messages, stripped=())
|
||||
return StrippedMessages(messages=kept, stripped=stripped)
|
||||
|
||||
|
||||
def strip_beta_gated_request_fields(
|
||||
request: Mapping[str, object], active_betas: Set[str], *, drop_params: object
|
||||
) -> StrippedRequest:
|
||||
"""Apply both strips to an Anthropic ``/v1/messages`` shaped body. Returns the same mapping when
|
||||
nothing is gated, otherwise a shallow copy with ``thinking`` and ``messages`` replaced."""
|
||||
thinking: Final = strip_beta_gated_thinking(request.get("thinking"), active_betas, drop_params=drop_params)
|
||||
messages: Final = request.get("messages")
|
||||
message_sequence: Final = _as_sequence(messages) if not isinstance(messages, str) else None
|
||||
system_messages: Final = (
|
||||
strip_beta_gated_system_messages(message_sequence, active_betas, drop_params=drop_params)
|
||||
if message_sequence is not None
|
||||
else StrippedMessages(messages=(), stripped=())
|
||||
)
|
||||
stripped: Final = (*thinking.stripped, *system_messages.stripped)
|
||||
if not stripped:
|
||||
return StrippedRequest(request=request, stripped=())
|
||||
rewritten: Final = {
|
||||
**request,
|
||||
**({"thinking": thinking.thinking} if thinking.stripped else {}),
|
||||
**({"messages": system_messages.messages} if system_messages.stripped else {}),
|
||||
}
|
||||
return StrippedRequest(request=rewritten, stripped=stripped)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class StrippedNativeRequest:
|
||||
thinking: object
|
||||
thinking_stripped: bool
|
||||
messages: Sequence[object]
|
||||
messages_stripped: bool
|
||||
stripped: tuple[str, ...]
|
||||
|
||||
|
||||
def unstripped_native_request(optional_params: object, messages: object) -> StrippedNativeRequest:
|
||||
"""The identity result for legs that strip later, once their final ``anthropic_beta`` list exists."""
|
||||
params: Final = _as_mapping(optional_params) or MappingProxyType({})
|
||||
message_sequence: Final = _as_sequence(messages) if not isinstance(messages, str) else None
|
||||
return StrippedNativeRequest(
|
||||
thinking=params.get("thinking"),
|
||||
thinking_stripped=False,
|
||||
messages=message_sequence if message_sequence is not None else (),
|
||||
messages_stripped=False,
|
||||
stripped=(),
|
||||
)
|
||||
|
||||
|
||||
def strip_beta_gated_native_request(
|
||||
headers: object, optional_params: object, messages: object, *, body_betas: object = (), drop_params: object
|
||||
) -> StrippedNativeRequest:
|
||||
"""The ``/v1/messages`` seam, fed the final (already filtered) headers, the final body ``anthropic_beta``
|
||||
list on legs that carry betas in the body, plus the body parts it may rewrite."""
|
||||
if not drops_unsupported_params(drop_params):
|
||||
return unstripped_native_request(optional_params, messages)
|
||||
active_betas: Final = active_beta_values(headers, body_betas)
|
||||
params: Final = _as_mapping(optional_params) or MappingProxyType({})
|
||||
thinking: Final = strip_beta_gated_thinking(params.get("thinking"), active_betas, drop_params=drop_params)
|
||||
message_sequence: Final = _as_sequence(messages) if not isinstance(messages, str) else None
|
||||
system_messages: Final = (
|
||||
strip_beta_gated_system_messages(message_sequence, active_betas, drop_params=drop_params)
|
||||
if message_sequence is not None
|
||||
else StrippedMessages(messages=(), stripped=())
|
||||
)
|
||||
return StrippedNativeRequest(
|
||||
thinking=thinking.thinking,
|
||||
thinking_stripped=bool(thinking.stripped),
|
||||
messages=system_messages.messages,
|
||||
messages_stripped=bool(system_messages.stripped),
|
||||
stripped=(*thinking.stripped, *system_messages.stripped),
|
||||
)
|
||||
|
||||
|
||||
def log_beta_gated_strips(stripped: Sequence[str], model: str | None, provider: str) -> None:
|
||||
if not stripped:
|
||||
return
|
||||
described: Final = ", ".join(
|
||||
f"{field} (unlocked by anthropic-beta {' or '.join(sorted(_FIELD_BETAS[field]))})"
|
||||
if field in _FIELD_BETAS
|
||||
else field
|
||||
for field in sorted(set(stripped))
|
||||
)
|
||||
verbose_logger.warning(
|
||||
"Dropped %s for model=%s provider=%s because drop_params is on and the provider would reject them on "
|
||||
"this request, either since this provider leg does not send that anthropic-beta (see "
|
||||
"anthropic_beta_headers_config.json) or since the model turns mid-conversation system messages into user "
|
||||
"messages.",
|
||||
described,
|
||||
model,
|
||||
provider,
|
||||
)
|
||||
|
||||
|
||||
def strip_beta_gated_request_body(
|
||||
request: Mapping[str, object], headers: object, *, drop_params: object, model: str | None, provider: str
|
||||
) -> Mapping[str, object]:
|
||||
gated: Final = strip_beta_gated_request_fields(
|
||||
request, active_beta_values(headers, request.get("anthropic_beta") or ()), drop_params=drop_params
|
||||
)
|
||||
log_beta_gated_strips(gated.stripped, model=model, provider=provider)
|
||||
return gated.request
|
||||
|
|
@ -248,6 +248,7 @@
|
|||
"text_editor_20241022": null,
|
||||
"text_editor_20250124": null,
|
||||
"thinking-binding-controls-2026-08-01": "thinking-binding-controls-2026-08-01",
|
||||
"thinking-display-updates-2026-08-18": "thinking-display-updates-2026-08-18",
|
||||
"token-efficient-tools-2025-02-19": "token-efficient-tools-2025-02-19",
|
||||
"web-fetch-2025-09-10": null,
|
||||
"web-search-2025-03-05": "web-search-2025-03-05"
|
||||
|
|
|
|||
|
|
@ -241,6 +241,13 @@ def get_provider_name(provider: str) -> str:
|
|||
return aliases.get(provider, provider)
|
||||
|
||||
|
||||
def has_beta_headers_config(provider: str) -> bool:
|
||||
config: Final = _bundled_beta_headers_config()
|
||||
aliases: Final = config.get("provider_aliases", {})
|
||||
canonical_provider: Final = (aliases.get(provider) if isinstance(aliases, Mapping) else None) or provider
|
||||
return canonical_provider in config
|
||||
|
||||
|
||||
def filter_and_transform_beta_headers(
|
||||
beta_headers: list[str],
|
||||
provider: str,
|
||||
|
|
@ -397,7 +404,7 @@ def update_request_with_filtered_beta(
|
|||
headers: dict,
|
||||
request_data: dict,
|
||||
provider: str,
|
||||
) -> tuple[dict, dict]:
|
||||
) -> tuple[dict[str, str], dict[str, object]]:
|
||||
"""
|
||||
Update both headers and request body beta fields based on provider support.
|
||||
Modifies both dicts in place and returns them.
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -455,17 +455,33 @@ class Cache:
|
|||
3. Else use the `model` passed in kwargs, plus `model_id` when set (e.g. a SageMaker inference component)
|
||||
"""
|
||||
litellm_params: Final[dict] = kwargs.get("litellm_params", {}) or {}
|
||||
metadata_sources: Final[tuple[dict, ...]] = (
|
||||
kwargs.get("metadata") or {},
|
||||
kwargs.get("litellm_metadata") or {},
|
||||
litellm_params.get("metadata") or {},
|
||||
litellm_params.get("litellm_metadata") or {},
|
||||
metadata_sources: Final = cast( # cast-ok: each source is the untyped request's metadata dict or empty
|
||||
"tuple[dict[str, object], dict[str, object], dict[str, object], dict[str, object]]",
|
||||
(
|
||||
kwargs.get("metadata") or {},
|
||||
kwargs.get("litellm_metadata") or {},
|
||||
litellm_params.get("metadata") or {},
|
||||
litellm_params.get("litellm_metadata") or {},
|
||||
),
|
||||
)
|
||||
model_group: Final[str | None] = next(
|
||||
(source["model_group"] for source in metadata_sources if source.get("model_group")), None
|
||||
metadata, litellm_metadata, params_metadata, params_litellm_metadata = metadata_sources
|
||||
call_model_info: Final[object] = kwargs.get("model_info")
|
||||
routed_source: Final = next(
|
||||
(
|
||||
source
|
||||
for source in (litellm_metadata, metadata, params_litellm_metadata, params_metadata)
|
||||
if call_model_info and source.get("model_group") and source.get("model_info") == call_model_info
|
||||
),
|
||||
None,
|
||||
)
|
||||
model_group: Final = (
|
||||
str(routed_source["model_group"])
|
||||
if routed_source is not None
|
||||
else next((str(source["model_group"]) for source in metadata_sources if source.get("model_group")), None)
|
||||
)
|
||||
caching_group_sources: Final = (routed_source,) if routed_source is not None else metadata_sources
|
||||
caching_group: Final = next(
|
||||
(group for source in metadata_sources if (group := self._get_caching_group(source, model_group))), None
|
||||
(group for source in caching_group_sources if (group := self._get_caching_group(source, model_group))), None
|
||||
)
|
||||
model_id_suffix: Final = (
|
||||
f" model_id: {kwargs['model_id']}" if "model_id" in kwargs and kwargs["model_id"] is not None else ""
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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]:
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -13,7 +13,7 @@ import asyncio
|
|||
import json
|
||||
import os
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Any, Final, Protocol, cast
|
||||
from typing import TYPE_CHECKING, Any, Final, Protocol
|
||||
|
||||
from pydantic import TypeAdapter
|
||||
|
||||
|
|
@ -292,10 +292,7 @@ class QdrantSemanticCache(BaseCache):
|
|||
prompt: Final = get_semantic_cache_prompt_from_messages(messages)
|
||||
|
||||
# create an embedding for prompt
|
||||
embedding_response: Final = cast(
|
||||
EmbeddingResponse,
|
||||
self._get_embedding(prompt, metadata=kwargs.get("metadata")),
|
||||
)
|
||||
embedding_response: Final = self._get_embedding(prompt, metadata=kwargs.get("metadata"))
|
||||
|
||||
# get the embedding
|
||||
embedding: Final = _EMBEDDING_VECTOR.validate_python(embedding_response["data"][0]["embedding"])
|
||||
|
|
@ -331,10 +328,7 @@ class QdrantSemanticCache(BaseCache):
|
|||
prompt: Final = get_semantic_cache_prompt_from_messages(messages)
|
||||
|
||||
# convert to embedding
|
||||
embedding_response: Final = cast(
|
||||
EmbeddingResponse,
|
||||
self._get_embedding(prompt, metadata=kwargs.get("metadata")),
|
||||
)
|
||||
embedding_response: Final = self._get_embedding(prompt, metadata=kwargs.get("metadata"))
|
||||
|
||||
# get the embedding
|
||||
embedding: Final = _EMBEDDING_VECTOR.validate_python(embedding_response["data"][0]["embedding"])
|
||||
|
|
|
|||
|
|
@ -14,7 +14,9 @@ import asyncio
|
|||
import json
|
||||
import os
|
||||
from collections.abc import Callable, Mapping
|
||||
from typing import TYPE_CHECKING, Any, Final, cast
|
||||
from typing import TYPE_CHECKING, Final, cast
|
||||
|
||||
from pydantic import ConfigDict, TypeAdapter
|
||||
|
||||
import litellm
|
||||
from litellm._logging import print_verbose, verbose_logger
|
||||
|
|
@ -38,6 +40,11 @@ if TYPE_CHECKING:
|
|||
from litellm.router import Router
|
||||
|
||||
|
||||
_CACHE_RESPONSE_ADAPTER: Final[TypeAdapter[str | bytes | None]] = TypeAdapter(
|
||||
str | bytes | None, config=ConfigDict(strict=True, hide_input_in_errors=True)
|
||||
)
|
||||
|
||||
|
||||
class RedisSemanticCache(BaseCache):
|
||||
"""
|
||||
Redis-backed semantic cache for LLM responses.
|
||||
|
|
@ -380,7 +387,7 @@ class RedisSemanticCache(BaseCache):
|
|||
)
|
||||
return embedding_response["data"][0]["embedding"]
|
||||
|
||||
def _get_cache_logic(self, cached_response: Any) -> object:
|
||||
def _get_cache_logic(self, cached_response: str | bytes | None) -> object:
|
||||
"""
|
||||
Process the cached response to prepare it for use.
|
||||
|
||||
|
|
@ -391,24 +398,23 @@ class RedisSemanticCache(BaseCache):
|
|||
The processed cache response, or None if input was None
|
||||
"""
|
||||
if cached_response is None:
|
||||
return cached_response
|
||||
return None
|
||||
|
||||
# Convert bytes to string if needed
|
||||
if isinstance(cached_response, bytes):
|
||||
cached_response = cached_response.decode("utf-8")
|
||||
response_text: Final = (
|
||||
cached_response.decode("utf-8") if isinstance(cached_response, bytes) else cached_response
|
||||
)
|
||||
|
||||
# Convert string representation to Python object
|
||||
try:
|
||||
cached_response = json.loads(cached_response)
|
||||
return json.loads(response_text)
|
||||
except json.JSONDecodeError:
|
||||
try:
|
||||
cached_response = ast.literal_eval(cached_response)
|
||||
return ast.literal_eval(response_text)
|
||||
except (ValueError, SyntaxError) as e:
|
||||
print_verbose(f"Error parsing cached response: {e}")
|
||||
return None
|
||||
|
||||
return cached_response
|
||||
|
||||
def set_cache(self, key: str, value: object, **kwargs) -> None:
|
||||
"""
|
||||
Store a value in the semantic cache.
|
||||
|
|
@ -494,7 +500,7 @@ class RedisSemanticCache(BaseCache):
|
|||
similarity: Final = 1 - vector_distance
|
||||
|
||||
cached_prompt: Final = cache_hit["prompt"]
|
||||
cached_response: Final = cache_hit["response"]
|
||||
cached_response: Final = _CACHE_RESPONSE_ADAPTER.validate_python(cache_hit["response"])
|
||||
|
||||
# update kwargs["metadata"] with similarity, don't rewrite the original metadata
|
||||
kwargs.setdefault("metadata", {})["semantic-similarity"] = similarity
|
||||
|
|
@ -645,7 +651,7 @@ class RedisSemanticCache(BaseCache):
|
|||
similarity: Final = 1 - vector_distance
|
||||
|
||||
cached_prompt: Final = cache_hit["prompt"]
|
||||
cached_response: Final = cache_hit["response"]
|
||||
cached_response: Final = _CACHE_RESPONSE_ADAPTER.validate_python(cache_hit["response"])
|
||||
|
||||
# update kwargs["metadata"] with similarity, don't rewrite the original metadata
|
||||
kwargs.setdefault("metadata", {})["semantic-similarity"] = similarity
|
||||
|
|
|
|||
|
|
@ -22,7 +22,7 @@ from litellm.constants import (
|
|||
OPENAI_LIVE_TERMINAL_RESPONSE_EVENT_TYPES,
|
||||
REALTIME_METERED_USAGE_KEY,
|
||||
)
|
||||
from litellm.litellm_core_utils.core_helpers import as_str_mapping
|
||||
from litellm.litellm_core_utils.core_helpers import as_str_mapping, billed_usage_cost
|
||||
from litellm.litellm_core_utils.cost_discount import resolve_cost_discount
|
||||
from litellm.litellm_core_utils.decisions_pricing import decisions_pricing_fields, is_decisions_call_type
|
||||
from litellm.litellm_core_utils.hidden_params import HIDDEN_PARAMS_ATTR
|
||||
|
|
@ -50,6 +50,7 @@ from litellm.litellm_core_utils.llm_cost_calc.utils import (
|
|||
parse_prompt_tokens_details,
|
||||
select_cost_metric_for_model,
|
||||
)
|
||||
from litellm.litellm_core_utils.llm_response_utils.response_metadata import response_ms_since_logging_start
|
||||
from litellm.llms.anthropic.cost_calculation import (
|
||||
cost_per_token as anthropic_cost_per_token,
|
||||
)
|
||||
|
|
@ -1396,9 +1397,9 @@ def _apply_cost_margin(
|
|||
|
||||
def _store_cost_breakdown_in_logging_obj(
|
||||
litellm_logging_obj: LitellmLoggingObject | None,
|
||||
prompt_tokens_cost_usd_dollar: float,
|
||||
completion_tokens_cost_usd_dollar: float,
|
||||
cost_for_built_in_tools_cost_usd_dollar: float,
|
||||
prompt_tokens_cost_usd_dollar: float | None,
|
||||
completion_tokens_cost_usd_dollar: float | None,
|
||||
cost_for_built_in_tools_cost_usd_dollar: float | None,
|
||||
total_cost_usd_dollar: float,
|
||||
additional_costs: dict | None = None,
|
||||
original_cost: float | None = None,
|
||||
|
|
@ -1500,9 +1501,45 @@ def _response_time_ms_for_cost(
|
|||
return float(stamped)
|
||||
if total_time:
|
||||
return total_time
|
||||
if litellm_logging_obj is not None:
|
||||
return litellm_logging_obj.get_response_ms()
|
||||
return 0.0
|
||||
if litellm_logging_obj is None:
|
||||
return 0.0
|
||||
reported_ms: Final = response_ms_since_logging_start(litellm_logging_obj)
|
||||
if reported_ms is not None and litellm_logging_obj.stream is not True:
|
||||
return reported_ms
|
||||
return litellm_logging_obj.get_response_ms()
|
||||
|
||||
|
||||
def _adjusted_provider_billed_cost(
|
||||
reported_cost: float,
|
||||
custom_llm_provider: str | None,
|
||||
model: str | None,
|
||||
region_name: str | None,
|
||||
litellm_logging_obj: LitellmLoggingObject | None,
|
||||
service_tier: str | None,
|
||||
data_residency: str | None,
|
||||
) -> float:
|
||||
discounted_cost, discount_percent, discount_amount = _apply_cost_discount(
|
||||
base_cost=reported_cost, custom_llm_provider=custom_llm_provider, model=model, region_name=region_name
|
||||
)
|
||||
adjusted_cost, margin_percent, margin_fixed_amount, margin_total_amount = _apply_cost_margin(
|
||||
base_cost=discounted_cost, custom_llm_provider=custom_llm_provider
|
||||
)
|
||||
_store_cost_breakdown_in_logging_obj(
|
||||
litellm_logging_obj=litellm_logging_obj,
|
||||
prompt_tokens_cost_usd_dollar=None,
|
||||
completion_tokens_cost_usd_dollar=None,
|
||||
cost_for_built_in_tools_cost_usd_dollar=None,
|
||||
total_cost_usd_dollar=adjusted_cost,
|
||||
original_cost=reported_cost,
|
||||
discount_percent=discount_percent,
|
||||
discount_amount=discount_amount,
|
||||
margin_percent=margin_percent,
|
||||
margin_fixed_amount=margin_fixed_amount,
|
||||
margin_total_amount=margin_total_amount,
|
||||
service_tier=service_tier,
|
||||
data_residency=data_residency,
|
||||
)
|
||||
return adjusted_cost
|
||||
|
||||
|
||||
def completion_cost(
|
||||
|
|
@ -1616,6 +1653,9 @@ def completion_cost(
|
|||
)
|
||||
|
||||
explicit_pricing: Final = custom_pricing is True or base_model is not None
|
||||
caller_set_prices: Final = (
|
||||
explicit_pricing or custom_cost_per_token is not None or custom_cost_per_second is not None
|
||||
)
|
||||
custom_priced_request_model: Final = model if custom_pricing is True else None
|
||||
configured_cost_per_query: Final = (
|
||||
_configured_search_cost_per_query(litellm_logging_obj) if custom_pricing else None
|
||||
|
|
@ -1666,6 +1706,7 @@ def completion_cost(
|
|||
raw_candidates if custom_pricing is True else _priced_candidates_before_unpriced_rules(raw_candidates)
|
||||
)
|
||||
|
||||
priced_response: Final[object] = completion_response
|
||||
for idx, (candidate_model, candidate_pricing_provider) in enumerate(potential_model_names):
|
||||
model = candidate_model
|
||||
try:
|
||||
|
|
@ -1984,6 +2025,19 @@ def completion_cost(
|
|||
total_time=total_time,
|
||||
)
|
||||
model_for_discount = custom_priced_request_model or model
|
||||
if (
|
||||
not caller_set_prices
|
||||
and (reported_cost := billed_usage_cost(priced_response, custom_llm_provider)) is not None
|
||||
):
|
||||
return _adjusted_provider_billed_cost(
|
||||
reported_cost=reported_cost,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
model=model_for_discount,
|
||||
region_name=region_name if isinstance(region_name, str) else None,
|
||||
litellm_logging_obj=litellm_logging_obj,
|
||||
service_tier=service_tier,
|
||||
data_residency=data_residency,
|
||||
)
|
||||
# Calculate cost based on prompt_tokens, completion_tokens
|
||||
if (
|
||||
"togethercomputer" in model or "together_ai" in model or custom_llm_provider == "together_ai"
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
@ -641,6 +651,8 @@ def image_generation(
|
|||
timeout=timeout,
|
||||
client=custom_client,
|
||||
)
|
||||
elif custom_llm_provider == LlmProviders.KYMA.value:
|
||||
raise LiteLLMUnknownProvider(model=model, custom_llm_provider=custom_llm_provider)
|
||||
|
||||
return model_response
|
||||
except Exception as e:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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())
|
||||
|
|
|
|||
|
|
@ -513,7 +513,7 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
|
|||
data: dict,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
call_type: CallTypesLiteral,
|
||||
) -> Any:
|
||||
) -> object:
|
||||
pass
|
||||
|
||||
async def async_post_call_streaming_hook(
|
||||
|
|
|
|||
|
|
@ -205,11 +205,12 @@ class DatadogCostManagementLogger(CustomBatchLogger):
|
|||
# see _RESERVED_TAG_KEYS for the rationale.
|
||||
if self.cost_tag_keys:
|
||||
allow: Final = set(self.cost_tag_keys)
|
||||
client_allow: Final = allow - {"is_silent_experiment"}
|
||||
for rt in log.get("request_tags") or []:
|
||||
if not isinstance(rt, str) or ":" not in rt:
|
||||
continue
|
||||
k, _, v = rt.partition(":")
|
||||
if k in allow and v:
|
||||
if k in client_allow and v:
|
||||
self._set_custom_tag(tags, k, v)
|
||||
for k, v in metadata.items():
|
||||
if k in allow and v is not None and not isinstance(v, (dict, list)):
|
||||
|
|
@ -218,7 +219,7 @@ class DatadogCostManagementLogger(CustomBatchLogger):
|
|||
nested = metadata.get(nested_key)
|
||||
if isinstance(nested, dict):
|
||||
for k, v in nested.items():
|
||||
if k in allow and v is not None and not isinstance(v, (dict, list)):
|
||||
if k in client_allow and v is not None and not isinstance(v, (dict, list)):
|
||||
self._set_custom_tag(tags, k, str(v))
|
||||
|
||||
return tags
|
||||
|
|
|
|||
|
|
@ -833,6 +833,12 @@ class LangFuseLogger:
|
|||
("vertex_location", vertex_location, bool(vertex_location)),
|
||||
("aws_region_name", aws_region_name, bool(aws_region_name)),
|
||||
("cache_hit", kwargs.get("cache_hit") or False, "cache_hit" in kwargs),
|
||||
(
|
||||
"is_silent_experiment",
|
||||
"is_silent_experiment" in allowlisted_metadata
|
||||
and allowlisted_metadata["is_silent_experiment"] is True,
|
||||
"is_silent_experiment" in allowlisted_metadata,
|
||||
),
|
||||
)
|
||||
enrichments: Final[Mapping[str, object]] = {
|
||||
key: value for key, value, include in candidate_enrichments if include
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -15,7 +15,7 @@ class BaseLLMObsOTELAttributes(ABC):
|
|||
pass
|
||||
|
||||
|
||||
def cast_as_primitive_value_type(value) -> str | bool | int | float:
|
||||
def cast_as_primitive_value_type(value: object) -> str | bool | int | float:
|
||||
"""
|
||||
Converts a value to an OTEL-supported primitive for Arize/Phoenix observability.
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -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():
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
@ -5148,6 +5149,12 @@ def _get_combined_custom_metadata_from_standard_logging_payload(
|
|||
**(requester_metadata if isinstance(requester_metadata, dict) else {}),
|
||||
**(user_api_key_auth_metadata if isinstance(user_api_key_auth_metadata, dict) else {}),
|
||||
**(spend_logs_metadata if isinstance(spend_logs_metadata, dict) else {}),
|
||||
"is_silent_experiment": (
|
||||
"true"
|
||||
if "is_silent_experiment" in standard_logging_metadata
|
||||
and standard_logging_metadata["is_silent_experiment"] is True
|
||||
else None
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -103,6 +103,7 @@ class LiteLLMResponsesInteractionsHandler:
|
|||
return LiteLLMResponsesInteractionsConfig.transform_responses_response_to_interactions_response(
|
||||
responses_response=responses_api_response,
|
||||
model=model,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
)
|
||||
|
||||
async def async_interactions_api_handler(
|
||||
|
|
@ -138,4 +139,5 @@ class LiteLLMResponsesInteractionsHandler:
|
|||
return LiteLLMResponsesInteractionsConfig.transform_responses_response_to_interactions_response(
|
||||
responses_response=responses_api_response,
|
||||
model=model,
|
||||
custom_llm_provider=responses_request.get("custom_llm_provider"),
|
||||
)
|
||||
|
|
|
|||
|
|
@ -12,6 +12,7 @@ from typing import Final, cast
|
|||
|
||||
from pydantic import BaseModel
|
||||
|
||||
from litellm.litellm_core_utils.core_helpers import billed_usage_cost
|
||||
from litellm.types.interactions import (
|
||||
InteractionInput,
|
||||
InteractionsAPIOptionalRequestParams,
|
||||
|
|
@ -190,6 +191,7 @@ class LiteLLMResponsesInteractionsConfig:
|
|||
def transform_responses_response_to_interactions_response(
|
||||
responses_response: ResponsesAPIResponse,
|
||||
model: str | None = None,
|
||||
custom_llm_provider: str | None = None,
|
||||
) -> InteractionsAPIResponse:
|
||||
"""
|
||||
Transform a Responses API response to an Interactions API response.
|
||||
|
|
@ -266,9 +268,11 @@ class LiteLLMResponsesInteractionsConfig:
|
|||
# (total_input_tokens, total_output_tokens)
|
||||
usage: Final[object] = getattr(responses_response, "usage", None)
|
||||
if usage:
|
||||
reported_cost: Final = billed_usage_cost(responses_response, custom_llm_provider)
|
||||
interactions_response_dict["usage"] = {
|
||||
"total_input_tokens": getattr(usage, "input_tokens", 0),
|
||||
"total_output_tokens": getattr(usage, "output_tokens", 0),
|
||||
**({} if reported_cost is None else {"cost": reported_cost}),
|
||||
}
|
||||
|
||||
# Add updated (same as created for now)
|
||||
|
|
|
|||
|
|
@ -2,6 +2,7 @@
|
|||
## Helper utilities
|
||||
import copy
|
||||
import logging
|
||||
import math
|
||||
import re
|
||||
from collections.abc import Collection, Iterable, Iterator, Mapping
|
||||
from types import MappingProxyType
|
||||
|
|
@ -1016,6 +1017,16 @@ def chunk_usage(chunk: object) -> object:
|
|||
return getattr(chunk, "usage", None)
|
||||
|
||||
|
||||
USAGE_COST_BILLED_PROVIDERS: Final = frozenset({"kyma"})
|
||||
|
||||
|
||||
def billed_usage_cost(response: object, custom_llm_provider: str | None) -> float | None:
|
||||
if custom_llm_provider not in USAGE_COST_BILLED_PROVIDERS:
|
||||
return None
|
||||
cost: Final = provider_reported_usage_cost(_usage_field(response, "usage"), custom_llm_provider)
|
||||
return cost if cost is not None and math.isfinite(cost) and cost >= 0 else None
|
||||
|
||||
|
||||
_MESSAGES_REPORTED_COST_PROVIDERS: Final = frozenset({"opencode_go"})
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
@ -23,7 +24,7 @@ from httpx import Response, Timeout
|
|||
from pydantic import BaseModel, JsonValue, TypeAdapter
|
||||
|
||||
import litellm
|
||||
from litellm._internal_context import post_response_phase
|
||||
from litellm._internal_context import in_mirrored_call, outside_mirrored_call, post_response_phase
|
||||
from litellm._logging import (
|
||||
is_debugging_on,
|
||||
redact_string,
|
||||
|
|
@ -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,
|
||||
|
|
@ -162,7 +164,7 @@ from litellm.types.llms.openai import (
|
|||
from litellm.types.mcp import MCPPostCallResponseObject
|
||||
from litellm.types.prompts.init_prompts import PromptSpec
|
||||
from litellm.types.rerank import RerankResponse
|
||||
from litellm.types.router import RouterRateLimitError
|
||||
from litellm.types.router import RouterRateLimitError, SilentExperimentFlag
|
||||
from litellm.types.utils import (
|
||||
DEPLOYMENT_SCOPED_PRICING_FIELDS,
|
||||
AgenticLoopParams,
|
||||
|
|
@ -315,7 +317,7 @@ else:
|
|||
_in_memory_loggers: Final[list[CustomLogger]] = []
|
||||
|
||||
_STANDARD_LOGGING_METADATA_RESOLVED_KEYS: Final[frozenset[str]] = frozenset(
|
||||
("used_client_oauth_token", "usage_object", KUBERNETES_POD_ROUTING_KEY)
|
||||
("used_client_oauth_token", "usage_object", KUBERNETES_POD_ROUTING_KEY, "is_silent_experiment")
|
||||
)
|
||||
_STANDARD_LOGGING_METADATA_KEYS: Final[frozenset[str]] = (
|
||||
frozenset(StandardLoggingMetadata.__annotations__.keys()) - _STANDARD_LOGGING_METADATA_RESOLVED_KEYS
|
||||
|
|
@ -809,9 +811,15 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
litellm_params = scrub_sensitive_keys_in_metadata(litellm_params)
|
||||
|
||||
self.litellm_params = litellm_params
|
||||
self.is_silent_experiment: Final = in_mirrored_call()
|
||||
|
||||
# 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
|
||||
|
||||
|
|
@ -873,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
|
||||
|
|
@ -1941,10 +1954,10 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
|
||||
def set_cost_breakdown(
|
||||
self,
|
||||
input_cost: float,
|
||||
output_cost: float,
|
||||
input_cost: float | None,
|
||||
output_cost: float | None,
|
||||
total_cost: float,
|
||||
cost_for_built_in_tools_cost_usd_dollar: float,
|
||||
cost_for_built_in_tools_cost_usd_dollar: float | None,
|
||||
additional_costs: dict | None = None,
|
||||
original_cost: float | None = None,
|
||||
discount_percent: float | None = None,
|
||||
|
|
@ -1983,14 +1996,17 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
|
||||
self.billed_token_rates = billed_token_rates
|
||||
self.cost_breakdown = CostBreakdown(
|
||||
input_cost=input_cost,
|
||||
output_cost=output_cost,
|
||||
total_cost=total_cost,
|
||||
tool_usage_cost=cost_for_built_in_tools_cost_usd_dollar,
|
||||
service_tier=service_tier,
|
||||
data_residency=data_residency,
|
||||
vertex_location=vertex_location,
|
||||
)
|
||||
if input_cost is not None:
|
||||
self.cost_breakdown["input_cost"] = input_cost
|
||||
if output_cost is not None:
|
||||
self.cost_breakdown["output_cost"] = output_cost
|
||||
if cost_for_built_in_tools_cost_usd_dollar is not None:
|
||||
self.cost_breakdown["tool_usage_cost"] = cost_for_built_in_tools_cost_usd_dollar
|
||||
if cache_read_cost is not None and cache_read_cost > 0:
|
||||
self.cost_breakdown["cache_read_cost"] = cache_read_cost
|
||||
if cache_creation_cost is not None and cache_creation_cost > 0:
|
||||
|
|
@ -2458,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
|
||||
|
|
@ -2808,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.
|
||||
|
|
@ -2932,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,
|
||||
|
|
@ -2943,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,
|
||||
|
|
@ -2964,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)
|
||||
|
|
@ -3003,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)
|
||||
|
|
@ -3033,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
|
||||
|
|
@ -3253,7 +3334,7 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
"""Restores trace_id/session_id contextvars once this attempt's own success
|
||||
logging (including any nested calls its callbacks trigger) is fully done."""
|
||||
try:
|
||||
with post_response_phase():
|
||||
with post_response_phase(), outside_mirrored_call():
|
||||
return self._success_handler_body(
|
||||
result=result,
|
||||
start_time=start_time,
|
||||
|
|
@ -3277,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,
|
||||
|
|
@ -3753,7 +3834,7 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
"""Restores trace_id/session_id contextvars once this attempt's own success
|
||||
logging (including any nested calls its callbacks trigger) is fully done."""
|
||||
try:
|
||||
with post_response_phase():
|
||||
with post_response_phase(), outside_mirrored_call():
|
||||
return await self._async_success_handler_body(
|
||||
result=result, start_time=start_time, end_time=end_time, cache_hit=cache_hit, **kwargs
|
||||
)
|
||||
|
|
@ -4200,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):
|
||||
|
|
@ -4270,13 +4353,14 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
``for_async_request`` marks the sync run of a failure whose async run also logs it, so a
|
||||
CustomLogger that gets ``async_log_failure_event`` from that async run skips its sync hook."""
|
||||
try:
|
||||
return self._failure_handler_body(
|
||||
exception=exception,
|
||||
traceback_exception=traceback_exception,
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
for_async_request=for_async_request,
|
||||
)
|
||||
with outside_mirrored_call():
|
||||
return self._failure_handler_body(
|
||||
exception=exception,
|
||||
traceback_exception=traceback_exception,
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
for_async_request=for_async_request,
|
||||
)
|
||||
finally:
|
||||
self._restore_correlation_context()
|
||||
|
||||
|
|
@ -4469,12 +4553,13 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
"""Restores trace_id/session_id contextvars once this attempt's own failure
|
||||
logging (including any nested calls its callbacks trigger) is fully done."""
|
||||
try:
|
||||
return await self._async_failure_handler_body(
|
||||
exception=exception,
|
||||
traceback_exception=traceback_exception,
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
)
|
||||
with outside_mirrored_call():
|
||||
return await self._async_failure_handler_body(
|
||||
exception=exception,
|
||||
traceback_exception=traceback_exception,
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
)
|
||||
finally:
|
||||
self._restore_correlation_context()
|
||||
|
||||
|
|
@ -5538,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)
|
||||
|
|
@ -6563,6 +6648,7 @@ class StandardLoggingPayloadSetup:
|
|||
start_time: dt_object | None = None,
|
||||
response_id: str | None = None,
|
||||
custom_llm_provider: str | None = None,
|
||||
is_silent_experiment: bool = False,
|
||||
) -> StandardLoggingMetadata:
|
||||
"""
|
||||
Clean and filter the metadata dictionary to include only the specified keys in StandardLoggingMetadata.
|
||||
|
|
@ -6636,6 +6722,7 @@ class StandardLoggingPayloadSetup:
|
|||
proxy_stamped_used_client_oauth_token(metadata, litellm_params),
|
||||
custom_llm_provider,
|
||||
),
|
||||
**(SilentExperimentFlag(is_silent_experiment=True) if is_silent_experiment else SilentExperimentFlag()),
|
||||
)
|
||||
if isinstance(metadata, dict):
|
||||
for key in metadata.keys() & _STANDARD_LOGGING_METADATA_KEYS:
|
||||
|
|
@ -7359,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,
|
||||
|
|
@ -7436,6 +7529,7 @@ def get_standard_logging_object_payload(
|
|||
start_time=start_time,
|
||||
response_id=id,
|
||||
custom_llm_provider=selected_provider if isinstance(selected_provider, str) else None,
|
||||
is_silent_experiment=logging_obj.is_silent_experiment is True,
|
||||
)
|
||||
_request_body: Final = proxy_server_request.get("body", {})
|
||||
end_user_id: Final = clean_metadata["user_api_key_end_user_id"] or _request_body.get(
|
||||
|
|
@ -7461,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)
|
||||
|
|
@ -7528,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 = (
|
||||
|
|
@ -7591,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,
|
||||
|
|
@ -7720,9 +7822,8 @@ def get_standard_logging_metadata(
|
|||
)
|
||||
if isinstance(metadata, dict):
|
||||
# Update the clean_metadata with values from input metadata that match StandardLoggingMetadata fields
|
||||
for key in StandardLoggingMetadata.__annotations__:
|
||||
if key in metadata:
|
||||
clean_metadata[key] = metadata[key]
|
||||
for key in (metadata.keys() & StandardLoggingMetadata.__annotations__.keys()) - {"is_silent_experiment"}:
|
||||
clean_metadata[key] = metadata[key]
|
||||
|
||||
if metadata.get("user_api_key") is not None:
|
||||
if is_valid_sha256_hash(str(metadata.get("user_api_key"))):
|
||||
|
|
|
|||
61
litellm/litellm_core_utils/llm_cost_calc/finalization.py
Normal file
61
litellm/litellm_core_utils/llm_cost_calc/finalization.py
Normal 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)
|
||||
|
|
@ -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
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -27,6 +27,15 @@ def _timing_window_start(
|
|||
return start_time, False
|
||||
|
||||
|
||||
def response_ms_since_logging_start(logging_obj: LiteLLMLoggingObject) -> float | None:
|
||||
reported_ms: Final = logging_obj.response_timing_metrics.get("_response_ms")
|
||||
if reported_ms is None:
|
||||
return None
|
||||
logging_start: Final = logging_obj.start_time
|
||||
window_start: Final = _timing_window_start(logging_start, logging_obj)[0]
|
||||
return reported_ms - max(0.0, logging_start.timestamp() - window_start.timestamp()) * 1000
|
||||
|
||||
|
||||
def _union_duration_ms(windows: object, lower: float, upper: float) -> float | None:
|
||||
if not isinstance(windows, (list, tuple)):
|
||||
return None
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -12,9 +12,14 @@ from typing_extensions import ReadOnly, Required
|
|||
|
||||
import litellm
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.litellm_core_utils.core_helpers import chunk_usage, provider_reported_usage_cost
|
||||
from litellm.litellm_core_utils.core_helpers import (
|
||||
USAGE_COST_BILLED_PROVIDERS,
|
||||
chunk_usage,
|
||||
provider_reported_usage_cost,
|
||||
)
|
||||
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
|
||||
)
|
||||
|
|
@ -1352,7 +1357,7 @@ def stream_chunk_builder_text_completion(
|
|||
return TextCompletionResponse(**response)
|
||||
|
||||
|
||||
_CALCULATOR_PRICED_REPORTED_COST_PROVIDERS: Final = frozenset({LlmProviders.XAI.value})
|
||||
_CALCULATOR_PRICED_REPORTED_COST_PROVIDERS: Final = frozenset({LlmProviders.XAI.value}) | USAGE_COST_BILLED_PROVIDERS
|
||||
|
||||
|
||||
def _logged_provider(logging_obj: Optional["Logging"]) -> object:
|
||||
|
|
@ -1696,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
|
||||
|
|
@ -1762,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)
|
||||
|
||||
|
|
|
|||
|
|
@ -432,27 +432,30 @@ class AnthropicMessagesHandler(BaseTranslation):
|
|||
message, _ = serialize_http_exception_detail(exc.detail)
|
||||
return tuple(anthropic_sse_error_frames(message))
|
||||
|
||||
def _standalone_block_chunks(self, exc: "ModifyResponseException") -> list[bytes]:
|
||||
def build_block_response(self, exc: "ModifyResponseException") -> "AnthropicMessagesResponse":
|
||||
import uuid
|
||||
|
||||
from litellm.llms.anthropic.messages.native.fake_stream_iterator import (
|
||||
FakeAnthropicMessagesStreamIterator,
|
||||
)
|
||||
from litellm.llms.base_llm.guardrail_translation.utils import (
|
||||
blocked_response_usage,
|
||||
)
|
||||
from litellm.types.utils import AnthropicMessagesResponse
|
||||
|
||||
block_response: Final = AnthropicMessagesResponse(
|
||||
return AnthropicMessagesResponse(
|
||||
id=f"msg_{uuid.uuid4()}",
|
||||
type="message",
|
||||
role="assistant",
|
||||
content=[{"type": "text", "text": exc.message}],
|
||||
model=exc.model,
|
||||
stop_reason="end_turn",
|
||||
usage=blocked_response_usage(getattr(exc, "original_response", None)),
|
||||
usage=blocked_response_usage(exc.original_response),
|
||||
)
|
||||
return list(FakeAnthropicMessagesStreamIterator(response=block_response))
|
||||
|
||||
def _standalone_block_chunks(self, exc: "ModifyResponseException") -> list[bytes]:
|
||||
from litellm.llms.anthropic.messages.native.fake_stream_iterator import (
|
||||
FakeAnthropicMessagesStreamIterator,
|
||||
)
|
||||
|
||||
return list(FakeAnthropicMessagesStreamIterator(response=self.build_block_response(exc)))
|
||||
|
||||
def _block_continuation_chunks(
|
||||
self, exc: "ModifyResponseException", responses_so_far: Sequence[object]
|
||||
|
|
|
|||
|
|
@ -17,6 +17,7 @@ import litellm
|
|||
import litellm.litellm_core_utils
|
||||
import litellm.types
|
||||
import litellm.types.utils
|
||||
from litellm.anthropic_beta_gated_fields import strip_beta_gated_request_body
|
||||
from litellm.anthropic_beta_headers_manager import (
|
||||
update_request_with_filtered_beta,
|
||||
)
|
||||
|
|
@ -506,11 +507,20 @@ class AnthropicChatCompletion(BaseLLM):
|
|||
The pair stays mutable because the streaming path rewrites it in
|
||||
place (`data["stream"] = True`) before sending.
|
||||
"""
|
||||
request_headers, data = update_request_with_filtered_beta(
|
||||
request_headers, filtered_data = update_request_with_filtered_beta(
|
||||
headers=headers,
|
||||
request_data=request_data,
|
||||
provider=custom_llm_provider,
|
||||
)
|
||||
data: Final = dict(
|
||||
strip_beta_gated_request_body(
|
||||
filtered_data,
|
||||
request_headers,
|
||||
drop_params=litellm_params.get("drop_params"), # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # litellm_params is the handler's bare dict
|
||||
model=model,
|
||||
provider=custom_llm_provider,
|
||||
)
|
||||
)
|
||||
logging_obj.pre_call(
|
||||
input=messages,
|
||||
api_key=api_key,
|
||||
|
|
|
|||
|
|
@ -51,6 +51,7 @@ from litellm.types.llms.anthropic import (
|
|||
ANTHROPIC_OAUTH_BETA_HEADER,
|
||||
ANTHROPIC_OAUTH_TOKEN_PREFIX,
|
||||
ANTHROPIC_SERVER_TOOL_USE_ID_PREFIX,
|
||||
ANTHROPIC_THINKING_BINDING_CONTROLS_BETA_HEADER,
|
||||
ANTHROPIC_THINKING_DISPLAY_UPDATES_BETA_HEADER,
|
||||
ANTHROPIC_TOOL_SEARCH_TOOL_TYPES,
|
||||
AllAnthropicToolsValues,
|
||||
|
|
@ -600,6 +601,20 @@ class AnthropicModelInfo(BaseLLMModelInfo):
|
|||
return False
|
||||
return thinking.get("type") in ("adaptive", "enabled") and thinking.get("display") == "updates"
|
||||
|
||||
def thinking_betas(self, thinking: AnthropicThinkingParam | None) -> tuple[str, ...]:
|
||||
return (
|
||||
*(
|
||||
(ANTHROPIC_THINKING_DISPLAY_UPDATES_BETA_HEADER,)
|
||||
if self.is_thinking_display_updates_used(thinking)
|
||||
else ()
|
||||
),
|
||||
*(
|
||||
(ANTHROPIC_THINKING_BINDING_CONTROLS_BETA_HEADER,)
|
||||
if isinstance(thinking, dict) and "block_binding" in thinking
|
||||
else ()
|
||||
),
|
||||
)
|
||||
|
||||
def is_mid_conversation_tool_change_used(self, messages: Sequence[object]) -> bool:
|
||||
for message in messages:
|
||||
if message_field(message, "role") != "system":
|
||||
|
|
@ -1180,7 +1195,7 @@ class AnthropicModelInfo(BaseLLMModelInfo):
|
|||
*,
|
||||
custom_llm_provider: str,
|
||||
is_mid_conversation_output_config_used: bool = False,
|
||||
is_thinking_display_updates_used: bool = False,
|
||||
thinking: AnthropicThinkingParam | None = None,
|
||||
is_mid_conversation_tool_change_used: bool = False,
|
||||
) -> list[str]:
|
||||
"""
|
||||
|
|
@ -1217,13 +1232,10 @@ class AnthropicModelInfo(BaseLLMModelInfo):
|
|||
if is_mid_conversation_output_config_used:
|
||||
betas.append(ANTHROPIC_MID_CONVERSATION_OUTPUT_CONFIG_BETA_HEADER)
|
||||
|
||||
thinking_display_betas: Final = (
|
||||
(ANTHROPIC_THINKING_DISPLAY_UPDATES_BETA_HEADER,) if is_thinking_display_updates_used else ()
|
||||
)
|
||||
tool_change_betas: Final = (
|
||||
(ANTHROPIC_MID_CONVERSATION_TOOL_CHANGES_BETA_HEADER,) if is_mid_conversation_tool_change_used else ()
|
||||
)
|
||||
return list(set(betas).union(thinking_display_betas, tool_change_betas))
|
||||
return list(set(betas).union(self.thinking_betas(thinking), tool_change_betas))
|
||||
|
||||
@staticmethod
|
||||
def _make_api_key_auth_header(
|
||||
|
|
@ -1285,7 +1297,7 @@ class AnthropicModelInfo(BaseLLMModelInfo):
|
|||
api_base: str | None = None,
|
||||
use_bearer_for_custom_base: bool = False,
|
||||
is_mid_conversation_output_config_used: bool = False,
|
||||
is_thinking_display_updates_used: bool = False,
|
||||
thinking: AnthropicThinkingParam | None = None,
|
||||
is_mid_conversation_tool_change_used: bool = False,
|
||||
wif_minted: bool = False,
|
||||
) -> dict:
|
||||
|
|
@ -1346,7 +1358,7 @@ class AnthropicModelInfo(BaseLLMModelInfo):
|
|||
betas.update(user_anthropic_beta_headers)
|
||||
|
||||
all_betas: Final = betas.union(
|
||||
(ANTHROPIC_THINKING_DISPLAY_UPDATES_BETA_HEADER,) if is_thinking_display_updates_used else (),
|
||||
self.thinking_betas(thinking),
|
||||
(ANTHROPIC_MID_CONVERSATION_TOOL_CHANGES_BETA_HEADER,) if is_mid_conversation_tool_change_used else (),
|
||||
)
|
||||
|
||||
|
|
@ -1433,7 +1445,7 @@ class AnthropicModelInfo(BaseLLMModelInfo):
|
|||
auth_token=auth_token,
|
||||
file_id_used=file_id_used,
|
||||
is_mid_conversation_output_config_used=is_mid_conversation_output_config_used,
|
||||
is_thinking_display_updates_used=self.is_thinking_display_updates_used(optional_params.get("thinking")),
|
||||
thinking=optional_params.get("thinking"),
|
||||
is_mid_conversation_tool_change_used=self.is_mid_conversation_tool_change_used(messages),
|
||||
web_search_tool_used=web_search_tool_used,
|
||||
is_vertex_request=optional_params.get("is_vertex_request", False),
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
@ -203,7 +206,7 @@ from litellm.types.llms.openai import (
|
|||
ChatCompletionUserMessage,
|
||||
ToolMessageContentPart,
|
||||
)
|
||||
from litellm.types.utils import Choices, ModelResponse, StreamingChoices, Usage
|
||||
from litellm.types.utils import Choices, LlmProviders, ModelResponse, StreamingChoices, Usage
|
||||
from litellm.utils import supports_mid_conversation_system
|
||||
|
||||
from .streaming_iterator import AnthropicSSEStream, AnthropicStreamWrapper
|
||||
|
|
@ -362,6 +365,9 @@ class AnthropicAdapter:
|
|||
_BlockT: Final = TypeVar("_BlockT", bound=Mapping[str, object])
|
||||
|
||||
|
||||
_CLAUDE_OVER_REASONING_EFFORT_PROVIDERS: Final = frozenset({LlmProviders.KYMA.value})
|
||||
|
||||
|
||||
class LiteLLMAnthropicMessagesAdapter:
|
||||
def __init__(self):
|
||||
pass
|
||||
|
|
@ -495,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:
|
||||
|
|
@ -528,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":
|
||||
|
|
@ -1128,7 +1131,8 @@ class LiteLLMAnthropicMessagesAdapter:
|
|||
"""Translate Anthropic thinking to either thinking or reasoning_effort.
|
||||
|
||||
A Claude-family target keeps ``thinking`` verbatim, since every bridged provider serving one
|
||||
speaks that param. Carrying its adaptive effort tier alongside takes two different params,
|
||||
speaks that param, except Kyma, which serves Claude over the OpenAI dialect and takes only
|
||||
``reasoning_effort``, so its Claude models get the non-Claude mapping. Carrying its adaptive effort tier alongside takes two different params,
|
||||
because the two are not interchangeable at the provider mapping below.
|
||||
|
||||
Bedrock takes ``output_config`` directly, which attaches the tier and leaves ``thinking``
|
||||
|
|
@ -1159,7 +1163,9 @@ class LiteLLMAnthropicMessagesAdapter:
|
|||
is_bedrock_target: Final = model.startswith(("bedrock/", "converse/", "invoke/")) or self.is_bedrock_arn_model(
|
||||
model
|
||||
)
|
||||
is_claude_target: Final = self.is_anthropic_claude_model(model) or self.is_bedrock_arn_model(model)
|
||||
is_claude_target: Final = (
|
||||
self.is_anthropic_claude_model(model) or self.is_bedrock_arn_model(model)
|
||||
) and custom_llm_provider not in _CLAUDE_OVER_REASONING_EFFORT_PROVIDERS
|
||||
output_config: Final = anthropic_message_request.get("output_config")
|
||||
|
||||
if is_claude_target:
|
||||
|
|
@ -1351,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 "")
|
||||
|
|
@ -1359,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)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -2,10 +2,21 @@ from collections.abc import AsyncIterator, Mapping, Sequence
|
|||
from typing import ClassVar, Final
|
||||
|
||||
import httpx
|
||||
from pydantic import TypeAdapter
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.anthropic_beta_headers_manager import get_unsupported_headers, is_beta_header_supported_by_bundled_config
|
||||
from litellm.anthropic_beta_gated_fields import (
|
||||
log_beta_gated_strips,
|
||||
strip_beta_gated_native_request,
|
||||
strip_beta_gated_system_messages,
|
||||
unstripped_native_request,
|
||||
)
|
||||
from litellm.anthropic_beta_headers_manager import (
|
||||
get_unsupported_headers,
|
||||
has_beta_headers_config,
|
||||
is_beta_header_supported_by_bundled_config,
|
||||
)
|
||||
from litellm.constants import ANTHROPIC_BILLING_METADATA_PREFIX
|
||||
from litellm.exceptions import AuthenticationError
|
||||
from litellm.litellm_core_utils.core_helpers import drops_unsupported_params
|
||||
|
|
@ -20,7 +31,6 @@ from litellm.types.llms.anthropic import (
|
|||
ANTHROPIC_ADVISOR_TOOL_TYPE,
|
||||
ANTHROPIC_BETA_HEADER_VALUES,
|
||||
ANTHROPIC_MID_CONVERSATION_TOOL_CHANGES_BETA_HEADER,
|
||||
ANTHROPIC_THINKING_DISPLAY_UPDATES_BETA_HEADER,
|
||||
AnthropicMessagesRequest,
|
||||
)
|
||||
from litellm.types.llms.anthropic_messages.anthropic_response import (
|
||||
|
|
@ -51,6 +61,8 @@ DEFAULT_ANTHROPIC_API_VERSION: Final = "2023-06-01"
|
|||
|
||||
_CALLER_CREDENTIAL_HEADERS: Final = frozenset({"x-api-key", "authorization"})
|
||||
|
||||
_MESSAGE_MAPPINGS: Final = TypeAdapter(tuple[Mapping[str, object], ...])
|
||||
|
||||
|
||||
def _carries_caller_credential(headers: Mapping[str, str]) -> bool:
|
||||
"""Whether the caller sent their own Anthropic credential, in which case this passthrough
|
||||
|
|
@ -85,6 +97,9 @@ def _ensure_stop_sequence(response: AnthropicMessagesResponse) -> AnthropicMessa
|
|||
return response
|
||||
|
||||
|
||||
_MESSAGES_ADAPTER: Final = TypeAdapter(list[dict[str, object]])
|
||||
|
||||
|
||||
class AnthropicMessagesConfig(BaseAnthropicMessagesConfig):
|
||||
_workload_identity_eligible: ClassVar[bool] = True
|
||||
|
||||
|
|
@ -229,7 +244,18 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig):
|
|||
else:
|
||||
return system_param
|
||||
|
||||
def _normalize_system_role_messages(self, anthropic_messages_request: dict, model: str) -> None:
|
||||
def _strip_beta_gated_system_turns(
|
||||
self, messages: Sequence[Mapping[str, object]], model: str, drop_params: object
|
||||
) -> Sequence[Mapping[str, object]]:
|
||||
if not has_beta_headers_config(self._resolved_provider):
|
||||
return messages
|
||||
gated: Final = strip_beta_gated_system_messages(messages, frozenset(), drop_params=drop_params)
|
||||
log_beta_gated_strips(gated.stripped, model=model, provider=self._resolved_provider)
|
||||
return _MESSAGE_MAPPINGS.validate_python(gated.messages) if gated.stripped else messages
|
||||
|
||||
def _normalize_system_role_messages(
|
||||
self, anthropic_messages_request: dict, model: str, drop_params: object
|
||||
) -> None:
|
||||
"""Normalize ``role: "system"`` entries in ``messages`` per the Anthropic
|
||||
``/v1/messages`` contract, which the first-party API, Bedrock Invoke,
|
||||
Vertex, and Azure Foundry all enforce identically.
|
||||
|
|
@ -277,7 +303,15 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig):
|
|||
custom_llm_provider=self.custom_llm_provider,
|
||||
key="supports_mid_conversation_system",
|
||||
)
|
||||
else list(convert_mid_conversation_system_turns(messages[leading_count:]))
|
||||
else list(
|
||||
convert_mid_conversation_system_turns(
|
||||
self._strip_beta_gated_system_turns(
|
||||
messages[leading_count:], # pyright: ignore[reportUnknownArgumentType] # messages comes from the bare request dict
|
||||
model=model,
|
||||
drop_params=drop_params,
|
||||
)
|
||||
)
|
||||
)
|
||||
)
|
||||
if hoisted or remaining != messages:
|
||||
anthropic_messages_request["messages"] = remaining
|
||||
|
|
@ -336,7 +370,7 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig):
|
|||
model=model,
|
||||
),
|
||||
)
|
||||
return self._finalize_messages_headers(headers, optional_params, messages), api_base
|
||||
return self.finalize_messages_headers(headers, optional_params, messages), api_base
|
||||
|
||||
async def avalidate_anthropic_messages_environment(
|
||||
self,
|
||||
|
|
@ -376,7 +410,7 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig):
|
|||
model=model,
|
||||
),
|
||||
)
|
||||
return self._finalize_messages_headers(oauth_headers, optional_params, messages), api_base
|
||||
return self.finalize_messages_headers(oauth_headers, optional_params, messages), api_base
|
||||
|
||||
def _require_auth_header(self, auth_header: Mapping[str, str] | None, model: str) -> Mapping[str, str]:
|
||||
if auth_header is None:
|
||||
|
|
@ -410,7 +444,7 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig):
|
|||
|
||||
return config_allows_workload_identity(self)
|
||||
|
||||
def _finalize_messages_headers(
|
||||
def finalize_messages_headers(
|
||||
self,
|
||||
headers: dict, # mutable-ok: out-param
|
||||
optional_params: dict, # mutable-ok: out-param
|
||||
|
|
@ -619,6 +653,12 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig):
|
|||
if thinking_enabled or effort_enabled:
|
||||
optional_params.pop("temperature", None)
|
||||
|
||||
def strips_beta_gated_fields_from_final_headers(self) -> bool:
|
||||
"""False on legs whose final ``anthropic_beta`` list is only known after this transform (Bedrock Invoke),
|
||||
which strip beta-gated fields themselves once that list exists, and on providers with no row in
|
||||
``anthropic_beta_headers_config.json``, whose filter drops every beta and so cannot say which fields fail."""
|
||||
return has_beta_headers_config(self._resolved_provider)
|
||||
|
||||
def transform_anthropic_messages_request(
|
||||
self,
|
||||
model: str,
|
||||
|
|
@ -720,18 +760,37 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig):
|
|||
if not _has_advisor:
|
||||
messages = strip_advisor_blocks_from_messages(messages)
|
||||
|
||||
gated: Final = (
|
||||
strip_beta_gated_native_request(
|
||||
headers, # pyright: ignore[reportUnknownArgumentType] # the base signature types this as a bare dict
|
||||
anthropic_messages_optional_request_params, # pyright: ignore[reportUnknownArgumentType] # the base signature types this as a bare dict
|
||||
messages, # pyright: ignore[reportUnknownArgumentType] # the base signature types this as a bare list
|
||||
drop_params=getattr(litellm_params, "drop_params", None),
|
||||
)
|
||||
if self.strips_beta_gated_fields_from_final_headers()
|
||||
else unstripped_native_request(anthropic_messages_optional_request_params, messages) # pyright: ignore[reportUnknownArgumentType] # the base signature types these as bare dicts
|
||||
)
|
||||
log_beta_gated_strips(gated.stripped, model=model, provider=self._resolved_provider)
|
||||
anthropic_messages_request: Final[AnthropicMessagesRequest] = AnthropicMessagesRequest(
|
||||
messages=strip_encrypted_reasoning_blocks_from_anthropic_messages(
|
||||
drop_non_string_tool_result_texts(messages, getattr(litellm_params, "drop_params", None))
|
||||
drop_non_string_tool_result_texts(
|
||||
_MESSAGES_ADAPTER.validate_python(gated.messages) if gated.messages_stripped else messages,
|
||||
getattr(litellm_params, "drop_params", None),
|
||||
)
|
||||
),
|
||||
max_tokens=max_tokens,
|
||||
model=model,
|
||||
**anthropic_messages_optional_request_params,
|
||||
)
|
||||
request: Final = (
|
||||
{**dict(anthropic_messages_request), "thinking": gated.thinking}
|
||||
if gated.thinking_stripped
|
||||
else dict(anthropic_messages_request) # pyright: ignore[reportUnknownVariableType] # built from the base signature's bare dict
|
||||
)
|
||||
if not isinstance(tools, list):
|
||||
return dict(anthropic_messages_request)
|
||||
return request
|
||||
return {
|
||||
**anthropic_messages_request,
|
||||
**request,
|
||||
"tools": self.tools_for_upstream(
|
||||
OBJECT_LIST_ADAPTER.validate_python(tools), model, litellm_params.drop_params
|
||||
),
|
||||
|
|
@ -866,17 +925,13 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig):
|
|||
if AnthropicModelInfo().is_tool_search_used(tools):
|
||||
beta_values.add(get_tool_search_beta_header(custom_llm_provider))
|
||||
|
||||
thinking_display_betas: Final = (
|
||||
(ANTHROPIC_THINKING_DISPLAY_UPDATES_BETA_HEADER,)
|
||||
if AnthropicModelInfo().is_thinking_display_updates_used(optional_params.get("thinking"))
|
||||
else ()
|
||||
)
|
||||
thinking_betas: Final = AnthropicModelInfo().thinking_betas(optional_params.get("thinking"))
|
||||
tool_change_betas: Final = (
|
||||
(ANTHROPIC_MID_CONVERSATION_TOOL_CHANGES_BETA_HEADER,)
|
||||
if AnthropicModelInfo().is_mid_conversation_tool_change_used(messages)
|
||||
else ()
|
||||
)
|
||||
all_beta_values: Final = beta_values.union(thinking_display_betas, tool_change_betas)
|
||||
all_beta_values: Final = beta_values.union(thinking_betas, tool_change_betas)
|
||||
dangerous_tool_use: Final = ANTHROPIC_BETA_HEADER_VALUES.DANGEROUS_TOOL_USE_2026_09_03.value
|
||||
forwarded_beta_values: Final = ( # pyright: ignore[reportUnknownVariableType] # derived from the bare beta set
|
||||
all_beta_values - {dangerous_tool_use}
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -661,12 +661,13 @@ def prepare_native_baseline_body(request: Mapping[str, object], model: str) -> M
|
|||
TypeAdapter(tuple[str, ...]).validate_python(owned.get("additional_drop_params") or ()),
|
||||
dict(options),
|
||||
)
|
||||
body: Final = AnthropicMessagesConfig().transform_anthropic_messages_request(
|
||||
config: Final = AnthropicMessagesConfig()
|
||||
body: Final = config.transform_anthropic_messages_request(
|
||||
model=resolved_model,
|
||||
messages=strip_provider_specific_fields_from_anthropic_messages(messages),
|
||||
anthropic_messages_optional_request_params=filtered,
|
||||
litellm_params=GenericLiteLLMParams(),
|
||||
headers={},
|
||||
headers=config.finalize_messages_headers({}, dict(filtered), list(messages)), # pyright: ignore[reportUnknownMemberType] # the shared helper takes and returns bare dicts
|
||||
)
|
||||
return MappingProxyType(_JSON_OBJECT.validate_python(body))
|
||||
except Exception: # noqa: BLE001 # an unsupported hypothetical request is unavailable, never an inference failure
|
||||
|
|
|
|||
|
|
@ -4,11 +4,12 @@ Azure Anthropic handler - reuses AnthropicChatCompletion logic with Azure authen
|
|||
|
||||
import copy
|
||||
import json
|
||||
from collections.abc import Callable
|
||||
from collections.abc import Callable, Mapping
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
|
||||
from litellm.anthropic_beta_gated_fields import strip_beta_gated_request_body
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.llms.anthropic.chat.handler import AnthropicChatCompletion
|
||||
from litellm.llms.anthropic.chat.transformation import JsonModeRequest
|
||||
|
|
@ -75,13 +76,22 @@ class AzureAnthropicChatCompletion(AnthropicChatCompletion):
|
|||
litellm_params=litellm_params,
|
||||
)
|
||||
|
||||
data: Final = config.transform_request(
|
||||
request_body: Final[Mapping[str, object]] = config.transform_request(
|
||||
model=model,
|
||||
messages=messages,
|
||||
optional_params=optional_params,
|
||||
litellm_params=litellm_params,
|
||||
headers=headers,
|
||||
)
|
||||
data: Final = dict(
|
||||
strip_beta_gated_request_body(
|
||||
request_body,
|
||||
headers, # pyright: ignore[reportUnknownArgumentType] # validate_environment returns a bare dict
|
||||
drop_params=litellm_params.get("drop_params"), # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # litellm_params is the handler's bare dict
|
||||
model=model,
|
||||
provider=custom_llm_provider,
|
||||
)
|
||||
)
|
||||
|
||||
## LOGGING
|
||||
logging_obj.pre_call(
|
||||
|
|
|
|||
|
|
@ -145,6 +145,8 @@ class AzureAnthropicMessagesConfig(AnthropicMessagesConfig):
|
|||
litellm_params=litellm_params,
|
||||
headers=headers,
|
||||
)
|
||||
self._normalize_system_role_messages(anthropic_messages_request, model=model)
|
||||
self._normalize_system_role_messages(
|
||||
anthropic_messages_request, model=model, drop_params=getattr(litellm_params, "drop_params", None)
|
||||
)
|
||||
self._remove_scope_from_cache_control(anthropic_messages_request)
|
||||
return anthropic_messages_request
|
||||
|
|
|
|||
|
|
@ -107,9 +107,11 @@ class AzureAnthropicConfig(AnthropicConfig):
|
|||
is_vertex_request=optional_params.get("is_vertex_request", False),
|
||||
user_anthropic_beta_headers=user_anthropic_beta_headers,
|
||||
mcp_server_used=mcp_server_used,
|
||||
thinking=optional_params.get("thinking"), # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # optional_params is the base config's bare dict
|
||||
)
|
||||
# Merge headers - Azure auth (api-key or Authorization) takes precedence
|
||||
headers = {**anthropic_headers, **headers}
|
||||
merged_beta_header: Final = {k: v for k, v in anthropic_headers.items() if k == "anthropic-beta"}
|
||||
headers = {**anthropic_headers, **headers, **merged_beta_header}
|
||||
|
||||
# Ensure anthropic-version header is set
|
||||
if "anthropic-version" not in headers:
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -5,8 +5,9 @@ from types import MappingProxyType
|
|||
from typing import TYPE_CHECKING, Final
|
||||
|
||||
import httpx
|
||||
from pydantic import TypeAdapter, ValidationError
|
||||
from pydantic import ConfigDict, TypeAdapter, ValidationError
|
||||
|
||||
from litellm._lazy_imports import get_default_encoding
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.llms.azure_ai.common_utils import (
|
||||
AzureFoundryModelInfo,
|
||||
|
|
@ -22,10 +23,15 @@ from litellm.llms.base_llm.passthrough.transformation import (
|
|||
relayed_body,
|
||||
strip_leading_model_segment,
|
||||
)
|
||||
from litellm.llms.cohere.embed.transformation import (
|
||||
CohereEmbeddingRowShape,
|
||||
CohereImageTokenBilling,
|
||||
cohere_embedding_response,
|
||||
)
|
||||
from litellm.rust_bridge.ocr.entrypoints import NATIVE_OCR_PASSTHROUGH_RESPONSE, NativeOcrPassthroughResponse
|
||||
from litellm.types.llms.openai import AllMessageValues
|
||||
from litellm.types.rerank import RerankResponse
|
||||
from litellm.types.utils import CallTypes, ImageResponse, StandardPassThroughResponseObject
|
||||
from litellm.types.utils import CallTypes, EmbeddingResponse, ImageResponse, StandardPassThroughResponseObject
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from httpx import URL, Response
|
||||
|
|
@ -35,6 +41,7 @@ if TYPE_CHECKING:
|
|||
|
||||
|
||||
EMPTY_QUERY: Final[Mapping[str, object]] = MappingProxyType({})
|
||||
ECHOED_TEXTS: Final = TypeAdapter(tuple[str, ...], config=ConfigDict(hide_input_in_errors=True))
|
||||
|
||||
|
||||
def api_version_from(litellm_params: Mapping[str, object]) -> str | None:
|
||||
|
|
@ -85,9 +92,29 @@ def relay_query_params(
|
|||
return MappingProxyType({**(request_query_params or EMPTY_QUERY), "api-version": api_version})
|
||||
|
||||
|
||||
def echoed_texts(body: Mapping[str, object]) -> tuple[str, ...]:
|
||||
try:
|
||||
return ECHOED_TEXTS.validate_python(body.get("texts"))
|
||||
except ValidationError:
|
||||
return ()
|
||||
|
||||
|
||||
def cohere_embed_relay_response(body: Mapping[str, object]) -> EmbeddingResponse:
|
||||
return cohere_embedding_response(
|
||||
body,
|
||||
None,
|
||||
get_default_encoding(),
|
||||
echoed_texts(body),
|
||||
CohereImageTokenBilling.FOLDED_INTO_INPUT_TOKENS,
|
||||
CohereEmbeddingRowShape.TYPED,
|
||||
)
|
||||
|
||||
|
||||
FOUNDRY_RELAY_SHAPES: Final = (
|
||||
RelayShape("/rerank", CallTypes.arerank, RerankResponse.model_validate),
|
||||
RelayShape("/providers/blackforestlabs/v1/flux-2-pro", CallTypes.aimage_generation, ImageResponse.model_validate),
|
||||
RelayShape("/providers/cohere/v1/embed", CallTypes.aembedding, cohere_embed_relay_response),
|
||||
RelayShape("/providers/cohere/v2/embed", CallTypes.aembedding, cohere_embed_relay_response),
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
Some files were not shown because too many files have changed in this diff Show more
Loading…
Add table
Reference in a new issue