merge upstream main: refresh stale CI results

This commit is contained in:
Artem Burei 2026-09-01 02:25:48 +02:00
commit b4a2a47ea5
2392 changed files with 157166 additions and 21211 deletions

View file

@ -430,7 +430,7 @@ jobs:
--verbose \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
-vv \
--cov=./litellm \
--cov=./litellm --cov=./enterprise/litellm_enterprise \
--cov-report=xml \
--junitxml=test-results/junit.xml \
--durations=20 \
@ -504,7 +504,7 @@ jobs:
--verbose \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
-vv \
--cov=./litellm \
--cov=./litellm --cov=./enterprise/litellm_enterprise \
--cov-report=xml \
--junitxml=test-results/junit.xml \
--durations=20 \
@ -631,7 +631,7 @@ jobs:
--verbose \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
-v -x \
--cov=./litellm --cov-report=xml \
--cov=./litellm --cov=./enterprise/litellm_enterprise --cov-report=xml \
--junitxml=test-results/junit.xml \
--durations=5 \
-n 2"
@ -651,126 +651,6 @@ jobs:
- auth_ui_unit_tests_coverage.xml
- auth_ui_unit_tests_coverage
proxy_behavior_tests:
docker:
- *python312_image
- image: cimg/postgres:16.0@sha256:b125148bc76e8e8eee5eb3ad6020a3a14110a14e8192f1c645128afebe2e2f84
environment:
POSTGRES_USER: postgres
POSTGRES_PASSWORD: postgres
POSTGRES_DB: litellm_test
working_directory: ~/project
environment:
DATABASE_URL: "postgresql://postgres:postgres@localhost:5432/litellm_test"
steps:
- checkout
- skip_if_unrelated_changes
- setup_google_dns
- install_uv
- install_rust
- run:
name: Install Dependencies
command: |
uv sync --frozen --all-groups --all-extras --python 3.12
- wait_for_service:
url: tcp://localhost:5432
timeout: "60"
- run:
name: Seed DB schema via prisma db push
command: |
uv run --no-sync prisma db push --schema litellm/proxy/schema.prisma --accept-data-loss
- run:
name: Generate Prisma Client
command: uv run --no-sync python -m prisma generate
- run:
name: Run proxy management behavior tests
command: |
mkdir -p test-results
uv run --no-sync python -m pytest tests/proxy_behavior \
-v --junitxml=test-results/junit.xml --durations=10
no_output_timeout: 15m
- store_test_results:
path: test-results
proxy_security_tests:
docker:
- *python312_image
- image: cimg/postgres:16.0@sha256:b125148bc76e8e8eee5eb3ad6020a3a14110a14e8192f1c645128afebe2e2f84
environment:
POSTGRES_USER: postgres
POSTGRES_PASSWORD: postgres
POSTGRES_DB: litellm_test
working_directory: ~/project
environment:
DATABASE_URL: "postgresql://postgres:postgres@localhost:5432/litellm_test"
steps:
- checkout
- skip_if_unrelated_changes
- setup_google_dns
- install_uv
- install_rust
- run:
name: Install Dependencies
command: |
uv sync --frozen --all-groups --all-extras --python 3.12
- wait_for_service:
url: tcp://localhost:5432
timeout: "60"
- run:
name: Seed DB schema via prisma db push
command: |
uv run --no-sync prisma db push --schema litellm/proxy/schema.prisma --accept-data-loss
- run:
name: Generate Prisma Client
command: uv run --no-sync python -m prisma generate
- run:
name: Run proxy security tests
command: |
mkdir -p test-results
uv run --no-sync python -m pytest tests/proxy_security_tests \
-v --junitxml=test-results/junit.xml --durations=10
no_output_timeout: 15m
- store_test_results:
path: test-results
schema_migration_check:
docker:
- *python312_image
- image: cimg/postgres:16.0@sha256:b125148bc76e8e8eee5eb3ad6020a3a14110a14e8192f1c645128afebe2e2f84
environment:
POSTGRES_USER: postgres
POSTGRES_PASSWORD: postgres
POSTGRES_DB: litellm_test
working_directory: ~/project
environment:
# An empty database; the test applies every committed migration itself.
DATABASE_URL: "postgresql://postgres:postgres@localhost:5432/litellm_test"
steps:
- checkout
- skip_if_unrelated_changes
- setup_google_dns
- install_uv
- install_rust
- run:
name: Install Dependencies
command: |
uv sync --frozen --all-groups --all-extras --python 3.12
- wait_for_service:
url: tcp://localhost:5432
timeout: "60"
- run:
name: Generate Prisma Client
command: uv run --no-sync python -m prisma generate
- run:
name: Check schema.prisma is in sync with committed migrations
command: |
mkdir -p test-results
uv run --no-sync python -m pytest tests/proxy_migration_tests \
-v --junitxml=test-results/junit.xml --durations=10
no_output_timeout: 15m
- store_test_results:
path: test-results
litellm_router_testing: # Runs all tests with the "router" keyword
docker:
- *python312_image
@ -858,7 +738,7 @@ jobs:
--verbose \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
-v -x \
--cov=./litellm --cov-report=xml \
--cov=./litellm --cov=./enterprise/litellm_enterprise --cov-report=xml \
--junitxml=test-results/junit.xml \
--durations=5 \
-n 4"
@ -985,7 +865,7 @@ jobs:
--verbose \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
-vv \
--cov=./litellm --cov-report=xml \
--cov=./litellm --cov=./enterprise/litellm_enterprise --cov-report=xml \
--junitxml=test-results/junit.xml \
--durations=20 \
-n 4 \
@ -1030,7 +910,7 @@ jobs:
--verbose \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
-vv -x -s \
--cov=./litellm --cov-report=xml \
--cov=./litellm --cov=./enterprise/litellm_enterprise --cov-report=xml \
--junitxml=test-results/junit.xml \
--durations=5"
no_output_timeout: 15m
@ -1074,7 +954,7 @@ jobs:
--verbose \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
-vv \
--cov=./litellm --cov-report=xml \
--cov=./litellm --cov=./enterprise/litellm_enterprise --cov-report=xml \
--junitxml=test-results/junit.xml \
--durations=5 \
-n 2 \
@ -1120,7 +1000,7 @@ jobs:
--verbose \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
-vv -x -s \
--cov=./litellm --cov-report=xml \
--cov=./litellm --cov=./enterprise/litellm_enterprise --cov-report=xml \
--junitxml=test-results/junit.xml \
--durations=5 \
--retries 3 --retry-delay 5"
@ -1211,7 +1091,7 @@ jobs:
--verbose \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
-vv -x \
--cov=./litellm --cov-report=xml \
--cov=./litellm --cov=./enterprise/litellm_enterprise --cov-report=xml \
--junitxml=test-results/junit.xml \
--durations=5 \
-n 4"
@ -1255,7 +1135,7 @@ jobs:
--verbose \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
-vv -x \
--cov=./litellm --cov-report=xml \
--cov=./litellm --cov=./enterprise/litellm_enterprise --cov-report=xml \
--junitxml=test-results/junit.xml \
--durations=5 \
-n 4"
@ -1274,40 +1154,6 @@ jobs:
paths:
- search_coverage.xml
- search_coverage
litellm_mapped_enterprise_tests:
docker:
- *python312_image
working_directory: ~/project
resource_class: large
steps:
- checkout
- skip_if_unrelated_changes
- setup_google_dns
- install_uv
- install_rust
- run:
name: Install Dependencies
command: |
uv sync --frozen --all-groups --all-extras --python 3.12
- setup_litellm_enterprise_pip
- run:
name: Run enterprise tests
command: |
uv run --no-sync python -m prisma generate
mkdir -p test-results
TEST_FILES=$(circleci tests glob "tests/enterprise/**/test_*.py")
echo "$TEST_FILES" | circleci tests run \
--verbose \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
-v -x \
--junitxml=test-results/junit-enterprise.xml \
--durations=10 \
-n 4"
no_output_timeout: 15m
# Store test results
- store_test_results:
path: test-results
batches_testing:
docker:
- *python312_image
@ -1333,7 +1179,7 @@ jobs:
--verbose \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
-vv -x -s \
--cov=./litellm --cov-report=xml \
--cov=./litellm --cov=./enterprise/litellm_enterprise --cov-report=xml \
--junitxml=test-results/junit.xml \
--durations=5 \
-n 2"
@ -1377,7 +1223,7 @@ jobs:
--verbose \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
-vv -x -s \
--cov=./litellm --cov-report=xml \
--cov=./litellm --cov=./enterprise/litellm_enterprise --cov-report=xml \
--junitxml=test-results/junit.xml \
--durations=5 \
-n 2"
@ -1422,7 +1268,7 @@ jobs:
--verbose \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
-vv -x \
--cov=./litellm --cov-report=xml \
--cov=./litellm --cov=./enterprise/litellm_enterprise --cov-report=xml \
--junitxml=test-results/junit.xml \
--durations=5 \
-n 4"
@ -1501,7 +1347,7 @@ jobs:
--verbose \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
-vv \
--cov=./litellm --cov-report=xml \
--cov=./litellm --cov=./enterprise/litellm_enterprise --cov-report=xml \
-n 4 \
--junitxml=test-results/junit.xml \
--durations=5 \
@ -1546,7 +1392,7 @@ jobs:
--verbose \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
-vv -x -s \
--cov=./litellm --cov-report=xml \
--cov=./litellm --cov=./enterprise/litellm_enterprise --cov-report=xml \
--junitxml=test-results/junit.xml \
--durations=5"
no_output_timeout: 15m
@ -1599,7 +1445,7 @@ jobs:
--verbose \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \
-vv -x -s \
--cov=./litellm --cov-report=xml \
--cov=./litellm --cov=./enterprise/litellm_enterprise --cov-report=xml \
--junitxml=test-results/junit.xml \
--durations=5 -n 2 \
--reruns 2 --reruns-delay 1"
@ -2575,45 +2421,6 @@ jobs:
- wait_for_service:
url: http://localhost:4000
timeout: "300"
# Add Ruby installation and testing before the existing Node.js and Python tests
- run:
name: Install Ruby and Bundler
command: |
# Clone RVM at pinned tag and verify the commit SHA matches the
# published tag before running its install script.
RVM_VERSION="1.29.12"
RVM_EXPECTED_SHA="6bfc9213c9d6914fe756f524eb034a403d51db81"
git clone --depth 1 --branch "$RVM_VERSION" https://github.com/rvm/rvm.git /tmp/rvm
RVM_ACTUAL_SHA="$(git -C /tmp/rvm rev-parse HEAD)"
if [ "$RVM_ACTUAL_SHA" != "$RVM_EXPECTED_SHA" ]; then
echo "RVM tag $RVM_VERSION resolved to $RVM_ACTUAL_SHA; expected $RVM_EXPECTED_SHA" >&2
exit 1
fi
# Import RVM signing keys (used by `rvm install` to verify Ruby tarballs)
gpg --keyserver hkp://keyserver.ubuntu.com --recv-keys 409B6B1796C275462A1703113804BB82D39DC0E3 7D2BAF1CF37B13E2069D6956105BD0E739499BDB
# Install RVM from the verified checkout. The install script
# sources `scripts/functions/installer` using paths relative to
# its own working directory, so it must be run from /tmp/rvm.
(cd /tmp/rvm && ./install --path "$HOME/.rvm")
source "$HOME/.rvm/scripts/rvm"
# Install Ruby 3.2.2 (RVM verifies the tarball PGP signature)
rvm install 3.2.2
rvm use 3.2.2 --default
# Install latest Bundler
gem install bundler
- run:
name: Run Ruby tests
command: |
source $HOME/.rvm/scripts/rvm
cd tests/pass_through_tests/ruby_passthrough_tests
bundle install
bundle exec rspec
no_output_timeout: 30m
# Install Node.js directly from nodejs.org with SHA256 verification,
# instead of piping NodeSource's setup_24.x apt-repo installer into
# sudo bash (which runs a mutable upstream script unattended).
@ -3105,12 +2912,6 @@ workflows:
filters: *main_branches
- auth_ui_unit_tests:
filters: *main_branches
- proxy_behavior_tests:
filters: *main_branches
- proxy_security_tests:
filters: *main_branches
- schema_migration_check:
filters: *main_branches
- build_docker_database_image:
filters: *main_branches
- e2e_ui_testing:
@ -3167,8 +2968,6 @@ workflows:
filters: *main_branches
- search_testing:
filters: *main_branches
- litellm_mapped_enterprise_tests:
filters: *main_branches
- batches_testing:
filters: *main_branches
- litellm_utils_testing:
@ -3191,7 +2990,6 @@ workflows:
- guardrails_testing
- ocr_testing
- search_testing
- litellm_mapped_enterprise_tests
- batches_testing
- litellm_utils_testing
- pass_through_unit_testing

3
.github/CODEOWNERS vendored
View file

@ -1,6 +1,9 @@
/ui/ @yuneng-berri @ryan-crabbe-berri
/litellm/proxy/_experimental/out/ @yuneng-berri @ryan-crabbe-berri
/ui/Dockerfile
/ui/nginx.conf
/ui/litellm-dashboard/src/lib/http/schema.d.ts
/ui/litellm-dashboard/tsconfig.tsbuildinfo
/model_prices_and_context_window.json @mateo-berri
/litellm/model_prices_and_context_window_backup.json @mateo-berri
/litellm-proxy-extras/litellm_proxy_extras/migrations/ @yuneng-berri @ryan-crabbe-berri

View file

@ -5,24 +5,21 @@ description: >-
test_paths:
- reason: >-
The caching suite in tests/local_testing, which runs nowhere. Every job that globs that
directory either deselects it (local_testing_part1 and part2 carry `-k "... and not caching
and not cache"`) or keeps only another keyword (langfuse, router, assistants), and no job
names these files the way redis_caching_unit_tests names test_dual_cache.py. Measured
2026-08-20 by collecting the directory under each job's own selector: 118 tests across
these eight files are selected by none of them. Listed so the gap is a decision rather
than an accident, and so the --slices guard has a baseline to ratchet down from. Revisit
when tests/local_testing is ported off CircleCI, where the keyless part of this suite
belongs in a real job
What is left of the caching suite in tests/local_testing that runs nowhere. Every job that
globs that directory either deselects it (local_testing_part1 and part2 carry `-k "... and
not caching and not cache"`) or keeps only another keyword (langfuse, router, assistants),
and no job names these files the way redis_caching_unit_tests names test_dual_cache.py.
The gap was eight files and 118 tests when measured 2026-08-20; the five keyless ones now
run in the caching-local shard, leaving these three. Measured 2026-08-21 with no provider
credentials and no Redis: test_caching.py needs both (37 of 65 fail without them),
test_disk_cache_unit_tests.py needs OPENAI_API_KEY for 2 of its 4, and
test_gcs_cache_unit_tests.py needs GCS credentials for all 4. They want the keyless/live
split that porting tests/local_testing off CircleCI will force, not a job that is red by
construction
paths:
- tests/local_testing/test_cache_preset_key.py
- tests/local_testing/test_caching.py
- tests/local_testing/test_caching_handler.py
- tests/local_testing/test_disk_cache_unit_tests.py
- tests/local_testing/test_gcs_cache_unit_tests.py
- tests/local_testing/test_prompt_caching.py
- tests/local_testing/test_responses_stream_cache_keys.py
- tests/local_testing/test_unit_test_caching.py
- reason: >-
The end-to-end suite runs against a deployed proxy from its own in-cluster rig rather than
from a pull request; it needs a live gateway and provider credentials no PR job holds
@ -92,17 +89,6 @@ test_paths:
- tests/integration/sandbox/test_e2b_sandbox.py
- tests/integration/test_oci_integration.py
- tests/integration/test_oci_proxy_integration.py
- reason: >-
A unit test for the proxy-extras package that no job invokes, while the package's other
tests live under tests/proxy_migration_tests. Measured 2026-08-20: 24 of its 28 tests pass
and the 4 in TestMigrationSQLIdempotency fail, because 13 migrations from 2026-03 onward use
bare CREATE TABLE, ADD COLUMN, CREATE INDEX and ADD CONSTRAINT rather than the guarded forms
this file requires. It also matches those keywords inside SQL comments, so two further
migrations are reported that are in fact fine. Wiring it up means deciding what to do about
the 13 first, and they cannot simply be edited: Prisma checksums an applied migration, so a
changed one breaks migrate deploy for existing installs
paths:
- tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py
dockerfiles:
- reason: >-

5
.github/mutmut-coverage.rc vendored Normal file
View file

@ -0,0 +1,5 @@
# mutmut's gather_coverage() looks covered lines up by absolute path, so the
# repo's `relative_files = true` makes every lookup miss and mutmut generates
# zero mutants. Point COVERAGE_RCFILE here for mutation runs only.
[run]
relative_files = false

View file

@ -1,7 +1,10 @@
<!-- The whole description's target audience is humans, not AI agents: write it in plain, simple,
everyday engineering language, extremely parsable and readable at a glance. This goes double for
the TLDR, User Flow, and Caveats sections -->
## TLDR
<!-- Fill in the bullets below and keep each one short and concrete: one line per bullet, roughly 10 words max
This section must be extremely human parsable, comprehensible, and readable: its target audience is humans, not AI agents -->
<!-- Fill in the bullets below and keep each one short and concrete: one line per bullet, roughly 10 words max -->
Problem this solves:
@ -110,8 +113,20 @@ If you're seeing a delay in your PR being merged, ping the LiteLLM Team on [Slac
## Caveats (if any)
<!-- Short bullet points, just like the TLDR: one line per bullet, roughly 10 words max
<!-- Group caveats under severity subheadings (### Severe, ### High, ### Medium, ### Low), with
short bullet points inside each, just like the TLDR: one line per bullet, roughly 10 words max
Call out known limitations, follow-up work, or anything a reviewer should watch out for
Include only the tiers that have caveats; drop the empty ones
- Severe: inherent to what the PR deliberately ships, there even when the code works as intended:
it can degrade or take down a running deployment (e.g. a slow or table-locking boot migration),
rewrite data by design, break an existing workflow on purpose, or change auth behavior. An
operator must plan around it before rollout
- High: an unintended hole: a correctness, security, data-loss, or backward-compatibility bug,
unsafe to ship as is
- Medium: a real gap someone can hit, but with a workaround or a narrow blast radius
- Low: anything else worth noting: naming, cleanup, an edge case nobody hits
Nest bullets as deep as helps: hierarchy beats one long line when it makes things clearer to a
human reader
Leave this section empty if there are none -->
## QA runbook
@ -134,6 +149,6 @@ Example checklists:
- [ ] Sanity check: this test makes sense to add and is not hand-wavey (e.g., assert actual expected spend instead of just spend > 0) or potentially flaky
-->
### Final Attestation
## Final Attestation
- [ ] The tests check the right things, including the edge cases, and regressions in the respective real-world customer use-cases are not possible after this PR

View file

@ -312,8 +312,23 @@ def _matchable_names(relative_path: str) -> frozenset[str]:
)
def _workflow_named_tokens() -> frozenset[str]:
"""Test tokens a GitHub Actions job names directly.
A CircleCI `-k` that deselects a file no longer means the file runs nowhere once a
workflow names it, so the slice check has to credit those the same way the census does.
"""
return _invoked_test_tokens(
scalar
for path in _config_files()
if path != CIRCLECI_CONFIG
for scalar in _scalars(yaml.safe_load(path.read_text(encoding="utf-8")), path.name)
)
def _deselected_everywhere(allowlist: Allowlist) -> tuple[Finding, ...]:
slices: Final = _slices()
named_by_workflow: Final = _workflow_named_tokens()
globbed: Final = tuple(
path
for path in _test_files()
@ -326,6 +341,7 @@ def _deselected_everywhere(allowlist: Allowlist) -> tuple[Finding, ...]:
)
for path in globbed
if not allowlist.covers_test(path)
and not any(_token_covers(token, path) for token in named_by_workflow)
and not any(slice_.claims(path, _matchable_names(path)) for slice_ in slices)
)

198
.github/scripts/e2e_egress_sentinel.py vendored Executable file
View file

@ -0,0 +1,198 @@
"""Prove an e2e replay run makes zero outbound provider calls, by counting them.
`serve` pins each provider host (`--host`) to a local sink address in the hosts
file and binds a counting listener on that address, so any connection the proxy
or the record/replay edge opens to a real provider is redirected to the sink,
recorded as one line in `--hits-file`, and never leaves the box. The record and
replay edge only ever dials `127.0.0.1:<edge-port>` (a different host than the
pinned provider names), so in a clean replay the sink sees nothing; a single hit
means a provider call escaped the bundle. `assert-empty` turns that hit file into
the pass/fail check.
Stdlib only, so CI runs it under the system interpreter as root (binding :443 and
editing the hosts file both need root); `--sink-address`, `--port`, and
`--hosts-file` are injectable so it runs unprivileged against a temp hosts file on
a high port under test.
"""
# ruff: noqa: T201 # CLI script: its stdout/stderr progress and results are the interface
from __future__ import annotations
import argparse
import json
import os
import signal
import socket
import sys
import threading
import time
from dataclasses import dataclass
from pathlib import Path
from types import FrameType
from typing import Final
_BLOCK_BEGIN: Final = "# BEGIN e2e-egress-sentinel"
_BLOCK_END: Final = "# END e2e-egress-sentinel"
@dataclass(frozen=True, slots=True)
class ServeConfig:
hosts: tuple[str, ...]
sink_address: str
ports: tuple[int, ...]
hits_file: Path
hosts_file: Path
ready_file: Path | None
pid_file: Path | None
def _pin_block(sink_address: str, hosts: tuple[str, ...]) -> str:
lines = "\n".join(f"{sink_address}\t{host}" for host in hosts)
return f"\n{_BLOCK_BEGIN}\n{lines}\n{_BLOCK_END}\n"
def _install_pins(hosts_file: Path, sink_address: str, hosts: tuple[str, ...]) -> bytes:
original = hosts_file.read_bytes() if hosts_file.exists() else b""
hosts_file.write_bytes(original + _pin_block(sink_address, hosts).encode())
return original
def _restore_pins(hosts_file: Path, original: bytes) -> None:
hosts_file.write_bytes(original)
def _bind(sink_address: str, port: int) -> socket.socket:
listener = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
listener.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
listener.bind((sink_address, port))
listener.listen(128)
return listener
@dataclass(frozen=True, slots=True)
class _HitLog:
path: Path
_lock: threading.Lock
def record(self, *, port: int, peer: tuple[str, int]) -> None:
entry = json.dumps({"ts": time.time(), "port": port, "peer": list(peer)})
with self._lock:
with self.path.open("a", encoding="utf-8") as handle:
handle.write(entry + "\n")
def _serve_socket(listener: socket.socket, port: int, hits: _HitLog, stop: threading.Event) -> None:
while not stop.is_set():
try:
conn, peer = listener.accept()
except OSError:
return
hits.record(port=port, peer=(peer[0], peer[1]))
try:
conn.close()
except OSError:
pass
def serve(config: ServeConfig) -> int:
config.hits_file.write_text("", encoding="utf-8")
original_hosts = _install_pins(config.hosts_file, config.sink_address, config.hosts)
try:
listeners = tuple(_bind(config.sink_address, port) for port in config.ports)
except OSError as exc:
_restore_pins(config.hosts_file, original_hosts)
print(f"egress sentinel could not bind a sink: {exc}", file=sys.stderr)
return 1
stop = threading.Event()
hits = _HitLog(path=config.hits_file, _lock=threading.Lock())
threads = tuple(
threading.Thread(target=_serve_socket, args=(listener, port, hits, stop), daemon=True)
for listener, port in zip(listeners, config.ports)
)
for thread in threads:
thread.start()
def _handle(_signum: int, _frame: FrameType | None) -> None:
stop.set()
for listener in listeners:
try:
listener.close()
except OSError:
pass
signal.signal(signal.SIGTERM, _handle)
signal.signal(signal.SIGINT, _handle)
if config.pid_file is not None:
config.pid_file.write_text(str(os.getpid()), encoding="utf-8")
if config.ready_file is not None:
config.ready_file.write_text("ready", encoding="utf-8")
print(
f"egress sentinel up: pinned {', '.join(config.hosts)} to {config.sink_address} "
f"on port(s) {', '.join(str(p) for p in config.ports)}",
flush=True,
)
stop.wait()
_restore_pins(config.hosts_file, original_hosts)
if config.ready_file is not None and config.ready_file.exists():
config.ready_file.unlink()
if config.pid_file is not None and config.pid_file.exists():
config.pid_file.unlink()
return 0
def assert_empty(hits_file: Path) -> int:
if not hits_file.exists():
print(f"egress sentinel recorded no provider calls ({hits_file} absent): zero egress")
return 0
hits = [line for line in hits_file.read_text(encoding="utf-8").splitlines() if line.strip()]
if not hits:
print("egress sentinel recorded no provider calls: zero egress")
return 0
print(f"egress sentinel recorded {len(hits)} provider call(s); replay was not hermetic:", file=sys.stderr)
for line in hits:
print(f" {line}", file=sys.stderr)
return 1
def _serve_from_args(args: argparse.Namespace) -> int:
config = ServeConfig(
hosts=tuple(args.host),
sink_address=args.sink_address,
ports=tuple(args.port),
hits_file=Path(args.hits_file),
hosts_file=Path(args.hosts_file),
ready_file=Path(args.ready_file) if args.ready_file else None,
pid_file=Path(args.pid_file) if args.pid_file else None,
)
return serve(config)
def main(argv: tuple[str, ...]) -> int:
parser = argparse.ArgumentParser(description="count outbound provider calls during an e2e replay")
sub = parser.add_subparsers(dest="command", required=True)
serve_parser = sub.add_parser("serve", help="pin provider hosts and count connection attempts")
serve_parser.add_argument("--host", action="append", required=True, help="provider host to pin and watch")
serve_parser.add_argument("--sink-address", default="127.0.0.1")
serve_parser.add_argument("--port", action="append", type=int, default=None)
serve_parser.add_argument("--hits-file", required=True)
serve_parser.add_argument("--hosts-file", default="/etc/hosts")
serve_parser.add_argument("--ready-file", default=None)
serve_parser.add_argument("--pid-file", default=None)
assert_parser = sub.add_parser("assert-empty", help="exit non-zero if any provider call was recorded")
assert_parser.add_argument("--hits-file", required=True)
args = parser.parse_args(argv)
if args.command == "serve":
if args.port is None:
args.port = [443]
return _serve_from_args(args)
return assert_empty(Path(args.hits_file))
if __name__ == "__main__":
raise SystemExit(main(tuple(sys.argv[1:])))

55
.github/scripts/e2e_fetch_fixture_bundle.sh vendored Executable file
View file

@ -0,0 +1,55 @@
#!/usr/bin/env bash
set -euo pipefail
REPO="${1:-${GITHUB_REPOSITORY:?REPO required}}"
ARTIFACT_NAME="${2:-e2e-fixtures-bundle}"
BASE_BRANCH="${3:?base branch required}"
DEST_DIR="${4:?destination bundle dir required}"
: "${GH_TOKEN:?GH_TOKEN required to query and download artifacts}"
WORKDIR="$(mktemp -d)"
trap 'rm -rf "${WORKDIR}"' EXIT
echo "resolving newest non-expired '${ARTIFACT_NAME}' artifact on ${REPO}@${BASE_BRANCH}"
SELECTED="$(
gh api "repos/${REPO}/actions/artifacts" -X GET -f per_page=100 --paginate \
--jq ".artifacts[] | select(.name == \"${ARTIFACT_NAME}\" and .expired == false and .workflow_run.head_branch == \"${BASE_BRANCH}\") | {id, digest, created_at, run_id: .workflow_run.id, run_number: .workflow_run.run_number}" \
| jq -s 'sort_by(.created_at) | reverse | .[0] // empty'
)"
if [[ -z "${SELECTED}" ]]; then
echo "no usable '${ARTIFACT_NAME}' artifact on ${BASE_BRANCH}: the last record run produced none (a red Saturday), so there is nothing fresh to replay; failing loudly instead of replaying a stale bundle" >&2
exit 1
fi
RUN_ID="$(echo "${SELECTED}" | jq -r '.run_id')"
RUN_NUMBER="$(echo "${SELECTED}" | jq -r '.run_number')"
ARTIFACT_ID="$(echo "${SELECTED}" | jq -r '.id')"
GH_DIGEST="$(echo "${SELECTED}" | jq -r '.digest // "unknown"')"
CREATED_AT="$(echo "${SELECTED}" | jq -r '.created_at')"
echo "pinned bundle: run #${RUN_NUMBER} (run_id=${RUN_ID}, artifact_id=${ARTIFACT_ID}), recorded ${CREATED_AT}, github digest ${GH_DIGEST}"
gh run download "${RUN_ID}" --repo "${REPO}" -n "${ARTIFACT_NAME}" -D "${WORKDIR}"
TARBALL="$(find "${WORKDIR}" -name '*.tar.gz' -type f | head -n 1)"
if [[ -z "${TARBALL}" ]]; then
echo "downloaded artifact contained no tarball" >&2
exit 1
fi
SIDECAR="${TARBALL}.sha256"
if [[ ! -f "${SIDECAR}" ]]; then
echo "downloaded artifact has no ${SIDECAR}: cannot verify the bundle digest" >&2
exit 1
fi
echo "verifying bundle against its recorded sha256 digest"
( cd "$(dirname "${TARBALL}")" && sha256sum -c "$(basename "${SIDECAR}")" )
mkdir -p "${DEST_DIR}"
tar xzf "${TARBALL}" -C "${DEST_DIR}"
echo "extracted bundle into ${DEST_DIR}"
python3 -c "import json,sys; m=json.load(open(sys.argv[1])); print(' recorded_at', m['recorded_at'], 'harness', m['harness_version'], 'format_version', m['format_version'])" "${DEST_DIR}/manifest.json"

36
.github/scripts/e2e_pack_fixture_bundle.sh vendored Executable file
View file

