Merge remote-tracking branch 'github/main' into litellm_rust_bridge_declarative_route_catalog

# Conflicts:
#	tests/e2e/access_control/test_model_access_group_e2e.py
This commit is contained in:
Yujong Lee 2026-09-17 11:08:08 -07:00
commit b26935416a
327 changed files with 23015 additions and 2667 deletions

View file

@ -257,7 +257,7 @@ commands:
- install_rust
- restore_cache:
keys:
- v1-uv-cache-{{ checksum "uv.lock" }}
- v3-integration-uv-cache-{{ checksum "uv.lock" }}
- run:
name: Install Dependencies
command: |
@ -266,7 +266,7 @@ commands:
- save_cache:
paths:
- ~/.cache/uv
key: v1-uv-cache-{{ checksum "uv.lock" }}
key: v3-integration-uv-cache-{{ checksum "uv.lock" }}
jobs:
# Add Windows testing job
@ -2955,6 +2955,32 @@ jobs:
working_directory: ~/project
steps:
- setup_litellm_test_deps
- when:
condition:
equal: [browser, << parameters.suite >>]
steps:
- install_node
- restore_cache:
keys:
- integration-ui-v3-{{ checksum "ui/litellm-dashboard/package-lock.json" }}-{{ checksum "tests/e2e/ui/package-lock.json" }}
- run:
name: Install locked browser dependencies
command: |
cd ui/litellm-dashboard
npm ci
cd ../../tests/e2e/ui
npm ci
sudo env PATH="$PATH" DEBIAN_FRONTEND=noninteractive NEEDRESTART_MODE=l \
timeout --signal=TERM --kill-after=20s 6m node node_modules/@playwright/test/cli.js install-deps chromium
timeout --signal=TERM --kill-after=20s 3m node node_modules/@playwright/test/cli.js install chromium
- save_cache:
key: integration-ui-v3-{{ checksum "ui/litellm-dashboard/package-lock.json" }}-{{ checksum "tests/e2e/ui/package-lock.json" }}
paths:
- ~/.npm
- ~/.cache/ms-playwright
- run:
name: Build the candidate dashboard
command: cd ui/litellm-dashboard && NEXT_TELEMETRY_DISABLED=1 npm run build
- start_postgres:
image: postgres:16@sha256:e17e86066e5ef83e0952a9347f5c792b7ece00972e2aa787a6986f471b3dd3d5
- start_redis
@ -2983,7 +3009,7 @@ workflows:
name: integration-<< matrix.suite >>
matrix:
parameters:
suite: [management, accounting, providers]
suite: [management, accounting, database, providers, extensions, sdk, browser]
filters:
branches:
only:

View file

@ -1,12 +1,14 @@
#!/usr/bin/env bash
set -uo pipefail
category="${1:?usage: classify_changes.sh <backend|client|ui|provider-harness>}"
category="${1:?usage: classify_changes.sh <backend|client|ui|provider-harness|cost-map-only>}"
has_client=false
has_backend=false
has_ci=false
has_provider_harness=false
has_cost_map=false
outside_cost_map_set=false
while IFS= read -r file || [ -n "$file" ]; do
[ -n "$file" ] || continue
case "$file" in
@ -20,9 +22,18 @@ while IFS= read -r file || [ -n "$file" ]; do
.github/* | .circleci/*) has_ci=true; has_backend=true ;;
*) has_backend=true ;;
esac
case "$file" in
model_prices_and_context_window.json | litellm/model_prices_and_context_window_backup.json | model_prices_and_context_window.schema.json)
has_cost_map=true ;;
tests/test_litellm/* | tests/proxy_unit_tests/*) : ;;
*) outside_cost_map_set=true ;;
esac
done
case "$category" in
cost-map-only)
{ [ "$has_cost_map" = true ] && [ "$outside_cost_map_set" = false ]; } && echo run || echo skip
;;
provider-harness)
[ "$has_provider_harness" = true ] && echo run || echo skip
;;

View file

@ -1,6 +1,11 @@
#!/usr/bin/env bash
set -euo pipefail
if [ "${GITHUB_ACTIONS:-}" = true ]; then
echo "Integration contracts are owned by CircleCI" >&2
exit 1
fi
suite="${1:?integration suite required}"
results="test-results/integration-${suite}"
mkdir -p "$results"
@ -65,7 +70,13 @@ export INTEGRATION_PROXY_URL=http://127.0.0.1:4000
export INTEGRATION_PEER_URL=""
export INTEGRATION_UPSTREAM_URL=http://127.0.0.1:8190
export INTEGRATION_MASTER_KEY="$LITELLM_MASTER_KEY"
export INTEGRATION_SEED="$((16#$(git rev-parse --short=8 HEAD)))"
export LITELLM_UI_PATH="$PWD/litellm/proxy/_experimental/out"
if [ "$suite" = browser ]; then
export LITELLM_UI_PATH="$PWD/ui/litellm-dashboard/out"
test -f "$LITELLM_UI_PATH/index.html"
fi
export INTEGRATION_SEED="$(.venv/bin/python -c 'import hashlib,os; print(int(hashlib.sha256((os.environ.get("CIRCLE_SHA1", "local") + os.environ.get("CIRCLE_WORKFLOW_ID", "local")).encode()).hexdigest()[:8],16))')"
export INTEGRATION_ORDER_SEED="$INTEGRATION_SEED"
uv run --no-sync prisma generate --schema litellm/proxy/schema.prisma > "$results/prisma-generate.log" 2>&1
@ -102,7 +113,7 @@ start_proxy() {
local log_name="$2"
setsid env -i PATH="$PATH" HOME="$HOME" PYTHONPATH="$PYTHONPATH" INTEGRATION_RUN_ID="$integration_identity" \
DATABASE_URL="$DATABASE_URL" REDIS_HOST="$REDIS_HOST" REDIS_PORT="$REDIS_PORT" \
LITELLM_MASTER_KEY="$LITELLM_MASTER_KEY" LITELLM_SALT_KEY="$LITELLM_SALT_KEY" \
LITELLM_MASTER_KEY="$LITELLM_MASTER_KEY" LITELLM_SALT_KEY="$LITELLM_SALT_KEY" LITELLM_UI_PATH="$LITELLM_UI_PATH" PROXY_BASE_URL="http://127.0.0.1:$port" \
LITELLM_MODE=PRODUCTION LITELLM_LOCAL_MODEL_COST_MAP=True STORE_MODEL_IN_DB=True \
AWS_EC2_METADATA_DISABLED=true DO_NOT_TRACK=1 \
.venv/bin/python -m integration._support.proxy --config tests/integration/proxy_config.yaml \
@ -131,6 +142,19 @@ if [ "$suite" = providers ]; then
--junitxml="$results/replay-controls.xml"
fi
if [ "$suite" = browser ]; then
export E2E_UI_BASE_URL="$INTEGRATION_PROXY_URL" E2E_UI_ARTIFACT_DIR="$PWD/$results"
export INTEGRATION_PYTHON="$PWD/.venv/bin/python"
timeout --signal=TERM --kill-after=20s 3m env -i PATH="$PATH" HOME="$HOME" PYTHONPATH="$PYTHONPATH" \
INTEGRATION_RUN_ID="$integration_identity" DATABASE_URL="$DATABASE_URL" \
INTEGRATION_UPSTREAM_URL="$INTEGRATION_UPSTREAM_URL" INTEGRATION_PYTHON="$INTEGRATION_PYTHON" \
E2E_UI_BASE_URL="$E2E_UI_BASE_URL" E2E_UI_ARTIFACT_DIR="$E2E_UI_ARTIFACT_DIR" \
LITELLM_MASTER_KEY="$LITELLM_MASTER_KEY" CI=true \
node tests/e2e/ui/node_modules/@playwright/test/cli.js test --config tests/e2e/ui/integration.config.ts
.venv/bin/python .circleci/scripts/verify_integration_browser.py "$results/browser-results.json"
exit 0
fi
timeout --signal=TERM --kill-after=20s 11m env -i PATH="$PATH" HOME="$HOME" PYTHONPATH="$PYTHONPATH" \
INTEGRATION_RUN_ID="$integration_identity" \
DATABASE_URL="$DATABASE_URL" REDIS_HOST="$REDIS_HOST" REDIS_PORT="$REDIS_PORT" \
@ -138,5 +162,6 @@ timeout --signal=TERM --kill-after=20s 11m env -i PATH="$PATH" HOME="$HOME" PYTH
INTEGRATION_UPSTREAM_URL="$INTEGRATION_UPSTREAM_URL" \
INTEGRATION_MASTER_KEY="$INTEGRATION_MASTER_KEY" LITELLM_MODE=PRODUCTION \
INTEGRATION_SEED="$INTEGRATION_SEED" \
INTEGRATION_ORDER_SEED="$INTEGRATION_ORDER_SEED" \
LITELLM_LOCAL_MODEL_COST_MAP=True AWS_EC2_METADATA_DISABLED=true DO_NOT_TRACK=1 \
.venv/bin/python tests/integration/run.py "$suite" --results "$results"

View file

@ -0,0 +1,60 @@
import json
import sys
from pathlib import Path
from typing import Final
from pydantic import TypeAdapter
from typing_extensions import NotRequired, ReadOnly, TypedDict
class BrowserAttempt(TypedDict):
status: ReadOnly[str]
retry: ReadOnly[int]
class BrowserTest(TypedDict):
results: ReadOnly[list[BrowserAttempt]]
class BrowserSpec(TypedDict):
file: ReadOnly[str]
title: ReadOnly[str]
tests: ReadOnly[list[BrowserTest]]
class BrowserSuite(TypedDict):
specs: NotRequired[ReadOnly[list[BrowserSpec]]]
suites: NotRequired[ReadOnly[list["BrowserSuite"]]]
def main() -> None:
result: Final = json.loads(Path(sys.argv[1]).read_text())
assert not result.get("errors"), result.get("errors")
expected: Final = json.loads(
(Path(__file__).resolve().parents[2] / "tests/integration/contracts.json").read_text()
)["browser"]
assert expected and result["stats"]["expected"] == len(expected)
assert all(result["stats"][name] == 0 for name in ("unexpected", "flaky", "skipped"))
def cases(suite: BrowserSuite) -> tuple[BrowserSpec, ...]:
return tuple(suite.get("specs", ())) + tuple(spec for child in suite.get("suites", ()) for spec in cases(child))
suites: Final = TypeAdapter(list[BrowserSuite]).validate_python(result["suites"], strict=True)
specs: Final = tuple(spec for suite in suites for spec in cases(suite))
repository: Final = Path(__file__).resolve().parents[2]
report_root: Final = Path(result["config"]["rootDir"])
assert report_root.is_absolute(), "Playwright rootDir must be explicit"
observed: Final = tuple(
str((report_root / spec["file"]).resolve().relative_to(repository)) + "::" + spec["title"] for spec in specs
)
assert sorted(observed) == sorted(expected)
for spec in specs:
tests: Final = spec["tests"]
assert len(tests) == 1 and len(tests[0]["results"]) == 1
assert tests[0]["results"][0]["status"] == "passed" and tests[0]["results"][0]["retry"] == 0
sys.stdout.write("One canonical browser contract passed once without skips or retries\n")
if __name__ == "__main__":
main()

View file

@ -505,6 +505,7 @@ def _integration_ownership(repo_root: pathlib.Path = REPO_ROOT) -> tuple[frozens
return frozenset(), ()
entries: Final = json.loads(manifest.read_text())
paths: Final = frozenset(node.split("::", 1)[0] for node in entries["tests"])
browser_paths: Final = frozenset(node.split("::", 1)[0] for node in entries.get("browser", {}))
circle_path: Final = repo_root / ".circleci/config.yml"
circle: Final = yaml.safe_load(circle_path.read_text()) if circle_path.exists() else {}
steps: Final = circle.get("jobs", {}).get("integration_contracts", {}).get("steps", ())
@ -523,7 +524,7 @@ def _integration_ownership(repo_root: pathlib.Path = REPO_ROOT) -> tuple[frozens
.get("suite", (job["integration_contracts"].get("suite"),))
if isinstance(suite, str)
)
required: Final = frozenset(
required: Final = (frozenset({"browser"}) if browser_paths else frozenset()) | frozenset(
group
for group, folders in entries["groups"].items()
if any(any(path.startswith(f"tests/integration/{folder}/") for folder in folders) for path in paths)
@ -551,6 +552,40 @@ def _integration_ownership(repo_root: pathlib.Path = REPO_ROOT) -> tuple[frozens
for path in paths
if not (repo_root / path).is_file()
)
browser_commands: Final = tuple(
scalar.value
for path in (repo_root / ".github/workflows").glob("*.y*ml")
for scalar in _scalars(yaml.safe_load(path.read_text()), path.name)
if scalar.key in {"run", "command"}
)
browser_findings: Final = tuple(
Finding(path, "browser integration contract is explicitly selected by GitHub Actions")
for path in browser_paths
if any(
path in command
or pathlib.Path(path).name in command
or "integrationCritical" in command
or "integration.config.ts" in command
or ("run_integration.sh" in command and "browser" in command)
for command in browser_commands
)
) + tuple(
Finding(path, "canonical browser integration file is missing")
for path in browser_paths
if not (repo_root / path).is_file()
)
default_browser: Final = repo_root / "tests/e2e/ui/playwright.config.ts"
exclusion_findings: Final = (
(
Finding(
str(default_browser.relative_to(repo_root)),
"default Playwright selection must exclude integrationCritical",
),
)
if browser_paths
and (not default_browser.exists() or "**/integrationCritical/**" not in default_browser.read_text())
else ()
)
group_findings: Final = tuple(
Finding(group, "canonical integration group is not scheduled by CircleCI")
for group in sorted(required - scheduled)
@ -559,7 +594,7 @@ def _integration_ownership(repo_root: pathlib.Path = REPO_ROOT) -> tuple[frozens
return frozenset(), findings + (
Finding(str(manifest.relative_to(repo_root)), "dedicated CircleCI runner is missing"),
)
return paths, findings + group_findings
return paths | browser_paths, findings + group_findings + browser_findings + exclusion_findings
def main() -> int:

465
.github/scripts/auto_merge_price_sync.py vendored Normal file
View file

@ -0,0 +1,465 @@
"""Auto-merge the provider-info-sync bot's cost-map pull requests.
Evaluates every gate (author allowlist, cost-map-only diff, required and
non-required checks, Greptile confidence, Bugbot review, human reviews) and
merges with a merge commit when all of them hold. Every hold reason is
logged; the process exits 0 on hold and 1 only on API or programming errors.
``DRY_RUN=1`` prints the verdict without calling the merge endpoint.
"""
from __future__ import annotations
import json
import os
import re
import subprocess
import sys
import time
import urllib.error
import urllib.request
from collections.abc import Callable, Mapping, Sequence
from dataclasses import dataclass
from datetime import datetime, timezone
from typing import Final
REPO_ROOT: Final = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
CLASSIFY_SCRIPT: Final = os.path.join(REPO_ROOT, ".circleci", "scripts", "classify_changes.sh")
API_ROOT: Final = "https://api.github.com"
CHANGED_FILE_CEILING: Final = 3000
OK_CHECK_CONCLUSIONS: Final = frozenset({"success", "skipped", "neutral"})
GREPTILE_LOGIN: Final = "greptile-apps[bot]"
BUGBOT_LOGIN: Final = "cursor[bot]"
GREPTILE_SCORE_RE: Final = re.compile(r"Confidence Score:\s*(\d)/5")
BUGBOT_REVIEW_MARKER: Final = "<!-- BUGBOT_REVIEW -->"
BUGBOT_STALE_MARKER: Final = "<!-- BUGBOT_REVIEW_STALE -->"
BUGBOT_CLEAN: Final = "found no new issues"
@dataclass(frozen=True, slots=True)
class PullRequest:
number: int
title: str
author_login: str
state: str
draft: bool
mergeable: bool | None
mergeable_state: str
head_sha: str
@dataclass(frozen=True, slots=True)
class CheckRun:
name: str
status: str
conclusion: str | None
@dataclass(frozen=True, slots=True)
class CommitStatus:
context: str
state: str
@dataclass(frozen=True, slots=True)
class IssueComment:
author_login: str
body: str
updated_at: datetime
@dataclass(frozen=True, slots=True)
class Review:
author_login: str
state: str
body: str
commit_id: str
submitted_at: datetime
@dataclass(frozen=True, slots=True)
class Verdict:
merge: bool
reasons: tuple[str, ...]
@dataclass(frozen=True, slots=True)
class EvaluationInputs:
pr: PullRequest
changed_files: tuple[str, ...]
required_contexts: frozenset[str]
check_runs: tuple[CheckRun, ...]
statuses: tuple[CommitStatus, ...]
comments: tuple[IssueComment, ...]
reviews: tuple[Review, ...]
head_commit_date: datetime
self_check_name: str
author_allowlist: frozenset[str]
def _is_bot_login(login: str) -> bool:
return login.lower().endswith("[bot]")
def _classify(changed_files: Sequence[str]) -> str:
result: Final = subprocess.run(
["bash", CLASSIFY_SCRIPT, "cost-map-only"],
input="\n".join(changed_files),
capture_output=True,
text=True,
check=False,
)
if result.returncode != 0:
return "error"
return result.stdout.strip()
def evaluate(
inputs: EvaluationInputs,
*,
classify: Callable[[Sequence[str]], str] = _classify,
) -> Verdict:
pr: Final = inputs.pr
reasons: list[str] = []
if pr.author_login.lower() not in {login.lower() for login in inputs.author_allowlist}:
reasons.append(f"author {pr.author_login!r} not in allowlist")
if pr.state != "open":
reasons.append("pr not open")
if pr.draft:
reasons.append("pr is a draft")
if pr.mergeable is None:
reasons.append("mergeability unknown")
elif not pr.mergeable:
reasons.append("pr not mergeable")
if pr.mergeable_state == "dirty":
reasons.append("pr has merge conflicts")
if len(inputs.changed_files) > CHANGED_FILE_CEILING:
reasons.append(f"changed file count {len(inputs.changed_files)} over {CHANGED_FILE_CEILING} ceiling")
else:
decision: Final = classify(inputs.changed_files)
if decision != "run":
reasons.append("changed files outside the cost-map-only set")
green_runs: Final = frozenset(run.name for run in inputs.check_runs if run.conclusion in OK_CHECK_CONCLUSIONS)
green_statuses: Final = frozenset(status.context for status in inputs.statuses if status.state == "success")
for context in sorted(inputs.required_contexts):
if context not in green_runs and context not in green_statuses:
reasons.append(f"required check {context!r} not green")
for run in inputs.check_runs:
if run.name == inputs.self_check_name:
continue
if run.status != "completed" or run.conclusion not in OK_CHECK_CONCLUSIONS:
reasons.append(f"check run {run.name!r} is {run.status}/{run.conclusion}")
for status in inputs.statuses:
if status.state != "success":
reasons.append(f"commit status {status.context!r} is {status.state}")
greptile: Final = tuple(
comment
for comment in inputs.comments
if comment.author_login == GREPTILE_LOGIN and GREPTILE_SCORE_RE.search(comment.body)
)
if not greptile:
reasons.append("greptile score not available")
else:
latest: Final = max(greptile, key=lambda comment: comment.updated_at)
match: Final = GREPTILE_SCORE_RE.search(latest.body)
score: Final = int(match.group(1)) if match else 0
if latest.updated_at < inputs.head_commit_date:
reasons.append("greptile score older than head commit")
elif score != 5:
reasons.append(f"greptile score {score}/5 below 5")
bugbot: Final = tuple(
review
for review in inputs.reviews
if review.author_login == BUGBOT_LOGIN
and BUGBOT_REVIEW_MARKER in review.body
and BUGBOT_STALE_MARKER not in review.body
and review.commit_id == pr.head_sha
)
if not bugbot:
reasons.append("bugbot review not available")
else:
latest_review: Final = max(bugbot, key=lambda review: review.submitted_at)
if BUGBOT_CLEAN not in latest_review.body:
reasons.append("bugbot reported issues")
latest_state_by_reviewer: Final[dict[str, str]] = {}
for review in sorted(inputs.reviews, key=lambda review: review.submitted_at):
if _is_bot_login(review.author_login):
continue
latest_state_by_reviewer[review.author_login] = review.state
for reviewer, state in latest_state_by_reviewer.items():
if state == "CHANGES_REQUESTED":
reasons.append(f"changes requested by {reviewer}")
return Verdict(merge=not reasons, reasons=tuple(reasons))
def _request(token: str, method: str, path: str, body: Mapping[str, object] | None = None) -> object:
url: Final = path if path.startswith("http") else f"{API_ROOT}{path}"
data: Final = None if body is None else json.dumps(body).encode("utf-8")
request: Final = urllib.request.Request(
url,
data=data,
method=method,
headers={
"Accept": "application/vnd.github+json",
"Authorization": f"Bearer {token}",
"X-GitHub-Api-Version": "2022-11-28",
},
)
with urllib.request.urlopen(request) as response:
return json.loads(response.read().decode("utf-8"))
def _request_allow_fail(
token: str, method: str, path: str, body: Mapping[str, object] | None = None
) -> tuple[int, object | None]:
url: Final = path if path.startswith("http") else f"{API_ROOT}{path}"
data: Final = None if body is None else json.dumps(body).encode("utf-8")
request: Final = urllib.request.Request(
url,
data=data,
method=method,
headers={
"Accept": "application/vnd.github+json",
"Authorization": f"Bearer {token}",
"X-GitHub-Api-Version": "2022-11-28",
},
)
try:
with urllib.request.urlopen(request) as response:
return response.status, json.loads(response.read().decode("utf-8"))
except urllib.error.HTTPError as exc:
return exc.code, None
def _items(payload: object, key: str | None = None) -> tuple[object, ...]:
source: Final = payload.get(key) if key and isinstance(payload, Mapping) else payload
if not isinstance(source, list):
return ()
return tuple(source)
def _paginate(token: str, path: str, key: str | None = None) -> list[object]:
separator: Final = "&" if "?" in path else "?"
results: list[object] = []
for page in range(1, 10_000):
batch: Final = _items(_request(token, "GET", f"{path}{separator}per_page=100&page={page}"), key)
results.extend(batch)
if len(batch) < 100:
return results
return results
def _text(value: object) -> str:
return value if isinstance(value, str) else ""
def _int(value: object) -> int:
return value if isinstance(value, int) else 0
def _bool(value: object) -> bool:
return value is True
def _nested(value: object, *keys: str) -> object:
current: object = value
for key in keys:
if not isinstance(current, Mapping):
return None
current = current.get(key)
return current
def _parse_time(value: object) -> datetime:
text: Final = _text(value)
if not text:
return datetime.min.replace(tzinfo=timezone.utc)
return datetime.fromisoformat(text.replace("Z", "+00:00"))
def _load_pr(token: str, repo: str, number: int) -> PullRequest:
data: Final = _request(token, "GET", f"/repos/{repo}/pulls/{number}")
if not isinstance(data, Mapping):
raise RuntimeError(f"unexpected pull payload for #{number}")
return PullRequest(
number=number,
title=_text(data.get("title")),
author_login=_text(_nested(data, "user", "login")),
state=_text(data.get("state")),
draft=_bool(data.get("draft")),
mergeable=data.get("mergeable") if isinstance(data.get("mergeable"), bool) else None,
mergeable_state=_text(data.get("mergeable_state")),
head_sha=_text(_nested(data, "head", "sha")),
)
def _list_candidate_prs(token: str, repo: str, base: str, allowlist: frozenset[str]) -> list[int]:
candidates: Final = _paginate(token, f"/repos/{repo}/pulls?state=open&base={base}")
return [
_int(item.get("number"))
for item in candidates
if isinstance(item, Mapping) and _text(_nested(item, "user", "login")).lower() in allowlist
]
def _changed_files(token: str, repo: str, number: int) -> tuple[str, ...]:
files: Final = _paginate(token, f"/repos/{repo}/pulls/{number}/files")
return tuple(_text(item.get("filename")) for item in files if isinstance(item, Mapping))
def _required_contexts(token: str, repo: str, base: str) -> frozenset[str]:
payload: Final = _request(token, "GET", f"/repos/{repo}/rules/branches/{base}")
contexts: set[str] = set()
for rule in _items(payload):
if not isinstance(rule, Mapping) or rule.get("type") != "required_status_checks":
continue
checks: Final = _nested(rule, "parameters", "required_status_checks")
for check in _items(checks):
if isinstance(check, Mapping):
context: Final = _text(check.get("context"))
if context:
contexts.add(context)
return frozenset(contexts)
def _check_runs(token: str, repo: str, sha: str) -> tuple[CheckRun, ...]:
runs: Final = _paginate(token, f"/repos/{repo}/commits/{sha}/check-runs", key="check_runs")
return tuple(
CheckRun(
name=_text(item.get("name")),
status=_text(item.get("status")),
conclusion=item.get("conclusion") if isinstance(item.get("conclusion"), str) else None,
)
for item in runs
if isinstance(item, Mapping)
)
def _statuses(token: str, repo: str, sha: str) -> tuple[CommitStatus, ...]:
payload: Final = _request(token, "GET", f"/repos/{repo}/commits/{sha}/status")
return tuple(
CommitStatus(context=_text(item.get("context")), state=_text(item.get("state")))
for item in _items(payload, "statuses")
if isinstance(item, Mapping)
)
def _comments(token: str, repo: str, number: int) -> tuple[IssueComment, ...]:
comments: Final = _paginate(token, f"/repos/{repo}/issues/{number}/comments")
return tuple(
IssueComment(
author_login=_text(_nested(item, "user", "login")),
body=_text(item.get("body")),
updated_at=_parse_time(item.get("updated_at")),
)
for item in comments
if isinstance(item, Mapping)
)
def _reviews(token: str, repo: str, number: int) -> tuple[Review, ...]:
reviews: Final = _paginate(token, f"/repos/{repo}/pulls/{number}/reviews")
return tuple(
Review(
author_login=_text(_nested(item, "user", "login")),
state=_text(item.get("state")),
body=_text(item.get("body")),
commit_id=_text(item.get("commit_id")),
submitted_at=_parse_time(item.get("submitted_at")),
)
for item in reviews
if isinstance(item, Mapping)
)
def _head_commit_date(token: str, repo: str, number: int) -> datetime:
commits: Final = _paginate(token, f"/repos/{repo}/pulls/{number}/commits")
if not commits:
return datetime.min.replace(tzinfo=timezone.utc)
last: Final = commits[-1]
if not isinstance(last, Mapping):
return datetime.min.replace(tzinfo=timezone.utc)
return _parse_time(_nested(last, "commit", "committer", "date"))
def _mergeable_or_refetch(token: str, repo: str, pr: PullRequest) -> PullRequest:
if pr.mergeable is not None:
return pr
time.sleep(5)
return _load_pr(token, repo, pr.number)
def _gather_inputs(
token: str,
repo: str,
number: int,
base: str,
self_check_name: str,
allowlist: frozenset[str],
) -> EvaluationInputs:
pr: Final = _mergeable_or_refetch(token, repo, _load_pr(token, repo, number))
return EvaluationInputs(
pr=pr,
changed_files=_changed_files(token, repo, number),
required_contexts=_required_contexts(token, repo, base),
check_runs=_check_runs(token, repo, pr.head_sha),
statuses=_statuses(token, repo, pr.head_sha),
comments=_comments(token, repo, number),
reviews=_reviews(token, repo, number),
head_commit_date=_head_commit_date(token, repo, number),
self_check_name=self_check_name,
author_allowlist=allowlist,
)
def merge_request_body(pr: PullRequest) -> dict[str, str]:
return {"merge_method": "merge", "commit_title": f"{pr.title} (#{pr.number})", "sha": pr.head_sha}
def _merge(token: str, repo: str, pr: PullRequest) -> None:
status, _ = _request_allow_fail(token, "PUT", f"/repos/{repo}/pulls/{pr.number}/merge", merge_request_body(pr))
if status in (200, 405, 409):
print(f"auto-merge-price-sync: PR #{pr.number} merge call returned {status}")
return
raise RuntimeError(f"merge call for PR #{pr.number} returned {status}")
def main() -> int:
token: Final = os.environ.get("GH_TOKEN", "")
repo: Final = os.environ.get("REPO", "")
base: Final = os.environ.get("BASE_BRANCH", "main")
dry_run: Final = os.environ.get("DRY_RUN", "") != ""
self_check_name: Final = os.environ.get("SELF_CHECK_NAME", "auto-merge-price-sync")
allowlist: Final = frozenset(login.lower() for login in os.environ.get("PR_AUTHOR_ALLOWLIST", "").split() if login)
if not token:
print("auto-merge-price-sync: app credentials not configured")
return 0
if not repo:
print("auto-merge-price-sync: REPO not set", file=sys.stderr)
return 1
pr_number_env: Final = os.environ.get("PR_NUMBER", "")
candidates: Final = [int(pr_number_env)] if pr_number_env else _list_candidate_prs(token, repo, base, allowlist)
for number in candidates:
inputs: Final = _gather_inputs(token, repo, number, base, self_check_name, allowlist)
verdict: Final = evaluate(inputs)
for reason in verdict.reasons:
print(f"auto-merge-price-sync: PR #{number} hold: {reason}")
if not verdict.merge:
continue
print(f"auto-merge-price-sync: PR #{number} all gates green")
if dry_run:
print(f"auto-merge-price-sync: DRY_RUN merge suppressed for PR #{number}")
continue
_merge(token, repo, inputs.pr)
return 0
if __name__ == "__main__":
sys.exit(main())

View file

@ -0,0 +1,61 @@
name: auto-merge-price-sync
on:
issue_comment:
types: [created, edited]
check_suite:
types: [completed]
status: {}
schedule:
- cron: "*/30 * * * *"
workflow_dispatch:
inputs:
pr-number:
description: "Evaluate only this PR number (empty = scan all open sync-bot PRs)"
required: false
default: ""
permissions:
contents: read
pull-requests: read
checks: read
statuses: read
concurrency:
group: auto-merge-price-sync
cancel-in-progress: false
jobs:
auto-merge-price-sync:
runs-on: ubuntu-latest
timeout-minutes: 15
env:
PROVIDER_INFO_SYNC_APP_ID: ${{ secrets.PROVIDER_INFO_SYNC_APP_ID }}
PROVIDER_INFO_SYNC_APP_PRIVATE_KEY: ${{ secrets.PROVIDER_INFO_SYNC_APP_PRIVATE_KEY }}
steps:
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
persist-credentials: false
- name: Set up Python
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
with:
python-version: "3.12"
- name: Mint app token
id: app-token
if: ${{ env.PROVIDER_INFO_SYNC_APP_ID != '' && env.PROVIDER_INFO_SYNC_APP_PRIVATE_KEY != '' }}
uses: actions/create-github-app-token@bcd2ba49218906704ab6c1aa796996da409d3eb1 # v3.2.0
with:
app-id: ${{ secrets.PROVIDER_INFO_SYNC_APP_ID }}
private-key: ${{ secrets.PROVIDER_INFO_SYNC_APP_PRIVATE_KEY }}
- name: Auto-merge eligible sync PRs
env:
GH_TOKEN: ${{ steps.app-token.outputs.token }}
REPO: ${{ github.repository }}
PR_NUMBER: ${{ (github.event.issue.pull_request && github.event.issue.number) || github.event.inputs.pr-number || '' }}
BASE_BRANCH: main
PR_AUTHOR_ALLOWLIST: "berriai-litellm-provider-info-sync[bot]"
SELF_CHECK_NAME: auto-merge-price-sync
run: python3 .github/scripts/auto_merge_price_sync.py

