scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-14 23:21:35 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 102

merge: bring litellm_internal_staging into litellm_bedrock_mantle_govcloud_cost_row

Browse source
This commit is contained in:
mateo-berri 2026-09-10 07:08:58 -07:00
commit 0308b05c7a
2100 changed files with 119350 additions and 23210 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

27
.circleci/config.yml
View file

@ -1440,6 +1440,7 @@ jobs:
TEST_FILES=$(printf "%s\n" \
tests/local_testing/test_dual_cache.py \
tests/local_testing/test_redis_batch_optimizations.py \
tests/local_testing/test_redis_increment_with_floor.py \
tests/local_testing/test_router_utils.py)
echo "$TEST_FILES" | circleci tests run \
--verbose \
@ -2648,6 +2649,19 @@ jobs:
name: Start mock LLM server
command: uv run --no-sync python tests/e2e/ui/fixtures/mock_llm_server/server.py
background: true
- run:
name: Start mock Presidio server
command: uv run --no-sync python tests/e2e/ui/fixtures/mock_presidio_server/server.py
background: true
- run:
name: Wait for mock Presidio server
command: |
for i in $(seq 1 30); do
if curl -sf http://127.0.0.1:8091/health >/dev/null 2>&1; then exit 0; fi
sleep 1
done
echo "Mock Presidio server never answered /health on port 8091" >&2
exit 1
- run:
name: Start LiteLLM proxy
environment:
@ -2778,6 +2792,19 @@ jobs:
name: Start mock LLM server
command: uv run --no-sync python tests/e2e/ui/fixtures/mock_llm_server/server.py
background: true
- run:
name: Start mock Presidio server
command: uv run --no-sync python tests/e2e/ui/fixtures/mock_presidio_server/server.py
background: true
- run:
name: Wait for mock Presidio server
command: |
for i in $(seq 1 30); do
if curl -sf http://127.0.0.1:8091/health >/dev/null 2>&1; then exit 0; fi
sleep 1
done
echo "Mock Presidio server never answered /health on port 8091" >&2
exit 1
- run:
name: Start LiteLLM proxy under a server root path
environment:

38
.github/e2e-stack/assert_tests_ran.py vendored Normal file
View file

@ -0,0 +1,38 @@
import sys
import xml.etree.ElementTree as ET
from pathlib import Path
from typing import Final
def main() -> int:
selected: Final = tuple(sys.argv[2:])
try:
report: Final = ET.parse(Path(sys.argv[1])).getroot()
except (ET.ParseError, OSError):
_ = sys.stdout.write("::error::could not read the test execution report\n")
return 1
cases: Final = tuple(report.iter("testcase"))
passed: Final = frozenset(
case.get("file") for case in cases if all(case.find(tag) is None for tag in ("skipped", "failure", "error"))
)
missing: Final = tuple(path for path in selected if path not in passed)
for path in selected:
collected: Final = sum(case.get("file") == path for case in cases)
skipped: Final = sum(case.get("file") == path and case.find("skipped") is not None for case in cases)
_ = sys.stdout.write(f"{path}: {collected} collected, {skipped} skipped\n")
for case in cases:
if case.get("file") != path or all(case.find(tag) is None for tag in ("failure", "error")):
continue
_ = sys.stdout.write(f" failed: {case.get('classname', '')}::{case.get('name', '')}\n")
if (
selected
and not missing
and not any(case.find(tag) is not None for case in cases for tag in ("failure", "error"))
):
return 0
_ = sys.stdout.write("::error::every selected file must execute a passing test, with no failures or errors\n")
return 1
if __name__ == "__main__":
sys.exit(main())

17
.github/e2e-stack/down.sh vendored Executable file
View file