@ -0,0 +1,36 @@
#!/usr/bin/env bash
set -euo pipefail
if [[ $# -ne 2 ]]; then
echo "usage: $0 <bundle-dir> <out-tarball>" >&2
exit 2
fi
BUNDLE_DIR="$1"
OUT_TARBALL="$2"
MANIFEST="${BUNDLE_DIR}/manifest.json"
if [[ ! -f "${MANIFEST}" ]]; then
echo "no ${MANIFEST}: refusing to publish a bundle with no manifest (record produced nothing)" >&2
exit 1
fi
echo "packing fixture bundle from ${BUNDLE_DIR}"
python3 -c "import json,sys; m=json.load(open(sys.argv[1])); print(' format_version', m['format_version'], 'recorded_at', m['recorded_at'], 'harness', m['harness_version'])" "${MANIFEST}"
TEST_DIRS=$(find "${BUNDLE_DIR}" -mindepth 1 -maxdepth 1 -type d | wc -l | tr -d ' ')
if [[ "${TEST_DIRS}" -eq 0 ]]; then
echo "bundle at ${BUNDLE_DIR} has a manifest but no recorded interactions; refusing to publish an empty bundle" >&2
exit 1
fi
echo " ${TEST_DIRS} recorded test director(ies)"
mkdir -p "$(dirname "${OUT_TARBALL}")"
tar czf "${OUT_TARBALL}" -C "${BUNDLE_DIR}" .
OUT_DIR="$(cd "$(dirname "${OUT_TARBALL}")" && pwd)"
OUT_BASE="$(basename "${OUT_TARBALL}")"
( cd "${OUT_DIR}" && sha256sum "${OUT_BASE}" > "${OUT_BASE}.sha256" )
echo "wrote ${OUT_TARBALL} ($(du -h "${OUT_TARBALL}" | cut -f1)) and ${OUT_BASE}.sha256"
cat "${OUT_DIR}/${OUT_BASE}.sha256"

View file

@ -149,7 +149,7 @@ jobs:
--reruns "${RERUNS}" \
--reruns-delay 1 \
--durations=20 \
--cov=./litellm \
--cov=./litellm --cov=./enterprise/litellm_enterprise \
--cov-report=xml:coverage.xml \
--cov-config=pyproject.toml
else
@ -161,7 +161,7 @@ jobs:
--reruns-delay 1 \
--dist="${DIST}" \
--durations=20 \
--cov=./litellm \
--cov=./litellm --cov=./enterprise/litellm_enterprise \
--cov-report=xml:coverage.xml \
--cov-config=pyproject.toml
fi

View file

@ -83,6 +83,24 @@ jobs:
if: steps.changes.outputs.relevant == 'true'
run: uv run --no-sync prisma generate --schema litellm/proxy/schema.prisma
- name: Regenerate the lazy OpenAPI snapshot
if: steps.changes.outputs.relevant == 'true'
run: uv run --no-sync python -m litellm.proxy._lazy_openapi_snapshot
- name: Fail if the lazy OpenAPI snapshot is stale
if: steps.changes.outputs.relevant == 'true'
run: |
if ! git diff --exit-code -- litellm/proxy/_lazy_openapi_snapshot.json; then
echo "::error file=litellm/proxy/_lazy_openapi_snapshot.json::The lazy OpenAPI snapshot is out of sync with the lazily loaded routes."
echo ""
echo "A lazily loaded route or model changed without regenerating the snapshot that /openapi.json serves for unloaded features."
echo "To fix, run from the repo root:"
echo " uv run python -m litellm.proxy._lazy_openapi_snapshot"
echo "then run npm run gen:api from ui/litellm-dashboard and commit both files."
exit 1
fi
echo "_lazy_openapi_snapshot.json is in sync with the lazily loaded routes."
- name: Set up Node.js
if: steps.changes.outputs.relevant == 'true'
uses: actions/setup-node@a0853c24544627f65ddf259abe73b1d18a591444 # v5.0.0

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:

237
.github/workflows/e2e_record_replay.yml vendored Normal file
View file

@ -0,0 +1,237 @@
name: "E2E Record and Replay"
on:
schedule:
- cron: "0 8 * * 6"
- cron: "0 8 * * 1-5"
workflow_dispatch:
inputs:
mode:
description: "record (hits real providers and publishes a fresh bundle) or replay (bundle only, zero provider egress)"
type: choice
options:
- record
- replay
default: record
permissions:
contents: read
jobs:
record:
name: "Record the e2e suite against real providers"
if: >-
(github.event_name != 'schedule' || github.repository == 'BerriAI/litellm') &&
(github.event.schedule == '0 8 * * 6' ||
(github.event_name == 'workflow_dispatch' && github.event.inputs.mode == 'record'))
runs-on: ubuntu-latest
timeout-minutes: 45
services:
postgres:
image: postgres:16.6
env:
POSTGRES_USER: llmproxy
POSTGRES_PASSWORD: dbpassword9090
POSTGRES_DB: litellm
ports:
- 5432:5432
options: >-
--health-cmd "pg_isready -U llmproxy"
--health-interval 5s
--health-timeout 5s
--health-retries 10
env:
DATABASE_URL: postgresql://llmproxy:dbpassword9090@localhost:5432/litellm
LITELLM_MASTER_KEY: sk-e2e-record-replay
LITELLM_LOCAL_MODEL_COST_MAP: "True"
OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
ANTHROPIC_API_KEY: ${{ secrets.ANTHROPIC_API_KEY }}
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: Cache the Rust build
uses: ./.github/actions/cache-cargo-build
- name: Install dependencies
run: |
.github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra proxy
- name: Cache Prisma binaries
uses: ./.github/actions/cache-prisma-binaries
- name: Generate Prisma client
run: |
uv run --no-sync prisma generate --schema litellm/proxy/schema.prisma
- name: Start the proxy
run: |
nohup uv run --no-sync litellm --config tests/e2e/gateway/record_replay_ci_config.yml --port 4000 > proxy.log 2>&1 &
for _ in $(seq 1 90); do
if curl -fs http://localhost:4000/health/liveliness > /dev/null; then
exit 0
fi
sleep 2
done
echo "proxy never became live"
tail -n 100 proxy.log
exit 1
- name: Record the replayable e2e lane
env:
E2E_FIXTURE_MODE: record
run: |
uv run --no-sync pytest tests/e2e -m replayable --reruns 0 -v --tb=short -rA
- name: Pack the fixture bundle
run: |
.github/scripts/e2e_pack_fixture_bundle.sh tests/e2e/.fixtures "${RUNNER_TEMP}/bundle/e2e-fixtures.tar.gz"
- name: Publish the fixture bundle
uses: actions/upload-artifact@4cec3d8aa04e39d1a68397de0c4cd6fb9dce8ec1 # v4.6.1
with:
name: e2e-fixtures-bundle
path: |
${{ runner.temp }}/bundle/e2e-fixtures.tar.gz
${{ runner.temp }}/bundle/e2e-fixtures.tar.gz.sha256
if-no-files-found: error
retention-days: 30
- name: Show proxy log on failure
if: failure()
run: tail -n 300 proxy.log
replay:
name: "Replay the e2e suite from the pinned bundle with zero egress"
if: >-
(github.event_name != 'schedule' || github.repository == 'BerriAI/litellm') &&
(github.event.schedule == '0 8 * * 1-5' ||
(github.event_name == 'workflow_dispatch' && github.event.inputs.mode == 'replay'))
runs-on: ubuntu-latest
timeout-minutes: 45
permissions:
contents: read
actions: read
services:
postgres:
image: postgres:16.6
env:
POSTGRES_USER: llmproxy
POSTGRES_PASSWORD: dbpassword9090
POSTGRES_DB: litellm
ports:
- 5432:5432
options: >-
--health-cmd "pg_isready -U llmproxy"
--health-interval 5s
--health-timeout 5s
--health-retries 10
env:
DATABASE_URL: postgresql://llmproxy:dbpassword9090@localhost:5432/litellm
LITELLM_MASTER_KEY: sk-e2e-record-replay
LITELLM_LOCAL_MODEL_COST_MAP: "True"
GH_TOKEN: ${{ github.token }}
OPENAI_API_KEY: sk-replay-must-never-reach-a-provider
ANTHROPIC_API_KEY: sk-ant-replay-must-never-reach-a-provider
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: Cache the Rust build
uses: ./.github/actions/cache-cargo-build
- name: Install dependencies
run: |
.github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra proxy
- name: Cache Prisma binaries
uses: ./.github/actions/cache-prisma-binaries
- name: Generate Prisma client
run: |
uv run --no-sync prisma generate --schema litellm/proxy/schema.prisma
- name: Fetch the pinned fixture bundle by digest
env:
BASE_BRANCH: ${{ github.ref_name }}
run: |
.github/scripts/e2e_fetch_fixture_bundle.sh \
"${GITHUB_REPOSITORY}" \
e2e-fixtures-bundle \
"${BASE_BRANCH}" \
tests/e2e/.fixtures
- name: Start the proxy
run: |
nohup uv run --no-sync litellm --config tests/e2e/gateway/record_replay_ci_config.yml --port 4000 > proxy.log 2>&1 &
for _ in $(seq 1 90); do
if curl -fs http://localhost:4000/health/liveliness > /dev/null; then
exit 0
fi
sleep 2
done
echo "proxy never became live"
tail -n 100 proxy.log
exit 1
- name: Start the egress sentinel
run: |
# shellcheck disable=SC2024 # the log redirect is deliberately the runner user's, so a later non-sudo cat can read it
sudo python3 .github/scripts/e2e_egress_sentinel.py serve \
--host api.openai.com \
--host api.anthropic.com \
--hits-file "${RUNNER_TEMP}/egress-hits.jsonl" \
--ready-file "${RUNNER_TEMP}/egress-ready" \
--pid-file "${RUNNER_TEMP}/egress.pid" \
> "${RUNNER_TEMP}/egress-sentinel.log" 2>&1 &
for _ in $(seq 1 30); do
if [[ -f "${RUNNER_TEMP}/egress-ready" ]]; then
cat "${RUNNER_TEMP}/egress-sentinel.log"
exit 0
fi
sleep 1
done
echo "egress sentinel never became ready"
cat "${RUNNER_TEMP}/egress-sentinel.log"
exit 1
- name: Replay the replayable e2e lane
env:
E2E_FIXTURE_MODE: replay
run: |
uv run --no-sync pytest tests/e2e -m replayable --reruns 0 -v --tb=short -rA
- name: Stop the egress sentinel and assert zero provider egress
if: always()
run: |
if [[ -f "${RUNNER_TEMP}/egress.pid" ]]; then
sudo kill -TERM "$(cat "${RUNNER_TEMP}/egress.pid")" 2>/dev/null || true
sleep 2
fi
python3 .github/scripts/e2e_egress_sentinel.py assert-empty --hits-file "${RUNNER_TEMP}/egress-hits.jsonl"
- name: Show proxy log on failure
if: failure()
run: tail -n 300 proxy.log

View file

@ -23,6 +23,8 @@ on:
- tests/proxy_migration_tests/**
- uv.lock
- ui/litellm-dashboard/package-lock.json
- ui/Dockerfile
- ui/nginx.conf
- .github/workflows/image-scan.yml
schedule:
- cron: "41 6 * * *"
@ -185,6 +187,35 @@ jobs:
python -m pip install "pytest==9.0.3"
python -m pytest tests/proxy_migration_tests/test_component_image_serves_offline.py -v
ui-image:
name: ui-image
runs-on: ubuntu-latest
if: >-
github.event_name != 'pull_request' ||
github.event.pull_request.head.repo.full_name == github.repository
timeout-minutes: 30
permissions:
contents: read
steps:
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
persist-credentials: false
- name: Build UI image
run: docker build -f ui/Dockerfile -t litellm-ui-scan:${{ github.sha }} .
- name: Set up Python
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
with:
python-version: "3.12"
- name: Verify the UI serves offline as an arbitrary uid with a read-only root fs
env:
LITELLM_IMAGE: litellm-ui-scan:${{ github.sha }}
run: |
python -m pip install "pytest==9.0.3"
python -m pytest tests/proxy_migration_tests/test_ui_image_serves_offline.py -v
backend-image:
name: backend-image
runs-on: ubuntu-latest

View file

@ -87,11 +87,20 @@ jobs:
run: |
uv pip uninstall pytest-retry || true
# Ends before the job's own deadline so a run that outlasts the budget is
# still followed by the report and upload steps. mutmut saves after every
# mutant result, to mutants/<source path>.meta, so an interrupted run
# still scores the mutants it finished and export-cicd-stats can read
# them; a cancelled job skips those steps and publishes nothing at all.
- name: Run mutmut
timeout-minutes: 300
env:
# Make the mutants/ sandbox win over site-packages on sys.path so the
# trampolined files are imported instead of the installed copy.
PYTHONPATH: ${{ github.workspace }}/mutants
# Without this mutmut finds no covered lines and generates 0 mutants.
# See the file itself for why.
COVERAGE_RCFILE: ${{ github.workspace }}/.github/mutmut-coverage.rc
run: |
set -o pipefail
mkdir -p mutants
@ -130,6 +139,7 @@ jobs:
mutmut-run.log
mutants/mutmut-stats.json
mutants/mutmut-cicd-stats.json
mutants/**/*.meta
mutants/litellm/proxy/management_endpoints/**/*.py
if-no-files-found: warn
retention-days: 14

View file

@ -0,0 +1,68 @@
name: Sync Together AI model registry
on:
schedule:
- cron: "30 6 * * *"
workflow_dispatch:
permissions:
contents: write
pull-requests: write
jobs:
sync_together_ai_models:
if: github.repository == 'BerriAI/litellm'
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
ref: litellm_internal_staging
persist-credentials: false
- name: Set up uv
uses: ./.github/actions/setup-uv-with-retries
with:
version: "0.10.9"
- name: Look for an already-open sync PR
id: existing
run: |
open_pr="$(gh pr list --repo "$GITHUB_REPOSITORY" --state open --limit 1000 --json headRefName \
--jq '[.[].headRefName | select(startswith("litellm_together_registry_sync_"))] | first // empty')"
echo "open_pr=$open_pr" >> "$GITHUB_OUTPUT"
if [ -n "$open_pr" ]; then
echo "An open sync PR already exists on branch $open_pr; skipping this run."
fi
env:
GH_TOKEN: ${{ secrets.GH_TOKEN || github.token }}
- name: Run the sync
if: steps.existing.outputs.open_pr == ''
run: |
uv run --frozen python scripts/sync_together_ai_models.py --write --pr-body-file "$RUNNER_TEMP/pr_body.md"
env:
TOGETHER_API_KEY: ${{ secrets.TOGETHER_API_KEY }}
- name: Regenerate the JSON schema
if: steps.existing.outputs.open_pr == ''
run: |
uv run --frozen python ci_cd/generate_model_prices_schema.py
- name: Create a pull request when the registry changed
if: steps.existing.outputs.open_pr == ''
run: |
if git diff --quiet; then
echo "Registry already in sync; no PR needed."
exit 0
fi
branch="litellm_together_registry_sync_$(date +'%Y-%m-%d')"
git config user.name "github-actions[bot]"
git config user.email "41898282+github-actions[bot]@users.noreply.github.com"
git checkout -b "$branch"
git add model_prices_and_context_window.json \
litellm/model_prices_and_context_window_backup.json \
model_prices_and_context_window.schema.json
git commit -m "feat(models): sync together_ai model registry $(date +'%Y-%m-%d')"
gh auth setup-git
git push origin "$branch"
gh pr create --title "feat(models): sync together_ai model registry" \
--body-file "$RUNNER_TEMP/pr_body.md" \
--head "$branch" \
--base litellm_internal_staging
env:
GH_TOKEN: ${{ secrets.GH_TOKEN || github.token }}

View file

@ -131,6 +131,9 @@ jobs:
- name: check_e2e_no_raw_requests
run: uv run --no-sync python ./tests/code_coverage_tests/check_e2e_no_raw_requests.py
- name: check_migrations_no_data_rewrites
run: uv run --no-sync python ./tests/code_coverage_tests/check_migrations_no_data_rewrites.py
- name: memory_test
run: uv run --no-sync python ./tests/code_coverage_tests/memory_test.py

View file

@ -67,6 +67,17 @@ jobs:
with:
version: "0.10.9"
- name: Cache uv dependencies
if: steps.changes.outputs.decision != 'skip'
uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0
with:
path: |
~/.cache/uv
.venv
key: ${{ runner.os }}-uv-lint-${{ hashFiles('uv.lock') }}
restore-keys: |
${{ runner.os }}-uv-lint-
- name: Clean Python cache
if: steps.changes.outputs.decision != 'skip'
run: |

View file

@ -60,4 +60,4 @@ jobs:
- name: Run MCP tests
if: steps.changes.outputs.decision != 'skip'
run: |
uv run --no-sync pytest tests/mcp_tests -x -vv -n 4 --cov=./litellm --cov-report=xml --durations=5
uv run --no-sync pytest tests/mcp_tests -x -vv -n 4 --cov=./litellm --cov=./enterprise/litellm_enterprise --cov-report=xml --durations=5

145
.github/workflows/test-postgres.yml vendored Normal file
View file

@ -0,0 +1,145 @@
name: "Postgres Tests"
on:
pull_request:
branches:
- main
- litellm_internal_staging
- litellm_oss_staging
- "litellm_**"
push:
branches:
- main
- litellm_internal_staging
workflow_dispatch:
permissions:
contents: read
concurrency:
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}
cancel-in-progress: ${{ github.event_name == 'pull_request' }}
jobs:
postgres:
name: ${{ matrix.shard }}
runs-on: ubuntu-latest
timeout-minutes: ${{ matrix.job-timeout-minutes }}
permissions:
contents: read
services:
postgres:
image: postgres:16@sha256:e17e86066e5ef83e0952a9347f5c792b7ece00972e2aa787a6986f471b3dd3d5
env:
POSTGRES_USER: postgres
POSTGRES_PASSWORD: postgres
POSTGRES_DB: litellm_test
ports:
- 5432:5432
options: >-
--health-cmd pg_isready
--health-interval 10s
--health-timeout 5s
--health-retries 10
strategy:
fail-fast: false
matrix:
include:
- shard: proxy-behavior
test-path: "tests/proxy_behavior"
seed: db-push
workers: 0
timeout-minutes: 25
job-timeout-minutes: 50
- shard: proxy-security
test-path: "tests/proxy_security_tests"
seed: db-push
workers: 0
timeout-minutes: 15
job-timeout-minutes: 40
- shard: schema-migration
test-path: "tests/proxy_migration_tests"
seed: none
workers: 0
timeout-minutes: 20
job-timeout-minutes: 45
env:
DATABASE_URL: "postgresql://postgres:postgres@localhost:5432/litellm_test"
steps:
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
timeout-minutes: 3
with:
persist-credentials: false
- name: Detect relevant changes
id: changes
timeout-minutes: 2
uses: ./.github/actions/detect-changes
- name: Set up Python
if: steps.changes.outputs.decision != 'skip'
timeout-minutes: 3
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
with:
python-version: "3.12"
- name: Set up uv
if: steps.changes.outputs.decision != 'skip'
timeout-minutes: 3
uses: ./.github/actions/setup-uv-with-retries
with:
version: "0.10.9"
- name: Cache uv dependencies
if: steps.changes.outputs.decision != 'skip'
timeout-minutes: 5
uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0
with:
path: |
~/.cache/uv
.venv
key: ${{ runner.os }}-uv-postgres-${{ hashFiles('uv.lock') }}
restore-keys: |
${{ runner.os }}-uv-postgres-
- name: Install dependencies
if: steps.changes.outputs.decision != 'skip'
timeout-minutes: 12
run: |
.github/scripts/uv_sync_with_retries.sh --frozen --all-groups --all-extras
- name: Cache Prisma binaries
if: steps.changes.outputs.decision != 'skip'
timeout-minutes: 3
uses: ./.github/actions/cache-prisma-binaries
- name: Generate Prisma client
if: steps.changes.outputs.decision != 'skip'
timeout-minutes: 5
run: |
uv run --no-sync prisma generate --schema litellm/proxy/schema.prisma
- name: Seed database schema
if: steps.changes.outputs.decision != 'skip' && matrix.seed != 'none'
timeout-minutes: 10
run: |
uv run --no-sync prisma db push --schema litellm/proxy/schema.prisma --accept-data-loss
- name: Run tests
if: steps.changes.outputs.decision != 'skip'
timeout-minutes: ${{ matrix.timeout-minutes }}
env:
TEST_PATH: ${{ matrix.test-path }}
WORKERS: ${{ matrix.workers }}
run: |
if [ "${WORKERS}" = "0" ]; then
uv run --no-sync pytest ${TEST_PATH:?} -vv --tb=short --durations=10
else
uv run --no-sync pytest ${TEST_PATH:?} -vv --tb=short --durations=10 -n "${WORKERS}"
fi

View file

@ -114,4 +114,4 @@ jobs:
- name: Audit provider endpoints against the schema
working-directory: terraform/provider
run: go run ./tools/endpointaudit -provider-dir ./litellm -spec "${RUNNER_TEMP}/openapi.json"
run: go run ./tools/endpointaudit -provider-dir ./litellm -spec "${RUNNER_TEMP}/openapi.json" -coverage-allowlist ./tools/endpointaudit/coverage_allowlist.txt

View file

@ -141,6 +141,7 @@ jobs:
test-path: >-
tests/test_litellm/proxy/analytics_endpoints
tests/test_litellm/proxy/management_endpoints
tests/test_litellm/proxy/list_api
tests/test_litellm/proxy/memory
tests/test_litellm/proxy/guardrails
tests/test_litellm/proxy/management_helpers
@ -164,11 +165,12 @@ jobs:
tests/test_litellm/proxy/public_endpoints
tests/test_litellm/proxy/prompts
tests/test_litellm/proxy/rag_endpoints
tests/test_litellm/proxy/rerank_endpoints
tests/test_litellm/proxy/realtime_endpoints
tests/test_litellm/proxy/ui_crud_endpoints
tests/test_litellm/proxy/config_resolvers
tests/test_litellm/proxy/utils
workers: 2
workers: 4
reruns: 2
timeout-minutes: 20
job-timeout-minutes: 60
@ -195,11 +197,40 @@ jobs:
tests/test_litellm/proxy/types_utils
tests/test_litellm/proxy/logging_endpoints
tests/test_litellm/proxy/test_*.py
workers: 4
reruns: 2
timeout-minutes: 20
job-timeout-minutes: 60
- shard: caching-local
artifact-name: caching-local
test-path: >-
tests/local_testing/test_cache_preset_key.py
tests/local_testing/test_caching_handler.py
tests/local_testing/test_prompt_caching.py
tests/local_testing/test_responses_stream_cache_keys.py
tests/local_testing/test_unit_test_caching.py
workers: 2
reruns: 2
timeout-minutes: 20
job-timeout-minutes: 60
- shard: proxy-extras
artifact-name: proxy-extras
test-path: "tests/litellm-proxy-extras"
workers: 2
reruns: 2
timeout-minutes: 20
job-timeout-minutes: 60
- shard: enterprise-package
artifact-name: enterprise-package
test-path: "tests/enterprise"
workers: 4
reruns: 2
timeout-minutes: 20
job-timeout-minutes: 60
- shard: responses-caching-types
artifact-name: responses-caching-types
test-path: >-

2
.gitignore vendored
View file

@ -3,6 +3,8 @@
tests/e2e/.fixtures/
.venv-typecheck
.venv_policy_test
.venv-mutmut
mutants/
.env
.claude
CLAUDE.local.md

View file

@ -37,13 +37,14 @@ If you're resolving a linear ticket, in the "## Linear ticket" section of the PR
Never use `pytest` commands or the like as "Screenshots / Proof of Fix". We prefer curl'ing a live proxy instance running on localhost:4000 (I like to run it with `python litellm/proxy/proxy_cli.py --config litellm/proxy/dev_config.yaml --detailed_debug --reload --use_v2_migration_resolver 2>&1 | tee litellm.log`; the Admin UI dev server is `npm run dev` in `ui/litellm-dashboard`, served on port 3000) and showing both the command run and the output. Also, it should hit real LLM provider APIs, not mocks, and cost real $$$ because that is the most realistic test. The proof of fix should be exactly what the end user / customer would see / do. The run logs in PR #27703 is a prime example of how to do it (not a huge fan of using a python test script that future me and the team will have no visibility into; I prefer just curl commands or a short list of bash commands (e.g., using `for`)). If it's a UI thing, just tell me which URLs to go to (e.g., http://localhost:4000/ui/?page=logs), where to click, what fields to fill out, etc. along with the other commands to run in an ordered list, and I'll do it myself and post the screenshots after you make the PR
If you ever make public-facing PR descriptions, comments, issues, commit messages, etc., always follow these guidelines to sound less AI-y:
If you ever write any human-facing text (pull requests, issues, commit messages, discussion posts, github comments, release notes, docs, etc.), always follow these guidelines to sound less AI-y:
- don't use emojis
- don't use "—". Instead, reach for ",", ".", conjunction words, ":", ";", etc. in descending order of preference: vary among them, weighted toward the front of the list, and skip "," where it would cause a comma splice or the sentence is getting long. Overusing any one of them, ";" especially, also feels AI-y. A word cap does not penalize you for adding more sentences: when writing under tight word budgets, prefer a period split or a conjunction over ";", and keep to at most one ";" per message
- don't use the pattern "It's not X, it's Y", "You're not X, you're Y", etc.
- don't use bulleted or numbered lists unless it would be nonsensical not to. Instead, prefer prose
- unless explicitly asked, don't use bulleted or numbered lists unless it would be nonsensical not to. Instead, prefer prose
- don't add a trailing "." at the end of paragraphs (just like this file). That means every paragraph, not just the last one (of the markdown file, PR description, GitHub comment, etc.). Rule of thumb: if you're adding new line(s) before the next sentence, don't add a "."
- don't use →. Instead, prefer not to use arrows, and if need be, use -> instead
- use plain, simple, everyday engineering language: the common phrase engineers actually say over rare compact phrasing, in grammatically complete sentences. When explicitly asked to use bullets or ordered lists and structure legitimately helps the reader, prefer nested bullets (any depth is fine) over dense lines in a flat structure
Don't hesitate to use values in .env to get needed API keys and other secrets, as long as you never add them to conversation history, commit them, or include them in GitHub issues / PRs
@ -65,6 +66,8 @@ Commit and push your work when you're done without asking
When referencing or running models (coding, QA'ing, writing docs, writing tests, etc.), use the latest model in that model family unless otherwise specified; treat your training knowledge, memories, configs, and tests as stale, and determine the family's latest with model_prices_and_context_window.json or the web
Always pull before starting any work. The checkout or worktree may be sitting on a stale branch
If you're an internal contributor, when creating a new PR, the typical flow is to branch off litellm_internal_staging and create a branch prefixed with litellm_. Do not create a branch prefixed with claude/ and generally do not have / in your branch names
Do not add `Co-Authored-By: Claude` or any Claude attribution to commit messages. Never use a `claude/` prefix or put a `/` in a branch name. Do not add "Generated with Claude Code" (or any similar attribution) to PR descriptions or comments. Do not create a new PR/branch off the existing PR to fix/add something that is related and could've just been committed directly to the existing PR's branch
@ -79,6 +82,8 @@ Do not put names of customers or customer company names in code, PR descriptions
CI supply-chain safety: Never pipe a remote script into a shell (`curl ... | bash`, `wget ... | sh`); download the artifact to a file, verify its SHA-256 checksum, then install. Pin every external tool to a specific version with a full URL (not `latest` or `stable`). Verify checksums for all downloaded binaries, using the provider's official `.sha256` / `.sha256sum` sidecar when available. These rules apply to every download in CI
Prisma migrations apply synchronously at proxy boot, before it serves traffic, so a migration must only change schema, never rewrite rows. No `UPDATE`, `DELETE` or `MERGE`, and no `INSERT ... SELECT`: on a spend-log-sized table any of those is minutes of downtime plus a doubled heap that plain autovacuum won't give back. `tests/code_coverage_tests/check_migrations_no_data_rewrites.py` enforces this. When a rewrite is genuinely bounded and has to ship inside the migration, mark the statement `-- data-migration-ok: <what bounds it>`
Follow these coding conventions for new/updated code (a three-line fix in a legacy file shouldn't trigger huge drive-by refactors):
- Composition over inheritance

View file

@ -1,36 +1,36 @@
{
"reportAny": {
"limit": 19955
"limit": 17270
},
"reportArgumentType": {
"limit": 2566
"limit": 2538
},
"reportAssignmentType": {
"limit": 320
"limit": 319
},
"reportAttributeAccessIssue": {
"limit": 488
"limit": 480
},
"reportCallIssue": {
"limit": 114
"limit": 112
},
"reportConstantRedefinition": {
"limit": 40
},
"reportDeprecated": {
"limit": 213
"limit": 212
},
"reportDuplicateImport": {
"limit": 19
},
"reportExplicitAny": {
"limit": 6049
"limit": 5485
},
"reportFunctionMemberAccess": {
"limit": 7
},
"reportGeneralTypeIssues": {
"limit": 154
"limit": 101
},
"reportIncompatibleMethodOverride": {
"limit": 56
@ -45,7 +45,7 @@
"limit": 35
},
"reportInvalidTypeForm": {
"limit": 35
"limit": 34
},
"reportInvalidTypeVarUse": {
"limit": 2
@ -54,10 +54,10 @@
"limit": 0
},
"reportMissingParameterType": {
"limit": 5663
"limit": 5658
},
"reportMissingTypeArgument": {
"limit": 15555
"limit": 15425
},
"reportMissingTypeStubs": {
"limit": 40
@ -72,7 +72,7 @@
"limit": 0
},
"reportOptionalMemberAccess": {
"limit": 1061
"limit": 1055
},
"reportOptionalOperand": {
"limit": 0
@ -84,7 +84,7 @@
"limit": 56
},
"reportPrivateUsage": {
"limit": 1822
"limit": 1808
},
"reportRedeclaration": {
"limit": 8
@ -93,37 +93,37 @@
"limit": 213
},
"reportTypedDictNotRequiredAccess": {
"limit": 26
"limit": 25
},
"reportUndefinedVariable": {
"limit": 0
},
"reportUnknownArgumentType": {
"limit": 44655
"limit": 44526
},
"reportUnknownLambdaType": {
"limit": 109
},
"reportUnknownMemberType": {
"limit": 39011
"limit": 38721
},
"reportUnknownParameterType": {
"limit": 19885
"limit": 19778
},
"reportUnknownVariableType": {
"limit": 30569
"limit": 30290
},
"reportUnnecessaryCast": {
"limit": 117
},
"reportUnnecessaryComparison": {
"limit": 699
"limit": 697
},
"reportUnnecessaryContains": {
"limit": 5
},
"reportUnnecessaryIsInstance": {
"limit": 836
"limit": 829
},
"reportUntypedBaseClass": {
"limit": 0
@ -135,12 +135,12 @@
"limit": 21
},
"reportUnusedFunction": {
"limit": 139
"limit": 138
},
"reportUnusedImport": {
"limit": 545
"limit": 543
},
"reportUnusedVariable": {
"limit": 146
"limit": 139
}
}

View file

@ -73,6 +73,11 @@ ARRAY_KEYS: dict[str, JsonSchema] = {
"description": "Output modalities the model can produce.",
"items": {"type": "string", "enum": ["text", "image", "audio", "video", "code"]},
},
"reasoning_effort_levels": {
"type": "array",
"description": "Exact reasoning_effort levels this deployment accepts; wins over supports_* flags.",
"items": {"type": "string", "enum": ["none", "minimal", "low", "medium", "high", "xhigh", "max"]},
},
"supported_regions": {
"type": "array",
"description": "Cloud regions the model is available in ('global' or region ids).",
@ -157,6 +162,9 @@ COST_DESCRIPTIONS: dict[str, str] = {
"input_cost_per_token": "USD per prompt token.",
"output_cost_per_token": "USD per generated token.",
"output_cost_per_reasoning_token": "USD per reasoning/thinking token, when billed separately.",
"google_maps_grounding_cost_per_query": (
"USD per Grounding with Google Maps request; billed per query or per prompt per web_search_billing_unit."
),
"cache_creation_input_token_cost": "USD per token written to the provider's prompt cache.",
"cache_read_input_token_cost": "USD per prompt token served from the provider's prompt cache.",
"input_cost_per_token_batches": "USD per prompt token via the provider's batch API.",
@ -212,6 +220,15 @@ def string_key_schemas(modes: tuple) -> dict[str, JsonSchema]:
"description": "Highest reasoning effort the Bedrock output_config accepts for this model.",
"enum": ["low", "medium", "high", "max", "xhigh"],
},
"default_reasoning_effort": {
"type": "string",
"description": (
"Reasoning effort the provider applies when the request omits reasoning_effort. "
"Gates whether a non-default temperature or the top_p/logprobs sampling params are "
"accepted, which hold only when the effort resolves to 'none'."
),
"enum": ["none", "minimal", "low", "medium", "high", "xhigh"],
},
"comment": STRING,
"audio_transcription_config": STRING,
}

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:

View file

@ -10,6 +10,11 @@
-- partitioned, so existing installs are unaffected until you run this.
--
-- IMPORTANT
-- * After partitioning, `prisma db push` (including the proxy's
-- --use_prisma_db_push startup mode) is NOT supported: it tries to rewrite
-- the primary key back to ("request_id"), which Postgres rejects on a
-- partitioned table. The proxy detects this and exits with guidance.
-- Use the default startup path (`prisma migrate deploy`) instead.
-- * Test on a staging copy first and take a backup.
-- * Postgres cannot convert a populated table to partitioned in place, so this
-- renames the old table aside and creates a fresh partitioned table.

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, Optional
#### 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.
@ -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())

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, Any, Dict, Final, List, Literal, Optional, Tuple, cast
from litellm._logging import verbose_proxy_logger
from litellm._uuid import uuid
@ -14,6 +15,8 @@ from litellm.constants import (
)
if TYPE_CHECKING:
from prisma import models as prisma_models
from litellm.integrations.prometheus import PrometheusLogger
from litellm.proxy._types import LiteLLM_ManagedObjectTable
from litellm.proxy.utils import PrismaClient, ProxyLogging
@ -351,7 +354,7 @@ class CheckBatchCost:
return isinstance(error, (NotFoundError, openai.NotFoundError)) and output_file_id in str(error)
async def _finalize_unbilled_terminal_job(
self, job: "LiteLLM_ManagedObjectTable", response: "LiteLLMBatch"
self, job: "prisma_models.LiteLLM_ManagedObjectTable", response: "LiteLLMBatch"
) -> None:
"""Persist a terminal batch that has nothing billable, converting any raw
provider file ids to managed ids, and take it out of the poll page."""
@ -624,6 +627,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,
)
@ -759,16 +763,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",
@ -800,9 +821,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)

View file

@ -1,6 +1,8 @@
"""
Polls LiteLLM_ManagedObjectTable to check if the response is complete.
Cost tracking is handled automatically by the get-responses call.
Cost tracking is handled by the get-responses call, which prices normally only because the
poll stamps itself with BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN; user-facing reads of the
same route are non-inference and free.
"""
from datetime import datetime, timedelta, timezone
@ -9,12 +11,14 @@ from typing import TYPE_CHECKING, Dict, Optional, cast
import litellm
from litellm._logging import verbose_proxy_logger
from litellm.constants import (
INTERNAL_CALL_ORIGIN_METADATA_KEY,
MANAGED_OBJECT_STALENESS_CUTOFF_DAYS,
MAX_OBJECTS_PER_POLL_CYCLE,
STALE_OBJECT_CLEANUP_BATCH_SIZE,
)
from litellm.responses.utils import ResponsesAPIRequestUtils
from litellm.types.llms.openai import ResponsesAPIResponse
from litellm.types.utils import BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN
if TYPE_CHECKING:
from litellm.proxy.utils import PrismaClient, ProxyLogging
@ -113,7 +117,8 @@ class CheckResponsesCost:
Check if background responses are complete and track their cost.
- Get all status="queued" or "in_progress" and file_purpose="response" jobs
- Query the provider to check if response is complete
- Cost is automatically tracked by the get-responses call
- Cost is tracked by the get-responses call, billed because the poll is stamped
with BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN
- Mark responses in a terminal state as complete in the database
"""
try:
@ -153,6 +158,7 @@ class CheckResponsesCost:
# Prepare metadata with model information for cost tracking
litellm_metadata = {
"user_api_key_user_id": job.created_by or "default-user-id",
INTERNAL_CALL_ORIGIN_METADATA_KEY: BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN,
}
# Add model information if available

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,
@ -222,7 +223,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 +239,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 +343,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 +474,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 +534,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,

View file

@ -12,9 +12,10 @@ Endpoints for /project operations
import json
from collections.abc import Sequence
from typing import TYPE_CHECKING
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,37 +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]:
@ -205,6 +219,114 @@ def _check_team_project_limits(
)
def _project_models_missing_positive_quota(
models: list[str] | None,
rpm_limits: Mapping[str, object] | None,
tpm_limits: Mapping[str, object] | None,
) -> list[str]:
"""Return the models that lack a positive `rpm` AND `tpm` quota.
A valid quota is a positive integer; null, zero, and negative are rejected
because downstream rate limiters treat a non-positive limit as immediately
exhausted (every request blocked).
"""
def _is_positive(value: object) -> bool:
return isinstance(value, int) and not isinstance(value, bool) and value > 0
rpm = rpm_limits or {}
tpm = tpm_limits or {}
return [model for model in (models or []) if not _is_positive(rpm.get(model)) or not _is_positive(tpm.get(model))]
def _router_access_group_names(llm_router: "Router | None") -> frozenset[str]:
return frozenset(llm_router.get_model_access_groups()) if llm_router is not None else frozenset()
def _project_models_expanding_at_request_time(
models: Sequence[str] | None, access_group_names: frozenset[str]
) -> tuple[str, ...]:
"""Entries project auth expands to many concrete models (`all-proxy-models`, `*` patterns,
access groups). The rate limiter looks quotas up by the exact requested model name, so a
quota keyed on one of these entries is never applied."""
return tuple(
model
for model in (models or ())
if model == SpecialModelNames.all_proxy_models.value or "*" in model or model in access_group_names
)
def _raise_on_project_models_expanding_at_request_time(
models: Sequence[str] | None, access_group_names: frozenset[str]
) -> None:
expanding: Final = _project_models_expanding_at_request_time(models, access_group_names)
if not expanding:
return
raise HTTPException(
status_code=400,
detail={
"error": f"models {list(expanding)} expand to multiple models at request time, so a per-model rpm/tpm quota cannot be enforced for them while 'enforce_project_model_quota' is enabled. List concrete model names instead."
},
)
def _raise_on_missing_project_model_quota(
data: NewProjectRequest | UpdateProjectRequest, access_group_names: frozenset[str] = frozenset()
) -> None:
"""Require a positive `rpm`/`tpm` quota for every model on project CREATE.
`model_rpm_limit`/`model_tpm_limit` are relocated into `metadata` by the request
model's `set_model_info` validator, so they are read from there.
Only invoked when `general_settings.enforce_project_model_quota` is enabled
(default off), so it is opt-in and does not change behavior for existing users.
"""
_raise_on_project_models_expanding_at_request_time(data.models, access_group_names)
metadata = data.metadata or {}
missing = _project_models_missing_positive_quota(
data.models, metadata.get("model_rpm_limit"), metadata.get("model_tpm_limit")
)
if not missing:
return
raise HTTPException(
status_code=400,
detail={
"error": f"models {missing} added to project without a positive rpm/tpm quota. Set a positive model_rpm_limit and model_tpm_limit for each model."
},
)
def _raise_on_missing_project_model_quota_on_update(
data: UpdateProjectRequest, existing_project: object, access_group_names: frozenset[str] = frozenset()
) -> None:
"""Require a positive `rpm`/`tpm` quota over the RESULTING state on project UPDATE.
`/project/update` replaces `models` and `metadata` when they are provided, so the
check runs on what the project WILL look like: a partial update that doesn't touch
models/quota keeps the existing values, while one that adds a model or clears a
model's quota must leave every resulting model with a positive limit.
Only invoked when `general_settings.enforce_project_model_quota` is enabled
(default off), so it is opt-in and does not change behavior for existing users.
"""
resulting_models = data.models if data.models is not None else (getattr(existing_project, "models", None) or [])
resulting_metadata = (
data.metadata if data.metadata is not None else (getattr(existing_project, "metadata", None) or {})
)
_raise_on_project_models_expanding_at_request_time(resulting_models, access_group_names)
missing = _project_models_missing_positive_quota(
resulting_models, resulting_metadata.get("model_rpm_limit"), resulting_metadata.get("model_tpm_limit")
)
if not missing:
return
raise HTTPException(
status_code=400,
detail={
"error": f"models {missing} would be left on the project without a positive rpm/tpm quota. Set a positive model_rpm_limit and model_tpm_limit for each model."
},
)
async def _create_budget_for_project(
data: NewProjectRequest,
user_id: str | None,
@ -219,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,
@ -242,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
@ -352,7 +472,9 @@ async def new_project(
```
"""
from litellm.proxy.proxy_server import (
general_settings,
litellm_proxy_admin_name,
llm_router,
premium_user,
prisma_client,
)
@ -399,6 +521,10 @@ async def new_project(
data=data,
)
# Opt-in (default off): require rpm/tpm for every model added to the project.
if general_settings.get("enforce_project_model_quota", False):
_raise_on_missing_project_model_quota(data, _router_access_group_names(llm_router))
# Check if user has permission to create projects for this team
# only team admins can create projects for their team
has_permission = await _check_user_permission_for_project(
@ -470,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},
)
@ -538,7 +662,9 @@ async def update_project(
```
"""
from litellm.proxy.proxy_server import (
general_settings,
litellm_proxy_admin_name,
llm_router,
premium_user,
prisma_client,
user_api_key_cache,
@ -642,6 +768,12 @@ async def update_project(
data=data,
)
# Opt-in (default off): require rpm/tpm for every model the update would leave on the project.
if general_settings.get("enforce_project_model_quota", False):
_raise_on_missing_project_model_quota_on_update(
data, existing_project, _router_access_group_names(llm_router)
)
# Prepare update data
update_data = _jsonified(prisma_client, data.model_dump(exclude_none=True, exclude={"project_id"}))
update_data["updated_by"] = user_api_key_dict.user_id or litellm_proxy_admin_name
@ -652,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,
@ -667,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
@ -694,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},
@ -934,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 []

View file

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

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

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"

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.
#
@ -428,9 +449,11 @@ ui:
maxUnavailable: ""
podAnnotations: {}
# Same shape as the gateway blocks of the same name. The nginx runtime
# writes its pid, cache, and proxy temp files under the image's root
# filesystem, so `securityContext.readOnlyRootFilesystem: true` here needs
# emptyDir volumes mounted over those paths.
# writes its pid, cache, and proxy temp files under /tmp, so it boots as
# any (arbitrary, non-root) uid; `securityContext.readOnlyRootFilesystem:
# true` here needs an emptyDir volume mounted over /tmp. Images before
# the /tmp move instead need emptyDirs over /var/cache/nginx and /run to
# run as a non-root uid at all.
podLabels: {}
podSecurityContext: {}
securityContext: {}

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")
);

