scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-10 22:41:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 129

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

Browse source 
# Conflicts:
#	tests/test_litellm/test_router.py
This commit is contained in:
mateo-berri 2026-09-01 17:17:24 -07:00
commit 7ed2e8acff
1771 changed files with 77250 additions and 15740 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

23
.github/actions/cache-cargo-build/action.yml vendored
View file

@ -4,17 +4,16 @@ description: >-
so only the first job on a given Cargo.lock compiles the bridge from scratch.
litellm builds through maturin, which compiles litellm-rust/crates/python-bridge
in release mode before it can produce a wheel. `uv sync` therefore pays a full
build in every job that installs the workspace: measured at 2m40s per unit shard
on 2026-08-21, more than the whole unit tier spends running tests. Nothing caught
it, because the uv cache holds wheels uv downloads rather than wheels it builds,
and a path dependency whose source moves every commit could never hit that cache
anyway. Cargo rebuilds only what changed when its target directory survives, so a
warm job pays for the bridge crate alone.
in the dev profile for editable installs. `uv sync` therefore pays a full build
in every job that installs the workspace. Nothing caught it, because the uv cache
holds wheels uv downloads rather than wheels it builds, and a path dependency
whose source moves every commit could never hit that cache anyway. Cargo rebuilds
only what changed when its target directory survives, so a warm job pays for the
bridge crate alone.
The key namespace is separate from test-rust.yml's. Both cache the same directory,
but that workflow fills it with debug and clippy artifacts, which a release build
cannot reuse, and a shared key would let whichever ran first deny the other a save.
The key namespace is separate from test-rust.yml's check and release caches. They
cache the same directory for different workloads, and a shared key would let
whichever ran first deny the others a save.
runs:
using: composite
@ -26,6 +25,6 @@ runs:
~/.cargo/registry
~/.cargo/git
litellm-rust/target
key: ${{ runner.os }}-cargo-release-${{ hashFiles('litellm-rust/Cargo.lock') }}
key: ${{ runner.os }}-maturin-dev-${{ hashFiles('litellm-rust/Cargo.lock') }}
restore-keys: |
${{ runner.os }}-cargo-release-
${{ runner.os }}-maturin-dev-

230
.github/scripts/close_duplicate_issues.py vendored
View file

@ -1,230 +0,0 @@
#!/usr/bin/env python3
"""
Detect and close duplicate GitHub issues using title similarity.
Modes:
--scan Compare all open issues against each other (batch)
--issue-number N Check a single issue against older open issues
Requires the `gh` CLI to be authenticated.
"""
import argparse
import difflib
import json
import re
import subprocess
import sys
def normalize_title(title: str) -> str:
"""Strip common prefixes, lowercase, and collapse whitespace."""
title = re.sub(
r"^\[?(bug|feature request|enhancement|question|docs)[:\]]?\s*",
"",
title,
flags=re.IGNORECASE,
)
return " ".join(title.lower().split())
def gh(*args: str) -> str:
"""Run a gh CLI command and return stdout."""
result = subprocess.run(
["gh", *args],
capture_output=True,
text=True,
check=True,
)
return result.stdout
def fetch_open_issues(repo: str | None) -> list[dict]:
"""Fetch all open issues (excluding PRs) via gh api --paginate."""
if repo:
endpoint = (
f"repos/{repo}/issues?state=open&per_page=100&sort=created&direction=asc"
)
else:
endpoint = "repos/{owner}/{repo}/issues?state=open&per_page=100&sort=created&direction=asc"
cmd = ["api", "--paginate", endpoint]
raw = gh(*cmd)
# gh --paginate concatenates JSON arrays, so we may get multiple arrays
issues = []
for line in raw.strip().splitlines():
line = line.strip()
if not line:
continue
parsed = json.loads(line)
if isinstance(parsed, list):
issues.extend(parsed)
else:
issues.append(parsed)
# Filter out pull requests (they also appear in the issues endpoint)
return [i for i in issues if "pull_request" not in i]
def close_as_duplicate(
issue_number: int, duplicate_of: int, repo: str | None, dry_run: bool
) -> None:
"""Close an issue as duplicate of another, adding a comment and label."""
repo_args = ["--repo", repo] if repo else []
if dry_run:
print(
f" [DRY RUN] Would close #{issue_number} as duplicate of #{duplicate_of}"
)
return
# Add comment
comment_body = (
f"Closing as duplicate of #{duplicate_of}.\n\n"
"If you believe this is not a duplicate, please reopen and add context "
"explaining how this differs."
)
gh("issue", "comment", str(issue_number), "--body", comment_body, *repo_args)
# Add label
gh("issue", "edit", str(issue_number), "--add-label", "duplicate", *repo_args)
# Close with not_planned reason
gh(
"api",
f"repos/{repo or '{owner}/{repo}'}/issues/{issue_number}",
"-X",
"PATCH",
"-f",
"state=closed",
"-f",
"state_reason=not_planned",
)
print(f" Closed #{issue_number} as duplicate of #{duplicate_of}")
def find_duplicate(
issue: dict, candidates: list[dict], threshold: float
) -> dict | None:
"""Return the first candidate whose normalized title is above threshold."""
norm = normalize_title(issue["title"])
for candidate in candidates:
if candidate["number"] == issue["number"]:
continue
cand_norm = normalize_title(candidate["title"])
ratio = difflib.SequenceMatcher(None, norm, cand_norm).ratio()
if ratio >= threshold:
return candidate
return None
def scan_all(
issues: list[dict], threshold: float, repo: str | None, dry_run: bool
) -> int:
"""Compare every issue against all older issues. Returns count of duplicates found."""
# Sort oldest first
issues.sort(key=lambda i: i["number"])
closed_count = 0
for idx, issue in enumerate(issues):
older = issues[:idx]
if not older:
continue
dup = find_duplicate(issue, older, threshold)
if dup:
ratio = difflib.SequenceMatcher(
None,
normalize_title(issue["title"]),
normalize_title(dup["title"]),
).ratio()
print(
f"#{issue['number']}: \"{issue['title']}\"\n"
f" -> duplicate of #{dup['number']}: \"{dup['title']}\" "
f"({ratio:.0%} similar)"
)
close_as_duplicate(issue["number"], dup["number"], repo, dry_run)
closed_count += 1
return closed_count
def check_single(
issue_number: int,
issues: list[dict],
threshold: float,
repo: str | None,
dry_run: bool,
) -> bool:
"""Check a single issue against all older open issues. Returns True if duplicate found."""
target = None
for i in issues:
if i["number"] == issue_number:
target = i
break
if target is None:
print(f"Issue #{issue_number} not found among open issues.")
return False
older = [i for i in issues if i["number"] < issue_number]
dup = find_duplicate(target, older, threshold)
if dup:
ratio = difflib.SequenceMatcher(
None,
normalize_title(target["title"]),
normalize_title(dup["title"]),
).ratio()
print(
f"#{target['number']}: \"{target['title']}\"\n"
f" -> duplicate of #{dup['number']}: \"{dup['title']}\" "
f"({ratio:.0%} similar)"
)
close_as_duplicate(issue_number, dup["number"], repo, dry_run)
return True
print(f"#{issue_number}: no duplicate found above threshold {threshold}")
return False
def main() -> None:
parser = argparse.ArgumentParser(
description="Detect and close duplicate GitHub issues"
)
mode = parser.add_mutually_exclusive_group(required=True)
mode.add_argument("--scan", action="store_true", help="Scan all open issues")
mode.add_argument("--issue-number", type=int, help="Check a single issue number")
parser.add_argument(
"--threshold", type=float, default=0.85, help="Similarity threshold (0-1)"
)
parser.add_argument(
"--close",
action="store_true",
help="Actually close duplicates (default is dry-run)",
)
parser.add_argument(
"--repo", type=str, help="Repository (owner/repo). Auto-detected if omitted."
)
args = parser.parse_args()
dry_run = not args.close
if dry_run:
print("=== DRY RUN MODE (pass --close to actually close issues) ===\n")
print("Fetching open issues...")
issues = fetch_open_issues(args.repo)
print(f"Found {len(issues)} open issues.\n")
if args.scan:
count = scan_all(issues, args.threshold, args.repo, dry_run)
print(f"\nTotal duplicates {'found' if dry_run else 'closed'}: {count}")
else:
found = check_single(
args.issue_number, issues, args.threshold, args.repo, dry_run
)
sys.exit(0 if found else 0) # Always exit 0; finding no dup is not an error
if __name__ == "__main__":
main()

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

@ -0,0 +1,69 @@
name: Auto-close duplicate issues
on:
schedule:
- cron: "0 9 * * *"
workflow_dispatch:
inputs:
dry_run:
description: Log which issues would close without closing anything
type: boolean
default: true
grace_period_days:
description: Days a duplicate notice must go unanswered before the close
type: number
default: 3
pull_request:
paths:
- .github/workflows/auto-close-duplicates.yml
- scripts/auto-close-duplicates.ts
- scripts/auto-close-duplicates.test.ts
permissions: {}
jobs:
test:
if: github.event_name == 'pull_request'
runs-on: ubuntu-latest
timeout-minutes: 5
permissions:
contents: read
steps:
- name: Checkout repository
uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
persist-credentials: false
- name: Setup Bun
uses: oven-sh/setup-bun@0c5077e51419868618aeaa5fe8019c62421857d6 # v2.2.0
with:
bun-version: "1.4.0"
- name: Test the sweep
run: bun test scripts/auto-close-duplicates.test.ts
sweep:
if: github.event_name != 'pull_request' && github.repository == 'BerriAI/litellm'
runs-on: ubuntu-latest
timeout-minutes: 10
permissions:
contents: read
issues: write
steps:
- name: Checkout repository
uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
persist-credentials: false
- name: Setup Bun
uses: oven-sh/setup-bun@0c5077e51419868618aeaa5fe8019c62421857d6 # v2.2.0
with:
# Exact version, never latest: the next step holds an issues: write token
bun-version: "1.4.0"
- name: Close unanswered duplicates, reopen ones the reporter answered
run: bun run scripts/auto-close-duplicates.ts
env:
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
DRY_RUN: ${{ inputs.dry_run == true }}
GRACE_PERIOD_DAYS: ${{ inputs.grace_period_days }}

40
.github/workflows/check_duplicate_issues.yml vendored
View file

@ -1,12 +1,19 @@
name: Check Duplicate Issues
# Flagging only. "Auto-close duplicate issues" closes a flagged issue 3 days later,
# and only when its title is identical to an older open issue and nobody replied.
# The HTML marker below is the handshake between the two, so keep it in the template.
on:
issues:
types: [opened, edited]
permissions: {}
jobs:
check-duplicate:
runs-on: ubuntu-latest
timeout-minutes: 5
permissions:
issues: write
contents: read
@ -19,35 +26,12 @@ jobs:
threshold: 0.6
reaction: eyes
comment: |
**⚠️ Potential duplicate detected**
<!-- litellm:potential-duplicate candidates={{#issues}}{{number}},{{/issues}} -->
**Potential duplicate detected**
This issue appears similar to existing issue(s):
This looks similar to:
{{#issues}}
- [#{{number}}]({{html_url}}) - {{title}} ({{accuracy}}% similar)
- #{{number}} - {{title}}
{{/issues}}
Please review the linked issue(s) to see if they address your concern. If this is not a duplicate, please provide additional context to help us understand the difference.
- name: Checkout close script
if: github.event.action == 'opened'
uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
sparse-checkout: .github/scripts
persist-credentials: false
- name: Set up Python
if: github.event.action == 'opened'
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
with:
python-version: "3.12"
- name: Auto-close if high-confidence duplicate
if: github.event.action == 'opened'
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
run: |
python3 .github/scripts/close_duplicate_issues.py \
--issue-number ${{ github.event.issue.number }} \
--repo ${{ github.repository }} \
--threshold 0.85 \
--close
If this is a duplicate, add a thumbs-up reaction to the existing issue and follow along there. When the title is identical to an older open issue, this issue closes automatically in 3 days unless someone responds. If it is not a duplicate, comment here or add a thumbs-down reaction to this comment and it stays open.

22
.github/workflows/codspeed.yml vendored
View file

@ -12,6 +12,7 @@ on:
- "uv.lock"
- ".github/workflows/codspeed.yml"
- ".github/actions/setup-uv-with-retries/**"
- ".github/actions/cache-cargo-build/**"
pull_request:
branches:
- main
@ -23,6 +24,7 @@ on:
- "uv.lock"
- ".github/workflows/codspeed.yml"
- ".github/actions/setup-uv-with-retries/**"
- ".github/actions/cache-cargo-build/**"
# Allow CodSpeed to trigger backtest performance analysis
# in order to generate initial data
workflow_dispatch:
@ -55,6 +57,26 @@ jobs:
with:
version: "0.10.9"
- name: Cache the Rust build
uses: ./.github/actions/cache-cargo-build
# Build the wheel and resolve every dependency outside the CodSpeed
# runner: the same maturin build took 42 minutes inside `codspeed run`
# versus under 3 minutes as a plain step (LIT-6183)
- name: Build environment
run: >
env PYTEST_DISABLE_PLUGIN_AUTOLOAD=1
uv run --frozen --no-default-groups
--with pytest==8.3.5
--with pytest-codspeed==4.3.0
--with "mcp>=1.26.0,<2.0"
--with "a2a-sdk>=1.1.0,<2.0"
pytest
-p pytest_codspeed.plugin
tests/benchmarks/
--codspeed
--collect-only -q
- name: Run benchmarks
uses: CodSpeedHQ/action@1c8ae4843586d3ba879736b7f6b7b0c990757fab # v4.12.1
with:

77
.github/workflows/test-redis-compat.yml vendored Normal file
View file

@ -0,0 +1,77 @@
name: "Unit Tests: Redis Client Version Compatibility"
on:
pull_request:
branches:
- main
- litellm_internal_staging
- litellm_oss_staging
- "litellm_**"
paths:
- "litellm/_redis.py"
- "litellm/_redis_credential_provider.py"
- "tests/test_litellm/test_redis.py"
- "tests/test_litellm/caching/test_redis_connection_pool.py"
- ".github/workflows/test-redis-compat.yml"
- "pyproject.toml"
- "uv.lock"
permissions:
contents: read
concurrency:
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
cancel-in-progress: true
jobs:
redis-compat:
name: "redis-py ${{ matrix.redis-version }}"
runs-on: ubuntu-latest
timeout-minutes: 15
strategy:
fail-fast: false
matrix:
# 5.3.1 is the version pinned in uv.lock (redisvl caps it below 6); the
# newer legs prove the inspect.signature introspection in litellm/_redis.py
# keeps extracting kwargs on the redis-py releases people actually run now.
# Only the exact release 6.0.0 is skipped: rq (pulled by the proxy extra)
# specifies `redis != 6`, which excludes 6.0.0 alone, so 6.4.0 stands in
# for the 6.x line.
redis-version: ["5.3.1", "6.4.0", "7.4.1", "8.0.1"]
steps:
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
persist-credentials: false
- name: Set up Python
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
with:
python-version: "3.12"
- name: Set up uv
uses: ./.github/actions/setup-uv-with-retries
with:
version: "0.10.9"
- name: Install dependencies
run: |
.github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router
- name: Pin redis-py to the matrix version
env:
REDIS_VERSION: ${{ matrix.redis-version }}
run: |
uv pip install "redis==${REDIS_VERSION:?}"
uv run --no-sync python -c "import redis; assert redis.__version__ == '${REDIS_VERSION:?}', redis.__version__; print('redis-py', redis.__version__)"
- name: Run redis unit tests
run: |
uv run --no-sync pytest \
tests/test_litellm/test_redis.py \
tests/test_litellm/caching/test_redis_connection_pool.py \
--tb=short -vv \
--reruns 2 \
--reruns-delay 1 \
--durations=20

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

@ -103,6 +103,7 @@ jobs:
tests/test_litellm/completion_extras
tests/test_litellm/compression
tests/test_litellm/containers
tests/test_litellm/endpoints
tests/test_litellm/experimental_mcp_client
tests/test_litellm/models
tests/test_litellm/repositories

2
CLAUDE.md
View file

@ -23,6 +23,8 @@ When adding new features, add meaningful tests. Don't add tests that don't check
Same thing for bug fixes. The tests should make it so that this specific bug can never happen again without failing tests (i.e., regression)
Never test structure of code only function of it
`tests/test_litellm/` mirrors `litellm/` in a parallel path (see `tests/test_litellm/readme.md`). Name tests `test_<filename>.py`, but always match the existing test file in the directory you touch — many provider dirs use longer descriptive names (e.g. `test_anthropic_chat_transformation.py`) to avoid ambiguity across sibling folders. For bug fixes, extend the existing mapped test file rather than creating a new one. Only create a new test file for a new feature (provider, endpoint, or transformation module) that has no mapped test yet, following that directory's naming convention (or `test_<filename>.py` if you're the first test there). One focused regression test beats many shallow ones
End-to-end tests belong in `tests/e2e/` and must follow the harness conventions documented in that directory's `CLAUDE.md`

21
Dockerfile
View file

@ -1,10 +1,10 @@
# syntax=docker/dockerfile:1.7
# Base image for building
ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:a31344ab2cb8618db84f535eec56f76f6178b142cb92cb2e48676cc2dcebea72
ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:e624c5d5e42382ce7165ddafcbbf8e6769a24cbd02ea6114b880b05ae5ba2a8d
# Runtime image
ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:a31344ab2cb8618db84f535eec56f76f6178b142cb92cb2e48676cc2dcebea72
ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:e624c5d5e42382ce7165ddafcbbf8e6769a24cbd02ea6114b880b05ae5ba2a8d
ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a
# Pinned by digest like the other base images; bump explicitly on Node upgrades.
ARG UI_BUILD_IMAGE=node:24.19-alpine3.24@sha256:d32cdf619f63fe0471182d08996dd516c6275bb5fd31ae06e55a570bd9e1ad43
@ -40,8 +40,8 @@ COPY --from=uvbin /uvx /usr/local/bin/uvx
RUN apk add --no-cache \
bash \
gcc \
python3 \
python3-dev \
python-3.13 \
python-3.13-dev \
rust \
openssl \
openssl-dev \
@ -51,6 +51,7 @@ RUN apk add --no-cache \
ENV UV_PROJECT_ENVIRONMENT=/app/.venv \
UV_LINK_MODE=copy \
UV_PYTHON_DOWNLOADS=0 \
PATH="/app/.venv/bin:${PATH}"
# Copy dependency metadata first for layer caching
@ -65,7 +66,7 @@ RUN uv sync --frozen --no-install-project --no-install-workspace --no-default-gr
--extra extra_proxy \
--extra semantic-router \
--extra saml \
--python python3
--python python3.13
# Copy full source tree
COPY . .
@ -86,7 +87,7 @@ RUN uv sync --frozen --no-default-groups --no-editable \
--extra extra_proxy \
--extra semantic-router \
--extra saml \
--python python3
--python python3.13
RUN HOME=/opt/prisma XDG_CACHE_HOME=/opt/prisma/.cache PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries \
npm_config_cache=/root/.npm \
@ -100,8 +101,14 @@ FROM $LITELLM_RUNTIME_IMAGE AS runtime
USER root
# The base image only configures Chainguard's authenticated apk repo, which
# requires an enterprise subscription. Add the public Wolfi repo so `apk add`
# also works for anyone installing extra packages into a running container.
# https://github.com/BerriAI/litellm/issues/33518
RUN echo "https://packages.wolfi.dev/os" >> /etc/apk/repositories
# node (without npm) is required by the prisma CLI at runtime
RUN apk add --no-cache bash openssl tzdata nodejs python3 libsndfile
RUN apk add --no-cache bash openssl tzdata nodejs python-3.13 libsndfile
WORKDIR /app
ENV PATH="/app/.venv/bin:${PATH}" \

2
README.md
View file

@ -354,6 +354,8 @@ curl -X POST 'http://0.0.0.0:4000/v1/chat/completions' \
| [Petals (`petals`)](https://docs.litellm.ai/docs/providers/petals) | ✅ | ✅ | ✅ | | | | | | | |
| [Pinstripes (`pinstripes`)](https://docs.litellm.ai/docs/providers/pinstripes) | ✅ | ✅ | ✅ | | | | | | | |
| [Predibase (`predibase`)](https://docs.litellm.ai/docs/providers/predibase) | ✅ | ✅ | ✅ | | | | | | | |
| [Qwen AI Platform (`qwen_ai_platform`)](https://docs.litellm.ai/docs/providers/qwencloud) | ✅ | ✅ | ✅ | ✅ | ✅ | | | | | ✅ |
| [QwenCloud (`qwencloud`)](https://docs.litellm.ai/docs/providers/qwencloud) | ✅ | ✅ | ✅ | ✅ | ✅ | | | | | ✅ |
| [Recraft (`recraft`)](https://docs.litellm.ai/docs/providers/recraft) | | | | | ✅ | | | | | |
| [Replicate (`replicate`)](https://docs.litellm.ai/docs/providers/replicate) | ✅ | ✅ | ✅ | | | | | | | |
| [Sagemaker Chat (`sagemaker_chat`)](https://docs.litellm.ai/docs/providers/aws_sagemaker) | ✅ | ✅ | ✅ | | | | | | | |

12
backend/Dockerfile
View file

@ -1,5 +1,5 @@
ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:a31344ab2cb8618db84f535eec56f76f6178b142cb92cb2e48676cc2dcebea72
ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:a31344ab2cb8618db84f535eec56f76f6178b142cb92cb2e48676cc2dcebea72
ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:e624c5d5e42382ce7165ddafcbbf8e6769a24cbd02ea6114b880b05ae5ba2a8d
ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:e624c5d5e42382ce7165ddafcbbf8e6769a24cbd02ea6114b880b05ae5ba2a8d
ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a
FROM $UV_IMAGE AS uvbin
@ -16,7 +16,7 @@ COPY --from=uvbin /uv /uvx /usr/local/bin/
# instead of nodeenv downloading one whose dynamic deps may not be in Wolfi
# (e.g. Node 26.2.0 needs libatomic). Retry for transient apk.cgr.dev flakes.
RUN for i in 1 2 3; do \
apk add --no-cache bash gcc python3 python3-dev openssl openssl-dev libsndfile nodejs npm && break; \
apk add --no-cache bash gcc python-3.13 python-3.13-dev openssl openssl-dev libsndfile nodejs npm && break; \
[ $i = 3 ] && { echo "apk add failed after 3 retries" >&2; exit 1; }; \
sleep 5; \
done
@ -46,7 +46,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
--extra proxy-runtime \
--extra extra_proxy \
--extra semantic-router \
--python python3
--python python3.13
# Stage 2 — copy source and install the project + workspace members.
COPY . .
@ -57,7 +57,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
--extra proxy-runtime \
--extra extra_proxy \
--extra semantic-router \
--python python3
--python python3.13
RUN HOME=/opt/prisma XDG_CACHE_HOME=/opt/prisma/.cache PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries \
npm_config_cache=/root/.npm \
@ -71,7 +71,7 @@ FROM $LITELLM_RUNTIME_IMAGE AS runtime
USER root
RUN for i in 1 2 3; do \
apk add --no-cache bash openssl tzdata python3 libsndfile libatomic && break; \
apk add --no-cache bash openssl tzdata python-3.13 libsndfile libatomic && break; \
[ $i = 3 ] && { echo "apk add failed after 3 retries" >&2; exit 1; }; \
sleep 5; \
done

42
basedpyright-code-budget.json
View file

@ -1,9 +1,9 @@
{
"reportAny": {
"limit": 18483
"limit": 14076
},
"reportArgumentType": {
"limit": 2557
"limit": 2216
},
"reportAssignmentType": {
"limit": 319
@ -12,25 +12,25 @@
"limit": 480
},
"reportCallIssue": {
"limit": 113
"limit": 112
},
"reportConstantRedefinition": {
"limit": 40
},
"reportDeprecated": {
"limit": 213
"limit": 211
},
"reportDuplicateImport": {
"limit": 19
},
"reportExplicitAny": {
"limit": 5960
"limit": 4128
},
"reportFunctionMemberAccess": {
"limit": 7
},
"reportGeneralTypeIssues": {
"limit": 105
"limit": 101
},
"reportIncompatibleMethodOverride": {
"limit": 56
@ -42,7 +42,7 @@
"limit": 12
},
"reportIndexIssue": {
"limit": 35
"limit": 25
},
"reportInvalidTypeForm": {
"limit": 34
@ -54,10 +54,10 @@
"limit": 0
},
"reportMissingParameterType": {
"limit": 5659
"limit": 5601
},
"reportMissingTypeArgument": {
"limit": 15482
"limit": 15306
},
"reportMissingTypeStubs": {
"limit": 40
@ -72,7 +72,7 @@
"limit": 0
},
"reportOptionalMemberAccess": {
"limit": 1058
"limit": 0
},
"reportOptionalOperand": {
"limit": 0
@ -90,40 +90,40 @@
"limit": 8
},
"reportReturnType": {
"limit": 213
"limit": 181
},
"reportTypedDictNotRequiredAccess": {
"limit": 26
"limit": 24
},
"reportUndefinedVariable": {
"limit": 0
},
"reportUnknownArgumentType": {
"limit": 44526
"limit": 44364
},
"reportUnknownLambdaType": {
"limit": 109
},
"reportUnknownMemberType": {
"limit": 38779
"limit": 38350
},
"reportUnknownParameterType": {
"limit": 19827
"limit": 19626
},
"reportUnknownVariableType": {
"limit": 30348
"limit": 29890
},
"reportUnnecessaryCast": {
"limit": 117
"limit": 111
},
"reportUnnecessaryComparison": {
"limit": 697
"limit": 695
},
"reportUnnecessaryContains": {
"limit": 5
},
"reportUnnecessaryIsInstance": {
"limit": 831
"limit": 826
},
"reportUntypedBaseClass": {
"limit": 0
@ -138,9 +138,9 @@
"limit": 138
},
"reportUnusedImport": {
"limit": 544
"limit": 543
},
"reportUnusedVariable": {
"limit": 145
"limit": 137
}
}

2
codecov.yaml
View file

@ -25,6 +25,8 @@ flag_management:
carryforward: false
- name: proxy-db-schema-migration
carryforward: false
- name: circleci
carryforward: false
component_management:
individual_components:

15
docker/Dockerfile.database
View file

@ -1,10 +1,10 @@
# syntax=docker/dockerfile:1.7
# Base image for building
ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:a31344ab2cb8618db84f535eec56f76f6178b142cb92cb2e48676cc2dcebea72
ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:e624c5d5e42382ce7165ddafcbbf8e6769a24cbd02ea6114b880b05ae5ba2a8d
# Runtime image
ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:a31344ab2cb8618db84f535eec56f76f6178b142cb92cb2e48676cc2dcebea72
ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:e624c5d5e42382ce7165ddafcbbf8e6769a24cbd02ea6114b880b05ae5ba2a8d
ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a
# Pinned by digest like the other base images; bump explicitly on Node upgrades.
ARG UI_BUILD_IMAGE=node:24.19-alpine3.24@sha256:d32cdf619f63fe0471182d08996dd516c6275bb5fd31ae06e55a570bd9e1ad43
@ -39,8 +39,8 @@ COPY --from=uvbin /uvx /usr/local/bin/uvx
RUN apk add --no-cache \
bash \
gcc \
python3 \
python3-dev \
python-3.13 \
python-3.13-dev \
openssl \
openssl-dev \
nodejs \
@ -49,6 +49,7 @@ RUN apk add --no-cache \
ENV UV_PROJECT_ENVIRONMENT=/app/.venv \
UV_LINK_MODE=copy \
UV_PYTHON_DOWNLOADS=0 \
PATH="/app/.venv/bin:${PATH}"
# Copy dependency metadata first for layer caching
@ -63,7 +64,7 @@ RUN uv sync --frozen --no-install-project --no-install-workspace --no-default-gr
--extra extra_proxy \
--extra semantic-router \
--extra saml \
--python python3
--python python3.13
# Copy full source tree
COPY . .
@ -84,7 +85,7 @@ RUN uv sync --frozen --no-default-groups --no-editable \
--extra extra_proxy \
--extra semantic-router \
--extra saml \
--python python3
--python python3.13
RUN HOME=/opt/prisma XDG_CACHE_HOME=/opt/prisma/.cache PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries \
npm_config_cache=/root/.npm \
@ -98,7 +99,7 @@ FROM $LITELLM_RUNTIME_IMAGE AS runtime
USER root
# node (without npm) is required by the prisma CLI at runtime
RUN apk add --no-cache bash openssl tzdata nodejs python3 libsndfile
RUN apk add --no-cache bash openssl tzdata nodejs python-3.13 libsndfile
WORKDIR /app
ENV PATH="/app/.venv/bin:${PATH}" \

17
docker/Dockerfile.non_root
View file

@ -1,8 +1,8 @@
# syntax=docker/dockerfile:1.7
# Base images
ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:a31344ab2cb8618db84f535eec56f76f6178b142cb92cb2e48676cc2dcebea72
ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:a31344ab2cb8618db84f535eec56f76f6178b142cb92cb2e48676cc2dcebea72
ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:e624c5d5e42382ce7165ddafcbbf8e6769a24cbd02ea6114b880b05ae5ba2a8d
ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:e624c5d5e42382ce7165ddafcbbf8e6769a24cbd02ea6114b880b05ae5ba2a8d
ARG PROXY_EXTRAS_SOURCE=published
ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a
# Pinned by digest like the other base images; bump explicitly on Node upgrades.
@ -37,8 +37,8 @@ COPY --from=uvbin /uvx /usr/local/bin/uvx
RUN for i in 1 2 3; do \
apk add --no-cache \
python3 \
python3-dev \
python-3.13 \
python-3.13-dev \
gcc \
rust \
bash \
@ -52,6 +52,7 @@ RUN for i in 1 2 3; do \
ENV UV_PROJECT_ENVIRONMENT=/app/.venv \
UV_LINK_MODE=copy \
UV_PYTHON_DOWNLOADS=0 \
PATH="/app/.venv/bin:${PATH}" \
LITELLM_NON_ROOT=true \
XDG_CACHE_HOME=/app/.cache
@ -69,7 +70,7 @@ RUN --mount=type=cache,target=/app/.cache/uv,id=litellm-uv-cache \
--extra extra_proxy \
--extra semantic-router \
--extra saml \
--python python3
--python python3.13
# Copy full source tree
COPY . .
@ -96,7 +97,7 @@ RUN --mount=type=cache,target=/app/.cache/uv,id=litellm-uv-cache \
--extra extra_proxy \
--extra semantic-router \
--extra saml \
--python python3 \
--python python3.13 \
--no-sources-package litellm-proxy-extras; \
else \
uv sync --frozen --no-default-groups --no-editable \
@ -105,7 +106,7 @@ RUN --mount=type=cache,target=/app/.cache/uv,id=litellm-uv-cache \
--extra extra_proxy \
--extra semantic-router \
--extra saml \
--python python3; \
--python python3.13; \
fi
RUN HOME=/opt/prisma XDG_CACHE_HOME=/opt/prisma/.cache PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries \
@ -124,7 +125,7 @@ RUN for i in 1 2 3; do \
apk upgrade --no-cache && break || sleep 5; \
done && \
for i in 1 2 3; do \
apk add --no-cache python3 bash openssl tzdata libsndfile nodejs && break || sleep 5; \
apk add --no-cache python-3.13 bash openssl tzdata libsndfile nodejs && break || sleep 5; \
done
# Copy only what runtime needs. The application is installed inside the venv;

101
enterprise/litellm_enterprise/proxy/audit_logging_endpoints.py
View file

@ -7,7 +7,7 @@ GET - /audit/{id} - Get audit log by id
GET - /audit - Get all audit logs
"""
from typing import Any, Dict, List, Optional
from typing import TYPE_CHECKING, Final
#### AUDIT LOGGING ####
from fastapi import APIRouter, Depends, HTTPException, Query
@ -18,11 +18,16 @@ from litellm_enterprise.types.proxy.audit_logging_endpoints import (
from litellm.proxy._types import CommonProxyErrors, UserAPIKeyAuth
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.repositories.prisma_protocols import TableActions
from litellm.repositories.table_repositories import AuditLogRepository
if TYPE_CHECKING:
from prisma import models as prisma_models
router = APIRouter()
def _build_json_field_or_condition(json_key: str, value: str) -> Dict[str, Any]:
def _build_json_field_or_condition(json_key: str, value: str) -> dict[str, object]:
"""
Build an OR condition that matches a value inside a JSON column at the
given key, checking both before_value and updated_values.
@ -53,33 +58,33 @@ async def get_audit_logs(
page: int = Query(1, ge=1),
page_size: int = Query(10, ge=1, le=100),
# Filter parameters
changed_by: Optional[str] = Query(
changed_by: str | None = Query(
None, description="Filter by user or system that performed the action"
),
changed_by_api_key: Optional[str] = Query(
changed_by_api_key: str | None = Query(
None, description="Filter by API key hash that performed the action"
),
action: Optional[str] = Query(
action: str | None = Query(
None, description="Filter by action type (create, update, delete)"
),
table_name: Optional[str] = Query(
table_name: str | None = Query(
None, description="Filter by table name that was modified"
),
object_id: Optional[str] = Query(
object_id: str | None = Query(
None, description="Filter by ID of the object that was modified"
),
start_date: Optional[str] = Query(None, description="Filter logs after this date"),
end_date: Optional[str] = Query(None, description="Filter logs before this date"),
object_team_id: Optional[str] = Query(
start_date: str | None = Query(None, description="Filter logs after this date"),
end_date: str | None = Query(None, description="Filter logs before this date"),
object_team_id: str | None = Query(
None,
description="Filter by team_id present in before_value or updated_values JSON (PostgreSQL only)",
),
object_key_hash: Optional[str] = Query(
object_key_hash: str | None = Query(
None,
description="Filter by token (key hash) present in before_value or updated_values JSON (PostgreSQL only)",
),
# Sorting parameters
sort_by: Optional[str] = Query(
sort_by: str | None = Query(
None,
description="Column to sort by (e.g. 'updated_at', 'action', 'table_name')",
),
@ -101,46 +106,37 @@ async def get_audit_logs(
detail={"message": CommonProxyErrors.db_not_connected_error.value},
)
# Build filter conditions
where_conditions: Dict[str, Any] = {}
if changed_by:
where_conditions["changed_by"] = changed_by
if changed_by_api_key:
where_conditions["changed_by_api_key"] = changed_by_api_key
if action:
where_conditions["action"] = action
if table_name:
where_conditions["table_name"] = table_name
if object_id:
where_conditions["object_id"] = object_id
if start_date or end_date:
date_filter: Dict[str, Any] = {}
if start_date:
date_filter["gte"] = start_date
if end_date:
date_filter["lte"] = end_date
where_conditions["updated_at"] = date_filter
date_filter: Final[dict[str, str]] = {
**({"gte": start_date} if start_date else {}),
**({"lte": end_date} if end_date else {}),
}
# JSON field filters (PostgreSQL only) — each filter is AND'd with the
# others, but checks both before_value and updated_values internally (OR).
if object_team_id:
where_conditions["AND"] = where_conditions.get("AND", []) + [
_build_json_field_or_condition("team_id", object_team_id)
]
if object_key_hash:
where_conditions["AND"] = where_conditions.get("AND", []) + [
_build_json_field_or_condition("token", object_key_hash)
]
json_field_conditions: Final[list[dict[str, object]]] = [
*([_build_json_field_or_condition("team_id", object_team_id)] if object_team_id else []),
*([_build_json_field_or_condition("token", object_key_hash)] if object_key_hash else []),
]
# Build sort conditions
order_by: Dict[str, Any] = {}
if sort_by and isinstance(sort_by, str):
order_by[sort_by] = sort_order
else:
order_by["updated_at"] = sort_order # Default sort by updated_at
# Build filter conditions
where_conditions: Final[dict[str, object]] = {
**({"changed_by": changed_by} if changed_by else {}),
**({"changed_by_api_key": changed_by_api_key} if changed_by_api_key else {}),
**({"action": action} if action else {}),
**({"table_name": table_name} if table_name else {}),
**({"object_id": object_id} if object_id else {}),
**({"updated_at": date_filter} if start_date or end_date else {}),
**({"AND": json_field_conditions} if json_field_conditions else {}),
}
order_by: Final[dict[str, str]] = (
{sort_by: sort_order} if sort_by and isinstance(sort_by, str) else {"updated_at": sort_order}
)
audit_log_table: Final[TableActions["prisma_models.LiteLLM_AuditLog"]] = AuditLogRepository(prisma_client).table
# Get paginated results
audit_logs = await prisma_client.db.litellm_auditlog.find_many(
audit_logs: Final = await audit_log_table.find_many(
where=where_conditions,
order=order_by,
skip=(page - 1) * page_size,
@ -148,13 +144,14 @@ async def get_audit_logs(
)
# Get total count for pagination
total_count = await prisma_client.db.litellm_auditlog.count(where=where_conditions)
total_pages = -(-total_count // page_size) # Ceiling division
total_count: Final = await audit_log_table.count(where=where_conditions)
total_pages: Final = -(-total_count // page_size) # Ceiling division
# Return paginated response
return PaginatedAuditLogResponse(
audit_logs=[
AuditLogResponse(**audit_log.model_dump()) for audit_log in audit_logs
AuditLogResponse.model_validate(audit_log.model_dump())
for audit_log in audit_logs
]
if audit_logs
else [],
@ -198,8 +195,10 @@ async def get_audit_log_by_id(
detail={"message": CommonProxyErrors.db_not_connected_error.value},
)
audit_log_table: Final[TableActions["prisma_models.LiteLLM_AuditLog"]] = AuditLogRepository(prisma_client).table
# Get the audit log by ID
audit_log = await prisma_client.db.litellm_auditlog.find_unique(where={"id": id})
audit_log: Final = await audit_log_table.find_unique(where={"id": id})
if audit_log is None:
raise HTTPException(
@ -207,4 +206,4 @@ async def get_audit_log_by_id(
)
# Convert to response model
return AuditLogResponse(**audit_log.model_dump())
return AuditLogResponse.model_validate(audit_log.model_dump())

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

@ -2,9 +2,10 @@
Polls LiteLLM_ManagedObjectTable to check if the batch job is complete, and if the cost has been tracked.
"""
from dataclasses import replace as dataclasses_replace
from datetime import datetime, timedelta, timezone
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Dict, Final, List, Optional, Tuple
from typing import TYPE_CHECKING, Final, List, Literal, Optional, Tuple, cast
from litellm._logging import verbose_proxy_logger
from litellm._uuid import uuid
@ -86,7 +87,7 @@ class CheckBatchCost:
return
self.batch_processed_support_confirmed = True
async def _get_user_info(self, batch_id: str, user_id: Optional[str]) -> Dict[str, Any]:
async def _get_user_info(self, batch_id: str, user_id: Optional[str]) -> dict[str, str | None]:
"""
Look up user email and key alias by user_id for enriching the S3 callback metadata.
Returns a dict with user_api_key_user_email and user_api_key_alias (both may be None).
@ -96,8 +97,10 @@ class CheckBatchCost:
if not user_id:
return {}
try:
user_row = await self.prisma_client.db.litellm_usertable.find_unique(
where={"user_id": user_id}
user_row: prisma_models.LiteLLM_UserTable | None = (
await self.prisma_client.db.litellm_usertable.find_unique(
where={"user_id": user_id}
)
)
if user_row is None:
return {}
@ -114,8 +117,10 @@ class CheckBatchCost:
if not api_key:
return None
try:
key_row = await self.prisma_client.db.litellm_verificationtoken.find_unique(
where={"token": api_key}
key_row: prisma_models.LiteLLM_VerificationToken | None = (
await self.prisma_client.db.litellm_verificationtoken.find_unique(
where={"token": api_key}
)
)
return getattr(key_row, "key_alias", None) if key_row is not None else None
except Exception as e:
@ -127,8 +132,10 @@ class CheckBatchCost:
if not team_id:
return None
try:
team_row = await self.prisma_client.db.litellm_teamtable.find_unique(
where={"team_id": team_id}
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, "team_alias", None) if team_row is not None else None
except Exception as e:
@ -137,7 +144,7 @@ class CheckBatchCost:
async def _build_creator_attribution_metadata(
self, job: "LiteLLM_ManagedObjectTable", batch_id: str
) -> Dict[str, Any]:
) -> dict[str, object]:
"""
Rebuild the spend-tracking metadata for the key, team, and tags that created the
batch so the batch-cost spend log is attributed the same way a non-batch request
@ -151,7 +158,7 @@ class CheckBatchCost:
team_id = getattr(job, "team_id", None)
request_tags = getattr(job, "request_tags", None)
metadata: Dict[str, Any] = {
metadata: dict[str, object] = {
"user_api_key_user_id": job.created_by,
"user_api_key": api_key,
"user_api_key_team_id": team_id,
@ -626,6 +633,7 @@ class CheckBatchCost:
later poll.
"""
from litellm.batches.batch_utils import (
count_error_file_failed_requests,
_get_file_content_as_dictionary,
calculate_batch_cost_and_usage,
)
@ -761,16 +769,33 @@ class CheckBatchCost:
model_id=model_id,
deployment_model=litellm_model_name,
)
batch_cost, batch_usage, batch_models = (
await calculate_batch_cost_and_usage(
file_content_dictionary=file_content_as_dict,
custom_llm_provider=llm_provider, # type: ignore
model_name=model_name,
model_info=deployment_model_info,
batch_file_provider: Final = cast(
Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"], llm_provider
)
output_file_result: Final = await calculate_batch_cost_and_usage(
file_content_dictionary=file_content_as_dict,
custom_llm_provider=batch_file_provider,
model_name=model_name,
model_info=deployment_model_info,
)
error_file_failed_requests: Final = await count_error_file_failed_requests(
response,
custom_llm_provider=batch_file_provider,
litellm_params={
**credentials,
"_litellm_internal_model_credentials": MappingProxyType(dict(credentials)),
},
)
batch_result: Final = (
output_file_result
if not error_file_failed_requests
else dataclasses_replace(
output_file_result,
failed_requests=output_file_result.failed_requests + error_file_failed_requests,
)
)
logging_obj = LiteLLMLogging(
model=batch_models[0],
model=batch_result.models[0],
messages=[{"role": "user", "content": "<retrieve_batch>"}],
stream=False,
call_type="aretrieve_batch",
@ -802,9 +827,11 @@ class CheckBatchCost:
try:
await logging_obj.async_success_handler(
result=response,
batch_cost=batch_cost,
batch_usage=batch_usage,
batch_models=batch_models,
batch_cost=batch_result.cost,
batch_usage=batch_result.usage,
batch_models=batch_result.models,
batch_successful_requests=batch_result.successful_requests,
batch_failed_requests=batch_result.failed_requests,
)
except Exception:
await self._release_job_claim(job)

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

@ -36,6 +36,7 @@ from litellm.llms.base_llm.managed_resources.isolation import (
build_list_page,
build_owner_filter,
can_access_resource,
resolve_resource_owner_id,
)
from litellm.proxy._types import (
CallTypes,
@ -181,6 +182,10 @@ class _ManagedObjectTableActions(Protocol):
async def update_many(self, where: Mapping[str, object], data: Mapping[str, object]) -> int: ...
class _SchedulerWithJobLookup(Protocol):
def get_job(self, job_id: str) -> object: ...
class _CursorPageArgs(TypedDict, total=False):
cursor: Mapping[str, str]
skip: int
@ -222,7 +227,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
file_object=file_object,
model_mappings=model_mappings,
flat_model_file_ids=list(model_mappings.values()),
created_by=user_api_key_dict.user_id,
created_by=resolve_resource_owner_id(user_api_key_dict),
team_id=user_api_key_dict.team_id,
updated_by=user_api_key_dict.user_id,
)
@ -238,7 +243,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
"unified_file_id": file_id,
"model_mappings": json.dumps(model_mappings),
"flat_model_file_ids": list(model_mappings.values()),
"created_by": user_api_key_dict.user_id,
"created_by": resolve_resource_owner_id(user_api_key_dict),
"team_id": user_api_key_dict.team_id,
"updated_by": user_api_key_dict.user_id,
}
@ -342,7 +347,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
"file_object": file_object.model_dump_json(),
"model_object_id": model_object_id,
"file_purpose": file_purpose,
"created_by": user_api_key_dict.user_id,
"created_by": resolve_resource_owner_id(user_api_key_dict),
"team_id": user_api_key_dict.team_id,
"updated_by": user_api_key_dict.user_id,
"status": file_object.status,
@ -473,19 +478,56 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
)
page_size: Final = min(limit or 20, 100)
cursor_args: _CursorPageArgs = {"cursor": {"unified_object_id": after}, "skip": 1} if after else {}
batches = await _managed_object_table(self.prisma_client).find_many(
where=where_clause,
take=page_size + 1,
order=[{"created_at": "desc"}, {"unified_object_id": "desc"}],
**cursor_args,
matches: Final = await self._collect_listed_batches(
where_clause=where_clause,
after=after,
wanted=page_size + 1,
user_api_key_dict=user_api_key_dict,
)
return build_list_page(list(matches[:page_size]), has_more=len(matches) > page_size)
has_more = len(batches) > page_size
async def _collect_listed_batches(
self,
where_clause: Mapping[str, object],
after: Optional[str],
wanted: int,
user_api_key_dict: UserAPIKeyAuth,
) -> tuple[LiteLLMBatch, ...]:
"""Read chunks newest-first until ``wanted`` batches survive parsing and
file-id resolution or the caller's rows run out, so a run of rows that will
not parse refills the page instead of emptying it. The first chunk is
``wanted`` rows, so a healthy page still costs one query; a scan that has to
continue widens to ``FILE_LIST_CONTINUATION_CHUNK_SIZE`` like ``afile_list``,
and every chunk advances the keyset cursor, so the walk ends once the
caller's rows are exhausted."""
matches: tuple[LiteLLMBatch, ...] = () # rebind-ok: accumulates survivors across chunks
cursor_id: Optional[str] = after # rebind-ok: keyset cursor advances to each chunk's last row
chunk_size: int = wanted # rebind-ok: widens once a scan has to continue past the first chunk
while len(matches) < wanted:
cursor_args: _CursorPageArgs = {"cursor": {"unified_object_id": cursor_id}, "skip": 1} if cursor_id else {}
chunk = await _managed_object_table(self.prisma_client).find_many(
where=where_clause,
take=chunk_size,
order=[{"created_at": "desc"}, {"unified_object_id": "desc"}],
**cursor_args,
)
matches = matches + await self._resolve_listed_rows(
rows=chunk, wanted=wanted - len(matches), user_api_key_dict=user_api_key_dict
)
if len(chunk) < chunk_size:
break
cursor_id = chunk[-1].unified_object_id
chunk_size = max(chunk_size, FILE_LIST_CONTINUATION_CHUNK_SIZE)
return matches
async def _resolve_listed_rows(
self,
rows: "Sequence[PrismaManagedObjectRow]",
wanted: int,
user_api_key_dict: UserAPIKeyAuth,
) -> tuple[LiteLLMBatch, ...]:
parsed_rows: Final = tuple(
(row, batch_obj) for row in batches[:page_size] if (batch_obj := _parse_managed_batch_row(row)) is not None
(row, batch_obj) for row in rows if (batch_obj := _parse_managed_batch_row(row)) is not None
)
unified_id_by_raw_id: Final = await map_raw_file_ids_to_unified(
raw_file_ids=frozenset(
@ -496,19 +538,19 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
),
prisma_client=self.prisma_client,
)
resolved_batches: Final = [
await self._resolve_listed_batch(
resolved: Final[list[LiteLLMBatch]] = [] # mutable-ok: resolution stops as soon as the page is full
for row, batch_obj in parsed_rows:
if len(resolved) == wanted:
break
resolved_batch = await self._resolve_listed_batch(
row=row,
batch_obj=batch_obj,
unified_id_by_raw_id=unified_id_by_raw_id,
user_api_key_dict=user_api_key_dict,
)
for row, batch_obj in parsed_rows
]
return build_list_page(
[batch_obj for batch_obj in resolved_batches if batch_obj is not None],
has_more=has_more,
)
if resolved_batch is not None:
resolved.append(resolved_batch)
return tuple(resolved)
async def _resolve_listed_batch(
self,
@ -815,7 +857,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
file_ids.append(file_id)
return file_ids
def get_file_ids_from_responses_input(self, input: Union[str, List[Dict[str, Any]]]) -> List[str]:
def get_file_ids_from_responses_input(self, input: Union[str, List[Dict[str, object]]]) -> List[str]:
"""
Gets file ids from responses API input.
@ -840,7 +882,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
# Check for direct input_file type
if item.get("type") == "input_file":
file_id = item.get("file_id")
if file_id:
if isinstance(file_id, str) and file_id:
file_ids.append(file_id)
# Check for input_file in content array
@ -849,7 +891,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
for content_item in content:
if isinstance(content_item, dict) and content_item.get("type") == "input_file":
file_id = content_item.get("file_id")
if file_id:
if isinstance(file_id, str) and file_id:
file_ids.append(file_id)
return file_ids
@ -1189,7 +1231,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
# Handle both output_file_id and error_file_id
for file_attr in ["output_file_id", "error_file_id"]:
file_id_value = getattr(response, file_attr, None)
file_id_value: str | None = getattr(response, file_attr, None)
if file_id_value and model_id:
decoded_output_file_id = _is_base64_encoded_unified_file_id(file_id_value)
if decoded_output_file_id and "llm_output_file_id," in decoded_output_file_id:
@ -1458,7 +1500,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
import litellm.proxy.proxy_server as proxy_server_module
# Check if the scheduler has the batch cost checking job registered
scheduler = getattr(proxy_server_module, "scheduler", None)
scheduler: Final[_SchedulerWithJobLookup | None] = getattr(proxy_server_module, "scheduler", None)
if scheduler is None:
return False
@ -1504,7 +1546,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
)
MAX_MATCHES_TO_RETURN = 10
batches = await self.prisma_client.db.litellm_managedobjecttable.find_many(
batches = await _managed_object_table(self.prisma_client).find_many(
where={
"file_purpose": "batch",
"batch_processed": False,
@ -1514,11 +1556,14 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
order={"created_at": "desc"},
)
referencing_batches = []
referencing_batches: Final[list[dict[str, object]]] = []
for batch in batches:
try:
# Parse the batch file_object to check for file references
batch_data = json.loads(batch.file_object) if isinstance(batch.file_object, str) else batch.file_object
decoded_file_object = _decode_json_blob(batch.file_object)
batch_data: Mapping[str, object] = (
decoded_file_object if isinstance(decoded_file_object, Mapping) else {}
)
# Extract file IDs from batch
# Batches typically reference the unified file ID in input_file_id

81
enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py
View file

@ -15,6 +15,7 @@ from collections.abc import Sequence
from typing import TYPE_CHECKING, Final
from fastapi import APIRouter, Depends, HTTPException, Request
from pydantic import TypeAdapter
from litellm._logging import verbose_proxy_logger
from litellm._uuid import uuid
@ -26,39 +27,50 @@ from litellm.proxy.management_helpers.utils import (
management_endpoint_wrapper,
)
from litellm.proxy.utils import PrismaClient, handle_exception_on_proxy
from litellm.repositories.budget_repository import BudgetRepository
from litellm.repositories.object_permission_repository import ObjectPermissionRepository
from litellm.repositories.prisma_protocols import TableActions
from litellm.repositories.project_repository import ProjectRepository
from litellm.repositories.team_repository import TeamRepository
from litellm.repositories.user_repository import UserRepository
from litellm.repositories.verification_token_repository import VerificationTokenRepository
if TYPE_CHECKING:
from prisma import models as prisma_models
from prisma.actions import (
LiteLLM_ProjectTableActions,
LiteLLM_TeamTableActions,
LiteLLM_VerificationTokenActions,
)
from litellm import Router
router = APIRouter()
def _team_table(prisma_client: PrismaClient) -> "LiteLLM_TeamTableActions[prisma_models.LiteLLM_TeamTable]":
team_table: LiteLLM_TeamTableActions[prisma_models.LiteLLM_TeamTable] = prisma_client.db.litellm_teamtable
return team_table
_OBJECT_PERMISSION_PAYLOAD: Final = TypeAdapter(dict[str, object])
def _project_table(prisma_client: PrismaClient) -> "LiteLLM_ProjectTableActions[prisma_models.LiteLLM_ProjectTable]":
project_table: LiteLLM_ProjectTableActions[prisma_models.LiteLLM_ProjectTable] = (
prisma_client.db.litellm_projecttable
)
return project_table
def _team_table(prisma_client: PrismaClient) -> TableActions["prisma_models.LiteLLM_TeamTable"]:
return TeamRepository(prisma_client).table
def _project_table(prisma_client: PrismaClient) -> TableActions["prisma_models.LiteLLM_ProjectTable"]:
return ProjectRepository(prisma_client).table
def _verification_token_table(
prisma_client: PrismaClient,
) -> "LiteLLM_VerificationTokenActions[prisma_models.LiteLLM_VerificationToken]":
verification_token_table: LiteLLM_VerificationTokenActions[prisma_models.LiteLLM_VerificationToken] = (
prisma_client.db.litellm_verificationtoken
)
return verification_token_table
) -> TableActions["prisma_models.LiteLLM_VerificationToken"]:
return VerificationTokenRepository(prisma_client).table
def _budget_table(prisma_client: PrismaClient) -> TableActions["prisma_models.LiteLLM_BudgetTable"]:
return BudgetRepository(prisma_client).table
def _object_permission_table(
prisma_client: PrismaClient,
) -> TableActions["prisma_models.LiteLLM_ObjectPermissionTable"]:
return ObjectPermissionRepository(prisma_client).table
def _user_table(prisma_client: PrismaClient) -> TableActions["prisma_models.LiteLLM_UserTable"]:
return UserRepository(prisma_client).table
def _jsonified(prisma_client: PrismaClient, payload: dict[str, object]) -> dict[str, object]:
@ -329,7 +341,7 @@ async def _create_budget_for_project(
new_budget = _jsonified(prisma_client, budget_row.model_dump(exclude_none=True))
_budget: prisma_models.LiteLLM_BudgetTable = await prisma_client.db.litellm_budgettable.create(
_budget: Final = await _budget_table(prisma_client).create(
data={
**new_budget,
"created_by": user_id or litellm_proxy_admin_name,
@ -352,10 +364,8 @@ async def _set_project_object_permission(
return None
if data.object_permission is not None:
created_object_permission: prisma_models.LiteLLM_ObjectPermissionTable = (
await prisma_client.db.litellm_objectpermissiontable.create(
data=data.object_permission.model_dump(exclude_none=True),
)
created_object_permission: Final = await _object_permission_table(prisma_client).create(
data=data.object_permission.model_dump(exclude_none=True),
)
del data.object_permission
return created_object_permission.object_permission_id
@ -586,10 +596,8 @@ async def new_project(
new_project_row = _remove_budget_fields_from_project_data(new_project_row)
verbose_proxy_logger.info(f"new_project_row: {json.dumps(new_project_row, indent=2)}")
response: prisma_models.LiteLLM_ProjectTable = await prisma_client.db.litellm_projecttable.create(
data={
**new_project_row, # type: ignore
},
response: Final = await _project_table(prisma_client).create(
data={**new_project_row},
include={"litellm_budget_table": True},
)
@ -776,7 +784,7 @@ async def update_project(
if budget_updates and existing_project.budget_id:
# Update existing budget
await prisma_client.db.litellm_budgettable.update(
await _budget_table(prisma_client).update(
where={"budget_id": existing_project.budget_id},
data={
**budget_updates,
@ -791,18 +799,17 @@ async def update_project(
if "object_permission" in update_data:
object_permission_data = update_data.pop("object_permission")
if object_permission_data:
object_permission_payload: Final = _OBJECT_PERMISSION_PAYLOAD.validate_python(object_permission_data)
if existing_project.object_permission_id:
# Update existing permission
await prisma_client.db.litellm_objectpermissiontable.update(
await _object_permission_table(prisma_client).update(
where={"object_permission_id": existing_project.object_permission_id},
data=object_permission_data,
data=object_permission_payload,
)
else:
# Create new permission
created_permission: prisma_models.LiteLLM_ObjectPermissionTable = (
await prisma_client.db.litellm_objectpermissiontable.create(
data=object_permission_data,
)
created_permission: Final = await _object_permission_table(prisma_client).create(
data=object_permission_payload,
)
update_data["object_permission_id"] = created_permission.object_permission_id
@ -818,7 +825,7 @@ async def update_project(
update_data = _remove_budget_fields_from_project_data(update_data)
# Update project
updated_project: prisma_models.LiteLLM_ProjectTable | None = await prisma_client.db.litellm_projecttable.update(
updated_project: Final = await _project_table(prisma_client).update(
where={"project_id": data.project_id},
data=update_data,
include={"litellm_budget_table": True, "object_permission": True},
@ -1058,7 +1065,7 @@ async def list_projects(
# Look up the user's team memberships via the reverse-index on
# LiteLLM_UserTable.teams (maintained by team_member_add alongside
# members_with_roles). This avoids a full scan of all team rows.
user_record: prisma_models.LiteLLM_UserTable | None = await prisma_client.db.litellm_usertable.find_unique(
user_record: Final = await _user_table(prisma_client).find_unique(
where={"user_id": user_api_key_dict.user_id},
)
user_team_ids: list[str] = user_record.teams if user_record is not None and user_record.teams else []

4
enterprise/pyproject.toml
View file

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

12
gateway/Dockerfile
View file

@ -1,5 +1,5 @@
ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:a31344ab2cb8618db84f535eec56f76f6178b142cb92cb2e48676cc2dcebea72
ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:a31344ab2cb8618db84f535eec56f76f6178b142cb92cb2e48676cc2dcebea72
ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:e624c5d5e42382ce7165ddafcbbf8e6769a24cbd02ea6114b880b05ae5ba2a8d
ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:e624c5d5e42382ce7165ddafcbbf8e6769a24cbd02ea6114b880b05ae5ba2a8d
ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a
FROM $UV_IMAGE AS uvbin
@ -16,7 +16,7 @@ COPY --from=uvbin /uv /uvx /usr/local/bin/
# instead of nodeenv downloading one whose dynamic deps may not be in Wolfi
# (e.g. Node 26.2.0 needs libatomic). Retry for transient apk.cgr.dev flakes.
RUN for i in 1 2 3; do \
apk add --no-cache bash gcc python3 python3-dev openssl openssl-dev libsndfile nodejs npm && break; \
apk add --no-cache bash gcc python-3.13 python-3.13-dev openssl openssl-dev libsndfile nodejs npm && break; \
[ $i = 3 ] && { echo "apk add failed after 3 retries" >&2; exit 1; }; \
sleep 5; \
done
@ -47,7 +47,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
--extra extra_proxy \
--extra semantic-router \
--extra bedrock-realtime \
--python python3
--python python3.13
# Stage 2 — copy source and install the project + workspace members.
COPY . .
@ -59,7 +59,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
--extra extra_proxy \
--extra semantic-router \
--extra bedrock-realtime \
--python python3
--python python3.13
RUN HOME=/opt/prisma XDG_CACHE_HOME=/opt/prisma/.cache PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries \
npm_config_cache=/root/.npm \
@ -73,7 +73,7 @@ FROM $LITELLM_RUNTIME_IMAGE AS runtime
USER root
RUN for i in 1 2 3; do \
apk add --no-cache bash openssl tzdata python3 libsndfile libatomic && break; \
apk add --no-cache bash openssl tzdata python-3.13 libsndfile libatomic && break; \
[ $i = 3 ] && { echo "apk add failed after 3 retries" >&2; exit 1; }; \
sleep 5; \
done

1
gateway/routes/allowlist.py
View file

@ -86,6 +86,7 @@ GATEWAY_PATH_PREFIXES: tuple[str, ...] = (
"/comprehendmedical",
"/cohere/",
"/gemini/",
"/gigachat/",
"/google/",
"/vertex_ai/",
"/vertex-ai/",

4
helm/litellm/templates/backend/deployment.yaml
View file

@ -7,6 +7,10 @@ metadata:
{{- include "litellm.commonLabels" . | nindent 4 }}
app.kubernetes.io/component: backend
spec:
{{- with .Values.backend.strategy }}
strategy:
{{- toYaml . | nindent 4 }}
{{- end }}
selector:
matchLabels:
{{- include "litellm.backend.selectorLabels" . | nindent 6 }}

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

@ -7,6 +7,10 @@ metadata:
{{- include "litellm.commonLabels" . | nindent 4 }}
app.kubernetes.io/component: gateway
spec:
{{- with .Values.gateway.strategy }}
strategy:
{{- toYaml . | nindent 4 }}
{{- end }}
selector:
matchLabels:
{{- include "litellm.gateway.selectorLabels" . | nindent 6 }}

148
helm/litellm/templates/ingress.yaml
View file

@ -5,6 +5,41 @@
{{- $gatewayPort := .Values.gateway.service.port -}}
{{- $backendPort := .Values.backend.service.port -}}
{{- $uiPort := .Values.ui.service.port -}}
{{/*
Backends addressable from ingress.extraPaths, keyed by the `service` field.
*/}}
{{- $extraPathBackends := dict
"gateway" (dict "name" $gatewayName "port" $gatewayPort)
"backend" (dict "name" $backendName "port" $backendPort)
"ui" (dict "name" $uiName "port" $uiPort)
-}}
{{/*
UI paths (Next.js static export).
/ui/* is where the SPA serves its login + dashboard routes (e.g. /ui/login).
Without it, /ui/* falls into the catch-all → backend → 404.
The App Router (output: "export", basePath: "") emits the RSC/flight payload
for every route as a ROOT-level <route>.txt (/index.txt, /teams.txt,
/__next._tree.txt, ...). The client router fetches these on every soft
navigation / prefetch as <route>.txt?_rsc=<hash> (the query string is
irrelevant to path matching). They are not under /ui, /_next, or
/litellm-asset-prefix, so without /*.txt they fall to the backend catch-all
→ 404 → client-side navigation never settles and the login flow spins in an
infinite redirect loop (/ ⇄ /ui/login). ui/nginx.conf already serves *.txt
from the export; the rule only routes the request to it. Needs an ingress
controller whose ImplementationSpecific path is a wildcard pattern
(AWS ALB: `*` = 0+ chars); this chart targets the AWS Load Balancer
Controller.
*/}}
{{- $uiPaths := list
(dict "path" "/" "pathType" "Exact")
(dict "path" "/favicon.ico" "pathType" "Exact")
(dict "path" "/litellm-asset-prefix" "pathType" "Prefix")
(dict "path" "/_next" "pathType" "Prefix")
(dict "path" "/ui" "pathType" "Prefix")
(dict "path" "/*.txt" "pathType" "ImplementationSpecific")
-}}
{{/*
Gateway data-plane prefixes — must mirror gateway/routes/allowlist.py.
Versioned paths are listed explicitly to avoid routing management routes
@ -39,6 +74,21 @@
routes at startup -> 404. So /test is rendered as a standalone Exact path
and /test/* falls through to the backend catch-all.
*/}}
{{/*
Every "<path>|<pathType>" this template renders on its own. An
ingress.extraPaths entry that repeats one of these is rejected: duplicates
in a single rule are resolved by position or by controller-specific tie
breaking, so the operator entry could take over a built-in route (an entry
at "/" Prefix would swallow the whole backend management API) instead of
adding to it.
*/}}
{{- $builtinPathKeys := list "/test|Exact" "/|Prefix" -}}
{{- range $uiPaths }}
{{- $builtinPathKeys = append $builtinPathKeys (printf "%s|%s" .path .pathType) }}
{{- end }}
{{- range $gatewayPrefixes }}
{{- $builtinPathKeys = append $builtinPathKeys (printf "%s|Prefix" .) }}
{{- end }}
apiVersion: networking.k8s.io/v1
kind: Ingress
metadata:
@ -64,65 +114,15 @@ spec:
http:
paths:
# --- UI (Next.js static export) ---
- path: /
pathType: Exact
backend:
service:
name: {{ $uiName }}
port:
number: {{ $uiPort }}
- path: /favicon.ico
pathType: Exact
backend:
service:
name: {{ $uiName }}
port:
number: {{ $uiPort }}
- path: /litellm-asset-prefix
pathType: Prefix
backend:
service:
name: {{ $uiName }}
port:
number: {{ $uiPort }}
- path: /_next
pathType: Prefix
backend:
service:
name: {{ $uiName }}
port:
number: {{ $uiPort }}
# /ui/* is where the Next.js SPA serves its login + dashboard
# routes (e.g. /ui/login). Without this, /ui/* falls into the
# catch-all → backend → 404.
- path: /ui
pathType: Prefix
backend:
service:
name: {{ $uiName }}
port:
number: {{ $uiPort }}
# Next.js App Router (output: "export", basePath: "") emits the
# RSC/flight payload for every route as a ROOT-level <route>.txt
# (/index.txt, /teams.txt, /__next._tree.txt, ...). The client
# router fetches these on every soft navigation / prefetch as
# <route>.txt?_rsc=<hash> (the query string is irrelevant to path
# matching). They are not under /ui, /_next, or
# /litellm-asset-prefix, so without this rule they fall to the
# backend catch-all → 404 → client-side navigation never settles
# and the login flow spins in an infinite redirect loop
# (/ ⇄ /ui/login). ui/nginx.conf already serves *.txt from the
# export; this rule only routes the request to it. Needs an
# ingress controller whose ImplementationSpecific path is a
# wildcard pattern (AWS ALB: `*` = 0+ chars); this chart targets
# the AWS Load Balancer Controller.
- path: /*.txt
pathType: ImplementationSpecific
{{- range $uiPaths }}
- path: {{ .path }}
pathType: {{ .pathType }}
backend:
service:
name: {{ $uiName }}
port:
number: {{ $uiPort }}
{{- end }}
# --- Gateway data plane ---
# Exact /test only (see the $gatewayPrefixes comment above);
# /test/* MCP management endpoints fall to the backend catch-all.
@ -142,6 +142,46 @@ spec:
port:
number: {{ $gatewayPort }}
{{- end }}
{{- /*
--- Operator-supplied extra paths (ingress.extraPaths) ---
Rendered after every built-in path so an entry can never take
precedence over a default, and before the backend catch-all.
Position only decides the match on controllers that honour manifest
order: the AWS Load Balancer Controller this chart targets sorts
Exact paths first and Prefix paths longest-first, but keeps
ImplementationSpecific paths in manifest order, which is what the
/*.txt rule above already depends on.
*/}}
{{- range $idx, $extra := .Values.ingress.extraPaths }}
{{- if not (kindIs "map" $extra) }}
{{- fail (printf "ingress.extraPaths[%d]: each entry must be a mapping with a 'path' key" $idx) }}
{{- end }}
{{- if not $extra.path }}
{{- fail (printf "ingress.extraPaths[%d]: 'path' is required" $idx) }}
{{- end }}
{{- $service := $extra.service | default "gateway" }}
{{- $target := get $extraPathBackends $service }}
{{- if not $target }}
{{- fail (printf "ingress.extraPaths[%d] (path %s): unknown service %q, expected one of backend, gateway, ui" $idx $extra.path $service) }}
{{- end }}
{{- $pathType := $extra.pathType | default "Prefix" }}
{{- if not (has $pathType (list "Prefix" "Exact" "ImplementationSpecific")) }}
{{- fail (printf "ingress.extraPaths[%d] (path %s): unknown pathType %q, expected one of Exact, ImplementationSpecific, Prefix" $idx $extra.path $pathType) }}
{{- end }}
{{- if eq $extra.path "/" }}
{{- fail (printf "ingress.extraPaths[%d]: path / is already routed in both directions, Exact to ui and Prefix to backend, so no pathType leaves a request for an entry here to capture" $idx) }}
{{- end }}
{{- if has (printf "%s|%s" $extra.path $pathType) $builtinPathKeys }}
{{- fail (printf "ingress.extraPaths[%d]: path %s with pathType %s is already routed by this chart, and a duplicate would take it over rather than add to it" $idx $extra.path $pathType) }}
{{- end }}
- path: {{ $extra.path | quote }}
pathType: {{ $pathType }}
backend:
service:
name: {{ $target.name }}
port:
number: {{ $target.port }}
{{- end }}
# --- Catch-all → backend (management API: /key/*, /user/*, /team/*, ...) ---
- path: /
pathType: Prefix

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

@ -7,6 +7,8 @@
#
# Running this pre-upgrade closes the window where new application pods would
# otherwise serve traffic against the previous release's unmigrated schema.
# Argo CD users can swap the Helm hook for a PreSync hook through
# `migrationJob.hooks`, which re-runs the Job on every sync.
apiVersion: batch/v1
kind: Job
metadata:
@ -14,10 +16,18 @@ metadata:
labels:
{{- include "litellm.commonLabels" . | nindent 4 }}
app.kubernetes.io/component: migrations
{{- if or .Values.migrationJob.hooks.helm.enabled .Values.migrationJob.hooks.argocd.enabled }}
annotations:
{{- if .Values.migrationJob.hooks.helm.enabled }}
helm.sh/hook: pre-install,pre-upgrade
helm.sh/hook-delete-policy: before-hook-creation
helm.sh/hook-weight: "0"
helm.sh/hook-weight: {{ .Values.migrationJob.hooks.helm.weight | default "0" | quote }}
{{- end }}
{{- if .Values.migrationJob.hooks.argocd.enabled }}
argocd.argoproj.io/hook: PreSync
argocd.argoproj.io/hook-delete-policy: BeforeHookCreation
{{- end }}
{{- end }}
spec:
backoffLimit: {{ .Values.migrationJob.backoffLimit }}
ttlSecondsAfterFinished: {{ .Values.migrationJob.ttlSecondsAfterFinished }}

4
helm/litellm/templates/ui/deployment.yaml
View file

@ -7,6 +7,10 @@ metadata:
{{- include "litellm.commonLabels" . | nindent 4 }}
app.kubernetes.io/component: ui
spec:
{{- with .Values.ui.strategy }}
strategy:
{{- toYaml . | nindent 4 }}
{{- end }}
selector:
matchLabels:
{{- include "litellm.ui.selectorLabels" . | nindent 6 }}

317
helm/litellm/tests/ingress_extra_paths_tests.yaml Normal file
View file

@ -0,0 +1,317 @@
suite: test ingress.extraPaths
templates:
- ingress.yaml
values:
- ./values/required.yaml
tests:
- it: renders nothing extra between the built-in gateway prefixes and the backend catch-all when unset
set:
ingress.enabled: true
asserts:
- equal:
path: spec.rules[0].http.paths[-1]
value:
path: /
pathType: Prefix
backend:
service:
name: RELEASE-NAME-litellm-backend
port:
number: 4001
- equal:
path: spec.rules[0].http.paths[-2]
value:
path: /metrics
pathType: Prefix
backend:
service:
name: RELEASE-NAME-litellm-gateway
port:
number: 4000
- it: routes an extra path to the gateway by default, immediately before the backend catch-all
set:
ingress.enabled: true
ingress.extraPaths:
- path: /watsonx
asserts:
- equal:
path: spec.rules[0].http.paths[-2]
value:
path: /watsonx
pathType: Prefix
backend:
service:
name: RELEASE-NAME-litellm-gateway
port:
number: 4000
- equal:
path: spec.rules[0].http.paths[-1]
value:
path: /
pathType: Prefix
backend:
service:
name: RELEASE-NAME-litellm-backend
port:
number: 4001
- it: keeps every built-in path when extra paths are supplied
set:
ingress.enabled: true
ingress.extraPaths:
- path: /watsonx
asserts:
- contains:
path: spec.rules[0].http.paths
content:
path: /
pathType: Exact
backend:
service:
name: RELEASE-NAME-litellm-ui
port:
number: 3000
- contains:
path: spec.rules[0].http.paths
content:
path: /ui
pathType: Prefix
backend:
service:
name: RELEASE-NAME-litellm-ui
port:
number: 3000
- contains:
path: spec.rules[0].http.paths
content:
path: /test
pathType: Exact
backend:
service:
name: RELEASE-NAME-litellm-gateway
port:
number: 4000
- contains:
path: spec.rules[0].http.paths
content:
path: /v1/chat
pathType: Prefix
backend:
service:
name: RELEASE-NAME-litellm-gateway
port:
number: 4000
- contains:
path: spec.rules[0].http.paths
content:
path: /vertex_ai
pathType: Prefix
backend:
service:
name: RELEASE-NAME-litellm-gateway
port:
number: 4000
- it: renders every entry in order and honours the service and pathType selectors
set:
ingress.enabled: true
ingress.extraPaths:
- path: /watsonx
service: gateway
- path: /my-passthrough
pathType: Exact
service: backend
- path: /brand.txt
pathType: ImplementationSpecific
service: ui
asserts:
- equal:
path: spec.rules[0].http.paths[-4]
value:
path: /watsonx
pathType: Prefix
backend:
service:
name: RELEASE-NAME-litellm-gateway
port:
number: 4000
- equal:
path: spec.rules[0].http.paths[-3]
value:
path: /my-passthrough
pathType: Exact
backend:
service:
name: RELEASE-NAME-litellm-backend
port:
number: 4001
- equal:
path: spec.rules[0].http.paths[-2]
value:
path: /brand.txt
pathType: ImplementationSpecific
backend:
service:
name: RELEASE-NAME-litellm-ui
port:
number: 3000
- it: addresses the component services by their configured ports
set:
ingress.enabled: true
gateway.service.port: 8000
backend.service.port: 8001
ui.service.port: 8080
ingress.extraPaths:
- path: /watsonx
- path: /my-passthrough
service: backend
- path: /brand.txt
service: ui
asserts:
- equal:
path: spec.rules[0].http.paths[-4].backend.service.port.number
value: 8000
- equal:
path: spec.rules[0].http.paths[-3].backend.service.port.number
value: 8001
- equal:
path: spec.rules[0].http.paths[-2].backend.service.port.number
value: 8080
- it: rejects an entry naming a service the chart does not deploy
set:
ingress.enabled: true
ingress.extraPaths:
- path: /watsonx
service: proxy
asserts:
- failedTemplate:
errorMessage: 'ingress.extraPaths[0] (path /watsonx): unknown service "proxy", expected one of backend, gateway, ui'
- it: rejects an entry whose pathType is not a kubernetes pathType
set:
ingress.enabled: true
ingress.extraPaths:
- path: /watsonx
pathType: prefix
asserts:
- failedTemplate:
errorMessage: 'ingress.extraPaths[0] (path /watsonx): unknown pathType "prefix", expected one of Exact, ImplementationSpecific, Prefix'
- it: rejects an entry with no path
set:
ingress.enabled: true
ingress.extraPaths:
- service: gateway
asserts:
- failedTemplate:
errorMessage: "ingress.extraPaths[0]: 'path' is required"
- it: rejects a root entry that would take over the backend catch-all
set:
ingress.enabled: true
ingress.extraPaths:
- path: /
service: gateway
asserts:
- failedTemplate:
errorMessage: "ingress.extraPaths[0]: path / is already routed in both directions, Exact to ui and Prefix to backend, so no pathType leaves a request for an entry here to capture"
- it: rejects a root entry that would take over the UI root
set:
ingress.enabled: true
ingress.extraPaths:
- path: /
pathType: Exact
service: gateway
asserts:
- failedTemplate:
errorMessage: "ingress.extraPaths[0]: path / is already routed in both directions, Exact to ui and Prefix to backend, so no pathType leaves a request for an entry here to capture"
# A root ImplementationSpecific entry duplicates no built-in pair, so the
# duplicate check alone would admit it. It is still dead: the built-in
# Exact / sorts ahead of it on the AWS Load Balancer Controller and claims
# the only request its pattern matches, so it renders and never routes.
- it: rejects a root entry that would render but never match
set:
ingress.enabled: true
ingress.extraPaths:
- path: /
pathType: ImplementationSpecific
service: gateway
asserts:
- failedTemplate:
errorMessage: "ingress.extraPaths[0]: path / is already routed in both directions, Exact to ui and Prefix to backend, so no pathType leaves a request for an entry here to capture"
- it: rejects an entry that would take over a UI prefix
set:
ingress.enabled: true
ingress.extraPaths:
- path: /ui
service: gateway
asserts:
- failedTemplate:
errorMessage: "ingress.extraPaths[0]: path /ui with pathType Prefix is already routed by this chart, and a duplicate would take it over rather than add to it"
- it: rejects an entry that would take over the UI RSC payload rule
set:
ingress.enabled: true
ingress.extraPaths:
- path: /*.txt
pathType: ImplementationSpecific
service: backend
asserts:
- failedTemplate:
errorMessage: "ingress.extraPaths[0]: path /*.txt with pathType ImplementationSpecific is already routed by this chart, and a duplicate would take it over rather than add to it"
- it: rejects an entry that would take over a gateway data-plane prefix
set:
ingress.enabled: true
ingress.extraPaths:
- path: /v1/chat
service: backend
asserts:
- failedTemplate:
errorMessage: "ingress.extraPaths[0]: path /v1/chat with pathType Prefix is already routed by this chart, and a duplicate would take it over rather than add to it"
- it: rejects an entry that would take over the exact /test route
set:
ingress.enabled: true
ingress.extraPaths:
- path: /test
pathType: Exact
service: backend
asserts:
- failedTemplate:
errorMessage: "ingress.extraPaths[0]: path /test with pathType Exact is already routed by this chart, and a duplicate would take it over rather than add to it"
- it: allows a built-in path under a different pathType, which is a distinct rule
set:
ingress.enabled: true
ingress.extraPaths:
- path: /ui
pathType: Exact
service: ui
asserts:
- equal:
path: spec.rules[0].http.paths[-2]
value:
path: /ui
pathType: Exact
backend:
service:
name: RELEASE-NAME-litellm-ui
port:
number: 3000
- it: rejects a bare string entry instead of failing on template internals
set:
ingress.enabled: true
ingress.extraPaths:
- /watsonx
asserts:
- failedTemplate:
errorMessage: "ingress.extraPaths[0]: each entry must be a mapping with a 'path' key"

63
helm/litellm/tests/migration_job_hooks_tests.yaml Normal file
View file

@ -0,0 +1,63 @@
suite: test migrations Job hook annotations
templates:
- migrations-job.yaml
values:
- ./values/required.yaml
tests:
- it: runs as a Helm pre-install / pre-upgrade hook by default
asserts:
- equal:
path: metadata.annotations["helm.sh/hook"]
value: pre-install,pre-upgrade
- equal:
path: metadata.annotations["helm.sh/hook-delete-policy"]
value: before-hook-creation
- equal:
path: metadata.annotations["helm.sh/hook-weight"]
value: "0"
- notExists:
path: metadata.annotations["argocd.argoproj.io/hook"]
- it: adds the Argo CD PreSync hook when asked
set:
migrationJob.hooks.argocd.enabled: true
asserts:
- equal:
path: metadata.annotations["argocd.argoproj.io/hook"]
value: PreSync
- equal:
path: metadata.annotations["argocd.argoproj.io/hook-delete-policy"]
value: BeforeHookCreation
- it: drops the Helm hook so Argo CD owns the Job
set:
migrationJob.hooks.argocd.enabled: true
migrationJob.hooks.helm.enabled: false
asserts:
- equal:
path: metadata.annotations["argocd.argoproj.io/hook"]
value: PreSync
- notExists:
path: metadata.annotations["helm.sh/hook"]
- notExists:
path: metadata.annotations["helm.sh/hook-delete-policy"]
- notExists:
path: metadata.annotations["helm.sh/hook-weight"]
- it: renders an ordinary Job when both hooks are disabled
set:
migrationJob.hooks.helm.enabled: false
asserts:
- notExists:
path: metadata.annotations
- equal:
path: kind
value: Job
- it: honours a custom Helm hook weight
set:
migrationJob.hooks.helm.weight: "-5"
asserts:
- equal:
path: metadata.annotations["helm.sh/hook-weight"]
value: "-5"

66
helm/litellm/tests/rollout_strategy_tests.yaml Normal file
View file

@ -0,0 +1,66 @@
suite: test rolling update strategy on the component deployments
templates:
- gateway/deployment.yaml
- gateway/configmap.yaml
- backend/deployment.yaml
- ui/deployment.yaml
values:
- ./values/required.yaml
tests:
- it: leaves the strategy to Kubernetes defaults when unset
asserts:
- notExists:
path: spec.strategy
- it: renders the configured strategy on each deployment
set:
gateway.strategy:
type: RollingUpdate
rollingUpdate:
maxUnavailable: 0
maxSurge: 1
backend.strategy:
type: RollingUpdate
rollingUpdate:
maxUnavailable: "25%"
maxSurge: 2
ui.strategy:
type: Recreate
asserts:
- equal:
path: spec.strategy
value:
type: RollingUpdate
rollingUpdate:
maxUnavailable: 0
maxSurge: 1
template: gateway/deployment.yaml
- equal:
path: spec.strategy
value:
type: RollingUpdate
rollingUpdate:
maxUnavailable: 25%
maxSurge: 2
template: backend/deployment.yaml
- equal:
path: spec.strategy
value:
type: Recreate
template: ui/deployment.yaml
- it: keeps a component on the cluster default when only another one sets a strategy
set:
gateway.strategy:
type: Recreate
asserts:
- equal:
path: spec.strategy.type
value: Recreate
template: gateway/deployment.yaml
- notExists:
path: spec.strategy
template: backend/deployment.yaml
- notExists:
path: spec.strategy
template: ui/deployment.yaml

50
helm/litellm/values.yaml
View file

@ -13,6 +13,27 @@ ingress:
annotations: {}
host: "" # optional; if set, becomes the rule's host
tls: []
# Extra HTTP paths appended to the ingress rule. Additive: every built-in
# UI / gateway / backend path is still rendered, these entries are placed
# after them and before the backend catch-all, and an entry that repeats a
# path the chart already routes is rejected at render time rather than
# silently taking it over.
#
# The chart's built-in gateway prefix list is a snapshot of the data-plane
# surface at release time. Use extraPaths for passthrough routes it does not
# cover: a provider prefix added upstream after this chart version, or a
# custom general_settings.pass_through_endpoints route.
#
# path required; the HTTP path to route
# service which component serves it: gateway (default), backend, or ui
# pathType Prefix (default), Exact, or ImplementationSpecific
#
# The target component only answers paths its own route allowlist keeps, so
# a path here still has to be one that component serves.
extraPaths: []
# - path: /watsonx
# pathType: Prefix
# service: gateway
# Per-component ServiceAccounts for gateway, backend, and ui.
#
@ -54,6 +75,22 @@ serviceAccounts:
# generate` — the migration engine doesn't need the generated client.
migrationJob:
enabled: true
# Which controller is responsible for running the Job.
#
# `helm.enabled` renders the Helm pre-install / pre-upgrade hook, so the Job
# runs whenever `helm upgrade` sees a change to apply. `argocd.enabled`
# renders an Argo CD PreSync hook instead, which runs the Job on every sync
# even when the rendered manifests are unchanged: the way to re-run
# migrations on demand from a GitOps pipeline. Turning the Helm hook off
# while the Argo CD hook is on leaves the Job out of Helm's own upgrade
# path, which is what Argo CD users want since Argo, not Helm, applies the
# manifests.
hooks:
helm:
enabled: true
weight: "0"
argocd:
enabled: false
backoffLimit: 4
ttlSecondsAfterFinished: 120
# Wall-clock budget for the whole Job, shared across every `backoffLimit`
@ -236,6 +273,15 @@ gateway:
initialDelaySeconds: 5
periodSeconds: 10
timeoutSeconds: 10
# Rolling update tuning for the gateway Deployment. Empty by default, so
# Kubernetes applies its own RollingUpdate defaults (25% maxSurge /
# 25% maxUnavailable). Example, for a surge-only rollout behind a load
# balancer that must never lose capacity:
# type: RollingUpdate
# rollingUpdate:
# maxUnavailable: 0
# maxSurge: 1
strategy: {}
# Optional startupProbe. Empty by default, so existing installs are unchanged
# and liveness/readiness apply from container start. Set it to gate
# liveness/readiness until a slow cold start finishes — a high failureThreshold
@ -348,6 +394,8 @@ backend:
initialDelaySeconds: 5
periodSeconds: 10
timeoutSeconds: 10
# Same shape as gateway.strategy.
strategy: {}
# Optional startupProbe; same shape as gateway.startupProbe. Empty by default.
startupProbe: {}
hpa:
@ -412,6 +460,8 @@ ui:
httpGet: { path: /, port: http }
initialDelaySeconds: 2
periodSeconds: 10
# Same shape as gateway.strategy.
strategy: {}
# Optional startupProbe; same shape as gateway.startupProbe. Empty by default.
startupProbe: {}
hpa:

12
litellm-proxy-extras/litellm_proxy_extras/migrations/20260804162853_add_budget_window_spend_table/migration.sql Normal file
View file

@ -0,0 +1,12 @@
CREATE TABLE IF NOT EXISTS "LiteLLM_BudgetWindowSpend" (
"entity_type" TEXT NOT NULL,
"entity_id" TEXT NOT NULL,
"window_duration" TEXT NOT NULL,
"window_start" TIMESTAMP(3) NOT NULL,
"spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0,
"created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
"updated_at" TIMESTAMP(3) NOT NULL,
CONSTRAINT "LiteLLM_BudgetWindowSpend_pkey" PRIMARY KEY ("entity_type","entity_id","window_duration")
);

20
litellm-proxy-extras/litellm_proxy_extras/migrations/20260829000000_add_model_access_group_budget_table/migration.sql Normal file
View file

@ -0,0 +1,20 @@
-- CreateTable
CREATE TABLE IF NOT EXISTS "LiteLLM_ModelAccessGroupBudgetTable" (
"access_group_name" TEXT NOT NULL,
"spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0,
"budget_id" TEXT,
"created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
"created_by" TEXT,
"updated_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
"updated_by" TEXT,
CONSTRAINT "LiteLLM_ModelAccessGroupBudgetTable_pkey" PRIMARY KEY ("access_group_name")
);
-- AddForeignKey
DO $$
BEGIN
IF NOT EXISTS (SELECT 1 FROM pg_constraint WHERE conname = 'LiteLLM_ModelAccessGroupBudgetTable_budget_id_fkey') THEN
ALTER TABLE "LiteLLM_ModelAccessGroupBudgetTable" ADD CONSTRAINT "LiteLLM_ModelAccessGroupBudgetTable_budget_id_fkey" FOREIGN KEY ("budget_id") REFERENCES "LiteLLM_BudgetTable"("budget_id") ON DELETE SET NULL ON UPDATE CASCADE;
END IF;
END $$;

21
litellm-proxy-extras/litellm_proxy_extras/migrations/20260831000000_shadow_eval_typed_targets/migration.sql Normal file
View file

@ -0,0 +1,21 @@
DO $$
BEGIN
IF EXISTS (
SELECT 1 FROM information_schema.columns
WHERE table_name = 'LiteLLM_ShadowEvalJob' AND column_name = 'api_key_id'
) THEN
ALTER TABLE "LiteLLM_ShadowEvalJob" RENAME COLUMN "api_key_id" TO "target_id";
END IF;
END $$;
ALTER TABLE "LiteLLM_ShadowEvalJob" ADD COLUMN IF NOT EXISTS "target_type" TEXT NOT NULL DEFAULT 'key';
DROP INDEX IF EXISTS "LiteLLM_ShadowEvalJob_one_active_per_key_direction";
CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_ShadowEvalJob_one_active_per_target_direction"
ON "LiteLLM_ShadowEvalJob"("target_type", "target_id", "direction") WHERE "stopped_at" IS NULL;
DROP INDEX IF EXISTS "LiteLLM_ShadowEvalJob_api_key_id_idx";
CREATE INDEX IF NOT EXISTS "LiteLLM_ShadowEvalJob_target_type_target_id_idx"
ON "LiteLLM_ShadowEvalJob"("target_type", "target_id");

3
litellm-proxy-extras/litellm_proxy_extras/migrations/20260901000000_shadow_eval_multi_router/migration.sql Normal file
View file

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

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

@ -29,6 +29,7 @@ model LiteLLM_BudgetTable {
keys LiteLLM_VerificationToken[] // multiple keys can have the same budget
end_users LiteLLM_EndUserTable[] // multiple end-users can have the same budget
tags LiteLLM_TagTable[] // multiple tags can have the same budget
model_access_groups LiteLLM_ModelAccessGroupBudgetTable[] // multiple model access groups can have the same budget
team_membership LiteLLM_TeamMembership[] // budgets of Users within a Team
organization_membership LiteLLM_OrganizationMembership[] // budgets of Users within a Organization
}
@ -585,6 +586,20 @@ model LiteLLM_EndUserTable {
blocked Boolean @default(false)
}
// Budget and shared spend for a model access group. The groups themselves are not rows anywhere:
// they are free-text strings in LiteLLM_ProxyModelTable.model_info.access_groups, so a row here
// exists only once someone gives that group a budget.
model LiteLLM_ModelAccessGroupBudgetTable {
access_group_name String @id
spend Float @default(0.0)
budget_id String?
litellm_budget_table LiteLLM_BudgetTable? @relation(fields: [budget_id], references: [budget_id])
created_at DateTime @default(now()) @map("created_at")
created_by String?
updated_at DateTime @default(now()) @updatedAt
updated_by String?
}
// Track tags with budgets and spend
model LiteLLM_TagTable {
tag_name String @id
@ -649,6 +664,18 @@ model LiteLLM_SpendLogs {
@@index([session_id])
}
model LiteLLM_BudgetWindowSpend {
entity_type String
entity_id String
window_duration String
window_start DateTime
spend Float @default(0.0)
created_at DateTime @default(now())
updated_at DateTime @updatedAt
@@id([entity_type, entity_id, window_duration])
}
// View spend, model, api_key per request
model LiteLLM_ErrorLogs {
request_id String @id @default(uuid())
@ -1502,14 +1529,16 @@ model LiteLLM_AutoRouterSession {
model LiteLLM_ShadowEvalJob {
id String @id @default(cuid())
group_id String // legs of one job share this; the API's job id
api_key_id String // hashed virtual key whose traffic this leg shadows
router_name String // the auto-router under evaluation, in either direction
target_type String @default("key") // key | team | user
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)
direction String @default("forward") // forward | reverse
baseline_model String? // reverse only: the fixed model the router is judged against
judge_model String
shadow_percentage Float
max_turns Int // sample-count ceiling: the whole budget on pre-max_budget jobs, the error-loop valve otherwise
max_budget Float? // per-key USD cap on the eval's own shadow + judge spend; null on jobs from before spend budgets
max_budget Float? // per-target USD cap on the eval's own shadow + judge spend; null on jobs from before spend budgets
created_at DateTime @default(now())
created_by String?
ends_at DateTime
@ -1517,7 +1546,7 @@ model LiteLLM_ShadowEvalJob {
stopped_by String? // operator who stopped it early; null when it ended on its own
@@index([group_id])
@@index([api_key_id])
@@index([target_type, target_id])
@@index([created_at])
}
@ -1527,6 +1556,7 @@ model LiteLLM_ShadowEvalAttempt {
job_id String
request_id String // the judged real request
outcome String // real | shadow | tie | error
router_name String? // the arm this verdict scores; NULL on legacy rows, meaning the job's own router
tier String? // router's tier for the prompt, when classified
real_model String?
shadow_model String?

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

@ -512,6 +512,13 @@ class ProxyExtrasDBManager:
try:
import psycopg
except ImportError:
logger.warning(
"psycopg is not installed; skipping the LiteLLM_SpendLogs "
"partition check. If this table is partitioned (see "
"db_scripts/partition_spend_logs.sql), schema reconciliation "
"will try to rewrite its primary key and fail. Install the "
"litellm[extra_proxy] extra, which now includes psycopg."
)
return False
cleaned_url = ProxyExtrasDBManager._strip_prisma_query_params(database_url)

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

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

340
litellm-rust/Cargo.lock generated
View file

@ -2,6 +2,36 @@
# It is not intended for manual editing.
version = 4
[[package]]
name = "aho-corasick"
version = "1.1.5"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "c982642fa9e8606056828ee9a8505737230110bb1099153c79efe865c59d12ba"
dependencies = [
"memchr",
]
[[package]]
name = "alloca"
version = "0.4.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e5a7d05ea6aea7e9e64d25b9156ba2fee3fdd659e34e41063cd2fc7cd020d7f4"
dependencies = [
"cc",
]
[[package]]
name = "anes"
version = "0.1.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "4b46cbb362ab8752921c97e041f5e366ee6297bd428a31275b9fcf1e380f7299"
[[package]]
name = "anstyle"
version = "1.0.14"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "940b3a0ca603d1eade50a4846a2afffd5ef57a9feac2c0e2ec2e14f9ead76000"
[[package]]
name = "arc-swap"
version = "1.9.2"
@ -506,6 +536,12 @@ dependencies = [
"either",
]
[[package]]
name = "cast"
version = "0.3.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "37b2a672a2cb129a2e41c10b1224bb368f9f37a2b16b612598138befd7b37eb5"
[[package]]
name = "cc"
version = "1.3.0"
@ -541,6 +577,58 @@ dependencies = [
"rand_core 0.10.1",
]
[[package]]
name = "ciborium"
version = "0.2.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "42e69ffd6f0917f5c029256a24d0161db17cea3997d185db0d35926308770f0e"
dependencies = [
"ciborium-io",
"ciborium-ll",
"serde",
]
[[package]]
name = "ciborium-io"
version = "0.2.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "05afea1e0a06c9be33d539b876f1ce3692f4afea2cb41f740e7743225ed1c757"
[[package]]
name = "ciborium-ll"
version = "0.2.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "57663b653d948a338bfb3eeba9bb2fd5fcfaecb9e199e87e1eda4d9e8b240fd9"
dependencies = [
"ciborium-io",
"half",
]
[[package]]
name = "clap"
version = "4.6.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "473c7e07f409a8d772161724aa8db6a765a2532a70f9667eeb7b49d3d02fbdca"
dependencies = [
"clap_builder",
]
[[package]]
name = "clap_builder"
version = "4.6.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "7b48fea5a88e9ae728a2dcbedbfc0e730f7d60da42e1cb049a83c9fb8b789889"
dependencies = [
"anstyle",
"clap_lex",
]
[[package]]
name = "clap_lex"
version = "1.1.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "c8d4a3bb8b1e0c1050499d1815f5ab16d04f0959b233085fb31653fbfc9d98f9"
[[package]]
name = "cmake"
version = "0.1.58"
@ -596,6 +684,72 @@ dependencies = [
"libc",
]
[[package]]
name = "criterion"
version = "0.8.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "950046b2aa2492f9a536f5f4f9a3de7b9e2476e575e05bd6c333371add4d98f3"
dependencies = [
"alloca",
"anes",
"cast",
"ciborium",
"clap",
"criterion-plot",
"itertools",
"num-traits",
"oorandom",
"page_size",
"plotters",
"rayon",
"regex",
"serde",
"serde_json",
"tinytemplate",
"walkdir",
]
[[package]]
name = "criterion-plot"
version = "0.8.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "d8d80a2f4f5b554395e47b5d8305bc3d27813bacb73493eb1001e8f76dae29ea"
dependencies = [
"cast",
"itertools",
]
[[package]]
name = "crossbeam-deque"
version = "0.8.7"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "5181e0de7b61eb03a81e347d6dd8797bae9da5146707b51077e2d71a54ec0ceb"
dependencies = [
"crossbeam-epoch",
"crossbeam-utils",
]
[[package]]
name = "crossbeam-epoch"
version = "0.9.20"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "2d6914041f254d6e9176c01941b21115dcfb7089e55135a35411081bd106ef3f"
dependencies = [
"crossbeam-utils",
]
[[package]]
name = "crossbeam-utils"
version = "0.8.22"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "61803da095bee82a81bb1a452ecc25d3b2f1416d1897eb86430c6159ef717c17"
[[package]]
name = "crunchy"
version = "0.2.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "460fbee9c2c2f33933d720630a6a0bac33ba7053db5344fac858d4b8952d77d5"
[[package]]
name = "crypto-common"
version = "0.1.7"
@ -856,6 +1010,17 @@ dependencies = [
"tracing",
]
[[package]]
name = "half"
version = "2.7.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b"
dependencies = [
"cfg-if",
"crunchy",
"zerocopy",
]
[[package]]
name = "hashbrown"
version = "0.17.1"
@ -1179,6 +1344,15 @@ version = "2.12.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "d98f6fed1fde3f8c21bc40a1abb88dd75e67924f9cffc3ef95607bad8017f8e2"
[[package]]
name = "itertools"
version = "0.13.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "413ee7dfc52ee1a4949ceeb7dbc8a33f2d6c088194d9f922fb8318faf1f01186"
dependencies = [
"either",
]
[[package]]
name = "itoa"
version = "1.0.18"
@ -1255,10 +1429,13 @@ dependencies = [
name = "litellm-python-bridge"
version = "0.1.0"
dependencies = [
"criterion",
"litellm-ai-gateway",
"litellm-core",
"pyo3",
"pyo3-async-runtimes",
"pythonize",
"serde",
"serde_json",
"tokio",
]
@ -1340,6 +1517,12 @@ version = "1.21.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50"
[[package]]
name = "oorandom"
version = "11.1.5"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "d6790f58c7ff633d8771f42965289203411a5e5c68388703c06e14f24770b41e"
[[package]]
name = "openssl-probe"
version = "0.2.1"
@ -1352,6 +1535,16 @@ version = "0.5.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "1a80800c0488c3a21695ea981a54918fbb37abf04f4d0720c453632255e2ff0e"
[[package]]
name = "page_size"
version = "0.6.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "30d5b2194ed13191c1999ae0704b7839fb18384fa22e49b57eeaa97d79ce40da"
dependencies = [
"libc",
"winapi",
]
[[package]]
name = "percent-encoding"
version = "2.3.2"
@ -1376,6 +1569,34 @@ version = "0.3.33"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "19f132c84eca552bf34cab8ec81f1c1dcc229b811638f9d283dceabe58c5569e"
[[package]]
name = "plotters"
version = "0.3.7"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "5aeb6f403d7a4911efb1e33402027fc44f29b5bf6def3effcc22d7bb75f2b747"
dependencies = [
"num-traits",
"plotters-backend",
"plotters-svg",
"wasm-bindgen",
"web-sys",
]
[[package]]
name = "plotters-backend"
version = "0.3.7"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "df42e13c12958a16b3f7f4386b9ab1f3e7933914ecea48da7139435263a4172a"
[[package]]
name = "plotters-svg"
version = "0.3.7"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "51bae2ac328883f7acdfea3d66a7c35751187f870bc81f94563733a154d7a670"
dependencies = [
"plotters-backend",
]
[[package]]
name = "portable-atomic"
version = "1.14.0"
@ -1486,6 +1707,16 @@ dependencies = [
"syn 2.0.119",
]
[[package]]
name = "pythonize"
version = "0.29.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "6ec376e1216e0c929a74964ce2020012a1a39f32d80e78aa688721219ea7fb89"
dependencies = [
"pyo3",
"serde",
]
[[package]]
name = "quinn"
version = "0.11.11"
@ -1613,12 +1844,61 @@ dependencies = [
"rand_core 0.10.1",
]
[[package]]
name = "rayon"
version = "1.12.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "fb39b166781f92d482534ef4b4b1b2568f42613b53e5b6c160e24cfbfa30926d"
dependencies = [
"either",
"rayon-core",
]
[[package]]
name = "rayon-core"
version = "1.13.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "22e18b0f0062d30d4230b2e85ff77fdfe4326feb054b9783a3460d8435c8ab91"
dependencies = [
"crossbeam-deque",
"crossbeam-utils",
]
[[package]]
name = "regex"
version = "1.13.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "f020237b6c8eed93db2e2cb53c00c60a8e1bc73da7d073199a1180401450218d"
dependencies = [
"aho-corasick",
"memchr",
"regex-automata",
"regex-syntax",
]
[[package]]
name = "regex-automata"
version = "0.4.18"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "ad8553b9b26413251cbf30e620595c7a41b3887f03da04579c0e6b0d6a06b4b2"
dependencies = [
"aho-corasick",
"memchr",
"regex-syntax",
]
[[package]]
name = "regex-lite"
version = "0.1.9"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "cab834c73d247e67f4fae452806d17d3c7501756d98c8808d7c9c7aa7d18f973"
[[package]]
name = "regex-syntax"
version = "0.8.11"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "d6f6ff9a378485b298a5286656da665ba74413d36db0979633275d2e708145d4"
[[package]]
name = "reqwest"
version = "0.12.28"
@ -1774,6 +2054,15 @@ version = "1.0.23"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "9774ba4a74de5f7b1c1451ed6cd5285a32eddb5cccb8cc655a4e50009e06477f"
[[package]]
name = "same-file"
version = "1.0.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "93fc1dc3aaa9bfed95e02e6eadabb4baf7e3078b0bd1b4d7b6b0b68378900502"
dependencies = [
"winapi-util",
]
[[package]]
name = "schannel"
version = "0.1.29"
@ -2099,6 +2388,16 @@ dependencies = [
"zerovec",
]
[[package]]
name = "tinytemplate"
version = "1.2.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "be4d6b5f19ff7664e8c98d03e2139cb510db9b0a60b55f8e8709b689d939b6bc"
dependencies = [
"serde",
"serde_json",
]
[[package]]
name = "tinyvec"
version = "1.12.0"
@ -2363,6 +2662,16 @@ version = "0.8.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "5c3082ca00d5a5ef149bb8b555a72ae84c9c59f7250f013ac822ac2e49b19c64"
[[package]]
name = "walkdir"
version = "2.5.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "29790946404f91d9c5d06f9874efddea1dc06c5efe94541a7d6863108e3a5e4b"
dependencies = [
"same-file",
"winapi-util",
]
[[package]]
name = "want"
version = "0.3.1"
@ -2475,6 +2784,37 @@ dependencies = [
"rustls-pki-types",
]
[[package]]
name = "winapi"
version = "0.3.9"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "5c839a674fcd7a98952e593242ea400abe93992746761e38641405d28b00f419"
dependencies = [
"winapi-i686-pc-windows-gnu",
"winapi-x86_64-pc-windows-gnu",
]
[[package]]
name = "winapi-i686-pc-windows-gnu"
version = "0.4.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "ac3b87c63620426dd9b991e5ce0329eff545bccbbb34f3be09ff6fb6ab51b7b6"
[[package]]
name = "winapi-util"
version = "0.1.11"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22"
dependencies = [
"windows-sys 0.61.2",
]
[[package]]
name = "winapi-x86_64-pc-windows-gnu"
version = "0.4.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "712e227841d057c1ee1cd2fb22fa7e5a5461ae8e48fa2ca79ec42cfc1931183f"
[[package]]
name = "windows-link"
version = "0.2.1"

10
litellm-rust/Cargo.toml
View file

@ -18,6 +18,7 @@ litellm-ai-gateway = { path = "crates/ai-gateway", default-features = false }
axum = "0.7"
pyo3 = "0.29.0"
pyo3-async-runtimes = { version = "0.29.0", features = ["tokio-runtime"] }
pythonize = "0.29.0"
rand = "0.8"
reqwest = { version = "0.12", default-features = false, features = ["blocking", "json", "rustls-tls", "http2", "stream"] }
serde = { version = "1.0", features = ["derive"] }
@ -29,3 +30,12 @@ tokio = { version = "1", features = ["rt-multi-thread", "macros", "time", "net"]
tokio-tungstenite = { version = "0.24", default-features = false, features = ["connect", "rustls-tls-native-roots"] }
futures-util = { version = "0.3", default-features = false, features = ["sink", "std"] }
base64 = "0.22"
[profile.release]
opt-level = 3
lto = "thin"
codegen-units = 1
panic = "unwind"
debug = false
incremental = false
strip = "symbols"

16
litellm-rust/crates/python-bridge/Cargo.toml
View file

@ -9,10 +9,24 @@ repository.workspace = true
name = "_native"
crate-type = ["cdylib"]
[features]
default = ["abi3"]
abi3 = ["pyo3/abi3-py310"]
extension-module = ["pyo3/extension-module"]
[dependencies]
litellm-core = { workspace = true, features = ["bedrock-auth"] }
litellm-ai-gateway = { workspace = true, default-features = false }
pyo3 = { workspace = true, features = ["extension-module"] }
pyo3.workspace = true
pyo3-async-runtimes.workspace = true
pythonize.workspace = true
serde.workspace = true
serde_json.workspace = true
tokio.workspace = true
[dev-dependencies]
criterion = "0.8.2"
[[bench]]
name = "serialization"
harness = false

103
litellm-rust/crates/python-bridge/benches/serialization.rs Normal file
View file

@ -0,0 +1,103 @@
use std::hint::black_box;
use std::time::Duration;
use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main};
use pyo3::prelude::*;
use pyo3::types::PyDict;
use serde_json::{Value, json};
const PAYLOAD_SIZES: &[(&str, usize)] = &[
("1_KiB", 1024),
("64_KiB", 64 * 1024),
("1_MiB", 1024 * 1024),
("4_MiB", 4 * 1024 * 1024),
("16_MiB", 16 * 1024 * 1024),
];
fn former_json_roundtrip_from_py(py: Python<'_>, value: &Bound<'_, PyAny>) -> Value {
let json = py.import("json").expect("Python json module should import");
let encoded: String = json
.call_method1("dumps", (value,))
.expect("payload should serialize")
.extract()
.expect("json.dumps should return a string");
serde_json::from_str(&encoded).expect("serialized JSON should parse")
}
fn pythonize_from_py(value: &Bound<'_, PyAny>) -> Value {
pythonize::depythonize(value).expect("payload should depythonize")
}
fn former_json_roundtrip_to_py(py: Python<'_>, value: &Value) -> Py<PyAny> {
let json = py.import("json").expect("Python json module should import");
let encoded = serde_json::to_string(value).expect("response should serialize");
json.call_method1("loads", (encoded,))
.expect("serialized response should parse in Python")
.unbind()
}
fn pythonize_to_py(py: Python<'_>, value: &Value) -> Py<PyAny> {
pythonize::pythonize(py, value)
.expect("response should pythonize")
.unbind()
}
fn serialization(c: &mut Criterion) {
Python::initialize();
Python::attach(|py| {
for &(label, payload_bytes) in PAYLOAD_SIZES {
let data_uri = format!("data:image/png;base64,{}", "A".repeat(payload_bytes));
let document = PyDict::new(py);
document
.set_item("type", "image_url")
.expect("document type should be set");
document
.set_item("image_url", &data_uri)
.expect("document URL should be set");
let response = json!({
"pages": [{
"index": 0,
"markdown": "OCR text",
"images": [{"image_base64": data_uri}],
}],
"model": "mistral-ocr-latest",
"document_annotation": null,
"usage_info": {"pages_processed": 1},
"object": "ocr",
});
c.bench_with_input(
BenchmarkId::new("python_to_rust_json", label),
&document,
|b, document| {
b.iter(|| former_json_roundtrip_from_py(py, black_box(document.as_any())))
},
);
c.bench_with_input(
BenchmarkId::new("python_to_rust_pythonize", label),
&document,
|b, document| b.iter(|| pythonize_from_py(black_box(document.as_any()))),
);
c.bench_with_input(
BenchmarkId::new("rust_to_python_json", label),
&response,
|b, response| b.iter(|| former_json_roundtrip_to_py(py, black_box(response))),
);
c.bench_with_input(
BenchmarkId::new("rust_to_python_pythonize", label),
&response,
|b, response| b.iter(|| pythonize_to_py(py, black_box(response))),
);
}
});
}
criterion_group! {
name = benches;
config = Criterion::default()
.sample_size(20)
.warm_up_time(Duration::from_secs(1))
.measurement_time(Duration::from_secs(4));
targets = serialization
}
criterion_main!(benches);

48
litellm-rust/crates/python-bridge/src/lib.rs
View file

@ -19,6 +19,9 @@ use pyo3::types::{PyAny, PyDict};
use serde_json::{Map, Value};
mod gil;
mod marshal;
use marshal::{from_py, to_py};
pyo3::create_exception!(
_native,
@ -41,35 +44,18 @@ type MarshaledOcrInputs = (
Option<Duration>,
);
fn py_to_json(py: Python<'_>, value: &Bound<'_, PyAny>) -> PyResult<Value> {
let json = py.import("json")?;
let encoded: String = json.call_method1("dumps", (value,))?.extract()?;
serde_json::from_str(&encoded).map_err(|err| PyValueError::new_err(err.to_string()))
}
fn json_to_py(py: Python<'_>, value: Value) -> PyResult<Py<PyAny>> {
let json = py.import("json")?;
let encoded =
serde_json::to_string(&value).map_err(|err| PyValueError::new_err(err.to_string()))?;
Ok(json.call_method1("loads", (encoded,))?.unbind())
}
fn messages_response_to_py(
py: Python<'_>,
response: AnthropicMessagesResponse,
) -> PyResult<Py<PyAny>> {
let value =
serde_json::to_value(response).map_err(|err| PyValueError::new_err(err.to_string()))?;
json_to_py(py, value)
to_py(py, &response)
}
fn chat_completions_response_to_py(
py: Python<'_>,
response: ChatCompletionsResponse,
) -> PyResult<Py<PyAny>> {
let value =
serde_json::to_value(response).map_err(|err| PyValueError::new_err(err.to_string()))?;
json_to_py(py, value)
to_py(py, &response)
}
fn core_error_to_pyerr(err: CoreError) -> PyErr {
@ -116,7 +102,7 @@ fn optional_object_to_map(
value: Option<Py<PyAny>>,
) -> PyResult<Map<String, Value>> {
match value {
Some(value) => match py_to_json(py, value.bind(py))? {
Some(value) => match from_py(value.bind(py))? {
Value::Object(map) => Ok(map),
_ => Err(PyValueError::new_err(format!("{name} must be a dict"))),
},
@ -139,7 +125,7 @@ fn marshal_headers(
headers: Option<Py<PyAny>>,
) -> PyResult<HashMap<String, String>> {
let value = match headers {
Some(headers) => py_to_json(py, headers.bind(py))?,
Some(headers) => from_py(headers.bind(py))?,
None => Value::Object(Map::new()),
};
let Value::Object(headers) = value else {
@ -211,7 +197,7 @@ fn marshal_inputs(
optional_params: Option<Py<PyAny>>,
timeout_seconds: Option<f64>,
) -> PyResult<MarshaledOcrInputs> {
let document = py_to_json(py, document.bind(py))?;
let document = from_py(document.bind(py))?;
let extra_headers = match extra_headers {
Some(headers) => Some(optional_object_to_map(py, "extra_headers", Some(headers))?),
None => None,
@ -262,7 +248,7 @@ fn ocr(
});
match result {
Ok(value) => json_to_py(py, value),
Ok(value) => to_py(py, &value),
Err(err) => Err(core_error_to_pyerr(err)),
}
}
@ -307,7 +293,7 @@ fn aocr(
.await
.map_err(core_error_to_pyerr)?;
Python::attach(|py| json_to_py(py, value))
Python::attach(|py| to_py(py, &value))
})
}
@ -325,7 +311,7 @@ fn transcription(
optional_params: Option<Py<PyAny>>,
timeout_seconds: Option<f64>,
) -> PyResult<Py<PyAny>> {
let audio = py_to_json(py, audio.bind(py))?;
let audio = from_py(audio.bind(py))?;
let extra_headers = match extra_headers {
Some(headers) => Some(optional_object_to_map(py, "extra_headers", Some(headers))?),
None => None,
@ -351,7 +337,7 @@ fn transcription(
))
});
match result {
Ok(value) => json_to_py(py, value),
Ok(value) => to_py(py, &value),
Err(err) => Err(core_error_to_pyerr(err)),
}
}
@ -370,7 +356,7 @@ fn atranscription(
optional_params: Option<Py<PyAny>>,
timeout_seconds: Option<f64>,
) -> PyResult<Bound<'_, PyAny>> {
let audio = py_to_json(py, audio.bind(py))?;
let audio = from_py(audio.bind(py))?;
let extra_headers = match extra_headers {
Some(headers) => Some(optional_object_to_map(py, "extra_headers", Some(headers))?),
None => None,
@ -394,7 +380,7 @@ fn atranscription(
})
.await
.map_err(core_error_to_pyerr)?;
Python::attach(|py| json_to_py(py, value))
Python::attach(|py| to_py(py, &value))
})
}
@ -406,7 +392,7 @@ fn marshal_messages_inputs(
extra_headers: Option<Py<PyAny>>,
timeout_seconds: Option<f64>,
) -> PyResult<MarshaledMessagesInputs> {
let body = py_to_json(py, body.bind(py))?;
let body: Value = from_py(body.bind(py))?;
if !body.is_object() {
return Err(PyValueError::new_err("body must be a dict"));
}
@ -498,7 +484,7 @@ fn marshal_chat_completions_inputs(
extra_headers: Option<Py<PyAny>>,
timeout_seconds: Option<f64>,
) -> PyResult<MarshaledChatCompletionsInputs> {
let messages = py_to_json(py, messages.bind(py))?;
let messages: Value = from_py(messages.bind(py))?;
if !messages.is_array() {
return Err(PyValueError::new_err("messages must be a list"));
}
@ -527,7 +513,7 @@ fn chat_completions_decline(
optional_params: Option<Py<PyAny>>,
custom_llm_provider: Option<String>,
) -> PyResult<Option<String>> {
let messages = py_to_json(py, messages.bind(py))?;
let messages = from_py(messages.bind(py))?;
let optional_params = optional_object_to_map(py, "optional_params", optional_params)?;
Ok(chat_completions_decline_reason(
&model,

20
litellm-rust/crates/python-bridge/src/marshal.rs Normal file
View file

@ -0,0 +1,20 @@
use pyo3::exceptions::PyValueError;
use pyo3::prelude::*;
use serde::Serialize;
use serde::de::DeserializeOwned;
pub fn from_py<T>(value: &Bound<'_, PyAny>) -> PyResult<T>
where
T: DeserializeOwned,
{
pythonize::depythonize(value).map_err(|error| PyValueError::new_err(error.to_string()))
}
pub fn to_py<T>(py: Python<'_>, value: &T) -> PyResult<Py<PyAny>>
where
T: Serialize + ?Sized,
{
pythonize::pythonize(py, value)
.map(Bound::unbind)
.map_err(|error| PyValueError::new_err(error.to_string()))
}

52
litellm-rust/crates/python-bridge/tests/marshal_boundary.rs Normal file
View file

@ -0,0 +1,52 @@
use std::fs;
use std::path::{Path, PathBuf};
const DISALLOWED_OUTSIDE_MARSHAL: &[&str] = &[
"py.import(\"json\")",
"pythonize::",
"serde_json::to_string",
"serde_json::from_str",
];
fn source_root() -> PathBuf {
Path::new(env!("CARGO_MANIFEST_DIR")).join("src")
}
fn rust_sources(directory: &Path) -> Vec<PathBuf> {
fs::read_dir(directory)
.expect("bridge source directory should be readable")
.map(|entry| {
entry
.expect("bridge source entry should be readable")
.path()
})
.flat_map(|path| {
if path.is_dir() {
rust_sources(&path)
} else if path.extension().is_some_and(|extension| extension == "rs") {
vec![path]
} else {
Vec::new()
}
})
.collect()
}
#[test]
fn serialization_is_centralized_in_marshal_module() {
let root = source_root();
for path in rust_sources(&root) {
if path == root.join("marshal.rs") {
continue;
}
let source = fs::read_to_string(&path).expect("bridge source should be readable");
for disallowed in DISALLOWED_OUTSIDE_MARSHAL {
assert!(
!source.contains(disallowed),
"{} bypasses the typed marshal module with `{disallowed}`",
path.display()
);
}
}
}

31
litellm/__init__.py
View file

@ -7,6 +7,9 @@ warnings.filterwarnings("ignore", message=".*conflict with protected namespace.*
# Suppress Pydantic 2.11+ deprecation warning about accessing model_fields on instances
# This warning can accumulate during streaming and cause memory leaks
warnings.filterwarnings("ignore", message=".*Accessing the.*attribute on the instance is deprecated.*")
# ReadOnly on TypedDict fields is repo-wide static discipline (LIT012); pydantic warns it
# cannot enforce it at runtime, which floods proxy boot once such a type is schema-walked
warnings.filterwarnings("ignore", message=".*`ReadOnly` qualifier.*")
### INIT VARIABLES #########################
import threading
import os
@ -656,6 +659,8 @@ aiml_models: Set = set()
deepgram_models: Set = set()
elevenlabs_models: Set = set()
dashscope_models: Set = set()
qwencloud_models: Set = set()
qwen_ai_platform_models: Set = set()
moonshot_models: Set = set()
publicai_models: Set = set()
darkbloom_models: Set = set()
@ -906,6 +911,10 @@ def _populate_provider_model_sets(model_cost_map: Dict) -> None:
heroku_models.add(key)
elif value.get("litellm_provider") == "dashscope":
dashscope_models.add(key)
elif value.get("litellm_provider") == "qwencloud":
qwencloud_models.add(key)
elif value.get("litellm_provider") == "qwen_ai_platform":
qwen_ai_platform_models.add(key)
elif value.get("litellm_provider") == "modelscope":
modelscope_models.add(key)
elif value.get("litellm_provider") == "moonshot":
@ -1069,6 +1078,8 @@ model_list = list(
| deepgram_models
| elevenlabs_models
| dashscope_models
| qwencloud_models
| qwen_ai_platform_models
| moonshot_models
| publicai_models
| darkbloom_models
@ -1175,6 +1186,8 @@ def _build_models_by_provider() -> dict:
"elevenlabs": elevenlabs_models,
"heroku": heroku_models,
"dashscope": dashscope_models,
"qwencloud": qwencloud_models,
"qwen_ai_platform": qwen_ai_platform_models,
"modelscope": modelscope_models,
"moonshot": moonshot_models,
"publicai": publicai_models,
@ -2011,6 +2024,24 @@ if TYPE_CHECKING:
from .llms.dashscope.rerank.transformation import (
DashScopeRerankConfig as DashScopeRerankConfig,
)
from .llms.dashscope.qwencloud import (
QwenCloudChatConfig as QwenCloudChatConfig,
)
from .llms.dashscope.qwencloud import (
QwenCloudEmbeddingConfig as QwenCloudEmbeddingConfig,
)
from .llms.dashscope.qwencloud import (
QwenCloudRerankConfig as QwenCloudRerankConfig,
)
from .llms.dashscope.qwen_ai_platform import (
QwenAIPlatformChatConfig as QwenAIPlatformChatConfig,
)
from .llms.dashscope.qwen_ai_platform import (
QwenAIPlatformEmbeddingConfig as QwenAIPlatformEmbeddingConfig,
)
from .llms.dashscope.qwen_ai_platform import (
QwenAIPlatformRerankConfig as QwenAIPlatformRerankConfig,
)
from .llms.modelscope.chat.transformation import (
ModelScopeChatConfig as ModelScopeChatConfig,
)

92
litellm/_lazy_imports.py
View file

@ -17,8 +17,11 @@ until they're actually needed.
import importlib
import sys
from collections.abc import Callable
from typing import Any, Final, cast
from collections.abc import Callable, Mapping
from types import ModuleType
from typing import TYPE_CHECKING, Any, Final, cast
from typing_extensions import ReadOnly, TypedDict
# Import all the data structures that define what can be lazy-loaded
# These are just lists of names and maps of where to find them
@ -53,8 +56,12 @@ from ._lazy_imports_registry import (
UTILS_NAMES,
)
if TYPE_CHECKING:
import httpx
from tiktoken import Encoding
def get_litellm_globals() -> dict:
def get_litellm_globals() -> dict[str, object]:
"""
Get the globals dictionary of the litellm module.
@ -64,7 +71,7 @@ def get_litellm_globals() -> dict:
return sys.modules["litellm"].__dict__
def _get_utils_globals() -> dict:
def _get_utils_globals() -> dict[str, object]:
"""
Get the globals dictionary of the utils module.
@ -74,14 +81,19 @@ def _get_utils_globals() -> dict:
return sys.modules["litellm.utils"].__dict__
def _get_module_level_client_timeout(litellm_globals: Mapping[str, Any]) -> "float | httpx.Timeout | None":
"""Read the configured `litellm.request_timeout` used for the module level http clients."""
return litellm_globals.get("request_timeout")
# These are special lazy loaders for things that are used internally
# They're separate from the main lazy import system because they have specific use cases
# Lazy loader for default encoding - avoids importing heavy tiktoken library at startup
_default_encoding: Any | None = None
_default_encoding: "Encoding | None" = None
def _get_default_encoding() -> Any:
def _get_default_encoding() -> "Encoding":
"""
Lazily load and cache the default OpenAI encoding.
@ -100,10 +112,10 @@ def _get_default_encoding() -> Any:
# Lazy loader for get_modified_max_tokens to avoid importing token_counter at module import time
_get_modified_max_tokens_func: Any | None = None
_get_modified_max_tokens_func: "Callable[..., int | None] | None" = None
def _get_modified_max_tokens() -> Any:
def _get_modified_max_tokens() -> "Callable[..., int | None]":
"""
Lazily load and cache the get_modified_max_tokens function.
@ -124,10 +136,10 @@ def _get_modified_max_tokens() -> Any:
# Lazy loader for token_counter to avoid importing token_counter module at module import time
_token_counter_new_func: Any | None = None
_token_counter_new_func: "Callable[..., int] | None" = None
def _get_token_counter_new() -> Any:
def _get_token_counter_new() -> "Callable[..., int]":
"""
Lazily load and cache the token_counter function (aliased as token_counter_new).
@ -154,10 +166,10 @@ def _get_token_counter_new() -> Any:
# This registry maps attribute names (like "ModelResponse") to handler functions
# It's built once the first time someone accesses a lazy-loaded attribute
# Example: {"ModelResponse": _lazy_import_utils, "Cache": _lazy_import_caching, ...}
_LAZY_IMPORT_REGISTRY: dict[str, Callable[[str], Any]] | None = None
_LAZY_IMPORT_REGISTRY: dict[str, Callable[[str], object]] | None = None
def _get_lazy_import_registry() -> dict[str, Callable[[str], Any]]:
def _get_lazy_import_registry() -> dict[str, Callable[[str], object]]:
"""
Build the registry that maps attribute names to their handler functions.
@ -206,7 +218,18 @@ def _get_lazy_import_registry() -> dict[str, Callable[[str], Any]]:
return _LAZY_IMPORT_REGISTRY
def _generic_lazy_import(name: str, import_map: dict[str, tuple[str, str]], category: str) -> Any:
class _AttributeView(TypedDict):
"""Holds one module attribute so the lazily fetched value is read back as ``object``."""
value: ReadOnly[object]
def _module_attribute(module: ModuleType, attr_name: str) -> object:
attribute: Final[_AttributeView] = {"value": getattr(module, attr_name)}
return attribute["value"]
def _generic_lazy_import(name: str, import_map: dict[str, tuple[str, str]], category: str) -> object:
"""
Generic function that handles lazy importing for most attributes.
@ -255,7 +278,7 @@ def _generic_lazy_import(name: str, import_map: dict[str, tuple[str, str]], cate
# Step 6: Get the actual attribute from the module
# Example: getattr(utils_module, "ModelResponse") returns the ModelResponse class
value: Final = getattr(module, attr_name)
value: Final = _module_attribute(module, attr_name)
# Step 7: Cache it so we don't have to import again next time
_globals[name] = value
@ -272,62 +295,62 @@ def _generic_lazy_import(name: str, import_map: dict[str, tuple[str, str]], cate
# The registry (above) maps attribute names to these handler functions.
def _lazy_import_utils(name: str) -> Any:
def _lazy_import_utils(name: str) -> object:
"""Handler for utils module attributes (ModelResponse, token_counter, etc.)"""
return _generic_lazy_import(name, _UTILS_IMPORT_MAP, "Utils")
def _lazy_import_cost_calculator(name: str) -> Any:
def _lazy_import_cost_calculator(name: str) -> object:
"""Handler for cost calculator functions (completion_cost, cost_per_token, etc.)"""
return _generic_lazy_import(name, _COST_CALCULATOR_IMPORT_MAP, "Cost calculator")
def _lazy_import_token_counter(name: str) -> Any:
def _lazy_import_token_counter(name: str) -> object:
"""Handler for token counter utilities"""
return _generic_lazy_import(name, _TOKEN_COUNTER_IMPORT_MAP, "Token counter")
def _lazy_import_bedrock_types(name: str) -> Any:
def _lazy_import_bedrock_types(name: str) -> object:
"""Handler for Bedrock type aliases"""
return _generic_lazy_import(name, _BEDROCK_TYPES_IMPORT_MAP, "Bedrock types")
def _lazy_import_types_utils(name: str) -> Any:
def _lazy_import_types_utils(name: str) -> object:
"""Handler for types from litellm.types.utils (BudgetConfig, ImageObject, etc.)"""
return _generic_lazy_import(name, _TYPES_UTILS_IMPORT_MAP, "Types utils")
def _lazy_import_caching(name: str) -> Any:
def _lazy_import_caching(name: str) -> object:
"""Handler for caching classes (Cache, DualCache, RedisCache, etc.)"""
return _generic_lazy_import(name, _CACHING_IMPORT_MAP, "Caching")
def _lazy_import_dotprompt(name: str) -> Any:
def _lazy_import_dotprompt(name: str) -> object:
"""Handler for dotprompt integration globals"""
return _generic_lazy_import(name, _DOTPROMPT_IMPORT_MAP, "Dotprompt")
def _lazy_import_types(name: str) -> Any:
def _lazy_import_types(name: str) -> object:
"""Handler for type classes (GuardrailItem, etc.)"""
return _generic_lazy_import(name, _TYPES_IMPORT_MAP, "Types")
def _lazy_import_llm_configs(name: str) -> Any:
def _lazy_import_llm_configs(name: str) -> object:
"""Handler for LLM config classes (AnthropicConfig, OpenAILikeChatConfig, etc.)"""
return _generic_lazy_import(name, _LLM_CONFIGS_IMPORT_MAP, "LLM config")
def _lazy_import_litellm_logging(name: str) -> Any:
def _lazy_import_litellm_logging(name: str) -> object:
"""Handler for litellm_logging module (Logging, modify_integration)"""
return _generic_lazy_import(name, _LITELLM_LOGGING_IMPORT_MAP, "Litellm logging")
def _lazy_import_llm_provider_logic(name: str) -> Any:
def _lazy_import_llm_provider_logic(name: str) -> object:
"""Handler for LLM provider logic functions (get_llm_provider, etc.)"""
return _generic_lazy_import(name, _LLM_PROVIDER_LOGIC_IMPORT_MAP, "LLM provider logic")
def _lazy_import_utils_module(name: str) -> Any:
def _lazy_import_utils_module(name: str) -> object:
"""
Handler for utils module lazy imports.
@ -355,7 +378,7 @@ def _lazy_import_utils_module(name: str) -> Any:
module = importlib.import_module(module_path)
# Get the actual attribute from the module
value: Final = getattr(module, attr_name)
value: Final = _module_attribute(module, attr_name)
# Cache it so we don't have to import again next time
_globals[name] = value
@ -370,7 +393,7 @@ def _lazy_import_utils_module(name: str) -> Any:
# These handlers have custom logic that doesn't fit the generic pattern
def _lazy_import_llm_client_cache(name: str) -> Any:
def _lazy_import_llm_client_cache(name: str) -> object:
"""
Handler for LLM client cache - has special logic for singleton instance.
@ -386,8 +409,7 @@ def _lazy_import_llm_client_cache(name: str) -> Any:
return _globals[name]
# Import the class
module: Final = importlib.import_module("litellm.caching.llm_caching_handler")
LLMClientCache: Final = getattr(module, "LLMClientCache")
from litellm.caching.llm_caching_handler import LLMClientCache
# If they want the class itself, return it
if name == "LLMClientCache":
@ -403,7 +425,7 @@ def _lazy_import_llm_client_cache(name: str) -> Any:
raise AttributeError(f"LLM client cache lazy import: unknown attribute {name!r}")
def _lazy_import_http_handlers(name: str) -> Any:
def _lazy_import_http_handlers(name: str) -> object:
"""
Handler for HTTP clients - has special logic for creating client instances.
@ -419,8 +441,8 @@ def _lazy_import_http_handlers(name: str) -> Any:
from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
# Get timeout from module config (if set)
timeout = _globals.get("request_timeout")
params: Final = {"timeout": timeout, "client_alias": "module level aclient"}
async_timeout: Final = _get_module_level_client_timeout(_globals)
params: Final = {"timeout": async_timeout, "client_alias": "module level aclient"}
# Create the client instance
provider_id: Final = cast(Any, "litellm_module_level_client")
@ -437,8 +459,8 @@ def _lazy_import_http_handlers(name: str) -> Any:
# Create a sync HTTP client
from litellm.llms.custom_httpx.http_handler import HTTPHandler
timeout = _globals.get("request_timeout")
sync_client: Final = HTTPHandler(timeout=timeout)
sync_timeout: Final = _get_module_level_client_timeout(_globals)
sync_client: Final = HTTPHandler(timeout=sync_timeout)
# Cache it
_globals["module_level_client"] = sync_client

10
litellm/_lazy_imports_registry.py
View file

@ -310,6 +310,8 @@ LLM_CONFIG_NAMES: Final = (
"GigaChatConfig",
"GigaChatEmbeddingConfig",
"DashScopeChatConfig",
"QwenCloudChatConfig",
"QwenAIPlatformChatConfig",
"ModelScopeChatConfig",
"MoonshotChatConfig",
"DockerModelRunnerChatConfig",
@ -1172,6 +1174,14 @@ _LLM_CONFIGS_IMPORT_MAP: Final = {
".llms.dashscope.chat.transformation",
"DashScopeChatConfig",
),
"QwenCloudChatConfig": (
".llms.dashscope.qwencloud",
"QwenCloudChatConfig",
),
"QwenAIPlatformChatConfig": (
".llms.dashscope.qwen_ai_platform",
"QwenAIPlatformChatConfig",
),
"GDCGeminiConfig": (
".llms.gdc.chat.transformation",
"GDCGeminiConfig",

15
litellm/_logging.py
View file

@ -264,13 +264,17 @@ def _plain_log_format(stdout: TextIO | None, stderr: TextIO | None) -> str:
class LevelRoutingStreamHandler(logging.StreamHandler):
"""Writes records below WARNING to stdout and WARNING and above to stderr.
"""Writes records below WARNING and invalid-key warnings to stdout, others to stderr.
Collectors that derive severity from the stream report every stderr line as an error.
Invalid-key warnings route to stdout so LITELLM_LOG=ERROR can suppress them.
"""
def emit(self, record: logging.LogRecord) -> None:
preferred: Final = sys.stdout if record.levelno < logging.WARNING else sys.stderr
is_stdout_record: Final = record.levelno < logging.WARNING or (
record.levelno == logging.WARNING and record.name == verbose_proxy_stdout_logger.name
)
preferred: Final = sys.stdout if is_stdout_record else sys.stderr
if preferred is None or getattr(preferred, "closed", False):
self.stream = sys.stderr # rebind-ok: fall back to the pre-fix stream rather than raising per record
else:
@ -508,6 +512,9 @@ else:
handler.setFormatter(formatter)
verbose_proxy_logger = logging.getLogger("LiteLLM Proxy")
# Malformed virtual key rejections log through this child; LevelRoutingStreamHandler
# writes its WARNING records to stdout. It has no handler or level of its own.
verbose_proxy_stdout_logger: Final = verbose_proxy_logger.getChild("stdout")
verbose_router_logger = logging.getLogger("LiteLLM Router")
verbose_logger = logging.getLogger("LiteLLM")
@ -520,6 +527,7 @@ verbose_logger.addHandler(handler)
# handlers (JSON mode, uvicorn log config, a host app's root handler).
verbose_router_logger.addFilter(_stdout_truncation_filter)
verbose_proxy_logger.addFilter(_stdout_truncation_filter)
verbose_proxy_stdout_logger.addFilter(_stdout_truncation_filter)
verbose_logger.addFilter(_stdout_truncation_filter)
@ -683,6 +691,7 @@ def _turn_on_json():
- Adds a JSON formatter to all loggers
"""
handler: Final = LevelRoutingStreamHandler()
handler.setLevel(numeric_level)
handler.setFormatter(JsonFormatter())
_initialize_loggers_with_handler(handler)
# Set up exception handlers
@ -700,12 +709,14 @@ def _disable_debugging():
verbose_logger.disabled = True
verbose_router_logger.disabled = True
verbose_proxy_logger.disabled = True
verbose_proxy_stdout_logger.disabled = True
def _enable_debugging():
verbose_logger.disabled = False
verbose_router_logger.disabled = False
verbose_proxy_logger.disabled = False
verbose_proxy_stdout_logger.disabled = False
def print_verbose(print_statement):

125
litellm/_redis.py
View file

@ -13,6 +13,7 @@ import json
# s/o [@Frank Colson](https://www.linkedin.com/in/frank-colson-422b9b183/) for this redis implementation
import os
from collections.abc import Callable, Mapping
from types import MappingProxyType
from typing import Final
from urllib.parse import urlsplit, urlunsplit
@ -38,9 +39,25 @@ from ._logging import verbose_logger
AZURE_REDIS_SCOPE: Final = "https://redis.azure.com/.default"
def _get_redis_kwargs():
arg_spec: Final = inspect.getfullargspec(redis.Redis)
def _unwrapped_init_args(cls: type) -> frozenset[str]:
"""Every parameter on a single class's own ``__init__``, decorator-unwrapped.
Unlike ``_init_arg_names`` below, this does not walk the MRO: ``redis.Redis``
and ``redis.RedisCluster`` (sync and async) each declare every real
constructor parameter directly on their own ``__init__``, so MRO-walking is
unnecessary — and it actively breaks the several tests here that mock the
class with ``patch(..., autospec=True)``, since ``inspect.getmro`` needs a
real ``__mro__`` that an autospec'd stand-in for a class does not provide.
Still unwraps first: redis-py >= 7.4 decorates these ``__init__``s with
``@deprecated_args`` too, which the same class of bug as ``_init_arg_names``
would otherwise silently empty this allowlist through (see its docstring).
"""
spec: Final = inspect.getfullargspec(inspect.unwrap(cls.__init__))
return frozenset(spec.args + spec.kwonlyargs)
def _get_redis_kwargs():
# Only allow primitive arguments
exclude_args: Final = {
"self",
@ -60,7 +77,7 @@ def _get_redis_kwargs():
"azure_client_secret",
}
available_args: Final = {x for x in arg_spec.args if x not in exclude_args} | include_args
available_args: Final = {x for x in _unwrapped_init_args(redis.Redis) if x not in exclude_args} | include_args
return available_args
@ -120,15 +137,23 @@ def _get_redis_url_kwargs(client: type | None = None) -> tuple[str, ...]:
return tuple(x for x in _init_arg_names(connection_cls) if x not in exclude_args) + include_args
def _get_redis_cluster_kwargs(client=None):
def _get_redis_cluster_kwargs(client: type | None = None):
"""Config kwargs the target cluster client's constructor actually accepts.
Defaults to the sync ``redis.RedisCluster``, but the async cluster client
(``redis.asyncio.cluster.RedisCluster``) declares connection settings such as
``decode_responses`` on its own constructor, where the sync class takes them
through ``**kwargs`` and so never names them in its signature. Introspecting
only the sync class regardless of which client is actually built silently
drops those for every async cluster caller.
"""
if client is None:
client = redis.Redis.from_url
arg_spec: Final = inspect.getfullargspec(redis.RedisCluster)
client = redis.RedisCluster
# Only allow primitive arguments
exclude_args: Final = {"self", "connection_pool", "retry", "host", "port", "startup_nodes"}
available_args = {x for x in arg_spec.args if x not in exclude_args}
available_args = {x for x in _unwrapped_init_args(client) if x not in exclude_args}
available_args |= {
"password",
"username",
@ -161,6 +186,79 @@ def _get_redis_env_kwarg_mapping():
return {f"{PREFIX}{x.upper()}": x for x in _get_redis_kwargs() if x not in exclude_from_environment}
def _str_to_bool(value: str) -> bool:
return value.lower() in ("true", "1", "yes")
def _coerce_redis_kwargs_types(
redis_kwargs: Mapping[str, object],
client: type | tuple[type, ...] = redis.Redis,
) -> dict[str, object]: # mutable-ok: a caller mutates the returned kwargs before constructing its client
"""Coerces string values to the numeric/boolean type ``client``'s constructor
declares for that parameter. ``client`` may be a tuple of client classes; a
parameter's type is taken from the first signature that declares it, which
lets cluster callers coerce cluster-only kwargs such as
``cluster_error_retry_attempts`` alongside the shared connection kwargs.
Environment variables are always strings, and Helm ``--set`` stringifies values
too, so a config value like ``health_check_interval`` or ``socket_timeout``
can arrive as ``"30"``/``"5.5"`` rather than a real number. redis-py's own
connection-health-check arithmetic (``loop.time() + self.health_check_interval``)
then raises ``TypeError`` on every Redis operation instead of connecting.
``max_connections``, ``socket_timeout``, and ``socket_connect_timeout`` use an
explicit target type rather than the parameter's own signature default: redis-py
8.x changed the timeout defaults from ``None`` to int ``5``, so inferring the
type from the default would make a fractional ``"5.5"`` fail ``int()`` and get
silently dropped on 8.x while working on older versions. ``socket_keepalive``
is explicit too: its signature default is ``None``, which carries no type to
infer from, and leaving it a string makes ``"false"`` truthy.
"""
signatures: Final = tuple(inspect.signature(c) for c in (client if isinstance(client, tuple) else (client,)))
explicit_param_types: Final = MappingProxyType(
{
"max_connections": int,
"socket_timeout": float,
"socket_connect_timeout": float,
"socket_keepalive": bool,
}
)
result: Final = dict(redis_kwargs) # mutable-ok: per-key try/except coercion below needs to drop individual keys
for key, value in redis_kwargs.items():
if not isinstance(value, str):
continue
param = next((sig.parameters[key] for sig in signatures if key in sig.parameters), None)
if param is None:
continue
explicit_type = explicit_param_types.get(key)
if explicit_type is bool:
result[key] = _str_to_bool(value)
continue
if explicit_type is not None:
try:
result[key] = explicit_type(value)
except (ValueError, TypeError):
del result[key]
continue
default: object = param.default # pyright: ignore[reportAny] # inspect.Parameter.default is stubbed as Any
if default is inspect.Parameter.empty:
continue
# bool must be checked before int, since bool subclasses int
if isinstance(default, bool):
result[key] = _str_to_bool(value)
elif isinstance(default, int):
try:
result[key] = int(value)
except (ValueError, TypeError):
del result[key]
elif isinstance(default, float):
try:
result[key] = float(value)
except (ValueError, TypeError):
del result[key]
return result
def _redis_kwargs_from_environment():
mapping: Final = _get_redis_env_kwarg_mapping()
@ -505,7 +603,12 @@ def _get_redis_client_logic(**env_overrides):
raise ValueError("Either 'host' or 'url' must be specified for redis.")
# litellm.print_verbose(f"redis_kwargs: {redis_kwargs}")
return redis_kwargs
coercion_client: Final = (
(redis.Redis, redis.RedisCluster, async_redis.RedisCluster)
if redis_kwargs.get("startup_nodes")
else redis.Redis
)
return _coerce_redis_kwargs_types(redis_kwargs, client=coercion_client)
def init_redis_cluster(redis_kwargs) -> redis.RedisCluster:
@ -657,7 +760,9 @@ def get_redis_client(**env_overrides):
if "sentinel_nodes" in redis_kwargs and "service_name" in redis_kwargs:
return _init_redis_sentinel(redis_kwargs)
return redis.Redis(**redis_kwargs)
return redis.Redis( # pyright: ignore[reportCallIssue] # object-valued kwargs match no overload statically
**redis_kwargs, # pyright: ignore[reportArgumentType] # allow-listed and coerced against this signature
)
def get_redis_async_client(
@ -669,7 +774,7 @@ def get_redis_async_client(
if "startup_nodes" in redis_kwargs:
from redis.cluster import ClusterNode
args = _get_redis_cluster_kwargs()
args = _get_redis_cluster_kwargs(async_redis.RedisCluster)
cluster_kwargs: Final = {}
for arg in redis_kwargs:
if arg in args:

132
litellm/a2a_protocol/litellm_completion_bridge/transformation.py
View file

@ -17,11 +17,27 @@ A2A Streaming Events:
- Artifact update (kind: "artifact-update") - Content/artifact delivery
"""
from collections.abc import Mapping, MutableMapping, Sequence
from datetime import datetime, timezone
from typing import Any, Final
from typing import TYPE_CHECKING, Final
from uuid import uuid4
from pydantic import JsonValue, TypeAdapter, ValidationError
from litellm._logging import verbose_logger
from litellm.types.utils import ModelResponse
if TYPE_CHECKING:
from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
_STR_KEY_MAPPING_ADAPTER: Final = TypeAdapter(Mapping[str, object])
def _as_object_mapping(value: object) -> Mapping[str, object]:
try:
return _STR_KEY_MAPPING_ADAPTER.validate_python(value)
except ValidationError:
return {}
class A2AStreamingContext:
@ -30,7 +46,7 @@ class A2AStreamingContext:
Tracks task_id, context_id, and message accumulation.
"""
def __init__(self, request_id: str, input_message: dict[str, Any]):
def __init__(self, request_id: str, input_message: Mapping[str, JsonValue]):
self.request_id = request_id
self.task_id = str(uuid4())
self.context_id = str(uuid4())
@ -46,44 +62,46 @@ class A2ACompletionBridgeTransformation:
"""
@staticmethod
def _extract_text_from_a2a_parts(parts: list[dict[str, Any]]) -> str:
def _text_from_a2a_part(part: JsonValue) -> str | None:
if not isinstance(part, dict):
return None
text: Final = part.get("text")
if text is None:
return None
if part.get("kind") not in (None, "", "text"):
return None
return str(text)
@staticmethod
def _extract_text_from_a2a_parts(parts: Sequence[JsonValue]) -> str:
"""Extract text from A2A parts (with or without explicit ``kind``)."""
content_parts: Final[list[str]] = []
for part in parts:
if not isinstance(part, dict):
continue
kind = part.get("kind")
text = part.get("text")
if text is None:
continue
if kind in (None, "", "text"):
content_parts.append(str(text))
return "\n".join(content_parts)
extracted: Final = (A2ACompletionBridgeTransformation._text_from_a2a_part(part) for part in parts)
return "\n".join(text for text in extracted if text is not None)
@staticmethod
def get_forward_metadata(
a2a_message: dict[str, Any],
params: dict[str, Any] | None = None,
) -> dict[str, Any] | None:
a2a_message: Mapping[str, JsonValue],
params: Mapping[str, JsonValue] | None = None,
) -> Mapping[str, JsonValue] | None:
"""
Merge A2A metadata from MessageSendParams and the message for downstream providers.
Forwarded once on the LangGraph run payload (``metadata``), not duplicated on
each input message — see ``apply_forward_metadata_to_completion_params``.
"""
merged: Final[dict[str, Any]] = {}
if params and isinstance(params.get("metadata"), dict):
merged.update(params["metadata"])
params_metadata: Final = params.get("metadata") if params else None
message_metadata: Final = a2a_message.get("metadata")
if isinstance(message_metadata, dict):
merged.update(message_metadata)
merged: Final[dict[str, JsonValue]] = {
**(params_metadata if isinstance(params_metadata, dict) else {}),
**(message_metadata if isinstance(message_metadata, dict) else {}),
}
return merged or None
@staticmethod
def apply_forward_metadata_to_completion_params(
completion_params: dict[str, Any],
a2a_message: dict[str, Any],
params: dict[str, Any] | None = None,
completion_params: MutableMapping[str, object],
a2a_message: Mapping[str, JsonValue],
params: Mapping[str, JsonValue] | None = None,
) -> None:
"""
Attach A2A metadata to completion kwargs for provider bridges (e.g. LangGraph).
@ -97,24 +115,20 @@ class A2ACompletionBridgeTransformation:
if not forward_metadata:
return
extra_body = completion_params.get("extra_body")
if not isinstance(extra_body, dict):
extra_body = {}
extra_body: Final = _as_object_mapping(completion_params.get("extra_body"))
# Layer client-supplied A2A metadata under any agent-owner-configured
# ``extra_body.metadata`` so the configured keys remain authoritative
# and an A2A caller cannot overwrite server-set run metadata.
existing_metadata: Final = extra_body.get("metadata")
existing_dict: Final[dict[str, Any]] = existing_metadata if isinstance(existing_metadata, dict) else {}
merged_metadata: Final[dict[str, Any]] = {**forward_metadata, **existing_dict}
extra_body = {**extra_body, "metadata": merged_metadata}
completion_params["extra_body"] = extra_body
existing_dict: Final = _as_object_mapping(extra_body.get("metadata"))
merged_metadata: Final[dict[str, object]] = {**forward_metadata, **existing_dict}
completion_params["extra_body"] = {**extra_body, "metadata": merged_metadata}
verbose_logger.debug("A2A -> completion forward metadata keys=%s", list(forward_metadata.keys()))
@staticmethod
def a2a_message_to_openai_messages(
a2a_message: dict[str, Any],
) -> list[dict[str, Any]]:
a2a_message: Mapping[str, JsonValue],
) -> list[dict[str, object]]:
"""
Transform an A2A message to OpenAI message format.
@ -125,25 +139,19 @@ class A2ACompletionBridgeTransformation:
List of OpenAI-format messages
"""
role: Final = a2a_message.get("role", "user")
parts = a2a_message.get("parts", [])
raw_parts: Final = a2a_message.get("parts", [])
# Map A2A roles to OpenAI roles
openai_role = role
if role == "user":
openai_role = "user"
elif role == "assistant":
openai_role = "assistant"
elif role == "system":
openai_role = "system"
if not isinstance(parts, list):
parts = []
openai_role: Final = (
"user" if role == "user" else "assistant" if role == "assistant" else "system" if role == "system" else role
)
parts: Final = raw_parts if isinstance(raw_parts, list) else []
content: Final = A2ACompletionBridgeTransformation._extract_text_from_a2a_parts(parts)
# Do not attach A2A message.metadata here — the completion bridge forwards it
# once at run level via extra_body.metadata (LangGraph POST /runs/wait shape).
openai_message: Final[dict[str, Any]] = {"role": openai_role, "content": content}
openai_message: Final[dict[str, object]] = {"role": openai_role, "content": content}
verbose_logger.debug(
"A2A -> OpenAI transform: role=%s -> %s, content_length=%s", role, openai_role, len(content)
@ -151,11 +159,20 @@ class A2ACompletionBridgeTransformation:
return [openai_message]
@staticmethod
def _extract_response_content(response: "ModelResponse | CustomStreamWrapper") -> str:
if not isinstance(response, ModelResponse) or not response.choices:
return ""
choice: Final = response.choices[0]
if not choice.message:
return ""
return choice.message.content or ""
@staticmethod
def openai_response_to_a2a_response(
response: Any,
response: "ModelResponse | CustomStreamWrapper",
request_id: str | None = None,
) -> dict[str, Any]:
) -> dict[str, object]:
"""
Transform a LiteLLM ModelResponse to A2A SendMessageResponse format.
@ -166,12 +183,7 @@ class A2ACompletionBridgeTransformation:
Returns:
A2A SendMessageResponse dict
"""
# Extract content from response
content = ""
if hasattr(response, "choices") and response.choices:
choice: Final = response.choices[0]
if hasattr(choice, "message") and choice.message:
content = choice.message.content or ""
content: Final = A2ACompletionBridgeTransformation._extract_response_content(response)
# Build A2A message
a2a_message: Final = {
@ -182,7 +194,7 @@ class A2ACompletionBridgeTransformation:
}
# Build A2A response
a2a_response: Final = {
a2a_response: Final[dict[str, object]] = {
"jsonrpc": "2.0",
"id": request_id,
"result": a2a_message,
@ -200,7 +212,7 @@ class A2ACompletionBridgeTransformation:
@staticmethod
def create_task_event(
ctx: A2AStreamingContext,
) -> dict[str, Any]:
) -> dict[str, object]:
"""
Create the initial task event with status 'submitted'.
@ -235,7 +247,7 @@ class A2ACompletionBridgeTransformation:
state: str,
final: bool = False,
message_text: str | None = None,
) -> dict[str, Any]:
) -> dict[str, object]:
"""
Create a status update event.
@ -245,7 +257,7 @@ class A2ACompletionBridgeTransformation:
final: Whether this is the final event
message_text: Optional message text for 'working' status
"""
status: Final[dict[str, Any]] = {
status: Final[dict[str, object]] = {
"state": state,
"timestamp": A2ACompletionBridgeTransformation._get_timestamp(),
}
@ -277,7 +289,7 @@ class A2ACompletionBridgeTransformation:
def create_artifact_update_event(
ctx: A2AStreamingContext,
text: str,
) -> dict[str, Any]:
) -> dict[str, object]:
"""
Create an artifact update event with content.

28
litellm/a2a_protocol/main.py
View file

@ -86,7 +86,7 @@ A2ACardResolver: Final = LiteLLMA2ACardResolver
def _set_usage_on_logging_obj(
kwargs: dict[str, Any],
kwargs: Mapping[str, object],
prompt_tokens: int,
completion_tokens: int,
) -> None:
@ -99,7 +99,7 @@ def _set_usage_on_logging_obj(
completion_tokens: Number of output tokens
"""
litellm_logging_obj: Final = kwargs.get("litellm_logging_obj")
if litellm_logging_obj is not None:
if isinstance(litellm_logging_obj, Logging):
usage: Final = litellm.Usage(
prompt_tokens=prompt_tokens,
completion_tokens=completion_tokens,
@ -109,7 +109,7 @@ def _set_usage_on_logging_obj(
def _set_agent_id_on_logging_obj(
kwargs: dict[str, Any],
kwargs: Mapping[str, object],
agent_id: str | None,
) -> None:
"""
@ -123,7 +123,7 @@ def _set_agent_id_on_logging_obj(
return
litellm_logging_obj: Final = kwargs.get("litellm_logging_obj")
if litellm_logging_obj is not None:
if isinstance(litellm_logging_obj, Logging):
# Set agent_id directly on model_call_details (same pattern as custom_llm_provider)
litellm_logging_obj.model_call_details["agent_id"] = agent_id
@ -132,7 +132,7 @@ _A2A_COST_PARAM_KEYS: Final = ("cost_per_query", "input_cost_per_token", "output
def _set_litellm_params_on_logging_obj(
kwargs: dict[str, Any],
kwargs: Mapping[str, object],
litellm_params: Mapping[str, object],
) -> None:
"""
@ -144,18 +144,22 @@ def _set_litellm_params_on_logging_obj(
context, so merge the pricing keys in rather than replacing the dict.
"""
logging_obj: Final = kwargs.get("litellm_logging_obj")
if logging_obj is None:
if not isinstance(logging_obj, Logging):
return
cost_params = {key: litellm_params[key] for key in _A2A_COST_PARAM_KEYS if litellm_params.get(key) is not None}
cost_params: Final = {
key: litellm_params[key] for key in _A2A_COST_PARAM_KEYS if litellm_params.get(key) is not None
}
if not cost_params:
return
existing: Final = logging_obj.model_call_details.get("litellm_params") or {}
logging_obj.model_call_details["litellm_params"] = {**existing, **cost_params}
logging_obj.model_call_details["litellm_params"] = {
**(logging_obj.model_call_details.get("litellm_params") or {}),
**cost_params,
}
def _get_a2a_model_info(a2a_client: "A2AClientType", kwargs: dict[str, Any]) -> str:
def _get_a2a_model_info(a2a_client: "A2AClientType", kwargs: Mapping[str, object]) -> str:
"""
Extract agent info and set model/custom_llm_provider for cost tracking.
@ -175,7 +179,7 @@ def _get_a2a_model_info(a2a_client: "A2AClientType", kwargs: dict[str, Any]) ->
# Set on litellm_logging_obj if available (for standard logging payload)
litellm_logging_obj: Final = kwargs.get("litellm_logging_obj")
if litellm_logging_obj is not None:
if isinstance(litellm_logging_obj, Logging):
litellm_logging_obj.model = model
litellm_logging_obj.custom_llm_provider = custom_llm_provider
litellm_logging_obj.model_call_details["model"] = model
@ -498,7 +502,7 @@ async def asend_message(
response: Final = LiteLLMSendMessageResponse.from_a2a_response(a2a_response, request_id=str(request.id))
# Calculate token usage from request and response
response_dict: Final[dict[str, object]] = a2a_response.model_dump(mode="json", exclude_none=True)
response_dict: Final[dict[str, object]] = a2a_response.root.model_dump(mode="json", exclude_none=True)
(
prompt_tokens,
completion_tokens,

227
litellm/batches/batch_utils.py
View file

@ -1,6 +1,8 @@
import json
from collections.abc import Iterable, Iterator, Mapping
from dataclasses import dataclass
from dataclasses import replace as dataclasses_replace
from enum import Enum
from typing import Any, Final, Literal
import litellm
@ -12,12 +14,23 @@ from litellm.types.utils import CallTypes, ModelInfo, Usage
from litellm.utils import token_counter
@dataclass(frozen=True, slots=True)
class BatchCostUsageResult:
"""Aggregate cost, usage, and per-line pass/fail counts for a completed batch."""
cost: float
usage: Usage
models: list[str]
successful_requests: int
failed_requests: int
async def calculate_batch_cost_and_usage(
file_content_dictionary: list[dict],
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"],
model_name: str | None = None,
model_info: ModelInfo | None = None,
) -> tuple[float, Usage, list[str]]:
) -> BatchCostUsageResult:
"""
Calculate the cost and usage of a batch.
@ -32,8 +45,7 @@ async def calculate_batch_cost_and_usage(
and model_name
and getattr(litellm, "disable_vertex_batch_output_transformation", False)
):
batch_cost, batch_usage = calculate_vertex_ai_batch_cost_and_usage(file_content_dictionary, model_name)
return batch_cost, batch_usage, [model_name]
return calculate_vertex_ai_batch_cost_and_usage(file_content_dictionary, model_name)
return _aggregate_batch_cost_usage_models(
entries=file_content_dictionary,
@ -49,7 +61,7 @@ async def _handle_completed_batch(
model_name: str | None = None,
litellm_params: dict | None = None,
model_info: ModelInfo | None = None,
) -> tuple[float, Usage, list[str]]:
) -> BatchCostUsageResult:
"""Fetch a completed batch's output file and aggregate its cost, usage, and
models in a single pass over the JSONL lines, so the parsed file content is
never materialized in memory.
@ -72,27 +84,49 @@ async def _handle_completed_batch(
# The generic retrieval helper keeps raising for callers that explicitly ask
# for a missing output file.
if batch.output_file_id is None:
return 0.0, Usage(prompt_tokens=0, completion_tokens=0, total_tokens=0), []
return BatchCostUsageResult(
cost=0.0,
usage=Usage(prompt_tokens=0, completion_tokens=0, total_tokens=0),
models=[], # mutable-ok: no output file means no model was ever priced; BatchCostUsageResult.models requires list[str]
successful_requests=0,
failed_requests=await count_error_file_failed_requests(
batch, custom_llm_provider=custom_llm_provider, litellm_params=litellm_params
),
)
file_content = await _fetch_batch_output_file_content(batch, custom_llm_provider, litellm_params=litellm_params)
if (
custom_llm_provider == "vertex_ai"
and model_name
and getattr(litellm, "disable_vertex_batch_output_transformation", False)
):
batch_cost, batch_usage = calculate_vertex_ai_batch_cost_and_usage(
_get_file_content_as_dictionary(file_content), model_name
)
return batch_cost, batch_usage, [model_name]
return _aggregate_batch_cost_usage_models(
entries=_iter_batch_output_entries(file_content),
custom_llm_provider=custom_llm_provider,
model_name=model_name,
model_info=model_info,
error_file_failed_requests: Final = await count_error_file_failed_requests(
batch, custom_llm_provider=custom_llm_provider, litellm_params=litellm_params
)
output_file_result: Final = (
calculate_vertex_ai_batch_cost_and_usage(_get_file_content_as_dictionary(file_content), model_name)
if (
custom_llm_provider == "vertex_ai"
and model_name
and getattr(litellm, "disable_vertex_batch_output_transformation", False)
)
else _aggregate_batch_cost_usage_models(
entries=_iter_batch_output_entries(file_content),
custom_llm_provider=custom_llm_provider,
model_name=model_name,
model_info=model_info,
)
)
if not error_file_failed_requests:
return output_file_result
return dataclasses_replace(
output_file_result, failed_requests=output_file_result.failed_requests + error_file_failed_requests
)
class _LineOutcome(Enum):
"""A batch output line that yielded no billable stats."""
PROVIDER_FAILED = "provider_failed"
UNCOSTABLE = "uncostable"
@dataclass(frozen=True, slots=True)
class _BatchOutputLineStats:
@ -102,19 +136,27 @@ class _BatchOutputLineStats:
total_tokens: int
cache_read_tokens: int
cache_creation_tokens: int
reasoning_tokens: int
model: str | None
def _iter_successful_output_line_stats(
def _classify_output_line_stats(
entries: Iterable[dict],
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock"],
model_name: str | None,
model_info: ModelInfo | None,
) -> Iterator[_BatchOutputLineStats]:
) -> Iterator[_BatchOutputLineStats | _LineOutcome]:
"""Classify every output line in a single pass, so counting failures never needs
a second read of a potentially huge output file. A line the provider reported as
failed yields ``PROVIDER_FAILED``; a successful line litellm could not price
yields ``UNCOSTABLE`` and still counts as a successful request billed at $0, so
the counts stay reconcilable with the provider's own ``request_counts``."""
for entry in entries:
if not _batch_response_was_successful(entry, custom_llm_provider):
yield _LineOutcome.PROVIDER_FAILED
continue
stats = _safe_output_line_stats(entry, custom_llm_provider, model_name, model_info)
if stats is not None:
yield stats
yield stats if stats is not None else _LineOutcome.UNCOSTABLE
def _safe_output_line_stats(
@ -123,13 +165,11 @@ def _safe_output_line_stats(
model_name: str | None,
model_info: ModelInfo | None,
) -> _BatchOutputLineStats | None:
"""Return the stats for one batch output line, or None for a line that is
unsuccessful or cannot be costed, so a single bad line never aborts the
whole batch's cost accounting."""
"""Return the stats for one provider-successful batch output line, or None when
it cannot be costed, so a single bad line never aborts the whole batch's cost
accounting."""
custom_id: Final = entry.get("custom_id") if isinstance(entry, dict) else None
try:
if not _batch_response_was_successful(entry, custom_llm_provider):
return None
return _compute_output_line_stats(entry, custom_llm_provider, model_name, model_info)
except Exception as e: # noqa: BLE001 # any single line's costing failure must not abort the whole batch
verbose_logger.warning(
@ -152,6 +192,7 @@ def _compute_output_line_stats(
prompt_details: Final = parse_prompt_tokens_details(usage)
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
return _BatchOutputLineStats(
cost=_output_line_cost(
response_body=response_body,
@ -166,6 +207,7 @@ def _compute_output_line_stats(
total_tokens=usage.total_tokens,
cache_read_tokens=prompt_details["cache_hit_tokens"],
cache_creation_tokens=prompt_details["cache_creation_tokens"],
reasoning_tokens=(completion_details.reasoning_tokens if completion_details else None) or 0,
model=response_model,
)
@ -203,10 +245,14 @@ def _aggregate_batch_cost_usage_models(
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock"],
model_name: str | None = None,
model_info: ModelInfo | None = None,
) -> tuple[float, Usage, list[str]]:
"""Aggregate cost, usage, and models from batch output entries in a single
pass, holding one small stats record per line instead of the parsed file."""
line_stats: Final = tuple(_iter_successful_output_line_stats(entries, custom_llm_provider, model_name, model_info))
) -> BatchCostUsageResult:
"""Aggregate cost, usage, models, and pass/fail counts from batch output
entries in a single pass, holding one small stats record per line instead
of the parsed file."""
all_results: Final = tuple(_classify_output_line_stats(entries, custom_llm_provider, model_name, model_info))
line_stats: Final = tuple(result for result in all_results if isinstance(result, _BatchOutputLineStats))
failed_requests: Final = sum(1 for result in all_results if result is _LineOutcome.PROVIDER_FAILED)
successful_requests: Final = len(all_results) - failed_requests
cache_token_params: Final = {
key: tokens
@ -220,18 +266,32 @@ def _aggregate_batch_cost_usage_models(
total_tokens=sum(stats.total_tokens for stats in line_stats),
prompt_tokens=sum(stats.prompt_tokens for stats in line_stats),
completion_tokens=sum(stats.completion_tokens for stats in line_stats),
reasoning_tokens=sum(stats.reasoning_tokens for stats in line_stats),
**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)
verbose_logger.debug("batch output aggregate: cost=%s usage=%s models=%s", total_cost, batch_usage, batch_models)
return total_cost, batch_usage, batch_models
verbose_logger.debug(
"batch output aggregate: cost=%s usage=%s models=%s successful=%d failed=%d",
total_cost,
batch_usage,
batch_models,
successful_requests,
failed_requests,
)
return BatchCostUsageResult(
cost=total_cost,
usage=batch_usage,
models=batch_models,
successful_requests=successful_requests,
failed_requests=failed_requests,
)
def calculate_vertex_ai_batch_cost_and_usage(
vertex_ai_batch_responses: list[dict],
model_name: str | None = None,
) -> tuple[float, Usage]:
) -> BatchCostUsageResult:
"""
Calculate both cost and usage from raw Vertex AI batch responses.
@ -242,6 +302,10 @@ def calculate_vertex_ai_batch_cost_and_usage(
{"request": ..., "response": {"candidates": [...], "usageMetadata": {...}}}
usageMetadata contains promptTokenCount, candidatesTokenCount, totalTokenCount.
A row with no ``response`` is counted as failed - the same signal already
used to skip it from cost/usage aggregation, since Vertex batch prediction
output doesn't establish a distinct error shape in this (non-default) path.
"""
from litellm.cost_calculator import batch_cost_calculator
@ -249,12 +313,16 @@ def calculate_vertex_ai_batch_cost_and_usage(
total_tokens = 0
prompt_tokens = 0
completion_tokens = 0
successful_requests = 0 # rebind-ok: loop accumulator, matches total_cost/total_tokens above
failed_requests = 0 # rebind-ok: loop accumulator, matches total_cost/total_tokens above
actual_model_name: Final = model_name or "gemini-2.0-flash-001"
for response in vertex_ai_batch_responses:
response_body = response.get("response")
if response_body is None:
failed_requests += 1
continue
successful_requests += 1
usage_metadata = response_body.get("usageMetadata", {})
_prompt = usage_metadata.get("promptTokenCount", 0) or 0
@ -282,17 +350,25 @@ def calculate_vertex_ai_batch_cost_and_usage(
total_tokens += _total
verbose_logger.info(
"vertex_ai batch cost: cost=%s, prompt=%d, completion=%d, total=%d",
"vertex_ai batch cost: cost=%s, prompt=%d, completion=%d, total=%d, successful=%d, failed=%d",
total_cost,
prompt_tokens,
completion_tokens,
total_tokens,
successful_requests,
failed_requests,
)
return total_cost, Usage(
total_tokens=total_tokens,
prompt_tokens=prompt_tokens,
completion_tokens=completion_tokens,
return BatchCostUsageResult(
cost=total_cost,
usage=Usage(
total_tokens=total_tokens,
prompt_tokens=prompt_tokens,
completion_tokens=completion_tokens,
),
models=[actual_model_name],
successful_requests=successful_requests,
failed_requests=failed_requests,
)
@ -322,6 +398,36 @@ def _provider_output_file_id(output_file_id: str) -> str:
return extracted
async def _fetch_batch_managed_file_content(
file_id: str,
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"] = "openai",
litellm_params: dict | None = None,
) -> bytes:
"""
Fetch a batch's output or error file and return its raw JSONL bytes.
Args:
file_id: The provider or unified (litellm-managed) file id to fetch
custom_llm_provider: The LLM provider
litellm_params: Optional litellm parameters containing credentials (api_key, api_base, etc.)
Required for Azure and other providers that need authentication
"""
from litellm.files.main import afile_content
# Build kwargs for afile_content with credentials from litellm_params
file_content_kwargs: Final = {
"file_id": _provider_output_file_id(file_id),
"custom_llm_provider": custom_llm_provider,
}
# Extract and add credentials for file access
credentials: Final = _extract_file_access_credentials(litellm_params)
file_content_kwargs.update(credentials)
_file_content: Final = await afile_content(**file_content_kwargs)
return _file_content.content
async def _fetch_batch_output_file_content(
batch: Batch,
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"] = "openai",
@ -336,25 +442,36 @@ async def _fetch_batch_output_file_content(
litellm_params: Optional litellm parameters containing credentials (api_key, api_base, etc.)
Required for Azure and other providers that need authentication
"""
from litellm.files.main import afile_content
if batch.output_file_id is None:
raise ValueError("Output file id is None cannot retrieve file content")
file_id: Final = _provider_output_file_id(batch.output_file_id)
return await _fetch_batch_managed_file_content(
batch.output_file_id, custom_llm_provider=custom_llm_provider, litellm_params=litellm_params
)
# Build kwargs for afile_content with credentials from litellm_params
file_content_kwargs: Final = {
"file_id": file_id,
"custom_llm_provider": custom_llm_provider,
}
# Extract and add credentials for file access
credentials: Final = _extract_file_access_credentials(litellm_params)
file_content_kwargs.update(credentials)
async def count_error_file_failed_requests(
batch: Batch,
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"],
litellm_params: dict | None,
) -> int:
"""Count failed requests reported only in the batch's separate error file.
_file_content: Final = await afile_content(**file_content_kwargs)
return _file_content.content
OpenAI-shaped batch providers write successful lines to ``output_file_id``
and per-request failures (e.g. a rejected param) to a distinct
``error_file_id`` - they never appear in the output file at all, so
counting failures from the output file alone silently undercounts them.
"""
if batch.error_file_id is None:
return 0
try:
error_file_content = await _fetch_batch_managed_file_content(
batch.error_file_id, custom_llm_provider=custom_llm_provider, litellm_params=litellm_params
)
except Exception as e: # noqa: BLE001 # a failed/missing error file must not abort cost tracking for the batch
verbose_logger.debug("Failed to fetch batch error file %s: %s", batch.error_file_id, e)
return 0
return sum(1 for _ in _iter_batch_input_lines(error_file_content))
def _extract_file_access_credentials(litellm_params: dict | None) -> dict:

2
litellm/batches/main.py
View file

@ -390,7 +390,7 @@ def _handle_retrieve_batch_providers_without_provider_config(
custom_llm_provider: Literal[
"openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "litellm_proxy", "anthropic"
] = "openai",
logging_obj: Any | None = None,
logging_obj: LiteLLMLoggingObj | None = None,
):
api_base: str | None = None
if custom_llm_provider in OPENAI_COMPATIBLE_BATCH_AND_FILES_PROVIDERS:

9
litellm/caching/caching.py
View file

@ -12,6 +12,7 @@ import hashlib
import json
import time
import traceback
from collections.abc import Mapping
from enum import Enum
from typing import Any, Final
@ -506,7 +507,7 @@ class Cache:
def _get_cache_logic(
self,
cached_result: Any | None,
cached_result: object | None,
max_age: float | None,
):
"""
@ -538,8 +539,8 @@ class Cache:
return cached_result
@staticmethod
def _get_safe_cache_lookup_kwargs(kwargs: dict[str, Any]) -> dict[str, Any]:
cache_lookup_kwargs: Final[dict[str, Any]] = {}
def _get_safe_cache_lookup_kwargs(kwargs: Mapping[str, object]) -> dict[str, object]:
cache_lookup_kwargs: Final[dict[str, object]] = {}
for prompt_kwarg in ("messages", "input"):
if prompt_kwarg in kwargs:
cache_lookup_kwargs[prompt_kwarg] = kwargs[prompt_kwarg]
@ -552,7 +553,7 @@ class Cache:
@staticmethod
def _update_metadata_from_cache_lookup_kwargs(
original_kwargs: dict[str, Any], cache_lookup_kwargs: dict[str, Any]
original_kwargs: Mapping[str, object], cache_lookup_kwargs: Mapping[str, object]
) -> None:
original_metadata: Final = original_kwargs.get("metadata")
cache_lookup_metadata: Final = cache_lookup_kwargs.get("metadata")

18
litellm/caching/qdrant_semantic_cache.py
View file

@ -12,7 +12,7 @@ import ast
import asyncio
import json
import os
from typing import TYPE_CHECKING, Any, Final, cast
from typing import TYPE_CHECKING, Any, Final, Protocol, cast
import litellm
from litellm._logging import print_verbose
@ -39,6 +39,12 @@ if TYPE_CHECKING:
from litellm.router import Router
class _QdrantCollectionDetailsResponse(Protocol):
"""The qdrant `/collections/{name}` response, whose body is kept as an opaque JSON object."""
def json(self) -> dict[str, object]: ...
class QdrantSemanticCache(BaseCache):
CACHE_KEY_FIELD_NAME = "litellm_cache_key"
embedding_max_input_tokens: int | None = None
@ -115,15 +121,15 @@ class QdrantSemanticCache(BaseCache):
raise ValueError(f"Error from qdrant checking if /collections exist {collection_exists.text}")
if collection_exists.json()["result"]["exists"]:
collection_details = self.sync_client.get(
collection_details: _QdrantCollectionDetailsResponse = self.sync_client.get(
url=f"{self.qdrant_api_base}/collections/{self.collection_name}",
headers=self.headers,
)
self.collection_info = collection_details.json()
self.collection_info: dict[str, object] = collection_details.json()
print_verbose(f"Collection already exists.\nCollection details:{self.collection_info}")
self._ensure_cache_key_payload_index()
else:
quantization_params: dict[str, Any]
quantization_params: dict[str, dict[str, object]]
if quantization_config is None or quantization_config == "binary":
quantization_params = {
"binary": {
@ -214,7 +220,7 @@ class QdrantSemanticCache(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) -> EmbeddingResponse:
def _get_embedding(self, prompt: str, metadata: dict[str, object] | None = None) -> EmbeddingResponse:
"""Embed via the proxy Router when it serves the model, else direct."""
try:
from litellm.proxy.proxy_server import llm_model_list, llm_router
@ -241,7 +247,7 @@ class QdrantSemanticCache(BaseCache):
num_retries=0,
)
async def _get_async_embedding(self, prompt: str, metadata: dict[str, Any] | None = None) -> EmbeddingResponse:
async def _get_async_embedding(self, prompt: str, metadata: dict[str, object] | None = None) -> EmbeddingResponse:
try:
from litellm.proxy.proxy_server import llm_model_list, llm_router
except ImportError:

57
litellm/caching/redis_cache.py
View file

@ -18,7 +18,7 @@ import time
from collections.abc import Awaitable, Callable, Sequence
from contextvars import ContextVar
from datetime import timedelta
from typing import TYPE_CHECKING, Any, Final, TypeVar, cast
from typing import TYPE_CHECKING, Any, Final, Protocol, TypeVar, cast
import litellm
from litellm._logging import print_verbose, verbose_logger
@ -58,6 +58,26 @@ else:
Span = Any
class _AsyncRedisCommands(Protocol):
"""Async redis commands this cache issues.
redis-py's type stubs omit these methods on RedisCluster, so the union returned by
init_async_client() is untyped at every call site without this protocol.
"""
def ping(self) -> Awaitable[bool]: ...
def delete(self, *names: str) -> Awaitable[int]: ...
def ttl(self, name: str) -> Awaitable[int]: ...
def rpush(self, name: str, *values: str | bytes | float) -> Awaitable[int]: ...
def lpop(self, name: str, count: int | None = None) -> Awaitable[object]: ...
def pipeline(self, transaction: bool = True) -> "Pipeline[bytes]": ...
def _get_call_stack_info(num_frames: int = 2) -> str:
"""
Get the function names from the previous 1-2 functions in the call stack.
@ -429,6 +449,9 @@ class RedisCache(BaseCache):
self.redis_async_client = redis_async_client
return redis_async_client
def _async_commands(self) -> _AsyncRedisCommands:
return self.init_async_client()
def check_and_fix_namespace(self, key: str) -> str:
"""
Make sure each key starts with the given namespace
@ -1055,19 +1078,17 @@ class RedisCache(BaseCache):
await self.async_set_cache_pipeline(self.redis_batch_writing_buffer)
self.redis_batch_writing_buffer = []
def _get_cache_logic(self, cached_response: Any):
def _get_cache_logic(self, cached_response: bytes | str | None):
"""
Common 'get_cache_logic' across sync + async redis client implementations
"""
if cached_response is None:
return cached_response
# cached_response is in `b{} convert it to ModelResponse
cached_response = cached_response.decode("utf-8") # Convert bytes to string
return None
decoded: Final = cached_response.decode("utf-8") if isinstance(cached_response, bytes) else cached_response
try:
cached_response = json.loads(cached_response) # Convert string to dictionary
return json.loads(decoded)
except Exception:
cached_response = ast.literal_eval(cached_response)
return cached_response
return ast.literal_eval(decoded)
def get_cache(self, key, parent_otel_span: Span | None = None, **kwargs):
try:
@ -1314,8 +1335,7 @@ class RedisCache(BaseCache):
raise e
async def ping(self) -> bool:
# typed as Any, redis python lib has incomplete type stubs for RedisCluster and does not include `ping`
_redis_client: Final[Any] = self.init_async_client()
_redis_client: Final = self._async_commands()
start_time: Final = time.time()
print_verbose("Pinging Async Redis Cache")
try:
@ -1349,8 +1369,7 @@ class RedisCache(BaseCache):
@_redis_circuit_breaker_guard
async def delete_cache_keys(self, keys):
# typed as Any, redis python lib has incomplete type stubs for RedisCluster and does not include `delete`
_redis_client: Final[Any] = self.init_async_client()
_redis_client: Final = self._async_commands()
keys = [self.check_and_fix_namespace(key=key) for key in keys]
# keys is a list, unpack it so it gets passed as individual elements to delete
await _redis_client.delete(*keys)
@ -1415,8 +1434,7 @@ class RedisCache(BaseCache):
@_redis_circuit_breaker_guard
async def async_delete_cache(self, key: str):
# typed as Any, redis python lib has incomplete type stubs for RedisCluster and does not include `delete`
_redis_client: Final[Any] = self.init_async_client()
_redis_client: Final = self._async_commands()
key = self.check_and_fix_namespace(key=key)
# keys is str
return await _redis_client.delete(key)
@ -1523,8 +1541,7 @@ class RedisCache(BaseCache):
Redis ref: https://redis.io/docs/latest/commands/ttl/
"""
try:
# typed as Any, redis python lib has incomplete type stubs for RedisCluster and does not include `ttl`
_redis_client: Final[Any] = self.init_async_client()
_redis_client: Final = self._async_commands()
key = self.check_and_fix_namespace(key=key)
ttl: Final = await _redis_client.ttl(key)
if ttl <= -1: # -1 means the key does not exist, -2 key does not exist
@ -1554,7 +1571,7 @@ class RedisCache(BaseCache):
Returns:
int: The length of the list after the push operation
"""
_redis_client: Final[Any] = self.init_async_client()
_redis_client: Final = self._async_commands()
key = self.check_and_fix_namespace(key=key)
start_time: Final = time.time()
try:
@ -1621,7 +1638,7 @@ class RedisCache(BaseCache):
if len(rpush_list) == 0:
return []
_redis_client: Final[Any] = self.init_async_client()
_redis_client: Final = self._async_commands()
start_time: Final = time.time()
try:
@ -1678,7 +1695,7 @@ class RedisCache(BaseCache):
parent_otel_span: Span | None = None,
**kwargs,
) -> Any | list[Any]:
_redis_client: Final[Any] = self.init_async_client()
_redis_client: Final = self._async_commands()
key = self.check_and_fix_namespace(key=key)
start_time: Final = time.time()
print_verbose(f"LPOP from Redis list: key: {key}, count: {count}")
@ -1810,7 +1827,7 @@ class RedisCache(BaseCache):
if len(lpop_list) == 0:
return []
_redis_client: Final[Any] = self.init_async_client()
_redis_client: Final = self._async_commands()
start_time: Final = time.time()
try:

85
litellm/caching/valkey_semantic_cache.py
View file

@ -17,6 +17,7 @@ RedisSemanticCache since those are backend agnostic.
import asyncio
import hashlib
import os
from collections.abc import Mapping, Sequence
from dataclasses import dataclass
from typing import Any, Final
@ -64,7 +65,7 @@ class ValkeySemanticCache(RedisSemanticCache):
async_client: AsyncRedis | None = None,
embedding_max_input_tokens: int | None = None,
embedding_timeout: float | None = None,
**kwargs: Any,
**kwargs: object,
):
if similarity_threshold is None:
raise ValueError("similarity_threshold must be provided, passed None")
@ -87,11 +88,13 @@ class ValkeySemanticCache(RedisSemanticCache):
self.key_prefix = f"{self.index_name}:"
self._index_dim: int | None = None
resolved_url = None
if sync_client is None or async_client is None:
resolved_url = redis_url or self._build_valkey_url(host, port, password, ssl)
self.sync_client = sync_client if sync_client is not None else Redis.from_url(resolved_url)
self.async_client = async_client if async_client is not None else AsyncRedis.from_url(resolved_url)
if sync_client is not None and async_client is not None:
self.sync_client = sync_client
self.async_client = async_client
else:
resolved_url: Final = redis_url or self._build_valkey_url(host, port, password, ssl)
self.sync_client = sync_client if sync_client is not None else Redis.from_url(resolved_url)
self.async_client = async_client if async_client is not None else AsyncRedis.from_url(resolved_url)
print_verbose(f"Valkey semantic-cache initializing index - {self.index_name}")
@ -118,7 +121,7 @@ class ValkeySemanticCache(RedisSemanticCache):
return hashlib.sha256(str(key).encode("utf-8")).hexdigest()
@staticmethod
def _embedding_to_bytes(embedding: list[float]) -> bytes:
def _embedding_to_bytes(embedding: Sequence[float]) -> bytes:
return pack_vector(embedding)
def _index_schema(self, dim: int) -> tuple[TagField, VectorField]:
@ -192,7 +195,9 @@ class ValkeySemanticCache(RedisSemanticCache):
def _doc_key(self, key: str) -> str:
return f"{self.key_prefix}{self._scope_tag(key)}:{uuid.uuid4()}"
def _doc_mapping(self, key: str, prompt: str, value_str: str, embedding: list[float]) -> dict:
def _doc_mapping(
self, key: str, prompt: str, value_str: str, embedding: Sequence[float]
) -> Mapping[str | bytes, str | bytes]:
return {
self.CACHE_KEY_FIELD_NAME: self._scope_tag(key),
self.PROMPT_FIELD_NAME: prompt,
@ -208,30 +213,49 @@ class ValkeySemanticCache(RedisSemanticCache):
)
return Query(query_string).return_fields(self.RESPONSE_FIELD_NAME, self.DISTANCE_FIELD_NAME).dialect(2)
async def _async_search(self, key: str, embedding: Sequence[float]) -> object:
"""Run the KNN query on the async client, stopping the untyped search surface here."""
return await self.async_client.ft(self.index_name).search(
self._knn_query(key),
query_params={"vec": self._embedding_to_bytes(embedding)}, # pyright: ignore[reportArgumentType] # redis stubs omit bytes; KNN vectors are raw bytes at runtime
)
@classmethod
def _first_hit(cls, search_result: Any) -> _ValkeyCacheHit | None:
docs: Final = getattr(search_result, "docs", [])
def _first_hit(cls, search_result: object) -> _ValkeyCacheHit | None:
docs: Final[Sequence[object]] = getattr(search_result, "docs", [])
if not docs:
return None
doc: Final = docs[0]
response_field: Final[object] = getattr(doc, cls.RESPONSE_FIELD_NAME)
distance_field: Final[str | bytes | float] = getattr(doc, cls.DISTANCE_FIELD_NAME)
return _ValkeyCacheHit(
response=str(getattr(doc, cls.RESPONSE_FIELD_NAME)),
distance=float(getattr(doc, cls.DISTANCE_FIELD_NAME)),
response=str(response_field),
distance=float(distance_field),
)
def _resolve_hit(self, hit: _ValkeyCacheHit | None, key: str, **kwargs: Any) -> Any:
@staticmethod
def _record_similarity(kwargs: dict[str, Any], similarity: float) -> None:
"""Stamp the semantic-similarity score onto the request metadata carried in ``kwargs``."""
kwargs.setdefault("metadata", {})["semantic-similarity"] = similarity
@staticmethod
def _embedding_metadata(kwargs: dict[str, Any]) -> dict[str, Any] | None:
"""The request metadata forwarded to the embedding call."""
return kwargs.get("metadata")
def _resolve_hit(self, hit: _ValkeyCacheHit | None, key: str, **kwargs: object) -> object:
if hit is None:
kwargs.setdefault("metadata", {})["semantic-similarity"] = 0.0
self._record_similarity(kwargs, 0.0)
return None
similarity: Final = 1 - hit.distance
kwargs.setdefault("metadata", {})["semantic-similarity"] = similarity
self._record_similarity(kwargs, similarity)
if similarity < self.similarity_threshold:
return None
return self._get_cache_logic(cached_response=hit.response)
def set_cache(self, key: str, value: Any, **kwargs: Any) -> None:
def set_cache(self, key: str, value: object, **kwargs: object) -> None:
print_verbose(f"Valkey semantic-cache set_cache, kwargs: {kwargs}")
try:
prompt: Final = self._get_prompt_from_kwargs(**kwargs)
@ -250,12 +274,12 @@ class ValkeySemanticCache(RedisSemanticCache):
except Exception as e:
print_verbose(f"Error in Valkey semantic-cache set_cache: {e}")
def get_cache(self, key: str, **kwargs: Any) -> Any:
def get_cache(self, key: str, **kwargs: object) -> object:
print_verbose(f"Valkey semantic-cache get_cache, kwargs: {kwargs}")
try:
prompt: Final = self._get_prompt_from_kwargs(**kwargs)
if prompt is None:
kwargs.setdefault("metadata", {})["semantic-similarity"] = 0.0
self._record_similarity(kwargs, 0.0)
return None
embedding: Final = self._get_embedding(prompt)
@ -263,14 +287,14 @@ class ValkeySemanticCache(RedisSemanticCache):
search_result: Final = self.sync_client.ft(self.index_name).search(
self._knn_query(key),
query_params={"vec": self._embedding_to_bytes(embedding)},
query_params={"vec": self._embedding_to_bytes(embedding)}, # pyright: ignore[reportArgumentType] # redis stubs omit bytes; KNN vectors are raw bytes at runtime
)
return self._resolve_hit(self._first_hit(search_result), key, **kwargs)
except Exception as e:
print_verbose(f"Error in Valkey semantic-cache get_cache: {e}")
kwargs.setdefault("metadata", {})["semantic-similarity"] = 0.0
self._record_similarity(kwargs, 0.0)
async def async_set_cache(self, key: str, value: Any, **kwargs: Any) -> None:
async def async_set_cache(self, key: str, value: object, **kwargs: object) -> None:
print_verbose(f"Async Valkey semantic-cache set_cache, kwargs: {kwargs}")
try:
prompt: Final = self._get_prompt_from_kwargs(**kwargs)
@ -278,7 +302,7 @@ class ValkeySemanticCache(RedisSemanticCache):
print_verbose("No prompt provided for semantic caching")
return
embedding: Final = await self._get_async_embedding(prompt, metadata=kwargs.get("metadata"))
embedding: Final = await self._get_async_embedding(prompt, metadata=self._embedding_metadata(kwargs))
await self._ensure_index_async(len(embedding))
doc_key: Final = self._doc_key(key)
@ -289,31 +313,28 @@ class ValkeySemanticCache(RedisSemanticCache):
except Exception as e:
print_verbose(f"Error in async Valkey semantic-cache set_cache: {e}")
async def async_get_cache(self, key: str, **kwargs: Any) -> Any:
async def async_get_cache(self, key: str, **kwargs: object) -> object:
print_verbose(f"Async Valkey semantic-cache get_cache, kwargs: {kwargs}")
try:
prompt: Final = self._get_prompt_from_kwargs(**kwargs)
if prompt is None:
kwargs.setdefault("metadata", {})["semantic-similarity"] = 0.0
self._record_similarity(kwargs, 0.0)
return None
embedding: Final = await self._get_async_embedding(prompt, metadata=kwargs.get("metadata"))
embedding: Final = await self._get_async_embedding(prompt, metadata=self._embedding_metadata(kwargs))
await self._ensure_index_async(len(embedding))
search_result: Final = await self.async_client.ft(self.index_name).search(
self._knn_query(key),
query_params={"vec": self._embedding_to_bytes(embedding)},
)
search_result: Final[object] = await self._async_search(key, embedding)
return self._resolve_hit(self._first_hit(search_result), key, **kwargs)
except Exception as e:
print_verbose(f"Error in async Valkey semantic-cache get_cache: {e}")
kwargs.setdefault("metadata", {})["semantic-similarity"] = 0.0
self._record_similarity(kwargs, 0.0)
async def async_set_cache_pipeline(self, cache_list: list[tuple[str, Any]], **kwargs: Any) -> None:
async def async_set_cache_pipeline(self, cache_list: list[tuple[str, object]], **kwargs: object) -> None:
try:
await asyncio.gather(*[self.async_set_cache(key, value, **kwargs) for key, value in cache_list])
except Exception as e:
print_verbose(f"Error in Valkey semantic-cache async_set_cache_pipeline: {e}")
async def _index_info(self) -> dict:
async def _index_info(self) -> Mapping[str, object]:
return await self.async_client.ft(self.index_name).info()

14
litellm/completion_extras/litellm_responses_transformation/handler.py
View file

@ -45,14 +45,14 @@ class ResponsesToCompletionBridgeHandler:
return bool(stream)
@staticmethod
def _is_preformatted_cached_chat_stream(result: Any) -> bool:
def _is_preformatted_cached_chat_stream(result: object) -> bool:
from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
return isinstance(result, CustomStreamWrapper) and result.custom_llm_provider == "cached_response"
@staticmethod
def _coerce_response_object(
response_obj: Any,
response_obj: object,
hidden_params: dict | None,
) -> "ResponsesAPIResponse":
if isinstance(response_obj, ResponsesAPIResponse):
@ -78,8 +78,8 @@ class ResponsesToCompletionBridgeHandler:
for _ in stream_iter:
pass
completed: Final = getattr(stream_iter, "completed_response", None)
response_obj: Final = getattr(completed, "response", None) if completed else None
completed: Final[object] = getattr(stream_iter, "completed_response", None)
response_obj: Final[object] = getattr(completed, "response", None) if completed else None
if response_obj is None:
raise ValueError("Stream ended without a completed response")
@ -93,8 +93,8 @@ class ResponsesToCompletionBridgeHandler:
async for _ in stream_iter:
pass
completed: Final = getattr(stream_iter, "completed_response", None)
response_obj: Final = getattr(completed, "response", None) if completed else None
completed: Final[object] = getattr(stream_iter, "completed_response", None)
response_obj: Final[object] = getattr(completed, "response", None) if completed else None
if response_obj is None:
raise ValueError("Stream ended without a completed response")
@ -157,7 +157,7 @@ class ResponsesToCompletionBridgeHandler:
def completion(
self, *args, **kwargs
) -> Union[
Coroutine[Any, Any, Union["ModelResponse", "CustomStreamWrapper"]],
Coroutine[None, None, Union["ModelResponse", "CustomStreamWrapper"]],
"ModelResponse",
"CustomStreamWrapper",
]:

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

@ -212,7 +212,8 @@ def _tool_call_dict_from_output_item(item: Mapping[str, Any], index: int) -> _Ch
LiteLLMCompletionResponsesConfig,
)
is_custom: Final = item.get("type") == "custom_tool_call"
item_type: Final[object] = item.get("type")
is_custom: Final = item_type == "custom_tool_call"
arguments: Final = (item.get("input") if is_custom else item.get("arguments")) or ""
name: Final = item.get("name") or ("custom_tool" if is_custom else "")
function_chunk: Final = ChatCompletionToolCallFunctionChunk(name=name, arguments=arguments)
@ -222,7 +223,7 @@ def _tool_call_dict_from_output_item(item: Mapping[str, Any], index: int) -> _Ch
function=function_chunk,
index=index,
)
raw_provider_fields: Final = item.get("provider_specific_fields")
raw_provider_fields: Final[object] = item.get("provider_specific_fields")
if isinstance(raw_provider_fields, dict):
provider_specific_fields = raw_provider_fields
elif raw_provider_fields and hasattr(raw_provider_fields, "__dict__"):
@ -507,7 +508,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
def _merge_responses_api_request_into_request_data(
self,
request_data: dict[str, Any],
request_data: dict[str, object],
responses_api_request: "ResponsesAPIOptionalRequestParams",
instructions: str | None,
) -> None:

49
litellm/constants.py
View file

@ -13,6 +13,12 @@ DEFAULT_BATCH_SIZE: Final = int(os.getenv("DEFAULT_BATCH_SIZE", 512))
DEFAULT_FLUSH_INTERVAL_SECONDS: Final = int(os.getenv("DEFAULT_FLUSH_INTERVAL_SECONDS", 5))
DEFAULT_S3_FLUSH_INTERVAL_SECONDS: Final = int(os.getenv("DEFAULT_S3_FLUSH_INTERVAL_SECONDS", 10))
DEFAULT_S3_BATCH_SIZE: Final = int(os.getenv("DEFAULT_S3_BATCH_SIZE", 512))
# https://docs.aws.amazon.com/AmazonS3/latest/userguide/object-keys.html
MAX_S3_OBJECT_KEY_BYTES: Final = 1024
S3_BOUNDED_OBJECT_KEY_HEAD_BYTES: Final = 64
S3_PREFIX_DIGEST_CHARS: Final = 16
# s3 allows 2048 bytes of combined metadata headers, which Content-Disposition counts against
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))
DYNAMIC_RATE_LIMIT_ERROR_THRESHOLD_PER_MINUTE = int(os.getenv("DYNAMIC_RATE_LIMIT_ERROR_THRESHOLD_PER_MINUTE", 1))
@ -35,6 +41,7 @@ DEFAULT_COOLDOWN_TIME_SECONDS: Final = int(os.getenv("DEFAULT_COOLDOWN_TIME_SECO
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))
HF_CONFIG_FETCH_TIMEOUT_SECONDS: Final = 10.0
# Maximum wall-clock seconds a streaming response is allowed to run.
# Streams exceeding this duration are terminated with a Timeout error.
@ -129,6 +136,7 @@ MCP_CLIENT_TIMEOUT: Final = float(os.getenv("LITELLM_MCP_CLIENT_TIMEOUT", "60.0"
MCP_TOOL_LISTING_TIMEOUT: Final = float(os.getenv("LITELLM_MCP_TOOL_LISTING_TIMEOUT", "30.0"))
MCP_METADATA_TIMEOUT: Final = float(os.getenv("LITELLM_MCP_METADATA_TIMEOUT", "10.0"))
MCP_HEALTH_CHECK_TIMEOUT: Final = float(os.getenv("LITELLM_MCP_HEALTH_CHECK_TIMEOUT", "10.0"))
MCP_TOOL_LISTING_MAX_PAGES: Final = 1000
# Allowlist of commands permitted for MCP stdio transport.
# Prevents arbitrary command execution via /mcp-rest/test/* endpoints or server creation.
@ -288,6 +296,7 @@ REDIS_DAILY_ORG_SPEND_UPDATE_BUFFER_KEY: Final = "litellm_daily_org_spend_update
REDIS_DAILY_END_USER_SPEND_UPDATE_BUFFER_KEY: Final = "litellm_daily_end_user_spend_update_buffer"
REDIS_DAILY_AGENT_SPEND_UPDATE_BUFFER_KEY: Final = "litellm_daily_agent_spend_update_buffer"
REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY: Final = "litellm_daily_tag_spend_update_buffer"
REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY: Final = "litellm_window_spend_update_buffer"
MAX_REDIS_BUFFER_DEQUEUE_COUNT: Final = int(os.getenv("MAX_REDIS_BUFFER_DEQUEUE_COUNT", 100))
# Bounds asyncio.Queue() instances (log queues, spend update queues, etc.) to prevent unbounded memory growth
LITELLM_ASYNCIO_QUEUE_MAXSIZE: Final = int(os.getenv("LITELLM_ASYNCIO_QUEUE_MAXSIZE", 1000))
@ -482,6 +491,22 @@ FIREWORKS_AI_80_B: Final = int(os.getenv("FIREWORKS_AI_80_B", 80))
#### Logging callback constants ####
REDACTED_BY_LITELM_STRING: Final = "REDACTED_BY_LITELM"
MAX_LANGFUSE_INITIALIZED_CLIENTS: Final = int(os.getenv("MAX_LANGFUSE_INITIALIZED_CLIENTS", 50))
# Backpressure + lifetime bounds for the /v1/messages streaming relay (see
# BaseAnthropicMessagesStreamingIterator.async_sse_wrapper). The relay queue is
# bounded so a slow client throttles the upstream pump instead of letting it
# buffer the whole response in memory; the detached-drain cap bounds how many
# post-disconnect drains may run concurrently so client behavior can't create
# unbounded worker state.
ANTHROPIC_MESSAGES_STREAM_RELAY_QUEUE_MAXSIZE: Final = int(
os.getenv("ANTHROPIC_MESSAGES_STREAM_RELAY_QUEUE_MAXSIZE", "1024")
)
# Setting this to 0 disables detached draining entirely: every post-disconnect
# pump bills whatever partial output it has already collected and aborts the
# upstream stream immediately, instead of continuing to drain for the real
# terminal usage.
ANTHROPIC_MESSAGES_MAX_DETACHED_STREAM_DRAINS: Final = int(
os.getenv("ANTHROPIC_MESSAGES_MAX_DETACHED_STREAM_DRAINS", "100")
)
LOGGING_WORKER_CONCURRENCY: Final = int(os.getenv("LOGGING_WORKER_CONCURRENCY", 100)) # Must be above 0
LOGGING_WORKER_MAX_QUEUE_SIZE: Final = int(os.getenv("LOGGING_WORKER_MAX_QUEUE_SIZE", 50_000))
LOGGING_WORKER_MAX_TIME_PER_COROUTINE: Final = float(os.getenv("LOGGING_WORKER_MAX_TIME_PER_COROUTINE", 20.0))
@ -612,6 +637,8 @@ LITELLM_CHAT_PROVIDERS: Final = [
"nscale",
"nebius",
"dashscope",
"qwencloud",
"qwen_ai_platform",
"modelscope",
"moonshot",
"publicai",
@ -781,6 +808,7 @@ openai_compatible_endpoints: Final[list] = [
"inference.api.nscale.com/v1",
"api.studio.nebius.ai/v1",
"https://dashscope-intl.aliyuncs.com/compatible-mode/v1",
"https://dashscope.aliyuncs.com/compatible-mode/v1",
"https://api-inference.modelscope.cn/v1",
"https://api.moonshot.ai/v1",
"https://api.publicai.co/v1",
@ -804,6 +832,7 @@ openai_compatible_endpoints: Final[list] = [
"https://api.meta.ai/v1",
"https://api.cognition.ai/v1",
"https://api.scx.ai/v1",
"https://gigachat.devices.sberbank.ru/api/v1",
]
@ -853,6 +882,8 @@ openai_compatible_providers: Final[list] = [
"nscale",
"nebius",
"dashscope",
"qwencloud",
"qwen_ai_platform",
"modelscope",
"moonshot",
"v0",
@ -883,6 +914,8 @@ openai_text_completion_compatible_providers: Final[list] = [ # providers that s
"featherless_ai",
"nebius",
"dashscope",
"qwencloud",
"qwen_ai_platform",
"modelscope",
"moonshot",
"publicai",
@ -1090,7 +1123,7 @@ nebius_models: Final[set] = set(
]
)
dashscope_models: Final[set] = set(
dashscope_models: Final[frozenset] = frozenset(
[
"qwen-turbo",
"qwen-plus",
@ -1105,6 +1138,10 @@ dashscope_models: Final[set] = set(
]
)
qwencloud_models: Final[frozenset] = frozenset(dashscope_models)
qwen_ai_platform_models: Final[frozenset] = frozenset(dashscope_models)
nebius_embedding_models: Final[set] = set(
[
"BAAI/bge-en-icl",
@ -1221,6 +1258,7 @@ BEDROCK_CONVERSE_MODELS: Final = [
"openai.gpt-oss-120b-1:0",
"anthropic.claude-haiku-4-5-20251001-v1:0",
"anthropic.claude-sonnet-4-5-20250929-v1:0",
"anthropic.claude-fable-5-1",
"anthropic.claude-fable-5",
"anthropic.claude-sonnet-5",
"anthropic.claude-opus-5",
@ -1408,6 +1446,12 @@ DEFAULT_SOFT_BUDGET: Final = float(
) # by default all litellm proxy keys have a soft budget of 50.0
# makes it clear this is a rate limit error for a litellm virtual key
RATE_LIMIT_ERROR_MESSAGE_FOR_VIRTUAL_KEY: Final = "LiteLLM Virtual Key user_api_key_hash"
# Prefix of the 401 raised when a submitted virtual key is not shaped like one.
INVALID_VIRTUAL_KEY_ERROR_MESSAGE: Final = "LiteLLM Virtual Key expected"
# Attribute stamped on that 401 at its raise site so log routing recognises it by
# provenance. Message text is caller-influenceable on other 401s, so it must not
# be used to classify.
INVALID_VIRTUAL_KEY_ERROR_MARKER: Final = "_litellm_invalid_virtual_key_error"
# Python garbage collection threshold configuration
# Format: "gen0,gen1,gen2" e.g., "1000,50,50"
@ -1546,6 +1590,7 @@ KEY_ROTATION_JOB_NAME: Final = "litellm_key_rotation_job"
EXPIRED_UI_SESSION_KEY_CLEANUP_JOB_NAME: Final = "litellm_expired_ui_session_key_cleanup_job"
WEEKLY_SPEND_REPORT_JOB_ID: Final = "weekly_spend_report_job"
MONTHLY_SPEND_REPORT_JOB_ID: Final = "monthly_spend_report_job"
USER_SPEND_ALERTS_JOB_ID: Final = "user_spend_alerts_job"
PROMETHEUS_FALLBACK_STATS_JOB_ID: Final = "prometheus_fallback_stats_job"
SLACK_DAILY_REPORT_LOCK_ID: Final = "slack_daily_report"
SLACK_MODEL_DEPRECATION_LOCK_ID: Final = "slack_model_deprecation_warning"
@ -1681,6 +1726,7 @@ DEFAULT_MCP_NAMESPACE_CSV_MAX_TOKENS: Final = 16
# Ceilings on the cached auth registries; larger tables fall back to per-row lookups
# instead of holding an unbounded id set in every worker.
TAG_REGISTRY_MAX_SIZE: Final = 5000
MODEL_ACCESS_GROUP_REGISTRY_MAX_SIZE: Final = 5000
END_USER_RESTRICTED_REGISTRY_MAX_SIZE: Final = 5000
# How long a failed registry load is remembered as "unusable", so a degraded Postgres
# is not re-scanned on every request on top of the per-id lookups it falls back to.
@ -1725,6 +1771,7 @@ SENTRY_DENYLIST: Final = [
"jwt_token",
"private_key",
"SLACK_WEBHOOK_URL",
"ALERTING_WEBHOOK_URL",
"webhook_url",
"LANGFUSE_SECRET_KEY",
# Email Configuration

34
litellm/cost_calculator.py
View file

@ -641,12 +641,12 @@ def cost_per_token(
return xai_cost_per_token(model=model, usage=usage_block)
elif custom_llm_provider == "lemonade":
return lemonade_cost_per_token(model=model, usage=usage_block)
elif custom_llm_provider == "dashscope":
elif custom_llm_provider in ("dashscope", "qwencloud", "qwen_ai_platform"):
from litellm.llms.dashscope.cost_calculator import (
cost_per_token as dashscope_cost_per_token,
)
return dashscope_cost_per_token(model=model, usage=usage_block)
return dashscope_cost_per_token(model=model, usage=usage_block, custom_llm_provider=custom_llm_provider)
elif custom_llm_provider == "azure_ai":
return azure_ai_cost_per_token(
model=model,
@ -1910,12 +1910,15 @@ def ocr_cost(
if credits is not None and cost_per_credit is not None:
return cost_per_credit * credits, 0.0
ocr_cost_per_page: float | None = None
if model_info is not None:
ocr_cost_per_page = model_info.get("ocr_cost_per_page")
ocr_cost_per_page: Final = model_info.get("ocr_cost_per_page") if model_info is not None else None
annotation_cost_per_page: Final = model_info.get("annotation_cost_per_page") if model_info is not None else None
annotation_rate: Final = annotation_cost_per_page if annotation_cost_per_page is not None else ocr_cost_per_page
pages_processed: Final = response.usage_info.pages_processed
if pages_processed is None:
annotation_pages: Final = response.usage_info.pages_processed_annotation or 0
has_billable_annotation_pages: Final = annotation_rate is not None and annotation_pages > 0
if pages_processed is None and not has_billable_annotation_pages:
if cost_per_credit is not None or ocr_cost_per_page is None:
# Surface missing usage data instead of silently under-reporting
# cost. The previous behavior raised ValueError; we now return 0.0
@ -1931,7 +1934,7 @@ def ocr_cost(
return 0.0, 0.0
raise ValueError("OCR response pages_processed is None")
if ocr_cost_per_page is None:
if ocr_cost_per_page is None and not has_billable_annotation_pages:
# No per-page pricing configured. Either the model is on credit-based
# pricing (and credits weren't returned, so the credit branch above did
# not match) or the model has no OCR pricing entry at all. Surface a
@ -1947,8 +1950,9 @@ def ocr_cost(
)
return 0.0, 0.0
total_ocr_processing_cost: Final[float] = ocr_cost_per_page * pages_processed
return total_ocr_processing_cost, 0.0
ocr_pages_cost: Final = (ocr_cost_per_page or 0.0) * (pages_processed or 0)
annotation_pages_cost: Final = (annotation_rate or 0.0) * annotation_pages
return ocr_pages_cost + annotation_pages_cost, 0.0
def vector_store_search_cost(
@ -2268,6 +2272,10 @@ def batch_cost_calculator(
return total_prompt_cost, total_completion_cost
def _attribute_value(obj: object, name: str) -> object:
return getattr(obj, name)
def _summable_prompt_token_fields(prompt_tokens_details: BaseModel) -> list[str]:
field_names: Final = list(type(prompt_tokens_details).model_fields)
if getattr(prompt_tokens_details, "cache_write_tokens", None) is None:
@ -2293,7 +2301,7 @@ class BaseTokenUsageProcessor:
for usage in usage_objects:
# Handle direct attributes by checking what exists in the model
for attr in dir(usage):
if not attr.startswith("_") and not callable(getattr(usage, attr)):
if not attr.startswith("_") and not callable(_attribute_value(usage, attr)):
current_val = getattr(combined, attr, 0)
new_val = getattr(usage, attr, 0)
if (
@ -2313,7 +2321,7 @@ class BaseTokenUsageProcessor:
if (
hasattr(usage.prompt_tokens_details, attr)
and not attr.startswith("_")
and not callable(getattr(usage.prompt_tokens_details, attr))
and not callable(_attribute_value(usage.prompt_tokens_details, attr))
):
current_val = getattr(combined.prompt_tokens_details, attr, 0) or 0
new_val = getattr(usage.prompt_tokens_details, attr, 0) or 0
@ -2332,7 +2340,9 @@ class BaseTokenUsageProcessor:
# Check what keys exist in the model's completion_tokens_details
# Access model_fields on the class, not the instance, to avoid Pydantic 2.11+ deprecation warnings
for attr in type(usage.completion_tokens_details).model_fields:
if not attr.startswith("_") and not callable(getattr(usage.completion_tokens_details, attr)):
if not attr.startswith("_") and not callable(
_attribute_value(usage.completion_tokens_details, attr)
):
current_val = getattr(combined.completion_tokens_details, attr, 0) or 0
new_val = getattr(usage.completion_tokens_details, attr, 0) or 0
if isinstance(new_val, (int, float)):

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

@ -115,9 +115,11 @@ class SpeechToCompletionBridgeHandler:
**request_data,
)
requested_response_format: Final = optional_params.get("response_format")
if isinstance(result, ModelResponse):
return self.transformation_handler.transform_response(
model_response=result,
response_format=requested_response_format if isinstance(requested_response_format, str) else None,
)
else:
raise Exception(f"Unmapped response type. Got type: {type(result)}")

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

@ -1,10 +1,14 @@
from collections.abc import Mapping
from types import MappingProxyType
from typing import TYPE_CHECKING, Final, cast
from typing_extensions import NotRequired, ReadOnly, TypedDict
from litellm.constants import OPENAI_CHAT_COMPLETION_PARAMS
if TYPE_CHECKING:
from litellm import Logging as LiteLLMLoggingObj
from litellm.types.llms.openai import HttpxBinaryResponseContent
from litellm.types.llms.openai import ChatCompletionUserMessage, HttpxBinaryResponseContent
from litellm.types.utils import ModelResponse
@ -16,7 +20,64 @@ def _completion_response_cost(model_response: "ModelResponse") -> float | None:
return response_cost if isinstance(response_cost, float) else None
GEMINI_TTS_CHAT_AUDIO_FORMAT: Final = "pcm16"
GEMINI_TTS_RAW_RESPONSE_FORMAT: Final = "pcm"
GEMINI_TTS_SUPPORTED_RESPONSE_FORMATS: Final = frozenset({"wav", GEMINI_TTS_RAW_RESPONSE_FORMAT})
class ChatAudioParam(TypedDict):
voice: ReadOnly[str]
format: ReadOnly[NotRequired[str]]
class SpeechToCompletionBridgeTransformationHandler:
def _validate_response_format(
self, model: str, custom_llm_provider: str, optional_params: Mapping[str, object]
) -> None:
if not self._is_gemini_tts_model(model):
return
response_format: Final = optional_params.get("response_format")
if not isinstance(response_format, str) or response_format in GEMINI_TTS_SUPPORTED_RESPONSE_FORMATS:
return
from litellm.exceptions import BadRequestError
supported: Final = ", ".join(sorted(GEMINI_TTS_SUPPORTED_RESPONSE_FORMATS))
raise BadRequestError(
message=(
f"Gemini TTS only produces raw PCM16 audio, so response_format='{response_format}'"
f" is not supported. Supported response formats: {supported}."
),
model=model,
llm_provider=custom_llm_provider,
)
def _chat_completion_params(self, optional_params: Mapping[str, object]) -> Mapping[str, object]:
return MappingProxyType(
{
param: value
for param, value in optional_params.items()
if param in OPENAI_CHAT_COMPLETION_PARAMS and param != "response_format"
}
)
def _chat_audio_format(self, model: str, optional_params: Mapping[str, object]) -> str | None:
if self._is_gemini_tts_model(model):
return GEMINI_TTS_CHAT_AUDIO_FORMAT
response_format: Final = optional_params.get("response_format")
return response_format if isinstance(response_format, str) else None
def _chat_audio_param(
self, model: str, voice: str | Mapping[str, object] | None, optional_params: Mapping[str, object]
) -> ChatAudioParam | None:
if not isinstance(voice, str):
return None
audio_format: Final = self._chat_audio_format(model, optional_params)
if audio_format is None:
voice_only: Final[ChatAudioParam] = {"voice": voice}
return voice_only
audio: Final[ChatAudioParam] = {"voice": voice, "format": audio_format}
return audio
def transform_request(
self,
model: str,
@ -28,36 +89,20 @@ class SpeechToCompletionBridgeTransformationHandler:
litellm_logging_obj: "LiteLLMLoggingObj",
custom_llm_provider: str,
) -> dict:
passed_optional_params: Final = {}
for op in optional_params:
if op in OPENAI_CHAT_COMPLETION_PARAMS:
passed_optional_params[op] = optional_params[op]
if voice is not None:
if isinstance(voice, str):
passed_optional_params["audio"] = {"voice": voice}
if "response_format" in optional_params:
passed_optional_params["audio"]["format"] = optional_params["response_format"]
return_kwargs = {
self._validate_response_format(model, custom_llm_provider, optional_params)
user_message: Final[ChatCompletionUserMessage] = {"role": "user", "content": input}
return_kwargs: Final = {
"model": model,
"messages": [
{
"role": "user",
"content": input,
}
],
"messages": [user_message],
"modalities": ["audio"],
**passed_optional_params,
**self._chat_completion_params(optional_params),
"audio": self._chat_audio_param(model, voice, optional_params),
**litellm_params,
"headers": headers,
"litellm_logging_obj": litellm_logging_obj,
"custom_llm_provider": custom_llm_provider,
}
# filter out None values
return_kwargs = {k: v for k, v in return_kwargs.items() if v is not None}
return return_kwargs
return {k: v for k, v in return_kwargs.items() if v is not None}
def _convert_pcm16_to_wav(self, pcm_data: bytes, sample_rate: int = 24000, channels: int = 1) -> bytes:
"""
@ -103,7 +148,14 @@ class SpeechToCompletionBridgeTransformationHandler:
"""Check if the model is a Gemini TTS model that returns PCM16 data."""
return "gemini" in model.lower() and ("tts" in model.lower() or "preview-tts" in model.lower())
def transform_response(self, model_response: "ModelResponse") -> "HttpxBinaryResponseContent":
def _gemini_tts_response_body(self, decoded_audio: bytes, response_format: str | None) -> tuple[bytes, str]:
if response_format == GEMINI_TTS_RAW_RESPONSE_FORMAT:
return decoded_audio, "audio/pcm"
return self._convert_pcm16_to_wav(decoded_audio), "audio/wav"
def transform_response(
self, model_response: "ModelResponse", response_format: str | None
) -> "HttpxBinaryResponseContent":
import base64
import httpx
@ -114,23 +166,17 @@ class SpeechToCompletionBridgeTransformationHandler:
audio_part: Final = cast(Choices, model_response.choices[0]).message.audio
if audio_part is None:
raise ValueError("No audio part found in the response")
audio_content: Final = audio_part.data
decoded_audio: Final = base64.b64decode(audio_part.data)
# Decode base64 to get binary content
binary_data = base64.b64decode(audio_content)
# Check if this is a Gemini TTS model that returns raw PCM16 data
model: Final = getattr(model_response, "model", "")
headers: Final = {}
if self._is_gemini_tts_model(model):
# Convert PCM16 to WAV format for proper audio file playback
binary_data = self._convert_pcm16_to_wav(binary_data)
headers["Content-Type"] = "audio/wav"
else:
headers["Content-Type"] = "audio/mpeg"
# Create an httpx.Response object
response: Final = httpx.Response(status_code=200, content=binary_data, headers=headers)
content, content_type = (
self._gemini_tts_response_body(decoded_audio, response_format)
if self._is_gemini_tts_model(model)
else (decoded_audio, "audio/mpeg")
)
response: Final = httpx.Response(
status_code=200, content=content, headers=MappingProxyType({"Content-Type": content_type})
)
binary_response: Final = HttpxBinaryResponseContent(response)
binary_response.set_response_cost(_completion_response_cost(model_response))
return binary_response

20
litellm/experimental_mcp_client/client.py
View file

@ -7,6 +7,7 @@ import base64
import os
from collections.abc import Awaitable, Callable, Generator
from datetime import timedelta
from functools import partial
from importlib import metadata
from typing import Any, Final, TypeVar
@ -47,7 +48,8 @@ from mcp.types import Tool as MCPTool
from pydantic import AnyUrl
from litellm._logging import verbose_logger
from litellm.constants import MCP_CLIENT_TIMEOUT, MCP_NPM_CACHE_DIR
from litellm.constants import MCP_CLIENT_TIMEOUT, MCP_NPM_CACHE_DIR, MCP_TOOL_LISTING_TIMEOUT
from litellm.experimental_mcp_client.tools import list_tools_with_pagination
from litellm.llms.custom_httpx.http_handler import get_ssl_configuration
from litellm.types.llms.custom_http import VerifyTypes
from litellm.types.mcp import (
@ -603,17 +605,19 @@ class MCPClient:
"""
verbose_logger.debug("MCP client listing tools from %s", self.server_url or "stdio")
async def _list_tools_operation(session: ClientSession):
return await session.list_tools()
try:
result: Final = await self.run_with_session(_list_tools_operation, quiet_on_error=raise_on_error)
tool_count: Final = len(result.tools)
tool_names: Final = [tool.name for tool in result.tools]
# A per-server timeout above the global default extends the whole-walk deadline
listing_deadline: Final = max(self.timeout, MCP_TOOL_LISTING_TIMEOUT)
tools: Final = await self.run_with_session(
partial(list_tools_with_pagination, listing_deadline=listing_deadline),
quiet_on_error=raise_on_error,
)
tool_count: Final = len(tools)
tool_names: Final = tuple(tool.name for tool in tools)
verbose_logger.info(
"MCP client listed %s tools from %s: %s", tool_count, self.server_url or "stdio", tool_names
)
return result.tools
return tools
except asyncio.CancelledError:
verbose_logger.warning("MCP client list_tools was cancelled")
raise

74
litellm/experimental_mcp_client/tools.py
View file

@ -1,14 +1,22 @@
import json
from typing import Final, Literal
import anyio
from mcp import ClientSession
from mcp.types import CallToolRequestParams as MCPCallToolRequestParams
from mcp.types import CallToolResult as MCPCallToolResult
from mcp.types import PaginatedRequestParams
from mcp.types import Tool as MCPTool
from openai.types.chat import ChatCompletionToolParam
from openai.types.responses.function_tool_param import FunctionToolParam
from openai.types.shared_params.function_definition import FunctionDefinition
from litellm._logging import verbose_logger
from litellm.constants import (
MCP_CLIENT_TIMEOUT,
MCP_TOOL_LISTING_MAX_PAGES,
MCP_TOOL_LISTING_TIMEOUT,
)
from litellm.types.llms.anthropic import AnthropicMessagesTool
from litellm.types.utils import ChatCompletionMessageToolCall
@ -90,6 +98,64 @@ def transform_mcp_tool_to_anthropic_tool(mcp_tool: MCPTool) -> AnthropicMessages
)
async def list_tools_with_pagination(
session: ClientSession, listing_deadline: float | None = None
) -> list[MCPTool]: # mutable-ok: list return contract
"""Collect tools from every tools/list page by following nextCursor.
Stops and returns the tools collected so far when the upstream repeats a
cursor, the page cap is reached, or the whole-walk deadline expires, so a
buggy or slow upstream yields a partial catalog instead of an error.
listing_deadline overrides the default whole-walk deadline; callers with a
per-server timeout above the global default pass it through here.
"""
tools: Final[list[MCPTool]] = [] # mutable-ok: accumulates each page's tools
seen_cursors: Final[set[str]] = set() # mutable-ok: guards against cursor loops
cursor: str | None = None # rebind-ok: advances to each page's nextCursor
# The per-request session read timeout restarts on every page, so a multi-page
# walk needs its own overall deadline. max() keeps the pre-pagination guarantee
# that a single page slower than the listing timeout but within the client
# timeout still succeeds.
effective_deadline: Final = (
listing_deadline if listing_deadline is not None else max(MCP_CLIENT_TIMEOUT, MCP_TOOL_LISTING_TIMEOUT)
)
with anyio.move_on_after(effective_deadline):
for _ in range(MCP_TOOL_LISTING_MAX_PAGES):
result = (
await session.list_tools()
if cursor is None
else await session.list_tools(params=PaginatedRequestParams(cursor=cursor))
)
tools.extend(result.tools)
next_cursor = getattr(result, "nextCursor", None)
if not isinstance(next_cursor, str) or not next_cursor:
return tools
if next_cursor in seen_cursors:
verbose_logger.warning(
"MCP server repeated a tools/list cursor while listing tools; returning %s tools collected so far",
len(tools),
)
return tools
seen_cursors.add(next_cursor)
cursor = next_cursor
verbose_logger.warning(
"MCP server tools/list pagination exceeded the maximum of %s pages; returning %s tools collected so far",
MCP_TOOL_LISTING_MAX_PAGES,
len(tools),
)
return tools
verbose_logger.warning(
"MCP server tools/list pagination exceeded the %s second listing deadline; returning %s tools collected so far",
effective_deadline,
len(tools),
)
return tools
async def load_mcp_tools(
session: ClientSession, format: Literal["mcp", "openai"] = "mcp"
) -> list[MCPTool] | list[ChatCompletionToolParam]:
@ -103,10 +169,12 @@ async def load_mcp_tools(
If format is set to "openai", the tools are converted to OpenAI API compatible tools.
"""
tools: Final = await session.list_tools()
tools: Final = await list_tools_with_pagination(session)
if format == "openai":
return [transform_mcp_tool_to_openai_tool(mcp_tool=tool) for tool in tools.tools]
return tools.tools
return [ # mutable-ok: public API returns a list
transform_mcp_tool_to_openai_tool(mcp_tool=tool) for tool in tools
]
return tools
########################################################

168
litellm/google_genai/adapters/transformation.py
View file

@ -1,8 +1,9 @@
import json
from collections.abc import AsyncIterator, Iterator, Sequence
from typing import Any, Final, TypedDict, cast
from collections.abc import AsyncIterator, Callable, Iterator, Mapping, Sequence
from types import MappingProxyType
from typing import Any, Final, TypeAlias, cast
from typing_extensions import ReadOnly
from typing_extensions import ReadOnly, TypedDict
from litellm import verbose_logger
from litellm.litellm_core_utils.json_validation_rule import normalize_tool_schema
@ -11,7 +12,6 @@ from litellm.types.llms.openai import (
ChatCompletionAssistantMessage,
ChatCompletionAssistantToolCall,
ChatCompletionImageObject,
ChatCompletionRequest,
ChatCompletionSystemMessage,
ChatCompletionTextObject,
ChatCompletionToolCallFunctionChunk,
@ -23,35 +23,63 @@ from litellm.types.llms.openai import (
from litellm.types.router import GenericLiteLLMParams
from litellm.types.utils import (
AdapterCompletionStreamWrapper,
ChatCompletionDeltaCustomToolCall,
ChatCompletionDeltaToolCall,
ChatCompletionMessageCustomToolCall,
ChatCompletionMessageToolCall,
Choices,
Delta,
Function,
Message,
ModelResponse,
ModelResponseStream,
StreamingChoices,
Usage,
)
class _GenAITextPart(TypedDict, total=False):
text: ReadOnly[str]
_JsonDict: TypeAlias = dict[str, object]
_JsonDictList: TypeAlias = list[_JsonDict]
class _GenAISystemInstruction(TypedDict, total=False):
parts: ReadOnly[list[_GenAITextPart]]
class _ToolCallAccumulator(TypedDict):
name: ReadOnly[str]
arguments: ReadOnly[str]
class _GenAIFunctionCall(TypedDict):
name: ReadOnly[str]
args: ReadOnly[Mapping[str, object]]
class _GenAIPart(TypedDict, total=False):
text: ReadOnly[str]
functionCall: ReadOnly[dict[str, object]]
functionCall: ReadOnly[_GenAIFunctionCall]
class _GenAIFunctionResponse(TypedDict, total=False):
name: ReadOnly[str]
response: ReadOnly[object]
class _GenAIRequestFunctionCall(TypedDict, total=False):
name: ReadOnly[str]
args: ReadOnly[Mapping[str, object]]
class _GenAIContentPart(TypedDict, total=False):
text: ReadOnly[str]
inline_data: ReadOnly[Mapping[str, str]]
functionResponse: ReadOnly[_GenAIFunctionResponse]
functionCall: ReadOnly[_GenAIRequestFunctionCall]
class _GenAIFunctionDeclaration(TypedDict, total=False):
name: ReadOnly[str]
description: ReadOnly[str]
parametersJsonSchema: ReadOnly[dict[str, object]]
parametersJsonSchema: ReadOnly[object]
class _GenAITool(TypedDict, total=False):
functionDeclarations: ReadOnly[list[_GenAIFunctionDeclaration]]
functionDeclarations: ReadOnly[Sequence[_GenAIFunctionDeclaration]]
class _GenAIFunctionCallingConfig(TypedDict, total=False):
@ -62,9 +90,11 @@ class _GenAIToolConfig(TypedDict, total=False):
functionCallingConfig: ReadOnly[_GenAIFunctionCallingConfig]
def _decode_tool_call_arguments(raw_arguments: str) -> object:
"""Decode a tool call's JSON-encoded arguments into the value Google GenAI expects."""
return json.loads(raw_arguments)
class _GenAISystemInstruction(TypedDict, total=False):
parts: ReadOnly[Sequence[Mapping[str, str]]]
_EMPTY_STR_MAPPING: Final[Mapping[str, str]] = MappingProxyType({})
class GoogleGenAIStreamWrapper(AdapterCompletionStreamWrapper):
@ -74,12 +104,11 @@ class GoogleGenAIStreamWrapper(AdapterCompletionStreamWrapper):
"""
sent_first_chunk: bool = False
# State tracking for accumulating partial tool calls
accumulated_tool_calls: dict[int, dict[str, str]]
_parse_accumulated_args: Callable[[str], Mapping[str, object]] = staticmethod(json.loads)
def __init__(self, completion_stream: object):
self.sent_first_chunk = False
self.accumulated_tool_calls = {}
self.accumulated_tool_calls = dict[int, _ToolCallAccumulator]()
self._returned_response = False
super().__init__(completion_stream)
@ -124,7 +153,7 @@ class GoogleGenAIStreamWrapper(AdapterCompletionStreamWrapper):
# After the stream is exhausted, check for any remaining accumulated tool calls
if self.accumulated_tool_calls:
try:
parts: Final[list[_GenAIPart]] = []
parts: Final = list[_GenAIPart]()
for (
tool_call_index,
tool_call_data,
@ -132,7 +161,9 @@ class GoogleGenAIStreamWrapper(AdapterCompletionStreamWrapper):
try:
# For tool calls with no arguments, accumulated_args will be "", which is not valid JSON.
# We default to an empty JSON object in this case.
parsed_args = _decode_tool_call_arguments(tool_call_data["arguments"] or "{}")
parsed_args: Mapping[str, object] = self._parse_accumulated_args(
tool_call_data["arguments"] or "{}"
)
function_call_part: _GenAIPart = {
"functionCall": {
"name": tool_call_data["name"] or "undefined_tool_name",
@ -149,7 +180,7 @@ class GoogleGenAIStreamWrapper(AdapterCompletionStreamWrapper):
tool_call_data["arguments"],
)
if parts:
final_chunk: Final[dict[str, object]] = {
final_chunk: Final = {
"candidates": [
{
"content": {"parts": parts, "role": "model"},
@ -211,14 +242,16 @@ class GoogleGenAIStreamWrapper(AdapterCompletionStreamWrapper):
class GoogleGenAIAdapter:
"""Adapter for transforming Google GenAI generate_content requests to/from litellm.completion format"""
_parse_tool_call_args: Callable[[str], Mapping[str, object]] = staticmethod(json.loads)
def __init__(self) -> None:
pass
def translate_generate_content_to_completion(
self,
model: str,
contents: list[dict[str, Any]] | dict[str, Any],
config: dict[str, Any] | None = None,
contents: _JsonDictList | _JsonDict,
config: Mapping[str, object] | None = None,
litellm_params: GenericLiteLLMParams | None = None,
**kwargs,
) -> dict[str, Any]:
@ -250,7 +283,7 @@ class GoogleGenAIAdapter:
messages: Final = self._transform_contents_to_messages(contents_list, system_instruction=system_instruction)
# Create base request as dict (which is compatible with ChatCompletionRequest)
completion_request: Final[ChatCompletionRequest] = {
completion_request: Final[_JsonDict] = {
"model": model,
"messages": messages,
}
@ -312,9 +345,9 @@ class GoogleGenAIAdapter:
def _add_generic_litellm_params_to_request(
self,
completion_request_dict: dict[str, object],
completion_request_dict: _JsonDict,
litellm_params: GenericLiteLLMParams | None = None,
) -> dict[str, object]:
) -> _JsonDict:
"""Add generic litellm params to request. e.g add api_base, api_key, api_version, etc.
Args:
@ -326,7 +359,7 @@ class GoogleGenAIAdapter:
"""
allowed_fields: Final = GenericLiteLLMParams.model_fields.keys()
if litellm_params:
litellm_dict: Final = litellm_params.model_dump(exclude_none=True)
litellm_dict: Final[_JsonDict] = litellm_params.model_dump(exclude_none=True)
for key, value in litellm_dict.items():
if key in allowed_fields:
completion_request_dict[key] = value
@ -346,12 +379,12 @@ class GoogleGenAIAdapter:
tools: Sequence[_GenAITool],
) -> list[ChatCompletionToolParam]:
"""Transform Google GenAI tools to OpenAI tools format"""
openai_tools: Final[list[dict[str, object]]] = []
openai_tools: Final = list[_JsonDict]()
for tool in tools:
if "functionDeclarations" in tool:
for func_decl in tool["functionDeclarations"]:
function_chunk: dict[str, object] = {
function_chunk: _JsonDict = {
"name": func_decl.get("name", ""),
}
@ -360,7 +393,7 @@ class GoogleGenAIAdapter:
if "parametersJsonSchema" in func_decl:
function_chunk["parameters"] = func_decl["parametersJsonSchema"]
openai_tool: dict[str, object] = {"type": "function", "function": function_chunk}
openai_tool: _JsonDict = {"type": "function", "function": function_chunk}
openai_tools.append(openai_tool)
# normalize the tool schemas
@ -391,13 +424,13 @@ class GoogleGenAIAdapter:
# Handle system instruction
if system_instruction:
system_parts: Final = system_instruction.get("parts", [])
system_parts: Final[Sequence[Mapping[str, str]]] = system_instruction.get("parts", [])
if system_parts and "text" in system_parts[0]:
messages.append(ChatCompletionSystemMessage(role="system", content=system_parts[0]["text"]))
for content in contents:
role = content.get("role", "user")
parts = content.get("parts", [])
parts: Sequence[_GenAIContentPart | str | None] = content.get("parts", [])
if role == "user":
# Handle user messages with potential function responses
@ -500,7 +533,7 @@ class GoogleGenAIAdapter:
def translate_completion_to_generate_content(
self,
response: ModelResponse,
) -> dict[str, object]:
) -> _JsonDict:
"""
Transform litellm completion response to Google GenAI generate_content format
@ -523,13 +556,13 @@ class GoogleGenAIAdapter:
parts = self._transform_openai_message_to_google_genai_parts(choice.message)
else:
# Fallback for generic choice objects
message_content = getattr(choice, "message", {}).get("content", "") or getattr(choice, "delta", {}).get(
"content", ""
)
message_content: str = getattr(choice, "message", _EMPTY_STR_MAPPING).get("content", "") or getattr(
choice, "delta", _EMPTY_STR_MAPPING
).get("content", "")
parts = [{"text": message_content}] if message_content else []
# Create Google GenAI format response
generate_content_response: Final[dict[str, object]] = {
generate_content_response: Final[_JsonDict] = {
"candidates": [
{
"content": {"parts": parts, "role": "model"},
@ -563,7 +596,7 @@ class GoogleGenAIAdapter:
self,
response: ModelResponse | ModelResponseStream,
wrapper: GoogleGenAIStreamWrapper,
) -> dict[str, object] | None:
) -> Mapping[str, object] | None:
"""
Transform streaming litellm completion chunk to Google GenAI generate_content format
@ -590,7 +623,7 @@ class GoogleGenAIAdapter:
finish_reason: str | None = getattr(choice, "finish_reason", None)
else:
# Fallback for generic choice objects
message_content: Final = getattr(choice, "delta", {}).get("content", "")
message_content: Final[str] = getattr(choice, "delta", _EMPTY_STR_MAPPING).get("content", "")
parts = [{"text": message_content}] if message_content else []
finish_reason = getattr(choice, "finish_reason", None)
@ -599,7 +632,7 @@ class GoogleGenAIAdapter:
return None
# Create Google GenAI streaming format response
streaming_chunk: Final[dict[str, object]] = {
streaming_chunk: Final[_JsonDict] = {
"candidates": [
{
"content": {"parts": parts, "role": "model"},
@ -635,10 +668,10 @@ class GoogleGenAIAdapter:
def _transform_openai_message_to_google_genai_parts(
self,
message: Any,
) -> list[_GenAIPart]:
message: Message,
) -> Sequence[_GenAIPart]:
"""Transform OpenAI message to Google GenAI parts format"""
parts: Final[list[_GenAIPart]] = []
parts: Final = list[_GenAIPart]()
# Add text content if present
if hasattr(message, "content") and message.content:
@ -646,20 +679,22 @@ class GoogleGenAIAdapter:
# Add tool calls if present
if hasattr(message, "tool_calls") and message.tool_calls:
for tool_call in message.tool_calls:
if hasattr(tool_call, "function") and tool_call.function:
tool_calls: Final[Sequence[ChatCompletionMessageToolCall | ChatCompletionMessageCustomToolCall]] = (
message.tool_calls
)
for tool_call in tool_calls:
function: Function | None = getattr(tool_call, "function", None)
if function:
try:
args = (
_decode_tool_call_arguments(tool_call.function.arguments)
if tool_call.function.arguments
else {}
args: Mapping[str, object] = (
self._parse_tool_call_args(function.arguments) if function.arguments else {}
)
except json.JSONDecodeError:
args = {}
function_call_part: _GenAIPart = {
"functionCall": {
"name": tool_call.function.name or "undefined_tool_name",
"name": function.name or "undefined_tool_name",
"args": args,
}
}
@ -668,24 +703,26 @@ class GoogleGenAIAdapter:
return parts if parts else [{"text": ""}]
def _transform_openai_delta_to_google_genai_parts_with_accumulation(
self, delta: Any, wrapper: GoogleGenAIStreamWrapper
) -> list[_GenAIPart]:
self, delta: Delta, wrapper: GoogleGenAIStreamWrapper
) -> Sequence[_GenAIPart]:
"""Transforms OpenAI delta to Google GenAI parts, accumulating streaming tool calls."""
# 1. Initialize wrapper state if it doesn't exist
if not hasattr(wrapper, "accumulated_tool_calls"):
wrapper.accumulated_tool_calls = {}
parts: Final[list[_GenAIPart]] = []
parts: Final = list[_GenAIPart]()
if hasattr(delta, "content") and delta.content:
parts.append({"text": delta.content})
# 2. Ensure tool_calls is iterable
tool_calls: Final = delta.tool_calls or []
tool_calls: Final[Sequence[ChatCompletionDeltaToolCall | ChatCompletionDeltaCustomToolCall]] = (
delta.tool_calls or []
)
for tool_call in tool_calls:
if not hasattr(tool_call, "function"):
if not hasattr(tool_call, "function") or isinstance(tool_call, ChatCompletionDeltaCustomToolCall):
continue
# 3. Use `index` as the primary key for accumulation
@ -701,19 +738,20 @@ class GoogleGenAIAdapter:
}
# Accumulate name and arguments
function_name = getattr(tool_call.function, "name", None)
args_chunk = getattr(tool_call.function, "arguments", None)
delta_function: Function | None = getattr(tool_call, "function", None)
function_name: str | None = getattr(delta_function, "name", None)
args_chunk: str | None = getattr(delta_function, "arguments", None)
# Optimization: Skip chunks that have no new data
if not function_name and not args_chunk:
verbose_logger.debug("Skipping empty tool call chunk for index: %s", tool_call_index)
continue
if function_name:
wrapper.accumulated_tool_calls[tool_call_index]["name"] = function_name
if args_chunk:
wrapper.accumulated_tool_calls[tool_call_index]["arguments"] += args_chunk
previous_data: _ToolCallAccumulator = wrapper.accumulated_tool_calls[tool_call_index]
wrapper.accumulated_tool_calls[tool_call_index] = _ToolCallAccumulator(
name=function_name or previous_data["name"],
arguments=previous_data["arguments"] + (args_chunk or ""),
)
# Attempt to parse and emit a complete tool call
accumulated_data = wrapper.accumulated_tool_calls[tool_call_index]
@ -723,7 +761,7 @@ class GoogleGenAIAdapter:
# 5. Attempt to parse arguments even if name hasn't arrived.
try:
# Attempt to parse the accumulated arguments string
parsed_args = _decode_tool_call_arguments(accumulated_args)
parsed_args: Mapping[str, object] = self._parse_tool_call_args(accumulated_args)
# If parsing succeeds, but we don't have a name yet, wait.
# The part will be created by a later chunk that brings the name.
@ -757,7 +795,7 @@ class GoogleGenAIAdapter:
return mapping.get(finish_reason, "STOP")
def _map_usage(self, usage: Usage | None) -> dict[str, int]:
def _map_usage(self, usage: object) -> Mapping[str, int]:
"""Map OpenAI usage to Google GenAI usage format"""
return {
"promptTokenCount": getattr(usage, "prompt_tokens", 0) or 0,

30
litellm/google_genai/main.py
View file

@ -52,10 +52,10 @@ class GenerateContentSetupResult(BaseModel):
model_config: ClassVar[ConfigDict] = ConfigDict(arbitrary_types_allowed=True)
model: str
request_body: dict[str, Any]
request_body: dict[str, object]
custom_llm_provider: str
generate_content_provider_config: BaseGoogleGenAIGenerateContentConfig | None
generate_content_config_dict: dict[str, Any]
generate_content_config_dict: dict[str, object]
native_request_fields: dict[str, object]
litellm_params: GenericLiteLLMParams
litellm_logging_obj: LiteLLMLoggingObj
@ -68,7 +68,7 @@ class GenerateContentHelper:
@staticmethod
def mock_generate_content_response(
mock_response: str = "This is a mock response from Google GenAI generate_content.",
) -> dict[str, Any]:
) -> dict[str, object]:
"""Mock response for generate_content for testing purposes"""
return {
"text": mock_response,
@ -239,9 +239,9 @@ async def agenerate_content(
tools: ToolConfigDict | None = None,
# Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs.
# The extra values given here take precedence over values defined on the client or passed to this method.
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,
timeout: float | httpx.Timeout | None = None,
# LiteLLM specific params,
custom_llm_provider: str | None = None,
@ -307,9 +307,9 @@ def generate_content(
tools: ToolConfigDict | None = None,
# Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs.
# The extra values given here take precedence over values defined on the client or passed to this method.
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,
timeout: float | httpx.Timeout | None = None,
# LiteLLM specific params,
custom_llm_provider: str | None = None,
@ -397,9 +397,9 @@ async def agenerate_content_stream(
tools: ToolConfigDict | None = None,
# Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs.
# The extra values given here take precedence over values defined on the client or passed to this method.
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,
timeout: float | httpx.Timeout | None = None,
# LiteLLM specific params,
custom_llm_provider: str | None = None,
@ -492,9 +492,9 @@ def generate_content_stream(
tools: ToolConfigDict | None = None,
# Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs.
# The extra values given here take precedence over values defined on the client or passed to this method.
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,
timeout: float | httpx.Timeout | None = None,
# LiteLLM specific params,
custom_llm_provider: str | None = None,

26
litellm/images/main.py
View file

@ -3,7 +3,7 @@ import contextvars
import importlib
from collections.abc import Coroutine
from functools import partial
from typing import TYPE_CHECKING, Any, Final, Literal, Optional, cast, overload
from typing import TYPE_CHECKING, Final, Literal, Optional, cast, overload
if TYPE_CHECKING:
from litellm.images.utils import ImageEditRequestUtils
@ -151,7 +151,7 @@ def image_generation(
*,
aimg_generation: Literal[True],
**kwargs,
) -> Coroutine[Any, Any, ImageResponse]:
) -> Coroutine[object, object, ImageResponse]:
...
@ -197,7 +197,7 @@ def image_generation(
api_version: str | None = None,
custom_llm_provider=None,
**kwargs,
) -> ImageResponse | Coroutine[Any, Any, ImageResponse]:
) -> ImageResponse | Coroutine[object, object, ImageResponse]:
"""
Maps the https://api.openai.com/v1/images/generations endpoint.
@ -386,6 +386,8 @@ def image_generation(
litellm.LlmProviders.VERTEX_AI,
litellm.LlmProviders.OPENROUTER,
litellm.LlmProviders.DASHSCOPE,
litellm.LlmProviders.QWENCLOUD,
litellm.LlmProviders.QWEN_AI_PLATFORM,
):
if image_generation_config is None:
raise ValueError(f"image generation config is not supported for {custom_llm_provider}")
@ -723,14 +725,14 @@ def image_edit(
user: str | None = None,
# Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs.
# The extra values given here take precedence over values defined on the client or passed to this method.
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,
timeout: float | httpx.Timeout | None = None,
# LiteLLM specific params,
custom_llm_provider: str | None = None,
**kwargs,
) -> ImageResponse | Coroutine[Any, Any, ImageResponse]:
) -> ImageResponse | Coroutine[object, object, ImageResponse]:
"""
Maps the image edit functionality, similar to OpenAI's images/edits endpoint.
"""
@ -769,7 +771,7 @@ def image_edit(
images: Final = image if isinstance(image, list) else ([image] if image is not None else [])
headers_from_kwargs: Final = kwargs.get("headers")
merged_extra_headers: Final[dict[str, Any]] = {}
merged_extra_headers: Final[dict[str, object]] = {}
if isinstance(headers_from_kwargs, dict):
merged_extra_headers.update(headers_from_kwargs)
if isinstance(extra_headers, dict):
@ -974,9 +976,9 @@ async def aimage_edit(
user: str | None = None,
# Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs.
# The extra values given here take precedence over values defined on the client or passed to this method.
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,
timeout: float | httpx.Timeout | None = None,
# LiteLLM specific params,
custom_llm_provider: str | None = None,
@ -1044,7 +1046,7 @@ async def aimage_edit(
)
def __getattr__(name: str) -> Any:
def __getattr__(name: str) -> type["ImageEditRequestUtils"]:
"""Lazy import handler for images.main module"""
if name == "ImageEditRequestUtils":
# Lazy load ImageEditRequestUtils to avoid heavy import from images.utils at module load time

96
litellm/integrations/SlackAlerting/slack_alerting.py
View file

@ -68,6 +68,7 @@ from .utils import process_slack_alerting_variables
if TYPE_CHECKING:
from litellm.proxy.db.db_transaction_queue.pod_lock_manager import PodLockManager
from litellm.proxy.utils import PrismaClient
from litellm.router import Router as _Router
Router = _Router
@ -545,7 +546,6 @@ class SlackAlerting(CustomBatchLogger):
# Get the appropriate budget alert type handler
budget_alert_class: Final = get_budget_alert_type(type)
_id: Final = budget_alert_class.get_id(user_info)
user_info_json: Final = user_info.model_dump(exclude_none=True)
user_info_str: Final = self._get_user_info_str(user_info)
event_message = budget_alert_class.get_event_message()
@ -575,7 +575,22 @@ class SlackAlerting(CustomBatchLogger):
webhook_event = WebhookEvent(
event=event,
event_message=event_message,
**user_info_json,
spend=user_info.spend,
max_budget=user_info.max_budget,
soft_budget=user_info.soft_budget,
token=user_info.token,
customer_id=user_info.customer_id,
user_id=user_info.user_id,
team_id=user_info.team_id,
team_alias=user_info.team_alias,
organization_id=user_info.organization_id,
user_email=user_info.user_email,
key_alias=user_info.key_alias,
projected_exceeded_date=user_info.projected_exceeded_date,
projected_spend=user_info.projected_spend,
event_group=user_info.event_group,
alert_emails=user_info.alert_emails,
max_budget_alert_emails=user_info.max_budget_alert_emails,
)
await self.send_alert(
message=event_message + "\n\n" + user_info_str,
@ -657,7 +672,7 @@ class SlackAlerting(CustomBatchLogger):
"""
Create a standard message for a budget alert
"""
_all_fields_as_dict: Final = user_info.model_dump(exclude_none=True)
_all_fields_as_dict: Final[dict[str, object]] = user_info.model_dump(exclude_none=True)
_all_fields_as_dict.pop("token")
msg = ""
for k, v in _all_fields_as_dict.items():
@ -1006,7 +1021,7 @@ class SlackAlerting(CustomBatchLogger):
except Exception:
pass
async def model_added_alert(self, model_name: str, litellm_model_name: str, passed_model_info: Any):
async def model_added_alert(self, model_name: str, litellm_model_name: str, passed_model_info: object):
base_model_from_user: Final = getattr(passed_model_info, "base_model", None)
model_info = {}
base_model = ""
@ -1485,9 +1500,9 @@ Model Info:
elif self.default_webhook_url is not None:
_digest_webhook = self.default_webhook_url
else:
_digest_webhook = os.getenv("SLACK_WEBHOOK_URL", None)
_digest_webhook = os.getenv("SLACK_WEBHOOK_URL") or os.getenv("ALERTING_WEBHOOK_URL")
if _digest_webhook is None:
raise ValueError("Missing SLACK_WEBHOOK_URL from environment")
raise ValueError("Missing SLACK_WEBHOOK_URL / ALERTING_WEBHOOK_URL from environment")
digest_key: Final = f"{alert_type_name_str}:{request_model or ''}:{api_base or ''}"
@ -1516,10 +1531,10 @@ Model Info:
elif self.default_webhook_url is not None:
slack_webhook_url = self.default_webhook_url
else:
slack_webhook_url = os.getenv("SLACK_WEBHOOK_URL", None)
slack_webhook_url = os.getenv("SLACK_WEBHOOK_URL") or os.getenv("ALERTING_WEBHOOK_URL")
if slack_webhook_url is None:
raise ValueError("Missing SLACK_WEBHOOK_URL from environment")
raise ValueError("Missing SLACK_WEBHOOK_URL / ALERTING_WEBHOOK_URL from environment")
payload: Final = {"text": formatted_message}
headers: Final = {"Content-type": "application/json"}
@ -1930,6 +1945,69 @@ Model Info:
except Exception as e:
verbose_proxy_logger.exception("Error sending weekly spend report %s", e)
async def send_user_spend_alerts(self, prisma_client: "PrismaClient | None" = None) -> None:
"""Check per-user daily/monthly spend thresholds and spend anomalies, alerting once per user per period."""
if self.alerting is None or "slack" not in self.alerting:
return
thresholds_enabled: Final = AlertType.user_spend_thresholds in self.alert_types
anomalies_enabled: Final = AlertType.user_spend_anomalies in self.alert_types
if not thresholds_enabled and not anomalies_enabled:
return
if prisma_client is None:
from litellm.proxy.proxy_server import prisma_client as global_prisma_client
prisma_client = global_prisma_client # rebind-ok: fall back to the proxy's global client
if prisma_client is None:
return
from litellm.integrations.SlackAlerting.user_spend_alerts import (
evaluate_user_spend,
fetch_user_spend_rows,
)
try:
today: Final = datetime.datetime.now(datetime.timezone.utc).date()
rows: Final = await fetch_user_spend_rows(
prisma_client=prisma_client,
today=today,
baseline_days=self.alerting_args.spend_anomaly_baseline_days,
)
all_events: Final = tuple(
event
for row in rows
for event in evaluate_user_spend(
row=row,
args=self.alerting_args,
today=today,
thresholds_enabled=thresholds_enabled,
anomalies_enabled=anomalies_enabled,
)
)
cached_flags: Final = await asyncio.gather(
*(self.internal_usage_cache.async_get_cache(key=event.cache_key) for event in all_events)
)
new_events: Final = tuple(event for event, cached in zip(all_events, cached_flags) if not cached)
for alert_type in (AlertType.user_spend_thresholds, AlertType.user_spend_anomalies):
typed_events = tuple(event for event in new_events if event.alert_type == alert_type)
if not typed_events:
continue
await self.send_alert(
message="\n\n".join(event.message for event in typed_events),
level="High",
alert_type=alert_type,
alerting_metadata={}, # mutable-ok: send_alert takes a dict payload
)
for event in typed_events:
await self.internal_usage_cache.async_set_cache(
key=event.cache_key,
value="SENT",
ttl=event.cache_ttl,
)
except Exception as e: # noqa: BLE001 # background job must not crash the scheduler
verbose_proxy_logger.exception("Error sending user spend alerts: %s", e)
async def send_fallback_stats_from_prometheus(self):
"""
Helper to send fallback statistics from prometheus server -> to slack
@ -1973,7 +2051,7 @@ Model Info:
try:
message = f"`{event_name}`\n"
key_event_dict: Final = key_event.model_dump()
key_event_dict: Final[dict[str, object]] = key_event.model_dump()
# Add Created by information first
message += "*Action Done by:*\n"

139
litellm/integrations/SlackAlerting/user_spend_alerts.py Normal file
View file

@ -0,0 +1,139 @@
"""Per-user daily/monthly spend threshold alerts and spend anomaly detection."""
import datetime
from dataclasses import dataclass
from typing import TYPE_CHECKING, Final, Literal
from pydantic import TypeAdapter
from litellm.constants import HOURS_IN_A_DAY
from litellm.types.integrations.slack_alerting import AlertType, SlackAlertingArgs
if TYPE_CHECKING:
from litellm.proxy.utils import PrismaClient
DAY_SECONDS: Final = HOURS_IN_A_DAY * 60 * 60
MONTHLY_ALERT_TTL_SECONDS: Final = 32 * DAY_SECONDS
USER_SPEND_QUERY: Final = """
SELECT
user_id,
COALESCE(SUM(spend) FILTER (WHERE date = $1), 0)::float AS daily_spend,
COALESCE(SUM(spend) FILTER (WHERE date >= $2), 0)::float AS monthly_spend,
COALESCE(SUM(spend) FILTER (WHERE date >= $3 AND date < $1), 0)::float AS baseline_spend
FROM "LiteLLM_DailyUserSpend"
WHERE date >= LEAST($2, $3) AND user_id IS NOT NULL
GROUP BY user_id
HAVING COALESCE(SUM(spend) FILTER (WHERE date >= $2), 0) > 0
"""
@dataclass(frozen=True, slots=True)
class UserSpendRow:
user_id: str
daily_spend: float
monthly_spend: float
baseline_spend: float
@dataclass(frozen=True, slots=True)
class UserSpendAlertEvent:
kind: Literal["daily_threshold", "monthly_threshold", "anomaly"]
alert_type: AlertType
message: str
cache_key: str
cache_ttl: int
USER_SPEND_ROWS_ADAPTER: Final = TypeAdapter(tuple[UserSpendRow, ...])
async def fetch_user_spend_rows(
prisma_client: "PrismaClient",
today: datetime.date,
baseline_days: int,
) -> tuple[UserSpendRow, ...]:
today_str: Final = today.strftime("%Y-%m-%d")
month_start_str: Final = today.replace(day=1).strftime("%Y-%m-%d")
baseline_start_str: Final = (today - datetime.timedelta(days=max(baseline_days, 1))).strftime("%Y-%m-%d")
raw: Final = await prisma_client.db.query_raw(USER_SPEND_QUERY, today_str, month_start_str, baseline_start_str)
return USER_SPEND_ROWS_ADAPTER.validate_python(raw)
def _daily_threshold_event(row: UserSpendRow, args: SlackAlertingArgs, today_str: str) -> UserSpendAlertEvent | None:
threshold: Final = args.daily_spend_per_user_threshold
if threshold is None or row.daily_spend < threshold:
return None
return UserSpendAlertEvent(
kind="daily_threshold",
alert_type=AlertType.user_spend_thresholds,
message=(
f"User Daily Spend Threshold Crossed:\n"
f"User: `{row.user_id}`\n"
f"Spend Today: `${row.daily_spend:.2f}`\n"
f"Daily Threshold: `${threshold:.2f}`"
),
cache_key=f"user_spend_alert_daily_{row.user_id}_{today_str}",
cache_ttl=DAY_SECONDS,
)
def _monthly_threshold_event(row: UserSpendRow, args: SlackAlertingArgs, month_str: str) -> UserSpendAlertEvent | None:
threshold: Final = args.monthly_spend_per_user_threshold
if threshold is None or row.monthly_spend < threshold:
return None
return UserSpendAlertEvent(
kind="monthly_threshold",
alert_type=AlertType.user_spend_thresholds,
message=(
f"User Monthly Spend Threshold Crossed:\n"
f"User: `{row.user_id}`\n"
f"Spend This Month: `${row.monthly_spend:.2f}`\n"
f"Monthly Threshold: `${threshold:.2f}`"
),
cache_key=f"user_spend_alert_monthly_{row.user_id}_{month_str}",
cache_ttl=MONTHLY_ALERT_TTL_SECONDS,
)
def _anomaly_event(row: UserSpendRow, args: SlackAlertingArgs, today_str: str) -> UserSpendAlertEvent | None:
if row.daily_spend < args.spend_anomaly_min_spend:
return None
baseline_daily_avg: Final = row.baseline_spend / args.spend_anomaly_baseline_days
if row.baseline_spend > 0 and row.daily_spend <= args.spend_anomaly_multiplier * baseline_daily_avg:
return None
return UserSpendAlertEvent(
kind="anomaly",
alert_type=AlertType.user_spend_anomalies,
message=(
f"User Spend Anomaly Detected:\n"
f"User: `{row.user_id}`\n"
f"Spend Today: `${row.daily_spend:.2f}`\n"
f"Daily Average (last {args.spend_anomaly_baseline_days} days): `${baseline_daily_avg:.2f}`\n"
f"Trigger: spend above `{args.spend_anomaly_multiplier}x` the daily average "
f"(minimum `${args.spend_anomaly_min_spend:.2f}`)"
),
cache_key=f"user_spend_alert_anomaly_{row.user_id}_{today_str}",
cache_ttl=DAY_SECONDS,
)
def evaluate_user_spend(
row: UserSpendRow,
args: SlackAlertingArgs,
today: datetime.date,
thresholds_enabled: bool,
anomalies_enabled: bool,
) -> tuple[UserSpendAlertEvent, ...]:
today_str: Final = today.strftime("%Y-%m-%d")
month_str: Final = today.strftime("%Y-%m")
threshold_events: Final = (
(
_daily_threshold_event(row=row, args=args, today_str=today_str),
_monthly_threshold_event(row=row, args=args, month_str=month_str),
)
if thresholds_enabled
else ()
)
anomaly_events: Final = (_anomaly_event(row=row, args=args, today_str=today_str),) if anomalies_enabled else ()
return tuple(event for event in (*threshold_events, *anomaly_events) if event is not None)

72
litellm/integrations/arize/arize_phoenix_prompt_manager.py
View file

@ -3,10 +3,12 @@ Arize Phoenix prompt manager that integrates with LiteLLM's prompt management sy
Fetches prompt versions from Arize Phoenix and provides workspace-based access control.
"""
from typing import Any, Final
from collections.abc import Mapping, Sequence
from typing import Any, Final, cast
from jinja2 import DictLoader, select_autoescape
from jinja2.sandbox import ImmutableSandboxedEnvironment
from typing_extensions import ReadOnly, TypedDict
from litellm.integrations.custom_prompt_management import CustomPromptManagement
from litellm.integrations.prompt_management_base import (
@ -20,6 +22,31 @@ from litellm.types.utils import StandardCallbackDynamicParams
from .arize_phoenix_client import ArizePhoenixClient
class ArizePhoenixContentPart(TypedDict, total=False):
type: ReadOnly[str]
text: ReadOnly[str]
class ArizePhoenixTemplateMessage(TypedDict, total=False):
role: ReadOnly[str]
content: ReadOnly[Sequence[ArizePhoenixContentPart]]
class ArizePhoenixTemplateBody(TypedDict, total=False):
messages: ReadOnly[Sequence[ArizePhoenixTemplateMessage]]
class ArizePhoenixPromptMetadata(TypedDict):
model_name: ReadOnly[str | None]
model_provider: ReadOnly[str | None]
description: ReadOnly[str]
template_type: ReadOnly[str | None]
template_format: ReadOnly[str]
invocation_parameters: ReadOnly[Mapping[str, Mapping[str, object]]]
temperature: ReadOnly[float | None]
max_tokens: ReadOnly[int | None]
class ArizePhoenixPromptTemplate:
"""
Represents a prompt template loaded from Arize Phoenix.
@ -28,10 +55,10 @@ class ArizePhoenixPromptTemplate:
def __init__(
self,
template_id: str,
messages: list[dict[str, Any]],
metadata: dict[str, Any],
messages: Sequence[ArizePhoenixTemplateMessage],
metadata: ArizePhoenixPromptMetadata,
model: str | None = None,
):
) -> None:
self.template_id = template_id
self.messages = messages
self.metadata = metadata
@ -43,7 +70,7 @@ class ArizePhoenixPromptTemplate:
self.description = metadata.get("description", "")
self.template_format = metadata.get("template_format", "MUSTACHE")
def __repr__(self):
def __repr__(self) -> str:
return f"ArizePhoenixPromptTemplate(id='{self.template_id}', model='{self.model}')"
@ -109,7 +136,7 @@ class ArizePhoenixTemplateManager:
def _parse_prompt_data(self, data: dict[str, Any], prompt_version_id: str) -> ArizePhoenixPromptTemplate:
"""Parse Arize Phoenix prompt data and extract messages and metadata."""
template_data: Final = data.get("template", {})
template_data: Final[ArizePhoenixTemplateBody] = data.get("template", {})
messages: Final = template_data.get("messages", [])
# Extract invocation parameters
@ -129,7 +156,7 @@ class ArizePhoenixTemplateManager:
break
# Build metadata dictionary
metadata: Final = {
metadata: Final[ArizePhoenixPromptMetadata] = {
"model_name": data.get("model_name"),
"model_provider": data.get("model_provider"),
"description": data.get("description", ""),
@ -146,7 +173,9 @@ class ArizePhoenixTemplateManager:
metadata=metadata,
)
def render_template(self, template_id: str, variables: dict[str, Any] | None = None) -> list[AllMessageValues]:
def render_template(
self, template_id: str, variables: Mapping[str, object] | None = None
) -> list[AllMessageValues]:
"""Render a template with the given variables and return formatted messages."""
if template_id not in self.prompts:
raise ValueError(f"Template '{template_id}' not found")
@ -174,7 +203,9 @@ class ArizePhoenixTemplateManager:
# Combine rendered content
final_content = " ".join(rendered_content_parts)
rendered_messages.append({"role": role, "content": final_content})
rendered_messages.append(
cast("AllMessageValues", {"role": role, "content": final_content}) # cast-ok: Phoenix roles are OpenAI
)
return rendered_messages
@ -243,8 +274,8 @@ class ArizePhoenixPromptManager(CustomPromptManagement):
def get_prompt_template(
self,
prompt_id: str,
prompt_variables: dict[str, Any] | None = None,
) -> tuple[list[AllMessageValues], dict[str, Any]]:
prompt_variables: Mapping[str, object] | None = None,
) -> tuple[list[AllMessageValues], dict[str, object]]:
"""
Get a prompt template and render it with variables.
@ -263,7 +294,7 @@ class ArizePhoenixPromptManager(CustomPromptManagement):
rendered_messages: Final = self.prompt_manager.render_template(prompt_id, prompt_variables or {})
# Extract metadata
metadata: Final = {
metadata: Final[dict[str, object]] = {
"model": template.model,
"temperature": template.temperature,
"max_tokens": template.max_tokens,
@ -271,7 +302,7 @@ class ArizePhoenixPromptManager(CustomPromptManagement):
# Add additional invocation parameters
invocation_params: Final = template.invocation_parameters
provider_params = {}
provider_params: Mapping[str, object] = {}
if "openai" in invocation_params:
provider_params = invocation_params["openai"]
@ -289,12 +320,12 @@ class ArizePhoenixPromptManager(CustomPromptManagement):
self,
user_id: str | None,
messages: list[AllMessageValues],
function_call: dict[str, Any] | str | None = None,
litellm_params: dict[str, Any] | None = None,
function_call: dict[str, object] | str | None = None,
litellm_params: dict[str, object] | None = None,
prompt_id: str | None = None,
prompt_variables: dict[str, Any] | None = None,
prompt_variables: dict[str, object] | None = None,
**kwargs,
) -> tuple[list[AllMessageValues], dict[str, Any] | None]:
) -> tuple[list[AllMessageValues], dict[str, object] | None]:
"""
Pre-call hook that processes the prompt template before making the LLM call.
"""
@ -335,9 +366,9 @@ class ArizePhoenixPromptManager(CustomPromptManagement):
except Exception as e:
# Log error but don't fail the call
import litellm
from litellm._logging import verbose_proxy_logger
litellm._logging.verbose_proxy_logger.error("Error in Arize Phoenix prompt pre_call_hook: %s", e)
verbose_proxy_logger.error("Error in Arize Phoenix prompt pre_call_hook: %s", e)
return messages, litellm_params
def get_available_prompts(self) -> list[str]:
@ -393,7 +424,8 @@ class ArizePhoenixPromptManager(CustomPromptManagement):
rendered_messages, prompt_metadata = self.get_prompt_template(prompt_id, prompt_variables)
# Extract model from metadata (if specified)
template_model: Final = prompt_metadata.get("model")
raw_template_model: Final = prompt_metadata.get("model")
template_model: Final = raw_template_model if isinstance(raw_template_model, str) else None
# Extract optional parameters from metadata
optional_params: Final = {}

55
litellm/integrations/bitbucket/bitbucket_client.py
View file

@ -4,11 +4,38 @@ BitBucket API client for fetching .prompt files from BitBucket repositories.
import base64
import urllib.parse
from typing import Any, Final
from collections.abc import Mapping
from typing import Final, TypedDict
from typing_extensions import NotRequired, ReadOnly
from litellm.llms.custom_httpx.http_handler import HTTPHandler
class BitBucketSrcEntry(TypedDict):
path: ReadOnly[NotRequired[str]]
type: ReadOnly[NotRequired[str]]
class BitBucketSrcListing(TypedDict):
values: ReadOnly[NotRequired[list[BitBucketSrcEntry]]]
class BitBucketBranch(TypedDict):
name: ReadOnly[NotRequired[str]]
type: ReadOnly[NotRequired[str]]
class BitBucketBranchListing(TypedDict):
values: ReadOnly[NotRequired[list[BitBucketBranch]]]
class BitBucketFileMetadata(TypedDict):
content_type: ReadOnly[str | None]
content_length: ReadOnly[str | None]
last_modified: ReadOnly[str | None]
def _sanitize_file_path(file_path: str) -> str:
"""Reject path traversal and URL-encode each path segment."""
if "#" in file_path or "?" in file_path:
@ -31,7 +58,7 @@ class BitBucketClient:
- Branch-specific file fetching
"""
def __init__(self, config: dict[str, Any]):
def __init__(self, config: Mapping[str, object]):
"""
Initialize the BitBucket client.
@ -135,16 +162,12 @@ class BitBucketClient:
response: Final = self.http_handler.get(url, headers=self.headers)
response.raise_for_status()
data: Final = response.json()
files: Final = []
for item in data.get("values", []):
if item.get("type") == "commit_file":
file_path = item.get("path", "")
if file_path.endswith(file_extension):
files.append(file_path)
return files
data: Final[BitBucketSrcListing] = response.json()
return [
file_path
for item in data.get("values", [])
if item.get("type") == "commit_file" and (file_path := item.get("path", "")).endswith(file_extension)
]
except Exception as e:
# Check if it's an HTTP error
@ -162,7 +185,7 @@ class BitBucketClient:
else:
raise Exception(f"Error listing files in '{directory_path}': {e}")
def get_repository_info(self) -> dict[str, Any]:
def get_repository_info(self) -> Mapping[str, object]:
"""
Get information about the repository.
@ -191,7 +214,7 @@ class BitBucketClient:
except Exception:
return False
def get_branches(self) -> list[dict[str, Any]]:
def get_branches(self) -> list[BitBucketBranch]:
"""
Get list of branches in the repository.
@ -204,12 +227,12 @@ class BitBucketClient:
response: Final = self.http_handler.get(url, headers=self.headers)
response.raise_for_status()
data: Final = response.json()
data: Final[BitBucketBranchListing] = response.json()
return data.get("values", [])
except Exception as e:
raise Exception(f"Failed to get branches: {e}")
def get_file_metadata(self, file_path: str) -> dict[str, Any] | None:
def get_file_metadata(self, file_path: str) -> BitBucketFileMetadata | None:
"""
Get metadata about a file (size, last modified, etc.).

37
litellm/integrations/bitbucket/bitbucket_prompt_manager.py
View file

@ -3,6 +3,7 @@ BitBucket prompt manager that integrates with LiteLLM's prompt management system
Fetches .prompt files from BitBucket repositories and provides team-based access control.
"""
from collections.abc import Mapping
from typing import TYPE_CHECKING, Any, Final
from jinja2 import DictLoader, select_autoescape
@ -65,7 +66,7 @@ class BitBucketTemplateManager:
def __init__(
self,
bitbucket_config: dict[str, Any],
bitbucket_config: Mapping[str, object],
prompt_id: str | None = None,
):
self.bitbucket_config = bitbucket_config
@ -123,7 +124,7 @@ class BitBucketTemplateManager:
template_content = content
# Parse YAML frontmatter
metadata: dict[str, Any] = {}
metadata: dict[str, object] = {}
if frontmatter_str:
try:
import yaml
@ -141,9 +142,9 @@ class BitBucketTemplateManager:
metadata=metadata,
)
def _parse_yaml_basic(self, yaml_str: str) -> dict[str, Any]:
def _parse_yaml_basic(self, yaml_str: str) -> dict[str, object]:
"""Basic YAML parser for simple cases when PyYAML is not available."""
result: Final[dict[str, Any]] = {}
result: Final[dict[str, object]] = {}
for line in yaml_str.split("\n"):
line = line.strip()
if ":" in line and not line.startswith("#"):
@ -162,7 +163,7 @@ class BitBucketTemplateManager:
result[key] = value.strip("\"'")
return result
def render_template(self, template_id: str, variables: dict[str, Any] | None = None) -> str:
def render_template(self, template_id: str, variables: Mapping[str, object] | None = None) -> str:
"""Render a template with the given variables."""
if template_id not in self.prompts:
raise ValueError(f"Template '{template_id}' not found")
@ -209,7 +210,7 @@ class BitBucketPromptManager(CustomPromptManagement):
def __init__(
self,
bitbucket_config: dict[str, Any],
bitbucket_config: Mapping[str, object],
prompt_id: str | None = None,
):
self.bitbucket_config = bitbucket_config
@ -234,7 +235,7 @@ class BitBucketPromptManager(CustomPromptManagement):
def get_prompt_template(
self,
prompt_id: str,
prompt_variables: dict[str, Any] | None = None,
prompt_variables: Mapping[str, object] | None = None,
) -> tuple[str, dict[str, Any]]:
"""
Get a prompt template and render it with variables.
@ -267,12 +268,12 @@ class BitBucketPromptManager(CustomPromptManagement):
self,
user_id: str | None,
messages: list[AllMessageValues],
function_call: dict[str, Any] | str | None = None,
litellm_params: dict[str, Any] | None = None,
function_call: Mapping[str, object] | str | None = None,
litellm_params: dict[str, object] | None = None,
prompt_id: str | None = None,
prompt_variables: dict[str, Any] | None = None,
prompt_variables: Mapping[str, object] | None = None,
**kwargs,
) -> tuple[list[AllMessageValues], dict[str, Any] | None]:
) -> tuple[list[AllMessageValues], dict[str, object] | None]:
"""
Pre-call hook that processes the prompt template before making the LLM call.
"""
@ -316,9 +317,9 @@ class BitBucketPromptManager(CustomPromptManagement):
except Exception as e:
# Log error but don't fail the call
import litellm
from litellm._logging import verbose_proxy_logger
litellm._logging.verbose_proxy_logger.error("Error in BitBucket prompt pre_call_hook: %s", e)
verbose_proxy_logger.error("Error in BitBucket prompt pre_call_hook: %s", e)
return messages, litellm_params
def _parse_prompt_to_messages(self, prompt_content: str) -> list[AllMessageValues]:
@ -384,14 +385,14 @@ class BitBucketPromptManager(CustomPromptManagement):
def post_call_hook(
self,
user_id: str | None,
response: Any,
response: object,
input_messages: list[AllMessageValues],
function_call: dict[str, Any] | str | None = None,
litellm_params: dict[str, Any] | None = None,
function_call: Mapping[str, object] | str | None = None,
litellm_params: Mapping[str, object] | None = None,
prompt_id: str | None = None,
prompt_variables: dict[str, Any] | None = None,
prompt_variables: Mapping[str, object] | None = None,
**kwargs,
) -> Any:
) -> object:
"""
Post-call hook for any post-processing after the LLM call.
"""

27
litellm/integrations/cloudzero/transform.py
View file

@ -19,14 +19,29 @@
"""Transform LiteLLM data to CloudZero AnyCost CBF format."""
from datetime import datetime
from typing import Any, Final
from typing import Final, SupportsFloat, SupportsIndex, SupportsInt
import polars as pl
from typing_extensions import Buffer
from ...types.integrations.cloudzero import CBFRecord
from .cz_resource_names import CZEntityType, CZRNGenerator
def _as_int(value: object) -> int:
"""The integer form of a spend table cell, computed the way :func:`int` computes it."""
if isinstance(value, (str, Buffer, SupportsInt, SupportsIndex)):
return int(value)
raise TypeError(f"int() argument must be a string or a number, not {type(value).__name__!r}")
def _as_float(value: object) -> float:
"""The floating point form of a spend table cell, computed the way :func:`float` computes it."""
if isinstance(value, (str, Buffer, SupportsFloat, SupportsIndex)):
return float(value)
raise TypeError(f"float() argument must be a string or a number, not {type(value).__name__!r}")
class CBFTransformer:
"""Transform LiteLLM usage data to CloudZero Billing Format (CBF)."""
@ -82,15 +97,15 @@ class CBFTransformer:
return pl.DataFrame(cbf_data)
def _create_cbf_record(self, row: dict[str, Any]) -> CBFRecord:
def _create_cbf_record(self, row: dict[str, object]) -> CBFRecord:
"""Create a single CBF record from LiteLLM daily spend row."""
# Parse date (daily spend tables use date strings like '2025-04-19')
usage_date: Final = self._parse_date(row.get("date"))
# Calculate total tokens
prompt_tokens: Final = int(row.get("prompt_tokens", 0))
completion_tokens: Final = int(row.get("completion_tokens", 0))
prompt_tokens: Final = _as_int(row.get("prompt_tokens", 0))
completion_tokens: Final = _as_int(row.get("completion_tokens", 0))
total_tokens: Final = prompt_tokens + completion_tokens
# Create CloudZero Resource Name (CZRN) as resource_id
@ -154,7 +169,7 @@ class CBFTransformer:
"time/usage_start": (
usage_date.isoformat() if usage_date else None
), # Required: ISO-formatted UTC datetime
"cost/cost": float(row.get("spend", 0.0)), # Required: billed cost
"cost/cost": _as_float(row.get("spend", 0.0)), # Required: billed cost
"resource/id": resource_id, # CZRN (CloudZero Resource Name)
# Usage metrics for token consumption
"usage/amount": total_tokens, # Numeric value of tokens consumed
@ -187,7 +202,7 @@ class CBFTransformer:
return CBFRecord(cbf_record)
def _parse_date(self, date_str) -> datetime | None:
def _parse_date(self, date_str: object) -> datetime | None:
"""Parse date string from daily spend tables (e.g., '2025-04-19')."""
if date_str is None:
return None

84
litellm/integrations/compression_interception/handler.py
View file

@ -7,7 +7,10 @@ litellm_content_retrieve tool calls server-side via the typed agentic loop plan.
import time
import uuid
from typing import Any, ClassVar, Final, cast
from collections.abc import Mapping, Sequence
from typing import Any, ClassVar, Final, Protocol, cast
from typing_extensions import ReadOnly, TypedDict
from litellm._logging import verbose_logger
from litellm.compression import compress
@ -26,6 +29,19 @@ LITELLM_CONTENT_RETRIEVE_TOOL_NAME: Final = "litellm_content_retrieve"
_CACHE_TTL_SECONDS: Final = 15 * 60
class _AgenticLoopParams(TypedDict, total=False):
"""The ``agentic_loop_params`` entry the agentic loop driver records on the logging object."""
model: ReadOnly[str]
class _AgenticLoopLoggingObj(Protocol):
"""Logging object view exposing the untyped call details this handler reads."""
@property
def model_call_details(self) -> Mapping[str, _AgenticLoopParams]: ...
def _compression_savings_from_counts(
original_tokens: object, compressed_tokens: object
) -> CompressionSavingsMetadata | None:
@ -80,7 +96,7 @@ class CompressionInterceptionLogger(CustomLogger):
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: dict[str, object] | None = None,
):
super().__init__()
self.enabled = enabled
@ -103,7 +119,7 @@ class CompressionInterceptionLogger(CustomLogger):
@staticmethod
def initialize_from_proxy_config(
litellm_settings: dict[str, Any],
callback_specific_params: dict[str, Any],
callback_specific_params: Mapping[str, object],
) -> "CompressionInterceptionLogger":
compression_params: CompressionInterceptionConfig = {}
if "compression_interception_params" in litellm_settings:
@ -117,7 +133,9 @@ class CompressionInterceptionLogger(CustomLogger):
)
return CompressionInterceptionLogger.from_config_yaml(compression_params)
async def async_pre_call_deployment_hook(self, kwargs: dict[str, Any], call_type: CallTypes | None) -> dict | None:
async def async_pre_call_deployment_hook(
self, kwargs: dict[str, Any], call_type: CallTypes | None
) -> dict[str, object] | None:
if not self.enabled:
return None
if call_type is not None and call_type != CallTypes.anthropic_messages:
@ -147,7 +165,7 @@ class CompressionInterceptionLogger(CustomLogger):
cache: Final = cast(dict[str, str], compressed.get("cache", {}))
skip_reason: Final = cast(str | None, compressed.get("compression_skipped_reason"))
compressed_tools: Final = cast(list[dict[str, Any]], compressed.get("tools", []))
compressed_tools: Final = cast(list[dict[str, object]], compressed.get("tools", []))
# Only mutate kwargs when compression actually produced a result.
# If compression was a no-op (below trigger, invalid tool sequence, etc.),
@ -158,7 +176,7 @@ class CompressionInterceptionLogger(CustomLogger):
kwargs["messages"] = compressed["messages"]
if compressed_tools:
kwargs["tools"] = self._merge_tools(
existing_tools=cast(list[dict[str, Any]] | None, kwargs.get("tools")),
existing_tools=cast(list[dict[str, object]] | None, kwargs.get("tools")),
compressed_tools=compressed_tools,
)
call_id = cast(str | None, kwargs.get("litellm_call_id"))
@ -191,14 +209,14 @@ class CompressionInterceptionLogger(CustomLogger):
async def async_should_run_agentic_loop(
self,
response: Any,
response: object,
model: str,
messages: list[dict],
tools: list[dict] | None,
messages: Sequence[Mapping[str, object]],
tools: Sequence[Mapping[str, object]] | None,
stream: bool,
custom_llm_provider: str,
kwargs: dict,
) -> tuple[bool, dict]:
kwargs: Mapping[str, object],
) -> tuple[bool, dict[str, object]]:
if not self.enabled:
return False, {}
if not self._has_retrieval_tool(tools):
@ -216,19 +234,19 @@ class CompressionInterceptionLogger(CustomLogger):
async def async_build_agentic_loop_plan(
self,
tools: dict,
tools: Mapping[str, object],
model: str,
messages: list[dict],
response: Any,
anthropic_messages_provider_config: Any,
anthropic_messages_optional_request_params: dict,
logging_obj: Any,
messages: list[dict[str, object]],
response: object,
anthropic_messages_provider_config: object,
anthropic_messages_optional_request_params: Mapping[str, object],
logging_obj: _AgenticLoopLoggingObj | None,
stream: bool,
kwargs: dict,
kwargs: Mapping[str, object],
) -> AgenticLoopPlan:
self._prune_expired_cache()
tool_calls: Final = cast(list[dict[str, Any]], tools.get("tool_calls", []))
thinking_blocks: Final = cast(list[dict[str, Any]], tools.get("thinking_blocks", []))
tool_calls: Final = cast(list[dict[str, object]], tools.get("tool_calls", []))
thinking_blocks: Final = cast(list[dict[str, object]], tools.get("thinking_blocks", []))
call_id: Final = self._resolve_call_id(logging_obj=logging_obj, kwargs=kwargs)
cache: Final = self._get_cache(call_id=call_id)
@ -271,7 +289,7 @@ class CompressionInterceptionLogger(CustomLogger):
full_model_name = model
if logging_obj is not None:
agentic_params: Final = logging_obj.model_call_details.get("agentic_loop_params", {})
full_model_name = cast(str, agentic_params.get("model", model))
full_model_name = agentic_params.get("model", model)
request_patch: Final = AgenticLoopRequestPatch(
model=full_model_name,
@ -306,15 +324,15 @@ class CompressionInterceptionLogger(CustomLogger):
return {}
return cache_entry[0]
def _resolve_call_id(self, logging_obj: Any, kwargs: dict[str, Any]) -> str | None:
def _resolve_call_id(self, logging_obj: _AgenticLoopLoggingObj | None, kwargs: Mapping[str, object]) -> str | None:
if logging_obj is not None:
logging_call_id: Final = getattr(logging_obj, "litellm_call_id", None)
if isinstance(logging_call_id, str) and logging_call_id:
return logging_call_id
kwargs_call_id: Final = kwargs.get("litellm_call_id")
return cast(str | None, kwargs_call_id if isinstance(kwargs_call_id, str) else None)
return kwargs_call_id if isinstance(kwargs_call_id, str) else None
def _resolve_retrieval_content(self, tool_call: dict[str, Any], cache: dict[str, str]) -> str:
def _resolve_retrieval_content(self, tool_call: Mapping[str, object], cache: Mapping[str, str]) -> str:
raw_input: Final = tool_call.get("input", {})
key = ""
if isinstance(raw_input, dict):
@ -325,7 +343,9 @@ class CompressionInterceptionLogger(CustomLogger):
return cache[key]
return f"[compressed content key '{key}' not found]"
def _extract_retrieval_tool_calls(self, response: Any) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]:
def _extract_retrieval_tool_calls(
self, response: object
) -> tuple[list[dict[str, object]], list[dict[str, object]]]:
if isinstance(response, dict):
content = response.get("content", [])
else:
@ -334,8 +354,8 @@ class CompressionInterceptionLogger(CustomLogger):
if not isinstance(content, list):
return [], []
tool_calls: Final[list[dict[str, Any]]] = []
thinking_blocks: Final[list[dict[str, Any]]] = []
tool_calls: Final[list[dict[str, object]]] = []
thinking_blocks: Final[list[dict[str, object]]] = []
for block in content:
if isinstance(block, dict):
@ -382,13 +402,13 @@ class CompressionInterceptionLogger(CustomLogger):
return tool_calls, thinking_blocks
def _prepare_followup_kwargs(self, kwargs: dict[str, Any]) -> dict[str, Any]:
def _prepare_followup_kwargs(self, kwargs: Mapping[str, object]) -> dict[str, object]:
internal_keys: Final = {"litellm_logging_obj"}
return {
k: v for k, v in kwargs.items() if not k.startswith("_compression_interception") and k not in internal_keys
}
def _has_retrieval_tool(self, tools: Any) -> bool:
def _has_retrieval_tool(self, tools: object) -> bool:
if not isinstance(tools, list):
return False
for tool in tools:
@ -404,9 +424,9 @@ class CompressionInterceptionLogger(CustomLogger):
def _merge_tools(
self,
existing_tools: list[dict[str, Any]] | None,
compressed_tools: list[dict[str, Any]],
) -> list[dict[str, Any]]:
existing_tools: Sequence[Mapping[str, object]] | None,
compressed_tools: Sequence[Mapping[str, object]],
) -> list[Mapping[str, object]]:
merged: Final = list(existing_tools or [])
if self._has_retrieval_tool(merged):
return merged

13
litellm/integrations/custom_guardrail.py
View file

@ -2,6 +2,7 @@ import contextvars
import hashlib
import os
import secrets
from collections.abc import Mapping
from datetime import datetime
from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, Optional, get_args
@ -227,13 +228,13 @@ class CustomGuardrail(CustomLogger):
)
super().__init__(**kwargs)
def render_violation_message(self, default: str, context: dict[str, Any] | None = None) -> str:
def render_violation_message(self, default: str, context: Mapping[str, object] | None = None) -> str:
"""Return a custom violation message if template is configured."""
if not self.violation_message_template:
return default
format_context: Final[dict[str, Any]] = {"default_message": default}
format_context: Final[dict[str, object]] = {"default_message": default}
if context:
format_context.update(context)
try:
@ -661,7 +662,7 @@ class CustomGuardrail(CustomLogger):
value: Final = self._get_admin_metadata(data).get("opted_out_global_guardrails")
return value if isinstance(value, list) else []
def _is_valid_response_type(self, result: Any) -> bool:
def _is_valid_response_type(self, result: object) -> bool:
"""
Check if result is a valid LLMResponseTypes instance.
@ -722,7 +723,7 @@ class CustomGuardrail(CustomLogger):
return None
return f"{_PRE_CALL_EXECUTED_TOKEN}:{name}"
def mark_pre_call_hook_ran(self, data: dict[str, Any]) -> None:
def mark_pre_call_hook_ran(self, data: dict[str, object]) -> None:
"""
Record that this guardrail's ``async_pre_call_hook`` already ran for this
request, so the deployment-level hook does not run it a second time.
@ -747,7 +748,7 @@ class CustomGuardrail(CustomLogger):
return
data["metadata"] = {PRE_CALL_EXECUTED_GUARDRAILS_KEY: [marker]}
def _pre_call_hook_already_ran(self, data: dict[str, Any]) -> bool:
def _pre_call_hook_already_ran(self, data: dict[str, object]) -> bool:
marker: Final = self._pre_call_marker()
if marker is None:
return False
@ -1170,7 +1171,7 @@ class CustomGuardrail(CustomLogger):
This gets logged on downsteam Langfuse, DataDog, etc.
"""
# Convert None to empty dict to satisfy type requirements
guardrail_response: dict[str, Any] | str = {} if response is None else response
guardrail_response: dict[str, object] | str = {} if response is None else response
# For apply_guardrail functions in custom_code_guardrail scenario,
# simplify the logged response to "allow", "deny", or "mask"

75
litellm/integrations/custom_logger.py
View file

@ -2,7 +2,7 @@
# On success, logs events to Promptlayer
import re
import traceback
from collections.abc import AsyncGenerator, Mapping
from collections.abc import AsyncGenerator, Mapping, Sequence
from typing import TYPE_CHECKING, Any, ClassVar, Final, Optional
from pydantic import BaseModel
@ -31,6 +31,9 @@ if TYPE_CHECKING:
from litellm.caching.caching import DualCache
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.base_llm.anthropic_messages.transformation import (
BaseAnthropicMessagesConfig,
)
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.mcp import (
MCPPostCallResponseObject,
@ -39,7 +42,7 @@ if TYPE_CHECKING:
)
from litellm.types.router import PreRoutingHookResponse
Span = _Span | Any
Span = _Span
else:
Span = Any
LiteLLMLoggingObj = Any
@ -123,11 +126,11 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
return []
callbacks: Final = AllCallbacks()
callback_info: Final = getattr(callbacks, lookup_name, None)
callback_info: Final[object] = getattr(callbacks, lookup_name, None)
if callback_info is None:
return []
params: Final = getattr(callback_info, "litellm_callback_params", None)
params: Final[Sequence[str] | None] = getattr(callback_info, "litellm_callback_params", None)
if not params:
return []
@ -268,7 +271,9 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
) -> list[dict]:
return healthy_deployments
async def async_pre_call_deployment_hook(self, kwargs: dict[str, Any], call_type: CallTypes | None) -> dict | None:
async def async_pre_call_deployment_hook(
self, kwargs: dict[str, object], call_type: CallTypes | None
) -> dict | None:
"""
Allow modifying the request just before it's sent to the deployment.
@ -344,9 +349,9 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
async def async_post_call_streaming_deployment_hook(
self,
request_data: dict,
response_chunk: Any,
response_chunk: object,
call_type: CallTypes | None,
) -> Any | None:
) -> object | None:
"""
Allow modifying streaming chunks just before they're returned to the user.
@ -378,7 +383,7 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
"""
def translate_completion_output_params_streaming(
self, completion_stream: Any
self, completion_stream: object
) -> AdapterCompletionStreamWrapper | None:
"""
Translates the streaming chunk, from the OpenAI format to the custom format.
@ -418,9 +423,9 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
self,
data: dict,
user_api_key_dict: UserAPIKeyAuth,
response: Any,
response: object,
request_headers: dict[str, str] | None = None,
litellm_call_info: dict[str, Any] | None = None,
litellm_call_info: dict[str, object] | None = None,
) -> dict[str, str] | None:
"""
Called after an LLM API call (success or failure) to allow injecting custom HTTP response headers.
@ -471,11 +476,11 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
) -> Any:
pass
async def async_logging_hook(self, kwargs: dict, result: Any, call_type: str) -> tuple[dict, Any]:
async def async_logging_hook(self, kwargs: dict, result: object, call_type: str) -> tuple[dict, object]:
"""For masking logged request/response. Return a modified version of the request/result."""
return kwargs, result
def logging_hook(self, kwargs: dict, result: Any, call_type: str) -> tuple[dict, Any]:
def logging_hook(self, kwargs: dict, result: object, call_type: str) -> tuple[dict, object]:
"""For masking logged request/response. Return a modified version of the request/result."""
return kwargs, result
@ -581,7 +586,7 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
async def async_should_run_agentic_loop(
self,
response: Any,
response: object,
model: str,
messages: list[dict],
tools: list[dict] | None,
@ -642,8 +647,8 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
tools: dict,
model: str,
messages: list[dict],
response: Any,
anthropic_messages_provider_config: Any,
response: object,
anthropic_messages_provider_config: "BaseAnthropicMessagesConfig | None",
anthropic_messages_optional_request_params: dict,
logging_obj: "LiteLLMLoggingObj",
stream: bool,
@ -711,8 +716,8 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
tools: dict,
model: str,
messages: list[dict],
response: Any,
anthropic_messages_provider_config: Any,
response: object,
anthropic_messages_provider_config: "BaseAnthropicMessagesConfig | None",
anthropic_messages_optional_request_params: dict,
logging_obj: "LiteLLMLoggingObj",
stream: bool,
@ -728,7 +733,7 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
async def async_post_agentic_loop_response_hook(
self,
response: Any,
response: object,
plan: AgenticLoopPlan,
kwargs: dict,
) -> Any:
@ -767,7 +772,7 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
async def async_should_run_chat_completion_agentic_loop(
self,
response: Any,
response: object,
model: str,
messages: list[dict],
tools: list[dict] | None,
@ -785,12 +790,12 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
tools: dict,
model: str,
messages: list[dict],
response: Any,
response: object,
optional_params: dict,
logging_obj: "LiteLLMLoggingObj",
stream: bool,
kwargs: dict,
) -> Any:
) -> object:
"""
Hook to execute chat completion agentic loop based on context from should_run hook.
"""
@ -800,7 +805,7 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
tools: dict,
model: str,
messages: list[dict],
response: Any,
response: object,
optional_params: dict,
logging_obj: "LiteLLMLoggingObj",
stream: bool,
@ -851,7 +856,7 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
- Converting to string and then truncating the logged content catches this
2. We want to avoid modifying the original `messages`, `response`, and `error_str` in the logging payload since these are in kwargs and could be returned to the user
"""
field_value: Final = standard_logging_object.get(field_name)
field_value: Final[object] = standard_logging_object.get(field_name)
if field_value:
str_value: Final = str(field_value)
if len(str_value) > max_length:
@ -1005,8 +1010,8 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
• Keep untyped or text content.
• Recursively redact inline base64 blobs in *any* string field, at any depth.
"""
raw_messages: Final[Any] = payload.get("messages", [])
messages: Final[list[Any]] = raw_messages if isinstance(raw_messages, list) else []
raw_messages: Final[object] = payload.get("messages", [])
messages: Final[list[object]] = raw_messages if isinstance(raw_messages, list) else []
verbose_logger.debug("[CustomLogger] Stripping base64 from %s messages", len(messages))
if messages:
@ -1037,8 +1042,8 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
• Keep untyped or text content.
• Recursively redact inline base64 blobs in *any* string field, at any depth.
"""
raw_messages: Final[Any] = payload.get("messages", [])
messages: Final[list[Any]] = raw_messages if isinstance(raw_messages, list) else []
raw_messages: Final[object] = payload.get("messages", [])
messages: Final[list[object]] = raw_messages if isinstance(raw_messages, list) else []
verbose_logger.debug("[CustomLogger] Stripping base64 from %s messages", len(messages))
if messages:
@ -1056,10 +1061,10 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
def _redact_base64(
self,
value: Any,
value: object,
depth: int = 0,
max_depth: int = DEFAULT_MAX_RECURSE_DEPTH_SENSITIVE_DATA_MASKER,
) -> Any:
) -> object:
"""Recursively redact inline base64 from any nested structure with a max recursion depth limit."""
if depth > max_depth:
verbose_logger.warning("[CustomLogger] Max recursion depth %s reached while redacting base64", max_depth)
@ -1079,7 +1084,7 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
return value
def _should_keep_content(self, content: Any) -> bool:
def _should_keep_content(self, content: object) -> bool:
"""Return True if this content item should be retained."""
if not isinstance(content, dict):
return True
@ -1090,16 +1095,16 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
def _process_messages(
self,
messages: list[Any],
messages: list[object],
max_depth: int = DEFAULT_MAX_RECURSE_DEPTH_SENSITIVE_DATA_MASKER,
) -> list[dict[str, Any]]:
filtered_messages: Final[list[dict[str, Any]]] = []
) -> list[dict[str, object]]:
filtered_messages: Final[list[dict[str, object]]] = []
for msg in messages:
if not isinstance(msg, dict):
continue
contents: Any = msg.get("content")
contents: object = msg.get("content")
if isinstance(contents, list):
cleaned: list[Any] = []
cleaned: list[object] = []
for c in contents:
if self._should_keep_content(content=c):
cleaned.append(self._redact_base64(value=c, max_depth=max_depth))

44
litellm/integrations/datadog/datadog.py
View file

@ -20,10 +20,11 @@ import time
import traceback
from collections.abc import Sequence
from datetime import datetime as datetimeObj
from typing import Any, Final
from typing import TYPE_CHECKING, Any, Final
import httpx
from httpx import Response
from typing_extensions import ReadOnly, TypedDict
import litellm
from litellm._logging import verbose_logger
@ -62,6 +63,18 @@ from litellm.types.utils import StandardLoggingPayload
from ..additional_logging_utils import AdditionalLoggingUtils
if TYPE_CHECKING:
from fastapi import HTTPException
from litellm.proxy._types import UserAPIKeyAuth
class _DatadogLoggingKwargs(TypedDict, total=False):
"""The subset of logging ``kwargs`` that the Datadog payload builder reads."""
standard_logging_object: ReadOnly[StandardLoggingPayload | None]
# max number of logs DD API can accept
@ -87,6 +100,11 @@ def _resolve_dd_batch_size() -> int:
return max(1, min(value, DD_MAX_BATCH_SIZE))
def _span_attribute(span: object, name: str) -> object:
"""Read an optional attribute off whatever span object the active tracer hands back."""
return getattr(span, name, None)
class DataDogLogger(
CustomBatchLogger,
AdditionalLoggingUtils,
@ -271,9 +289,9 @@ class DataDogLogger(
self,
request_data: dict,
original_exception: Exception,
user_api_key_dict: Any,
user_api_key_dict: "UserAPIKeyAuth",
traceback_str: str | None = None,
) -> Any | None:
) -> "HTTPException | None":
"""
Log proxy-level failures (e.g. 401 auth, DB connection errors) to Datadog.
@ -297,7 +315,7 @@ class DataDogLogger(
status_code = int(_code)
# Use project-standard sanitized user context when running in proxy
user_context: dict[str, Any] = {}
user_context: dict[str, object] = {}
try:
from litellm.proxy.litellm_pre_call_utils import (
LiteLLMProxyRequestSetup,
@ -553,8 +571,8 @@ class DataDogLogger(
def create_datadog_logging_payload(
self,
kwargs: dict | Any,
response_obj: Any,
kwargs: _DatadogLoggingKwargs,
response_obj: object,
start_time: datetime.datetime,
end_time: datetime.datetime,
) -> DatadogPayload:
@ -562,8 +580,8 @@ class DataDogLogger(
Helper function to create a datadog payload for logging
Args:
kwargs (Union[dict, Any]): request kwargs
response_obj (Any): llm api response
kwargs: request kwargs, read for its standard logging object
response_obj: llm api response
start_time (datetime.datetime): start time of request
end_time (datetime.datetime): end time of request
@ -625,7 +643,7 @@ class DataDogLogger(
self,
payload: ServiceLoggerPayload,
error: str | None = "",
parent_otel_span: Any | None = None,
parent_otel_span: object = None,
start_time: datetimeObj | float | None = None,
end_time: float | datetimeObj | None = None,
event_metadata: dict | None = None,
@ -659,7 +677,7 @@ class DataDogLogger(
self,
payload: ServiceLoggerPayload,
error: str | None = "",
parent_otel_span: Any | None = None,
parent_otel_span: object = None,
start_time: datetimeObj | float | None = None,
end_time: float | datetimeObj | None = None,
event_metadata: dict | None = None,
@ -696,7 +714,7 @@ class DataDogLogger(
def _create_v0_logging_payload(
self,
kwargs: dict | Any,
kwargs: dict,
response_obj: Any,
start_time: datetime.datetime,
end_time: datetime.datetime,
@ -810,11 +828,11 @@ class DataDogLogger(
if current_span is None:
return None
trace_id: Final = getattr(current_span, "trace_id", None)
trace_id: Final = _span_attribute(current_span, "trace_id")
if trace_id is None:
return None
span_id: Final = getattr(current_span, "span_id", None)
span_id: Final = _span_attribute(current_span, "span_id")
trace_context: Final[dict[str, str]] = {"trace_id": str(trace_id)}
if span_id is not None:
trace_context["span_id"] = str(span_id)

19
litellm/integrations/datadog/datadog_llm_obs.py
View file

@ -9,6 +9,7 @@ API Reference: https://docs.datadoghq.com/llm_observability/setup/api/?tab=examp
import asyncio
import json
import os
from collections.abc import Mapping, Sequence
from datetime import datetime
from typing import Any, Final, Literal
@ -334,7 +335,7 @@ class DataDogLLMObsLogger(CustomBatchLogger):
def _get_response_messages(
self, standard_logging_payload: StandardLoggingPayload, call_type: str | None
) -> list[Any]:
) -> list[object]:
"""
Get the messages from the response object
@ -484,7 +485,7 @@ class DataDogLLMObsLogger(CustomBatchLogger):
# Default fallback for unknown or passthrough operations
return "llm"
def _ensure_string_content(self, messages: str | list[Any] | dict[Any, Any] | None) -> list[Any]:
def _ensure_string_content(self, messages: str | Sequence[object] | Mapping[object, object] | None) -> list[object]:
if messages is None:
return []
if isinstance(messages, str):
@ -495,11 +496,11 @@ class DataDogLLMObsLogger(CustomBatchLogger):
return [str(messages.get("content", ""))]
return []
def _get_dd_llm_obs_payload_metadata(self, standard_logging_payload: StandardLoggingPayload) -> dict[str, Any]:
def _get_dd_llm_obs_payload_metadata(self, standard_logging_payload: StandardLoggingPayload) -> dict[str, object]:
"""
Fields to track in DD LLM Observability metadata from litellm standard logging payload
"""
_metadata: Final[dict[str, Any]] = {
_metadata: Final[dict[str, object]] = {
"model_name": standard_logging_payload.get("model", "unknown"),
"model_provider": standard_logging_payload.get("custom_llm_provider", "unknown"),
"id": standard_logging_payload.get("id", "unknown"),
@ -647,7 +648,7 @@ class DataDogLLMObsLogger(CustomBatchLogger):
return spend_metrics
def _process_input_messages_preserving_tool_calls(self, messages: list[Any]) -> list[dict[str, Any]]:
def _process_input_messages_preserving_tool_calls(self, messages: Sequence[object]) -> list[dict[str, object]]:
"""
Process input messages while preserving tool_calls and tool message types.
@ -671,13 +672,13 @@ class DataDogLLMObsLogger(CustomBatchLogger):
return processed
@staticmethod
def _tool_calls_kv_pair(tool_calls: list[dict[str, Any]]) -> dict[str, Any]:
def _tool_calls_kv_pair(tool_calls: list[dict[str, Any]]) -> dict[str, object]:
"""
Extract tool call information into key-value pairs for Datadog metadata.
Similar to OpenTelemetry's implementation but adapted for Datadog's format.
"""
kv_pairs: Final[dict[str, Any]] = {}
kv_pairs: Final[dict[str, object]] = {}
for idx, tool_call in enumerate(tool_calls):
try:
# Extract tool call ID
@ -712,11 +713,11 @@ class DataDogLLMObsLogger(CustomBatchLogger):
return kv_pairs
def _extract_tool_call_metadata(self, standard_logging_payload: StandardLoggingPayload) -> dict[str, Any]:
def _extract_tool_call_metadata(self, standard_logging_payload: StandardLoggingPayload) -> dict[str, object]:
"""
Extract tool call information from both input messages and response for Datadog metadata.
"""
tool_call_metadata: Final[dict[str, Any]] = {}
tool_call_metadata: Final[dict[str, object]] = {}
try:
# Extract tool calls from input messages

25
litellm/integrations/dotprompt/prompt_manager.py
View file

@ -3,12 +3,21 @@ Based on Google's GenAI Kit dotprompt implementation: https://google.github.io/d
"""
import re
from collections.abc import Mapping
from pathlib import Path
from typing import Any, Final
import yaml
from jinja2 import DictLoader, select_autoescape
from jinja2.sandbox import ImmutableSandboxedEnvironment
from typing_extensions import NotRequired, ReadOnly, TypedDict
class _PromptFileJson(TypedDict):
"""JSON form of a .prompt file: rendered template text plus its frontmatter."""
content: ReadOnly[NotRequired[str]]
metadata: ReadOnly[NotRequired[dict[str, object]]]
def strip_version_suffix(prompt_id: str) -> str | None:
@ -167,7 +176,7 @@ class PromptManager:
template_id=prompt_id,
)
def _parse_frontmatter(self, content: str) -> tuple[dict[str, Any], str]:
def _parse_frontmatter(self, content: str) -> tuple[dict[str, object], str]:
"""Parse YAML frontmatter from prompt content."""
# Match YAML frontmatter between --- delimiters
frontmatter_pattern: Final = r"^---\s*\n(.*?)\n---\s*\n(.*)$"
@ -178,7 +187,7 @@ class PromptManager:
template_content = match.group(2)
try:
frontmatter = yaml.safe_load(frontmatter_yaml) or {}
frontmatter: dict[str, object] = yaml.safe_load(frontmatter_yaml) or {}
except yaml.YAMLError as e:
raise ValueError(f"Invalid YAML frontmatter: {e}")
else:
@ -191,7 +200,7 @@ class PromptManager:
def render(
self,
prompt_id: str,
prompt_variables: dict[str, Any] | None = None,
prompt_variables: Mapping[str, object] | None = None,
version: int | None = None,
) -> str:
"""
@ -231,7 +240,7 @@ class PromptManager:
except Exception as e:
raise ValueError(f"Error rendering template '{prompt_id}': {e}")
def _validate_input(self, variables: dict[str, Any], schema: dict[str, Any]) -> None:
def _validate_input(self, variables: Mapping[str, object], schema: Mapping[str, str]) -> None:
"""Basic validation of input variables against schema."""
for field_name, field_type in schema.items():
if field_name in variables:
@ -291,7 +300,7 @@ class PromptManager:
"""Get a list of all available prompt IDs."""
return list(self.prompts.keys())
def get_prompt_metadata(self, prompt_id: str) -> dict[str, Any] | None:
def get_prompt_metadata(self, prompt_id: str) -> dict[str, object] | None:
"""Get metadata for a specific prompt."""
template: Final = self.prompts.get(prompt_id)
return template.metadata if template else None
@ -302,12 +311,12 @@ class PromptManager:
if self.prompt_directory:
self._load_prompts()
def add_prompt(self, prompt_id: str, content: str, metadata: dict[str, Any] | None = None) -> None:
def add_prompt(self, prompt_id: str, content: str, metadata: dict[str, object] | None = None) -> None:
"""Add a prompt template programmatically."""
template: Final = PromptTemplate(content=content, metadata=metadata or {}, template_id=prompt_id)
self.prompts[prompt_id] = template
def prompt_file_to_json(self, file_path: str | Path) -> dict[str, Any]:
def prompt_file_to_json(self, file_path: str | Path) -> _PromptFileJson:
"""Convert a .prompt file to JSON format.
Args:
@ -324,7 +333,7 @@ class PromptManager:
return {"content": template_content.strip(), "metadata": frontmatter}
def json_to_prompt_file(self, prompt_data: dict[str, Any]) -> str:
def json_to_prompt_file(self, prompt_data: _PromptFileJson) -> str:
"""Convert JSON prompt data to .prompt file format.
Args:

55
litellm/integrations/focus/destinations/mavvrik_destination.py
View file

@ -9,9 +9,12 @@ Flow:
from __future__ import annotations
import gzip
from typing import Any, Final
from collections.abc import Mapping
from typing import Final, Protocol
from urllib.parse import urlparse
from typing_extensions import NotRequired, ReadOnly, TypedDict
from litellm._logging import verbose_logger
from litellm.llms.custom_httpx.http_handler import (
AsyncHTTPHandler,
@ -28,6 +31,34 @@ _MAVVRIK_ALLOWED_SUFFIXES: Final = (".mavvrik.dev", ".mavvrik.ai", ".mavvrik.app
_GCS_CHUNK_SIZE: Final = 8 * 1024 * 1024 # 8 MB
class MavvrikRegisterBody(TypedDict):
metricsMarker: ReadOnly[NotRequired[int | str]]
class MavvrikUploadUrlBody(TypedDict):
url: ReadOnly[NotRequired[str]]
class _RegisterResponse(Protocol):
def json(self) -> MavvrikRegisterBody: ...
class _UploadUrlResponse(Protocol):
def json(self) -> MavvrikUploadUrlBody: ...
def _register_body(response: _RegisterResponse) -> MavvrikRegisterBody:
return response.json()
def _upload_url_body(response: _UploadUrlResponse) -> MavvrikUploadUrlBody:
return response.json()
def _header_value(headers: Mapping[str, str], name: str) -> str | None:
return headers.get(name)
def _validate_api_endpoint(api_endpoint: str) -> None:
if not api_endpoint.startswith("https://"):
raise ValueError("MAVVRIK_API_ENDPOINT must be an HTTPS URL")
@ -56,12 +87,12 @@ class FocusMavvrikDestination(FocusDestination):
self,
*,
prefix: str,
config: dict[str, Any] | None = None,
config: Mapping[str, str] | None = None,
) -> None:
config = config or {}
api_key: Final = config.get("api_key")
api_endpoint: Final = config.get("api_endpoint")
connection_id: Final = config.get("connection_id")
resolved_config: Final[Mapping[str, str]] = config or {}
api_key: Final = resolved_config.get("api_key")
api_endpoint: Final = resolved_config.get("api_endpoint")
connection_id: Final = resolved_config.get("connection_id")
if not api_key:
raise ValueError(
@ -100,7 +131,7 @@ class FocusMavvrikDestination(FocusDestination):
def _auth_headers(self) -> dict[str, str]:
return {"Content-Type": "application/json", "x-api-key": self.api_key}
async def _ensure_registered(self) -> int | None:
async def _ensure_registered(self) -> int | str | None:
"""POST agent endpoint to register/initialize the connector (once per instance).
Returns metricsMarker from the Mavvrik response — the last date index
@ -127,7 +158,7 @@ class FocusMavvrikDestination(FocusDestination):
if resp.status_code >= 400:
raise RuntimeError(f"Mavvrik FOCUS destination: register failed ({resp.status_code}): {resp.text[:200]}")
self._registered = True
metrics_marker: Final = resp.json().get("metricsMarker", 0)
metrics_marker: Final = _register_body(resp).get("metricsMarker", 0)
verbose_logger.debug(
"Mavvrik FOCUS destination: connector registered (metricsMarker=%s)",
metrics_marker,
@ -148,7 +179,7 @@ class FocusMavvrikDestination(FocusDestination):
raise RuntimeError(
f"Mavvrik FOCUS destination: failed to get signed URL ({resp.status_code}): {resp.text[:200]}"
)
signed_url: Final = resp.json().get("url")
signed_url: Final = _upload_url_body(resp).get("url")
if not signed_url:
raise RuntimeError(f"Mavvrik FOCUS destination: response missing 'url' field: {resp.json()}")
_validate_gcs_url(signed_url, "signed URL")
@ -190,7 +221,7 @@ class FocusMavvrikDestination(FocusDestination):
f"Mavvrik FOCUS destination: GCS session init failed ({init_resp.status_code}): {init_resp.text[:400]}"
)
session_uri: Final = init_resp.headers.get("Location")
session_uri: Final = _header_value(init_resp.headers, "Location")
if not session_uri:
raise RuntimeError("Mavvrik FOCUS destination: GCS session init missing Location header")
_validate_gcs_url(session_uri, "session URI")
@ -264,7 +295,7 @@ class FocusMavvrikDestination(FocusDestination):
)
verbose_logger.debug("Mavvrik FOCUS destination: metricsMarker advanced to %s", date_epoch)
async def get_metrics_marker(self) -> int | None:
async def get_metrics_marker(self) -> int | str | None:
"""Register with Mavvrik and return the current metricsMarker.
Always calls the Mavvrik register API — unlike deliver() which skips
@ -287,7 +318,7 @@ class FocusMavvrikDestination(FocusDestination):
if resp.status_code >= 400:
raise RuntimeError(f"Mavvrik FOCUS destination: register failed ({resp.status_code}): {resp.text[:200]}")
self._registered = True
metrics_marker: Final = resp.json().get("metricsMarker", 0)
metrics_marker: Final = _register_body(resp).get("metricsMarker", 0)
verbose_logger.debug("Mavvrik FOCUS destination: got metricsMarker=%s", metrics_marker)
return metrics_marker

35
litellm/integrations/galileo.py
View file

@ -6,10 +6,11 @@ import re
import uuid
from collections.abc import Mapping, Sequence
from datetime import datetime, timezone, tzinfo
from typing import Any, Final, TypedDict, cast
from typing import Any, Final, Protocol, cast
import httpx
from pydantic import BaseModel, Field
from typing_extensions import ReadOnly, TypedDict
import litellm
from litellm._logging import verbose_logger
@ -35,6 +36,34 @@ GALILEO_CLOUD_API_BASE_URL: Final = "https://api.galileo.ai"
GALILEO_MAX_IN_MEMORY_RECORDS: Final = 1000
class _GalileoLoginBody(TypedDict):
"""Decoded body of the Galileo login response."""
access_token: ReadOnly[str]
class _GalileoLoginResponse(Protocol):
"""The login call's HTTP response, read for the access token it carries."""
def json(self) -> _GalileoLoginBody: ...
class _JsonResponse(Protocol):
"""An HTTP response read only for whatever JSON body it decodes to."""
def json(self) -> object: ...
def _login_access_token(response: _GalileoLoginResponse) -> str:
"""Read the bearer token out of a Galileo login response body."""
return response.json()["access_token"]
def _decoded_body(response: _JsonResponse) -> object:
"""Decode a response body without asserting anything about its shape."""
return response.json()
class GalileoStandardLoggingFields(TypedDict, total=False):
call_type: str
model: str
@ -156,7 +185,7 @@ class GalileoObserve(CustomLogger):
},
)
galileo_login_response.raise_for_status()
access_token: Final = galileo_login_response.json()["access_token"]
access_token: Final = _login_access_token(galileo_login_response)
self.headers = {
"accept": "application/json",
"Content-Type": "application/json",
@ -421,7 +450,7 @@ class GalileoObserve(CustomLogger):
try:
verbose_logger.debug(
"Galileo Logger HTTP error response json: %s",
response.json(),
_decoded_body(response),
)
except Exception:
pass

88
litellm/integrations/gitlab/gitlab_client.py
View file

@ -4,12 +4,80 @@ Now supports selecting a tag via `config["tag"]`; falls back to branch ("main").
"""
import base64
from typing import Any, Final
from collections.abc import Mapping, Sequence
from typing import Any, Final, Protocol, TypedDict
from urllib.parse import quote
from typing_extensions import ReadOnly
from litellm.llms.custom_httpx.http_handler import HTTPHandler
class GitLabFilePayload(TypedDict, total=False):
"""A repository-files API entry."""
content: ReadOnly[str]
encoding: ReadOnly[str]
class GitLabTreeEntry(TypedDict, total=False):
"""A repository-tree API entry."""
path: ReadOnly[str]
type: ReadOnly[str]
class GitLabBranch(TypedDict, total=False):
"""A repository-branches API entry."""
name: ReadOnly[str]
type: ReadOnly[str]
class GitLabFileMetadata(TypedDict):
"""The response headers a raw file request exposes as metadata."""
content_type: ReadOnly[str | None]
content_length: ReadOnly[str | None]
last_modified: ReadOnly[str | None]
class _FileJsonResponse(Protocol):
def json(self) -> GitLabFilePayload: ...
class _TreeJsonResponse(Protocol):
def json(self) -> Sequence[GitLabTreeEntry] | None: ...
class _ProjectJsonResponse(Protocol):
def json(self) -> Mapping[str, object]: ...
class _BranchesJsonResponse(Protocol):
def json(self) -> Sequence[GitLabBranch] | None: ...
def _file_payload(resp: _FileJsonResponse) -> GitLabFilePayload:
"""The JSON body of a repository-files response."""
return resp.json()
def _tree_entries(resp: _TreeJsonResponse) -> Sequence[GitLabTreeEntry]:
"""The entries of a repository-tree response."""
return resp.json() or []
def _project_info(resp: _ProjectJsonResponse) -> Mapping[str, object]:
"""The JSON body of a project response."""
return resp.json()
def _branch_entries(resp: _BranchesJsonResponse) -> Sequence[GitLabBranch] | None:
"""The JSON body of a repository-branches response."""
return resp.json()
class GitLabClient:
"""
Client for interacting with the GitLab API to fetch files.
@ -42,12 +110,12 @@ class GitLabClient:
self.project: str | int = project
self.access_token: str = str(access_token)
self.auth_method = config.get("auth_method", "token") # 'token' or 'oauth'
self.auth_method: str = config.get("auth_method", "token") # 'token' or 'oauth'
self.branch = config.get("branch", None)
if not self.branch:
self.branch = "main"
self.tag = config.get("tag")
self.base_url = config.get("base_url", "https://gitlab.com/api/v4")
self.base_url: str = config.get("base_url", "https://gitlab.com/api/v4")
if not all([self.project, self.access_token]):
raise ValueError("project and access_token are required")
@ -159,7 +227,7 @@ class GitLabClient:
if resp.status_code == 404:
return None
resp.raise_for_status()
data: Final = resp.json()
data: Final = _file_payload(resp)
content: Final = data.get("content")
encoding: Final = data.get("encoding", "")
if content and encoding == "base64":
@ -208,7 +276,7 @@ class GitLabClient:
return []
resp.raise_for_status()
data: Final = resp.json() or []
data: Final = _tree_entries(resp)
files: Final[list[str]] = []
for item in data:
if item.get("type") == "blob":
@ -229,13 +297,13 @@ class GitLabClient:
raise Exception("Authentication failed. Check your GitLab token and auth_method.")
raise Exception(f"Failed to list files in '{directory_path}': {e}")
def get_repository_info(self) -> dict[str, Any]:
def get_repository_info(self) -> Mapping[str, object]:
"""Get information about the project/repository."""
url: Final = f"{self.base_url}/projects/{self._project_enc}"
try:
resp: Final = self.http_handler.get(url, headers=self.headers)
resp.raise_for_status()
return resp.json()
return _project_info(resp)
except Exception as e:
raise Exception(f"Failed to get repository info: {e}")
@ -247,18 +315,18 @@ class GitLabClient:
except Exception:
return False
def get_branches(self) -> list[dict[str, Any]]:
def get_branches(self) -> list[GitLabBranch]:
"""Get list of branches in the repository."""
url: Final = f"{self.base_url}/projects/{self._project_enc}/repository/branches"
try:
resp: Final = self.http_handler.get(url, headers=self.headers)
resp.raise_for_status()
data: Final = resp.json()
data: Final = _branch_entries(resp)
return data if isinstance(data, list) else []
except Exception as e:
raise Exception(f"Failed to get branches: {e}")
def get_file_metadata(self, file_path: str, *, ref: str | None = None) -> dict[str, Any] | None:
def get_file_metadata(self, file_path: str, *, ref: str | None = None) -> GitLabFileMetadata | None:
"""
Get minimal metadata about a file via RAW endpoint headers at a given ref.

59
litellm/integrations/gitlab/gitlab_prompt_manager.py
View file

@ -2,10 +2,12 @@
GitLab prompt manager with configurable prompts folder.
"""
from typing import TYPE_CHECKING, Any, Final
from collections.abc import Mapping
from typing import TYPE_CHECKING, Any, Final, TypeVar
from jinja2 import DictLoader, select_autoescape
from jinja2.sandbox import ImmutableSandboxedEnvironment
from typing_extensions import ReadOnly, TypedDict
from litellm.integrations.custom_prompt_management import CustomPromptManagement
@ -24,6 +26,19 @@ from litellm.types.utils import StandardCallbackDynamicParams
GITLAB_PREFIX: Final = "gitlab::"
_ResponseT = TypeVar("_ResponseT")
class GitLabCachedPrompt(TypedDict):
id: ReadOnly[str]
path: ReadOnly[str]
content: ReadOnly[str]
metadata: ReadOnly[Mapping[str, object]]
model: ReadOnly[str | None]
temperature: ReadOnly[float | None]
max_tokens: ReadOnly[int | None]
optional_params: ReadOnly[Mapping[str, object]]
def encode_prompt_id(raw_id: str) -> str:
"""Convert GitLab path IDs like 'invoice/extract' → 'gitlab::invoice::extract'"""
@ -206,7 +221,7 @@ class GitLabTemplateManager:
result[key] = value.strip("\"'")
return result
def render_template(self, template_id: str, variables: dict[str, Any] | None = None) -> str:
def render_template(self, template_id: str, variables: Mapping[str, object] | None = None) -> str:
if template_id not in self.prompts:
raise ValueError(f"Template '{template_id}' not found")
template: Final = self.prompts[template_id]
@ -313,7 +328,7 @@ class GitLabPromptManager(CustomPromptManagement):
def get_prompt_template(
self,
prompt_id: str,
prompt_variables: dict[str, Any] | None = None,
prompt_variables: Mapping[str, object] | None = None,
*,
ref: str | None = None,
) -> tuple[str, dict[str, Any]]:
@ -338,13 +353,13 @@ class GitLabPromptManager(CustomPromptManagement):
self,
user_id: str | None,
messages: list[AllMessageValues],
function_call: dict[str, Any] | str | None = None,
litellm_params: dict[str, Any] | None = None,
function_call: Mapping[str, object] | str | None = None,
litellm_params: dict[str, object] | None = None,
prompt_id: str | None = None,
prompt_variables: dict[str, Any] | None = None,
prompt_variables: Mapping[str, object] | None = None,
prompt_version: str | None = None,
**kwargs,
) -> tuple[list[AllMessageValues], dict[str, Any] | None]:
) -> tuple[list[AllMessageValues], dict[str, object] | None]:
if not prompt_id:
return messages, litellm_params
try:
@ -377,9 +392,9 @@ class GitLabPromptManager(CustomPromptManagement):
return final_messages, litellm_params
except Exception as e:
import litellm
from litellm._logging import verbose_proxy_logger
litellm._logging.verbose_proxy_logger.error("Error in GitLab prompt pre_call_hook: %s", e)
verbose_proxy_logger.error("Error in GitLab prompt pre_call_hook: %s", e)
return messages, litellm_params
def _parse_prompt_to_messages(self, prompt_content: str) -> list[AllMessageValues]:
@ -435,14 +450,14 @@ class GitLabPromptManager(CustomPromptManagement):
def post_call_hook(
self,
user_id: str | None,
response: Any,
response: _ResponseT,
input_messages: list[AllMessageValues],
function_call: dict[str, Any] | str | None = None,
litellm_params: dict[str, Any] | None = None,
function_call: Mapping[str, object] | str | None = None,
litellm_params: Mapping[str, object] | None = None,
prompt_id: str | None = None,
prompt_variables: dict[str, Any] | None = None,
prompt_variables: Mapping[str, object] | None = None,
**kwargs,
) -> Any:
) -> _ResponseT:
return response
def get_available_prompts(self) -> list[str]:
@ -498,7 +513,7 @@ class GitLabPromptManager(CustomPromptManagement):
messages: Final = self._parse_prompt_to_messages(rendered_prompt)
template_model: Final = prompt_metadata.get("model")
optional_params: Final[dict[str, Any]] = {}
optional_params: Final[dict[str, object]] = {}
for param in [
"temperature",
"max_tokens",
@ -658,14 +673,14 @@ class GitLabPromptCache:
self.template_manager: GitLabTemplateManager = self.prompt_manager.prompt_manager
# In-memory stores
self._by_file: dict[str, dict[str, Any]] = {}
self._by_id: dict[str, dict[str, Any]] = {}
self._by_file: dict[str, GitLabCachedPrompt] = {}
self._by_id: dict[str, GitLabCachedPrompt] = {}
# -------------------------
# Public API
# -------------------------
def load_all(self, *, recursive: bool = True) -> dict[str, dict[str, Any]]:
def load_all(self, *, recursive: bool = True) -> dict[str, GitLabCachedPrompt]:
"""
Scan GitLab for all .prompt files under prompts_path, load and parse each,
and return the mapping of repo file path -> JSON-like dict.
@ -695,7 +710,7 @@ class GitLabPromptCache:
return self._by_id
def reload(self, *, recursive: bool = True) -> dict[str, dict[str, Any]]:
def reload(self, *, recursive: bool = True) -> dict[str, GitLabCachedPrompt]:
"""Clear the cache and re-load from GitLab."""
self._by_file.clear()
self._by_id.clear()
@ -709,11 +724,11 @@ class GitLabPromptCache:
"""Return the template IDs (relative to prompts_path, without extension) currently cached."""
return list(self._by_id.keys())
def get_by_file(self, file_path: str) -> dict[str, Any] | None:
def get_by_file(self, file_path: str) -> GitLabCachedPrompt | None:
"""Get a cached prompt JSON by repo file path."""
return self._by_file.get(file_path)
def get_by_id(self, prompt_id: str) -> dict[str, Any] | None:
def get_by_id(self, prompt_id: str) -> GitLabCachedPrompt | None:
"""Get a cached prompt JSON by prompt ID (relative to prompts_path)."""
if prompt_id in self._by_id:
return self._by_id[prompt_id]
@ -728,7 +743,7 @@ class GitLabPromptCache:
# Internals
# -------------------------
def _template_to_json(self, prompt_id: str, tmpl: GitLabPromptTemplate) -> dict[str, Any]:
def _template_to_json(self, prompt_id: str, tmpl: GitLabPromptTemplate) -> GitLabCachedPrompt:
"""
Normalize a GitLabPromptTemplate into a JSON-like dict that is easy to serialize.
"""

39
litellm/integrations/langfuse/langfuse.py
View file

@ -89,7 +89,7 @@ def _extract_cache_read_input_tokens(usage_obj) -> int:
# Check prompt_tokens_details.cached_tokens (used by Gemini and other providers)
if hasattr(usage_obj, "prompt_tokens_details"):
prompt_tokens_details: Final = getattr(usage_obj, "prompt_tokens_details", None)
prompt_tokens_details: Final[object] = getattr(usage_obj, "prompt_tokens_details", None)
if prompt_tokens_details is not None and hasattr(prompt_tokens_details, "cached_tokens"):
cached_tokens: Final = getattr(prompt_tokens_details, "cached_tokens", None)
if cached_tokens is not None and isinstance(cached_tokens, (int, float)) and cached_tokens > 0:
@ -623,9 +623,16 @@ class LangFuseLogger:
)
# Apply custom masking function if provided
if masking_function is not None and callable(masking_function):
input = self._apply_masking_function(input, masking_function)
output = self._apply_masking_function(output, masking_function)
masked_input: Final[object] = (
self._apply_masking_function(input, masking_function)
if masking_function is not None and callable(masking_function)
else input
)
masked_output: Final[object] = (
self._apply_masking_function(output, masking_function)
if masking_function is not None and callable(masking_function)
else output
)
clean_metadata = redact_user_api_key_info(metadata=clean_metadata)
@ -651,15 +658,15 @@ class LangFuseLogger:
# Special keys that are found in the function arguments and not the metadata
if "input" in update_trace_keys:
trace_params["input"] = input if not mask_input else "redacted-by-litellm"
trace_params["input"] = masked_input if not mask_input else "redacted-by-litellm"
if "output" in update_trace_keys:
trace_params["output"] = output if not mask_output else "redacted-by-litellm"
trace_params["output"] = masked_output if not mask_output else "redacted-by-litellm"
else: # don't overwrite an existing trace
trace_params = {
"id": trace_id,
"name": trace_name,
"session_id": session_id,
"input": input if not mask_input else "redacted-by-litellm",
"input": masked_input if not mask_input else "redacted-by-litellm",
"version": clean_metadata.pop(
"trace_version", clean_metadata.get("version", None)
), # If provided just version, it will applied to the trace as well, if applied a trace version it will take precedence
@ -669,9 +676,9 @@ class LangFuseLogger:
trace_params[key.replace("trace_", "")] = clean_metadata.pop(key, None)
if level == "ERROR":
trace_params["status_message"] = output
trace_params["status_message"] = masked_output
else:
trace_params["output"] = output if not mask_output else "redacted-by-litellm"
trace_params["output"] = masked_output if not mask_output else "redacted-by-litellm"
if debug is True or (isinstance(debug, str) and debug.lower() == "true"):
debug_metadata: Final = {
@ -708,7 +715,7 @@ class LangFuseLogger:
("aws_region_name", aws_region_name, bool(aws_region_name)),
("cache_hit", kwargs.get("cache_hit") or False, self._supports_tags() and "cache_hit" in kwargs),
)
enrichments: Final[Mapping[str, Any]] = {
enrichments: Final[Mapping[str, object]] = {
key: value for key, value, include in candidate_enrichments if include
}
@ -802,8 +809,8 @@ class LangFuseLogger:
"end_time": end_time,
"model": model_name,
"model_parameters": optional_params,
"input": input if not mask_input else "redacted-by-litellm",
"output": output if not mask_output else "redacted-by-litellm",
"input": masked_input if not mask_input else "redacted-by-litellm",
"output": masked_output if not mask_output else "redacted-by-litellm",
"usage": usage,
"usage_details": usage_details,
"metadata": {
@ -825,8 +832,8 @@ class LangFuseLogger:
prompt_management_metadata=prompt_management_metadata,
langfuse_client=self.Langfuse,
)
if output is not None and isinstance(output, str) and level == "ERROR":
generation_params["status_message"] = output
if masked_output is not None and isinstance(masked_output, str) and level == "ERROR":
generation_params["status_message"] = masked_output
if self._supports_completion_start_time():
generation_params["completion_start_time"] = kwargs.get("completion_start_time", None)
@ -935,7 +942,7 @@ class LangFuseLogger:
return Version(self.langfuse_sdk_version) >= Version("2.7.3")
@staticmethod
def _apply_masking_function(data: Any, masking_function: Callable[[Any], Any]) -> Any:
def _apply_masking_function(data: object, masking_function: Callable[[object], object]) -> object:
"""
Apply a masking function to data, handling different data types.
@ -1049,7 +1056,7 @@ def _add_prompt_to_generation_params(
generation_params: dict,
clean_metadata: dict,
prompt_management_metadata: StandardLoggingPromptManagementMetadata | None,
langfuse_client: Any,
langfuse_client: object,
) -> dict:
from langfuse import Langfuse
from langfuse.model import (

27
litellm/integrations/opik/opik.py
View file

@ -4,9 +4,12 @@ Opik Logger that logs LLM events to an Opik server
import asyncio
import traceback
from collections.abc import Mapping
from datetime import datetime
from typing import Any, Final
from typing_extensions import ReadOnly, TypedDict, Unpack
from litellm._logging import verbose_logger
from litellm.integrations.custom_batch_logger import CustomBatchLogger
from litellm.llms.custom_httpx.http_handler import (
@ -23,7 +26,7 @@ except Exception:
opik_client = None
def _should_skip_event(kwargs: dict[str, Any]) -> bool:
def _should_skip_event(kwargs: Mapping[str, object]) -> bool:
"""Check if event should be skipped due to missing standard_logging_object."""
if kwargs.get("standard_logging_object") is None:
verbose_logger.debug("OpikLogger skipping event; no standard_logging_object found")
@ -31,12 +34,24 @@ def _should_skip_event(kwargs: dict[str, Any]) -> bool:
return False
class _OpikLoggerKwargs(TypedDict, total=False):
"""Constructor options accepted by ``OpikLogger``."""
project_name: ReadOnly[str | None]
url: ReadOnly[str | None]
api_key: ReadOnly[str | None]
workspace: ReadOnly[str | None]
batch_size: ReadOnly[int | None]
flush_interval: ReadOnly[int | None]
max_queue_size: ReadOnly[int | None]
class OpikLogger(CustomBatchLogger):
"""
Opik Logger for logging events to an Opik Server
"""
def __init__(self, **kwargs: Any) -> None:
def __init__(self, **kwargs: Unpack[_OpikLoggerKwargs]) -> None:
self.async_httpx_client = get_async_httpx_client(llm_provider=httpxSpecialProvider.LoggingCallback)
self.sync_httpx_client = _get_httpx_client()
@ -95,7 +110,7 @@ class OpikLogger(CustomBatchLogger):
async def async_log_success_event(
self,
kwargs: dict[str, Any],
kwargs: dict[str, object],
response_obj: Any,
start_time: datetime,
end_time: datetime,
@ -163,7 +178,7 @@ class OpikLogger(CustomBatchLogger):
except Exception as e:
verbose_logger.exception("OpikLogger failed to log success event - %s\n%s", e, traceback.format_exc())
def _sync_send(self, url: str, headers: dict[str, str], batch: dict[str, Any]) -> None:
def _sync_send(self, url: str, headers: dict[str, str], batch: dict[str, object]) -> None:
try:
response: Final = self.sync_httpx_client.post(
url=url,
@ -178,7 +193,7 @@ class OpikLogger(CustomBatchLogger):
def log_success_event(
self,
kwargs: dict[str, Any],
kwargs: dict[str, object],
response_obj: Any,
start_time: datetime,
end_time: datetime,
@ -247,7 +262,7 @@ class OpikLogger(CustomBatchLogger):
except Exception as e:
verbose_logger.exception("OpikLogger failed to log success event - %s\n%s", e, traceback.format_exc())
async def _submit_batch(self, url: str, headers: dict[str, str], batch: dict[str, Any]) -> None:
async def _submit_batch(self, url: str, headers: dict[str, str], batch: dict[str, object]) -> None:
try:
response: Final = await self.async_httpx_client.post(
url=url,

21
litellm/integrations/opik/opik_payload_builder/extractors.py
View file

@ -1,6 +1,7 @@
"""Data extraction functions for Opik payload building."""
import json
from collections.abc import Mapping
from typing import Any, Final
from litellm import _logging
@ -35,8 +36,8 @@ def normalize_provider_name(provider: str | None) -> str | None:
def extract_opik_metadata(
litellm_metadata: dict[str, Any],
standard_logging_metadata: dict[str, Any],
litellm_metadata: Mapping[str, Any],
standard_logging_metadata: Mapping[str, Any],
) -> dict[str, Any]:
"""
Merge Opik metadata from three sources in increasing priority order:
@ -97,7 +98,7 @@ def extract_span_identifiers(
def extract_tags(
opik_metadata: dict[str, Any],
opik_metadata: Mapping[str, Any],
custom_llm_provider: str | None,
) -> list[str]:
"""
@ -122,7 +123,7 @@ def apply_proxy_header_overrides(
project_name: str,
tags: list[str],
thread_id: str | None,
proxy_headers: dict[str, Any],
proxy_headers: Mapping[str, str],
) -> tuple[str, list[str], str | None]:
"""
Apply overrides from proxy request headers (opik_* prefix).
@ -148,7 +149,7 @@ def apply_proxy_header_overrides(
thread_id = value
elif param_key == "tags":
try:
parsed_tags = json.loads(value)
parsed_tags: object = json.loads(value)
if isinstance(parsed_tags, list):
tags.extend(parsed_tags)
except (json.JSONDecodeError, TypeError):
@ -158,11 +159,11 @@ def apply_proxy_header_overrides(
def extract_and_build_metadata(
opik_metadata: dict[str, Any],
standard_logging_metadata: dict[str, Any],
standard_logging_object: dict[str, Any],
litellm_kwargs: dict[str, Any],
) -> dict[str, Any]:
opik_metadata: Mapping[str, object],
standard_logging_metadata: Mapping[str, object],
standard_logging_object: Mapping[str, object],
litellm_kwargs: Mapping[str, object],
) -> dict[str, object]:
"""
Build the complete metadata dictionary from all available sources.

12
litellm/integrations/opik/opik_payload_builder/payload_builders.py
View file

@ -17,12 +17,12 @@ def build_trace_payload(
end_time: datetime,
input_data: Any,
output_data: Any,
metadata: dict[str, Any],
metadata: dict[str, object],
tags: list[str],
thread_id: str | None,
) -> types.TracePayload:
"""Build a complete trace payload."""
trace_name: Final = response_obj.get("object", "unknown type")
trace_name: Final[str] = response_obj.get("object", "unknown type")
return types.TracePayload(
project_name=project_name,
@ -47,7 +47,7 @@ def build_span_payload(
end_time: datetime,
input_data: Any,
output_data: Any,
metadata: dict[str, Any],
metadata: dict[str, object],
tags: list[str],
usage: dict[str, int],
provider: str | None = None,
@ -56,9 +56,9 @@ def build_span_payload(
"""Build a complete span payload."""
span_id: Final = utils.create_uuid7()
model: Final = response_obj.get("model", "unknown-model")
obj_type: Final = response_obj.get("object", "unknown-object")
created: Final = response_obj.get("created", 0)
model: Final[str] = response_obj.get("model", "unknown-model")
obj_type: Final[str] = response_obj.get("object", "unknown-object")
created: Final[int] = response_obj.get("created", 0)
span_name: Final = f"{model}_{obj_type}_{created}"
_logging.verbose_logger.debug("OpikLogger creating span with id %s for trace %s", span_id, trace_id)

2
litellm/integrations/otel/mappers/genai.py
View file

@ -62,6 +62,8 @@ class GenAIMapper:
GenAI.RESPONSE_TIME_TO_FIRST_CHUNK: lambda d: d.time_to_first_chunk_seconds,
GenAI.USAGE_INPUT_TOKENS: lambda d: d.usage.input_tokens,
GenAI.USAGE_OUTPUT_TOKENS: lambda d: d.usage.output_tokens,
GenAI.USAGE_CACHE_CREATION_INPUT_TOKENS: lambda d: d.usage.cache_creation_input_tokens,
GenAI.USAGE_CACHE_READ_INPUT_TOKENS: lambda d: d.usage.cache_read_input_tokens,
Error.TYPE: lambda d: d.error.error_type if d.error else None,
Server.ADDRESS: lambda d: d.server.address if d.server else None,
Server.PORT: lambda d: d.server.port if d.server else None,

22
litellm/integrations/otel/model/payloads.py
View file

@ -95,6 +95,22 @@ class LLMUsage:
input_tokens: int | None = None
output_tokens: int | None = None
total_tokens: int | None = None
cache_creation_input_tokens: int | None = None
cache_read_input_tokens: int | None = None
@classmethod
def from_standard_logging_payload(cls, payload: StandardLoggingPayload) -> LLMUsage:
# Cache token counts only exist on the raw provider usage object under metadata
metadata: Final[Mapping[str, object]] = payload.get("metadata") or {}
raw_usage: Final = metadata.get("usage_object")
usage_object: Final[Mapping[str, object]] = raw_usage if isinstance(raw_usage, Mapping) else {}
return cls(
input_tokens=as_int(payload.get("prompt_tokens")),
output_tokens=as_int(payload.get("completion_tokens")),
total_tokens=as_int(payload.get("total_tokens")),
cache_creation_input_tokens=as_int(usage_object.get("cache_creation_input_tokens")),
cache_read_input_tokens=as_int(usage_object.get("cache_read_input_tokens")),
)
@dataclass(frozen=True)
@ -363,11 +379,7 @@ class LLMCallSpanData:
response_model=context.response_model,
response_id=as_str(response.get("id")),
request_params=LLMRequestParams.from_model_parameters(params),
usage=LLMUsage(
input_tokens=as_int(payload.get("prompt_tokens")),
output_tokens=as_int(payload.get("completion_tokens")),
total_tokens=as_int(payload.get("total_tokens")),
),
usage=LLMUsage.from_standard_logging_payload(payload),
finish_reasons=finish_reasons,
error=_parse_error(payload),
response_cost=as_float(payload.get("response_cost")),

2
litellm/integrations/otel/model/semconv.py
View file

@ -110,6 +110,8 @@ class GenAI:
# usage
USAGE_INPUT_TOKENS: Final = "gen_ai.usage.input_tokens"
USAGE_OUTPUT_TOKENS: Final = "gen_ai.usage.output_tokens"
USAGE_CACHE_CREATION_INPUT_TOKENS: Final = "gen_ai.usage.cache_creation.input_tokens"
USAGE_CACHE_READ_INPUT_TOKENS: Final = "gen_ai.usage.cache_read.input_tokens"
# content (opt-in, gated by capture mode)
INPUT_MESSAGES: Final = "gen_ai.input.messages"
OUTPUT_MESSAGES: Final = "gen_ai.output.messages"

45
litellm/integrations/otel/plumbing/metrics.py
View file

@ -11,9 +11,10 @@ identical metrics. The attribute cardinality filter is reused from v1 by import
from collections.abc import Mapping
from dataclasses import dataclass
from datetime import datetime
from typing import Any, Final, TypeAlias
from typing import Any, Final, Literal, Protocol, TypeAlias
from opentelemetry.metrics import Histogram, Meter
from typing_extensions import ReadOnly, TypedDict
import litellm
from litellm._logging import verbose_logger
@ -151,6 +152,29 @@ METRIC_ATTRIBUTE_CEILING: Final[frozenset[str]] = frozenset(
BOUNDED_HIDDEN_PARAM_KEYS: Final[tuple[str, ...]] = ("model_id",)
class _TokenUsage(TypedDict, total=False):
"""The token counts a response's ``usage`` carries, as the recorder reads them."""
prompt_tokens: ReadOnly[int]
completion_tokens: ReadOnly[int]
class _ResponseView(Protocol):
"""The one read the recorder makes on a litellm response object."""
def get(self, key: Literal["usage"], /) -> _TokenUsage | None: ...
class _MetricKwargs(TypedDict, total=False):
"""The logging kwargs the recorder reads directly."""
call_type: ReadOnly[str | None]
litellm_params: ReadOnly[Mapping[str, object] | None]
response_cost: ReadOnly[float | None]
completion_start_time: ReadOnly[datetime | float | str | None]
api_call_start_time: ReadOnly[datetime | float | str | None]
def resolve_error_type(kwargs: Mapping[str, Any]) -> str:
"""The ``error.type`` value for a failed request.
@ -192,8 +216,8 @@ class GenAIMetricRecorder:
def record(
self,
kwargs: Mapping[str, Any],
response_obj: Any,
kwargs: _MetricKwargs,
response_obj: _ResponseView | None,
start_time: datetime,
end_time: datetime,
) -> None:
@ -218,7 +242,7 @@ class GenAIMetricRecorder:
def record_failure(
self,
kwargs: Mapping[str, Any],
kwargs: _MetricKwargs,
start_time: datetime,
end_time: datetime,
) -> None:
@ -342,7 +366,7 @@ class GenAIMetricRecorder:
# Per-metric recording
# ------------------------------------------------------------------ #
def _record_token_usage(self, response_obj: Any, common_attrs: dict) -> None:
def _record_token_usage(self, response_obj: _ResponseView | None, common_attrs: dict) -> None:
if not response_obj:
return
usage: Final = response_obj.get("usage")
@ -353,7 +377,7 @@ class GenAIMetricRecorder:
self._metrics.token_usage.record(usage.get("prompt_tokens", 0), attributes=in_attrs)
self._metrics.token_usage.record(usage.get("completion_tokens", 0), attributes=out_attrs)
def _record_time_to_first_token(self, kwargs: Mapping[str, Any], common_attrs: dict) -> None:
def _record_time_to_first_token(self, kwargs: _MetricKwargs, common_attrs: dict) -> None:
time_to_first_chunk: Final = time_to_first_chunk_seconds(kwargs)
if time_to_first_chunk is None:
return
@ -361,15 +385,14 @@ class GenAIMetricRecorder:
def _record_time_per_output_token(
self,
kwargs: Mapping[str, Any],
response_obj: Any,
kwargs: _MetricKwargs,
response_obj: _ResponseView | None,
end_time: datetime,
duration_s: float,
common_attrs: dict,
) -> None:
completion_tokens = None
if response_obj and (usage := response_obj.get("usage")):
completion_tokens = usage.get("completion_tokens")
usage: Final = response_obj.get("usage") if response_obj else None
completion_tokens: Final = usage.get("completion_tokens") if usage else None
if completion_tokens is None or completion_tokens <= 0:
return

30
litellm/integrations/otel/plumbing/providers.py
View file

@ -1,7 +1,7 @@
"""Provider / exporter factory + the Baggage span processor."""
from collections.abc import Callable, Iterable
from typing import TYPE_CHECKING, Any, Final
from typing import TYPE_CHECKING, Any, Final, Literal
from opentelemetry import _logs, baggage, metrics
from opentelemetry._events import EventLogger
@ -135,14 +135,36 @@ def parse_headers(raw: str | None) -> dict[str, str]:
return dict(parse_env_headers(raw, liberal=True))
_IN_MEMORY_KINDS: Final = ("in_memory", "inmemory", "memory")
_OTLP_HTTP_KINDS: Final = ("otlp_http", "http", "http/protobuf", "http/json")
_OTLP_GRPC_KINDS: Final = ("otlp_grpc", "grpc")
def exporter_transport(kind: str) -> Literal["http", "grpc", "headerless"]:
"""How an exporter of this ``kind`` carries credentials, per ``_exporter_from_spec``.
``http``/``grpc`` exporters (and any registered factory, which builds an
OTLP exporter) stamp ``spec.headers``; ``console``, ``in_memory``, and any
unrecognized kind (which falls back to a header-ignoring console exporter)
are ``headerless``. Routability decisions must read this rather than a
denylist, so a typo'd or unavailable kind is not mistaken for OTLP.
"""
resolved: Final = kind.lower()
if resolved in _OTLP_HTTP_KINDS or resolved in _EXPORTER_FACTORIES:
return "http"
if resolved in _OTLP_GRPC_KINDS:
return "grpc"
return "headerless"
def _exporter_from_spec(spec: ExporterSpec) -> SpanExporter:
kind: Final = (spec.kind or "console").lower()
factory: Final = _EXPORTER_FACTORIES.get(kind)
if factory is not None:
return factory(spec)
if kind in ("in_memory", "inmemory", "memory"):
if kind in _IN_MEMORY_KINDS:
return InMemorySpanExporter()
if kind in ("otlp_http", "http", "http/protobuf", "http/json"):
if kind in _OTLP_HTTP_KINDS:
from opentelemetry.exporter.otlp.proto.http.trace_exporter import (
OTLPSpanExporter as HTTPExporter,
)
@ -151,7 +173,7 @@ def _exporter_from_spec(spec: ExporterSpec) -> SpanExporter:
endpoint=_otlp_traces_endpoint(spec.endpoint),
headers=parse_headers(spec.headers),
)
if kind in ("otlp_grpc", "grpc"):
if kind in _OTLP_GRPC_KINDS:
from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import (
OTLPSpanExporter as GRPCExporter,
)

62
litellm/integrations/otel/plumbing/routing.py
View file

@ -27,6 +27,7 @@ from litellm.constants import OTEL_SERVICE_NAME_METADATA_KEYS
from litellm.integrations.otel.model.config import ExporterSpec, OpenTelemetryV2Config
from litellm.integrations.otel.plumbing.providers import (
build_tracer_provider,
exporter_transport,
get_tracer,
)
from litellm.integrations.otel.presets import (
@ -121,13 +122,27 @@ def _encoded_header_string(headers: Mapping[str, str]) -> str:
class TenantRoute:
"""The tracer to create a span on, plus whether it must root its own trace.
``detached`` is True when project routing engaged. Phoenix assigns a whole
``detached`` is True when the routed span exports to a DIFFERENT backend
than the request's root span, which always exports through the default
tracer. A detached span roots a fresh trace with a link back to the request
trace for correlation, so the destination account is not left holding a
child whose parent it never received. It is driven by whether routing
headers were actually applied to an owned exporter, not merely requested:
a credential or project route whose callback owns no exporter those headers
can reach exports through the default backend unchanged, so it stays
parented like an unrouted span.
Credential routing (a team/key's own vendor account) is one detaching case:
the root, auth, and db spans stay on the operator's default backend while
the LLM-call span exports to the tenant's account, so parenting it into the
request trace makes the tenant account show a fragmented span with a missing
parent. Project routing (Phoenix) is the other: Phoenix assigns a whole
trace to one project by whichever of its spans arrives first, so a
project-routed span parented into the request trace gets dragged into the
project of the default-exported request spans and the header does nothing.
The span must therefore start a fresh trace (with a link back to the
request trace for correlation) — which is also how the v1 Phoenix logger
behaved, exporting each request under its own Phoenix-local parent span.
Both mirror the v1 loggers, which exported each request under its own
backend-local root. Service-name routing does NOT detach: it relabels
``service.name`` on the SAME operator backend, where the parent is present.
"""
tracer: Tracer
@ -161,11 +176,20 @@ class TenantTracerCache:
self._open_span_counts: dict[TracerProvider, int] = {} # mutable-ok: live refcount state
# Oldest-first so an overflow of draining providers sheds the stalest.
self._retired: OrderedDict[TracerProvider, None] = OrderedDict() # mutable-ok: draining evicted providers
self._project_routable = any(
spec.owner == callback_name and spec.kind.lower() not in (*_NON_OTLP_KINDS, *_GRPC_KINDS)
for spec in config.exporters
# An owned exporter is routable only when its kind actually resolves to a
# header-carrying OTLP exporter. A denylist would accept a typo'd or
# unavailable kind, which ``_exporter_from_spec`` falls back to a
# header-ignoring console exporter: detaching such a span would strand it
# on the operator's console, never reaching the tenant backend. Project
# headers are HTTP-only; credentials ride gRPC metadata too (Arize's
# default exporter is gRPC), so they accept either OTLP transport.
owned_transports: Final = tuple(
exporter_transport(spec.kind) for spec in config.exporters if spec.owner == callback_name
)
self._project_routable = "http" in owned_transports
self._credential_routable = "http" in owned_transports or "grpc" in owned_transports
self._warned_project_unroutable = False
self._warned_credential_unroutable = False
def release(self, provider: TracerProvider | None) -> None:
"""Drop one open-span count; shut a retired provider down once drained.
@ -207,7 +231,7 @@ class TenantTracerCache:
concurrent overflow eviction can't shut it down between selection and
the caller's span start. The caller must ``release`` it exactly once.
"""
credential_headers: Final = dynamic_otlp_headers(self._callback_name, dynamic_params) or _NO_HEADERS
credential_headers: Final = self._credential_headers(dynamic_params)
project_headers: Final = self._project_headers(auth_metadata)
service_name: Final = tenant_service_name(auth_metadata)
if not credential_headers and not project_headers and service_name is None:
@ -231,7 +255,7 @@ class TenantTracerCache:
_shutdown_provider(evicted)
return TenantRoute(
tracer=get_tracer(provider, self._tracer_name),
detached=bool(project_headers),
detached=bool(project_headers) or bool(credential_headers),
provider=provider,
)
@ -275,6 +299,26 @@ class TenantTracerCache:
self._open_span_counts.pop(overflowed, None)
return overflowed
def _credential_headers(self, dynamic_params: StandardCallbackDynamicParams | None) -> Mapping[str, str]:
"""The per-request dynamic OTLP credentials, if this cache can apply them.
A callback owning only a console/in_memory exporter has nowhere to stamp
them, so the span would export to the operator's default backend
unchanged; routing there and detaching would orphan it on the very
backend that holds its parent. Warn once and keep the default tracer.
"""
requested: Final = dynamic_otlp_headers(self._callback_name, dynamic_params) or _NO_HEADERS
if not requested or self._credential_routable:
return requested
if not self._warned_credential_unroutable:
self._warned_credential_unroutable = True
verbose_logger.warning(
"OTel V2: %s request carries dynamic credentials, but the callback owns no "
"OTLP exporter to stamp them onto; spans export to the default backend.",
self._callback_name,
)
return _NO_HEADERS
def _project_headers(self, auth_metadata: Mapping[str, str] | None) -> Mapping[str, str]:
"""The per-request project-routing headers, if this cache can apply them.

62
litellm/integrations/posthog.py
View file

@ -12,7 +12,10 @@ For batching specific details see CustomBatchLogger class
import asyncio
import atexit
import os
from typing import Any, Final
from collections.abc import Mapping, Sequence
from typing import Final
from typing_extensions import ReadOnly, TypedDict
from litellm._logging import verbose_logger
from litellm._uuid import uuid
@ -34,6 +37,21 @@ from litellm.types.integrations.posthog import (
from litellm.types.utils import StandardCallbackDynamicParams, StandardLoggingPayload
class PostHogBatchPayload(TypedDict):
api_key: ReadOnly[str]
batch: ReadOnly[Sequence[PostHogEventPayload]]
class PostHogLiteLLMParams(TypedDict, total=False):
metadata: ReadOnly[Mapping[str, object]]
class PostHogLogKwargs(TypedDict, total=False):
standard_logging_object: ReadOnly[StandardLoggingPayload]
standard_callback_dynamic_params: ReadOnly[StandardCallbackDynamicParams]
litellm_params: ReadOnly[PostHogLiteLLMParams]
class PostHogLogger(CustomBatchLogger):
def __init__(self, **kwargs):
"""
@ -137,7 +155,7 @@ class PostHogLogger(CustomBatchLogger):
if len(self.log_queue) >= self.batch_size:
await self.flush_queue()
def create_posthog_event_payload(self, kwargs: dict[str, Any]) -> PostHogEventPayload:
def create_posthog_event_payload(self, kwargs: PostHogLogKwargs) -> PostHogEventPayload:
"""
Helper function to create a PostHog event payload for logging
@ -171,11 +189,11 @@ class PostHogLogger(CustomBatchLogger):
def _create_posthog_properties(
self,
standard_logging_object: StandardLoggingPayload,
kwargs: dict[str, Any],
kwargs: PostHogLogKwargs,
event_name: str,
) -> dict[str, Any]:
) -> dict[str, object]:
"""Create PostHog properties following LLM Analytics spec"""
properties: Final = {}
properties: Final[dict[str, object]] = {}
# Core model information
properties["$ai_model"] = self._safe_get(standard_logging_object, "model", "")
@ -211,16 +229,19 @@ class PostHogLogger(CustomBatchLogger):
properties["$ai_error"] = error_str
# Add trace properties
self._add_trace_properties(properties, kwargs)
self._add_trace_properties(properties, standard_logging_object, kwargs)
# Add custom metadata fields
self._add_custom_metadata_properties(properties, kwargs)
return properties
def _add_trace_properties(self, properties: dict[str, Any], kwargs: dict[str, Any]):
standard_logging_object: Final = self._safe_get(kwargs, "standard_logging_object", {})
def _add_trace_properties(
self,
properties: dict[str, object],
standard_logging_object: StandardLoggingPayload,
kwargs: PostHogLogKwargs,
) -> None:
trace_id: Final = self._safe_get(standard_logging_object, "trace_id", self._safe_uuid())
properties["$ai_trace_id"] = trace_id
@ -232,7 +253,7 @@ class PostHogLogger(CustomBatchLogger):
if parent_id:
properties["$ai_parent_id"] = parent_id
def _add_custom_metadata_properties(self, properties: dict[str, Any], kwargs: dict[str, Any]):
def _add_custom_metadata_properties(self, properties: dict[str, object], kwargs: PostHogLogKwargs) -> None:
"""Add custom metadata fields to PostHog properties"""
metadata: Final = self._extract_metadata(kwargs)
if not isinstance(metadata, dict):
@ -277,7 +298,7 @@ class PostHogLogger(CustomBatchLogger):
if key not in litellm_internal_fields:
properties[key] = value
def _get_distinct_id(self, standard_logging_object: StandardLoggingPayload, kwargs: dict[str, Any]) -> str:
def _get_distinct_id(self, standard_logging_object: StandardLoggingPayload, kwargs: PostHogLogKwargs) -> str:
metadata: Final = self._extract_metadata(kwargs)
user_id: Final = self._safe_get(metadata, "user_id")
if user_id:
@ -291,7 +312,7 @@ class PostHogLogger(CustomBatchLogger):
return self._safe_uuid()
def _get_credentials_for_request(self, kwargs: dict[str, Any]) -> tuple[str | None, str | None]:
def _get_credentials_for_request(self, kwargs: PostHogLogKwargs) -> tuple[str | None, str | None]:
"""
Get PostHog credentials for this request.
@ -334,7 +355,7 @@ class PostHogLogger(CustomBatchLogger):
verbose_logger.debug("[POSTHOG MOCK] Mock mode enabled - API calls will be intercepted")
# Group events by credentials for batch sending
batches_by_credentials: Final[dict[tuple[str, str], list]] = {}
batches_by_credentials: Final[dict[tuple[str, str], list[PostHogEventPayload]]] = {}
for item in self.log_queue:
key = (item["api_key"], item["api_url"])
if key not in batches_by_credentials:
@ -380,18 +401,19 @@ class PostHogLogger(CustomBatchLogger):
verbose_logger.error("PostHog: Failed to initialize async components: %s", e)
raise
def _extract_metadata(self, kwargs: dict[str, Any]) -> dict[str, Any]:
litellm_params: Final = kwargs.get("litellm_params", {}) or {}
return litellm_params.get("metadata", {}) or {}
def _extract_metadata(self, kwargs: PostHogLogKwargs) -> Mapping[str, object]:
litellm_params: Final[PostHogLiteLLMParams] = kwargs.get("litellm_params", {}) or {}
metadata: Final[Mapping[str, object]] = litellm_params.get("metadata", {}) or {}
return metadata
def _safe_uuid(self) -> str:
return str(uuid.uuid4())
def _create_posthog_payload(self, events: list, api_key: str) -> dict[str, Any]:
def _create_posthog_payload(self, events: Sequence[PostHogEventPayload], api_key: str) -> PostHogBatchPayload:
return {"api_key": api_key, "batch": events}
def _safe_get(self, obj: Any, key: str, default: Any = None) -> Any:
if obj is None or not hasattr(obj, "get"):
def _safe_get(self, obj: Mapping[str, object] | None, key: str, default: object = None) -> object:
if not isinstance(obj, Mapping):
return default
return obj.get(key, default)
@ -412,7 +434,7 @@ class PostHogLogger(CustomBatchLogger):
try:
# Group events by credentials (same logic as async_send_batch)
batches_by_credentials: Final[dict[tuple[str, str], list]] = {}
batches_by_credentials: Final[dict[tuple[str, str], list[PostHogEventPayload]]] = {}
for item in self.log_queue:
key = (item["api_key"], item["api_url"])
if key not in batches_by_credentials:

52
litellm/integrations/prometheus.py
View file

@ -59,6 +59,8 @@ from litellm.types.utils import (
if TYPE_CHECKING:
from apscheduler.schedulers.asyncio import AsyncIOScheduler
from prometheus_client.metrics import MetricWrapperBase
from litellm.router import Router
else:
AsyncIOScheduler = Any
@ -67,6 +69,8 @@ _TableRowT: Final = TypeVar("_TableRowT", bound=BaseModel)
_DEFAULT_BUDGET_METRICS_PER_REQUEST_TIMEOUT: Final = 5.0
UNRECOGNIZED_REQUESTED_MODEL_LABEL: Final = "other"
_NON_ENUM_METRIC_LABELS: Final[frozenset[str]] = frozenset(
(
"guardrail_name",
@ -154,6 +158,44 @@ def _get_budget_metrics_per_request_timeout() -> float:
return parsed
def _get_proxy_llm_router() -> Router | None:
try:
from litellm.proxy.proxy_server import llm_router
except Exception:
return None
return llm_router
def _bounded_requested_model_label(requested_model: str | None, router_originated: bool = False) -> str | None:
"""
Bound ``requested_model`` label cardinality: names the router recognizes
(model names, deployment ids, aliases, routing groups, team public model
names) or matches via a global or team wildcard/pattern route keep their
own label value; any other client-supplied string collapses into the
single ``other`` bucket. With no proxy router to vouch for the string,
client-supplied values collapse to ``other`` while ``router_originated``
values (emitted by an SDK ``Router``'s own deployment failure and
fallback events, where the proxy router never exists) pass through.
"""
if not requested_model:
return requested_model
llm_router: Final = _get_proxy_llm_router()
if llm_router is None:
return requested_model if router_originated else UNRECOGNIZED_REQUESTED_MODEL_LABEL
if llm_router.is_recognized_model(requested_model):
return requested_model
if requested_model in llm_router.team_public_model_names:
return requested_model
if llm_router.pattern_router.route(requested_model) is not None:
return requested_model
if any(
team_pattern_router.route(requested_model) is not None
for team_pattern_router in llm_router.team_pattern_routers.values()
):
return requested_model
return UNRECOGNIZED_REQUESTED_MODEL_LABEL
class PrometheusLogger(CustomLogger):
# Class variables or attributes
@ -2407,7 +2449,7 @@ class PrometheusLogger(CustomLogger):
team_alias=user_api_key_dict.team_alias,
org_id=user_api_key_dict.org_id,
org_alias=user_api_key_dict.organization_alias,
requested_model=request_data.get("model", ""),
requested_model=_bounded_requested_model_label(request_data.get("model", "")),
status_code=str(status_code),
exception_status=str(status_code),
exception_class=self._get_exception_class_name(original_exception),
@ -2627,7 +2669,9 @@ class PrometheusLogger(CustomLogger):
label_model_id = ""
label_api_base = ""
label_api_provider = ""
label_requested_model = litellm_model_name or model_group or ""
label_requested_model = (
_bounded_requested_model_label(litellm_model_name or model_group, router_originated=True) or ""
)
enum_values: Final = UserAPIKeyLabelValues(
litellm_model_name=label_litellm_model_name,
@ -3186,7 +3230,7 @@ class PrometheusLogger(CustomLogger):
_tags: Final = cast(list[str], kwargs.get("tags") or [])
enum_values: Final = UserAPIKeyLabelValues(
requested_model=original_model_group,
requested_model=_bounded_requested_model_label(original_model_group, router_originated=True),
fallback_model=_new_model,
hashed_api_key=standard_metadata["user_api_key_hash"],
api_key_alias=standard_metadata["user_api_key_alias"],
@ -3227,7 +3271,7 @@ class PrometheusLogger(CustomLogger):
)
enum_values: Final = UserAPIKeyLabelValues(
requested_model=original_model_group,
requested_model=_bounded_requested_model_label(original_model_group, router_originated=True),
fallback_model=_new_model,
hashed_api_key=standard_metadata["user_api_key_hash"],
api_key_alias=standard_metadata["user_api_key_alias"],

66
litellm/integrations/prometheus_helpers/prometheus_api.py
View file

@ -7,6 +7,9 @@ import time
from datetime import datetime, timedelta
from typing import Final
from pydantic import BaseModel, TypeAdapter
from typing_extensions import ReadOnly, TypedDict
from litellm import get_secret
from litellm._logging import verbose_logger
from litellm.llms.custom_httpx.http_handler import (
@ -18,10 +21,32 @@ PROMETHEUS_URL: Final[str | None] = get_secret("PROMETHEUS_URL")
PROMETHEUS_SELECTED_INSTANCE: Final[str | None] = get_secret("PROMETHEUS_SELECTED_INSTANCE")
async_http_handler: Final = get_async_httpx_client(llm_provider=httpxSpecialProvider.LoggingCallback)
_RAW_JSON_PAYLOAD: Final = TypeAdapter(object)
class PrometheusRangeSample(BaseModel):
"""One ``matrix`` series of the Prometheus HTTP query API."""
metric: dict[str, object]
values: list[tuple[float, str]]
class PrometheusQueryData(BaseModel):
result: list[PrometheusRangeSample]
class PrometheusQueryResponse(BaseModel):
data: PrometheusQueryData
class PrometheusDailySpend(TypedDict):
date: ReadOnly[str]
spend: ReadOnly[float]
async def get_metric_from_prometheus(
metric_name: str,
):
) -> list[PrometheusRangeSample]:
# Get the start of the current day in Unix timestamp
if PROMETHEUS_URL is None:
raise ValueError("PROMETHEUS_URL not set please set 'PROMETHEUS_URL=<>' in .env")
@ -31,13 +56,13 @@ async def get_metric_from_prometheus(
response: Final = await async_http_handler.get(
f"{PROMETHEUS_URL}/api/v1/query", params={"query": query, "time": now}
) # End of the day
_json_response: Final = response.json()
_json_response: Final = _RAW_JSON_PAYLOAD.validate_python(response.json())
verbose_logger.debug("json response from prometheus /query api %s", _json_response)
results: Final = response.json()["data"]["result"]
results: Final = PrometheusQueryResponse.model_validate(_json_response).data.result
return results
async def get_fallback_metric_from_prometheus():
async def get_fallback_metric_from_prometheus() -> str:
"""
Gets fallback metrics from prometheus for the last 24 hours
"""
@ -55,17 +80,17 @@ async def get_fallback_metric_from_prometheus():
verbose_logger.debug("response json %s", response_json)
for result in response_json:
verbose_logger.debug("result= %s", result)
metric = result["metric"]
metric_values = result["values"]
metric_labels = result.metric
metric_values = result.values
most_recent_value = metric_values[0]
if PROMETHEUS_SELECTED_INSTANCE is not None:
if metric.get("instance") != PROMETHEUS_SELECTED_INSTANCE:
if metric_labels.get("instance") != PROMETHEUS_SELECTED_INSTANCE:
continue
value = int(float(most_recent_value[1])) # Convert value to integer
primary_model = metric.get("primary_model", "Unknown")
fallback_model = metric.get("fallback_model", "Unknown")
primary_model = metric_labels.get("primary_model", "Unknown")
fallback_model = metric_labels.get("fallback_model", "Unknown")
response_message += f"`{value} successful fallback requests` with primary model=`{primary_model}` -> fallback model=`{fallback_model}`"
response_message += "\n"
verbose_logger.debug("response message %s", response_message)
@ -96,7 +121,7 @@ def _quote_promql_string_literal(value: str) -> str:
return json.dumps(value, ensure_ascii=False)
async def get_daily_spend_from_prometheus(api_key: str | None):
async def get_daily_spend_from_prometheus(api_key: str | None) -> list[PrometheusDailySpend]:
"""
Expected Response Format:
[
@ -133,17 +158,16 @@ async def get_daily_spend_from_prometheus(api_key: str | None):
}
response: Final = await async_http_handler.get(url, params=params)
_json_response: Final = response.json()
_json_response: Final = _RAW_JSON_PAYLOAD.validate_python(response.json())
verbose_logger.debug("json response from prometheus /query api %s", _json_response)
results: Final = response.json()["data"]["result"]
formatted_results: Final = []
for result in results:
metric_data = result["values"]
for timestamp, value in metric_data:
# Convert timestamp to ISO 8601 string with UTC offset
date = datetime.fromtimestamp(float(timestamp)).isoformat() + "+00:00"
spend = float(value)
formatted_results.append({"date": date, "spend": spend})
results: Final = PrometheusQueryResponse.model_validate(_json_response).data.result
formatted_results: Final[list[PrometheusDailySpend]] = [
{
"date": datetime.fromtimestamp(float(timestamp)).isoformat() + "+00:00",
"spend": float(value),
}
for result in results
for timestamp, value in result.values
]
return formatted_results

211
litellm/integrations/rubrik.py
View file

@ -8,11 +8,12 @@ import uuid
from collections import Counter
from collections.abc import Awaitable, Mapping, Sequence
from dataclasses import dataclass
from datetime import datetime
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, Literal, Optional, TypedDict
from typing import TYPE_CHECKING, Final, Literal, Optional, Protocol, TypedDict, overload
import httpx
from typing_extensions import Never, ReadOnly
from typing_extensions import Never, ReadOnly, Required
from litellm._logging import verbose_logger
from litellm.integrations.custom_batch_logger import CustomBatchLogger
@ -30,6 +31,7 @@ from litellm.llms.custom_httpx.http_handler import (
httpxSpecialProvider,
)
from litellm.types.guardrails import GuardrailEventHooks
from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolCallChunk
from litellm.types.utils import (
ChatCompletionMessageToolCall,
Function,
@ -52,17 +54,102 @@ _DROP_WARNING_INTERVAL_SECONDS: Final = 60.0
_EMPTY_MAPPING: Final[Mapping[str, Never]] = MappingProxyType({})
class _ServiceToolCall(TypedDict):
id: ReadOnly[str]
class _ModerationToolCall(TypedDict, total=False):
id: ReadOnly[Required[str]]
class _ServiceMessage(TypedDict, total=False):
class _ModerationMessage(TypedDict, total=False):
content: ReadOnly[str | None]
tool_calls: ReadOnly[Sequence[_ModerationToolCall] | None]
class _ModerationChoice(TypedDict, total=False):
message: ReadOnly[_ModerationMessage | None]
class _ModerationResponse(TypedDict, total=False):
choices: ReadOnly[Sequence[_ModerationChoice]]
class _LogEventKwargs(TypedDict, total=False):
standard_logging_object: ReadOnly[Required[StandardLoggingPayload]]
litellm_call_id: ReadOnly[str]
class _HasCallId(Protocol):
def get(self, key: Literal["litellm_call_id"], /) -> str | None: ...
class _HasModelAttr(Protocol):
model: str | None
class _ResponseSource(Protocol):
def get(self, key: Literal["response"], /) -> "_HasModelAttr | None": ...
class _ModelSource(Protocol):
def get(self, key: Literal["model"], default: str, /) -> str: ...
class _FallbackSource(Protocol):
@overload
def get(self, key: Literal["start_time"], /) -> datetime | None: ...
@overload
def get(self, key: str, /) -> object | None: ...
class _RequestContextSource(Protocol):
@overload
def get(self, key: Literal["optional_params"], /) -> Mapping[str, object] | None: ...
@overload
def get(self, key: str, /) -> object | None: ...
def __contains__(self, key: object, /) -> bool: ...
def __getitem__(self, key: str, /) -> object: ...
class _ToolCallLike(Protocol):
id: str | None
type: str | None
function: Function
class _ModerationSourceToolCall(TypedDict, total=False):
function: ReadOnly[Mapping[str, object] | None]
class _ModerationSourceMessage(TypedDict, total=False):
role: ReadOnly[str]
function_call: ReadOnly[Mapping[str, object] | None]
tool_calls: ReadOnly[Sequence[_ModerationSourceToolCall | None] | None]
class _FlattenedModerationMessage(TypedDict):
role: ReadOnly[str | None]
content: ReadOnly[str]
tool_calls: ReadOnly[Sequence[_ServiceToolCall]]
class _ServiceChoice(TypedDict, total=False):
message: ReadOnly[_ServiceMessage]
class _CorrelatablePayload(TypedDict):
id: str # writable-ok: _apply_correlation_id overwrites the provider id on a deep-copied payload
class _SystemPromptCarrier(TypedDict, total=False):
messages: object # writable-ok: _prepend_system_prompt rebinds messages on the copied payload by design
class _BlockFailurePayload(TypedDict, total=False):
id: object # writable-ok: correlation id is pinned after copying the base payload
model: ReadOnly[object]
model_group: ReadOnly[object]
model_id: ReadOnly[str]
model_parameters: ReadOnly[object]
startTime: ReadOnly[float | None]
endTime: ReadOnly[float | None]
completionStartTime: ReadOnly[float | None]
messages: object # writable-ok: passed to _prepend_system_prompt, which rebinds messages
metadata: ReadOnly[StandardLoggingUserAPIKeyMetadata]
response: str # writable-ok: block failure text replaces the copied response
status: ReadOnly[str]
class _MalformedToolBlockingResponseError(Exception):
@ -385,7 +472,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
@staticmethod
def _stash_block_context(
logging_obj: Optional["LiteLLMLoggingObj"],
request_data: dict,
request_data: dict[str, object],
) -> None:
"""Stash signals so the deferred success-event skips this request and
``async_post_call_failure_hook`` can build the failure payload.
@ -414,12 +501,16 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
request_data["_rubrik_logging_obj"] = logging_obj
@staticmethod
def _normalize_tool_calls(tool_calls: Sequence[object]) -> tuple[ChatCompletionMessageToolCall, ...]:
def _normalize_tool_calls(
tool_calls: Sequence[ChatCompletionToolCallChunk | ChatCompletionMessageToolCall | _ToolCallLike],
) -> tuple[ChatCompletionMessageToolCall, ...]:
"""Convert tool_calls from inputs to ChatCompletionMessageToolCall objects."""
return tuple(RubrikLogger._normalize_tool_call(tc) for tc in tool_calls)
@staticmethod
def _normalize_tool_call(tc: Any) -> ChatCompletionMessageToolCall:
def _normalize_tool_call(
tc: ChatCompletionToolCallChunk | ChatCompletionMessageToolCall | _ToolCallLike,
) -> ChatCompletionMessageToolCall:
if isinstance(tc, ChatCompletionMessageToolCall):
return tc
if isinstance(tc, dict):
@ -460,12 +551,15 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
``content`` is sent so the webhook can moderate the response text;
``None`` when the assistant produced no text (tool-call-only response).
"""
message: Final[dict[str, object]] = {
message: Final[Mapping[str, object]] = {
"role": "assistant",
"content": content or None,
**(
{"tool_calls": tuple(tc.model_dump(exclude_none=True) for tc in tool_calls)}
if tool_calls
else _EMPTY_MAPPING
),
}
if tool_calls:
message["tool_calls"] = tuple(tc.model_dump(exclude_none=True) for tc in tool_calls)
return {
"id": request_id or f"chatcmpl-{uuid.uuid4()}",
"object": "chat.completion",
@ -481,7 +575,9 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
}
@staticmethod
def _flatten_messages_for_moderation(messages: Sequence[object] | None) -> tuple[Mapping[str, Any], ...]:
def _flatten_messages_for_moderation(
messages: Sequence[AllMessageValues | None] | None,
) -> tuple[_FlattenedModerationMessage, ...]:
"""Collapse each message's content to a plain string for the webhook.
litellm normalizes Anthropic ``/v1/messages`` requests to OpenAI shape,
@ -502,7 +598,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
)
@staticmethod
def _moderation_text_parts(message: Mapping[str, Any]) -> tuple[str, ...]:
def _moderation_text_parts(message: _ModerationSourceMessage) -> tuple[str, ...]:
"""Every attacker-controlled text segment of a message: its content plus
the arguments of any tool call or deprecated function call."""
fc: Final = message.get("function_call")
@ -530,16 +626,8 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
``/v1/messages`` requests too. Optional fields are sent only when
present so the payload stays clean.
"""
payload: Final[dict[str, object]] = {
"model": inputs.get("model") or request_data.get("model") or "",
"messages": RubrikLogger._flatten_messages_for_moderation(inputs.get("structured_messages")),
}
tools: Final = inputs.get("tools")
if tools is not None:
payload["tools"] = tools
user: Final = request_data.get("user")
if user:
payload["user"] = user
# Fall back to litellm_call_id, the stable cross-provider join key the
# response/tool path uses (see _correlation_id). LiteLLM does not
# populate request_data["correlation_key"]; it carries litellm_call_id.
@ -547,14 +635,18 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
# when correlation_key is empty, so without this the block fires but no
# log is ever written. An explicit correlation_key still wins.
correlation_key: Final = request_data.get("correlation_key") or request_data.get("litellm_call_id")
if correlation_key:
payload["correlation_key"] = correlation_key
return payload
return {
"model": inputs.get("model") or request_data.get("model") or "",
"messages": RubrikLogger._flatten_messages_for_moderation(inputs.get("structured_messages")),
**({"tools": tools} if tools is not None else _EMPTY_MAPPING),
**({"user": user} if user else _EMPTY_MAPPING),
**({"correlation_key": correlation_key} if correlation_key else _EMPTY_MAPPING),
}
@staticmethod
def _extract_request_data(
call_details: Mapping[str, Any],
request_data: Mapping[str, object] | None,
call_details: _RequestContextSource,
request_data: _RequestContextSource | None,
) -> Mapping[str, object]:
"""Extract original request data from model_call_details for the
response moderation service envelope.
@ -590,7 +682,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
}
@staticmethod
def _sanitize_proxy_server_request(proxy_server_request: object) -> object:
def _sanitize_proxy_server_request(proxy_server_request: Mapping[str, object] | str | None) -> object:
"""Allowlist only routing fields (``url``, ``method``) when forwarding
``proxy_server_request`` to an external webhook, dropping inbound
``headers`` (Authorization, Cookie, x-api-key, ...) and the raw
@ -600,18 +692,19 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
return {key: proxy_server_request[key] for key in ("url", "method") if key in proxy_server_request}
@staticmethod
def _resolve_model(request_data: Mapping[str, object], call_details: Mapping[str, str]) -> str:
def _resolve_model(request_data: _ResponseSource, call_details: _ModelSource) -> str:
"""Get the model name for the ModifyResponseException."""
response: Final = request_data.get("response")
if response and hasattr(response, "model"):
response_model: Final[str | None] = getattr(response, "model", None)
return response_model or "unknown"
return response.model or "unknown"
return call_details.get("model", "unknown")
# -- Logging hooks ---------------------------------------------------------
@staticmethod
def _correlation_id(call_details: Mapping[str, str], request_data: Mapping[str, str] | None = None) -> str | None:
def _correlation_id(
call_details: _HasCallId | _LogEventKwargs, request_data: _HasCallId | None = None
) -> str | None:
"""The id that joins a blocked request's two S3 logs by filename: the
moderation (``_blocking``) log and the failure (response) log.
@ -625,7 +718,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
return call_details.get("litellm_call_id") or (request_data or _EMPTY_MAPPING).get("litellm_call_id")
@classmethod
def _apply_correlation_id(cls, payload: dict[str, object], source: Mapping[str, str]) -> None:
def _apply_correlation_id(cls, payload: _CorrelatablePayload, source: _HasCallId | _LogEventKwargs) -> None:
"""Pin ``payload["id"]`` to ``litellm_call_id`` in place so this log
shares its S3 filename id with the moderation (``_blocking``) and
failure logs for the same request -- for every provider.
@ -645,7 +738,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
payload["id"] = correlated
@staticmethod
def _prepend_system_prompt(payload: dict[str, object], source: Mapping[str, object]) -> None:
def _prepend_system_prompt(payload: _SystemPromptCarrier, source: Mapping[str, object]) -> None:
"""Prepend ``source["system"]`` onto ``payload["messages"]``.
Builds a NEW messages list rather than mutating ``payload["messages"]``
@ -673,9 +766,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
exc_info=True,
)
async def _prepare_log_payload(
self, kwargs: Mapping[str, object], event_type: str
) -> StandardLoggingPayload | None:
async def _prepare_log_payload(self, kwargs: _LogEventKwargs, event_type: str) -> StandardLoggingPayload | None:
"""Shared logic for success logging (sampled)."""
if random.random() > self.sampling_rate:
verbose_logger.debug("Skipping Rubrik %s logging (sampling_rate=%s)", event_type, self.sampling_rate)
@ -684,12 +775,12 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
# Deep-copy so mutations don't affect other callbacks sharing this object
standard_logging_payload: Final[StandardLoggingPayload] = safe_deep_copy(kwargs["standard_logging_object"])
self._apply_correlation_id(standard_logging_payload, kwargs) # pyright: ignore[reportArgumentType] # StandardLoggingPayload is dict[str,Any] at runtime
self._apply_correlation_id(standard_logging_payload, kwargs)
self._prepend_system_prompt(standard_logging_payload, kwargs) # pyright: ignore[reportArgumentType] # StandardLoggingPayload is dict[str,Any] at runtime
return standard_logging_payload
async def _append_and_maybe_flush(self, payload) -> None:
async def _append_and_maybe_flush(self, payload: Mapping[str, object]) -> None:
self._ensure_periodic_flush_task()
self.log_queue.append(payload)
self._enforce_max_queue_size()
@ -714,7 +805,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
self._dropped_since_warning = 0
self._last_drop_warning_time = now
async def _enqueue_log_event(self, kwargs: Mapping[str, object], event_type: str):
async def _enqueue_log_event(self, kwargs: _LogEventKwargs, event_type: str):
try:
payload: Final = await self._prepare_log_payload(kwargs, event_type)
if payload is None:
@ -835,7 +926,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
logging_obj: "LiteLLMLoggingObj",
exception: "ModifyResponseException",
user_api_key_dict: "UserAPIKeyAuth",
) -> StandardLoggingPayload:
) -> _BlockFailurePayload:
"""Build a failure-style payload using the exception text as response.
Blocked-tool events are security-relevant and **bypass sampling**:
@ -877,9 +968,9 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
call_details: Final = logging_obj.model_call_details
exception_text: Final = f"{type(exception).__name__}: {exception.message}"
base: Final = call_details.get("standard_logging_object")
base: Final[StandardLoggingPayload | None] = call_details.get("standard_logging_object")
if base is not None:
payload: dict[str, object] = safe_deep_copy(base)
payload: _BlockFailurePayload = self._copy_block_payload_base(base)
else:
verbose_logger.debug(
"Rubrik: standard_logging_object not yet on model_call_details "
@ -901,6 +992,10 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
return payload
@staticmethod
def _copy_block_payload_base(base: StandardLoggingPayload) -> _BlockFailurePayload:
return safe_deep_copy(base)
@staticmethod
def _caller_metadata(user_api_key_dict: "UserAPIKeyAuth") -> StandardLoggingUserAPIKeyMetadata:
"""Identify the caller whose request was blocked.
@ -923,9 +1018,9 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
@classmethod
def _build_fallback_payload(
cls,
call_details: Mapping[str, Any],
call_details: _FallbackSource,
user_api_key_dict: "UserAPIKeyAuth",
) -> dict[str, object]:
) -> _BlockFailurePayload:
# Convert datetime to a Unix float so json.dumps can serialize it.
# httpx's json= parameter uses stdlib json.dumps with no custom encoder.
_raw_start: Final = call_details.get("start_time")
@ -959,7 +1054,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
response: Final = await self.async_httpx_client.post(
url=self.logging_endpoint,
json=data,
headers=self._headers,
headers=dict(self._headers),
)
response.raise_for_status()
except httpx.HTTPStatusError as e:
@ -1013,7 +1108,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
# -- Webhook services ------------------------------------------------------
async def _post_json(self, endpoint: str, payload: Mapping[str, object], service_name: str) -> Mapping[str, Any]:
async def _post_json(self, endpoint: str, payload: Mapping[str, object], service_name: str) -> _ModerationResponse:
"""POST ``payload`` to a Rubrik webhook and return its dict response.
Raises:
@ -1023,11 +1118,11 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
verbose_logger.debug("Sending request to %s: %s", service_name, endpoint)
http_response: Final = await self.moderation_client.post(
endpoint,
json=payload,
headers=self._headers,
json=dict(payload),
headers=dict(self._headers),
)
http_response.raise_for_status()
result: Final[object] = http_response.json()
result: Final[_ModerationResponse | None] = http_response.json()
if not isinstance(result, dict):
raise TypeError(
f"{service_name} returned non-dict JSON "
@ -1040,7 +1135,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
self,
response_data: Mapping[str, object],
request_data: Mapping[str, object],
) -> Mapping[str, Any]:
) -> _ModerationResponse:
"""Post the ``{request, response}`` envelope to the after_completion
webhook and return its (possibly rewritten) response.
@ -1056,7 +1151,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
"Response moderation service",
)
async def _post_to_prompt_moderation_endpoint(self, payload: Mapping[str, object]) -> Mapping[str, Any]:
async def _post_to_prompt_moderation_endpoint(self, payload: Mapping[str, object]) -> _ModerationResponse:
"""Post a bare OpenAI request to the before_prompt webhook.
Returns ``{}`` (passthrough) or a synthetic chat.completion (block).
@ -1064,14 +1159,14 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
return await self._post_json(self.prompt_moderation_endpoint, payload, "Prompt moderation service")
@staticmethod
def _extract_prompt_refusal(service_response: Mapping[str, Any]) -> str | None:
def _extract_prompt_refusal(service_response: _ModerationResponse) -> str | None:
"""Return the refusal text when the prompt was blocked, else None.
The before_prompt webhook returns ``{}`` (passthrough) or a synthetic
chat.completion whose ``choices[0].message.content`` is the refusal
explanation.
"""
choices: Final[Sequence[_ServiceChoice] | None] = service_response.get("choices")
choices: Final = service_response.get("choices")
if not choices:
return None
message: Final = choices[0].get("message") or _EMPTY_MAPPING
@ -1080,7 +1175,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
@staticmethod
def _extract_response_block(
service_response: Mapping[str, Any],
service_response: _ModerationResponse,
all_tool_calls: Sequence[ChatCompletionMessageToolCall],
sent_content: str,
) -> BlockedResponseResult | None:
@ -1103,7 +1198,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
Expects service_response in OpenAI chat completion format:
{"choices": [{"message": {"tool_calls": [...], "content": "..."}}]}
"""
choices: Final[Sequence[_ServiceChoice]] = service_response.get("choices") or ()
choices: Final = service_response.get("choices") or ()
if not choices:
raise _MalformedToolBlockingResponseError("Response moderation service returned empty response")

81
litellm/integrations/s3.py
View file

@ -1,11 +1,18 @@
#### What this does ####
# On success + failure, log events to Supabase
import hashlib
from datetime import datetime
from typing import Final, cast
import litellm
from litellm._logging import print_verbose, verbose_logger
from litellm.constants import (
MAX_S3_OBJECT_DOWNLOAD_FILENAME_BYTES,
MAX_S3_OBJECT_KEY_BYTES,
S3_BOUNDED_OBJECT_KEY_HEAD_BYTES,
S3_PREFIX_DIGEST_CHARS,
)
from litellm.types.utils import StandardLoggingPayload
@ -133,9 +140,7 @@ class S3Logger:
s3_file_name,
)
s3_object_download_filename: Final = (
"time-" + start_time.strftime("%Y-%m-%dT%H-%M-%S-%f") + "_" + payload["id"] + ".json"
)
s3_object_download_filename: Final = get_s3_object_download_filename(start_time, payload["id"])
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
@ -198,6 +203,47 @@ def resolve_sse_params(
return algorithm, valid_key_id
S3_MIN_BOUNDED_FILE_NAME_BYTES: Final = 64
def _truncate_to_utf8_bytes(value: str, max_bytes: int) -> str:
"""Trim `value` so its UTF-8 encoding fits `max_bytes`, never splitting a character."""
if max_bytes <= 0:
return ""
encoded: Final = value.encode("utf-8")
if len(encoded) <= max_bytes:
return value
return encoded[:max_bytes].decode("utf-8", errors="ignore")
def get_s3_object_download_filename(start_time: datetime, response_id: str) -> str:
"""Content-Disposition filename for the uploaded object, bounded to the metadata header cap."""
sanitized_response_id: Final = response_id.replace("/", "_").replace('"', "_")
file_name: Final = f"time-{start_time.strftime('%Y-%m-%dT%H-%M-%S-%f')}_{response_id}"
sanitized_file_name: Final = f"time-{start_time.strftime('%Y-%m-%dT%H-%M-%S-%f')}_{sanitized_response_id}"
budget: Final = MAX_S3_OBJECT_DOWNLOAD_FILENAME_BYTES - len(b".json")
if len(sanitized_file_name.encode("utf-8")) <= budget:
return sanitized_file_name + ".json"
return _bounded_s3_file_name(file_name, sanitized_file_name, budget) + ".json"
def _bounded_s3_file_name(s3_file_name: str, sanitized_s3_file_name: str, max_bytes: int) -> str:
"""As much of the file name as `max_bytes` allows, then the sha256 of the whole name."""
digest: Final = hashlib.sha256(s3_file_name.encode("utf-8")).hexdigest()
head_budget: Final = min(S3_BOUNDED_OBJECT_KEY_HEAD_BYTES, max_bytes - len(digest) - 1)
head: Final = _truncate_to_utf8_bytes(sanitized_s3_file_name, head_budget)
return f"{head}_{digest}" if head else digest
def _bounded_s3_prefix(configured_prefix: str, max_bytes: int) -> str:
"""As much of the configured prefix as fits, then a digest segment naming the full prefix."""
digest_segment: Final = hashlib.sha256(configured_prefix.encode("utf-8")).hexdigest()[:S3_PREFIX_DIGEST_CHARS] + "/"
if max_bytes < len(digest_segment):
return ""
head: Final = _truncate_to_utf8_bytes(configured_prefix, max_bytes - len(digest_segment) - 1).rstrip("/")
return f"{head}/{digest_segment}" if head else digest_segment
def get_s3_object_key(
s3_path: str,
prefix: str,
@ -205,12 +251,23 @@ def get_s3_object_key(
s3_file_name: str,
) -> str:
sanitized_s3_file_name: Final = s3_file_name.replace("/", "_")
s3_object_key = (
(s3_path.rstrip("/") + "/" if s3_path else "")
+ prefix
+ start_time.strftime("%Y-%m-%d")
+ "/"
+ sanitized_s3_file_name
) # we need the s3 key to include the time, so we log cache hits too
s3_object_key += ".json"
return s3_object_key
configured_prefix: Final = (s3_path.rstrip("/") + "/" if s3_path else "") + prefix
date_segment: Final = start_time.strftime("%Y-%m-%d") + "/"
# we need the s3 key to include the time, so we log cache hits too
s3_object_key: Final = configured_prefix + date_segment + sanitized_s3_file_name + ".json"
if len(s3_object_key.encode("utf-8")) <= MAX_S3_OBJECT_KEY_BYTES:
return s3_object_key
# shorten the response id first and only trim the configured prefix if that is what does not
# fit, so prefix scoped IAM policies and lifecycle rules keep matching
budget: Final = MAX_S3_OBJECT_KEY_BYTES - len(date_segment.encode("utf-8")) - len(b".json")
prefix_bytes: Final = len(configured_prefix.encode("utf-8"))
if prefix_bytes + S3_MIN_BOUNDED_FILE_NAME_BYTES <= budget:
bounded_file_name: Final = _bounded_s3_file_name(s3_file_name, sanitized_s3_file_name, budget - prefix_bytes)
return configured_prefix + date_segment + bounded_file_name + ".json"
shortest_file_name: Final = _bounded_s3_file_name(
s3_file_name, sanitized_s3_file_name, S3_MIN_BOUNDED_FILE_NAME_BYTES
)
bounded_prefix: Final = _bounded_s3_prefix(configured_prefix, budget - len(shortest_file_name.encode("utf-8")))
return bounded_prefix + date_segment + shortest_file_name + ".json"

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