mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-10 22:41:41 +00:00
Merge branch 'litellm_internal_staging' into feature/improve-gigachat-provider
This commit is contained in:
commit
9da1ccd12f
1623 changed files with 162715 additions and 30719 deletions
|
|
@ -111,6 +111,28 @@ commands:
|
|||
- wait_for_service:
|
||||
url: tcp://localhost:6379
|
||||
timeout: "60"
|
||||
start_openai_record_replay_proxy:
|
||||
description: "Start the record/replay proxy (tests/_openai_record_replay_proxy.py) on host port 8090 and wait until healthy. Models whose api_base points here replay recorded provider responses, so the E2E run neither pays for nor depends on the live provider. The default upstream is OpenAI; a non-OpenAI model must point its api_base at /__recorder_upstream/<host>/ so the recorder forwards there instead of defaulting to OpenAI. Run after uv deps are synced."
|
||||
steps:
|
||||
- run:
|
||||
name: Start record/replay proxy
|
||||
background: true
|
||||
command: |
|
||||
CASSETTE_REDIS_URL="$CASSETTE_REDIS_URL" \
|
||||
RECORDER_UPSTREAM_BASE_URL="https://api.openai.com" \
|
||||
uv run --no-sync python tests/_openai_record_replay_proxy.py --host 0.0.0.0 --port 8090
|
||||
- run:
|
||||
name: Wait for record/replay proxy
|
||||
command: |
|
||||
for i in $(seq 1 30); do
|
||||
if curl -sf http://localhost:8090/__recorder_health >/dev/null 2>&1; then
|
||||
echo "record/replay proxy is up"
|
||||
exit 0
|
||||
fi
|
||||
sleep 1
|
||||
done
|
||||
echo "record/replay proxy did not become ready" >&2
|
||||
exit 1
|
||||
setup_litellm_enterprise_pip:
|
||||
steps:
|
||||
- run:
|
||||
|
|
@ -452,6 +474,120 @@ 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
|
||||
- setup_google_dns
|
||||
- install_uv
|
||||
- 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
|
||||
- setup_google_dns
|
||||
- install_uv
|
||||
- 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
|
||||
- setup_google_dns
|
||||
- install_uv
|
||||
- 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
|
||||
|
|
@ -1511,6 +1647,7 @@ jobs:
|
|||
command: |
|
||||
zstd -d litellm-docker-database.tar.zst --stdout | docker load
|
||||
docker tag litellm-docker-database:ci my-app:latest
|
||||
- start_openai_record_replay_proxy
|
||||
- run:
|
||||
name: Run Docker container
|
||||
command: |
|
||||
|
|
@ -1541,6 +1678,7 @@ jobs:
|
|||
-e LANGFUSE_PROJECT2_PUBLIC=$LANGFUSE_PROJECT2_PUBLIC \
|
||||
-e LANGFUSE_PROJECT1_SECRET=$LANGFUSE_PROJECT1_SECRET \
|
||||
-e LANGFUSE_PROJECT2_SECRET=$LANGFUSE_PROJECT2_SECRET \
|
||||
-e RECORDER_OPENAI_BASE_URL=http://host.docker.internal:8090/v1 \
|
||||
--add-host host.docker.internal:host-gateway \
|
||||
--name my-app \
|
||||
-v $(pwd)/proxy_server_config.yaml:/app/config.yaml \
|
||||
|
|
@ -1678,6 +1816,7 @@ jobs:
|
|||
command: |
|
||||
zstd -d litellm-docker-database.tar.zst --stdout | docker load
|
||||
docker images | grep litellm-docker-database
|
||||
- start_openai_record_replay_proxy
|
||||
- run:
|
||||
name: Run Docker container
|
||||
# intentionally give bad redis credentials here
|
||||
|
|
@ -1701,6 +1840,7 @@ jobs:
|
|||
-e DD_SITE=$DD_SITE \
|
||||
-e AWS_REGION_NAME=$AWS_REGION_NAME \
|
||||
-e COHERE_API_KEY=$COHERE_API_KEY \
|
||||
-e RECORDER_COHERE_BASE_URL=http://host.docker.internal:8090/__recorder_upstream/api.cohere.com \
|
||||
-e GCS_FLUSH_INTERVAL="1" \
|
||||
--add-host host.docker.internal:host-gateway \
|
||||
--name my-app \
|
||||
|
|
@ -2266,6 +2406,7 @@ jobs:
|
|||
command: |
|
||||
zstd -d litellm-docker-database.tar.zst --stdout | docker load
|
||||
docker images | grep litellm-docker-database
|
||||
- start_openai_record_replay_proxy
|
||||
- run:
|
||||
name: Run Docker container with test config
|
||||
command: |
|
||||
|
|
@ -2274,6 +2415,7 @@ jobs:
|
|||
-e DATABASE_URL=postgresql://postgres:postgres@host.docker.internal:5432/circle_test \
|
||||
-e LITELLM_MASTER_KEY="sk-1234" \
|
||||
-e ANTHROPIC_API_KEY=$ANTHROPIC_API_KEY \
|
||||
-e RECORDER_ANTHROPIC_BASE_URL=http://host.docker.internal:8090/__recorder_upstream/api.anthropic.com \
|
||||
-e AWS_ACCESS_KEY_ID=$AWS_ACCESS_KEY_ID \
|
||||
-e AWS_SECRET_ACCESS_KEY=$AWS_SECRET_ACCESS_KEY \
|
||||
-e AWS_REGION_NAME="us-east-1" \
|
||||
|
|
@ -2548,6 +2690,122 @@ jobs:
|
|||
path: ui/litellm-dashboard/playwright-report
|
||||
destination: e2e-playwright-report
|
||||
|
||||
e2e_ui_testing_server_root_path:
|
||||
docker:
|
||||
- image: cimg/python:3.12-browsers@sha256:b432899af01c9a311bf74f4f22e9ada2e5306d4b1b4383f8d29e1228a5844ef2
|
||||
auth:
|
||||
username: ${DOCKERHUB_USERNAME}
|
||||
password: ${DOCKERHUB_PASSWORD}
|
||||
- image: cimg/postgres:16.0@sha256:b125148bc76e8e8eee5eb3ad6020a3a14110a14e8192f1c645128afebe2e2f84
|
||||
environment:
|
||||
POSTGRES_USER: e2euser
|
||||
POSTGRES_PASSWORD: e2epassword
|
||||
POSTGRES_DB: litellm_e2e
|
||||
resource_class: large
|
||||
working_directory: ~/project
|
||||
environment:
|
||||
DATABASE_URL: "postgresql://e2euser:e2epassword@localhost:5432/litellm_e2e"
|
||||
CI: "true"
|
||||
# The whole job exercises the proxy mounted under a prefix. SERVER_ROOT_PATH
|
||||
# is read both by the proxy at boot (to rewrite the built UI bundle in place)
|
||||
# and by migration.serverRootPath.config.ts, which refuses to run without it.
|
||||
SERVER_ROOT_PATH: "/litellm"
|
||||
steps:
|
||||
- checkout
|
||||
- setup_google_dns
|
||||
- install_uv
|
||||
- restore_cache:
|
||||
keys:
|
||||
- v1-uv-cache-{{ checksum "uv.lock" }}
|
||||
- run:
|
||||
name: Install Python dependencies
|
||||
command: |
|
||||
uv sync --frozen --all-groups --all-extras --python 3.12
|
||||
uv run --no-sync python -m prisma generate --schema litellm/proxy/schema.prisma
|
||||
- save_cache:
|
||||
key: v1-uv-cache-{{ checksum "uv.lock" }}
|
||||
paths:
|
||||
- ~/.cache/uv
|
||||
- restore_cache:
|
||||
keys:
|
||||
- ui-e2e-node-deps-v2-{{ checksum "ui/litellm-dashboard/package-lock.json" }}
|
||||
- run:
|
||||
name: Install Node dependencies and Playwright
|
||||
command: |
|
||||
cd ui/litellm-dashboard
|
||||
npm ci
|
||||
npx playwright install chromium
|
||||
- save_cache:
|
||||
key: ui-e2e-node-deps-v2-{{ checksum "ui/litellm-dashboard/package-lock.json" }}
|
||||
paths:
|
||||
- ui/litellm-dashboard/node_modules
|
||||
- ~/.cache/ms-playwright
|
||||
- run:
|
||||
name: Build UI from source
|
||||
command: |
|
||||
cd ui/litellm-dashboard
|
||||
npm run build
|
||||
rm -rf ../../litellm/proxy/_experimental/out
|
||||
mv out ../../litellm/proxy/_experimental/out
|
||||
find ../../litellm/proxy/_experimental/out -name '*.html' ! -name 'index.html' | while read -r f; do
|
||||
d="${f%.html}"; mkdir -p "$d"; mv "$f" "$d/index.html"
|
||||
done
|
||||
- wait_for_service:
|
||||
url: tcp://localhost:5432
|
||||
timeout: "30"
|
||||
- run:
|
||||
name: Push Prisma schema
|
||||
command: uv run --no-sync python -m prisma db push --schema litellm/proxy/schema.prisma --accept-data-loss
|
||||
- run:
|
||||
name: Seed database
|
||||
command: |
|
||||
PGPASSWORD=e2epassword psql -h localhost -p 5432 -U e2euser -d litellm_e2e \
|
||||
-f ui/litellm-dashboard/e2e_tests/fixtures/seed.sql
|
||||
- run:
|
||||
name: Start mock LLM server
|
||||
command: uv run --no-sync python ui/litellm-dashboard/e2e_tests/fixtures/mock_llm_server/server.py
|
||||
background: true
|
||||
- run:
|
||||
name: Start LiteLLM proxy under a server root path
|
||||
environment:
|
||||
LITELLM_MASTER_KEY: "sk-1234"
|
||||
MOCK_LLM_URL: "http://127.0.0.1:8090/v1"
|
||||
DISABLE_SCHEMA_UPDATE: "true"
|
||||
# Output flows to this step's own log, so a boot crash is visible here
|
||||
# rather than swallowed by a downstream readiness probe.
|
||||
command: |
|
||||
LITELLM_LICENSE="$LITELLM_LICENSE" \
|
||||
uv run --no-sync python -m litellm.proxy.proxy_cli \
|
||||
--config ui/litellm-dashboard/e2e_tests/fixtures/config.yml \
|
||||
--port 4000
|
||||
background: true
|
||||
- run:
|
||||
name: Wait for prefixed proxy to be ready
|
||||
command: |
|
||||
for i in $(seq 1 60); do
|
||||
HTTP_CODE=$(curl -s -o /dev/null -w "%{http_code}" --max-time 5 -H "Authorization: Bearer sk-1234" http://127.0.0.1:4000/litellm/health 2>/dev/null || true)
|
||||
if [ "$HTTP_CODE" = "200" ]; then
|
||||
echo "Prefixed proxy is ready"
|
||||
exit 0
|
||||
fi
|
||||
sleep 2
|
||||
done
|
||||
echo "Prefixed proxy failed to start; see the 'Start LiteLLM proxy under a server root path' step for the boot log"
|
||||
exit 1
|
||||
- run:
|
||||
name: Run migration smoke under SERVER_ROOT_PATH
|
||||
command: |
|
||||
cd ui/litellm-dashboard
|
||||
LITELLM_LICENSE="$LITELLM_LICENSE" \
|
||||
npx playwright test --config e2e_tests/migration.serverRootPath.config.ts
|
||||
no_output_timeout: 10m
|
||||
- store_artifacts:
|
||||
path: ui/litellm-dashboard/test-results
|
||||
destination: e2e-server-root-path-test-results
|
||||
- store_artifacts:
|
||||
path: ui/litellm-dashboard/playwright-report
|
||||
destination: e2e-server-root-path-playwright-report
|
||||
|
||||
build_docker_database_image:
|
||||
machine:
|
||||
image: ubuntu-2204:2024.04.1
|
||||
|
|
@ -2643,10 +2901,18 @@ 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:
|
||||
filters: *main_branches
|
||||
- e2e_ui_testing_server_root_path:
|
||||
filters: *main_branches
|
||||
- build_and_test:
|
||||
requires:
|
||||
- build_docker_database_image
|
||||
|
|
|
|||
|
|
@ -8,3 +8,6 @@
|
|||
|
||||
# Update pydantic code to fix warnings (GH-3600)
|
||||
876840e9957bc7e9f7d6a2b58c4d7c53dad16481
|
||||
|
||||
# style(ui): run prettier --write across the dashboard (#29622)
|
||||
7edf3a9cb55548b143df1692f4ed7c4681d7fcf7
|
||||
|
|
|
|||
3
.gitattributes
vendored
3
.gitattributes
vendored
|
|
@ -1 +1,2 @@
|
|||
*.ipynb linguist-vendored
|
||||
*.ipynb linguist-vendored
|
||||
ui/litellm-dashboard/src/lib/http/schema.d.ts linguist-generated
|
||||
75
.githooks/commit-msg
Executable file
75
.githooks/commit-msg
Executable file
|
|
@ -0,0 +1,75 @@
|
|||
#!/usr/bin/env bash
|
||||
#
|
||||
# commit-msg — enforce Conventional Commits 1.0.0
|
||||
# https://www.conventionalcommits.org/en/v1.0.0/
|
||||
#
|
||||
# Subject format: <type>(<scope>)!: <description>
|
||||
# - <type> must be one of the angular types (feat, fix, ...)
|
||||
# - (<scope>) is optional
|
||||
# - ! is optional and marks a breaking change
|
||||
# - <description> is mandatory and must be non-empty
|
||||
#
|
||||
# Bypass: commit with --no-verify.
|
||||
# Merge, revert, fixup!, squash!, and amend! messages are passed through.
|
||||
|
||||
set -eu
|
||||
|
||||
COMMIT_MSG_FILE="${1:-}"
|
||||
if [ -z "$COMMIT_MSG_FILE" ] || [ ! -f "$COMMIT_MSG_FILE" ]; then
|
||||
echo "commit-msg: missing commit message file" >&2
|
||||
exit 1
|
||||
fi
|
||||
|
||||
# First non-comment, non-empty line is the subject.
|
||||
subject=""
|
||||
while IFS= read -r line || [ -n "$line" ]; do
|
||||
case "$line" in
|
||||
''|'#'*) continue ;;
|
||||
esac
|
||||
subject="$line"
|
||||
break
|
||||
done < "$COMMIT_MSG_FILE"
|
||||
|
||||
if [ -z "$subject" ]; then
|
||||
echo "commit-msg: empty commit message" >&2
|
||||
exit 1
|
||||
fi
|
||||
|
||||
# Pass-through commits generated by git itself.
|
||||
case "$subject" in
|
||||
"Merge "*|"Revert \""*|"fixup! "*|"squash! "*|"amend! "*)
|
||||
exit 0
|
||||
;;
|
||||
esac
|
||||
|
||||
ALLOWED_TYPES="feat|fix|docs|style|refactor|perf|test|build|ci|chore|revert"
|
||||
# Description must not start with an uppercase letter — kept in sync with the
|
||||
# subjectPattern in .github/workflows/conventional-commits.yml so the local
|
||||
# hook is the strictly tighter of the two gates. (Without this guard, a commit
|
||||
# like "feat: Add thing" passes locally but fails the PR-title CI check.)
|
||||
PATTERN="^(${ALLOWED_TYPES})(\([^)]+\))?!?: [^A-Z].*"
|
||||
|
||||
if printf '%s' "$subject" | grep -Eq "$PATTERN"; then
|
||||
exit 0
|
||||
fi
|
||||
|
||||
cat >&2 <<EOF
|
||||
✗ Commit message does not follow Conventional Commits.
|
||||
|
||||
Got: $subject
|
||||
|
||||
Expected: <type>(<scope>)!: <description>
|
||||
(description must start with a lowercase letter)
|
||||
|
||||
Allowed types: feat, fix, docs, style, refactor, perf, test, build, ci, chore, revert
|
||||
Examples:
|
||||
feat(router): add weighted round-robin strategy
|
||||
fix(bedrock): decouple STS region from aws_region_name
|
||||
chore(deps): bump black to 26.3.1
|
||||
refactor!: drop Python 3.8 support
|
||||
|
||||
See https://www.conventionalcommits.org/en/v1.0.0/
|
||||
|
||||
To bypass (use sparingly): git commit --no-verify
|
||||
EOF
|
||||
exit 1
|
||||
92
.githooks/pre-push
Executable file
92
.githooks/pre-push
Executable file
|
|
@ -0,0 +1,92 @@
|
|||
#!/usr/bin/env bash
|
||||
#
|
||||
# pre-push — enforce Conventional Branches
|
||||
# https://conventional-branch.github.io/
|
||||
#
|
||||
# Branch format: <type>/<description>
|
||||
# <type> must be one of: feature, bugfix, hotfix, release, chore
|
||||
#
|
||||
# Protected branches (always allowed):
|
||||
# - main
|
||||
# - litellm_internal_staging
|
||||
# - dependabot/*
|
||||
# - gh-readonly-queue/*
|
||||
#
|
||||
# Tag pushes and branch deletions are skipped.
|
||||
# Bypass: git push --no-verify.
|
||||
|
||||
set -eu
|
||||
|
||||
ZERO_OID="0000000000000000000000000000000000000000"
|
||||
ZERO_OID_SHA256="0000000000000000000000000000000000000000000000000000000000000000"
|
||||
ALLOWED_TYPES="feature|bugfix|hotfix|release|chore"
|
||||
BRANCH_PATTERN="^(${ALLOWED_TYPES})/.+"
|
||||
|
||||
PROTECTED_NAMES="main litellm_internal_staging"
|
||||
PROTECTED_PREFIXES="dependabot/ gh-readonly-queue/"
|
||||
|
||||
is_protected() {
|
||||
branch="$1"
|
||||
for name in $PROTECTED_NAMES; do
|
||||
if [ "$branch" = "$name" ]; then
|
||||
return 0
|
||||
fi
|
||||
done
|
||||
for prefix in $PROTECTED_PREFIXES; do
|
||||
case "$branch" in "$prefix"*) return 0 ;; esac
|
||||
done
|
||||
return 1
|
||||
}
|
||||
|
||||
invalid=""
|
||||
|
||||
while read -r local_ref local_oid remote_ref remote_oid; do
|
||||
# Branch deletion (no local commit being pushed).
|
||||
if [ "$local_oid" = "$ZERO_OID" ] || [ "$local_oid" = "$ZERO_OID_SHA256" ]; then
|
||||
continue
|
||||
fi
|
||||
|
||||
# Only validate branch pushes; ignore tags and other ref namespaces.
|
||||
case "$remote_ref" in
|
||||
refs/heads/*) ;;
|
||||
*) continue ;;
|
||||
esac
|
||||
|
||||
branch="${remote_ref#refs/heads/}"
|
||||
|
||||
if is_protected "$branch"; then
|
||||
continue
|
||||
fi
|
||||
|
||||
if ! printf '%s' "$branch" | grep -Eq "$BRANCH_PATTERN"; then
|
||||
invalid="$invalid $branch"
|
||||
fi
|
||||
done
|
||||
|
||||
if [ -n "$invalid" ]; then
|
||||
cat >&2 <<EOF
|
||||
✗ Branch name does not follow Conventional Branches.
|
||||
|
||||
Invalid:$invalid
|
||||
|
||||
Expected: <type>/<description>
|
||||
|
||||
Allowed types: feature, bugfix, hotfix, release, chore
|
||||
Examples:
|
||||
feature/weighted-round-robin
|
||||
bugfix/streaming-empty-chunks
|
||||
chore/bump-deps
|
||||
hotfix/auth-bypass
|
||||
|
||||
Protected (always allowed): main, litellm_internal_staging,
|
||||
dependabot/*, gh-readonly-queue/*.
|
||||
|
||||
See https://conventional-branch.github.io/
|
||||
|
||||
Rename with: git branch -m <new-name>
|
||||
To bypass (use sparingly): git push --no-verify
|
||||
EOF
|
||||
exit 1
|
||||
fi
|
||||
|
||||
exit 0
|
||||
40
.github/workflows/_test-unit-base.yml
vendored
40
.github/workflows/_test-unit-base.yml
vendored
|
|
@ -27,6 +27,11 @@ on:
|
|||
required: false
|
||||
type: number
|
||||
default: 10
|
||||
dist:
|
||||
description: "pytest-xdist distribution mode (loadscope|load|worksteal|loadfile|no)"
|
||||
required: false
|
||||
type: string
|
||||
default: "loadscope"
|
||||
artifact-name:
|
||||
description: "Unique name for the coverage artifact (must be unique per run)"
|
||||
required: true
|
||||
|
|
@ -82,18 +87,31 @@ jobs:
|
|||
MAX_FAILURES: ${{ inputs.max-failures }}
|
||||
WORKERS: ${{ inputs.workers }}
|
||||
RERUNS: ${{ inputs.reruns }}
|
||||
DIST: ${{ inputs.dist }}
|
||||
run: |
|
||||
uv run --no-sync pytest ${TEST_PATH:?} \
|
||||
--tb=short -vv \
|
||||
--maxfail="${MAX_FAILURES}" \
|
||||
-n "${WORKERS}" \
|
||||
--reruns "${RERUNS}" \
|
||||
--reruns-delay 1 \
|
||||
--dist=loadscope \
|
||||
--durations=20 \
|
||||
--cov=./litellm \
|
||||
--cov-report=xml:coverage.xml \
|
||||
--cov-config=pyproject.toml
|
||||
if [ "${WORKERS}" = "0" ]; then
|
||||
uv run --no-sync pytest ${TEST_PATH:?} \
|
||||
--tb=short -vv \
|
||||
--maxfail="${MAX_FAILURES}" \
|
||||
--reruns "${RERUNS}" \
|
||||
--reruns-delay 1 \
|
||||
--durations=20 \
|
||||
--cov=./litellm \
|
||||
--cov-report=xml:coverage.xml \
|
||||
--cov-config=pyproject.toml
|
||||
else
|
||||
uv run --no-sync pytest ${TEST_PATH:?} \
|
||||
--tb=short -vv \
|
||||
--maxfail="${MAX_FAILURES}" \
|
||||
-n "${WORKERS}" \
|
||||
--reruns "${RERUNS}" \
|
||||
--reruns-delay 1 \
|
||||
--dist="${DIST}" \
|
||||
--durations=20 \
|
||||
--cov=./litellm \
|
||||
--cov-report=xml:coverage.xml \
|
||||
--cov-config=pyproject.toml
|
||||
fi
|
||||
|
||||
- name: Save coverage report
|
||||
if: always()
|
||||
|
|
|
|||
190
.github/workflows/_test-unit-services-base.yml
vendored
190
.github/workflows/_test-unit-services-base.yml
vendored
|
|
@ -1,190 +0,0 @@
|
|||
name: _Unit Test Services Base (Reusable)
|
||||
|
||||
on:
|
||||
workflow_call:
|
||||
inputs:
|
||||
test-path:
|
||||
description: "Pytest path(s) to run"
|
||||
required: true
|
||||
type: string
|
||||
workers:
|
||||
description: "Number of pytest-xdist workers (0 = no parallelism)"
|
||||
required: false
|
||||
type: number
|
||||
default: 2
|
||||
reruns:
|
||||
description: "Number of reruns for flaky tests"
|
||||
required: false
|
||||
type: number
|
||||
default: 2
|
||||
timeout-minutes:
|
||||
description: "Job timeout in minutes"
|
||||
required: false
|
||||
type: number
|
||||
default: 20
|
||||
max-failures:
|
||||
description: "Stop after this many failures"
|
||||
required: false
|
||||
type: number
|
||||
default: 10
|
||||
enable-postgres:
|
||||
description: "Start a local Postgres service container and run Prisma migrations"
|
||||
required: false
|
||||
type: boolean
|
||||
default: false
|
||||
dist:
|
||||
description: "pytest-xdist distribution mode (loadscope|load|worksteal|loadfile|no)"
|
||||
required: false
|
||||
type: string
|
||||
default: "loadscope"
|
||||
artifact-name:
|
||||
description: "Unique name for the coverage artifact (must be unique per run)"
|
||||
required: false
|
||||
type: string
|
||||
default: "run"
|
||||
|
||||
permissions:
|
||||
contents: read
|
||||
|
||||
# The postgres service container below is spawned per-job on localhost and
|
||||
# destroyed with the job. Nothing outside the runner can reach it. The
|
||||
# user/password/database here are not secrets — they're bootstrap values
|
||||
# for a throwaway container — so we hardcode them instead of attaching
|
||||
# every matrix shard to a GHA environment just to read three "secrets"
|
||||
# (which also produces a "temporarily deployed to …" notification on the
|
||||
# PR timeline per shard per push).
|
||||
jobs:
|
||||
run:
|
||||
name: Run tests
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: ${{ inputs.timeout-minutes }}
|
||||
|
||||
services:
|
||||
postgres:
|
||||
image: postgres@sha256:705a5d5b5836f3fcba0d02c4d281e6a7dd9ed2dd4078640f08a1e1e9896e097d # postgres:14
|
||||
env:
|
||||
POSTGRES_USER: litellm
|
||||
POSTGRES_PASSWORD: litellm
|
||||
POSTGRES_DB: litellm_test
|
||||
ports:
|
||||
- 5432:5432
|
||||
options: >-
|
||||
--health-cmd "pg_isready"
|
||||
--health-interval 10s
|
||||
--health-timeout 5s
|
||||
--health-retries 5
|
||||
|
||||
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: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7
|
||||
with:
|
||||
version: "0.10.9"
|
||||
|
||||
- name: Cache uv dependencies
|
||||
uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0
|
||||
with:
|
||||
path: |
|
||||
~/.cache/uv
|
||||
.venv
|
||||
key: ${{ runner.os }}-uv-services-${{ hashFiles('uv.lock') }}
|
||||
restore-keys: |
|
||||
${{ runner.os }}-uv-services-
|
||||
|
||||
- name: Install dependencies
|
||||
run: |
|
||||
uv sync --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router
|
||||
|
||||
- name: Generate Prisma client
|
||||
env:
|
||||
PRISMA_BINARY_CACHE_DIR: ${{ runner.temp }}/prisma-cache
|
||||
run: |
|
||||
uv run --no-sync prisma generate --schema litellm/proxy/schema.prisma
|
||||
|
||||
- name: Run Prisma migrations
|
||||
if: ${{ inputs.enable-postgres }}
|
||||
env:
|
||||
DATABASE_URL: "postgresql://litellm:litellm@localhost:5432/litellm_test"
|
||||
run: |
|
||||
uv run --no-sync prisma db push --schema litellm/proxy/schema.prisma --accept-data-loss
|
||||
|
||||
- name: Run tests
|
||||
env:
|
||||
TEST_PATH: ${{ inputs.test-path }}
|
||||
MAX_FAILURES: ${{ inputs.max-failures }}
|
||||
WORKERS: ${{ inputs.workers }}
|
||||
RERUNS: ${{ inputs.reruns }}
|
||||
DIST: ${{ inputs.dist }}
|
||||
DATABASE_URL: ${{ inputs.enable-postgres && 'postgresql://litellm:litellm@localhost:5432/litellm_test' || '' }}
|
||||
run: |
|
||||
if [ "${WORKERS}" = "0" ]; then
|
||||
uv run --no-sync pytest ${TEST_PATH:?} \
|
||||
--tb=short -vv \
|
||||
--maxfail="${MAX_FAILURES}" \
|
||||
--reruns "${RERUNS}" \
|
||||
--reruns-delay 1 \
|
||||
--durations=20 \
|
||||
--cov=./litellm \
|
||||
--cov-report=xml:coverage.xml \
|
||||
--cov-config=pyproject.toml
|
||||
else
|
||||
uv run --no-sync pytest ${TEST_PATH:?} \
|
||||
--tb=short -vv \
|
||||
--maxfail="${MAX_FAILURES}" \
|
||||
-n "${WORKERS}" \
|
||||
--reruns "${RERUNS}" \
|
||||
--reruns-delay 1 \
|
||||
--dist="${DIST}" \
|
||||
--durations=20 \
|
||||
--cov=./litellm \
|
||||
--cov-report=xml:coverage.xml \
|
||||
--cov-config=pyproject.toml
|
||||
fi
|
||||
|
||||
- name: Save coverage report
|
||||
if: always()
|
||||
uses: actions/upload-artifact@4cec3d8aa04e39d1a68397de0c4cd6fb9dce8ec1 # v4.6.1
|
||||
with:
|
||||
name: coverage-${{ inputs.artifact-name }}-${{ github.run_id }}-${{ github.run_attempt }}
|
||||
path: coverage.xml
|
||||
retention-days: 1
|
||||
|
||||
upload-coverage:
|
||||
name: Upload coverage to Codecov
|
||||
needs: run
|
||||
if: always()
|
||||
runs-on: ubuntu-latest
|
||||
permissions:
|
||||
contents: read
|
||||
id-token: write
|
||||
pull-requests: write
|
||||
|
||||
steps:
|
||||
- name: Checkout code
|
||||
uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
|
||||
with:
|
||||
persist-credentials: false
|
||||
|
||||
- name: Download coverage report
|
||||
uses: actions/download-artifact@95815c38cf2ff2164869cbab79da8d1f422bc89e # v4.2.1
|
||||
with:
|
||||
pattern: coverage-${{ inputs.artifact-name }}-${{ github.run_id }}-${{ github.run_attempt }}
|
||||
path: coverage-reports
|
||||
merge-multiple: true
|
||||
|
||||
- name: Upload to Codecov
|
||||
uses: codecov/codecov-action@75cd11691c0faa626561e295848008c8a7dddffe # v5.5.4
|
||||
with:
|
||||
use_oidc: true
|
||||
directory: coverage-reports
|
||||
root_dir: ${{ github.workspace }}
|
||||
flags: ${{ inputs.artifact-name }}
|
||||
fail_ci_if_error: false
|
||||
84
.github/workflows/check-ui-api-types.yml
vendored
Normal file
84
.github/workflows/check-ui-api-types.yml
vendored
Normal file
|
|
@ -0,0 +1,84 @@
|
|||
name: Check UI API Types Sync
|
||||
|
||||
on:
|
||||
pull_request:
|
||||
paths:
|
||||
- "litellm/proxy/**"
|
||||
- "litellm/types/**"
|
||||
- "ui/litellm-dashboard/src/lib/http/schema.d.ts"
|
||||
- "ui/litellm-dashboard/scripts/gen-api-types.mjs"
|
||||
- "ui/litellm-dashboard/package.json"
|
||||
- "ui/litellm-dashboard/package-lock.json"
|
||||
- ".github/workflows/check-ui-api-types.yml"
|
||||
|
||||
permissions:
|
||||
contents: read
|
||||
|
||||
jobs:
|
||||
check-sync:
|
||||
name: Verify schema.d.ts matches the proxy OpenAPI spec
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 15
|
||||
steps:
|
||||
- name: Checkout repository
|
||||
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: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7
|
||||
with:
|
||||
version: "0.10.9"
|
||||
|
||||
- name: Cache uv dependencies
|
||||
uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0
|
||||
with:
|
||||
path: |
|
||||
~/.cache/uv
|
||||
.venv
|
||||
key: ${{ runner.os }}-uv-${{ hashFiles('uv.lock') }}
|
||||
restore-keys: |
|
||||
${{ runner.os }}-uv-
|
||||
|
||||
- name: Install backend dependencies
|
||||
run: uv sync --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router
|
||||
|
||||
- name: Generate Prisma client
|
||||
env:
|
||||
PRISMA_BINARY_CACHE_DIR: ${{ runner.temp }}/prisma-cache
|
||||
run: uv run --no-sync prisma generate --schema litellm/proxy/schema.prisma
|
||||
|
||||
- name: Set up Node.js
|
||||
uses: actions/setup-node@a0853c24544627f65ddf259abe73b1d18a591444 # v5.0
|
||||
with:
|
||||
node-version: "20"
|
||||
cache: "npm"
|
||||
cache-dependency-path: ui/litellm-dashboard/package-lock.json
|
||||
|
||||
- name: Install dashboard dependencies
|
||||
working-directory: ui/litellm-dashboard
|
||||
run: npm ci
|
||||
|
||||
- name: Regenerate types from the live spec
|
||||
working-directory: ui/litellm-dashboard
|
||||
env:
|
||||
LITELLM_PYTHON: "uv run --no-sync python"
|
||||
run: npm run gen:api
|
||||
|
||||
- name: Fail if types are stale
|
||||
run: |
|
||||
if ! git diff --exit-code -- ui/litellm-dashboard/src/lib/http/schema.d.ts; then
|
||||
echo "::error file=ui/litellm-dashboard/src/lib/http/schema.d.ts::Generated API types are out of sync with the proxy OpenAPI spec."
|
||||
echo ""
|
||||
echo "A backend route or model changed without regenerating the dashboard types."
|
||||
echo "To fix, run from ui/litellm-dashboard:"
|
||||
echo " npm run gen:api"
|
||||
echo "then commit the updated src/lib/http/schema.d.ts."
|
||||
exit 1
|
||||
fi
|
||||
echo "schema.d.ts is in sync with the proxy OpenAPI spec."
|
||||
46
.github/workflows/conventional-commits.yml
vendored
Normal file
46
.github/workflows/conventional-commits.yml
vendored
Normal file
|
|
@ -0,0 +1,46 @@
|
|||
name: Conventional PR Title
|
||||
|
||||
# Squash-merge replaces the merge commit subject with the PR title, so
|
||||
# enforcing Conventional Commits at the PR-title level is what actually gates
|
||||
# the commits that land on the default branch. The local commit-msg hook
|
||||
# (.githooks/commit-msg) is a best-effort assist; this workflow is the gate.
|
||||
#
|
||||
# See https://www.conventionalcommits.org/en/v1.0.0/
|
||||
|
||||
on:
|
||||
pull_request:
|
||||
types: [opened, edited, reopened, synchronize, labeled, unlabeled]
|
||||
|
||||
permissions:
|
||||
pull-requests: read
|
||||
|
||||
jobs:
|
||||
lint-pr-title:
|
||||
name: Validate PR title
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- name: Check title against Conventional Commits
|
||||
uses: amannn/action-semantic-pull-request@48f256284bd46cdaab1048c3721360e808335d50 # v6.1.1
|
||||
env:
|
||||
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
|
||||
with:
|
||||
# Must mirror the type list in .githooks/commit-msg.
|
||||
types: |
|
||||
feat
|
||||
fix
|
||||
docs
|
||||
style
|
||||
refactor
|
||||
perf
|
||||
test
|
||||
build
|
||||
ci
|
||||
chore
|
||||
revert
|
||||
requireScope: false
|
||||
subjectPattern: ^(?![A-Z]).+$
|
||||
subjectPatternError: |
|
||||
The subject "{subject}" must start with a lowercase character.
|
||||
# Allow merges/reverts that GitHub generates automatically.
|
||||
ignoreLabels: |
|
||||
ignore-semantic-pull-request
|
||||
33
.github/workflows/create-release.yml
vendored
33
.github/workflows/create-release.yml
vendored
|
|
@ -52,6 +52,22 @@ jobs:
|
|||
// are stable maintenance releases, not pre-releases.
|
||||
const isPrerelease = /(?:rc|nightly|alpha|beta|[-.]dev)/i.test(tag);
|
||||
|
||||
// A stable release should only claim the repo "latest" badge when its
|
||||
// version is >= the current latest. Otherwise a backport (e.g. 1.84.6)
|
||||
// would steal "latest" from a newer line (e.g. 1.88.1).
|
||||
const versionKey = (rawTag) => {
|
||||
const m = String(rawTag).match(/^v?(\d+)\.(\d+)\.(\d+)/);
|
||||
if (!m) return null;
|
||||
const maintenance = String(rawTag).match(/(?:\.post|\.patch\.)(\d+)/i);
|
||||
return [Number(m[1]), Number(m[2]), Number(m[3]), maintenance ? Number(maintenance[1]) : 0];
|
||||
};
|
||||
const isAtLeast = (a, b) => {
|
||||
for (let i = 0; i < a.length; i++) {
|
||||
if (a[i] !== b[i]) return a[i] > b[i];
|
||||
}
|
||||
return true;
|
||||
};
|
||||
|
||||
const cosignSection = [
|
||||
`## Verify Docker Image Signature`,
|
||||
``,
|
||||
|
|
@ -90,6 +106,22 @@ jobs:
|
|||
].join('\n');
|
||||
|
||||
try {
|
||||
let makeLatest = "false";
|
||||
const newVersion = versionKey(tag);
|
||||
if (!isPrerelease && newVersion) {
|
||||
let latestVersion = null;
|
||||
try {
|
||||
const latest = await github.rest.repos.getLatestRelease({
|
||||
owner: context.repo.owner,
|
||||
repo: context.repo.repo,
|
||||
});
|
||||
latestVersion = versionKey(latest.data.tag_name);
|
||||
} catch (error) {
|
||||
if (error.status !== 404) throw error;
|
||||
}
|
||||
makeLatest = (!latestVersion || isAtLeast(newVersion, latestVersion)) ? "true" : "false";
|
||||
}
|
||||
|
||||
const response = await github.rest.repos.createRelease({
|
||||
draft: true,
|
||||
generate_release_notes: true,
|
||||
|
|
@ -108,6 +140,7 @@ jobs:
|
|||
release_id: response.data.id,
|
||||
body: updatedBody,
|
||||
draft: false,
|
||||
make_latest: makeLatest,
|
||||
});
|
||||
|
||||
} catch (error) {
|
||||
|
|
|
|||
2
.github/workflows/test-unit-misc.yml
vendored
2
.github/workflows/test-unit-misc.yml
vendored
|
|
@ -28,6 +28,8 @@ jobs:
|
|||
tests/test_litellm/completion_extras
|
||||
tests/test_litellm/containers
|
||||
tests/test_litellm/experimental_mcp_client
|
||||
tests/test_litellm/models
|
||||
tests/test_litellm/repositories
|
||||
tests/test_litellm/images
|
||||
tests/test_litellm/interactions
|
||||
tests/test_litellm/passthrough
|
||||
|
|
|
|||
25
.github/workflows/test-unit-proxy-db.yml
vendored
25
.github/workflows/test-unit-proxy-db.yml
vendored
|
|
@ -1,9 +1,10 @@
|
|||
name: "Unit Tests: Proxy DB Operations"
|
||||
|
||||
# Uses DATABASE_URL secret — only runs on trusted branches, not PRs.
|
||||
on:
|
||||
push:
|
||||
branches: [main, "litellm_**"]
|
||||
pull_request:
|
||||
branches:
|
||||
- main
|
||||
- litellm_internal_staging
|
||||
|
||||
permissions:
|
||||
contents: read
|
||||
|
|
@ -30,9 +31,6 @@ concurrency:
|
|||
# xdist balances its 188 parametrized cases across workers instead of
|
||||
# pinning the whole file to one worker (the default --dist=loadscope
|
||||
# behavior for single-file targets).
|
||||
# * test_db_schema_migration.py is isolated because one test in it
|
||||
# (test_aaaasschema_migration_check) takes ~170s — by itself it
|
||||
# determines the shard's wall-clock floor.
|
||||
jobs:
|
||||
# Fast guard — fails the workflow if a test_*.py file under
|
||||
# tests/proxy_unit_tests/ is not referenced by any matrix entry below.
|
||||
|
|
@ -166,18 +164,6 @@ jobs:
|
|||
dist: loadscope
|
||||
timeout: 15
|
||||
|
||||
# ---- db-and-spend: isolate the 170s schema-migration test ----
|
||||
# test_db_schema_migration.py has exactly one test, and that test
|
||||
# is mostly waiting on `prisma migrate deploy` / `prisma migrate
|
||||
# diff` subprocesses (~170s). It does no CPU-bound Python work
|
||||
# inside the test. Running with workers=0 (serial, no xdist)
|
||||
# skips the 4-worker cold-start cost we'd otherwise pay for a
|
||||
# single test, saving ~4 minutes of wall-clock.
|
||||
- test-group: schema-migration
|
||||
test-path: "tests/proxy_unit_tests/test_db_schema_migration.py"
|
||||
workers: 0
|
||||
dist: loadscope
|
||||
timeout: 15
|
||||
- test-group: db-and-spend
|
||||
test-path: >-
|
||||
tests/proxy_unit_tests/test_prisma_client_backoff_retry.py
|
||||
|
|
@ -232,12 +218,11 @@ jobs:
|
|||
workers: 4
|
||||
dist: loadscope
|
||||
timeout: 15
|
||||
uses: ./.github/workflows/_test-unit-services-base.yml
|
||||
uses: ./.github/workflows/_test-unit-base.yml
|
||||
with:
|
||||
test-path: ${{ matrix.test-path }}
|
||||
workers: ${{ matrix.workers }}
|
||||
reruns: 2
|
||||
timeout-minutes: ${{ matrix.timeout }}
|
||||
enable-postgres: true
|
||||
dist: ${{ matrix.dist }}
|
||||
artifact-name: proxy-db-${{ matrix.test-group }}
|
||||
|
|
|
|||
|
|
@ -1,34 +0,0 @@
|
|||
name: "Unit Tests: Proxy Management-Endpoint Behavior Pinning"
|
||||
|
||||
on:
|
||||
pull_request:
|
||||
branches:
|
||||
- main
|
||||
- litellm_internal_staging
|
||||
- litellm_oss_branch
|
||||
- "litellm_**"
|
||||
|
||||
permissions:
|
||||
contents: read
|
||||
id-token: write
|
||||
pull-requests: write
|
||||
|
||||
concurrency:
|
||||
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
|
||||
cancel-in-progress: true
|
||||
|
||||
jobs:
|
||||
proxy-mgmt-behavior:
|
||||
uses: ./.github/workflows/_test-unit-services-base.yml
|
||||
with:
|
||||
test-path: tests/proxy_behavior
|
||||
# workers=0 (no xdist): the world seed is a single shared Postgres
|
||||
# state — two xdist workers both call seed_world() and race on the
|
||||
# ``behavior-pin-budget`` row, producing UniqueViolation + cascading
|
||||
# missing-membership FK failures. The whole suite is ~7s sequentially,
|
||||
# so the cost of disabling parallelism here is negligible.
|
||||
workers: 0
|
||||
reruns: 0
|
||||
enable-postgres: true
|
||||
artifact-name: proxy-mgmt-behavior
|
||||
timeout-minutes: 15
|
||||
28
.github/workflows/test-unit-security.yml
vendored
28
.github/workflows/test-unit-security.yml
vendored
|
|
@ -1,28 +0,0 @@
|
|||
name: "Unit Tests: Security"
|
||||
|
||||
# Kept push-only (was previously required by DATABASE_URL secret scoping;
|
||||
# now the postgres credentials are ephemeral localhost values but the
|
||||
# push-trigger stays to match the proxy-db workflow cadence).
|
||||
on:
|
||||
push:
|
||||
branches: [main, "litellm_**"]
|
||||
|
||||
permissions:
|
||||
contents: read
|
||||
id-token: write
|
||||
pull-requests: write
|
||||
|
||||
concurrency:
|
||||
group: ${{ github.workflow }}-${{ github.ref }}
|
||||
cancel-in-progress: true
|
||||
|
||||
jobs:
|
||||
security:
|
||||
uses: ./.github/workflows/_test-unit-services-base.yml
|
||||
with:
|
||||
test-path: "tests/proxy_security_tests/"
|
||||
workers: 1
|
||||
reruns: 2
|
||||
timeout-minutes: 20
|
||||
enable-postgres: true
|
||||
artifact-name: security
|
||||
|
|
@ -240,6 +240,24 @@ graph LR
|
|||
7. `DBSpendUpdateWriter.update_database()` queues spend increments to Redis
|
||||
8. Background job `update_spend` flushes queued spend to PostgreSQL every 60s
|
||||
|
||||
### Data Access Layer (Models & Repositories)
|
||||
|
||||
Database entities and the operations on them live in two packages at the root of `litellm/` so both the gateway (`proxy/`) and the SDK can use them without importing proxy internals:
|
||||
|
||||
- `litellm/models/` holds the canonical Pydantic definitions for every persisted entity (`LiteLLM_VerificationToken`, `LiteLLM_TeamTable`, `LiteLLM_UserTable`, etc.). `proxy/_types.py` re-exports these for backwards compatibility, so existing imports keep working.
|
||||
- `litellm/repositories/` holds the data-access layer. `BaseRepository[T]` provides the generic CRUD (`find_by_id`, `find_many`, `create`, `update`, `delete`, `count`, `exists`); entity repositories such as `VerificationTokenRepository`, `TeamRepository`, and `UserRepository` add domain-specific queries and writes on top of it.
|
||||
|
||||
Conventions to follow when touching this layer:
|
||||
|
||||
| Concern | How it's handled |
|
||||
|---------|------------------|
|
||||
| JSON columns | Prisma `Json` columns are stored as JSON strings. Repositories `json.dumps()` on write and `json.loads()` on read (see `_to_model` and the `_build_*_data` helpers). |
|
||||
| Archive-then-delete | `delete_team` / `delete_token` copy the row into the `LiteLLM_Deleted*` table and delete the original inside a single `prisma_client.db.tx()` transaction. Archive payloads are built explicitly so only columns that exist on the archive table are written. |
|
||||
| Column vs. field names | Where a model field differs from its DB column (for example `org_id` maps to the `organization_id` column), the repository translates in both directions rather than relying on Pydantic to guess. |
|
||||
| Array mutations | Adds use Prisma's atomic `push` (`add_member`, `add_admin`, `add_models`) to avoid read-modify-write races. Removals fall back to read-modify-write because Prisma has no atomic array remove. |
|
||||
|
||||
To add a new entity, define the model under `litellm/models/`, re-export it from `proxy/_types.py` if existing code imports it from there, and add a repository under `litellm/repositories/` (subclass `BaseRepository` for plain CRUD, or add bespoke methods when the entity needs encryption, archiving, or atomic array updates). Mirror the tests in `tests/test_litellm/repositories/`.
|
||||
|
||||
---
|
||||
|
||||
## 2. SDK Request Flow
|
||||
|
|
|
|||
15
CLAUDE.md
15
CLAUDE.md
|
|
@ -42,7 +42,7 @@ When you must use real LLM models to, for example, write e2e tests, write a QA r
|
|||
|
||||
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. 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
|
||||
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
|
||||
|
||||
When working on a PR, keep the PR description in sync with new commits being made
|
||||
|
||||
|
|
@ -52,6 +52,19 @@ Do not put names of customers or customer company names in code, PRs, and issues
|
|||
|
||||
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
|
||||
|
||||
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
|
||||
- Never-nester: early returns over deep nesting
|
||||
- Don't throw; model failures as values (One function (e.g., raise_public) maps error union to existing public exception contracts via exhaustive match + assert_never)
|
||||
- No mutation; instead of mutable lists and dicts, prefer tuples, NamedTuples, frozen dataclasses, etc.
|
||||
- Use dependency injection
|
||||
- Fully typed; no `Any` or coarse types like dict[str, Any]. Every function parameter must be strongly typed
|
||||
- Use tagged unions + match
|
||||
- No monster files or god objects
|
||||
|
||||
Follow conventional commits for commit names and PR titles
|
||||
|
||||
## Think Before Coding
|
||||
|
||||
**Don't assume. Don't hide confusion. Surface tradeoffs**
|
||||
|
|
|
|||
|
|
@ -38,18 +38,25 @@ Before contributing code to LiteLLM, you must sign our [Contributor License Agre
|
|||
git clone https://github.com/YOUR_USERNAME/litellm.git
|
||||
cd litellm
|
||||
|
||||
# Create a new branch for your feature
|
||||
git checkout -b your-feature-branch
|
||||
# Create a new branch for your feature (see "Commit and Branch Conventions" below)
|
||||
git checkout -b feature/your-feature
|
||||
|
||||
# Install development dependencies
|
||||
make install-dev
|
||||
|
||||
# Install git hooks that enforce commit + branch conventions (one-time, opt-in)
|
||||
make install-hooks
|
||||
|
||||
# Verify your setup works
|
||||
make help
|
||||
```
|
||||
|
||||
That's it! Your local development environment is ready.
|
||||
|
||||
## Commit and Branch Conventions
|
||||
|
||||
Commits follow [Conventional Commits](https://www.conventionalcommits.org/en/v1.0.0/) and branches follow [Conventional Branches](https://conventional-branch.github.io/). Run `make install-hooks` once per clone to enable the local git hooks that enforce these — see the [contributor docs](https://docs.litellm.ai/docs/extras/contributing_code#commit-and-branch-conventions) for the full type list, examples, the protected-branch bypass list, and how to opt out.
|
||||
|
||||
### 2. Development Workflow
|
||||
|
||||
Here's the recommended workflow for making changes:
|
||||
|
|
@ -67,12 +74,12 @@ make lint
|
|||
# Run unit tests to ensure nothing is broken
|
||||
make test-unit
|
||||
|
||||
# Commit your changes
|
||||
# Commit your changes (must follow Conventional Commits — see above)
|
||||
git add .
|
||||
git commit -m "Your descriptive commit message"
|
||||
git commit -m "feat(scope): your descriptive commit message"
|
||||
|
||||
# Push and create a PR
|
||||
git push origin your-feature-branch
|
||||
# Push and create a PR (branch must follow Conventional Branches — see above)
|
||||
git push origin feature/your-feature
|
||||
```
|
||||
|
||||
## Adding Testing
|
||||
|
|
|
|||
26
Dockerfile
26
Dockerfile
|
|
@ -68,22 +68,24 @@ FROM $LITELLM_RUNTIME_IMAGE AS runtime
|
|||
|
||||
USER root
|
||||
|
||||
RUN apk add --no-cache bash openssl tzdata nodejs npm python3 libsndfile && \
|
||||
npm install -g npm@11.14.0 tar@7.5.11 glob@13.0.6 @isaacs/brace-expansion@5.0.1 brace-expansion@5.0.5 minimatch@10.2.4 diff@8.0.3 picomatch@4.0.4 && \
|
||||
GLOBAL="$(npm root -g)" && \
|
||||
for pkg in tar glob @isaacs/brace-expansion brace-expansion minimatch diff picomatch; do \
|
||||
name="${pkg##*/}"; \
|
||||
find "$GLOBAL/npm" -type d -name "$name" -path "*/node_modules/$pkg" | while read d; do \
|
||||
rm -rf "$d" && cp -rL "$GLOBAL/$pkg" "$d"; \
|
||||
done; \
|
||||
done && \
|
||||
npm cache clean --force && \
|
||||
{ apk del --no-cache npm 2>/dev/null || true; }
|
||||
# node (without npm) is required by the prisma CLI at runtime
|
||||
RUN apk add --no-cache bash openssl tzdata nodejs python3 libsndfile
|
||||
|
||||
WORKDIR /app
|
||||
ENV PATH="/app/.venv/bin:${PATH}"
|
||||
|
||||
COPY --from=builder /app /app
|
||||
# Copy only what runtime needs. The application is installed inside the venv;
|
||||
# the rest of the builder's /app is source and build metadata that must not
|
||||
# ship (manifest-scanning tools attribute everything in it to this image).
|
||||
# entrypoint.sh invokes litellm/proxy/prisma_migration.py by source path.
|
||||
COPY --from=builder /app/.venv /app/.venv
|
||||
COPY --from=builder /app/docker /app/docker
|
||||
COPY --from=builder /app/schema.prisma /app/schema.prisma
|
||||
COPY --from=builder /app/litellm/proxy/prisma_migration.py /app/litellm/proxy/prisma_migration.py
|
||||
# enterprise/ is imported by source path at runtime (proxy_cli puts the
|
||||
# working directory on sys.path; litellm/proxy/hooks resolves
|
||||
# enterprise.enterprise_hooks from it)
|
||||
COPY --from=builder /app/enterprise /app/enterprise
|
||||
# Prisma binaries live in $HOME/.cache (default prisma-python location),
|
||||
# which is /root/.cache here. Copy only the Prisma subdirs — copying the
|
||||
# whole /root/.cache drags in the uv build cache (~660 MB, includes a
|
||||
|
|
|
|||
8
Makefile
8
Makefile
|
|
@ -5,7 +5,7 @@
|
|||
test-unit-integrations test-unit-core-utils test-unit-other test-unit-root \
|
||||
test-proxy-unit-a test-proxy-unit-b test-integration test-unit-helm \
|
||||
info lint lint-dev format \
|
||||
install-dev install-proxy-dev install-test-deps \
|
||||
install-dev install-proxy-dev install-test-deps install-hooks \
|
||||
install-helm-unittest check-circular-imports check-import-safety
|
||||
|
||||
# Default target
|
||||
|
|
@ -17,6 +17,7 @@ help:
|
|||
@echo " make install-proxy-dev-ci - Install proxy dev dependencies (CI-compatible)"
|
||||
@echo " make install-test-deps - Install the full local test environment"
|
||||
@echo " make install-helm-unittest - Install helm unittest plugin"
|
||||
@echo " make install-hooks - Install git hooks (Conventional Commits + Branches)"
|
||||
@echo " make format - Apply Black code formatting"
|
||||
@echo " make format-check - Check Black code formatting (matches CI)"
|
||||
@echo " make lint - Run all linting (Ruff, MyPy, Black check, circular imports, import safety)"
|
||||
|
|
@ -68,6 +69,11 @@ install-test-deps: install-proxy-dev
|
|||
install-helm-unittest:
|
||||
helm plugin install https://github.com/helm-unittest/helm-unittest --version v0.4.4 || echo "ignore error if plugin exists"
|
||||
|
||||
# Install git hooks that enforce Conventional Commits and Conventional Branches.
|
||||
# Opt-in: not chained into install-dev.
|
||||
install-hooks:
|
||||
./scripts/install_git_hooks.sh
|
||||
|
||||
# Formatting
|
||||
format: install-dev
|
||||
cd litellm && $(UV_RUN) black . && cd ..
|
||||
|
|
|
|||
|
|
@ -407,7 +407,7 @@ Support for more providers. Missing a provider or LLM Platform, raise a [feature
|
|||
### Run in Developer Mode
|
||||
#### Services
|
||||
1. Setup .env file in root
|
||||
2. Run dependant services `docker-compose up db prometheus`
|
||||
2. Run dependent services `docker-compose up db prometheus`
|
||||
|
||||
#### Backend
|
||||
1. (In root) create virtual environment `python -m venv .venv`
|
||||
|
|
|
|||
99
db_scripts/partition_spend_logs.sql
Normal file
99
db_scripts/partition_spend_logs.sql
Normal file
|
|
@ -0,0 +1,99 @@
|
|||
-- Converts an existing LiteLLM_SpendLogs table into a native Postgres
|
||||
-- range-partitioned table keyed on "startTime".
|
||||
--
|
||||
-- Why: at high request volume, retention via DELETE leaves dead tuples that
|
||||
-- autovacuum cannot reclaim quickly enough, so the table keeps growing on disk
|
||||
-- (seen at 450GB+ after ~1 month). With partitioning, retention drops whole
|
||||
-- partitions, which is instant and returns disk to the OS immediately.
|
||||
--
|
||||
-- This is an opt-in, manual operation. The default LiteLLM schema is NOT
|
||||
-- partitioned, so existing installs are unaffected until you run this.
|
||||
--
|
||||
-- IMPORTANT
|
||||
-- * 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.
|
||||
-- * The partition key ("startTime") must be part of the primary key, so the
|
||||
-- PK becomes composite ("request_id", "startTime"). LiteLLM's write path uses
|
||||
-- INSERT ... ON CONFLICT DO NOTHING, which is compatible with this.
|
||||
-- * Choose a partition granularity ("day" is the recommended default for
|
||||
-- high-volume tables) and keep it consistent with SPEND_LOG_PARTITION_INTERVAL.
|
||||
--
|
||||
-- After running this, enable the feature and set a retention period in
|
||||
-- proxy_config.yaml:
|
||||
-- general_settings:
|
||||
-- use_spend_logs_partitioning: true
|
||||
-- maximum_spend_logs_retention_period: "30d"
|
||||
-- The spend-log cleanup job then verifies the table is partitioned and reclaims
|
||||
-- disk by dropping expired partitions instead of deleting rows. It also
|
||||
-- pre-creates upcoming partitions on each run. To roll back, see
|
||||
-- db_scripts/unpartition_spend_logs.sql.
|
||||
|
||||
BEGIN;
|
||||
|
||||
ALTER TABLE "LiteLLM_SpendLogs" RENAME TO "LiteLLM_SpendLogs_legacy";
|
||||
|
||||
-- Renaming a table does NOT rename its indexes, and index names are unique per
|
||||
-- schema. Move the legacy table's indexes aside so the CREATE INDEX statements
|
||||
-- below actually create indexes on the new partitioned table instead of being
|
||||
-- silently skipped by IF NOT EXISTS, and so the new PK keeps the canonical
|
||||
-- name instead of getting a "_pkey1" suffix.
|
||||
ALTER INDEX IF EXISTS "LiteLLM_SpendLogs_pkey"
|
||||
RENAME TO "LiteLLM_SpendLogs_legacy_pkey";
|
||||
ALTER INDEX IF EXISTS "LiteLLM_SpendLogs_startTime_idx"
|
||||
RENAME TO "LiteLLM_SpendLogs_legacy_startTime_idx";
|
||||
ALTER INDEX IF EXISTS "LiteLLM_SpendLogs_startTime_request_id_idx"
|
||||
RENAME TO "LiteLLM_SpendLogs_legacy_startTime_request_id_idx";
|
||||
ALTER INDEX IF EXISTS "LiteLLM_SpendLogs_end_user_idx"
|
||||
RENAME TO "LiteLLM_SpendLogs_legacy_end_user_idx";
|
||||
ALTER INDEX IF EXISTS "LiteLLM_SpendLogs_session_id_idx"
|
||||
RENAME TO "LiteLLM_SpendLogs_legacy_session_id_idx";
|
||||
|
||||
CREATE TABLE "LiteLLM_SpendLogs" (
|
||||
LIKE "LiteLLM_SpendLogs_legacy" INCLUDING DEFAULTS INCLUDING GENERATED
|
||||
) PARTITION BY RANGE ("startTime");
|
||||
|
||||
ALTER TABLE "LiteLLM_SpendLogs"
|
||||
ADD PRIMARY KEY ("request_id", "startTime");
|
||||
|
||||
-- Recreate every index Prisma defines on the table. LIKE ... INCLUDING DEFAULTS
|
||||
-- INCLUDING GENERATED copies columns and defaults but NOT indexes, so without
|
||||
-- these the admin-UI cost-reporting queries that filter by end_user/session_id
|
||||
-- fall back to sequential scans. On a partitioned parent these propagate to
|
||||
-- every current and future partition automatically.
|
||||
CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogs_startTime_idx"
|
||||
ON "LiteLLM_SpendLogs" ("startTime");
|
||||
|
||||
CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogs_startTime_request_id_idx"
|
||||
ON "LiteLLM_SpendLogs" ("startTime", "request_id");
|
||||
|
||||
CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogs_end_user_idx"
|
||||
ON "LiteLLM_SpendLogs" ("end_user");
|
||||
|
||||
CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogs_session_id_idx"
|
||||
ON "LiteLLM_SpendLogs" ("session_id");
|
||||
|
||||
-- Safety net: any row whose startTime has no explicit partition lands here so
|
||||
-- writes never fail. The cleanup job never drops the DEFAULT partition.
|
||||
CREATE TABLE IF NOT EXISTS "LiteLLM_SpendLogs_pdefault"
|
||||
PARTITION OF "LiteLLM_SpendLogs" DEFAULT;
|
||||
|
||||
COMMIT;
|
||||
|
||||
-- Backfill (optional). Rows route to the correct partition automatically.
|
||||
-- For large legacy tables, copy in time-bounded batches during a low-traffic
|
||||
-- window instead of one statement, or simply keep "LiteLLM_SpendLogs_legacy"
|
||||
-- read-only until its data ages past your retention, then DROP it.
|
||||
--
|
||||
-- Backfilled rows land in the DEFAULT partition until explicit partitions
|
||||
-- cover their dates. Postgres refuses to create a partition whose range
|
||||
-- overlaps rows already in DEFAULT, so the cleanup job may log a warning when
|
||||
-- pre-creating today's partition right after a backfill; it recovers on its
|
||||
-- own once those dates age out, and future partitions are unaffected because
|
||||
-- they are always created ahead of writes.
|
||||
--
|
||||
-- INSERT INTO "LiteLLM_SpendLogs"
|
||||
-- SELECT * FROM "LiteLLM_SpendLogs_legacy"
|
||||
-- WHERE "startTime" >= now() - interval '30 days';
|
||||
--
|
||||
-- DROP TABLE "LiteLLM_SpendLogs_legacy";
|
||||
69
db_scripts/unpartition_spend_logs.sql
Normal file
69
db_scripts/unpartition_spend_logs.sql
Normal file
|
|
@ -0,0 +1,69 @@
|
|||
-- Rolls back db_scripts/partition_spend_logs.sql: converts the native
|
||||
-- range-partitioned "LiteLLM_SpendLogs" table back into a plain,
|
||||
-- non-partitioned table matching the default LiteLLM schema.
|
||||
--
|
||||
-- When/why: run this if you want to stop using partition-based retention and
|
||||
-- return to DELETE-based cleanup, or to restore the original single-column
|
||||
-- primary key ("request_id") that the partitioned layout had to widen to a
|
||||
-- composite ("request_id", "startTime").
|
||||
--
|
||||
-- IMPORTANT
|
||||
-- * Test on a staging copy first and take a backup.
|
||||
-- * Postgres cannot convert a partitioned table back in place, so this
|
||||
-- renames the partitioned table aside and creates a fresh plain table.
|
||||
-- * The composite PK could in principle hold the same "request_id" in more
|
||||
-- than one partition, so rows are copied with ON CONFLICT DO NOTHING to
|
||||
-- restore the single-column PK without failing on such duplicates.
|
||||
-- * For large tables the INSERT ... SELECT copies every surviving row and may
|
||||
-- run long; do it during a low-traffic window.
|
||||
-- * Also remove use_spend_logs_partitioning from proxy_config.yaml (or set it
|
||||
-- to false) so the cleanup job returns to DELETE-based retention.
|
||||
|
||||
BEGIN;
|
||||
|
||||
ALTER TABLE "LiteLLM_SpendLogs" RENAME TO "LiteLLM_SpendLogs_partitioned";
|
||||
|
||||
-- Renaming a table does NOT rename its indexes, and index names are unique per
|
||||
-- schema. Move the partitioned table's indexes aside so the CREATE INDEX
|
||||
-- statements below actually create indexes on the new plain table instead of
|
||||
-- being silently skipped by IF NOT EXISTS, and so the new PK keeps the
|
||||
-- canonical name.
|
||||
ALTER INDEX IF EXISTS "LiteLLM_SpendLogs_pkey"
|
||||
RENAME TO "LiteLLM_SpendLogs_partitioned_pkey";
|
||||
ALTER INDEX IF EXISTS "LiteLLM_SpendLogs_pkey1"
|
||||
RENAME TO "LiteLLM_SpendLogs_partitioned_pkey1";
|
||||
ALTER INDEX IF EXISTS "LiteLLM_SpendLogs_startTime_idx"
|
||||
RENAME TO "LiteLLM_SpendLogs_partitioned_startTime_idx";
|
||||
ALTER INDEX IF EXISTS "LiteLLM_SpendLogs_startTime_request_id_idx"
|
||||
RENAME TO "LiteLLM_SpendLogs_partitioned_startTime_request_id_idx";
|
||||
ALTER INDEX IF EXISTS "LiteLLM_SpendLogs_end_user_idx"
|
||||
RENAME TO "LiteLLM_SpendLogs_partitioned_end_user_idx";
|
||||
ALTER INDEX IF EXISTS "LiteLLM_SpendLogs_session_id_idx"
|
||||
RENAME TO "LiteLLM_SpendLogs_partitioned_session_id_idx";
|
||||
|
||||
CREATE TABLE "LiteLLM_SpendLogs" (
|
||||
LIKE "LiteLLM_SpendLogs_partitioned" INCLUDING DEFAULTS INCLUDING GENERATED
|
||||
);
|
||||
|
||||
ALTER TABLE "LiteLLM_SpendLogs"
|
||||
ADD PRIMARY KEY ("request_id");
|
||||
|
||||
CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogs_startTime_idx"
|
||||
ON "LiteLLM_SpendLogs" ("startTime");
|
||||
|
||||
CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogs_startTime_request_id_idx"
|
||||
ON "LiteLLM_SpendLogs" ("startTime", "request_id");
|
||||
|
||||
CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogs_end_user_idx"
|
||||
ON "LiteLLM_SpendLogs" ("end_user");
|
||||
|
||||
CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogs_session_id_idx"
|
||||
ON "LiteLLM_SpendLogs" ("session_id");
|
||||
|
||||
INSERT INTO "LiteLLM_SpendLogs"
|
||||
SELECT * FROM "LiteLLM_SpendLogs_partitioned"
|
||||
ON CONFLICT ("request_id") DO NOTHING;
|
||||
|
||||
DROP TABLE "LiteLLM_SpendLogs_partitioned";
|
||||
|
||||
COMMIT;
|
||||
|
|
@ -66,36 +66,31 @@ FROM $LITELLM_RUNTIME_IMAGE AS runtime
|
|||
|
||||
USER root
|
||||
|
||||
RUN apk add --no-cache bash openssl tzdata nodejs npm python3 libsndfile && \
|
||||
npm install -g npm@11.12.1 tar@7.5.11 glob@11.1.0 @isaacs/brace-expansion@5.0.1 minimatch@10.2.4 diff@8.0.3 && \
|
||||
GLOBAL="$(npm root -g)" && \
|
||||
find "$GLOBAL/npm" -type d -name "tar" -path "*/node_modules/tar" | while read d; do \
|
||||
rm -rf "$d" && cp -rL "$GLOBAL/tar" "$d"; \
|
||||
done && \
|
||||
find "$GLOBAL/npm" -type d -name "glob" -path "*/node_modules/glob" | while read d; do \
|
||||
rm -rf "$d" && cp -rL "$GLOBAL/glob" "$d"; \
|
||||
done && \
|
||||
find "$GLOBAL/npm" -type d -name "brace-expansion" -path "*/node_modules/@isaacs/brace-expansion" | while read d; do \
|
||||
rm -rf "$d" && cp -rL "$GLOBAL/@isaacs/brace-expansion" "$d"; \
|
||||
done && \
|
||||
find "$GLOBAL/npm" -type d -name "minimatch" -path "*/node_modules/minimatch" | while read d; do \
|
||||
rm -rf "$d" && cp -rL "$GLOBAL/minimatch" "$d"; \
|
||||
done && \
|
||||
find "$GLOBAL/npm" -type d -name "diff" -path "*/node_modules/diff" | while read d; do \
|
||||
rm -rf "$d" && cp -rL "$GLOBAL/diff" "$d"; \
|
||||
done && \
|
||||
npm cache clean --force && \
|
||||
{ apk del --no-cache npm 2>/dev/null || true; }
|
||||
# node (without npm) is required by the prisma CLI at runtime
|
||||
RUN apk add --no-cache bash openssl tzdata nodejs python3 libsndfile
|
||||
|
||||
WORKDIR /app
|
||||
ENV PATH="/app/.venv/bin:${PATH}"
|
||||
|
||||
COPY --from=builder /app /app
|
||||
# Copy only what runtime needs. The application is installed inside the venv;
|
||||
# the rest of the builder's /app is source and build metadata that must not
|
||||
# ship (manifest-scanning tools attribute everything in it to this image).
|
||||
# entrypoint.sh invokes litellm/proxy/prisma_migration.py by source path.
|
||||
COPY --from=builder /app/.venv /app/.venv
|
||||
COPY --from=builder /app/docker /app/docker
|
||||
COPY --from=builder /app/schema.prisma /app/schema.prisma
|
||||
COPY --from=builder /app/litellm/proxy/prisma_migration.py /app/litellm/proxy/prisma_migration.py
|
||||
# enterprise/ is imported by source path at runtime (proxy_cli puts the
|
||||
# working directory on sys.path; litellm/proxy/hooks resolves
|
||||
# enterprise.enterprise_hooks from it)
|
||||
COPY --from=builder /app/enterprise /app/enterprise
|
||||
# Prisma binaries live in $HOME/.cache (default prisma-python location),
|
||||
# which is /root/.cache here. Copy them from the builder so they survive
|
||||
# deployments that volume-mount /app/.cache (e.g. readOnlyRootFilesystem
|
||||
# + emptyDir) — otherwise the mount would shadow the baked-in query engine.
|
||||
COPY --from=builder /root/.cache /root/.cache
|
||||
# Only the Prisma subdirs: the whole /root/.cache drags in the uv build cache.
|
||||
COPY --from=builder /root/.cache/prisma /root/.cache/prisma
|
||||
COPY --from=builder /root/.cache/prisma-python /root/.cache/prisma-python
|
||||
|
||||
RUN find /app/.venv -type f -path "*/tornado/test/*" -delete && \
|
||||
find /app/.venv -type d -path "*/tornado/test" -delete
|
||||
|
|
|
|||
|
|
@ -95,7 +95,21 @@ RUN for i in 1 2 3; do \
|
|||
apk add --no-cache python3 bash openssl tzdata libsndfile nodejs && break || sleep 5; \
|
||||
done
|
||||
|
||||
COPY --from=builder /app /app
|
||||
# Copy only what runtime needs. The application is installed inside the venv;
|
||||
# the rest of the builder's /app is source and build metadata that must not
|
||||
# ship (manifest-scanning tools attribute everything in it to this image).
|
||||
# entrypoint.sh invokes litellm/proxy/prisma_migration.py by source path.
|
||||
# Prisma caches live under /app/.cache here (XDG_CACHE_HOME /
|
||||
# PRISMA_BINARY_CACHE_DIR) so the runtime prisma generate finds them.
|
||||
COPY --from=builder /app/.venv /app/.venv
|
||||
COPY --from=builder /app/docker /app/docker
|
||||
COPY --from=builder /app/schema.prisma /app/schema.prisma
|
||||
COPY --from=builder /app/litellm/proxy/prisma_migration.py /app/litellm/proxy/prisma_migration.py
|
||||
# enterprise/ is imported by source path at runtime (proxy_cli puts the
|
||||
# working directory on sys.path; litellm/proxy/hooks resolves
|
||||
# enterprise.enterprise_hooks from it)
|
||||
COPY --from=builder /app/enterprise /app/enterprise
|
||||
COPY --from=builder /app/.cache /app/.cache
|
||||
COPY --from=builder /var/lib/litellm/ui /var/lib/litellm/ui
|
||||
COPY --from=builder /var/lib/litellm/assets /var/lib/litellm/assets
|
||||
|
||||
|
|
|
|||
|
|
@ -504,7 +504,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
if retrieve_file_id
|
||||
else False
|
||||
)
|
||||
if potential_file_id:
|
||||
if potential_file_id and "llm_output_file_id," in potential_file_id:
|
||||
model_id = self.get_model_id_from_unified_file_id(potential_file_id)
|
||||
if model_id:
|
||||
data["model"] = model_id
|
||||
|
|
@ -1058,7 +1058,12 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
return file_id.split("llm_output_file_model_id,")[1].split(";")[0]
|
||||
|
||||
def get_output_file_id_from_unified_file_id(self, file_id: str) -> str:
|
||||
return file_id.split("llm_output_file_id,")[1].split(";")[0]
|
||||
marker = "llm_output_file_id,"
|
||||
if marker not in file_id:
|
||||
raise ValueError(
|
||||
f"Unified id does not contain {marker!r}: {file_id[:80]!r}"
|
||||
)
|
||||
return file_id.split(marker, 1)[1].split(";")[0]
|
||||
|
||||
async def async_post_call_success_hook(
|
||||
self, data: Dict, user_api_key_dict: UserAPIKeyAuth, response: LLMResponseTypes
|
||||
|
|
@ -1099,13 +1104,33 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
for file_attr in ["output_file_id", "error_file_id"]:
|
||||
file_id_value = getattr(response, file_attr, None)
|
||||
if file_id_value and model_id:
|
||||
original_file_id = file_id_value
|
||||
unified_file_id = self.get_unified_output_file_id(
|
||||
output_file_id=original_file_id,
|
||||
model_id=model_id,
|
||||
model_name=resolved_model_name,
|
||||
decoded_output_file_id = _is_base64_encoded_unified_file_id(
|
||||
file_id_value
|
||||
)
|
||||
setattr(response, file_attr, unified_file_id)
|
||||
if (
|
||||
decoded_output_file_id
|
||||
and "llm_output_file_id," in decoded_output_file_id
|
||||
):
|
||||
provider_file_id = (
|
||||
self.get_output_file_id_from_unified_file_id(
|
||||
decoded_output_file_id
|
||||
)
|
||||
)
|
||||
unified_file_id = file_id_value
|
||||
elif decoded_output_file_id:
|
||||
verbose_logger.warning(
|
||||
f"Skipping {file_attr}={file_id_value!r}: "
|
||||
"unified id is not a managed file output id"
|
||||
)
|
||||
continue
|
||||
else:
|
||||
provider_file_id = file_id_value
|
||||
unified_file_id = self.get_unified_output_file_id(
|
||||
output_file_id=provider_file_id,
|
||||
model_id=model_id,
|
||||
model_name=resolved_model_name,
|
||||
)
|
||||
setattr(response, file_attr, unified_file_id)
|
||||
|
||||
# Use llm_router credentials when available. Without credentials,
|
||||
# Azure and other auth-required providers return 500/401.
|
||||
|
|
@ -1125,27 +1150,27 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
or {}
|
||||
)
|
||||
file_object = await litellm.afile_retrieve(
|
||||
file_id=original_file_id,
|
||||
file_id=provider_file_id,
|
||||
**_creds,
|
||||
)
|
||||
else:
|
||||
file_object = await litellm.afile_retrieve(
|
||||
custom_llm_provider=model_name.split("/")[0] if model_name and "/" in model_name else "openai", # type: ignore[arg-type]
|
||||
file_id=original_file_id,
|
||||
file_id=provider_file_id,
|
||||
)
|
||||
verbose_logger.debug(
|
||||
f"Successfully retrieved file object for {file_attr}={original_file_id}"
|
||||
f"Successfully retrieved file object for {file_attr}={provider_file_id}"
|
||||
)
|
||||
except Exception as e:
|
||||
verbose_logger.warning(
|
||||
f"Failed to retrieve file object for {file_attr}={original_file_id}: {str(e)}. Storing with None and will fetch on-demand."
|
||||
f"Failed to retrieve file object for {file_attr}={provider_file_id}: {str(e)}. Storing with None and will fetch on-demand."
|
||||
)
|
||||
|
||||
await self.store_unified_file_id(
|
||||
file_id=unified_file_id,
|
||||
file_object=file_object,
|
||||
litellm_parent_otel_span=user_api_key_dict.parent_otel_span,
|
||||
model_mappings={model_id: original_file_id},
|
||||
model_mappings={model_id: provider_file_id},
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
)
|
||||
await self.store_unified_object_id(
|
||||
|
|
|
|||
|
|
@ -12,9 +12,14 @@ spec:
|
|||
{{- include "litellm.backend.selectorLabels" . | nindent 6 }}
|
||||
template:
|
||||
metadata:
|
||||
{{- with .Values.backend.podAnnotations }}
|
||||
{{- if or .Values.gateway.config.create .Values.backend.podAnnotations }}
|
||||
annotations:
|
||||
{{- if .Values.gateway.config.create }}
|
||||
checksum/config: {{ include (print $.Template.BasePath "/gateway/configmap.yaml") . | sha256sum }}
|
||||
{{- end }}
|
||||
{{- with .Values.backend.podAnnotations }}
|
||||
{{- toYaml . | nindent 8 }}
|
||||
{{- end }}
|
||||
{{- end }}
|
||||
labels:
|
||||
{{- include "litellm.backend.selectorLabels" . | nindent 8 }}
|
||||
|
|
@ -35,7 +40,17 @@ spec:
|
|||
protocol: TCP
|
||||
env:
|
||||
{{- include "litellm.serverEnv" (dict "root" $ "component" .Values.backend) | nindent 12 }}
|
||||
{{- if .Values.gateway.config.create }}
|
||||
- name: CONFIG_FILE_PATH
|
||||
value: /app/config/config.yaml
|
||||
{{- end }}
|
||||
{{- include "litellm.envFrom" .Values.backend | nindent 10 }}
|
||||
{{- if .Values.gateway.config.create }}
|
||||
volumeMounts:
|
||||
- name: gateway-config
|
||||
mountPath: /app/config/config.yaml
|
||||
subPath: config.yaml
|
||||
{{- end }}
|
||||
{{- with .Values.backend.livenessProbe }}
|
||||
livenessProbe:
|
||||
{{- toYaml . | nindent 12 }}
|
||||
|
|
@ -46,6 +61,12 @@ spec:
|
|||
{{- end }}
|
||||
resources:
|
||||
{{- toYaml .Values.backend.resources | nindent 12 }}
|
||||
{{- if .Values.gateway.config.create }}
|
||||
volumes:
|
||||
- name: gateway-config
|
||||
configMap:
|
||||
name: {{ include "litellm.gateway.fullname" . }}-config
|
||||
{{- end }}
|
||||
{{- with .Values.backend.nodeSelector }}
|
||||
nodeSelector:
|
||||
{{- toYaml . | nindent 8 }}
|
||||
|
|
|
|||
|
|
@ -0,0 +1,23 @@
|
|||
-- AlterTable: add admin-configured env_vars to MCP server table
|
||||
ALTER TABLE "LiteLLM_MCPServerTable" ADD COLUMN IF NOT EXISTS "env_vars" JSONB DEFAULT '[]';
|
||||
|
||||
-- CreateTable: per-user env var values for MCP servers
|
||||
CREATE TABLE IF NOT EXISTS "LiteLLM_MCPUserEnvVars" (
|
||||
"id" TEXT NOT NULL,
|
||||
"user_id" TEXT NOT NULL,
|
||||
"server_id" TEXT NOT NULL,
|
||||
"values_b64" TEXT NOT NULL,
|
||||
"created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
||||
"updated_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
||||
|
||||
CONSTRAINT "LiteLLM_MCPUserEnvVars_pkey" PRIMARY KEY ("id")
|
||||
);
|
||||
|
||||
-- CreateIndex
|
||||
CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_MCPUserEnvVars_user_id_server_id_key" ON "LiteLLM_MCPUserEnvVars"("user_id", "server_id");
|
||||
|
||||
-- CreateIndex
|
||||
CREATE INDEX IF NOT EXISTS "LiteLLM_MCPUserEnvVars_user_id_idx" ON "LiteLLM_MCPUserEnvVars"("user_id");
|
||||
|
||||
-- CreateIndex
|
||||
CREATE INDEX IF NOT EXISTS "LiteLLM_MCPUserEnvVars_server_id_idx" ON "LiteLLM_MCPUserEnvVars"("server_id");
|
||||
|
|
@ -0,0 +1,2 @@
|
|||
-- AlterTable
|
||||
ALTER TABLE "LiteLLM_MCPServerTable" ADD COLUMN IF NOT EXISTS "oauth2_flow" TEXT;
|
||||
|
|
@ -0,0 +1,3 @@
|
|||
-- AlterTable
|
||||
ALTER TABLE "LiteLLM_MCPServerTable" ADD COLUMN "timeout" DOUBLE PRECISION;
|
||||
|
||||
|
|
@ -311,6 +311,11 @@ model LiteLLM_MCPServerTable {
|
|||
tool_name_to_description Json? @default("{}")
|
||||
extra_headers String[] @default([])
|
||||
static_headers Json? @default("{}")
|
||||
// Admin-configured environment variables interpolated into static_headers
|
||||
// via ${NAME} syntax. Stored as an array of
|
||||
// {name, value, scope, description}. scope is "global" (value used as-is)
|
||||
// or "user" (value supplied per-user via LiteLLM_MCPUserEnvVars).
|
||||
env_vars Json? @default("[]")
|
||||
// Health check status
|
||||
status String? @default("unknown")
|
||||
last_health_check DateTime?
|
||||
|
|
@ -322,6 +327,7 @@ model LiteLLM_MCPServerTable {
|
|||
authorization_url String?
|
||||
token_url String?
|
||||
registration_url String?
|
||||
oauth2_flow String?
|
||||
allow_all_keys Boolean @default(false)
|
||||
available_on_public_internet Boolean @default(true)
|
||||
delegate_auth_to_upstream Boolean @default(false)
|
||||
|
|
@ -330,6 +336,7 @@ model LiteLLM_MCPServerTable {
|
|||
byok_description String[] @default([])
|
||||
byok_api_key_help_url String?
|
||||
source_url String?
|
||||
timeout Float?
|
||||
// BYOM submission lifecycle
|
||||
approval_status String? @default("active")
|
||||
submitted_by String?
|
||||
|
|
@ -364,6 +371,21 @@ model LiteLLM_MCPUserCredentials {
|
|||
@@unique([user_id, server_id])
|
||||
}
|
||||
|
||||
// Per-user environment variable values for MCP servers.
|
||||
// values_b64 is an encrypted JSON object: {VAR_NAME: "value", ...}.
|
||||
model LiteLLM_MCPUserEnvVars {
|
||||
id String @id @default(uuid())
|
||||
user_id String
|
||||
server_id String
|
||||
values_b64 String
|
||||
created_at DateTime @default(now())
|
||||
updated_at DateTime @default(now()) @updatedAt
|
||||
|
||||
@@unique([user_id, server_id])
|
||||
@@index([user_id])
|
||||
@@index([server_id])
|
||||
}
|
||||
|
||||
// Generate Tokens for Proxy
|
||||
model LiteLLM_VerificationToken {
|
||||
token String @id
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
[project]
|
||||
name = "litellm-proxy-extras"
|
||||
version = "0.4.73"
|
||||
version = "0.4.74"
|
||||
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.73"
|
||||
version = "0.4.74"
|
||||
version_files = [
|
||||
"pyproject.toml:^version",
|
||||
"../pyproject.toml:litellm-proxy-extras==",
|
||||
|
|
|
|||
|
|
@ -16,8 +16,17 @@ import os
|
|||
# Load .env before any other litellm imports so env vars (e.g. LITELLM_UI_SESSION_DURATION) are available
|
||||
import dotenv as _dotenv
|
||||
|
||||
|
||||
def _dev_env_hot_reload_enabled() -> bool:
|
||||
"""The proxy exports this flag when started with ``--reload``. A reloaded
|
||||
worker is a fresh process that inherits the reloader's environment, so an
|
||||
edited ``.env`` value stays masked by the stale inherited one unless we
|
||||
let the file win; overriding makes the edit take effect on reload."""
|
||||
return os.getenv("LITELLM_DEV_ENV_HOT_RELOAD") == "True"
|
||||
|
||||
|
||||
if os.getenv("LITELLM_MODE", "DEV") == "DEV":
|
||||
_dotenv.load_dotenv()
|
||||
_dotenv.load_dotenv(override=_dev_env_hot_reload_enabled())
|
||||
|
||||
from typing import (
|
||||
Callable,
|
||||
|
|
@ -34,6 +43,7 @@ from typing import (
|
|||
Type,
|
||||
)
|
||||
from litellm.types.integrations.datadog import DatadogInitParams
|
||||
from litellm.types.integrations.newrelic import NewRelicInitParams
|
||||
from litellm._logging import (
|
||||
set_verbose,
|
||||
_turn_on_debug,
|
||||
|
|
@ -145,10 +155,12 @@ _custom_logger_compatible_callbacks_literal = Literal[
|
|||
"gitlab",
|
||||
"cloudzero",
|
||||
"focus",
|
||||
"mavvrik",
|
||||
"vantage",
|
||||
"posthog",
|
||||
"levo",
|
||||
"compression_interception",
|
||||
"newrelic",
|
||||
]
|
||||
cold_storage_custom_logger: Optional[_custom_logger_compatible_callbacks_literal] = None
|
||||
logged_real_time_event_types: Optional[Union[List[str], Literal["*"]]] = None
|
||||
|
|
@ -350,6 +362,9 @@ enable_gemini_default_thinking_level_low: bool = (
|
|||
####################
|
||||
logging: bool = True
|
||||
enable_loadbalancing_on_batch_endpoints: Optional[bool] = None
|
||||
require_managed_files: bool = (
|
||||
False # proxy only - require target_model_names on POST /v1/files
|
||||
)
|
||||
enable_caching_on_provider_specific_optional_params: bool = (
|
||||
False # feature-flag for caching on optional params - e.g. 'top_k'
|
||||
)
|
||||
|
|
@ -403,6 +418,7 @@ s3_callback_params: Optional[Dict] = None
|
|||
s3_audit_callback_params: Optional[Dict] = None
|
||||
datadog_llm_observability_params: Optional[Union[DatadogLLMObsInitParams, Dict]] = None
|
||||
datadog_params: Optional[Union[DatadogInitParams, Dict]] = None
|
||||
newrelic_params: Optional[Union[NewRelicInitParams, Dict]] = None
|
||||
aws_sqs_callback_params: Optional[Dict] = None
|
||||
generic_logger_headers: Optional[Dict] = None
|
||||
default_key_generate_params: Optional[Dict] = None
|
||||
|
|
@ -433,6 +449,13 @@ custom_prometheus_metadata_labels: List[str] = []
|
|||
custom_prometheus_tags: List[str] = []
|
||||
prometheus_metrics_config: Optional[List] = None
|
||||
prometheus_emit_stream_label: bool = False
|
||||
# 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
|
||||
# that metric keep matching after upgrade. Enable when downstream consumers
|
||||
# are ready to split 429s by source (vendor vs. litellm) and dimension
|
||||
# (RPM/TPM/concurrent/budget).
|
||||
prometheus_emit_rate_limit_labels: bool = False
|
||||
prometheus_user_budget_label_include_email_alias: bool = False
|
||||
prometheus_end_user_metrics_max_series_per_metric: Optional[int] = 10000
|
||||
prometheus_end_user_metrics_ttl_seconds: Optional[float] = 3600.0
|
||||
|
|
@ -612,6 +635,7 @@ cerebras_models: Set = set()
|
|||
galadriel_models: Set = set()
|
||||
nvidia_nim_models: Set = set()
|
||||
nvidia_riva_models: Set = set()
|
||||
soniox_models: Set = set()
|
||||
sambanova_models: Set = set()
|
||||
sambanova_embedding_models: Set = set()
|
||||
novita_models: Set = set()
|
||||
|
|
@ -844,6 +868,8 @@ def add_known_models(model_cost_map: Optional[Dict] = None):
|
|||
nvidia_nim_models.add(key)
|
||||
elif value.get("litellm_provider") == "nvidia_riva":
|
||||
nvidia_riva_models.add(key)
|
||||
elif value.get("litellm_provider") == "soniox":
|
||||
soniox_models.add(key)
|
||||
elif value.get("litellm_provider") == "sambanova":
|
||||
sambanova_models.add(key)
|
||||
elif value.get("litellm_provider") == "sambanova-embedding-models":
|
||||
|
|
@ -1009,6 +1035,7 @@ model_list = list(
|
|||
| galadriel_models
|
||||
| nvidia_nim_models
|
||||
| nvidia_riva_models
|
||||
| soniox_models
|
||||
| sambanova_models
|
||||
| azure_text_models
|
||||
| novita_models
|
||||
|
|
@ -1109,6 +1136,7 @@ models_by_provider: dict = {
|
|||
"galadriel": galadriel_models,
|
||||
"nvidia_nim": nvidia_nim_models,
|
||||
"nvidia_riva": nvidia_riva_models,
|
||||
"soniox": soniox_models,
|
||||
"sambanova": sambanova_models | sambanova_embedding_models,
|
||||
"novita": novita_models,
|
||||
"nebius": nebius_models | nebius_embedding_models,
|
||||
|
|
@ -1289,6 +1317,8 @@ from .exceptions import (
|
|||
NotFoundError,
|
||||
PermissionDeniedError,
|
||||
RateLimitError,
|
||||
RateLimitErrorCategory,
|
||||
RateLimitType,
|
||||
ServiceUnavailableError,
|
||||
BadGatewayError,
|
||||
OpenAIError,
|
||||
|
|
@ -1350,6 +1380,7 @@ from .search.main import *
|
|||
from .realtime_api.main import (
|
||||
_arealtime,
|
||||
acreate_realtime_client_secret,
|
||||
acreate_realtime_transcription_session,
|
||||
arealtime_calls,
|
||||
)
|
||||
from .responses.main import _aresponses_websocket
|
||||
|
|
@ -1740,6 +1771,9 @@ if TYPE_CHECKING:
|
|||
from .llms.openrouter.responses.transformation import (
|
||||
OpenRouterResponsesAPIConfig as OpenRouterResponsesAPIConfig,
|
||||
)
|
||||
from .llms.bedrock_mantle.responses.transformation import (
|
||||
BedrockMantleResponsesAPIConfig as BedrockMantleResponsesAPIConfig,
|
||||
)
|
||||
from .llms.gemini.interactions.transformation import (
|
||||
GoogleAIStudioInteractionsConfig as GoogleAIStudioInteractionsConfig,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -237,6 +237,7 @@ LLM_CONFIG_NAMES = (
|
|||
"PerplexityResponsesConfig",
|
||||
"DatabricksResponsesAPIConfig",
|
||||
"OpenRouterResponsesAPIConfig",
|
||||
"BedrockMantleResponsesAPIConfig",
|
||||
"GoogleAIStudioInteractionsConfig",
|
||||
"OpenAIOSeriesConfig",
|
||||
"AnthropicSkillsConfig",
|
||||
|
|
@ -320,6 +321,7 @@ LLM_CONFIG_NAMES = (
|
|||
"LemonadeChatConfig",
|
||||
"SnowflakeEmbeddingConfig",
|
||||
"AmazonNovaChatConfig",
|
||||
"SonioxAudioTranscriptionConfig",
|
||||
)
|
||||
|
||||
# Types that support lazy loading via _lazy_import_types
|
||||
|
|
@ -958,6 +960,10 @@ _LLM_CONFIGS_IMPORT_MAP = {
|
|||
".llms.openrouter.responses.transformation",
|
||||
"OpenRouterResponsesAPIConfig",
|
||||
),
|
||||
"BedrockMantleResponsesAPIConfig": (
|
||||
".llms.bedrock_mantle.responses.transformation",
|
||||
"BedrockMantleResponsesAPIConfig",
|
||||
),
|
||||
"GoogleAIStudioInteractionsConfig": (
|
||||
".llms.gemini.interactions.transformation",
|
||||
"GoogleAIStudioInteractionsConfig",
|
||||
|
|
@ -1190,6 +1196,10 @@ _LLM_CONFIGS_IMPORT_MAP = {
|
|||
".llms.amazon_nova.chat.transformation",
|
||||
"AmazonNovaChatConfig",
|
||||
),
|
||||
"SonioxAudioTranscriptionConfig": (
|
||||
".llms.soniox.audio_transcription.transformation",
|
||||
"SonioxAudioTranscriptionConfig",
|
||||
),
|
||||
}
|
||||
|
||||
# Import map for utils module lazy imports
|
||||
|
|
|
|||
|
|
@ -19,6 +19,7 @@ from litellm.a2a_protocol.litellm_completion_bridge.transformation import (
|
|||
A2AStreamingContext,
|
||||
)
|
||||
from litellm.a2a_protocol.providers.config_manager import A2AProviderConfigManager
|
||||
from litellm.interactions.agents.utils import merge_agent_headers
|
||||
|
||||
# litellm_params key carrying the authenticated principal (hashed virtual key) so
|
||||
# A2A provider configs can scope provider-side state (e.g. LangFlow session memory)
|
||||
|
|
@ -48,6 +49,7 @@ class A2ACompletionBridgeHandler:
|
|||
params: Dict[str, Any],
|
||||
litellm_params: Dict[str, Any],
|
||||
api_base: Optional[str] = None,
|
||||
agent_extra_headers: Optional[Dict[str, str]] = None,
|
||||
*,
|
||||
_skip_a2a_provider_routing: bool = False,
|
||||
) -> Dict[str, Any]:
|
||||
|
|
@ -59,6 +61,8 @@ class A2ACompletionBridgeHandler:
|
|||
params: A2A MessageSendParams containing the message
|
||||
litellm_params: Agent's litellm_params (custom_llm_provider, model, etc.)
|
||||
api_base: API base URL from agent_card_params
|
||||
agent_extra_headers: Per-request headers (from x-a2a-{agent}-* rewrite and
|
||||
admin extra_headers) to forward on the upstream HTTP call.
|
||||
|
||||
Returns:
|
||||
A2A SendMessageResponse dict
|
||||
|
|
@ -80,6 +84,7 @@ class A2ACompletionBridgeHandler:
|
|||
params=params,
|
||||
api_base=api_base,
|
||||
litellm_params=litellm_params,
|
||||
agent_extra_headers=agent_extra_headers,
|
||||
)
|
||||
|
||||
# Extract message from params
|
||||
|
|
@ -106,7 +111,7 @@ class A2ACompletionBridgeHandler:
|
|||
)
|
||||
|
||||
# Build completion params dict
|
||||
completion_params = {
|
||||
completion_params: Dict[str, Any] = {
|
||||
"model": full_model,
|
||||
"messages": openai_messages,
|
||||
"api_base": api_base,
|
||||
|
|
@ -128,6 +133,12 @@ class A2ACompletionBridgeHandler:
|
|||
params=params,
|
||||
)
|
||||
|
||||
if agent_extra_headers:
|
||||
completion_params["extra_headers"] = merge_agent_headers(
|
||||
dynamic_headers=agent_extra_headers,
|
||||
static_headers=completion_params.get("extra_headers"),
|
||||
)
|
||||
|
||||
# Call litellm.acompletion
|
||||
response = await litellm.acompletion(**completion_params)
|
||||
|
||||
|
|
@ -149,6 +160,7 @@ class A2ACompletionBridgeHandler:
|
|||
params: Dict[str, Any],
|
||||
litellm_params: Dict[str, Any],
|
||||
api_base: Optional[str] = None,
|
||||
agent_extra_headers: Optional[Dict[str, str]] = None,
|
||||
*,
|
||||
_skip_a2a_provider_routing: bool = False,
|
||||
) -> AsyncIterator[Dict[str, Any]]:
|
||||
|
|
@ -166,6 +178,8 @@ class A2ACompletionBridgeHandler:
|
|||
params: A2A MessageSendParams containing the message
|
||||
litellm_params: Agent's litellm_params (custom_llm_provider, model, etc.)
|
||||
api_base: API base URL from agent_card_params
|
||||
agent_extra_headers: Per-request headers (from x-a2a-{agent}-* rewrite and
|
||||
admin extra_headers) to forward on the upstream HTTP call.
|
||||
|
||||
Yields:
|
||||
A2A streaming response events
|
||||
|
|
@ -187,6 +201,7 @@ class A2ACompletionBridgeHandler:
|
|||
params=params,
|
||||
api_base=api_base,
|
||||
litellm_params=litellm_params,
|
||||
agent_extra_headers=agent_extra_headers,
|
||||
):
|
||||
yield chunk
|
||||
|
||||
|
|
@ -222,7 +237,7 @@ class A2ACompletionBridgeHandler:
|
|||
)
|
||||
|
||||
# Build completion params dict
|
||||
completion_params = {
|
||||
completion_params: Dict[str, Any] = {
|
||||
"model": full_model,
|
||||
"messages": openai_messages,
|
||||
"api_base": api_base,
|
||||
|
|
@ -244,6 +259,12 @@ class A2ACompletionBridgeHandler:
|
|||
params=params,
|
||||
)
|
||||
|
||||
if agent_extra_headers:
|
||||
completion_params["extra_headers"] = merge_agent_headers(
|
||||
dynamic_headers=agent_extra_headers,
|
||||
static_headers=completion_params.get("extra_headers"),
|
||||
)
|
||||
|
||||
# 1. Emit initial task event (kind: "task", status: "submitted")
|
||||
task_event = A2ACompletionBridgeTransformation.create_task_event(ctx)
|
||||
yield task_event
|
||||
|
|
@ -305,6 +326,7 @@ async def handle_a2a_completion(
|
|||
params: Dict[str, Any],
|
||||
litellm_params: Dict[str, Any],
|
||||
api_base: Optional[str] = None,
|
||||
agent_extra_headers: Optional[Dict[str, str]] = None,
|
||||
) -> Dict[str, Any]:
|
||||
"""Convenience function for non-streaming A2A completion."""
|
||||
return await A2ACompletionBridgeHandler.handle_non_streaming(
|
||||
|
|
@ -312,6 +334,7 @@ async def handle_a2a_completion(
|
|||
params=params,
|
||||
litellm_params=litellm_params,
|
||||
api_base=api_base,
|
||||
agent_extra_headers=agent_extra_headers,
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -320,6 +343,7 @@ async def handle_a2a_completion_streaming(
|
|||
params: Dict[str, Any],
|
||||
litellm_params: Dict[str, Any],
|
||||
api_base: Optional[str] = None,
|
||||
agent_extra_headers: Optional[Dict[str, str]] = None,
|
||||
) -> AsyncIterator[Dict[str, Any]]:
|
||||
"""Convenience function for streaming A2A completion."""
|
||||
async for chunk in A2ACompletionBridgeHandler.handle_streaming(
|
||||
|
|
@ -327,5 +351,6 @@ async def handle_a2a_completion_streaming(
|
|||
params=params,
|
||||
litellm_params=litellm_params,
|
||||
api_base=api_base,
|
||||
agent_extra_headers=agent_extra_headers,
|
||||
):
|
||||
yield chunk
|
||||
|
|
|
|||
|
|
@ -132,6 +132,7 @@ async def _send_message_via_completion_bridge(
|
|||
custom_llm_provider: str,
|
||||
api_base: Optional[str],
|
||||
litellm_params: Dict[str, Any],
|
||||
agent_extra_headers: Optional[Dict[str, str]] = None,
|
||||
) -> LiteLLMSendMessageResponse:
|
||||
"""
|
||||
Route a send_message through the LiteLLM completion bridge (e.g. LangGraph, Bedrock AgentCore).
|
||||
|
|
@ -157,6 +158,7 @@ async def _send_message_via_completion_bridge(
|
|||
params=params,
|
||||
litellm_params=litellm_params,
|
||||
api_base=api_base,
|
||||
agent_extra_headers=agent_extra_headers,
|
||||
)
|
||||
|
||||
return LiteLLMSendMessageResponse.from_dict(
|
||||
|
|
@ -283,6 +285,7 @@ async def asend_message(
|
|||
custom_llm_provider=custom_llm_provider,
|
||||
api_base=api_base,
|
||||
litellm_params=litellm_params,
|
||||
agent_extra_headers=agent_extra_headers,
|
||||
)
|
||||
|
||||
# Standard A2A client flow
|
||||
|
|
@ -509,6 +512,7 @@ async def asend_message_streaming( # noqa: PLR0915
|
|||
params=params,
|
||||
litellm_params=litellm_params,
|
||||
api_base=api_base,
|
||||
agent_extra_headers=agent_extra_headers,
|
||||
):
|
||||
yield chunk
|
||||
return
|
||||
|
|
|
|||
|
|
@ -37,6 +37,7 @@ class BedrockAgentCoreA2AConfig(BaseA2AProviderConfig):
|
|||
request_id=request_id,
|
||||
params=params,
|
||||
litellm_params=litellm_params,
|
||||
agent_extra_headers=kwargs.get("agent_extra_headers"),
|
||||
)
|
||||
|
||||
async def handle_streaming(
|
||||
|
|
@ -57,5 +58,6 @@ class BedrockAgentCoreA2AConfig(BaseA2AProviderConfig):
|
|||
request_id=request_id,
|
||||
params=params,
|
||||
litellm_params=litellm_params,
|
||||
agent_extra_headers=kwargs.get("agent_extra_headers"),
|
||||
):
|
||||
yield chunk
|
||||
|
|
|
|||
|
|
@ -6,7 +6,7 @@ completion bridge that would otherwise strip the envelope.
|
|||
"""
|
||||
|
||||
import json
|
||||
from typing import Any, AsyncIterator, Dict, cast
|
||||
from typing import Any, AsyncIterator, Dict, Optional, cast
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.a2a_protocol.providers.bedrock_agentcore.transformation import (
|
||||
|
|
@ -29,6 +29,7 @@ class BedrockAgentCoreA2AHandler:
|
|||
request_id: str,
|
||||
params: Dict[str, Any],
|
||||
litellm_params: Dict[str, Any],
|
||||
agent_extra_headers: Optional[Dict[str, str]] = None,
|
||||
) -> Dict[str, Any]:
|
||||
"""
|
||||
Handle non-streaming A2A request to AgentCore.
|
||||
|
|
@ -37,6 +38,8 @@ class BedrockAgentCoreA2AHandler:
|
|||
request_id: A2A JSON-RPC request ID
|
||||
params: A2A MessageSendParams containing the message
|
||||
litellm_params: Agent's litellm_params (model, api_key, etc.)
|
||||
agent_extra_headers: Per-request headers (from x-a2a-{agent}-* rewrite and
|
||||
admin extra_headers) to forward on the upstream HTTP call.
|
||||
|
||||
Returns:
|
||||
A2A JSON-RPC response dict from the AgentCore agent
|
||||
|
|
@ -47,6 +50,7 @@ class BedrockAgentCoreA2AHandler:
|
|||
params=params,
|
||||
litellm_params=litellm_params,
|
||||
method="message/send",
|
||||
agent_extra_headers=agent_extra_headers,
|
||||
)
|
||||
)
|
||||
|
||||
|
|
@ -77,6 +81,7 @@ class BedrockAgentCoreA2AHandler:
|
|||
request_id: str,
|
||||
params: Dict[str, Any],
|
||||
litellm_params: Dict[str, Any],
|
||||
agent_extra_headers: Optional[Dict[str, str]] = None,
|
||||
) -> AsyncIterator[Dict[str, Any]]:
|
||||
"""
|
||||
Handle streaming A2A request to AgentCore.
|
||||
|
|
@ -85,6 +90,8 @@ class BedrockAgentCoreA2AHandler:
|
|||
request_id: A2A JSON-RPC request ID
|
||||
params: A2A MessageSendParams containing the message
|
||||
litellm_params: Agent's litellm_params (model, api_key, etc.)
|
||||
agent_extra_headers: Per-request headers (from x-a2a-{agent}-* rewrite and
|
||||
admin extra_headers) to forward on the upstream HTTP call.
|
||||
|
||||
Yields:
|
||||
A2A streaming response events from the AgentCore agent
|
||||
|
|
@ -96,6 +103,7 @@ class BedrockAgentCoreA2AHandler:
|
|||
litellm_params=litellm_params,
|
||||
method="message/send",
|
||||
stream=True,
|
||||
agent_extra_headers=agent_extra_headers,
|
||||
)
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -6,11 +6,66 @@ and signs requests via AmazonAgentCoreConfig (SigV4 or JWT).
|
|||
"""
|
||||
|
||||
import json
|
||||
from typing import Any, AsyncIterator, Dict, Tuple
|
||||
from typing import Any, AsyncIterator, Dict, Mapping, Optional, Tuple
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.llms.bedrock.chat.agentcore.transformation import AmazonAgentCoreConfig
|
||||
|
||||
# Reserved outbound header names that must never be sourced from per-request
|
||||
# ``agent_extra_headers`` for AgentCore requests. ``agent_extra_headers`` carries
|
||||
# values rewritten from the client-controlled ``x-a2a-{agent}-*`` convention, so
|
||||
# allowing these would let any caller with access to the agent spoof the AWS
|
||||
# request identity / SigV4 metadata by overwriting headers the proxy sets from
|
||||
# trusted server-side config.
|
||||
#
|
||||
# The runtime headers (session / user id) are derived server-side from
|
||||
# ``runtimeSessionId`` / ``runtimeUserId`` in the agent's ``litellm_params``;
|
||||
# ``authorization`` is set by the AgentCore signer (JWT or SigV4); ``host`` and
|
||||
# the ``x-amz-*`` family are owned by SigV4 itself.
|
||||
_RESERVED_EXACT_HEADERS = frozenset(
|
||||
{
|
||||
"authorization",
|
||||
"host",
|
||||
}
|
||||
)
|
||||
_RESERVED_PREFIX_HEADERS: Tuple[str, ...] = (
|
||||
"x-amzn-bedrock-agentcore-runtime-",
|
||||
"x-amz-",
|
||||
)
|
||||
|
||||
|
||||
def _filter_reserved_headers(
|
||||
agent_extra_headers: Optional[Mapping[str, str]],
|
||||
) -> Optional[Dict[str, str]]:
|
||||
"""
|
||||
Strip reserved AWS / AgentCore headers from caller-supplied
|
||||
``agent_extra_headers`` before they are merged into the signed request.
|
||||
|
||||
Returns ``None`` if the result is empty.
|
||||
"""
|
||||
if not agent_extra_headers:
|
||||
return None
|
||||
|
||||
filtered: Dict[str, str] = {}
|
||||
dropped: list = []
|
||||
for k, v in agent_extra_headers.items():
|
||||
k_lower = k.lower()
|
||||
if k_lower in _RESERVED_EXACT_HEADERS or any(
|
||||
k_lower.startswith(prefix) for prefix in _RESERVED_PREFIX_HEADERS
|
||||
):
|
||||
dropped.append(k)
|
||||
continue
|
||||
filtered[k] = v
|
||||
|
||||
if dropped:
|
||||
verbose_logger.warning(
|
||||
"BedrockAgentCore A2A: dropping reserved header(s) from "
|
||||
"agent_extra_headers (not forwarded to AgentCore): %s",
|
||||
sorted(dropped),
|
||||
)
|
||||
|
||||
return filtered or None
|
||||
|
||||
|
||||
class BedrockAgentCoreA2ATransformation:
|
||||
"""
|
||||
|
|
@ -27,6 +82,7 @@ class BedrockAgentCoreA2ATransformation:
|
|||
litellm_params: Dict[str, Any],
|
||||
method: str = "message/send",
|
||||
stream: bool = False,
|
||||
agent_extra_headers: Optional[Dict[str, str]] = None,
|
||||
) -> Tuple[str, dict, bytes]:
|
||||
"""
|
||||
Build the AgentCore URL, construct a JSON-RPC envelope, and sign the request.
|
||||
|
|
@ -37,6 +93,15 @@ class BedrockAgentCoreA2ATransformation:
|
|||
litellm_params: Agent's litellm_params (model, api_key, etc.)
|
||||
method: JSON-RPC method name (default: "message/send")
|
||||
stream: Whether this is a streaming request
|
||||
agent_extra_headers: Per-request headers (from x-a2a-{agent}-* rewrite and
|
||||
admin extra_headers) to forward on the upstream HTTP call. Merged into
|
||||
the headers dict before signing so SigV4 includes them in the signature.
|
||||
Reserved AWS / AgentCore identity headers (``authorization``, ``host``,
|
||||
``x-amzn-bedrock-agentcore-runtime-*``, ``x-amz-*``) are filtered out
|
||||
here to prevent a caller-controlled ``x-a2a-{agent}-*`` header from
|
||||
spoofing the AgentCore runtime user id or other SigV4 metadata. Use
|
||||
``api_key`` / ``runtimeUserId`` / ``runtimeSessionId`` in litellm_params
|
||||
(not ``agent_extra_headers``) to override those values.
|
||||
|
||||
Returns:
|
||||
Tuple of (url, signed_headers, signed_body_bytes)
|
||||
|
|
@ -85,6 +150,13 @@ class BedrockAgentCoreA2ATransformation:
|
|||
if runtime_user_id:
|
||||
headers["X-Amzn-Bedrock-AgentCore-Runtime-User-Id"] = runtime_user_id
|
||||
|
||||
# Merge per-request agent headers before signing so SigV4 covers them.
|
||||
# Reserved headers are stripped first to prevent client-controlled values
|
||||
# from spoofing the AgentCore runtime identity / SigV4 metadata.
|
||||
safe_extra_headers = _filter_reserved_headers(agent_extra_headers)
|
||||
if safe_extra_headers:
|
||||
headers.update(safe_extra_headers)
|
||||
|
||||
# Sign the request (SigV4 or JWT depending on api_key presence)
|
||||
signed_headers, signed_body = agentcore_config.sign_request(
|
||||
headers=headers,
|
||||
|
|
|
|||
|
|
@ -31,6 +31,7 @@ class PydanticAIProviderConfig(BaseA2AProviderConfig):
|
|||
params=params,
|
||||
api_base=api_base,
|
||||
timeout=kwargs.get("timeout", 60.0),
|
||||
agent_extra_headers=kwargs.get("agent_extra_headers"),
|
||||
)
|
||||
|
||||
async def handle_streaming(
|
||||
|
|
@ -50,5 +51,6 @@ class PydanticAIProviderConfig(BaseA2AProviderConfig):
|
|||
timeout=kwargs.get("timeout", 60.0),
|
||||
chunk_size=kwargs.get("chunk_size", 50),
|
||||
delay_ms=kwargs.get("delay_ms", 10),
|
||||
agent_extra_headers=kwargs.get("agent_extra_headers"),
|
||||
):
|
||||
yield chunk
|
||||
|
|
|
|||
|
|
@ -28,6 +28,7 @@ class PydanticAIHandler:
|
|||
params: Dict[str, Any],
|
||||
api_base: Optional[str] = None,
|
||||
timeout: float = 60.0,
|
||||
agent_extra_headers: Optional[Dict[str, str]] = None,
|
||||
) -> Dict[str, Any]:
|
||||
"""
|
||||
Handle non-streaming request to Pydantic AI agent.
|
||||
|
|
@ -37,6 +38,8 @@ class PydanticAIHandler:
|
|||
params: A2A MessageSendParams containing the message
|
||||
api_base: Base URL of the Pydantic AI agent
|
||||
timeout: Request timeout in seconds
|
||||
agent_extra_headers: Per-request headers (from x-a2a-{agent}-* rewrite and
|
||||
admin extra_headers) to forward on the upstream HTTP call.
|
||||
|
||||
Returns:
|
||||
A2A SendMessageResponse dict
|
||||
|
|
@ -51,6 +54,7 @@ class PydanticAIHandler:
|
|||
request_id=request_id,
|
||||
params=params,
|
||||
timeout=timeout,
|
||||
agent_extra_headers=agent_extra_headers,
|
||||
)
|
||||
|
||||
return response_data
|
||||
|
|
@ -63,6 +67,7 @@ class PydanticAIHandler:
|
|||
timeout: float = 60.0,
|
||||
chunk_size: int = 50,
|
||||
delay_ms: int = 10,
|
||||
agent_extra_headers: Optional[Dict[str, str]] = None,
|
||||
) -> AsyncIterator[Dict[str, Any]]:
|
||||
"""
|
||||
Handle streaming request to Pydantic AI agent with fake streaming.
|
||||
|
|
@ -78,6 +83,8 @@ class PydanticAIHandler:
|
|||
timeout: Request timeout in seconds
|
||||
chunk_size: Number of characters per chunk
|
||||
delay_ms: Delay between chunks in milliseconds
|
||||
agent_extra_headers: Per-request headers (from x-a2a-{agent}-* rewrite and
|
||||
admin extra_headers) to forward on the upstream HTTP call.
|
||||
|
||||
Yields:
|
||||
A2A streaming response events
|
||||
|
|
@ -94,6 +101,7 @@ class PydanticAIHandler:
|
|||
request_id=request_id,
|
||||
params=params,
|
||||
timeout=timeout,
|
||||
agent_extra_headers=agent_extra_headers,
|
||||
)
|
||||
|
||||
# Convert raw task response to fake streaming chunks
|
||||
|
|
|
|||
|
|
@ -6,7 +6,7 @@ This module provides fake streaming by converting non-streaming responses into s
|
|||
"""
|
||||
|
||||
import asyncio
|
||||
from typing import Any, AsyncIterator, Dict, cast
|
||||
from typing import Any, AsyncIterator, Dict, Optional, cast
|
||||
from uuid import uuid4
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
|
|
@ -86,6 +86,7 @@ class PydanticAITransformation:
|
|||
request_id: str,
|
||||
max_attempts: int = 30,
|
||||
poll_interval: float = 0.5,
|
||||
agent_extra_headers: Optional[Dict[str, str]] = None,
|
||||
) -> Dict[str, Any]:
|
||||
"""
|
||||
Poll for task completion using tasks/get method.
|
||||
|
|
@ -112,7 +113,10 @@ class PydanticAITransformation:
|
|||
response = await client.post(
|
||||
endpoint,
|
||||
json=poll_request,
|
||||
headers={"Content-Type": "application/json"},
|
||||
headers={
|
||||
**(agent_extra_headers or {}),
|
||||
"Content-Type": "application/json",
|
||||
},
|
||||
)
|
||||
response.raise_for_status()
|
||||
poll_data = response.json()
|
||||
|
|
@ -142,6 +146,7 @@ class PydanticAITransformation:
|
|||
request_id: str,
|
||||
params: Any,
|
||||
timeout: float = 60.0,
|
||||
agent_extra_headers: Optional[Dict[str, str]] = None,
|
||||
) -> Dict[str, Any]:
|
||||
"""
|
||||
Send a request to Pydantic AI agent and return the raw task response.
|
||||
|
|
@ -189,7 +194,10 @@ class PydanticAITransformation:
|
|||
response = await client.post(
|
||||
endpoint,
|
||||
json=a2a_request,
|
||||
headers={"Content-Type": "application/json"},
|
||||
headers={
|
||||
**(agent_extra_headers or {}),
|
||||
"Content-Type": "application/json",
|
||||
},
|
||||
)
|
||||
response.raise_for_status()
|
||||
response_data = response.json()
|
||||
|
|
@ -211,6 +219,7 @@ class PydanticAITransformation:
|
|||
endpoint=endpoint,
|
||||
task_id=task_id,
|
||||
request_id=request_id,
|
||||
agent_extra_headers=agent_extra_headers,
|
||||
)
|
||||
|
||||
verbose_logger.info(
|
||||
|
|
@ -225,6 +234,7 @@ class PydanticAITransformation:
|
|||
request_id: str,
|
||||
params: Any,
|
||||
timeout: float = 60.0,
|
||||
agent_extra_headers: Optional[Dict[str, str]] = None,
|
||||
) -> Dict[str, Any]:
|
||||
"""
|
||||
Send a non-streaming A2A request to Pydantic AI agent and wait for completion.
|
||||
|
|
@ -234,6 +244,7 @@ class PydanticAITransformation:
|
|||
request_id: A2A JSON-RPC request ID
|
||||
params: A2A MessageSendParams containing the message (dict or Pydantic model)
|
||||
timeout: Request timeout in seconds
|
||||
agent_extra_headers: Per-request headers to forward on the upstream HTTP call.
|
||||
|
||||
Returns:
|
||||
Standard A2A non-streaming response format with message
|
||||
|
|
@ -244,6 +255,7 @@ class PydanticAITransformation:
|
|||
request_id=request_id,
|
||||
params=params,
|
||||
timeout=timeout,
|
||||
agent_extra_headers=agent_extra_headers,
|
||||
)
|
||||
|
||||
# Transform to standard A2A non-streaming format
|
||||
|
|
@ -258,6 +270,7 @@ class PydanticAITransformation:
|
|||
request_id: str,
|
||||
params: Any,
|
||||
timeout: float = 60.0,
|
||||
agent_extra_headers: Optional[Dict[str, str]] = None,
|
||||
) -> Dict[str, Any]:
|
||||
"""
|
||||
Send a request to Pydantic AI agent and return the raw task response.
|
||||
|
|
@ -269,6 +282,7 @@ class PydanticAITransformation:
|
|||
request_id: A2A JSON-RPC request ID
|
||||
params: A2A MessageSendParams containing the message
|
||||
timeout: Request timeout in seconds
|
||||
agent_extra_headers: Per-request headers to forward on the upstream HTTP call.
|
||||
|
||||
Returns:
|
||||
Raw Pydantic AI task response (with history/artifacts)
|
||||
|
|
@ -278,6 +292,7 @@ class PydanticAITransformation:
|
|||
request_id=request_id,
|
||||
params=params,
|
||||
timeout=timeout,
|
||||
agent_extra_headers=agent_extra_headers,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
|
|
|
|||
|
|
@ -129,7 +129,7 @@
|
|||
"bash_20241022": null,
|
||||
"bash_20250124": null,
|
||||
"code-execution-2025-08-25": null,
|
||||
"compact-2026-01-12": null,
|
||||
"compact-2026-01-12": "compact-2026-01-12",
|
||||
"computer-use-2025-01-24": "computer-use-2025-01-24",
|
||||
"computer-use-2025-11-24": "computer-use-2025-11-24",
|
||||
"context-1m-2025-08-07": "context-1m-2025-08-07",
|
||||
|
|
|
|||
|
|
@ -309,9 +309,13 @@ class Cache:
|
|||
param_value = kwargs[param]
|
||||
cache_key += f"{str(param)}: {str(param_value)}"
|
||||
|
||||
verbose_logger.debug("\nCreated cache key: %s", cache_key)
|
||||
hashed_cache_key = Cache._get_hashed_cache_key(cache_key)
|
||||
hashed_cache_key = self._add_namespace_to_cache_key(hashed_cache_key, **kwargs)
|
||||
verbose_logger.debug(
|
||||
"\nCreated cache key: %s (source material length: %d)",
|
||||
hashed_cache_key,
|
||||
len(cache_key),
|
||||
)
|
||||
# Remove preset_cache_key from kwargs to avoid "got multiple values" TypeError
|
||||
# when kwargs already contains preset_cache_key from upstream callers
|
||||
kwargs_for_preset = {k: v for k, v in kwargs.items() if k != "preset_cache_key"}
|
||||
|
|
@ -497,6 +501,34 @@ class Cache:
|
|||
return cached_response
|
||||
return cached_result
|
||||
|
||||
@staticmethod
|
||||
def _get_safe_cache_lookup_kwargs(kwargs: Dict[str, Any]) -> Dict[str, Any]:
|
||||
cache_lookup_kwargs: Dict[str, Any] = {}
|
||||
for prompt_kwarg in ("messages", "input"):
|
||||
if prompt_kwarg in kwargs:
|
||||
cache_lookup_kwargs[prompt_kwarg] = kwargs[prompt_kwarg]
|
||||
|
||||
if isinstance(kwargs.get("metadata"), dict):
|
||||
cache_lookup_kwargs["metadata"] = {}
|
||||
|
||||
return cache_lookup_kwargs
|
||||
|
||||
@staticmethod
|
||||
def _update_metadata_from_cache_lookup_kwargs(
|
||||
original_kwargs: Dict[str, Any], cache_lookup_kwargs: Dict[str, Any]
|
||||
) -> None:
|
||||
original_metadata = original_kwargs.get("metadata")
|
||||
cache_lookup_metadata = cache_lookup_kwargs.get("metadata")
|
||||
if not isinstance(original_metadata, dict) or not isinstance(
|
||||
cache_lookup_metadata, dict
|
||||
):
|
||||
return
|
||||
|
||||
if "semantic-similarity" in cache_lookup_metadata:
|
||||
original_metadata["semantic-similarity"] = cache_lookup_metadata[
|
||||
"semantic-similarity"
|
||||
]
|
||||
|
||||
def get_cache(self, dynamic_cache_object: Optional[BaseCache] = None, **kwargs):
|
||||
"""
|
||||
Retrieves the cached result for the given arguments.
|
||||
|
|
@ -511,7 +543,6 @@ class Cache:
|
|||
try: # never block execution
|
||||
if self.should_use_cache(**kwargs) is not True:
|
||||
return
|
||||
messages = kwargs.get("messages", [])
|
||||
if "cache_key" in kwargs:
|
||||
cache_key = kwargs["cache_key"]
|
||||
else:
|
||||
|
|
@ -523,12 +554,19 @@ class Cache:
|
|||
or cache_control_args.get("s-max-age")
|
||||
or float("inf")
|
||||
)
|
||||
cache_lookup_kwargs = self._get_safe_cache_lookup_kwargs(kwargs)
|
||||
if dynamic_cache_object is not None:
|
||||
cached_result = dynamic_cache_object.get_cache(
|
||||
cache_key, messages=messages
|
||||
cache_key, **cache_lookup_kwargs
|
||||
)
|
||||
else:
|
||||
cached_result = self.cache.get_cache(cache_key, messages=messages)
|
||||
cached_result = self.cache.get_cache(
|
||||
cache_key, **cache_lookup_kwargs
|
||||
)
|
||||
self._update_metadata_from_cache_lookup_kwargs(
|
||||
original_kwargs=kwargs,
|
||||
cache_lookup_kwargs=cache_lookup_kwargs,
|
||||
)
|
||||
return self._get_cache_logic(
|
||||
cached_result=cached_result, max_age=max_age
|
||||
)
|
||||
|
|
@ -549,7 +587,6 @@ class Cache:
|
|||
if self.should_use_cache(**kwargs) is not True:
|
||||
return
|
||||
|
||||
kwargs.get("messages", [])
|
||||
if "cache_key" in kwargs:
|
||||
cache_key = kwargs["cache_key"]
|
||||
else:
|
||||
|
|
@ -654,6 +691,7 @@ class Cache:
|
|||
self,
|
||||
embedding_response: Any,
|
||||
model: Optional[str],
|
||||
prompt_tokens: Optional[int] = None,
|
||||
prompt_tokens_details: Optional[dict] = None,
|
||||
) -> CachedEmbedding:
|
||||
"""
|
||||
|
|
@ -666,6 +704,7 @@ class Cache:
|
|||
"index": embedding_response.get("index"),
|
||||
"object": embedding_response.get("object"),
|
||||
"model": model,
|
||||
"prompt_tokens": prompt_tokens,
|
||||
"prompt_tokens_details": prompt_tokens_details,
|
||||
}
|
||||
elif hasattr(embedding_response, "model_dump"):
|
||||
|
|
@ -675,6 +714,7 @@ class Cache:
|
|||
"index": data.get("index"),
|
||||
"object": data.get("object"),
|
||||
"model": model,
|
||||
"prompt_tokens": prompt_tokens,
|
||||
"prompt_tokens_details": prompt_tokens_details,
|
||||
}
|
||||
else:
|
||||
|
|
@ -684,6 +724,7 @@ class Cache:
|
|||
"index": data.get("index"),
|
||||
"object": data.get("object"),
|
||||
"model": model,
|
||||
"prompt_tokens": prompt_tokens,
|
||||
"prompt_tokens_details": prompt_tokens_details,
|
||||
}
|
||||
except KeyError as e:
|
||||
|
|
@ -732,6 +773,29 @@ class Cache:
|
|||
per_item[key] = value
|
||||
return per_item if per_item else None
|
||||
|
||||
def _get_per_item_prompt_tokens(
|
||||
self,
|
||||
result: EmbeddingResponse,
|
||||
idx_in_result_data: int,
|
||||
) -> Optional[int]:
|
||||
"""
|
||||
Extract the per-item prompt_tokens from a response for caching.
|
||||
|
||||
Single-item responses store the full usage.prompt_tokens. Multi-item
|
||||
responses distribute it evenly (with remainder) so that summing all
|
||||
per-item values on retrieval reconstructs the original total.
|
||||
"""
|
||||
if result.usage is None or result.usage.prompt_tokens is None:
|
||||
return None
|
||||
|
||||
total = result.usage.prompt_tokens
|
||||
num_items = len(result.data)
|
||||
if num_items <= 1:
|
||||
return total
|
||||
|
||||
quotient, remainder = divmod(total, num_items)
|
||||
return quotient + (1 if idx_in_result_data < remainder else 0)
|
||||
|
||||
def add_embedding_response_to_cache(
|
||||
self,
|
||||
result: EmbeddingResponse,
|
||||
|
|
@ -743,7 +807,11 @@ class Cache:
|
|||
kwargs["cache_key"] = preset_cache_key
|
||||
embedding_response = result.data[idx_in_result_data]
|
||||
|
||||
# Extract per-item prompt_tokens_details from response usage
|
||||
# Extract per-item prompt_tokens + details from response usage
|
||||
prompt_tokens = self._get_per_item_prompt_tokens(
|
||||
result=result,
|
||||
idx_in_result_data=idx_in_result_data,
|
||||
)
|
||||
prompt_tokens_details = self._get_per_item_prompt_tokens_details(
|
||||
result=result,
|
||||
idx_in_result_data=idx_in_result_data,
|
||||
|
|
@ -754,6 +822,7 @@ class Cache:
|
|||
embedding_dict: CachedEmbedding = self._convert_to_cached_embedding(
|
||||
embedding_response,
|
||||
model_name,
|
||||
prompt_tokens=prompt_tokens,
|
||||
prompt_tokens_details=prompt_tokens_details,
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -394,7 +394,7 @@ class LLMCachingHandler:
|
|||
return cr["model"]
|
||||
return None
|
||||
|
||||
def _process_async_embedding_cached_response(
|
||||
def _process_async_embedding_cached_response( # noqa: PLR0915
|
||||
self,
|
||||
final_embedding_cached_response: Optional[EmbeddingResponse],
|
||||
cached_result: List[Optional[CachedEmbedding]],
|
||||
|
|
@ -456,7 +456,10 @@ class LLMCachingHandler:
|
|||
index=idx,
|
||||
object="embedding",
|
||||
)
|
||||
if isinstance(kwargs_input_as_list[idx], str):
|
||||
cached_prompt_tokens = cr.get("prompt_tokens")
|
||||
if cached_prompt_tokens is not None:
|
||||
prompt_tokens += cached_prompt_tokens
|
||||
elif isinstance(kwargs_input_as_list[idx], str):
|
||||
from litellm.utils import token_counter
|
||||
|
||||
prompt_tokens += token_counter(
|
||||
|
|
|
|||
|
|
@ -22,6 +22,7 @@ import litellm
|
|||
from litellm._logging import print_verbose, verbose_logger
|
||||
from litellm.constants import (
|
||||
DEFAULT_REDIS_MAJOR_VERSION,
|
||||
REDIS_CIRCUIT_BREAKER_ENABLED,
|
||||
REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD,
|
||||
REDIS_CIRCUIT_BREAKER_RECOVERY_TIMEOUT,
|
||||
)
|
||||
|
|
@ -114,15 +115,23 @@ class RedisCircuitBreaker:
|
|||
OPEN = "open"
|
||||
HALF_OPEN = "half_open"
|
||||
|
||||
def __init__(self, failure_threshold: int, recovery_timeout: int) -> None:
|
||||
def __init__(
|
||||
self,
|
||||
failure_threshold: int,
|
||||
recovery_timeout: int,
|
||||
enabled: bool = True,
|
||||
) -> None:
|
||||
self.failure_threshold = failure_threshold
|
||||
self.recovery_timeout = recovery_timeout
|
||||
self.enabled = enabled
|
||||
self._failure_count = 0
|
||||
self._opened_at: Optional[float] = None
|
||||
self._state = self.CLOSED
|
||||
|
||||
def is_open(self) -> bool:
|
||||
"""Returns True if Redis calls should be skipped."""
|
||||
if not self.enabled:
|
||||
return False
|
||||
if self._state == self.HALF_OPEN:
|
||||
# Probe already in flight — fast-fail all concurrent requests.
|
||||
# Only the one call that caused the OPEN→HALF_OPEN transition
|
||||
|
|
@ -136,6 +145,8 @@ class RedisCircuitBreaker:
|
|||
return False
|
||||
|
||||
def record_failure(self) -> None:
|
||||
if not self.enabled:
|
||||
return
|
||||
self._failure_count += 1
|
||||
self._opened_at = time.time()
|
||||
if self._failure_count >= self.failure_threshold:
|
||||
|
|
@ -149,6 +160,8 @@ class RedisCircuitBreaker:
|
|||
self._state = self.OPEN
|
||||
|
||||
def record_success(self) -> None:
|
||||
if not self.enabled:
|
||||
return
|
||||
if self._state == self.HALF_OPEN:
|
||||
verbose_logger.info("Redis circuit breaker CLOSED — Redis recovered")
|
||||
self._failure_count = 0
|
||||
|
|
@ -243,6 +256,7 @@ class RedisCache(BaseCache):
|
|||
self._circuit_breaker = RedisCircuitBreaker(
|
||||
failure_threshold=REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD,
|
||||
recovery_timeout=REDIS_CIRCUIT_BREAKER_RECOVERY_TIMEOUT,
|
||||
enabled=REDIS_CIRCUIT_BREAKER_ENABLED,
|
||||
)
|
||||
|
||||
self._setup_health_pings()
|
||||
|
|
|
|||
|
|
@ -213,6 +213,78 @@ class RedisSemanticCache(BaseCache):
|
|||
ttl = int(ttl)
|
||||
return ttl
|
||||
|
||||
@classmethod
|
||||
def _get_prompt_from_kwargs(cls, **kwargs) -> Optional[str]:
|
||||
"""
|
||||
Extract a semantic-cache prompt from chat or Responses API request kwargs.
|
||||
"""
|
||||
messages = kwargs.get("messages")
|
||||
if messages:
|
||||
return get_str_from_messages(messages)
|
||||
|
||||
if "input" not in kwargs:
|
||||
return None
|
||||
|
||||
prompt_parts: List[str] = []
|
||||
cls._collect_responses_input_text(kwargs.get("input"), prompt_parts)
|
||||
prompt = "\n".join(prompt_parts).strip()
|
||||
return prompt or None
|
||||
|
||||
@classmethod
|
||||
def _collect_responses_input_text(cls, value: Any, prompt_parts: List[str]) -> None:
|
||||
value = cls._coerce_response_input_value(value)
|
||||
if value is None:
|
||||
return
|
||||
|
||||
if isinstance(value, str):
|
||||
stripped_value = value.strip()
|
||||
if stripped_value:
|
||||
prompt_parts.append(stripped_value)
|
||||
return
|
||||
|
||||
if isinstance(value, (list, tuple)):
|
||||
for item in value:
|
||||
cls._collect_responses_input_text(item, prompt_parts)
|
||||
return
|
||||
|
||||
if isinstance(value, dict):
|
||||
content = value.get("content")
|
||||
if content is not None:
|
||||
cls._collect_responses_input_text(content, prompt_parts)
|
||||
return
|
||||
|
||||
for text_key in ("text", "output", "input_text", "output_text"):
|
||||
text_value = value.get(text_key)
|
||||
if isinstance(text_value, str):
|
||||
stripped_text = text_value.strip()
|
||||
if stripped_text:
|
||||
prompt_parts.append(stripped_text)
|
||||
return
|
||||
return
|
||||
|
||||
content = getattr(value, "content", None)
|
||||
if content is not None:
|
||||
cls._collect_responses_input_text(content, prompt_parts)
|
||||
return
|
||||
|
||||
for text_key in ("text", "output", "input_text", "output_text"):
|
||||
text_value = getattr(value, text_key, None)
|
||||
if isinstance(text_value, str):
|
||||
stripped_text = text_value.strip()
|
||||
if stripped_text:
|
||||
prompt_parts.append(stripped_text)
|
||||
return
|
||||
|
||||
@staticmethod
|
||||
def _coerce_response_input_value(value: Any) -> Any:
|
||||
model_dump = getattr(value, "model_dump", None)
|
||||
if callable(model_dump):
|
||||
return model_dump()
|
||||
dict_method = getattr(value, "dict", None)
|
||||
if callable(dict_method):
|
||||
return dict_method()
|
||||
return value
|
||||
|
||||
def _get_embedding(self, prompt: str) -> List[float]:
|
||||
"""
|
||||
Generate an embedding vector for the given prompt using the configured embedding model.
|
||||
|
|
@ -278,13 +350,11 @@ class RedisSemanticCache(BaseCache):
|
|||
|
||||
value_str: Optional[str] = None
|
||||
try:
|
||||
# Extract the prompt from messages
|
||||
messages = kwargs.get("messages", [])
|
||||
if not messages:
|
||||
print_verbose("No messages provided for semantic caching")
|
||||
prompt = self._get_prompt_from_kwargs(**kwargs)
|
||||
if prompt is None:
|
||||
print_verbose("No prompt provided for semantic caching")
|
||||
return
|
||||
|
||||
prompt = get_str_from_messages(messages)
|
||||
value_str = str(value)
|
||||
|
||||
store_kwargs: Dict[str, Any] = {
|
||||
|
|
@ -315,14 +385,12 @@ class RedisSemanticCache(BaseCache):
|
|||
print_verbose(f"Redis semantic-cache get_cache, kwargs: {kwargs}")
|
||||
|
||||
try:
|
||||
# Extract the prompt from messages
|
||||
messages = kwargs.get("messages", [])
|
||||
if not messages:
|
||||
print_verbose("No messages provided for semantic cache lookup")
|
||||
prompt = self._get_prompt_from_kwargs(**kwargs)
|
||||
if prompt is None:
|
||||
print_verbose("No prompt provided for semantic cache lookup")
|
||||
kwargs.setdefault("metadata", {})["semantic-similarity"] = 0.0
|
||||
return None
|
||||
|
||||
prompt = get_str_from_messages(messages)
|
||||
# Check the cache for semantically similar prompts in this exact
|
||||
# LiteLLM cache-key scope.
|
||||
check_kwargs: Dict[str, Any] = {
|
||||
|
|
@ -428,13 +496,11 @@ class RedisSemanticCache(BaseCache):
|
|||
print_verbose(f"Async Redis semantic-cache set_cache, kwargs: {kwargs}")
|
||||
|
||||
try:
|
||||
# Extract the prompt from messages
|
||||
messages = kwargs.get("messages", [])
|
||||
if not messages:
|
||||
print_verbose("No messages provided for semantic caching")
|
||||
prompt = self._get_prompt_from_kwargs(**kwargs)
|
||||
if prompt is None:
|
||||
print_verbose("No prompt provided for semantic caching")
|
||||
return
|
||||
|
||||
prompt = get_str_from_messages(messages)
|
||||
value_str = str(value)
|
||||
|
||||
# Generate embedding for the value (response) to cache
|
||||
|
|
@ -471,15 +537,12 @@ class RedisSemanticCache(BaseCache):
|
|||
print_verbose(f"Async Redis semantic-cache get_cache, kwargs: {kwargs}")
|
||||
|
||||
try:
|
||||
# Extract the prompt from messages
|
||||
messages = kwargs.get("messages", [])
|
||||
if not messages:
|
||||
print_verbose("No messages provided for semantic cache lookup")
|
||||
prompt = self._get_prompt_from_kwargs(**kwargs)
|
||||
if prompt is None:
|
||||
print_verbose("No prompt provided for semantic cache lookup")
|
||||
kwargs.setdefault("metadata", {})["semantic-similarity"] = 0.0
|
||||
return None
|
||||
|
||||
prompt = get_str_from_messages(messages)
|
||||
|
||||
# Generate embedding for the prompt
|
||||
prompt_embedding = await self._get_async_embedding(prompt, **kwargs)
|
||||
|
||||
|
|
|
|||
|
|
@ -171,6 +171,8 @@ class ResponsesToCompletionBridgeHandler:
|
|||
model_response = validated_kwargs["model_response"]
|
||||
logging_obj = validated_kwargs["logging_obj"]
|
||||
custom_llm_provider = validated_kwargs["custom_llm_provider"]
|
||||
if kwargs.get("stream") is True and "stream" not in optional_params:
|
||||
optional_params = {**optional_params, "stream": True}
|
||||
|
||||
request_data = self.transformation_handler.transform_request(
|
||||
model=model,
|
||||
|
|
@ -263,6 +265,8 @@ class ResponsesToCompletionBridgeHandler:
|
|||
model_response = validated_kwargs["model_response"]
|
||||
logging_obj = validated_kwargs["logging_obj"]
|
||||
custom_llm_provider = validated_kwargs["custom_llm_provider"]
|
||||
if kwargs.get("stream") is True and "stream" not in optional_params:
|
||||
optional_params = {**optional_params, "stream": True}
|
||||
|
||||
try:
|
||||
request_data = self.transformation_handler.transform_request(
|
||||
|
|
|
|||
|
|
@ -402,6 +402,20 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
|
|||
instructions,
|
||||
) = self.convert_chat_completion_messages_to_responses_api(messages)
|
||||
|
||||
# OpenAI's Responses API rejects an empty input. For a system-only
|
||||
# request, carry the system message as a system-role input item instead
|
||||
# of instructions, mirroring how non-string system content is already
|
||||
# handled in convert_chat_completion_messages_to_responses_api.
|
||||
if not input_items and instructions is not None:
|
||||
input_items = [
|
||||
{
|
||||
"type": "message",
|
||||
"role": "system",
|
||||
"content": [{"type": "input_text", "text": instructions}],
|
||||
}
|
||||
]
|
||||
instructions = None
|
||||
|
||||
optional_params = self._extract_extra_body_params(optional_params)
|
||||
|
||||
# Build responses API request using the reverse transformation logic
|
||||
|
|
|
|||
|
|
@ -398,6 +398,9 @@ REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD = int(
|
|||
REDIS_CIRCUIT_BREAKER_RECOVERY_TIMEOUT = int(
|
||||
os.getenv("REDIS_CIRCUIT_BREAKER_RECOVERY_TIMEOUT", 60)
|
||||
)
|
||||
REDIS_CIRCUIT_BREAKER_ENABLED = (
|
||||
os.getenv("REDIS_CIRCUIT_BREAKER_ENABLED", "true").lower() == "true"
|
||||
)
|
||||
# Default Redis major version to assume when version cannot be determined
|
||||
# Using 7 as it's the modern version that supports LPOP with count parameter
|
||||
DEFAULT_REDIS_MAJOR_VERSION = int(os.getenv("DEFAULT_REDIS_MAJOR_VERSION", 7))
|
||||
|
|
@ -418,6 +421,7 @@ REPLICATE_POLLING_DELAY_SECONDS = float(
|
|||
DEFAULT_ANTHROPIC_CHAT_MAX_TOKENS = int(
|
||||
os.getenv("DEFAULT_ANTHROPIC_CHAT_MAX_TOKENS", 4096)
|
||||
)
|
||||
DEFAULT_OCI_CHAT_MAX_TOKENS = 4096
|
||||
TOGETHER_AI_4_B = int(os.getenv("TOGETHER_AI_4_B", 4))
|
||||
TOGETHER_AI_8_B = int(os.getenv("TOGETHER_AI_8_B", 8))
|
||||
TOGETHER_AI_21_B = int(os.getenv("TOGETHER_AI_21_B", 21))
|
||||
|
|
@ -831,6 +835,7 @@ openai_compatible_providers: List = [
|
|||
"nano-gpt", # Nano-GPT - JSON-configured provider
|
||||
"poe", # Poe - JSON-configured provider
|
||||
"chutes", # Chutes - JSON-configured provider
|
||||
"parasail", # Parasail - JSON-configured provider
|
||||
"featherless_ai",
|
||||
"nscale",
|
||||
"nebius",
|
||||
|
|
@ -1157,6 +1162,7 @@ BEDROCK_CONVERSE_MODELS = [
|
|||
"openai.gpt-oss-120b-1:0",
|
||||
"anthropic.claude-haiku-4-5-20251001-v1:0",
|
||||
"anthropic.claude-sonnet-4-5-20250929-v1:0",
|
||||
"anthropic.claude-fable-5",
|
||||
"anthropic.claude-opus-4-8",
|
||||
"anthropic.claude-opus-4-7",
|
||||
"anthropic.claude-opus-4-6-v1:0",
|
||||
|
|
@ -1478,6 +1484,7 @@ DB_SPEND_UPDATE_JOB_NAME = "db_spend_update_job"
|
|||
DB_DAILY_TAG_SPEND_UPDATE_JOB_NAME = "db_daily_tag_spend_update_job"
|
||||
PROMETHEUS_EMIT_BUDGET_METRICS_JOB_NAME = "prometheus_emit_budget_metrics"
|
||||
CLOUDZERO_EXPORT_USAGE_DATA_JOB_NAME = "cloudzero_export_usage_data"
|
||||
MAVVRIK_FOCUS_EXPORT_JOB_NAME = "mavvrik_focus_export_usage_data"
|
||||
CLOUDZERO_MAX_FETCHED_DATA_RECORDS = int(
|
||||
os.getenv("CLOUDZERO_MAX_FETCHED_DATA_RECORDS", 50000)
|
||||
)
|
||||
|
|
@ -1492,6 +1499,10 @@ SPEND_LOG_CLEANUP_MAX_CONSECUTIVE_BATCH_FAILURES = int(
|
|||
SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS = float(
|
||||
os.getenv("SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS", 0.5)
|
||||
)
|
||||
SPEND_LOG_PARTITION_INTERVAL = os.getenv("SPEND_LOG_PARTITION_INTERVAL", "day")
|
||||
SPEND_LOG_PARTITION_PRECREATE_AHEAD = int(
|
||||
os.getenv("SPEND_LOG_PARTITION_PRECREATE_AHEAD", 7)
|
||||
)
|
||||
SPEND_LOG_QUEUE_SIZE_THRESHOLD = int(os.getenv("SPEND_LOG_QUEUE_SIZE_THRESHOLD", 100))
|
||||
SPEND_LOG_QUEUE_POLL_INTERVAL = float(os.getenv("SPEND_LOG_QUEUE_POLL_INTERVAL", 2.0))
|
||||
SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE = int(
|
||||
|
|
|
|||
|
|
@ -2425,12 +2425,11 @@ class BaseTokenUsageProcessor:
|
|||
if not attr.startswith("_") and not callable(
|
||||
getattr(usage.completion_tokens_details, attr)
|
||||
):
|
||||
current_val = getattr(
|
||||
combined.completion_tokens_details, attr, 0
|
||||
current_val = (
|
||||
getattr(combined.completion_tokens_details, attr, 0) or 0
|
||||
)
|
||||
new_val = getattr(usage.completion_tokens_details, attr, 0)
|
||||
|
||||
if new_val is not None and current_val is not None:
|
||||
new_val = getattr(usage.completion_tokens_details, attr, 0) or 0
|
||||
if isinstance(new_val, (int, float)):
|
||||
setattr(
|
||||
combined.completion_tokens_details,
|
||||
attr,
|
||||
|
|
@ -2489,6 +2488,11 @@ class RealtimeAPITokenUsageProcessor(BaseTokenUsageProcessor):
|
|||
)
|
||||
|
||||
|
||||
_TRANSCRIPTION_COMPLETED_EVENT_TYPE = (
|
||||
"conversation.item.input_audio_transcription.completed"
|
||||
)
|
||||
|
||||
|
||||
def handle_realtime_stream_cost_calculation(
|
||||
results: OpenAIRealtimeStreamList,
|
||||
combined_usage_object: Usage,
|
||||
|
|
@ -2534,4 +2538,99 @@ def handle_realtime_stream_cost_calculation(
|
|||
break # exit if we find a valid model
|
||||
total_cost = input_cost_per_token + output_cost_per_token
|
||||
|
||||
if any(r.get("type") == _TRANSCRIPTION_COMPLETED_EVENT_TYPE for r in results):
|
||||
total_cost += handle_realtime_transcription_cost_calculation(
|
||||
results=results,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
litellm_model_name=litellm_model_name,
|
||||
)
|
||||
|
||||
return total_cost
|
||||
|
||||
|
||||
def handle_realtime_transcription_cost_calculation(
|
||||
results: OpenAIRealtimeStreamList,
|
||||
custom_llm_provider: str,
|
||||
litellm_model_name: str,
|
||||
) -> float:
|
||||
"""
|
||||
Cost for realtime transcription sessions (e.g. gpt-realtime-whisper).
|
||||
|
||||
Transcription sessions emit no `response.done` events; instead each
|
||||
`conversation.item.input_audio_transcription.completed` event carries a
|
||||
`usage` object billed by the ASR model. The usage is one of:
|
||||
- {"type": "duration", "seconds": <float>} → priced via input_cost_per_second
|
||||
- {"type": "tokens", "input_tokens": ...} → priced via input/audio token cost
|
||||
"""
|
||||
completed_events = [
|
||||
cast(dict, result)
|
||||
for result in results
|
||||
if result.get("type") == _TRANSCRIPTION_COMPLETED_EVENT_TYPE
|
||||
]
|
||||
if not completed_events:
|
||||
return 0.0
|
||||
|
||||
model_name = (
|
||||
_get_transcription_model_name_from_results(results) or litellm_model_name
|
||||
)
|
||||
try:
|
||||
model_info = litellm.get_model_info(
|
||||
model=model_name, custom_llm_provider=custom_llm_provider
|
||||
)
|
||||
except Exception:
|
||||
model_info = None
|
||||
|
||||
total_cost = 0.0
|
||||
for event in completed_events:
|
||||
usage = event.get("usage") or {}
|
||||
total_cost += _transcription_usage_cost(usage, model_info)
|
||||
return total_cost
|
||||
|
||||
|
||||
def _get_transcription_model_name_from_results(
|
||||
results: OpenAIRealtimeStreamList,
|
||||
) -> Optional[str]:
|
||||
"""Resolve the ASR model from a transcription_session.* / session.* event."""
|
||||
for result in results:
|
||||
if result.get("type") in (
|
||||
"transcription_session.created",
|
||||
"transcription_session.updated",
|
||||
"session.created",
|
||||
"session.updated",
|
||||
):
|
||||
session = cast(dict, result).get("session", {}) or {}
|
||||
transcription = (
|
||||
(session.get("audio", {}) or {}).get("input", {}) or {}
|
||||
).get("transcription", {}) or session.get("input_audio_transcription", {})
|
||||
model = (transcription or {}).get("model") or session.get("model")
|
||||
if model:
|
||||
return model
|
||||
return None
|
||||
|
||||
|
||||
def _transcription_usage_cost(usage: dict, model_info: Optional[ModelInfo]) -> float:
|
||||
if model_info is None:
|
||||
return 0.0
|
||||
usage_type = usage.get("type")
|
||||
if usage_type == "duration":
|
||||
seconds = usage.get("seconds") or 0.0
|
||||
per_second = model_info.get("input_cost_per_second") or 0.0
|
||||
return float(seconds) * float(per_second)
|
||||
if usage_type == "tokens":
|
||||
input_token_details = usage.get("input_token_details") or {}
|
||||
audio_tokens = input_token_details.get("audio_tokens") or 0
|
||||
text_tokens = input_token_details.get("text_tokens") or 0
|
||||
output_tokens = usage.get("output_tokens") or 0
|
||||
audio_cost = float(audio_tokens) * float(
|
||||
model_info.get("input_cost_per_audio_token")
|
||||
or model_info.get("input_cost_per_token")
|
||||
or 0.0
|
||||
)
|
||||
text_cost = float(text_tokens) * float(
|
||||
model_info.get("input_cost_per_token") or 0.0
|
||||
)
|
||||
output_cost = float(output_tokens) * float(
|
||||
model_info.get("output_cost_per_token") or 0.0
|
||||
)
|
||||
return audio_cost + text_cost + output_cost
|
||||
return 0.0
|
||||
|
|
|
|||
|
|
@ -9,13 +9,109 @@
|
|||
|
||||
## LiteLLM versions of the OpenAI Exception Types
|
||||
|
||||
from typing import Any, Dict, Optional
|
||||
import enum
|
||||
from typing import Any, Dict, Optional, Union
|
||||
|
||||
import httpx
|
||||
import openai
|
||||
|
||||
from litellm.types.utils import LiteLLMCommonStrings
|
||||
|
||||
|
||||
class RateLimitErrorCategory(str, enum.Enum):
|
||||
"""
|
||||
Category of a rate limit error, allowing callers to distinguish where the rate
|
||||
limit originated. Exposed on every :class:`RateLimitError` instance via the
|
||||
``category`` attribute.
|
||||
|
||||
Use these values to switch on the rate limit source, e.g.::
|
||||
|
||||
try:
|
||||
...
|
||||
except litellm.RateLimitError as e:
|
||||
if e.category == RateLimitErrorCategory.LITELLM_RATE_LIMIT:
|
||||
... # litellm's own limiter (key/team/user/model RPM/TPM/budget)
|
||||
elif e.category == RateLimitErrorCategory.VENDOR_RATE_LIMIT:
|
||||
... # the upstream LLM provider returned 429
|
||||
"""
|
||||
|
||||
VENDOR_RATE_LIMIT = "vendor_rate_limit"
|
||||
"""The upstream LLM provider returned a rate-limit response (e.g. OpenAI 429)."""
|
||||
|
||||
VENDOR_BATCH_RATE_LIMIT = "vendor_batch_rate_limit"
|
||||
"""The upstream LLM provider returned a rate-limit response on a batch endpoint."""
|
||||
|
||||
LITELLM_RATE_LIMIT = "litellm_rate_limit"
|
||||
"""LiteLLM's own rate limiter (key/team/user/model RPM/TPM, budget, parallel-requests, etc.) blocked the request."""
|
||||
|
||||
LITELLM_BATCH_RATE_LIMIT = "litellm_batch_rate_limit"
|
||||
"""LiteLLM's own batch rate limiter (token/request budget across a batch input file) blocked the request."""
|
||||
|
||||
|
||||
class RateLimitType(str, enum.Enum):
|
||||
"""
|
||||
The dimension that was exceeded when a rate-limit error fired.
|
||||
|
||||
This is orthogonal to :class:`RateLimitErrorCategory` — *category* tells
|
||||
callers **who** rate-limited the request (the upstream vendor vs. one of
|
||||
litellm's own limiters), while *type* tells them **which limit dimension**
|
||||
was exceeded (an RPM ceiling, a TPM ceiling, a max-parallel-requests
|
||||
ceiling, a budget cap, or a max-iterations cap).
|
||||
|
||||
Surfaced both on every :class:`RateLimitError` instance via the
|
||||
``rate_limit_type`` attribute and on the structured
|
||||
``StandardLoggingPayload.error_information.error_rate_limit_type`` field
|
||||
so custom callbacks / metrics consumers can split rate-limit failures by
|
||||
cause without parsing free-text error messages.
|
||||
"""
|
||||
|
||||
REQUESTS = "requests"
|
||||
"""Requests-per-minute (RPM) or requests-per-window ceiling exceeded."""
|
||||
|
||||
TOKENS = "tokens"
|
||||
"""Tokens-per-minute (TPM) or tokens-per-window ceiling exceeded."""
|
||||
|
||||
CONCURRENT_REQUESTS = "concurrent_requests"
|
||||
"""``max_parallel_requests`` — too many in-flight requests at once."""
|
||||
|
||||
BUDGET = "budget"
|
||||
"""Spend budget cap reached (key, team, user, or per-session)."""
|
||||
|
||||
MAX_ITERATIONS = "max_iterations"
|
||||
"""Per-session max-iterations cap reached (agent-style flows)."""
|
||||
|
||||
|
||||
_RATE_LIMIT_CATEGORY_VALUES = frozenset(c.value for c in RateLimitErrorCategory)
|
||||
_RATE_LIMIT_TYPE_VALUES = frozenset(t.value for t in RateLimitType)
|
||||
|
||||
|
||||
def validate_rate_limit_category(value: Any) -> Optional[str]:
|
||||
"""Return ``value`` only if it matches a known :class:`RateLimitErrorCategory`.
|
||||
|
||||
Used at duck-typed read sites (StandardLoggingPayload extraction, Prometheus
|
||||
labels) to reject `.category` strings set by unrelated third-party exceptions
|
||||
— otherwise those would leak into custom-callback payloads and Prometheus
|
||||
label cardinality.
|
||||
"""
|
||||
if isinstance(value, RateLimitErrorCategory):
|
||||
return value.value
|
||||
if isinstance(value, str) and value in _RATE_LIMIT_CATEGORY_VALUES:
|
||||
return value
|
||||
return None
|
||||
|
||||
|
||||
def validate_rate_limit_type(value: Any) -> Optional[str]:
|
||||
"""Return ``value`` only if it matches a known :class:`RateLimitType`.
|
||||
|
||||
See :func:`validate_rate_limit_category` for the rationale.
|
||||
"""
|
||||
if isinstance(value, RateLimitType):
|
||||
return value.value
|
||||
if isinstance(value, str) and value in _RATE_LIMIT_TYPE_VALUES:
|
||||
return value
|
||||
return None
|
||||
|
||||
|
||||
_MINIMAL_ERROR_RESPONSE: Optional[httpx.Response] = None
|
||||
|
||||
|
||||
|
|
@ -321,6 +417,18 @@ class PermissionDeniedError(openai.PermissionDeniedError): # type: ignore
|
|||
|
||||
|
||||
class RateLimitError(openai.RateLimitError): # type: ignore
|
||||
"""
|
||||
Unified rate-limit error.
|
||||
|
||||
Every rate-limit condition surfaced by litellm — whether it originated from
|
||||
an upstream LLM provider, a vendor batch endpoint, or one of litellm's own
|
||||
proxy-side limiters (parallel-requests, dynamic-rate, batch-rate, budget,
|
||||
max-iterations, etc.) — is raised as an instance of this class.
|
||||
|
||||
The :attr:`category` attribute lets callers distinguish the source. See
|
||||
:class:`RateLimitErrorCategory` for the available values.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
message,
|
||||
|
|
@ -330,6 +438,12 @@ class RateLimitError(openai.RateLimitError): # type: ignore
|
|||
litellm_debug_info: Optional[str] = None,
|
||||
max_retries: Optional[int] = None,
|
||||
num_retries: Optional[int] = None,
|
||||
category: Union[str, RateLimitErrorCategory] = (
|
||||
RateLimitErrorCategory.VENDOR_RATE_LIMIT
|
||||
),
|
||||
rate_limit_type: Optional[Union[str, RateLimitType]] = None,
|
||||
headers: Optional[Dict[str, str]] = None,
|
||||
detail: Any = None,
|
||||
):
|
||||
self.status_code = 429
|
||||
self.message = "litellm.RateLimitError: {}".format(message)
|
||||
|
|
@ -338,9 +452,39 @@ class RateLimitError(openai.RateLimitError): # type: ignore
|
|||
self.litellm_debug_info = litellm_debug_info
|
||||
self.max_retries = max_retries
|
||||
self.num_retries = num_retries
|
||||
self.category = (
|
||||
category.value if isinstance(category, RateLimitErrorCategory) else category
|
||||
)
|
||||
# Which dimension was exceeded — request count, token count, parallel
|
||||
# requests, budget, max iterations. None when the source didn't
|
||||
# classify the failure (e.g. legacy vendor 429 with no header hints).
|
||||
self.rate_limit_type: Optional[str] = (
|
||||
rate_limit_type.value
|
||||
if isinstance(rate_limit_type, RateLimitType)
|
||||
else rate_limit_type
|
||||
)
|
||||
# Headers explicitly attached to the error (e.g. retry-after,
|
||||
# rate_limit_type, reset_at). Preserved across the proxy boundary so
|
||||
# clients can react appropriately.
|
||||
#
|
||||
# IMPORTANT: we deliberately do NOT auto-populate self.headers from
|
||||
# response.headers when only `response` is provided. A vendor 429 can
|
||||
# set arbitrary response headers (Set-Cookie, CORS overrides, …); if
|
||||
# those leaked into e.headers and a downstream proxy serializer
|
||||
# forwarded them to the client, a malicious upstream could inject
|
||||
# browser-interpreted headers for the proxy origin. Vendor response
|
||||
# headers stay reachable on `e.response.headers` for callers that
|
||||
# explicitly want them; only the proxy-supplied `headers=` kwarg
|
||||
# makes it onto `self.headers`.
|
||||
_response_headers = (
|
||||
getattr(response, "headers", None) if response is not None else None
|
||||
)
|
||||
self.headers: Optional[Dict[str, str]] = (
|
||||
{k: str(v) for k, v in headers.items()} if headers else None
|
||||
)
|
||||
# Mirrors FastAPI HTTPException.detail so the same instance can be
|
||||
# serialized through both the ProxyException and HTTPException paths.
|
||||
self.detail = detail if detail is not None else self.message
|
||||
self.response = httpx.Response(
|
||||
status_code=429,
|
||||
headers=_response_headers,
|
||||
|
|
@ -843,11 +987,24 @@ LITELLM_EXCEPTION_TYPES = [
|
|||
|
||||
class BudgetExceededError(Exception):
|
||||
def __init__(
|
||||
self, current_cost: float, max_budget: float, message: Optional[str] = None
|
||||
self,
|
||||
current_cost: float,
|
||||
max_budget: float,
|
||||
message: Optional[str] = None,
|
||||
llm_provider: Optional[str] = None,
|
||||
):
|
||||
self.current_cost = current_cost
|
||||
self.max_budget = max_budget
|
||||
self.status_code = 429
|
||||
self.llm_provider = llm_provider or ""
|
||||
# Surface unified rate-limit fields without joining the RateLimitError
|
||||
# hierarchy so existing `except BudgetExceededError:` handlers keep
|
||||
# working; custom callbacks reading StandardLoggingPayload pick these
|
||||
# up via the same `category` / `rate_limit_type` attributes the rest
|
||||
# of the unified rate-limit error path uses. Stored as plain strings
|
||||
# to match the normalization RateLimitError.__init__ performs.
|
||||
self.category: str = RateLimitErrorCategory.LITELLM_RATE_LIMIT.value
|
||||
self.rate_limit_type: str = RateLimitType.BUDGET.value
|
||||
message = (
|
||||
message
|
||||
or f"Budget has been exceeded! Current cost: {current_cost}, Max budget: {max_budget}"
|
||||
|
|
@ -1062,3 +1219,37 @@ class GuardrailInterventionNormalStringError(
|
|||
|
||||
def __repr__(self):
|
||||
return self.__str__()
|
||||
|
||||
|
||||
class SensitiveDataRouteException(Exception):
|
||||
"""
|
||||
Exception raised when a guardrail detects sensitive data and wants to reroute the request.
|
||||
|
||||
Instead of blocking the request, this exception signals that the request should be
|
||||
routed to a different model (typically an on-premise model for data privacy).
|
||||
|
||||
The proxy catches this exception and:
|
||||
1. Reroutes the current request to the specified model
|
||||
2. When sticky_session_routing is True, stores the routing decision in session
|
||||
cache so all subsequent requests in the same session are routed to the same model
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
route_to_model: str,
|
||||
session_id: str,
|
||||
guardrail_name: Optional[str] = None,
|
||||
detection_info: Optional[Dict[str, Any]] = None,
|
||||
message: Optional[str] = None,
|
||||
sticky_session_routing: bool = True,
|
||||
):
|
||||
self.route_to_model = route_to_model
|
||||
self.session_id = session_id
|
||||
self.guardrail_name = guardrail_name
|
||||
self.detection_info = detection_info or {}
|
||||
self.sticky_session_routing = sticky_session_routing
|
||||
self.message = (
|
||||
message
|
||||
or f"Sensitive data detected by {guardrail_name}. Routing to model: {route_to_model}"
|
||||
)
|
||||
super().__init__(self.message)
|
||||
|
|
|
|||
|
|
@ -4,6 +4,7 @@ LiteLLM Proxy uses this MCP Client to connnect to other MCP servers.
|
|||
|
||||
import asyncio
|
||||
import base64
|
||||
import os
|
||||
from typing import (
|
||||
Any,
|
||||
Awaitable,
|
||||
|
|
@ -16,7 +17,6 @@ from typing import (
|
|||
TypeVar,
|
||||
Union,
|
||||
)
|
||||
|
||||
import httpx
|
||||
from mcp import ClientSession, ReadResourceResult, Resource, StdioServerParameters
|
||||
from mcp.client.sse import sse_client
|
||||
|
|
@ -42,9 +42,8 @@ from mcp.types import (
|
|||
)
|
||||
from mcp.types import Tool as MCPTool
|
||||
from pydantic import AnyUrl
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.constants import MCP_CLIENT_TIMEOUT
|
||||
from litellm.constants import MCP_CLIENT_TIMEOUT, MCP_NPM_CACHE_DIR
|
||||
from litellm.llms.custom_httpx.http_handler import get_ssl_configuration
|
||||
from litellm.types.llms.custom_http import VerifyTypes
|
||||
from litellm.types.mcp import (
|
||||
|
|
@ -61,13 +60,33 @@ def to_basic_auth(auth_value: str) -> str:
|
|||
return base64.b64encode(auth_value.encode("utf-8")).decode()
|
||||
|
||||
|
||||
def _strip_header_whitespace(headers: Dict[str, str]) -> Dict[str, str]:
|
||||
return {
|
||||
(key.strip() if isinstance(key, str) else key): (
|
||||
value.strip() if isinstance(value, str) else value
|
||||
)
|
||||
for key, value in headers.items()
|
||||
}
|
||||
|
||||
|
||||
def _first_non_cancelled_cause(exc: BaseException) -> Optional[BaseException]:
|
||||
queue: List[BaseException] = [exc]
|
||||
while queue:
|
||||
current = queue.pop(0)
|
||||
nested = getattr(current, "exceptions", None)
|
||||
if nested:
|
||||
queue.extend(nested)
|
||||
elif not isinstance(current, asyncio.CancelledError):
|
||||
return current
|
||||
return None
|
||||
|
||||
|
||||
TSessionResult = TypeVar("TSessionResult")
|
||||
|
||||
|
||||
class MCPSigV4Auth(httpx.Auth):
|
||||
"""
|
||||
httpx Auth class that signs each request with AWS SigV4.
|
||||
|
||||
This is used for MCP servers that require AWS SigV4 authentication,
|
||||
such as AWS Bedrock AgentCore MCP servers. httpx calls auth_flow()
|
||||
for every outgoing request, enabling per-request signature computation.
|
||||
|
|
@ -92,10 +111,8 @@ class MCPSigV4Auth(httpx.Auth):
|
|||
"Missing botocore to use AWS SigV4 authentication. "
|
||||
"Run 'pip install boto3'."
|
||||
)
|
||||
|
||||
self.service_name = aws_service_name or "bedrock-agentcore"
|
||||
self.region_name = aws_region_name or "us-east-1"
|
||||
|
||||
# Note: os.environ/ prefixed values are already resolved by
|
||||
# ProxyConfig._check_for_os_environ_vars() at config load time.
|
||||
# Values arrive here as plain strings.
|
||||
|
|
@ -143,20 +160,17 @@ class MCPSigV4Auth(httpx.Auth):
|
|||
session_name = (
|
||||
aws_session_name or f"litellm-mcp-{int(__import__('time').time())}"
|
||||
)
|
||||
|
||||
sts_kwargs: dict = {"region_name": aws_region_name}
|
||||
if aws_access_key_id and aws_secret_access_key:
|
||||
sts_kwargs["aws_access_key_id"] = aws_access_key_id
|
||||
sts_kwargs["aws_secret_access_key"] = aws_secret_access_key
|
||||
if aws_session_token:
|
||||
sts_kwargs["aws_session_token"] = aws_session_token
|
||||
|
||||
sts_client = boto3.client("sts", **sts_kwargs)
|
||||
sts_response = sts_client.assume_role(
|
||||
RoleArn=aws_role_name,
|
||||
RoleSessionName=session_name,
|
||||
)
|
||||
|
||||
sts_creds = sts_response["Credentials"]
|
||||
return Credentials(
|
||||
access_key=sts_creds["AccessKeyId"],
|
||||
|
|
@ -178,17 +192,14 @@ class MCPSigV4Auth(httpx.Auth):
|
|||
data=request.content,
|
||||
headers=dict(request.headers),
|
||||
)
|
||||
|
||||
# Sign the request — SigV4Auth.add_auth() adds Authorization,
|
||||
# X-Amz-Date, and X-Amz-Security-Token (if session token present).
|
||||
# Host header is derived automatically from the URL.
|
||||
sigv4 = SigV4Auth(self.credentials, self.service_name, self.region_name)
|
||||
sigv4.add_auth(aws_request)
|
||||
|
||||
# Copy SigV4 headers back to the httpx request
|
||||
for header_name, header_value in aws_request.headers.items():
|
||||
request.headers[header_name] = header_value
|
||||
|
||||
yield request
|
||||
|
||||
|
||||
|
|
@ -198,6 +209,8 @@ class MCPClient:
|
|||
SSE and HTTP transports
|
||||
Authentication via Bearer token, Basic Auth, or API Key
|
||||
Tool calling with error handling and result parsing
|
||||
Sampling callbacks for upstream server LLM requests
|
||||
Elicitation callbacks for upstream server user-input requests
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
|
|
@ -211,6 +224,9 @@ class MCPClient:
|
|||
extra_headers: Optional[Dict[str, str]] = None,
|
||||
ssl_verify: Optional[VerifyTypes] = None,
|
||||
aws_auth: Optional[httpx.Auth] = None,
|
||||
sampling_callback: Optional[Callable] = None,
|
||||
elicitation_callback: Optional[Callable] = None,
|
||||
logging_callback: Optional[Callable] = None,
|
||||
):
|
||||
self.server_url: str = server_url
|
||||
self.transport_type: MCPTransport = transport_type
|
||||
|
|
@ -222,6 +238,9 @@ class MCPClient:
|
|||
self.ssl_verify: Optional[VerifyTypes] = ssl_verify
|
||||
self._aws_auth: Optional[httpx.Auth] = aws_auth
|
||||
self._last_initialize_instructions: Optional[str] = None
|
||||
self._sampling_callback: Optional[Callable] = sampling_callback
|
||||
self._elicitation_callback: Optional[Callable] = elicitation_callback
|
||||
self._logging_callback: Optional[Callable] = logging_callback
|
||||
# handle the basic auth value if provided
|
||||
if auth_value:
|
||||
self.update_auth_value(auth_value)
|
||||
|
|
@ -231,23 +250,20 @@ class MCPClient:
|
|||
) -> Tuple[Any, Optional[httpx.AsyncClient]]:
|
||||
"""
|
||||
Create the appropriate transport context based on transport type.
|
||||
|
||||
Returns:
|
||||
Tuple of (transport_context, http_client).
|
||||
http_client is only set for HTTP transport and needs cleanup.
|
||||
"""
|
||||
http_client: Optional[httpx.AsyncClient] = None
|
||||
|
||||
if self.transport_type == MCPTransport.stdio:
|
||||
if not self.stdio_config:
|
||||
raise ValueError("stdio_config is required for stdio transport")
|
||||
server_params = StdioServerParameters(
|
||||
command=self.stdio_config.get("command", ""),
|
||||
args=self.stdio_config.get("args", []),
|
||||
env=self.stdio_config.get("env", {}),
|
||||
env=self._get_safe_stdio_env(self.stdio_config.get("env")),
|
||||
)
|
||||
return stdio_client(server_params), None
|
||||
|
||||
if self.transport_type == MCPTransport.sse:
|
||||
headers = self._get_auth_headers()
|
||||
httpx_client_factory = self._create_httpx_client_factory()
|
||||
|
|
@ -260,14 +276,12 @@ class MCPClient:
|
|||
),
|
||||
None,
|
||||
)
|
||||
|
||||
# HTTP transport (default)
|
||||
if streamable_http_client is None:
|
||||
raise ImportError(
|
||||
"streamable_http_client is not available. "
|
||||
"Please install mcp with HTTP support."
|
||||
)
|
||||
|
||||
headers = self._get_auth_headers()
|
||||
httpx_client_factory = self._create_httpx_client_factory()
|
||||
verbose_logger.debug("litellm headers for streamable_http_client: %s", headers)
|
||||
|
|
@ -281,6 +295,54 @@ class MCPClient:
|
|||
)
|
||||
return transport_ctx, http_client
|
||||
|
||||
def _get_safe_stdio_env(
|
||||
self, provided_env: Optional[Dict[str, str]]
|
||||
) -> Optional[Dict[str, str]]:
|
||||
"""
|
||||
Return a safe environment for the stdio subprocess.
|
||||
|
||||
If provided_env is set, we use it as-is.
|
||||
If provided_env is None, we return a minimal allowlist from the parent environment
|
||||
to avoid leaking sensitive LiteLLM keys (OPENAI_API_KEY, etc.) to sub-processes.
|
||||
"""
|
||||
if provided_env is not None:
|
||||
return provided_env
|
||||
|
||||
# Minimal allowlist of safe/standard environment variables
|
||||
safe_keys = {
|
||||
"PATH",
|
||||
"HOME",
|
||||
"USER",
|
||||
"LOGNAME",
|
||||
"TMPDIR",
|
||||
"TMP",
|
||||
"TEMP",
|
||||
"SHELL",
|
||||
"LANG",
|
||||
"LC_ALL",
|
||||
# Node/Package manager caches
|
||||
"NPM_CONFIG_CACHE",
|
||||
"PNPM_HOME",
|
||||
"XDG_CACHE_HOME",
|
||||
"XDG_CONFIG_HOME",
|
||||
"XDG_DATA_HOME",
|
||||
# System info
|
||||
"SYSTEMROOT",
|
||||
"COMSPEC",
|
||||
"PATHEXT",
|
||||
"WINDIR",
|
||||
}
|
||||
|
||||
safe_env = {}
|
||||
for key in safe_keys:
|
||||
if key in os.environ:
|
||||
safe_env[key] = os.environ[key]
|
||||
|
||||
if "NPM_CONFIG_CACHE" not in safe_env:
|
||||
safe_env["NPM_CONFIG_CACHE"] = MCP_NPM_CACHE_DIR
|
||||
|
||||
return safe_env
|
||||
|
||||
async def _execute_session_operation(
|
||||
self,
|
||||
transport_ctx: Any,
|
||||
|
|
@ -288,13 +350,24 @@ class MCPClient:
|
|||
) -> TSessionResult:
|
||||
"""
|
||||
Execute an operation within a transport and session context.
|
||||
|
||||
Handles entering/exiting contexts and running the operation.
|
||||
Passes sampling/elicitation/logging callbacks to the ClientSession
|
||||
so that upstream MCP servers can request LLM inference (sampling),
|
||||
user input (elicitation), or send log messages.
|
||||
"""
|
||||
transport = await transport_ctx.__aenter__()
|
||||
in_flight_error: Optional[BaseException] = None
|
||||
try:
|
||||
read_stream, write_stream = transport[0], transport[1]
|
||||
session_ctx = ClientSession(read_stream, write_stream)
|
||||
# Build session kwargs with optional callbacks
|
||||
session_kwargs: Dict[str, Any] = {}
|
||||
if self._sampling_callback is not None:
|
||||
session_kwargs["sampling_callback"] = self._sampling_callback
|
||||
if self._elicitation_callback is not None:
|
||||
session_kwargs["elicitation_callback"] = self._elicitation_callback
|
||||
if self._logging_callback is not None:
|
||||
session_kwargs["logging_callback"] = self._logging_callback
|
||||
session_ctx = ClientSession(read_stream, write_stream, **session_kwargs)
|
||||
session = await session_ctx.__aenter__()
|
||||
try:
|
||||
init_result = await session.initialize()
|
||||
|
|
@ -309,11 +382,21 @@ class MCPClient:
|
|||
await session_ctx.__aexit__(None, None, None)
|
||||
except BaseException as e:
|
||||
verbose_logger.debug(f"Error during session context exit: {e}")
|
||||
except BaseException as e:
|
||||
in_flight_error = e
|
||||
raise
|
||||
finally:
|
||||
try:
|
||||
await transport_ctx.__aexit__(None, None, None)
|
||||
except BaseException as e:
|
||||
verbose_logger.debug(f"Error during transport context exit: {e}")
|
||||
except BaseException as exit_error:
|
||||
verbose_logger.debug(
|
||||
f"Error during transport context exit: {exit_error}"
|
||||
)
|
||||
root_cause = _first_non_cancelled_cause(exit_error)
|
||||
if root_cause is not None and isinstance(
|
||||
in_flight_error, asyncio.CancelledError
|
||||
):
|
||||
raise root_cause from in_flight_error
|
||||
|
||||
async def run_with_session(
|
||||
self, operation: Callable[[ClientSession], Awaitable[TSessionResult]]
|
||||
|
|
@ -351,7 +434,6 @@ class MCPClient:
|
|||
def _get_auth_headers(self) -> dict:
|
||||
"""Generate authentication headers based on auth type."""
|
||||
headers = {}
|
||||
|
||||
if self._mcp_auth_value:
|
||||
if isinstance(self._mcp_auth_value, str):
|
||||
if self.auth_type == MCPAuth.bearer_token:
|
||||
|
|
@ -373,17 +455,14 @@ class MCPClient:
|
|||
# Note: aws_sigv4 auth is not handled here — SigV4 requires per-request
|
||||
# signing (including the body hash), so it uses httpx.Auth flow instead
|
||||
# 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)
|
||||
|
||||
return headers
|
||||
return _strip_header_whitespace(headers)
|
||||
|
||||
def _create_httpx_client_factory(self) -> Callable[..., httpx.AsyncClient]:
|
||||
"""
|
||||
Create a custom httpx client factory that uses LiteLLM's SSL configuration.
|
||||
|
||||
This factory follows the same CA bundle path logic as http_handler.py:
|
||||
1. Check ssl_verify parameter (can be SSLContext, bool, or path to CA bundle)
|
||||
2. Check SSL_VERIFY environment variable
|
||||
|
|
@ -400,17 +479,14 @@ class MCPClient:
|
|||
"""Create an httpx.AsyncClient with LiteLLM's SSL configuration."""
|
||||
# Get unified SSL configuration using the same logic as http_handler.py
|
||||
ssl_config = get_ssl_configuration(self.ssl_verify)
|
||||
|
||||
verbose_logger.debug(
|
||||
f"MCP client using SSL configuration: {type(ssl_config).__name__}"
|
||||
)
|
||||
|
||||
# Use SigV4 auth if configured and no explicit auth provided.
|
||||
# The MCP SDK's sse_client and streamable_http_client call this
|
||||
# factory without passing auth=, so self._aws_auth is used.
|
||||
# For non-SigV4 clients, self._aws_auth is None — no behavior change.
|
||||
effective_auth = auth if auth is not None else self._aws_auth
|
||||
|
||||
return httpx.AsyncClient(
|
||||
headers=headers,
|
||||
timeout=timeout,
|
||||
|
|
@ -458,7 +534,6 @@ class MCPClient:
|
|||
f"Server: {self.server_url or 'stdio'}, "
|
||||
f"Transport: {self.transport_type}"
|
||||
)
|
||||
|
||||
# Check if it's a stream/connection error
|
||||
if "BrokenResourceError" in error_type or "Broken" in error_type:
|
||||
verbose_logger.error(
|
||||
|
|
@ -491,7 +566,6 @@ class MCPClient:
|
|||
f"MCP Tool '{call_tool_request_params.name}' progress: "
|
||||
f"{progress}/{total} ({percentage:.0f}%) - {message or ''}"
|
||||
)
|
||||
|
||||
# Forward to Host if callback provided
|
||||
if host_progress_callback:
|
||||
try:
|
||||
|
|
@ -514,14 +588,15 @@ class MCPClient:
|
|||
)
|
||||
return tool_result
|
||||
except asyncio.CancelledError:
|
||||
verbose_logger.warning("MCP client tool call was cancelled")
|
||||
verbose_logger.warning(
|
||||
f"MCP client tool call timed out after {self.timeout}s for {self.server_url}"
|
||||
)
|
||||
raise
|
||||
except Exception as e:
|
||||
import traceback
|
||||
|
||||
error_trace = traceback.format_exc()
|
||||
verbose_logger.debug(f"MCP client tool call traceback:\n{error_trace}")
|
||||
|
||||
# Log detailed error information
|
||||
error_type = type(e).__name__
|
||||
verbose_logger.error(
|
||||
|
|
@ -532,14 +607,12 @@ class MCPClient:
|
|||
f"Server: {self.server_url or 'stdio'}, "
|
||||
f"Transport: {self.transport_type}"
|
||||
)
|
||||
|
||||
# Check if it's a stream/connection error
|
||||
if "BrokenResourceError" in error_type or "Broken" in error_type:
|
||||
verbose_logger.error(
|
||||
"MCP client detected broken connection/stream - "
|
||||
"the MCP server may have crashed, disconnected, or timed out."
|
||||
)
|
||||
|
||||
# Return a default error result instead of raising
|
||||
return MCPCallToolResult(
|
||||
content=[
|
||||
|
|
@ -577,14 +650,12 @@ class MCPClient:
|
|||
f"Server: {self.server_url or 'stdio'}, "
|
||||
f"Transport: {self.transport_type}"
|
||||
)
|
||||
|
||||
# Check if it's a stream/connection error
|
||||
if "BrokenResourceError" in error_type or "Broken" in error_type:
|
||||
verbose_logger.error(
|
||||
"MCP client detected broken connection/stream during list_tools - "
|
||||
"the MCP server may have crashed, disconnected, or timed out"
|
||||
)
|
||||
|
||||
# Return empty list instead of raising to allow graceful degradation
|
||||
return []
|
||||
|
||||
|
|
@ -617,7 +688,6 @@ class MCPClient:
|
|||
|
||||
error_trace = traceback.format_exc()
|
||||
verbose_logger.debug(f"MCP client get_prompt traceback:\n{error_trace}")
|
||||
|
||||
# Log detailed error information
|
||||
error_type = type(e).__name__
|
||||
verbose_logger.error(
|
||||
|
|
@ -628,14 +698,12 @@ class MCPClient:
|
|||
f"Server: {self.server_url or 'stdio'}, "
|
||||
f"Transport: {self.transport_type}"
|
||||
)
|
||||
|
||||
# Check if it's a stream/connection error
|
||||
if "BrokenResourceError" in error_type or "Broken" in error_type:
|
||||
verbose_logger.error(
|
||||
"MCP client detected broken connection/stream during get_prompt - "
|
||||
"the MCP server may have crashed, disconnected, or timed out."
|
||||
)
|
||||
|
||||
raise
|
||||
|
||||
async def list_resources(self) -> list[Resource]:
|
||||
|
|
@ -667,14 +735,12 @@ class MCPClient:
|
|||
f"Server: {self.server_url or 'stdio'}, "
|
||||
f"Transport: {self.transport_type}"
|
||||
)
|
||||
|
||||
# Check if it's a stream/connection error
|
||||
if "BrokenResourceError" in error_type or "Broken" in error_type:
|
||||
verbose_logger.error(
|
||||
"MCP client detected broken connection/stream during list_resources - "
|
||||
"the MCP server may have crashed, disconnected, or timed out"
|
||||
)
|
||||
|
||||
# Return empty list instead of raising to allow graceful degradation
|
||||
return []
|
||||
|
||||
|
|
@ -709,14 +775,12 @@ class MCPClient:
|
|||
f"Server: {self.server_url or 'stdio'}, "
|
||||
f"Transport: {self.transport_type}"
|
||||
)
|
||||
|
||||
# Check if it's a stream/connection error
|
||||
if "BrokenResourceError" in error_type or "Broken" in error_type:
|
||||
verbose_logger.error(
|
||||
"MCP client detected broken connection/stream during list_resource_templates - "
|
||||
"the MCP server may have crashed, disconnected, or timed out"
|
||||
)
|
||||
|
||||
# Return empty list instead of raising to allow graceful degradation
|
||||
return []
|
||||
|
||||
|
|
@ -742,7 +806,6 @@ class MCPClient:
|
|||
|
||||
error_trace = traceback.format_exc()
|
||||
verbose_logger.debug(f"MCP client read_resource traceback:\n{error_trace}")
|
||||
|
||||
# Log detailed error information
|
||||
error_type = type(e).__name__
|
||||
verbose_logger.error(
|
||||
|
|
@ -753,12 +816,10 @@ class MCPClient:
|
|||
f"Server: {self.server_url or 'stdio'}, "
|
||||
f"Transport: {self.transport_type}"
|
||||
)
|
||||
|
||||
# Check if it's a stream/connection error
|
||||
if "BrokenResourceError" in error_type or "Broken" in error_type:
|
||||
verbose_logger.error(
|
||||
"MCP client detected broken connection/stream during read_resource - "
|
||||
"the MCP server may have crashed, disconnected, or timed out."
|
||||
)
|
||||
|
||||
raise
|
||||
|
|
|
|||
|
|
@ -18,6 +18,42 @@ else:
|
|||
GLOBAL_PASS_THROUGH_SUCCESS_HANDLER_OBJ = PassThroughEndpointLogging()
|
||||
|
||||
|
||||
def _encode_google_genai_sse_event(event_lines: List[str]) -> bytes:
|
||||
return ("\n".join(event_lines) + "\n\n").encode("utf-8")
|
||||
|
||||
|
||||
def _next_google_genai_sse_chunk(line_iter) -> bytes:
|
||||
event_lines: List[str] = []
|
||||
while True:
|
||||
try:
|
||||
line = next(line_iter)
|
||||
except StopIteration:
|
||||
if event_lines:
|
||||
return _encode_google_genai_sse_event(event_lines)
|
||||
raise
|
||||
if line == "":
|
||||
if event_lines:
|
||||
return _encode_google_genai_sse_event(event_lines)
|
||||
continue
|
||||
event_lines.append(line)
|
||||
|
||||
|
||||
async def _anext_google_genai_sse_chunk(line_iter) -> bytes:
|
||||
event_lines: List[str] = []
|
||||
while True:
|
||||
try:
|
||||
line = await line_iter.__anext__()
|
||||
except StopAsyncIteration:
|
||||
if event_lines:
|
||||
return _encode_google_genai_sse_event(event_lines)
|
||||
raise
|
||||
if line == "":
|
||||
if event_lines:
|
||||
return _encode_google_genai_sse_event(event_lines)
|
||||
continue
|
||||
event_lines.append(line)
|
||||
|
||||
|
||||
class BaseGoogleGenAIGenerateContentStreamingIterator:
|
||||
"""
|
||||
Base class for Google GenAI Generate Content streaming iterators that provides common logic
|
||||
|
|
@ -91,18 +127,17 @@ class GoogleGenAIGenerateContentStreamingIterator(
|
|||
self.generate_content_provider_config = generate_content_provider_config
|
||||
self.litellm_metadata = litellm_metadata
|
||||
self.custom_llm_provider = custom_llm_provider
|
||||
# Store the iterator once to avoid multiple stream consumption
|
||||
self.stream_iterator = response.iter_bytes()
|
||||
# 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()
|
||||
|
||||
def __iter__(self):
|
||||
return self
|
||||
|
||||
def __next__(self):
|
||||
try:
|
||||
# Get the next chunk from the stored iterator
|
||||
chunk = next(self.stream_iterator)
|
||||
chunk = _next_google_genai_sse_chunk(self.stream_iterator)
|
||||
self.collected_chunks.append(chunk)
|
||||
# Just yield raw bytes
|
||||
return chunk
|
||||
except StopIteration:
|
||||
raise StopIteration
|
||||
|
|
@ -147,18 +182,17 @@ class AsyncGoogleGenAIGenerateContentStreamingIterator(
|
|||
self.generate_content_provider_config = generate_content_provider_config
|
||||
self.litellm_metadata = litellm_metadata
|
||||
self.custom_llm_provider = custom_llm_provider
|
||||
# Store the async iterator once to avoid multiple stream consumption
|
||||
self.stream_iterator = response.aiter_bytes()
|
||||
# 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()
|
||||
|
||||
def __aiter__(self):
|
||||
return self
|
||||
|
||||
async def __anext__(self):
|
||||
try:
|
||||
# Get the next chunk from the stored async iterator
|
||||
chunk = await self.stream_iterator.__anext__()
|
||||
chunk = await _anext_google_genai_sse_chunk(self.stream_iterator)
|
||||
self.collected_chunks.append(chunk)
|
||||
# Just yield raw bytes
|
||||
return chunk
|
||||
except StopAsyncIteration:
|
||||
await self._handle_async_streaming_logging()
|
||||
|
|
|
|||
|
|
@ -1,7 +1,7 @@
|
|||
from abc import ABC, abstractmethod
|
||||
from typing import Literal
|
||||
|
||||
from litellm.proxy._types import CallInfo
|
||||
from litellm.proxy._types import CallInfo, Litellm_EntityType
|
||||
|
||||
|
||||
class BaseBudgetAlertType(ABC):
|
||||
|
|
@ -31,6 +31,8 @@ class SoftBudgetAlert(BaseBudgetAlertType):
|
|||
return "Soft Budget Crossed: "
|
||||
|
||||
def get_id(self, user_info: CallInfo) -> str:
|
||||
if user_info.event_group == Litellm_EntityType.TEAM:
|
||||
return user_info.team_id or "default_id"
|
||||
return user_info.token or "default_id"
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -37,6 +37,8 @@ from litellm.proxy._types import (
|
|||
VirtualKeyEvent,
|
||||
WebhookEvent,
|
||||
)
|
||||
from litellm.repositories.team_repository import TeamRepository
|
||||
from litellm.repositories.user_repository import UserRepository
|
||||
from litellm.types.integrations.slack_alerting import *
|
||||
|
||||
from ..email_templates.templates import *
|
||||
|
|
@ -1231,7 +1233,7 @@ Model Info:
|
|||
and recipient_user_id is not None
|
||||
and prisma_client is not None
|
||||
):
|
||||
user_row = await prisma_client.db.litellm_usertable.find_unique(
|
||||
user_row = await UserRepository(prisma_client).table.find_unique(
|
||||
where={"user_id": recipient_user_id}
|
||||
)
|
||||
|
||||
|
|
@ -1263,7 +1265,7 @@ Model Info:
|
|||
team_id = webhook_event.team_id
|
||||
team_name = "Default Team"
|
||||
if team_id is not None and prisma_client is not None:
|
||||
team_row = await prisma_client.db.litellm_teamtable.find_unique(
|
||||
team_row = await TeamRepository(prisma_client).table.find_unique(
|
||||
where={"team_id": team_id}
|
||||
)
|
||||
if team_row is not None:
|
||||
|
|
|
|||
|
|
@ -104,6 +104,51 @@
|
|||
},
|
||||
"description": "Datadog Custom Metrics Integration"
|
||||
},
|
||||
{
|
||||
"id": "galileo",
|
||||
"displayName": "Galileo",
|
||||
"logo": "galileo.ico",
|
||||
"supports_key_team_logging": false,
|
||||
"dynamic_params": {
|
||||
"GALILEO_API_KEY": {
|
||||
"type": "password",
|
||||
"ui_name": "API Key",
|
||||
"description": "Galileo Cloud API key (app.galileo.ai). Omit for enterprise username/password auth.",
|
||||
"required": false
|
||||
},
|
||||
"GALILEO_PROJECT_ID": {
|
||||
"type": "text",
|
||||
"ui_name": "Project ID",
|
||||
"description": "Galileo project ID to log traces to",
|
||||
"required": true
|
||||
},
|
||||
"GALILEO_LOG_STREAM_ID": {
|
||||
"type": "text",
|
||||
"ui_name": "Log Stream ID",
|
||||
"description": "Galileo log stream ID for v2 spans logging (optional)",
|
||||
"required": false
|
||||
},
|
||||
"GALILEO_BASE_URL": {
|
||||
"type": "text",
|
||||
"ui_name": "Base URL",
|
||||
"description": "Galileo API base URL (e.g. https://api.galileo.ai for Cloud, or your enterprise API URL)",
|
||||
"required": false
|
||||
},
|
||||
"GALILEO_USERNAME": {
|
||||
"type": "text",
|
||||
"ui_name": "Username",
|
||||
"description": "Galileo enterprise username (legacy Observe auth; use instead of API key)",
|
||||
"required": false
|
||||
},
|
||||
"GALILEO_PASSWORD": {
|
||||
"type": "password",
|
||||
"ui_name": "Password",
|
||||
"description": "Galileo enterprise password (legacy Observe auth)",
|
||||
"required": false
|
||||
}
|
||||
},
|
||||
"description": "Galileo AI Observability Integration"
|
||||
},
|
||||
{
|
||||
"id": "datadog_cost_management",
|
||||
"displayName": "Datadog Cost Management",
|
||||
|
|
@ -245,6 +290,21 @@
|
|||
},
|
||||
"description": "Langsmith Logging Integration"
|
||||
},
|
||||
{
|
||||
"id": "newrelic",
|
||||
"displayName": "New Relic",
|
||||
"logo": "newrelic.png",
|
||||
"supports_key_team_logging": false,
|
||||
"dynamic_params": {
|
||||
"NEW_RELIC_AI_MONITORING_RECORD_CONTENT_ENABLED": {
|
||||
"type": "text",
|
||||
"ui_name": "Record AI Content (default: true)",
|
||||
"description": "Whether to record AI message content. Set to false to disable.",
|
||||
"required": false
|
||||
}
|
||||
},
|
||||
"description": "New Relic AI Monitoring Integration"
|
||||
},
|
||||
{
|
||||
"id": "openmeter",
|
||||
"displayName": "OpenMeter",
|
||||
|
|
|
|||
|
|
@ -72,8 +72,13 @@ class CompressionInterceptionLogger(CustomLogger):
|
|||
compression_params: CompressionInterceptionConfig = {}
|
||||
if "compression_interception_params" in litellm_settings:
|
||||
compression_params = litellm_settings["compression_interception_params"]
|
||||
elif "compression_interception" in callback_specific_params:
|
||||
compression_params = callback_specific_params["compression_interception"]
|
||||
elif "compression_interception" in callback_specific_params and isinstance(
|
||||
callback_specific_params["compression_interception"], dict
|
||||
):
|
||||
compression_params = cast(
|
||||
CompressionInterceptionConfig,
|
||||
callback_specific_params["compression_interception"],
|
||||
)
|
||||
return CompressionInterceptionLogger.from_config_yaml(compression_params)
|
||||
|
||||
async def async_pre_call_deployment_hook(
|
||||
|
|
|
|||
|
|
@ -47,9 +47,29 @@ from litellm.exceptions import (
|
|||
BlockedPiiEntityError,
|
||||
GuardrailRaisedException,
|
||||
ModifyResponseException,
|
||||
SensitiveDataRouteException,
|
||||
)
|
||||
|
||||
|
||||
def get_session_id_from_request_data(request_data: Dict[str, Any]) -> Optional[str]:
|
||||
"""Extract session_id from request data (litellm_session_id or metadata)."""
|
||||
session_id = request_data.get("litellm_session_id")
|
||||
if session_id:
|
||||
return str(session_id)
|
||||
|
||||
metadata = request_data.get("metadata") or {}
|
||||
session_id = metadata.get("session_id")
|
||||
if session_id:
|
||||
return str(session_id)
|
||||
|
||||
litellm_metadata = request_data.get("litellm_metadata") or {}
|
||||
session_id = litellm_metadata.get("session_id")
|
||||
if session_id:
|
||||
return str(session_id)
|
||||
|
||||
return None
|
||||
|
||||
|
||||
class CustomGuardrail(CustomLogger):
|
||||
# If True, during_call runs async_moderation_hook instead of the unified apply_guardrail path.
|
||||
use_native_during_call_hook: ClassVar[bool] = False
|
||||
|
|
@ -68,6 +88,9 @@ class CustomGuardrail(CustomLogger):
|
|||
end_session_after_n_fails: Optional[int] = None,
|
||||
on_violation: Optional[str] = None,
|
||||
realtime_violation_message: Optional[str] = None,
|
||||
on_sensitive_data: Optional[str] = None,
|
||||
sensitive_data_route_to_model: Optional[str] = None,
|
||||
sticky_session_routing: bool = True,
|
||||
**kwargs,
|
||||
):
|
||||
"""
|
||||
|
|
@ -83,6 +106,9 @@ class CustomGuardrail(CustomLogger):
|
|||
end_session_after_n_fails: For /v1/realtime sessions, end the session after this many violations
|
||||
on_violation: For /v1/realtime sessions, 'warn' or 'end_session'
|
||||
realtime_violation_message: Message the bot speaks aloud when a /v1/realtime guardrail fires
|
||||
on_sensitive_data: Action when sensitive data is detected. 'block' (default) or 'route'
|
||||
sensitive_data_route_to_model: Model to route to when on_sensitive_data='route'
|
||||
sticky_session_routing: When True, all subsequent requests in the session use the same model
|
||||
"""
|
||||
self.guardrail_name = guardrail_name
|
||||
self.supported_event_hooks = supported_event_hooks
|
||||
|
|
@ -96,6 +122,11 @@ class CustomGuardrail(CustomLogger):
|
|||
self.end_session_after_n_fails: Optional[int] = end_session_after_n_fails
|
||||
self.on_violation: Optional[str] = on_violation
|
||||
self.realtime_violation_message: Optional[str] = realtime_violation_message
|
||||
self.on_sensitive_data: Optional[str] = on_sensitive_data
|
||||
self.sensitive_data_route_to_model: Optional[str] = (
|
||||
sensitive_data_route_to_model
|
||||
)
|
||||
self.sticky_session_routing: bool = sticky_session_routing
|
||||
|
||||
if supported_event_hooks:
|
||||
## validate event_hook is in supported_event_hooks
|
||||
|
|
@ -167,6 +198,108 @@ class CustomGuardrail(CustomLogger):
|
|||
detection_info=detection_info,
|
||||
)
|
||||
|
||||
def raise_sensitive_data_route_exception(
|
||||
self,
|
||||
route_to_model: str,
|
||||
request_data: Dict[str, Any],
|
||||
detection_info: Optional[Dict[str, Any]] = None,
|
||||
) -> None:
|
||||
"""
|
||||
Raise an exception to reroute the request to a different model.
|
||||
|
||||
Use this when sensitive data is detected and the guardrail is configured
|
||||
to route to an on-premise model instead of blocking.
|
||||
|
||||
The exception will reroute this request to the specified model. When
|
||||
sticky_session_routing is enabled (the default), it also stores the
|
||||
routing decision so subsequent requests in this session reuse the model.
|
||||
|
||||
Args:
|
||||
route_to_model: The model to route this request (and session) to
|
||||
request_data: The original request data dictionary
|
||||
detection_info: Optional non-sensitive detection metadata (e.g. matched
|
||||
entity types, rule ids, scores). This is surfaced in request metadata
|
||||
and logs, so it must not contain the raw detected sensitive values.
|
||||
|
||||
Raises:
|
||||
SensitiveDataRouteException: Always raises to trigger rerouting
|
||||
"""
|
||||
session_id = self._get_session_id_from_request_data(request_data)
|
||||
if not session_id:
|
||||
raise ValueError(
|
||||
"Cannot route sensitive data without a session_id. "
|
||||
"Ensure the request includes a session_id in metadata or headers."
|
||||
)
|
||||
|
||||
raise SensitiveDataRouteException(
|
||||
route_to_model=route_to_model,
|
||||
session_id=session_id,
|
||||
guardrail_name=self.guardrail_name,
|
||||
detection_info=detection_info,
|
||||
sticky_session_routing=self.sticky_session_routing,
|
||||
)
|
||||
|
||||
def _get_session_id_from_request_data(
|
||||
self, request_data: Dict[str, Any]
|
||||
) -> Optional[str]:
|
||||
"""Extract session_id from request data."""
|
||||
return get_session_id_from_request_data(request_data)
|
||||
|
||||
def should_route_on_sensitive_data(self) -> bool:
|
||||
"""
|
||||
Returns True if this guardrail is configured to route requests
|
||||
to a different model when sensitive data is detected.
|
||||
"""
|
||||
return (
|
||||
self.on_sensitive_data == "route"
|
||||
and self.sensitive_data_route_to_model is not None
|
||||
)
|
||||
|
||||
def handle_sensitive_data_detection(
|
||||
self,
|
||||
request_data: Dict[str, Any],
|
||||
detection_info: Optional[Dict[str, Any]] = None,
|
||||
) -> None:
|
||||
"""
|
||||
Handle sensitive data detection based on guardrail configuration.
|
||||
|
||||
If on_sensitive_data='route', raises SensitiveDataRouteException to reroute.
|
||||
Otherwise, raises GuardrailRaisedException to block. When routing is
|
||||
configured but the request carries no session_id, routing is not possible
|
||||
so the request falls back to a graceful block.
|
||||
|
||||
Args:
|
||||
request_data: The request data dictionary
|
||||
detection_info: Optional non-sensitive detection metadata. When routing,
|
||||
this is surfaced in request metadata and logs, so it must not contain
|
||||
the raw detected sensitive values.
|
||||
|
||||
Raises:
|
||||
SensitiveDataRouteException: When configured to route and a session_id is present
|
||||
GuardrailRaisedException: When configured to block, or when routing is
|
||||
configured but no session_id is available
|
||||
"""
|
||||
if self.should_route_on_sensitive_data():
|
||||
try:
|
||||
self.raise_sensitive_data_route_exception(
|
||||
route_to_model=self.sensitive_data_route_to_model, # type: ignore
|
||||
request_data=request_data,
|
||||
detection_info=detection_info,
|
||||
)
|
||||
except ValueError:
|
||||
raise GuardrailRaisedException(
|
||||
message=(
|
||||
f"Sensitive data detected by {self.guardrail_name} "
|
||||
"(routing skipped: request has no session_id)"
|
||||
),
|
||||
guardrail_name=self.guardrail_name,
|
||||
)
|
||||
else:
|
||||
raise GuardrailRaisedException(
|
||||
message=f"Sensitive data detected by {self.guardrail_name}",
|
||||
guardrail_name=self.guardrail_name,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def get_config_model() -> Optional[Type["GuardrailConfigModel"]]:
|
||||
"""
|
||||
|
|
@ -753,12 +886,20 @@ class CustomGuardrail(CustomLogger):
|
|||
Guardrails signal intentional blocks by raising:
|
||||
- GuardrailRaisedException (generic guardrail API, tool permission)
|
||||
- BlockedPiiEntityError (Presidio PII detection)
|
||||
- SensitiveDataRouteException (sensitive-data reroute to on-premise model)
|
||||
- HTTPException with status 400 (content policy violation)
|
||||
- ModifyResponseException (passthrough mode violation)
|
||||
"""
|
||||
if isinstance(e, ModifyResponseException):
|
||||
return True
|
||||
if isinstance(e, (GuardrailRaisedException, BlockedPiiEntityError)):
|
||||
if isinstance(
|
||||
e,
|
||||
(
|
||||
GuardrailRaisedException,
|
||||
BlockedPiiEntityError,
|
||||
SensitiveDataRouteException,
|
||||
),
|
||||
):
|
||||
return True
|
||||
if (
|
||||
HTTPException is not None
|
||||
|
|
|
|||
|
|
@ -7,6 +7,7 @@ from typing import List, Optional
|
|||
|
||||
from litellm._logging import verbose_logger, verbose_proxy_logger
|
||||
from litellm.proxy._types import WebhookEvent
|
||||
from litellm.repositories.team_repository import TeamRepository
|
||||
|
||||
# we use this for the email header, please send a test email if you change this. verify it looks good on email
|
||||
LITELLM_LOGO_URL = "https://litellm-listing.s3.amazonaws.com/litellm_logo.png"
|
||||
|
|
@ -24,7 +25,7 @@ async def get_all_team_member_emails(team_id: Optional[str] = None) -> list:
|
|||
if prisma_client is None:
|
||||
raise Exception("Not connected to DB!")
|
||||
|
||||
team_row = await prisma_client.db.litellm_teamtable.find_unique(
|
||||
team_row = await TeamRepository(prisma_client).table.find_unique(
|
||||
where={
|
||||
"team_id": team_id,
|
||||
}
|
||||
|
|
|
|||
|
|
@ -80,11 +80,15 @@ class FocusLiteLLMDatabase:
|
|||
vt.team_id,
|
||||
vt.key_alias as api_key_alias,
|
||||
tt.team_alias,
|
||||
ut.user_email as user_email
|
||||
ut.user_email as user_email,
|
||||
COALESCE(vt.organization_id, tt.organization_id) as organization_id,
|
||||
ot.organization_alias as organization_alias
|
||||
FROM "LiteLLM_DailyUserSpend" dus
|
||||
LEFT JOIN "LiteLLM_VerificationToken" vt ON dus.api_key = vt.token
|
||||
LEFT JOIN "LiteLLM_TeamTable" tt ON vt.team_id = tt.team_id
|
||||
LEFT JOIN "LiteLLM_UserTable" ut ON dus.user_id = ut.user_id
|
||||
LEFT JOIN "LiteLLM_OrganizationTable" ot
|
||||
ON ot.organization_id = COALESCE(vt.organization_id, tt.organization_id)
|
||||
{where_clause}
|
||||
ORDER BY dus.date DESC, dus.created_at DESC
|
||||
{limit_clause}
|
||||
|
|
|
|||
|
|
@ -2,13 +2,17 @@
|
|||
|
||||
from .base import FocusDestination, FocusTimeWindow
|
||||
from .factory import FocusDestinationFactory
|
||||
from .gcs_destination import FocusGCSDestination
|
||||
from .s3_destination import FocusS3Destination
|
||||
from .mavvrik_destination import FocusMavvrikDestination
|
||||
from .vantage_destination import FocusVantageDestination
|
||||
|
||||
__all__ = [
|
||||
"FocusDestination",
|
||||
"FocusDestinationFactory",
|
||||
"FocusGCSDestination",
|
||||
"FocusTimeWindow",
|
||||
"FocusS3Destination",
|
||||
"FocusMavvrikDestination",
|
||||
"FocusVantageDestination",
|
||||
]
|
||||
|
|
|
|||
|
|
@ -6,7 +6,9 @@ import os
|
|||
from typing import Any, Dict, Optional
|
||||
|
||||
from .base import FocusDestination
|
||||
from .gcs_destination import FocusGCSDestination
|
||||
from .s3_destination import FocusS3Destination
|
||||
from .mavvrik_destination import FocusMavvrikDestination
|
||||
from .vantage_destination import FocusVantageDestination
|
||||
|
||||
|
||||
|
|
@ -29,6 +31,10 @@ class FocusDestinationFactory:
|
|||
return FocusS3Destination(prefix=prefix, config=normalized_config)
|
||||
if provider_lower == "vantage":
|
||||
return FocusVantageDestination(prefix=prefix, config=normalized_config)
|
||||
if provider_lower == "gcs":
|
||||
return FocusGCSDestination(prefix=prefix, config=normalized_config)
|
||||
if provider_lower == "mavvrik":
|
||||
return FocusMavvrikDestination(prefix=prefix, config=normalized_config)
|
||||
raise NotImplementedError(
|
||||
f"Provider '{provider}' not supported for Focus export"
|
||||
)
|
||||
|
|
@ -72,6 +78,27 @@ class FocusDestinationFactory:
|
|||
"VANTAGE_INTEGRATION_TOKEN must be provided for Vantage exports"
|
||||
)
|
||||
return {k: v for k, v in resolved.items() if v is not None}
|
||||
if provider == "gcs":
|
||||
resolved = {
|
||||
"bucket_name": overrides.get("bucket_name")
|
||||
or os.getenv("FOCUS_GCS_BUCKET_NAME"),
|
||||
"service_account_json": overrides.get("service_account_json")
|
||||
or os.getenv("FOCUS_GCS_PATH_SERVICE_ACCOUNT"),
|
||||
}
|
||||
if not resolved.get("bucket_name"):
|
||||
raise ValueError(
|
||||
"FOCUS_GCS_BUCKET_NAME must be provided for GCS exports"
|
||||
)
|
||||
return {k: v for k, v in resolved.items() if v is not None}
|
||||
if provider == "mavvrik":
|
||||
resolved = {
|
||||
"api_key": overrides.get("api_key") or os.getenv("MAVVRIK_API_KEY"),
|
||||
"api_endpoint": overrides.get("api_endpoint")
|
||||
or os.getenv("MAVVRIK_API_ENDPOINT"),
|
||||
"connection_id": overrides.get("connection_id")
|
||||
or os.getenv("MAVVRIK_CONNECTION_ID"),
|
||||
}
|
||||
return {k: v for k, v in resolved.items() if v is not None}
|
||||
raise NotImplementedError(
|
||||
f"Provider '{provider}' not supported for Focus export configuration"
|
||||
)
|
||||
|
|
|
|||
74
litellm/integrations/focus/destinations/gcs_destination.py
Normal file
74
litellm/integrations/focus/destinations/gcs_destination.py
Normal file
|
|
@ -0,0 +1,74 @@
|
|||
"""GCS destination for Focus export — reuses GCSBucketBase auth and httpx client."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from datetime import timezone
|
||||
from typing import Any, Optional
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.integrations.gcs_bucket.gcs_bucket_base import GCSBucketBase
|
||||
from litellm.litellm_core_utils.cloud_storage_security import (
|
||||
encode_gcs_object_name_for_url,
|
||||
)
|
||||
|
||||
from .base import FocusDestination, FocusTimeWindow
|
||||
|
||||
|
||||
class FocusGCSDestination(GCSBucketBase, FocusDestination):
|
||||
"""Upload serialized Focus exports to GCS using the GCS JSON API."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
prefix: str,
|
||||
config: Optional[dict[str, Any]] = None,
|
||||
) -> None:
|
||||
config = config or {}
|
||||
bucket_name = config.get("bucket_name")
|
||||
if not bucket_name:
|
||||
raise ValueError("bucket_name must be provided for GCS destination")
|
||||
super().__init__(bucket_name=bucket_name)
|
||||
service_account_json = config.get("service_account_json")
|
||||
if service_account_json is not None:
|
||||
self.path_service_account_json = service_account_json
|
||||
self.prefix = prefix.rstrip("/")
|
||||
|
||||
async def deliver(
|
||||
self,
|
||||
*,
|
||||
content: bytes,
|
||||
time_window: FocusTimeWindow,
|
||||
filename: str,
|
||||
) -> None:
|
||||
object_name = self._build_object_key(time_window=time_window, filename=filename)
|
||||
headers = await self.construct_request_headers(
|
||||
service_account_json=self.path_service_account_json
|
||||
)
|
||||
headers["Content-Type"] = "application/octet-stream"
|
||||
encoded_name = encode_gcs_object_name_for_url(object_name)
|
||||
url = (
|
||||
f"https://storage.googleapis.com/upload/storage/v1/b/"
|
||||
f"{self.BUCKET_NAME}/o?uploadType=media&name={encoded_name}"
|
||||
)
|
||||
response = await self.async_httpx_client.post(
|
||||
url=url, headers=headers, data=content
|
||||
)
|
||||
if response.status_code != 200:
|
||||
raise RuntimeError(
|
||||
f"GCS upload failed: status={response.status_code} body={response.text}"
|
||||
)
|
||||
verbose_logger.debug(
|
||||
"Focus GCS: uploaded %d bytes to gs://%s/%s",
|
||||
len(content),
|
||||
self.BUCKET_NAME,
|
||||
object_name,
|
||||
)
|
||||
|
||||
def _build_object_key(self, *, time_window: FocusTimeWindow, filename: str) -> str:
|
||||
start_utc = time_window.start_time.astimezone(timezone.utc)
|
||||
date_component = f"date={start_utc.strftime('%Y-%m-%d')}"
|
||||
parts = [self.prefix, date_component]
|
||||
if time_window.frequency == "hourly":
|
||||
parts.append(f"hour={start_utc.strftime('%H')}")
|
||||
key_prefix = "/".join(filter(None, parts))
|
||||
return f"{key_prefix}/{filename}" if key_prefix else filename
|
||||
345
litellm/integrations/focus/destinations/mavvrik_destination.py
Normal file
345
litellm/integrations/focus/destinations/mavvrik_destination.py
Normal file
|
|
@ -0,0 +1,345 @@
|
|||
"""Mavvrik GCS destination for FOCUS export.
|
||||
|
||||
Flow:
|
||||
1. GET /metrics/agent/ai/{connection_id}/upload-url → GCS signed URL
|
||||
2. PUT <signed_url> with CSV content
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import gzip
|
||||
from typing import Any, Optional
|
||||
from urllib.parse import urlparse
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.llms.custom_httpx.http_handler import (
|
||||
AsyncHTTPHandler,
|
||||
get_async_httpx_client,
|
||||
httpxSpecialProvider,
|
||||
)
|
||||
|
||||
from .base import FocusDestination, FocusTimeWindow
|
||||
|
||||
_MAVVRIK_ALLOWED_SUFFIXES = (".mavvrik.dev", ".mavvrik.ai", ".mavvrik.app")
|
||||
|
||||
# GCS requires intermediate chunks to be a multiple of 256 KB.
|
||||
# 8 MB gives a good balance between round-trips and memory pressure.
|
||||
_GCS_CHUNK_SIZE = 8 * 1024 * 1024 # 8 MB
|
||||
|
||||
|
||||
def _validate_api_endpoint(api_endpoint: str) -> None:
|
||||
if not api_endpoint.startswith("https://"):
|
||||
raise ValueError("MAVVRIK_API_ENDPOINT must be an HTTPS URL")
|
||||
hostname = (urlparse(api_endpoint).hostname or "").lower()
|
||||
if not any(hostname.endswith(suffix) for suffix in _MAVVRIK_ALLOWED_SUFFIXES):
|
||||
raise ValueError(
|
||||
"MAVVRIK_API_ENDPOINT host must be a Mavvrik domain "
|
||||
"(e.g. https://api.mavvrik.dev/<tenant_id>)"
|
||||
)
|
||||
|
||||
|
||||
def _validate_gcs_url(url: str, label: str) -> None:
|
||||
parsed = urlparse(url)
|
||||
if parsed.scheme != "https":
|
||||
raise ValueError(
|
||||
f"Mavvrik FOCUS destination: {label} must be HTTPS, got scheme '{parsed.scheme}'"
|
||||
)
|
||||
hostname = (parsed.hostname or "").lower()
|
||||
if not (
|
||||
hostname == "storage.googleapis.com"
|
||||
or hostname.endswith(".storage.googleapis.com")
|
||||
):
|
||||
raise ValueError(
|
||||
f"Mavvrik FOCUS destination: {label} must be a GCS endpoint "
|
||||
f"(storage.googleapis.com), got '{hostname}'"
|
||||
)
|
||||
|
||||
|
||||
class FocusMavvrikDestination(FocusDestination):
|
||||
"""Upload FOCUS CSV exports to Mavvrik via GCS signed URL."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
prefix: str,
|
||||
config: Optional[dict[str, Any]] = None,
|
||||
) -> None:
|
||||
config = config or {}
|
||||
api_key = config.get("api_key")
|
||||
api_endpoint = config.get("api_endpoint")
|
||||
connection_id = config.get("connection_id")
|
||||
|
||||
if not api_key:
|
||||
raise ValueError(
|
||||
"MAVVRIK_API_KEY must be provided for Mavvrik FOCUS destination "
|
||||
"(set MAVVRIK_API_KEY env var or pass in destination_config)"
|
||||
)
|
||||
if not api_endpoint:
|
||||
raise ValueError(
|
||||
"MAVVRIK_API_ENDPOINT must be provided for Mavvrik FOCUS destination "
|
||||
"(set MAVVRIK_API_ENDPOINT env var or pass in destination_config)"
|
||||
)
|
||||
if not connection_id:
|
||||
raise ValueError(
|
||||
"MAVVRIK_CONNECTION_ID must be provided for Mavvrik FOCUS destination "
|
||||
"(set MAVVRIK_CONNECTION_ID env var or pass in destination_config)"
|
||||
)
|
||||
|
||||
_validate_api_endpoint(api_endpoint)
|
||||
|
||||
self.api_key = api_key
|
||||
self.api_endpoint = api_endpoint.rstrip("/")
|
||||
self.connection_id = connection_id
|
||||
self.prefix = prefix
|
||||
self._http: AsyncHTTPHandler = get_async_httpx_client(
|
||||
llm_provider=httpxSpecialProvider.LoggingCallback
|
||||
)
|
||||
self._registered = False
|
||||
|
||||
@property
|
||||
def _agent_url(self) -> str:
|
||||
return f"{self.api_endpoint}/metrics/agent/ai/{self.connection_id}"
|
||||
|
||||
@property
|
||||
def _upload_url_endpoint(self) -> str:
|
||||
return f"{self.api_endpoint}/metrics/agent/ai/{self.connection_id}/upload-url"
|
||||
|
||||
@property
|
||||
def _auth_headers(self) -> dict[str, str]:
|
||||
return {"Content-Type": "application/json", "x-api-key": self.api_key}
|
||||
|
||||
async def _ensure_registered(self) -> Optional[int]:
|
||||
"""POST agent endpoint to register/initialize the connector (once per instance).
|
||||
|
||||
Returns metricsMarker from the Mavvrik response — the last date index
|
||||
Mavvrik has successfully processed. Used by the logger to catch up any
|
||||
dates that were missed due to previous export failures.
|
||||
|
||||
Returns None if the connector was already registered (cached).
|
||||
"""
|
||||
if self._registered:
|
||||
return None
|
||||
resp = await self._http.client.request(
|
||||
method="POST",
|
||||
url=self._agent_url,
|
||||
headers=self._auth_headers,
|
||||
json={"name": self.connection_id},
|
||||
timeout=30.0,
|
||||
)
|
||||
if resp.status_code == 410:
|
||||
# Connector has been disconnected in Mavvrik — reset flag so next
|
||||
# delivery attempt re-registers after it becomes active again.
|
||||
self._registered = False
|
||||
raise RuntimeError(
|
||||
"Mavvrik FOCUS destination: connector is disconnected (410). "
|
||||
"Re-enable the connection in the Mavvrik dashboard."
|
||||
)
|
||||
if resp.status_code >= 400:
|
||||
raise RuntimeError(
|
||||
f"Mavvrik FOCUS destination: register failed "
|
||||
f"({resp.status_code}): {resp.text[:200]}"
|
||||
)
|
||||
self._registered = True
|
||||
metrics_marker = resp.json().get("metricsMarker", 0)
|
||||
verbose_logger.debug(
|
||||
"Mavvrik FOCUS destination: connector registered (metricsMarker=%s)",
|
||||
metrics_marker,
|
||||
)
|
||||
return metrics_marker
|
||||
|
||||
async def _get_signed_url(self, date_str: str) -> str:
|
||||
"""GET upload-url endpoint → GCS signed URL for the given date."""
|
||||
params = {"name": date_str, "type": "metrics", "datetime": date_str}
|
||||
resp = await self._http.client.request(
|
||||
method="GET",
|
||||
url=self._upload_url_endpoint,
|
||||
headers=self._auth_headers,
|
||||
params=params,
|
||||
timeout=30.0,
|
||||
)
|
||||
if resp.status_code >= 400:
|
||||
raise RuntimeError(
|
||||
f"Mavvrik FOCUS destination: failed to get signed URL "
|
||||
f"({resp.status_code}): {resp.text[:200]}"
|
||||
)
|
||||
signed_url = resp.json().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")
|
||||
verbose_logger.debug(
|
||||
"Mavvrik FOCUS destination: got signed URL for date %s", date_str
|
||||
)
|
||||
return signed_url
|
||||
|
||||
async def _upload_to_gcs(self, signed_url: str, content: bytes) -> None:
|
||||
"""Upload gzip-compressed CSV to GCS via chunked resumable upload.
|
||||
|
||||
The full CSV is gzip-compressed first, then uploaded in _GCS_CHUNK_SIZE
|
||||
chunks using the GCS resumable upload protocol. GCS assembles the chunks
|
||||
server-side into a single complete object — the bucket receives one file
|
||||
regardless of how many chunks were sent.
|
||||
|
||||
Intermediate chunks: Content-Range: bytes X-Y/* → expect 308
|
||||
Final chunk: Content-Range: bytes X-Y/T → expect 200/201
|
||||
|
||||
This handles exports larger than available memory for a single PUT while
|
||||
keeping the destination code self-contained (no changes to the FOCUS
|
||||
pipeline upstream).
|
||||
"""
|
||||
gzip_bytes = gzip.compress(content)
|
||||
total = len(gzip_bytes)
|
||||
|
||||
# Step 1: initiate resumable upload session
|
||||
metadata = b'{"contentEncoding":"gzip","contentDisposition":"attachment"}'
|
||||
init_resp = await self._http.client.request(
|
||||
method="POST",
|
||||
url=signed_url,
|
||||
headers={
|
||||
"Content-Type": "application/gzip",
|
||||
"x-goog-resumable": "start",
|
||||
},
|
||||
content=metadata,
|
||||
timeout=30.0,
|
||||
)
|
||||
if init_resp.status_code not in (200, 201):
|
||||
raise RuntimeError(
|
||||
f"Mavvrik FOCUS destination: GCS session init failed "
|
||||
f"({init_resp.status_code}): {init_resp.text[:400]}"
|
||||
)
|
||||
|
||||
session_uri = init_resp.headers.get("Location")
|
||||
if not session_uri:
|
||||
raise RuntimeError(
|
||||
"Mavvrik FOCUS destination: GCS session init missing Location header"
|
||||
)
|
||||
_validate_gcs_url(session_uri, "session URI")
|
||||
|
||||
verbose_logger.debug(
|
||||
"Mavvrik FOCUS destination: GCS session started, uploading %d gzip bytes "
|
||||
"in %d chunk(s)",
|
||||
total,
|
||||
max(1, -(-total // _GCS_CHUNK_SIZE)), # ceiling division
|
||||
)
|
||||
|
||||
# Step 2: upload in chunks; cancel session on any failure to avoid
|
||||
# lingering GCS sessions (they stay open for ~1 week otherwise).
|
||||
offset = 0
|
||||
try:
|
||||
while offset < total:
|
||||
chunk = gzip_bytes[offset : offset + _GCS_CHUNK_SIZE]
|
||||
chunk_end = offset + len(chunk) - 1
|
||||
is_final = (offset + len(chunk)) >= total
|
||||
content_range = (
|
||||
f"bytes {offset}-{chunk_end}/{total}"
|
||||
if is_final
|
||||
else f"bytes {offset}-{chunk_end}/*"
|
||||
)
|
||||
expected_statuses = {200, 201} if is_final else {308}
|
||||
|
||||
resp = await self._http.client.request(
|
||||
method="PUT",
|
||||
url=session_uri,
|
||||
headers={
|
||||
"Content-Type": "application/gzip",
|
||||
"Content-Range": content_range,
|
||||
},
|
||||
content=chunk,
|
||||
timeout=120.0,
|
||||
)
|
||||
if resp.status_code not in expected_statuses:
|
||||
raise RuntimeError(
|
||||
f"Mavvrik FOCUS destination: GCS chunk upload failed "
|
||||
f"(chunk offset={offset}, expected={expected_statuses}, "
|
||||
f"got={resp.status_code}): {resp.text[:400]}"
|
||||
)
|
||||
offset += len(chunk)
|
||||
verbose_logger.debug(
|
||||
"Mavvrik FOCUS destination: uploaded chunk offset=%d/%d",
|
||||
offset,
|
||||
total,
|
||||
)
|
||||
except Exception:
|
||||
# Cancel the open GCS session so it doesn't linger for up to 1 week.
|
||||
try:
|
||||
await self._http.client.request(
|
||||
method="DELETE", url=session_uri, timeout=10.0
|
||||
)
|
||||
verbose_logger.debug(
|
||||
"Mavvrik FOCUS destination: cancelled GCS session after error"
|
||||
)
|
||||
except Exception:
|
||||
pass
|
||||
raise
|
||||
|
||||
async def get_metrics_marker(self) -> Optional[int]:
|
||||
"""Register with Mavvrik and return the current metricsMarker.
|
||||
|
||||
The metricsMarker is a Unix timestamp (seconds) representing the last
|
||||
date Mavvrik has successfully ingested. Called on every scheduled run
|
||||
so the logger can detect and catch up any dates missed due to previous
|
||||
export failures.
|
||||
|
||||
Always calls the Mavvrik register API — unlike deliver() which skips
|
||||
registration once _registered is True, catch-up requires a fresh
|
||||
marker value on every run.
|
||||
"""
|
||||
resp = await self._http.client.request(
|
||||
method="POST",
|
||||
url=self._agent_url,
|
||||
headers=self._auth_headers,
|
||||
json={"name": self.connection_id},
|
||||
timeout=30.0,
|
||||
)
|
||||
if resp.status_code == 410:
|
||||
self._registered = False
|
||||
raise RuntimeError(
|
||||
"Mavvrik FOCUS destination: connector is disconnected (410). "
|
||||
"Re-enable the connection in the Mavvrik dashboard."
|
||||
)
|
||||
if resp.status_code >= 400:
|
||||
raise RuntimeError(
|
||||
f"Mavvrik FOCUS destination: register failed "
|
||||
f"({resp.status_code}): {resp.text[:200]}"
|
||||
)
|
||||
self._registered = True
|
||||
metrics_marker = resp.json().get("metricsMarker", 0)
|
||||
verbose_logger.debug(
|
||||
"Mavvrik FOCUS destination: got metricsMarker=%s", metrics_marker
|
||||
)
|
||||
return metrics_marker
|
||||
|
||||
async def deliver(
|
||||
self,
|
||||
*,
|
||||
content: bytes,
|
||||
time_window: FocusTimeWindow,
|
||||
filename: str,
|
||||
) -> None:
|
||||
"""Upload FOCUS CSV to Mavvrik via GCS signed URL.
|
||||
|
||||
Uses the start date of the time window as the object date key.
|
||||
"""
|
||||
if not content:
|
||||
verbose_logger.debug(
|
||||
"Mavvrik FOCUS destination: empty content, skipping upload"
|
||||
)
|
||||
return
|
||||
|
||||
date_str = time_window.start_time.strftime("%Y-%m-%d")
|
||||
|
||||
verbose_logger.debug(
|
||||
"Mavvrik FOCUS destination: uploading %d bytes for date=%s (%s)",
|
||||
len(content),
|
||||
date_str,
|
||||
filename,
|
||||
)
|
||||
|
||||
await self._ensure_registered()
|
||||
signed_url = await self._get_signed_url(date_str)
|
||||
await self._upload_to_gcs(signed_url, content)
|
||||
|
||||
verbose_logger.debug(
|
||||
"Mavvrik FOCUS destination: upload complete for date=%s", date_str
|
||||
)
|
||||
|
|
@ -12,6 +12,8 @@ from .schema import FOCUS_NORMALIZED_SCHEMA
|
|||
_TAG_KEYS = (
|
||||
"team_id",
|
||||
"team_alias",
|
||||
"organization_id",
|
||||
"organization_alias",
|
||||
"user_id",
|
||||
"user_email",
|
||||
"api_key_alias",
|
||||
|
|
|
|||
|
|
@ -1,8 +1,13 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
from typing import Any, Dict, List, Optional, Tuple, cast
|
||||
import uuid
|
||||
from datetime import datetime, timezone
|
||||
from typing import Any, Dict, List, Optional, Tuple, Union, cast
|
||||
|
||||
import httpx
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
import litellm
|
||||
|
|
@ -12,11 +17,16 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import (
|
|||
convert_content_list_to_str,
|
||||
get_content_from_model_response,
|
||||
)
|
||||
from litellm.types.llms.openai import (
|
||||
AllMessageValues,
|
||||
HttpxBinaryResponseContent,
|
||||
ResponsesAPIResponse,
|
||||
)
|
||||
from litellm.llms.custom_httpx.http_handler import (
|
||||
get_async_httpx_client,
|
||||
httpxSpecialProvider,
|
||||
)
|
||||
from litellm.types.llms.openai import AllMessageValues
|
||||
from litellm.types.integrations.base_health_check import IntegrationHealthCheckStatus
|
||||
|
||||
GALILEO_CLOUD_API_BASE_URL = "https://api.galileo.ai"
|
||||
# Cap the in-memory buffer so persistent flush failures (e.g. Galileo
|
||||
|
|
@ -33,6 +43,11 @@ class LLMResponse(BaseModel):
|
|||
model: str
|
||||
num_input_tokens: int
|
||||
num_output_tokens: int
|
||||
num_total_tokens: int
|
||||
cost: Optional[float] = Field(
|
||||
default=None,
|
||||
description="Total cost of the LLM call in USD as computed by LiteLLM.",
|
||||
)
|
||||
output_logprobs: Optional[Dict[str, Any]] = Field(
|
||||
default=None,
|
||||
description="Optional. When available, logprobs are used to compute Uncertainty.",
|
||||
|
|
@ -75,6 +90,52 @@ class GalileoObserve(CustomLogger):
|
|||
return bool(self.api_key)
|
||||
return bool(self.username and self.password)
|
||||
|
||||
async def async_health_check(self) -> IntegrationHealthCheckStatus:
|
||||
try:
|
||||
if not self.project_id:
|
||||
return IntegrationHealthCheckStatus(
|
||||
status="unhealthy",
|
||||
error_message="GALILEO_PROJECT_ID environment variable not set",
|
||||
)
|
||||
|
||||
if not self.base_url:
|
||||
return IntegrationHealthCheckStatus(
|
||||
status="unhealthy",
|
||||
error_message="GALILEO_BASE_URL environment variable not set",
|
||||
)
|
||||
|
||||
if not self.use_v2_api and (not self.username or not self.password):
|
||||
return IntegrationHealthCheckStatus(
|
||||
status="unhealthy",
|
||||
error_message=(
|
||||
"GALILEO_API_KEY or GALILEO_USERNAME and GALILEO_PASSWORD "
|
||||
"environment variables must be set"
|
||||
),
|
||||
)
|
||||
|
||||
if not await self._ensure_headers():
|
||||
return IntegrationHealthCheckStatus(
|
||||
status="unhealthy",
|
||||
error_message="Galileo authentication failed",
|
||||
)
|
||||
|
||||
response = await self.async_httpx_handler.get(
|
||||
url=f"{self.base_url}/current_user",
|
||||
headers=self.headers,
|
||||
)
|
||||
if response.status_code >= 400:
|
||||
return IntegrationHealthCheckStatus(
|
||||
status="unhealthy",
|
||||
error_message=(f"Galileo API returned HTTP {response.status_code}"),
|
||||
)
|
||||
|
||||
return IntegrationHealthCheckStatus(status="healthy", error_message=None)
|
||||
except Exception as e:
|
||||
return IntegrationHealthCheckStatus(
|
||||
status="unhealthy",
|
||||
error_message=f"Galileo health check failed: {str(e)}",
|
||||
)
|
||||
|
||||
async def async_set_galileo_headers(self) -> None:
|
||||
galileo_login_response = await self.async_httpx_handler.post(
|
||||
url=f"{self.base_url}/login",
|
||||
|
|
@ -121,10 +182,14 @@ class GalileoObserve(CustomLogger):
|
|||
|
||||
@staticmethod
|
||||
def _galileo_input_messages(
|
||||
messages: Optional[List[Any]], input_text: str
|
||||
messages: Optional[Any], input_text: str
|
||||
) -> List[Dict[str, str]]:
|
||||
if isinstance(messages, dict):
|
||||
messages = messages.get("messages")
|
||||
if not messages:
|
||||
return [{"role": "user", "content": input_text}]
|
||||
if not isinstance(messages, list):
|
||||
return [{"role": "user", "content": input_text}]
|
||||
|
||||
galileo_messages: List[Dict[str, str]] = []
|
||||
for message in messages:
|
||||
|
|
@ -147,13 +212,59 @@ class GalileoObserve(CustomLogger):
|
|||
return [{"role": "user", "content": input_text}]
|
||||
|
||||
@staticmethod
|
||||
def _record_to_v2_span(record: Dict[str, Any]) -> Dict[str, Any]:
|
||||
created_at = record.get("created_at", "")
|
||||
def _local_timezone():
|
||||
return datetime.now().astimezone().tzinfo or timezone.utc
|
||||
|
||||
@staticmethod
|
||||
def _format_created_at(dt: Union[datetime, Any]) -> str:
|
||||
"""Serialize timestamps as UTC ISO-8601 for Galileo."""
|
||||
if not isinstance(dt, datetime):
|
||||
return str(dt)
|
||||
|
||||
if dt.tzinfo is None:
|
||||
# LiteLLM often passes naive datetimes in local time; convert to UTC
|
||||
# instead of appending Z to local time (which shifts Traces tab sorting).
|
||||
dt = dt.replace(tzinfo=GalileoObserve._local_timezone())
|
||||
|
||||
return dt.astimezone(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")
|
||||
|
||||
@staticmethod
|
||||
def _normalize_created_at(created_at: str) -> str:
|
||||
if created_at and not re.search(r"(Z|[+-]\d{2}:?\d{2})$", created_at):
|
||||
created_at = f"{created_at}Z"
|
||||
return f"{created_at}Z"
|
||||
return created_at
|
||||
|
||||
@staticmethod
|
||||
def _token_metrics_from_record(record: Dict[str, Any]) -> Dict[str, Any]:
|
||||
num_input_tokens = int(record.get("num_input_tokens") or 0)
|
||||
num_output_tokens = int(record.get("num_output_tokens") or 0)
|
||||
num_total_tokens = int(record.get("num_total_tokens") or 0)
|
||||
if num_total_tokens == 0 and (num_input_tokens or num_output_tokens):
|
||||
num_total_tokens = num_input_tokens + num_output_tokens
|
||||
metrics: Dict[str, Any] = {
|
||||
"num_input_tokens": num_input_tokens,
|
||||
"num_output_tokens": num_output_tokens,
|
||||
"num_total_tokens": num_total_tokens,
|
||||
}
|
||||
cost = record.get("cost")
|
||||
if cost is not None:
|
||||
metrics["cost"] = float(cost)
|
||||
return metrics
|
||||
|
||||
@staticmethod
|
||||
def _record_to_v2_span(
|
||||
record: Dict[str, Any],
|
||||
*,
|
||||
trace_id: str,
|
||||
span_id: str,
|
||||
) -> Dict[str, Any]:
|
||||
created_at = GalileoObserve._normalize_created_at(record.get("created_at", ""))
|
||||
|
||||
span: Dict[str, Any] = {
|
||||
"type": "llm",
|
||||
"id": span_id,
|
||||
"trace_id": trace_id,
|
||||
"parent_id": trace_id,
|
||||
"name": record.get("node_type", "litellm"),
|
||||
"created_at": created_at,
|
||||
"input": GalileoObserve._galileo_input_messages(
|
||||
|
|
@ -167,14 +278,49 @@ class GalileoObserve(CustomLogger):
|
|||
"model": record.get("model"),
|
||||
"metrics": {
|
||||
"duration_ns": int(record.get("latency_ms", 0)) * 1_000_000,
|
||||
"num_input_tokens": record.get("num_input_tokens"),
|
||||
"num_output_tokens": record.get("num_output_tokens"),
|
||||
**GalileoObserve._token_metrics_from_record(record),
|
||||
},
|
||||
}
|
||||
if record.get("tags"):
|
||||
span["tags"] = record["tags"]
|
||||
return span
|
||||
|
||||
@staticmethod
|
||||
def _record_to_v2_trace(record: Dict[str, Any]) -> Dict[str, Any]:
|
||||
trace_id = str(uuid.uuid4())
|
||||
span_id = str(uuid.uuid4())
|
||||
created_at = GalileoObserve._normalize_created_at(record.get("created_at", ""))
|
||||
|
||||
return {
|
||||
"type": "trace",
|
||||
"id": trace_id,
|
||||
"name": record.get("node_type", "litellm"),
|
||||
"created_at": created_at,
|
||||
"input": record.get("input_text", ""),
|
||||
"output": record.get("output_text", ""),
|
||||
"status_code": record.get("status_code", 200),
|
||||
"metrics": {
|
||||
"duration_ns": int(record.get("latency_ms", 0)) * 1_000_000,
|
||||
**GalileoObserve._token_metrics_from_record(record),
|
||||
},
|
||||
"spans": [
|
||||
GalileoObserve._record_to_v2_span(
|
||||
record, trace_id=trace_id, span_id=span_id
|
||||
)
|
||||
],
|
||||
}
|
||||
|
||||
def _build_traces_payload(self, records: List[dict]) -> Dict[str, Any]:
|
||||
payload: Dict[str, Any] = {
|
||||
"traces": [self._record_to_v2_trace(record) for record in records],
|
||||
"logging_method": "api_direct",
|
||||
"reliable": False,
|
||||
"is_complete": True,
|
||||
}
|
||||
if self.log_stream_id:
|
||||
payload["log_stream_id"] = self.log_stream_id
|
||||
return payload
|
||||
|
||||
def _get_ingest_request(self) -> Optional[Tuple[str, Dict[str, Any]]]:
|
||||
if not self.base_url or not self.project_id:
|
||||
return None
|
||||
|
|
@ -184,105 +330,453 @@ class GalileoObserve(CustomLogger):
|
|||
# flush_in_memory_records) aren't silently dropped when we later clear
|
||||
# the in-memory buffer.
|
||||
records = list(self.in_memory_records)
|
||||
payload = self._build_traces_payload(records)
|
||||
|
||||
if self.use_v2_api:
|
||||
payload: Dict[str, Any] = {
|
||||
"spans": [self._record_to_v2_span(record) for record in records],
|
||||
"reliable": False,
|
||||
}
|
||||
if self.log_stream_id:
|
||||
payload["log_stream_id"] = self.log_stream_id
|
||||
return (
|
||||
f"{self.base_url}/v2/projects/{self.project_id}/spans",
|
||||
f"{self.base_url}/ingest/traces/{self.project_id}",
|
||||
payload,
|
||||
)
|
||||
|
||||
# Username/password auth logs in for a JWT and uses the standard v2 traces API.
|
||||
return (
|
||||
f"{self.base_url}/projects/{self.project_id}/observe/ingest",
|
||||
{"records": records},
|
||||
f"{self.base_url}/v2/projects/{self.project_id}/traces",
|
||||
payload,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _redact_headers(headers: Optional[Dict[str, str]]) -> Dict[str, str]:
|
||||
if not headers:
|
||||
return {}
|
||||
redacted: Dict[str, str] = {}
|
||||
for key, value in headers.items():
|
||||
if key.lower() in {"authorization", "galileo-api-key"} and value:
|
||||
redacted[key] = (
|
||||
f"{value[:8]}...{value[-4:]}" if len(value) > 12 else "***"
|
||||
)
|
||||
else:
|
||||
redacted[key] = value
|
||||
return redacted
|
||||
|
||||
def _log_flush_config(self) -> None:
|
||||
verbose_logger.debug(
|
||||
"Galileo Logger flush config: use_v2_api=%s base_url=%s project_id=%s "
|
||||
"log_stream_id=%s api_key_set=%s username_set=%s record_count=%s",
|
||||
self.use_v2_api,
|
||||
self.base_url,
|
||||
self.project_id,
|
||||
self.log_stream_id,
|
||||
bool(self.api_key),
|
||||
bool(self.username),
|
||||
len(self.in_memory_records),
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _log_v2_payload_validation(payload: Dict[str, Any]) -> None:
|
||||
missing_fields: List[str] = []
|
||||
traces = payload.get("traces", [])
|
||||
if not traces:
|
||||
missing_fields.append("traces")
|
||||
|
||||
for trace_index, trace in enumerate(traces):
|
||||
if not isinstance(trace, dict):
|
||||
continue
|
||||
for field in ("id", "type", "spans"):
|
||||
if field not in trace:
|
||||
missing_fields.append(f"traces[{trace_index}].{field}")
|
||||
|
||||
trace_id = trace.get("id")
|
||||
for span_index, span in enumerate(trace.get("spans", [])):
|
||||
if not isinstance(span, dict):
|
||||
continue
|
||||
for field in ("id", "trace_id", "parent_id"):
|
||||
if field not in span:
|
||||
missing_fields.append(
|
||||
f"traces[{trace_index}].spans[{span_index}].{field}"
|
||||
)
|
||||
if trace_id and span.get("trace_id") != trace_id:
|
||||
missing_fields.append(
|
||||
f"traces[{trace_index}].spans[{span_index}].trace_id mismatch"
|
||||
)
|
||||
|
||||
if missing_fields:
|
||||
verbose_logger.debug(
|
||||
"Galileo Logger: ingest /traces payload validation issues: %s",
|
||||
missing_fields,
|
||||
)
|
||||
|
||||
def _log_flush_payload(self, url: str, payload: Dict[str, Any]) -> None:
|
||||
traces = payload.get("traces", [])
|
||||
verbose_logger.debug(
|
||||
"Galileo Logger flush URL: %s trace_count=%s",
|
||||
url,
|
||||
len(traces) if isinstance(traces, list) else 0,
|
||||
)
|
||||
if self.use_v2_api and "/ingest/traces/" in url:
|
||||
self._log_v2_payload_validation(payload)
|
||||
|
||||
@staticmethod
|
||||
def _log_http_status_error(error: httpx.HTTPStatusError, url: str) -> None:
|
||||
response = error.response
|
||||
verbose_logger.debug(
|
||||
"Galileo Logger HTTP error: status=%s url=%s",
|
||||
response.status_code,
|
||||
url,
|
||||
)
|
||||
verbose_logger.debug(
|
||||
"Galileo Logger HTTP error response body: %s",
|
||||
response.text,
|
||||
)
|
||||
try:
|
||||
verbose_logger.debug(
|
||||
"Galileo Logger HTTP error response json: %s",
|
||||
response.json(),
|
||||
)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
@staticmethod
|
||||
def _build_prompt(kwargs: Dict[str, Any]) -> Dict[str, Any]:
|
||||
optional_params = kwargs.get("optional_params", {}) or {}
|
||||
prompt: Dict[str, Any] = {"messages": kwargs.get("messages")}
|
||||
if optional_params.get("functions") is not None:
|
||||
prompt["functions"] = optional_params["functions"]
|
||||
if optional_params.get("tools") is not None:
|
||||
prompt["tools"] = optional_params["tools"]
|
||||
return prompt
|
||||
|
||||
@staticmethod
|
||||
def _serialize_galileo_output(value: Any) -> str:
|
||||
if value is None:
|
||||
return ""
|
||||
if isinstance(value, str):
|
||||
return value
|
||||
|
||||
def _json_default(obj: Any) -> Any:
|
||||
if hasattr(obj, "model_dump"):
|
||||
return obj.model_dump()
|
||||
return str(obj)
|
||||
|
||||
return json.dumps(value, default=_json_default)
|
||||
|
||||
@staticmethod
|
||||
def _prompt_to_input_text(prompt: Dict[str, Any]) -> str:
|
||||
messages = prompt.get("messages")
|
||||
if messages is not None:
|
||||
text = GalileoObserve._input_text_from_messages(messages)
|
||||
if text:
|
||||
return text
|
||||
return json.dumps(prompt, default=str)
|
||||
|
||||
@staticmethod
|
||||
def _get_chat_content_for_galileo(response_obj: litellm.ModelResponse) -> Any:
|
||||
if response_obj.choices and len(response_obj.choices) > 0:
|
||||
message = response_obj["choices"][0]["message"]
|
||||
if hasattr(message, "json"):
|
||||
message_json = message.json()
|
||||
if isinstance(message_json, str):
|
||||
return json.loads(message_json)
|
||||
return message_json
|
||||
return message
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def _get_text_completion_content_for_galileo(
|
||||
response_obj: litellm.TextCompletionResponse,
|
||||
) -> Optional[str]:
|
||||
if response_obj.choices and len(response_obj.choices) > 0:
|
||||
return response_obj.choices[0].text
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def _get_responses_api_content_for_galileo(
|
||||
response_obj: ResponsesAPIResponse,
|
||||
) -> Any:
|
||||
if hasattr(response_obj, "output") and response_obj.output:
|
||||
return response_obj.output
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def _langfuse_style_rerank_prompt(kwargs: Dict[str, Any]) -> Dict[str, Any]:
|
||||
"""Match Langfuse rerank input: prompt = {"messages": kwargs.get("messages")}."""
|
||||
return {"messages": kwargs.get("messages")}
|
||||
|
||||
def _get_galileo_input_output_content(
|
||||
self,
|
||||
kwargs: Dict[str, Any],
|
||||
response_obj: Any,
|
||||
level: str = "DEFAULT",
|
||||
status_message: Optional[str] = None,
|
||||
) -> Tuple[str, str, Any]:
|
||||
"""
|
||||
Mirror Langfuse _get_langfuse_input_output_content for Galileo ingest.
|
||||
|
||||
Returns (input_text, output_text, messages_for_span).
|
||||
"""
|
||||
call_type = kwargs.get("call_type")
|
||||
prompt = self._build_prompt(kwargs)
|
||||
|
||||
if (
|
||||
level == "ERROR"
|
||||
and status_message is not None
|
||||
and isinstance(status_message, str)
|
||||
):
|
||||
return self._prompt_to_input_text(prompt), status_message, prompt
|
||||
|
||||
if response_obj is not None and (
|
||||
call_type in ("embedding", "aembedding")
|
||||
or isinstance(response_obj, litellm.EmbeddingResponse)
|
||||
):
|
||||
# Match Langfuse OTEL: log embeddings without serializing vectors.
|
||||
return self._prompt_to_input_text(prompt), "embedding-output", prompt
|
||||
|
||||
if response_obj is not None and isinstance(response_obj, litellm.ModelResponse):
|
||||
output = self._get_chat_content_for_galileo(response_obj)
|
||||
return (
|
||||
self._prompt_to_input_text(prompt),
|
||||
self._serialize_galileo_output(output),
|
||||
kwargs.get("messages") or [],
|
||||
)
|
||||
|
||||
if response_obj is not None and isinstance(
|
||||
response_obj, HttpxBinaryResponseContent
|
||||
):
|
||||
return self._prompt_to_input_text(prompt), "speech-output", prompt
|
||||
|
||||
if response_obj is not None and isinstance(
|
||||
response_obj, litellm.TextCompletionResponse
|
||||
):
|
||||
output = self._get_text_completion_content_for_galileo(response_obj)
|
||||
return (
|
||||
self._prompt_to_input_text(prompt),
|
||||
self._serialize_galileo_output(output),
|
||||
kwargs.get("messages") or [],
|
||||
)
|
||||
|
||||
if response_obj is not None and isinstance(response_obj, litellm.ImageResponse):
|
||||
output = response_obj.get("data", None)
|
||||
return (
|
||||
self._prompt_to_input_text(prompt),
|
||||
self._serialize_galileo_output(output),
|
||||
prompt,
|
||||
)
|
||||
|
||||
if response_obj is not None and isinstance(
|
||||
response_obj, litellm.TranscriptionResponse
|
||||
):
|
||||
output = response_obj.get("text", None)
|
||||
return (
|
||||
self._prompt_to_input_text(prompt),
|
||||
self._serialize_galileo_output(output),
|
||||
prompt,
|
||||
)
|
||||
|
||||
if response_obj is not None and isinstance(
|
||||
response_obj, litellm.RerankResponse
|
||||
):
|
||||
output = response_obj.results
|
||||
rerank_prompt = self._langfuse_style_rerank_prompt(kwargs)
|
||||
return (
|
||||
json.dumps(rerank_prompt, default=str),
|
||||
self._serialize_galileo_output(output),
|
||||
rerank_prompt,
|
||||
)
|
||||
|
||||
if response_obj is not None and isinstance(response_obj, ResponsesAPIResponse):
|
||||
output = self._get_responses_api_content_for_galileo(response_obj)
|
||||
return (
|
||||
self._prompt_to_input_text(prompt),
|
||||
self._serialize_galileo_output(output),
|
||||
kwargs.get("messages") or [],
|
||||
)
|
||||
|
||||
if (
|
||||
call_type == "_arealtime"
|
||||
and response_obj is not None
|
||||
and isinstance(response_obj, list)
|
||||
):
|
||||
input_val = kwargs.get("input")
|
||||
return (
|
||||
self._serialize_galileo_output(input_val),
|
||||
self._serialize_galileo_output(response_obj),
|
||||
input_val,
|
||||
)
|
||||
|
||||
if (
|
||||
call_type == "pass_through_endpoint"
|
||||
and response_obj is not None
|
||||
and isinstance(response_obj, dict)
|
||||
):
|
||||
output = response_obj.get("response", "")
|
||||
return (
|
||||
self._prompt_to_input_text(prompt),
|
||||
self._serialize_galileo_output(output),
|
||||
prompt,
|
||||
)
|
||||
|
||||
if response_obj is not None and isinstance(response_obj, dict):
|
||||
output = get_content_from_model_response(response_obj)
|
||||
return (
|
||||
self._prompt_to_input_text(prompt),
|
||||
self._serialize_galileo_output(output),
|
||||
kwargs.get("messages") or [],
|
||||
)
|
||||
|
||||
return self._prompt_to_input_text(prompt), "", kwargs.get("messages") or []
|
||||
|
||||
def get_output_str_from_response(
|
||||
self, response_obj: Any, kwargs: Dict[str, Any]
|
||||
) -> Optional[str]:
|
||||
if response_obj is None:
|
||||
return None
|
||||
if kwargs.get("call_type", None) == "embedding" or isinstance(
|
||||
response_obj, litellm.EmbeddingResponse
|
||||
):
|
||||
return None
|
||||
if isinstance(response_obj, litellm.TextCompletionResponse):
|
||||
return response_obj.choices[0].text
|
||||
if isinstance(response_obj, litellm.ImageResponse):
|
||||
return json.dumps(response_obj["data"], default=str)
|
||||
if isinstance(response_obj, (litellm.ModelResponse, dict)):
|
||||
return get_content_from_model_response(response_obj)
|
||||
return None
|
||||
) -> str:
|
||||
_, output_text, _ = self._get_galileo_input_output_content(
|
||||
kwargs=kwargs, response_obj=response_obj
|
||||
)
|
||||
return output_text
|
||||
|
||||
@staticmethod
|
||||
def _input_text_from_messages(messages: Any) -> str:
|
||||
"""Return a plain-string summary of the input suitable for the trace-level input field."""
|
||||
if isinstance(messages, str):
|
||||
return messages
|
||||
if not isinstance(messages, list):
|
||||
return ""
|
||||
# Use the last user/human message so the trace table shows the actual prompt
|
||||
for msg in reversed(messages):
|
||||
if not isinstance(msg, dict):
|
||||
continue
|
||||
if str(msg.get("role", "")).lower() in ("user", "human"):
|
||||
content = msg.get("content") or ""
|
||||
if isinstance(content, list):
|
||||
content = " ".join(
|
||||
b.get("text", "") if isinstance(b, dict) else str(b)
|
||||
for b in content
|
||||
)
|
||||
if content:
|
||||
return str(content)
|
||||
# Fallback: first non-empty content of any role
|
||||
for msg in messages:
|
||||
if isinstance(msg, dict):
|
||||
content = msg.get("content") or ""
|
||||
if isinstance(content, list):
|
||||
content = " ".join(
|
||||
b.get("text", "") if isinstance(b, dict) else str(b)
|
||||
for b in content
|
||||
)
|
||||
if content:
|
||||
return str(content)
|
||||
return ""
|
||||
|
||||
async def async_log_success_event(
|
||||
self, kwargs: Any, response_obj: Any, start_time: Any, end_time: Any
|
||||
):
|
||||
verbose_logger.debug("On Async Success")
|
||||
try:
|
||||
await self._async_log_success_event_impl(
|
||||
kwargs=kwargs,
|
||||
response_obj=response_obj,
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
)
|
||||
except Exception:
|
||||
verbose_logger.exception(
|
||||
"Galileo Logger: unexpected error in async_log_success_event"
|
||||
)
|
||||
|
||||
async def _async_log_success_event_impl(
|
||||
self, kwargs: Any, response_obj: Any, start_time: Any, end_time: Any
|
||||
):
|
||||
if not self._is_configured():
|
||||
verbose_logger.debug(
|
||||
"Galileo Logger: skipping flush — set GALILEO_PROJECT_ID and "
|
||||
"either GALILEO_API_KEY (hosted) or GALILEO_USERNAME/GALILEO_PASSWORD "
|
||||
"(enterprise Observe)."
|
||||
"Galileo Logger: skipping — GALILEO_PROJECT_ID=%s GALILEO_API_KEY=%s GALILEO_BASE_URL=%s",
|
||||
bool(self.project_id),
|
||||
bool(self.api_key),
|
||||
bool(self.base_url),
|
||||
)
|
||||
return
|
||||
|
||||
_latency_ms = int((end_time - start_time).total_seconds() * 1000)
|
||||
_call_type = kwargs.get("call_type", "litellm")
|
||||
input_text = litellm.utils.get_formatted_prompt(
|
||||
data=kwargs, call_type=_call_type
|
||||
slo: Optional[Dict[str, Any]] = kwargs.get("standard_logging_object")
|
||||
if slo is None:
|
||||
verbose_logger.debug(
|
||||
"Galileo Logger: no standard_logging_object in kwargs, skipping"
|
||||
)
|
||||
return
|
||||
|
||||
_call_type: str = str(
|
||||
slo.get("call_type") or kwargs.get("call_type") or "litellm"
|
||||
)
|
||||
|
||||
_usage = response_obj.get("usage", {}) or {}
|
||||
num_input_tokens = _usage.get("prompt_tokens", 0)
|
||||
num_output_tokens = _usage.get("completion_tokens", 0)
|
||||
|
||||
output_text = self.get_output_str_from_response(
|
||||
response_obj=response_obj, kwargs=kwargs
|
||||
input_text, output_text, messages = self._get_galileo_input_output_content(
|
||||
kwargs=kwargs, response_obj=response_obj
|
||||
)
|
||||
|
||||
if output_text is not None:
|
||||
request_record = LLMResponse(
|
||||
latency_ms=_latency_ms,
|
||||
status_code=200,
|
||||
input_text=input_text,
|
||||
output_text=output_text,
|
||||
node_type=_call_type,
|
||||
model=kwargs.get("model", "-"),
|
||||
num_input_tokens=num_input_tokens,
|
||||
num_output_tokens=num_output_tokens,
|
||||
created_at=start_time.strftime(
|
||||
"%Y-%m-%dT%H:%M:%S"
|
||||
), # timestamp str constructed in "%Y-%m-%dT%H:%M:%S" format
|
||||
raw_start = slo.get("startTime")
|
||||
raw_end = slo.get("endTime")
|
||||
if raw_start is None or raw_end is None:
|
||||
verbose_logger.debug(
|
||||
"Galileo Logger: standard_logging_object missing startTime/endTime, "
|
||||
"falling back to start_time/end_time params"
|
||||
)
|
||||
if not isinstance(start_time, datetime) or not isinstance(
|
||||
end_time, datetime
|
||||
):
|
||||
return
|
||||
start_ts = start_time
|
||||
end_ts = end_time
|
||||
if start_ts.tzinfo is None:
|
||||
start_ts = start_ts.replace(tzinfo=GalileoObserve._local_timezone())
|
||||
if end_ts.tzinfo is None:
|
||||
end_ts = end_ts.replace(tzinfo=GalileoObserve._local_timezone())
|
||||
start_ts = start_ts.astimezone(timezone.utc)
|
||||
end_ts = end_ts.astimezone(timezone.utc)
|
||||
else:
|
||||
start_ts = datetime.fromtimestamp(float(raw_start), tz=timezone.utc)
|
||||
end_ts = datetime.fromtimestamp(float(raw_end), tz=timezone.utc)
|
||||
_latency_ms = max(0, int((end_ts - start_ts).total_seconds() * 1000))
|
||||
num_input_tokens = int(slo.get("prompt_tokens") or 0)
|
||||
num_output_tokens = int(slo.get("completion_tokens") or 0)
|
||||
num_total_tokens = int(slo.get("total_tokens") or 0)
|
||||
if num_total_tokens == 0 and (num_input_tokens or num_output_tokens):
|
||||
num_total_tokens = num_input_tokens + num_output_tokens
|
||||
|
||||
request_record = LLMResponse(
|
||||
latency_ms=_latency_ms,
|
||||
status_code=200,
|
||||
input_text=input_text,
|
||||
output_text=output_text,
|
||||
node_type=_call_type,
|
||||
model=str(slo.get("model") or kwargs.get("model") or "-"),
|
||||
num_input_tokens=num_input_tokens,
|
||||
num_output_tokens=num_output_tokens,
|
||||
num_total_tokens=num_total_tokens,
|
||||
cost=slo.get("response_cost"),
|
||||
created_at=GalileoObserve._format_created_at(start_ts),
|
||||
)
|
||||
|
||||
request_dict = request_record.model_dump()
|
||||
if isinstance(messages, dict):
|
||||
messages = messages.get("messages")
|
||||
if isinstance(messages, list) and messages:
|
||||
request_dict["messages"] = messages
|
||||
self.in_memory_records.append(request_dict)
|
||||
verbose_logger.debug(
|
||||
"Galileo Logger: queued record, in_memory=%d", len(self.in_memory_records)
|
||||
)
|
||||
|
||||
# Bound the buffer so persistent flush failures cannot grow it
|
||||
# without limit. Drop the oldest records once we exceed the cap.
|
||||
if len(self.in_memory_records) > GALILEO_MAX_IN_MEMORY_RECORDS:
|
||||
dropped = len(self.in_memory_records) - GALILEO_MAX_IN_MEMORY_RECORDS
|
||||
self.in_memory_records = self.in_memory_records[
|
||||
-GALILEO_MAX_IN_MEMORY_RECORDS:
|
||||
]
|
||||
verbose_logger.warning(
|
||||
"Galileo Logger: in-memory buffer exceeded %s records; "
|
||||
"dropped %s oldest record(s). Check Galileo connectivity/credentials.",
|
||||
GALILEO_MAX_IN_MEMORY_RECORDS,
|
||||
dropped,
|
||||
)
|
||||
|
||||
request_dict = request_record.model_dump()
|
||||
messages = kwargs.get("messages")
|
||||
if messages:
|
||||
request_dict["messages"] = messages
|
||||
self.in_memory_records.append(request_dict)
|
||||
|
||||
# Bound the buffer so persistent flush failures cannot grow it
|
||||
# without limit. Drop the oldest records once we exceed the cap.
|
||||
if len(self.in_memory_records) > GALILEO_MAX_IN_MEMORY_RECORDS:
|
||||
dropped = len(self.in_memory_records) - GALILEO_MAX_IN_MEMORY_RECORDS
|
||||
self.in_memory_records = self.in_memory_records[
|
||||
-GALILEO_MAX_IN_MEMORY_RECORDS:
|
||||
]
|
||||
verbose_logger.warning(
|
||||
"Galileo Logger: in-memory buffer exceeded %s records; "
|
||||
"dropped %s oldest record(s). Check Galileo connectivity/credentials.",
|
||||
GALILEO_MAX_IN_MEMORY_RECORDS,
|
||||
dropped,
|
||||
)
|
||||
|
||||
if len(self.in_memory_records) >= self.batch_size:
|
||||
await self.flush_in_memory_records()
|
||||
if len(self.in_memory_records) >= self.batch_size:
|
||||
await self.flush_in_memory_records()
|
||||
|
||||
async def flush_in_memory_records(self):
|
||||
if not self.in_memory_records:
|
||||
|
|
@ -296,15 +790,23 @@ class GalileoObserve(CustomLogger):
|
|||
ingest_request = self._get_ingest_request()
|
||||
if ingest_request is None:
|
||||
verbose_logger.debug(
|
||||
"Galileo Logger: missing GALILEO_BASE_URL or GALILEO_PROJECT_ID"
|
||||
"Galileo Logger: missing GALILEO_BASE_URL or GALILEO_PROJECT_ID — skipping flush"
|
||||
)
|
||||
return
|
||||
|
||||
if not await self._ensure_headers():
|
||||
verbose_logger.debug("Galileo Logger: could not set request headers")
|
||||
verbose_logger.debug(
|
||||
"Galileo Logger: could not set request headers — skipping flush"
|
||||
)
|
||||
return
|
||||
|
||||
url, payload = ingest_request
|
||||
self._log_flush_config()
|
||||
self._log_flush_payload(url=url, payload=payload)
|
||||
verbose_logger.debug(
|
||||
"Galileo Logger flush headers: %s",
|
||||
self._redact_headers(self.headers),
|
||||
)
|
||||
verbose_logger.debug("flushing in memory records to %s", url)
|
||||
|
||||
try:
|
||||
|
|
@ -313,6 +815,12 @@ class GalileoObserve(CustomLogger):
|
|||
headers=self.headers,
|
||||
json=payload,
|
||||
)
|
||||
except httpx.HTTPStatusError as e:
|
||||
self._log_http_status_error(error=e, url=url)
|
||||
verbose_logger.debug(
|
||||
"Galileo Logger: failed to flush in memory records: %s", e
|
||||
)
|
||||
return
|
||||
except Exception as e:
|
||||
verbose_logger.debug(
|
||||
"Galileo Logger: failed to flush in memory records: %s", e
|
||||
|
|
@ -323,6 +831,11 @@ class GalileoObserve(CustomLogger):
|
|||
verbose_logger.debug(
|
||||
"Galileo Logger: successfully flushed in memory records"
|
||||
)
|
||||
verbose_logger.debug(
|
||||
"Galileo Logger flush response: status=%s body=%s",
|
||||
response.status_code,
|
||||
response.text,
|
||||
)
|
||||
del self.in_memory_records[:records_in_payload]
|
||||
else:
|
||||
verbose_logger.debug("Galileo Logger: failed to flush in memory records")
|
||||
|
|
|
|||
|
|
@ -102,6 +102,18 @@ def langfuse_client_init(
|
|||
if Version(langfuse.version.__version__) >= Version("2.6.0"):
|
||||
parameters["sdk_integration"] = "litellm"
|
||||
|
||||
if Version(langfuse.version.__version__) >= Version("2.7.3"):
|
||||
import httpx
|
||||
|
||||
import litellm
|
||||
|
||||
from ...llms.custom_httpx.http_handler import get_ssl_configuration
|
||||
|
||||
parameters["httpx_client"] = httpx.Client(
|
||||
verify=get_ssl_configuration(),
|
||||
cert=os.getenv("SSL_CERTIFICATE", litellm.ssl_certificate),
|
||||
)
|
||||
|
||||
client = Langfuse(**parameters)
|
||||
|
||||
return client
|
||||
|
|
|
|||
0
litellm/integrations/mavvrik_focus/__init__.py
Normal file
0
litellm/integrations/mavvrik_focus/__init__.py
Normal file
272
litellm/integrations/mavvrik_focus/mavvrik_focus_logger.py
Normal file
272
litellm/integrations/mavvrik_focus/mavvrik_focus_logger.py
Normal file
|
|
@ -0,0 +1,272 @@
|
|||
"""MavvrikFocusLogger — FOCUS-based Mavvrik export logger.
|
||||
|
||||
Usage in config.yaml:
|
||||
litellm_settings:
|
||||
callbacks: ["mavvrik"]
|
||||
|
||||
Required env vars:
|
||||
MAVVRIK_API_KEY
|
||||
MAVVRIK_API_ENDPOINT
|
||||
MAVVRIK_CONNECTION_ID
|
||||
|
||||
Optional env vars:
|
||||
MAVVRIK_FOCUS_MAX_ROWS — row cap per export window (default: 500000)
|
||||
|
||||
Only daily frequency is supported. The Mavvrik ingestion protocol stores one
|
||||
file per calendar date (metrics/YYYY-MM-DD). Hourly or interval exports would
|
||||
overwrite each other within the same day, producing incomplete data.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from typing import TYPE_CHECKING, Any, List, Optional
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.constants import MAVVRIK_FOCUS_EXPORT_JOB_NAME
|
||||
from litellm.integrations.focus.destinations.base import FocusTimeWindow
|
||||
from litellm.integrations.focus.focus_logger import FocusLogger
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from apscheduler.schedulers.asyncio import AsyncIOScheduler
|
||||
else:
|
||||
AsyncIOScheduler = Any
|
||||
|
||||
|
||||
def _parse_metrics_marker(
|
||||
marker: Optional[object],
|
||||
) -> Optional[datetime]:
|
||||
"""Parse metricsMarker from Mavvrik register response into a UTC datetime.
|
||||
|
||||
Handles both formats Mavvrik may return:
|
||||
- Unix timestamp (int/float): e.g. 1749340800
|
||||
- ISO date string: e.g. "2026-06-09" or "2026-06-09T00:00:00Z"
|
||||
|
||||
Returns None for falsy values (0, None, empty string) which indicate
|
||||
no data has been ingested yet.
|
||||
"""
|
||||
if not marker:
|
||||
return None
|
||||
try:
|
||||
if isinstance(marker, (int, float)):
|
||||
return datetime.fromtimestamp(float(marker), tz=timezone.utc).replace(
|
||||
hour=0, minute=0, second=0, microsecond=0
|
||||
)
|
||||
if isinstance(marker, str):
|
||||
marker = marker.strip()
|
||||
if not marker:
|
||||
return None
|
||||
# Try ISO date first (YYYY-MM-DD), then full ISO datetime
|
||||
for fmt in ("%Y-%m-%d", "%Y-%m-%dT%H:%M:%SZ", "%Y-%m-%dT%H:%M:%S"):
|
||||
try:
|
||||
return datetime.strptime(marker, fmt).replace(tzinfo=timezone.utc)
|
||||
except ValueError:
|
||||
continue
|
||||
except Exception:
|
||||
pass
|
||||
verbose_proxy_logger.warning(
|
||||
"Mavvrik FOCUS: could not parse metricsMarker %r — skipping catch-up", marker
|
||||
)
|
||||
return None
|
||||
|
||||
|
||||
class MavvrikFocusLogger(FocusLogger):
|
||||
"""FOCUS-based export logger that routes to the Mavvrik destination."""
|
||||
|
||||
def __init__(self, **kwargs: Any) -> None:
|
||||
frequency = os.getenv("MAVVRIK_FOCUS_FREQUENCY", "daily").lower()
|
||||
if frequency != "daily":
|
||||
raise ValueError(
|
||||
f"MAVVRIK_FOCUS_FREQUENCY='{frequency}' is not supported. "
|
||||
"Only 'daily' is allowed -- the Mavvrik ingestion protocol stores one "
|
||||
"file per calendar date (metrics/YYYY-MM-DD). Hourly or interval "
|
||||
"exports would overwrite each other within the same day."
|
||||
)
|
||||
super().__init__(
|
||||
provider="mavvrik",
|
||||
export_format="csv",
|
||||
frequency="daily",
|
||||
prefix="mavvrik_focus_exports",
|
||||
destination_config={
|
||||
"api_key": os.getenv("MAVVRIK_API_KEY"),
|
||||
"api_endpoint": os.getenv("MAVVRIK_API_ENDPOINT"),
|
||||
"connection_id": os.getenv("MAVVRIK_CONNECTION_ID"),
|
||||
},
|
||||
**kwargs,
|
||||
)
|
||||
raw = os.getenv("MAVVRIK_FOCUS_MAX_ROWS")
|
||||
self._max_rows: Optional[int] = int(raw) if raw else 500_000
|
||||
|
||||
async def _export_window(
|
||||
self,
|
||||
*,
|
||||
window: FocusTimeWindow,
|
||||
limit: Optional[int],
|
||||
) -> None:
|
||||
"""Export with Mavvrik row cap applied when no explicit limit is passed."""
|
||||
effective_limit = limit if limit is not None else self._max_rows
|
||||
engine = self._ensure_engine()
|
||||
data = await engine._database.get_usage_data(
|
||||
limit=effective_limit,
|
||||
start_time_utc=window.start_time,
|
||||
end_time_utc=window.end_time,
|
||||
)
|
||||
if effective_limit is not None and len(data) >= effective_limit:
|
||||
verbose_proxy_logger.warning(
|
||||
"Mavvrik FOCUS export: row cap reached (%d rows). "
|
||||
"Some data for window %s→%s may be excluded. "
|
||||
"Increase MAVVRIK_FOCUS_MAX_ROWS to export all rows.",
|
||||
effective_limit,
|
||||
window.start_time.date(),
|
||||
window.end_time.date(),
|
||||
)
|
||||
if data.is_empty():
|
||||
verbose_proxy_logger.debug(
|
||||
"Mavvrik FOCUS export: no usage data for window %s", window
|
||||
)
|
||||
return
|
||||
normalized = engine._transformer.transform(data)
|
||||
if normalized.is_empty():
|
||||
return
|
||||
payload = engine._serializer.serialize(normalized)
|
||||
if not payload:
|
||||
return
|
||||
await engine._destination.deliver(
|
||||
content=payload,
|
||||
time_window=window,
|
||||
filename=engine._build_filename(window),
|
||||
)
|
||||
|
||||
# Maximum number of days to catch up in a single run. Prevents runaway
|
||||
# loops if the connector was disabled for a long time, and avoids querying
|
||||
# data that has likely been cleaned up from LiteLLM_DailyUserSpend.
|
||||
_MAX_CATCHUP_DAYS = 7
|
||||
|
||||
async def _run_scheduled_export(self) -> None:
|
||||
"""Export today's window, catching up any dates Mavvrik has not yet received.
|
||||
|
||||
On each run:
|
||||
1. Register with Mavvrik → get metricsMarker (last successfully ingested date)
|
||||
2. If metricsMarker is behind yesterday, catch up missed dates (capped at
|
||||
_MAX_CATCHUP_DAYS to avoid runaway loops on long outages)
|
||||
3. Export yesterday (today's daily window)
|
||||
|
||||
This ensures a failed export on day N is automatically retried on day N+1
|
||||
without any manual intervention.
|
||||
"""
|
||||
engine = self._ensure_engine()
|
||||
from litellm.integrations.focus.destinations.mavvrik_destination import ( # noqa: PLC0415
|
||||
FocusMavvrikDestination,
|
||||
)
|
||||
|
||||
destination = engine._destination
|
||||
if not isinstance(destination, FocusMavvrikDestination):
|
||||
await super()._run_scheduled_export()
|
||||
return
|
||||
|
||||
# Register and get the last date Mavvrik has processed.
|
||||
# metricsMarker may be a Unix timestamp (int/float) or an ISO date string.
|
||||
marker = await destination.get_metrics_marker()
|
||||
|
||||
now = datetime.now(timezone.utc)
|
||||
yesterday = now.replace(hour=0, minute=0, second=0, microsecond=0) - timedelta(
|
||||
days=1
|
||||
)
|
||||
|
||||
last_ingested = _parse_metrics_marker(marker)
|
||||
|
||||
# Catch up missed dates, capped at _MAX_CATCHUP_DAYS
|
||||
if last_ingested and last_ingested < yesterday:
|
||||
# Never go further back than _MAX_CATCHUP_DAYS from yesterday
|
||||
earliest_catchup = yesterday - timedelta(days=self._MAX_CATCHUP_DAYS - 1)
|
||||
catch_up_date = max(last_ingested + timedelta(days=1), earliest_catchup)
|
||||
|
||||
if last_ingested + timedelta(days=1) < earliest_catchup:
|
||||
verbose_proxy_logger.warning(
|
||||
"Mavvrik FOCUS export: metricsMarker is more than %d days behind "
|
||||
"(%s). Catching up from %s only; earlier data will not be re-exported.",
|
||||
self._MAX_CATCHUP_DAYS,
|
||||
last_ingested.date(),
|
||||
catch_up_date.date(),
|
||||
)
|
||||
|
||||
while catch_up_date < yesterday:
|
||||
verbose_proxy_logger.info(
|
||||
"Mavvrik FOCUS export: catching up missed date %s",
|
||||
catch_up_date.date(),
|
||||
)
|
||||
window = FocusTimeWindow(
|
||||
start_time=catch_up_date,
|
||||
end_time=catch_up_date + timedelta(days=1),
|
||||
frequency="daily",
|
||||
)
|
||||
await self._export_window(window=window, limit=None)
|
||||
catch_up_date += timedelta(days=1)
|
||||
|
||||
# Export yesterday's window (the normal daily run)
|
||||
window = FocusTimeWindow(
|
||||
start_time=yesterday,
|
||||
end_time=yesterday + timedelta(days=1),
|
||||
frequency="daily",
|
||||
)
|
||||
await self._export_window(window=window, limit=None)
|
||||
|
||||
async def initialize_mavvrik_focus_export_job(self) -> None:
|
||||
"""Scheduler entry point — uses Mavvrik-specific pod-lock key."""
|
||||
from litellm.proxy.proxy_server import proxy_logging_obj # noqa: PLC0415
|
||||
|
||||
pod_lock_manager = None
|
||||
if proxy_logging_obj is not None:
|
||||
writer = getattr(proxy_logging_obj, "db_spend_update_writer", None)
|
||||
if writer is not None:
|
||||
pod_lock_manager = getattr(writer, "pod_lock_manager", None)
|
||||
|
||||
if pod_lock_manager and pod_lock_manager.redis_cache:
|
||||
acquired = await pod_lock_manager.acquire_lock(
|
||||
cronjob_id=MAVVRIK_FOCUS_EXPORT_JOB_NAME
|
||||
)
|
||||
if not acquired:
|
||||
verbose_proxy_logger.debug(
|
||||
"Mavvrik FOCUS export: unable to acquire pod lock"
|
||||
)
|
||||
return
|
||||
try:
|
||||
await self._run_scheduled_export()
|
||||
finally:
|
||||
await pod_lock_manager.release_lock(
|
||||
cronjob_id=MAVVRIK_FOCUS_EXPORT_JOB_NAME
|
||||
)
|
||||
else:
|
||||
await self._run_scheduled_export()
|
||||
|
||||
@staticmethod
|
||||
async def init_mavvrik_focus_background_job(
|
||||
scheduler: AsyncIOScheduler,
|
||||
) -> None:
|
||||
"""Register the Mavvrik FOCUS export job on the provided scheduler."""
|
||||
loggers: List[MavvrikFocusLogger] = [
|
||||
cb
|
||||
for cb in litellm.logging_callback_manager.get_custom_loggers_for_type(
|
||||
callback_type=MavvrikFocusLogger
|
||||
)
|
||||
if type(cb) is MavvrikFocusLogger
|
||||
]
|
||||
if not loggers:
|
||||
verbose_proxy_logger.debug(
|
||||
"No MavvrikFocusLogger registered; skipping scheduler"
|
||||
)
|
||||
return
|
||||
|
||||
logger = loggers[0]
|
||||
trigger_kwargs = logger._build_scheduler_trigger()
|
||||
scheduler.add_job( # type: ignore[attr-defined]
|
||||
logger.initialize_mavvrik_focus_export_job,
|
||||
id=MAVVRIK_FOCUS_EXPORT_JOB_NAME,
|
||||
replace_existing=True,
|
||||
**trigger_kwargs,
|
||||
)
|
||||
verbose_proxy_logger.info(
|
||||
"mavvrik_focus: background export job scheduled (%s)", trigger_kwargs
|
||||
)
|
||||
10
litellm/integrations/newrelic/__init__.py
Normal file
10
litellm/integrations/newrelic/__init__.py
Normal file
|
|
@ -0,0 +1,10 @@
|
|||
"""
|
||||
New Relic AI Monitoring Integration for LiteLLM
|
||||
|
||||
This module provides integration with New Relic's AI Monitoring feature to track
|
||||
LLM requests, responses, and usage metrics.
|
||||
"""
|
||||
|
||||
from litellm.integrations.newrelic.newrelic import NewRelicLogger
|
||||
|
||||
__all__ = ["NewRelicLogger"]
|
||||
926
litellm/integrations/newrelic/newrelic.py
Normal file
926
litellm/integrations/newrelic/newrelic.py
Normal file
|
|
@ -0,0 +1,926 @@
|
|||
"""
|
||||
New Relic AI Monitoring Integration for LiteLLM
|
||||
|
||||
This module provides integration with New Relic's AI Monitoring feature to track
|
||||
LLM requests, responses, and usage metrics.
|
||||
|
||||
Environment Variables (consumed by the New Relic agent at process bootstrap -
|
||||
set via container env, or before invoking `newrelic-admin run-program`):
|
||||
NEW_RELIC_LICENSE_KEY: Your New Relic license key (required)
|
||||
NEW_RELIC_APP_NAME: Your application name (required)
|
||||
|
||||
UI- and runtime-toggleable:
|
||||
NEW_RELIC_AI_MONITORING_RECORD_CONTENT_ENABLED: Whether to record message
|
||||
content (optional, default: true)
|
||||
|
||||
Configuration:
|
||||
Message logging can be controlled via (both must agree to record):
|
||||
1. turn_off_message_logging parameter - pass via callback initialization or config YAML
|
||||
2. NEW_RELIC_AI_MONITORING_RECORD_CONTENT_ENABLED env var
|
||||
|
||||
Default behavior: Messages ARE recorded unless explicitly disabled by either method
|
||||
Either method can disable recording - both must enable for recording to occur
|
||||
|
||||
Usage - Python SDK:
|
||||
import litellm
|
||||
litellm.callbacks = ["newrelic"]
|
||||
|
||||
# Or with explicit configuration:
|
||||
from litellm.integrations.newrelic import NewRelicLogger
|
||||
litellm.callbacks = [NewRelicLogger(turn_off_message_logging=True)]
|
||||
|
||||
Usage - Proxy Server (config.yaml):
|
||||
litellm_settings:
|
||||
callbacks: ["newrelic"]
|
||||
newrelic_params:
|
||||
turn_off_message_logging: true # Disable message content recording
|
||||
|
||||
# Or disable via environment variable:
|
||||
# export NEW_RELIC_AI_MONITORING_RECORD_CONTENT_ENABLED=false
|
||||
|
||||
# Ensure New Relic agent is initialized (use newrelic-admin or initialize manually)
|
||||
# newrelic-admin run-program python your_app.py
|
||||
"""
|
||||
|
||||
import json
|
||||
import os
|
||||
import threading
|
||||
import time
|
||||
import uuid
|
||||
from typing import Any, Dict, List, Optional, Tuple, Union
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.litellm_core_utils.redact_messages import should_redact_message_logging
|
||||
from litellm.types.integrations.newrelic import NewRelicInitParams
|
||||
from litellm.types.integrations.base_health_check import IntegrationHealthCheckStatus
|
||||
from litellm.types.utils import ModelResponse, Message, StandardLoggingPayload
|
||||
|
||||
try:
|
||||
import newrelic.agent as _newrelic_agent
|
||||
except ImportError:
|
||||
_newrelic_agent = None # type: ignore
|
||||
|
||||
|
||||
class NewRelicLogger(CustomLogger):
|
||||
"""
|
||||
New Relic logger for LiteLLM to send AI monitoring events.
|
||||
|
||||
This logger creates two types of New Relic custom events:
|
||||
1. LlmChatCompletionSummary - One per completion request
|
||||
2. LlmChatCompletionMessage - One per message (request and response)
|
||||
"""
|
||||
|
||||
# Class-level state for supportability metric emission, shared across all instances.
|
||||
# Protected by _metric_lock to ensure thread-safe access.
|
||||
_last_metric_emission_time: float = 0.0
|
||||
_metric_lock = threading.Lock()
|
||||
|
||||
def __init__(self, **kwargs):
|
||||
#########################################################
|
||||
# Handle newrelic_params set as litellm.newrelic_params
|
||||
#########################################################
|
||||
dict_newrelic_params = self._get_newrelic_params()
|
||||
|
||||
# Use setdefault so constructor kwargs take priority over global params.
|
||||
# model_dump() always returns all fields (including defaults), so update()
|
||||
# would silently overwrite explicit constructor args like turn_off_message_logging=True.
|
||||
for k, v in dict_newrelic_params.items():
|
||||
kwargs.setdefault(k, v)
|
||||
|
||||
# CustomLogger.__init__ will set self.turn_off_message_logging from kwargs
|
||||
super().__init__(**kwargs)
|
||||
|
||||
# Check for required environment variables
|
||||
self.license_key = os.getenv("NEW_RELIC_LICENSE_KEY")
|
||||
self.app_name = os.getenv("NEW_RELIC_APP_NAME")
|
||||
|
||||
# Validate configuration
|
||||
if not self.license_key or not self.app_name:
|
||||
verbose_logger.warning(
|
||||
"New Relic integration requires NEW_RELIC_LICENSE_KEY and "
|
||||
"NEW_RELIC_APP_NAME environment variables. Integration will be disabled."
|
||||
)
|
||||
self.enabled = False
|
||||
elif _newrelic_agent is None:
|
||||
verbose_logger.error(
|
||||
"New Relic Python agent not installed. Review the New Relic integration documentation at https://docs.litellm.ai/docs/observability/newrelic."
|
||||
)
|
||||
self.enabled = False
|
||||
else:
|
||||
try:
|
||||
# timeout=0 forces non-blocking startup: the agent connects in a
|
||||
# background thread regardless of newrelic.ini / NEW_RELIC_STARTUP_TIMEOUT.
|
||||
_newrelic_agent.register_application(timeout=0)
|
||||
|
||||
self.enabled = True
|
||||
verbose_logger.info(
|
||||
f"New Relic AI Monitoring initialized for app: {self.app_name}, "
|
||||
f"content recording: {self.record_content}"
|
||||
)
|
||||
except Exception as e:
|
||||
verbose_logger.error(
|
||||
f"Failed to initialize New Relic agent: {e}. "
|
||||
"Integration will be disabled."
|
||||
)
|
||||
self.enabled = False
|
||||
|
||||
def _get_newrelic_params(self) -> Dict:
|
||||
"""
|
||||
Get the newrelic_params from litellm.newrelic_params
|
||||
|
||||
These are params specific to initializing the NewRelicLogger e.g. turn_off_message_logging
|
||||
"""
|
||||
dict_newrelic_params: Dict = {}
|
||||
if litellm.newrelic_params is not None:
|
||||
if isinstance(litellm.newrelic_params, NewRelicInitParams):
|
||||
dict_newrelic_params = litellm.newrelic_params.model_dump()
|
||||
elif isinstance(litellm.newrelic_params, Dict):
|
||||
# only allow params that are of NewRelicInitParams
|
||||
dict_newrelic_params = NewRelicInitParams(
|
||||
**litellm.newrelic_params
|
||||
).model_dump()
|
||||
return dict_newrelic_params
|
||||
|
||||
@property
|
||||
def record_content(self) -> bool:
|
||||
"""Whether to record message content in New Relic.
|
||||
|
||||
Both turn_off_message_logging param AND NEW_RELIC_AI_MONITORING_RECORD_CONTENT_ENABLED
|
||||
env var must agree to record content. If either disables recording, content will not
|
||||
be recorded. Read at call time so UI config changes take effect without a restart.
|
||||
Default: True (record content) unless explicitly disabled by either method.
|
||||
"""
|
||||
return (not self.turn_off_message_logging) and self._parse_bool_env(
|
||||
"NEW_RELIC_AI_MONITORING_RECORD_CONTENT_ENABLED", True
|
||||
)
|
||||
|
||||
def _parse_bool_env(self, var_name: str, default: bool = False) -> bool:
|
||||
"""Parse a boolean environment variable.
|
||||
|
||||
Accepts true/false, 1/0, yes/no, on/off (case-insensitive,
|
||||
whitespace-tolerant) — matching the convention used in
|
||||
``litellm/__init__.py`` and the standard library's
|
||||
``configparser.BOOLEAN_STATES``. Unrecognised values log a
|
||||
warning and fall back to ``default`` rather than silently
|
||||
flipping user intent.
|
||||
"""
|
||||
raw = os.getenv(var_name)
|
||||
if not raw:
|
||||
return default
|
||||
value = raw.strip().lower()
|
||||
if value in ("1", "true", "yes", "on"):
|
||||
return True
|
||||
if value in ("0", "false", "no", "off"):
|
||||
return False
|
||||
verbose_logger.warning(
|
||||
f"{var_name}={raw!r} is not a recognised boolean "
|
||||
f"(accepts true/false, 1/0, yes/no, on/off). "
|
||||
f"Falling back to default ({default})."
|
||||
)
|
||||
return default
|
||||
|
||||
def _get_litellm_version(self) -> str:
|
||||
"""
|
||||
Get litellm version for supportability metrics.
|
||||
|
||||
Returns:
|
||||
Version string (e.g., "1.80.0") or "unknown" if unable to determine
|
||||
"""
|
||||
try:
|
||||
from importlib.metadata import version
|
||||
|
||||
return version("litellm")
|
||||
except Exception as e:
|
||||
verbose_logger.warning(f"Unable to determine litellm version: {e}")
|
||||
return "unknown"
|
||||
|
||||
def _emit_supportability_metric(self):
|
||||
"""
|
||||
Emit New Relic supportability metric for LiteLLM usage.
|
||||
|
||||
Per spec, this metric should be emitted at least once every 27 hours
|
||||
to indicate the library is in use. Format:
|
||||
Supportability/Python/ML/LiteLLM/{version}
|
||||
|
||||
This method updates _last_metric_emission_time and should
|
||||
be called within a lock when checking periodic emission.
|
||||
"""
|
||||
try:
|
||||
litellm_version = self._get_litellm_version()
|
||||
metric_name = f"Supportability/Python/ML/LiteLLM/{litellm_version}"
|
||||
|
||||
# Record metric with value of 1 (will be aggregated by New Relic)
|
||||
app = _newrelic_agent.application()
|
||||
|
||||
# Always update the timestamp so the 27-hour back-off applies
|
||||
# regardless of whether the app is ready, preventing lock contention
|
||||
# on every request when the agent is slow to register or never starts.
|
||||
NewRelicLogger._last_metric_emission_time = time.time()
|
||||
|
||||
if app and app.enabled:
|
||||
app.record_custom_metric(metric_name, 1)
|
||||
verbose_logger.info(
|
||||
f"Emitted New Relic supportability metric: {metric_name}"
|
||||
)
|
||||
else:
|
||||
verbose_logger.info(
|
||||
"New Relic application is not enabled; skipping metric recording."
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
verbose_logger.warning(f"Failed to emit supportability metric: {e}")
|
||||
|
||||
def _check_and_emit_periodic_metric(self):
|
||||
"""
|
||||
Check if 27 hours have passed since last metric emission and re-emit if needed.
|
||||
|
||||
Uses a mutex to ensure only one thread emits the metric even if multiple
|
||||
requests are being processed concurrently.
|
||||
"""
|
||||
# Quick check without lock to avoid unnecessary locking
|
||||
current_time = time.time()
|
||||
time_since_last_emission = (
|
||||
current_time - NewRelicLogger._last_metric_emission_time
|
||||
)
|
||||
|
||||
if time_since_last_emission >= 97200: # 27 hours = 97200 seconds
|
||||
# Acquire lock to ensure only one thread emits
|
||||
with NewRelicLogger._metric_lock:
|
||||
# Double-check inside lock in case another thread just emitted
|
||||
current_time = time.time()
|
||||
time_since_last_emission = (
|
||||
current_time - NewRelicLogger._last_metric_emission_time
|
||||
)
|
||||
|
||||
if time_since_last_emission >= 97200:
|
||||
self._emit_supportability_metric()
|
||||
|
||||
def _get_trace_context(
|
||||
self,
|
||||
kwargs: Dict,
|
||||
standard_logging_object: Optional[StandardLoggingPayload] = None,
|
||||
) -> str:
|
||||
"""
|
||||
Get the New Relic trace ID for AI monitoring events.
|
||||
|
||||
This integration runs in LiteLLM's async logging worker, outside the
|
||||
New Relic agent's current transaction. Because we can't call
|
||||
`newrelic.agent.current_trace_id()` to let the agent populate the
|
||||
trace_id on AIM custom events, we manually simulate what the agent
|
||||
would do. An AIM event without a trace_id is malformed per the NR
|
||||
schema, so this method always returns a valid string.
|
||||
|
||||
Resolution order:
|
||||
1. W3C traceparent header (litellm_params.metadata.headers.traceparent) -
|
||||
what the agent would link to if we were in-transaction.
|
||||
2. StandardLoggingPayload.trace_id - LiteLLM's internal trace for
|
||||
retry/fallback grouping.
|
||||
3. Generated UUID - synthetic grouping key when upstream context is
|
||||
absent or parsing it fails.
|
||||
|
||||
Span IDs are intentionally not emitted: any span ID recoverable from
|
||||
the inbound traceparent is the caller's parent span, not ours.
|
||||
|
||||
Returns:
|
||||
trace_id: always a non-empty string.
|
||||
"""
|
||||
trace_id: Optional[str] = None
|
||||
try:
|
||||
litellm_params = kwargs.get("litellm_params") or {}
|
||||
metadata = litellm_params.get("metadata") or {}
|
||||
headers = metadata.get("headers") or {}
|
||||
# Normalize header key lookup to be case-insensitive per W3C spec
|
||||
traceparent = next(
|
||||
(v for k, v in headers.items() if k.lower() == "traceparent"), None
|
||||
)
|
||||
|
||||
if traceparent:
|
||||
# Extract trace_id from traceparent header if available
|
||||
# traceparent format: "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-00"
|
||||
parts = traceparent.split("-")
|
||||
if len(parts) == 4:
|
||||
trace_id = parts[1]
|
||||
|
||||
if not trace_id and standard_logging_object:
|
||||
slo_trace_id = standard_logging_object.get("trace_id")
|
||||
if slo_trace_id:
|
||||
trace_id = slo_trace_id
|
||||
|
||||
except Exception as e:
|
||||
verbose_logger.warning(
|
||||
f"Unable to parse New Relic trace context from upstream sources: {e}"
|
||||
)
|
||||
|
||||
if not trace_id:
|
||||
trace_id = uuid.uuid4().hex
|
||||
verbose_logger.debug(
|
||||
f"New Relic trace_id not available from distributed tracing headers or "
|
||||
f"StandardLoggingPayload. Generated trace_id={trace_id} for AI monitoring "
|
||||
f"event grouping."
|
||||
)
|
||||
|
||||
return trace_id
|
||||
|
||||
def _extract_completion_id(self, kwargs: Dict, response_obj: ModelResponse) -> str:
|
||||
"""
|
||||
Extract completion ID from kwargs or response_obj, or generate one.
|
||||
"""
|
||||
completion_id = None
|
||||
|
||||
if response_obj:
|
||||
completion_id = response_obj.get("id")
|
||||
|
||||
if not completion_id:
|
||||
completion_id = kwargs.get("litellm_call_id")
|
||||
|
||||
# If still not found, generate UUID and log warning per spec
|
||||
if not completion_id:
|
||||
completion_id = str(uuid.uuid4())
|
||||
|
||||
return completion_id
|
||||
|
||||
def _get_vendor(
|
||||
self,
|
||||
kwargs: Dict,
|
||||
standard_logging_object: Optional[StandardLoggingPayload] = None,
|
||||
) -> str:
|
||||
"""Extract vendor/provider, preferring StandardLoggingPayload."""
|
||||
if standard_logging_object:
|
||||
vendor = standard_logging_object.get("custom_llm_provider")
|
||||
if vendor:
|
||||
return vendor
|
||||
litellm_params = kwargs.get("litellm_params", {}) or {}
|
||||
return litellm_params.get("custom_llm_provider") or "litellm"
|
||||
|
||||
def _get_model_names(
|
||||
self,
|
||||
kwargs: Dict,
|
||||
response_obj: ModelResponse,
|
||||
standard_logging_object: Optional[StandardLoggingPayload] = None,
|
||||
) -> Tuple[str, str]:
|
||||
"""
|
||||
Extract request and response model names, preferring StandardLoggingPayload
|
||||
for the request model.
|
||||
|
||||
Returns:
|
||||
Tuple of (request_model, response_model)
|
||||
"""
|
||||
request_model = None
|
||||
if standard_logging_object:
|
||||
slo_model = standard_logging_object.get("model")
|
||||
if slo_model:
|
||||
request_model = str(slo_model)
|
||||
if not request_model:
|
||||
request_model = str(kwargs.get("model") or "unknown")
|
||||
response_model: str = str(response_obj.get("model") or request_model)
|
||||
return request_model, response_model
|
||||
|
||||
def _extract_usage(
|
||||
self,
|
||||
response_obj: ModelResponse,
|
||||
standard_logging_object: Optional[StandardLoggingPayload] = None,
|
||||
) -> Dict[str, int]:
|
||||
"""Extract usage statistics, preferring StandardLoggingPayload."""
|
||||
if standard_logging_object:
|
||||
prompt = standard_logging_object.get("prompt_tokens")
|
||||
completion = standard_logging_object.get("completion_tokens")
|
||||
total = standard_logging_object.get("total_tokens")
|
||||
if any(x is not None for x in [prompt, completion, total]):
|
||||
return {
|
||||
"prompt_tokens": prompt or 0,
|
||||
"completion_tokens": completion or 0,
|
||||
"total_tokens": total or 0,
|
||||
}
|
||||
|
||||
usage = response_obj.get("usage", None)
|
||||
if not usage:
|
||||
return {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}
|
||||
|
||||
return {
|
||||
"prompt_tokens": usage.get("prompt_tokens") or 0,
|
||||
"completion_tokens": usage.get("completion_tokens") or 0,
|
||||
"total_tokens": usage.get("total_tokens") or 0,
|
||||
}
|
||||
|
||||
def _get_finish_reason(self, response_obj: ModelResponse) -> str:
|
||||
"""
|
||||
Extract finish reason from first choice in the response.
|
||||
|
||||
Returns "unknown" if choices are not present or finish_reason is not found.
|
||||
"""
|
||||
choices = response_obj.get("choices") or []
|
||||
if choices and len(choices) > 0:
|
||||
return choices[0].get("finish_reason") or "unknown"
|
||||
return "unknown"
|
||||
|
||||
def _to_epoch_ms(self, t: Any) -> float:
|
||||
"""Convert a datetime or float timestamp to epoch milliseconds."""
|
||||
if hasattr(t, "timestamp"):
|
||||
return t.timestamp() * 1000.0
|
||||
return float(t) * 1000.0
|
||||
|
||||
def _get_duration(
|
||||
self,
|
||||
kwargs: Dict,
|
||||
start_time: Any,
|
||||
end_time: Any,
|
||||
standard_logging_object: Optional[StandardLoggingPayload] = None,
|
||||
) -> Optional[float]:
|
||||
"""
|
||||
Extract duration in milliseconds.
|
||||
|
||||
Resolution order:
|
||||
1. StandardLoggingPayload.response_time (already computed by LiteLLM)
|
||||
2. llm_api_duration_ms from kwargs
|
||||
3. Calculated from start_time and end_time
|
||||
"""
|
||||
if standard_logging_object:
|
||||
response_time = standard_logging_object.get("response_time")
|
||||
if response_time is not None:
|
||||
return (
|
||||
float(response_time) * 1000.0
|
||||
) # SLO stores seconds; convert to ms
|
||||
|
||||
duration_ms = kwargs.get("llm_api_duration_ms")
|
||||
if duration_ms is not None:
|
||||
return float(duration_ms)
|
||||
|
||||
if start_time is not None and end_time is not None:
|
||||
return self._to_epoch_ms(end_time) - self._to_epoch_ms(start_time)
|
||||
|
||||
return None
|
||||
|
||||
def _get_request_params(
|
||||
self,
|
||||
kwargs: Dict,
|
||||
standard_logging_object: Optional[StandardLoggingPayload] = None,
|
||||
) -> Dict[str, Any]:
|
||||
"""
|
||||
Extract request parameters like temperature and max_tokens, preferring
|
||||
StandardLoggingPayload.model_parameters.
|
||||
|
||||
Returns dict with available parameters, omitting those not present.
|
||||
"""
|
||||
if standard_logging_object:
|
||||
source_params = standard_logging_object.get("model_parameters") or {}
|
||||
else:
|
||||
source_params = kwargs.get("optional_params") or {}
|
||||
|
||||
params = {}
|
||||
|
||||
temperature = source_params.get("temperature")
|
||||
if temperature is not None:
|
||||
params["temperature"] = temperature
|
||||
|
||||
max_tokens = source_params.get("max_tokens")
|
||||
if max_tokens is not None:
|
||||
params["max_tokens"] = max_tokens
|
||||
|
||||
return params
|
||||
|
||||
def _extract_message_content(self, message: Union[Message, Dict]) -> str:
|
||||
"""
|
||||
Extract content from a message, handling various formats.
|
||||
|
||||
Handles tool calls, multimodal content (as JSON), and standard text content.
|
||||
Returns empty string if content is None or missing.
|
||||
"""
|
||||
content = message.get("content")
|
||||
|
||||
# Handle tool calls
|
||||
if message.get("tool_calls"):
|
||||
try:
|
||||
return json.dumps(message["tool_calls"])
|
||||
except Exception:
|
||||
return str(message["tool_calls"])
|
||||
|
||||
# Handle None or missing content
|
||||
if content is None:
|
||||
return ""
|
||||
|
||||
# Handle list content (multimodal)
|
||||
if isinstance(content, list):
|
||||
try:
|
||||
return json.dumps(content)
|
||||
except Exception:
|
||||
return str(content)
|
||||
|
||||
# Handle non-string content
|
||||
if not isinstance(content, str):
|
||||
return str(content)
|
||||
|
||||
return content
|
||||
|
||||
def _extract_all_messages(
|
||||
self,
|
||||
kwargs: Dict,
|
||||
response_obj: ModelResponse,
|
||||
response_model: str,
|
||||
vendor: str,
|
||||
standard_logging_object: Optional[StandardLoggingPayload] = None,
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Extract all messages (request + response) with sequence numbers and timestamps.
|
||||
|
||||
Processes request messages from StandardLoggingPayload.messages (preferred) or
|
||||
kwargs["messages"] (fallback), and response messages from response_obj["choices"].
|
||||
Assigns sequential numbers starting at 0.
|
||||
Adds timestamps from StandardLoggingPayload (preferred) or kwargs if available
|
||||
(converted to epoch milliseconds).
|
||||
"""
|
||||
messages = []
|
||||
sequence = 0
|
||||
|
||||
# Extract timestamps, preferring StandardLoggingPayload
|
||||
start_time = None
|
||||
if standard_logging_object:
|
||||
start_time = standard_logging_object.get("startTime")
|
||||
if not start_time:
|
||||
start_time = kwargs.get("start_time")
|
||||
|
||||
end_time = None
|
||||
if standard_logging_object:
|
||||
end_time = standard_logging_object.get("endTime")
|
||||
if not end_time:
|
||||
end_time = kwargs.get("end_time")
|
||||
|
||||
# Content is recorded only when the NR-specific switches allow it AND
|
||||
# LiteLLM's wider redaction decision (turn_off_message_logging, dynamic
|
||||
# params, headers) does not require redaction. Async streaming hands the
|
||||
# callback an unredacted async_complete_streaming_response, so without
|
||||
# this gate generated content would still reach NR even when the user
|
||||
# has globally disabled message logging.
|
||||
record_content = self.record_content and not should_redact_message_logging(
|
||||
kwargs
|
||||
)
|
||||
|
||||
# Extract request messages, preferring StandardLoggingPayload.
|
||||
# SLO messages can be a string (serialized/redacted), so only use it when it's a list.
|
||||
slo_messages = (
|
||||
standard_logging_object.get("messages") if standard_logging_object else None
|
||||
)
|
||||
if isinstance(slo_messages, list):
|
||||
request_messages = slo_messages
|
||||
else:
|
||||
request_messages = kwargs.get("messages") or []
|
||||
for msg in request_messages:
|
||||
message_data = {
|
||||
"role": msg.get("role") or "user",
|
||||
"sequence": sequence,
|
||||
"response.model": response_model,
|
||||
"vendor": vendor,
|
||||
}
|
||||
|
||||
# Add timestamp for request message if available (convert to milliseconds)
|
||||
if start_time is not None:
|
||||
message_data["timestamp"] = int(self._to_epoch_ms(start_time))
|
||||
|
||||
if record_content:
|
||||
message_data["content"] = self._extract_message_content(msg)
|
||||
|
||||
messages.append(message_data)
|
||||
sequence += 1
|
||||
|
||||
# Extract response messages from choices
|
||||
choices = response_obj.get("choices") or []
|
||||
if choices and len(choices) > 0:
|
||||
for choice in choices:
|
||||
# Prefer "message" (non-streaming); fall back to "delta" (streaming-assembled)
|
||||
message = choice.get("message", None) or choice.get("delta", None)
|
||||
if message:
|
||||
message_data = {
|
||||
"role": message.get("role") or "assistant",
|
||||
"sequence": sequence,
|
||||
"response.model": response_model,
|
||||
"vendor": vendor,
|
||||
"is_response": True,
|
||||
}
|
||||
|
||||
# Add timestamp for response message if available (convert to milliseconds)
|
||||
if end_time is not None:
|
||||
message_data["timestamp"] = int(self._to_epoch_ms(end_time))
|
||||
|
||||
if record_content:
|
||||
message_data["content"] = self._extract_message_content(message)
|
||||
|
||||
messages.append(message_data)
|
||||
sequence += 1
|
||||
|
||||
return messages
|
||||
|
||||
def _record_summary_event(
|
||||
self,
|
||||
request_id: str,
|
||||
trace_id: Optional[str],
|
||||
request_model: str,
|
||||
response_model: str,
|
||||
vendor: str,
|
||||
finish_reason: str,
|
||||
num_messages: int,
|
||||
usage: Dict[str, int],
|
||||
duration: Optional[float] = None,
|
||||
request_params: Optional[Dict[str, Any]] = None,
|
||||
):
|
||||
"""Record LlmChatCompletionSummary event to New Relic."""
|
||||
try:
|
||||
event_data = {
|
||||
"id": request_id,
|
||||
"request_id": request_id,
|
||||
"request.model": request_model,
|
||||
"response.model": response_model,
|
||||
"response.choices.finish_reason": finish_reason,
|
||||
"response.number_of_messages": num_messages,
|
||||
"vendor": vendor,
|
||||
"ingest_source": "litellm",
|
||||
"response.usage.prompt_tokens": usage["prompt_tokens"],
|
||||
"response.usage.completion_tokens": usage["completion_tokens"],
|
||||
"response.usage.total_tokens": usage["total_tokens"],
|
||||
}
|
||||
|
||||
# Add optional attributes if present
|
||||
if trace_id:
|
||||
event_data["trace_id"] = trace_id
|
||||
|
||||
if duration is not None:
|
||||
event_data["duration"] = duration
|
||||
|
||||
# Add request parameters if present
|
||||
if request_params:
|
||||
if "temperature" in request_params:
|
||||
event_data["request.temperature"] = request_params["temperature"]
|
||||
if "max_tokens" in request_params:
|
||||
event_data["request.max_tokens"] = request_params["max_tokens"]
|
||||
|
||||
app = _newrelic_agent.application()
|
||||
|
||||
if app and app.enabled:
|
||||
app.record_custom_event("LlmChatCompletionSummary", event_data)
|
||||
else:
|
||||
verbose_logger.warning(
|
||||
"New Relic application is not enabled; skipping summary event recording."
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
verbose_logger.warning(f"Failed to record New Relic summary event: {e}")
|
||||
self.handle_callback_failure("newrelic")
|
||||
|
||||
def _record_message_events(
|
||||
self,
|
||||
request_id: str,
|
||||
llm_response_id: str,
|
||||
trace_id: Optional[str],
|
||||
messages: List[Dict[str, Any]],
|
||||
):
|
||||
"""Record LlmChatCompletionMessage events to New Relic.
|
||||
|
||||
Args:
|
||||
request_id: Agent-generated UUID that links to Summary event's id
|
||||
llm_response_id: LLM's response ID (e.g., "chatcmpl-...") for message id format
|
||||
trace_id: Trace ID for distributed tracing (None if not available)
|
||||
messages: List of message dicts to record
|
||||
"""
|
||||
try:
|
||||
app = _newrelic_agent.application()
|
||||
|
||||
if not (app and app.enabled):
|
||||
verbose_logger.warning(
|
||||
"New Relic application is not enabled; skipping message event recording."
|
||||
)
|
||||
return
|
||||
|
||||
for message in messages:
|
||||
sequence = message["sequence"]
|
||||
event_data = {
|
||||
"id": f"{llm_response_id}-{sequence}",
|
||||
"request_id": request_id,
|
||||
"completion_id": request_id,
|
||||
"role": message["role"],
|
||||
"sequence": sequence,
|
||||
"response.model": message["response.model"],
|
||||
"vendor": message["vendor"],
|
||||
"ingest_source": "litellm",
|
||||
"token_count": 0, # Per-message token counts are not available from LiteLLM
|
||||
}
|
||||
|
||||
# Add trace context if available
|
||||
if trace_id:
|
||||
event_data["trace_id"] = trace_id
|
||||
|
||||
# Add content only if it was included in the message data
|
||||
if "content" in message:
|
||||
event_data["content"] = message["content"]
|
||||
|
||||
# Add is_response only if True (per spec, omit for request messages)
|
||||
if message.get("is_response"):
|
||||
event_data["is_response"] = True
|
||||
|
||||
# Forward actual request/response timestamp (ms) so NR uses the
|
||||
# real LLM call window rather than the async-logger fire time.
|
||||
# Requires newrelic>=11.2.0 which reads params["timestamp"] as
|
||||
# the intrinsic event timestamp.
|
||||
if "timestamp" in message:
|
||||
event_data["timestamp"] = message["timestamp"]
|
||||
|
||||
app.record_custom_event("LlmChatCompletionMessage", event_data)
|
||||
|
||||
except Exception as e:
|
||||
verbose_logger.warning(f"Failed to record New Relic message events: {e}")
|
||||
self.handle_callback_failure("newrelic")
|
||||
|
||||
def _record_error_metric(self):
|
||||
"""Record error metric to New Relic."""
|
||||
try:
|
||||
if not self.enabled:
|
||||
return
|
||||
|
||||
self._check_and_emit_periodic_metric()
|
||||
|
||||
app = _newrelic_agent.application()
|
||||
if app and app.enabled:
|
||||
app.record_custom_metric("LLM/LiteLLM/Error", 1)
|
||||
except Exception as e:
|
||||
verbose_logger.warning(f"Failed to record New Relic error metric: {e}")
|
||||
self.handle_callback_failure("newrelic")
|
||||
|
||||
def _process_success(
|
||||
self,
|
||||
kwargs: Dict,
|
||||
response_obj: ModelResponse,
|
||||
start_time: Optional[float] = None,
|
||||
end_time: Optional[float] = None,
|
||||
):
|
||||
"""
|
||||
Core logic for processing successful LLM calls.
|
||||
Used by both sync and async success event handlers.
|
||||
"""
|
||||
# Early exit if not enabled
|
||||
if not self.enabled:
|
||||
return
|
||||
|
||||
# Check and emit periodic supportability metric if 27 hours have passed
|
||||
self._check_and_emit_periodic_metric()
|
||||
|
||||
# Use StandardLoggingPayload where available for normalized, pre-computed values
|
||||
standard_logging_object: Optional[StandardLoggingPayload] = kwargs.get(
|
||||
"standard_logging_object"
|
||||
)
|
||||
|
||||
# Get trace context
|
||||
trace_id = self._get_trace_context(kwargs, standard_logging_object)
|
||||
|
||||
# Generate unique request ID for this request (used as Summary event id)
|
||||
request_id = str(uuid.uuid4())
|
||||
|
||||
# Extract data from response
|
||||
llm_response_id = self._extract_completion_id(kwargs, response_obj)
|
||||
vendor = self._get_vendor(kwargs, standard_logging_object)
|
||||
request_model, response_model = self._get_model_names(
|
||||
kwargs, response_obj, standard_logging_object
|
||||
)
|
||||
usage = self._extract_usage(response_obj, standard_logging_object)
|
||||
finish_reason = self._get_finish_reason(response_obj)
|
||||
|
||||
# Extract additional summary event fields
|
||||
duration = self._get_duration(
|
||||
kwargs, start_time, end_time, standard_logging_object
|
||||
)
|
||||
request_params = self._get_request_params(kwargs, standard_logging_object)
|
||||
|
||||
# Extract all messages
|
||||
messages = self._extract_all_messages(
|
||||
kwargs, response_obj, response_model, vendor, standard_logging_object
|
||||
)
|
||||
|
||||
# Record summary event
|
||||
self._record_summary_event(
|
||||
request_id=request_id,
|
||||
trace_id=trace_id,
|
||||
request_model=request_model,
|
||||
response_model=response_model,
|
||||
vendor=vendor,
|
||||
finish_reason=finish_reason,
|
||||
num_messages=len(messages),
|
||||
usage=usage,
|
||||
duration=duration,
|
||||
request_params=request_params,
|
||||
)
|
||||
|
||||
# Record message events
|
||||
self._record_message_events(
|
||||
request_id=request_id,
|
||||
llm_response_id=llm_response_id,
|
||||
trace_id=trace_id,
|
||||
messages=messages,
|
||||
)
|
||||
|
||||
async def async_health_check(self) -> IntegrationHealthCheckStatus:
|
||||
"""
|
||||
Check if the New Relic integration is healthy.
|
||||
|
||||
Verifies that the integration is enabled and the New Relic agent
|
||||
has an active, connected application, then records a small
|
||||
`LiteLLMConnectionTest` custom event so the user can confirm the
|
||||
end-to-end pipeline in the New Relic UI via NRQL:
|
||||
`SELECT * FROM LiteLLMConnectionTest SINCE 1 hour ago`.
|
||||
|
||||
The `LiteLLMConnectionTest` event type is intentionally outside the
|
||||
`Llm*` family that AI Monitoring queries, so test events do not
|
||||
appear in AI Monitoring dashboards.
|
||||
"""
|
||||
if not self.enabled:
|
||||
return IntegrationHealthCheckStatus(
|
||||
status="unhealthy",
|
||||
error_message="New Relic integration is disabled. Check that "
|
||||
"NEW_RELIC_LICENSE_KEY and NEW_RELIC_APP_NAME are set and the "
|
||||
"newrelic package is installed.",
|
||||
)
|
||||
|
||||
try:
|
||||
app = _newrelic_agent.application()
|
||||
if not (app and app.enabled):
|
||||
return IntegrationHealthCheckStatus(
|
||||
status="unhealthy",
|
||||
error_message=(
|
||||
"New Relic Python agent not installed. Review the New Relic integration documentation at https://docs.litellm.ai/docs/observability/newrelic."
|
||||
),
|
||||
)
|
||||
|
||||
app.record_custom_event(
|
||||
"LiteLLMConnectionTest",
|
||||
{
|
||||
"is_test_event": True,
|
||||
"app_name": self.app_name,
|
||||
"source": "litellm-proxy",
|
||||
"timestamp": time.time(),
|
||||
},
|
||||
)
|
||||
return IntegrationHealthCheckStatus(status="healthy", error_message=None)
|
||||
except Exception as e:
|
||||
return IntegrationHealthCheckStatus(
|
||||
status="unhealthy",
|
||||
error_message=str(e),
|
||||
)
|
||||
|
||||
# CustomLogger interface implementation
|
||||
|
||||
def log_pre_api_call(self, model, messages, kwargs):
|
||||
"""Unused per spec."""
|
||||
pass
|
||||
|
||||
def log_post_api_call(self, kwargs, response_obj, start_time, end_time):
|
||||
"""Unused per spec."""
|
||||
pass
|
||||
|
||||
def log_success_event(self, kwargs, response_obj, start_time, end_time):
|
||||
"""
|
||||
Main success path for non-streaming requests.
|
||||
|
||||
Note: New Relic's record_custom_event is synchronous but non-blocking
|
||||
(in-memory operation), so it's safe to call from sync context.
|
||||
"""
|
||||
try:
|
||||
self._process_success(kwargs, response_obj, start_time, end_time)
|
||||
except Exception as e:
|
||||
verbose_logger.warning(f"Error in New Relic log_success_event: {e}")
|
||||
self.handle_callback_failure("newrelic")
|
||||
|
||||
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
|
||||
"""
|
||||
Main success path for async/streaming requests.
|
||||
|
||||
Note: New Relic's SDK is thread-safe and record_custom_event is fast,
|
||||
so we can call it directly without asyncio.to_thread().
|
||||
"""
|
||||
try:
|
||||
self._process_success(kwargs, response_obj, start_time, end_time)
|
||||
except Exception as e:
|
||||
verbose_logger.warning(f"Error in New Relic async_log_success_event: {e}")
|
||||
self.handle_callback_failure("newrelic")
|
||||
|
||||
def log_failure_event(self, kwargs, response_obj, start_time, end_time):
|
||||
"""
|
||||
Log error metric for failed LLM calls (sync).
|
||||
|
||||
Per spec: Do not send AI events on failure, only record error metric.
|
||||
"""
|
||||
try:
|
||||
self._record_error_metric()
|
||||
|
||||
except Exception as e:
|
||||
verbose_logger.warning(f"Error in New Relic log_failure_event: {e}")
|
||||
self.handle_callback_failure("newrelic")
|
||||
|
||||
async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time):
|
||||
"""
|
||||
Log error metric for failed LLM calls (async).
|
||||
|
||||
Per spec: Do not send AI events on failure, only record error metric.
|
||||
"""
|
||||
try:
|
||||
self._record_error_metric()
|
||||
|
||||
except Exception as e:
|
||||
verbose_logger.warning(f"Error in New Relic async_log_failure_event: {e}")
|
||||
self.handle_callback_failure("newrelic")
|
||||
|
|
@ -65,7 +65,15 @@ class OpenMeterLogger(CustomLogger):
|
|||
"total_tokens": response_obj["usage"].get("total_tokens"),
|
||||
}
|
||||
|
||||
user_param = kwargs.get("user", None) # end-user passed in via 'user' param
|
||||
# OPENMETER_TRUST_REQUEST_USER (default "true"): when set to "false",
|
||||
# the request-supplied `user` field is ignored and the subject is
|
||||
# resolved solely from the key-bound user_api_key_user_id. Proxies
|
||||
# serving multi-tenant traffic enable this to prevent clients from
|
||||
# forging attribution by setting `user` in the request body.
|
||||
trust_request_user = (
|
||||
os.getenv("OPENMETER_TRUST_REQUEST_USER", "true").lower() != "false"
|
||||
)
|
||||
user_param = kwargs.get("user", None) if trust_request_user else None
|
||||
|
||||
# If no user provided directly, try to get it from token user_id
|
||||
if user_param is None:
|
||||
|
|
|
|||
|
|
@ -15,6 +15,7 @@ from litellm.integrations.otel.model.baggage import promoted_baggage
|
|||
from litellm.integrations.otel.model.config import OpenTelemetryV2Config
|
||||
from litellm.integrations.otel.plumbing.context import (
|
||||
is_recordable_span,
|
||||
request_root_span,
|
||||
resolve_parent_context,
|
||||
resolve_request_span_context,
|
||||
set_request_baggage,
|
||||
|
|
@ -435,8 +436,12 @@ class OpenTelemetryV2(CustomLogger):
|
|||
attach(set_request_baggage(bag, context=get_current()))
|
||||
# The server span was started by the instrumentor before this ran,
|
||||
# so the Baggage processor (which only fires at span start) won't
|
||||
# backfill it — stamp identity on it directly.
|
||||
server_span = get_current_span()
|
||||
# backfill it — stamp identity on it directly. Prefer the anchored
|
||||
# root span over the ambient one so identity still lands on the
|
||||
# server span when seeding from inside the live ``auth`` phase span
|
||||
# (the auth-failure path), where ``get_current_span`` is the phase
|
||||
# span, not the request's root.
|
||||
server_span = request_root_span() or get_current_span()
|
||||
if is_recordable_span(server_span):
|
||||
# Re-capture the anchor here too: this runs post-auth with the
|
||||
# server span active and covers entrypoints that bypass
|
||||
|
|
|
|||
|
|
@ -24,14 +24,18 @@ from typing import (
|
|||
|
||||
import litellm
|
||||
from litellm._logging import print_verbose, verbose_logger
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.integrations.prometheus_helpers.bounded_prometheus_series_tracker import (
|
||||
BoundedPrometheusSeriesTracker,
|
||||
from litellm.exceptions import (
|
||||
validate_rate_limit_category,
|
||||
validate_rate_limit_type,
|
||||
)
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.integrations.prometheus_helpers import (
|
||||
PrometheusLabelFactoryContext,
|
||||
_get_cached_end_user_id_for_cost_tracking,
|
||||
)
|
||||
from litellm.integrations.prometheus_helpers.bounded_prometheus_series_tracker import (
|
||||
BoundedPrometheusSeriesTracker,
|
||||
)
|
||||
from litellm.litellm_core_utils.core_helpers import (
|
||||
get_litellm_metadata_from_kwargs,
|
||||
get_metadata_variable_name_from_kwargs,
|
||||
|
|
@ -42,6 +46,9 @@ from litellm.proxy._types import (
|
|||
LiteLLM_UserTable,
|
||||
UserAPIKeyAuth,
|
||||
)
|
||||
from litellm.repositories.organization_repository import OrganizationRepository
|
||||
from litellm.repositories.team_repository import TeamRepository
|
||||
from litellm.repositories.user_repository import UserRepository
|
||||
from litellm.types.integrations.prometheus import *
|
||||
from litellm.types.integrations.prometheus import (
|
||||
_sanitize_prometheus_label_name,
|
||||
|
|
@ -78,6 +85,20 @@ class PrometheusLogger(CustomLogger):
|
|||
# Always initialize label_filters, even for non-premium users
|
||||
self.label_filters = self._parse_prometheus_config()
|
||||
|
||||
# Cache resolved label sets per metric. Several entries in
|
||||
# ``PrometheusMetricLabels.get_labels`` read module-level toggles
|
||||
# (e.g. ``litellm.prometheus_emit_stream_label``,
|
||||
# ``litellm.prometheus_emit_rate_limit_labels``) that can be
|
||||
# changed at runtime. Prometheus counters/gauges/histograms are
|
||||
# created with a *fixed* ``labelnames`` set; if a runtime call
|
||||
# to ``get_labels_for_metric`` returned a different set, the
|
||||
# subsequent ``counter.labels(**_labels)`` would raise a
|
||||
# ``ValueError`` from the prometheus client. Snapshotting at
|
||||
# logger init time pins the label set for the lifetime of the
|
||||
# logger so toggling these flags only takes effect after a
|
||||
# restart, keeping init-time and runtime label sets in sync.
|
||||
self._cached_metric_labels: Dict[str, List[str]] = {}
|
||||
|
||||
_custom_buckets = litellm.prometheus_latency_buckets
|
||||
self.latency_buckets = (
|
||||
tuple(_custom_buckets)
|
||||
|
|
@ -1033,13 +1054,27 @@ class PrometheusLogger(CustomLogger):
|
|||
self, metric_name: DEFINED_PROMETHEUS_METRICS
|
||||
) -> List[str]:
|
||||
"""
|
||||
Get the labels for a metric, filtered if configured
|
||||
Get the labels for a metric, filtered if configured.
|
||||
|
||||
The result is cached on the instance so the label set used to
|
||||
construct each Prometheus metric at ``__init__`` time stays in lock
|
||||
step with the label set passed to ``counter.labels(...)`` at
|
||||
runtime, even if the underlying module-level toggles consulted by
|
||||
:meth:`PrometheusMetricLabels.get_labels` (e.g.
|
||||
``litellm.prometheus_emit_rate_limit_labels``,
|
||||
``litellm.prometheus_emit_stream_label``) are flipped after the
|
||||
logger has been created.
|
||||
"""
|
||||
cached = self._cached_metric_labels.get(metric_name)
|
||||
if cached is not None:
|
||||
return cached
|
||||
|
||||
# Get default labels for this metric from PrometheusMetricLabels
|
||||
default_labels = PrometheusMetricLabels.get_labels(metric_name)
|
||||
|
||||
# If no label filtering is configured for this metric, use default labels
|
||||
if metric_name not in self.label_filters:
|
||||
self._cached_metric_labels[metric_name] = default_labels
|
||||
return default_labels
|
||||
|
||||
# Get configured labels for this metric
|
||||
|
|
@ -1050,6 +1085,7 @@ class PrometheusLogger(CustomLogger):
|
|||
label for label in default_labels if label in configured_labels
|
||||
]
|
||||
|
||||
self._cached_metric_labels[metric_name] = filtered_labels
|
||||
return filtered_labels
|
||||
|
||||
def _track_end_user_metric_series(
|
||||
|
|
@ -2029,14 +2065,8 @@ class PrometheusLogger(CustomLogger):
|
|||
|
||||
Proxy level tracking - failed client side requests
|
||||
|
||||
labelnames=[
|
||||
"end_user",
|
||||
"hashed_api_key",
|
||||
"api_key_alias",
|
||||
REQUESTED_MODEL,
|
||||
"team",
|
||||
"team_alias",
|
||||
] + EXCEPTION_LABELS,
|
||||
See :attr:`PrometheusMetricLabels.litellm_proxy_failed_requests_metric`
|
||||
for the authoritative list of labels emitted on this metric.
|
||||
"""
|
||||
from litellm.litellm_core_utils.litellm_logging import (
|
||||
StandardLoggingPayloadSetup,
|
||||
|
|
@ -2059,6 +2089,9 @@ class PrometheusLogger(CustomLogger):
|
|||
model_id = _metadata.get("model_info", {}).get("id") or request_data.get(
|
||||
"model_info", {}
|
||||
).get("id")
|
||||
rate_limit_category, rate_limit_type = self._extract_rate_limit_labels(
|
||||
original_exception
|
||||
)
|
||||
enum_values = UserAPIKeyLabelValues(
|
||||
end_user=user_api_key_dict.end_user_id,
|
||||
user=user_api_key_dict.user_id,
|
||||
|
|
@ -2073,6 +2106,8 @@ class PrometheusLogger(CustomLogger):
|
|||
status_code=str(status_code),
|
||||
exception_status=str(status_code),
|
||||
exception_class=self._get_exception_class_name(original_exception),
|
||||
rate_limit_category=rate_limit_category,
|
||||
rate_limit_type=rate_limit_type,
|
||||
tags=_tags,
|
||||
route=user_api_key_dict.request_route,
|
||||
client_ip=_metadata.get("requester_ip_address"),
|
||||
|
|
@ -2690,7 +2725,7 @@ class PrometheusLogger(CustomLogger):
|
|||
Args:
|
||||
guardrail_name: Name of the guardrail
|
||||
latency_seconds: Execution latency in seconds
|
||||
status: "success" or "error"
|
||||
status: "success", "error", or "intervened"
|
||||
error_type: Type of error if any, None otherwise
|
||||
hook_type: "pre_call", "during_call", or "post_call"
|
||||
"""
|
||||
|
|
@ -2843,6 +2878,33 @@ class PrometheusLogger(CustomLogger):
|
|||
|
||||
@staticmethod
|
||||
def _get_exception_class_name(exception: Exception) -> str:
|
||||
# Some exception types pin the ``exception_class`` label to a legacy
|
||||
# value for back-compat with existing dashboards (e.g. proxy-side 429s
|
||||
# keep reporting as "HTTPException"). Honor that opt-in marker before
|
||||
# deriving the label from the runtime class name. Reading it via
|
||||
# ``getattr`` keeps this core integrations module free of a transitive
|
||||
# ``fastapi`` dependency.
|
||||
legacy_class_name = getattr(exception, "prometheus_exception_class_name", None)
|
||||
if isinstance(legacy_class_name, str) and legacy_class_name:
|
||||
return legacy_class_name
|
||||
|
||||
# Same back-compat reasoning for ``BudgetExceededError``: the unified
|
||||
# rate-limit error work attached ``.llm_provider`` to budget errors
|
||||
# too (so callbacks reading ``StandardLoggingPayload`` get provider
|
||||
# attribution). Without this short-circuit, the provider prefix below
|
||||
# would silently flip the label from "BudgetExceededError" to e.g.
|
||||
# "Openai.BudgetExceededError" and break dashboards keyed on the
|
||||
# original value.
|
||||
try:
|
||||
from litellm.exceptions import BudgetExceededError
|
||||
except ImportError:
|
||||
BudgetExceededError = None # type: ignore[assignment,misc]
|
||||
|
||||
if BudgetExceededError is not None and isinstance(
|
||||
exception, BudgetExceededError
|
||||
):
|
||||
return "BudgetExceededError"
|
||||
|
||||
exception_class_name = ""
|
||||
if hasattr(exception, "llm_provider"):
|
||||
exception_class_name = getattr(exception, "llm_provider") or ""
|
||||
|
|
@ -2857,6 +2919,27 @@ class PrometheusLogger(CustomLogger):
|
|||
exception_class_name += exception.__class__.__name__
|
||||
return exception_class_name
|
||||
|
||||
@staticmethod
|
||||
def _extract_rate_limit_labels(
|
||||
exception: Optional[Exception],
|
||||
) -> Tuple[Optional[str], Optional[str]]:
|
||||
"""
|
||||
Pull the unified ``category`` / ``rate_limit_type`` fields off any
|
||||
exception that declares them (``litellm.RateLimitError`` and bare-
|
||||
Exception subclasses like ``BudgetExceededError``).
|
||||
|
||||
Values are validated against the :class:`RateLimitErrorCategory` /
|
||||
:class:`RateLimitType` enums so unrelated third-party exceptions that
|
||||
happen to declare ``.category`` / ``.rate_limit_type`` string attributes
|
||||
can't leak garbage into Prometheus label cardinality.
|
||||
"""
|
||||
if exception is None:
|
||||
return None, None
|
||||
return (
|
||||
validate_rate_limit_category(getattr(exception, "category", None)),
|
||||
validate_rate_limit_type(getattr(exception, "rate_limit_type", None)),
|
||||
)
|
||||
|
||||
async def log_success_fallback_event(
|
||||
self, original_model_group: str, kwargs: dict, original_exception: Exception
|
||||
):
|
||||
|
|
@ -3198,12 +3281,12 @@ class PrometheusLogger(CustomLogger):
|
|||
page_size: int, page: int
|
||||
) -> Tuple[List[LiteLLM_UserTable], Optional[int]]:
|
||||
skip = (page - 1) * page_size
|
||||
users = await prisma_client.db.litellm_usertable.find_many(
|
||||
users = await UserRepository(prisma_client).table.find_many(
|
||||
skip=skip,
|
||||
take=page_size,
|
||||
order={"created_at": "desc"},
|
||||
)
|
||||
total_count = await prisma_client.db.litellm_usertable.count()
|
||||
total_count = await UserRepository(prisma_client).table.count()
|
||||
return users, total_count
|
||||
|
||||
await self._initialize_budget_metrics(
|
||||
|
|
@ -3226,13 +3309,13 @@ class PrometheusLogger(CustomLogger):
|
|||
|
||||
async def fetch_orgs(page_size: int, page: int) -> Tuple[list, Optional[int]]:
|
||||
skip = (page - 1) * page_size
|
||||
orgs = await prisma_client.db.litellm_organizationtable.find_many(
|
||||
orgs = await OrganizationRepository(prisma_client).table.find_many(
|
||||
skip=skip,
|
||||
take=page_size,
|
||||
order={"created_at": "desc"},
|
||||
include={"litellm_budget_table": True},
|
||||
)
|
||||
total_count = await prisma_client.db.litellm_organizationtable.count()
|
||||
total_count = await OrganizationRepository(prisma_client).table.count()
|
||||
return orgs, total_count
|
||||
|
||||
await self._initialize_budget_metrics(
|
||||
|
|
@ -3300,14 +3383,14 @@ class PrometheusLogger(CustomLogger):
|
|||
|
||||
try:
|
||||
# Get total user count
|
||||
total_users = await prisma_client.db.litellm_usertable.count()
|
||||
total_users = await UserRepository(prisma_client).table.count()
|
||||
self.litellm_total_users_metric.set(total_users)
|
||||
verbose_logger.debug(
|
||||
f"Prometheus: set litellm_total_users to {total_users}"
|
||||
)
|
||||
|
||||
# Get total team count
|
||||
total_teams = await prisma_client.db.litellm_teamtable.count()
|
||||
total_teams = await TeamRepository(prisma_client).table.count()
|
||||
self.litellm_teams_count_metric.set(total_teams)
|
||||
verbose_logger.debug(
|
||||
f"Prometheus: set litellm_teams_count to {total_teams}"
|
||||
|
|
|
|||
|
|
@ -244,6 +244,9 @@ search_tools:
|
|||
- search_tool_name: "my-tavily-tool"
|
||||
litellm_params:
|
||||
search_provider: "tavily"
|
||||
- search_tool_name: "my-you-com-tool"
|
||||
litellm_params:
|
||||
search_provider: "you_com"
|
||||
```
|
||||
|
||||
---
|
||||
|
|
|
|||
|
|
@ -1339,8 +1339,13 @@ class WebSearchInterceptionLogger(CustomLogger):
|
|||
websearch_params: WebSearchInterceptionConfig = {}
|
||||
if "websearch_interception_params" in litellm_settings:
|
||||
websearch_params = litellm_settings["websearch_interception_params"]
|
||||
elif "websearch_interception" in callback_specific_params:
|
||||
websearch_params = callback_specific_params["websearch_interception"]
|
||||
elif "websearch_interception" in callback_specific_params and isinstance(
|
||||
callback_specific_params["websearch_interception"], dict
|
||||
):
|
||||
websearch_params = cast(
|
||||
WebSearchInterceptionConfig,
|
||||
callback_specific_params["websearch_interception"],
|
||||
)
|
||||
|
||||
# Use classmethod to initialize from config
|
||||
return WebSearchInterceptionLogger.from_config_yaml(websearch_params)
|
||||
|
|
|
|||
|
|
@ -2,11 +2,40 @@
|
|||
Utility functions for the Agents API SDK.
|
||||
"""
|
||||
|
||||
from typing import Optional
|
||||
from typing import Dict, Mapping, Optional
|
||||
|
||||
from litellm.llms.base_llm.agents.transformation import BaseAgentsAPIConfig
|
||||
|
||||
|
||||
def merge_agent_headers(
|
||||
*,
|
||||
dynamic_headers: Optional[Mapping[str, str]] = None,
|
||||
static_headers: Optional[Mapping[str, str]] = None,
|
||||
) -> Optional[Dict[str, str]]:
|
||||
"""Merge outbound HTTP headers for A2A agent calls.
|
||||
|
||||
Merge rules:
|
||||
- Start with ``dynamic_headers`` (values extracted from the incoming client request).
|
||||
- Overlay ``static_headers`` (admin-configured per agent).
|
||||
- Comparison is case-insensitive (HTTP headers are case-insensitive), so a
|
||||
static ``Authorization`` strips any dynamic ``authorization`` before the
|
||||
static value is written. The static side's casing is preserved.
|
||||
|
||||
If both contain the same header (case-insensitively), ``static_headers`` wins.
|
||||
"""
|
||||
merged: Dict[str, str] = {}
|
||||
|
||||
if dynamic_headers:
|
||||
merged.update({str(k): str(v) for k, v in dynamic_headers.items()})
|
||||
|
||||
if static_headers:
|
||||
static_lower = {str(k).lower() for k in static_headers}
|
||||
merged = {k: v for k, v in merged.items() if k.lower() not in static_lower}
|
||||
merged.update({str(k): str(v) for k, v in static_headers.items()})
|
||||
|
||||
return merged or None
|
||||
|
||||
|
||||
def get_provider_agents_api_config(
|
||||
custom_llm_provider: Optional[str],
|
||||
) -> Optional[BaseAgentsAPIConfig]:
|
||||
|
|
|
|||
|
|
@ -37,7 +37,7 @@ def get_litellm_gateway_api_key(
|
|||
"""
|
||||
Get the stored CLI API key for use with LiteLLM SDK.
|
||||
|
||||
This function reads the token file created by `litellm-proxy login`
|
||||
This function reads the token file created by `lite login`
|
||||
and returns the API key for use in Python scripts.
|
||||
|
||||
Args:
|
||||
|
|
|
|||
|
|
@ -25,6 +25,7 @@ from litellm.integrations.datadog.datadog_metrics import DatadogMetricsLogger
|
|||
from litellm.integrations.deepeval import DeepEvalLogger
|
||||
from litellm.integrations.dotprompt import DotpromptManager
|
||||
from litellm.integrations.focus.focus_logger import FocusLogger
|
||||
from litellm.integrations.mavvrik_focus.mavvrik_focus_logger import MavvrikFocusLogger
|
||||
from litellm.integrations.vantage.vantage_logger import VantageLogger
|
||||
from litellm.integrations.galileo import GalileoObserve
|
||||
from litellm.integrations.gcs_bucket.gcs_bucket import GCSBucketLogger
|
||||
|
|
@ -39,6 +40,7 @@ from litellm.integrations.langsmith import LangsmithLogger
|
|||
from litellm.integrations.litellm_agent import LiteLLMAgentModelResolver
|
||||
from litellm.integrations.literal_ai import LiteralAILogger
|
||||
from litellm.integrations.mlflow import MlflowLogger
|
||||
from litellm.integrations.newrelic import NewRelicLogger
|
||||
from litellm.integrations.openmeter import OpenMeterLogger
|
||||
from litellm.integrations.opentelemetry import OpenTelemetry
|
||||
from litellm.integrations.opik.opik import OpikLogger
|
||||
|
|
@ -102,8 +104,10 @@ class CustomLoggerRegistry:
|
|||
"gitlab": GitLabPromptManager,
|
||||
"cloudzero": CloudZeroLogger,
|
||||
"focus": FocusLogger,
|
||||
"mavvrik": MavvrikFocusLogger,
|
||||
"vantage": VantageLogger,
|
||||
"posthog": PostHogLogger,
|
||||
"newrelic": NewRelicLogger,
|
||||
}
|
||||
|
||||
try:
|
||||
|
|
|
|||
|
|
@ -131,6 +131,8 @@ def get_next_standardized_reset_time(
|
|||
# Handle different time units
|
||||
if unit == "d":
|
||||
return _handle_day_reset(current_time, base_midnight, value, tz)
|
||||
elif unit == "w":
|
||||
return _handle_day_reset(current_time, base_midnight, value * 7, tz)
|
||||
elif unit == "h":
|
||||
return _handle_hour_reset(current_time, base_midnight, value)
|
||||
elif unit == "m":
|
||||
|
|
|
|||
|
|
@ -655,7 +655,11 @@ def exception_type( # type: ignore # noqa: PLR0915
|
|||
custom_llm_provider == "anthropic"
|
||||
or custom_llm_provider == "anthropic_text"
|
||||
): # one of the anthropics
|
||||
if "prompt is too long" in error_str or "prompt: length" in error_str:
|
||||
if (
|
||||
"prompt is too long" in error_str
|
||||
or "prompt: length" in error_str
|
||||
or ExceptionCheckers.is_error_str_context_window_exceeded(error_str)
|
||||
):
|
||||
exception_mapping_worked = True
|
||||
raise ContextWindowExceededError(
|
||||
message="AnthropicError - {}".format(error_str),
|
||||
|
|
|
|||
|
|
@ -4,7 +4,7 @@ from litellm.llms.openai.data_residency import infer_openai_data_residency
|
|||
|
||||
# Pre-define optional kwargs keys as frozenset for O(1) lookups
|
||||
# These are extracted from kwargs only if present, avoiding unnecessary .get() calls
|
||||
_OPTIONAL_KWARGS_KEYS = frozenset(
|
||||
OPTIONAL_KWARGS_KEYS = frozenset(
|
||||
{
|
||||
"azure_ad_token",
|
||||
"tenant_id",
|
||||
|
|
@ -32,14 +32,19 @@ _OPTIONAL_KWARGS_KEYS = frozenset(
|
|||
"aws_sts_endpoint",
|
||||
"aws_external_id",
|
||||
"aws_bedrock_runtime_endpoint",
|
||||
"aws_bedrock_project_id",
|
||||
"gigachat_scope",
|
||||
"gigachat_auth_url",
|
||||
"gigachat_access_token",
|
||||
"tpm",
|
||||
"rpm",
|
||||
"use_xai_oauth",
|
||||
}
|
||||
)
|
||||
|
||||
# Backward-compatible alias for existing imports/tests.
|
||||
_OPTIONAL_KWARGS_KEYS = OPTIONAL_KWARGS_KEYS
|
||||
|
||||
|
||||
def _get_base_model_from_litellm_call_metadata(
|
||||
metadata: Optional[dict],
|
||||
|
|
@ -167,7 +172,7 @@ def get_litellm_params(
|
|||
|
||||
# Sparse extraction: only add kwargs keys that are actually present
|
||||
if kwargs:
|
||||
for key in _OPTIONAL_KWARGS_KEYS:
|
||||
for key in OPTIONAL_KWARGS_KEYS:
|
||||
if key in kwargs:
|
||||
litellm_params[key] = kwargs[key]
|
||||
|
||||
|
|
|
|||
|
|
@ -1,5 +1,5 @@
|
|||
import re
|
||||
from typing import Optional, Tuple
|
||||
from typing import Optional, Tuple, cast
|
||||
from urllib.parse import urlparse
|
||||
|
||||
import litellm
|
||||
|
|
@ -7,7 +7,7 @@ from litellm.constants import REPLICATE_MODEL_NAME_WITH_ID_LENGTH
|
|||
from litellm.llms.openai_like.json_loader import JSONProviderRegistry
|
||||
from litellm.secret_managers.main import get_secret, get_secret_str
|
||||
|
||||
from ..types.router import LiteLLM_Params
|
||||
from ..types.router import GenericLiteLLMParams, LiteLLM_Params
|
||||
|
||||
|
||||
def _endpoint_matches_api_base(endpoint: str, api_base: str) -> bool:
|
||||
|
|
@ -159,7 +159,7 @@ def get_llm_provider( # noqa: PLR0915
|
|||
custom_llm_provider: Optional[str] = None,
|
||||
api_base: Optional[str] = None,
|
||||
api_key: Optional[str] = None,
|
||||
litellm_params: Optional[LiteLLM_Params] = None,
|
||||
litellm_params: Optional[GenericLiteLLMParams] = None,
|
||||
) -> Tuple[str, str, Optional[str], Optional[str]]:
|
||||
"""
|
||||
Returns the provider for a given model name - e.g. 'azure/chatgpt-v-2' -> 'azure'
|
||||
|
|
@ -178,7 +178,7 @@ def get_llm_provider( # noqa: PLR0915
|
|||
)
|
||||
|
||||
if litellm.LiteLLMProxyChatConfig._should_use_litellm_proxy_by_default(
|
||||
litellm_params=litellm_params
|
||||
litellm_params=cast(Optional[LiteLLM_Params], litellm_params)
|
||||
):
|
||||
return litellm.LiteLLMProxyChatConfig.litellm_proxy_get_custom_llm_provider_info(
|
||||
model=model, api_base=api_base, api_key=api_key
|
||||
|
|
@ -186,12 +186,10 @@ def get_llm_provider( # noqa: PLR0915
|
|||
|
||||
## IF LITELLM PARAMS GIVEN ##
|
||||
if litellm_params:
|
||||
assert (
|
||||
custom_llm_provider is None and api_base is None and api_key is None
|
||||
), "Either pass in litellm_params or the custom_llm_provider/api_base/api_key. Otherwise, these values will be overriden."
|
||||
custom_llm_provider = litellm_params.custom_llm_provider
|
||||
api_base = litellm_params.api_base
|
||||
api_key = litellm_params.api_key
|
||||
if custom_llm_provider is None and api_base is None and api_key is None:
|
||||
custom_llm_provider = litellm_params.custom_llm_provider
|
||||
api_base = litellm_params.api_base
|
||||
api_key = litellm_params.api_key
|
||||
|
||||
dynamic_api_key = None
|
||||
# check if llm provider provided
|
||||
|
|
@ -235,6 +233,7 @@ def get_llm_provider( # noqa: PLR0915
|
|||
api_base=api_base,
|
||||
api_key=api_key,
|
||||
dynamic_api_key=dynamic_api_key,
|
||||
litellm_params=litellm_params,
|
||||
)
|
||||
|
||||
# check if llm provider part of model name
|
||||
|
|
@ -250,6 +249,7 @@ def get_llm_provider( # noqa: PLR0915
|
|||
api_base=api_base,
|
||||
api_key=api_key,
|
||||
dynamic_api_key=dynamic_api_key,
|
||||
litellm_params=litellm_params,
|
||||
)
|
||||
elif model.split("/", 1)[0] in litellm.provider_list:
|
||||
custom_llm_provider = model.split("/", 1)[0]
|
||||
|
|
@ -575,6 +575,7 @@ def _get_openai_compatible_provider_info( # noqa: PLR0915
|
|||
api_base: Optional[str],
|
||||
api_key: Optional[str],
|
||||
dynamic_api_key: Optional[str],
|
||||
litellm_params: Optional[GenericLiteLLMParams] = None,
|
||||
) -> Tuple[str, str, Optional[str], Optional[str]]:
|
||||
"""
|
||||
Returns:
|
||||
|
|
@ -642,7 +643,7 @@ def _get_openai_compatible_provider_info( # noqa: PLR0915
|
|||
api_base,
|
||||
dynamic_api_key,
|
||||
) = litellm.BedrockMantleChatConfig()._get_openai_compatible_provider_info(
|
||||
api_base, api_key
|
||||
api_base, api_key, litellm_params=litellm_params
|
||||
)
|
||||
elif custom_llm_provider == "nvidia_nim":
|
||||
# nvidia_nim is openai compatible, we just need to set this to custom_openai and have the api_base be https://api.endpoints.anyscale.com/v1
|
||||
|
|
@ -664,6 +665,11 @@ def _get_openai_compatible_provider_info( # noqa: PLR0915
|
|||
or get_secret_str("NVIDIA_RIVA_API_KEY")
|
||||
or get_secret_str("NVIDIA_NIM_API_KEY")
|
||||
)
|
||||
elif custom_llm_provider == "soniox":
|
||||
api_base = (
|
||||
api_base or get_secret_str("SONIOX_API_BASE") or "https://api.soniox.com"
|
||||
)
|
||||
dynamic_api_key = api_key or get_secret_str("SONIOX_API_KEY")
|
||||
elif custom_llm_provider == "cerebras":
|
||||
api_base = (
|
||||
api_base or get_secret("CEREBRAS_API_BASE") or "https://api.cerebras.ai/v1"
|
||||
|
|
|
|||
|
|
@ -22,9 +22,11 @@ def get_supported_openai_params( # noqa: PLR0915
|
|||
```
|
||||
|
||||
Args:
|
||||
base_model: For Azure, the true underlying model (e.g. ``"azure/gpt-5.2"``)
|
||||
when the deployment name differs. Used for model-type detection so that
|
||||
non-standard deployment names route to the correct config.
|
||||
base_model: An optional capability hint for deployments whose ``model``
|
||||
label isn't recognized on its own (e.g. an Azure deployment name, or a
|
||||
friendly Bedrock alias). It is additive: the result is the union of the
|
||||
params supported by ``model`` and by ``base_model``, so a hint can only
|
||||
add capabilities, never strip ones the real model already supports.
|
||||
|
||||
Returns:
|
||||
- List if custom_llm_provider is mapped
|
||||
|
|
@ -52,7 +54,15 @@ def get_supported_openai_params( # noqa: PLR0915
|
|||
provider_config = None
|
||||
|
||||
if provider_config and request_type == "chat_completion":
|
||||
return provider_config.get_supported_openai_params(model=base_model or model)
|
||||
supported_params = provider_config.get_supported_openai_params(model=model)
|
||||
if base_model and base_model != model:
|
||||
base_model_params = provider_config.get_supported_openai_params(
|
||||
model=base_model
|
||||
)
|
||||
supported_params = list(
|
||||
dict.fromkeys([*supported_params, *base_model_params])
|
||||
)
|
||||
return supported_params
|
||||
|
||||
if custom_llm_provider == "bedrock":
|
||||
return litellm.AmazonConverseConfig().get_supported_openai_params(model=model)
|
||||
|
|
@ -331,6 +341,11 @@ def get_supported_openai_params( # noqa: PLR0915
|
|||
return ElevenLabsAudioTranscriptionConfig().get_supported_openai_params(
|
||||
model=model
|
||||
)
|
||||
elif custom_llm_provider == "soniox":
|
||||
if request_type == "transcription":
|
||||
return litellm.SonioxAudioTranscriptionConfig().get_supported_openai_params(
|
||||
model=model
|
||||
)
|
||||
elif custom_llm_provider in litellm._custom_providers:
|
||||
if request_type == "chat_completion":
|
||||
provider_config = litellm.ProviderConfigManager.get_provider_chat_config(
|
||||
|
|
|
|||
|
|
@ -37,6 +37,10 @@ from litellm import (
|
|||
turn_off_message_logging,
|
||||
)
|
||||
from litellm._logging import _is_debugging_on, _redact_string, verbose_logger
|
||||
from litellm.exceptions import (
|
||||
validate_rate_limit_category,
|
||||
validate_rate_limit_type,
|
||||
)
|
||||
from litellm._uuid import uuid
|
||||
from litellm.batches.batch_utils import _handle_completed_batch
|
||||
from litellm.caching.caching import DualCache, InMemoryCache
|
||||
|
|
@ -154,6 +158,7 @@ from ..integrations.litellm_agent import LiteLLMAgentModelResolver
|
|||
from ..integrations.literal_ai import LiteralAILogger
|
||||
from ..integrations.logfire_logger import LogfireLevel, LogfireLogger
|
||||
from ..integrations.lunary import LunaryLogger
|
||||
from ..integrations.newrelic import NewRelicLogger
|
||||
from ..integrations.openmeter import OpenMeterLogger
|
||||
from ..integrations.opik.opik import OpikLogger
|
||||
from ..integrations.posthog import PostHogLogger
|
||||
|
|
@ -3528,6 +3533,14 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
elif isinstance(result, ModelResponse):
|
||||
return result
|
||||
|
||||
if isinstance(
|
||||
result,
|
||||
(ResponseCompletedEvent, ResponseIncompleteEvent, ResponseFailedEvent),
|
||||
):
|
||||
result = result.response
|
||||
if isinstance(result, ResponsesAPIResponse):
|
||||
return self._translate_responses_api_response_to_model_response(result)
|
||||
|
||||
httpx_response = self.model_call_details.get("httpx_response", None)
|
||||
if httpx_response and isinstance(httpx_response, httpx.Response):
|
||||
result = litellm.AnthropicConfig().transform_response(
|
||||
|
|
@ -3560,6 +3573,55 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
)
|
||||
return result
|
||||
|
||||
def _translate_responses_api_response_to_model_response(
|
||||
self, result: ResponsesAPIResponse
|
||||
) -> ModelResponse:
|
||||
"""
|
||||
Convert a Responses API response into a ModelResponse for spend_logs.
|
||||
|
||||
The proxy UI parses spend_log rows expecting chat-completion shape
|
||||
(response.choices[0].message); a raw ResponsesAPIResponse dump (output[...])
|
||||
would render as empty in the Logs tab. Translation also yields full
|
||||
choices/message detail downstream consumers can rely on.
|
||||
"""
|
||||
from litellm.completion_extras.litellm_responses_transformation.transformation import (
|
||||
LiteLLMResponsesTransformationHandler,
|
||||
)
|
||||
|
||||
try:
|
||||
return LiteLLMResponsesTransformationHandler().transform_response(
|
||||
model=self.model,
|
||||
raw_response=result,
|
||||
model_response=litellm.ModelResponse(),
|
||||
logging_obj=self,
|
||||
request_data={},
|
||||
messages=[],
|
||||
optional_params={},
|
||||
litellm_params={},
|
||||
encoding=litellm.encoding,
|
||||
)
|
||||
except Exception as e:
|
||||
verbose_logger.debug(
|
||||
"Responses API -> ModelResponse translation failed for "
|
||||
"anthropic_messages logging (%s); falling back to minimal "
|
||||
"usage-only ModelResponse to keep the spend_logs row.",
|
||||
str(e),
|
||||
)
|
||||
model_response = litellm.ModelResponse()
|
||||
model_response.model = self.model
|
||||
usage = getattr(result, "usage", None)
|
||||
if usage is not None and ResponseAPILoggingUtils._is_response_api_usage(
|
||||
usage
|
||||
):
|
||||
setattr(
|
||||
model_response,
|
||||
"usage",
|
||||
ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(
|
||||
usage
|
||||
),
|
||||
)
|
||||
return model_response
|
||||
|
||||
def _handle_non_streaming_google_genai_generate_content_response_logging(
|
||||
self, result: Any
|
||||
) -> ModelResponse:
|
||||
|
|
@ -4114,6 +4176,17 @@ def _init_custom_logger_compatible_class( # noqa: PLR0915
|
|||
focus_logger = FocusLogger()
|
||||
_in_memory_loggers.append(focus_logger)
|
||||
return focus_logger # type: ignore
|
||||
elif logging_integration == "mavvrik":
|
||||
from litellm.integrations.mavvrik_focus.mavvrik_focus_logger import (
|
||||
MavvrikFocusLogger,
|
||||
)
|
||||
|
||||
for callback in _in_memory_loggers:
|
||||
if type(callback) is MavvrikFocusLogger:
|
||||
return callback # type: ignore
|
||||
mavvrik_focus_logger = MavvrikFocusLogger()
|
||||
_in_memory_loggers.append(mavvrik_focus_logger)
|
||||
return mavvrik_focus_logger # type: ignore
|
||||
elif logging_integration == "vantage":
|
||||
from litellm.integrations.vantage.vantage_logger import VantageLogger
|
||||
|
||||
|
|
@ -4409,6 +4482,13 @@ def _init_custom_logger_compatible_class( # noqa: PLR0915
|
|||
gitlab_logger = GitLabPromptManager(gitlab_config=gitlab_config)
|
||||
_in_memory_loggers.append(gitlab_logger)
|
||||
return gitlab_logger # type: ignore
|
||||
elif logging_integration == "newrelic":
|
||||
for callback in _in_memory_loggers:
|
||||
if isinstance(callback, NewRelicLogger):
|
||||
return callback # type: ignore
|
||||
newrelic_logger = NewRelicLogger()
|
||||
_in_memory_loggers.append(newrelic_logger)
|
||||
return newrelic_logger # type: ignore
|
||||
return None
|
||||
except Exception as e:
|
||||
verbose_logger.exception(
|
||||
|
|
@ -4710,6 +4790,10 @@ def get_custom_logger_compatible_class( # noqa: PLR0915
|
|||
for callback in _in_memory_loggers:
|
||||
if isinstance(callback, SMTPEmailLogger):
|
||||
return callback
|
||||
elif logging_integration == "newrelic":
|
||||
for callback in _in_memory_loggers:
|
||||
if isinstance(callback, NewRelicLogger):
|
||||
return callback
|
||||
return None
|
||||
|
||||
except Exception as e:
|
||||
|
|
@ -5312,12 +5396,27 @@ class StandardLoggingPayloadSetup:
|
|||
else str(original_exception)
|
||||
)
|
||||
|
||||
# Duck-typed read so bare-Exception subclasses like
|
||||
# `litellm.BudgetExceededError` can participate without joining the
|
||||
# RateLimitError hierarchy (which would break `except BudgetExceededError`).
|
||||
# Validated against the enum value sets so a third-party exception that
|
||||
# happens to declare a `.category` or `.rate_limit_type` string attribute
|
||||
# can't leak garbage into the payload or Prometheus label cardinality.
|
||||
rate_limit_category = validate_rate_limit_category(
|
||||
getattr(original_exception, "category", None)
|
||||
)
|
||||
rate_limit_type = validate_rate_limit_type(
|
||||
getattr(original_exception, "rate_limit_type", None)
|
||||
)
|
||||
|
||||
return StandardLoggingPayloadErrorInformation(
|
||||
error_code=error_status,
|
||||
error_class=error_class,
|
||||
llm_provider=_llm_provider_in_exception,
|
||||
traceback=traceback_info,
|
||||
error_message=error_message if original_exception else "",
|
||||
error_rate_limit_category=rate_limit_category,
|
||||
error_rate_limit_type=rate_limit_type,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
|
|
|
|||
|
|
@ -6,6 +6,7 @@ from typing import Any, Dict, List, Literal, Optional, Tuple
|
|||
|
||||
import litellm
|
||||
from litellm.constants import OPENAI_FILE_SEARCH_COST_PER_1K_CALLS
|
||||
from litellm.litellm_core_utils.llm_cost_calc.utils import _get_web_search_requests
|
||||
from litellm.types.llms.openai import (
|
||||
FileSearchTool,
|
||||
ResponsesAPIResponse,
|
||||
|
|
@ -339,8 +340,7 @@ class StandardBuiltInToolCostTracking:
|
|||
# and _handle_web_search_cost() is never called.
|
||||
if (
|
||||
hasattr(usage, "server_tool_use")
|
||||
and usage.server_tool_use is not None
|
||||
and usage.server_tool_use.web_search_requests is not None
|
||||
and _get_web_search_requests(usage.server_tool_use) is not None
|
||||
):
|
||||
return True
|
||||
return False
|
||||
|
|
@ -352,8 +352,7 @@ class StandardBuiltInToolCostTracking:
|
|||
elif usage is not None:
|
||||
if (
|
||||
hasattr(usage, "server_tool_use")
|
||||
and usage.server_tool_use is not None
|
||||
and usage.server_tool_use.web_search_requests is not None
|
||||
and _get_web_search_requests(usage.server_tool_use) is not None
|
||||
):
|
||||
return True
|
||||
elif (
|
||||
|
|
|
|||
|
|
@ -1,7 +1,7 @@
|
|||
# What is this?
|
||||
## Helper utilities for cost_per_token()
|
||||
|
||||
from typing import Literal, Optional, Tuple, TypedDict, cast
|
||||
from typing import Any, Literal, Optional, Tuple, TypedDict, cast
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_logger
|
||||
|
|
@ -42,6 +42,26 @@ def _get_token_detail_value(details: object, key: str) -> Optional[int]:
|
|||
return value if isinstance(value, int) else None
|
||||
|
||||
|
||||
def _get_web_search_requests(server_tool_use: Any) -> Optional[int]:
|
||||
"""
|
||||
Tolerantly read ``web_search_requests`` from a ``server_tool_use`` value
|
||||
that may be ``None``, a ``dict``, a ``ServerToolUse`` pydantic instance,
|
||||
or any other object supporting attribute access.
|
||||
|
||||
Returns ``None`` when the value cannot be resolved — callers can
|
||||
distinguish "absent" from "zero" using ``is None``.
|
||||
|
||||
See https://github.com/BerriAI/litellm/issues/26153 — ``stream_chunk_builder``
|
||||
historically left this as a plain ``dict``, which broke direct attribute
|
||||
access in cost calculation.
|
||||
"""
|
||||
if server_tool_use is None:
|
||||
return None
|
||||
if isinstance(server_tool_use, dict):
|
||||
return server_tool_use.get("web_search_requests")
|
||||
return getattr(server_tool_use, "web_search_requests", None)
|
||||
|
||||
|
||||
def _is_above_128k(tokens: float) -> bool:
|
||||
if tokens > 128000:
|
||||
return True
|
||||
|
|
@ -929,6 +949,43 @@ def calculate_image_response_cost_from_usage(
|
|||
return prompt_cost + completion_cost
|
||||
|
||||
|
||||
def calculate_image_response_web_search_cost(
|
||||
image_response: ImageResponse,
|
||||
custom_llm_provider: str,
|
||||
model_info: ModelInfo,
|
||||
) -> float:
|
||||
"""
|
||||
Cost of Google Search grounding performed during image generation.
|
||||
|
||||
The grounding request count is carried on the image usage object by the
|
||||
provider transformers; it is billed with the same per-request accounting
|
||||
used for chat completions.
|
||||
"""
|
||||
usage = image_response.usage
|
||||
if usage is None:
|
||||
return 0.0
|
||||
|
||||
web_search_requests = getattr(usage, "web_search_requests", None)
|
||||
if not web_search_requests:
|
||||
return 0.0
|
||||
|
||||
from litellm.llms import get_cost_for_web_search_request
|
||||
|
||||
synthetic_usage = Usage(
|
||||
prompt_tokens_details=PromptTokensDetailsWrapper(
|
||||
web_search_requests=web_search_requests
|
||||
)
|
||||
)
|
||||
return (
|
||||
get_cost_for_web_search_request(
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
usage=synthetic_usage,
|
||||
model_info=model_info,
|
||||
)
|
||||
or 0.0
|
||||
)
|
||||
|
||||
|
||||
class CostCalculatorUtils:
|
||||
@staticmethod
|
||||
def _call_type_has_image_response(call_type: str) -> bool:
|
||||
|
|
|
|||
|
|
@ -633,11 +633,6 @@ def convert_to_model_response_object( # noqa: PLR0915
|
|||
thinking_blocks = choice["message"]["thinking_blocks"]
|
||||
provider_specific_fields["thinking_blocks"] = thinking_blocks
|
||||
|
||||
if reasoning_content:
|
||||
provider_specific_fields["reasoning_content"] = (
|
||||
reasoning_content
|
||||
)
|
||||
|
||||
message = Message(
|
||||
content=content,
|
||||
role=choice["message"]["role"] or "assistant",
|
||||
|
|
|
|||
|
|
@ -3653,17 +3653,13 @@ from litellm.types.llms.bedrock import ContentBlock as BedrockContentBlock
|
|||
from litellm.types.llms.bedrock import DocumentBlock as BedrockDocumentBlock
|
||||
from litellm.types.llms.bedrock import ImageBlock as BedrockImageBlock
|
||||
from litellm.types.llms.bedrock import SourceBlock as BedrockSourceBlock
|
||||
from litellm.types.llms.bedrock import BedrockToolSpec
|
||||
from litellm.types.llms.bedrock import ToolBlock as BedrockToolBlock
|
||||
from litellm.types.llms.bedrock import (
|
||||
ToolInputSchemaBlock as BedrockToolInputSchemaBlock,
|
||||
)
|
||||
from litellm.types.llms.bedrock import ToolJsonSchemaBlock as BedrockToolJsonSchemaBlock
|
||||
from litellm.types.llms.bedrock import SearchResultBlock
|
||||
from litellm.types.llms.bedrock import ToolResultBlock as BedrockToolResultBlock
|
||||
from litellm.types.llms.bedrock import (
|
||||
ToolResultContentBlock as BedrockToolResultContentBlock,
|
||||
)
|
||||
from litellm.types.llms.bedrock import ToolSpecBlock as BedrockToolSpecBlock
|
||||
from litellm.types.llms.bedrock import ToolUseBlock as BedrockToolUseBlock
|
||||
from litellm.types.llms.bedrock import VideoBlock as BedrockVideoBlock
|
||||
|
||||
|
|
@ -4294,6 +4290,49 @@ def _deduplicate_bedrock_tool_content(
|
|||
return _deduplicate_bedrock_content_blocks(tool_content, "toolResult")
|
||||
|
||||
|
||||
def _rename_duplicate_bedrock_document_names(
|
||||
contents: List[BedrockMessageBlock],
|
||||
) -> List[BedrockMessageBlock]:
|
||||
"""
|
||||
Rename duplicate document names across all messages in a Bedrock request.
|
||||
|
||||
Document names are derived from a content hash, so the same file appearing
|
||||
in multiple conversation turns produces identical names and Bedrock rejects
|
||||
the request with "Messages can not contain duplicate document names". The
|
||||
first occurrence keeps its original name so prompt-cache prefixes stay
|
||||
stable; later occurrences get a deterministic positional suffix
|
||||
(``_2``, ``_3``, ...), bumped further if the suffixed name already
|
||||
belongs to another document (e.g. an organic name ending in ``_2``).
|
||||
"""
|
||||
used_names: Set[str] = set()
|
||||
for message in contents:
|
||||
for block in message.get("content") or []:
|
||||
document = block.get("document")
|
||||
if isinstance(document, dict) and document.get("name"):
|
||||
used_names.add(document["name"])
|
||||
|
||||
name_counts: Dict[str, int] = {}
|
||||
for message in contents:
|
||||
for block in message.get("content") or []:
|
||||
document = block.get("document")
|
||||
if not isinstance(document, dict):
|
||||
continue
|
||||
name = document.get("name")
|
||||
if not name:
|
||||
continue
|
||||
count = name_counts.get(name, 0) + 1
|
||||
name_counts[name] = count
|
||||
if count > 1:
|
||||
suffix = count
|
||||
new_name = f"{name}_{suffix}"
|
||||
while new_name in used_names:
|
||||
suffix += 1
|
||||
new_name = f"{name}_{suffix}"
|
||||
used_names.add(new_name)
|
||||
document["name"] = new_name
|
||||
return contents
|
||||
|
||||
|
||||
def _sort_bedrock_assistant_content_blocks(
|
||||
blocks: List[BedrockContentBlock],
|
||||
) -> List[BedrockContentBlock]:
|
||||
|
|
@ -4702,6 +4741,12 @@ class BedrockConverseMessagesProcessor:
|
|||
guardContent={"text": {"text": element["text"]}}
|
||||
)
|
||||
_parts.append(_part)
|
||||
elif element["type"] in ("grounding_source", "query"):
|
||||
# Contextual grounding tags are guardrail metadata; the
|
||||
# model only needs the underlying text, so render them
|
||||
# as plain text on the generate path.
|
||||
_part = BedrockContentBlock(text=element["text"])
|
||||
_parts.append(_part)
|
||||
elif element["type"] == "image_url":
|
||||
format: Optional[str] = None
|
||||
if isinstance(element["image_url"], dict):
|
||||
|
|
@ -4942,7 +4987,7 @@ class BedrockConverseMessagesProcessor:
|
|||
llm_provider=llm_provider,
|
||||
)
|
||||
|
||||
return contents
|
||||
return _rename_duplicate_bedrock_document_names(contents)
|
||||
|
||||
@staticmethod
|
||||
def translate_thinking_blocks_to_reasoning_content_blocks(
|
||||
|
|
@ -5134,6 +5179,12 @@ def _bedrock_converse_messages_pt( # noqa: PLR0915
|
|||
guardContent={"text": {"text": element["text"]}}
|
||||
)
|
||||
_parts.append(_part)
|
||||
elif element["type"] in ("grounding_source", "query"):
|
||||
# Contextual grounding tags are guardrail metadata; the
|
||||
# model only needs the underlying text, so render them as
|
||||
# plain text on the generate path.
|
||||
_part = BedrockContentBlock(text=element["text"])
|
||||
_parts.append(_part)
|
||||
elif element["type"] == "image_url":
|
||||
format: Optional[str] = None
|
||||
if isinstance(element["image_url"], dict):
|
||||
|
|
@ -5364,7 +5415,7 @@ def _bedrock_converse_messages_pt( # noqa: PLR0915
|
|||
llm_provider=llm_provider,
|
||||
)
|
||||
|
||||
return contents
|
||||
return _rename_duplicate_bedrock_document_names(contents)
|
||||
|
||||
|
||||
def make_valid_bedrock_tool_name(input_tool_name: str) -> str:
|
||||
|
|
@ -5496,6 +5547,7 @@ def _bedrock_tools_pt(
|
|||
]
|
||||
"""
|
||||
from litellm.llms.bedrock.common_utils import (
|
||||
get_bedrock_base_model,
|
||||
normalize_json_schema_custom_types_to_object,
|
||||
)
|
||||
from litellm.litellm_core_utils.prompt_templates.common_utils import unpack_defs
|
||||
|
|
@ -5503,6 +5555,11 @@ def _bedrock_tools_pt(
|
|||
_valid_json_schema_root_types = frozenset(
|
||||
("array", "boolean", "integer", "null", "number", "object", "string")
|
||||
)
|
||||
# Only Claude on Bedrock honours strict tool schemas; other families
|
||||
# (Nova, Llama, GPT-OSS) reject the strict field outright.
|
||||
supports_strict_tools = bool(
|
||||
model and get_bedrock_base_model(model).startswith("anthropic")
|
||||
)
|
||||
tool_block_list: List[BedrockToolBlock] = []
|
||||
for tool_idx, tool in enumerate(tools):
|
||||
# Check if tool is already a BedrockToolBlock (e.g., systemTool for Nova grounding)
|
||||
|
|
@ -5548,17 +5605,16 @@ def _bedrock_tools_pt(
|
|||
normalize_json_schema_custom_types_to_object(parameters)
|
||||
if parameters.get("type") not in _valid_json_schema_root_types:
|
||||
parameters["type"] = "object"
|
||||
tool_input_schema = BedrockToolInputSchemaBlock(
|
||||
json=BedrockToolJsonSchemaBlock(
|
||||
type=parameters["type"],
|
||||
properties=parameters.get("properties", {}),
|
||||
required=parameters.get("required", []),
|
||||
)
|
||||
tool_block = cast(
|
||||
BedrockToolBlock,
|
||||
BedrockToolSpec(
|
||||
name=name,
|
||||
description=description,
|
||||
parameters=parameters,
|
||||
strict=tool.get("function", {}).get("strict", None),
|
||||
supports_strict_tools=supports_strict_tools,
|
||||
),
|
||||
)
|
||||
tool_spec = BedrockToolSpecBlock(
|
||||
inputSchema=tool_input_schema, name=name, description=description
|
||||
)
|
||||
tool_block = BedrockToolBlock(toolSpec=tool_spec)
|
||||
tool_block_list.append(tool_block)
|
||||
|
||||
## ADD CACHE POINT TOOL BLOCK ##
|
||||
|
|
|
|||
|
|
@ -47,6 +47,7 @@ class RealTimeStreaming:
|
|||
user_api_key_dict: Optional[Any] = None,
|
||||
request_data: Optional[Dict] = None,
|
||||
backend_uses_beta_protocol: Optional[bool] = None,
|
||||
force_transcription_model: Optional[str] = None,
|
||||
):
|
||||
self.websocket = websocket
|
||||
self.backend_ws = backend_ws
|
||||
|
|
@ -100,6 +101,11 @@ class RealTimeStreaming:
|
|||
self._flushing_pending_messages_until_setup: bool = False
|
||||
self._pending_messages_until_setup: List[str] = []
|
||||
self._pending_messages_byte_total: int = 0
|
||||
# Whether this is a transcription-only session (session.type == "transcription",
|
||||
# e.g. gpt-realtime-whisper). Such sessions must not be sent response.create and
|
||||
# their input_audio_transcription.completed usage drives duration-based cost.
|
||||
self._force_transcription_model = force_transcription_model
|
||||
self._is_transcription_session: bool = force_transcription_model is not None
|
||||
|
||||
# Per-connection caps for pre-setup audio frames (message count + total bytes).
|
||||
_MAX_BUFFERED_MESSAGES: int = 200
|
||||
|
|
@ -111,8 +117,12 @@ class RealTimeStreaming:
|
|||
"input_audio_buffer.append",
|
||||
"input_audio_buffer.commit",
|
||||
"input_audio_buffer.clear",
|
||||
"input_audio_buffer.end",
|
||||
]
|
||||
)
|
||||
_CLIENT_AUDIO_BUFFER_COMMIT_TYPES = frozenset(
|
||||
["input_audio_buffer.commit", "input_audio_buffer.end"]
|
||||
)
|
||||
_AUDIO_FORMAT_MAP: Dict[str, Dict[str, Any]] = {
|
||||
"pcm16": {"type": "audio/pcm", "rate": 24000},
|
||||
"g711_ulaw": {"type": "audio/G711-ulaw", "rate": 8000},
|
||||
|
|
@ -144,7 +154,7 @@ class RealTimeStreaming:
|
|||
return True
|
||||
return False
|
||||
|
||||
def store_message(self, message: Union[str, bytes, OpenAIRealtimeEvents]):
|
||||
def store_message(self, message: Union[str, bytes, dict, OpenAIRealtimeEvents]):
|
||||
"""Store message in list"""
|
||||
if isinstance(message, bytes):
|
||||
message = message.decode("utf-8")
|
||||
|
|
@ -154,22 +164,20 @@ class RealTimeStreaming:
|
|||
else:
|
||||
message_obj = cast(Dict[str, Any], json.loads(cast(str, message)))
|
||||
self._collect_tool_calls_from_response_done(cast(dict, message_obj))
|
||||
if not self._should_store_message(message_obj):
|
||||
return
|
||||
try:
|
||||
event_type = message_obj.get("type", "")
|
||||
if event_type in self._SESSION_EVENT_TYPES:
|
||||
typed_obj = OpenAIRealtimeStreamSessionEvents(**message_obj) # type: ignore
|
||||
typed_obj: OpenAIRealtimeEvents = OpenAIRealtimeStreamSessionEvents(**message_obj) # type: ignore
|
||||
else:
|
||||
# Use the base object as a safe catch-all for all other event types
|
||||
# (both beta and GA), so unknown/new event names never raise here.
|
||||
# Catch-all base object so unknown/new event names never raise.
|
||||
typed_obj = OpenAIRealtimeStreamResponseBaseObject(**message_obj) # type: ignore
|
||||
except Exception as e:
|
||||
verbose_logger.debug(f"Error parsing message for logging: {e}")
|
||||
# Don't re-raise — a parse failure must not drop or delay the message
|
||||
if self._should_store_message(message_obj):
|
||||
self.messages.append(message_obj) # type: ignore[arg-type]
|
||||
self.messages.append(message_obj) # type: ignore[arg-type]
|
||||
return
|
||||
if self._should_store_message(typed_obj):
|
||||
self.messages.append(typed_obj)
|
||||
self.messages.append(typed_obj)
|
||||
|
||||
def _collect_user_input_from_client_event(self, message: Union[str, dict]) -> None:
|
||||
"""Extract user text content from client WebSocket events for spend logging."""
|
||||
|
|
@ -209,6 +217,8 @@ class RealTimeStreaming:
|
|||
self.session_tools = tools
|
||||
# GA: session.type is required; log it for traceability but no action needed
|
||||
verbose_logger.debug(f"Realtime session.type: {session.get('type')}")
|
||||
if session.get("type") == "transcription":
|
||||
self._is_transcription_session = True
|
||||
except (json.JSONDecodeError, AttributeError, TypeError):
|
||||
pass
|
||||
|
||||
|
|
@ -225,6 +235,55 @@ class RealTimeStreaming:
|
|||
except (AttributeError, TypeError):
|
||||
pass
|
||||
|
||||
def _detect_transcription_session_from_backend(
|
||||
self, event_obj: Union[dict, OpenAIRealtimeEvents]
|
||||
) -> None:
|
||||
"""Flag transcription-only sessions from backend session events."""
|
||||
try:
|
||||
event_type = event_obj.get("type", "")
|
||||
if event_type in (
|
||||
"transcription_session.created",
|
||||
"transcription_session.updated",
|
||||
):
|
||||
self._is_transcription_session = True
|
||||
elif event_type in ("session.created", "session.updated"):
|
||||
session = cast(dict, event_obj).get("session", {}) or {}
|
||||
if session.get("type") == "transcription":
|
||||
self._is_transcription_session = True
|
||||
except (AttributeError, TypeError):
|
||||
pass
|
||||
|
||||
def _capture_transcription_usage(
|
||||
self, event_obj: Union[dict, OpenAIRealtimeEvents]
|
||||
) -> None:
|
||||
"""
|
||||
Append a usage-only transcription completed event to the logged results so
|
||||
the cost calculator can bill it by audio duration. The default logged event
|
||||
types exclude this event, so it is captured here directly for transcription
|
||||
sessions rather than widening logging for every realtime session. Only the
|
||||
type and usage are kept — the transcript is already captured separately in
|
||||
input_messages, so it is not duplicated into the response log here.
|
||||
"""
|
||||
try:
|
||||
usage = event_obj.get("usage")
|
||||
if usage is None:
|
||||
return
|
||||
# If this event type is already captured by store_message (e.g. the user
|
||||
# logs all realtime events), don't append a second copy.
|
||||
if self._should_store_message(event_obj):
|
||||
return
|
||||
self.messages.append(
|
||||
cast(
|
||||
OpenAIRealtimeEvents,
|
||||
{
|
||||
"type": "conversation.item.input_audio_transcription.completed",
|
||||
"usage": usage,
|
||||
},
|
||||
)
|
||||
)
|
||||
except (AttributeError, TypeError):
|
||||
pass
|
||||
|
||||
def _collect_tool_calls_from_response_done(
|
||||
self, event_obj: Union[dict, OpenAIRealtimeEvents]
|
||||
) -> None:
|
||||
|
|
@ -285,6 +344,7 @@ class RealTimeStreaming:
|
|||
backend, False if the provider transformation produced no output and
|
||||
the message was effectively dropped.
|
||||
"""
|
||||
message = self._enforce_transcription_session_model(message)
|
||||
if self.provider_config:
|
||||
transformed = self.provider_config.transform_realtime_request(
|
||||
message, self.model, self.session_configuration_request
|
||||
|
|
@ -304,12 +364,128 @@ class RealTimeStreaming:
|
|||
await self.backend_ws.send(message) # type: ignore[union-attr, attr-defined]
|
||||
return True
|
||||
|
||||
def _enforce_transcription_session_model(self, message: str) -> str:
|
||||
"""Force client transcription session updates to the authorized model.
|
||||
|
||||
`/v1/realtime?intent=transcription` may intentionally omit `model` from
|
||||
the upstream URL for Azure compatibility, but the proxy still authorizes
|
||||
a resolved LiteLLM model before opening the backend websocket. If a
|
||||
client later sends a transcription `session.update`, any model embedded
|
||||
in that update must be rewritten to the same authorized model instead of
|
||||
allowing a post-auth model/deployment switch.
|
||||
|
||||
Normal realtime sessions keep their independent nested transcription
|
||||
model behavior because `_force_transcription_model` is only set for
|
||||
transcription-intent websocket routes.
|
||||
"""
|
||||
if self._force_transcription_model is None:
|
||||
return message
|
||||
|
||||
try:
|
||||
message_obj = json.loads(message)
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
return message
|
||||
|
||||
if message_obj.get("type") not in (
|
||||
"session.update",
|
||||
"transcription_session.update",
|
||||
):
|
||||
return message
|
||||
|
||||
session = message_obj.get("session")
|
||||
if not isinstance(session, dict):
|
||||
return message
|
||||
|
||||
if session.get("type") == "transcription":
|
||||
self._is_transcription_session = True
|
||||
|
||||
authorized_model = self._force_transcription_model
|
||||
changed = False
|
||||
|
||||
transcription = session.get("input_audio_transcription")
|
||||
if (
|
||||
isinstance(transcription, dict)
|
||||
and transcription.get("model") != authorized_model
|
||||
):
|
||||
session["input_audio_transcription"] = {
|
||||
**transcription,
|
||||
"model": authorized_model,
|
||||
}
|
||||
changed = True
|
||||
|
||||
audio = session.get("audio")
|
||||
if isinstance(audio, dict):
|
||||
audio_input = audio.get("input")
|
||||
if isinstance(audio_input, dict):
|
||||
nested_transcription = audio_input.get("transcription")
|
||||
if (
|
||||
isinstance(nested_transcription, dict)
|
||||
and nested_transcription.get("model") != authorized_model
|
||||
):
|
||||
session["audio"] = {
|
||||
**audio,
|
||||
"input": {
|
||||
**audio_input,
|
||||
"transcription": {
|
||||
**nested_transcription,
|
||||
"model": authorized_model,
|
||||
},
|
||||
},
|
||||
}
|
||||
changed = True
|
||||
|
||||
if not changed:
|
||||
return message
|
||||
return json.dumps(message_obj)
|
||||
|
||||
def _uses_deferred_backend_setup(self) -> bool:
|
||||
"""True when setup is deferred until the client's first session.update."""
|
||||
if self.provider_config is None:
|
||||
return False
|
||||
return not self.provider_config.requires_session_configuration()
|
||||
|
||||
@staticmethod
|
||||
def _collapse_buffered_audio_messages(messages: List[str]) -> List[str]:
|
||||
"""Apply ``input_audio_buffer.clear`` semantics before replaying buffered frames.
|
||||
|
||||
During deferred Gemini Live setup, ``clear`` is buffered alongside appends.
|
||||
On flush each append becomes a provider ``realtimeInput``; ``clear`` must
|
||||
drop preceding uncommitted appends instead of being forwarded as a no-op.
|
||||
"""
|
||||
collapsed: List[str] = []
|
||||
pending_appends: List[str] = []
|
||||
|
||||
for message in messages:
|
||||
try:
|
||||
msg_type = json.loads(message).get("type")
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
collapsed.extend(pending_appends)
|
||||
pending_appends = []
|
||||
collapsed.append(message)
|
||||
continue
|
||||
|
||||
if msg_type == "input_audio_buffer.append":
|
||||
pending_appends.append(message)
|
||||
elif msg_type == "input_audio_buffer.clear":
|
||||
pending_appends = []
|
||||
elif msg_type in RealTimeStreaming._CLIENT_AUDIO_BUFFER_COMMIT_TYPES:
|
||||
collapsed.extend(pending_appends)
|
||||
pending_appends = []
|
||||
collapsed.append(message)
|
||||
else:
|
||||
collapsed.extend(pending_appends)
|
||||
pending_appends = []
|
||||
collapsed.append(message)
|
||||
|
||||
collapsed.extend(pending_appends)
|
||||
return collapsed
|
||||
|
||||
def _sync_pending_messages_byte_total(self) -> None:
|
||||
self._pending_messages_byte_total = sum(
|
||||
len(message.encode("utf-8"))
|
||||
for message in self._pending_messages_until_setup
|
||||
)
|
||||
|
||||
def _should_buffer_client_message_until_setup(self, message: str) -> bool:
|
||||
if not self._uses_deferred_backend_setup():
|
||||
return False
|
||||
|
|
@ -325,6 +501,18 @@ class RealTimeStreaming:
|
|||
return msg_obj.get("type") in RealTimeStreaming._CLIENT_AUDIO_BUFFER_TYPES
|
||||
|
||||
def _buffer_pending_message_until_setup(self, message: str) -> None:
|
||||
try:
|
||||
msg_type = json.loads(message).get("type")
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
msg_type = None
|
||||
|
||||
if msg_type == "input_audio_buffer.clear":
|
||||
self._pending_messages_until_setup = self._collapse_buffered_audio_messages(
|
||||
self._pending_messages_until_setup + [message]
|
||||
)
|
||||
self._sync_pending_messages_byte_total()
|
||||
return
|
||||
|
||||
msg_bytes = len(message.encode("utf-8"))
|
||||
if (
|
||||
len(self._pending_messages_until_setup)
|
||||
|
|
@ -342,7 +530,9 @@ class RealTimeStreaming:
|
|||
)
|
||||
|
||||
async def _flush_pending_messages_until_setup(self) -> bool:
|
||||
pending = self._pending_messages_until_setup
|
||||
pending = self._collapse_buffered_audio_messages(
|
||||
self._pending_messages_until_setup
|
||||
)
|
||||
self._pending_messages_until_setup = []
|
||||
self._pending_messages_byte_total = 0
|
||||
for idx, message in enumerate(pending):
|
||||
|
|
@ -358,8 +548,7 @@ class RealTimeStreaming:
|
|||
for msg in self._pending_messages_until_setup
|
||||
)
|
||||
verbose_logger.debug(
|
||||
"Failed to flush buffered client message after setup: %s "
|
||||
"(%d buffered message(s) retained)",
|
||||
"Failed to flush buffered client message after setup: %s (%d buffered message(s) retained)",
|
||||
e,
|
||||
len(unsent),
|
||||
)
|
||||
|
|
@ -376,8 +565,7 @@ class RealTimeStreaming:
|
|||
return True
|
||||
except Exception as e:
|
||||
verbose_logger.warning(
|
||||
"Failed to translate %s to beta protocol, forwarding "
|
||||
"untranslated event to client: %s",
|
||||
"Failed to translate %s to beta protocol, forwarding untranslated event to client: %s",
|
||||
event.get("type"),
|
||||
e,
|
||||
)
|
||||
|
|
@ -429,16 +617,13 @@ class RealTimeStreaming:
|
|||
if sent:
|
||||
self._guardrail_turn_detection_update_sent = True
|
||||
|
||||
def _has_realtime_guardrails(self) -> bool:
|
||||
"""Return True if any callback is registered for realtime guardrail event types."""
|
||||
def _has_realtime_guardrails_for_event_hooks(
|
||||
self,
|
||||
event_hooks: List[Any],
|
||||
) -> bool:
|
||||
"""Return True if any callback would run for one of ``event_hooks``."""
|
||||
from litellm.integrations.custom_guardrail import CustomGuardrail
|
||||
from litellm.types.guardrails import GuardrailEventHooks
|
||||
|
||||
_realtime_event_types = [
|
||||
GuardrailEventHooks.realtime_input_transcription,
|
||||
GuardrailEventHooks.pre_call,
|
||||
GuardrailEventHooks.post_call,
|
||||
]
|
||||
return any(
|
||||
isinstance(cb, CustomGuardrail)
|
||||
and any(
|
||||
|
|
@ -446,31 +631,45 @@ class RealTimeStreaming:
|
|||
data=self.request_data,
|
||||
event_type=et,
|
||||
)
|
||||
for et in _realtime_event_types
|
||||
for et in event_hooks
|
||||
)
|
||||
for cb in litellm.callbacks
|
||||
)
|
||||
|
||||
def _has_realtime_guardrails(self) -> bool:
|
||||
"""Return True if any callback is registered for realtime guardrail event types."""
|
||||
from litellm.types.guardrails import GuardrailEventHooks
|
||||
|
||||
return self._has_realtime_guardrails_for_event_hooks(
|
||||
[
|
||||
GuardrailEventHooks.realtime_input_transcription,
|
||||
GuardrailEventHooks.pre_call,
|
||||
GuardrailEventHooks.post_call,
|
||||
]
|
||||
)
|
||||
|
||||
def _has_audio_transcription_guardrails(self) -> bool:
|
||||
"""Return True if any callback needs to run on audio transcriptions (VAD path).
|
||||
"""Return True when a guardrail is configured for the audio/VAD transcript path.
|
||||
|
||||
When this returns True, we inject a session.update to disable the LLM's
|
||||
auto-response so the guardrail can gate it first.
|
||||
|
||||
Must match the same hook criteria as run_realtime_guardrails() so that
|
||||
any guardrail that would actually check the transcript also disables
|
||||
auto-response before the transcript arrives.
|
||||
Only ``realtime_input_transcription`` hooks disable ``server_vad`` auto-response.
|
||||
``pre_call`` / ``post_call`` guardrails (e.g. Model Armor on chat completions)
|
||||
must not override ``turn_detection.create_response`` on realtime sessions.
|
||||
"""
|
||||
return self._has_realtime_guardrails()
|
||||
from litellm.types.guardrails import GuardrailEventHooks
|
||||
|
||||
return self._has_realtime_guardrails_for_event_hooks(
|
||||
[GuardrailEventHooks.realtime_input_transcription]
|
||||
)
|
||||
|
||||
async def run_realtime_guardrails(
|
||||
self,
|
||||
transcript: str,
|
||||
item_id: Optional[str] = None,
|
||||
pre_block_backend_message: Optional[str] = None,
|
||||
event_hooks: Optional[List[Any]] = None,
|
||||
) -> bool:
|
||||
"""
|
||||
Run registered guardrails on a completed speech transcription.
|
||||
Run registered guardrails on realtime text (transcript, user message, tool output).
|
||||
|
||||
Returns True if blocked (synthetic warning already sent to client).
|
||||
Returns False if clean (caller should send response.create to the backend).
|
||||
|
|
@ -481,15 +680,17 @@ class RealTimeStreaming:
|
|||
specific message to be sent first — e.g. Gemini Live requires a
|
||||
matching ``toolResponse`` immediately after a ``toolCall`` before any
|
||||
other client messages can be accepted.
|
||||
|
||||
``event_hooks`` selects which guardrail modes to evaluate. Audio/VAD
|
||||
transcript completion uses ``realtime_input_transcription`` only;
|
||||
typed user messages and tool outputs use ``pre_call``.
|
||||
"""
|
||||
from litellm.integrations.custom_guardrail import CustomGuardrail
|
||||
from litellm.types.guardrails import GuardrailEventHooks
|
||||
|
||||
_realtime_event_types = [
|
||||
GuardrailEventHooks.realtime_input_transcription,
|
||||
GuardrailEventHooks.pre_call,
|
||||
GuardrailEventHooks.post_call,
|
||||
]
|
||||
if event_hooks is None:
|
||||
event_hooks = [GuardrailEventHooks.realtime_input_transcription]
|
||||
_realtime_event_types = event_hooks
|
||||
_check_data = {**self.request_data, "transcript": transcript}
|
||||
_already_run: set = set()
|
||||
|
||||
|
|
@ -705,48 +906,58 @@ class RealTimeStreaming:
|
|||
self.store_message(event_str)
|
||||
await self._send_event_to_client(event, event_str)
|
||||
|
||||
async def _handle_raw_backend_message(self, raw_response) -> bool:
|
||||
@staticmethod
|
||||
def _parse_backend_event(raw_response: str) -> Optional[dict]:
|
||||
"""Parse a backend frame once. Returns None for non-JSON or non-object frames."""
|
||||
try:
|
||||
event = json.loads(raw_response)
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
return None
|
||||
return event if isinstance(event, dict) else None
|
||||
|
||||
async def _handle_raw_backend_message(
|
||||
self, event_obj: dict, raw_response: str
|
||||
) -> bool:
|
||||
"""Process a backend message without provider_config (raw path).
|
||||
|
||||
Returns True if the caller should skip the default store+forward (i.e. continue the loop).
|
||||
"""
|
||||
try:
|
||||
event_obj = json.loads(raw_response)
|
||||
event_type = event_obj.get("type")
|
||||
|
||||
# For audio/VAD guardrail path: once the session is ready, tell the backend
|
||||
# not to auto-respond after VAD detects end-of-speech. We send the
|
||||
# session.created to the client FIRST so the client is always in sync, then
|
||||
# inject the session.update so a potential error from the backend doesn't
|
||||
# arrive before the client sees session.created.
|
||||
if (
|
||||
event_obj.get("type") == "session.created"
|
||||
and self._has_audio_transcription_guardrails()
|
||||
):
|
||||
self.store_message(raw_response)
|
||||
await self.websocket.send_text(raw_response)
|
||||
await self._send_to_backend(self._make_disable_auto_response_message())
|
||||
self._detect_transcription_session_from_backend(event_obj)
|
||||
|
||||
# Send session.created to the client FIRST so it stays in sync, then inject
|
||||
# the disable-auto-response session.update; otherwise a backend error could
|
||||
# reach the client before it sees session.created.
|
||||
if (
|
||||
event_type == "session.created"
|
||||
and self._has_audio_transcription_guardrails()
|
||||
):
|
||||
self.store_message(event_obj)
|
||||
await self.websocket.send_text(raw_response)
|
||||
await self._send_to_backend(self._make_disable_auto_response_message())
|
||||
return True
|
||||
|
||||
if event_type == "conversation.item.input_audio_transcription.completed":
|
||||
transcript = event_obj.get("transcript", "")
|
||||
self._collect_user_input_from_backend_event(event_obj)
|
||||
self.store_message(event_obj)
|
||||
await self.websocket.send_text(raw_response)
|
||||
|
||||
# Transcription-only sessions (e.g. gpt-realtime-whisper) have no
|
||||
# assistant turn: capture audio-duration usage for cost and never
|
||||
# trigger response.create.
|
||||
if self._is_transcription_session:
|
||||
self._capture_transcription_usage(event_obj)
|
||||
return True
|
||||
|
||||
if (
|
||||
event_obj.get("type")
|
||||
== "conversation.item.input_audio_transcription.completed"
|
||||
):
|
||||
transcript = event_obj.get("transcript", "")
|
||||
self._collect_user_input_from_backend_event(event_obj)
|
||||
## LOGGING — must happen before continue below
|
||||
self.store_message(raw_response)
|
||||
# Forward transcript to client so user sees what they said
|
||||
await self.websocket.send_text(raw_response)
|
||||
blocked = await self.run_realtime_guardrails(
|
||||
transcript,
|
||||
item_id=event_obj.get("item_id"),
|
||||
)
|
||||
if not blocked:
|
||||
# Clean — trigger LLM response
|
||||
await self._send_to_backend(json.dumps({"type": "response.create"}))
|
||||
return True
|
||||
except (json.JSONDecodeError, AttributeError):
|
||||
pass
|
||||
blocked = await self.run_realtime_guardrails(
|
||||
transcript,
|
||||
item_id=event_obj.get("item_id"),
|
||||
)
|
||||
if not blocked:
|
||||
await self._send_to_backend(json.dumps({"type": "response.create"}))
|
||||
return True
|
||||
return False
|
||||
|
||||
async def backend_to_client_send_messages(self):
|
||||
|
|
@ -779,25 +990,25 @@ class RealTimeStreaming:
|
|||
)
|
||||
continue
|
||||
else:
|
||||
handled = await self._handle_raw_backend_message(raw_response)
|
||||
if handled:
|
||||
continue
|
||||
## LOGGING
|
||||
self.store_message(raw_response)
|
||||
|
||||
# If the client opted into beta protocol, translate GA event
|
||||
# names/shapes back to the beta equivalents before forwarding.
|
||||
if self._client_wants_beta:
|
||||
try:
|
||||
event_dict = json.loads(raw_response)
|
||||
translated = self._translate_event_to_beta(event_dict)
|
||||
if translated is None:
|
||||
continue # drop GA-only events (e.g. conversation.item.done)
|
||||
await self.websocket.send_text(json.dumps(translated))
|
||||
except Exception:
|
||||
await self.websocket.send_text(raw_response)
|
||||
else:
|
||||
event = self._parse_backend_event(raw_response)
|
||||
if event is None:
|
||||
await self.websocket.send_text(raw_response)
|
||||
continue
|
||||
|
||||
if await self._handle_raw_backend_message(event, raw_response):
|
||||
continue
|
||||
self.store_message(event)
|
||||
|
||||
if not self._client_wants_beta:
|
||||
await self.websocket.send_text(raw_response)
|
||||
continue
|
||||
|
||||
translated = self._translate_event_to_beta(event)
|
||||
if translated is None:
|
||||
continue
|
||||
await self.websocket.send_text(
|
||||
raw_response if translated is event else json.dumps(translated)
|
||||
)
|
||||
|
||||
except websockets.exceptions.ConnectionClosed as e: # type: ignore
|
||||
verbose_logger.exception(
|
||||
|
|
@ -927,41 +1138,43 @@ class RealTimeStreaming:
|
|||
def _translate_event_to_beta(event: dict) -> Optional[dict]:
|
||||
"""Translate a single GA event dict to its beta equivalent.
|
||||
|
||||
Returns None if the event should be dropped entirely (e.g. the GA-only
|
||||
conversation.item.done has no beta counterpart).
|
||||
Returns the (possibly mutated copy of the) event otherwise.
|
||||
Returns None when the event must be dropped (the GA-only
|
||||
conversation.item.done has no beta counterpart). Returns the original
|
||||
event object unchanged when no translation applies, so the caller can
|
||||
forward the raw frame without re-serializing; otherwise returns a
|
||||
translated copy.
|
||||
"""
|
||||
event_type = event.get("type", "")
|
||||
|
||||
# conversation.item.done has no beta equivalent — the client already
|
||||
# received conversation.item.created (translated from .added).
|
||||
if event_type == "conversation.item.done":
|
||||
return None
|
||||
|
||||
# Shallow-copy so we don't mutate the stored message
|
||||
renamed_type = RealTimeStreaming._GA_TO_BETA_EVENT_TYPES.get(event_type)
|
||||
has_item = isinstance(event.get("item"), dict)
|
||||
response = event.get("response")
|
||||
has_response_output = isinstance(response, dict) and isinstance(
|
||||
response.get("output"), list
|
||||
)
|
||||
if renamed_type is None and not has_item and not has_response_output:
|
||||
return event
|
||||
|
||||
translated = dict(event)
|
||||
|
||||
# Rename the type field
|
||||
if event_type in RealTimeStreaming._GA_TO_BETA_EVENT_TYPES:
|
||||
translated["type"] = RealTimeStreaming._GA_TO_BETA_EVENT_TYPES[event_type]
|
||||
|
||||
# Fix content block types inside items (response.done output list,
|
||||
# conversation.item.created item content, etc.)
|
||||
if "item" in translated and isinstance(translated["item"], dict):
|
||||
if renamed_type is not None:
|
||||
translated["type"] = renamed_type
|
||||
if has_item:
|
||||
translated["item"] = RealTimeStreaming._translate_item_content_types(
|
||||
dict(translated["item"])
|
||||
)
|
||||
if "response" in translated and isinstance(translated["response"], dict):
|
||||
if has_response_output:
|
||||
resp = dict(translated["response"])
|
||||
if "output" in resp and isinstance(resp["output"], list):
|
||||
resp["output"] = [
|
||||
(
|
||||
RealTimeStreaming._translate_item_content_types(dict(o))
|
||||
if isinstance(o, dict)
|
||||
else o
|
||||
)
|
||||
for o in resp["output"]
|
||||
]
|
||||
resp["output"] = [
|
||||
(
|
||||
RealTimeStreaming._translate_item_content_types(dict(o))
|
||||
if isinstance(o, dict)
|
||||
else o
|
||||
)
|
||||
for o in resp["output"]
|
||||
]
|
||||
translated["response"] = resp
|
||||
|
||||
return translated
|
||||
|
|
@ -994,6 +1207,8 @@ class RealTimeStreaming:
|
|||
guardrail_turn_detection_injected = False
|
||||
msg_type: Optional[str] = None
|
||||
try:
|
||||
from litellm.types.guardrails import GuardrailEventHooks
|
||||
|
||||
msg_obj = json.loads(message)
|
||||
msg_type = msg_obj.get("type")
|
||||
|
||||
|
|
@ -1046,6 +1261,7 @@ class RealTimeStreaming:
|
|||
blocked = await self.run_realtime_guardrails(
|
||||
output_text,
|
||||
pre_block_backend_message=sanitized_msg,
|
||||
event_hooks=[GuardrailEventHooks.pre_call],
|
||||
)
|
||||
if blocked:
|
||||
# ``_pending_guardrail_message`` is
|
||||
|
|
@ -1071,7 +1287,8 @@ class RealTimeStreaming:
|
|||
combined_text = " ".join(texts)
|
||||
if combined_text:
|
||||
blocked = await self.run_realtime_guardrails(
|
||||
combined_text
|
||||
combined_text,
|
||||
event_hooks=[GuardrailEventHooks.pre_call],
|
||||
)
|
||||
if blocked:
|
||||
# Store the guardrail reason so the next response.create
|
||||
|
|
|
|||
|
|
@ -17,6 +17,10 @@ from litellm.integrations.custom_logger import CustomLogger
|
|||
from litellm.litellm_core_utils.core_helpers import (
|
||||
get_metadata_variable_name_from_kwargs,
|
||||
)
|
||||
from litellm.llms.vertex_ai.common_utils import (
|
||||
redact_vertex_ai_metadata_from_litellm_params,
|
||||
redact_vertex_ai_metadata_from_logged_object,
|
||||
)
|
||||
from litellm.secret_managers.main import str_to_bool
|
||||
from litellm.types.utils import StandardCallbackDynamicParams
|
||||
|
||||
|
|
@ -119,10 +123,12 @@ def _redact_standard_logging_object(model_call_details: dict):
|
|||
# ResponsesAPIResponse format - redact content in output items
|
||||
if isinstance(response.get("output"), list):
|
||||
_redact_responses_api_output_dict(response["output"], redacted_str)
|
||||
redact_vertex_ai_metadata_from_logged_object(response)
|
||||
elif isinstance(response, dict) and "choices" in response:
|
||||
# ModelResponse dict format - redact content in choices
|
||||
if isinstance(response.get("choices"), list):
|
||||
_redact_model_response_dict_choices(response["choices"], redacted_str)
|
||||
redact_vertex_ai_metadata_from_logged_object(response)
|
||||
elif isinstance(response, str):
|
||||
standard_logging_object["response"] = redacted_str
|
||||
else:
|
||||
|
|
@ -164,6 +170,7 @@ def perform_redaction(model_call_details: dict, result):
|
|||
model_call_details["prompt"] = ""
|
||||
model_call_details["input"] = ""
|
||||
_redact_standard_logging_object(model_call_details)
|
||||
redact_vertex_ai_metadata_from_litellm_params(model_call_details)
|
||||
|
||||
# Redact streaming response
|
||||
if (
|
||||
|
|
@ -174,6 +181,7 @@ def perform_redaction(model_call_details: dict, result):
|
|||
if hasattr(_streaming_response, "choices"):
|
||||
for choice in _streaming_response.choices:
|
||||
_redact_choice_content(choice)
|
||||
redact_vertex_ai_metadata_from_logged_object(_streaming_response)
|
||||
elif hasattr(_streaming_response, "output"):
|
||||
_redact_responses_api_output(_streaming_response.output)
|
||||
# Redact reasoning field in ResponsesAPIResponse
|
||||
|
|
@ -200,12 +208,14 @@ def perform_redaction(model_call_details: dict, result):
|
|||
if hasattr(_result, "choices") and _result.choices is not None:
|
||||
for choice in _result.choices:
|
||||
_redact_choice_content(choice)
|
||||
redact_vertex_ai_metadata_from_logged_object(_result)
|
||||
elif isinstance(_result, dict) and "choices" in _result:
|
||||
# Handle dict representation of ModelResponse (e.g., from model_dump())
|
||||
if _result.get("choices") is not None:
|
||||
_redact_model_response_dict_choices(
|
||||
_result["choices"], "redacted-by-litellm"
|
||||
)
|
||||
redact_vertex_ai_metadata_from_logged_object(_result)
|
||||
elif isinstance(_result, dict) and "output" in _result:
|
||||
if isinstance(_result.get("output"), list):
|
||||
_redact_responses_api_output_dict(
|
||||
|
|
|
|||
|
|
@ -20,6 +20,7 @@ from litellm.types.utils import (
|
|||
ServerToolUse,
|
||||
Usage,
|
||||
)
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.utils import print_verbose, token_counter
|
||||
|
||||
if TYPE_CHECKING:
|
||||
|
|
@ -79,6 +80,54 @@ class ChunkProcessor:
|
|||
model_response._hidden_params = chunk.get("_hidden_params", {})
|
||||
return model_response
|
||||
|
||||
@staticmethod
|
||||
def apply_provider_assembled_streaming_metadata(
|
||||
response: ModelResponse,
|
||||
chunks: List[Any],
|
||||
logging_obj: Optional[Any] = None,
|
||||
) -> None:
|
||||
if not chunks:
|
||||
return
|
||||
|
||||
model = getattr(response, "model", None)
|
||||
if not model:
|
||||
return
|
||||
|
||||
custom_llm_provider = None
|
||||
if logging_obj is not None:
|
||||
custom_llm_provider = logging_obj.model_call_details.get(
|
||||
"custom_llm_provider"
|
||||
)
|
||||
|
||||
try:
|
||||
from litellm.litellm_core_utils.get_llm_provider_logic import (
|
||||
get_llm_provider,
|
||||
)
|
||||
from litellm.types.utils import LlmProviders
|
||||
from litellm.utils import ProviderConfigManager
|
||||
|
||||
if custom_llm_provider:
|
||||
provider = LlmProviders(custom_llm_provider)
|
||||
else:
|
||||
_, provider_str, _, _ = get_llm_provider(model)
|
||||
provider = LlmProviders(provider_str)
|
||||
|
||||
provider_config = ProviderConfigManager.get_provider_chat_config(
|
||||
model=model,
|
||||
provider=provider,
|
||||
)
|
||||
if provider_config is not None:
|
||||
provider_config.apply_assembled_streaming_response_metadata(
|
||||
response=response,
|
||||
chunks=chunks,
|
||||
)
|
||||
except Exception as e:
|
||||
verbose_logger.debug(
|
||||
"apply_provider_assembled_streaming_metadata failed for model=%s: %s",
|
||||
model,
|
||||
e,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _get_chunk_id(chunks: List[Dict[str, Any]]) -> str:
|
||||
"""
|
||||
|
|
@ -588,7 +637,18 @@ class ChunkProcessor:
|
|||
hasattr(usage_chunk, "server_tool_use")
|
||||
and usage_chunk.server_tool_use is not None
|
||||
):
|
||||
server_tool_use = usage_chunk.server_tool_use
|
||||
# Coerce dict to ServerToolUse so downstream cost-calc code
|
||||
# (which accesses .web_search_requests as an attribute)
|
||||
# doesn't raise AttributeError. Some providers / streaming
|
||||
# paths leave server_tool_use as a plain dict on the chunk.
|
||||
if isinstance(usage_chunk.server_tool_use, dict):
|
||||
server_tool_use = ServerToolUse(**usage_chunk.server_tool_use)
|
||||
elif isinstance(usage_chunk.server_tool_use, ServerToolUse):
|
||||
server_tool_use = usage_chunk.server_tool_use
|
||||
else:
|
||||
server_tool_use = ServerToolUse.model_validate(
|
||||
usage_chunk.server_tool_use
|
||||
)
|
||||
if (
|
||||
usage_chunk_dict["prompt_tokens_details"] is not None
|
||||
and getattr(
|
||||
|
|
|
|||
|
|
@ -1149,6 +1149,32 @@ class CustomStreamWrapper:
|
|||
completion_obj: Dict[str, Any] = {"content": ""}
|
||||
from litellm.types.utils import GenericStreamingChunk as GChunk
|
||||
|
||||
if (
|
||||
isinstance(chunk, ModelResponseStream)
|
||||
and self.custom_llm_provider is not None
|
||||
and self.custom_llm_provider in litellm._custom_providers
|
||||
):
|
||||
_has_content = bool(
|
||||
chunk.choices
|
||||
and chunk.choices[0].delta is not None
|
||||
and (
|
||||
chunk.choices[0].delta.content
|
||||
or chunk.choices[0].delta.tool_calls
|
||||
)
|
||||
)
|
||||
if self.received_finish_reason is not None:
|
||||
if not _has_content:
|
||||
raise StopIteration
|
||||
if chunk.choices and chunk.choices[0].finish_reason:
|
||||
self.received_finish_reason = chunk.choices[0].finish_reason
|
||||
if not _has_content:
|
||||
return None
|
||||
# Strip finish_reason from the content chunk so it appears
|
||||
# only on the trailing empty-delta chunk (OpenAI spec).
|
||||
# finish_reason_handler() will emit the proper terminal chunk.
|
||||
chunk.choices[0].finish_reason = None # type: ignore[assignment]
|
||||
return chunk
|
||||
|
||||
if (
|
||||
isinstance(chunk, dict)
|
||||
and generic_chunk_has_all_required_fields(
|
||||
|
|
|
|||
|
|
@ -81,7 +81,6 @@ from litellm.types.utils import (
|
|||
from litellm.utils import (
|
||||
ModelResponse,
|
||||
Usage,
|
||||
_supports_factory,
|
||||
add_dummy_tool,
|
||||
any_assistant_message_has_thinking_blocks,
|
||||
get_max_tokens,
|
||||
|
|
@ -337,50 +336,6 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
|
|||
v in model_lower for v in ("opus-4-7", "opus_4_7", "opus-4.7", "opus_4.7")
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _supports_model_capability(model: str, key: str) -> bool:
|
||||
"""Check a boolean capability ``key`` in the model map.
|
||||
|
||||
Strips bedrock/vertex prefixes so a provider-routed Claude still
|
||||
resolves to the Anthropic model-map entry.
|
||||
"""
|
||||
try:
|
||||
if _supports_factory(
|
||||
model=model,
|
||||
custom_llm_provider="anthropic",
|
||||
key=key,
|
||||
):
|
||||
return True
|
||||
except Exception:
|
||||
pass
|
||||
candidates = [model]
|
||||
for prefix in (
|
||||
"bedrock/converse/",
|
||||
"bedrock/invoke/",
|
||||
"bedrock/",
|
||||
"vertex_ai/",
|
||||
):
|
||||
if model.startswith(prefix):
|
||||
candidates.append(model[len(prefix) :])
|
||||
try:
|
||||
from litellm.llms.bedrock.common_utils import BedrockModelInfo
|
||||
|
||||
base = BedrockModelInfo.get_base_model(model)
|
||||
if base:
|
||||
candidates.append(base)
|
||||
candidates.append(f"bedrock/{base}")
|
||||
except Exception:
|
||||
pass
|
||||
try:
|
||||
for cand in candidates:
|
||||
if cand in litellm.model_cost and (
|
||||
litellm.model_cost[cand].get(key) is True
|
||||
):
|
||||
return True
|
||||
except Exception:
|
||||
pass
|
||||
return False
|
||||
|
||||
@staticmethod
|
||||
def _supports_effort_level(model: str, level: str) -> bool:
|
||||
"""Check ``supports_{level}_reasoning_effort`` in the model map."""
|
||||
|
|
@ -918,7 +873,39 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
|
|||
anthropic_tools = []
|
||||
mcp_servers = []
|
||||
for tool in tools:
|
||||
if "input_schema" in tool: # assume in anthropic format
|
||||
if tool.get("type") == "namespace":
|
||||
# Namespace is a grouping container (e.g. codex's multi_agent_v1).
|
||||
# Extract its nested tools and map them individually.
|
||||
for nested in tool.get("tools") or []:
|
||||
if "input_schema" in nested:
|
||||
# Already in Anthropic format.
|
||||
anthropic_tools.append(nested)
|
||||
elif "function" not in nested and "name" in nested:
|
||||
# Flat format: {type, name, description, parameters, ...}.
|
||||
# Normalize to OpenAI-wrapped format before mapping.
|
||||
wrapped = cast(
|
||||
ChatCompletionToolParam,
|
||||
{
|
||||
"type": nested.get("type", "function"),
|
||||
"function": {
|
||||
k: v for k, v in nested.items() if k != "type"
|
||||
},
|
||||
},
|
||||
)
|
||||
nested_tool, nested_mcp = self._map_tool_helper(wrapped)
|
||||
if nested_tool is not None:
|
||||
anthropic_tools.append(nested_tool)
|
||||
if nested_mcp is not None:
|
||||
mcp_servers.append(nested_mcp)
|
||||
elif "function" in nested:
|
||||
nested_tool, nested_mcp = self._map_tool_helper(
|
||||
cast(ChatCompletionToolParam, nested)
|
||||
)
|
||||
if nested_tool is not None:
|
||||
anthropic_tools.append(nested_tool)
|
||||
if nested_mcp is not None:
|
||||
mcp_servers.append(nested_mcp)
|
||||
elif "input_schema" in tool: # assume in anthropic format
|
||||
anthropic_tools.append(tool)
|
||||
else: # assume openai tool call
|
||||
new_tool, mcp_server_tool = self._map_tool_helper(tool)
|
||||
|
|
@ -1468,10 +1455,15 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
|
|||
_value = self._map_stop_sequences(value)
|
||||
if _value is not None:
|
||||
optional_params["stop_sequences"] = _value
|
||||
elif param == "temperature":
|
||||
optional_params["temperature"] = value
|
||||
elif param == "top_p":
|
||||
optional_params["top_p"] = value
|
||||
elif param == "temperature" or param == "top_p":
|
||||
AnthropicConfig._apply_sampling_param(
|
||||
optional_params=optional_params,
|
||||
model=model,
|
||||
param=param,
|
||||
value=value,
|
||||
drop_params=drop_params,
|
||||
output_key=param,
|
||||
)
|
||||
elif param == "response_format" and isinstance(value, dict):
|
||||
if any(
|
||||
substring in model
|
||||
|
|
@ -1620,6 +1612,15 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
|
|||
)
|
||||
return _tool
|
||||
|
||||
def should_strip_billing_metadata(self) -> bool:
|
||||
"""
|
||||
Whether to drop x-anthropic-billing-header system blocks before sending upstream.
|
||||
|
||||
The first-party Anthropic API uses these blocks for Claude Code attribution, so the
|
||||
base config keeps them. Providers that reject them (e.g. Bedrock) override this to True.
|
||||
"""
|
||||
return False
|
||||
|
||||
def translate_system_message(
|
||||
self, messages: List[AllMessageValues]
|
||||
) -> List[AnthropicSystemMessageContent]:
|
||||
|
|
@ -1627,7 +1628,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
|
|||
Translate system message to anthropic format.
|
||||
|
||||
Removes system message from the original list and returns a new list of anthropic system message content.
|
||||
Filters out system messages containing x-anthropic-billing-header metadata.
|
||||
When should_strip_billing_metadata() is True, x-anthropic-billing-header system blocks are dropped.
|
||||
"""
|
||||
system_prompt_indices = []
|
||||
anthropic_system_message_list: List[AnthropicSystemMessageContent] = []
|
||||
|
|
@ -1639,10 +1640,9 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
|
|||
# Skip empty text blocks - Anthropic API raises errors for empty text
|
||||
if not system_message_block["content"]:
|
||||
continue
|
||||
# Skip system messages containing x-anthropic-billing-header metadata
|
||||
if system_message_block["content"].startswith(
|
||||
"x-anthropic-billing-header:"
|
||||
):
|
||||
if self.should_strip_billing_metadata() and system_message_block[
|
||||
"content"
|
||||
].startswith("x-anthropic-billing-header:"):
|
||||
continue
|
||||
anthropic_system_message_content = AnthropicSystemMessageContent(
|
||||
type="text",
|
||||
|
|
@ -1661,9 +1661,9 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
|
|||
text_value = _content.get("text")
|
||||
if _content.get("type") == "text" and not text_value:
|
||||
continue
|
||||
# Skip system messages containing x-anthropic-billing-header metadata
|
||||
if (
|
||||
_content.get("type") == "text"
|
||||
self.should_strip_billing_metadata()
|
||||
and _content.get("type") == "text"
|
||||
and text_value
|
||||
and text_value.startswith("x-anthropic-billing-header:")
|
||||
):
|
||||
|
|
@ -1978,6 +1978,21 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
|
|||
|
||||
# Remove internal LiteLLM parameters that should not be sent to Anthropic API
|
||||
optional_params.pop("is_vertex_request", None)
|
||||
optional_params.pop("client_metadata", None)
|
||||
|
||||
# ``top_k`` is a provider-specific kwarg that bypasses
|
||||
# ``map_openai_params``; gate it here, the single boundary shared by
|
||||
# the direct Anthropic, Bedrock invoke, Vertex, and Azure paths.
|
||||
top_k = optional_params.pop("top_k", None)
|
||||
if top_k is not None:
|
||||
AnthropicConfig._apply_sampling_param(
|
||||
optional_params=optional_params,
|
||||
model=model,
|
||||
param="top_k",
|
||||
value=top_k,
|
||||
drop_params=litellm_params.get("drop_params") is True,
|
||||
output_key="top_k",
|
||||
)
|
||||
|
||||
data = {
|
||||
"model": model,
|
||||
|
|
|
|||
|
|
@ -272,19 +272,133 @@ class AnthropicModelInfo(BaseLLMModelInfo):
|
|||
)
|
||||
|
||||
@staticmethod
|
||||
def _is_adaptive_thinking_model(model: str) -> bool:
|
||||
"""Claude 4.6+ models use adaptive thinking with ``output_config.effort``."""
|
||||
def _supports_sampling_params(model: str) -> bool:
|
||||
"""Claude 4.7+ (Opus 4.7/4.8, Fable 5) removed sampling params: the API
|
||||
rejects ``top_p``, ``top_k``, and any ``temperature`` other than 1 with
|
||||
a 400 ("`temperature` is deprecated for this model").
|
||||
|
||||
Driven by the ``supports_sampling_params`` flag in the model map; the
|
||||
name check remains only as a fallback for provider-routed ids whose
|
||||
map entries predate the flag."""
|
||||
flag = AnthropicModelInfo._get_model_capability(
|
||||
model, "supports_sampling_params"
|
||||
)
|
||||
if flag is not None:
|
||||
return flag
|
||||
model_lower = model.lower()
|
||||
return not any(
|
||||
v in model_lower
|
||||
for v in (
|
||||
"fable",
|
||||
"opus-4-7",
|
||||
"opus_4_7",
|
||||
"opus-4.7",
|
||||
"opus_4.7",
|
||||
"opus-4-8",
|
||||
"opus_4_8",
|
||||
"opus-4.8",
|
||||
"opus_4.8",
|
||||
)
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _apply_sampling_param(
|
||||
optional_params: dict,
|
||||
model: str,
|
||||
param: str,
|
||||
value: Any,
|
||||
drop_params: bool,
|
||||
output_key: str,
|
||||
) -> None:
|
||||
"""Forward ``temperature``/``top_p``/``top_k`` to
|
||||
``optional_params[output_key]`` unless the model removed sampling
|
||||
params, in which case drop the param (with drop_params) or raise a
|
||||
clean client-side 400."""
|
||||
if AnthropicModelInfo._supports_sampling_params(model) or (
|
||||
param == "temperature" and value == 1
|
||||
):
|
||||
optional_params[output_key] = value
|
||||
elif not (litellm.drop_params or drop_params):
|
||||
supported_hint = (
|
||||
"Only temperature=1 is supported. " if param == "temperature" else ""
|
||||
)
|
||||
raise litellm.utils.UnsupportedParamsError(
|
||||
message=(
|
||||
f"{model} does not support {param}={value}. {supported_hint}"
|
||||
"To drop unsupported params, set `litellm.drop_params = True`."
|
||||
),
|
||||
status_code=400,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _model_map_lookup_candidates(model: str) -> List[str]:
|
||||
"""Model-map keys to try for ``model``, stripping bedrock/vertex
|
||||
prefixes so a provider-routed Claude still resolves to its entry."""
|
||||
candidates = [model]
|
||||
for prefix in (
|
||||
"bedrock/converse/",
|
||||
"bedrock/invoke/",
|
||||
"bedrock/",
|
||||
"vertex_ai/",
|
||||
):
|
||||
if model.startswith(prefix):
|
||||
candidates.append(model[len(prefix) :])
|
||||
try:
|
||||
from litellm.llms.bedrock.common_utils import BedrockModelInfo
|
||||
|
||||
base = BedrockModelInfo.get_base_model(model)
|
||||
if base:
|
||||
candidates.append(base)
|
||||
candidates.append(f"bedrock/{base}")
|
||||
except Exception:
|
||||
pass
|
||||
return candidates
|
||||
|
||||
@staticmethod
|
||||
def _get_model_capability(model: str, key: str) -> Optional[bool]:
|
||||
"""Read boolean capability ``key`` from the model map, or None when
|
||||
no entry declares it."""
|
||||
try:
|
||||
for cand in AnthropicModelInfo._model_map_lookup_candidates(model):
|
||||
value = litellm.model_cost.get(cand, {}).get(key)
|
||||
if isinstance(value, bool):
|
||||
return value
|
||||
except Exception:
|
||||
pass
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def _supports_model_capability(model: str, key: str) -> bool:
|
||||
"""Check a boolean capability ``key`` in the model map.
|
||||
|
||||
Strips bedrock/vertex prefixes so a provider-routed Claude still
|
||||
resolves to the Anthropic model-map entry.
|
||||
"""
|
||||
from litellm.utils import _supports_factory
|
||||
|
||||
try:
|
||||
if _supports_factory(
|
||||
model=model,
|
||||
custom_llm_provider=None,
|
||||
key="supports_adaptive_thinking",
|
||||
custom_llm_provider="anthropic",
|
||||
key=key,
|
||||
):
|
||||
return True
|
||||
except Exception:
|
||||
pass
|
||||
return AnthropicModelInfo._get_model_capability(model, key) is True
|
||||
|
||||
@staticmethod
|
||||
def _is_adaptive_thinking_model(model: str) -> bool:
|
||||
"""Claude 4.6+ models use adaptive thinking with ``output_config.effort``.
|
||||
|
||||
Driven by the ``supports_adaptive_thinking`` flag in the model map; the
|
||||
4.6/4.7 name checks remain only as a fallback for provider-routed ids
|
||||
whose map entries predate the flag.
|
||||
"""
|
||||
if AnthropicModelInfo._supports_model_capability(
|
||||
model, "supports_adaptive_thinking"
|
||||
):
|
||||
return True
|
||||
return AnthropicModelInfo._is_claude_4_6_model(
|
||||
model
|
||||
) or AnthropicModelInfo._is_claude_4_7_model(model)
|
||||
|
|
|
|||
|
|
@ -7,6 +7,7 @@ from typing import TYPE_CHECKING, Optional, Tuple
|
|||
|
||||
from litellm.litellm_core_utils.llm_cost_calc.utils import (
|
||||
_get_token_base_cost,
|
||||
_get_web_search_requests,
|
||||
_parse_prompt_tokens_details,
|
||||
calculate_cache_writing_cost,
|
||||
generic_cost_per_token,
|
||||
|
|
@ -110,11 +111,12 @@ def get_cost_for_anthropic_web_search(
|
|||
if model_info is None:
|
||||
return 0.0
|
||||
|
||||
if (
|
||||
usage is None
|
||||
or usage.server_tool_use is None
|
||||
or usage.server_tool_use.web_search_requests is None
|
||||
):
|
||||
if usage is None:
|
||||
return 0.0
|
||||
web_search_requests = _get_web_search_requests(
|
||||
getattr(usage, "server_tool_use", None)
|
||||
)
|
||||
if web_search_requests is None:
|
||||
return 0.0
|
||||
|
||||
## Get the cost per web search request
|
||||
|
|
@ -128,5 +130,5 @@ def get_cost_for_anthropic_web_search(
|
|||
return 0.0
|
||||
|
||||
## Calculate the total cost
|
||||
total_cost = cost_per_web_search_request * usage.server_tool_use.web_search_requests
|
||||
total_cost = cost_per_web_search_request * web_search_requests
|
||||
return total_cost
|
||||
|
|
|
|||
|
|
@ -1,5 +1,6 @@
|
|||
# What is this?
|
||||
## Translates OpenAI call to Anthropic `/v1/messages` format
|
||||
import copy
|
||||
import json
|
||||
import traceback
|
||||
from collections import deque
|
||||
|
|
@ -29,6 +30,98 @@ if TYPE_CHECKING:
|
|||
from litellm.types.utils import ModelResponseStream
|
||||
|
||||
|
||||
class _CombinedChunkSplitter:
|
||||
"""
|
||||
Splits a streaming chunk that carries BOTH response content and a
|
||||
``finish_reason`` into two chunks: a content-only chunk followed by a
|
||||
finish-only chunk.
|
||||
|
||||
``AnthropicStreamWrapper`` (via ``translate_streaming_openai_response_to_anthropic``)
|
||||
assumes content and ``finish_reason`` never arrive in the same chunk — true for
|
||||
real provider streams, but false for fake-streamed providers (e.g. Vertex AI
|
||||
Gemma ``:predict``) where ``MockResponseIterator`` collapses the entire response
|
||||
into a single chunk. Without this split the assumption causes all content to be
|
||||
silently dropped (only the ``message_delta`` stop event is emitted).
|
||||
|
||||
Supports both sync and async iteration, since ``AnthropicStreamWrapper`` exposes
|
||||
both ``__next__`` and ``__anext__``. An instance is single-mode: callers must
|
||||
iterate it either synchronously or asynchronously, never both — the two modes
|
||||
hold independent iterator references on the upstream stream and mixing them
|
||||
would advance them out of sync.
|
||||
"""
|
||||
|
||||
def __init__(self, completion_stream: Any):
|
||||
self._stream = completion_stream
|
||||
self._sync_iter: Optional[Iterator[Any]] = None
|
||||
self._async_iter: Optional[AsyncIterator[Any]] = None
|
||||
self._buffer: deque = deque()
|
||||
|
||||
@staticmethod
|
||||
def _is_combined(chunk: Any) -> bool:
|
||||
"""True if ``chunk`` carries response content AND a finish_reason."""
|
||||
choices = getattr(chunk, "choices", None)
|
||||
if not choices:
|
||||
return False
|
||||
choice = choices[0]
|
||||
if getattr(choice, "finish_reason", None) is None:
|
||||
return False
|
||||
delta = getattr(choice, "delta", None)
|
||||
if delta is None:
|
||||
return False
|
||||
return bool(
|
||||
getattr(delta, "content", None)
|
||||
or getattr(delta, "tool_calls", None)
|
||||
or getattr(delta, "reasoning_content", None)
|
||||
or getattr(delta, "thinking_blocks", None)
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _split(chunk: Any) -> List[Any]:
|
||||
"""Return ``[chunk]``, or ``[content_chunk, finish_chunk]`` if combined."""
|
||||
if not _CombinedChunkSplitter._is_combined(chunk):
|
||||
return [chunk]
|
||||
|
||||
# Content chunk: keep the delta payload, clear the finish_reason.
|
||||
content_chunk = copy.deepcopy(chunk)
|
||||
content_chunk.choices[0].finish_reason = None
|
||||
|
||||
# Finish chunk: keep finish_reason (and usage), clear the delta payload.
|
||||
finish_chunk = copy.deepcopy(chunk)
|
||||
finish_delta = finish_chunk.choices[0].delta
|
||||
finish_delta.content = None
|
||||
if hasattr(finish_delta, "tool_calls"):
|
||||
finish_delta.tool_calls = None
|
||||
if hasattr(finish_delta, "reasoning_content"):
|
||||
finish_delta.reasoning_content = None
|
||||
if hasattr(finish_delta, "thinking_blocks"):
|
||||
finish_delta.thinking_blocks = None
|
||||
return [content_chunk, finish_chunk]
|
||||
|
||||
def __iter__(self) -> "Iterator[Any]":
|
||||
return self
|
||||
|
||||
def __next__(self) -> Any:
|
||||
if self._buffer:
|
||||
return self._buffer.popleft()
|
||||
if self._sync_iter is None:
|
||||
self._sync_iter = iter(self._stream)
|
||||
chunk = next(self._sync_iter) # propagates StopIteration when exhausted
|
||||
self._buffer.extend(self._split(chunk))
|
||||
return self._buffer.popleft()
|
||||
|
||||
def __aiter__(self) -> "AsyncIterator[Any]":
|
||||
return self
|
||||
|
||||
async def __anext__(self) -> Any:
|
||||
if self._buffer:
|
||||
return self._buffer.popleft()
|
||||
if self._async_iter is None:
|
||||
self._async_iter = self._stream.__aiter__()
|
||||
chunk = await self._async_iter.__anext__() # propagates StopAsyncIteration
|
||||
self._buffer.extend(self._split(chunk))
|
||||
return self._buffer.popleft()
|
||||
|
||||
|
||||
class AnthropicStreamWrapper(AdapterCompletionStreamWrapper):
|
||||
"""
|
||||
- first chunk return 'message_start'
|
||||
|
|
@ -62,7 +155,10 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper):
|
|||
compaction_block: Optional[CompactionBlock] = None,
|
||||
iterations_usage: Optional[List[UsageIteration]] = None,
|
||||
):
|
||||
super().__init__(completion_stream)
|
||||
# Wrap the upstream stream so chunks that carry both content and a
|
||||
# finish_reason (fake-streamed providers) are split into two — see
|
||||
# _CombinedChunkSplitter.
|
||||
super().__init__(_CombinedChunkSplitter(completion_stream))
|
||||
self.model = model
|
||||
# Mapping of truncated tool names to original names (for OpenAI's 64-char limit)
|
||||
self.tool_name_mapping = tool_name_mapping or {}
|
||||
|
|
@ -373,12 +469,16 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper):
|
|||
|
||||
if should_start_new_block and not self.sent_content_block_finish:
|
||||
# Queue the sequence: content_block_stop -> content_block_start
|
||||
# For text blocks the trigger chunk is not emitted as a separate
|
||||
# delta because content_block_start carries the information.
|
||||
# For tool_use blocks we must also emit the trigger chunk's delta
|
||||
# when it carries input_json_delta data, because some providers
|
||||
# (e.g. xAI, Gemini) include tool arguments in the same streaming
|
||||
# chunk as the function name/id.
|
||||
# -> (optionally) the trigger chunk's delta.
|
||||
#
|
||||
# The synthesized content_block_start always carries an
|
||||
# empty body, so the chunk that *triggered* the transition
|
||||
# also carries the new block's first delta. It must be
|
||||
# re-emitted or the first token of the new block is lost.
|
||||
# This applies to text_delta and thinking_delta (the first
|
||||
# non-empty text/thinking token) as well as input_json_delta
|
||||
# (providers like xAI/Gemini bundle tool arguments with the
|
||||
# function name/id in a single chunk).
|
||||
|
||||
# 1. Stop current content block
|
||||
self.chunk_queue.append(
|
||||
|
|
@ -397,14 +497,9 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper):
|
|||
}
|
||||
)
|
||||
|
||||
# 3. If the trigger chunk carries tool argument data, queue it
|
||||
# so the input_json_delta is not silently dropped.
|
||||
if (
|
||||
processed_chunk.get("type") == "content_block_delta"
|
||||
and isinstance(processed_chunk.get("delta"), dict)
|
||||
and processed_chunk["delta"].get("type") == "input_json_delta"
|
||||
and processed_chunk["delta"].get("partial_json")
|
||||
):
|
||||
# 3. If the trigger chunk carries delta content, queue it
|
||||
# so the first delta of the new block is not silently dropped.
|
||||
if self._trigger_delta_has_content(processed_chunk):
|
||||
self.chunk_queue.append(processed_chunk)
|
||||
|
||||
self.sent_content_block_finish = False
|
||||
|
|
@ -615,12 +710,16 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper):
|
|||
if not self.queued_usage_chunk:
|
||||
if should_start_new_block and not self.sent_content_block_finish:
|
||||
# Queue the sequence: content_block_stop -> content_block_start
|
||||
# For text blocks the trigger chunk is not emitted as a separate
|
||||
# delta because content_block_start carries the information.
|
||||
# For tool_use blocks we must also emit the trigger chunk's delta
|
||||
# when it carries input_json_delta data, because some providers
|
||||
# (e.g. xAI, Gemini) include tool arguments in the same streaming
|
||||
# chunk as the function name/id.
|
||||
# -> (optionally) the trigger chunk's delta.
|
||||
#
|
||||
# The synthesized content_block_start always carries an
|
||||
# empty body, so the chunk that *triggered* the transition
|
||||
# also carries the new block's first delta. It must be
|
||||
# re-emitted or the first token of the new block is lost.
|
||||
# This applies to text_delta and thinking_delta (the
|
||||
# first non-empty text/thinking token) as well as
|
||||
# input_json_delta (providers like xAI/Gemini bundle tool
|
||||
# arguments with the function name/id in a single chunk).
|
||||
|
||||
# 1. Stop current content block
|
||||
self.chunk_queue.append(
|
||||
|
|
@ -637,15 +736,9 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper):
|
|||
}
|
||||
)
|
||||
|
||||
# 3. If the trigger chunk carries tool argument data, queue it
|
||||
# so the input_json_delta is not silently dropped.
|
||||
if (
|
||||
processed_chunk.get("type") == "content_block_delta"
|
||||
and isinstance(processed_chunk.get("delta"), dict)
|
||||
and processed_chunk["delta"].get("type")
|
||||
== "input_json_delta"
|
||||
and processed_chunk["delta"].get("partial_json")
|
||||
):
|
||||
# 3. If the trigger chunk carries delta content, queue it
|
||||
# so the first delta of the new block is not silently dropped.
|
||||
if self._trigger_delta_has_content(processed_chunk):
|
||||
self.chunk_queue.append(processed_chunk)
|
||||
|
||||
# Reset state for new block
|
||||
|
|
@ -802,6 +895,38 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper):
|
|||
def _increment_content_block_index(self):
|
||||
self.current_content_block_index += 1
|
||||
|
||||
@staticmethod
|
||||
def _trigger_delta_has_content(processed_chunk: Dict[str, Any]) -> bool:
|
||||
"""Return True if a translated trigger chunk carries a non-empty
|
||||
``content_block_delta`` payload that must be re-emitted after a
|
||||
block transition.
|
||||
|
||||
When an upstream chunk both *triggers* a new content block (its type
|
||||
differs from the active block) and *carries* delta content, that
|
||||
content belongs to the new block. The synthesized
|
||||
``content_block_start`` only ever carries an empty body — see
|
||||
``_translate_streaming_openai_chunk_to_anthropic_content_block``,
|
||||
which returns an empty ``TextBlock``/``ToolUseBlock``/thinking block —
|
||||
so the trigger chunk's delta must be re-queued or the first token of
|
||||
the new block (the first non-empty text/thinking delta, or bundled
|
||||
tool arguments) is silently dropped.
|
||||
"""
|
||||
if processed_chunk.get("type") != "content_block_delta":
|
||||
return False
|
||||
delta = processed_chunk.get("delta")
|
||||
if not isinstance(delta, dict):
|
||||
return False
|
||||
delta_type = delta.get("type")
|
||||
if delta_type == "text_delta":
|
||||
return bool(delta.get("text"))
|
||||
if delta_type == "input_json_delta":
|
||||
return bool(delta.get("partial_json"))
|
||||
if delta_type == "thinking_delta":
|
||||
return bool(delta.get("thinking"))
|
||||
if delta_type == "signature_delta":
|
||||
return bool(delta.get("signature"))
|
||||
return False
|
||||
|
||||
def _should_start_new_content_block(self, chunk: "ModelResponseStream") -> bool:
|
||||
"""
|
||||
Determine if we should start a new content block based on the processed chunk.
|
||||
|
|
|
|||
|
|
@ -84,6 +84,15 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig):
|
|||
if isinstance(content, list):
|
||||
_process_content_list(content)
|
||||
|
||||
def should_strip_billing_metadata(self) -> bool:
|
||||
"""
|
||||
Whether to drop x-anthropic-billing-header system blocks before sending upstream.
|
||||
|
||||
The first-party Anthropic API uses these blocks for Claude Code attribution, so the
|
||||
base config keeps them. Providers that reject them override this to True.
|
||||
"""
|
||||
return False
|
||||
|
||||
@staticmethod
|
||||
def _filter_billing_headers_from_system(system_param):
|
||||
"""
|
||||
|
|
@ -286,14 +295,12 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig):
|
|||
optional_params=anthropic_messages_optional_request_params,
|
||||
)
|
||||
|
||||
# Filter out x-anthropic-billing-header from system messages
|
||||
system_param = anthropic_messages_optional_request_params.get("system")
|
||||
if system_param is not None:
|
||||
if self.should_strip_billing_metadata() and system_param is not None:
|
||||
filtered_system = self._filter_billing_headers_from_system(system_param)
|
||||
if filtered_system is not None and len(filtered_system) > 0:
|
||||
anthropic_messages_optional_request_params["system"] = filtered_system
|
||||
else:
|
||||
# Remove system parameter if all content was filtered out
|
||||
anthropic_messages_optional_request_params.pop("system", None)
|
||||
|
||||
# Transform context_management from OpenAI format to Anthropic format if needed
|
||||
|
|
|
|||
|
|
@ -102,9 +102,9 @@ def _build_responses_kwargs(
|
|||
from litellm.types.utils import CallTypes
|
||||
|
||||
if isinstance(value, LiteLLMLoggingObject):
|
||||
# Reclassify as acompletion so the success handler doesn't try to
|
||||
# validate the Responses API event as an AnthropicResponse.
|
||||
# (Mirrors the pattern used in LiteLLMMessagesToCompletionTransformationHandler.)
|
||||
# Keep call_type as anthropic_messages so spend_logs are billed
|
||||
# against /v1/messages; the success handler translates the
|
||||
# Responses API result back to a ModelResponse for the row.
|
||||
setattr(value, "call_type", CallTypes.anthropic_messages.value)
|
||||
responses_kwargs[key] = value
|
||||
elif key not in excluded and key not in responses_kwargs and value is not None:
|
||||
|
|
|
|||
|
|
@ -155,10 +155,24 @@ class AnthropicResponsesStreamWrapper:
|
|||
event.get("delta", "") if isinstance(event, dict) else ""
|
||||
)
|
||||
block_idx = (
|
||||
self._item_id_to_block_index.get(item_id, self._current_block_index)
|
||||
self._item_id_to_block_index.get(item_id, -1)
|
||||
if item_id
|
||||
else self._current_block_index
|
||||
)
|
||||
if block_idx < 0:
|
||||
# Some providers (e.g. LMStudio) skip response.output_item.added,
|
||||
# so no text block is open yet; synthesize content_block_start
|
||||
# instead of emitting a delta with index -1
|
||||
block_idx = self._next_block_index()
|
||||
if item_id:
|
||||
self._item_id_to_block_index[item_id] = block_idx
|
||||
self._chunk_queue.append(
|
||||
{
|
||||
"type": "content_block_start",
|
||||
"index": block_idx,
|
||||
"content_block": {"type": "text", "text": ""},
|
||||
}
|
||||
)
|
||||
self._chunk_queue.append(
|
||||
{
|
||||
"type": "content_block_delta",
|
||||
|
|
|
|||
Some files were not shown because too many files have changed in this diff Show more
Loading…
Add table
Reference in a new issue