View file

@ -0,0 +1,18 @@
-- AlterTable
ALTER TABLE "LiteLLM_DailyUserSpend" ADD COLUMN IF NOT EXISTS "gateway_injected_caching_savings_spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0;
-- AlterTable
ALTER TABLE "LiteLLM_DailyOrganizationSpend" ADD COLUMN IF NOT EXISTS "gateway_injected_caching_savings_spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0;
-- AlterTable
ALTER TABLE "LiteLLM_DailyEndUserSpend" ADD COLUMN IF NOT EXISTS "gateway_injected_caching_savings_spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0;
-- AlterTable
ALTER TABLE "LiteLLM_DailyAgentSpend" ADD COLUMN IF NOT EXISTS "gateway_injected_caching_savings_spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0;
-- AlterTable
ALTER TABLE "LiteLLM_DailyTeamSpend" ADD COLUMN IF NOT EXISTS "gateway_injected_caching_savings_spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0;
-- AlterTable
ALTER TABLE "LiteLLM_DailyTagSpend" ADD COLUMN IF NOT EXISTS "gateway_injected_caching_savings_spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0;

View file

@ -0,0 +1,22 @@
-- AlterTable
ALTER TABLE "LiteLLM_ShadowEvalAttempt" ADD COLUMN IF NOT EXISTS "real_cost" DOUBLE PRECISION;
-- AlterTable
ALTER TABLE "LiteLLM_ShadowEvalAttempt" ADD COLUMN IF NOT EXISTS "real_classifier_cost" DOUBLE PRECISION NOT NULL DEFAULT 0;
-- AlterTable
ALTER TABLE "LiteLLM_ShadowEvalAttempt" ADD COLUMN IF NOT EXISTS "shadow_classifier_cost" DOUBLE PRECISION NOT NULL DEFAULT 0;
-- AlterTable
ALTER TABLE "LiteLLM_ShadowEvalAttempt" ADD COLUMN IF NOT EXISTS "real_cache_hit" BOOLEAN NOT NULL DEFAULT false;
-- CreateTable
CREATE TABLE IF NOT EXISTS "LiteLLM_ShadowEvalFunnel" (
"job_id" TEXT NOT NULL,
"not_sampled" INTEGER NOT NULL DEFAULT 0,
"unjudgeable" INTEGER NOT NULL DEFAULT 0,
"shed" INTEGER NOT NULL DEFAULT 0,
"withheld" INTEGER NOT NULL DEFAULT 0,
CONSTRAINT "LiteLLM_ShadowEvalFunnel_pkey" PRIMARY KEY ("job_id")
);

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 $$;

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())
@ -754,6 +781,7 @@ model LiteLLM_DailyUserSpend {
compression_saved_tokens BigInt @default(0)
compression_savings_spend Float @default(0.0)
prompt_caching_savings_spend Float @default(0.0)
gateway_injected_caching_savings_spend Float @default(0.0)
autorouter_savings_spend Float @default(0.0)
spend Float @default(0.0)
api_requests BigInt @default(0)
@ -789,6 +817,7 @@ model LiteLLM_DailyOrganizationSpend {
compression_saved_tokens BigInt @default(0)
compression_savings_spend Float @default(0.0)
prompt_caching_savings_spend Float @default(0.0)
gateway_injected_caching_savings_spend Float @default(0.0)
autorouter_savings_spend Float @default(0.0)
spend Float @default(0.0)
api_requests BigInt @default(0)
@ -824,6 +853,7 @@ model LiteLLM_DailyEndUserSpend {
compression_saved_tokens BigInt @default(0)
compression_savings_spend Float @default(0.0)
prompt_caching_savings_spend Float @default(0.0)
gateway_injected_caching_savings_spend Float @default(0.0)
autorouter_savings_spend Float @default(0.0)
spend Float @default(0.0)
api_requests BigInt @default(0)
@ -858,6 +888,7 @@ model LiteLLM_DailyAgentSpend {
compression_saved_tokens BigInt @default(0)
compression_savings_spend Float @default(0.0)
prompt_caching_savings_spend Float @default(0.0)
gateway_injected_caching_savings_spend Float @default(0.0)
autorouter_savings_spend Float @default(0.0)
spend Float @default(0.0)
api_requests BigInt @default(0)
@ -892,6 +923,7 @@ model LiteLLM_DailyTeamSpend {
compression_saved_tokens BigInt @default(0)
compression_savings_spend Float @default(0.0)
prompt_caching_savings_spend Float @default(0.0)
gateway_injected_caching_savings_spend Float @default(0.0)
autorouter_savings_spend Float @default(0.0)
spend Float @default(0.0)
api_requests BigInt @default(0)
@ -929,6 +961,7 @@ model LiteLLM_DailyTagSpend {
compression_saved_tokens BigInt @default(0)
compression_savings_spend Float @default(0.0)
prompt_caching_savings_spend Float @default(0.0)
gateway_injected_caching_savings_spend Float @default(0.0)
autorouter_savings_spend Float @default(0.0)
spend Float @default(0.0)
api_requests BigInt @default(0)
@ -1527,12 +1560,27 @@ model LiteLLM_ShadowEvalAttempt {
confidence Float?
judge_cost Float @default(0)
shadow_cost Float @default(0)
real_cost Float? // NULL = row predates cost measurement; comparisons read only measured rows
real_classifier_cost Float @default(0)
shadow_classifier_cost Float @default(0)
real_cache_hit Boolean @default(false)
error String?
created_at DateTime @default(now())
@@index([job_id])
}
// Per-leg sampling funnel counters the attempt rows cannot derive: requests an
// admitting job saw but did not judge. attempted = the leg's attempt rows; the
// leg's eligible traffic = not_sampled + unjudgeable + shed + withheld + attempted.
model LiteLLM_ShadowEvalFunnel {
job_id String @id
not_sampled Int @default(0)
unjudgeable Int @default(0)
shed Int @default(0)
withheld Int @default(0)
}
// ---------------------------------------------------------------------------
// Workflow Run Tracking
//

View file

@ -40,6 +40,65 @@ def _get_prisma_env() -> dict:
_MIGRATION_TS_RE = re.compile(r"^(\d{14})_")
_SPEND_LOGS_ALTER_RE = re.compile(r'^ALTER\s+TABLE\s+"LiteLLM_SpendLogs"\s', re.IGNORECASE)
_SPEND_LOGS_ARTIFACT_DROP_RE = re.compile(
r'^DROP\s+TABLE\s+"LiteLLM_SpendLogs_[^"]*"', re.IGNORECASE
)
_SPEND_LOGS_PK_CLAUSE_RE = re.compile(
r'^(?:DROP\s+CONSTRAINT\s+"[^"]*_pkey"'
r'|ADD\s+(?:CONSTRAINT\s+"[^"]*"\s+)?PRIMARY\s+KEY\s*\([^)]*\))$',
re.IGNORECASE,
)
PARTITIONED_SPEND_LOGS_PUSH_ERROR = (
"LiteLLM_SpendLogs is a partitioned table (see db_scripts/partition_spend_logs.sql), "
"so its primary key must include the partition key (\"startTime\"). `prisma db push` "
"reconciles the database against schema.prisma, which declares the unpartitioned "
"primary key (\"request_id\"), and Postgres rejects that rewrite with: unique "
"constraint on partitioned table must include all partitioning columns. Start the "
"proxy without --use_prisma_db_push so it uses `prisma migrate deploy`, which only "
"applies shipped migrations and leaves the partitioned primary key alone."
)
def _without_sql_comments(statement: str) -> str:
return "\n".join(
line
for line in statement.splitlines()
if line.strip() and not line.strip().startswith("--")
).strip()
def _without_spend_logs_pk_clauses(statement: str) -> Optional[str]:
prefix_match = _SPEND_LOGS_ALTER_RE.match(statement)
if not prefix_match:
return statement
kept = tuple(
clause.strip()
for clause in statement[prefix_match.end():].split(",\n")
if not _SPEND_LOGS_PK_CLAUSE_RE.match(clause.strip())
)
if not kept:
return None
return statement[: prefix_match.end()] + ",\n".join(kept)
def filter_partitioned_spend_logs_diff(diff_sql: str) -> str:
"""Drop statements from a `prisma migrate diff` script that fight the
SpendLogs partitioning runbook (db_scripts/partition_spend_logs.sql): the
primary-key rewrite on "LiteLLM_SpendLogs", which Postgres rejects on a
partitioned table, and drops of runbook artifacts such as
"LiteLLM_SpendLogs_legacy"."""
kept = tuple(
filtered
for statement in diff_sql.split(";")
for bare in (_without_sql_comments(statement),)
if bare and not _SPEND_LOGS_ARTIFACT_DROP_RE.match(bare)
for filtered in (_without_spend_logs_pk_clauses(bare),)
if filtered is not None
)
return "".join(f"{statement};\n\n" for statement in kept)
def _migration_timestamp(name: str) -> int:
"""Extract the leading `YYYYMMDDHHMMSS` timestamp from a migration name.
@ -355,7 +414,24 @@ class ProxyExtrasDBManager:
return
logger.info(f"Migration diff created at {diff_sql_path}")
if ProxyExtrasDBManager.spend_logs_is_partitioned():
filtered_sql = filter_partitioned_spend_logs_diff(
diff_sql_path.read_text()
)
diff_sql_path.write_text(filtered_sql)
logger.info(
"LiteLLM_SpendLogs is partitioned; removed its primary-key "
"rewrite and partitioning artifacts from the drift script"
)
if not filtered_sql.strip():
logger.info("Drift script is empty after filtering; nothing to apply")
if not mark_all_applied:
return
ProxyExtrasDBManager._mark_migrations_applied(migrations_dir)
return
# 2. Run prisma db execute to apply the migration
applied_ok = False
try:
logger.info("Running prisma db execute to apply the migration diff...")
result = subprocess.run(
@ -376,6 +452,7 @@ class ProxyExtrasDBManager:
)
logger.info(f"prisma db execute stdout: {result.stdout}")
logger.info("✅ Migration diff applied successfully")
applied_ok = True
except subprocess.CalledProcessError as e:
logger.warning(f"Failed to apply migration diff: {e.stderr}")
except subprocess.TimeoutExpired:
@ -384,6 +461,16 @@ class ProxyExtrasDBManager:
# 3. Mark all migrations as applied
if not mark_all_applied:
return
if not applied_ok:
logger.warning(
"Drift script failed to apply; NOT marking migrations as "
"applied so a later migration run can retry them"
)
return
ProxyExtrasDBManager._mark_migrations_applied(migrations_dir)
@staticmethod
def _mark_migrations_applied(migrations_dir: str) -> None:
migration_names = ProxyExtrasDBManager._get_migration_names(migrations_dir)
logger.info(f"Resolving {len(migration_names)} migrations")
for migration_name in migration_names:
@ -410,6 +497,55 @@ class ProxyExtrasDBManager:
f"Failed to resolve migration {migration_name}: {e.stderr}"
)
@staticmethod
def spend_logs_is_partitioned() -> bool:
"""True when the connected database's LiteLLM_SpendLogs is a
partitioned table in Prisma's target schema (the `schema` URL param,
falling back to Prisma's default target, public), i.e. the operator
ran db_scripts/partition_spend_logs.sql. Returns False when psycopg is
unavailable or the database cannot be reached, preserving the
pre-existing behavior in those cases."""
database_url = os.getenv("DATABASE_URL")
if not database_url:
return False
try:
import psycopg
except ImportError:
return False
cleaned_url = ProxyExtrasDBManager._strip_prisma_query_params(database_url)
try:
with psycopg.connect(
cleaned_url, connect_timeout=10, autocommit=True
) as conn:
row = conn.execute(
"SELECT 1 "
"FROM pg_partitioned_table pt "
"JOIN pg_class c ON c.oid = pt.partrelid "
"JOIN pg_namespace n ON n.oid = c.relnamespace "
"WHERE c.relname = 'LiteLLM_SpendLogs' "
" AND n.nspname = %s",
(
ProxyExtrasDBManager._prisma_schema_param(database_url)
or "public",
),
).fetchone()
except (psycopg.OperationalError, psycopg.DatabaseError):
return False
return row is not None
@staticmethod
def _prisma_schema_param(url: str) -> Optional[str]:
"""The `schema` query param Prisma uses to pick its target schema,
or None when the URL does not set one."""
from urllib.parse import urlparse, parse_qsl
return next(
(v for k, v in parse_qsl(urlparse(url).query) if k == "schema"),
None,
)
@staticmethod
def _strip_prisma_query_params(url: str) -> str:
"""Remove Prisma-specific query params (connection_limit, pool_timeout,
@ -528,7 +664,8 @@ class ProxyExtrasDBManager:
migrations_dir = ProxyExtrasDBManager._get_prisma_dir()
if not use_migrate:
# Preserve `prisma db push` path unchanged.
if ProxyExtrasDBManager.spend_logs_is_partitioned():
raise RuntimeError(PARTITIONED_SPEND_LOGS_PUSH_ERROR)
original_dir = os.getcwd()
os.chdir(migrations_dir)
try:
@ -972,6 +1109,8 @@ class ProxyExtrasDBManager:
)
raise
else:
if ProxyExtrasDBManager.spend_logs_is_partitioned():
raise RuntimeError(PARTITIONED_SPEND_LOGS_PUSH_ERROR)
# Use prisma db push with increased timeout
subprocess.run(
[_get_prisma_command(), "db", "push", "--accept-data-loss"],

View file

@ -1,6 +1,6 @@
[project]
name = "litellm-proxy-extras"
version = "0.4.89"
version = "0.4.91"
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.89"
version = "0.4.91"
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"

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"] }

View file

@ -9,10 +9,23 @@ repository.workspace = true
name = "_native"
crate-type = ["cdylib"]
[features]
default = ["extension-module"]
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

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);

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,

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()))
}

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()
);
}
}
}

View file

