scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-24 00:52:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 90

Merge remote-tracking branch 'origin/main' into litellm_team_audit_lifecycle

Browse source
This commit is contained in:
yucheng 2026-09-19 22:41:02 +00:00
commit 2ade97f80e
687 changed files with 70165 additions and 16503 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

8
.circleci/config.yml
View file

@ -1785,6 +1785,12 @@ jobs:
- wait_for_service:
url: http://localhost:4000
timeout: "300"
- run:
name: Seed the routing strategy through /config/update
command: |
curl --noproxy '*' -sSf -X POST http://localhost:4000/config/update \
-H 'Authorization: Bearer sk-1234' -H 'Content-Type: application/json' \
-d '{"router_settings": {"routing_strategy": "usage-based-routing-v2"}}'
- run:
name: Run tests
command: |
@ -3009,7 +3015,7 @@ workflows:
name: integration-<< matrix.suite >>
matrix:
parameters:
suite: [management, accounting, database, providers, extensions, sdk, browser]
suite: [management, accounting, database, providers, extensions, sdk, cost, browser]
filters:
branches:
only:

11
.circleci/scripts/classify_changes.sh
View file

@ -1,16 +1,22 @@
#!/usr/bin/env bash
set -uo pipefail
category="${1:?usage: classify_changes.sh <backend|client|ui|provider-harness|cost-map-only>}"
category="${1:?usage: classify_changes.sh <backend|client|ui|provider-harness|cost-map-only|mcp-dependencies>}"
has_client=false
has_backend=false
has_ci=false
has_provider_harness=false
has_cost_map=false
has_mcp_dependencies=false
outside_cost_map_set=false
while IFS= read -r file || [ -n "$file" ]; do
[ -n "$file" ] || continue
case "$file" in
*.md | *.mdx) : ;;
pyproject.toml | */pyproject.toml | uv.lock | uv.toml | .python-version | rust-toolchain.toml | litellm-rust/* | litellm/__init__.py | litellm/proxy/proxy_server.py | litellm/*mcp* | tests/*mcp* | litellm/integrations/arize/* | tests/base_sdk_tests/* | scripts/check_mcp_sdk_install.py | .github/workflows/test-mcp-dependency-resolution.yml | .github/actions/detect-changes/* | .github/actions/setup-uv-with-retries/* | .github/actions/cache-cargo-build/* | .github/scripts/detect_changes.sh | .github/scripts/uv_sync_with_retries.sh | .circleci/scripts/classify_changes.sh | tests/test_litellm/test_circleci_path_filter.py | tests/test_litellm/test_detect_changes.py)
has_mcp_dependencies=true ;;
esac
case "$file" in
tests/e2e/*/*.py) : ;;
tests/e2e/*.py | tests/code_coverage_tests/test_provider_cache.py | tests/code_coverage_tests/test_provider_replay_harness.py | tests/test_litellm/test_circleci_path_filter.py | .circleci/* | pyproject.toml | uv.lock)
@ -31,6 +37,9 @@ while IFS= read -r file || [ -n "$file" ]; do
done
case "$category" in
mcp-dependencies)
[ "$has_mcp_dependencies" = true ] && echo run || echo skip
;;
cost-map-only)
{ [ "$has_cost_map" = true ] && [ "$outside_cost_map_set" = false ]; } && echo run || echo skip
;;

22
.circleci/scripts/run_integration.sh
View file

@ -9,6 +9,7 @@ fi
suite="${1:?integration suite required}"
results="test-results/integration-${suite}"
mkdir -p "$results"
shard_timeout=11m
integration_identity="$(.venv/bin/python -c 'import uuid; print(uuid.uuid4().hex)')"
upstream_pid=""
proxy_pid=""
@ -108,13 +109,26 @@ awk '$3 == "REJECT" && $1 > 0 { rejected=1 } END { exit !rejected }' "$results/e
setsid env -i PATH="$PATH" HOME="$HOME" PYTHONPATH="$PYTHONPATH" INTEGRATION_RUN_ID="$integration_identity" \
.venv/bin/python -m integration._support.upstream > "$results/upstream.log" 2>&1 &
upstream_pid=$!
if [ "$suite" = cost ]; then
export INTEGRATION_WORKERS=8
fi
start_proxy() {
local port="$1"
local log_name="$2"
local -a cost_map_env
if [ "$suite" = cost ]; then
cost_map_env=(
"LITELLM_MODEL_COST_MAP_URL=$INTEGRATION_UPSTREAM_URL/_cost_map"
"MODEL_COST_MAP_MIN_MODEL_COUNT=1"
"MODEL_COST_MAP_MAX_SHRINK_RATIO=0"
)
else
cost_map_env=("LITELLM_LOCAL_MODEL_COST_MAP=True")
fi
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_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 \
LITELLM_MODE=PRODUCTION STORE_MODEL_IN_DB=True "${cost_map_env[@]}" \
AWS_EC2_METADATA_DISABLED=true DO_NOT_TRACK=1 \
.venv/bin/python -m integration._support.proxy --config tests/integration/proxy_config.yaml \
--host 127.0.0.1 --port "$port" --num_workers 1 --telemetry False \
@ -125,6 +139,9 @@ start_proxy() {
start_proxy 4000 proxy.log
proxy_pid="$launched_pid"
.venv/bin/python .circleci/scripts/wait_integration_services.py
curl --noproxy '*' -sSf -X POST "$INTEGRATION_PROXY_URL/config/update" \
-H "Authorization: Bearer $LITELLM_MASTER_KEY" -H 'Content-Type: application/json' \
-d '{"router_settings": {"num_retries": 0}}' > "$results/seed-router-settings.json"
if [ "$suite" = management ]; then
export INTEGRATION_PEER_URL=http://127.0.0.1:4001
start_proxy 4001 peer.log
@ -155,11 +172,12 @@ if [ "$suite" = browser ]; then
exit 0
fi
timeout --signal=TERM --kill-after=20s 11m env -i PATH="$PATH" HOME="$HOME" PYTHONPATH="$PYTHONPATH" \
timeout --signal=TERM --kill-after=20s "$shard_timeout" 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" \
INTEGRATION_PROXY_URL="$INTEGRATION_PROXY_URL" INTEGRATION_PEER_URL="$INTEGRATION_PEER_URL" \
INTEGRATION_UPSTREAM_URL="$INTEGRATION_UPSTREAM_URL" \
INTEGRATION_WORKERS="${INTEGRATION_WORKERS:-1}" \
INTEGRATION_MASTER_KEY="$INTEGRATION_MASTER_KEY" LITELLM_MODE=PRODUCTION \
INTEGRATION_SEED="$INTEGRATION_SEED" \
INTEGRATION_ORDER_SEED="$INTEGRATION_ORDER_SEED" \

2
.github/actions/detect-changes/action.yml vendored
View file

@ -14,7 +14,7 @@ description: >-
inputs:
category:
description: "Which classification to apply: backend, client or ui"
description: "Which classification to apply: backend, client, ui, provider-harness, cost-map-only or mcp-dependencies"
required: false
default: backend
github-token:

50
.github/scripts/_agent_shin_actions.py vendored
View file

@ -1,50 +0,0 @@
"""Dry-run wrapper(s) around Agent Shin GitHub mutations.
The rollout scripts currently need only one mutation wrapped, so this module
exposes a single ``maybe_post_comment`` helper. It takes a ``dry_run: bool``
keyword argument and the body is intentionally trivial:
if dry_run:
print(...) # log what we would do, return
return
real_mutation(...) # otherwise, actually do it
That shape means a dry-run preview differs from the real run in exactly one
line per side effect: the call site. So when you `python3 script.py` locally
without ``--close``, you can be confident the actions printed are the ones the
GitHub Action would have performed (modulo ordering on retry/error paths,
which are deliberately simple). Any further mutation a rollout script needs
should get the same ``maybe_*`` treatment instead of calling the raw
``triage_with_llm`` mutation directly.
Importing from this module pulls in the real mutation from ``triage_with_llm``
— call sites in the rollout scripts should NEVER import ``post_comment``
directly; that would skip the dry-run gate and is the bug class this module
exists to prevent.
"""
from __future__ import annotations
import sys
import textwrap
# Import the module itself rather than the bare names so monkeypatching
# `triage_with_llm.post_comment` (or any of the other mutations) in tests is
# reflected here — `from triage_with_llm import post_comment` would bind the
# original function to a local name and bypass the patch, defeating the whole
# point of these wrappers.
import triage_with_llm
def _log(line: str) -> None:
"""Print a single dry-run line to stdout (one log statement per side effect)."""
print(line, file=sys.stdout, flush=True)
def maybe_post_comment(repo: str, number: int, body: str, *, dry_run: bool) -> None:
"""Post a comment on ``repo#number`` — or, in dry-run, log what we would post."""
if dry_run:
_log(f"[DRY RUN] comment {repo}#{number}:")
_log(textwrap.indent(body, " "))
return
triage_with_llm.post_comment(repo, number, body)

211
.github/scripts/agent_shin_shared.py vendored
View file

@ -1,211 +0,0 @@
"""Constants and helpers shared by Agent Shin's triage scripts.
Both `triage_with_llm.py` (the LLM-judge entrypoint) and
`close_low_quality_prs.py` (the daily Greptile-score sweep) need to
agree on the same notions of:
* What counts as a Greptile-authored review comment
(``GREPTILE_BOT_LOGINS``) and how to extract a confidence score from
its body (``SCORE_PATTERN`` / :func:`extract_greptile_score`).
* How long the 2-hour grace window is (``GRACE_PERIOD_SECONDS``) and
the HTML marker stamped into a grace-warning comment so the *other*
script can see "Agent Shin already warned" and behave accordingly
(``GRACE_COMMENT_MARKER``).
* Who Agent Shin is on GitHub (``AGENT_SHIN_DEFAULT_BOT_LOGIN``).
* How GitHub-style ISO-8601 timestamps round-trip into timezone-aware
:class:`datetime.datetime` (:func:`parse_iso8601`).
Keeping these in one module means a future change (new Greptile output
format, a longer grace window, a new allowlisted account) is a single edit
instead of two — the original split version had to call out in comments
that the two copies "must stay in sync" precisely because nothing
enforced it.
"""
from __future__ import annotations
import datetime as dt
import json
import os
import re
import subprocess
from typing import Iterable
GREPTILE_BOT_LOGINS = frozenset({"greptile-apps", "greptile-apps[bot]"})
SCORE_PATTERN = re.compile(
r"confidence\s*score\s*[:\-]?\s*(\d+)\s*/\s*5",
re.IGNORECASE,
)
GRACE_COMMENT_MARKER = "<!-- agent-shin:grace-warning -->"
# Hidden HTML marker stamped on every Agent Shin auto-close comment (the LLM
# judge's grace/review-gate close and the daily Greptile sweep's close).
# `was_closed_by_agent_shin` requires this marker — not just the closing actor —
# before `@agent-shin reconsider` may reopen, because the `github-actions[bot]`
# identity is shared with every other workflow in the repo and is not unique to
# Agent Shin. Both close paths must stamp it or the reconsider path silently
# rejects the contributor.
AGENT_SHIN_CLOSE_MARKER = "<!-- agent-shin:closed -->"
# 2 hours between the grace warning and the auto-close. Short enough to
# dogfood the "fix it before it closes" loop in one sitting; bump back up
# (e.g. 86400 for a day) for the public rollout.
GRACE_PERIOD_SECONDS = 7200
AGENT_SHIN_DEFAULT_BOT_LOGIN = "github-actions[bot]"
def _logins(*names: str) -> frozenset[str]:
"""Build a login set normalized for case-insensitive membership checks.
Callers compare via ``login.lower() in <set>``, so the stored values
must be lowercase. Normalizing here lets the literals keep each
account's canonical GitHub casing (e.g. ``SwiftWinds``) for
readability without breaking the lookup.
"""
return frozenset(name.lower() for name in names)
# Dogfood rollout gate. While this set is non-empty, Agent Shin acts ONLY on
# PRs/issues authored by these logins and skips everyone else. For an
# allowlisted author the usual internal/external classification is bypassed, so
# an internal account (e.g. a maintainer's own work login) still gets triaged
# while the bot is being tested on a small set of accounts. Empty the set to
# lift the restriction and restore full triage for the public rollout. Logins
# are compared case-insensitively.
ALLOWLIST_LOGINS = _logins("mateo-berri", "SwiftWinds")
# `gh {pr,issue} list` has no "fetch everything" flag — `--limit` is the only
# control and it defaults to 30. Pass a ceiling far above any realistic open
# backlog (low thousands today) so gh paginates the API until the queue is
# exhausted rather than silently truncating. The bulk sweeps MUST see the whole
# backlog: gh lists newest-first, so a low cap drops the *oldest* PRs/issues —
# exactly the stale ones a low-quality sweep is meant to catch.
GH_LIST_ALL_LIMIT = 100_000
def extract_greptile_score(comments: Iterable[dict]) -> tuple[int, dict] | None:
"""Return (score, comment) for the most recent Greptile-authored comment
that contains a "Confidence Score: X/5". Returns None if no such comment.
"Most recent" is determined by the comment's `updated_at` (falling back to
`created_at`), so re-reviews override earlier passes.
"""
candidates: list[tuple[str, int, dict]] = []
for comment in comments:
user = (comment.get("user") or {}).get("login", "")
if user not in GREPTILE_BOT_LOGINS:
continue
body = comment.get("body") or ""
match = SCORE_PATTERN.search(body)
if not match:
continue
score = int(match.group(1))
timestamp = comment.get("updated_at") or comment.get("created_at") or ""
candidates.append((timestamp, score, comment))
if not candidates:
return None
candidates.sort(key=lambda triple: triple[0])
_, score, comment = candidates[-1]
return score, comment
def parse_iso8601(value: str) -> dt.datetime:
"""Parse a GitHub ISO-8601 timestamp into a timezone-aware datetime."""
return dt.datetime.fromisoformat(value.replace("Z", "+00:00"))
def gh(*args: str) -> str:
"""Run a `gh` CLI command and return stdout. Raises on non-zero exit.
Shared by both Agent Shin entrypoints so a future change here
(timeout handling, logging, retry on transient failures) only needs
to be made once.
"""
result = subprocess.run(
["gh", *args],
capture_output=True,
text=True,
check=True,
)
return result.stdout
def list_open_items(kind: str, *, repo: str | None, fields: str) -> list[dict]:
"""Return EVERY open PR (``kind="pr"``) or issue (``kind="issue"``) in ``repo``.
Wraps ``gh {pr,issue} list`` with ``--limit GH_LIST_ALL_LIMIT`` so the full
backlog is fetched instead of the default 30 (or any other arbitrary cap).
Both bulk sweeps — the daily Greptile closer and the one-shot rollout
heads-up — rely on this seeing the whole queue, including the oldest items.
``fields`` is the comma-separated ``--json`` field list the caller needs
(e.g. ``"number"`` for the rollout, the full set for the closer).
"""
if kind not in ("pr", "issue"):
raise ValueError(f"kind must be 'pr' or 'issue', got {kind!r}")
repo_args = ["--repo", repo] if repo else []
raw = gh(
kind,
"list",
"--state",
"open",
"--limit",
str(GH_LIST_ALL_LIMIT),
"--json",
fields,
*repo_args,
)
return json.loads(raw)
def seconds_since_latest_marker_comment(
comments: Iterable[dict],
*,
marker: str,
bot_login: str | None = None,
now: dt.datetime | None = None,
) -> float | None:
"""Return seconds since the bot's most recent comment containing ``marker``.
Filters comments by author so a contributor who quotes the HTML
marker (e.g. via GitHub's "Quote reply" feature, which preserves
HTML comments in the raw markdown of the quoted text) is not
mistaken for a bot warning — that would silently reset cooldown
timers and suppress legitimate notifications.
``bot_login`` defaults to the `AGENT_SHIN_BOT_LOGIN` env override or
``AGENT_SHIN_DEFAULT_BOT_LOGIN`` so callers normally don't need to
pass it. ``now`` is injectable for tests / callers (like the daily
sweep) that want every age calculation pinned to one snapshot.
"""
expected_login = (
bot_login
or os.environ.get("AGENT_SHIN_BOT_LOGIN")
or AGENT_SHIN_DEFAULT_BOT_LOGIN
).lower()
latest: dt.datetime | None = None
for comment in comments:
author = ((comment.get("user") or {}).get("login") or "").lower()
if author != expected_login:
continue
body = comment.get("body") or ""
if marker not in body:
continue
created = comment.get("created_at")
if not created:
continue
try:
ts = parse_iso8601(created)
except ValueError:
continue
if latest is None or ts > latest:
latest = ts
if latest is None:
return None
reference = now if now is not None else dt.datetime.now(dt.timezone.utc)
return (reference - latest).total_seconds()

393
.github/scripts/auto_merge_price_sync.py vendored
View file

@ -1,393 +0,0 @@
"""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, 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 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"})
@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 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, ...]
reviews: tuple[Review, ...]
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}")
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 _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 _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),
reviews=_reviews(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())

573
.github/scripts/close_low_quality_prs.py vendored
View file

@ -1,573 +0,0 @@
#!/usr/bin/env python3
"""
Auto-close low-quality pull requests.
Closes open PRs (including drafts, regardless of age) that satisfy ALL of:
1. Have a Greptile (`greptile-apps`) review comment whose latest
"Confidence Score: X/5" is below the configured threshold (default: 4).
2. Are authored by an external OSS contributor (internal BerriAI
contributors are exempt).
3. Do not carry an opt-out label (default: "do not close").
`--min-age-days` is retained as an opt-in safety net for one-off backfill
runs (default: 0). The team's intent is that the count of open PRs equals
the count of PRs internal collaborators need to action on, so neither age
nor draft status acts as a free pass.
For each match, the script posts an explanatory comment and closes the PR.
Because OSS contributors *cannot* reopen a PR closed by the bot/maintainer
(GitHub limitation), the close-comment instructs them to push their fixes
and **open a fresh PR**, or to comment `@agent-shin reconsider` on the
closed PR to have the LLM judge re-evaluate (and reopen on pass).
Requires the `gh` CLI to be authenticated.
Usage examples:
# Dry run (default) - prints what would be closed
python3 close_low_quality_prs.py
# Actually close matching PRs
python3 close_low_quality_prs.py --close
# Restrict to PRs at least N days old (one-off backfill safety net)
python3 close_low_quality_prs.py --min-age-days 7 --min-score 4 --close
"""
from __future__ import annotations
import argparse
import datetime as dt
import json
import os
import subprocess
import sys
from typing import Iterable
# Add this script's directory to `sys.path` so the sibling
# `agent_shin_shared` module is importable when the script is invoked
# directly (e.g. `python3 .github/scripts/close_low_quality_prs.py ...`).
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from agent_shin_shared import ( # noqa: E402 -- sys.path adjusted above
AGENT_SHIN_CLOSE_MARKER,
ALLOWLIST_LOGINS,
GRACE_COMMENT_MARKER,
GRACE_PERIOD_SECONDS,
GREPTILE_BOT_LOGINS,
SCORE_PATTERN,
extract_greptile_score,
gh,
list_open_items,
parse_iso8601,
seconds_since_latest_marker_comment,
)
# `GREPTILE_BOT_LOGINS` and `SCORE_PATTERN` (Greptile's GitHub App login
# variants and the "Confidence Score: X/5" regex) are imported from
# `agent_shin_shared` so the LLM judge in `triage_with_llm.py` and this
# daily Greptile sweep read the score through the same set of logins
# and the same regex.
# `author_association` values for internal BerriAI contributors who should be
# exempt from auto-triage.
INTERNAL_AUTHOR_ASSOCIATIONS = frozenset({"OWNER", "MEMBER", "COLLABORATOR"})
# Default labels that exempt a PR from auto-close. Defined at module scope (not
# as a mutable argparse default) so that `--optout-label foo` REPLACES the
# defaults instead of appending to them — the argparse `action="append"` +
# `default=[...]` combination silently mutates the shared default list.
DEFAULT_OPTOUT_LABELS = ("do not close", "keep open", "wip")
# `GRACE_COMMENT_MARKER` (HTML marker appended to grace-period warning
# comments — used by either script to recognize that a warning was
# already posted) and `GRACE_PERIOD_SECONDS` (length of the grace
# period between the warning and the actual auto-close, 2 hours) are
# imported from `agent_shin_shared` so the Agent Shin LLM judge and
# this daily Greptile sweep agree on the same marker and duration.
def fetch_open_prs(repo: str | None) -> list[dict]:
"""Fetch all open PRs (number, createdAt, isDraft, labels, author).
Includes drafts: `gh pr list --state open` returns both ready-for-review
and draft PRs by default. This is the desired behavior — drafts are not
a free pass; the internal-collaborator open-PR queue should reflect every
PR that needs human attention regardless of draft status.
"""
fields = "number,title,createdAt,isDraft,labels,author,url"
return list_open_items("pr", repo=repo, fields=fields)
def fetch_pr_author_association(pr_number: int, repo: str | None) -> str:
"""Return the GitHub `author_association` for a PR, uppercase.
Values: OWNER, MEMBER, COLLABORATOR, CONTRIBUTOR, FIRST_TIME_CONTRIBUTOR,
FIRST_TIMER, MANNEQUIN, NONE. Returns "" on lookup failure.
"""
endpoint = (
f"repos/{repo}/pulls/{pr_number}"
if repo
else f"repos/{{owner}}/{{repo}}/pulls/{pr_number}"
)
try:
data = json.loads(gh("api", endpoint))
except subprocess.CalledProcessError:
return ""
return (data.get("author_association") or "").upper()
def is_external_pr_author(pr: dict, repo: str | None) -> bool:
"""Return True if the PR author is an external OSS contributor.
Internal = `OWNER` / `MEMBER` / `COLLABORATOR` association, or a bot login.
"""
login = ((pr.get("author") or {}).get("login") or "").lower()
if login.endswith("[bot]") or login in {"dependabot", "github-actions"}:
return False
association = fetch_pr_author_association(pr["number"], repo)
# Fail-safe: if the API lookup failed (empty string), treat the author as
# internal so we don't auto-close their PR. Auto-close is destructive, so
# an unknown association should never make a PR eligible for closing.
if not association or association in INTERNAL_AUTHOR_ASSOCIATIONS:
return False
return True
def fetch_pr_comments(pr_number: int, repo: str | None) -> list[dict]:
"""Fetch issue-level comments on a PR (where Greptile posts its summary)."""
endpoint = (
f"repos/{repo}/issues/{pr_number}/comments?per_page=100"
if repo
else f"repos/{{owner}}/{{repo}}/issues/{pr_number}/comments?per_page=100"
)
raw = gh("api", "--paginate", endpoint)
comments: list[dict] = []
for line in raw.strip().splitlines():
line = line.strip()
if not line:
continue
try:
parsed = json.loads(line)
except json.JSONDecodeError:
# A malformed line should not blow up the whole sweep. Skip and
# carry on so the remaining PRs in this run still get evaluated.
continue
if isinstance(parsed, list):
comments.extend(parsed)
else:
comments.append(parsed)
return comments
def has_optout_label(pr: dict, optout_labels: set[str]) -> bool:
labels = {label.get("name", "").lower() for label in pr.get("labels", [])}
return bool(labels & {lbl.lower() for lbl in optout_labels})
def seconds_since_last_grace_warning(
comments: Iterable[dict],
*,
bot_login: str | None = None,
now: dt.datetime | None = None,
) -> float | None:
"""Return seconds since the bot's most recent grace-period warning, or
None if no such warning has ever been posted on this PR.
Thin wrapper over
`agent_shin_shared.seconds_since_latest_marker_comment` — the
centralized helper handles the bot-author filter, marker match,
timestamp parsing, and `now` injection. Keeping this wrapper
preserves the closer's "already-fetched comments + injectable now"
interface so callers (and tests) don't need to change.
"""
return seconds_since_latest_marker_comment(
comments,
marker=GRACE_COMMENT_MARKER,
bot_login=bot_login,
now=now,
)
def format_grace_warning_comment(score: int, threshold: int) -> str:
"""Comment posted on the FIRST low-Greptile-score detection — gives
the contributor a 2-hour grace window before the auto-close fires on
the next daily cron run.
Mirrors `format_grace_warning_pr_comment` in
`triage_with_llm.py` in spirit (2-hour grace + escape hatches), but
framed around Greptile's confidence score instead of the LLM judge's
rubric since the close trigger here is the Greptile signal.
"""
return (
"🚅 Hi, thanks for the PR! I'm **Agent Shin**, the automated triage bot for this "
"repository.\n"
"\n"
"Heads up: Greptile's most recent review scored this PR "
f"**{score}/5**, below our merge bar of **{threshold}/5**.\n"
"\n"
"If the score isn't lifted in the next **2 hours**, I'll auto-close this PR. That's "
"**not** us saying the change isn't worthwhile. We want the open-PR list to mirror "
"what a maintainer can act on *right now*, so contributors like you don't get lost in "
"a backlog. Take your time; everything below still works after the close.\n"
"\n"
"**During the grace period:** push fixes that address Greptile's feedback, then comment "
"`@greptileai` to request a fresh review. If "
f"the new score is **{threshold}/5 or higher**, the PR stays open and no further "
"action is needed on your side.\n"
"\n"
"**If the PR does get auto-closed in 2 hours, you still have an easy recovery path:**\n"
"\n"
"- Comment `@greptileai` to request a fresh review. **This still works even after "
f"the PR is closed**, and a score of {threshold}/5 or higher is one of the signals "
"that lifts the PR back into the review queue. A low Greptile score isn't a blocker.\n"
"- Comment `@agent-shin reconsider` after pushing fixes; I'll re-run the rubric and "
"reopen the PR if both gates (description rubric + Greptile score) now pass.\n"
"\n"
f"{GRACE_COMMENT_MARKER}"
)
def post_grace_warning(
pr: dict,
score: int,
threshold: int,
repo: str | None,
dry_run: bool,
) -> None:
"""Post the 2-hour grace-period warning comment on `pr`.
The warning carries `GRACE_COMMENT_MARKER` so subsequent runs can
detect that the contributor has already been told about the
pending close. Does NOT close the PR — the close happens on the
next eligible run after `GRACE_PERIOD_SECONDS` elapses (handled
by `close_pr`).
"""
pr_number = pr["number"]
repo_args = ["--repo", repo] if repo else []
if dry_run:
print(
f" [DRY RUN] Would post grace warning to PR #{pr_number} "
f"(greptile={score}/5): {pr['title']}"
)
return
comment_body = format_grace_warning_comment(score, threshold)
gh("pr", "comment", str(pr_number), "--body", comment_body, *repo_args)
print(f" Posted grace warning on PR #{pr_number} (greptile={score}/5)")
def format_close_comment(score: int, threshold: int) -> str:
"""Comment posted when a low-Greptile-score PR is auto-closed.
Carries `AGENT_SHIN_CLOSE_MARKER` so the `@agent-shin reconsider` path
(guarded by `was_closed_by_agent_shin`) recognizes this as an Agent Shin
close and is allowed to reopen the PR once it passes again; without the
marker that recovery path the comment advertises silently rejects the
contributor.
"""
score_sentence = (
f"Greptile's most recent review scored this PR **{score}/5**, below "
f"our merge bar of **{threshold}/5**, and the 2-hour grace period since "
"the warning has elapsed.\n\n"
)
return (
f"Closing as part of automated PR triage.\n\n"
f"{score_sentence}"
"We close low-confidence PRs aggressively to keep the review queue "
"manageable for maintainers and contributors alike. **This is not a "
"rejection of the idea.** To bring this back:\n\n"
"1. Push the fixes that address Greptile's feedback (continue using "
"your existing branch is fine).\n"
"2. **Open a new PR** with the updated branch. Greptile will review "
"it again, and if it scores "
f"**{threshold}/5 or higher** a maintainer will take another look.\n\n"
"_Why open a new PR instead of reopening this one?_ GitHub does not "
"let external contributors reopen a PR that was closed by a bot or "
"maintainer, so a fresh PR is the most reliable path forward. If you "
"would prefer this exact PR re-evaluated, comment "
"`@agent-shin reconsider` once you've pushed the fixes; Agent Shin "
"will re-run triage and reopen this PR if it now meets the bar. "
"You can also comment `@greptileai` to request a fresh Greptile "
"review; that works **even after the PR is closed**.\n\n"
"Thanks for contributing to LiteLLM. We know auto-closures can sting; "
"the goal is to keep the project healthy, not to dismiss your work."
f"\n\n{AGENT_SHIN_CLOSE_MARKER}"
)
def close_pr(
pr: dict,
score: int,
threshold: int,
age_days: int,
repo: str | None,
dry_run: bool,
label: str | None,
) -> None:
"""Post the explanatory comment and close the PR."""
pr_number = pr["number"]
repo_args = ["--repo", repo] if repo else []
if dry_run:
print(
f" [DRY RUN] Would close PR #{pr_number} "
f"(age={age_days}d, greptile={score}/5): {pr['title']}"
)
return
comment_body = format_close_comment(score, threshold)
gh("pr", "comment", str(pr_number), "--body", comment_body, *repo_args)
if label:
try:
gh("pr", "edit", str(pr_number), "--add-label", label, *repo_args)
except subprocess.CalledProcessError as exc:
stderr = (exc.stderr or "").strip()
print(f" warn: failed to add label '{label}' to #{pr_number}: {stderr}")
gh("pr", "close", str(pr_number), *repo_args)
print(f" Closed PR #{pr_number} (greptile={score}/5, age={age_days}d)")
def evaluate_pr(
pr: dict,
now: dt.datetime,
min_age_days: int,
min_score: int,
repo: str | None,
optout_labels: set[str],
allowlist: frozenset[str] = ALLOWLIST_LOGINS,
) -> tuple[str, int | None, int | None]:
"""Decide what to do with `pr` on this triage run.
Returns (action, score_or_none, age_days_or_none) where action is one of:
"skip-too-young", "skip-optout-label", "skip-not-allowlisted",
"skip-internal", "skip-no-greptile-score", "skip-score-ok",
"warn-grace", "skip-in-grace-period", or "close".
Drafts are NOT skipped — the goal is "open PR count == PRs internal
collaborators need to action on", and a draft that Greptile scored <4/5
is still in that queue. Authors can opt out via the `wip` label (see
`DEFAULT_OPTOUT_LABELS`) if they need to keep a long-lived draft open.
Grace-period semantics: the first time a PR fails the rubric, the
action is `warn-grace` — the caller should post a warning comment but
NOT close the PR. On a subsequent run, if the warning is still less
than `GRACE_PERIOD_SECONDS` old AND the PR still fails, the action is
`skip-in-grace-period`. Once the warning ages out and the rubric is
still failing, the action is `close`.
"""
if has_optout_label(pr, optout_labels):
return ("skip-optout-label", None, None)
created = parse_iso8601(pr["createdAt"])
age_days = (now - created).days
# `min_age_days` defaults to 0 (close as soon as Greptile scores low).
# Set a positive value via --min-age-days for one-off backfill runs that
# want to skip very-young PRs.
if min_age_days > 0 and age_days < min_age_days:
return ("skip-too-young", None, age_days)
# While the allowlist is active it is the sole author gate: only those
# logins are acted on and the external-only restriction is bypassed for
# them. Otherwise auto-close only external OSS contributors — internal
# contributors (BerriAI org members) handle their own backlog.
login = ((pr.get("author") or {}).get("login") or "").lower()
if allowlist:
if login not in allowlist:
return ("skip-not-allowlisted", None, age_days)
elif not is_external_pr_author(pr, repo):
return ("skip-internal", None, age_days)
comments = fetch_pr_comments(pr["number"], repo)
extraction = extract_greptile_score(comments)
if extraction is None:
return ("skip-no-greptile-score", None, age_days)
score, _ = extraction
if score >= min_score:
return ("skip-score-ok", score, age_days)
grace_age = seconds_since_last_grace_warning(comments, now=now)
if grace_age is None:
return ("warn-grace", score, age_days)
if grace_age < GRACE_PERIOD_SECONDS:
return ("skip-in-grace-period", score, age_days)
return ("close", score, age_days)
def main() -> int:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument(
"--repo",
type=str,
default=None,
help="Repository (owner/repo). Auto-detected if omitted.",
)
parser.add_argument(
"--min-age-days",
type=int,
default=0,
help=(
"Minimum age (in days) before a PR is eligible. Default 0 = "
"close as soon as Greptile flags it. Set a positive value for "
"one-off backfill runs that want to spare very-young PRs."
),
)
parser.add_argument(
"--min-score",
type=int,
default=4,
choices=range(1, 6),
help="Greptile score below which a PR is closed (default: 4 -> closes <4/5).",
)
parser.add_argument(
"--optout-label",
action="append",
default=None,
help=(
"Label(s) that exempt a PR from auto-close. Repeat to add more. "
"Case-insensitive. When omitted, defaults to "
f"{list(DEFAULT_OPTOUT_LABELS)!r}; passing this flag REPLACES the "
"defaults (argparse `append` with a mutable default would append "
"instead, which we explicitly avoid)."
),
)
parser.add_argument(
"--close-label",
type=str,
default=None,
help=(
"Optional label to add to PRs that get auto-closed "
"(e.g. 'auto-closed-low-quality'). Must already exist on the repo."
),
)
parser.add_argument(
"--close",
action="store_true",
help="Actually close matching PRs (default is dry-run).",
)
parser.add_argument(
"--limit",
type=int,
default=None,
help="Maximum number of PRs to close in one run (safety net).",
)
args = parser.parse_args()
dry_run = not args.close
if dry_run:
print("=== DRY RUN MODE (pass --close to actually close PRs) ===\n")
print("Fetching open PRs...")
prs = fetch_open_prs(args.repo)
print(f"Found {len(prs)} open PRs.\n")
now = dt.datetime.now(dt.timezone.utc)
optout_labels = set(args.optout_label or DEFAULT_OPTOUT_LABELS)
closed = 0
summary = {
"close": 0,
"warn-grace": 0,
"skip-in-grace-period": 0,
"skip-too-young": 0,
"skip-optout-label": 0,
"skip-not-allowlisted": 0,
"skip-internal": 0,
"skip-no-greptile-score": 0,
"skip-score-ok": 0,
}
# `warned` tracks grace-warning comments posted in this run so the
# `--limit` safety net bounds *all* destructive write actions, not
# just closures. Without this cap, a backlog of PRs failing the
# threshold simultaneously could flood contributors with comments.
warned = 0
for pr in sorted(prs, key=lambda p: p["createdAt"]):
try:
action, score, age_days = evaluate_pr(
pr,
now,
args.min_age_days,
args.min_score,
args.repo,
optout_labels,
)
summary[action] = summary.get(action, 0) + 1
if action == "warn-grace":
assert score is not None
print(
f"#{pr['number']}: \"{pr['title']}\" "
f"(age={age_days}d, greptile={score}/5) -> warn-grace"
)
post_grace_warning(
pr,
score=score,
threshold=args.min_score,
repo=args.repo,
dry_run=dry_run,
)
if not dry_run:
warned += 1
if args.limit is not None and (warned + closed) >= args.limit:
print(
f"\nReached --limit={args.limit} "
f"(closed={closed}, warned={warned}); stopping."
)
break
continue
if action != "close":
continue
assert score is not None and age_days is not None
print(
f"#{pr['number']}: \"{pr['title']}\" "
f"(age={age_days}d, greptile={score}/5) -> close"
)
close_pr(
pr,
score=score,
threshold=args.min_score,
age_days=age_days,
repo=args.repo,
dry_run=dry_run,
label=args.close_label,
)
if not dry_run:
closed += 1
if args.limit is not None and (warned + closed) >= args.limit:
print(
f"\nReached --limit={args.limit} "
f"(closed={closed}, warned={warned}); stopping."
)
break
except Exception as exc: # noqa: BLE001 - per-PR errors don't abort the sweep
summary["error"] = summary.get("error", 0) + 1
print(
f"!! PR #{pr.get('number')}: {exc}",
file=sys.stderr,
)
continue
print("\n=== Summary ===")
for key, value in summary.items():
print(f" {key:28s} {value}")
if dry_run:
print(f"\nTotal would close: {summary['close']}")
else:
print(f"\nTotal closed: {closed}")
print(
f"Total {'would warn (grace)' if dry_run else 'warned (grace)'}: "
f"{summary['warn-grace']}"
)
return 0
if __name__ == "__main__":
sys.exit(main())

282
.github/scripts/triage-requirements.txt vendored
View file

@ -1,282 +0,0 @@
# Hash-pinned dependency set for the Agent Shin triage scripts.
# Installed in privileged triage workflows, so every package is pinned to an
# exact version with SHA-256 hashes and installed with pip --require-hashes.
#
# Regenerate after bumping openai:
# echo 'openai==<version>' \
# | uv pip compile - --generate-hashes --python-version 3.12 \
# --no-annotate --no-header -o .github/scripts/triage-requirements.txt
annotated-types==0.7.0 \
--hash=sha256:1f02e8b43a8fbbc3f3e0d4f0f4bfc8131bcb4eebe8849b8e5c773f3a1c582a53 \
--hash=sha256:aff07c09a53a08bc8cfccb9c85b05f1aa9a2a6f23728d790723543408344ce89
anyio==4.14.0 \
--hash=sha256:b47c1f9ccf73e67021df785332508f99379c68fa7d0684e8e3492cb1d4b23f89 \
--hash=sha256:dd9b7a2a9799ed6552fde617b2c5df02b7fdd7d88392fc48101e51bae46164d9
certifi==2026.6.17 \
--hash=sha256:024c88eeec92ca068db80f02b8b07c9cef7b9fe261d1d535abfd5abd6f6af432 \
--hash=sha256:2227dcbaafe0d2f59279d1762ddddc37783ed4354594f194ffc31d20f41fc3db
distro==1.9.0 \
--hash=sha256:2fa77c6fd8940f116ee1d6b94a2f90b13b5ea8d019b98bc8bafdcabcdd9bdbed \
--hash=sha256:7bffd925d65168f85027d8da9af6bddab658135b840670a223589bc0c8ef02b2
h11==0.16.0 \
--hash=sha256:4e35b956cf45792e4caa5885e69fba00bdbc6ffafbfa020300e549b208ee5ff1 \
--hash=sha256:63cf8bbe7522de3bf65932fda1d9c2772064ffb3dae62d55932da54b31cb6c86
httpcore==1.0.9 \
--hash=sha256:2d400746a40668fc9dec9810239072b40b4484b640a8c38fd654a024c7a1bf55 \
--hash=sha256:6e34463af53fd2ab5d807f399a9b45ea31c3dfa2276f15a2c3f00afff6e176e8
httpx==0.28.1 \
--hash=sha256:75e98c5f16b0f35b567856f597f06ff2270a374470a5c2392242528e3e3e42fc \
--hash=sha256:d909fcccc110f8c7faf814ca82a9a4d816bc5a6dbfea25d6591d6985b8ba59ad
idna==3.18 \
--hash=sha256:7f952cbe720b688055e3f87de14f5c3e5fdaa8bc3928985c4077ca689de849a2 \
--hash=sha256:ffb385a7e039654cef1ab9ef32c6fafe283c0c0467bba1d9029738ce4a14a848
jiter==0.15.0 \
--hash=sha256:01a8222cf05ab1128e239421156c207949808acaaea2bdfd33130ae666786e86 \
--hash=sha256:032396229564bca02440396bd327710719f724f5e7b7e9f7a8eb3faa4a2c2281 \
--hash=sha256:04b400bbf8c9efb03d9bdd976475c919c1d85593b04b9fff7ae234065daf87ae \
--hash=sha256:05906b93d72f03339e6bb7cf8dc10ebda64a0266126eed6beba79e20abcf5fd4 \
--hash=sha256:066f8f33f18b2419cd8213b2436fa7fbc9c499f315971cfa3ce1f9820c001b1b \
--hash=sha256:0ab068bce62a45aa3e7367eceaffb5dde60b7eb853be8dece45132e3d0ff4879 \
--hash=sha256:0be6f5ad41a809f303f416d17cec92a7a725902fb9b4f3de3d19362ac0ef8554 \
--hash=sha256:0e90a1c315a0226ec822d973817967f9223b7701546c8c2a7913e7ab0926294d \
--hash=sha256:0f862193b8696249d22ec433e85fd2ab0ad9596bc3e45e6c0bc55e8aeba97be2 \
--hash=sha256:1303d4d68a9b051ea90502402063ecf3807da00ad2affa19ca1ae3b90b3c5f67 \
--hash=sha256:144f8e72cb53dab146347b91cceac01f5481237f2b93b4a339a1ee8f8878b67c \
--hash=sha256:182226cbc930c9fab81bc2e41a4da672f89539906dadb05e75670ac07b94f71f \
--hash=sha256:1c11465f97e2abf45a014b83b730222f8f1c5335e802c7055a67d50de6f1f4e3 \
--hash=sha256:1c15024a3d892223b18f597c86d59387249dc396590844ce6b9f6131d1093bae \
--hash=sha256:1d54fb5b31dea401a41af3f8a7d2512e9b6a6a005491e6166c7e4ffab9639a9c \
--hash=sha256:25ffbe229aa8cd98c28879d8aa1a6e34ae77992ab984a65fba800859dab16269 \
--hash=sha256:2a77aadd57cac1682e4401a72724d2796d89a4ba129b1a5812aa94ee480826eb \
--hash=sha256:2ae901f3a55bfafdde31d289590fa25e3245735a2b1e8c7cc15871710a002871 \
--hash=sha256:2b0074e2f56eb2dacca1689760fd2852a068f85a0547a157b82cb4cafeb6768b \
--hash=sha256:2c8aea7781d2a372227871de4e1a1332aa96f5a89fd76c5e835dafdbad102887 \
--hash=sha256:2c9cb907439d20bd0c7d7565ca01ee52234203208433749bae5b516907526928 \
--hash=sha256:2fb6a5d26af81fc0f00f9360a891e05cf755e149bba391c4d563adc54812973d \
--hash=sha256:2fd73e3da91a0a722d67165e849ce2cdc10de0e0d48738c142be8c6c5f310f4c \
--hash=sha256:30ce1a5d16b5641dc935d50ef775af6a0871e3d14ab05d6fc54dff371b78e558 \
--hash=sha256:30ce785d2adb8e32c3f7741442370a74834ec4c01f3c48f0750227a0b4ef27d6 \
--hash=sha256:30f2218e6a9e5c18bc10fe6d41ac189c442c88eacf11bad9f28ef95a9bef00e6 \
--hash=sha256:351a341c2105aa430b7047e30f1bf7975f6313b00165d3fc07be2edaf741f279 \
--hash=sha256:37a10c377ce3a4a85f4a67f28b7afe093154cde77eaf248a72e856aa08b4d865 \
--hash=sha256:392b8ab019e5502d08aff85c6272209c24bc2cbe706ea82a56368f524236614a \
--hash=sha256:3e4540b8e74e4268811ac05db226a6a128ff572e7e0ce3f1163b693cadb184cd \
--hash=sha256:40b2c7e92c44a84d748d21706c68dc6ff8161d80b59c99d774721a0d2317d7c7 \
--hash=sha256:411fa4dfa5a7ae3d11491027ffb9beadec3996010a986862db70d91abba1c750 \
--hash=sha256:4251acc80e2b7c9b7b8823456ea0fceeb0734dac2df7636d3c711b38476b5a76 \
--hash=sha256:42bfb257930800cf43e7c62c832402c704ab60797c992faf88d20e903eac8f32 \
--hash=sha256:4363818355dbc70ae1a8e9eaba9de350d93ede4ff6992b8f8eb8cbb6e5122d42 \
--hash=sha256:4ab395feec8d249ec4044e228e98a7033f043426a265df439dc3698823f0a4e4 \
--hash=sha256:50164d7610c00e7cd913a873fce30b6beeebf4b37e53983e33f22de4c900f6b8 \
--hash=sha256:50e51156192722a9c58db112837d3f8ef96fb3c5ecc14e95f409134b08b158ec \
--hash=sha256:510c8b3c17a0ed9ac69850c0438dada3c9b82d9c4d589fcb62002a5a9cf3a866 \
--hash=sha256:5157de9f76eb4bc5ea74a1219366a25f945ad305641d74e04f59c54087091aa9 \
--hash=sha256:54d5d6090cdc1b7c9e780dfb04949a990adb1e301a2fc0bbcee7de4638d33f9a \
--hash=sha256:553fcac2ef2cb990877f9fc0833b8b629a3e6a5670b6b5fd58219b41a653ddc4 \
--hash=sha256:5607e6013ed7e6b0ec9661e467b7ffde0aa7ab36833a04850f26fcf88ed4845b \
--hash=sha256:5d6a60072b44c3c2b797a7ddcbcbbf2b34ea3cfd4721580fbfd2a09d9d9b84ba \
--hash=sha256:5f30bae8bc1c2d613e28e5af3e8cceb09b742f1c8a8a5f839fb67afaffc03b61 \
--hash=sha256:62ebd14e47e9aed9df4472afcb2663668ce4d74891cd54f86bf6e44029d6dc89 \
--hash=sha256:631f13a3d04e97d4e083993b10f4b99530e3a10d953e2eb5e196b7dc7f812ce0 \
--hash=sha256:6550fa135c7deb8ead6af49ed7ff648532ea8334a1447fe34a36315ef79c5c29 \
--hash=sha256:66b1880df2d01e206e8339769d1c7c1753bcb653efd6289e203f6f24ebada0c0 \
--hash=sha256:6eac374c5c975709b69c10f09afd199df74150172156ad10c8d4fd785b7da995 \
--hash=sha256:71683c38c825452999b5717fcae07ea708e8c93003e808be4319c1b02e3d176e \
--hash=sha256:7553333dd0930c104a5a0db8df72bf7219fe663d731383b576bb6ed6351c984d \
--hash=sha256:75e8a04e91432dde9f1838373cf93d23726c79d3e908d319acf0e796f85592e7 \
--hash=sha256:773b6eb282ce11ee19f05f6b2d4404fa308e5bbd353b0b80a0262caad6db2cd7 \
--hash=sha256:774f93f65031856bf14ad9f59bdcab8b8cad501e5ceabd51ba3525f76937a25b \
--hash=sha256:7c468136b8bd6bb18c8786e4236a1fa27362f24cb23450ba0cb204ab379b8e6f \
--hash=sha256:7ce8902f939970048b233087082e7bb829db29375811c7ad50687b8624c6fd08 \
--hash=sha256:7d3d6683288c11cbab50e865f2e2f13950179aa45410e30b2cfbd3fb7b0177bf \
--hash=sha256:7f6163c0f10b055245f814dcc59f4818da60dfe72f3e72ab89fc24b6bd5e9c52 \
--hash=sha256:8020c99ec13a7db2b6f96cbe82ef4721c88b426a4892f27478044af0284615ef \
--hash=sha256:813dfbb17d65328bf86e5f0905dd277ba2265d3ca20556e86c0c7035b7182e5a \
--hash=sha256:860a74063284a2ae9bfedd694f299cc2c68e2696c5f3d440cc9d18bb81b9dd04 \
--hash=sha256:8c9004af7c8d67cce7f1aae1026fb55607f4aa600710d08ede3a3ce4aeefe7e0 \
--hash=sha256:8d2c0c44d569ce0f2850f5c926f8caeb5f245fbc84475aeb36efccc2103e6dbd \
--hash=sha256:8f7e9bc0f1135039b22ee6eab588d42df1ce55842b30740a352885eb267bd941 \
--hash=sha256:90c5db5527c221249a876160663ab891ace358c17f7b9c93ec1478b7f0550e5c \
--hash=sha256:9100ddbec09741cc66feb0fc6773f8bdbd0e3c345689368f260082ff85dcc0cd \
--hash=sha256:913d02d29c9606643418d9ccfc3b72492ab25a6bf7889934e09a3490f8d3438b \
--hash=sha256:980c256edb05b78a111b99c4de3b1d32e31634b867fd1fc2cf726e7b7bba9854 \
--hash=sha256:9f924585cdacf631cd382b657966847bb537bf9ed0a6f9b991da5f05a631480f \
--hash=sha256:a254e10b593624d230c365b6d616b22ca0ad65e63a16e6631c2b3466022e6ba8 \
--hash=sha256:a2a438005b6f22d0273413484d6094d7c2c5d10ec1b3a3bf128e0d1d3ba53258 \
--hash=sha256:a97261f1fccb8e50ecd2890a96e46efdc3f57c80a197324c6777827231eca712 \
--hash=sha256:ab596fa3837e91e7e6a31b5f639988bfc6a35d1f915ac3932d946062219d588f \
--hash=sha256:abbf258599526ad0326fe51e252e24f2bd6f24f1852681b4b78feda3808f1d18 \
--hash=sha256:ac0d9ddea4350974be7a221fc25895f251a8fee748c889bdced2141c0fec1a49 \
--hash=sha256:acf4ee4d1fc55917239fe72972fb292dd773055d05eb040d36f4326e02cc2c0e \
--hash=sha256:ae1b0d82ac2d987f9ea512b1c9adfcc71a28de3dea3a6039b54d76cffda9901e \
--hash=sha256:b15741f501469009ae0ae90b7147958a664a7dede40aa7ff174a8a4645f546d0 \
--hash=sha256:b15d3ec9b0449c40e85319bdb4caa8b77ab526e74f5532ed94bec15e2f66822c \
--hash=sha256:b3b3b775e33d3bfaec9899edc526ae97b0da0bf9d071a46124ba419149a414f8 \
--hash=sha256:b6c0ffae686c39bf3737be60793783267628783ea42545632c10b291105aee45 \
--hash=sha256:c210f8b35dc6f30aafd4b4365ca89b9d1189f21ab49b8e68fa6322a847aef138 \
--hash=sha256:c2f6bb8b5216ab9e7873bc08b5d7bef2b8abbb578a3069bf1cd14a45d71d771d \
--hash=sha256:c60e71b6d10cfc284c9bf36bd885e8d44c46f688ce50aa91b5edd90181dea687 \
--hash=sha256:c6694a173ecabc12eb60efbc0b474464ead1951ff65cd8b1e72100715c64512b \
--hash=sha256:c77496cb10bd7549690fbbab3e5ec05857b83e49276f4a9423a766ddd2afcd4c \
--hash=sha256:c84c1b7be454b0c16f8499b4ebfbfd82ea5cca6527cceefcbbc06a7557b5ed2e \
--hash=sha256:cc0bc345cf2df9d1c00ac443f50d543c1ccfa8b0422cb85b1ab70d681c0b255b \
--hash=sha256:ceb8fc27d38793f9c97149be8302720c5b22e5c195a37bf2c45dc36c4600a512 \
--hash=sha256:cf4bd113a69c0a740e27cb962ce10630c36d2b8f59d759a651b955ee9d18a823 \
--hash=sha256:d1aa62e277fc1cbd80e6deacae6f4d983b41b3d7728e0645c5d741a6149bba45 \
--hash=sha256:d1e7b1776f0797956c509e123d0952d10d293a9492dea9f288ab9570ec01d1a5 \
--hash=sha256:d636d5095155afd364247f65070fab7beda13498d7ff4de331046e704ab9657f \
--hash=sha256:d726e3ceeb337191324b49de298142f27c3ad10886341555d1d5315b5f252c6a \
--hash=sha256:d72d8af5c1013656a8870c866660627d1a75bc185814ee022c8533caa1de88ae \
--hash=sha256:d8d2955167274e15d79a7a020afdd9b39c990eb80b2d89fca695d92dcfdd38ec \
--hash=sha256:d92a5cd21fdb083931d546c207aa29633787c5dc5b02daab2d32b843f88a2c53 \
--hash=sha256:e58585a58209d72691ce2d62a9147445f5a87beb0bde97fde284c96ae392a3d1 \
--hash=sha256:e7196e56f1cd69af1dbb07dff02dcfb260a50b45a82d409d92a06fedb32473b5 \
--hash=sha256:eda3071db3346334beae1360b46da4606da57bf3528c167b3c38533afaf9f2c5 \
--hash=sha256:edebcf7d1f601199084bb6e844d7dc67e03e04f6ac786b0332d616635c4ff7a4 \
--hash=sha256:ef1fd24d9413f6209e00d3d5a453e67acfe004a25cc6c8e8484faed4311ab9e8 \
--hash=sha256:f0b271b462769543716f92d3a4f90527df6ef5ed05ee95ec4137f513e21e1b77 \
--hash=sha256:f18f85e4218d1b40f000f42a92239a7a61a902cd42c65e6c360dbd17dcb20894 \
--hash=sha256:f1e1754960f38ec40613a07e5e372df67acb3b890fb383b6fb3de3e49ddbf3c7 \
--hash=sha256:f2143ab06181d2b029eedcb6af3cebe95f11bbac62441781860f98ee9330a6a6 \
--hash=sha256:f3d37768fce7f88dd2a8c6091f2325dea27d30d30d5c6e7a1c0f0af77723b708 \
--hash=sha256:fa248c9eb220197d363f688818dac2fd4b2f0cd7d843ca7105d652034823427d
openai==2.33.0 \
--hash=sha256:03ac37d70e8c9e3a8124214e3afa785e2cbc12e627fbd98177a086ef2fd87ad5 \
--hash=sha256:f850c435e2a4685bba3295bd54912dd26315d9c1b7733068186134d6e0599f9a
pydantic==2.13.4 \
--hash=sha256:45a282cde31d808236fd7ea9d919b128653c8b38b393d1c4ab335c62924d9aba \
--hash=sha256:c40756b57adaa8b1efeeced5c196f3f3b7c435f90e84ea7f443901bec8099ef6
pydantic-core==2.46.4 \
--hash=sha256:00c603d540afdd6b80eb39f078f33ebd46211f02f33e34a32d9f053bba711de0 \
--hash=sha256:0186750b482eefa11d7f435892b09c5c606193ef3375bcf94aa00ae6bfb66262 \
--hash=sha256:041bde0a48fd37cf71cab1c9d56d3e8625a3793fef1f7dd232b3ff37e978ecda \
--hash=sha256:0c563b08bca408dc7f65f700633d8442fffb2421fc47b8101377e9fd65051ff0 \
--hash=sha256:0cbe8b01f948de4286c74cdd6c667aceb38f5c1e26f0693b3983d9d74887c65e \
--hash=sha256:0ce40cd7b21210e99342afafbd4d0f76d784eb5b1d60f3bdc566be4983c6c73b \
--hash=sha256:0e96592440881c74a213e5ad528e2b24d3d4f940de2766bed9010ab1d9e51594 \
--hash=sha256:10e17cbb10a330363733efc4d7c4d0dd827ac0909b8f6a6542298fed1ea62f29 \
--hash=sha256:133878133d271ade3d41d1bfb2a45ec38dbdbda40bc065921c6b04e4630127e2 \
--hash=sha256:14d4edf427bdcf950a8a02d7cb44a08614388dd6e1bdcbf4f67504fa7887da9c \
--hash=sha256:14f4c5d6db102bd796a627bbb3a17b4cf4574b9ae861d8b7c9a9661c6dd3362d \
--hash=sha256:17299feefe090f2caa5b8e37222bb5f663e4935a8bfa6931d4102e5df1a9f398 \
--hash=sha256:184c081504d17f1c1066e430e117142b2c77d9448a97f7b65c6ac9fd9aee238d \
--hash=sha256:18e5ceec2ab67e6d5f1a9085e5a24c9c4e2ac4545730bfe668680bca05e555f3 \
--hash=sha256:19e51f073cd3df251856a8a4189fbdf1de4012c3ebacfb1884f94f1eb406079f \
--hash=sha256:1a7dd0b3ee80d90150e3495a3a13ac34dbcbfd4f012996a6a1d8900e91b5c0fb \
--hash=sha256:1d8ba486450b14f3b1d63bc521d410ec7565e52f887b9fb671791886436a42f7 \
--hash=sha256:2108ba5c1c1eca18030634489dc544844144ee36357f2f9f780b93e7ddbb44b5 \
--hash=sha256:228ee9bae8bef5b1e97ec58302f80357c37199e0d0a99174e138d28e6957b9d9 \
--hash=sha256:23ace664830ee0bfe014a0c7bc248b1f7f25ed7ad103852c317624a1083af462 \
--hash=sha256:2412e734dcb48da14d4e4006b82b46b74f2518b8a26ee7e58c6844a6cd6d03c4 \
--hash=sha256:29c61fc04a3d840155ff08e475a04809278972fe6aef51e2720554e96367e34b \
--hash=sha256:2f84c03c8607173d16b5a854ec68a2f9079ae03237a54fb506d13af47e1d018d \
--hash=sha256:3009f12e4e90b7f88b4f9adb1b0c4a3d58fe7820f3238c190047209d148026df \
--hash=sha256:3245406455a5d98187ec35530fd772b1d799b26667980872c8d4614991e2c4a2 \
--hash=sha256:3447661d99f75a3683a4cf5c87da72f2161964611864dbbeac7fbb118bb4bfc0 \
--hash=sha256:372429a130e469c9cd698925ce5fc50940b7a1336b0d82038e63d5bbc4edc519 \
--hash=sha256:395aebd9183f9d112f569aeb5b2214d1a10a33bec8456447f7fbdfa51d38d4cd \
--hash=sha256:3a233125ac121aa3ffba9a2b59edfc4a985a76092dc8279586ab4b71390875e7 \
--hash=sha256:3be77f45df024d789a672ae34f8b06fb346c4f9f46ea714956660ea4862e89ac \
--hash=sha256:3bf92c5d0e00fefaab325a4d27828fe6b6e2a21848686b5b60d2d9eeb09d76c6 \
--hash=sha256:3ecbc122d18468d06ca279dc26a8c2e2d5acb10943bb35e36ae92096dc3b5565 \
--hash=sha256:3fb702cd90b0446a3a1c5e470bfa0dd23c0233b676a9099ddcc964fa6ca13898 \
--hash=sha256:428e04521a40150c85216fc8b85e8d39fece235a9cf5e383761238c7fa9b96fb \
--hash=sha256:432c179df7874eeb73307aad2df0755e1ae0efa61ff0ea89b93e194411ae3928 \
--hash=sha256:4a05d69cba51d852c5c3e92758653245a50c0b646ced0cf05bd793ed592839d6 \
--hash=sha256:4c63ebc82684aa89d9a3bcbd13d515b3be44250dc68dd3bd81526c1cb31286c3 \
--hash=sha256:4fc73cb559bdb54b1134a706a2802a4cddd27a0633f5abb7e53056268751ac6a \
--hash=sha256:4fcbe087dbc2068af7eda3aa87634eba216dbda64d1ae73c8684b621d33f6596 \
--hash=sha256:56cb4851bcaf3d117eddcef4fe66afd750a50274b0da8e22be256d10e5611987 \
--hash=sha256:5855698a4856556d86e8e6cd8434bc3ac0314ee8e12089ae0e143f64c6256e4e \
--hash=sha256:5a4330cdbc57162e4b3aa303f588ba752257694c9c9be3e7ebb11b4aca659b5d \
--hash=sha256:5b712b53160b79a5850310b912a5ef8e57e56947c8ad690c227f5c9d7e561712 \
--hash=sha256:5d5902252db0d3cedf8d4a1bc68f70eeb430f7e4c7104c8c476753519b423008 \
--hash=sha256:617d7e2ca7dcb8c5cf6bcb8c59b8832c94b36196bbf1cbd1bfb56ed341905edd \
--hash=sha256:62f875393d7f270851f20523dd2e29f082bcc82292d66db2b64ea71f64b6e1c1 \
--hash=sha256:633147d34cf4550417f12e2b1a0383973bdf5cdfde212cb09e9a581cf10820be \
--hash=sha256:66ce7632c22d837c95301830e111ad0128a32b8207533b60896a96c4915192ea \
--hash=sha256:6b3ace8194b0e5204818c92802dcdca7fc6d88aabbb799d7c795540d9cd6d292 \
--hash=sha256:6f2eeda33a839975441c86a4119e1383c50b47faf0cbb5176985565c6bb02c33 \
--hash=sha256:7027560ee92211647d0d34e3f7cd6f50da56399d26a9c8ad0da286d3869a53f3 \
--hash=sha256:7283d57845ecf5a163403eb0702dfc220cc4fbdd18919cb5ccea4f95ee1cdab4 \
--hash=sha256:7a5f930472650a82629163023e630d160863fce524c616f4e5186e5de9d9a49b \
--hash=sha256:7bfb192b3f4b9e8a89b6277b6ce787564f62cfd272055f6e685726b111dc7826 \
--hash=sha256:811ff8e9c313ab425368bcbb36e5c4ebd7108c2bbf4e4089cfbb0b01eff63fac \
--hash=sha256:8233f2947cf85404441fd7e0085f53b10c93e0ee78611099b5c7237e36aacbf7 \
--hash=sha256:82cf5301172168103724d49a1444d3378cb20cdee30b116a1bd6031236298a5d \
--hash=sha256:8358a950c8909158e3df31538a7e4edc2d7265a7c54b47f0864d9e5bae9dcebf \
--hash=sha256:85bb3611ff1802f3ee7fdd7dbff26b56f343fb432d57a4728fdd49b6ef35e2f4 \
--hash=sha256:86e1a4418c6cd97d60c95c71164158eaf7324fae7b0923264016baa993eba6fc \
--hash=sha256:8b9bab013d1c7a79d3501ff86d0bc9c31bf587db4551677b96bec07df78c6b15 \
--hash=sha256:8c5dac79fa1614d1e06ca695109c6105923bd9c7d1d6c918d4e637b7e6b32fd3 \
--hash=sha256:8d0820e8192167f80d88d64038e609c31452eeca865b4e1d9950a27a4609b00b \
--hash=sha256:8daafc69c93ee8a0204506a3b6b30f586ef54028f52aeeeb5c4cfc5184fd5914 \
--hash=sha256:9037063db01f09b09e237c282b6792bd4da634b5402c4e7f0c61effed7701a04 \
--hash=sha256:905a0ed8ea6f2d61c1738835f99b699348d7857379083e5fc497fa0c967a407c \
--hash=sha256:90884113d8b48f760e9587002789ddd741e76ab9f89518cd1e43b1f1a52ec44b \
--hash=sha256:91a06d2e259ecfbd8c901d70c3c507900458498142b3026a296b7de4d1322cc9 \
--hash=sha256:926c9541b14b12b1681dca8a0b75feb510b06c6341b70a8e500c2fdcff837cce \
--hash=sha256:9401557acd873c3a7f3eb9383edef8ac4968f9510e340f4808d427e75667e7b4 \
--hash=sha256:9551187363ffc0de2a00b2e47c25aeaeb1020b69b668762966df15fc5659dd5a \
--hash=sha256:962ccbab7b642487b1d8b7df90ef677e03134cf1fd8880bf698649b22a69371f \
--hash=sha256:97e7cf2be5c77b7d1a9713a05605d49460d02c6078d38d8bef3cbe323c548424 \
--hash=sha256:9aa768456404a8bf48a4406685ac2bec8e72b62c69313734fa3b73cf33b3a894 \
--hash=sha256:9bc519fbf2b7578398853d815009ae5e4d4603d12f4e3f91da8c06852d3da3e9 \
--hash=sha256:9d56801be94b86a9da183e5f3766e6310752b99ff647e38b09a9500d88e46e76 \
--hash=sha256:9f444c499b3eefd3a92e348059471ea0c3a6e303d9c1cec09fa748fd9f895201 \
--hash=sha256:9fa8ae11da9e2b3126c6426f147e0fba88d96d65921799bb30c6abd1cb2c97fb \
--hash=sha256:a0f62d0a58f4e7da165457e995725421e0064f2255d8eccebc49f41bbc23b109 \
--hash=sha256:a396dcc17e5a0b164dbe026896245a4fa9ff402edca1dff0be3d53a517f74de4 \
--hash=sha256:aaa2a54443eff1950ba5ddc6b6ccda0d9c84a364276a62f969bdf2a390650848 \
--hash=sha256:ad785e92e6dc634c21555edc8bd6b64957ab844541bcb96a1366c202951ae526 \
--hash=sha256:af8244b2bef6aaad6d92cda81372de7f8c8d36c9f0c3ea36e827c60e7d9467a0 \
--hash=sha256:b078afbc25f3a1436c7a1d2cd3e322497ee99615ba97c563566fdf46aff1ee01 \
--hash=sha256:b2f69dec1725e79a012d920df1707de5caf7ed5e08f3be4435e25803efc47458 \
--hash=sha256:b8458003118a712e66286df6a707db01c52c0f52f7db8e4a38f0da1d3b94fc4e \
--hash=sha256:bb63e0198ca18aad131c089b9204c23079c3afa95487e561f4c522d519e55aba \
--hash=sha256:bfec22eab3c8cc2ceec0248aec886624116dc079afa027ecc8ad4a7e62010f8a \
--hash=sha256:c1747f85cee84c26985853c6f3d9bd3e75da5212912443fa111c113b9c246f39 \
--hash=sha256:c1b3f518abeca3aa13c712fd202306e145abf59a18b094a6bafb2d2bbf59192c \
--hash=sha256:c50f2528cf200c5eed56faf3f4e22fcd5f38c157a8b78576e6ba3168ec35f000 \
--hash=sha256:c68fcd102d71ea85c5b2dfac3f4f8476eff42a9e078fd5faefff6d145063536b \
--hash=sha256:c7a7bd4e39e8e4c12c39cd480356842b6a8a06e41b23a55a5e3e191718838ddf \
--hash=sha256:c94f0688e7b8d0a67abf40e57a7eaaecd17cc9586706a31b76c031f63df052b4 \
--hash=sha256:cbaf13819775b7f769bf4a1f066cb6df7a28d4480081a589828ef190226881cd \
--hash=sha256:cd2213145bcc2ba85884d0ac63d222fece9209678f77b9b4d76f054c561adb28 \
--hash=sha256:ce5c1d2a8b27468f433ca974829c44060b8097eedc39933e3c206a90ee49c4a9 \
--hash=sha256:d396ec2b979760aaf3218e76c24e65bd0aca24983298653b3a9d7a45f9e47b30 \
--hash=sha256:d51026d73fcfd93610abc7b27789c26b313920fcfb20e27462d74a7f8b06e983 \
--hash=sha256:d80ee3d731373b24cebbc10d689ca4ee1875caf0d5703a245db18efd4dd37fc1 \
--hash=sha256:d995260fdf4e1db774581b4900e0f832abe3c7c84996726bbc161b19c8f29e76 \
--hash=sha256:da4b951fe36dc7c3a1ccb4e3cd1747c3542b8c9ceede8fc86cae054e764485f5 \
--hash=sha256:daa27d92c36f24388fe3ad306b174781c747627f134452e4f128ea00ce1fe8c4 \
--hash=sha256:db06ffe51636ffe9ca531fe9023dd64bdd794be8754cb5df57c5498ae5b518a7 \
--hash=sha256:e0d65b8c354be7fb5f720c3caa8bc940bc2d20ce749c8e06135f07f8ed95dd7c \
--hash=sha256:e68b7a074f65a2fd746c52a7ce6142ab7006074ac269ace0c25cd8ba171f8066 \
--hash=sha256:e739fee756ba1010f8bcccb534252e85a35fe45ae92c295a06059ce58b74ccd3 \
--hash=sha256:e846ae7835bf0703ae43f534ab79a867146dadd59dc9ca5c8b53d5c8f7c9ef02 \
--hash=sha256:e9c26f834c65f5752f3f06cb08cb86a913ceb7274d0db6e267808a708b46bc89 \
--hash=sha256:ea793e075b70290d89d8142074262885d3f7da19634845135751bd6344f73b50 \
--hash=sha256:f027324c56cd5406ca49c124b0db10e56c69064fec039acc571c29020cc87c76 \
--hash=sha256:f13a646d65d09fbf1bc6b3a9635d30095c8e7e5cc419ff35ecc563c5fd04cd49 \
--hash=sha256:f47286a97f0bc9b8859519809077b91b2cefe4ae47fcbf5e466a009c1c5d742b \
--hash=sha256:f747929cf940cddb5b3668a390056ddd5ba2e5010615ea2dcf4f9c4f3ab8791d \
--hash=sha256:f99626688942fb746e545232e7726926f3be91b5975f8b55327665fafda991c7 \
--hash=sha256:f9fa868638bf362d3d138ea55829cefb3d5f4b0d7f142234382a15e2485dbec4 \
--hash=sha256:fbdb89b3e1c94a30cc5edfce477c6e6a5dc4d8f84665b455c27582f211a1c72c \
--hash=sha256:fc010ab034c8c7452522748bf937df58020d256ccae0874463d1f4d01758af8e \
--hash=sha256:fc3e9034a63de20e15e8ade85358bc6efc614008cab72898b4b4952bea0509ff \
--hash=sha256:fd8b3d9fd264be37976686c7f65cd52a83f5e84f4bfd2adf9c1d469676bbb6ae
sniffio==1.3.1 \
--hash=sha256:2f6da418d1f1e0fddd844478f41680e794e6051915791a034ff65e5f100525a2 \
--hash=sha256:f4324edc670a0f49750a81b895f35c3adb843cca46f0530f79fc1babb23789dc
tqdm==4.68.3 \
--hash=sha256:00dfa48452b6b6cfae3dd9885636c23d3422d1ec97c66d96818cbd5e0821d482 \
--hash=sha256:39832cc2def2789a6f29df83f172db7416cea70052c0907a57801c5f2fdccb03
typing-extensions==4.15.0 \
--hash=sha256:0cea48d173cc12fa28ecabc3b837ea3cf6f38c6d1136f85cbaaf598984861466 \
--hash=sha256:f0fa19c6845758ab08074a0cfa8b7aecb71c999ca73d62883bc25cc018c4e548
typing-inspection==0.4.2 \
--hash=sha256:4ed1cacbdc298c220f1bd249ed5287caa16f34d44ef4e9c3d0cbad5b521545e7 \
--hash=sha256:ba561c48a67c5958007083d386c3295464928b01faa735ab8547c5692e87f464

1797
.github/scripts/triage_with_llm.py vendored
View file

File diff suppressed because it is too large Load diff

29
.github/workflows/_test-unit-base.yml vendored
View file

@ -37,6 +37,18 @@ on:
required: false
type: number
default: 60
test-timeout-seconds:
description: >-
Per-test ceiling enforced by pytest-timeout, covering fixture setup and
teardown as well as the test body. A test that hangs fails with a
traceback of where it was stuck instead of idling the shard until
`timeout-minutes` cancels it. Timed-out tests are excluded from reruns
because pytest-timeout arms its timer once per test and
pytest-rerunfailures reruns inside that same window, so a rerun of a
timed-out test would run with no timer at all.
required: false
type: number
default: 120
max-failures:
description: "Stop after this many failures"
required: false
@ -51,6 +63,11 @@ on:
description: "Unique name for the coverage artifact (must be unique per run)"
required: true
type: string
legacy-mcp-peer:
description: "Install the isolated SDK1 peer for MCP compatibility tests"
required: false
type: boolean
default: false
permissions:
contents: read
@ -113,10 +130,17 @@ jobs:
- name: Install dependencies
if: steps.changes.outputs.decision != 'skip'
timeout-minutes: 8
env:
LEGACY_MCP_PEER: ${{ inputs.legacy-mcp-peer }}
run: |
diff -u model_prices_and_context_window.json litellm/model_prices_and_context_window_backup.json
.github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra saml
uv run --no-sync python -c 'import os, sys; print(sys.version); assert f"{sys.version_info.major}.{sys.version_info.minor}" == os.environ["UV_PYTHON"]'
if [ "$LEGACY_MCP_PEER" = "true" ]; then
uv venv --python "${UV_PYTHON}" .venv-mcp-peer
uv pip install --python .venv-mcp-peer 'mcp==1.28.1' 'langchain-mcp-adapters==0.2.1'
echo "MCP_TEST_PEER_PYTHON=$GITHUB_WORKSPACE/.venv-mcp-peer/bin/python" >> "$GITHUB_ENV"
fi
- name: Cache Prisma binaries
if: steps.changes.outputs.decision != 'skip'
@ -137,6 +161,7 @@ jobs:
MAX_FAILURES: ${{ inputs.max-failures }}
WORKERS: ${{ inputs.workers }}
RERUNS: ${{ inputs.reruns }}
TEST_TIMEOUT_SECONDS: ${{ inputs.test-timeout-seconds }}
DIST: ${{ inputs.dist }}
COVERAGE_CORE: sysmon
run: |
@ -146,6 +171,8 @@ jobs:
--maxfail="${MAX_FAILURES}" \
--reruns "${RERUNS}" \
--reruns-delay 1 \
--timeout="${TEST_TIMEOUT_SECONDS}" \
--rerun-except "from pytest-timeout" \
--durations=20 \
--cov=./litellm --cov=./enterprise/litellm_enterprise \
--cov-report=xml:coverage.xml \
@ -157,6 +184,8 @@ jobs:
-n "${WORKERS}" \
--reruns "${RERUNS}" \
--reruns-delay 1 \
--timeout="${TEST_TIMEOUT_SECONDS}" \
--rerun-except "from pytest-timeout" \
--dist="${DIST}" \
--durations=20 \
--cov=./litellm --cov=./enterprise/litellm_enterprise \

61
.github/workflows/auto-merge-price-sync.yml vendored
View file

@ -1,61 +0,0 @@
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

92
.github/workflows/close_low_quality_prs.yml vendored
View file

@ -1,92 +0,0 @@
name: Close Low-Quality PRs
# Auto-close any open PR (including drafts, regardless of age) authored by an
# external OSS contributor that Greptile reviewed with a confidence score
# below 4/5. Closures are explained in a comment that tells the contributor
# to push fixes and open a fresh PR (since OSS authors cannot reopen a PR
# closed by a bot/maintainer) or comment `@agent-shin reconsider` to have
# Agent Shin re-evaluate.
#
# Manual one-off run:
# gh workflow run "Close Low-Quality PRs" -f close=true
#
# Dry-run preview (no PRs are touched):
# gh workflow run "Close Low-Quality PRs" -f close=false
on:
schedule:
# Daily at 09:00 UTC. Pairs well with the stale-issue workflow at midnight.
- cron: "0 9 * * *"
workflow_dispatch:
inputs:
close:
description: "Actually close matching PRs (false = dry run)."
required: false
default: "false"
type: choice
options:
- "true"
- "false"
min_age_days:
description: "Minimum PR age in days (default 0 = no age filter)."
required: false
default: "0"
min_score:
description: "Greptile score below which a PR is closed (1-5)."
required: false
default: "4"
limit:
description: "Maximum number of PRs to close in a single run."
required: false
default: "25"
permissions:
contents: read
pull-requests: write
issues: write
jobs:
close-low-quality-prs:
if: github.repository == 'BerriAI/litellm'
runs-on: ubuntu-latest
steps:
- name: Checkout triage script
uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
sparse-checkout: .github/scripts
persist-credentials: false
- name: Set up Python
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
with:
python-version: "3.12"
- name: Run low-quality PR closer
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
# Scheduled runs are ALWAYS dry-run, even when AGENT_SHIN_ENABLED is
# "true", so the team can QA the closer's verdicts in step summaries
# before any contributor sees a PR closed. Real closures only happen
# on manual workflow_dispatch with close=true (and the variable set).
CLOSE_FLAG: ${{ github.event.inputs.close || 'false' }}
AGENT_SHIN_ENABLED: ${{ vars.AGENT_SHIN_ENABLED }}
MIN_AGE_DAYS: ${{ github.event.inputs.min_age_days || '0' }}
MIN_SCORE: ${{ github.event.inputs.min_score || '4' }}
LIMIT: ${{ github.event.inputs.limit || '25' }}
run: |
set -euo pipefail
ARGS=(
--repo "${{ github.repository }}"
--min-age-days "${MIN_AGE_DAYS}"
--min-score "${MIN_SCORE}"
--limit "${LIMIT}"
)
if [ "${AGENT_SHIN_ENABLED:-false}" != "true" ]; then
echo "::notice::AGENT_SHIN_ENABLED is not 'true' -> forcing dry-run regardless of close input."
elif [ "${GITHUB_EVENT_NAME:-}" = "workflow_dispatch" ] && [ "${CLOSE_FLAG}" = "true" ]; then
ARGS+=(--close)
echo "::notice::Running in close-on-fail mode."
else
echo "::notice::AGENT_SHIN_ENABLED is true but this trigger is dry-run (scheduled event or close=false)."
fi
python3 .github/scripts/close_low_quality_prs.py "${ARGS[@]}"

4
.github/workflows/codspeed.yml vendored
View file

@ -69,7 +69,7 @@ jobs:
uv run --frozen --no-default-groups
--with pytest==8.3.5
--with pytest-codspeed==4.3.0
--with "mcp>=1.26.0,<2.0"
--with "mcp>=2.2.0,<3.0"
--with "a2a-sdk>=1.1.0,<2.0"
pytest
-p pytest_codspeed.plugin
@ -86,7 +86,7 @@ jobs:
uv run --frozen --no-default-groups
--with pytest==8.3.5
--with pytest-codspeed==4.3.0
--with "mcp>=1.26.0,<2.0"
--with "mcp>=2.2.0,<3.0"
--with "a2a-sdk>=1.1.0,<2.0"
pytest
-p pytest_codspeed.plugin

28
.github/workflows/create_daily_oss_agent_shin_branch.yml vendored
View file

@ -1,28 +0,0 @@
name: Create Daily oss-agent-shin Branch
on:
schedule:
- cron: "0 0 * * *" # Runs every day at midnight UTC
workflow_dispatch: # Allow manual trigger
jobs:
create-oss-agent-shin-branch:
if: github.repository == 'BerriAI/litellm'
runs-on: ubuntu-latest
permissions:
contents: write
steps:
- name: Create daily oss-agent-shin branch
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
run: |
BRANCH_NAME="litellm_oss_agent_shin_$(date +'%m_%d_%Y')"
echo "Creating branch: $BRANCH_NAME"
if gh api "repos/${{ github.repository }}/git/ref/heads/$BRANCH_NAME" --silent 2>/dev/null; then
echo "Branch $BRANCH_NAME already exists. Skipping creation."
exit 0
fi
MAIN_SHA=$(gh api "repos/${{ github.repository }}/git/ref/heads/main" --jq '.object.sha')
gh api "repos/${{ github.repository }}/git/refs" -f ref="refs/heads/$BRANCH_NAME" -f sha="$MAIN_SHA" --silent
echo "Successfully created branch: $BRANCH_NAME at $MAIN_SHA"

11
.github/workflows/duplicate_issue_check.yml vendored
View file

@ -93,12 +93,13 @@ jobs:
responses-api-endpoint: ${{ vars.LITELLM_API_BASE }}/v1/responses
prompt-file: .github/prompts/duplicate-issue-check.md
output-schema-file: .github/prompts/duplicate-issue-check.schema.json
sandbox: read-only
# read-only denies network, and the whole method is searching the tracker with gh
codex-args: '["-c", "sandbox_permissions=[\"network-full-access\"]"]'
sandbox: workspace-write
# The whole method is searching the tracker with gh, and network is only switchable in workspace-write
codex-args: '["-c", "sandbox_workspace_write.network_access=true"]'
model: ${{ vars.DUPLICATE_CHECK_MODEL }}
# Issue authors have no write access and the action refuses them by default; the
# prompt is fixed, the sandbox read-only, and the only token is read-only on a public repo
codex-version: "0.154.0"
# Issue authors have no write access and the action refuses them by default; the prompt is
# fixed, writes stay inside the throwaway checkout, and the only token is read-only on a public repo
allow-users: "*"
- name: Summary

71
.github/workflows/issue_fixed_comment.yml vendored Normal file
View file

@ -0,0 +1,71 @@
name: Issue fixed comment
on:
issues:
types: [closed]
workflow_dispatch:
inputs:
issue_number:
description: "Closed issue number to comment on manually."
required: true
pull_request:
paths:
- .github/workflows/issue_fixed_comment.yml
- scripts/comment-fixed-issue.ts
- scripts/comment-fixed-issue.test.ts
- scripts/auto-close-duplicates.ts
permissions: {}
concurrency:
group: issue-fixed-comment-${{ github.event.issue.number || github.event.inputs.issue_number || github.run_id }}
cancel-in-progress: false
jobs:
comment-fixed-issue-tests:
if: github.event_name == 'pull_request'
runs-on: ubuntu-latest
timeout-minutes: 5
permissions:
contents: read
steps:
- name: Checkout repository
uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
persist-credentials: false
- name: Setup Bun
uses: oven-sh/setup-bun@0c5077e51419868618aeaa5fe8019c62421857d6 # v2.2.0
with:
bun-version: "1.4.0"
- name: Test the closer lookup, the release placement and the comment
run: bun test scripts/comment-fixed-issue.test.ts
comment-fixed-issue:
if: github.event_name != 'pull_request' && github.repository == 'BerriAI/litellm'
runs-on: ubuntu-latest
timeout-minutes: 5
permissions:
contents: read
issues: write
steps:
- name: Checkout scripts
uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
sparse-checkout: scripts
persist-credentials: false
- name: Setup Bun
uses: oven-sh/setup-bun@0c5077e51419868618aeaa5fe8019c62421857d6 # v2.2.0
with:
# Exact version, never latest: the next step holds an issues: write token
bun-version: "1.4.0"
- name: Name the release that carries the fix
run: bun run scripts/comment-fixed-issue.ts | tee -a "${GITHUB_STEP_SUMMARY}"
env:
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
ISSUE_NUMBER: ${{ github.event.issue.number || github.event.inputs.issue_number }}
DEFAULT_BRANCH: ${{ github.event.repository.default_branch }}
DRY_RUN: ${{ vars.ISSUE_FIXED_COMMENT_ENABLED != 'true' }}

98
.github/workflows/test-mcp-dependency-resolution.yml vendored Normal file
View file

@ -0,0 +1,98 @@
name: LiteLLM MCP Dependency Resolution
on:
pull_request:
branches:
- main
- litellm_internal_staging
- litellm_oss_staging
- "litellm_**"
permissions:
contents: read
pull-requests: read
concurrency:
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}
cancel-in-progress: ${{ github.event_name == 'pull_request' }}
jobs:
resolve:
runs-on: ubuntu-latest
timeout-minutes: 15
strategy:
fail-fast: false
matrix:
python-version: ["3.10", "3.11", "3.12", "3.13", "3.14"]
steps:
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
persist-credentials: false
- name: Detect relevant changes
id: changes
uses: ./.github/actions/detect-changes
with:
category: mcp-dependencies
- name: Set up Python
if: steps.changes.outputs.decision != 'skip'
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
with:
python-version: ${{ matrix.python-version }}
- name: Set up uv
if: steps.changes.outputs.decision != 'skip'
uses: ./.github/actions/setup-uv-with-retries
with:
version: "0.10.9"
- name: Cache the Rust build
if: steps.changes.outputs.decision != 'skip'
uses: ./.github/actions/cache-cargo-build
- name: Verify lockfile
if: steps.changes.outputs.decision != 'skip'
run: |
uv lock --check
- name: Check locked runtime installations
if: steps.changes.outputs.decision != 'skip'
run: |
for extra in core mcp proxy; do
args=()
if [ "$extra" != core ]; then args=(--extra "$extra"); fi
UV_PROJECT_ENVIRONMENT=".venv-$extra" .github/scripts/uv_sync_with_retries.sh --frozen --no-dev --no-editable --python ${{ matrix.python-version }} "${args[@]}"
uv pip check --python ".venv-$extra"
if [ "$extra" = core ]; then
checker=("$GITHUB_WORKSPACE/tests/base_sdk_tests/check_base_sdk_install.py")
else
checker=("$GITHUB_WORKSPACE/scripts/check_mcp_sdk_install.py" --extra "$extra")
fi
(cd "$RUNNER_TEMP" && "$GITHUB_WORKSPACE/.venv-$extra/bin/python" "${checker[@]}")
done
- name: Build the public wheel
if: steps.changes.outputs.decision != 'skip'
run: uv build --all-packages --wheel --out-dir dist/mcp-check
- name: Check lowest direct runtime installations
if: steps.changes.outputs.decision != 'skip'
run: |
wheel=$(realpath dist/mcp-check/litellm-[0-9]*.whl)
for extra in core mcp proxy; do
args=()
if [ "$extra" != core ]; then args=(--extra "$extra"); fi
uv pip compile pyproject.toml --no-sources --find-links dist/mcp-check "${args[@]}" --python-version ${{ matrix.python-version }} --resolution lowest-direct -o "lowest-$extra.txt"
uv venv --python ${{ matrix.python-version }} ".venv-lowest-$extra"
uv pip sync --find-links dist/mcp-check --python ".venv-lowest-$extra" "lowest-$extra.txt"
uv pip install --python ".venv-lowest-$extra" --no-deps "$wheel"
uv pip check --python ".venv-lowest-$extra"
if [ "$extra" = core ]; then
checker=("$GITHUB_WORKSPACE/tests/base_sdk_tests/check_base_sdk_install.py")
else
checker=("$GITHUB_WORKSPACE/scripts/check_mcp_sdk_install.py" --extra "$extra")
fi
(cd "$RUNNER_TEMP" && "$GITHUB_WORKSPACE/.venv-lowest-$extra/bin/python" "${checker[@]}")
done

63
.github/workflows/test-mcp.yml vendored
View file

@ -1,63 +0,0 @@
name: LiteLLM MCP Tests (folder - tests/mcp_tests)
on:
pull_request:
branches:
- main
- litellm_internal_staging
- litellm_oss_staging
- "litellm_**"
permissions:
contents: read
pull-requests: read
concurrency:
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}
cancel-in-progress: ${{ github.event_name == 'pull_request' }}
jobs:
test:
runs-on: ubuntu-latest
timeout-minutes: 25
steps:
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
persist-credentials: false
- name: Detect relevant changes
id: changes
uses: ./.github/actions/detect-changes
- name: Thank You Message
run: |
echo "### 🙏 Thank you for contributing to LiteLLM!" >> $GITHUB_STEP_SUMMARY
echo "Your PR is being tested now. We appreciate your help in making LiteLLM better!" >> $GITHUB_STEP_SUMMARY
- name: Set up Python
if: steps.changes.outputs.decision != 'skip'
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
with:
python-version: "3.12"
- name: Set up uv
if: steps.changes.outputs.decision != 'skip'
uses: ./.github/actions/setup-uv-with-retries
with:
version: "0.10.9"
- name: Cache the Rust build
if: steps.changes.outputs.decision != 'skip'
uses: ./.github/actions/cache-cargo-build
- name: Install dependencies
if: steps.changes.outputs.decision != 'skip'
run: |
uv lock --check
.github/scripts/uv_sync_with_retries.sh --frozen --group proxy-dev --extra proxy --extra semantic-router
- name: Run MCP tests
if: steps.changes.outputs.decision != 'skip'
run: |
uv run --no-sync pytest tests/mcp_tests -x -vv -n 4 --cov=./litellm --cov=./enterprise/litellm_enterprise --cov-report=xml --durations=5

9
.github/workflows/test-unit.yml vendored
View file

@ -49,6 +49,14 @@ jobs:
fail-fast: false
matrix:
include:
- shard: mcp-integration
artifact-name: mcp-integration
test-path: "tests/mcp_tests"
workers: 2
reruns: 0
timeout-minutes: 20
job-timeout-minutes: 60
- shard: core-utils
artifact-name: core-utils
test-path: "tests/test_litellm/litellm_core_utils"
@ -254,3 +262,4 @@ jobs:
timeout-minutes: ${{ matrix.timeout-minutes }}
job-timeout-minutes: ${{ matrix.job-timeout-minutes }}
artifact-name: ${{ matrix.artifact-name }}
legacy-mcp-peer: ${{ matrix.shard == 'mcp-integration' }}

172
.github/workflows/triage_reconsider.yml vendored
View file

@ -1,172 +0,0 @@
name: Agent Shin — reconsider
# Comment-trigger workflow: when the PR/issue author (or an internal
# collaborator) comments `@agent-shin reconsider` on a CLOSED PR/issue,
# Agent Shin re-runs LLM-judge triage on the current title+body and:
#
# - on PASS: posts a "re-evaluated and reopened" comment + reopens.
# - on FAIL: posts a "still missing X" comment and leaves it closed,
# so the contributor can iterate again.
#
# This exists because GitHub does NOT let an external (non-write-access)
# OSS contributor reopen a PR/issue closed by a bot or maintainer. Without
# this comment trigger, a contributor whose PR Agent Shin auto-closed
# would have no path back into the review queue except opening a fresh PR
# (which loses the original PR's history). The bot, on the other hand,
# has write access via GH_TOKEN and can reopen on their behalf.
#
# DRY-RUN BY DEFAULT — gated on `vars.AGENT_SHIN_ENABLED == 'true'` just
# like the other Agent Shin workflows. The workflow also gates on the
# commenter being either the PR/issue author or an internal collaborator
# (OWNER/MEMBER/COLLABORATOR) so random commenters cannot DOS the LLM
# judge or force a reopen.
on:
issue_comment:
types: [created]
permissions:
contents: read
issues: write
pull-requests: write
jobs:
reconsider:
if: |
github.repository == 'BerriAI/litellm'
&& contains(github.event.comment.body, '@agent-shin reconsider')
runs-on: ubuntu-latest
steps:
- name: Authorize commenter
# Only the PR/issue author OR an internal collaborator may trigger
# a reconsider. Outside random commenters could otherwise spam the
# phrase to burn LLM budget or, if a fail-open bug were ever
# introduced, force a reopen on someone else's behalf.
#
# We expose the authorization decision as a step output and gate
# every subsequent (potentially destructive) step on it. A `run:`
# step with `exit 0` would NOT stop the job — only `if:` gating
# on a known-true output is safe here.
id: auth
env:
COMMENTER: ${{ github.event.comment.user.login }}
AUTHOR: ${{ github.event.issue.user.login }}
ASSOCIATION: ${{ github.event.comment.author_association }}
run: |
set -euo pipefail
if [ "${COMMENTER}" = "${AUTHOR}" ]; then
echo "::notice::Authorized: commenter is the PR/issue author."
echo "authorized=true" >> "$GITHUB_OUTPUT"
exit 0
fi
case "${ASSOCIATION}" in
OWNER|MEMBER|COLLABORATOR)
echo "::notice::Authorized: commenter is an internal collaborator (${ASSOCIATION})."
echo "authorized=true" >> "$GITHUB_OUTPUT"
;;
*)
echo "::notice::Commenter '${COMMENTER}' (${ASSOCIATION}) is not authorized to trigger reconsider; skipping subsequent steps."
echo "authorized=false" >> "$GITHUB_OUTPUT"
;;
esac
- name: React 👀 to acknowledge the reconsider
# Add an eyes reaction to the triggering comment the moment we accept
# it, so the contributor gets instant feedback that the bot saw their
# `@agent-shin reconsider` before the slower triage steps run. Gated on
# AGENT_SHIN_ENABLED so dry-run leaves no visible trace. Best-effort:
# a reactions API hiccup must never fail the actual reconsider.
if: steps.auth.outputs.authorized == 'true' && vars.AGENT_SHIN_ENABLED == 'true'
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
COMMENT_ID: ${{ github.event.comment.id }}
run: |
set -euo pipefail
gh api --method POST \
-H "Accept: application/vnd.github+json" \
"repos/${{ github.repository }}/issues/comments/${COMMENT_ID}/reactions" \
-f content=eyes \
|| echo "::warning::failed to add 👀 reaction (non-fatal)"
- name: Checkout triage script
if: steps.auth.outputs.authorized == 'true'
uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
sparse-checkout: .github/scripts
persist-credentials: false
- name: Set up Python
if: steps.auth.outputs.authorized == 'true'
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
with:
python-version: "3.12"
- name: Install LLM client
if: steps.auth.outputs.authorized == 'true'
run: pip install --no-cache-dir --require-hashes -r .github/scripts/triage-requirements.txt
- name: Run Agent Shin reconsider
if: steps.auth.outputs.authorized == 'true'
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
# Only expose the LLM key when the bot is enabled, so a PR/issue
# author can't force paid LLM calls by spamming `@agent-shin
# reconsider` while the bot is still in dry-run. The Python script
# calls the LLM whenever this var is set (regardless of `--close`);
# stripping `--close` doesn't suppress the API call, only the
# destructive side effects. Mirror the gating used by every other
# Agent Shin workflow (triage_pr_with_llm.yml, review_gate.yml, ...).
OPENAI_API_KEY: ${{ vars.AGENT_SHIN_ENABLED == 'true' && secrets.OPENAI_API_KEY || '' }}
OPENAI_BASE_URL: ${{ vars.OPENAI_BASE_URL }}
TRIAGE_MODEL: ${{ vars.TRIAGE_MODEL }}
AGENT_SHIN_ENABLED: ${{ vars.AGENT_SHIN_ENABLED }}
# `issue_comment` events fire for both issues and PR comments.
# `issue.pull_request` is set iff this is a PR comment, so we use
# its presence to decide whether to invoke `--pr N` or `--issue N`.
IS_PR: ${{ github.event.issue.pull_request != null }}
NUMBER: ${{ github.event.issue.number }}
run: |
set -euo pipefail
if [ "${IS_PR}" = "true" ]; then
ARGS=(--repo "${{ github.repository }}" --pr "${NUMBER}" --reconsider)
else
ARGS=(--repo "${{ github.repository }}" --issue "${NUMBER}" --reconsider)
fi
# Reconsider's destructive actions (post comment + reopen) are
# gated on `--close`, mirroring the regular triage workflows.
# When AGENT_SHIN_ENABLED is not the EXACT string "true", we
# still run the script so its verdict + would-X action lands in
# the step summary for QA — but without `--close`, the script
# returns `would-reopen` / `would-reconsider-still-failing`
# instead of touching GitHub state.
#
# Use the positive `= "true"` gate (not `!= "true" -> exit`) so
# the workflow guardrails in
# tests/test_litellm/test_github_triage_workflows.py see the
# canonical fail-safe enable pattern. Unknown values like
# "True", "yes", "1", or typos fall through to the dry-run
# branch, which is the safe default.
if [ "${AGENT_SHIN_ENABLED:-false}" = "true" ]; then
ARGS+=(--close)
echo "::notice::Agent Shin reconsider ENABLED — running real triage (close=true)."
else
echo "::notice::AGENT_SHIN_ENABLED is not 'true' -> reconsider stays in dry-run (no comment, no reopen)."
fi
python3 .github/scripts/triage_with_llm.py "${ARGS[@]}"
- name: React 👍 when the reconsider finishes
# Once the reconsider run has completed successfully, add a thumbs-up so
# the contributor sees the bot is done (the 👀 stays, signalling
# seen -> handled). `success()` keeps this from firing if the run
# errored, and the AGENT_SHIN_ENABLED gate keeps dry-run inert.
if: success() && steps.auth.outputs.authorized == 'true' && vars.AGENT_SHIN_ENABLED == 'true'
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
COMMENT_ID: ${{ github.event.comment.id }}
run: |
set -euo pipefail
gh api --method POST \
-H "Accept: application/vnd.github+json" \
"repos/${{ github.repository }}/issues/comments/${COMMENT_ID}/reactions" \
-f content=+1 \
|| echo "::warning::failed to add 👍 reaction (non-fatal)"

130
AGENTS.md
View file

@ -1,3 +1,131 @@
Read @CLAUDE.md for coding guidelines
Do not write comments unless they are any of:
- absolutely necessary to explain some very complex business logic (in which case, keep it concise and clear)
- used as an input for tools to read and act on. For example:
- entries in `.git-blame-ignore-revs` saying which commit is excluded from git blame
- a lint or type checker suppression like `# mutable-ok` or `# pyright: ignore[reportArgumentType] # <reason>` when introducing a truly unavoidable violation
- a TODO or FIXME
- Not great to have those, but if it's unavoidable, make sure to include a strong, concise reason for why it's there or, better yet, link to a GitHub issue for the follow-up work
Explanation: The point of this rule is to keep out AI slop comments. AI writes way too many and way too verbose comments. Code comments are, in a way, a violation of DRY code. You must update logic in two locations to change the code, and "hard to change" is literally the definition of tech debt. We should instead aim to write code that is intuitive and clear, even at a glance, to the reader, being both easy to maintain and high performance
Don't assume that the existing code is correct or the right way of doing things / good coding patterns. In fact, there are a lot of bad coding practices, overly complex code, code smells, etc. If something doesn't look right, speak up. Feel free to break existing patterns or question weird existing code to make new code high quality, as in:
- correct
- secure
- performant
- readable
- easy to maintain/change
- modern
In descending order of importance
When adding new features, add meaningful tests. Don't add tests that don't check anything substantial and is there just to make the code coverage pass. Yes, code coverage is important, but I'd rather have no signal whether the code is working than tests that don't fail when code is broken. The goal is to have tests that would fail before the feature was added/if the code was mutated in a way that breaks the feature and succeed only when the feature is fully working. I should run mutation testing and see > 90% kill rate
Same thing for bug fixes. The tests should make it so that this specific bug can never happen again without failing tests (i.e., regression)
Never test structure of code only function of it
A test must only fail when litellm code changes. Never pin facts we don't own (a vendor's price, a third party's field, an upstream default, today's date) as literals or as "X must be absent"; assert the invariant our code guarantees instead, e.g. two rows agree, a value is within range, a field is derived from another. If an outside fact is truly load-bearing, cite its source and date next to the assertion so a reader can tell stale from broken
`tests/test_litellm/` mirrors `litellm/` in a parallel path (see `tests/test_litellm/readme.md`). Name tests `test_<filename>.py`, but always match the existing test file in the directory you touch — many provider dirs use longer descriptive names (e.g. `test_anthropic_chat_transformation.py`) to avoid ambiguity across sibling folders. For bug fixes, extend the existing mapped test file rather than creating a new one. Only create a new test file for a new feature (provider, endpoint, or transformation module) that has no mapped test yet, following that directory's naming convention (or `test_<filename>.py` if you're the first test there). One focused regression test beats many shallow ones
End-to-end tests belong in `tests/e2e/` and must follow the harness conventions documented in that directory's `AGENTS.md`
When creating PRs, target the repository's current default branch for both internal and external / OSS contributions. Check it with `python3 scripts/default_branch.py --branch` instead of assuming a branch name or relying on cached `origin/HEAD`
When writing a PR body, treat the comments and imperative instructions inside .github/pull_request_template.md as rules to follow, not just layout. Agent harnesses may strip HTML comments from copies of that file injected into context, so read .github/pull_request_template.md from disk before writing a PR body to make sure you see every comment rule
Same applies for filing bug reports and feature requests, with .github/ISSUE_TEMPLATE/bug_report.yml and .github/ISSUE_TEMPLATE/feature_request.yml, respectively
If you're resolving a linear ticket, in the "## Linear ticket" section of the PR, say "Resolves LIT-1234", replacing "LIT-1234" with the actual ticket id that you're resolving. If you don't have the ticket id, don't make one up or search for it. Just leave the section blank
Never use `pytest` commands or the like as "Screenshots / Proof of Fix". We prefer curl'ing a live proxy instance running on localhost:4000 (I like to run it with `python litellm/proxy/proxy_cli.py --config litellm/proxy/dev_config.yaml --detailed_debug --reload --use_v2_migration_resolver 2>&1 | tee litellm.log`; the Admin UI dev server is `npm run dev` in `ui/litellm-dashboard`, served on port 3000) and showing both the command run and the output. Also, it should hit real LLM provider APIs, not mocks, and cost real $$$ because that is the most realistic test. The proof of fix should be exactly what the end user / customer would see / do. The run logs in PR #27703 is a prime example of how to do it (not a huge fan of using a python test script that future me and the team will have no visibility into; I prefer just curl commands or a short list of bash commands (e.g., using `for`)). If it's a UI thing, or the main use case runs through a headful agentic coding tool like Claude Code or Codex, drive that surface yourself and embed your own before and after screenshots of it in the PR (the Admin UI page, or what the coding tool shows), next to an ordered list of the URLs to go to (e.g., http://localhost:4000/ui/?page=logs), where to click, and what fields to fill out so a reviewer can reproduce it
If you ever write any human-facing text (pull requests, issues, commit messages, discussion posts, github comments, release notes, docs, etc.), always follow these guidelines to sound less AI-y:
- don't use emojis
- don't use "—". Instead, reach for ",", ".", conjunction words, ":", ";", etc. in descending order of preference: vary among them, weighted toward the front of the list, and skip "," where it would cause a comma splice or the sentence is getting long. Overusing any one of them, ";" especially, also feels AI-y. A word cap does not penalize you for adding more sentences: when writing under tight word budgets, prefer a period split or a conjunction over ";", and keep to at most one ";" per message
- don't use the pattern "It's not X, it's Y", "You're not X, you're Y", etc.
- unless explicitly asked, don't use bulleted or numbered lists unless it would be nonsensical not to. Instead, prefer prose
- don't add a trailing "." at the end of paragraphs (just like this file). That means every paragraph, not just the last one (of the markdown file, PR description, GitHub comment, etc.). Rule of thumb: if you're adding new line(s) before the next sentence, don't add a "."
- don't use →. Instead, prefer not to use arrows, and if need be, use -> instead
- use plain, simple, everyday engineering language: the common phrase engineers actually say over rare compact phrasing, in grammatically complete sentences. When explicitly asked to use bullets or ordered lists and structure legitimately helps the reader, prefer nested bullets (any depth is fine) over dense lines in a flat structure
Don't hesitate to use values in .env to get needed API keys and other secrets, as long as you never add them to conversation history, commit them, or include them in GitHub issues / PRs
Python max line length is 120, not 88
Never edit or commit `ruff-strict-budget.json`, `type-discipline-budget.json`, `basedpyright-code-budget.json`, or `test-quality-budget.json` on a PR branch, and don't run `make lint-budget-update` there. A scheduled Devin automation lowers the limits on the default branch in its own PR by exactly what landed since the last ratchet, so concurrent PRs don't fight over the same `"limit"` lines. Keep the hosted automation's target in sync when the repository default changes. If your branch already carries a budget edit, drop it before opening the PR
`make check` (f.k.a. `make pre-commit`, which still works identically as an alias) saves its complete output to a log file in .git (overwriting previous logs) and prints that path as its first and last output lines. To inspect a run, read or grep that log instead of re-running the multi-minute checks just to see a different slice
`make check`, `make lint`, `scripts/pre_commit_lint.sh`, and the standalone budget gates (`scripts/ruff_strict_gate.py`, `scripts/type_discipline_gate.py`, `scripts/type_check_gate.py`) each hold one of 2 machine-wide slots, so when other sessions or worktrees on the same box are already running heavy work, yours prints "all N machine-wide slots are busy; queueing" and then stays quiet until a slot frees. Give the command a long timeout and let it wait rather than killing it, retrying it, or assuming it hung. Don't change the # of machine-wide slots or make it unlimited by setting `LITELLM_GATE_SLOTS=0`
If you're trying to create a new function that relies on untyped stuff, instead of adding more Any's and pushing `reportAny` / `reportExplicitAny` closer to their basedpyright ceilings, just validate it in the caller with Pydantic (a model or `TypeAdapter` that returns the typed thing or raises will do) and then pass the now typed variable in
If you get an LIT001 or LIT002 fail, refactor the code to follow functional programming best practices rather than introducing mutable data structures. For example, build values in one shot with comprehensions or generators wrapped in `tuple()` / `MappingProxyType()` / `frozenset()` instead of seeding an empty `list`/`dict`/`set` and mutating it over time. Ideally, `# mutable-ok` is never used; reach for it only as a genuine last resort when an immutable rewrite is truly impossible, and always pair it with a real reason
Every lint or type suppression must name the exact rule inside brackets and carry a reason comment, e.g. `# pyright: ignore[reportArgumentType] # stubs lack async overload` or `# noqa: TID251 # <reason>`. `# type: ignore` is banned (LIT009): pyrightconfig.json sets `enableTypeIgnoreComments` to false, so it silently does nothing
Commit and push your work when you're done without asking
When referencing or running models (coding, QA'ing, writing docs, writing tests, etc.), use the latest model in that model family unless otherwise specified; treat your training knowledge, memories, configs, and tests as stale, and determine the family's latest with model_prices_and_context_window.json or the web
Always pull before starting any work. The checkout or worktree may be sitting on a stale branch
If you're an internal contributor, when creating a new PR, the typical flow is to branch off the repository's current default branch and create a branch prefixed with litellm_. Do not create a branch prefixed with claude/ and generally do not have / in your branch names
Do not add `Co-Authored-By: Claude` or any Claude attribution to commit messages. Never use a `claude/` prefix or put a `/` in a branch name. Do not add "Generated with Claude Code" (or any similar attribution) to PR descriptions or comments. Do not create a new PR/branch off the existing PR to fix/add something that is related and could've just been committed directly to the existing PR's branch
When working on a PR, keep the PR description in sync with new commits being made
All GitHub comments must be human-readable and 15-25 words max
Monkeypatching attributes of a class to do testing is an anti-pattern. Prefer dependency-injecting things into classes. That way, at unit test time, you can pass a mocked dependency in
Do not put names of customers or customer company names in code, PR descriptions, issue bodies, etc. This means never mention literally any company name. Especially if you're about to say a sentence mentioning that the reason the PR exists was a feature/model/bug fix/etc. requested by a company. That's the indication that you should replace that company name with "the customer". e.g. not "Model request from Acme (Pylon #1234)" but "Model request from a customer (Pylon #1234)". This is because the codebase is public. The only exception is for publicly known providers or vendors such as OpenAI, Anthropic, AWS Bedrock, etc. only IF we're adding support for that provider/vendor in general and NOT if that PR or whatnot was a request by one of them, and they're actually one of our customers.
CI supply-chain safety: Never pipe a remote script into a shell (`curl ... | bash`, `wget ... | sh`); download the artifact to a file, verify its SHA-256 checksum, then install. Pin every external tool to a specific version with a full URL (not `latest` or `stable`). Verify checksums for all downloaded binaries, using the provider's official `.sha256` / `.sha256sum` sidecar when available. These rules apply to every download in CI
Prisma migrations apply synchronously at proxy boot, before it serves traffic, so a migration must only change schema, never rewrite rows. No `UPDATE`, `DELETE` or `MERGE`, and no `INSERT ... SELECT`: on a spend-log-sized table any of those is minutes of downtime plus a doubled heap that plain autovacuum won't give back. `tests/code_coverage_tests/check_migrations_no_data_rewrites.py` enforces this. When a rewrite is genuinely bounded and has to ship inside the migration, mark the statement `-- data-migration-ok: <what bounds it>`
Follow these coding conventions for new/updated code (a three-line fix in a legacy file shouldn't trigger huge drive-by refactors):
- Composition over inheritance
- Never-nester: early returns over deep nesting
- Don't throw; model failures as values (One function (e.g., raise_public) maps error union to existing public exception contracts via exhaustive match + assert_never)
- No mutation; don't reassign variables, global or local. Instead of mutable lists and dicts, prefer tuples, frozen dataclasses (with slots=True), `MappingProxyType`, etc.
- Annotate every variable with `: Final` (LIT010). Unpacking and walrus targets cannot carry the annotation, so they are implicitly final. Don't rebind them. Never rebind or mutate function parameters (LIT011); `self`/`cls` attribute stores are the exception. If rebinding or in-place mutation is truly unavoidable, suppress with `# rebind-ok: <reason>`
- Qualify every TypedDict field with `ReadOnly[...]` (LIT012), which nests freely with `Required` / `NotRequired` / `Annotated` in any order. If making the key writable is truly unavoidable, suppress with `# writable-ok: <reason>`
- Use dependency injection
- Fully typed; no `Any` or coarse types like `dict[str, Any]` or just `dict`. Every function parameter must be strongly typed
- Use tagged unions + match
- No monster files or god objects
- No file sprawl: deliberate file and folder structure
- Standard over hand-rolled: use the official SDK or a library where one exists; where none does, follow industry standards instead of inventing local conventions
- API-fragmentation-aware: when logic must branch on which API surface produced or consumes data (e.g. chat completions vs Anthropic Messages vs Responses API shapes), proactively look for an existing shared helper (e.g. `litellm_core_utils/prompt_templates/factory.py`) before writing per-surface parsing in the new module; if none exists, add one there instead of duplicating the same format-detection logic in every new guardrail/integration
Follow conventional commits for commit names and PR titles
## Think Before Coding
**Don't assume. Don't hide confusion. Surface tradeoffs**
Before implementing:
- State your assumptions explicitly. If uncertain, ask
- If multiple interpretations exist, present them. Don't pick silently
- If a simpler approach exists, say so. Push back when warranted
- If something is unclear, stop. Name what's confusing. Ask
## Simplicity First
**Minimum code that solves the problem. Nothing speculative**
- No features beyond what was asked
- No abstractions for single-use code
- No "flexibility" or "configurability" that wasn't requested
- No error handling for impossible scenarios
- If you write 200 lines and it could be 50, rewrite it
Ask yourself: "Would a senior engineer say this is overcomplicated?" If yes, simplify
Before requesting maintainer review, verify the current PR tip passes required CI and code coverage, meets Greptile confidence of at least 4/5, and has acceptable Veria and Bugbot reviews. Inspect warnings and findings, fix actionable issues, and rerun the affected checks and reviewers after changes. Record evidence for any false positive or unavailable review; never treat a pending or missing bot result as a pass. Do not lower coverage thresholds or lint budgets to satisfy a check

129
CLAUDE.md
View file

@ -1,129 +0,0 @@
Do not write comments unless they are any of:
- absolutely necessary to explain some very complex business logic (in which case, keep it concise and clear)
- used as an input for tools to read and act on. For example:
- entries in `.git-blame-ignore-revs` saying which commit is excluded from git blame
- a lint or type checker suppression like `# mutable-ok` or `# pyright: ignore[reportArgumentType] # <reason>` when introducing a truly unavoidable violation
- a TODO or FIXME
- Not great to have those, but if it's unavoidable, make sure to include a strong, concise reason for why it's there or, better yet, link to a GitHub issue for the follow-up work
Explanation: The point of this rule is to keep out AI slop comments. AI writes way too many and way too verbose comments. Code comments are, in a way, a violation of DRY code. You must update logic in two locations to change the code, and "hard to change" is literally the definition of tech debt. We should instead aim to write code that is intuitive and clear, even at a glance, to the reader, being both easy to maintain and high performance
Don't assume that the existing code is correct or the right way of doing things / good coding patterns. In fact, there are a lot of bad coding practices, overly complex code, code smells, etc. If something doesn't look right, speak up. Feel free to break existing patterns or question weird existing code to make new code high quality, as in:
- correct
- secure
- performant
- readable
- easy to maintain/change
- modern
In descending order of importance
When adding new features, add meaningful tests. Don't add tests that don't check anything substantial and is there just to make the code coverage pass. Yes, code coverage is important, but I'd rather have no signal whether the code is working than tests that don't fail when code is broken. The goal is to have tests that would fail before the feature was added/if the code was mutated in a way that breaks the feature and succeed only when the feature is fully working. I should run mutation testing and see > 90% kill rate
Same thing for bug fixes. The tests should make it so that this specific bug can never happen again without failing tests (i.e., regression)
Never test structure of code only function of it
A test must only fail when litellm code changes. Never pin facts we don't own (a vendor's price, a third party's field, an upstream default, today's date) as literals or as "X must be absent"; assert the invariant our code guarantees instead, e.g. two rows agree, a value is within range, a field is derived from another. If an outside fact is truly load-bearing, cite its source and date next to the assertion so a reader can tell stale from broken
`tests/test_litellm/` mirrors `litellm/` in a parallel path (see `tests/test_litellm/readme.md`). Name tests `test_<filename>.py`, but always match the existing test file in the directory you touch — many provider dirs use longer descriptive names (e.g. `test_anthropic_chat_transformation.py`) to avoid ambiguity across sibling folders. For bug fixes, extend the existing mapped test file rather than creating a new one. Only create a new test file for a new feature (provider, endpoint, or transformation module) that has no mapped test yet, following that directory's naming convention (or `test_<filename>.py` if you're the first test there). One focused regression test beats many shallow ones
End-to-end tests belong in `tests/e2e/` and must follow the harness conventions documented in that directory's `CLAUDE.md`
When creating PRs, target the repository's current default branch for both internal and external / OSS contributions. Check it with `python3 scripts/default_branch.py --branch` instead of assuming a branch name or relying on cached `origin/HEAD`
When writing a PR body, treat the comments and imperative instructions inside .github/pull_request_template.md as rules to follow, not just layout. Agent harnesses may strip HTML comments from copies of that file injected into context, so read .github/pull_request_template.md from disk before writing a PR body to make sure you see every comment rule
Same applies for filing bug reports and feature requests, with .github/ISSUE_TEMPLATE/bug_report.yml and .github/ISSUE_TEMPLATE/feature_request.yml, respectively
If you're resolving a linear ticket, in the "## Linear ticket" section of the PR, say "Resolves LIT-1234", replacing "LIT-1234" with the actual ticket id that you're resolving. If you don't have the ticket id, don't make one up or search for it. Just leave the section blank
Never use `pytest` commands or the like as "Screenshots / Proof of Fix". We prefer curl'ing a live proxy instance running on localhost:4000 (I like to run it with `python litellm/proxy/proxy_cli.py --config litellm/proxy/dev_config.yaml --detailed_debug --reload --use_v2_migration_resolver 2>&1 | tee litellm.log`; the Admin UI dev server is `npm run dev` in `ui/litellm-dashboard`, served on port 3000) and showing both the command run and the output. Also, it should hit real LLM provider APIs, not mocks, and cost real $$$ because that is the most realistic test. The proof of fix should be exactly what the end user / customer would see / do. The run logs in PR #27703 is a prime example of how to do it (not a huge fan of using a python test script that future me and the team will have no visibility into; I prefer just curl commands or a short list of bash commands (e.g., using `for`)). If it's a UI thing, or the main use case runs through a headful agentic coding tool like Claude Code or Codex, drive that surface yourself and embed your own before and after screenshots of it in the PR (the Admin UI page, or what the coding tool shows), next to an ordered list of the URLs to go to (e.g., http://localhost:4000/ui/?page=logs), where to click, and what fields to fill out so a reviewer can reproduce it
If you ever write any human-facing text (pull requests, issues, commit messages, discussion posts, github comments, release notes, docs, etc.), always follow these guidelines to sound less AI-y:
- don't use emojis
- don't use "—". Instead, reach for ",", ".", conjunction words, ":", ";", etc. in descending order of preference: vary among them, weighted toward the front of the list, and skip "," where it would cause a comma splice or the sentence is getting long. Overusing any one of them, ";" especially, also feels AI-y. A word cap does not penalize you for adding more sentences: when writing under tight word budgets, prefer a period split or a conjunction over ";", and keep to at most one ";" per message
- don't use the pattern "It's not X, it's Y", "You're not X, you're Y", etc.
- unless explicitly asked, don't use bulleted or numbered lists unless it would be nonsensical not to. Instead, prefer prose
- don't add a trailing "." at the end of paragraphs (just like this file). That means every paragraph, not just the last one (of the markdown file, PR description, GitHub comment, etc.). Rule of thumb: if you're adding new line(s) before the next sentence, don't add a "."
- don't use →. Instead, prefer not to use arrows, and if need be, use -> instead
- use plain, simple, everyday engineering language: the common phrase engineers actually say over rare compact phrasing, in grammatically complete sentences. When explicitly asked to use bullets or ordered lists and structure legitimately helps the reader, prefer nested bullets (any depth is fine) over dense lines in a flat structure
Don't hesitate to use values in .env to get needed API keys and other secrets, as long as you never add them to conversation history, commit them, or include them in GitHub issues / PRs
Python max line length is 120, not 88
Never edit or commit `ruff-strict-budget.json`, `type-discipline-budget.json`, `basedpyright-code-budget.json`, or `test-quality-budget.json` on a PR branch, and don't run `make lint-budget-update` there. A scheduled Devin automation lowers the limits on the default branch in its own PR by exactly what landed since the last ratchet, so concurrent PRs don't fight over the same `"limit"` lines. Keep the hosted automation's target in sync when the repository default changes. If your branch already carries a budget edit, drop it before opening the PR
`make check` (f.k.a. `make pre-commit`, which still works identically as an alias) saves its complete output to a log file in .git (overwriting previous logs) and prints that path as its first and last output lines. To inspect a run, read or grep that log instead of re-running the multi-minute checks just to see a different slice
`make check`, `make lint`, `scripts/pre_commit_lint.sh`, and the standalone budget gates (`scripts/ruff_strict_gate.py`, `scripts/type_discipline_gate.py`, `scripts/type_check_gate.py`) each hold one of 2 machine-wide slots, so when other sessions or worktrees on the same box are already running heavy work, yours prints "all N machine-wide slots are busy; queueing" and then stays quiet until a slot frees. Give the command a long timeout and let it wait rather than killing it, retrying it, or assuming it hung. Don't change the # of machine-wide slots or make it unlimited by setting `LITELLM_GATE_SLOTS=0`
If you're trying to create a new function that relies on untyped stuff, instead of adding more Any's and pushing `reportAny` / `reportExplicitAny` closer to their basedpyright ceilings, just validate it in the caller with Pydantic (a model or `TypeAdapter` that returns the typed thing or raises will do) and then pass the now typed variable in
If you get an LIT001 or LIT002 fail, refactor the code to follow functional programming best practices rather than introducing mutable data structures. For example, build values in one shot with comprehensions or generators wrapped in `tuple()` / `MappingProxyType()` / `frozenset()` instead of seeding an empty `list`/`dict`/`set` and mutating it over time. Ideally, `# mutable-ok` is never used; reach for it only as a genuine last resort when an immutable rewrite is truly impossible, and always pair it with a real reason
Every lint or type suppression must name the exact rule inside brackets and carry a reason comment, e.g. `# pyright: ignore[reportArgumentType] # stubs lack async overload` or `# noqa: TID251 # <reason>`. `# type: ignore` is banned (LIT009): pyrightconfig.json sets `enableTypeIgnoreComments` to false, so it silently does nothing
Commit and push your work when you're done without asking
When referencing or running models (coding, QA'ing, writing docs, writing tests, etc.), use the latest model in that model family unless otherwise specified; treat your training knowledge, memories, configs, and tests as stale, and determine the family's latest with model_prices_and_context_window.json or the web
Always pull before starting any work. The checkout or worktree may be sitting on a stale branch
If you're an internal contributor, when creating a new PR, the typical flow is to branch off the repository's current default branch and create a branch prefixed with litellm_. Do not create a branch prefixed with claude/ and generally do not have / in your branch names
Do not add `Co-Authored-By: Claude` or any Claude attribution to commit messages. Never use a `claude/` prefix or put a `/` in a branch name. Do not add "Generated with Claude Code" (or any similar attribution) to PR descriptions or comments. Do not create a new PR/branch off the existing PR to fix/add something that is related and could've just been committed directly to the existing PR's branch
When working on a PR, keep the PR description in sync with new commits being made
All GitHub comments must be human-readable and 15-25 words max
Monkeypatching attributes of a class to do testing is an anti-pattern. Prefer dependency-injecting things into classes. That way, at unit test time, you can pass a mocked dependency in
Do not put names of customers or customer company names in code, PR descriptions, issue bodies, etc. This means never mention literally any company name. Especially if you're about to say a sentence mentioning that the reason the PR exists was a feature/model/bug fix/etc. requested by a company. That's the indication that you should replace that company name with "the customer". e.g. not "Model request from Acme (Pylon #1234)" but "Model request from a customer (Pylon #1234)". This is because the codebase is public. The only exception is for publicly known providers or vendors such as OpenAI, Anthropic, AWS Bedrock, etc. only IF we're adding support for that provider/vendor in general and NOT if that PR or whatnot was a request by one of them, and they're actually one of our customers.
CI supply-chain safety: Never pipe a remote script into a shell (`curl ... | bash`, `wget ... | sh`); download the artifact to a file, verify its SHA-256 checksum, then install. Pin every external tool to a specific version with a full URL (not `latest` or `stable`). Verify checksums for all downloaded binaries, using the provider's official `.sha256` / `.sha256sum` sidecar when available. These rules apply to every download in CI
Prisma migrations apply synchronously at proxy boot, before it serves traffic, so a migration must only change schema, never rewrite rows. No `UPDATE`, `DELETE` or `MERGE`, and no `INSERT ... SELECT`: on a spend-log-sized table any of those is minutes of downtime plus a doubled heap that plain autovacuum won't give back. `tests/code_coverage_tests/check_migrations_no_data_rewrites.py` enforces this. When a rewrite is genuinely bounded and has to ship inside the migration, mark the statement `-- data-migration-ok: <what bounds it>`
Follow these coding conventions for new/updated code (a three-line fix in a legacy file shouldn't trigger huge drive-by refactors):
- Composition over inheritance
- Never-nester: early returns over deep nesting
- Don't throw; model failures as values (One function (e.g., raise_public) maps error union to existing public exception contracts via exhaustive match + assert_never)
- No mutation; don't reassign variables, global or local. Instead of mutable lists and dicts, prefer tuples, frozen dataclasses (with slots=True), `MappingProxyType`, etc.
- Annotate every variable with `: Final` (LIT010). Unpacking and walrus targets cannot carry the annotation, so they are implicitly final. Don't rebind them. Never rebind or mutate function parameters (LIT011); `self`/`cls` attribute stores are the exception. If rebinding or in-place mutation is truly unavoidable, suppress with `# rebind-ok: <reason>`
- Qualify every TypedDict field with `ReadOnly[...]` (LIT012), which nests freely with `Required` / `NotRequired` / `Annotated` in any order. If making the key writable is truly unavoidable, suppress with `# writable-ok: <reason>`
- Use dependency injection
- Fully typed; no `Any` or coarse types like `dict[str, Any]` or just `dict`. Every function parameter must be strongly typed
- Use tagged unions + match
- No monster files or god objects
- No file sprawl: deliberate file and folder structure
- Standard over hand-rolled: use the official SDK or a library where one exists; where none does, follow industry standards instead of inventing local conventions
- API-fragmentation-aware: when logic must branch on which API surface produced or consumes data (e.g. chat completions vs Anthropic Messages vs Responses API shapes), proactively look for an existing shared helper (e.g. `litellm_core_utils/prompt_templates/factory.py`) before writing per-surface parsing in the new module; if none exists, add one there instead of duplicating the same format-detection logic in every new guardrail/integration
Follow conventional commits for commit names and PR titles
## Think Before Coding
**Don't assume. Don't hide confusion. Surface tradeoffs**
Before implementing:
- State your assumptions explicitly. If uncertain, ask
- If multiple interpretations exist, present them. Don't pick silently
- If a simpler approach exists, say so. Push back when warranted
- If something is unclear, stop. Name what's confusing. Ask
## Simplicity First
**Minimum code that solves the problem. Nothing speculative**
- No features beyond what was asked
- No abstractions for single-use code
- No "flexibility" or "configurability" that wasn't requested
- No error handling for impossible scenarios
- If you write 200 lines and it could be 50, rewrite it
Ask yourself: "Would a senior engineer say this is overcomplicated?" If yes, simplify

17
CONTRIBUTING.md
View file

@ -148,7 +148,7 @@ make lint
Individual linting commands:
```bash
make format-check # Check Black formatting
make format-check # Check ruff format formatting
make lint-ruff # Run Ruff linting
make lint-basedpyright # Run basedpyright type checking
make check-circular-imports # Check for circular imports
@ -160,14 +160,14 @@ Apply formatting (auto-fixes issues):
make format
```
> **Black formatting is enforced in CI.** All PRs must pass the Black formatting check.
> **Formatting is enforced in CI.** All PRs must pass the `ruff format --check` step.
>
> - **AI coding agents** (Claude Code, Copilot, Cursor, etc.): `AGENTS.md` and `CLAUDE.md` instruct agents to run `poetry run black .` before committing.
> - **VS Code users**: Install the [Black Formatter extension](https://marketplace.visualstudio.com/items?itemName=ms-python.black-formatter) and enable format-on-save:
> - **AI coding agents** (Claude Code, Copilot, Cursor, etc.): follow `AGENTS.md` and run `make format` before committing.
> - **VS Code users**: Install the [Ruff extension](https://marketplace.visualstudio.com/items?itemName=charliermarsh.ruff) and enable format-on-save:
> ```json
> {
> "[python]": {
> "editor.defaultFormatter": "ms-python.black-formatter",
> "editor.defaultFormatter": "charliermarsh.ruff",
> "editor.formatOnSave": true
> }
> }
@ -197,8 +197,8 @@ make help # Show all available commands
make install-dev # Install development dependencies
make install-proxy-dev # Install proxy development dependencies
make install-test-deps # Install the full local test environment
make format # Apply Black code formatting
make format-check # Check Black formatting (matches CI)
make format # Apply ruff format code formatting
make format-check # Check ruff format formatting (matches CI)
make lint # Run all linting checks
make test-unit # Run unit tests
make test-integration # Run integration tests
@ -210,8 +210,7 @@ make test-unit-helm # Run Helm unit tests
LiteLLM follows the [Google Python Style Guide](https://google.github.io/styleguide/pyguide.html).
Our automated quality checks include:
- **Black** for consistent code formatting
- **Ruff** for linting and code quality
- **Ruff** for formatting, linting, and code quality
- **basedpyright** for static type checking
- **Circular import detection**
- **Import safety validation**

2
GEMINI.md
View file

@ -1 +1 @@
Read @CLAUDE.md for coding guidelines
Read @AGENTS.md for coding guidelines

5
README.md
View file

@ -633,9 +633,8 @@ For detailed contributing guidelines, see [CONTRIBUTING.md](CONTRIBUTING.md).
LiteLLM follows the [Google Python Style Guide](https://google.github.io/styleguide/pyguide.html).
Our automated checks include:
- **Black** for code formatting
- **Ruff** for linting and code quality
- **MyPy** for type checking
- **Ruff** for formatting, linting, and code quality
- **basedpyright** for type checking
- **Circular import detection**
- **Import safety checks**

54
ci_cd/generate_model_prices_schema.py
View file

@ -3,7 +3,7 @@ from __future__ import annotations
import json
import sys
from pathlib import Path
from typing import Optional
from typing import Final, Optional
import jsonschema
@ -19,6 +19,10 @@ NONNEG_NUMBER: JsonSchema = {"type": "number", "minimum": 0}
NONNEG_INTEGER: JsonSchema = {"type": "integer", "minimum": 0}
BOOLEAN: JsonSchema = {"type": "boolean"}
STRING: JsonSchema = {"type": "string"}
TIME_WINDOW: Final[JsonSchema] = {"type": "string", "pattern": r"^([01]\d|2[0-3]):[0-5]\d-([01]\d|2[0-3]):[0-5]\d$"}
WEEKDAY_PATTERN: Final = (
r"(?i)^(mon|monday|tue|tues|tuesday|wed|wednesday|thu|thur|thurs|thursday|fri|friday|sat|saturday|sun|sunday)$"
)
EXTRA_BOOLEAN_KEYS = frozenset(
{
@ -31,7 +35,51 @@ EXTRA_BOOLEAN_KEYS = frozenset(
}
)
HOURS_UTC: Final[JsonSchema] = {
"description": 'UTC "HH:MM-HH:MM" window, or a list of them; a window may wrap past midnight.',
"oneOf": [TIME_WINDOW, {"type": "array", "items": TIME_WINDOW, "minItems": 1}],
}
OFF_PEAK_WINDOW: Final[JsonSchema] = {
"type": "object",
"properties": {
"hours_utc": HOURS_UTC,
"weekdays": {
"type": "array",
"description": "ISO-8601 weekday numbers (1 = Monday .. 7 = Sunday) or English day names the window applies on.",
"items": {
"oneOf": [
{"type": "integer", "minimum": 1, "maximum": 7},
{"type": "string", "pattern": WEEKDAY_PATTERN},
]
},
"minItems": 1,
},
},
"required": ["hours_utc"],
"additionalProperties": False,
}
OBJECT_KEYS: dict[str, JsonSchema] = {
"off_peak_pricing": {
"type": "object",
"description": "Rates that replace the same-named base fields while the request falls inside the stated UTC windows.",
"properties": {
"hours_utc": HOURS_UTC,
"windows": {"type": "array", "items": OFF_PEAK_WINDOW, "minItems": 1},
"weekday_timezone": {
"type": "string",
"description": "IANA zone the weekdays of each window are read on; defaults to UTC.",
},
"input_cost_per_token": NONNEG_NUMBER,
"output_cost_per_token": NONNEG_NUMBER,
"output_cost_per_reasoning_token": NONNEG_NUMBER,
"cache_read_input_token_cost": NONNEG_NUMBER,
"cache_creation_input_token_cost": NONNEG_NUMBER,
},
"anyOf": [{"required": ["hours_utc"]}, {"required": ["windows"]}],
"additionalProperties": False,
},
"search_context_cost_per_query": {
"type": "object",
"description": "USD cost per web search query, keyed by search context size.",
@ -327,9 +375,7 @@ def render(schema: JsonSchema) -> str:
def validation_errors(prices: dict, schema: JsonSchema) -> tuple:
validator = jsonschema.Draft202012Validator(
schema, format_checker=jsonschema.Draft202012Validator.FORMAT_CHECKER
)
validator = jsonschema.Draft202012Validator(schema, format_checker=jsonschema.Draft202012Validator.FORMAT_CHECKER)
return tuple(
f"{'.'.join(str(part) for part in error.absolute_path)}: {error.message}"
for error in validator.iter_errors(prices)

9
enterprise/enterprise_hooks/openai_moderation.py
View file

@ -17,6 +17,7 @@ from fastapi import HTTPException
import litellm
from litellm._logging import verbose_proxy_logger
from litellm.constants import DEFAULT_OPENAI_MODERATIONS_MODEL
from litellm.integrations.custom_logger import CustomLogger
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.guardrails._content_utils import iter_message_text
@ -24,11 +25,9 @@ from litellm.types.utils import CallTypesLiteral
class _ENTERPRISE_OpenAI_Moderation(CustomLogger):
def __init__(self):
self.model_name = (
litellm.openai_moderations_model_name or "text-moderation-latest"
) # pass the model_name you initialized on litellm.Router()
pass
@property
def model_name(self) -> str:
return litellm.openai_moderations_model_name or DEFAULT_OPENAI_MODERATIONS_MODEL
#### CALL HOOKS - proxy only ####

37
enterprise/litellm_enterprise/enterprise_callbacks/llm_guard.py
View file

@ -8,7 +8,7 @@
## This provides an LLM Guard Integration for content moderation on the proxy
import asyncio
from typing import Optional
from typing import Final, Optional
import aiohttp
from fastapi import HTTPException
@ -137,15 +137,20 @@ class _ENTERPRISE_LLMGuard(CustomLogger):
return
self.print_verbose("Makes LLM Guard Check")
if call_type not in [
accepted_call_types: Final = (
"completion",
"acompletion",
"text_completion",
"atext_completion",
"embeddings",
"embedding",
"aembedding",
"image_generation",
"moderation",
"audio_transcription",
]:
"aimage_generation",
)
if call_type not in accepted_call_types:
self.print_verbose(
f"Call Type - {call_type}, not in accepted list - ['completion','embeddings','image_generation','moderation','audio_transcription']"
f"Call Type - {call_type}, not in accepted list - {accepted_call_types}"
)
return data
@ -163,16 +168,14 @@ class _ENTERPRISE_LLMGuard(CustomLogger):
*(self._moderate_message(message) for message in messages)
)
)
return data
input_ = data.get("input")
if input_ is not None:
data["input"] = await self._moderate_input(input_)
return data
data["input"] = await self._moderate_text_or_list(input_)
prompt = data.get("prompt")
if isinstance(prompt, str):
data["prompt"] = await self.moderation_check(text=prompt)
if prompt is not None:
data["prompt"] = await self._moderate_text_or_list(prompt)
return data
async def _moderate_message(self, message: dict) -> dict:
@ -195,17 +198,17 @@ class _ENTERPRISE_LLMGuard(CustomLogger):
return {**part, "text": await self.moderation_check(text=part["text"])}
return part
async def _moderate_input(self, input_: object) -> object:
if isinstance(input_, str):
return await self.moderation_check(text=input_)
if isinstance(input_, list):
async def _moderate_text_or_list(self, value: object) -> object:
if isinstance(value, str):
return await self.moderation_check(text=value)
if isinstance(value, list):
return [
await self.moderation_check(text=item)
if isinstance(item, str)
else item
for item in input_
for item in value
]
return input_
return value
async def async_post_call_streaming_hook(
self, user_api_key_dict: UserAPIKeyAuth, response: str

9
enterprise/litellm_enterprise/enterprise_callbacks/send_emails/endpoints.py
View file

@ -60,6 +60,11 @@ async def _get_email_settings(prisma_client) -> Dict[str, bool]:
async def _save_email_settings(prisma_client, settings: Dict[str, bool]):
"""Helper function to save email settings to general_settings in db"""
from litellm.proxy.proxy_server import proxy_config
proxy_config.reject_config_owned_writes(
section_name="general_settings", changed_keys={"email_settings": settings}
)
try:
verbose_proxy_logger.debug(
f"Saving email settings to general_settings: {settings}"
@ -168,6 +173,8 @@ async def update_event_settings(
await _save_email_settings(prisma_client, settings_dict)
return {"message": "Email event settings updated successfully"}
except HTTPException:
raise
except Exception as e:
verbose_proxy_logger.exception(f"Error updating email settings: {str(e)}")
raise HTTPException(status_code=500, detail=str(e))
@ -197,6 +204,8 @@ async def reset_event_settings(
await _save_email_settings(prisma_client, default_settings)
return {"message": "Email event settings reset to defaults"}
except HTTPException:
raise
except Exception as e:
verbose_proxy_logger.exception(f"Error resetting email settings: {str(e)}")
raise HTTPException(status_code=500, detail=str(e))

43
enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py
View file

@ -11,7 +11,7 @@ Endpoints for /project operations
#### PROJECT MANAGEMENT ####
import json
from collections.abc import Sequence
from collections.abc import Mapping, Sequence
from typing import TYPE_CHECKING, Final
from fastapi import APIRouter, Depends, HTTPException, Request
@ -22,7 +22,11 @@ from litellm._uuid import uuid
from litellm.proxy._types import *
from litellm.proxy.auth.auth_checks import delete_cached_project_object
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.management_endpoints.common_utils import _set_object_metadata_field
from litellm.proxy.management_endpoints.common_utils import (
_is_user_team_admin, # pyright: ignore[reportPrivateUsage] # shared owner of team-admin membership
_set_object_metadata_field,
)
from litellm.proxy.management_endpoints.team_admin_field_permissions import team_admin_may_manage_projects
from litellm.proxy.management_helpers.utils import (
management_endpoint_wrapper,
)
@ -82,37 +86,38 @@ async def _check_user_permission_for_project(
user_api_key_dict: UserAPIKeyAuth,
team_id: str | None,
prisma_client: PrismaClient,
general_settings: Mapping[str, object],
require_admin: bool = False,
team_object: LiteLLM_TeamTable | None = None,
) -> bool:
"""
Check if user has permission to manage a project.
Returns True if user is proxy admin or team admin (when team_id provided).
Returns True if user is proxy admin, or a team admin of ``team_id`` when the
``team_admin_editable_team_fields`` setting grants team admins the ``projects`` permission.
If require_admin=True, only proxy admins are allowed.
If team_object is provided, it will be used instead of fetching from DB
(avoids duplicate DB queries when team was already fetched for validation).
"""
is_proxy_admin = user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN
is_proxy_admin: Final = user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN
if require_admin:
if require_admin or is_proxy_admin:
return is_proxy_admin
if is_proxy_admin:
return True
if not team_id or not user_api_key_dict.user_id:
if not team_id or not user_api_key_dict.user_id or not team_admin_may_manage_projects(general_settings):
return False
team = team_object
if team is None:
team = await _team_table(prisma_client).find_unique(where={"team_id": team_id})
team_row: Final = (
team_object
if team_object is not None
else await _team_table(prisma_client).find_unique(where={"team_id": team_id})
)
if team_row is None:
return False
if team and team.admins:
return user_api_key_dict.user_id in team.admins
return False
team: Final = LiteLLM_TeamTable.model_validate(team_row.model_dump())
return _is_user_team_admin(user_api_key_dict, team) or user_api_key_dict.user_id in (team.admins or [])
async def _validate_team_exists(
@ -531,6 +536,7 @@ async def new_project(
user_api_key_dict=user_api_key_dict,
team_id=data.team_id,
prisma_client=prisma_client,
general_settings=general_settings,
team_object=LiteLLM_TeamTable.model_validate(team_object.model_dump()),
)
@ -735,6 +741,7 @@ async def update_project(
user_api_key_dict=user_api_key_dict,
team_id=existing_project.team_id,
prisma_client=prisma_client,
general_settings=general_settings,
)
if not has_permission:
@ -751,6 +758,7 @@ async def update_project(
user_api_key_dict=user_api_key_dict,
team_id=data.team_id,
prisma_client=prisma_client,
general_settings=general_settings,
team_object=(
LiteLLM_TeamTable.model_validate(target_team_obj.model_dump()) if target_team_obj else None
),
@ -877,7 +885,7 @@ async def delete_project(
}'
```
"""
from litellm.proxy.proxy_server import premium_user, prisma_client, user_api_key_cache
from litellm.proxy.proxy_server import general_settings, premium_user, prisma_client, user_api_key_cache
try:
if not premium_user:
@ -899,6 +907,7 @@ async def delete_project(
user_api_key_dict=user_api_key_dict,
team_id=None,
prisma_client=prisma_client,
general_settings=general_settings,
require_admin=True,
)

4
enterprise/pyproject.toml
View file

@ -1,6 +1,6 @@
[project]
name = "litellm-enterprise"
version = "0.1.68"
version = "0.1.69"
description = "Package for LiteLLM Enterprise features"
readme = "README.md"
requires-python = ">=3.9"
@ -26,7 +26,7 @@ required-version = ">=0.10.9"
module-root = ""
[tool.commitizen]
version = "0.1.68"
version = "0.1.69"
version_files = [
"pyproject.toml:^version",
"../pyproject.toml:litellm-enterprise==",

1
gateway/routes/allowlist.py
View file

@ -95,6 +95,7 @@ GATEWAY_PATH_PREFIXES: tuple[str, ...] = (
"/vertex-ai/",
"/assemblyai/",
"/eu.assemblyai/",
"/deepgram/",
"/langfuse/",
"/vllm/",
"/mistral/",

35
litellm-proxy-extras/litellm_proxy_extras/migrations/20260915000000_add_daily_global_spend/migration.sql Normal file
View file

@ -0,0 +1,35 @@
-- CreateTable
CREATE TABLE IF NOT EXISTS "LiteLLM_DailyGlobalSpend" (
"id" TEXT NOT NULL,
"date" TEXT NOT NULL,
"model" TEXT,
"model_group" TEXT,
"custom_llm_provider" TEXT,
"mcp_namespaced_tool_name" TEXT,
"endpoint" TEXT,
"prompt_tokens" BIGINT NOT NULL DEFAULT 0,
"completion_tokens" BIGINT NOT NULL DEFAULT 0,
"cache_read_input_tokens" BIGINT NOT NULL DEFAULT 0,
"cache_creation_input_tokens" BIGINT NOT NULL DEFAULT 0,
"compression_saved_tokens" BIGINT NOT NULL DEFAULT 0,
"compression_savings_spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0,
"prompt_caching_savings_spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0,
"gateway_injected_caching_savings_spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0,
"autorouter_savings_spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0,
"spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0,
"api_requests" BIGINT NOT NULL DEFAULT 0,
"successful_requests" BIGINT NOT NULL DEFAULT 0,
"failed_requests" BIGINT NOT NULL DEFAULT 0,
"total_response_time_ms" BIGINT NOT NULL DEFAULT 0,
"timed_requests" BIGINT NOT NULL DEFAULT 0,
"created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
"updated_at" TIMESTAMP(3) NOT NULL,
CONSTRAINT "LiteLLM_DailyGlobalSpend_pkey" PRIMARY KEY ("id")
);
-- CreateIndex
CREATE INDEX IF NOT EXISTS "LiteLLM_DailyGlobalSpend_date_idx" ON "LiteLLM_DailyGlobalSpend"("date");
-- CreateIndex
CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_DailyGlobalSpend_date_model_model_group_custom_llm__key" ON "LiteLLM_DailyGlobalSpend"("date", "model", "model_group", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint");

31
litellm-proxy-extras/litellm_proxy_extras/schema.prisma
View file

@ -820,6 +820,37 @@ model LiteLLM_DailyUserSpend {
@@index([endpoint])
}
// Key-free daily rollup of LiteLLM_DailyUserSpend, read by the global usage view
model LiteLLM_DailyGlobalSpend {
id String @id @default(uuid())
date String
model String?
model_group String?
custom_llm_provider String?
mcp_namespaced_tool_name String?
endpoint String?
prompt_tokens BigInt @default(0)
completion_tokens BigInt @default(0)
cache_read_input_tokens BigInt @default(0)
cache_creation_input_tokens BigInt @default(0)
compression_saved_tokens BigInt @default(0)
compression_savings_spend Float @default(0.0)
prompt_caching_savings_spend Float @default(0.0)
gateway_injected_caching_savings_spend Float @default(0.0)
autorouter_savings_spend Float @default(0.0)
spend Float @default(0.0)
api_requests BigInt @default(0)
successful_requests BigInt @default(0)
failed_requests BigInt @default(0)
total_response_time_ms BigInt @default(0)
timed_requests BigInt @default(0)
created_at DateTime @default(now())
updated_at DateTime @updatedAt
@@unique([date, model, model_group, custom_llm_provider, mcp_namespaced_tool_name, endpoint])
@@index([date])
}
// Track daily organization spend metrics per model and key
model LiteLLM_DailyOrganizationSpend {
id String @id @default(uuid())

4
litellm-proxy-extras/pyproject.toml
View file

@ -1,6 +1,6 @@
[project]
name = "litellm-proxy-extras"
version = "0.4.99"
version = "0.4.100"
description = "Additional files for the LiteLLM Proxy. Reduces the size of the main litellm package."
readme = "README.md"
requires-python = ">=3.9"
@ -26,7 +26,7 @@ required-version = ">=0.10.9"
module-root = ""
[tool.commitizen]
version = "0.4.99"
version = "0.4.100"
version_files = [
"pyproject.toml:^version",
"../pyproject.toml:litellm-proxy-extras==",

119
litellm-rust/Cargo.lock generated
View file

@ -1960,6 +1960,7 @@ dependencies = [
"aws-smithy-runtime-api",
"aws-types",
"litellm-auth",
"litellm-http",
"moka",
"reqwest 0.12.28",
"serde_json",
@ -1976,6 +1977,7 @@ dependencies = [
"azure_identity",
"litellm-auth",
"moka",
"rstest",
"serde_json",
"sha2 0.10.9",
"strum",
@ -2027,24 +2029,18 @@ dependencies = [
"tokio",
]
[[package]]
name = "litellm-callbacks"
version = "0.1.0"
dependencies = [
"rstest",
"serde_json",
"tokio",
]
[[package]]
name = "litellm-callbacks-legacy"
version = "0.1.0"
dependencies = [
"litellm-callbacks",
"litellm-auth",
"litellm-host",
"litellm-host-python",
"proptest",
"pyo3",
"rstest",
"serde_json",
"strum",
]
[[package]]
@ -2056,8 +2052,10 @@ dependencies = [
"futures-util",
"litellm-auth",
"litellm-auth-aws",
"litellm-callbacks",
"litellm-auth-gcp",
"litellm-core-utils",
"litellm-host",
"litellm-http",
"litellm-llms",
"litellm-types",
"mime_guess",
@ -2110,12 +2108,22 @@ dependencies = [
"tokio",
]
[[package]]
name = "litellm-host"
version = "0.1.0"
dependencies = [
"litellm-auth",
"rstest",
"serde_json",
"tokio",
]
[[package]]
name = "litellm-host-python"
version = "0.1.0"
dependencies = [
"futures-util",
"litellm-callbacks",
"litellm-host",
"pyo3",
"pyo3-async-runtimes",
"pythonize",
@ -2125,6 +2133,24 @@ dependencies = [
"tokio",
]
[[package]]
name = "litellm-http"
version = "0.1.0"
dependencies = [
"http 1.4.2",
"hyper-util",
"litellm-core-utils",
"reqwest 0.12.28",
"rstest",
"rustls 0.23.42",
"serde",
"serde_json",
"thiserror 2.0.19",
"tokio",
"veil",
"webpki-roots",
]
[[package]]
name = "litellm-llms"
version = "0.1.0"
@ -2139,9 +2165,10 @@ dependencies = [
"litellm-auth-aws",
"litellm-auth-azure",
"litellm-auth-gcp",
"litellm-callbacks",
"litellm-core-utils",
"litellm-framing",
"litellm-host",
"litellm-http",
"litellm-types",
"reqwest 0.12.28",
"rstest",
@ -2149,6 +2176,7 @@ dependencies = [
"serde_json",
"serde_path_to_error",
"serde_with",
"strum",
"thiserror 2.0.19",
"time",
"tokio",
@ -2163,9 +2191,12 @@ dependencies = [
"criterion",
"futures-util",
"litellm-auth",
"litellm-auth-gcp",
"litellm-callbacks-legacy",
"litellm-core",
"litellm-core-utils",
"litellm-host-python",
"litellm-http",
"litellm-llms",
"litellm-token-counter",
"litellm-types",
@ -2598,6 +2629,25 @@ dependencies = [
"unicode-ident",
]
[[package]]
name = "proptest"
version = "1.11.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "4b45fcc2344c680f5025fe57779faef368840d0bd1f42f216291f0dc4ace4744"
dependencies = [
"bit-set",
"bit-vec",
"bitflags",
"num-traits",
"rand 0.9.5",
"rand_chacha 0.9.0",
"rand_xorshift",
"regex-syntax",
"rusty-fork",
"tempfile",
"unarray",
]
[[package]]
name = "pyo3"
version = "0.29.2"
@ -2679,6 +2729,12 @@ dependencies = [
"serde",
]
[[package]]
name = "quick-error"
version = "1.2.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "a1d01941d82fa2ab50be1e79e6714289dd7cde78eba4c074bc5a4374f650dfe0"
[[package]]
name = "quinn"
version = "0.11.11"
@ -2842,6 +2898,15 @@ dependencies = [
"rand_core 0.10.1",
]
[[package]]
name = "rand_xorshift"
version = "0.4.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "513962919efc330f829edb2535844d1b912b0fbe2ca165d613e4e8788bb05a5a"
dependencies = [
"rand_core 0.9.5",
]
[[package]]
name = "rayon"
version = "1.12.0"
@ -2981,7 +3046,6 @@ checksum = "eddd3ca559203180a307f12d114c268abf583f59b03cb906fd0b3ff8646c1147"
dependencies = [
"base64 0.22.1",
"bytes",
"futures-channel",
"futures-core",
"futures-util",
"h2 0.4.15",
@ -3241,6 +3305,18 @@ version = "1.0.23"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "cf54715a573b99ac80df0bc206da022bcd442c974952c7b9720069370852e21f"
[[package]]
name = "rusty-fork"
version = "0.3.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "cc6bf79ff24e648f6da1f8d1f011e9cac26491b619e6b9280f2b47f1774e6ee2"
dependencies = [
"fnv",
"quick-error",
"tempfile",
"wait-timeout",
]
[[package]]
name = "ryu"
version = "1.0.23"
@ -4096,6 +4172,12 @@ dependencies = [
"syn 2.0.119",
]
[[package]]
name = "unarray"
version = "0.1.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "eaea85b334db583fe3274d12b4cd1880032beab409c0d774be044d4480ab9a94"
[[package]]
name = "unicase"
version = "2.9.0"
@ -4209,6 +4291,15 @@ version = "0.8.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "5c3082ca00d5a5ef149bb8b555a72ae84c9c59f7250f013ac822ac2e49b19c64"
[[package]]
name = "wait-timeout"
version = "0.2.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "09ac3b126d3914f9849036f826e054cbabdc8519970b8998ddaf3b5bd3c65f11"
dependencies = [
"libc",
]
[[package]]
name = "walkdir"
version = "2.5.0"

9
litellm-rust/Cargo.toml
View file

@ -10,13 +10,14 @@ repository = "https://github.com/BerriAI/litellm"
[workspace.dependencies]
litellm-core = { path = "crates/core" }
litellm-callbacks = { path = "crates/callbacks" }
litellm-host = { path = "crates/host" }
litellm-callbacks-legacy = { path = "crates/callbacks-legacy" }
litellm-framing = { path = "crates/framer" }
litellm-auth = { path = "crates/auth" }
litellm-auth-aws = { path = "crates/auth-aws" }
litellm-auth-azure = { path = "crates/auth-azure" }
litellm-auth-gcp = { path = "crates/auth-gcp" }
litellm-http = { path = "crates/http" }
litellm-llms = { path = "crates/llms" }
litellm-types = { path = "crates/types" }
litellm-core-utils = { path = "crates/core-utils" }
@ -26,11 +27,14 @@ litellm-token-counter = { path = "crates/token-counter" }
litellm-host-python = { path = "crates/host-python" }
bytes = "1"
http = "1"
hyper-util = { version = "0.1.20", default-features = false, features = ["client-proxy"] }
proptest = "1.7.0"
pyo3 = "0.29.2"
pyo3-async-runtimes = { version = "0.29.0", features = ["tokio-runtime"] }
pythonize = "0.29.0"
rand = "0.8"
reqwest = { version = "0.12", default-features = false, features = ["blocking", "json", "multipart", "rustls-tls", "http2", "stream"] }
reqwest = { version = "0.12", default-features = false, features = ["json", "multipart", "rustls-tls", "http2", "stream"] }
rstest = "0.26.1"
rstest_reuse = "0.7.0"
rustls = { version = "0.23", default-features = false, features = ["ring", "std", "tls12"] }
@ -48,6 +52,7 @@ base64 = "0.22"
moka = { version = "0.12.16", features = ["future"] }
strum = { version = "0.28.0", features = ["derive"] }
url = "2.5.8"
webpki-roots = "1"
time = { version = "0.3.53", features = ["parsing"] }
criterion = "0.8.2"
fancy-regex = "0.19.2"

10
litellm-rust/clippy.toml Normal file
View file

@ -0,0 +1,10 @@
# The Tokio runtime is reached only through `host-python/src/execution.rs`, whose fork gate
# must see every entry. Going around it makes a fork-after-use hang instead of raising.
disallowed-methods = [
{ path = "pyo3_async_runtimes::tokio::get_runtime", reason = "use litellm_host_python::run_sync / run_sync_value" },
{ path = "pyo3_async_runtimes::tokio::future_into_py", reason = "use litellm_host_python::run_async / run_async_value" },
{ path = "pyo3_async_runtimes::tokio::future_into_py_with_locals", reason = "use litellm_host_python::run_async / run_async_value" },
{ path = "pyo3_async_runtimes::tokio::local_future_into_py", reason = "use litellm_host_python::run_async / run_async_value" },
{ path = "pyo3_async_runtimes::tokio::run", reason = "use litellm_host_python::run_sync / run_sync_value" },
{ path = "pyo3_async_runtimes::tokio::run_until_complete", reason = "use litellm_host_python::run_sync / run_sync_value" },
]

1
litellm-rust/crates/auth-aws/Cargo.toml
View file

@ -7,6 +7,7 @@ repository.workspace = true
[dependencies]
litellm-auth.workspace = true
litellm-http.workspace = true
moka = { workspace = true, features = ["sync"] }
serde_json.workspace = true

73
litellm-rust/crates/auth-aws/src/aws.rs
View file

@ -20,8 +20,7 @@ use super::constants::{
AWS_ACCESS_KEY_ID, AWS_EXTERNAL_ID, AWS_PROFILE_NAME, AWS_REGION, AWS_REGION_NAME,
AWS_ROLE_ARN, AWS_ROLE_NAME, AWS_SECRET_ACCESS_KEY, AWS_SESSION_NAME, AWS_SESSION_TOKEN,
AWS_SIGNED_HEADER_NAMES, AWS_STS_ENDPOINT, AWS_WEB_IDENTITY_TOKEN, AWS_WEB_IDENTITY_TOKEN_FILE,
BEDROCK_SERVICE, DEFAULT_BEDROCK_REGION, DEFAULT_SESSION_NAME_PREFIX,
SIGV4_COMPUTED_HEADER_NAMES,
DEFAULT_BEDROCK_REGION, DEFAULT_SESSION_NAME_PREFIX, SIGV4_COMPUTED_HEADER_NAMES,
};
const STATIC_CREDENTIALS_TTL: Duration = Duration::from_secs(3600 - 60);
@ -451,11 +450,12 @@ pub fn is_sigv4_computed_header(name: &str) -> bool {
SIGV4_COMPUTED_HEADER_NAMES.contains(&name.to_ascii_lowercase().as_str())
}
pub fn sign_bedrock_post(
pub fn sign_post(
url: &str,
body: &[u8],
headers: &BTreeMap<String, String>,
region: &str,
service: &str,
credentials: &Credentials,
signing_time: SystemTime,
) -> Result<BTreeMap<String, String>, Error> {
@ -463,7 +463,7 @@ pub fn sign_bedrock_post(
let params = v4::SigningParams::builder()
.identity(&identity)
.region(region)
.name(BEDROCK_SERVICE)
.name(service)
.time(signing_time)
.settings(SigningSettings::default())
.build()
@ -534,22 +534,28 @@ fn is_bedrock_region(value: &str) -> bool {
.all(|char| char.is_ascii_alphanumeric() || char == '-')
}
/// The region a caller configured: `aws_region_name`, then the model's own
/// region, then the environment. Each service decides what a missing one means.
pub fn resolve_aws_region(
model_region: Option<&str>,
optional_params: &Map<String, Value>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Option<String> {
optional_params
.get("aws_region_name")
.and_then(Value::as_str)
.or(model_region)
.map(str::to_string)
.or_else(|| env_lookup(AWS_REGION_NAME))
.or_else(|| env_lookup(AWS_REGION))
}
pub fn resolve_bedrock_region(
model_region: Option<&str>,
optional_params: &Map<String, Value>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> String {
if let Some(region) = optional_params
.get("aws_region_name")
.and_then(Value::as_str)
{
return region.to_string();
}
if let Some(region) = model_region {
return region.to_string();
}
env_lookup(AWS_REGION_NAME)
.or_else(|| env_lookup(AWS_REGION))
resolve_aws_region(model_region, optional_params, env_lookup)
.unwrap_or_else(|| DEFAULT_BEDROCK_REGION.to_string())
}
@ -609,11 +615,36 @@ pub fn host_supplied_credentials(optional_params: &Map<String, Value>) -> Option
#[cfg(test)]
mod tests {
use super::*;
use crate::constants::BEDROCK_SERVICE;
fn no_env(_: &str) -> Option<String> {
None
}
#[test]
fn a_region_comes_from_the_call_then_the_model_then_the_environment() {
let params = Map::from_iter([("aws_region_name".to_string(), Value::from("eu-west-1"))]);
let region_name = |key: &str| (key == AWS_REGION_NAME).then(|| "ap-south-1".to_string());
let region = |key: &str| (key == AWS_REGION).then(|| "sa-east-1".to_string());
let resolved = [
resolve_aws_region(Some("us-east-2"), &params, &region_name),
resolve_aws_region(Some("us-east-2"), &Map::new(), &region_name),
resolve_aws_region(None, &Map::new(), &region_name),
resolve_aws_region(None, &Map::new(), &region),
resolve_aws_region(None, &Map::new(), &no_env),
];
assert_eq!(
resolved.map(|region| region.unwrap_or_else(|| "none".into())),
["eu-west-1", "us-east-2", "ap-south-1", "sa-east-1", "none"]
);
assert_eq!(
resolve_bedrock_region(None, &Map::new(), &no_env),
DEFAULT_BEDROCK_REGION
);
}
fn parity_inputs() -> (String, Vec<u8>, BTreeMap<String, String>) {
(
"https://bedrock-runtime.us-east-1.amazonaws.com/model/amazon.titan-text-express-v1/invoke"
@ -811,11 +842,12 @@ mod tests {
None,
"test",
);
let signed = sign_bedrock_post(
let signed = sign_post(
&url,
&body,
&signable,
"us-east-1",
BEDROCK_SERVICE,
&credentials,
SystemTime::UNIX_EPOCH,
)
@ -843,11 +875,12 @@ mod tests {
None,
"test",
);
let signed = sign_bedrock_post(
let signed = sign_post(
&url,
&body,
&headers,
"us-east-1",
BEDROCK_SERVICE,
&credentials,
UNIX_EPOCH + std::time::Duration::from_secs(1_704_164_645),
)
@ -878,11 +911,12 @@ mod tests {
None,
"test",
);
let signed = sign_bedrock_post(
let signed = sign_post(
&url,
&body,
&headers,
"us-east-1",
BEDROCK_SERVICE,
&credentials,
UNIX_EPOCH + std::time::Duration::from_secs(1_704_164_645),
)
@ -915,11 +949,12 @@ mod tests {
let url = format!(
"https://bedrock-runtime.{region}.amazonaws.com/model/us.anthropic.claude-opus-4-8/invoke"
);
let signed_headers = sign_bedrock_post(
let signed_headers = sign_post(
&url,
&body,
&headers,
region,
BEDROCK_SERVICE,
&credentials,
SystemTime::now(),
)?;

3
litellm-rust/crates/auth-aws/src/lib.rs
View file

@ -1,6 +1,9 @@
mod aws;
pub mod constants;
mod error;
mod signer;
pub use aws::*;
pub use aws_credential_types::Credentials;
pub use error::Error;
pub use signer::SigV4Signer;

178
litellm-rust/crates/auth-aws/src/signer.rs Normal file
View file

@ -0,0 +1,178 @@
use std::{collections::BTreeMap, time::SystemTime};
use aws_credential_types::Credentials;
use litellm_http::outbound::{RequestSigner, UnsignedRequest};
use serde_json::{Map, Value};
use crate::{
Error, aws_auth_config, aws_signature_headers, host_supplied_credentials,
is_sigv4_computed_header, resolve_credentials, sign_post,
};
#[derive(Clone, Debug)]
pub struct SigV4Signer {
region: String,
service: &'static str,
credentials: Credentials,
clock: fn() -> SystemTime,
}
impl SigV4Signer {
pub fn new(region: String, service: &'static str, credentials: Credentials) -> Self {
Self {
region,
service,
credentials,
clock: SystemTime::now,
}
}
pub fn with_clock(self, clock: fn() -> SystemTime) -> Self {
Self { clock, ..self }
}
pub async fn resolve(
region: String,
service: &'static str,
optional_params: &Map<String, Value>,
env_lookup: &(dyn Fn(&str) -> Option<String> + Sync),
) -> Result<Self, Error> {
let credentials = match host_supplied_credentials(optional_params) {
Some(credentials) => credentials,
None => {
resolve_credentials(aws_auth_config(optional_params, env_lookup), env_lookup)
.await?
}
};
Ok(Self::new(region, service, credentials))
}
}
impl RequestSigner for SigV4Signer {
fn sign(
&self,
request: UnsignedRequest<'_>,
) -> Result<Vec<(String, String)>, litellm_http::Error> {
if let Some((name, _)) = request
.headers
.iter()
.find(|(name, _)| is_sigv4_computed_header(name))
{
return Err(litellm_http::Error::ComputedHeader(name.clone()));
}
let headers: BTreeMap<String, String> = request.headers.iter().cloned().collect();
sign_post(
request.url,
request.body,
&aws_signature_headers(&headers),
&self.region,
self.service,
&self.credentials,
(self.clock)(),
)
.map(|signature| signature.into_iter().collect())
.map_err(|error| litellm_http::Error::Signature(error.to_string()))
}
}
#[cfg(test)]
mod tests {
use std::time::{Duration, UNIX_EPOCH};
use litellm_http::outbound::OutboundRequest;
use serde_json::json;
use super::*;
fn fixed_clock() -> SystemTime {
UNIX_EPOCH + Duration::from_secs(1_700_000_000)
}
fn signer(service: &'static str) -> SigV4Signer {
SigV4Signer::new(
"us-east-1".into(),
service,
Credentials::new("AKIDEXAMPLE", "secret", None, None, "test"),
)
.with_clock(fixed_clock)
}
fn authorization(body: &Value, service: &'static str) -> String {
OutboundRequest::signed_json(
"https://textract.us-east-1.amazonaws.com/".into(),
vec![("X-Amz-Target".into(), "Textract.DetectDocumentText".into())],
body,
None,
&signer(service),
)
.unwrap()
.header("Authorization")
.unwrap()
.to_string()
}
#[test]
fn the_signature_verifies_against_the_bytes_that_are_sent() {
let sent = OutboundRequest::signed_json(
"https://textract.us-east-1.amazonaws.com/".into(),
vec![("X-Amz-Target".into(), "Textract.DetectDocumentText".into())],
&json!({"Document": {"Bytes": "aGk="}}),
None,
&signer("textract"),
)
.unwrap();
let unsigned: BTreeMap<String, String> = sent
.headers()
.iter()
.filter(|(name, _)| !is_sigv4_computed_header(name))
.cloned()
.collect();
let recomputed = sign_post(
sent.url(),
sent.body(),
&aws_signature_headers(&unsigned),
"us-east-1",
"textract",
&Credentials::new("AKIDEXAMPLE", "secret", None, None, "test"),
fixed_clock(),
)
.unwrap();
assert_eq!(
sent.header("Authorization"),
Some(recomputed["Authorization"].as_str())
);
}
#[test]
fn the_signature_depends_on_the_body_and_the_service() {
let original = authorization(&json!({"text": "card 4111"}), "textract");
assert_ne!(
original,
authorization(&json!({"text": "card [REDACTED]"}), "textract")
);
assert_ne!(
original,
authorization(&json!({"text": "card 4111"}), "bedrock")
);
assert!(original.contains("/us-east-1/textract/aws4_request"));
}
#[test]
fn a_forwarded_computed_header_is_refused_instead_of_sent_twice() {
let error = OutboundRequest::signed_json(
"https://textract.us-east-1.amazonaws.com/".into(),
vec![("authorization".into(), "Bearer caller".into())],
&json!({}),
None,
&signer("textract"),
)
.unwrap_err();
assert_eq!(
error,
litellm_http::Error::ComputedHeader("authorization".into())
);
}
}

1
litellm-rust/crates/auth-azure/Cargo.toml
View file

@ -18,4 +18,5 @@ azure_core = "1.0.0"
azure_identity = { version = "1.0.0", features = ["tokio"] }
[dev-dependencies]
rstest.workspace = true
tokio.workspace = true

48
litellm-rust/crates/auth-azure/src/types.rs
View file

@ -1,11 +1,10 @@
use serde_json::{Map, Value};
use std::collections::BTreeMap;
use strum::EnumString;
use litellm_auth::Error;
use litellm_auth::{
CredentialResolverHandle, InputSource, SecretValue, Sourced, TokenProviderHandle,
CredentialResolverHandle, Error, InputSource, SecretValue, Sourced, TokenProviderHandle,
};
use serde_json::{Map, Value};
use strum::EnumString;
pub const DEFAULT_AZURE_SCOPE: &str = "https://cognitiveservices.azure.com/.default";
@ -52,6 +51,16 @@ pub struct AzureAuthInputs {
}
impl AzureAuthInputs {
pub fn or_configured_token_refresh(self, enabled: bool) -> Self {
if *self.enable_azure_ad_token_refresh.value() || !enabled {
return self;
}
Self {
enable_azure_ad_token_refresh: Sourced::new(true, InputSource::Deployment),
..self
}
}
#[cfg(test)]
pub fn from_optional_params(params: &Map<String, Value>) -> Result<Self, Error> {
Self::from_sourced_optional_params(params, &BTreeMap::new())
@ -115,12 +124,12 @@ fn source_for(sources: &BTreeMap<String, InputSource>, name: &str) -> InputSourc
#[cfg(test)]
mod tests {
use serde_json::json;
use std::collections::BTreeMap;
use super::{AzureAuthInputs, AzureCredentialType, ConfigValue};
use litellm_auth::{InputSource, Sourced};
use serde_json::json;
use super::{AzureAuthInputs, AzureCredentialType, ConfigValue};
#[test]
fn selector_parsing_is_exact() {
@ -189,4 +198,29 @@ mod tests {
assert!(!debug.contains("token-value"));
assert!(!debug.contains("secret-value"));
}
#[rstest::rstest]
#[case::global_turns_refresh_on(json!({}), true, true, InputSource::Deployment)]
#[case::global_overrides_a_call_false_like_python(json!({"enable_azure_ad_token_refresh": false}), true, true, InputSource::Deployment)]
#[case::call_true_survives_a_global_false(json!({"enable_azure_ad_token_refresh": true}), false, true, InputSource::Request)]
#[case::both_off(json!({}), false, false, InputSource::Request)]
fn token_refresh_follows_the_configured_global(
#[case] params: serde_json::Value,
#[case] global: bool,
#[case] enabled: bool,
#[case] source: InputSource,
) {
let sources = BTreeMap::from([(
"enable_azure_ad_token_refresh".to_string(),
InputSource::Request,
)]);
let inputs =
AzureAuthInputs::from_sourced_optional_params(params.as_object().unwrap(), &sources)
.unwrap()
.or_configured_token_refresh(global);
assert_eq!(
inputs.enable_azure_ad_token_refresh,
Sourced::new(enabled, source)
);
}
}

47
litellm-rust/crates/auth-gcp/src/lib.rs
View file

@ -1,17 +1,13 @@
use std::collections::BTreeMap;
use std::future::Future;
use std::path::Path;
use std::pin::Pin;
use std::sync::Arc;
use std::{collections::BTreeMap, future::Future, path::Path, pin::Pin, sync::Arc};
use gcp_auth::{CustomServiceAccount, TokenProvider};
use litellm_auth::{
CredentialPlacement, Error, InputSource, SecretValue, Sourced, http::apply_credential,
};
use moka::future::Cache;
use serde_json::{Map, Value};
use sha2::{Digest, Sha256};
use litellm_auth::http::apply_credential;
use litellm_auth::{CredentialPlacement, Error, InputSource, SecretValue, Sourced};
const CLOUD_PLATFORM_SCOPE: &str = "https://www.googleapis.com/auth/cloud-platform";
const GOOGLE_OAUTH_TOKEN_ENDPOINT: &str = "https://oauth2.googleapis.com/token";
const GOOGLE_APPLICATION_CREDENTIALS_ENV: &str = "GOOGLE_APPLICATION_CREDENTIALS";
@ -45,6 +41,16 @@ impl VertexConfig {
})
}
pub fn or_configured(self, project_id: Option<&str>, location: Option<&str>) -> Self {
let configured =
|value: Option<&str>| value.filter(|value| !value.is_empty()).map(str::to_string);
Self {
project_id: self.project_id.or_else(|| configured(project_id)),
location: self.location.or_else(|| configured(location)),
..self
}
}
pub fn project_id(&self) -> Option<&str> {
self.project_id.as_deref()
}
@ -571,4 +577,29 @@ mod tests {
assert_eq!(loads.load(Ordering::SeqCst), 1);
assert_eq!(calls.load(Ordering::SeqCst), 4);
}
#[test]
fn configured_defaults_sit_between_call_params_and_the_environment() {
let env = |name: &str| Some(format!("env-{name}"));
let from_config =
VertexConfig::default().or_configured(Some("global-project"), Some("global-location"));
assert_eq!(
get_vertex_ai_project(&from_config, &env).as_deref(),
Some("global-project")
);
assert_eq!(
get_vertex_ai_location(&from_config, &env).as_deref(),
Some("global-location")
);
let from_call =
config(json!({"vertex_project":"call-project","vertex_location":"call-location"}))
.or_configured(Some("global-project"), Some("global-location"));
assert_eq!(from_call.project_id(), Some("call-project"));
assert_eq!(from_call.location(), Some("call-location"));
let empty_global = VertexConfig::default().or_configured(Some(""), None);
assert_eq!(
get_vertex_ai_project(&empty_global, &env).as_deref(),
Some("env-VERTEXAI_PROJECT")
);
}
}

19
litellm-rust/crates/auth/src/http.rs
View file

@ -40,13 +40,22 @@ pub fn apply_credential(
)
}
/// How the upstream call is authenticated. API-key strategies are resolved in
/// `prepare`; SigV4 needs the serialized body, so the handler signs it.
/// How the upstream call is authenticated. API-key strategies become headers
/// in `prepare`; SigV4 covers the serialized body, so it is applied where the
/// outbound request is built.
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum RequestAuth {
Header { name: &'static str, value: String },
Bearer { token: String },
AwsSigV4 { region: String },
Header {
name: &'static str,
value: String,
},
Bearer {
token: String,
},
AwsSigV4 {
region: String,
service: &'static str,
},
}
#[cfg(test)]

4
litellm-rust/crates/auth/src/secret.rs
View file

@ -1,6 +1,8 @@
use serde::Deserialize;
use veil::Redact;
#[derive(Redact, Clone)]
#[derive(Redact, Clone, Deserialize)]
#[serde(transparent)]
pub struct SecretValue(#[redact(with = "[REDACTED]")] String);
impl SecretValue {

14
litellm-rust/crates/callbacks-legacy/AGENTS.md
View file

@ -1,15 +1,17 @@
- Target invariants, not completion claims
- Keep this crate the legacy `@client` wrapper as the native call sees it, and nothing else: the `Logging` contract (`function_setup`, the deployment hooks, `pre_call`/`post_call`, the sync and async success and failure fan-out, the deferred proxy release, the argument sharing those callbacks rely on) plus the kwargs rewrites the wrapper makes on the way in (credential-name inheritance, the budget and retry-count limits)
- The driver in `litellm-host-python`, the routes and core see one `CallbackAdapter`; they never learn which Python objects consume a call
- The driver in `litellm-host-python`, the routes and core see one `PythonLifecycle`; they never learn which Python objects consume a call
- Rust drives the call; every litellm Python internal it still borrows is a variant of `LegacyPython`, grouped by subsystem (`Wrapper`, `Logging`, `DeploymentHooks`)
- The enum only shrinks: when Rust owns a subsystem, delete its group rather than adding a Rust path beside it
- Calling a user's own callback directly is permanent Python surface and gets its own type outside `LegacyPython`
- `PublicCall` is the caller's call as `Logging` sees it: the positional arguments, the keyword view as the legacy path rewrites it (setup, deployment hook, prepare) and the bound request object whose attributes back keywords the caller omitted; routes hand it over through `run_legacy_call` and keep no copy
- `setup` decides once who owns the `Logging` instance and returns it as `CallSetup.bridge_owned`; `PythonLogger` carries it and nothing on the instance records it
- A logger the caller passed as `litellm_logging_obj` is caller-owned and observed in full, because the caller reads it after the call; the proxy is the live case
- A logger `function_setup` built for this call is bridge-owned, so each fan-out phase is skipped when `callbacks_needed` finds no registry, dynamic callback, `logger_fn` or debug switch for it; cost, timing and response metadata still run
- `setup` reuses a `Logging` the caller passed as `litellm_logging_obj` (the proxy and Router are the live cases) and otherwise builds one through `function_setup`, as `@client` does
- Either way every phase calls the same `Logging` method the Python path calls; which callbacks run is `Logging`'s decision, never this crate's
- Callbacks receive the caller's own objects and may mutate them; this crate alone carries that obligation
- Retain complete boundary arguments, opaque unknown values, aliases, omitted/default distinctions and deliberate copies; preserve the established deployment-hook kwargs view
- Re-alias every `passthrough_fields` body key to the caller's object before `pre_call`; a keyword wins over the request attribute even when it is an explicit `None`
- Before `pre_call`, re-alias every body key whose value equals the caller's argument to the caller's own object; this crate compares the two itself, and the argument is resolved by `litellm_host_python::lookup`
- Retain independently captured body/header roots from `pre_call` to `post_call`; in-place mutation reaches the wire, envelope field replacement is visible to later callbacks only
- A later kind of callback host (WASM, in-process Rust) has none of these obligations, so they stay out of `litellm-callbacks`, `litellm-host-python` and the bridge; the only facts that cross from the route are the prepared keyword view and `RequestContext.passthrough_fields`
- A later kind of callback host (WASM, in-process Rust) has none of these obligations, so they stay out of `litellm-host`, `litellm-host-python` and the bridge; the only fact that crosses from the route is the prepared keyword view
- Success and failure handlers receive the exact selected public response or exception; logging projections, redaction and snapshots keep their own copy contracts
- Ordinary failure-handler errors cannot suppress the other eligible family or replace the mapped provider error; a cancellation ends the call with no further dispatch
- Dispatch errors never replay provider work or trigger the opposite outcome; the proxy's acceptance or rejection releases deferred success at most once

7
litellm-rust/crates/callbacks-legacy/Cargo.toml
View file

@ -7,10 +7,15 @@ repository.workspace = true
autotests = false
[dependencies]
litellm-callbacks.workspace = true
litellm-host.workspace = true
litellm-host-python.workspace = true
pyo3.workspace = true
strum.workspace = true
serde_json.workspace = true
[dev-dependencies]
litellm-auth.workspace = true
proptest.workspace = true
rstest.workspace = true
serde_json.workspace = true

118
litellm-rust/crates/callbacks-legacy/python_contract.json Normal file
View file

@ -0,0 +1,118 @@
{
"setup": [
"call_type",
"args",
"kwargs",
"start_time",
"asynchronous"
],
"check_limits": [
"kwargs"
],
"finalize": [
"response",
"logger",
"kwargs",
"start_time",
"end_time"
],
"update_logging": [
"logger",
"kwargs",
"model",
"optional_params",
"litellm_params",
"custom_llm_provider"
],
"pre_call": [
"logger",
"input",
"api_key",
"additional_args"
],
"post_call": [
"logger",
"original_response",
"api_key",
"additional_args"
],
"defers_async_logging": [
"logger"
],
"defer_success": [
"logger",
"pending"
],
"sync_success_for_async_call": [
"logger",
"response",
"start",
"end"
],
"failure_handler": [
"logger",
"error",
"start",
"end",
"asynchronous"
],
"submit_success": [
"logger",
"response",
"start",
"end"
],
"async_success_handler": [
"logger",
"response",
"start",
"end"
],
"enqueue_logging": [
"coroutine"
],
"restore_context": [
"logger"
],
"custom_pricing_fields": [],
"is_internal_call": [],
"credential_list": [],
"warn_unknown_credential": [
"name",
"loaded"
],
"before_deployment_call": [
"kwargs",
"call_type"
],
"after_deployment_success": [
"kwargs",
"response",
"call_type"
],
"after_deployment_failure": [
"kwargs",
"error",
"call_type"
],
"stream_opened": [
"logger"
],
"stream_success": [
"logger",
"url_route",
"endpoint_type",
"request_body",
"chunks",
"start",
"end",
"first_chunk"
],
"stream_failure": [
"logger",
"endpoint_type",
"request_body",
"chunks",
"error"
]
}

246
litellm-rust/crates/callbacks-legacy/src/adapter.rs
View file

@ -2,21 +2,26 @@
//! raises is answered with the same `Logging` calls, in the same order, as the Python
//! `@client` path makes them.
use litellm_callbacks::event::{CallEvent, FailureOrigin, RequestContext, Timing, WireRequest};
use litellm_host::event::{
FailureOrigin, MachineEvent, RequestContext, Timing, WireRequest, epoch_seconds,
};
use litellm_host_python::{
AdapterStep, CallbackAdapter, PublicValue, from_py, missing_state, to_py,
LifecycleEvent, LifecycleStep, PythonLifecycle, from_py, missing_state, to_py,
};
use pyo3::{
exceptions::{PyBaseException, PyException},
gc::{PyTraverseError, PyVisit},
prelude::*,
types::PyDict,
types::{PyDateTime, PyDict, PyList},
};
use serde_json::Value;
use crate::{
DeploymentHooks, LegacyCallbacks, PublicCall, PythonLogger,
deferred::{PendingLogging, PendingSuccess},
finalize, is_internal_call, prepare, setup,
finalize, is_internal_call,
legacy_python::Streaming,
prepare, setup,
};
/// What the legacy contract needs to know about the route it is logging.
@ -25,6 +30,22 @@ pub struct LegacySurface {
pub call_type: &'static str,
/// What `Logging.pre_call` is told the input was.
pub input_description: &'static str,
/// How a streamed response is billed; `None` for a route that never streams.
pub stream: Option<PassThroughStream>,
}
/// The pass-through billing a streamed response goes through once its chunks are in.
#[derive(Clone, Copy, Debug)]
pub struct PassThroughStream {
pub url_route: &'static str,
/// A value of Python's `EndpointType`.
pub endpoint_type: &'static str,
}
/// What the Messages stream iterator keeps for its end-of-stream billing.
struct DeliveredStream {
chunks: Py<PyList>,
first_chunk: Option<Py<PyAny>>,
}
enum Pending {
@ -44,16 +65,15 @@ pub struct LegacyLogging {
error: Option<Py<PyBaseException>>,
body: Option<Py<PyDict>>,
headers: Option<Py<PyDict>>,
context: Option<RequestContext>,
stream: Option<DeliveredStream>,
asynchronous: bool,
internal: bool,
pending: Option<Pending>,
}
fn datetime(py: Python<'_>, epoch_seconds: f64) -> PyResult<Py<PyAny>> {
py.import("datetime")?
.getattr("datetime")?
.call_method1("fromtimestamp", (epoch_seconds,))
.map(Bound::unbind)
PyDateTime::from_timestamp(py, epoch_seconds, None).map(|value| value.into_any().unbind())
}
fn is_cancellation(py: Python<'_>, error: &PyErr) -> bool {
@ -77,6 +97,8 @@ impl LegacyLogging {
error: None,
body: None,
headers: None,
context: None,
stream: None,
asynchronous,
internal: false,
pending: None,
@ -85,8 +107,8 @@ impl LegacyLogging {
/// Deployment hooks are awaited, and Python's synchronous `@client` wrapper never
/// runs them.
fn deployment_hooks(&self, py: Python<'_>) -> PyResult<bool> {
Ok(self.asynchronous && DeploymentHooks::needed(py)?)
fn runs_deployment_hooks(&self) -> bool {
self.asynchronous
}
fn logger(&self) -> PyResult<&PythonLogger> {
@ -95,13 +117,13 @@ impl LegacyLogging {
})
}
fn prepare(&mut self, py: Python<'_>) -> PyResult<AdapterStep> {
fn prepare(&mut self, py: Python<'_>) -> PyResult<LifecycleStep> {
let prepared = prepare(py, self.call.kwargs().bind(py), self.logger()?)?.unbind();
self.call.set_kwargs(prepared);
Ok(AdapterStep::Arguments(self.call.kwargs().clone_ref(py)))
Ok(LifecycleStep::Arguments(self.call.kwargs().clone_ref(py)))
}
fn finalize(&mut self, py: Python<'_>) -> PyResult<AdapterStep> {
fn finalize(&mut self, py: Python<'_>) -> PyResult<LifecycleStep> {
finalize(
py,
&self.response,
@ -112,7 +134,7 @@ impl LegacyLogging {
)?;
self.response
.as_ref()
.map(|response| AdapterStep::Response(response.clone_ref(py)))
.map(|response| LifecycleStep::Response(response.clone_ref(py)))
.ok_or_else(missing_state)
}
@ -145,9 +167,7 @@ impl LegacyLogging {
.get_item("fallbacks")?
.is_none_or(|value| value.is_none())
{
if !logger.callbacks_needed(py, "async_success")? {
logger.success_bookkeeping(py, &self.response, &self.start, &self.end, true)?;
} else if logger.defers_async_logging(py) {
if logger.defers_async_logging(py) {
let pending = Py::new(
py,
PendingLogging {
@ -162,15 +182,72 @@ impl LegacyLogging {
logger.sync_success_for_async_call(py, &self.response, &self.start, &self.end)
}
fn stream_success(&self, py: Python<'_>, stream: &DeliveredStream) -> PyResult<()> {
let logger = self.logger()?;
let billing = self.surface.stream.ok_or_else(missing_state)?;
let billed = Streaming::Success.call(
py,
(
logger.object(py),
billing.url_route,
billing.endpoint_type,
&self.body,
&stream.chunks,
&self.start,
&self.end,
&stream.first_chunk,
),
);
match billed {
Err(error) if error.is_instance_of::<PyException>(py) => {
error.write_unraisable(py, Some(logger.object(py)));
Ok(())
}
result => result.map(|_| ()),
}
}
/// A failure after the stream reached the caller bills the delivered chunks as
/// partial usage. The sync path has no loop to schedule that on, so it falls back to
/// the plain failure handler.
fn stream_failure(&mut self, py: Python<'_>) -> PyResult<LifecycleStep> {
let (Some(logger), Some(error), Some(stream), Some(billing)) =
(&self.logger, &self.error, &self.stream, self.surface.stream)
else {
return Ok(LifecycleStep::Done);
};
if !self.asynchronous {
return self.dispatch_failure(py);
}
let scheduled = Streaming::Failure.call(
py,
(
logger.object(py),
billing.endpoint_type,
&self.body,
&stream.chunks,
error,
),
);
match scheduled {
Ok(awaitable) => {
self.pending = Some(Pending::AsyncFailure);
Ok(LifecycleStep::Await(awaitable.unbind()))
}
Err(failure) if is_cancellation(py, &failure) => Err(failure),
Err(_) => Ok(LifecycleStep::Done),
}
}
/// The sync failure handler, then the async one for async calls. Ordinary handler
/// errors never replace the selected failure or suppress the other family; a
/// cancellation does end the call.
fn dispatch_failure(&mut self, py: Python<'_>) -> PyResult<AdapterStep> {
fn dispatch_failure(&mut self, py: Python<'_>) -> PyResult<LifecycleStep> {
let (Some(logger), Some(error)) = (&self.logger, &self.error) else {
return Ok(AdapterStep::Done);
return Ok(LifecycleStep::Done);
};
if self.asynchronous && self.internal {
return Ok(AdapterStep::Done);
return Ok(LifecycleStep::Done);
}
if let Err(failure) = logger.failure(py, error, &self.start, &self.end, false)
&& is_cancellation(py, &failure)
@ -178,27 +255,27 @@ impl LegacyLogging {
return Err(failure);
}
if !self.asynchronous {
return Ok(AdapterStep::Done);
return Ok(LifecycleStep::Done);
}
match logger.failure(py, error, &self.start, &self.end, true) {
Ok(Some(awaitable)) => {
self.pending = Some(Pending::AsyncFailure);
Ok(AdapterStep::Await(awaitable))
Ok(LifecycleStep::Await(awaitable))
}
Ok(None) => Ok(AdapterStep::Done),
Ok(None) => Ok(LifecycleStep::Done),
Err(failure) if is_cancellation(py, &failure) => Err(failure),
Err(_) => Ok(AdapterStep::Done),
Err(_) => Ok(LifecycleStep::Done),
}
}
}
impl CallbackAdapter for LegacyLogging {
impl PythonLifecycle for LegacyLogging {
fn begin(
&mut self,
py: Python<'_>,
arguments: Py<PyDict>,
started_at: f64,
) -> PyResult<AdapterStep> {
) -> PyResult<LifecycleStep> {
self.call.set_kwargs(arguments);
self.start = datetime(py, started_at)?;
self.internal = is_internal_call(py)?;
@ -212,9 +289,9 @@ impl CallbackAdapter for LegacyLogging {
)?;
self.logger = Some(result.logger()?);
self.call.set_kwargs(result.kwargs()?);
if self.deployment_hooks(py)? {
if self.runs_deployment_hooks() {
self.pending = Some(Pending::DeploymentPreCall);
return Ok(AdapterStep::Await(DeploymentHooks::before_call(
return Ok(LifecycleStep::Await(DeploymentHooks::before_call(
py,
self.call.kwargs(),
self.surface.call_type,
@ -228,18 +305,16 @@ impl CallbackAdapter for LegacyLogging {
py: Python<'_>,
wire: Box<WireRequest>,
context: &RequestContext,
) -> PyResult<AdapterStep> {
) -> PyResult<LifecycleStep> {
let logger = self.logger()?;
logger.update_from_kwargs(py, self.call.kwargs(), &wire, context)?;
if !logger.callbacks_needed(py, "payload")? {
logger.record_api_call_start(py)?;
return Ok(AdapterStep::Wire(wire));
}
let body = to_py(py, &wire.body)?
.into_bound(py)
.cast_into::<PyDict>()?;
for name in context.passthrough_fields.iter() {
if let Some(value) = self.call.lookup(py, name)? {
for (name, sent) in wire.body.as_object().into_iter().flatten() {
if let Some(value) = self.call.lookup(py, name)?
&& from_py::<Value>(&value).is_ok_and(|caller| caller == *sent)
{
body.set_item(name, value)?;
}
}
@ -249,11 +324,11 @@ impl CallbackAdapter for LegacyLogging {
}
self.body = Some(body.clone().unbind());
self.headers = Some(headers.clone().unbind());
let api_key = self.call.lookup(py, "api_key")?;
self.context = Some(context.clone());
self.logger()?.pre_call(
py,
self.surface.input_description,
api_key.as_ref(),
context.api_key.as_ref().map(|api_key| api_key.expose()),
&body,
&headers,
&wire.url,
@ -262,7 +337,7 @@ impl CallbackAdapter for LegacyLogging {
.iter()
.map(|(name, value)| Ok((name.extract::<String>()?, value.extract::<String>()?)))
.collect::<PyResult<Vec<_>>>()?;
Ok(AdapterStep::Wire(Box::new(WireRequest {
Ok(LifecycleStep::Wire(Box::new(WireRequest {
body: from_py(&body)?,
headers,
..*wire
@ -274,12 +349,12 @@ impl CallbackAdapter for LegacyLogging {
py: Python<'_>,
response: Py<PyAny>,
timing: Timing,
) -> PyResult<AdapterStep> {
) -> PyResult<LifecycleStep> {
self.end = Some(datetime(py, timing.end_time)?);
self.response = Some(response);
if self.deployment_hooks(py)? {
if self.runs_deployment_hooks() {
self.pending = Some(Pending::DeploymentPostCall);
return Ok(AdapterStep::Await(DeploymentHooks::after_success(
return Ok(LifecycleStep::Await(DeploymentHooks::after_success(
py,
self.call.kwargs(),
&self.response,
@ -289,36 +364,50 @@ impl CallbackAdapter for LegacyLogging {
self.finalize(py)
}
fn emit(
&mut self,
py: Python<'_>,
event: &CallEvent,
public: Option<PublicValue<'_>>,
) -> PyResult<AdapterStep> {
match (event, public) {
(CallEvent::ResponseReceived { raw }, _) => {
let logger = self.logger()?;
if logger.callbacks_needed(py, "payload")? {
logger.post_call(py, &raw.body, self.body.as_ref(), self.headers.as_ref())?;
}
Ok(AdapterStep::Done)
fn emit(&mut self, py: Python<'_>, event: LifecycleEvent<'_>) -> PyResult<LifecycleStep> {
match event {
LifecycleEvent::Started { .. } => Ok(LifecycleStep::Done),
LifecycleEvent::Machine(MachineEvent::ResponseReceived { raw }) => {
let api_key = self
.context
.as_ref()
.and_then(|context| context.api_key.as_ref())
.map(|api_key| api_key.expose());
self.logger()?.post_call(
py,
&raw.body,
api_key,
self.body.as_ref(),
self.headers.as_ref(),
)?;
Ok(LifecycleStep::Done)
}
(CallEvent::Succeeded { timing }, Some(PublicValue::Response(response))) => {
LifecycleEvent::Succeeded { timing, response } => {
self.end = Some(datetime(py, timing.end_time)?);
self.response = Some(response.clone_ref(py));
self.dispatch_success(py)?;
Ok(AdapterStep::Done)
match &self.stream {
Some(stream) => self.stream_success(py, stream)?,
None => self.dispatch_success(py)?,
}
Ok(LifecycleStep::Done)
}
(CallEvent::Failed { timing, origin }, Some(PublicValue::Error(error))) => {
LifecycleEvent::Failed {
timing,
origin,
error,
} => {
self.end = Some(datetime(py, timing.end_time)?);
self.error = Some(error.clone_ref(py).into_value(py));
if *origin == FailureOrigin::Call
if self.stream.is_some() {
return self.stream_failure(py);
}
if origin == FailureOrigin::Call
&& self.logger.is_some()
&& self.deployment_hooks(py)?
&& self.runs_deployment_hooks()
{
let error = self.error.as_ref().ok_or_else(missing_state)?;
self.pending = Some(Pending::DeploymentFailure);
return Ok(AdapterStep::Await(DeploymentHooks::after_failure(
return Ok(LifecycleStep::Await(DeploymentHooks::after_failure(
py,
self.call.kwargs(),
error,
@ -327,11 +416,30 @@ impl CallbackAdapter for LegacyLogging {
}
self.dispatch_failure(py)
}
_ => Err(missing_state()),
}
}
fn resume(&mut self, py: Python<'_>, result: PyResult<Py<PyAny>>) -> PyResult<AdapterStep> {
fn opened(&mut self, py: Python<'_>) -> PyResult<()> {
if self.surface.stream.is_none() {
return Err(missing_state());
}
Streaming::Opened.call(py, (self.logger()?.object(py),))?;
self.stream = Some(DeliveredStream {
chunks: PyList::empty(py).unbind(),
first_chunk: None,
});
Ok(())
}
fn delivered(&mut self, py: Python<'_>, chunk: &Py<PyAny>) -> PyResult<()> {
let stream = self.stream.as_mut().ok_or_else(missing_state)?;
if stream.first_chunk.is_none() {
stream.first_chunk = Some(datetime(py, epoch_seconds())?);
}
stream.chunks.bind(py).append(chunk)
}
fn resume(&mut self, py: Python<'_>, result: PyResult<Py<PyAny>>) -> PyResult<LifecycleStep> {
match self.pending.take().ok_or_else(missing_state)? {
Pending::DeploymentPreCall => {
self.call
@ -345,7 +453,7 @@ impl CallbackAdapter for LegacyLogging {
Pending::DeploymentFailure => self.dispatch_failure(py),
Pending::AsyncFailure => match result {
Err(failure) if is_cancellation(py, &failure) => Err(failure),
_ => Ok(AdapterStep::Done),
_ => Ok(LifecycleStep::Done),
},
}
}
@ -357,7 +465,8 @@ impl CallbackAdapter for LegacyLogging {
error.write_unraisable(py, None);
}
self.body = None;
self.headers = None;
self.context = None;
self.stream = None;
}
fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
@ -369,8 +478,11 @@ impl CallbackAdapter for LegacyLogging {
visit.call(&self.end)?;
visit.call(&self.response)?;
visit.call(&self.error)?;
visit.call(&self.body)?;
visit.call(&self.headers)
if let Some(stream) = &self.stream {
visit.call(&stream.chunks)?;
visit.call(&stream.first_chunk)?;
}
visit.call(&self.body)
}
}

45
litellm-rust/crates/callbacks-legacy/src/call.rs
View file

@ -3,8 +3,8 @@
//! lifetime. No other callback host has that obligation, which is why nothing outside
//! this crate holds them.
use litellm_callbacks::{machine::Machine, route::Route};
use litellm_host_python::{RouteHost, run_call};
use litellm_host::{machine::Machine, route::Route};
use litellm_host_python::{RouteHost, lookup, run_call};
use pyo3::{
gc::{PyTraverseError, PyVisit},
prelude::*,
@ -63,21 +63,6 @@ impl PublicCall {
}
}
/// The caller's own object for a public argument, as every legacy reader resolves it: the
/// keyword if given, even an explicit `None`, else the bound request's attribute. A route
/// host projecting from the prepared keyword view uses the same rule, so the callbacks
/// and the provider see one object per argument.
pub fn lookup<'py>(
kwargs: &Bound<'py, PyDict>,
request: &Bound<'py, PyAny>,
name: &str,
) -> PyResult<Option<Bound<'py, PyAny>>> {
if let Some(value) = kwargs.get_item(name)? {
return Ok(Some(value));
}
request.getattr_opt(name)
}
/// Runs one native call under the legacy `Logging` contract: the route host projects from
/// the keyword view the contract prepares, and the contract observes the call.
pub fn run_legacy_call<H, M>(
@ -121,32 +106,6 @@ mod tests {
(call, locals)
}
#[test]
fn lookup_prefers_the_keyword_even_when_none_and_falls_back_to_the_request() {
Python::initialize();
Python::attach(|py| {
let (call, locals) = capture(
py,
c"
key = object()
document = {'type': 'document_url'}
class Request:
api_key = 'from-request'
api_base = 'from-request'
document = document
request = Request()
kwargs = {'api_key': key, 'api_base': None}
",
);
let key = locals.get_item("key").unwrap().unwrap();
let document = locals.get_item("document").unwrap().unwrap();
assert!(call.lookup(py, "api_key").unwrap().unwrap().is(&key));
assert!(call.lookup(py, "api_base").unwrap().unwrap().is_none());
assert!(call.lookup(py, "document").unwrap().unwrap().is(&document));
assert!(call.lookup(py, "model").unwrap().is_none());
});
}
#[test]
fn capture_copies_the_keyword_dict_without_copying_its_values() {
Python::initialize();

241
litellm-rust/crates/callbacks-legacy/src/callbacks.rs
View file

@ -2,15 +2,14 @@
//! the deferred and worker-submitted success paths, and the sync-callbacks-for-async-calls
//! duplication. All of it expires with the legacy callback contract.
use litellm_callbacks::event::{RequestContext, WireRequest};
use litellm_host::event::{RequestContext, WireRequest};
use litellm_host_python::to_py;
use pyo3::{exceptions::PyBaseException, prelude::*, types::PyDict};
use crate::legacy_python::{Logging, Wrapper};
use crate::logger::PythonLogger;
pub trait LegacyCallbacks {
fn callbacks_needed(&self, py: Python<'_>, phase: &str) -> PyResult<bool>;
/// `Logging.update_from_kwargs`: what the logger is told about the request it is
/// about to see, with consumed credentials redacted.
fn update_from_kwargs(
@ -21,24 +20,23 @@ pub trait LegacyCallbacks {
context: &RequestContext,
) -> PyResult<()>;
fn record_api_call_start(&self, py: Python<'_>) -> PyResult<()>;
/// `Logging.pre_call`, or its payload-free shortcut when no input callback listens.
/// `Logging.pre_call`.
fn pre_call(
&self,
py: Python<'_>,
input: &str,
api_key: Option<&Bound<'_, PyAny>>,
api_key: Option<&str>,
body: &Bound<'_, PyDict>,
headers: &Bound<'_, PyDict>,
url: &str,
) -> PyResult<()>;
/// `Logging.post_call`, or its payload-free shortcut when no input callback listens.
/// `Logging.post_call`.
fn post_call(
&self,
py: Python<'_>,
original_response: &str,
api_key: Option<&str>,
body: Option<&Py<PyDict>>,
headers: Option<&Py<PyDict>>,
) -> PyResult<()>;
@ -82,16 +80,6 @@ pub trait LegacyCallbacks {
}
impl LegacyCallbacks for PythonLogger {
fn callbacks_needed(&self, py: Python<'_>, phase: &str) -> PyResult<bool> {
if !self.bridge_owned() {
return Ok(true);
}
py.import("litellm.rust_bridge.legacy_callbacks")?
.getattr("callbacks_needed")?
.call1((self.object(py), phase))?
.extract()
}
fn update_from_kwargs(
&self,
py: Python<'_>,
@ -100,18 +88,13 @@ impl LegacyCallbacks for PythonLogger {
context: &RequestContext,
) -> PyResult<()> {
let secret_fields: Vec<&str> = context.secret_fields.iter().map(String::as_str).collect();
let update = PyDict::new(py);
update.set_item("kwargs", redact(py, kwargs.bind(py), &secret_fields)?)?;
update.set_item("model", &context.model)?;
update.set_item(
"optional_params",
redact(
py,
&to_py(py, &context.optional_params)?
.into_bound(py)
.cast_into::<PyDict>()?,
&secret_fields,
)?,
let redacted_kwargs = redact(py, kwargs.bind(py), &secret_fields)?;
let optional_params = redact(
py,
&to_py(py, &context.optional_params)?
.into_bound(py)
.cast_into::<PyDict>()?,
&secret_fields,
)?;
let params = PyDict::new(py);
params.set_item(
@ -131,15 +114,17 @@ impl LegacyCallbacks for PythonLogger {
params.set_item(name, value)?;
}
}
update.set_item("litellm_params", params)?;
update.set_item("custom_llm_provider", &context.custom_llm_provider)?;
self.object(py)
.call_method("update_from_kwargs", (), Some(&update))?;
Ok(())
}
fn record_api_call_start(&self, py: Python<'_>) -> PyResult<()> {
self.object(py).call_method0("record_api_call_start_time")?;
Logging::Update.call(
py,
(
self.object(py),
redacted_kwargs,
&context.model,
optional_params,
params,
&context.custom_llm_provider,
),
)?;
Ok(())
}
@ -147,7 +132,7 @@ impl LegacyCallbacks for PythonLogger {
&self,
py: Python<'_>,
input: &str,
api_key: Option<&Bound<'_, PyAny>>,
api_key: Option<&str>,
body: &Bound<'_, PyDict>,
headers: &Bound<'_, PyDict>,
url: &str,
@ -156,17 +141,7 @@ impl LegacyCallbacks for PythonLogger {
additional.set_item("complete_input_dict", body)?;
additional.set_item("headers", headers)?;
additional.set_item("api_base", url)?;
let kwargs = PyDict::new(py);
kwargs.set_item("input", input)?;
kwargs.set_item("api_key", api_key)?;
kwargs.set_item("additional_args", &additional)?;
if self.callbacks_needed(py, "input")? {
self.object(py).call_method("pre_call", (), Some(&kwargs))?;
} else {
self.object(py)
.call_method("_pre_call", (), Some(&kwargs))?;
self.record_api_call_start(py)?;
}
Logging::PreCall.call(py, (self.object(py), input, api_key, &additional))?;
Ok(())
}
@ -174,37 +149,30 @@ impl LegacyCallbacks for PythonLogger {
&self,
py: Python<'_>,
original_response: &str,
api_key: Option<&str>,
body: Option<&Py<PyDict>>,
headers: Option<&Py<PyDict>>,
) -> PyResult<()> {
let additional = PyDict::new(py);
additional.set_item("complete_input_dict", body)?;
additional.set_item("headers", headers)?;
if self.callbacks_needed(py, "input")? {
let kwargs = PyDict::new(py);
kwargs.set_item("original_response", original_response)?;
kwargs.set_item("additional_args", &additional)?;
self.object(py)
.call_method("post_call", (), Some(&kwargs))?;
} else {
let response = py
.import("json")?
.call_method1("dumps", (original_response,))?;
self.object(py).call_method1(
"record_post_call",
(response, py.None(), py.None(), additional),
)?;
}
Logging::PostCall.call(
py,
(self.object(py), original_response, api_key, &additional),
)?;
Ok(())
}
fn defers_async_logging(&self, py: Python<'_>) -> bool {
self.object(py)
.getattr("_defer_async_logging")
.is_ok_and(|value| value.is_truthy().unwrap_or(false))
Logging::DefersAsync
.call(py, (self.object(py),))
.and_then(|value| value.extract())
.unwrap_or(false)
}
fn defer_success(&self, py: Python<'_>, pending: &Bound<'_, PyAny>) -> PyResult<()> {
self.object(py).setattr("_native_pending_logging", pending)
Logging::DeferSuccess.call(py, (self.object(py), pending))?;
Ok(())
}
fn sync_success_for_async_call(
@ -214,13 +182,7 @@ impl LegacyCallbacks for PythonLogger {
start: &Py<PyAny>,
end: &Option<Py<PyAny>>,
) -> PyResult<()> {
if !self.callbacks_needed(py, "sync_success_async")? {
return Ok(());
}
self.object(py).call_method1(
"handle_sync_success_callbacks_for_async_calls",
(response, start, end),
)?;
Logging::SyncSuccessForAsyncCall.call(py, (self.object(py), response, start, end))?;
Ok(())
}
@ -232,34 +194,11 @@ impl LegacyCallbacks for PythonLogger {
end: &Option<Py<PyAny>>,
asynchronous: bool,
) -> PyResult<Option<Py<PyAny>>> {
if !self.callbacks_needed(
py,
if asynchronous {
"async_failure"
} else {
"sync_failure"
},
)? {
py.import("litellm.rust_bridge.legacy_callbacks")?
.getattr("failure_bookkeeping")?
.call1((self.object(py), error, start, end, asynchronous))?;
return Ok(None);
}
let trace = py
.import("traceback")?
.getattr("format_exception")?
.call1((error,))?;
let trace = pyo3::types::PyString::new(py, "").call_method1("join", (trace,))?;
let value = self.object(py).call_method1(
if asynchronous {
"async_failure_handler"
} else {
"failure_handler"
},
(error, trace, start, end),
)?;
let value =
Logging::FailureHandler.call(py, (self.object(py), error, start, end, asynchronous))?;
Ok(asynchronous.then(|| value.unbind()))
}
fn submit_success(
&self,
py: Python<'_>,
@ -267,22 +206,7 @@ impl LegacyCallbacks for PythonLogger {
start: &Py<PyAny>,
end: &Option<Py<PyAny>>,
) -> PyResult<()> {
if !self.callbacks_needed(py, "sync_success")? {
return self.success_bookkeeping(py, response, start, end, false);
}
let context = py.import("contextvars")?.call_method0("copy_context")?;
py.import("litellm.litellm_core_utils.litellm_logging")?
.getattr("executor")?
.call_method1(
"submit",
(
context.getattr("run")?,
self.object(py).getattr("success_handler")?,
response,
start,
end,
),
)?;
Logging::SubmitSuccess.call(py, (self.object(py), response, start, end))?;
Ok(())
}
@ -293,18 +217,9 @@ impl LegacyCallbacks for PythonLogger {
start: &Py<PyAny>,
end: &Option<Py<PyAny>>,
) -> PyResult<()> {
if !self.callbacks_needed(py, "async_success")? {
return self.success_bookkeeping(py, response, start, end, true);
}
let context = py.import("contextvars")?.call_method0("copy_context")?;
let worker = py
.import("litellm.litellm_core_utils.logging_worker")?
.getattr("GLOBAL_LOGGING_WORKER")?
.getattr("ensure_initialized_and_enqueue")?;
let coroutine = self
.object(py)
.call_method1("async_success_handler", (response, start, end))?;
let enqueue = context.call_method1("run", (worker, &coroutine));
let coroutine =
Logging::AsyncSuccessHandler.call(py, (self.object(py), response, start, end))?;
let enqueue = Logging::Enqueue.call(py, (&coroutine,));
if enqueue.is_err()
&& let Err(error) = coroutine.call_method0("close")
{
@ -315,14 +230,7 @@ impl LegacyCallbacks for PythonLogger {
}
fn custom_pricing_fields(py: Python<'_>) -> PyResult<Vec<String>> {
py.import("litellm.types.utils")?
.getattr("CustomPricingLiteLLMParams")?
.getattr("model_fields")?
.cast_into::<PyDict>()?
.keys()
.iter()
.map(|name| name.extract::<String>())
.collect()
Logging::CustomPricingFields.call(py, ())?.extract()
}
fn redact(
@ -347,58 +255,5 @@ fn redact(
/// Proxy-internal calls skip the legacy success fan-out.
pub fn is_internal_call(py: Python<'_>) -> PyResult<bool> {
py.import("litellm._internal_context")?
.getattr("is_internal_call")?
.call_method0("get")?
.extract()
}
#[cfg(test)]
mod tests {
use pyo3::types::PyDict;
use super::*;
fn logger_whose_registries_need_no_input(py: Python<'_>, bridge_owned: bool) -> PythonLogger {
let locals = PyDict::new(py);
py.run(
c"
import sys
import types
for name in ('litellm', 'litellm.rust_bridge', 'litellm.rust_bridge.legacy_callbacks'):
sys.modules.setdefault(name, types.ModuleType(name))
legacy = sys.modules['litellm.rust_bridge.legacy_callbacks']
legacy.callbacks_needed = lambda logger, phase: logger.needed.get(phase, True)
class Logger:
needed = {'input': False}
logger = Logger()
",
Some(&locals),
Some(&locals),
)
.unwrap();
PythonLogger::new(
locals.get_item("logger").unwrap().unwrap().unbind(),
bridge_owned,
)
}
#[test]
fn a_caller_owned_logger_is_observed_in_full() {
Python::initialize();
Python::attach(|py| {
let logger = logger_whose_registries_need_no_input(py, false);
assert!(logger.callbacks_needed(py, "input").unwrap());
});
}
#[test]
fn a_bridge_owned_logger_is_elided_where_no_registry_needs_it() {
Python::initialize();
Python::attach(|py| {
let logger = logger_whose_registries_need_no_input(py, true);
assert!(!logger.callbacks_needed(py, "input").unwrap());
assert!(logger.callbacks_needed(py, "payload").unwrap());
});
}
Wrapper::IsInternalCall.call(py, ())?.extract()
}

183
litellm-rust/crates/callbacks-legacy/src/legacy_python.rs Normal file
View file

@ -0,0 +1,183 @@
use pyo3::prelude::*;
use strum::{IntoStaticStr, VariantArray};
const MODULE: &str = "litellm.rust_bridge.legacy_callbacks";
/// Every litellm Python internal the native call still borrows, grouped by the subsystem it
/// belongs to. Rust drives the call; these exist only so behaviour that Python owns today
/// (span tracking, the standard logging payload, spend, callback fan-out) keeps working.
/// A group is deleted once Rust owns that subsystem, so this enum only shrinks. Calling a
/// user's own callback is not borrowing and does not belong here.
///
/// `litellm/rust_bridge/legacy_callbacks.py` is the only Python module behind it, and
/// `python_contract.json` pins each function's parameters on both sides.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(crate) enum LegacyPython {
Wrapper(Wrapper),
Logging(Logging),
DeploymentHooks(DeploymentHooks),
Streaming(Streaming),
}
/// The `@client` wrapper around the call: `function_setup`, limits, credentials,
/// response metadata and the correlation context.
#[derive(Clone, Copy, Debug, IntoStaticStr, PartialEq, Eq, VariantArray)]
pub(crate) enum Wrapper {
#[strum(serialize = "setup")]
Setup,
#[strum(serialize = "check_limits")]
CheckLimits,
#[strum(serialize = "credential_list")]
CredentialList,
#[strum(serialize = "warn_unknown_credential")]
WarnUnknownCredential,
#[strum(serialize = "is_internal_call")]
IsInternalCall,
#[strum(serialize = "finalize")]
Finalize,
#[strum(serialize = "restore_context")]
RestoreContext,
}
/// litellm's `Logging` object and the sync and async callback fan-out behind it.
#[derive(Clone, Copy, Debug, IntoStaticStr, PartialEq, Eq, VariantArray)]
pub(crate) enum Logging {
#[strum(serialize = "custom_pricing_fields")]
CustomPricingFields,
#[strum(serialize = "update_logging")]
Update,
#[strum(serialize = "pre_call")]
PreCall,
#[strum(serialize = "post_call")]
PostCall,
#[strum(serialize = "defers_async_logging")]
DefersAsync,
#[strum(serialize = "defer_success")]
DeferSuccess,
#[strum(serialize = "sync_success_for_async_call")]
SyncSuccessForAsyncCall,
#[strum(serialize = "submit_success")]
SubmitSuccess,
#[strum(serialize = "async_success_handler")]
AsyncSuccessHandler,
#[strum(serialize = "enqueue_logging")]
Enqueue,
#[strum(serialize = "failure_handler")]
FailureHandler,
}
/// The `litellm.utils` fan-outs that run every callback's deployment hook.
#[derive(Clone, Copy, Debug, IntoStaticStr, PartialEq, Eq, VariantArray)]
pub(crate) enum DeploymentHooks {
#[strum(serialize = "before_deployment_call")]
BeforeDeploymentCall,
#[strum(serialize = "after_deployment_success")]
AfterDeploymentSuccess,
#[strum(serialize = "after_deployment_failure")]
AfterDeploymentFailure,
}
/// The Messages stream iterator's logging: the stream flag, the end-of-stream billing
/// from the delivered chunks, and the partial-usage failure path.
#[derive(Clone, Copy, Debug, IntoStaticStr, PartialEq, Eq, VariantArray)]
pub(crate) enum Streaming {
#[strum(serialize = "stream_opened")]
Opened,
#[strum(serialize = "stream_success")]
Success,
#[strum(serialize = "stream_failure")]
Failure,
}
impl LegacyPython {
fn name(self) -> &'static str {
match self {
Self::Wrapper(function) => function.into(),
Self::Logging(function) => function.into(),
Self::DeploymentHooks(function) => function.into(),
Self::Streaming(function) => function.into(),
}
}
pub(crate) fn call<'py, A>(self, py: Python<'py>, args: A) -> PyResult<Bound<'py, PyAny>>
where
A: pyo3::call::PyCallArgs<'py>,
{
py.import(MODULE)?.getattr(self.name())?.call1(args)
}
}
impl Wrapper {
pub(crate) fn call<'py, A>(self, py: Python<'py>, args: A) -> PyResult<Bound<'py, PyAny>>
where
A: pyo3::call::PyCallArgs<'py>,
{
LegacyPython::Wrapper(self).call(py, args)
}
}
impl Logging {
pub(crate) fn call<'py, A>(self, py: Python<'py>, args: A) -> PyResult<Bound<'py, PyAny>>
where
A: pyo3::call::PyCallArgs<'py>,
{
LegacyPython::Logging(self).call(py, args)
}
}
impl Streaming {
pub(crate) fn call<'py, A>(self, py: Python<'py>, args: A) -> PyResult<Bound<'py, PyAny>>
where
A: pyo3::call::PyCallArgs<'py>,
{
LegacyPython::Streaming(self).call(py, args)
}
}
impl DeploymentHooks {
pub(crate) fn call<'py, A>(self, py: Python<'py>, args: A) -> PyResult<Bound<'py, PyAny>>
where
A: pyo3::call::PyCallArgs<'py>,
{
LegacyPython::DeploymentHooks(self).call(py, args)
}
}
#[cfg(test)]
mod tests {
use std::collections::BTreeSet;
use strum::VariantArray;
use super::{DeploymentHooks, LegacyPython, Logging, Streaming, Wrapper};
use crate::test_support::PYTHON_CONTRACT;
#[test]
fn every_borrowed_function_is_in_the_python_contract() {
let contract: serde_json::Map<String, serde_json::Value> =
serde_json::from_str(PYTHON_CONTRACT).unwrap();
let declared: BTreeSet<&str> = contract.keys().map(String::as_str).collect();
let called: Vec<&str> = Wrapper::VARIANTS
.iter()
.map(|&function| LegacyPython::Wrapper(function))
.chain(
Logging::VARIANTS
.iter()
.map(|&function| LegacyPython::Logging(function)),
)
.chain(
DeploymentHooks::VARIANTS
.iter()
.map(|&function| LegacyPython::DeploymentHooks(function)),
)
.chain(
Streaming::VARIANTS
.iter()
.map(|&function| LegacyPython::Streaming(function)),
)
.map(LegacyPython::name)
.collect();
assert_eq!(called.len(), declared.len(), "a function is borrowed twice");
assert_eq!(called.into_iter().collect::<BTreeSet<_>>(), declared);
}
}

7
litellm-rust/crates/callbacks-legacy/src/lib.rs
View file

@ -2,7 +2,7 @@
//! sync and async callback registries it fans out to, the deployment hooks, the deferred
//! proxy release, and the kwargs rewrites the wrapper makes on the way in (credential-name
//! inheritance, budget and retry-count limits). All of it sits behind one
//! [`CallbackAdapter`](litellm_host_python::CallbackAdapter), so the driver, the routes and
//! [`PythonLifecycle`](litellm_host_python::PythonLifecycle), so the driver, the routes and
//! core never learn which Python object is on the other end.
//!
//! Legacy callbacks receive the caller's own objects and may mutate them. [`PublicCall`]
@ -13,6 +13,7 @@ mod adapter;
mod call;
mod callbacks;
mod deferred;
mod legacy_python;
mod logger;
mod preparation;
#[cfg(test)]
@ -20,8 +21,8 @@ mod preparation;
mod test_support;
pub(crate) use adapter::LegacyLogging;
pub use adapter::LegacySurface;
pub use call::{PublicCall, lookup, run_legacy_call};
pub use adapter::{LegacySurface, PassThroughStream};
pub use call::{PublicCall, run_legacy_call};
pub(crate) use callbacks::{LegacyCallbacks, is_internal_call};
pub(crate) use logger::{DeploymentHooks, PythonLogger, finalize, setup};
pub(crate) use preparation::prepare;

91
litellm-rust/crates/callbacks-legacy/src/logger.rs
View file

@ -5,34 +5,25 @@ use pyo3::{
types::{PyDict, PyTuple},
};
/// The `Logging` instance one call fans out through, and who owns it. A logger the caller
/// handed in is observed in full, because the caller reads it after the call; one this
/// crate built through `function_setup` is elided wherever no registry needs it.
use crate::legacy_python::{self, Wrapper};
/// The `Logging` instance one call fans out through.
pub struct PythonLogger {
object: Py<PyAny>,
bridge_owned: bool,
}
impl PythonLogger {
pub(crate) fn new(object: Py<PyAny>, bridge_owned: bool) -> Self {
Self {
object,
bridge_owned,
}
pub(crate) fn new(object: Py<PyAny>) -> Self {
Self { object }
}
pub(crate) fn object<'py>(&self, py: Python<'py>) -> &Bound<'py, PyAny> {
self.object.bind(py)
}
pub(crate) fn bridge_owned(&self) -> bool {
self.bridge_owned
}
pub fn clone_ref(&self, py: Python<'_>) -> Self {
Self {
object: self.object.clone_ref(py),
bridge_owned: self.bridge_owned,
}
}
@ -40,34 +31,17 @@ impl PythonLogger {
visit.call(&self.object)
}
pub fn success_bookkeeping(
&self,
py: Python<'_>,
response: &Option<Py<PyAny>>,
start: &Py<PyAny>,
end: &Option<Py<PyAny>>,
asynchronous: bool,
) -> PyResult<()> {
py.import("litellm.rust_bridge.legacy_callbacks")?
.getattr("success_bookkeeping")?
.call1((self.object(py), response, start, end, asynchronous))?;
Ok(())
}
pub fn restore_context(&self, py: Python<'_>) -> PyResult<()> {
py.import("litellm.utils")?
.getattr("_restore_correlation_context_if_supported")?
.call1((self.object(py),))?;
Wrapper::RestoreContext.call(py, (self.object(py),))?;
Ok(())
}
}
/// A bare Python object was not obtained from `setup`, so it is caller-owned.
impl FromPyObject<'_, '_> for PythonLogger {
type Error = PyErr;
fn extract(object: Borrowed<'_, '_, PyAny>) -> PyResult<Self> {
Ok(Self::new(object.to_owned().unbind(), false))
Ok(Self::new(object.to_owned().unbind()))
}
}
@ -75,9 +49,7 @@ pub struct SetupResult<'py>(Bound<'py, PyAny>);
impl SetupResult<'_> {
pub fn logger(&self) -> PyResult<PythonLogger> {
let object = self.0.getattr("logger")?.unbind();
let bridge_owned = self.0.getattr("bridge_owned")?.extract()?;
Ok(PythonLogger::new(object, bridge_owned))
Ok(PythonLogger::new(self.0.getattr("logger")?.unbind()))
}
pub fn kwargs(&self) -> PyResult<Py<PyDict>> {
@ -93,9 +65,8 @@ pub fn setup<'py>(
start: &Py<PyAny>,
asynchronous: bool,
) -> PyResult<SetupResult<'py>> {
py.import("litellm.rust_bridge.legacy_callbacks")?
.getattr("setup")?
.call1((call_type, args, kwargs, start, asynchronous))
Wrapper::Setup
.call(py, (call_type, args, kwargs, start, asynchronous))
.map(SetupResult)
}
@ -107,30 +78,20 @@ pub fn finalize(
start: &Py<PyAny>,
end: &Option<Py<PyAny>>,
) -> PyResult<()> {
py.import("litellm.rust_bridge.legacy_callbacks")?
.getattr("finalize")?
.call1((response, logger.object(py), kwargs, start, end))?;
Wrapper::Finalize.call(py, (response, logger.object(py), kwargs, start, end))?;
Ok(())
}
pub struct DeploymentHooks;
impl DeploymentHooks {
pub fn needed(py: Python<'_>) -> PyResult<bool> {
py.import("litellm.rust_bridge.legacy_callbacks")?
.getattr("deployment_callbacks_needed")?
.call0()?
.extract()
}
pub fn before_call(
py: Python<'_>,
kwargs: &Py<PyDict>,
call_type: &str,
) -> PyResult<Py<PyAny>> {
py.import("litellm.utils")?
.getattr("async_pre_call_deployment_hook")?
.call1((kwargs, call_type))
legacy_python::DeploymentHooks::BeforeDeploymentCall
.call(py, (kwargs, call_type))
.map(Bound::unbind)
}
@ -140,9 +101,8 @@ impl DeploymentHooks {
response: &Option<Py<PyAny>>,
call_type: &str,
) -> PyResult<Py<PyAny>> {
py.import("litellm.utils")?
.getattr("async_post_call_success_deployment_hook")?
.call1((kwargs, response, call_type))
legacy_python::DeploymentHooks::AfterDeploymentSuccess
.call(py, (kwargs, response, call_type))
.map(Bound::unbind)
}
@ -152,9 +112,8 @@ impl DeploymentHooks {
error: &Py<PyBaseException>,
call_type: &str,
) -> PyResult<Py<PyAny>> {
py.import("litellm.utils")?
.getattr("async_post_call_failure_deployment_hook")?
.call1((kwargs, error, call_type))
legacy_python::DeploymentHooks::AfterDeploymentFailure
.call(py, (kwargs, error, call_type))
.map(Bound::unbind)
}
}
@ -185,10 +144,6 @@ class Setup:
reads.append('logger')
return logger
@property
def bridge_owned(self):
reads.append('bridge_owned')
return True
@property
def kwargs(self):
reads.append('kwargs')
return []
@ -206,7 +161,6 @@ result = Setup()
.object(py)
.is(locals.get_item("logger").unwrap().unwrap())
);
assert!(logger.bridge_owned());
assert!(
result
.kwargs()
@ -220,17 +174,8 @@ result = Setup()
.unwrap()
.extract::<Vec<String>>()
.unwrap(),
["logger", "bridge_owned", "kwargs"]
["logger", "kwargs"]
);
});
}
#[test]
fn a_logger_extracted_from_a_bare_object_is_caller_owned() {
Python::initialize();
Python::attach(|py| {
let logger: PythonLogger = py.None().into_bound(py).extract().unwrap();
assert!(!logger.bridge_owned());
});
}
}

44
litellm-rust/crates/callbacks-legacy/src/preparation.rs
View file

@ -3,6 +3,8 @@ use pyo3::{
types::{PyDict, PyList},
};
use crate::legacy_python::Wrapper;
struct CredentialEntry<'py>(Bound<'py, PyAny>);
impl<'py> CredentialEntry<'py> {
@ -22,18 +24,19 @@ pub fn prepare<'py>(
) -> PyResult<Bound<'py, PyDict>> {
let arguments = kwargs.copy()?;
arguments.set_item("litellm_logging_obj", logger.object(py))?;
let litellm = py.import("litellm")?;
inherit_credentials(py, &litellm, &arguments)?;
py.import("litellm.rust_bridge.legacy_callbacks")?
.getattr("check_limits")?
.call1((&arguments,))?;
inherit_credentials(py, &arguments, || {
Ok(Wrapper::CredentialList
.call(py, ())?
.cast_into::<PyList>()?)
})?;
Wrapper::CheckLimits.call(py, (&arguments,))?;
Ok(arguments)
}
fn inherit_credentials(
py: Python<'_>,
litellm: &Bound<'_, PyModule>,
arguments: &Bound<'_, PyDict>,
fn inherit_credentials<'py>(
py: Python<'py>,
arguments: &Bound<'py, PyDict>,
credential_list: impl FnOnce() -> PyResult<Bound<'py, PyList>>,
) -> PyResult<()> {
let Some(requested) = arguments
.get_item("litellm_credential_name")?
@ -45,16 +48,13 @@ fn inherit_credentials(
return Ok(());
}
let requested: String = requested.extract()?;
let credentials = litellm.getattr("credential_list")?.cast_into::<PyList>()?;
let credentials = credential_list()?;
let names = credentials
.iter()
.map(|credential| CredentialEntry(credential).name())
.collect::<PyResult<Vec<_>>>()?;
let Some(index) = names.iter().position(|name| *name == requested) else {
py.import("litellm._logging")?.getattr("verbose_logger")?.call_method1(
"warning",
("litellm_credential_name=%s matched none of the %d loaded credentials; the request runs without it", requested, names.len()),
)?;
Wrapper::WarnUnknownCredential.call(py, (requested, names.len()))?;
return Ok(());
};
let selected = CredentialEntry(credentials.get_item(index)?);
@ -80,19 +80,19 @@ mod tests {
}
fn inherit(py: Python<'_>, locals: &Bound<'_, PyDict>) -> PyResult<()> {
let litellm = PyModule::new(py, "credential_host")?;
litellm.setattr(
"credential_list",
locals.get_item("credentials").unwrap().unwrap(),
)?;
inherit_credentials(
py,
&litellm,
&locals
.get_item("arguments")
.unwrap()
.unwrap()
.cast_into::<PyDict>()?,
|| {
Ok(locals
.get_item("credentials")?
.unwrap()
.cast_into::<PyList>()?)
},
)
}
@ -304,11 +304,11 @@ arguments = {'litellm_credential_name': 'ocr-test'}
fn falsy_credential_names_return_before_loading_credentials() {
Python::initialize();
Python::attach(|py| {
let litellm = PyModule::new(py, "credential_host").unwrap();
for name in [py.None(), py.eval(c"''", None, None).unwrap().unbind()] {
let arguments = PyDict::new(py);
arguments.set_item("litellm_credential_name", name).unwrap();
inherit_credentials(py, &litellm, &arguments).unwrap();
inherit_credentials(py, &arguments, || panic!("credentials must not be loaded"))
.unwrap();
}
});
}

18
litellm-rust/crates/callbacks-legacy/tests/deferred.rs
View file

@ -16,7 +16,7 @@ fn defer<'py>(py: Python<'py>, script: &CStr) -> Bound<'py, PyDict> {
py,
PendingLogging {
pending: Some(PendingSuccess {
logger: PythonLogger::new(local(&locals, "logger").unbind(), true),
logger: PythonLogger::new(local(&locals, "logger").unbind()),
response: Some(local(&locals, "response").unbind()),
start: py.None(),
end: Some(py.None()),
@ -79,22 +79,6 @@ assert logger.calls == [], logger.calls
});
}
#[test]
fn a_release_after_the_async_callbacks_went_away_only_keeps_the_books() {
Python::initialize();
Python::attach(|py| {
let locals = defer(py, c"logger.needed = {'async_success': False}");
run(
py,
&locals,
c"
pending.release(True)
assert logger.calls == [('success_bookkeeping', True)], logger.calls
",
);
});
}
#[rstest]
#[case::ordinary_error(c"RuntimeError('queue full')", false)]
#[case::cancellation(c"asyncio.CancelledError()", true)]

64
litellm-rust/crates/callbacks-legacy/tests/deployment_hooks.rs
View file

@ -1,7 +1,7 @@
use std::ffi::CStr;
use litellm_callbacks::event::{CallEvent, FailureOrigin, Timing};
use litellm_host_python::{AdapterStep, CallbackAdapter, PublicValue};
use litellm_host::event::{FailureOrigin, Timing};
use litellm_host_python::{LifecycleEvent, LifecycleStep, PythonLifecycle};
use pyo3::exceptions::asyncio::CancelledError;
use pyo3::prelude::*;
use pyo3::types::PyDict;
@ -24,7 +24,7 @@ fn begin<'py>(
py: Python<'py>,
locals: &Bound<'py, PyDict>,
asynchronous: bool,
) -> (LegacyLogging, AdapterStep) {
) -> (LegacyLogging, LifecycleStep) {
let mut logging = legacy_call(py, locals, asynchronous);
let kwargs = local(locals, "kwargs")
.cast_into::<PyDict>()
@ -34,15 +34,15 @@ fn begin<'py>(
(logging, step)
}
fn arguments<'py>(py: Python<'py>, step: AdapterStep) -> Bound<'py, PyDict> {
let AdapterStep::Arguments(arguments) = step else {
fn arguments<'py>(py: Python<'py>, step: LifecycleStep) -> Bound<'py, PyDict> {
let LifecycleStep::Arguments(arguments) = step else {
panic!("expected the prepared arguments");
};
arguments.into_bound(py)
}
fn awaits_deployment_hook(step: &AdapterStep) -> bool {
matches!(step, AdapterStep::Await(_))
fn awaits_deployment_hook(step: &LifecycleStep) -> bool {
matches!(step, LifecycleStep::Await(_))
}
#[rstest]
@ -97,6 +97,43 @@ assert checked is prepared
});
}
#[rstest]
#[case::synchronous(false)]
#[case::asynchronous(true)]
fn a_keyword_the_bridge_never_reads_reaches_every_reader_as_the_callers_object(
#[case] asynchronous: bool,
) {
Python::initialize();
Python::attach(|py| {
let locals = namespace(
py,
c"
opaque = object()
hooked = []
logger.hooks = {'pre': lambda kwargs: hooked.append(kwargs['vendor_extension']) or kwargs}
kwargs = {'logger': logger, 'vendor_extension': opaque}
",
);
let (mut logging, step) = begin(py, &locals, asynchronous);
let step = match step {
LifecycleStep::Await(hook_result) => logging.resume(py, Ok(hook_result)).unwrap(),
step => step,
};
locals.set_item("prepared", arguments(py, step)).unwrap();
locals.set_item("asynchronous", asynchronous).unwrap();
run(
py,
&locals,
c"
assert prepared['vendor_extension'] is opaque
[checked] = [value for name, value in logger.calls if name == 'check_limits']
assert checked['vendor_extension'] is opaque
assert hooked == ([opaque] if asynchronous else []), hooked
",
);
});
}
#[test]
fn response_returned_by_the_post_call_hook_is_finalized_and_returned() {
Python::initialize();
@ -121,7 +158,7 @@ logger.hooks = {'pre': lambda kwargs: kwargs}
let step = logging
.resume(py, Ok(local(&locals, "replacement").unbind()))
.unwrap();
let AdapterStep::Response(returned) = step else {
let LifecycleStep::Response(returned) = step else {
panic!("expected the finalized response");
};
assert!(returned.bind(py).is(local(&locals, "replacement")));
@ -180,13 +217,12 @@ fn failure_callbacks_run_after_the_failure_hook_however_it_ends(#[case] cancelle
.resume(py, Ok(local(&locals, "kwargs").unbind()))
.unwrap();
let failure = PyErr::from_value(local(&locals, "failure"));
let failed = CallEvent::Failed {
let failed = LifecycleEvent::Failed {
timing: TIMING,
origin: FailureOrigin::Call,
error: &failure,
};
let step = logging
.emit(py, &failed, Some(PublicValue::Error(&failure)))
.unwrap();
let step = logging.emit(py, failed).unwrap();
assert!(awaits_deployment_hook(&step));
let hook_result = if cancelled {
Err(CancelledError::new_err("cancelled"))
@ -195,7 +231,7 @@ fn failure_callbacks_run_after_the_failure_hook_however_it_ends(#[case] cancelle
};
assert!(matches!(
logging.resume(py, hook_result).unwrap(),
AdapterStep::Await(_)
LifecycleStep::Await(_)
));
run(
py,
@ -237,7 +273,7 @@ kwargs = {'logger': logger}
.unwrap()
.unbind();
let result = logging.begin(py, kwargs, 0.0).and_then(|step| match step {
AdapterStep::Await(_) => logging.resume(py, Ok(local(&locals, "kwargs").unbind())),
LifecycleStep::Await(_) => logging.resume(py, Ok(local(&locals, "kwargs").unbind())),
step => Ok(step),
});
let error = result.err().unwrap();

292
litellm-rust/crates/callbacks-legacy/tests/payload.rs
View file

@ -1,10 +1,12 @@
use std::ffi::CStr;
use litellm_callbacks::event::{CallEvent, Passthrough, RawResponse, RequestContext, WireRequest};
use litellm_host_python::{AdapterStep, CallbackAdapter};
use litellm_auth::SecretValue;
use litellm_host::event::{MachineEvent, RawResponse, RequestContext, WireRequest};
use litellm_host_python::{LifecycleEvent, LifecycleStep, PythonLifecycle, to_py};
use proptest::prelude::*;
use pyo3::prelude::*;
use rstest::rstest;
use serde_json::{Value, json};
use serde_json::{Map, Value, json};
use super::LegacyLogging;
use crate::PythonLogger;
@ -23,20 +25,12 @@ class PayloadLogger(StubLogger):
def pre_call(self, input, api_key, additional_args):
self.record('pre_call', None)
self.pre = additional_args
self.pre_api_key = api_key
on_pre_call(additional_args)
def _pre_call(self, input, api_key, additional_args):
self.record('_pre_call', None)
def record_api_call_start_time(self):
self.record('record_api_call_start_time', None)
def post_call(self, original_response, additional_args):
def post_call(self, original_response, api_key, additional_args):
self.record('post_call', None)
self.post = (original_response, additional_args)
def record_post_call(self, response, *rest):
self.record('record_post_call', response)
self.post = (original_response, api_key, additional_args)
request = Request()
kwargs = {}
@ -52,33 +46,47 @@ fn document(source: &str) -> Value {
json!({"type": "document_url", "document_url": source})
}
fn before_send(script: &CStr, caller: Value, body: Value) -> WireRequest {
before_send_with_secrets(script, caller, body, &[])
fn before_send(script: &CStr, body: Value) -> WireRequest {
before_send_with_secrets(script, json!({}), body, &[])
}
/// Runs `before_send` over `body` for a caller whose route-side view is `caller`, with the
/// Python objects `script` binds, then delivers the provider's raw response the way the
/// Runs `before_send` over `body` for a route whose parameters are `optional_params`, with
/// the Python objects `script` binds, then delivers the provider's raw response the way the
/// driver does and runs the script's `check()`.
fn before_send_with_secrets(
script: &CStr,
caller: Value,
optional_params: Value,
body: Value,
secret_fields: &[&str],
) -> WireRequest {
before_send_bound(&[], script, optional_params, body, secret_fields)
}
/// [`before_send_with_secrets`] with `bindings` placed in the namespace before `script` runs.
fn before_send_bound(
bindings: &[(&str, &Value)],
script: &CStr,
optional_params: Value,
body: Value,
secret_fields: &[&str],
) -> WireRequest {
Python::initialize();
Python::attach(|py| {
let locals = namespace(py, PAYLOAD_LOGGER);
for &(name, value) in bindings {
locals.set_item(name, to_py(py, value).unwrap()).unwrap();
}
run(py, &locals, script);
let mut logging = LegacyLogging {
logger: Some(PythonLogger::new(local(&locals, "logger").unbind(), true)),
logger: Some(PythonLogger::new(local(&locals, "logger").unbind())),
..legacy_call(py, &locals, false)
};
let context = RequestContext {
model: "model".into(),
custom_llm_provider: "provider".into(),
optional_params: caller.clone(),
passthrough_fields: Passthrough::unchanged(caller.as_object().unwrap(), &body),
optional_params,
secret_fields: secret_fields.iter().map(|name| name.to_string()).collect(),
api_key: Some(SecretValue::new("route-key")),
};
let wire = WireRequest {
url: "https://provider.invalid/ocr".into(),
@ -86,17 +94,17 @@ fn before_send_with_secrets(
body,
};
let step = logging.before_send(py, Box::new(wire), &context).unwrap();
let raw = CallEvent::ResponseReceived {
let raw = MachineEvent::ResponseReceived {
raw: RawResponse {
body: "raw response".into(),
},
};
assert!(matches!(
logging.emit(py, &raw, None).unwrap(),
AdapterStep::Done
logging.emit(py, LifecycleEvent::Machine(&raw)).unwrap(),
LifecycleStep::Done
));
run(py, &locals, c"check()");
let AdapterStep::Wire(wire) = step else {
let LifecycleStep::Wire(wire) = step else {
panic!("before_send did not hand back the wire request");
};
*wire
@ -129,11 +137,7 @@ def check():
")]
fn passthrough_keys_reach_pre_call_as_the_callers_own_objects(#[case] script: &CStr) {
let body = json!({"model": "model", "document": document(DOCUMENT), "pages": [0]});
let wire = before_send(
script,
json!({"document": document(DOCUMENT), "pages": [0]}),
body.clone(),
);
let wire = before_send(script, body.clone());
assert_eq!(wire.body, body);
}
@ -149,7 +153,6 @@ def check():
assert document['document_url'] == 'data:application/pdf;base64,ZWRpdGVk'
",
json!({"document": document(DOCUMENT)}),
json!({"document": document(DOCUMENT)}),
);
assert_eq!(wire.body["document"], document(EDITED));
}
@ -168,7 +171,6 @@ def check():
assert observed == [False], observed
assert document == {'type': 'document_url', 'document_url': 'https://example.invalid/scan.pdf'}
",
json!({"document": document("https://example.invalid/scan.pdf")}),
json!({"document": document(DOCUMENT)}),
);
assert_eq!(
@ -177,6 +179,23 @@ def check():
);
}
#[test]
fn a_caller_value_with_no_json_form_is_left_out_of_realiasing() {
let body = json!({"pages": [0]});
let wire = before_send(
c"
opaque = object()
kwargs = {'pages': opaque}
observed = []
on_pre_call = lambda args: observed.append(args['complete_input_dict']['pages'])
def check():
assert observed == [[0]], observed
",
body.clone(),
);
assert_eq!(wire.body, body);
}
#[rstest]
#[case::body(
c"
@ -192,7 +211,7 @@ def on_pre_call(args):
)]
fn rebinding_the_payload_envelope_does_not_reach_the_wire(#[case] script: &CStr) {
let body = json!({"document": document(DOCUMENT)});
let wire = before_send(script, json!({}), body.clone());
let wire = before_send(script, body.clone());
assert_eq!(wire.body, body);
assert_eq!(wire.headers, [("x-route".to_string(), "route".to_string())]);
}
@ -205,7 +224,6 @@ def on_pre_call(args):
args['headers']['x-callback'] = 'edited'
",
json!({}),
json!({}),
);
assert_eq!(
wire.headers,
@ -288,7 +306,7 @@ def on_pre_call(args):
)]
fn pre_call_body_edits_reach_the_wire(#[case] script: &CStr, #[case] expected: Value) {
let body = json!({"document": document(DOCUMENT)});
let wire = before_send(script, json!({"document": document(DOCUMENT)}), body);
let wire = before_send(script, body);
assert_eq!(wire.body, expected);
}
@ -302,7 +320,6 @@ def on_pre_call(args):
retained['x-retained'] = 'sent'
",
json!({}),
json!({}),
);
assert_eq!(
wire.headers,
@ -314,52 +331,193 @@ def on_pre_call(args):
}
#[test]
fn post_call_receives_the_raw_response_and_the_payload_dicts_pre_call_saw() {
fn post_call_receives_the_raw_response_the_route_key_and_the_body_and_headers_pre_call_saw() {
before_send(
c"
def check():
original_response, additional_args = logger.post
original_response, api_key, additional_args = logger.post
assert original_response == 'raw response', original_response
assert api_key == logger.pre_api_key == 'route-key', (api_key, logger.pre_api_key)
assert additional_args == {
'complete_input_dict': logger.pre['complete_input_dict'],
'headers': logger.pre['headers'],
}, additional_args
assert additional_args['complete_input_dict'] is logger.pre['complete_input_dict']
assert additional_args['headers'] is logger.pre['headers']
",
json!({}),
json!({"document": document(DOCUMENT)}),
);
}
#[rstest]
#[case::every_phase_listens(c"{}", &["pre_call", "post_call"])]
#[case::no_input_callback(
c"{'input': False}",
&["_pre_call", "record_api_call_start_time", "record_post_call"]
)]
#[case::no_payload_consumer(c"{'payload': False}", &["record_api_call_start_time"])]
fn payload_callbacks_run_only_for_the_phases_someone_listens_to(
#[case] needed: &CStr,
#[case] expected_calls: &[&str],
) {
let script = std::ffi::CString::new(format!(
"
logger.needed = {needed}
#[test]
fn every_request_runs_the_full_pre_call_and_post_call() {
let wire = before_send(
c"
def on_pre_call(args):
args['complete_input_dict']['include_image_base64'] = True
def check():
assert logger.names() == {expected_calls:?}, logger.calls
assert logger.names() == ['pre_call', 'post_call'], logger.calls
",
needed = needed.to_str().unwrap(),
expected_calls = expected_calls,
))
.unwrap();
let body = json!({"document": document(DOCUMENT)});
let wire = before_send(&script, json!({}), body.clone());
let edited = json!({"document": document(DOCUMENT), "include_image_base64": true});
json!({"document": document(DOCUMENT)}),
);
assert_eq!(
wire.body,
if expected_calls.contains(&"pre_call") {
edited
} else {
body
}
json!({"document": document(DOCUMENT), "include_image_base64": true})
);
}
/// What one pre-call callback does to the payload it is handed.
#[derive(Clone, Debug)]
enum Edit {
Nothing,
Set(String, Value),
Remove(String),
Rebind(Value),
RebindThenSetRetained(String, Value),
}
impl Edit {
fn script(&self) -> Value {
match self {
Self::Nothing => json!({"kind": "nothing"}),
Self::Set(key, value) => json!({"kind": "set", "key": key, "value": value}),
Self::Remove(key) => json!({"kind": "remove", "key": key}),
Self::Rebind(value) => json!({"kind": "rebind", "value": value}),
Self::RebindThenSetRetained(key, value) => {
json!({"kind": "rebind_then_set_retained", "key": key, "value": value})
}
}
}
/// The legacy contract: the provider is sent the body object `pre_call` received, as
/// the callback left it. Rebinding the envelope's key points the envelope elsewhere and
/// leaves that object alone.
fn sent(&self, body: &Map<String, Value>) -> Value {
let mut sent = body.clone();
match self {
Self::Nothing | Self::Rebind(_) => {}
Self::Set(key, value) | Self::RebindThenSetRetained(key, value) => {
sent.insert(key.clone(), value.clone());
}
Self::Remove(key) => {
sent.remove(key);
}
}
Value::Object(sent)
}
}
/// How the caller's keyword for a body key relates to what the route sends under it.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
enum Caller {
PassedUnchanged,
RewrittenByTheRoute,
NotPassed,
}
const MODEL: &CStr = c"
aliased = {}
def on_pre_call(args):
body = args['complete_input_dict']
aliased.update({name: body[name] is kwargs[name] for name in unchanged})
kind = edit['kind']
if kind == 'set':
body[edit['key']] = edit['value']
elif kind == 'remove':
body.pop(edit['key'], None)
elif kind == 'rebind':
args['complete_input_dict'] = edit['value']
elif kind == 'rebind_then_set_retained':
args['complete_input_dict'] = {}
body[edit['key']] = edit['value']
def check():
assert aliased == {name: True for name in unchanged}, aliased
assert logger.names() == ['pre_call', 'post_call'], logger.calls
";
fn json_value() -> impl Strategy<Value = Value> {
let leaf = prop_oneof![
Just(Value::Null),
any::<bool>().prop_map(Value::from),
any::<i64>().prop_map(Value::from),
any::<f64>()
.prop_filter("JSON has no NaN or infinity", |number| number.is_finite())
.prop_map(Value::from),
".{0,8}".prop_map(Value::from),
];
leaf.prop_recursive(3, 24, 4, |inner| {
prop_oneof![
prop::collection::vec(inner.clone(), 0..4).prop_map(Value::from),
prop::collection::btree_map(key(), inner, 0..4)
.prop_map(|fields| Value::Object(fields.into_iter().collect())),
]
})
}
fn key() -> impl Strategy<Value = String> {
"[a-z]{1,6}"
}
fn caller() -> impl Strategy<Value = Caller> {
prop_oneof![
Just(Caller::PassedUnchanged),
Just(Caller::RewrittenByTheRoute),
Just(Caller::NotPassed),
]
}
fn edit() -> impl Strategy<Value = Edit> {
prop_oneof![
Just(Edit::Nothing),
(key(), json_value()).prop_map(|(key, value)| Edit::Set(key, value)),
key().prop_map(Edit::Remove),
json_value().prop_map(Edit::Rebind),
(key(), json_value()).prop_map(|(key, value)| Edit::RebindThenSetRetained(key, value)),
]
}
proptest! {
#![proptest_config(ProptestConfig::with_cases(128))]
/// For any body, any caller keywords and any callback edit: every keyword the route
/// sends unchanged reaches `pre_call` as the caller's own object, and the provider is
/// sent exactly what the model says, so a callback that edits nothing changes nothing.
#[test]
fn the_wire_is_the_body_pre_call_received_as_the_callback_left_it(
fields in prop::collection::btree_map(key(), (json_value(), caller()), 0..5),
edit in edit(),
) {
let body: Map<String, Value> = fields
.iter()
.map(|(name, (value, _))| (name.clone(), value.clone()))
.collect();
let kwargs: Map<String, Value> = fields
.iter()
.filter_map(|(name, (value, caller))| match caller {
Caller::PassedUnchanged => Some((name.clone(), value.clone())),
Caller::RewrittenByTheRoute => Some((name.clone(), json!([value]))),
Caller::NotPassed => None,
})
.collect();
let unchanged: Value = fields
.iter()
.filter(|(_, (_, caller))| *caller == Caller::PassedUnchanged)
.map(|(name, _)| Value::from(name.clone()))
.collect();
let wire = before_send_bound(
&[
("kwargs", &Value::Object(kwargs)),
("unchanged", &unchanged),
("edit", &edit.script()),
],
MODEL,
json!({}),
Value::Object(body.clone()),
&[],
);
prop_assert_eq!(wire.body, edit.sent(&body));
prop_assert_eq!(wire.headers, [("x-route".to_string(), "route".to_string())]);
}
}

141
litellm-rust/crates/callbacks-legacy/tests/support.rs
View file

@ -5,64 +5,94 @@ use pyo3::types::{PyDict, PyTuple};
use crate::{LegacyLogging, LegacySurface, PublicCall};
/// Stand-ins for every litellm function the legacy contract calls. Tests share one
/// interpreter and run concurrently, so each stub is installed idempotently and forwards to
/// the per-test `StubLogger` it is handed (directly, or as `kwargs['logger']`).
/// The parameters of every `legacy_callbacks` function, as the real module declares them.
/// `tests/test_litellm/rust_bridge/test_legacy_callbacks.py` pins this file to the Python
/// signatures, and [`namespace`] binds every fake call against it.
pub(crate) const PYTHON_CONTRACT: &str = include_str!("../python_contract.json");
/// Stand-ins for `legacy_callbacks`, the only Python module the crate calls. Tests
/// share one interpreter and run concurrently, so each fake is installed idempotently and
/// forwards to the per-test `StubLogger` it is handed (directly, or as `kwargs['logger']`).
/// Every fake is bound against the contract first, so a call the real module would reject
/// fails here too.
const STUBS: &CStr = c"
import contextvars
import inspect
import json
import sys
import traceback
import types
for name in (
'litellm',
'litellm.utils',
'litellm.types',
'litellm.types.utils',
'litellm._internal_context',
'litellm.litellm_core_utils',
'litellm.litellm_core_utils.logging_worker',
'litellm.litellm_core_utils.litellm_logging',
'litellm.rust_bridge',
'litellm.rust_bridge.legacy_callbacks',
):
for name in ('litellm', 'litellm.rust_bridge', 'litellm.rust_bridge.legacy_callbacks'):
sys.modules.setdefault(name, types.ModuleType(name))
legacy = sys.modules['litellm.rust_bridge.legacy_callbacks']
legacy.setup = lambda call_type, args, kwargs, start, asynchronous: types.SimpleNamespace(
logger=kwargs['logger_factory'](kwargs) if 'logger_factory' in kwargs else kwargs['logger'],
kwargs=kwargs,
bridge_owned=True,
)
legacy.deployment_callbacks_needed = lambda: True
legacy.check_limits = lambda arguments: arguments['logger'].check_limits(arguments)
legacy.callbacks_needed = lambda logger, phase: logger.needed.get(phase, True)
legacy.success_bookkeeping = lambda logger, response, start, end, asynchronous: logger.record(
'success_bookkeeping', asynchronous
)
legacy.failure_bookkeeping = lambda logger, error, start, end, asynchronous: logger.record(
'failure_bookkeeping', asynchronous
)
legacy.finalize = lambda response, logger, kwargs, start, end: logger.record('finalize', response)
CONTRACT = json.loads(python_contract)
utils = sys.modules['litellm.utils']
utils.async_pre_call_deployment_hook = lambda kwargs, call_type: kwargs['logger'].hook(
'pre', kwargs, call_type
)
utils.async_post_call_success_deployment_hook = lambda kwargs, response, call_type: kwargs[
'logger'
].hook('success', response, call_type)
utils.async_post_call_failure_deployment_hook = lambda kwargs, error, call_type: kwargs[
'logger'
].hook('failure', error, call_type)
utils._restore_correlation_context_if_supported = lambda logger: logger.record('restore', None)
internal = sys.modules['litellm._internal_context']
if not hasattr(internal, 'is_internal_call'):
internal.is_internal_call = contextvars.ContextVar('is_internal_call', default=False)
def contracted(name, fake):
signature = inspect.Signature(
[inspect.Parameter(parameter, inspect.Parameter.POSITIONAL_OR_KEYWORD) for parameter in CONTRACT[name]]
)
sys.modules['litellm.types.utils'].CustomPricingLiteLLMParams = type(
'CustomPricingLiteLLMParams', (), {'model_fields': {'ocr_cost_per_page': None}}
)
def checked(*args, **kwargs):
signature.bind(*args, **kwargs)
return fake(*args, **kwargs)
return checked
if not hasattr(legacy, 'is_internal'):
legacy.is_internal = contextvars.ContextVar('is_internal_call', default=False)
FAKES = {
'setup': lambda call_type, args, kwargs, start, asynchronous: types.SimpleNamespace(
logger=kwargs['logger_factory'](kwargs) if 'logger_factory' in kwargs else kwargs['logger'],
kwargs=kwargs,
),
'check_limits': lambda arguments: arguments['logger'].check_limits(arguments),
'finalize': lambda response, logger, kwargs, start, end: logger.record('finalize', response),
'update_logging': lambda logger, kwargs, model, optional_params, litellm_params, provider: logger.update_from_kwargs(
kwargs=kwargs,
model=model,
optional_params=optional_params,
litellm_params=litellm_params,
custom_llm_provider=provider,
),
'pre_call': lambda logger, input, api_key, additional_args: logger.pre_call(input, api_key, additional_args),
'post_call': lambda logger, original_response, api_key, additional_args: logger.post_call(
original_response, api_key, additional_args
),
'defers_async_logging': lambda logger: bool(getattr(logger, '_defer_async_logging', False)),
'defer_success': lambda logger, pending: setattr(logger, '_native_pending_logging', pending),
'sync_success_for_async_call': lambda logger, response, start, end: logger.handle_sync_success_callbacks_for_async_calls(
response, start, end
),
'failure_handler': lambda logger, error, start, end, asynchronous: (
logger.async_failure_handler if asynchronous else logger.failure_handler
)(error, ''.join(traceback.format_exception(error)), start, end),
'submit_success': lambda logger, response, start, end: logger.record('submit', (response, start, end)),
'async_success_handler': lambda logger, response, start, end: logger.async_success_handler(response, start, end),
'enqueue_logging': lambda coroutine: coroutine.enqueue(),
'restore_context': lambda logger: logger.record('restore', None),
'custom_pricing_fields': lambda: ('ocr_cost_per_page',),
'is_internal_call': lambda: legacy.is_internal.get(),
'credential_list': lambda: [],
'warn_unknown_credential': lambda name, loaded: None,
'before_deployment_call': lambda kwargs, call_type: kwargs['logger'].hook('pre', kwargs, call_type),
'after_deployment_success': lambda kwargs, response, call_type: kwargs['logger'].hook(
'success', response, call_type
),
'after_deployment_failure': lambda kwargs, error, call_type: kwargs['logger'].hook('failure', error, call_type),
'stream_opened': lambda logger: logger.record('stream_opened', None),
'stream_success': lambda logger, request_body, chunks, start, end, first_chunk: logger.record(
'stream_success', list(chunks)
),
'stream_failure': lambda logger, request_body, chunks, error: logger.record('stream_failure', error),
}
assert FAKES.keys() == CONTRACT.keys(), sorted(FAKES.keys() ^ CONTRACT.keys())
for name, fake in FAKES.items():
setattr(legacy, name, contracted(name, fake))
unraisable = sys.modules.setdefault(
@ -77,20 +107,6 @@ def unraisable_from(owner):
return [error for source, error in unraisable.events if source is owner]
class Worker:
def ensure_initialized_and_enqueue(self, coroutine):
return coroutine.enqueue()
class Executor:
def submit(self, run, handler, *args):
handler.__self__.record('submit', args)
sys.modules['litellm.litellm_core_utils.logging_worker'].GLOBAL_LOGGING_WORKER = Worker()
sys.modules['litellm.litellm_core_utils.litellm_logging'].executor = Executor()
class StubCoroutine:
def __init__(self, logger):
self.logger = logger
@ -106,7 +122,6 @@ class StubCoroutine:
class StubLogger:
def __init__(self):
self.calls = []
self.needed = {}
self.hooks = {}
self.on_enqueue = lambda coroutine: None
@ -147,6 +162,7 @@ logger = StubLogger()
/// A namespace with the stubs, `StubLogger` and a fresh `logger`, after `script` ran in it.
pub(crate) fn namespace<'py>(py: Python<'py>, script: &CStr) -> Bound<'py, PyDict> {
let locals = PyDict::new(py);
locals.set_item("python_contract", PYTHON_CONTRACT).unwrap();
py.run(STUBS, Some(&locals), Some(&locals)).unwrap();
py.run(script, Some(&locals), Some(&locals)).unwrap();
locals
@ -181,6 +197,7 @@ pub(crate) fn legacy_call(
LegacySurface {
call_type: "test",
input_description: "test input",
stream: None,
},
call,
asynchronous,

58
litellm-rust/crates/callbacks-legacy/tests/terminal.rs
View file

@ -1,7 +1,7 @@
use std::ffi::CStr;
use litellm_callbacks::event::{CallEvent, FailureOrigin, Timing};
use litellm_host_python::{AdapterStep, CallbackAdapter, PublicValue};
use litellm_host::event::{FailureOrigin, Timing};
use litellm_host_python::{LifecycleEvent, LifecycleStep, PythonLifecycle};
use pyo3::exceptions::PyRuntimeError;
use pyo3::exceptions::asyncio::CancelledError;
use pyo3::prelude::*;
@ -19,52 +19,52 @@ const TIMING: Timing = Timing {
fn logged(py: Python<'_>, locals: &Bound<'_, PyDict>, asynchronous: bool) -> LegacyLogging {
LegacyLogging {
logger: Some(PythonLogger::new(local(locals, "logger").unbind(), true)),
logger: Some(PythonLogger::new(local(locals, "logger").unbind())),
..legacy_call(py, locals, asynchronous)
}
}
fn succeed(py: Python<'_>, locals: &Bound<'_, PyDict>, logging: &mut LegacyLogging) -> AdapterStep {
fn succeed(
py: Python<'_>,
locals: &Bound<'_, PyDict>,
logging: &mut LegacyLogging,
) -> LifecycleStep {
let response = local(locals, "response").unbind();
logging
.emit(
py,
&CallEvent::Succeeded { timing: TIMING },
Some(PublicValue::Response(&response)),
LifecycleEvent::Succeeded {
timing: TIMING,
response: &response,
},
)
.unwrap()
}
fn fail(py: Python<'_>, locals: &Bound<'_, PyDict>, logging: &mut LegacyLogging) -> AdapterStep {
fn fail(py: Python<'_>, locals: &Bound<'_, PyDict>, logging: &mut LegacyLogging) -> LifecycleStep {
let failure = PyErr::from_value(local(locals, "failure"));
logging
.emit(
py,
&CallEvent::Failed {
LifecycleEvent::Failed {
timing: TIMING,
origin: FailureOrigin::Host,
error: &failure,
},
Some(PublicValue::Error(&failure)),
)
.unwrap()
}
#[rstest]
#[case::sync_listened(false, c"", &["submit"])]
#[case::sync_unlistened(false, c"logger.needed = {'sync_success': False}", &["success_bookkeeping"])]
#[case::async_listened(
true,
c"",
&["async_success_handler", "enqueued", "sync_success_for_async_call"]
)]
#[case::async_unlistened(
true,
c"logger.needed = {'async_success': False, 'sync_success_async': False}",
&["success_bookkeeping"]
)]
#[case::async_deferred(true, c"logger._defer_async_logging = True", &["sync_success_for_async_call"])]
#[case::async_with_fallbacks(true, c"kwargs = {'fallbacks': ['other']}", &["sync_success_for_async_call"])]
fn success_reaches_only_the_callbacks_that_listen(
fn success_reaches_the_logging_handlers(
#[case] asynchronous: bool,
#[case] script: &CStr,
#[case] expected: &[&str],
@ -76,7 +76,7 @@ fn success_reaches_only_the_callbacks_that_listen(
let mut logging = logged(py, &locals, asynchronous);
assert!(matches!(
succeed(py, &locals, &mut logging),
AdapterStep::Done
LifecycleStep::Done
));
let names: Vec<String> = local(&locals, "logger")
.call_method0("names")
@ -109,7 +109,10 @@ fn internal_calls_skip_failure_callbacks_only_when_asynchronous(
internal: true,
..logged(py, &locals, asynchronous)
};
assert!(matches!(fail(py, &locals, &mut logging), AdapterStep::Done));
assert!(matches!(
fail(py, &locals, &mut logging),
LifecycleStep::Done
));
let names: Vec<String> = local(&locals, "logger")
.call_method0("names")
.unwrap()
@ -157,7 +160,7 @@ logger = FailingLogger()
let mut logging = logged(py, &locals, true);
assert!(matches!(
succeed(py, &locals, &mut logging),
AdapterStep::Done
LifecycleStep::Done
));
assert!(
logging
@ -173,14 +176,8 @@ logger = FailingLogger()
#[rstest]
#[case::sync_listened(false, c"", &["failure_handler"])]
#[case::sync_unlistened(false, c"logger.needed = {'sync_failure': False}", &["failure_bookkeeping"])]
#[case::async_listened(true, c"", &["failure_handler", "async_failure_handler"])]
#[case::async_unlistened(
true,
c"logger.needed = {'sync_failure': False, 'async_failure': False}",
&["failure_bookkeeping", "failure_bookkeeping"]
)]
fn failure_reaches_only_the_callbacks_that_listen(
fn failure_reaches_the_logging_handlers(
#[case] asynchronous: bool,
#[case] script: &CStr,
#[case] expected: &[&str],
@ -192,7 +189,10 @@ fn failure_reaches_only_the_callbacks_that_listen(
let mut logging = logged(py, &locals, asynchronous);
let step = fail(py, &locals, &mut logging);
let awaits_async_handler = expected.contains(&"async_failure_handler");
assert_eq!(matches!(step, AdapterStep::Await(_)), awaits_async_handler);
assert_eq!(
matches!(step, LifecycleStep::Await(_)),
awaits_async_handler
);
let names: Vec<String> = local(&locals, "logger")
.call_method0("names")
.unwrap()
@ -227,7 +227,7 @@ logger = FailingLogger()
let mut logging = logged(py, &locals, true);
assert!(matches!(
fail(py, &locals, &mut logging),
AdapterStep::Await(_)
LifecycleStep::Await(_)
));
assert!(
logging
@ -265,7 +265,7 @@ fn the_async_failure_handler_ends_the_call_unless_it_was_cancelled(
};
let expected = result.as_ref().err().map(|error| error.value(py).clone());
match logging.resume(py, result) {
Ok(step) => assert!(done && matches!(step, AdapterStep::Done)),
Ok(step) => assert!(done && matches!(step, LifecycleStep::Done)),
Err(propagated) => {
assert!(!done);
assert!(propagated.value(py).is(expected.unwrap()));

135
litellm-rust/crates/callbacks/src/event.rs
View file

@ -1,135 +0,0 @@
use std::time::{SystemTime, UNIX_EPOCH};
use serde_json::{Map, Value};
/// Seconds since the Unix epoch, on one clock for every host.
pub fn epoch_seconds() -> f64 {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.map(|duration| duration.as_secs_f64())
.unwrap_or(0.0)
}
#[derive(Clone, Copy, Debug, PartialEq)]
pub struct Timing {
pub start_time: f64,
pub end_time: f64,
}
/// The provider request as it is about to leave, offered to the host for rewriting.
#[derive(Clone, Debug, PartialEq)]
pub struct WireRequest {
pub url: String,
pub headers: Vec<(String, String)>,
pub body: Value,
}
/// What the route knows about the request it is sending, for a host that logs it. The
/// route owns these facts; a host reads them beside the wire request and never rewrites
/// them.
#[derive(Clone, Debug, PartialEq)]
pub struct RequestContext {
pub model: String,
pub custom_llm_provider: String,
/// The route's parameters before the provider transformation.
pub optional_params: Value,
pub passthrough_fields: Passthrough,
/// Optional-param names that carry credentials and must be redacted when logged.
pub secret_fields: Vec<String>,
}
/// Body keys whose values are the caller's inputs, unchanged by the route. The only way to
/// build one is to compare the two, so a route cannot name a key it rewrote.
#[derive(Clone, Debug, Default, PartialEq, Eq)]
pub struct Passthrough(Vec<String>);
impl Passthrough {
pub fn unchanged(caller: &Map<String, Value>, body: &Value) -> Self {
Self(
caller
.iter()
.filter(|(name, value)| body.get(name.as_str()) == Some(*value))
.map(|(name, _)| name.clone())
.collect(),
)
}
pub fn iter(&self) -> impl Iterator<Item = &str> {
self.0.iter().map(String::as_str)
}
pub fn contains(&self, name: &str) -> bool {
self.0.iter().any(|field| field == name)
}
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct RawResponse {
pub body: String,
}
/// Whether a failure surfaced inside the call, including a host op the call asked for,
/// or in a host step around it (preparing the arguments, finalizing the response).
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum FailureOrigin {
Call,
Host,
}
#[derive(Clone, Debug, PartialEq)]
pub enum CallEvent {
ResponseReceived {
raw: RawResponse,
},
Succeeded {
timing: Timing,
},
Failed {
timing: Timing,
origin: FailureOrigin,
},
}
#[cfg(test)]
mod tests {
use rstest::rstest;
use serde_json::json;
use super::*;
#[rstest]
#[case::unchanged_scalar(json!({"pages": [0]}), json!({"pages": [0]}), &["pages"])]
#[case::unchanged_explicit_null(json!({"pages": null}), json!({"pages": null}), &["pages"])]
#[case::unchanged_nested_object(
json!({"document": {"type": "document_url", "document_url": "https://a/b.pdf"}}),
json!({"document": {"type": "document_url", "document_url": "https://a/b.pdf"}, "model": "m"}),
&["document"]
)]
#[case::rewritten_value(
json!({"document": {"type": "document_url", "document_url": "https://a/b.pdf"}}),
json!({"document": {"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"}}),
&[]
)]
#[case::dropped_nested_field(
json!({"document": {"type": "image_url", "image_url": "https://a/b.png", "document_name": "b.png"}}),
json!({"document": {"type": "image_url", "image_url": "https://a/b.png"}}),
&[]
)]
#[case::added_nested_field(
json!({"document": {"type": "image_url", "image_url": "https://a/b.png"}}),
json!({"document": {"type": "image_url", "image_url": "https://a/b.png", "detail": "high"}}),
&[]
)]
#[case::reordered_array(json!({"pages": [0, 1]}), json!({"pages": [1, 0]}), &[])]
#[case::consumed_by_the_route(json!({"api_key": "k", "pages": [0]}), json!({"pages": [0]}), &["pages"])]
#[case::added_by_the_route(json!({}), json!({"model": "m"}), &[])]
#[case::non_object_body(json!({"pages": [0]}), json!([{"pages": [0]}]), &[])]
fn passthrough_is_exactly_the_callers_unchanged_keys(
#[case] caller: Value,
#[case] body: Value,
#[case] expected: &[&str],
) {
let passthrough = Passthrough::unchanged(caller.as_object().unwrap(), &body);
assert_eq!(passthrough.iter().collect::<Vec<_>>(), expected);
}
}

1
litellm-rust/crates/core-utils/src/lib.rs
View file

@ -6,4 +6,5 @@ pub mod params;
pub mod prompt_templates;
pub mod secret_redaction;
pub mod serde_compat;
pub mod settings;
pub mod url_utils;

144
litellm-rust/crates/core-utils/src/settings.rs Normal file
View file

@ -0,0 +1,144 @@
use std::str::FromStr;
pub trait Lookup {
fn get(&self, name: &str) -> Option<String>;
fn truthy(&self, name: &str) -> Option<String> {
self.get(name).filter(|value| !value.is_empty())
}
fn enabled(&self, name: &str) -> Option<bool> {
self.get(name)
.is_some_and(|value| value.trim().eq_ignore_ascii_case("true"))
.then_some(true)
}
fn parsed<T: FromStr>(&self, name: &str) -> Option<T>
where
Self: Sized,
{
self.get(name).and_then(|value| value.trim().parse().ok())
}
}
impl<F: Fn(&str) -> Option<String>> Lookup for F {
fn get(&self, name: &str) -> Option<String> {
self(name)
}
}
pub struct ProcessEnvironment;
impl Lookup for ProcessEnvironment {
fn get(&self, name: &str) -> Option<String> {
std::env::var(name).ok()
}
}
pub trait Layer: Default {
fn or(self, lower: Self) -> Self;
}
pub fn merge<L: Layer>(highest_precedence_first: impl IntoIterator<Item = L>) -> L {
highest_precedence_first
.into_iter()
.reduce(L::or)
.unwrap_or_default()
}
#[cfg(test)]
mod tests {
use super::*;
fn env_of(values: &'static [(&'static str, &'static str)]) -> impl Fn(&str) -> Option<String> {
move |name| {
values
.iter()
.find(|(key, _)| *key == name)
.map(|(_, value)| value.to_string())
}
}
#[test]
fn get_keeps_a_present_empty_value_like_os_getenv_with_a_fallback() {
let env = env_of(&[("EMPTY", "")]);
assert_eq!(env.get("EMPTY"), Some(String::new()));
assert_eq!(env.get("ABSENT"), None);
}
#[test]
fn truthy_drops_an_empty_value_like_a_python_or_chain() {
let env = env_of(&[("EMPTY", ""), ("SET", "value")]);
assert_eq!(env.truthy("EMPTY"), None);
assert_eq!(env.truthy("SET").as_deref(), Some("value"));
}
#[test]
fn enabled_only_switches_on_for_true_and_never_forces_off() {
let env = env_of(&[
("LOWER", "true"),
("PADDED", " True "),
("OFF", "false"),
("ONE", "1"),
]);
assert_eq!(env.enabled("LOWER"), Some(true));
assert_eq!(env.enabled("PADDED"), Some(true));
assert_eq!(env.enabled("OFF"), None);
assert_eq!(env.enabled("ONE"), None);
assert_eq!(env.enabled("ABSENT"), None);
}
#[test]
fn parsed_trims_and_skips_values_that_do_not_parse() {
let env = env_of(&[("PADDED", " 45 "), ("WORD", "soon"), ("FRACTION", "0.5")]);
assert_eq!(env.parsed::<u32>("PADDED"), Some(45));
assert_eq!(env.parsed::<u32>("WORD"), None);
assert_eq!(env.parsed::<f64>("FRACTION"), Some(0.5));
assert_eq!(env.parsed::<u32>("ABSENT"), None);
}
#[derive(Debug, Default, PartialEq)]
struct Pair {
first: Option<u8>,
second: Option<u8>,
}
impl Layer for Pair {
fn or(self, lower: Self) -> Self {
Self {
first: self.first.or(lower.first),
second: self.second.or(lower.second),
}
}
}
#[test]
fn merge_takes_each_field_from_the_highest_layer_that_sets_it() {
let merged = merge([
Pair {
first: Some(1),
second: None,
},
Pair {
first: Some(2),
second: Some(2),
},
Pair {
first: Some(3),
second: Some(3),
},
]);
assert_eq!(
merged,
Pair {
first: Some(1),
second: Some(2),
}
);
}
#[test]
fn merging_no_layers_yields_the_empty_layer() {
assert_eq!(merge(Vec::<Pair>::new()), Pair::default());
}
}

7
litellm-rust/crates/core/AGENTS.md
View file

@ -5,10 +5,11 @@ litellm-core is the LiteLLM SDK in Rust — it makes the LLM call. Each top-leve
Each crate mirrors one top-level Python package, so a Rust path reads as its Python path with the crate name in place of the package directory. Dependencies only point down:
- `litellm-types` mirrors `litellm/types/`: pure serde data, no I/O
- `litellm-core-utils` mirrors `litellm/litellm_core_utils/`: pure helpers (provider resolution, prompt factory, call arguments), no network I/O
- `litellm-llms` mirrors `litellm/llms/`: `base_llm/<api>/transformation.rs`, `<provider>/<api>/transformation.rs`, and `custom_httpx/` (HTTP helpers and the OCR request handler)
- `litellm-core-utils` mirrors `litellm/litellm_core_utils/`: pure helpers (provider resolution, prompt factory, call arguments, settings lookup and layer merge), no network I/O
- `litellm-http` is Rust-only and route-neutral: settings resolution, the pooled `reqwest` clients, TLS, proxies, the SSRF-safe media fetcher, request and header helpers, and transport errors. Python's `litellm/llms/custom_httpx/` is split by responsibility instead of mirrored: its transport half lives here, its OCR handler in `litellm-llms`
- `litellm-llms` mirrors `litellm/llms/`: `base_llm/<api>/transformation.rs`, `<provider>/<api>/transformation.rs`, and `base_llm/ocr/handler.rs` (the OCR request handler)
- `litellm-core` mirrors the route packages (`litellm/ocr/`, `litellm/messages/`, ...): entrypoints, route request types, provider dispatch, the route machine, and hooks
A route module owns the call entrypoint, route request types (`*Request<'a>`), credential fallback, provider dispatch, and the handler glue that runs a provider config. Provider code never imports from core; when it needs the caller's hooks mid-call it goes through `litellm_llms::custom_httpx::llm_http_handler::CallHooks`, which each route implements over its host. Import every item from its canonical path. Never re-export another crate's items or give an item a second public path; the only re-export allowed is a private submodule surfacing its item at its module root (`mod error; pub use error::Error;`). Handlers belong in core or llms, never in a host crate
A route module owns the call entrypoint, route request types (`*Request<'a>`), credential fallback, provider dispatch, and the handler glue that runs a provider config. Provider code never imports from core; when it needs the caller's hooks mid-call it goes through `litellm_llms::base_llm::ocr::handler::CallHooks`, which each route implements over its host. Import every item from its canonical path. Never re-export another crate's items or give an item a second public path; the only re-export allowed is a private submodule surfacing its item at its module root (`mod error; pub use error::Error;`). Handlers belong in core or llms, never in a host crate
Not here: serving HTTP (axum routes, extractors), config file reading, rollout state, databases, or callback execution of any kind. Core runs each route as a machine that yields host operations and call events; which integrations consume those events is the host's business.

4
litellm-rust/crates/core/Cargo.toml
View file

@ -9,12 +9,13 @@ autotests = false
[dependencies]
litellm-types.workspace = true
litellm-core-utils.workspace = true
litellm-callbacks.workspace = true
litellm-host.workspace = true
bytes.workspace = true
futures-util.workspace = true
base64.workspace = true
litellm-auth.workspace = true
litellm-auth-aws.workspace = true
litellm-http.workspace = true
litellm-llms.workspace = true
moka.workspace = true
mime_guess = "2.0.5"
@ -35,6 +36,7 @@ url.workspace = true
veil.workspace = true
[dev-dependencies]
litellm-auth-gcp.workspace = true
litellm-llms = { workspace = true, features = ["test-support"] }
rstest.workspace = true
rstest_reuse.workspace = true

6
litellm-rust/crates/core/src/audio_transcription/error.rs
View file

@ -20,9 +20,11 @@ pub enum Error {
#[error(transparent)]
Auth(#[from] litellm_auth::Error),
#[error(transparent)]
Transport(#[from] litellm_llms::custom_httpx::transport::Error),
Transport(#[from] litellm_http::transport::Error),
#[error(transparent)]
Headers(#[from] litellm_llms::custom_httpx::http_handler::HeaderError),
Headers(#[from] litellm_http::request::HeaderError),
#[error(transparent)]
Http(#[from] litellm_http::Error),
#[error(transparent)]
Aws(#[from] litellm_auth_aws::Error),
}

73
litellm-rust/crates/core/src/audio_transcription/handler.rs
View file

@ -1,4 +1,4 @@
use litellm_llms::custom_httpx::http_handler::{http_request, truncate_error_body};
use litellm_http::request::truncate_error_body;
use serde_json::Value;
use super::{Error, client::http_client};
@ -7,34 +7,29 @@ use crate::audio_transcription::types::ProviderAudioTranscriptionRequest;
pub async fn execute_audio_transcription_provider_call(
request: ProviderAudioTranscriptionRequest,
) -> Result<Value, Error> {
let body = serde_json::to_vec(&request.body)
.map_err(|error| Error::InvalidRequest(format!("invalid audio request body: {error}")))?;
let headers = signed_headers(&request, &body).await?;
let mut request_builder = http_client().post(&request.url).body(body);
for (key, value) in headers {
request_builder = request_builder.header(key, value);
}
if let Some(duration) = request.timeout {
request_builder = request_builder.timeout(duration);
}
let response = http_request(request_builder).await.map_err(|error| {
Error::Transport(litellm_llms::custom_httpx::transport::Error::Network(
error.to_string(),
))
let response = crate::outbound::outbound_request::<Error>(
&request.auth,
request.url.clone(),
request.upstream_headers.clone(),
&request.body,
request.timeout,
&request.optional_params,
)
.await?
.send(http_client())
.await
.map_err(|error| {
Error::Transport(litellm_http::transport::Error::Network(error.to_string()))
})?;
let status = response.status();
let text = response.text().await.map_err(|error| {
Error::Transport(litellm_llms::custom_httpx::transport::Error::Network(
error.to_string(),
))
Error::Transport(litellm_http::transport::Error::Network(error.to_string()))
})?;
if !status.is_success() {
return Err(Error::Transport(
litellm_llms::custom_httpx::transport::Error::Http {
status: status.as_u16(),
body: truncate_error_body(&text),
},
));
return Err(Error::Transport(litellm_http::transport::Error::Http {
status: status.as_u16(),
body: truncate_error_body(&text),
}));
}
let response_json = serde_json::from_str(&text)
.map_err(|error| Error::InvalidResponse(format!("invalid audio response JSON: {error}")))?;
@ -43,33 +38,3 @@ pub async fn execute_audio_transcription_provider_call(
.transform_audio_transcription_response(&request.model, response_json)?
.into_json())
}
async fn signed_headers(
request: &ProviderAudioTranscriptionRequest,
body: &[u8],
) -> Result<Vec<(String, String)>, Error> {
use std::{collections::BTreeMap, time::SystemTime};
use litellm_auth_aws::{aws_auth_config, resolve_credentials, sign_bedrock_post};
use litellm_llms::base_llm::audio_transcription::transformation::AudioTranscriptionAuth;
let AudioTranscriptionAuth::AwsSigV4 { region, .. } = &request.auth else {
return Ok(request.upstream_headers.clone());
};
let env_lookup = |key: &str| std::env::var(key).ok();
let credentials = resolve_credentials(
aws_auth_config(&request.optional_params, &env_lookup),
&env_lookup,
)
.await?;
let unsigned: BTreeMap<String, String> = request.upstream_headers.iter().cloned().collect();
let signature = sign_bedrock_post(
&request.url,
body,
&unsigned,
region,
&credentials,
SystemTime::now(),
)?;
Ok(unsigned.into_iter().chain(signature).collect())
}

19
litellm-rust/crates/core/src/audio_transcription/prepare.rs
View file

@ -1,10 +1,8 @@
use litellm_core_utils::get_llm_provider_logic::{CustomLlmProvider, get_custom_llm_provider};
use litellm_http::request::{has_header, string_headers};
use litellm_llms::{
base_llm::audio_transcription::transformation::{
AudioTranscriptionAuth, BaseAudioTranscriptionConfig,
},
base_llm::audio_transcription::transformation::{BaseAudioTranscriptionConfig, RequestAuth},
bedrock::audio_transcription::BEDROCK_AUDIO_TRANSCRIPTION_CONFIG,
custom_httpx::http_handler::{has_header, string_headers},
};
use super::Error;
@ -43,11 +41,14 @@ pub fn prepare_audio_transcription_provider_call(
let env_lookup = |key: &str| std::env::var(key).ok();
let mut headers = string_headers("audio transcription", request.extra_headers)?;
let auth = config.auth_strategy(&model, &request.optional_params, &env_lookup)?;
if matches!(auth, AudioTranscriptionAuth::Bearer)
&& !has_header(&headers, "authorization")
&& let Some(api_key) = request.api_key
{
headers.push(("Authorization".to_string(), format!("Bearer {api_key}")));
match &auth {
RequestAuth::Bearer { token } if !has_header(&headers, "authorization") => {
headers.push(("Authorization".to_string(), format!("Bearer {token}")));
}
RequestAuth::Header { name, value } if !has_header(&headers, name) => {
headers.push(((*name).to_string(), value.clone()));
}
RequestAuth::Bearer { .. } | RequestAuth::Header { .. } | RequestAuth::AwsSigV4 { .. } => {}
}
if !has_header(&headers, "content-type") {
headers.push(("Content-Type".to_string(), "application/json".to_string()));

4
litellm-rust/crates/core/src/audio_transcription/types.rs
View file

@ -1,7 +1,7 @@
use std::time::Duration;
use litellm_llms::base_llm::audio_transcription::transformation::{
AudioTranscriptionAuth, BaseAudioTranscriptionConfig,
BaseAudioTranscriptionConfig, RequestAuth,
};
use serde_json::{Map, Value};
@ -24,7 +24,7 @@ pub struct ProviderAudioTranscriptionRequest {
pub url: String,
pub body: Value,
pub upstream_headers: Vec<(String, String)>,
pub auth: AudioTranscriptionAuth,
pub auth: RequestAuth,
pub optional_params: Map<String, Value>,
pub timeout: Option<Duration>,
}

2
litellm-rust/crates/core/src/chat_completions/common_utils.rs
View file

@ -1,8 +1,8 @@
use litellm_http::request::string_headers as shared_string_headers;
use litellm_llms::{
anthropic::chat::transformation::ANTHROPIC_CHAT_COMPLETIONS_CONFIG,
base_llm::chat::transformation::BaseConfig,
bedrock::chat::converse_transformation::BEDROCK_CHAT_COMPLETIONS_CONFIG,
custom_httpx::http_handler::string_headers as shared_string_headers,
};
use serde_json::{Map, Value};

6
litellm-rust/crates/core/src/chat_completions/error.rs
View file

@ -20,9 +20,11 @@ pub enum Error {
#[error(transparent)]
Auth(#[from] litellm_auth::Error),
#[error(transparent)]
Transport(#[from] litellm_llms::custom_httpx::transport::Error),
Transport(#[from] litellm_http::transport::Error),
#[error(transparent)]
Headers(#[from] litellm_llms::custom_httpx::http_handler::HeaderError),
Headers(#[from] litellm_http::request::HeaderError),
#[error(transparent)]
Http(#[from] litellm_http::Error),
#[error(transparent)]
Aws(#[from] litellm_auth_aws::Error),
}

118
litellm-rust/crates/core/src/chat_completions/handler.rs
View file

@ -1,7 +1,5 @@
use litellm_llms::{
base_llm::chat::transformation::{ChatCompletionsAuth, ProviderChatResponseData},
custom_httpx::http_handler::{http_request, truncate_error_body},
};
use litellm_http::{outbound::OutboundRequest, request::truncate_error_body};
use litellm_llms::base_llm::chat::transformation::ProviderChatResponseData;
use litellm_types::utils::ChatCompletionsResponse;
use serde_json::Value;
@ -14,50 +12,29 @@ pub(super) async fn execute_chat_completions_provider_call(
request: ResolvedChatCompletionsRequest<'_>,
) -> Result<ChatCompletionsResponse, Error> {
let request = prepare_provider_request(request)?;
let body = serde_json::to_vec(&request.body).map_err(|err| {
Error::InvalidRequest(format!(
"failed to serialize chat completions request: {err}"
))
})?;
let headers = signed_headers(&request, &body).await?;
let outbound = outbound_request(&request).await?;
let mut request_builder = http_client().post(&request.url).body(body);
for (key, value) in &headers {
request_builder = request_builder.header(key, value);
}
if let Some(duration) = request.timeout {
request_builder = request_builder.timeout(duration);
}
let response = http_request(request_builder).await.map_err(|err| {
let response = outbound.send(http_client()).await.map_err(|err| {
// Failing to establish the connection means the request never went out,
// so the host can still serve it. Everything else here, a timeout
// above all, may have reached the provider and been answered.
if err.is_connect() || err.is_builder() {
Error::Transport(litellm_llms::custom_httpx::transport::Error::Connect(
err.to_string(),
))
Error::Transport(litellm_http::transport::Error::Connect(err.to_string()))
} else {
Error::Transport(litellm_llms::custom_httpx::transport::Error::Network(
err.to_string(),
))
Error::Transport(litellm_http::transport::Error::Network(err.to_string()))
}
})?;
let status = response.status();
let text = response.text().await.map_err(|err| {
Error::Transport(litellm_llms::custom_httpx::transport::Error::Network(
err.to_string(),
))
Error::Transport(litellm_http::transport::Error::Network(err.to_string()))
})?;
if !status.is_success() {
return Err(Error::Transport(
litellm_llms::custom_httpx::transport::Error::Http {
status: status.as_u16(),
body: truncate_error_body(&text),
},
));
return Err(Error::Transport(litellm_http::transport::Error::Http {
status: status.as_u16(),
body: truncate_error_body(&text),
}));
}
let body: Value = serde_json::from_str(&text).map_err(|err| {
@ -82,64 +59,29 @@ pub(super) async fn execute_chat_completions_provider_call(
pub(super) fn as_response_error(err: Error) -> Error {
match err {
already @ (Error::InvalidResponse(_)
| Error::Transport(litellm_llms::custom_httpx::transport::Error::Http {
..
})) => already,
| Error::Transport(litellm_http::transport::Error::Http { .. })) => already,
other => Error::InvalidResponse(other.to_string()),
}
}
pub(super) async fn signed_headers(
pub(super) async fn outbound_request(
request: &ProviderChatCompletionsRequest,
body: &[u8],
) -> Result<Vec<(String, String)>, Error> {
use std::{collections::BTreeMap, time::SystemTime};
use litellm_auth_aws::{
aws_auth_config, aws_signature_headers, host_supplied_credentials,
is_sigv4_computed_header, resolve_credentials, sign_bedrock_post,
};
let ChatCompletionsAuth::AwsSigV4 { region } = &request.auth else {
return Ok(request.upstream_headers.clone());
};
// Reattaching a header the signer also emits would put both copies on the
// wire, and Bedrock rejects that pair. Python instead drops the caller's
// copy and prefers a forwarded Authorization over the signature, so leave
// the request to Python rather than serving it a different way here.
if request
.upstream_headers
.iter()
.any(|(name, _)| is_sigv4_computed_header(name))
{
return Err(Error::Unsupported(
"request forwards a header AWS SigV4 computes",
));
}
let env_lookup = |key: &str| std::env::var(key).ok();
let unsigned: BTreeMap<String, String> = request.upstream_headers.iter().cloned().collect();
// A host with its own resolution chain hands the result down; only fall
// back to deriving credentials here when it supplied none.
let credentials = match host_supplied_credentials(&request.optional_params) {
Some(credentials) => credentials,
None => {
resolve_credentials(
aws_auth_config(&request.optional_params, &env_lookup),
&env_lookup,
)
.await?
) -> Result<OutboundRequest, Error> {
crate::outbound::outbound_request(
&request.auth,
request.url.clone(),
request.upstream_headers.clone(),
&request.body,
request.timeout,
&request.optional_params,
)
.await
.map_err(|error| match error {
// Python drops the caller's copy and prefers a forwarded Authorization
// over the signature, so leave the request to it.
Error::Http(litellm_http::Error::ComputedHeader(_)) => {
Error::Unsupported("request forwards a header AWS SigV4 computes")
}
};
let signature = sign_bedrock_post(
&request.url,
body,
&aws_signature_headers(&unsigned),
region,
&credentials,
SystemTime::now(),
)?;
// Every original header goes back on the wire alongside the computed ones,
// as Python reattaches them. The guard above already rejected the names
// that would collide, so no name appears twice.
Ok(unsigned.into_iter().chain(signature).collect())
other => other,
})
}

14
litellm-rust/crates/core/src/chat_completions/prepare.rs
View file

@ -1,8 +1,6 @@
use litellm_core_utils::get_llm_provider_logic::{CustomLlmProvider, get_custom_llm_provider};
use litellm_llms::{
base_llm::chat::transformation::{BaseConfig, ChatCompletionsAuth},
custom_httpx::http_handler::has_header,
};
use litellm_http::request::has_header;
use litellm_llms::base_llm::chat::transformation::{BaseConfig, RequestAuth};
use litellm_types::llms::openai::ChatMessage;
use serde_json::Value;
@ -69,7 +67,7 @@ fn validate_environment(
request: &ResolvedChatCompletionsRequest<'_>,
model: &str,
config: &dyn BaseConfig,
) -> Result<(Vec<(String, String)>, ChatCompletionsAuth), Error> {
) -> Result<(Vec<(String, String)>, RequestAuth), Error> {
let env_lookup = |key: &str| std::env::var(key).ok();
let mut headers = string_headers(request.extra_headers.clone())?;
let auth = config.auth(
@ -79,7 +77,7 @@ fn validate_environment(
&env_lookup,
)?;
match &auth {
ChatCompletionsAuth::Header { name, value } => {
RequestAuth::Header { name, value } => {
// The deployment's credential replaces whatever the caller forwarded
// under the same name, mirroring Python's
// `{**headers, **anthropic_headers}`: letting a request header win
@ -94,7 +92,7 @@ fn validate_environment(
headers.push(((*name).to_string(), value.clone()));
}
}
ChatCompletionsAuth::Bearer { token } => {
RequestAuth::Bearer { token } => {
// Bedrock's `get_request_headers` assigns `headers["Authorization"]`
// unconditionally once a bearer token resolves, so the deployment's
// identity outranks whatever the caller forwarded. Keeping the
@ -107,7 +105,7 @@ fn validate_environment(
headers.push(("authorization".to_string(), format!("Bearer {token}")));
}
// SigV4 signs the serialized body, so the handler adds its headers.
ChatCompletionsAuth::AwsSigV4 { .. } => {}
RequestAuth::AwsSigV4 { .. } => {}
}
for (name, value) in config.default_headers() {

45
litellm-rust/crates/core/src/chat_completions/tests.rs
View file

@ -1,4 +1,4 @@
use litellm_llms::base_llm::chat::transformation::ChatCompletionsAuth;
use litellm_llms::base_llm::chat::transformation::RequestAuth;
use serde_json::{Map, Value, json};
use super::{
@ -90,7 +90,7 @@ fn adds_the_auth_and_default_headers() {
);
assert!(matches!(
prepared.auth,
ChatCompletionsAuth::Header {
RequestAuth::Header {
name: "x-api-key",
..
}
@ -265,7 +265,7 @@ fn rejects_non_string_extra_headers() {
call.extra_headers = Some(Map::from_iter([("x-trace".to_string(), json!(7))]));
assert_eq!(
decline(call),
Error::Headers(litellm_llms::custom_httpx::http_handler::HeaderError {
Error::Headers(litellm_http::request::HeaderError {
context: "chat completions",
name: "x-trace".to_string(),
actual: "number",
@ -289,8 +289,9 @@ fn prepares_a_bedrock_call_without_resolving_credentials() {
);
assert_eq!(
prepared.auth,
ChatCompletionsAuth::AwsSigV4 {
region: "us-east-1".to_string()
RequestAuth::AwsSigV4 {
region: "us-east-1".to_string(),
service: "bedrock",
}
);
// SigV4 signs the serialized body, so prepare must not have added an
@ -326,15 +327,14 @@ async fn a_forwarded_client_header_does_not_enter_the_bedrock_signature() {
json!("abc-123"),
)]));
let prepared = prepare_chat_completions_call(call).expect("prepares");
let signed = super::handler::signed_headers(&prepared, br#"{"a":1}"#)
let signed = super::handler::outbound_request(&prepared)
.await
.expect("signs");
let authorization = signed
.iter()
.find(|(name, _)| name.eq_ignore_ascii_case("authorization"))
.map(|(_, value)| value.clone())
.expect("carries an authorization header");
.header("authorization")
.expect("carries an authorization header")
.to_string();
assert!(
authorization.starts_with("AWS4-HMAC-SHA256"),
"expected a SigV4 signature, got {authorization}"
@ -346,6 +346,7 @@ async fn a_forwarded_client_header_does_not_enter_the_bedrock_signature() {
// It still goes on the wire, it is just not part of the signature.
assert!(
signed
.headers()
.iter()
.any(|(name, value)| name == "x-request-id" && value == "abc-123"),
"forwarded header was dropped instead of reattached"
@ -376,7 +377,7 @@ async fn a_forwarded_header_the_signer_computes_declines_to_python() {
call.api_key = None;
call.extra_headers = Some(Map::from_iter([(forwarded.to_string(), json!("forged"))]));
let prepared = prepare_chat_completions_call(call).expect("prepares");
let error = super::handler::signed_headers(&prepared, br#"{"a":1}"#)
let error = super::handler::outbound_request(&prepared)
.await
.expect_err("{forwarded} should decline instead of being signed");
assert!(
@ -466,7 +467,7 @@ fn a_bedrock_api_key_is_sent_as_a_bearer_token_instead_of_being_signed() {
.expect("prepares");
assert_eq!(
prepared.auth,
ChatCompletionsAuth::Bearer {
RequestAuth::Bearer {
token: "sk-test".to_string()
}
);
@ -771,10 +772,7 @@ mod round_trip {
assert!(
matches!(
err,
Error::Transport(litellm_llms::custom_httpx::transport::Error::Http {
status: 429,
..
})
Error::Transport(litellm_http::transport::Error::Http { status: 429, .. })
),
"expected a 429, got {err:?}"
);
@ -801,7 +799,7 @@ mod round_trip {
assert!(
matches!(
err,
Error::Transport(litellm_llms::custom_httpx::transport::Error::Connect(_))
Error::Transport(litellm_http::transport::Error::Connect(_))
),
"expected a pre-send connect failure, got {err:?}"
);
@ -825,16 +823,11 @@ mod round_trip {
}
// An upstream status is already unambiguous, so it survives intact.
assert!(matches!(
as_response_error(Error::Transport(
litellm_llms::custom_httpx::transport::Error::Http {
status: 500,
body: "boom".to_string()
}
)),
Error::Transport(litellm_llms::custom_httpx::transport::Error::Http {
as_response_error(Error::Transport(litellm_http::transport::Error::Http {
status: 500,
..
})
body: "boom".to_string()
})),
Error::Transport(litellm_http::transport::Error::Http { status: 500, .. })
));
}
}

4
litellm-rust/crates/core/src/chat_completions/types.rs
View file

@ -1,6 +1,6 @@
use std::time::Duration;
use litellm_llms::base_llm::chat::transformation::{BaseConfig, ChatCompletionsAuth};
use litellm_llms::base_llm::chat::transformation::{BaseConfig, RequestAuth};
use litellm_types::llms::openai::ChatMessage;
use serde_json::{Map, Value};
@ -38,7 +38,7 @@ pub struct ProviderChatCompletionsRequest {
pub url: String,
pub body: Value,
pub upstream_headers: Vec<(String, String)>,
pub auth: ChatCompletionsAuth,
pub auth: RequestAuth,
pub optional_params: Map<String, Value>,
pub timeout: Option<Duration>,
}

2
litellm-rust/crates/core/src/lib.rs
View file

@ -2,9 +2,9 @@ pub mod audio_transcription;
pub mod chat_completions;
pub mod constants;
pub mod error;
pub mod machine;
pub mod messages;
pub mod ocr;
mod outbound;
pub mod responses;
pub use error::Error;

6
litellm-rust/crates/core/src/messages/common_utils.rs
View file

@ -1,11 +1,9 @@
pub(super) use litellm_llms::custom_httpx::http_handler::{
has_bearer_auth, has_header, truncate_error_body,
};
use litellm_http::request::string_headers as shared_string_headers;
pub(super) use litellm_http::request::{has_bearer_auth, has_header, truncate_error_body};
use litellm_llms::{
anthropic::experimental_pass_through::messages::transformation::ANTHROPIC_MESSAGES_CONFIG,
azure_ai::anthropic::messages_transformation::AZURE_ANTHROPIC_MESSAGES_CONFIG,
base_llm::anthropic_messages::transformation::BaseAnthropicMessagesConfig,
custom_httpx::http_handler::string_headers as shared_string_headers,
};
use serde_json::{Map, Value};

4
litellm-rust/crates/core/src/messages/error.rs
View file

@ -15,9 +15,9 @@ pub enum Error {
#[error(transparent)]
Auth(#[from] litellm_auth::Error),
#[error(transparent)]
Transport(#[from] litellm_llms::custom_httpx::transport::Error),
Transport(#[from] litellm_http::transport::Error),
#[error(transparent)]
Headers(#[from] litellm_llms::custom_httpx::http_handler::HeaderError),
Headers(#[from] litellm_http::request::HeaderError),
}
impl From<LlmError> for Error {

126
litellm-rust/crates/core/src/messages/handler.rs
View file

@ -1,88 +1,52 @@
use litellm_llms::custom_httpx::http_handler::http_request;
use std::time::Duration;
use litellm_http::{request::http_request, transport::Error as TransportError};
use litellm_llms::base_llm::anthropic_messages::transformation::BaseAnthropicMessagesConfig;
use litellm_types::llms::anthropic_messages::anthropic_response::AnthropicMessagesResponse;
use serde_json::Value;
use super::{
Error, client::http_client, common_utils::truncate_error_body,
prepare::prepare_provider_request,
};
use crate::{constants::ANTHROPIC_MESSAGES_PROVIDER, messages::types::MessagesRequest};
use super::{Error, client::http_client, common_utils::truncate_error_body};
pub(super) async fn execute_messages_provider_call(
request: MessagesRequest<'_>,
pub(super) fn network(error: reqwest::Error) -> Error {
Error::Transport(TransportError::Network(error.to_string()))
}
pub(super) async fn send(
url: &str,
headers: &[(String, String)],
body: &Value,
timeout: Option<Duration>,
) -> Result<reqwest::Response, Error> {
let builder = headers.iter().fold(
http_client().post(url).json(body),
|builder, (key, value)| builder.header(key, value),
);
let builder = match timeout {
Some(duration) => builder.timeout(duration),
None => builder,
};
http_request(builder).await.map_err(network)
}
pub(super) async fn provider_error(response: reqwest::Response) -> Error {
let status = response.status().as_u16();
match response.text().await {
Ok(text) => Error::Transport(TransportError::Http {
status,
body: truncate_error_body(&text),
}),
Err(error) => network(error),
}
}
pub(super) fn decode_response(
config: &dyn BaseAnthropicMessagesConfig,
model: &str,
text: &str,
) -> Result<AnthropicMessagesResponse, Error> {
let request = prepare_provider_request(request)?;
let mut request_builder = http_client().post(&request.url).json(&request.body);
for (key, value) in &request.upstream_headers {
request_builder = request_builder.header(key, value);
}
if let Some(duration) = request.timeout {
request_builder = request_builder.timeout(duration);
}
let response = http_request(request_builder).await.map_err(|err| {
Error::Transport(litellm_llms::custom_httpx::transport::Error::Network(
err.to_string(),
))
})?;
let status = response.status();
let text = response.text().await.map_err(|err| {
Error::Transport(litellm_llms::custom_httpx::transport::Error::Network(
err.to_string(),
))
})?;
if !status.is_success() {
return Err(Error::Transport(
litellm_llms::custom_httpx::transport::Error::Http {
status: status.as_u16(),
body: truncate_error_body(&text),
},
));
}
let response = serde_json::from_str(&text)
let response = serde_json::from_str(text)
.map_err(|err| Error::InvalidResponse(format!("invalid messages response JSON: {err}")))?;
request
.config
.transform_anthropic_messages_response(&request.model, response)
config
.transform_anthropic_messages_response(model, response)
.map_err(Error::from)
}
pub(super) async fn execute_messages_provider_stream(
request: MessagesRequest<'_>,
) -> Result<reqwest::Response, Error> {
let request = prepare_provider_request(request)?;
if request.provider != ANTHROPIC_MESSAGES_PROVIDER {
return Err(Error::Unsupported("streaming messages for this provider"));
}
let mut request_builder = http_client().post(&request.url).json(&request.body);
for (key, value) in &request.upstream_headers {
request_builder = request_builder.header(key, value);
}
if let Some(duration) = request.timeout {
request_builder = request_builder.timeout(duration);
}
let response = http_request(request_builder).await.map_err(|err| {
Error::Transport(litellm_llms::custom_httpx::transport::Error::Network(
err.to_string(),
))
})?;
let status = response.status();
if !status.is_success() {
let text = response.text().await.map_err(|err| {
Error::Transport(litellm_llms::custom_httpx::transport::Error::Network(
err.to_string(),
))
})?;
return Err(Error::Transport(
litellm_llms::custom_httpx::transport::Error::Http {
status: status.as_u16(),
body: truncate_error_body(&text),
},
));
}
Ok(response)
}

36
litellm-rust/crates/core/src/messages/mod.rs
View file

@ -1,11 +1,8 @@
//! The Anthropic Messages call, the Rust equivalent of Python's
//! `litellm.messages()`.
//!
//! [`messages`] is the top-level entrypoint: give it a model, a body, and
//! credentials, and it resolves the provider, transforms the request, calls the
//! provider, and returns a typed non-streaming response. [`messages_stream`]
//! is the streaming variant; it hands the raw upstream response back so a host
//! can splice the event stream to its own caller.
//! [`route`] is the call as a machine a host drives, streaming or not. [`messages`] runs
//! it in process for a caller that already holds the request and wants the message.
mod error;
pub mod types;
@ -14,17 +11,34 @@ mod client;
mod common_utils;
mod handler;
mod prepare;
use handler::{execute_messages_provider_call, execute_messages_provider_stream};
pub mod route;
use litellm_types::llms::anthropic_messages::anthropic_response::AnthropicMessagesResponse;
use route::{LocalMessagesHost, MessagesCall, MessagesOutput, messages_machine};
use serde_json::Value;
use crate::messages::types::MessagesRequest;
pub async fn messages(request: MessagesRequest<'_>) -> Result<AnthropicMessagesResponse, Error> {
execute_messages_provider_call(request).await
}
pub async fn messages_stream(request: MessagesRequest<'_>) -> Result<reqwest::Response, Error> {
execute_messages_provider_stream(request).await
let Value::Object(body) = request.body else {
return Err(Error::InvalidRequest(
"messages body must be an object".into(),
));
};
let call = MessagesCall {
model: request.model.into(),
body,
api_key: request.api_key.map(Into::into),
api_base: request.api_base.map(Into::into),
custom_llm_provider: request.custom_llm_provider.map(Into::into),
extra_headers: request.extra_headers,
timeout: request.timeout,
};
match litellm_host::run::run(messages_machine(), &LocalMessagesHost::new(call)).await? {
MessagesOutput::Message(message) => Ok(*message),
MessagesOutput::Streamed => Err(Error::Unsupported(
"streamed responses need a streaming host",
)),
}
}
#[cfg(test)]

11
litellm-rust/crates/core/src/messages/prepare.rs
View file

@ -2,6 +2,7 @@ use litellm_core_utils::get_llm_provider_logic::{CustomLlmProvider, get_custom_l
use litellm_llms::base_llm::anthropic_messages::transformation::{
BaseAnthropicMessagesConfig, MessagesAuthStrategy,
};
use litellm_types::llms::anthropic_messages::anthropic_request::AnthropicMessagesRequest;
use serde_json::{Map, Value};
use super::{
@ -37,10 +38,14 @@ pub(super) fn prepare_provider_request(
let headers =
validate_environment(config, request.extra_headers, request.api_key, &env_lookup)?;
let typed_request = serde_json::from_value(request.body).map_err(|err| {
Error::InvalidRequest(format!("invalid Anthropic messages request: {err}"))
let typed_request: AnthropicMessagesRequest =
serde_json::from_value(request.body).map_err(|err| {
Error::InvalidRequest(format!("invalid Anthropic messages request: {err}"))
})?;
let transformed = config.transform_anthropic_messages_request(AnthropicMessagesRequest {
model: model.clone(),
..typed_request
})?;
let transformed = config.transform_anthropic_messages_request(typed_request)?;
let body = serde_json::to_value(transformed).map_err(|err| {
Error::InvalidRequest(format!(
"failed to serialize Anthropic messages request: {err}"

196
litellm-rust/crates/core/src/messages/route.rs Normal file
View file

@ -0,0 +1,196 @@
use std::{sync::Mutex, time::Duration};
use bytes::Bytes;
use litellm_auth::SecretValue;
use litellm_core_utils::get_llm_provider_logic::get_custom_llm_provider;
use litellm_host::{
event::{MachineEvent, RawResponse, RequestContext, WireRequest},
host::{Demand, Host},
machine::{HostChannel, MachineFault, RouteMachine},
route::Route,
};
use litellm_types::llms::anthropic_messages::anthropic_response::AnthropicMessagesResponse;
use serde_json::{Map, Value};
use super::{
Error,
common_utils::messages_provider_config,
handler::{decode_response, network, provider_error, send},
prepare::prepare_provider_request,
types::MessagesRequest,
};
use crate::constants::ANTHROPIC_MESSAGES_PROVIDER;
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum MessagesOp {
ProjectRequest,
}
pub enum MessagesOpResult {
Request(Box<MessagesCall>),
}
/// The caller's request as the host projects it.
pub struct MessagesCall {
pub model: String,
pub body: Map<String, Value>,
pub api_key: Option<String>,
pub api_base: Option<String>,
pub custom_llm_provider: Option<String>,
pub extra_headers: Option<Map<String, Value>>,
pub timeout: Option<Duration>,
}
impl MessagesCall {
fn streams(&self) -> bool {
self.body.get("stream").and_then(Value::as_bool) == Some(true)
}
}
pub enum MessagesOutput {
Message(Box<AnthropicMessagesResponse>),
/// Every chunk already reached the host through `Deliver`.
Streamed,
}
pub struct Messages;
impl Route for Messages {
type Response = MessagesOutput;
type Error = Error;
type Op = MessagesOp;
type OpResult = MessagesOpResult;
type Chunk = Bytes;
type StreamHead = ();
}
impl From<MachineFault> for Error {
fn from(fault: MachineFault) -> Self {
Self::InvalidRequest(match fault {
MachineFault::Abandoned => "messages host driver was abandoned".into(),
MachineFault::Protocol(message) => format!("messages {message}"),
MachineFault::Mismatch => "invalid messages host operation result".into(),
})
}
}
pub type MessagesHost = HostChannel<Messages>;
pub type MessagesMachine = RouteMachine<Messages>;
/// Whether this route serves the request, decided before any callback runs so a host
/// can still run its own path.
pub fn supports(model: &str, custom_llm_provider: Option<&str>, stream: bool) -> bool {
let provider = get_custom_llm_provider(model, custom_llm_provider)
.map(|resolved| resolved.custom_llm_provider)
.or(custom_llm_provider);
match provider {
Some(ANTHROPIC_MESSAGES_PROVIDER) => true,
Some(provider) => !stream && messages_provider_config(provider).is_some(),
None => false,
}
}
/// The in-process host for a request already in hand. It answers projection once and
/// observes nothing.
pub struct LocalMessagesHost {
call: Mutex<Option<MessagesCall>>,
}
impl LocalMessagesHost {
pub fn new(call: MessagesCall) -> Self {
Self {
call: Mutex::new(Some(call)),
}
}
}
impl Host<Messages> for LocalMessagesHost {
async fn route(&self, op: MessagesOp) -> Result<MessagesOpResult, Error> {
match op {
MessagesOp::ProjectRequest => self
.call
.lock()
.unwrap_or_else(|error| error.into_inner())
.take()
.map(|call| MessagesOpResult::Request(Box::new(call)))
.ok_or_else(|| {
Error::InvalidRequest("messages request was already projected".into())
}),
}
}
}
pub fn messages_machine() -> MessagesMachine {
RouteMachine::new(|host| Box::pin(execute(host)))
}
async fn execute(host: MessagesHost) -> Result<MessagesOutput, Error> {
let MessagesOpResult::Request(call) = host.route(MessagesOp::ProjectRequest).await?;
let stream = call.streams();
let request = prepare_provider_request(MessagesRequest {
model: &call.model,
body: Value::Object(call.body.clone()),
api_key: call.api_key.as_deref(),
api_base: call.api_base.as_deref(),
custom_llm_provider: call.custom_llm_provider.as_deref(),
extra_headers: call.extra_headers.clone(),
timeout: call.timeout,
})?;
if stream && request.provider != ANTHROPIC_MESSAGES_PROVIDER {
return Err(Error::Unsupported("streaming messages for this provider"));
}
let context = RequestContext {
model: request.model.clone(),
custom_llm_provider: request.provider.clone(),
optional_params: Value::Object(
call.body
.iter()
.filter(|(name, _)| !matches!(name.as_str(), "model" | "messages"))
.map(|(name, value)| (name.clone(), value.clone()))
.collect(),
),
secret_fields: Vec::new(),
api_key: call.api_key.clone().map(SecretValue::new),
};
let wire = host
.before_send(
WireRequest {
url: request.url,
headers: request.upstream_headers,
body: request.body,
},
context,
)
.await?;
let response = send(&wire.url, &wire.headers, &wire.body, request.timeout).await?;
if !response.status().is_success() {
return Err(provider_error(response).await);
}
if stream {
return relay(&host, response).await;
}
let text = response.text().await.map_err(network)?;
host.emit(MachineEvent::ResponseReceived {
raw: RawResponse { body: text.clone() },
})
.await?;
decode_response(request.config, &request.model, &text)
.map(|message| MessagesOutput::Message(Box::new(message)))
}
/// Hands each upstream chunk to the caller as it arrives. A caller that stops reading
/// ends the upstream read, and the call completes with what it delivered.
async fn relay(
host: &MessagesHost,
mut response: reqwest::Response,
) -> Result<MessagesOutput, Error> {
if host.open(()).await? == Demand::Detached {
return Ok(MessagesOutput::Streamed);
}
while let Some(chunk) = response.chunk().await.map_err(network)? {
if host.deliver(chunk).await? == Demand::Detached {
break;
}
}
Ok(MessagesOutput::Streamed)
}

4
litellm-rust/crates/core/src/messages/tests.rs
View file

@ -82,7 +82,7 @@ fn string_headers_rejects_non_string_values() {
let err = string_headers(Some(headers)).expect_err("non-string header rejected");
assert_eq!(
err,
Error::Headers(litellm_llms::custom_httpx::http_handler::HeaderError {
Error::Headers(litellm_http::request::HeaderError {
context: "messages",
name: "x-count".to_string(),
actual: "number",
@ -432,7 +432,7 @@ async fn messages_maps_provider_error_status_to_http_error() {
assert!(matches!(
err,
Error::Transport(litellm_llms::custom_httpx::transport::Error::Http { status: 401, .. })
Error::Transport(litellm_http::transport::Error::Http { status: 401, .. })
));
}

16
litellm-rust/crates/core/src/ocr/arguments.rs
View file

@ -15,6 +15,18 @@ const AZURE_AUTH_OPTION_FIELDS: &[&str] = &[
"azure_federated_token_file",
"enable_azure_ad_token_refresh",
];
const AWS_AUTH_OPTION_FIELDS: &[&str] = &[
"aws_access_key_id",
"aws_secret_access_key",
"aws_session_token",
"aws_region_name",
"aws_session_name",
"aws_profile_name",
"aws_role_name",
"aws_web_identity_token",
"aws_sts_endpoint",
"aws_external_id",
];
const VERTEX_AUTH_OPTION_FIELDS: &[&str] = &[
"vertex_credentials",
"vertex_ai_credentials",
@ -35,6 +47,7 @@ pub fn consumed_optional_param_names(
let (model, config) = resolve_provider_config(model, custom_llm_provider)?;
let provider_fields = config.get_supported_ocr_params(&model);
let auth_fields: &[&str] = match config {
OcrConfigKind::AwsTextract | OcrConfigKind::AwsTextractAnalyze => AWS_AUTH_OPTION_FIELDS,
OcrConfigKind::AzureAi
| OcrConfigKind::AzureDocumentIntelligence
| OcrConfigKind::AzureCohere => AZURE_AUTH_OPTION_FIELDS,
@ -57,6 +70,9 @@ pub(crate) fn is_secret_param(name: &str) -> bool {
| "azure_federated_token_file"
| "vertex_credentials"
| "vertex_ai_credentials"
| "aws_secret_access_key"
| "aws_session_token"
| "aws_web_identity_token"
)
}

11
litellm-rust/crates/core/src/ocr/client.rs
View file

@ -1,6 +1,5 @@
use litellm_llms::{
base_llm::ocr::{error::Error, transformation::LiteLLMOcrResponse},
custom_httpx::llm_http_handler::OcrClient,
use litellm_llms::base_llm::ocr::{
error::Error, handler::OcrClient, transformation::LiteLLMOcrResponse,
};
use crate::ocr::{
@ -12,9 +11,5 @@ pub async fn perform(
client: &OcrClient,
request: LiteLLMOcrRequest,
) -> Result<LiteLLMOcrResponse, Error> {
litellm_callbacks::run::run(ocr_machine(client.clone()), &LocalOcrHost::new(request)).await
}
pub async fn ocr(request: LiteLLMOcrRequest) -> Result<LiteLLMOcrResponse, Error> {
perform(&OcrClient::shared()?, request).await
litellm_host::run::run(ocr_machine(client.clone()), &LocalOcrHost::new(request)).await
}

27
litellm-rust/crates/core/src/ocr/handler.rs
View file

@ -1,11 +1,10 @@
use futures_util::future::BoxFuture;
use litellm_callbacks::event::{CallEvent, Passthrough, RawResponse, RequestContext, WireRequest};
use litellm_llms::{
base_llm::ocr::{
error::Error,
transformation::{LiteLLMOcrResponse, PreparedOcrRequest},
},
custom_httpx::llm_http_handler::{CallHooks, OcrClient},
use litellm_auth::SecretValue;
use litellm_host::event::{MachineEvent, RawResponse, RequestContext, WireRequest};
use litellm_llms::base_llm::ocr::{
error::Error,
handler::{CallHooks, OcrClient},
transformation::{LiteLLMOcrResponse, PreparedOcrRequest},
};
use serde_json::Value;
@ -23,7 +22,7 @@ pub(crate) async fn perform_ocr_request(
) -> Result<LiteLLMOcrResponse, Error> {
request.response_format()?;
let config = request.config;
let request = prepare_request(request, caller_document);
let request = prepare_request(request, caller_document, client);
let hooks = OcrCallHooks::new(host.clone(), &request, config);
config.ocr(client, &request, &hooks).await
}
@ -36,6 +35,7 @@ pub(crate) struct OcrCallHooks {
custom_llm_provider: &'static str,
optional_params: Value,
secret_fields: Vec<String>,
api_key: Option<SecretValue>,
}
impl OcrCallHooks {
@ -51,28 +51,25 @@ impl OcrCallHooks {
.filter(|name| is_secret_param(name))
.cloned()
.collect(),
api_key: request.connection.api_key.clone(),
}
}
}
impl CallHooks<Error> for OcrCallHooks {
fn before_send(
&self,
wire: WireRequest,
passthrough_fields: Passthrough,
) -> BoxFuture<'_, Result<WireRequest, Error>> {
fn before_send(&self, wire: WireRequest) -> BoxFuture<'_, Result<WireRequest, Error>> {
let context = RequestContext {
model: self.model.clone(),
custom_llm_provider: self.custom_llm_provider.into(),
optional_params: self.optional_params.clone(),
passthrough_fields,
secret_fields: self.secret_fields.clone(),
api_key: self.api_key.clone(),
};
Box::pin(self.host.before_send(wire, context))
}
fn response_received<'a>(&'a self, body: &'a [u8]) -> BoxFuture<'a, Result<(), Error>> {
Box::pin(self.host.emit(CallEvent::ResponseReceived {
Box::pin(self.host.emit(MachineEvent::ResponseReceived {
raw: RawResponse {
body: String::from_utf8_lossy(body).into_owned(),
},

8
litellm-rust/crates/core/src/ocr/mod.rs
View file

@ -8,6 +8,10 @@ pub mod route;
pub mod types;
pub mod wire;
#[cfg(test)]
#[path = "../../tests/aws_textract_ocr.rs"]
mod aws_textract_tests;
#[cfg(test)]
#[path = "../../tests/azure_ai_ocr.rs"]
mod azure_ai_tests;
@ -21,8 +25,8 @@ mod cohere_tests;
#[path = "../../tests/deepseek_ocr.rs"]
mod deepseek_tests;
#[cfg(test)]
#[path = "../../tests/ocr/passthrough.rs"]
mod passthrough_tests;
#[path = "../../tests/ocr/document.rs"]
mod document_tests;
#[cfg(test)]
#[path = "../../tests/reducto_ocr.rs"]
mod reducto_tests;

48
litellm-rust/crates/core/src/ocr/prepare.rs
View file

@ -1,6 +1,7 @@
use litellm_auth::{InputSource, Sourced};
use litellm_llms::base_llm::ocr::transformation::{
OcrConnection, OcrCredentialInputs, PreparedOcrRequest, credential_env,
use litellm_auth::{InputSource, SecretValue, Sourced};
use litellm_llms::base_llm::ocr::{
handler::OcrClient,
transformation::{OcrConnection, OcrCredentialInputs, PreparedOcrRequest},
};
use super::provider_config::OcrProvider;
@ -9,26 +10,34 @@ use crate::ocr::types::{LiteLLMOcrRequest, ResolvedOcrRequest};
pub(crate) fn prepare_request(
request: ResolvedOcrRequest,
caller_document: bool,
client: &OcrClient,
) -> PreparedOcrRequest {
let credentials = request.credentials.clone();
let api_base_env = match request.config.provider() {
OcrProvider::Mistral => Some("MISTRAL_API_BASE"),
OcrProvider::AzureAi => Some("AZURE_AI_API_BASE"),
OcrProvider::Cohere | OcrProvider::Reducto | OcrProvider::VertexAi => None,
let (preferred_api_key_env, api_base_env) = match request.config.provider() {
OcrProvider::Mistral => (
Some("MISTRAL_AZURE_API_KEY"),
Some("MISTRAL_AZURE_API_BASE"),
),
OcrProvider::AzureAi => (None, Some("AZURE_AI_API_BASE")),
OcrProvider::AwsTextract
| OcrProvider::Cohere
| OcrProvider::Reducto
| OcrProvider::VertexAi => (None, None),
};
let secret = |name: &str| client.secrets().truthy(name);
let dynamic_api_key = credentials.dynamic_api_key.or_else(|| {
credentials.api_key.clone().or_else(|| {
request
.config
.get_api_key_env_var()
.and_then(credential_env)
.map(|value| Sourced::new(value, InputSource::Environment))
preferred_api_key_env
.into_iter()
.chain(request.config.get_api_key_env_var())
.find_map(secret)
.map(|value| Sourced::new(SecretValue::new(value), InputSource::Environment))
})
});
let dynamic_api_base = credentials.dynamic_api_base.or_else(|| {
credentials.api_base.clone().or_else(|| {
api_base_env
.and_then(credential_env)
.and_then(secret)
.map(|value| Sourced::new(value, InputSource::Environment))
})
});
@ -51,7 +60,12 @@ pub(crate) fn prepare_request(
PreparedOcrRequest {
model,
document,
connection: OcrConnection::new(resolved, transport),
connection: OcrConnection::new(
resolved,
transport,
client.settings().clone(),
client.secrets().clone(),
),
caller_document,
optional_params,
input_sources,
@ -61,7 +75,11 @@ pub(crate) fn prepare_request(
#[cfg(test)]
pub(crate) fn prepare_request_for_test(request: ResolvedOcrRequest) -> PreparedOcrRequest {
prepare_request(request, true)
prepare_request(
request,
true,
&OcrClient::for_test(reqwest::Client::new(), reqwest::Client::new()),
)
}
#[cfg(test)]

86
litellm-rust/crates/core/src/ocr/provider_config.rs
View file

@ -1,5 +1,9 @@
use litellm_core_utils::get_llm_provider_logic::{CustomLlmProvider, get_custom_llm_provider};
use litellm_llms::{
aws_textract::ocr::{
analyze_transformation::TextractAnalyzeDocumentConfig, common_utils::TextractOperation,
transformation::TextractDetectTextConfig,
},
azure_ai::ocr::{
cohere_parse_transformation::AzureAICohereParseConfig,
document_intelligence::transformation::AzureDocumentIntelligenceOcrConfig,
@ -7,13 +11,13 @@ use litellm_llms::{
},
base_llm::ocr::{
error::Error,
handler::{self, CallHooks, OcrClient},
transformation::{
BaseOcrConfig, LiteLLMOcrResponse, OcrCredentialInputs, OcrDocument,
PreparedOcrRequest, ResolvedOcrCredentials,
},
},
cohere::ocr::transformation::CohereParseConfig,
custom_httpx::llm_http_handler::{self, CallHooks, OcrClient},
mistral::ocr::transformation::MistralOcrConfig,
reducto::ocr::transformation::{ReductoParseLegacyConfig, ReductoParseV3Config},
vertex_ai::ocr::{
@ -25,6 +29,14 @@ use strum::{EnumString, IntoStaticStr};
macro_rules! with_config {
($kind:expr, $config:ident => $body:expr) => {
match $kind {
OcrConfigKind::AwsTextract => {
let $config = TextractDetectTextConfig;
$body
}
OcrConfigKind::AwsTextractAnalyze => {
let $config = TextractAnalyzeDocumentConfig;
$body
}
OcrConfigKind::Cohere => {
let $config = CohereParseConfig;
$body
@ -67,6 +79,8 @@ macro_rules! with_config {
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(crate) enum OcrConfigKind {
AwsTextract,
AwsTextractAnalyze,
Cohere,
Mistral,
AzureAi,
@ -81,6 +95,7 @@ pub(crate) enum OcrConfigKind {
impl OcrConfigKind {
pub(crate) const fn provider(self) -> OcrProvider {
match self {
Self::AwsTextract | Self::AwsTextractAnalyze => OcrProvider::AwsTextract,
Self::Cohere => OcrProvider::Cohere,
Self::Mistral => OcrProvider::Mistral,
Self::AzureAi | Self::AzureCohere | Self::AzureDocumentIntelligence => {
@ -116,7 +131,7 @@ impl OcrConfigKind {
request: &PreparedOcrRequest,
hooks: &dyn CallHooks<Error>,
) -> Result<LiteLLMOcrResponse, Error> {
with_config!(self, config => llm_http_handler::ocr(&config, client, request, hooks).await)
with_config!(self, config => handler::ocr(&config, client, request, hooks).await)
}
}
@ -141,6 +156,7 @@ pub fn get_health_check_document(
#[derive(Clone, Copy, Debug, EnumString, IntoStaticStr, PartialEq, Eq)]
#[strum(serialize_all = "snake_case")]
pub(crate) enum OcrProvider {
AwsTextract,
Cohere,
Mistral,
AzureAi,
@ -162,6 +178,10 @@ pub(crate) fn resolve_provider_config(
.parse::<OcrProvider>()
.map_err(|_| Error::InvalidProvider(provider.custom_llm_provider.to_string()))?;
let config = match ocr_provider {
OcrProvider::AwsTextract => match TextractOperation::from_model(provider.model)? {
TextractOperation::DetectDocumentText => OcrConfigKind::AwsTextract,
TextractOperation::AnalyzeDocument => OcrConfigKind::AwsTextractAnalyze,
},
OcrProvider::Cohere => OcrConfigKind::Cohere,
OcrProvider::Mistral => OcrConfigKind::Mistral,
OcrProvider::AzureAi if is_document_intelligence_model(provider.model) => {
@ -277,12 +297,18 @@ mod tests {
#[test]
fn connection_resolution_preserves_dynamic_precedence_and_input_sources() {
let connection = OcrConfigKind::Mistral.resolve_connection_params(OcrCredentialInputs {
api_key: Some(Sourced::new("explicit-key".into(), InputSource::Deployment)),
api_key: Some(Sourced::new(
litellm_auth::SecretValue::new("explicit-key"),
InputSource::Deployment,
)),
api_base: Some(Sourced::new(
"https://explicit.test".into(),
InputSource::Deployment,
)),
dynamic_api_key: Some(Sourced::new("dynamic-key".into(), InputSource::Environment)),
dynamic_api_key: Some(Sourced::new(
litellm_auth::SecretValue::new("dynamic-key"),
InputSource::Environment,
)),
dynamic_api_base: Some(Sourced::new(
"https://dynamic.test".into(),
InputSource::Request,
@ -292,7 +318,7 @@ mod tests {
connection
.api_key
.as_ref()
.map(|value| value.value().as_str()),
.map(|value| value.value().expose()),
Some("dynamic-key")
);
assert_eq!(
@ -318,22 +344,31 @@ mod tests {
fn empty_or_missing_dynamic_credentials_preserve_explicit_values(
#[case] dynamic_value: Option<&str>,
) {
let dynamic =
let dynamic_key = dynamic_value.map(|value| {
Sourced::new(
litellm_auth::SecretValue::new(value),
InputSource::Environment,
)
});
let dynamic_base =
dynamic_value.map(|value| Sourced::new(value.into(), InputSource::Environment));
let connection = OcrConfigKind::Mistral.resolve_connection_params(OcrCredentialInputs {
api_key: Some(Sourced::new("explicit-key".into(), InputSource::Deployment)),
api_key: Some(Sourced::new(
litellm_auth::SecretValue::new("explicit-key"),
InputSource::Deployment,
)),
api_base: Some(Sourced::new(
"https://explicit.test".into(),
InputSource::Deployment,
)),
dynamic_api_key: dynamic.clone(),
dynamic_api_base: dynamic,
dynamic_api_key: dynamic_key,
dynamic_api_base: dynamic_base,
});
assert_eq!(
connection
.api_key
.as_ref()
.map(|value| value.value().as_str()),
.map(|value| value.value().expose()),
Some("explicit-key")
);
assert_eq!(
@ -356,11 +391,18 @@ mod tests {
) {
let connection = OcrConfigKind::AzureDocumentIntelligence.resolve_connection_params(
OcrCredentialInputs {
api_key: explicit_key
.map(|value| Sourced::new(value.into(), InputSource::Deployment)),
api_key: explicit_key.map(|value| {
Sourced::new(
litellm_auth::SecretValue::new(value),
InputSource::Deployment,
)
}),
api_base: explicit_base
.map(|value| Sourced::new(value.into(), InputSource::Deployment)),
dynamic_api_key: Some(Sourced::new("dynamic-key".into(), InputSource::Environment)),
dynamic_api_key: Some(Sourced::new(
litellm_auth::SecretValue::new("dynamic-key"),
InputSource::Environment,
)),
dynamic_api_base: Some(Sourced::new(
"https://dynamic.test".into(),
InputSource::Deployment,
@ -371,7 +413,7 @@ mod tests {
connection
.api_key
.as_ref()
.map(|value| value.value().as_str()),
.map(|value| value.value().expose()),
explicit_key.map(|_| "dynamic-key")
);
assert_eq!(
@ -397,6 +439,22 @@ mod tests {
}
#[rstest]
#[case::misspelled_operation("aws_textract/analyse-document")]
#[case::operation_name_from_the_api("aws_textract/AnalyzeDocument")]
fn textract_models_outside_its_two_operations_are_refused(#[case] model: &str) {
assert!(matches!(
resolve_provider_config(model, None),
Err(Error::InvalidModel {
provider: "aws_textract",
..
})
));
}
#[rstest]
#[case("aws_textract/detect-document-text", OcrConfigKind::AwsTextract)]
#[case("aws_textract/analyze-document", OcrConfigKind::AwsTextractAnalyze)]
#[case("aws_textract/Analyze-Document", OcrConfigKind::AwsTextractAnalyze)]
#[case("reducto/parse-legacy", OcrConfigKind::ReductoLegacy)]
#[case("reducto/future-parse-model", OcrConfigKind::ReductoV3)]
#[case("azure_ai/Cohere-parse-v5", OcrConfigKind::AzureCohere)]

27
litellm-rust/crates/core/src/ocr/route.rs
View file

@ -1,20 +1,17 @@
use std::sync::{Arc, Mutex};
use litellm_auth::ResolvedCredential;
use litellm_callbacks::{
use litellm_host::{
event::{CallEvent, RequestContext, WireRequest},
machine::{HostChannel, HostTokenProvider, MachineFault, RouteMachine, TokenRoute},
route::Route,
};
use litellm_llms::{
base_llm::ocr::{error::Error, transformation::LiteLLMOcrResponse},
custom_httpx::llm_http_handler::OcrClient,
use litellm_llms::base_llm::ocr::{
error::Error, handler::OcrClient, transformation::LiteLLMOcrResponse,
};
use super::handler::perform_ocr_request;
use crate::{
machine::{HostChannel, HostTokenProvider, MachineFault, RouteMachine, TokenRoute},
ocr::types::{LiteLLMOcrRequest, OcrDocumentInput, OcrFileContent, ResolvedOcrRequest},
};
use crate::ocr::types::{LiteLLMOcrRequest, OcrDocumentInput, OcrFileContent, ResolvedOcrRequest};
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum OcrOp {
@ -39,6 +36,8 @@ impl Route for Ocr {
type Error = Error;
type Op = OcrOp;
type OpResult = OcrOpResult;
type Chunk = std::convert::Infallible;
type StreamHead = std::convert::Infallible;
}
impl TokenRoute for Ocr {
@ -54,16 +53,6 @@ impl TokenRoute for Ocr {
}
}
impl From<MachineFault> for Error {
fn from(fault: MachineFault) -> Self {
Self::InvalidRequest(match fault {
MachineFault::Abandoned => "OCR host driver was abandoned".into(),
MachineFault::Protocol(message) => format!("OCR {message}"),
MachineFault::Mismatch => "invalid OCR host operation result".into(),
})
}
}
pub type OcrHost = HostChannel<Ocr>;
pub type OcrMachine = RouteMachine<Ocr>;
@ -173,7 +162,7 @@ impl LocalOcrHost {
}
}
impl litellm_callbacks::host::Host<Ocr> for LocalOcrHost {
impl litellm_host::host::Host<Ocr> for LocalOcrHost {
async fn route(&self, op: OcrOp) -> Result<OcrOpResult, Error> {
match op {
OcrOp::ProjectRequest => self

20
litellm-rust/crates/core/src/ocr/types.rs
View file

@ -1,7 +1,7 @@
use std::{collections::BTreeMap, path::PathBuf, time::Duration};
use bytes::Bytes;
use litellm_auth::{InputSource, TokenProviderHandle};
use litellm_auth::{InputSource, SecretValue, TokenProviderHandle};
use litellm_core_utils::call_arguments::CallArguments;
use litellm_llms::base_llm::ocr::{
error::Error,
@ -56,7 +56,7 @@ pub struct OcrFileContent {
/// credentials, and per-field provenance in `input_sources`.
#[derive(Clone, Debug, Default)]
pub struct OcrConnectionInputs {
pub api_key: Option<String>,
pub api_key: Option<SecretValue>,
pub api_base: Option<String>,
pub extra_headers: Map<String, Value>,
pub timeout: Option<Duration>,
@ -237,6 +237,16 @@ mod tests {
.unwrap()
}
#[test]
fn connection_inputs_debug_hides_the_api_key() {
let inputs = OcrConnectionInputs {
api_key: Some(SecretValue::new("caller-api-key")),
..OcrConnectionInputs::default()
};
assert!(!format!("{inputs:?}").contains("caller-api-key"));
}
#[test]
fn from_inputs_applies_connection_overrides_with_field_sources() {
let request = LiteLLMOcrRequest::from_inputs(
@ -245,7 +255,7 @@ mod tests {
None,
Default::default(),
OcrConnectionInputs {
api_key: Some(" key ".into()),
api_key: Some(SecretValue::new(" key ")),
api_base: Some("".into()),
extra_headers: json!({"x-a": "1"}).as_object().unwrap().clone(),
timeout: Some(Duration::from_secs(7)),
@ -259,7 +269,7 @@ mod tests {
.unwrap();
let api_key = request.credentials.api_key.as_ref().unwrap();
assert_eq!(api_key.clone().into_value(), "key");
assert_eq!(api_key.value().expose(), "key");
assert_eq!(api_key.source(), InputSource::Request);
assert!(request.credentials.api_base.is_none());
assert_eq!(
@ -267,7 +277,7 @@ mod tests {
vec![("x-a".to_string(), "1".to_string())]
);
assert_eq!(request.transport.extra_headers_source, InputSource::Request);
assert_eq!(request.transport.timeout, Duration::from_secs(7));
assert_eq!(request.transport.timeout, Some(Duration::from_secs(7)));
assert_eq!(request.input_sources.len(), 2);
let defaulted = LiteLLMOcrRequest::from_inputs(

4
litellm-rust/crates/core/src/ocr/wire.rs
View file

@ -1,6 +1,6 @@
use std::{collections::BTreeMap, time::Duration};
use litellm_auth::InputSource;
use litellm_auth::{InputSource, SecretValue};
use litellm_llms::base_llm::ocr::{
error::Error,
transformation::{OcrDocument, decode_request_value},
@ -44,7 +44,7 @@ pub fn consumed_optional_param_names(
pub struct OcrWireRequest<D = Value> {
pub model: String,
pub document: D,
pub api_key: Option<String>,
pub api_key: Option<SecretValue>,
pub api_base: Option<String>,
pub custom_llm_provider: Option<String>,
pub extra_headers: Option<Map<String, Value>>,

30
litellm-rust/crates/core/src/outbound.rs Normal file
View file

@ -0,0 +1,30 @@
use std::time::Duration;
use litellm_auth::RequestAuth;
use litellm_auth_aws::SigV4Signer;
use litellm_http::outbound::OutboundRequest;
use serde_json::{Map, Value};
/// Header credentials are already in `headers`; SigV4 is applied here, over the
/// bytes that are sent.
pub(crate) async fn outbound_request<E>(
auth: &RequestAuth,
url: String,
headers: Vec<(String, String)>,
body: &Value,
timeout: Option<Duration>,
optional_params: &Map<String, Value>,
) -> Result<OutboundRequest, E>
where
E: From<litellm_http::Error> + From<litellm_auth_aws::Error>,
{
let RequestAuth::AwsSigV4 { region, service } = auth else {
return Ok(OutboundRequest::json(url, headers, body, timeout)?);
};
let env_lookup = |key: &str| std::env::var(key).ok();
let signer =
SigV4Signer::resolve(region.clone(), service, optional_params, &env_lookup).await?;
Ok(OutboundRequest::signed_json(
url, headers, body, timeout, &signer,
)?)
}

4
litellm-rust/crates/core/src/responses/error.rs
View file

@ -11,7 +11,7 @@ pub enum Error {
#[error(transparent)]
Auth(#[from] litellm_auth::Error),
#[error(transparent)]
Transport(#[from] litellm_llms::custom_httpx::transport::Error),
Transport(#[from] litellm_http::transport::Error),
#[error(transparent)]
Headers(#[from] litellm_llms::custom_httpx::http_handler::HeaderError),
Headers(#[from] litellm_http::request::HeaderError),
}

34
litellm-rust/crates/core/src/responses/websocket.rs
View file

@ -95,9 +95,7 @@ impl ResponsesWebSocketConnection {
timeout: Option<Duration>,
) -> Result<Self, Error> {
let mut request = url.into_client_request().map_err(|error| {
Error::Transport(litellm_llms::custom_httpx::transport::Error::Network(
error.to_string(),
))
Error::Transport(litellm_http::transport::Error::Network(error.to_string()))
})?;
for (name, value) in headers {
let header_name = name
@ -110,7 +108,7 @@ impl ResponsesWebSocketConnection {
let connect = connect_upstream(request);
let result = match timeout {
Some(timeout) => tokio::time::timeout(timeout, connect).await.map_err(|_| {
Error::Transport(litellm_llms::custom_httpx::transport::Error::Network(
Error::Transport(litellm_http::transport::Error::Network(
"Responses WebSocket connection timed out".into(),
))
})?,
@ -118,14 +116,12 @@ impl ResponsesWebSocketConnection {
};
let (socket, _) = result.map_err(|error| match *error {
tokio_tungstenite::tungstenite::Error::Http(response) => {
Error::Transport(litellm_llms::custom_httpx::transport::Error::Http {
Error::Transport(litellm_http::transport::Error::Http {
status: response.status().as_u16(),
body: String::new(),
})
}
other => Error::Transport(litellm_llms::custom_httpx::transport::Error::Network(
other.to_string(),
)),
other => Error::Transport(litellm_http::transport::Error::Network(other.to_string())),
})?;
Ok(Self {
socket: Arc::new(Mutex::new(Some(socket))),
@ -135,16 +131,12 @@ impl ResponsesWebSocketConnection {
pub async fn send_text(&self, text: String) -> Result<(), Error> {
let mut socket = self.socket.lock().await;
let Some(socket) = socket.as_mut() else {
return Err(Error::Transport(
litellm_llms::custom_httpx::transport::Error::Network(
"Responses WebSocket is closed".into(),
),
));
return Err(Error::Transport(litellm_http::transport::Error::Network(
"Responses WebSocket is closed".into(),
)));
};
socket.send(Message::Text(text)).await.map_err(|error| {
Error::Transport(litellm_llms::custom_httpx::transport::Error::Network(
error.to_string(),
))
Error::Transport(litellm_http::transport::Error::Network(error.to_string()))
})
}
@ -160,9 +152,9 @@ impl ResponsesWebSocketConnection {
.map_err(|error| Error::InvalidResponse(error.to_string())),
Some(Ok(Message::Close(_))) | None => Ok(None),
Some(Ok(_)) => Ok(None),
Some(Err(error)) => Err(Error::Transport(
litellm_llms::custom_httpx::transport::Error::Network(error.to_string()),
)),
Some(Err(error)) => Err(Error::Transport(litellm_http::transport::Error::Network(
error.to_string(),
))),
}
}
@ -170,9 +162,7 @@ impl ResponsesWebSocketConnection {
let mut socket = self.socket.lock().await;
if let Some(socket) = socket.as_mut() {
socket.close(None).await.map_err(|error| {
Error::Transport(litellm_llms::custom_httpx::transport::Error::Network(
error.to_string(),
))
Error::Transport(litellm_http::transport::Error::Network(error.to_string()))
})?;
}
*socket = None;

193
litellm-rust/crates/core/tests/aws_textract_ocr.rs Normal file
View file

@ -0,0 +1,193 @@
use std::{collections::BTreeMap, time::SystemTime};
use litellm_auth_aws::{Credentials, aws_signature_headers, sign_post};
use litellm_llms::base_llm::ocr::error::Error;
use serde_json::{Value, json};
use time::{PrimitiveDateTime, format_description};
use crate::ocr::{
route::LocalOcrHost,
test_support::{
MockResponse, header, mock_server, perform_ocr_with, request_body,
wire_request_with_document,
},
types::LiteLLMOcrRequest,
};
const ACCESS_KEY_ID: &str = "AKIDEXAMPLE";
const SECRET_ACCESS_KEY: &str = "wJalrXUtnFEMI/K7MDENG+bPxRfiCYEXAMPLEKEY";
fn textract_request(base: &str) -> LiteLLMOcrRequest {
textract_request_for("aws_textract/detect-document-text", base)
}
fn textract_request_for(model: &str, base: &str) -> LiteLLMOcrRequest {
wire_request_with_document(
model,
&format!("{base}/"),
json!({"type": "image_url", "image_url": "data:image/png;base64,b3JpZ2luYWw="}),
json!({
"aws_access_key_id": ACCESS_KEY_ID,
"aws_secret_access_key": SECRET_ACCESS_KEY,
"aws_region_name": "eu-west-1"
}),
)
}
fn textract_response() -> MockResponse {
MockResponse::json(json!({
"DocumentMetadata": {"Pages": 1},
"Blocks": [{"BlockType": "PAGE"}, {"BlockType": "LINE", "Text": "Invoice 12345"}]
}))
}
/// Recomputes SigV4 over the bytes the server received, at the time the client claimed.
fn expected_authorization(url: &str, raw_request: &str) -> String {
let format =
format_description::parse_borrowed::<2>("[year][month][day]T[hour][minute][second]Z")
.unwrap();
let signed_at: SystemTime =
PrimitiveDateTime::parse(header(raw_request, "x-amz-date").unwrap(), &format)
.unwrap()
.assume_utc()
.into();
let headers: BTreeMap<String, String> = ["content-type", "x-amz-target"]
.into_iter()
.map(|name| {
(
name.to_string(),
header(raw_request, name).unwrap().to_string(),
)
})
.collect();
let body = raw_request.split_once("\r\n\r\n").unwrap().1;
sign_post(
url,
body.as_bytes(),
&aws_signature_headers(&headers),
"eu-west-1",
"textract",
&Credentials::new(ACCESS_KEY_ID, SECRET_ACCESS_KEY, None, None, "test"),
signed_at,
)
.unwrap()["Authorization"]
.clone()
}
#[tokio::test]
async fn the_request_is_signed_for_textract_and_lines_become_the_page() {
let (base, seen, server) = mock_server(vec![textract_response()]).await;
let response = perform_ocr_with(LocalOcrHost::new(textract_request(&base)))
.await
.unwrap();
server.await.unwrap();
let raw = seen.lock().unwrap()[0].clone();
assert_eq!(
header(&raw, "x-amz-target"),
Some("Textract.DetectDocumentText")
);
assert_eq!(
header(&raw, "content-type"),
Some("application/x-amz-json-1.1")
);
assert_eq!(
request_body(&raw),
json!({"Document": {"Bytes": "b3JpZ2luYWw="}})
);
assert_eq!(
header(&raw, "authorization"),
Some(expected_authorization(&format!("{base}/"), &raw).as_str())
);
assert_eq!(response.pages[0].markdown, "Invoice 12345");
assert_eq!(response.usage_info.unwrap().pages_processed, Some(1));
}
#[tokio::test]
async fn a_body_rewritten_by_before_send_is_what_gets_signed_and_sent() {
let (base, seen, server) = mock_server(vec![textract_response()]).await;
let host = LocalOcrHost::new(textract_request(&base)).with_before_send(|mut wire, _| {
assert!(
!wire
.headers
.iter()
.any(|(name, _)| name.eq_ignore_ascii_case("authorization")),
"the hook ran after signing"
);
wire.body["Document"]["Bytes"] = Value::from("cmVkYWN0ZWQ=");
Ok(wire)
});
perform_ocr_with(host).await.unwrap();
server.await.unwrap();
let raw = seen.lock().unwrap()[0].clone();
assert_eq!(
request_body(&raw),
json!({"Document": {"Bytes": "cmVkYWN0ZWQ="}})
);
assert_eq!(
header(&raw, "authorization"),
Some(expected_authorization(&format!("{base}/"), &raw).as_str())
);
}
#[tokio::test]
async fn a_multi_page_rejection_reaches_the_caller_with_the_single_page_limit() {
let (base, _, server) = mock_server(vec![MockResponse {
status: 400,
headers: vec![],
body: json!({
"__type": "UnsupportedDocumentException",
"Message": "Request has unsupported document format"
}),
}])
.await;
let error = perform_ocr_with(LocalOcrHost::new(textract_request(&base)))
.await
.unwrap_err();
server.await.unwrap();
let Error::Provider { status, body, .. } = error else {
panic!("expected a provider error, got {error:?}");
};
assert_eq!(status, 400);
assert!(
body.contains("multi-page documents are not supported"),
"{body}"
);
}
#[tokio::test]
async fn analyze_document_asks_for_layout_and_tables_and_returns_markdown() {
let (base, seen, server) = mock_server(vec![MockResponse::json(json!({
"DocumentMetadata": {"Pages": 1},
"Blocks": [
{"Id": "l1", "BlockType": "LINE", "Text": "Quarterly Report"},
{"Id": "t", "BlockType": "LAYOUT_TITLE",
"Relationships": [{"Type": "CHILD", "Ids": ["l1"]}]}
]
}))])
.await;
let request = textract_request_for("aws_textract/analyze-document", &base);
let response = perform_ocr_with(LocalOcrHost::new(request)).await.unwrap();
server.await.unwrap();
let raw = seen.lock().unwrap()[0].clone();
assert_eq!(
header(&raw, "x-amz-target"),
Some("Textract.AnalyzeDocument")
);
assert_eq!(
request_body(&raw)["FeatureTypes"],
json!(["LAYOUT", "TABLES"])
);
assert_eq!(
header(&raw, "authorization"),
Some(expected_authorization(&format!("{base}/"), &raw).as_str())
);
assert_eq!(response.pages[0].markdown, "# Quarterly Report");
}

70
litellm-rust/crates/core/tests/azure_document_intelligence_ocr.rs
View file

@ -1,10 +1,12 @@
use litellm_callbacks::event::CallEvent;
use litellm_llms::base_llm::ocr::error::Error;
use litellm_host::event::{CallEvent, MachineEvent};
use litellm_llms::base_llm::ocr::{error::Error, settings::OcrSettings};
use rstest::rstest;
use serde_json::{Value, json};
use super::{
test_support::{MockResponse, mock_server, perform_ocr, perform_ocr_with, wire_request},
test_support::{
MockResponse, mock_server, ocr_client, perform_ocr, perform_ocr_with, wire_request,
},
wire::{OcrWireRequest, decode_request},
};
use crate::ocr::route::LocalOcrHost;
@ -69,7 +71,7 @@ async fn rejects_invalid_pages_features_and_format(
let result = decode_request(OcrWireRequest {
model: "azure_ai/doc-intelligence/prebuilt-read".into(),
document: json!({"type":"document_url","document_url":"https://example.com/a.pdf"}),
api_key: Some("key".into()),
api_key: Some(litellm_auth::SecretValue::new("key")),
api_base: Some(base),
custom_llm_provider: None,
extra_headers: None,
@ -200,6 +202,42 @@ async fn immediate_response_normalizes_pages_and_preserves_native() {
);
}
#[tokio::test]
async fn client_settings_choose_the_api_version_and_the_inch_to_pixel_dpi() {
let (base, seen, server) = mock_server(vec![MockResponse::json(json!({
"status":"succeeded",
"analyzeResult":{"pages":[{"pageNumber":1,"width":8.5,"height":11,"unit":"inch"}]}
}))])
.await;
let client = ocr_client().with_settings(OcrSettings {
document_intelligence_api_version: "2099-01-01".into(),
document_intelligence_dpi: 72,
..OcrSettings::default()
});
let result = crate::ocr::client::perform(
&client,
wire_request("azure_ai/doc-intelligence/prebuilt-read", &base, json!({})),
)
.await
.unwrap();
server.await.unwrap();
let target = seen.lock().unwrap()[0]
.split_whitespace()
.nth(1)
.unwrap()
.to_string();
assert_eq!(
query_value(&format!("{base}{target}"), "api-version").as_deref(),
Some("2099-01-01")
);
assert_eq!(
serde_json::to_value(&result.pages[0].dimensions).unwrap(),
json!({"width":612,"height":792,"dpi":72})
);
}
#[tokio::test]
async fn accepted_response_polls_to_success_with_only_credentials() {
let operation = json!({"status":"succeeded","analyzeResult":{"pages":[]}});
@ -263,7 +301,7 @@ async fn accepted_response_emits_response_received_before_polling() {
json!({}),
))
.with_observer(move |event| {
let CallEvent::ResponseReceived { raw } = event else {
let CallEvent::Machine(MachineEvent::ResponseReceived { raw }) = event else {
return;
};
match request_count.lock().unwrap().len() {
@ -425,13 +463,19 @@ async fn polling_deadline_bounds_retry_delay() {
},
])
.await;
let mut request = wire_request("azure_ai/doc-intelligence/prebuilt-read", &base, json!({}));
request.transport.poll_timeout = std::time::Duration::from_millis(100);
let request = wire_request("azure_ai/doc-intelligence/prebuilt-read", &base, json!({}));
let client = ocr_client().with_settings(OcrSettings {
poll_timeout: std::time::Duration::from_millis(100),
..OcrSettings::default()
});
let error = tokio::time::timeout(std::time::Duration::from_secs(1), perform_ocr(request))
.await
.unwrap()
.unwrap_err();
let error = tokio::time::timeout(
std::time::Duration::from_secs(1),
crate::ocr::client::perform(&client, request),
)
.await
.unwrap()
.unwrap_err();
server.await.unwrap();
assert!(error.to_string().contains("timed out"));
}
@ -466,7 +510,7 @@ async fn model_id_is_encoded_and_dot_segments_are_rejected() {
mod transformation {
use std::sync::{Arc, Mutex};
use litellm_callbacks::event::CallEvent;
use litellm_host::event::{CallEvent, MachineEvent};
use litellm_llms::base_llm::ocr::transformation::OcrDocument;
use serde_json::{Value, json};
@ -646,7 +690,7 @@ mod transformation {
json!({}),
))
.with_observer(move |event| {
if let CallEvent::ResponseReceived { raw } = event {
if let CallEvent::Machine(MachineEvent::ResponseReceived { raw }) = event {
observed
.lock()
.unwrap()

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