mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-07 08:26:10 +00:00
Merge remote-tracking branch 'origin/litellm_internal_staging' into litellm_internal_copy_38013
This commit is contained in:
commit
c813cc71ae
719 changed files with 50219 additions and 8333 deletions
23
.github/actions/cache-cargo-build/action.yml
vendored
23
.github/actions/cache-cargo-build/action.yml
vendored
|
|
@ -4,17 +4,16 @@ description: >-
|
|||
so only the first job on a given Cargo.lock compiles the bridge from scratch.
|
||||
|
||||
litellm builds through maturin, which compiles litellm-rust/crates/python-bridge
|
||||
in release mode before it can produce a wheel. `uv sync` therefore pays a full
|
||||
build in every job that installs the workspace: measured at 2m40s per unit shard
|
||||
on 2026-08-21, more than the whole unit tier spends running tests. Nothing caught
|
||||
it, because the uv cache holds wheels uv downloads rather than wheels it builds,
|
||||
and a path dependency whose source moves every commit could never hit that cache
|
||||
anyway. Cargo rebuilds only what changed when its target directory survives, so a
|
||||
warm job pays for the bridge crate alone.
|
||||
in the dev profile for editable installs. `uv sync` therefore pays a full build
|
||||
in every job that installs the workspace. Nothing caught it, because the uv cache
|
||||
holds wheels uv downloads rather than wheels it builds, and a path dependency
|
||||
whose source moves every commit could never hit that cache anyway. Cargo rebuilds
|
||||
only what changed when its target directory survives, so a warm job pays for the
|
||||
bridge crate alone.
|
||||
|
||||
The key namespace is separate from test-rust.yml's. Both cache the same directory,
|
||||
but that workflow fills it with debug and clippy artifacts, which a release build
|
||||
cannot reuse, and a shared key would let whichever ran first deny the other a save.
|
||||
The key namespace is separate from test-rust.yml's check and release caches. They
|
||||
cache the same directory for different workloads, and a shared key would let
|
||||
whichever ran first deny the others a save.
|
||||
|
||||
runs:
|
||||
using: composite
|
||||
|
|
@ -26,6 +25,6 @@ runs:
|
|||
~/.cargo/registry
|
||||
~/.cargo/git
|
||||
litellm-rust/target
|
||||
key: ${{ runner.os }}-cargo-release-${{ hashFiles('litellm-rust/Cargo.lock') }}
|
||||
key: ${{ runner.os }}-maturin-dev-${{ hashFiles('litellm-rust/Cargo.lock') }}
|
||||
restore-keys: |
|
||||
${{ runner.os }}-cargo-release-
|
||||
${{ runner.os }}-maturin-dev-
|
||||
|
|
|
|||
230
.github/scripts/close_duplicate_issues.py
vendored
230
.github/scripts/close_duplicate_issues.py
vendored
|
|
@ -1,230 +0,0 @@
|
|||
#!/usr/bin/env python3
|
||||
"""
|
||||
Detect and close duplicate GitHub issues using title similarity.
|
||||
|
||||
Modes:
|
||||
--scan Compare all open issues against each other (batch)
|
||||
--issue-number N Check a single issue against older open issues
|
||||
|
||||
Requires the `gh` CLI to be authenticated.
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import difflib
|
||||
import json
|
||||
import re
|
||||
import subprocess
|
||||
import sys
|
||||
|
||||
|
||||
def normalize_title(title: str) -> str:
|
||||
"""Strip common prefixes, lowercase, and collapse whitespace."""
|
||||
title = re.sub(
|
||||
r"^\[?(bug|feature request|enhancement|question|docs)[:\]]?\s*",
|
||||
"",
|
||||
title,
|
||||
flags=re.IGNORECASE,
|
||||
)
|
||||
return " ".join(title.lower().split())
|
||||
|
||||
|
||||
def gh(*args: str) -> str:
|
||||
"""Run a gh CLI command and return stdout."""
|
||||
result = subprocess.run(
|
||||
["gh", *args],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
check=True,
|
||||
)
|
||||
return result.stdout
|
||||
|
||||
|
||||
def fetch_open_issues(repo: str | None) -> list[dict]:
|
||||
"""Fetch all open issues (excluding PRs) via gh api --paginate."""
|
||||
if repo:
|
||||
endpoint = (
|
||||
f"repos/{repo}/issues?state=open&per_page=100&sort=created&direction=asc"
|
||||
)
|
||||
else:
|
||||
endpoint = "repos/{owner}/{repo}/issues?state=open&per_page=100&sort=created&direction=asc"
|
||||
cmd = ["api", "--paginate", endpoint]
|
||||
|
||||
raw = gh(*cmd)
|
||||
# gh --paginate concatenates JSON arrays, so we may get multiple arrays
|
||||
issues = []
|
||||
for line in raw.strip().splitlines():
|
||||
line = line.strip()
|
||||
if not line:
|
||||
continue
|
||||
parsed = json.loads(line)
|
||||
if isinstance(parsed, list):
|
||||
issues.extend(parsed)
|
||||
else:
|
||||
issues.append(parsed)
|
||||
|
||||
# Filter out pull requests (they also appear in the issues endpoint)
|
||||
return [i for i in issues if "pull_request" not in i]
|
||||
|
||||
|
||||
def close_as_duplicate(
|
||||
issue_number: int, duplicate_of: int, repo: str | None, dry_run: bool
|
||||
) -> None:
|
||||
"""Close an issue as duplicate of another, adding a comment and label."""
|
||||
repo_args = ["--repo", repo] if repo else []
|
||||
|
||||
if dry_run:
|
||||
print(
|
||||
f" [DRY RUN] Would close #{issue_number} as duplicate of #{duplicate_of}"
|
||||
)
|
||||
return
|
||||
|
||||
# Add comment
|
||||
comment_body = (
|
||||
f"Closing as duplicate of #{duplicate_of}.\n\n"
|
||||
"If you believe this is not a duplicate, please reopen and add context "
|
||||
"explaining how this differs."
|
||||
)
|
||||
gh("issue", "comment", str(issue_number), "--body", comment_body, *repo_args)
|
||||
|
||||
# Add label
|
||||
gh("issue", "edit", str(issue_number), "--add-label", "duplicate", *repo_args)
|
||||
|
||||
# Close with not_planned reason
|
||||
gh(
|
||||
"api",
|
||||
f"repos/{repo or '{owner}/{repo}'}/issues/{issue_number}",
|
||||
"-X",
|
||||
"PATCH",
|
||||
"-f",
|
||||
"state=closed",
|
||||
"-f",
|
||||
"state_reason=not_planned",
|
||||
)
|
||||
|
||||
print(f" Closed #{issue_number} as duplicate of #{duplicate_of}")
|
||||
|
||||
|
||||
def find_duplicate(
|
||||
issue: dict, candidates: list[dict], threshold: float
|
||||
) -> dict | None:
|
||||
"""Return the first candidate whose normalized title is above threshold."""
|
||||
norm = normalize_title(issue["title"])
|
||||
for candidate in candidates:
|
||||
if candidate["number"] == issue["number"]:
|
||||
continue
|
||||
cand_norm = normalize_title(candidate["title"])
|
||||
ratio = difflib.SequenceMatcher(None, norm, cand_norm).ratio()
|
||||
if ratio >= threshold:
|
||||
return candidate
|
||||
return None
|
||||
|
||||
|
||||
def scan_all(
|
||||
issues: list[dict], threshold: float, repo: str | None, dry_run: bool
|
||||
) -> int:
|
||||
"""Compare every issue against all older issues. Returns count of duplicates found."""
|
||||
# Sort oldest first
|
||||
issues.sort(key=lambda i: i["number"])
|
||||
closed_count = 0
|
||||
|
||||
for idx, issue in enumerate(issues):
|
||||
older = issues[:idx]
|
||||
if not older:
|
||||
continue
|
||||
dup = find_duplicate(issue, older, threshold)
|
||||
if dup:
|
||||
ratio = difflib.SequenceMatcher(
|
||||
None,
|
||||
normalize_title(issue["title"]),
|
||||
normalize_title(dup["title"]),
|
||||
).ratio()
|
||||
print(
|
||||
f"#{issue['number']}: \"{issue['title']}\"\n"
|
||||
f" -> duplicate of #{dup['number']}: \"{dup['title']}\" "
|
||||
f"({ratio:.0%} similar)"
|
||||
)
|
||||
close_as_duplicate(issue["number"], dup["number"], repo, dry_run)
|
||||
closed_count += 1
|
||||
|
||||
return closed_count
|
||||
|
||||
|
||||
def check_single(
|
||||
issue_number: int,
|
||||
issues: list[dict],
|
||||
threshold: float,
|
||||
repo: str | None,
|
||||
dry_run: bool,
|
||||
) -> bool:
|
||||
"""Check a single issue against all older open issues. Returns True if duplicate found."""
|
||||
target = None
|
||||
for i in issues:
|
||||
if i["number"] == issue_number:
|
||||
target = i
|
||||
break
|
||||
|
||||
if target is None:
|
||||
print(f"Issue #{issue_number} not found among open issues.")
|
||||
return False
|
||||
|
||||
older = [i for i in issues if i["number"] < issue_number]
|
||||
dup = find_duplicate(target, older, threshold)
|
||||
if dup:
|
||||
ratio = difflib.SequenceMatcher(
|
||||
None,
|
||||
normalize_title(target["title"]),
|
||||
normalize_title(dup["title"]),
|
||||
).ratio()
|
||||
print(
|
||||
f"#{target['number']}: \"{target['title']}\"\n"
|
||||
f" -> duplicate of #{dup['number']}: \"{dup['title']}\" "
|
||||
f"({ratio:.0%} similar)"
|
||||
)
|
||||
close_as_duplicate(issue_number, dup["number"], repo, dry_run)
|
||||
return True
|
||||
|
||||
print(f"#{issue_number}: no duplicate found above threshold {threshold}")
|
||||
return False
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Detect and close duplicate GitHub issues"
|
||||
)
|
||||
mode = parser.add_mutually_exclusive_group(required=True)
|
||||
mode.add_argument("--scan", action="store_true", help="Scan all open issues")
|
||||
mode.add_argument("--issue-number", type=int, help="Check a single issue number")
|
||||
parser.add_argument(
|
||||
"--threshold", type=float, default=0.85, help="Similarity threshold (0-1)"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--close",
|
||||
action="store_true",
|
||||
help="Actually close duplicates (default is dry-run)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--repo", type=str, help="Repository (owner/repo). Auto-detected if omitted."
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
dry_run = not args.close
|
||||
|
||||
if dry_run:
|
||||
print("=== DRY RUN MODE (pass --close to actually close issues) ===\n")
|
||||
|
||||
print("Fetching open issues...")
|
||||
issues = fetch_open_issues(args.repo)
|
||||
print(f"Found {len(issues)} open issues.\n")
|
||||
|
||||
if args.scan:
|
||||
count = scan_all(issues, args.threshold, args.repo, dry_run)
|
||||
print(f"\nTotal duplicates {'found' if dry_run else 'closed'}: {count}")
|
||||
else:
|
||||
found = check_single(
|
||||
args.issue_number, issues, args.threshold, args.repo, dry_run
|
||||
)
|
||||
sys.exit(0 if found else 0) # Always exit 0; finding no dup is not an error
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
69
.github/workflows/auto-close-duplicates.yml
vendored
Normal file
69
.github/workflows/auto-close-duplicates.yml
vendored
Normal file
|
|
@ -0,0 +1,69 @@
|
|||
name: Auto-close duplicate issues
|
||||
|
||||
on:
|
||||
schedule:
|
||||
- cron: "0 9 * * *"
|
||||
workflow_dispatch:
|
||||
inputs:
|
||||
dry_run:
|
||||
description: Log which issues would close without closing anything
|
||||
type: boolean
|
||||
default: true
|
||||
grace_period_days:
|
||||
description: Days a duplicate notice must go unanswered before the close
|
||||
type: number
|
||||
default: 3
|
||||
pull_request:
|
||||
paths:
|
||||
- .github/workflows/auto-close-duplicates.yml
|
||||
- scripts/auto-close-duplicates.ts
|
||||
- scripts/auto-close-duplicates.test.ts
|
||||
|
||||
permissions: {}
|
||||
|
||||
jobs:
|
||||
test:
|
||||
if: github.event_name == 'pull_request'
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 5
|
||||
permissions:
|
||||
contents: read
|
||||
steps:
|
||||
- name: Checkout repository
|
||||
uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
|
||||
with:
|
||||
persist-credentials: false
|
||||
|
||||
- name: Setup Bun
|
||||
uses: oven-sh/setup-bun@0c5077e51419868618aeaa5fe8019c62421857d6 # v2.2.0
|
||||
with:
|
||||
bun-version: "1.4.0"
|
||||
|
||||
- name: Test the sweep
|
||||
run: bun test scripts/auto-close-duplicates.test.ts
|
||||
|
||||
sweep:
|
||||
if: github.event_name != 'pull_request' && github.repository == 'BerriAI/litellm'
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 10
|
||||
permissions:
|
||||
contents: read
|
||||
issues: write
|
||||
steps:
|
||||
- name: Checkout repository
|
||||
uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
|
||||
with:
|
||||
persist-credentials: false
|
||||
|
||||
- name: Setup Bun
|
||||
uses: oven-sh/setup-bun@0c5077e51419868618aeaa5fe8019c62421857d6 # v2.2.0
|
||||
with:
|
||||
# Exact version, never latest: the next step holds an issues: write token
|
||||
bun-version: "1.4.0"
|
||||
|
||||
- name: Close unanswered duplicates, reopen ones the reporter answered
|
||||
run: bun run scripts/auto-close-duplicates.ts
|
||||
env:
|
||||
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
|
||||
DRY_RUN: ${{ inputs.dry_run == true }}
|
||||
GRACE_PERIOD_DAYS: ${{ inputs.grace_period_days }}
|
||||
40
.github/workflows/check_duplicate_issues.yml
vendored
40
.github/workflows/check_duplicate_issues.yml
vendored
|
|
@ -1,12 +1,19 @@
|
|||
name: Check Duplicate Issues
|
||||
|
||||
# Flagging only. "Auto-close duplicate issues" closes a flagged issue 3 days later,
|
||||
# and only when its title is identical to an older open issue and nobody replied.
|
||||
# The HTML marker below is the handshake between the two, so keep it in the template.
|
||||
|
||||
on:
|
||||
issues:
|
||||
types: [opened, edited]
|
||||
|
||||
permissions: {}
|
||||
|
||||
jobs:
|
||||
check-duplicate:
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 5
|
||||
permissions:
|
||||
issues: write
|
||||
contents: read
|
||||
|
|
@ -19,35 +26,12 @@ jobs:
|
|||
threshold: 0.6
|
||||
reaction: eyes
|
||||
comment: |
|
||||
**⚠️ Potential duplicate detected**
|
||||
<!-- litellm:potential-duplicate candidates={{#issues}}{{number}},{{/issues}} -->
|
||||
**Potential duplicate detected**
|
||||
|
||||
This issue appears similar to existing issue(s):
|
||||
This looks similar to:
|
||||
{{#issues}}
|
||||
- [#{{number}}]({{html_url}}) - {{title}} ({{accuracy}}% similar)
|
||||
- #{{number}} - {{title}}
|
||||
{{/issues}}
|
||||
|
||||
Please review the linked issue(s) to see if they address your concern. If this is not a duplicate, please provide additional context to help us understand the difference.
|
||||
|
||||
- name: Checkout close script
|
||||
if: github.event.action == 'opened'
|
||||
uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
|
||||
with:
|
||||
sparse-checkout: .github/scripts
|
||||
persist-credentials: false
|
||||
|
||||
- name: Set up Python
|
||||
if: github.event.action == 'opened'
|
||||
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
|
||||
with:
|
||||
python-version: "3.12"
|
||||
|
||||
- name: Auto-close if high-confidence duplicate
|
||||
if: github.event.action == 'opened'
|
||||
env:
|
||||
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
|
||||
run: |
|
||||
python3 .github/scripts/close_duplicate_issues.py \
|
||||
--issue-number ${{ github.event.issue.number }} \
|
||||
--repo ${{ github.repository }} \
|
||||
--threshold 0.85 \
|
||||
--close
|
||||
If this is a duplicate, add a thumbs-up reaction to the existing issue and follow along there. When the title is identical to an older open issue, this issue closes automatically in 3 days unless someone responds. If it is not a duplicate, comment here or add a thumbs-down reaction to this comment and it stays open.
|
||||
|
|
|
|||
6
.github/workflows/image-scan.yml
vendored
6
.github/workflows/image-scan.yml
vendored
|
|
@ -80,7 +80,7 @@ jobs:
|
|||
LITELLM_IMAGE: litellm-image-scan:${{ github.sha }}
|
||||
run: |
|
||||
python -m pip install "pytest==9.0.3"
|
||||
python -m pytest tests/proxy_migration_tests/test_offline_image_migration.py -v
|
||||
python -m pytest tests/proxy_migration_tests/test_offline_image_migration.py tests/proxy_migration_tests/test_image_bedrock_realtime_extra.py -v
|
||||
|
||||
# Scans the whole shipped artifact: OS/apk plus every language package
|
||||
# baked into the image, including ones no lockfile declares (e.g. prisma's
|
||||
|
|
@ -124,7 +124,7 @@ jobs:
|
|||
LITELLM_IMAGE: litellm-runtime-scan:${{ github.sha }}
|
||||
run: |
|
||||
python -m pip install "pytest==9.0.3"
|
||||
python -m pytest tests/proxy_migration_tests/test_offline_image_migration.py -v
|
||||
python -m pytest tests/proxy_migration_tests/test_offline_image_migration.py tests/proxy_migration_tests/test_image_bedrock_realtime_extra.py -v
|
||||
|
||||
migrations-image:
|
||||
name: migrations-image
|
||||
|
|
@ -185,7 +185,7 @@ jobs:
|
|||
LITELLM_COMPONENT_PORT: "4000"
|
||||
run: |
|
||||
python -m pip install "pytest==9.0.3"
|
||||
python -m pytest tests/proxy_migration_tests/test_component_image_serves_offline.py -v
|
||||
python -m pytest tests/proxy_migration_tests/test_component_image_serves_offline.py tests/proxy_migration_tests/test_image_bedrock_realtime_extra.py -v
|
||||
|
||||
ui-image:
|
||||
name: ui-image
|
||||
|
|
|
|||
77
.github/workflows/test-redis-compat.yml
vendored
Normal file
77
.github/workflows/test-redis-compat.yml
vendored
Normal file
|
|
@ -0,0 +1,77 @@
|
|||
name: "Unit Tests: Redis Client Version Compatibility"
|
||||
|
||||
on:
|
||||
pull_request:
|
||||
branches:
|
||||
- main
|
||||
- litellm_internal_staging
|
||||
- litellm_oss_staging
|
||||
- "litellm_**"
|
||||
paths:
|
||||
- "litellm/_redis.py"
|
||||
- "litellm/_redis_credential_provider.py"
|
||||
- "tests/test_litellm/test_redis.py"
|
||||
- "tests/test_litellm/caching/test_redis_connection_pool.py"
|
||||
- ".github/workflows/test-redis-compat.yml"
|
||||
- "pyproject.toml"
|
||||
- "uv.lock"
|
||||
|
||||
permissions:
|
||||
contents: read
|
||||
|
||||
concurrency:
|
||||
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
|
||||
cancel-in-progress: true
|
||||
|
||||
jobs:
|
||||
redis-compat:
|
||||
name: "redis-py ${{ matrix.redis-version }}"
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 15
|
||||
|
||||
strategy:
|
||||
fail-fast: false
|
||||
matrix:
|
||||
# 5.3.1 is the version pinned in uv.lock (redisvl caps it below 6); the
|
||||
# newer legs prove the inspect.signature introspection in litellm/_redis.py
|
||||
# keeps extracting kwargs on the redis-py releases people actually run now.
|
||||
# Only the exact release 6.0.0 is skipped: rq (pulled by the proxy extra)
|
||||
# specifies `redis != 6`, which excludes 6.0.0 alone, so 6.4.0 stands in
|
||||
# for the 6.x line.
|
||||
redis-version: ["5.3.1", "6.4.0", "7.4.1", "8.0.1"]
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
|
||||
with:
|
||||
persist-credentials: false
|
||||
|
||||
- name: Set up Python
|
||||
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
|
||||
with:
|
||||
python-version: "3.12"
|
||||
|
||||
- name: Set up uv
|
||||
uses: ./.github/actions/setup-uv-with-retries
|
||||
with:
|
||||
version: "0.10.9"
|
||||
|
||||
- name: Install dependencies
|
||||
run: |
|
||||
.github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router
|
||||
|
||||
- name: Pin redis-py to the matrix version
|
||||
env:
|
||||
REDIS_VERSION: ${{ matrix.redis-version }}
|
||||
run: |
|
||||
uv pip install "redis==${REDIS_VERSION:?}"
|
||||
uv run --no-sync python -c "import redis; assert redis.__version__ == '${REDIS_VERSION:?}', redis.__version__; print('redis-py', redis.__version__)"
|
||||
|
||||
- name: Run redis unit tests
|
||||
run: |
|
||||
uv run --no-sync pytest \
|
||||
tests/test_litellm/test_redis.py \
|
||||
tests/test_litellm/caching/test_redis_connection_pool.py \
|
||||
--tb=short -vv \
|
||||
--reruns 2 \
|
||||
--reruns-delay 1 \
|
||||
--durations=20
|
||||
1
.github/workflows/test-unit.yml
vendored
1
.github/workflows/test-unit.yml
vendored
|
|
@ -103,6 +103,7 @@ jobs:
|
|||
tests/test_litellm/completion_extras
|
||||
tests/test_litellm/compression
|
||||
tests/test_litellm/containers
|
||||
tests/test_litellm/endpoints
|
||||
tests/test_litellm/experimental_mcp_client
|
||||
tests/test_litellm/models
|
||||
tests/test_litellm/repositories
|
||||
|
|
|
|||
|
|
@ -23,6 +23,8 @@ When adding new features, add meaningful tests. Don't add tests that don't check
|
|||
|
||||
Same thing for bug fixes. The tests should make it so that this specific bug can never happen again without failing tests (i.e., regression)
|
||||
|
||||
Never test structure of code only function of it
|
||||
|
||||
`tests/test_litellm/` mirrors `litellm/` in a parallel path (see `tests/test_litellm/readme.md`). Name tests `test_<filename>.py`, but always match the existing test file in the directory you touch — many provider dirs use longer descriptive names (e.g. `test_anthropic_chat_transformation.py`) to avoid ambiguity across sibling folders. For bug fixes, extend the existing mapped test file rather than creating a new one. Only create a new test file for a new feature (provider, endpoint, or transformation module) that has no mapped test yet, following that directory's naming convention (or `test_<filename>.py` if you're the first test there). One focused regression test beats many shallow ones
|
||||
|
||||
End-to-end tests belong in `tests/e2e/` and must follow the harness conventions documented in that directory's `CLAUDE.md`
|
||||
|
|
|
|||
23
Dockerfile
23
Dockerfile
|
|
@ -1,10 +1,10 @@
|
|||
# syntax=docker/dockerfile:1.7
|
||||
|
||||
# Base image for building
|
||||
ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:a31344ab2cb8618db84f535eec56f76f6178b142cb92cb2e48676cc2dcebea72
|
||||
ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:e624c5d5e42382ce7165ddafcbbf8e6769a24cbd02ea6114b880b05ae5ba2a8d
|
||||
|
||||
# Runtime image
|
||||
ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:a31344ab2cb8618db84f535eec56f76f6178b142cb92cb2e48676cc2dcebea72
|
||||
ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:e624c5d5e42382ce7165ddafcbbf8e6769a24cbd02ea6114b880b05ae5ba2a8d
|
||||
ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a
|
||||
# Pinned by digest like the other base images; bump explicitly on Node upgrades.
|
||||
ARG UI_BUILD_IMAGE=node:24.19-alpine3.24@sha256:d32cdf619f63fe0471182d08996dd516c6275bb5fd31ae06e55a570bd9e1ad43
|
||||
|
|
@ -40,8 +40,8 @@ COPY --from=uvbin /uvx /usr/local/bin/uvx
|
|||
RUN apk add --no-cache \
|
||||
bash \
|
||||
gcc \
|
||||
python3 \
|
||||
python3-dev \
|
||||
python-3.13 \
|
||||
python-3.13-dev \
|
||||
rust \
|
||||
openssl \
|
||||
openssl-dev \
|
||||
|
|
@ -51,6 +51,7 @@ RUN apk add --no-cache \
|
|||
|
||||
ENV UV_PROJECT_ENVIRONMENT=/app/.venv \
|
||||
UV_LINK_MODE=copy \
|
||||
UV_PYTHON_DOWNLOADS=0 \
|
||||
PATH="/app/.venv/bin:${PATH}"
|
||||
|
||||
# Copy dependency metadata first for layer caching
|
||||
|
|
@ -65,7 +66,8 @@ RUN uv sync --frozen --no-install-project --no-install-workspace --no-default-gr
|
|||
--extra extra_proxy \
|
||||
--extra semantic-router \
|
||||
--extra saml \
|
||||
--python python3
|
||||
--extra bedrock-realtime \
|
||||
--python python3.13
|
||||
|
||||
# Copy full source tree
|
||||
COPY . .
|
||||
|
|
@ -86,7 +88,8 @@ RUN uv sync --frozen --no-default-groups --no-editable \
|
|||
--extra extra_proxy \
|
||||
--extra semantic-router \
|
||||
--extra saml \
|
||||
--python python3
|
||||
--extra bedrock-realtime \
|
||||
--python python3.13
|
||||
|
||||
RUN HOME=/opt/prisma XDG_CACHE_HOME=/opt/prisma/.cache PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries \
|
||||
npm_config_cache=/root/.npm \
|
||||
|
|
@ -100,8 +103,14 @@ FROM $LITELLM_RUNTIME_IMAGE AS runtime
|
|||
|
||||
USER root
|
||||
|
||||
# The base image only configures Chainguard's authenticated apk repo, which
|
||||
# requires an enterprise subscription. Add the public Wolfi repo so `apk add`
|
||||
# also works for anyone installing extra packages into a running container.
|
||||
# https://github.com/BerriAI/litellm/issues/33518
|
||||
RUN echo "https://packages.wolfi.dev/os" >> /etc/apk/repositories
|
||||
|
||||
# node (without npm) is required by the prisma CLI at runtime
|
||||
RUN apk add --no-cache bash openssl tzdata nodejs python3 libsndfile
|
||||
RUN apk add --no-cache bash openssl tzdata nodejs python-3.13 libsndfile
|
||||
|
||||
WORKDIR /app
|
||||
ENV PATH="/app/.venv/bin:${PATH}" \
|
||||
|
|
|
|||
|
|
@ -354,6 +354,8 @@ curl -X POST 'http://0.0.0.0:4000/v1/chat/completions' \
|
|||
| [Petals (`petals`)](https://docs.litellm.ai/docs/providers/petals) | ✅ | ✅ | ✅ | | | | | | | |
|
||||
| [Pinstripes (`pinstripes`)](https://docs.litellm.ai/docs/providers/pinstripes) | ✅ | ✅ | ✅ | | | | | | | |
|
||||
| [Predibase (`predibase`)](https://docs.litellm.ai/docs/providers/predibase) | ✅ | ✅ | ✅ | | | | | | | |
|
||||
| [Qwen AI Platform (`qwen_ai_platform`)](https://docs.litellm.ai/docs/providers/qwencloud) | ✅ | ✅ | ✅ | ✅ | ✅ | | | | | ✅ |
|
||||
| [QwenCloud (`qwencloud`)](https://docs.litellm.ai/docs/providers/qwencloud) | ✅ | ✅ | ✅ | ✅ | ✅ | | | | | ✅ |
|
||||
| [Recraft (`recraft`)](https://docs.litellm.ai/docs/providers/recraft) | | | | | ✅ | | | | | |
|
||||
| [Replicate (`replicate`)](https://docs.litellm.ai/docs/providers/replicate) | ✅ | ✅ | ✅ | | | | | | | |
|
||||
| [Sagemaker Chat (`sagemaker_chat`)](https://docs.litellm.ai/docs/providers/aws_sagemaker) | ✅ | ✅ | ✅ | | | | | | | |
|
||||
|
|
|
|||
|
|
@ -1,5 +1,5 @@
|
|||
ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:a31344ab2cb8618db84f535eec56f76f6178b142cb92cb2e48676cc2dcebea72
|
||||
ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:a31344ab2cb8618db84f535eec56f76f6178b142cb92cb2e48676cc2dcebea72
|
||||
ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:e624c5d5e42382ce7165ddafcbbf8e6769a24cbd02ea6114b880b05ae5ba2a8d
|
||||
ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:e624c5d5e42382ce7165ddafcbbf8e6769a24cbd02ea6114b880b05ae5ba2a8d
|
||||
ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a
|
||||
|
||||
FROM $UV_IMAGE AS uvbin
|
||||
|
|
@ -16,7 +16,7 @@ COPY --from=uvbin /uv /uvx /usr/local/bin/
|
|||
# instead of nodeenv downloading one whose dynamic deps may not be in Wolfi
|
||||
# (e.g. Node 26.2.0 needs libatomic). Retry for transient apk.cgr.dev flakes.
|
||||
RUN for i in 1 2 3; do \
|
||||
apk add --no-cache bash gcc python3 python3-dev openssl openssl-dev libsndfile nodejs npm && break; \
|
||||
apk add --no-cache bash gcc python-3.13 python-3.13-dev openssl openssl-dev libsndfile nodejs npm && break; \
|
||||
[ $i = 3 ] && { echo "apk add failed after 3 retries" >&2; exit 1; }; \
|
||||
sleep 5; \
|
||||
done
|
||||
|
|
@ -46,7 +46,8 @@ RUN --mount=type=cache,target=/root/.cache/uv \
|
|||
--extra proxy-runtime \
|
||||
--extra extra_proxy \
|
||||
--extra semantic-router \
|
||||
--python python3
|
||||
--extra saml \
|
||||
--python python3.13
|
||||
|
||||
# Stage 2 — copy source and install the project + workspace members.
|
||||
COPY . .
|
||||
|
|
@ -57,7 +58,8 @@ RUN --mount=type=cache,target=/root/.cache/uv \
|
|||
--extra proxy-runtime \
|
||||
--extra extra_proxy \
|
||||
--extra semantic-router \
|
||||
--python python3
|
||||
--extra saml \
|
||||
--python python3.13
|
||||
|
||||
RUN HOME=/opt/prisma XDG_CACHE_HOME=/opt/prisma/.cache PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries \
|
||||
npm_config_cache=/root/.npm \
|
||||
|
|
@ -71,7 +73,7 @@ FROM $LITELLM_RUNTIME_IMAGE AS runtime
|
|||
USER root
|
||||
|
||||
RUN for i in 1 2 3; do \
|
||||
apk add --no-cache bash openssl tzdata python3 libsndfile libatomic && break; \
|
||||
apk add --no-cache bash openssl tzdata python-3.13 libsndfile libatomic && break; \
|
||||
[ $i = 3 ] && { echo "apk add failed after 3 retries" >&2; exit 1; }; \
|
||||
sleep 5; \
|
||||
done
|
||||
|
|
|
|||
|
|
@ -1,9 +1,9 @@
|
|||
{
|
||||
"reportAny": {
|
||||
"limit": 17270
|
||||
"limit": 14076
|
||||
},
|
||||
"reportArgumentType": {
|
||||
"limit": 2538
|
||||
"limit": 2216
|
||||
},
|
||||
"reportAssignmentType": {
|
||||
"limit": 319
|
||||
|
|
@ -18,13 +18,13 @@
|
|||
"limit": 40
|
||||
},
|
||||
"reportDeprecated": {
|
||||
"limit": 212
|
||||
"limit": 211
|
||||
},
|
||||
"reportDuplicateImport": {
|
||||
"limit": 19
|
||||
},
|
||||
"reportExplicitAny": {
|
||||
"limit": 5485
|
||||
"limit": 4128
|
||||
},
|
||||
"reportFunctionMemberAccess": {
|
||||
"limit": 7
|
||||
|
|
@ -42,7 +42,7 @@
|
|||
"limit": 12
|
||||
},
|
||||
"reportIndexIssue": {
|
||||
"limit": 35
|
||||
"limit": 25
|
||||
},
|
||||
"reportInvalidTypeForm": {
|
||||
"limit": 34
|
||||
|
|
@ -54,10 +54,10 @@
|
|||
"limit": 0
|
||||
},
|
||||
"reportMissingParameterType": {
|
||||
"limit": 5658
|
||||
"limit": 5601
|
||||
},
|
||||
"reportMissingTypeArgument": {
|
||||
"limit": 15425
|
||||
"limit": 15306
|
||||
},
|
||||
"reportMissingTypeStubs": {
|
||||
"limit": 40
|
||||
|
|
@ -72,7 +72,7 @@
|
|||
"limit": 0
|
||||
},
|
||||
"reportOptionalMemberAccess": {
|
||||
"limit": 1055
|
||||
"limit": 0
|
||||
},
|
||||
"reportOptionalOperand": {
|
||||
"limit": 0
|
||||
|
|
@ -90,40 +90,40 @@
|
|||
"limit": 8
|
||||
},
|
||||
"reportReturnType": {
|
||||
"limit": 213
|
||||
"limit": 181
|
||||
},
|
||||
"reportTypedDictNotRequiredAccess": {
|
||||
"limit": 25
|
||||
"limit": 24
|
||||
},
|
||||
"reportUndefinedVariable": {
|
||||
"limit": 0
|
||||
},
|
||||
"reportUnknownArgumentType": {
|
||||
"limit": 44526
|
||||
"limit": 44364
|
||||
},
|
||||
"reportUnknownLambdaType": {
|
||||
"limit": 109
|
||||
},
|
||||
"reportUnknownMemberType": {
|
||||
"limit": 38721
|
||||
"limit": 38350
|
||||
},
|
||||
"reportUnknownParameterType": {
|
||||
"limit": 19778
|
||||
"limit": 19626
|
||||
},
|
||||
"reportUnknownVariableType": {
|
||||
"limit": 30290
|
||||
"limit": 29890
|
||||
},
|
||||
"reportUnnecessaryCast": {
|
||||
"limit": 117
|
||||
"limit": 111
|
||||
},
|
||||
"reportUnnecessaryComparison": {
|
||||
"limit": 697
|
||||
"limit": 692
|
||||
},
|
||||
"reportUnnecessaryContains": {
|
||||
"limit": 5
|
||||
},
|
||||
"reportUnnecessaryIsInstance": {
|
||||
"limit": 829
|
||||
"limit": 826
|
||||
},
|
||||
"reportUntypedBaseClass": {
|
||||
"limit": 0
|
||||
|
|
@ -141,6 +141,6 @@
|
|||
"limit": 543
|
||||
},
|
||||
"reportUnusedVariable": {
|
||||
"limit": 139
|
||||
"limit": 137
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,10 +1,10 @@
|
|||
# syntax=docker/dockerfile:1.7
|
||||
|
||||
# Base image for building
|
||||
ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:a31344ab2cb8618db84f535eec56f76f6178b142cb92cb2e48676cc2dcebea72
|
||||
ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:e624c5d5e42382ce7165ddafcbbf8e6769a24cbd02ea6114b880b05ae5ba2a8d
|
||||
|
||||
# Runtime image
|
||||
ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:a31344ab2cb8618db84f535eec56f76f6178b142cb92cb2e48676cc2dcebea72
|
||||
ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:e624c5d5e42382ce7165ddafcbbf8e6769a24cbd02ea6114b880b05ae5ba2a8d
|
||||
ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a
|
||||
# Pinned by digest like the other base images; bump explicitly on Node upgrades.
|
||||
ARG UI_BUILD_IMAGE=node:24.19-alpine3.24@sha256:d32cdf619f63fe0471182d08996dd516c6275bb5fd31ae06e55a570bd9e1ad43
|
||||
|
|
@ -39,8 +39,8 @@ COPY --from=uvbin /uvx /usr/local/bin/uvx
|
|||
RUN apk add --no-cache \
|
||||
bash \
|
||||
gcc \
|
||||
python3 \
|
||||
python3-dev \
|
||||
python-3.13 \
|
||||
python-3.13-dev \
|
||||
openssl \
|
||||
openssl-dev \
|
||||
nodejs \
|
||||
|
|
@ -49,6 +49,7 @@ RUN apk add --no-cache \
|
|||
|
||||
ENV UV_PROJECT_ENVIRONMENT=/app/.venv \
|
||||
UV_LINK_MODE=copy \
|
||||
UV_PYTHON_DOWNLOADS=0 \
|
||||
PATH="/app/.venv/bin:${PATH}"
|
||||
|
||||
# Copy dependency metadata first for layer caching
|
||||
|
|
@ -63,7 +64,8 @@ RUN uv sync --frozen --no-install-project --no-install-workspace --no-default-gr
|
|||
--extra extra_proxy \
|
||||
--extra semantic-router \
|
||||
--extra saml \
|
||||
--python python3
|
||||
--extra bedrock-realtime \
|
||||
--python python3.13
|
||||
|
||||
# Copy full source tree
|
||||
COPY . .
|
||||
|
|
@ -84,7 +86,8 @@ RUN uv sync --frozen --no-default-groups --no-editable \
|
|||
--extra extra_proxy \
|
||||
--extra semantic-router \
|
||||
--extra saml \
|
||||
--python python3
|
||||
--extra bedrock-realtime \
|
||||
--python python3.13
|
||||
|
||||
RUN HOME=/opt/prisma XDG_CACHE_HOME=/opt/prisma/.cache PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries \
|
||||
npm_config_cache=/root/.npm \
|
||||
|
|
@ -98,7 +101,7 @@ FROM $LITELLM_RUNTIME_IMAGE AS runtime
|
|||
USER root
|
||||
|
||||
# node (without npm) is required by the prisma CLI at runtime
|
||||
RUN apk add --no-cache bash openssl tzdata nodejs python3 libsndfile
|
||||
RUN apk add --no-cache bash openssl tzdata nodejs python-3.13 libsndfile
|
||||
|
||||
WORKDIR /app
|
||||
ENV PATH="/app/.venv/bin:${PATH}" \
|
||||
|
|
|
|||
|
|
@ -1,8 +1,8 @@
|
|||
# syntax=docker/dockerfile:1.7
|
||||
|
||||
# Base images
|
||||
ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:a31344ab2cb8618db84f535eec56f76f6178b142cb92cb2e48676cc2dcebea72
|
||||
ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:a31344ab2cb8618db84f535eec56f76f6178b142cb92cb2e48676cc2dcebea72
|
||||
ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:e624c5d5e42382ce7165ddafcbbf8e6769a24cbd02ea6114b880b05ae5ba2a8d
|
||||
ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:e624c5d5e42382ce7165ddafcbbf8e6769a24cbd02ea6114b880b05ae5ba2a8d
|
||||
ARG PROXY_EXTRAS_SOURCE=published
|
||||
ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a
|
||||
# Pinned by digest like the other base images; bump explicitly on Node upgrades.
|
||||
|
|
@ -37,8 +37,8 @@ COPY --from=uvbin /uvx /usr/local/bin/uvx
|
|||
|
||||
RUN for i in 1 2 3; do \
|
||||
apk add --no-cache \
|
||||
python3 \
|
||||
python3-dev \
|
||||
python-3.13 \
|
||||
python-3.13-dev \
|
||||
gcc \
|
||||
rust \
|
||||
bash \
|
||||
|
|
@ -52,6 +52,7 @@ RUN for i in 1 2 3; do \
|
|||
|
||||
ENV UV_PROJECT_ENVIRONMENT=/app/.venv \
|
||||
UV_LINK_MODE=copy \
|
||||
UV_PYTHON_DOWNLOADS=0 \
|
||||
PATH="/app/.venv/bin:${PATH}" \
|
||||
LITELLM_NON_ROOT=true \
|
||||
XDG_CACHE_HOME=/app/.cache
|
||||
|
|
@ -69,7 +70,8 @@ RUN --mount=type=cache,target=/app/.cache/uv,id=litellm-uv-cache \
|
|||
--extra extra_proxy \
|
||||
--extra semantic-router \
|
||||
--extra saml \
|
||||
--python python3
|
||||
--extra bedrock-realtime \
|
||||
--python python3.13
|
||||
|
||||
# Copy full source tree
|
||||
COPY . .
|
||||
|
|
@ -96,7 +98,8 @@ RUN --mount=type=cache,target=/app/.cache/uv,id=litellm-uv-cache \
|
|||
--extra extra_proxy \
|
||||
--extra semantic-router \
|
||||
--extra saml \
|
||||
--python python3 \
|
||||
--extra bedrock-realtime \
|
||||
--python python3.13 \
|
||||
--no-sources-package litellm-proxy-extras; \
|
||||
else \
|
||||
uv sync --frozen --no-default-groups --no-editable \
|
||||
|
|
@ -105,7 +108,8 @@ RUN --mount=type=cache,target=/app/.cache/uv,id=litellm-uv-cache \
|
|||
--extra extra_proxy \
|
||||
--extra semantic-router \
|
||||
--extra saml \
|
||||
--python python3; \
|
||||
--extra bedrock-realtime \
|
||||
--python python3.13; \
|
||||
fi
|
||||
|
||||
RUN HOME=/opt/prisma XDG_CACHE_HOME=/opt/prisma/.cache PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries \
|
||||
|
|
@ -124,7 +128,7 @@ RUN for i in 1 2 3; do \
|
|||
apk upgrade --no-cache && break || sleep 5; \
|
||||
done && \
|
||||
for i in 1 2 3; do \
|
||||
apk add --no-cache python3 bash openssl tzdata libsndfile nodejs && break || sleep 5; \
|
||||
apk add --no-cache python-3.13 bash openssl tzdata libsndfile nodejs && break || sleep 5; \
|
||||
done
|
||||
|
||||
# Copy only what runtime needs. The application is installed inside the venv;
|
||||
|
|
|
|||
|
|
@ -7,7 +7,7 @@ GET - /audit/{id} - Get audit log by id
|
|||
GET - /audit - Get all audit logs
|
||||
"""
|
||||
|
||||
from typing import TYPE_CHECKING, Final, Optional
|
||||
from typing import TYPE_CHECKING, Final
|
||||
|
||||
#### AUDIT LOGGING ####
|
||||
from fastapi import APIRouter, Depends, HTTPException, Query
|
||||
|
|
@ -58,33 +58,33 @@ async def get_audit_logs(
|
|||
page: int = Query(1, ge=1),
|
||||
page_size: int = Query(10, ge=1, le=100),
|
||||
# Filter parameters
|
||||
changed_by: Optional[str] = Query(
|
||||
changed_by: str | None = Query(
|
||||
None, description="Filter by user or system that performed the action"
|
||||
),
|
||||
changed_by_api_key: Optional[str] = Query(
|
||||
changed_by_api_key: str | None = Query(
|
||||
None, description="Filter by API key hash that performed the action"
|
||||
),
|
||||
action: Optional[str] = Query(
|
||||
action: str | None = Query(
|
||||
None, description="Filter by action type (create, update, delete)"
|
||||
),
|
||||
table_name: Optional[str] = Query(
|
||||
table_name: str | None = Query(
|
||||
None, description="Filter by table name that was modified"
|
||||
),
|
||||
object_id: Optional[str] = Query(
|
||||
object_id: str | None = Query(
|
||||
None, description="Filter by ID of the object that was modified"
|
||||
),
|
||||
start_date: Optional[str] = Query(None, description="Filter logs after this date"),
|
||||
end_date: Optional[str] = Query(None, description="Filter logs before this date"),
|
||||
object_team_id: Optional[str] = Query(
|
||||
start_date: str | None = Query(None, description="Filter logs after this date"),
|
||||
end_date: str | None = Query(None, description="Filter logs before this date"),
|
||||
object_team_id: str | None = Query(
|
||||
None,
|
||||
description="Filter by team_id present in before_value or updated_values JSON (PostgreSQL only)",
|
||||
),
|
||||
object_key_hash: Optional[str] = Query(
|
||||
object_key_hash: str | None = Query(
|
||||
None,
|
||||
description="Filter by token (key hash) present in before_value or updated_values JSON (PostgreSQL only)",
|
||||
),
|
||||
# Sorting parameters
|
||||
sort_by: Optional[str] = Query(
|
||||
sort_by: str | None = Query(
|
||||
None,
|
||||
description="Column to sort by (e.g. 'updated_at', 'action', 'table_name')",
|
||||
),
|
||||
|
|
|
|||
|
|
@ -5,7 +5,7 @@ Polls LiteLLM_ManagedObjectTable to check if the batch job is complete, and if t
|
|||
from dataclasses import replace as dataclasses_replace
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Any, Dict, Final, List, Literal, Optional, Tuple, cast
|
||||
from typing import TYPE_CHECKING, Final, List, Literal, Optional, Tuple, cast
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm._uuid import uuid
|
||||
|
|
@ -87,7 +87,7 @@ class CheckBatchCost:
|
|||
return
|
||||
self.batch_processed_support_confirmed = True
|
||||
|
||||
async def _get_user_info(self, batch_id: str, user_id: Optional[str]) -> Dict[str, Any]:
|
||||
async def _get_user_info(self, batch_id: str, user_id: Optional[str]) -> dict[str, str | None]:
|
||||
"""
|
||||
Look up user email and key alias by user_id for enriching the S3 callback metadata.
|
||||
Returns a dict with user_api_key_user_email and user_api_key_alias (both may be None).
|
||||
|
|
@ -97,8 +97,10 @@ class CheckBatchCost:
|
|||
if not user_id:
|
||||
return {}
|
||||
try:
|
||||
user_row = await self.prisma_client.db.litellm_usertable.find_unique(
|
||||
where={"user_id": user_id}
|
||||
user_row: prisma_models.LiteLLM_UserTable | None = (
|
||||
await self.prisma_client.db.litellm_usertable.find_unique(
|
||||
where={"user_id": user_id}
|
||||
)
|
||||
)
|
||||
if user_row is None:
|
||||
return {}
|
||||
|
|
@ -115,8 +117,10 @@ class CheckBatchCost:
|
|||
if not api_key:
|
||||
return None
|
||||
try:
|
||||
key_row = await self.prisma_client.db.litellm_verificationtoken.find_unique(
|
||||
where={"token": api_key}
|
||||
key_row: prisma_models.LiteLLM_VerificationToken | None = (
|
||||
await self.prisma_client.db.litellm_verificationtoken.find_unique(
|
||||
where={"token": api_key}
|
||||
)
|
||||
)
|
||||
return getattr(key_row, "key_alias", None) if key_row is not None else None
|
||||
except Exception as e:
|
||||
|
|
@ -128,8 +132,10 @@ class CheckBatchCost:
|
|||
if not team_id:
|
||||
return None
|
||||
try:
|
||||
team_row = await self.prisma_client.db.litellm_teamtable.find_unique(
|
||||
where={"team_id": team_id}
|
||||
team_row: prisma_models.LiteLLM_TeamTable | None = (
|
||||
await self.prisma_client.db.litellm_teamtable.find_unique(
|
||||
where={"team_id": team_id}
|
||||
)
|
||||
)
|
||||
return getattr(team_row, "team_alias", None) if team_row is not None else None
|
||||
except Exception as e:
|
||||
|
|
@ -138,7 +144,7 @@ class CheckBatchCost:
|
|||
|
||||
async def _build_creator_attribution_metadata(
|
||||
self, job: "LiteLLM_ManagedObjectTable", batch_id: str
|
||||
) -> Dict[str, Any]:
|
||||
) -> dict[str, object]:
|
||||
"""
|
||||
Rebuild the spend-tracking metadata for the key, team, and tags that created the
|
||||
batch so the batch-cost spend log is attributed the same way a non-batch request
|
||||
|
|
@ -152,7 +158,7 @@ class CheckBatchCost:
|
|||
team_id = getattr(job, "team_id", None)
|
||||
request_tags = getattr(job, "request_tags", None)
|
||||
|
||||
metadata: Dict[str, Any] = {
|
||||
metadata: dict[str, object] = {
|
||||
"user_api_key_user_id": job.created_by,
|
||||
"user_api_key": api_key,
|
||||
"user_api_key_team_id": team_id,
|
||||
|
|
|
|||
|
|
@ -182,6 +182,10 @@ class _ManagedObjectTableActions(Protocol):
|
|||
async def update_many(self, where: Mapping[str, object], data: Mapping[str, object]) -> int: ...
|
||||
|
||||
|
||||
class _SchedulerWithJobLookup(Protocol):
|
||||
def get_job(self, job_id: str) -> object: ...
|
||||
|
||||
|
||||
class _CursorPageArgs(TypedDict, total=False):
|
||||
cursor: Mapping[str, str]
|
||||
skip: int
|
||||
|
|
@ -853,7 +857,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
file_ids.append(file_id)
|
||||
return file_ids
|
||||
|
||||
def get_file_ids_from_responses_input(self, input: Union[str, List[Dict[str, Any]]]) -> List[str]:
|
||||
def get_file_ids_from_responses_input(self, input: Union[str, List[Dict[str, object]]]) -> List[str]:
|
||||
"""
|
||||
Gets file ids from responses API input.
|
||||
|
||||
|
|
@ -878,7 +882,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
# Check for direct input_file type
|
||||
if item.get("type") == "input_file":
|
||||
file_id = item.get("file_id")
|
||||
if file_id:
|
||||
if isinstance(file_id, str) and file_id:
|
||||
file_ids.append(file_id)
|
||||
|
||||
# Check for input_file in content array
|
||||
|
|
@ -887,7 +891,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
for content_item in content:
|
||||
if isinstance(content_item, dict) and content_item.get("type") == "input_file":
|
||||
file_id = content_item.get("file_id")
|
||||
if file_id:
|
||||
if isinstance(file_id, str) and file_id:
|
||||
file_ids.append(file_id)
|
||||
|
||||
return file_ids
|
||||
|
|
@ -1227,7 +1231,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
|
||||
# Handle both output_file_id and error_file_id
|
||||
for file_attr in ["output_file_id", "error_file_id"]:
|
||||
file_id_value = getattr(response, file_attr, None)
|
||||
file_id_value: str | None = getattr(response, file_attr, None)
|
||||
if file_id_value and model_id:
|
||||
decoded_output_file_id = _is_base64_encoded_unified_file_id(file_id_value)
|
||||
if decoded_output_file_id and "llm_output_file_id," in decoded_output_file_id:
|
||||
|
|
@ -1496,7 +1500,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
import litellm.proxy.proxy_server as proxy_server_module
|
||||
|
||||
# Check if the scheduler has the batch cost checking job registered
|
||||
scheduler = getattr(proxy_server_module, "scheduler", None)
|
||||
scheduler: Final[_SchedulerWithJobLookup | None] = getattr(proxy_server_module, "scheduler", None)
|
||||
if scheduler is None:
|
||||
return False
|
||||
|
||||
|
|
@ -1542,7 +1546,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
)
|
||||
MAX_MATCHES_TO_RETURN = 10
|
||||
|
||||
batches = await self.prisma_client.db.litellm_managedobjecttable.find_many(
|
||||
batches = await _managed_object_table(self.prisma_client).find_many(
|
||||
where={
|
||||
"file_purpose": "batch",
|
||||
"batch_processed": False,
|
||||
|
|
@ -1552,11 +1556,14 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
order={"created_at": "desc"},
|
||||
)
|
||||
|
||||
referencing_batches = []
|
||||
referencing_batches: Final[list[dict[str, object]]] = []
|
||||
for batch in batches:
|
||||
try:
|
||||
# Parse the batch file_object to check for file references
|
||||
batch_data = json.loads(batch.file_object) if isinstance(batch.file_object, str) else batch.file_object
|
||||
decoded_file_object = _decode_json_blob(batch.file_object)
|
||||
batch_data: Mapping[str, object] = (
|
||||
decoded_file_object if isinstance(decoded_file_object, Mapping) else {}
|
||||
)
|
||||
|
||||
# Extract file IDs from batch
|
||||
# Batches typically reference the unified file ID in input_file_id
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
[project]
|
||||
name = "litellm-enterprise"
|
||||
version = "0.1.62"
|
||||
version = "0.1.63"
|
||||
description = "Package for LiteLLM Enterprise features"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.9"
|
||||
|
|
@ -26,7 +26,7 @@ required-version = ">=0.10.9"
|
|||
module-root = ""
|
||||
|
||||
[tool.commitizen]
|
||||
version = "0.1.62"
|
||||
version = "0.1.63"
|
||||
version_files = [
|
||||
"pyproject.toml:^version",
|
||||
"../pyproject.toml:litellm-enterprise==",
|
||||
|
|
|
|||
|
|
@ -1,5 +1,5 @@
|
|||
ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:a31344ab2cb8618db84f535eec56f76f6178b142cb92cb2e48676cc2dcebea72
|
||||
ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:a31344ab2cb8618db84f535eec56f76f6178b142cb92cb2e48676cc2dcebea72
|
||||
ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:e624c5d5e42382ce7165ddafcbbf8e6769a24cbd02ea6114b880b05ae5ba2a8d
|
||||
ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:e624c5d5e42382ce7165ddafcbbf8e6769a24cbd02ea6114b880b05ae5ba2a8d
|
||||
ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a
|
||||
|
||||
FROM $UV_IMAGE AS uvbin
|
||||
|
|
@ -16,7 +16,7 @@ COPY --from=uvbin /uv /uvx /usr/local/bin/
|
|||
# instead of nodeenv downloading one whose dynamic deps may not be in Wolfi
|
||||
# (e.g. Node 26.2.0 needs libatomic). Retry for transient apk.cgr.dev flakes.
|
||||
RUN for i in 1 2 3; do \
|
||||
apk add --no-cache bash gcc python3 python3-dev openssl openssl-dev libsndfile nodejs npm && break; \
|
||||
apk add --no-cache bash gcc python-3.13 python-3.13-dev openssl openssl-dev libsndfile nodejs npm && break; \
|
||||
[ $i = 3 ] && { echo "apk add failed after 3 retries" >&2; exit 1; }; \
|
||||
sleep 5; \
|
||||
done
|
||||
|
|
@ -47,7 +47,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
|
|||
--extra extra_proxy \
|
||||
--extra semantic-router \
|
||||
--extra bedrock-realtime \
|
||||
--python python3
|
||||
--python python3.13
|
||||
|
||||
# Stage 2 — copy source and install the project + workspace members.
|
||||
COPY . .
|
||||
|
|
@ -59,7 +59,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
|
|||
--extra extra_proxy \
|
||||
--extra semantic-router \
|
||||
--extra bedrock-realtime \
|
||||
--python python3
|
||||
--python python3.13
|
||||
|
||||
RUN HOME=/opt/prisma XDG_CACHE_HOME=/opt/prisma/.cache PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries \
|
||||
npm_config_cache=/root/.npm \
|
||||
|
|
@ -73,7 +73,7 @@ FROM $LITELLM_RUNTIME_IMAGE AS runtime
|
|||
USER root
|
||||
|
||||
RUN for i in 1 2 3; do \
|
||||
apk add --no-cache bash openssl tzdata python3 libsndfile libatomic && break; \
|
||||
apk add --no-cache bash openssl tzdata python-3.13 libsndfile libatomic && break; \
|
||||
[ $i = 3 ] && { echo "apk add failed after 3 retries" >&2; exit 1; }; \
|
||||
sleep 5; \
|
||||
done
|
||||
|
|
|
|||
|
|
@ -86,6 +86,7 @@ GATEWAY_PATH_PREFIXES: tuple[str, ...] = (
|
|||
"/comprehendmedical",
|
||||
"/cohere/",
|
||||
"/gemini/",
|
||||
"/gigachat/",
|
||||
"/google/",
|
||||
"/vertex_ai/",
|
||||
"/vertex-ai/",
|
||||
|
|
|
|||
|
|
@ -18,7 +18,7 @@ type: application
|
|||
# This is the chart version. This version number should be incremented each time you make changes
|
||||
# to the chart and its templates, including the app version.
|
||||
# Versions are expected to follow Semantic Versioning (https://semver.org/)
|
||||
version: 1.1.2
|
||||
version: 1.1.3
|
||||
|
||||
# This is the version number of the application being deployed. This version number should be
|
||||
# incremented each time you make changes to the application. Versions are not expected to
|
||||
|
|
|
|||
|
|
@ -26,7 +26,7 @@ If `db.useStackgresOperator` is used (not yet implemented):
|
|||
| `replicaCount` | The number of LiteLLM Proxy pods to be deployed | `1` |
|
||||
| `masterkeySecretName` | The name of the Kubernetes Secret that contains the Master API Key for LiteLLM. If not specified, use the generated secret name. | N/A |
|
||||
| `masterkeySecretKey` | The key within the Kubernetes Secret that contains the Master API Key for LiteLLM. If not specified, use `masterkey` as the key. | N/A |
|
||||
| `masterkey` | The Master API Key for LiteLLM. If not specified, a random key in the `sk-...` format is generated. | N/A |
|
||||
| `masterkey` | The Master API Key for LiteLLM. If not specified, a random key in the `sk-...` format is generated on first install and reused on upgrades. | N/A |
|
||||
| `environmentSecrets` | An optional array of Secret object names. The keys and values in these secrets will be presented to the LiteLLM proxy pod as environment variables. See below for an example Secret object. | `[]` |
|
||||
| `environmentConfigMaps` | An optional array of ConfigMap object names. The keys and values in these configmaps will be presented to the LiteLLM proxy pod as environment variables. See below for an example Secret object. | `[]` |
|
||||
| `image.repository` | LiteLLM Proxy image repository | `ghcr.io/berriai/litellm` |
|
||||
|
|
@ -212,6 +212,8 @@ service, the **Proxy Endpoint** should be set to `http://<RELEASE>-litellm:4000`
|
|||
The **Proxy Key** is the value specified for `masterkey` or, if a `masterkey`
|
||||
was not provided to the helm command line, the `masterkey` is a randomly
|
||||
generated string in the `sk-...` format stored in the `<RELEASE>-litellm-masterkey` Kubernetes Secret.
|
||||
The key is generated once on the first install; later `helm upgrade` runs reuse the
|
||||
value already in that Secret, so upgrading never rotates the master key.
|
||||
|
||||
```bash
|
||||
kubectl -n litellm get secret <RELEASE>-litellm-masterkey -o jsonpath="{.data.masterkey}"
|
||||
|
|
|
|||
|
|
@ -1,9 +1,11 @@
|
|||
{{- if not .Values.masterkeySecretName }}
|
||||
{{ $masterkey := (.Values.masterkey | default (printf "sk-%s" (randAlphaNum 18))) }}
|
||||
{{- $secretName := printf "%s-masterkey" (include "litellm.fullname" .) }}
|
||||
{{- $existing := lookup "v1" "Secret" .Release.Namespace $secretName }}
|
||||
{{- $masterkey := .Values.masterkey | default (dig "data" "masterkey" "" $existing | b64dec) | default (printf "sk-%s" (randAlphaNum 18)) }}
|
||||
apiVersion: v1
|
||||
kind: Secret
|
||||
metadata:
|
||||
name: {{ include "litellm.fullname" . }}-masterkey
|
||||
name: {{ $secretName }}
|
||||
data:
|
||||
masterkey: {{ $masterkey | b64enc }}
|
||||
type: Opaque
|
||||
|
|
|
|||
|
|
@ -1,4 +1,4 @@
|
|||
suite: "hpa with behavior"
|
||||
suite: "hpa"
|
||||
templates:
|
||||
- hpa.yaml
|
||||
tests:
|
||||
|
|
@ -23,14 +23,44 @@ tests:
|
|||
- equal: { path: spec.behavior.scaleUp.stabilizationWindowSeconds, value: 60 }
|
||||
- equal: { path: spec.behavior.scaleDown.stabilizationWindowSeconds, value: 90 }
|
||||
|
||||
---
|
||||
suite: "hpa without behavior"
|
||||
templates:
|
||||
- hpa.yaml
|
||||
tests:
|
||||
- it: "does not render behavior when not set"
|
||||
set:
|
||||
autoscaling.enabled: true
|
||||
asserts:
|
||||
- isKind: { of: HorizontalPodAutoscaler }
|
||||
- isNull: { path: spec.behavior }
|
||||
|
||||
- it: "scales on cpu at the documented 60 percent by default"
|
||||
set:
|
||||
autoscaling.enabled: true
|
||||
asserts:
|
||||
- isKind: { of: HorizontalPodAutoscaler }
|
||||
- equal: { path: "spec.metrics[0].resource.name", value: cpu }
|
||||
- equal: { path: "spec.metrics[0].resource.target.type", value: Utilization }
|
||||
- equal: { path: "spec.metrics[0].resource.target.averageUtilization", value: 60 }
|
||||
|
||||
- it: "does not scale on memory by default"
|
||||
set:
|
||||
autoscaling.enabled: true
|
||||
asserts:
|
||||
- lengthEqual: { path: spec.metrics, count: 1 }
|
||||
|
||||
- it: "honours an explicit cpu target override"
|
||||
set:
|
||||
autoscaling.enabled: true
|
||||
autoscaling.targetCPUUtilizationPercentage: 75
|
||||
asserts:
|
||||
- equal: { path: "spec.metrics[0].resource.target.averageUtilization", value: 75 }
|
||||
|
||||
- it: "renders a memory metric only when a memory target is set"
|
||||
set:
|
||||
autoscaling.enabled: true
|
||||
autoscaling.targetMemoryUtilizationPercentage: 80
|
||||
asserts:
|
||||
- lengthEqual: { path: spec.metrics, count: 2 }
|
||||
- equal: { path: "spec.metrics[1].resource.name", value: memory }
|
||||
- equal: { path: "spec.metrics[1].resource.target.averageUtilization", value: 80 }
|
||||
|
||||
- it: "renders no hpa when autoscaling is disabled"
|
||||
asserts:
|
||||
- hasDocuments: { count: 0 }
|
||||
|
|
|
|||
|
|
@ -15,6 +15,53 @@ tests:
|
|||
# Note: The masterkey is generated as "sk-<18-random-chars>" in plain text,
|
||||
# but stored as base64 encoded in Kubernetes secret (requirement).
|
||||
# "sk-" base64 encodes to "c2st", so we check for "^c2st" pattern.
|
||||
- it: should reuse the master key already stored in the cluster instead of generating a new one on upgrade
|
||||
template: secret-masterkey.yaml
|
||||
set:
|
||||
masterkeySecretName: ""
|
||||
kubernetesProvider:
|
||||
scheme:
|
||||
"v1/Secret":
|
||||
gvr:
|
||||
version: "v1"
|
||||
resource: "secrets"
|
||||
namespaced: true
|
||||
objects:
|
||||
- kind: Secret
|
||||
apiVersion: v1
|
||||
metadata:
|
||||
name: RELEASE-NAME-litellm-masterkey
|
||||
namespace: NAMESPACE
|
||||
data:
|
||||
masterkey: c2stZXhpc3Rpbmcta2V5
|
||||
asserts:
|
||||
- equal:
|
||||
path: data.masterkey
|
||||
value: c2stZXhpc3Rpbmcta2V5
|
||||
- it: should let an explicit masterkey value override the one already stored in the cluster
|
||||
template: secret-masterkey.yaml
|
||||
set:
|
||||
masterkeySecretName: ""
|
||||
masterkey: sk-explicit
|
||||
kubernetesProvider:
|
||||
scheme:
|
||||
"v1/Secret":
|
||||
gvr:
|
||||
version: "v1"
|
||||
resource: "secrets"
|
||||
namespaced: true
|
||||
objects:
|
||||
- kind: Secret
|
||||
apiVersion: v1
|
||||
metadata:
|
||||
name: RELEASE-NAME-litellm-masterkey
|
||||
namespace: NAMESPACE
|
||||
data:
|
||||
masterkey: c2stZXhpc3Rpbmcta2V5
|
||||
asserts:
|
||||
- equal:
|
||||
path: data.masterkey
|
||||
value: c2stZXhwbGljaXQ=
|
||||
- it: should not create a secret if masterkeySecretName is set
|
||||
template: secret-masterkey.yaml
|
||||
set:
|
||||
|
|
|
|||
|
|
@ -200,7 +200,16 @@ autoscaling:
|
|||
enabled: false
|
||||
minReplicas: 1
|
||||
maxReplicas: 100
|
||||
targetCPUUtilizationPercentage: 80
|
||||
# 60 is the documented recommendation. See "Recommended Machine Specifications"
|
||||
# in https://docs.litellm.ai/docs/proxy/prod. A new replica clears the startupProbe
|
||||
# above only after up to failureThreshold x periodSeconds = 300 seconds, so a target
|
||||
# high enough to trip near saturation adds capacity minutes after it was needed.
|
||||
targetCPUUtilizationPercentage: 60
|
||||
# Deliberately left unset rather than given a value. The prisma query engine's
|
||||
# resident memory is a high-water mark that ratchets to the pod's worst-ever write
|
||||
# and is never returned, so a memory target reads the largest write a pod ever did
|
||||
# rather than what it is doing now, and replicas ratchet up without scaling back in.
|
||||
# Memory is a floor to provision under 'resources', not a signal to scale on.
|
||||
# targetMemoryUtilizationPercentage: 80
|
||||
# behavior: {}
|
||||
|
||||
|
|
|
|||
|
|
@ -7,6 +7,10 @@ metadata:
|
|||
{{- include "litellm.commonLabels" . | nindent 4 }}
|
||||
app.kubernetes.io/component: backend
|
||||
spec:
|
||||
{{- with .Values.backend.strategy }}
|
||||
strategy:
|
||||
{{- toYaml . | nindent 4 }}
|
||||
{{- end }}
|
||||
selector:
|
||||
matchLabels:
|
||||
{{- include "litellm.backend.selectorLabels" . | nindent 6 }}
|
||||
|
|
|
|||
|
|
@ -7,6 +7,10 @@ metadata:
|
|||
{{- include "litellm.commonLabels" . | nindent 4 }}
|
||||
app.kubernetes.io/component: gateway
|
||||
spec:
|
||||
{{- with .Values.gateway.strategy }}
|
||||
strategy:
|
||||
{{- toYaml . | nindent 4 }}
|
||||
{{- end }}
|
||||
selector:
|
||||
matchLabels:
|
||||
{{- include "litellm.gateway.selectorLabels" . | nindent 6 }}
|
||||
|
|
|
|||
|
|
@ -7,6 +7,8 @@
|
|||
#
|
||||
# Running this pre-upgrade closes the window where new application pods would
|
||||
# otherwise serve traffic against the previous release's unmigrated schema.
|
||||
# Argo CD users can swap the Helm hook for a PreSync hook through
|
||||
# `migrationJob.hooks`, which re-runs the Job on every sync.
|
||||
apiVersion: batch/v1
|
||||
kind: Job
|
||||
metadata:
|
||||
|
|
@ -14,10 +16,18 @@ metadata:
|
|||
labels:
|
||||
{{- include "litellm.commonLabels" . | nindent 4 }}
|
||||
app.kubernetes.io/component: migrations
|
||||
{{- if or .Values.migrationJob.hooks.helm.enabled .Values.migrationJob.hooks.argocd.enabled }}
|
||||
annotations:
|
||||
{{- if .Values.migrationJob.hooks.helm.enabled }}
|
||||
helm.sh/hook: pre-install,pre-upgrade
|
||||
helm.sh/hook-delete-policy: before-hook-creation
|
||||
helm.sh/hook-weight: "0"
|
||||
helm.sh/hook-weight: {{ .Values.migrationJob.hooks.helm.weight | default "0" | quote }}
|
||||
{{- end }}
|
||||
{{- if .Values.migrationJob.hooks.argocd.enabled }}
|
||||
argocd.argoproj.io/hook: PreSync
|
||||
argocd.argoproj.io/hook-delete-policy: BeforeHookCreation
|
||||
{{- end }}
|
||||
{{- end }}
|
||||
spec:
|
||||
backoffLimit: {{ .Values.migrationJob.backoffLimit }}
|
||||
ttlSecondsAfterFinished: {{ .Values.migrationJob.ttlSecondsAfterFinished }}
|
||||
|
|
|
|||
|
|
@ -7,6 +7,10 @@ metadata:
|
|||
{{- include "litellm.commonLabels" . | nindent 4 }}
|
||||
app.kubernetes.io/component: ui
|
||||
spec:
|
||||
{{- with .Values.ui.strategy }}
|
||||
strategy:
|
||||
{{- toYaml . | nindent 4 }}
|
||||
{{- end }}
|
||||
selector:
|
||||
matchLabels:
|
||||
{{- include "litellm.ui.selectorLabels" . | nindent 6 }}
|
||||
|
|
|
|||
63
helm/litellm/tests/migration_job_hooks_tests.yaml
Normal file
63
helm/litellm/tests/migration_job_hooks_tests.yaml
Normal file
|
|
@ -0,0 +1,63 @@
|
|||
suite: test migrations Job hook annotations
|
||||
templates:
|
||||
- migrations-job.yaml
|
||||
values:
|
||||
- ./values/required.yaml
|
||||
tests:
|
||||
- it: runs as a Helm pre-install / pre-upgrade hook by default
|
||||
asserts:
|
||||
- equal:
|
||||
path: metadata.annotations["helm.sh/hook"]
|
||||
value: pre-install,pre-upgrade
|
||||
- equal:
|
||||
path: metadata.annotations["helm.sh/hook-delete-policy"]
|
||||
value: before-hook-creation
|
||||
- equal:
|
||||
path: metadata.annotations["helm.sh/hook-weight"]
|
||||
value: "0"
|
||||
- notExists:
|
||||
path: metadata.annotations["argocd.argoproj.io/hook"]
|
||||
|
||||
- it: adds the Argo CD PreSync hook when asked
|
||||
set:
|
||||
migrationJob.hooks.argocd.enabled: true
|
||||
asserts:
|
||||
- equal:
|
||||
path: metadata.annotations["argocd.argoproj.io/hook"]
|
||||
value: PreSync
|
||||
- equal:
|
||||
path: metadata.annotations["argocd.argoproj.io/hook-delete-policy"]
|
||||
value: BeforeHookCreation
|
||||
|
||||
- it: drops the Helm hook so Argo CD owns the Job
|
||||
set:
|
||||
migrationJob.hooks.argocd.enabled: true
|
||||
migrationJob.hooks.helm.enabled: false
|
||||
asserts:
|
||||
- equal:
|
||||
path: metadata.annotations["argocd.argoproj.io/hook"]
|
||||
value: PreSync
|
||||
- notExists:
|
||||
path: metadata.annotations["helm.sh/hook"]
|
||||
- notExists:
|
||||
path: metadata.annotations["helm.sh/hook-delete-policy"]
|
||||
- notExists:
|
||||
path: metadata.annotations["helm.sh/hook-weight"]
|
||||
|
||||
- it: renders an ordinary Job when both hooks are disabled
|
||||
set:
|
||||
migrationJob.hooks.helm.enabled: false
|
||||
asserts:
|
||||
- notExists:
|
||||
path: metadata.annotations
|
||||
- equal:
|
||||
path: kind
|
||||
value: Job
|
||||
|
||||
- it: honours a custom Helm hook weight
|
||||
set:
|
||||
migrationJob.hooks.helm.weight: "-5"
|
||||
asserts:
|
||||
- equal:
|
||||
path: metadata.annotations["helm.sh/hook-weight"]
|
||||
value: "-5"
|
||||
66
helm/litellm/tests/rollout_strategy_tests.yaml
Normal file
66
helm/litellm/tests/rollout_strategy_tests.yaml
Normal file
|
|
@ -0,0 +1,66 @@
|
|||
suite: test rolling update strategy on the component deployments
|
||||
templates:
|
||||
- gateway/deployment.yaml
|
||||
- gateway/configmap.yaml
|
||||
- backend/deployment.yaml
|
||||
- ui/deployment.yaml
|
||||
values:
|
||||
- ./values/required.yaml
|
||||
tests:
|
||||
- it: leaves the strategy to Kubernetes defaults when unset
|
||||
asserts:
|
||||
- notExists:
|
||||
path: spec.strategy
|
||||
|
||||
- it: renders the configured strategy on each deployment
|
||||
set:
|
||||
gateway.strategy:
|
||||
type: RollingUpdate
|
||||
rollingUpdate:
|
||||
maxUnavailable: 0
|
||||
maxSurge: 1
|
||||
backend.strategy:
|
||||
type: RollingUpdate
|
||||
rollingUpdate:
|
||||
maxUnavailable: "25%"
|
||||
maxSurge: 2
|
||||
ui.strategy:
|
||||
type: Recreate
|
||||
asserts:
|
||||
- equal:
|
||||
path: spec.strategy
|
||||
value:
|
||||
type: RollingUpdate
|
||||
rollingUpdate:
|
||||
maxUnavailable: 0
|
||||
maxSurge: 1
|
||||
template: gateway/deployment.yaml
|
||||
- equal:
|
||||
path: spec.strategy
|
||||
value:
|
||||
type: RollingUpdate
|
||||
rollingUpdate:
|
||||
maxUnavailable: 25%
|
||||
maxSurge: 2
|
||||
template: backend/deployment.yaml
|
||||
- equal:
|
||||
path: spec.strategy
|
||||
value:
|
||||
type: Recreate
|
||||
template: ui/deployment.yaml
|
||||
|
||||
- it: keeps a component on the cluster default when only another one sets a strategy
|
||||
set:
|
||||
gateway.strategy:
|
||||
type: Recreate
|
||||
asserts:
|
||||
- equal:
|
||||
path: spec.strategy.type
|
||||
value: Recreate
|
||||
template: gateway/deployment.yaml
|
||||
- notExists:
|
||||
path: spec.strategy
|
||||
template: backend/deployment.yaml
|
||||
- notExists:
|
||||
path: spec.strategy
|
||||
template: ui/deployment.yaml
|
||||
|
|
@ -75,6 +75,22 @@ serviceAccounts:
|
|||
# generate` — the migration engine doesn't need the generated client.
|
||||
migrationJob:
|
||||
enabled: true
|
||||
# Which controller is responsible for running the Job.
|
||||
#
|
||||
# `helm.enabled` renders the Helm pre-install / pre-upgrade hook, so the Job
|
||||
# runs whenever `helm upgrade` sees a change to apply. `argocd.enabled`
|
||||
# renders an Argo CD PreSync hook instead, which runs the Job on every sync
|
||||
# even when the rendered manifests are unchanged: the way to re-run
|
||||
# migrations on demand from a GitOps pipeline. Turning the Helm hook off
|
||||
# while the Argo CD hook is on leaves the Job out of Helm's own upgrade
|
||||
# path, which is what Argo CD users want since Argo, not Helm, applies the
|
||||
# manifests.
|
||||
hooks:
|
||||
helm:
|
||||
enabled: true
|
||||
weight: "0"
|
||||
argocd:
|
||||
enabled: false
|
||||
backoffLimit: 4
|
||||
ttlSecondsAfterFinished: 120
|
||||
# Wall-clock budget for the whole Job, shared across every `backoffLimit`
|
||||
|
|
@ -257,6 +273,15 @@ gateway:
|
|||
initialDelaySeconds: 5
|
||||
periodSeconds: 10
|
||||
timeoutSeconds: 10
|
||||
# Rolling update tuning for the gateway Deployment. Empty by default, so
|
||||
# Kubernetes applies its own RollingUpdate defaults (25% maxSurge /
|
||||
# 25% maxUnavailable). Example, for a surge-only rollout behind a load
|
||||
# balancer that must never lose capacity:
|
||||
# type: RollingUpdate
|
||||
# rollingUpdate:
|
||||
# maxUnavailable: 0
|
||||
# maxSurge: 1
|
||||
strategy: {}
|
||||
# Optional startupProbe. Empty by default, so existing installs are unchanged
|
||||
# and liveness/readiness apply from container start. Set it to gate
|
||||
# liveness/readiness until a slow cold start finishes — a high failureThreshold
|
||||
|
|
@ -369,6 +394,8 @@ backend:
|
|||
initialDelaySeconds: 5
|
||||
periodSeconds: 10
|
||||
timeoutSeconds: 10
|
||||
# Same shape as gateway.strategy.
|
||||
strategy: {}
|
||||
# Optional startupProbe; same shape as gateway.startupProbe. Empty by default.
|
||||
startupProbe: {}
|
||||
hpa:
|
||||
|
|
@ -433,6 +460,8 @@ ui:
|
|||
httpGet: { path: /, port: http }
|
||||
initialDelaySeconds: 2
|
||||
periodSeconds: 10
|
||||
# Same shape as gateway.strategy.
|
||||
strategy: {}
|
||||
# Optional startupProbe; same shape as gateway.startupProbe. Empty by default.
|
||||
startupProbe: {}
|
||||
hpa:
|
||||
|
|
|
|||
|
|
@ -0,0 +1,21 @@
|
|||
DO $$
|
||||
BEGIN
|
||||
IF EXISTS (
|
||||
SELECT 1 FROM information_schema.columns
|
||||
WHERE table_name = 'LiteLLM_ShadowEvalJob' AND column_name = 'api_key_id'
|
||||
) THEN
|
||||
ALTER TABLE "LiteLLM_ShadowEvalJob" RENAME COLUMN "api_key_id" TO "target_id";
|
||||
END IF;
|
||||
END $$;
|
||||
|
||||
ALTER TABLE "LiteLLM_ShadowEvalJob" ADD COLUMN IF NOT EXISTS "target_type" TEXT NOT NULL DEFAULT 'key';
|
||||
|
||||
DROP INDEX IF EXISTS "LiteLLM_ShadowEvalJob_one_active_per_key_direction";
|
||||
|
||||
CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_ShadowEvalJob_one_active_per_target_direction"
|
||||
ON "LiteLLM_ShadowEvalJob"("target_type", "target_id", "direction") WHERE "stopped_at" IS NULL;
|
||||
|
||||
DROP INDEX IF EXISTS "LiteLLM_ShadowEvalJob_api_key_id_idx";
|
||||
|
||||
CREATE INDEX IF NOT EXISTS "LiteLLM_ShadowEvalJob_target_type_target_id_idx"
|
||||
ON "LiteLLM_ShadowEvalJob"("target_type", "target_id");
|
||||
|
|
@ -0,0 +1,3 @@
|
|||
ALTER TABLE "LiteLLM_ShadowEvalJob" ADD COLUMN IF NOT EXISTS "router_names" TEXT[] NOT NULL DEFAULT ARRAY[]::TEXT[];
|
||||
|
||||
ALTER TABLE "LiteLLM_ShadowEvalAttempt" ADD COLUMN IF NOT EXISTS "router_name" TEXT;
|
||||
|
|
@ -1529,14 +1529,16 @@ model LiteLLM_AutoRouterSession {
|
|||
model LiteLLM_ShadowEvalJob {
|
||||
id String @id @default(cuid())
|
||||
group_id String // legs of one job share this; the API's job id
|
||||
api_key_id String // hashed virtual key whose traffic this leg shadows
|
||||
router_name String // the auto-router under evaluation, in either direction
|
||||
target_type String @default("key") // key | team | user
|
||||
target_id String // hashed virtual key, team_id, or user_id whose traffic this leg shadows
|
||||
router_name String // first (often only) auto-router under evaluation; router_names is the full set
|
||||
router_names String[] @default([]) // all routers this job runs as shadow arms; empty on legacy rows, whose set is (router_name)
|
||||
direction String @default("forward") // forward | reverse
|
||||
baseline_model String? // reverse only: the fixed model the router is judged against
|
||||
judge_model String
|
||||
shadow_percentage Float
|
||||
max_turns Int // sample-count ceiling: the whole budget on pre-max_budget jobs, the error-loop valve otherwise
|
||||
max_budget Float? // per-key USD cap on the eval's own shadow + judge spend; null on jobs from before spend budgets
|
||||
max_budget Float? // per-target USD cap on the eval's own shadow + judge spend; null on jobs from before spend budgets
|
||||
created_at DateTime @default(now())
|
||||
created_by String?
|
||||
ends_at DateTime
|
||||
|
|
@ -1544,7 +1546,7 @@ model LiteLLM_ShadowEvalJob {
|
|||
stopped_by String? // operator who stopped it early; null when it ended on its own
|
||||
|
||||
@@index([group_id])
|
||||
@@index([api_key_id])
|
||||
@@index([target_type, target_id])
|
||||
@@index([created_at])
|
||||
}
|
||||
|
||||
|
|
@ -1554,6 +1556,7 @@ model LiteLLM_ShadowEvalAttempt {
|
|||
job_id String
|
||||
request_id String // the judged real request
|
||||
outcome String // real | shadow | tie | error
|
||||
router_name String? // the arm this verdict scores; NULL on legacy rows, meaning the job's own router
|
||||
tier String? // router's tier for the prompt, when classified
|
||||
real_model String?
|
||||
shadow_model String?
|
||||
|
|
|
|||
|
|
@ -512,6 +512,13 @@ class ProxyExtrasDBManager:
|
|||
try:
|
||||
import psycopg
|
||||
except ImportError:
|
||||
logger.warning(
|
||||
"psycopg is not installed; skipping the LiteLLM_SpendLogs "
|
||||
"partition check. If this table is partitioned (see "
|
||||
"db_scripts/partition_spend_logs.sql), schema reconciliation "
|
||||
"will try to rewrite its primary key and fail. Install the "
|
||||
"litellm[extra_proxy] extra, which now includes psycopg."
|
||||
)
|
||||
return False
|
||||
|
||||
cleaned_url = ProxyExtrasDBManager._strip_prisma_query_params(database_url)
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
[project]
|
||||
name = "litellm-proxy-extras"
|
||||
version = "0.4.91"
|
||||
version = "0.4.92"
|
||||
description = "Additional files for the LiteLLM Proxy. Reduces the size of the main litellm package."
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.9"
|
||||
|
|
@ -26,7 +26,7 @@ required-version = ">=0.10.9"
|
|||
module-root = ""
|
||||
|
||||
[tool.commitizen]
|
||||
version = "0.4.91"
|
||||
version = "0.4.92"
|
||||
version_files = [
|
||||
"pyproject.toml:^version",
|
||||
"../pyproject.toml:litellm-proxy-extras==",
|
||||
|
|
|
|||
|
|
@ -30,3 +30,12 @@ tokio = { version = "1", features = ["rt-multi-thread", "macros", "time", "net"]
|
|||
tokio-tungstenite = { version = "0.24", default-features = false, features = ["connect", "rustls-tls-native-roots"] }
|
||||
futures-util = { version = "0.3", default-features = false, features = ["sink", "std"] }
|
||||
base64 = "0.22"
|
||||
|
||||
[profile.release]
|
||||
opt-level = 3
|
||||
lto = "thin"
|
||||
codegen-units = 1
|
||||
panic = "unwind"
|
||||
debug = false
|
||||
incremental = false
|
||||
strip = "symbols"
|
||||
|
|
|
|||
|
|
@ -10,7 +10,8 @@ name = "_native"
|
|||
crate-type = ["cdylib"]
|
||||
|
||||
[features]
|
||||
default = ["extension-module"]
|
||||
default = ["abi3"]
|
||||
abi3 = ["pyo3/abi3-py310"]
|
||||
extension-module = ["pyo3/extension-module"]
|
||||
|
||||
[dependencies]
|
||||
|
|
|
|||
|
|
@ -7,6 +7,9 @@ warnings.filterwarnings("ignore", message=".*conflict with protected namespace.*
|
|||
# Suppress Pydantic 2.11+ deprecation warning about accessing model_fields on instances
|
||||
# This warning can accumulate during streaming and cause memory leaks
|
||||
warnings.filterwarnings("ignore", message=".*Accessing the.*attribute on the instance is deprecated.*")
|
||||
# ReadOnly on TypedDict fields is repo-wide static discipline (LIT012); pydantic warns it
|
||||
# cannot enforce it at runtime, which floods proxy boot once such a type is schema-walked
|
||||
warnings.filterwarnings("ignore", message=".*`ReadOnly` qualifier.*")
|
||||
### INIT VARIABLES #########################
|
||||
import threading
|
||||
import os
|
||||
|
|
@ -656,6 +659,8 @@ aiml_models: Set = set()
|
|||
deepgram_models: Set = set()
|
||||
elevenlabs_models: Set = set()
|
||||
dashscope_models: Set = set()
|
||||
qwencloud_models: Set = set()
|
||||
qwen_ai_platform_models: Set = set()
|
||||
moonshot_models: Set = set()
|
||||
publicai_models: Set = set()
|
||||
darkbloom_models: Set = set()
|
||||
|
|
@ -906,6 +911,10 @@ def _populate_provider_model_sets(model_cost_map: Dict) -> None:
|
|||
heroku_models.add(key)
|
||||
elif value.get("litellm_provider") == "dashscope":
|
||||
dashscope_models.add(key)
|
||||
elif value.get("litellm_provider") == "qwencloud":
|
||||
qwencloud_models.add(key)
|
||||
elif value.get("litellm_provider") == "qwen_ai_platform":
|
||||
qwen_ai_platform_models.add(key)
|
||||
elif value.get("litellm_provider") == "modelscope":
|
||||
modelscope_models.add(key)
|
||||
elif value.get("litellm_provider") == "moonshot":
|
||||
|
|
@ -1069,6 +1078,8 @@ model_list = list(
|
|||
| deepgram_models
|
||||
| elevenlabs_models
|
||||
| dashscope_models
|
||||
| qwencloud_models
|
||||
| qwen_ai_platform_models
|
||||
| moonshot_models
|
||||
| publicai_models
|
||||
| darkbloom_models
|
||||
|
|
@ -1175,6 +1186,8 @@ def _build_models_by_provider() -> dict:
|
|||
"elevenlabs": elevenlabs_models,
|
||||
"heroku": heroku_models,
|
||||
"dashscope": dashscope_models,
|
||||
"qwencloud": qwencloud_models,
|
||||
"qwen_ai_platform": qwen_ai_platform_models,
|
||||
"modelscope": modelscope_models,
|
||||
"moonshot": moonshot_models,
|
||||
"publicai": publicai_models,
|
||||
|
|
@ -2011,6 +2024,24 @@ if TYPE_CHECKING:
|
|||
from .llms.dashscope.rerank.transformation import (
|
||||
DashScopeRerankConfig as DashScopeRerankConfig,
|
||||
)
|
||||
from .llms.dashscope.qwencloud import (
|
||||
QwenCloudChatConfig as QwenCloudChatConfig,
|
||||
)
|
||||
from .llms.dashscope.qwencloud import (
|
||||
QwenCloudEmbeddingConfig as QwenCloudEmbeddingConfig,
|
||||
)
|
||||
from .llms.dashscope.qwencloud import (
|
||||
QwenCloudRerankConfig as QwenCloudRerankConfig,
|
||||
)
|
||||
from .llms.dashscope.qwen_ai_platform import (
|
||||
QwenAIPlatformChatConfig as QwenAIPlatformChatConfig,
|
||||
)
|
||||
from .llms.dashscope.qwen_ai_platform import (
|
||||
QwenAIPlatformEmbeddingConfig as QwenAIPlatformEmbeddingConfig,
|
||||
)
|
||||
from .llms.dashscope.qwen_ai_platform import (
|
||||
QwenAIPlatformRerankConfig as QwenAIPlatformRerankConfig,
|
||||
)
|
||||
from .llms.modelscope.chat.transformation import (
|
||||
ModelScopeChatConfig as ModelScopeChatConfig,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -17,7 +17,7 @@ until they're actually needed.
|
|||
|
||||
import importlib
|
||||
import sys
|
||||
from collections.abc import Callable
|
||||
from collections.abc import Callable, Mapping
|
||||
from types import ModuleType
|
||||
from typing import TYPE_CHECKING, Any, Final, cast
|
||||
|
||||
|
|
@ -57,10 +57,11 @@ from ._lazy_imports_registry import (
|
|||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
import httpx
|
||||
from tiktoken import Encoding
|
||||
|
||||
|
||||
def get_litellm_globals() -> dict:
|
||||
def get_litellm_globals() -> dict[str, object]:
|
||||
"""
|
||||
Get the globals dictionary of the litellm module.
|
||||
|
||||
|
|
@ -70,7 +71,7 @@ def get_litellm_globals() -> dict:
|
|||
return sys.modules["litellm"].__dict__
|
||||
|
||||
|
||||
def _get_utils_globals() -> dict:
|
||||
def _get_utils_globals() -> dict[str, object]:
|
||||
"""
|
||||
Get the globals dictionary of the utils module.
|
||||
|
||||
|
|
@ -80,6 +81,11 @@ def _get_utils_globals() -> dict:
|
|||
return sys.modules["litellm.utils"].__dict__
|
||||
|
||||
|
||||
def _get_module_level_client_timeout(litellm_globals: Mapping[str, Any]) -> "float | httpx.Timeout | None":
|
||||
"""Read the configured `litellm.request_timeout` used for the module level http clients."""
|
||||
return litellm_globals.get("request_timeout")
|
||||
|
||||
|
||||
# These are special lazy loaders for things that are used internally
|
||||
# They're separate from the main lazy import system because they have specific use cases
|
||||
|
||||
|
|
@ -435,8 +441,8 @@ def _lazy_import_http_handlers(name: str) -> object:
|
|||
from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
|
||||
|
||||
# Get timeout from module config (if set)
|
||||
timeout = _globals.get("request_timeout")
|
||||
params: Final = {"timeout": timeout, "client_alias": "module level aclient"}
|
||||
async_timeout: Final = _get_module_level_client_timeout(_globals)
|
||||
params: Final = {"timeout": async_timeout, "client_alias": "module level aclient"}
|
||||
|
||||
# Create the client instance
|
||||
provider_id: Final = cast(Any, "litellm_module_level_client")
|
||||
|
|
@ -453,8 +459,8 @@ def _lazy_import_http_handlers(name: str) -> object:
|
|||
# Create a sync HTTP client
|
||||
from litellm.llms.custom_httpx.http_handler import HTTPHandler
|
||||
|
||||
timeout = _globals.get("request_timeout")
|
||||
sync_client: Final = HTTPHandler(timeout=timeout)
|
||||
sync_timeout: Final = _get_module_level_client_timeout(_globals)
|
||||
sync_client: Final = HTTPHandler(timeout=sync_timeout)
|
||||
|
||||
# Cache it
|
||||
_globals["module_level_client"] = sync_client
|
||||
|
|
|
|||
|
|
@ -310,6 +310,8 @@ LLM_CONFIG_NAMES: Final = (
|
|||
"GigaChatConfig",
|
||||
"GigaChatEmbeddingConfig",
|
||||
"DashScopeChatConfig",
|
||||
"QwenCloudChatConfig",
|
||||
"QwenAIPlatformChatConfig",
|
||||
"ModelScopeChatConfig",
|
||||
"MoonshotChatConfig",
|
||||
"DockerModelRunnerChatConfig",
|
||||
|
|
@ -1172,6 +1174,14 @@ _LLM_CONFIGS_IMPORT_MAP: Final = {
|
|||
".llms.dashscope.chat.transformation",
|
||||
"DashScopeChatConfig",
|
||||
),
|
||||
"QwenCloudChatConfig": (
|
||||
".llms.dashscope.qwencloud",
|
||||
"QwenCloudChatConfig",
|
||||
),
|
||||
"QwenAIPlatformChatConfig": (
|
||||
".llms.dashscope.qwen_ai_platform",
|
||||
"QwenAIPlatformChatConfig",
|
||||
),
|
||||
"GDCGeminiConfig": (
|
||||
".llms.gdc.chat.transformation",
|
||||
"GDCGeminiConfig",
|
||||
|
|
|
|||
|
|
@ -264,13 +264,17 @@ def _plain_log_format(stdout: TextIO | None, stderr: TextIO | None) -> str:
|
|||
|
||||
|
||||
class LevelRoutingStreamHandler(logging.StreamHandler):
|
||||
"""Writes records below WARNING to stdout and WARNING and above to stderr.
|
||||
"""Writes records below WARNING and invalid-key warnings to stdout, others to stderr.
|
||||
|
||||
Collectors that derive severity from the stream report every stderr line as an error.
|
||||
Invalid-key warnings route to stdout so LITELLM_LOG=ERROR can suppress them.
|
||||
"""
|
||||
|
||||
def emit(self, record: logging.LogRecord) -> None:
|
||||
preferred: Final = sys.stdout if record.levelno < logging.WARNING else sys.stderr
|
||||
is_stdout_record: Final = record.levelno < logging.WARNING or (
|
||||
record.levelno == logging.WARNING and record.name == verbose_proxy_stdout_logger.name
|
||||
)
|
||||
preferred: Final = sys.stdout if is_stdout_record else sys.stderr
|
||||
if preferred is None or getattr(preferred, "closed", False):
|
||||
self.stream = sys.stderr # rebind-ok: fall back to the pre-fix stream rather than raising per record
|
||||
else:
|
||||
|
|
@ -508,6 +512,9 @@ else:
|
|||
handler.setFormatter(formatter)
|
||||
|
||||
verbose_proxy_logger = logging.getLogger("LiteLLM Proxy")
|
||||
# Malformed virtual key rejections log through this child; LevelRoutingStreamHandler
|
||||
# writes its WARNING records to stdout. It has no handler or level of its own.
|
||||
verbose_proxy_stdout_logger: Final = verbose_proxy_logger.getChild("stdout")
|
||||
verbose_router_logger = logging.getLogger("LiteLLM Router")
|
||||
verbose_logger = logging.getLogger("LiteLLM")
|
||||
|
||||
|
|
@ -520,6 +527,7 @@ verbose_logger.addHandler(handler)
|
|||
# handlers (JSON mode, uvicorn log config, a host app's root handler).
|
||||
verbose_router_logger.addFilter(_stdout_truncation_filter)
|
||||
verbose_proxy_logger.addFilter(_stdout_truncation_filter)
|
||||
verbose_proxy_stdout_logger.addFilter(_stdout_truncation_filter)
|
||||
verbose_logger.addFilter(_stdout_truncation_filter)
|
||||
|
||||
|
||||
|
|
@ -683,6 +691,7 @@ def _turn_on_json():
|
|||
- Adds a JSON formatter to all loggers
|
||||
"""
|
||||
handler: Final = LevelRoutingStreamHandler()
|
||||
handler.setLevel(numeric_level)
|
||||
handler.setFormatter(JsonFormatter())
|
||||
_initialize_loggers_with_handler(handler)
|
||||
# Set up exception handlers
|
||||
|
|
@ -700,12 +709,14 @@ def _disable_debugging():
|
|||
verbose_logger.disabled = True
|
||||
verbose_router_logger.disabled = True
|
||||
verbose_proxy_logger.disabled = True
|
||||
verbose_proxy_stdout_logger.disabled = True
|
||||
|
||||
|
||||
def _enable_debugging():
|
||||
verbose_logger.disabled = False
|
||||
verbose_router_logger.disabled = False
|
||||
verbose_proxy_logger.disabled = False
|
||||
verbose_proxy_stdout_logger.disabled = False
|
||||
|
||||
|
||||
def print_verbose(print_statement):
|
||||
|
|
|
|||
|
|
@ -13,6 +13,7 @@ import json
|
|||
# s/o [@Frank Colson](https://www.linkedin.com/in/frank-colson-422b9b183/) for this redis implementation
|
||||
import os
|
||||
from collections.abc import Callable, Mapping
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
from urllib.parse import urlsplit, urlunsplit
|
||||
|
||||
|
|
@ -38,9 +39,25 @@ from ._logging import verbose_logger
|
|||
AZURE_REDIS_SCOPE: Final = "https://redis.azure.com/.default"
|
||||
|
||||
|
||||
def _get_redis_kwargs():
|
||||
arg_spec: Final = inspect.getfullargspec(redis.Redis)
|
||||
def _unwrapped_init_args(cls: type) -> frozenset[str]:
|
||||
"""Every parameter on a single class's own ``__init__``, decorator-unwrapped.
|
||||
|
||||
Unlike ``_init_arg_names`` below, this does not walk the MRO: ``redis.Redis``
|
||||
and ``redis.RedisCluster`` (sync and async) each declare every real
|
||||
constructor parameter directly on their own ``__init__``, so MRO-walking is
|
||||
unnecessary — and it actively breaks the several tests here that mock the
|
||||
class with ``patch(..., autospec=True)``, since ``inspect.getmro`` needs a
|
||||
real ``__mro__`` that an autospec'd stand-in for a class does not provide.
|
||||
|
||||
Still unwraps first: redis-py >= 7.4 decorates these ``__init__``s with
|
||||
``@deprecated_args`` too, which the same class of bug as ``_init_arg_names``
|
||||
would otherwise silently empty this allowlist through (see its docstring).
|
||||
"""
|
||||
spec: Final = inspect.getfullargspec(inspect.unwrap(cls.__init__))
|
||||
return frozenset(spec.args + spec.kwonlyargs)
|
||||
|
||||
|
||||
def _get_redis_kwargs():
|
||||
# Only allow primitive arguments
|
||||
exclude_args: Final = {
|
||||
"self",
|
||||
|
|
@ -60,7 +77,7 @@ def _get_redis_kwargs():
|
|||
"azure_client_secret",
|
||||
}
|
||||
|
||||
available_args: Final = {x for x in arg_spec.args if x not in exclude_args} | include_args
|
||||
available_args: Final = {x for x in _unwrapped_init_args(redis.Redis) if x not in exclude_args} | include_args
|
||||
|
||||
return available_args
|
||||
|
||||
|
|
@ -120,15 +137,23 @@ def _get_redis_url_kwargs(client: type | None = None) -> tuple[str, ...]:
|
|||
return tuple(x for x in _init_arg_names(connection_cls) if x not in exclude_args) + include_args
|
||||
|
||||
|
||||
def _get_redis_cluster_kwargs(client=None):
|
||||
def _get_redis_cluster_kwargs(client: type | None = None):
|
||||
"""Config kwargs the target cluster client's constructor actually accepts.
|
||||
|
||||
Defaults to the sync ``redis.RedisCluster``, but the async cluster client
|
||||
(``redis.asyncio.cluster.RedisCluster``) declares connection settings such as
|
||||
``decode_responses`` on its own constructor, where the sync class takes them
|
||||
through ``**kwargs`` and so never names them in its signature. Introspecting
|
||||
only the sync class regardless of which client is actually built silently
|
||||
drops those for every async cluster caller.
|
||||
"""
|
||||
if client is None:
|
||||
client = redis.Redis.from_url
|
||||
arg_spec: Final = inspect.getfullargspec(redis.RedisCluster)
|
||||
client = redis.RedisCluster
|
||||
|
||||
# Only allow primitive arguments
|
||||
exclude_args: Final = {"self", "connection_pool", "retry", "host", "port", "startup_nodes"}
|
||||
|
||||
available_args = {x for x in arg_spec.args if x not in exclude_args}
|
||||
available_args = {x for x in _unwrapped_init_args(client) if x not in exclude_args}
|
||||
available_args |= {
|
||||
"password",
|
||||
"username",
|
||||
|
|
@ -161,6 +186,79 @@ def _get_redis_env_kwarg_mapping():
|
|||
return {f"{PREFIX}{x.upper()}": x for x in _get_redis_kwargs() if x not in exclude_from_environment}
|
||||
|
||||
|
||||
def _str_to_bool(value: str) -> bool:
|
||||
return value.lower() in ("true", "1", "yes")
|
||||
|
||||
|
||||
def _coerce_redis_kwargs_types(
|
||||
redis_kwargs: Mapping[str, object],
|
||||
client: type | tuple[type, ...] = redis.Redis,
|
||||
) -> dict[str, object]: # mutable-ok: a caller mutates the returned kwargs before constructing its client
|
||||
"""Coerces string values to the numeric/boolean type ``client``'s constructor
|
||||
declares for that parameter. ``client`` may be a tuple of client classes; a
|
||||
parameter's type is taken from the first signature that declares it, which
|
||||
lets cluster callers coerce cluster-only kwargs such as
|
||||
``cluster_error_retry_attempts`` alongside the shared connection kwargs.
|
||||
|
||||
Environment variables are always strings, and Helm ``--set`` stringifies values
|
||||
too, so a config value like ``health_check_interval`` or ``socket_timeout``
|
||||
can arrive as ``"30"``/``"5.5"`` rather than a real number. redis-py's own
|
||||
connection-health-check arithmetic (``loop.time() + self.health_check_interval``)
|
||||
then raises ``TypeError`` on every Redis operation instead of connecting.
|
||||
|
||||
``max_connections``, ``socket_timeout``, and ``socket_connect_timeout`` use an
|
||||
explicit target type rather than the parameter's own signature default: redis-py
|
||||
8.x changed the timeout defaults from ``None`` to int ``5``, so inferring the
|
||||
type from the default would make a fractional ``"5.5"`` fail ``int()`` and get
|
||||
silently dropped on 8.x while working on older versions. ``socket_keepalive``
|
||||
is explicit too: its signature default is ``None``, which carries no type to
|
||||
infer from, and leaving it a string makes ``"false"`` truthy.
|
||||
"""
|
||||
signatures: Final = tuple(inspect.signature(c) for c in (client if isinstance(client, tuple) else (client,)))
|
||||
explicit_param_types: Final = MappingProxyType(
|
||||
{
|
||||
"max_connections": int,
|
||||
"socket_timeout": float,
|
||||
"socket_connect_timeout": float,
|
||||
"socket_keepalive": bool,
|
||||
}
|
||||
)
|
||||
result: Final = dict(redis_kwargs) # mutable-ok: per-key try/except coercion below needs to drop individual keys
|
||||
for key, value in redis_kwargs.items():
|
||||
if not isinstance(value, str):
|
||||
continue
|
||||
param = next((sig.parameters[key] for sig in signatures if key in sig.parameters), None)
|
||||
if param is None:
|
||||
continue
|
||||
explicit_type = explicit_param_types.get(key)
|
||||
if explicit_type is bool:
|
||||
result[key] = _str_to_bool(value)
|
||||
continue
|
||||
if explicit_type is not None:
|
||||
try:
|
||||
result[key] = explicit_type(value)
|
||||
except (ValueError, TypeError):
|
||||
del result[key]
|
||||
continue
|
||||
default: object = param.default # pyright: ignore[reportAny] # inspect.Parameter.default is stubbed as Any
|
||||
if default is inspect.Parameter.empty:
|
||||
continue
|
||||
# bool must be checked before int, since bool subclasses int
|
||||
if isinstance(default, bool):
|
||||
result[key] = _str_to_bool(value)
|
||||
elif isinstance(default, int):
|
||||
try:
|
||||
result[key] = int(value)
|
||||
except (ValueError, TypeError):
|
||||
del result[key]
|
||||
elif isinstance(default, float):
|
||||
try:
|
||||
result[key] = float(value)
|
||||
except (ValueError, TypeError):
|
||||
del result[key]
|
||||
return result
|
||||
|
||||
|
||||
def _redis_kwargs_from_environment():
|
||||
mapping: Final = _get_redis_env_kwarg_mapping()
|
||||
|
||||
|
|
@ -505,7 +603,12 @@ def _get_redis_client_logic(**env_overrides):
|
|||
raise ValueError("Either 'host' or 'url' must be specified for redis.")
|
||||
|
||||
# litellm.print_verbose(f"redis_kwargs: {redis_kwargs}")
|
||||
return redis_kwargs
|
||||
coercion_client: Final = (
|
||||
(redis.Redis, redis.RedisCluster, async_redis.RedisCluster)
|
||||
if redis_kwargs.get("startup_nodes")
|
||||
else redis.Redis
|
||||
)
|
||||
return _coerce_redis_kwargs_types(redis_kwargs, client=coercion_client)
|
||||
|
||||
|
||||
def init_redis_cluster(redis_kwargs) -> redis.RedisCluster:
|
||||
|
|
@ -657,7 +760,9 @@ def get_redis_client(**env_overrides):
|
|||
if "sentinel_nodes" in redis_kwargs and "service_name" in redis_kwargs:
|
||||
return _init_redis_sentinel(redis_kwargs)
|
||||
|
||||
return redis.Redis(**redis_kwargs)
|
||||
return redis.Redis( # pyright: ignore[reportCallIssue] # object-valued kwargs match no overload statically
|
||||
**redis_kwargs, # pyright: ignore[reportArgumentType] # allow-listed and coerced against this signature
|
||||
)
|
||||
|
||||
|
||||
def get_redis_async_client(
|
||||
|
|
@ -669,7 +774,7 @@ def get_redis_async_client(
|
|||
if "startup_nodes" in redis_kwargs:
|
||||
from redis.cluster import ClusterNode
|
||||
|
||||
args = _get_redis_cluster_kwargs()
|
||||
args = _get_redis_cluster_kwargs(async_redis.RedisCluster)
|
||||
cluster_kwargs: Final = {}
|
||||
for arg in redis_kwargs:
|
||||
if arg in args:
|
||||
|
|
|
|||
|
|
@ -12,6 +12,7 @@ import hashlib
|
|||
import json
|
||||
import time
|
||||
import traceback
|
||||
from collections.abc import Mapping
|
||||
from enum import Enum
|
||||
from typing import Any, Final
|
||||
|
||||
|
|
@ -506,7 +507,7 @@ class Cache:
|
|||
|
||||
def _get_cache_logic(
|
||||
self,
|
||||
cached_result: Any | None,
|
||||
cached_result: object | None,
|
||||
max_age: float | None,
|
||||
):
|
||||
"""
|
||||
|
|
@ -538,8 +539,8 @@ class Cache:
|
|||
return cached_result
|
||||
|
||||
@staticmethod
|
||||
def _get_safe_cache_lookup_kwargs(kwargs: dict[str, Any]) -> dict[str, Any]:
|
||||
cache_lookup_kwargs: Final[dict[str, Any]] = {}
|
||||
def _get_safe_cache_lookup_kwargs(kwargs: Mapping[str, object]) -> dict[str, object]:
|
||||
cache_lookup_kwargs: Final[dict[str, object]] = {}
|
||||
for prompt_kwarg in ("messages", "input"):
|
||||
if prompt_kwarg in kwargs:
|
||||
cache_lookup_kwargs[prompt_kwarg] = kwargs[prompt_kwarg]
|
||||
|
|
@ -552,7 +553,7 @@ class Cache:
|
|||
|
||||
@staticmethod
|
||||
def _update_metadata_from_cache_lookup_kwargs(
|
||||
original_kwargs: dict[str, Any], cache_lookup_kwargs: dict[str, Any]
|
||||
original_kwargs: Mapping[str, object], cache_lookup_kwargs: Mapping[str, object]
|
||||
) -> None:
|
||||
original_metadata: Final = original_kwargs.get("metadata")
|
||||
cache_lookup_metadata: Final = cache_lookup_kwargs.get("metadata")
|
||||
|
|
|
|||
|
|
@ -12,7 +12,7 @@ import ast
|
|||
import asyncio
|
||||
import json
|
||||
import os
|
||||
from typing import TYPE_CHECKING, Any, Final, cast
|
||||
from typing import TYPE_CHECKING, Any, Final, Protocol, cast
|
||||
|
||||
import litellm
|
||||
from litellm._logging import print_verbose
|
||||
|
|
@ -39,6 +39,12 @@ if TYPE_CHECKING:
|
|||
from litellm.router import Router
|
||||
|
||||
|
||||
class _QdrantCollectionDetailsResponse(Protocol):
|
||||
"""The qdrant `/collections/{name}` response, whose body is kept as an opaque JSON object."""
|
||||
|
||||
def json(self) -> dict[str, object]: ...
|
||||
|
||||
|
||||
class QdrantSemanticCache(BaseCache):
|
||||
CACHE_KEY_FIELD_NAME = "litellm_cache_key"
|
||||
embedding_max_input_tokens: int | None = None
|
||||
|
|
@ -115,15 +121,15 @@ class QdrantSemanticCache(BaseCache):
|
|||
raise ValueError(f"Error from qdrant checking if /collections exist {collection_exists.text}")
|
||||
|
||||
if collection_exists.json()["result"]["exists"]:
|
||||
collection_details = self.sync_client.get(
|
||||
collection_details: _QdrantCollectionDetailsResponse = self.sync_client.get(
|
||||
url=f"{self.qdrant_api_base}/collections/{self.collection_name}",
|
||||
headers=self.headers,
|
||||
)
|
||||
self.collection_info = collection_details.json()
|
||||
self.collection_info: dict[str, object] = collection_details.json()
|
||||
print_verbose(f"Collection already exists.\nCollection details:{self.collection_info}")
|
||||
self._ensure_cache_key_payload_index()
|
||||
else:
|
||||
quantization_params: dict[str, Any]
|
||||
quantization_params: dict[str, dict[str, object]]
|
||||
if quantization_config is None or quantization_config == "binary":
|
||||
quantization_params = {
|
||||
"binary": {
|
||||
|
|
@ -214,7 +220,7 @@ class QdrantSemanticCache(BaseCache):
|
|||
resolve_embedding_max_input_tokens(self.embedding_max_input_tokens, self.embedding_model, router),
|
||||
)
|
||||
|
||||
def _get_embedding(self, prompt: str, metadata: dict[str, Any] | None = None) -> EmbeddingResponse:
|
||||
def _get_embedding(self, prompt: str, metadata: dict[str, object] | None = None) -> EmbeddingResponse:
|
||||
"""Embed via the proxy Router when it serves the model, else direct."""
|
||||
try:
|
||||
from litellm.proxy.proxy_server import llm_model_list, llm_router
|
||||
|
|
@ -241,7 +247,7 @@ class QdrantSemanticCache(BaseCache):
|
|||
num_retries=0,
|
||||
)
|
||||
|
||||
async def _get_async_embedding(self, prompt: str, metadata: dict[str, Any] | None = None) -> EmbeddingResponse:
|
||||
async def _get_async_embedding(self, prompt: str, metadata: dict[str, object] | None = None) -> EmbeddingResponse:
|
||||
try:
|
||||
from litellm.proxy.proxy_server import llm_model_list, llm_router
|
||||
except ImportError:
|
||||
|
|
|
|||
|
|
@ -18,7 +18,7 @@ import time
|
|||
from collections.abc import Awaitable, Callable, Sequence
|
||||
from contextvars import ContextVar
|
||||
from datetime import timedelta
|
||||
from typing import TYPE_CHECKING, Any, Final, TypeVar, cast
|
||||
from typing import TYPE_CHECKING, Any, Final, Protocol, TypeVar, cast
|
||||
|
||||
import litellm
|
||||
from litellm._logging import print_verbose, verbose_logger
|
||||
|
|
@ -58,6 +58,26 @@ else:
|
|||
Span = Any
|
||||
|
||||
|
||||
class _AsyncRedisCommands(Protocol):
|
||||
"""Async redis commands this cache issues.
|
||||
|
||||
redis-py's type stubs omit these methods on RedisCluster, so the union returned by
|
||||
init_async_client() is untyped at every call site without this protocol.
|
||||
"""
|
||||
|
||||
def ping(self) -> Awaitable[bool]: ...
|
||||
|
||||
def delete(self, *names: str) -> Awaitable[int]: ...
|
||||
|
||||
def ttl(self, name: str) -> Awaitable[int]: ...
|
||||
|
||||
def rpush(self, name: str, *values: str | bytes | float) -> Awaitable[int]: ...
|
||||
|
||||
def lpop(self, name: str, count: int | None = None) -> Awaitable[object]: ...
|
||||
|
||||
def pipeline(self, transaction: bool = True) -> "Pipeline[bytes]": ...
|
||||
|
||||
|
||||
def _get_call_stack_info(num_frames: int = 2) -> str:
|
||||
"""
|
||||
Get the function names from the previous 1-2 functions in the call stack.
|
||||
|
|
@ -429,6 +449,9 @@ class RedisCache(BaseCache):
|
|||
self.redis_async_client = redis_async_client
|
||||
return redis_async_client
|
||||
|
||||
def _async_commands(self) -> _AsyncRedisCommands:
|
||||
return self.init_async_client()
|
||||
|
||||
def check_and_fix_namespace(self, key: str) -> str:
|
||||
"""
|
||||
Make sure each key starts with the given namespace
|
||||
|
|
@ -1055,19 +1078,17 @@ class RedisCache(BaseCache):
|
|||
await self.async_set_cache_pipeline(self.redis_batch_writing_buffer)
|
||||
self.redis_batch_writing_buffer = []
|
||||
|
||||
def _get_cache_logic(self, cached_response: Any):
|
||||
def _get_cache_logic(self, cached_response: bytes | str | None):
|
||||
"""
|
||||
Common 'get_cache_logic' across sync + async redis client implementations
|
||||
"""
|
||||
if cached_response is None:
|
||||
return cached_response
|
||||
# cached_response is in `b{} convert it to ModelResponse
|
||||
cached_response = cached_response.decode("utf-8") # Convert bytes to string
|
||||
return None
|
||||
decoded: Final = cached_response.decode("utf-8") if isinstance(cached_response, bytes) else cached_response
|
||||
try:
|
||||
cached_response = json.loads(cached_response) # Convert string to dictionary
|
||||
return json.loads(decoded)
|
||||
except Exception:
|
||||
cached_response = ast.literal_eval(cached_response)
|
||||
return cached_response
|
||||
return ast.literal_eval(decoded)
|
||||
|
||||
def get_cache(self, key, parent_otel_span: Span | None = None, **kwargs):
|
||||
try:
|
||||
|
|
@ -1314,8 +1335,7 @@ class RedisCache(BaseCache):
|
|||
raise e
|
||||
|
||||
async def ping(self) -> bool:
|
||||
# typed as Any, redis python lib has incomplete type stubs for RedisCluster and does not include `ping`
|
||||
_redis_client: Final[Any] = self.init_async_client()
|
||||
_redis_client: Final = self._async_commands()
|
||||
start_time: Final = time.time()
|
||||
print_verbose("Pinging Async Redis Cache")
|
||||
try:
|
||||
|
|
@ -1349,8 +1369,7 @@ class RedisCache(BaseCache):
|
|||
|
||||
@_redis_circuit_breaker_guard
|
||||
async def delete_cache_keys(self, keys):
|
||||
# typed as Any, redis python lib has incomplete type stubs for RedisCluster and does not include `delete`
|
||||
_redis_client: Final[Any] = self.init_async_client()
|
||||
_redis_client: Final = self._async_commands()
|
||||
keys = [self.check_and_fix_namespace(key=key) for key in keys]
|
||||
# keys is a list, unpack it so it gets passed as individual elements to delete
|
||||
await _redis_client.delete(*keys)
|
||||
|
|
@ -1415,8 +1434,7 @@ class RedisCache(BaseCache):
|
|||
|
||||
@_redis_circuit_breaker_guard
|
||||
async def async_delete_cache(self, key: str):
|
||||
# typed as Any, redis python lib has incomplete type stubs for RedisCluster and does not include `delete`
|
||||
_redis_client: Final[Any] = self.init_async_client()
|
||||
_redis_client: Final = self._async_commands()
|
||||
key = self.check_and_fix_namespace(key=key)
|
||||
# keys is str
|
||||
return await _redis_client.delete(key)
|
||||
|
|
@ -1523,8 +1541,7 @@ class RedisCache(BaseCache):
|
|||
Redis ref: https://redis.io/docs/latest/commands/ttl/
|
||||
"""
|
||||
try:
|
||||
# typed as Any, redis python lib has incomplete type stubs for RedisCluster and does not include `ttl`
|
||||
_redis_client: Final[Any] = self.init_async_client()
|
||||
_redis_client: Final = self._async_commands()
|
||||
key = self.check_and_fix_namespace(key=key)
|
||||
ttl: Final = await _redis_client.ttl(key)
|
||||
if ttl <= -1: # -1 means the key does not exist, -2 key does not exist
|
||||
|
|
@ -1554,7 +1571,7 @@ class RedisCache(BaseCache):
|
|||
Returns:
|
||||
int: The length of the list after the push operation
|
||||
"""
|
||||
_redis_client: Final[Any] = self.init_async_client()
|
||||
_redis_client: Final = self._async_commands()
|
||||
key = self.check_and_fix_namespace(key=key)
|
||||
start_time: Final = time.time()
|
||||
try:
|
||||
|
|
@ -1621,7 +1638,7 @@ class RedisCache(BaseCache):
|
|||
if len(rpush_list) == 0:
|
||||
return []
|
||||
|
||||
_redis_client: Final[Any] = self.init_async_client()
|
||||
_redis_client: Final = self._async_commands()
|
||||
start_time: Final = time.time()
|
||||
|
||||
try:
|
||||
|
|
@ -1678,7 +1695,7 @@ class RedisCache(BaseCache):
|
|||
parent_otel_span: Span | None = None,
|
||||
**kwargs,
|
||||
) -> Any | list[Any]:
|
||||
_redis_client: Final[Any] = self.init_async_client()
|
||||
_redis_client: Final = self._async_commands()
|
||||
key = self.check_and_fix_namespace(key=key)
|
||||
start_time: Final = time.time()
|
||||
print_verbose(f"LPOP from Redis list: key: {key}, count: {count}")
|
||||
|
|
@ -1810,7 +1827,7 @@ class RedisCache(BaseCache):
|
|||
if len(lpop_list) == 0:
|
||||
return []
|
||||
|
||||
_redis_client: Final[Any] = self.init_async_client()
|
||||
_redis_client: Final = self._async_commands()
|
||||
start_time: Final = time.time()
|
||||
|
||||
try:
|
||||
|
|
|
|||
|
|
@ -17,6 +17,7 @@ RedisSemanticCache since those are backend agnostic.
|
|||
import asyncio
|
||||
import hashlib
|
||||
import os
|
||||
from collections.abc import Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Final
|
||||
|
||||
|
|
@ -64,7 +65,7 @@ class ValkeySemanticCache(RedisSemanticCache):
|
|||
async_client: AsyncRedis | None = None,
|
||||
embedding_max_input_tokens: int | None = None,
|
||||
embedding_timeout: float | None = None,
|
||||
**kwargs: Any,
|
||||
**kwargs: object,
|
||||
):
|
||||
if similarity_threshold is None:
|
||||
raise ValueError("similarity_threshold must be provided, passed None")
|
||||
|
|
@ -87,11 +88,13 @@ class ValkeySemanticCache(RedisSemanticCache):
|
|||
self.key_prefix = f"{self.index_name}:"
|
||||
self._index_dim: int | None = None
|
||||
|
||||
resolved_url = None
|
||||
if sync_client is None or async_client is None:
|
||||
resolved_url = redis_url or self._build_valkey_url(host, port, password, ssl)
|
||||
self.sync_client = sync_client if sync_client is not None else Redis.from_url(resolved_url)
|
||||
self.async_client = async_client if async_client is not None else AsyncRedis.from_url(resolved_url)
|
||||
if sync_client is not None and async_client is not None:
|
||||
self.sync_client = sync_client
|
||||
self.async_client = async_client
|
||||
else:
|
||||
resolved_url: Final = redis_url or self._build_valkey_url(host, port, password, ssl)
|
||||
self.sync_client = sync_client if sync_client is not None else Redis.from_url(resolved_url)
|
||||
self.async_client = async_client if async_client is not None else AsyncRedis.from_url(resolved_url)
|
||||
|
||||
print_verbose(f"Valkey semantic-cache initializing index - {self.index_name}")
|
||||
|
||||
|
|
@ -118,7 +121,7 @@ class ValkeySemanticCache(RedisSemanticCache):
|
|||
return hashlib.sha256(str(key).encode("utf-8")).hexdigest()
|
||||
|
||||
@staticmethod
|
||||
def _embedding_to_bytes(embedding: list[float]) -> bytes:
|
||||
def _embedding_to_bytes(embedding: Sequence[float]) -> bytes:
|
||||
return pack_vector(embedding)
|
||||
|
||||
def _index_schema(self, dim: int) -> tuple[TagField, VectorField]:
|
||||
|
|
@ -192,7 +195,9 @@ class ValkeySemanticCache(RedisSemanticCache):
|
|||
def _doc_key(self, key: str) -> str:
|
||||
return f"{self.key_prefix}{self._scope_tag(key)}:{uuid.uuid4()}"
|
||||
|
||||
def _doc_mapping(self, key: str, prompt: str, value_str: str, embedding: list[float]) -> dict:
|
||||
def _doc_mapping(
|
||||
self, key: str, prompt: str, value_str: str, embedding: Sequence[float]
|
||||
) -> Mapping[str | bytes, str | bytes]:
|
||||
return {
|
||||
self.CACHE_KEY_FIELD_NAME: self._scope_tag(key),
|
||||
self.PROMPT_FIELD_NAME: prompt,
|
||||
|
|
@ -208,30 +213,49 @@ class ValkeySemanticCache(RedisSemanticCache):
|
|||
)
|
||||
return Query(query_string).return_fields(self.RESPONSE_FIELD_NAME, self.DISTANCE_FIELD_NAME).dialect(2)
|
||||
|
||||
async def _async_search(self, key: str, embedding: Sequence[float]) -> object:
|
||||
"""Run the KNN query on the async client, stopping the untyped search surface here."""
|
||||
return await self.async_client.ft(self.index_name).search(
|
||||
self._knn_query(key),
|
||||
query_params={"vec": self._embedding_to_bytes(embedding)}, # pyright: ignore[reportArgumentType] # redis stubs omit bytes; KNN vectors are raw bytes at runtime
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def _first_hit(cls, search_result: Any) -> _ValkeyCacheHit | None:
|
||||
docs: Final = getattr(search_result, "docs", [])
|
||||
def _first_hit(cls, search_result: object) -> _ValkeyCacheHit | None:
|
||||
docs: Final[Sequence[object]] = getattr(search_result, "docs", [])
|
||||
if not docs:
|
||||
return None
|
||||
doc: Final = docs[0]
|
||||
response_field: Final[object] = getattr(doc, cls.RESPONSE_FIELD_NAME)
|
||||
distance_field: Final[str | bytes | float] = getattr(doc, cls.DISTANCE_FIELD_NAME)
|
||||
return _ValkeyCacheHit(
|
||||
response=str(getattr(doc, cls.RESPONSE_FIELD_NAME)),
|
||||
distance=float(getattr(doc, cls.DISTANCE_FIELD_NAME)),
|
||||
response=str(response_field),
|
||||
distance=float(distance_field),
|
||||
)
|
||||
|
||||
def _resolve_hit(self, hit: _ValkeyCacheHit | None, key: str, **kwargs: Any) -> Any:
|
||||
@staticmethod
|
||||
def _record_similarity(kwargs: dict[str, Any], similarity: float) -> None:
|
||||
"""Stamp the semantic-similarity score onto the request metadata carried in ``kwargs``."""
|
||||
kwargs.setdefault("metadata", {})["semantic-similarity"] = similarity
|
||||
|
||||
@staticmethod
|
||||
def _embedding_metadata(kwargs: dict[str, Any]) -> dict[str, Any] | None:
|
||||
"""The request metadata forwarded to the embedding call."""
|
||||
return kwargs.get("metadata")
|
||||
|
||||
def _resolve_hit(self, hit: _ValkeyCacheHit | None, key: str, **kwargs: object) -> object:
|
||||
if hit is None:
|
||||
kwargs.setdefault("metadata", {})["semantic-similarity"] = 0.0
|
||||
self._record_similarity(kwargs, 0.0)
|
||||
return None
|
||||
|
||||
similarity: Final = 1 - hit.distance
|
||||
kwargs.setdefault("metadata", {})["semantic-similarity"] = similarity
|
||||
self._record_similarity(kwargs, similarity)
|
||||
|
||||
if similarity < self.similarity_threshold:
|
||||
return None
|
||||
return self._get_cache_logic(cached_response=hit.response)
|
||||
|
||||
def set_cache(self, key: str, value: Any, **kwargs: Any) -> None:
|
||||
def set_cache(self, key: str, value: object, **kwargs: object) -> None:
|
||||
print_verbose(f"Valkey semantic-cache set_cache, kwargs: {kwargs}")
|
||||
try:
|
||||
prompt: Final = self._get_prompt_from_kwargs(**kwargs)
|
||||
|
|
@ -250,12 +274,12 @@ class ValkeySemanticCache(RedisSemanticCache):
|
|||
except Exception as e:
|
||||
print_verbose(f"Error in Valkey semantic-cache set_cache: {e}")
|
||||
|
||||
def get_cache(self, key: str, **kwargs: Any) -> Any:
|
||||
def get_cache(self, key: str, **kwargs: object) -> object:
|
||||
print_verbose(f"Valkey semantic-cache get_cache, kwargs: {kwargs}")
|
||||
try:
|
||||
prompt: Final = self._get_prompt_from_kwargs(**kwargs)
|
||||
if prompt is None:
|
||||
kwargs.setdefault("metadata", {})["semantic-similarity"] = 0.0
|
||||
self._record_similarity(kwargs, 0.0)
|
||||
return None
|
||||
|
||||
embedding: Final = self._get_embedding(prompt)
|
||||
|
|
@ -263,14 +287,14 @@ class ValkeySemanticCache(RedisSemanticCache):
|
|||
|
||||
search_result: Final = self.sync_client.ft(self.index_name).search(
|
||||
self._knn_query(key),
|
||||
query_params={"vec": self._embedding_to_bytes(embedding)},
|
||||
query_params={"vec": self._embedding_to_bytes(embedding)}, # pyright: ignore[reportArgumentType] # redis stubs omit bytes; KNN vectors are raw bytes at runtime
|
||||
)
|
||||
return self._resolve_hit(self._first_hit(search_result), key, **kwargs)
|
||||
except Exception as e:
|
||||
print_verbose(f"Error in Valkey semantic-cache get_cache: {e}")
|
||||
kwargs.setdefault("metadata", {})["semantic-similarity"] = 0.0
|
||||
self._record_similarity(kwargs, 0.0)
|
||||
|
||||
async def async_set_cache(self, key: str, value: Any, **kwargs: Any) -> None:
|
||||
async def async_set_cache(self, key: str, value: object, **kwargs: object) -> None:
|
||||
print_verbose(f"Async Valkey semantic-cache set_cache, kwargs: {kwargs}")
|
||||
try:
|
||||
prompt: Final = self._get_prompt_from_kwargs(**kwargs)
|
||||
|
|
@ -278,7 +302,7 @@ class ValkeySemanticCache(RedisSemanticCache):
|
|||
print_verbose("No prompt provided for semantic caching")
|
||||
return
|
||||
|
||||
embedding: Final = await self._get_async_embedding(prompt, metadata=kwargs.get("metadata"))
|
||||
embedding: Final = await self._get_async_embedding(prompt, metadata=self._embedding_metadata(kwargs))
|
||||
await self._ensure_index_async(len(embedding))
|
||||
|
||||
doc_key: Final = self._doc_key(key)
|
||||
|
|
@ -289,31 +313,28 @@ class ValkeySemanticCache(RedisSemanticCache):
|
|||
except Exception as e:
|
||||
print_verbose(f"Error in async Valkey semantic-cache set_cache: {e}")
|
||||
|
||||
async def async_get_cache(self, key: str, **kwargs: Any) -> Any:
|
||||
async def async_get_cache(self, key: str, **kwargs: object) -> object:
|
||||
print_verbose(f"Async Valkey semantic-cache get_cache, kwargs: {kwargs}")
|
||||
try:
|
||||
prompt: Final = self._get_prompt_from_kwargs(**kwargs)
|
||||
if prompt is None:
|
||||
kwargs.setdefault("metadata", {})["semantic-similarity"] = 0.0
|
||||
self._record_similarity(kwargs, 0.0)
|
||||
return None
|
||||
|
||||
embedding: Final = await self._get_async_embedding(prompt, metadata=kwargs.get("metadata"))
|
||||
embedding: Final = await self._get_async_embedding(prompt, metadata=self._embedding_metadata(kwargs))
|
||||
await self._ensure_index_async(len(embedding))
|
||||
|
||||
search_result: Final = await self.async_client.ft(self.index_name).search(
|
||||
self._knn_query(key),
|
||||
query_params={"vec": self._embedding_to_bytes(embedding)},
|
||||
)
|
||||
search_result: Final[object] = await self._async_search(key, embedding)
|
||||
return self._resolve_hit(self._first_hit(search_result), key, **kwargs)
|
||||
except Exception as e:
|
||||
print_verbose(f"Error in async Valkey semantic-cache get_cache: {e}")
|
||||
kwargs.setdefault("metadata", {})["semantic-similarity"] = 0.0
|
||||
self._record_similarity(kwargs, 0.0)
|
||||
|
||||
async def async_set_cache_pipeline(self, cache_list: list[tuple[str, Any]], **kwargs: Any) -> None:
|
||||
async def async_set_cache_pipeline(self, cache_list: list[tuple[str, object]], **kwargs: object) -> None:
|
||||
try:
|
||||
await asyncio.gather(*[self.async_set_cache(key, value, **kwargs) for key, value in cache_list])
|
||||
except Exception as e:
|
||||
print_verbose(f"Error in Valkey semantic-cache async_set_cache_pipeline: {e}")
|
||||
|
||||
async def _index_info(self) -> dict:
|
||||
async def _index_info(self) -> Mapping[str, object]:
|
||||
return await self.async_client.ft(self.index_name).info()
|
||||
|
|
|
|||
|
|
@ -45,14 +45,14 @@ class ResponsesToCompletionBridgeHandler:
|
|||
return bool(stream)
|
||||
|
||||
@staticmethod
|
||||
def _is_preformatted_cached_chat_stream(result: Any) -> bool:
|
||||
def _is_preformatted_cached_chat_stream(result: object) -> bool:
|
||||
from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
|
||||
|
||||
return isinstance(result, CustomStreamWrapper) and result.custom_llm_provider == "cached_response"
|
||||
|
||||
@staticmethod
|
||||
def _coerce_response_object(
|
||||
response_obj: Any,
|
||||
response_obj: object,
|
||||
hidden_params: dict | None,
|
||||
) -> "ResponsesAPIResponse":
|
||||
if isinstance(response_obj, ResponsesAPIResponse):
|
||||
|
|
@ -78,8 +78,8 @@ class ResponsesToCompletionBridgeHandler:
|
|||
for _ in stream_iter:
|
||||
pass
|
||||
|
||||
completed: Final = getattr(stream_iter, "completed_response", None)
|
||||
response_obj: Final = getattr(completed, "response", None) if completed else None
|
||||
completed: Final[object] = getattr(stream_iter, "completed_response", None)
|
||||
response_obj: Final[object] = getattr(completed, "response", None) if completed else None
|
||||
if response_obj is None:
|
||||
raise ValueError("Stream ended without a completed response")
|
||||
|
||||
|
|
@ -93,8 +93,8 @@ class ResponsesToCompletionBridgeHandler:
|
|||
async for _ in stream_iter:
|
||||
pass
|
||||
|
||||
completed: Final = getattr(stream_iter, "completed_response", None)
|
||||
response_obj: Final = getattr(completed, "response", None) if completed else None
|
||||
completed: Final[object] = getattr(stream_iter, "completed_response", None)
|
||||
response_obj: Final[object] = getattr(completed, "response", None) if completed else None
|
||||
if response_obj is None:
|
||||
raise ValueError("Stream ended without a completed response")
|
||||
|
||||
|
|
@ -157,7 +157,7 @@ class ResponsesToCompletionBridgeHandler:
|
|||
def completion(
|
||||
self, *args, **kwargs
|
||||
) -> Union[
|
||||
Coroutine[Any, Any, Union["ModelResponse", "CustomStreamWrapper"]],
|
||||
Coroutine[None, None, Union["ModelResponse", "CustomStreamWrapper"]],
|
||||
"ModelResponse",
|
||||
"CustomStreamWrapper",
|
||||
]:
|
||||
|
|
|
|||
|
|
@ -212,7 +212,8 @@ def _tool_call_dict_from_output_item(item: Mapping[str, Any], index: int) -> _Ch
|
|||
LiteLLMCompletionResponsesConfig,
|
||||
)
|
||||
|
||||
is_custom: Final = item.get("type") == "custom_tool_call"
|
||||
item_type: Final[object] = item.get("type")
|
||||
is_custom: Final = item_type == "custom_tool_call"
|
||||
arguments: Final = (item.get("input") if is_custom else item.get("arguments")) or ""
|
||||
name: Final = item.get("name") or ("custom_tool" if is_custom else "")
|
||||
function_chunk: Final = ChatCompletionToolCallFunctionChunk(name=name, arguments=arguments)
|
||||
|
|
@ -222,7 +223,7 @@ def _tool_call_dict_from_output_item(item: Mapping[str, Any], index: int) -> _Ch
|
|||
function=function_chunk,
|
||||
index=index,
|
||||
)
|
||||
raw_provider_fields: Final = item.get("provider_specific_fields")
|
||||
raw_provider_fields: Final[object] = item.get("provider_specific_fields")
|
||||
if isinstance(raw_provider_fields, dict):
|
||||
provider_specific_fields = raw_provider_fields
|
||||
elif raw_provider_fields and hasattr(raw_provider_fields, "__dict__"):
|
||||
|
|
@ -507,7 +508,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
|
|||
|
||||
def _merge_responses_api_request_into_request_data(
|
||||
self,
|
||||
request_data: dict[str, Any],
|
||||
request_data: dict[str, object],
|
||||
responses_api_request: "ResponsesAPIOptionalRequestParams",
|
||||
instructions: str | None,
|
||||
) -> None:
|
||||
|
|
|
|||
|
|
@ -13,6 +13,12 @@ DEFAULT_BATCH_SIZE: Final = int(os.getenv("DEFAULT_BATCH_SIZE", 512))
|
|||
DEFAULT_FLUSH_INTERVAL_SECONDS: Final = int(os.getenv("DEFAULT_FLUSH_INTERVAL_SECONDS", 5))
|
||||
DEFAULT_S3_FLUSH_INTERVAL_SECONDS: Final = int(os.getenv("DEFAULT_S3_FLUSH_INTERVAL_SECONDS", 10))
|
||||
DEFAULT_S3_BATCH_SIZE: Final = int(os.getenv("DEFAULT_S3_BATCH_SIZE", 512))
|
||||
# https://docs.aws.amazon.com/AmazonS3/latest/userguide/object-keys.html
|
||||
MAX_S3_OBJECT_KEY_BYTES: Final = 1024
|
||||
S3_BOUNDED_OBJECT_KEY_HEAD_BYTES: Final = 64
|
||||
S3_PREFIX_DIGEST_CHARS: Final = 16
|
||||
# s3 allows 2048 bytes of combined metadata headers, which Content-Disposition counts against
|
||||
MAX_S3_OBJECT_DOWNLOAD_FILENAME_BYTES: Final = 1024
|
||||
DEFAULT_SQS_FLUSH_INTERVAL_SECONDS: Final = int(os.getenv("DEFAULT_SQS_FLUSH_INTERVAL_SECONDS", 10))
|
||||
DEFAULT_NUM_WORKERS_LITELLM_PROXY: Final = int(os.getenv("DEFAULT_NUM_WORKERS_LITELLM_PROXY", 1))
|
||||
DYNAMIC_RATE_LIMIT_ERROR_THRESHOLD_PER_MINUTE = int(os.getenv("DYNAMIC_RATE_LIMIT_ERROR_THRESHOLD_PER_MINUTE", 1))
|
||||
|
|
@ -130,6 +136,7 @@ MCP_CLIENT_TIMEOUT: Final = float(os.getenv("LITELLM_MCP_CLIENT_TIMEOUT", "60.0"
|
|||
MCP_TOOL_LISTING_TIMEOUT: Final = float(os.getenv("LITELLM_MCP_TOOL_LISTING_TIMEOUT", "30.0"))
|
||||
MCP_METADATA_TIMEOUT: Final = float(os.getenv("LITELLM_MCP_METADATA_TIMEOUT", "10.0"))
|
||||
MCP_HEALTH_CHECK_TIMEOUT: Final = float(os.getenv("LITELLM_MCP_HEALTH_CHECK_TIMEOUT", "10.0"))
|
||||
MCP_TOOL_LISTING_MAX_PAGES: Final = 1000
|
||||
|
||||
# Allowlist of commands permitted for MCP stdio transport.
|
||||
# Prevents arbitrary command execution via /mcp-rest/test/* endpoints or server creation.
|
||||
|
|
@ -484,6 +491,22 @@ FIREWORKS_AI_80_B: Final = int(os.getenv("FIREWORKS_AI_80_B", 80))
|
|||
#### Logging callback constants ####
|
||||
REDACTED_BY_LITELM_STRING: Final = "REDACTED_BY_LITELM"
|
||||
MAX_LANGFUSE_INITIALIZED_CLIENTS: Final = int(os.getenv("MAX_LANGFUSE_INITIALIZED_CLIENTS", 50))
|
||||
# Backpressure + lifetime bounds for the /v1/messages streaming relay (see
|
||||
# BaseAnthropicMessagesStreamingIterator.async_sse_wrapper). The relay queue is
|
||||
# bounded so a slow client throttles the upstream pump instead of letting it
|
||||
# buffer the whole response in memory; the detached-drain cap bounds how many
|
||||
# post-disconnect drains may run concurrently so client behavior can't create
|
||||
# unbounded worker state.
|
||||
ANTHROPIC_MESSAGES_STREAM_RELAY_QUEUE_MAXSIZE: Final = int(
|
||||
os.getenv("ANTHROPIC_MESSAGES_STREAM_RELAY_QUEUE_MAXSIZE", "1024")
|
||||
)
|
||||
# Setting this to 0 disables detached draining entirely: every post-disconnect
|
||||
# pump bills whatever partial output it has already collected and aborts the
|
||||
# upstream stream immediately, instead of continuing to drain for the real
|
||||
# terminal usage.
|
||||
ANTHROPIC_MESSAGES_MAX_DETACHED_STREAM_DRAINS: Final = int(
|
||||
os.getenv("ANTHROPIC_MESSAGES_MAX_DETACHED_STREAM_DRAINS", "100")
|
||||
)
|
||||
LOGGING_WORKER_CONCURRENCY: Final = int(os.getenv("LOGGING_WORKER_CONCURRENCY", 100)) # Must be above 0
|
||||
LOGGING_WORKER_MAX_QUEUE_SIZE: Final = int(os.getenv("LOGGING_WORKER_MAX_QUEUE_SIZE", 50_000))
|
||||
LOGGING_WORKER_MAX_TIME_PER_COROUTINE: Final = float(os.getenv("LOGGING_WORKER_MAX_TIME_PER_COROUTINE", 20.0))
|
||||
|
|
@ -614,6 +637,8 @@ LITELLM_CHAT_PROVIDERS: Final = [
|
|||
"nscale",
|
||||
"nebius",
|
||||
"dashscope",
|
||||
"qwencloud",
|
||||
"qwen_ai_platform",
|
||||
"modelscope",
|
||||
"moonshot",
|
||||
"publicai",
|
||||
|
|
@ -783,6 +808,7 @@ openai_compatible_endpoints: Final[list] = [
|
|||
"inference.api.nscale.com/v1",
|
||||
"api.studio.nebius.ai/v1",
|
||||
"https://dashscope-intl.aliyuncs.com/compatible-mode/v1",
|
||||
"https://dashscope.aliyuncs.com/compatible-mode/v1",
|
||||
"https://api-inference.modelscope.cn/v1",
|
||||
"https://api.moonshot.ai/v1",
|
||||
"https://api.publicai.co/v1",
|
||||
|
|
@ -806,6 +832,7 @@ openai_compatible_endpoints: Final[list] = [
|
|||
"https://api.meta.ai/v1",
|
||||
"https://api.cognition.ai/v1",
|
||||
"https://api.scx.ai/v1",
|
||||
"https://gigachat.devices.sberbank.ru/api/v1",
|
||||
]
|
||||
|
||||
|
||||
|
|
@ -855,6 +882,8 @@ openai_compatible_providers: Final[list] = [
|
|||
"nscale",
|
||||
"nebius",
|
||||
"dashscope",
|
||||
"qwencloud",
|
||||
"qwen_ai_platform",
|
||||
"modelscope",
|
||||
"moonshot",
|
||||
"v0",
|
||||
|
|
@ -885,6 +914,8 @@ openai_text_completion_compatible_providers: Final[list] = [ # providers that s
|
|||
"featherless_ai",
|
||||
"nebius",
|
||||
"dashscope",
|
||||
"qwencloud",
|
||||
"qwen_ai_platform",
|
||||
"modelscope",
|
||||
"moonshot",
|
||||
"publicai",
|
||||
|
|
@ -1092,7 +1123,7 @@ nebius_models: Final[set] = set(
|
|||
]
|
||||
)
|
||||
|
||||
dashscope_models: Final[set] = set(
|
||||
dashscope_models: Final[frozenset] = frozenset(
|
||||
[
|
||||
"qwen-turbo",
|
||||
"qwen-plus",
|
||||
|
|
@ -1107,6 +1138,10 @@ dashscope_models: Final[set] = set(
|
|||
]
|
||||
)
|
||||
|
||||
qwencloud_models: Final[frozenset] = frozenset(dashscope_models)
|
||||
|
||||
qwen_ai_platform_models: Final[frozenset] = frozenset(dashscope_models)
|
||||
|
||||
nebius_embedding_models: Final[set] = set(
|
||||
[
|
||||
"BAAI/bge-en-icl",
|
||||
|
|
@ -1223,6 +1258,7 @@ BEDROCK_CONVERSE_MODELS: Final = [
|
|||
"openai.gpt-oss-120b-1:0",
|
||||
"anthropic.claude-haiku-4-5-20251001-v1:0",
|
||||
"anthropic.claude-sonnet-4-5-20250929-v1:0",
|
||||
"anthropic.claude-fable-5-1",
|
||||
"anthropic.claude-fable-5",
|
||||
"anthropic.claude-sonnet-5",
|
||||
"anthropic.claude-opus-5",
|
||||
|
|
@ -1410,6 +1446,12 @@ DEFAULT_SOFT_BUDGET: Final = float(
|
|||
) # by default all litellm proxy keys have a soft budget of 50.0
|
||||
# makes it clear this is a rate limit error for a litellm virtual key
|
||||
RATE_LIMIT_ERROR_MESSAGE_FOR_VIRTUAL_KEY: Final = "LiteLLM Virtual Key user_api_key_hash"
|
||||
# Prefix of the 401 raised when a submitted virtual key is not shaped like one.
|
||||
INVALID_VIRTUAL_KEY_ERROR_MESSAGE: Final = "LiteLLM Virtual Key expected"
|
||||
# Attribute stamped on that 401 at its raise site so log routing recognises it by
|
||||
# provenance. Message text is caller-influenceable on other 401s, so it must not
|
||||
# be used to classify.
|
||||
INVALID_VIRTUAL_KEY_ERROR_MARKER: Final = "_litellm_invalid_virtual_key_error"
|
||||
|
||||
# Python garbage collection threshold configuration
|
||||
# Format: "gen0,gen1,gen2" e.g., "1000,50,50"
|
||||
|
|
@ -1548,6 +1590,7 @@ KEY_ROTATION_JOB_NAME: Final = "litellm_key_rotation_job"
|
|||
EXPIRED_UI_SESSION_KEY_CLEANUP_JOB_NAME: Final = "litellm_expired_ui_session_key_cleanup_job"
|
||||
WEEKLY_SPEND_REPORT_JOB_ID: Final = "weekly_spend_report_job"
|
||||
MONTHLY_SPEND_REPORT_JOB_ID: Final = "monthly_spend_report_job"
|
||||
USER_SPEND_ALERTS_JOB_ID: Final = "user_spend_alerts_job"
|
||||
PROMETHEUS_FALLBACK_STATS_JOB_ID: Final = "prometheus_fallback_stats_job"
|
||||
SLACK_DAILY_REPORT_LOCK_ID: Final = "slack_daily_report"
|
||||
SLACK_MODEL_DEPRECATION_LOCK_ID: Final = "slack_model_deprecation_warning"
|
||||
|
|
@ -1728,6 +1771,7 @@ SENTRY_DENYLIST: Final = [
|
|||
"jwt_token",
|
||||
"private_key",
|
||||
"SLACK_WEBHOOK_URL",
|
||||
"ALERTING_WEBHOOK_URL",
|
||||
"webhook_url",
|
||||
"LANGFUSE_SECRET_KEY",
|
||||
# Email Configuration
|
||||
|
|
|
|||
|
|
@ -641,12 +641,12 @@ def cost_per_token(
|
|||
return xai_cost_per_token(model=model, usage=usage_block)
|
||||
elif custom_llm_provider == "lemonade":
|
||||
return lemonade_cost_per_token(model=model, usage=usage_block)
|
||||
elif custom_llm_provider == "dashscope":
|
||||
elif custom_llm_provider in ("dashscope", "qwencloud", "qwen_ai_platform"):
|
||||
from litellm.llms.dashscope.cost_calculator import (
|
||||
cost_per_token as dashscope_cost_per_token,
|
||||
)
|
||||
|
||||
return dashscope_cost_per_token(model=model, usage=usage_block)
|
||||
return dashscope_cost_per_token(model=model, usage=usage_block, custom_llm_provider=custom_llm_provider)
|
||||
elif custom_llm_provider == "azure_ai":
|
||||
return azure_ai_cost_per_token(
|
||||
model=model,
|
||||
|
|
@ -1910,12 +1910,15 @@ def ocr_cost(
|
|||
if credits is not None and cost_per_credit is not None:
|
||||
return cost_per_credit * credits, 0.0
|
||||
|
||||
ocr_cost_per_page: float | None = None
|
||||
if model_info is not None:
|
||||
ocr_cost_per_page = model_info.get("ocr_cost_per_page")
|
||||
ocr_cost_per_page: Final = model_info.get("ocr_cost_per_page") if model_info is not None else None
|
||||
annotation_cost_per_page: Final = model_info.get("annotation_cost_per_page") if model_info is not None else None
|
||||
annotation_rate: Final = annotation_cost_per_page if annotation_cost_per_page is not None else ocr_cost_per_page
|
||||
|
||||
pages_processed: Final = response.usage_info.pages_processed
|
||||
if pages_processed is None:
|
||||
annotation_pages: Final = response.usage_info.pages_processed_annotation or 0
|
||||
has_billable_annotation_pages: Final = annotation_rate is not None and annotation_pages > 0
|
||||
|
||||
if pages_processed is None and not has_billable_annotation_pages:
|
||||
if cost_per_credit is not None or ocr_cost_per_page is None:
|
||||
# Surface missing usage data instead of silently under-reporting
|
||||
# cost. The previous behavior raised ValueError; we now return 0.0
|
||||
|
|
@ -1931,7 +1934,7 @@ def ocr_cost(
|
|||
return 0.0, 0.0
|
||||
raise ValueError("OCR response pages_processed is None")
|
||||
|
||||
if ocr_cost_per_page is None:
|
||||
if ocr_cost_per_page is None and not has_billable_annotation_pages:
|
||||
# No per-page pricing configured. Either the model is on credit-based
|
||||
# pricing (and credits weren't returned, so the credit branch above did
|
||||
# not match) or the model has no OCR pricing entry at all. Surface a
|
||||
|
|
@ -1947,8 +1950,9 @@ def ocr_cost(
|
|||
)
|
||||
return 0.0, 0.0
|
||||
|
||||
total_ocr_processing_cost: Final[float] = ocr_cost_per_page * pages_processed
|
||||
return total_ocr_processing_cost, 0.0
|
||||
ocr_pages_cost: Final = (ocr_cost_per_page or 0.0) * (pages_processed or 0)
|
||||
annotation_pages_cost: Final = (annotation_rate or 0.0) * annotation_pages
|
||||
return ocr_pages_cost + annotation_pages_cost, 0.0
|
||||
|
||||
|
||||
def vector_store_search_cost(
|
||||
|
|
@ -2268,6 +2272,10 @@ def batch_cost_calculator(
|
|||
return total_prompt_cost, total_completion_cost
|
||||
|
||||
|
||||
def _attribute_value(obj: object, name: str) -> object:
|
||||
return getattr(obj, name)
|
||||
|
||||
|
||||
def _summable_prompt_token_fields(prompt_tokens_details: BaseModel) -> list[str]:
|
||||
field_names: Final = list(type(prompt_tokens_details).model_fields)
|
||||
if getattr(prompt_tokens_details, "cache_write_tokens", None) is None:
|
||||
|
|
@ -2293,7 +2301,7 @@ class BaseTokenUsageProcessor:
|
|||
for usage in usage_objects:
|
||||
# Handle direct attributes by checking what exists in the model
|
||||
for attr in dir(usage):
|
||||
if not attr.startswith("_") and not callable(getattr(usage, attr)):
|
||||
if not attr.startswith("_") and not callable(_attribute_value(usage, attr)):
|
||||
current_val = getattr(combined, attr, 0)
|
||||
new_val = getattr(usage, attr, 0)
|
||||
if (
|
||||
|
|
@ -2313,7 +2321,7 @@ class BaseTokenUsageProcessor:
|
|||
if (
|
||||
hasattr(usage.prompt_tokens_details, attr)
|
||||
and not attr.startswith("_")
|
||||
and not callable(getattr(usage.prompt_tokens_details, attr))
|
||||
and not callable(_attribute_value(usage.prompt_tokens_details, attr))
|
||||
):
|
||||
current_val = getattr(combined.prompt_tokens_details, attr, 0) or 0
|
||||
new_val = getattr(usage.prompt_tokens_details, attr, 0) or 0
|
||||
|
|
@ -2332,7 +2340,9 @@ class BaseTokenUsageProcessor:
|
|||
# Check what keys exist in the model's completion_tokens_details
|
||||
# Access model_fields on the class, not the instance, to avoid Pydantic 2.11+ deprecation warnings
|
||||
for attr in type(usage.completion_tokens_details).model_fields:
|
||||
if not attr.startswith("_") and not callable(getattr(usage.completion_tokens_details, attr)):
|
||||
if not attr.startswith("_") and not callable(
|
||||
_attribute_value(usage.completion_tokens_details, attr)
|
||||
):
|
||||
current_val = getattr(combined.completion_tokens_details, attr, 0) or 0
|
||||
new_val = getattr(usage.completion_tokens_details, attr, 0) or 0
|
||||
if isinstance(new_val, (int, float)):
|
||||
|
|
|
|||
|
|
@ -115,9 +115,11 @@ class SpeechToCompletionBridgeHandler:
|
|||
**request_data,
|
||||
)
|
||||
|
||||
requested_response_format: Final = optional_params.get("response_format")
|
||||
if isinstance(result, ModelResponse):
|
||||
return self.transformation_handler.transform_response(
|
||||
model_response=result,
|
||||
response_format=requested_response_format if isinstance(requested_response_format, str) else None,
|
||||
)
|
||||
else:
|
||||
raise Exception(f"Unmapped response type. Got type: {type(result)}")
|
||||
|
|
|
|||
|
|
@ -1,10 +1,14 @@
|
|||
from collections.abc import Mapping
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Final, cast
|
||||
|
||||
from typing_extensions import NotRequired, ReadOnly, TypedDict
|
||||
|
||||
from litellm.constants import OPENAI_CHAT_COMPLETION_PARAMS
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm import Logging as LiteLLMLoggingObj
|
||||
from litellm.types.llms.openai import HttpxBinaryResponseContent
|
||||
from litellm.types.llms.openai import ChatCompletionUserMessage, HttpxBinaryResponseContent
|
||||
from litellm.types.utils import ModelResponse
|
||||
|
||||
|
||||
|
|
@ -16,7 +20,64 @@ def _completion_response_cost(model_response: "ModelResponse") -> float | None:
|
|||
return response_cost if isinstance(response_cost, float) else None
|
||||
|
||||
|
||||
GEMINI_TTS_CHAT_AUDIO_FORMAT: Final = "pcm16"
|
||||
GEMINI_TTS_RAW_RESPONSE_FORMAT: Final = "pcm"
|
||||
GEMINI_TTS_SUPPORTED_RESPONSE_FORMATS: Final = frozenset({"wav", GEMINI_TTS_RAW_RESPONSE_FORMAT})
|
||||
|
||||
|
||||
class ChatAudioParam(TypedDict):
|
||||
voice: ReadOnly[str]
|
||||
format: ReadOnly[NotRequired[str]]
|
||||
|
||||
|
||||
class SpeechToCompletionBridgeTransformationHandler:
|
||||
def _validate_response_format(
|
||||
self, model: str, custom_llm_provider: str, optional_params: Mapping[str, object]
|
||||
) -> None:
|
||||
if not self._is_gemini_tts_model(model):
|
||||
return
|
||||
response_format: Final = optional_params.get("response_format")
|
||||
if not isinstance(response_format, str) or response_format in GEMINI_TTS_SUPPORTED_RESPONSE_FORMATS:
|
||||
return
|
||||
from litellm.exceptions import BadRequestError
|
||||
|
||||
supported: Final = ", ".join(sorted(GEMINI_TTS_SUPPORTED_RESPONSE_FORMATS))
|
||||
raise BadRequestError(
|
||||
message=(
|
||||
f"Gemini TTS only produces raw PCM16 audio, so response_format='{response_format}'"
|
||||
f" is not supported. Supported response formats: {supported}."
|
||||
),
|
||||
model=model,
|
||||
llm_provider=custom_llm_provider,
|
||||
)
|
||||
|
||||
def _chat_completion_params(self, optional_params: Mapping[str, object]) -> Mapping[str, object]:
|
||||
return MappingProxyType(
|
||||
{
|
||||
param: value
|
||||
for param, value in optional_params.items()
|
||||
if param in OPENAI_CHAT_COMPLETION_PARAMS and param != "response_format"
|
||||
}
|
||||
)
|
||||
|
||||
def _chat_audio_format(self, model: str, optional_params: Mapping[str, object]) -> str | None:
|
||||
if self._is_gemini_tts_model(model):
|
||||
return GEMINI_TTS_CHAT_AUDIO_FORMAT
|
||||
response_format: Final = optional_params.get("response_format")
|
||||
return response_format if isinstance(response_format, str) else None
|
||||
|
||||
def _chat_audio_param(
|
||||
self, model: str, voice: str | Mapping[str, object] | None, optional_params: Mapping[str, object]
|
||||
) -> ChatAudioParam | None:
|
||||
if not isinstance(voice, str):
|
||||
return None
|
||||
audio_format: Final = self._chat_audio_format(model, optional_params)
|
||||
if audio_format is None:
|
||||
voice_only: Final[ChatAudioParam] = {"voice": voice}
|
||||
return voice_only
|
||||
audio: Final[ChatAudioParam] = {"voice": voice, "format": audio_format}
|
||||
return audio
|
||||
|
||||
def transform_request(
|
||||
self,
|
||||
model: str,
|
||||
|
|
@ -28,36 +89,20 @@ class SpeechToCompletionBridgeTransformationHandler:
|
|||
litellm_logging_obj: "LiteLLMLoggingObj",
|
||||
custom_llm_provider: str,
|
||||
) -> dict:
|
||||
passed_optional_params: Final = {}
|
||||
for op in optional_params:
|
||||
if op in OPENAI_CHAT_COMPLETION_PARAMS:
|
||||
passed_optional_params[op] = optional_params[op]
|
||||
|
||||
if voice is not None:
|
||||
if isinstance(voice, str):
|
||||
passed_optional_params["audio"] = {"voice": voice}
|
||||
if "response_format" in optional_params:
|
||||
passed_optional_params["audio"]["format"] = optional_params["response_format"]
|
||||
|
||||
return_kwargs = {
|
||||
self._validate_response_format(model, custom_llm_provider, optional_params)
|
||||
user_message: Final[ChatCompletionUserMessage] = {"role": "user", "content": input}
|
||||
return_kwargs: Final = {
|
||||
"model": model,
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": input,
|
||||
}
|
||||
],
|
||||
"messages": [user_message],
|
||||
"modalities": ["audio"],
|
||||
**passed_optional_params,
|
||||
**self._chat_completion_params(optional_params),
|
||||
"audio": self._chat_audio_param(model, voice, optional_params),
|
||||
**litellm_params,
|
||||
"headers": headers,
|
||||
"litellm_logging_obj": litellm_logging_obj,
|
||||
"custom_llm_provider": custom_llm_provider,
|
||||
}
|
||||
|
||||
# filter out None values
|
||||
return_kwargs = {k: v for k, v in return_kwargs.items() if v is not None}
|
||||
return return_kwargs
|
||||
return {k: v for k, v in return_kwargs.items() if v is not None}
|
||||
|
||||
def _convert_pcm16_to_wav(self, pcm_data: bytes, sample_rate: int = 24000, channels: int = 1) -> bytes:
|
||||
"""
|
||||
|
|
@ -103,7 +148,14 @@ class SpeechToCompletionBridgeTransformationHandler:
|
|||
"""Check if the model is a Gemini TTS model that returns PCM16 data."""
|
||||
return "gemini" in model.lower() and ("tts" in model.lower() or "preview-tts" in model.lower())
|
||||
|
||||
def transform_response(self, model_response: "ModelResponse") -> "HttpxBinaryResponseContent":
|
||||
def _gemini_tts_response_body(self, decoded_audio: bytes, response_format: str | None) -> tuple[bytes, str]:
|
||||
if response_format == GEMINI_TTS_RAW_RESPONSE_FORMAT:
|
||||
return decoded_audio, "audio/pcm"
|
||||
return self._convert_pcm16_to_wav(decoded_audio), "audio/wav"
|
||||
|
||||
def transform_response(
|
||||
self, model_response: "ModelResponse", response_format: str | None
|
||||
) -> "HttpxBinaryResponseContent":
|
||||
import base64
|
||||
|
||||
import httpx
|
||||
|
|
@ -114,23 +166,17 @@ class SpeechToCompletionBridgeTransformationHandler:
|
|||
audio_part: Final = cast(Choices, model_response.choices[0]).message.audio
|
||||
if audio_part is None:
|
||||
raise ValueError("No audio part found in the response")
|
||||
audio_content: Final = audio_part.data
|
||||
decoded_audio: Final = base64.b64decode(audio_part.data)
|
||||
|
||||
# Decode base64 to get binary content
|
||||
binary_data = base64.b64decode(audio_content)
|
||||
|
||||
# Check if this is a Gemini TTS model that returns raw PCM16 data
|
||||
model: Final = getattr(model_response, "model", "")
|
||||
headers: Final = {}
|
||||
if self._is_gemini_tts_model(model):
|
||||
# Convert PCM16 to WAV format for proper audio file playback
|
||||
binary_data = self._convert_pcm16_to_wav(binary_data)
|
||||
headers["Content-Type"] = "audio/wav"
|
||||
else:
|
||||
headers["Content-Type"] = "audio/mpeg"
|
||||
|
||||
# Create an httpx.Response object
|
||||
response: Final = httpx.Response(status_code=200, content=binary_data, headers=headers)
|
||||
content, content_type = (
|
||||
self._gemini_tts_response_body(decoded_audio, response_format)
|
||||
if self._is_gemini_tts_model(model)
|
||||
else (decoded_audio, "audio/mpeg")
|
||||
)
|
||||
response: Final = httpx.Response(
|
||||
status_code=200, content=content, headers=MappingProxyType({"Content-Type": content_type})
|
||||
)
|
||||
binary_response: Final = HttpxBinaryResponseContent(response)
|
||||
binary_response.set_response_cost(_completion_response_cost(model_response))
|
||||
return binary_response
|
||||
|
|
|
|||
|
|
@ -7,6 +7,7 @@ import base64
|
|||
import os
|
||||
from collections.abc import Awaitable, Callable, Generator
|
||||
from datetime import timedelta
|
||||
from functools import partial
|
||||
from importlib import metadata
|
||||
from typing import Any, Final, TypeVar
|
||||
|
||||
|
|
@ -47,7 +48,8 @@ from mcp.types import Tool as MCPTool
|
|||
from pydantic import AnyUrl
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.constants import MCP_CLIENT_TIMEOUT, MCP_NPM_CACHE_DIR
|
||||
from litellm.constants import MCP_CLIENT_TIMEOUT, MCP_NPM_CACHE_DIR, MCP_TOOL_LISTING_TIMEOUT
|
||||
from litellm.experimental_mcp_client.tools import list_tools_with_pagination
|
||||
from litellm.llms.custom_httpx.http_handler import get_ssl_configuration
|
||||
from litellm.types.llms.custom_http import VerifyTypes
|
||||
from litellm.types.mcp import (
|
||||
|
|
@ -603,17 +605,19 @@ class MCPClient:
|
|||
"""
|
||||
verbose_logger.debug("MCP client listing tools from %s", self.server_url or "stdio")
|
||||
|
||||
async def _list_tools_operation(session: ClientSession):
|
||||
return await session.list_tools()
|
||||
|
||||
try:
|
||||
result: Final = await self.run_with_session(_list_tools_operation, quiet_on_error=raise_on_error)
|
||||
tool_count: Final = len(result.tools)
|
||||
tool_names: Final = [tool.name for tool in result.tools]
|
||||
# A per-server timeout above the global default extends the whole-walk deadline
|
||||
listing_deadline: Final = max(self.timeout, MCP_TOOL_LISTING_TIMEOUT)
|
||||
tools: Final = await self.run_with_session(
|
||||
partial(list_tools_with_pagination, listing_deadline=listing_deadline),
|
||||
quiet_on_error=raise_on_error,
|
||||
)
|
||||
tool_count: Final = len(tools)
|
||||
tool_names: Final = tuple(tool.name for tool in tools)
|
||||
verbose_logger.info(
|
||||
"MCP client listed %s tools from %s: %s", tool_count, self.server_url or "stdio", tool_names
|
||||
)
|
||||
return result.tools
|
||||
return tools
|
||||
except asyncio.CancelledError:
|
||||
verbose_logger.warning("MCP client list_tools was cancelled")
|
||||
raise
|
||||
|
|
|
|||
|
|
@ -1,14 +1,22 @@
|
|||
import json
|
||||
from typing import Final, Literal
|
||||
|
||||
import anyio
|
||||
from mcp import ClientSession
|
||||
from mcp.types import CallToolRequestParams as MCPCallToolRequestParams
|
||||
from mcp.types import CallToolResult as MCPCallToolResult
|
||||
from mcp.types import PaginatedRequestParams
|
||||
from mcp.types import Tool as MCPTool
|
||||
from openai.types.chat import ChatCompletionToolParam
|
||||
from openai.types.responses.function_tool_param import FunctionToolParam
|
||||
from openai.types.shared_params.function_definition import FunctionDefinition
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.constants import (
|
||||
MCP_CLIENT_TIMEOUT,
|
||||
MCP_TOOL_LISTING_MAX_PAGES,
|
||||
MCP_TOOL_LISTING_TIMEOUT,
|
||||
)
|
||||
from litellm.types.llms.anthropic import AnthropicMessagesTool
|
||||
from litellm.types.utils import ChatCompletionMessageToolCall
|
||||
|
||||
|
|
@ -90,6 +98,64 @@ def transform_mcp_tool_to_anthropic_tool(mcp_tool: MCPTool) -> AnthropicMessages
|
|||
)
|
||||
|
||||
|
||||
async def list_tools_with_pagination(
|
||||
session: ClientSession, listing_deadline: float | None = None
|
||||
) -> list[MCPTool]: # mutable-ok: list return contract
|
||||
"""Collect tools from every tools/list page by following nextCursor.
|
||||
|
||||
Stops and returns the tools collected so far when the upstream repeats a
|
||||
cursor, the page cap is reached, or the whole-walk deadline expires, so a
|
||||
buggy or slow upstream yields a partial catalog instead of an error.
|
||||
listing_deadline overrides the default whole-walk deadline; callers with a
|
||||
per-server timeout above the global default pass it through here.
|
||||
"""
|
||||
tools: Final[list[MCPTool]] = [] # mutable-ok: accumulates each page's tools
|
||||
seen_cursors: Final[set[str]] = set() # mutable-ok: guards against cursor loops
|
||||
cursor: str | None = None # rebind-ok: advances to each page's nextCursor
|
||||
# The per-request session read timeout restarts on every page, so a multi-page
|
||||
# walk needs its own overall deadline. max() keeps the pre-pagination guarantee
|
||||
# that a single page slower than the listing timeout but within the client
|
||||
# timeout still succeeds.
|
||||
effective_deadline: Final = (
|
||||
listing_deadline if listing_deadline is not None else max(MCP_CLIENT_TIMEOUT, MCP_TOOL_LISTING_TIMEOUT)
|
||||
)
|
||||
|
||||
with anyio.move_on_after(effective_deadline):
|
||||
for _ in range(MCP_TOOL_LISTING_MAX_PAGES):
|
||||
result = (
|
||||
await session.list_tools()
|
||||
if cursor is None
|
||||
else await session.list_tools(params=PaginatedRequestParams(cursor=cursor))
|
||||
)
|
||||
tools.extend(result.tools)
|
||||
|
||||
next_cursor = getattr(result, "nextCursor", None)
|
||||
if not isinstance(next_cursor, str) or not next_cursor:
|
||||
return tools
|
||||
if next_cursor in seen_cursors:
|
||||
verbose_logger.warning(
|
||||
"MCP server repeated a tools/list cursor while listing tools; returning %s tools collected so far",
|
||||
len(tools),
|
||||
)
|
||||
return tools
|
||||
seen_cursors.add(next_cursor)
|
||||
cursor = next_cursor
|
||||
|
||||
verbose_logger.warning(
|
||||
"MCP server tools/list pagination exceeded the maximum of %s pages; returning %s tools collected so far",
|
||||
MCP_TOOL_LISTING_MAX_PAGES,
|
||||
len(tools),
|
||||
)
|
||||
return tools
|
||||
|
||||
verbose_logger.warning(
|
||||
"MCP server tools/list pagination exceeded the %s second listing deadline; returning %s tools collected so far",
|
||||
effective_deadline,
|
||||
len(tools),
|
||||
)
|
||||
return tools
|
||||
|
||||
|
||||
async def load_mcp_tools(
|
||||
session: ClientSession, format: Literal["mcp", "openai"] = "mcp"
|
||||
) -> list[MCPTool] | list[ChatCompletionToolParam]:
|
||||
|
|
@ -103,10 +169,12 @@ async def load_mcp_tools(
|
|||
|
||||
If format is set to "openai", the tools are converted to OpenAI API compatible tools.
|
||||
"""
|
||||
tools: Final = await session.list_tools()
|
||||
tools: Final = await list_tools_with_pagination(session)
|
||||
if format == "openai":
|
||||
return [transform_mcp_tool_to_openai_tool(mcp_tool=tool) for tool in tools.tools]
|
||||
return tools.tools
|
||||
return [ # mutable-ok: public API returns a list
|
||||
transform_mcp_tool_to_openai_tool(mcp_tool=tool) for tool in tools
|
||||
]
|
||||
return tools
|
||||
|
||||
|
||||
########################################################
|
||||
|
|
|
|||
|
|
@ -1,8 +1,9 @@
|
|||
import json
|
||||
from collections.abc import AsyncIterator, Iterator, Sequence
|
||||
from typing import Any, Final, TypedDict, cast
|
||||
from collections.abc import AsyncIterator, Callable, Iterator, Mapping, Sequence
|
||||
from types import MappingProxyType
|
||||
from typing import Any, Final, TypeAlias, cast
|
||||
|
||||
from typing_extensions import ReadOnly
|
||||
from typing_extensions import ReadOnly, TypedDict
|
||||
|
||||
from litellm import verbose_logger
|
||||
from litellm.litellm_core_utils.json_validation_rule import normalize_tool_schema
|
||||
|
|
@ -11,7 +12,6 @@ from litellm.types.llms.openai import (
|
|||
ChatCompletionAssistantMessage,
|
||||
ChatCompletionAssistantToolCall,
|
||||
ChatCompletionImageObject,
|
||||
ChatCompletionRequest,
|
||||
ChatCompletionSystemMessage,
|
||||
ChatCompletionTextObject,
|
||||
ChatCompletionToolCallFunctionChunk,
|
||||
|
|
@ -23,35 +23,63 @@ from litellm.types.llms.openai import (
|
|||
from litellm.types.router import GenericLiteLLMParams
|
||||
from litellm.types.utils import (
|
||||
AdapterCompletionStreamWrapper,
|
||||
ChatCompletionDeltaCustomToolCall,
|
||||
ChatCompletionDeltaToolCall,
|
||||
ChatCompletionMessageCustomToolCall,
|
||||
ChatCompletionMessageToolCall,
|
||||
Choices,
|
||||
Delta,
|
||||
Function,
|
||||
Message,
|
||||
ModelResponse,
|
||||
ModelResponseStream,
|
||||
StreamingChoices,
|
||||
Usage,
|
||||
)
|
||||
|
||||
|
||||
class _GenAITextPart(TypedDict, total=False):
|
||||
text: ReadOnly[str]
|
||||
_JsonDict: TypeAlias = dict[str, object]
|
||||
_JsonDictList: TypeAlias = list[_JsonDict]
|
||||
|
||||
|
||||
class _GenAISystemInstruction(TypedDict, total=False):
|
||||
parts: ReadOnly[list[_GenAITextPart]]
|
||||
class _ToolCallAccumulator(TypedDict):
|
||||
name: ReadOnly[str]
|
||||
arguments: ReadOnly[str]
|
||||
|
||||
|
||||
class _GenAIFunctionCall(TypedDict):
|
||||
name: ReadOnly[str]
|
||||
args: ReadOnly[Mapping[str, object]]
|
||||
|
||||
|
||||
class _GenAIPart(TypedDict, total=False):
|
||||
text: ReadOnly[str]
|
||||
functionCall: ReadOnly[dict[str, object]]
|
||||
functionCall: ReadOnly[_GenAIFunctionCall]
|
||||
|
||||
|
||||
class _GenAIFunctionResponse(TypedDict, total=False):
|
||||
name: ReadOnly[str]
|
||||
response: ReadOnly[object]
|
||||
|
||||
|
||||
class _GenAIRequestFunctionCall(TypedDict, total=False):
|
||||
name: ReadOnly[str]
|
||||
args: ReadOnly[Mapping[str, object]]
|
||||
|
||||
|
||||
class _GenAIContentPart(TypedDict, total=False):
|
||||
text: ReadOnly[str]
|
||||
inline_data: ReadOnly[Mapping[str, str]]
|
||||
functionResponse: ReadOnly[_GenAIFunctionResponse]
|
||||
functionCall: ReadOnly[_GenAIRequestFunctionCall]
|
||||
|
||||
|
||||
class _GenAIFunctionDeclaration(TypedDict, total=False):
|
||||
name: ReadOnly[str]
|
||||
description: ReadOnly[str]
|
||||
parametersJsonSchema: ReadOnly[dict[str, object]]
|
||||
parametersJsonSchema: ReadOnly[object]
|
||||
|
||||
|
||||
class _GenAITool(TypedDict, total=False):
|
||||
functionDeclarations: ReadOnly[list[_GenAIFunctionDeclaration]]
|
||||
functionDeclarations: ReadOnly[Sequence[_GenAIFunctionDeclaration]]
|
||||
|
||||
|
||||
class _GenAIFunctionCallingConfig(TypedDict, total=False):
|
||||
|
|
@ -62,9 +90,11 @@ class _GenAIToolConfig(TypedDict, total=False):
|
|||
functionCallingConfig: ReadOnly[_GenAIFunctionCallingConfig]
|
||||
|
||||
|
||||
def _decode_tool_call_arguments(raw_arguments: str) -> object:
|
||||
"""Decode a tool call's JSON-encoded arguments into the value Google GenAI expects."""
|
||||
return json.loads(raw_arguments)
|
||||
class _GenAISystemInstruction(TypedDict, total=False):
|
||||
parts: ReadOnly[Sequence[Mapping[str, str]]]
|
||||
|
||||
|
||||
_EMPTY_STR_MAPPING: Final[Mapping[str, str]] = MappingProxyType({})
|
||||
|
||||
|
||||
class GoogleGenAIStreamWrapper(AdapterCompletionStreamWrapper):
|
||||
|
|
@ -74,12 +104,11 @@ class GoogleGenAIStreamWrapper(AdapterCompletionStreamWrapper):
|
|||
"""
|
||||
|
||||
sent_first_chunk: bool = False
|
||||
# State tracking for accumulating partial tool calls
|
||||
accumulated_tool_calls: dict[int, dict[str, str]]
|
||||
_parse_accumulated_args: Callable[[str], Mapping[str, object]] = staticmethod(json.loads)
|
||||
|
||||
def __init__(self, completion_stream: object):
|
||||
self.sent_first_chunk = False
|
||||
self.accumulated_tool_calls = {}
|
||||
self.accumulated_tool_calls = dict[int, _ToolCallAccumulator]()
|
||||
self._returned_response = False
|
||||
super().__init__(completion_stream)
|
||||
|
||||
|
|
@ -124,7 +153,7 @@ class GoogleGenAIStreamWrapper(AdapterCompletionStreamWrapper):
|
|||
# After the stream is exhausted, check for any remaining accumulated tool calls
|
||||
if self.accumulated_tool_calls:
|
||||
try:
|
||||
parts: Final[list[_GenAIPart]] = []
|
||||
parts: Final = list[_GenAIPart]()
|
||||
for (
|
||||
tool_call_index,
|
||||
tool_call_data,
|
||||
|
|
@ -132,7 +161,9 @@ class GoogleGenAIStreamWrapper(AdapterCompletionStreamWrapper):
|
|||
try:
|
||||
# For tool calls with no arguments, accumulated_args will be "", which is not valid JSON.
|
||||
# We default to an empty JSON object in this case.
|
||||
parsed_args = _decode_tool_call_arguments(tool_call_data["arguments"] or "{}")
|
||||
parsed_args: Mapping[str, object] = self._parse_accumulated_args(
|
||||
tool_call_data["arguments"] or "{}"
|
||||
)
|
||||
function_call_part: _GenAIPart = {
|
||||
"functionCall": {
|
||||
"name": tool_call_data["name"] or "undefined_tool_name",
|
||||
|
|
@ -149,7 +180,7 @@ class GoogleGenAIStreamWrapper(AdapterCompletionStreamWrapper):
|
|||
tool_call_data["arguments"],
|
||||
)
|
||||
if parts:
|
||||
final_chunk: Final[dict[str, object]] = {
|
||||
final_chunk: Final = {
|
||||
"candidates": [
|
||||
{
|
||||
"content": {"parts": parts, "role": "model"},
|
||||
|
|
@ -211,14 +242,16 @@ class GoogleGenAIStreamWrapper(AdapterCompletionStreamWrapper):
|
|||
class GoogleGenAIAdapter:
|
||||
"""Adapter for transforming Google GenAI generate_content requests to/from litellm.completion format"""
|
||||
|
||||
_parse_tool_call_args: Callable[[str], Mapping[str, object]] = staticmethod(json.loads)
|
||||
|
||||
def __init__(self) -> None:
|
||||
pass
|
||||
|
||||
def translate_generate_content_to_completion(
|
||||
self,
|
||||
model: str,
|
||||
contents: list[dict[str, Any]] | dict[str, Any],
|
||||
config: dict[str, Any] | None = None,
|
||||
contents: _JsonDictList | _JsonDict,
|
||||
config: Mapping[str, object] | None = None,
|
||||
litellm_params: GenericLiteLLMParams | None = None,
|
||||
**kwargs,
|
||||
) -> dict[str, Any]:
|
||||
|
|
@ -250,7 +283,7 @@ class GoogleGenAIAdapter:
|
|||
messages: Final = self._transform_contents_to_messages(contents_list, system_instruction=system_instruction)
|
||||
|
||||
# Create base request as dict (which is compatible with ChatCompletionRequest)
|
||||
completion_request: Final[ChatCompletionRequest] = {
|
||||
completion_request: Final[_JsonDict] = {
|
||||
"model": model,
|
||||
"messages": messages,
|
||||
}
|
||||
|
|
@ -312,9 +345,9 @@ class GoogleGenAIAdapter:
|
|||
|
||||
def _add_generic_litellm_params_to_request(
|
||||
self,
|
||||
completion_request_dict: dict[str, object],
|
||||
completion_request_dict: _JsonDict,
|
||||
litellm_params: GenericLiteLLMParams | None = None,
|
||||
) -> dict[str, object]:
|
||||
) -> _JsonDict:
|
||||
"""Add generic litellm params to request. e.g add api_base, api_key, api_version, etc.
|
||||
|
||||
Args:
|
||||
|
|
@ -326,7 +359,7 @@ class GoogleGenAIAdapter:
|
|||
"""
|
||||
allowed_fields: Final = GenericLiteLLMParams.model_fields.keys()
|
||||
if litellm_params:
|
||||
litellm_dict: Final = litellm_params.model_dump(exclude_none=True)
|
||||
litellm_dict: Final[_JsonDict] = litellm_params.model_dump(exclude_none=True)
|
||||
for key, value in litellm_dict.items():
|
||||
if key in allowed_fields:
|
||||
completion_request_dict[key] = value
|
||||
|
|
@ -346,12 +379,12 @@ class GoogleGenAIAdapter:
|
|||
tools: Sequence[_GenAITool],
|
||||
) -> list[ChatCompletionToolParam]:
|
||||
"""Transform Google GenAI tools to OpenAI tools format"""
|
||||
openai_tools: Final[list[dict[str, object]]] = []
|
||||
openai_tools: Final = list[_JsonDict]()
|
||||
|
||||
for tool in tools:
|
||||
if "functionDeclarations" in tool:
|
||||
for func_decl in tool["functionDeclarations"]:
|
||||
function_chunk: dict[str, object] = {
|
||||
function_chunk: _JsonDict = {
|
||||
"name": func_decl.get("name", ""),
|
||||
}
|
||||
|
||||
|
|
@ -360,7 +393,7 @@ class GoogleGenAIAdapter:
|
|||
if "parametersJsonSchema" in func_decl:
|
||||
function_chunk["parameters"] = func_decl["parametersJsonSchema"]
|
||||
|
||||
openai_tool: dict[str, object] = {"type": "function", "function": function_chunk}
|
||||
openai_tool: _JsonDict = {"type": "function", "function": function_chunk}
|
||||
openai_tools.append(openai_tool)
|
||||
|
||||
# normalize the tool schemas
|
||||
|
|
@ -391,13 +424,13 @@ class GoogleGenAIAdapter:
|
|||
|
||||
# Handle system instruction
|
||||
if system_instruction:
|
||||
system_parts: Final = system_instruction.get("parts", [])
|
||||
system_parts: Final[Sequence[Mapping[str, str]]] = system_instruction.get("parts", [])
|
||||
if system_parts and "text" in system_parts[0]:
|
||||
messages.append(ChatCompletionSystemMessage(role="system", content=system_parts[0]["text"]))
|
||||
|
||||
for content in contents:
|
||||
role = content.get("role", "user")
|
||||
parts = content.get("parts", [])
|
||||
parts: Sequence[_GenAIContentPart | str | None] = content.get("parts", [])
|
||||
|
||||
if role == "user":
|
||||
# Handle user messages with potential function responses
|
||||
|
|
@ -500,7 +533,7 @@ class GoogleGenAIAdapter:
|
|||
def translate_completion_to_generate_content(
|
||||
self,
|
||||
response: ModelResponse,
|
||||
) -> dict[str, object]:
|
||||
) -> _JsonDict:
|
||||
"""
|
||||
Transform litellm completion response to Google GenAI generate_content format
|
||||
|
||||
|
|
@ -523,13 +556,13 @@ class GoogleGenAIAdapter:
|
|||
parts = self._transform_openai_message_to_google_genai_parts(choice.message)
|
||||
else:
|
||||
# Fallback for generic choice objects
|
||||
message_content = getattr(choice, "message", {}).get("content", "") or getattr(choice, "delta", {}).get(
|
||||
"content", ""
|
||||
)
|
||||
message_content: str = getattr(choice, "message", _EMPTY_STR_MAPPING).get("content", "") or getattr(
|
||||
choice, "delta", _EMPTY_STR_MAPPING
|
||||
).get("content", "")
|
||||
parts = [{"text": message_content}] if message_content else []
|
||||
|
||||
# Create Google GenAI format response
|
||||
generate_content_response: Final[dict[str, object]] = {
|
||||
generate_content_response: Final[_JsonDict] = {
|
||||
"candidates": [
|
||||
{
|
||||
"content": {"parts": parts, "role": "model"},
|
||||
|
|
@ -563,7 +596,7 @@ class GoogleGenAIAdapter:
|
|||
self,
|
||||
response: ModelResponse | ModelResponseStream,
|
||||
wrapper: GoogleGenAIStreamWrapper,
|
||||
) -> dict[str, object] | None:
|
||||
) -> Mapping[str, object] | None:
|
||||
"""
|
||||
Transform streaming litellm completion chunk to Google GenAI generate_content format
|
||||
|
||||
|
|
@ -590,7 +623,7 @@ class GoogleGenAIAdapter:
|
|||
finish_reason: str | None = getattr(choice, "finish_reason", None)
|
||||
else:
|
||||
# Fallback for generic choice objects
|
||||
message_content: Final = getattr(choice, "delta", {}).get("content", "")
|
||||
message_content: Final[str] = getattr(choice, "delta", _EMPTY_STR_MAPPING).get("content", "")
|
||||
parts = [{"text": message_content}] if message_content else []
|
||||
finish_reason = getattr(choice, "finish_reason", None)
|
||||
|
||||
|
|
@ -599,7 +632,7 @@ class GoogleGenAIAdapter:
|
|||
return None
|
||||
|
||||
# Create Google GenAI streaming format response
|
||||
streaming_chunk: Final[dict[str, object]] = {
|
||||
streaming_chunk: Final[_JsonDict] = {
|
||||
"candidates": [
|
||||
{
|
||||
"content": {"parts": parts, "role": "model"},
|
||||
|
|
@ -635,10 +668,10 @@ class GoogleGenAIAdapter:
|
|||
|
||||
def _transform_openai_message_to_google_genai_parts(
|
||||
self,
|
||||
message: Any,
|
||||
) -> list[_GenAIPart]:
|
||||
message: Message,
|
||||
) -> Sequence[_GenAIPart]:
|
||||
"""Transform OpenAI message to Google GenAI parts format"""
|
||||
parts: Final[list[_GenAIPart]] = []
|
||||
parts: Final = list[_GenAIPart]()
|
||||
|
||||
# Add text content if present
|
||||
if hasattr(message, "content") and message.content:
|
||||
|
|
@ -646,20 +679,22 @@ class GoogleGenAIAdapter:
|
|||
|
||||
# Add tool calls if present
|
||||
if hasattr(message, "tool_calls") and message.tool_calls:
|
||||
for tool_call in message.tool_calls:
|
||||
if hasattr(tool_call, "function") and tool_call.function:
|
||||
tool_calls: Final[Sequence[ChatCompletionMessageToolCall | ChatCompletionMessageCustomToolCall]] = (
|
||||
message.tool_calls
|
||||
)
|
||||
for tool_call in tool_calls:
|
||||
function: Function | None = getattr(tool_call, "function", None)
|
||||
if function:
|
||||
try:
|
||||
args = (
|
||||
_decode_tool_call_arguments(tool_call.function.arguments)
|
||||
if tool_call.function.arguments
|
||||
else {}
|
||||
args: Mapping[str, object] = (
|
||||
self._parse_tool_call_args(function.arguments) if function.arguments else {}
|
||||
)
|
||||
except json.JSONDecodeError:
|
||||
args = {}
|
||||
|
||||
function_call_part: _GenAIPart = {
|
||||
"functionCall": {
|
||||
"name": tool_call.function.name or "undefined_tool_name",
|
||||
"name": function.name or "undefined_tool_name",
|
||||
"args": args,
|
||||
}
|
||||
}
|
||||
|
|
@ -668,24 +703,26 @@ class GoogleGenAIAdapter:
|
|||
return parts if parts else [{"text": ""}]
|
||||
|
||||
def _transform_openai_delta_to_google_genai_parts_with_accumulation(
|
||||
self, delta: Any, wrapper: GoogleGenAIStreamWrapper
|
||||
) -> list[_GenAIPart]:
|
||||
self, delta: Delta, wrapper: GoogleGenAIStreamWrapper
|
||||
) -> Sequence[_GenAIPart]:
|
||||
"""Transforms OpenAI delta to Google GenAI parts, accumulating streaming tool calls."""
|
||||
|
||||
# 1. Initialize wrapper state if it doesn't exist
|
||||
if not hasattr(wrapper, "accumulated_tool_calls"):
|
||||
wrapper.accumulated_tool_calls = {}
|
||||
|
||||
parts: Final[list[_GenAIPart]] = []
|
||||
parts: Final = list[_GenAIPart]()
|
||||
|
||||
if hasattr(delta, "content") and delta.content:
|
||||
parts.append({"text": delta.content})
|
||||
|
||||
# 2. Ensure tool_calls is iterable
|
||||
tool_calls: Final = delta.tool_calls or []
|
||||
tool_calls: Final[Sequence[ChatCompletionDeltaToolCall | ChatCompletionDeltaCustomToolCall]] = (
|
||||
delta.tool_calls or []
|
||||
)
|
||||
|
||||
for tool_call in tool_calls:
|
||||
if not hasattr(tool_call, "function"):
|
||||
if not hasattr(tool_call, "function") or isinstance(tool_call, ChatCompletionDeltaCustomToolCall):
|
||||
continue
|
||||
|
||||
# 3. Use `index` as the primary key for accumulation
|
||||
|
|
@ -701,19 +738,20 @@ class GoogleGenAIAdapter:
|
|||
}
|
||||
|
||||
# Accumulate name and arguments
|
||||
function_name = getattr(tool_call.function, "name", None)
|
||||
args_chunk = getattr(tool_call.function, "arguments", None)
|
||||
delta_function: Function | None = getattr(tool_call, "function", None)
|
||||
function_name: str | None = getattr(delta_function, "name", None)
|
||||
args_chunk: str | None = getattr(delta_function, "arguments", None)
|
||||
|
||||
# Optimization: Skip chunks that have no new data
|
||||
if not function_name and not args_chunk:
|
||||
verbose_logger.debug("Skipping empty tool call chunk for index: %s", tool_call_index)
|
||||
continue
|
||||
|
||||
if function_name:
|
||||
wrapper.accumulated_tool_calls[tool_call_index]["name"] = function_name
|
||||
|
||||
if args_chunk:
|
||||
wrapper.accumulated_tool_calls[tool_call_index]["arguments"] += args_chunk
|
||||
previous_data: _ToolCallAccumulator = wrapper.accumulated_tool_calls[tool_call_index]
|
||||
wrapper.accumulated_tool_calls[tool_call_index] = _ToolCallAccumulator(
|
||||
name=function_name or previous_data["name"],
|
||||
arguments=previous_data["arguments"] + (args_chunk or ""),
|
||||
)
|
||||
|
||||
# Attempt to parse and emit a complete tool call
|
||||
accumulated_data = wrapper.accumulated_tool_calls[tool_call_index]
|
||||
|
|
@ -723,7 +761,7 @@ class GoogleGenAIAdapter:
|
|||
# 5. Attempt to parse arguments even if name hasn't arrived.
|
||||
try:
|
||||
# Attempt to parse the accumulated arguments string
|
||||
parsed_args = _decode_tool_call_arguments(accumulated_args)
|
||||
parsed_args: Mapping[str, object] = self._parse_tool_call_args(accumulated_args)
|
||||
|
||||
# If parsing succeeds, but we don't have a name yet, wait.
|
||||
# The part will be created by a later chunk that brings the name.
|
||||
|
|
@ -757,7 +795,7 @@ class GoogleGenAIAdapter:
|
|||
|
||||
return mapping.get(finish_reason, "STOP")
|
||||
|
||||
def _map_usage(self, usage: Usage | None) -> dict[str, int]:
|
||||
def _map_usage(self, usage: object) -> Mapping[str, int]:
|
||||
"""Map OpenAI usage to Google GenAI usage format"""
|
||||
return {
|
||||
"promptTokenCount": getattr(usage, "prompt_tokens", 0) or 0,
|
||||
|
|
|
|||
|
|
@ -52,10 +52,10 @@ class GenerateContentSetupResult(BaseModel):
|
|||
model_config: ClassVar[ConfigDict] = ConfigDict(arbitrary_types_allowed=True)
|
||||
|
||||
model: str
|
||||
request_body: dict[str, Any]
|
||||
request_body: dict[str, object]
|
||||
custom_llm_provider: str
|
||||
generate_content_provider_config: BaseGoogleGenAIGenerateContentConfig | None
|
||||
generate_content_config_dict: dict[str, Any]
|
||||
generate_content_config_dict: dict[str, object]
|
||||
native_request_fields: dict[str, object]
|
||||
litellm_params: GenericLiteLLMParams
|
||||
litellm_logging_obj: LiteLLMLoggingObj
|
||||
|
|
@ -68,7 +68,7 @@ class GenerateContentHelper:
|
|||
@staticmethod
|
||||
def mock_generate_content_response(
|
||||
mock_response: str = "This is a mock response from Google GenAI generate_content.",
|
||||
) -> dict[str, Any]:
|
||||
) -> dict[str, object]:
|
||||
"""Mock response for generate_content for testing purposes"""
|
||||
return {
|
||||
"text": mock_response,
|
||||
|
|
@ -239,9 +239,9 @@ async def agenerate_content(
|
|||
tools: ToolConfigDict | None = None,
|
||||
# Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs.
|
||||
# The extra values given here take precedence over values defined on the client or passed to this method.
|
||||
extra_headers: dict[str, Any] | None = None,
|
||||
extra_query: dict[str, Any] | None = None,
|
||||
extra_body: dict[str, Any] | None = None,
|
||||
extra_headers: dict[str, object] | None = None,
|
||||
extra_query: dict[str, object] | None = None,
|
||||
extra_body: dict[str, object] | None = None,
|
||||
timeout: float | httpx.Timeout | None = None,
|
||||
# LiteLLM specific params,
|
||||
custom_llm_provider: str | None = None,
|
||||
|
|
@ -307,9 +307,9 @@ def generate_content(
|
|||
tools: ToolConfigDict | None = None,
|
||||
# Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs.
|
||||
# The extra values given here take precedence over values defined on the client or passed to this method.
|
||||
extra_headers: dict[str, Any] | None = None,
|
||||
extra_query: dict[str, Any] | None = None,
|
||||
extra_body: dict[str, Any] | None = None,
|
||||
extra_headers: dict[str, object] | None = None,
|
||||
extra_query: dict[str, object] | None = None,
|
||||
extra_body: dict[str, object] | None = None,
|
||||
timeout: float | httpx.Timeout | None = None,
|
||||
# LiteLLM specific params,
|
||||
custom_llm_provider: str | None = None,
|
||||
|
|
@ -397,9 +397,9 @@ async def agenerate_content_stream(
|
|||
tools: ToolConfigDict | None = None,
|
||||
# Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs.
|
||||
# The extra values given here take precedence over values defined on the client or passed to this method.
|
||||
extra_headers: dict[str, Any] | None = None,
|
||||
extra_query: dict[str, Any] | None = None,
|
||||
extra_body: dict[str, Any] | None = None,
|
||||
extra_headers: dict[str, object] | None = None,
|
||||
extra_query: dict[str, object] | None = None,
|
||||
extra_body: dict[str, object] | None = None,
|
||||
timeout: float | httpx.Timeout | None = None,
|
||||
# LiteLLM specific params,
|
||||
custom_llm_provider: str | None = None,
|
||||
|
|
@ -492,9 +492,9 @@ def generate_content_stream(
|
|||
tools: ToolConfigDict | None = None,
|
||||
# Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs.
|
||||
# The extra values given here take precedence over values defined on the client or passed to this method.
|
||||
extra_headers: dict[str, Any] | None = None,
|
||||
extra_query: dict[str, Any] | None = None,
|
||||
extra_body: dict[str, Any] | None = None,
|
||||
extra_headers: dict[str, object] | None = None,
|
||||
extra_query: dict[str, object] | None = None,
|
||||
extra_body: dict[str, object] | None = None,
|
||||
timeout: float | httpx.Timeout | None = None,
|
||||
# LiteLLM specific params,
|
||||
custom_llm_provider: str | None = None,
|
||||
|
|
|
|||
|
|
@ -3,7 +3,7 @@ import contextvars
|
|||
import importlib
|
||||
from collections.abc import Coroutine
|
||||
from functools import partial
|
||||
from typing import TYPE_CHECKING, Any, Final, Literal, Optional, cast, overload
|
||||
from typing import TYPE_CHECKING, Final, Literal, Optional, cast, overload
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.images.utils import ImageEditRequestUtils
|
||||
|
|
@ -151,7 +151,7 @@ def image_generation(
|
|||
*,
|
||||
aimg_generation: Literal[True],
|
||||
**kwargs,
|
||||
) -> Coroutine[Any, Any, ImageResponse]:
|
||||
) -> Coroutine[object, object, ImageResponse]:
|
||||
...
|
||||
|
||||
|
||||
|
|
@ -197,7 +197,7 @@ def image_generation(
|
|||
api_version: str | None = None,
|
||||
custom_llm_provider=None,
|
||||
**kwargs,
|
||||
) -> ImageResponse | Coroutine[Any, Any, ImageResponse]:
|
||||
) -> ImageResponse | Coroutine[object, object, ImageResponse]:
|
||||
"""
|
||||
Maps the https://api.openai.com/v1/images/generations endpoint.
|
||||
|
||||
|
|
@ -386,6 +386,8 @@ def image_generation(
|
|||
litellm.LlmProviders.VERTEX_AI,
|
||||
litellm.LlmProviders.OPENROUTER,
|
||||
litellm.LlmProviders.DASHSCOPE,
|
||||
litellm.LlmProviders.QWENCLOUD,
|
||||
litellm.LlmProviders.QWEN_AI_PLATFORM,
|
||||
):
|
||||
if image_generation_config is None:
|
||||
raise ValueError(f"image generation config is not supported for {custom_llm_provider}")
|
||||
|
|
@ -723,14 +725,14 @@ def image_edit(
|
|||
user: str | None = None,
|
||||
# Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs.
|
||||
# The extra values given here take precedence over values defined on the client or passed to this method.
|
||||
extra_headers: dict[str, Any] | None = None,
|
||||
extra_query: dict[str, Any] | None = None,
|
||||
extra_body: dict[str, Any] | None = None,
|
||||
extra_headers: dict[str, object] | None = None,
|
||||
extra_query: dict[str, object] | None = None,
|
||||
extra_body: dict[str, object] | None = None,
|
||||
timeout: float | httpx.Timeout | None = None,
|
||||
# LiteLLM specific params,
|
||||
custom_llm_provider: str | None = None,
|
||||
**kwargs,
|
||||
) -> ImageResponse | Coroutine[Any, Any, ImageResponse]:
|
||||
) -> ImageResponse | Coroutine[object, object, ImageResponse]:
|
||||
"""
|
||||
Maps the image edit functionality, similar to OpenAI's images/edits endpoint.
|
||||
"""
|
||||
|
|
@ -769,7 +771,7 @@ def image_edit(
|
|||
images: Final = image if isinstance(image, list) else ([image] if image is not None else [])
|
||||
|
||||
headers_from_kwargs: Final = kwargs.get("headers")
|
||||
merged_extra_headers: Final[dict[str, Any]] = {}
|
||||
merged_extra_headers: Final[dict[str, object]] = {}
|
||||
if isinstance(headers_from_kwargs, dict):
|
||||
merged_extra_headers.update(headers_from_kwargs)
|
||||
if isinstance(extra_headers, dict):
|
||||
|
|
@ -974,9 +976,9 @@ async def aimage_edit(
|
|||
user: str | None = None,
|
||||
# Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs.
|
||||
# The extra values given here take precedence over values defined on the client or passed to this method.
|
||||
extra_headers: dict[str, Any] | None = None,
|
||||
extra_query: dict[str, Any] | None = None,
|
||||
extra_body: dict[str, Any] | None = None,
|
||||
extra_headers: dict[str, object] | None = None,
|
||||
extra_query: dict[str, object] | None = None,
|
||||
extra_body: dict[str, object] | None = None,
|
||||
timeout: float | httpx.Timeout | None = None,
|
||||
# LiteLLM specific params,
|
||||
custom_llm_provider: str | None = None,
|
||||
|
|
@ -1044,7 +1046,7 @@ async def aimage_edit(
|
|||
)
|
||||
|
||||
|
||||
def __getattr__(name: str) -> Any:
|
||||
def __getattr__(name: str) -> type["ImageEditRequestUtils"]:
|
||||
"""Lazy import handler for images.main module"""
|
||||
if name == "ImageEditRequestUtils":
|
||||
# Lazy load ImageEditRequestUtils to avoid heavy import from images.utils at module load time
|
||||
|
|
|
|||
|
|
@ -68,6 +68,7 @@ from .utils import process_slack_alerting_variables
|
|||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.proxy.db.db_transaction_queue.pod_lock_manager import PodLockManager
|
||||
from litellm.proxy.utils import PrismaClient
|
||||
from litellm.router import Router as _Router
|
||||
|
||||
Router = _Router
|
||||
|
|
@ -545,7 +546,6 @@ class SlackAlerting(CustomBatchLogger):
|
|||
# Get the appropriate budget alert type handler
|
||||
budget_alert_class: Final = get_budget_alert_type(type)
|
||||
_id: Final = budget_alert_class.get_id(user_info)
|
||||
user_info_json: Final = user_info.model_dump(exclude_none=True)
|
||||
user_info_str: Final = self._get_user_info_str(user_info)
|
||||
event_message = budget_alert_class.get_event_message()
|
||||
|
||||
|
|
@ -575,7 +575,22 @@ class SlackAlerting(CustomBatchLogger):
|
|||
webhook_event = WebhookEvent(
|
||||
event=event,
|
||||
event_message=event_message,
|
||||
**user_info_json,
|
||||
spend=user_info.spend,
|
||||
max_budget=user_info.max_budget,
|
||||
soft_budget=user_info.soft_budget,
|
||||
token=user_info.token,
|
||||
customer_id=user_info.customer_id,
|
||||
user_id=user_info.user_id,
|
||||
team_id=user_info.team_id,
|
||||
team_alias=user_info.team_alias,
|
||||
organization_id=user_info.organization_id,
|
||||
user_email=user_info.user_email,
|
||||
key_alias=user_info.key_alias,
|
||||
projected_exceeded_date=user_info.projected_exceeded_date,
|
||||
projected_spend=user_info.projected_spend,
|
||||
event_group=user_info.event_group,
|
||||
alert_emails=user_info.alert_emails,
|
||||
max_budget_alert_emails=user_info.max_budget_alert_emails,
|
||||
)
|
||||
await self.send_alert(
|
||||
message=event_message + "\n\n" + user_info_str,
|
||||
|
|
@ -657,7 +672,7 @@ class SlackAlerting(CustomBatchLogger):
|
|||
"""
|
||||
Create a standard message for a budget alert
|
||||
"""
|
||||
_all_fields_as_dict: Final = user_info.model_dump(exclude_none=True)
|
||||
_all_fields_as_dict: Final[dict[str, object]] = user_info.model_dump(exclude_none=True)
|
||||
_all_fields_as_dict.pop("token")
|
||||
msg = ""
|
||||
for k, v in _all_fields_as_dict.items():
|
||||
|
|
@ -1006,7 +1021,7 @@ class SlackAlerting(CustomBatchLogger):
|
|||
except Exception:
|
||||
pass
|
||||
|
||||
async def model_added_alert(self, model_name: str, litellm_model_name: str, passed_model_info: Any):
|
||||
async def model_added_alert(self, model_name: str, litellm_model_name: str, passed_model_info: object):
|
||||
base_model_from_user: Final = getattr(passed_model_info, "base_model", None)
|
||||
model_info = {}
|
||||
base_model = ""
|
||||
|
|
@ -1485,9 +1500,9 @@ Model Info:
|
|||
elif self.default_webhook_url is not None:
|
||||
_digest_webhook = self.default_webhook_url
|
||||
else:
|
||||
_digest_webhook = os.getenv("SLACK_WEBHOOK_URL", None)
|
||||
_digest_webhook = os.getenv("SLACK_WEBHOOK_URL") or os.getenv("ALERTING_WEBHOOK_URL")
|
||||
if _digest_webhook is None:
|
||||
raise ValueError("Missing SLACK_WEBHOOK_URL from environment")
|
||||
raise ValueError("Missing SLACK_WEBHOOK_URL / ALERTING_WEBHOOK_URL from environment")
|
||||
|
||||
digest_key: Final = f"{alert_type_name_str}:{request_model or ''}:{api_base or ''}"
|
||||
|
||||
|
|
@ -1516,10 +1531,10 @@ Model Info:
|
|||
elif self.default_webhook_url is not None:
|
||||
slack_webhook_url = self.default_webhook_url
|
||||
else:
|
||||
slack_webhook_url = os.getenv("SLACK_WEBHOOK_URL", None)
|
||||
slack_webhook_url = os.getenv("SLACK_WEBHOOK_URL") or os.getenv("ALERTING_WEBHOOK_URL")
|
||||
|
||||
if slack_webhook_url is None:
|
||||
raise ValueError("Missing SLACK_WEBHOOK_URL from environment")
|
||||
raise ValueError("Missing SLACK_WEBHOOK_URL / ALERTING_WEBHOOK_URL from environment")
|
||||
payload: Final = {"text": formatted_message}
|
||||
headers: Final = {"Content-type": "application/json"}
|
||||
|
||||
|
|
@ -1930,6 +1945,69 @@ Model Info:
|
|||
except Exception as e:
|
||||
verbose_proxy_logger.exception("Error sending weekly spend report %s", e)
|
||||
|
||||
async def send_user_spend_alerts(self, prisma_client: "PrismaClient | None" = None) -> None:
|
||||
"""Check per-user daily/monthly spend thresholds and spend anomalies, alerting once per user per period."""
|
||||
if self.alerting is None or "slack" not in self.alerting:
|
||||
return
|
||||
|
||||
thresholds_enabled: Final = AlertType.user_spend_thresholds in self.alert_types
|
||||
anomalies_enabled: Final = AlertType.user_spend_anomalies in self.alert_types
|
||||
if not thresholds_enabled and not anomalies_enabled:
|
||||
return
|
||||
|
||||
if prisma_client is None:
|
||||
from litellm.proxy.proxy_server import prisma_client as global_prisma_client
|
||||
|
||||
prisma_client = global_prisma_client # rebind-ok: fall back to the proxy's global client
|
||||
if prisma_client is None:
|
||||
return
|
||||
|
||||
from litellm.integrations.SlackAlerting.user_spend_alerts import (
|
||||
evaluate_user_spend,
|
||||
fetch_user_spend_rows,
|
||||
)
|
||||
|
||||
try:
|
||||
today: Final = datetime.datetime.now(datetime.timezone.utc).date()
|
||||
rows: Final = await fetch_user_spend_rows(
|
||||
prisma_client=prisma_client,
|
||||
today=today,
|
||||
baseline_days=self.alerting_args.spend_anomaly_baseline_days,
|
||||
)
|
||||
all_events: Final = tuple(
|
||||
event
|
||||
for row in rows
|
||||
for event in evaluate_user_spend(
|
||||
row=row,
|
||||
args=self.alerting_args,
|
||||
today=today,
|
||||
thresholds_enabled=thresholds_enabled,
|
||||
anomalies_enabled=anomalies_enabled,
|
||||
)
|
||||
)
|
||||
cached_flags: Final = await asyncio.gather(
|
||||
*(self.internal_usage_cache.async_get_cache(key=event.cache_key) for event in all_events)
|
||||
)
|
||||
new_events: Final = tuple(event for event, cached in zip(all_events, cached_flags) if not cached)
|
||||
for alert_type in (AlertType.user_spend_thresholds, AlertType.user_spend_anomalies):
|
||||
typed_events = tuple(event for event in new_events if event.alert_type == alert_type)
|
||||
if not typed_events:
|
||||
continue
|
||||
await self.send_alert(
|
||||
message="\n\n".join(event.message for event in typed_events),
|
||||
level="High",
|
||||
alert_type=alert_type,
|
||||
alerting_metadata={}, # mutable-ok: send_alert takes a dict payload
|
||||
)
|
||||
for event in typed_events:
|
||||
await self.internal_usage_cache.async_set_cache(
|
||||
key=event.cache_key,
|
||||
value="SENT",
|
||||
ttl=event.cache_ttl,
|
||||
)
|
||||
except Exception as e: # noqa: BLE001 # background job must not crash the scheduler
|
||||
verbose_proxy_logger.exception("Error sending user spend alerts: %s", e)
|
||||
|
||||
async def send_fallback_stats_from_prometheus(self):
|
||||
"""
|
||||
Helper to send fallback statistics from prometheus server -> to slack
|
||||
|
|
@ -1973,7 +2051,7 @@ Model Info:
|
|||
try:
|
||||
message = f"`{event_name}`\n"
|
||||
|
||||
key_event_dict: Final = key_event.model_dump()
|
||||
key_event_dict: Final[dict[str, object]] = key_event.model_dump()
|
||||
|
||||
# Add Created by information first
|
||||
message += "*Action Done by:*\n"
|
||||
|
|
|
|||
139
litellm/integrations/SlackAlerting/user_spend_alerts.py
Normal file
139
litellm/integrations/SlackAlerting/user_spend_alerts.py
Normal file
|
|
@ -0,0 +1,139 @@
|
|||
"""Per-user daily/monthly spend threshold alerts and spend anomaly detection."""
|
||||
|
||||
import datetime
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING, Final, Literal
|
||||
|
||||
from pydantic import TypeAdapter
|
||||
|
||||
from litellm.constants import HOURS_IN_A_DAY
|
||||
from litellm.types.integrations.slack_alerting import AlertType, SlackAlertingArgs
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.proxy.utils import PrismaClient
|
||||
|
||||
DAY_SECONDS: Final = HOURS_IN_A_DAY * 60 * 60
|
||||
MONTHLY_ALERT_TTL_SECONDS: Final = 32 * DAY_SECONDS
|
||||
|
||||
USER_SPEND_QUERY: Final = """
|
||||
SELECT
|
||||
user_id,
|
||||
COALESCE(SUM(spend) FILTER (WHERE date = $1), 0)::float AS daily_spend,
|
||||
COALESCE(SUM(spend) FILTER (WHERE date >= $2), 0)::float AS monthly_spend,
|
||||
COALESCE(SUM(spend) FILTER (WHERE date >= $3 AND date < $1), 0)::float AS baseline_spend
|
||||
FROM "LiteLLM_DailyUserSpend"
|
||||
WHERE date >= LEAST($2, $3) AND user_id IS NOT NULL
|
||||
GROUP BY user_id
|
||||
HAVING COALESCE(SUM(spend) FILTER (WHERE date >= $2), 0) > 0
|
||||
"""
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class UserSpendRow:
|
||||
user_id: str
|
||||
daily_spend: float
|
||||
monthly_spend: float
|
||||
baseline_spend: float
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class UserSpendAlertEvent:
|
||||
kind: Literal["daily_threshold", "monthly_threshold", "anomaly"]
|
||||
alert_type: AlertType
|
||||
message: str
|
||||
cache_key: str
|
||||
cache_ttl: int
|
||||
|
||||
|
||||
USER_SPEND_ROWS_ADAPTER: Final = TypeAdapter(tuple[UserSpendRow, ...])
|
||||
|
||||
|
||||
async def fetch_user_spend_rows(
|
||||
prisma_client: "PrismaClient",
|
||||
today: datetime.date,
|
||||
baseline_days: int,
|
||||
) -> tuple[UserSpendRow, ...]:
|
||||
today_str: Final = today.strftime("%Y-%m-%d")
|
||||
month_start_str: Final = today.replace(day=1).strftime("%Y-%m-%d")
|
||||
baseline_start_str: Final = (today - datetime.timedelta(days=max(baseline_days, 1))).strftime("%Y-%m-%d")
|
||||
raw: Final = await prisma_client.db.query_raw(USER_SPEND_QUERY, today_str, month_start_str, baseline_start_str)
|
||||
return USER_SPEND_ROWS_ADAPTER.validate_python(raw)
|
||||
|
||||
|
||||
def _daily_threshold_event(row: UserSpendRow, args: SlackAlertingArgs, today_str: str) -> UserSpendAlertEvent | None:
|
||||
threshold: Final = args.daily_spend_per_user_threshold
|
||||
if threshold is None or row.daily_spend < threshold:
|
||||
return None
|
||||
return UserSpendAlertEvent(
|
||||
kind="daily_threshold",
|
||||
alert_type=AlertType.user_spend_thresholds,
|
||||
message=(
|
||||
f"User Daily Spend Threshold Crossed:\n"
|
||||
f"User: `{row.user_id}`\n"
|
||||
f"Spend Today: `${row.daily_spend:.2f}`\n"
|
||||
f"Daily Threshold: `${threshold:.2f}`"
|
||||
),
|
||||
cache_key=f"user_spend_alert_daily_{row.user_id}_{today_str}",
|
||||
cache_ttl=DAY_SECONDS,
|
||||
)
|
||||
|
||||
|
||||
def _monthly_threshold_event(row: UserSpendRow, args: SlackAlertingArgs, month_str: str) -> UserSpendAlertEvent | None:
|
||||
threshold: Final = args.monthly_spend_per_user_threshold
|
||||
if threshold is None or row.monthly_spend < threshold:
|
||||
return None
|
||||
return UserSpendAlertEvent(
|
||||
kind="monthly_threshold",
|
||||
alert_type=AlertType.user_spend_thresholds,
|
||||
message=(
|
||||
f"User Monthly Spend Threshold Crossed:\n"
|
||||
f"User: `{row.user_id}`\n"
|
||||
f"Spend This Month: `${row.monthly_spend:.2f}`\n"
|
||||
f"Monthly Threshold: `${threshold:.2f}`"
|
||||
),
|
||||
cache_key=f"user_spend_alert_monthly_{row.user_id}_{month_str}",
|
||||
cache_ttl=MONTHLY_ALERT_TTL_SECONDS,
|
||||
)
|
||||
|
||||
|
||||
def _anomaly_event(row: UserSpendRow, args: SlackAlertingArgs, today_str: str) -> UserSpendAlertEvent | None:
|
||||
if row.daily_spend < args.spend_anomaly_min_spend:
|
||||
return None
|
||||
baseline_daily_avg: Final = row.baseline_spend / args.spend_anomaly_baseline_days
|
||||
if row.baseline_spend > 0 and row.daily_spend <= args.spend_anomaly_multiplier * baseline_daily_avg:
|
||||
return None
|
||||
return UserSpendAlertEvent(
|
||||
kind="anomaly",
|
||||
alert_type=AlertType.user_spend_anomalies,
|
||||
message=(
|
||||
f"User Spend Anomaly Detected:\n"
|
||||
f"User: `{row.user_id}`\n"
|
||||
f"Spend Today: `${row.daily_spend:.2f}`\n"
|
||||
f"Daily Average (last {args.spend_anomaly_baseline_days} days): `${baseline_daily_avg:.2f}`\n"
|
||||
f"Trigger: spend above `{args.spend_anomaly_multiplier}x` the daily average "
|
||||
f"(minimum `${args.spend_anomaly_min_spend:.2f}`)"
|
||||
),
|
||||
cache_key=f"user_spend_alert_anomaly_{row.user_id}_{today_str}",
|
||||
cache_ttl=DAY_SECONDS,
|
||||
)
|
||||
|
||||
|
||||
def evaluate_user_spend(
|
||||
row: UserSpendRow,
|
||||
args: SlackAlertingArgs,
|
||||
today: datetime.date,
|
||||
thresholds_enabled: bool,
|
||||
anomalies_enabled: bool,
|
||||
) -> tuple[UserSpendAlertEvent, ...]:
|
||||
today_str: Final = today.strftime("%Y-%m-%d")
|
||||
month_str: Final = today.strftime("%Y-%m")
|
||||
threshold_events: Final = (
|
||||
(
|
||||
_daily_threshold_event(row=row, args=args, today_str=today_str),
|
||||
_monthly_threshold_event(row=row, args=args, month_str=month_str),
|
||||
)
|
||||
if thresholds_enabled
|
||||
else ()
|
||||
)
|
||||
anomaly_events: Final = (_anomaly_event(row=row, args=args, today_str=today_str),) if anomalies_enabled else ()
|
||||
return tuple(event for event in (*threshold_events, *anomaly_events) if event is not None)
|
||||
|
|
@ -3,10 +3,12 @@ Arize Phoenix prompt manager that integrates with LiteLLM's prompt management sy
|
|||
Fetches prompt versions from Arize Phoenix and provides workspace-based access control.
|
||||
"""
|
||||
|
||||
from typing import Any, Final
|
||||
from collections.abc import Mapping, Sequence
|
||||
from typing import Any, Final, cast
|
||||
|
||||
from jinja2 import DictLoader, select_autoescape
|
||||
from jinja2.sandbox import ImmutableSandboxedEnvironment
|
||||
from typing_extensions import ReadOnly, TypedDict
|
||||
|
||||
from litellm.integrations.custom_prompt_management import CustomPromptManagement
|
||||
from litellm.integrations.prompt_management_base import (
|
||||
|
|
@ -20,6 +22,31 @@ from litellm.types.utils import StandardCallbackDynamicParams
|
|||
from .arize_phoenix_client import ArizePhoenixClient
|
||||
|
||||
|
||||
class ArizePhoenixContentPart(TypedDict, total=False):
|
||||
type: ReadOnly[str]
|
||||
text: ReadOnly[str]
|
||||
|
||||
|
||||
class ArizePhoenixTemplateMessage(TypedDict, total=False):
|
||||
role: ReadOnly[str]
|
||||
content: ReadOnly[Sequence[ArizePhoenixContentPart]]
|
||||
|
||||
|
||||
class ArizePhoenixTemplateBody(TypedDict, total=False):
|
||||
messages: ReadOnly[Sequence[ArizePhoenixTemplateMessage]]
|
||||
|
||||
|
||||
class ArizePhoenixPromptMetadata(TypedDict):
|
||||
model_name: ReadOnly[str | None]
|
||||
model_provider: ReadOnly[str | None]
|
||||
description: ReadOnly[str]
|
||||
template_type: ReadOnly[str | None]
|
||||
template_format: ReadOnly[str]
|
||||
invocation_parameters: ReadOnly[Mapping[str, Mapping[str, object]]]
|
||||
temperature: ReadOnly[float | None]
|
||||
max_tokens: ReadOnly[int | None]
|
||||
|
||||
|
||||
class ArizePhoenixPromptTemplate:
|
||||
"""
|
||||
Represents a prompt template loaded from Arize Phoenix.
|
||||
|
|
@ -28,10 +55,10 @@ class ArizePhoenixPromptTemplate:
|
|||
def __init__(
|
||||
self,
|
||||
template_id: str,
|
||||
messages: list[dict[str, Any]],
|
||||
metadata: dict[str, Any],
|
||||
messages: Sequence[ArizePhoenixTemplateMessage],
|
||||
metadata: ArizePhoenixPromptMetadata,
|
||||
model: str | None = None,
|
||||
):
|
||||
) -> None:
|
||||
self.template_id = template_id
|
||||
self.messages = messages
|
||||
self.metadata = metadata
|
||||
|
|
@ -43,7 +70,7 @@ class ArizePhoenixPromptTemplate:
|
|||
self.description = metadata.get("description", "")
|
||||
self.template_format = metadata.get("template_format", "MUSTACHE")
|
||||
|
||||
def __repr__(self):
|
||||
def __repr__(self) -> str:
|
||||
return f"ArizePhoenixPromptTemplate(id='{self.template_id}', model='{self.model}')"
|
||||
|
||||
|
||||
|
|
@ -109,7 +136,7 @@ class ArizePhoenixTemplateManager:
|
|||
|
||||
def _parse_prompt_data(self, data: dict[str, Any], prompt_version_id: str) -> ArizePhoenixPromptTemplate:
|
||||
"""Parse Arize Phoenix prompt data and extract messages and metadata."""
|
||||
template_data: Final = data.get("template", {})
|
||||
template_data: Final[ArizePhoenixTemplateBody] = data.get("template", {})
|
||||
messages: Final = template_data.get("messages", [])
|
||||
|
||||
# Extract invocation parameters
|
||||
|
|
@ -129,7 +156,7 @@ class ArizePhoenixTemplateManager:
|
|||
break
|
||||
|
||||
# Build metadata dictionary
|
||||
metadata: Final = {
|
||||
metadata: Final[ArizePhoenixPromptMetadata] = {
|
||||
"model_name": data.get("model_name"),
|
||||
"model_provider": data.get("model_provider"),
|
||||
"description": data.get("description", ""),
|
||||
|
|
@ -146,7 +173,9 @@ class ArizePhoenixTemplateManager:
|
|||
metadata=metadata,
|
||||
)
|
||||
|
||||
def render_template(self, template_id: str, variables: dict[str, Any] | None = None) -> list[AllMessageValues]:
|
||||
def render_template(
|
||||
self, template_id: str, variables: Mapping[str, object] | None = None
|
||||
) -> list[AllMessageValues]:
|
||||
"""Render a template with the given variables and return formatted messages."""
|
||||
if template_id not in self.prompts:
|
||||
raise ValueError(f"Template '{template_id}' not found")
|
||||
|
|
@ -174,7 +203,9 @@ class ArizePhoenixTemplateManager:
|
|||
# Combine rendered content
|
||||
final_content = " ".join(rendered_content_parts)
|
||||
|
||||
rendered_messages.append({"role": role, "content": final_content})
|
||||
rendered_messages.append(
|
||||
cast("AllMessageValues", {"role": role, "content": final_content}) # cast-ok: Phoenix roles are OpenAI
|
||||
)
|
||||
|
||||
return rendered_messages
|
||||
|
||||
|
|
@ -243,8 +274,8 @@ class ArizePhoenixPromptManager(CustomPromptManagement):
|
|||
def get_prompt_template(
|
||||
self,
|
||||
prompt_id: str,
|
||||
prompt_variables: dict[str, Any] | None = None,
|
||||
) -> tuple[list[AllMessageValues], dict[str, Any]]:
|
||||
prompt_variables: Mapping[str, object] | None = None,
|
||||
) -> tuple[list[AllMessageValues], dict[str, object]]:
|
||||
"""
|
||||
Get a prompt template and render it with variables.
|
||||
|
||||
|
|
@ -263,7 +294,7 @@ class ArizePhoenixPromptManager(CustomPromptManagement):
|
|||
rendered_messages: Final = self.prompt_manager.render_template(prompt_id, prompt_variables or {})
|
||||
|
||||
# Extract metadata
|
||||
metadata: Final = {
|
||||
metadata: Final[dict[str, object]] = {
|
||||
"model": template.model,
|
||||
"temperature": template.temperature,
|
||||
"max_tokens": template.max_tokens,
|
||||
|
|
@ -271,7 +302,7 @@ class ArizePhoenixPromptManager(CustomPromptManagement):
|
|||
|
||||
# Add additional invocation parameters
|
||||
invocation_params: Final = template.invocation_parameters
|
||||
provider_params = {}
|
||||
provider_params: Mapping[str, object] = {}
|
||||
|
||||
if "openai" in invocation_params:
|
||||
provider_params = invocation_params["openai"]
|
||||
|
|
@ -289,12 +320,12 @@ class ArizePhoenixPromptManager(CustomPromptManagement):
|
|||
self,
|
||||
user_id: str | None,
|
||||
messages: list[AllMessageValues],
|
||||
function_call: dict[str, Any] | str | None = None,
|
||||
litellm_params: dict[str, Any] | None = None,
|
||||
function_call: dict[str, object] | str | None = None,
|
||||
litellm_params: dict[str, object] | None = None,
|
||||
prompt_id: str | None = None,
|
||||
prompt_variables: dict[str, Any] | None = None,
|
||||
prompt_variables: dict[str, object] | None = None,
|
||||
**kwargs,
|
||||
) -> tuple[list[AllMessageValues], dict[str, Any] | None]:
|
||||
) -> tuple[list[AllMessageValues], dict[str, object] | None]:
|
||||
"""
|
||||
Pre-call hook that processes the prompt template before making the LLM call.
|
||||
"""
|
||||
|
|
@ -335,9 +366,9 @@ class ArizePhoenixPromptManager(CustomPromptManagement):
|
|||
|
||||
except Exception as e:
|
||||
# Log error but don't fail the call
|
||||
import litellm
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
|
||||
litellm._logging.verbose_proxy_logger.error("Error in Arize Phoenix prompt pre_call_hook: %s", e)
|
||||
verbose_proxy_logger.error("Error in Arize Phoenix prompt pre_call_hook: %s", e)
|
||||
return messages, litellm_params
|
||||
|
||||
def get_available_prompts(self) -> list[str]:
|
||||
|
|
@ -393,7 +424,8 @@ class ArizePhoenixPromptManager(CustomPromptManagement):
|
|||
rendered_messages, prompt_metadata = self.get_prompt_template(prompt_id, prompt_variables)
|
||||
|
||||
# Extract model from metadata (if specified)
|
||||
template_model: Final = prompt_metadata.get("model")
|
||||
raw_template_model: Final = prompt_metadata.get("model")
|
||||
template_model: Final = raw_template_model if isinstance(raw_template_model, str) else None
|
||||
|
||||
# Extract optional parameters from metadata
|
||||
optional_params: Final = {}
|
||||
|
|
|
|||
|
|
@ -163,15 +163,11 @@ class BitBucketClient:
|
|||
response.raise_for_status()
|
||||
|
||||
data: Final[BitBucketSrcListing] = response.json()
|
||||
files: Final[list[str]] = []
|
||||
|
||||
for item in data.get("values", []):
|
||||
if item.get("type") == "commit_file":
|
||||
file_path = item.get("path", "")
|
||||
if file_path.endswith(file_extension):
|
||||
files.append(file_path)
|
||||
|
||||
return files
|
||||
return [
|
||||
file_path
|
||||
for item in data.get("values", [])
|
||||
if item.get("type") == "commit_file" and (file_path := item.get("path", "")).endswith(file_extension)
|
||||
]
|
||||
|
||||
except Exception as e:
|
||||
# Check if it's an HTTP error
|
||||
|
|
|
|||
|
|
@ -3,6 +3,7 @@ BitBucket prompt manager that integrates with LiteLLM's prompt management system
|
|||
Fetches .prompt files from BitBucket repositories and provides team-based access control.
|
||||
"""
|
||||
|
||||
from collections.abc import Mapping
|
||||
from typing import TYPE_CHECKING, Any, Final
|
||||
|
||||
from jinja2 import DictLoader, select_autoescape
|
||||
|
|
@ -65,7 +66,7 @@ class BitBucketTemplateManager:
|
|||
|
||||
def __init__(
|
||||
self,
|
||||
bitbucket_config: dict[str, Any],
|
||||
bitbucket_config: Mapping[str, object],
|
||||
prompt_id: str | None = None,
|
||||
):
|
||||
self.bitbucket_config = bitbucket_config
|
||||
|
|
@ -123,7 +124,7 @@ class BitBucketTemplateManager:
|
|||
template_content = content
|
||||
|
||||
# Parse YAML frontmatter
|
||||
metadata: dict[str, Any] = {}
|
||||
metadata: dict[str, object] = {}
|
||||
if frontmatter_str:
|
||||
try:
|
||||
import yaml
|
||||
|
|
@ -141,9 +142,9 @@ class BitBucketTemplateManager:
|
|||
metadata=metadata,
|
||||
)
|
||||
|
||||
def _parse_yaml_basic(self, yaml_str: str) -> dict[str, Any]:
|
||||
def _parse_yaml_basic(self, yaml_str: str) -> dict[str, object]:
|
||||
"""Basic YAML parser for simple cases when PyYAML is not available."""
|
||||
result: Final[dict[str, Any]] = {}
|
||||
result: Final[dict[str, object]] = {}
|
||||
for line in yaml_str.split("\n"):
|
||||
line = line.strip()
|
||||
if ":" in line and not line.startswith("#"):
|
||||
|
|
@ -162,7 +163,7 @@ class BitBucketTemplateManager:
|
|||
result[key] = value.strip("\"'")
|
||||
return result
|
||||
|
||||
def render_template(self, template_id: str, variables: dict[str, Any] | None = None) -> str:
|
||||
def render_template(self, template_id: str, variables: Mapping[str, object] | None = None) -> str:
|
||||
"""Render a template with the given variables."""
|
||||
if template_id not in self.prompts:
|
||||
raise ValueError(f"Template '{template_id}' not found")
|
||||
|
|
@ -209,7 +210,7 @@ class BitBucketPromptManager(CustomPromptManagement):
|
|||
|
||||
def __init__(
|
||||
self,
|
||||
bitbucket_config: dict[str, Any],
|
||||
bitbucket_config: Mapping[str, object],
|
||||
prompt_id: str | None = None,
|
||||
):
|
||||
self.bitbucket_config = bitbucket_config
|
||||
|
|
@ -234,7 +235,7 @@ class BitBucketPromptManager(CustomPromptManagement):
|
|||
def get_prompt_template(
|
||||
self,
|
||||
prompt_id: str,
|
||||
prompt_variables: dict[str, Any] | None = None,
|
||||
prompt_variables: Mapping[str, object] | None = None,
|
||||
) -> tuple[str, dict[str, Any]]:
|
||||
"""
|
||||
Get a prompt template and render it with variables.
|
||||
|
|
@ -267,12 +268,12 @@ class BitBucketPromptManager(CustomPromptManagement):
|
|||
self,
|
||||
user_id: str | None,
|
||||
messages: list[AllMessageValues],
|
||||
function_call: dict[str, Any] | str | None = None,
|
||||
litellm_params: dict[str, Any] | None = None,
|
||||
function_call: Mapping[str, object] | str | None = None,
|
||||
litellm_params: dict[str, object] | None = None,
|
||||
prompt_id: str | None = None,
|
||||
prompt_variables: dict[str, Any] | None = None,
|
||||
prompt_variables: Mapping[str, object] | None = None,
|
||||
**kwargs,
|
||||
) -> tuple[list[AllMessageValues], dict[str, Any] | None]:
|
||||
) -> tuple[list[AllMessageValues], dict[str, object] | None]:
|
||||
"""
|
||||
Pre-call hook that processes the prompt template before making the LLM call.
|
||||
"""
|
||||
|
|
@ -316,9 +317,9 @@ class BitBucketPromptManager(CustomPromptManagement):
|
|||
|
||||
except Exception as e:
|
||||
# Log error but don't fail the call
|
||||
import litellm
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
|
||||
litellm._logging.verbose_proxy_logger.error("Error in BitBucket prompt pre_call_hook: %s", e)
|
||||
verbose_proxy_logger.error("Error in BitBucket prompt pre_call_hook: %s", e)
|
||||
return messages, litellm_params
|
||||
|
||||
def _parse_prompt_to_messages(self, prompt_content: str) -> list[AllMessageValues]:
|
||||
|
|
@ -384,14 +385,14 @@ class BitBucketPromptManager(CustomPromptManagement):
|
|||
def post_call_hook(
|
||||
self,
|
||||
user_id: str | None,
|
||||
response: Any,
|
||||
response: object,
|
||||
input_messages: list[AllMessageValues],
|
||||
function_call: dict[str, Any] | str | None = None,
|
||||
litellm_params: dict[str, Any] | None = None,
|
||||
function_call: Mapping[str, object] | str | None = None,
|
||||
litellm_params: Mapping[str, object] | None = None,
|
||||
prompt_id: str | None = None,
|
||||
prompt_variables: dict[str, Any] | None = None,
|
||||
prompt_variables: Mapping[str, object] | None = None,
|
||||
**kwargs,
|
||||
) -> Any:
|
||||
) -> object:
|
||||
"""
|
||||
Post-call hook for any post-processing after the LLM call.
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -19,14 +19,29 @@
|
|||
"""Transform LiteLLM data to CloudZero AnyCost CBF format."""
|
||||
|
||||
from datetime import datetime
|
||||
from typing import Any, Final
|
||||
from typing import Final, SupportsFloat, SupportsIndex, SupportsInt
|
||||
|
||||
import polars as pl
|
||||
from typing_extensions import Buffer
|
||||
|
||||
from ...types.integrations.cloudzero import CBFRecord
|
||||
from .cz_resource_names import CZEntityType, CZRNGenerator
|
||||
|
||||
|
||||
def _as_int(value: object) -> int:
|
||||
"""The integer form of a spend table cell, computed the way :func:`int` computes it."""
|
||||
if isinstance(value, (str, Buffer, SupportsInt, SupportsIndex)):
|
||||
return int(value)
|
||||
raise TypeError(f"int() argument must be a string or a number, not {type(value).__name__!r}")
|
||||
|
||||
|
||||
def _as_float(value: object) -> float:
|
||||
"""The floating point form of a spend table cell, computed the way :func:`float` computes it."""
|
||||
if isinstance(value, (str, Buffer, SupportsFloat, SupportsIndex)):
|
||||
return float(value)
|
||||
raise TypeError(f"float() argument must be a string or a number, not {type(value).__name__!r}")
|
||||
|
||||
|
||||
class CBFTransformer:
|
||||
"""Transform LiteLLM usage data to CloudZero Billing Format (CBF)."""
|
||||
|
||||
|
|
@ -82,15 +97,15 @@ class CBFTransformer:
|
|||
|
||||
return pl.DataFrame(cbf_data)
|
||||
|
||||
def _create_cbf_record(self, row: dict[str, Any]) -> CBFRecord:
|
||||
def _create_cbf_record(self, row: dict[str, object]) -> CBFRecord:
|
||||
"""Create a single CBF record from LiteLLM daily spend row."""
|
||||
|
||||
# Parse date (daily spend tables use date strings like '2025-04-19')
|
||||
usage_date: Final = self._parse_date(row.get("date"))
|
||||
|
||||
# Calculate total tokens
|
||||
prompt_tokens: Final = int(row.get("prompt_tokens", 0))
|
||||
completion_tokens: Final = int(row.get("completion_tokens", 0))
|
||||
prompt_tokens: Final = _as_int(row.get("prompt_tokens", 0))
|
||||
completion_tokens: Final = _as_int(row.get("completion_tokens", 0))
|
||||
total_tokens: Final = prompt_tokens + completion_tokens
|
||||
|
||||
# Create CloudZero Resource Name (CZRN) as resource_id
|
||||
|
|
@ -154,7 +169,7 @@ class CBFTransformer:
|
|||
"time/usage_start": (
|
||||
usage_date.isoformat() if usage_date else None
|
||||
), # Required: ISO-formatted UTC datetime
|
||||
"cost/cost": float(row.get("spend", 0.0)), # Required: billed cost
|
||||
"cost/cost": _as_float(row.get("spend", 0.0)), # Required: billed cost
|
||||
"resource/id": resource_id, # CZRN (CloudZero Resource Name)
|
||||
# Usage metrics for token consumption
|
||||
"usage/amount": total_tokens, # Numeric value of tokens consumed
|
||||
|
|
@ -187,7 +202,7 @@ class CBFTransformer:
|
|||
|
||||
return CBFRecord(cbf_record)
|
||||
|
||||
def _parse_date(self, date_str) -> datetime | None:
|
||||
def _parse_date(self, date_str: object) -> datetime | None:
|
||||
"""Parse date string from daily spend tables (e.g., '2025-04-19')."""
|
||||
if date_str is None:
|
||||
return None
|
||||
|
|
|
|||
|
|
@ -7,7 +7,10 @@ litellm_content_retrieve tool calls server-side via the typed agentic loop plan.
|
|||
|
||||
import time
|
||||
import uuid
|
||||
from typing import TYPE_CHECKING, Any, ClassVar, Final, cast
|
||||
from collections.abc import Mapping, Sequence
|
||||
from typing import Any, ClassVar, Final, Protocol, cast
|
||||
|
||||
from typing_extensions import ReadOnly, TypedDict
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.compression import compress
|
||||
|
|
@ -22,13 +25,23 @@ from litellm.types.integrations.custom_logger import (
|
|||
)
|
||||
from litellm.types.utils import CallTypes
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
|
||||
LITELLM_CONTENT_RETRIEVE_TOOL_NAME: Final = "litellm_content_retrieve"
|
||||
_CACHE_TTL_SECONDS: Final = 15 * 60
|
||||
|
||||
|
||||
class _AgenticLoopParams(TypedDict, total=False):
|
||||
"""The ``agentic_loop_params`` entry the agentic loop driver records on the logging object."""
|
||||
|
||||
model: ReadOnly[str]
|
||||
|
||||
|
||||
class _AgenticLoopLoggingObj(Protocol):
|
||||
"""Logging object view exposing the untyped call details this handler reads."""
|
||||
|
||||
@property
|
||||
def model_call_details(self) -> Mapping[str, _AgenticLoopParams]: ...
|
||||
|
||||
|
||||
def _compression_savings_from_counts(
|
||||
original_tokens: object, compressed_tokens: object
|
||||
) -> CompressionSavingsMetadata | None:
|
||||
|
|
@ -83,7 +96,7 @@ class CompressionInterceptionLogger(CustomLogger):
|
|||
compression_trigger: int = 200_000,
|
||||
compression_target: int | None = None,
|
||||
embedding_model: str | None = None,
|
||||
embedding_model_params: dict[str, Any] | None = None,
|
||||
embedding_model_params: dict[str, object] | None = None,
|
||||
):
|
||||
super().__init__()
|
||||
self.enabled = enabled
|
||||
|
|
@ -106,7 +119,7 @@ class CompressionInterceptionLogger(CustomLogger):
|
|||
@staticmethod
|
||||
def initialize_from_proxy_config(
|
||||
litellm_settings: dict[str, Any],
|
||||
callback_specific_params: dict[str, Any],
|
||||
callback_specific_params: Mapping[str, object],
|
||||
) -> "CompressionInterceptionLogger":
|
||||
compression_params: CompressionInterceptionConfig = {}
|
||||
if "compression_interception_params" in litellm_settings:
|
||||
|
|
@ -120,7 +133,9 @@ class CompressionInterceptionLogger(CustomLogger):
|
|||
)
|
||||
return CompressionInterceptionLogger.from_config_yaml(compression_params)
|
||||
|
||||
async def async_pre_call_deployment_hook(self, kwargs: dict[str, Any], call_type: CallTypes | None) -> dict | None:
|
||||
async def async_pre_call_deployment_hook(
|
||||
self, kwargs: dict[str, Any], call_type: CallTypes | None
|
||||
) -> dict[str, object] | None:
|
||||
if not self.enabled:
|
||||
return None
|
||||
if call_type is not None and call_type != CallTypes.anthropic_messages:
|
||||
|
|
@ -150,7 +165,7 @@ class CompressionInterceptionLogger(CustomLogger):
|
|||
|
||||
cache: Final = cast(dict[str, str], compressed.get("cache", {}))
|
||||
skip_reason: Final = cast(str | None, compressed.get("compression_skipped_reason"))
|
||||
compressed_tools: Final = cast(list[dict[str, Any]], compressed.get("tools", []))
|
||||
compressed_tools: Final = cast(list[dict[str, object]], compressed.get("tools", []))
|
||||
|
||||
# Only mutate kwargs when compression actually produced a result.
|
||||
# If compression was a no-op (below trigger, invalid tool sequence, etc.),
|
||||
|
|
@ -161,7 +176,7 @@ class CompressionInterceptionLogger(CustomLogger):
|
|||
kwargs["messages"] = compressed["messages"]
|
||||
if compressed_tools:
|
||||
kwargs["tools"] = self._merge_tools(
|
||||
existing_tools=cast(list[dict[str, Any]] | None, kwargs.get("tools")),
|
||||
existing_tools=cast(list[dict[str, object]] | None, kwargs.get("tools")),
|
||||
compressed_tools=compressed_tools,
|
||||
)
|
||||
call_id = cast(str | None, kwargs.get("litellm_call_id"))
|
||||
|
|
@ -194,14 +209,14 @@ class CompressionInterceptionLogger(CustomLogger):
|
|||
|
||||
async def async_should_run_agentic_loop(
|
||||
self,
|
||||
response: Any,
|
||||
response: object,
|
||||
model: str,
|
||||
messages: list[dict],
|
||||
tools: list[dict] | None,
|
||||
messages: Sequence[Mapping[str, object]],
|
||||
tools: Sequence[Mapping[str, object]] | None,
|
||||
stream: bool,
|
||||
custom_llm_provider: str,
|
||||
kwargs: dict,
|
||||
) -> tuple[bool, dict]:
|
||||
kwargs: Mapping[str, object],
|
||||
) -> tuple[bool, dict[str, object]]:
|
||||
if not self.enabled:
|
||||
return False, {}
|
||||
if not self._has_retrieval_tool(tools):
|
||||
|
|
@ -219,19 +234,19 @@ class CompressionInterceptionLogger(CustomLogger):
|
|||
|
||||
async def async_build_agentic_loop_plan(
|
||||
self,
|
||||
tools: dict,
|
||||
tools: Mapping[str, object],
|
||||
model: str,
|
||||
messages: list[dict],
|
||||
response: Any,
|
||||
anthropic_messages_provider_config: Any,
|
||||
anthropic_messages_optional_request_params: dict,
|
||||
logging_obj: "LiteLLMLoggingObj | None",
|
||||
messages: list[dict[str, object]],
|
||||
response: object,
|
||||
anthropic_messages_provider_config: object,
|
||||
anthropic_messages_optional_request_params: Mapping[str, object],
|
||||
logging_obj: _AgenticLoopLoggingObj | None,
|
||||
stream: bool,
|
||||
kwargs: dict,
|
||||
kwargs: Mapping[str, object],
|
||||
) -> AgenticLoopPlan:
|
||||
self._prune_expired_cache()
|
||||
tool_calls: Final = cast(list[dict[str, Any]], tools.get("tool_calls", []))
|
||||
thinking_blocks: Final = cast(list[dict[str, Any]], tools.get("thinking_blocks", []))
|
||||
tool_calls: Final = cast(list[dict[str, object]], tools.get("tool_calls", []))
|
||||
thinking_blocks: Final = cast(list[dict[str, object]], tools.get("thinking_blocks", []))
|
||||
|
||||
call_id: Final = self._resolve_call_id(logging_obj=logging_obj, kwargs=kwargs)
|
||||
cache: Final = self._get_cache(call_id=call_id)
|
||||
|
|
@ -274,7 +289,7 @@ class CompressionInterceptionLogger(CustomLogger):
|
|||
full_model_name = model
|
||||
if logging_obj is not None:
|
||||
agentic_params: Final = logging_obj.model_call_details.get("agentic_loop_params", {})
|
||||
full_model_name = cast(str, agentic_params.get("model", model))
|
||||
full_model_name = agentic_params.get("model", model)
|
||||
|
||||
request_patch: Final = AgenticLoopRequestPatch(
|
||||
model=full_model_name,
|
||||
|
|
@ -309,15 +324,15 @@ class CompressionInterceptionLogger(CustomLogger):
|
|||
return {}
|
||||
return cache_entry[0]
|
||||
|
||||
def _resolve_call_id(self, logging_obj: Any, kwargs: dict[str, Any]) -> str | None:
|
||||
def _resolve_call_id(self, logging_obj: _AgenticLoopLoggingObj | None, kwargs: Mapping[str, object]) -> str | None:
|
||||
if logging_obj is not None:
|
||||
logging_call_id: Final = getattr(logging_obj, "litellm_call_id", None)
|
||||
if isinstance(logging_call_id, str) and logging_call_id:
|
||||
return logging_call_id
|
||||
kwargs_call_id: Final = kwargs.get("litellm_call_id")
|
||||
return cast(str | None, kwargs_call_id if isinstance(kwargs_call_id, str) else None)
|
||||
return kwargs_call_id if isinstance(kwargs_call_id, str) else None
|
||||
|
||||
def _resolve_retrieval_content(self, tool_call: dict[str, Any], cache: dict[str, str]) -> str:
|
||||
def _resolve_retrieval_content(self, tool_call: Mapping[str, object], cache: Mapping[str, str]) -> str:
|
||||
raw_input: Final = tool_call.get("input", {})
|
||||
key = ""
|
||||
if isinstance(raw_input, dict):
|
||||
|
|
@ -328,7 +343,9 @@ class CompressionInterceptionLogger(CustomLogger):
|
|||
return cache[key]
|
||||
return f"[compressed content key '{key}' not found]"
|
||||
|
||||
def _extract_retrieval_tool_calls(self, response: Any) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]:
|
||||
def _extract_retrieval_tool_calls(
|
||||
self, response: object
|
||||
) -> tuple[list[dict[str, object]], list[dict[str, object]]]:
|
||||
if isinstance(response, dict):
|
||||
content = response.get("content", [])
|
||||
else:
|
||||
|
|
@ -337,8 +354,8 @@ class CompressionInterceptionLogger(CustomLogger):
|
|||
if not isinstance(content, list):
|
||||
return [], []
|
||||
|
||||
tool_calls: Final[list[dict[str, Any]]] = []
|
||||
thinking_blocks: Final[list[dict[str, Any]]] = []
|
||||
tool_calls: Final[list[dict[str, object]]] = []
|
||||
thinking_blocks: Final[list[dict[str, object]]] = []
|
||||
|
||||
for block in content:
|
||||
if isinstance(block, dict):
|
||||
|
|
@ -385,13 +402,13 @@ class CompressionInterceptionLogger(CustomLogger):
|
|||
|
||||
return tool_calls, thinking_blocks
|
||||
|
||||
def _prepare_followup_kwargs(self, kwargs: dict[str, Any]) -> dict[str, Any]:
|
||||
def _prepare_followup_kwargs(self, kwargs: Mapping[str, object]) -> dict[str, object]:
|
||||
internal_keys: Final = {"litellm_logging_obj"}
|
||||
return {
|
||||
k: v for k, v in kwargs.items() if not k.startswith("_compression_interception") and k not in internal_keys
|
||||
}
|
||||
|
||||
def _has_retrieval_tool(self, tools: Any) -> bool:
|
||||
def _has_retrieval_tool(self, tools: object) -> bool:
|
||||
if not isinstance(tools, list):
|
||||
return False
|
||||
for tool in tools:
|
||||
|
|
@ -407,9 +424,9 @@ class CompressionInterceptionLogger(CustomLogger):
|
|||
|
||||
def _merge_tools(
|
||||
self,
|
||||
existing_tools: list[dict[str, Any]] | None,
|
||||
compressed_tools: list[dict[str, Any]],
|
||||
) -> list[dict[str, Any]]:
|
||||
existing_tools: Sequence[Mapping[str, object]] | None,
|
||||
compressed_tools: Sequence[Mapping[str, object]],
|
||||
) -> list[Mapping[str, object]]:
|
||||
merged: Final = list(existing_tools or [])
|
||||
if self._has_retrieval_tool(merged):
|
||||
return merged
|
||||
|
|
|
|||
|
|
@ -1,7 +1,9 @@
|
|||
import contextvars
|
||||
import copy
|
||||
import hashlib
|
||||
import os
|
||||
import secrets
|
||||
from collections.abc import Mapping
|
||||
from datetime import datetime
|
||||
from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, Optional, get_args
|
||||
|
||||
|
|
@ -38,6 +40,7 @@ except ImportError:
|
|||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation
|
||||
dc: Final = DualCache()
|
||||
|
||||
|
||||
|
|
@ -227,13 +230,13 @@ class CustomGuardrail(CustomLogger):
|
|||
)
|
||||
super().__init__(**kwargs)
|
||||
|
||||
def render_violation_message(self, default: str, context: dict[str, Any] | None = None) -> str:
|
||||
def render_violation_message(self, default: str, context: Mapping[str, object] | None = None) -> str:
|
||||
"""Return a custom violation message if template is configured."""
|
||||
|
||||
if not self.violation_message_template:
|
||||
return default
|
||||
|
||||
format_context: Final[dict[str, Any]] = {"default_message": default}
|
||||
format_context: Final[dict[str, object]] = {"default_message": default}
|
||||
if context:
|
||||
format_context.update(context)
|
||||
try:
|
||||
|
|
@ -661,7 +664,7 @@ class CustomGuardrail(CustomLogger):
|
|||
value: Final = self._get_admin_metadata(data).get("opted_out_global_guardrails")
|
||||
return value if isinstance(value, list) else []
|
||||
|
||||
def _is_valid_response_type(self, result: Any) -> bool:
|
||||
def _is_valid_response_type(self, result: object) -> bool:
|
||||
"""
|
||||
Check if result is a valid LLMResponseTypes instance.
|
||||
|
||||
|
|
@ -722,7 +725,7 @@ class CustomGuardrail(CustomLogger):
|
|||
return None
|
||||
return f"{_PRE_CALL_EXECUTED_TOKEN}:{name}"
|
||||
|
||||
def mark_pre_call_hook_ran(self, data: dict[str, Any]) -> None:
|
||||
def mark_pre_call_hook_ran(self, data: dict[str, object]) -> None:
|
||||
"""
|
||||
Record that this guardrail's ``async_pre_call_hook`` already ran for this
|
||||
request, so the deployment-level hook does not run it a second time.
|
||||
|
|
@ -747,7 +750,7 @@ class CustomGuardrail(CustomLogger):
|
|||
return
|
||||
data["metadata"] = {PRE_CALL_EXECUTED_GUARDRAILS_KEY: [marker]}
|
||||
|
||||
def _pre_call_hook_already_ran(self, data: dict[str, Any]) -> bool:
|
||||
def _pre_call_hook_already_ran(self, data: dict[str, object]) -> bool:
|
||||
marker: Final = self._pre_call_marker()
|
||||
if marker is None:
|
||||
return False
|
||||
|
|
@ -851,6 +854,69 @@ class CustomGuardrail(CustomLogger):
|
|||
|
||||
return result
|
||||
|
||||
async def async_logging_hook(
|
||||
self,
|
||||
kwargs: dict, # mutable-ok: CustomLogger.async_logging_hook contract
|
||||
result: object,
|
||||
call_type: str,
|
||||
) -> tuple[dict, object]: # mutable-ok: CustomLogger.async_logging_hook contract
|
||||
"""logging_only: run apply_guardrail on copies of the logged request/response and record the verdict."""
|
||||
from litellm.llms import get_guardrail_translation_mapping
|
||||
|
||||
if not self.uses_apply_guardrail_interface() or self.use_native_lifecycle_hooks:
|
||||
return kwargs, result
|
||||
try:
|
||||
translation: Final = get_guardrail_translation_mapping(CallTypes(call_type))()
|
||||
except ValueError:
|
||||
verbose_logger.debug(
|
||||
"Guardrail %s: no guardrail translation for call_type=%s, skipping logging_only scan",
|
||||
self.guardrail_name,
|
||||
call_type,
|
||||
)
|
||||
return kwargs, result
|
||||
litellm_params: Final = kwargs.get("litellm_params") or {}
|
||||
scratch_metadata: Final = {
|
||||
key: value
|
||||
for key, value in (litellm_params.get("metadata") or {}).items()
|
||||
if key != "standard_logging_guardrail_information"
|
||||
}
|
||||
try:
|
||||
await self._scan_logged_call(kwargs, result, translation, scratch_metadata)
|
||||
except Exception as e:
|
||||
verbose_logger.warning("Guardrail %s: logging_only scan raised: %s", self.guardrail_name, e)
|
||||
recorded: Final = scratch_metadata.get("standard_logging_guardrail_information")
|
||||
standard_logging_object: Final = kwargs.get("standard_logging_object")
|
||||
if not recorded or not isinstance(standard_logging_object, dict):
|
||||
return kwargs, result
|
||||
entries: Final = recorded if isinstance(recorded, list) else [recorded]
|
||||
existing: Final = standard_logging_object.get("guardrail_information") or []
|
||||
return {
|
||||
**kwargs,
|
||||
"standard_logging_object": {**standard_logging_object, "guardrail_information": [*existing, *entries]},
|
||||
}, result
|
||||
|
||||
async def _scan_logged_call(
|
||||
self,
|
||||
kwargs: dict, # mutable-ok: CustomLogger.async_logging_hook contract
|
||||
result: object,
|
||||
translation: "BaseTranslation",
|
||||
scratch_metadata: dict, # mutable-ok: apply_guardrail records its verdict into request metadata
|
||||
) -> None:
|
||||
optional_params: Final = kwargs.get("optional_params") or {}
|
||||
scratch_input: Final = copy.deepcopy(kwargs.get("messages") or kwargs.get("input"))
|
||||
scratch_request: Final = {
|
||||
"model": kwargs.get("model"),
|
||||
"messages": scratch_input,
|
||||
"input": scratch_input,
|
||||
"tools": copy.deepcopy(optional_params.get("tools")),
|
||||
"litellm_call_id": kwargs.get("litellm_call_id"),
|
||||
"metadata": scratch_metadata,
|
||||
}
|
||||
await translation.process_input_messages(data=scratch_request, guardrail_to_apply=self)
|
||||
await translation.process_output_response(
|
||||
response=copy.deepcopy(result), guardrail_to_apply=self, request_data=scratch_request
|
||||
)
|
||||
|
||||
def supports_scan_only_tool_results(self) -> bool:
|
||||
"""Whether this guardrail can scan tool-result content.
|
||||
|
||||
|
|
@ -1170,7 +1236,7 @@ class CustomGuardrail(CustomLogger):
|
|||
This gets logged on downsteam Langfuse, DataDog, etc.
|
||||
"""
|
||||
# Convert None to empty dict to satisfy type requirements
|
||||
guardrail_response: dict[str, Any] | str = {} if response is None else response
|
||||
guardrail_response: dict[str, object] | str = {} if response is None else response
|
||||
|
||||
# For apply_guardrail functions in custom_code_guardrail scenario,
|
||||
# simplify the logged response to "allow", "deny", or "mask"
|
||||
|
|
|
|||
|
|
@ -2,7 +2,7 @@
|
|||
# On success, logs events to Promptlayer
|
||||
import re
|
||||
import traceback
|
||||
from collections.abc import AsyncGenerator, Mapping
|
||||
from collections.abc import AsyncGenerator, Mapping, Sequence
|
||||
from typing import TYPE_CHECKING, Any, ClassVar, Final, Optional
|
||||
|
||||
from pydantic import BaseModel
|
||||
|
|
@ -31,6 +31,9 @@ if TYPE_CHECKING:
|
|||
|
||||
from litellm.caching.caching import DualCache
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.llms.base_llm.anthropic_messages.transformation import (
|
||||
BaseAnthropicMessagesConfig,
|
||||
)
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.types.mcp import (
|
||||
MCPPostCallResponseObject,
|
||||
|
|
@ -39,7 +42,7 @@ if TYPE_CHECKING:
|
|||
)
|
||||
from litellm.types.router import PreRoutingHookResponse
|
||||
|
||||
Span = _Span | Any
|
||||
Span = _Span
|
||||
else:
|
||||
Span = Any
|
||||
LiteLLMLoggingObj = Any
|
||||
|
|
@ -123,11 +126,11 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
|
|||
return []
|
||||
|
||||
callbacks: Final = AllCallbacks()
|
||||
callback_info: Final = getattr(callbacks, lookup_name, None)
|
||||
callback_info: Final[object] = getattr(callbacks, lookup_name, None)
|
||||
if callback_info is None:
|
||||
return []
|
||||
|
||||
params: Final = getattr(callback_info, "litellm_callback_params", None)
|
||||
params: Final[Sequence[str] | None] = getattr(callback_info, "litellm_callback_params", None)
|
||||
if not params:
|
||||
return []
|
||||
|
||||
|
|
@ -268,7 +271,9 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
|
|||
) -> list[dict]:
|
||||
return healthy_deployments
|
||||
|
||||
async def async_pre_call_deployment_hook(self, kwargs: dict[str, Any], call_type: CallTypes | None) -> dict | None:
|
||||
async def async_pre_call_deployment_hook(
|
||||
self, kwargs: dict[str, object], call_type: CallTypes | None
|
||||
) -> dict | None:
|
||||
"""
|
||||
Allow modifying the request just before it's sent to the deployment.
|
||||
|
||||
|
|
@ -344,9 +349,9 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
|
|||
async def async_post_call_streaming_deployment_hook(
|
||||
self,
|
||||
request_data: dict,
|
||||
response_chunk: Any,
|
||||
response_chunk: object,
|
||||
call_type: CallTypes | None,
|
||||
) -> Any | None:
|
||||
) -> object | None:
|
||||
"""
|
||||
Allow modifying streaming chunks just before they're returned to the user.
|
||||
|
||||
|
|
@ -378,7 +383,7 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
|
|||
"""
|
||||
|
||||
def translate_completion_output_params_streaming(
|
||||
self, completion_stream: Any
|
||||
self, completion_stream: object
|
||||
) -> AdapterCompletionStreamWrapper | None:
|
||||
"""
|
||||
Translates the streaming chunk, from the OpenAI format to the custom format.
|
||||
|
|
@ -418,9 +423,9 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
|
|||
self,
|
||||
data: dict,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
response: Any,
|
||||
response: object,
|
||||
request_headers: dict[str, str] | None = None,
|
||||
litellm_call_info: dict[str, Any] | None = None,
|
||||
litellm_call_info: dict[str, object] | None = None,
|
||||
) -> dict[str, str] | None:
|
||||
"""
|
||||
Called after an LLM API call (success or failure) to allow injecting custom HTTP response headers.
|
||||
|
|
@ -471,11 +476,11 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
|
|||
) -> Any:
|
||||
pass
|
||||
|
||||
async def async_logging_hook(self, kwargs: dict, result: Any, call_type: str) -> tuple[dict, Any]:
|
||||
async def async_logging_hook(self, kwargs: dict, result: object, call_type: str) -> tuple[dict, object]:
|
||||
"""For masking logged request/response. Return a modified version of the request/result."""
|
||||
return kwargs, result
|
||||
|
||||
def logging_hook(self, kwargs: dict, result: Any, call_type: str) -> tuple[dict, Any]:
|
||||
def logging_hook(self, kwargs: dict, result: object, call_type: str) -> tuple[dict, object]:
|
||||
"""For masking logged request/response. Return a modified version of the request/result."""
|
||||
return kwargs, result
|
||||
|
||||
|
|
@ -581,7 +586,7 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
|
|||
|
||||
async def async_should_run_agentic_loop(
|
||||
self,
|
||||
response: Any,
|
||||
response: object,
|
||||
model: str,
|
||||
messages: list[dict],
|
||||
tools: list[dict] | None,
|
||||
|
|
@ -642,8 +647,8 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
|
|||
tools: dict,
|
||||
model: str,
|
||||
messages: list[dict],
|
||||
response: Any,
|
||||
anthropic_messages_provider_config: Any,
|
||||
response: object,
|
||||
anthropic_messages_provider_config: "BaseAnthropicMessagesConfig | None",
|
||||
anthropic_messages_optional_request_params: dict,
|
||||
logging_obj: "LiteLLMLoggingObj",
|
||||
stream: bool,
|
||||
|
|
@ -711,8 +716,8 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
|
|||
tools: dict,
|
||||
model: str,
|
||||
messages: list[dict],
|
||||
response: Any,
|
||||
anthropic_messages_provider_config: Any,
|
||||
response: object,
|
||||
anthropic_messages_provider_config: "BaseAnthropicMessagesConfig | None",
|
||||
anthropic_messages_optional_request_params: dict,
|
||||
logging_obj: "LiteLLMLoggingObj",
|
||||
stream: bool,
|
||||
|
|
@ -728,7 +733,7 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
|
|||
|
||||
async def async_post_agentic_loop_response_hook(
|
||||
self,
|
||||
response: Any,
|
||||
response: object,
|
||||
plan: AgenticLoopPlan,
|
||||
kwargs: dict,
|
||||
) -> Any:
|
||||
|
|
@ -767,7 +772,7 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
|
|||
|
||||
async def async_should_run_chat_completion_agentic_loop(
|
||||
self,
|
||||
response: Any,
|
||||
response: object,
|
||||
model: str,
|
||||
messages: list[dict],
|
||||
tools: list[dict] | None,
|
||||
|
|
@ -785,12 +790,12 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
|
|||
tools: dict,
|
||||
model: str,
|
||||
messages: list[dict],
|
||||
response: Any,
|
||||
response: object,
|
||||
optional_params: dict,
|
||||
logging_obj: "LiteLLMLoggingObj",
|
||||
stream: bool,
|
||||
kwargs: dict,
|
||||
) -> Any:
|
||||
) -> object:
|
||||
"""
|
||||
Hook to execute chat completion agentic loop based on context from should_run hook.
|
||||
"""
|
||||
|
|
@ -800,7 +805,7 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
|
|||
tools: dict,
|
||||
model: str,
|
||||
messages: list[dict],
|
||||
response: Any,
|
||||
response: object,
|
||||
optional_params: dict,
|
||||
logging_obj: "LiteLLMLoggingObj",
|
||||
stream: bool,
|
||||
|
|
@ -851,7 +856,7 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
|
|||
- Converting to string and then truncating the logged content catches this
|
||||
2. We want to avoid modifying the original `messages`, `response`, and `error_str` in the logging payload since these are in kwargs and could be returned to the user
|
||||
"""
|
||||
field_value: Final = standard_logging_object.get(field_name)
|
||||
field_value: Final[object] = standard_logging_object.get(field_name)
|
||||
if field_value:
|
||||
str_value: Final = str(field_value)
|
||||
if len(str_value) > max_length:
|
||||
|
|
@ -1005,8 +1010,8 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
|
|||
• Keep untyped or text content.
|
||||
• Recursively redact inline base64 blobs in *any* string field, at any depth.
|
||||
"""
|
||||
raw_messages: Final[Any] = payload.get("messages", [])
|
||||
messages: Final[list[Any]] = raw_messages if isinstance(raw_messages, list) else []
|
||||
raw_messages: Final[object] = payload.get("messages", [])
|
||||
messages: Final[list[object]] = raw_messages if isinstance(raw_messages, list) else []
|
||||
verbose_logger.debug("[CustomLogger] Stripping base64 from %s messages", len(messages))
|
||||
|
||||
if messages:
|
||||
|
|
@ -1037,8 +1042,8 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
|
|||
• Keep untyped or text content.
|
||||
• Recursively redact inline base64 blobs in *any* string field, at any depth.
|
||||
"""
|
||||
raw_messages: Final[Any] = payload.get("messages", [])
|
||||
messages: Final[list[Any]] = raw_messages if isinstance(raw_messages, list) else []
|
||||
raw_messages: Final[object] = payload.get("messages", [])
|
||||
messages: Final[list[object]] = raw_messages if isinstance(raw_messages, list) else []
|
||||
verbose_logger.debug("[CustomLogger] Stripping base64 from %s messages", len(messages))
|
||||
|
||||
if messages:
|
||||
|
|
@ -1056,10 +1061,10 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
|
|||
|
||||
def _redact_base64(
|
||||
self,
|
||||
value: Any,
|
||||
value: object,
|
||||
depth: int = 0,
|
||||
max_depth: int = DEFAULT_MAX_RECURSE_DEPTH_SENSITIVE_DATA_MASKER,
|
||||
) -> Any:
|
||||
) -> object:
|
||||
"""Recursively redact inline base64 from any nested structure with a max recursion depth limit."""
|
||||
if depth > max_depth:
|
||||
verbose_logger.warning("[CustomLogger] Max recursion depth %s reached while redacting base64", max_depth)
|
||||
|
|
@ -1079,7 +1084,7 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
|
|||
|
||||
return value
|
||||
|
||||
def _should_keep_content(self, content: Any) -> bool:
|
||||
def _should_keep_content(self, content: object) -> bool:
|
||||
"""Return True if this content item should be retained."""
|
||||
if not isinstance(content, dict):
|
||||
return True
|
||||
|
|
@ -1090,16 +1095,16 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
|
|||
|
||||
def _process_messages(
|
||||
self,
|
||||
messages: list[Any],
|
||||
messages: list[object],
|
||||
max_depth: int = DEFAULT_MAX_RECURSE_DEPTH_SENSITIVE_DATA_MASKER,
|
||||
) -> list[dict[str, Any]]:
|
||||
filtered_messages: Final[list[dict[str, Any]]] = []
|
||||
) -> list[dict[str, object]]:
|
||||
filtered_messages: Final[list[dict[str, object]]] = []
|
||||
for msg in messages:
|
||||
if not isinstance(msg, dict):
|
||||
continue
|
||||
contents: Any = msg.get("content")
|
||||
contents: object = msg.get("content")
|
||||
if isinstance(contents, list):
|
||||
cleaned: list[Any] = []
|
||||
cleaned: list[object] = []
|
||||
for c in contents:
|
||||
if self._should_keep_content(content=c):
|
||||
cleaned.append(self._redact_base64(value=c, max_depth=max_depth))
|
||||
|
|
|
|||
|
|
@ -20,10 +20,11 @@ import time
|
|||
import traceback
|
||||
from collections.abc import Sequence
|
||||
from datetime import datetime as datetimeObj
|
||||
from typing import Any, Final
|
||||
from typing import TYPE_CHECKING, Any, Final
|
||||
|
||||
import httpx
|
||||
from httpx import Response
|
||||
from typing_extensions import ReadOnly, TypedDict
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_logger
|
||||
|
|
@ -62,6 +63,18 @@ from litellm.types.utils import StandardLoggingPayload
|
|||
|
||||
from ..additional_logging_utils import AdditionalLoggingUtils
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from fastapi import HTTPException
|
||||
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
|
||||
|
||||
class _DatadogLoggingKwargs(TypedDict, total=False):
|
||||
"""The subset of logging ``kwargs`` that the Datadog payload builder reads."""
|
||||
|
||||
standard_logging_object: ReadOnly[StandardLoggingPayload | None]
|
||||
|
||||
|
||||
# max number of logs DD API can accept
|
||||
|
||||
|
||||
|
|
@ -87,6 +100,11 @@ def _resolve_dd_batch_size() -> int:
|
|||
return max(1, min(value, DD_MAX_BATCH_SIZE))
|
||||
|
||||
|
||||
def _span_attribute(span: object, name: str) -> object:
|
||||
"""Read an optional attribute off whatever span object the active tracer hands back."""
|
||||
return getattr(span, name, None)
|
||||
|
||||
|
||||
class DataDogLogger(
|
||||
CustomBatchLogger,
|
||||
AdditionalLoggingUtils,
|
||||
|
|
@ -271,9 +289,9 @@ class DataDogLogger(
|
|||
self,
|
||||
request_data: dict,
|
||||
original_exception: Exception,
|
||||
user_api_key_dict: Any,
|
||||
user_api_key_dict: "UserAPIKeyAuth",
|
||||
traceback_str: str | None = None,
|
||||
) -> Any | None:
|
||||
) -> "HTTPException | None":
|
||||
"""
|
||||
Log proxy-level failures (e.g. 401 auth, DB connection errors) to Datadog.
|
||||
|
||||
|
|
@ -297,7 +315,7 @@ class DataDogLogger(
|
|||
status_code = int(_code)
|
||||
|
||||
# Use project-standard sanitized user context when running in proxy
|
||||
user_context: dict[str, Any] = {}
|
||||
user_context: dict[str, object] = {}
|
||||
try:
|
||||
from litellm.proxy.litellm_pre_call_utils import (
|
||||
LiteLLMProxyRequestSetup,
|
||||
|
|
@ -553,8 +571,8 @@ class DataDogLogger(
|
|||
|
||||
def create_datadog_logging_payload(
|
||||
self,
|
||||
kwargs: dict | Any,
|
||||
response_obj: Any,
|
||||
kwargs: _DatadogLoggingKwargs,
|
||||
response_obj: object,
|
||||
start_time: datetime.datetime,
|
||||
end_time: datetime.datetime,
|
||||
) -> DatadogPayload:
|
||||
|
|
@ -562,8 +580,8 @@ class DataDogLogger(
|
|||
Helper function to create a datadog payload for logging
|
||||
|
||||
Args:
|
||||
kwargs (Union[dict, Any]): request kwargs
|
||||
response_obj (Any): llm api response
|
||||
kwargs: request kwargs, read for its standard logging object
|
||||
response_obj: llm api response
|
||||
start_time (datetime.datetime): start time of request
|
||||
end_time (datetime.datetime): end time of request
|
||||
|
||||
|
|
@ -625,7 +643,7 @@ class DataDogLogger(
|
|||
self,
|
||||
payload: ServiceLoggerPayload,
|
||||
error: str | None = "",
|
||||
parent_otel_span: Any | None = None,
|
||||
parent_otel_span: object = None,
|
||||
start_time: datetimeObj | float | None = None,
|
||||
end_time: float | datetimeObj | None = None,
|
||||
event_metadata: dict | None = None,
|
||||
|
|
@ -659,7 +677,7 @@ class DataDogLogger(
|
|||
self,
|
||||
payload: ServiceLoggerPayload,
|
||||
error: str | None = "",
|
||||
parent_otel_span: Any | None = None,
|
||||
parent_otel_span: object = None,
|
||||
start_time: datetimeObj | float | None = None,
|
||||
end_time: float | datetimeObj | None = None,
|
||||
event_metadata: dict | None = None,
|
||||
|
|
@ -696,7 +714,7 @@ class DataDogLogger(
|
|||
|
||||
def _create_v0_logging_payload(
|
||||
self,
|
||||
kwargs: dict | Any,
|
||||
kwargs: dict,
|
||||
response_obj: Any,
|
||||
start_time: datetime.datetime,
|
||||
end_time: datetime.datetime,
|
||||
|
|
@ -810,11 +828,11 @@ class DataDogLogger(
|
|||
if current_span is None:
|
||||
return None
|
||||
|
||||
trace_id: Final = getattr(current_span, "trace_id", None)
|
||||
trace_id: Final = _span_attribute(current_span, "trace_id")
|
||||
if trace_id is None:
|
||||
return None
|
||||
|
||||
span_id: Final = getattr(current_span, "span_id", None)
|
||||
span_id: Final = _span_attribute(current_span, "span_id")
|
||||
trace_context: Final[dict[str, str]] = {"trace_id": str(trace_id)}
|
||||
if span_id is not None:
|
||||
trace_context["span_id"] = str(span_id)
|
||||
|
|
|
|||
|
|
@ -9,7 +9,9 @@ API Reference: https://docs.datadoghq.com/llm_observability/setup/api/?tab=examp
|
|||
import asyncio
|
||||
import json
|
||||
import os
|
||||
from collections.abc import Mapping, Sequence
|
||||
from datetime import datetime
|
||||
from types import MappingProxyType
|
||||
from typing import Any, Final, Literal
|
||||
|
||||
import httpx
|
||||
|
|
@ -29,12 +31,16 @@ from litellm.integrations.datadog.datadog_mock_client import (
|
|||
)
|
||||
from litellm.litellm_core_utils.dd_tracing import tracer
|
||||
from litellm.litellm_core_utils.prompt_templates.common_utils import (
|
||||
convert_content_list_to_str,
|
||||
handle_any_messages_to_chat_completion_str_messages_conversion,
|
||||
)
|
||||
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
|
||||
from litellm.litellm_core_utils.safe_json_loads import safe_json_loads
|
||||
from litellm.llms.custom_httpx.http_handler import (
|
||||
get_async_httpx_client,
|
||||
httpxSpecialProvider,
|
||||
)
|
||||
from litellm.proxy.spend_tracking.savings import extract_cache_creation_tokens, extract_cache_read_tokens
|
||||
from litellm.types.integrations.datadog_llm_obs import *
|
||||
from litellm.types.utils import (
|
||||
CallTypes,
|
||||
|
|
@ -43,6 +49,189 @@ from litellm.types.utils import (
|
|||
StandardLoggingPayloadErrorInformation,
|
||||
)
|
||||
|
||||
_EMPTY_MAPPING: Final[Mapping[str, Any]] = MappingProxyType({})
|
||||
_EMPTY_MESSAGE: Final[Message] = {"role": "", "content": ""}
|
||||
_MAX_PARSED_TOOL_ARGUMENT_CHARS: Final = 256 * 1024
|
||||
|
||||
|
||||
def _mapping_field(source: Mapping[str, Any], key: str) -> Mapping[str, Any]:
|
||||
"""The value at `key` when it is a mapping, else an empty one."""
|
||||
value: Final = source.get(key)
|
||||
return value if isinstance(value, dict) else _EMPTY_MAPPING
|
||||
|
||||
|
||||
def _content_blocks(message: Mapping[str, Any]) -> tuple[Mapping[str, Any], ...]:
|
||||
content: Final = message.get("content")
|
||||
if not isinstance(content, list):
|
||||
return ()
|
||||
return tuple(block for block in content if isinstance(block, dict))
|
||||
|
||||
|
||||
def _to_dd_arguments(raw_arguments: object) -> dict[str, Any] | str:
|
||||
"""
|
||||
Arguments as the object LLM Obs types them as, or the raw string when they are not one.
|
||||
|
||||
Strings past the size bound ship unparsed: decoding multiplies memory on hostile compact
|
||||
JSON, and the raw string is what the intake receives either way.
|
||||
"""
|
||||
if not isinstance(raw_arguments, str):
|
||||
return raw_arguments if isinstance(raw_arguments, dict) else str(raw_arguments)
|
||||
if len(raw_arguments) > _MAX_PARSED_TOOL_ARGUMENT_CHARS:
|
||||
return raw_arguments
|
||||
parsed: Final = safe_json_loads(raw_arguments)
|
||||
return parsed if isinstance(parsed, dict) else raw_arguments
|
||||
|
||||
|
||||
def _to_dd_tool_calls(message: Mapping[str, Any]) -> tuple[ToolCall, ...]:
|
||||
"""
|
||||
The tool calls a message carries, in LLM Obs' ToolCall schema, from either dialect.
|
||||
|
||||
OpenAI puts them in `tool_calls` with the callee nested under `function` and `arguments`
|
||||
serialized; Anthropic puts them in `content` as `tool_use` blocks with `input` already an
|
||||
object. LLM Obs reads `name` / `arguments` / `tool_id` either way.
|
||||
"""
|
||||
raw_tool_calls: Final = message.get("tool_calls")
|
||||
openai_calls: Final = tuple(
|
||||
ToolCall(
|
||||
name=function.get("name", ""),
|
||||
arguments=_to_dd_arguments(function.get("arguments", "")),
|
||||
tool_id=tool_call.get("id", ""),
|
||||
type=tool_call.get("type", "function"),
|
||||
)
|
||||
for tool_call in (raw_tool_calls if isinstance(raw_tool_calls, list) else ())
|
||||
if isinstance(tool_call, dict)
|
||||
for function in [_mapping_field(tool_call, "function")]
|
||||
)
|
||||
anthropic_calls: Final = tuple(
|
||||
ToolCall(
|
||||
name=block.get("name", ""),
|
||||
arguments=_to_dd_arguments(block.get("input") or {}),
|
||||
tool_id=block.get("id", ""),
|
||||
type="tool_use",
|
||||
)
|
||||
for block in _content_blocks(message)
|
||||
if block.get("type") == "tool_use"
|
||||
)
|
||||
return openai_calls + anthropic_calls
|
||||
|
||||
|
||||
def _to_dd_tool_results(message: Mapping[str, Any], tool_call_names: Mapping[str, str]) -> tuple[ToolResult, ...]:
|
||||
"""
|
||||
The tool results a message carries, linked back to the call each answers.
|
||||
|
||||
OpenAI models a result as a whole `role: "tool"` message keyed by `tool_call_id`;
|
||||
Anthropic nests `tool_result` blocks inside a user message, keyed by `tool_use_id`.
|
||||
"""
|
||||
|
||||
def to_result(tool_id: str, result: object) -> ToolResult:
|
||||
return ToolResult(
|
||||
name=tool_call_names.get(tool_id, ""),
|
||||
result=result if isinstance(result, str) else safe_dumps(result),
|
||||
tool_id=tool_id,
|
||||
type="function",
|
||||
)
|
||||
|
||||
if message.get("role") == "tool":
|
||||
return (to_result(str(message.get("tool_call_id", "")), message.get("content") or ""),)
|
||||
return tuple(
|
||||
to_result(str(block.get("tool_use_id", "")), block.get("content") or "")
|
||||
for block in _content_blocks(message)
|
||||
if block.get("type") == "tool_result"
|
||||
)
|
||||
|
||||
|
||||
def _tool_call_names_by_id(messages: Sequence[object]) -> Mapping[str, str]:
|
||||
"""Ids to tool names for result linking; reads names structurally and parses nothing."""
|
||||
openai_pairs: Final = tuple(
|
||||
(tool_call.get("id"), function.get("name", ""))
|
||||
for message in messages
|
||||
if isinstance(message, dict) and isinstance(message.get("tool_calls"), list)
|
||||
for tool_call in message["tool_calls"]
|
||||
if isinstance(tool_call, dict)
|
||||
for function in [_mapping_field(tool_call, "function")]
|
||||
)
|
||||
anthropic_pairs: Final = tuple(
|
||||
(block.get("id"), block.get("name", ""))
|
||||
for message in messages
|
||||
if isinstance(message, dict)
|
||||
for block in _content_blocks(message)
|
||||
if block.get("type") == "tool_use"
|
||||
)
|
||||
return MappingProxyType({str(tool_id): str(name) for tool_id, name in openai_pairs + anthropic_pairs if tool_id})
|
||||
|
||||
|
||||
def _to_dd_message(message: object, tool_call_names: Mapping[str, str]) -> Message:
|
||||
"""
|
||||
Map one chat message onto LLM Obs' Message schema, adding fields and never destroying content.
|
||||
|
||||
Content collapses to its text only when it has text; a content list with none (tool blocks,
|
||||
images) rides along unchanged so nothing the caller logged is lost. Tool calls and results
|
||||
move into the fields the LLM Obs Tools panel reads, from both the OpenAI and Anthropic shapes.
|
||||
"""
|
||||
if not isinstance(message, dict):
|
||||
converted: Final = handle_any_messages_to_chat_completion_str_messages_conversion(message)
|
||||
return converted[0] if converted else _EMPTY_MESSAGE
|
||||
|
||||
text: Final = convert_content_list_to_str(message) # pyright: ignore[reportArgumentType] # caller-supplied dict
|
||||
original_content: Final = message.get("content")
|
||||
content: Final = (
|
||||
text if text or not isinstance(original_content, list) or not original_content else original_content
|
||||
)
|
||||
reasoning: Final = message.get("reasoning_content")
|
||||
tool_calls: Final = _to_dd_tool_calls(message)
|
||||
tool_results: Final = _to_dd_tool_results(message, tool_call_names)
|
||||
dd_message: Final[Message] = {
|
||||
"role": message.get("role", ""),
|
||||
"content": content,
|
||||
**({"reasoning_content": reasoning} if reasoning is not None else {}),
|
||||
**({"tool_calls": tool_calls} if tool_calls else {}),
|
||||
**({"tool_results": tool_results} if tool_results else {}),
|
||||
}
|
||||
return dd_message
|
||||
|
||||
|
||||
def _to_dd_messages(messages: object) -> tuple[Message, ...]:
|
||||
"""Map a whole conversation, resolving each tool result against the calls that precede it."""
|
||||
if messages is None:
|
||||
return ()
|
||||
if not isinstance(messages, list):
|
||||
return tuple(handle_any_messages_to_chat_completion_str_messages_conversion(messages))
|
||||
tool_call_names: Final = _tool_call_names_by_id(messages)
|
||||
return tuple(_to_dd_message(message, tool_call_names) for message in messages)
|
||||
|
||||
|
||||
def _to_dd_tool_definition(entry: Mapping[str, Any]) -> ToolDefinition | None:
|
||||
function: Final = entry.get("function")
|
||||
declared: Final[Mapping[str, Any]] = function if isinstance(function, dict) else entry
|
||||
name: Final = declared.get("name")
|
||||
if not name:
|
||||
return None
|
||||
schema: Final = declared.get("parameters") or declared.get("input_schema")
|
||||
description: Final = declared.get("description", "")
|
||||
if not isinstance(schema, dict):
|
||||
return ToolDefinition(name=name, description=description)
|
||||
return ToolDefinition(name=name, description=description, schema=schema)
|
||||
|
||||
|
||||
def _to_dd_tool_definitions(model_parameters: object) -> tuple[ToolDefinition, ...]:
|
||||
"""
|
||||
Map the request's declared tools onto LLM Obs' ToolDefinition schema.
|
||||
|
||||
Handles the wrapped chat-completions shape and the bare shape the Anthropic and
|
||||
Responses surfaces use, since both reach this logger through `model_parameters`.
|
||||
"""
|
||||
if not isinstance(model_parameters, dict):
|
||||
return ()
|
||||
raw_tools: Final = model_parameters.get("tools") or model_parameters.get("functions")
|
||||
if not isinstance(raw_tools, list):
|
||||
return ()
|
||||
return tuple(
|
||||
definition
|
||||
for entry in raw_tools
|
||||
if isinstance(entry, dict)
|
||||
if (definition := _to_dd_tool_definition(entry)) is not None
|
||||
)
|
||||
|
||||
|
||||
class DataDogLLMObsLogger(CustomBatchLogger):
|
||||
def __init__(self, **kwargs):
|
||||
|
|
@ -221,12 +410,9 @@ class DataDogLLMObsLogger(CustomBatchLogger):
|
|||
if standard_logging_payload is None:
|
||||
raise Exception("DataDogLLMObs: standard_logging_object is not set")
|
||||
|
||||
messages = standard_logging_payload["messages"]
|
||||
messages = self._ensure_string_content(messages=messages)
|
||||
|
||||
metadata: Final = kwargs.get("litellm_params", {}).get("metadata", {})
|
||||
|
||||
input_meta: Final = InputMeta(messages=handle_any_messages_to_chat_completion_str_messages_conversion(messages))
|
||||
input_meta: Final = InputMeta(messages=_to_dd_messages(standard_logging_payload["messages"]))
|
||||
output_meta: Final = OutputMeta(
|
||||
messages=self._get_response_messages(
|
||||
standard_logging_payload=standard_logging_payload,
|
||||
|
|
@ -240,22 +426,20 @@ class DataDogLLMObsLogger(CustomBatchLogger):
|
|||
if isinstance(metadata, dict):
|
||||
metadata_parent_id = metadata.get("parent_id")
|
||||
|
||||
meta: Final = Meta(
|
||||
kind=self._get_datadog_span_kind(standard_logging_payload.get("call_type"), metadata_parent_id),
|
||||
input=input_meta,
|
||||
output=output_meta,
|
||||
metadata=self._get_dd_llm_obs_payload_metadata(standard_logging_payload),
|
||||
error=error_info,
|
||||
)
|
||||
tool_definitions: Final = _to_dd_tool_definitions(standard_logging_payload.get("model_parameters"))
|
||||
span_kind: Final = self._get_datadog_span_kind(standard_logging_payload.get("call_type"), metadata_parent_id)
|
||||
payload_metadata: Final = self._get_dd_llm_obs_payload_metadata(standard_logging_payload)
|
||||
|
||||
# Calculate metrics (you may need to adjust these based on available data)
|
||||
metrics: Final = LLMMetrics(
|
||||
input_tokens=float(standard_logging_payload.get("prompt_tokens", 0)),
|
||||
output_tokens=float(standard_logging_payload.get("completion_tokens", 0)),
|
||||
total_tokens=float(standard_logging_payload.get("total_tokens", 0)),
|
||||
total_cost=float(standard_logging_payload.get("response_cost", 0)),
|
||||
time_to_first_token=self._get_time_to_first_token_seconds(standard_logging_payload),
|
||||
)
|
||||
meta: Final[Meta] = {
|
||||
"kind": span_kind,
|
||||
"input": input_meta,
|
||||
"output": output_meta,
|
||||
"metadata": payload_metadata,
|
||||
"error": error_info,
|
||||
**({"tool_definitions": tool_definitions} if tool_definitions else {}),
|
||||
}
|
||||
|
||||
metrics: Final = self._assemble_metrics(standard_logging_payload)
|
||||
|
||||
payload: Final[LLMObsPayload] = LLMObsPayload(
|
||||
parent_id=metadata_parent_id if metadata_parent_id else "undefined",
|
||||
|
|
@ -313,6 +497,45 @@ class DataDogLLMObsLogger(CustomBatchLogger):
|
|||
)
|
||||
return error_info
|
||||
|
||||
def _assemble_metrics(self, standard_logging_payload: StandardLoggingPayload) -> LLMMetrics:
|
||||
"""
|
||||
Build the span metrics, including the prompt-cache counts LLM Obs charts cache savings from.
|
||||
|
||||
Cache counts resolve through the same owners the savings dashboard uses, so every provider
|
||||
spelling is covered, and `non_cached_input_tokens` subtracts BOTH cache categories because
|
||||
litellm's normalized prompt count includes both (the invariant the cost calculator's custom
|
||||
pricing helper documents). A zero residual on a fully cached request is real data and is
|
||||
emitted; a zero read or write count is absence and is not.
|
||||
"""
|
||||
prompt_tokens: Final = float(standard_logging_payload.get("prompt_tokens", 0))
|
||||
completion_tokens: Final = float(standard_logging_payload.get("completion_tokens", 0))
|
||||
total_tokens: Final = float(standard_logging_payload.get("total_tokens", 0))
|
||||
total_cost: Final = float(standard_logging_payload.get("response_cost", 0))
|
||||
time_to_first_token: Final = self._get_time_to_first_token_seconds(standard_logging_payload)
|
||||
|
||||
raw_usage: Final = (standard_logging_payload.get("metadata") or {}).get("usage_object")
|
||||
usage_object: Final = raw_usage if isinstance(raw_usage, dict) else None
|
||||
cache_read: Final = float(extract_cache_read_tokens(usage_object))
|
||||
cache_write: Final = float(extract_cache_creation_tokens(usage_object))
|
||||
|
||||
metrics: Final[LLMMetrics] = {
|
||||
"input_tokens": prompt_tokens,
|
||||
"output_tokens": completion_tokens,
|
||||
"total_tokens": total_tokens,
|
||||
"total_cost": total_cost,
|
||||
"time_to_first_token": time_to_first_token,
|
||||
**(
|
||||
{
|
||||
**({"cache_read_input_tokens": cache_read} if cache_read else {}),
|
||||
**({"cache_write_input_tokens": cache_write} if cache_write else {}),
|
||||
"non_cached_input_tokens": max(prompt_tokens - cache_read - cache_write, 0.0),
|
||||
}
|
||||
if cache_read or cache_write
|
||||
else {}
|
||||
),
|
||||
}
|
||||
return metrics
|
||||
|
||||
def _get_time_to_first_token_seconds(self, standard_logging_payload: StandardLoggingPayload) -> float:
|
||||
"""
|
||||
Get the time to first token in seconds
|
||||
|
|
@ -334,7 +557,7 @@ class DataDogLLMObsLogger(CustomBatchLogger):
|
|||
|
||||
def _get_response_messages(
|
||||
self, standard_logging_payload: StandardLoggingPayload, call_type: str | None
|
||||
) -> list[Any]:
|
||||
) -> tuple[Message, ...]:
|
||||
"""
|
||||
Get the messages from the response object
|
||||
|
||||
|
|
@ -343,7 +566,7 @@ class DataDogLLMObsLogger(CustomBatchLogger):
|
|||
|
||||
response_obj = standard_logging_payload.get("response")
|
||||
if response_obj is None:
|
||||
return []
|
||||
return ()
|
||||
|
||||
# edge case: handle response_obj is a string representation of a dict
|
||||
if isinstance(response_obj, str):
|
||||
|
|
@ -356,7 +579,7 @@ class DataDogLLMObsLogger(CustomBatchLogger):
|
|||
# fallback to json parsing
|
||||
response_obj = json.loads(str(response_obj))
|
||||
except json.JSONDecodeError:
|
||||
return []
|
||||
return ()
|
||||
|
||||
if call_type in [
|
||||
CallTypes.completion.value,
|
||||
|
|
@ -374,12 +597,12 @@ class DataDogLLMObsLogger(CustomBatchLogger):
|
|||
if isinstance(response_obj, dict) and "choices" in response_obj:
|
||||
choices: Final = response_obj["choices"]
|
||||
if choices and len(choices) > 0 and "message" in choices[0]:
|
||||
return [choices[0]["message"]]
|
||||
return []
|
||||
return _to_dd_messages([choices[0]["message"]])
|
||||
return ()
|
||||
except (KeyError, IndexError, TypeError):
|
||||
# In case of any error accessing the response structure, return empty list
|
||||
return []
|
||||
return []
|
||||
return ()
|
||||
return ()
|
||||
|
||||
def _get_datadog_span_kind(
|
||||
self, call_type: str | None, parent_id: str | None = None
|
||||
|
|
@ -484,22 +707,11 @@ class DataDogLLMObsLogger(CustomBatchLogger):
|
|||
# Default fallback for unknown or passthrough operations
|
||||
return "llm"
|
||||
|
||||
def _ensure_string_content(self, messages: str | list[Any] | dict[Any, Any] | None) -> list[Any]:
|
||||
if messages is None:
|
||||
return []
|
||||
if isinstance(messages, str):
|
||||
return [messages]
|
||||
elif isinstance(messages, list):
|
||||
return [message for message in messages]
|
||||
elif isinstance(messages, dict):
|
||||
return [str(messages.get("content", ""))]
|
||||
return []
|
||||
|
||||
def _get_dd_llm_obs_payload_metadata(self, standard_logging_payload: StandardLoggingPayload) -> dict[str, Any]:
|
||||
def _get_dd_llm_obs_payload_metadata(self, standard_logging_payload: StandardLoggingPayload) -> dict[str, object]:
|
||||
"""
|
||||
Fields to track in DD LLM Observability metadata from litellm standard logging payload
|
||||
"""
|
||||
_metadata: Final[dict[str, Any]] = {
|
||||
_metadata: Final[dict[str, object]] = {
|
||||
"model_name": standard_logging_payload.get("model", "unknown"),
|
||||
"model_provider": standard_logging_payload.get("custom_llm_provider", "unknown"),
|
||||
"id": standard_logging_payload.get("id", "unknown"),
|
||||
|
|
@ -523,10 +735,6 @@ class DataDogLLMObsLogger(CustomBatchLogger):
|
|||
spend_metrics: Final = self._get_spend_metrics(standard_logging_payload)
|
||||
_metadata.update({"spend_metrics": dict(spend_metrics)})
|
||||
|
||||
## extract tool calls and add to metadata
|
||||
tool_call_metadata: Final = self._extract_tool_call_metadata(standard_logging_payload)
|
||||
_metadata.update(tool_call_metadata)
|
||||
|
||||
_standard_logging_metadata: Final[dict] = dict(standard_logging_payload.get("metadata", {})) or {}
|
||||
_metadata.update(_standard_logging_metadata)
|
||||
return _metadata
|
||||
|
|
@ -646,107 +854,3 @@ class DataDogLLMObsLogger(CustomBatchLogger):
|
|||
verbose_logger.debug("Original value: %s", user_api_key_budget_reset_at)
|
||||
|
||||
return spend_metrics
|
||||
|
||||
def _process_input_messages_preserving_tool_calls(self, messages: list[Any]) -> list[dict[str, Any]]:
|
||||
"""
|
||||
Process input messages while preserving tool_calls and tool message types.
|
||||
|
||||
This bypasses the lossy string conversion when tool calls are present,
|
||||
allowing complex nested tool_calls objects to be preserved for Datadog.
|
||||
"""
|
||||
processed: Final = []
|
||||
for msg in messages:
|
||||
if isinstance(msg, dict):
|
||||
# Preserve messages with tool_calls or tool role as-is
|
||||
if "tool_calls" in msg or msg.get("role") == "tool":
|
||||
processed.append(msg)
|
||||
else:
|
||||
# For regular messages, still apply string conversion
|
||||
converted = handle_any_messages_to_chat_completion_str_messages_conversion([msg])
|
||||
processed.extend(converted)
|
||||
else:
|
||||
# For non-dict messages, apply string conversion
|
||||
converted = handle_any_messages_to_chat_completion_str_messages_conversion([msg])
|
||||
processed.extend(converted)
|
||||
return processed
|
||||
|
||||
@staticmethod
|
||||
def _tool_calls_kv_pair(tool_calls: list[dict[str, Any]]) -> dict[str, Any]:
|
||||
"""
|
||||
Extract tool call information into key-value pairs for Datadog metadata.
|
||||
|
||||
Similar to OpenTelemetry's implementation but adapted for Datadog's format.
|
||||
"""
|
||||
kv_pairs: Final[dict[str, Any]] = {}
|
||||
for idx, tool_call in enumerate(tool_calls):
|
||||
try:
|
||||
# Extract tool call ID
|
||||
tool_id = tool_call.get("id")
|
||||
if tool_id:
|
||||
kv_pairs[f"tool_calls.{idx}.id"] = tool_id
|
||||
|
||||
# Extract tool call type
|
||||
tool_type = tool_call.get("type")
|
||||
if tool_type:
|
||||
kv_pairs[f"tool_calls.{idx}.type"] = tool_type
|
||||
|
||||
# Extract function information
|
||||
function = tool_call.get("function")
|
||||
if function:
|
||||
function_name = function.get("name")
|
||||
if function_name:
|
||||
kv_pairs[f"tool_calls.{idx}.function.name"] = function_name
|
||||
|
||||
function_arguments = function.get("arguments")
|
||||
if function_arguments:
|
||||
# Store arguments as JSON string for Datadog
|
||||
if isinstance(function_arguments, str):
|
||||
kv_pairs[f"tool_calls.{idx}.function.arguments"] = function_arguments
|
||||
else:
|
||||
import json
|
||||
|
||||
kv_pairs[f"tool_calls.{idx}.function.arguments"] = json.dumps(function_arguments)
|
||||
except (KeyError, TypeError, ValueError) as e:
|
||||
verbose_logger.debug("DataDogLLMObs: Error processing tool call %s: %s", idx, e)
|
||||
continue
|
||||
|
||||
return kv_pairs
|
||||
|
||||
def _extract_tool_call_metadata(self, standard_logging_payload: StandardLoggingPayload) -> dict[str, Any]:
|
||||
"""
|
||||
Extract tool call information from both input messages and response for Datadog metadata.
|
||||
"""
|
||||
tool_call_metadata: Final[dict[str, Any]] = {}
|
||||
|
||||
try:
|
||||
# Extract tool calls from input messages
|
||||
messages: Final = standard_logging_payload.get("messages", [])
|
||||
if messages and isinstance(messages, list):
|
||||
for message in messages:
|
||||
if isinstance(message, dict) and "tool_calls" in message:
|
||||
tool_calls = message.get("tool_calls")
|
||||
if tool_calls:
|
||||
input_tool_calls_kv = self._tool_calls_kv_pair(tool_calls)
|
||||
# Prefix with "input_" to distinguish from response tool calls
|
||||
for key, value in input_tool_calls_kv.items():
|
||||
tool_call_metadata[f"input_{key}"] = value
|
||||
|
||||
# Extract tool calls from response
|
||||
response_obj: Final = standard_logging_payload.get("response")
|
||||
if response_obj and isinstance(response_obj, dict):
|
||||
choices: Final = response_obj.get("choices", [])
|
||||
for choice in choices:
|
||||
if isinstance(choice, dict):
|
||||
message = choice.get("message")
|
||||
if message and isinstance(message, dict):
|
||||
tool_calls = message.get("tool_calls")
|
||||
if tool_calls:
|
||||
response_tool_calls_kv = self._tool_calls_kv_pair(tool_calls)
|
||||
# Prefix with "output_" to distinguish from input tool calls
|
||||
for key, value in response_tool_calls_kv.items():
|
||||
tool_call_metadata[f"output_{key}"] = value
|
||||
|
||||
except Exception as e:
|
||||
verbose_logger.debug("DataDogLLMObs: Error extracting tool call metadata: %s", e)
|
||||
|
||||
return tool_call_metadata
|
||||
|
|
|
|||
|
|
@ -3,12 +3,21 @@ Based on Google's GenAI Kit dotprompt implementation: https://google.github.io/d
|
|||
"""
|
||||
|
||||
import re
|
||||
from collections.abc import Mapping
|
||||
from pathlib import Path
|
||||
from typing import Any, Final
|
||||
|
||||
import yaml
|
||||
from jinja2 import DictLoader, select_autoescape
|
||||
from jinja2.sandbox import ImmutableSandboxedEnvironment
|
||||
from typing_extensions import NotRequired, ReadOnly, TypedDict
|
||||
|
||||
|
||||
class _PromptFileJson(TypedDict):
|
||||
"""JSON form of a .prompt file: rendered template text plus its frontmatter."""
|
||||
|
||||
content: ReadOnly[NotRequired[str]]
|
||||
metadata: ReadOnly[NotRequired[dict[str, object]]]
|
||||
|
||||
|
||||
def strip_version_suffix(prompt_id: str) -> str | None:
|
||||
|
|
@ -167,7 +176,7 @@ class PromptManager:
|
|||
template_id=prompt_id,
|
||||
)
|
||||
|
||||
def _parse_frontmatter(self, content: str) -> tuple[dict[str, Any], str]:
|
||||
def _parse_frontmatter(self, content: str) -> tuple[dict[str, object], str]:
|
||||
"""Parse YAML frontmatter from prompt content."""
|
||||
# Match YAML frontmatter between --- delimiters
|
||||
frontmatter_pattern: Final = r"^---\s*\n(.*?)\n---\s*\n(.*)$"
|
||||
|
|
@ -178,7 +187,7 @@ class PromptManager:
|
|||
template_content = match.group(2)
|
||||
|
||||
try:
|
||||
frontmatter = yaml.safe_load(frontmatter_yaml) or {}
|
||||
frontmatter: dict[str, object] = yaml.safe_load(frontmatter_yaml) or {}
|
||||
except yaml.YAMLError as e:
|
||||
raise ValueError(f"Invalid YAML frontmatter: {e}")
|
||||
else:
|
||||
|
|
@ -191,7 +200,7 @@ class PromptManager:
|
|||
def render(
|
||||
self,
|
||||
prompt_id: str,
|
||||
prompt_variables: dict[str, Any] | None = None,
|
||||
prompt_variables: Mapping[str, object] | None = None,
|
||||
version: int | None = None,
|
||||
) -> str:
|
||||
"""
|
||||
|
|
@ -231,7 +240,7 @@ class PromptManager:
|
|||
except Exception as e:
|
||||
raise ValueError(f"Error rendering template '{prompt_id}': {e}")
|
||||
|
||||
def _validate_input(self, variables: dict[str, Any], schema: dict[str, Any]) -> None:
|
||||
def _validate_input(self, variables: Mapping[str, object], schema: Mapping[str, str]) -> None:
|
||||
"""Basic validation of input variables against schema."""
|
||||
for field_name, field_type in schema.items():
|
||||
if field_name in variables:
|
||||
|
|
@ -291,7 +300,7 @@ class PromptManager:
|
|||
"""Get a list of all available prompt IDs."""
|
||||
return list(self.prompts.keys())
|
||||
|
||||
def get_prompt_metadata(self, prompt_id: str) -> dict[str, Any] | None:
|
||||
def get_prompt_metadata(self, prompt_id: str) -> dict[str, object] | None:
|
||||
"""Get metadata for a specific prompt."""
|
||||
template: Final = self.prompts.get(prompt_id)
|
||||
return template.metadata if template else None
|
||||
|
|
@ -302,12 +311,12 @@ class PromptManager:
|
|||
if self.prompt_directory:
|
||||
self._load_prompts()
|
||||
|
||||
def add_prompt(self, prompt_id: str, content: str, metadata: dict[str, Any] | None = None) -> None:
|
||||
def add_prompt(self, prompt_id: str, content: str, metadata: dict[str, object] | None = None) -> None:
|
||||
"""Add a prompt template programmatically."""
|
||||
template: Final = PromptTemplate(content=content, metadata=metadata or {}, template_id=prompt_id)
|
||||
self.prompts[prompt_id] = template
|
||||
|
||||
def prompt_file_to_json(self, file_path: str | Path) -> dict[str, Any]:
|
||||
def prompt_file_to_json(self, file_path: str | Path) -> _PromptFileJson:
|
||||
"""Convert a .prompt file to JSON format.
|
||||
|
||||
Args:
|
||||
|
|
@ -324,7 +333,7 @@ class PromptManager:
|
|||
|
||||
return {"content": template_content.strip(), "metadata": frontmatter}
|
||||
|
||||
def json_to_prompt_file(self, prompt_data: dict[str, Any]) -> str:
|
||||
def json_to_prompt_file(self, prompt_data: _PromptFileJson) -> str:
|
||||
"""Convert JSON prompt data to .prompt file format.
|
||||
|
||||
Args:
|
||||
|
|
|
|||
|
|
@ -6,10 +6,11 @@ import re
|
|||
import uuid
|
||||
from collections.abc import Mapping, Sequence
|
||||
from datetime import datetime, timezone, tzinfo
|
||||
from typing import Any, Final, TypedDict, cast
|
||||
from typing import Any, Final, Protocol, cast
|
||||
|
||||
import httpx
|
||||
from pydantic import BaseModel, Field
|
||||
from typing_extensions import ReadOnly, TypedDict
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_logger
|
||||
|
|
@ -35,6 +36,34 @@ GALILEO_CLOUD_API_BASE_URL: Final = "https://api.galileo.ai"
|
|||
GALILEO_MAX_IN_MEMORY_RECORDS: Final = 1000
|
||||
|
||||
|
||||
class _GalileoLoginBody(TypedDict):
|
||||
"""Decoded body of the Galileo login response."""
|
||||
|
||||
access_token: ReadOnly[str]
|
||||
|
||||
|
||||
class _GalileoLoginResponse(Protocol):
|
||||
"""The login call's HTTP response, read for the access token it carries."""
|
||||
|
||||
def json(self) -> _GalileoLoginBody: ...
|
||||
|
||||
|
||||
class _JsonResponse(Protocol):
|
||||
"""An HTTP response read only for whatever JSON body it decodes to."""
|
||||
|
||||
def json(self) -> object: ...
|
||||
|
||||
|
||||
def _login_access_token(response: _GalileoLoginResponse) -> str:
|
||||
"""Read the bearer token out of a Galileo login response body."""
|
||||
return response.json()["access_token"]
|
||||
|
||||
|
||||
def _decoded_body(response: _JsonResponse) -> object:
|
||||
"""Decode a response body without asserting anything about its shape."""
|
||||
return response.json()
|
||||
|
||||
|
||||
class GalileoStandardLoggingFields(TypedDict, total=False):
|
||||
call_type: str
|
||||
model: str
|
||||
|
|
@ -156,7 +185,7 @@ class GalileoObserve(CustomLogger):
|
|||
},
|
||||
)
|
||||
galileo_login_response.raise_for_status()
|
||||
access_token: Final = galileo_login_response.json()["access_token"]
|
||||
access_token: Final = _login_access_token(galileo_login_response)
|
||||
self.headers = {
|
||||
"accept": "application/json",
|
||||
"Content-Type": "application/json",
|
||||
|
|
@ -421,7 +450,7 @@ class GalileoObserve(CustomLogger):
|
|||
try:
|
||||
verbose_logger.debug(
|
||||
"Galileo Logger HTTP error response json: %s",
|
||||
response.json(),
|
||||
_decoded_body(response),
|
||||
)
|
||||
except Exception:
|
||||
pass
|
||||
|
|
|
|||
|
|
@ -4,12 +4,80 @@ Now supports selecting a tag via `config["tag"]`; falls back to branch ("main").
|
|||
"""
|
||||
|
||||
import base64
|
||||
from typing import Any, Final
|
||||
from collections.abc import Mapping, Sequence
|
||||
from typing import Any, Final, Protocol, TypedDict
|
||||
from urllib.parse import quote
|
||||
|
||||
from typing_extensions import ReadOnly
|
||||
|
||||
from litellm.llms.custom_httpx.http_handler import HTTPHandler
|
||||
|
||||
|
||||
class GitLabFilePayload(TypedDict, total=False):
|
||||
"""A repository-files API entry."""
|
||||
|
||||
content: ReadOnly[str]
|
||||
encoding: ReadOnly[str]
|
||||
|
||||
|
||||
class GitLabTreeEntry(TypedDict, total=False):
|
||||
"""A repository-tree API entry."""
|
||||
|
||||
path: ReadOnly[str]
|
||||
type: ReadOnly[str]
|
||||
|
||||
|
||||
class GitLabBranch(TypedDict, total=False):
|
||||
"""A repository-branches API entry."""
|
||||
|
||||
name: ReadOnly[str]
|
||||
type: ReadOnly[str]
|
||||
|
||||
|
||||
class GitLabFileMetadata(TypedDict):
|
||||
"""The response headers a raw file request exposes as metadata."""
|
||||
|
||||
content_type: ReadOnly[str | None]
|
||||
content_length: ReadOnly[str | None]
|
||||
last_modified: ReadOnly[str | None]
|
||||
|
||||
|
||||
class _FileJsonResponse(Protocol):
|
||||
def json(self) -> GitLabFilePayload: ...
|
||||
|
||||
|
||||
class _TreeJsonResponse(Protocol):
|
||||
def json(self) -> Sequence[GitLabTreeEntry] | None: ...
|
||||
|
||||
|
||||
class _ProjectJsonResponse(Protocol):
|
||||
def json(self) -> Mapping[str, object]: ...
|
||||
|
||||
|
||||
class _BranchesJsonResponse(Protocol):
|
||||
def json(self) -> Sequence[GitLabBranch] | None: ...
|
||||
|
||||
|
||||
def _file_payload(resp: _FileJsonResponse) -> GitLabFilePayload:
|
||||
"""The JSON body of a repository-files response."""
|
||||
return resp.json()
|
||||
|
||||
|
||||
def _tree_entries(resp: _TreeJsonResponse) -> Sequence[GitLabTreeEntry]:
|
||||
"""The entries of a repository-tree response."""
|
||||
return resp.json() or []
|
||||
|
||||
|
||||
def _project_info(resp: _ProjectJsonResponse) -> Mapping[str, object]:
|
||||
"""The JSON body of a project response."""
|
||||
return resp.json()
|
||||
|
||||
|
||||
def _branch_entries(resp: _BranchesJsonResponse) -> Sequence[GitLabBranch] | None:
|
||||
"""The JSON body of a repository-branches response."""
|
||||
return resp.json()
|
||||
|
||||
|
||||
class GitLabClient:
|
||||
"""
|
||||
Client for interacting with the GitLab API to fetch files.
|
||||
|
|
@ -42,12 +110,12 @@ class GitLabClient:
|
|||
|
||||
self.project: str | int = project
|
||||
self.access_token: str = str(access_token)
|
||||
self.auth_method = config.get("auth_method", "token") # 'token' or 'oauth'
|
||||
self.auth_method: str = config.get("auth_method", "token") # 'token' or 'oauth'
|
||||
self.branch = config.get("branch", None)
|
||||
if not self.branch:
|
||||
self.branch = "main"
|
||||
self.tag = config.get("tag")
|
||||
self.base_url = config.get("base_url", "https://gitlab.com/api/v4")
|
||||
self.base_url: str = config.get("base_url", "https://gitlab.com/api/v4")
|
||||
|
||||
if not all([self.project, self.access_token]):
|
||||
raise ValueError("project and access_token are required")
|
||||
|
|
@ -159,7 +227,7 @@ class GitLabClient:
|
|||
if resp.status_code == 404:
|
||||
return None
|
||||
resp.raise_for_status()
|
||||
data: Final = resp.json()
|
||||
data: Final = _file_payload(resp)
|
||||
content: Final = data.get("content")
|
||||
encoding: Final = data.get("encoding", "")
|
||||
if content and encoding == "base64":
|
||||
|
|
@ -208,7 +276,7 @@ class GitLabClient:
|
|||
return []
|
||||
resp.raise_for_status()
|
||||
|
||||
data: Final = resp.json() or []
|
||||
data: Final = _tree_entries(resp)
|
||||
files: Final[list[str]] = []
|
||||
for item in data:
|
||||
if item.get("type") == "blob":
|
||||
|
|
@ -229,13 +297,13 @@ class GitLabClient:
|
|||
raise Exception("Authentication failed. Check your GitLab token and auth_method.")
|
||||
raise Exception(f"Failed to list files in '{directory_path}': {e}")
|
||||
|
||||
def get_repository_info(self) -> dict[str, Any]:
|
||||
def get_repository_info(self) -> Mapping[str, object]:
|
||||
"""Get information about the project/repository."""
|
||||
url: Final = f"{self.base_url}/projects/{self._project_enc}"
|
||||
try:
|
||||
resp: Final = self.http_handler.get(url, headers=self.headers)
|
||||
resp.raise_for_status()
|
||||
return resp.json()
|
||||
return _project_info(resp)
|
||||
except Exception as e:
|
||||
raise Exception(f"Failed to get repository info: {e}")
|
||||
|
||||
|
|
@ -247,18 +315,18 @@ class GitLabClient:
|
|||
except Exception:
|
||||
return False
|
||||
|
||||
def get_branches(self) -> list[dict[str, Any]]:
|
||||
def get_branches(self) -> list[GitLabBranch]:
|
||||
"""Get list of branches in the repository."""
|
||||
url: Final = f"{self.base_url}/projects/{self._project_enc}/repository/branches"
|
||||
try:
|
||||
resp: Final = self.http_handler.get(url, headers=self.headers)
|
||||
resp.raise_for_status()
|
||||
data: Final = resp.json()
|
||||
data: Final = _branch_entries(resp)
|
||||
return data if isinstance(data, list) else []
|
||||
except Exception as e:
|
||||
raise Exception(f"Failed to get branches: {e}")
|
||||
|
||||
def get_file_metadata(self, file_path: str, *, ref: str | None = None) -> dict[str, Any] | None:
|
||||
def get_file_metadata(self, file_path: str, *, ref: str | None = None) -> GitLabFileMetadata | None:
|
||||
"""
|
||||
Get minimal metadata about a file via RAW endpoint headers at a given ref.
|
||||
|
||||
|
|
|
|||
|
|
@ -2,10 +2,12 @@
|
|||
GitLab prompt manager with configurable prompts folder.
|
||||
"""
|
||||
|
||||
from typing import TYPE_CHECKING, Any, Final
|
||||
from collections.abc import Mapping
|
||||
from typing import TYPE_CHECKING, Any, Final, TypeVar
|
||||
|
||||
from jinja2 import DictLoader, select_autoescape
|
||||
from jinja2.sandbox import ImmutableSandboxedEnvironment
|
||||
from typing_extensions import ReadOnly, TypedDict
|
||||
|
||||
from litellm.integrations.custom_prompt_management import CustomPromptManagement
|
||||
|
||||
|
|
@ -24,6 +26,19 @@ from litellm.types.utils import StandardCallbackDynamicParams
|
|||
|
||||
GITLAB_PREFIX: Final = "gitlab::"
|
||||
|
||||
_ResponseT = TypeVar("_ResponseT")
|
||||
|
||||
|
||||
class GitLabCachedPrompt(TypedDict):
|
||||
id: ReadOnly[str]
|
||||
path: ReadOnly[str]
|
||||
content: ReadOnly[str]
|
||||
metadata: ReadOnly[Mapping[str, object]]
|
||||
model: ReadOnly[str | None]
|
||||
temperature: ReadOnly[float | None]
|
||||
max_tokens: ReadOnly[int | None]
|
||||
optional_params: ReadOnly[Mapping[str, object]]
|
||||
|
||||
|
||||
def encode_prompt_id(raw_id: str) -> str:
|
||||
"""Convert GitLab path IDs like 'invoice/extract' → 'gitlab::invoice::extract'"""
|
||||
|
|
@ -206,7 +221,7 @@ class GitLabTemplateManager:
|
|||
result[key] = value.strip("\"'")
|
||||
return result
|
||||
|
||||
def render_template(self, template_id: str, variables: dict[str, Any] | None = None) -> str:
|
||||
def render_template(self, template_id: str, variables: Mapping[str, object] | None = None) -> str:
|
||||
if template_id not in self.prompts:
|
||||
raise ValueError(f"Template '{template_id}' not found")
|
||||
template: Final = self.prompts[template_id]
|
||||
|
|
@ -313,7 +328,7 @@ class GitLabPromptManager(CustomPromptManagement):
|
|||
def get_prompt_template(
|
||||
self,
|
||||
prompt_id: str,
|
||||
prompt_variables: dict[str, Any] | None = None,
|
||||
prompt_variables: Mapping[str, object] | None = None,
|
||||
*,
|
||||
ref: str | None = None,
|
||||
) -> tuple[str, dict[str, Any]]:
|
||||
|
|
@ -338,13 +353,13 @@ class GitLabPromptManager(CustomPromptManagement):
|
|||
self,
|
||||
user_id: str | None,
|
||||
messages: list[AllMessageValues],
|
||||
function_call: dict[str, Any] | str | None = None,
|
||||
litellm_params: dict[str, Any] | None = None,
|
||||
function_call: Mapping[str, object] | str | None = None,
|
||||
litellm_params: dict[str, object] | None = None,
|
||||
prompt_id: str | None = None,
|
||||
prompt_variables: dict[str, Any] | None = None,
|
||||
prompt_variables: Mapping[str, object] | None = None,
|
||||
prompt_version: str | None = None,
|
||||
**kwargs,
|
||||
) -> tuple[list[AllMessageValues], dict[str, Any] | None]:
|
||||
) -> tuple[list[AllMessageValues], dict[str, object] | None]:
|
||||
if not prompt_id:
|
||||
return messages, litellm_params
|
||||
try:
|
||||
|
|
@ -377,9 +392,9 @@ class GitLabPromptManager(CustomPromptManagement):
|
|||
|
||||
return final_messages, litellm_params
|
||||
except Exception as e:
|
||||
import litellm
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
|
||||
litellm._logging.verbose_proxy_logger.error("Error in GitLab prompt pre_call_hook: %s", e)
|
||||
verbose_proxy_logger.error("Error in GitLab prompt pre_call_hook: %s", e)
|
||||
return messages, litellm_params
|
||||
|
||||
def _parse_prompt_to_messages(self, prompt_content: str) -> list[AllMessageValues]:
|
||||
|
|
@ -435,14 +450,14 @@ class GitLabPromptManager(CustomPromptManagement):
|
|||
def post_call_hook(
|
||||
self,
|
||||
user_id: str | None,
|
||||
response: Any,
|
||||
response: _ResponseT,
|
||||
input_messages: list[AllMessageValues],
|
||||
function_call: dict[str, Any] | str | None = None,
|
||||
litellm_params: dict[str, Any] | None = None,
|
||||
function_call: Mapping[str, object] | str | None = None,
|
||||
litellm_params: Mapping[str, object] | None = None,
|
||||
prompt_id: str | None = None,
|
||||
prompt_variables: dict[str, Any] | None = None,
|
||||
prompt_variables: Mapping[str, object] | None = None,
|
||||
**kwargs,
|
||||
) -> Any:
|
||||
) -> _ResponseT:
|
||||
return response
|
||||
|
||||
def get_available_prompts(self) -> list[str]:
|
||||
|
|
@ -498,7 +513,7 @@ class GitLabPromptManager(CustomPromptManagement):
|
|||
messages: Final = self._parse_prompt_to_messages(rendered_prompt)
|
||||
template_model: Final = prompt_metadata.get("model")
|
||||
|
||||
optional_params: Final[dict[str, Any]] = {}
|
||||
optional_params: Final[dict[str, object]] = {}
|
||||
for param in [
|
||||
"temperature",
|
||||
"max_tokens",
|
||||
|
|
@ -658,14 +673,14 @@ class GitLabPromptCache:
|
|||
self.template_manager: GitLabTemplateManager = self.prompt_manager.prompt_manager
|
||||
|
||||
# In-memory stores
|
||||
self._by_file: dict[str, dict[str, Any]] = {}
|
||||
self._by_id: dict[str, dict[str, Any]] = {}
|
||||
self._by_file: dict[str, GitLabCachedPrompt] = {}
|
||||
self._by_id: dict[str, GitLabCachedPrompt] = {}
|
||||
|
||||
# -------------------------
|
||||
# Public API
|
||||
# -------------------------
|
||||
|
||||
def load_all(self, *, recursive: bool = True) -> dict[str, dict[str, Any]]:
|
||||
def load_all(self, *, recursive: bool = True) -> dict[str, GitLabCachedPrompt]:
|
||||
"""
|
||||
Scan GitLab for all .prompt files under prompts_path, load and parse each,
|
||||
and return the mapping of repo file path -> JSON-like dict.
|
||||
|
|
@ -695,7 +710,7 @@ class GitLabPromptCache:
|
|||
|
||||
return self._by_id
|
||||
|
||||
def reload(self, *, recursive: bool = True) -> dict[str, dict[str, Any]]:
|
||||
def reload(self, *, recursive: bool = True) -> dict[str, GitLabCachedPrompt]:
|
||||
"""Clear the cache and re-load from GitLab."""
|
||||
self._by_file.clear()
|
||||
self._by_id.clear()
|
||||
|
|
@ -709,11 +724,11 @@ class GitLabPromptCache:
|
|||
"""Return the template IDs (relative to prompts_path, without extension) currently cached."""
|
||||
return list(self._by_id.keys())
|
||||
|
||||
def get_by_file(self, file_path: str) -> dict[str, Any] | None:
|
||||
def get_by_file(self, file_path: str) -> GitLabCachedPrompt | None:
|
||||
"""Get a cached prompt JSON by repo file path."""
|
||||
return self._by_file.get(file_path)
|
||||
|
||||
def get_by_id(self, prompt_id: str) -> dict[str, Any] | None:
|
||||
def get_by_id(self, prompt_id: str) -> GitLabCachedPrompt | None:
|
||||
"""Get a cached prompt JSON by prompt ID (relative to prompts_path)."""
|
||||
if prompt_id in self._by_id:
|
||||
return self._by_id[prompt_id]
|
||||
|
|
@ -728,7 +743,7 @@ class GitLabPromptCache:
|
|||
# Internals
|
||||
# -------------------------
|
||||
|
||||
def _template_to_json(self, prompt_id: str, tmpl: GitLabPromptTemplate) -> dict[str, Any]:
|
||||
def _template_to_json(self, prompt_id: str, tmpl: GitLabPromptTemplate) -> GitLabCachedPrompt:
|
||||
"""
|
||||
Normalize a GitLabPromptTemplate into a JSON-like dict that is easy to serialize.
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -89,7 +89,7 @@ def _extract_cache_read_input_tokens(usage_obj) -> int:
|
|||
|
||||
# Check prompt_tokens_details.cached_tokens (used by Gemini and other providers)
|
||||
if hasattr(usage_obj, "prompt_tokens_details"):
|
||||
prompt_tokens_details: Final = getattr(usage_obj, "prompt_tokens_details", None)
|
||||
prompt_tokens_details: Final[object] = getattr(usage_obj, "prompt_tokens_details", None)
|
||||
if prompt_tokens_details is not None and hasattr(prompt_tokens_details, "cached_tokens"):
|
||||
cached_tokens: Final = getattr(prompt_tokens_details, "cached_tokens", None)
|
||||
if cached_tokens is not None and isinstance(cached_tokens, (int, float)) and cached_tokens > 0:
|
||||
|
|
@ -623,9 +623,16 @@ class LangFuseLogger:
|
|||
)
|
||||
|
||||
# Apply custom masking function if provided
|
||||
if masking_function is not None and callable(masking_function):
|
||||
input = self._apply_masking_function(input, masking_function)
|
||||
output = self._apply_masking_function(output, masking_function)
|
||||
masked_input: Final[object] = (
|
||||
self._apply_masking_function(input, masking_function)
|
||||
if masking_function is not None and callable(masking_function)
|
||||
else input
|
||||
)
|
||||
masked_output: Final[object] = (
|
||||
self._apply_masking_function(output, masking_function)
|
||||
if masking_function is not None and callable(masking_function)
|
||||
else output
|
||||
)
|
||||
|
||||
clean_metadata = redact_user_api_key_info(metadata=clean_metadata)
|
||||
|
||||
|
|
@ -651,15 +658,15 @@ class LangFuseLogger:
|
|||
|
||||
# Special keys that are found in the function arguments and not the metadata
|
||||
if "input" in update_trace_keys:
|
||||
trace_params["input"] = input if not mask_input else "redacted-by-litellm"
|
||||
trace_params["input"] = masked_input if not mask_input else "redacted-by-litellm"
|
||||
if "output" in update_trace_keys:
|
||||
trace_params["output"] = output if not mask_output else "redacted-by-litellm"
|
||||
trace_params["output"] = masked_output if not mask_output else "redacted-by-litellm"
|
||||
else: # don't overwrite an existing trace
|
||||
trace_params = {
|
||||
"id": trace_id,
|
||||
"name": trace_name,
|
||||
"session_id": session_id,
|
||||
"input": input if not mask_input else "redacted-by-litellm",
|
||||
"input": masked_input if not mask_input else "redacted-by-litellm",
|
||||
"version": clean_metadata.pop(
|
||||
"trace_version", clean_metadata.get("version", None)
|
||||
), # If provided just version, it will applied to the trace as well, if applied a trace version it will take precedence
|
||||
|
|
@ -669,9 +676,9 @@ class LangFuseLogger:
|
|||
trace_params[key.replace("trace_", "")] = clean_metadata.pop(key, None)
|
||||
|
||||
if level == "ERROR":
|
||||
trace_params["status_message"] = output
|
||||
trace_params["status_message"] = masked_output
|
||||
else:
|
||||
trace_params["output"] = output if not mask_output else "redacted-by-litellm"
|
||||
trace_params["output"] = masked_output if not mask_output else "redacted-by-litellm"
|
||||
|
||||
if debug is True or (isinstance(debug, str) and debug.lower() == "true"):
|
||||
debug_metadata: Final = {
|
||||
|
|
@ -708,7 +715,7 @@ class LangFuseLogger:
|
|||
("aws_region_name", aws_region_name, bool(aws_region_name)),
|
||||
("cache_hit", kwargs.get("cache_hit") or False, self._supports_tags() and "cache_hit" in kwargs),
|
||||
)
|
||||
enrichments: Final[Mapping[str, Any]] = {
|
||||
enrichments: Final[Mapping[str, object]] = {
|
||||
key: value for key, value, include in candidate_enrichments if include
|
||||
}
|
||||
|
||||
|
|
@ -802,8 +809,8 @@ class LangFuseLogger:
|
|||
"end_time": end_time,
|
||||
"model": model_name,
|
||||
"model_parameters": optional_params,
|
||||
"input": input if not mask_input else "redacted-by-litellm",
|
||||
"output": output if not mask_output else "redacted-by-litellm",
|
||||
"input": masked_input if not mask_input else "redacted-by-litellm",
|
||||
"output": masked_output if not mask_output else "redacted-by-litellm",
|
||||
"usage": usage,
|
||||
"usage_details": usage_details,
|
||||
"metadata": {
|
||||
|
|
@ -825,8 +832,8 @@ class LangFuseLogger:
|
|||
prompt_management_metadata=prompt_management_metadata,
|
||||
langfuse_client=self.Langfuse,
|
||||
)
|
||||
if output is not None and isinstance(output, str) and level == "ERROR":
|
||||
generation_params["status_message"] = output
|
||||
if masked_output is not None and isinstance(masked_output, str) and level == "ERROR":
|
||||
generation_params["status_message"] = masked_output
|
||||
|
||||
if self._supports_completion_start_time():
|
||||
generation_params["completion_start_time"] = kwargs.get("completion_start_time", None)
|
||||
|
|
@ -935,7 +942,7 @@ class LangFuseLogger:
|
|||
return Version(self.langfuse_sdk_version) >= Version("2.7.3")
|
||||
|
||||
@staticmethod
|
||||
def _apply_masking_function(data: Any, masking_function: Callable[[Any], Any]) -> Any:
|
||||
def _apply_masking_function(data: object, masking_function: Callable[[object], object]) -> object:
|
||||
"""
|
||||
Apply a masking function to data, handling different data types.
|
||||
|
||||
|
|
@ -1049,7 +1056,7 @@ def _add_prompt_to_generation_params(
|
|||
generation_params: dict,
|
||||
clean_metadata: dict,
|
||||
prompt_management_metadata: StandardLoggingPromptManagementMetadata | None,
|
||||
langfuse_client: Any,
|
||||
langfuse_client: object,
|
||||
) -> dict:
|
||||
from langfuse import Langfuse
|
||||
from langfuse.model import (
|
||||
|
|
|
|||
|
|
@ -4,9 +4,12 @@ Opik Logger that logs LLM events to an Opik server
|
|||
|
||||
import asyncio
|
||||
import traceback
|
||||
from collections.abc import Mapping
|
||||
from datetime import datetime
|
||||
from typing import Any, Final
|
||||
|
||||
from typing_extensions import ReadOnly, TypedDict, Unpack
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.integrations.custom_batch_logger import CustomBatchLogger
|
||||
from litellm.llms.custom_httpx.http_handler import (
|
||||
|
|
@ -23,7 +26,7 @@ except Exception:
|
|||
opik_client = None
|
||||
|
||||
|
||||
def _should_skip_event(kwargs: dict[str, Any]) -> bool:
|
||||
def _should_skip_event(kwargs: Mapping[str, object]) -> bool:
|
||||
"""Check if event should be skipped due to missing standard_logging_object."""
|
||||
if kwargs.get("standard_logging_object") is None:
|
||||
verbose_logger.debug("OpikLogger skipping event; no standard_logging_object found")
|
||||
|
|
@ -31,12 +34,24 @@ def _should_skip_event(kwargs: dict[str, Any]) -> bool:
|
|||
return False
|
||||
|
||||
|
||||
class _OpikLoggerKwargs(TypedDict, total=False):
|
||||
"""Constructor options accepted by ``OpikLogger``."""
|
||||
|
||||
project_name: ReadOnly[str | None]
|
||||
url: ReadOnly[str | None]
|
||||
api_key: ReadOnly[str | None]
|
||||
workspace: ReadOnly[str | None]
|
||||
batch_size: ReadOnly[int | None]
|
||||
flush_interval: ReadOnly[int | None]
|
||||
max_queue_size: ReadOnly[int | None]
|
||||
|
||||
|
||||
class OpikLogger(CustomBatchLogger):
|
||||
"""
|
||||
Opik Logger for logging events to an Opik Server
|
||||
"""
|
||||
|
||||
def __init__(self, **kwargs: Any) -> None:
|
||||
def __init__(self, **kwargs: Unpack[_OpikLoggerKwargs]) -> None:
|
||||
self.async_httpx_client = get_async_httpx_client(llm_provider=httpxSpecialProvider.LoggingCallback)
|
||||
self.sync_httpx_client = _get_httpx_client()
|
||||
|
||||
|
|
@ -95,7 +110,7 @@ class OpikLogger(CustomBatchLogger):
|
|||
|
||||
async def async_log_success_event(
|
||||
self,
|
||||
kwargs: dict[str, Any],
|
||||
kwargs: dict[str, object],
|
||||
response_obj: Any,
|
||||
start_time: datetime,
|
||||
end_time: datetime,
|
||||
|
|
@ -163,7 +178,7 @@ class OpikLogger(CustomBatchLogger):
|
|||
except Exception as e:
|
||||
verbose_logger.exception("OpikLogger failed to log success event - %s\n%s", e, traceback.format_exc())
|
||||
|
||||
def _sync_send(self, url: str, headers: dict[str, str], batch: dict[str, Any]) -> None:
|
||||
def _sync_send(self, url: str, headers: dict[str, str], batch: dict[str, object]) -> None:
|
||||
try:
|
||||
response: Final = self.sync_httpx_client.post(
|
||||
url=url,
|
||||
|
|
@ -178,7 +193,7 @@ class OpikLogger(CustomBatchLogger):
|
|||
|
||||
def log_success_event(
|
||||
self,
|
||||
kwargs: dict[str, Any],
|
||||
kwargs: dict[str, object],
|
||||
response_obj: Any,
|
||||
start_time: datetime,
|
||||
end_time: datetime,
|
||||
|
|
@ -247,7 +262,7 @@ class OpikLogger(CustomBatchLogger):
|
|||
except Exception as e:
|
||||
verbose_logger.exception("OpikLogger failed to log success event - %s\n%s", e, traceback.format_exc())
|
||||
|
||||
async def _submit_batch(self, url: str, headers: dict[str, str], batch: dict[str, Any]) -> None:
|
||||
async def _submit_batch(self, url: str, headers: dict[str, str], batch: dict[str, object]) -> None:
|
||||
try:
|
||||
response: Final = await self.async_httpx_client.post(
|
||||
url=url,
|
||||
|
|
|
|||
|
|
@ -1,6 +1,7 @@
|
|||
"""Data extraction functions for Opik payload building."""
|
||||
|
||||
import json
|
||||
from collections.abc import Mapping
|
||||
from typing import Any, Final
|
||||
|
||||
from litellm import _logging
|
||||
|
|
@ -35,8 +36,8 @@ def normalize_provider_name(provider: str | None) -> str | None:
|
|||
|
||||
|
||||
def extract_opik_metadata(
|
||||
litellm_metadata: dict[str, Any],
|
||||
standard_logging_metadata: dict[str, Any],
|
||||
litellm_metadata: Mapping[str, Any],
|
||||
standard_logging_metadata: Mapping[str, Any],
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
Merge Opik metadata from three sources in increasing priority order:
|
||||
|
|
@ -97,7 +98,7 @@ def extract_span_identifiers(
|
|||
|
||||
|
||||
def extract_tags(
|
||||
opik_metadata: dict[str, Any],
|
||||
opik_metadata: Mapping[str, Any],
|
||||
custom_llm_provider: str | None,
|
||||
) -> list[str]:
|
||||
"""
|
||||
|
|
@ -122,7 +123,7 @@ def apply_proxy_header_overrides(
|
|||
project_name: str,
|
||||
tags: list[str],
|
||||
thread_id: str | None,
|
||||
proxy_headers: dict[str, Any],
|
||||
proxy_headers: Mapping[str, str],
|
||||
) -> tuple[str, list[str], str | None]:
|
||||
"""
|
||||
Apply overrides from proxy request headers (opik_* prefix).
|
||||
|
|
@ -148,7 +149,7 @@ def apply_proxy_header_overrides(
|
|||
thread_id = value
|
||||
elif param_key == "tags":
|
||||
try:
|
||||
parsed_tags = json.loads(value)
|
||||
parsed_tags: object = json.loads(value)
|
||||
if isinstance(parsed_tags, list):
|
||||
tags.extend(parsed_tags)
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
|
|
@ -158,11 +159,11 @@ def apply_proxy_header_overrides(
|
|||
|
||||
|
||||
def extract_and_build_metadata(
|
||||
opik_metadata: dict[str, Any],
|
||||
standard_logging_metadata: dict[str, Any],
|
||||
standard_logging_object: dict[str, Any],
|
||||
litellm_kwargs: dict[str, Any],
|
||||
) -> dict[str, Any]:
|
||||
opik_metadata: Mapping[str, object],
|
||||
standard_logging_metadata: Mapping[str, object],
|
||||
standard_logging_object: Mapping[str, object],
|
||||
litellm_kwargs: Mapping[str, object],
|
||||
) -> dict[str, object]:
|
||||
"""
|
||||
Build the complete metadata dictionary from all available sources.
|
||||
|
||||
|
|
|
|||
|
|
@ -17,12 +17,12 @@ def build_trace_payload(
|
|||
end_time: datetime,
|
||||
input_data: Any,
|
||||
output_data: Any,
|
||||
metadata: dict[str, Any],
|
||||
metadata: dict[str, object],
|
||||
tags: list[str],
|
||||
thread_id: str | None,
|
||||
) -> types.TracePayload:
|
||||
"""Build a complete trace payload."""
|
||||
trace_name: Final = response_obj.get("object", "unknown type")
|
||||
trace_name: Final[str] = response_obj.get("object", "unknown type")
|
||||
|
||||
return types.TracePayload(
|
||||
project_name=project_name,
|
||||
|
|
@ -47,7 +47,7 @@ def build_span_payload(
|
|||
end_time: datetime,
|
||||
input_data: Any,
|
||||
output_data: Any,
|
||||
metadata: dict[str, Any],
|
||||
metadata: dict[str, object],
|
||||
tags: list[str],
|
||||
usage: dict[str, int],
|
||||
provider: str | None = None,
|
||||
|
|
@ -56,9 +56,9 @@ def build_span_payload(
|
|||
"""Build a complete span payload."""
|
||||
span_id: Final = utils.create_uuid7()
|
||||
|
||||
model: Final = response_obj.get("model", "unknown-model")
|
||||
obj_type: Final = response_obj.get("object", "unknown-object")
|
||||
created: Final = response_obj.get("created", 0)
|
||||
model: Final[str] = response_obj.get("model", "unknown-model")
|
||||
obj_type: Final[str] = response_obj.get("object", "unknown-object")
|
||||
created: Final[int] = response_obj.get("created", 0)
|
||||
span_name: Final = f"{model}_{obj_type}_{created}"
|
||||
|
||||
_logging.verbose_logger.debug("OpikLogger creating span with id %s for trace %s", span_id, trace_id)
|
||||
|
|
|
|||
|
|
@ -62,6 +62,8 @@ class GenAIMapper:
|
|||
GenAI.RESPONSE_TIME_TO_FIRST_CHUNK: lambda d: d.time_to_first_chunk_seconds,
|
||||
GenAI.USAGE_INPUT_TOKENS: lambda d: d.usage.input_tokens,
|
||||
GenAI.USAGE_OUTPUT_TOKENS: lambda d: d.usage.output_tokens,
|
||||
GenAI.USAGE_CACHE_CREATION_INPUT_TOKENS: lambda d: d.usage.cache_creation_input_tokens,
|
||||
GenAI.USAGE_CACHE_READ_INPUT_TOKENS: lambda d: d.usage.cache_read_input_tokens,
|
||||
Error.TYPE: lambda d: d.error.error_type if d.error else None,
|
||||
Server.ADDRESS: lambda d: d.server.address if d.server else None,
|
||||
Server.PORT: lambda d: d.server.port if d.server else None,
|
||||
|
|
|
|||
|
|
@ -6,6 +6,7 @@ import json
|
|||
from collections.abc import Mapping
|
||||
from dataclasses import dataclass, field
|
||||
from enum import Enum
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, ClassVar, Final, cast
|
||||
from urllib.parse import urlsplit
|
||||
|
||||
|
|
@ -62,6 +63,31 @@ if TYPE_CHECKING:
|
|||
# --- typed sub-structures ---------------------------------------------------- #
|
||||
|
||||
|
||||
def _cache_token_value(*values: object) -> int | None:
|
||||
explicit_zero = False
|
||||
invalid_before_zero = False
|
||||
for raw_value in values:
|
||||
if raw_value is None:
|
||||
continue
|
||||
if isinstance(raw_value, bool):
|
||||
parsed = None
|
||||
else:
|
||||
try:
|
||||
parsed = as_int(raw_value)
|
||||
except (OverflowError, ValueError):
|
||||
parsed = None
|
||||
if parsed is None:
|
||||
if not explicit_zero:
|
||||
invalid_before_zero = True
|
||||
elif parsed > 0:
|
||||
return parsed
|
||||
elif parsed == 0:
|
||||
explicit_zero = True
|
||||
elif not explicit_zero:
|
||||
invalid_before_zero = True
|
||||
return 0 if explicit_zero and not invalid_before_zero else None
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class LLMRequestParams:
|
||||
temperature: float | None = None
|
||||
|
|
@ -95,6 +121,35 @@ class LLMUsage:
|
|||
input_tokens: int | None = None
|
||||
output_tokens: int | None = None
|
||||
total_tokens: int | None = None
|
||||
cache_creation_input_tokens: int | None = None
|
||||
cache_read_input_tokens: int | None = None
|
||||
|
||||
@classmethod
|
||||
def from_standard_logging_payload(cls, payload: StandardLoggingPayload) -> LLMUsage:
|
||||
# Cache token counts only exist on the raw provider usage object under metadata
|
||||
metadata: Final[Mapping[str, object]] = payload.get("metadata") or {}
|
||||
raw_usage: Final = metadata.get("usage_object")
|
||||
usage_object: Final[Mapping[str, object]] = raw_usage if isinstance(raw_usage, Mapping) else {}
|
||||
raw_details: Final = usage_object.get("prompt_tokens_details")
|
||||
prompt_details: Final[Mapping[str, object]] = (
|
||||
raw_details if isinstance(raw_details, Mapping) else MappingProxyType({})
|
||||
)
|
||||
return cls(
|
||||
input_tokens=as_int(payload.get("prompt_tokens")),
|
||||
output_tokens=as_int(payload.get("completion_tokens")),
|
||||
total_tokens=as_int(payload.get("total_tokens")),
|
||||
cache_creation_input_tokens=_cache_token_value(
|
||||
usage_object.get("cache_creation_input_tokens"),
|
||||
prompt_details.get("cache_write_tokens"),
|
||||
prompt_details.get("cache_creation_tokens"),
|
||||
prompt_details.get("cache_creation_input_tokens"),
|
||||
),
|
||||
cache_read_input_tokens=_cache_token_value(
|
||||
usage_object.get("cache_read_input_tokens"),
|
||||
prompt_details.get("cached_tokens"),
|
||||
usage_object.get("prompt_cache_hit_tokens"),
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
|
|
@ -363,11 +418,7 @@ class LLMCallSpanData:
|
|||
response_model=context.response_model,
|
||||
response_id=as_str(response.get("id")),
|
||||
request_params=LLMRequestParams.from_model_parameters(params),
|
||||
usage=LLMUsage(
|
||||
input_tokens=as_int(payload.get("prompt_tokens")),
|
||||
output_tokens=as_int(payload.get("completion_tokens")),
|
||||
total_tokens=as_int(payload.get("total_tokens")),
|
||||
),
|
||||
usage=LLMUsage.from_standard_logging_payload(payload),
|
||||
finish_reasons=finish_reasons,
|
||||
error=_parse_error(payload),
|
||||
response_cost=as_float(payload.get("response_cost")),
|
||||
|
|
|
|||
|
|
@ -110,6 +110,8 @@ class GenAI:
|
|||
# usage
|
||||
USAGE_INPUT_TOKENS: Final = "gen_ai.usage.input_tokens"
|
||||
USAGE_OUTPUT_TOKENS: Final = "gen_ai.usage.output_tokens"
|
||||
USAGE_CACHE_CREATION_INPUT_TOKENS: Final = "gen_ai.usage.cache_creation.input_tokens"
|
||||
USAGE_CACHE_READ_INPUT_TOKENS: Final = "gen_ai.usage.cache_read.input_tokens"
|
||||
# content (opt-in, gated by capture mode)
|
||||
INPUT_MESSAGES: Final = "gen_ai.input.messages"
|
||||
OUTPUT_MESSAGES: Final = "gen_ai.output.messages"
|
||||
|
|
|
|||
|
|
@ -11,9 +11,10 @@ identical metrics. The attribute cardinality filter is reused from v1 by import
|
|||
from collections.abc import Mapping
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime
|
||||
from typing import Any, Final, TypeAlias
|
||||
from typing import Any, Final, Literal, Protocol, TypeAlias
|
||||
|
||||
from opentelemetry.metrics import Histogram, Meter
|
||||
from typing_extensions import ReadOnly, TypedDict
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_logger
|
||||
|
|
@ -151,6 +152,29 @@ METRIC_ATTRIBUTE_CEILING: Final[frozenset[str]] = frozenset(
|
|||
BOUNDED_HIDDEN_PARAM_KEYS: Final[tuple[str, ...]] = ("model_id",)
|
||||
|
||||
|
||||
class _TokenUsage(TypedDict, total=False):
|
||||
"""The token counts a response's ``usage`` carries, as the recorder reads them."""
|
||||
|
||||
prompt_tokens: ReadOnly[int]
|
||||
completion_tokens: ReadOnly[int]
|
||||
|
||||
|
||||
class _ResponseView(Protocol):
|
||||
"""The one read the recorder makes on a litellm response object."""
|
||||
|
||||
def get(self, key: Literal["usage"], /) -> _TokenUsage | None: ...
|
||||
|
||||
|
||||
class _MetricKwargs(TypedDict, total=False):
|
||||
"""The logging kwargs the recorder reads directly."""
|
||||
|
||||
call_type: ReadOnly[str | None]
|
||||
litellm_params: ReadOnly[Mapping[str, object] | None]
|
||||
response_cost: ReadOnly[float | None]
|
||||
completion_start_time: ReadOnly[datetime | float | str | None]
|
||||
api_call_start_time: ReadOnly[datetime | float | str | None]
|
||||
|
||||
|
||||
def resolve_error_type(kwargs: Mapping[str, Any]) -> str:
|
||||
"""The ``error.type`` value for a failed request.
|
||||
|
||||
|
|
@ -192,8 +216,8 @@ class GenAIMetricRecorder:
|
|||
|
||||
def record(
|
||||
self,
|
||||
kwargs: Mapping[str, Any],
|
||||
response_obj: Any,
|
||||
kwargs: _MetricKwargs,
|
||||
response_obj: _ResponseView | None,
|
||||
start_time: datetime,
|
||||
end_time: datetime,
|
||||
) -> None:
|
||||
|
|
@ -218,7 +242,7 @@ class GenAIMetricRecorder:
|
|||
|
||||
def record_failure(
|
||||
self,
|
||||
kwargs: Mapping[str, Any],
|
||||
kwargs: _MetricKwargs,
|
||||
start_time: datetime,
|
||||
end_time: datetime,
|
||||
) -> None:
|
||||
|
|
@ -342,7 +366,7 @@ class GenAIMetricRecorder:
|
|||
# Per-metric recording
|
||||
# ------------------------------------------------------------------ #
|
||||
|
||||
def _record_token_usage(self, response_obj: Any, common_attrs: dict) -> None:
|
||||
def _record_token_usage(self, response_obj: _ResponseView | None, common_attrs: dict) -> None:
|
||||
if not response_obj:
|
||||
return
|
||||
usage: Final = response_obj.get("usage")
|
||||
|
|
@ -353,7 +377,7 @@ class GenAIMetricRecorder:
|
|||
self._metrics.token_usage.record(usage.get("prompt_tokens", 0), attributes=in_attrs)
|
||||
self._metrics.token_usage.record(usage.get("completion_tokens", 0), attributes=out_attrs)
|
||||
|
||||
def _record_time_to_first_token(self, kwargs: Mapping[str, Any], common_attrs: dict) -> None:
|
||||
def _record_time_to_first_token(self, kwargs: _MetricKwargs, common_attrs: dict) -> None:
|
||||
time_to_first_chunk: Final = time_to_first_chunk_seconds(kwargs)
|
||||
if time_to_first_chunk is None:
|
||||
return
|
||||
|
|
@ -361,15 +385,14 @@ class GenAIMetricRecorder:
|
|||
|
||||
def _record_time_per_output_token(
|
||||
self,
|
||||
kwargs: Mapping[str, Any],
|
||||
response_obj: Any,
|
||||
kwargs: _MetricKwargs,
|
||||
response_obj: _ResponseView | None,
|
||||
end_time: datetime,
|
||||
duration_s: float,
|
||||
common_attrs: dict,
|
||||
) -> None:
|
||||
completion_tokens = None
|
||||
if response_obj and (usage := response_obj.get("usage")):
|
||||
completion_tokens = usage.get("completion_tokens")
|
||||
usage: Final = response_obj.get("usage") if response_obj else None
|
||||
completion_tokens: Final = usage.get("completion_tokens") if usage else None
|
||||
if completion_tokens is None or completion_tokens <= 0:
|
||||
return
|
||||
|
||||
|
|
|
|||
|
|
@ -12,7 +12,10 @@ For batching specific details see CustomBatchLogger class
|
|||
import asyncio
|
||||
import atexit
|
||||
import os
|
||||
from typing import Any, Final
|
||||
from collections.abc import Mapping, Sequence
|
||||
from typing import Final
|
||||
|
||||
from typing_extensions import ReadOnly, TypedDict
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm._uuid import uuid
|
||||
|
|
@ -34,6 +37,21 @@ from litellm.types.integrations.posthog import (
|
|||
from litellm.types.utils import StandardCallbackDynamicParams, StandardLoggingPayload
|
||||
|
||||
|
||||
class PostHogBatchPayload(TypedDict):
|
||||
api_key: ReadOnly[str]
|
||||
batch: ReadOnly[Sequence[PostHogEventPayload]]
|
||||
|
||||
|
||||
class PostHogLiteLLMParams(TypedDict, total=False):
|
||||
metadata: ReadOnly[Mapping[str, object]]
|
||||
|
||||
|
||||
class PostHogLogKwargs(TypedDict, total=False):
|
||||
standard_logging_object: ReadOnly[StandardLoggingPayload]
|
||||
standard_callback_dynamic_params: ReadOnly[StandardCallbackDynamicParams]
|
||||
litellm_params: ReadOnly[PostHogLiteLLMParams]
|
||||
|
||||
|
||||
class PostHogLogger(CustomBatchLogger):
|
||||
def __init__(self, **kwargs):
|
||||
"""
|
||||
|
|
@ -137,7 +155,7 @@ class PostHogLogger(CustomBatchLogger):
|
|||
if len(self.log_queue) >= self.batch_size:
|
||||
await self.flush_queue()
|
||||
|
||||
def create_posthog_event_payload(self, kwargs: dict[str, Any]) -> PostHogEventPayload:
|
||||
def create_posthog_event_payload(self, kwargs: PostHogLogKwargs) -> PostHogEventPayload:
|
||||
"""
|
||||
Helper function to create a PostHog event payload for logging
|
||||
|
||||
|
|
@ -171,11 +189,11 @@ class PostHogLogger(CustomBatchLogger):
|
|||
def _create_posthog_properties(
|
||||
self,
|
||||
standard_logging_object: StandardLoggingPayload,
|
||||
kwargs: dict[str, Any],
|
||||
kwargs: PostHogLogKwargs,
|
||||
event_name: str,
|
||||
) -> dict[str, Any]:
|
||||
) -> dict[str, object]:
|
||||
"""Create PostHog properties following LLM Analytics spec"""
|
||||
properties: Final = {}
|
||||
properties: Final[dict[str, object]] = {}
|
||||
|
||||
# Core model information
|
||||
properties["$ai_model"] = self._safe_get(standard_logging_object, "model", "")
|
||||
|
|
@ -211,16 +229,19 @@ class PostHogLogger(CustomBatchLogger):
|
|||
properties["$ai_error"] = error_str
|
||||
|
||||
# Add trace properties
|
||||
self._add_trace_properties(properties, kwargs)
|
||||
self._add_trace_properties(properties, standard_logging_object, kwargs)
|
||||
|
||||
# Add custom metadata fields
|
||||
self._add_custom_metadata_properties(properties, kwargs)
|
||||
|
||||
return properties
|
||||
|
||||
def _add_trace_properties(self, properties: dict[str, Any], kwargs: dict[str, Any]):
|
||||
standard_logging_object: Final = self._safe_get(kwargs, "standard_logging_object", {})
|
||||
|
||||
def _add_trace_properties(
|
||||
self,
|
||||
properties: dict[str, object],
|
||||
standard_logging_object: StandardLoggingPayload,
|
||||
kwargs: PostHogLogKwargs,
|
||||
) -> None:
|
||||
trace_id: Final = self._safe_get(standard_logging_object, "trace_id", self._safe_uuid())
|
||||
properties["$ai_trace_id"] = trace_id
|
||||
|
||||
|
|
@ -232,7 +253,7 @@ class PostHogLogger(CustomBatchLogger):
|
|||
if parent_id:
|
||||
properties["$ai_parent_id"] = parent_id
|
||||
|
||||
def _add_custom_metadata_properties(self, properties: dict[str, Any], kwargs: dict[str, Any]):
|
||||
def _add_custom_metadata_properties(self, properties: dict[str, object], kwargs: PostHogLogKwargs) -> None:
|
||||
"""Add custom metadata fields to PostHog properties"""
|
||||
metadata: Final = self._extract_metadata(kwargs)
|
||||
if not isinstance(metadata, dict):
|
||||
|
|
@ -277,7 +298,7 @@ class PostHogLogger(CustomBatchLogger):
|
|||
if key not in litellm_internal_fields:
|
||||
properties[key] = value
|
||||
|
||||
def _get_distinct_id(self, standard_logging_object: StandardLoggingPayload, kwargs: dict[str, Any]) -> str:
|
||||
def _get_distinct_id(self, standard_logging_object: StandardLoggingPayload, kwargs: PostHogLogKwargs) -> str:
|
||||
metadata: Final = self._extract_metadata(kwargs)
|
||||
user_id: Final = self._safe_get(metadata, "user_id")
|
||||
if user_id:
|
||||
|
|
@ -291,7 +312,7 @@ class PostHogLogger(CustomBatchLogger):
|
|||
|
||||
return self._safe_uuid()
|
||||
|
||||
def _get_credentials_for_request(self, kwargs: dict[str, Any]) -> tuple[str | None, str | None]:
|
||||
def _get_credentials_for_request(self, kwargs: PostHogLogKwargs) -> tuple[str | None, str | None]:
|
||||
"""
|
||||
Get PostHog credentials for this request.
|
||||
|
||||
|
|
@ -334,7 +355,7 @@ class PostHogLogger(CustomBatchLogger):
|
|||
verbose_logger.debug("[POSTHOG MOCK] Mock mode enabled - API calls will be intercepted")
|
||||
|
||||
# Group events by credentials for batch sending
|
||||
batches_by_credentials: Final[dict[tuple[str, str], list]] = {}
|
||||
batches_by_credentials: Final[dict[tuple[str, str], list[PostHogEventPayload]]] = {}
|
||||
for item in self.log_queue:
|
||||
key = (item["api_key"], item["api_url"])
|
||||
if key not in batches_by_credentials:
|
||||
|
|
@ -380,18 +401,19 @@ class PostHogLogger(CustomBatchLogger):
|
|||
verbose_logger.error("PostHog: Failed to initialize async components: %s", e)
|
||||
raise
|
||||
|
||||
def _extract_metadata(self, kwargs: dict[str, Any]) -> dict[str, Any]:
|
||||
litellm_params: Final = kwargs.get("litellm_params", {}) or {}
|
||||
return litellm_params.get("metadata", {}) or {}
|
||||
def _extract_metadata(self, kwargs: PostHogLogKwargs) -> Mapping[str, object]:
|
||||
litellm_params: Final[PostHogLiteLLMParams] = kwargs.get("litellm_params", {}) or {}
|
||||
metadata: Final[Mapping[str, object]] = litellm_params.get("metadata", {}) or {}
|
||||
return metadata
|
||||
|
||||
def _safe_uuid(self) -> str:
|
||||
return str(uuid.uuid4())
|
||||
|
||||
def _create_posthog_payload(self, events: list, api_key: str) -> dict[str, Any]:
|
||||
def _create_posthog_payload(self, events: Sequence[PostHogEventPayload], api_key: str) -> PostHogBatchPayload:
|
||||
return {"api_key": api_key, "batch": events}
|
||||
|
||||
def _safe_get(self, obj: Any, key: str, default: Any = None) -> Any:
|
||||
if obj is None or not hasattr(obj, "get"):
|
||||
def _safe_get(self, obj: Mapping[str, object] | None, key: str, default: object = None) -> object:
|
||||
if not isinstance(obj, Mapping):
|
||||
return default
|
||||
return obj.get(key, default)
|
||||
|
||||
|
|
@ -412,7 +434,7 @@ class PostHogLogger(CustomBatchLogger):
|
|||
|
||||
try:
|
||||
# Group events by credentials (same logic as async_send_batch)
|
||||
batches_by_credentials: Final[dict[tuple[str, str], list]] = {}
|
||||
batches_by_credentials: Final[dict[tuple[str, str], list[PostHogEventPayload]]] = {}
|
||||
for item in self.log_queue:
|
||||
key = (item["api_key"], item["api_url"])
|
||||
if key not in batches_by_credentials:
|
||||
|
|
|
|||
|
|
@ -8,6 +8,7 @@ import math
|
|||
import os
|
||||
import sys
|
||||
from collections.abc import Awaitable, Callable, Mapping, Sequence
|
||||
from dataclasses import replace
|
||||
from datetime import datetime, timedelta
|
||||
from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, TypeVar, cast
|
||||
|
||||
|
|
@ -58,7 +59,10 @@ from litellm.types.utils import (
|
|||
|
||||
if TYPE_CHECKING:
|
||||
from apscheduler.schedulers.asyncio import AsyncIOScheduler
|
||||
from prometheus_client import Gauge
|
||||
from prometheus_client.metrics import MetricWrapperBase
|
||||
|
||||
from litellm.router import Router
|
||||
else:
|
||||
AsyncIOScheduler = Any
|
||||
|
||||
|
|
@ -67,6 +71,8 @@ _TableRowT: Final = TypeVar("_TableRowT", bound=BaseModel)
|
|||
|
||||
_DEFAULT_BUDGET_METRICS_PER_REQUEST_TIMEOUT: Final = 5.0
|
||||
|
||||
UNRECOGNIZED_REQUESTED_MODEL_LABEL: Final = "other"
|
||||
|
||||
_NON_ENUM_METRIC_LABELS: Final[frozenset[str]] = frozenset(
|
||||
(
|
||||
"guardrail_name",
|
||||
|
|
@ -154,6 +160,44 @@ def _get_budget_metrics_per_request_timeout() -> float:
|
|||
return parsed
|
||||
|
||||
|
||||
def _get_proxy_llm_router() -> Router | None:
|
||||
try:
|
||||
from litellm.proxy.proxy_server import llm_router
|
||||
except Exception:
|
||||
return None
|
||||
return llm_router
|
||||
|
||||
|
||||
def _bounded_requested_model_label(requested_model: str | None, router_originated: bool = False) -> str | None:
|
||||
"""
|
||||
Bound ``requested_model`` label cardinality: names the router recognizes
|
||||
(model names, deployment ids, aliases, routing groups, team public model
|
||||
names) or matches via a global or team wildcard/pattern route keep their
|
||||
own label value; any other client-supplied string collapses into the
|
||||
single ``other`` bucket. With no proxy router to vouch for the string,
|
||||
client-supplied values collapse to ``other`` while ``router_originated``
|
||||
values (emitted by an SDK ``Router``'s own deployment failure and
|
||||
fallback events, where the proxy router never exists) pass through.
|
||||
"""
|
||||
if not requested_model:
|
||||
return requested_model
|
||||
llm_router: Final = _get_proxy_llm_router()
|
||||
if llm_router is None:
|
||||
return requested_model if router_originated else UNRECOGNIZED_REQUESTED_MODEL_LABEL
|
||||
if llm_router.is_recognized_model(requested_model):
|
||||
return requested_model
|
||||
if requested_model in llm_router.team_public_model_names:
|
||||
return requested_model
|
||||
if llm_router.pattern_router.route(requested_model) is not None:
|
||||
return requested_model
|
||||
if any(
|
||||
team_pattern_router.route(requested_model) is not None
|
||||
for team_pattern_router in llm_router.team_pattern_routers.values()
|
||||
):
|
||||
return requested_model
|
||||
return UNRECOGNIZED_REQUESTED_MODEL_LABEL
|
||||
|
||||
|
||||
class PrometheusLogger(CustomLogger):
|
||||
# Class variables or attributes
|
||||
|
||||
|
|
@ -434,6 +478,30 @@ class PrometheusLogger(CustomLogger):
|
|||
labelnames=self.get_labels_for_metric("litellm_remaining_api_key_tokens_for_model"),
|
||||
)
|
||||
|
||||
self.litellm_api_key_rate_limit_allowed_metric = self._gauge_factory(
|
||||
"litellm_api_key_rate_limit_allowed_metric",
|
||||
"Configured rate limit for the API Key in the current window (rpm_limit / tpm_limit), by rate_limit_type",
|
||||
labelnames=self.get_labels_for_metric("litellm_api_key_rate_limit_allowed_metric"),
|
||||
)
|
||||
|
||||
self.litellm_api_key_rate_limit_used_metric = self._gauge_factory(
|
||||
"litellm_api_key_rate_limit_used_metric",
|
||||
"Requests or tokens the API Key has consumed in the current rate limit window, by rate_limit_type",
|
||||
labelnames=self.get_labels_for_metric("litellm_api_key_rate_limit_used_metric"),
|
||||
)
|
||||
|
||||
self.litellm_team_rate_limit_allowed_metric = self._gauge_factory(
|
||||
"litellm_team_rate_limit_allowed_metric",
|
||||
"Configured rate limit for the Team in the current window (team rpm_limit / tpm_limit), by rate_limit_type",
|
||||
labelnames=self.get_labels_for_metric("litellm_team_rate_limit_allowed_metric"),
|
||||
)
|
||||
|
||||
self.litellm_team_rate_limit_used_metric = self._gauge_factory(
|
||||
"litellm_team_rate_limit_used_metric",
|
||||
"Requests or tokens the Team has consumed in the current rate limit window, by rate_limit_type",
|
||||
labelnames=self.get_labels_for_metric("litellm_team_rate_limit_used_metric"),
|
||||
)
|
||||
|
||||
########################################
|
||||
# LLM API Deployment Metrics / analytics
|
||||
########################################
|
||||
|
|
@ -1433,6 +1501,11 @@ class PrometheusLogger(CustomLogger):
|
|||
model_id=enum_values.model_id,
|
||||
)
|
||||
|
||||
self._set_key_and_team_rate_limit_metrics(
|
||||
standard_logging_payload=standard_logging_payload, # pyright: ignore[reportArgumentType] # isinstance(dict) above narrows the TypedDict to dict[Unknown, Unknown]
|
||||
enum_values=enum_values,
|
||||
)
|
||||
|
||||
# set latency metrics
|
||||
self._set_latency_metrics(
|
||||
kwargs=kwargs,
|
||||
|
|
@ -1960,17 +2033,102 @@ class PrometheusLogger(CustomLogger):
|
|||
"""
|
||||
if standard_logging_payload is None:
|
||||
return None
|
||||
return PrometheusLogger._get_int_from_v3_rate_limit_headers(
|
||||
standard_logging_payload=standard_logging_payload,
|
||||
header_name=f"x-ratelimit-model_per_key-remaining-{rate_limit_type}",
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _get_int_from_v3_rate_limit_headers(
|
||||
standard_logging_payload: StandardLoggingPayload,
|
||||
header_name: str,
|
||||
) -> int | None:
|
||||
hidden_params: Final = standard_logging_payload.get("hidden_params")
|
||||
if hidden_params is None:
|
||||
return None
|
||||
additional_headers: Final = hidden_params.get("additional_headers")
|
||||
additional_headers: Final[Mapping[str, object] | None] = hidden_params.get("additional_headers")
|
||||
if additional_headers is None:
|
||||
return None
|
||||
value: Final = dict(additional_headers).get(f"x-ratelimit-model_per_key-remaining-{rate_limit_type}")
|
||||
value: Final = additional_headers.get(header_name)
|
||||
if isinstance(value, bool) or not isinstance(value, int):
|
||||
return None
|
||||
return value
|
||||
|
||||
def _set_key_and_team_rate_limit_metrics(
|
||||
self,
|
||||
standard_logging_payload: StandardLoggingPayload,
|
||||
enum_values: UserAPIKeyLabelValues,
|
||||
) -> None:
|
||||
"""
|
||||
Export the key-level and team-level RPM / TPM limit and current window
|
||||
usage from the ``x-ratelimit-{api_key,team}-{limit,remaining}-*``
|
||||
headers the v3 rate limiter mirrors into the logging payload. The
|
||||
limiter already read these counters (from Redis when configured) on
|
||||
the request path, so no extra store lookup happens here. Descriptors
|
||||
without a configured limit emit no header, so their series is removed
|
||||
rather than left at the value from before the limit was dropped.
|
||||
"""
|
||||
descriptor_gauges: Final[
|
||||
tuple[tuple[Literal["api_key", "team"], DEFINED_PROMETHEUS_METRICS, Gauge, Gauge], ...]
|
||||
] = (
|
||||
(
|
||||
"api_key",
|
||||
"litellm_api_key_rate_limit_allowed_metric",
|
||||
self.litellm_api_key_rate_limit_allowed_metric,
|
||||
self.litellm_api_key_rate_limit_used_metric,
|
||||
),
|
||||
(
|
||||
"team",
|
||||
"litellm_team_rate_limit_allowed_metric",
|
||||
self.litellm_team_rate_limit_allowed_metric,
|
||||
self.litellm_team_rate_limit_used_metric,
|
||||
),
|
||||
)
|
||||
for descriptor_key, metric_name, allowed_gauge, used_gauge in descriptor_gauges:
|
||||
for rate_limit_type in ("requests", "tokens"):
|
||||
self._set_rate_limit_allowed_and_used_gauges(
|
||||
standard_logging_payload=standard_logging_payload,
|
||||
enum_values=enum_values,
|
||||
descriptor_key=descriptor_key,
|
||||
metric_name=metric_name,
|
||||
allowed_gauge=allowed_gauge,
|
||||
used_gauge=used_gauge,
|
||||
rate_limit_type=rate_limit_type,
|
||||
)
|
||||
|
||||
def _set_rate_limit_allowed_and_used_gauges(
|
||||
self,
|
||||
standard_logging_payload: StandardLoggingPayload,
|
||||
enum_values: UserAPIKeyLabelValues,
|
||||
descriptor_key: Literal["api_key", "team"],
|
||||
metric_name: DEFINED_PROMETHEUS_METRICS,
|
||||
allowed_gauge: Gauge,
|
||||
used_gauge: Gauge,
|
||||
rate_limit_type: Literal["requests", "tokens"],
|
||||
) -> None:
|
||||
limit: Final = self._get_int_from_v3_rate_limit_headers(
|
||||
standard_logging_payload=standard_logging_payload,
|
||||
header_name=f"x-ratelimit-{descriptor_key}-limit-{rate_limit_type}",
|
||||
)
|
||||
remaining: Final = self._get_int_from_v3_rate_limit_headers(
|
||||
standard_logging_payload=standard_logging_payload,
|
||||
header_name=f"x-ratelimit-{descriptor_key}-remaining-{rate_limit_type}",
|
||||
)
|
||||
labelled_values: Final = replace(enum_values, rate_limit_type=rate_limit_type)
|
||||
labelnames: Final = self.get_labels_for_metric(metric_name)
|
||||
labels: Final = prometheus_label_factory(
|
||||
supported_enum_labels=labelnames,
|
||||
enum_values=labelled_values,
|
||||
label_context=PrometheusLabelFactoryContext(labelled_values),
|
||||
)
|
||||
if limit is None or remaining is None:
|
||||
label_values: Final = tuple(labels.get(label) for label in labelnames)
|
||||
self._bounded_prometheus_series_tracker.remove_series(allowed_gauge, label_values)
|
||||
self._bounded_prometheus_series_tracker.remove_series(used_gauge, label_values)
|
||||
return
|
||||
allowed_gauge.labels(**labels).set(limit)
|
||||
used_gauge.labels(**labels).set(limit - remaining)
|
||||
|
||||
def _set_virtual_key_rate_limit_metrics(
|
||||
self,
|
||||
user_api_key: str | None,
|
||||
|
|
@ -2407,7 +2565,7 @@ class PrometheusLogger(CustomLogger):
|
|||
team_alias=user_api_key_dict.team_alias,
|
||||
org_id=user_api_key_dict.org_id,
|
||||
org_alias=user_api_key_dict.organization_alias,
|
||||
requested_model=request_data.get("model", ""),
|
||||
requested_model=_bounded_requested_model_label(request_data.get("model", "")),
|
||||
status_code=str(status_code),
|
||||
exception_status=str(status_code),
|
||||
exception_class=self._get_exception_class_name(original_exception),
|
||||
|
|
@ -2627,7 +2785,9 @@ class PrometheusLogger(CustomLogger):
|
|||
label_model_id = ""
|
||||
label_api_base = ""
|
||||
label_api_provider = ""
|
||||
label_requested_model = litellm_model_name or model_group or ""
|
||||
label_requested_model = (
|
||||
_bounded_requested_model_label(litellm_model_name or model_group, router_originated=True) or ""
|
||||
)
|
||||
|
||||
enum_values: Final = UserAPIKeyLabelValues(
|
||||
litellm_model_name=label_litellm_model_name,
|
||||
|
|
@ -3186,7 +3346,7 @@ class PrometheusLogger(CustomLogger):
|
|||
_tags: Final = cast(list[str], kwargs.get("tags") or [])
|
||||
|
||||
enum_values: Final = UserAPIKeyLabelValues(
|
||||
requested_model=original_model_group,
|
||||
requested_model=_bounded_requested_model_label(original_model_group, router_originated=True),
|
||||
fallback_model=_new_model,
|
||||
hashed_api_key=standard_metadata["user_api_key_hash"],
|
||||
api_key_alias=standard_metadata["user_api_key_alias"],
|
||||
|
|
@ -3227,7 +3387,7 @@ class PrometheusLogger(CustomLogger):
|
|||
)
|
||||
|
||||
enum_values: Final = UserAPIKeyLabelValues(
|
||||
requested_model=original_model_group,
|
||||
requested_model=_bounded_requested_model_label(original_model_group, router_originated=True),
|
||||
fallback_model=_new_model,
|
||||
hashed_api_key=standard_metadata["user_api_key_hash"],
|
||||
api_key_alias=standard_metadata["user_api_key_alias"],
|
||||
|
|
|
|||
|
|
@ -60,6 +60,10 @@ class BoundedPrometheusSeriesTracker:
|
|||
break
|
||||
del series[tracked_label_values]
|
||||
|
||||
def remove_series(self, metric: object, label_values: tuple[str | None, ...]) -> bool:
|
||||
"""Drop one child series, True when it is gone (removed or never existed)."""
|
||||
return self._remove_metric_child(metric, label_values)
|
||||
|
||||
def _should_run_ttl_cleanup(
|
||||
self,
|
||||
metric_name: str,
|
||||
|
|
|
|||
|
|
@ -8,11 +8,12 @@ import uuid
|
|||
from collections import Counter
|
||||
from collections.abc import Awaitable, Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Any, Final, Literal, Optional, TypedDict
|
||||
from typing import TYPE_CHECKING, Final, Literal, Optional, Protocol, TypedDict, overload
|
||||
|
||||
import httpx
|
||||
from typing_extensions import Never, ReadOnly
|
||||
from typing_extensions import Never, ReadOnly, Required
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.integrations.custom_batch_logger import CustomBatchLogger
|
||||
|
|
@ -30,6 +31,7 @@ from litellm.llms.custom_httpx.http_handler import (
|
|||
httpxSpecialProvider,
|
||||
)
|
||||
from litellm.types.guardrails import GuardrailEventHooks
|
||||
from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolCallChunk
|
||||
from litellm.types.utils import (
|
||||
ChatCompletionMessageToolCall,
|
||||
Function,
|
||||
|
|
@ -52,17 +54,102 @@ _DROP_WARNING_INTERVAL_SECONDS: Final = 60.0
|
|||
_EMPTY_MAPPING: Final[Mapping[str, Never]] = MappingProxyType({})
|
||||
|
||||
|
||||
class _ServiceToolCall(TypedDict):
|
||||
id: ReadOnly[str]
|
||||
class _ModerationToolCall(TypedDict, total=False):
|
||||
id: ReadOnly[Required[str]]
|
||||
|
||||
|
||||
class _ServiceMessage(TypedDict, total=False):
|
||||
class _ModerationMessage(TypedDict, total=False):
|
||||
content: ReadOnly[str | None]
|
||||
tool_calls: ReadOnly[Sequence[_ModerationToolCall] | None]
|
||||
|
||||
|
||||
class _ModerationChoice(TypedDict, total=False):
|
||||
message: ReadOnly[_ModerationMessage | None]
|
||||
|
||||
|
||||
class _ModerationResponse(TypedDict, total=False):
|
||||
choices: ReadOnly[Sequence[_ModerationChoice]]
|
||||
|
||||
|
||||
class _LogEventKwargs(TypedDict, total=False):
|
||||
standard_logging_object: ReadOnly[Required[StandardLoggingPayload]]
|
||||
litellm_call_id: ReadOnly[str]
|
||||
|
||||
|
||||
class _HasCallId(Protocol):
|
||||
def get(self, key: Literal["litellm_call_id"], /) -> str | None: ...
|
||||
|
||||
|
||||
class _HasModelAttr(Protocol):
|
||||
model: str | None
|
||||
|
||||
|
||||
class _ResponseSource(Protocol):
|
||||
def get(self, key: Literal["response"], /) -> "_HasModelAttr | None": ...
|
||||
|
||||
|
||||
class _ModelSource(Protocol):
|
||||
def get(self, key: Literal["model"], default: str, /) -> str: ...
|
||||
|
||||
|
||||
class _FallbackSource(Protocol):
|
||||
@overload
|
||||
def get(self, key: Literal["start_time"], /) -> datetime | None: ...
|
||||
@overload
|
||||
def get(self, key: str, /) -> object | None: ...
|
||||
|
||||
|
||||
class _RequestContextSource(Protocol):
|
||||
@overload
|
||||
def get(self, key: Literal["optional_params"], /) -> Mapping[str, object] | None: ...
|
||||
@overload
|
||||
def get(self, key: str, /) -> object | None: ...
|
||||
def __contains__(self, key: object, /) -> bool: ...
|
||||
def __getitem__(self, key: str, /) -> object: ...
|
||||
|
||||
|
||||
class _ToolCallLike(Protocol):
|
||||
id: str | None
|
||||
type: str | None
|
||||
function: Function
|
||||
|
||||
|
||||
class _ModerationSourceToolCall(TypedDict, total=False):
|
||||
function: ReadOnly[Mapping[str, object] | None]
|
||||
|
||||
|
||||
class _ModerationSourceMessage(TypedDict, total=False):
|
||||
role: ReadOnly[str]
|
||||
function_call: ReadOnly[Mapping[str, object] | None]
|
||||
tool_calls: ReadOnly[Sequence[_ModerationSourceToolCall | None] | None]
|
||||
|
||||
|
||||
class _FlattenedModerationMessage(TypedDict):
|
||||
role: ReadOnly[str | None]
|
||||
content: ReadOnly[str]
|
||||
tool_calls: ReadOnly[Sequence[_ServiceToolCall]]
|
||||
|
||||
|
||||
class _ServiceChoice(TypedDict, total=False):
|
||||
message: ReadOnly[_ServiceMessage]
|
||||
class _CorrelatablePayload(TypedDict):
|
||||
id: str # writable-ok: _apply_correlation_id overwrites the provider id on a deep-copied payload
|
||||
|
||||
|
||||
class _SystemPromptCarrier(TypedDict, total=False):
|
||||
messages: object # writable-ok: _prepend_system_prompt rebinds messages on the copied payload by design
|
||||
|
||||
|
||||
class _BlockFailurePayload(TypedDict, total=False):
|
||||
id: object # writable-ok: correlation id is pinned after copying the base payload
|
||||
model: ReadOnly[object]
|
||||
model_group: ReadOnly[object]
|
||||
model_id: ReadOnly[str]
|
||||
model_parameters: ReadOnly[object]
|
||||
startTime: ReadOnly[float | None]
|
||||
endTime: ReadOnly[float | None]
|
||||
completionStartTime: ReadOnly[float | None]
|
||||
messages: object # writable-ok: passed to _prepend_system_prompt, which rebinds messages
|
||||
metadata: ReadOnly[StandardLoggingUserAPIKeyMetadata]
|
||||
response: str # writable-ok: block failure text replaces the copied response
|
||||
status: ReadOnly[str]
|
||||
|
||||
|
||||
class _MalformedToolBlockingResponseError(Exception):
|
||||
|
|
@ -385,7 +472,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
|
|||
@staticmethod
|
||||
def _stash_block_context(
|
||||
logging_obj: Optional["LiteLLMLoggingObj"],
|
||||
request_data: dict,
|
||||
request_data: dict[str, object],
|
||||
) -> None:
|
||||
"""Stash signals so the deferred success-event skips this request and
|
||||
``async_post_call_failure_hook`` can build the failure payload.
|
||||
|
|
@ -414,12 +501,16 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
|
|||
request_data["_rubrik_logging_obj"] = logging_obj
|
||||
|
||||
@staticmethod
|
||||
def _normalize_tool_calls(tool_calls: Sequence[object]) -> tuple[ChatCompletionMessageToolCall, ...]:
|
||||
def _normalize_tool_calls(
|
||||
tool_calls: Sequence[ChatCompletionToolCallChunk | ChatCompletionMessageToolCall | _ToolCallLike],
|
||||
) -> tuple[ChatCompletionMessageToolCall, ...]:
|
||||
"""Convert tool_calls from inputs to ChatCompletionMessageToolCall objects."""
|
||||
return tuple(RubrikLogger._normalize_tool_call(tc) for tc in tool_calls)
|
||||
|
||||
@staticmethod
|
||||
def _normalize_tool_call(tc: Any) -> ChatCompletionMessageToolCall:
|
||||
def _normalize_tool_call(
|
||||
tc: ChatCompletionToolCallChunk | ChatCompletionMessageToolCall | _ToolCallLike,
|
||||
) -> ChatCompletionMessageToolCall:
|
||||
if isinstance(tc, ChatCompletionMessageToolCall):
|
||||
return tc
|
||||
if isinstance(tc, dict):
|
||||
|
|
@ -460,12 +551,15 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
|
|||
``content`` is sent so the webhook can moderate the response text;
|
||||
``None`` when the assistant produced no text (tool-call-only response).
|
||||
"""
|
||||
message: Final[dict[str, object]] = {
|
||||
message: Final[Mapping[str, object]] = {
|
||||
"role": "assistant",
|
||||
"content": content or None,
|
||||
**(
|
||||
{"tool_calls": tuple(tc.model_dump(exclude_none=True) for tc in tool_calls)}
|
||||
if tool_calls
|
||||
else _EMPTY_MAPPING
|
||||
),
|
||||
}
|
||||
if tool_calls:
|
||||
message["tool_calls"] = tuple(tc.model_dump(exclude_none=True) for tc in tool_calls)
|
||||
return {
|
||||
"id": request_id or f"chatcmpl-{uuid.uuid4()}",
|
||||
"object": "chat.completion",
|
||||
|
|
@ -481,7 +575,9 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
|
|||
}
|
||||
|
||||
@staticmethod
|
||||
def _flatten_messages_for_moderation(messages: Sequence[object] | None) -> tuple[Mapping[str, Any], ...]:
|
||||
def _flatten_messages_for_moderation(
|
||||
messages: Sequence[AllMessageValues | None] | None,
|
||||
) -> tuple[_FlattenedModerationMessage, ...]:
|
||||
"""Collapse each message's content to a plain string for the webhook.
|
||||
|
||||
litellm normalizes Anthropic ``/v1/messages`` requests to OpenAI shape,
|
||||
|
|
@ -502,7 +598,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
|
|||
)
|
||||
|
||||
@staticmethod
|
||||
def _moderation_text_parts(message: Mapping[str, Any]) -> tuple[str, ...]:
|
||||
def _moderation_text_parts(message: _ModerationSourceMessage) -> tuple[str, ...]:
|
||||
"""Every attacker-controlled text segment of a message: its content plus
|
||||
the arguments of any tool call or deprecated function call."""
|
||||
fc: Final = message.get("function_call")
|
||||
|
|
@ -530,16 +626,8 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
|
|||
``/v1/messages`` requests too. Optional fields are sent only when
|
||||
present so the payload stays clean.
|
||||
"""
|
||||
payload: Final[dict[str, object]] = {
|
||||
"model": inputs.get("model") or request_data.get("model") or "",
|
||||
"messages": RubrikLogger._flatten_messages_for_moderation(inputs.get("structured_messages")),
|
||||
}
|
||||
tools: Final = inputs.get("tools")
|
||||
if tools is not None:
|
||||
payload["tools"] = tools
|
||||
user: Final = request_data.get("user")
|
||||
if user:
|
||||
payload["user"] = user
|
||||
# Fall back to litellm_call_id, the stable cross-provider join key the
|
||||
# response/tool path uses (see _correlation_id). LiteLLM does not
|
||||
# populate request_data["correlation_key"]; it carries litellm_call_id.
|
||||
|
|
@ -547,14 +635,18 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
|
|||
# when correlation_key is empty, so without this the block fires but no
|
||||
# log is ever written. An explicit correlation_key still wins.
|
||||
correlation_key: Final = request_data.get("correlation_key") or request_data.get("litellm_call_id")
|
||||
if correlation_key:
|
||||
payload["correlation_key"] = correlation_key
|
||||
return payload
|
||||
return {
|
||||
"model": inputs.get("model") or request_data.get("model") or "",
|
||||
"messages": RubrikLogger._flatten_messages_for_moderation(inputs.get("structured_messages")),
|
||||
**({"tools": tools} if tools is not None else _EMPTY_MAPPING),
|
||||
**({"user": user} if user else _EMPTY_MAPPING),
|
||||
**({"correlation_key": correlation_key} if correlation_key else _EMPTY_MAPPING),
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def _extract_request_data(
|
||||
call_details: Mapping[str, Any],
|
||||
request_data: Mapping[str, object] | None,
|
||||
call_details: _RequestContextSource,
|
||||
request_data: _RequestContextSource | None,
|
||||
) -> Mapping[str, object]:
|
||||
"""Extract original request data from model_call_details for the
|
||||
response moderation service envelope.
|
||||
|
|
@ -590,7 +682,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
|
|||
}
|
||||
|
||||
@staticmethod
|
||||
def _sanitize_proxy_server_request(proxy_server_request: object) -> object:
|
||||
def _sanitize_proxy_server_request(proxy_server_request: Mapping[str, object] | str | None) -> object:
|
||||
"""Allowlist only routing fields (``url``, ``method``) when forwarding
|
||||
``proxy_server_request`` to an external webhook, dropping inbound
|
||||
``headers`` (Authorization, Cookie, x-api-key, ...) and the raw
|
||||
|
|
@ -600,18 +692,19 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
|
|||
return {key: proxy_server_request[key] for key in ("url", "method") if key in proxy_server_request}
|
||||
|
||||
@staticmethod
|
||||
def _resolve_model(request_data: Mapping[str, object], call_details: Mapping[str, str]) -> str:
|
||||
def _resolve_model(request_data: _ResponseSource, call_details: _ModelSource) -> str:
|
||||
"""Get the model name for the ModifyResponseException."""
|
||||
response: Final = request_data.get("response")
|
||||
if response and hasattr(response, "model"):
|
||||
response_model: Final[str | None] = getattr(response, "model", None)
|
||||
return response_model or "unknown"
|
||||
return response.model or "unknown"
|
||||
return call_details.get("model", "unknown")
|
||||
|
||||
# -- Logging hooks ---------------------------------------------------------
|
||||
|
||||
@staticmethod
|
||||
def _correlation_id(call_details: Mapping[str, str], request_data: Mapping[str, str] | None = None) -> str | None:
|
||||
def _correlation_id(
|
||||
call_details: _HasCallId | _LogEventKwargs, request_data: _HasCallId | None = None
|
||||
) -> str | None:
|
||||
"""The id that joins a blocked request's two S3 logs by filename: the
|
||||
moderation (``_blocking``) log and the failure (response) log.
|
||||
|
||||
|
|
@ -625,7 +718,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
|
|||
return call_details.get("litellm_call_id") or (request_data or _EMPTY_MAPPING).get("litellm_call_id")
|
||||
|
||||
@classmethod
|
||||
def _apply_correlation_id(cls, payload: dict[str, object], source: Mapping[str, str]) -> None:
|
||||
def _apply_correlation_id(cls, payload: _CorrelatablePayload, source: _HasCallId | _LogEventKwargs) -> None:
|
||||
"""Pin ``payload["id"]`` to ``litellm_call_id`` in place so this log
|
||||
shares its S3 filename id with the moderation (``_blocking``) and
|
||||
failure logs for the same request -- for every provider.
|
||||
|
|
@ -645,7 +738,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
|
|||
payload["id"] = correlated
|
||||
|
||||
@staticmethod
|
||||
def _prepend_system_prompt(payload: dict[str, object], source: Mapping[str, object]) -> None:
|
||||
def _prepend_system_prompt(payload: _SystemPromptCarrier, source: Mapping[str, object]) -> None:
|
||||
"""Prepend ``source["system"]`` onto ``payload["messages"]``.
|
||||
|
||||
Builds a NEW messages list rather than mutating ``payload["messages"]``
|
||||
|
|
@ -673,9 +766,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
|
|||
exc_info=True,
|
||||
)
|
||||
|
||||
async def _prepare_log_payload(
|
||||
self, kwargs: Mapping[str, object], event_type: str
|
||||
) -> StandardLoggingPayload | None:
|
||||
async def _prepare_log_payload(self, kwargs: _LogEventKwargs, event_type: str) -> StandardLoggingPayload | None:
|
||||
"""Shared logic for success logging (sampled)."""
|
||||
if random.random() > self.sampling_rate:
|
||||
verbose_logger.debug("Skipping Rubrik %s logging (sampling_rate=%s)", event_type, self.sampling_rate)
|
||||
|
|
@ -684,12 +775,12 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
|
|||
# Deep-copy so mutations don't affect other callbacks sharing this object
|
||||
standard_logging_payload: Final[StandardLoggingPayload] = safe_deep_copy(kwargs["standard_logging_object"])
|
||||
|
||||
self._apply_correlation_id(standard_logging_payload, kwargs) # pyright: ignore[reportArgumentType] # StandardLoggingPayload is dict[str,Any] at runtime
|
||||
self._apply_correlation_id(standard_logging_payload, kwargs)
|
||||
self._prepend_system_prompt(standard_logging_payload, kwargs) # pyright: ignore[reportArgumentType] # StandardLoggingPayload is dict[str,Any] at runtime
|
||||
|
||||
return standard_logging_payload
|
||||
|
||||
async def _append_and_maybe_flush(self, payload) -> None:
|
||||
async def _append_and_maybe_flush(self, payload: Mapping[str, object]) -> None:
|
||||
self._ensure_periodic_flush_task()
|
||||
self.log_queue.append(payload)
|
||||
self._enforce_max_queue_size()
|
||||
|
|
@ -714,7 +805,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
|
|||
self._dropped_since_warning = 0
|
||||
self._last_drop_warning_time = now
|
||||
|
||||
async def _enqueue_log_event(self, kwargs: Mapping[str, object], event_type: str):
|
||||
async def _enqueue_log_event(self, kwargs: _LogEventKwargs, event_type: str):
|
||||
try:
|
||||
payload: Final = await self._prepare_log_payload(kwargs, event_type)
|
||||
if payload is None:
|
||||
|
|
@ -835,7 +926,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
|
|||
logging_obj: "LiteLLMLoggingObj",
|
||||
exception: "ModifyResponseException",
|
||||
user_api_key_dict: "UserAPIKeyAuth",
|
||||
) -> StandardLoggingPayload:
|
||||
) -> _BlockFailurePayload:
|
||||
"""Build a failure-style payload using the exception text as response.
|
||||
|
||||
Blocked-tool events are security-relevant and **bypass sampling**:
|
||||
|
|
@ -877,9 +968,9 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
|
|||
call_details: Final = logging_obj.model_call_details
|
||||
exception_text: Final = f"{type(exception).__name__}: {exception.message}"
|
||||
|
||||
base: Final = call_details.get("standard_logging_object")
|
||||
base: Final[StandardLoggingPayload | None] = call_details.get("standard_logging_object")
|
||||
if base is not None:
|
||||
payload: dict[str, object] = safe_deep_copy(base)
|
||||
payload: _BlockFailurePayload = self._copy_block_payload_base(base)
|
||||
else:
|
||||
verbose_logger.debug(
|
||||
"Rubrik: standard_logging_object not yet on model_call_details "
|
||||
|
|
@ -901,6 +992,10 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
|
|||
|
||||
return payload
|
||||
|
||||
@staticmethod
|
||||
def _copy_block_payload_base(base: StandardLoggingPayload) -> _BlockFailurePayload:
|
||||
return safe_deep_copy(base)
|
||||
|
||||
@staticmethod
|
||||
def _caller_metadata(user_api_key_dict: "UserAPIKeyAuth") -> StandardLoggingUserAPIKeyMetadata:
|
||||
"""Identify the caller whose request was blocked.
|
||||
|
|
@ -923,9 +1018,9 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
|
|||
@classmethod
|
||||
def _build_fallback_payload(
|
||||
cls,
|
||||
call_details: Mapping[str, Any],
|
||||
call_details: _FallbackSource,
|
||||
user_api_key_dict: "UserAPIKeyAuth",
|
||||
) -> dict[str, object]:
|
||||
) -> _BlockFailurePayload:
|
||||
# Convert datetime to a Unix float so json.dumps can serialize it.
|
||||
# httpx's json= parameter uses stdlib json.dumps with no custom encoder.
|
||||
_raw_start: Final = call_details.get("start_time")
|
||||
|
|
@ -959,7 +1054,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
|
|||
response: Final = await self.async_httpx_client.post(
|
||||
url=self.logging_endpoint,
|
||||
json=data,
|
||||
headers=self._headers,
|
||||
headers=dict(self._headers),
|
||||
)
|
||||
response.raise_for_status()
|
||||
except httpx.HTTPStatusError as e:
|
||||
|
|
@ -1013,7 +1108,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
|
|||
|
||||
# -- Webhook services ------------------------------------------------------
|
||||
|
||||
async def _post_json(self, endpoint: str, payload: Mapping[str, object], service_name: str) -> Mapping[str, Any]:
|
||||
async def _post_json(self, endpoint: str, payload: Mapping[str, object], service_name: str) -> _ModerationResponse:
|
||||
"""POST ``payload`` to a Rubrik webhook and return its dict response.
|
||||
|
||||
Raises:
|
||||
|
|
@ -1023,11 +1118,11 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
|
|||
verbose_logger.debug("Sending request to %s: %s", service_name, endpoint)
|
||||
http_response: Final = await self.moderation_client.post(
|
||||
endpoint,
|
||||
json=payload,
|
||||
headers=self._headers,
|
||||
json=dict(payload),
|
||||
headers=dict(self._headers),
|
||||
)
|
||||
http_response.raise_for_status()
|
||||
result: Final[object] = http_response.json()
|
||||
result: Final[_ModerationResponse | None] = http_response.json()
|
||||
if not isinstance(result, dict):
|
||||
raise TypeError(
|
||||
f"{service_name} returned non-dict JSON "
|
||||
|
|
@ -1040,7 +1135,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
|
|||
self,
|
||||
response_data: Mapping[str, object],
|
||||
request_data: Mapping[str, object],
|
||||
) -> Mapping[str, Any]:
|
||||
) -> _ModerationResponse:
|
||||
"""Post the ``{request, response}`` envelope to the after_completion
|
||||
webhook and return its (possibly rewritten) response.
|
||||
|
||||
|
|
@ -1056,7 +1151,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
|
|||
"Response moderation service",
|
||||
)
|
||||
|
||||
async def _post_to_prompt_moderation_endpoint(self, payload: Mapping[str, object]) -> Mapping[str, Any]:
|
||||
async def _post_to_prompt_moderation_endpoint(self, payload: Mapping[str, object]) -> _ModerationResponse:
|
||||
"""Post a bare OpenAI request to the before_prompt webhook.
|
||||
|
||||
Returns ``{}`` (passthrough) or a synthetic chat.completion (block).
|
||||
|
|
@ -1064,14 +1159,14 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
|
|||
return await self._post_json(self.prompt_moderation_endpoint, payload, "Prompt moderation service")
|
||||
|
||||
@staticmethod
|
||||
def _extract_prompt_refusal(service_response: Mapping[str, Any]) -> str | None:
|
||||
def _extract_prompt_refusal(service_response: _ModerationResponse) -> str | None:
|
||||
"""Return the refusal text when the prompt was blocked, else None.
|
||||
|
||||
The before_prompt webhook returns ``{}`` (passthrough) or a synthetic
|
||||
chat.completion whose ``choices[0].message.content`` is the refusal
|
||||
explanation.
|
||||
"""
|
||||
choices: Final[Sequence[_ServiceChoice] | None] = service_response.get("choices")
|
||||
choices: Final = service_response.get("choices")
|
||||
if not choices:
|
||||
return None
|
||||
message: Final = choices[0].get("message") or _EMPTY_MAPPING
|
||||
|
|
@ -1080,7 +1175,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
|
|||
|
||||
@staticmethod
|
||||
def _extract_response_block(
|
||||
service_response: Mapping[str, Any],
|
||||
service_response: _ModerationResponse,
|
||||
all_tool_calls: Sequence[ChatCompletionMessageToolCall],
|
||||
sent_content: str,
|
||||
) -> BlockedResponseResult | None:
|
||||
|
|
@ -1103,7 +1198,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
|
|||
Expects service_response in OpenAI chat completion format:
|
||||
{"choices": [{"message": {"tool_calls": [...], "content": "..."}}]}
|
||||
"""
|
||||
choices: Final[Sequence[_ServiceChoice]] = service_response.get("choices") or ()
|
||||
choices: Final = service_response.get("choices") or ()
|
||||
if not choices:
|
||||
raise _MalformedToolBlockingResponseError("Response moderation service returned empty response")
|
||||
|
||||
|
|
|
|||
|
|
@ -1,11 +1,18 @@
|
|||
#### What this does ####
|
||||
# On success + failure, log events to Supabase
|
||||
|
||||
import hashlib
|
||||
from datetime import datetime
|
||||
from typing import Final, cast
|
||||
|
||||
import litellm
|
||||
from litellm._logging import print_verbose, verbose_logger
|
||||
from litellm.constants import (
|
||||
MAX_S3_OBJECT_DOWNLOAD_FILENAME_BYTES,
|
||||
MAX_S3_OBJECT_KEY_BYTES,
|
||||
S3_BOUNDED_OBJECT_KEY_HEAD_BYTES,
|
||||
S3_PREFIX_DIGEST_CHARS,
|
||||
)
|
||||
from litellm.types.utils import StandardLoggingPayload
|
||||
|
||||
|
||||
|
|
@ -133,9 +140,7 @@ class S3Logger:
|
|||
s3_file_name,
|
||||
)
|
||||
|
||||
s3_object_download_filename: Final = (
|
||||
"time-" + start_time.strftime("%Y-%m-%dT%H-%M-%S-%f") + "_" + payload["id"] + ".json"
|
||||
)
|
||||
s3_object_download_filename: Final = get_s3_object_download_filename(start_time, payload["id"])
|
||||
|
||||
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
|
||||
|
||||
|
|
@ -198,6 +203,47 @@ def resolve_sse_params(
|
|||
return algorithm, valid_key_id
|
||||
|
||||
|
||||
S3_MIN_BOUNDED_FILE_NAME_BYTES: Final = 64
|
||||
|
||||
|
||||
def _truncate_to_utf8_bytes(value: str, max_bytes: int) -> str:
|
||||
"""Trim `value` so its UTF-8 encoding fits `max_bytes`, never splitting a character."""
|
||||
if max_bytes <= 0:
|
||||
return ""
|
||||
encoded: Final = value.encode("utf-8")
|
||||
if len(encoded) <= max_bytes:
|
||||
return value
|
||||
return encoded[:max_bytes].decode("utf-8", errors="ignore")
|
||||
|
||||
|
||||
def get_s3_object_download_filename(start_time: datetime, response_id: str) -> str:
|
||||
"""Content-Disposition filename for the uploaded object, bounded to the metadata header cap."""
|
||||
sanitized_response_id: Final = response_id.replace("/", "_").replace('"', "_")
|
||||
file_name: Final = f"time-{start_time.strftime('%Y-%m-%dT%H-%M-%S-%f')}_{response_id}"
|
||||
sanitized_file_name: Final = f"time-{start_time.strftime('%Y-%m-%dT%H-%M-%S-%f')}_{sanitized_response_id}"
|
||||
budget: Final = MAX_S3_OBJECT_DOWNLOAD_FILENAME_BYTES - len(b".json")
|
||||
if len(sanitized_file_name.encode("utf-8")) <= budget:
|
||||
return sanitized_file_name + ".json"
|
||||
return _bounded_s3_file_name(file_name, sanitized_file_name, budget) + ".json"
|
||||
|
||||
|
||||
def _bounded_s3_file_name(s3_file_name: str, sanitized_s3_file_name: str, max_bytes: int) -> str:
|
||||
"""As much of the file name as `max_bytes` allows, then the sha256 of the whole name."""
|
||||
digest: Final = hashlib.sha256(s3_file_name.encode("utf-8")).hexdigest()
|
||||
head_budget: Final = min(S3_BOUNDED_OBJECT_KEY_HEAD_BYTES, max_bytes - len(digest) - 1)
|
||||
head: Final = _truncate_to_utf8_bytes(sanitized_s3_file_name, head_budget)
|
||||
return f"{head}_{digest}" if head else digest
|
||||
|
||||
|
||||
def _bounded_s3_prefix(configured_prefix: str, max_bytes: int) -> str:
|
||||
"""As much of the configured prefix as fits, then a digest segment naming the full prefix."""
|
||||
digest_segment: Final = hashlib.sha256(configured_prefix.encode("utf-8")).hexdigest()[:S3_PREFIX_DIGEST_CHARS] + "/"
|
||||
if max_bytes < len(digest_segment):
|
||||
return ""
|
||||
head: Final = _truncate_to_utf8_bytes(configured_prefix, max_bytes - len(digest_segment) - 1).rstrip("/")
|
||||
return f"{head}/{digest_segment}" if head else digest_segment
|
||||
|
||||
|
||||
def get_s3_object_key(
|
||||
s3_path: str,
|
||||
prefix: str,
|
||||
|
|
@ -205,12 +251,23 @@ def get_s3_object_key(
|
|||
s3_file_name: str,
|
||||
) -> str:
|
||||
sanitized_s3_file_name: Final = s3_file_name.replace("/", "_")
|
||||
s3_object_key = (
|
||||
(s3_path.rstrip("/") + "/" if s3_path else "")
|
||||
+ prefix
|
||||
+ start_time.strftime("%Y-%m-%d")
|
||||
+ "/"
|
||||
+ sanitized_s3_file_name
|
||||
) # we need the s3 key to include the time, so we log cache hits too
|
||||
s3_object_key += ".json"
|
||||
return s3_object_key
|
||||
configured_prefix: Final = (s3_path.rstrip("/") + "/" if s3_path else "") + prefix
|
||||
date_segment: Final = start_time.strftime("%Y-%m-%d") + "/"
|
||||
# we need the s3 key to include the time, so we log cache hits too
|
||||
s3_object_key: Final = configured_prefix + date_segment + sanitized_s3_file_name + ".json"
|
||||
if len(s3_object_key.encode("utf-8")) <= MAX_S3_OBJECT_KEY_BYTES:
|
||||
return s3_object_key
|
||||
|
||||
# shorten the response id first and only trim the configured prefix if that is what does not
|
||||
# fit, so prefix scoped IAM policies and lifecycle rules keep matching
|
||||
budget: Final = MAX_S3_OBJECT_KEY_BYTES - len(date_segment.encode("utf-8")) - len(b".json")
|
||||
prefix_bytes: Final = len(configured_prefix.encode("utf-8"))
|
||||
if prefix_bytes + S3_MIN_BOUNDED_FILE_NAME_BYTES <= budget:
|
||||
bounded_file_name: Final = _bounded_s3_file_name(s3_file_name, sanitized_s3_file_name, budget - prefix_bytes)
|
||||
return configured_prefix + date_segment + bounded_file_name + ".json"
|
||||
|
||||
shortest_file_name: Final = _bounded_s3_file_name(
|
||||
s3_file_name, sanitized_s3_file_name, S3_MIN_BOUNDED_FILE_NAME_BYTES
|
||||
)
|
||||
bounded_prefix: Final = _bounded_s3_prefix(configured_prefix, budget - len(shortest_file_name.encode("utf-8")))
|
||||
return bounded_prefix + date_segment + shortest_file_name + ".json"
|
||||
|
|
|
|||
|
|
@ -16,7 +16,11 @@ from urllib.parse import quote
|
|||
import litellm
|
||||
from litellm._logging import print_verbose, verbose_logger
|
||||
from litellm.constants import DEFAULT_S3_BATCH_SIZE, DEFAULT_S3_FLUSH_INTERVAL_SECONDS
|
||||
from litellm.integrations.s3 import get_s3_object_key, resolve_sse_params
|
||||
from litellm.integrations.s3 import (
|
||||
get_s3_object_download_filename,
|
||||
get_s3_object_key,
|
||||
resolve_sse_params,
|
||||
)
|
||||
from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix
|
||||
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
|
||||
from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker
|
||||
|
|
@ -259,11 +263,11 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM):
|
|||
now: Final = datetime.now(timezone.utc)
|
||||
audit_log_id: Final = audit_log.get("id", "unknown")
|
||||
|
||||
s3_path = cast(str | None, self.s3_path) or ""
|
||||
s3_path = s3_path.rstrip("/") + "/" if s3_path else ""
|
||||
|
||||
s3_object_key: Final = (
|
||||
f"{s3_path}audit_logs/{now.strftime('%Y-%m-%d')}/{now.strftime('%H-%M-%S')}_{audit_log_id}.json"
|
||||
s3_object_key: Final = get_s3_object_key(
|
||||
cast(str | None, self.s3_path) or "",
|
||||
"audit_logs/",
|
||||
now,
|
||||
f"{now.strftime('%H-%M-%S')}_{audit_log_id}",
|
||||
)
|
||||
|
||||
element: Final = s3BatchLoggingElement(
|
||||
|
|
@ -463,9 +467,7 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM):
|
|||
)
|
||||
verbose_logger.debug("s3_object_key=%s", s3_object_key)
|
||||
|
||||
s3_object_download_filename: Final = (
|
||||
f"time-{start_time.strftime('%Y-%m-%dT%H-%M-%S-%f')}_{standard_logging_payload['id']}.json"
|
||||
)
|
||||
s3_object_download_filename: Final = get_s3_object_download_filename(start_time, standard_logging_payload["id"])
|
||||
|
||||
return s3BatchLoggingElement(
|
||||
payload=dict(standard_logging_payload),
|
||||
|
|
|
|||
|
|
@ -1,8 +1,11 @@
|
|||
"""Shadow Eval Logger: samples a shadowed key's successful LLM requests (chat completions,
|
||||
Anthropic Messages, and Responses API surfaces, each normalized to chat shape), duplicates
|
||||
each against the job's other arm in a detached task (the auto-router for a forward job, the
|
||||
fixed baseline model for a reverse one), blind-judges real vs shadow, and appends one
|
||||
``LiteLLM_ShadowEvalAttempt`` row (verdict or error) as the feature's only hot-path write.
|
||||
each through every shadow arm in one detached task (each candidate auto-router for a
|
||||
forward job, the fixed baseline model for a reverse one), blind-judges real vs each arm,
|
||||
and appends one ``LiteLLM_ShadowEvalAttempt`` row per arm (verdict or error) as the
|
||||
feature's only hot-path write. A multi-router job's arms therefore score the identical
|
||||
sampled requests against the identical real responses, which is what makes their win
|
||||
rates comparable head-to-head.
|
||||
Counts, status, and spend derive from those rows at read time, so nothing can disagree
|
||||
across pods or stop races; the hook reads active jobs through a short-TTL cache."""
|
||||
|
||||
|
|
@ -498,12 +501,16 @@ def _decision_classifier_cost(metadata: Mapping[str, object]) -> float:
|
|||
return float(raw) if isinstance(raw, (int, float)) else 0.0
|
||||
|
||||
|
||||
def _request_was_routed_by(request_metadata: Mapping[str, object], router_name: str) -> bool:
|
||||
"""Whether the router under evaluation served this request, which is what decides
|
||||
the direction it belongs to. A forward job skips its own router's traffic, since
|
||||
duplicating it would compare the router to itself: guaranteed ties, judge spend for
|
||||
zero information. A reverse job samples exactly that traffic and nothing else."""
|
||||
return _routing_decision(request_metadata).get("router_model_name") == router_name
|
||||
def _direction_admits(request_metadata: Mapping[str, object], job: "ActiveShadowEvalJob") -> bool:
|
||||
"""Whether this request belongs to the job's direction. A forward job skips traffic
|
||||
any of its candidate routers served: duplicating a router's own request compares it
|
||||
to itself (guaranteed ties), and judging a sibling against another candidate's live
|
||||
response would score candidates against each other instead of against the incumbent.
|
||||
A reverse job samples exactly its one router's traffic and nothing else."""
|
||||
routed_by: Final = _routing_decision(request_metadata).get("router_model_name")
|
||||
if job.direction == "reverse":
|
||||
return routed_by == job.router_name
|
||||
return routed_by not in job.arm_router_names
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
|
|
@ -546,6 +553,7 @@ class ActiveShadowEvalJob(BaseModel):
|
|||
|
||||
id: str
|
||||
router_name: str
|
||||
router_names: tuple[str, ...] = ()
|
||||
direction: ShadowEvalDirection = "forward"
|
||||
baseline_model: str | None = None
|
||||
shadow_percentage: float
|
||||
|
|
@ -567,12 +575,25 @@ class ActiveShadowEvalJob(BaseModel):
|
|||
raise ValueError("baseline_model is set for exactly the reverse jobs")
|
||||
return self
|
||||
|
||||
@model_validator(mode="after")
|
||||
def _reverse_evaluates_one_router(self) -> "ActiveShadowEvalJob":
|
||||
"""A reverse row naming several routers is unsamplable (there is no one traffic
|
||||
slice they share) and fails closed."""
|
||||
if self.direction == "reverse" and len(self.arm_router_names) > 1:
|
||||
raise ValueError("a reverse job evaluates exactly one router")
|
||||
return self
|
||||
|
||||
@property
|
||||
def shadow_target(self) -> str:
|
||||
"""The model the duplicated arm calls: the router itself for a forward job, the
|
||||
fixed baseline for a reverse one. Total because the validator above pins
|
||||
def arm_router_names(self) -> tuple[str, ...]:
|
||||
"""The job's full router set; rows from before router_names existed hold it in
|
||||
router_name alone. The one place that reading lives on the sampling side."""
|
||||
return self.router_names or (self.router_name,)
|
||||
|
||||
def arm_target(self, arm_router: str) -> str:
|
||||
"""The model one duplicated arm calls: the candidate router itself for a forward
|
||||
job, the fixed baseline for a reverse one. Total because the validator above pins
|
||||
baseline_model to reverse jobs and only those."""
|
||||
return self.baseline_model or self.router_name
|
||||
return self.baseline_model or arm_router
|
||||
|
||||
|
||||
def _as_active_job(record: object, attempts: int, spend: float) -> ActiveShadowEvalJob | None:
|
||||
|
|
@ -592,7 +613,12 @@ _JOBS_CACHE_KEY: Final = "shadow_eval:active_jobs"
|
|||
|
||||
|
||||
class ShadowEvalLogger(CustomLogger):
|
||||
"""Fires blind pairwise shadow evaluations for keys with an active shadow-eval job."""
|
||||
"""Fires blind pairwise shadow evaluations for targets with an active shadow-eval job.
|
||||
|
||||
A job targets a virtual key, a team, or a user; a request qualifies for a job when
|
||||
any of its resolved identities (key hash, team id, user id) matches the job's
|
||||
target, so team and user jobs cover JWT-authenticated traffic, which carries no
|
||||
key hash at all."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
|
|
@ -617,10 +643,10 @@ class ShadowEvalLogger(CustomLogger):
|
|||
# generation; the refill absorbs written rows and resets.
|
||||
self._job_starts: dict[str, int] = {} # mutable-ok: per-generation counter
|
||||
|
||||
async def _active_jobs(self) -> Mapping[str, tuple[ActiveShadowEvalJob, ...]]:
|
||||
"""Active jobs by api_key_id, cache-first. A key holds at most one job per
|
||||
direction, so the value is a collection. A DB fault returns empty without
|
||||
caching, so sampling pauses for that request and the next one retries."""
|
||||
async def _active_jobs(self) -> Mapping[tuple[str, str], tuple[ActiveShadowEvalJob, ...]]:
|
||||
"""Active jobs by (target_type, target_id), cache-first. A target holds at most
|
||||
one job per direction, so the value is a collection. A DB fault returns empty
|
||||
without caching, so sampling pauses for that request and the next one retries."""
|
||||
cached: Final = await self._jobs_cache.async_get_cache(_JOBS_CACHE_KEY)
|
||||
if cached is not None:
|
||||
return cached # pyright: ignore[reportReturnType] # cache stores exactly this mapping shape
|
||||
|
|
@ -652,10 +678,10 @@ class ShadowEvalLogger(CustomLogger):
|
|||
)
|
||||
for row in grouped or []
|
||||
}
|
||||
by_key: Final = tuple(
|
||||
by_target: Final = tuple(
|
||||
sorted(
|
||||
(
|
||||
(str(record.api_key_id), job)
|
||||
((str(record.target_type), str(record.target_id)), job)
|
||||
for record in records or []
|
||||
if (job := _as_active_job(record, *attempt_stats.get(str(record.id), (0, 0.0)))) is not None
|
||||
),
|
||||
|
|
@ -663,7 +689,7 @@ class ShadowEvalLogger(CustomLogger):
|
|||
)
|
||||
)
|
||||
jobs: Final = MappingProxyType(
|
||||
{key: tuple(job for _, job in group) for key, group in groupby(by_key, key=itemgetter(0))}
|
||||
{target: tuple(job for _, job in group) for target, group in groupby(by_target, key=itemgetter(0))}
|
||||
)
|
||||
await self._jobs_cache.async_set_cache(_JOBS_CACHE_KEY, jobs)
|
||||
self._job_starts = {} # rebind-ok: new generation, counts absorbed into the fill
|
||||
|
|
@ -691,7 +717,7 @@ class ShadowEvalLogger(CustomLogger):
|
|||
now >= job.ends_at
|
||||
or job.attempts + self._job_starts.get(job.id, 0) >= job.max_turns
|
||||
or (job.max_budget is not None and job.spend >= job.max_budget)
|
||||
or _request_was_routed_by(request_metadata, job.router_name) != (job.direction == "reverse")
|
||||
or not _direction_admits(request_metadata, job)
|
||||
):
|
||||
continue
|
||||
if not _sample_hits(request_id, job.id, job.shadow_percentage):
|
||||
|
|
@ -720,8 +746,18 @@ class ShadowEvalLogger(CustomLogger):
|
|||
if should_redact_message_logging(dict(kwargs)): # mutable-ok: predicate takes a plain dict
|
||||
return
|
||||
metadata: Final = payload.get("metadata") or _EMPTY_METADATA
|
||||
api_key_hash: Final = metadata.get("user_api_key_hash")
|
||||
if not api_key_hash:
|
||||
# Each identity the request resolved to is a candidate target; JWT-auth
|
||||
# requests carry no key hash but do carry a team and user.
|
||||
targets: Final = tuple(
|
||||
(target_type, str(value))
|
||||
for target_type, value in (
|
||||
("key", metadata.get("user_api_key_hash")),
|
||||
("team", metadata.get("user_api_key_team_id")),
|
||||
("user", metadata.get("user_api_key_user_id")),
|
||||
)
|
||||
if value
|
||||
)
|
||||
if not targets:
|
||||
return
|
||||
request_id: Final = payload.get("id") or ""
|
||||
if not request_id:
|
||||
|
|
@ -731,8 +767,11 @@ class ShadowEvalLogger(CustomLogger):
|
|||
return # only surfaces this table can normalize are comparable; unknown types fail closed
|
||||
if ops.wire_params and _request_mutating_guardrail_ran(request_metadata):
|
||||
return # the wire-body snapshot predates the rewrite; replaying it would resurrect stripped content
|
||||
active_jobs: Final = await self._active_jobs()
|
||||
eligible: Final = self._sampled_jobs(
|
||||
(await self._active_jobs()).get(str(api_key_hash), ()), request_metadata, request_id
|
||||
tuple(job for target in targets for job in active_jobs.get(target, ())),
|
||||
request_metadata,
|
||||
request_id,
|
||||
)
|
||||
if not eligible:
|
||||
return
|
||||
|
|
@ -755,7 +794,10 @@ class ShadowEvalLogger(CustomLogger):
|
|||
if self._inflight_shadow_tasks >= _MAX_CONCURRENT_SHADOW_TASKS:
|
||||
self._record_funnel(job.id, "shed")
|
||||
continue
|
||||
self._job_starts[job.id] = self._job_starts.get(job.id, 0) + 1
|
||||
# One start writes one attempt row per arm, and max_turns is a row
|
||||
# ceiling, so admission must pre-count every arm or a multi-router
|
||||
# job overshoots the valve N-fold within a cache generation.
|
||||
self._job_starts[job.id] = self._job_starts.get(job.id, 0) + len(job.arm_router_names)
|
||||
self._inflight_shadow_tasks += 1
|
||||
asyncio.create_task(
|
||||
self._run_shadow_eval(
|
||||
|
|
@ -794,32 +836,74 @@ class ShadowEvalLogger(CustomLogger):
|
|||
shadow_params: Mapping[str, object],
|
||||
parent_metadata: Mapping[str, object],
|
||||
) -> None:
|
||||
"""Budget gate -> shadow call -> blind judge -> one attempt row, and every exit
|
||||
in exactly one coverage bucket: the gates that decline to spend on an admitted
|
||||
sample (no DB to record into, an over-budget key, an unverifiable or exhausted
|
||||
eval budget) count it withheld, so eligible traffic still reconciles as
|
||||
not_sampled + unjudgeable + shed + withheld + attempt rows. The prisma gate sits
|
||||
above the dispatch so no provider spend happens without a place to record the
|
||||
outcome, and the budget read lives here rather than in the success hook."""
|
||||
"""Budget gates once per sampled request, then every router arm in turn: shadow
|
||||
call -> blind judge -> one attempt row stamped with the arm. The gates that
|
||||
decline to spend on an admitted sample (no DB to record into, an over-budget key,
|
||||
an unverifiable or exhausted eval budget) count the REQUEST withheld before any
|
||||
arm runs, so funnel counters stay per-request and a leg's eligible traffic still
|
||||
reconciles as not_sampled + unjudgeable + shed + withheld + sampled requests,
|
||||
where each sampled request writes one attempt row per arm. A budget crossed
|
||||
mid-loop lets the remaining arms overshoot by one round, the same class of
|
||||
overshoot as the samples already in flight when the cap is crossed. The prisma
|
||||
gate sits above the dispatch so no provider spend happens without a place to
|
||||
record the outcome, and the budget read lives here rather than in the success
|
||||
hook."""
|
||||
prisma: Final = self._prisma_provider()
|
||||
if prisma is None:
|
||||
self._record_funnel(job.id, "withheld")
|
||||
return
|
||||
if await _key_or_team_is_over_budget(parent_metadata):
|
||||
self._record_funnel(job.id, "withheld")
|
||||
return
|
||||
if job.max_budget is not None:
|
||||
try:
|
||||
spend: Final = await self._read_job_spend(_job_spend_counter_key(job.id), job.spend, job.max_budget)
|
||||
except Exception as e: # noqa: BLE001 # unverifiable budget: skip the sample rather than spend on it
|
||||
verbose_logger.warning("shadow_eval: budget unverifiable for %s, sample skipped: %s", job.id, e)
|
||||
self._record_funnel(job.id, "withheld")
|
||||
return
|
||||
if spend >= job.max_budget:
|
||||
self._record_funnel(job.id, "withheld")
|
||||
return
|
||||
for arm_router in job.arm_router_names:
|
||||
await self._run_shadow_arm(
|
||||
prisma=prisma,
|
||||
job=job,
|
||||
arm_router=arm_router,
|
||||
request_id=request_id,
|
||||
messages=messages,
|
||||
real_text=real_text,
|
||||
real_model=real_model,
|
||||
real_cost=real_cost,
|
||||
real_classifier_cost=real_classifier_cost,
|
||||
real_cache_hit=real_cache_hit,
|
||||
control_tier=control_tier,
|
||||
shadow_params=shadow_params,
|
||||
parent_metadata=parent_metadata,
|
||||
)
|
||||
|
||||
async def _run_shadow_arm(
|
||||
self,
|
||||
prisma: "PrismaClient",
|
||||
job: ActiveShadowEvalJob,
|
||||
arm_router: str,
|
||||
request_id: str,
|
||||
messages: Sequence[Mapping[str, object]],
|
||||
real_text: str,
|
||||
real_model: str,
|
||||
real_cost: float,
|
||||
real_classifier_cost: float,
|
||||
real_cache_hit: bool,
|
||||
control_tier: str | None,
|
||||
shadow_params: Mapping[str, object],
|
||||
parent_metadata: Mapping[str, object],
|
||||
) -> None:
|
||||
"""One arm's pipeline: shadow call -> blind judge -> one attempt row, every exit
|
||||
recording this arm's outcome, so one arm's fault never silences a sibling arm."""
|
||||
try:
|
||||
if prisma is None:
|
||||
self._record_funnel(job.id, "withheld")
|
||||
return
|
||||
if await _key_or_team_is_over_budget(parent_metadata):
|
||||
self._record_funnel(job.id, "withheld")
|
||||
return
|
||||
if job.max_budget is not None:
|
||||
try:
|
||||
spend: Final = await self._read_job_spend(_job_spend_counter_key(job.id), job.spend, job.max_budget)
|
||||
except Exception as e: # noqa: BLE001 # unverifiable budget: skip the sample rather than spend on it
|
||||
verbose_logger.warning("shadow_eval: budget unverifiable for %s, sample skipped: %s", job.id, e)
|
||||
self._record_funnel(job.id, "withheld")
|
||||
return
|
||||
if spend >= job.max_budget:
|
||||
self._record_funnel(job.id, "withheld")
|
||||
return
|
||||
shadow: Final = await self._call_router_shadow(job.shadow_target, messages, shadow_params, parent_metadata)
|
||||
shadow: Final = await self._call_router_shadow(
|
||||
job.arm_target(arm_router), messages, shadow_params, parent_metadata
|
||||
)
|
||||
except Exception as e: # noqa: BLE001 # detached task: nothing billed yet, record and never raise
|
||||
verbose_logger.debug("shadow_eval: pipeline failed for %s: %s", request_id, e)
|
||||
await self._record_attempt(
|
||||
|
|
@ -827,6 +911,7 @@ class ShadowEvalLogger(CustomLogger):
|
|||
job,
|
||||
request_id,
|
||||
control_tier,
|
||||
router_name=arm_router,
|
||||
outcome="error",
|
||||
error=f"pipeline error: {e}",
|
||||
real_cost=real_cost,
|
||||
|
|
@ -840,6 +925,7 @@ class ShadowEvalLogger(CustomLogger):
|
|||
job,
|
||||
request_id,
|
||||
control_tier,
|
||||
router_name=arm_router,
|
||||
outcome="error",
|
||||
error=shadow.error,
|
||||
shadow_cost=shadow.cost,
|
||||
|
|
@ -864,6 +950,7 @@ class ShadowEvalLogger(CustomLogger):
|
|||
job,
|
||||
request_id,
|
||||
control_tier,
|
||||
router_name=arm_router,
|
||||
outcome="error",
|
||||
error=verdict.error,
|
||||
shadow=shadow,
|
||||
|
|
@ -880,6 +967,7 @@ class ShadowEvalLogger(CustomLogger):
|
|||
job,
|
||||
request_id,
|
||||
control_tier,
|
||||
router_name=arm_router,
|
||||
outcome=verdict.preference,
|
||||
shadow=shadow,
|
||||
real_model=real_model,
|
||||
|
|
@ -898,6 +986,7 @@ class ShadowEvalLogger(CustomLogger):
|
|||
job,
|
||||
request_id,
|
||||
control_tier,
|
||||
router_name=arm_router,
|
||||
outcome="error",
|
||||
error=f"pipeline error: {e}",
|
||||
shadow=shadow,
|
||||
|
|
@ -915,6 +1004,7 @@ class ShadowEvalLogger(CustomLogger):
|
|||
request_id: str,
|
||||
control_tier: str | None,
|
||||
*,
|
||||
router_name: str,
|
||||
outcome: str,
|
||||
real_cost: float,
|
||||
real_classifier_cost: float,
|
||||
|
|
@ -937,6 +1027,7 @@ class ShadowEvalLogger(CustomLogger):
|
|||
data={ # mutable-ok: Prisma payload
|
||||
"job_id": job.id,
|
||||
"request_id": request_id,
|
||||
"router_name": router_name,
|
||||
"outcome": outcome,
|
||||
"tier": control_tier if job.direction == "reverse" else (shadow.tier if shadow else None),
|
||||
"real_model": real_model or None,
|
||||
|
|
@ -1056,7 +1147,7 @@ class ShadowEvalLogger(CustomLogger):
|
|||
)
|
||||
|
||||
|
||||
_EMPTY_JOBS: Final[Mapping[str, tuple[ActiveShadowEvalJob, ...]]] = MappingProxyType({})
|
||||
_EMPTY_JOBS: Final[Mapping[tuple[str, str], tuple[ActiveShadowEvalJob, ...]]] = MappingProxyType({})
|
||||
|
||||
|
||||
def _default_prisma_provider() -> "PrismaClient | None":
|
||||
|
|
|
|||
|
|
@ -13,7 +13,7 @@ from litellm._logging import verbose_logger
|
|||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.types.llms.openai import AllMessageValues, ChatCompletionUserMessage
|
||||
from litellm.types.prompts.init_prompts import PromptSpec
|
||||
from litellm.types.utils import StandardCallbackDynamicParams
|
||||
from litellm.types.utils import CallTypes, StandardCallbackDynamicParams
|
||||
from litellm.types.vector_stores import (
|
||||
LiteLLM_ManagedVectorStore,
|
||||
VectorStoreResultContent,
|
||||
|
|
@ -226,7 +226,7 @@ class VectorStorePreCallHook(CustomLogger):
|
|||
self,
|
||||
request_data: dict,
|
||||
response: Any,
|
||||
call_type: Any | None,
|
||||
call_type: CallTypes | None,
|
||||
) -> Any | None:
|
||||
"""
|
||||
Add search results to the response after successful LLM call.
|
||||
|
|
@ -283,7 +283,7 @@ class VectorStorePreCallHook(CustomLogger):
|
|||
self,
|
||||
request_data: dict,
|
||||
response_chunk: Any,
|
||||
call_type: Any | None,
|
||||
call_type: CallTypes | None,
|
||||
) -> Any | None:
|
||||
"""
|
||||
Add search results to the final streaming chunk.
|
||||
|
|
|
|||
|
|
@ -10,7 +10,7 @@ import asyncio
|
|||
import math
|
||||
import uuid
|
||||
from collections.abc import AsyncIterator, Mapping, Sequence
|
||||
from typing import TYPE_CHECKING, Any, Final, TypedDict, cast
|
||||
from typing import TYPE_CHECKING, Any, Final, Literal, Never, TypedDict, TypeVar, cast
|
||||
|
||||
from typing_extensions import ReadOnly
|
||||
|
||||
|
|
@ -46,7 +46,13 @@ from litellm.types.integrations.websearch_interception import (
|
|||
AnthropicServerToolUseBlock,
|
||||
WebSearchInterceptionConfig,
|
||||
)
|
||||
from litellm.types.llms.openai import AllMessageValues
|
||||
from litellm.types.llms.anthropic import AnthropicThinkingParam
|
||||
from litellm.types.llms.openai import (
|
||||
AllMessageValues,
|
||||
ChatCompletionAudioParam,
|
||||
ChatCompletionPredictionContentParam,
|
||||
OpenAIWebSearchOptions,
|
||||
)
|
||||
from litellm.types.utils import (
|
||||
AgenticLoopParams,
|
||||
CallTypes,
|
||||
|
|
@ -56,6 +62,8 @@ from litellm.types.utils import (
|
|||
from litellm.utils import ProviderConfigManager
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from aiohttp import ClientSession
|
||||
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.llms.base_llm.anthropic_messages.transformation import (
|
||||
BaseAnthropicMessagesConfig,
|
||||
|
|
@ -77,6 +85,10 @@ WEBSEARCH_EMIT_NATIVE_BLOCKS_KEY: Final = "_websearch_interception_emit_native_b
|
|||
# ``web_search_tool_result`` blocks to inject into the final response.
|
||||
WEBSEARCH_NATIVE_BLOCKS_METADATA_KEY: Final = "websearch_native_blocks"
|
||||
|
||||
_RESPONSE_CONTENT_FIELD: Final = "content"
|
||||
|
||||
_ResponseT: Final = TypeVar("_ResponseT")
|
||||
|
||||
|
||||
class _PlanMetadataView(TypedDict):
|
||||
websearch_native_blocks: Sequence[Mapping[str, object]] | None
|
||||
|
|
@ -90,23 +102,98 @@ class _WebSearchSettingsView(TypedDict):
|
|||
websearch_interception_params: WebSearchInterceptionConfig
|
||||
|
||||
|
||||
class _SearchToolLitellmParams(TypedDict, total=False):
|
||||
search_provider: ReadOnly[str | None]
|
||||
|
||||
|
||||
class _SearchToolConfig(TypedDict, total=False):
|
||||
search_tool_name: str
|
||||
litellm_params: Mapping[str, object] | None
|
||||
litellm_params: ReadOnly[_SearchToolLitellmParams | None]
|
||||
|
||||
|
||||
class _DeploymentKwargsView(TypedDict):
|
||||
"""Typed reads of the untyped request kwargs seen by the deployment hook."""
|
||||
|
||||
class _LitellmParamsProviderView(TypedDict, total=False):
|
||||
custom_llm_provider: ReadOnly[str]
|
||||
litellm_params: ReadOnly[Mapping[str, object]]
|
||||
|
||||
|
||||
class _DeploymentCallKwargsView(TypedDict):
|
||||
custom_llm_provider: ReadOnly[str]
|
||||
litellm_params: ReadOnly[_LitellmParamsProviderView]
|
||||
model: ReadOnly[str]
|
||||
|
||||
|
||||
class _UserAuthView(TypedDict):
|
||||
"""Typed read of the optional team attached to the caller's auth object."""
|
||||
class _AcreateNamedParams(TypedDict, total=False):
|
||||
metadata: ReadOnly[Never]
|
||||
stop_sequences: ReadOnly[Never]
|
||||
stream: ReadOnly[bool | None]
|
||||
system: ReadOnly[str | None]
|
||||
temperature: ReadOnly[float | None]
|
||||
thinking: ReadOnly[Never]
|
||||
tool_choice: ReadOnly[Never]
|
||||
tools: ReadOnly[Never]
|
||||
top_k: ReadOnly[int | None]
|
||||
top_p: ReadOnly[float | None]
|
||||
container: ReadOnly[Never]
|
||||
|
||||
team_id: ReadOnly[str | None]
|
||||
|
||||
class _AsearchNamedParams(TypedDict, total=False):
|
||||
max_results: ReadOnly[int | None]
|
||||
search_domain_filter: ReadOnly[Never]
|
||||
max_tokens_per_page: ReadOnly[int | None]
|
||||
country: ReadOnly[str | None]
|
||||
api_key: ReadOnly[str | None]
|
||||
api_base: ReadOnly[str | None]
|
||||
timeout: ReadOnly[float | None]
|
||||
extra_headers: ReadOnly[Never]
|
||||
|
||||
|
||||
class _AcompletionNamedParams(TypedDict, total=False):
|
||||
functions: ReadOnly[Never]
|
||||
function_call: ReadOnly[str | None]
|
||||
timeout: ReadOnly[float | None]
|
||||
temperature: ReadOnly[float | None]
|
||||
top_p: ReadOnly[float | None]
|
||||
n: ReadOnly[int | None]
|
||||
stream: ReadOnly[bool | None]
|
||||
stream_options: ReadOnly[Never]
|
||||
stop: ReadOnly[Never]
|
||||
max_tokens: ReadOnly[int | None]
|
||||
max_completion_tokens: ReadOnly[int | None]
|
||||
modalities: ReadOnly[Never]
|
||||
prediction: ReadOnly[ChatCompletionPredictionContentParam | None]
|
||||
audio: ReadOnly[ChatCompletionAudioParam | None]
|
||||
presence_penalty: ReadOnly[float | None]
|
||||
frequency_penalty: ReadOnly[float | None]
|
||||
logit_bias: ReadOnly[Never]
|
||||
user: ReadOnly[str | None]
|
||||
response_format: ReadOnly[Never]
|
||||
seed: ReadOnly[int | None]
|
||||
tools: ReadOnly[Never]
|
||||
tool_choice: ReadOnly[Never]
|
||||
parallel_tool_calls: ReadOnly[bool | None]
|
||||
logprobs: ReadOnly[bool | None]
|
||||
top_logprobs: ReadOnly[int | None]
|
||||
deployment_id: ReadOnly[str | None]
|
||||
reasoning_effort: ReadOnly[Literal["none", "minimal", "low", "medium", "high", "xhigh", "default"] | None]
|
||||
verbosity: ReadOnly[Literal["low", "medium", "high"] | None]
|
||||
safety_identifier: ReadOnly[str | None]
|
||||
service_tier: ReadOnly[str | None]
|
||||
store: ReadOnly[bool | None]
|
||||
prompt_cache_key: ReadOnly[str | None]
|
||||
base_url: ReadOnly[str | None]
|
||||
api_version: ReadOnly[str | None]
|
||||
api_key: ReadOnly[str | None]
|
||||
model_list: ReadOnly[Never]
|
||||
extra_headers: ReadOnly[Never]
|
||||
thinking: ReadOnly[AnthropicThinkingParam | None]
|
||||
web_search_options: ReadOnly[OpenAIWebSearchOptions | None]
|
||||
include_server_side_tool_invocations: ReadOnly[bool | None]
|
||||
shared_session: ReadOnly["ClientSession | None"]
|
||||
enable_json_schema_validation: ReadOnly[bool | None]
|
||||
|
||||
|
||||
_NO_ACREATE_NAMED: Final[_AcreateNamedParams] = {}
|
||||
_NO_ASEARCH_NAMED: Final[_AsearchNamedParams] = {}
|
||||
_NO_ACOMPLETION_NAMED: Final[_AcompletionNamedParams] = {}
|
||||
|
||||
|
||||
class WebSearchInterceptionLogger(CustomLogger):
|
||||
|
|
@ -308,17 +395,17 @@ class WebSearchInterceptionLogger(CustomLogger):
|
|||
"""
|
||||
# Check if this is for an enabled provider
|
||||
# Try top-level kwargs first, then nested litellm_params, then derive from model name
|
||||
kwargs_view: Final[_DeploymentKwargsView] = {
|
||||
call_kwargs_view: Final[_DeploymentCallKwargsView] = {
|
||||
"custom_llm_provider": kwargs.get("custom_llm_provider", ""),
|
||||
"litellm_params": kwargs.get("litellm_params", {}),
|
||||
"model": kwargs.get("model", ""),
|
||||
}
|
||||
custom_llm_provider = kwargs_view["custom_llm_provider"] or kwargs_view["litellm_params"].get(
|
||||
custom_llm_provider = call_kwargs_view["custom_llm_provider"] or call_kwargs_view["litellm_params"].get(
|
||||
"custom_llm_provider", ""
|
||||
)
|
||||
if not custom_llm_provider:
|
||||
try:
|
||||
_, custom_llm_provider, _, _ = litellm.get_llm_provider(model=kwargs_view["model"])
|
||||
_, custom_llm_provider, _, _ = litellm.get_llm_provider(model=call_kwargs_view["model"])
|
||||
except Exception:
|
||||
custom_llm_provider = ""
|
||||
if custom_llm_provider not in self.enabled_providers:
|
||||
|
|
@ -948,17 +1035,17 @@ class WebSearchInterceptionLogger(CustomLogger):
|
|||
)
|
||||
|
||||
@staticmethod
|
||||
def _inject_native_blocks(response: Any, native_blocks: Sequence[Mapping[str, object]]) -> Any:
|
||||
def _inject_native_blocks(response: _ResponseT, native_blocks: Sequence[Mapping[str, object]]) -> _ResponseT:
|
||||
"""Prepend native blocks to response content, dict or object form."""
|
||||
if not native_blocks:
|
||||
return response
|
||||
if isinstance(response, dict):
|
||||
existing = response.get("content") or []
|
||||
response["content"] = list(native_blocks) + list(existing)
|
||||
existing = response.get(_RESPONSE_CONTENT_FIELD) or []
|
||||
response[_RESPONSE_CONTENT_FIELD] = list(native_blocks) + list(existing)
|
||||
return response
|
||||
existing = getattr(response, "content", None) or []
|
||||
existing = getattr(response, _RESPONSE_CONTENT_FIELD, None) or []
|
||||
try:
|
||||
response.content = list(native_blocks) + list(existing)
|
||||
setattr(response, _RESPONSE_CONTENT_FIELD, list(native_blocks) + list(existing))
|
||||
except (AttributeError, TypeError):
|
||||
# Object refused write — fall through and leave the response
|
||||
# untouched rather than crash the request.
|
||||
|
|
@ -1214,10 +1301,10 @@ class WebSearchInterceptionLogger(CustomLogger):
|
|||
messages: list[dict],
|
||||
tool_calls: list[dict],
|
||||
thinking_blocks: list[dict],
|
||||
anthropic_messages_optional_request_params: dict,
|
||||
anthropic_messages_optional_request_params: Mapping[str, object],
|
||||
logging_obj: "LiteLLMLoggingObj | None",
|
||||
stream: bool,
|
||||
kwargs: dict,
|
||||
kwargs: Mapping[str, object],
|
||||
) -> "AnthropicMessagesResponse | AsyncIterator[object]":
|
||||
"""Legacy path: execute search + build patch + run follow-up call."""
|
||||
request_patch, structured_results = await self._build_anthropic_request_patch(
|
||||
|
|
@ -1225,9 +1312,9 @@ class WebSearchInterceptionLogger(CustomLogger):
|
|||
messages=messages,
|
||||
tool_calls=tool_calls,
|
||||
thinking_blocks=thinking_blocks,
|
||||
anthropic_messages_optional_request_params=anthropic_messages_optional_request_params,
|
||||
anthropic_messages_optional_request_params=dict[str, object](anthropic_messages_optional_request_params),
|
||||
logging_obj=logging_obj,
|
||||
kwargs=kwargs,
|
||||
kwargs=dict[str, object](kwargs),
|
||||
)
|
||||
if request_patch.messages is None:
|
||||
raise ValueError("WebSearchInterception: missing follow-up messages")
|
||||
|
|
@ -1242,12 +1329,14 @@ class WebSearchInterceptionLogger(CustomLogger):
|
|||
if max_tokens is None:
|
||||
max_tokens = cast(int, kwargs.get("max_tokens", 1024))
|
||||
|
||||
patch_kwargs: Final = dict[str, object](request_patch.kwargs)
|
||||
response: AnthropicMessagesResponse | AsyncIterator[object] = await anthropic_messages.acreate(
|
||||
max_tokens=max_tokens,
|
||||
messages=request_patch.messages,
|
||||
model=request_patch.model or model,
|
||||
**_NO_ACREATE_NAMED,
|
||||
**optional_params,
|
||||
**request_patch.kwargs,
|
||||
**patch_kwargs,
|
||||
)
|
||||
|
||||
# Legacy path: the new path goes through the typed plan + core
|
||||
|
|
@ -1389,12 +1478,13 @@ class WebSearchInterceptionLogger(CustomLogger):
|
|||
|
||||
search_tool: Final = self._select_search_tool_from_router(llm_router=llm_router)
|
||||
search_provider: str | None = None
|
||||
search_litellm_params: dict[str, Any] = {}
|
||||
search_litellm_params: Mapping[str, object] = {}
|
||||
search_tool_name: Final = self._selected_search_tool_name(search_tool=search_tool)
|
||||
if search_tool is not None:
|
||||
await self._authorize_search_tool(search_tool=search_tool, kwargs=kwargs)
|
||||
search_litellm_params = dict(search_tool.get("litellm_params", {}) or {})
|
||||
search_provider = search_litellm_params.get("search_provider")
|
||||
tool_params: Final[_SearchToolLitellmParams] = search_tool.get("litellm_params", {}) or {}
|
||||
search_litellm_params = dict[str, object](tool_params)
|
||||
search_provider = tool_params.get("search_provider")
|
||||
|
||||
# Fallback to perplexity if no router or no search tools configured
|
||||
if not search_provider:
|
||||
|
|
@ -1422,12 +1512,15 @@ class WebSearchInterceptionLogger(CustomLogger):
|
|||
if key != "search_provider" and value is not None
|
||||
}
|
||||
result: Final = (
|
||||
await litellm.asearch(query=query, search_provider=search_provider, **search_kwargs)
|
||||
await litellm.asearch(
|
||||
query=query, search_provider=search_provider, **_NO_ASEARCH_NAMED, **search_kwargs
|
||||
)
|
||||
if search_metadata is None
|
||||
else await litellm.asearch(
|
||||
query=query,
|
||||
search_provider=search_provider,
|
||||
litellm_metadata=search_metadata,
|
||||
**_NO_ASEARCH_NAMED,
|
||||
**search_kwargs,
|
||||
)
|
||||
)
|
||||
|
|
@ -1467,8 +1560,7 @@ class WebSearchInterceptionLogger(CustomLogger):
|
|||
valid_token=user_api_key_auth,
|
||||
)
|
||||
|
||||
auth_view: Final[_UserAuthView] = {"team_id": getattr(user_api_key_auth, "team_id", None)}
|
||||
team_id: Final = auth_view["team_id"]
|
||||
team_id: Final[str | None] = getattr(user_api_key_auth, "team_id", None)
|
||||
if team_id:
|
||||
from litellm.proxy.proxy_server import (
|
||||
prisma_client,
|
||||
|
|
@ -1541,16 +1633,18 @@ class WebSearchInterceptionLogger(CustomLogger):
|
|||
def _select_search_tool_from_router(self, llm_router: object) -> "_SearchToolConfig | None":
|
||||
if llm_router is None or not hasattr(llm_router, "search_tools"):
|
||||
return None
|
||||
search_tools: Final = list(getattr(llm_router, "search_tools") or [])
|
||||
search_tools: Final = tuple(getattr(llm_router, "search_tools", None) or ())
|
||||
return self._select_search_tool_from_list(search_tools=search_tools, source="router")
|
||||
|
||||
def _select_search_tool_from_list(
|
||||
self,
|
||||
search_tools: list[_SearchToolConfig],
|
||||
search_tools: Sequence[_SearchToolConfig],
|
||||
source: str,
|
||||
) -> "_SearchToolConfig | None":
|
||||
if self.search_tool_name:
|
||||
matching_tools = [tool for tool in search_tools if tool.get("search_tool_name") == self.search_tool_name]
|
||||
matching_tools: Final = tuple(
|
||||
tool for tool in search_tools if tool.get("search_tool_name") == self.search_tool_name
|
||||
)
|
||||
if matching_tools:
|
||||
search_provider = (matching_tools[0].get("litellm_params", {}) or {}).get("search_provider")
|
||||
verbose_logger.debug(
|
||||
|
|
@ -1583,10 +1677,10 @@ class WebSearchInterceptionLogger(CustomLogger):
|
|||
model: str,
|
||||
messages: list[dict],
|
||||
tool_calls: list[dict],
|
||||
optional_params: dict,
|
||||
optional_params: Mapping[str, object],
|
||||
logging_obj: "LiteLLMLoggingObj | None",
|
||||
stream: bool,
|
||||
kwargs: dict,
|
||||
kwargs: Mapping[str, object],
|
||||
response_format: str = "openai",
|
||||
) -> "ModelResponse | CustomStreamWrapper":
|
||||
"""Legacy path: execute search + build patch + run follow-up call."""
|
||||
|
|
@ -1594,8 +1688,8 @@ class WebSearchInterceptionLogger(CustomLogger):
|
|||
model=model,
|
||||
messages=messages,
|
||||
tool_calls=tool_calls,
|
||||
optional_params=optional_params,
|
||||
kwargs=kwargs,
|
||||
optional_params=dict[str, object](optional_params),
|
||||
kwargs=dict[str, object](kwargs),
|
||||
response_format=response_format,
|
||||
)
|
||||
if request_patch.messages is None:
|
||||
|
|
@ -1603,11 +1697,13 @@ class WebSearchInterceptionLogger(CustomLogger):
|
|||
params: Final = dict(optional_params)
|
||||
params.update(request_patch.optional_params)
|
||||
params.pop("tool_choice", None)
|
||||
patch_kwargs: Final = dict[str, object](request_patch.kwargs)
|
||||
return await litellm.acompletion(
|
||||
model=request_patch.model or model,
|
||||
messages=request_patch.messages,
|
||||
**_NO_ACOMPLETION_NAMED,
|
||||
**params,
|
||||
**request_patch.kwargs,
|
||||
**patch_kwargs,
|
||||
)
|
||||
|
||||
async def _build_chat_completion_request_patch(
|
||||
|
|
|
|||
|
|
@ -7,7 +7,13 @@ import os
|
|||
from dataclasses import dataclass
|
||||
from typing import Final
|
||||
|
||||
from litellm.types.files import get_file_mime_type_from_extension
|
||||
from litellm.types.files import (
|
||||
AUDIO_FILE_TYPES,
|
||||
FILE_EXTENSIONS,
|
||||
FILE_MIME_TYPES,
|
||||
FileType,
|
||||
get_file_mime_type_from_extension,
|
||||
)
|
||||
from litellm.types.utils import FileTypes
|
||||
|
||||
|
||||
|
|
@ -323,3 +329,75 @@ def calculate_request_duration(file: FileTypes) -> float | None:
|
|||
except Exception:
|
||||
# Silently fail if duration extraction fails
|
||||
return None
|
||||
|
||||
|
||||
DEFAULT_SPEECH_MEDIA_TYPE: Final = "audio/mpeg"
|
||||
|
||||
|
||||
def _speech_media_type_for_response_format(response_format: str) -> str | None:
|
||||
file_type: Final = next(
|
||||
(candidate for candidate, extensions in FILE_EXTENSIONS.items() if response_format.lower() in extensions),
|
||||
None,
|
||||
)
|
||||
if file_type is None or file_type not in AUDIO_FILE_TYPES:
|
||||
return None
|
||||
return FILE_MIME_TYPES[file_type]
|
||||
|
||||
|
||||
def resolve_speech_media_type(upstream_content_type: str | None, response_format: str | None) -> str:
|
||||
upstream_media_type: Final = (upstream_content_type or "").split(";", 1)[0].strip().lower()
|
||||
if upstream_media_type.startswith("audio/"):
|
||||
return upstream_media_type
|
||||
requested_media_type: Final = (
|
||||
None if response_format is None else _speech_media_type_for_response_format(response_format)
|
||||
)
|
||||
return requested_media_type or DEFAULT_SPEECH_MEDIA_TYPE
|
||||
|
||||
|
||||
_OGG_OPUS_HEAD_WINDOW: Final = 64
|
||||
_ADTS_SYNC_AND_LAYER_MASK: Final = 0xF6
|
||||
_ADTS_SYNC_AND_LAYER: Final = 0xF0
|
||||
_ADTS_SAMPLE_RATE_INDEX_LIMIT: Final = 13
|
||||
_MPEG_SYNC_MASK: Final = 0xE0
|
||||
_MPEG_LAYER_MASK: Final = 0x06
|
||||
_MPEG_RESERVED_VERSION: Final = 0x01
|
||||
_MPEG_INVALID_BITRATE_INDEX: Final = 0x0F
|
||||
_MPEG_RESERVED_SAMPLE_RATE_INDEX: Final = 0x03
|
||||
|
||||
|
||||
def _adts_aac_frame_media_type(header: bytes) -> str | None:
|
||||
sample_rate_index: Final = (header[2] >> 2) & 0x0F
|
||||
return FILE_MIME_TYPES[FileType.AAC] if sample_rate_index < _ADTS_SAMPLE_RATE_INDEX_LIMIT else None
|
||||
|
||||
|
||||
def _mpeg_audio_frame_media_type(header: bytes) -> str | None:
|
||||
version: Final = (header[1] >> 3) & 0x03
|
||||
layer: Final = header[1] & _MPEG_LAYER_MASK
|
||||
bitrate_index: Final = header[2] >> 4
|
||||
sample_rate_index: Final = (header[2] >> 2) & 0x03
|
||||
if (
|
||||
(header[1] & _MPEG_SYNC_MASK) != _MPEG_SYNC_MASK
|
||||
or version == _MPEG_RESERVED_VERSION
|
||||
or layer == 0
|
||||
or bitrate_index == _MPEG_INVALID_BITRATE_INDEX
|
||||
or sample_rate_index == _MPEG_RESERVED_SAMPLE_RATE_INDEX
|
||||
):
|
||||
return None
|
||||
return FILE_MIME_TYPES[FileType.MP3]
|
||||
|
||||
|
||||
def speech_media_type_from_audio_bytes(audio: bytes) -> str | None:
|
||||
if audio[:4] == b"RIFF" and audio[8:12] == b"WAVE":
|
||||
return FILE_MIME_TYPES[FileType.WAV]
|
||||
if audio[:4] == b"fLaC":
|
||||
return FILE_MIME_TYPES[FileType.FLAC]
|
||||
if audio[:4] == b"OggS":
|
||||
is_opus: Final = b"OpusHead" in audio[:_OGG_OPUS_HEAD_WINDOW]
|
||||
return FILE_MIME_TYPES[FileType.OPUS if is_opus else FileType.OGG]
|
||||
if audio[:3] == b"ID3":
|
||||
return FILE_MIME_TYPES[FileType.MP3]
|
||||
if len(audio) < 3 or audio[0] != 0xFF:
|
||||
return None
|
||||
if (audio[1] & _ADTS_SYNC_AND_LAYER_MASK) == _ADTS_SYNC_AND_LAYER:
|
||||
return _adts_aac_frame_media_type(audio)
|
||||
return _mpeg_audio_frame_media_type(audio)
|
||||
|
|
|
|||
|
|
@ -84,6 +84,9 @@ OPTIONAL_KWARGS_KEYS: Final = (
|
|||
"vertex_ai_project",
|
||||
"vertex_ai_location",
|
||||
"vertex_ai_credentials",
|
||||
"gigachat_scope",
|
||||
"gigachat_auth_url",
|
||||
"gigachat_access_token",
|
||||
"tpm",
|
||||
"rpm",
|
||||
"itpm",
|
||||
|
|
|
|||
|
|
@ -127,7 +127,7 @@ def handle_anthropic_text_model_custom_llm_provider(
|
|||
return model, custom_llm_provider
|
||||
|
||||
|
||||
def declared_authenticating_provider(model: str, custom_llm_provider: str | None = None) -> str | None:
|
||||
def declared_authenticating_provider(model: str | None, custom_llm_provider: str | None = None) -> str | None:
|
||||
"""The authenticating provider this pair already names, or None.
|
||||
|
||||
get_llm_provider runs the OAuth device flow for github_copilot and chatgpt, because their
|
||||
|
|
@ -135,7 +135,7 @@ def declared_authenticating_provider(model: str, custom_llm_provider: str | None
|
|||
and for a declared pair the resolver's answer is the declaration itself, so metadata callers
|
||||
adopt the declaration instead of resolving.
|
||||
"""
|
||||
declared: Final = custom_llm_provider or model.split("/", 1)[0]
|
||||
declared: Final = custom_llm_provider or (model.split("/", 1)[0] if model and "/" in model else None)
|
||||
return declared if declared in PROVIDERS_THAT_AUTHENTICATE_ON_PROVIDER_INFO else None
|
||||
|
||||
|
||||
|
|
@ -369,6 +369,9 @@ def get_llm_provider(
|
|||
elif endpoint == "https://api.meta.ai/v1":
|
||||
custom_llm_provider = "meta"
|
||||
dynamic_api_key = get_secret_str("META_API_KEY")
|
||||
elif endpoint == "https://gigachat.devices.sberbank.ru/api/v1":
|
||||
custom_llm_provider = "gigachat"
|
||||
dynamic_api_key = get_secret_str("GIGACHAT_API_KEY")
|
||||
elif (json_provider := JSONProviderRegistry.get_by_base_url(endpoint)) is not None:
|
||||
custom_llm_provider = json_provider.slug
|
||||
dynamic_api_key = api_key if api_key is not None else get_secret_str(json_provider.api_key_env)
|
||||
|
|
@ -533,6 +536,14 @@ def get_llm_provider(
|
|||
)
|
||||
|
||||
|
||||
def _dashscope_family_chat_config(custom_llm_provider: str) -> "litellm.DashScopeChatConfig":
|
||||
if custom_llm_provider == "qwencloud":
|
||||
return litellm.QwenCloudChatConfig()
|
||||
if custom_llm_provider == "qwen_ai_platform":
|
||||
return litellm.QwenAIPlatformChatConfig()
|
||||
return litellm.DashScopeChatConfig()
|
||||
|
||||
|
||||
def _get_openai_compatible_provider_info(
|
||||
model: str,
|
||||
api_base: str | None,
|
||||
|
|
@ -782,11 +793,11 @@ def _get_openai_compatible_provider_info(
|
|||
api_base,
|
||||
dynamic_api_key,
|
||||
) = litellm.HerokuChatConfig()._get_openai_compatible_provider_info(api_base, api_key)
|
||||
elif custom_llm_provider == "dashscope":
|
||||
elif custom_llm_provider in ("dashscope", "qwencloud", "qwen_ai_platform"):
|
||||
(
|
||||
api_base,
|
||||
dynamic_api_key,
|
||||
) = litellm.DashScopeChatConfig()._get_openai_compatible_provider_info(api_base, api_key)
|
||||
) = _dashscope_family_chat_config(custom_llm_provider)._get_openai_compatible_provider_info(api_base, api_key)
|
||||
elif custom_llm_provider == "modelscope":
|
||||
(
|
||||
api_base,
|
||||
|
|
@ -867,6 +878,9 @@ def _get_openai_compatible_provider_info(
|
|||
# Manus is OpenAI compatible for responses API
|
||||
api_base = api_base or get_secret_str("MANUS_API_BASE") or "https://api.manus.im"
|
||||
dynamic_api_key = api_key or get_secret_str("MANUS_API_KEY")
|
||||
elif custom_llm_provider == "gigachat":
|
||||
api_base = api_base or get_secret_str("GIGACHAT_API_BASE") or "https://gigachat.devices.sberbank.ru/api/v1"
|
||||
dynamic_api_key = api_key or get_secret_str("GIGACHAT_API_KEY")
|
||||
|
||||
if api_base is not None and not isinstance(api_base, str):
|
||||
raise Exception(f"api base needs to be a string. api_base={api_base}")
|
||||
|
|
|
|||
|
|
@ -12,6 +12,7 @@ import asyncio
|
|||
import json
|
||||
import os
|
||||
import random
|
||||
import time
|
||||
from collections.abc import Awaitable, Callable
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime, timezone
|
||||
|
|
@ -154,18 +155,6 @@ class GetModelCostMap:
|
|||
|
||||
return True
|
||||
|
||||
@staticmethod
|
||||
def fetch_remote_model_cost_map(url: str, timeout: int = 5) -> dict:
|
||||
"""
|
||||
Fetch the model cost map from a remote URL.
|
||||
|
||||
Returns the parsed JSON dict. Raises on network/parse errors
|
||||
(caller is expected to handle).
|
||||
"""
|
||||
response: Final = httpx.get(url, timeout=timeout)
|
||||
response.raise_for_status()
|
||||
return response.json()
|
||||
|
||||
|
||||
RETRYABLE_FETCH_STATUS_CODES: Final = frozenset({429, 500, 502, 503, 504})
|
||||
MODEL_COST_MAP_FETCH_MAX_ATTEMPTS: Final = 3
|
||||
|
|
@ -212,6 +201,13 @@ class _AsyncGetClient(Protocol):
|
|||
def get(self, url: str, *, timeout: float | None = None) -> Awaitable[httpx.Response]: ...
|
||||
|
||||
|
||||
class _SyncGetClient(Protocol):
|
||||
def get(self, url: str, *, timeout: float | None = None) -> httpx.Response: ...
|
||||
|
||||
|
||||
_FetchAttemptOutcome = ModelCostMapReloaded | ModelCostMapReloadUnavailable | _FetchAttemptRetryable
|
||||
|
||||
|
||||
def _default_reload_client() -> _AsyncGetClient:
|
||||
from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
|
||||
from litellm.types.llms.custom_http import httpxSpecialProvider
|
||||
|
|
@ -219,13 +215,30 @@ def _default_reload_client() -> _AsyncGetClient:
|
|||
return get_async_httpx_client(llm_provider=httpxSpecialProvider.ModelCostMap)
|
||||
|
||||
|
||||
async def _attempt_fetch(
|
||||
client: _AsyncGetClient, url: str, timeout: int
|
||||
) -> ModelCostMapReloaded | ModelCostMapReloadUnavailable | _FetchAttemptRetryable:
|
||||
def _classify_fetch_error(error: httpx.HTTPError | httpx.InvalidURL, url: str) -> _FetchAttemptOutcome:
|
||||
reason: Final = f"{type(error).__name__} fetching {url}: {error}"
|
||||
if isinstance(error, (httpx.InvalidURL, httpx.UnsupportedProtocol)):
|
||||
return ModelCostMapReloadUnavailable(reason=reason)
|
||||
return _FetchAttemptRetryable(reason=reason, retry_after_seconds=None)
|
||||
|
||||
|
||||
async def _attempt_fetch(client: _AsyncGetClient, url: str, timeout: int) -> _FetchAttemptOutcome:
|
||||
try:
|
||||
response: Final = await client.get(url, timeout=timeout)
|
||||
except httpx.HTTPError as e:
|
||||
return _FetchAttemptRetryable(reason=f"{type(e).__name__} fetching {url}: {e}", retry_after_seconds=None)
|
||||
except (httpx.HTTPError, httpx.InvalidURL) as e:
|
||||
return _classify_fetch_error(e, url)
|
||||
return _classify_fetch_response(response, url)
|
||||
|
||||
|
||||
def _attempt_fetch_sync(client: _SyncGetClient, url: str, timeout: int) -> _FetchAttemptOutcome:
|
||||
try:
|
||||
response: Final = client.get(url, timeout=timeout)
|
||||
except (httpx.HTTPError, httpx.InvalidURL) as e:
|
||||
return _classify_fetch_error(e, url)
|
||||
return _classify_fetch_response(response, url)
|
||||
|
||||
|
||||
def _classify_fetch_response(response: httpx.Response, url: str) -> _FetchAttemptOutcome:
|
||||
if response.status_code in RETRYABLE_FETCH_STATUS_CODES:
|
||||
return _FetchAttemptRetryable(
|
||||
reason=f"HTTP {response.status_code} from {url}",
|
||||
|
|
@ -242,6 +255,22 @@ async def _attempt_fetch(
|
|||
return ModelCostMapReloaded(model_cost_map=parsed)
|
||||
|
||||
|
||||
def _next_retry_wait(
|
||||
outcome: _FetchAttemptRetryable, attempt: int, max_attempts: int, rng: random.Random
|
||||
) -> float | ModelCostMapReloadUnavailable:
|
||||
if attempt == max_attempts:
|
||||
return ModelCostMapReloadUnavailable(reason=f"{outcome.reason} (after {max_attempts} attempts)")
|
||||
wait_seconds: Final = _retry_wait_seconds(outcome=outcome, attempt=attempt, rng=rng)
|
||||
verbose_logger.warning(
|
||||
"LiteLLM: model cost map fetch attempt %d/%d failed (%s); retrying in %.1fs",
|
||||
attempt,
|
||||
max_attempts,
|
||||
outcome.reason,
|
||||
wait_seconds,
|
||||
)
|
||||
return wait_seconds
|
||||
|
||||
|
||||
async def _fetch_remote_model_cost_map_with_retry(
|
||||
url: str,
|
||||
timeout: int,
|
||||
|
|
@ -254,20 +283,32 @@ async def _fetch_remote_model_cost_map_with_retry(
|
|||
outcome = await _attempt_fetch(client=client, url=url, timeout=timeout)
|
||||
if not isinstance(outcome, _FetchAttemptRetryable):
|
||||
return outcome
|
||||
if attempt == max_attempts:
|
||||
return ModelCostMapReloadUnavailable(reason=f"{outcome.reason} (after {max_attempts} attempts)")
|
||||
wait_seconds = _retry_wait_seconds(outcome=outcome, attempt=attempt, rng=rng)
|
||||
verbose_logger.warning(
|
||||
"LiteLLM: model cost map fetch attempt %d/%d failed (%s); retrying in %.1fs",
|
||||
attempt,
|
||||
max_attempts,
|
||||
outcome.reason,
|
||||
wait_seconds,
|
||||
)
|
||||
wait_seconds = _next_retry_wait(outcome=outcome, attempt=attempt, max_attempts=max_attempts, rng=rng)
|
||||
if isinstance(wait_seconds, ModelCostMapReloadUnavailable):
|
||||
return wait_seconds
|
||||
await sleep(wait_seconds)
|
||||
return ModelCostMapReloadUnavailable(reason="model cost map fetch failed")
|
||||
|
||||
|
||||
def _fetch_remote_model_cost_map_with_retry_sync(
|
||||
url: str,
|
||||
timeout: int,
|
||||
max_attempts: int,
|
||||
sleep: Callable[[float], None],
|
||||
rng: random.Random,
|
||||
client: _SyncGetClient,
|
||||
) -> ModelCostMapReloadResult:
|
||||
for attempt in range(1, max_attempts + 1):
|
||||
outcome = _attempt_fetch_sync(client=client, url=url, timeout=timeout)
|
||||
if not isinstance(outcome, _FetchAttemptRetryable):
|
||||
return outcome
|
||||
wait_seconds = _next_retry_wait(outcome=outcome, attempt=attempt, max_attempts=max_attempts, rng=rng)
|
||||
if isinstance(wait_seconds, ModelCostMapReloadUnavailable):
|
||||
return wait_seconds
|
||||
sleep(wait_seconds)
|
||||
return ModelCostMapReloadUnavailable(reason="model cost map fetch failed")
|
||||
|
||||
|
||||
async def refetch_model_cost_map(
|
||||
url: str,
|
||||
timeout: int = 5,
|
||||
|
|
@ -423,13 +464,21 @@ def _finalize_model_cost_map(model_cost: dict) -> dict:
|
|||
return _expand_model_aliases(model_cost)
|
||||
|
||||
|
||||
def get_model_cost_map(url: str) -> dict:
|
||||
def get_model_cost_map(
|
||||
url: str,
|
||||
timeout: int = 5,
|
||||
max_attempts: int = MODEL_COST_MAP_FETCH_MAX_ATTEMPTS,
|
||||
sleep: Callable[[float], None] = time.sleep,
|
||||
rng: random.Random | None = None,
|
||||
client: "_SyncGetClient | None" = None,
|
||||
) -> dict:
|
||||
"""
|
||||
Public entry point — returns the model cost map dict.
|
||||
|
||||
1. If ``LITELLM_LOCAL_MODEL_COST_MAP`` is set, uses the local backup only.
|
||||
2. Otherwise fetches from ``url``, validates integrity, and falls back
|
||||
to the local backup on any failure.
|
||||
2. Otherwise fetches from ``url``, retrying transient HTTP errors
|
||||
(429/5xx/transport) with Retry-After-aware backoff, validates
|
||||
integrity, and falls back to the local backup on any failure.
|
||||
|
||||
Only the backup model count is cached (a single int) for validation.
|
||||
The full backup dict is only parsed when it must be *returned* as a
|
||||
|
|
@ -448,17 +497,24 @@ def get_model_cost_map(url: str) -> dict:
|
|||
_cost_map_source_info.url = url
|
||||
_cost_map_source_info.is_env_forced = False
|
||||
|
||||
try:
|
||||
content: Final = GetModelCostMap.fetch_remote_model_cost_map(url)
|
||||
except Exception as e:
|
||||
result: Final = _fetch_remote_model_cost_map_with_retry_sync(
|
||||
url=url,
|
||||
timeout=timeout,
|
||||
max_attempts=max_attempts,
|
||||
sleep=sleep,
|
||||
rng=rng if rng is not None else random.Random(),
|
||||
client=client if client is not None else httpx,
|
||||
)
|
||||
if isinstance(result, ModelCostMapReloadUnavailable):
|
||||
verbose_logger.warning(
|
||||
"LiteLLM: Failed to fetch remote model cost map from %s: %s. Falling back to local backup.",
|
||||
url,
|
||||
str(e),
|
||||
result.reason,
|
||||
)
|
||||
_cost_map_source_info.source = "local"
|
||||
_cost_map_source_info.fallback_reason = f"Remote fetch failed: {e}"
|
||||
_cost_map_source_info.fallback_reason = f"Remote fetch failed: {result.reason}"
|
||||
return _finalize_model_cost_map(GetModelCostMap.load_local_model_cost_map())
|
||||
content: Final = result.model_cost_map
|
||||
|
||||
# Validate using cached count (cheap int comparison, no file I/O)
|
||||
if not GetModelCostMap.validate_model_cost_map(
|
||||
|
|
|
|||
|
|
@ -111,6 +111,7 @@ from litellm.types.mcp import MCPPostCallResponseObject
|
|||
from litellm.types.prompts.init_prompts import PromptSpec
|
||||
from litellm.types.rerank import RerankResponse
|
||||
from litellm.types.utils import (
|
||||
DEPLOYMENT_SCOPED_PRICING_FIELDS,
|
||||
CachingDetails,
|
||||
CallTypes,
|
||||
CostBreakdown,
|
||||
|
|
@ -255,6 +256,7 @@ _STANDARD_LOGGING_METADATA_KEYS: Final[frozenset[str]] = frozenset(StandardLoggi
|
|||
|
||||
# Cache custom pricing keys as frozenset for O(1) lookups instead of looping through 49 keys
|
||||
_CUSTOM_PRICING_KEYS: Final[frozenset[str]] = frozenset(CustomPricingLiteLLMParams.model_fields.keys())
|
||||
_MODEL_INFO_CUSTOM_PRICING_KEYS: Final[frozenset[str]] = _CUSTOM_PRICING_KEYS | DEPLOYMENT_SCOPED_PRICING_FIELDS
|
||||
|
||||
sentry_sdk_instance = None
|
||||
capture_exception = None
|
||||
|
|
@ -2141,6 +2143,9 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
|
||||
logging_result: Final = self.normalize_logging_result(result=result)
|
||||
|
||||
if isinstance(result, Response) and isinstance(logging_result, (ModelResponse, EmbeddingResponse)):
|
||||
result = logging_result
|
||||
|
||||
if standard_logging_object is None and result is not None and self.stream is not True:
|
||||
if self._is_recognized_call_type_for_logging(logging_result=logging_result) or isinstance(
|
||||
logging_result, (dict, list)
|
||||
|
|
@ -2954,13 +2959,25 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
"Model=%s not found in completion cost map. Setting 'response_cost' to None", self.model
|
||||
)
|
||||
self.model_call_details["response_cost"] = None
|
||||
except Exception: # noqa: BLE001 # cost calculation must never block later callbacks (slot release)
|
||||
verbose_logger.exception(
|
||||
"Error calculating streaming response cost for model=%s. Setting 'response_cost' to None",
|
||||
self.model,
|
||||
)
|
||||
self.model_call_details["response_cost"] = None
|
||||
|
||||
self._merge_hidden_params_from_response_into_metadata(complete_streaming_response)
|
||||
|
||||
## STANDARDIZED LOGGING PAYLOAD
|
||||
self.model_call_details["standard_logging_object"] = self._build_standard_logging_payload(
|
||||
complete_streaming_response, start_time, end_time
|
||||
)
|
||||
try:
|
||||
self.model_call_details["standard_logging_object"] = self._build_standard_logging_payload(
|
||||
complete_streaming_response, start_time, end_time
|
||||
)
|
||||
except Exception: # noqa: BLE001 # payload build must never block later callbacks (slot release)
|
||||
verbose_logger.exception(
|
||||
"LiteLLM.LoggingError: [Non-Blocking] Exception building the standard logging payload "
|
||||
"for a streaming response; callbacks still run without it"
|
||||
)
|
||||
|
||||
# print standard logging payload
|
||||
if (standard_logging_payload := self.model_call_details.get("standard_logging_object")) is not None:
|
||||
|
|
@ -3000,32 +3017,39 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
## LOGGING HOOK ##
|
||||
|
||||
for callback in callbacks:
|
||||
if isinstance(callback, CustomGuardrail):
|
||||
from litellm.types.guardrails import GuardrailEventHooks
|
||||
try:
|
||||
if isinstance(callback, CustomGuardrail):
|
||||
from litellm.types.guardrails import GuardrailEventHooks
|
||||
|
||||
if (
|
||||
callback.should_run_guardrail(
|
||||
data=self.model_call_details,
|
||||
event_type=GuardrailEventHooks.logging_only,
|
||||
if (
|
||||
callback.should_run_guardrail(
|
||||
data=self.model_call_details,
|
||||
event_type=GuardrailEventHooks.logging_only,
|
||||
)
|
||||
is not True
|
||||
):
|
||||
continue
|
||||
|
||||
self.model_call_details, result = await callback.async_logging_hook(
|
||||
kwargs=self.model_call_details,
|
||||
result=result,
|
||||
call_type=self.call_type,
|
||||
)
|
||||
is not True
|
||||
):
|
||||
continue
|
||||
|
||||
self.model_call_details, result = await callback.async_logging_hook(
|
||||
kwargs=self.model_call_details,
|
||||
result=result,
|
||||
call_type=self.call_type,
|
||||
)
|
||||
elif isinstance(callback, CustomLogger):
|
||||
result = redact_message_input_output_from_custom_logger(
|
||||
result=result, litellm_logging_obj=self, custom_logger=callback
|
||||
)
|
||||
self.model_call_details, result = await callback.async_logging_hook(
|
||||
kwargs=self.model_call_details,
|
||||
result=result,
|
||||
call_type=self.call_type,
|
||||
elif isinstance(callback, CustomLogger):
|
||||
result = redact_message_input_output_from_custom_logger(
|
||||
result=result, litellm_logging_obj=self, custom_logger=callback
|
||||
)
|
||||
self.model_call_details, result = await callback.async_logging_hook(
|
||||
kwargs=self.model_call_details,
|
||||
result=result,
|
||||
call_type=self.call_type,
|
||||
)
|
||||
except Exception: # noqa: BLE001 # one failing hook must not skip later callbacks (slot release)
|
||||
verbose_logger.error(
|
||||
"LiteLLM.LoggingError: [Non-Blocking] Exception occurred in async_logging_hook %s",
|
||||
traceback.format_exc(),
|
||||
)
|
||||
self._handle_callback_failure(callback=callback)
|
||||
|
||||
self.has_run_logging(event_type="async_success")
|
||||
|
||||
|
|
@ -5030,7 +5054,9 @@ def use_custom_pricing_for_model(litellm_params: dict | None) -> bool:
|
|||
"""
|
||||
Check if the model uses custom pricing
|
||||
|
||||
Returns True if any of `SPECIAL_MODEL_INFO_PARAMS` are present in `litellm_params` or `model_info`
|
||||
Returns True if any custom pricing field is present in `litellm_params`, or if
|
||||
any custom pricing or deployment-scoped pricing field (such as
|
||||
``off_peak_pricing``) is present in the metadata ``model_info``
|
||||
"""
|
||||
if litellm_params is None:
|
||||
return False
|
||||
|
|
@ -5048,7 +5074,7 @@ def use_custom_pricing_for_model(litellm_params: dict | None) -> bool:
|
|||
model_info: dict = metadata.get("model_info", {}) or {}
|
||||
|
||||
if model_info:
|
||||
matching_keys = _CUSTOM_PRICING_KEYS & model_info.keys()
|
||||
matching_keys = _MODEL_INFO_CUSTOM_PRICING_KEYS & model_info.keys()
|
||||
for key in matching_keys:
|
||||
if model_info.get(key) is not None:
|
||||
return True
|
||||
|
|
@ -6152,7 +6178,10 @@ def get_standard_logging_object_payload(
|
|||
|
||||
def emit_standard_logging_payload(payload: StandardLoggingPayload):
|
||||
if os.getenv("LITELLM_PRINT_STANDARD_LOGGING_PAYLOAD"):
|
||||
print(json.dumps(payload, indent=4), flush=True) # noqa: T201
|
||||
try:
|
||||
print(json.dumps(payload, indent=4, default=str), flush=True) # noqa: T201
|
||||
except Exception as e: # noqa: BLE001 # Safe catch-all for verbose logging
|
||||
verbose_logger.exception("Error serializing standard logging payload for debug output: %s", e)
|
||||
|
||||
|
||||
def get_standard_logging_metadata(
|
||||
|
|
|
|||
|
|
@ -3,7 +3,7 @@ Helper utilities for tracking the cost of built-in tools.
|
|||
"""
|
||||
|
||||
from collections.abc import Mapping
|
||||
from typing import Any, Final, Literal
|
||||
from typing import Final, Literal
|
||||
|
||||
import litellm
|
||||
from litellm.constants import OPENAI_FILE_SEARCH_COST_PER_1K_CALLS
|
||||
|
|
@ -16,6 +16,7 @@ from litellm.types.llms.openai import (
|
|||
WebSearchOptions,
|
||||
)
|
||||
from litellm.types.utils import (
|
||||
ChatCompletionAnnotation,
|
||||
Message,
|
||||
ModelInfo,
|
||||
ModelResponse,
|
||||
|
|
@ -49,7 +50,7 @@ class StandardBuiltInToolCostTracking:
|
|||
@staticmethod
|
||||
def get_cost_for_built_in_tools(
|
||||
model: str,
|
||||
response_object: Any,
|
||||
response_object: object,
|
||||
usage: Usage | None = None,
|
||||
custom_llm_provider: str | None = None,
|
||||
standard_built_in_tools_params: StandardBuiltInToolsParams | None = None,
|
||||
|
|
@ -201,8 +202,7 @@ class StandardBuiltInToolCostTracking:
|
|||
model_info: Final = StandardBuiltInToolCostTracking._safe_get_model_info(
|
||||
model=model, custom_llm_provider=custom_llm_provider
|
||||
)
|
||||
file_search_raw: Final[Any] = standard_built_in_tools_params.get("file_search", {})
|
||||
file_search_usage: Final[FileSearchTool | None] = FileSearchTool(**file_search_raw) if file_search_raw else None
|
||||
file_search_usage: Final[FileSearchTool | None] = standard_built_in_tools_params.get("file_search") or None
|
||||
|
||||
# Convert model_info to dict and extract usage parameters
|
||||
model_info_dict: Final = dict(model_info) if model_info is not None else None
|
||||
|
|
@ -245,7 +245,7 @@ class StandardBuiltInToolCostTracking:
|
|||
|
||||
@staticmethod
|
||||
def _extract_file_search_params(
|
||||
file_search_usage: Any,
|
||||
file_search_usage: object,
|
||||
) -> tuple[float | None, float | None]:
|
||||
"""Extract and convert file search parameters safely."""
|
||||
storage_gb = None
|
||||
|
|
@ -335,7 +335,7 @@ class StandardBuiltInToolCostTracking:
|
|||
|
||||
@staticmethod
|
||||
def _extract_token_counts(
|
||||
computer_use_usage: Any,
|
||||
computer_use_usage: object,
|
||||
) -> tuple[int | None, int | None]:
|
||||
"""Extract and convert token counts safely."""
|
||||
input_tokens = None
|
||||
|
|
@ -351,9 +351,9 @@ class StandardBuiltInToolCostTracking:
|
|||
return input_tokens, output_tokens
|
||||
|
||||
@staticmethod
|
||||
def _safe_convert_to_int(value: Any) -> int | None:
|
||||
def _safe_convert_to_int(value: object) -> int | None:
|
||||
"""Safely convert a value to int."""
|
||||
if value is not None:
|
||||
if isinstance(value, (int, float, str)):
|
||||
try:
|
||||
return int(value)
|
||||
except (TypeError, ValueError):
|
||||
|
|
@ -381,7 +381,7 @@ class StandardBuiltInToolCostTracking:
|
|||
return usage.model_copy(update={"server_tool_use": server_tool_use})
|
||||
|
||||
@staticmethod
|
||||
def response_object_includes_web_search_call(response_object: Any, usage: Usage | None = None) -> bool:
|
||||
def response_object_includes_web_search_call(response_object: object, usage: Usage | None = None) -> bool:
|
||||
"""
|
||||
Check if the response object includes a web search call.
|
||||
|
||||
|
|
@ -446,7 +446,7 @@ class StandardBuiltInToolCostTracking:
|
|||
|
||||
@staticmethod
|
||||
def response_object_includes_file_search_call(
|
||||
response_object: Any,
|
||||
response_object: object,
|
||||
) -> bool:
|
||||
"""
|
||||
Check if the response object includes a file search call.
|
||||
|
|
@ -477,11 +477,11 @@ class StandardBuiltInToolCostTracking:
|
|||
message: Message | None = getattr(choice, "message", None)
|
||||
if message is None:
|
||||
continue
|
||||
if annotations := getattr(message, "annotations", None):
|
||||
if len(annotations) > 0:
|
||||
for annotation in annotations:
|
||||
if annotation.get("type", None) == annotation_type:
|
||||
return True
|
||||
annotations: list[ChatCompletionAnnotation] | None = getattr(message, "annotations", None)
|
||||
if annotations:
|
||||
for annotation in annotations:
|
||||
if annotation.get("type", None) == annotation_type:
|
||||
return True
|
||||
return False
|
||||
|
||||
@staticmethod
|
||||
|
|
@ -522,10 +522,8 @@ class StandardBuiltInToolCostTracking:
|
|||
if model_info is None:
|
||||
return 0.0
|
||||
|
||||
search_context_raw: Final[Any] = model_info.get("search_context_cost_per_query", {})
|
||||
search_context_pricing: Final[SearchContextCostPerQuery] = (
|
||||
SearchContextCostPerQuery(**search_context_raw) if search_context_raw else SearchContextCostPerQuery()
|
||||
)
|
||||
search_context_raw: Final = model_info.get("search_context_cost_per_query")
|
||||
search_context_pricing: Final[SearchContextCostPerQuery] = search_context_raw or SearchContextCostPerQuery()
|
||||
if web_search_options.get("search_context_size", None) == "low":
|
||||
return search_context_pricing.get("search_context_size_low", 0.0)
|
||||
elif web_search_options.get("search_context_size", None) == "medium":
|
||||
|
|
@ -545,10 +543,8 @@ class StandardBuiltInToolCostTracking:
|
|||
"""
|
||||
if model_info is None:
|
||||
return 0.0
|
||||
search_context_raw: Final[Any] = model_info.get("search_context_cost_per_query", {}) or {}
|
||||
search_context_pricing: Final[SearchContextCostPerQuery] = (
|
||||
SearchContextCostPerQuery(**search_context_raw) if search_context_raw else SearchContextCostPerQuery()
|
||||
)
|
||||
search_context_raw: Final = model_info.get("search_context_cost_per_query")
|
||||
search_context_pricing: Final[SearchContextCostPerQuery] = search_context_raw or SearchContextCostPerQuery()
|
||||
return search_context_pricing.get("search_context_size_medium", 0.0)
|
||||
|
||||
@staticmethod
|
||||
|
|
@ -714,7 +710,7 @@ class StandardBuiltInToolCostTracking:
|
|||
response_object: ModelResponse,
|
||||
) -> bool:
|
||||
for _choice in response_object.choices:
|
||||
message = getattr(_choice, "message", None)
|
||||
message: Message | None = getattr(_choice, "message", None)
|
||||
if (
|
||||
message is not None
|
||||
and hasattr(message, "annotations")
|
||||
|
|
|
|||
|
|
@ -2,10 +2,12 @@
|
|||
## Helper utilities for cost_per_token()
|
||||
|
||||
import re
|
||||
from collections.abc import Mapping
|
||||
from collections.abc import Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime, timezone, tzinfo
|
||||
from types import MappingProxyType
|
||||
from typing import Any, Final, Literal, TypedDict, cast
|
||||
from zoneinfo import ZoneInfo, ZoneInfoNotFoundError
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_logger
|
||||
|
|
@ -290,10 +292,187 @@ def _get_tiered_base_costs(model_info: ModelInfo, usage: Usage) -> tuple[float,
|
|||
)
|
||||
|
||||
|
||||
def _is_within_off_peak_window(off_peak_hours_utc: str | Sequence[str], current_time: datetime | None = None) -> bool:
|
||||
"""Return True if current_time (UTC, defaulting to now) falls inside any off-peak window.
|
||||
|
||||
off_peak_hours_utc is a "HH:MM-HH:MM" string in UTC, or a list of such strings for providers
|
||||
with multiple daily windows (e.g. ["16:30-00:30", "04:00-06:00"]). A window may wrap past
|
||||
midnight, and a window whose start equals its end covers the whole day. The start is
|
||||
inclusive and the end is exclusive; malformed windows are ignored.
|
||||
|
||||
An aware current_time is converted to UTC. A naive one is taken to already be UTC rather
|
||||
than being localised, so callers must pass datetime.now(timezone.utc), never datetime.now(),
|
||||
or every window shifts by the host's offset.
|
||||
"""
|
||||
reference: Final = current_time if current_time is not None else datetime.now(timezone.utc)
|
||||
now: Final = (reference.astimezone(timezone.utc) if reference.tzinfo is not None else reference).time()
|
||||
windows: Final = (off_peak_hours_utc,) if isinstance(off_peak_hours_utc, str) else off_peak_hours_utc
|
||||
for window in windows:
|
||||
try:
|
||||
start_str, end_str = window.split("-")
|
||||
start = datetime.strptime(start_str.strip(), "%H:%M").replace(tzinfo=timezone.utc).time()
|
||||
end = datetime.strptime(end_str.strip(), "%H:%M").replace(tzinfo=timezone.utc).time()
|
||||
except (ValueError, AttributeError):
|
||||
continue
|
||||
if start < end:
|
||||
if start <= now < end:
|
||||
return True
|
||||
elif now >= start or now < end:
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
_WEEKDAY_NUMBERS: Final = MappingProxyType(
|
||||
{
|
||||
"mon": 1,
|
||||
"monday": 1,
|
||||
"tue": 2,
|
||||
"tues": 2,
|
||||
"tuesday": 2,
|
||||
"wed": 3,
|
||||
"wednesday": 3,
|
||||
"thu": 4,
|
||||
"thur": 4,
|
||||
"thurs": 4,
|
||||
"thursday": 4,
|
||||
"fri": 5,
|
||||
"friday": 5,
|
||||
"sat": 6,
|
||||
"saturday": 6,
|
||||
"sun": 7,
|
||||
"sunday": 7,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _normalize_weekday(value: object) -> int | None:
|
||||
if isinstance(value, bool):
|
||||
return None
|
||||
if isinstance(value, int):
|
||||
return value if 1 <= value <= 7 else None
|
||||
if isinstance(value, str):
|
||||
return _WEEKDAY_NUMBERS.get(value.strip().lower())
|
||||
return None
|
||||
|
||||
|
||||
def _weekday_calendar(weekday_timezone: object) -> tzinfo:
|
||||
if isinstance(weekday_timezone, str) and weekday_timezone.strip():
|
||||
try:
|
||||
return ZoneInfo(weekday_timezone.strip())
|
||||
except (ValueError, ZoneInfoNotFoundError):
|
||||
return timezone.utc
|
||||
return timezone.utc
|
||||
|
||||
|
||||
def _matches_weekdays(reference_utc: datetime, weekdays: object, weekday_timezone: object) -> bool:
|
||||
"""Return True when reference_utc falls on one of the rule's weekdays, read on the calendar
|
||||
named by weekday_timezone (default UTC). An absent weekdays means every day. The calendar
|
||||
matters even when UTC and vendor-local weekdays agree at every currently priced hour: a
|
||||
window past 16:00 UTC is where an Asia/Shanghai weekday diverges from the UTC one.
|
||||
"""
|
||||
if weekdays is None:
|
||||
return True
|
||||
if isinstance(weekdays, str) or not isinstance(weekdays, Sequence):
|
||||
return False
|
||||
allowed: Final = frozenset(day for day in map(_normalize_weekday, weekdays) if day is not None)
|
||||
return reference_utc.astimezone(_weekday_calendar(weekday_timezone)).isoweekday() in allowed
|
||||
|
||||
|
||||
def _as_window_strings(value: object) -> tuple[str, ...]:
|
||||
if isinstance(value, str):
|
||||
return (value,)
|
||||
if isinstance(value, Sequence):
|
||||
return tuple(entry for entry in value if isinstance(entry, str))
|
||||
return ()
|
||||
|
||||
|
||||
def _is_off_peak(off_peak: Mapping[str, object], current_time: datetime | None = None) -> bool:
|
||||
"""Return True when current_time (UTC, defaulting to now) is off-peak under the block's
|
||||
rules: the flat hours_utc windows, which apply every day, or any entry in windows, whose
|
||||
hours apply only on its weekdays.
|
||||
"""
|
||||
reference: Final = current_time if current_time is not None else datetime.now(timezone.utc)
|
||||
reference_utc: Final = (
|
||||
reference.astimezone(timezone.utc) if reference.tzinfo is not None else reference.replace(tzinfo=timezone.utc)
|
||||
)
|
||||
flat_windows: Final = _as_window_strings(off_peak.get("hours_utc"))
|
||||
if flat_windows and _is_within_off_peak_window(flat_windows, reference_utc):
|
||||
return True
|
||||
windows: Final = off_peak.get("windows")
|
||||
if isinstance(windows, str) or not isinstance(windows, Sequence):
|
||||
return False
|
||||
weekday_timezone: Final = off_peak.get("weekday_timezone")
|
||||
for rule in windows:
|
||||
if not isinstance(rule, Mapping):
|
||||
continue
|
||||
rule_windows = _as_window_strings(rule.get("hours_utc"))
|
||||
if not rule_windows:
|
||||
continue
|
||||
if not _matches_weekdays(reference_utc, rule.get("weekdays"), weekday_timezone):
|
||||
continue
|
||||
if _is_within_off_peak_window(rule_windows, reference_utc):
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def _coerce_off_peak_rate(value: object, default: float) -> float:
|
||||
if isinstance(value, bool):
|
||||
return default
|
||||
if isinstance(value, (int, float)):
|
||||
return float(value)
|
||||
if isinstance(value, str):
|
||||
try:
|
||||
return float(value)
|
||||
except ValueError:
|
||||
return default
|
||||
return default
|
||||
|
||||
|
||||
def _apply_off_peak_pricing(
|
||||
model_info: ModelInfo,
|
||||
current_time: datetime | None,
|
||||
prompt_base_cost: float,
|
||||
completion_base_cost: float,
|
||||
cache_read_cost: float,
|
||||
) -> tuple[float, float, float]:
|
||||
"""Swap in off-peak per-token rates when the current UTC time is inside one of the model's
|
||||
off_peak_pricing rules, the every-day hours_utc windows or a day-of-week-qualified entry in
|
||||
windows. An off-peak rate replaces the rate that would otherwise apply rather than
|
||||
discounting it, so a model that also has tiered or above-threshold pricing bills the flat
|
||||
off-peak rate for the whole request while the window is open. Any rate left unset in
|
||||
off_peak_pricing falls back to the standard rate.
|
||||
"""
|
||||
off_peak: Final = model_info.get("off_peak_pricing")
|
||||
if not isinstance(off_peak, Mapping) or not _is_off_peak(off_peak, current_time):
|
||||
return prompt_base_cost, completion_base_cost, cache_read_cost
|
||||
return (
|
||||
_coerce_off_peak_rate(off_peak.get("input_cost_per_token"), prompt_base_cost),
|
||||
_coerce_off_peak_rate(off_peak.get("output_cost_per_token"), completion_base_cost),
|
||||
_coerce_off_peak_rate(off_peak.get("cache_read_input_token_cost"), cache_read_cost),
|
||||
)
|
||||
|
||||
|
||||
def _apply_off_peak_to_base_costs(
|
||||
model_info: ModelInfo,
|
||||
current_time: datetime | None,
|
||||
base_costs: tuple[float, float, float, float, float],
|
||||
) -> tuple[float, float, float, float, float]:
|
||||
"""Apply off-peak rates to an already-resolved set of base costs, whichever pricing path
|
||||
produced them. Cache-creation rates are passed through untouched, since off_peak_pricing
|
||||
has no field for them.
|
||||
"""
|
||||
prompt, completion, cache_creation, cache_creation_above_1hr, cache_read = base_costs
|
||||
off_peak_prompt, off_peak_completion, off_peak_cache_read = _apply_off_peak_pricing(
|
||||
model_info, current_time, prompt, completion, cache_read
|
||||
)
|
||||
return (off_peak_prompt, off_peak_completion, cache_creation, cache_creation_above_1hr, off_peak_cache_read)
|
||||
|
||||
|
||||
def _get_token_base_cost(
|
||||
model_info: ModelInfo,
|
||||
usage: Usage,
|
||||
service_tier: str | None = None,
|
||||
current_time: datetime | None = None,
|
||||
*,
|
||||
threshold_is_inclusive: bool = False,
|
||||
) -> tuple[float, float, float, float, float]:
|
||||
|
|
@ -311,7 +490,7 @@ def _get_token_base_cost(
|
|||
"""
|
||||
tiered_base_costs: Final = _get_tiered_base_costs(model_info=model_info, usage=usage)
|
||||
if tiered_base_costs is not None:
|
||||
return tiered_base_costs
|
||||
return _apply_off_peak_to_base_costs(model_info, current_time, tiered_base_costs)
|
||||
|
||||
# Get service tier aware cost keys
|
||||
input_cost_key: Final = _get_service_tier_cost_key("input_cost_per_token", service_tier)
|
||||
|
|
@ -345,12 +524,16 @@ def _get_token_base_cost(
|
|||
k for k in model_info if k.startswith("input_cost_per_token_above_") and not k.endswith(_SERVICE_TIER_SUFFIXES)
|
||||
]
|
||||
if not threshold_keys:
|
||||
return (
|
||||
prompt_base_cost,
|
||||
completion_base_cost,
|
||||
cache_creation_cost,
|
||||
cache_creation_cost_above_1hr,
|
||||
cache_read_cost,
|
||||
return _apply_off_peak_to_base_costs(
|
||||
model_info,
|
||||
current_time,
|
||||
(
|
||||
prompt_base_cost,
|
||||
completion_base_cost,
|
||||
cache_creation_cost,
|
||||
cache_creation_cost_above_1hr,
|
||||
cache_read_cost,
|
||||
),
|
||||
)
|
||||
|
||||
# Only sort the threshold keys (typically 1-2 keys instead of 66+)
|
||||
|
|
@ -451,12 +634,16 @@ def _get_token_base_cost(
|
|||
except Exception:
|
||||
continue
|
||||
|
||||
return (
|
||||
prompt_base_cost,
|
||||
completion_base_cost,
|
||||
cache_creation_cost,
|
||||
cache_creation_cost_above_1hr,
|
||||
cache_read_cost,
|
||||
return _apply_off_peak_to_base_costs(
|
||||
model_info,
|
||||
current_time,
|
||||
(
|
||||
prompt_base_cost,
|
||||
completion_base_cost,
|
||||
cache_creation_cost,
|
||||
cache_creation_cost_above_1hr,
|
||||
cache_read_cost,
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -2,6 +2,8 @@ import datetime
|
|||
from collections.abc import Mapping
|
||||
from typing import Any, Final
|
||||
|
||||
import httpx
|
||||
|
||||
from litellm.constants import LITELLM_DETAILED_TIMING
|
||||
from litellm.litellm_core_utils.core_helpers import process_response_headers
|
||||
from litellm.litellm_core_utils.llm_response_utils.get_api_base import get_api_base
|
||||
|
|
@ -59,11 +61,7 @@ class ResponseMetadata:
|
|||
@property
|
||||
def supports_response_time(self) -> bool:
|
||||
"""Check if response type supports timing metrics"""
|
||||
return (
|
||||
isinstance(self.result, ModelResponse)
|
||||
or isinstance(self.result, EmbeddingResponse)
|
||||
or isinstance(self.result, TranscriptionResponse)
|
||||
)
|
||||
return isinstance(self.result, (ModelResponse, EmbeddingResponse, TranscriptionResponse))
|
||||
|
||||
def set_hidden_params(self, logging_obj: LiteLLMLoggingObject, model: str | None, kwargs: dict) -> None:
|
||||
"""Set hidden parameters on the response"""
|
||||
|
|
@ -79,7 +77,7 @@ class ResponseMetadata:
|
|||
result=self.result, litellm_model_name=model, router_model_id=model_id
|
||||
),
|
||||
"additional_headers": process_response_headers(
|
||||
self._get_value_from_hidden_params("additional_headers") or {},
|
||||
self._get_additional_headers_from_hidden_params() or {},
|
||||
preserve_litellm_internal_headers=True,
|
||||
),
|
||||
"litellm_model_name": model,
|
||||
|
|
@ -98,12 +96,12 @@ class ResponseMetadata:
|
|||
for key, value in new_params.items():
|
||||
setattr(self._hidden_params, key, value)
|
||||
|
||||
def _get_value_from_hidden_params(self, key: str) -> Any | None:
|
||||
"""Get value from hidden params - handles when self._hidden_params is a dict or HiddenParams object"""
|
||||
def _get_additional_headers_from_hidden_params(self) -> httpx.Headers | dict[str, str] | None:
|
||||
"""Get `additional_headers` from hidden params - handles when self._hidden_params is a dict or HiddenParams object"""
|
||||
if isinstance(self._hidden_params, dict):
|
||||
return self._hidden_params.get(key, None)
|
||||
return self._hidden_params.get("additional_headers", None)
|
||||
elif isinstance(self._hidden_params, HiddenParams):
|
||||
return getattr(self._hidden_params, key, None)
|
||||
return getattr(self._hidden_params, "additional_headers", None)
|
||||
|
||||
def set_timing_metrics(
|
||||
self,
|
||||
|
|
@ -129,7 +127,7 @@ class ResponseMetadata:
|
|||
#########################################################
|
||||
# 2. Add callback processing duration
|
||||
#########################################################
|
||||
callback_duration_ms: Final = getattr(logging_obj, "callback_duration_ms", None)
|
||||
callback_duration_ms: Final[float | None] = getattr(logging_obj, "callback_duration_ms", None)
|
||||
if callback_duration_ms is not None:
|
||||
self._update_hidden_params(
|
||||
{
|
||||
|
|
@ -142,17 +140,17 @@ class ResponseMetadata:
|
|||
#########################################################
|
||||
llm_api_duration_ms: Final = logging_obj.model_call_details.get("llm_api_duration_ms")
|
||||
if LITELLM_DETAILED_TIMING and llm_api_duration_ms is not None:
|
||||
detailed: Final[dict] = {
|
||||
detailed: Final[dict[str, float]] = {
|
||||
"timing_llm_api_ms": round(llm_api_duration_ms, 4),
|
||||
}
|
||||
|
||||
# message copy time from Logging.__init__()
|
||||
msg_copy_ms: Final = getattr(logging_obj, "message_copy_duration_ms", None)
|
||||
msg_copy_ms: Final[float | None] = getattr(logging_obj, "message_copy_duration_ms", None)
|
||||
if msg_copy_ms is not None:
|
||||
detailed["timing_message_copy_ms"] = round(msg_copy_ms, 4)
|
||||
|
||||
# pre-processing = time from request start to LLM API call start
|
||||
api_call_start: Final = logging_obj.model_call_details.get("api_call_start_time")
|
||||
api_call_start: Final[datetime.datetime | None] = logging_obj.model_call_details.get("api_call_start_time")
|
||||
if api_call_start is not None and start_time is not None:
|
||||
pre_ms: Final = (api_call_start - start_time).total_seconds() * 1000
|
||||
detailed["timing_pre_processing_ms"] = round(pre_ms, 4)
|
||||
|
|
|
|||
Some files were not shown because too many files have changed in this diff Show more
Loading…
Add table
Reference in a new issue