scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-06 08:16:43 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 78

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

Browse source
This commit is contained in:
mateo-berri 2026-08-31 13:16:41 -07:00
commit 59732f068b
124 changed files with 5380 additions and 1971 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

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

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

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

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

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

2
CLAUDE.md
View file

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

32
basedpyright-code-budget.json
View file

@ -1,9 +1,9 @@
{
"reportAny": {
"limit": 17270
"limit": 16171
},
"reportArgumentType": {
"limit": 2538
"limit": 2226
},
"reportAssignmentType": {
"limit": 319
@ -18,13 +18,13 @@
"limit": 40
},
"reportDeprecated": {
"limit": 212
"limit": 211
},
"reportDuplicateImport": {
"limit": 19
},
"reportExplicitAny": {
"limit": 5485
"limit": 5199
},
"reportFunctionMemberAccess": {
"limit": 7
@ -54,10 +54,10 @@
"limit": 0
},
"reportMissingParameterType": {
"limit": 5658
"limit": 5611
},
"reportMissingTypeArgument": {
"limit": 15425
"limit": 15350
},
"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": 44368
},
"reportUnknownLambdaType": {
"limit": 109
},
"reportUnknownMemberType": {
"limit": 38721
"limit": 38468
},
"reportUnknownParameterType": {
"limit": 19778
"limit": 19665
},
"reportUnknownVariableType": {
"limit": 30290
"limit": 30066
},
"reportUnnecessaryCast": {
"limit": 117
"limit": 111
},
"reportUnnecessaryComparison": {
"limit": 697
"limit": 695
},
"reportUnnecessaryContains": {
"limit": 5
},
"reportUnnecessaryIsInstance": {
"limit": 829
"limit": 828
},
"reportUntypedBaseClass": {
"limit": 0

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

@ -7,7 +7,7 @@ GET - /audit/{id} - Get audit log by id
GET - /audit - Get all audit logs
"""
from typing import 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')",
),

20
litellm/_lazy_imports.py
View file

@ -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

85
litellm/caching/valkey_semantic_cache.py
View file

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

1
litellm/constants.py
View file

@ -1729,6 +1729,7 @@ SENTRY_DENYLIST: Final = [
"jwt_token",
"private_key",
"SLACK_WEBHOOK_URL",
"ALERTING_WEBHOOK_URL",
"webhook_url",
"LANGFUSE_SECRET_KEY",
# Email Configuration

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

@ -1,8 +1,9 @@
import json
from collections.abc import AsyncIterator, Iterator, Sequence
from typing import Any, Final, TypedDict, cast
from collections.abc import AsyncIterator, Callable, Iterator, Mapping, Sequence
from types import MappingProxyType
from typing import Any, Final, TypeAlias, cast
from typing_extensions import ReadOnly
from typing_extensions import ReadOnly, TypedDict
from litellm import verbose_logger
from litellm.litellm_core_utils.json_validation_rule import normalize_tool_schema
@ -11,7 +12,6 @@ from litellm.types.llms.openai import (
ChatCompletionAssistantMessage,
ChatCompletionAssistantToolCall,
ChatCompletionImageObject,
ChatCompletionRequest,
ChatCompletionSystemMessage,
ChatCompletionTextObject,
ChatCompletionToolCallFunctionChunk,
@ -23,35 +23,63 @@ from litellm.types.llms.openai import (
from litellm.types.router import GenericLiteLLMParams
from litellm.types.utils import (
AdapterCompletionStreamWrapper,
ChatCompletionDeltaCustomToolCall,
ChatCompletionDeltaToolCall,
ChatCompletionMessageCustomToolCall,
ChatCompletionMessageToolCall,
Choices,
Delta,
Function,
Message,
ModelResponse,
ModelResponseStream,
StreamingChoices,
Usage,
)
class _GenAITextPart(TypedDict, total=False):
text: ReadOnly[str]
_JsonDict: TypeAlias = dict[str, object]
_JsonDictList: TypeAlias = list[_JsonDict]
class _GenAISystemInstruction(TypedDict, total=False):
parts: ReadOnly[list[_GenAITextPart]]
class _ToolCallAccumulator(TypedDict):
name: ReadOnly[str]
arguments: ReadOnly[str]
class _GenAIFunctionCall(TypedDict):
name: ReadOnly[str]
args: ReadOnly[Mapping[str, object]]
class _GenAIPart(TypedDict, total=False):
text: ReadOnly[str]
functionCall: ReadOnly[dict[str, object]]
functionCall: ReadOnly[_GenAIFunctionCall]
class _GenAIFunctionResponse(TypedDict, total=False):
name: ReadOnly[str]
response: ReadOnly[object]
class _GenAIRequestFunctionCall(TypedDict, total=False):
name: ReadOnly[str]
args: ReadOnly[Mapping[str, object]]
class _GenAIContentPart(TypedDict, total=False):
text: ReadOnly[str]
inline_data: ReadOnly[Mapping[str, str]]
functionResponse: ReadOnly[_GenAIFunctionResponse]
functionCall: ReadOnly[_GenAIRequestFunctionCall]
class _GenAIFunctionDeclaration(TypedDict, total=False):
name: ReadOnly[str]
description: ReadOnly[str]
parametersJsonSchema: ReadOnly[dict[str, object]]
parametersJsonSchema: ReadOnly[object]
class _GenAITool(TypedDict, total=False):
functionDeclarations: ReadOnly[list[_GenAIFunctionDeclaration]]
functionDeclarations: ReadOnly[Sequence[_GenAIFunctionDeclaration]]
class _GenAIFunctionCallingConfig(TypedDict, total=False):
@ -62,9 +90,11 @@ class _GenAIToolConfig(TypedDict, total=False):
functionCallingConfig: ReadOnly[_GenAIFunctionCallingConfig]
def _decode_tool_call_arguments(raw_arguments: str) -> object:
"""Decode a tool call's JSON-encoded arguments into the value Google GenAI expects."""
return json.loads(raw_arguments)
class _GenAISystemInstruction(TypedDict, total=False):
parts: ReadOnly[Sequence[Mapping[str, str]]]
_EMPTY_STR_MAPPING: Final[Mapping[str, str]] = MappingProxyType({})
class GoogleGenAIStreamWrapper(AdapterCompletionStreamWrapper):
@ -74,12 +104,12 @@ 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 = {}
# State tracking for accumulating partial tool calls
self.accumulated_tool_calls = dict[int, _ToolCallAccumulator]()
self._returned_response = False
super().__init__(completion_stream)
@ -124,7 +154,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 +162,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 +181,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 +243,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 +284,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 +346,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 +360,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 +380,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 +394,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 +425,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 +534,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 +557,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 +597,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 +624,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 +633,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 +669,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 +680,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,21 +704,23 @@ 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"):
@ -701,19 +739,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 +762,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 +796,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,

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

@ -1485,9 +1485,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 +1516,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"}

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

@ -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

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

@ -7,7 +7,10 @@ litellm_content_retrieve tool calls server-side via the typed agentic loop plan.
import time
import uuid
from typing import 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

28
litellm/integrations/custom_logger.py
View file

@ -2,7 +2,7 @@
# On success, logs events to Promptlayer
import re
import traceback
from collections.abc import AsyncGenerator, Mapping
from collections.abc import AsyncGenerator, Mapping, Sequence
from typing import TYPE_CHECKING, Any, ClassVar, Final, Optional
from pydantic import BaseModel
@ -123,11 +123,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 []
@ -851,7 +851,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 +1005,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 +1037,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:
@ -1059,7 +1059,7 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
value: Any,
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)
@ -1090,16 +1090,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))

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

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

211
litellm/integrations/rubrik.py
View file

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

162
litellm/integrations/websearch_interception/handler.py
View file

@ -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,
@ -1583,10 +1675,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 +1686,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 +1695,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(

26
litellm/litellm_core_utils/llm_response_utils/response_metadata.py
View file

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

8
litellm/litellm_core_utils/streaming_handler.py
View file

@ -93,7 +93,7 @@ def print_verbose(print_statement: object):
@dataclass(frozen=True, slots=True)
class _ProviderChunkParsed:
response_obj: dict[str, Any]
response_obj: dict[str, object]
@dataclass(frozen=True, slots=True)
@ -1288,7 +1288,7 @@ class CustomStreamWrapper:
for key, value in anthropic_response_obj["provider_specific_fields"].items():
setattr(model_response, key, value)
response_obj = cast(dict[str, Any], anthropic_response_obj)
response_obj = cast(dict[str, object], anthropic_response_obj)
elif self.model == "replicate" or self.custom_llm_provider == "replicate":
response_obj = self.handle_replicate_chunk(chunk)
completion_obj["content"] = response_obj["text"]
@ -1444,7 +1444,7 @@ class CustomStreamWrapper:
if not isinstance(chunk, str):
raise ValueError(f"chunk is not a string: {chunk}")
response_obj = cast(
dict[str, Any],
dict[str, object],
litellm.CodestralTextCompletionConfig()._chunk_parser(chunk),
)
completion_obj["content"] = response_obj["text"]
@ -2551,7 +2551,7 @@ def calculate_total_usage(chunks: list[ModelResponse]) -> Usage:
prompt_tokens: int = 0
completion_tokens: int = 0
latest_usage_chunk = None
latest_usage_chunk: Usage | Mapping[str, int] | None = None
prompt_tokens_details: PromptTokensDetailsWrapper | None = None
completion_tokens_details: CompletionTokensDetailsWrapper | None = None
cache_creation_token_details: CacheCreationTokenDetails | None = None

4
litellm/litellm_core_utils/url_utils.py
View file

@ -417,7 +417,7 @@ def _extract_redirect_url(response: httpx.Response, request_url: str) -> str:
return str(httpx.URL(request_url).join(location))
def safe_get(client: Any, url: str, **kwargs: Any) -> Any:
def safe_get(client: Any, url: str, **kwargs: Any) -> httpx.Response:
"""
Fetch a user-supplied URL with SSRF protection on every redirect hop.
@ -460,7 +460,7 @@ def safe_get(client: Any, url: str, **kwargs: Any) -> Any:
raise SSRFError("Too many redirects")
async def async_safe_get(client: Any, url: str, **kwargs: Any) -> Any:
async def async_safe_get(client: Any, url: str, **kwargs: Any) -> httpx.Response:
"""Async version of safe_get."""
if not getattr(litellm, "user_url_validation", True):
kwargs.setdefault("follow_redirects", True)

59
litellm/llms/anthropic/batches/transformation.py
View file

@ -1,9 +1,11 @@
import json
import time
from collections.abc import Mapping
from typing import TYPE_CHECKING, Any, Final, Literal, cast
import httpx
from httpx import Headers, Response
from typing_extensions import ReadOnly, TypedDict
from litellm.litellm_core_utils.url_utils import encode_url_path_segment
from litellm.llms.base_llm.batches.transformation import BaseBatchesConfig
@ -21,6 +23,29 @@ else:
LoggingClass = Any
class AnthropicBatchRequestCounts(TypedDict, total=False):
"""The ``request_counts`` object of an Anthropic Message Batch."""
processing: ReadOnly[int]
succeeded: ReadOnly[int]
errored: ReadOnly[int]
canceled: ReadOnly[int]
expired: ReadOnly[int]
class AnthropicMessageBatch(TypedDict, total=False):
"""The fields of an Anthropic Message Batch that map onto an OpenAI Batch."""
id: ReadOnly[str]
processing_status: ReadOnly[str]
created_at: ReadOnly[str | None]
ended_at: ReadOnly[str | None]
expires_at: ReadOnly[str | None]
cancel_initiated_at: ReadOnly[str | None]
archived_at: ReadOnly[str | None]
request_counts: ReadOnly[AnthropicBatchRequestCounts]
class AnthropicBatchesConfig(BaseBatchesConfig):
def __init__(self):
from ..chat.transformation import AnthropicConfig
@ -85,7 +110,7 @@ class AnthropicBatchesConfig(BaseBatchesConfig):
create_batch_data: CreateBatchRequest,
optional_params: dict,
litellm_params: dict,
) -> bytes | str | dict[str, Any]:
) -> bytes | str | dict[str, object]:
"""
Transform the batch creation request to Anthropic format.
@ -135,7 +160,7 @@ class AnthropicBatchesConfig(BaseBatchesConfig):
batch_id: str,
optional_params: dict,
litellm_params: dict,
) -> bytes | str | dict[str, Any]:
) -> bytes | str | dict[str, object]:
"""
Transform batch retrieval request for Anthropic.
@ -154,7 +179,7 @@ class AnthropicBatchesConfig(BaseBatchesConfig):
) -> LiteLLMBatch:
"""Transform Anthropic MessageBatch retrieval response to LiteLLM format."""
try:
response_data: Final = raw_response.json()
response_data: Final[AnthropicMessageBatch] = raw_response.json()
except Exception as e:
raise ValueError(f"Failed to parse Anthropic batch response: {e}")
@ -163,18 +188,20 @@ class AnthropicBatchesConfig(BaseBatchesConfig):
processing_status: Final = response_data.get("processing_status", "in_progress")
# Map Anthropic processing_status to OpenAI status
status_mapping: dict[
str,
Literal[
"validating",
"failed",
"in_progress",
"finalizing",
"completed",
"expired",
"cancelling",
"cancelled",
],
status_mapping: Final[
Mapping[
str,
Literal[
"validating",
"failed",
"in_progress",
"finalizing",
"completed",
"expired",
"cancelling",
"cancelled",
],
]
] = {
"in_progress": "in_progress",
"canceling": "cancelling",
@ -281,7 +308,7 @@ class AnthropicBatchesConfig(BaseBatchesConfig):
if not line:
continue
try:
response_json = json.loads(line)
response_json: Mapping[str, Mapping[str, dict[str, object]]] = json.loads(line)
# Update model_response with the parsed JSON
completion_response = response_json["result"]["message"]
transformed_response = self.anthropic_chat_config.transform_parsed_response(

186
litellm/llms/anthropic/chat/guardrail_translation/handler.py
View file

@ -16,9 +16,9 @@ import json
from collections.abc import Mapping, Sequence
from copy import deepcopy
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any, Final, cast
from typing import TYPE_CHECKING, Any, Final, Protocol, cast, overload, runtime_checkable
from typing_extensions import assert_never
from typing_extensions import ReadOnly, TypedDict, assert_never
from litellm._logging import verbose_proxy_logger
from litellm.llms.anthropic.chat.transformation import AnthropicConfig
@ -58,6 +58,8 @@ from litellm.types.utils import (
)
if TYPE_CHECKING:
from fastapi import HTTPException
from litellm.integrations.custom_guardrail import (
CustomGuardrail,
ModifyResponseException,
@ -98,6 +100,48 @@ InputWriteBackTarget = (
)
class _SSEDelta(TypedDict, total=False):
type: ReadOnly[str]
text: ReadOnly[str]
stop_reason: ReadOnly[str | None]
class _SSEEventData(TypedDict, total=False):
delta: ReadOnly[_SSEDelta]
def _as_str_mapping(value: Mapping[str, object]) -> Mapping[str, object]:
return value
def _content_block_at(blocks: Sequence[object], index: int) -> object:
return blocks[index]
@runtime_checkable
class _ModelDumpBlock(Protocol):
def model_dump(self) -> Mapping[str, object]: ...
@runtime_checkable
class _TextAttrBlock(Protocol):
text: str
class _WritableMessage(Protocol):
@overload
def get(self, key: str, /) -> object | None: ...
@overload
def get(self, key: str, default: object, /) -> object: ...
def __setitem__(self, key: str, value: object, /) -> None: ...
def _as_writable(value: _WritableMessage) -> _WritableMessage:
return value
@dataclass(frozen=True, slots=True)
class ScannedText:
text: str
@ -126,7 +170,7 @@ class AnthropicMessagesHandler(BaseTranslation):
@staticmethod
def _build_streaming_usage_response(
responses_so_far: list[object],
responses_so_far: Sequence[object],
request_data: dict | None,
) -> ModelResponse | None:
chunks: Final = tuple(response for response in responses_so_far if isinstance(response, (str, bytes)))
@ -144,7 +188,7 @@ class AnthropicMessagesHandler(BaseTranslation):
self,
exc: "ModifyResponseException",
stream_started: bool = False,
responses_so_far: list[object] | None = None,
responses_so_far: Sequence[object] | None = None,
) -> list[bytes]:
"""
Build an Anthropic SSE sequence delivering the guardrail block message
@ -162,9 +206,22 @@ class AnthropicMessagesHandler(BaseTranslation):
would make Anthropic clients reject the stream.
"""
if stream_started:
return self._block_continuation_chunks(exc, responses_so_far or [])
return list(self._block_continuation_chunks(exc, responses_so_far or []))
return self._standalone_block_chunks(exc)
def build_stream_error_items(
self,
exc: "HTTPException",
responses_so_far: Sequence[Any] | None = None,
) -> Sequence[Any] | None:
from litellm.proxy.common_request_processing import (
serialize_http_exception_detail,
)
from litellm.proxy.guardrails.anthropic_sse import anthropic_sse_error_frames
message, _ = serialize_http_exception_detail(exc.detail)
return tuple(anthropic_sse_error_frames(message))
def _standalone_block_chunks(self, exc: "ModifyResponseException") -> list[bytes]:
import uuid
@ -187,7 +244,9 @@ class AnthropicMessagesHandler(BaseTranslation):
)
return list(FakeAnthropicMessagesStreamIterator(response=block_response))
def _block_continuation_chunks(self, exc: "ModifyResponseException", responses_so_far: list[object]) -> list[bytes]:
def _block_continuation_chunks(
self, exc: "ModifyResponseException", responses_so_far: Sequence[object]
) -> Sequence[bytes]:
"""Continue an already-started message: close the open content block,
append the block message as a new text block, then end the message --
without a second message_start."""
@ -199,7 +258,7 @@ class AnthropicMessagesHandler(BaseTranslation):
def _sse(event_type: str, payload: dict) -> bytes:
return f"event: {event_type}\ndata: {json.dumps(payload)}\n\n".encode()
output_tokens: Final = blocked_response_usage(getattr(exc, "original_response", None))["output_tokens"]
output_tokens: Final = blocked_response_usage(getattr(exc, "original_response", None)).get("output_tokens", 0)
open_index, max_index = self._content_block_state(responses_so_far)
new_index: Final = (max_index + 1) if max_index is not None else 0
chunks: list[bytes] = []
@ -237,7 +296,7 @@ class AnthropicMessagesHandler(BaseTranslation):
@staticmethod
def _content_block_state(
responses_so_far: list[object],
responses_so_far: Sequence[object],
) -> tuple[int | None, int | None]:
"""From the SSE chunks already sent to the client, return (open
content-block index or None, highest content-block index seen or None).
@ -263,7 +322,20 @@ class AnthropicMessagesHandler(BaseTranslation):
return open_index, max_index
@staticmethod
def _iter_sse_events(item: object) -> list[dict[str, object]]:
def _parse_sse_data_line(raw_line: str) -> tuple[Mapping[str, object], ...]:
line: Final = raw_line.strip()
if not line.startswith("data:"):
return ()
try:
parsed: Final[object] = json.loads(line[len("data:") :].strip())
except json.JSONDecodeError:
return ()
if not isinstance(parsed, dict):
return ()
return (_as_str_mapping(parsed),)
@staticmethod
def _iter_sse_events(item: object) -> Sequence[Mapping[str, object]]:
"""Yield the event-data dicts in one stream chunk.
Handles both formats this stream can carry (see
@ -271,24 +343,15 @@ class AnthropicMessagesHandler(BaseTranslation):
several events separated by a blank line -- and an already-parsed event
``dict``."""
if isinstance(item, dict):
return [item]
return (_as_str_mapping(item),)
if not isinstance(item, (bytes, bytearray)):
return []
events: Final[list[dict[str, object]]] = []
for block in item.decode("utf-8", errors="replace").split("\n\n"):
for line in block.split("\n"):
line = line.strip()
if not line.startswith("data:"):
continue
try:
parsed: str | int | float | bool | None | Sequence[object] | Mapping[str, object] = json.loads(
line[len("data:") :].strip()
)
except json.JSONDecodeError:
continue
if isinstance(parsed, dict):
events.append(parsed)
return events
return ()
return tuple(
event
for block in item.decode("utf-8", errors="replace").split("\n\n")
for line in block.split("\n")
for event in AnthropicMessagesHandler._parse_sse_data_line(line)
)
def _translate_to_openai(self, data: dict) -> ChatCompletionRequest:
"""Translate Anthropic request to OpenAI chat completion format."""
@ -321,7 +384,7 @@ class AnthropicMessagesHandler(BaseTranslation):
data: dict,
guardrail_to_apply: "CustomGuardrail",
litellm_logging_obj: "LiteLLMLoggingObj | None" = None,
) -> Any:
) -> Mapping[str, object]:
"""
Process input messages by applying guardrails to text content.
"""
@ -481,7 +544,7 @@ class AnthropicMessagesHandler(BaseTranslation):
@staticmethod
def _openai_system_message_to_anthropic(
message: dict[str, object],
message: Mapping[str, object],
) -> dict[str, object] | None: # mutable-ok: API message payload
"""Convert an OpenAI system message to the client's Anthropic-shaped entry."""
content: Final = message.get("content")
@ -561,7 +624,7 @@ class AnthropicMessagesHandler(BaseTranslation):
@staticmethod
def _defer_systems_inside_tool_exchanges(
structured_messages: list, # mutable-ok: API message payload
structured_messages: Sequence[Mapping[str, object]],
) -> list:
"""Hold a system row until the tool exchange around it completes so the call/result pair converts together."""
from litellm.litellm_core_utils.prompt_templates.factory import group_tool_exchanges
@ -755,7 +818,7 @@ class AnthropicMessagesHandler(BaseTranslation):
if scan_only_tool_results:
return EMPTY_EXTRACTED_INPUT
text_str: Final = content_item.get("text", None)
text_str: Final[str | None] = content_item.get("text")
return ExtractedInput(
scanned=(
() if text_str is None else (ScannedText(text_str, ContentBlockTextTarget(msg_idx, content_idx)),)
@ -805,7 +868,7 @@ class AnthropicMessagesHandler(BaseTranslation):
async def _apply_guardrail_responses_to_input(
self,
messages: list[dict[str, object]],
messages: Sequence[_WritableMessage],
responses: list[str],
scanned: tuple[ScannedText, ...],
) -> None:
@ -931,7 +994,7 @@ class AnthropicMessagesHandler(BaseTranslation):
litellm_logging_obj: "LiteLLMLoggingObj | None" = None,
user_api_key_dict: "UserAPIKeyAuth | None" = None,
request_data: dict | None = None,
) -> list[Any]:
) -> Sequence[object]:
"""
Process output streaming response by applying guardrails to text content.
@ -1027,7 +1090,7 @@ class AnthropicMessagesHandler(BaseTranslation):
return request_data
@staticmethod
def _get_response_content(response: object) -> list[Any]:
def _get_response_content(response: object) -> Sequence[object]:
"""Extract content list from a dict or object response."""
if isinstance(response, dict):
return response.get("content", []) or []
@ -1037,7 +1100,7 @@ class AnthropicMessagesHandler(BaseTranslation):
def _extract_from_content_blocks(
self,
response_content: list[Any],
response_content: Sequence[object],
texts_to_check: list[str],
images_to_check: list[str],
task_mappings: list[tuple[int, int | None]],
@ -1045,21 +1108,10 @@ class AnthropicMessagesHandler(BaseTranslation):
) -> None:
"""Extract text, images, and tool calls from content blocks."""
for content_idx, content_block in enumerate(response_content):
block_dict: dict[str, object] = {}
if isinstance(content_block, dict):
block_type = content_block.get("type")
block_dict = cast(dict[str, object], content_block)
elif hasattr(content_block, "type"):
block_type = getattr(content_block, "type", None)
if hasattr(content_block, "model_dump"):
block_dict = content_block.model_dump()
else:
block_dict = {
"type": block_type,
"text": getattr(content_block, "text", None),
}
else:
fields = self._output_block_fields(content_block)
if fields is None:
continue
block_type, block_dict = fields
if block_type in ["text", "tool_use"]:
self._extract_output_text_and_images(
@ -1071,6 +1123,21 @@ class AnthropicMessagesHandler(BaseTranslation):
tool_calls_to_check=tool_calls_to_check,
)
@staticmethod
def _output_block_fields(content_block: object) -> "tuple[object, Mapping[str, object]] | None":
if isinstance(content_block, dict):
block_dict: Final = _as_str_mapping(content_block)
return block_dict.get("type"), block_dict
if not hasattr(content_block, "type"):
return None
block_type: Final = getattr(content_block, "type", None)
if isinstance(content_block, _ModelDumpBlock):
return block_type, content_block.model_dump()
return block_type, {
"type": block_type,
"text": getattr(content_block, "text", None),
}
@staticmethod
def _build_guardrail_inputs(
texts_to_check: list[str],
@ -1093,7 +1160,7 @@ class AnthropicMessagesHandler(BaseTranslation):
inputs["model"] = response_model
return inputs
def get_streaming_string_so_far(self, responses_so_far: list[Any]) -> str:
def get_streaming_string_so_far(self, responses_so_far: Sequence[object]) -> str:
"""
Parse streaming responses and extract accumulated text content.
@ -1164,7 +1231,7 @@ class AnthropicMessagesHandler(BaseTranslation):
# Only process content_block_delta events
if event_type == "content_block_delta" and data_line:
try:
data = json.loads(data_line)
data: _SSEEventData = json.loads(data_line)
delta = data.get("delta", {})
if delta.get("type") == "text_delta":
text += delta.get("text", "")
@ -1176,7 +1243,7 @@ class AnthropicMessagesHandler(BaseTranslation):
return text
def _check_streaming_has_ended(self, responses_so_far: list[Any]) -> bool:
def _check_streaming_has_ended(self, responses_so_far: Sequence[object]) -> bool:
"""
Check if streaming response has ended by looking for non-null stop_reason.
@ -1227,7 +1294,7 @@ class AnthropicMessagesHandler(BaseTranslation):
# Check for message_delta event with stop_reason
if event_type == "message_delta" and data_line:
try:
data = json.loads(data_line)
data: _SSEEventData = json.loads(data_line)
delta = data.get("delta", {})
stop_reason = delta.get("stop_reason")
if stop_reason is not None:
@ -1271,7 +1338,7 @@ class AnthropicMessagesHandler(BaseTranslation):
def _extract_output_text_and_images(
self,
content_block: dict[str, object],
content_block: Mapping[str, object],
content_idx: int,
texts_to_check: list[str],
images_to_check: list[str],
@ -1294,7 +1361,7 @@ class AnthropicMessagesHandler(BaseTranslation):
task_mappings.append((content_idx, None))
# Extract tool calls
elif content_type == "tool_use":
elif content_type == "tool_use" and isinstance(content_block, dict):
tool_call: Final = AnthropicConfig.convert_tool_use_to_openai_format(
anthropic_tool_content=content_block,
index=content_idx,
@ -1319,7 +1386,7 @@ class AnthropicMessagesHandler(BaseTranslation):
content_idx = cast(int, mapping[0])
# Handle both dict and object responses
response_content: list[Any] = []
response_content: Sequence[object] = []
if isinstance(response, dict):
response_content = response.get("content", []) or []
elif hasattr(response, "content"):
@ -1335,14 +1402,15 @@ class AnthropicMessagesHandler(BaseTranslation):
if content_idx >= len(response_content):
continue
content_block = response_content[content_idx]
content_block = _content_block_at(response_content, content_idx)
# Verify it's a text block and update the text field
# Handle both dict and Pydantic object content blocks
if isinstance(content_block, dict):
if content_block.get("type") == "text":
cast(dict[str, object], content_block)["text"] = guardrail_response
block = _as_writable(content_block)
if block.get("type") == "text":
block["text"] = guardrail_response
elif hasattr(content_block, "type") and getattr(content_block, "type", None) == "text":
# Update Pydantic object's text attribute
if hasattr(content_block, "text"):
if isinstance(content_block, _TextAttrBlock):
content_block.text = guardrail_response

243
litellm/llms/anthropic/experimental_pass_through/context_management/editors/compact.py
View file

@ -13,10 +13,10 @@ Mirrors Anthropic's native ``compact_20260112`` for non-Anthropic providers:
"""
import re
from collections.abc import Mapping, Sequence
from typing import TYPE_CHECKING, Any, Final, Literal, NotRequired, Optional, TypedDict, Union, cast
from collections.abc import Awaitable, Mapping, Sequence
from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Protocol, TypeVar, Union, cast
from typing_extensions import ReadOnly
from typing_extensions import NotRequired, ReadOnly, TypedDict, Unpack
import litellm
from litellm._logging import verbose_logger
@ -29,6 +29,7 @@ from litellm.types.llms.anthropic import (
if TYPE_CHECKING:
from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.hooks.parallel_request_limiter_v3 import RateLimitDescriptor, RateLimitResponse
from litellm.router import Router
from litellm.types.llms.anthropic import (
AllAnthropicPassThroughMessageValues,
@ -84,6 +85,77 @@ _PROPAGATED_METADATA_KEYS: Final = (
_SUMMARY_TAG_RE: Final = re.compile(r"<summary>(.*?)</summary>", re.IGNORECASE | re.DOTALL)
_MsgT: Final = TypeVar("_MsgT", bound=Mapping[str, object])
def _as_object(value: object) -> object:
return value
def _is_tool_result_block(block: object) -> bool:
return isinstance(block, dict) and block.get("type") in ("tool_result",)
class _SummaryCallKwargs(TypedDict):
model: ReadOnly[str]
max_tokens: ReadOnly[int]
timeout: ReadOnly[float]
litellm_metadata: ReadOnly[Mapping[str, object]]
user: ReadOnly[NotRequired[str]]
allowed_model_region: ReadOnly[NotRequired[str]]
class _SummaryOptionalKwargs(TypedDict, total=False):
user: ReadOnly[str]
allowed_model_region: ReadOnly[str]
class _SummaryAcompletion(Protocol):
def __call__(
self,
*,
messages: Sequence[Mapping[str, object]],
**kwargs: Unpack[_SummaryCallKwargs], # kwargs-ok: forwarded verbatim to acompletion, which owns them
) -> "Awaitable[ModelResponse | CustomStreamWrapper]": ...
class _CreateRateLimitDescriptors(Protocol):
def __call__(
self,
*,
user_api_key_dict: "UserAPIKeyAuth",
data: Mapping[str, str],
rpm_limit_type: object,
tpm_limit_type: object,
model_has_failures: bool,
) -> "Sequence[RateLimitDescriptor]": ...
class _AddModelRateLimitDescriptor(Protocol):
def __call__(
self,
*,
user_api_key_dict: "UserAPIKeyAuth",
requested_model: str,
descriptors: "Sequence[RateLimitDescriptor]",
) -> None: ...
class _CreateOrgRateLimitDescriptors(Protocol):
def __call__(
self, user_api_key_dict: "UserAPIKeyAuth", requested_model: str | None = None
) -> "Sequence[RateLimitDescriptor]": ...
class _ShouldRateLimit(Protocol):
def __call__(
self,
*,
descriptors: "Sequence[RateLimitDescriptor]",
parent_otel_span: object,
read_only: bool,
) -> "Awaitable[RateLimitResponse]": ...
def _read_summary_model_setting() -> str | None:
"""Look up the configured summarization model from proxy general_settings."""
@ -159,11 +231,11 @@ async def _check_summary_model_access(
return True
key_models: Final = list(getattr(user_api_key_auth, "models", None) or [])
team_id: Final = getattr(user_api_key_auth, "team_id", None)
team_id: Final[str | None] = getattr(user_api_key_auth, "team_id", None)
team_model_aliases: Final = getattr(user_api_key_auth, "team_model_aliases", None)
team_models: Final = list(getattr(user_api_key_auth, "team_models", None) or [])
user_id: Final = getattr(user_api_key_auth, "user_id", None)
project_id: Final = getattr(user_api_key_auth, "project_id", None)
user_id: Final[str | None] = getattr(user_api_key_auth, "user_id", None)
project_id: Final[str | None] = getattr(user_api_key_auth, "project_id", None)
checks: Final[tuple[tuple[Literal["key", "team"], list[str]], ...]] = (
("key", key_models),
@ -372,7 +444,7 @@ async def _check_summary_model_budget(
return False
end_user_model_max_budget: Final = getattr(user_api_key_auth, "end_user_model_max_budget", None)
end_user_id: Final = getattr(user_api_key_auth, "end_user_id", None)
end_user_id: Final[str | None] = getattr(user_api_key_auth, "end_user_id", None)
if isinstance(end_user_model_max_budget, dict) and end_user_model_max_budget and end_user_id is not None:
try:
await model_max_budget_limiter.is_end_user_within_model_budget(
@ -424,40 +496,57 @@ async def _check_summary_model_rate_limit(
except Exception:
return True
limiter: Final = getattr(proxy_logging_obj, "max_parallel_request_limiter", None)
limiter: Final[object] = getattr(proxy_logging_obj, "max_parallel_request_limiter", None)
should_rate_limit_check: Final[_ShouldRateLimit | None] = getattr(limiter, "should_rate_limit", None)
create_descriptors: Final[_CreateRateLimitDescriptors | None] = getattr(
limiter, "_create_rate_limit_descriptors", None
)
add_team_descriptor: Final[_AddModelRateLimitDescriptor | None] = getattr(
limiter, "_add_team_model_rate_limit_descriptor_from_metadata", None
)
add_project_descriptor: Final[_AddModelRateLimitDescriptor | None] = getattr(
limiter, "_add_project_model_rate_limit_descriptor_from_metadata", None
)
create_org_descriptors: Final[_CreateOrgRateLimitDescriptors | None] = getattr(
limiter, "create_organization_rate_limit_descriptor", None
)
if (
limiter is None
or not hasattr(limiter, "should_rate_limit")
or not hasattr(limiter, "_create_rate_limit_descriptors")
or should_rate_limit_check is None
or create_descriptors is None
or add_team_descriptor is None
or add_project_descriptor is None
or create_org_descriptors is None
):
return True
try:
metadata: Final = getattr(user_api_key_auth, "metadata", None) or {}
metadata: Final[Mapping[str, object]] = getattr(user_api_key_auth, "metadata", None) or {}
data: Final = {"model": summary_model}
descriptors: Final = limiter._create_rate_limit_descriptors(
base_descriptors: Final = create_descriptors(
user_api_key_dict=user_api_key_auth,
data=data,
rpm_limit_type=metadata.get("rpm_limit_type"),
tpm_limit_type=metadata.get("tpm_limit_type"),
model_has_failures=False,
)
limiter._add_team_model_rate_limit_descriptor_from_metadata(
add_team_descriptor(
user_api_key_dict=user_api_key_auth,
requested_model=summary_model,
descriptors=descriptors,
descriptors=base_descriptors,
)
limiter._add_project_model_rate_limit_descriptor_from_metadata(
add_project_descriptor(
user_api_key_dict=user_api_key_auth,
requested_model=summary_model,
descriptors=descriptors,
descriptors=base_descriptors,
)
descriptors.extend(limiter.create_organization_rate_limit_descriptor(user_api_key_auth, summary_model))
descriptors: Final = (*base_descriptors, *create_org_descriptors(user_api_key_auth, summary_model))
if not descriptors:
return True
response: Final = await limiter.should_rate_limit(
parent_otel_span: Final[object] = getattr(user_api_key_auth, "parent_otel_span", None)
response: Final[RateLimitResponse] = await should_rate_limit_check(
descriptors=descriptors,
parent_otel_span=getattr(user_api_key_auth, "parent_otel_span", None),
parent_otel_span=parent_otel_span,
read_only=True,
)
except Exception as e:
@ -471,7 +560,7 @@ async def _check_summary_model_rate_limit(
def _find_latest_compaction_index(
messages: list[dict[str, object]],
messages: Sequence[Mapping[str, object]],
) -> tuple[int | None, int | None]:
"""Return (message_index, block_index) of the most recent compaction block.
@ -490,8 +579,8 @@ def _find_latest_compaction_index(
def _slice_around_compaction_block(
messages: list[dict[str, Any]],
) -> tuple[list[dict[str, object]], dict[str, object] | None]:
messages: Sequence[_MsgT],
) -> tuple[Sequence[_MsgT | dict[str, object]], dict[str, object] | None]:
"""Apply Anthropic's "drop everything before the compaction block" rule.
Returns ``(sliced_messages_with_compaction_block, compaction_block_dict)``
@ -506,19 +595,21 @@ def _slice_around_compaction_block(
original_msg: Final = messages[msg_idx]
original_content: Final = original_msg["content"]
compaction_block: Final = cast(dict[str, object], original_content[blk_idx])
if not isinstance(original_content, list):
return messages, None
original_blocks: Final = cast("Sequence[dict[str, object]]", original_content)
compaction_block: Final = original_blocks[blk_idx]
# Per Anthropic's contract everything before the compaction block is
# dropped, including earlier blocks within the same assistant message.
sliced_content: Final = list(original_content[blk_idx:])
sliced_content: Final = list(original_blocks[blk_idx:])
sliced_messages: Final[list[dict[str, object]]] = [{**original_msg, "content": sliced_content}]
sliced_messages.extend(messages[msg_idx + 1 :])
sliced_messages: Final = [{**original_msg, "content": sliced_content}, *messages[msg_idx + 1 :]]
return sliced_messages, compaction_block
def _strip_compaction_blocks(
messages: list[dict[str, object]],
messages: Sequence[dict[str, object]],
) -> list[dict[str, object]]:
"""Drop any ``compaction`` content blocks from messages.
@ -625,7 +716,7 @@ def _propagate_metadata(
def _count_effective_tokens(
model: str,
effective_messages: list[dict[str, object]],
effective_messages: Sequence[dict[str, object]],
compaction_block: CompactionBlock | None,
tools: list[dict[str, object]] | None,
system: str | list[dict[str, object]] | None = None,
@ -704,17 +795,18 @@ def _system_to_text(
return ""
if isinstance(system, str):
return system
parts: Final[list[str]] = []
for block in system:
if isinstance(block, dict) and block.get("type") == "text":
text = block.get("text")
if isinstance(text, str) and text:
parts.append(text)
return "\n".join(parts)
return "\n".join(
text
for block in system
if isinstance(block, dict)
and block.get("type") == "text"
and isinstance(text := block.get("text"), str)
and text
)
def _select_last_user_question(
messages: list[dict[str, object]],
messages: Sequence[dict[str, object]],
) -> list[dict[str, object]]:
"""Pick the most recent ``user`` turn that is a real question.
@ -729,16 +821,18 @@ def _select_last_user_question(
turns, or contained no user turns at all). The downstream call always
needs a non-empty user message.
"""
blocks: Sequence[object]
for msg in reversed(messages):
if msg.get("role") != "user":
continue
content = msg.get("content")
if isinstance(content, list):
filtered = [blk for blk in content if not (isinstance(blk, dict) and blk.get("type") == "tool_result")]
blocks = [*map(_as_object, content)]
filtered = [blk for blk in blocks if not _is_tool_result_block(blk)]
if not filtered:
# Purely tool_result — skip and look for an earlier turn.
continue
if len(filtered) < len(content):
if len(filtered) < len(blocks):
return [{**msg, "content": filtered}]
return [msg]
return [
@ -761,7 +855,7 @@ def _extract_summary_text(raw: str | None) -> str | None:
def _system_to_openai_message(
system: str | list[dict[str, Any]] | None,
) -> dict[str, object] | None:
) -> Mapping[str, object] | None:
"""Translate Anthropic-shaped ``system`` to an OpenAI system message.
Accepts a bare string or a list of Anthropic content blocks; returns
@ -772,17 +866,19 @@ def _system_to_openai_message(
if isinstance(system, str):
return {"role": "system", "content": system} if system else None
if isinstance(system, list):
parts = [block.get("text", "") for block in system if isinstance(block, dict) and block.get("type") == "text"]
parts: Final[tuple[str, ...]] = tuple(
block.get("text", "") for block in system if isinstance(block, dict) and block.get("type") == "text"
)
joined: Final = "\n\n".join(part for part in parts if part)
return {"role": "system", "content": joined} if joined else None
return None
def _build_summary_messages(
effective_messages: list[dict[str, object]],
effective_messages: Sequence[dict[str, object]],
prompt: str,
system: str | list[dict[str, object]] | None = None,
) -> list[dict[str, object]]:
) -> Sequence[Mapping[str, object]]:
"""Build the OpenAI-shape message list for the summary call.
The caller's ``system`` prompt is prepended (the default summarization
@ -810,7 +906,7 @@ def _build_summary_messages(
)
openai_messages = stripped
summary_messages: Final[list[dict[str, object]]] = []
summary_messages: Final[list[Mapping[str, object]]] = []
system_message: Final = _system_to_openai_message(system)
if system_message is not None:
summary_messages.append(system_message)
@ -845,35 +941,17 @@ def _append_text_to_content(content: object, extra_text: str) -> object:
if isinstance(content, str):
return f"{content}\n\n{extra_text}"
if isinstance(content, list):
appended: Final[list[object]] = [*content, {"type": "text", "text": extra_text}]
appended: Final[Sequence[object]] = [*map(_as_object, content), {"type": "text", "text": extra_text}]
return appended
return [content, {"type": "text", "text": extra_text}]
class _SummaryCallUserKwarg(TypedDict, total=False):
user: ReadOnly[object]
class _SummaryCallRegionKwarg(TypedDict, total=False):
allowed_model_region: ReadOnly[str]
class _SummaryCallKwargs(TypedDict):
model: ReadOnly[str]
messages: ReadOnly[list[dict[str, object]]]
max_tokens: ReadOnly[int]
timeout: ReadOnly[float]
litellm_metadata: ReadOnly[Mapping[str, object]]
user: NotRequired[ReadOnly[object]]
allowed_model_region: NotRequired[ReadOnly[str]]
async def _call_summary_model(
*,
summary_model: str,
summary_messages: list[dict[str, object]],
summary_messages: Sequence[Mapping[str, object]],
metadata: Mapping[str, object],
llm_router: Any,
llm_router: object,
allowed_model_region: str | None = None,
max_tokens: int = COMPACT_SUMMARY_MAX_TOKENS,
) -> Union["ModelResponse", "CustomStreamWrapper"]:
@ -909,28 +987,37 @@ async def _call_summary_model(
# than from ``litellm_metadata``, so without it the summary tokens would not
# debit the caller's end-user counters.
end_user_id: Final = metadata.get("user_api_key_end_user_id")
user_kwargs: Final = (
_SummaryOptionalKwargs(user=end_user_id)
if isinstance(end_user_id, str) and end_user_id
else _SummaryOptionalKwargs()
)
region_kwargs: Final = (
_SummaryOptionalKwargs(allowed_model_region=allowed_model_region)
if allowed_model_region is not None
else _SummaryOptionalKwargs()
)
call_kwargs: Final[_SummaryCallKwargs] = {
"model": summary_model,
"messages": summary_messages,
"max_tokens": max_tokens,
"timeout": COMPACT_SUMMARY_TIMEOUT_SECONDS,
"litellm_metadata": metadata,
**(_SummaryCallUserKwarg(user=end_user_id) if end_user_id else _SummaryCallUserKwarg()),
**(
_SummaryCallRegionKwarg(allowed_model_region=allowed_model_region)
if allowed_model_region is not None
else _SummaryCallRegionKwarg()
),
**user_kwargs,
**region_kwargs,
}
if llm_router is not None and hasattr(llm_router, "acompletion"):
return await llm_router.acompletion(**call_kwargs)
return await litellm.acompletion(**call_kwargs)
router_acompletion: Final[_SummaryAcompletion | None] = getattr(llm_router, "acompletion", None)
if llm_router is not None and router_acompletion is not None:
return await router_acompletion(messages=summary_messages, **call_kwargs)
return await litellm.acompletion(messages=[*summary_messages], **call_kwargs)
def _extract_response_text(response: Any) -> str | None:
def _extract_response_text(response: object) -> str | None:
try:
choice: Final = response.choices[0]
message: Final = choice.message
choices: Final[Sequence[object] | None] = getattr(response, "choices", None)
if choices is None:
return None
choice: Final = choices[0]
message: Final = getattr(choice, "message", None)
content: Final = getattr(message, "content", None)
if isinstance(content, str):
return content
@ -946,7 +1033,7 @@ def _extract_response_text(response: Any) -> str | None:
def _extract_usage(response: object) -> tuple[int, int]:
usage: Final = getattr(response, "usage", None)
usage: Final[object] = getattr(response, "usage", None)
if usage is None:
return 0, 0
return (

4
litellm/llms/anthropic/experimental_pass_through/responses_adapters/handler.py
View file

@ -160,7 +160,7 @@ class LiteLLMMessagesToResponsesAPIHandler:
top_k: int | None = None,
top_p: float | None = None,
output_format: AnthropicOutputSchema | None = None,
**kwargs,
**kwargs: object,
) -> AnthropicMessagesResponse | AsyncIterator[bytes]:
responses_kwargs: Final = _build_responses_kwargs(
max_tokens=max_tokens,
@ -214,7 +214,7 @@ class LiteLLMMessagesToResponsesAPIHandler:
top_p: float | None = None,
output_format: AnthropicOutputSchema | None = None,
_is_async: bool = False,
**kwargs,
**kwargs: object,
) -> (
AnthropicMessagesResponse
| AsyncIterator[bytes]

48
litellm/llms/base_llm/guardrail_translation/base_translation.py
View file

@ -1,8 +1,11 @@
from abc import ABC, abstractmethod
from collections.abc import Sequence
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any, Final, Optional
if TYPE_CHECKING:
from fastapi import HTTPException
from litellm.integrations.custom_guardrail import (
CustomGuardrail,
ModifyResponseException,
@ -73,6 +76,31 @@ class BaseTranslation(ABC):
return transformed
@staticmethod
def merge_user_api_key_metadata_into_request(
request_data: dict[str, Any], # mutable-ok: proxy hooks share and mutate the request payload dict in place
user_api_key_dict: Optional["UserAPIKeyAuth"],
) -> None:
"""
Add the prefixed ``user_api_key_*`` metadata to the request's resolved
metadata bucket without overwriting existing keys.
Writes must go through ``get_or_create_metadata_bucket``: creating a
``litellm_metadata`` key on a route whose bucket is ``metadata`` (chat
completions) flips the bucket for every later metadata write, and spend
logging never sees those writes (e.g. guardrail_information).
"""
from litellm.litellm_core_utils.core_helpers import (
get_or_create_metadata_bucket,
)
user_metadata: Final = BaseTranslation.transform_user_api_key_dict_to_metadata(user_api_key_dict)
if not user_metadata:
return
_, metadata_bucket = get_or_create_metadata_bucket(request_data)
for key, value in user_metadata.items():
metadata_bucket.setdefault(key, value)
@abstractmethod
async def process_input_messages(
self,
@ -147,6 +175,26 @@ class BaseTranslation(ABC):
"""
return None
def build_stream_error_items(
self,
exc: "HTTPException",
responses_so_far: Sequence[Any] | None = None,
) -> Sequence[Any] | None:
"""
Build the stream items that surface a guardrail HTTPException (a block
with the default exception-on-block config, or a failed scan) after the
response has already started streaming, in this endpoint's wire format.
Called only once chunks have been sent: the HTTP status is gone, so the
failure must travel as an in-stream error frame. ``responses_so_far``
holds the chunks the client has already received, for formats whose
error frame continues the stream (e.g. sequence numbers).
Returns None when the format has no in-stream error frame; the caller
then re-raises ``exc``. Override in endpoint subclasses.
"""
return None
def get_structured_messages(self, data: dict) -> list["AllMessageValues"] | None:
"""
Convert request data to OpenAI-spec structured messages.

9
litellm/llms/bedrock/chat/converse_transformation.py
View file

@ -1631,6 +1631,11 @@ class AmazonConverseConfig(BaseConfig):
bedrock_tool_config["toolChoice"] = tool_choice_values
self._drop_tool_choice_type_conflicting_with_tool_config(additional_request_params)
config_block_entries: Final = tuple(
(config_name, config_class, inference_params.pop(config_name, None))
for config_name, config_class in self.get_config_blocks().items()
)
data: Final[CommonRequestObject] = {
"inferenceConfig": self._transform_inference_params(inference_params=inference_params),
}
@ -1641,9 +1646,7 @@ class AmazonConverseConfig(BaseConfig):
if system_content_blocks:
data["system"] = system_content_blocks
# Handle all config blocks
for config_name, config_class in self.get_config_blocks().items():
config_value = inference_params.pop(config_name, None)
for config_name, config_class, config_value in config_block_entries:
if config_value is not None:
data[config_name] = config_class(**config_value)

18
litellm/llms/bedrock/realtime/handler.py
View file

@ -94,7 +94,7 @@ class BedrockRealtime(BaseAWSLLM):
aws_sts_endpoint: str | None = None,
aws_bedrock_runtime_endpoint: str | None = None,
aws_external_id: str | None = None,
**kwargs,
**kwargs: object,
):
"""
Establish bidirectional streaming connection with Bedrock Nova Sonic.
@ -166,13 +166,16 @@ class BedrockRealtime(BaseAWSLLM):
)
bedrock_client: Final = BedrockRuntimeClient(config=config)
async def open_bidirectional_stream() -> BedrockBidirectionalStream:
return await bedrock_client.invoke_model_with_bidirectional_stream(
InvokeModelWithBidirectionalStreamOperationInput(model_id=model)
)
transformation_config: Final = BedrockRealtimeConfig()
try:
# Initialize the bidirectional stream
bedrock_stream: Final = await bedrock_client.invoke_model_with_bidirectional_stream(
InvokeModelWithBidirectionalStreamOperationInput(model_id=model)
)
bedrock_stream: Final = await open_bidirectional_stream()
verbose_proxy_logger.debug("Bedrock Realtime: Bidirectional stream established")
@ -243,10 +246,11 @@ class BedrockRealtime(BaseAWSLLM):
InvokeModelWithBidirectionalStreamInputChunk,
)
def build_input_chunk(payload: bytes) -> object:
return InvokeModelWithBidirectionalStreamInputChunk(value=BidirectionalInputPayloadPart(bytes_=payload))
async def send_to_bedrock(bedrock_message: str) -> None:
event: Final = InvokeModelWithBidirectionalStreamInputChunk(
value=BidirectionalInputPayloadPart(bytes_=bedrock_message.encode("utf-8"))
)
event: Final = build_input_chunk(bedrock_message.encode("utf-8"))
await bedrock_stream.input_stream.send(event)
verbose_proxy_logger.debug("Bedrock Realtime: Sent to Bedrock: %s", bedrock_message[:200])

4
litellm/llms/chatgpt/chat/streaming_utils.py
View file

@ -49,10 +49,10 @@ class ChatGPTToolCallNormalizer:
def __getattr__(self, name: str) -> object:
return getattr(self._stream, name)
def __iter__(self):
def __iter__(self) -> "ChatGPTToolCallNormalizer":
return self
def __aiter__(self):
def __aiter__(self) -> "ChatGPTToolCallNormalizer":
return self
def __next__(self) -> ModelResponseStream:

31
litellm/llms/compactifai/chat/transformation.py
View file

@ -2,13 +2,16 @@
CompactifAI chat completion transformation
"""
from collections.abc import Mapping, Sequence
from typing import TYPE_CHECKING, Any, Final
import httpx
from typing_extensions import ReadOnly, TypedDict
from litellm.llms.base_llm.chat.transformation import BaseLLMException
from litellm.llms.openai.common_utils import OpenAIError
from litellm.secret_managers.main import get_secret_str
from litellm.types.llms.openai import AllMessageValues
from litellm.types.utils import ModelResponse
from ...openai.chat.gpt_transformation import OpenAIGPTConfig
@ -23,6 +26,18 @@ else:
LiteLLMLoggingObj = Any
class CompactifAIResponseFields(TypedDict, total=False):
"""The chat completion fields of a CompactifAI response body."""
id: ReadOnly[str]
choices: ReadOnly[Sequence[Mapping[str, object]]]
created: ReadOnly[int]
model: ReadOnly[str]
system_fingerprint: ReadOnly[str | None]
usage: ReadOnly[Mapping[str, object]]
object: ReadOnly[str]
class CompactifAIChatConfig(OpenAIGPTConfig):
"""
Configuration class for CompactifAI chat completions.
@ -47,10 +62,10 @@ class CompactifAIChatConfig(OpenAIGPTConfig):
raw_response: httpx.Response,
model_response: ModelResponse,
logging_obj: LiteLLMLoggingObj,
request_data: dict,
messages: list,
optional_params: dict,
litellm_params: dict,
request_data: Mapping[str, object],
messages: Sequence[AllMessageValues],
optional_params: Mapping[str, object],
litellm_params: Mapping[str, object],
encoding: "tiktoken.Encoding | None",
api_key: str | None = None,
json_mode: bool | None = None,
@ -81,14 +96,18 @@ class CompactifAIChatConfig(OpenAIGPTConfig):
message["content"] = tool_calls[0]["function"].get("arguments", "")
message["tool_calls"] = None
returned_response: Final = ModelResponse(**response_json)
response_fields: Final[CompactifAIResponseFields] = response_json
returned_response: Final = ModelResponse(**response_fields)
# Set model name with provider prefix
returned_response.model = f"compactifai/{model}"
return returned_response
def get_error_class(self, error_message: str, status_code: int, headers: dict | httpx.Headers) -> BaseLLMException:
def get_error_class(
self, error_message: str, status_code: int, headers: dict[str, str] | httpx.Headers
) -> BaseLLMException:
"""
Get the appropriate error class for CompactifAI errors.
Since CompactifAI is OpenAI-compatible, we use OpenAI error handling.

179
litellm/llms/custom_httpx/container_handler.py
View file

@ -6,11 +6,12 @@ endpoint defined in endpoints.json, eliminating the need for individual handler
"""
import json
from collections.abc import Coroutine
from collections.abc import Coroutine, Mapping, Sequence
from pathlib import Path
from typing import TYPE_CHECKING, Any, Final
import httpx
from typing_extensions import NotRequired, ReadOnly, TypedDict
import litellm
from litellm.litellm_core_utils.url_utils import encode_url_path_segment
@ -32,26 +33,58 @@ if TYPE_CHECKING:
from litellm.llms.base_llm.containers.transformation import BaseContainerConfig
class EndpointConfig(TypedDict):
"""One endpoint entry of ``litellm/containers/endpoints.json``."""
name: ReadOnly[str]
async_name: ReadOnly[str]
path: ReadOnly[str]
method: ReadOnly[str]
path_params: ReadOnly[Sequence[str]]
query_params: ReadOnly[Sequence[str]]
response_type: ReadOnly[str]
is_multipart: NotRequired[ReadOnly[bool]]
returns_binary: NotRequired[ReadOnly[bool]]
class EndpointsConfig(TypedDict):
"""The parsed ``litellm/containers/endpoints.json`` document."""
endpoints: ReadOnly[Sequence[EndpointConfig]]
class ContainerErrorDetail(TypedDict, total=False):
"""The ``error`` object of a container API error body."""
message: ReadOnly[str]
class ContainerResponseBody(TypedDict, total=False):
"""The fields this handler reads off a container API JSON body."""
error: ReadOnly[ContainerErrorDetail]
_ContainerResponseModel = ContainerFileListResponse | ContainerFileObject | DeleteContainerFileResponse
# Response type mapping
RESPONSE_TYPES: Final[dict[str, type]] = {
RESPONSE_TYPES: Final[Mapping[str, type[_ContainerResponseModel]]] = {
"ContainerFileListResponse": ContainerFileListResponse,
"ContainerFileObject": ContainerFileObject,
"DeleteContainerFileResponse": DeleteContainerFileResponse,
}
ContainerEndpointResponse = (
ContainerFileListResponse | ContainerFileObject | DeleteContainerFileResponse | bytes | dict[str, object]
)
ContainerEndpointResponse = _ContainerResponseModel | bytes | ContainerResponseBody
def _load_endpoints_config() -> dict:
def _load_endpoints_config() -> EndpointsConfig:
"""Load the endpoints configuration from JSON file."""
config_path: Final = Path(__file__).parent.parent.parent / "containers" / "endpoints.json"
with open(config_path) as f:
return json.load(f)
def _get_endpoint_config(endpoint_name: str) -> dict | None:
def _get_endpoint_config(endpoint_name: str) -> EndpointConfig | None:
"""Get config for a specific endpoint by name."""
config: Final = _load_endpoints_config()
for endpoint in config["endpoints"]:
@ -60,10 +93,15 @@ def _get_endpoint_config(endpoint_name: str) -> dict | None:
return None
def _response_model(response_type_name: str) -> type[_ContainerResponseModel] | None:
"""The pydantic model a container endpoint's ``response_type`` names."""
return RESPONSE_TYPES.get(response_type_name)
def _build_url(
api_base: str,
path_template: str,
path_params: dict[str, str],
path_params: Mapping[str, object],
) -> str:
"""Build the full URL by substituting path parameters.
@ -93,16 +131,12 @@ def _build_url(
def _build_query_params(
query_param_names: list,
kwargs: dict[str, Any],
) -> dict[str, str]:
query_param_names: Sequence[str],
kwargs: Mapping[str, object],
) -> dict[str, object]:
"""Build query parameters from kwargs."""
params: Final = {}
for param_name in query_param_names:
value = kwargs.get(param_name)
if value is not None:
params[param_name] = str(value) if not isinstance(value, str) else value
return params
supplied: Final = ((param_name, kwargs.get(param_name)) for param_name in query_param_names)
return {name: value if isinstance(value, str) else str(value) for name, value in supplied if value is not None}
def _error_message_from_response(response: httpx.Response) -> str:
@ -136,24 +170,24 @@ def _transform_response(
if returns_binary:
return response.content
response_json: Final = response.json()
response_json: Final[ContainerResponseBody] = response.json()
if "error" in response_json:
raise BaseLLMException(
status_code=response.status_code,
message=response_json.get("error", {}).get("message", str(response_json)),
message=response_json["error"].get("message", str(response_json)),
headers=dict(response.headers),
)
response_type: Final = RESPONSE_TYPES.get(response_type_name)
response_type: Final = _response_model(response_type_name)
if response_type:
return response_type(**response_json)
return response_type.model_validate(response_json)
return response_json
def _prepare_multipart_file_upload(
file: Any,
headers: dict[str, Any],
) -> tuple:
headers: dict[str, object],
) -> tuple[dict[str, tuple[str, bytes, str]], dict[str, object]]:
"""
Prepare file and headers for multipart upload.
@ -178,6 +212,52 @@ def _prepare_multipart_file_upload(
return files, headers_copy
def _request_headers(
container_provider_config: "BaseContainerConfig",
extra_headers: dict[str, object] | None,
litellm_params: GenericLiteLLMParams,
) -> dict[str, object]:
"""The provider auth headers for a container request."""
return container_provider_config.validate_environment(
headers=extra_headers or {},
api_key=litellm_params.get("api_key", None),
)
def _request_api_base(
container_provider_config: "BaseContainerConfig",
litellm_params: GenericLiteLLMParams,
) -> str:
"""The provider base URL for a container request."""
return container_provider_config.get_complete_url(
api_base=litellm_params.get("api_base", None),
litellm_params=dict(litellm_params),
)
def _sync_http_client(
client: HTTPHandler | AsyncHTTPHandler | None,
litellm_params: GenericLiteLLMParams,
) -> HTTPHandler:
"""The sync HTTP client for a container request, reusing the caller's when usable."""
if client is None or not isinstance(client, HTTPHandler):
return _get_httpx_client(params={"ssl_verify": litellm_params.get("ssl_verify", None)})
return client
def _async_http_client(
client: HTTPHandler | AsyncHTTPHandler | None,
litellm_params: GenericLiteLLMParams,
) -> AsyncHTTPHandler:
"""The async HTTP client for a container request, reusing the caller's when usable."""
if client is None or not isinstance(client, AsyncHTTPHandler):
return get_async_httpx_client(
llm_provider=litellm.LlmProviders.OPENAI,
params={"ssl_verify": litellm_params.get("ssl_verify", None)},
)
return client
class GenericContainerHandler:
"""
Generic handler for container file API endpoints.
@ -192,13 +272,13 @@ class GenericContainerHandler:
container_provider_config: "BaseContainerConfig",
litellm_params: GenericLiteLLMParams,
logging_obj: "LiteLLMLoggingObj",
extra_headers: dict[str, Any] | None = None,
extra_query: dict[str, Any] | None = None,
extra_headers: dict[str, object] | None = None,
extra_query: dict[str, object] | None = None,
timeout: float | httpx.Timeout = 600,
_is_async: bool = False,
client: HTTPHandler | AsyncHTTPHandler | None = None,
**kwargs,
) -> Any | Coroutine[Any, Any, Any]:
**kwargs: object,
) -> Any | Coroutine[object, object, Any]:
"""
Generic handler for any container file endpoint.
@ -245,11 +325,11 @@ class GenericContainerHandler:
container_provider_config: "BaseContainerConfig",
litellm_params: GenericLiteLLMParams,
logging_obj: "LiteLLMLoggingObj",
extra_headers: dict[str, Any] | None = None,
extra_query: dict[str, Any] | None = None,
extra_headers: dict[str, object] | None = None,
extra_query: dict[str, object] | None = None,
timeout: float | httpx.Timeout = 600,
client: HTTPHandler | AsyncHTTPHandler | None = None,
**kwargs,
**kwargs: object,
) -> Any:
"""Synchronous request handler."""
endpoint_config: Final = _get_endpoint_config(endpoint_name)
@ -257,23 +337,14 @@ class GenericContainerHandler:
raise ValueError(f"Unknown endpoint: {endpoint_name}")
# Get HTTP client
if client is None or not isinstance(client, HTTPHandler):
http_client = _get_httpx_client(params={"ssl_verify": litellm_params.get("ssl_verify", None)})
else:
http_client = client
http_client: Final = _sync_http_client(client, litellm_params)
# Build request
headers = container_provider_config.validate_environment(
headers=extra_headers or {},
api_key=litellm_params.get("api_key", None),
)
headers = _request_headers(container_provider_config, extra_headers, litellm_params)
if extra_headers:
headers.update(extra_headers)
api_base: Final = container_provider_config.get_complete_url(
api_base=litellm_params.get("api_base", None),
litellm_params=dict(litellm_params),
)
api_base: Final = _request_api_base(container_provider_config, litellm_params)
# Build URL with path params
path_params: Final = {p: kwargs.get(p, "") for p in endpoint_config.get("path_params", [])}
@ -334,11 +405,11 @@ class GenericContainerHandler:
container_provider_config: "BaseContainerConfig",
litellm_params: GenericLiteLLMParams,
logging_obj: "LiteLLMLoggingObj",
extra_headers: dict[str, Any] | None = None,
extra_query: dict[str, Any] | None = None,
extra_headers: dict[str, object] | None = None,
extra_query: dict[str, object] | None = None,
timeout: float | httpx.Timeout = 600,
client: HTTPHandler | AsyncHTTPHandler | None = None,
**kwargs,
**kwargs: object,
) -> Any:
"""Asynchronous request handler."""
endpoint_config: Final = _get_endpoint_config(endpoint_name)
@ -346,26 +417,14 @@ class GenericContainerHandler:
raise ValueError(f"Unknown endpoint: {endpoint_name}")
# Get HTTP client
if client is None or not isinstance(client, AsyncHTTPHandler):
http_client = get_async_httpx_client(
llm_provider=litellm.LlmProviders.OPENAI,
params={"ssl_verify": litellm_params.get("ssl_verify", None)},
)
else:
http_client = client
http_client: Final = _async_http_client(client, litellm_params)
# Build request
headers = container_provider_config.validate_environment(
headers=extra_headers or {},
api_key=litellm_params.get("api_key", None),
)
headers = _request_headers(container_provider_config, extra_headers, litellm_params)
if extra_headers:
headers.update(extra_headers)
api_base: Final = container_provider_config.get_complete_url(
api_base=litellm_params.get("api_base", None),
litellm_params=dict(litellm_params),
)
api_base: Final = _request_api_base(container_provider_config, litellm_params)
# Build URL with path params
path_params: Final = {p: kwargs.get(p, "") for p in endpoint_config.get("path_params", [])}

6
litellm/llms/custom_httpx/http_handler.py
View file

@ -9,7 +9,7 @@ import threading
import time
from collections.abc import AsyncIterable, Callable, Iterable, Mapping
from http.cookiejar import CookieJar, DefaultCookiePolicy
from typing import TYPE_CHECKING, Any, ClassVar, Final, Optional, TypeAlias, TypedDict
from typing import TYPE_CHECKING, Any, ClassVar, Final, NoReturn, Optional, TypeAlias, TypedDict
import certifi
import httpx
@ -447,7 +447,7 @@ def _safe_read_response(response: httpx.Response, timeout: float | None = None)
return b""
def _raise_masked_sync_error(e: httpx.HTTPStatusError, stream: bool) -> None:
def _raise_masked_sync_error(e: httpx.HTTPStatusError, stream: bool) -> NoReturn:
"""Raise a MaskedHTTPStatusError for sync HTTP handlers."""
if stream:
try:
@ -467,7 +467,7 @@ def _raise_masked_sync_error(e: httpx.HTTPStatusError, stream: bool) -> None:
raise MaskedHTTPStatusError(e, message=_text, text=_text) from None
async def _raise_masked_async_error(e: httpx.HTTPStatusError, stream: bool) -> None:
async def _raise_masked_async_error(e: httpx.HTTPStatusError, stream: bool) -> NoReturn:
"""Raise a MaskedHTTPStatusError for async HTTP handlers."""
if stream:
try:

35
litellm/llms/e2b/sandbox/transformation.py
View file

@ -8,7 +8,7 @@ Talks to e2b's REST API directly over httpx (no e2b SDK dependency):
"""
import json
from typing import Final, cast
from typing import Final
import httpx
@ -68,13 +68,10 @@ class E2BSandboxConfig(BaseSandboxConfig):
if metadata:
body["metadata"] = metadata
response: Final = cast(
httpx.Response,
await self._http(client).post(
url=f"{base}/sandboxes",
headers={"X-API-Key": key, "Content-Type": "application/json"},
json=body,
),
response: Final = await self._http(client).post(
url=f"{base}/sandboxes",
headers={"X-API-Key": key, "Content-Type": "application/json"},
json=body,
)
data: Final = response.json()
@ -117,14 +114,11 @@ class E2BSandboxConfig(BaseSandboxConfig):
headers["E2B-Traffic-Access-Token"] = traffic_token
url: Final = f"https://{JUPYTER_PORT}-{handle.id}.{handle.domain}/execute"
response: Final = cast(
httpx.Response,
await self._http(client).post(
url=url,
headers=headers,
json={"code": code, "context_id": None, "env_vars": env_vars},
stream=True,
),
response: Final = await self._http(client).post(
url=url,
headers=headers,
json={"code": code, "context_id": None, "env_vars": env_vars},
stream=True,
)
lines: Final = await self._read_capped_lines(response)
return self._parse_lines(lines)
@ -142,12 +136,9 @@ class E2BSandboxConfig(BaseSandboxConfig):
key: Final = api_key or handle._hidden_params.get("api_key") or self.validate_environment()
base: Final = api_base or handle._hidden_params.get("api_base") or E2B_API_BASE
try:
response: Final = cast(
httpx.Response,
await self._http(client).delete(
url=f"{base}/sandboxes/{handle.id}",
headers={"X-API-Key": key},
),
response: Final = await self._http(client).delete(
url=f"{base}/sandboxes/{handle.id}",
headers={"X-API-Key": key},
)
except httpx.HTTPStatusError as e:
if e.response.status_code == 404:

45
litellm/llms/gdc/chat/transformation.py
View file

@ -6,11 +6,25 @@ import json
import os
import re
import threading
from typing import Any, Final
from collections.abc import Callable
from typing import Any, Final, Protocol
from urllib.parse import urlsplit
import litellm
from litellm.llms.openai_like.chat.transformation import OpenAILikeChatConfig
from litellm.types.llms.openai import AllMessageValues
class _GDCHAudienceCredentials(Protocol):
"""A GDCH service account credential already bound to an audience, ready to mint a bearer token."""
@property
def valid(self) -> bool: ...
@property
def token(self) -> str: ...
def refresh(self, request: object) -> None: ...
class GDCGeminiConfig(OpenAILikeChatConfig):
@ -21,7 +35,7 @@ class GDCGeminiConfig(OpenAILikeChatConfig):
def __init__(self, **kwargs: Any) -> None:
super().__init__(**kwargs)
self._creds_lock = threading.Lock()
self._gdch_creds_cache: dict = {}
self._gdch_creds_cache: dict[tuple[str, str], _GDCHAudienceCredentials] = {}
def get_supported_openai_params(self, model: str) -> list:
return [
@ -110,7 +124,7 @@ class GDCGeminiConfig(OpenAILikeChatConfig):
return f"{api_base}/v1/projects/{project}/locations/{location}/chat/completions"
def _read_env_bool(self, val: Any, env_var: str, default: bool = True) -> bool | str:
def _read_env_bool(self, val: bool | str | None, env_var: str, default: bool = True) -> bool | str:
def _parse(s: str) -> bool | str:
cleaned: Final = s.strip().lower()
if cleaned in ("false", "0", "no", "off"):
@ -129,7 +143,7 @@ class GDCGeminiConfig(OpenAILikeChatConfig):
return default
return _parse(_env_val)
def _fetch_auth(self, gdch_creds: Any, ssl_verify: bool | str) -> None:
def _fetch_auth(self, gdch_creds: _GDCHAudienceCredentials, ssl_verify: bool | str) -> None:
import requests
from google.auth.transport import requests as auth_requests
@ -138,13 +152,24 @@ class GDCGeminiConfig(OpenAILikeChatConfig):
auth_request: Final = auth_requests.Request(session=auth_session)
gdch_creds.refresh(auth_request)
def _cached_fetch_token(self, creds: Any, audience: str, ssl_verify: bool | str, api_key: str | None = None) -> str:
def _with_gdch_audience(self, creds: object, audience: str) -> _GDCHAudienceCredentials:
"""The credential rebound to ``audience``, which GDCH requires before a token refresh."""
bind_audience: Final[Callable[[str], _GDCHAudienceCredentials] | None] = getattr(
creds, "with_gdch_audience", None
)
if bind_audience is None:
raise AttributeError("GDC credentials must expose with_gdch_audience to be bound to a request audience")
return bind_audience(audience)
def _cached_fetch_token(
self, creds: object, audience: str, ssl_verify: bool | str, api_key: str | None = None
) -> str:
# Key cache by both audience and credential identity to prevent cross-caller contamination
cache_key: Final = (audience.rstrip("/"), api_key or str(id(creds)))
with self._creds_lock:
if cache_key not in self._gdch_creds_cache:
self._gdch_creds_cache[cache_key] = creds.with_gdch_audience(audience.rstrip("/"))
self._gdch_creds_cache[cache_key] = self._with_gdch_audience(creds, audience.rstrip("/"))
gdch_creds: Final = self._gdch_creds_cache[cache_key]
@ -155,7 +180,7 @@ class GDCGeminiConfig(OpenAILikeChatConfig):
return token
def _load_creds_from_key(self, api_key: str) -> tuple[Any, bool]:
def _load_creds_from_key(self, api_key: str) -> tuple[object | None, bool]:
import google.auth
try:
@ -175,7 +200,7 @@ class GDCGeminiConfig(OpenAILikeChatConfig):
self,
headers: dict,
model: str,
messages: list[Any],
messages: list[AllMessageValues],
optional_params: dict,
litellm_params: dict,
api_key: str | None = None,
@ -230,7 +255,7 @@ class GDCGeminiConfig(OpenAILikeChatConfig):
if self._read_env_bool(litellm_params.get("gdc_token_caching"), "GDC_TOKEN_CACHING", default=False):
token = self._cached_fetch_token(creds, audience, ssl_verify, api_key)
else:
gdch_creds: Final = creds.with_gdch_audience(audience)
gdch_creds: Final = self._with_gdch_audience(creds, audience)
self._fetch_auth(gdch_creds, ssl_verify)
token = gdch_creds.token
headers["Authorization"] = f"Bearer {token}"
@ -252,7 +277,7 @@ class GDCGeminiConfig(OpenAILikeChatConfig):
def transform_request(
self,
model: str,
messages: list[Any],
messages: list[AllMessageValues],
optional_params: dict,
litellm_params: dict,
headers: dict,

30
litellm/llms/infinity/rerank/transformation.py
View file

@ -4,10 +4,11 @@ Transformation logic from Cohere's /v1/rerank format to Infinity's `/v1/rerank`
Why separate file? Make it easy to see how transformation works
"""
from collections.abc import Mapping
from collections.abc import Mapping, Sequence
from typing import Final
import httpx
from typing_extensions import NotRequired, ReadOnly, TypedDict
import litellm
from litellm._uuid import uuid
@ -26,6 +27,31 @@ from litellm.types.rerank import (
from ..common_utils import InfinityError
class _InfinityRerankUsage(TypedDict, extra_items=ReadOnly[int]):
"""The token counters Infinity reports in the ``usage`` block of a rerank response."""
class _InfinityRerankResult(TypedDict):
"""One scored document in an Infinity ``/v1/rerank`` response."""
index: ReadOnly[int]
relevance_score: ReadOnly[float]
document: ReadOnly[str]
class _InfinityRerankResponse(TypedDict):
"""The JSON body returned by Infinity's ``/v1/rerank`` endpoint."""
id: ReadOnly[NotRequired[str]]
usage: ReadOnly[NotRequired[_InfinityRerankUsage]]
results: ReadOnly[Sequence[_InfinityRerankResult]]
def _parse_rerank_response(raw_response: httpx.Response) -> _InfinityRerankResponse:
"""Read the untyped JSON body of an Infinity rerank response."""
return raw_response.json()
class InfinityRerankConfig(CohereRerankConfig):
def get_complete_url(
self,
@ -82,7 +108,7 @@ class InfinityRerankConfig(CohereRerankConfig):
No transformation required, Infinity follows Cohere API response format
"""
try:
raw_response_json: Final = raw_response.json()
raw_response_json: Final = _parse_rerank_response(raw_response)
except Exception:
raise InfinityError(message=raw_response.text, status_code=raw_response.status_code)

76
litellm/llms/litellm_proxy/skills/code_execution.py
View file

@ -13,12 +13,57 @@ Generated files are returned directly in the response - no separate storage need
import base64
import json
from collections.abc import Sequence
from enum import Enum
from typing import Any, Final
from typing import Any, Final, Protocol
from typing_extensions import NotRequired, ReadOnly, TypedDict
from litellm._logging import verbose_logger
class _ToolCallFunction(Protocol):
"""Function payload of an assistant tool call."""
name: str | None
arguments: str
class _ToolCall(Protocol):
"""Tool call requested by the assistant on a chat completion choice."""
id: str
function: _ToolCallFunction
class _AssistantMessage(Protocol):
"""Assistant message carried by a chat completion choice."""
content: str | None
tool_calls: Sequence[_ToolCall] | None
class _CompletionChoice(Protocol):
"""Single choice of a chat completion response."""
finish_reason: str
message: _AssistantMessage
class _SandboxFile(TypedDict):
"""File generated inside the sandbox during a code execution run."""
name: ReadOnly[str]
mime_type: ReadOnly[str]
content_base64: ReadOnly[str]
class _CodeExecutionArguments(TypedDict):
"""Arguments the model passes to the `litellm_code_execution` tool."""
code: NotRequired[ReadOnly[str]]
class LiteLLMInternalTools(str, Enum):
"""
Enum for internal LiteLLM tools that are injected into requests.
@ -30,7 +75,7 @@ class LiteLLMInternalTools(str, Enum):
CODE_EXECUTION = "litellm_code_execution"
def get_litellm_code_execution_tool() -> dict[str, Any]:
def get_litellm_code_execution_tool() -> dict[str, object]:
"""
Returns the litellm_code_execution tool definition in OpenAI format.
@ -51,7 +96,7 @@ def get_litellm_code_execution_tool() -> dict[str, Any]:
}
def get_litellm_code_execution_tool_anthropic() -> dict[str, Any]:
def get_litellm_code_execution_tool_anthropic() -> dict[str, object]:
"""
Returns the litellm_code_execution tool definition in Anthropic/messages API format.
@ -103,7 +148,7 @@ class CodeExecutionHandler:
skill_files: dict[str, bytes],
skill_id: str | None = None,
**kwargs,
) -> dict[str, Any]:
) -> dict[str, object]:
"""
Execute an LLM call with automatic code execution handling.
@ -134,8 +179,8 @@ class CodeExecutionHandler:
)
current_messages: Final = list(messages)
generated_files: Final[list[dict[str, Any]]] = [] # Files returned directly
execution_results: Final[list[dict]] = []
generated_files: Final[list[dict[str, object]]] = [] # Files returned directly
execution_results: Final[list[dict[str, object]]] = []
executor: Final = SkillsSandboxExecutor(timeout=self.sandbox_timeout)
response: Any = None # Initialize to avoid possibly unbound error
@ -151,11 +196,12 @@ class CodeExecutionHandler:
**kwargs,
)
assistant_message = response.choices[0].message
stop_reason = response.choices[0].finish_reason
choice: _CompletionChoice = response.choices[0]
assistant_message = choice.message
stop_reason: str = choice.finish_reason
# Build assistant message for conversation history
assistant_msg_dict: dict[str, Any] = {
assistant_msg_dict: dict[str, object] = {
"role": "assistant",
"content": assistant_message.content,
}
@ -190,8 +236,8 @@ class CodeExecutionHandler:
if tool_name == LiteLLMInternalTools.CODE_EXECUTION.value:
# Execute code in sandbox
try:
args = json.loads(tool_call.function.arguments)
code = args.get("code", "")
args: _CodeExecutionArguments = json.loads(tool_call.function.arguments)
code: str = args.get("code", "")
verbose_logger.debug("CodeExecutionHandler: Executing code (%s chars)", len(code))
@ -202,13 +248,15 @@ class CodeExecutionHandler:
verbose_logger.debug("CodeExecutionHandler: Execution result: %s", exec_result)
sandbox_files: Sequence[_SandboxFile] = exec_result["files"]
execution_results.append(
{
"iteration": iteration,
"success": exec_result["success"],
"output": exec_result["output"],
"error": exec_result["error"],
"files": [f["name"] for f in exec_result["files"]],
"files": [f["name"] for f in sandbox_files],
}
)
@ -216,9 +264,9 @@ class CodeExecutionHandler:
tool_result = exec_result["output"] or ""
# Collect generated files (returned directly, no storage)
if exec_result["files"]:
if sandbox_files:
tool_result += "\n\nGenerated files:"
for f in exec_result["files"]:
for f in sandbox_files:
file_content = base64.b64decode(f["content_base64"])
# Add to generated files list (returned in response)
generated_files.append(

36
litellm/llms/ollama/completion/handler.py
View file

@ -4,16 +4,32 @@ Ollama /chat/completion calls handled in llm_http_handler.py
[TODO]: migrate embeddings to a base handler as well.
"""
from typing import Any, Final
from collections.abc import Mapping, Sequence
from typing import Any, Final, Protocol, TypedDict
from typing_extensions import NotRequired, ReadOnly
import litellm
from litellm.types.utils import EmbeddingResponse
class TokenEncoder(Protocol):
"""The tokenizer surface used to estimate prompt tokens."""
def encode(self, text: str, /) -> Sequence[int]: ...
class OllamaEmbeddingResponse(TypedDict):
"""Body of an Ollama ``/api/embed`` response."""
embeddings: ReadOnly[list[list[float]]]
prompt_eval_count: ReadOnly[NotRequired[int]]
def _prepare_ollama_embedding_payload(
model: str, prompts: list[str], optional_params: dict[str, Any]
) -> dict[str, Any]:
data: Final[dict[str, Any]] = {"model": model, "input": prompts}
model: str, prompts: list[str], optional_params: Mapping[str, object]
) -> dict[str, object]:
data: Final[dict[str, object]] = {"model": model, "input": prompts}
special_optional_params: Final = ["truncate", "options", "keep_alive", "dimensions"]
for k, v in optional_params.items():
@ -27,12 +43,12 @@ def _prepare_ollama_embedding_payload(
def _process_ollama_embedding_response(
response_json: dict,
response_json: OllamaEmbeddingResponse,
prompts: list[str],
model: str,
model_response: EmbeddingResponse,
logging_obj: Any,
encoding: Any,
encoding: TokenEncoder | None,
) -> EmbeddingResponse:
output_data: Final = []
embeddings: Final[list[list[float]]] = response_json["embeddings"]
@ -72,7 +88,7 @@ async def ollama_aembeddings(
model_response: EmbeddingResponse,
optional_params: dict,
logging_obj: Any,
encoding: Any,
encoding: TokenEncoder | None,
):
if not api_base.endswith("/api/embed"):
api_base += "/api/embed"
@ -80,7 +96,7 @@ async def ollama_aembeddings(
data: Final = _prepare_ollama_embedding_payload(model, prompts, optional_params)
response: Final = await litellm.module_level_aclient.post(url=api_base, json=data)
response_json: Final = response.json()
response_json: Final[OllamaEmbeddingResponse] = response.json()
return _process_ollama_embedding_response(
response_json=response_json,
@ -99,7 +115,7 @@ def ollama_embeddings(
optional_params: dict,
model_response: EmbeddingResponse,
logging_obj: Any,
encoding: Any = None,
encoding: TokenEncoder | None = None,
):
if not api_base.endswith("/api/embed"):
api_base += "/api/embed"
@ -107,7 +123,7 @@ def ollama_embeddings(
data: Final = _prepare_ollama_embedding_payload(model, prompts, optional_params)
response: Final = litellm.module_level_client.post(url=api_base, json=data)
response_json: Final = response.json()
response_json: Final[OllamaEmbeddingResponse] = response.json()
return _process_ollama_embedding_response(
response_json=response_json,

32
litellm/llms/openai/chat/guardrail_translation/handler.py
View file

@ -14,6 +14,7 @@ Pattern Overview:
This pattern can be replicated for other message formats (e.g., Anthropic).
"""
from collections.abc import Sequence
from typing import TYPE_CHECKING, Any, Final, Union, cast
import litellm
@ -46,6 +47,8 @@ from litellm.types.utils import (
)
if TYPE_CHECKING:
from fastapi import HTTPException
from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
@ -382,11 +385,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
if "response" not in request_data:
request_data["response"] = response
# Add user API key metadata with prefixed keys
if "litellm_metadata" not in request_data:
user_metadata: Final = self.transform_user_api_key_dict_to_metadata(user_api_key_dict)
if user_metadata:
request_data["litellm_metadata"] = user_metadata
self.merge_user_api_key_metadata_into_request(request_data, user_api_key_dict)
inputs: Final = GenericGuardrailAPIInputs(texts=texts_to_check)
if images_to_check:
@ -555,11 +554,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
if "responses" not in request_data:
request_data["responses"] = responses_so_far
# Add user API key metadata with prefixed keys
if "litellm_metadata" not in request_data:
user_metadata: Final = self.transform_user_api_key_dict_to_metadata(user_api_key_dict)
if user_metadata:
request_data["litellm_metadata"] = user_metadata
self.merge_user_api_key_metadata_into_request(request_data, user_api_key_dict)
inputs: Final = GenericGuardrailAPIInputs(texts=texts_to_check)
if images_to_check:
@ -591,6 +586,18 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
return responses_so_far
def build_stream_error_items(
self,
exc: "HTTPException",
responses_so_far: Sequence[Any] | None = None,
) -> Sequence[Any] | None:
import json
from litellm.proxy.common_request_processing import sse_error_payload
_, error_obj = sse_error_payload(exc)
return (f'data: {{"error": {json.dumps(error_obj)}}}\n\n'.encode(),)
@staticmethod
def _accumulate_string_content_by_choice_index(
responses_so_far: list["ModelResponseStream"],
@ -653,10 +660,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
request_data = {"responses": responses_so_far}
elif "responses" not in request_data:
request_data["responses"] = responses_so_far
if "litellm_metadata" not in request_data:
user_metadata: Final = self.transform_user_api_key_dict_to_metadata(user_api_key_dict)
if user_metadata:
request_data["litellm_metadata"] = user_metadata
self.merge_user_api_key_metadata_into_request(request_data, user_api_key_dict)
inputs: Final = GenericGuardrailAPIInputs(texts=texts_to_check)
if responses_so_far and getattr(responses_so_far[0], "model", None):

72
litellm/llms/openai/containers/transformation.py
View file

@ -1,6 +1,8 @@
from typing import TYPE_CHECKING, Any, Final
from collections.abc import Mapping
from typing import TYPE_CHECKING, Any, Final, Literal
import httpx
from typing_extensions import ReadOnly, TypedDict
import litellm
from litellm.litellm_core_utils.llm_cost_calc.tool_call_cost_tracking import (
@ -11,9 +13,11 @@ from litellm.secret_managers.main import get_secret_str
from litellm.types.containers.main import (
ContainerCreateOptionalRequestParams,
ContainerFileListResponse,
ContainerFileObject,
ContainerListResponse,
ContainerObject,
DeleteContainerResult,
ExpiresAfter,
)
from litellm.types.router import GenericLiteLLMParams
@ -32,6 +36,46 @@ else:
BaseLLMException = Any
class OpenAIContainerPayload(TypedDict):
"""The JSON body OpenAI returns for a single container."""
id: ReadOnly[str]
object: ReadOnly[Literal["container"]]
created_at: ReadOnly[int]
status: ReadOnly[str]
expires_after: ReadOnly[ExpiresAfter | None]
last_active_at: ReadOnly[int | None]
name: ReadOnly[str | None]
class OpenAIContainerListPayload(TypedDict):
"""The JSON body OpenAI returns for a page of containers."""
object: ReadOnly[Literal["list"]]
data: ReadOnly[list[ContainerObject]]
first_id: ReadOnly[str | None]
last_id: ReadOnly[str | None]
has_more: ReadOnly[bool]
class OpenAIContainerDeletedPayload(TypedDict):
"""The JSON body OpenAI returns for a deleted container."""
id: ReadOnly[str]
object: ReadOnly[Literal["container.deleted"]]
deleted: ReadOnly[bool]
class OpenAIContainerFileListPayload(TypedDict):
"""The JSON body OpenAI returns for a page of container files."""
object: ReadOnly[Literal["list"]]
data: ReadOnly[list[ContainerFileObject]]
first_id: ReadOnly[str | None]
last_id: ReadOnly[str | None]
has_more: ReadOnly[bool]
class OpenAIContainerConfig(BaseContainerConfig):
"""Configuration class for OpenAI container API."""
@ -87,7 +131,7 @@ class OpenAIContainerConfig(BaseContainerConfig):
def transform_container_create_request(
self,
name: str,
container_create_optional_request_params: dict,
container_create_optional_request_params: Mapping[str, object],
litellm_params: GenericLiteLLMParams,
headers: dict,
) -> dict:
@ -111,7 +155,7 @@ class OpenAIContainerConfig(BaseContainerConfig):
logging_obj: LiteLLMLoggingObj,
) -> ContainerObject:
"""Transform the OpenAI container creation response."""
response_data: Final = raw_response.json()
response_data: Final[OpenAIContainerPayload] = raw_response.json()
# Transform the response data
container_obj: Final = ContainerObject(**response_data)
@ -140,7 +184,7 @@ class OpenAIContainerConfig(BaseContainerConfig):
after: str | None = None,
limit: int | None = None,
order: str | None = None,
extra_query: dict[str, Any] | None = None,
extra_query: Mapping[str, object] | None = None,
) -> tuple[str, dict]:
"""Transform the container list request for OpenAI API.
@ -151,7 +195,7 @@ class OpenAIContainerConfig(BaseContainerConfig):
url: Final = api_base
# Prepare query parameters
params: Final = {}
params: Final[dict[str, object]] = {}
if after is not None:
params["after"] = after
if limit is not None:
@ -171,7 +215,7 @@ class OpenAIContainerConfig(BaseContainerConfig):
logging_obj: LiteLLMLoggingObj,
) -> ContainerListResponse:
"""Transform the OpenAI container list response."""
response_data: Final = raw_response.json()
response_data: Final[OpenAIContainerListPayload] = raw_response.json()
# Transform the response data
container_list: Final = ContainerListResponse(**response_data)
@ -191,7 +235,7 @@ class OpenAIContainerConfig(BaseContainerConfig):
url: Final = join_container_api_base_path(api_base, f"/{encoded_container_id}")
# No additional data needed for GET request
data: Final[dict[str, Any]] = {}
data: Final[dict[str, object]] = {}
return url, data
@ -201,7 +245,7 @@ class OpenAIContainerConfig(BaseContainerConfig):
logging_obj: LiteLLMLoggingObj,
) -> ContainerObject:
"""Transform the OpenAI container retrieve response."""
response_data: Final = raw_response.json()
response_data: Final[OpenAIContainerPayload] = raw_response.json()
# Transform the response data
container_obj: Final = ContainerObject(**response_data)
@ -224,7 +268,7 @@ class OpenAIContainerConfig(BaseContainerConfig):
url: Final = join_container_api_base_path(api_base, f"/{encoded_container_id}")
# No data needed for DELETE request
data: Final[dict[str, Any]] = {}
data: Final[dict[str, object]] = {}
return url, data
@ -234,7 +278,7 @@ class OpenAIContainerConfig(BaseContainerConfig):
logging_obj: LiteLLMLoggingObj,
) -> DeleteContainerResult:
"""Transform the OpenAI container delete response."""
response_data: Final = raw_response.json()
response_data: Final[OpenAIContainerDeletedPayload] = raw_response.json()
# Transform the response data
delete_result: Final = DeleteContainerResult(**response_data)
@ -250,7 +294,7 @@ class OpenAIContainerConfig(BaseContainerConfig):
after: str | None = None,
limit: int | None = None,
order: str | None = None,
extra_query: dict[str, Any] | None = None,
extra_query: Mapping[str, object] | None = None,
) -> tuple[str, dict]:
"""Transform the container file list request for OpenAI API.
@ -262,7 +306,7 @@ class OpenAIContainerConfig(BaseContainerConfig):
url: Final = join_container_api_base_path(api_base, f"/{encoded_container_id}/files")
# Prepare query parameters
params: Final[dict[str, Any]] = {}
params: Final[dict[str, object]] = {}
if after is not None:
params["after"] = after
if limit is not None:
@ -282,7 +326,7 @@ class OpenAIContainerConfig(BaseContainerConfig):
logging_obj: LiteLLMLoggingObj,
) -> ContainerFileListResponse:
"""Transform the OpenAI container file list response."""
response_data: Final = raw_response.json()
response_data: Final[OpenAIContainerFileListPayload] = raw_response.json()
# Transform the response data
file_list: Final = ContainerFileListResponse(**response_data)
@ -308,7 +352,7 @@ class OpenAIContainerConfig(BaseContainerConfig):
url: Final = join_container_api_base_path(api_base, f"/{encoded_container_id}/files/{encoded_file_id}/content")
# No query parameters needed
params: Final[dict[str, Any]] = {}
params: Final[dict[str, object]] = {}
return url, params

35
litellm/llms/openai/responses/guardrail_translation/handler.py
View file

@ -48,6 +48,8 @@ from litellm.types.llms.openai import (
AllMessageValues,
ChatCompletionToolCallChunk,
ChatCompletionToolParam,
ErrorEvent,
ErrorEventError,
OpenAIMcpServerTool,
ResponsesAPIStreamEvents,
)
@ -59,6 +61,8 @@ from litellm.types.responses.main import (
from litellm.types.utils import GenericGuardrailAPIInputs
if TYPE_CHECKING:
from fastapi import HTTPException
from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.proxy._types import UserAPIKeyAuth
@ -80,6 +84,14 @@ class ResponsesStreamChunk(TypedDict, total=False):
text: ReadOnly[str]
def _next_stream_sequence_number(responses_so_far: Sequence[Any] | None) -> int:
sequence_numbers: Final = (
item.get("sequence_number") if isinstance(item, dict) else getattr(item, "sequence_number", None)
for item in reversed(responses_so_far or ())
)
return next((n + 1 for n in sequence_numbers if isinstance(n, int)), 0)
class OpenAIResponsesHandler(BaseTranslation):
"""
Handler for processing OpenAI Responses API with guardrails.
@ -620,6 +632,29 @@ class OpenAIResponsesHandler(BaseTranslation):
}
return responses_so_far[-1].get("type") in terminal_types
def build_stream_error_items(
self,
exc: "HTTPException",
responses_so_far: Sequence[Any] | None = None,
) -> Sequence[Any] | None:
from litellm.proxy.common_request_processing import (
serialize_http_exception_detail,
)
message, _ = serialize_http_exception_detail(exc.detail)
return (
ErrorEvent(
type=ResponsesAPIStreamEvents.ERROR,
sequence_number=_next_stream_sequence_number(responses_so_far),
error=ErrorEventError(
type="guardrail_error",
code=str(exc.status_code),
message=message,
param=None,
),
),
)
def get_streaming_string_so_far(self, responses_so_far: Sequence[ResponsesStreamChunk]) -> str:
"""
Get the string so far from the responses so far.

57
litellm/llms/opensandbox/sandbox/transformation.py
View file

@ -1,7 +1,7 @@
import asyncio
import json
import time
from typing import Final, cast
from typing import Final
import httpx
@ -86,13 +86,10 @@ class OpenSandboxSandboxConfig(BaseSandboxConfig):
secure_access=secure_access,
)
response: Final = cast(
httpx.Response,
await self._http(client).post(
url=f"{base}/sandboxes",
headers=self._lifecycle_headers(key),
json=body,
),
response: Final = await self._http(client).post(
url=f"{base}/sandboxes",
headers=self._lifecycle_headers(key),
json=body,
)
data: Final = response.json()
sandbox_id: Final = str(data["id"])
@ -182,12 +179,9 @@ class OpenSandboxSandboxConfig(BaseSandboxConfig):
base: Final = str(handle._hidden_params.get("api_base") or self._api_base(api_base))
key: Final = self._api_key(api_key=api_key, handle=handle)
try:
response: Final = cast(
httpx.Response,
await self._http(client).delete(
url=f"{base}/sandboxes/{handle.id}",
headers=self._lifecycle_headers(key),
),
response: Final = await self._http(client).delete(
url=f"{base}/sandboxes/{handle.id}",
headers=self._lifecycle_headers(key),
)
except httpx.HTTPStatusError as e:
if e.response.status_code == 404:
@ -245,12 +239,9 @@ class OpenSandboxSandboxConfig(BaseSandboxConfig):
) -> None:
deadline: Final = time.monotonic() + ready_timeout
while True:
response = cast(
httpx.Response,
await self._http(client).get(
url=f"{api_base}/sandboxes/{sandbox_id}",
headers=headers,
),
response = await self._http(client).get(
url=f"{api_base}/sandboxes/{sandbox_id}",
headers=headers,
)
data = response.json()
state = self._sandbox_state(data)
@ -306,13 +297,10 @@ class OpenSandboxSandboxConfig(BaseSandboxConfig):
use_server_proxy: bool,
client: AsyncHTTPHandler | None,
) -> tuple[str, dict[str, str]]:
response: Final = cast(
httpx.Response,
await self._http(client).get(
url=f"{api_base}/sandboxes/{sandbox_id}/endpoints/{OPEN_SANDBOX_EXECD_PORT}",
headers=headers,
params={"use_server_proxy": use_server_proxy},
),
response: Final = await self._http(client).get(
url=f"{api_base}/sandboxes/{sandbox_id}/endpoints/{OPEN_SANDBOX_EXECD_PORT}",
headers=headers,
params={"use_server_proxy": use_server_proxy},
)
data: Final = response.json()
endpoint: Final = data.get("endpoint")
@ -329,15 +317,12 @@ class OpenSandboxSandboxConfig(BaseSandboxConfig):
client: AsyncHTTPHandler | None,
) -> list[str]:
timeout: Final = httpx.Timeout(connect=30.0, read=None, write=30.0, pool=None)
response: Final = cast(
httpx.Response,
await self._http(client).post(
url=url,
headers=headers,
timeout=timeout,
json=body,
stream=True,
),
response: Final = await self._http(client).post(
url=url,
headers=headers,
timeout=timeout,
json=body,
stream=True,
)
return await self._read_capped_lines(response)

67
litellm/llms/runwayml/videos/transformation.py
View file

@ -117,6 +117,10 @@ class RunwayMLVideoConfig(BaseVideoConfig):
def __init__(self):
super().__init__()
@staticmethod
def _parse_task_response(raw_response: httpx.Response) -> _RunwayTaskResponse:
return raw_response.json()
def get_supported_openai_params(self, model: str) -> list:
"""
Get the list of supported OpenAI parameters for video generation.
@ -141,7 +145,7 @@ class RunwayMLVideoConfig(BaseVideoConfig):
video_create_optional_params: VideoCreateOptionalRequestParams,
model: str,
drop_params: bool,
) -> dict:
) -> dict[str, object]:
"""
Map OpenAI parameters to RunwayML format.
@ -151,37 +155,44 @@ class RunwayMLVideoConfig(BaseVideoConfig):
- size -> ratio (convert "WIDTHxHEIGHT" to "WIDTH:HEIGHT")
- seconds -> duration (convert to integer)
"""
mapped_params: Final[dict[str, object]] = {}
supported_openai_params: Final = self.get_supported_openai_params(model)
return {
**self._prompt_image_param(video_create_optional_params),
**self._ratio_param(video_create_optional_params),
**self._duration_param(video_create_optional_params),
# Pass through other parameters that aren't OpenAI-specific
**{key: value for key, value in video_create_optional_params.items() if key not in supported_openai_params},
}
@staticmethod
def _prompt_image_param(video_create_optional_params: VideoCreateOptionalRequestParams) -> Mapping[str, object]:
# Handle input_reference parameter - map to promptImage
# RunwayML supports URLs and data URIs directly
if "input_reference" in video_create_optional_params:
input_reference: Final = video_create_optional_params["input_reference"]
# RunwayML supports URLs and data URIs directly
mapped_params["promptImage"] = input_reference
return {"promptImage": video_create_optional_params["input_reference"]}
return {}
@staticmethod
def _ratio_param(video_create_optional_params: VideoCreateOptionalRequestParams) -> Mapping[str, str]:
# Handle size parameter - convert "1280x720" to "1280:720"
if "size" in video_create_optional_params:
size: Final = video_create_optional_params["size"]
if isinstance(size, str) and "x" in size:
mapped_params["ratio"] = size.replace("x", ":")
return {"ratio": size.replace("x", ":")}
return {}
@staticmethod
def _duration_param(video_create_optional_params: VideoCreateOptionalRequestParams) -> Mapping[str, int]:
# Handle seconds parameter - convert to integer
if "seconds" in video_create_optional_params:
seconds: Final = video_create_optional_params["seconds"]
if seconds is not None:
try:
mapped_params["duration"] = int(float(seconds)) if isinstance(seconds, str) else int(seconds)
return {"duration": int(float(seconds)) if isinstance(seconds, str) else int(seconds)}
except (ValueError, TypeError):
# If conversion fails, use default duration
pass
# Pass through other parameters that aren't OpenAI-specific
supported_openai_params: Final = self.get_supported_openai_params(model)
for key, value in video_create_optional_params.items():
if key not in supported_openai_params:
mapped_params[key] = value
return mapped_params
return {}
def validate_environment(
self,
@ -236,7 +247,7 @@ class RunwayMLVideoConfig(BaseVideoConfig):
model: str,
prompt: str,
api_base: str,
video_create_optional_request_params: dict,
video_create_optional_request_params: dict[str, object],
litellm_params: GenericLiteLLMParams,
headers: dict,
) -> tuple[dict, RequestFiles, str]:
@ -406,20 +417,18 @@ class RunwayMLVideoConfig(BaseVideoConfig):
# Get task status to retrieve video URL
url: Final = f"{api_base}/tasks/{encoded_video_id}"
params: Final[dict[str, str]] = {}
return url, dict[str, str]()
return url, params
def _extract_video_url_from_response(self, response_data: dict[str, Any]) -> str:
def _extract_video_url_from_response(self, response_data: _RunwayTaskResponse) -> str:
"""
Helper method to extract video URL from RunwayML response.
Shared between sync and async transforms.
"""
# Extract video URL from the output field
video_url = None
if "output" in response_data and response_data["output"]:
output: Final = response_data["output"]
video_url = output[0] if isinstance(output, list) else output
raw_output: Final = response_data.get("output")
if raw_output:
video_url = raw_output if isinstance(raw_output, str) else raw_output[0]
if not video_url:
# Check if the video generation failed or is still processing
@ -453,7 +462,7 @@ class RunwayMLVideoConfig(BaseVideoConfig):
"output":["https://dnznrvs05pmza.cloudfront.net/.../video.mp4?_jwt=..."]
}
"""
response_data: Final = raw_response.json()
response_data: Final[_RunwayTaskResponse] = self._parse_task_response(raw_response)
video_url: Final = self._extract_video_url_from_response(response_data)
# Download the video from the CloudFront URL synchronously
@ -482,7 +491,7 @@ class RunwayMLVideoConfig(BaseVideoConfig):
"output":["https://dnznrvs05pmza.cloudfront.net/.../video.mp4?_jwt=..."]
}
"""
response_data: Final = raw_response.json()
response_data: Final[_RunwayTaskResponse] = self._parse_task_response(raw_response)
video_url: Final = self._extract_video_url_from_response(response_data)
# Download the video from the CloudFront URL asynchronously
@ -564,9 +573,7 @@ class RunwayMLVideoConfig(BaseVideoConfig):
# Construct the URL for task cancellation
url: Final = f"{api_base}/tasks/{encoded_video_id}/cancel"
data: Final[dict[str, str]] = {}
return url, data
return url, dict[str, str]()
def transform_video_delete_response(
self,
@ -604,9 +611,7 @@ class RunwayMLVideoConfig(BaseVideoConfig):
url: Final = f"{api_base}/tasks/{encoded_video_id}"
# Empty dict for GET request (no body)
data: Final[dict[str, str]] = {}
return url, data
return url, dict[str, str]()
def transform_video_status_retrieve_response(
self,

8
litellm/llms/vertex_ai/vector_stores/rag_api/transformation.py
View file

@ -97,7 +97,7 @@ class VertexVectorStoreConfig(BaseVectorStoreConfig, VertexBase):
def __init__(self):
super().__init__()
def get_auth_credentials(self, litellm_params: dict) -> BaseVectorStoreAuthCredentials:
def get_auth_credentials(self, litellm_params: Mapping[str, object]) -> BaseVectorStoreAuthCredentials:
# Get credentials and project info
vertex_credentials: Final = self.get_vertex_ai_credentials(dict(litellm_params))
vertex_project: Final = self.get_vertex_ai_project(dict(litellm_params))
@ -122,7 +122,9 @@ class VertexVectorStoreConfig(BaseVectorStoreConfig, VertexBase):
"write": [("POST", "/ragCorpora")],
}
def validate_environment(self, headers: dict, litellm_params: GenericLiteLLMParams | None) -> dict:
def validate_environment(
self, headers: dict[str, str], litellm_params: GenericLiteLLMParams | None
) -> dict[str, str]:
"""
Validate and set up authentication for Vertex AI RAG API
"""
@ -135,7 +137,7 @@ class VertexVectorStoreConfig(BaseVectorStoreConfig, VertexBase):
def get_complete_url(
self,
api_base: str | None,
litellm_params: dict,
litellm_params: dict[str, object],
) -> str:
"""
Get the Base endpoint for Vertex AI RAG API

20
litellm/llms/vertex_ai/vertex_llm_base.py
View file

@ -27,6 +27,15 @@ from .common_utils import (
get_vertex_base_url,
)
def _graft_default_vertex_path(api_base: str, default_url: str) -> str:
parsed_api_base: Final = urlparse(api_base)
default_segments: Final = urlparse(default_url).path.lstrip("/").split("/")
graft_segments: Final = default_segments[1:] if default_segments[0] in ("v1", "v1beta1") else default_segments
grafted_path: Final = parsed_api_base.path.rstrip("/") + "/" + "/".join(graft_segments)
return parsed_api_base._replace(path=grafted_path).geturl()
GOOGLE_IMPORT_ERROR_MESSAGE: Final = (
"Google Cloud SDK not found. Install it with: pip install 'litellm[google]' or pip install google-cloud-aiplatform"
)
@ -621,8 +630,9 @@ class VertexBase:
Handles custom api_base for:
1. Gemini (Google AI Studio) - constructs /models/{model}:{endpoint}
2. Vertex AI with standard proxies - constructs {api_base}:{endpoint};
if api_base has no path (bare host), grafts the default vertex URL path onto it
2. Vertex AI with standard proxies - grafts the default vertex URL path onto the
api_base when its path is empty or only an API version (/v1, /v1beta1);
otherwise constructs {api_base}:{endpoint}
3. Vertex AI with PSC endpoints - constructs full path structure
{api_base}/v1/projects/{project}/locations/{location}/endpoints/{model}:{endpoint}
(only when use_psc_endpoint_format=True)
@ -669,10 +679,14 @@ class VertexBase:
)
elif urlparse(api_base).path in ("", "/"):
url = api_base.rstrip("/") + urlparse(url).path
elif urlparse(api_base).path.rstrip("/") in ("/v1", "/v1beta1") and "/projects/" in urlparse(url).path:
url = _graft_default_vertex_path(api_base=api_base, default_url=url)
else:
url = f"{api_base}:{endpoint}"
if stream is True:
url = url + "?alt=sse"
parsed_stream_url: Final = urlparse(url)
stream_query: Final = f"{parsed_stream_url.query}&alt=sse" if parsed_stream_url.query else "alt=sse"
url = parsed_stream_url._replace(query=stream_query).geturl()
return auth_header, url
def _get_token_and_url(

267
litellm/model_prices_and_context_window_backup.json
View file

File diff suppressed because it is too large Load diff

2
litellm/models/base.py
View file

@ -33,6 +33,6 @@ class DomainModel(BaseModel):
return cls(**record.dict())
return cls(**dict(record))
def to_db_dict(self, exclude_unset: bool = False) -> dict[str, Any]:
def to_db_dict(self, exclude_unset: bool = False) -> dict[str, object]:
"""Convert domain model to a dictionary for database operations."""
return self.model_dump(exclude_none=True, exclude_unset=exclude_unset)

16
litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py
View file

@ -1096,8 +1096,7 @@ async def exchange_token_with_server(
headers={"Accept": "application/json", **token_request.headers},
data=token_data,
)
if response is not None:
response.raise_for_status()
response.raise_for_status()
except httpx.HTTPStatusError as exc:
fault: Final = classify_upstream_token_rejection(
exc.response,
@ -1119,11 +1118,6 @@ async def exchange_token_with_server(
)
return _bridge_mint_error_response("invalid_refresh")
return render_token_fault(fault)
if response is None:
raise HTTPException(
status_code=502,
detail="MCP upstream token endpoint returned no response",
)
token_response = response.json()
# Validate token response against server-configured rules before any storage.
@ -1536,16 +1530,10 @@ async def _post_dcr_registration(
headers=headers,
json=register_data,
)
if response is not None:
response.raise_for_status()
response.raise_for_status()
except httpx.HTTPStatusError as exc:
status_code, detail = dcr_fault_detail(classify_upstream_dcr_rejection(exc.response, log_context=server_id))
raise HTTPException(status_code=status_code, detail=detail) from exc
if response is None:
raise HTTPException(
status_code=502,
detail="MCP upstream registration endpoint returned no response",
)
return response

63
litellm/proxy/_experimental/mcp_server/oauth2_flow_backfill.py
View file

@ -36,7 +36,10 @@ a healed fleet has no null rows and the backfill exits after one query.
import json
from collections import Counter
from typing import Any, Final, Literal
from collections.abc import Mapping, Sequence
from typing import Final, Literal, Protocol
from pydantic import JsonValue
from litellm._logging import verbose_proxy_logger
from litellm.proxy._experimental.mcp_server.db import _decode_oauth_payload, decrypt_credentials
@ -55,9 +58,59 @@ BackfillRule = Literal[
_BACKFILL_AUDIT_ACTOR: Final = "oauth2_flow_backfill"
def _decrypted_credentials(raw_credentials: Any) -> MCPCredentials | None:
class _MCPServerRow(Protocol):
"""The ``LiteLLM_MCPServerTable`` columns this backfill reads."""
@property
def server_id(self) -> str: ...
@property
def authorization_url(self) -> str | None: ...
@property
def registration_url(self) -> str | None: ...
@property
def token_url(self) -> str | None: ...
@property
def credentials(self) -> str | Mapping[str, JsonValue] | None: ...
class _MCPUserCredentialRow(Protocol):
"""The ``LiteLLM_MCPUserCredentials`` columns this backfill reads."""
@property
def server_id(self) -> str: ...
@property
def credential_b64(self) -> str: ...
class _MCPServerTable(Protocol):
async def find_many(self, *, where: Mapping[str, object]) -> Sequence[_MCPServerRow]: ...
async def update_many(self, *, where: Mapping[str, object], data: Mapping[str, str]) -> object: ...
class _MCPUserCredentialsTable(Protocol):
async def find_many(self, *, where: Mapping[str, object]) -> Sequence[_MCPUserCredentialRow]: ...
def _mcp_server_table(prisma_client: PrismaClient) -> _MCPServerTable:
"""The MCP server table, typed so the untyped prisma client surface stops here."""
return prisma_client.db.litellm_mcpservertable
def _mcp_user_credentials_table(prisma_client: PrismaClient) -> _MCPUserCredentialsTable:
"""The per-user MCP credential table, typed so the untyped prisma client surface stops here."""
return prisma_client.db.litellm_mcpusercredentials
def _decrypted_credentials(raw_credentials: str | Mapping[str, JsonValue] | None) -> MCPCredentials | None:
if raw_credentials is None:
return None
parsed: JsonValue | Mapping[str, JsonValue]
if isinstance(raw_credentials, str):
try:
parsed = json.loads(raw_credentials)
@ -92,14 +145,14 @@ def classify_null_flow_row(
async def backfill_null_oauth2_flows(prisma_client: PrismaClient) -> dict[BackfillRule, int]:
"""Classify every ``auth_type=oauth2`` row whose ``oauth2_flow`` is null; stamp the provable
ones, warn on the ambiguous ones, and return counts per rule."""
null_rows: Final[list[Any]] = await prisma_client.db.litellm_mcpservertable.find_many(
null_rows: Final[Sequence[_MCPServerRow]] = await _mcp_server_table(prisma_client).find_many(
where={"auth_type": "oauth2", "oauth2_flow": None},
)
if not null_rows:
return {}
server_ids: Final = [row.server_id for row in null_rows]
token_rows: Final[list[Any]] = await prisma_client.db.litellm_mcpusercredentials.find_many(
token_rows: Final[Sequence[_MCPUserCredentialRow]] = await _mcp_user_credentials_table(prisma_client).find_many(
where={"server_id": {"in": server_ids}},
)
server_ids_with_oauth_tokens: Final[set[str]] = {
@ -141,7 +194,7 @@ async def backfill_null_oauth2_flows(prisma_client: PrismaClient) -> dict[Backfi
stamped_flows: Final = {flow for _, (flow, _) in classified if flow is not None}
for stamped_flow in stamped_flows:
server_ids_for_flow = [row.server_id for row, (row_flow, _) in classified if row_flow == stamped_flow]
await prisma_client.db.litellm_mcpservertable.update_many(
await _mcp_server_table(prisma_client).update_many(
where={"server_id": {"in": server_ids_for_flow}, "oauth2_flow": None},
data={"oauth2_flow": stamped_flow, "updated_by": _BACKFILL_AUDIT_ACTOR},
)

17
litellm/proxy/_experimental/mcp_server/outbound_credentials/client_credentials.py
View file

@ -19,9 +19,9 @@ Implements the client-credentials behavior contract for the v2 resolver:
identity.
The token-endpoint POST is injected (``M2MTokenEndpointPost``) so the grant orchestration is
testable without a live IdP; ``post_client_credentials_grant`` is the httpx edge and the one
place the untyped response boundary is contained. Failures are values: the source returns
``Result[OAuthToken, CredError]``; only the httpx edge touches exceptions.
testable without a live IdP; ``post_client_credentials_grant`` is the httpx edge. Failures are
values: the source returns ``Result[OAuthToken, CredError]``; only the httpx edge touches
exceptions.
"""
from __future__ import annotations
@ -95,18 +95,17 @@ async def post_client_credentials_grant(
) -> TokenEndpointOutcome:
"""POST the grant to the token endpoint and classify the transport outcome.
The httpx edge: litellm's handler is partially typed (and raises ``HTTPStatusError`` itself on
a 4xx/5xx), so the untyped boundary is contained here and every field the caller reads comes
out of a validated ``TokenEndpointOutcome``.
The httpx edge: litellm's handler raises ``HTTPStatusError`` itself on a 4xx/5xx, and every
field the caller reads comes out of a validated ``TokenEndpointOutcome``.
"""
from litellm.llms.custom_httpx.http_handler import ( # noqa: PLC0415 # defer heavy handler import to call time
get_async_httpx_client, # pyright: ignore[reportUnknownVariableType] # handler is partially typed
get_async_httpx_client, # pyright: ignore[reportUnknownVariableType] # handler factory params are coarsely typed
)
from litellm.types.llms.custom_http import httpxSpecialProvider # noqa: PLC0415 # deferred with the handler import
try:
client: Final = get_async_httpx_client(llm_provider=httpxSpecialProvider.Oauth2Check)
response = await client.post( # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # handler is partially typed
response: Final = await client.post( # pyright: ignore[reportUnknownMemberType] # handler params are coarsely typed
url, headers={"Accept": "application/json", **headers}, data=form
)
except httpx.HTTPStatusError as status_err:
@ -114,8 +113,6 @@ async def post_client_credentials_grant(
return TokenEndpointDenied(status_code=status_code, detail=f"token endpoint returned HTTP {status_code}")
except Exception as exc: # noqa: BLE001 # any transport failure is the same outcome: unreachable
return TokenEndpointUnreachable(detail=str(exc))
if not isinstance(response, httpx.Response):
return TokenEndpointUnreachable(detail="token endpoint returned no response")
try:
body: Final = _TOKEN_BODY_ADAPTER.validate_json(response.content)
except ValidationError:

7
litellm/proxy/_experimental/mcp_server/outbound_credentials/token_endpoint.py
View file

@ -111,9 +111,6 @@ class TokenEndpointClient:
return Error(
CredError.of_upstream_unavailable("token exchange failed: token endpoint returned a non-JSON response")
)
if raw is None:
verbose_proxy_logger.warning("MCP token endpoint %s returned no response", endpoint)
return Error(CredError.of_upstream_unavailable("token exchange failed: no response from token endpoint"))
try:
parsed: Final = _TokenEndpointResponse.model_validate(raw)
except ValidationError:
@ -199,7 +196,7 @@ def _cache_ttl_seconds(expires_in: int | None) -> int:
)
async def _post_form(endpoint: str, data: dict[str, str]) -> object | None:
async def _post_form(endpoint: str, data: dict[str, str]) -> object:
# litellm's httpx handler and httpx.Response are only partially typed; the token endpoint
# returns a JSON object that `_TokenEndpointResponse` validates, so the untyped boundary is
# contained here. A non-2xx raises `httpx.HTTPStatusError`, an unreachable endpoint raises
@ -208,8 +205,6 @@ async def _post_form(endpoint: str, data: dict[str, str]) -> object | None:
# each to a CredError.
client = get_async_httpx_client(llm_provider=httpxSpecialProvider.MCP) # pyright: ignore[reportUnknownVariableType] # litellm http handler is untyped
response = await client.post(endpoint, data=data) # pyright: ignore[reportUnknownMemberType,reportUnknownVariableType] # litellm http handler is untyped
if response is None:
return None
response.raise_for_status()
return response.json() # pyright: ignore[reportAny] # untyped JSON; validated by _TokenEndpointResponse in fetch

93
litellm/proxy/_experimental/mcp_server/toolset_db.py
View file

@ -1,5 +1,9 @@
import json
from typing import Final
from collections.abc import Mapping, Sequence
from datetime import datetime
from typing import Final, Protocol
from typing_extensions import NotRequired, ReadOnly, TypedDict
from litellm._logging import verbose_proxy_logger
from litellm._uuid import uuid
@ -7,18 +11,73 @@ from litellm.proxy.utils import PrismaClient
from litellm.repositories.table_repositories import MCPToolsetRepository
from litellm.types.mcp_server.mcp_toolset import (
MCPToolset,
MCPToolsetTool,
NewMCPToolsetRequest,
UpdateMCPToolsetRequest,
)
def _toolset_from_row(row) -> MCPToolset:
class MCPToolsetFields(TypedDict):
"""The ``MCPToolset`` constructor keywords a toolset row expands into."""
toolset_id: ReadOnly[str]
toolset_name: ReadOnly[str]
description: NotRequired[ReadOnly[str | None]]
tools: NotRequired[ReadOnly[list[MCPToolsetTool]]]
created_at: NotRequired[ReadOnly[datetime | None]]
created_by: NotRequired[ReadOnly[str | None]]
updated_at: NotRequired[ReadOnly[datetime | None]]
updated_by: NotRequired[ReadOnly[str | None]]
class MCPToolsetRowData(TypedDict):
"""A toolset table row, whose ``tools`` column is stored as JSON."""
toolset_id: ReadOnly[str]
toolset_name: ReadOnly[str]
description: NotRequired[ReadOnly[str | None]]
tools: NotRequired[ReadOnly[str | list[MCPToolsetTool]]]
created_at: NotRequired[ReadOnly[datetime | None]]
created_by: NotRequired[ReadOnly[str | None]]
updated_at: NotRequired[ReadOnly[datetime | None]]
updated_by: NotRequired[ReadOnly[str | None]]
class MCPToolsetRow(Protocol):
"""A row of the toolset table, as the prisma client returns it."""
def model_dump(self) -> MCPToolsetRowData: ...
class MCPToolsetTable(Protocol):
"""The prisma table actions this module runs against the toolset table."""
async def create(self, data: Mapping[str, object]) -> MCPToolsetRow: ...
async def find_unique(self, where: Mapping[str, object]) -> MCPToolsetRow | None: ...
async def find_first(self, where: Mapping[str, object]) -> MCPToolsetRow | None: ...
async def find_many(self, where: Mapping[str, object]) -> Sequence[MCPToolsetRow]: ...
async def update(self, where: Mapping[str, object], data: Mapping[str, object]) -> MCPToolsetRow: ...
async def delete(self, where: Mapping[str, object]) -> MCPToolsetRow: ...
def _toolset_table(prisma_client: PrismaClient) -> MCPToolsetTable:
"""The toolset table actions of the prisma client."""
return MCPToolsetRepository(prisma_client).table
def _toolset_from_row(row: MCPToolsetRow) -> MCPToolset:
data: Final = row.model_dump()
tools = data.get("tools") or []
if isinstance(tools, str):
tools = json.loads(tools)
data["tools"] = tools
return MCPToolset(**data)
tools: Final = data.get("tools") or []
resolved: Final[MCPToolsetFields] = {
**data,
"tools": json.loads(tools) if isinstance(tools, str) else tools,
}
return MCPToolset(**resolved)
async def create_mcp_toolset(
@ -31,7 +90,7 @@ async def create_mcp_toolset(
data_dict["tools"] = json.dumps(data_dict.get("tools", []))
data_dict["created_by"] = touched_by
data_dict["updated_by"] = touched_by
row: Final = await MCPToolsetRepository(prisma_client).table.create(data=data_dict)
row: Final = await _toolset_table(prisma_client).create(data=data_dict)
return _toolset_from_row(row)
@ -39,7 +98,7 @@ async def get_mcp_toolset(
prisma_client: PrismaClient,
toolset_id: str,
) -> MCPToolset | None:
row: Final = await MCPToolsetRepository(prisma_client).table.find_unique(where={"toolset_id": toolset_id})
row: Final = await _toolset_table(prisma_client).find_unique(where={"toolset_id": toolset_id})
if row is None:
return None
return _toolset_from_row(row)
@ -47,13 +106,11 @@ async def get_mcp_toolset(
async def list_mcp_toolsets(
prisma_client: PrismaClient,
toolset_ids: list[str] | None = None,
) -> list[MCPToolset]:
toolset_ids: Sequence[str] | None = None,
) -> Sequence[MCPToolset]:
try:
where = {}
if toolset_ids is not None:
where = {"toolset_id": {"in": toolset_ids}}
rows: Final = await MCPToolsetRepository(prisma_client).table.find_many(where=where)
where: Final[Mapping[str, object]] = {} if toolset_ids is None else {"toolset_id": {"in": toolset_ids}}
rows: Final = await _toolset_table(prisma_client).find_many(where=where)
return [_toolset_from_row(r) for r in rows]
except Exception as e:
verbose_proxy_logger.warning("litellm.proxy._experimental.mcp_server.toolset_db::list_mcp_toolsets - %s", e)
@ -64,7 +121,7 @@ async def get_mcp_toolset_by_name(
prisma_client: PrismaClient,
toolset_name: str,
) -> MCPToolset | None:
row: Final = await MCPToolsetRepository(prisma_client).table.find_first(where={"toolset_name": toolset_name})
row: Final = await _toolset_table(prisma_client).find_first(where={"toolset_name": toolset_name})
if row is None:
return None
return _toolset_from_row(row)
@ -80,7 +137,7 @@ async def update_mcp_toolset(
data_dict["tools"] = json.dumps(data_dict["tools"])
data_dict["updated_by"] = touched_by
try:
row: Final = await MCPToolsetRepository(prisma_client).table.update(
row: Final = await _toolset_table(prisma_client).update(
where={"toolset_id": data.toolset_id},
data=data_dict,
)
@ -98,7 +155,7 @@ async def delete_mcp_toolset(
toolset_id: str,
) -> MCPToolset | None:
try:
row: Final = await MCPToolsetRepository(prisma_client).table.delete(where={"toolset_id": toolset_id})
row: Final = await _toolset_table(prisma_client).delete(where={"toolset_id": toolset_id})
except Exception as e:
from prisma.errors import RecordNotFoundError

16
litellm/proxy/_types.py
View file

@ -2542,7 +2542,7 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase):
)
alerting: list | None = Field(
None,
description="List of alerting integrations. Today, just slack - `alerting: ['slack']`",
description="List of alerting integrations - e.g. `alerting: ['slack', 'webhook', 'email']`. 'slack' posts Slack-format messages to any Slack-compatible webhook (Slack, Rocket.Chat, Mattermost); 'webhook' posts structured JSON budget alerts to WEBHOOK_URL",
)
alert_types: list[AlertType] | None = Field(
None,
@ -3550,6 +3550,19 @@ class AllCallbacks(LiteLLMPydanticObjectBase):
)
class SpendLogsRouterMetadata(TypedDict):
"""
Router provenance stamped on spend logs for deployments flagged with
model_info.internal_router_model, correlating the requested model group
with the provider deployment that served the call
"""
requested_model: ReadOnly[str | None]
selected_model: ReadOnly[str | None]
selected_provider: ReadOnly[str | None]
router_correlation_id: ReadOnly[str | None]
class SpendLogsMetadata(TypedDict):
"""
Specific metadata k,v pairs logged to spendlogs for easier cost tracking
@ -3592,6 +3605,7 @@ class SpendLogsMetadata(TypedDict):
compression_savings: CompressionSavingsMetadata | None
autorouter_savings: ReadOnly[float | None] # stamped by the logging payload; None = not auto-routed
litellm_gateway_injected_cache: ReadOnly[str | None]
router_metadata: ReadOnly[SpendLogsRouterMetadata | None] # None = deployment not flagged internal_router_model
class SpendLogsPayload(TypedDict):

10
litellm/proxy/agent_endpoints/agent_registry.py
View file

@ -401,12 +401,10 @@ class AgentRegistry:
The patched agent
"""
try:
existing_row: Final = await AgentsRepository(prisma_client).table.find_unique(
where={"agent_id": agent_id} # mutable-ok: prisma filters are plain dicts
)
if existing_row is None:
existing_record: Final = await agents_table(prisma_client).find_unique(where={"agent_id": agent_id})
if existing_record is None:
raise Exception(f"Agent with ID {agent_id} not found")
existing_agent: Final = dict(existing_row)
existing_agent: Final[Mapping[str, object]] = dict(existing_record)
augment_agent: Final = {**existing_agent, **agent}
update_data: Final[dict[str, object]] = {}
@ -433,7 +431,7 @@ class AgentRegistry:
update_data["extra_headers"] = extra_headers_value if extra_headers_value is not None else []
if agent.get("object_permission") is not None:
agent_copy: Final = dict(augment_agent)
existing_object_permission_id: Final = existing_agent.get("object_permission_id")
existing_object_permission_id: Final = existing_record.object_permission_id
object_permission_id: Final = await handle_update_object_permission_common(
agent_copy,
existing_object_permission_id,

6
litellm/proxy/client/cli/README.md
View file

@ -489,7 +489,7 @@ lite codex exec "summarize the repo"
Each command resolves your LiteLLM key (logging in via SSO when none is stored and you are at a terminal; otherwise it expects `LITELLM_PROXY_API_KEY` or `--api-key`), checks the key against the proxy so bad credentials fail immediately instead of deep inside the agent, exports the environment variables the agent reads, then replaces itself with the agent process.
The right variables are picked per agent. Claude Code gets `ANTHROPIC_BASE_URL` (the proxy root, so it appends `/v1/messages`) and `ANTHROPIC_AUTH_TOKEN`, with any stray `ANTHROPIC_API_KEY` cleared so the proxy token wins. Codex and OpenCode get `OPENAI_BASE_URL` (the proxy plus `/v1`) and `OPENAI_API_KEY`. Codex ignores `OPENAI_BASE_URL`, so it is additionally pointed at the proxy through a custom provider passed as `-c` config overrides (HTTP/SSE Responses transport, since the proxy does not speak the Responses WebSocket protocol).
The right variables are picked per agent. Claude Code gets `ANTHROPIC_BASE_URL` (the proxy root, so it appends `/v1/messages`) and `ANTHROPIC_AUTH_TOKEN`, with any stray `ANTHROPIC_API_KEY` cleared so the proxy token wins, and `ENABLE_TOOL_SEARCH=true` (unless you already set it) so Claude Code keeps tool search on even though the base URL is a proxy rather than a first-party Anthropic host. Codex and OpenCode get `OPENAI_BASE_URL` (the proxy plus `/v1`) and `OPENAI_API_KEY`. Codex ignores `OPENAI_BASE_URL`, so it is additionally pointed at the proxy through a custom provider passed as `-c` config overrides (HTTP/SSE Responses transport, since the proxy does not speak the Responses WebSocket protocol).
Options (these belong to the wrapper, so put them before the agent's own flags):
@ -505,7 +505,7 @@ The credential is short-lived by design (default 24h, configurable via `LITELLM_
### Route Every Claude Code Session Through the Proxy
`lite claude` wraps a single invocation, but `lite up` goes further: it patches `~/.claude/settings.json`, Claude Code's own config file, so that every Claude Code session started afterward -- from any terminal, launched normally with just `claude`, no wrapper needed -- routes through your LiteLLM proxy. It sets `env.ANTHROPIC_BASE_URL` to the proxy URL and `apiKeyHelper` to a `lite auth print-token` invocation, drops any stray static `ANTHROPIC_API_KEY` so the helper-issued token wins, and leaves every other setting in the file untouched. It backs up the original file before patching it.
`lite claude` wraps a single invocation, but `lite up` goes further: it patches `~/.claude/settings.json`, Claude Code's own config file, so that every Claude Code session started afterward -- from any terminal, launched normally with just `claude`, no wrapper needed -- routes through your LiteLLM proxy. It sets `env.ANTHROPIC_BASE_URL` to the proxy URL, `env.ENABLE_TOOL_SEARCH` to `true` when that key is missing, and `apiKeyHelper` to a `lite auth print-token` invocation, drops any stray static `ANTHROPIC_API_KEY` so the helper-issued token wins, and leaves every other setting in the file untouched. It backs up the original file before patching it.
Two things need to already be true: you've run `lite login` (or `lite login --pkce`, whose key the helper renews on its own), since the apiKeyHelper depends on that stored token, and the proxy is already reachable, since `lite up` does not start one for you.
@ -529,7 +529,7 @@ Cursor is not supported: it has no equivalent file-based config to hot-patch thi
lite --base-url https://your-proxy.example.com login --config-claude
```
It writes the same two settings `lite up` does, `env.ANTHROPIC_BASE_URL` and `apiKeyHelper`, but persistently: there is no backup, nothing to restore, and no foreground process to keep alive. Every other key in `~/.claude/settings.json` is preserved, the file is created if it does not exist, and it is written atomically with owner-only permissions. Plain `lite login` is unchanged; nothing happens to your Claude Code config unless you pass the flag.
It writes the same settings `lite up` does, `env.ANTHROPIC_BASE_URL`, `env.ENABLE_TOOL_SEARCH`, and `apiKeyHelper`, but persistently: there is no backup, nothing to restore, and no foreground process to keep alive. Every other key in `~/.claude/settings.json` is preserved, the file is created if it does not exist, and it is written atomically with owner-only permissions. Plain `lite login` is unchanged; nothing happens to your Claude Code config unless you pass the flag.
Because the credential is reached through `apiKeyHelper` rather than copied into the file, a later `lite login` refreshes it with no further action: Claude Code re-runs the helper on every request and picks up whatever token the most recent login stored. Nothing secret is written to `settings.json`.

9
litellm/proxy/client/cli/commands/agents.py
View file

@ -13,6 +13,8 @@ from .auth import context_secret_vault, get_stored_api_key, login
ANTHROPIC_BASE_URL_ENV: Final = "ANTHROPIC_BASE_URL"
ANTHROPIC_AUTH_TOKEN_ENV: Final = "ANTHROPIC_AUTH_TOKEN"
ANTHROPIC_API_KEY_ENV: Final = "ANTHROPIC_API_KEY"
ENABLE_TOOL_SEARCH_ENV: Final = "ENABLE_TOOL_SEARCH"
ENABLE_TOOL_SEARCH_VALUE: Final = "true"
OPENAI_BASE_URL_ENV: Final = "OPENAI_BASE_URL"
OPENAI_API_KEY_ENV: Final = "OPENAI_API_KEY"
@ -61,7 +63,10 @@ def build_agent_env(
Anthropic clients (Claude Code) append /v1/messages to ANTHROPIC_BASE_URL,
so it stays the bare proxy root; OpenAI clients (Codex, OpenCode) expect the
/v1 suffix on OPENAI_BASE_URL. ANTHROPIC_API_KEY is dropped so a stray
Anthropic key cannot win over the bearer token we set.
Anthropic key cannot win over the bearer token we set. ENABLE_TOOL_SEARCH
defaults to true because Claude Code turns tool search off when
ANTHROPIC_BASE_URL is not a first-party Anthropic host; a value already in
the environment is left alone.
"""
env: Final = dict(base_env)
root: Final = base_url.rstrip("/")
@ -69,6 +74,8 @@ def build_agent_env(
env[ANTHROPIC_BASE_URL_ENV] = root
env[ANTHROPIC_AUTH_TOKEN_ENV] = api_key
env.pop(ANTHROPIC_API_KEY_ENV, None)
if ENABLE_TOOL_SEARCH_ENV not in env:
env[ENABLE_TOOL_SEARCH_ENV] = ENABLE_TOOL_SEARCH_VALUE
if PROFILE_OPENAI in profiles:
env[OPENAI_BASE_URL_ENV] = root + "/v1"
env[OPENAI_API_KEY_ENV] = api_key

3
litellm/proxy/client/cli/commands/autoroute/settings.py
View file

@ -9,6 +9,8 @@ API_KEY_HELPER_KEY: Final = "apiKeyHelper"
ANTHROPIC_API_KEY_KEY: Final = "ANTHROPIC_API_KEY"
ANTHROPIC_AUTH_TOKEN_KEY: Final = "ANTHROPIC_AUTH_TOKEN"
ANTHROPIC_BASE_URL_KEY: Final = "ANTHROPIC_BASE_URL"
ENABLE_TOOL_SEARCH_KEY: Final = "ENABLE_TOOL_SEARCH"
ENABLE_TOOL_SEARCH_VALUE: Final = "true"
# Force every one of Claude Code's own model tiers to request the auto-router by name.
# Router's auto-router registry is keyed by the literal requested model string
# (litellm/router.py:10711-10717) with no wildcard/pattern resolution, so a bare "*"
@ -34,6 +36,7 @@ def merge_claude_settings_static_token(
raw_env: Final = settings.get(ENV_KEY, {})
base_env: Final = raw_env if isinstance(raw_env, dict) else {}
env: Final[dict[str, JsonValue]] = {
ENABLE_TOOL_SEARCH_KEY: ENABLE_TOOL_SEARCH_VALUE,
**base_env,
ANTHROPIC_BASE_URL_KEY: base_url.rstrip("/"),
ANTHROPIC_AUTH_TOKEN_KEY: auth_token,

11
litellm/proxy/client/cli/commands/claude_settings.py
View file

@ -21,6 +21,8 @@ ENV_KEY: Final = "env"
API_KEY_HELPER_KEY: Final = "apiKeyHelper"
ANTHROPIC_BASE_URL_KEY: Final = "ANTHROPIC_BASE_URL"
ANTHROPIC_API_KEY_KEY: Final = "ANTHROPIC_API_KEY"
ENABLE_TOOL_SEARCH_KEY: Final = "ENABLE_TOOL_SEARCH"
ENABLE_TOOL_SEARCH_VALUE: Final = "true"
CLAUDE_SETTINGS_PATH: Final = Path.home() / ".claude" / "settings.json"
BACKUP_PATH: Final = Path.home() / ".litellm" / "claude_settings_backup.json"
@ -70,12 +72,15 @@ def merge_claude_settings(
Only env.ANTHROPIC_BASE_URL and the top-level apiKeyHelper are overridden; a
stray env.ANTHROPIC_API_KEY is dropped so it cannot outrank the helper-issued
token (same reasoning as build_agent_env in agents.py). Every other key is
preserved untouched.
token (same reasoning as build_agent_env in agents.py). ENABLE_TOOL_SEARCH
defaults to true because Claude Code turns tool search off when
ANTHROPIC_BASE_URL is not a first-party Anthropic host; an existing value is
left alone. Every other key is preserved untouched.
"""
raw_env: Final = settings.get(ENV_KEY, {})
base_env: Final = raw_env if isinstance(raw_env, dict) else {}
env: Final = {
ENABLE_TOOL_SEARCH_KEY: ENABLE_TOOL_SEARCH_VALUE,
**{key: value for key, value in base_env.items() if key != ANTHROPIC_API_KEY_KEY},
ANTHROPIC_BASE_URL_KEY: base_url.rstrip("/"),
}
@ -144,6 +149,8 @@ __all__ = (
"AUTOROUTE_BACKUP_PATH",
"BACKUP_PATH",
"CLAUDE_SETTINGS_PATH",
"ENABLE_TOOL_SEARCH_KEY",
"ENABLE_TOOL_SEARCH_VALUE",
"ENV_KEY",
"SETTINGS_FILE_OWNERS",
"ClaudeSettingsError",

2
litellm/proxy/client/cli/commands/teams.py
View file

@ -34,7 +34,7 @@ def teams():
"""Manage teams and team assignments"""
def display_teams_table(teams: list[dict[str, Any]]) -> None:
def display_teams_table(teams: Sequence[dict[str, Any]]) -> None:
"""Display teams in a formatted table"""
console: Final = Console()

151
litellm/proxy/common_request_processing.py
View file

@ -502,7 +502,7 @@ def _as_success_dispatcher(logging_obj: _DispatchesSuccessHandlers) -> _Dispatch
return logging_obj
def _serialize_http_exception_detail(
def serialize_http_exception_detail(
detail: object,
) -> tuple[str, dict | None]:
"""
@ -535,7 +535,7 @@ def _serialize_http_exception_detail(
def proxy_exception_from_http_exception(exc: HTTPException, headers: dict[str, str]) -> ProxyException:
raw_detail: Final = _getattr_object(exc, "detail", str(exc))
message, structured_fields = _serialize_http_exception_detail(raw_detail)
message, structured_fields = serialize_http_exception_detail(raw_detail)
existing_fields: Final = getattr(exc, "provider_specific_fields", None) or {}
merged_fields: Final = {**existing_fields, **structured_fields} if structured_fields else (existing_fields or None)
return ProxyException(
@ -818,7 +818,7 @@ async def _buffer_first_chunk_honoring_disconnect(
raise _ClientDisconnectedBeforeFirstChunk()
def _sse_error_payload(exc: BaseException) -> tuple[int, Mapping[str, object]]:
def sse_error_payload(exc: BaseException) -> tuple[int, Mapping[str, object]]:
"""Build the ProxyException-shaped ``{"error": ...}`` body used in SSE error frames.
Matches ``ProxyException.to_dict()`` so streaming and non-streaming error frames
@ -827,7 +827,7 @@ def _sse_error_payload(exc: BaseException) -> tuple[int, Mapping[str, object]]:
# Preserve status code from HTTPException (e.g. guardrail blocks)
error_status: Final = getattr(exc, "status_code", status.HTTP_500_INTERNAL_SERVER_ERROR)
raw_detail: Final = _getattr_object(exc, "detail", "Error processing stream start")
message, structured_fields = _serialize_http_exception_detail(raw_detail)
message, structured_fields = serialize_http_exception_detail(raw_detail)
existing_fields: Final = getattr(exc, "provider_specific_fields", None) or {}
merged_fields: Final = {**existing_fields, **structured_fields} if structured_fields else (existing_fields or None)
@ -942,7 +942,7 @@ async def create_response(
# Unexpected error consuming first chunk.
verbose_proxy_logger.exception("Error consuming first chunk from generator: %s", e)
error_status, error_obj = _sse_error_payload(e)
error_status, error_obj = sse_error_payload(e)
async def error_gen_message() -> AsyncGenerator[str, None]:
for frame in _sse_error_frames(error_obj):
@ -1119,7 +1119,7 @@ async def open_sse_before_first_byte(
# would never fire and the failure would go unaudited. The hook
# also gets to sanitize what reaches the client, by returning or
# raising a replacement, so its answer decides the frame.
_, error_obj = _sse_error_payload(await _sanitized_late_failure(exc, on_late_failure))
_, error_obj = sse_error_payload(await _sanitized_late_failure(exc, on_late_failure))
for frame in _sse_error_frames(error_obj):
yield frame.encode()
return
@ -2409,53 +2409,14 @@ class ProxyBaseLLMRequestProcessing:
if requested_model_from_client:
self.data["_litellm_client_requested_model"] = requested_model_from_client
# Streaming: attach a closure that fires after all guardrail
# end-of-stream blocks complete. CSW.__anext__ stores the
# assembled response on logging_obj; the outer consumer
# (ProxyLogging._fire_deferred_stream_logging) fires the
# closure after the full streaming pipeline finishes.
# The closure runs non-apply_guardrail hooks on the
# assembled response, then fires success logging.
# Only for CustomStreamWrapper — raw async generators from
# passthrough routes bypass CSW and would orphan the closure.
from litellm.litellm_core_utils.streaming_handler import (
CustomStreamWrapper,
)
if _post_call_guardrails_active and isinstance(response, CustomStreamWrapper):
# Intentionally a live reference (not a copy) — mirrors
# ProxyLogging.post_call_success_hook which also mutates
# data["guardrail_to_apply"] during iteration.
_captured_data: Final = self.data
_captured_user_api_key_dict: Final = user_api_key_dict
_captured_logging_obj: Final = logging_obj
async def _on_deferred_stream_complete(assembled_response: object, cache_hit: object) -> None:
await ProxyBaseLLMRequestProcessing._run_deferred_stream_guardrails(
captured_data=_captured_data,
captured_user_api_key_dict=_captured_user_api_key_dict,
captured_logging_obj=_captured_logging_obj,
assembled_response=assembled_response,
cache_hit=cache_hit,
)
logging_obj._on_deferred_stream_complete = _on_deferred_stream_complete
elif (
_post_call_guardrails_active
and route_type == "anthropic_messages"
and self._is_streaming_response(response)
):
from litellm.litellm_core_utils.logging_worker import (
GLOBAL_LOGGING_WORKER,
if _post_call_guardrails_active:
self._arm_deferred_stream_dispatch(
response=response,
route_type=route_type,
user_api_key_dict=user_api_key_dict,
logging_obj=logging_obj,
)
async def _on_deferred_native_stream_complete(
logging_coroutine: Coroutine[object, object, object],
) -> None:
GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(async_coroutine=logging_coroutine)
logging_obj._on_deferred_stream_complete = _on_deferred_native_stream_complete
if route_type == "allm_passthrough_route":
upstream_response_headers: Final = getattr(response, "headers", None)
streaming_headers: Final = (
@ -3135,6 +3096,94 @@ class ProxyBaseLLMRequestProcessing:
except Exception as e:
verbose_proxy_logger.exception("Error firing deferred logging: %s", e)
def _arm_deferred_stream_dispatch(
self,
response: object,
route_type: str,
user_api_key_dict: "UserAPIKeyAuth",
logging_obj: LiteLLMLoggingObj,
) -> None:
"""
Streaming with post-call guardrails active: attach a closure that
ProxyLogging._fire_deferred_stream_logging fires after all guardrail
end-of-stream blocks complete, so the spend log sees
guardrail_information.
Three closure shapes, matching who owns logging for the stream:
- CustomStreamWrapper (chat completions) stores
(assembled_response, cache_hit); the closure also runs
non-apply_guardrail post-call hooks via
_run_deferred_stream_guardrails.
- Bridged /v1/responses (LiteLLMCompletionStreamingIterator) shares
its inner CustomStreamWrapper's logging_obj, so it stores the same
(assembled_response, cache_hit) shape; the closure only dispatches
success logging, matching the route's pre-existing hook surface.
- Native anthropic_messages/aresponses iterators store a single
ready-made logging coroutine to enqueue.
Raw async generators from passthrough routes bypass all three and
would orphan the closure, so they are not armed here.
The router wraps iterators that cannot carry _hidden_params in
HiddenParamsAsyncIteratorWrapper, so class sniffing runs on the
unwrapped inner iterator.
"""
from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
from litellm.router_utils.add_retry_fallback_headers import HiddenParamsAsyncIteratorWrapper
unwrapped: Final = response._inner if isinstance(response, HiddenParamsAsyncIteratorWrapper) else response
if isinstance(unwrapped, CustomStreamWrapper):
# Intentionally a live reference (not a copy) — mirrors
# ProxyLogging.post_call_success_hook which also mutates
# data["guardrail_to_apply"] during iteration.
_captured_data: Final = self.data
_captured_user_api_key_dict: Final = user_api_key_dict
_captured_logging_obj: Final = logging_obj
async def _on_deferred_stream_complete(assembled_response: object, cache_hit: object) -> None:
await ProxyBaseLLMRequestProcessing._run_deferred_stream_guardrails(
captured_data=_captured_data,
captured_user_api_key_dict=_captured_user_api_key_dict,
captured_logging_obj=_captured_logging_obj,
assembled_response=assembled_response,
cache_hit=cache_hit,
)
logging_obj._on_deferred_stream_complete = _on_deferred_stream_complete
return
if route_type not in ("anthropic_messages", "aresponses") or not self._is_streaming_response(response):
return
from litellm.responses.litellm_completion_transformation.streaming_iterator import (
LiteLLMCompletionStreamingIterator,
)
if isinstance(unwrapped, LiteLLMCompletionStreamingIterator):
_captured_bridge_logging_obj: Final = logging_obj
async def _on_deferred_bridged_stream_complete(assembled_response: object, cache_hit: object) -> None:
await _as_success_dispatcher(_captured_bridge_logging_obj).dispatch_success_handlers(
assembled_response,
cache_hit=cache_hit,
start_time=None,
end_time=None,
prefer_async_handlers=True,
)
logging_obj._on_deferred_stream_complete = _on_deferred_bridged_stream_complete
return
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
async def _on_deferred_native_stream_complete(
logging_coroutine: Coroutine[object, object, object],
) -> None:
GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(async_coroutine=logging_coroutine)
logging_obj._on_deferred_stream_complete = _on_deferred_native_stream_complete
@staticmethod
async def _run_deferred_stream_guardrails(
captured_data: dict,

14
litellm/proxy/common_utils/callback_utils.py
View file

@ -1,6 +1,6 @@
import copy
import os
from collections.abc import Callable, Iterable
from collections.abc import Callable, Iterable, Mapping
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any, Final, Literal, NoReturn, Optional, TypeAlias
@ -525,8 +525,8 @@ LITELLM_PROXY_INTERNAL_METADATA_KEYS: Final = frozenset(
def sanitize_openai_provider_metadata(
metadata: dict[str, Any] | None,
) -> dict[str, str] | None:
metadata: Mapping[str, object] | None,
) -> Mapping[str, object] | None:
"""
Keep only provider-safe OpenAI metadata entries (string keys -> string values).
@ -644,7 +644,7 @@ def process_callback(_callback: str, callback_type: str, environment_variables:
return {"name": _callback, "variables": env_vars_dict, "type": callback_type}
def normalize_callback_names(callbacks: Iterable[Any]) -> list[Any]:
def normalize_callback_names(callbacks: Iterable[object] | None) -> list[object]:
if callbacks is None:
return []
return [c.lower() if isinstance(c, str) else c for c in callbacks]
@ -674,7 +674,7 @@ def decrypt_callback_vars(metadata: Any) -> Any:
return _transform_callback_vars(metadata, _decrypt_or_passthrough)
def _transform_callback_vars(metadata: Any, transform: Callable[[str, Any], Any]) -> Any:
def _transform_callback_vars(metadata: object, transform: Callable[[str, Any], Any]) -> object:
if not isinstance(metadata, dict):
return metadata
out: Final = copy.deepcopy(metadata)
@ -704,7 +704,7 @@ def is_sensitive_callback_key(
return _CALLBACK_VAR_MASKER.is_sensitive_key(key)
def _encrypt_if_plaintext(key: str, value: Any) -> Any:
def _encrypt_if_plaintext(key: str, value: object) -> object:
if not isinstance(value, str) or not value:
return value
if not is_sensitive_callback_key(key):
@ -725,7 +725,7 @@ def _encrypt_if_plaintext(key: str, value: Any) -> Any:
return value
def _decrypt_or_passthrough(key: str, value: Any) -> Any:
def _decrypt_or_passthrough(key: str, value: object) -> object:
if not isinstance(value, str) or not value:
return value
if not value.startswith(_CALLBACK_VAR_ENCRYPTED_PREFIX):

48
litellm/proxy/common_utils/user_api_key_cache.py
View file

@ -40,30 +40,31 @@ class UserApiKeyCache(DualCache):
@overload
def get_cache(
self,
key: Any,
parent_otel_span: Any = None,
key: object,
parent_otel_span: object = None,
local_only: bool = False,
*,
model_type: type[T],
**kwargs: Any,
**kwargs: object,
) -> T | None: ...
@overload
def get_cache(
self,
key: Any,
parent_otel_span: Any = None,
key: object,
parent_otel_span: object = None,
local_only: bool = False,
**kwargs: Any,
model_type: None = None,
**kwargs: object,
) -> Any: ...
def get_cache(
self,
key,
parent_otel_span=None,
key: object,
parent_otel_span: object = None,
local_only: bool = False,
model_type: type[BaseModel] | None = None,
**kwargs,
**kwargs: object,
) -> Any | BaseModel | None:
if model_type is None and "model_type" in kwargs:
model_type = cast(type[BaseModel] | None, kwargs.pop("model_type", None))
@ -85,30 +86,31 @@ class UserApiKeyCache(DualCache):
@overload
async def async_get_cache(
self,
key: Any,
parent_otel_span: Any = None,
key: object,
parent_otel_span: object = None,
local_only: bool = False,
*,
model_type: type[T],
**kwargs: Any,
**kwargs: object,
) -> T | None: ...
@overload
async def async_get_cache(
self,
key: Any,
parent_otel_span: Any = None,
key: object,
parent_otel_span: object = None,
local_only: bool = False,
**kwargs: Any,
model_type: None = None,
**kwargs: object,
) -> Any: ...
async def async_get_cache(
self,
key,
parent_otel_span=None,
key: object,
parent_otel_span: object = None,
local_only: bool = False,
model_type: type[BaseModel] | None = None,
**kwargs,
**kwargs: object,
) -> Any | BaseModel | None:
if model_type is None and "model_type" in kwargs:
model_type = cast(type[BaseModel] | None, kwargs.pop("model_type", None))
@ -129,17 +131,17 @@ class UserApiKeyCache(DualCache):
return None
return decoded
def set_cache(self, key, value, local_only: bool = False, **kwargs):
def set_cache(self, key: object, value: object, local_only: bool = False, **kwargs: object):
model_type: Final = cast(type[BaseModel] | None, kwargs.pop("model_type", None))
payload: Final = CacheCodec.serialize(value, model_type=model_type)
payload: Final[object] = CacheCodec.serialize(value, model_type=model_type)
return super().set_cache(key=key, value=payload, local_only=local_only, **kwargs)
async def async_set_cache(self, key, value, local_only: bool = False, **kwargs):
async def async_set_cache(self, key: object, value: object, local_only: bool = False, **kwargs: object):
model_type: Final = cast(type[BaseModel] | None, kwargs.pop("model_type", None))
payload: Final = CacheCodec.serialize(value, model_type=model_type)
payload: Final[object] = CacheCodec.serialize(value, model_type=model_type)
return await super().async_set_cache(key=key, value=payload, local_only=local_only, **kwargs)
async def async_set_cache_pipeline(self, cache_list: list, local_only: bool = False, **kwargs) -> None:
async def async_set_cache_pipeline(self, cache_list: list, local_only: bool = False, **kwargs: object) -> None:
"""
Batch writes with the same Codec boundary as ``async_set_cache`` without
``model_type``: ``BaseModel`` values become JSON-safe dicts; dicts/scalars unchanged.

101
litellm/proxy/db/tool_registry_writer.py
View file

@ -6,9 +6,11 @@ Admins use the management endpoints to read and update input_policy / output_pol
"""
import uuid
from collections.abc import Mapping
from collections.abc import Mapping, Sequence
from datetime import datetime, timezone
from typing import TYPE_CHECKING, Any, Final
from typing import TYPE_CHECKING, Final, Protocol
from pydantic import TypeAdapter
from litellm._logging import verbose_proxy_logger
from litellm.proxy._types import ToolDiscoveryQueueItem
@ -27,6 +29,13 @@ if TYPE_CHECKING:
from litellm.proxy.utils import PrismaClient
class _ModelDumpMethod(Protocol):
def __call__(self) -> Mapping: ...
_ROW_DICT: Final = TypeAdapter(dict)
def _tool_table_actions(prisma_client: "PrismaClient") -> "TableActions[prisma_db_models.LiteLLM_ToolTable]":
table: Final[TableActions[prisma_db_models.LiteLLM_ToolTable]] = ToolRepository(prisma_client).table
return table
@ -41,33 +50,35 @@ def _object_permission_table_actions(
return table
def _row_to_model(row: dict | Any) -> LiteLLM_ToolTableRow:
def _row_to_model(row: object) -> LiteLLM_ToolTableRow:
"""Convert a Prisma model instance or dict to LiteLLM_ToolTableRow."""
model_dump: Final = getattr(row, "model_dump", None)
model_dump: Final[_ModelDumpMethod | None] = getattr(row, "model_dump", None)
if callable(model_dump):
row = model_dump()
elif not isinstance(row, dict):
row = {
k: getattr(row, k, None)
for k in (
"tool_id",
"tool_name",
"origin",
"input_policy",
"output_policy",
"call_count",
"assignments",
"key_hash",
"team_id",
"key_alias",
"user_agent",
"last_used_at",
"created_at",
"updated_at",
"created_by",
"updated_by",
)
}
row = _ROW_DICT.validate_python(
{
k: getattr(row, k, None)
for k in (
"tool_id",
"tool_name",
"origin",
"input_policy",
"output_policy",
"call_count",
"assignments",
"key_hash",
"team_id",
"key_alias",
"user_agent",
"last_used_at",
"created_at",
"updated_at",
"created_by",
"updated_by",
)
}
)
return LiteLLM_ToolTableRow(
tool_id=row.get("tool_id", ""),
tool_name=row.get("tool_name", ""),
@ -190,7 +201,7 @@ async def update_tool_policy(
_updated_by: Final = updated_by or "system"
now: Final = datetime.now(timezone.utc)
create_data: Final[dict[str, object]] = {
create_data: Final[Mapping[str, str | datetime]] = {
"tool_id": str(uuid.uuid4()),
"tool_name": tool_name,
"input_policy": input_policy or "untrusted",
@ -200,14 +211,16 @@ async def update_tool_policy(
"created_at": now,
"updated_at": now,
}
update_data: Final[dict[str, object]] = {
"updated_by": _updated_by,
"updated_at": now,
update_data: Final[Mapping[str, str | datetime]] = {
key: value
for key, value in (
("updated_by", _updated_by),
("updated_at", now),
("input_policy", input_policy),
("output_policy", output_policy),
)
if value is not None
}
if input_policy is not None:
update_data["input_policy"] = input_policy
if output_policy is not None:
update_data["output_policy"] = output_policy
await _tool_table_actions(prisma_client).upsert(
where={"tool_name": tool_name},
@ -338,7 +351,7 @@ class ToolPolicyRegistry:
self._blocked_tools_by_op_id = {}
for row in perms:
op_id = getattr(row, "object_permission_id", None)
blocked = getattr(row, "blocked_tools", None) or []
blocked: Sequence[str] = getattr(row, "blocked_tools", None) or []
if op_id:
self._blocked_tools_by_op_id[op_id] = list(blocked)
@ -370,10 +383,12 @@ class ToolPolicyRegistry:
"""
if not tool_names:
return {}
blocked: Final[set[str]] = set()
for op_id in (object_permission_id, team_object_permission_id):
if op_id and op_id.strip():
blocked.update(self._blocked_tools_by_op_id.get(op_id.strip(), []))
blocked: Final[frozenset[str]] = frozenset(
tool
for op_id in (object_permission_id, team_object_permission_id)
if op_id and op_id.strip()
for tool in self._blocked_tools_by_op_id.get(op_id.strip(), [])
)
result: Final[dict[str, str]] = {}
for name in tool_names:
if name in blocked:
@ -408,13 +423,12 @@ async def add_tool_to_object_permission_blocked(
)
if row is None:
return False
current: Final = list(getattr(row, "blocked_tools", []) or [])
current: Final[Sequence[str]] = getattr(row, "blocked_tools", []) or []
if tool_name in current:
return True
current.append(tool_name)
await _object_permission_table_actions(prisma_client).update(
where={"object_permission_id": object_permission_id},
data={"blocked_tools": current},
data={"blocked_tools": [*current, tool_name]},
)
return True
except Exception as e:
@ -436,13 +450,12 @@ async def remove_tool_from_object_permission_blocked(
)
if row is None:
return False
current = list(getattr(row, "blocked_tools", []) or [])
current: Final[Sequence[str]] = getattr(row, "blocked_tools", []) or []
if tool_name not in current:
return False
current = [t for t in current if t != tool_name]
await _object_permission_table_actions(prisma_client).update(
where={"object_permission_id": object_permission_id},
data={"blocked_tools": current},
data={"blocked_tools": [t for t in current if t != tool_name]},
)
return True
except Exception as e:

53
litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py
View file

@ -18,6 +18,7 @@ import time
from collections.abc import AsyncGenerator, Mapping, Sequence
from datetime import datetime, timezone
from itertools import accumulate, groupby
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, NamedTuple, Optional, cast
import httpx
@ -42,7 +43,7 @@ from litellm.llms.custom_httpx.http_handler import (
httpxSpecialProvider,
)
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.common_request_processing import _serialize_http_exception_detail
from litellm.proxy.common_request_processing import serialize_http_exception_detail
from litellm.proxy.common_utils.sse_keepalive import keepalive_ping_has_fired
from litellm.proxy.guardrails.anthropic_sse import (
anthropic_sse_chunks_from_response,
@ -52,7 +53,12 @@ from litellm.proxy.guardrails.anthropic_sse import (
model_response_text,
)
from litellm.secret_managers.main import get_secret_str
from litellm.types.guardrails import BedrockChecksConfigModel, GuardrailEventHooks
from litellm.types.guardrails import (
BedrockChecksConfigModel,
BedrockGuardrailStreamingParams,
GuardrailEventHooks,
LitellmParams,
)
from litellm.types.llms.openai import AllMessageValues, ChatCompletionUserMessage
from litellm.types.proxy.guardrails.guardrail_hooks.bedrock_guardrails import (
BedrockChecksMessage,
@ -221,9 +227,23 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
prompt_attack_threshold: float | None = 0.5,
pii_confidence_threshold: float | None = 0.5,
chunk_budget_chars: int = BEDROCK_APPLY_GUARDRAIL_CHUNK_BUDGET_CHARS,
streaming_buffer_until_moderated: bool | None = None,
streaming_sampling_rate: int | None = None,
streaming_end_of_stream_only: bool | None = None,
**kwargs,
):
self.async_handler = get_async_httpx_client(llm_provider=httpxSpecialProvider.GuardrailCallback)
self._set_streaming_params(
BedrockGuardrailStreamingParams.from_extras(
MappingProxyType(
{
"streaming_buffer_until_moderated": streaming_buffer_until_moderated,
"streaming_sampling_rate": streaming_sampling_rate,
"streaming_end_of_stream_only": streaming_end_of_stream_only,
}
)
)
)
self.guardrailIdentifier = guardrailIdentifier
self.guardrailVersion = guardrailVersion
self.guardrail_provider = "bedrock"
@ -278,6 +298,18 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
list(self.checks.keys()) if self.checks else None,
)
def _set_streaming_params(self, streaming_params: BedrockGuardrailStreamingParams) -> None:
self.streaming_buffer_until_moderated = streaming_params.streaming_buffer_until_moderated
self.streaming_sampling_rate = streaming_params.streaming_sampling_rate
self.streaming_end_of_stream_only = streaming_params.streaming_end_of_stream_only
def update_in_memory_litellm_params(self, litellm_params: LitellmParams) -> None:
super().update_in_memory_litellm_params(litellm_params)
self._set_streaming_params(BedrockGuardrailStreamingParams.from_extras(litellm_params.model_extra))
def _streams_incrementally(self) -> bool:
return not self.streaming_buffer_until_moderated and not self.mask_response_content
@classmethod
def get_supported_event_hooks(cls) -> list[GuardrailEventHooks]:
return [
@ -2660,6 +2692,21 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
Collect content from the stream and run the bedrock OUTPUT scan
(post_call only validates the response).
"""
if self._streams_incrementally():
from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import (
UnifiedLLMGuardrails,
)
async for streamed_chunk in UnifiedLLMGuardrails().async_post_call_streaming_iterator_hook(
user_api_key_dict=user_api_key_dict,
response=response,
request_data=request_data,
guardrail_to_apply=self,
buffer_until_moderated_default=False,
):
yield streamed_chunk
return
# Import here to avoid circular imports
from litellm.llms.base_llm.base_model_iterator import MockResponseIterator
from litellm.main import stream_chunk_builder
@ -2716,7 +2763,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
)
if not raw_sse or (not is_block and not headers_flushed):
raise
block_message, _ = _serialize_http_exception_detail(block_detail)
block_message, _ = serialize_http_exception_detail(block_detail)
for error_frame in anthropic_sse_error_frames(
block_message if is_block else f"{block_exc.status_code}: {block_message}"
):

6
litellm/proxy/guardrails/guardrail_hooks/compresr/compresr.py
View file

@ -748,7 +748,7 @@ class CompresrGuardrail(CustomGuardrail):
}
try:
raw_response: HttpxResponse | None = await self.async_handler.post( # pyright: ignore[reportUnknownMemberType] # AsyncHTTPHandler.post is untyped
raw_response: HttpxResponse = await self.async_handler.post( # pyright: ignore[reportUnknownMemberType] # AsyncHTTPHandler.post is untyped
url=url,
json=payload,
headers=self._request_headers(),
@ -778,11 +778,11 @@ class CompresrGuardrail(CustomGuardrail):
{"detail": str(e)},
)
return None
if raw_response is None or not 200 <= raw_response.status_code < 300:
if not 200 <= raw_response.status_code < 300:
self._handle_compress_failure(
"Compresr compression service returned an error",
{
"status_code": getattr(raw_response, "status_code", None),
"status_code": raw_response.status_code,
"body": _safe_response_text(raw_response),
},
)

16
litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py
View file

@ -433,7 +433,7 @@ class HeadroomGuardrail(CustomGuardrail):
payload["model"] = model
try:
raw_response: HttpxResponse | None = await self.async_handler.post( # pyright: ignore[reportUnknownMemberType]
raw_response: HttpxResponse = await self.async_handler.post( # pyright: ignore[reportUnknownMemberType]
url=f"{self.headroom_api_base}/v1/compress",
json=payload,
headers=self._request_headers(),
@ -458,16 +458,6 @@ class HeadroomGuardrail(CustomGuardrail):
False,
{},
)
if raw_response is None:
return (
self._handle_compress_failure(
messages,
"Headroom compression service returned no response",
{},
),
False,
{},
)
response: Final[HttpxResponse] = raw_response
if response.status_code != 200:
@ -580,7 +570,7 @@ class HeadroomGuardrail(CustomGuardrail):
params["query"] = query
try:
raw_response: HttpxResponse | None = await self.async_handler.get( # pyright: ignore[reportUnknownMemberType]
raw_response: HttpxResponse = await self.async_handler.get( # pyright: ignore[reportUnknownMemberType]
url=f"{self.headroom_api_base}/v1/retrieve/{hash_value}",
params=params,
headers=self._request_headers(),
@ -589,7 +579,7 @@ class HeadroomGuardrail(CustomGuardrail):
verbose_proxy_logger.warning("Headroom: retrieve failed for hash=%s: %s", hash_value, e)
return f"[Headroom: retrieval failed for hash={hash_value}]"
if raw_response is None or raw_response.status_code == 404:
if raw_response.status_code == 404:
return f"[Headroom: hash={hash_value} not found or expired]"
if raw_response.status_code != 200:

58
litellm/proxy/guardrails/guardrail_hooks/hiddenlayer/hiddenlayer.py
View file

@ -2,7 +2,7 @@ from __future__ import annotations
import os
from collections.abc import Mapping, Sequence
from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, TypedDict
from typing import TYPE_CHECKING, Final, Literal, Protocol
from urllib.parse import urlparse
from uuid import uuid4
@ -11,7 +11,7 @@ import requests
from fastapi import HTTPException
from httpx import HTTPStatusError
from requests.auth import HTTPBasicAuth
from typing_extensions import ReadOnly
from typing_extensions import ReadOnly, TypedDict, Unpack
from litellm._logging import verbose_proxy_logger
from litellm.integrations.custom_guardrail import (
@ -40,14 +40,21 @@ if TYPE_CHECKING:
_AUTH_TIMEOUT_SECONDS: Final[float] = 30.0
class _CustomGuardrailOptions(TypedDict, total=False, extra_items=object):
"""Base-class constructor options carried by this guardrail's forwarded keyword arguments."""
guardrail_name: ReadOnly[str | None]
supported_event_hooks: list[GuardrailEventHooks] | None
class _HiddenlayerEvaluation(TypedDict, total=False):
action: str
threat_level: str
action: ReadOnly[str]
threat_level: ReadOnly[str]
class _HiddenlayerAnalysisEntry(TypedDict, total=False):
name: str
detected: bool
name: ReadOnly[str]
detected: ReadOnly[bool]
class _HiddenlayerModifiedMessage(TypedDict):
@ -59,9 +66,9 @@ class _HiddenlayerModifiedSide(TypedDict):
class _HiddenlayerResponse(TypedDict, total=False):
evaluation: _HiddenlayerEvaluation
analysis: Sequence[_HiddenlayerAnalysisEntry]
modified_data: Mapping[str, _HiddenlayerModifiedSide]
evaluation: ReadOnly[_HiddenlayerEvaluation]
analysis: ReadOnly[Sequence[_HiddenlayerAnalysisEntry]]
modified_data: ReadOnly[Mapping[str, _HiddenlayerModifiedSide]]
class _ProxyServerRequest(TypedDict, total=False):
@ -149,6 +156,31 @@ def _header_value(headers: Mapping[str, str], key: str, default: str) -> str:
return headers.get(key, default)
def _is_image_part(item: object) -> bool:
"""Whether a structured-message content part carries an image rather than text."""
if not isinstance(item, Mapping):
return False
part: Final[Mapping[object, object]] = item
return part.get("type") == "image_url"
def _scannable_text(content: object) -> str:
"""Flatten a structured message's content into the single string the v1 detection endpoint takes.
Image parts are dropped: the endpoint accepts one string, so an image would only reach it as
its stringified source (a base64 blob or a URL), which is not text the scanner can evaluate.
"""
if not isinstance(content, list):
return str(content or "")
parts: Final[Sequence[object]] = content
text_parts: Final = [item for item in parts if not _is_image_part(item)] # mutable-ok: sent as a list repr
return str(text_parts or "")
def is_saas(host: str) -> bool:
"""Checks whether the connection is to the SaaS platform"""
@ -194,7 +226,7 @@ class HiddenlayerGuardrail(CustomGuardrail):
api_key: str | None = None,
api_base: str | None = None,
auth_url: str | None = None,
**kwargs: Any,
**kwargs: Unpack[_CustomGuardrailOptions],
) -> None:
kwargs.setdefault("supported_event_hooks", list(self.get_supported_event_hooks()))
self.hiddenlayer_client_id = api_id or os.getenv("HIDDENLAYER_CLIENT_ID")
@ -263,7 +295,7 @@ class HiddenlayerGuardrail(CustomGuardrail):
"messages": [
{
"role": last_msg.get("role", "user"),
"content": str(last_msg.get("content", "")),
"content": _scannable_text(last_msg.get("content")),
}
]
},
@ -399,7 +431,7 @@ class HiddenlayerGuardrailV2(CustomGuardrail):
api_key: str | None = None,
api_base: str | None = None,
auth_url: str | None = None,
**kwargs: Any,
**kwargs: Unpack[_CustomGuardrailOptions],
) -> None:
self.hiddenlayer_client_id = api_id or os.getenv("HIDDENLAYER_CLIENT_ID")
self.hiddenlayer_client_secret = api_key or os.getenv("HIDDENLAYER_CLIENT_SECRET")
@ -530,7 +562,7 @@ class HiddenlayerGuardrailV2(CustomGuardrail):
self,
payload: _HiddenlayerV2Payload,
input_type: Literal["request", "response"],
hl_headers: dict[str, str],
hl_headers: Mapping[str, str],
) -> httpx.Response:
if input_type == "request":
path = "detection/v2/request-evaluations"

26
litellm/proxy/guardrails/guardrail_hooks/noma/noma_v2.py
View file

@ -7,8 +7,9 @@
import enum
import json
import os
from collections.abc import Callable, Mapping
from datetime import datetime
from typing import TYPE_CHECKING, Any, Final, Literal, Optional, cast
from typing import TYPE_CHECKING, Any, Final, Literal, Optional, TypeAlias, cast
from urllib.parse import urlparse
from litellm._logging import verbose_proxy_logger
@ -36,6 +37,8 @@ _AIDR_SCAN_ENDPOINT: Final = "/litellm/guardrail"
_INTERVENED_INPUT_FIELDS: Final = ("texts", "images", "tools", "tool_calls")
_DEFAULT_API_BASE_HOSTNAME: Final = urlparse(_DEFAULT_API_BASE).hostname
_GuardrailJsonResponse: TypeAlias = Exception | str | dict[str, object]
_KEYS_DUPLICATING_SCAN_INPUTS: Final = ("messages", "input")
_LOGGING_KEYS_DUPLICATING_SCAN_INPUTS: Final = _KEYS_DUPLICATING_SCAN_INPUTS + (
"additional_args",
@ -119,7 +122,7 @@ class NomaV2Guardrail(CustomGuardrail):
def _resolve_action_from_response(
self,
response_json: dict,
response_json: Mapping[str, object],
) -> _Action:
action: Final = response_json.get("action")
if isinstance(action, str):
@ -165,10 +168,11 @@ class NomaV2Guardrail(CustomGuardrail):
@staticmethod
def _sanitize_payload_for_transport(payload: dict) -> dict:
def _default(obj: Any) -> Any:
if hasattr(obj, "model_dump"):
def _default(obj: object) -> object:
model_dump: Final[Callable[[], Mapping[str, object]] | None] = getattr(obj, "model_dump", None)
if model_dump is not None:
try:
return obj.model_dump()
return model_dump()
except Exception:
pass
return str(obj)
@ -178,7 +182,7 @@ class NomaV2Guardrail(CustomGuardrail):
except (ValueError, TypeError):
json_str = safe_dumps(payload)
safe_payload: Final = safe_json_loads(json_str, default={})
safe_payload: Final[object] = safe_json_loads(json_str, default={})
if safe_payload == {} and payload:
verbose_proxy_logger.warning(
"Noma v2 guardrail: payload serialization failed, falling back to empty payload"
@ -196,7 +200,7 @@ class NomaV2Guardrail(CustomGuardrail):
async def _call_noma_scan(
self,
payload: dict,
) -> dict:
) -> dict[str, object]:
headers: Final[dict[str, str]] = {"Content-Type": "application/json"}
authorization_header: Final = self._get_authorization_header()
if authorization_header:
@ -215,7 +219,7 @@ class NomaV2Guardrail(CustomGuardrail):
response.text,
)
response.raise_for_status()
response_json: Final = response.json()
response_json: Final[dict[str, object]] = response.json()
verbose_proxy_logger.debug(
"Noma v2 AIDR response parsed: %s",
json.dumps(response_json, default=str),
@ -227,7 +231,7 @@ class NomaV2Guardrail(CustomGuardrail):
request_data: dict,
start_time: datetime,
guardrail_status: GuardrailStatus,
guardrail_json_response: Any,
guardrail_json_response: _GuardrailJsonResponse,
) -> None:
end_time: Final = datetime.now()
duration: Final = (end_time - start_time).total_seconds()
@ -270,11 +274,11 @@ class NomaV2Guardrail(CustomGuardrail):
) -> GenericGuardrailAPIInputs:
start_time: Final = datetime.now()
guardrail_status: GuardrailStatus = "success"
guardrail_json_response: Any = {}
guardrail_json_response: _GuardrailJsonResponse = {}
dynamic_params = self.get_guardrail_dynamic_request_body_params(request_data)
if not isinstance(dynamic_params, dict):
dynamic_params = {}
response_json: dict | None = None
response_json: dict[str, object] | None = None
# Per-request dynamic params can override configured application context.
application_id = self._get_non_empty_str(dynamic_params.get("application_id"))

1
litellm/proxy/guardrails/guardrail_hooks/prompt_security/__init__.py
View file

@ -20,6 +20,7 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail"
guardrail_name=guardrail.get("guardrail_name", ""),
event_hook=litellm_params.mode,
default_on=litellm_params.default_on,
file_sanitization_fail_open=getattr(litellm_params, "file_sanitization_fail_open", None),
)
litellm.logging_callback_manager.add_litellm_callback(_prompt_security_callback)

49
litellm/proxy/guardrails/guardrail_hooks/prompt_security/prompt_security.py
View file

@ -4,10 +4,12 @@ import os
from collections.abc import Mapping, Sequence
from typing import TYPE_CHECKING, Final, Literal, Optional
import httpx
from fastapi import HTTPException
from typing_extensions import ReadOnly, TypedDict
from litellm._logging import verbose_proxy_logger
from litellm.exceptions import Timeout as LiteLLMTimeout
from litellm.integrations.custom_guardrail import (
CustomGuardrail,
log_guardrail_information,
@ -24,6 +26,9 @@ if TYPE_CHECKING:
from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel
_SANITIZE_FILE_FAIL_OPEN_TIMEOUT_SECONDS: Final = 30.0
class PromptSecurityGuardrailMissingSecrets(Exception):
pass
@ -63,6 +68,13 @@ class _SanitizeStatusResponse(TypedDict, total=False):
metadata: ReadOnly[_SanitizeMetadata]
class _SanitizeResult(TypedDict):
action: ReadOnly[str]
content: ReadOnly[str | None]
metadata: ReadOnly[_SanitizeMetadata]
violations: ReadOnly[Sequence[str]]
class PromptSecurityGuardrail(CustomGuardrail):
@classmethod
def get_supported_event_hooks(cls) -> list[GuardrailEventHooks]:
@ -79,6 +91,8 @@ class PromptSecurityGuardrail(CustomGuardrail):
user: str | None = None,
system_prompt: str | None = None,
check_tool_results: bool | None = None,
file_sanitization_timeout: float = _SANITIZE_FILE_FAIL_OPEN_TIMEOUT_SECONDS,
file_sanitization_fail_open: bool | None = None,
**kwargs,
):
kwargs.setdefault("supported_event_hooks", list(self.get_supported_event_hooks()))
@ -108,6 +122,8 @@ class PromptSecurityGuardrail(CustomGuardrail):
# Configuration for file sanitization
self.max_poll_attempts = 30 # Maximum number of polling attempts
self.poll_interval = 2 # Seconds between polling attempts
self.file_sanitization_timeout = file_sanitization_timeout
self.file_sanitization_fail_open = file_sanitization_fail_open is not False
super().__init__(**kwargs)
@ -397,6 +413,39 @@ class PromptSecurityGuardrail(CustomGuardrail):
Sanitize file content using Prompt Security API.
Returns: dict with keys 'action', 'content', 'metadata'
"""
try:
return await asyncio.wait_for(
self._sanitize_file_content(file_data, filename, user_api_key_alias),
timeout=self.file_sanitization_timeout,
)
except (asyncio.TimeoutError, httpx.TimeoutException, LiteLLMTimeout) as exc:
if not self.file_sanitization_fail_open:
verbose_proxy_logger.error(
"Prompt Security Guardrail: file sanitization for %s timed out with %s; failing closed",
filename,
type(exc).__name__,
)
raise HTTPException(status_code=408, detail="File sanitization timeout") from exc
verbose_proxy_logger.error(
"Prompt Security Guardrail: file sanitization for %s timed out with %s; failing open",
filename,
type(exc).__name__,
)
fail_open_result: Final[_SanitizeResult] = {
"action": "allow",
"content": None,
"metadata": {},
"violations": (),
}
return fail_open_result
async def _sanitize_file_content(
self,
file_data: bytes,
filename: str,
user_api_key_alias: str | None,
) -> _SanitizeResult:
headers: Final = {"APP-ID": self.api_key}
if user_api_key_alias:
headers["X-LiteLLM-Key-Alias"] = user_api_key_alias

5
litellm/proxy/guardrails/guardrail_hooks/repelloai/repelloai.py
View file

@ -197,14 +197,11 @@ class RepelloAIGuardrail(CustomGuardrail):
repelloai_response: RepelloAIAnalyzeResponse | None = None
try:
verbose_proxy_logger.debug("RepelloAI Argus request: %s", request)
raw_response: HttpxResponse | None = await self.async_handler.post( # pyright: ignore[reportUnknownMemberType]
response: Final[HttpxResponse] = await self.async_handler.post( # pyright: ignore[reportUnknownMemberType]
url=endpoint,
headers={"X-API-Key": self.repelloai_api_key},
json=request,
)
if raw_response is None:
raise ValueError("RepelloAI Argus returned no response")
response: Final[HttpxResponse] = raw_response
self._raise_for_config_error(response)
response.raise_for_status()
try:

67
litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py
View file

@ -57,6 +57,9 @@ class _EndpointTranslation(Protocol):
@property
def build_block_sse_chunks(self) -> "Callable[..., Sequence[bytes] | None]": ...
@property
def build_stream_error_items(self) -> "Callable[..., Sequence[object] | None]": ...
def _as_endpoint_translation(translation: _EndpointTranslation) -> _EndpointTranslation:
return translation
@ -408,14 +411,32 @@ class UnifiedLLMGuardrails(CustomLogger):
call_type: str | None,
responses_so_far: Sequence[object],
request_data: dict,
endpoint_translation: _EndpointTranslation | None = None,
stream_started: bool = False,
responses_yielded: Sequence[object] | None = None,
) -> AsyncGenerator[object, None]:
"""Surface a mid-stream HTTPException. For A2A call types the response has
already started, so emit an in-stream JSON-RPC error chunk; otherwise
re-raise so the proxy can report it.
"""Surface a mid-stream HTTPException (a guardrail block with the default
exception-on-block config, or a failed scan).
A2A call types emit an in-stream JSON-RPC error chunk. For other call
types, once chunks have already reached the client the HTTP status is
gone, so the failure is delegated to the endpoint translation's
``build_stream_error_items`` and travels as an in-stream error frame in
that endpoint's wire format. Before the first chunk (or when the format
has no in-stream error frame) the exception is re-raised so the proxy
can report it with a real HTTP status.
"""
if call_type is not None and CallTypes(call_type) in A2A_CALL_TYPES:
yield _a2a_jsonrpc_error_chunk(exc, _get_a2a_request_id(responses_so_far, request_data))
return
if stream_started and endpoint_translation is not None:
error_items: Final = endpoint_translation.build_stream_error_items(
exc, responses_so_far=tuple(responses_yielded) if responses_yielded is not None else None
)
if error_items is not None:
for error_item in error_items:
yield error_item
return
raise exc
def _build_transform_chunk(
@ -586,7 +607,15 @@ class UnifiedLLMGuardrails(CustomLogger):
yield block_chunk
raise _StreamTerminated()
except HTTPException as e:
async for error_item in self._emit_streaming_http_error(e, call_type, responses_so_far, request_data):
async for error_item in self._emit_streaming_http_error(
e,
call_type,
responses_so_far,
request_data,
endpoint_translation=endpoint_translation,
stream_started=bool(responses_yielded),
responses_yielded=responses_yielded,
):
yield error_item
raise _StreamTerminated()
@ -1070,11 +1099,17 @@ class UnifiedLLMGuardrails(CustomLogger):
return
except HTTPException as e:
# Response already started (we already yielded chunks); cannot send 400.
# For A2A, yield an in-stream JSON-RPC error so the client sees it.
if call_type is not None and CallTypes(call_type) in A2A_CALL_TYPES:
yield _a2a_jsonrpc_error_chunk(e, _get_a2a_request_id(responses_so_far, request_data))
return
raise
async for error_item in self._emit_streaming_http_error(
e,
call_type,
responses_so_far,
request_data,
endpoint_translation=endpoint_translation,
stream_started=chunks_yielded,
responses_yielded=responses_yielded,
):
yield error_item
return
chunks_yielded = True
responses_yielded.append(original_item)
yield original_item
@ -1133,7 +1168,13 @@ class UnifiedLLMGuardrails(CustomLogger):
yield block_chunk
return
except HTTPException as e:
if call_type is not None and CallTypes(call_type) in A2A_CALL_TYPES:
yield _a2a_jsonrpc_error_chunk(e, _get_a2a_request_id(responses_so_far, request_data))
else:
raise
async for error_item in self._emit_streaming_http_error(
e,
call_type,
responses_so_far,
request_data,
endpoint_translation=endpoint_translation,
stream_started=bool(responses_yielded),
responses_yielded=responses_yielded,
):
yield error_item

4
litellm/proxy/guardrails/guardrail_initializers.py
View file

@ -11,6 +11,7 @@ def initialize_bedrock(litellm_params: LitellmParams, guardrail: Guardrail):
BedrockGuardrail,
)
streaming_params: Final = BedrockGuardrailStreamingParams.from_extras(litellm_params.model_extra)
_bedrock_callback: Final = BedrockGuardrail(
guardrail_name=guardrail.get("guardrail_name", ""),
event_hook=litellm_params.mode,
@ -38,6 +39,9 @@ def initialize_bedrock(litellm_params: LitellmParams, guardrail: Guardrail):
aws_bedrock_runtime_endpoint=litellm_params.aws_bedrock_runtime_endpoint,
experimental_use_latest_role_message_only=litellm_params.experimental_use_latest_role_message_only,
only_scan_new_messages=litellm_params.only_scan_new_messages or False,
streaming_buffer_until_moderated=streaming_params.streaming_buffer_until_moderated,
streaming_sampling_rate=streaming_params.streaming_sampling_rate,
streaming_end_of_stream_only=streaming_params.streaming_end_of_stream_only,
)
litellm.logging_callback_manager.add_litellm_callback(_bedrock_callback)
return _bedrock_callback

103
litellm/proxy/management_endpoints/internal_user_endpoints.py
View file

@ -17,7 +17,7 @@ import json
import traceback
from collections.abc import Awaitable, Mapping, Sequence
from datetime import datetime, timezone
from typing import Any, Final, Literal, cast
from typing import Any, Final, Literal, Protocol, cast, overload
import fastapi
from fastapi import APIRouter, Depends, Header, HTTPException, Request, status
@ -735,10 +735,44 @@ def _enforce_user_info_access(user_id: str | None, user_api_key_dict: UserAPIKey
)
async def _get_user_info_teams(
prisma_client: Any,
class _UserInfoDataClient(Protocol):
@overload
async def get_data(self, *, user_id: str) -> "prisma_models.LiteLLM_UserTable | None": ...
@overload
async def get_data(
self,
*,
user_id: str | None,
table_name: Literal["key"],
query_type: Literal["find_all"],
) -> "Sequence[LiteLLM_VerificationToken] | None": ...
@overload
async def get_data(
self,
*,
team_id_list: list[str],
table_name: Literal["team"],
query_type: Literal["find_all"],
) -> "Sequence[TeamListResponseObject] | None": ...
async def _get_user_info_keys(
prisma_client: "_UserInfoDataClient",
user_id: str | None,
user_info: Any | None,
) -> "Sequence[LiteLLM_VerificationToken] | None":
return await prisma_client.get_data(
user_id=user_id,
table_name="key",
query_type="find_all",
)
async def _get_user_info_teams(
prisma_client: "_UserInfoDataClient",
user_id: str | None,
user_info: "prisma_models.LiteLLM_UserTable",
user_api_key_dict: UserAPIKeyAuth,
) -> tuple[list[TeamListResponseObject], list[TeamListResponseObject] | None]:
"""Fetch and merge teams from membership + user.teams field."""
@ -759,7 +793,7 @@ async def _get_user_info_teams(
team_list = teams_1
team_id_list = [team.team_id for team in teams_1]
teams_2: list[TeamListResponseObject] | None = None
teams_2: Sequence[TeamListResponseObject] | None = None
target_team_ids: Final = getattr(user_info, "teams", None)
if target_team_ids and isinstance(target_team_ids, list):
@ -769,8 +803,8 @@ async def _get_user_info_teams(
query_type="find_all",
)
elif user_api_key_dict.user_id is not None and user_id is None:
caller_user_info: Final[object] = await prisma_client.get_data(user_id=user_api_key_dict.user_id)
caller_team_ids: Final = getattr(caller_user_info, "teams", None)
caller_user_info: Final = await prisma_client.get_data(user_id=user_api_key_dict.user_id)
caller_team_ids: Final = caller_user_info.teams if caller_user_info is not None else None
if caller_team_ids:
teams_2 = await prisma_client.get_data(
team_id_list=caller_team_ids,
@ -807,7 +841,7 @@ def _redact_scim_enterprise_metadata(
def _build_user_info_response(
user_id: str | None,
user_info: Any | None,
keys: list[LiteLLM_VerificationToken] | None,
keys: Sequence[LiteLLM_VerificationToken] | None,
team_list: list[TeamListResponseObject],
teams_1: list[TeamListResponseObject] | None,
model_max_budget_usage: dict[str, dict[str, object]] | None = None,
@ -894,11 +928,7 @@ async def user_info(
)
## GET ALL KEYS ##
keys: Final = await prisma_client.get_data(
user_id=user_id,
table_name="key",
query_type="find_all",
)
keys: Final = await _get_user_info_keys(prisma_client, user_id)
response_data: Final = _build_user_info_response(
user_id=user_id,
@ -997,6 +1027,14 @@ async def user_info_v2(
This is the v2 replacement for /user/info, designed to avoid the "god endpoint" problem
where the old endpoint loaded all keys and teams into memory.
Note on `spend`: this is the user's running budget counter, which the budget reset job
resets whenever `budget_reset_at` elapses (see `budget_duration`): to zero by default,
or to the overage above `max_budget` when `budget_rollover` is enabled. It is NOT
lifetime or per-period historical spend. For historical spend over a date range, use
`/user/daily/activity` or `/user/daily/activity/aggregated`, which read daily spend
records that only ever accumulate and are never reset. The two values are expected to
diverge once a budget reset has occurred within the queried period.
Access control:
- Proxy admins can query any user
- Team admins can query users within their teams
@ -1077,6 +1115,12 @@ async def user_info_v2(
raise handle_exception_on_proxy(e)
async def _fetch_admin_teams_and_keys_rows(
prisma_client: "PrismaClient", sql_query: str
) -> Sequence[Mapping[str, Sequence[Mapping[str, object]] | None]]:
return await prisma_client.db.query_raw(sql_query)
async def _get_user_info_for_proxy_admin(user_api_key_dict: UserAPIKeyAuth):
"""
Admin UI Endpoint - Returns All Teams and Keys when Proxy Admin is querying
@ -1100,22 +1144,25 @@ async def _get_user_info_for_proxy_admin(user_api_key_dict: UserAPIKeyAuth):
"Database not connected. Connect a database to your proxy - https://docs.litellm.ai/docs/simple_proxy#managing-auth---virtual-keys"
)
results: Final = await prisma_client.db.query_raw(sql_query)
results: Final = await _fetch_admin_teams_and_keys_rows(prisma_client, sql_query)
verbose_proxy_logger.debug("results_keys: %s", results)
_keys_in_db: Final[Sequence[dict[str, object]]] = results[0]["keys"] or []
_keys_in_db: Final[Sequence[Mapping[str, object]]] = results[0]["keys"] or []
# cast all keys to LiteLLM_VerificationToken
keys_in_db: Final = []
for key in _keys_in_db:
if key.get("models") is None:
key["models"] = []
keys_in_db.append(LiteLLM_VerificationToken.model_validate(key))
key_payload = dict[str, object](key)
if key_payload.get("models") is None:
key_payload["models"] = []
keys_in_db.append(LiteLLM_VerificationToken.model_validate(key_payload))
# cast all teams to LiteLLM_TeamTable
_teams_in_db: list[LiteLLM_TeamTable] = results[0]["teams"] or []
_teams_in_db = [LiteLLM_TeamTable.model_validate(team) for team in _teams_in_db]
_teams_in_db.sort(key=lambda x: getattr(x, "team_alias", "") or "")
_teams_rows: Final[Sequence[Mapping[str, object]]] = results[0]["teams"] or []
_teams_in_db: Final = sorted(
(LiteLLM_TeamTable.model_validate(team) for team in _teams_rows),
key=lambda x: getattr(x, "team_alias", "") or "",
)
returned_keys: Final = _process_keys_for_user_info(keys=keys_in_db, all_teams=_teams_in_db)
# Get admin's own user_id and user_info
@ -1140,7 +1187,7 @@ async def _get_user_info_for_proxy_admin(user_api_key_dict: UserAPIKeyAuth):
def _process_keys_for_user_info(
keys: list[LiteLLM_VerificationToken] | None,
keys: Sequence[LiteLLM_VerificationToken] | None,
all_teams: list[LiteLLM_TeamTable] | list[TeamListResponseObject] | None,
):
from litellm.constants import UI_SESSION_TOKEN_TEAM_ID
@ -1231,7 +1278,7 @@ def _update_internal_user_params(data_json: dict, data: UpdateUserRequest | Upda
async def _schedule_user_update_audit_log(
response: dict[str, Any],
response: Mapping[str, object],
existing_user_row: BaseModel | None,
litellm_changed_by: str | None,
user_api_key_dict: UserAPIKeyAuth,
@ -2687,6 +2734,11 @@ async def get_user_daily_activity(
Meant to optimize querying spend data for analytics for a user.
Reads daily spend records that only ever accumulate and are never affected by budget
resets. Their total can legitimately exceed the `spend` field returned by
`/v2/user/info`, which is a running budget counter that every budget reset sets back
to zero (or to the overage above `max_budget` when `budget_rollover` is enabled).
Returns:
(by date)
- spend
@ -2800,6 +2852,11 @@ async def get_user_daily_activity_aggregated(
"""
Aggregated analytics for a user's daily activity without pagination.
Returns the same response shape as the paginated endpoint with page metadata set to single-page.
Reads daily spend records that only ever accumulate and are never affected by budget
resets. Their total can legitimately exceed the `spend` field returned by
`/v2/user/info`, which is a running budget counter that every budget reset sets back
to zero (or to the overage above `max_budget` when `budget_rollover` is enabled).
"""
from litellm.proxy.proxy_server import prisma_client

76
litellm/proxy/management_endpoints/jwt_key_mapping_endpoints.py
View file

@ -1,4 +1,6 @@
from typing import Final
from collections.abc import Mapping, Sequence
from datetime import datetime
from typing import Final, Protocol
from fastapi import APIRouter, Depends, HTTPException, Query
@ -18,7 +20,59 @@ from litellm.repositories.table_repositories import JWTKeyMappingRepository
router: Final = APIRouter()
def _to_response(mapping) -> JWTKeyMappingResponse:
class _JWTKeyMappingRecord(Protocol):
"""A ``LiteLLM_JWTKeyMapping`` row, viewed through the columns these endpoints read."""
@property
def id(self) -> str: ...
@property
def jwt_claim_name(self) -> str: ...
@property
def jwt_claim_value(self) -> str: ...
@property
def description(self) -> str | None: ...
@property
def is_active(self) -> bool: ...
@property
def created_at(self) -> datetime: ...
@property
def updated_at(self) -> datetime: ...
@property
def created_by(self) -> str | None: ...
@property
def updated_by(self) -> str | None: ...
class _JWTKeyMappingTable(Protocol):
"""The Prisma table actions these endpoints issue against the JWT key mapping table."""
async def create(self, *, data: Mapping[str, object]) -> _JWTKeyMappingRecord: ...
async def find_unique(self, *, where: Mapping[str, object]) -> _JWTKeyMappingRecord | None: ...
async def update(self, *, where: Mapping[str, object], data: Mapping[str, object]) -> _JWTKeyMappingRecord: ...
async def delete(self, *, where: Mapping[str, object]) -> _JWTKeyMappingRecord | None: ...
async def find_many(self, *, skip: int, take: int, order: Mapping[str, str]) -> Sequence[_JWTKeyMappingRecord]: ...
async def count(self) -> int: ...
def _mapping_table(prisma_client: object) -> _JWTKeyMappingTable:
"""View the JWT key mapping repository's untyped Prisma table through the actions used here."""
return JWTKeyMappingRepository(prisma_client).table
def _to_response(mapping: _JWTKeyMappingRecord) -> JWTKeyMappingResponse:
"""Convert a Prisma mapping object to a safe response (no hashed token)."""
return JWTKeyMappingResponse(
id=mapping.id,
@ -62,7 +116,7 @@ async def create_jwt_key_mapping(
if data.description is not None:
create_data["description"] = data.description
new_mapping: Final = await JWTKeyMappingRepository(prisma_client).table.create(data=create_data)
new_mapping: Final = await _mapping_table(prisma_client).create(data=create_data)
# Invalidate cache
cache_key: Final = f"jwt_key_mapping:{data.jwt_claim_name}:{data.jwt_claim_value}"
@ -110,7 +164,7 @@ async def update_jwt_key_mapping(
try:
# Get old mapping for cache invalidation
old_mapping: Final = await JWTKeyMappingRepository(prisma_client).table.find_unique(where={"id": data.id})
old_mapping: Final = await _mapping_table(prisma_client).find_unique(where={"id": data.id})
if old_mapping is None:
raise HTTPException(status_code=404, detail="Mapping not found")
@ -118,9 +172,7 @@ async def update_jwt_key_mapping(
cache_key = f"jwt_key_mapping:{old_mapping.jwt_claim_name}:{old_mapping.jwt_claim_value}"
await user_api_key_cache.async_delete_cache(cache_key)
updated_mapping: Final = await JWTKeyMappingRepository(prisma_client).table.update(
where={"id": data.id}, data=update_data
)
updated_mapping: Final = await _mapping_table(prisma_client).update(where={"id": data.id}, data=update_data)
if updated_mapping is None:
raise HTTPException(status_code=404, detail="Mapping not found")
@ -162,7 +214,7 @@ async def delete_jwt_key_mapping(
try:
# Get old mapping for cache invalidation
old_mapping: Final = await JWTKeyMappingRepository(prisma_client).table.find_unique(where={"id": data.id})
old_mapping: Final = await _mapping_table(prisma_client).find_unique(where={"id": data.id})
if old_mapping is None:
raise HTTPException(status_code=404, detail="Mapping not found")
@ -170,7 +222,7 @@ async def delete_jwt_key_mapping(
cache_key: Final = f"jwt_key_mapping:{old_mapping.jwt_claim_name}:{old_mapping.jwt_claim_value}"
await user_api_key_cache.async_delete_cache(cache_key)
await JWTKeyMappingRepository(prisma_client).table.delete(where={"id": data.id})
await _mapping_table(prisma_client).delete(where={"id": data.id})
return {"status": "success"}
except HTTPException:
raise
@ -198,12 +250,12 @@ async def list_jwt_key_mappings(
try:
skip: Final = (page - 1) * size
mappings: Final = await JWTKeyMappingRepository(prisma_client).table.find_many(
mappings: Final = await _mapping_table(prisma_client).find_many(
skip=skip,
take=size,
order={"created_at": "desc"},
)
total_count: Final = await JWTKeyMappingRepository(prisma_client).table.count()
total_count: Final = await _mapping_table(prisma_client).count()
return {
"mappings": [_to_response(m) for m in mappings],
"total_count": total_count,
@ -235,7 +287,7 @@ async def info_jwt_key_mapping(
raise HTTPException(status_code=500, detail="Database not connected")
try:
mapping: Final = await JWTKeyMappingRepository(prisma_client).table.find_unique(where={"id": id})
mapping: Final = await _mapping_table(prisma_client).find_unique(where={"id": id})
if mapping is None:
raise HTTPException(status_code=404, detail="Mapping not found")
return _to_response(mapping)

100
litellm/proxy/management_endpoints/key_management_endpoints.py
View file

@ -18,7 +18,8 @@ import os
import re
import secrets
import traceback
from collections.abc import Awaitable, Callable, Mapping, Sequence
from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence
from contextlib import AbstractAsyncContextManager
from datetime import datetime, timedelta, timezone
from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Protocol, TypeVar, cast
@ -230,6 +231,54 @@ def _config_table(prisma_client: PrismaClient) -> _ConfigTableActions:
)
class _CustomKeyHooksModule(Protocol):
user_custom_key_generate: Callable[..., Awaitable[Mapping[str, object]]] | None
user_custom_key_update: Callable[..., Awaitable[Mapping[str, object]]] | None
def _custom_key_generate_hook(
hooks: _CustomKeyHooksModule,
) -> Callable[..., Awaitable[Mapping[str, object]]] | None:
return hooks.user_custom_key_generate
def _custom_key_update_hook(
hooks: _CustomKeyHooksModule,
) -> Callable[..., Awaitable[Mapping[str, object]]] | None:
return hooks.user_custom_key_update
class _LegacyDumpable(Protocol):
def dict(self) -> Mapping[str, object]: ...
def _legacy_model_dict(row: _LegacyDumpable) -> Mapping[str, object]:
return row.dict()
def _as_object_dict(values: Mapping[str, object]) -> Mapping[str, object]:
return values
def _model_items(model: BaseModel) -> Iterator[tuple[str, object]]:
return iter(model)
class _EnvVarsParam(Protocol):
@property
def param_value(self) -> Mapping[str, str] | None: ...
def _env_vars_param_value(param: _EnvVarsParam) -> Mapping[str, str] | None:
return param.param_value
def _tx_tables_context(
open_tx: Callable[[], AbstractAsyncContextManager[_TxTables]],
) -> AbstractAsyncContextManager[_TxTables]:
return open_tx()
async def _check_custom_key_allowed(custom_key_value: str | None) -> None:
"""Raise 403 if custom API keys are disabled and a custom key was provided."""
if custom_key_value is None:
@ -910,7 +959,7 @@ async def _common_key_generation_helper(
# check if user set default key/generate params on config.yaml
if litellm.default_key_generate_params is not None:
for elem in data:
for elem in _model_items(data):
key, value = elem
if (
value is None
@ -1692,11 +1741,11 @@ async def generate_key_fn(
- user_id: (str) Unique user id - used for tracking spend across multiple keys for same user id.
"""
try:
from litellm.proxy import proxy_server
from litellm.proxy._types import CommonProxyErrors
from litellm.proxy.proxy_server import (
prisma_client,
user_api_key_cache,
user_custom_key_generate,
)
if prisma_client is None:
@ -1723,7 +1772,7 @@ async def generate_key_fn(
)
custom_key_generate_hook: Final[Callable[..., Awaitable[Mapping[str, object]]] | None] = (
user_custom_key_generate
_custom_key_generate_hook(proxy_server)
)
if custom_key_generate_hook is not None:
if inspect.iscoroutinefunction(custom_key_generate_hook):
@ -1892,11 +1941,11 @@ async def generate_service_account_key_fn(
- user_id: (str) Unique user id - used for tracking spend across multiple keys for same user id.
"""
from litellm.proxy import proxy_server
from litellm.proxy._types import CommonProxyErrors
from litellm.proxy.proxy_server import (
prisma_client,
user_api_key_cache,
user_custom_key_generate,
)
if prisma_client is None:
@ -1924,7 +1973,9 @@ async def generate_service_account_key_fn(
verbose_proxy_logger.debug("entered /key/generate")
custom_key_generate_hook: Final[Callable[..., Awaitable[Mapping[str, object]]] | None] = user_custom_key_generate
custom_key_generate_hook: Final[Callable[..., Awaitable[Mapping[str, object]]] | None] = _custom_key_generate_hook(
proxy_server
)
if custom_key_generate_hook is not None:
if inspect.iscoroutinefunction(custom_key_generate_hook):
result: Final = await custom_key_generate_hook(data)
@ -1998,7 +2049,7 @@ def prepare_metadata_fields(data: BaseModel, non_default_values: dict, existing_
)
casted_metadata[reserved_field] = existing_value
data_json: Final[Mapping[str, object]] = data.model_dump(exclude_unset=True, exclude_none=True)
data_json: Final = _as_object_dict(data.model_dump(exclude_unset=True, exclude_none=True))
try:
for k, v in data_json.items():
@ -2805,13 +2856,13 @@ async def update_key_fn(
}'
```
"""
from litellm.proxy import proxy_server
from litellm.proxy.proxy_server import (
llm_router,
premium_user,
prisma_client,
proxy_logging_obj,
user_api_key_cache,
user_custom_key_update,
)
try:
@ -2842,7 +2893,9 @@ async def update_key_fn(
)
# Custom key update hook
custom_key_update_hook: Final[Callable[..., Awaitable[Mapping[str, object]]] | None] = user_custom_key_update
custom_key_update_hook: Final[Callable[..., Awaitable[Mapping[str, object]]] | None] = _custom_key_update_hook(
proxy_server
)
if custom_key_update_hook is not None:
if inspect.iscoroutinefunction(custom_key_update_hook):
result: Final = await custom_key_update_hook(data)
@ -3004,14 +3057,16 @@ async def bulk_update_keys(
}'
```
"""
from litellm.proxy import proxy_server
from litellm.proxy.proxy_server import (
llm_router,
prisma_client,
proxy_logging_obj,
user_api_key_cache,
user_custom_key_update,
)
custom_key_update_hook: Final = _custom_key_update_hook(proxy_server)
if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value:
raise HTTPException(
status_code=403,
@ -3057,7 +3112,7 @@ async def bulk_update_keys(
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
llm_router=llm_router,
user_custom_key_update=user_custom_key_update,
user_custom_key_update=custom_key_update_hook,
)
successful_updates.append(
@ -3135,7 +3190,7 @@ def _build_failed_team_key_update(
if hasattr(existing_key_row, "model_dump"):
key_info = existing_key_row.model_dump()
elif hasattr(existing_key_row, "dict"):
key_info = existing_key_row.dict()
key_info = dict[str, object](_legacy_model_dict(existing_key_row))
if key_info:
key_info.pop("token", None)
@ -3166,14 +3221,16 @@ async def bulk_update_team_keys(
Callable by proxy admins, or by team admins with `KEY_UPDATE` permission.
"""
from litellm.proxy import proxy_server
from litellm.proxy.proxy_server import (
llm_router,
prisma_client,
proxy_logging_obj,
user_api_key_cache,
user_custom_key_update,
)
custom_key_update_hook: Final = _custom_key_update_hook(proxy_server)
if prisma_client is None:
raise HTTPException(
status_code=500,
@ -3302,7 +3359,7 @@ async def bulk_update_team_keys(
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
llm_router=llm_router,
user_custom_key_update=user_custom_key_update,
user_custom_key_update=custom_key_update_hook,
existing_key_row=existing_by_token[db_token],
)
@ -4301,7 +4358,7 @@ def _transform_verification_tokens_to_deleted_records(
"litellm_changed_by": litellm_changed_by,
}
)
record = deleted_record.model_dump()
record = dict[str, object](_as_object_dict(deleted_record.model_dump()))
# Map org_id to organization_id (model uses org_id, but schema expects organization_id)
org_id_value: object = record.pop("org_id", None)
@ -4437,13 +4494,12 @@ async def _rotate_master_key(
should_create_model_in_db=False,
)
if new_model:
_dumped = new_model.model_dump(exclude_none=True)
_dumped = dict[str, object](_as_object_dict(new_model.model_dump(exclude_none=True)))
_dumped["litellm_params"] = prisma.Json(_dumped["litellm_params"])
_dumped["model_info"] = prisma.Json(_dumped["model_info"])
new_models.append(_dumped)
verbose_proxy_logger.debug("Resetting proxy model table")
async with prisma_client.db.tx() as tx_ctx:
tx: Final[_TxTables] = tx_ctx
async with _tx_tables_context(prisma_client.db.tx) as tx:
await tx.litellm_proxymodeltable.delete_many()
verbose_proxy_logger.debug("Creating %s models", len(new_models))
await tx.litellm_proxymodeltable.create_many(
@ -4458,14 +4514,14 @@ async def _rotate_master_key(
if config:
"""If environment_variables is found, decrypt it and encrypt it with the new master key"""
environment_variables_dict = {}
environment_variables_dict: Mapping[str, str] | None = {}
for c in config:
if c.param_name == "environment_variables":
environment_variables_dict = c.param_value
environment_variables_dict = _env_vars_param_value(c)
if environment_variables_dict:
decrypted_env_vars: Final = proxy_config._decrypt_and_set_db_env_variables(
environment_variables=environment_variables_dict
environment_variables=dict[str, str](environment_variables_dict)
)
encrypted_env_vars: Final = proxy_config._encrypt_env_variables(
environment_variables=decrypted_env_vars,
@ -4531,7 +4587,7 @@ async def _rotate_master_key(
updated_patch=decrypted_cred,
new_encryption_key=new_master_key,
)
_cred_data = encrypted_cred.model_dump(exclude_none=True)
_cred_data = dict[str, object](_as_object_dict(encrypted_cred.model_dump(exclude_none=True)))
if "credential_values" in _cred_data:
_cred_data["credential_values"] = prisma.Json(_cred_data["credential_values"])
if "credential_info" in _cred_data:

4
litellm/proxy/management_endpoints/team_callback_endpoints.py
View file

@ -252,7 +252,7 @@ async def add_team_callbacks(
Use this if if you want different teams to have different success/failure callbacks
Parameters:
- callback_name (Literal["langfuse", "langsmith", "gcs"], required): The name of the callback to add
- callback_name (str, required): The name of the callback to add, e.g. "langfuse", "langsmith", "gcs", "newrelic". The value is validated against the callbacks that support team-scoped credentials
- callback_type (Literal["success", "failure", "success_and_failure"], required): The type of callback to add. One of:
- "success": Callback for successful LLM calls
- "failure": Callback for failed LLM calls
@ -268,6 +268,8 @@ async def add_team_callbacks(
- langsmith_api_key: The API key for the Langsmith callback
- langsmith_project: The project for the Langsmith callback
- langsmith_base_url: The base URL for the Langsmith callback
- newrelic_api_key: The ingest license key for the team's New Relic account; routes both LLM/agent traces and cost metrics to that account. Requires the proxy to run with LITELLM_OTEL_V2=true, otherwise this callback is rejected with a 400
- newrelic_region: The New Relic region for the team's account ("us" or "eu"), riding the team's own key
Example curl:
```

41
litellm/proxy/pass_through_endpoints/streaming_handler.py
View file

@ -65,6 +65,19 @@ class PassThroughStreamingHandler:
route_streaming_logging or PassThroughStreamingHandler._route_streaming_logging_to_handler
)
raw_bytes: Final[list[bytes]] = []
def _build_logging_coroutine() -> Coroutine[None, None, None]:
return resolved_route_streaming_logging(
litellm_logging_obj=litellm_logging_obj,
passthrough_success_handler_obj=passthrough_success_handler_obj,
url_route=url_route,
request_body=request_body or {},
endpoint_type=endpoint_type,
start_time=start_time,
raw_bytes=raw_bytes,
end_time=datetime.now(),
)
logging_scheduled = False
model_name: Final = PassThroughStreamingHandler._extract_model_for_cost_injection(
request_body=request_body,
@ -114,6 +127,21 @@ class PassThroughStreamingHandler:
)
if pending:
yield pending
# Stream completed cleanly. When the proxy armed deferred
# dispatch (post-call guardrails active), park the logging
# coroutine on logging_obj instead of enqueueing now, so
# ProxyLogging._fire_deferred_stream_logging fires it after
# guardrail end-of-stream blocks populate guardrail_information.
# Disconnect/exception paths skip this and fall through to the
# immediate enqueue in ``finally`` to keep partial billing
# (LIT-2642).
if (
getattr(litellm_logging_obj, "_on_deferred_stream_complete", None) is not None
and raw_bytes
and response.status_code < 400
):
logging_scheduled = True
litellm_logging_obj._deferred_stream_complete_args = (_build_logging_coroutine(),)
except Exception as e:
verbose_proxy_logger.error("Error in chunk_processor: %s", e)
raise
@ -128,18 +156,7 @@ class PassThroughStreamingHandler:
if not logging_scheduled and raw_bytes and response.status_code < 400:
logging_scheduled = True
try:
GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(
async_coroutine=resolved_route_streaming_logging(
litellm_logging_obj=litellm_logging_obj,
passthrough_success_handler_obj=passthrough_success_handler_obj,
url_route=url_route,
request_body=request_body or {},
endpoint_type=endpoint_type,
start_time=start_time,
raw_bytes=raw_bytes,
end_time=datetime.now(),
)
)
GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(async_coroutine=_build_logging_coroutine())
except Exception as e:
verbose_proxy_logger.error("Error scheduling chunk_processor logging: %s", e)

4
litellm/proxy/prompts/prompt_endpoints.py
View file

@ -6,7 +6,7 @@ import tempfile
from collections.abc import Awaitable, Mapping, Sequence
from datetime import datetime
from pathlib import Path
from typing import TYPE_CHECKING, Any, Final, Protocol, cast
from typing import TYPE_CHECKING, Final, Protocol, cast
from fastapi import (
APIRouter,
@ -1317,7 +1317,7 @@ async def test_prompt(
async def convert_prompt_file_to_json(
file: UploadFile = File(...),
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
) -> dict[str, Any]:
) -> Mapping[str, object]:
"""
Convert a .prompt file to JSON format.

3
litellm/proxy/proxy_server.py
View file

@ -1363,7 +1363,7 @@ _OPENAPI_HTTP_METHODS: Final = {
# the UI. Kept here at module scope to match the analogous descriptor
# `is_secret` flags in litellm.proxy.config_resolvers and the
# `_CACHE_SENSITIVE_FIELDS` constant in the cache endpoint file.
_ALERTING_SENSITIVE_VARS: Final[set[str]] = {"SLACK_WEBHOOK_URL", "SMTP_PASSWORD"}
_ALERTING_SENSITIVE_VARS: Final[set[str]] = {"ALERTING_WEBHOOK_URL", "SLACK_WEBHOOK_URL", "SMTP_PASSWORD"}
def _strip_operation_id_method_suffix(operation_id: str) -> str:
@ -16566,6 +16566,7 @@ async def create_config_audit_log(
_EXTRA_SECRET_CALLBACK_ENV_VARS: Final = frozenset(
{
"ALERTING_WEBHOOK_URL",
"GALILEO_USERNAME",
"GENERIC_LOGGER_HEADERS",
"OTEL_HEADERS",

133
litellm/proxy/response_polling/background_streaming.py
View file

@ -10,12 +10,12 @@ https://platform.openai.com/docs/api-reference/responses-streaming
import asyncio
import json
from collections.abc import Sequence
from typing import TYPE_CHECKING, Final, TypedDict, cast
from collections.abc import Callable, Mapping, Sequence
from typing import TYPE_CHECKING, Final, TypeAlias
from fastapi import Request, Response
from fastapi.responses import StreamingResponse
from typing_extensions import ReadOnly
from typing_extensions import ReadOnly, TypedDict
from litellm._logging import verbose_proxy_logger
from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth
@ -29,23 +29,59 @@ if TYPE_CHECKING:
from litellm.router import Router
class _StreamContentPart(TypedDict, total=False):
text: ReadOnly[str]
_JsonDict: TypeAlias = dict[str, object]
_JsonList: TypeAlias = list[object]
class _StreamOutputItem(TypedDict, total=False):
class _OutputItem(TypedDict, total=False):
id: ReadOnly[str]
content: ReadOnly[Sequence[_StreamContentPart | None]]
content: ReadOnly[Sequence[object]]
class _TerminalResponse(TypedDict, total=False):
status: ReadOnly[ResponsesAPIStatus]
error: ReadOnly[_JsonDict]
usage: ReadOnly[_JsonDict]
reasoning: ReadOnly[_JsonDict]
tool_choice: ReadOnly[object]
tools: ReadOnly[_JsonList]
model: ReadOnly[str]
instructions: ReadOnly[str]
temperature: ReadOnly[float]
top_p: ReadOnly[float]
max_output_tokens: ReadOnly[int]
previous_response_id: ReadOnly[str]
text: ReadOnly[_JsonDict]
truncation: ReadOnly[str]
parallel_tool_calls: ReadOnly[bool]
user: ReadOnly[str]
store: ReadOnly[bool]
incomplete_details: ReadOnly[_JsonDict]
output: ReadOnly[Sequence[_OutputItem]]
class _StreamEvent(TypedDict, total=False):
type: ReadOnly[str]
item: ReadOnly[_OutputItem]
item_id: ReadOnly[str]
content_index: ReadOnly[int]
delta: ReadOnly[str]
part: ReadOnly[object]
response: ReadOnly[_TerminalResponse]
class _StreamEventParser:
parse: Callable[[str], _StreamEvent] = staticmethod(json.loads)
async def background_streaming_task(
polling_id: str,
data,
data: dict,
polling_handler: ResponsePollingHandler,
request: Request,
fastapi_response: Response,
user_api_key_dict: UserAPIKeyAuth,
general_settings,
general_settings: dict,
llm_router: "Router | None",
proxy_config: "ProxyConfig",
proxy_logging_obj: "ProxyLogging",
@ -108,9 +144,10 @@ async def background_streaming_task(
# Process streaming response following OpenAI events format
# https://platform.openai.com/docs/api-reference/responses-streaming
output_items: Final[dict[str, _StreamOutputItem]] = {} # Track output items by ID
# Track accumulated text deltas by (item_id, content_index)
accumulated_text: Final[dict[tuple[str, int], str]] = {}
output_items: Final = dict[str, _OutputItem]() # Track output items by ID
accumulated_text: Final = dict[
tuple[str, int], str
]() # Track accumulated text deltas by (item_id, content_index)
# ResponsesAPIResponse fields to extract from response.completed
usage_data = None
@ -139,7 +176,7 @@ async def background_streaming_task(
None # Will be set by response.completed/failed/incomplete/cancelled
)
terminal_error = None
_event_to_status: Final = {
_event_to_status: Final[Mapping[str, ResponsesAPIStatus]] = {
"response.completed": "completed",
"response.failed": "failed",
"response.incomplete": "incomplete",
@ -180,7 +217,7 @@ async def background_streaming_task(
break
try:
event = json.loads(chunk_data)
event: _StreamEvent = _StreamEventParser.parse(chunk_data)
event_type = event.get("type", "")
# Process different event types based on OpenAI streaming spec
@ -199,19 +236,18 @@ async def background_streaming_task(
if item_id and item_id in output_items:
# Update the output item with new content
current_item = output_items[item_id]
appended_item: _StreamOutputItem = {
**current_item,
"content": (*current_item.get("content", ()), content_part),
added_item = output_items[item_id]
output_items[item_id] = {
**added_item,
"content": (*added_item.get("content", ()), content_part),
}
output_items[item_id] = appended_item
state_dirty = True
elif event_type == "response.output_text.delta":
# Text delta - accumulate text content
# https://platform.openai.com/docs/api-reference/responses-streaming/response-text-delta
item_id = event.get("item_id")
content_index: int = event.get("content_index", 0)
content_index = event.get("content_index", 0)
delta = event.get("delta", "")
if item_id and item_id in output_items:
@ -222,24 +258,14 @@ async def background_streaming_task(
accumulated_text[key] += delta
# Update the content in output_items
current_item = output_items[item_id]
content_list: Sequence[_StreamContentPart | None] = current_item.get("content", ())
if content_index < len(content_list):
# Update existing content part with accumulated text
content_entry = content_list[content_index]
if isinstance(content_entry, dict):
delta_part: _StreamContentPart = {
**content_entry,
"text": accumulated_text[key],
}
delta_item: _StreamOutputItem = {
**current_item,
"content": tuple(
delta_part if index == content_index else entry
for index, entry in enumerate(content_list)
),
}
output_items[item_id] = delta_item
delta_item = output_items[item_id]
if "content" in delta_item:
content_list = delta_item["content"]
if content_index < len(content_list):
# Update existing content part with accumulated text
content_entry = content_list[content_index]
if isinstance(content_entry, dict):
content_entry["text"] = accumulated_text[key]
state_dirty = True
elif event_type == "response.content_part.done":
@ -250,17 +276,17 @@ async def background_streaming_task(
if item_id and item_id in output_items:
# Update with final content from event
current_item = output_items[item_id]
content_list = current_item.get("content", ())
if content_index < len(content_list):
finalized_item: _StreamOutputItem = {
**current_item,
"content": tuple(
content_part if index == content_index else entry
for index, entry in enumerate(content_list)
),
}
output_items[item_id] = finalized_item
done_item = output_items[item_id]
if "content" in done_item:
content_list = done_item["content"]
if content_index < len(content_list):
output_items[item_id] = {
**done_item,
"content": tuple(
content_part if part_index == content_index else existing_part
for part_index, existing_part in enumerate(content_list)
),
}
state_dirty = True
elif event_type == "response.output_item.done":
@ -288,12 +314,9 @@ async def background_streaming_task(
# Terminal event - extract all ResponsesAPIResponse fields
# https://platform.openai.com/docs/api-reference/responses-streaming
response_data = event.get("response", {})
terminal_status = cast(
ResponsesAPIStatus,
response_data.get(
"status",
_event_to_status.get(event_type, "completed"),
),
terminal_status = response_data.get(
"status",
_event_to_status.get(event_type, "completed"),
)
# Extract error for failed and incomplete responses

57
litellm/proxy/search_endpoints/search_tool_registry.py
View file

@ -2,8 +2,9 @@
Search Tool Registry for managing search tool configurations.
"""
from collections.abc import Iterator, Mapping, Sequence
from datetime import datetime, timezone
from typing import Final
from typing import Final, Protocol
from litellm._logging import verbose_proxy_logger
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
@ -13,6 +14,40 @@ from litellm.repositories.table_repositories import SearchToolsRepository
from litellm.types.search import SearchTool
class SearchToolRecord(Protocol):
search_tool_id: str
search_tool_name: str
created_at: datetime
updated_at: datetime
def __iter__(self) -> Iterator[tuple[str, object]]: ...
class SearchToolTableClient(Protocol):
async def create(self, data: Mapping[str, object]) -> SearchToolRecord: ...
async def find_unique(self, where: Mapping[str, object]) -> SearchToolRecord | None: ...
async def find_many(self, order: Mapping[str, str] | None = None) -> Sequence[SearchToolRecord]: ...
async def update(self, where: Mapping[str, object], data: Mapping[str, object]) -> SearchToolRecord: ...
async def delete(self, where: Mapping[str, object]) -> SearchToolRecord: ...
class _SearchToolsRepositoryView(Protocol):
@property
def table(self) -> SearchToolTableClient: ...
def _search_tools_table_of(repository: _SearchToolsRepositoryView) -> SearchToolTableClient:
return repository.table
def _search_tools_table(prisma_client: PrismaClient) -> SearchToolTableClient:
return _search_tools_table_of(SearchToolsRepository(prisma_client))
class SearchToolRegistry:
"""
Handles adding, removing, and getting search tools in DB + in memory.
@ -22,7 +57,7 @@ class SearchToolRegistry:
pass
@staticmethod
def _convert_prisma_to_dict(prisma_obj) -> dict:
def _convert_prisma_to_dict(prisma_obj: SearchToolRecord) -> dict:
"""
Convert Prisma result to dict with datetime objects as ISO format strings.
@ -35,9 +70,9 @@ class SearchToolRegistry:
result: Final = dict(prisma_obj)
# Convert datetime objects to ISO format strings
if "created_at" in result and result["created_at"]:
result["created_at"] = result["created_at"].isoformat()
result["created_at"] = prisma_obj.created_at.isoformat()
if "updated_at" in result and result["updated_at"]:
result["updated_at"] = result["updated_at"].isoformat()
result["updated_at"] = prisma_obj.updated_at.isoformat()
return result
###########################################################
@ -61,7 +96,7 @@ class SearchToolRegistry:
search_tool_info: Final[str] = safe_dumps(search_tool.get("search_tool_info", {}))
# Create search tool in DB
created_search_tool: Final = await SearchToolsRepository(prisma_client).table.create(
created_search_tool: Final = await _search_tools_table(prisma_client).create(
data={
"search_tool_name": search_tool_name,
"litellm_params": litellm_params,
@ -95,7 +130,7 @@ class SearchToolRegistry:
"""
try:
# Get search tool before deletion for response
existing_tool: Final = await SearchToolsRepository(prisma_client).table.find_unique(
existing_tool: Final = await _search_tools_table(prisma_client).find_unique(
where={"search_tool_id": search_tool_id}
)
@ -103,7 +138,7 @@ class SearchToolRegistry:
raise Exception(f"Search tool with ID {search_tool_id} not found")
# Delete from DB
await SearchToolsRepository(prisma_client).table.delete(where={"search_tool_id": search_tool_id})
await _search_tools_table(prisma_client).delete(where={"search_tool_id": search_tool_id})
return {
"message": f"Search tool {search_tool_id} deleted successfully",
@ -131,7 +166,7 @@ class SearchToolRegistry:
search_tool_info: Final[str] = safe_dumps(search_tool.get("search_tool_info", {}))
# Update in DB
updated_search_tool: Final = await SearchToolsRepository(prisma_client).table.update(
updated_search_tool: Final = await _search_tools_table(prisma_client).update(
where={"search_tool_id": search_tool_id},
data={
"search_tool_name": search_tool_name,
@ -163,7 +198,7 @@ class SearchToolRegistry:
try:
search_tools_from_db: Final = await call_with_db_reconnect_retry(
prisma_client,
lambda: SearchToolsRepository(prisma_client).table.find_many(
lambda: _search_tools_table(prisma_client).find_many(
order={"created_at": "desc"},
),
reason="get_all_search_tools_from_db_lookup_failure",
@ -194,7 +229,7 @@ class SearchToolRegistry:
Search tool configuration or None if not found
"""
try:
search_tool: Final = await SearchToolsRepository(prisma_client).table.find_unique(
search_tool: Final = await _search_tools_table(prisma_client).find_unique(
where={"search_tool_id": search_tool_id}
)
@ -222,7 +257,7 @@ class SearchToolRegistry:
Search tool configuration or None if not found
"""
try:
search_tool: Final = await SearchToolsRepository(prisma_client).table.find_unique(
search_tool: Final = await _search_tools_table(prisma_client).find_unique(
where={"search_tool_name": search_tool_name}
)

60
litellm/proxy/spend_tracking/spend_tracking_utils.py
View file

@ -30,7 +30,7 @@ from litellm.litellm_core_utils.litellm_logging import (
request_model_access_groups_from_litellm_params,
)
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps, strip_null_bytes
from litellm.proxy._types import SpendLogsMetadata, SpendLogsPayload
from litellm.proxy._types import SpendLogsMetadata, SpendLogsPayload, SpendLogsRouterMetadata
from litellm.proxy.spend_tracking.spend_log_error_logger import spend_log_error
from litellm.proxy.utils import PrismaClient, hash_token
from litellm.types.utils import (
@ -93,6 +93,24 @@ def _redact_logged_api_key(value: str | None, *, already_redacted: bool = False)
return hash_token(stripped)
def _get_router_metadata_for_spend_log(
metadata: Mapping[str, object] | None,
requested_model: str | None,
selected_model: str | None,
selected_provider: str | None,
router_correlation_id: str | None,
) -> SpendLogsRouterMetadata | None:
model_info: Final = metadata.get("model_info") if metadata is not None else None
if not isinstance(model_info, Mapping) or model_info.get("internal_router_model") is not True:
return None
return SpendLogsRouterMetadata(
requested_model=requested_model or None,
selected_model=selected_model or None,
selected_provider=selected_provider or None,
router_correlation_id=router_correlation_id,
)
def _get_spend_logs_metadata(
metadata: dict | None,
applied_guardrails: list[str] | None = None,
@ -109,6 +127,7 @@ def _get_spend_logs_metadata(
cost_breakdown: CostBreakdown | None = None,
litellm_call_id: str | None = None,
autorouter_savings: float | None = None,
router_metadata: SpendLogsRouterMetadata | None = None,
) -> SpendLogsMetadata:
if metadata is None:
return SpendLogsMetadata(
@ -148,13 +167,17 @@ def _get_spend_logs_metadata(
autorouter_savings=autorouter_savings,
litellm_gateway_injected_cache=None,
litellm_call_id=litellm_call_id,
router_metadata=router_metadata,
)
verbose_proxy_logger.debug(
"getting payload for SpendLogs, available keys in metadata: " + str(list(metadata.keys()))
)
# Filter the metadata dictionary to include only the specified keys
clean_metadata: Final = SpendLogsMetadata(**{key: metadata.get(key) for key in SpendLogsMetadata.__annotations__})
clean_metadata: Final = SpendLogsMetadata(
**{key: metadata.get(key) for key in SpendLogsMetadata.__annotations__ if key != "router_metadata"},
router_metadata=router_metadata,
)
_raw_key: Final = clean_metadata.get("user_api_key")
_trusted_hash: Final = metadata.get("user_api_key_hash")
_already_redacted: Final = (
@ -375,6 +398,20 @@ def get_logging_payload(kwargs, response_obj, start_time, end_time) -> SpendLogs
hidden_params: Final = standard_logging_payload.get("hidden_params", {})
litellm_overhead_time_ms = hidden_params.get("litellm_overhead_time_ms")
custom_llm_provider: Final = (
kwargs.get("custom_llm_provider")
or _sl_attribution_fallback(standard_logging_payload, "custom_llm_provider")
or None
)
raw_model: Final = cast(str, kwargs.get("model") or "")
model_name: Final = (
standard_logging_payload.get("model") if standard_logging_payload is not None else None
) or reconstruct_model_name(raw_model, custom_llm_provider, metadata or {})
litellm_call_id: Final = cast(
str | None,
kwargs.get("litellm_call_id") or litellm_params.get("litellm_call_id"),
)
# clean up litellm metadata
clean_metadata = _get_spend_logs_metadata(
metadata,
@ -433,9 +470,13 @@ def get_logging_payload(kwargs, response_obj, start_time, end_time) -> SpendLogs
autorouter_savings=(
standard_logging_payload.get("autorouter_savings", None) if standard_logging_payload is not None else None
),
litellm_call_id=cast(
str | None,
kwargs.get("litellm_call_id") or litellm_params.get("litellm_call_id"),
litellm_call_id=litellm_call_id,
router_metadata=_get_router_metadata_for_spend_log(
metadata=metadata,
requested_model=_model_group,
selected_model=model_name,
selected_provider=custom_llm_provider,
router_correlation_id=litellm_call_id,
),
)
@ -480,15 +521,6 @@ def get_logging_payload(kwargs, response_obj, start_time, end_time) -> SpendLogs
# Extract agent_id for A2A requests (set directly on model_call_details)
agent_id: Final[str | None] = kwargs.get("agent_id") or metadata.get("agent_id")
custom_llm_provider: Final = (
kwargs.get("custom_llm_provider")
or _sl_attribution_fallback(standard_logging_payload, "custom_llm_provider")
or None
)
raw_model: Final = cast(str, kwargs.get("model") or "")
model_name: Final = (
standard_logging_payload.get("model") if standard_logging_payload is not None else None
) or reconstruct_model_name(raw_model, custom_llm_provider, metadata or {})
try:
payload: Final[SpendLogsPayload] = SpendLogsPayload(

6
litellm/proxy/utils.py
View file

@ -645,7 +645,7 @@ class ProxyLogging:
self.max_parallel_request_limiter = _PROXY_MaxParallelRequestsHandler(self.internal_usage_cache)
self.max_budget_limiter = _PROXY_MaxBudgetLimiter()
self.cache_control_check = _PROXY_CacheControlCheck()
self.alerting: list | None = None
self.alerting: list[str] | None = None
self.alerting_threshold: float = 300 # default to 5 min. threshold
self.alert_types: list[AlertType] = DEFAULT_ALERT_TYPES
self.alert_to_webhook_url: dict | None = None
@ -2364,7 +2364,9 @@ class ProxyLogging:
# do nothing if alerting is not switched on (unless it's a soft_budget alert with team-specific emails)
return
if self.alerting is not None and ("slack" in self.alerting or "ms_teams" in self.alerting):
if self.alerting is not None and (
"slack" in self.alerting or "ms_teams" in self.alerting or "webhook" in self.alerting
):
if self.slack_alerting_instance is not None:
await self.slack_alerting_instance.budget_alerts(
type=type,

65
litellm/rag/rag_query.py
View file

@ -1,11 +1,45 @@
from collections.abc import Sequence
from typing import Any, Final
from typing_extensions import NotRequired, ReadOnly, TypedDict
from litellm.types.llms.openai import AllMessageValues, ChatCompletionUserMessage
from litellm.types.utils import ModelResponse
from litellm.types.vector_stores import (
VectorStoreResultContent,
VectorStoreSearchResponse,
)
from litellm.types.vector_stores import VectorStoreSearchResponse
class _ResultContentView(TypedDict):
"""Content entry carried by a vector store search result."""
type: ReadOnly[NotRequired[str]]
text: ReadOnly[str]
class _SearchResultView(TypedDict):
"""Vector store search result, as far as :class:`RAGQuery` reads it."""
content: ReadOnly[NotRequired[Sequence[_ResultContentView]]]
text: ReadOnly[NotRequired[str]]
class _SearchDataView(TypedDict):
results: ReadOnly[Sequence[_SearchResultView]]
class _ContextChunksView(TypedDict):
chunks: ReadOnly[Sequence[_SearchResultView | str | None]]
class _RerankResultView(TypedDict):
index: ReadOnly[NotRequired[int]]
class _RerankResultsView(TypedDict):
results: ReadOnly[Sequence[_RerankResultView]]
class _MessageView(TypedDict):
message: ReadOnly[object]
class RAGQuery:
@ -42,9 +76,10 @@ class RAGQuery:
"""
context_content = RAGQuery.CONTENT_PREFIX_STRING
for chunk in context_chunks:
chunks: Final[_ContextChunksView] = {"chunks": context_chunks}
for chunk in chunks["chunks"]:
if isinstance(chunk, dict):
result_content: list[VectorStoreResultContent] | None = chunk.get("content")
result_content: Sequence[_ResultContentView] | None = chunk.get("content")
if result_content:
for content_item in result_content:
content_text: str | None = content_item.get("text")
@ -64,14 +99,15 @@ class RAGQuery:
def add_search_results_to_response(
response: ModelResponse,
search_results: VectorStoreSearchResponse,
rerank_results: Any | None = None,
rerank_results: object = None,
) -> ModelResponse:
"""
Add search results to the response choices.
"""
if hasattr(response, "choices") and response.choices:
for choice in response.choices:
message = getattr(choice, "message", None)
message_view: _MessageView = {"message": getattr(choice, "message", None)}
message = message_view["message"]
if message is not None:
# Get existing provider_specific_fields or create new dict
provider_fields = getattr(message, "provider_specific_fields", None) or {}
@ -91,7 +127,8 @@ class RAGQuery:
) -> list[str | dict[str, Any]]:
"""Extract text documents from vector store search response."""
documents: Final[list[str | dict[str, Any]]] = []
for result in search_response.get("data", []):
search_data: Final[_SearchDataView] = {"results": search_response.get("data", [])}
for result in search_data["results"]:
content_list = result.get("content", [])
for content in content_list:
if content.get("type") == "text" and content.get("text"):
@ -99,11 +136,13 @@ class RAGQuery:
return documents
@staticmethod
def get_top_chunks_from_rerank(search_response: Any, rerank_response: Any) -> list[Any]:
def get_top_chunks_from_rerank(search_response: Any, rerank_response: Any) -> list[_SearchResultView]:
"""Get the original search results corresponding to the top reranked results."""
top_chunks: Final = []
original_results: Final = search_response.get("data", [])
for result in rerank_response.get("results", []):
top_chunks: Final[list[_SearchResultView]] = []
search_data: Final[_SearchDataView] = {"results": search_response.get("data", [])}
original_results: Final = search_data["results"]
reranked: Final[_RerankResultsView] = {"results": rerank_response.get("results", [])}
for result in reranked["results"]:
index = result.get("index")
if index is not None and index < len(original_results):
top_chunks.append(original_results[index])

2
litellm/repositories/base_repository.py
View file

@ -40,7 +40,7 @@ def record_to_dict(record: DbRecord) -> Mapping[str, object]:
class BaseRepository(ABC, Generic[T]):
"""Abstract base class for all repositories."""
def __init__(self, prisma_client: Any): # any-ok: PrismaClient is an untyped runtime wrapper
def __init__(self, prisma_client: object):
self._prisma_client = prisma_client
@property

33
litellm/repositories/team_repository.py
View file

@ -3,9 +3,9 @@ Team repository for database operations on LiteLLM_TeamTable.
"""
import json
from collections.abc import Mapping
from collections.abc import Mapping, Sequence
from datetime import datetime
from typing import TYPE_CHECKING, Final
from typing import TYPE_CHECKING, Final, Protocol
from pydantic import TypeAdapter
@ -21,6 +21,25 @@ if TYPE_CHECKING:
from prisma import Prisma
from prisma import models as prisma_models
class _TeamArrays(Protocol):
"""The string array columns of a team row, which the domain model leaves untyped."""
@property
def members(self) -> Sequence[str]: ...
@property
def admins(self) -> Sequence[str]: ...
@property
def models(self) -> Sequence[str]: ...
def _team_arrays(team: LiteLLM_TeamTable) -> _TeamArrays:
"""View a team's untyped list columns as sequences of ids."""
return team
_MEMBERS_WITH_ROLES_ADAPTER: Final = TypeAdapter(list[Member])
_JSON_ENCODED_TEAM_FIELDS: Final = (
"metadata",
@ -80,8 +99,8 @@ class TeamRepository(BaseRepository[LiteLLM_TeamTable]):
)
if not rows:
return None
raw_value: Final = rows[0]["members_with_roles"]
parsed: Final = json.loads(raw_value) if isinstance(raw_value, str) else raw_value
raw_value: Final[object] = rows[0]["members_with_roles"]
parsed: Final[object] = json.loads(raw_value) if isinstance(raw_value, str) else raw_value
if not parsed:
return []
return _MEMBERS_WITH_ROLES_ADAPTER.validate_python(parsed)
@ -315,7 +334,7 @@ class TeamRepository(BaseRepository[LiteLLM_TeamTable]):
if team is None:
return None
members: Final = [m for m in team.members if m != user_id]
members: Final = [m for m in _team_arrays(team).members if m != user_id]
return await self.update(team_id, {"members": members}, id_field="team_id")
async def add_admin(self, team_id: str, user_id: str) -> LiteLLM_TeamTable | None:
@ -340,7 +359,7 @@ class TeamRepository(BaseRepository[LiteLLM_TeamTable]):
if team is None:
return None
admins: Final = [a for a in team.admins if a != user_id]
admins: Final = [a for a in _team_arrays(team).admins if a != user_id]
return await self.update(team_id, {"admins": admins}, id_field="team_id")
async def add_models(self, team_id: str, models: list[str]) -> LiteLLM_TeamTable | None:
@ -365,5 +384,5 @@ class TeamRepository(BaseRepository[LiteLLM_TeamTable]):
if team is None:
return None
current_models: Final = [m for m in team.models if m not in models]
current_models: Final = [m for m in _team_arrays(team).models if m not in models]
return await self.update(team_id, {"models": current_models}, id_field="team_id")

226
litellm/responses/streaming_iterator.py
View file

@ -5,11 +5,11 @@ import json
import time
import traceback
import uuid
from collections.abc import Awaitable, Callable, Mapping, Sequence
from collections.abc import Awaitable, Callable, Iterable, Mapping, Sequence
from datetime import datetime
from functools import lru_cache
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, runtime_checkable
from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, overload, runtime_checkable
import httpx
from openai._streaming import SSEDecoder
@ -42,27 +42,14 @@ from litellm.types.utils import CallTypes
from litellm.utils import async_post_call_success_deployment_hook
if TYPE_CHECKING:
from litellm.caching.caching_handler import LLMCachingHandler
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.responses.streaming_websocket import (
PresidioGuardrailCallback,
ResponsesBackendWebSocket,
ResponsesClientWebSocket,
)
class _StreamCachingHandler(Protocol):
"""The ``_llm_caching_handler`` attached to a logging object, as this module uses it."""
original_function: Callable[..., object]
def _should_store_result_in_cache(
self, original_function: Callable[..., object], kwargs: Mapping[str, object]
) -> bool: ...
class PiiUnmaskingGuardrailCallback(PresidioGuardrailCallback, Protocol):
"""Guardrail callback that can also reverse its own masking, selected by
``llm_http_handler`` on exactly this attribute."""
def _unmask_pii_text(self, text: str, pii_tokens: Mapping[str, str]) -> str: ...
from litellm.types.router import LiteLLM_Params
class ProjectQuotaCallback(Protocol):
@ -94,6 +81,60 @@ def _is_str_mapping(value: object) -> TypeIs[dict[str, str]]: # guard-ok: verif
return _is_json_object(value) and all(isinstance(item, str) for item in value.values())
class _MutableJsonObject(Protocol):
@overload
def get(self, key: str, /) -> object | None: ...
@overload
def get(self, key: str, default: object, /) -> object: ...
def __getitem__(self, key: str, /) -> object: ...
def __setitem__(self, key: str, value: object, /) -> None: ...
def __contains__(self, key: object, /) -> bool: ...
def items(self) -> Iterable[tuple[str, object]]: ...
class _GetsLitellmParams(Protocol):
def __call__(self, key: str, default: Mapping[str, object], /) -> LiteLLM_Params: ...
class _PopsOptionalStr(Protocol):
def __call__(self, key: str, default: None, /) -> str | None: ...
class _UnmasksPiiText(Protocol):
def __call__(self, text: str, pii_tokens: Mapping[str, str]) -> str: ...
class _ShouldStoreResultInCache(Protocol):
def __call__(self, *, original_function: Callable[..., object] | None, kwargs: Mapping[str, object]) -> bool: ...
class _PostStreamingDeploymentHook(Protocol):
def __call__(
self,
*,
request_data: Mapping[str, object],
response_chunk: ResponsesAPIStreamingResponse,
call_type: CallTypes | None,
) -> Awaitable[ResponsesAPIStreamingResponse | None]: ...
@runtime_checkable
class _HasPostStreamingDeploymentHook(Protocol):
async_post_call_streaming_deployment_hook: _PostStreamingDeploymentHook
def _typed_gets_litellm_params(fn: _GetsLitellmParams) -> _GetsLitellmParams:
return fn
def _typed_pops_optional_str(fn: _PopsOptionalStr) -> _PopsOptionalStr:
return fn
_SHOULD_STORE_RESULT_IN_CACHE_ATTR: Final = "_should_store_result_in_cache"
_UNMASK_PII_TEXT_ATTR: Final = "_unmask_pii_text"
def _load_json_object(payload: str | bytes) -> dict[str, object]:
"""Parse a JSON payload that the caller consumes as an object."""
return json.loads(payload)
@ -220,7 +261,7 @@ class BaseResponsesAPIStreamingIterator:
# This matches the stream wrapper in litellm/litellm_core_utils/streaming_handler.py
_api_base: Final = get_api_base(
model=model or "",
optional_params=self.logging_obj.model_call_details.get("litellm_params", {}),
optional_params=_typed_gets_litellm_params(self.logging_obj.model_call_details.get)("litellm_params", {}),
)
self._hidden_params: dict[str, object] = {
"model_id": _model_id_from_metadata(litellm_metadata),
@ -422,15 +463,20 @@ class BaseResponsesAPIStreamingIterator:
end_time: Final = datetime.now()
if is_async:
asyncio.create_task(
self.logging_obj.dispatch_success_handlers(
logging_response,
start_time=self.start_time,
end_time=end_time,
cache_hit=self._completed_response_cache_hit,
prefer_async_handlers=True,
)
logging_coroutine: Final = self.logging_obj.dispatch_success_handlers(
logging_response,
start_time=self.start_time,
end_time=end_time,
cache_hit=self._completed_response_cache_hit,
prefer_async_handlers=True,
)
deferred_dispatch_armed: Final = getattr(self.logging_obj, "_on_deferred_stream_complete", None) is not None
if deferred_dispatch_armed:
# End-of-stream guardrail scans write guardrail_information after
# the terminal event; dispatching now would snapshot metadata early.
self.logging_obj._deferred_stream_complete_args = (logging_coroutine,)
else:
asyncio.create_task(logging_coroutine)
else:
run_async_function(
async_function=self.logging_obj.async_success_handler,
@ -549,7 +595,7 @@ class BaseResponsesAPIStreamingIterator:
if response_obj is None:
return
caching_handler: Final[_StreamCachingHandler | None] = getattr(self.logging_obj, "_llm_caching_handler", None)
caching_handler: Final[LLMCachingHandler | None] = getattr(self.logging_obj, "_llm_caching_handler", None)
if caching_handler is None:
return
@ -567,8 +613,11 @@ class BaseResponsesAPIStreamingIterator:
if preset_cache_key is not None:
request_kwargs["cache_key"] = preset_cache_key
if not caching_handler._should_store_result_in_cache( # pyright: ignore[reportPrivateUsage] # no public API
original_function=caching_handler.original_function,
should_store_result_in_cache: Final[_ShouldStoreResultInCache] = getattr(
caching_handler, _SHOULD_STORE_RESULT_IN_CACHE_ATTR
)
if not should_store_result_in_cache(
original_function=getattr(caching_handler, "original_function", None),
kwargs=request_kwargs,
):
return
@ -624,12 +673,15 @@ class BaseResponsesAPIStreamingIterator:
typed_call_type = None
request_data: Final = self.request_data or getattr(self.logging_obj, "model_call_details", {})
callbacks: Final = getattr(litellm, "callbacks", None) or []
callbacks: Final[Sequence[object]] = getattr(litellm, "callbacks", None) or []
hooks_ran = False
for callback in callbacks:
if hasattr(callback, "async_post_call_streaming_deployment_hook"):
if isinstance(callback, _HasPostStreamingDeploymentHook):
hooks_ran = True
result = await callback.async_post_call_streaming_deployment_hook(
post_streaming_hook: _PostStreamingDeploymentHook = (
callback.async_post_call_streaming_deployment_hook
)
result = await post_streaming_hook(
request_data=request_data,
response_chunk=chunk,
call_type=typed_call_type,
@ -1083,7 +1135,7 @@ class _HasModelDumpJson(Protocol):
def model_dump_json(self, *, exclude_none: bool = ...) -> str: ...
def _dump_response_object(obj: object) -> dict[str, Any]:
def _dump_response_object(obj: object) -> Mapping[str, object]:
if isinstance(obj, _HasModelDump):
return obj.model_dump()
if _is_json_object(obj):
@ -1113,21 +1165,20 @@ def _build_content_part_done_event(
item_id: str,
output_index: int,
content_index: int,
part_payload: dict[str, Any],
part_payload: Mapping[str, object],
) -> ResponsesAPIStreamingResponse | None:
openai_types: Final = _get_openai_response_types()
part_type: Final = part_payload.get("type")
part: PART_UNION_TYPES
if part_type == "output_text":
annotations: Final = [
openai_types.BaseLiteLLMOpenAIResponseObject(**annotation)
for annotation in part_payload.get("annotations", []) or []
]
part = openai_types.ContentPartDonePartOutputText(
type="output_text",
text=str(part_payload.get("text") or ""),
annotations=annotations,
logprobs=part_payload.get("logprobs"),
raw_annotations: Final[object] = part_payload.get("annotations", []) or []
part = openai_types.ContentPartDonePartOutputText.model_validate(
{
"type": "output_text",
"text": str(part_payload.get("text") or ""),
"annotations": raw_annotations,
"logprobs": part_payload.get("logprobs"),
}
)
elif part_type == "refusal":
part = openai_types.ContentPartDonePartRefusal(
@ -1157,7 +1208,7 @@ def _add_text_like_part_events(
item_id: str,
output_index: int,
content_index: int,
part_payload: dict[str, Any],
part_payload: Mapping[str, object],
chunk_size: int,
) -> None:
openai_types: Final = _get_openai_response_types()
@ -1174,16 +1225,19 @@ def _add_text_like_part_events(
delta=text[i : i + chunk_size],
)
)
annotations_payload: Final[Sequence[dict[str, object]]] = part_payload.get("annotations", []) or []
for annotation_index, annotation in enumerate(annotations_payload):
raw_annotation_items: Final = part_payload.get("annotations")
annotation_items: Final[Sequence[object]] = raw_annotation_items if _is_json_array(raw_annotation_items) else []
for annotation_index, annotation in enumerate(annotation_items):
events.append(
openai_types.OutputTextAnnotationAddedEvent(
type=openai_types.ResponsesAPIStreamEvents.OUTPUT_TEXT_ANNOTATION_ADDED,
item_id=item_id,
output_index=output_index,
content_index=content_index,
annotation_index=annotation_index,
annotation=annotation,
openai_types.OutputTextAnnotationAddedEvent.model_validate(
{
"type": openai_types.ResponsesAPIStreamEvents.OUTPUT_TEXT_ANNOTATION_ADDED,
"item_id": item_id,
"output_index": output_index,
"content_index": content_index,
"annotation_index": annotation_index,
"annotation": annotation,
}
)
)
events.append(
@ -1256,7 +1310,8 @@ def _build_synthetic_response_events(
)
if item_type == "message":
content_parts: Sequence[object] = output_item_payload.get("content", []) or []
raw_content_parts = output_item_payload.get("content")
content_parts: Sequence[object] = raw_content_parts if _is_json_array(raw_content_parts) else []
for content_index, part in enumerate(content_parts):
part_payload = _dump_response_object(part)
events.append(
@ -1304,8 +1359,9 @@ def _build_synthetic_response_events(
)
)
elif item_type == "reasoning":
summaries: Sequence[object] = output_item_payload.get("summary", []) or []
for summary_index, summary in enumerate(summaries):
raw_summary_items = output_item_payload.get("summary")
summary_items: Sequence[object] = raw_summary_items if _is_json_array(raw_summary_items) else []
for summary_index, summary in enumerate(summary_items):
summary_payload = _dump_response_object(summary)
summary_text = str(summary_payload.get("text") or "")
for i in range(0, len(summary_text), chunk_size):
@ -1476,7 +1532,7 @@ class ResponsesWebSocketStreaming:
user_api_key_dict: UserAPIKeyAuth | None = None,
request_data: dict[str, object] | None = None,
first_message: str | None = None,
guardrail_callbacks: list[PiiUnmaskingGuardrailCallback] | None = None,
guardrail_callbacks: Sequence[PresidioGuardrailCallback] | None = None,
output_guardrail_callbacks: list[PresidioGuardrailCallback] | None = None,
quota_callbacks: Sequence[ProjectQuotaCallback] | None = None,
authorized_model: str | None = None,
@ -1486,17 +1542,17 @@ class ResponsesWebSocketStreaming:
self.logging_obj = logging_obj
self.user_api_key_dict = user_api_key_dict
self.request_data: dict[str, object] = request_data or {}
self.messages: list[dict[str, object]] = []
self.messages: list[_MutableJsonObject] = []
self.input_messages: list[dict[str, object]] = []
self.first_message = first_message
self.guardrail_callbacks: list[PiiUnmaskingGuardrailCallback] = guardrail_callbacks or []
self.guardrail_callbacks: Sequence[PresidioGuardrailCallback] = guardrail_callbacks or []
self.output_guardrail_callbacks: list[PresidioGuardrailCallback] = output_guardrail_callbacks or []
self.quota_callbacks: tuple[ProjectQuotaCallback, ...] = tuple(quota_callbacks) if quota_callbacks else ()
# Model name authorized at connection time; enforced on every
# response.create frame to prevent deployment-substitution attacks.
self.authorized_model: str | None = authorized_model
def _should_store_event(self, event_obj: Mapping[str, object]) -> bool:
def _should_store_event(self, event_obj: _MutableJsonObject) -> bool:
return event_obj.get("type") in RESPONSES_WS_LOGGED_EVENT_TYPES
def _store_event(self, event: str | bytes | dict[str, object]) -> None:
@ -1610,7 +1666,7 @@ class ResponsesWebSocketStreaming:
finally:
await self._log_messages()
def _enforce_authorized_model(self, msg_obj: dict[str, object]) -> bool:
def _enforce_authorized_model(self, msg_obj: _MutableJsonObject) -> bool:
"""
Overwrite any ``model`` field in a ``response.create`` frame with the
connection-authorized model to prevent deployment-substitution attacks.
@ -1679,7 +1735,7 @@ class ResponsesWebSocketStreaming:
# forwarded unmasked regardless of where the client places it.
nested_candidate = msg_obj.get("response")
nested_response = nested_candidate if _is_json_object(nested_candidate) else None
text_containers: list[tuple[dict[str, object], str]] = []
text_containers: list[tuple[_MutableJsonObject, str]] = []
for container in (msg_obj, nested_response):
if container is None:
continue
@ -1786,6 +1842,7 @@ class ResponsesWebSocketStreaming:
return response_str
cb: Final = self.guardrail_callbacks[0]
unmask_pii_text: Final[_UnmasksPiiText] = getattr(cb, _UNMASK_PII_TEXT_ATTR)
event_type: Final = evt_obj.get("type")
if event_type == "response.completed":
@ -1805,9 +1862,7 @@ class ResponsesWebSocketStreaming:
continue
text = content_block.get("text")
if isinstance(text, str):
unmasked = cb._unmask_pii_text( # pyright: ignore[reportPrivateUsage] # no public unmasker
text, pii_tokens
)
unmasked = unmask_pii_text(text, pii_tokens)
if unmasked != text:
content_block["text"] = unmasked
modified = True
@ -1816,9 +1871,7 @@ class ResponsesWebSocketStreaming:
if event_type in self._DELTA_EVENT_TYPES:
delta: Final = evt_obj.get("delta")
if isinstance(delta, str):
unmasked = cb._unmask_pii_text( # pyright: ignore[reportPrivateUsage] # no public unmasker
delta, pii_tokens
)
unmasked = unmask_pii_text(delta, pii_tokens)
if unmasked != delta:
evt_obj["delta"] = unmasked
return json.dumps(evt_obj)
@ -2020,7 +2073,7 @@ class ManagedResponsesWebSocketHandler:
model: str,
logging_obj: LiteLLMLoggingObj,
user_api_key_dict: UserAPIKeyAuth | None = None,
litellm_metadata: dict[str, Any] | None = None,
litellm_metadata: Mapping[str, object] | None = None,
api_key: str | None = None,
api_base: str | None = None,
timeout: float | None = None,
@ -2033,10 +2086,11 @@ class ManagedResponsesWebSocketHandler:
self.model = model
self.logging_obj = logging_obj
self.user_api_key_dict = user_api_key_dict
self.litellm_metadata: dict[str, Any] = litellm_metadata or {}
self.model_group: str | None = self.litellm_metadata.get("model_group") or self.litellm_metadata.get(
self.litellm_metadata: Mapping[str, object] = litellm_metadata or {}
raw_model_group: Final = self.litellm_metadata.get("model_group") or self.litellm_metadata.get(
"deployment_model_name"
)
self.model_group: str | None = raw_model_group if isinstance(raw_model_group, str) else None
self.api_key = api_key
self.api_base = api_base
self.timeout = timeout
@ -2057,7 +2111,7 @@ class ManagedResponsesWebSocketHandler:
# ------------------------------------------------------------------
@staticmethod
def _serialize_chunk(chunk: Any) -> str | None:
def _serialize_chunk(chunk: object) -> str | None:
"""Serialize a streaming chunk to a JSON string for WebSocket transmission."""
try:
if isinstance(chunk, _HasModelDumpJson):
@ -2100,7 +2154,7 @@ class ManagedResponsesWebSocketHandler:
self._session_history[response_id] = messages
@staticmethod
def _extract_response_id(completed_event: dict[str, object]) -> str | None:
def _extract_response_id(completed_event: _MutableJsonObject) -> str | None:
"""
Pull the raw (decoded) response ID out of a ``response.completed`` event.
Returns *None* if the event doesn't contain a usable ID.
@ -2115,7 +2169,7 @@ class ManagedResponsesWebSocketHandler:
@staticmethod
def _extract_output_messages(
completed_event: dict[str, object],
completed_event: _MutableJsonObject,
) -> list[dict[str, object]]:
"""
Convert the output items in a ``response.completed`` event into
@ -2172,7 +2226,7 @@ class ManagedResponsesWebSocketHandler:
# _process_response_create sub-methods
# ------------------------------------------------------------------
async def _parse_message(self, raw_message: str) -> dict[str, object] | None:
async def _parse_message(self, raw_message: str) -> _MutableJsonObject | None:
"""Parse raw WS text; return the message dict or None (JSON error / ignored type)."""
try:
msg_obj: Final = _load_json_object(raw_message)
@ -2185,7 +2239,7 @@ class ManagedResponsesWebSocketHandler:
return msg_obj
@staticmethod
def _is_warmup_frame(msg_obj: dict[str, object]) -> bool:
def _is_warmup_frame(msg_obj: _MutableJsonObject) -> bool:
"""Return True for a response.create whose generate flag is false."""
nested: Final = msg_obj.get("response")
source: Final = nested if _is_json_object(nested) and nested else msg_obj
@ -2201,13 +2255,13 @@ class ManagedResponsesWebSocketHandler:
return str(raw_id).startswith(_WARMUP_RESPONSE_ID_PREFIX)
@staticmethod
def _warmup_source_params(msg_obj: dict[str, object]) -> dict[str, object]:
def _warmup_source_params(msg_obj: _MutableJsonObject) -> dict[str, object]:
nested: Final = msg_obj.get("response")
if _is_json_object(nested) and nested:
return nested
return {k: v for k, v in msg_obj.items() if k != "type"}
def _build_warmup_response(self, msg_obj: dict[str, object]) -> dict[str, object]:
def _build_warmup_response(self, msg_obj: _MutableJsonObject) -> dict[str, object]:
"""Build a minimal completed Responses API object for a warmup ack."""
source: Final = self._warmup_source_params(msg_obj)
wire_model: Final = source.get("model") or self.model_group or self.model
@ -2225,7 +2279,7 @@ class ManagedResponsesWebSocketHandler:
},
}
async def _send_warmup_ack(self, msg_obj: dict[str, object]) -> None:
async def _send_warmup_ack(self, msg_obj: _MutableJsonObject) -> None:
"""
Acknowledge a generate=false prewarm without calling the provider.
@ -2248,7 +2302,7 @@ class ManagedResponsesWebSocketHandler:
await self.websocket.send_text(serialized)
@staticmethod
def _build_base_call_kwargs(msg_obj: dict[str, object]) -> dict[str, Any]:
def _build_base_call_kwargs(msg_obj: _MutableJsonObject) -> dict[str, Any]:
"""
Extract Responses API params from the event, handling both wire formats:
Nested: {"type": "response.create", "response": {"input": [...], ...}}
@ -2357,7 +2411,7 @@ class ManagedResponsesWebSocketHandler:
call_kwargs.setdefault("litellm_params", {})
call_kwargs["litellm_params"]["proxy_server_request"] = proxy_server_request
async def _stream_and_forward(self, model: str, call_kwargs: dict[str, Any]) -> dict[str, object] | None:
async def _stream_and_forward(self, model: str, call_kwargs: dict[str, Any]) -> _MutableJsonObject | None:
"""
Stream ``litellm.aresponses`` and forward every chunk over the WebSocket.
@ -2365,7 +2419,7 @@ class ManagedResponsesWebSocketHandler:
directly (before serialization) to avoid a redundant JSON round-trip on
every chunk. Returns the completed event dict, or ``None``.
"""
completed_event: dict[str, object] | None = (
completed_event: _MutableJsonObject | None = (
None # rebind-ok: captures the completed event once the stream yields it
)
stream_response: Final = await litellm.aresponses(model=model, **call_kwargs)
@ -2391,7 +2445,7 @@ class ManagedResponsesWebSocketHandler:
def _save_turn_history(
self,
completed_event: dict[str, object] | None,
completed_event: _MutableJsonObject | None,
prior_history: list[dict[str, object]],
current_messages: list[dict[str, object]],
) -> None:
@ -2464,12 +2518,14 @@ class ManagedResponsesWebSocketHandler:
# reuse the router-resolved self.model; passing the alias raw to
# litellm.aresponses fails in get_llm_provider. A genuinely different
# provider-prefixed per-frame model is still honored.
requested_model: Final[str | None] = call_kwargs.pop("model", None)
requested_model: Final[str | None] = _typed_pops_optional_str(call_kwargs.pop)("model", None)
model: Final[str] = (
self.model if requested_model is None or requested_model == self.model_group else requested_model
)
previous_response_id: Final[str | None] = call_kwargs.pop("previous_response_id", None)
previous_response_id: Final[str | None] = _typed_pops_optional_str(call_kwargs.pop)(
"previous_response_id", None
)
current_messages: Final = self._input_to_messages(call_kwargs.get("input"))
# Fetch history once; reused in both _apply_history and _save_turn_history

212
litellm/router_strategy/tag_based_routing.py
View file

@ -8,16 +8,14 @@ Use this to route requests between Teams
"""
import re
from collections.abc import Iterable, Mapping, Sequence
from collections.abc import Mapping, Sequence
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, Literal, TypedDict
from typing_extensions import ReadOnly
from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, overload
from litellm._logging import verbose_logger
from litellm.constants import CONSUMED_REQUEST_TAGS_METADATA_KEY
from litellm.litellm_core_utils.core_helpers import get_metadata_variable_name_from_kwargs
from litellm.types.router import ConsumedRequestTagsStamp, RouterErrors
from litellm.types.router import ConsumedRequestTagsStamp, DeploymentTypedDict, RouterErrors
if TYPE_CHECKING:
from litellm.router import Router as _Router
@ -27,34 +25,63 @@ else:
LitellmRouter = Any
class _TagRoutingLitellmParams(TypedDict, total=False):
tags: ReadOnly[Sequence[str] | None]
tag_regex: ReadOnly[Sequence[str] | None]
class _TagLitellmParamsLike(Protocol):
@overload
def get(self, key: Literal["tags"], /) -> Sequence[str] | None: ...
@overload
def get(self, key: Literal["tags"], default: Sequence[str], /) -> Sequence[str]: ...
@overload
def get(self, key: Literal["tag_regex"], /) -> Sequence[str] | None: ...
class _TagRoutingDeployment(TypedDict, total=False):
model_name: ReadOnly[str]
litellm_params: ReadOnly[_TagRoutingLitellmParams]
model_info: ReadOnly[Mapping[str, object] | None]
class _ModelInfoLike(Protocol):
@overload
def get(self, key: Literal["allow_fail_open"], /) -> bool | None: ...
@overload
def get(self, key: Literal["enable_tag_filtering"], /) -> bool | None: ...
class _TagRoutingMatchStamp(TypedDict):
matched_deployment: ReadOnly[str | None]
matched_via: ReadOnly[str]
matched_value: ReadOnly[str]
request_tags: ReadOnly[Sequence[str]]
user_agent: ReadOnly[str]
class _DeploymentLike(Protocol):
@overload
def get(self, key: Literal["litellm_params"], default: Mapping[str, object], /) -> _TagLitellmParamsLike: ...
@overload
def get(self, key: Literal["model_info"], /) -> _ModelInfoLike | None: ...
@overload
def get(self, key: Literal["model_name"], /) -> object: ...
class _TagRoutingMetadata(TypedDict, total=False):
tags: ReadOnly[Sequence[str] | None]
inherited_tags: ReadOnly[Sequence[str] | None]
user_agent: ReadOnly[str]
tag_routing: ReadOnly[_TagRoutingMatchStamp]
_consumed_request_tags: ReadOnly[object]
class _MetadataLike(Protocol):
@overload
def get(self, key: Literal["tags"], /) -> Sequence[str] | None: ...
@overload
def get(self, key: Literal["tags"], default: Sequence[str], /) -> Sequence[str] | None: ...
@overload
def get(self, key: Literal["user_agent"], default: str, /) -> str: ...
@overload
def get(self, key: Literal["inherited_tags"], /) -> object: ...
def __contains__(self, key: object, /) -> bool: ...
def __setitem__(self, key: Literal["tag_routing"], value: Mapping[str, object], /) -> None: ...
_EMPTY_MODEL_INFO: Final[Mapping[str, object]] = MappingProxyType({})
class _NestedLitellmParamsLike(Protocol):
def get(
self, key: Literal["metadata", "litellm_metadata"], default: Mapping[str, object], /
) -> _MetadataLike | None: ...
class _RequestKwargsLike(Protocol):
@overload
def get(self, key: Literal["enable_tag_filtering"], /) -> bool | None: ...
@overload
def get(self, key: Literal["metadata", "litellm_metadata"], /) -> _MetadataLike | None: ...
def __contains__(self, key: object, /) -> bool: ...
@overload
def __getitem__(self, key: Literal["metadata", "litellm_metadata"], /) -> _MetadataLike: ...
@overload
def __getitem__(self, key: Literal["litellm_params"], /) -> _NestedLitellmParamsLike: ...
_DeploymentPool = Sequence[_DeploymentLike] | Mapping[_DeploymentLike, object]
def _is_valid_deployment_tag_regex(
@ -109,11 +136,11 @@ def is_valid_deployment_tag(
def _match_deployment(
deployment: _TagRoutingDeployment,
request_tags: Sequence[str] | None,
header_strings: Sequence[str],
deployment: _DeploymentLike,
request_tags: list[str] | None,
header_strings: list[str],
match_any: bool,
) -> Mapping[str, str] | None:
) -> dict[str, str] | None:
"""
Determine whether *deployment* matches the current request.
@ -198,38 +225,38 @@ def _split_tags(tags: Sequence[str]) -> tuple[tuple[str, ...], list[str], tuple[
def _exclude_deployments(
deployments: Iterable[_TagRoutingDeployment],
deployments: _DeploymentPool,
excluded_set: frozenset[str],
) -> list[_TagRoutingDeployment]:
) -> Sequence[_DeploymentLike]:
if not excluded_set:
return list(deployments)
return [d for d in deployments if not excluded_set.intersection(d.get("litellm_params", {}).get("tags") or [])]
def _require_all_tags(
deployments: Iterable[_TagRoutingDeployment],
deployments: _DeploymentPool,
required_set: frozenset[str],
) -> tuple[_TagRoutingDeployment, ...]:
) -> tuple[_DeploymentLike, ...]:
if not required_set:
return tuple(deployments)
return tuple(d for d in deployments if required_set.issubset(d.get("litellm_params", {}).get("tags") or []))
def _default_tagged_pool(
deployments: Iterable[_TagRoutingDeployment],
) -> tuple[_TagRoutingDeployment, ...]:
deployments: _DeploymentPool,
) -> tuple[_DeploymentLike, ...]:
defaults: Final = tuple(d for d in deployments if "default" in (d.get("litellm_params", {}).get("tags") or []))
return defaults if defaults else tuple(deployments)
def _known_tag_values(deployments: Iterable[_TagRoutingDeployment]) -> frozenset[str]:
def _known_tag_values(deployments: _DeploymentPool) -> frozenset[str]:
return frozenset(
tag for d in deployments for tag in (d.get("litellm_params", _TagRoutingLitellmParams()).get("tags") or ())
tag for d in deployments for tag in (d.get("litellm_params", MappingProxyType({})).get("tags") or ())
)
def _unknown_required_tag_hides_an_answer(
healthy_deployments: Iterable[_TagRoutingDeployment],
healthy_deployments: _DeploymentPool,
excluded_set: frozenset[str],
required_set: frozenset[str],
routing_confirmed: frozenset[str],
@ -253,23 +280,23 @@ def _unknown_required_tag_hides_an_answer(
def _chain_allows_fail_open(
healthy_deployments: Iterable[_TagRoutingDeployment],
healthy_deployments: _DeploymentPool,
excluded_set: frozenset[str],
required_set: frozenset[str],
routing_confirmed: frozenset[str],
) -> bool:
if _unknown_required_tag_hides_an_answer(healthy_deployments, excluded_set, required_set, routing_confirmed):
return False
return any((d.get("model_info") or _EMPTY_MODEL_INFO).get("allow_fail_open") is True for d in healthy_deployments)
return any((d.get("model_info") or {}).get("allow_fail_open") is True for d in healthy_deployments)
def _trusted_only_pool(
healthy_deployments: Iterable[_TagRoutingDeployment],
healthy_deployments: _DeploymentPool,
excluded_set: frozenset[str],
required_set: frozenset[str],
inherited_excluded_set: frozenset[str] | None,
inherited_required_set: frozenset[str] | None,
) -> tuple[_TagRoutingDeployment, ...]:
) -> tuple[_DeploymentLike, ...]:
# inherited_*_set is None only when this request carries no origin information
# at all (e.g. direct SDK Router usage, bypassing the proxy layer that
# populates metadata.inherited_tags) -- treat every constraint as
@ -296,8 +323,8 @@ def _trusted_only_pool(
def _resolve_or_fail_open(
pool: Sequence[_TagRoutingDeployment],
healthy_deployments: Iterable[_TagRoutingDeployment],
pool: Sequence[_DeploymentLike],
healthy_deployments: _DeploymentPool,
excluded_set: frozenset[str],
required_set: frozenset[str],
inherited_excluded_set: frozenset[str] | None,
@ -305,7 +332,7 @@ def _resolve_or_fail_open(
routing_confirmed: frozenset[str],
model: str,
request_tags: object,
) -> tuple[_TagRoutingDeployment, ...]:
) -> tuple[_DeploymentLike, ...]:
if pool:
return tuple(pool)
if _chain_allows_fail_open(healthy_deployments, excluded_set, required_set, routing_confirmed):
@ -325,7 +352,7 @@ def _resolve_or_fail_open(
def _resolve_constraint_only_pool(
healthy_deployments: Iterable[_TagRoutingDeployment],
healthy_deployments: _DeploymentPool,
excluded_set: frozenset[str],
required_set: frozenset[str],
inherited_excluded_set: frozenset[str] | None,
@ -333,7 +360,7 @@ def _resolve_constraint_only_pool(
routing_confirmed: frozenset[str],
model: str,
request_tags: object,
) -> tuple[_TagRoutingDeployment, ...]:
) -> tuple[_DeploymentLike, ...]:
pool: Final = (
_require_all_tags(_exclude_deployments(healthy_deployments, excluded_set), required_set)
if required_set
@ -355,8 +382,8 @@ def _resolve_constraint_only_pool(
def _all_deployments_or_fallback(
llm_router_instance: LitellmRouter,
model: str,
fallback: Iterable[_TagRoutingDeployment],
) -> Iterable[_TagRoutingDeployment]:
fallback: _DeploymentPool,
) -> Sequence[_DeploymentLike | DeploymentTypedDict] | Mapping[_DeploymentLike, object]:
try:
return llm_router_instance._get_all_deployments(model_name=model)
except Exception: # noqa: BLE001 # fail safe toward today's healthy-only behavior on lookup errors
@ -366,8 +393,8 @@ def _all_deployments_or_fallback(
def _chain_tag_filtering_override(
llm_router_instance: LitellmRouter,
model: str,
healthy_deployments: Iterable[_TagRoutingDeployment],
) -> object:
healthy_deployments: _DeploymentPool,
) -> bool | None:
# Resolved from every deployment configured for this model group, not just the
# ones that survived cooldown/health filtering (async_get_healthy_deployments
# filters cooldowns before calling get_deployments_for_tag) -- otherwise the
@ -379,14 +406,14 @@ def _chain_tag_filtering_override(
# than crashing the request.
all_deployments: Final = _all_deployments_or_fallback(llm_router_instance, model, healthy_deployments)
for d in all_deployments:
value = (d.get("model_info") or _EMPTY_MODEL_INFO).get("enable_tag_filtering")
value = (d.get("model_info") or MappingProxyType({})).get("enable_tag_filtering")
if value is not None:
return value
return None
def _inherited_constraint_sets(
inherited_tags: Sequence[str] | None, routing_prefix: str
inherited_tags: object, routing_prefix: str
) -> tuple[frozenset[str] | None, frozenset[str] | None]:
# None means no origin information is available at all (e.g. this request
# bypassed the proxy layer that populates metadata.inherited_tags, as direct
@ -417,43 +444,42 @@ def _tag_known_to_group(
if tag_set & routing_confirmed:
return True
try:
all_deployments: Final[Sequence[_TagRoutingDeployment]] = llm_router_instance._get_all_deployments(
model_name=model
)
all_deployments: Final = llm_router_instance._get_all_deployments(model_name=model)
except Exception: # noqa: BLE001 # fail safe toward "unrecognized" so lookup errors preserve the existing silent-fallback behavior
return False
return any(
tag_set.intersection(d.get("litellm_params", _TagRoutingLitellmParams()).get("tags") or ())
for d in all_deployments
tag_set.intersection(d.get("litellm_params", MappingProxyType({})).get("tags") or ()) for d in all_deployments
)
def _request_tags_after_router_consumption(metadata: _TagRoutingMetadata, model: str) -> Sequence[str] | None:
def _request_tags_after_router_consumption(metadata: object, model: str) -> Sequence[str] | None:
# The pre-routing hook stamps which tags selected the router it rewrote the request
# to: those tags already did their job and must not also constrain deployment choice
# inside the routed group. The request's other tags still apply there, on top of the
# inherited_tags snapshot that keeps key/team policy applying. Every other model
# group keeps the full list.
stamp: Final = metadata.get(CONSUMED_REQUEST_TAGS_METADATA_KEY)
if not isinstance(metadata, Mapping):
return None
typed_metadata: Final[Mapping[str, object]] = metadata
request_tags: Final = _tags_in_metadata(typed_metadata)
stamp: Final = typed_metadata.get(CONSUMED_REQUEST_TAGS_METADATA_KEY)
if not isinstance(stamp, ConsumedRequestTagsStamp) or stamp.model_group != model:
return metadata.get("tags")
request_tags: Final = metadata.get("tags")
leftover: Final = tuple(
tag for tag in (request_tags if isinstance(request_tags, (list, tuple)) else ()) if tag not in stamp.tags
)
inherited_tags: Final = metadata.get("inherited_tags")
return request_tags
leftover: Final = tuple(tag for tag in request_tags if tag not in stamp.tags)
inherited_tags: Final = typed_metadata.get("inherited_tags")
if not isinstance(inherited_tags, (list, tuple)):
return leftover or None
return tuple(dict.fromkeys((*leftover, *inherited_tags)))
typed_inherited_tags: Final[Sequence[object]] = inherited_tags
return tuple(dict.fromkeys((*leftover, *(tag for tag in typed_inherited_tags if isinstance(tag, str)))))
async def get_deployments_for_tag(
llm_router_instance: LitellmRouter,
model: str, # used to raise the correct error
healthy_deployments: list[Any] | dict[Any, Any],
request_kwargs: dict[Any, Any] | None = None,
healthy_deployments: _DeploymentPool,
request_kwargs: _RequestKwargsLike | None = None,
metadata_variable_name: Literal["metadata", "litellm_metadata"] = "metadata",
):
) -> _DeploymentPool:
"""
Returns a list of deployments that match the requested model and tags in the request.
@ -486,8 +512,7 @@ async def get_deployments_for_tag(
verbose_logger.debug("request metadata: %s", request_kwargs.get(metadata_variable_name))
if metadata_variable_name in request_kwargs:
metadata: Final[_TagRoutingMetadata] = request_kwargs[metadata_variable_name]
stampable_metadata: Final[dict[str, object]] = request_kwargs[metadata_variable_name]
metadata: Final = request_kwargs[metadata_variable_name]
request_tags: Final = _request_tags_after_router_consumption(metadata, model)
match_any: Final = llm_router_instance.tag_filtering_match_any
routing_prefix: Final = llm_router_instance.tag_routing_prefix or ""
@ -532,25 +557,25 @@ async def get_deployments_for_tag(
request_tags,
)
new_healthy_deployments: Final[list[_TagRoutingDeployment]] = []
default_deployments: Final[list[_TagRoutingDeployment]] = []
if has_positive_filter:
verbose_logger.debug(
"get_deployments_for_tag routing: request_tags=%s user_agent=%s",
request_tags,
user_agent,
)
for deployment in candidates:
deployment_tags = deployment.get("litellm_params", {}).get("tags")
match_result = _match_deployment(
deployment=deployment,
request_tags=positive_tags,
header_strings=header_strings,
match_any=match_any,
deployment_matches: Final = tuple(
(
deployment,
_match_deployment(
deployment=deployment,
request_tags=positive_tags,
header_strings=header_strings,
match_any=match_any,
),
)
for deployment in candidates
)
for deployment, match_result in deployment_matches:
if match_result is not None:
verbose_logger.debug(
"tag routing match: deployment=%s matched_via=%s matched_value=%s",
@ -559,17 +584,17 @@ async def get_deployments_for_tag(
match_result["matched_value"],
)
if "tag_routing" not in metadata:
stampable_metadata["tag_routing"] = {
metadata["tag_routing"] = {
"matched_deployment": deployment.get("model_name"),
"matched_via": match_result["matched_via"],
"matched_value": match_result["matched_value"],
"request_tags": request_tags or [],
"user_agent": user_agent,
}
new_healthy_deployments.append(deployment)
if deployment_tags and "default" in deployment_tags:
default_deployments.append(deployment)
new_healthy_deployments: Final = [d for d, result in deployment_matches if result is not None]
default_deployments: Final = [
d for d, _ in deployment_matches if "default" in (d.get("litellm_params", {}).get("tags") or ())
]
if len(new_healthy_deployments) == 0 and len(default_deployments) == 0:
return _resolve_or_fail_open(
@ -604,10 +629,11 @@ async def get_deployments_for_tag(
return new_healthy_deployments if len(new_healthy_deployments) > 0 else default_deployments
# for Untagged requests use default deployments if set
_default_deployments_with_tags: Final[list[_TagRoutingDeployment]] = []
for deployment in healthy_deployments:
if "default" in deployment.get("litellm_params", {}).get("tags", []):
_default_deployments_with_tags.append(deployment)
_default_deployments_with_tags: Final = [
deployment
for deployment in healthy_deployments
if "default" in deployment.get("litellm_params", {}).get("tags", [])
]
if len(_default_deployments_with_tags) > 0:
return _default_deployments_with_tags

36
litellm/types/guardrails.py
View file

@ -1,5 +1,7 @@
from collections.abc import Mapping
from datetime import datetime
from enum import Enum
from types import MappingProxyType
from typing import Any, Final, Literal
from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator
@ -550,6 +552,40 @@ class BedrockGuardrailConfigModel(BaseModel):
)
class BedrockGuardrailStreamingParams(BaseModel):
streaming_buffer_until_moderated: bool = Field(
default=True,
description="If True (default), withhold every streamed chunk until the end-of-stream "
"ApplyGuardrail scan passes, so no flagged content reaches the client before a block. "
"If False, chunks stream through unbuffered, so flagged content can reach the client "
"before the scan finishes; a flagged scan still ends the stream, with a block message "
"when disable_exception_on_block is true and an in-stream error frame otherwise.",
)
streaming_sampling_rate: int = Field(
default=5,
ge=1,
description="When not buffering and not end-of-stream-only, scan the accumulated response "
"every Nth streamed chunk. Each sampled scan is a full ApplyGuardrail call that delays "
"that chunk, so lower values add latency and AWS text-unit cost.",
)
streaming_end_of_stream_only: bool = Field(
default=False,
description="When not buffering, skip per-chunk sampling and run one ApplyGuardrail scan "
"on the assembled response at end of stream. Combined with "
"streaming_buffer_until_moderated=false the full response streams live before the scan "
"and the scan result lands in guardrail_information; a flagged response still ends the "
"stream with a block message (disable_exception_on_block=true) or an error frame.",
)
@classmethod
def from_extras(cls, extras: Mapping[str, object] | None) -> "BedrockGuardrailStreamingParams":
if not extras:
return cls()
return cls.model_validate(
MappingProxyType({name: extras[name] for name in cls.model_fields if extras.get(name) is not None})
)
class LakeraV2GuardrailConfigModel(BaseModel):
"""Configuration parameters for the Lakera AI v2 guardrail"""

4
litellm/types/proxy/guardrails/guardrail_hooks/prompt_security.py
View file

@ -12,6 +12,10 @@ class PromptSecurityGuardrailConfigModel(GuardrailConfigModel):
default=None,
description="The API base for the Prompt Security guardrail. If not provided, the `PROMPT_SECURITY_API_BASE` environment variable is used.",
)
file_sanitization_fail_open: bool = Field(
default=True,
description="Whether file sanitization timeouts allow the original file through instead of blocking the request.",
)
@staticmethod
def ui_friendly_name() -> str:

5
litellm/types/router.py
View file

@ -189,6 +189,11 @@ class ModelInfo(MirroredPricingParams):
# router-wide default.
enable_tag_filtering: bool | None = None
# when True, calls routed to this deployment persist a router_metadata block
# (requested model group, selected model + provider, router correlation id)
# in the spend log row's metadata. Set it on every deployment of the group.
internal_router_model: bool | None = None
def __init__(self, id: str | int | None = None, **params) -> None:
if id is None:
id = str(uuid.uuid4()) # Generate a UUID if id is None or not provided

267
model_prices_and_context_window.json
View file

File diff suppressed because it is too large Load diff

22
ruff-strict-budget.json
View file

@ -1,21 +1,21 @@
{
"ANN001": {
"limit": 3012
"limit": 2995
},
"ANN002": {
"limit": 71
},
"ANN003": {
"limit": 827
"limit": 809
},
"ANN201": {
"limit": 2003
"limit": 2002
},
"ANN202": {
"limit": 845
},
"ANN204": {
"limit": 702
"limit": 698
},
"ANN205": {
"limit": 112
@ -24,7 +24,7 @@
"limit": 133
},
"ANN401": {
"limit": 654
"limit": 587
},
"ASYNC230": {
"limit": 11
@ -123,7 +123,7 @@
"limit": 12
},
"PERF403": {
"limit": 34
"limit": 33
},
"PIE804": {
"limit": 18
@ -177,7 +177,7 @@
"limit": 8
},
"RUF019": {
"limit": 32
"limit": 31
},
"RUF046": {
"limit": 4
@ -195,10 +195,10 @@
"limit": 22
},
"SIM101": {
"limit": 58
"limit": 56
},
"SIM102": {
"limit": 315
"limit": 314
},
"SIM103": {
"limit": 119
@ -231,7 +231,7 @@
"limit": 5
},
"TID251": {
"limit": 1116
"limit": 1108
},
"TRY002": {
"limit": 524
@ -246,7 +246,7 @@
"limit": 113
},
"TRY300": {
"limit": 857
"limit": 855
},
"UP028": {
"limit": 2

348
scripts/auto-close-duplicates.test.ts Normal file
View file

@ -0,0 +1,348 @@
import { describe, expect, test } from "bun:test";
import {
CLOSED_MARKER,
REOPEN_COMMENT,
candidateNumbers,
duplicateTarget,
normalizeTitle,
pendingNotice,
readConfig,
reopenTarget,
sweepClosedIssue,
sweepIssue,
type Comment,
type GitHubApi,
type Issue,
type Reaction,
type SweepConfig,
} from "./auto-close-duplicates";
const NOW = new Date("2026-09-04T09:00:00Z");
const DAY_MS = 24 * 60 * 60 * 1000;
const daysAgo = (days: number): string => new Date(NOW.getTime() - days * DAY_MS).toISOString();
const issue = (number: number, title: string, overrides: Partial<Issue> = {}): Issue => ({
number,
title,
state: "open",
user: { login: "reporter" },
...overrides,
});
const notice = (candidates: readonly number[], createdAt: string, overrides: Partial<Comment> = {}): Comment => ({
id: 900,
body: `<!-- litellm:potential-duplicate candidates=${candidates.join(",")}, -->\n**Potential duplicate detected**`,
created_at: createdAt,
user: { type: "Bot", login: "github-actions[bot]" },
...overrides,
});
const humanComment = (createdAt: string, body = "It is not the same thing", login = "reporter"): Comment => ({
id: 901,
body,
created_at: createdAt,
user: { type: "User", login },
});
const config: SweepConfig = { repo: "BerriAI/litellm", graceDays: 3, dryRun: false, now: NOW };
describe("normalizeTitle", () => {
test("drops the template prefix, case, and punctuation", () => {
expect(normalizeTitle("[Bug]: Gemma 4-e4b fails on Vertex!")).toBe("gemma 4 e4b fails on vertex");
expect(normalizeTitle("[Feature]: ")).toBe("");
});
});
describe("candidateNumbers", () => {
test("reads only the marker field, keeps older issues, sorted ascending and deduplicated", () => {
const body = "<!-- litellm:potential-duplicate candidates=40,10,30,10, -->\n- #1 - see #1 (100% similar)";
expect(candidateNumbers(body, 35)).toEqual([10, 30]);
});
test("returns nothing without the marker", () => {
expect(candidateNumbers("- #1 - looks like #1", 35)).toEqual([]);
});
});
describe("pendingNotice", () => {
test("waits out the grace period from the latest notice", () => {
const fresh = pendingNotice(issue(35, "t"), [notice([10], daysAgo(2.9))], config);
expect(fresh.kind).toBe("skip");
const aged = pendingNotice(issue(35, "t"), [notice([10], daysAgo(3.1))], config);
expect(aged.kind).toBe("pending");
const reposted = pendingNotice(
issue(35, "t"),
[notice([10], daysAgo(6)), notice([10], daysAgo(1), { id: 902 })],
config,
);
expect(reposted.kind).toBe("skip");
});
test("an objection posted before a re-posted notice still keeps the issue open", () => {
const verdict = pendingNotice(
issue(35, "t"),
[notice([10], daysAgo(10)), humanComment(daysAgo(7)), notice([10], daysAgo(4), { id: 902 })],
config,
);
expect(verdict).toEqual({ kind: "skip", reason: "someone replied after the notice" });
});
test("a zero-day grace period acts on the notice at once", () => {
const verdict = pendingNotice(issue(35, "t"), [notice([10], daysAgo(0.01))], { ...config, graceDays: 0 });
expect(verdict.kind).toBe("pending");
});
test("a human reply after the notice keeps the issue open, a bot reply does not", () => {
const human = pendingNotice(issue(35, "t"), [notice([10], daysAgo(5)), humanComment(daysAgo(4))], config);
expect(human).toEqual({ kind: "skip", reason: "someone replied after the notice" });
const bot = pendingNotice(
issue(35, "t"),
[notice([10], daysAgo(5)), { id: 903, body: "triage", created_at: daysAgo(4), user: { type: "Bot", login: "triage[bot]" } }],
config,
);
expect(bot.kind).toBe("pending");
});
test("a human quoting the marker is not a notice", () => {
const quoted = pendingNotice(issue(35, "t"), [notice([10], daysAgo(5), { user: { type: "User", login: "reporter" } })], config);
expect(quoted).toEqual({ kind: "skip", reason: "carries no duplicate notice" });
});
test("never closes an issue twice: a reopened issue is left alone", () => {
const reopened = pendingNotice(
issue(35, "t"),
[notice([10], daysAgo(9)), { id: 904, body: `Closed automatically\n\n${CLOSED_MARKER}`, created_at: daysAgo(5), user: { type: "Bot", login: "github-actions[bot]" } }],
config,
);
expect(reopened).toEqual({ kind: "skip", reason: "was reopened after an automatic close" });
});
test("skips pull requests and issues whose only candidates are newer", () => {
expect(pendingNotice(issue(35, "t", { pull_request: {} }), [notice([10], daysAgo(5))], config).kind).toBe("skip");
expect(pendingNotice(issue(35, "t"), [notice([40], daysAgo(5))], config)).toEqual({
kind: "skip",
reason: "no candidate is older than this issue",
});
});
});
describe("duplicateTarget", () => {
const reporter = issue(35, "[Bug]: Gemma 4-e4b fails on Vertex");
test("closes only against the earliest open issue with the identical normalized title", () => {
const verdict = duplicateTarget(
reporter,
[issue(10, "[Bug]: Gemma 4-e4n fails on Vertex"), issue(20, "[bug]: gemma 4-e4b fails on vertex"), issue(30, "[Bug]: Gemma 4-e4b fails on Vertex")],
[],
);
expect(verdict).toEqual({ kind: "close", duplicateOf: 20 });
});
test("a near miss in the title is not a duplicate", () => {
const verdict = duplicateTarget(reporter, [issue(10, "[Bug]: Gemma 4-e4n fails on Vertex")], []);
expect(verdict).toEqual({ kind: "skip", reason: "no older open issue has the identical title" });
});
test("bare template titles never match each other", () => {
const verdict = duplicateTarget(issue(35, "[Bug]: "), [issue(10, "[Bug]: ")], []);
expect(verdict.kind).toBe("skip");
expect(verdict.kind === "skip" && verdict.reason).toContain("too short");
});
test("a closed candidate or a pull request is never the target", () => {
expect(duplicateTarget(reporter, [issue(10, reporter.title, { state: "closed" })], []).kind).toBe("skip");
expect(duplicateTarget(reporter, [issue(10, reporter.title, { pull_request: {} })], []).kind).toBe("skip");
});
test("a thumbs down on the notice keeps the issue open", () => {
const verdict = duplicateTarget(reporter, [issue(10, reporter.title)], [{ content: "+1" }, { content: "-1" }]);
expect(verdict).toEqual({ kind: "skip", reason: "someone gave the notice a thumbs down" });
});
});
describe("sweepIssue", () => {
const reporter = issue(35, "[Bug]: Gemma 4-e4b fails on Vertex");
const original = issue(10, "[Bug]: Gemma 4-e4b fails on Vertex");
function fakeApi(
comments: readonly Comment[] = [notice([10], daysAgo(5))],
reactionsByNotice: Readonly<Record<number, readonly Reaction[]>> = {},
): { readonly api: GitHubApi; readonly writes: readonly string[] } {
const writes: string[] = [];
const api: GitHubApi = {
request: async <T>(method: string, path: string, body?: object): Promise<T> => {
if (method !== "GET") {
writes.push(`${method} ${path} ${JSON.stringify(body)}`);
return {} as T;
}
if (path.startsWith("/repos/BerriAI/litellm/issues/35/comments")) {
return comments as T;
}
const reactionsPath = path.match(/^\/repos\/BerriAI\/litellm\/issues\/comments\/(\d+)\/reactions/);
if (reactionsPath) {
return (reactionsByNotice[Number(reactionsPath[1])] ?? []) as T;
}
if (path === "/repos/BerriAI/litellm/issues/10") {
return original as T;
}
throw new Error(`unexpected GET ${path}`);
},
};
return { api, writes };
}
test("a dry run reports the close and writes nothing", async () => {
const { api, writes } = fakeApi();
const verdict = await sweepIssue(api, { ...config, dryRun: true }, reporter);
expect(verdict).toEqual({ kind: "close", duplicateOf: 10 });
expect(writes).toEqual([]);
});
test("a thumbs down on an earlier notice still keeps the issue open", async () => {
const { api, writes } = fakeApi([notice([10], daysAgo(9)), notice([10], daysAgo(5), { id: 902 })], { 900: [{ content: "-1" }] });
const verdict = await sweepIssue(api, config, reporter);
expect(verdict).toEqual({ kind: "skip", reason: "someone gave the notice a thumbs down" });
expect(writes).toEqual([]);
});
test("a real run comments, labels, then closes with the duplicate reason", async () => {
const { api, writes } = fakeApi();
const verdict = await sweepIssue(api, config, reporter);
expect(verdict).toEqual({ kind: "close", duplicateOf: 10 });
expect(writes.map((write) => write.split(" ").slice(0, 2).join(" "))).toEqual([
"POST /repos/BerriAI/litellm/issues/35/comments",
"POST /repos/BerriAI/litellm/issues/35/labels",
"PATCH /repos/BerriAI/litellm/issues/35",
]);
expect(writes[0]).toContain("duplicate of #10");
expect(writes[0]).toContain("unanswered for 3 days");
expect(writes[0]).toContain(CLOSED_MARKER);
expect(writes[1]).toContain('{"labels":["duplicate"]}');
expect(writes[2]).toContain('{"state":"closed","state_reason":"duplicate"}');
});
});
describe("reopenTarget", () => {
const closedByBot = (overrides: Partial<Issue> = {}): Issue =>
issue(35, "t", { state: "closed", closed_by: { type: "Bot" }, ...overrides });
const closeMarker = (createdAt: string): Comment => ({
id: 905,
body: `Closed automatically as a duplicate of #10.\n\n${CLOSED_MARKER}`,
created_at: createdAt,
user: { type: "Bot", login: "github-actions[bot]" },
});
test("a reporter reply after the automatic close reopens", () => {
const verdict = reopenTarget(closedByBot(), [closeMarker(daysAgo(2)), humanComment(daysAgo(1))]);
expect(verdict).toEqual({ kind: "reopen" });
});
test("an issue closed by a person stays closed", () => {
const verdict = reopenTarget(closedByBot({ closed_by: { type: "User" } }), [
closeMarker(daysAgo(2)),
humanComment(daysAgo(1)),
]);
expect(verdict).toEqual({ kind: "skip", reason: "was closed by a person" });
});
test("without the automatic-close marker nothing reopens", () => {
const verdict = reopenTarget(closedByBot(), [humanComment(daysAgo(1))]);
expect(verdict).toEqual({ kind: "skip", reason: "carries no automatic-close marker" });
});
test("a maintainer reply alone does not reopen", () => {
const verdict = reopenTarget(closedByBot(), [
closeMarker(daysAgo(2)),
humanComment(daysAgo(1), "Confirmed duplicate", "maintainer"),
]);
expect(verdict).toEqual({ kind: "skip", reason: "the reporter has not replied since the close" });
});
test("a reporter comment from before the close does not reopen", () => {
const verdict = reopenTarget(closedByBot(), [humanComment(daysAgo(3)), closeMarker(daysAgo(2))]);
expect(verdict).toEqual({ kind: "skip", reason: "the reporter has not replied since the close" });
});
test("a pull request never reopens", () => {
const verdict = reopenTarget(closedByBot({ pull_request: {} }), [closeMarker(daysAgo(2)), humanComment(daysAgo(1))]);
expect(verdict).toEqual({ kind: "skip", reason: "is a pull request" });
});
});
describe("sweepClosedIssue", () => {
function fakeApi(issueBody: Issue, comments: readonly Comment[]): { readonly api: GitHubApi; readonly writes: readonly string[] } {
const writes: string[] = [];
const api: GitHubApi = {
request: async <T>(method: string, path: string, body?: object): Promise<T> => {
if (method !== "GET") {
writes.push(`${method} ${path} ${JSON.stringify(body)}`);
return {} as T;
}
if (path.startsWith("/repos/BerriAI/litellm/issues/35/comments")) {
return comments as T;
}
if (path === "/repos/BerriAI/litellm/issues/35") {
return issueBody as T;
}
throw new Error(`unexpected GET ${path}`);
},
};
return { api, writes };
}
const closedByBot = issue(35, "t", { state: "closed", closed_by: { type: "Bot" } });
const closeMarker: Comment = {
id: 905,
body: `Closed automatically as a duplicate of #10.\n\n${CLOSED_MARKER}`,
created_at: daysAgo(2),
user: { type: "Bot", login: "github-actions[bot]" },
};
test("a real run unlabels, reopens, then explains", async () => {
const { api, writes } = fakeApi(closedByBot, [closeMarker, humanComment(daysAgo(1))]);
const verdict = await sweepClosedIssue(api, config, 35);
expect(verdict).toEqual({ kind: "reopen" });
expect(writes).toEqual([
"DELETE /repos/BerriAI/litellm/issues/35/labels/duplicate undefined",
'PATCH /repos/BerriAI/litellm/issues/35 {"state":"open"}',
`POST /repos/BerriAI/litellm/issues/35/comments {"body":"${REOPEN_COMMENT}"}`,
]);
});
test("a dry run reports the reopen and writes nothing", async () => {
const { api, writes } = fakeApi(closedByBot, [closeMarker, humanComment(daysAgo(1))]);
const verdict = await sweepClosedIssue(api, { ...config, dryRun: true }, 35);
expect(verdict).toEqual({ kind: "reopen" });
expect(writes).toEqual([]);
});
});
describe("readConfig", () => {
test("defaults to a real run with a 3-day grace period", () => {
const parsed = readConfig({ GITHUB_TOKEN: "t", GITHUB_REPOSITORY: "BerriAI/litellm" }, NOW);
expect(parsed).toEqual({ token: "t", repo: "BerriAI/litellm", graceDays: 3, dryRun: false, now: NOW });
});
test("honors DRY_RUN and GRACE_PERIOD_DAYS overrides", () => {
const parsed = readConfig(
{ GITHUB_TOKEN: "t", GITHUB_REPOSITORY: "o/r", DRY_RUN: "true", GRACE_PERIOD_DAYS: "0" },
NOW,
);
expect(parsed.dryRun).toBe(true);
expect(parsed.graceDays).toBe(0);
});
test("an empty GRACE_PERIOD_DAYS, as a schedule run renders it, means the default", () => {
const parsed = readConfig({ GITHUB_TOKEN: "t", GITHUB_REPOSITORY: "o/r", GRACE_PERIOD_DAYS: "" }, NOW);
expect(parsed.graceDays).toBe(3);
});
test("refuses a missing token, a malformed repository, or a bad grace period", () => {
expect(() => readConfig({ GITHUB_REPOSITORY: "o/r" }, NOW)).toThrow("GITHUB_TOKEN");
expect(() => readConfig({ GITHUB_TOKEN: "t", GITHUB_REPOSITORY: "litellm" }, NOW)).toThrow("owner/repo");
expect(() => readConfig({ GITHUB_TOKEN: "t", GITHUB_REPOSITORY: "o/r", GRACE_PERIOD_DAYS: "-1" }, NOW)).toThrow(
"GRACE_PERIOD_DAYS",
);
});
});

300
scripts/auto-close-duplicates.ts Normal file
View file

@ -0,0 +1,300 @@
#!/usr/bin/env bun
declare const process: { readonly env: Readonly<Record<string, string | undefined>> };
export interface Issue {
readonly number: number;
readonly title: string;
readonly state: string;
readonly user: { readonly login: string };
readonly closed_by?: { readonly type: string } | null;
readonly pull_request?: unknown;
}
export interface Comment {
readonly id: number;
readonly body: string;
readonly created_at: string;
readonly user: { readonly type: string; readonly login: string };
}
export interface Reaction {
readonly content: string;
}
export interface GitHubApi {
readonly request: <T>(method: "GET" | "POST" | "PATCH" | "DELETE", path: string, body?: object) => Promise<T>;
}
export interface SweepConfig {
readonly repo: string;
readonly graceDays: number;
readonly dryRun: boolean;
readonly now: Date;
}
export type NoticeVerdict =
| { readonly kind: "pending"; readonly notices: readonly Comment[]; readonly candidates: readonly number[] }
| { readonly kind: "skip"; readonly reason: string };
export type CloseVerdict =
| { readonly kind: "close"; readonly duplicateOf: number }
| { readonly kind: "skip"; readonly reason: string };
export type ReopenVerdict =
| { readonly kind: "reopen" }
| { readonly kind: "skip"; readonly reason: string };
export const FLAG_LABEL = "potential-duplicate";
export const CLOSED_MARKER = "<!-- litellm:closed-as-duplicate -->";
export const DEFAULT_GRACE_DAYS = 3;
export const REOPEN_COMMENT =
"Reopened automatically: the reporter replied after the duplicate close, so this needs a human look.";
const NOTICE_MARKER = /<!-- litellm:potential-duplicate candidates=([\d,]*) -->/;
const MIN_TITLE_WORDS = 3;
const PAGE_SIZE = 100;
const DAY_MS = 24 * 60 * 60 * 1000;
const REOPEN_LOOKBACK_DAYS = 30;
const skip = (reason: string): { readonly kind: "skip"; readonly reason: string } => ({ kind: "skip", reason });
export function normalizeTitle(title: string): string {
return title
.toLowerCase()
.replace(/^\s*\[[^\]]*\]\s*:?/, "")
.replace(/[^a-z0-9]+/g, " ")
.trim();
}
export function candidateNumbers(noticeBody: string, issueNumber: number): readonly number[] {
const field = noticeBody.match(NOTICE_MARKER);
if (!field) {
return [];
}
const older = field[1]
.split(",")
.filter((value) => value !== "")
.map(Number)
.filter((candidate) => candidate < issueNumber);
return [...new Set(older)].sort((a, b) => a - b);
}
export function pendingNotice(
issue: Issue,
comments: readonly Comment[],
config: Pick<SweepConfig, "graceDays" | "now">,
): NoticeVerdict {
if (issue.pull_request !== undefined) {
return skip("is a pull request");
}
if (comments.some((comment) => comment.body.includes(CLOSED_MARKER))) {
return skip("was reopened after an automatic close");
}
const notices = comments.filter((comment) => comment.user.type === "Bot" && NOTICE_MARKER.test(comment.body));
const first = notices[0];
const latest = notices[notices.length - 1];
if (first === undefined || latest === undefined) {
return skip("carries no duplicate notice");
}
const ageDays = (config.now.getTime() - new Date(latest.created_at).getTime()) / DAY_MS;
if (ageDays < config.graceDays) {
return skip(`notice is ${ageDays.toFixed(1)} days old, grace period is ${config.graceDays}`);
}
const firstNoticeAt = new Date(first.created_at);
if (comments.some((comment) => comment.user.type !== "Bot" && new Date(comment.created_at) > firstNoticeAt)) {
return skip("someone replied after the notice");
}
const candidates = candidateNumbers(latest.body, issue.number);
if (candidates.length === 0) {
return skip("no candidate is older than this issue");
}
return { kind: "pending", notices, candidates };
}
export function duplicateTarget(
issue: Issue,
candidates: readonly Issue[],
reactions: readonly Reaction[],
): CloseVerdict {
if (reactions.some((reaction) => reaction.content === "-1")) {
return skip("someone gave the notice a thumbs down");
}
const title = normalizeTitle(issue.title);
if (title.split(" ").length < MIN_TITLE_WORDS) {
return skip(`title "${issue.title}" is too short to match on`);
}
const original = candidates.find(
(candidate) =>
candidate.state === "open" && candidate.pull_request === undefined && normalizeTitle(candidate.title) === title,
);
if (original === undefined) {
return skip("no older open issue has the identical title");
}
return { kind: "close", duplicateOf: original.number };
}
export function reopenTarget(issue: Issue, comments: readonly Comment[]): ReopenVerdict {
if (issue.pull_request !== undefined) {
return skip("is a pull request");
}
if (issue.closed_by?.type !== "Bot") {
return skip("was closed by a person");
}
const marker = comments.find((comment) => comment.body.includes(CLOSED_MARKER));
if (marker === undefined) {
return skip("carries no automatic-close marker");
}
const markerAt = new Date(marker.created_at);
if (!comments.some((comment) => comment.user.login === issue.user.login && new Date(comment.created_at) > markerAt)) {
return skip("the reporter has not replied since the close");
}
return { kind: "reopen" };
}
export function closingComment(duplicateOf: number, graceDays: number): string {
return `Closed automatically as a duplicate of #${duplicateOf}. Its title is identical to that older open issue and the duplicate notice above went unanswered for ${graceDays} days. If this is wrong, comment here with how it differs from #${duplicateOf} and this issue will be reopened automatically within a day.
${CLOSED_MARKER}`;
}
async function listAll<T>(api: GitHubApi, path: string, page = 1): Promise<readonly T[]> {
const separator = path.includes("?") ? "&" : "?";
const batch = await api.request<readonly T[]>("GET", `${path}${separator}per_page=${PAGE_SIZE}&page=${page}`);
return batch.length < PAGE_SIZE ? batch : [...batch, ...(await listAll<T>(api, path, page + 1))];
}
async function closeAsDuplicate(
api: GitHubApi,
config: SweepConfig,
issueNumber: number,
duplicateOf: number,
): Promise<void> {
const issuePath = `/repos/${config.repo}/issues/${issueNumber}`;
await api.request("POST", `${issuePath}/comments`, { body: closingComment(duplicateOf, config.graceDays) });
await api.request("POST", `${issuePath}/labels`, { labels: ["duplicate"] });
await api.request("PATCH", issuePath, { state: "closed", state_reason: "duplicate" });
}
async function reopenForReporter(api: GitHubApi, config: SweepConfig, issueNumber: number): Promise<void> {
const issuePath = `/repos/${config.repo}/issues/${issueNumber}`;
await api.request("DELETE", `${issuePath}/labels/duplicate`);
await api.request("PATCH", issuePath, { state: "open" });
await api.request("POST", `${issuePath}/comments`, { body: REOPEN_COMMENT });
}
export async function sweepClosedIssue(api: GitHubApi, config: SweepConfig, issueNumber: number): Promise<ReopenVerdict> {
const issue = await api.request<Issue>("GET", `/repos/${config.repo}/issues/${issueNumber}`);
const comments = await listAll<Comment>(api, `/repos/${config.repo}/issues/${issueNumber}/comments`);
const verdict = reopenTarget(issue, comments);
if (verdict.kind === "reopen" && !config.dryRun) {
await reopenForReporter(api, config, issueNumber);
}
return verdict;
}
export async function sweepIssue(api: GitHubApi, config: SweepConfig, issue: Issue): Promise<CloseVerdict> {
const comments = await listAll<Comment>(api, `/repos/${config.repo}/issues/${issue.number}/comments`);
const pending = pendingNotice(issue, comments, config);
if (pending.kind === "skip") {
return pending;
}
const reactions = (
await Promise.all(
pending.notices.map((notice) => listAll<Reaction>(api, `/repos/${config.repo}/issues/comments/${notice.id}/reactions`)),
)
).flat();
const candidates = await Promise.all(
pending.candidates.map((candidate) => api.request<Issue>("GET", `/repos/${config.repo}/issues/${candidate}`)),
);
const verdict = duplicateTarget(issue, candidates, reactions);
if (verdict.kind === "close" && !config.dryRun) {
await closeAsDuplicate(api, config, issue.number, verdict.duplicateOf);
}
return verdict;
}
function describe(issue: Issue, verdict: CloseVerdict, dryRun: boolean): string {
if (verdict.kind === "skip") {
return `#${issue.number}: skipped, ${verdict.reason}`;
}
return `#${issue.number}: ${dryRun ? "would close" : "closed"} as a duplicate of #${verdict.duplicateOf}`;
}
export async function sweep(api: GitHubApi, config: SweepConfig): Promise<readonly CloseVerdict[]> {
const issues = await listAll<Issue>(api, `/repos/${config.repo}/issues?state=open&labels=${FLAG_LABEL}`);
console.log(`${issues.length} open issues carry the ${FLAG_LABEL} label in ${config.repo}${config.dryRun ? " (dry run)" : ""}`);
return issues.reduce<Promise<readonly CloseVerdict[]>>(async (previous, issue) => {
const verdicts = await previous;
const verdict = await sweepIssue(api, config, issue);
console.log(describe(issue, verdict, config.dryRun));
return [...verdicts, verdict];
}, Promise.resolve([]));
}
function describeReopen(issueNumber: number, verdict: ReopenVerdict, dryRun: boolean): string {
if (verdict.kind === "skip") {
return `#${issueNumber}: skipped, ${verdict.reason}`;
}
return `#${issueNumber}: ${dryRun ? "would reopen" : "reopened"} for the reporter's reply`;
}
export async function reopenSweep(api: GitHubApi, config: SweepConfig): Promise<readonly ReopenVerdict[]> {
const since = new Date(config.now.getTime() - REOPEN_LOOKBACK_DAYS * DAY_MS).toISOString();
const closedPath = `/repos/${config.repo}/issues?state=closed&labels=duplicate,${FLAG_LABEL}&since=${encodeURIComponent(since)}`;
const issues = await listAll<Issue>(api, closedPath);
console.log(`${issues.length} recently closed issues carry the duplicate and ${FLAG_LABEL} labels in ${config.repo}${config.dryRun ? " (dry run)" : ""}`);
return issues.reduce<Promise<readonly ReopenVerdict[]>>(async (previous, issue) => {
const verdicts = await previous;
const verdict = await sweepClosedIssue(api, config, issue.number);
console.log(describeReopen(issue.number, verdict, config.dryRun));
return [...verdicts, verdict];
}, Promise.resolve([]));
}
export function readConfig(env: Readonly<Record<string, string | undefined>>, now: Date): SweepConfig & { readonly token: string } {
const token = env.GITHUB_TOKEN;
const repo = env.GITHUB_REPOSITORY;
if (!token || !repo || !/^[\w.-]+\/[\w.-]+$/.test(repo)) {
throw new Error("GITHUB_TOKEN and GITHUB_REPOSITORY (owner/repo) are required");
}
const rawGraceDays = env.GRACE_PERIOD_DAYS?.trim();
const graceDays = rawGraceDays === undefined || rawGraceDays === "" ? DEFAULT_GRACE_DAYS : Number(rawGraceDays);
if (!Number.isFinite(graceDays) || graceDays < 0) {
throw new Error(`GRACE_PERIOD_DAYS must be a non-negative number, got "${env.GRACE_PERIOD_DAYS}"`);
}
return { token, repo, graceDays, dryRun: env.DRY_RUN === "true", now };
}
export function githubApi(token: string): GitHubApi {
return {
request: async <T>(method: "GET" | "POST" | "PATCH" | "DELETE", path: string, body?: object): Promise<T> => {
const response = await fetch(`https://api.github.com${path}`, {
method,
headers: {
Authorization: `Bearer ${token}`,
Accept: "application/vnd.github+json",
"X-GitHub-Api-Version": "2022-11-28",
"User-Agent": "litellm-auto-close-duplicates",
...(body === undefined ? {} : { "Content-Type": "application/json" }),
},
body: body === undefined ? undefined : JSON.stringify(body),
});
if (!response.ok) {
throw new Error(`${method} ${path} failed: ${response.status} ${response.statusText}`);
}
return (await response.json()) as T;
},
};
}
if (import.meta.main) {
const { token, ...config } = readConfig(process.env, new Date());
const api = githubApi(token);
const closeVerdicts = await sweep(api, config);
const reopenVerdicts = await reopenSweep(api, config);
const closed = closeVerdicts.filter((verdict) => verdict.kind === "close").length;
const reopened = reopenVerdicts.filter((verdict) => verdict.kind === "reopen").length;
console.log(
`${config.dryRun ? "Would close" : "Closed"} ${closed} of ${closeVerdicts.length} flagged issues, ${config.dryRun ? "would reopen" : "reopened"} ${reopened}`,
);
}

56
tests/proxy_unit_tests/test_proxy_utils.py
View file

@ -1317,6 +1317,62 @@ def test_proxy_config_state_post_init_callback_call(monkeypatch):
assert config["litellm_settings"]["default_team_settings"][0]["team_id"] == "test"
@pytest.mark.asyncio
async def test_default_team_settings_newrelic_resolves_traces_and_metrics():
"""Static `default_team_settings` is the config-file twin of POST /team/callback.
A team pinned to New Relic through `default_team_settings` must reach the
same two loggers the dynamic path does: the per-team metrics logger (cost
and usage) and the trace logger (LLM/agent spans). This proves the static
path resolves both, not just one, so the config-file customer gets the
same per-team routing as the API customer.
"""
from litellm.litellm_core_utils.litellm_logging import Logging
from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup
from litellm.proxy.proxy_server import ProxyConfig
pc = ProxyConfig()
pc.config = {
"litellm_settings": {
"default_team_settings": [
{
"team_id": "team-a",
"success_callback": ["newrelic"],
"newrelic_api_key": "team-a-ingest-key",
"newrelic_region": "eu",
}
]
}
}
callback_metadata = LiteLLMProxyRequestSetup.add_team_based_callbacks_from_config(
team_id="team-a",
proxy_config=pc,
)
assert callback_metadata is not None
assert callback_metadata.success_callback == ["newrelic"]
assert callback_metadata.callback_vars == {
"newrelic_api_key": "team-a-ingest-key",
"newrelic_region": "eu",
}
logging_obj = Logging(
model="gpt-3.5-turbo",
messages=[{"role": "user", "content": "hi"}],
stream=False,
call_type="completion",
start_time=None,
litellm_call_id="static-nr-1",
function_id="static-nr-1",
)
logging_obj._trusted_callback_vars = tuple(callback_metadata.callback_vars.items())
resolved = logging_obj._resolve_dynamic_callback_string("newrelic")
resolved_names = {type(logger).__name__ for logger in resolved}
assert resolved_names == {"NewRelicMetricsLogger", "NewRelicLogger"}
def test_proxy_config_state_get_config_state_error():
"""
Ensures that get_config_state does not raise an error when the config is not a valid dictionary

55
tests/test_litellm/integrations/SlackAlerting/test_slack_alerting.py
View file

@ -12,7 +12,7 @@ import litellm
from litellm.caching.caching import DualCache
from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting
from litellm.proxy._types import CallInfo, Litellm_EntityType
from litellm.types.integrations.slack_alerting import SlackAlertingCacheKeys
from litellm.types.integrations.slack_alerting import AlertType, SlackAlertingCacheKeys
class TestSlackAlerting(unittest.TestCase):
@ -366,3 +366,56 @@ async def test_scheduled_daily_report_threads_the_pod_lock_manager_through():
_, kwargs = slack_alerting._run_scheduler_helper.await_args
assert kwargs["pod_lock_manager"] is pod_lock_manager
def _slack_alerting_with_env_resolution() -> SlackAlerting:
slack_alerting: Final = SlackAlerting(alerting=["slack"], internal_usage_cache=DualCache())
slack_alerting.periodic_started = True
return slack_alerting
@pytest.mark.asyncio
async def test_send_alert_falls_back_to_alerting_webhook_url_env(monkeypatch):
monkeypatch.delenv("SLACK_WEBHOOK_URL", raising=False)
monkeypatch.setenv("ALERTING_WEBHOOK_URL", "https://chat.example.com/hooks/abc")
slack_alerting: Final = _slack_alerting_with_env_resolution()
await slack_alerting.send_alert(
message="budget crossed",
level="High",
alert_type=AlertType.budget_alerts,
alerting_metadata={},
)
assert slack_alerting.log_queue[0]["url"] == "https://chat.example.com/hooks/abc"
@pytest.mark.asyncio
async def test_send_alert_prefers_slack_webhook_url_over_fallback(monkeypatch):
monkeypatch.setenv("SLACK_WEBHOOK_URL", "https://hooks.slack.com/services/T0/B0/X0")
monkeypatch.setenv("ALERTING_WEBHOOK_URL", "https://chat.example.com/hooks/abc")
slack_alerting: Final = _slack_alerting_with_env_resolution()
await slack_alerting.send_alert(
message="budget crossed",
level="High",
alert_type=AlertType.budget_alerts,
alerting_metadata={},
)
assert slack_alerting.log_queue[0]["url"] == "https://hooks.slack.com/services/T0/B0/X0"
@pytest.mark.asyncio
async def test_send_alert_raises_when_no_webhook_url_configured(monkeypatch):
monkeypatch.delenv("SLACK_WEBHOOK_URL", raising=False)
monkeypatch.delenv("ALERTING_WEBHOOK_URL", raising=False)
slack_alerting: Final = _slack_alerting_with_env_resolution()
with pytest.raises(ValueError, match="SLACK_WEBHOOK_URL / ALERTING_WEBHOOK_URL"):
await slack_alerting.send_alert(
message="budget crossed",
level="High",
alert_type=AlertType.budget_alerts,
alerting_metadata={},
)

17
tests/test_litellm/integrations/SlackAlerting/test_slack_alerting_digest.py
View file

@ -79,6 +79,23 @@ class TestDigestMode(unittest.IsolatedAsyncioTestCase):
self.assertEqual(len(self.slack_alerting.digest_buckets), 2)
async def test_digest_falls_back_to_alerting_webhook_url_env(self):
"""With SLACK_WEBHOOK_URL unset, the digest entry resolves ALERTING_WEBHOOK_URL instead."""
env = {k: v for k, v in os.environ.items() if k != "SLACK_WEBHOOK_URL"}
env["ALERTING_WEBHOOK_URL"] = "https://chat.example.com/hooks/abc"
with unittest.mock.patch.dict(os.environ, env, clear=True):
await self.slack_alerting.send_alert(
message="`Requests are hanging`",
level="Medium",
alert_type=AlertType.llm_requests_hanging,
alerting_metadata={},
request_model="gemini-2.5-flash",
api_base="None",
)
bucket = list(self.slack_alerting.digest_buckets.values())[0]
self.assertEqual(bucket["webhook_url"], "https://chat.example.com/hooks/abc")
async def test_non_digest_alert_goes_to_queue(self):
"""Alert types without digest enabled should go straight to the log queue."""
message = "Budget exceeded"

30
tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py
View file

@ -957,6 +957,28 @@ def test_transform_request_helper_includes_anthropic_beta_and_tools():
assert fields["tools"][0]["type"] == "computer_20250124"
def test_config_blocks_do_not_leak_into_inference_config():
"""Regression: inferenceConfig was built before the config blocks were popped, so a dead
nested copy of each block (guardrailConfig, performanceConfig, serviceTier) rode inside
inferenceConfig alongside the real top-level one."""
data = AmazonConverseConfig()._transform_request_helper(
model="anthropic.claude-haiku-4-5-20251001-v1:0",
system_content_blocks=[],
optional_params={
"maxTokens": 100,
"guardrailConfig": {"guardrailIdentifier": "gr-id", "guardrailVersion": "DRAFT"},
"performanceConfig": {"latency": "optimized"},
"serviceTier": {"type": "priority"},
},
messages=[{"role": "user", "content": "hi"}],
)
assert data["inferenceConfig"] == {"maxTokens": 100}
assert data["guardrailConfig"] == {"guardrailIdentifier": "gr-id", "guardrailVersion": "DRAFT"}
assert data["performanceConfig"] == {"latency": "optimized"}
assert data["serviceTier"] == {"type": "priority"}
def test_parallel_tool_calls_config_kept_for_sonnet_5(monkeypatch):
old_env = os.environ.get("LITELLM_LOCAL_MODEL_COST_MAP")
old_cost = litellm.model_cost
@ -2853,17 +2875,11 @@ def test_guarded_text_guardrail_config_preserved():
headers={},
)
# GuardrailConfig should be present at top level
assert "guardrailConfig" in result
assert result["guardrailConfig"]["guardrailIdentifier"] == "gr-abc123"
# GuardrailConfig should also be in inferenceConfig
assert "inferenceConfig" in result
assert "guardrailConfig" in result["inferenceConfig"]
assert (
result["inferenceConfig"]["guardrailConfig"]["guardrailIdentifier"]
== "gr-abc123"
)
assert "guardrailConfig" not in result["inferenceConfig"]
def test_auto_convert_last_user_message_to_guarded_text():

19
tests/test_litellm/llms/gdc/chat/test_gdc_chat_transformation.py
View file

@ -310,6 +310,25 @@ class TestGDCGeminiConfig:
api_base=TEST_API_BASE,
)
def test_validate_environment_credentials_missing_audience_binding_are_named(self):
config = GDCGeminiConfig()
creds_without_audience_binding = MagicMock(spec=[])
with patch(
"google.auth.load_credentials_from_dict",
return_value=(creds_without_audience_binding, None),
):
with pytest.raises(AttributeError, match="must expose with_gdch_audience"):
config.validate_environment(
headers={},
model=TEST_MODEL,
messages=[],
optional_params={},
litellm_params={"vertex_project": TEST_PROJECT},
api_key=TEST_API_KEY,
api_base=TEST_API_BASE,
)
def test_validate_environment_string_false_disables_token_caching(self):
config = GDCGeminiConfig()
mock_creds = MagicMock()

110
tests/test_litellm/llms/vertex_ai/test_vertex_llm_base.py
View file

@ -982,6 +982,116 @@ class TestVertexBase:
assert result_url == f"{gateway_api_base}:embedContent"
def test_check_custom_proxy_vertex_api_base_with_version_path_grafts_default_path(self):
vertex_base = VertexBase()
_, result_url = vertex_base._check_custom_proxy(
api_base="https://aiplatform.googleapis.com/v1beta1",
custom_llm_provider="vertex_ai",
gemini_api_key=None,
endpoint="generateContent",
stream=None,
auth_header="Bearer token123",
url="https://us-central1-aiplatform.googleapis.com/v1/projects/test-project/locations/us-central1/publishers/google/models/gemini-3.5-flash-lite:generateContent",
model="gemini-3.5-flash-lite",
)
assert (
result_url
== "https://aiplatform.googleapis.com/v1beta1/projects/test-project/locations/us-central1/publishers/google/models/gemini-3.5-flash-lite:generateContent"
)
def test_check_custom_proxy_vertex_api_base_with_version_path_trailing_slash_grafts_default_path(self):
vertex_base = VertexBase()
_, result_url = vertex_base._check_custom_proxy(
api_base="https://internal-gateway.example.com/v1/",
custom_llm_provider="vertex_ai",
gemini_api_key=None,
endpoint="generateContent",
stream=None,
auth_header="Bearer token123",
url="https://us-central1-aiplatform.googleapis.com/v1/projects/test-project/locations/us-central1/publishers/google/models/gemini-3.5-flash-lite:generateContent",
model="gemini-3.5-flash-lite",
)
assert (
result_url
== "https://internal-gateway.example.com/v1/projects/test-project/locations/us-central1/publishers/google/models/gemini-3.5-flash-lite:generateContent"
)
def test_check_custom_proxy_vertex_api_base_with_version_path_and_query_grafts_before_query(self):
vertex_base = VertexBase()
_, result_url = vertex_base._check_custom_proxy(
api_base="https://internal-gateway.example.com/v1beta1?key=abc",
custom_llm_provider="vertex_ai",
gemini_api_key=None,
endpoint="generateContent",
stream=None,
auth_header="Bearer token123",
url="https://us-central1-aiplatform.googleapis.com/v1/projects/test-project/locations/us-central1/publishers/google/models/gemini-3.5-flash-lite:generateContent",
model="gemini-3.5-flash-lite",
)
assert (
result_url
== "https://internal-gateway.example.com/v1beta1/projects/test-project/locations/us-central1/publishers/google/models/gemini-3.5-flash-lite:generateContent?key=abc"
)
def test_check_custom_proxy_vertex_api_base_with_version_path_and_query_streaming_appends_alt_sse(self):
vertex_base = VertexBase()
_, result_url = vertex_base._check_custom_proxy(
api_base="https://internal-gateway.example.com/v1beta1?key=abc",
custom_llm_provider="vertex_ai",
gemini_api_key=None,
endpoint="streamGenerateContent",
stream=True,
auth_header="Bearer token123",
url="https://us-central1-aiplatform.googleapis.com/v1/projects/test-project/locations/us-central1/publishers/google/models/gemini-3.5-flash-lite:streamGenerateContent",
model="gemini-3.5-flash-lite",
)
assert (
result_url
== "https://internal-gateway.example.com/v1beta1/projects/test-project/locations/us-central1/publishers/google/models/gemini-3.5-flash-lite:streamGenerateContent?key=abc&alt=sse"
)
def test_check_custom_proxy_vertex_api_base_with_non_version_path_keeps_endpoint_append(self):
vertex_base = VertexBase()
gateway_api_base = "https://gateway.example.com/vertex-proxy"
_, result_url = vertex_base._check_custom_proxy(
api_base=gateway_api_base,
custom_llm_provider="vertex_ai",
gemini_api_key=None,
endpoint="generateContent",
stream=None,
auth_header="Bearer token123",
url="https://us-central1-aiplatform.googleapis.com/v1/projects/test-project/locations/us-central1/publishers/google/models/gemini-3.5-flash-lite:generateContent",
model="gemini-3.5-flash-lite",
)
assert result_url == f"{gateway_api_base}:generateContent"
def test_check_custom_proxy_vertex_api_base_without_projects_in_default_url_keeps_endpoint_append(self):
vertex_base = VertexBase()
gemma_api_base = "https://example.com/custom/gemma-deployment"
_, result_url = vertex_base._check_custom_proxy(
api_base=gemma_api_base,
custom_llm_provider="vertex_ai",
gemini_api_key=None,
endpoint="predict",
stream=False,
auth_header=None,
url=gemma_api_base,
model="gemma-3-27b-it",
)
assert result_url == f"{gemma_api_base}:predict"
def test_check_custom_proxy_vertex_bare_host_streaming_keeps_single_alt_sse(self):
vertex_base = VertexBase()

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