View file

@ -45,7 +45,7 @@ sequenceDiagram
ProxyServer->>Auth: user_api_key_auth()
Auth->>Redis: Check API key cache
Redis-->>Auth: Key info + spend limits
ProxyServer->>Hooks: max_budget_limiter, parallel_request_limiter
ProxyServer->>Hooks: parallel_request_limiter, cache_control_check
Hooks->>Redis: Check/increment rate limit counters
ProxyServer->>Router: route_request()
Router->>Main: litellm.acompletion()
@ -145,7 +145,6 @@ graph TD
| Hook | File | Purpose |
|------|------|---------|
| `max_budget_limiter` | `proxy/hooks/max_budget_limiter.py` | Enforce budget limits |
| `parallel_request_limiter` | `proxy/hooks/parallel_request_limiter_v3.py` | Rate limiting per key/user |
| `cache_control_check` | `proxy/hooks/cache_control_check.py` | Cache validation |
| `responses_id_security` | `proxy/hooks/responses_id_security.py` | Response ID validation |

143
litellm-rust/Cargo.lock generated
View file

@ -70,6 +70,12 @@ dependencies = [
"rustversion",
]
[[package]]
name = "arcstr"
version = "1.2.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "03918c3dbd7701a85c6b9887732e2921175f26c350b4563841d0958c21d57e6d"
[[package]]
name = "async-compression"
version = "0.4.46"
@ -262,6 +268,17 @@ dependencies = [
"tokio",
]
[[package]]
name = "aws-smithy-eventstream"
version = "0.61.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "5a9381123ab62d20c13082b151f30f962a3b112b727345394536dfa39a482944"
dependencies = [
"aws-smithy-types",
"bytes",
"crc32fast",
]
[[package]]
name = "aws-smithy-http"
version = "0.64.0"
@ -1837,6 +1854,12 @@ version = "0.2.186"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "68ab91017fe16c622486840e4c83c9a37afeff978bd239b5293d61ece587de66"
[[package]]
name = "linux-raw-sys"
version = "0.12.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "32a66949e030da00e8c7d4434b251670a91556f4144941d37452769c25d58a53"
[[package]]
name = "litellm-auth"
version = "0.1.0"
@ -1915,10 +1938,23 @@ dependencies = [
"tokio",
]
[[package]]
name = "litellm-cache-redis"
version = "0.1.0"
dependencies = [
"litellm-cache",
"redis",
"redis-test",
"serde_json",
"tokio",
]
[[package]]
name = "litellm-core"
version = "0.1.0"
dependencies = [
"aws-smithy-eventstream",
"aws-smithy-types",
"base64 0.22.1",
"bytes",
"data-url",
@ -1927,6 +1963,7 @@ dependencies = [
"litellm-auth-aws",
"litellm-auth-azure",
"litellm-auth-gcp",
"litellm-framing",
"mime_guess",
"moka",
"rand 0.8.7",
@ -1941,12 +1978,27 @@ dependencies = [
"strum",
"subtle",
"thiserror 2.0.19",
"time",
"tokio",
"tokio-tungstenite",
"url",
"veil",
]
[[package]]
name = "litellm-framing"
version = "0.1.0"
dependencies = [
"aws-smithy-eventstream",
"aws-smithy-types",
"bytes",
"futures-util",
"rstest",
"sse-stream",
"thiserror 2.0.19",
"tokio",
]
[[package]]
name = "litellm-python-bridge"
version = "0.1.0"
@ -2140,6 +2192,16 @@ dependencies = [
"minimal-lexical",
]
[[package]]
name = "num-bigint"
version = "0.5.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "93e7820bc0a80a0238e650327316f929ba18d5be054b647490a3a6a339f3e7c0"
dependencies = [
"num-integer",
"num-traits",
]
[[package]]
name = "num-conv"
version = "0.2.2"
@ -2656,6 +2718,36 @@ dependencies = [
"crossbeam-utils",
]
[[package]]
name = "redis"
version = "1.7.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "2acbc41a996f7652b2ddd9dfd98cc4ff602cfd742ae35382f07f608405ab50ed"
dependencies = [
"arcstr",
"combine",
"itoa",
"num-bigint",
"percent-encoding",
"ryu",
"sha1_smol",
"socket2 0.6.5",
"url",
"xxhash-rust",
]
[[package]]
name = "redis-test"
version = "1.0.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "804d36862e4323b69f96440cbb13c9894fc90176abdeaf91264e21d5d77f6aca"
dependencies = [
"rand 0.9.5",
"redis",
"socket2 0.6.5",
"tempfile",
]
[[package]]
name = "redox_syscall"
version = "0.5.18"
@ -2846,6 +2938,19 @@ dependencies = [
"semver",
]
[[package]]
name = "rustix"
version = "1.1.5"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "891efababe418670775f199f0d233d84843c227a0949a883ce15b37c78d6629d"
dependencies = [
"bitflags",
"errno",
"libc",
"linux-raw-sys",
"windows-sys 0.61.2",
]
[[package]]
name = "rustls"
version = "0.21.12"
@ -3096,6 +3201,12 @@ dependencies = [
"digest 0.10.7",
]
[[package]]
name = "sha1_smol"
version = "1.0.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "bbfa15b3dddfee50a0fff136974b3e1bde555604ba463834a7eb7deb6417705d"
[[package]]
name = "sha2"
version = "0.10.9"
@ -3200,6 +3311,19 @@ dependencies = [
"unicode-segmentation",
]
[[package]]
name = "sse-stream"
version = "0.2.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "c25ac7aff0abd1dbc474536e40416e1102c7dd9bfba0b9861c6d357f835dcfb4"
dependencies = [
"bytes",
"futures-util",
"http-body 1.1.0",
"http-body-util",
"pin-project-lite",
]
[[package]]
name = "stable_deref_trait"
version = "1.2.1"
@ -3299,6 +3423,19 @@ version = "0.13.5"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "adb6935a6f5c20170eeceb1a3835a49e12e19d792f6dd344ccc76a985ca5a6ca"
[[package]]
name = "tempfile"
version = "3.27.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "32497e9a4c7b38532efcdebeef879707aa9f794296a4f0244f6f69e9bc8574bd"
dependencies = [
"fastrand",
"getrandom 0.4.3",
"once_cell",
"rustix",
"windows-sys 0.61.2",
]
[[package]]
name = "thiserror"
version = "1.0.69"
@ -4182,6 +4319,12 @@ version = "0.13.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "66fee0b777b0f5ac1c69bb06d361268faafa61cd4682ae064a171c16c433e9e4"
[[package]]
name = "xxhash-rust"
version = "0.8.18"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "aee1b19627c7c60102ab80d3a9cbe18de90bfe03bfa6c3715447681f0e8c8af6"
[[package]]
name = "yoke"
version = "0.8.3"

View file

@ -11,6 +11,7 @@ repository = "https://github.com/BerriAI/litellm"
[workspace.dependencies]
bytes = "1"
litellm-core = { path = "crates/core" }
litellm-framing = { path = "crates/framer" }
litellm-auth = { path = "crates/auth" }
litellm-auth-aws = { path = "crates/auth-aws" }
litellm-auth-azure = { path = "crates/auth-azure" }
@ -39,6 +40,7 @@ base64 = "0.22"
moka = { version = "0.12.16", features = ["future"] }
strum = { version = "0.28.0", features = ["derive"] }
url = "2.5.8"
time = { version = "0.3.53", features = ["parsing"] }
criterion = "0.8.2"
veil = "0.3.0"

View file

@ -0,0 +1,15 @@
[package]
name = "litellm-cache-redis"
version = "0.1.0"
edition.workspace = true
license.workspace = true
repository.workspace = true
[dependencies]
litellm-cache.workspace = true
redis = "1.7.0"
serde_json.workspace = true
tokio.workspace = true
[dev-dependencies]
redis-test = "1.0.4"

View file

@ -0,0 +1,315 @@
use std::sync::{Arc, Mutex, MutexGuard};
use std::time::Duration;
use litellm_cache::{
BaseCache, CacheConnectionResult, CacheConnectionStatus, CacheEntry, CacheFuture, CacheKwargs,
Error,
};
use redis::Commands;
const DEFAULT_TTL: Duration = Duration::from_secs(600);
const KEY_PREFIX: &str = "litellm-cache:";
pub struct RedisCache<C = redis::Connection> {
connection: Arc<Mutex<C>>,
default_ttl: Duration,
}
impl RedisCache<redis::Connection> {
pub fn new(url: &str, default_ttl: Option<Duration>) -> Result<Self, Error> {
let client = redis::Client::open(url).map_err(|_| Error::Unavailable)?;
let connection = client.get_connection().map_err(|_| Error::Unavailable)?;
Ok(Self::with_connection(connection, default_ttl))
}
}
impl<C> RedisCache<C>
where
C: redis::ConnectionLike + Send + 'static,
{
fn with_connection(connection: C, default_ttl: Option<Duration>) -> Self {
Self {
connection: Arc::new(Mutex::new(connection)),
default_ttl: default_ttl.unwrap_or(DEFAULT_TTL),
}
}
fn connection(&self) -> Result<MutexGuard<'_, C>, Error> {
self.connection.lock().map_err(|_| Error::Unavailable)
}
fn namespaced_key(key: &str) -> String {
format!("{KEY_PREFIX}{key}")
}
fn namespaced_pattern() -> &'static str {
const PATTERN: &str = "litellm-cache:*";
PATTERN
}
fn encode(value: &CacheEntry) -> Result<Vec<u8>, Error> {
serde_json::to_vec(value).map_err(|_| Error::InvalidEntry)
}
fn decode(value: Vec<u8>) -> Result<CacheEntry, Error> {
serde_json::from_slice(&value).map_err(|_| Error::InvalidEntry)
}
fn ttl_seconds(ttl: Duration) -> u64 {
ttl.as_secs()
.saturating_add(u64::from(ttl.subsec_nanos() > 0))
.max(1)
}
fn run_blocking<T, F>(connection: Arc<Mutex<C>>, operation: F) -> CacheFuture<'static, T>
where
T: Send + 'static,
F: FnOnce(&mut C) -> Result<T, Error> + Send + 'static,
{
Box::pin(async move {
tokio::task::spawn_blocking(move || {
let mut connection = connection.lock().map_err(|_| Error::Unavailable)?;
operation(&mut connection)
})
.await
.map_err(|_| Error::Unavailable)?
})
}
}
impl<C> BaseCache for RedisCache<C>
where
C: redis::ConnectionLike + Send + 'static,
{
type Value = CacheEntry;
fn default_ttl(&self) -> Duration {
self.default_ttl
}
fn set_cache(&self, key: &str, value: Self::Value, kwargs: CacheKwargs) -> Result<(), Error> {
let payload = Self::encode(&value)?;
let ttl = Self::ttl_seconds(self.get_ttl(&kwargs));
self.connection()?
.set_ex::<_, _, ()>(Self::namespaced_key(key), payload, ttl)
.map_err(|_| Error::Unavailable)
}
fn get_cache(&self, key: &str, _: &CacheKwargs) -> Result<Option<Self::Value>, Error> {
self.connection()?
.get::<_, Option<Vec<u8>>>(Self::namespaced_key(key))
.map_err(|_| Error::Unavailable)?
.map(Self::decode)
.transpose()
}
fn delete_cache(&self, key: &str) -> Result<(), Error> {
self.connection()?
.del::<_, ()>(Self::namespaced_key(key))
.map_err(|_| Error::Unavailable)
}
fn flush_cache(&self) -> Result<(), Error> {
let mut connection = self.connection()?;
let keys = connection
.scan_match(Self::namespaced_pattern())
.map_err(|_| Error::Unavailable)?
.collect::<redis::RedisResult<Vec<String>>>()
.map_err(|_| Error::Unavailable)?;
if keys.is_empty() {
return Ok(());
}
connection
.del::<_, usize>(keys)
.map(|_| ())
.map_err(|_| Error::Unavailable)
}
fn async_set_cache<'a>(
&'a self,
key: &'a str,
value: Self::Value,
kwargs: CacheKwargs,
) -> CacheFuture<'a, ()> {
let payload = Self::encode(&value);
let key = Self::namespaced_key(key);
let ttl = Self::ttl_seconds(self.get_ttl(&kwargs));
Self::run_blocking(Arc::clone(&self.connection), move |connection| {
connection
.set_ex::<_, _, ()>(key, payload?, ttl)
.map_err(|_| Error::Unavailable)
})
}
fn async_get_cache<'a>(
&'a self,
key: &'a str,
_: &'a CacheKwargs,
) -> CacheFuture<'a, Option<Self::Value>> {
let key = Self::namespaced_key(key);
Box::pin(async move {
Self::run_blocking(Arc::clone(&self.connection), move |connection| {
connection
.get::<_, Option<Vec<u8>>>(key)
.map_err(|_| Error::Unavailable)
})
.await?
.map(Self::decode)
.transpose()
})
}
fn async_set_cache_pipeline<'a>(
&'a self,
cache_list: Vec<(String, Self::Value)>,
kwargs: CacheKwargs,
) -> CacheFuture<'a, ()> {
let entries = cache_list
.into_iter()
.map(|(key, value)| {
Self::encode(&value).map(|payload| (Self::namespaced_key(&key), payload))
})
.collect::<Result<Vec<_>, _>>();
let ttl = Self::ttl_seconds(self.get_ttl(&kwargs));
Self::run_blocking(Arc::clone(&self.connection), move |connection| {
for (key, payload) in entries? {
connection
.set_ex::<_, _, ()>(key, payload, ttl)
.map_err(|_| Error::Unavailable)?;
}
Ok(())
})
}
fn async_delete_cache<'a>(&'a self, key: &'a str) -> CacheFuture<'a, ()> {
let key = Self::namespaced_key(key);
Self::run_blocking(Arc::clone(&self.connection), move |connection| {
connection.del::<_, ()>(key).map_err(|_| Error::Unavailable)
})
}
fn disconnect(&self) -> CacheFuture<'_, ()> {
Box::pin(async { Ok(()) })
}
fn test_connection(&self) -> CacheFuture<'_, CacheConnectionResult> {
Box::pin(async move {
Self::run_blocking(Arc::clone(&self.connection), |connection| {
redis::cmd("PING")
.query::<String>(connection)
.map_err(|_| Error::Unavailable)
})
.await?;
Ok(CacheConnectionResult {
status: CacheConnectionStatus::Success,
message: "Redis cache connection test successful".into(),
error: None,
})
})
}
}
#[cfg(test)]
mod tests {
use super::RedisCache;
use litellm_cache::{BaseCache, CacheEntry, CacheKwargs};
use redis_test::{MockCmd, MockRedisConnection};
use serde_json::json;
use std::time::Duration;
fn entry() -> CacheEntry {
CacheEntry {
timestamp: 123.0,
response: json!({"choices": [{"text": "cached"}]}),
}
}
#[test]
fn cache_entries_round_trip_through_json() {
let entry = entry();
let encoded = RedisCache::<redis::Connection>::encode(&entry).unwrap();
assert_eq!(
RedisCache::<redis::Connection>::decode(encoded).unwrap(),
entry
);
}
#[test]
fn invalid_json_is_rejected() {
assert!(RedisCache::<redis::Connection>::decode(b"not json".to_vec()).is_err());
}
#[test]
fn ttl_seconds_rounds_up_and_keeps_expiration_positive() {
assert_eq!(
RedisCache::<redis::Connection>::ttl_seconds(Duration::ZERO),
1
);
assert_eq!(
RedisCache::<redis::Connection>::ttl_seconds(Duration::from_millis(1500)),
2
);
assert_eq!(
RedisCache::<redis::Connection>::ttl_seconds(Duration::from_secs(15)),
15
);
}
#[test]
fn redis_commands_round_trip_entries_and_delete_only_namespaced_keys() {
let value = entry();
let payload = RedisCache::<redis::Connection>::encode(&value).unwrap();
let connection = MockRedisConnection::new([
MockCmd::new(
redis::cmd("SETEX")
.arg("litellm-cache:key")
.arg(600)
.arg(payload.clone()),
Ok("OK"),
),
MockCmd::new(redis::cmd("GET").arg("litellm-cache:key"), Ok(payload)),
MockCmd::new(redis::cmd("DEL").arg("litellm-cache:key"), Ok(1u32)),
])
.assert_all_commands_consumed();
let cache = RedisCache::with_connection(connection, None);
cache
.set_cache("key", value.clone(), CacheKwargs::default())
.unwrap();
assert_eq!(
cache.get_cache("key", &CacheKwargs::default()).unwrap(),
Some(value)
);
cache.delete_cache("key").unwrap();
}
#[test]
fn flush_scans_and_deletes_only_cache_keys() {
let connection = MockRedisConnection::new([
MockCmd::new(
redis::cmd("SCAN")
.cursor_arg(0)
.arg("MATCH")
.arg("litellm-cache:*"),
Ok(redis_test::redis_value!(["0", ["litellm-cache:key"]])),
),
MockCmd::new(redis::cmd("DEL").arg("litellm-cache:key"), Ok(1u32)),
])
.assert_all_commands_consumed();
let cache = RedisCache::with_connection(connection, None);
cache.flush_cache().unwrap();
}
#[tokio::test]
async fn test_connection_runs_ping_off_executor() {
let connection = MockRedisConnection::new([MockCmd::new(redis::cmd("PING"), Ok("PONG"))])
.assert_all_commands_consumed();
let cache = RedisCache::with_connection(connection, None);
assert_eq!(
cache.test_connection().await.unwrap().status,
litellm_cache::CacheConnectionStatus::Success
);
}
}

View file

@ -0,0 +1,3 @@
mod cache;
pub use cache::RedisCache;

View file

@ -0,0 +1,6 @@
use litellm_cache_redis::RedisCache;
#[test]
fn constructor_rejects_invalid_urls() {
assert!(RedisCache::new("not a redis url", None).is_err());
}

View file

@ -15,6 +15,7 @@ litellm-auth.workspace = true
litellm-auth-aws.workspace = true
litellm-auth-azure.workspace = true
litellm-auth-gcp.workspace = true
litellm-framing.workspace = true
moka.workspace = true
mime_guess = "2.0.5"
rand.workspace = true
@ -29,9 +30,12 @@ subtle.workspace = true
tokio = { workspace = true, features = ["sync"] }
tokio-tungstenite.workspace = true
thiserror.workspace = true
time.workspace = true
sha2.workspace = true
url.workspace = true
veil.workspace = true
[dev-dependencies]
aws-smithy-eventstream = "=0.61.1"
aws-smithy-types = "1.6.1"
rstest.workspace = true

View file

@ -14,6 +14,7 @@ pub mod conversation;
pub(crate) mod handler;
mod prepare;
pub mod response_utils;
pub mod streaming;
pub mod transformation;
pub mod types;

View file

@ -0,0 +1,9 @@
pub trait StreamTransformer {
type Input;
type Output;
type Error;
fn transform(&mut self, input: Self::Input) -> Result<Vec<Self::Output>, Self::Error>;
fn finish(&mut self) -> Result<Vec<Self::Output>, Self::Error>;
}

View file