@ -0,0 +1,17 @@
#!/usr/bin/env bash
set -uo pipefail
STACK_DIR="${E2E_STACK_DIR:-${RUNNER_TEMP:-/tmp}/litellm-e2e-stack}"
for pid_file in "${STACK_DIR}"/pids/*.pid; do
[[ -f "${pid_file}" ]] || continue
pkill -TERM -P "$(cat "${pid_file}")" 2>/dev/null
kill -TERM "$(cat "${pid_file}")" 2>/dev/null
rm -f "${pid_file}"
done
for container in e2e-nginx e2e-valkey e2e-jaeger e2e-postgres; do
docker rm -f "${container}" >/dev/null 2>&1
done
exit 0

49
.github/e2e-stack/secrets_to_env.py vendored Normal file
View file

@ -0,0 +1,49 @@
import os
import re
import sys
from pathlib import Path
from typing import Final
from pydantic import TypeAdapter, ValidationError
secrets_adapter: Final[TypeAdapter[dict[str, str]]] = TypeAdapter(dict[str, str])
ENV_NAME: Final = re.compile(r"[A-Za-z_][A-Za-z0-9_]*")
MIN_MASKED_LENGTH: Final = 8
def main() -> int:
env_path: Final = Path(sys.argv[1])
try:
secrets: Final = {
key: value.rstrip("\r\n") for key, value in secrets_adapter.validate_json(sys.stdin.read()).items()
}
except (ValidationError, UnicodeError):
_ = sys.stderr.write("expected a JSON object containing string environment values\n")
return 1
unusable: Final = tuple(
key
for key, value in secrets.items()
if ENV_NAME.fullmatch(key) is None or any(char in value for char in "'\n\r\0")
)
if unusable:
_ = sys.stderr.write(
f"these names or values cannot be represented in both bash and dotenv: {' '.join(sorted(unusable))}\n"
)
return 1
for value in secrets.values():
if len(value) >= MIN_MASKED_LENGTH:
_ = sys.stdout.write(f"::add-mask::{value.replace('%', '%25')}\n")
sys.stdout.flush()
lines: Final = tuple(f"{key}='{value}'" for key, value in secrets.items() if value)
try:
with os.fdopen(os.open(env_path, os.O_WRONLY | os.O_APPEND | os.O_CREAT | os.O_NOFOLLOW, 0o600), "w") as handle:
os.fchmod(handle.fileno(), 0o600)
_ = handle.write("\n".join(lines) + "\n")
except OSError:
_ = sys.stderr.write("could not write the environment file\n")
return 1
return 0
if __name__ == "__main__":
sys.exit(main())

44
.github/e2e-stack/select_tests.py vendored Normal file
View file

@ -0,0 +1,44 @@
import re
import sys
from typing import Final
SELECTABLE: Final = re.compile(r"^tests/e2e/([A-Za-z0-9_.-]+/)*test_[A-Za-z0-9_.-]+\.py$")
UNSUPPORTED: Final = re.compile(
r"^tests/e2e/(ui|claude_code|load)/"
r"|^tests/e2e/llm_translation/realtime/test_realtime_pipecat_audio_e2e\.py$"
r"|^tests/e2e/batches/test_managed_files_enforcement_e2e\.py$"
r"|^tests/e2e/guardrails/test_presidio_masking_e2e\.py$"
)
HARNESS: Final = re.compile(
r"^tests/e2e/[A-Za-z0-9_.-]+\.(py|ini)$"
r"|^tests/e2e/gateway/"
r"|^\.github/e2e-stack/"
r"|^\.github/workflows/test-e2e-changed\.yml$"
)
UNEXPANDED: Final = re.compile(r"[*?\[]")
def is_selectable(path: str) -> bool:
return SELECTABLE.match(path) is not None and UNSUPPORTED.match(path) is None
def select(changed: tuple[str, ...], canary: tuple[str, ...]) -> tuple[str, ...]:
direct: Final = frozenset(path for path in changed if is_selectable(path))
harness_changed: Final = any(HARNESS.match(path) for path in changed)
canary_tests: Final = frozenset(path for path in canary if harness_changed and is_selectable(path))
return tuple(sorted(direct | canary_tests))
def main() -> int:
canary: Final = tuple(sys.argv[1:])
unexpanded: Final = tuple(path for path in canary if UNEXPANDED.search(path))
if unexpanded:
_ = sys.stderr.write(f"the canary paths reached the selector unexpanded: {' '.join(unexpanded)}\n")
return 1
changed: Final = tuple(line.strip() for line in sys.stdin if line.strip())
_ = sys.stdout.write(" ".join(select(changed, canary)) + "\n")
return 0
if __name__ == "__main__":
sys.exit(main())

207
.github/e2e-stack/up.sh vendored Executable file
View file

@ -0,0 +1,207 @@
#!/usr/bin/env bash
set -euo pipefail
umask 077
REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)"
STACK_DIR="${E2E_STACK_DIR:-${RUNNER_TEMP:-/tmp}/litellm-e2e-stack}"
CERTS_DIR="${STACK_DIR}/certs"
LOGS_DIR="${STACK_DIR}/logs"
PIDS_DIR="${STACK_DIR}/pids"
POSTGRES_IMAGE="${E2E_POSTGRES_IMAGE:-postgres:16.6}"
VALKEY_IMAGE="${E2E_VALKEY_IMAGE:-valkey/valkey:8.1.4@sha256:81db6d39e1bba3b3ff32bd3a1b19a6d69690f94a3954ec131277b9a26b95b3aa}"
JAEGER_IMAGE="${E2E_JAEGER_IMAGE:-jaegertracing/jaeger:2.10.0}"
NGINX_IMAGE="${E2E_NGINX_IMAGE:-nginx:1.29.1-alpine@sha256:42a516af16b852e33b7682d5ef8acbd5d13fe08fecadc7ed98605ba5e3b26ab8}"
LB_PORT="${E2E_LB_PORT:-4000}"
GATEWAY_PORT_1="${E2E_GATEWAY_PORT_1:-4010}"
GATEWAY_PORT_2="${E2E_GATEWAY_PORT_2:-4011}"
BACKEND_PORT="${E2E_BACKEND_PORT:-4001}"
REDIS_PORT="${E2E_REDIS_PORT:-6379}"
DATABASE_HOST="${E2E_DATABASE_HOST:-127.0.0.1}"
DATABASE_PORT="${E2E_DATABASE_PORT:-5432}"
DATABASE_USER="${E2E_DATABASE_USER:-litellm}"
DATABASE_PASSWORD="${E2E_DATABASE_PASSWORD:-dbpassword9090}"
DATABASE_NAME="${E2E_DATABASE_NAME:-litellm}"
JAEGER_OTLP_PORT="${E2E_JAEGER_OTLP_PORT:-4318}"
JAEGER_QUERY_PORT="${E2E_JAEGER_QUERY_PORT:-16686}"
MASTER_KEY="${LITELLM_MASTER_KEY:-sk-e2e-$(openssl rand -hex 16)}"
mkdir -p "${CERTS_DIR}" "${LOGS_DIR}" "${PIDS_DIR}"
chmod 700 "${STACK_DIR}" "${LOGS_DIR}" "${PIDS_DIR}"
chmod 755 "${CERTS_DIR}"
log() { printf 'e2e-stack: %s\n' "$*"; }
port_open() { (exec 3<>"/dev/tcp/127.0.0.1/$1") 2>/dev/null; }
wait_for() {
local label="$1" check="$2" deadline=$((SECONDS + ${3:-120}))
until eval "${check}"; do
if ((SECONDS >= deadline)); then
log "timed out waiting for ${label}"
exit 1
fi
sleep 2
done
log "${label} is up"
}
if [[ -f "${REPO_ROOT}/tests/e2e/.env" ]]; then
set -a
source "${REPO_ROOT}/tests/e2e/.env"
set +a
fi
if [[ -z "${DD_API_KEY:-}" ]]; then
log "DD_API_KEY is empty; the gateway config enables the datadog callback, so put a Datadog API key in tests/e2e/.env"
exit 1
fi
export DD_SITE="${DD_SITE:-datadoghq.com}"
if ! port_open "${DATABASE_PORT}"; then
docker run -d --name e2e-postgres -p "${DATABASE_PORT}:5432" \
-e "POSTGRES_USER=${DATABASE_USER}" -e "POSTGRES_PASSWORD=${DATABASE_PASSWORD}" -e "POSTGRES_DB=${DATABASE_NAME}" \
"${POSTGRES_IMAGE}" >/dev/null
fi
wait_for "postgres" "port_open ${DATABASE_PORT}"
if ! port_open "${JAEGER_QUERY_PORT}"; then
docker run -d --name e2e-jaeger -p "${JAEGER_OTLP_PORT}:4318" -p "${JAEGER_QUERY_PORT}:16686" \
"${JAEGER_IMAGE}" >/dev/null
fi
wait_for "jaeger" "curl -fs http://127.0.0.1:${JAEGER_QUERY_PORT}/api/services >/dev/null"
openssl genrsa -out "${CERTS_DIR}/ca.key" 2048 2>/dev/null
openssl req -x509 -new -nodes -key "${CERTS_DIR}/ca.key" -sha256 -days 7 \
-subj "/CN=litellm-e2e-ca" \
-addext "basicConstraints=critical,CA:TRUE" -addext "keyUsage=critical,keyCertSign,cRLSign" \
-out "${CERTS_DIR}/ca.crt" 2>/dev/null
openssl genrsa -out "${CERTS_DIR}/server.key" 2048 2>/dev/null
openssl req -new -key "${CERTS_DIR}/server.key" -subj "/CN=localhost" -out "${CERTS_DIR}/server.csr" 2>/dev/null
openssl x509 -req -in "${CERTS_DIR}/server.csr" -CA "${CERTS_DIR}/ca.crt" -CAkey "${CERTS_DIR}/ca.key" \
-CAcreateserial -days 7 -sha256 \
-extfile <(printf 'basicConstraints=CA:FALSE\nkeyUsage=critical,digitalSignature,keyEncipherment\nextendedKeyUsage=serverAuth\nsubjectAltName=DNS:localhost,IP:127.0.0.1\n') \
-out "${CERTS_DIR}/server.crt" 2>/dev/null
chmod 644 "${CERTS_DIR}"/*.key "${CERTS_DIR}"/*.crt
CERTIFI_BUNDLE="$(cd "${REPO_ROOT}" && uv run --no-sync python -c 'import certifi; print(certifi.where())')"
cat "${CERTIFI_BUNDLE}" "${CERTS_DIR}/ca.crt" > "${CERTS_DIR}/ca-bundle.pem"
docker rm -f e2e-valkey >/dev/null 2>&1 || true
docker run -d --name e2e-valkey -p "${REDIS_PORT}:${REDIS_PORT}" -v "${CERTS_DIR}:/certs:ro" \
"${VALKEY_IMAGE}" valkey-server \
--cluster-enabled yes --port 0 --tls-port "${REDIS_PORT}" \
--tls-cert-file /certs/server.crt --tls-key-file /certs/server.key --tls-ca-cert-file /certs/ca.crt \
--tls-auth-clients no --cluster-announce-ip 127.0.0.1 >/dev/null
VALKEY_CLI="docker exec e2e-valkey valkey-cli --tls --cacert /certs/ca.crt -h 127.0.0.1 -p ${REDIS_PORT}"
wait_for "valkey" "${VALKEY_CLI} ping 2>/dev/null | grep -q PONG"
${VALKEY_CLI} cluster addslotsrange 0 16383 >/dev/null
wait_for "valkey cluster" "${VALKEY_CLI} cluster info 2>/dev/null | grep -q cluster_state:ok"
CONFIG_SOURCE="${REPO_ROOT}/tests/e2e/gateway/stage_mirror_ci_config.yml"
CONFIG_PATH="${CONFIG_SOURCE}"
if [[ "${REDIS_PORT}" != "6379" ]]; then
CONFIG_PATH="${STACK_DIR}/litellm-config.yml"
sed "s/port: 6379/port: ${REDIS_PORT}/" "${CONFIG_SOURCE}" > "${CONFIG_PATH}"
fi
SERVER_ENV=(
"LITELLM_MASTER_KEY=${MASTER_KEY}"
"DATABASE_HOST=${DATABASE_HOST}"
"DATABASE_PORT=${DATABASE_PORT}"
"DATABASE_USER=${DATABASE_USER}"
"DATABASE_PASSWORD=${DATABASE_PASSWORD}"
"DATABASE_NAME=${DATABASE_NAME}"
"DISABLE_SCHEMA_UPDATE=true"
"REDIS_HOST=127.0.0.1"
"REDIS_PORT=${REDIS_PORT}"
"REDIS_CLUSTER_NODES=[{\"host\":\"127.0.0.1\",\"port\":${REDIS_PORT}}]"
"CONFIG_FILE_PATH=${CONFIG_PATH}"
"STORE_MODEL_IN_DB=True"
"OTEL_EXPORTER_OTLP_PROTOCOL=http/protobuf"
"OTEL_EXPORTER_OTLP_ENDPOINT=http://127.0.0.1:${JAEGER_OTLP_PORT}"
"SSL_CERT_FILE=${CERTS_DIR}/ca-bundle.pem"
"PYTHONPATH=${REPO_ROOT}"
)
if [[ -n "${VERTEXAI_CREDENTIALS:-}" ]]; then
printf '%s' "${VERTEXAI_CREDENTIALS}" > "${STACK_DIR}/vertex-adc.json"
SERVER_ENV+=("GOOGLE_APPLICATION_CREDENTIALS=${STACK_DIR}/vertex-adc.json")
fi
cd "${REPO_ROOT}"
log "running migrations"
env "${SERVER_ENV[@]}" uv run --no-sync python migrations/run.py >"${LOGS_DIR}/migrations.log" 2>&1
start_server() {
local name="$1"; shift
env "${SERVER_ENV[@]}" "$@" >"${LOGS_DIR}/${name}.log" 2>&1 &
echo $! > "${PIDS_DIR}/${name}.pid"
}
start_server backend uv run --no-sync uvicorn backend.main:app --host 0.0.0.0 --port "${BACKEND_PORT}"
start_server gateway-1 uv run --no-sync uvicorn gateway.main:app --workers 1 --host 0.0.0.0 --port "${GATEWAY_PORT_1}"
start_server gateway-2 uv run --no-sync uvicorn gateway.main:app --workers 1 --host 0.0.0.0 --port "${GATEWAY_PORT_2}"
if [[ "$(uname)" == "Linux" ]]; then
NGINX_UPSTREAM_HOST=127.0.0.1
NGINX_DOCKER_ARGS=(--network host)
else
NGINX_UPSTREAM_HOST=host.docker.internal
NGINX_DOCKER_ARGS=(-p "${LB_PORT}:${LB_PORT}")
fi
cat > "${STACK_DIR}/nginx.conf" <<EOF
events {}
http {
map \$http_upgrade \$connection_upgrade {
default upgrade;
'' close;
}
upstream litellm_gateways {
server ${NGINX_UPSTREAM_HOST}:${GATEWAY_PORT_1};
server ${NGINX_UPSTREAM_HOST}:${GATEWAY_PORT_2};
}
server {
listen ${LB_PORT};
client_max_body_size 100m;
location / {
proxy_pass http://litellm_gateways;
proxy_http_version 1.1;
proxy_set_header Host \$host;
proxy_set_header X-Forwarded-For \$proxy_add_x_forwarded_for;
proxy_set_header X-Forwarded-Proto \$scheme;
proxy_set_header Upgrade \$http_upgrade;
proxy_set_header Connection \$connection_upgrade;
proxy_buffering off;
proxy_read_timeout 600s;
proxy_send_timeout 600s;
}
}
}
EOF
docker rm -f e2e-nginx >/dev/null 2>&1 || true
docker run -d --name e2e-nginx "${NGINX_DOCKER_ARGS[@]}" \
-v "${STACK_DIR}/nginx.conf:/etc/nginx/nginx.conf:ro" "${NGINX_IMAGE}" >/dev/null
wait_for "backend" "curl -fs http://127.0.0.1:${BACKEND_PORT}/health/liveliness >/dev/null" 300
wait_for "gateway-1" "curl -fs http://127.0.0.1:${GATEWAY_PORT_1}/health/liveliness >/dev/null" 300
wait_for "gateway-2" "curl -fs http://127.0.0.1:${GATEWAY_PORT_2}/health/liveliness >/dev/null" 300
wait_for "load balancer" "curl -fs http://127.0.0.1:${LB_PORT}/health/liveliness >/dev/null" 60
cat > "${STACK_DIR}/stack.env" <<EOF
LITELLM_PROXY_URL=http://127.0.0.1:${LB_PORT}
LITELLM_CONTROL_PLANE_URL=http://127.0.0.1:${BACKEND_PORT}
LITELLM_PROXY_REPLICA_URLS=http://127.0.0.1:${GATEWAY_PORT_1},http://127.0.0.1:${GATEWAY_PORT_2}
LITELLM_MASTER_KEY=${MASTER_KEY}
REDIS_HOST=127.0.0.1
REDIS_PORT=${REDIS_PORT}
E2E_OTEL_QUERY_URL=http://127.0.0.1:${JAEGER_QUERY_PORT}
SSL_CERT_FILE=${CERTS_DIR}/ca-bundle.pem
DATABASE_URL=postgresql://${DATABASE_USER}:${DATABASE_PASSWORD}@${DATABASE_HOST}:${DATABASE_PORT}/${DATABASE_NAME}
EOF
log "stack is up; pytest env written to ${STACK_DIR}/stack.env"

21
.github/workflows/_test-unit-base.yml vendored
View file

@ -55,17 +55,14 @@ on:
permissions:
contents: read
env:
UV_PYTHON: "3.12"
jobs:
run:
name: ${{ matrix.python-version == '3.12' && 'Run tests' || format('Run tests (Python {0})', matrix.python-version) }}
name: Run tests
runs-on: ubuntu-latest
timeout-minutes: ${{ inputs.job-timeout-minutes }}
strategy:
fail-fast: false
matrix:
python-version: ["3.10", "3.11", "3.12", "3.13", "3.14"]
env:
UV_PYTHON: ${{ matrix.python-version }}
permissions:
contents: read
pull-requests: read
@ -88,7 +85,7 @@ jobs:
timeout-minutes: 3
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
with:
python-version: ${{ matrix.python-version }}
python-version: ${{ env.UV_PYTHON }}
- name: Set up uv
if: steps.changes.outputs.decision != 'skip'
@ -103,9 +100,9 @@ jobs:
uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0
with:
path: ${{ env.UV_CACHE_DIR }}
key: ${{ runner.os }}-uv-downloads-py${{ matrix.python-version }}-${{ hashFiles('uv.lock') }}
key: ${{ runner.os }}-uv-downloads-py${{ env.UV_PYTHON }}-${{ hashFiles('uv.lock') }}
restore-keys: |
${{ runner.os }}-uv-downloads-py${{ matrix.python-version }}-
${{ runner.os }}-uv-downloads-py${{ env.UV_PYTHON }}-
- name: Cache the Rust build
if: steps.changes.outputs.decision != 'skip'
@ -139,7 +136,7 @@ jobs:
WORKERS: ${{ inputs.workers }}
RERUNS: ${{ inputs.reruns }}
DIST: ${{ inputs.dist }}
COVERAGE_CORE: ${{ contains(fromJSON('["3.10", "3.11"]'), matrix.python-version) && 'ctrace' || 'sysmon' }}
COVERAGE_CORE: sysmon
run: |
if [ "${WORKERS}" = "0" ]; then
uv run --no-sync pytest ${TEST_PATH:?} \
@ -166,7 +163,7 @@ jobs:
fi
- name: Save coverage report
if: always() && matrix.python-version == '3.12' && steps.changes.outputs.decision != 'skip'
if: always() && steps.changes.outputs.decision != 'skip'
uses: actions/upload-artifact@4cec3d8aa04e39d1a68397de0c4cd6fb9dce8ec1 # v4.6.1
with:
name: coverage-${{ inputs.artifact-name }}-${{ github.run_id }}-${{ github.run_attempt }}

2
.github/workflows/auto-close-duplicates.yml vendored
View file

@ -22,7 +22,7 @@ on:
permissions: {}
jobs:
test:
sweep-tests:
if: github.event_name == 'pull_request'
runs-on: ubuntu-latest
timeout-minutes: 5

45
.github/workflows/cost-map-guard.yml vendored Normal file
View file

@ -0,0 +1,45 @@
name: Cost map guard
on: # zizmor: ignore[dangerous-triggers] runs the base branch's code only; the PR's cost map files are read as data and never executed
pull_request_target:
branches:
- main
- litellm_internal_staging
- litellm_oss_staging
- "litellm_**"
permissions:
contents: read
concurrency:
group: ${{ github.workflow }}-${{ github.event.pull_request.number }}
cancel-in-progress: true
jobs:
cost-map-guard:
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
persist-credentials: false
- name: Fetch the pull request head and its merge base
id: revisions
env:
GH_TOKEN: ${{ github.token }}
BASE_SHA: ${{ github.event.pull_request.base.sha }}
HEAD_SHA: ${{ github.event.pull_request.head.sha }}
run: |
merge_base="$(gh api "repos/${GITHUB_REPOSITORY}/compare/${BASE_SHA}...${HEAD_SHA}" --jq '.merge_base_commit.sha')"
git fetch --no-tags --depth=1 origin "$merge_base" "$HEAD_SHA"
echo "merge_base=$merge_base" >> "$GITHUB_OUTPUT"
- name: Set up uv
uses: ./.github/actions/setup-uv-with-retries
with:
version: "0.10.9"
- name: Run the guard
env:
MERGE_BASE: ${{ steps.revisions.outputs.merge_base }}
HEAD_SHA: ${{ github.event.pull_request.head.sha }}
HEAD_REF: ${{ github.event.pull_request.head.ref }}
run: |
uv run --frozen python ci_cd/cost_map_guard.py --base "$MERGE_BASE" --head "$HEAD_SHA" --head-ref "$HEAD_REF"

6
.github/workflows/publish-basedpyright-base-counts.yml vendored
View file

@ -1,6 +1,6 @@
name: Publish basedpyright base counts
# Every commit on litellm_internal_staging is some branch's future merge-base.
# Every commit on main or litellm_internal_staging can become a future merge-base.
# Publishing its per-rule basedpyright counts as an artifact lets
# scripts/type_check_gate.py download them in seconds instead of paying a
# 60-110s second basedpyright pass on every fresh worktree or moved merge-base.
@ -10,13 +10,13 @@ name: Publish basedpyright base counts
on:
push:
branches:
- main
- litellm_internal_staging
workflow_dispatch:
inputs:
ref:
description: "Ref to compute and publish base counts for"
description: "Ref to compute and publish base counts for (defaults to the workflow run's commit)"
required: false
default: litellm_internal_staging
permissions:
contents: read

130
.github/workflows/report-rust-release-wheel.yml vendored
View file

@ -1,130 +0,0 @@
name: Report LiteLLM Rust release wheel
on: # zizmor: ignore[dangerous-triggers] reporter executes no PR code and consumes no PR artifacts or outputs
workflow_run:
workflows:
- LiteLLM Rust
types:
- completed
permissions: {}
concurrency:
group: ${{ github.workflow }}-${{ github.event.workflow_run.pull_requests[0].number || github.event.workflow_run.id }}
cancel-in-progress: false
jobs:
report-release-wheel:
name: report release wheel
if: >-
github.event.workflow_run.event == 'pull_request' &&
github.event.workflow_run.path == '.github/workflows/test-rust.yml' &&
github.event.workflow_run.head_repository.full_name == github.repository &&
github.event.workflow_run.pull_requests[0].number != null
runs-on: ubuntu-latest
timeout-minutes: 5
permissions:
issues: write # PR comments use the issues API
pull-requests: read # Current-head validation rejects stale workflow runs
steps:
- name: Link release wheel report on PR
uses: actions/github-script@f28e40c7f34bde8b3046d885e986cb6290c5673b # v7.1.0
env:
COMMENT_MARKER: "<!-- litellm-release-wheel-size -->"
with:
script: |
const marker = process.env.COMMENT_MARKER;
const workflowRun = context.payload.workflow_run;
const allowedConclusions = new Set([
"action_required",
"cancelled",
"failure",
"neutral",
"skipped",
"stale",
"startup_failure",
"success",
"timed_out",
]);
if (
!allowedConclusions.has(workflowRun.conclusion) ||
workflowRun.event !== "pull_request" ||
workflowRun.path !== ".github/workflows/test-rust.yml" ||
workflowRun.head_repository?.full_name !==
`${context.repo.owner}/${context.repo.repo}` ||
workflowRun.pull_requests?.length !== 1
) {
throw new Error("unexpected source workflow");
}
const pullRequest = workflowRun.pull_requests[0];
const pullRequestNumber = pullRequest.number;
const headSha = workflowRun.head_sha;
const runId = workflowRun.id;
if (
!Number.isSafeInteger(pullRequestNumber) ||
pullRequestNumber <= 0 ||
!Number.isSafeInteger(runId) ||
runId <= 0 ||
!/^[0-9a-f]{40}$/.test(headSha) ||
pullRequest.head?.sha !== headSha
) {
throw new Error("invalid source workflow metadata");
}
const runUrl =
`${context.serverUrl}/${context.repo.owner}/${context.repo.repo}` +
`/actions/runs/${runId}`;
const result =
workflowRun.conclusion === "success"
? "successfully"
: `with \`${workflowRun.conclusion}\``;
const body = [
marker,
"## LiteLLM Rust workflow",
"",
`Workflow completed ${result} for \`${headSha}\``,
"",
`[View workflow run](${runUrl})`,
].join("\n");
const comments = await github.paginate(github.rest.issues.listComments, {
owner: context.repo.owner,
repo: context.repo.repo,
issue_number: pullRequestNumber,
per_page: 100,
});
const existing = comments.find(
(comment) =>
comment.user?.login === "github-actions[bot]" &&
comment.body?.startsWith(marker),
);
const currentPullRequest = (
await github.rest.pulls.get({
owner: context.repo.owner,
repo: context.repo.repo,
pull_number: pullRequestNumber,
})
).data;
if (
currentPullRequest.state !== "open" ||
currentPullRequest.head.repo?.full_name !==
`${context.repo.owner}/${context.repo.repo}` ||
currentPullRequest.head.sha !== headSha
) {
core.info("source workflow no longer matches the current pull request head");
return;
}
if (existing) {
await github.rest.issues.updateComment({
owner: context.repo.owner,
repo: context.repo.repo,
comment_id: existing.id,
body,
});
} else {
await github.rest.issues.createComment({
owner: context.repo.owner,
repo: context.repo.repo,
issue_number: pullRequestNumber,
body,
});
}

6
.github/workflows/sync-together-ai-models.yml vendored
View file

@ -13,10 +13,12 @@ jobs:
sync_together_ai_models:
if: github.repository == 'BerriAI/litellm'
runs-on: ubuntu-latest
env:
BASE_BRANCH: ${{ github.event.repository.default_branch }}
steps:
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
ref: litellm_internal_staging
ref: ${{ env.BASE_BRANCH }}
persist-credentials: false
- name: Set up uv
uses: ./.github/actions/setup-uv-with-retries
@ -63,6 +65,6 @@ jobs:
gh pr create --title "feat(models): sync together_ai model registry" \
--body-file "$RUNNER_TEMP/pr_body.md" \
--head "$branch" \
--base litellm_internal_staging
--base "$BASE_BRANCH"
env:
GH_TOKEN: ${{ secrets.GH_TOKEN || github.token }}

9
.github/workflows/test-code-quality.yml vendored
View file

@ -74,6 +74,15 @@ jobs:
- name: check_workflow_startup_safety
run: uv run --no-sync python ./tests/code_coverage_tests/check_workflow_startup_safety.py
- name: check_workflow_job_name_collisions
run: uv run --no-sync python ./tests/code_coverage_tests/check_workflow_job_name_collisions.py
- name: test_workflow_job_name_collisions
run: uv run --no-sync pytest -q --noconftest -p no:cacheprovider -c /dev/null tests/code_coverage_tests/test_workflow_job_name_collisions.py
- name: test_e2e_changed_gate
run: uv run --no-sync pytest -q --noconftest -p no:cacheprovider -c /dev/null tests/code_coverage_tests/test_e2e_changed_gate.py
- name: router_code_coverage
run: uv run --no-sync python ./tests/code_coverage_tests/router_code_coverage.py

240
.github/workflows/test-e2e-changed.yml vendored Normal file
View file

@ -0,0 +1,240 @@
name: e2e-changed-tests
on:
pull_request:
concurrency:
group: e2e-changed-${{ github.event.pull_request.number }}
cancel-in-progress: true
permissions: {}
jobs:
detect:
name: Detect changed e2e tests
runs-on: ubuntu-latest
timeout-minutes: 5
permissions:
contents: read
pull-requests: read
outputs:
tests: ${{ steps.changed.outputs.tests }}
any: ${{ steps.changed.outputs.any }}
steps:
- name: Checkout the selector and the canary suite
uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
sparse-checkout: |
.github/e2e-stack
tests/e2e/access_control
persist-credentials: false
ref: ${{ github.sha }}
- name: List the e2e test files this PR added or modified
id: changed
env:
GH_TOKEN: ${{ github.token }}
REPO: ${{ github.repository }}
PR_NUMBER: ${{ github.event.pull_request.number }}
HEAD_SHA: ${{ github.event.pull_request.head.sha }}
run: |
gh api "repos/${REPO}/pulls/${PR_NUMBER}" \
--jq 'select(.head.sha == env.HEAD_SHA and .changed_files < 3000) | .head.sha' \
| grep -Fxq "${HEAD_SHA}"
files="$(gh api "repos/${REPO}/pulls/${PR_NUMBER}/files" --paginate \
--jq '.[] | select(.status != "removed") | .filename')"
gh api "repos/${REPO}/pulls/${PR_NUMBER}" --jq '.head.sha' | grep -Fxq "${HEAD_SHA}"
tests="$(printf '%s\n' "${files}" \
| python3 .github/e2e-stack/select_tests.py tests/e2e/access_control/test_*.py)"
echo "tests=${tests}" >> "${GITHUB_OUTPUT}"
if [ -n "${tests}" ]; then
echo "any=true" >> "${GITHUB_OUTPUT}"
echo "selected e2e tests: ${tests}"
else
echo "any=false" >> "${GITHUB_OUTPUT}"
echo "no changed e2e test files supported by this stack; nothing to run"
fi
run:
name: Run changed e2e tests against the stage-mirror stack
needs: detect
if: needs.detect.outputs.any == 'true' && github.event.pull_request.head.repo.full_name == github.repository
runs-on: ubuntu-latest
timeout-minutes: 90
environment: e2e-changed
permissions:
contents: read
id-token: write
services:
postgres:
image: postgres:16.6
env:
POSTGRES_USER: litellm
POSTGRES_PASSWORD: dbpassword9090
POSTGRES_DB: litellm
ports:
- 5432:5432
options: >-
--health-cmd "pg_isready -U litellm"
--health-interval 5s
--health-timeout 5s
--health-retries 10
jaeger:
image: jaegertracing/jaeger:2.10.0
ports:
- 4318:4318
- 16686:16686
steps:
- name: Validate configuration
env:
ROLE: ${{ vars.E2E_AWS_ROLE_TO_ASSUME }}
run: test -n "${ROLE}" || { echo "::error::Set repo variable E2E_AWS_ROLE_TO_ASSUME to an OIDC role with read access to the e2e secrets"; exit 1; }
- name: Checkout
uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
persist-credentials: false
ref: ${{ github.sha }}
- name: Set up Python
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
with:
python-version: "3.13"
- name: Set up uv
uses: ./.github/actions/setup-uv-with-retries
with:
version: "0.10.9"
- name: Cache the Rust build
uses: ./.github/actions/cache-cargo-build
- name: Install dependencies
run: |
.github/scripts/uv_sync_with_retries.sh --frozen \
--extra proxy --extra proxy-runtime --extra extra_proxy \
--extra semantic-router --extra bedrock-realtime \
--group ci --group proxy-dev --group e2e-dev
uv pip install "pipecat-ai[openai]==1.4.0"
- name: Cache Prisma binaries
uses: ./.github/actions/cache-prisma-binaries
- name: Generate Prisma client
run: uv run --no-sync prisma generate --schema litellm/proxy/schema.prisma
- name: Install Playwright chromium
run: uv run --no-sync playwright install --with-deps chromium
- name: Configure AWS credentials
id: aws
uses: aws-actions/configure-aws-credentials@e7f100cf4c008499ea8adda475de1042d6975c7b # v6.2.0
with:
role-to-assume: ${{ vars.E2E_AWS_ROLE_TO_ASSUME }}
aws-region: us-east-1
role-session-name: litellm-e2e-changed-${{ github.run_id }}
role-duration-seconds: 900
output-env-credentials: false
output-credentials: true
- name: Fetch provider credentials from AWS Secrets Manager
env:
AWS_ACCESS_KEY_ID: ${{ steps.aws.outputs.aws-access-key-id }}
AWS_SECRET_ACCESS_KEY: ${{ steps.aws.outputs.aws-secret-access-key }}
AWS_SESSION_TOKEN: ${{ steps.aws.outputs.aws-session-token }}
AWS_DEFAULT_REGION: us-east-1
run: |
umask 077
aws secretsmanager get-secret-value --secret-id litellm-e2e-changed-provider-keys \
--query SecretString --output text \
| uv run --no-sync python .github/e2e-stack/secrets_to_env.py tests/e2e/.env
aws secretsmanager get-secret-value --secret-id litellm-e2e-changed-license \
--query SecretString --output text \
| jq -R -s '{"LITELLM_LICENSE": .}' \
| uv run --no-sync python .github/e2e-stack/secrets_to_env.py tests/e2e/.env
- name: Boot the stage-mirror stack
id: boot
run: |
umask 077
if ! bash .github/e2e-stack/up.sh > "${RUNNER_TEMP}/e2e-boot.log" 2>&1; then
echo "::error::stage-mirror stack failed to boot; raw logs are not published"
exit 1
fi
- name: Export stack environment
run: |
master_key="$(grep '^LITELLM_MASTER_KEY=' "${RUNNER_TEMP}/litellm-e2e-stack/stack.env" | cut -d= -f2-)"
echo "::add-mask::${master_key}"
cat "${RUNNER_TEMP}/litellm-e2e-stack/stack.env" >> "${GITHUB_ENV}"
- name: Run the selected tests three times
env:
TESTS: ${{ needs.detect.outputs.tests }}
E2E_FIXTURE_MODE: live
run: |
umask 077
read -r -a test_files <<< "${TESTS}"
for pass in 1 2 3; do
report="${RUNNER_TEMP}/e2e-pass-${pass}.xml"
log="${RUNNER_TEMP}/e2e-pass-${pass}.log"
echo "::group::pass ${pass} of 3"
set +e
uv run --no-sync pytest "${test_files[@]}" --rootdir=. -v -p no:cacheprovider \
-o junit_family=xunit1 --junitxml="${report}" > "${log}" 2>&1
status=$?
uv run --no-sync python .github/e2e-stack/assert_tests_ran.py "${report}" "${test_files[@]}"
verified=$?
set -e
grep -E '^=+ .* in [0-9.]+s( \([0-9:]+\))? =+$' "${log}" | tail -n 1
echo "::endgroup::"
if [ "${status}" = "5" ]; then
echo "::error::the selected files collected no runnable tests, so nothing was verified"
exit 1
fi
if [ "${status}" != "0" ]; then
echo "::error::pass ${pass} of 3 failed with exit code ${status}"
exit "${status}"
fi
if [ "${verified}" != "0" ]; then
echo "::error::pass ${pass} of 3 did not verify every selected file"
exit 1
fi
echo "pass ${pass} of 3 passed"
done
- name: Stop the stack
if: always() && steps.boot.outcome != 'skipped'
run: bash .github/e2e-stack/down.sh
- name: Remove credentials and raw output
if: always()
run: |
rm -f tests/e2e/.env "${RUNNER_TEMP}/e2e-boot.log" "${RUNNER_TEMP}"/e2e-pass-*.log "${RUNNER_TEMP}"/e2e-pass-*.xml
rm -rf "${RUNNER_TEMP}/litellm-e2e-stack"
gate:
name: e2e-changed-tests
needs: [detect, run]
if: always()
runs-on: ubuntu-latest
timeout-minutes: 5
steps:
- name: Require three successful passes when tests changed
env:
DETECT_RESULT: ${{ needs.detect.result }}
ANY_TESTS: ${{ needs.detect.outputs.any }}
RUN_RESULT: ${{ needs.run.result }}
run: |
if [ "${DETECT_RESULT}" != "success" ]; then
echo "::error::changed-test detection did not succeed"
exit 1
fi
if [ "${ANY_TESTS}" = "false" ]; then
echo "no changed e2e test files supported by this stack; nothing to run"
exit 0
fi
if [ "${ANY_TESTS}" != "true" ] || [ "${RUN_RESULT}" != "success" ]; then
echo "::error::selected e2e tests require an approved, successful run; fork PRs must run from a reviewed same-repository branch"
exit 1
fi

5
.github/workflows/test-litellm-ui-unit.yml vendored
View file

@ -12,6 +12,7 @@ on:
- "litellm_**"
push:
branches:
- main
- litellm_internal_staging
concurrency:
@ -65,7 +66,7 @@ jobs:
BASE_SHA: ${{ github.event.pull_request.base.sha }}
HEAD_SHA: ${{ github.event.pull_request.head.sha }}
run: |
full_suite() { npm run test -- --run --pool forks --poolOptions.forks.maxForks=14; }
full_suite() { npm run test -- --run --pool forks --maxWorkers=14; }
if [ -z "$BASE_SHA" ]; then
echo "Push to $GITHUB_REF_NAME: running the full suite"
@ -94,4 +95,4 @@ jobs:
echo "Pull request: running tests related to ${#changed_files[@]} changed UI files"
npm run test -- related "${changed_files[@]}" --run --passWithNoTests \
--pool forks --poolOptions.forks.maxForks=14
--pool forks --maxWorkers=14

37
.github/workflows/test-model-map.yml vendored
View file

@ -1,37 +0,0 @@
name: Validate model_prices_and_context_window.json
on:
pull_request:
branches:
- main
- litellm_internal_staging
- litellm_oss_staging
- "litellm_**"
permissions:
contents: read
concurrency:
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}
cancel-in-progress: ${{ github.event_name == 'pull_request' }}
jobs:
validate-model-prices-json:
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
persist-credentials: false
- name: Validate model_prices_and_context_window.json
run: |
jq empty model_prices_and_context_window.json
- name: Set up uv
uses: ./.github/actions/setup-uv-with-retries
with:
version: "0.10.9"
- name: Check model_prices_and_context_window.schema.json is in sync
run: |
uv run --frozen python ci_cd/generate_model_prices_schema.py --check

114
.github/workflows/test-rust.yml vendored
View file

@ -7,6 +7,7 @@ on:
- ".cargo/**"
- "pyproject.toml"
- "rust-toolchain.toml"
- ".github/actions/setup-uv-with-retries/**"
- ".github/scripts/smoke_test_native_wheel.py"
- ".github/scripts/verify_linux_native_wheel.py"
- "tests/test_litellm/rust_bridge/native_route_wheel_test.py"
@ -22,6 +23,7 @@ on:
- ".cargo/**"
- "pyproject.toml"
- "rust-toolchain.toml"
- ".github/actions/setup-uv-with-retries/**"
- ".github/scripts/smoke_test_native_wheel.py"
- ".github/scripts/verify_linux_native_wheel.py"
- "tests/test_litellm/rust_bridge/native_route_wheel_test.py"
@ -34,102 +36,92 @@ concurrency:
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
cancel-in-progress: true
env:
CARGO_TERM_COLOR: always
jobs:
rust-checks:
name: rustfmt, clippy, test
rust-lint:
runs-on: ubuntu-latest
timeout-minutes: 10
defaults:
run:
working-directory: litellm-rust
env:
CARGO_TERM_COLOR: always
steps:
- name: Checkout repository
uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
persist-credentials: false
- name: Set up Rust
run: rustup toolchain install
- run: rustup toolchain install --no-self-update
- name: Cache Cargo registry and target
uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0
- run: cargo fmt --check
- uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0
with:
path: |
~/.cargo/registry
~/.cargo/git
litellm-rust/target
key: ${{ runner.os }}-cargo-${{ hashFiles('rust-toolchain.toml', 'litellm-rust/Cargo.lock') }}
key: ${{ runner.os }}-cargo-${{ github.job }}-${{ hashFiles('rust-toolchain.toml', '.cargo/**', 'litellm-rust/Cargo.lock') }}
restore-keys: |
${{ runner.os }}-cargo-
${{ runner.os }}-cargo-${{ github.job }}-
- name: Check Rust formatting
run: cargo fmt --check
- run: cargo clippy --workspace --all-targets --locked -- -D warnings
- name: Run Clippy
run: cargo clippy --workspace --all-targets --locked -- -D warnings
- run: cargo clippy -p litellm-core --all-targets --features bedrock-auth --locked -- -D warnings
- name: Run Clippy with Bedrock auth
run: cargo clippy -p litellm-core --all-targets --features bedrock-auth --locked -- -D warnings
- run: cargo clippy -p litellm-ai-gateway --all-targets --all-features --locked -- -D warnings
- name: Run Clippy with all gateway features
run: cargo clippy -p litellm-ai-gateway --all-targets --all-features --locked -- -D warnings
- name: Run Rust tests
run: cargo test --workspace --locked
- name: Run core tests with Bedrock auth
run: cargo test -p litellm-core --features bedrock-auth --locked
# Not --all-features: python-config links libpython, which this job does not install.
- name: Run gateway tests with the server feature
run: cargo test -p litellm-ai-gateway --features server --locked
release-wheel:
name: release wheel
rust-test:
runs-on: ubuntu-latest
timeout-minutes: 20
permissions:
contents: read
env:
CARGO_TERM_COLOR: always
timeout-minutes: 30
steps:
- name: Checkout repository
uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
persist-credentials: false
- name: Set up Python
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
- uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
with:
python-version: "3.12"
- name: Set up uv
uses: ./.github/actions/setup-uv-with-retries
- uses: ./.github/actions/setup-uv-with-retries
with:
version: "0.10.9"
- name: Set up Rust
run: rustup toolchain install
- run: rustup toolchain install --no-self-update
- name: Build release wheel
run: uv build --wheel --out-dir dist
- uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0
with:
path: |
~/.cargo/registry
~/.cargo/git
litellm-rust/target
key: ${{ runner.os }}-cargo-${{ github.job }}-${{ hashFiles('rust-toolchain.toml', '.cargo/**', 'litellm-rust/Cargo.lock') }}
restore-keys: |
${{ runner.os }}-cargo-${{ github.job }}-
- name: Build panic contract wheel
run: >-
- run: cargo test --workspace --locked
working-directory: litellm-rust
- run: cargo test -p litellm-core --features bedrock-auth --locked
working-directory: litellm-rust
- run: cargo test -p litellm-ai-gateway --features server --locked
working-directory: litellm-rust
- run: uv build --wheel --out-dir dist
- run: python .github/scripts/verify_linux_native_wheel.py dist/*.whl
env:
RELEASE_WHEEL_COMMIT_SHA: ${{ github.event.pull_request.head.sha || github.sha }}
- run: python tests/test_litellm/rust_bridge/native_route_wheel_test.py dist/*.whl
- name: Run pytest tests/test_litellm_rust with the compiled extension
run: make test-rust-extension
- run: >-
uv build --wheel --out-dir panic-dist
--config-setting "maturin.build-args=--features panic-test,extension-module"
- name: Smoke-test native panic unwinding
run: python .github/scripts/smoke_test_native_wheel.py panic-dist/*.whl
- name: Verify stripped native extension
env:
RELEASE_WHEEL_COMMIT_SHA: ${{ github.event.pull_request.head.sha || github.sha }}
run: python .github/scripts/verify_linux_native_wheel.py dist/*.whl
- name: Test native route wheel
run: python tests/test_litellm/rust_bridge/native_route_wheel_test.py dist/*.whl
- run: python .github/scripts/smoke_test_native_wheel.py panic-dist/*.whl

31
.github/workflows/test-terraform-modules.yml vendored
View file

@ -4,6 +4,7 @@ on:
push:
paths:
- "terraform/litellm/aws/**"
- "terraform/litellm/gcp/**"
- ".github/workflows/test-terraform-modules.yml"
pull_request:
branches:
@ -13,6 +14,7 @@ on:
- "litellm_**"
paths:
- "terraform/litellm/aws/**"
- "terraform/litellm/gcp/**"
- ".github/workflows/test-terraform-modules.yml"
permissions:
@ -52,3 +54,32 @@ jobs:
# Plan-only, mock_provider-backed: no AWS credentials, no API calls.
- name: test
run: terraform test
gcp-module:
name: fmt, validate, test (gcp)
runs-on: ubuntu-latest
timeout-minutes: 15
defaults:
run:
working-directory: terraform/litellm/gcp
steps:
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
persist-credentials: false
- uses: hashicorp/setup-terraform@b9cd54a3c349d3f38e8881555d616ced269862dd # v3.1.2
with:
terraform_version: 1.13.3
terraform_wrapper: false
- name: fmt
run: terraform fmt -recursive -check -diff
- name: init
run: terraform init -backend=false -input=false
- name: validate
run: terraform validate
- name: test
run: terraform test

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

@ -116,6 +116,7 @@ jobs:
tests/test_litellm/rerank_api
tests/test_litellm/rust_bridge
tests/test_litellm/sandbox
tests/test_litellm/skills
tests/test_litellm/test_router
tests/test_litellm/vector_stores
tests/test_litellm/videos

16
ARCHITECTURE.md
View file

@ -149,7 +149,7 @@ graph TD
| `parallel_request_limiter` | `proxy/hooks/parallel_request_limiter_v3.py` | Rate limiting per key/user |
| `cache_control_check` | `proxy/hooks/cache_control_check.py` | Cache validation |
| `responses_id_security` | `proxy/hooks/responses_id_security.py` | Response ID validation |
| `litellm_skills` | `proxy/hooks/skills_injection.py` | Skills injection |
| `litellm_skills` | `proxy/hooks/litellm_skills/main.py` | Skills injection |
To add a new proxy hook, implement `CustomLogger` and register in `PROXY_HOOKS`.
@ -220,20 +220,20 @@ graph LR
| Job | Interval | Purpose | Key Files |
|-----|----------|---------|-----------|
| `update_spend` | 60s | Batch write spend logs to PostgreSQL | `proxy/db/db_spend_update_writer.py` |
| `reset_budget` | 10-12min | Reset budgets for keys/users/teams | `proxy/management_helpers/budget_reset_job.py` |
| `reset_budget` | 10-12min | Reset budgets for keys/users/teams | `proxy/common_utils/reset_budget_job.py` |
| `add_deployment` | 10s | Sync new model deployments from DB | `proxy/proxy_server.py` (`ProxyConfig`) |
| `cleanup_old_spend_logs` | cron/interval | Delete old spend logs | `proxy/management_helpers/spend_log_cleanup.py` |
| `check_batch_cost` | 30min | Calculate costs for batch jobs | `proxy/management_helpers/check_batch_cost_job.py` |
| `check_responses_cost` | 30min | Calculate costs for responses API | `proxy/management_helpers/check_responses_cost_job.py` |
| `process_rotations` | 1hr | Auto-rotate API keys | `proxy/management_helpers/key_rotation_manager.py` |
| `cleanup_old_spend_logs` | cron/interval | Delete old spend logs | `proxy/db/db_transaction_queue/spend_log_cleanup.py` |
| `check_batch_cost` | 30min | Calculate costs for batch jobs | `enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py` |
| `check_responses_cost` | 30min | Calculate costs for responses API | `enterprise/litellm_enterprise/proxy/common_utils/check_responses_cost.py` |
| `process_rotations` | 1hr | Auto-rotate API keys | `proxy/common_utils/key_rotation_manager.py` |
| `_run_background_health_check` | continuous | Health check model deployments | `proxy/proxy_server.py` |
| `send_weekly_spend_report` | weekly | Slack spend alerts | `proxy/utils.py` (`SlackAlerting`) |
| `send_monthly_spend_report` | monthly | Slack spend alerts | `proxy/utils.py` (`SlackAlerting`) |
**Cost Attribution Flow:**
1. LLM response returns to `utils.py` wrapper after `litellm.acompletion()` completes
2. `update_response_metadata()` (`llm_response_utils/response_metadata.py`) is called
3. `logging_obj._response_cost_calculator()` (`litellm_logging.py`) calculates cost via `litellm.completion_cost()` (`cost_calculator.py`)
2. `update_response_metadata()` (`litellm_core_utils/llm_response_utils/response_metadata.py`) is called
3. `logging_obj._response_cost_calculator()` (`litellm_core_utils/litellm_logging.py`) calculates cost via `litellm.completion_cost()` (`cost_calculator.py`)
4. Cost is stored in `response._hidden_params["response_cost"]`
5. `proxy/common_request_processing.py` extracts cost from `hidden_params` and adds to response headers (`x-litellm-response-cost`)
6. `logging_obj.async_success_handler()` triggers callbacks including `_ProxyDBLogger.async_log_success_event()`

8
CLAUDE.md
View file

@ -29,7 +29,7 @@ Never test structure of code only function of it
End-to-end tests belong in `tests/e2e/` and must follow the harness conventions documented in that directory's `CLAUDE.md`
When creating PRs, don't set base to `main`. `litellm_internal_staging` is the default base branch and serves that purpose for both internal and external / OSS contributions
When creating PRs, target the repository's current default branch for both internal and external / OSS contributions. Check it with `python3 scripts/default_branch.py --branch` instead of assuming a branch name or relying on cached `origin/HEAD`
When writing a PR body, treat the comments and imperative instructions inside .github/pull_request_template.md as rules to follow, not just layout. Agent harnesses may strip HTML comments from copies of that file injected into context, so read .github/pull_request_template.md from disk before writing a PR body to make sure you see every comment rule
@ -37,7 +37,7 @@ Same applies for filing bug reports and feature requests, with .github/ISSUE_TEM
If you're resolving a linear ticket, in the "## Linear ticket" section of the PR, say "Resolves LIT-1234", replacing "LIT-1234" with the actual ticket id that you're resolving. If you don't have the ticket id, don't make one up or search for it. Just leave the section blank
Never use `pytest` commands or the like as "Screenshots / Proof of Fix". We prefer curl'ing a live proxy instance running on localhost:4000 (I like to run it with `python litellm/proxy/proxy_cli.py --config litellm/proxy/dev_config.yaml --detailed_debug --reload --use_v2_migration_resolver 2>&1 | tee litellm.log`; the Admin UI dev server is `npm run dev` in `ui/litellm-dashboard`, served on port 3000) and showing both the command run and the output. Also, it should hit real LLM provider APIs, not mocks, and cost real $$$ because that is the most realistic test. The proof of fix should be exactly what the end user / customer would see / do. The run logs in PR #27703 is a prime example of how to do it (not a huge fan of using a python test script that future me and the team will have no visibility into; I prefer just curl commands or a short list of bash commands (e.g., using `for`)). If it's a UI thing, just tell me which URLs to go to (e.g., http://localhost:4000/ui/?page=logs), where to click, what fields to fill out, etc. along with the other commands to run in an ordered list, and I'll do it myself and post the screenshots after you make the PR
Never use `pytest` commands or the like as "Screenshots / Proof of Fix". We prefer curl'ing a live proxy instance running on localhost:4000 (I like to run it with `python litellm/proxy/proxy_cli.py --config litellm/proxy/dev_config.yaml --detailed_debug --reload --use_v2_migration_resolver 2>&1 | tee litellm.log`; the Admin UI dev server is `npm run dev` in `ui/litellm-dashboard`, served on port 3000) and showing both the command run and the output. Also, it should hit real LLM provider APIs, not mocks, and cost real $$$ because that is the most realistic test. The proof of fix should be exactly what the end user / customer would see / do. The run logs in PR #27703 is a prime example of how to do it (not a huge fan of using a python test script that future me and the team will have no visibility into; I prefer just curl commands or a short list of bash commands (e.g., using `for`)). If it's a UI thing, or the main use case runs through a headful agentic coding tool like Claude Code or Codex, drive that surface yourself and embed your own before and after screenshots of it in the PR (the Admin UI page, or what the coding tool shows), next to an ordered list of the URLs to go to (e.g., http://localhost:4000/ui/?page=logs), where to click, and what fields to fill out so a reviewer can reproduce it
If you ever write any human-facing text (pull requests, issues, commit messages, discussion posts, github comments, release notes, docs, etc.), always follow these guidelines to sound less AI-y:
- don't use emojis
@ -52,7 +52,7 @@ Don't hesitate to use values in .env to get needed API keys and other secrets, a
Python max line length is 120, not 88
When you fix violations gated by `ruff-strict-budget.json`, `type-discipline-budget.json`, or `basedpyright-code-budget.json`, run `make lint-budget-update` and commit the lowered limits so the ceilings ratchet down instead of leaving stale headroom. It measures the working tree, so it must contain exactly the fixes you're committing
Never edit or commit `ruff-strict-budget.json`, `type-discipline-budget.json`, `basedpyright-code-budget.json`, or `test-quality-budget.json` on a PR branch, and don't run `make lint-budget-update` there. A scheduled Devin automation lowers the limits on the default branch in its own PR by exactly what landed since the last ratchet, so concurrent PRs don't fight over the same `"limit"` lines. Keep the hosted automation's target in sync when the repository default changes. If your branch already carries a budget edit, drop it before opening the PR
`make check` (f.k.a. `make pre-commit`, which still works identically as an alias) saves its complete output to a log file in .git (overwriting previous logs) and prints that path as its first and last output lines. To inspect a run, read or grep that log instead of re-running the multi-minute checks just to see a different slice
@ -70,7 +70,7 @@ When referencing or running models (coding, QA'ing, writing docs, writing tests,
Always pull before starting any work. The checkout or worktree may be sitting on a stale branch
If you're an internal contributor, when creating a new PR, the typical flow is to branch off litellm_internal_staging and create a branch prefixed with litellm_. Do not create a branch prefixed with claude/ and generally do not have / in your branch names
If you're an internal contributor, when creating a new PR, the typical flow is to branch off the repository's current default branch and create a branch prefixed with litellm_. Do not create a branch prefixed with claude/ and generally do not have / in your branch names
Do not add `Co-Authored-By: Claude` or any Claude attribution to commit messages. Never use a `claude/` prefix or put a `/` in a branch name. Do not add "Generated with Claude Code" (or any similar attribution) to PR descriptions or comments. Do not create a new PR/branch off the existing PR to fix/add something that is related and could've just been committed directly to the existing PR's branch

4
CONTRIBUTING.md
View file

@ -315,10 +315,12 @@ Ensure the UI builds successfully before submitting your PR:
npm run build
```
Local lint and budget checks follow origin's current default branch. They refresh it from the remote instead of trusting cached `origin/HEAD`. For an intentional comparison against another branch or commit, use `make check BASE_REF=<ref>` or the standalone gate's `--base <ref>` option. An explicit ref can also be used offline once it has been fetched locally. Without an override, unavailable remote metadata stops the check
## Submitting Your PR
1. **Push your branch**: `git push origin your-feature-branch`
2. **Create a PR**: Go to GitHub and open a pull request against [`litellm_internal_staging`](https://github.com/BerriAI/litellm/tree/litellm_internal_staging), which is the default base branch. Do not target `main`.
2. **Create a PR**: Go to GitHub and open a pull request against the repository's current default branch. Run `python3 scripts/default_branch.py --branch` to check its name
3. **Fill out the PR template**: Provide clear description of changes
4. **Wait for review**: Maintainers will review and provide feedback
5. **Address feedback**: Make requested changes and push updates

73
Makefile
View file

@ -4,6 +4,7 @@
.PHONY: help test test-unit test-unit-llms test-unit-proxy-guardrails test-unit-proxy-core test-unit-proxy-misc \
test-unit-integrations test-unit-core-utils test-unit-other test-unit-root \
test-proxy-unit-a test-proxy-unit-b test-integration test-unit-helm \
test-rust-extension \
info lint lint-inner lint-dev lint-checks format \
lint-basedpyright lint-e2e-basedpyright lint-basedpyright-budget-update lint-type-discipline lint-type-discipline-budget-update \
lint-ruff-budget lint-ruff-budget-update lint-budget-update lint-gate \
@ -34,7 +35,7 @@ help:
@echo " make lint-basedpyright-budget-update - Ratchet basedpyright limits down by what this branch fixed"
@echo " make lint-format - Check ruff format formatting (matches CI)"
@echo " make lint-ruff-budget - Gate the codebase total of each strict ruff rule against its limit"
@echo " make lint-gate - Strict ruff gate in CI-parity mode (fetches staging, simulates the merge)"
@echo " make lint-gate - Strict ruff gate in CI-parity mode (fetches the default branch, simulates the merge)"
@echo " make lint-ruff-budget-update - Ratchet ruff-strict-budget.json limits down by what this branch fixed"
@echo " make lint-test-quality - Gate the test suite against test-quality-budget.json"
@echo " make lint-budget-update - Ratchet all budgets down (ruff + type-discipline + test quality + basedpyright)"
@ -54,12 +55,16 @@ help:
@echo " make test-proxy-unit-b - Run proxy_unit_tests (p-z, ~28 files)"
@echo " make test-integration - Run integration tests"
@echo " make test-unit-helm - Run helm unit tests"
@echo " make test-rust-extension - Build the Rust extension and run its public Python tests"
@echo ""
@echo "Heavy targets (check, lint) queue for LITELLM_GATE_SLOTS machine-wide"
@echo "slots (default 2; 0 disables) so parallel sessions don't thrash one machine."
UV := uv
UV_RUN := $(UV) run --no-sync
BASE_REF ?=
export BASE_REF
RESOLVE_BASE = python3 scripts/default_branch.py --base "$(BASE_REF)"
# Machine-wide slot queue for the heavy targets below; python3 + stdlib only, so
# it runs before any venv exists. See scripts/gate_slot_lock.py.
@ -67,7 +72,7 @@ GATE_SLOT_LOCK := python3 scripts/gate_slot_lock.py
LINT_DEP_INSTALL ?= install-dev
LINT_E2E_DEP_INSTALL ?= lint-install
LINT_DEP_BASE ?= lint-fetch-base
LINT_DEP_BASE ?=
LINT_JOBS := $(shell sysctl -n hw.ncpu 2>/dev/null || nproc 2>/dev/null || echo 4)
LINT_OUTPUT_SYNC := $(if $(filter output-sync,$(.FEATURES)),--output-sync=target,)
@ -130,10 +135,8 @@ format: install-dev
format-check: install-dev
cd litellm && $(UV_RUN) ruff format --check --exclude '/enterprise/' . && cd ..
# Single fetch of the PR base so the delta-based gates below share one network round
# trip instead of each re-fetching when chained from `lint`.
lint-fetch-base:
git fetch origin litellm_internal_staging
@$(RESOLVE_BASE)
# Mirror test-linting.yml's lint job environment: the proxy-dev group plus a generated
# Prisma client, so `basedpyright tests/e2e` resolves the same modules CI does. The
@ -150,7 +153,9 @@ lint-install:
# recursively, so 'litellm/*.py' covers nested modules and the top-level files that
# CI's 'litellm/**/*.py' skips, which makes this target a superset of the CI step.
lint-format-check-changed: $(LINT_DEP_INSTALL) $(LINT_DEP_BASE)
@files=$$(git diff --name-only --diff-filter=ACMR origin/litellm_internal_staging...HEAD -- 'litellm/*.py' | grep -v '^litellm/enterprise/' || true); \
@base_ref=$$($(RESOLVE_BASE)) && \
changed=$$(git diff --name-only --diff-filter=ACMR "$$base_ref...HEAD" -- 'litellm/*.py') && \
files=$$(printf '%s\n' "$$changed" | grep -v '^litellm/enterprise/' || true) || exit $$?; \
if [ -z "$$files" ]; then \
echo "No changed litellm Python files to format-check."; \
else \
@ -167,7 +172,9 @@ lint-ruff: $(LINT_DEP_INSTALL)
# https://github.com/astral-sh/ruff/discussions/10977
# https://github.com/astral-sh/ruff/discussions/4049
lint-format-changed: install-dev
@git diff origin/main --unified=0 --no-color -- '*.py' | \
@base_ref=$$($(RESOLVE_BASE)) && \
diff=$$(git diff "$$base_ref" --unified=0 --no-color -- '*.py') && \
printf '%s\n' "$$diff" | \
perl -ne '\
if (/^diff --git a\/(.*) b\//) { $$file = $$1; } \
if (/^@@ .* \+(\d+)(?:,(\d+))? @@/) { \
@ -182,20 +189,22 @@ lint-format-changed: install-dev
done
lint-ruff-dev: install-dev
@tmpfile=$$(mktemp /tmp/ruff-dev.XXXXXX) && \
@base_ref=$$($(RESOLVE_BASE)) || exit $$?; \
tmpfile=$$(mktemp /tmp/ruff-dev.XXXXXX) && \
cd litellm && \
($(UV_RUN) ruff check . --output-format=pylint || true) > "$$tmpfile" && \
$(UV_RUN) diff-quality --violations=pylint "$$tmpfile" --compare-branch=origin/main && \
$(UV_RUN) diff-quality --violations=pylint "$$tmpfile" --compare-branch="$$base_ref" && \
cd .. ; \
rm -f "$$tmpfile"
lint-ruff-FULL-dev: install-dev
@files=$$(git diff --name-only origin/main -- '*.py'); \
@base_ref=$$($(RESOLVE_BASE)) && \
files=$$(git diff --name-only "$$base_ref" -- '*.py') || exit $$?; \
if [ -n "$$files" ]; then echo "$$files" | xargs $(UV_RUN) ruff check; \
else echo "No changed .py files to check."; fi
lint-basedpyright: $(LINT_DEP_INSTALL) $(LINT_DEP_BASE)
$(UV_RUN) python scripts/type_check_gate.py --base origin/litellm_internal_staging
$(UV_RUN) python scripts/type_check_gate.py --base "$(BASE_REF)"
lint-e2e-basedpyright: $(LINT_E2E_DEP_INSTALL)
$(UV_RUN) basedpyright tests/e2e
@ -203,37 +212,37 @@ lint-e2e-basedpyright: $(LINT_E2E_DEP_INSTALL)
# Type-discipline budget (mutable collections / casts / type guards / kwargs /
# unexplained suppressions), the test-linting.yml step `make lint` used to omit.
lint-type-discipline: $(LINT_DEP_INSTALL) $(LINT_DEP_BASE)
$(UV_RUN) python scripts/type_discipline_gate.py --base origin/litellm_internal_staging
$(UV_RUN) python scripts/type_discipline_gate.py --base "$(BASE_REF)"
# Test-quality budget (zero-assert / mock-echo tests, sys.path.insert, raw env writes,
# litellm module-global mutation, credential-gated skips, conftest snapshot
# inventory), counted across tests/ the same delta-vs-base way.
lint-test-quality: $(LINT_DEP_INSTALL) $(LINT_DEP_BASE)
$(UV_RUN) python scripts/test_quality_gate.py --base origin/litellm_internal_staging
$(UV_RUN) python scripts/test_quality_gate.py --base "$(BASE_REF)"
# --update lowers each limit by what this branch fixed since its branch point, so
# it needs the base ref fetched to resolve the merge-base.
lint-basedpyright-budget-update: install-dev lint-fetch-base
$(UV_RUN) python scripts/type_check_gate.py --update
lint-basedpyright-budget-update: install-dev
$(UV_RUN) python scripts/type_check_gate.py --update --base "$(BASE_REF)"
lint-format: format-check
lint-ruff-budget: install-dev
$(UV_RUN) python scripts/ruff_strict_gate.py
$(UV_RUN) python scripts/ruff_strict_gate.py --base "$(BASE_REF)"
# Strict gate, invoked the same way CI does in test-linting.yml so a local pass
# means the CI check will pass too.
lint-gate: $(LINT_DEP_INSTALL) $(LINT_DEP_BASE)
$(UV_RUN) python scripts/ruff_strict_gate.py --base origin/litellm_internal_staging
$(UV_RUN) python scripts/ruff_strict_gate.py --base "$(BASE_REF)"
lint-ruff-budget-update: install-dev lint-fetch-base
$(UV_RUN) python scripts/ruff_strict_gate.py --update
lint-ruff-budget-update: install-dev
$(UV_RUN) python scripts/ruff_strict_gate.py --update --base "$(BASE_REF)"
lint-type-discipline-budget-update: install-dev lint-fetch-base
$(UV_RUN) python scripts/type_discipline_gate.py --update
lint-type-discipline-budget-update: install-dev
$(UV_RUN) python scripts/type_discipline_gate.py --update --base "$(BASE_REF)"
lint-test-quality-budget-update: install-dev lint-fetch-base
$(UV_RUN) python scripts/test_quality_gate.py --update
lint-test-quality-budget-update: install-dev
$(UV_RUN) python scripts/test_quality_gate.py --update --base "$(BASE_REF)"
# Ratchet all budgets in one shot (ruff strict + type-discipline + test quality + basedpyright)
lint-budget-update: lint-ruff-budget-update lint-type-discipline-budget-update lint-test-quality-budget-update lint-basedpyright-budget-update
@ -249,14 +258,15 @@ check-import-safety: $(LINT_DEP_INSTALL)
# runs the diff-scoped ruff format check, whole-tree ruff check, the strict-rule /
# type-discipline / basedpyright budgets as a delta vs the base, then the circular-import
# and import-safety checks. Steps that compare against the base resolve it the same way CI
# does (merge-base with origin/litellm_internal_staging). Setup (env sync, Prisma client,
# does (merge-base with origin's current default branch). Setup (env sync, Prisma client,
# base fetch) runs once up front; the checks themselves are independent, so a sub-make
# fans them out with -j and the fast ones finish under basedpyright's shadow.
lint:
@$(GATE_SLOT_LOCK) $(MAKE) lint-inner
lint-inner: lint-install lint-fetch-base
$(MAKE) -j $(LINT_JOBS) $(LINT_OUTPUT_SYNC) LINT_DEP_INSTALL= LINT_E2E_DEP_INSTALL= LINT_DEP_BASE= lint-checks
lint-inner: lint-install
@base_ref=$$($(RESOLVE_BASE)) && \
$(MAKE) BASE_REF="$$base_ref" -j $(LINT_JOBS) $(LINT_OUTPUT_SYNC) LINT_DEP_INSTALL= LINT_E2E_DEP_INSTALL= LINT_DEP_BASE= lint-checks
lint-checks: lint-format-check-changed lint-ruff lint-gate lint-type-discipline lint-test-quality lint-basedpyright lint-e2e-basedpyright check-circular-imports check-import-safety
@ -281,6 +291,17 @@ pre-commit:
@$(MAKE) check
# Testing targets
test-rust-extension:
@temporary=$$(mktemp -d) && \
trap 'rm -rf "$$temporary"' EXIT HUP INT TERM && \
$(UV) build --python 3.12 --wheel --out-dir "$$temporary/wheels" && \
set -- "$$temporary"/wheels/*.whl && \
[ "$$#" -eq 1 ] && \
UV_PROJECT_ENVIRONMENT="$$temporary/venv" $(UV) sync --python 3.12 --frozen --no-install-project --all-groups --all-extras && \
$(UV) pip install --python "$$temporary/venv/bin/python" --no-deps "$$1" && \
LITELLM_RUST=1 LITELLM_LOCAL_MODEL_COST_MAP=True \
"$$temporary/venv/bin/python" -I -m pytest --import-mode=importlib -m requires_rust_extension tests/test_litellm_rust
test: install-test-deps
$(UV_RUN) pytest tests/

28
basedpyright-code-budget.json
View file

@ -1,9 +1,9 @@
{
"reportAny": {
"limit": 14074
"limit": 13429
},
"reportArgumentType": {
"limit": 2206
"limit": 2198
},
"reportAssignmentType": {
"limit": 319
@ -24,7 +24,7 @@
"limit": 19
},
"reportExplicitAny": {
"limit": 4124
"limit": 3369
},
"reportFunctionMemberAccess": {
"limit": 7
@ -48,16 +48,16 @@
"limit": 30
},
"reportInvalidTypeVarUse": {
"limit": 2
"limit": 1
},
"reportMatchNotExhaustive": {
"limit": 0
},
"reportMissingParameterType": {
"limit": 5601
"limit": 5570
},
"reportMissingTypeArgument": {
"limit": 15284
"limit": 15281
},
"reportMissingTypeStubs": {
"limit": 40
@ -84,13 +84,13 @@
"limit": 56
},
"reportPrivateUsage": {
"limit": 1808
"limit": 1804
},
"reportRedeclaration": {
"limit": 8
},
"reportReturnType": {
"limit": 181
"limit": 180
},
"reportTypedDictNotRequiredAccess": {
"limit": 22
@ -105,16 +105,16 @@
"limit": 109
},
"reportUnknownMemberType": {
"limit": 38309
"limit": 38269
},
"reportUnknownParameterType": {
"limit": 19621
"limit": 19584
},
"reportUnknownVariableType": {
"limit": 29844
"limit": 29814
},
"reportUnnecessaryCast": {
"limit": 111
"limit": 110
},
"reportUnnecessaryComparison": {
"limit": 687
@ -123,7 +123,7 @@
"limit": 4
},
"reportUnnecessaryIsInstance": {
"limit": 819
"limit": 816
},
"reportUntypedBaseClass": {
"limit": 0
@ -135,7 +135,7 @@
"limit": 21
},
"reportUnusedFunction": {
"limit": 138
"limit": 136
},
"reportUnusedImport": {
"limit": 542

146
ci_cd/cost_map_guard.py Normal file
View file

@ -0,0 +1,146 @@
"""Guard the cost map on pull requests.
Every pull request gets the file checks: the three cost map files parse, the backup copy matches the root file,
and the JSON schema is in sync and validates the map. Pull requests from the cost map sync bot (branches named
litellm_cost_map_sync_*) additionally may only touch those three files and may only add or update models.
"""
from __future__ import annotations
import argparse
import json
import subprocess
import sys
from collections.abc import Sequence
from dataclasses import dataclass
from typing import Final
from generate_model_prices_schema import SPECIAL_ROOT_KEYS, build_schema, render, validation_errors
COST_MAP_PATH: Final = "model_prices_and_context_window.json"
BACKUP_PATH: Final = "litellm/model_prices_and_context_window_backup.json"
SCHEMA_PATH: Final = "model_prices_and_context_window.schema.json"
GUARDED_PATHS: Final = (COST_MAP_PATH, BACKUP_PATH, SCHEMA_PATH)
BOT_BRANCH_PREFIX: Final = "litellm_cost_map_sync_"
CostMap = dict[str, object]
@dataclass(frozen=True, slots=True)
class Snapshot:
cost_map: str
backup: str
schema: str
def _parse_object(text: str, path: str) -> CostMap | str:
try:
parsed: Final = json.loads(text)
except json.JSONDecodeError as error:
return f"{path} is not valid JSON: {error}"
return parsed if isinstance(parsed, dict) else f"{path} must be a JSON object at the root"
def _rendered_schema(cost_map: CostMap) -> str:
try:
return render(build_schema(cost_map))
except SystemExit as error:
return str(error)
def _file_failures(head: Snapshot, head_map: CostMap) -> tuple[str, ...]:
schema_text: Final = _rendered_schema(head_map)
if not schema_text.startswith("{"):
return (schema_text,)
backup_failure: Final = (
()
if head.backup == head.cost_map
else (f"{BACKUP_PATH} differs from {COST_MAP_PATH}; copy the root file over it",)
)
schema_failure: Final = (
()
if head.schema == schema_text
else (
f"{SCHEMA_PATH} is out of sync with {COST_MAP_PATH}; "
"run `python ci_cd/generate_model_prices_schema.py` and commit the result",
)
)
return (
*backup_failure,
*schema_failure,
*(
f"{COST_MAP_PATH} does not validate against its schema: {error}"
for error in validation_errors(head_map, json.loads(schema_text))[:20]
),
)
def _entries(cost_map: CostMap) -> dict[str, dict[str, object]]:
return {key: entry for key, entry in cost_map.items() if isinstance(entry, dict)}
def _bot_failures(base: Snapshot, head_map: CostMap, changed_files: Sequence[str]) -> tuple[str, ...]:
base_map: Final = _parse_object(base.cost_map, COST_MAP_PATH)
if isinstance(base_map, str):
return (f"merge base: {base_map}",)
base_entries: Final = _entries(base_map)
head_entries: Final = _entries(head_map)
removed_fields: Final = tuple(
f"{key}.{field}"
for key, entry in base_entries.items()
if key in head_entries
for field in entry
if field not in head_entries[key]
)
return (
*(
f"bot PRs may only change the cost map files, not {path}"
for path in changed_files
if path not in GUARDED_PATHS
),
*(f"bot PRs may not remove models: {key}" for key in base_map if key not in head_map),
*(f"bot PRs may not remove fields: {ref}" for ref in removed_fields),
*(
f"bot PRs may not change {key}"
for key in sorted(SPECIAL_ROOT_KEYS)
if base_map.get(key) != head_map.get(key)
),
)
def guard_failures(base: Snapshot, head: Snapshot, changed_files: Sequence[str], bot: bool) -> tuple[str, ...]:
head_map: Final = _parse_object(head.cost_map, COST_MAP_PATH)
if isinstance(head_map, str):
return (head_map,)
return (*_file_failures(head, head_map), *(_bot_failures(base, head_map, changed_files) if bot else ()))
def _git(*args: str) -> str:
result: Final = subprocess.run(("git", *args), check=False, capture_output=True, text=True)
return result.stdout if result.returncode == 0 else ""
def snapshot(revision: str) -> Snapshot:
return Snapshot(*(_git("show", f"{revision}:{path}") for path in GUARDED_PATHS))
def main(argv: Sequence[str]) -> int:
parser: Final = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--base", required=True, help="merge base of the pull request")
parser.add_argument("--head", required=True, help="head commit of the pull request")
parser.add_argument("--head-ref", required=True, help="head branch name of the pull request")
args: Final = parser.parse_args(argv)
bot: Final = args.head_ref.startswith(BOT_BRANCH_PREFIX)
changed_files: Final = tuple(_git("diff", "--name-only", args.base, args.head).splitlines())
failures: Final = guard_failures(snapshot(args.base), snapshot(args.head), changed_files, bot)
contract: Final = "bot contract enforced" if bot else "human PR, file checks only"
if failures:
print(f"cost map guard failed ({contract}):")
print("\n".join(f"- {failure}" for failure in failures))
return 1
print(f"cost map guard passed ({contract})")
return 0
if __name__ == "__main__":
sys.exit(main(sys.argv[1:]))

9
ci_cd/generate_model_prices_schema.py
View file

@ -58,6 +58,11 @@ OBJECT_KEYS: dict[str, JsonSchema] = {
}
ARRAY_KEYS: dict[str, JsonSchema] = {
"supported_audio_formats": {
"type": "array",
"description": "Audio container formats the model can return.",
"items": {"type": "string", "enum": ["mp3", "wav"]},
},
"supported_endpoints": {
"type": "array",
"description": "OpenAI-style API routes this model can be called through, e.g. /v1/chat/completions.",
@ -231,6 +236,10 @@ def string_key_schemas(modes: tuple) -> dict[str, JsonSchema]:
},
"comment": STRING,
"audio_transcription_config": STRING,
"vertex_ai_audio_api": {
"type": "string",
"enum": ["lyria_predict", "lyria_interactions"],
},
}

62
ci_cd/run_migration.py
View file

@ -6,12 +6,9 @@ import subprocess
import sys
from datetime import datetime
from pathlib import Path
import testing.postgresql
from typing import Final
DESTRUCTIVE_PATTERN = re.compile(r"\bDROP\s+(COLUMN|TABLE|INDEX)\b", re.IGNORECASE)
DEFAULT_BASE_BRANCH = "litellm_internal_staging"
def _find_destructive_statements(sql: str) -> list:
@ -94,31 +91,57 @@ def _print_stale_branch_refusal(base_branch: str, behind: int) -> None:
print(banner, file=out)
def _check_branch_freshness(root_dir: Path, base_branch: str) -> None:
def _default_base_branch(root_dir: Path) -> str:
try:
result: Final = subprocess.run(
[
sys.executable,
str(Path(__file__).resolve().parents[1] / "scripts" / "default_branch.py"),
"--repo-root",
str(root_dir),
"--branch",
],
check=True,
capture_output=True,
text=True,
timeout=90,
)
except (OSError, subprocess.SubprocessError) as exc:
_print_freshness_failure(
"default branch",
"Could not discover origin's default branch. Pass --base-branch <name> to choose one.",
exc.stderr if isinstance(exc, subprocess.CalledProcessError) else str(exc),
)
sys.exit(3)
return result.stdout.strip()
def _check_branch_freshness(root_dir: Path, base_branch: str | None = None) -> None:
"""Fetch origin/<base_branch> and exit 3 if HEAD is behind it."""
resolved_branch: Final = base_branch or _default_base_branch(root_dir)
cwd = str(root_dir)
try:
subprocess.run(
["git", "fetch", "origin", base_branch],
["git", "fetch", "origin", f"+refs/heads/{resolved_branch}:refs/remotes/origin/{resolved_branch}"],
check=True,
capture_output=True,
text=True,
cwd=cwd,
)
except FileNotFoundError:
_print_freshness_failure(base_branch, "git executable not found on PATH")
_print_freshness_failure(resolved_branch, "git executable not found on PATH")
sys.exit(3)
except subprocess.CalledProcessError as e:
_print_freshness_failure(
base_branch,
f"`git fetch origin {base_branch}` failed",
resolved_branch,
f"`git fetch origin {resolved_branch}` failed",
e.stderr or "",
)
sys.exit(3)
try:
result = subprocess.run(
["git", "rev-list", "--count", f"HEAD..origin/{base_branch}"],
["git", "rev-list", "--count", f"HEAD..origin/{resolved_branch}"],
check=True,
capture_output=True,
text=True,
@ -127,23 +150,23 @@ def _check_branch_freshness(root_dir: Path, base_branch: str) -> None:
behind = int(result.stdout.strip())
except subprocess.CalledProcessError as e:
_print_freshness_failure(
base_branch,
f"`git rev-list HEAD..origin/{base_branch}` failed",
resolved_branch,
f"`git rev-list HEAD..origin/{resolved_branch}` failed",
e.stderr or "",
)
sys.exit(3)
except ValueError:
_print_freshness_failure(
base_branch,
resolved_branch,
"could not parse commit count from `git rev-list`",
)
sys.exit(3)
if behind > 0:
_print_stale_branch_refusal(base_branch, behind)
_print_stale_branch_refusal(resolved_branch, behind)
sys.exit(3)
print(f"Branch freshness OK: up to date with origin/{base_branch}.")
print(f"Branch freshness OK: up to date with origin/{resolved_branch}.")
def _print_destructive_refusal(destructive_lines: list) -> None:
@ -198,7 +221,7 @@ def _print_destructive_refusal(destructive_lines: list) -> None:
def create_migration(
migration_name: str = None,
allow_destructive: bool = False,
base_branch: str = DEFAULT_BASE_BRANCH,
base_branch: str | None = None,
skip_freshness_check: bool = False,
):
"""
@ -211,7 +234,7 @@ def create_migration(
DROP COLUMN, DROP TABLE, or DROP INDEX statements. Without this
flag, the script exits non-zero and prints guidance.
base_branch (str): Branch to check freshness against
(default: "litellm_internal_staging").
(default: origin's current default branch).
skip_freshness_check (bool): Skip the "branch is up to date" check.
Only for intentional migrations against an older base.
"""
@ -225,6 +248,8 @@ def create_migration(
else:
_check_branch_freshness(root_dir, base_branch)
import testing.postgresql
try:
migrations_dir = (
root_dir / "litellm-proxy-extras" / "litellm_proxy_extras" / "migrations"
@ -342,9 +367,8 @@ if __name__ == "__main__":
)
parser.add_argument(
"--base-branch",
default=DEFAULT_BASE_BRANCH,
help=(
f"Branch to check freshness against (default: {DEFAULT_BASE_BRANCH}). "
"Branch to check freshness against (default: origin's current default branch). "
"The script fetches origin/<base-branch> and refuses to run if HEAD "
"is behind it."
),

6
docker/component_entrypoint.sh
View file

@ -1,5 +1,11 @@
#!/bin/sh
# stale samples from a previous container incarnation would be summed into the aggregate
if [ -n "$PROMETHEUS_MULTIPROC_DIR" ]; then
mkdir -p "$PROMETHEUS_MULTIPROC_DIR"
rm -f "$PROMETHEUS_MULTIPROC_DIR"/*.db
fi
case "$USE_DDTRACE" in
[Tt][Rr][Uu][Ee])
export DD_TRACE_OPENAI_ENABLED="False"

27
enterprise/litellm_enterprise/enterprise_callbacks/secret_detection.py
View file

@ -433,9 +433,9 @@ _default_detect_secrets_config = {
"name": "ZendeskSecretKeyDetector",
"path": _custom_plugins_path + "/zendesk_secret_key.py",
},
{"name": "Base64HighEntropyString", "limit": 3.0},
{"name": "Base64HighEntropyString", "limit": 4.5},
{"name": "HexHighEntropyString", "limit": 3.0},
]
],
}
@ -466,16 +466,19 @@ class _ENTERPRISE_SecretDetection(CustomGuardrail):
os.remove(temp_file.name)
detected_secrets = []
for file in secrets.files:
for found_secret in secrets[file]:
if found_secret.secret_value is None:
continue
detected_secrets.append(
{"type": found_secret.type, "value": found_secret.secret_value}
)
return detected_secrets
return [
{"type": found_secret.type, "value": found_secret.secret_value}
for file in sorted(secrets.files)
for found_secret in sorted(
secrets[file],
key=lambda secret: (
-len(secret.secret_value or ""),
secret.type,
secret.secret_value or "",
),
)
if found_secret.secret_value is not None
]
def redact_text(self, text: str, source: str = "message") -> str:
"""Replace every detected secret in ``text`` with ``[REDACTED]`` and

15
enterprise/litellm_enterprise/enterprise_callbacks/secrets_plugins/openai_api_key.py
View file

@ -3,6 +3,7 @@ This plugin searches for OpenAI API Keys.
"""
import re
from collections.abc import Generator
from detect_secrets.plugins.base import RegexBasedDetector
@ -16,4 +17,16 @@ class OpenAIApiKeyDetector(RegexBasedDetector):
@property
def denylist(self) -> list[re.Pattern]:
return [re.compile(r"""(sk-[a-zA-Z0-9]{5,})""")]
return [
re.compile(
r"((?:(?<![a-zA-Z0-9])|(?<=%[0-9A-Fa-f]{2}))"
r"sk[-_]"
r"[a-zA-Z0-9_-]{5,}"
r"(?![a-zA-Z0-9_-]))"
)
]
def analyze_string(self, string: str) -> Generator[str, None, None]:
# the digit check lives outside the regex: a lookahead re-scans the token
# from every `sk` inside it, which is quadratic on `-sk-sk-sk-...` input
yield from (match for match in super().analyze_string(string) if re.search(r"[0-9]", match))

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

@ -142,6 +142,40 @@ class CheckBatchCost:
verbose_proxy_logger.error(f"CheckBatchCost: could not look up team alias for team {team_id}: {e}")
return None
async def _get_org_id(self, job: "LiteLLM_ManagedObjectTable", batch_id: str) -> str | None:
org_id = getattr(job, "org_id", None)
if org_id:
return org_id
api_key = getattr(job, "api_key", None)
team_id = getattr(job, "team_id", None)
if api_key:
try:
key_row: prisma_models.LiteLLM_VerificationToken | None = (
await self.prisma_client.db.litellm_verificationtoken.find_unique(
where={"token": api_key}
)
)
key_org_id = getattr(key_row, "organization_id", None) if key_row is not None else None
if key_org_id:
return key_org_id
except Exception as e:
verbose_proxy_logger.error(
f"CheckBatchCost: could not resolve the key's org for batch {batch_id}, "
f"still trying the team's: {e}"
)
if not team_id:
return None
try:
team_row: prisma_models.LiteLLM_TeamTable | None = (
await self.prisma_client.db.litellm_teamtable.find_unique(
where={"team_id": team_id}
)
)
return getattr(team_row, "organization_id", None) if team_row is not None else None
except Exception as e:
verbose_proxy_logger.error(f"CheckBatchCost: could not resolve the team's org for batch {batch_id}: {e}")
return None
async def _build_creator_attribution_metadata(
self, job: "LiteLLM_ManagedObjectTable", batch_id: str
) -> dict[str, object]:
@ -153,6 +187,10 @@ class CheckBatchCost:
user_api_key_alias; when it has no alias, or the key has since been rotated or
deleted, the field keeps the creating user's alias that _get_user_info filled in,
because a resolvable name is more useful on the spend row than a null.
user_api_key_org_id must be resolved here too: the spend update writer reads it
off this metadata to increment organization spend, so leaving it out silently
drops batch cost from org accounting for keys and teams that belong to one.
"""
api_key = getattr(job, "api_key", None)
team_id = getattr(job, "team_id", None)
@ -172,6 +210,9 @@ class CheckBatchCost:
team_alias = await self._get_team_alias(team_id)
if team_alias is not None:
metadata["user_api_key_team_alias"] = team_alias
org_id: Final = await self._get_org_id(job, batch_id)
if org_id is not None:
metadata["user_api_key_org_id"] = org_id
if isinstance(request_tags, list) and request_tags:
metadata["tags"] = [tag for tag in request_tags if isinstance(tag, str)]
@ -641,7 +682,7 @@ class CheckBatchCost:
from litellm.files.main import afile_content
from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
from litellm.litellm_core_utils.litellm_logging import deployment_pricing_model_info
from litellm.litellm_core_utils.litellm_logging import deployment_pricing_model_info, mask_api_base_credentials
from litellm.proxy.openai_files_endpoints.common_utils import (
_is_base64_encoded_unified_file_id,
)
@ -805,6 +846,7 @@ class CheckBatchCost:
function_id=str(uuid.uuid4()),
)
deployment_api_base: Final = deployment_info.litellm_params.api_base
logging_obj.update_environment_variables(
litellm_params={
# set the user-agent header so that S3 callback consumers can easily identify CheckBatchCost callbacks
@ -813,9 +855,17 @@ class CheckBatchCost:
"user-agent": CHECK_BATCH_COST_USER_AGENT,
}
},
"metadata": await self._build_creator_attribution_metadata(job, batch_id),
**({"api_base": mask_api_base_credentials(deployment_api_base)} if deployment_api_base else {}),
"metadata": {
**(await self._build_creator_attribution_metadata(job, batch_id)),
# spend logs read the deployment identity off these metadata keys, so
# without them the batch cost row carries no model_id or model_group
"model_info": {"id": model_id},
"model_group": deployment_info.model_name,
},
},
optional_params={},
custom_llm_provider=str(llm_provider) if llm_provider else None,
)
if not await self._claim_job_for_costing(job):
@ -833,6 +883,8 @@ class CheckBatchCost:
batch_models=batch_result.models,
batch_successful_requests=batch_result.successful_requests,
batch_failed_requests=batch_result.failed_requests,
batch_prompt_cost=batch_result.prompt_cost,
batch_completion_cost=batch_result.completion_cost,
)
except Exception:
await self._release_job_claim(job)

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

@ -46,6 +46,7 @@ from litellm.proxy._types import (
UserAPIKeyAuth,
)
from litellm.proxy.openai_files_endpoints.common_utils import (
BATCH_CREATE_HIDDEN_PARAM,
FILE_LIST_CONTINUATION_CHUNK_SIZE,
MAX_FILE_LIST_LIMIT,
_is_base64_encoded_unified_file_id,
@ -279,6 +280,27 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
)
verbose_logger.debug(f"LiteLLM Managed File object with id={file_id} stored in db: {result}")
async def _resolve_creator_org_id(self, user_api_key_dict: UserAPIKeyAuth) -> Optional[str]:
if user_api_key_dict.org_id:
return user_api_key_dict.org_id
if not user_api_key_dict.team_id:
return None
from litellm.proxy.auth.auth_checks import get_team_object
from litellm.proxy.proxy_server import proxy_logging_obj, user_api_key_cache
try:
team: Final = await get_team_object(
team_id=user_api_key_dict.team_id,
prisma_client=self.prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=user_api_key_dict.parent_otel_span,
proxy_logging_obj=proxy_logging_obj,
)
return team.organization_id
except Exception as e:
verbose_logger.warning(f"could not resolve org for managed object attribution: {e}")
return None
async def store_unified_object_id(
self,
unified_object_id: str,
@ -351,6 +373,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
"file_purpose": file_purpose,
"created_by": resolve_resource_owner_id(user_api_key_dict),
"team_id": user_api_key_dict.team_id,
"org_id": await self._resolve_creator_org_id(user_api_key_dict),
"updated_by": user_api_key_dict.user_id,
"status": file_object.status,
**attribution_columns,
@ -1321,7 +1344,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
## Check if unified_file_id is in the response
unified_file_id = response._hidden_params.get("unified_file_id") # managed file id
unified_batch_id = response._hidden_params.get("unified_batch_id") # managed batch id
is_batch_create: Final = unified_file_id is not None
is_batch_create: Final = response._hidden_params.get(BATCH_CREATE_HIDDEN_PARAM) is True
model_id = cast(Optional[str], response._hidden_params.get("model_id"))
model_name = cast(Optional[str], response._hidden_params.get("model_name"))
@ -1410,10 +1433,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
litellm_parent_otel_span=user_api_key_dict.parent_otel_span,
)
# Only record batch creation metric on actual create (not retrieve/cancel).
# unified_file_id in _hidden_params is only set by the create_batch endpoint.
original_unified_file_id = response._hidden_params.get("unified_file_id")
if original_unified_file_id:
if is_batch_create:
prom_logger = self._get_prometheus_logger()
if prom_logger:
batch_provider = ""
@ -1781,7 +1801,16 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
# Remove conflicting keys from data to avoid duplicate keyword arguments
filtered_data = {k: v for k, v in data.items() if k not in ("model", "file_id")}
for model_id, model_file_id in specific_model_file_id_mapping.items():
delete_response = await llm_router.afile_delete(model=model_id, file_id=model_file_id, **filtered_data) # type: ignore
credentials = llm_router.get_deployment_credentials_with_provider(model_id=model_id)
delete_data = {
**{k: v for k, v in filtered_data.items() if k != "_litellm_internal_model_credentials"},
**(
{"_litellm_internal_model_credentials": MappingProxyType(dict(credentials))}
if credentials is not None
else {}
),
}
delete_response = await llm_router.afile_delete(model=model_id, file_id=model_file_id, **delete_data)
stored_file_object = await self.delete_unified_file_id(file_id, litellm_parent_otel_span)
@ -1792,7 +1821,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
prom_logger.record_managed_file_deleted(result="success")
if stored_file_object:
return stored_file_object
return OpenAIFileObject.model_validate(stored_file_object).model_copy(update={"id": file_id})
elif delete_response:
delete_response.id = file_id
return delete_response

12
enterprise/litellm_enterprise/proxy/hooks/managed_vector_stores.py
View file

@ -16,6 +16,7 @@ from litellm.llms.base_llm.managed_resources.utils import (
is_base64_encoded_unified_id,
)
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.utils import LLMResponseTypes
from litellm.types.vector_stores import (
VectorStoreCreateOptionalRequestParams,
VectorStoreCreateResponse,
@ -24,6 +25,7 @@ from litellm.types.vector_stores import (
if TYPE_CHECKING:
from opentelemetry.trace import Span as _Span
from litellm.caching.caching import DualCache
from litellm.proxy.utils import InternalUsageCache as _InternalUsageCache
from litellm.proxy.utils import PrismaClient as _PrismaClient
@ -156,7 +158,7 @@ class _PROXY_LiteLLMManagedVectorStores(
# Create vector store for each model
# Convert TypedDict to Dict[str, Any] for base class compatibility
request_data_dict: Dict[str, Any] = dict(create_request)
request_data_dict: Dict[str, object] = dict(create_request)
responses = await self.create_resource_for_each_model(
llm_router=llm_router,
request_data=request_data_dict,
@ -209,7 +211,7 @@ class _PROXY_LiteLLMManagedVectorStores(
limit: Optional[int] = None,
after: Optional[str] = None,
order: Optional[str] = None,
) -> Dict[str, Any]:
) -> Dict[str, object]:
"""
List vector stores created by a user.
@ -301,7 +303,7 @@ class _PROXY_LiteLLMManagedVectorStores(
async def async_pre_call_hook(
self,
user_api_key_dict: UserAPIKeyAuth,
cache: Any,
cache: "DualCache",
data: Dict,
call_type: str,
) -> Union[Exception, str, Dict, None]:
@ -403,8 +405,8 @@ class _PROXY_LiteLLMManagedVectorStores(
self,
data: Dict,
user_api_key_dict: UserAPIKeyAuth,
response: Any,
) -> Any:
response: LLMResponseTypes,
) -> LLMResponseTypes:
"""
Post-call hook to transform responses.

4
enterprise/pyproject.toml
View file

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

12
helm/litellm-helm/templates/deployment.yaml
View file

@ -152,6 +152,13 @@ spec:
{{- if .Values.billingMetrics.enabled }}
{{- include "litellm.billingMetricsEnv" . | nindent 12 }}
{{- end }}
{{- if .Values.metricsServer.enabled }}
{{- if eq (int .Values.metricsServer.port) (int .Values.service.port) }}
{{- fail "metricsServer.port must differ from service.port" }}
{{- end }}
- name: PROMETHEUS_METRICS_PORT
value: {{ .Values.metricsServer.port | quote }}
{{- end }}
{{- if .Values.migrationJob.enabled }}
# Schema updates are owned by the dedicated migrations Job; skip
# the proxy's startup `prisma db push` so N replicas don't race
@ -189,6 +196,11 @@ spec:
- name: http
containerPort: {{ .Values.service.port }}
protocol: TCP
{{- if .Values.metricsServer.enabled }}
- name: metrics
containerPort: {{ .Values.metricsServer.port }}
protocol: TCP
{{- end }}
livenessProbe:
httpGet:
path: {{ .Values.livenessProbe.path | quote }}

17
helm/litellm-helm/templates/service-metrics.yaml Normal file
View file

@ -0,0 +1,17 @@
{{- if .Values.metricsServer.enabled }}
apiVersion: v1
kind: Service
metadata:
name: {{ include "litellm.fullname" . }}-metrics
labels:
{{- include "litellm.labels" . | nindent 4 }}
spec:
type: ClusterIP
ports:
- port: {{ .Values.metricsServer.port }}
targetPort: metrics
protocol: TCP
name: metrics
selector:
{{- include "litellm.selectorLabels" . | nindent 4 }}
{{- end }}

2
helm/litellm-helm/templates/servicemonitor.yaml
View file

@ -26,7 +26,7 @@ spec:
{{- toYaml .namespaceSelector.matchNames | nindent 4 }}
{{- end }}
endpoints:
- port: http
- port: {{ ternary "metrics" "http" $.Values.metricsServer.enabled }}
path: /metrics/
interval: {{ .interval }}
scrapeTimeout: {{ .scrapeTimeout }}

106
helm/litellm-helm/tests/metrics_server_tests.yaml Normal file
View file

@ -0,0 +1,106 @@
suite: separate metrics server
templates:
- configmap-litellm.yaml
- deployment.yaml
- service.yaml
- service-metrics.yaml
- servicemonitor.yaml
tests:
- it: should not expose a metrics port or PROMETHEUS_METRICS_PORT by default
asserts:
- notContains:
path: spec.template.spec.containers[0].ports
content:
name: metrics
any: true
template: deployment.yaml
- notContains:
path: spec.template.spec.containers[0].env
content:
name: PROMETHEUS_METRICS_PORT
any: true
template: deployment.yaml
- lengthEqual:
path: spec.ports
count: 1
template: service.yaml
- hasDocuments:
count: 0
template: service-metrics.yaml
- it: should scrape the proxy port when the metrics server is disabled
template: servicemonitor.yaml
set:
serviceMonitor.enabled: true
asserts:
- equal:
path: spec.endpoints[0].port
value: http
- it: should wire the separate metrics server through container, a ClusterIP metrics service and servicemonitor
set:
metricsServer.enabled: true
metricsServer.port: 4101
serviceMonitor.enabled: true
service.type: LoadBalancer
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: PROMETHEUS_METRICS_PORT
value: "4101"
template: deployment.yaml
- contains:
path: spec.template.spec.containers[0].ports
content:
name: metrics
containerPort: 4101
protocol: TCP
template: deployment.yaml
- lengthEqual:
path: spec.ports
count: 1
template: service.yaml
- equal:
path: spec.type
value: LoadBalancer
template: service.yaml
- equal:
path: metadata.name
value: RELEASE-NAME-litellm-metrics
template: service-metrics.yaml
- equal:
path: spec.type
value: ClusterIP
template: service-metrics.yaml
- equal:
path: spec.ports
value:
- port: 4101
targetPort: metrics
protocol: TCP
name: metrics
template: service-metrics.yaml
- equal:
path: spec.selector
value:
app.kubernetes.io/name: litellm
app.kubernetes.io/instance: RELEASE-NAME
template: service-metrics.yaml
- equal:
path: spec.endpoints[0].port
value: metrics
template: servicemonitor.yaml
- equal:
path: spec.endpoints[0].path
value: /metrics/
template: servicemonitor.yaml
- it: should reject a metrics port equal to the proxy port
template: deployment.yaml
set:
metricsServer.enabled: true
metricsServer.port: 4000
asserts:
- failedTemplate:
errorMessage: metricsServer.port must differ from service.port

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

@ -180,6 +180,16 @@ proxy_config:
general_settings:
master_key: os.environ/PROXY_MASTER_KEY
# Serve Prometheus /metrics from a separate process (PROMETHEUS_METRICS_PORT)
# so a scrape never runs on an inference worker. Adds a `metrics` port to the
# container and a dedicated ClusterIP `<release>-metrics` Service, and the
# ServiceMonitor scrapes it instead of the proxy port. The separate port has
# no virtual-key auth: keep it off public ingress. Needs the proxy image
# v1.101.0 or newer.
metricsServer:
enabled: false
port: 4001
resources:
{}
# Unset by default so the chart installs on small clusters such as Minikube, and so an

2
helm/litellm/templates/_helpers.tpl
View file

@ -441,3 +441,5 @@ ImplementationSpecific
{{- .pathType -}}
{{- end -}}
{{- end -}}
{{- define "litellm.gateway.prometheusMultiprocDir" -}}/tmp/litellm_prometheus_multiproc{{- end -}}

53
helm/litellm/templates/gateway/deployment.yaml
View file

@ -64,14 +64,25 @@ spec:
{{- if .Values.billingMetrics.enabled }}
{{- include "litellm.billingMetricsEnv" . | nindent 12 }}
{{- end }}
{{- if .Values.gateway.metricsServer.enabled }}
{{- if eq (int .Values.gateway.metricsServer.port) 4000 }}
{{- fail "gateway.metricsServer.port must differ from the gateway port 4000" }}
{{- end }}
- name: PROMETHEUS_MULTIPROC_DIR
value: {{ include "litellm.gateway.prometheusMultiprocDir" . }}
{{- end }}
{{- include "litellm.envFrom" .Values.gateway | nindent 10 }}
{{- if or .Values.gateway.config.create .Values.gateway.volumeMounts .Values.billingMetrics.enabled }}
{{- if or .Values.gateway.config.create .Values.gateway.volumeMounts .Values.billingMetrics.enabled .Values.gateway.metricsServer.enabled }}
volumeMounts:
{{- if .Values.gateway.config.create }}
- name: gateway-config
mountPath: /app/config/config.yaml
subPath: config.yaml
{{- end }}
{{- if .Values.gateway.metricsServer.enabled }}
- name: prometheus-multiproc
mountPath: {{ include "litellm.gateway.prometheusMultiprocDir" . }}
{{- end }}
{{- if .Values.billingMetrics.enabled }}
{{- include "litellm.billingMetricsVolumeMounts" . | nindent 12 }}
{{- end }}
@ -97,16 +108,54 @@ spec:
{{- end }}
resources:
{{- toYaml .Values.gateway.resources | nindent 12 }}
{{- if .Values.gateway.metricsServer.enabled }}
- name: metrics
image: "{{ .Values.gateway.image.repository }}:{{ .Values.gateway.image.tag | default .Chart.AppVersion }}"
imagePullPolicy: {{ .Values.gateway.image.pullPolicy }}
{{- with .Values.gateway.securityContext }}
securityContext:
{{- toYaml . | nindent 12 }}
{{- end }}
command:
- python
- -m
- litellm.proxy.prometheus_metrics_server
- --port
- {{ .Values.gateway.metricsServer.port | quote }}
env:
- name: PROMETHEUS_MULTIPROC_DIR
value: {{ include "litellm.gateway.prometheusMultiprocDir" . }}
ports:
- name: metrics
containerPort: {{ .Values.gateway.metricsServer.port }}
protocol: TCP
volumeMounts:
- name: prometheus-multiproc
mountPath: {{ include "litellm.gateway.prometheusMultiprocDir" . }}
readinessProbe:
tcpSocket: { port: metrics }
periodSeconds: 10
livenessProbe:
tcpSocket: { port: metrics }
periodSeconds: 15
failureThreshold: 6
resources:
{{- toYaml .Values.gateway.metricsServer.resources | nindent 12 }}
{{- end }}
{{- with .Values.gateway.extraContainers }}
{{- tpl (toYaml .) $ | nindent 8 }}
{{- end }}
{{- if or .Values.gateway.config.create .Values.gateway.volumes .Values.billingMetrics.enabled }}
{{- if or .Values.gateway.config.create .Values.gateway.volumes .Values.billingMetrics.enabled .Values.gateway.metricsServer.enabled }}
volumes:
{{- if .Values.gateway.config.create }}
- name: gateway-config
configMap:
name: {{ include "litellm.gateway.fullname" . }}-config
{{- end }}
{{- if .Values.gateway.metricsServer.enabled }}
- name: prometheus-multiproc
emptyDir: {}
{{- end }}
{{- if .Values.billingMetrics.enabled }}
{{- include "litellm.billingMetricsVolumes" . | nindent 8 }}
{{- end }}

18
helm/litellm/templates/gateway/service-metrics.yaml Normal file
View file

@ -0,0 +1,18 @@
{{- if and .Values.gateway.enabled .Values.gateway.metricsServer.enabled }}
apiVersion: v1
kind: Service
metadata:
name: {{ include "litellm.gateway.fullname" . }}-metrics
labels:
{{- include "litellm.commonLabels" . | nindent 4 }}
app.kubernetes.io/component: gateway
spec:
type: ClusterIP
ports:
- port: {{ .Values.gateway.metricsServer.port }}
targetPort: metrics
protocol: TCP
name: metrics
selector:
{{- include "litellm.gateway.selectorLabels" . | nindent 4 }}
{{- end }}

12
helm/litellm/templates/migrations-job.yaml
View file

@ -77,4 +77,16 @@ spec:
volumes:
{{- toYaml . | nindent 8 }}
{{- end }}
{{- with .Values.migrationJob.nodeSelector }}
nodeSelector:
{{- toYaml . | nindent 8 }}
{{- end }}
{{- with .Values.migrationJob.affinity }}
affinity:
{{- toYaml . | nindent 8 }}
{{- end }}
{{- with .Values.migrationJob.tolerations }}
tolerations:
{{- toYaml . | nindent 8 }}
{{- end }}
{{- end }}

148
helm/litellm/tests/metrics_server_tests.yaml Normal file
View file

@ -0,0 +1,148 @@
suite: test gateway metrics sidecar
templates:
- gateway/configmap.yaml
- gateway/deployment.yaml
- gateway/service.yaml
- gateway/service-metrics.yaml
values:
- ./values/required.yaml
tests:
- it: adds no sidecar, volume, env or service port when the metrics server is off
asserts:
- lengthEqual:
path: spec.template.spec.containers
count: 1
template: gateway/deployment.yaml
- notContains:
path: spec.template.spec.containers[0].env
content:
name: PROMETHEUS_MULTIPROC_DIR
any: true
template: gateway/deployment.yaml
- notContains:
path: spec.template.spec.volumes
content:
name: prometheus-multiproc
any: true
template: gateway/deployment.yaml
- lengthEqual:
path: spec.ports
count: 1
template: gateway/service.yaml
- hasDocuments:
count: 0
template: gateway/service-metrics.yaml
- it: runs the metrics server as a sidecar over a shared multiproc dir and exposes it on a ClusterIP metrics service
set:
gateway.metricsServer.enabled: true
gateway.metricsServer.port: 4101
gateway.service.type: LoadBalancer
gateway.image.tag: v1.101.0
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: PROMETHEUS_MULTIPROC_DIR
value: /tmp/litellm_prometheus_multiproc
template: gateway/deployment.yaml
- contains:
path: spec.template.spec.containers[0].volumeMounts
content:
name: prometheus-multiproc
mountPath: /tmp/litellm_prometheus_multiproc
template: gateway/deployment.yaml
- equal:
path: spec.template.spec.containers[1].name
value: metrics
template: gateway/deployment.yaml
- equal:
path: spec.template.spec.containers[1].image
value: ghcr.io/berriai/litellm-gateway:v1.101.0
template: gateway/deployment.yaml
- equal:
path: spec.template.spec.containers[1].command
value:
- python
- -m
- litellm.proxy.prometheus_metrics_server
- --port
- "4101"
template: gateway/deployment.yaml
- equal:
path: spec.template.spec.containers[1].env
value:
- name: PROMETHEUS_MULTIPROC_DIR
value: /tmp/litellm_prometheus_multiproc
template: gateway/deployment.yaml
- equal:
path: spec.template.spec.containers[1].ports
value:
- name: metrics
containerPort: 4101
protocol: TCP
template: gateway/deployment.yaml
- equal:
path: spec.template.spec.containers[1].volumeMounts
value:
- name: prometheus-multiproc
mountPath: /tmp/litellm_prometheus_multiproc
template: gateway/deployment.yaml
- equal:
path: spec.template.spec.containers[1].readinessProbe.tcpSocket.port
value: metrics
template: gateway/deployment.yaml
- equal:
path: spec.template.spec.containers[1].livenessProbe.tcpSocket.port
value: metrics
template: gateway/deployment.yaml
- equal:
path: spec.template.spec.containers[1].resources.requests.cpu
value: 50m
template: gateway/deployment.yaml
- contains:
path: spec.template.spec.volumes
content:
name: prometheus-multiproc
emptyDir: {}
template: gateway/deployment.yaml
- lengthEqual:
path: spec.ports
count: 1
template: gateway/service.yaml
- equal:
path: spec.type
value: LoadBalancer
template: gateway/service.yaml
- equal:
path: metadata.name
value: RELEASE-NAME-litellm-gateway-metrics
template: gateway/service-metrics.yaml
- equal:
path: spec.type
value: ClusterIP
template: gateway/service-metrics.yaml
- equal:
path: spec.ports
value:
- port: 4101
targetPort: metrics
protocol: TCP
name: metrics
template: gateway/service-metrics.yaml
- equal:
path: spec.selector
value:
app.kubernetes.io/name: litellm
app.kubernetes.io/instance: RELEASE-NAME
app.kubernetes.io/component: gateway
template: gateway/service-metrics.yaml
- it: rejects a metrics port equal to the gateway port
template: gateway/deployment.yaml
set:
gateway.metricsServer.enabled: true
gateway.metricsServer.port: 4000
asserts:
- failedTemplate:
errorMessage: gateway.metricsServer.port must differ from the gateway port 4000

68
helm/litellm/tests/migration_job_tests.yaml
View file

@ -1,4 +1,4 @@
suite: test migrations Job ServiceAccount resolution and pod hardening
suite: test migrations Job ServiceAccount resolution, pod hardening, and scheduling
templates:
- migrations-job.yaml
values:
@ -188,3 +188,69 @@ tests:
asserts:
- notExists:
path: spec.activeDeadlineSeconds
- it: renders no scheduling fields by default
asserts:
- isNull:
path: spec.template.spec.nodeSelector
- isNull:
path: spec.template.spec.tolerations
- isNull:
path: spec.template.spec.affinity
- it: renders nodeSelector, tolerations, and affinity from the migrationJob values
set:
migrationJob.nodeSelector:
intent: no-csi-nodes
migrationJob.tolerations:
- key: intent
operator: Equal
value: no-csi-nodes
effect: NoSchedule
migrationJob.affinity:
nodeAffinity:
requiredDuringSchedulingIgnoredDuringExecution:
nodeSelectorTerms:
- matchExpressions:
- key: intent
operator: In
values:
- no-csi-nodes
asserts:
- equal:
path: spec.template.spec.nodeSelector
value:
intent: no-csi-nodes
- equal:
path: spec.template.spec.tolerations
value:
- key: intent
operator: Equal
value: no-csi-nodes
effect: NoSchedule
- equal:
path: spec.template.spec.affinity
value:
nodeAffinity:
requiredDuringSchedulingIgnoredDuringExecution:
nodeSelectorTerms:
- matchExpressions:
- key: intent
operator: In
values:
- no-csi-nodes
- it: does not inherit the gateway's scheduling values
set:
gateway.nodeSelector:
intent: no-csi-nodes
gateway.tolerations:
- key: intent
operator: Equal
value: no-csi-nodes
effect: NoSchedule
asserts:
- isNull:
path: spec.template.spec.nodeSelector
- isNull:
path: spec.template.spec.tolerations

23
helm/litellm/values.yaml
View file

@ -152,6 +152,13 @@ migrationJob:
# the writable scratch space a read-only root filesystem needs.
volumes: []
volumeMounts: []
# Scheduling for the Job pod, same shape as gateway.nodeSelector /
# gateway.tolerations / gateway.affinity. The Job does not inherit the other
# components' scheduling values: a migration usually needs a larger node
# than the gateway, so pin it here explicitly.
nodeSelector: {}
tolerations: []
affinity: {}
image:
repository: ghcr.io/berriai/litellm-migrations
tag: "" # defaults to .Chart.AppVersion
@ -261,6 +268,22 @@ gateway:
config:
create: true
proxy_config: {}
# Serve Prometheus /metrics from a `metrics` sidecar container (same image,
# `python -m litellm.proxy.prometheus_metrics_server`) that aggregates the
# workers' PROMETHEUS_MULTIPROC_DIR samples over a shared emptyDir, so a
# scrape never runs on an inference worker. Adds a `metrics` port to the pod
# and a dedicated ClusterIP `<gateway>-metrics` Service; point your scrape
# config at it. The port has no virtual-key auth: keep it off public ingress.
# Needs the gateway image v1.101.0 or newer.
metricsServer:
enabled: false
port: 4001
resources:
requests:
cpu: 50m
memory: 128Mi
limits:
memory: 512Mi
image:
repository: ghcr.io/berriai/litellm-gateway
tag: "" # defaults to .Chart.AppVersion

15
litellm-proxy-extras/litellm_proxy_extras/migrations/20260717000000_cascade_delete_jwt_key_mapping_on_token_delete/migration.sql Normal file
View file

@ -0,0 +1,15 @@
-- DropForeignKey
DO $$
BEGIN
IF EXISTS (SELECT 1 FROM pg_constraint WHERE conname = 'LiteLLM_JWTKeyMapping_token_fkey') THEN
ALTER TABLE "LiteLLM_JWTKeyMapping" DROP CONSTRAINT "LiteLLM_JWTKeyMapping_token_fkey";
END IF;
END $$;
-- AddForeignKey
DO $$
BEGIN
IF NOT EXISTS (SELECT 1 FROM pg_constraint WHERE conname = 'LiteLLM_JWTKeyMapping_token_fkey') THEN
ALTER TABLE "LiteLLM_JWTKeyMapping" ADD CONSTRAINT "LiteLLM_JWTKeyMapping_token_fkey" FOREIGN KEY ("token") REFERENCES "LiteLLM_VerificationToken"("token") ON DELETE CASCADE ON UPDATE CASCADE;
END IF;
END $$;

4
litellm-proxy-extras/litellm_proxy_extras/migrations/20260903230000_add_org_id_to_managed_object_table/migration.sql Normal file
View file

@ -0,0 +1,4 @@
-- Add org_id column to LiteLLM_ManagedObjectTable
-- Snapshots the creating key's organization at submission time, like team_id,
-- so CheckBatchCost can bill organization spend hours later without re-resolving
ALTER TABLE "LiteLLM_ManagedObjectTable" ADD COLUMN IF NOT EXISTS "org_id" TEXT;

1
litellm-proxy-extras/litellm_proxy_extras/migrations/20260904000000_shadow_eval_model_scope/migration.sql Normal file
View file

@ -0,0 +1 @@
ALTER TABLE "LiteLLM_ShadowEvalJob" ADD COLUMN IF NOT EXISTS "models" TEXT[] NOT NULL DEFAULT ARRAY[]::TEXT[];

2
litellm-proxy-extras/litellm_proxy_extras/migrations/20260905120000_add_per_server_oauth_discovery_to_mcp_servers/migration.sql Normal file
View file

@ -0,0 +1,2 @@
-- AlterTable
ALTER TABLE "LiteLLM_MCPServerTable" ADD COLUMN IF NOT EXISTS "per_server_oauth_discovery" BOOLEAN NOT NULL DEFAULT false;

3
litellm-proxy-extras/litellm_proxy_extras/migrations/20260907000000_add_autorouter_classifier_cost/migration.sql Normal file
View file

@ -0,0 +1,3 @@
ALTER TABLE "LiteLLM_AutoRouterSession"
ADD COLUMN IF NOT EXISTS "classifier_cost" DOUBLE PRECISION NOT NULL DEFAULT 0,
ADD COLUMN IF NOT EXISTS "classifier_cost_recorded_turns" INTEGER NOT NULL DEFAULT 0;

7
litellm-proxy-extras/litellm_proxy_extras/schema.prisma
View file

@ -343,6 +343,7 @@ model LiteLLM_MCPServerTable {
delegate_auth_to_upstream Boolean @default(false)
oauth_passthrough Boolean @default(false)
dcr_bridge Boolean?
per_server_oauth_discovery Boolean @default(false)
is_byok Boolean @default(false)
byok_description String[] @default([])
byok_api_key_help_url String?
@ -491,7 +492,7 @@ model LiteLLM_JWTKeyMapping {
updated_at DateTime @default(now()) @updatedAt
updated_by String?
litellm_verification_token LiteLLM_VerificationToken @relation(fields: [token], references: [token])
litellm_verification_token LiteLLM_VerificationToken @relation(fields: [token], references: [token], onDelete: Cascade)
@@unique([jwt_claim_name, jwt_claim_value])
@@index([jwt_claim_name, jwt_claim_value, is_active])
@ -1035,6 +1036,7 @@ model LiteLLM_ManagedObjectTable { // for batches or finetuning jobs which use t
created_at DateTime @default(now())
created_by String?
team_id String?
org_id String? // creating key's organization at submission time; CheckBatchCost bills org spend against it
api_key String?
request_tags Json? @default("[]")
updated_at DateTime @updatedAt
@ -1508,6 +1510,8 @@ model LiteLLM_AutoRouterSession {
total_tokens BigInt @default(0)
spend Float @default(0)
saved_spend Float @default(0)
classifier_cost Float @default(0)
classifier_cost_recorded_turns Int @default(0)
tier_turns Json @default("{}")
@@id([api_key, session_id, router_name])
@ -1536,6 +1540,7 @@ model LiteLLM_ShadowEvalJob {
target_id String // hashed virtual key, team_id, or user_id whose traffic this leg shadows
router_name String // first (often only) auto-router under evaluation; router_names is the full set
router_names String[] @default([]) // all routers this job runs as shadow arms; empty on legacy rows, whose set is (router_name)
models String[] @default([]) // model groups the sampled traffic is narrowed to; empty samples every model
direction String @default("forward") // forward | reverse
baseline_model String? // reverse only: the fixed model the router is judged against
judge_model String

130
litellm-proxy-extras/litellm_proxy_extras/utils.py
View file

@ -8,14 +8,10 @@ import tempfile
import time
from dataclasses import dataclass, replace
from pathlib import Path
from typing import Optional
from typing import TYPE_CHECKING, Final, Optional
from litellm_proxy_extras import prisma_toolchain
from litellm_proxy_extras._logging import logger
from litellm_proxy_extras.replica_identity import (
REPLICA_IDENTITY_FULL_ENV_VAR,
apply_replica_identity_full,
)
from litellm_proxy_extras.prisma_toolchain import (
PRISMA_COMMAND_TIMEOUT_ENV_VAR,
PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR,
@ -23,6 +19,14 @@ from litellm_proxy_extras.prisma_toolchain import (
prisma_command_timeout,
prisma_migrate_deploy_timeout,
)
from litellm_proxy_extras.replica_identity import (
REPLICA_IDENTITY_FULL_ENV_VAR,
apply_replica_identity_full,
)
if TYPE_CHECKING:
import psycopg
import psycopg.sql
def str_to_bool(value: Optional[str]) -> bool:
@ -46,6 +50,28 @@ def _get_prisma_env() -> dict:
_MIGRATION_TS_RE = re.compile(r"^(\d{14})_")
_MIGRATION_DEADLOCK_MARKER = "deadlock detected"
INDEX_REPAIR_ADVISORY_LOCK_KEY: Final = int.from_bytes(b"litellm", "big")
_TRANSIENT_INDEX_SUFFIX_RE: Final = re.compile(r"_cc(?:new|old)\d*$")
_INVALID_LITELLM_INDEXES_SQL: Final = (
"SELECT n.nspname, c.relname, pg_size_pretty(pg_table_size(t.oid)) "
"FROM pg_index i "
"JOIN pg_class c ON c.oid = i.indexrelid "
"JOIN pg_class t ON t.oid = i.indrelid "
"JOIN pg_namespace n ON n.oid = t.relnamespace "
"WHERE NOT i.indisvalid "
" AND c.relkind = 'i' "
" AND n.nspname = %s "
" AND t.relname LIKE %s "
" AND NOT EXISTS (SELECT 1 FROM pg_constraint k WHERE k.conindid = i.indexrelid) "
"ORDER BY c.relname"
)
@dataclass(frozen=True, slots=True)
class _InvalidIndex:
schema: str
name: str
table_size: str
MAX_MIGRATE_DEPLOY_ATTEMPTS = 4
@ -624,7 +650,7 @@ class ProxyExtrasDBManager:
def _strip_prisma_query_params(url: str) -> str:
"""Remove Prisma-specific query params (connection_limit, pool_timeout,
schema, etc.) from DATABASE_URL so psycopg can parse it."""
from urllib.parse import urlparse, urlunparse, parse_qsl, urlencode
from urllib.parse import parse_qsl, quote, urlencode, urlparse, urlunparse
parsed = urlparse(url)
if not parsed.query:
@ -645,7 +671,7 @@ class ProxyExtrasDBManager:
"target_session_attrs",
}
kept = [(k, v) for k, v in parse_qsl(parsed.query) if k in libpq_params]
return urlunparse(parsed._replace(query=urlencode(kept)))
return urlunparse(parsed._replace(query=urlencode(kept, quote_via=quote)))
@staticmethod
def _warn_if_db_ahead_of_head(migrations_dir: str) -> None:
@ -719,6 +745,95 @@ class ProxyExtrasDBManager:
", ".join(sorted_hostile[:5]) + (" ..." if len(sorted_hostile) > 5 else ""),
)
@staticmethod
def _invalid_litellm_indexes(
conn: "psycopg.Connection[tuple[str, str, str]]", schema: str
) -> tuple[_InvalidIndex, ...]:
rows: Final = conn.execute(_INVALID_LITELLM_INDEXES_SQL, (schema, "LiteLLM\\_%")).fetchall()
return tuple(_InvalidIndex(*row) for row in rows)
@staticmethod
def _index_repair(index: _InvalidIndex) -> tuple["psycopg.sql.Composed", str]:
from psycopg import sql
target: Final = sql.Identifier(index.schema, index.name)
if _TRANSIENT_INDEX_SUFFIX_RE.search(index.name):
return sql.SQL("DROP INDEX CONCURRENTLY IF EXISTS {}").format(target), "Dropped leftover"
return sql.SQL("REINDEX INDEX CONCURRENTLY {}").format(target), "Rebuilt"
@staticmethod
def _repair_index(conn: "psycopg.Connection[tuple[str, str, str]]", index: _InvalidIndex) -> None:
import psycopg
statement, action = ProxyExtrasDBManager._index_repair(index)
try:
conn.execute(statement)
except psycopg.Error as e:
logger.warning(
"Could not repair invalid index %s.%s, will retry on the next startup. "
"If this keeps happening, run `%s` by hand as the index owner. Error: %s",
index.schema,
index.name,
statement.as_string(conn),
e,
)
return
logger.info("%s invalid index %s.%s", action, index.schema, index.name)
@staticmethod
def repair_invalid_indexes(lock_timeout: str = "30s") -> bool:
"""Rebuild LiteLLM indexes an interrupted CREATE INDEX CONCURRENTLY left
INVALID (a migration deadlock between replicas is the usual cause; the
retried migration skips them because of IF NOT EXISTS). Never raises:
returns True when no invalid index remains, False when the repair was
skipped or failed and will be retried on the next startup. Looks in the
schema DATABASE_URL names, the only URL Prisma migrates through, but
connects over DIRECT_URL when set: the session settings, the advisory
lock and REINDEX CONCURRENTLY all need one server session, which a
transaction pooler does not give."""
prisma_url: Final = os.getenv("DATABASE_URL")
if not prisma_url:
return False
try:
import psycopg
from psycopg import sql
except ImportError:
logger.warning(
"psycopg is not installed; skipping the invalid index check. "
"Install the litellm[extra_proxy] extra, which includes psycopg."
)
return False
schema: Final = ProxyExtrasDBManager._prisma_schema_param(prisma_url) or "public"
cleaned_url: Final = ProxyExtrasDBManager._strip_prisma_query_params(os.getenv("DIRECT_URL") or prisma_url)
try:
with psycopg.connect(cleaned_url, connect_timeout=10, autocommit=True) as conn:
conn.execute("SET statement_timeout = 0")
conn.execute(sql.SQL("SET lock_timeout = {}").format(sql.Literal(lock_timeout)))
found: Final = ProxyExtrasDBManager._invalid_litellm_indexes(conn, schema)
if not found:
return True
logger.warning(
"Found %d invalid index(es) left by an interrupted CREATE INDEX "
"CONCURRENTLY, rebuilding: %s",
len(found),
", ".join(f"{index.name} (table size {index.table_size})" for index in found),
)
lock_row: Final = conn.execute(
"SELECT pg_try_advisory_lock(%s)", (INDEX_REPAIR_ADVISORY_LOCK_KEY,)
).fetchone()
if lock_row is None or not lock_row[0]:
logger.info("Another replica is already rebuilding the invalid indexes, skipping")
return False
for index in ProxyExtrasDBManager._invalid_litellm_indexes(conn, schema):
ProxyExtrasDBManager._repair_index(conn, index)
remaining: Final = ProxyExtrasDBManager._invalid_litellm_indexes(conn, schema)
except psycopg.Error as e:
logger.warning("Could not check for invalid indexes, will retry on the next startup. Error: %s", e)
return False
return not remaining
@staticmethod
def _setup_database_v2(use_migrate: bool) -> bool:
"""
@ -994,6 +1109,7 @@ class ProxyExtrasDBManager:
use_migrate=use_migrate, use_v2_resolver=use_v2_resolver
)
if migrated:
ProxyExtrasDBManager.repair_invalid_indexes()
ProxyExtrasDBManager.apply_replica_identity_full_if_requested()
return migrated

11
litellm-proxy-extras/migration_runbook.md
View file

@ -48,7 +48,7 @@ uv run --with testing.postgresql python ci_cd/run_migration.py "your_migration_n
## What It Does
1. **Verifies the current branch is up to date with `origin/litellm_internal_staging`** (see [Branch freshness](#branch-freshness-check))
1. **Verifies the current branch is up to date with origin's current default branch** (see [Branch freshness](#branch-freshness-check))
2. Creates temp PostgreSQL DB
3. Applies existing migrations
4. Compares with `schema.prisma`
@ -57,11 +57,11 @@ uv run --with testing.postgresql python ci_cd/run_migration.py "your_migration_n
## Branch Freshness Check
Before generating anything, `run_migration.py` runs `git fetch origin <base>` and refuses to proceed if `HEAD` is behind `origin/<base>`. Default base is `litellm_internal_staging` (the branch PRs target). A previous incident saw a stale branch silently drop production columns; freshness is the first-line defense.
Before generating anything, `run_migration.py` runs `git fetch origin <base>` and refuses to proceed if `HEAD` is behind `origin/<base>`. The default base is discovered from origin's advertised HEAD on each run, so an existing clone follows a default-branch change without trusting cached `origin/HEAD`. If discovery or fetching fails, migration generation stops. A previous incident saw a stale branch silently drop production columns; freshness is the first-line defense.
Flags:
- `--base-branch <name>` — check against a different base (e.g. `main`). Default is `litellm_internal_staging`.
- `--base-branch <name>` — check against a different base (e.g. a release branch). Defaults to origin's current default branch
- `--skip-freshness-check` — bypass entirely. Only for intentional migrations against an older base.
When the guard fires:
@ -69,8 +69,9 @@ When the guard fires:
1. Update your branch:
```bash
git fetch origin && git rebase origin/litellm_internal_staging
# or git merge origin/litellm_internal_staging — whichever matches your workflow
base_branch=$(python3 scripts/default_branch.py --branch) &&
git fetch origin "+refs/heads/$base_branch:refs/remotes/origin/$base_branch" &&
git rebase "origin/$base_branch"
```
2. Re-run `run_migration.py`.

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

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

2
litellm-rust/Cargo.lock generated
View file

@ -1415,6 +1415,8 @@ dependencies = [
"litellm-config",
"litellm-core",
"reqwest",
"rustls 0.23.42",
"rustls-native-certs",
"serde",
"serde_json",
"sha2 0.10.9",

2
litellm-rust/Cargo.toml
View file

@ -28,6 +28,8 @@ pythonize = "0.29.0"
rand = "0.8"
reqwest = { version = "0.12", default-features = false, features = ["blocking", "json", "multipart", "rustls-tls", "http2", "stream"] }
rstest = "0.26.1"
rustls = { version = "0.23", default-features = false, features = ["ring", "std", "tls12"] }
rustls-native-certs = "0.8"
serde = { version = "1.0", features = ["derive"] }
serde_json = { version = "1.0", features = ["float_roundtrip"] }
sha2 = "0.10"

4
litellm-rust/crates/ai-gateway/Cargo.toml
View file

@ -20,6 +20,10 @@ litellm-config.workspace = true
# reqwest (rustls + json) is used by io/ocr and ships realtime logs to the
# Python proxy callbacks API.
reqwest.workspace = true
# rustls and its root store are direct dependencies so `io::tls` can build the
# one TLS config the outbound dials use; see that module for why it has to.
rustls.workspace = true
rustls-native-certs.workspace = true
# `sync` powers the bounded mpsc channel the realtime logger drains.
tokio = { workspace = true, features = ["rt-multi-thread", "macros", "net", "time", "sync"] }
tokio-tungstenite.workspace = true

1
litellm-rust/crates/ai-gateway/src/io/mod.rs
View file

@ -3,3 +3,4 @@ pub mod ocr;
pub mod realtime;
pub mod realtime_pool;
pub mod responses_ws;
pub(crate) mod tls;

33
litellm-rust/crates/ai-gateway/src/io/realtime.rs
View file

@ -23,10 +23,12 @@ use tokio_tungstenite::tungstenite::Message;
use tokio_tungstenite::tungstenite::client::IntoClientRequest;
use tokio_tungstenite::tungstenite::http::HeaderValue;
use tokio_tungstenite::tungstenite::http::header::AUTHORIZATION;
use tokio_tungstenite::{MaybeTlsStream, WebSocketStream, connect_async};
use tokio_tungstenite::{MaybeTlsStream, WebSocketStream};
use litellm_core::providers::openai::realtime::transformation::OPENAI_REALTIME_CONFIG;
use crate::io::tls::connect_upstream;
/// Environment variable holding the OpenAI API key (last-resort fallback).
const OPENAI_API_KEY_ENV: &str = "OPENAI_API_KEY";
@ -84,7 +86,7 @@ pub(crate) async fn dial_upstream(
.map_err(|err| Error::Auth(err.to_string()))?,
);
let (upstream, _response) = connect_async(request)
let (upstream, _response) = connect_upstream(request)
.await
.map_err(|err| Error::Network(err.to_string()))?;
Ok(upstream)
@ -284,6 +286,33 @@ mod tests {
serde_json::from_str(raw).expect("valid event json")
}
/// The realtime dial has to reach a `wss://` upstream without a process-wide
/// crypto provider installed, which is what dialing through `io::tls` buys.
#[tokio::test]
async fn dial_upstream_over_wss_reports_an_error_instead_of_panicking() {
let listener = tokio::net::TcpListener::bind("127.0.0.1:0")
.await
.expect("bind a loopback port");
let port = listener
.local_addr()
.expect("read the bound address")
.port();
tokio::spawn(async move {
while let Ok((stream, _peer)) = listener.accept().await {
drop(stream);
}
});
let result = dial_upstream(
"gpt-realtime",
"sk-test",
Some(&format!("wss://127.0.0.1:{port}")),
)
.await;
assert!(matches!(result, Err(Error::Network(_))));
}
#[test]
fn resolve_api_key_prefers_param_then_blank_falls_through() {
assert_eq!(resolve_api_key(Some("sk-test")).unwrap(), "sk-test");

35
litellm-rust/crates/ai-gateway/src/io/responses_ws.rs
View file

@ -14,7 +14,9 @@ use tokio_tungstenite::tungstenite::Message;
use tokio_tungstenite::tungstenite::client::IntoClientRequest;
use tokio_tungstenite::tungstenite::http::HeaderValue;
use tokio_tungstenite::tungstenite::http::header::{AUTHORIZATION, HeaderName};
use tokio_tungstenite::{MaybeTlsStream, WebSocketStream, connect_async};
use tokio_tungstenite::{MaybeTlsStream, WebSocketStream};
use crate::io::tls::connect_upstream;
use crate::constants::{
DEFAULT_RESPONSES_WS_CONNECT_TIMEOUT_SECS, DEFAULT_RESPONSES_WS_IDLE_TIMEOUT_SECS,
@ -49,14 +51,14 @@ impl ResponsesWebSocketConnection {
.map_err(|error| Error::InvalidRequest(error.to_string()))?;
request.headers_mut().insert(header_name, header_value);
}
let connect = connect_async(request);
let connect = connect_upstream(request);
let result = match timeout {
Some(timeout) => tokio::time::timeout(timeout, connect).await.map_err(|_| {
Error::Network("Responses WebSocket connection timed out".to_string())
})?,
None => connect.await,
};
let (socket, _) = result.map_err(|error| match error {
let (socket, _) = result.map_err(|error| match *error {
tokio_tungstenite::tungstenite::Error::Http(response) => Error::Http {
status: response.status().as_u16(),
body: String::new(),
@ -138,13 +140,13 @@ async fn dial_upstream(
);
let result = tokio::time::timeout(
Duration::from_secs(DEFAULT_RESPONSES_WS_CONNECT_TIMEOUT_SECS),
connect_async(request),
connect_upstream(request),
)
.await
.map_err(|_| Error::Network("Responses WebSocket connection timed out".to_string()))?;
result
.map(|(socket, _)| socket)
.map_err(|error| match error {
.map_err(|error| match *error {
tokio_tungstenite::tungstenite::Error::Http(response) => Error::Http {
status: response.status().as_u16(),
body: String::new(),
@ -324,6 +326,29 @@ mod tests {
use tokio::net::TcpListener;
use tokio_tungstenite::accept_async;
/// The Responses dial has to reach a `wss://` upstream without a process-wide
/// crypto provider installed, which is what dialing through `io::tls` buys.
#[tokio::test]
async fn dial_upstream_over_wss_reports_an_error_instead_of_panicking() {
let listener = TcpListener::bind("127.0.0.1:0")
.await
.expect("bind a loopback port");
let port = listener
.local_addr()
.expect("read the bound address")
.port();
tokio::spawn(async move {
while let Ok((stream, _peer)) = listener.accept().await {
drop(stream);
}
});
let result =
dial_upstream("gpt-5", "sk-test", Some(&format!("wss://127.0.0.1:{port}"))).await;
assert!(matches!(result, Err(Error::Network(_))));
}
async fn websocket_base() -> (String, tokio::task::JoinHandle<()>) {
let listener = TcpListener::bind("127.0.0.1:0").await.expect("bind");
let address = listener.local_addr().expect("local address");

80
litellm-rust/crates/ai-gateway/src/io/tls.rs Normal file
View file

@ -0,0 +1,80 @@
//! Outbound WebSocket dials over a TLS config this crate builds once and owns.
//!
//! `reqwest/rustls-tls` enables `rustls/ring` and `litellm-core`'s `bedrock-auth`
//! enables `rustls/aws-lc-rs`, so the bare `ClientConfig::builder()` that
//! `tokio-tungstenite` uses when handed no connector panics rather than guess
//! between them. Naming ring on a connector of our own settles that for these
//! dials without touching the process-wide default, and building the config
//! once keeps the platform trust store, which `tokio-tungstenite` would
//! otherwise re-read on every dial, off the dial path.
use std::io;
use std::sync::{Arc, OnceLock};
use rustls::{ClientConfig, RootCertStore};
use tokio::net::TcpStream;
use tokio_tungstenite::tungstenite::Error;
use tokio_tungstenite::tungstenite::client::IntoClientRequest;
use tokio_tungstenite::tungstenite::error::TlsError;
use tokio_tungstenite::tungstenite::handshake::client::Response;
use tokio_tungstenite::{
Connector, MaybeTlsStream, WebSocketStream, connect_async_tls_with_config,
};
static TLS_CONFIG: OnceLock<Arc<ClientConfig>> = OnceLock::new();
fn build_config() -> Result<ClientConfig, Box<Error>> {
let native = rustls_native_certs::load_native_certs();
let roots = {
let mut store = RootCertStore::empty();
let (added, _ignored) = store.add_parsable_certificates(native.certs);
if added == 0 {
return Err(Box::new(Error::Io(io::Error::other(format!(
"no usable native root certificates: {:?}",
native.errors
)))));
}
store
};
ClientConfig::builder_with_provider(Arc::new(rustls::crypto::ring::default_provider()))
.with_safe_default_protocol_versions()
.map(|builder| builder.with_root_certificates(roots).with_no_client_auth())
.map_err(|error| Box::new(Error::Tls(TlsError::Rustls(error))))
}
fn tls_config() -> Result<Arc<ClientConfig>, Box<Error>> {
if let Some(config) = TLS_CONFIG.get() {
return Ok(Arc::clone(config));
}
let built = Arc::new(build_config()?);
Ok(Arc::clone(TLS_CONFIG.get_or_init(|| built)))
}
pub(crate) async fn connect_upstream<R>(
request: R,
) -> Result<(WebSocketStream<MaybeTlsStream<TcpStream>>, Response), Box<Error>>
where
R: IntoClientRequest + Unpin,
{
let request = request.into_client_request().map_err(Box::new)?;
let connector = match request.uri().scheme_str() {
Some("wss") => Some(Connector::Rustls(tls_config()?)),
_ => None,
};
connect_async_tls_with_config(request, None, false, connector)
.await
.map_err(Box::new)
}
#[cfg(test)]
mod tests {
use super::build_config;
#[test]
fn builds_a_usable_config_with_both_provider_features_enabled() {
let config = build_config().expect("a client config");
assert!(!config.crypto_provider().cipher_suites.is_empty());
}
}

12
litellm-rust/crates/ai-gateway/src/ocr/hooks.rs
View file

@ -265,6 +265,12 @@ impl CallLifecycleHooks<PreparedOcrRequest, PreparedOcrRequest, Value> for OcrLi
Box::pin(async move { Ok(request) })
}
#[tracing::instrument(
name = "success_callback",
target = "litellm::function_trace",
level = "trace",
skip_all
)]
fn async_log_success_event<'a>(
&'a self,
context: &'a CallLifecycleContext,
@ -288,6 +294,12 @@ impl CallLifecycleHooks<PreparedOcrRequest, PreparedOcrRequest, Value> for OcrLi
})
}
#[tracing::instrument(
name = "failure_callback",
target = "litellm::function_trace",
level = "trace",
skip_all
)]
fn async_log_failure_event<'a>(
&'a self,
context: &'a CallLifecycleContext,

48
litellm-rust/crates/ai-gateway/tests/crypto_provider_wiring.rs Normal file
View file

@ -0,0 +1,48 @@
//! Guards the wiring, not just the helper: a `wss://` dial through the public
//! API has to resolve its own crypto provider, in a test binary where nothing
//! has installed a process-wide one, and has to leave it uninstalled.
use std::collections::HashMap;
use std::time::Duration;
use litellm_ai_gateway::io::responses_ws::ResponsesWebSocketConnection;
use tokio::net::TcpListener;
async fn dead_tls_server() -> u16 {
let listener = TcpListener::bind("127.0.0.1:0")
.await
.expect("bind a loopback port");
let port = listener
.local_addr()
.expect("read the bound address")
.port();
tokio::spawn(async move {
while let Ok((stream, _peer)) = listener.accept().await {
drop(stream);
}
});
port
}
#[tokio::test]
async fn dialing_wss_returns_an_error_instead_of_panicking() {
let port = dead_tls_server().await;
let result = ResponsesWebSocketConnection::connect_url(
&format!("wss://127.0.0.1:{port}/"),
&HashMap::new(),
Some(Duration::from_secs(10)),
)
.await;
assert!(
result.is_err(),
"a plain TCP server cannot finish a TLS handshake"
);
assert!(
rustls::crypto::CryptoProvider::get_default().is_none(),
"the dial settles its provider on its own connector, not process-wide"
);
}

52
litellm-rust/crates/ai-gateway/tests/ocr_lifecycle.rs
View file

@ -11,9 +11,13 @@ use litellm_ai_gateway::integrations::custom_logger::{
use litellm_ai_gateway::integrations::types::RequestMetadata;
use litellm_ai_gateway::ocr::{OcrRequest, ocr};
use litellm_core::error::Error;
#[cfg(feature = "trace-parity")]
use litellm_core::observability::FunctionTrace;
use serde_json::{Map, Value, json};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::{TcpListener, TcpStream};
#[cfg(feature = "trace-parity")]
use tracing::instrument::WithSubscriber;
async fn read_http_headers(socket: &mut TcpStream) -> String {
let mut request = Vec::new();
@ -320,14 +324,17 @@ async fn ocr_lifecycle_runs_pre_during_and_success_hooks() {
GuardrailEventHook::PreCall,
GuardrailEventHook::DuringCall,
]));
let response = ocr(OcrRequest {
#[cfg(feature = "trace-parity")]
let trace = FunctionTrace::default();
let api_base = format!("http://{addr}");
let call = ocr(OcrRequest {
model: "mistral-ocr-latest",
document: json!({
"type": "document_url",
"document_url": "https://example.com/doc.pdf"
}),
api_key: Some("sk-test"),
api_base: Some(&format!("http://{addr}")),
api_base: Some(&api_base),
custom_llm_provider: Some("mistral"),
extra_headers: None,
optional_params: Map::new(),
@ -339,9 +346,10 @@ async fn ocr_lifecycle_runs_pre_during_and_success_hooks() {
..Default::default()
},
litellm_call_id: Some("ocr-call-1"),
})
.await
.expect("ocr request succeeds");
});
#[cfg(feature = "trace-parity")]
let call = call.with_subscriber(trace.dispatcher());
let response = call.await.expect("ocr request succeeds");
assert_eq!(response["pages"][0]["markdown"], "ok");
assert_eq!(
@ -359,6 +367,16 @@ async fn ocr_lifecycle_runs_pre_during_and_success_hooks() {
error_kind: None,
}]
);
#[cfg(feature = "trace-parity")]
assert_eq!(
trace
.events()
.iter()
.filter(|event| event.function.ends_with("_callback"))
.map(|event| event.function)
.collect::<Vec<_>>(),
vec!["success_callback"]
);
let request = server.await.expect("server task completes");
assert!(request.contains(r#""guarded_pre":true"#), "{request}");
@ -388,14 +406,17 @@ async fn ocr_lifecycle_runs_failure_hook_on_provider_error() {
});
let logger = Arc::new(RecordingOcrLogger::default());
let err = ocr(OcrRequest {
#[cfg(feature = "trace-parity")]
let trace = FunctionTrace::default();
let api_base = format!("http://{addr}");
let call = ocr(OcrRequest {
model: "mistral-ocr-latest",
document: json!({
"type": "document_url",
"document_url": "https://example.com/doc.pdf"
}),
api_key: Some("sk-test"),
api_base: Some(&format!("http://{addr}")),
api_base: Some(&api_base),
custom_llm_provider: Some("mistral"),
extra_headers: None,
optional_params: Map::new(),
@ -404,9 +425,10 @@ async fn ocr_lifecycle_runs_failure_hook_on_provider_error() {
guardrails: Vec::new(),
request_metadata: RequestMetadata::default(),
litellm_call_id: Some("ocr-call-2"),
})
.await
.expect_err("provider error propagates");
});
#[cfg(feature = "trace-parity")]
let call = call.with_subscriber(trace.dispatcher());
let err = call.await.expect_err("provider error propagates");
assert!(matches!(err, Error::Http { status: 500, .. }));
server.await.expect("server task completes");
@ -421,6 +443,16 @@ async fn ocr_lifecycle_runs_failure_hook_on_provider_error() {
error_kind: Some("HttpError".to_string()),
}]
);
#[cfg(feature = "trace-parity")]
assert_eq!(
trace
.events()
.iter()
.filter(|event| event.function.ends_with("_callback"))
.map(|event| event.function)
.collect::<Vec<_>>(),
vec!["failure_callback"]
);
}
#[tokio::test]

28
litellm-rust/crates/python-bridge/src/function_trace.rs
View file

@ -1,3 +1,4 @@
use std::fmt::Display;
use std::future::Future;
use litellm_core::observability::{FunctionTrace, FunctionTraceEvent};
@ -6,17 +7,32 @@ use tracing::instrument::WithSubscriber;
#[derive(Serialize)]
pub(crate) struct TracedResponse<T> {
response: T,
#[serde(skip_serializing_if = "Option::is_none")]
response: Option<T>,
#[serde(skip_serializing_if = "Option::is_none")]
error: Option<String>,
trace: Vec<FunctionTraceEvent>,
}
pub(crate) async fn capture<T, E>(
future: impl Future<Output = Result<T, E>>,
) -> Result<TracedResponse<T>, E> {
) -> Result<TracedResponse<T>, E>
where
E: Display,
{
let trace = FunctionTrace::default();
let response = future.with_subscriber(trace.dispatcher()).await?;
Ok(TracedResponse {
response,
trace: trace.events(),
let result = future.with_subscriber(trace.dispatcher()).await;
let events = trace.events();
Ok(match result {
Ok(response) => TracedResponse {
response: Some(response),
error: None,
trace: events,
},
Err(error) => TracedResponse {
response: None,
error: Some(error.to_string()),
trace: events,
},
})
}

9
litellm-rust/crates/python-bridge/src/routes/definition.rs
View file

@ -486,10 +486,11 @@ asyncio.run(exercise())
let code = CString::new(
r#"
result = routes.echo("traced")
assert result == {
"response": "traced",
"trace": [{"function": "execute_echo", "depth": 0}],
}
assert result["response"] == "traced", result
assert [event["function"] for event in result["trace"]] == ["execute_echo"], result
failure = routes.echo("error")
assert failure["error"] == "invalid request: synthetic error", failure
assert [event["function"] for event in failure["trace"]] == ["execute_echo"], failure
"#,
)
.expect("Python source should not contain null bytes");

15
litellm/__init__.py
View file

@ -47,6 +47,7 @@ from typing import (
)
from litellm.types.integrations.datadog import DatadogInitParams
from litellm.types.integrations.newrelic import NewRelicInitParams
from litellm.litellm_core_utils.core_helpers import drop_params_env_flag
from litellm._logging import (
set_verbose,
_turn_on_debug,
@ -238,7 +239,7 @@ token: Optional[str] = (
)
telemetry = True
max_tokens: int = DEFAULT_MAX_TOKENS # OpenAI Defaults
drop_params = bool(os.getenv("LITELLM_DROP_PARAMS", False))
drop_params = drop_params_env_flag(os.environ, verbose_logger)
modify_params = bool(os.getenv("LITELLM_MODIFY_PARAMS", False))
use_chat_completions_url_for_anthropic_messages: bool = bool(
os.getenv("LITELLM_USE_CHAT_COMPLETIONS_URL_FOR_ANTHROPIC_MESSAGES", False)
@ -325,6 +326,9 @@ ssl_certificate: Optional[str] = None
user_url_validation: bool = True
user_url_allowed_hosts: List[str] = []
provider_url_destination_allowed_hosts: List[str] = []
#: "override" (default) or "additive": whether a key or team destination replaces
#: the operator's exporter for that backend or exports alongside it.
otel_tenant_destination_mode: str | None = None
ssl_ecdh_curve: Optional[str] = None # Set to 'X25519' to disable PQC and improve performance
disable_streaming_logging: bool = False
disable_token_counter: bool = False
@ -471,6 +475,7 @@ prometheus_metrics_config: Optional[List] = None
prometheus_exclude_metrics: Optional[List[str]] = None
prometheus_exclude_labels: Optional[List[str]] = None
prometheus_emit_stream_label: bool = False
prometheus_emit_input_sequence_length_label: bool = False
prometheus_deployment_and_latency_caller_identity: Literal[
"api_key_alias",
"user_email",
@ -495,6 +500,7 @@ public_model_groups: Optional[List[str]] = None
public_agent_groups: Optional[List[str]] = None
agent_search_embedding_model: Optional[str] = None
mcp_tool_search: Optional[Mapping[str, object]] = None
skill_search_embedding_model: Optional[str] = None
# Supports both old format (Dict[str, str]) and new format (Dict[str, Dict[str, Any]])
# New format: { "displayName": { "url": "...", "index": 0 } }
# Old format: { "displayName": "url" } (for backward compatibility)
@ -541,7 +547,7 @@ _key_management_system: Optional["KeyManagementSystem"] = None
#### PII MASKING ####
output_parse_pii: bool = False
#############################################
from litellm.litellm_core_utils.get_model_cost_map import get_model_cost_map
from litellm.litellm_core_utils.get_model_cost_map import get_model_cost_map, mark_litellm_import_complete
model_cost = get_model_cost_map(url=model_cost_map_url)
cost_discount_config: Dict[str, float] = {} # Provider-specific cost discounts {"vertex_ai": 0.05} = 5% discount
@ -2001,6 +2007,9 @@ if TYPE_CHECKING:
from .llms.hosted_vllm.responses.transformation import (
HostedVLLMResponsesAPIConfig as HostedVLLMResponsesAPIConfig,
)
from .llms.fireworks_ai.responses.transformation import (
FireworksAIResponsesAPIConfig as FireworksAIResponsesAPIConfig,
)
from .llms.github_copilot.chat.transformation import (
GithubCopilotConfig as GithubCopilotConfig,
)
@ -2397,3 +2406,5 @@ def __getattr__(name: str) -> Any:
# ALL_LITELLM_RESPONSE_TYPES is lazy-loaded via __getattr__ to avoid loading utils at import time
mark_litellm_import_complete()

24
litellm/_internal_context.py
View file

@ -6,9 +6,33 @@ be settable from user input. Context variables are scoped to the current
asyncio task and cannot be injected via HTTP request bodies.
"""
from collections.abc import Generator
from contextlib import contextmanager
from contextvars import ContextVar
from datetime import datetime, timezone
from typing import Final
# When True, suppresses async logging and billing for internal sub-calls
# (e.g., emulated file-search steps that make nested LLM calls).
is_internal_call: Final[ContextVar[bool]] = ContextVar("is_internal_call", default=False)
# One request prices its totals, its per-token-type lines and the rates it reports on
# separate code paths. Each reads the clock for off-peak pricing, so without a pinned
# moment they can land on either side of a window boundary and disagree with each other.
_billing_time: Final[ContextVar[datetime | None]] = ContextVar("billing_time", default=None)
@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."""
token: Final = _billing_time.set(moment)
try:
yield
finally:
_billing_time.reset(token)
def current_billing_time() -> datetime:
"""The pinned billing moment, or now in UTC outside a pinned block."""
pinned: Final = _billing_time.get()
return pinned if pinned is not None else datetime.now(timezone.utc)

2
litellm/_lazy_imports.py
View file

@ -229,7 +229,7 @@ def _module_attribute(module: ModuleType, attr_name: str) -> object:
return attribute["value"]
def _generic_lazy_import(name: str, import_map: dict[str, tuple[str, str]], category: str) -> object:
def _generic_lazy_import(name: str, import_map: Mapping[str, tuple[str, str]], category: str) -> object:
"""
Generic function that handles lazy importing for most attributes.

5
litellm/_lazy_imports_registry.py
View file

@ -237,6 +237,7 @@ LLM_CONFIG_NAMES: Final = (
"XAIResponsesAPIConfig",
"LiteLLMProxyResponsesAPIConfig",
"HostedVLLMResponsesAPIConfig",
"FireworksAIResponsesAPIConfig",
"VolcEngineResponsesAPIConfig",
"PerplexityResponsesConfig",
"DatabricksResponsesAPIConfig",
@ -957,6 +958,10 @@ _LLM_CONFIGS_IMPORT_MAP: Final = {
".llms.hosted_vllm.responses.transformation",
"HostedVLLMResponsesAPIConfig",
),
"FireworksAIResponsesAPIConfig": (
".llms.fireworks_ai.responses.transformation",
"FireworksAIResponsesAPIConfig",
),
"VolcEngineResponsesAPIConfig": (
".llms.volcengine.responses.transformation",
"VolcEngineResponsesAPIConfig",

17
litellm/_redis_credential_provider.py
View file

@ -1,7 +1,7 @@
import asyncio
import threading
import time
from typing import Any, Final
from typing import Final, Protocol
from redis.credentials import CredentialProvider
@ -18,6 +18,19 @@ _token_cache: Final[dict[str, tuple[str, float]]] = {}
_token_cache_lock: Final = threading.Lock()
class AzureAccessToken(Protocol):
"""The ``azure.core.credentials.AccessToken`` shape this module reads."""
@property
def token(self) -> str: ...
class AzureCredential(Protocol):
"""The ``azure-identity`` credential surface this module calls."""
def get_token(self, *scopes: str) -> AzureAccessToken: ...
def _generate_gcp_iam_access_token(service_account: str) -> str:
"""
Generate GCP IAM access token for Redis authentication.
@ -115,7 +128,7 @@ class AzureADCredentialProvider(CredentialProvider):
fail authentication after the initial token expired (~1 hour TTL).
"""
def __init__(self, credential: Any, username: str | None = None) -> None:
def __init__(self, credential: AzureCredential, username: str | None = None) -> None:
self._credential = credential
self._username = username

46
litellm/_service_logger.py
View file

@ -1,7 +1,7 @@
import asyncio
from collections.abc import Callable, Coroutine
from datetime import datetime, timedelta
from typing import TYPE_CHECKING, Any, Final
from typing import TYPE_CHECKING, Any, Final, Protocol
import litellm
from litellm._logging import verbose_logger
@ -25,7 +25,30 @@ else:
UserAPIKeyAuth = Any
def _get_otel_v2_class() -> type | None:
class _ServiceSpanLogger(Protocol):
"""The OTel logger surface this module drives: the two service-span hooks it calls."""
async def async_service_success_hook(
self,
payload: ServiceLoggerPayload,
parent_otel_span: Span | None = None,
start_time: datetime | float | None = None,
end_time: datetime | float | None = None,
event_metadata: dict | None = None,
) -> None: ...
async def async_service_failure_hook(
self,
payload: ServiceLoggerPayload,
error: str | None = "",
parent_otel_span: Span | None = None,
start_time: datetime | float | None = None,
end_time: datetime | float | None = None,
event_metadata: dict | None = None,
) -> None: ...
def _get_otel_v2_class() -> type[_ServiceSpanLogger] | None:
"""Return the ``OpenTelemetryV2`` class, or ``None`` if the OTel SDK is absent.
Imported lazily: ``litellm.integrations.otel.logger`` imports the OpenTelemetry
@ -55,7 +78,7 @@ class ServiceLogging(CustomLogger):
if "prometheus_system" in litellm.service_callback:
self.prometheusServicesLogger = PrometheusServicesLogger()
def _resolve_otel_service_logger(self, callback: Any) -> Any | None:
def _resolve_otel_service_logger(self, callback: object) -> _ServiceSpanLogger | None:
"""Resolve the OTel logger (legacy or V2) to emit a service span on.
Returns the logger instance whose ``async_service_*_hook`` should fire for
@ -70,18 +93,21 @@ class ServiceLogging(CustomLogger):
"""
otel_v2_cls: Final = _get_otel_v2_class()
def _is_otel_logger(obj: Any) -> bool:
def _as_otel_logger(obj: object) -> _ServiceSpanLogger | None:
if isinstance(obj, OpenTelemetry):
return True
return otel_v2_cls is not None and isinstance(obj, otel_v2_cls)
return obj
if otel_v2_cls is not None and isinstance(obj, otel_v2_cls):
return obj
return None
if _is_otel_logger(callback):
return callback
resolved_callback: Final = _as_otel_logger(callback)
if resolved_callback is not None:
return resolved_callback
if callback == "otel":
from litellm.proxy.proxy_server import open_telemetry_logger
if open_telemetry_logger is not None and _is_otel_logger(open_telemetry_logger):
return open_telemetry_logger
if open_telemetry_logger is not None:
return _as_otel_logger(open_telemetry_logger)
return None
@staticmethod

3
litellm/a2a_protocol/card_resolver.py
View file

@ -4,6 +4,7 @@ Custom A2A Card Resolver for LiteLLM.
Extends the A2A SDK's card resolver to support multiple well-known paths.
"""
from collections.abc import Mapping
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final
@ -152,7 +153,7 @@ class LiteLLMA2ACardResolver(_A2ACardResolver):
async def get_agent_card(
self,
relative_card_path: str | None = None,
http_kwargs: dict[str, Any] | None = None,
http_kwargs: Mapping[str, object] | None = None,
) -> "AgentCard":
"""
Fetch the agent card, trying multiple well-known paths.

19
litellm/a2a_protocol/providers/bedrock_agentcore/handler.py
View file

@ -6,13 +6,14 @@ completion bridge that would otherwise strip the envelope.
"""
import json
from collections.abc import AsyncIterator
from typing import Any, Final, cast
from collections.abc import AsyncIterator, Mapping
from typing import Any, Final
from litellm._logging import verbose_logger
from litellm.a2a_protocol.providers.bedrock_agentcore.transformation import (
BedrockAgentCoreA2ATransformation,
)
from litellm.llms.bedrock.base_aws_llm import run_aws_signing
from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
from litellm.types.llms.custom_http import httpxSpecialProvider
@ -28,7 +29,7 @@ class BedrockAgentCoreA2AHandler:
@staticmethod
async def handle_non_streaming(
request_id: str,
params: dict[str, Any],
params: Mapping[str, object],
litellm_params: dict[str, Any],
agent_extra_headers: dict[str, str] | None = None,
) -> dict[str, Any]:
@ -45,7 +46,8 @@ class BedrockAgentCoreA2AHandler:
Returns:
A2A JSON-RPC response dict from the AgentCore agent
"""
url, headers, body = BedrockAgentCoreA2ATransformation.get_url_and_signed_request(
url, headers, body = await run_aws_signing(
BedrockAgentCoreA2ATransformation.get_url_and_signed_request,
request_id=request_id,
params=params,
litellm_params=litellm_params,
@ -56,7 +58,7 @@ class BedrockAgentCoreA2AHandler:
verbose_logger.info("BedrockAgentCore A2A: Sending non-streaming request to %s", url)
client: Final = get_async_httpx_client(
llm_provider=cast(Any, httpxSpecialProvider.A2AProvider),
llm_provider=httpxSpecialProvider.A2AProvider,
)
response: Final = await client.post(
url,
@ -74,7 +76,7 @@ class BedrockAgentCoreA2AHandler:
@staticmethod
async def handle_streaming(
request_id: str,
params: dict[str, Any],
params: Mapping[str, object],
litellm_params: dict[str, Any],
agent_extra_headers: dict[str, str] | None = None,
) -> AsyncIterator[dict[str, Any]]:
@ -91,7 +93,8 @@ class BedrockAgentCoreA2AHandler:
Yields:
A2A streaming response events from the AgentCore agent
"""
url, headers, body = BedrockAgentCoreA2ATransformation.get_url_and_signed_request(
url, headers, body = await run_aws_signing(
BedrockAgentCoreA2ATransformation.get_url_and_signed_request,
request_id=request_id,
params=params,
litellm_params=litellm_params,
@ -103,7 +106,7 @@ class BedrockAgentCoreA2AHandler:
verbose_logger.info("BedrockAgentCore A2A: Sending streaming request to %s", url)
client: Final = get_async_httpx_client(
llm_provider=cast(Any, httpxSpecialProvider.A2AProvider),
llm_provider=httpxSpecialProvider.A2AProvider,
)
response: Final = await client.post(
url,

12
litellm/a2a_protocol/providers/bedrock_agentcore/transformation.py
View file

@ -7,7 +7,7 @@ and signs requests via AmazonAgentCoreConfig (SigV4 or JWT).
import json
from collections.abc import AsyncIterator, Mapping
from typing import Any, Final
from typing import Any, Final, Protocol
from litellm._logging import verbose_logger
from litellm.a2a_protocol.litellm_completion_bridge.handler import (
@ -47,6 +47,12 @@ _RESERVED_PREFIX_HEADERS: Final[tuple[str, ...]] = (
)
class _SSELineSource(Protocol):
"""Minimal streaming-response surface used to read SSE lines."""
def aiter_lines(self) -> AsyncIterator[str]: ...
def _filter_reserved_headers(
agent_extra_headers: Mapping[str, str] | None,
) -> dict[str, str] | None:
@ -114,7 +120,7 @@ class BedrockAgentCoreA2ATransformation:
@staticmethod
def get_url_and_signed_request(
request_id: str,
params: dict[str, Any],
params: Mapping[str, object],
litellm_params: dict[str, Any],
method: str = "message/send",
stream: bool = False,
@ -213,7 +219,7 @@ class BedrockAgentCoreA2ATransformation:
return url, signed_headers, signed_body
@staticmethod
async def parse_sse_events(response: Any) -> AsyncIterator[dict[str, Any]]:
async def parse_sse_events(response: _SSELineSource) -> AsyncIterator[dict[str, Any]]:
"""
Parse SSE events from an httpx streaming response.

14
litellm/a2a_protocol/providers/watsonx_orchestrate/transformation.py
View file

@ -8,7 +8,7 @@ WXO uses a REST API (not A2A/JSON-RPC) with an async-poll execution model:
"""
import asyncio
from collections.abc import AsyncIterator
from collections.abc import AsyncIterator, Mapping
from typing import Any, Final
from uuid import uuid4
@ -51,9 +51,9 @@ class WatsonxOrchestrateTransformation:
wxo_agent_id: str,
text: str,
thread_id: str | None = None,
) -> dict[str, Any]:
) -> dict[str, object]:
"""Build the WXO POST /v1/orchestrate/runs request body."""
body: Final[dict[str, Any]] = {
body: Final[dict[str, object]] = {
"agent_id": wxo_agent_id,
"message": {
"role": "user",
@ -70,7 +70,7 @@ class WatsonxOrchestrateTransformation:
return body
@staticmethod
def extract_text_from_wxo_result(result: Any) -> str:
def extract_text_from_wxo_result(result: object) -> str:
"""
Extract response text from a WXO run result.
@ -103,7 +103,7 @@ class WatsonxOrchestrateTransformation:
return ""
@staticmethod
def extract_text_from_a2a_message_response(a2a_response: dict[str, Any]) -> str:
def extract_text_from_a2a_message_response(a2a_response: Mapping[str, object]) -> str:
result: Final = a2a_response.get("result")
if not isinstance(result, dict):
verbose_logger.warning("WXO: A2A response missing result object")
@ -119,7 +119,7 @@ class WatsonxOrchestrateTransformation:
return ""
@staticmethod
def build_a2a_message_response(request_id: str, text: str) -> dict[str, Any]:
def build_a2a_message_response(request_id: str, text: str) -> dict[str, object]:
"""
Build a standard A2A non-streaming SendMessageResponse (kind=message).
"""
@ -140,7 +140,7 @@ class WatsonxOrchestrateTransformation:
request_id: str,
chunk_size: int = 50,
delay_ms: int = 10,
) -> AsyncIterator[dict[str, Any]]:
) -> AsyncIterator[dict[str, object]]:
"""
Emit standard A2A streaming events from a completed text response.

4
litellm/a2a_protocol/streaming_iterator.py
View file

@ -148,9 +148,9 @@ class A2AStreamingIterator:
except Exception as e:
verbose_logger.debug("Error in A2A streaming completion handler: %s", e)
def _build_logging_result(self, usage: litellm.Usage) -> dict[str, Any]:
def _build_logging_result(self, usage: litellm.Usage) -> dict[str, object]:
"""Build a result dict for logging."""
result: Final[dict[str, Any]] = {
result: Final[dict[str, object]] = {
"id": getattr(self.request, "id", "unknown"),
"jsonrpc": "2.0",
"usage": (usage.model_dump() if hasattr(usage, "model_dump") else dict(usage)),

6
litellm/a2a_protocol/utils.py
View file

@ -48,7 +48,7 @@ class A2ARequestUtils:
return " ".join(text_parts)
@staticmethod
def extract_text_from_response(response_dict: dict[str, Any]) -> str:
def extract_text_from_response(response_dict: Mapping[str, object]) -> str:
"""
Extract text content from A2A response result.
@ -111,7 +111,7 @@ class A2ARequestUtils:
@staticmethod
def calculate_usage_from_request_response(
request: "SendMessageRequest | SendStreamingMessageRequest",
response_dict: dict[str, Any],
response_dict: Mapping[str, object],
) -> tuple[int, int, int]:
"""
Calculate token usage from A2A request and response.
@ -170,5 +170,5 @@ def extract_text_from_a2a_message(message: Any) -> str:
return A2ARequestUtils.extract_text_from_message(message)
def extract_text_from_a2a_response(response_dict: dict[str, Any]) -> str:
def extract_text_from_a2a_response(response_dict: Mapping[str, object]) -> str:
return A2ARequestUtils.extract_text_from_response(response_dict)

45
litellm/assistants/utils.py
View file

@ -1,3 +1,4 @@
from collections.abc import Mapping, Sequence
from typing import Final
import litellm
@ -10,20 +11,22 @@ def get_optional_params_add_message(
role: str | None,
content: str | list[MessageContentTextObject | MessageContentImageFileObject | MessageContentImageURLObject] | None,
attachments: list[Attachment] | None,
metadata: dict | None,
metadata: Mapping[str, object] | None,
custom_llm_provider: str,
**kwargs,
):
**kwargs: object,
) -> dict[str, object]:
"""
Azure doesn't support 'attachments' for creating a message
Reference - https://learn.microsoft.com/en-us/azure/ai-services/openai/assistants-reference-messages?tabs=python#create-message
"""
passed_params: Final = locals()
custom_llm_provider = passed_params.pop("custom_llm_provider")
special_params: Final = passed_params.pop("kwargs")
for k, v in special_params.items():
passed_params[k] = v
passed_params: Final[Mapping[str, object]] = {
"role": role,
"content": content,
"attachments": attachments,
"metadata": metadata,
**kwargs,
}
default_params: Final = {
"role": None,
@ -33,10 +36,10 @@ def get_optional_params_add_message(
}
non_default_params = {k: v for k, v in passed_params.items() if (k in default_params and v != default_params[k])}
optional_params = {}
optional_params: dict[str, object] = {}
## raise exception if non-default value passed for non-openai/azure embedding calls
def _check_valid_arg(supported_params):
def _check_valid_arg(supported_params: Sequence[str]) -> Mapping[str, object] | None:
if len(non_default_params.keys()) > 0:
keys: Final = list(non_default_params.keys())
for k in keys:
@ -71,14 +74,18 @@ def get_optional_params_image_gen(
style: str | None = None,
user: str | None = None,
custom_llm_provider: str | None = None,
**kwargs,
):
**kwargs: object,
) -> dict[str, object]:
# retrieve all parameters passed to the function
passed_params: Final = locals()
custom_llm_provider = passed_params.pop("custom_llm_provider")
special_params: Final = passed_params.pop("kwargs")
for k, v in special_params.items():
passed_params[k] = v
passed_params: Final[Mapping[str, object]] = {
"n": n,
"quality": quality,
"response_format": response_format,
"size": size,
"style": style,
"user": user,
**kwargs,
}
default_params: Final = {
"n": None,
@ -90,10 +97,10 @@ def get_optional_params_image_gen(
}
non_default_params = {k: v for k, v in passed_params.items() if (k in default_params and v != default_params[k])}
optional_params = {}
optional_params: dict[str, object] = {}
## raise exception if non-default value passed for non-openai/azure embedding calls
def _check_valid_arg(supported_params):
def _check_valid_arg(supported_params: Sequence[str]) -> Mapping[str, object] | None:
if len(non_default_params.keys()) > 0:
keys: Final = list(non_default_params.keys())
for k in keys:

83
litellm/batches/batch_utils.py
View file

@ -10,7 +10,7 @@ from litellm._logging import verbose_logger
from litellm.litellm_core_utils.get_litellm_params import AWS_CREDENTIAL_KWARGS_KEYS
from litellm.litellm_core_utils.llm_cost_calc.utils import parse_prompt_tokens_details
from litellm.types.llms.openai import Batch
from litellm.types.utils import CallTypes, ModelInfo, Usage
from litellm.types.utils import ModelInfo, Usage
from litellm.utils import token_counter
@ -23,6 +23,29 @@ class BatchCostUsageResult:
models: list[str]
successful_requests: int
failed_requests: int
prompt_cost: float = 0.0
completion_cost: float = 0.0
_COMPLETED_BATCH_STATUSES: Final = frozenset({"completed", "complete"})
_TERMINAL_BATCH_STATUSES: Final = _COMPLETED_BATCH_STATUSES | frozenset({"failed", "cancelled", "expired"})
def batch_cost_is_final(batch: Batch) -> bool:
"""Whether this retrieve of the batch is the one to account its cost from.
A batch still in flight has nothing to price, and a "completed" batch can report
no output_file_id for a moment before the output populates; pricing either records
$0 under the batch's single spend row and pins it there. Final means a completed
batch whose output file has arrived or whose counts prove no line succeeded, or
any other terminal status (failed, cancelled, expired).
"""
if batch.status not in _TERMINAL_BATCH_STATUSES:
return False
if batch.status not in _COMPLETED_BATCH_STATUSES or batch.output_file_id is not None:
return True
request_counts: Final = batch.request_counts
return request_counts is not None and request_counts.total > 0 and request_counts.completed == 0
async def calculate_batch_cost_and_usage(
@ -130,7 +153,8 @@ class _LineOutcome(Enum):
@dataclass(frozen=True, slots=True)
class _BatchOutputLineStats:
cost: float
prompt_cost: float
completion_cost: float
prompt_tokens: int
completion_tokens: int
total_tokens: int
@ -160,7 +184,7 @@ def _classify_output_line_stats(
def _safe_output_line_stats(
entry: Mapping[str, Any],
entry: Mapping[str, object],
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock"],
model_name: str | None,
model_info: ModelInfo | None,
@ -182,7 +206,7 @@ def _safe_output_line_stats(
def _compute_output_line_stats(
entry: Mapping[str, Any],
entry: Mapping[str, object],
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock"],
model_name: str | None,
model_info: ModelInfo | None,
@ -193,15 +217,16 @@ def _compute_output_line_stats(
raw_model: Final = response_body.get("model")
response_model: Final = raw_model if isinstance(raw_model, str) and raw_model else None
completion_details: Final = usage.completion_tokens_details
line_prompt_cost, line_completion_cost = _output_line_cost(
usage=usage,
custom_llm_provider=custom_llm_provider,
model_name=model_name,
response_model=response_model,
model_info=model_info,
)
return _BatchOutputLineStats(
cost=_output_line_cost(
response_body=response_body,
usage=usage,
custom_llm_provider=custom_llm_provider,
model_name=model_name,
response_model=response_model,
model_info=model_info,
),
prompt_cost=line_prompt_cost,
completion_cost=line_completion_cost,
prompt_tokens=usage.prompt_tokens,
completion_tokens=usage.completion_tokens,
total_tokens=usage.total_tokens,
@ -213,31 +238,24 @@ def _compute_output_line_stats(
def _output_line_cost(
response_body: Mapping[str, Any],
usage: Usage,
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock"],
model_name: str | None,
response_model: str | None,
model_info: ModelInfo | None,
) -> float:
) -> tuple[float, float]:
"""(prompt_cost, completion_cost) for one output line, priced at batch rates."""
from litellm.cost_calculator import batch_cost_calculator
if model_info is None and custom_llm_provider not in ("anthropic", "bedrock"):
return litellm.completion_cost(
completion_response=response_body,
custom_llm_provider=custom_llm_provider,
call_type=CallTypes.aretrieve_batch.value,
)
cost_model: Final = (
model_name if custom_llm_provider == "bedrock" and model_name else response_model or model_name or ""
)
prompt_cost, completion_cost = batch_cost_calculator(
return batch_cost_calculator(
usage=usage,
model=cost_model,
custom_llm_provider=custom_llm_provider,
model_info=model_info,
)
return prompt_cost + completion_cost
def _aggregate_batch_cost_usage_models(
@ -270,7 +288,9 @@ def _aggregate_batch_cost_usage_models(
**cache_token_params,
)
batch_models: Final = [model_name] if model_name else [stats.model for stats in line_stats if stats.model]
total_cost: Final = sum((stats.cost for stats in line_stats), 0.0)
total_prompt_cost: Final = sum((stats.prompt_cost for stats in line_stats), 0.0)
total_completion_cost: Final = sum((stats.completion_cost for stats in line_stats), 0.0)
total_cost: Final = total_prompt_cost + total_completion_cost
verbose_logger.debug(
"batch output aggregate: cost=%s usage=%s models=%s successful=%d failed=%d",
total_cost,
@ -285,6 +305,8 @@ def _aggregate_batch_cost_usage_models(
models=batch_models,
successful_requests=successful_requests,
failed_requests=failed_requests,
prompt_cost=total_prompt_cost,
completion_cost=total_completion_cost,
)
@ -309,7 +331,8 @@ def calculate_vertex_ai_batch_cost_and_usage(
"""
from litellm.cost_calculator import batch_cost_calculator
total_cost = 0.0
total_prompt_cost = 0.0 # rebind-ok: loop accumulator, matches total_tokens below
total_completion_cost = 0.0 # rebind-ok: loop accumulator, matches total_tokens below
total_tokens = 0
prompt_tokens = 0
completion_tokens = 0
@ -341,7 +364,8 @@ def calculate_vertex_ai_batch_cost_and_usage(
model=actual_model_name,
custom_llm_provider="vertex_ai",
)
total_cost += p_cost + c_cost
total_prompt_cost += p_cost
total_completion_cost += c_cost
except Exception as e:
verbose_logger.debug("vertex_ai batch cost calculation error for line: %s", str(e))
@ -349,6 +373,7 @@ def calculate_vertex_ai_batch_cost_and_usage(
completion_tokens += _completion
total_tokens += _total
total_cost: Final = total_prompt_cost + total_completion_cost
verbose_logger.info(
"vertex_ai batch cost: cost=%s, prompt=%d, completion=%d, total=%d, successful=%d, failed=%d",
total_cost,
@ -369,6 +394,8 @@ def calculate_vertex_ai_batch_cost_and_usage(
models=[actual_model_name],
successful_requests=successful_requests,
failed_requests=failed_requests,
prompt_cost=total_prompt_cost,
completion_cost=total_completion_cost,
)
@ -556,7 +583,7 @@ def _iter_batch_output_entries(file_content: bytes) -> Iterator[dict]:
def _parse_batch_output_line(line: bytes) -> dict | None:
try:
parsed: Final = json.loads(line)
parsed: Final[object] = json.loads(line)
except ValueError as e:
verbose_logger.warning("skipping malformed batch output line: %s", str(e))
return None
@ -601,7 +628,7 @@ def _count_entry_tokens(
return 0
def _count_prompt_or_input_tokens(model: str, value: Any) -> int:
def _count_prompt_or_input_tokens(model: str, value: object) -> int:
"""Token-count a ``prompt`` / ``input`` field that the OpenAI batch
schema allows in four shapes:
@ -680,7 +707,7 @@ def _get_anthropic_result_from_batch_results_line(batch_results_line: Mapping[st
def _get_response_from_batch_job_output_file(
batch_job_output_file: Mapping[str, Any], custom_llm_provider: str = "openai"
) -> Mapping[str, Any]:
) -> Mapping[str, object]:
"""
Get the response from the batch job output file
"""

1
litellm/batches/main.py
View file

@ -319,6 +319,7 @@ def create_batch(
timeout=timeout,
max_retries=optional_params.max_retries,
create_batch_data=_create_batch_request,
custom_endpoint=optional_params.get("custom_endpoint"),
)
else:
raise litellm.exceptions.BadRequestError(

4
litellm/caching/caching_handler.py
View file

@ -672,7 +672,7 @@ class LLMCachingHandler:
def _async_log_cache_hit_on_callbacks(
self,
logging_obj: LiteLLMLoggingObj,
cached_result: Any,
cached_result: object,
start_time: datetime.datetime,
end_time: datetime.datetime,
cache_hit: bool,
@ -1184,7 +1184,7 @@ class LLMCachingHandler:
logging_obj: LiteLLMLoggingObj,
model: str,
kwargs: dict[str, Any],
cached_result: Any,
cached_result: object,
is_async: bool,
is_embedding: bool = False,
custom_llm_provider: str | None = None,

2
litellm/caching/dual_cache.py
View file

@ -257,7 +257,7 @@ class DualCache(BaseCache):
self,
current_time: float,
keys: list[str],
result: Sequence[Any],
result: Sequence[object],
) -> tuple[list[str], dict[str, float | None]]:
"""
Atomically choose keys to fetch from Redis and reserve their access time.

65
litellm/caching/redis_cache.py
View file

@ -20,6 +20,8 @@ from contextvars import ContextVar
from datetime import timedelta
from typing import TYPE_CHECKING, Any, Final, Protocol, TypeVar, cast
from pydantic import TypeAdapter
import litellm
from litellm._logging import print_verbose, verbose_logger
from litellm.constants import (
@ -80,11 +82,29 @@ class _AsyncRedisCommands(Protocol):
def pipeline(self, transaction: bool = True) -> "Pipeline[bytes]": ...
def eval(self, script: str, numkeys: int, *keys_and_args: str | bytes | float) -> Awaitable[object]: ...
_BREAKER_GUARD_FRAME_NAMES: Final = frozenset(
{"<lambda>", "wrapper", "_run_under_circuit_breaker", "_run_under_circuit_breaker_sync"}
)
_INCREMENT_WITH_FLOOR_LUA: Final = (
"local count = redis.call('INCRBY', KEYS[1], ARGV[1]) "
"if count < 0 then count = redis.call('INCRBY', KEYS[1], -count) end "
"if redis.call('TTL', KEYS[1]) < 0 then redis.call('EXPIRE', KEYS[1], ARGV[2]) end "
"return count"
)
_LUA_COUNT: Final = TypeAdapter(int)
_OPTIONAL_COUNTS: Final = TypeAdapter(tuple[int | None, ...])
def _decoded_counts(values: Sequence[bytes | str | None]) -> tuple[int | None, ...]:
return _OPTIONAL_COUNTS.validate_python(
tuple(value.decode("utf-8") if isinstance(value, bytes) else value for value in values)
)
def _get_call_stack_info(num_frames: int = 2) -> str:
"""
@ -736,6 +756,43 @@ class RedisCache(BaseCache):
)
raise e
@_redis_circuit_breaker_guard_sync
def increment_with_floor(self, key: str, value: int, ttl: int) -> int:
"""Add ``value`` to ``key``, clamp the result at zero, and give a new key ``ttl``, in one Lua call.
A counter whose key expired while a request was still in flight would otherwise be
recreated negative by that request's decrement. Clamping inside the same call is what
keeps it safe: a separate corrective write could land after another pod's increment and
erase it.
The TTL is set only on a key that has none, so a counter expires ``ttl`` after it was
created rather than ``ttl`` after it was last touched. Refreshing it on every touch
would keep a count a dead worker never decremented alive for as long as the group
takes traffic. Returns the resulting count.
"""
namespaced_key: Final = self.check_and_fix_namespace(key=key)
count: Final[object] = self.redis_client.eval( # pyright: ignore[reportAttributeAccessIssue] # stubs omit eval
_INCREMENT_WITH_FLOOR_LUA, 1, namespaced_key, value, ttl
)
return _LUA_COUNT.validate_python(count)
@_redis_circuit_breaker_guard_sync
def batch_get_counts(self, key_list: list[str]) -> tuple[int | None, ...]:
"""Read integer counters for ``key_list``, in order, raising when Redis cannot answer.
``batch_get_cache`` swallows every failure and returns an empty dict, which the caller
cannot tell apart from "every counter is unset". A caller that has to fall back to its
own numbers when Redis is unreachable needs the failure, not a dict of zeros.
"""
namespaced_keys: Final = [self.check_and_fix_namespace(key=key) for key in key_list]
return _decoded_counts(self._run_redis_mget_operation(keys=namespaced_keys))
@_redis_circuit_breaker_guard
async def async_batch_get_counts(self, key_list: list[str]) -> tuple[int | None, ...]:
"""Async twin of ``batch_get_counts``, raising on failure the same way."""
namespaced_keys: Final = [self.check_and_fix_namespace(key=key) for key in key_list]
return _decoded_counts(await self._async_run_redis_mget_operation(keys=namespaced_keys))
@_redis_circuit_breaker_guard
async def async_scan_iter(self, pattern: str, count: int = 100) -> list:
start_time: Final = time.time()
@ -1241,6 +1298,14 @@ class RedisCache(BaseCache):
result = result.decode()
return float(result)
@_redis_circuit_breaker_guard
async def async_increment_with_floor(self, key: str, value: int, ttl: int) -> int:
"""Async twin of ``increment_with_floor``, sharing its Lua script and its guarantees."""
_redis_client: Final = self._async_commands()
namespaced_key: Final = self.check_and_fix_namespace(key=key)
count: Final = await _redis_client.eval(_INCREMENT_WITH_FLOOR_LUA, 1, namespaced_key, value, ttl)
return _LUA_COUNT.validate_python(count)
async def flush_cache_buffer(self):
print_verbose(f"flushing to redis....reached size of buffer {len(self.redis_batch_writing_buffer)}")
await self.async_set_cache_pipeline(self.redis_batch_writing_buffer)

14
litellm/caching/redis_semantic_cache.py
View file

@ -116,7 +116,7 @@ class RedisSemanticCache(BaseCache):
password = password or os.environ["REDIS_PASSWORD"]
except KeyError as e:
# Raise a more informative exception if any of the required keys are missing
missing_var: Final = e.args[0]
missing_var: Final[object] = e.args[0]
raise ValueError(
f"Missing required Redis configuration: {missing_var}. Provide {missing_var} or redis_url."
) from e
@ -273,7 +273,7 @@ class RedisSemanticCache(BaseCache):
return prompt or None
@classmethod
def _collect_responses_input_text(cls, value: Any, prompt_parts: list[str]) -> None:
def _collect_responses_input_text(cls, value: object, prompt_parts: list[str]) -> None:
value = cls._coerce_response_input_value(value)
if value is None:
return
@ -334,7 +334,7 @@ class RedisSemanticCache(BaseCache):
resolve_embedding_max_input_tokens(self.embedding_max_input_tokens, self.embedding_model, router),
)
def _get_embedding(self, prompt: str, metadata: dict[str, Any] | None = None) -> list[float]:
def _get_embedding(self, prompt: str, metadata: dict[str, object] | None = None) -> list[float]:
"""
Routes through the proxy Router when the embedding model is a Router
deployment so per-deployment auth (e.g. Bedrock aws_role_name) applies,
@ -425,7 +425,7 @@ class RedisSemanticCache(BaseCache):
prompt_embedding: Final = self._get_embedding(prompt, metadata=kwargs.get("metadata"))
store_kwargs: Final[dict[str, Any]] = {
store_kwargs: Final[dict[str, object]] = {
"vector": prompt_embedding,
"filters": self._get_cache_filters(key),
}
@ -504,7 +504,7 @@ class RedisSemanticCache(BaseCache):
print_verbose(f"Error retrieving from Redis semantic cache: {e}")
kwargs.setdefault("metadata", {})["semantic-similarity"] = 0.0
async def _get_async_embedding(self, prompt: str, metadata: dict[str, Any] | None = None) -> list[float]:
async def _get_async_embedding(self, prompt: str, metadata: dict[str, object] | None = None) -> list[float]:
"""
Asynchronously generate an embedding for the given prompt.
@ -571,7 +571,7 @@ class RedisSemanticCache(BaseCache):
# Generate embedding for the value (response) to cache
prompt_embedding: Final = await self._get_async_embedding(prompt, metadata=kwargs.get("metadata"))
store_kwargs: Final[dict[str, Any]] = {
store_kwargs: Final[dict[str, object]] = {
"vector": prompt_embedding,
"filters": self._get_cache_filters(key),
}
@ -665,7 +665,7 @@ class RedisSemanticCache(BaseCache):
aindex: Final = await self.llmcache._get_async_index()
return await aindex.info()
async def async_set_cache_pipeline(self, cache_list: list[tuple[str, Any]], **kwargs: object) -> None:
async def async_set_cache_pipeline(self, cache_list: list[tuple[str, object]], **kwargs: object) -> None:
"""
Asynchronously store multiple values in the semantic cache.

63
litellm/completion_extras/litellm_responses_transformation/transformation.py
View file

@ -5,8 +5,11 @@ Handler for transforming /chat/completions api requests to litellm.responses req
import json
import os
from collections.abc import AsyncIterator, Callable, Iterable, Iterator, Mapping, Sequence
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, Literal, TypedDict, Union, cast, get_args
from openai.types.chat import ChatCompletion
from openai.types.responses import Response
from openai.types.responses.custom_tool_param import CustomToolParam
from openai.types.responses.response_input_param import (
FunctionCallOutput,
@ -22,7 +25,7 @@ import litellm
from litellm import ModelResponse
from litellm._logging import verbose_logger
from litellm.litellm_core_utils.prompt_templates.common_utils import (
responses_reasoning_item_from_thinking_blocks,
responses_reasoning_items_from_thinking_blocks,
)
from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator
from litellm.llms.base_llm.bridges.completion_transformation import (
@ -33,7 +36,7 @@ from litellm.responses.sse_output_recovery import (
record_output_item_chunk,
record_output_text_chunk,
)
from litellm.responses.utils import normalize_responses_api_stream_options
from litellm.responses.utils import ResponsesAPIRequestUtils, normalize_responses_api_stream_options
from litellm.types.llms.openai import (
REASONING_EFFORT,
ChatCompletionAnnotation,
@ -43,6 +46,7 @@ from litellm.types.llms.openai import (
ChatCompletionToolParamFunctionChunk,
Reasoning,
ResponsesAPIOptionalRequestParams,
ResponsesAPIResponse,
ResponsesAPIStreamEvents,
)
from litellm.types.utils import GenericStreamingChunk, ModelResponseStream
@ -54,7 +58,7 @@ if TYPE_CHECKING:
)
from pydantic import BaseModel
from litellm import LiteLLMLoggingObj, ModelResponse
from litellm import LiteLLMLoggingObj
from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator
from litellm.types.llms.openai import (
ALL_RESPONSES_API_TOOL_PARAMS,
@ -69,6 +73,28 @@ if TYPE_CHECKING:
from litellm.types.utils import Choices
_CHAT_COMPLETION_FIELDS: Final = frozenset((*ModelResponse.model_fields, "usage"))
_RESPONSES_API_ONLY_FIELDS: Final = frozenset((*Response.model_fields, *ResponsesAPIResponse.model_fields)) - frozenset(
ChatCompletion.model_fields
)
def _provider_metadata(response_fields: Mapping[str, object] | None) -> Mapping[str, object]:
return MappingProxyType(
{
key: value
for key, value in (response_fields.items() if response_fields else ())
if value is not None and key not in _CHAT_COMPLETION_FIELDS and key not in _RESPONSES_API_ONLY_FIELDS
}
)
def _upstream_response_id(response_id: str | None) -> str | None:
if response_id is None:
return None
return ResponsesAPIRequestUtils.decode_previous_response_id_to_original_previous_response_id(response_id)
class _ReasoningSummaryText(TypedDict):
type: str
text: str
@ -103,8 +129,8 @@ def _reasoning_input_items(msg: "AllMessageValues") -> list[dict[str, object]]:
return stored
raw_blocks: Final = msg.get("thinking_blocks") or ()
blocks: Final = cast("Iterable[ChatCompletionThinkingBlock]", raw_blocks) # cast-ok: untyped client json
from_thinking: Final = responses_reasoning_item_from_thinking_blocks(blocks)
return [] if from_thinking is None else [dict(from_thinking)] # mutable-ok: API message payload
replayed: Final = responses_reasoning_items_from_thinking_blocks(blocks)
return [dict(item) for item in replayed] # mutable-ok: API message payload
def _build_reasoning_item(
@ -201,7 +227,7 @@ class _ChatToolCallDict(ChatCompletionToolCallChunk, total=False):
provider_specific_fields: Mapping[str, object]
def _tool_call_dict_from_output_item(item: Mapping[str, Any], index: int) -> _ChatToolCallDict:
def tool_call_dict_from_output_item(item: Mapping[str, Any], index: int) -> _ChatToolCallDict:
"""Convert a ``function_call`` or ``custom_tool_call`` output item dict to a chat
completions tool_call dict. Custom (grammar/freeform) tool calls carry their raw
string payload in ``input`` rather than ``arguments``; both map to
@ -344,7 +370,12 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
and isinstance(tool_call.get("custom"), dict)
)
for msg in messages:
leading_system_count: Final = next(
(index for index, msg in enumerate(messages) if msg.get("role") != "system"),
len(messages),
)
for index, msg in enumerate(messages):
role = msg.get("role")
content = msg.get("content", "")
tool_calls = msg.get("tool_calls")
@ -352,7 +383,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
if role == "system":
# Extract system message as instructions
if isinstance(content, str):
if isinstance(content, str) and index < leading_system_count:
if instructions:
# Concatenate multiple system prompts with a space
instructions = f"{instructions} {content}"
@ -724,7 +755,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
# Tool calls accumulate into the single trailing tool_calls choice
# like the typed branches above; a choice per call would hide every
# call after choices[0] from chat clients
accumulated_tool_calls.append(_tool_call_dict_from_output_item(raw_item, tool_call_index))
accumulated_tool_calls.append(tool_call_dict_from_output_item(raw_item, tool_call_index))
tool_call_index += 1
elif handle_raw_dict_callback is not None:
choice, index = handle_raw_dict_callback(item=raw_item, index=index)
@ -904,6 +935,10 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(raw_response.usage),
)
model_response.id = _upstream_response_id(raw_response.id) or raw_response.id
for key, value in _provider_metadata(raw_response.model_extra).items():
setattr(model_response, key, value)
# Preserve hidden params from the ResponsesAPIResponse, especially the headers
# which contain important provider information like x-request-id
raw_response_hidden_params: Final = getattr(raw_response, "_hidden_params", {})
@ -1359,20 +1394,22 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator):
if event_type == "response.created":
# Initial response creation event
verbose_logger.debug("Chat provider: response.created -> %s", parsed_chunk)
created_response: Final = parsed_chunk.get("response")
return ModelResponseStream(
id=_upstream_response_id(created_response.get("id")) if created_response else None,
choices=[
StreamingChoices(
index=0,
delta=Delta(content=""),
finish_reason=None,
)
]
],
)
elif event_type == "response.output_item.added":
# New output item added
output_item = parsed_chunk.get("item", {})
if output_item.get("type") in ("function_call", "custom_tool_call"):
converted: Final = _tool_call_dict_from_output_item(output_item, parsed_chunk.get("output_index", 0))
converted: Final = tool_call_dict_from_output_item(output_item, parsed_chunk.get("output_index", 0))
provider_specific_fields: Final = converted.get("provider_specific_fields")
function_chunk: Final = ChatCompletionToolCallFunctionChunk(
@ -1447,7 +1484,7 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator):
index=0,
delta=Delta(
tool_calls=(
_tool_call_dict_from_output_item(
tool_call_dict_from_output_item(
output_item, parsed_chunk.get("output_index", 0)
),
)
@ -1534,6 +1571,7 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator):
from litellm.responses.utils import ResponseAPILoggingUtils
usage = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(response_data.get("usage"))
provider_metadata: Final = _provider_metadata(response_data)
return ModelResponseStream(
choices=[
StreamingChoices(
@ -1546,6 +1584,7 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator):
)
],
usage=usage,
provider_specific_fields=dict(provider_metadata) or None, # mutable-ok: field is typed dict
)
else:
pass

12
litellm/compression/compress.py
View file

@ -66,7 +66,7 @@ def _build_retrieval_tools(keys: list[str], call_type: str) -> list[dict]:
return cast(list[dict], anthropic_tools)
def _content_to_text(content: Any) -> str:
def _content_to_text(content: object) -> str:
"""
Convert OpenAI/Anthropic message content blocks to plain text.
@ -78,7 +78,7 @@ def _content_to_text(content: Any) -> str:
Implemented iteratively (stack-based) to avoid unbounded recursion.
"""
parts: Final[list[str]] = []
stack: Final[list[Any]] = [content]
stack: Final[list[object]] = [content]
while stack:
item = stack.pop()
if isinstance(item, str):
@ -111,7 +111,7 @@ def _normalize_messages_for_compression(
f"Unsupported call_type={call_type!r} for compression. Expected one of: {sorted(_SUPPORTED_CALL_TYPES)}."
)
original_messages: Final[list[dict[str, Any]]] = [dict(m) for m in messages]
original_messages: Final[list[dict[str, object]]] = [dict(m) for m in messages]
normalized_messages: Final[list[dict]] = []
for msg in original_messages:
@ -132,7 +132,7 @@ def _extract_last_user_message(messages: list[dict]) -> str:
return ""
def _extract_tool_use_ids(content: Any) -> list[str]:
def _extract_tool_use_ids(content: object) -> list[str]:
if not isinstance(content, list):
return []
tool_use_ids: Final[list[str]] = []
@ -147,7 +147,7 @@ def _extract_tool_use_ids(content: Any) -> list[str]:
return tool_use_ids
def _extract_tool_result_ids(content: Any) -> set[str]:
def _extract_tool_result_ids(content: object) -> set[str]:
if not isinstance(content, list):
return set()
tool_result_ids: Final[set[str]] = set()
@ -337,7 +337,7 @@ def compress(
compression_trigger: int = 200_000,
compression_target: int | None = None,
embedding_model: str | None = None,
embedding_model_params: dict[str, Any] | None = None,
embedding_model_params: Mapping[str, object] | None = None,
compression_cache: DualCache | None = None,
) -> CompressedResult:
"""

3
litellm/compression/scoring/embedding_scorer.py
View file

@ -5,6 +5,7 @@ Computes cosine similarity between the query embedding and each message embeddin
"""
import math
from collections.abc import Mapping
from typing import Any, Final
from litellm.caching.dual_cache import DualCache
@ -49,7 +50,7 @@ def embedding_score_messages(
messages: list[dict],
model: str,
cache: DualCache | None = None,
embedding_model_params: dict[str, Any] | None = None,
embedding_model_params: Mapping[str, object] | None = None,
) -> list[float]:
"""
Score each message's semantic similarity to the query using embeddings.

45
litellm/constants.py
View file

@ -40,7 +40,7 @@ ROUTER_SETTINGS_MANAGED_OUTSIDE_CONFIG: Final[frozenset[str]] = frozenset(
"router_general_settings",
"ignore_invalid_deployments",
"fallback_access_check",
"heuristic_v2_router_limit",
"auto_router_capability_limit",
}
)
DEFAULT_BATCH_SIZE: Final = int(os.getenv("DEFAULT_BATCH_SIZE", 512))
@ -55,6 +55,7 @@ S3_PREFIX_DIGEST_CHARS: Final = 16
MAX_S3_OBJECT_DOWNLOAD_FILENAME_BYTES: Final = 1024
DEFAULT_SQS_FLUSH_INTERVAL_SECONDS: Final = int(os.getenv("DEFAULT_SQS_FLUSH_INTERVAL_SECONDS", 10))
DEFAULT_NUM_WORKERS_LITELLM_PROXY: Final = int(os.getenv("DEFAULT_NUM_WORKERS_LITELLM_PROXY", 1))
budget_reservation_disabled_info_emitted = False
DYNAMIC_RATE_LIMIT_ERROR_THRESHOLD_PER_MINUTE = int(os.getenv("DYNAMIC_RATE_LIMIT_ERROR_THRESHOLD_PER_MINUTE", 1))
DEFAULT_SQS_BATCH_SIZE: Final = int(os.getenv("DEFAULT_SQS_BATCH_SIZE", 512))
SQS_SEND_MESSAGE_ACTION: Final = "SendMessage"
@ -72,6 +73,9 @@ DEFAULT_MAX_TOKENS: Final = int(os.getenv("DEFAULT_MAX_TOKENS", 4096))
DEFAULT_ALLOWED_FAILS: Final = int(os.getenv("DEFAULT_ALLOWED_FAILS", 3))
DEFAULT_REDIS_SYNC_INTERVAL: Final = int(os.getenv("DEFAULT_REDIS_SYNC_INTERVAL", 1))
DEFAULT_COOLDOWN_TIME_SECONDS: Final = int(os.getenv("DEFAULT_COOLDOWN_TIME_SECONDS", 5))
DEFAULT_COOLDOWN_REDIS_READ_INTERVAL_SECONDS: Final = float(
os.getenv("DEFAULT_COOLDOWN_REDIS_READ_INTERVAL_SECONDS", "1")
)
DEFAULT_REPLICATE_POLLING_RETRIES: Final = int(os.getenv("DEFAULT_REPLICATE_POLLING_RETRIES", 5))
DEFAULT_REPLICATE_POLLING_DELAY_SECONDS: Final = int(os.getenv("DEFAULT_REPLICATE_POLLING_DELAY_SECONDS", 1))
DEFAULT_IMAGE_TOKEN_COUNT: Final = int(os.getenv("DEFAULT_IMAGE_TOKEN_COUNT", 250))
@ -89,6 +93,7 @@ LITELLM_MAX_STREAMING_DURATION_SECONDS: Final = (
# Data URIs exceeding this are replaced with a size placeholder.
# Set to 0 to disable truncation.
MAX_BASE64_LENGTH_FOR_LOGGING: Final = int(os.getenv("MAX_BASE64_LENGTH_FOR_LOGGING", 64))
BASE64_TRUNCATION_OFFLOAD_THRESHOLD_CHARS: Final = 256 * 1024
REDACTED_BY_LITELLM: Final = "redacted-by-litellm"
# in-memory stand-in handed to provider converters for redacted arguments; never stored
REDACTED_TOOL_CALL_ARGUMENTS_PLACEHOLDER: Final = "{}"
@ -138,6 +143,7 @@ DEFAULT_MCP_SEMANTIC_FILTER_SIMILARITY_THRESHOLD: Final = float(
os.getenv("DEFAULT_MCP_SEMANTIC_FILTER_SIMILARITY_THRESHOLD", 0.3)
)
MAX_MCP_SEMANTIC_FILTER_TOOLS_HEADER_LENGTH: Final = int(os.getenv("MAX_MCP_SEMANTIC_FILTER_TOOLS_HEADER_LENGTH", 150))
MAX_GUARDRAIL_SCAN_METADATA_HEADER_LENGTH: Final = 2048
DEFAULT_AUTO_ROUTER_MAX_INPUT_CHARS: Final = 2000
@ -192,6 +198,7 @@ LITELLM_UI_ALLOW_HEADERS: Final = [
"x-litellm-adaptive-router-model",
"x-litellm-applied-guardrails",
"x-litellm-guardrail-scan-id",
"x-litellm-guardrail-scan-metadata",
"x-litellm-cache-key",
]
@ -215,6 +222,9 @@ MAX_CALLBACKS: Final = get_env_int("LITELLM_MAX_CALLBACKS", 100)
# so the deployment-level hook does not re-run them for the same request
PRE_CALL_EXECUTED_GUARDRAILS_KEY: Final = "_pre_call_executed_guardrails"
# Attribute stamped on log_guardrail_information wrappers so __init_subclass__ does not wrap them again
LOGS_GUARDRAIL_INFORMATION_MARKER: Final = "_litellm_logs_guardrail_information"
# Generic fallback for unknown models
DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET: Final = int(
os.getenv("DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET", 128)
@ -325,6 +335,7 @@ DEFAULT_SSL_CIPHERS: Final = os.getenv(
########### v2 Architecture constants for managing writing updates to the database ###########
REDIS_UPDATE_BUFFER_KEY: Final = "litellm_spend_update_buffer"
REDIS_GATEWAY_REQUESTS_BUFFER_KEY: Final = "litellm_gateway_requests_buffer"
REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY: Final = "litellm_daily_spend_update_buffer"
REDIS_DAILY_TEAM_SPEND_UPDATE_BUFFER_KEY: Final = "litellm_daily_team_spend_update_buffer"
REDIS_DAILY_ORG_SPEND_UPDATE_BUFFER_KEY: Final = "litellm_daily_org_spend_update_buffer"
@ -387,6 +398,18 @@ TIKTOKEN_ENCODE_CHUNK_SIZE_CHARS: Final = get_env_int_in_range(
minimum=1,
maximum=TIKTOKEN_ENCODE_MAX_CHUNK_SIZE_CHARS,
)
TOKEN_COUNTER_MAX_EXACT_CHARS: Final = get_env_int_in_range(
"TOKEN_COUNTER_MAX_EXACT_CHARS",
default=4_000_000,
minimum=1,
maximum=1_000_000_000,
)
TOKEN_COUNTER_MAX_CONCURRENT_COUNTS: Final = get_env_int_in_range(
"TOKEN_COUNTER_MAX_CONCURRENT_COUNTS",
default=4,
minimum=1,
maximum=256,
)
MAX_TILE_WIDTH: Final = int(os.getenv("MAX_TILE_WIDTH", 512))
MAX_TILE_HEIGHT: Final = int(os.getenv("MAX_TILE_HEIGHT", 512))
OPENAI_FILE_SEARCH_COST_PER_1K_CALLS: Final = float(os.getenv("OPENAI_FILE_SEARCH_COST_PER_1K_CALLS", 2.5 / 1000))
@ -559,6 +582,7 @@ LOGGING_WORKER_AGGRESSIVE_CLEAR_COOLDOWN_SECONDS: Final = float(
LOGGING_EXECUTOR_MAX_THREADS: Final = get_env_int("LOGGING_EXECUTOR_MAX_THREADS", 100)
LOGGING_EXECUTOR_MAX_PENDING_TASKS: Final = get_env_int("LOGGING_EXECUTOR_MAX_PENDING_TASKS", 10_000)
LOGGING_EXECUTOR_DROPPED_TASK_LOG_INTERVAL_SECONDS: Final = 30.0
AWS_SIGNING_MAX_THREADS: Final = 16
DD_TRACER_STREAMING_CHUNK_YIELD_RESOURCE: Final = os.getenv(
"DD_TRACER_STREAMING_CHUNK_YIELD_RESOURCE", "streaming.chunk.yield"
)
@ -1366,6 +1390,7 @@ bedrock_embedding_models: Final[set] = set(
"cohere.embed-multilingual-v3",
"cohere.embed-v4:0",
"twelvelabs.marengo-embed-2-7-v1:0",
"twelvelabs.marengo-embed-3-0-v1:0",
]
)
@ -1453,7 +1478,10 @@ LITELLM_METADATA_FIELD: Final = "litellm_metadata"
OLD_LITELLM_METADATA_FIELD: Final = "metadata"
RETURN_RAW_MODEL_NAME_METADATA_KEY: Final = "_complexity_router_return_raw_model_name"
SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY: Final = "_session_deployment_affinity_ttl"
OUTPUT_TOKEN_CEILING_PARAMS: Final = frozenset({"max_tokens", "max_completion_tokens", "max_output_tokens"})
CLIENT_OUTPUT_CEILING_METADATA_KEY: Final = "_client_output_ceiling"
CONSUMED_REQUEST_TAGS_METADATA_KEY: Final = "_consumed_request_tags"
ROUTING_REQUEST_TAGS_METADATA_KEY: Final = "_routing_request_tags"
INTERNAL_CALL_ORIGIN_METADATA_KEY: Final = "internal_call_origin"
SESSION_ID_GENERATED_METADATA_KEY: Final = "litellm_session_id_generated"
SESSION_ID_OMITTED_METADATA_KEY: Final = "litellm_session_id_omitted"
@ -1654,6 +1682,7 @@ SPEND_LOG_QUEUE_POLL_INTERVAL: Final = float(os.getenv("SPEND_LOG_QUEUE_POLL_INT
RESPONSES_SESSION_LOOKUP_MAX_ATTEMPTS: Final = max(1, int(os.getenv("RESPONSES_SESSION_LOOKUP_MAX_ATTEMPTS", "3")))
RESPONSES_SESSION_LOOKUP_RETRY_INTERVAL: Final = float(os.getenv("RESPONSES_SESSION_LOOKUP_RETRY_INTERVAL", "0.2"))
SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE: Final = int(os.getenv("SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE", 10000))
PROXY_DB_LOOKUP_MAX_CONCURRENCY: Final = max(1, int(os.getenv("PROXY_DB_LOOKUP_MAX_CONCURRENCY", "25")))
DEFAULT_CRON_JOB_LOCK_TTL_SECONDS: Final = int(os.getenv("DEFAULT_CRON_JOB_LOCK_TTL_SECONDS", 60)) # 1 minute
PROXY_BUDGET_RESCHEDULER_MIN_TIME: Final = int(os.getenv("PROXY_BUDGET_RESCHEDULER_MIN_TIME", 597))
RESET_BUDGET_JOB_BATCH_SIZE: Final = max(1, int(os.getenv("RESET_BUDGET_JOB_BATCH_SIZE", "500")))
@ -1755,6 +1784,10 @@ LITELLM_SETTINGS_SAFE_DB_OVERRIDES: Final = [
SPECIAL_LITELLM_AUTH_TOKEN: Final = ["ui-token"]
DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL = int(os.getenv("DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL", 60))
DEFAULT_ACCESS_GROUP_CACHE_TTL: Final = int(os.getenv("DEFAULT_ACCESS_GROUP_CACHE_TTL", 600))
SPEND_LOG_KEY_METADATA_CACHE_TTL: Final = 600
SPEND_LOG_KEY_METADATA_MISS_CACHE_TTL: Final = 30
SPEND_LOG_KEY_METADATA_CACHE_MAX_ITEMS: Final = 10000
SPEND_LOG_KEY_METADATA_QUERY_TIMEOUT_MS: Final = 5000
# Short TTL for negative MCP access-group existence lookups. Keeps unauthenticated
# callers from forcing a DB query per request for unknown names, while bounding
# staleness so a transient DB error (which surfaces as an empty list) cannot
@ -1897,6 +1930,16 @@ HTTP_FRAMING_HEADERS: Final[frozenset[str]] = frozenset(
}
)
PROVIDER_REQUEST_ID_HEADERS: Final[tuple[str, ...]] = (
"x-amzn-requestid",
"x-request-id",
"request-id",
"x-ms-request-id",
"apim-request-id",
"x-goog-request-id",
"cf-ray",
)
# Browser-facing security headers that a malicious or misconfigured upstream
# provider must not be able to set on the proxy's own response.
BROWSER_SECURITY_HEADERS: Final[frozenset[str]] = frozenset(

14
litellm/containers/endpoint_factory.py
View file

@ -11,7 +11,7 @@ import json
from collections.abc import Callable
from functools import partial
from pathlib import Path
from typing import Any, Final, Literal
from typing import Final, Literal
import litellm
from litellm.constants import request_timeout as DEFAULT_REQUEST_TIMEOUT
@ -56,9 +56,9 @@ def create_sync_endpoint_function(endpoint_config: dict) -> Callable:
def endpoint_func(
timeout: int = 600,
custom_llm_provider: Literal["openai", "azure", "azure_text"] = "openai",
extra_headers: dict[str, Any] | None = None,
extra_query: dict[str, Any] | None = None,
extra_body: dict[str, Any] | None = None,
extra_headers: dict[str, object] | None = None,
extra_query: dict[str, object] | None = None,
extra_body: dict[str, object] | None = None,
**kwargs,
):
local_vars: Final = locals()
@ -145,9 +145,9 @@ def create_async_endpoint_function(
async def async_endpoint_func(
timeout: int = 600,
custom_llm_provider: Literal["openai", "azure", "azure_text"] = "openai",
extra_headers: dict[str, Any] | None = None,
extra_query: dict[str, Any] | None = None,
extra_body: dict[str, Any] | None = None,
extra_headers: dict[str, object] | None = None,
extra_query: dict[str, object] | None = None,
extra_body: dict[str, object] | None = None,
**kwargs,
):
local_vars: Final = locals()

61
litellm/cost_calculator.py
View file

@ -25,6 +25,7 @@ from litellm.litellm_core_utils.llm_cost_calc.usage_object_transformation import
TranscriptionUsageObjectTransformation,
)
from litellm.litellm_core_utils.llm_cost_calc.utils import (
BilledTokenRates,
CostCalculatorUtils,
_generic_cost_per_character,
_get_regional_uplift_multiplier,
@ -45,6 +46,9 @@ from litellm.llms.azure.cost_calculation import (
from litellm.llms.azure_ai.cost_calculator import (
cost_per_token as azure_ai_cost_per_token,
)
from litellm.llms.azure_ai.cost_calculator import (
is_azure_model_router as azure_ai_is_model_router_name,
)
from litellm.llms.base_llm.search.transformation import SearchResponse
from litellm.llms.bedrock.cost_calculation import (
cost_per_token as bedrock_cost_per_token,
@ -81,6 +85,7 @@ from litellm.llms.together_ai.cost_calculator import (
get_model_params_and_category,
has_together_registry_pricing,
)
from litellm.llms.vertex_ai.common_utils import get_vertex_ai_lyria_generation_cost
from litellm.llms.vertex_ai.cost_calculator import (
cost_per_character as google_cost_per_character,
)
@ -497,6 +502,13 @@ def cost_per_token(
# see this https://learn.microsoft.com/en-us/azure/ai-services/openai/concepts/models
if call_type == "speech" or call_type == "aspeech":
lyria_generation_cost: Final = (
get_vertex_ai_lyria_generation_cost(model=model_without_prefix)
if custom_llm_provider in ("vertex_ai", "vertex_ai_beta")
else None
)
if lyria_generation_cost is not None:
return 0.0, lyria_generation_cost
speech_model_info = litellm.get_model_info(model=model_without_prefix, custom_llm_provider=custom_llm_provider)
cost_metric: Final = select_cost_metric_for_model(speech_model_info)
prompt_cost: float = 0.0
@ -1118,6 +1130,7 @@ def _store_cost_breakdown_in_logging_obj(
service_tier: str | None = None,
data_residency: str | None = None,
vertex_location: str | None = None,
billed_token_rates: BilledTokenRates | None = None,
) -> None:
"""
Helper function to store cost breakdown in the logging object.
@ -1162,6 +1175,7 @@ def _store_cost_breakdown_in_logging_obj(
service_tier=service_tier,
data_residency=data_residency,
vertex_location=vertex_location,
billed_token_rates=billed_token_rates,
)
except Exception as breakdown_error:
@ -1656,11 +1670,10 @@ def completion_cost(
data_residency=data_residency,
vertex_location=vertex_location,
response=completion_response,
request_model=request_model_for_cost,
)
# Get additional costs from provider (e.g., routing fees, infrastructure costs)
if custom_llm_provider == "azure_ai":
if custom_llm_provider == "azure_ai" and not azure_ai_is_model_router_name(model):
model_for_additional_costs = request_model_for_cost
if completion_response is not None:
hidden_params = getattr(completion_response, "_hidden_params", None) or {}
@ -1732,6 +1745,7 @@ def completion_cost(
_reasoning_cost: float | None = None
_cache_read_cost: float | None = None
_cache_creation_cost: float | None = None
_billed_token_rates: BilledTokenRates | None = None
if cost_per_token_usage_object is not None and model:
_breakdown_provider: str | None = (
custom_llm_provider if isinstance(custom_llm_provider, str) else None
@ -1743,10 +1757,12 @@ def completion_cost(
service_tier=service_tier,
data_residency=data_residency,
vertex_location=vertex_location,
custom_cost_per_token=custom_cost_per_token,
)
_reasoning_cost = _token_type_breakdown.reasoning_cost
_cache_read_cost = _token_type_breakdown.cache_read_cost
_cache_creation_cost = _token_type_breakdown.cache_creation_cost
_billed_token_rates = _token_type_breakdown.rates
_store_cost_breakdown_in_logging_obj(
litellm_logging_obj=litellm_logging_obj,
prompt_tokens_cost_usd_dollar=prompt_tokens_cost_usd_dollar,
@ -1766,6 +1782,7 @@ def completion_cost(
service_tier=service_tier,
data_residency=data_residency,
vertex_location=vertex_location,
billed_token_rates=_billed_token_rates,
)
return _final_cost
@ -2413,6 +2430,46 @@ class RealtimeAPITokenUsageProcessor(BaseTokenUsageProcessor):
)
_RESPONSES_WS_BILLABLE_EVENT_TYPES: Final = frozenset({"response.completed", "response.incomplete"})
class _ResponsesWsEventResponse(BaseModel):
usage: Mapping[str, object] | None = None
class _ResponsesWsEvent(BaseModel):
type: str = ""
response: _ResponsesWsEventResponse | None = None
class ResponsesWebSocketTokenUsageProcessor(BaseTokenUsageProcessor):
@staticmethod
def collect_usage_from_responses_ws_results(
results: Sequence[Mapping[str, object]],
) -> tuple[Usage, ...]:
events: Final = tuple(_ResponsesWsEvent.model_validate(result) for result in results)
return tuple(
ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage( # pyright: ignore[reportPrivateUsage] # same shared transform the realtime processor uses
event.response.usage
)
for event in events
if event.type in _RESPONSES_WS_BILLABLE_EVENT_TYPES
and event.response is not None
and event.response.usage is not None
)
@staticmethod
def collect_and_combine_usage_from_responses_ws_results(
results: Sequence[Mapping[str, object]],
) -> Usage:
collected_usage_objects: Final = ResponsesWebSocketTokenUsageProcessor.collect_usage_from_responses_ws_results(
results
)
return ResponsesWebSocketTokenUsageProcessor.combine_usage_objects(
list(collected_usage_objects) # mutable-ok: combine_usage_objects requires a list parameter
)
_TRANSCRIPTION_COMPLETED_EVENT_TYPE: Final = "conversation.item.input_audio_transcription.completed"

7
litellm/exceptions.py
View file

@ -85,7 +85,7 @@ _RATE_LIMIT_CATEGORY_VALUES: Final = frozenset(c.value for c in RateLimitErrorCa
_RATE_LIMIT_TYPE_VALUES: Final = frozenset(t.value for t in RateLimitType)
def validate_rate_limit_category(value: Any) -> str | None:
def validate_rate_limit_category(value: object) -> str | None:
"""Return ``value`` only if it matches a known :class:`RateLimitErrorCategory`.
Used at duck-typed read sites (StandardLoggingPayload extraction, Prometheus
@ -100,7 +100,7 @@ def validate_rate_limit_category(value: Any) -> str | None:
return None
def validate_rate_limit_type(value: Any) -> str | None:
def validate_rate_limit_type(value: object) -> str | None:
"""Return ``value`` only if it matches a known :class:`RateLimitType`.
See :func:`validate_rate_limit_category` for the rationale.
@ -338,6 +338,7 @@ class Timeout(openai.APITimeoutError):
num_retries: int | None = None,
headers: dict | None = None,
exception_status_code: int | None = None,
response: httpx.Response | None = None,
):
request: Final = httpx.Request(
method="POST",
@ -352,6 +353,8 @@ class Timeout(openai.APITimeoutError):
self.max_retries = max_retries
self.num_retries = num_retries
self.headers = headers
if response is not None:
self.response = response
# custom function to convert to str
def __str__(self):

113
litellm/experimental_mcp_client/client.py
View file

@ -6,17 +6,36 @@ import asyncio
import base64
import os
from collections.abc import Awaitable, Callable, Generator
from contextlib import AbstractAsyncContextManager
from datetime import timedelta
from functools import partial
from importlib import metadata
from typing import Any, Final, TypeVar
from typing import Any, Final, Protocol, TypeAlias, TypeVar
import httpx
from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
from mcp import ClientSession, McpError, ReadResourceResult, Resource, StdioServerParameters
from mcp.client.sse import sse_client
from mcp.client.stdio import stdio_client
from mcp.shared.message import SessionMessage
from mcp.shared.session import RequestResponder
from typing_extensions import Unpack
streamable_http_client: Any | None = None
_TransportStreams: TypeAlias = tuple[
MemoryObjectReceiveStream[SessionMessage | Exception],
MemoryObjectSendStream[SessionMessage],
Unpack[tuple[object, ...]],
]
_TransportContext: TypeAlias = AbstractAsyncContextManager[_TransportStreams]
class _StreamableHttpClientFactory(Protocol):
"""The ``streamable_http_client`` entry point this module calls on the installed MCP SDK."""
def __call__(self, *, url: str, http_client: httpx.AsyncClient | None) -> _TransportContext: ...
streamable_http_client: _StreamableHttpClientFactory | None = None
try:
import mcp.client.streamable_http as streamable_http_module
@ -35,15 +54,22 @@ def missing_streamable_http_client_error() -> ImportError:
)
from mcp.types import CallToolRequestParams as MCPCallToolRequestParams
from mcp.types import CallToolResult as MCPCallToolResult
from mcp.types import (
METHOD_NOT_FOUND,
ClientResult,
GetPromptRequestParams,
GetPromptResult,
ListPromptsResult,
ListResourcesResult,
ListResourceTemplatesResult,
Prompt,
ResourceTemplate,
ServerNotification,
ServerRequest,
TextContent,
)
from mcp.types import CallToolRequestParams as MCPCallToolRequestParams
from mcp.types import CallToolResult as MCPCallToolResult
from mcp.types import Tool as MCPTool
from pydantic import AnyUrl
@ -128,8 +154,8 @@ _SDK_READ_TIMEOUT_CODE: Final = int(httpx.codes.REQUEST_TIMEOUT)
otherwise carries JSON-RPC error codes."""
def _as_read_timeout(exc: BaseException) -> TimeoutError | None:
"""The session read timeout elapsing, re-expressed as a ``TimeoutError``, or ``None``.
def as_mcp_read_timeout(exc: BaseException) -> TimeoutError | None:
"""Normalize an MCP SDK read timeout for client and gateway diagnostics, or return ``None``.
The SDK reports its own elapsed read timeout as ``McpError`` carrying an HTTP status code in a
field that otherwise holds JSON-RPC error codes, and it relays an upstream's JSON-RPC error
@ -216,10 +242,12 @@ class MCPSigV4Auth(httpx.Auth):
aws_region_name: str,
):
"""Call STS AssumeRole and return temporary credentials."""
import time
import boto3
from botocore.credentials import Credentials
session_name: Final = aws_session_name or f"litellm-mcp-{int(__import__('time').time())}"
session_name: Final = aws_session_name or f"litellm-mcp-{int(time.time())}"
sts_kwargs: Final[dict] = {"region_name": aws_region_name}
if aws_access_key_id and aws_secret_access_key:
sts_kwargs["aws_access_key_id"] = aws_access_key_id
@ -315,7 +343,7 @@ class MCPClient:
def _create_transport_context(
self,
) -> tuple[Any, httpx.AsyncClient | None]:
) -> tuple[_TransportContext, httpx.AsyncClient | None]:
"""
Create the appropriate transport context based on transport type.
Returns:
@ -408,7 +436,7 @@ class MCPClient:
async def _execute_session_operation(
self,
transport_ctx: Any,
transport_ctx: _TransportContext,
operation: Callable[[ClientSession], Awaitable[TSessionResult]],
) -> TSessionResult:
"""
@ -422,6 +450,18 @@ class MCPClient:
in_flight_error: BaseException | None = None
try:
read_stream, write_stream = transport[0], transport[1]
stream_error: Final[asyncio.Future[Exception]] = asyncio.get_running_loop().create_future()
async def receive_message(
message: RequestResponder[ServerRequest, ClientResult] | ServerNotification | Exception,
) -> None:
if not isinstance(message, (ValueError, httpx.RequestError, OSError)):
return
if not stream_error.done():
stream_error.set_result(message)
# The SDK closes pending requests when its message handler raises.
raise RuntimeError("MCP response stream failed")
# Build session kwargs with optional callbacks
session_kwargs: Final[dict[str, Any]] = {}
if self._sampling_callback is not None:
@ -436,6 +476,7 @@ class MCPClient:
read_stream,
write_stream,
read_timeout_seconds=timedelta(seconds=self.timeout),
message_handler=receive_message,
**session_kwargs,
)
session: Final = await session_ctx.__aenter__()
@ -447,6 +488,10 @@ class MCPClient:
if isinstance(ins, str) and ins.strip():
self._last_initialize_instructions = ins.strip()
return await operation(session)
except McpError:
if stream_error.done():
raise stream_error.result()
raise
finally:
try:
await session_ctx.__aexit__(None, None, None)
@ -481,11 +526,10 @@ class MCPClient:
transport_ctx, http_client = self._create_transport_context()
return await self._execute_session_operation(transport_ctx, operation)
except Exception as e:
read_timeout: Final = _as_read_timeout(e)
read_timeout: Final = as_mcp_read_timeout(e)
if read_timeout is not None:
verbose_logger.warning(
"MCP client timed out after %ss waiting for %s to answer; the server accepted the "
"request and ended its response stream without a JSON-RPC reply",
"MCP client timed out after %ss waiting for a valid MCP response from %s",
self.timeout,
self.server_url or "stdio",
)
@ -737,8 +781,19 @@ class MCPClient:
"""List available prompts from the server."""
verbose_logger.debug("MCP client listing tools from %s", self.server_url or "stdio")
async def _list_prompts_operation(session: ClientSession):
return await session.list_prompts()
async def _list_prompts_operation(session: ClientSession) -> ListPromptsResult:
capabilities: Final = session.get_server_capabilities()
if capabilities is not None and capabilities.prompts is None:
return ListPromptsResult(prompts=[])
try:
return await session.list_prompts()
except McpError as error:
if error.error.code != METHOD_NOT_FOUND:
raise
verbose_logger.debug(
"MCP client list_prompts is unsupported by %s: %s", self.server_url or "stdio", error
)
return ListPromptsResult(prompts=[])
try:
result: Final = await self.run_with_session(_list_prompts_operation)
@ -814,8 +869,19 @@ class MCPClient:
"""List available resources from the server."""
verbose_logger.debug("MCP client listing resources from %s", self.server_url or "stdio")
async def _list_resources_operation(session: ClientSession):
return await session.list_resources()
async def _list_resources_operation(session: ClientSession) -> ListResourcesResult:
capabilities: Final = session.get_server_capabilities()
if capabilities is not None and capabilities.resources is None:
return ListResourcesResult(resources=[])
try:
return await session.list_resources()
except McpError as error:
if error.error.code != METHOD_NOT_FOUND:
raise
verbose_logger.debug(
"MCP client list_resources is unsupported by %s: %s", self.server_url or "stdio", error
)
return ListResourcesResult(resources=[])
try:
result: Final = await self.run_with_session(_list_resources_operation)
@ -850,8 +916,19 @@ class MCPClient:
"""List available resource templates from the server."""
verbose_logger.debug("MCP client listing resource templates from %s", self.server_url or "stdio")
async def _list_resource_templates_operation(session: ClientSession):
return await session.list_resource_templates()
async def _list_resource_templates_operation(session: ClientSession) -> ListResourceTemplatesResult:
capabilities: Final = session.get_server_capabilities()
if capabilities is not None and capabilities.resources is None:
return ListResourceTemplatesResult(resourceTemplates=[])
try:
return await session.list_resource_templates()
except McpError as error:
if error.error.code != METHOD_NOT_FOUND:
raise
verbose_logger.debug(
"MCP client list_resource_templates is unsupported by %s: %s", self.server_url or "stdio", error
)
return ListResourceTemplatesResult(resourceTemplates=[])
try:
result: Final = await self.run_with_session(_list_resource_templates_operation)

13
litellm/files/main.py
View file

@ -14,6 +14,7 @@ from functools import partial
from typing import Any, Final, Literal, cast
import httpx
from openai import AsyncOpenAI, OpenAI
# Type aliases for provider parameters
FileCreateProvider = Literal[
@ -30,7 +31,7 @@ FileCreateProvider = Literal[
FileRetrieveProvider = Literal[
"openai", "azure", "gemini", "vertex_ai", "hosted_vllm", "litellm_proxy", "manus", "anthropic"
]
FileDeleteProvider = Literal["openai", "azure", "gemini", "litellm_proxy", "manus", "anthropic"]
FileDeleteProvider = Literal["openai", "azure", "gemini", "bedrock", "litellm_proxy", "manus", "anthropic"]
FileListProvider = Literal["openai", "azure", "litellm_proxy", "manus", "anthropic"]
import litellm
from litellm import get_secret_str
@ -431,7 +432,7 @@ async def afile_delete(
extra_headers: dict[str, str] | None = None,
extra_body: dict[str, str] | None = None,
**kwargs,
) -> Coroutine[Any, Any, FileObject]:
) -> Coroutine[object, object, FileObject]:
"""
Async: Delete file
@ -1002,8 +1003,8 @@ def file_content_streaming(
timeout: float | httpx.Timeout,
logging_obj: LiteLLMLoggingObj | None,
_is_async: bool,
client: Any | None,
) -> FileContentStreamingResult | Coroutine[Any, Any, FileContentStreamingResult]:
client: OpenAI | AsyncOpenAI | None,
) -> FileContentStreamingResult | Coroutine[object, object, FileContentStreamingResult]:
if logging_obj is not None:
logging_obj.model = model or ""
logging_obj.model_call_details["model"] = model or ""
@ -1028,8 +1029,8 @@ def file_content_streaming(
headers=response.headers,
)
response: FileContentStreamingResult | Coroutine[Any, Any, FileContentStreamingResult] = FileContentStreamingResult(
stream_iterator=iter(()), headers={}
response: FileContentStreamingResult | Coroutine[object, object, FileContentStreamingResult] = (
FileContentStreamingResult(stream_iterator=iter(()), headers={})
)
if custom_llm_provider in OPENAI_COMPATIBLE_BATCH_AND_FILES_PROVIDERS:
openai_creds: Final = get_openai_credentials(

14
litellm/fine_tuning/main.py
View file

@ -11,7 +11,7 @@ https://platform.openai.com/docs/api-reference/fine-tuning
import asyncio
import contextvars
import os
from collections.abc import Coroutine
from collections.abc import Coroutine, Mapping
from functools import partial
from typing import Any, Final, Literal
@ -37,8 +37,8 @@ vertex_fine_tuning_apis_instance: Final = VertexFineTuningAPI()
def _prepare_azure_extra_body(
extra_body: dict[str, Any] | None,
kwargs: dict[str, Any],
azure_specific_hyperparams: dict[str, Any],
kwargs: Mapping[str, object],
azure_specific_hyperparams: Mapping[str, object],
) -> dict[str, Any]:
"""
Prepare extra_body for Azure fine-tuning API by combining Azure-specific parameters.
@ -138,7 +138,7 @@ def _build_fine_tuning_job_data(model, training_file, hyperparameters, suffix, v
def _resolve_fine_tuning_timeout(
timeout: Any,
timeout: float | str | httpx.Timeout | None,
custom_llm_provider: str,
) -> float | httpx.Timeout:
"""Normalise a raw timeout value to a float (seconds) or httpx.Timeout for fine-tuning calls."""
@ -163,7 +163,7 @@ def create_fine_tuning_job(
extra_headers: dict[str, str] | None = None,
extra_body: dict[str, str] | None = None,
**kwargs,
) -> LiteLLMFineTuningJob | Coroutine[Any, Any, LiteLLMFineTuningJob]:
) -> LiteLLMFineTuningJob | Coroutine[object, object, LiteLLMFineTuningJob]:
"""
Creates a fine-tuning job which begins the process of creating a new model from a given dataset.
@ -375,7 +375,7 @@ def cancel_fine_tuning_job(
extra_headers: dict[str, str] | None = None,
extra_body: dict[str, str] | None = None,
**kwargs,
) -> LiteLLMFineTuningJob | Coroutine[Any, Any, LiteLLMFineTuningJob]:
) -> LiteLLMFineTuningJob | Coroutine[object, object, LiteLLMFineTuningJob]:
"""
Immediately cancel a fine-tune job.
@ -682,7 +682,7 @@ def retrieve_fine_tuning_job(
extra_headers: dict[str, str] | None = None,
extra_body: dict[str, str] | None = None,
**kwargs,
) -> LiteLLMFineTuningJob | Coroutine[Any, Any, LiteLLMFineTuningJob]:
) -> LiteLLMFineTuningJob | Coroutine[object, object, LiteLLMFineTuningJob]:
"""
Get info about a fine-tuning job.
"""

3
litellm/images/utils.py
View file

@ -1,3 +1,4 @@
from collections.abc import Mapping
from io import BufferedReader, BytesIO
from typing import Any, Final, cast, get_type_hints
@ -61,7 +62,7 @@ class ImageEditRequestUtils:
@staticmethod
def get_requested_image_edit_optional_param(
params: dict[str, Any],
params: Mapping[str, object],
) -> ImageEditOptionalRequestParams:
"""
Filter parameters to only include those defined in ImageEditOptionalRequestParams.

82
litellm/integrations/anthropic_cache_control_hook.py
View file

@ -16,6 +16,8 @@ from collections.abc import Iterable, Mapping, Sequence
from typing import TYPE_CHECKING, Any, Final, cast
from urllib.parse import urlparse
from pydantic import TypeAdapter, ValidationError
from litellm._logging import verbose_logger
from litellm.integrations.custom_logger import CustomLogger
from litellm.integrations.custom_prompt_management import CustomPromptManagement
@ -23,6 +25,7 @@ from litellm.integrations.prompt_management_base import PromptManagementClient
from litellm.litellm_core_utils.prompt_templates.common_utils import (
with_prompt_cache_breakpoint,
)
from litellm.llms.anthropic.common_utils import is_claude_code_one_shot_subagent_request
from litellm.types.integrations.anthropic_cache_control_hook import (
GATEWAY_INJECTED_CACHE_METADATA_KEY,
GATEWAY_INJECTED_FOR_EVERY_DEPLOYMENT,
@ -62,10 +65,26 @@ OPENAI_PROMPT_CACHE_BREAKPOINT_BLOCK_TYPES: Final = frozenset(
)
OPENAI_API_HOST: Final = "api.openai.com"
OPENAI_API_BASE_ENV_VARS: Final = ("OPENAI_BASE_URL", "OPENAI_API_BASE")
_OBJECT_MAPPING_ADAPTER: Final = TypeAdapter(dict[object, object])
_OBJECT_LIST_ADAPTER: Final = TypeAdapter(list[object])
AllToolParamValues = ChatCompletionToolParam | AllAnthropicToolsValues
def _validated_object_mapping(value: object) -> dict[object, object] | None:
try:
return _OBJECT_MAPPING_ADAPTER.validate_python(value)
except ValidationError:
return None
def _validated_object_list(value: object) -> list[object] | None:
try:
return _OBJECT_LIST_ADAPTER.validate_python(value)
except ValidationError:
return None
def supports_openai_prompt_cache_breakpoint(model: str) -> bool:
model_map_flag: Final = _model_map_prompt_cache_breakpoint_flag(model)
if model_map_flag is not None:
@ -114,6 +133,36 @@ CARRY_UNMATCHED_MESSAGE_POINTS: Final = "_litellm_carry_unmatched_cache_control_
class AnthropicCacheControlHook(CustomPromptManagement):
@staticmethod
def _request_value(request_kwargs: object, key: str) -> object:
request_mapping: Final = _validated_object_mapping(request_kwargs)
if request_mapping is None:
return None
return request_mapping.get(key)
@staticmethod
def _request_user_agent(request_kwargs: object) -> str | None:
proxy_server_request: Final = AnthropicCacheControlHook._request_value(request_kwargs, "proxy_server_request")
proxy_server_request_mapping: Final = _validated_object_mapping(proxy_server_request)
if proxy_server_request_mapping is None:
return None
headers: Final = proxy_server_request_mapping.get("headers")
headers_mapping: Final = _validated_object_mapping(headers)
if headers_mapping is None:
return None
user_agent: Final = next(
(value for key, value in headers_mapping.items() if isinstance(key, str) and key.lower() == "user-agent"),
None,
)
return user_agent if isinstance(user_agent, str) else None
@staticmethod
def _request_system(request_kwargs: object) -> str | list[object] | None:
system: Final = AnthropicCacheControlHook._request_value(request_kwargs, "system")
if isinstance(system, str):
return system
return _validated_object_list(system)
def get_chat_completion_prompt(
self,
model: str,
@ -520,12 +569,13 @@ class AnthropicCacheControlHook(CustomPromptManagement):
points: Sequence[CacheControlInjectionPoint],
messages: list[AllMessageValues],
tools: list[object] | None,
cache_control: object,
model: str,
custom_llm_provider: str | None,
api_base: object,
prompt_cache_options: object,
) -> Sequence[Mapping[str, object]] | None:
if AnthropicCacheControlHook._should_stand_down(points, messages, None, tools):
if AnthropicCacheControlHook._should_stand_down(points, messages, None, tools, cache_control):
return None
return AnthropicCacheControlHook._stamped_with_dialect(
points, model, custom_llm_provider, api_base, prompt_cache_options
@ -561,6 +611,7 @@ class AnthropicCacheControlHook(CustomPromptManagement):
messages: list[AllMessageValues],
system: str | list | None,
tools: list | None,
cache_control: object = None,
) -> bool:
"""Whether configured injection points must yield to client-set cache_control.
@ -573,13 +624,14 @@ class AnthropicCacheControlHook(CustomPromptManagement):
"""
if all(point.get("_litellm_judged") for point in points):
return False
return AnthropicCacheControlHook._request_has_cache_control(messages, system, tools)
return AnthropicCacheControlHook._request_has_cache_control(messages, system, tools, cache_control)
@staticmethod
def _request_has_cache_control(
messages: list[AllMessageValues],
system: str | list | None,
tools: list | None = None,
cache_control: object = None,
) -> bool:
"""Return True if the request already carries any client-supplied cache_control.
@ -591,6 +643,8 @@ class AnthropicCacheControlHook(CustomPromptManagement):
carry the mark either at the top level (Anthropic shape) or nested under
``function`` (OpenAI shape); the Anthropic chat transform accepts both.
"""
if cache_control is not None:
return True
if AnthropicCacheControlHook.count_request_cache_breakpoints(messages, system) > 0:
return True
if tools is not None:
@ -612,6 +666,8 @@ class AnthropicCacheControlHook(CustomPromptManagement):
custom_llm_provider: str | None,
tools: list | None = None,
enable_prompt_caching: bool | None = None,
cache_control: object = None,
request_kwargs: object = None,
) -> list[CacheControlInjectionPoint]:
"""Default breakpoints when ``litellm.enable_anthropic_prompt_caching`` is on.
@ -649,7 +705,12 @@ class AnthropicCacheControlHook(CustomPromptManagement):
if not supports_prompt_caching(model=model, custom_llm_provider=provider):
return []
if AnthropicCacheControlHook._request_has_cache_control(messages, system, tools):
if AnthropicCacheControlHook._request_has_cache_control(messages, system, tools, cache_control):
return []
if is_claude_code_one_shot_subagent_request(
messages, system, tools, AnthropicCacheControlHook._request_user_agent(request_kwargs)
):
return []
control: Final = AnthropicCacheControlHook._default_control()
@ -665,6 +726,7 @@ class AnthropicCacheControlHook(CustomPromptManagement):
models: Iterable[str],
tools: list[AllToolParamValues] | None = None,
enable_prompt_caching: bool | None = None,
request_kwargs: object = None,
) -> list[AllMessageValues]:
"""Return the messages auto prompt caching will send, default breakpoints included.
@ -681,11 +743,13 @@ class AnthropicCacheControlHook(CustomPromptManagement):
for candidate in (
AnthropicCacheControlHook.get_default_injection_points(
messages=messages,
system=None,
model=model,
custom_llm_provider=None,
tools=tools,
enable_prompt_caching=enable_prompt_caching,
system=AnthropicCacheControlHook._request_system(request_kwargs),
cache_control=AnthropicCacheControlHook._request_value(request_kwargs, "cache_control"),
request_kwargs=request_kwargs,
)
for model in models
)
@ -730,6 +794,7 @@ class AnthropicCacheControlHook(CustomPromptManagement):
non_default_params["cache_control_injection_points"],
messages,
tools,
non_default_params.get("cache_control"),
model,
custom_llm_provider,
api_base,
@ -747,6 +812,8 @@ class AnthropicCacheControlHook(CustomPromptManagement):
custom_llm_provider=custom_llm_provider,
tools=tools,
enable_prompt_caching=enable_prompt_caching,
cache_control=non_default_params.get("cache_control"),
request_kwargs=non_default_params,
)
if points:
non_default_params["cache_control_injection_points"] = points
@ -853,10 +920,13 @@ class AnthropicCacheControlHook(CustomPromptManagement):
enable_prompt_caching: Final = cast( # cast-ok: kwargs is untyped; key stamped as bool by the proxy
bool | None, kwargs.pop("enable_prompt_caching", None)
)
cache_control: Final = kwargs.get("cache_control")
configured: Final = cast( # cast-ok: kwargs is untyped; this key only holds the documented injection-point list
list[CacheControlInjectionPoint] | None, kwargs.pop("cache_control_injection_points", None)
)
if configured and AnthropicCacheControlHook._should_stand_down(configured, typed_messages, system, tools):
if configured and AnthropicCacheControlHook._should_stand_down(
configured, typed_messages, system, tools, cache_control
):
return messages, system
injection_points: list[CacheControlInjectionPoint] = configured or []
if not injection_points and model is not None:
@ -867,6 +937,8 @@ class AnthropicCacheControlHook(CustomPromptManagement):
model=model,
custom_llm_provider=custom_llm_provider,
enable_prompt_caching=enable_prompt_caching,
cache_control=cache_control,
request_kwargs=kwargs,
)
if not injection_points:
return messages, system

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

@ -2,7 +2,7 @@ import json
from collections.abc import Mapping
from typing import TYPE_CHECKING, Any, Final
from typing_extensions import override
from typing_extensions import ReadOnly, TypedDict, override
from litellm._logging import verbose_logger
from litellm.integrations.opentelemetry_utils.base_otel_llm_obs_attributes import (
@ -492,12 +492,12 @@ def _sanitize_optional_params(optional_params: dict | None) -> dict:
return optional_params
def _set_metadata_attributes(span: "Span", metadata: Any | None, span_attrs) -> None:
def _set_metadata_attributes(span: "Span", metadata: object | None, span_attrs) -> None:
if metadata is not None:
safe_set_attribute(span, span_attrs.METADATA, safe_dumps(metadata))
def _extract_metadata_tools(metadata: Any | None) -> list | None:
def _extract_metadata_tools(metadata: object | None) -> list | None:
if not isinstance(metadata, dict):
return None
llm_obj: Final = metadata.get("llm")
@ -670,7 +670,22 @@ def _get_tool_calls(message) -> list | None:
return tool_calls if isinstance(tool_calls, list) and tool_calls else None
def _normalize_tool_call(raw_tc) -> dict[str, Any] | None:
class _NormalizedToolCallFunction(TypedDict):
"""The ``function`` sub-object of a normalized tool call."""
name: ReadOnly[object]
arguments: ReadOnly[object]
class _NormalizedToolCall(TypedDict):
"""A tool call reduced to the stable shape the OpenInference emitters read."""
id: ReadOnly[object]
type: ReadOnly[object]
function: ReadOnly[_NormalizedToolCallFunction]
def _normalize_tool_call(raw_tc) -> _NormalizedToolCall | None:
"""Normalize a single tool_call (dict or Pydantic) into a stable shape:
{"id": str|None, "type": str, "function": {"name": str|None, "arguments": str|None}}

172
litellm/integrations/azure_sentinel/azure_sentinel.py
View file

@ -16,23 +16,34 @@ import asyncio
import os
import time
import traceback
from collections.abc import Mapping
from collections.abc import Mapping, Sequence
from types import MappingProxyType
from typing import Final
from typing import Final, TypeVar
from urllib.parse import urlparse
import httpx
from litellm._logging import verbose_logger
from litellm.integrations.batch_utils import (
BatchSendCancelled,
send_batch_with_413_split,
undelivered_after_http_error,
)
from litellm.integrations.custom_batch_logger import CustomBatchLogger
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
from litellm.llms.custom_httpx.http_handler import (
MaskedHTTPStatusError,
get_async_httpx_client,
httpxSpecialProvider,
)
from litellm.types.integrations.azure_sentinel import AZURE_SENTINEL_MAX_PAYLOAD_SIZE_BYTES
from litellm.types.utils import StandardAuditLogPayload, StandardLoggingPayload
DEFAULT_AZURE_AUTHORITY_HOST: Final = "https://login.microsoftonline.com"
DEFAULT_AZURE_MONITOR_SCOPE: Final = "https://monitor.azure.com/.default"
_QueuedPayload = TypeVar("_QueuedPayload", StandardLoggingPayload, StandardAuditLogPayload)
MONITOR_SCOPE_BY_AUTHORITY_HOST: Final[Mapping[str, str]] = MappingProxyType(
{
"login.microsoftonline.com": DEFAULT_AZURE_MONITOR_SCOPE,
@ -153,6 +164,8 @@ class AzureSentinelLogger(CustomBatchLogger):
asyncio.create_task(self.periodic_flush())
self.log_queue: list[StandardLoggingPayload] = []
self.audit_log_queue: list[StandardAuditLogPayload] = []
self.logs_awaiting_retry = False
self.audit_logs_awaiting_retry = False
@staticmethod
def _normalize_authority_host(authority_host: str) -> str:
@ -245,8 +258,8 @@ class AzureSentinelLogger(CustomBatchLogger):
self.log_queue.append(standard_logging_payload)
if len(self.log_queue) >= self.batch_size:
await self.async_send_batch()
if len(self.log_queue) >= self.batch_size and not self.logs_awaiting_retry:
await self._threshold_send_logs()
except Exception as e:
verbose_logger.exception("Azure Sentinel Layer Error - %s\n%s", e, traceback.format_exc())
@ -275,8 +288,8 @@ class AzureSentinelLogger(CustomBatchLogger):
self.log_queue.append(standard_logging_payload)
if len(self.log_queue) >= self.batch_size:
await self.async_send_batch()
if len(self.log_queue) >= self.batch_size and not self.logs_awaiting_retry:
await self._threshold_send_logs()
except Exception as e:
verbose_logger.exception("Azure Sentinel Layer Error - %s\n%s", e, traceback.format_exc())
@ -298,12 +311,24 @@ class AzureSentinelLogger(CustomBatchLogger):
self.audit_log_queue.append(audit_log)
if len(self.audit_log_queue) >= self.batch_size:
await self.async_send_audit_batch()
if len(self.audit_log_queue) >= self.batch_size and not self.audit_logs_awaiting_retry:
await self._threshold_send_audit_logs()
except Exception as e:
verbose_logger.exception("Azure Sentinel Audit Log Layer Error - %s\n%s", e, traceback.format_exc())
async def _threshold_send_logs(self) -> None:
async with self.flush_lock:
if self.logs_awaiting_retry or len(self.log_queue) < self.batch_size:
return
await self.async_send_batch()
async def _threshold_send_audit_logs(self) -> None:
async with self.flush_lock:
if self.audit_logs_awaiting_retry or len(self.audit_log_queue) < self.batch_size:
return
await self.async_send_audit_batch()
async def async_send_batch(self):
"""
Sends the batch of logs to Azure Monitor Logs Ingestion API
@ -311,67 +336,110 @@ class AzureSentinelLogger(CustomBatchLogger):
Raises:
Raises a NON Blocking verbose_logger.exception if an error occurs
"""
await self._async_send_batch_to_api(
log_queue=self.log_queue,
api_endpoint=self.api_endpoint,
log_type="logs",
)
batch_to_send: Final = tuple(self.log_queue)
self.log_queue = [] # mutable-ok: queue ownership is detached before the async send
try:
undelivered: Final = await self._async_send_batch_to_api(
log_queue=batch_to_send,
api_endpoint=self.api_endpoint,
log_type="logs",
)
except BatchSendCancelled as cancelled:
self.log_queue = self._requeue(cancelled.undelivered, self.log_queue, "logs")
self.logs_awaiting_retry = bool(self.log_queue)
raise asyncio.CancelledError() from cancelled
except asyncio.CancelledError:
self.log_queue = self._requeue(batch_to_send, self.log_queue, "logs")
self.logs_awaiting_retry = bool(self.log_queue)
raise
self.log_queue = self._requeue(undelivered, self.log_queue, "logs")
self.logs_awaiting_retry = bool(undelivered) and bool(self.log_queue)
async def async_send_audit_batch(self):
"""
Sends the batch of audit logs to Azure Monitor Logs Ingestion API
"""
await self._async_send_batch_to_api(
log_queue=self.audit_log_queue,
api_endpoint=self.audit_api_endpoint,
log_type="audit logs",
batch_to_send: Final = tuple(self.audit_log_queue)
self.audit_log_queue = [] # mutable-ok: queue ownership is detached before the async send
try:
undelivered: Final = await self._async_send_batch_to_api(
log_queue=batch_to_send,
api_endpoint=self.audit_api_endpoint,
log_type="audit logs",
)
except BatchSendCancelled as cancelled:
self.audit_log_queue = self._requeue(cancelled.undelivered, self.audit_log_queue, "audit logs")
self.audit_logs_awaiting_retry = bool(self.audit_log_queue)
raise asyncio.CancelledError() from cancelled
except asyncio.CancelledError:
self.audit_log_queue = self._requeue(batch_to_send, self.audit_log_queue, "audit logs")
self.audit_logs_awaiting_retry = bool(self.audit_log_queue)
raise
self.audit_log_queue = self._requeue(undelivered, self.audit_log_queue, "audit logs")
self.audit_logs_awaiting_retry = bool(undelivered) and bool(self.audit_log_queue)
def _requeue(
self,
undelivered: tuple[_QueuedPayload, ...],
queue: list[_QueuedPayload],
log_type: str,
) -> list[_QueuedPayload]:
merged: Final = [*undelivered, *queue] # mutable-ok: queue trimming returns a mutable logger queue
overflow: Final = len(merged) - self.max_queue_size
if overflow <= 0:
return merged
verbose_logger.warning(
"Azure Sentinel: %s queue exceeded max_queue_size=%s, dropped %s oldest records",
log_type,
self.max_queue_size,
overflow,
)
return merged[overflow:]
async def _async_send_batch_to_api(
self,
log_queue: list[StandardLoggingPayload | StandardAuditLogPayload],
log_queue: tuple[_QueuedPayload, ...],
api_endpoint: str,
log_type: str,
) -> None:
) -> tuple[_QueuedPayload, ...]:
if not log_queue:
return ()
verbose_logger.debug("Azure Sentinel - about to flush %s %s", len(log_queue), log_type)
try:
if not log_queue:
return
verbose_logger.debug("Azure Sentinel - about to flush %s %s", len(log_queue), log_type)
# Get OAuth2 token
bearer_token: Final = await self._get_oauth_token()
except MaskedHTTPStatusError as e:
return undelivered_after_http_error(log_queue, e.status_code, "Azure Sentinel OAuth token", str(e))
except Exception as e:
verbose_logger.exception("Azure Sentinel Error getting OAuth token - %s", e)
return tuple(log_queue)
# Convert log queue to JSON array format expected by Logs Ingestion API
# Each log entry should be a JSON object in the array
body: Final = safe_dumps(log_queue)
headers: Final = {
"Authorization": f"Bearer {bearer_token}",
"Content-Type": "application/json",
}
# Set headers for Logs Ingestion API
headers: Final = {
"Authorization": f"Bearer {bearer_token}",
"Content-Type": "application/json",
}
# Send the request
response = await self.async_httpx_client.post(url=api_endpoint, data=body.encode("utf-8"), headers=headers)
if response.status_code not in [200, 204]:
verbose_logger.error(
"Azure Sentinel API error: status_code=%s, response=%s",
response.status_code,
response.text,
)
raise Exception(f"Failed to send logs to Azure Sentinel: {response.status_code} - {response.text}")
verbose_logger.debug(
"Azure Sentinel: Response from API status_code: %s",
response.status_code,
async def _send_batch(batch: Sequence[_QueuedPayload]) -> httpx.Response:
body: Final = safe_dumps(batch)
return await self.async_httpx_client.post(
url=api_endpoint,
data=body.encode("utf-8"),
headers=headers,
)
except Exception as e:
verbose_logger.exception("Azure Sentinel Error sending batch API - %s\n%s", e, traceback.format_exc())
finally:
log_queue.clear()
return await send_batch_with_413_split(
batch=log_queue,
send_batch=_send_batch,
exceeds_limits=lambda batch: (
len(batch) > self.batch_size
or len(safe_dumps(batch).encode("utf-8")) > AZURE_SENTINEL_MAX_PAYLOAD_SIZE_BYTES
),
success_status_codes=frozenset({200, 204}),
integration_name="Azure Sentinel",
drop_error_message="Azure Sentinel API Error - Payload too large for a single record",
non_success_handler=undelivered_after_http_error,
)
async def flush_queue(self):
if self.flush_lock is None:

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