@ -199,6 +199,7 @@ standard_logging_payload_excluded_fields: Optional[List[str]] = (
None # Fields to exclude from StandardLoggingPayload before callbacks receive it
)
log_raw_request_response: bool = False
log_client_error_tracebacks: bool = False
request_correlation_in_logs: bool = False
redact_messages_in_exceptions: Optional[bool] = False
redact_user_api_key_info: Optional[bool] = False
@ -273,7 +274,6 @@ databricks_key: Optional[str] = None
openai_like_key: Optional[str] = None
azure_key: Optional[str] = None
anthropic_key: Optional[str] = None
autorouter_savings_baseline_model: Optional[str] = None
replicate_key: Optional[str] = None
bytez_key: Optional[str] = None
gdc_key: Optional[str] = None
@ -444,6 +444,7 @@ max_ui_session_budget: Optional[float] = (
1.0 # USD budget for each dashboard login session (playground, test connection)
)
internal_user_budget_duration: Optional[str] = None
budget_rollover: bool = False # carry spend beyond max_budget into the next window instead of zeroing it
tag_budget_config: Optional[Dict[str, "BudgetConfig"]] = None
max_end_user_budget: Optional[float] = None
max_end_user_budget_id: Optional[str] = None
@ -463,6 +464,11 @@ prometheus_metrics_config: Optional[List] = None
prometheus_exclude_metrics: Optional[List[str]] = None
prometheus_exclude_labels: Optional[List[str]] = None
prometheus_emit_stream_label: bool = False
prometheus_deployment_and_latency_caller_identity: Literal[
"api_key_alias",
"user_email",
"both",
] = "api_key_alias"
# Opt-in: emit `rate_limit_category` and `rate_limit_type` labels on
# `litellm_proxy_failed_requests_metric`. Off by default to preserve the
# pre-unification label set so existing dashboards / recording rules keyed on
@ -480,6 +486,7 @@ public_mcp_servers: Optional[List[str]] = None
public_mcp_hub_strict_whitelist: bool = True
public_model_groups: Optional[List[str]] = None
public_agent_groups: Optional[List[str]] = None
agent_search_embedding_model: Optional[str] = None
# Supports both old format (Dict[str, str]) and new format (Dict[str, Dict[str, Any]])
# New format: { "displayName": { "url": "...", "index": 0 } }
# Old format: { "displayName": "url" } (for backward compatibility)
@ -1628,6 +1635,9 @@ if TYPE_CHECKING:
AmazonMantleMessagesConfig as AmazonMantleMessagesConfig,
)
from .llms.together_ai.chat import TogetherAIConfig as TogetherAIConfig
from .llms.together_ai.chat.transformation import (
TogetherAIChatConfig as TogetherAIChatConfig,
)
from .llms.nlp_cloud.chat.handler import NLPCloudConfig as NLPCloudConfig
from .llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
VertexGeminiConfig as VertexGeminiConfig,
@ -1801,6 +1811,9 @@ if TYPE_CHECKING:
from .llms.gemini.interactions.transformation import (
GoogleAIStudioInteractionsConfig as GoogleAIStudioInteractionsConfig,
)
from .llms.vertex_ai.interactions.transformation import (
VertexAIInteractionsConfig as VertexAIInteractionsConfig,
)
from .llms.openai.chat.o_series_transformation import (
OpenAIOSeriesConfig as OpenAIOSeriesConfig,
OpenAIOSeriesConfig as OpenAIO1Config,

View file

@ -18,7 +18,10 @@ until they're actually needed.
import importlib
import sys
from collections.abc import Callable
from typing import Any, Final, cast
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,6 +56,9 @@ from ._lazy_imports_registry import (
UTILS_NAMES,
)
if TYPE_CHECKING:
from tiktoken import Encoding
def get_litellm_globals() -> dict:
"""
@ -78,10 +84,10 @@ def _get_utils_globals() -> dict:
# 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 +106,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 +130,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 +160,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 +212,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 +272,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 +289,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 +372,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 +387,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 +403,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 +419,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.

View file

@ -177,6 +177,7 @@ LLM_CONFIG_NAMES: Final = (
"AmazonAnthropicClaudeMessagesConfig",
"AmazonMantleMessagesConfig",
"TogetherAIConfig",
"TogetherAIChatConfig",
"NLPCloudConfig",
"VertexGeminiConfig",
"GoogleAIStudioGeminiConfig",
@ -242,6 +243,7 @@ LLM_CONFIG_NAMES: Final = (
"OpenRouterResponsesAPIConfig",
"BedrockMantleResponsesAPIConfig",
"GoogleAIStudioInteractionsConfig",
"VertexAIInteractionsConfig",
"OpenAIOSeriesConfig",
"AnthropicSkillsConfig",
"BaseSkillsAPIConfig",
@ -740,6 +742,10 @@ _LLM_CONFIGS_IMPORT_MAP: Final = {
"AmazonMantleMessagesConfig",
),
"TogetherAIConfig": (".llms.together_ai.chat", "TogetherAIConfig"),
"TogetherAIChatConfig": (
".llms.together_ai.chat.transformation",
"TogetherAIChatConfig",
),
"NLPCloudConfig": (".llms.nlp_cloud.chat.handler", "NLPCloudConfig"),
"VertexGeminiConfig": (
".llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini",
@ -977,6 +983,10 @@ _LLM_CONFIGS_IMPORT_MAP: Final = {
".llms.gemini.interactions.transformation",
"GoogleAIStudioInteractionsConfig",
),
"VertexAIInteractionsConfig": (
".llms.vertex_ai.interactions.transformation",
"VertexAIInteractionsConfig",
),
"OpenAIOSeriesConfig": (
".llms.openai.chat.o_series_transformation",
"OpenAIOSeriesConfig",

View file

@ -5,7 +5,7 @@ import os
import sys
from datetime import datetime
from logging import Formatter
from typing import Any, Final
from typing import Any, Final, TextIO
import litellm
from litellm.constants import (
@ -234,11 +234,65 @@ class CorrelationContextFilter(logging.Filter):
_correlation_filter: Final = CorrelationContextFilter()
json_logs = bool(os.getenv("JSON_LOGS", False))
_LOG_FORMAT_PREFIX: Final = "%(asctime)s - %(name)s:%(levelname)s"
_LOG_FORMAT_SUFFIX: Final = ": %(filename)s:%(lineno)s - %(message)s"
_PLAIN_LOG_FORMAT: Final = _LOG_FORMAT_PREFIX + _LOG_FORMAT_SUFFIX
_COLOR_LOG_FORMAT: Final = f"\033[92m{_LOG_FORMAT_PREFIX}\033[0m{_LOG_FORMAT_SUFFIX}"
def _stream_is_tty(stream: TextIO | None) -> bool:
"""True when the stream is an open interactive terminal; never raises.
A stream can be None (pythonw/embedded interpreters), lack isatty entirely
(GUI log-redirect shims), or be closed; import must survive all three.
"""
try:
return stream is not None and stream.isatty()
except (AttributeError, ValueError):
return False
def _plain_log_format(stdout: TextIO | None, stderr: TextIO | None) -> str:
"""The plain-text log format, colorized only when both streams are an interactive terminal.
Honors the NO_COLOR convention from no-color.org: color is disabled when
NO_COLOR is present with a non-empty value.
"""
if os.environ.get("NO_COLOR"):
return _PLAIN_LOG_FORMAT
return _COLOR_LOG_FORMAT if _stream_is_tty(stdout) and _stream_is_tty(stderr) else _PLAIN_LOG_FORMAT
class LevelRoutingStreamHandler(logging.StreamHandler):
"""Writes records below WARNING to stdout and WARNING and above to stderr.
Collectors that derive severity from the stream report every stderr line as an error.
"""
def emit(self, record: logging.LogRecord) -> None:
preferred: Final = sys.stdout if record.levelno < logging.WARNING 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:
self.stream = preferred # rebind-ok: StreamHandler.emit writes self.stream under the handler lock
super().emit(record)
def _parse_json_logs_env(value: str | None) -> bool:
"""Strict opt-in parse for the JSON_LOGS env var: only "true" (any case) enables JSON logs.
Matches the reader in litellm-proxy-extras/_logging.py. The previous
bool(os.getenv(...)) treated any non-empty value, including "false" and "0",
as enabled.
"""
return (value or "").lower() == "true"
json_logs: Final = _parse_json_logs_env(os.getenv("JSON_LOGS"))
# Create a handler for the logger (you may need to adapt this based on your needs)
log_level: Final = os.getenv("LITELLM_LOG", "DEBUG")
numeric_level: Final[str] = getattr(logging, log_level.upper())
handler: Final = logging.StreamHandler()
handler: Final = LevelRoutingStreamHandler()
handler.setLevel(numeric_level)
handler.addFilter(_secret_filter)
handler.addFilter(_correlation_filter)
@ -447,7 +501,7 @@ if json_logs:
_setup_json_exception_handlers(JsonFormatter())
else:
formatter: Final = CorrelationPlainFormatter(
"\033[92m%(asctime)s - %(name)s:%(levelname)s\033[0m: %(filename)s:%(lineno)s - %(message)s",
_plain_log_format(sys.stdout, sys.stderr),
datefmt="%H:%M:%S",
)
@ -628,7 +682,7 @@ def _turn_on_json():
- Adds a JSON formatter to all loggers
"""
handler: Final = logging.StreamHandler()
handler: Final = LevelRoutingStreamHandler()
handler.setFormatter(JsonFormatter())
_initialize_loggers_with_handler(handler)
# Set up exception handlers

View file

@ -12,8 +12,9 @@ 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
from collections.abc import Callable, Mapping
from typing import Final
from urllib.parse import urlsplit, urlunsplit
import redis
import redis.asyncio as async_redis
@ -50,6 +51,7 @@ def _get_redis_kwargs():
include_args: Final = {
"url",
"redis_connect_func",
"credential_provider",
"gcp_service_account",
"gcp_ssl_ca_certs",
"azure_redis_ad_token",
@ -155,7 +157,8 @@ def _get_redis_cluster_kwargs(client=None):
def _get_redis_env_kwarg_mapping():
PREFIX: Final = "REDIS_"
return {f"{PREFIX}{x.upper()}": x for x in _get_redis_kwargs()}
exclude_from_environment: Final = frozenset({"credential_provider"})
return {f"{PREFIX}{x.upper()}": x for x in _get_redis_kwargs() if x not in exclude_from_environment}
def _redis_kwargs_from_environment():
@ -353,6 +356,12 @@ def get_redis_url_from_environment():
return f"{redis_protocol}://{auth_part}{os.environ['REDIS_HOST']}:{os.environ['REDIS_PORT']}"
def _url_without_userinfo(url: str) -> str:
parts: Final = urlsplit(url)
netloc: Final = parts.netloc.rsplit("@", 1)[-1]
return urlunsplit((parts.scheme, netloc, parts.path, parts.query, parts.fragment))
def _get_redis_client_logic(**env_overrides):
"""
Common functionality across sync + async redis client implementations
@ -410,54 +419,58 @@ def _get_redis_client_logic(**env_overrides):
if _service_name is not None:
redis_kwargs["service_name"] = _service_name
# Handle GCP IAM authentication
_gcp_service_account: Final = redis_kwargs.get("gcp_service_account") or get_secret_str("REDIS_GCP_SERVICE_ACCOUNT")
_gcp_ssl_ca_certs: Final = redis_kwargs.get("gcp_ssl_ca_certs") or get_secret_str("REDIS_GCP_SSL_CA_CERTS")
if _gcp_service_account is not None:
verbose_logger.debug("Setting up GCP IAM authentication for Redis with service account.")
redis_kwargs["redis_connect_func"] = create_gcp_iam_redis_connect_func(
service_account=_gcp_service_account, ssl_ca_certs=_gcp_ssl_ca_certs
if redis_kwargs.get("credential_provider") is None:
# Handle GCP IAM authentication
_gcp_service_account: Final = redis_kwargs.get("gcp_service_account") or get_secret_str(
"REDIS_GCP_SERVICE_ACCOUNT"
)
# Store GCP service account in redis_connect_func for async cluster access
redis_kwargs["redis_connect_func"]._gcp_service_account = _gcp_service_account
_gcp_ssl_ca_certs: Final = redis_kwargs.get("gcp_ssl_ca_certs") or get_secret_str("REDIS_GCP_SSL_CA_CERTS")
# Remove GCP-specific kwargs that shouldn't be passed to Redis client
redis_kwargs.pop("gcp_service_account", None)
redis_kwargs.pop("gcp_ssl_ca_certs", None)
if _gcp_service_account is not None:
verbose_logger.debug("Setting up GCP IAM authentication for Redis with service account.")
redis_kwargs["redis_connect_func"] = create_gcp_iam_redis_connect_func(
service_account=_gcp_service_account, ssl_ca_certs=_gcp_ssl_ca_certs
)
# Store GCP service account in redis_connect_func for async cluster access
redis_kwargs["redis_connect_func"]._gcp_service_account = _gcp_service_account
# Only enable SSL if explicitly requested AND SSL CA certs are provided
if _gcp_ssl_ca_certs and redis_kwargs.get("ssl", False):
redis_kwargs["ssl_ca_certs"] = _gcp_ssl_ca_certs
# Only enable SSL if explicitly requested AND SSL CA certs are provided
if _gcp_ssl_ca_certs and redis_kwargs.get("ssl", False):
redis_kwargs["ssl_ca_certs"] = _gcp_ssl_ca_certs
# Handle Azure AD authentication (after GCP IAM block)
_azure_redis_ad_token: Final = redis_kwargs.get("azure_redis_ad_token") or get_secret("REDIS_AZURE_AD_TOKEN")
# Handle Azure AD authentication (after GCP IAM block)
_azure_redis_ad_token: Final = redis_kwargs.get("azure_redis_ad_token") or get_secret("REDIS_AZURE_AD_TOKEN")
_azure_ad_enabled: Final = _azure_redis_ad_token is not None and str(_azure_redis_ad_token).lower() == "true"
_azure_ad_enabled: Final = _azure_redis_ad_token is not None and str(_azure_redis_ad_token).lower() == "true"
if _azure_ad_enabled and _gcp_service_account is not None:
verbose_logger.warning(
"Both GCP IAM (gcp_service_account) and Azure AD (azure_redis_ad_token) are configured for Redis. "
"Using GCP IAM. Remove one to avoid misconfiguration."
)
if _azure_ad_enabled and _gcp_service_account is not None:
verbose_logger.warning(
"Both GCP IAM (gcp_service_account) and Azure AD (azure_redis_ad_token) are configured for Redis. "
"Using GCP IAM. Remove one to avoid misconfiguration."
)
if _azure_ad_enabled and _gcp_service_account is None:
_azure_client_id: Final = redis_kwargs.get("azure_client_id") or get_secret_str("AZURE_CLIENT_ID")
_azure_tenant_id: Final = redis_kwargs.get("azure_tenant_id") or get_secret_str("AZURE_TENANT_ID")
_azure_client_secret: Final = redis_kwargs.get("azure_client_secret") or get_secret_str("AZURE_CLIENT_SECRET")
if _azure_ad_enabled and _gcp_service_account is None:
_azure_client_id: Final = redis_kwargs.get("azure_client_id") or get_secret_str("AZURE_CLIENT_ID")
_azure_tenant_id: Final = redis_kwargs.get("azure_tenant_id") or get_secret_str("AZURE_TENANT_ID")
_azure_client_secret: Final = redis_kwargs.get("azure_client_secret") or get_secret_str(
"AZURE_CLIENT_SECRET"
)
verbose_logger.debug("Setting up Azure AD authentication for Redis.")
redis_kwargs["redis_connect_func"] = create_azure_ad_redis_connect_func(
azure_client_id=_azure_client_id,
azure_tenant_id=_azure_tenant_id,
azure_client_secret=_azure_client_secret,
)
# Marker for async paths to detect Azure AD auth. The live credential
# object is attached separately as `_azure_credential` by
# `create_azure_ad_redis_connect_func`; the raw client_id/tenant_id/secret
# are intentionally NOT exposed on the function to avoid leaking
# credentials via inspection or logging.
redis_kwargs["redis_connect_func"]._azure_redis_ad_token = True
verbose_logger.debug("Setting up Azure AD authentication for Redis.")
redis_kwargs["redis_connect_func"] = create_azure_ad_redis_connect_func(
azure_client_id=_azure_client_id,
azure_tenant_id=_azure_tenant_id,
azure_client_secret=_azure_client_secret,
)
# Marker for async paths to detect Azure AD auth. The live credential
# object is attached separately as `_azure_credential` by
# `create_azure_ad_redis_connect_func`; the raw client_id/tenant_id/secret
# are intentionally NOT exposed on the function to avoid leaking
# credentials via inspection or logging.
redis_kwargs["redis_connect_func"]._azure_redis_ad_token = True
redis_kwargs.pop("gcp_service_account", None)
redis_kwargs.pop("gcp_ssl_ca_certs", None)
# Always remove Azure-specific kwargs that shouldn't be passed to Redis client
redis_kwargs.pop("azure_redis_ad_token", None)
@ -465,6 +478,13 @@ def _get_redis_client_logic(**env_overrides):
redis_kwargs.pop("azure_tenant_id", None)
redis_kwargs.pop("azure_client_secret", None)
if redis_kwargs.get("credential_provider") is not None:
redis_kwargs.pop("redis_connect_func", None)
redis_kwargs.pop("username", None)
redis_kwargs.pop("password", None)
if redis_kwargs.get("url") is not None:
redis_kwargs["url"] = _url_without_userinfo(redis_kwargs["url"])
if "url" in redis_kwargs and redis_kwargs["url"] is not None:
# Only strip host/port/db/password when not routing to a cluster.
# When startup_nodes is also present the cluster path takes priority and
@ -532,8 +552,7 @@ def _init_redis_sentinel(redis_kwargs) -> redis.Redis:
service_name: Final = redis_kwargs.get("service_name")
connection_kwargs: Final = _get_redis_sentinel_connection_kwargs(redis_kwargs)
connection_kwargs.setdefault("socket_timeout", REDIS_SOCKET_TIMEOUT)
sentinel_kwargs: Final = dict(connection_kwargs)
sentinel_kwargs["password"] = sentinel_password
sentinel_kwargs: Final = _sentinel_auth_kwargs(connection_kwargs, sentinel_password)
if not sentinel_nodes or not service_name:
raise ValueError("Both 'sentinel_nodes' and 'service_name' are required for Redis Sentinel.")
@ -605,7 +624,12 @@ def _async_credential_provider(redis_connect_func: object | None) -> CredentialP
def _async_auth_kwargs(redis_kwargs: dict) -> dict:
"""Swaps a connect func an async path cannot run for the equivalent credential provider,
which supersedes any static username or password redis-py would otherwise reject it with."""
credential_provider: Final = _async_credential_provider(redis_kwargs.get("redis_connect_func"))
explicit_provider: Final = redis_kwargs.get("credential_provider")
credential_provider: Final = (
explicit_provider
if explicit_provider is not None
else _async_credential_provider(redis_kwargs.get("redis_connect_func"))
)
if credential_provider is None:
return redis_kwargs
@ -738,8 +762,20 @@ def get_redis_connection_pool(
return async_redis.BlockingConnectionPool(timeout=REDIS_CONNECTION_POOL_TIMEOUT, **redis_kwargs)
def _redis_kwargs_for_logging(redis_kwargs: Mapping[str, object]) -> Mapping[str, object]:
return {
key: "<credential provider>"
if key == "credential_provider" and value is not None
else "<redis connect function>"
if key == "redis_connect_func" and value is not None
else value
for key, value in redis_kwargs.items()
}
def _pretty_print_redis_config(redis_kwargs: dict) -> None:
"""Pretty print the Redis configuration using rich with sensitive data masking"""
redis_kwargs_for_logging: Final = _redis_kwargs_for_logging(redis_kwargs)
try:
import logging
@ -757,7 +793,7 @@ def _pretty_print_redis_config(redis_kwargs: dict) -> None:
masker = SensitiveDataMasker()
# Mask sensitive data in redis_kwargs
masked_redis_kwargs = masker.mask_dict(redis_kwargs)
masked_redis_kwargs = masker.mask_dict(redis_kwargs_for_logging)
# Create main panel title
title: Final = Text("Redis Configuration", style="bold blue")
@ -820,7 +856,7 @@ def _pretty_print_redis_config(redis_kwargs: dict) -> None:
except ImportError:
# Fallback to simple logging if rich is not available
masker = SensitiveDataMasker()
masked_redis_kwargs = masker.mask_dict(redis_kwargs)
masked_redis_kwargs = masker.mask_dict(redis_kwargs_for_logging)
verbose_logger.info("Redis configuration: %s", masked_redis_kwargs)
except Exception as e:
verbose_logger.error("Error pretty printing Redis configuration: %s", e)

View file

@ -4,6 +4,7 @@ Custom A2A Card Resolver for LiteLLM.
Extends the A2A SDK's card resolver to support multiple well-known paths.
"""
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final
from litellm._logging import verbose_logger
@ -48,6 +49,43 @@ def is_localhost_or_internal_url(url: str | None) -> bool:
return any(pattern in url_lower for pattern in LOCALHOST_URL_PATTERNS)
_CANONICAL_PROTOCOL_BINDINGS: Final = MappingProxyType(
{
"jsonrpc": "JSONRPC",
"http+json": "HTTP+JSON",
"grpc": "GRPC",
}
)
_LEGACY_PROTOCOL_VERSION: Final = "0.3"
def normalize_agent_card_interfaces(agent_card: "AgentCard") -> "AgentCard":
"""
Canonicalize the supported interfaces of spec-adjacent agent cards.
Some A2A servers (e.g. LangGraph Platform) serve agent cards with lowercase
bindings like "jsonrpc", but a2a-sdk's ClientFactory matches bindings
case-sensitively against its uppercase TransportProtocol constants and fails
with "no compatible transports found." for spec-adjacent casings.
The same servers also speak the A2A 0.3 JSON dialect ("kind"-discriminated
payloads) while declaring protocolVersion "1.0", which a2a-sdk's strict v1
proto parsing rejects. A mis-cased binding fingerprints such a server, so its
declared version is downgraded to 0.3 to route the SDK's ClientFactory onto
its v0.3 compat transport, which speaks that dialect.
"""
normalized: Final = type(agent_card)()
normalized.CopyFrom(agent_card)
for interface in normalized.supported_interfaces:
canonical: str | None = _CANONICAL_PROTOCOL_BINDINGS.get(interface.protocol_binding.lower())
if canonical is None or canonical == interface.protocol_binding:
continue
interface.protocol_binding = canonical
interface.protocol_version = _LEGACY_PROTOCOL_VERSION
return normalized
def get_agent_card_url(agent_card: "AgentCard") -> str | None:
"""Return the agent endpoint URL from the resolved SDK card."""
url: Final = getattr(agent_card, "url", None)

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.

View file

@ -73,6 +73,7 @@ except ImportError:
from litellm.a2a_protocol.card_resolver import (
LiteLLMA2ACardResolver,
get_agent_card_url,
normalize_agent_card_interfaces,
)
from litellm.a2a_protocol.exception_mapping_utils import (
handle_a2a_localhost_retry,
@ -85,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:
@ -98,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,
@ -108,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:
"""
@ -122,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
@ -131,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:
"""
@ -143,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.
@ -174,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
@ -497,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,
@ -782,13 +787,17 @@ async def create_a2a_client(
if extra_headers:
verbose_proxy_logger.debug("A2A client created with extra_headers=%s", list(extra_headers.keys()))
resolver: Final = A2ACardResolver(httpx_client=httpx_client, base_url=base_url)
agent_card: Final = normalize_agent_card_interfaces(
await resolver.get_agent_card(http_kwargs={"headers": extra_headers} if extra_headers else None)
)
a2a_client: Final = await create_client( # pyright: ignore[reportOptionalCall]
base_url,
agent_card,
client_config=ClientConfig( # pyright: ignore[reportOptionalCall]
httpx_client=httpx_client,
streaming=streaming,
),
resolver_http_kwargs={"headers": extra_headers} if extra_headers else None,
)
# Stash LiteLLM-owned handles on the client so the localhost-retry path can reuse
# the configured httpx client and this agent's headers without excavating
@ -799,9 +808,7 @@ async def create_a2a_client(
if extra_headers
else None
)
agent_card: Final = getattr(a2a_client, "_card", None)
if agent_card is not None:
a2a_client._litellm_agent_card = agent_card
a2a_client._litellm_agent_card = agent_card
verbose_logger.info("A2A client created for %s", base_url)

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:
@ -551,7 +668,7 @@ def _get_batch_job_usage_from_response_body(
return usage
def _get_anthropic_result_from_batch_results_line(batch_results_line: Mapping[str, Any]) -> dict:
def _get_anthropic_result_from_batch_results_line(batch_results_line: Mapping[str, Any]) -> Mapping[str, Any]:
"""
Get the ``result`` object from a line of an Anthropic message batch results JSONL file.
@ -563,7 +680,7 @@ def _get_anthropic_result_from_batch_results_line(batch_results_line: Mapping[st
def _get_response_from_batch_job_output_file(
batch_job_output_file: Mapping[str, Any], custom_llm_provider: str = "openai"
) -> Any:
) -> Mapping[str, Any]:
"""
Get the response from the batch job output file
"""

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:

View file

@ -18,7 +18,7 @@ import asyncio
import datetime
import inspect
import time
from collections.abc import AsyncGenerator, AsyncIterator, Callable, Generator, Mapping
from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Callable, Generator, Mapping
from typing import TYPE_CHECKING, Any, Final, Optional, TypeVar
from pydantic import BaseModel
@ -27,6 +27,7 @@ import litellm
from litellm._logging import print_verbose, verbose_logger
from litellm.caching import InMemoryCache
from litellm.caching.caching import S3Cache
from litellm.constants import CACHE_WRITE_SHUTDOWN_FLUSH_TIMEOUT_SECONDS
from litellm.litellm_core_utils.llm_response_utils.response_metadata import (
update_response_metadata,
)
@ -124,6 +125,29 @@ def _prompt_tokens_details_as_mapping(details: "PromptTokensDetailsWrapper") ->
return details.model_dump(exclude_none=True) if hasattr(details, "model_dump") else {}
_PENDING_CACHE_WRITES: Final[set["asyncio.Task[None]"]] = set() # mutable-ok: strong refs to pending write tasks
async def _complete_cache_write_despite_cancellation(write_factory: Callable[[], Awaitable[None]]) -> None:
try:
await write_factory()
except asyncio.CancelledError:
try:
await asyncio.wait_for(write_factory(), timeout=CACHE_WRITE_SHUTDOWN_FLUSH_TIMEOUT_SECONDS)
except Exception as flush_error: # noqa: BLE001 # shutdown flush failures are logged, never raised
verbose_logger.warning(
"LiteLLM Cache: pending cache write failed during event loop shutdown: %s", flush_error
)
raise
def create_cache_write_task(write_factory: Callable[[], Awaitable[None]]) -> "asyncio.Task[None]":
task: Final = asyncio.create_task(_complete_cache_write_despite_cancellation(write_factory))
_PENDING_CACHE_WRITES.add(task)
task.add_done_callback(_PENDING_CACHE_WRITES.discard)
return task
def _request_cache_key(request_kwargs: Mapping[str, Any]) -> str | None:
"""Read the caller-supplied ``cache_key`` off the request kwargs."""
return request_kwargs.get("cache_key", None)
@ -983,6 +1007,7 @@ class LLMCachingHandler:
if litellm.cache is None:
return
cache: Final = litellm.cache
new_kwargs: Final = kwargs.copy()
new_kwargs.update(
@ -1004,24 +1029,24 @@ class LLMCachingHandler:
):
if (
isinstance(result, EmbeddingResponse)
and litellm.cache is not None
and not isinstance(litellm.cache.cache, S3Cache) # s3 doesn't support bulk writing. Exclude.
and not isinstance(cache.cache, S3Cache) # s3 doesn't support bulk writing. Exclude.
):
asyncio.create_task(
litellm.cache.async_add_cache_pipeline(
create_cache_write_task(
lambda: cache.async_add_cache_pipeline(
result, dynamic_cache_object=self.dual_cache, **new_kwargs
)
)
else:
asyncio.create_task(
litellm.cache.async_add_cache(
result.model_dump_json(),
result_json: Final = result.model_dump_json()
create_cache_write_task(
lambda: cache.async_add_cache(
result_json,
dynamic_cache_object=self.dual_cache,
**new_kwargs,
)
)
else:
asyncio.create_task(litellm.cache.async_add_cache(result, **new_kwargs))
create_cache_write_task(lambda: cache.async_add_cache(result, **new_kwargs))
def sync_set_cache(
self,

View file

@ -175,6 +175,10 @@ _RedisCallResult = TypeVar("_RedisCallResult")
_swallowed_redis_failures: Final[ContextVar[int]] = ContextVar("litellm_swallowed_redis_failures", default=0)
def _opaque_kwarg_key(value: object) -> str:
return f"{type(value).__name__}-{id(value)}"
@functools.lru_cache(maxsize=1)
def _redis_health_error_types() -> tuple[type, ...]:
"""Exception types that mean the Redis backend itself is unhealthy.
@ -399,10 +403,9 @@ class RedisCache(BaseCache):
Generate a cache key for the async Redis client based on connection parameters.
This ensures different Redis configurations use different cached clients.
"""
# Create a stable representation of redis_kwargs for hashing
# Sort keys to ensure consistent hash regardless of parameter order
sorted_kwargs: Final = sorted(self.redis_kwargs.items())
kwargs_str: Final = json.dumps(sorted_kwargs, sort_keys=True)
kwargs_str: Final = json.dumps(sorted_kwargs, sort_keys=True, default=_opaque_kwarg_key)
kwargs_hash: Final = hashlib.sha256(kwargs_str.encode()).hexdigest()[:16]
return f"async-redis-client-{kwargs_hash}"
@ -432,7 +435,7 @@ class RedisCache(BaseCache):
"""
if key is None:
return key
if self.namespace is not None and not key.startswith(self.namespace):
if self.namespace and not key.startswith(self.namespace + ":"):
key = self.namespace + ":" + key
return key
@ -1384,10 +1387,10 @@ class RedisCache(BaseCache):
dict: {"status": "success" | "failed", "message": str, "error": Optional[str]}
"""
try:
import redis.asyncio as redis_async
from .._redis import get_redis_async_client
# Create a fresh Redis client with current settings
redis_client: Final = redis_async.Redis(**self.redis_kwargs)
redis_client: Final = get_redis_async_client(**self.redis_kwargs)
# Test the connection
ping_result: Final = await redis_client.ping()

View file

@ -64,22 +64,9 @@ class RedisClusterCache(RedisCache):
dict: {"status": "success" | "failed", "message": str, "error": Optional[str]}
"""
try:
import redis.asyncio as redis_async
from redis.cluster import ClusterNode
from .._redis import get_redis_async_client
# Create ClusterNode objects from startup_nodes
cluster_kwargs: Final = self.redis_kwargs.copy()
startup_nodes: Final = cluster_kwargs.pop("startup_nodes", [])
new_startup_nodes: Final[list[ClusterNode]] = []
for item in startup_nodes:
new_startup_nodes.append(ClusterNode(**item))
# Create a fresh Redis Cluster client with current settings
redis_client: Final = redis_async.RedisCluster(
startup_nodes=new_startup_nodes,
**cluster_kwargs,
)
redis_client: Final = get_redis_async_client(**self.redis_kwargs)
# Test the connection
ping_result: Final = await redis_client.ping()

View file

@ -18,6 +18,14 @@ already does when one of its pooled connections errors), leaving every other nod
connections untouched. Every other branch (MOVED, ASK, CLUSTERDOWN, slot-not-covered,
retry-exhaustion) is unchanged from upstream, since those already carry real evidence the
topology changed.
redis-py 8.x fixed this upstream with gentler machinery than this override's
``node.disconnect()`` (which also kills connections other coroutines are mid-operation
on, so one timeout cascades into a reconnect storm and, with TLS, a fresh handshake per
killed connection): it marks in-use connections for reconnect only after their current
operation completes, disconnects only the idle pooled ones, and defers reinitialization
to the outer retry loop. When the installed ``ClusterNode`` has that per-connection
recovery API, the factory returns the base ``RedisCluster`` unmodified.
"""
import asyncio
@ -72,8 +80,16 @@ class _ClusterAttrs(Protocol):
_VERIFIED_REDIS_VERSIONS: Final = frozenset({"5.3.1"})
def get_litellm_async_redis_cluster_class() -> type["_AsyncRedisClusterType"]:
"""Builds the ``RedisCluster`` subclass with the per-node isolation fix.
def get_litellm_async_redis_cluster_class(
cluster_node_class: type | None = None,
) -> type["_AsyncRedisClusterType"]:
"""Returns the base ``RedisCluster`` when the installed redis-py already recovers a
node-level connection error per-connection (8.x+), else builds the ``RedisCluster``
subclass with the per-node isolation fix for older versions whose upstream branch
tears down the whole cluster client.
``cluster_node_class`` exists for dependency injection in tests; production callers
leave it unset and the installed ``ClusterNode`` is used.
Imported lazily because this module is reachable from a base ``import litellm`` while
redis is not a base dependency. Cheap to call repeatedly: the underlying redis
@ -81,7 +97,10 @@ def get_litellm_async_redis_cluster_class() -> type["_AsyncRedisClusterType"]:
"""
import redis
from redis.asyncio.cluster import (
RedisCluster as _BaseAsyncRedisCluster, # pyright: ignore[reportUnknownVariableType] # redis-py ships no resolvable stub for this class under the repo's current (stale) types-redis pin
ClusterNode as _AsyncClusterNode, # pyright: ignore[reportUnknownVariableType] # redis-py ships no resolvable stub for this class under the repo's current (stale) types-redis pin
)
from redis.asyncio.cluster import (
RedisCluster as _BaseAsyncRedisCluster, # pyright: ignore[reportUnknownVariableType] # same stale-stub gap as the import above
)
from redis.cluster import get_node_name
from redis.commands import READ_COMMANDS
@ -98,6 +117,15 @@ def get_litellm_async_redis_cluster_class() -> type["_AsyncRedisClusterType"]:
from redis.exceptions import ConnectionError as _RedisConnectionError
from redis.exceptions import TimeoutError as _RedisTimeoutError
node_class: Final = cluster_node_class if cluster_node_class is not None else _AsyncClusterNode
if hasattr(node_class, "update_active_connections_for_reconnect"):
verbose_logger.debug(
"redis-py %s recovers a node-level connection error per-connection upstream; "
"using the base RedisCluster without litellm's node-isolation override.",
redis.__version__,
)
return _BaseAsyncRedisCluster
if redis.__version__ not in _VERIFIED_REDIS_VERSIONS:
verbose_logger.warning(
"redis-py %s is not in the set this cluster-teardown-storm fix was verified "

View file

@ -5,7 +5,7 @@ Handler for transforming /chat/completions api requests to litellm.responses req
import json
import os
from collections.abc import AsyncIterator, Callable, Iterable, Iterator, Mapping, Sequence
from typing import TYPE_CHECKING, Any, Final, Literal, TypedDict, Union, cast
from typing import TYPE_CHECKING, Any, Final, Literal, TypedDict, Union, cast, get_args
from openai.types.responses.custom_tool_param import CustomToolParam
from openai.types.responses.response_input_param import (
@ -21,6 +21,9 @@ from pydantic import BaseModel
import litellm
from litellm import ModelResponse
from litellm._logging import verbose_logger
from litellm.litellm_core_utils.prompt_templates.common_utils import (
responses_reasoning_item_from_thinking_blocks,
)
from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator
from litellm.llms.base_llm.bridges.completion_transformation import (
CompletionTransformationBridge,
@ -32,6 +35,7 @@ from litellm.responses.sse_output_recovery import (
)
from litellm.responses.utils import normalize_responses_api_stream_options
from litellm.types.llms.openai import (
REASONING_EFFORT,
ChatCompletionAnnotation,
ChatCompletionReasoningItem,
ChatCompletionToolCallChunk,
@ -55,9 +59,11 @@ if TYPE_CHECKING:
from litellm.types.llms.openai import (
ALL_RESPONSES_API_TOOL_PARAMS,
AllMessageValues,
ChatCompletionFileObject,
ChatCompletionImageObject,
ChatCompletionRedactedThinkingBlock,
ChatCompletionThinkingBlock,
ChatCompletionToolReferenceObject,
OpenAIMessageContentListBlock,
)
from litellm.types.utils import Choices
@ -85,6 +91,22 @@ def _get_reasoning_items(
return []
def _reasoning_input_items(msg: "AllMessageValues") -> list[dict[str, object]]: # mutable-ok: API message payload
"""Reasoning input items for an assistant message.
Stored reasoning items win because they carry an id the Responses API minted; thinking
blocks are the fallback for turns that arrived over another API surface.
"""
items: Final = _get_reasoning_items(msg)
stored: Final = [_reasoning_item_to_response_input(item) for item in items] # mutable-ok: API message payload
if stored:
return stored
raw_blocks: Final = msg.get("thinking_blocks") or ()
blocks: Final = cast("Iterable[ChatCompletionThinkingBlock]", raw_blocks) # cast-ok: untyped client json
from_thinking: Final = responses_reasoning_item_from_thinking_blocks(blocks)
return [] if from_thinking is None else [dict(from_thinking)] # mutable-ok: API message payload
def _build_reasoning_item(
item_id: str,
encrypted_content: str | None,
@ -155,6 +177,16 @@ def _map_incomplete_reason_to_finish_reason(incomplete_reason: str | None) -> Li
return "length"
def _input_file_from_file_value(file_value: object) -> dict[str, object]:
if not isinstance(file_value, dict):
return {"type": "input_file"}
file_dict: Final = cast("dict[str, object]", file_value) # cast-ok: runtime dict checked
return {
"type": "input_file",
**{key: file_dict[key] for key in ("file_id", "file_data", "filename") if key in file_dict},
}
def _incomplete_reason_from_response_payload(response_payload: object) -> str | None:
if not isinstance(response_payload, Mapping):
return None
@ -372,8 +404,15 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
)
)
elif role == "assistant" and tool_calls and isinstance(tool_calls, list):
for r_item in _get_reasoning_items(msg):
input_items.append(_reasoning_item_to_response_input(r_item))
input_items.extend(_reasoning_input_items(msg))
if content:
input_items.append(
{ # mutable-ok: API message payload
"type": "message",
"role": "assistant",
"content": self._convert_content_to_responses_format(content, "assistant"),
}
)
for tool_call in tool_calls:
function = tool_call.get("function")
custom = tool_call.get("custom")
@ -400,15 +439,16 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
raise ValueError(f"tool call not supported: {tool_call}")
elif content is not None:
if role == "assistant":
for r_item in _get_reasoning_items(msg):
input_items.append(_reasoning_item_to_response_input(r_item))
input_items.extend(_reasoning_input_items(msg))
input_items.append(
{
{ # mutable-ok: API message payload
"type": "message",
"role": role,
"content": self._convert_content_to_responses_format(content, cast(str, role)),
}
)
elif role == "assistant":
input_items.extend(_reasoning_input_items(msg))
return input_items, instructions
@ -929,7 +969,12 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
content: str
| list[object]
| Iterable[
Union["OpenAIMessageContentListBlock", "ChatCompletionThinkingBlock", "ChatCompletionRedactedThinkingBlock"]
Union[
"OpenAIMessageContentListBlock",
"ChatCompletionThinkingBlock",
"ChatCompletionRedactedThinkingBlock",
"ChatCompletionToolReferenceObject",
]
]
| None,
role: str,
@ -978,17 +1023,15 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
result.append(converted)
verbose_logger.debug("Chat provider: image -> %s", converted)
elif item_type == "file":
# Map Chat Completion file to Responses API input_file
# {"type": "file", "file": {"file_data": "...", "filename": "..."}}
# -> {"type": "input_file", "file_data": "...", "filename": "..."}
file_data = item.get("file", {})
converted = {"type": "input_file"}
if isinstance(file_data, dict):
for key in ["file_id", "file_data", "filename"]:
if key in file_data:
converted[key] = file_data[key]
converted = _input_file_from_file_value(
cast("ChatCompletionFileObject", item).get("file"), # cast-ok: type tag checked
)
result.append(converted)
verbose_logger.debug("Chat provider: file -> %s", converted)
elif item_type == "tool_reference":
verbose_logger.debug(
"Chat provider: tool_reference has no responses API equivalent; skipped"
)
elif item_type in [
"input_text",
"input_image",
@ -1086,22 +1129,11 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
litellm.reasoning_auto_summary or os.getenv("LITELLM_REASONING_AUTO_SUMMARY", "false").lower() == "true"
)
# If string is passed, map with optional summary based on flag/env var
if reasoning_effort == "none":
return Reasoning(effort="none", summary="detailed") if auto_summary_enabled else Reasoning(effort="none")
elif reasoning_effort == "high":
return Reasoning(effort="high", summary="detailed") if auto_summary_enabled else Reasoning(effort="high")
elif reasoning_effort == "xhigh":
return Reasoning(effort="xhigh", summary="detailed") if auto_summary_enabled else Reasoning(effort="xhigh")
elif reasoning_effort == "medium":
if reasoning_effort in get_args(REASONING_EFFORT):
return (
Reasoning(effort="medium", summary="detailed") if auto_summary_enabled else Reasoning(effort="medium")
)
elif reasoning_effort == "low":
return Reasoning(effort="low", summary="detailed") if auto_summary_enabled else Reasoning(effort="low")
elif reasoning_effort == "minimal":
return (
Reasoning(effort="minimal", summary="detailed") if auto_summary_enabled else Reasoning(effort="minimal")
Reasoning(effort=reasoning_effort, summary="detailed")
if auto_summary_enabled
else Reasoning(effort=reasoning_effort)
)
return None

View file

@ -35,6 +35,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.
@ -48,6 +49,9 @@ LITELLM_MAX_STREAMING_DURATION_SECONDS: Final = (
# Data URIs exceeding this are replaced with a size placeholder.
# Set to 0 to disable truncation.
MAX_BASE64_LENGTH_FOR_LOGGING: Final = int(os.getenv("MAX_BASE64_LENGTH_FOR_LOGGING", 64))
REDACTED_BY_LITELLM: Final = "redacted-by-litellm"
# in-memory stand-in handed to provider converters for redacted arguments; never stored
REDACTED_TOOL_CALL_ARGUMENTS_PLACEHOLDER: Final = "{}"
MAX_STRING_LENGTH_STDOUT_LOG: Final = get_env_int("MAX_STRING_LENGTH_STDOUT_LOG", 4096)
@ -146,6 +150,7 @@ LITELLM_UI_ALLOW_HEADERS: Final = [
"x-litellm-adaptive-router-model",
"x-litellm-applied-guardrails",
"x-litellm-guardrail-scan-id",
"x-litellm-cache-key",
]
# Gemini model-specific minimal thinking budget constants
@ -284,6 +289,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))
@ -292,6 +298,9 @@ GUARDRAIL_SCANNED_MESSAGES_CACHE_TTL_SECONDS: Final = int(
os.getenv("GUARDRAIL_SCANNED_MESSAGES_CACHE_TTL_SECONDS", 24 * 60 * 60)
)
BEDROCK_APPLY_GUARDRAIL_CHUNK_BUDGET_CHARS: Final = 25_000
DEFAULT_PRESIDIO_ANALYZE_CHUNK_SIZE_BYTES: Final = 500_000
PRESIDIO_ANALYZE_CHUNK_OVERLAP_CHARS: Final = 4096
PRESIDIO_ANALYZE_CHUNK_CONCURRENCY: Final = 8
# Aggregation threshold: default to 80% of the asyncio queue maxsize so the check can always trigger.
# Must be < LITELLM_ASYNCIO_QUEUE_MAXSIZE; if set higher the aggregation logic will never fire.
MAX_SIZE_IN_MEMORY_QUEUE: Final = int(os.getenv("MAX_SIZE_IN_MEMORY_QUEUE", int(LITELLM_ASYNCIO_QUEUE_MAXSIZE * 0.8)))
@ -377,6 +386,7 @@ AZURE_OPERATION_POLLING_TIMEOUT: Final = int(os.getenv("AZURE_OPERATION_POLLING_
AZURE_DOCUMENT_INTELLIGENCE_API_VERSION: Final = str(os.getenv("AZURE_DOCUMENT_INTELLIGENCE_API_VERSION", "2024-11-30"))
AZURE_DOCUMENT_INTELLIGENCE_DEFAULT_DPI: Final = int(os.getenv("AZURE_DOCUMENT_INTELLIGENCE_DEFAULT_DPI", 96))
REDIS_SOCKET_TIMEOUT: Final = float(os.getenv("REDIS_SOCKET_TIMEOUT", 0.1))
CACHE_WRITE_SHUTDOWN_FLUSH_TIMEOUT_SECONDS: Final[float] = 5.0
REDIS_CONNECTION_POOL_TIMEOUT: Final = int(os.getenv("REDIS_CONNECTION_POOL_TIMEOUT", 5))
REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD: Final = int(os.getenv("REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD", 5))
REDIS_CIRCUIT_BREAKER_RECOVERY_TIMEOUT: Final = int(os.getenv("REDIS_CIRCUIT_BREAKER_RECOVERY_TIMEOUT", 60))
@ -460,6 +470,8 @@ CONNECTION_ERROR_PATTERNS: Final[list[str]] = [
]
STREAM_SSE_DONE_STRING: Final[str] = "[DONE]"
STREAM_SSE_DATA_PREFIX: Final[str] = "data: "
STREAM_SSE_KEEPALIVE_PING_CHUNK: Final[str] = 'event: ping\ndata: {"type": "ping"}\n\n'
STREAM_SSE_KEEPALIVE_PING_BYTES: Final[bytes] = STREAM_SSE_KEEPALIVE_PING_CHUNK.encode("utf-8")
### SPEND TRACKING ###
DEFAULT_REPLICATE_GPU_PRICE_PER_SECOND: Final = float(
os.getenv("DEFAULT_REPLICATE_GPU_PRICE_PER_SECOND", 0.001400)
@ -619,6 +631,15 @@ LITELLM_CHAT_PROVIDERS: Final = [
"amazon_nova",
]
# Resolving these providers runs an OAuth device flow (their provider info IS the login), so any
# metadata or capability lookup against them can block for minutes waiting on a human.
PROVIDERS_THAT_AUTHENTICATE_ON_PROVIDER_INFO: Final = frozenset(
{
"github_copilot",
"chatgpt",
}
)
LITELLM_EMBEDDING_PROVIDERS_SUPPORTING_INPUT_ARRAY_OF_TOKENS: Final = [
"openai",
"azure",
@ -749,6 +770,7 @@ openai_compatible_endpoints: Final[list] = [
"api.groq.com/openai/v1",
"https://integrate.api.nvidia.com/v1",
"api.deepseek.com/v1",
"api.together.ai/v1",
"api.together.xyz/v1",
"app.empower.dev/api/v1",
"https://api.friendli.ai/serverless/v1",
@ -1359,8 +1381,6 @@ X_LITELLM_DISABLE_CALLBACKS: Final = "x-litellm-disable-callbacks"
LITELLM_METADATA_FIELD: Final = "litellm_metadata"
OLD_LITELLM_METADATA_FIELD: Final = "metadata"
RETURN_RAW_MODEL_NAME_METADATA_KEY: Final = "_complexity_router_return_raw_model_name"
AUTO_ROUTED_REQUEST_METADATA_KEY: Final = "_auto_routed_request"
ROUTER_MODEL_NAME_RESPONSE_FIELD: Final = "router_model_name"
SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY: Final = "_session_deployment_affinity_ttl"
CONSUMED_REQUEST_TAGS_METADATA_KEY: Final = "_consumed_request_tags"
INTERNAL_CALL_ORIGIN_METADATA_KEY: Final = "internal_call_origin"
@ -1470,6 +1490,12 @@ LITELLM_PROXY_MASTER_KEY_ALIAS: Final = "litellm_proxy_master_key"
# ``ProxyLogging._handle_logging_proxy_only_error``.
LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL: Final = "litellm_no_upstream_llm_call"
# Key/team metadata fields naming the OTel Resource ``service.name``, highest
# precedence first. Shared between the OTel v2 tenant router (which reads them
# out of ``user_api_key_auth_metadata``) and proxy request setup (which re-applies
# the key's values after the team metadata merge so a key outranks its team).
OTEL_SERVICE_NAME_METADATA_KEYS: Final = ("otel_service_name_override", "otel_service_name")
# Key Rotation Constants
LITELLM_KEY_ROTATION_ENABLED: Final = os.getenv("LITELLM_KEY_ROTATION_ENABLED", "false")
LITELLM_KEY_ROTATION_CHECK_INTERVAL_SECONDS: Final = int(
@ -1545,6 +1571,8 @@ SPEND_LOG_WRITE_BATCH_MAX_ROWS: Final = max(1, int(os.getenv("SPEND_LOG_WRITE_BA
SPEND_LOG_QUEUE_SIZE_THRESHOLD: Final = int(os.getenv("SPEND_LOG_QUEUE_SIZE_THRESHOLD", 100))
SPEND_LOG_QUEUE_MAX_BYTES: Final = max(1, int(os.getenv("SPEND_LOG_QUEUE_MAX_BYTES", "64000000")))
SPEND_LOG_QUEUE_POLL_INTERVAL: Final = float(os.getenv("SPEND_LOG_QUEUE_POLL_INTERVAL", 2.0))
RESPONSES_SESSION_LOOKUP_MAX_ATTEMPTS: Final = max(1, int(os.getenv("RESPONSES_SESSION_LOOKUP_MAX_ATTEMPTS", "3")))
RESPONSES_SESSION_LOOKUP_RETRY_INTERVAL: Final = float(os.getenv("RESPONSES_SESSION_LOOKUP_RETRY_INTERVAL", "0.2"))
SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE: Final = int(os.getenv("SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE", 10000))
DEFAULT_CRON_JOB_LOCK_TTL_SECONDS: Final = int(os.getenv("DEFAULT_CRON_JOB_LOCK_TTL_SECONDS", 60)) # 1 minute
PROXY_BUDGET_RESCHEDULER_MIN_TIME: Final = int(os.getenv("PROXY_BUDGET_RESCHEDULER_MIN_TIME", 597))
@ -1564,6 +1592,19 @@ STALE_OBJECT_CLEANUP_BATCH_SIZE: Final = max(1, int(os.getenv("STALE_OBJECT_CLEA
# installations with large numbers of stale managed objects).
_batch_polling_env: Final = os.getenv("PROXY_BATCH_POLLING_ENABLED", "true").lower()
PROXY_BATCH_POLLING_ENABLED: Final = _batch_polling_env == "true"
BACKGROUND_INTERACTION_COST_POLL_INITIAL_INTERVAL_SECONDS: Final = float(
os.getenv("BACKGROUND_INTERACTION_COST_POLL_INITIAL_INTERVAL_SECONDS", "5")
)
BACKGROUND_INTERACTION_COST_POLL_MAX_INTERVAL_SECONDS: Final = float(
os.getenv("BACKGROUND_INTERACTION_COST_POLL_MAX_INTERVAL_SECONDS", "60")
)
BACKGROUND_INTERACTION_COST_POLL_TIMEOUT_SECONDS: Final = float(
os.getenv("BACKGROUND_INTERACTION_COST_POLL_TIMEOUT_SECONDS", "3600")
)
_background_interaction_cost_polling_env: Final = os.getenv(
"BACKGROUND_INTERACTION_COST_POLLING_ENABLED", "true"
).lower()
BACKGROUND_INTERACTION_COST_POLLING_ENABLED: Final = _background_interaction_cost_polling_env == "true"
PROXY_BUDGET_RESCHEDULER_MAX_TIME: Final = int(os.getenv("PROXY_BUDGET_RESCHEDULER_MAX_TIME", 605))
PROXY_BATCH_WRITE_AT: Final = int(os.getenv("PROXY_BATCH_WRITE_AT", 10)) # in seconds, increased from 10
PROXY_CONFIG_RELOAD_INTERVAL_SECONDS: Final = get_env_int("PROXY_CONFIG_RELOAD_INTERVAL_SECONDS", 30)
@ -1628,6 +1669,7 @@ LITELLM_SETTINGS_SAFE_DB_OVERRIDES: Final = [
"enable_anthropic_prompt_caching",
"anthropic_prompt_caching_ttl",
"max_ui_session_budget",
"budget_rollover",
]
SPECIAL_LITELLM_AUTH_TOKEN: Final = ["ui-token"]
DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL = int(os.getenv("DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL", 60))
@ -1644,6 +1686,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.
@ -1794,6 +1837,43 @@ BROWSER_SECURITY_HEADERS: Final[frozenset[str]] = frozenset(
UNSAFE_PROXY_RESPONSE_HEADERS: Final[frozenset[str]] = HTTP_FRAMING_HEADERS | BROWSER_SECURITY_HEADERS
# A retrieved response replays the usage of the call that created it, so pricing these
# read/management routes like inference bills the same tokens twice.
NON_INFERENCE_CALL_TYPES: Final[frozenset[str]] = frozenset(
{
"get_responses",
"aget_responses",
"delete_responses",
"adelete_responses",
"cancel_responses",
"acancel_responses",
"list_input_items",
"alist_input_items",
"vector_store_create",
"avector_store_create",
"vector_store_retrieve",
"avector_store_retrieve",
"vector_store_list",
"avector_store_list",
"vector_store_update",
"avector_store_update",
"vector_store_delete",
"avector_store_delete",
"vector_store_file_create",
"avector_store_file_create",
"vector_store_file_list",
"avector_store_file_list",
"vector_store_file_retrieve",
"avector_store_file_retrieve",
"vector_store_file_content",
"avector_store_file_content",
"vector_store_file_update",
"avector_store_file_update",
"vector_store_file_delete",
"avector_store_file_delete",
}
)
# PTU reservation rollup writes rows to LiteLLM_DailyTeamSpend with this
# sentinel api_key so PTU flat cost stays distinguishable from real per-request
# spend under the table's composite unique constraint.

View file

@ -2,6 +2,7 @@
## File for 'response_cost' calculation in Logging
import logging
import time
from collections.abc import Mapping, Sequence
from functools import lru_cache
from typing import TYPE_CHECKING, Any, Final, Literal, cast
@ -19,6 +20,7 @@ from litellm.litellm_core_utils.llm_cost_calc.tool_call_cost_tracking import (
StandardBuiltInToolCostTracking,
)
from litellm.litellm_core_utils.llm_cost_calc.usage_object_transformation import (
InteractionsUsageObjectTransformation,
TranscriptionUsageObjectTransformation,
)
from litellm.litellm_core_utils.llm_cost_calc.utils import (
@ -74,7 +76,10 @@ from litellm.llms.perplexity.cost_calculator import (
from litellm.llms.tencent.cost_calculator import (
cost_per_token as tencent_cost_per_token,
)
from litellm.llms.together_ai.cost_calculator import get_model_params_and_category
from litellm.llms.together_ai.cost_calculator import (
get_model_params_and_category,
has_together_registry_pricing,
)
from litellm.llms.vertex_ai.cost_calculator import (
cost_per_character as google_cost_per_character,
)
@ -150,6 +155,7 @@ _VIDEO_CALL_TYPES: Final = frozenset(
}
)
_SPEECH_CALL_TYPES: Final = frozenset(
{
CallTypes.speech.value,
@ -554,9 +560,10 @@ def cost_per_token(
)
elif call_type == "atranscription" or call_type == "transcription":
if _transcription_usage_has_token_details(usage_block):
return openai_cost_per_token(
return generic_cost_per_token(
model=model_without_prefix,
usage=usage_block,
custom_llm_provider=custom_llm_provider,
service_tier=service_tier,
data_residency=data_residency,
)
@ -589,6 +596,7 @@ def cost_per_token(
prompt_characters=prompt_characters,
completion_characters=completion_characters,
usage=usage_block,
service_tier=service_tier,
vertex_location=vertex_location,
)
elif cost_router == "cost_per_token":
@ -731,6 +739,13 @@ def _get_provider_for_cost_calc(
return custom_llm_provider
def _get_hidden_str_for_cost_calc(hidden_params: object, key: str) -> str | None:
if not isinstance(hidden_params, Mapping):
return None
value: Final[object] = hidden_params.get(key)
return value if isinstance(value, str) and value else None
def _select_model_name_for_cost_calc(
model: str | None,
completion_response: object | None,
@ -747,7 +762,6 @@ def _select_model_name_for_cost_calc(
"""
return_model: str | None = None
region_name: str | None = None
custom_llm_provider = _get_provider_for_cost_calc(model=model, custom_llm_provider=custom_llm_provider)
completion_response_model: str | None = None
@ -757,6 +771,14 @@ def _select_model_name_for_cost_calc(
elif isinstance(completion_response, dict):
completion_response_model = completion_response.get("model", None)
hidden_params: Final[dict | None] = getattr(completion_response, "_hidden_params", None)
provider_response_model: Final = _get_hidden_str_for_cost_calc(hidden_params, "provider_response_model")
explicit_pricing: Final = custom_pricing is True or base_model is not None
priced_from_response: Final = provider_response_model is not None or completion_response_model is not None
region_name: Final = (
_get_hidden_str_for_cost_calc(hidden_params, "region_name")
if not explicit_pricing and priced_from_response
else None
)
if custom_pricing is True:
if router_model_id is not None and router_model_id in litellm.model_cost:
@ -772,14 +794,12 @@ def _select_model_name_for_cost_calc(
else:
return_model = model
elif base_model is not None:
return_model = base_model
elif base_model is not None or provider_response_model is not None:
return_model = base_model if base_model is not None else provider_response_model
elif completion_response_model is None and hidden_params is not None:
if hidden_params.get("model", None) is not None and len(hidden_params["model"]) > 0:
return_model = hidden_params.get("model", model)
elif hidden_params is not None and hidden_params.get("region_name", None) is not None:
region_name = hidden_params.get("region_name", None)
if return_model is None and completion_response_model is not None:
return_model = completion_response_model
@ -792,14 +812,27 @@ def _select_model_name_for_cost_calc(
and custom_llm_provider is not None
and not _model_contains_known_llm_provider(return_model)
): # add provider prefix if not already present, to match model_cost
if region_name is not None:
return_model = f"{custom_llm_provider}/{region_name}/{return_model}"
else:
return_model = f"{custom_llm_provider}/{return_model}"
provider_prefix: Final = custom_llm_provider if region_name is None else f"{custom_llm_provider}/{region_name}"
return_model = _strip_unregistered_leading_segments(f"{provider_prefix}/{return_model}", region_name)
return return_model
def _strip_unregistered_leading_segments(model: str, region_name: str | None) -> str:
"""Resolve a provider-prefixed slash alias like "vertex_ai/vertex/claude-opus-5" to the
registered cost key ("vertex_ai/claude-opus-5"), keeping the model unchanged when it already
resolves downstream (custom-priced router ids) or no stripped candidate is registered (#38069)."""
segments: Final = model.split("/")
if "/".join(segments[1:]) in litellm.model_cost:
return model
head_len: Final = 2 if region_name is not None and len(segments) > 2 and segments[1] == region_name else 1
head: Final = "/".join(segments[:head_len])
tail: Final = segments[head_len:]
strippable: Final = next((index for index, segment in enumerate(tail) if segment in LlmProvidersSet), len(tail))
candidates: Final = (f"{head}/{'/'.join(tail[start:])}" for start in range(min(strippable, len(tail) - 1) + 1))
return next((candidate for candidate in candidates if candidate in litellm.model_cost), model)
@lru_cache(maxsize=DEFAULT_MAX_LRU_CACHE_SIZE)
def _model_contains_known_llm_provider(model: str) -> bool:
"""
@ -830,9 +863,11 @@ def _get_response_model(completion_response: object) -> str | None:
_GEMINI_TRAFFIC_TYPE_TO_SERVICE_TIER: Final[dict] = {
# ON_DEMAND_PRIORITY maps to "priority" — selects input_cost_per_token_priority, etc.
"ON_DEMAND_PRIORITY": "priority",
# FLEX / BATCH maps to "flex" — selects input_cost_per_token_flex, etc.
# FLEX / BATCH / ON_DEMAND_FLEX maps to "flex" — selects input_cost_per_token_flex, etc.
# Vertex AI reports flex/shared-capacity traffic as ON_DEMAND_FLEX, not FLEX.
"FLEX": "flex",
"BATCH": "flex",
"ON_DEMAND_FLEX": "flex",
# ON_DEMAND is standard pricing — no service_tier suffix applied
"ON_DEMAND": None,
}
@ -847,9 +882,9 @@ def _map_traffic_type_to_service_tier(traffic_type: str | None) -> str | None:
trafficType values seen in practice
------------------------------------
ON_DEMAND -> standard pricing (service_tier = None)
ON_DEMAND_PRIORITY -> priority pricing (service_tier = "priority")
FLEX / BATCH -> batch/flex pricing (service_tier = "flex")
ON_DEMAND -> standard pricing (service_tier = None)
ON_DEMAND_PRIORITY -> priority pricing (service_tier = "priority")
FLEX / BATCH / ON_DEMAND_FLEX -> batch/flex pricing (service_tier = "flex")
"""
if traffic_type is None:
return None
@ -912,6 +947,8 @@ def _get_usage_object(
usage_obj,
)
)
elif isinstance(usage_obj, dict) and InteractionsUsageObjectTransformation.is_interactions_usage_object(usage_obj):
return InteractionsUsageObjectTransformation.transform_interactions_usage_object(usage_obj)
elif isinstance(usage_obj, dict):
return Usage(**usage_obj)
elif isinstance(usage_obj, BaseModel):
@ -1288,6 +1325,10 @@ def completion_cost(
)
if tr_usage is not None:
_usage = tr_usage.model_dump()
elif InteractionsUsageObjectTransformation.is_interactions_usage_object(_usage):
_usage = InteractionsUsageObjectTransformation.transform_interactions_usage_object(
_usage
).model_dump()
else:
_usage = _usage
@ -1372,23 +1413,36 @@ def completion_cost(
if custom_pricing and litellm_logging_obj is not None:
_litellm_params = getattr(litellm_logging_obj, "litellm_params", None)
if _litellm_params is not None:
_metadata = _litellm_params.get("metadata", {}) or {}
_video_model_info = _metadata.get("model_info", None)
_video_model_info = next(
(
model_info
for _metadata_key in ("metadata", "litellm_metadata")
if (model_info := (_litellm_params.get(_metadata_key) or {}).get("model_info"))
is not None
),
None,
)
usage_obj = getattr(completion_response, "usage", None)
duration_seconds: float | None = None
video_resolution: str | None = None
provider_reported_cost: float | None = None
if completion_response is not None and usage_obj:
# Handle both dict and Pydantic Usage object
if isinstance(usage_obj, dict):
duration_seconds = usage_obj.get("duration_seconds", None)
_vr = usage_obj.get("video_resolution", None)
provider_reported_cost = usage_obj.get("provider_reported_cost_usd", None)
else:
duration_seconds = getattr(usage_obj, "duration_seconds", None)
_vr = getattr(usage_obj, "video_resolution", None)
provider_reported_cost = getattr(usage_obj, "provider_reported_cost_usd", None)
if _vr is not None:
video_resolution = str(_vr).strip().lower()
if _video_model_info is None and provider_reported_cost is not None:
return float(provider_reported_cost)
if duration_seconds is not None:
# Calculate cost based on video duration using video-specific cost calculation
from litellm.llms.openai.cost_calculation import (
@ -1530,10 +1584,9 @@ def completion_cost(
return MCPCostCalculator.calculate_mcp_tool_call_cost(litellm_logging_obj=litellm_logging_obj)
# Calculate cost based on prompt_tokens, completion_tokens
if "togethercomputer" in model or "together_ai" in model or custom_llm_provider == "together_ai":
# together ai prices based on size of llm
# get_model_params_and_category takes a model name and returns the category of LLM size it is in model_prices_and_context_window.json
if (
"togethercomputer" in model or "together_ai" in model or custom_llm_provider == "together_ai"
) and not has_together_registry_pricing(model, litellm.model_cost):
model = get_model_params_and_category(model, call_type=CallTypes(call_type))
# replicate llms are calculate based on time for request running
@ -2336,6 +2389,64 @@ class RealtimeAPITokenUsageProcessor(BaseTokenUsageProcessor):
_TRANSCRIPTION_COMPLETED_EVENT_TYPE: Final = "conversation.item.input_audio_transcription.completed"
def _candidate_realtime_token_costs(
model_name: str,
combined_usage_object: Usage,
custom_llm_provider: str,
data_residency: str | None,
) -> tuple[float, float] | None:
try:
return generic_cost_per_token(
model=model_name,
usage=combined_usage_object,
custom_llm_provider=custom_llm_provider,
data_residency=data_residency,
)
except Exception:
return None
def _cost_map_entry_declares_pricing(model_name: str, custom_llm_provider: str) -> bool:
entries: Final = (
litellm.model_cost.get(model_name),
litellm.model_cost.get(f"{custom_llm_provider}/{model_name}"),
)
return any(
entry is not None and any("cost_per" in field and value is not None for field, value in entry.items())
for entry in entries
)
def _first_priced_realtime_token_costs(
potential_model_names: Sequence[str | None],
combined_usage_object: Usage,
custom_llm_provider: str,
data_residency: str | None,
) -> tuple[float, float]:
candidate_costs: Final = (
(model_name, costs)
for model_name in potential_model_names
if model_name is not None
and (
costs := _candidate_realtime_token_costs(
model_name=model_name,
combined_usage_object=combined_usage_object,
custom_llm_provider=custom_llm_provider,
data_residency=data_residency,
)
)
is not None
)
return next(
(
costs
for model_name, costs in candidate_costs
if sum(costs) > 0 or _cost_map_entry_declares_pricing(model_name, custom_llm_provider)
),
(0.0, 0.0),
)
def handle_realtime_stream_cost_calculation(
results: OpenAIRealtimeStreamList,
combined_usage_object: Usage,
@ -2360,24 +2471,12 @@ def handle_realtime_stream_cost_calculation(
potential_model_names.append(received_model)
potential_model_names.append(litellm_model_name)
input_cost_per_token = 0.0
output_cost_per_token = 0.0
for model_name in potential_model_names:
try:
if model_name is None:
continue
_input_cost_per_token, _output_cost_per_token = generic_cost_per_token(
model=model_name,
usage=combined_usage_object,
custom_llm_provider=custom_llm_provider,
data_residency=data_residency,
)
except Exception:
continue
input_cost_per_token += _input_cost_per_token
output_cost_per_token += _output_cost_per_token
break # exit if we find a valid model
input_cost_per_token, output_cost_per_token = _first_priced_realtime_token_costs(
potential_model_names=potential_model_names,
combined_usage_object=combined_usage_object,
custom_llm_provider=custom_llm_provider,
data_residency=data_residency,
)
transcription_cost: Final = (
handle_realtime_transcription_cost_calculation(
results=results,

View file

@ -8,6 +8,14 @@ if TYPE_CHECKING:
from litellm.types.utils import ModelResponse
def _completion_response_cost(model_response: "ModelResponse") -> float | None:
hidden_params: Final = getattr(model_response, "_hidden_params", None)
if not isinstance(hidden_params, dict):
return None
response_cost: Final = hidden_params.get("response_cost")
return response_cost if isinstance(response_cost, float) else None
class SpeechToCompletionBridgeTransformationHandler:
def transform_request(
self,
@ -123,4 +131,6 @@ class SpeechToCompletionBridgeTransformationHandler:
# Create an httpx.Response object
response: Final = httpx.Response(status_code=200, content=binary_data, headers=headers)
return HttpxBinaryResponseContent(response)
binary_response: Final = HttpxBinaryResponseContent(response)
binary_response.set_response_cost(_completion_response_cost(model_response))
return binary_response

View file

@ -7,6 +7,7 @@ import base64
import os
from collections.abc import Awaitable, Callable, Generator
from datetime import timedelta
from importlib import metadata
from typing import Any, Final, TypeVar
import httpx
@ -21,6 +22,18 @@ try:
streamable_http_client = getattr(streamable_http_module, "streamable_http_client", None)
except ImportError:
pass
MCP_STREAMABLE_HTTP_REQUIREMENT: Final = "mcp>=1.28.1"
def missing_streamable_http_client_error() -> ImportError:
return ImportError(
f"MCP streamable HTTP transport requires {MCP_STREAMABLE_HTTP_REQUIREMENT}, but the installed "
f"mcp {metadata.version('mcp')} does not provide streamable_http_client. "
"Fix with: pip install 'litellm[mcp]' (or upgrade mcp directly: pip install -U mcp)"
)
from mcp.types import CallToolRequestParams as MCPCallToolRequestParams
from mcp.types import CallToolResult as MCPCallToolResult
from mcp.types import (
@ -43,6 +56,9 @@ from litellm.types.mcp import (
MCPStdioConfig,
MCPTransport,
MCPTransportType,
credential_redirect_hook,
has_header,
without_header,
)
@ -260,6 +276,7 @@ class MCPClient:
transport_type: MCPTransportType = MCPTransport.http,
auth_type: MCPAuthType = None,
auth_value: str | dict[str, str] | None = None,
auth_header_name: str | None = None,
timeout: float | None = None,
stdio_config: MCPStdioConfig | None = None,
extra_headers: dict[str, str] | None = None,
@ -275,6 +292,11 @@ class MCPClient:
self.auth_type: MCPAuthType = auth_type
self.timeout: float = timeout if timeout is not None else MCP_CLIENT_TIMEOUT
self._mcp_auth_value: str | dict[str, str] | None = None
# The one place this client decides which header its credential occupies: the operator's
# configured slot on the v1 path, or the slot the v2 resolver's auth object already owns.
# Every consumer reads this rather than re-deriving it, since each re-derivation so far
# picked up a different bug.
self._credential_slot: str | None = auth_header_name or getattr(resolved_auth, "header_name", None)
self.stdio_config: MCPStdioConfig | None = stdio_config
self.extra_headers: dict[str, str] | None = extra_headers
self.ssl_verify: VerifyTypes | None = ssl_verify
@ -323,7 +345,7 @@ class MCPClient:
)
# HTTP transport (default)
if streamable_http_client is None:
raise ImportError("streamable_http_client is not available. Please install mcp with HTTP support.")
raise missing_streamable_http_client_error()
headers = self._get_auth_headers()
httpx_client_factory = self._create_httpx_client_factory()
verbose_logger.debug("litellm headers for streamable_http_client: %s", headers)
@ -488,26 +510,33 @@ class MCPClient:
else:
self._mcp_auth_value = mcp_auth_value
def _header_slot(self, default: str) -> str:
return self._credential_slot or default
def _get_auth_headers(self) -> dict:
"""Generate authentication headers based on auth type."""
headers: Final = {}
if self._mcp_auth_value:
if isinstance(self._mcp_auth_value, str):
if self.auth_type == MCPAuth.bearer_token:
headers["Authorization"] = f"Bearer {strip_auth_scheme(self._mcp_auth_value, 'Bearer')}"
static_bearer: Final = strip_auth_scheme(self._mcp_auth_value, "Bearer")
headers[self._header_slot("Authorization")] = f"Bearer {static_bearer}"
elif self.auth_type == MCPAuth.basic:
headers["Authorization"] = f"Basic {self._mcp_auth_value}"
headers[self._header_slot("Authorization")] = f"Basic {self._mcp_auth_value}"
elif self.auth_type == MCPAuth.api_key:
headers["X-API-Key"] = self._mcp_auth_value
headers[self._header_slot("X-API-Key")] = self._mcp_auth_value
elif self.auth_type == MCPAuth.authorization:
# This auth type means the caller owns the whole header value.
headers["Authorization"] = self._mcp_auth_value
headers[self._header_slot("Authorization")] = self._mcp_auth_value
elif self.auth_type == MCPAuth.oauth2:
headers["Authorization"] = f"Bearer {strip_auth_scheme(self._mcp_auth_value, 'Bearer')}"
oauth2_bearer: Final = strip_auth_scheme(self._mcp_auth_value, "Bearer")
headers[self._header_slot("Authorization")] = f"Bearer {oauth2_bearer}"
elif self.auth_type == MCPAuth.token:
headers["Authorization"] = f"token {strip_auth_scheme(self._mcp_auth_value, 'token')}"
scheme_token: Final = strip_auth_scheme(self._mcp_auth_value, "token")
headers[self._header_slot("Authorization")] = f"token {scheme_token}"
elif self.auth_type == MCPAuth.oauth2_token_exchange:
headers["Authorization"] = f"Bearer {strip_auth_scheme(self._mcp_auth_value, 'Bearer')}"
exchanged_bearer: Final = strip_auth_scheme(self._mcp_auth_value, "Bearer")
headers[self._header_slot("Authorization")] = f"Bearer {exchanged_bearer}"
elif isinstance(self._mcp_auth_value, dict):
headers.update(self._mcp_auth_value)
# Note: aws_sigv4 auth is not handled here — SigV4 requires per-request
@ -515,7 +544,14 @@ class MCPClient:
# of static headers. See MCPSigV4Auth and _create_httpx_client_factory().
# update the headers with the extra headers
if self.extra_headers:
headers.update(self.extra_headers)
# Mirrors _resolve_v2_auth: when the operator named a slot for the credential the
# gateway resolved, no injected header may shadow it, case-insensitively, since HTTP
# header names are. Without a configured slot the old precedence stands unchanged.
slot: Final = self._credential_slot
injected: Final = (
without_header(self.extra_headers, slot) if slot and has_header(headers, slot) else self.extra_headers
)
headers.update(injected or {})
return _strip_header_whitespace(headers)
def _create_httpx_client_factory(self) -> Callable[..., httpx.AsyncClient]:
@ -543,12 +579,14 @@ class MCPClient:
# SigV4 aws_auth. Both are None for the common case — no behavior change.
fallback_auth: Final = self._resolved_auth if self._resolved_auth is not None else self._aws_auth
effective_auth: Final = auth if auth is not None else fallback_auth
guard: Final = credential_redirect_hook(self.server_url, self._credential_slot)
return httpx.AsyncClient(
headers=headers,
timeout=timeout,
auth=effective_auth,
verify=ssl_config,
follow_redirects=True,
event_hooks={"request": [guard]} if guard else {},
)
return factory

View file

@ -2,6 +2,7 @@ import asyncio
from datetime import datetime
from typing import TYPE_CHECKING, Any, Final
import litellm
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.proxy.pass_through_endpoints.success_handler import (
PassThroughEndpointLogging,
@ -65,6 +66,7 @@ class BaseGoogleGenAIGenerateContentStreamingIterator:
litellm_logging_obj: LiteLLMLoggingObj,
request_body: dict,
model: str,
custom_llm_provider: str,
hidden_params: dict[str, Any] | None = None,
):
self.litellm_logging_obj = litellm_logging_obj
@ -72,6 +74,10 @@ class BaseGoogleGenAIGenerateContentStreamingIterator:
self.start_time = datetime.now()
self.collected_chunks: list[bytes] = []
self.model = model
self.custom_llm_provider = custom_llm_provider
self.endpoint_type: Final = (
EndpointType.GEMINI if custom_llm_provider == litellm.LlmProviders.GEMINI.value else EndpointType.VERTEX_AI
)
self._hidden_params: dict[str, Any] = hidden_params or {}
async def _handle_async_streaming_logging(
@ -89,7 +95,7 @@ class BaseGoogleGenAIGenerateContentStreamingIterator:
passthrough_success_handler_obj=GLOBAL_PASS_THROUGH_SUCCESS_HANDLER_OBJ,
url_route="/v1/generateContent",
request_body=self.request_body or {},
endpoint_type=EndpointType.VERTEX_AI,
endpoint_type=self.endpoint_type,
start_time=self.start_time,
raw_bytes=self.collected_chunks,
end_time=end_time,
@ -118,13 +124,13 @@ class GoogleGenAIGenerateContentStreamingIterator(BaseGoogleGenAIGenerateContent
litellm_logging_obj=logging_obj,
request_body=request_body or {},
model=model,
custom_llm_provider=custom_llm_provider,
hidden_params=hidden_params,
)
self.response = response
self.model = model
self.generate_content_provider_config = generate_content_provider_config
self.litellm_metadata = litellm_metadata
self.custom_llm_provider = custom_llm_provider
# Gemini streamGenerateContent uses SSE line framing; iter_lines keeps
# large inlineData payloads (e.g. image/jpeg) intact within one event.
self.stream_iterator = response.iter_lines()
@ -169,13 +175,13 @@ class AsyncGoogleGenAIGenerateContentStreamingIterator(BaseGoogleGenAIGenerateCo
litellm_logging_obj=logging_obj,
request_body=request_body or {},
model=model,
custom_llm_provider=custom_llm_provider,
hidden_params=hidden_params,
)
self.response = response
self.model = model
self.generate_content_provider_config = generate_content_provider_config
self.litellm_metadata = litellm_metadata
self.custom_llm_provider = custom_llm_provider
# Gemini streamGenerateContent uses SSE line framing; aiter_lines keeps
# large inlineData payloads (e.g. image/jpeg) intact within one event.
self.stream_iterator = response.aiter_lines()

View file

@ -19,6 +19,7 @@ from litellm.constants import request_timeout as DEFAULT_REQUEST_TIMEOUT
from litellm.exceptions import LiteLLMUnknownProvider
from litellm.litellm_core_utils.litellm_logging import Logging
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.litellm_core_utils.llm_request_utils import flatten_form_field_values
from litellm.litellm_core_utils.mock_functions import mock_image_generation
from litellm.llms.base_llm import BaseImageEditConfig, BaseImageGenerationConfig
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
@ -422,24 +423,32 @@ def image_generation(
aimg_generation=aimg_generation,
)
elif custom_llm_provider == "azure_ai":
from litellm.llms.azure_ai.common_utils import AzureFoundryModelInfo
from litellm.llms.azure_ai.common_utils import (
AzureFoundryModelInfo,
get_azure_ai_auth_headers,
)
api_base = AzureFoundryModelInfo.get_api_base(api_base)
api_key = AzureFoundryModelInfo.get_api_key(api_key)
if extra_headers is not None:
optional_params["extra_headers"] = extra_headers
default_headers = {
caller_header_names = frozenset(name.lower() for name in headers)
caller_set_auth = "api-key" in caller_header_names or "authorization" in caller_header_names
auth_headers = (
headers
if caller_set_auth
else get_azure_ai_auth_headers(
api_key=api_key,
litellm_params=litellm_params_dict,
api_key_header="api-key",
)
)
request_headers: Final = {
"Content-Type": "application/json",
**auth_headers,
**headers,
}
# Only add api-key header if api_key is not None
# Azure AD authentication will use Authorization header instead
if api_key is not None:
default_headers["api-key"] = api_key
for k, v in default_headers.items():
if k not in headers:
headers[k] = v
model_response = azure_chat_completions.image_generation(
model=model,
@ -455,7 +464,7 @@ def image_generation(
api_version=api_version,
aimg_generation=aimg_generation,
client=client,
headers=headers,
headers=request_headers,
litellm_params=litellm_params_dict,
)
elif (
@ -846,6 +855,18 @@ def image_edit(
additional_drop_params=kwargs.get("additional_drop_params"),
)
if (
custom_llm_provider == "openai"
or custom_llm_provider == "azure"
or custom_llm_provider in litellm.openai_compatible_providers
):
image_edit_request_params.update(
flatten_form_field_values(
non_default_params,
extra_body if isinstance(extra_body, dict) else None,
)
)
# Pre Call logging
litellm_logging_obj.update_from_kwargs(
kwargs=kwargs,
@ -995,6 +1016,9 @@ async def aimage_edit(
response_format=response_format,
size=size,
user=user,
extra_headers=extra_headers,
extra_query=extra_query,
extra_body=extra_body,
timeout=timeout,
custom_llm_provider=custom_llm_provider,
**kwargs,

View file

@ -10,6 +10,8 @@ from typing import TYPE_CHECKING, Any, Final
from litellm._logging import verbose_proxy_logger
from .ms_teams import MS_TEAMS_ALERTING_DESTINATION, build_ms_teams_payload
if TYPE_CHECKING:
from .slack_alerting import SlackAlerting as _SlackAlerting
@ -62,14 +64,17 @@ async def send_to_webhook(slackAlertingInstance: SlackAlertingType, item, count)
if count > 1:
payload["text"] = f"[Num Alerts: {count}]\n\n{payload['text']}"
request_body: Final = (
build_ms_teams_payload(payload["text"]) if item.get("format") == MS_TEAMS_ALERTING_DESTINATION else payload
)
response: Final = await slackAlertingInstance.async_http_handler.post(
url=item["url"],
headers=item["headers"],
data=json.dumps(payload),
data=json.dumps(request_body),
)
if response.status_code != 200:
verbose_proxy_logger.debug("Error sending slack alert to url=%s. Error=%s", item["url"], response.text)
verbose_proxy_logger.debug("Error sending alert to url=%s. Error=%s", item["url"], response.text)
except Exception as e:
verbose_proxy_logger.debug("Error sending slack alert: %s", e)
verbose_proxy_logger.debug("Error sending alert: %s", e)
finally:
_print_alerting_payload_warning(payload, slackAlertingInstance=slackAlertingInstance)

View file

@ -0,0 +1,75 @@
"""Microsoft Teams alert delivery helpers.
Teams incoming webhooks (Workflows and legacy connectors) accept an Adaptive
Card wrapped in a message attachment, so alert text is delivered as a single
wrapped TextBlock.
"""
import os
from collections.abc import Mapping
from types import MappingProxyType
from typing import Final
from typing_extensions import ReadOnly, TypedDict
from litellm.types.integrations.slack_alerting import AlertType
MS_TEAMS_WEBHOOK_URL_ENV: Final = "MS_TEAMS_WEBHOOK_URL"
MS_TEAMS_ALERTING_DESTINATION: Final = "ms_teams"
MS_TEAMS_ALERT_HEADERS: Final[Mapping[str, str]] = MappingProxyType({"Content-type": "application/json"})
class MSTeamsTextBlock(TypedDict):
type: ReadOnly[str]
text: ReadOnly[str]
wrap: ReadOnly[bool]
class MSTeamsAdaptiveCard(TypedDict):
type: ReadOnly[str]
version: ReadOnly[str]
body: ReadOnly[tuple[MSTeamsTextBlock, ...]]
class MSTeamsAttachment(TypedDict):
contentType: ReadOnly[str]
content: ReadOnly[MSTeamsAdaptiveCard]
class MSTeamsMessage(TypedDict):
type: ReadOnly[str]
attachments: ReadOnly[tuple[MSTeamsAttachment, ...]]
class MSTeamsAlertText(TypedDict):
text: ReadOnly[str]
class MSTeamsQueueItem(TypedDict):
url: ReadOnly[str]
headers: ReadOnly[Mapping[str, str]]
payload: ReadOnly[MSTeamsAlertText]
alert_type: ReadOnly[AlertType]
format: ReadOnly[str]
def get_ms_teams_webhook_url() -> str | None:
return os.getenv(MS_TEAMS_WEBHOOK_URL_ENV)
def build_ms_teams_payload(text: str) -> MSTeamsMessage:
return MSTeamsMessage(
type="message",
attachments=(
MSTeamsAttachment(
contentType="application/vnd.microsoft.card.adaptive",
content=MSTeamsAdaptiveCard(
type="AdaptiveCard",
version="1.4",
body=(MSTeamsTextBlock(type="TextBlock", text=text, wrap=True),),
),
),
),
)

View file

@ -57,6 +57,13 @@ from litellm.types.proxy.model_deprecation import (
from ..email_templates.templates import *
from .batching_handler import send_to_webhook, squash_payloads
from .ms_teams import (
MS_TEAMS_ALERT_HEADERS,
MS_TEAMS_ALERTING_DESTINATION,
MSTeamsAlertText,
MSTeamsQueueItem,
get_ms_teams_webhook_url,
)
from .utils import process_slack_alerting_variables
if TYPE_CHECKING:
@ -1431,13 +1438,43 @@ Model Info:
# only send budget alerts over Email
await self.send_email_alert_using_smtp(webhook_event=user_info, alert_type=alert_type)
if "slack" not in self.alerting:
send_to_slack: Final = "slack" in self.alerting
send_to_ms_teams: Final = MS_TEAMS_ALERTING_DESTINATION in self.alerting
if not send_to_slack and not send_to_ms_teams:
return
if alert_type not in self.alert_types:
return
from datetime import datetime
current_time: Final = datetime.now().strftime("%H:%M:%S")
_proxy_base_url: Final = os.getenv("PROXY_BASE_URL", None)
alert_type_name: Final = getattr(alert_type, "name", alert_type)
alert_type_formatted: Final = f"Alert type: `{alert_type_name}`"
if alert_type == "daily_reports" or alert_type == "new_model_added":
formatted_message = alert_type_formatted + message
else:
formatted_message = (
f"{alert_type_formatted}\nLevel: `{level}`\nTimestamp: `{current_time}`\n\nMessage: {message}"
)
if kwargs:
for key, value in kwargs.items():
formatted_message += f"\n\n{key}: `{value}`\n\n"
if alerting_metadata:
for key, value in alerting_metadata.items():
formatted_message += f"\n\n*Alerting Metadata*: \n{key}: `{value}`\n\n"
if _proxy_base_url is not None:
formatted_message += f"\n\nProxy URL: `{_proxy_base_url}`"
if send_to_ms_teams:
self._enqueue_ms_teams_alert(formatted_message=formatted_message, alert_type=alert_type)
if not send_to_slack:
if len(self.log_queue) >= self.batch_size:
await self.flush_queue()
return
# Check if digest mode is enabled for this alert type
alert_type_name_str: Final = getattr(alert_type, "value", str(alert_type))
_atc: Final = self.alert_type_config.get(alert_type_name_str)
@ -1473,28 +1510,6 @@ Model Info:
)
return # Suppress immediate alert; will be emitted by _flush_digest_buckets
# Get the current timestamp
current_time: Final = datetime.now().strftime("%H:%M:%S")
_proxy_base_url: Final = os.getenv("PROXY_BASE_URL", None)
# Use .name if it's an enum, otherwise use as is
alert_type_name: Final = getattr(alert_type, "name", alert_type)
alert_type_formatted: Final = f"Alert type: `{alert_type_name}`"
if alert_type == "daily_reports" or alert_type == "new_model_added":
formatted_message = alert_type_formatted + message
else:
formatted_message = (
f"{alert_type_formatted}\nLevel: `{level}`\nTimestamp: `{current_time}`\n\nMessage: {message}"
)
if kwargs:
for key, value in kwargs.items():
formatted_message += f"\n\n{key}: `{value}`\n\n"
if alerting_metadata:
for key, value in alerting_metadata.items():
formatted_message += f"\n\n*Alerting Metadata*: \n{key}: `{value}`\n\n"
if _proxy_base_url is not None:
formatted_message += f"\n\nProxy URL: `{_proxy_base_url}`"
# check if we find the slack webhook url in self.alert_to_webhook_url
if self.alert_to_webhook_url is not None and alert_type in self.alert_to_webhook_url:
slack_webhook_url: str | list[str] | None = self.alert_to_webhook_url[alert_type]
@ -1531,6 +1546,24 @@ Model Info:
if len(self.log_queue) >= self.batch_size:
await self.flush_queue()
def _enqueue_ms_teams_alert(self, formatted_message: str, alert_type: AlertType) -> None:
ms_teams_webhook_url: Final = get_ms_teams_webhook_url()
if ms_teams_webhook_url is None:
verbose_proxy_logger.error(
"MS Teams alerting is enabled but MS_TEAMS_WEBHOOK_URL is not set. Dropping alert type=%s",
alert_type,
)
return
payload: Final[MSTeamsAlertText] = {"text": formatted_message}
item: Final[MSTeamsQueueItem] = {
"url": ms_teams_webhook_url,
"headers": MS_TEAMS_ALERT_HEADERS,
"payload": payload,
"alert_type": alert_type,
"format": MS_TEAMS_ALERTING_DESTINATION,
}
self.log_queue.append(item)
async def async_send_batch(self):
if not self.log_queue:
return

View file

@ -24,6 +24,8 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import (
with_prompt_cache_breakpoint,
)
from litellm.types.integrations.anthropic_cache_control_hook import (
GATEWAY_INJECTED_CACHE_METADATA_KEY,
GATEWAY_INJECTED_FOR_EVERY_DEPLOYMENT,
CacheControlInjectionPoint,
CacheControlMessageInjectionPoint,
)
@ -104,6 +106,13 @@ def _accepts_prompt_cache_breakpoint(block: object) -> bool:
return isinstance(block, dict) and block.get("type") in OPENAI_PROMPT_CACHE_BREAKPOINT_BLOCK_TYPES
# Set by a caller whose message list is not the one that goes upstream -- today the
# Responses API layer, whose `instructions` only becomes a system message further down.
# Tells this hook to hand role-targeted points to the pass holding the final messages
# rather than spending them on a list that is still missing some of their targets.
CARRY_UNMATCHED_MESSAGE_POINTS: Final = "_litellm_carry_unmatched_cache_control_points"
class AnthropicCacheControlHook(CustomPromptManagement):
def get_chat_completion_prompt(
self,
@ -128,6 +137,7 @@ class AnthropicCacheControlHook(CustomPromptManagement):
- non_default_params: dict - params with any global cache controls
"""
# Extract cache control injection points
carry_unmatched: Final = bool(non_default_params.pop(CARRY_UNMATCHED_MESSAGE_POINTS, False))
injection_points: Final[list[CacheControlInjectionPoint]] = non_default_params.pop(
"cache_control_injection_points", []
)
@ -161,26 +171,44 @@ class AnthropicCacheControlHook(CustomPromptManagement):
non_default_params.get("prompt_cache_options"),
)
)
# A provisional message list defers every role-targeted point to the pass holding
# the final one: a role with no message here may have one there, and settling all
# of them in one pass is what lets config order decide the shared breakpoint
# budget. An ordinal names a different message once a later layer builds its own
# list, so it is placed here or not at all.
carried_message_points: Final[Sequence[CacheControlMessageInjectionPoint]] = (
tuple(point for point in message_points if point.get("index") is None) if carry_unmatched else ()
)
applied_message_points: Final[Sequence[CacheControlMessageInjectionPoint]] = (
tuple(point for point in message_points if point.get("index") is not None)
if carry_unmatched
else tuple(message_points)
)
reserved_blocks: Final = (
1 if not openai_dialect and any(p.get("location") == "tool_config" for p in remaining_points) else 0
)
breakpoints_before: Final = AnthropicCacheControlHook._count_request_cache_breakpoints(processed_messages)
breakpoints_before: Final = AnthropicCacheControlHook.count_request_cache_breakpoints(processed_messages)
processed_messages = self._apply_message_injections(
points=message_points,
points=applied_message_points,
messages=processed_messages,
max_blocks=MAX_CACHE_CONTROL_BLOCKS - reserved_blocks,
openai_dialect=openai_dialect,
)
if (
openai_dialect
and AnthropicCacheControlHook._count_request_cache_breakpoints(processed_messages) > breakpoints_before
and AnthropicCacheControlHook.count_request_cache_breakpoints(processed_messages) > breakpoints_before
):
non_default_params.setdefault("prompt_cache_options", PromptCacheOptions(mode="explicit"))
# Pass through non-message injection points for provider-specific handling
if remaining_points:
# Points this pass did not place: non-message ones for the provider transform, and
# the deferred role-targeted ones. Deferring is what reaches the Responses API's
# `instructions`, which is only a system message once the bridge builds one. The
# judged stamp is what makes it safe: the next pass must not re-judge points
# against messages this pass already marked (see `_should_stand_down`).
carried_points: Final[Sequence[CacheControlInjectionPoint]] = (*remaining_points, *carried_message_points)
if carried_points:
non_default_params["cache_control_injection_points"] = AnthropicCacheControlHook._stamped_as_judged(
remaining_points
carried_points
)
return model, processed_messages, non_default_params
@ -210,7 +238,7 @@ class AnthropicCacheControlHook(CustomPromptManagement):
return provider
@staticmethod
def _count_request_cache_breakpoints(messages: Iterable[object], system: object = None) -> int:
def count_request_cache_breakpoints(messages: Iterable[object], system: object = None) -> int:
system_blocks: Final = (
sum(1 for block in system if _carries_cache_breakpoint(block)) if isinstance(system, list) else 0
)
@ -218,7 +246,7 @@ class AnthropicCacheControlHook(CustomPromptManagement):
@staticmethod
def _apply_message_injections(
points: list[CacheControlMessageInjectionPoint],
points: Sequence[CacheControlMessageInjectionPoint],
messages: list[AllMessageValues],
max_blocks: int,
openai_dialect: bool = False,
@ -232,7 +260,7 @@ class AnthropicCacheControlHook(CustomPromptManagement):
``max_blocks`` is reached. Injection points are honored in config order,
so earlier points win when slots are scarce.
"""
used_blocks = AnthropicCacheControlHook._count_request_cache_breakpoints(messages)
used_blocks = AnthropicCacheControlHook.count_request_cache_breakpoints(messages)
limit_reached = False
for point in points:
@ -350,7 +378,7 @@ class AnthropicCacheControlHook(CustomPromptManagement):
# 2. list of objects - only apply to last item per Anthropic spec
elif isinstance(message_content, list):
if len(message_content) > 0 and isinstance(message_content[-1], dict):
message_content[-1]["cache_control"] = control
message_content[-1]["cache_control"] = control # pyright: ignore[reportGeneralTypeIssues] # loose runtime dict
return message
@staticmethod
@ -428,8 +456,8 @@ class AnthropicCacheControlHook(CustomPromptManagement):
)
max_blocks: Final = MAX_CACHE_CONTROL_BLOCKS - reserved_blocks
message_blocks: Final = AnthropicCacheControlHook._count_request_cache_breakpoints(processed_messages)
system_blocks = AnthropicCacheControlHook._count_request_cache_breakpoints((), processed_system)
message_blocks: Final = AnthropicCacheControlHook.count_request_cache_breakpoints(processed_messages)
system_blocks = AnthropicCacheControlHook.count_request_cache_breakpoints((), processed_system)
if system_points and processed_system is not None and message_blocks + system_blocks < max_blocks:
system_already_has_cc: Final = isinstance(processed_system, list) and any(
@ -563,7 +591,7 @@ class AnthropicCacheControlHook(CustomPromptManagement):
carry the mark either at the top level (Anthropic shape) or nested under
``function`` (OpenAI shape); the Anthropic chat transform accepts both.
"""
if AnthropicCacheControlHook._count_request_cache_breakpoints(messages, system) > 0:
if AnthropicCacheControlHook.count_request_cache_breakpoints(messages, system) > 0:
return True
if tools is not None:
return any(
@ -723,6 +751,64 @@ class AnthropicCacheControlHook(CustomPromptManagement):
if points:
non_default_params["cache_control_injection_points"] = points
@staticmethod
def record_gateway_injection(
request_kwargs: Mapping[str, object],
added: int,
) -> None:
"""Name the deployment whose payload the gateway, not the client, put breakpoints on.
Spend accounting only asks whether litellm acted, so what it needs is which
deployment, not a count. Recording that is what makes the mark attempt-scoped: the
metadata bucket is one dict shared by every retry, failover and fallback of a
request, and ``litellm_call_id`` is shared with it, so anything request-scoped
written by one attempt is read by all of them and each boundary would have to
remember to strip it. The deployment is the part that actually changes when the
request moves, so a leg that injected nothing is never credited for one that did.
It also makes a zero delta (hook re-entry) and a negative one (a prompt manager
replacing the messages) harmless, since neither rewrites an earlier mark.
A pass that runs before a deployment is chosen, which is what the proxy does for
prompt templates, injects into the payload every leg goes on to send, so it marks
the request for all of them rather than for one.
Only what this pass actually placed counts. A ``tool_config`` point is placed by
the Bedrock converse transform, and only when the request carries tools, so the
presence of one here says nothing about whether a breakpoint reaches the wire;
claiming it marked three request shapes out of four that inject nothing. Missing
that Bedrock credit is the fail-closed direction, and the alternative is a
provider transform that carries spend-attribution state.
Reads whichever bucket the request actually carries rather than asking the shared
name resolver, which answers on key presence: ``litellm_params`` declares
``litellm_metadata`` as None on every request, so the resolver names a bucket that
is not there and the mark is dropped.
Never CREATES the bucket. The proxy seeds it on every request and is the marker's
only reader, so a request without one is a bare SDK call nothing would consume it
from. Creating it would also add a key to a dict call sites splat as ``**kwargs``,
and on the Responses API ``metadata`` is both this bucket's default name and an
explicit parameter, so the splat collides with the caller's own value.
"""
if added <= 0:
return
bucket: Final = next(
(
candidate
for candidate in (request_kwargs.get("litellm_metadata"), request_kwargs.get("metadata"))
if isinstance(candidate, dict)
),
None,
)
if bucket is not None:
model_info: Final = request_kwargs.get("model_info")
bucket[GATEWAY_INJECTED_CACHE_METADATA_KEY] = (
model_info.get("id", GATEWAY_INJECTED_FOR_EVERY_DEPLOYMENT)
if isinstance(model_info, dict)
else GATEWAY_INJECTED_FOR_EVERY_DEPLOYMENT
)
@staticmethod
def maybe_inject_cache_control(
messages: list[dict],
@ -772,17 +858,18 @@ class AnthropicCacheControlHook(CustomPromptManagement):
openai_dialect: Final = AnthropicCacheControlHook._targets_openai_prompt_cache_breakpoint(
model, custom_llm_provider, api_base, kwargs.get("prompt_cache_options")
)
breakpoints_before: Final = AnthropicCacheControlHook._count_request_cache_breakpoints(messages, system)
breakpoints_before: Final = AnthropicCacheControlHook.count_request_cache_breakpoints(messages, system)
messages, system, remaining = AnthropicCacheControlHook.apply_to_anthropic_messages_request(
messages=messages,
system=system,
injection_points=injection_points,
openai_dialect=openai_dialect,
)
if (
openai_dialect
and AnthropicCacheControlHook._count_request_cache_breakpoints(messages, system) > breakpoints_before
):
breakpoints_added: Final = (
AnthropicCacheControlHook.count_request_cache_breakpoints(messages, system) - breakpoints_before
)
AnthropicCacheControlHook.record_gateway_injection(kwargs, breakpoints_added)
if openai_dialect and breakpoints_added > 0:
kwargs.setdefault("prompt_cache_options", PromptCacheOptions(mode="explicit"))
if remaining:
kwargs["cache_control_injection_points"] = AnthropicCacheControlHook._stamped_as_judged(remaining)

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,8 +162,8 @@ class BitBucketClient:
response: Final = self.http_handler.get(url, headers=self.headers)
response.raise_for_status()
data: Final = response.json()
files: Final = []
data: Final[BitBucketSrcListing] = response.json()
files: Final[list[str]] = []
for item in data.get("values", []):
if item.get("type") == "commit_file":
@ -162,7 +189,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 +218,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 +231,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.).

View file

@ -220,6 +220,12 @@
"ui_name": "Host URL",
"description": "Langfuse host URL (default: https://cloud.langfuse.com)",
"required": false
},
"langfuse_environment": {
"type": "text",
"ui_name": "Tracing Environment",
"description": "Langfuse tracing environment (lowercase; falls back to LANGFUSE_TRACING_ENVIRONMENT)",
"required": false
}
},
"description": "Langfuse v2 Logging Integration"
@ -247,6 +253,12 @@
"ui_name": "Host URL",
"description": "Langfuse host URL (default: https://cloud.langfuse.com)",
"required": false
},
"langfuse_environment": {
"type": "text",
"ui_name": "Tracing Environment",
"description": "Langfuse tracing environment (lowercase; falls back to LANGFUSE_TRACING_ENVIRONMENT)",
"required": false
}
},
"description": "Langfuse v3 OTEL Logging Integration"
@ -294,12 +306,18 @@
"id": "newrelic",
"displayName": "New Relic",
"logo": "newrelic.png",
"supports_key_team_logging": false,
"supports_key_team_logging": true,
"dynamic_params": {
"NEW_RELIC_AI_MONITORING_RECORD_CONTENT_ENABLED": {
"newrelic_api_key": {
"type": "password",
"ui_name": "New Relic Ingest License Key",
"description": "Per-team ingest (license) key. Team traces export to this key's New Relic account over OTLP.",
"required": false
},
"newrelic_region": {
"type": "text",
"ui_name": "Record AI Content (default: true)",
"description": "Whether to record AI message content. Set to false to disable.",
"ui_name": "New Relic Region (us or eu)",
"description": "Data center region for this team's account. Defaults to us.",
"required": false
}
},

View file

@ -7,7 +7,7 @@ litellm_content_retrieve tool calls server-side via the typed agentic loop plan.
import time
import uuid
from typing import Any, Final, cast
from typing import TYPE_CHECKING, Any, ClassVar, Final, cast
from litellm._logging import verbose_logger
from litellm.compression import compress
@ -22,6 +22,9 @@ from litellm.types.integrations.custom_logger import (
)
from litellm.types.utils import CallTypes
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
LITELLM_CONTENT_RETRIEVE_TOOL_NAME: Final = "litellm_content_retrieve"
_CACHE_TTL_SECONDS: Final = 15 * 60
@ -72,6 +75,8 @@ class CompressionInterceptionLogger(CustomLogger):
4. Build typed rerun plan with tool_result blocks from the compressed cache.
"""
server_fulfilled_tool_names: ClassVar[frozenset[str]] = frozenset({LITELLM_CONTENT_RETRIEVE_TOOL_NAME})
def __init__(
self,
enabled: bool = True,
@ -220,7 +225,7 @@ class CompressionInterceptionLogger(CustomLogger):
response: Any,
anthropic_messages_provider_config: Any,
anthropic_messages_optional_request_params: dict,
logging_obj: Any,
logging_obj: "LiteLLMLoggingObj | None",
stream: bool,
kwargs: dict,
) -> AgenticLoopPlan:

View file

@ -45,7 +45,7 @@ class CustomBatchLogger(CustomLogger):
super().__init__(**kwargs)
async def periodic_flush(self):
async def periodic_flush(self) -> None:
while True:
await asyncio.sleep(self.flush_interval)
verbose_logger.debug("CustomLogger periodic flush after %s seconds", self.flush_interval)

View file

@ -60,6 +60,12 @@ _PRE_CALL_EXECUTED_TOKEN: Final = secrets.token_hex(16)
_GUARDRAIL_BLOCK_STATUS_CODES: Final = frozenset({400, 403, 422})
DEFAULT_ADVISORY_MESSAGE: Final = (
"The user's latest message was flagged for {reason} by a content safety "
"guardrail. This may be a false positive. Use your judgment: respond "
"helpfully if the request is legitimate, or decline if it is not."
)
_guardrail_self_recorded: Final[contextvars.ContextVar[bool]] = contextvars.ContextVar(
"litellm_guardrail_self_recorded", default=False
)
@ -158,6 +164,7 @@ class CustomGuardrail(CustomLogger):
sensitive_data_route_to_model: str | None = None,
sticky_session_routing: bool = True,
run_in_parallel: bool = False,
scan_raw_request: bool = False,
only_scan_new_messages: bool = False,
**kwargs,
):
@ -180,6 +187,13 @@ class CustomGuardrail(CustomLogger):
run_in_parallel: When True, this pre_call or post_call guardrail runs concurrently with
other opted-in guardrails of the same hook. Only safe for block-only guardrails that
do not mutate the request or response.
scan_raw_request: When True, this pre_call guardrail always evaluates the request as it
was before any guardrail in this hook ran, regardless of where it's declared in the
guardrails list -- so an earlier guardrail that masks/rewrites content (e.g. PII
redaction) can never hide a violation from this one. Only safe for block-only
guardrails: any data this guardrail returns is discarded, matching run_in_parallel's
contract, since applying its mutations on top of a stale snapshot would silently
undo whatever later guardrails already did to the live request.
"""
self.guardrail_name = guardrail_name
self.supported_event_hooks = supported_event_hooks
@ -195,6 +209,7 @@ class CustomGuardrail(CustomLogger):
self.sensitive_data_route_to_model: str | None = sensitive_data_route_to_model
self.sticky_session_routing: bool = sticky_session_routing
self.run_in_parallel: bool = run_in_parallel
self.scan_raw_request: bool = scan_raw_request
self.only_scan_new_messages: bool = only_scan_new_messages
if supported_event_hooks:
@ -281,6 +296,82 @@ class CustomGuardrail(CustomLogger):
original_response=original_response,
)
def inject_advisory_message(
self,
data: dict[str, Any], # mutable-ok: caller's dict is mutated in place, matching mark_pre_call_hook_ran
message: str,
) -> bool:
"""
Append an advisory system message to the request in place, so the LLM
itself can weigh a possible false-positive guardrail flag rather than
the request being hard-blocked or silently allowed.
Unlike raise_passthrough_exception, this does NOT short-circuit the LLM
call; the request proceeds normally with the extra message appended.
Guardrails should call this from on_flagged handling analogous to how
passthrough-supporting guardrails call raise_passthrough_exception.
Args:
data: The request data dictionary, mutated in place to append the
advisory message to its "messages" list and/or "input"/
"instructions" text.
message: The formatted advisory message to append as a system message.
Returns:
True if the advisory was actually written somewhere the model will
see it. False if ``data["input"]`` is a structured Responses-API
list (not a plain string) -- the Responses API reads only
``input``, so appending to ``messages`` would be inert regardless
of whether a ``messages`` list also happens to be present, and
there is no field this helper can safely append into. The caller
must treat this like any other case where the mitigation can't
land and degrade to blocking instead of silently letting the
flagged request through unmodified.
"""
advisory_message: Final = {"role": "system", "content": message} # mutable-ok: plain dict for live request
existing_messages: Final = data.get("messages")
existing_input: Final = data.get("input")
existing_instructions: Final = data.get("instructions")
if isinstance(existing_instructions, str):
# Responses API "instructions" is the privileged, developer-set
# system-level field the model treats as authoritative -- unlike
# "input", which the caller controls and could use to tell the
# model to disregard a trailing warning. Prefer it over "input"
# whenever present.
if isinstance(existing_messages, list):
messages_with_instructions_note: Final = [ # mutable-ok: fresh list
*existing_messages,
advisory_message,
]
data["messages"] = messages_with_instructions_note # rebind-ok: mutates caller's dict by design
data["instructions"] = f"{existing_instructions}\n\n{message}" # rebind-ok: mutates caller's dict by design
return True
if isinstance(existing_input, str):
# A plain-string "input" doesn't rule out "messages" also being a
# real, read field (e.g. a chat-completions call carrying a stray
# "input"), so write to both when both are present.
if isinstance(existing_messages, list):
messages_with_input_note: Final = [*existing_messages, advisory_message] # mutable-ok: fresh list
data["messages"] = messages_with_input_note # rebind-ok: mutates caller's dict by design
# The Responses API reads "input", not "messages" -- appending only to
# "messages" would leave the advisory unreachable for that endpoint.
data["input"] = f"{existing_input}\n\n{message}" # rebind-ok: mutates caller's dict by design
return True
if existing_input is not None:
# existing_input is a structured (non-string) Responses-API item
# list. That endpoint reads only "input", so appending to
# "messages" -- even if "messages" also happens to be present --
# would never reach the model. Leave data untouched and report
# non-delivery so the caller degrades to blocking.
return False
if isinstance(existing_messages, list):
messages_without_input_note: Final = [*existing_messages, advisory_message] # mutable-ok: fresh list
data["messages"] = messages_without_input_note # rebind-ok: mutates caller's dict by design
return True
sole_message: Final = [advisory_message] # mutable-ok: plain list for the live JSON request
data["messages"] = sole_message # rebind-ok: mutates caller's dict by design
return True
def raise_sensitive_data_route_exception(
self,
route_to_model: str,

View file

@ -2,8 +2,8 @@
# On success, logs events to Promptlayer
import re
import traceback
from collections.abc import AsyncGenerator
from typing import TYPE_CHECKING, Any, Final, Optional
from collections.abc import AsyncGenerator, Mapping
from typing import TYPE_CHECKING, Any, ClassVar, Final, Optional
from pydantic import BaseModel
@ -60,6 +60,7 @@ _BASE64_INLINE_PATTERN: Final = re.compile(
class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callback#callback-class
# Class variables or attributes
server_fulfilled_tool_names: ClassVar[frozenset[str]] = frozenset()
enforces_request_content: bool = False
"""
@ -292,6 +293,54 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
Allow modifying / reviewing the response just after it's received from the deployment.
"""
async def async_post_call_failure_deployment_hook(
self,
request_data: Mapping[str, object],
exception: Exception,
call_type: CallTypes | None,
fallback_depth: int | None = None,
) -> None:
"""
Called once per failed deployment attempt - attempt 1, every retry, and
every fallback chain step - because the router re-invokes the wrapped
function on each attempt, re-entering this hook's call site fresh
every time.
This is a DEPLOYMENT-LEVEL signal, distinct from the REQUEST-LEVEL
``async_log_failure_event``, which fires once per logical client
request behind a dedup gate. ``request_data`` is mostly this
attempt's own kwargs, with one exception: it omits
``attempted_targets``, the router's own bookkeeping of which fallback
targets this request has already tried, since that one object *is*
shared by reference across every hop of the live fallback walk.
Pairs with ``async_pre_call_deployment_hook`` and
``async_post_call_success_deployment_hook`` to complete the
pre-call/success/failure lifecycle for a single deployment attempt.
``fallback_depth`` is best-effort: ``None`` on the first attempt and on
any call made without a ``Router`` (a bare SDK call has no fallback
chain to be at a depth in), ``1`` on the first fallback hop, ``2`` on
the second, and so on. It reflects ``Router``'s own internal fallback
bookkeeping (``kwargs["fallback_depth"]``), not a value this hook
computes or guarantees the shape of across versions. It tracks
fallback hops only, not retries within the same model group - a
retry-only failure (no fallback yet) also reports ``None``. If an
override predates this field it's simply never passed, rather than
raising - safe to leave off an override written before it existed.
``exception`` is a same-class snapshot, not the exact object about to
be re-raised to the real caller: read it freely, but setting an
attribute on it (e.g. ``status_code``) has no effect on what the
caller actually receives.
Default: no-op. Opt in by overriding. Keep overrides fast - this
runs on the request's exception path, so a slow implementation
delays error propagation to the caller. The reported failure
duration is captured before this hook runs, so a slow override
doesn't inflate that metric, but the caller still waits for it.
"""
async def async_post_call_streaming_deployment_hook(
self,
request_data: dict,

View file

@ -62,12 +62,16 @@ def prompt_initializer(litellm_params: "PromptLiteLLMParams", prompt_spec: "Prom
if dotprompt_content and not prompt_data and not prompt_file:
prompt_data = _get_prompt_data_from_dotprompt_content(dotprompt_content)
from .prompt_manager import strip_version_suffix
registration_prompt_id: Final = prompt_id or strip_version_suffix(prompt_spec.prompt_id) or prompt_spec.prompt_id
try:
dot_prompt_manager: Final = DotpromptManager(
prompt_directory=prompt_directory,
prompt_data=prompt_data,
prompt_file=prompt_file,
prompt_id=prompt_id,
prompt_id=registration_prompt_id,
)
return dot_prompt_manager

View file

@ -96,7 +96,7 @@ class DotpromptManager(CustomPromptManagement):
if prompt_id is None:
return False
try:
return prompt_id in self.prompt_manager.list_prompts()
return self.prompt_manager.get_prompt(prompt_id) is not None
except Exception:
# If there's any error accessing prompts, don't run prompt management
return False
@ -209,6 +209,8 @@ class DotpromptManager(CustomPromptManagement):
prompt_spec=prompt_spec,
prompt_label=prompt_label,
prompt_version=prompt_version,
ignore_prompt_manager_model=ignore_prompt_manager_model,
ignore_prompt_manager_optional_params=ignore_prompt_manager_optional_params,
)
async def async_get_chat_completion_prompt(

View file

@ -11,6 +11,13 @@ from jinja2 import DictLoader, select_autoescape
from jinja2.sandbox import ImmutableSandboxedEnvironment
def strip_version_suffix(prompt_id: str) -> str | None:
base, separator, version = prompt_id.rpartition(".v")
if separator and base and version.isdigit():
return base
return None
class PromptTemplate:
"""Represents a single prompt template with metadata and content."""
@ -124,11 +131,13 @@ class PromptManager:
"content": "template content",
"metadata": {"model": "gpt-4", "temperature": 0.7, ...}
} + prompt_id
"""
if prompt_id:
prompt_data = {prompt_id: prompt_data}
for prompt_id, prompt_info in prompt_data.items():
A dict carrying a "content" key is a single flat template registered under
prompt_id; anything else is treated as already keyed by template ID.
"""
keyed_prompts: Final = {prompt_id: prompt_data} if prompt_id and "content" in prompt_data else prompt_data
for template_id, prompt_info in keyed_prompts.items():
try:
content = prompt_info.get("content", "")
metadata = prompt_info.get("metadata", {})
@ -136,11 +145,10 @@ class PromptManager:
template = PromptTemplate(
content=content,
metadata=metadata,
template_id=prompt_id,
template_id=template_id,
)
self.prompts[prompt_id] = template
self.prompts[template_id] = template
except Exception:
# Optional: print(f"Error loading prompt from JSON: {prompt_id}")
pass
def _load_prompt_file(self, file_path: str | Path, prompt_id: str) -> PromptTemplate:
@ -272,8 +280,12 @@ class PromptManager:
if versioned_id in self.prompts:
return self.prompts[versioned_id]
# Fall back to base prompt_id
return self.prompts.get(prompt_id)
direct_match: Final = self.prompts.get(prompt_id)
if direct_match is not None:
return direct_match
base_prompt_id: Final = strip_version_suffix(prompt_id)
return self.prompts.get(base_prompt_id) if base_prompt_id else None
def list_prompts(self) -> list[str]:
"""Get a list of all available prompt IDs."""

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

View file

@ -416,17 +416,8 @@ class GenericPromptManager(CustomPromptManagement):
tools=tools,
prompt_label=prompt_label,
prompt_version=prompt_version,
ignore_prompt_manager_model=(
ignore_prompt_manager_model or prompt_spec.litellm_params.ignore_prompt_manager_model
if prompt_spec
else False
),
ignore_prompt_manager_optional_params=(
ignore_prompt_manager_optional_params
or prompt_spec.litellm_params.ignore_prompt_manager_optional_params
if prompt_spec
else False
),
ignore_prompt_manager_model=ignore_prompt_manager_model,
ignore_prompt_manager_optional_params=ignore_prompt_manager_optional_params,
)
def get_chat_completion_prompt(
@ -457,17 +448,8 @@ class GenericPromptManager(CustomPromptManagement):
prompt_spec=prompt_spec,
prompt_label=prompt_label,
prompt_version=prompt_version,
ignore_prompt_manager_model=(
ignore_prompt_manager_model or prompt_spec.litellm_params.ignore_prompt_manager_model
if prompt_spec
else False
),
ignore_prompt_manager_optional_params=(
ignore_prompt_manager_optional_params
or prompt_spec.litellm_params.ignore_prompt_manager_optional_params
if prompt_spec
else False
),
ignore_prompt_manager_model=ignore_prompt_manager_model,
ignore_prompt_manager_optional_params=ignore_prompt_manager_optional_params,
)
def clear_cache(self) -> None:

View file

@ -1,9 +1,11 @@
#### What this does ####
# On success, logs events to Langfuse
import inspect
import os
import traceback
from collections.abc import Callable, Iterable, Mapping
from datetime import datetime
from functools import lru_cache
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, cast
@ -21,6 +23,9 @@ from litellm.litellm_core_utils.core_helpers import (
reconstruct_model_name,
safe_deep_copy,
)
from litellm.litellm_core_utils.initialize_dynamic_callback_params import (
validate_langfuse_environment_value,
)
from litellm.litellm_core_utils.redact_messages import redact_user_api_key_info
from litellm.llms.custom_httpx.http_handler import _get_httpx_client
from litellm.secret_managers.main import str_to_bool
@ -133,6 +138,16 @@ def resolve_langfuse_credentials(
return public_key, secret_key, resolved_host
@lru_cache(maxsize=8)
def _warn_invalid_deployment_environment(raw_value: str, error: str) -> None:
verbose_logger.warning(
"Ignoring invalid LANGFUSE_TRACING_ENVIRONMENT=%r for the langfuse callback: %s. "
"Traces will be sent to Langfuse's default environment.",
raw_value,
error,
)
class LangFuseLogger:
# Class variables or attributes
def __init__(
@ -140,6 +155,7 @@ class LangFuseLogger:
langfuse_public_key=None,
langfuse_secret=None,
langfuse_host=None,
langfuse_environment: str | None = None,
flush_interval=1,
allow_env_credentials: bool = True,
):
@ -159,6 +175,12 @@ class LangFuseLogger:
if not (self.langfuse_host.startswith("http://") or self.langfuse_host.startswith("https://")):
# add http:// if unset, assume communicating over private network - e.g. render
self.langfuse_host = "http://" + self.langfuse_host
_env_override: Final = str(langfuse_environment).strip() if langfuse_environment is not None else None
if _env_override:
validate_langfuse_environment_value(_env_override)
self.langfuse_environment: str | None = _env_override
else:
self.langfuse_environment = self.resolve_deployment_environment()
self.langfuse_release = os.getenv("LANGFUSE_RELEASE")
self.langfuse_debug = os.getenv("LANGFUSE_DEBUG")
self.langfuse_flush_interval = LangFuseLogger._get_langfuse_flush_interval(flush_interval)
@ -182,6 +204,8 @@ class LangFuseLogger:
}
self.langfuse_sdk_version: str = langfuse.version.__version__
if "environment" in inspect.signature(Langfuse.__init__).parameters:
parameters["environment"] = self.langfuse_environment
if Version(self.langfuse_sdk_version) >= Version("2.6.0"):
parameters["sdk_integration"] = "litellm"
self.Langfuse: Langfuse = self.safe_init_langfuse_client(parameters)
@ -942,6 +966,20 @@ class LangFuseLogger:
verbose_logger.warning("Failed to apply masking function: %s. Returning original data.", e)
return data
@staticmethod
def resolve_deployment_environment() -> str | None:
"""Resolve LANGFUSE_TRACING_ENVIRONMENT: stripped value, "default" plus a warning when invalid, None when unset."""
raw: Final = os.getenv("LANGFUSE_TRACING_ENVIRONMENT")
if not raw:
return None
value: Final = raw.strip()
try:
validate_langfuse_environment_value(value)
except ValueError as e:
_warn_invalid_deployment_environment(raw, str(e))
return "default"
return value
@staticmethod
def _get_langfuse_flush_interval(flush_interval: int) -> int:
"""

View file

@ -6,6 +6,7 @@ Used to get the LangFuseLogger for a given request
Handles Key/Team Based Langfuse Logging
"""
import os
from typing import TYPE_CHECKING, Any, Final
from litellm.litellm_core_utils.litellm_logging import StandardCallbackDynamicParams
@ -108,6 +109,7 @@ class LangFuseHandler:
langfuse_public_key=credentials.get("langfuse_public_key"),
langfuse_secret=credentials.get("langfuse_secret") or credentials.get("langfuse_secret_key"),
langfuse_host=credentials.get("langfuse_host"),
langfuse_environment=credentials.get("langfuse_environment"),
allow_env_credentials=credentials.get("langfuse_host") is None,
)
in_memory_dynamic_logger_cache.set_cache(
@ -135,8 +137,33 @@ class LangFuseHandler:
or standard_callback_dynamic_params.get("langfuse_secret_key"),
langfuse_public_key=standard_callback_dynamic_params.get("langfuse_public_key"),
langfuse_host=standard_callback_dynamic_params.get("langfuse_host"),
langfuse_environment=LangFuseHandler._meaningful_dynamic_environment(standard_callback_dynamic_params),
)
@staticmethod
def _meaningful_dynamic_environment(
standard_callback_dynamic_params: StandardCallbackDynamicParams,
) -> str | None:
"""Return the per-request environment only when it changes behavior.
Empty/whitespace values and values equal to the deployment-wide
LANGFUSE_TRACING_ENVIRONMENT fallback are treated as absent so an
environment-only override that matches the default does not mint a
duplicate SDK client (each client costs threads and counts against
MAX_LANGFUSE_INITIALIZED_CLIENTS).
"""
raw = standard_callback_dynamic_params.get("langfuse_environment")
if raw is None:
return None
value = str(raw).strip()
if (
not value
or value == os.getenv("LANGFUSE_TRACING_ENVIRONMENT")
or value == LangFuseLogger.resolve_deployment_environment()
):
return None
return value
@staticmethod
def _dynamic_langfuse_credentials_are_passed(
standard_callback_dynamic_params: StandardCallbackDynamicParams,
@ -153,6 +180,7 @@ class LangFuseHandler:
or standard_callback_dynamic_params.get("langfuse_public_key") is not None
or standard_callback_dynamic_params.get("langfuse_secret") is not None
or standard_callback_dynamic_params.get("langfuse_secret_key") is not None
or LangFuseHandler._meaningful_dynamic_environment(standard_callback_dynamic_params) is not None
):
return True
return False

View file

@ -10,6 +10,7 @@ from litellm.integrations.langfuse.langfuse_otel_attributes import (
LangfuseLLMObsOTELAttributes,
)
from litellm.integrations.opentelemetry import OpenTelemetry, OpenTelemetryConfig
from litellm.litellm_core_utils.safe_json_loads import safe_json_loads
from litellm.types.integrations.langfuse_otel import (
LangfuseSpanAttributes,
)
@ -197,7 +198,11 @@ class LangfuseOtelLogger(OpenTelemetry):
)
elif item_type == "function_call":
arguments_str = getattr(item, "arguments", "{}")
arguments_obj = json.loads(arguments_str) if isinstance(arguments_str, str) else arguments_str
arguments_obj = (
safe_json_loads(arguments_str, default={})
if isinstance(arguments_str, str)
else arguments_str
)
langfuse_tool_call = {
"id": getattr(item, "id", ""),
"name": getattr(item, "name", ""),
@ -226,7 +231,10 @@ class LangfuseOtelLogger(OpenTelemetry):
from litellm.integrations.arize._utils import safe_set_attribute
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
langfuse_environment: Final = os.environ.get("LANGFUSE_TRACING_ENVIRONMENT")
dynamic_params: Final = kwargs.get("standard_callback_dynamic_params")
langfuse_environment: Final = (
dynamic_params.get("langfuse_environment") if dynamic_params else None
) or os.environ.get("LANGFUSE_TRACING_ENVIRONMENT")
if langfuse_environment:
safe_set_attribute(
span,

View file

@ -2,6 +2,7 @@
Call Hook for LiteLLM Proxy which allows Langfuse prompt management.
"""
import inspect
import os
from functools import lru_cache
from typing import TYPE_CHECKING, Any, Final, Literal, TypeAlias, cast
@ -109,6 +110,9 @@ def langfuse_client_init(
cert=os.getenv("SSL_CERTIFICATE", litellm.ssl_certificate),
)
if "environment" in inspect.signature(Langfuse.__init__).parameters:
parameters["environment"] = LangFuseLogger.resolve_deployment_environment()
client: Final = Langfuse(**parameters)
return client

View file

@ -168,17 +168,20 @@ class LangsmithLogger(CustomBatchLogger):
return outputs
def _ensure_required_ids(self, data: dict, run_id: str | None):
resolved_id: Final = run_id or str(uuid.uuid4())
if "id" not in data or data["id"] is None:
run_id = str(uuid.uuid4())
data["id"] = run_id
data["id"] = resolved_id
if "trace_id" not in data or data["trace_id"] is None:
if run_id is not None and isinstance(run_id, str):
data["trace_id"] = run_id
# LangSmith rejects the whole ingest batch unless a root run's trace_id
# equals the run id embedded in the first segment of dotted_order
posts_as_root: Final = ("parent_run_id" not in data or data["parent_run_id"] is None) and (
"dotted_order" not in data or data["dotted_order"] is None
)
if posts_as_root or "trace_id" not in data or data["trace_id"] is None:
data["trace_id"] = resolved_id
if "dotted_order" not in data or data["dotted_order"] is None:
if run_id is not None and isinstance(run_id, str):
data["dotted_order"] = self.make_dot_order(run_id=run_id)
data["dotted_order"] = self.make_dot_order(run_id=resolved_id)
def _prepare_log_data(
self,
@ -193,6 +196,11 @@ class LangsmithLogger(CustomBatchLogger):
metadata = _litellm_params.get("metadata", {}) or {}
fields: Final = self._extract_metadata_fields(metadata, credentials)
# the proxy header fan-out mirrors one value into both keys, and LangSmith
# rejects the whole ingest batch when run-body session_id is not an
# existing tracer-session uuid
if fields["session_id"] == fields["trace_id"]:
fields["session_id"] = None
verbose_logger.debug(
"Langsmith Logging - project_name: %s, run_name %s", fields["project_name"], fields["run_name"]
)

View file

@ -0,0 +1,395 @@
"""
New Relic Metric API Integration - sends per-team cost/usage metrics to /metric/v1
NR Reference API: https://docs.newrelic.com/docs/data-apis/ingest-apis/metric-api/introduction-metric-api/
`async_log_success_event` / `async_log_failure_event` queue one record per request;
at flush the queue is aggregated by (team, model group, model, provider, status)
into count/summary metrics. `interval.ms` is the real window between flushes,
computed at flush time.
Team-scoped by construction: the ingest key is injected explicitly and there is
deliberately no environment-variable fallback, so a team's metrics are never sent
with the proxy operator's credentials (mirrors ``allow_env_credentials=False`` on
the Datadog team logger).
Error policy on flush: 4xx drops the batch (a retry would fail identically; 403
is a permanent credential failure), 5xx/network re-queues capped at
``max_queue_size`` records with the oldest dropped.
For batching specific details see CustomBatchLogger class
"""
import asyncio
import gzip
import time
import traceback
from collections.abc import Mapping
from math import ceil
from types import MappingProxyType
from typing import Final
from httpx import HTTPStatusError, Response
from litellm._logging import verbose_logger
from litellm.integrations.custom_batch_logger import CustomBatchLogger
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
from litellm.llms.custom_httpx.http_handler import (
get_async_httpx_client,
httpxSpecialProvider,
)
from litellm.types.integrations.newrelic import (
NEWRELIC_DEFAULT_REGION,
NEWRELIC_METRIC_ATTRIBUTE_MAX_LEN,
NEWRELIC_METRIC_COMPLETION_TOKENS,
NEWRELIC_METRIC_COST_USD,
NEWRELIC_METRIC_ENDPOINT_BY_REGION,
NEWRELIC_METRIC_PROMPT_TOKENS,
NEWRELIC_METRIC_REQUEST_DURATION_MS,
NEWRELIC_METRIC_REQUESTS,
NEWRELIC_METRIC_TOTAL_TOKENS,
NEWRELIC_METRICS_MAX_BATCH_SIZE,
NEWRELIC_METRICS_MAX_DRAIN_PASSES,
NEWRELIC_METRICS_MAX_RETRY_QUEUE_SIZE,
NewRelicCountMetric,
NewRelicMetric,
NewRelicMetricCommon,
NewRelicMetricEnvelope,
NewRelicMetricRecord,
NewRelicSummaryMetric,
NewRelicSummaryValue,
)
from litellm.types.utils import StandardLoggingPayload
# 408 (request timeout) and 429 (rate limit) are transient client errors the
# Metric API expects a retry on, unlike 400/403 which a retry would only repeat.
_RETRYABLE_CLIENT_STATUSES: Final = frozenset({408, 429})
def resolve_newrelic_metric_endpoint(newrelic_region: str | None) -> str:
if not newrelic_region:
return NEWRELIC_METRIC_ENDPOINT_BY_REGION[NEWRELIC_DEFAULT_REGION]
endpoint: Final = NEWRELIC_METRIC_ENDPOINT_BY_REGION.get(newrelic_region.lower())
if endpoint is None:
verbose_logger.warning(
"New Relic: unknown newrelic_region %r; supported regions: %s. Using the default (US) endpoint.",
newrelic_region,
", ".join(sorted(NEWRELIC_METRIC_ENDPOINT_BY_REGION)),
)
return NEWRELIC_METRIC_ENDPOINT_BY_REGION[NEWRELIC_DEFAULT_REGION]
return endpoint
def _metric_record_from_payload(standard_logging_object: StandardLoggingPayload) -> NewRelicMetricRecord:
metadata: Final = standard_logging_object.get("metadata")
team_id: Final = ((metadata.get("user_api_key_team_id") or metadata.get("team_id")) if metadata else None) or ""
team_alias: Final = (
(metadata.get("user_api_key_team_alias") or metadata.get("team_alias")) if metadata else None
) or ""
return NewRelicMetricRecord(
team_id=team_id,
team_alias=team_alias,
model_group=standard_logging_object.get("model_group") or "",
model=standard_logging_object.get("model") or "",
custom_llm_provider=standard_logging_object.get("custom_llm_provider") or "",
status=str(standard_logging_object.get("status") or "success"),
response_cost=float(standard_logging_object.get("response_cost") or 0.0),
prompt_tokens=int(standard_logging_object.get("prompt_tokens") or 0),
completion_tokens=int(standard_logging_object.get("completion_tokens") or 0),
total_tokens=int(standard_logging_object.get("total_tokens") or 0),
duration_ms=float(standard_logging_object.get("response_time") or 0.0) * 1000.0,
)
def _bucket_metrics(bucket_records: tuple[NewRelicMetricRecord, ...]) -> tuple[NewRelicMetric, ...]:
first: Final = bucket_records[0]
attributes: Final[Mapping[str, str]] = { # mutable-ok: JSON leaf; safe_dumps stringifies MappingProxyType
key: value[:NEWRELIC_METRIC_ATTRIBUTE_MAX_LEN]
for key, value in (
("team_id", first.team_id),
("team_alias", first.team_alias),
("model_group", first.model_group),
("model", first.model),
("custom_llm_provider", first.custom_llm_provider),
("status", first.status),
)
if value
}
durations: Final = tuple(record.duration_ms for record in bucket_records)
counts: Final[tuple[tuple[str, float], ...]] = (
(NEWRELIC_METRIC_REQUESTS, float(len(bucket_records))),
(NEWRELIC_METRIC_COST_USD, sum(record.response_cost for record in bucket_records)),
(NEWRELIC_METRIC_PROMPT_TOKENS, float(sum(record.prompt_tokens for record in bucket_records))),
(NEWRELIC_METRIC_COMPLETION_TOKENS, float(sum(record.completion_tokens for record in bucket_records))),
(NEWRELIC_METRIC_TOTAL_TOKENS, float(sum(record.total_tokens for record in bucket_records))),
)
count_metrics: Final[tuple[NewRelicMetric, ...]] = tuple(
NewRelicCountMetric(name=name, type="count", value=value, attributes=attributes) for name, value in counts
)
summary_metric: Final = NewRelicSummaryMetric(
name=NEWRELIC_METRIC_REQUEST_DURATION_MS,
type="summary",
value=NewRelicSummaryValue(
count=len(durations),
sum=sum(durations),
min=min(durations),
max=max(durations),
),
attributes=attributes,
)
return (*count_metrics, summary_metric)
def build_metric_payload(
records: tuple[NewRelicMetricRecord, ...],
*,
window_start: float,
now: float,
) -> tuple[NewRelicMetricEnvelope, ...]:
"""Aggregates records into one Metric API envelope for the flush window."""
interval_ms: Final = max(1, int((now - window_start) * 1000))
bucket_keys: Final = tuple(dict.fromkeys(record.bucket_key for record in records))
metrics: Final = tuple(
metric
for key in bucket_keys
for metric in _bucket_metrics(tuple(record for record in records if record.bucket_key == key))
)
common: Final[NewRelicMetricCommon] = {
"timestamp": int(window_start * 1000),
"interval.ms": interval_ms,
}
return (NewRelicMetricEnvelope(common=common, metrics=metrics),)
class NewRelicMetricsLogger(CustomBatchLogger):
def __init__(
self,
newrelic_api_key: str,
newrelic_region: str | None = None,
) -> None:
if not newrelic_api_key:
raise ValueError(
"newrelic_api_key is required for NewRelicMetricsLogger; "
"team-scoped metrics never fall back to environment credentials"
)
self.newrelic_api_key: Final = newrelic_api_key
self.metric_api_url: Final = resolve_newrelic_metric_endpoint(newrelic_region)
self.async_client = get_async_httpx_client(llm_provider=httpxSpecialProvider.LoggingCallback)
self._stopped: bool = False
self._drain_lock = asyncio.Lock()
asyncio.create_task(self.periodic_flush())
self.flush_lock = asyncio.Lock()
super().__init__(
flush_lock=self.flush_lock,
batch_size=NEWRELIC_METRICS_MAX_BATCH_SIZE,
max_queue_size=NEWRELIC_METRICS_MAX_RETRY_QUEUE_SIZE,
)
def stop(self) -> None:
"""Ends the periodic flush loop; called on DynamicLoggingCache eviction.
Schedules one final drain of anything still queued, so eviction never
silently discards records. Guarded so it can never raise into the
cache's eviction path.
"""
self._stopped = True
try:
asyncio.get_running_loop().create_task(self._final_drain())
except Exception: # noqa: BLE001 # no running loop / shutdown; the periodic loop's final drain still runs
verbose_logger.debug("New Relic Metrics: could not schedule final drain on stop()", exc_info=True)
async def _drain_with_retry(self) -> None:
"""Deliver everything queued on a stopped logger, or drop it with a log.
A stopped logger has no periodic loop left, so every post-stop path
funnels through here. ``_drain_lock`` serializes drains: a callback that
appends and starts its own drain queues behind the running one instead
of racing it. Each pass attempts the whole current queue in
``batch_size`` chunks, unlike the periodic path it does not stop at the
first failing chunk, so a persistently failing head never starves the
tail. Only after ``_MAX_DRAIN_PASSES`` against a permanently failing
destination is the remainder dropped, and then only the records that were
queued when this drain began, so every dropped record got the full retry
budget: a record a callback appended mid-drain is not in that snapshot,
so it is left for its own serialized drain rather than dropped after
fewer attempts, and is never stranded.
"""
async with self._drain_lock:
attempted: Final = tuple(self.log_queue)
for _pass in range(NEWRELIC_METRICS_MAX_DRAIN_PASSES):
await self._drain_flush_once()
if not self.log_queue:
return
if _pass < NEWRELIC_METRICS_MAX_DRAIN_PASSES - 1:
await asyncio.sleep(2**_pass)
async with self.flush_lock:
tried_ids: Final = frozenset(id(record) for record in attempted)
survivors: Final = tuple(record for record in self.log_queue if id(record) not in tried_ids)
dropped: Final = len(self.log_queue) - len(survivors)
if dropped:
verbose_logger.warning(
"New Relic Metrics: dropping %s records after %s drain passes",
dropped,
NEWRELIC_METRICS_MAX_DRAIN_PASSES,
)
self.log_queue[:] = list(survivors) # mutable-ok: leave late arrivals for the next serialized drain
async def _drain_flush_once(self) -> None:
"""Attempt every queued record once, in ``batch_size`` chunks, without
stopping at the first failing chunk so a persistently failing head does
not starve the tail (the periodic ``flush_queue`` deliberately stops
instead). Takes the queue under ``flush_lock`` and re-queues only the
chunks a 5xx/network error left undelivered, so records a concurrent
request appends during the sends survive for the next pass."""
async with self.flush_lock:
pending: Final = tuple(self.log_queue)
window_start: Final = self.last_flush_time
self.last_flush_time = time.time()
del self.log_queue[:]
if not pending:
return
chunks: Final = tuple(
pending[start : start + self.batch_size] for start in range(0, len(pending), self.batch_size)
)
delivered: Final = tuple([await self._classify_and_send(chunk, window_start) for chunk in chunks])
failed: Final = tuple(record for chunk, ok in zip(chunks, delivered) for record in (() if ok else chunk))
if failed:
self._requeue(failed)
async def _final_drain(self) -> None:
await self._drain_with_retry()
async def periodic_flush(self) -> None:
while not self._stopped:
await asyncio.sleep(self.flush_interval)
if self._stopped:
break
await self.flush_queue()
await self._final_drain()
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time) -> None:
try:
await self._log_async_event(standard_logging_object=kwargs.get("standard_logging_object", None))
except Exception as e: # noqa: BLE001 # logging must never break the request path
verbose_logger.exception("New Relic Metrics Layer Error - %s\n%s", e, traceback.format_exc())
async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time) -> None:
try:
await self._log_async_event(standard_logging_object=kwargs.get("standard_logging_object", None))
except Exception as e: # noqa: BLE001 # logging must never break the request path
verbose_logger.exception("New Relic Metrics Layer Error - %s\n%s", e, traceback.format_exc())
async def _log_async_event(self, standard_logging_object: StandardLoggingPayload | None) -> None:
if standard_logging_object is None:
raise ValueError("standard_logging_object not found in kwargs")
self.log_queue.append(_metric_record_from_payload(standard_logging_object))
if self._stopped:
# A stopped logger has no periodic loop left; an in-flight callback
# that appends after the eviction drain delivers its own record.
await self._drain_with_retry()
return
if len(self.log_queue) >= self.batch_size:
await self.flush_queue()
async def flush_queue(self) -> None:
async with self.flush_lock:
window_start: Final = self.last_flush_time
self.last_flush_time = time.time()
queued: Final = len(self.log_queue)
if not queued:
return
verbose_logger.debug("New Relic Metrics: Flushing %s queued records", queued)
# Bounded by what is queued now: records appended mid-flush belong to
# the next window, and looping until empty would never end under load.
for _chunk in range(ceil(queued / self.batch_size)):
if not await self.async_send_batch(window_start=window_start):
return
async def async_send_batch(self, window_start: float | None = None) -> bool:
"""Sends the oldest ``batch_size`` records only, so a queue grown past that
by re-queues cannot breach the Metric API data point cap in one request.
Returns False once a chunk fails and is re-queued, so the caller stops."""
if not self.log_queue:
return False
batch_to_send: Final[tuple[NewRelicMetricRecord, ...]] = tuple(self.log_queue[: self.batch_size])
del self.log_queue[: len(batch_to_send)]
delivered: Final = await self._classify_and_send(
batch_to_send, window_start if window_start is not None else self.last_flush_time
)
if not delivered:
self._requeue(batch_to_send)
return delivered
async def _classify_and_send(self, batch: tuple[NewRelicMetricRecord, ...], window_start: float) -> bool:
"""Send one chunk and classify the outcome, never touching the queue.
Returns True when the batch is done with (delivered on any 2xx, or a 4xx
a retry would only repeat, 403 being a permanent bad-key rejection), and
False when a 5xx or network error means the caller should re-queue it.
``AsyncHTTPHandler.post`` raises ``HTTPStatusError`` on any non-2xx, so a
4xx never returns a response here; the status is read off the raised
error to keep the client-error path (drop) distinct from 5xx (retry)."""
payload: Final = build_metric_payload(records=batch, window_start=window_start, now=time.time())
try:
status = (
await self.async_send_compressed_data(payload)
).status_code # rebind-ok: reassigned from the raised HTTPStatusError below
except HTTPStatusError as e:
status = e.response.status_code
except Exception as e: # noqa: BLE001 # transport/network failure re-queues the batch
verbose_logger.warning(
"New Relic Metrics: network error sending %s records, will retry - %s",
len(batch),
e,
)
return False
if 200 <= status < 300:
return True
if 400 <= status < 500 and status not in _RETRYABLE_CLIENT_STATUSES:
verbose_logger.warning(
"New Relic Metrics: %s from Metric API%s, dropping %s records.",
status,
" (permanent credential failure: invalid or revoked team ingest key)" if status == 403 else "",
len(batch),
)
return True
verbose_logger.warning(
"New Relic Metrics: %s from Metric API, will retry %s records",
status,
len(batch),
)
return False
def _requeue(self, batch: tuple[NewRelicMetricRecord, ...]) -> None:
"""Prepends ``batch`` in place (never by assignment: records appended by
concurrent requests during the flush await must survive), keeping
chronological order so the cap drops the oldest records first."""
self.log_queue[:0] = batch
overflow: Final = len(self.log_queue) - self.max_queue_size
if overflow > 0:
del self.log_queue[:overflow]
verbose_logger.warning(
"New Relic Metrics: retry queue exceeded max_queue_size=%s; dropped %s oldest records.",
self.max_queue_size,
overflow,
)
async def async_send_compressed_data(self, payload: tuple[NewRelicMetricEnvelope, ...]) -> Response:
compressed_data: Final = gzip.compress(safe_dumps(payload).encode("utf-8"))
headers: Final[Mapping[str, str]] = MappingProxyType(
{
"Content-Type": "application/json",
"Content-Encoding": "gzip",
"Api-Key": self.newrelic_api_key,
}
)
return await self.async_client.post(
url=self.metric_api_url,
data=compressed_data,
headers=headers,
)

View file

@ -0,0 +1,90 @@
"""
New Relic Team Handler
Used to get the NewRelicMetricsLogger for a given request.
Handles Key/Team Based New Relic metrics, following the same pattern as DataDogHandler.
"""
from typing import TYPE_CHECKING, Final
from typing_extensions import ReadOnly, TypedDict
from litellm._logging import verbose_logger
from litellm.litellm_core_utils.litellm_logging import StandardCallbackDynamicParams
from .newrelic_metrics import NewRelicMetricsLogger
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import DynamicLoggingCache
class NewRelicLoggingConfig(TypedDict):
newrelic_api_key: ReadOnly[str | None]
newrelic_region: ReadOnly[str | None]
class NewRelicHandler:
@staticmethod
def get_newrelic_logger_for_request(
standard_callback_dynamic_params: StandardCallbackDynamicParams,
in_memory_dynamic_logger_cache: "DynamicLoggingCache",
) -> NewRelicMetricsLogger:
"""
Get a team-scoped NewRelicMetricsLogger for a given request.
Resolves and caches per-team NewRelicMetricsLogger instances using
DynamicLoggingCache, keyed by the team's New Relic credentials. Each unique
set of credentials gets its own logger instance with its own batch/flush loop.
Note: This handler is only called when a team-scoped newrelic_api_key is
present. The trace logger for the ``newrelic`` callback (OTel v2 / legacy
agent) is managed separately by _init_custom_logger_compatible_class via
_in_memory_loggers.
"""
_credentials: Final = NewRelicHandler.get_dynamic_newrelic_logging_config(
standard_callback_dynamic_params=standard_callback_dynamic_params,
)
temp_newrelic_logger = in_memory_dynamic_logger_cache.get_cache(
credentials=_credentials, service_name="newrelic"
)
if temp_newrelic_logger is None:
temp_newrelic_logger = NewRelicHandler._create_newrelic_logger_from_credentials(
credentials=_credentials,
in_memory_dynamic_logger_cache=in_memory_dynamic_logger_cache,
)
return temp_newrelic_logger
@staticmethod
def _create_newrelic_logger_from_credentials(
credentials: NewRelicLoggingConfig,
in_memory_dynamic_logger_cache: "DynamicLoggingCache",
) -> NewRelicMetricsLogger:
newrelic_logger: Final = NewRelicMetricsLogger(
newrelic_api_key=credentials.get("newrelic_api_key") or "",
newrelic_region=credentials.get("newrelic_region"),
)
in_memory_dynamic_logger_cache.set_cache(
credentials=credentials,
service_name="newrelic",
logging_obj=newrelic_logger,
)
verbose_logger.debug("New Relic: Created and cached new NewRelicMetricsLogger for team-scoped credentials")
return newrelic_logger
@staticmethod
def get_dynamic_newrelic_logging_config(
standard_callback_dynamic_params: StandardCallbackDynamicParams,
) -> NewRelicLoggingConfig:
return NewRelicLoggingConfig(
newrelic_api_key=standard_callback_dynamic_params.get("newrelic_api_key"),
newrelic_region=standard_callback_dynamic_params.get("newrelic_region"),
)
@staticmethod
def _dynamic_newrelic_credentials_are_passed(
standard_callback_dynamic_params: StandardCallbackDynamicParams,
) -> bool:
return standard_callback_dynamic_params.get("newrelic_api_key") is not None

View file

@ -22,6 +22,7 @@ from litellm.integrations.opentelemetry_utils.gen_ai_semconv import (
)
from litellm.integrations.otel.model.db_endpoint import db_span_attributes
from litellm.integrations.otel.model.semconv import Metric
from litellm.litellm_core_utils.internal_call_metadata import is_unbilled_non_inference_call_from_params
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
from litellm.litellm_core_utils.secret_redaction import redact_string
from litellm.litellm_core_utils.service_tier_utils import (
@ -1643,7 +1644,12 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
if self._operation_duration_histogram:
self._operation_duration_histogram.record(duration_s, attributes=common_attrs)
if response_obj and (usage := response_obj.get("usage")) and self._token_usage_histogram:
if (
self._token_usage_histogram
and response_obj
and not is_unbilled_non_inference_call_from_params(kwargs.get("call_type"), params, response_obj)
and (usage := response_obj.get("usage"))
):
in_attrs: Final = {**common_attrs, TOKEN_TYPE_ATTRIBUTE: "input"}
out_attrs: Final = {**common_attrs, TOKEN_TYPE_ATTRIBUTE: "output"}
self._token_usage_histogram.record(usage.get("prompt_tokens", 0), attributes=in_attrs)
@ -1719,6 +1725,11 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
if not self._time_per_output_token_histogram:
return
if is_unbilled_non_inference_call_from_params(
kwargs.get("call_type"), kwargs.get("litellm_params"), response_obj
):
return
# Get completion tokens from response_obj
completion_tokens = None
if response_obj and (usage := response_obj.get("usage")):
@ -2049,6 +2060,26 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
# serialise to JSON once so set_attribute never coerces.
guardrail_span.set_attribute("guardrail_violation_categories", safe_dumps(violation_categories))
# Billable usage counters and USD cost stamped by the provider hook
# (e.g. Azure Prompt Shield text records, Bedrock policy units).
guardrail_usage = guardrail_information.get("guardrail_usage")
if guardrail_usage is not None:
guardrail_span.set_attribute("guardrail_usage", safe_dumps(guardrail_usage))
guardrail_cost = guardrail_information.get("guardrail_cost")
if guardrail_cost is not None:
self.safe_set_attribute(
span=guardrail_span,
key="guardrail_cost",
value=guardrail_cost,
)
guardrail_cost_in_spend = guardrail_information.get("guardrail_cost_in_spend")
if isinstance(guardrail_cost_in_spend, bool):
self.safe_set_attribute(
span=guardrail_span,
key="guardrail_cost_in_spend",
value=guardrail_cost_in_spend,
)
self._set_team_attributes_from_kwargs(guardrail_span, kwargs)
guardrail_span.end(end_time=self._to_ns(end_time_datetime))
@ -2468,7 +2499,14 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
self._set_service_tier_attributes(span=span, standard_logging_payload=standard_logging_payload)
usage: Final = response_obj and response_obj.get("usage")
usage: Final = (
response_obj.get("usage")
if response_obj
and not is_unbilled_non_inference_call_from_params(
kwargs.get("call_type"), litellm_params, response_obj
)
else None
)
if usage:
self.safe_set_attribute(
span=span,

View file

@ -146,7 +146,7 @@ class SpanEmitter:
For callers that own and manage their own span lifecycle. ``tracer``
overrides the bound tracer for this span only, used for per-request
multi-tenant credential routing. ``links`` records related-but-not-parent
spans (e.g. the transport span of an MCP message, per MCP semconv).
spans (e.g. the trace context an MCP client propagated in ``params._meta``).
"""
return (tracer or self._tracer).start_span(
name,
@ -156,6 +156,12 @@ class SpanEmitter:
links=list(links) if links else None,
)
def mark_emitted(self, dedup_key: str | None, role: SpanRole) -> None:
"""Register a span emitted outside :meth:`emit` (the boundary-opened
LLM-call span closed via :meth:`finish_span`) so a later :meth:`emit`
for the same ``(dedup_key, role)`` deduplicates against it."""
self._seen(dedup_key, role)
def _seen(self, dedup_key: str | None, role: SpanRole) -> bool:
"""Return True once a ``(dedup_key, role)`` pair has been emitted.
@ -190,8 +196,8 @@ class SpanEmitter:
Return the span, or ``None`` if it was deduplicated away. ``tracer``
overrides the bound tracer for this span, used for per-request routing.
``links`` records related-but-not-parent spans (the transport span of an
MCP message).
``links`` records related-but-not-parent spans (e.g. the trace context an
MCP client propagated in ``params._meta``).
"""
# LLM-call and MCP tool-call spans carry a dedup key (their request's
# call id), so a sync+async double-firing coalesces. ``isinstance`` narrows

View file

@ -390,10 +390,10 @@ class OpenTelemetryV2(CustomLogger):
MCP tool calls reach the success/failure callbacks like any other request
(with ``call_type`` ``call_mcp_tool``), but they are not LLM calls and have
no ``pre_call`` carrier so they get their own CLIENT span here. Per the MCP
semconv it parents to the trace context the client propagated in
``params._meta`` (or starts a new root) and links the transport span, rather
than nesting under the HTTP/session span. Returns whether it handled the
no ``pre_call`` carrier so they get their own CLIENT span here. It nests
under the transport span of the request carrying this message, and trace
context the client propagated in ``params._meta`` is recorded as a span
link (see ``resolve_mcp_span_context``). Returns whether it handled the
event, so the caller skips the LLM-call path. The whole span is emitted at
once (there is no boundary to open it at), deduped on the call id.
"""
@ -436,9 +436,9 @@ class OpenTelemetryV2(CustomLogger):
Like a tool call, listing reaches the success/failure callbacks (here with
``call_type`` ``list_mcp_tools``) with no ``pre_call`` carrier, so it gets its
own CLIENT span. Per the MCP semconv it parents to the ``params._meta`` trace
context (or starts a new root) and links the transport span, rather than
nesting under the HTTP/session span. Returns whether it handled the event so
own CLIENT span, nested under the transport span of the request carrying
this message with any ``params._meta`` trace context recorded as a span
link (see ``resolve_mcp_span_context``). Returns whether it handled the event so
the caller skips the LLM-call path.
"""
raw_payload: Final = kwargs.get("standard_logging_object")
@ -484,10 +484,15 @@ class OpenTelemetryV2(CustomLogger):
# ``pop`` is the dedup: this method runs from both the success and failure
# paths, and whichever fires first removes the carrier and closes the span.
carrier: Final = self._open_llm_calls.pop(call_id, None) if call_id else None
if carrier is None:
# A missing carrier does not always mean nothing happened: a team/key-scoped
# logger is a success/failure callback only, so ``pre_call`` never reaches it
# and no carrier exists. The payload plus the request-level provider-handoff
# stamp (``upstream_started``) is the affirmative signal of a real call; a
# gate rejection carries ``is_no_upstream_call`` and gets no span.
if carrier is None and (call.is_no_upstream_call or not call.upstream_started or call.payload is None):
return None
try:
return self._finish_carrier(carrier, call, end_time)
return self._finish_carrier(carrier, call, start_time, end_time)
finally:
# After the span has ended, so a release-triggered provider shutdown
# force-flushes it out rather than racing its enqueue.
@ -497,8 +502,11 @@ class OpenTelemetryV2(CustomLogger):
"""Remember an in-flight LLM call, evicting the oldest if over budget.
A call that opens but never closes (a stream that only fires stream
events) would linger otherwise; the evicted span is simply dropped
(never exported).
events) would linger otherwise. Eviction only drops the boundary carrier,
not the call: if that call later closes as a real completed call, it still
emits through the deferred branch in ``_close_llm_call`` (the same path a
team/key-scoped logger uses, since it never opens a carrier), deduplicated
by call id. Only a call that is evicted and never closes goes unexported.
"""
self._open_llm_calls[call_id] = carrier
if len(self._open_llm_calls) > _OPEN_CALLS_MAX:
@ -512,15 +520,20 @@ class OpenTelemetryV2(CustomLogger):
def _finish_carrier(
self,
carrier: _LLMCallSpan,
carrier: "_LLMCallSpan | None",
call: LLMCallEvent,
start_time: datetime | float | None,
end_time: datetime | float | None,
) -> Span | None:
payload: Final = call.payload
call_id: Final = call.call_id
if payload is None:
if carrier.span is not None:
if carrier is not None and carrier.span is not None:
# Opened at the boundary but the payload never materialized — end
# it (named provisionally) so it isn't leaked as an open span.
# it (named provisionally) so it isn't leaked as an open span, and
# register the dedup marker so a later payload-carrying close for
# the same call id cannot re-emit through the deferred branch.
self._emitter.mark_emitted(call_id, SpanRole.LLM_CALL)
carrier.span.end(end_time=to_ns(end_time))
return None
data: Final = LLMCallSpanData.from_standard_logging_payload(
@ -529,10 +542,13 @@ class OpenTelemetryV2(CustomLogger):
time_to_first_chunk_seconds=call.time_to_first_chunk_seconds,
)
end_time_ns: Final = to_ns(end_time)
if carrier.span is not None:
if carrier is not None and carrier.span is not None:
# Born at the boundary: stamp attributes from the typed payload, set
# status, and end it. Its parent (the server span) was captured at
# creation from real ambient context.
# creation from real ambient context. Register the dedup marker so a
# second close for the same call id (success then failure on one
# logging object) cannot re-emit through the deferred branch.
self._emitter.mark_emitted(call_id, SpanRole.LLM_CALL)
self._emitter.finish_span(SpanRole.LLM_CALL, carrier.span, data, end_time_ns=end_time_ns)
return carrier.span
# Deferred: ``pre_call`` saw no recordable parent, so create the span now.
@ -549,7 +565,7 @@ class OpenTelemetryV2(CustomLogger):
SpanRole.LLM_CALL,
data,
parent_context=(set_span_in_context(INVALID_SPAN, parent_ctx) if route.detached else parent_ctx),
start_time_ns=carrier.start_time_ns,
start_time_ns=(carrier.start_time_ns if carrier is not None else to_ns(start_time)),
end_time_ns=end_time_ns,
tracer=route.tracer,
links=_request_trace_links(parent_ctx) if route.detached else None,

View file

@ -136,6 +136,9 @@ class GenAIMapper:
LiteLLM.GUARDRAIL_ID: lambda d: d.guardrail_id,
LiteLLM.GUARDRAIL_POLICY_TEMPLATE: lambda d: d.policy_template,
LiteLLM.GUARDRAIL_DETECTION_METHOD: lambda d: d.detection_method,
LiteLLM.GUARDRAIL_USAGE: lambda d: d.usage_json,
LiteLLM.GUARDRAIL_COST: lambda d: d.cost,
LiteLLM.GUARDRAIL_COST_IN_SPEND: lambda d: d.cost_in_spend,
}
_SERVICE_ATTRS: dict[str, Callable[[ServiceSpanData], AttrValue | None]] = {

View file

@ -39,6 +39,7 @@ class ExporterOwner(str, Enum):
WEAVE_OTEL = "weave_otel"
LEVO = "levo"
AGENTOPS = "agentops"
NEWRELIC = "newrelic"
class _OTelV2Flag(BaseSettings):
@ -97,6 +98,15 @@ class ExporterSpec(BaseModel):
"auto (Simple for console/in_memory, Batch otherwise)."
),
)
requires_headers: bool = Field(
default=False,
description=(
"Skip this exporter when no headers are resolved. For destinations "
"that reject unauthenticated exports (e.g. New Relic), a spec kept "
"only as the per-request credential-stamping target would otherwise "
"export keyless traffic and produce a 4xx for every span batch."
),
)
class OpenTelemetryV2Config(BaseSettings):

View file

@ -203,6 +203,11 @@ class LLMCallEvent:
# True for synthetic proxy-gate logs (auth / rate-limit rejections): they fire
# the ``pre_call`` hook but never made an upstream call, so they get no span.
is_no_upstream_call: bool
# True once the request handed off to a provider (``pre_call`` stamped
# ``api_call_start_time``). The affirmative signal that an LLM call was
# actually attempted — router pre-call rejections, SDK failures before the
# provider handoff, and standalone guardrail runs all lack it.
upstream_started: bool
# A best-effort ``"{operation} {model}"`` name known at ``pre_call`` time. The
# span is renamed from the typed payload at close (``finish_span``); this only
# needs to be reasonable for a span that never gets closed (a leak).
@ -221,6 +226,7 @@ class LLMCallEvent:
dynamic_params=kwargs.get("standard_callback_dynamic_params"),
auth_metadata=auth_metadata(payload, kwargs),
is_no_upstream_call=bool(kwargs.get(LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL)),
upstream_started=kwargs.get("api_call_start_time") is not None,
provisional_span_name=f"{operation.value} {model}".strip(),
time_to_first_chunk_seconds=time_to_first_chunk_seconds(kwargs),
)

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