@ -120,3 +120,83 @@ pub struct ChatCompletionsResponse {
pub choices: Vec<ChatCompletionsChoice>,
pub usage: ChatCompletionsUsage,
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct ChatCompletionToolCallFunctionChunk {
#[serde(default, skip_serializing_if = "Option::is_none")]
pub name: Option<String>,
pub arguments: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub provider_specific_fields: Option<Map<String, Value>>,
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct ChatCompletionToolCallChunk {
#[serde(default, skip_serializing_if = "Option::is_none")]
pub id: Option<String>,
#[serde(rename = "type")]
pub tool_type: String,
pub function: ChatCompletionToolCallFunctionChunk,
pub index: i64,
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum ChatCompletionThinkingBlock {
Thinking {
#[serde(default, skip_serializing_if = "Option::is_none")]
thinking: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
signature: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
cache_control: Option<Value>,
},
RedactedThinking {
#[serde(default, skip_serializing_if = "Option::is_none")]
data: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
cache_control: Option<Value>,
},
}
#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
pub struct ChatCompletionDelta {
#[serde(default, skip_serializing_if = "Option::is_none")]
pub content: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub role: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub tool_calls: Option<Vec<ChatCompletionToolCallChunk>>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub reasoning_content: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub thinking_blocks: Option<Vec<ChatCompletionThinkingBlock>>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub provider_specific_fields: Option<Map<String, Value>>,
#[serde(flatten)]
pub extra: Map<String, Value>,
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct ChatCompletionStreamingChoice {
pub index: u64,
pub delta: ChatCompletionDelta,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub finish_reason: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub logprobs: Option<Value>,
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct ChatCompletionChunk {
pub id: String,
pub created: u64,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub model: Option<String>,
pub object: String,
pub choices: Vec<ChatCompletionStreamingChoice>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub usage: Option<ChatCompletionsUsage>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub provider_specific_fields: Option<Map<String, Value>>,
}

View file

@ -2,16 +2,54 @@
pub enum Error {
#[error("invalid provider: {0}")]
InvalidProvider(String),
#[error("missing required field: {0}")]
MissingField(&'static str),
#[error("invalid request: {0}")]
InvalidRequest(String),
#[error("invalid response: {0}")]
InvalidResponse(String),
#[error("routing error: {0}")]
Routing(String),
#[error("unsupported by the Rust messages route: {0}")]
Unsupported(&'static str),
#[error(transparent)]
Auth(#[from] litellm_auth::Error),
#[error(transparent)]
Transport(#[from] crate::transport::Error),
#[error(transparent)]
Headers(#[from] crate::http_utils::HeaderError),
#[error("stream framing failed: {0}")]
StreamFraming(String),
#[error("Anthropic SSE frame has no data")]
MissingStreamData,
#[error("Anthropic stream event is invalid: {0}")]
InvalidStreamEvent(String),
#[error("Bedrock event payload is invalid: {0}")]
InvalidBedrockPayload(String),
#[error("Bedrock event payload has invalid base64: {0}")]
InvalidBedrockBase64(String),
}
impl Error {
pub fn is_request(&self) -> bool {
match self {
Self::InvalidProvider(_)
| Self::MissingField(_)
| Self::InvalidRequest(_)
| Self::Unsupported(_)
| Self::Headers(_) => true,
Self::Auth(error) => !matches!(error, litellm_auth::Error::MissingApiKey { .. }),
_ => false,
}
}
pub fn is_response(&self) -> bool {
matches!(
self,
Self::InvalidResponse(_)
| Self::StreamFraming(_)
| Self::MissingStreamData
| Self::InvalidStreamEvent(_)
| Self::InvalidBedrockPayload(_)
| Self::InvalidBedrockBase64(_)
)
}
}

View file

@ -46,9 +46,7 @@ pub(super) async fn execute_messages_provider_stream(
) -> Result<reqwest::Response, Error> {
let request = prepare_provider_request(request)?;
if request.provider != ANTHROPIC_MESSAGES_PROVIDER {
return Err(Error::InvalidRequest(
"streaming messages is not supported for this provider".to_string(),
));
return Err(Error::Unsupported("streaming messages for this provider"));
}
let mut request_builder = http_client().post(&request.url).json(&request.body);

View file

@ -1 +1,2 @@
pub mod streaming;
pub mod transformation;

View file

@ -0,0 +1,164 @@
use std::collections::HashMap;
use serde_json::Value;
use crate::chat_completions::Error;
use crate::chat_completions::streaming::StreamTransformer;
use crate::chat_completions::types::{
ChatCompletionChunk, ChatCompletionThinkingBlock, ChatCompletionToolCallChunk,
ChatCompletionsUsage,
};
use crate::providers::anthropic::messages::streaming::{
AnthropicContentBlock, AnthropicContentBlockDelta, AnthropicMessagesStreamEvent,
AnthropicStreamUsage,
};
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum AnthropicJsonChunkType {
ValidJson,
AccumulatedJson,
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub enum AnthropicContentBlockType {
Text,
ToolUse,
ServerToolUse,
Thinking,
RedactedThinking,
Compaction,
ToolResult(String),
Other(String),
}
#[derive(Clone, Debug, PartialEq)]
pub struct AnthropicContentBlockDeltaEvent {
pub index: u64,
pub delta: AnthropicContentBlockDelta,
}
pub struct AnthropicChatCompletionsStreamTransformer {
pub content_blocks: Vec<AnthropicContentBlockDeltaEvent>,
pub tool_index: i64,
pub json_mode: bool,
pub speed: Option<String>,
pub tool_name_reverse_map: HashMap<String, String>,
pub response_id: String,
pub served_model: Option<String>,
pub is_response_format_tool: bool,
pub converted_response_format_tool: bool,
pub accumulated_json: String,
pub chunk_type: AnthropicJsonChunkType,
pub current_content_block_type: Option<AnthropicContentBlockType>,
pub web_search_results: Vec<Value>,
pub web_search_calls: HashMap<String, Value>,
pub compaction_blocks: Vec<Value>,
pub reasoning_content_chunks: Vec<String>,
pub server_tool_inputs: HashMap<String, Value>,
pub tool_results: Vec<Value>,
pub current_server_tool_id: Option<String>,
pub container_id: Option<String>,
}
impl AnthropicChatCompletionsStreamTransformer {
pub fn new(
_json_mode: bool,
_speed: Option<String>,
_tool_name_reverse_map: HashMap<String, String>,
) -> Self {
todo!()
}
pub fn check_empty_tool_call_args(&self) -> bool {
todo!()
}
pub fn handle_usage(&mut self, _usage: AnthropicStreamUsage) -> ChatCompletionsUsage {
todo!()
}
pub fn handle_content_block_delta(
&mut self,
_index: u64,
_delta: AnthropicContentBlockDelta,
) -> (
String,
Option<ChatCompletionToolCallChunk>,
Vec<ChatCompletionThinkingBlock>,
Option<Value>,
Option<String>,
) {
todo!()
}
pub fn handle_content_block_start(
&mut self,
_index: u64,
_content_block: AnthropicContentBlock,
) -> Result<ChatCompletionChunk, Error> {
todo!()
}
pub fn handle_json_mode_chunk(
&mut self,
_text: String,
_tool_use: Option<ChatCompletionToolCallChunk>,
) -> (String, Option<ChatCompletionToolCallChunk>) {
todo!()
}
pub fn handle_accumulated_json_chunk(
&mut self,
_data: &str,
_is_final: bool,
) -> Result<Option<ChatCompletionChunk>, Error> {
todo!()
}
pub fn handle_redacted_thinking_content(
&mut self,
_content_block: &AnthropicContentBlock,
) -> Vec<ChatCompletionThinkingBlock> {
todo!()
}
pub fn web_search_call_snapshot(&self) -> HashMap<String, Value> {
todo!()
}
pub fn complete_web_search_call(&mut self, _result: Value) {
todo!()
}
pub fn build_code_interpreter_results(&self) -> Vec<Value> {
todo!()
}
pub fn handle_message_delta(
&mut self,
_event: AnthropicMessagesStreamEvent,
) -> (Option<String>, Option<ChatCompletionsUsage>, Option<Value>) {
todo!()
}
pub fn chunk_parser(
&mut self,
_event: AnthropicMessagesStreamEvent,
) -> Result<ChatCompletionChunk, Error> {
todo!()
}
}
impl StreamTransformer for AnthropicChatCompletionsStreamTransformer {
type Input = AnthropicMessagesStreamEvent;
type Output = ChatCompletionChunk;
type Error = Error;
fn transform(&mut self, _input: Self::Input) -> Result<Vec<Self::Output>, Self::Error> {
todo!()
}
fn finish(&mut self) -> Result<Vec<Self::Output>, Self::Error> {
todo!()
}
}

View file

@ -0,0 +1,338 @@
use serde::{Deserialize, Serialize};
use serde_json::Value;
use time::OffsetDateTime;
use url::Url;
use crate::messages::Error;
use crate::messages::types::AnthropicMessagesResponse;
use crate::providers::anthropic::messages::transformation::resolve_anthropic_api_base;
const BATCHES_PATH_SUFFIX: &str = "/v1/messages/batches";
#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct AnthropicBatchRequestCounts {
#[serde(default)]
pub processing: u64,
#[serde(default)]
pub succeeded: u64,
#[serde(default)]
pub errored: u64,
#[serde(default)]
pub canceled: u64,
#[serde(default)]
pub expired: u64,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct AnthropicMessageBatch {
#[serde(default)]
pub id: String,
#[serde(default = "default_processing_status")]
pub processing_status: String,
pub created_at: Option<String>,
pub ended_at: Option<String>,
pub expires_at: Option<String>,
pub cancel_initiated_at: Option<String>,
pub archived_at: Option<String>,
#[serde(default)]
pub request_counts: AnthropicBatchRequestCounts,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum BatchStatus {
InProgress,
Cancelling,
Completed,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct BatchRequestCounts {
pub total: u64,
pub completed: u64,
pub failed: u64,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct LiteLlmMessageBatch {
pub id: String,
pub object: String,
pub endpoint: String,
pub input_file_id: String,
pub completion_window: String,
pub status: BatchStatus,
pub output_file_id: String,
pub created_at: i64,
pub in_progress_at: Option<i64>,
pub expires_at: Option<i64>,
pub completed_at: Option<i64>,
pub expired_at: Option<i64>,
pub cancelling_at: Option<i64>,
pub cancelled_at: Option<i64>,
pub request_counts: BatchRequestCounts,
}
pub trait AnthropicBatchesConfig {
fn create_batch_url(
&self,
api_base: Option<&str>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Result<String, Error>;
fn transform_create_batch_request(&self) -> Result<Value, Error>;
fn transform_create_batch_response(
&self,
response: AnthropicMessageBatch,
now: i64,
) -> Result<LiteLlmMessageBatch, Error>;
fn retrieve_batch_url(
&self,
api_base: Option<&str>,
batch_id: &str,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Result<String, Error>;
fn transform_retrieve_batch_request(&self) -> Value;
fn transform_retrieve_batch_response(
&self,
response: AnthropicMessageBatch,
now: i64,
) -> LiteLlmMessageBatch;
fn transform_batch_results(&self, body: &str) -> Result<Vec<AnthropicMessagesResponse>, Error>;
}
pub struct AnthropicBatchesTransformation;
pub const ANTHROPIC_BATCHES_TRANSFORMATION: AnthropicBatchesTransformation =
AnthropicBatchesTransformation;
fn default_processing_status() -> String {
"in_progress".into()
}
fn timestamp(value: Option<&str>) -> Option<i64> {
value
.and_then(|value| {
OffsetDateTime::parse(value, &time::format_description::well_known::Rfc3339).ok()
})
.map(OffsetDateTime::unix_timestamp)
}
fn batches_base_url(
api_base: Option<&str>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Result<Url, Error> {
let api_base = resolve_anthropic_api_base(api_base, env_lookup);
let api_base = api_base.trim_end_matches('/');
let complete_url = if api_base.ends_with(BATCHES_PATH_SUFFIX) {
api_base.to_string()
} else if let Some(base) = api_base.strip_suffix("/v1/messages") {
format!("{base}{BATCHES_PATH_SUFFIX}")
} else {
format!("{api_base}{BATCHES_PATH_SUFFIX}")
};
Url::parse(&complete_url)
.map_err(|error| Error::InvalidRequest(format!("invalid Anthropic API base: {error}")))
}
impl AnthropicBatchesConfig for AnthropicBatchesTransformation {
fn create_batch_url(
&self,
api_base: Option<&str>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Result<String, Error> {
Ok(batches_base_url(api_base, env_lookup)?.into())
}
fn transform_create_batch_request(&self) -> Result<Value, Error> {
Err(Error::Unsupported("Anthropic message batch creation"))
}
fn transform_create_batch_response(
&self,
_response: AnthropicMessageBatch,
_now: i64,
) -> Result<LiteLlmMessageBatch, Error> {
Err(Error::Unsupported("Anthropic message batch creation"))
}
fn retrieve_batch_url(
&self,
api_base: Option<&str>,
batch_id: &str,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Result<String, Error> {
if batch_id.is_empty() {
return Err(Error::MissingField("batch_id"));
}
let mut url = batches_base_url(api_base, env_lookup)?;
url.path_segments_mut()
.map_err(|_| Error::InvalidRequest("Anthropic API base cannot be a base URL".into()))?
.push(batch_id);
Ok(url.into())
}
fn transform_retrieve_batch_request(&self) -> Value {
Value::Object(Default::default())
}
fn transform_retrieve_batch_response(
&self,
response: AnthropicMessageBatch,
now: i64,
) -> LiteLlmMessageBatch {
let created_at = timestamp(response.created_at.as_deref());
let ended_at = timestamp(response.ended_at.as_deref());
let expires_at = timestamp(response.expires_at.as_deref());
let cancel_initiated_at = timestamp(response.cancel_initiated_at.as_deref());
let archived_at = timestamp(response.archived_at.as_deref());
let status = match response.processing_status.as_str() {
"canceling" => BatchStatus::Cancelling,
"ended" => BatchStatus::Completed,
_ => BatchStatus::InProgress,
};
let request_counts = BatchRequestCounts {
total: response.request_counts.processing
+ response.request_counts.succeeded
+ response.request_counts.errored
+ response.request_counts.canceled
+ response.request_counts.expired,
completed: response.request_counts.succeeded,
failed: response.request_counts.errored,
};
LiteLlmMessageBatch {
id: response.id.clone(),
object: "batch".into(),
endpoint: "/v1/messages".into(),
input_file_id: "None".into(),
completion_window: "24h".into(),
status,
output_file_id: response.id,
created_at: created_at.unwrap_or(now),
in_progress_at: (response.processing_status == "in_progress")
.then_some(created_at)
.flatten(),
expires_at,
completed_at: (response.processing_status == "ended")
.then_some(ended_at)
.flatten(),
expired_at: archived_at,
cancelling_at: (response.processing_status == "canceling")
.then_some(cancel_initiated_at)
.flatten(),
cancelled_at: (response.processing_status == "canceling")
.then_some(ended_at)
.flatten(),
request_counts,
}
}
fn transform_batch_results(&self, body: &str) -> Result<Vec<AnthropicMessagesResponse>, Error> {
body.lines()
.filter(|line| !line.trim().is_empty())
.filter_map(|line| serde_json::from_str::<Value>(line.trim()).ok())
.map(|record| {
serde_json::from_value(record["result"]["message"].clone()).map_err(|error| {
Error::InvalidResponse(format!("invalid Anthropic batch result: {error}"))
})
})
.collect()
}
}
#[cfg(test)]
mod tests {
use serde_json::json;
use super::*;
#[test]
fn builds_and_encodes_message_batch_urls() {
assert_eq!(
ANTHROPIC_BATCHES_TRANSFORMATION
.create_batch_url(None, &|_| None)
.unwrap(),
"https://api.anthropic.com/v1/messages/batches"
);
assert_eq!(
ANTHROPIC_BATCHES_TRANSFORMATION
.create_batch_url(Some("https://proxy.test/v1/messages/batches"), &|_| None)
.unwrap(),
"https://proxy.test/v1/messages/batches"
);
assert_eq!(
ANTHROPIC_BATCHES_TRANSFORMATION
.retrieve_batch_url(Some("https://proxy.test"), "batch/id ?", &|_| None)
.unwrap(),
"https://proxy.test/v1/messages/batches/batch%2Fid%20%3F"
);
assert_eq!(
ANTHROPIC_BATCHES_TRANSFORMATION.transform_retrieve_batch_request(),
json!({})
);
}
#[test]
fn maps_retrieved_batch_status_counts_and_timestamps_like_python() {
let response: AnthropicMessageBatch = serde_json::from_value(json!({
"id": "msgbatch_1",
"processing_status": "ended",
"created_at": "2025-01-01T00:00:00Z",
"ended_at": "2025-01-01T00:01:00Z",
"expires_at": "not-a-timestamp",
"request_counts": {
"processing": 1,
"succeeded": 2,
"errored": 3,
"canceled": 4,
"expired": 5
}
}))
.unwrap();
let batch = ANTHROPIC_BATCHES_TRANSFORMATION.transform_retrieve_batch_response(response, 7);
assert_eq!(batch.status, BatchStatus::Completed);
assert_eq!(batch.created_at, 1_735_689_600);
assert_eq!(batch.completed_at, Some(1_735_689_660));
assert_eq!(batch.expires_at, None);
assert_eq!(
batch.request_counts,
BatchRequestCounts {
total: 15,
completed: 2,
failed: 3
}
);
}
#[test]
fn extracts_message_responses_from_ndjson_and_skips_non_json_lines() {
let body = r#"not-json
{"result":{"message":{"id":"msg_1","type":"message","role":"assistant","model":"claude-test","content":[],"stop_reason":"end_turn","stop_sequence":null}}}
"#;
let messages = ANTHROPIC_BATCHES_TRANSFORMATION
.transform_batch_results(body)
.unwrap();
assert_eq!(messages.len(), 1);
assert_eq!(messages[0].id, "msg_1");
}
#[test]
fn preserves_python_placeholder_for_batch_creation() {
assert!(matches!(
ANTHROPIC_BATCHES_TRANSFORMATION.transform_create_batch_request(),
Err(Error::Unsupported("Anthropic message batch creation"))
));
let response: AnthropicMessageBatch = serde_json::from_value(json!({})).unwrap();
assert!(matches!(
ANTHROPIC_BATCHES_TRANSFORMATION.transform_create_batch_response(response, 0),
Err(Error::Unsupported("Anthropic message batch creation"))
));
}
}

View file

@ -0,0 +1,168 @@
use serde::{Deserialize, Serialize};
use serde_json::Value;
use crate::constants::ANTHROPIC_OAUTH_TOKEN_PREFIX;
use crate::messages::Error;
use crate::messages::types::{AnthropicMessage, SystemPrompt};
const COUNT_TOKENS_ENDPOINT: &str = "https://api.anthropic.com/v1/messages/count_tokens";
const TOKEN_COUNTING_BETA: &str = "token-counting-2024-11-01";
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct AnthropicCountTokensRequest {
pub model: String,
pub messages: Vec<AnthropicMessage>,
#[serde(skip_serializing_if = "Option::is_none")]
pub tools: Option<Vec<Value>>,
#[serde(skip_serializing_if = "Option::is_none")]
pub system: Option<SystemPrompt>,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct AnthropicCountTokensResponse {
pub input_tokens: u64,
}
pub trait AnthropicCountTokensConfig {
fn endpoint(&self) -> &'static str;
fn validate_request(&self, model: &str, messages: &[AnthropicMessage]) -> Result<(), Error>;
fn transform_request(
&self,
model: &str,
messages: Vec<AnthropicMessage>,
tools: Option<Vec<Value>>,
system: Option<SystemPrompt>,
) -> Result<AnthropicCountTokensRequest, Error>;
fn required_headers(&self, api_key: &str) -> Vec<(&'static str, String)>;
}
pub struct AnthropicCountTokensTransformation;
pub const ANTHROPIC_COUNT_TOKENS_TRANSFORMATION: AnthropicCountTokensTransformation =
AnthropicCountTokensTransformation;
impl AnthropicCountTokensConfig for AnthropicCountTokensTransformation {
fn endpoint(&self) -> &'static str {
COUNT_TOKENS_ENDPOINT
}
fn transform_request(
&self,
model: &str,
messages: Vec<AnthropicMessage>,
tools: Option<Vec<Value>>,
system: Option<SystemPrompt>,
) -> Result<AnthropicCountTokensRequest, Error> {
self.validate_request(model, &messages)?;
Ok(AnthropicCountTokensRequest {
model: model.to_string(),
messages,
tools,
system,
})
}
fn validate_request(&self, model: &str, messages: &[AnthropicMessage]) -> Result<(), Error> {
if model.is_empty() {
return Err(Error::MissingField("model"));
}
if messages.is_empty() {
return Err(Error::MissingField("messages"));
}
Ok(())
}
fn required_headers(&self, api_key: &str) -> Vec<(&'static str, String)> {
let auth = if api_key.starts_with(ANTHROPIC_OAUTH_TOKEN_PREFIX) {
("authorization", format!("Bearer {api_key}"))
} else {
("x-api-key", api_key.to_string())
};
vec![
("content-type", "application/json".to_string()),
auth,
("anthropic-version", "2023-06-01".to_string()),
("anthropic-beta", TOKEN_COUNTING_BETA.to_string()),
]
}
}
#[cfg(test)]
mod tests {
use serde_json::{Map, json};
use super::*;
use crate::messages::types::MessageContent;
fn message() -> AnthropicMessage {
AnthropicMessage {
role: "user".into(),
content: MessageContent::Text("hello".into()),
extra: Map::new(),
}
}
#[test]
fn maps_the_python_count_tokens_contract() {
let request = ANTHROPIC_COUNT_TOKENS_TRANSFORMATION
.transform_request(
"claude-test",
vec![message()],
Some(vec![json!({"name": "lookup"})]),
Some(SystemPrompt::Text("system".into())),
)
.unwrap();
assert_eq!(
serde_json::to_value(request).unwrap(),
json!({
"model": "claude-test",
"messages": [{"role": "user", "content": "hello"}],
"tools": [{"name": "lookup"}],
"system": "system"
})
);
assert_eq!(
ANTHROPIC_COUNT_TOKENS_TRANSFORMATION.endpoint(),
COUNT_TOKENS_ENDPOINT
);
}
#[test]
fn rejects_the_invalid_requests_python_rejects() {
assert!(matches!(
ANTHROPIC_COUNT_TOKENS_TRANSFORMATION.transform_request(
"",
vec![message()],
None,
None
),
Err(Error::MissingField("model"))
));
assert!(matches!(
ANTHROPIC_COUNT_TOKENS_TRANSFORMATION.transform_request(
"claude-test",
vec![],
None,
None
),
Err(Error::MissingField("messages"))
));
}
#[test]
fn uses_api_key_or_oauth_headers_without_combining_credentials() {
let api_key = ANTHROPIC_COUNT_TOKENS_TRANSFORMATION.required_headers("sk-ant-api");
assert!(api_key.contains(&("x-api-key", "sk-ant-api".into())));
assert!(!api_key.iter().any(|(name, _)| *name == "authorization"));
let oauth = ANTHROPIC_COUNT_TOKENS_TRANSFORMATION.required_headers("sk-ant-oat-test");
assert!(oauth.contains(&("authorization", "Bearer sk-ant-oat-test".into())));
assert!(!oauth.iter().any(|(name, _)| *name == "x-api-key"));
assert!(oauth.contains(&("anthropic-beta", TOKEN_COUNTING_BETA.into())));
}
}

View file

@ -1 +1,4 @@
pub mod batches;
pub mod count_tokens;
pub mod streaming;
pub mod transformation;

View file

@ -0,0 +1,282 @@
use base64::Engine;
use bytes::Buf;
use futures_util::{Stream, StreamExt};
use litellm_framing::Framer;
use litellm_framing::aws_event_stream::{AwsEventStreamFrame, AwsEventStreamFramer};
use litellm_framing::sse::{SseFrame, SseFramer};
use serde::{Deserialize, Serialize};
use serde_json::{Map, Value};
use crate::messages::Error;
#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
pub struct AnthropicStreamUsage {
#[serde(default)]
pub input_tokens: u64,
#[serde(default)]
pub output_tokens: u64,
#[serde(default)]
pub cache_creation_input_tokens: u64,
#[serde(default)]
pub cache_read_input_tokens: u64,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub server_tool_use: Option<Value>,
#[serde(flatten)]
pub extra: Map<String, Value>,
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct AnthropicStreamMessage {
pub id: String,
#[serde(rename = "type")]
pub message_type: String,
pub role: String,
pub model: String,
pub content: Vec<Value>,
pub stop_reason: Option<String>,
pub stop_sequence: Option<String>,
pub usage: AnthropicStreamUsage,
#[serde(flatten)]
pub extra: Map<String, Value>,
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum AnthropicContentBlockDelta {
TextDelta {
text: String,
},
InputJsonDelta {
partial_json: String,
},
#[serde(rename = "citations_delta")]
Citations {
citation: Value,
},
ThinkingDelta {
thinking: String,
},
SignatureDelta {
signature: String,
},
CompactionDelta {
content: String,
},
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct AnthropicContentBlock {
#[serde(rename = "type")]
pub block_type: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub id: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub name: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub text: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub input: Option<Value>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub thinking: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub signature: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub data: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub content: Option<Value>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub caller: Option<Value>,
#[serde(flatten)]
pub extra: Map<String, Value>,
}
#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
pub struct AnthropicMessageDelta {
#[serde(default, skip_serializing_if = "Option::is_none")]
pub stop_reason: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub stop_sequence: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub stop_details: Option<Value>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub container: Option<Value>,
#[serde(flatten)]
pub extra: Map<String, Value>,
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct AnthropicStreamError {
#[serde(rename = "type")]
pub error_type: String,
pub message: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub details: Option<Value>,
#[serde(flatten)]
pub extra: Map<String, Value>,
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum AnthropicMessagesStreamEvent {
MessageStart {
message: AnthropicStreamMessage,
},
ContentBlockStart {
index: u64,
content_block: AnthropicContentBlock,
},
ContentBlockDelta {
index: u64,
delta: AnthropicContentBlockDelta,
},
ContentBlockStop {
index: u64,
},
MessageDelta {
delta: AnthropicMessageDelta,
#[serde(default, skip_serializing_if = "Option::is_none")]
usage: Option<AnthropicStreamUsage>,
#[serde(default, skip_serializing_if = "Option::is_none")]
context_management: Option<Value>,
},
MessageStop,
Ping,
Error {
error: AnthropicStreamError,
},
}
#[derive(Deserialize)]
struct BedrockChunkPayload {
bytes: String,
}
pub fn decode_anthropic_sse_frame(frame: SseFrame) -> Result<AnthropicMessagesStreamEvent, Error> {
let data = frame.data.ok_or(Error::MissingStreamData)?;
serde_json::from_str(&data).map_err(|error| Error::InvalidStreamEvent(error.to_string()))
}
pub fn decode_bedrock_anthropic_frame(
frame: AwsEventStreamFrame,
) -> Result<AnthropicMessagesStreamEvent, Error> {
let payload: BedrockChunkPayload = serde_json::from_slice(&frame.payload)
.map_err(|error| Error::InvalidBedrockPayload(error.to_string()))?;
let event = base64::engine::general_purpose::STANDARD
.decode(payload.bytes)
.map_err(|error| Error::InvalidBedrockBase64(error.to_string()))?;
serde_json::from_slice(&event).map_err(|error| Error::InvalidStreamEvent(error.to_string()))
}
pub fn direct_anthropic_event_stream<S, B, E>(
input: S,
) -> impl Stream<Item = Result<AnthropicMessagesStreamEvent, Error>> + Send
where
S: Stream<Item = Result<B, E>> + Send,
B: Buf + Send,
E: std::error::Error + Send + Sync + 'static,
{
SseFramer.frame(input).map(|frame| {
let frame = frame.map_err(|error| Error::StreamFraming(error.to_string()))?;
decode_anthropic_sse_frame(frame)
})
}
pub fn bedrock_anthropic_event_stream<S, B, E>(
input: S,
) -> impl Stream<Item = Result<AnthropicMessagesStreamEvent, Error>> + Send
where
S: Stream<Item = Result<B, E>> + Send,
B: Buf + Send,
E: std::error::Error + Send + Sync + 'static,
{
AwsEventStreamFramer.frame(input).map(|frame| {
let frame = frame.map_err(|error| Error::StreamFraming(error.to_string()))?;
decode_bedrock_anthropic_frame(frame)
})
}
#[cfg(test)]
mod tests {
use std::io;
use aws_smithy_eventstream::frame::write_message_to;
use aws_smithy_types::event_stream::{Header, HeaderValue, Message};
use base64::engine::general_purpose::STANDARD;
use bytes::Bytes;
use futures_util::TryStreamExt;
use super::*;
const TEXT_DELTA: &str =
r#"{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"hello"}}"#;
#[tokio::test]
async fn direct_anthropic_sse_frames_into_typed_events() {
let wire = format!("event: content_block_delta\ndata: {TEXT_DELTA}\n\n");
let events = direct_anthropic_event_stream(futures_util::stream::iter(
wire.as_bytes().chunks(3).map(Ok::<_, io::Error>),
))
.try_collect::<Vec<_>>()
.await
.unwrap();
assert_eq!(
events,
vec![AnthropicMessagesStreamEvent::ContentBlockDelta {
index: 0,
delta: AnthropicContentBlockDelta::TextDelta {
text: "hello".into(),
},
}]
);
}
#[test]
fn decodes_citations_delta_events() {
let event = decode_anthropic_sse_frame(SseFrame {
event: Some("content_block_delta".into()),
data: Some(
r#"{"type":"content_block_delta","index":0,"delta":{"type":"citations_delta","citation":{"type":"char_location"}}}"#
.into(),
),
id: None,
retry: None,
})
.unwrap();
assert!(matches!(
event,
AnthropicMessagesStreamEvent::ContentBlockDelta {
delta: AnthropicContentBlockDelta::Citations { .. },
..
}
));
}
#[tokio::test]
async fn bedrock_aws_frames_into_the_same_typed_events() {
let payload = serde_json::json!({"bytes": STANDARD.encode(TEXT_DELTA)});
let message = Message::new(Bytes::from(serde_json::to_vec(&payload).unwrap())).add_header(
Header::new(":event-type", HeaderValue::String("chunk".into())),
);
let mut wire = Vec::new();
write_message_to(&message, &mut wire).unwrap();
let events = bedrock_anthropic_event_stream(futures_util::stream::iter(
wire.chunks(3).map(Ok::<_, io::Error>),
))
.try_collect::<Vec<_>>()
.await
.unwrap();
assert_eq!(
events,
vec![AnthropicMessagesStreamEvent::ContentBlockDelta {
index: 0,
delta: AnthropicContentBlockDelta::TextDelta {
text: "hello".into(),
},
}]
);
}
}

View file

@ -31,10 +31,7 @@ pub fn complete_anthropic_url(
api_base: Option<&str>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> String {
let api_base = non_empty(api_base)
.map(str::to_string)
.or_else(|| env_lookup(ANTHROPIC_API_BASE_ENV).filter(|value| !value.trim().is_empty()))
.unwrap_or_else(|| DEFAULT_ANTHROPIC_API_BASE.to_string());
let api_base = resolve_anthropic_api_base(api_base, env_lookup);
let api_base = api_base.trim_end_matches('/');
if api_base.ends_with(MESSAGES_PATH_SUFFIX) {
@ -43,6 +40,16 @@ pub fn complete_anthropic_url(
format!("{api_base}{MESSAGES_PATH_SUFFIX}")
}
pub fn resolve_anthropic_api_base(
api_base: Option<&str>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> String {
non_empty(api_base)
.map(str::to_string)
.or_else(|| env_lookup(ANTHROPIC_API_BASE_ENV).filter(|value| !value.trim().is_empty()))
.unwrap_or_else(|| DEFAULT_ANTHROPIC_API_BASE.to_string())
}
impl AnthropicMessagesProviderConfig for AnthropicMessagesConfig {
fn complete_url(
&self,

View file

@ -0,0 +1,23 @@
[package]
name = "litellm-framing"
version = "0.1.0"
edition.workspace = true
license.workspace = true
repository.workspace = true
[features]
default = ["aws", "sse"]
aws = ["dep:aws-smithy-eventstream", "dep:aws-smithy-types"]
sse = ["dep:sse-stream"]
[dependencies]
aws-smithy-eventstream = { version = "=0.61.1", optional = true }
aws-smithy-types = { version = "1.6.1", optional = true }
bytes = "1"
futures-util.workspace = true
sse-stream = { version = "=0.2.6", optional = true }
thiserror.workspace = true
[dev-dependencies]
rstest.workspace = true
tokio.workspace = true

View file

@ -0,0 +1,66 @@
use bytes::{Buf, Bytes, BytesMut};
use futures_util::{Stream, StreamExt};
use aws_smithy_eventstream::frame::read_message_from;
use aws_smithy_types::event_stream::Header;
use crate::{Error, Framer};
const MAX_FRAME_BYTES: usize = 16 * 1024 * 1024;
#[derive(Clone, Debug, PartialEq)]
pub struct AwsEventStreamFrame {
pub headers: Vec<Header>,
pub payload: Bytes,
}
#[derive(Clone, Copy, Debug, Default)]
pub struct AwsEventStreamFramer;
impl Framer for AwsEventStreamFramer {
type Frame = AwsEventStreamFrame;
fn frame<S, B, E>(self, input: S) -> impl Stream<Item = Result<Self::Frame, Error>> + Send
where
S: Stream<Item = Result<B, E>> + Send,
B: Buf + Send,
E: std::error::Error + Send + Sync + 'static,
{
futures_util::stream::try_unfold(
(Box::pin(input), BytesMut::new()),
|(mut input, mut buffer)| async move {
loop {
if buffer.len() >= 4 {
let length = (&buffer[..4]).get_u32() as usize;
if !(16..=MAX_FRAME_BYTES).contains(&length) {
return Err(Error::InvalidLength(length));
}
if buffer.len() >= length {
let raw = buffer.split_to(length).freeze();
let message = read_message_from(raw)?;
let frame = AwsEventStreamFrame {
headers: message.headers().to_vec(),
payload: message.payload().clone(),
};
return Ok(Some((frame, (input, buffer))));
}
}
match input.next().await {
Some(Ok(mut chunk)) => {
while chunk.has_remaining() {
let bytes = chunk.chunk();
buffer.extend_from_slice(bytes);
let length = bytes.len();
chunk.advance(length);
}
}
Some(Err(error)) => return Err(Error::Body(Box::new(error))),
None if buffer.is_empty() => return Ok(None),
None => return Err(Error::Truncated),
}
}
},
)
.fuse()
}
}

View file

@ -0,0 +1,17 @@
#[derive(Debug, thiserror::Error)]
pub enum Error {
#[cfg(feature = "sse")]
#[error("SSE framing failed: {0}")]
Sse(#[from] sse_stream::Error),
#[cfg(feature = "aws")]
#[error("AWS EventStream framing failed: {0}")]
Aws(#[from] aws_smithy_eventstream::error::Error),
#[error("body stream failed: {0}")]
Body(#[source] Box<dyn std::error::Error + Send + Sync>),
#[cfg(feature = "aws")]
#[error("invalid AWS EventStream frame length: {0}")]
InvalidLength(usize),
#[cfg(feature = "aws")]
#[error("truncated AWS EventStream frame")]
Truncated,
}

View file

@ -0,0 +1,13 @@
use futures_util::Stream;
use crate::Error;
pub trait Framer: Send {
type Frame: Send;
fn frame<S, B, E>(self, input: S) -> impl Stream<Item = Result<Self::Frame, Error>> + Send
where
S: Stream<Item = Result<B, E>> + Send,
B: bytes::Buf + Send,
E: std::error::Error + Send + Sync + 'static;
}

View file

@ -0,0 +1,10 @@
mod error;
mod framer;
pub use error::*;
pub use framer::*;
#[cfg(feature = "aws")]
pub mod aws_event_stream;
#[cfg(feature = "sse")]
pub mod sse;

View file

@ -0,0 +1,43 @@
use futures_util::{Stream, StreamExt};
use crate::{Error, Framer};
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct SseFrame {
pub event: Option<String>,
pub data: Option<String>,
pub id: Option<String>,
pub retry: Option<u64>,
}
#[derive(Clone, Copy, Debug, Default)]
pub struct SseFramer;
impl Framer for SseFramer {
type Frame = SseFrame;
fn frame<S, B, E>(self, input: S) -> impl Stream<Item = Result<SseFrame, Error>> + Send
where
S: Stream<Item = Result<B, E>> + Send,
B: bytes::Buf + Send,
E: std::error::Error + Send + Sync + 'static,
{
let frames = Box::pin(sse_stream::SseStream::from_bytes_stream(input));
futures_util::stream::try_unfold(frames, |mut frames| async move {
let Some(frame) = frames.next().await else {
return Ok(None);
};
let frame = frame?;
Ok(Some((
SseFrame {
event: frame.event,
data: frame.data,
id: frame.id,
retry: frame.retry,
},
frames,
)))
})
.fuse()
}
}

View file

@ -0,0 +1,92 @@
#![cfg(feature = "aws")]
mod support;
use std::io;
use futures_util::TryStreamExt;
use litellm_framing::aws_event_stream::{AwsEventStreamFrame, AwsEventStreamFramer};
use litellm_framing::{Error, Framer};
use rstest::{fixture, rstest};
use support::encode;
async fn collect_aws(bytes: &[u8], chunk_size: usize) -> Result<Vec<AwsEventStreamFrame>, Error> {
AwsEventStreamFramer
.frame(futures_util::stream::iter(
bytes.chunks(chunk_size).map(Ok::<_, io::Error>),
))
.try_collect()
.await
}
#[fixture]
fn two_frames() -> Vec<u8> {
[encode(b"\xff\x00"), encode(b"second")].concat()
}
#[fixture]
fn payload_frame() -> Vec<u8> {
encode(b"payload")
}
#[rstest]
#[case(1)]
#[case(3)]
#[case(12)]
#[case(usize::MAX)]
#[tokio::test]
async fn fragmented_and_coalesced_frames_preserve_typed_headers_and_binary_payloads(
two_frames: Vec<u8>,
#[case] chunk_size: usize,
) {
let chunk_size = chunk_size.min(two_frames.len());
let frames = collect_aws(&two_frames, chunk_size).await.unwrap();
assert_eq!(frames.len(), 2);
assert_eq!(frames[0].payload, &b"\xff\x00"[..]);
assert_eq!(frames[1].payload, "second");
assert_eq!(
frames[0].headers[0].value().as_string().unwrap().as_str(),
"payload"
);
assert_eq!(frames[0].headers[1].value().as_int32(), Ok(7));
}
#[rstest]
#[case(8)]
#[case(0)]
#[tokio::test]
async fn rejects_corrupt_crcs(payload_frame: Vec<u8>, #[case] index: usize) {
let corrupt_index = if index == 0 {
payload_frame.len() - 1
} else {
index
};
let mut corrupt = payload_frame;
corrupt[corrupt_index] ^= 1;
assert!(matches!(collect_aws(&corrupt, 3).await, Err(Error::Aws(_))));
}
#[rstest]
#[case(0_u32)]
#[case(15)]
#[case(u32::MAX)]
#[tokio::test]
async fn rejects_invalid_lengths(#[case] length: u32) {
assert!(matches!(
collect_aws(&length.to_be_bytes(), 1).await,
Err(Error::InvalidLength(_))
));
}
#[rstest]
#[case(1)]
#[case(3)]
#[case(5)]
#[tokio::test]
async fn rejects_truncation(payload_frame: Vec<u8>, #[case] end: usize) {
assert!(matches!(
collect_aws(&payload_frame[..end], 1).await,
Err(Error::Truncated)
));
}

View file

@ -0,0 +1,29 @@
#![cfg(all(feature = "aws", feature = "sse"))]
mod support;
use std::io;
use futures_util::TryStreamExt;
use litellm_framing::Framer;
use litellm_framing::aws_event_stream::{AwsEventStreamFrame, AwsEventStreamFramer};
use litellm_framing::sse::SseFramer;
use support::encode;
#[tokio::test]
async fn hosting_payloads_feed_the_same_sse_framer_across_envelope_boundaries() {
let bytes = [encode(b"event: delta\ndata: hel"), encode(b"lo\nid: 7\n\n")].concat();
let envelopes = AwsEventStreamFramer.frame(futures_util::stream::iter(
bytes.chunks(3).map(Ok::<_, io::Error>),
));
let frames = SseFramer
.frame(envelopes.map_ok(|frame: AwsEventStreamFrame| frame.payload))
.try_collect::<Vec<_>>()
.await
.unwrap();
assert_eq!(frames.len(), 1);
assert_eq!(frames[0].event.as_deref(), Some("delta"));
assert_eq!(frames[0].data.as_deref(), Some("hello"));
assert_eq!(frames[0].id.as_deref(), Some("7"));
}

View file

@ -0,0 +1,67 @@
#![cfg(feature = "sse")]
use std::io;
use futures_util::{StreamExt, TryStreamExt};
use litellm_framing::sse::{SseFrame, SseFramer};
use litellm_framing::{Error, Framer};
use rstest::rstest;
async fn collect_sse(chunks: &[&[u8]]) -> Result<Vec<SseFrame>, Error> {
SseFramer
.frame(futures_util::stream::iter(
chunks.iter().copied().map(Ok::<_, io::Error>),
))
.try_collect()
.await
}
#[rstest]
#[case(
&[&b":ping\r\nevent: delta\r\nid: 7\r\nretry: 10\r\ndata: \xe2"[..], &b"\x82"[..], &b"\xac\r"[..], &b"\ndata: next\r\n\r"[..], &b"\ndata: [DONE]\n\n"[..]],
vec![
SseFrame {
event: Some("delta".into()),
data: Some("\nnext".into()),
id: Some("7".into()),
retry: Some(10),
},
SseFrame {
event: None,
data: Some("[DONE]".into()),
id: None,
retry: None,
},
]
)]
#[tokio::test]
async fn fragmented_utf8_crlf_and_multiline_data_retain_metadata_and_sentinel(
#[case] chunks: &[&[u8]],
#[case] expected: Vec<SseFrame>,
) {
assert_eq!(collect_sse(chunks).await.unwrap(), expected);
}
#[tokio::test]
async fn eof_does_not_dispatch_an_unterminated_frame() {
assert!(collect_sse(&[b"data: partial\n"]).await.unwrap().is_empty());
}
#[rstest]
#[case(io::ErrorKind::ConnectionReset)]
#[case(io::ErrorKind::UnexpectedEof)]
#[tokio::test]
async fn framing_errors_terminate_and_preserve_input_error_causes(#[case] kind: io::ErrorKind) {
let mut frames = Box::pin(SseFramer.frame(futures_util::stream::iter([
Err(io::Error::new(kind, "reset")),
Ok(&b"data: later\n\n"[..]),
])));
let error = frames.next().await.unwrap().unwrap_err();
assert!(matches!(
error,
Error::Sse(sse_stream::Error::Body(ref cause))
if cause.downcast_ref::<io::Error>().unwrap().kind() == kind
));
assert!(frames.next().await.is_none());
assert!(frames.next().await.is_none());
}

View file

@ -0,0 +1,15 @@
use aws_smithy_eventstream::frame::write_message_to;
use aws_smithy_types::event_stream::{Header, HeaderValue, Message};
use bytes::Bytes;
pub fn encode(payload: &'static [u8]) -> Vec<u8> {
let message = Message::new(Bytes::from_static(payload))
.add_header(Header::new(
":event-type",
HeaderValue::String("payload".into()),
))
.add_header(Header::new("sequence", HeaderValue::Int32(7)));
let mut bytes = Vec::new();
write_message_to(&message, &mut bytes).unwrap();
bytes
}

View file

@ -46,10 +46,7 @@ pub(crate) fn core_error_to_pyerr(error: Error) -> PyErr {
),
Error::Messages(error) => match error {
messages::Error::Auth(source) => auth_is_value_error(source),
messages::Error::InvalidProvider(_)
| messages::Error::InvalidRequest(_)
| messages::Error::Headers(_) => true,
_ => false,
_ => error.is_request(),
},
Error::AudioTranscription(error) => match error {
audio_transcription::Error::Auth(source) => auth_is_value_error(source),

View file

@ -244,6 +244,7 @@ telemetry = True
max_tokens: int = DEFAULT_MAX_TOKENS # OpenAI Defaults
drop_params = drop_params_env_flag(os.environ, verbose_logger)
modify_params = bool(os.getenv("LITELLM_MODIFY_PARAMS", False))
bedrock_neutralize_orphaned_tool_blocks: bool = True
use_chat_completions_url_for_anthropic_messages: bool = bool(
os.getenv("LITELLM_USE_CHAT_COMPLETIONS_URL_FOR_ANTHROPIC_MESSAGES", False)
) # When True, routes OpenAI /v1/messages requests to chat/completions instead of the Responses API
@ -1833,6 +1834,9 @@ if TYPE_CHECKING:
from .llms.azure.responses.o_series_transformation import (
AzureOpenAIOSeriesResponsesAPIConfig as AzureOpenAIOSeriesResponsesAPIConfig,
)
from .llms.azure_ai.responses.transformation import (
AzureAIResponsesAPIConfig as AzureAIResponsesAPIConfig,
)
from .llms.xai.responses.transformation import (
XAIResponsesAPIConfig as XAIResponsesAPIConfig,
)

View file

@ -234,6 +234,7 @@ LLM_CONFIG_NAMES: Final = (
"OpenAIResponsesAPIConfig",
"AzureOpenAIResponsesAPIConfig",
"AzureOpenAIOSeriesResponsesAPIConfig",
"AzureAIResponsesAPIConfig",
"XAIResponsesAPIConfig",
"LiteLLMProxyResponsesAPIConfig",
"HostedVLLMResponsesAPIConfig",
@ -946,6 +947,10 @@ _LLM_CONFIGS_IMPORT_MAP: Final = {
".llms.azure.responses.o_series_transformation",
"AzureOpenAIOSeriesResponsesAPIConfig",
),
"AzureAIResponsesAPIConfig": (
".llms.azure_ai.responses.transformation",
"AzureAIResponsesAPIConfig",
),
"XAIResponsesAPIConfig": (
".llms.xai.responses.transformation",
"XAIResponsesAPIConfig",

View file

@ -3,6 +3,7 @@ import contextvars
import functools
import logging
import os
import re
import sys
from datetime import datetime
from logging import Formatter
@ -13,10 +14,11 @@ import litellm
from litellm.constants import (
LITELLM_TRUNCATED_PAYLOAD_FIELD,
LITELLM_TRUNCATION_STDOUT_SAFEGUARD_NOTE,
MAX_BASE64_LENGTH_STDOUT_LOG,
MAX_STRING_LENGTH_STDOUT_LOG,
)
from litellm.litellm_core_utils.env_utils import get_env_int
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
from litellm.litellm_core_utils.safe_json_dumps import UNSERIALIZABLE_OBJECT, safe_dumps, safe_json_structure
from litellm.litellm_core_utils.safe_json_loads import safe_json_loads
from litellm.litellm_core_utils.secret_redaction import (
redact_internal_details,
@ -77,6 +79,37 @@ def _redact_structured_value(key: str | None, value: str) -> str:
return redact_structured_value(key, value)
_REDACTED_RECORD_ATTR: Final = "litellm_redacted"
_REDACTED_STAMP: Final = object()
_UNREDACTED_SCALAR_TYPES: Final = (bool, int, float, type(None))
def _is_redacted(record: logging.LogRecord) -> bool:
return getattr(record, _REDACTED_RECORD_ATTR, None) is _REDACTED_STAMP
def _scrubbing_changed_nothing(scrubbed: object, original: object) -> bool:
try:
return bool(scrubbed == original)
except Exception:
return False
def _plain_text(value: object) -> str:
try:
return str(value)
except Exception:
return UNSERIALIZABLE_OBJECT
def _redact_extra_value(key: str, value: object) -> object:
try:
scrubbed: Final = safe_json_structure(value, value_transform=_redact_structured_value, key=key)
except Exception:
return _redact_string(_plain_text(value))
return value if _scrubbing_changed_nothing(scrubbed, value) else scrubbed
def redact_secrets(value: str) -> str:
"""Public API: redact known secret/credential patterns from an arbitrary string.
@ -126,7 +159,7 @@ class SecretRedactionFilter(logging.Filter):
_formatter = logging.Formatter()
def filter(self, record: logging.LogRecord) -> bool:
if not _ENABLE_SECRET_REDACTION:
if not _ENABLE_SECRET_REDACTION or _is_redacted(record):
return True
# Runs before args are cleared, and before the extra-field loop below
@ -149,11 +182,19 @@ class SecretRedactionFilter(logging.Filter):
except Exception:
pass
if isinstance(record.stack_info, str):
record.stack_info = _redact_string(record.stack_info) # rebind-ok: a Filter scrubs records in place
# Redact extra fields passed via logger.debug("msg", extra={...})
for key, value in list(record.__dict__.items()):
if key not in _STANDARD_RECORD_ATTRS and isinstance(value, str):
setattr(record, key, _redact_string(value))
if key in _STANDARD_RECORD_ATTRS:
continue
if isinstance(value, str):
setattr(record, key, _redact_structured_value(key, value))
elif not isinstance(value, _UNREDACTED_SCALAR_TYPES):
setattr(record, key, _redact_extra_value(key, value))
setattr(record, _REDACTED_RECORD_ATTR, _REDACTED_STAMP)
return True
@ -277,6 +318,51 @@ def _truncate_for_stdout_log(text: str, limit: int) -> str:
return f"{text[:head_chars]}{_stdout_truncation_marker(len(text) - kept_chars)}{text[-tail_chars:]}"
_BYTES_PER_KIB: Final = 1024
_BYTES_PER_MIB: Final = 1024 * 1024
def format_base64_size(num_chars: int) -> str:
"""Return a human-readable byte-size estimate from a base64 character count."""
num_bytes: Final = num_chars * 3 / 4
if num_bytes >= _BYTES_PER_MIB:
return f"{num_bytes / _BYTES_PER_MIB:.2f}MB"
if num_bytes >= _BYTES_PER_KIB:
return f"{num_bytes / _BYTES_PER_KIB:.1f}KB"
return f"{int(num_bytes)}B"
def _get_max_base64_length_stdout_log() -> int:
return get_env_int("MAX_BASE64_LENGTH_STDOUT_LOG", MAX_BASE64_LENGTH_STDOUT_LOG)
@functools.lru_cache(maxsize=8)
def _base64_run_pattern(min_chars: int) -> "re.Pattern[str]":
return re.compile(rf"(?<![A-Za-z0-9+/])[A-Za-z0-9+/]{{{min_chars},}}={{0,2}}")
_LOWER_HEX_DIGITS: Final = "0123456789abcdef"
_UPPER_HEX_DIGITS: Final = "0123456789ABCDEF"
def _looks_like_base64(run: str) -> bool:
unpadded: Final = run.rstrip("=")
is_hex_or_decimal: Final = not unpadded.strip(_LOWER_HEX_DIGITS) or not unpadded.strip(_UPPER_HEX_DIGITS)
is_one_repeated_char: Final = not unpadded.strip(unpadded[0])
return not is_hex_or_decimal or is_one_repeated_char
def _replace_base64_run(match: "re.Match[str]") -> str:
run: Final = match.group(0)
if not _looks_like_base64(run):
return run
return f"[base64_data truncated: {format_base64_size(len(run))}]"
def _collapse_base64_runs(text: str, limit: int) -> str:
return _base64_run_pattern(limit + 1).sub(_replace_base64_run, text)
class StdoutLogTruncationFilter(logging.Filter):
"""Bounds how much of an oversized log line reaches stdout.
@ -284,36 +370,42 @@ class StdoutLogTruncationFilter(logging.Filter):
request writes hundreds of KB to stdout, repeatedly as the exception propagates from
the router to the proxy handler and into its traceback, all inline on the event loop.
DEBUG records pass through untouched, since dumping full payloads is the point of
At every level, in the message and in the traceback alike, a base64 run longer than
MAX_BASE64_LENGTH_STDOUT_LOG collapses to a size placeholder first: a multi-megabyte
document upload otherwise costs seconds of event-loop time per DEBUG line in the
secret regex alone. Hex and decimal runs (digests, numeric ids) are left alone unless
they are one repeated character, which is what a zero-filled payload encodes to.
The text around a run stays, since dumping payloads is the point of
`--detailed_debug`, and logging callbacks (OTEL, Datadog, etc.) don't run through
logging filters at all, so they still get the untruncated error.
logging filters at all, so they still get the untouched record.
"""
_formatter = logging.Formatter()
def filter(self, record: logging.LogRecord) -> bool:
if record.levelno < logging.INFO:
return True
limit: Final = _get_max_string_length_stdout_log()
if limit <= 0:
return True
try:
message: Final = record.getMessage()
except (TypeError, ValueError):
return True
if len(message) > limit:
record.msg = _truncate_for_stdout_log(message, limit) # rebind-ok: the Filter interface mutates the record
record.args = None # rebind-ok: args are consumed by the truncated message above
base64_limit: Final = _get_max_base64_length_stdout_log()
collapsed: Final = _collapse_base64_runs(message, base64_limit) if base64_limit > 0 else message
limit: Final = _get_max_string_length_stdout_log() if record.levelno >= logging.INFO else 0
bounded: Final = _truncate_for_stdout_log(collapsed, limit) if 0 < limit < len(collapsed) else collapsed
if bounded != message:
record.msg = bounded # rebind-ok: the Filter interface mutates the record
record.args = None # rebind-ok: args are consumed by the rewritten message above
if isinstance(record.exc_info, tuple):
exc_text: Final = record.exc_text or self._formatter.formatException(record.exc_info)
if len(exc_text) > limit:
record.exc_text = _truncate_for_stdout_log( # rebind-ok: the Filter interface mutates the record
exc_text, limit
)
if not isinstance(record.exc_info, tuple):
return True
exc_text: Final = record.exc_text or self._formatter.formatException(record.exc_info)
collapsed_exc: Final = _collapse_base64_runs(exc_text, base64_limit) if base64_limit > 0 else exc_text
bounded_exc: Final = (
_truncate_for_stdout_log(collapsed_exc, limit) if 0 < limit < len(collapsed_exc) else collapsed_exc
)
if bounded_exc != exc_text:
record.exc_text = bounded_exc # rebind-ok: the Filter interface mutates the record
return True
@ -474,6 +566,7 @@ def _get_standard_record_attrs() -> frozenset:
_STANDARD_RECORD_ATTRS: Final = _get_standard_record_attrs()
_NON_EXTRA_RECORD_ATTRS: Final = _STANDARD_RECORD_ATTRS | {_REDACTED_RECORD_ATTR}
# CorrelationContextFilter is the only legitimate source for these two JSON fields;
# see JsonFormatter.format() for why they're excluded from the generic message-content
@ -514,7 +607,7 @@ class JsonFormatter(Formatter):
# Include extra attributes passed via logger.debug("msg", extra={...})
for key, value in record.__dict__.items():
if key not in _STANDARD_RECORD_ATTRS and key not in json_record:
if key not in _NON_EXTRA_RECORD_ATTRS and key not in json_record:
json_record[key] = value
# trace_id/session_id are reserved: CorrelationContextFilter is the only
@ -538,7 +631,7 @@ class JsonFormatter(Formatter):
if record.exc_info:
json_record["stacktrace"] = record.exc_text or self.formatException(record.exc_info)
return safe_dumps(json_record, value_transform=_redact_structured_value)
return safe_dumps(json_record, value_transform=None if _is_redacted(record) else _redact_structured_value)
class CorrelationPlainFormatter(logging.Formatter):
@ -549,7 +642,8 @@ class CorrelationPlainFormatter(logging.Formatter):
"""
def format(self, record: logging.LogRecord) -> str:
formatted: Final = _redact_string(super().format(record))
rendered: Final = super().format(record)
formatted: Final = rendered if _is_redacted(record) else _redact_string(rendered)
trace_id: Final = getattr(record, "trace_id", None)
session_id: Final = getattr(record, "session_id", None)
if not trace_id and not session_id:
@ -567,8 +661,8 @@ def _setup_json_exception_handlers(formatter):
# Create a handler with JSON formatting for exceptions
error_handler: Final = logging.StreamHandler()
error_handler.setFormatter(formatter)
error_handler.addFilter(_secret_filter)
error_handler.addFilter(_stdout_truncation_filter)
error_handler.addFilter(_secret_filter)
error_handler.addFilter(_correlation_filter)
# Setup excepthook for uncaught exceptions

View file

@ -502,7 +502,12 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
elif key == "response_format":
text_format = self._transform_response_format_to_text_format(value)
if text_format:
responses_api_request["text"] = text_format
responses_api_request["text"] = self._merge_text(responses_api_request, text_format)
elif key == "verbosity":
responses_api_request["text"] = self._merge_text(
responses_api_request,
MappingProxyType({"verbosity": value}), # pyright: ignore[reportUnknownArgumentType] # untyped value
)
elif key == "tool_choice":
responses_api_request["tool_choice"] = self._normalize_tool_choice_for_responses_api(value)
elif key == "stream_options":
@ -518,6 +523,19 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
elif key == "web_search_options":
self._add_web_search_tool(responses_api_request, value)
@staticmethod
def _merge_text(
responses_api_request: "ResponsesAPIOptionalRequestParams", update: Mapping[str, object]
) -> "ResponseText":
existing: Final = cast( # cast-ok: text field is a ResponseText | dict[str, Any] | None union
"dict[str, object]",
dict(responses_api_request).get("text") or {}, # mutable-ok: one-shot merge seed
)
return cast( # cast-ok: merged mapping is a valid ResponseText shape
"ResponseText",
{**existing, **update}, # mutable-ok: one-shot merged payload
)
def _build_sanitized_litellm_params(self, litellm_params: dict) -> dict[str, object]:
"""Build sanitized litellm_params with merged metadata."""
responses_optional_param_keys: Final = set(ResponsesAPIOptionalRequestParams.__annotations__.keys())

View file

@ -102,6 +102,7 @@ REDACTED_BY_LITELLM: Final = "redacted-by-litellm"
REDACTED_TOOL_CALL_ARGUMENTS_PLACEHOLDER: Final = "{}"
MAX_STRING_LENGTH_STDOUT_LOG: Final = get_env_int("MAX_STRING_LENGTH_STDOUT_LOG", 4096)
MAX_BASE64_LENGTH_STDOUT_LOG: Final = get_env_int("MAX_BASE64_LENGTH_STDOUT_LOG", 4096)
# When true, adds detailed per-phase timing breakdown headers to responses.
# Headers: x-litellm-timing-{pre-processing,llm-api,post-processing,message-copy}-ms

View file

@ -346,13 +346,17 @@ class MCPClient:
self.update_auth_value(auth_value)
async def discovery_auth_fingerprint(self) -> str:
return self._hash_discovery_auth(await self.prepare_request_auth())
async def prepare_request_auth(self) -> httpx.Request:
"""Preview the authenticated request without sending it, closing the auth flow afterwards."""
request: Final = httpx.Request("POST", self.server_url or "http://localhost/", headers=self._get_auth_headers())
if self._resolved_auth is None:
return self._hash_discovery_auth(request)
return request
flow: Final = self._resolved_auth.async_auth_flow(request)
try:
authenticated: Final = await flow.__anext__()
return self._hash_discovery_auth(authenticated)
return authenticated
finally:
await flow.aclose()

View file

@ -16,6 +16,7 @@ from litellm.litellm_core_utils.core_helpers import (
get_or_create_metadata_bucket,
redact_nested_match_and_regex_keys,
)
from litellm.llms.base_llm.guardrail_translation.base_translation import REQUEST_SCAN_CONTEXT_KEY
from litellm.secret_managers.main import str_to_bool
from litellm.types.guardrails import (
DynamicGuardrailParams,
@ -949,9 +950,28 @@ class CustomGuardrail(CustomLogger):
await translation.process_input_messages(data=scratch_request, guardrail_to_apply=self)
if response is None:
return
await output_translation.process_output_response(
response=copy.deepcopy(response), guardrail_to_apply=self, request_data=scratch_request
output_request: Final = (
scratch_request
if type(output_translation) is type(translation)
else self._chat_shaped_request(scratch_request, translation)
)
await output_translation.process_output_response(
response=copy.deepcopy(response), guardrail_to_apply=self, request_data=output_request
)
def _chat_shaped_request(
self,
scratch_request: Mapping[str, object],
translation: "BaseTranslation",
) -> dict[str, object]: # mutable-ok: BaseTranslation.process_output_response contract
"""The logged request in OpenAI chat shape, for an output scan whose translation differs from the input's."""
context: Final = translation.request_scan_context(scratch_request, self)
return {
**scratch_request,
"messages": list(context.structured_messages),
"tools": list(context.tools),
REQUEST_SCAN_CONTEXT_KEY: context,
}
def supports_scan_only_tool_results(self) -> bool:
"""Whether this guardrail can scan tool-result content.
@ -1379,8 +1399,9 @@ class CustomGuardrail(CustomLogger):
raise e
def _inputs_were_modified(self, original_inputs: Mapping[str, object], response: Mapping[str, object]) -> bool:
"""True when any key of either mapping differs between them (mask), False otherwise (allow)."""
return any(original_inputs.get(key) != response.get(key) for key in original_inputs.keys() | response.keys())
"""True when any content key of either mapping differs between them (mask), False otherwise (allow)."""
compared_keys: Final = (original_inputs.keys() | response.keys()) - _STREAM_CONTROL_KEYS
return any(original_inputs.get(key) != response.get(key) for key in compared_keys)
def mask_content_in_string(
self,
@ -1490,6 +1511,7 @@ def _sync_guardrail_info_to_logging_obj(request_data: dict, logging_obj: object)
_PRE_CALL_CONTENT_KEYS: Final = frozenset(
{"messages", "input", "prompt", "system", "instructions", "tools", "functions", "function_call", "tool_choice"}
)
_STREAM_CONTROL_KEYS: Final = frozenset({"stream_holdback_chars"})
def _original_inputs_for(

View file

@ -2965,16 +2965,13 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
)
propagator: Final = TraceContextTextMapPropagator()
carrier: Final = {"traceparent": _traceparent}
carrier: Final = {key: headers[key] for key in ("traceparent", "tracestate") if headers.get(key) is not None}
_parent_context: Final = propagator.extract(carrier=carrier)
return _parent_context
def _get_span_context(self, kwargs, default_span: Span | None = None):
from opentelemetry import context, trace
from opentelemetry.trace.propagation.tracecontext import (
TraceContextTextMapPropagator,
)
litellm_params: Final = kwargs.get("litellm_params", {}) or {}
proxy_server_request: Final = litellm_params.get("proxy_server_request", {}) or {}
@ -2998,11 +2995,7 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
# Priority 2: HTTP traceparent header
if traceparent is not None:
verbose_logger.debug("OpenTelemetry: Using traceparent header for context propagation")
carrier: Final = {"traceparent": traceparent}
return (
TraceContextTextMapPropagator().extract(carrier=carrier),
None,
)
return self.get_traceparent_from_header(headers=headers), None
# Priority 3: Active span from global context (auto-detection)
try:

View file

@ -1,15 +1,19 @@
"""The span engine: dedup, start, run the mapper chain, set status, end."""
from collections import OrderedDict
from collections.abc import Callable, Sequence
from collections.abc import Callable, Mapping, Sequence
from types import MappingProxyType
from typing import Final
from opentelemetry.context import Context
from opentelemetry.sdk.trace import ReadableSpan, SpanLimits
from opentelemetry.sdk.trace import Span as SdkSpan
from opentelemetry.trace import Link, Span, Tracer
from opentelemetry.trace.status import Status, StatusCode
from litellm.integrations.otel.mappers import resolve_mappers
from litellm.integrations.otel.mappers.base import AttributeMapper, SpanData
from litellm.integrations.otel.mappers.base import AttributeMapper, AttrValue, SpanData
from litellm.integrations.otel.mappers.openinference import fit_indexed_messages
from litellm.integrations.otel.model.config import OpenTelemetryV2Config
from litellm.integrations.otel.model.payloads import (
GuardrailSpanData,
@ -52,25 +56,48 @@ _NAME_BUILDERS: Final[dict[SpanRole, Callable[..., str]]] = {
_DEDUP_CACHE_MAX: Final = 10_000
def _stamp_otel_error_attributes(span: Span, error_type: str, resolved_message: str) -> None:
"""Stamp the OTel-semconv error attributes (``error.type`` + ``error.message``).
``error_type`` and ``resolved_message`` are ``finish_span``'s already-computed
fallback chains, so the pair on the status, event, and attributes stays in
lockstep."""
span.set_attribute(Error.TYPE, error_type)
span.set_attribute(Error.MESSAGE, resolved_message)
def _resolve_error(error: SpanError) -> tuple[str, str] | None:
"""The ``(error_type, message)`` fallback chain shared by the status, the event and the attributes, or
``None`` when ``error`` carries neither a type nor a message."""
if not (error.error_type or error.message):
return None
return error.error_type or "error", error.message or error.error_type or "error"
def _stamp_litellm_error_attributes(span: Span, error: SpanError) -> None:
"""Stamp litellm-specific error detail attributes. Emitted only when the
corresponding field is populated so guardrail-shape errors carrying only a
message aren't polluted with empty detail keys."""
if error.code:
span.set_attribute(LiteLLMError.CODE, error.code)
if error.stack_trace:
span.set_attribute(LiteLLMError.STACK_TRACE, error.stack_trace)
if error.llm_provider:
span.set_attribute(LiteLLMError.LLM_PROVIDER, error.llm_provider)
_NO_ATTRIBUTES: Final[Mapping[str, AttrValue]] = MappingProxyType({})
def error_attributes(error: SpanError) -> Mapping[str, AttrValue]:
"""The v2 error attribute set: the OTel-semconv ``error.*`` pair plus the litellm detail keys that are
populated, so guardrail-shape errors carrying only a message aren't polluted with empty detail keys."""
resolved: Final = _resolve_error(error)
if resolved is None:
return _NO_ATTRIBUTES
error_type, message = resolved
pairs: Final = (
(Error.TYPE, error_type),
(Error.MESSAGE, message),
(LiteLLMError.CODE, error.code),
(LiteLLMError.STACK_TRACE, error.stack_trace),
(LiteLLMError.LLM_PROVIDER, error.llm_provider),
)
return MappingProxyType({key: value for key, value in pairs if value})
def span_attribute_limit(span: Span) -> int | None:
"""The attribute count limit ``span`` was built with, ``None`` when unbounded."""
if not isinstance(span, SdkSpan):
return SpanLimits().max_span_attributes
return span._limits.max_span_attributes # pyright: ignore[reportPrivateUsage] # SDK has no public getter
def attribute_budget(span: Span, reserved: int) -> int | None:
"""How many mapped attributes fit on ``span`` next to what it already carries and ``reserved`` more."""
limit: Final = span_attribute_limit(span)
if limit is None:
return None
on_span: Final = len(span.attributes or ()) if isinstance(span, ReadableSpan) else 0
return limit - on_span - reserved
def stamp_error(
@ -93,12 +120,12 @@ def stamp_error(
``set_status`` are opt-outs for callers whose span lifecycle (``use_span``) or
owner (the FastAPI instrumentor) already records the event or the status.
"""
if not (error.error_type or error.message):
resolved: Final = _resolve_error(error)
if resolved is None:
return None
error_type: Final = error.error_type or "error"
message: Final = error.message or error.error_type or "error"
_stamp_otel_error_attributes(span, error_type, message)
_stamp_litellm_error_attributes(span, error)
error_type, message = resolved
for key, value in error_attributes(error).items():
span.set_attribute(key, value)
if set_status:
span.set_status(Status(StatusCode.ERROR, message))
if record_event:
@ -238,9 +265,6 @@ class SpanEmitter:
data, since the boundary opener only has a provisional name.
"""
span.update_name(_NAME_BUILDERS[role](data))
for mapper in self._mappers:
for key, value in mapper.map(data).items():
span.set_attribute(key, value)
error: Final = (
data.error
if isinstance(
@ -255,6 +279,13 @@ class SpanEmitter:
)
else None
)
mapped: Final = MappingProxyType(
{key: value for mapper in self._mappers for key, value in mapper.map(data).items()}
)
stamped_later: Final = error_attributes(error) if error else _NO_ATTRIBUTES
reserved: Final = len(stamped_later.keys() - mapped.keys())
for key, value in fit_indexed_messages(mapped, attribute_budget(span, reserved)).items():
span.set_attribute(key, value)
if error:
stamped: Final = stamp_error(span, error)
if stamped is not None and self._event_recorder is not None and role is SpanRole.LLM_CALL:

View file

@ -6,10 +6,10 @@ from litellm.integrations.otel.logger import OpenTelemetryV2
from litellm.integrations.otel.mappers.langfuse import (
LANGFUSE_OBSERVATION_INPUT,
LANGFUSE_OBSERVATION_OUTPUT,
LANGFUSE_TRACE_NAME,
LangfuseMapper,
)
from litellm.integrations.otel.model.metadata import caller_trace_name
from litellm.integrations.otel.model.request_io import request_input, response_output, stream_output
from litellm.integrations.otel.model.trace_controls import caller_trace_controls
from litellm.integrations.otel.plumbing.context import request_root_span
if TYPE_CHECKING:
@ -18,14 +18,13 @@ if TYPE_CHECKING:
class LangfuseOpenTelemetryV2(OpenTelemetryV2):
"""Names the trace from the request. Langfuse reads ``langfuse.trace.name`` off the root observation,
and the proxy's root span is still recording when the LLM call starts."""
"""Stamps the caller's trace controls (name, user, session, tags) on the request. Langfuse reads them off
the root observation, and the proxy's root span is still recording when the LLM call starts."""
def log_pre_api_call(self, model: str, messages: object, kwargs: Mapping[str, object]) -> None:
root: Final = request_root_span()
name: Final = caller_trace_name(kwargs)
if root is not None and root.is_recording() and name is not None:
root.set_attribute(LANGFUSE_TRACE_NAME, name)
if root is not None and root.is_recording():
root.set_attributes(LangfuseMapper.trace_attributes(caller_trace_controls(kwargs)))
super().log_pre_api_call(model, messages, kwargs)

View file

@ -555,7 +555,7 @@ class OpenTelemetryV2(CustomLogger):
capture_content=self.config.capture_span_content,
time_to_first_chunk_seconds=call.time_to_first_chunk_seconds,
request_route=request_root_http_route(),
trace_name=call.trace_name,
trace=call.trace,
)
end_time_ns: Final = to_ns(end_time)
if carrier is not None and carrier.span is not None:

View file

@ -6,7 +6,8 @@ Langfuse ingests OTLP spans and reads from its own vendor namespace
Every attribute is declared as a ``key -> extractor`` table entry (one callable
per mapping operation): ``_LLM_CALL_ATTRS`` for scalars and ``_BLOB_ATTRS`` for
the JSON-serialized payloads. ``_llm_call`` just applies both tables.
the JSON-serialized payloads. ``trace_attributes`` maps the caller's trace controls
(shared with the root observation); ``_llm_call`` applies both tables plus it.
"""
import json
@ -16,6 +17,7 @@ from typing import Final
from litellm.integrations.otel.mappers.base import AttributeMap, AttrValue, SpanData
from litellm.integrations.otel.mappers.utils import (
collect,
drop_none_pairs,
json_if,
output_messages,
serialize_messages,
@ -25,10 +27,14 @@ from litellm.integrations.otel.model.payloads import (
LLMRequestParams,
LLMUsage,
)
from litellm.integrations.otel.model.trace_controls import TraceControls
LANGFUSE_OBSERVATION_INPUT: Final = "langfuse.observation.input"
LANGFUSE_OBSERVATION_OUTPUT: Final = "langfuse.observation.output"
LANGFUSE_TRACE_NAME: Final = "langfuse.trace.name"
LANGFUSE_TRACE_USER_ID: Final = "user.id"
LANGFUSE_TRACE_SESSION_ID: Final = "session.id"
LANGFUSE_TRACE_TAGS: Final = "langfuse.trace.tags"
class LangfuseMapper:
@ -37,7 +43,6 @@ class LangfuseMapper:
"langfuse.observation.model.name": lambda d: d.request_model or None,
"langfuse.observation.metadata.provider": lambda d: d.provider or None,
"langfuse.observation.id": lambda d: d.identity.call_id or None,
LANGFUSE_TRACE_NAME: lambda d: d.trace_name or None,
"langfuse.trace.metadata.team_id": lambda d: d.identity.team_id or None,
"langfuse.trace.metadata.team_alias": lambda d: d.identity.team_alias or None,
}
@ -77,9 +82,21 @@ class LangfuseMapper:
case _:
return {}
@staticmethod
def trace_attributes(trace: TraceControls) -> AttributeMap:
return drop_none_pairs(
(
(LANGFUSE_TRACE_NAME, trace.name or None),
(LANGFUSE_TRACE_USER_ID, trace.user_id or None),
(LANGFUSE_TRACE_SESSION_ID, trace.session_id or None),
(LANGFUSE_TRACE_TAGS, trace.tags or None),
)
)
@classmethod
def _llm_call(cls, data: LLMCallSpanData) -> AttributeMap:
return {
**collect(cls._LLM_CALL_ATTRS, data),
**cls.trace_attributes(data.trace),
**collect(cls._BLOB_ATTRS, data),
}

View file

@ -7,12 +7,13 @@ Phoenix + any other OpenInference-aware backend simultaneously.
"""
import json
from collections.abc import Callable, Sequence
from collections.abc import Callable, Mapping, Sequence
from itertools import accumulate, chain, groupby
from types import MappingProxyType
from typing import Final
from litellm.integrations.otel.mappers.base import AttributeMap, AttrValue, SpanData
from litellm.integrations.otel.mappers.utils import (
MAX_MESSAGE_ATTRS_PER_SPAN,
MAX_TOOL_DEFINITION_ATTRS_PER_SPAN,
collect,
drop_none,
@ -27,7 +28,53 @@ from litellm.integrations.otel.model.payloads import (
ToolDefinition,
)
_MAX_INDEXED_MESSAGES: Final = MAX_MESSAGE_ATTRS_PER_SPAN // 2
_INPUT_MESSAGES: Final = "llm.input_messages"
_OUTPUT_MESSAGES: Final = "llm.output_messages"
_MESSAGE_FAMILIES: Final = (_INPUT_MESSAGES, _OUTPUT_MESSAGES)
def _message_key_groups(attrs: Mapping[str, AttrValue]) -> Mapping[tuple[str, int], tuple[str, ...]]:
"""Per-index message keys in ``attrs`` grouped by ``(family, index)``."""
tagged: Final = sorted(
(family, int(key.split(".")[2]), key)
for key in attrs
for family in _MESSAGE_FAMILIES
if key.startswith(f"{family}.")
)
return MappingProxyType(
{group: tuple(key for _, _, key in keys) for group, keys in groupby(tagged, key=lambda tag: tag[:2])}
)
def _shed_order(groups: Mapping[tuple[str, int], tuple[str, ...]]) -> tuple[tuple[str, int], ...]:
"""Message groups least valuable first: middle prompt turns, extra choices, then the opener, the newest turn
and the first choice."""
inputs: Final = sorted(idx for family, idx in groups if family == _INPUT_MESSAGES)
outputs: Final = sorted(idx for family, idx in groups if family == _OUTPUT_MESSAGES)
pinned_inputs: Final = tuple(dict.fromkeys((*inputs[:1], *inputs[-1:])))
return (
*((_INPUT_MESSAGES, idx) for idx in inputs[1:-1]),
*((_OUTPUT_MESSAGES, idx) for idx in reversed(outputs[1:])),
*((_INPUT_MESSAGES, idx) for idx in pinned_inputs),
*((_OUTPUT_MESSAGES, idx) for idx in outputs[:1]),
)
def fit_indexed_messages(attrs: Mapping[str, AttrValue], budget: int | None) -> Mapping[str, AttrValue]:
"""``attrs`` with whole per-index messages shed, least valuable first, until at most ``budget`` keys remain.
``None`` means the span has no attribute count limit. Every message still rides the ``input.value`` and
``output.value`` blobs, so shedding a per-index pair loses no content.
"""
if budget is None or len(attrs) <= budget:
return attrs
groups: Final = _message_key_groups(attrs)
order: Final = _shed_order(groups)
running: Final = tuple(accumulate(len(groups[group]) for group in order))
excess: Final = len(attrs) - budget
shed_count: Final = next((n + 1 for n, total in enumerate(running) if total >= excess), len(order))
shed: Final = frozenset(chain.from_iterable(groups[group] for group in order[:shed_count]))
return MappingProxyType({key: value for key, value in attrs.items() if key not in shed})
class OpenInferenceMapper:
@ -87,42 +134,22 @@ class OpenInferenceMapper:
return {}
def _llm_call(self, data: LLMCallSpanData) -> AttributeMap:
outputs: Final = output_messages(data)
indexed_in, indexed_out = self._indexed_split(len(data.messages_in), len(outputs))
return {
**collect(self._LLM_CALL_ATTRS, data),
**collect(self._BLOB_ATTRS, data),
**self._messages(
"llm.input_messages",
"input.value",
data.messages_in,
self._prompt_positions(len(data.messages_in), indexed_in),
),
**self._messages("llm.output_messages", "output.value", outputs, range(indexed_out)),
**self._messages(_INPUT_MESSAGES, "input.value", data.messages_in),
**self._messages(_OUTPUT_MESSAGES, "output.value", output_messages(data)),
**self._tools(data),
}
@staticmethod
def _indexed_split(inputs: int, outputs: int) -> tuple[int, int]:
"""Prompt and response share one allowance; the response is reserved at least half of it."""
indexed_out: Final = min(outputs, max(_MAX_INDEXED_MESSAGES // 2, _MAX_INDEXED_MESSAGES - inputs))
return _MAX_INDEXED_MESSAGES - indexed_out, indexed_out
@staticmethod
def _prompt_positions(total: int, indexed: int) -> tuple[int, ...]:
"""Prompt messages that get per-index attributes: message 0 and the most recent turns."""
if total <= indexed:
return tuple(range(total))
return (0, *range(total - indexed + 1, total))
@staticmethod
def _messages(prefix: str, value_key: str, messages: Sequence[object], positions: Sequence[int]) -> AttributeMap:
"""``{prefix}.{idx}.message.*`` keys for the messages at ``positions`` + the ``value_key`` blob of all."""
def _messages(prefix: str, value_key: str, messages: Sequence[object]) -> AttributeMap:
"""``{prefix}.{idx}.message.*`` keys for every message + the ``value_key`` blob of all of them."""
parsed: Final = [(m.get("role") if isinstance(m, dict) else None, message_content(m)) for m in messages]
attrs: Final = drop_none(
{
key: value
for idx, (role, content) in ((idx, parsed[idx]) for idx in positions)
for idx, (role, content) in enumerate(parsed)
for key, value in (
(f"{prefix}.{idx}.message.role", role if isinstance(role, str) else None),
(f"{prefix}.{idx}.message.content", content),

View file

@ -6,7 +6,7 @@ they live in one place.
"""
import json
from collections.abc import Callable, Mapping, Sequence
from collections.abc import Callable, Iterable, Mapping, Sequence
from typing import Final
from litellm.integrations.otel.mappers.base import AttributeMap, AttrValue
@ -32,14 +32,6 @@ core telemetry no matter how many vocabularies are configured.
"""
MAX_MESSAGE_ATTRS_PER_SPAN: Final = DEFAULT_SPAN_ATTRIBUTE_LIMIT // 8
"""Span-wide ceiling on per-index chat message attributes, prompt and response together.
An eighth is the largest share that still fits beside the tool ceiling and the core
of every vocabulary at once. The complete conversation still rides the JSON blobs.
"""
def tool_attr_budget(vocabularies: int) -> int:
"""Split the span-wide tool-definition ceiling across active vocabularies."""
return MAX_TOOL_DEFINITION_ATTRS_PER_SPAN // max(vocabularies, 1)
@ -47,7 +39,12 @@ def tool_attr_budget(vocabularies: int) -> int:
def drop_none(values: Mapping[str, AttrValue | None]) -> AttributeMap:
"""Return ``values`` with ``None``-valued entries removed."""
return {k: v for k, v in values.items() if v is not None}
return drop_none_pairs(values.items())
def drop_none_pairs(pairs: Iterable[tuple[str, AttrValue | None]]) -> AttributeMap:
"""Return ``pairs`` as a map with ``None``-valued entries removed."""
return {k: v for k, v in pairs if v is not None}
def tool_definition_attrs(

View file

@ -43,12 +43,12 @@ from typing import TYPE_CHECKING, Any, Final, cast
from litellm.constants import LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL
from litellm.integrations.otel.model.semconv import resolve_operation
from litellm.integrations.otel.model.utils import as_str, to_seconds
from litellm.integrations.otel.model.trace_controls import TraceControls, caller_trace_controls
from litellm.integrations.otel.model.utils import as_str, as_str_mapping, to_seconds
if TYPE_CHECKING:
from litellm.types.utils import StandardLoggingPayload
LANGFUSE_TRACE_NAME_HEADER: Final = "langfuse_trace_name"
REQUESTER_METADATA_KEY: Final = "requester_metadata"
REQUESTER_METADATA_PATH: Final = f"{REQUESTER_METADATA_KEY}."
@ -225,7 +225,7 @@ class LLMCallEvent:
# needs to be reasonable for a span that never gets closed (a leak).
provisional_span_name: str
time_to_first_chunk_seconds: float | None
trace_name: str | None
trace: TraceControls
@classmethod
def from_dict(cls, kwargs: Mapping[str, Any]) -> LLMCallEvent:
@ -242,30 +242,10 @@ class LLMCallEvent:
upstream_started=kwargs.get("api_call_start_time") is not None,
provisional_span_name=f"{operation.value} {model}".strip(),
time_to_first_chunk_seconds=time_to_first_chunk_seconds(kwargs),
trace_name=caller_trace_name(kwargs),
trace=caller_trace_controls(kwargs),
)
def caller_trace_name(kwargs: Mapping[str, object]) -> str | None:
request: Final = _as_str_mapping(kwargs.get("litellm_params"))
if request is None:
return None
proxy_request: Final = _as_str_mapping(request.get("proxy_server_request"))
headers: Final = _as_str_mapping(proxy_request.get("headers")) if proxy_request is not None else None
from_header: Final = as_str(headers.get(LANGFUSE_TRACE_NAME_HEADER)) if headers is not None else None
if from_header:
return from_header
return next(
(
name
for key in ("metadata", "litellm_metadata")
if (metadata := _as_str_mapping(request.get(key))) is not None
and (name := as_str(metadata.get("trace_name")))
),
None,
)
def time_to_first_chunk_seconds(kwargs: Mapping[str, Any]) -> float | None:
"""Seconds from the upstream request being issued (``api_call_start_time``)
to the first streamed chunk (``completion_start_time``); ``None`` for
@ -300,15 +280,8 @@ def auth_metadata(payload: StandardLoggingPayload | None, kwargs: Mapping[str, o
)
def _as_str_mapping(value: object) -> Mapping[str, object] | None:
"""A read-only view of ``value`` when it is a mapping, else ``None``."""
if not isinstance(value, Mapping):
return None
return cast("Mapping[str, object]", value) # cast-ok: isinstance-guarded, JSON metadata has str keys
def _string_entries(value: object) -> Mapping[str, str] | None:
entries: Final = _as_str_mapping(value)
entries: Final = as_str_mapping(value)
if entries is None:
return None
typed: Final = MappingProxyType({key: item for key, item in entries.items() if isinstance(item, str)})
@ -324,18 +297,18 @@ def _metadata_dicts(
litellm copies it onto ``metadata``, but both are yielded so a route that
populates only one is still covered.
"""
payload_view: Final = _as_str_mapping(payload)
payload_view: Final = as_str_mapping(payload)
if payload_view is not None:
payload_metadata: Final = _as_str_mapping(payload_view.get("metadata"))
payload_metadata: Final = as_str_mapping(payload_view.get("metadata"))
if payload_metadata is not None:
yield payload_metadata
params: Final = _as_str_mapping(kwargs.get("litellm_params"))
params: Final = as_str_mapping(kwargs.get("litellm_params"))
if params is None:
return
yield from (
metadata
for key in ("metadata", "litellm_metadata")
if (metadata := _as_str_mapping(params.get(key))) is not None
if (metadata := as_str_mapping(params.get(key))) is not None
)
@ -365,14 +338,14 @@ def metadata_from_request_data(data: object) -> Mapping[str, object] | None:
The proxy stores it under ``metadata`` or ``litellm_metadata`` depending on the route;
the proxy-owned siblings (``user_api_key_*``, ``requester_ip_address``) are not read.
"""
top: Final = _as_str_mapping(data)
top: Final = as_str_mapping(data)
if top is None:
return None
snapshots: Final = tuple(
snapshot
for name in ("metadata", "litellm_metadata")
if (nested := _as_str_mapping(top.get(name))) is not None
and (snapshot := _as_str_mapping(nested.get(REQUESTER_METADATA_KEY))) is not None
if (nested := as_str_mapping(top.get(name))) is not None
and (snapshot := as_str_mapping(nested.get(REQUESTER_METADATA_KEY))) is not None
)
return MappingProxyType({REQUESTER_METADATA_KEY: snapshots[0]}) if snapshots else None
@ -382,7 +355,7 @@ def flatten_metadata(raw: Mapping[str, object]) -> Iterator[tuple[str, str]]:
stack: Final = list(tuple(raw.items())[::-1]) # mutable-ok: iterative worklist keeps the walk off the call stack
while stack:
key, value = stack.pop()
if (nested := _as_str_mapping(value)) is not None:
if (nested := as_str_mapping(value)) is not None:
stack.extend(tuple((f"{key}.{sub_key}", sub_value) for sub_key, sub_value in nested.items())[::-1])
elif isinstance(value, (str, bool, int, float)):
yield key, str(value)

View file

@ -10,10 +10,7 @@ from types import MappingProxyType
from typing import TYPE_CHECKING, ClassVar, Final, cast
from urllib.parse import urlsplit
from litellm.integrations.otel.model.metadata import (
RequestContext,
RequestIdentity,
)
from litellm.integrations.otel.model.metadata import RequestContext, RequestIdentity
from litellm.integrations.otel.model.semconv import (
GenAIOperation,
GenAIOutputType,
@ -22,6 +19,7 @@ from litellm.integrations.otel.model.semconv import (
resolve_output_type,
resolve_provider,
)
from litellm.integrations.otel.model.trace_controls import TraceControls
from litellm.integrations.otel.model.utils import (
as_bool,
as_float,
@ -387,7 +385,7 @@ class LLMCallSpanData:
output_type: GenAIOutputType | None = None
call_type: str | None = None
request_route: str | None = None
trace_name: str | None = None
trace: TraceControls = field(default_factory=TraceControls)
@classmethod
def from_standard_logging_payload(
@ -396,7 +394,7 @@ class LLMCallSpanData:
capture_content: bool = False,
time_to_first_chunk_seconds: float | None = None,
request_route: str | None = None,
trace_name: str | None = None,
trace: TraceControls | None = None,
) -> LLMCallSpanData:
params: Final = cast(Mapping[str, object], payload.get("model_parameters") or {})
# The single parse of the request's metadata — the request-vs-provider
@ -438,7 +436,7 @@ class LLMCallSpanData:
output_type=resolve_output_type(call_type),
call_type=call_type or None,
request_route=request_route or context.identity.request_route,
trace_name=trace_name,
trace=trace or TraceControls(),
)

View file

@ -0,0 +1,61 @@
"""The caller's Langfuse trace controls, parsed from the live callback kwargs."""
from __future__ import annotations
from collections.abc import Mapping
from dataclasses import dataclass
from typing import Final
from pydantic import TypeAdapter, ValidationError
from litellm.integrations.otel.model.utils import as_str, as_str_mapping
LANGFUSE_HEADER_PREFIX: Final = "langfuse_"
_ITEMS: Final = TypeAdapter(tuple[object, ...])
@dataclass(frozen=True, slots=True)
class TraceControls:
"""The caller's trace-level Langfuse controls: ``metadata.trace_name`` / ``trace_user_id`` / ``session_id`` /
``tags`` on the request (SDK or proxy body), with the proxy's ``langfuse_<control>`` headers winning over the
body for the scalar ones. Mutation controls (``trace_id``, ``existing_trace_id``, ``update_trace_keys``) are
deliberately not carried."""
name: str | None = None
user_id: str | None = None
session_id: str | None = None
tags: tuple[str, ...] = ()
def caller_trace_controls(kwargs: Mapping[str, object]) -> TraceControls:
request: Final = as_str_mapping(kwargs.get("litellm_params"))
if request is None:
return TraceControls()
proxy_request: Final = as_str_mapping(request.get("proxy_server_request"))
headers: Final = as_str_mapping(proxy_request.get("headers")) if proxy_request is not None else None
bodies: Final = tuple(
metadata
for key in ("metadata", "litellm_metadata")
if (metadata := as_str_mapping(request.get(key))) is not None
)
def scalar(control: str) -> str | None:
from_header: Final = as_str(headers.get(f"{LANGFUSE_HEADER_PREFIX}{control}")) if headers is not None else None
if from_header:
return from_header
return next((value for body in bodies if (value := as_str(body.get(control)))), None)
return TraceControls(
name=scalar("trace_name"),
user_id=scalar("trace_user_id"),
session_id=scalar("session_id"),
tags=next((tags for body in bodies if (tags := _str_items(body.get("tags")))), ()),
)
def _str_items(value: object) -> tuple[str, ...]:
try:
items: Final = _ITEMS.validate_python(value)
except ValidationError:
return ()
return tuple(item for item in items if isinstance(item, str) and item)

View file

@ -8,7 +8,13 @@ parsing lives in :mod:`litellm.integrations.otel.plumbing.providers` instead,
because it delegates to the OTel SDK's own W3C Baggage parser.
"""
from collections.abc import Mapping
from datetime import datetime
from typing import Final
from pydantic import TypeAdapter, ValidationError
_STR_MAPPING: Final = TypeAdapter(Mapping[str, object])
def as_str(value: object) -> str | None:
@ -55,6 +61,13 @@ def as_bool(value: object) -> bool | None:
return bool(value)
def as_str_mapping(value: object) -> Mapping[str, object] | None:
try:
return _STR_MAPPING.validate_python(value)
except ValidationError:
return None
def as_str_tuple(value: object) -> tuple[str, ...] | None:
if value is None:
return None

View file

@ -26,6 +26,7 @@ if TYPE_CHECKING:
from litellm.integrations.otel.model.destination import OtelDestination
_PROPAGATOR: Final = TraceContextTextMapPropagator()
_W3C_TRACE_HEADERS: Final = frozenset(("traceparent", "tracestate"))
# The request's root span — the FastAPI-owned SERVER span — captured ONCE when the
# proxy first resolves it, so request-level spans (the LLM call, guardrails) can
@ -310,6 +311,37 @@ def extract_traceparent(headers: Mapping[str, str]) -> Context | None:
return _PROPAGATOR.extract(carrier)
def _outgoing_trace_context(parent_span: object) -> Context | None:
if isinstance(parent_span, Span) and is_recordable_span(parent_span):
return context_from_span(parent_span)
root: Final = request_root_span()
if root is not None:
return context_from_span(root)
current: Final = get_current()
if is_recordable_span(get_current_span(current)):
return current
return None
def inject_trace_context(headers: Mapping[str, str], parent_span: object = None) -> dict[str, str]:
"""``headers`` plus W3C ``traceparent``/``tracestate`` for this request's span.
Parent preference: ``parent_span`` (the request span auth stashed on the key), then
the anchored request root span, then the ambient active span. Only trace context is
injected, never Baggage. Unchanged when no valid span exists anywhere.
"""
context: Final = _outgoing_trace_context(parent_span)
if context is None:
return dict(headers) # mutable-ok: OpenTelemetry propagator requires a mutable carrier
carrier: Final = { # mutable-ok: OpenTelemetry propagator requires a mutable carrier
key: value for key, value in headers.items() if key.lower() not in _W3C_TRACE_HEADERS
}
_PROPAGATOR.inject(carrier, context=context)
return carrier
# The OTLP destinations this request's key or team pointed its traces at, resolved
# once during auth. A ``ContextVar`` for the same reason the root span above is one:
# it rides the request task's context into the ``asyncio.create_task`` children that

View file

@ -160,7 +160,7 @@ def get_llm_provider(
if model is None:
raise ValueError("model parameter is required but was None. Please provide a valid model name.")
if litellm.LiteLLMProxyChatConfig._should_use_litellm_proxy_by_default(
if litellm.LiteLLMProxyChatConfig.should_use_litellm_proxy_by_default(
litellm_params=cast(LiteLLM_Params | None, litellm_params)
):
return litellm.LiteLLMProxyChatConfig.litellm_proxy_get_custom_llm_provider_info(

View file

@ -7,7 +7,7 @@ from collections.abc import Iterator, Mapping, Sequence
from datetime import datetime
from typing import TYPE_CHECKING, Any, Final
from litellm._logging import verbose_logger
from litellm._logging import format_base64_size, verbose_logger
from litellm.constants import (
BASE64_TRUNCATION_OFFLOAD_THRESHOLD_CHARS,
MAX_BASE64_LENGTH_FOR_LOGGING,
@ -40,9 +40,6 @@ import litellm
Helper utils used for logging callbacks
"""
_BYTES_PER_KIB: Final = 1024
_BYTES_PER_MIB: Final = 1024 * 1024
# Regex matching data-URI base64 content: "data:<mime>;base64,<payload>"
# Captures: group(1)=mime_type, group(2)=base64_payload
_DATA_URI_RE: Final = re.compile(r"data:([^;]+);base64,([A-Za-z0-9+/=]+)")
@ -52,23 +49,13 @@ _DATA_URI_RE: Final = re.compile(r"data:([^;]+);base64,([A-Za-z0-9+/=]+)")
_MAX_TRUNCATION_DEPTH: Final = 20
def _format_base64_size(num_chars: int) -> str:
"""Return a human-readable byte-size estimate from a base64 character count."""
num_bytes: Final = num_chars * 3 / 4
if num_bytes >= _BYTES_PER_MIB:
return f"{num_bytes / _BYTES_PER_MIB:.2f}MB"
if num_bytes >= _BYTES_PER_KIB:
return f"{num_bytes / _BYTES_PER_KIB:.1f}KB"
return f"{int(num_bytes)}B"
def _base64_data_uri_replacer(match: re.Match) -> str:
"""Replace a single base64 data-URI match with a size placeholder if too long."""
mime_type: Final = match.group(1)
payload: Final = match.group(2)
if len(payload) <= MAX_BASE64_LENGTH_FOR_LOGGING:
return match.group(0)
size_str: Final = _format_base64_size(len(payload))
size_str: Final = format_base64_size(len(payload))
return f"data:{mime_type};base64,[base64_data truncated: {size_str}]"

View file

@ -6,25 +6,29 @@ from pydantic import BaseModel
from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH
UNSERIALIZABLE_OBJECT: Final = "Unserializable Object"
def strip_null_bytes(value: str) -> str:
"""Strip NUL bytes, which PostgreSQL text/jsonb columns reject (error 22P05)."""
return value.replace("\x00", "")
def safe_dumps(
data: Any,
def safe_json_structure(
data: object,
max_depth: int = DEFAULT_MAX_RECURSE_DEPTH,
value_transform: Callable[[str | None, str], str] | None = None,
) -> str:
key: str | None = None,
) -> object:
"""
Recursively serialize data while detecting circular references.
Rebuild data out of JSON-native pieces while detecting circular references.
If a circular reference is detected then a marker string is returned.
NUL bytes are stripped from strings to prevent PostgreSQL 22P05 errors.
value_transform, when given, is applied to every string leaf (and to the
str() fallback for non-serializable objects) with the mapping key the leaf
was reached under, so callers can rewrite values without touching structure.
key is the mapping key data itself was reached under, when the caller has one.
"""
def _transform(key: str | None, value: str) -> str:
@ -75,7 +79,15 @@ def safe_dumps(
try:
return _transform(key, strip_null_bytes(str(obj)))
except Exception:
return "Unserializable Object"
return UNSERIALIZABLE_OBJECT
safe_data: Final = _serialize(data, set(), 0)
return json.dumps(safe_data, default=str)
return _serialize(data, set(), 0, key)
def safe_dumps(
data: Any,
max_depth: int = DEFAULT_MAX_RECURSE_DEPTH,
value_transform: Callable[[str | None, str], str] | None = None,
) -> str:
"""Serialize data to JSON text through safe_json_structure."""
return json.dumps(safe_json_structure(data, max_depth, value_transform), default=str)

View file

@ -148,6 +148,7 @@ class _ToolCallChunk(TypedDict):
class _UsageBearingChunk(TypedDict, total=False):
usage: Usage | None
_hidden_params: Mapping[str, str]
choices: ReadOnly[Sequence[StreamingChoices | Mapping[str, object]]]
class _UsageSummary(TypedDict):
@ -921,21 +922,22 @@ class ChunkProcessor:
prompt_tokens_details = attach_cache_creation_token_details(prompt_tokens_details, cache_creation_token_details)
completion_tokens = self._reset_anthropic_cursor_completion_tokens(
recovered_completion_tokens: Final = self._reset_anthropic_cursor_completion_tokens(
chunks=chunks,
completion_tokens=completion_tokens,
completion_usage_updates=completion_usage_updates,
)
cursor_was_reset: Final = recovered_completion_tokens != completion_tokens
return UsagePerChunk(
prompt_tokens=prompt_tokens,
completion_tokens=completion_tokens,
completion_tokens=recovered_completion_tokens,
cache_creation_input_tokens=cache_creation_input_tokens,
cache_read_input_tokens=cache_read_input_tokens,
server_tool_use=server_tool_use,
web_search_requests=web_search_requests,
google_maps_grounding_requests=google_maps_grounding_requests,
completion_tokens_details=completion_tokens_details,
completion_tokens_details=None if cursor_was_reset else completion_tokens_details,
prompt_tokens_details=prompt_tokens_details,
cost=cost,
inference_geo=self._last_provider_pricing_field(chunks, "inference_geo"),
@ -960,6 +962,30 @@ class ChunkProcessor:
]
return values[-1] if values else None
@staticmethod
def _finish_reason_of_choice(choice: object) -> str | None:
match choice:
case StreamingChoices(finish_reason=reason) | Choices(finish_reason=reason):
return reason
case {"finish_reason": str() as reason}:
return reason
case _:
return None
@staticmethod
def _chunk_choices(chunk: "_UsageBearingChunk | ModelResponse | ModelResponseStream") -> Sequence[object]:
if isinstance(chunk, dict):
return chunk.get("choices", ())
return getattr(chunk, "choices", ())
@staticmethod
def _saw_finish_reason(chunks: Sequence["_UsageBearingChunk | ModelResponse"]) -> bool:
return any(
ChunkProcessor._finish_reason_of_choice(choice) is not None
for chunk in chunks
for choice in ChunkProcessor._chunk_choices(chunk)
)
@staticmethod
def _reset_anthropic_cursor_completion_tokens(
chunks: Sequence["_UsageBearingChunk | ModelResponse"],
@ -970,18 +996,18 @@ class ChunkProcessor:
See the ``completion_usage_updates`` comment in
``_calculate_usage_per_chunk``. The accumulated value is NOT a stale
cursor when either it is > 1 (definitely not a placeholder) or we saw
>= 2 completion-bearing usage events (positive evidence ``message_delta``
arrived). Otherwise the only completion update we ever saw was the
Anthropic ``message_start`` cursor (=1) reset to 0 so
``calculate_usage()``'s ``or token_counter(text=...)`` fallback estimates
from the actually-received completion text instead of trusting the
placeholder. Gated on ``custom_llm_provider == "anthropic"`` so the
heuristic (which encodes Anthropic's specific message_start SSE shape)
does not silently affect other providers that may legitimately report
``completion_tokens=1`` from a single usage event.
cursor when we saw >= 2 completion-bearing usage events or any chunk
carried a ``finish_reason`` (positive evidence ``message_delta``
arrived). Otherwise the only completion update we ever saw was the
Anthropic ``message_start`` cursor, a small placeholder whose magnitude
varies per request (1 and 8 both observed live), so reset to 0 and let
``calculate_usage()``'s ``or token_counter(...)`` fallback estimate from
the actually-received text and reasoning instead. Gated on
``custom_llm_provider == "anthropic"`` so the heuristic (which encodes
Anthropic's specific message_start SSE shape) does not silently affect
other providers that legitimately report usage from a single event.
"""
saw_non_cursor_completion: Final = completion_tokens > 1 or completion_usage_updates >= 2
saw_non_cursor_completion: Final = completion_usage_updates >= 2 or ChunkProcessor._saw_finish_reason(chunks)
if saw_non_cursor_completion:
return completion_tokens
@ -995,7 +1021,7 @@ class ChunkProcessor:
if isinstance(hp, dict):
custom_llm_provider = hp.get("custom_llm_provider")
if custom_llm_provider == "anthropic" and completion_tokens == 1:
if custom_llm_provider == "anthropic":
return 0
return completion_tokens
@ -1039,10 +1065,13 @@ class ChunkProcessor:
returned_usage.prompt_tokens = 0
returned_usage.completion_tokens = (
completion_tokens
or token_counter(
model=model,
text=completion_output,
count_response_tokens=True, # count_response_tokens is a Flag to tell token counter this is a response, No need to add extra tokens we do for input messages
or (
token_counter(
model=model,
text=completion_output,
count_response_tokens=True, # count_response_tokens is a Flag to tell token counter this is a response, No need to add extra tokens we do for input messages
)
+ (reasoning_tokens or 0)
)
)
returned_usage.total_tokens = returned_usage.prompt_tokens + returned_usage.completion_tokens
@ -1066,15 +1095,16 @@ class ChunkProcessor:
returned_usage.completion_tokens_details = completion_tokens_details
if reasoning_tokens is not None:
capped_reasoning_tokens: Final = min(max(0, reasoning_tokens), returned_usage.completion_tokens)
if returned_usage.completion_tokens_details is None:
returned_usage.completion_tokens_details = CompletionTokensDetailsWrapper(
reasoning_tokens=reasoning_tokens
reasoning_tokens=capped_reasoning_tokens,
text_tokens=returned_usage.completion_tokens - capped_reasoning_tokens,
)
elif (
returned_usage.completion_tokens_details is not None
and returned_usage.completion_tokens_details.reasoning_tokens is None
):
capped_reasoning_tokens: Final = min(max(0, reasoning_tokens), returned_usage.completion_tokens)
returned_usage.completion_tokens_details.reasoning_tokens = capped_reasoning_tokens
if returned_usage.completion_tokens_details.text_tokens is None:
returned_usage.completion_tokens_details.text_tokens = (

View file

@ -31,6 +31,7 @@ from litellm.llms.anthropic.experimental_pass_through.adapters.transformation im
)
from litellm.llms.base_llm.guardrail_translation.base_translation import (
BaseTranslation,
RequestScanContext,
StreamingScanKey,
StreamTransformSink,
)
@ -527,6 +528,26 @@ class AnthropicMessagesHandler(BaseTranslation):
)
return result if result else None
def request_scan_context(
self, data: Mapping[str, object], guardrail_to_apply: "CustomGuardrail"
) -> RequestScanContext:
if data.get("messages") is None:
return RequestScanContext()
translated: Final = self._translate_to_openai(
{key: value for key, value in data.items() if key != "system"} # mutable-ok: API message payload
)
hoisted_system_message: Final = (
None
if effective_skip_system_message_for_guardrail(guardrail_to_apply)
else self._hoisted_top_level_system_message(data)
)
return RequestScanContext.scoped(
(*(() if hoisted_system_message is None else (hoisted_system_message,)), *translated["messages"]),
tuple(tool for tool in translated.get("tools") or () if not is_provider_native_tool_dict(tool)),
guardrail_to_apply,
skip_system=False,
)
async def process_input_messages(
self,
data: dict,
@ -696,9 +717,7 @@ class AnthropicMessagesHandler(BaseTranslation):
return data
def _hoisted_top_level_system_message(
self, data: dict
) -> AllMessageValues | None: # mutable-ok: API message payload
def _hoisted_top_level_system_message(self, data: Mapping[str, object]) -> AllMessageValues | None:
"""Return the system message produced by translating the top-level prompt."""
system: Final = data.get("system")
if not system:
@ -1200,7 +1219,7 @@ class AnthropicMessagesHandler(BaseTranslation):
)
guardrailed_inputs: Final = await guardrail_to_apply.apply_guardrail(
inputs=inputs,
inputs=self.with_response_context(inputs, request_data, guardrail_to_apply),
request_data=request_data,
input_type="response",
logging_obj=litellm_logging_obj,
@ -1273,7 +1292,7 @@ class AnthropicMessagesHandler(BaseTranslation):
key="response",
)
_guardrailed_inputs = await guardrail_to_apply.apply_guardrail(
inputs=guardrail_inputs,
inputs=self.with_response_context(guardrail_inputs, prepared_request_data, guardrail_to_apply),
request_data=prepared_request_data,
input_type="response",
logging_obj=litellm_logging_obj,
@ -1319,7 +1338,11 @@ class AnthropicMessagesHandler(BaseTranslation):
key="responses",
)
_guardrailed_inputs = await guardrail_to_apply.apply_guardrail(
inputs={"texts": [string_so_far]},
inputs=self.with_response_context(
GenericGuardrailAPIInputs(texts=[string_so_far]), # mutable-ok: guardrail inputs want a list
prepared_request_data,
guardrail_to_apply,
),
request_data=prepared_request_data,
input_type="response",
logging_obj=litellm_logging_obj,

View file

@ -1180,7 +1180,7 @@ class LiteLLMAnthropicMessagesAdapter:
self._add_system_message_to_messages(new_messages, anthropic_message_request)
new_kwargs: Final[ChatCompletionRequest] = {
"model": anthropic_message_request["model"],
"model": anthropic_message_request.get("model", ""),
"messages": new_messages,
}
## CONVERT METADATA (user_id + litellm metadata)

View file

@ -49,12 +49,7 @@ class AzureOpenAIResponsesAPIConfig(OpenAIResponsesAPIConfig):
return BaseAzureLLM._base_validate_azure_environment(headers=headers, litellm_params=litellm_params)
def get_stripped_model_name(self, model: str) -> str:
# if "responses/" is in the model name, remove it
if "responses/" in model:
model = model.replace("responses/", "")
if "o_series" in model:
model = model.replace("o_series/", "")
return model
return model.replace("responses/", "").replace("o_series/", "").replace("azure_ai/", "")
def _handle_reasoning_item(self, item: dict[str, Any]) -> dict[str, Any]:
"""

View file

@ -9,6 +9,7 @@ from litellm.types.llms.openai import AllMessageValues
from litellm.types.router import GenericLiteLLMParams
AzureAIApiKeyHeader = Literal["Authorization", "api-key", "Api-Key", "Ocp-Apim-Subscription-Key"]
AZURE_OPENAI_V1_HOST_SUFFIXES: Final = (".services.ai.azure.com", ".openai.azure.com")
def is_foundry_model_inference_base(api_base: str) -> bool:
@ -19,11 +20,13 @@ def is_foundry_model_inference_base(api_base: str) -> bool:
return "/openai/deployments" not in parsed.path
def api_key_header_for_base(api_base: str | None) -> AzureAIApiKeyHeader:
def is_azure_openai_v1_host(api_base: str | None) -> bool:
host: Final = urlparse(api_base).hostname if api_base else None
if host and (host.endswith(".services.ai.azure.com") or host.endswith(".openai.azure.com")):
return "api-key"
return "Authorization"
return host is not None and host.endswith(AZURE_OPENAI_V1_HOST_SUFFIXES)
def api_key_header_for_base(api_base: str | None) -> AzureAIApiKeyHeader:
return "api-key" if is_azure_openai_v1_host(api_base) else "Authorization"
def get_azure_ai_entra_token(litellm_params: Mapping[str, object] | None = None) -> str | None:
@ -70,6 +73,17 @@ def get_azure_ai_auth_headers(
AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY: Final = "azure_model_router_selected_model"
def azure_ai_supports_native_responses(model: str | None, api_base: str | None) -> bool:
resolved_base: Final = AzureFoundryModelInfo.get_api_base(api_base)
if resolved_base is not None and not is_azure_openai_v1_host(resolved_base):
return False
if model is None:
return True
if "claude" in model.lower():
return False
return AzureFoundryModelInfo.get_azure_ai_route(model) == "default"
class AzureFoundryModelInfo(BaseLLMModelInfo):
"""Model info for Azure AI / Azure Foundry models."""

View file

@ -0,0 +1,53 @@
from typing import Final
import httpx
from litellm.llms.azure.responses.transformation import AzureOpenAIResponsesAPIConfig
from litellm.llms.azure_ai.common_utils import (
AzureFoundryModelInfo,
api_key_header_for_base,
get_azure_ai_auth_headers,
)
from litellm.types.router import GenericLiteLLMParams
from litellm.types.utils import LlmProviders
_PROJECT_PATH_PREFIX: Final = ("api", "projects")
_RESPONSES_PATH: Final = ("openai", "v1", "responses")
def _responses_url(api_base: str) -> str:
base_url: Final = httpx.URL(api_base)
segments: Final = tuple(segment for segment in base_url.path.split("/") if segment)
project_root: Final = segments[:3] if segments[:2] == _PROJECT_PATH_PREFIX else ()
return str(base_url.copy_with(path="/" + "/".join((*project_root, *_RESPONSES_PATH)), query=None))
class AzureAIResponsesAPIConfig(AzureOpenAIResponsesAPIConfig):
@property
def custom_llm_provider(self) -> LlmProviders:
return LlmProviders.AZURE_AI
def validate_environment(self, headers: dict, model: str, litellm_params: GenericLiteLLMParams | None) -> dict:
params: Final = litellm_params or GenericLiteLLMParams()
auth_headers: Final = get_azure_ai_auth_headers(
api_key=AzureFoundryModelInfo.get_api_key(params.api_key),
litellm_params=params.model_dump(),
api_key_header=api_key_header_for_base(AzureFoundryModelInfo.get_api_base(params.api_base)),
)
return { # mutable-ok: the handler updates the returned headers in place per the dict contract
**headers,
**auth_headers,
"Content-Type": "application/json",
}
def supports_native_websocket(self) -> bool:
return False
def get_complete_url(self, api_base: str | None, litellm_params: dict) -> str:
resolved_base: Final = AzureFoundryModelInfo.get_api_base(api_base)
if resolved_base is None:
raise ValueError(
"api_base is required for the Azure AI Foundry Responses API. "
"Set the api_base parameter or the AZURE_AI_API_BASE environment variable."
)
return _responses_url(resolved_base)

View file

@ -1,8 +1,17 @@
from abc import ABC, abstractmethod
from collections.abc import Sequence
from collections.abc import Mapping, Sequence
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any, ClassVar, Final, Optional
from litellm.llms.base_llm.guardrail_translation.utils import (
effective_scan_only_tool_results_for_guardrail,
effective_skip_system_message_for_guardrail,
effective_skip_tool_message_for_guardrail,
request_tools,
response_assistant_turn,
scoped_structured_message_indices,
)
if TYPE_CHECKING:
from fastapi import HTTPException
@ -12,7 +21,43 @@ if TYPE_CHECKING:
)
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.llms.openai import AllMessageValues
from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolParam
from litellm.types.utils import GenericGuardrailAPIInputs
@dataclass(frozen=True, slots=True)
class RequestScanContext:
"""The scoped request turns and tool definitions a guardrail's request scan sees, in OpenAI chat shape."""
structured_messages: tuple["AllMessageValues", ...] = ()
tools: tuple["ChatCompletionToolParam", ...] = ()
conversation_supplied: bool = False
@staticmethod
def scoped(
structured_messages: Sequence["AllMessageValues"],
tools: Sequence["ChatCompletionToolParam"],
guardrail_to_apply: "CustomGuardrail",
*,
skip_system: bool | None = None,
) -> "RequestScanContext":
scan_only_tool_results: Final = effective_scan_only_tool_results_for_guardrail(guardrail_to_apply)
scoped_indices: Final = scoped_structured_message_indices(
structured_messages,
scan_only_tool_results=scan_only_tool_results,
skip_system=(
effective_skip_system_message_for_guardrail(guardrail_to_apply) if skip_system is None else skip_system
),
skip_tool=effective_skip_tool_message_for_guardrail(guardrail_to_apply),
)
return RequestScanContext(
structured_messages=tuple(structured_messages[index] for index in scoped_indices),
tools=() if scan_only_tool_results else tuple(tools),
conversation_supplied=bool(structured_messages),
)
REQUEST_SCAN_CONTEXT_KEY: Final = "litellm_request_scan_context"
@dataclass(slots=True)
@ -257,6 +302,50 @@ class BaseTranslation(ABC):
"""
return None
def request_scan_context(
self, data: Mapping[str, object], guardrail_to_apply: "CustomGuardrail"
) -> RequestScanContext:
"""Override wherever ``process_input_messages`` scopes or translates the request differently."""
structured_messages: Final = self.get_structured_messages(
dict(data) # mutable-ok: get_structured_messages takes the request as a dict
)
return RequestScanContext.scoped(
structured_messages or (), request_tools(data.get("tools")), guardrail_to_apply
)
def with_response_context(
self,
inputs: "GenericGuardrailAPIInputs",
request_data: Mapping[str, object] | None,
guardrail_to_apply: "CustomGuardrail",
) -> "GenericGuardrailAPIInputs":
"""``inputs`` plus the scoped request conversation, closed by the scanned reply, and the request tools."""
if request_data is None:
return inputs
precomputed: Final = request_data.get(REQUEST_SCAN_CONTEXT_KEY)
context: Final = (
precomputed
if isinstance(precomputed, RequestScanContext)
else self.request_scan_context(request_data, guardrail_to_apply)
)
if not context.conversation_supplied:
return inputs
assistant_turn: Final = response_assistant_turn(inputs.get("texts") or (), inputs.get("tool_calls") or ())
contextual_inputs: Final[GenericGuardrailAPIInputs] = {
**inputs,
"structured_messages": [ # mutable-ok: GenericGuardrailAPIInputs fields are lists
*context.structured_messages,
*(() if assistant_turn is None else (assistant_turn,)),
],
}
if not context.tools:
return contextual_inputs
with_tools: Final[GenericGuardrailAPIInputs] = {
**contextual_inputs,
"tools": list(context.tools), # mutable-ok: GenericGuardrailAPIInputs fields are lists
}
return with_tools
def extract_request_tool_names(self, data: dict) -> list[str]:
"""
Extract tool names from the request body for allowlist/policy checks.

View file

@ -2,12 +2,24 @@ from __future__ import annotations
import json
from collections.abc import Callable, Iterator, Mapping, Sequence
from typing import Final, TypeVar, cast # noqa: TID251 # a rebuilt chat row has no typed constructor across roles
from typing import TYPE_CHECKING, Final, TypeVar, cast # noqa: TID251 # a rebuilt chat row has no typed constructor
from pydantic import BaseModel
from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicUsage
from litellm.types.llms.openai import AllMessageValues, ResponseAPIUsage
from litellm.types.llms.openai import (
AllMessageValues,
ChatCompletionAssistantMessage,
ChatCompletionAssistantToolCall,
ChatCompletionTextObject,
ChatCompletionToolCallChunk,
ChatCompletionToolCallFunctionChunk,
ChatCompletionToolParam,
ResponseAPIUsage,
)
if TYPE_CHECKING:
from litellm.types.utils import ChatCompletionMessageToolCall
def _anthropic_stream_chunk_events(item: object) -> list[dict]:
@ -278,9 +290,57 @@ def scoped_structured_message_indices(
)
def _assistant_tool_call(
tool_call: ChatCompletionToolCallChunk | ChatCompletionMessageToolCall,
) -> ChatCompletionAssistantToolCall:
function: Final = stream_item_field(tool_call, "function")
tool_call_id: Final = stream_item_field(tool_call, "id")
name: Final = stream_item_field(function, "name")
arguments: Final = stream_item_field(function, "arguments")
return ChatCompletionAssistantToolCall(
id=tool_call_id if isinstance(tool_call_id, str) else None,
type="function",
function=ChatCompletionToolCallFunctionChunk(
name=name if isinstance(name, str) else None,
arguments=arguments if isinstance(arguments, str) else "",
),
)
def response_assistant_turn(
texts: Sequence[str],
tool_calls: Sequence[ChatCompletionToolCallChunk] | Sequence[ChatCompletionMessageToolCall],
) -> ChatCompletionAssistantMessage | None:
"""The scanned reply as the assistant turn closing the request conversation."""
assistant_tool_calls: Final = tuple(_assistant_tool_call(tool_call) for tool_call in tool_calls)
if not texts and not assistant_tool_calls:
return None
content: Final = (
texts[0]
if len(texts) == 1
else tuple(ChatCompletionTextObject(type="text", text=text) for text in texts) or None
)
if not assistant_tool_calls:
return ChatCompletionAssistantMessage(role="assistant", content=content)
return ChatCompletionAssistantMessage(
role="assistant",
content=content,
tool_calls=list(assistant_tool_calls), # mutable-ok: the assistant message type takes a list
)
ToolT = TypeVar("ToolT")
def request_tools(raw_tools: object) -> tuple[ChatCompletionToolParam, ...]:
"""The request's ``tools`` list, as the chat completion request model already validated it upstream."""
if not isinstance(raw_tools, list):
return ()
return tuple(
cast(Sequence[ChatCompletionToolParam], raw_tools) # cast-ok: the request model validated tools upstream
)
def openai_tool_name(tool: object) -> str | None:
if not isinstance(tool, dict):
return None

View file

@ -53,6 +53,7 @@ from litellm.types.llms.openai import (
AllMessageValues,
ChatCompletionAnnotation,
ChatCompletionAssistantMessage,
ChatCompletionAssistantToolCall,
ChatCompletionRedactedThinkingBlock,
ChatCompletionResponseMessage,
ChatCompletionSystemMessage,
@ -205,6 +206,84 @@ class AmazonConverseConfig(BaseConfig):
return messages_copy
@staticmethod
def _has_orphaned_tool_blocks(messages: list[AllMessageValues]) -> bool:
return any(
(m.get("role") == "assistant" and m.get("tool_calls")) or m.get("role") in ("tool", "function")
for m in messages
)
@staticmethod
def _neutralize_orphaned_tool_blocks(
messages: list[AllMessageValues], optional_params: dict
) -> list[AllMessageValues]:
if optional_params.get("tools") or not AmazonConverseConfig._has_orphaned_tool_blocks(messages):
return messages
from litellm.litellm_core_utils.prompt_templates.common_utils import (
convert_content_list_to_str,
)
def _tool_call_text(tool_call: ChatCompletionAssistantToolCall) -> str:
function = tool_call.get("function") or {}
name = function.get("name") or "unknown_tool"
arguments = function.get("arguments") or ""
call_id = tool_call.get("id")
label = f"tool call {call_id}" if call_id else "tool call"
return f"[{label}: {name}({arguments})]"
def _result_text(message: AllMessageValues) -> str:
rendered = convert_content_list_to_str(message).strip()
return rendered or "<non-text tool result omitted>"
guardrail_active: Final = "guardrailConfig" in optional_params
def _rewrite(message: AllMessageValues) -> AllMessageValues:
role = message.get("role")
tool_calls = message.get("tool_calls")
if role == "assistant" and tool_calls:
base_text: Final = convert_content_list_to_str(message)
call_texts: Final = tuple(_tool_call_text(call) for call in tool_calls)
text: Final = "\n".join(part for part in (base_text, *call_texts) if part)
return ChatCompletionAssistantMessage(role="assistant", content=text)
if role in ("tool", "function"):
tool_call_id = message.get("tool_call_id")
name = message.get("name")
label = f"tool result for {tool_call_id or name or 'unknown'}"
result_text: Final = f"[{label}: {_result_text(message)}]"
# Tool results are externally controlled, so guard them wherever they
# land in history; _convert_consecutive_user_messages_to_guarded_text
# only covers the trailing user turn.
content: Final = [{"type": "guarded_text", "text": result_text}] if guardrail_active else result_text
return ChatCompletionUserMessage(role="user", content=content)
return message
verbose_logger.warning(
"litellm.bedrock: request has tool blocks in message history but no "
"`tools=` param; neutralizing orphaned tool blocks to text so Bedrock "
"accepts the request without a toolConfig. Non-text tool-result "
"payloads are dropped. Pass `tools=` to preserve structured tool calling."
)
return [_rewrite(message) for message in messages]
@staticmethod
def _handle_orphaned_tool_blocks(messages: list[AllMessageValues], optional_params: dict) -> list[AllMessageValues]:
if litellm.bedrock_neutralize_orphaned_tool_blocks:
return AmazonConverseConfig._neutralize_orphaned_tool_blocks(messages, optional_params)
if "tools" in optional_params or not has_tool_call_blocks(messages):
return messages
if litellm.modify_params:
optional_params["tools"] = add_dummy_tool(custom_llm_provider="bedrock_converse")
return messages
raise litellm.utils.UnsupportedParamsError(
message="Bedrock doesn't support tool calling without `tools=` param specified. Pass `tools=` param OR set `litellm.modify_params = True` // `litellm_settings::modify_params: True` to add dummy tool to the request.",
model="",
llm_provider="bedrock",
)
@classmethod
def get_config(cls):
return {
@ -1609,20 +1688,6 @@ class AmazonConverseConfig(BaseConfig):
drop_params: bool = False,
litellm_params: Mapping[str, object] | None = None,
) -> CommonRequestObject:
## VALIDATE REQUEST
"""
Bedrock doesn't support tool calling without `tools=` param specified.
"""
if "tools" not in optional_params and messages is not None and has_tool_call_blocks(messages):
if litellm.modify_params:
optional_params["tools"] = add_dummy_tool(custom_llm_provider="bedrock_converse")
else:
raise litellm.UnsupportedParamsError(
message="Bedrock doesn't support tool calling without `tools=` param specified. Pass `tools=` param OR set `litellm.modify_params = True` // `litellm_settings::modify_params: True` to add dummy tool to the request.",
model="",
llm_provider="bedrock",
)
# Drop thinking param if thinking is enabled but thinking_blocks are missing
# This prevents the error: "Expected thinking or redacted_thinking, but found tool_use"
#
@ -1735,7 +1800,9 @@ class AmazonConverseConfig(BaseConfig):
messages, system_content_blocks = self._transform_system_message(messages, model=model)
# Convert last user message to guarded_text if guardrailConfig is present
messages = self._convert_consecutive_user_messages_to_guarded_text(messages, optional_params)
messages = self._convert_consecutive_user_messages_to_guarded_text(
self._handle_orphaned_tool_blocks(messages, optional_params), optional_params
)
## TRANSFORMATION ##
_data: Final[CommonRequestObject] = self._transform_request_helper(
@ -1796,7 +1863,9 @@ class AmazonConverseConfig(BaseConfig):
messages, system_content_blocks = self._transform_system_message(messages, model=model)
# Convert last user message to guarded_text if guardrailConfig is present
messages = self._convert_consecutive_user_messages_to_guarded_text(messages, optional_params)
messages = self._convert_consecutive_user_messages_to_guarded_text(
self._handle_orphaned_tool_blocks(messages, optional_params), optional_params
)
_data: Final[CommonRequestObject] = self._transform_request_helper(
model=model,

View file

@ -22,6 +22,7 @@ from litellm.llms.bedrock_mantle.common_utils import (
BEDROCK_MANTLE_DEFAULT_REGION,
BedrockMantleAuthMixin,
)
from litellm.llms.openai.chat.gpt_5_transformation import is_gpt_reasoning_series_name
from litellm.secret_managers.main import get_secret_str
from litellm.types.llms.openai import AllMessageValues
from litellm.types.router import GenericLiteLLMParams
@ -108,13 +109,22 @@ class BedrockMantleChatConfig(BedrockMantleAuthMixin, OpenAILikeChatConfig):
def get_supported_openai_params(self, model: str) -> list:
base_params: Final = super().get_supported_openai_params(model)
extra_params: Final = tuple(
param
for param, supported in (
("verbosity", is_gpt_reasoning_series_name(model)),
("reasoning_effort", self._supports_reasoning(model)),
)
if supported and param not in base_params
)
return [*base_params, *extra_params] # mutable-ok: fresh list required by the inherited signature
def _supports_reasoning(self, model: str) -> bool:
try:
if litellm.supports_reasoning(model=model, custom_llm_provider=self.custom_llm_provider):
if "reasoning_effort" not in base_params:
base_params.append("reasoning_effort")
return litellm.supports_reasoning(model=model, custom_llm_provider=self.custom_llm_provider)
except Exception as e:
verbose_logger.debug("BedrockMantleChatConfig: error checking reasoning support: %s", e)
return base_params
return False
def get_model_response_iterator(
self,

View file

@ -12,6 +12,12 @@ from ...openai.chat.gpt_transformation import OpenAIGPTConfig
class DashScopeChatConfig(OpenAIGPTConfig):
def get_supported_openai_params(self, model: str) -> list[str]: # mutable-ok: base class contract returns a list
return [ # mutable-ok: base class contract returns a list
*super().get_supported_openai_params(model=model),
"reasoning_effort",
]
def remove_cache_control_flag_from_messages_and_tools(
self,
model: str,

View file

@ -54,7 +54,7 @@ class LiteLLMProxyChatConfig(OpenAIGPTConfig):
return api_key or get_secret_str("LITELLM_PROXY_API_KEY")
@staticmethod
def _should_use_litellm_proxy_by_default(
def should_use_litellm_proxy_by_default(
litellm_params: LiteLLM_Params | None = None,
):
"""

View file

@ -4,9 +4,9 @@ from typing import Final
import litellm
from litellm.utils import (
_is_explicitly_disabled_factory,
_supports_factory,
declared_value_factory,
is_explicitly_disabled_factory,
)
from .gpt_transformation import OpenAIGPTConfig
@ -192,7 +192,7 @@ class OpenAIGPT5Config(OpenAIGPTConfig):
Use this for opt-out checks where unknown models should be allowed through.
"""
return _is_explicitly_disabled_factory(
return is_explicitly_disabled_factory(
model=cls._model_map_lookup_name(model),
custom_llm_provider=None,
key=f"supports_{level}_reasoning_effort",

View file

@ -452,7 +452,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
inputs["model"] = response.model
guardrailed_inputs: Final = await guardrail_to_apply.apply_guardrail(
inputs=inputs,
inputs=self.with_response_context(inputs, request_data, guardrail_to_apply),
request_data=request_data,
input_type="response",
logging_obj=litellm_logging_obj,
@ -615,7 +615,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
if responses_so_far and hasattr(responses_so_far[0], "model") and responses_so_far[0].model:
inputs["model"] = responses_so_far[0].model
guardrailed_inputs: Final = await guardrail_to_apply.apply_guardrail(
inputs=inputs,
inputs=self.with_response_context(inputs, request_data, guardrail_to_apply),
request_data=request_data,
input_type="response",
logging_obj=litellm_logging_obj,
@ -760,7 +760,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
if responses_so_far and getattr(responses_so_far[0], "model", None):
inputs["model"] = responses_so_far[0].model
guardrailed_inputs: Final = await guardrail_to_apply.apply_guardrail(
inputs=inputs,
inputs=self.with_response_context(inputs, request_data, guardrail_to_apply),
request_data=request_data,
input_type="response",
logging_obj=litellm_logging_obj,

View file

@ -48,6 +48,7 @@ from litellm.completion_extras.litellm_responses_transformation.transformation i
)
from litellm.llms.base_llm.guardrail_translation.base_translation import (
BaseTranslation,
RequestScanContext,
StreamingScanKey,
StreamTransformSink,
)
@ -451,6 +452,28 @@ class OpenAIResponsesHandler(BaseTranslation):
)
return cast(list[AllMessageValues], messages) if messages else None
def request_scan_context(
self, data: Mapping[str, object], guardrail_to_apply: "CustomGuardrail"
) -> RequestScanContext:
raw_tools: Final = data.get("tools")
structured_messages: Final = tuple(
self.get_structured_messages(
dict(data) # mutable-ok: get_structured_messages takes the request as a dict
)
or ()
)
return RequestScanContext(
structured_messages=structured_messages,
tools=tuple(
cast(ChatCompletionToolParam, tool) # cast-ok: mcp tools ride along in the guardrail's tool list
for form in LiteLLMCompletionResponsesConfig.responses_tools_to_chat_forms(
tuple(raw_tools) if isinstance(raw_tools, list) else ()
)
for tool in form.chat_tools
),
conversation_supplied=bool(structured_messages),
)
async def process_input_messages(
self,
data: dict,
@ -754,7 +777,7 @@ class OpenAIResponsesHandler(BaseTranslation):
pre_guardrail_tool_calls: Final = _tool_call_shapes(tool_calls_to_check)
guardrailed_inputs: Final = await guardrail_to_apply.apply_guardrail(
inputs=inputs,
inputs=self.with_response_context(inputs, request_data, guardrail_to_apply),
request_data=request_data,
input_type="response",
logging_obj=litellm_logging_obj,
@ -867,7 +890,7 @@ class OpenAIResponsesHandler(BaseTranslation):
pre_guardrail_tool_calls: Final = _tool_call_shapes(tool_calls_to_check)
guardrailed_inputs: Final = await guardrail_to_apply.apply_guardrail(
inputs=inputs,
inputs=self.with_response_context(inputs, request_data, guardrail_to_apply),
request_data=request_data,
input_type="response",
logging_obj=litellm_logging_obj,
@ -926,7 +949,7 @@ class OpenAIResponsesHandler(BaseTranslation):
if hasattr(model_response_stream, "model") and model_response_stream.model:
inputs["model"] = model_response_stream.model
await guardrail_to_apply.apply_guardrail(
inputs=inputs,
inputs=self.with_response_context(inputs, request_data, guardrail_to_apply),
request_data=request_data if request_data is not None else {},
input_type="response",
logging_obj=litellm_logging_obj,
@ -949,7 +972,7 @@ class OpenAIResponsesHandler(BaseTranslation):
if response_model:
fallback_inputs["model"] = response_model
fallback_outputs: Final = await guardrail_to_apply.apply_guardrail(
inputs=fallback_inputs,
inputs=self.with_response_context(fallback_inputs, request_data, guardrail_to_apply),
request_data=request_data if request_data is not None else {},
input_type="response",
logging_obj=litellm_logging_obj,

View file

@ -0,0 +1,92 @@
import hashlib
import json
from collections.abc import Mapping
from types import MappingProxyType
from typing import Annotated, Final, TypeAlias
import httpx
from pydantic import BaseModel, BeforeValidator, ConfigDict
from litellm._logging import verbose_logger
from litellm.caching.in_memory_cache import InMemoryCache
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
from litellm.utils import _add_path_to_api_base # pyright: ignore[reportPrivateUsage] # shared provider URL helper
MODEL_INFO_REFRESH_SECONDS: Final = 300
MODEL_INFO_REFRESH_CONCURRENCY: Final = 8
MODEL_INFO_DISCOVERY_PROVIDERS: Final = frozenset({"hosted_vllm", "openai", "text-completion-openai", "openai_like"})
_EMPTY_LIMITS: Final[Mapping[str, int]] = MappingProxyType({})
def _positive_limit(value: object) -> int | None:
return value if isinstance(value, int) and not isinstance(value, bool) and value > 0 else None
_TokenLimit: TypeAlias = Annotated[int | None, BeforeValidator(_positive_limit)]
class _ModelCard(BaseModel):
model_config = ConfigDict(frozen=True)
id: str
max_model_len: _TokenLimit = None
context_length: _TokenLimit = None
max_input_tokens: _TokenLimit = None
max_output_tokens: _TokenLimit = None
def token_limits(self) -> Mapping[str, int]:
context: Final = self.max_model_len or self.context_length
input_limit: Final = self.max_input_tokens or context
output_limit: Final = self.max_output_tokens or context
return MappingProxyType(
{
key: value
for key, value in (
("max_tokens", context),
("max_input_tokens", min(input_limit, context) if input_limit and context else input_limit),
("max_output_tokens", min(output_limit, context) if output_limit and context else output_limit),
)
if value is not None
}
)
class _ModelList(BaseModel):
model_config = ConfigDict(frozen=True)
data: tuple[_ModelCard, ...] = ()
async def get_openai_compatible_model_info(
*,
model: str,
api_base: str,
headers: Mapping[str, str],
client: AsyncHTTPHandler,
cache: InMemoryCache,
) -> Mapping[str, int]:
url: Final = _add_path_to_api_base(api_base, "/v1/models")
cache_key: Final = (
"upstream_model_info:" + hashlib.sha256(json.dumps((url, sorted(headers.items()))).encode()).hexdigest()
)
cached: Final[object] = cache.get_cache(cache_key)
if isinstance(cached, _ModelList):
return next((card.token_limits() for card in cached.data if card.id == model), _EMPTY_LIMITS)
try:
response: Final = await client.get(
url=url,
headers=dict(headers), # mutable-ok: AsyncHTTPHandler requires a concrete dict
timeout=httpx.Timeout(5.0),
follow_redirects=False,
max_response_bytes=2 * 1024 * 1024,
)
response.raise_for_status()
models: Final = _ModelList.model_validate_json(response.content)
except Exception: # noqa: BLE001 # optional upstream metadata must not interrupt proxy refresh
verbose_logger.debug("Could not discover upstream model token limits")
cache.set_cache(cache_key, _ModelList(), ttl=60)
return _EMPTY_LIMITS
cache.set_cache(cache_key, models, ttl=MODEL_INFO_REFRESH_SECONDS)
return next((card.token_limits() for card in models.data if card.id == model), _EMPTY_LIMITS)

View file

@ -79,6 +79,7 @@ from litellm.utils import (
CustomStreamWrapper,
ModelResponse,
is_base64_encoded,
is_explicitly_disabled_factory,
supports_reasoning,
)
@ -866,6 +867,14 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
else:
raise _unsupported_reasoning_effort(reasoning_effort)
@staticmethod
def _supports_minimal_thinking_level(model: str) -> bool:
lowered: Final = model.lower()
is_gemini3flash: Final = "gemini-3" in lowered and "flash" in lowered
return is_gemini3flash and not is_explicitly_disabled_factory(
model=model, custom_llm_provider=None, key="supports_minimal_reasoning_effort"
)
@staticmethod
def _map_reasoning_effort_to_thinking_level(
reasoning_effort: str,
@ -880,13 +889,11 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
Returns:
GeminiThinkingConfig with thinkingLevel and includeThoughts
"""
# Check if this is gemini-3-flash which supports MINIMAL thinking level
# Covers gemini-3-flash, gemini-3-flash-preview, gemini-3.1-flash, gemini-3.1-flash-lite-preview,
# gemini-3.5-flash, and any future 3.x-flash variants.
is_gemini3flash: Final = model and ("flash" in model.lower() and "gemini-3" in model.lower())
supports_minimal: Final = bool(model) and VertexGeminiConfig._supports_minimal_thinking_level(model)
is_gemini31pro: Final = model and ("gemini-3.1-pro-preview" in model.lower())
if reasoning_effort == "minimal":
if is_gemini3flash:
if supports_minimal:
return {"thinkingLevel": "minimal", "includeThoughts": True}
else:
return {"thinkingLevel": "low", "includeThoughts": True}
@ -899,18 +906,11 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
return {"thinkingLevel": "high", "includeThoughts": True}
elif reasoning_effort == "high":
return {"thinkingLevel": "high", "includeThoughts": True}
elif reasoning_effort == "disable":
# Gemini 3 cannot fully disable thinking, so we use "minimal" for gemini-3-flash-preview, "low" for others
if is_gemini3flash:
return {"thinkingLevel": "minimal", "includeThoughts": False}
else:
return {"thinkingLevel": "low", "includeThoughts": False}
elif reasoning_effort == "none":
# For gemini-3-flash-preview, use "minimal" instead of "low"
if is_gemini3flash:
return {"thinkingLevel": "minimal", "includeThoughts": False}
else:
return {"thinkingLevel": "low", "includeThoughts": False}
elif reasoning_effort in ("disable", "none"):
return {
"thinkingLevel": "minimal" if supports_minimal else "low",
"includeThoughts": False,
}
else:
raise _unsupported_reasoning_effort(reasoning_effort)
@ -977,8 +977,9 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
params["includeThoughts"] = True
# Follow provider defaults unless explicitly opted into legacy behavior.
if litellm.enable_gemini_default_thinking_level_low is True:
is_gemini3flash: Final = "gemini-3" in model.lower() and "flash" in model.lower()
params["thinkingLevel"] = "minimal" if is_gemini3flash else "low"
params["thinkingLevel"] = (
"minimal" if VertexGeminiConfig._supports_minimal_thinking_level(model) else "low"
)
else:
# Thinking disabled
params["includeThoughts"] = False

View file

@ -377,8 +377,7 @@
"supports_prompt_caching": true,
"supports_response_schema": true,
"supports_vision": true,
"supports_tool_choice": true,
"cache_read_input_token_cost": 1.5e-08
"supports_tool_choice": true
},
"amazon.nova-2-lite-v1:0": {
"cache_read_input_token_cost": 7.5e-08,
@ -561,8 +560,7 @@
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"cache_read_input_token_cost": 8.75e-09
"supports_tool_choice": true
},
"amazon.nova-pro-v1:0": {
"cache_read_input_token_cost": 2e-07,
@ -578,8 +576,7 @@
"supports_prompt_caching": true,
"supports_response_schema": true,
"supports_vision": true,
"supports_tool_choice": true,
"cache_read_input_token_cost": 2e-07
"supports_tool_choice": true
},
"amazon.nova-sonic-v1:0": {
"deprecation_date": "2026-09-14",
@ -26108,6 +26105,7 @@
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_minimal_reasoning_effort": false,
"supports_response_schema": true,
"supports_system_messages": true,
"supports_tool_choice": true,
@ -26165,6 +26163,7 @@
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_minimal_reasoning_effort": false,
"supports_response_schema": true,
"supports_system_messages": true,
"supports_tool_choice": true,
@ -28114,6 +28113,7 @@
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_minimal_reasoning_effort": false,
"supports_response_schema": true,
"supports_system_messages": true,
"supports_tool_choice": true,
@ -28173,6 +28173,7 @@
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_minimal_reasoning_effort": false,
"supports_response_schema": true,
"supports_system_messages": true,
"supports_tool_choice": true,
@ -28595,6 +28596,7 @@
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_minimal_reasoning_effort": false,
"supports_response_schema": true,
"supports_system_messages": true,
"supports_tool_choice": true,
@ -28652,6 +28654,7 @@
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_minimal_reasoning_effort": false,
"supports_response_schema": true,
"supports_system_messages": true,
"supports_tool_choice": true,
@ -42310,6 +42313,20 @@
"max_tokens": 128000,
"mode": "chat"
},
"openrouter/stealth/union-alpha": {
"input_cost_per_token": 0,
"output_cost_per_token": 0,
"litellm_provider": "openrouter",
"max_input_tokens": 262144,
"max_output_tokens": 131072,
"max_tokens": 131072,
"mode": "chat",
"source": "https://openrouter.ai/stealth/union-alpha",
"supports_function_calling": true,
"supports_tool_choice": true,
"supports_response_schema": true,
"supports_vision": true
},
"ovhcloud/DeepSeek-R1-Distill-Llama-70B": {
"input_cost_per_token": 6.7e-07,
"litellm_provider": "ovhcloud",
@ -45794,8 +45811,7 @@
"supports_prompt_caching": true,
"supports_response_schema": true,
"supports_vision": true,
"supports_tool_choice": true,
"cache_read_input_token_cost": 1.5e-08
"supports_tool_choice": true
},
"us.amazon.nova-micro-v1:0": {
"cache_read_input_token_cost": 8.75e-09,
@ -45809,8 +45825,7 @@
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"cache_read_input_token_cost": 8.75e-09
"supports_tool_choice": true
},
"us.amazon.nova-premier-v1:0": {
"deprecation_date": "2026-09-14",
@ -45842,8 +45857,7 @@
"supports_prompt_caching": true,
"supports_response_schema": true,
"supports_vision": true,
"supports_tool_choice": true,
"cache_read_input_token_cost": 2e-07
"supports_tool_choice": true
},
"us.anthropic.claude-3-5-haiku-20241022-v1:0": {
"cache_creation_input_token_cost": 1e-06,

View file

@ -102,6 +102,7 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials import (
UpstreamCredentialProvider,
)
from litellm.proxy._experimental.mcp_server.outbound_credentials.adapter import (
prepare_mcp_client,
raise_public,
raise_token_exchange_challenge,
raise_user_oauth_challenge,
@ -2804,6 +2805,8 @@ class MCPServerManager:
headers=headers,
server_label=server.name or server.server_name or server.alias or server.server_id,
relays_upstream_auth=server.is_client_forwarded_token,
auth_type=server.auth_type,
upstream_token_header=server.upstream_token_header,
)
tool_func.__name__ = prefixed_tool_name
tool_func.__doc__ = description
@ -4259,15 +4262,20 @@ class MCPServerManager:
user_api_key_auth=user_api_key_auth,
extra_headers=extra_headers,
)
return MCPClient(
server_url=server_url,
transport_type=transport,
auth_type=resolved_server.auth_type,
timeout=(resolved_server.timeout if resolved_server.timeout is not None else MCP_CLIENT_TIMEOUT),
extra_headers=extra_headers,
resolved_auth=resolved_auth,
sampling_callback=sampling_cb,
elicitation_callback=elicitation_cb,
return await prepare_mcp_client(
resolved_server,
MCPClient(
server_url=server_url,
transport_type=transport,
auth_type=resolved_server.auth_type,
timeout=(
resolved_server.timeout if resolved_server.timeout is not None else MCP_CLIENT_TIMEOUT
),
extra_headers=extra_headers,
resolved_auth=resolved_auth,
sampling_callback=sampling_cb,
elicitation_callback=elicitation_cb,
),
)
# Create SigV4 auth if configured
@ -4297,17 +4305,20 @@ class MCPServerManager:
else AuthResolution.no_auth
)
record_auth_resolution(server.server_id, legacy_source)
return MCPClient(
server_url=server_url,
transport_type=transport,
auth_type=resolved_server.auth_type,
auth_value=auth_value,
auth_header_name=auth_header_name,
timeout=(resolved_server.timeout if resolved_server.timeout is not None else MCP_CLIENT_TIMEOUT),
extra_headers=extra_headers,
aws_auth=aws_auth,
sampling_callback=sampling_cb,
elicitation_callback=elicitation_cb,
return await prepare_mcp_client(
resolved_server,
MCPClient(
server_url=server_url,
transport_type=transport,
auth_type=resolved_server.auth_type,
auth_value=auth_value,
auth_header_name=auth_header_name,
timeout=(resolved_server.timeout if resolved_server.timeout is not None else MCP_CLIENT_TIMEOUT),
extra_headers=extra_headers,
aws_auth=aws_auth,
sampling_callback=sampling_cb,
elicitation_callback=elicitation_cb,
),
)
async def _get_tools_from_server(

View file

@ -54,7 +54,7 @@ from litellm.llms.custom_httpx.http_handler import (
from litellm.proxy._experimental.mcp_server.tool_registry import (
global_mcp_tool_registry,
)
from litellm.types.mcp import credential_redirect_hook, custom_credential_slot
from litellm.types.mcp import MCPAuthType, credential_redirect_hook, custom_credential_slot
class _OpenAPIJSONSchema(TypedDict, total=False):
@ -471,6 +471,8 @@ def create_tool_function(
headers: dict[str, str] | None = None,
server_label: str | None = None,
relays_upstream_auth: bool = False,
auth_type: MCPAuthType = None,
upstream_token_header: str | None = None,
):
"""Create a tool function for an OpenAPI operation.
@ -503,6 +505,18 @@ def create_tool_function(
by using **kwargs instead of named parameters.
"""
effective_headers: Final = _merge_openapi_tool_request_headers(headers)
if auth_type is not None:
from litellm.proxy._experimental.mcp_server.outbound_credentials.adapter import (
raise_public,
validate_static_credential,
)
from litellm.proxy._experimental.mcp_server.outbound_credentials.result import Error, Ok
match validate_static_credential(auth_type, effective_headers, upstream_token_header, headers or ()):
case Error(error):
raise_public(error)
case Ok():
pass
# Build URL from base_url and path
url = base_url + path

View file

@ -13,15 +13,17 @@ from __future__ import annotations
import base64
import os
from collections.abc import Iterable, Mapping
from typing import TYPE_CHECKING, Final, Literal, NoReturn
from fastapi import HTTPException
from pydantic import SecretStr
from typing_extensions import assert_never
from litellm.experimental_mcp_client.client import strip_auth_scheme, to_basic_credentials
from litellm.experimental_mcp_client.client import MCPClient, strip_auth_scheme, to_basic_credentials
from litellm.proxy._experimental.mcp_server.exceptions import MCPServerURLCredentialsError
from litellm.proxy._experimental.mcp_server.oauth_utils import resolve_upstream_resource
from litellm.proxy._experimental.mcp_server.outbound_credentials.result import Error, Ok, Result
from litellm.proxy._experimental.mcp_server.outbound_credentials.types import (
DEFAULT_CREDENTIAL_HEADER,
ApiKeyConfig,
@ -39,7 +41,7 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.types import (
Subject,
TokenExchangeConfig,
)
from litellm.types.mcp import DEFAULT_SUBJECT_TOKEN_TYPE, MCPAuth
from litellm.types.mcp import DEFAULT_SUBJECT_TOKEN_TYPE, MCPAuth, MCPAuthType, MCPTransport
if TYPE_CHECKING:
from litellm.proxy._types import UserAPIKeyAuth
@ -79,7 +81,7 @@ def to_server_spec(server: MCPServer) -> ServerSpec | None:
BYOK is the per-user source of the ``api_key`` mode; its scheme rides on ``auth_type`` just
like a shared key, but the value is per-user and not migrated yet, so a BYOK server defers
to v1 regardless of ``auth_type`` (this guard is the seam the BYOK arm replaces later).
to v1 for its static schemes. Declared OBO always stays with the exchange arm.
Dispatches on the declared ``auth_type``. The match is exhaustive over ``MCPAuthType`` with
an ``assert_never`` tail, so a newly added auth mode fails the type gate here until it is
@ -90,8 +92,8 @@ def to_server_spec(server: MCPServer) -> ServerSpec | None:
modes ``true_passthrough`` / ``oauth_delegate`` (``PassthroughConfig``); delegated/passthrough
oauth2 and SigV4 return None and stay on v1.
"""
if server.is_byok:
return None # per-user BYOK source not migrated yet -> defer to v1 (any auth_type)
if server.is_byok and server.auth_type != MCPAuth.oauth2_token_exchange:
return None # per-user BYOK source not migrated yet -> defer to v1
resource: Final = server.url or server.server_id
auth_type: Final = server.auth_type
match auth_type:
@ -165,21 +167,9 @@ def _client_credentials_spec(server: MCPServer, resource: str) -> ServerSpec:
)
def _token_exchange_spec(server: MCPServer, resource: str) -> ServerSpec | None:
"""Build a token_exchange (OBO) spec, or defer (None) when it is not OBO-configured.
An OBO server with ``client_id``/``client_secret`` is owned by the v2 arm even if the
``token_exchange_endpoint``/``token_url`` is absent: a missing endpoint then fails closed (412) at
the exchanger rather than silently deferring to v1 and connecting unauthenticated, since the
gateway must not guess the IdP or fall back to a weaker source. Without client credentials there is
nothing to own, so the server stays on v1 (parity-safe). ``profile`` selects the wire dialect
(``rfc8693`` default, ``entra_obo`` for Microsoft Entra On-Behalf-Of); an unrecognized value
normalizes to ``rfc8693`` so a bad config value cannot crash spec-building. ``audience`` is
forwarded only when the operator set it; a missing one is omitted, not derived.
"""
def _token_exchange_spec(server: MCPServer, resource: str) -> ServerSpec:
"""Keep declared OBO owned by the resolver, including incomplete client configuration."""
endpoint: Final = server.token_exchange_endpoint or server.effective_token_url
if not server.client_id or not server.client_secret:
return None
profile: Final[Literal["rfc8693", "entra_obo"]] = (
"entra_obo" if server.token_exchange_profile == "entra_obo" else "rfc8693"
)
@ -193,7 +183,7 @@ def _token_exchange_spec(server: MCPServer, resource: str) -> ServerSpec | None:
token_exchange_endpoint=endpoint,
audience=server.audience,
client_id=server.client_id,
client_secret=SecretStr(server.client_secret),
client_secret=SecretStr(server.client_secret) if server.client_secret else None,
token_endpoint_auth_method=server.token_endpoint_auth_method,
scopes=tuple(server.scopes or ()),
),
@ -397,3 +387,74 @@ def raise_token_exchange_challenge(
detail="Unauthorized",
headers={"WWW-Authenticate": www_authenticate},
)
_STATIC_MODES: Final = frozenset(
(MCPAuth.api_key, MCPAuth.bearer_token, MCPAuth.basic, MCPAuth.token, MCPAuth.authorization)
)
def _usable_credential_value(auth_type: MCPAuthType, name: str, value: str) -> bool:
if not value:
return False
if auth_type == MCPAuth.api_key and name != "authorization":
return True
if value.lower() in ("bearer", "basic", "token", "apikey"):
return False
if auth_type == MCPAuth.api_key:
api_scheme: Final = value.split(None, 1)[0]
if api_scheme.lower() in ("bearer", "token", "apikey"):
api_credential: Final = strip_auth_scheme(value, api_scheme).strip()
return api_credential.lower() != api_scheme.lower()
if auth_type in (MCPAuth.bearer_token, MCPAuth.token):
scheme: Final = "Bearer" if auth_type == MCPAuth.bearer_token else "token"
credential: Final = strip_auth_scheme(value, scheme).strip()
return bool(credential) and credential.lower() != scheme.lower()
if auth_type == MCPAuth.basic:
parts: Final = value.split(None, 1)
if len(parts) != 2 or parts[0].lower() != "basic":
return False
try:
decoded: Final = base64.b64decode(parts[1], validate=True).strip()
return b":" in decoded
except ValueError:
return False
return True
def validate_static_credential(
auth_type: MCPAuthType,
headers: Mapping[str, str],
upstream_token_header: str | None = None,
static_header_names: Iterable[str] = (),
) -> Result[None, CredError]:
if auth_type not in _STATIC_MODES:
return Ok(None)
default_slot: Final = "X-API-Key" if auth_type == MCPAuth.api_key else "Authorization"
admin_chosen_slots: Final = tuple(static_header_names) if auth_type == MCPAuth.api_key else ()
slots: Final = frozenset(
name.lower()
for name in (
upstream_token_header or default_slot,
default_slot,
"Authorization",
*admin_chosen_slots,
)
)
values: Final = tuple((name.lower(), value.strip()) for name, value in headers.items() if name.lower() in slots)
if any(_usable_credential_value(auth_type, name, value) for name, value in values):
return Ok(None)
return Error(CredError.of_misconfigured(f"{auth_type} requires a usable upstream credential"))
async def prepare_mcp_client(server: MCPServer, client: MCPClient) -> MCPClient:
if server.auth_type not in _STATIC_MODES or client.transport_type == MCPTransport.stdio:
return client
request: Final = await client.prepare_request_auth()
match validate_static_credential(
server.auth_type, request.headers, server.upstream_token_header, server.static_headers or ()
):
case Error(error):
raise_public(error)
case Ok():
return client

View file

@ -844,6 +844,9 @@ class LiteLLMRoutes(enum.Enum):
)
self_managed_routes = [
# update_team resolves proxy/org/team admin itself and filters team admins
# through the team_admin_editable_team_fields setting
"/team/update",
"/team/member_add",
"/team/member_delete",
"/management/v1/teams/{team_id}/members/bulk_delete",
@ -4467,6 +4470,29 @@ class TeamInfoMember(Member):
user_alias: str | None = None
class TeamEditUnrestricted(BaseModel):
kind: Literal["unrestricted"] = "unrestricted"
class TeamEditAsTeamAdmin(BaseModel):
kind: Literal["team_admin"] = "team_admin"
editable_fields: tuple[str, ...]
class TeamEditAsTeamAdminDisabled(BaseModel):
kind: Literal["team_admin_disabled"] = "team_admin_disabled"
class TeamEditNone(BaseModel):
kind: Literal["none"] = "none"
TeamEditAccess = Annotated[
TeamEditUnrestricted | TeamEditAsTeamAdmin | TeamEditAsTeamAdminDisabled | TeamEditNone,
Field(discriminator="kind"),
]
class TeamInfoResponseObjectTeamTable(LiteLLM_TeamTable):
members_with_roles: tuple[TeamInfoMember, ...] = ()
team_member_budget_table: LiteLLM_BudgetTableFull | None = None
@ -4479,6 +4505,7 @@ class TeamInfoResponseObjectTeamTable(LiteLLM_TeamTable):
# None = no org or not a manager; [] or ["all-proxy-models"] = no ceiling.
organization_models: list[str] | None = None
model_max_budget_usage: Mapping[str, Mapping[str, object]] | None = None
caller_edit_access: TeamEditAccess = Field(default_factory=TeamEditNone)
class TeamInfoResponseObject(TypedDict):

View file

@ -8,7 +8,6 @@ from fastapi import APIRouter, Depends, HTTPException, Request, Response
from fastapi.responses import JSONResponse
import litellm
from litellm._logging import verbose_proxy_logger
from litellm.anthropic_interface.exceptions import AnthropicErrorResponse, AnthropicExceptionMapping
from litellm.integrations.custom_guardrail import ModifyResponseException
from litellm.llms.anthropic.experimental_pass_through.context_management import (
@ -22,13 +21,16 @@ from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.common_request_processing import (
ProxyBaseLLMRequestProcessing,
create_response,
log_llm_api_exception,
proxy_exception_from_http_exception,
resolve_litellm_call_id,
)
from litellm.proxy.common_utils.http_parsing_utils import _read_request_body
from litellm.proxy.common_utils.openai_error_payload import (
error_status_code,
openai_error_param,
openai_error_type,
with_litellm_call_id,
)
from litellm.types.utils import TokenCountResponse
@ -218,10 +220,12 @@ async def anthropic_response(
await proxy_logging_obj.post_call_failure_hook(
user_api_key_dict=user_api_key_dict, original_exception=e, request_data=base_llm_response_processor.data
)
verbose_proxy_logger.exception("litellm.proxy.proxy_server.anthropic_response(): Exception occured - %s", e)
log_llm_api_exception(e, base_llm_response_processor.litellm_call_id)
if isinstance(e, ProxyException):
return _anthropic_error_json_response(e, request)
return _anthropic_error_json_response(
with_litellm_call_id(e, base_llm_response_processor.litellm_call_id), request
)
# Extract model_id from request metadata (same as success path)
litellm_metadata: Final = data.get("litellm_metadata", {}) or {}
@ -231,7 +235,7 @@ async def anthropic_response(
# Get headers
headers: Final = ProxyBaseLLMRequestProcessing.get_custom_headers(
user_api_key_dict=user_api_key_dict,
call_id=data.get("litellm_call_id", ""),
call_id=base_llm_response_processor.litellm_call_id,
model_id=model_id,
version=version,
response_cost=0,
@ -288,6 +292,7 @@ async def count_tokens(
"""
from litellm.proxy.proxy_server import token_counter as internal_token_counter
litellm_call_id: Final = resolve_litellm_call_id(request.headers.get("x-litellm-call-id"))
try:
request_data: Final = await _read_request_body(request=request)
data: Final[dict] = {**request_data}
@ -339,7 +344,7 @@ async def count_tokens(
detail=detail,
)
except Exception as e:
verbose_proxy_logger.exception("litellm.proxy.anthropic_endpoints.count_tokens(): Exception occurred - %s", e)
log_llm_api_exception(e, litellm_call_id)
raise HTTPException(status_code=500, detail={"error": f"Internal server error: {e}"})

View file

@ -856,6 +856,16 @@ BUDGET_ENFORCED_SIDE_EFFECT_ROUTES: Final = frozenset(
)
def route_skips_budget_checks(route: str) -> bool:
return route not in BUDGET_ENFORCED_SIDE_EFFECT_ROUTES and (
route in MODEL_DISCOVERY_ROUTES or not RouteChecks.is_llm_api_route(route=route)
)
def request_skips_budget_checks(route: str, model: str | list[str] | None, llm_router: Router | None) -> bool:
return route_skips_budget_checks(route=route) or _is_model_cost_zero(model=model, llm_router=llm_router)
async def common_checks(
request_body: dict,
team_object: LiteLLM_TeamTable | None,
@ -903,10 +913,7 @@ async def common_checks(
team_id=valid_token.team_id if valid_token is not None else None,
)
skip_all_budget_checks: Final = skip_budget_checks or (
route not in BUDGET_ENFORCED_SIDE_EFFECT_ROUTES
and (route in MODEL_DISCOVERY_ROUTES or not RouteChecks.is_llm_api_route(route=route))
)
skip_all_budget_checks: Final = skip_budget_checks or route_skips_budget_checks(route=route)
membership_user_id: Final = (
valid_token.user_id if valid_token is not None and (bool(_model) or not skip_all_budget_checks) else None
@ -2104,7 +2111,7 @@ async def _fetch_uncached_tags(
@log_db_metrics
async def get_tag_objects_batch(
tag_names: list[str],
tag_names: Sequence[str],
prisma_client: PrismaClient | None,
user_api_key_cache: UserApiKeyCache,
parent_otel_span: Span | None = None,
@ -5863,15 +5870,25 @@ async def _tag_max_budget_check(
"""
from litellm.proxy.common_utils.http_parsing_utils import get_tags_from_request_body
if prisma_client is None:
await tag_max_budget_check_for_tags(
tags=get_tags_from_request_body(request_body=request_body),
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
valid_token=valid_token,
)
async def tag_max_budget_check_for_tags(
tags: Sequence[str],
prisma_client: PrismaClient | None,
user_api_key_cache: UserApiKeyCache,
proxy_logging_obj: ProxyLogging,
valid_token: UserAPIKeyAuth | None,
) -> None:
if prisma_client is None or not tags:
return
# Get tags from request metadata
tags: Final = get_tags_from_request_body(request_body=request_body)
if not tags:
return
# Batch fetch all tags in one go
tag_objects: Final = await get_tag_objects_batch(
tag_names=tags,
prisma_client=prisma_client,

View file

@ -7,6 +7,7 @@
import asyncio
import os
from collections.abc import Mapping
from types import MappingProxyType
from typing import Any, Final, cast
from fastapi import APIRouter, Depends, HTTPException, Path, Request, Response
@ -17,7 +18,11 @@ from litellm.batches.main import CancelBatchRequest, RetrieveBatchRequest
from litellm.proxy._types import *
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.batches_endpoints.common_utils import validate_batch_list_limit
from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
from litellm.proxy.common_request_processing import (
ProxyBaseLLMRequestProcessing,
log_llm_api_exception,
request_litellm_call_id,
)
from litellm.proxy.common_utils.callback_utils import sanitize_openai_provider_metadata
from litellm.proxy.common_utils.http_parsing_utils import _read_request_body
from litellm.proxy.common_utils.openai_endpoint_utils import (
@ -383,8 +388,9 @@ async def create_batch(
await proxy_logging_obj.post_call_failure_hook(
user_api_key_dict=user_api_key_dict, original_exception=e, request_data=data
)
verbose_proxy_logger.exception("litellm.proxy.proxy_server.create_batch(): Exception occured - %s", e)
raise handle_exception_on_proxy(e)
litellm_call_id: Final = request_litellm_call_id(data)
log_llm_api_exception(e, litellm_call_id)
raise handle_exception_on_proxy(e, litellm_call_id)
@router.get(
@ -674,8 +680,9 @@ async def retrieve_batch(
await proxy_logging_obj.post_call_failure_hook(
user_api_key_dict=user_api_key_dict, original_exception=e, request_data=data
)
verbose_proxy_logger.exception("litellm.proxy.proxy_server.retrieve_batch(): Exception occured - %s", e)
raise handle_exception_on_proxy(e)
litellm_call_id: Final = request_litellm_call_id(data)
log_llm_api_exception(e, litellm_call_id)
raise handle_exception_on_proxy(e, litellm_call_id)
@router.get(
@ -725,6 +732,7 @@ async def list_batches(
)
verbose_proxy_logger.debug("GET /v1/batches after=%s limit=%s", after, limit)
data: Mapping[str, object] = MappingProxyType({})
try:
if llm_router is None:
raise HTTPException(
@ -854,10 +862,11 @@ async def list_batches(
await proxy_logging_obj.post_call_failure_hook(
user_api_key_dict=user_api_key_dict,
original_exception=e,
request_data={"after": after, "limit": limit},
request_data={**data, "after": after, "limit": limit},
)
verbose_proxy_logger.error("litellm.proxy.proxy_server.retrieve_batch(): Exception occured - %s", e)
raise handle_exception_on_proxy(e)
litellm_call_id: Final = request_litellm_call_id(data)
log_llm_api_exception(e, litellm_call_id)
raise handle_exception_on_proxy(e, litellm_call_id)
@router.post(
@ -1079,8 +1088,9 @@ async def cancel_batch(
await proxy_logging_obj.post_call_failure_hook(
user_api_key_dict=user_api_key_dict, original_exception=e, request_data=data
)
verbose_proxy_logger.exception("litellm.proxy.proxy_server.create_batch(): Exception occured - %s", e)
raise handle_exception_on_proxy(e)
litellm_call_id: Final = request_litellm_call_id(data)
log_llm_api_exception(e, litellm_call_id)
raise handle_exception_on_proxy(e, litellm_call_id)
######################################################################

View file

@ -7,14 +7,25 @@ from collections.abc import AsyncGenerator, Awaitable, Callable, Mapping, Sequen
from datetime import datetime
from functools import lru_cache
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, Literal, NamedTuple, Protocol, TypeAlias, TypeVar, overload
from typing import (
TYPE_CHECKING,
Any,
Final,
Literal,
NamedTuple,
Protocol,
TypeAlias,
TypeVar,
overload,
runtime_checkable,
)
import anyio
import httpx
import orjson
from fastapi import HTTPException, Request, status
from fastapi.responses import JSONResponse, Response, StreamingResponse
from pydantic import ValidationError
from pydantic import TypeAdapter, ValidationError
from starlette.types import Receive, Scope, Send
import litellm
@ -34,7 +45,11 @@ from litellm.constants import (
UNSAFE_PROXY_RESPONSE_HEADERS,
)
from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.litellm_core_utils.core_helpers import get_or_create_metadata_bucket, is_expected_client_error
from litellm.litellm_core_utils.core_helpers import (
get_or_create_metadata_bucket,
independent_snapshot,
is_expected_client_error,
)
from litellm.litellm_core_utils.dd_tracing import NullTracer, tracer
from litellm.litellm_core_utils.get_supported_openai_params import (
get_supported_openai_params,
@ -49,14 +64,21 @@ from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
from litellm.litellm_core_utils.streaming_handler import (
backfill_missing_cache_usage_fields,
)
from litellm.proxy._types import ProxyException, UserAPIKeyAuth
from litellm.proxy.auth.auth_checks import can_key_call_resolved_model
from litellm.proxy.auth.auth_utils import check_response_size_is_safe
from litellm.proxy._types import ProxyErrorTypes, ProxyException, UserAPIKeyAuth
from litellm.proxy.auth.auth_checks import (
can_key_call_resolved_model,
request_skips_budget_checks,
tag_max_budget_check_for_tags,
)
from litellm.proxy.auth.auth_utils import check_response_size_is_safe, get_request_route
from litellm.proxy.common_utils.callback_utils import (
get_logging_caching_headers,
get_remaining_tokens_and_requests_from_request_data,
)
from litellm.proxy.common_utils.http_parsing_utils import get_client_requested_model
from litellm.proxy.common_utils.http_parsing_utils import (
get_client_requested_model,
get_tags_from_request_body,
)
from litellm.proxy.common_utils.openai_error_payload import (
attribute_of,
error_status_code,
@ -643,6 +665,48 @@ async def _resolve_per_request_model_group_alias(
return target
_REQUEST_MODEL: Final[TypeAdapter[str | list[str] | None]] = TypeAdapter(str | list[str] | None)
def _request_model(data: Mapping[str, object]) -> str | list[str] | None:
try:
return _REQUEST_MODEL.validate_python(data.get("model"), strict=True)
except ValidationError:
return None
async def _enforce_guardrail_added_tag_budgets(
data: Mapping[str, object],
tags_before_guardrails: frozenset[str],
route: str,
llm_router: Router | None,
user_api_key_dict: UserAPIKeyAuth,
proxy_logging_obj: ProxyLogging,
) -> None:
added_tags: Final = tuple(
tag for tag in get_tags_from_request_body(request_body=data) if tag not in tags_before_guardrails
)
if not added_tags or request_skips_budget_checks(route=route, model=_request_model(data), llm_router=llm_router):
return
from litellm.proxy.proxy_server import prisma_client, user_api_key_cache
try:
await tag_max_budget_check_for_tags(
tags=added_tags,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
valid_token=user_api_key_dict,
)
except litellm.BudgetExceededError as e:
raise ProxyException(
message=e.message,
type=ProxyErrorTypes.budget_exceeded,
param=None,
code=e.status_code,
) from e
async def _parse_event_data_for_error(event_line: str | bytes) -> int | None:
"""Parses an event line and returns an error code if present, else None."""
event_line = event_line.decode("utf-8") if isinstance(event_line, bytes) else event_line
@ -1452,7 +1516,19 @@ def _has_attribute_error_in_chain(exc: Exception) -> bool:
_CLIENT_DISCONNECT_DETAIL: Final = "Client disconnected the request"
def _log_llm_api_exception(e: Exception, litellm_call_id: str | None) -> None:
@runtime_checkable
class _CarriesLitellmCallId(Protocol):
litellm_call_id: str | None
def request_litellm_call_id(data: Mapping[str, object]) -> str | None:
logging_obj: Final = data.get("litellm_logging_obj")
logged_id: Final = logging_obj.litellm_call_id if isinstance(logging_obj, _CarriesLitellmCallId) else None
call_id: Final = logged_id or data.get("litellm_call_id")
return call_id if isinstance(call_id, str) else None
def log_llm_api_exception(e: Exception, litellm_call_id: str | None) -> None:
if getattr(e, "status_code", None) == 499 and getattr(e, "detail", None) == _CLIENT_DISCONNECT_DETAIL:
verbose_proxy_logger.info(
"litellm.proxy.proxy_server._handle_llm_api_exception(): client disconnected, "
@ -1531,6 +1607,11 @@ def _timing_values(
class ProxyBaseLLMRequestProcessing:
def __init__(self, data: dict):
self.data = data
self._tags_before_guardrails: frozenset[str] | None = None
@property
def litellm_call_id(self) -> str | None:
return request_litellm_call_id(self.data)
@staticmethod
def _merge_passthrough_streaming_headers(
@ -2020,11 +2101,21 @@ class ProxyBaseLLMRequestProcessing:
# to run below.
await _arm_auto_router_compression(data=self.data, llm_router=llm_router)
if self._tags_before_guardrails is None:
self._tags_before_guardrails = frozenset(get_tags_from_request_body(request_body=self.data))
self.data = await proxy_logging_obj.pre_call_hook(
user_api_key_dict=user_api_key_dict,
data=self.data,
call_type=route_type,
)
await _enforce_guardrail_added_tag_budgets(
data=self.data,
tags_before_guardrails=self._tags_before_guardrails,
route=get_request_route(request=request),
llm_router=llm_router,
user_api_key_dict=user_api_key_dict,
proxy_logging_obj=proxy_logging_obj,
)
if route_type == "aget_responses":
attach_post_call_pipelines_to_retrieval(
data=self.data,
@ -2062,6 +2153,13 @@ class ProxyBaseLLMRequestProcessing:
) -> tuple[dict, LiteLLMLoggingObj]:
from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError
configured_fallbacks: Final = (
self._configured_fallbacks(llm_router=llm_router, user_api_key_dict=user_api_key_dict)
if llm_router is not None and not self.data.get("disable_fallbacks")
else None
)
pristine: Final = independent_snapshot(self.data) if configured_fallbacks else None
try:
return await self.common_processing_pre_call_logic(
request=request,
@ -2080,14 +2178,19 @@ class ProxyBaseLLMRequestProcessing:
llm_router=llm_router,
)
except ProxyRateLimitError as original_exc:
original_model: Final = self.data.get("model")
if not original_model or not llm_router or self.data.get("disable_fallbacks"):
rate_limited_data: Final = self.data
original_model: Final = rate_limited_data.get("model")
if (
pristine is None
or not configured_fallbacks
or rate_limited_data.get("disable_fallbacks")
or not isinstance(original_model, str)
):
raise
fallback_models: Final = self._resolve_fallback_models(
model=original_model,
llm_router=llm_router,
user_api_key_dict=user_api_key_dict,
fallbacks=configured_fallbacks,
)
if not fallback_models:
raise
@ -2102,6 +2205,7 @@ class ProxyBaseLLMRequestProcessing:
for fallback_model in fallback_models:
if fallback_model == original_model:
continue
self.data = independent_snapshot(pristine)
self.data["model"] = fallback_model
try:
return await self.common_processing_pre_call_logic(
@ -2123,39 +2227,30 @@ class ProxyBaseLLMRequestProcessing:
except ProxyRateLimitError:
continue
except BaseException:
self.data["model"] = original_model
self.data = rate_limited_data
raise
self.data["model"] = original_model
self.data = rate_limited_data
raise original_exc
def _resolve_fallback_models(
self,
model: str,
llm_router: Router,
user_api_key_dict: UserAPIKeyAuth,
) -> list | None:
from litellm.router_utils.fallback_event_handlers import get_fallback_model_group
fallbacks = None
@staticmethod
def _configured_fallbacks(llm_router: Router, user_api_key_dict: UserAPIKeyAuth) -> list | None:
key_router_settings: Final = user_api_key_dict.router_settings
if isinstance(key_router_settings, dict) and "fallbacks" in key_router_settings:
fallbacks = key_router_settings["fallbacks"]
key_fallbacks: Final = key_router_settings.get("fallbacks") if isinstance(key_router_settings, dict) else None
fallbacks: Final = key_fallbacks if key_fallbacks is not None else llm_router.fallbacks
return fallbacks if isinstance(fallbacks, list) and fallbacks else None
if fallbacks is None:
fallbacks = llm_router.fallbacks
if not fallbacks:
return None
@staticmethod
def _resolve_fallback_models(model: str, fallbacks: list) -> list | None:
from litellm.router_utils.fallback_event_handlers import get_fallback_model_group
fallback_model_group, generic_fallback_idx = get_fallback_model_group(
fallbacks=fallbacks,
model_group=model,
)
if fallback_model_group is None and generic_fallback_idx is not None:
fallback_model_group = fallbacks[generic_fallback_idx]["*"]
return fallback_model_group
if fallback_model_group is not None:
return fallback_model_group
return fallbacks[generic_fallback_idx]["*"] if generic_fallback_idx is not None else None
@staticmethod
def _get_model_id_from_response(hidden_params: Mapping[str, object], data: Mapping[str, object]) -> str:
@ -3429,11 +3524,7 @@ class ProxyBaseLLMRequestProcessing:
version: str | None = None,
):
"""Raises ProxyException (OpenAI API compatible) if an exception is raised"""
logging_obj: Final[LiteLLMLoggingObj | None] = self.data.get("litellm_logging_obj", None)
_log_llm_api_exception(
e,
(logging_obj.litellm_call_id if logging_obj is not None else None) or self.data.get("litellm_call_id"),
)
log_llm_api_exception(e, self.litellm_call_id)
# Allow callbacks to transform the error response
transformed_exception: Final = await proxy_logging_obj.post_call_failure_hook(
user_api_key_dict=user_api_key_dict,
@ -3463,9 +3554,7 @@ class ProxyBaseLLMRequestProcessing:
custom_headers: Final = ProxyBaseLLMRequestProcessing.get_custom_headers(
user_api_key_dict=user_api_key_dict,
call_id=(
_litellm_logging_obj.litellm_call_id if _litellm_logging_obj else self.data.get("litellm_call_id")
),
call_id=self.litellm_call_id,
model_id=model_id,
version=version,
response_cost=0,

View file

@ -9,6 +9,9 @@ from typing import Final
from fastapi import status
from litellm.constants import STRINGIFIED_NONE
from litellm.proxy._types import ProxyException
LITELLM_CALL_ID_HEADER: Final = "x-litellm-call-id"
_OPENAI_ERROR_TYPE_BY_STATUS: Final[Mapping[int, str]] = MappingProxyType(
{
@ -52,3 +55,23 @@ def openai_error_param(exc: object) -> str | None:
serializes as JSON ``null``."""
carried: Final = attribute_of(exc, "param")
return carried if isinstance(carried, str) and carried != STRINGIFIED_NONE else None
def litellm_call_id_headers(litellm_call_id: str | None) -> dict[str, str] | None: # mutable-ok: ProxyException.headers
if litellm_call_id is None:
return None
return {LITELLM_CALL_ID_HEADER: litellm_call_id} # mutable-ok: ProxyException mutates its headers dict
def with_litellm_call_id(exc: ProxyException, litellm_call_id: str | None) -> ProxyException:
"""The same error object, answering with ``x-litellm-call-id`` when it was raised without one."""
if litellm_call_id is not None:
exc.headers.setdefault(LITELLM_CALL_ID_HEADER, litellm_call_id)
return exc
def headers_with_litellm_call_id(headers: Mapping[str, str] | None, litellm_call_id: str) -> Mapping[str, str]:
"""``headers`` plus ``x-litellm-call-id``, keeping the value they already carry under that name."""
if headers is None:
return MappingProxyType({LITELLM_CALL_ID_HEADER: litellm_call_id})
return MappingProxyType({LITELLM_CALL_ID_HEADER: litellm_call_id, **headers})

View file

@ -11,7 +11,7 @@ exception types:
an upstream LLM provider returns 429.
* :class:`fastapi.HTTPException` (status 429) raised directly by proxy hooks
such as ``parallel_request_limiter``, ``dynamic_rate_limiter``,
``batch_rate_limiter``, ``max_budget_limiter``, ``max_iterations_limiter``,
``batch_rate_limiter``, ``max_iterations_limiter``,
etc.
* :class:`litellm.llms.base_llm.chat.transformation.BaseLLMException` (status
429) raised by some provider transports.

View file

@ -257,8 +257,8 @@ class DBSpendUpdateWriter:
# Completion object fields
kwargs: dict | None,
completion_response: object,
start_time: datetime | None,
end_time: datetime | None,
start_time: datetime,
end_time: datetime,
response_cost: float | None,
) -> bool:
"""Record the request's spend, answering whether its cost still needs charging.
@ -299,6 +299,7 @@ class DBSpendUpdateWriter:
response_obj=completion_response,
start_time=start_time,
end_time=end_time,
llm_router=get_llm_router(),
)
payload["spend"] = response_cost or 0.0
if isinstance(payload["startTime"], datetime):

View file

@ -232,7 +232,8 @@ class AktoGuardrail(CustomGuardrail):
"""
request_path: Final = self.extract_request_path(request_data)
request_headers: Final = self.build_request_headers(request_data)
request_body: Final = self.build_request_body(inputs, request_data)
request_inputs: Final = GenericGuardrailAPIInputs(model=inputs.get("model")) if include_response else inputs
request_body: Final = self.build_request_body(request_inputs, request_data)
tag: Final = self.build_tag_metadata(request_data)
response_payload = json.dumps({}) # Empty body wrapper when no response yet

View file

@ -425,10 +425,7 @@ class CrowdStrikeAIDRHandler(CustomGuardrail):
def _build_guard_input_for_response(self, inputs: GenericGuardrailAPIInputs) -> _GuardInput:
output_texts: Final[list[str]] = inputs.get("texts", [])
return _GuardInput(
messages=[_Message(role="assistant", content=text) for text in output_texts],
tools=inputs.get("tools", []),
)
return _GuardInput(messages=[_Message(role="assistant", content=text) for text in output_texts], tools=[])
def _extract_transformed_texts(self, guard_output: _GuardInput, num_assistant_messages: int) -> list[str]:
tail: Final = guard_output.messages[-num_assistant_messages:] if num_assistant_messages > 0 else []

View file

@ -286,7 +286,7 @@ class HiddenlayerGuardrail(CustomGuardrail):
hl_request_metadata["requester_id"] = headers.get("hl-requester-id") or "LiteLLM"
project_id: Final = headers.get("hl-project-id")
if scan_params := inputs.get("structured_messages"):
if input_type == "request" and (scan_params := inputs.get("structured_messages")):
last_msg: Final = scan_params[-1]
result: _HiddenlayerResponse = await self._call_hiddenlayer(
project_id,

View file

@ -197,7 +197,7 @@ class OpenAIModerationGuardrail(OpenAIGuardrailBase, CustomGuardrail):
text_to_moderate: str | None = None
# Prefer structured_messages if available (has role context)
if structured_messages := inputs.get("structured_messages"):
if input_type == "request" and (structured_messages := inputs.get("structured_messages")):
text_to_moderate = self.get_user_prompt(structured_messages)
# Fall back to texts

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,
streaming_transform_mode=getattr(litellm_params, "streaming_transform_mode", None),
file_sanitization_fail_open=getattr(litellm_params, "file_sanitization_fail_open", None),
block_on_file_modify=getattr(litellm_params, "block_on_file_modify", None),
)

View file

@ -38,6 +38,11 @@ class PromptSecurityGuardrailMissingSecrets(Exception):
pass
def _modified_or_original(text: str, verdict: "_ProtectVerdict") -> str:
modified_text: Final = verdict.get("modified_text") if verdict.get("action") == "modify" else None
return text if modified_text is None else modified_text
def _inputs_with_structured_messages(
inputs: GenericGuardrailAPIInputs, rewritten_messages: Sequence[AllMessageValues] | None
) -> GenericGuardrailAPIInputs:
@ -119,6 +124,7 @@ class PromptSecurityGuardrail(CustomGuardrail):
user: str | None = None,
system_prompt: str | None = None,
check_tool_results: bool | None = None,
streaming_transform_mode: Literal["block_only", "incremental_diff"] | None = None,
file_sanitization_timeout: float = _SANITIZE_FILE_FAIL_OPEN_TIMEOUT_SECONDS,
file_sanitization_fail_open: bool | None = None,
block_on_file_modify: bool | None = None,
@ -148,6 +154,10 @@ class PromptSecurityGuardrail(CustomGuardrail):
)
raise PromptSecurityGuardrailMissingSecrets(msg)
self.streaming_transform_mode: Literal["block_only", "incremental_diff"] = (
"block_only" if streaming_transform_mode is None else streaming_transform_mode
)
# Configuration for file sanitization
self.max_poll_attempts = 30 # Maximum number of polling attempts
self.poll_interval = 2 # Seconds between polling attempts
@ -342,16 +352,46 @@ class PromptSecurityGuardrail(CustomGuardrail):
texts: list[str],
user_api_key_alias: str | None,
) -> GenericGuardrailAPIInputs:
"""Handle response-side guardrail checks."""
"""Handle response-side guardrail checks, one protect verdict per text.
Prompt Security rewrites a single string, so texts from several choices must be scanned separately
or one ``modified_text`` cannot be mapped back onto the choice it came from. It also returns no span
offsets, so on a stream every text is held back in full until the final verdict: a value the vendor
redacts later may start anywhere in text that looked clean so far, and streamed bytes cannot be recalled.
"""
if not texts:
return inputs
# Combine all texts for response checking
combined_text: Final = "\n".join(texts)
verdicts: Final = await asyncio.gather(
*(self._protect_response_text(text, user_api_key_alias) for text in texts)
)
violations: Final = tuple(
violation
for verdict in verdicts
if verdict.get("action") == "block"
for violation in verdict.get("violations", ())
)
if any(verdict.get("action") == "block" for verdict in verdicts):
raise HTTPException(
status_code=400,
detail="Blocked by Prompt Security, Violations: " + ", ".join(violations),
)
returned_texts: Final = [ # mutable-ok: GenericGuardrailAPIInputs.texts is list[str]
_modified_or_original(text, verdict) for text, verdict in zip(texts, verdicts, strict=True)
]
patched: Final[GenericGuardrailAPIInputs] = {
**inputs,
"texts": returned_texts,
"stream_holdback_chars": [ # mutable-ok: GenericGuardrailAPIInputs.stream_holdback_chars is list[int]
len(text) for text in returned_texts
],
}
return patched
async def _protect_response_text(self, text: str, user_api_key_alias: str | None) -> _ProtectVerdict:
headers: Final = self._build_headers(user_api_key_alias)
payload: Final = {
"response": combined_text,
"response": text,
"user": user_api_key_alias or self.user,
"system_prompt": self.system_prompt,
}
@ -360,7 +400,7 @@ class PromptSecurityGuardrail(CustomGuardrail):
method="POST",
url=f"{self.api_base}/api/protect",
headers=headers,
payload={"response_length": len(combined_text)},
payload={"response_length": len(text)},
)
response: Final = await self.async_handler.post(
@ -377,26 +417,8 @@ class PromptSecurityGuardrail(CustomGuardrail):
payload={"result": res.get("result")},
)
result: Final = res.get("result", {}).get("response", {})
if result is None:
return inputs
action: Final = result.get("action")
violations: Final = result.get("violations", [])
if action == "block":
raise HTTPException(
status_code=400,
detail="Blocked by Prompt Security, Violations: " + ", ".join(violations),
)
elif action == "modify":
modified_text: Final = result.get("modified_text")
if modified_text is not None:
# If we combined multiple texts, return the modified version as single text
# The framework will handle distributing it back
inputs["texts"] = [modified_text]
return inputs
verdict: Final = res.get("result", {}).get("response", {})
return {} if verdict is None else verdict
def _extract_texts_from_messages(self, messages: Sequence[Mapping[str, object]]) -> list[str]:
return [text for message in messages for text in message_slot_texts(message)]

View file

@ -121,7 +121,7 @@ class PromptGuardGuardrail(CustomGuardrail):
) -> GenericGuardrailAPIInputs:
texts: Final = inputs.get("texts", [])
images: Final = inputs.get("images", [])
structured_messages: Final = inputs.get("structured_messages", [])
structured_messages: Final = inputs.get("structured_messages") if input_type == "request" else None
model: Final = inputs.get("model")
if structured_messages:

View file

@ -452,7 +452,7 @@ class QualifireGuardrail(CustomGuardrail):
dynamic_params: Final = self.get_guardrail_dynamic_request_body_params(request_data=request_data)
# Extract messages from structured_messages or request_data
messages: list[AllMessageValues] | None = inputs.get("structured_messages")
messages: list[AllMessageValues] | None = inputs.get("structured_messages") if input_type == "request" else None
if not messages:
messages = request_data.get("messages")

View file

@ -380,11 +380,12 @@ class StraikerGuardrail(CustomGuardrail):
call_id: Final = getattr(logging_obj, "litellm_call_id", None) if logging_obj else None
event_id: Final = f"{call_id or 'litellm'}:{input_type}"
is_request: Final = input_type == "request"
content: Final = StraikerWebhookContent(
texts=list(inputs.get("texts") or []),
images=list(inputs.get("images") or []),
structured_messages=_opaque_dict_list(inputs.get("structured_messages")),
tools=_opaque_dict_list(inputs.get("tools")),
structured_messages=_opaque_dict_list(inputs.get("structured_messages")) if is_request else None,
tools=_opaque_dict_list(inputs.get("tools")) if is_request else None,
tool_calls=_opaque_dict_list(inputs.get("tool_calls")),
)

View file

@ -104,6 +104,10 @@ def _chunk_choices(item: object) -> Sequence[object]:
return choices
def _held_choices(held_chars_per_choice: Mapping[int, int]) -> frozenset[int]:
return frozenset(idx for idx, held in held_chars_per_choice.items() if held > 0)
def _is_redundant_scan(scan_key: "StreamingScanKey | None", last_scan_key: "StreamingScanKey | None") -> bool:
if scan_key is None:
return False
@ -472,6 +476,7 @@ class UnifiedLLMGuardrails(CustomLogger):
emitted_text_per_choice: dict[int, str],
holdback_per_choice: dict[int, int],
finish_reason_per_choice: dict[int, str | None],
held_chars_per_choice: dict[int, int],
is_final: bool,
) -> ModelResponseStream | None:
"""Build the synthetic chunk carrying the newly-guardrailed deltas.
@ -479,7 +484,9 @@ class UnifiedLLMGuardrails(CustomLogger):
For each choice, the new delta is the mutated accumulated text past what
has already been emitted, minus a trailing holdback (forced to 0 on the
final flush). ``emitted_text_per_choice`` holds the exact bytes already
sent per choice and is extended in place. Returns None when there is no
sent per choice and is extended in place; ``held_chars_per_choice`` is
updated in place with how many mutated chars per choice are still withheld
after this round. Returns None when there is no
text to emit (e.g. a tool-call-only turn) or nothing new and this is not
the final chunk.
@ -536,6 +543,7 @@ class UnifiedLLMGuardrails(CustomLogger):
holdback = 0 if is_final else max(0, holdback_per_choice.get(choice_idx, 0))
end = max(len(already), len(text) - holdback)
deltas[choice_idx] = text[len(already) : end]
held_chars_per_choice[choice_idx] = len(text) - end
# Iterate the mutated choices (not just those in reference_chunk) so a
# choice with pending text is never dropped for n > 1. finish_reason is
@ -590,6 +598,7 @@ class UnifiedLLMGuardrails(CustomLogger):
responses_yielded: list[object],
emitted_text_per_choice: dict[int, str],
finish_reason_per_choice: dict[int, str | None],
held_chars_per_choice: dict[int, int],
is_final: bool,
) -> AsyncGenerator[object, None]:
"""Run one guardrail processing round and emit the resulting diff chunk.
@ -618,6 +627,7 @@ class UnifiedLLMGuardrails(CustomLogger):
emitted_text_per_choice=emitted_text_per_choice,
holdback_per_choice=sink.holdback_per_choice,
finish_reason_per_choice=finish_reason_per_choice,
held_chars_per_choice=held_chars_per_choice,
is_final=is_final,
)
except ModifyResponseException as e:
@ -673,6 +683,7 @@ class UnifiedLLMGuardrails(CustomLogger):
responses_yielded: Final[list[object]] = []
emitted_text_per_choice: Final[dict[int, str]] = {}
finish_reason_per_choice: Final[dict[int, str | None]] = {}
held_chars_per_choice: Final[dict[int, int]] = {}
chunk_counter = 0
last_chunk: object | None = None
@ -688,6 +699,7 @@ class UnifiedLLMGuardrails(CustomLogger):
responses_yielded=responses_yielded,
emitted_text_per_choice=emitted_text_per_choice,
finish_reason_per_choice=finish_reason_per_choice,
held_chars_per_choice=held_chars_per_choice,
is_final=is_final,
)
@ -724,12 +736,18 @@ class UnifiedLLMGuardrails(CustomLogger):
# finish_reason to the final text terminator (see the
# _tool_call_passthrough_chunk docstring).
tool_only = self._tool_call_passthrough_chunk(
item, finish_reason_per_choice=finish_reason_per_choice
item,
finish_reason_per_choice=finish_reason_per_choice,
held_choices=_held_choices(held_chars_per_choice),
)
responses_yielded.append(tool_only)
yield tool_only
continue
if self._is_trailing_metadata_chunk(item):
responses_so_far.append(item)
continue
chunk_counter += 1
responses_so_far.append(item)
last_chunk = item
@ -773,12 +791,33 @@ class UnifiedLLMGuardrails(CustomLogger):
):
yield out
if last_chunk is not None:
async for out in _round(last_chunk, is_final=True):
yield out
async for out in self._emit_stream_tail(
last_chunk=last_chunk,
final_round=_round,
responses_so_far=responses_so_far,
responses_yielded=responses_yielded,
):
yield out
except _StreamTerminated:
return
async def _emit_stream_tail(
self,
*,
last_chunk: object | None,
final_round: Callable[[object, bool], AsyncGenerator[object, None]],
responses_so_far: Sequence[object],
responses_yielded: list[object],
) -> AsyncGenerator[object, None]:
"""Flush the held text with holdback 0, then replay metadata-only chunks
(usage) so they land after the text and its finish_reason, as upstream sent them."""
if last_chunk is not None:
async for out in final_round(last_chunk, True):
yield out
for trailing in self._trailing_metadata_chunks(responses_so_far):
responses_yielded.append(trailing)
yield trailing
async def _inspect_full_response_for_block(
self,
*,
@ -829,6 +868,23 @@ class UnifiedLLMGuardrails(CustomLogger):
return True
return False
@classmethod
def _is_trailing_metadata_chunk(cls, item: object) -> bool:
"""True for a chunk that carries only stream metadata (no choices, or a
``usage`` chunk whose deltas are empty); such chunks are replayed after
the final text flush instead of being folded into the transform."""
if not _chunk_choices(item):
return True
return (
getattr(item, "usage", None) is not None
and not cls._chunk_carries_text(item)
and not cls._chunk_has_finish_reason(item)
)
@classmethod
def _trailing_metadata_chunks(cls, items: Sequence[object]) -> tuple[object, ...]:
return tuple(item for item in items if cls._is_trailing_metadata_chunk(item))
@staticmethod
def _chunk_carries_text(item: object) -> bool:
"""True if any choice in this chunk has non-empty string ``delta.content``."""
@ -843,6 +899,7 @@ class UnifiedLLMGuardrails(CustomLogger):
def _tool_call_passthrough_chunk(
item: object,
finish_reason_per_choice: "dict[int, str | None] | None" = None,
held_choices: frozenset[int] = frozenset(),
) -> ModelResponseStream:
"""Copy of a chunk carrying tool calls with all text content stripped.
@ -851,8 +908,9 @@ class UnifiedLLMGuardrails(CustomLogger):
transform instead). Applies per choice so an n>1 chunk mixing a text
choice and a tool-call choice does not leak the text choice.
For a choice that carries BOTH text content AND tool_calls, ``finish_reason``
is suppressed on the passthrough and recorded on
For a choice that carries BOTH text content AND tool_calls, or whose earlier
text is still withheld (``held_choices``), ``finish_reason`` is suppressed on
the passthrough and recorded on
``finish_reason_per_choice`` (when provided) so the final synthetic text
chunk delivers it. Emitting the passthrough's ``finish_reason`` before the
text flush would let a spec-compliant SSE client stop reading at
@ -865,7 +923,8 @@ class UnifiedLLMGuardrails(CustomLogger):
idx = getattr(choice, "index", 0) or 0
original_finish = getattr(choice, "finish_reason", None)
has_text = isinstance(getattr(delta, "content", None), str) and getattr(delta, "content", "") != ""
if has_text and original_finish is not None and finish_reason_per_choice is not None:
text_pending = has_text or idx in held_choices
if text_pending and original_finish is not None and finish_reason_per_choice is not None:
finish_reason_per_choice[idx] = original_finish
passthrough_finish: str | None = None
else:

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