mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-11 22:51:28 +00:00
Merge branch 'litellm_internal_staging' into litellm_realtime_cost_metrics
Some checks failed
LiteLLM Rust / rustfmt, clippy, test (push) Has been cancelled
Some checks failed
LiteLLM Rust / rustfmt, clippy, test (push) Has been cancelled
This commit is contained in:
commit
f883d6b134
1939 changed files with 134788 additions and 26036 deletions
|
|
@ -133,6 +133,26 @@ commands:
|
|||
done
|
||||
echo "record/replay proxy did not become ready" >&2
|
||||
exit 1
|
||||
start_fake_openai_endpoint:
|
||||
description: "Start the canned OpenAI mock (tests/_fake_openai_endpoint_server.py) on host port 8190 and wait until healthy. Models whose api_base points here (via FAKE_OPENAI_API_BASE) get well-formed chat/text/embedding responses with realistic usage, so the E2E run neither pays for nor depends on the live provider. A request whose model is '429' returns HTTP 429 for rate-limit/cooldown tests. Run after uv deps are synced."
|
||||
steps:
|
||||
- run:
|
||||
name: Start fake OpenAI endpoint
|
||||
background: true
|
||||
command: |
|
||||
uv run --no-sync python tests/_fake_openai_endpoint_server.py --host 0.0.0.0 --port 8190
|
||||
- run:
|
||||
name: Wait for fake OpenAI endpoint
|
||||
command: |
|
||||
for i in $(seq 1 30); do
|
||||
if curl -sf http://localhost:8190/health >/dev/null 2>&1; then
|
||||
echo "fake OpenAI endpoint is up"
|
||||
exit 0
|
||||
fi
|
||||
sleep 1
|
||||
done
|
||||
echo "fake OpenAI endpoint did not become ready" >&2
|
||||
exit 1
|
||||
setup_litellm_enterprise_pip:
|
||||
steps:
|
||||
- run:
|
||||
|
|
@ -168,6 +188,8 @@ jobs:
|
|||
name: win/default
|
||||
shell: powershell.exe
|
||||
working_directory: ~/project
|
||||
environment:
|
||||
UV_PYTHON: "3.11"
|
||||
steps:
|
||||
- checkout
|
||||
- run:
|
||||
|
|
@ -200,7 +222,7 @@ jobs:
|
|||
if (-not (Select-String -Path $PROFILE -SimpleMatch $uvBin -Quiet)) {
|
||||
Add-Content -Path $PROFILE -Value "`$env:Path = `"$uvBin;`$env:Path`""
|
||||
}
|
||||
uv sync --frozen --group dev --python (Get-Command python).Source
|
||||
uv sync --frozen --group dev --python 3.11
|
||||
- run:
|
||||
name: Run Windows-specific test
|
||||
command: |
|
||||
|
|
@ -594,6 +616,8 @@ jobs:
|
|||
working_directory: ~/project
|
||||
resource_class: large
|
||||
parallelism: 4
|
||||
environment:
|
||||
FAKE_OPENAI_API_BASE: http://127.0.0.1:8190
|
||||
steps:
|
||||
- checkout
|
||||
- setup_google_dns
|
||||
|
|
@ -609,6 +633,7 @@ jobs:
|
|||
paths:
|
||||
- ~/.cache/uv
|
||||
key: v1-uv-cache-{{ checksum "uv.lock" }}
|
||||
- start_fake_openai_endpoint
|
||||
# Run pytest and generate JUnit XML report
|
||||
- setup_litellm_enterprise_pip
|
||||
- run:
|
||||
|
|
@ -1549,6 +1574,7 @@ jobs:
|
|||
name: Install Dependencies
|
||||
command: |
|
||||
uv sync --frozen --all-groups --all-extras --python 3.12
|
||||
- start_fake_openai_endpoint
|
||||
- start_postgres:
|
||||
db_name: litellm_test
|
||||
- attach_workspace:
|
||||
|
|
@ -1586,6 +1612,7 @@ jobs:
|
|||
-e DATABASE_URL="postgresql://postgres:postgres@host.docker.internal:5432/litellm_test" \
|
||||
-e DEFAULT_NUM_WORKERS_LITELLM_PROXY=1 \
|
||||
-e DISABLE_SCHEMA_UPDATE="True" \
|
||||
-e FAKE_OPENAI_API_BASE=http://host.docker.internal:8190 \
|
||||
--name my-app \
|
||||
--add-host=host.docker.internal:host-gateway \
|
||||
-v $(pwd)/litellm/proxy/example_config_yaml/bad_schema.prisma:/app/schema.prisma \
|
||||
|
|
@ -1648,6 +1675,7 @@ jobs:
|
|||
zstd -d litellm-docker-database.tar.zst --stdout | docker load
|
||||
docker tag litellm-docker-database:ci my-app:latest
|
||||
- start_openai_record_replay_proxy
|
||||
- start_fake_openai_endpoint
|
||||
- run:
|
||||
name: Run Docker container
|
||||
command: |
|
||||
|
|
@ -1655,6 +1683,7 @@ jobs:
|
|||
-p 4000:4000 \
|
||||
-e DATABASE_URL=postgresql://postgres:postgres@host.docker.internal:5432/circle_test \
|
||||
-e USE_PRISMA_MIGRATE=True \
|
||||
-e FAKE_OPENAI_API_BASE=http://host.docker.internal:8190 \
|
||||
-e AZURE_API_KEY=$AZURE_API_KEY \
|
||||
-e REDIS_HOST=$REDIS_HOST \
|
||||
-e REDIS_PASSWORD=$REDIS_PASSWORD \
|
||||
|
|
@ -1817,6 +1846,7 @@ jobs:
|
|||
zstd -d litellm-docker-database.tar.zst --stdout | docker load
|
||||
docker images | grep litellm-docker-database
|
||||
- start_openai_record_replay_proxy
|
||||
- start_fake_openai_endpoint
|
||||
- run:
|
||||
name: Run Docker container
|
||||
# intentionally give bad redis credentials here
|
||||
|
|
@ -1830,6 +1860,7 @@ jobs:
|
|||
-e REDIS_PORT=$REDIS_PORT \
|
||||
-e LITELLM_MASTER_KEY="sk-1234" \
|
||||
-e OPENAI_API_KEY=$OPENAI_API_KEY \
|
||||
-e FAKE_OPENAI_API_BASE=http://host.docker.internal:8190 \
|
||||
-e LITELLM_LICENSE=$LITELLM_LICENSE \
|
||||
-e OTEL_EXPORTER="in_memory" \
|
||||
-e AWS_ACCESS_KEY_ID=$AWS_ACCESS_KEY_ID \
|
||||
|
|
@ -1889,6 +1920,7 @@ jobs:
|
|||
-e REDIS_PORT=$REDIS_PORT \
|
||||
-e LITELLM_MASTER_KEY="sk-1234" \
|
||||
-e OPENAI_API_KEY=$OPENAI_API_KEY \
|
||||
-e FAKE_OPENAI_API_BASE=http://host.docker.internal:8190 \
|
||||
-e LITELLM_LICENSE="bad-license" \
|
||||
--add-host host.docker.internal:host-gateway \
|
||||
--name my-app-3 \
|
||||
|
|
@ -1938,6 +1970,7 @@ jobs:
|
|||
uv sync --frozen --all-groups --all-extras --python 3.12
|
||||
- start_postgres
|
||||
- start_redis
|
||||
- start_fake_openai_endpoint
|
||||
- attach_workspace:
|
||||
at: ~/project
|
||||
- run:
|
||||
|
|
@ -1961,6 +1994,7 @@ jobs:
|
|||
-e REDIS_PORT=6379 \
|
||||
-e LITELLM_MASTER_KEY="sk-1234" \
|
||||
-e OPENAI_API_KEY=$OPENAI_API_KEY \
|
||||
-e FAKE_OPENAI_API_BASE=http://host.docker.internal:8190 \
|
||||
-e LITELLM_LICENSE=$LITELLM_LICENSE \
|
||||
-e AWS_ACCESS_KEY_ID=$AWS_ACCESS_KEY_ID \
|
||||
-e AWS_SECRET_ACCESS_KEY=$AWS_SECRET_ACCESS_KEY \
|
||||
|
|
@ -2020,6 +2054,7 @@ jobs:
|
|||
command: |
|
||||
uv sync --frozen --all-groups --all-extras --python 3.12
|
||||
- start_postgres
|
||||
- start_fake_openai_endpoint
|
||||
- attach_workspace:
|
||||
at: ~/project
|
||||
- run:
|
||||
|
|
@ -2039,6 +2074,7 @@ jobs:
|
|||
-e REDIS_PASSWORD=$REDIS_PASSWORD \
|
||||
-e REDIS_PORT=$REDIS_PORT \
|
||||
-e LITELLM_MASTER_KEY="sk-1234" \
|
||||
-e FAKE_OPENAI_API_BASE=http://host.docker.internal:8190 \
|
||||
-e LITELLM_LICENSE=$LITELLM_LICENSE \
|
||||
-e USE_DDTRACE=True \
|
||||
-e DD_API_KEY=$DD_API_KEY \
|
||||
|
|
@ -2060,6 +2096,7 @@ jobs:
|
|||
-e REDIS_PASSWORD=$REDIS_PASSWORD \
|
||||
-e REDIS_PORT=$REDIS_PORT \
|
||||
-e LITELLM_MASTER_KEY="sk-1234" \
|
||||
-e FAKE_OPENAI_API_BASE=http://host.docker.internal:8190 \
|
||||
-e LITELLM_LICENSE=$LITELLM_LICENSE \
|
||||
-e USE_DDTRACE=True \
|
||||
-e DD_API_KEY=$DD_API_KEY \
|
||||
|
|
@ -2112,6 +2149,7 @@ jobs:
|
|||
command: |
|
||||
uv sync --frozen --all-groups --all-extras --python 3.12
|
||||
- start_postgres
|
||||
- start_fake_openai_endpoint
|
||||
- attach_workspace:
|
||||
at: ~/project
|
||||
- run:
|
||||
|
|
@ -2129,6 +2167,7 @@ jobs:
|
|||
-e DATABASE_URL=postgresql://postgres:postgres@host.docker.internal:5432/circle_test \
|
||||
-e STORE_MODEL_IN_DB="True" \
|
||||
-e LITELLM_MASTER_KEY="sk-1234" \
|
||||
-e FAKE_OPENAI_API_BASE=http://host.docker.internal:8190 \
|
||||
-e LITELLM_LICENSE=$LITELLM_LICENSE \
|
||||
--add-host host.docker.internal:host-gateway \
|
||||
--name my-app \
|
||||
|
|
@ -2187,6 +2226,7 @@ jobs:
|
|||
command: |
|
||||
docker build -t my-app:latest -f docker/build_from_pip/Dockerfile.build_from_pip .
|
||||
- start_postgres
|
||||
- start_fake_openai_endpoint
|
||||
- run:
|
||||
name: Run Docker container
|
||||
# intentionally give bad redis credentials here
|
||||
|
|
@ -2200,6 +2240,7 @@ jobs:
|
|||
-e REDIS_PORT=$REDIS_PORT \
|
||||
-e LITELLM_MASTER_KEY="sk-1234" \
|
||||
-e OPENAI_API_KEY=$OPENAI_API_KEY \
|
||||
-e FAKE_OPENAI_API_BASE=http://host.docker.internal:8190 \
|
||||
-e LITELLM_LICENSE=$LITELLM_LICENSE \
|
||||
-e OTEL_EXPORTER="in_memory" \
|
||||
-e AWS_ACCESS_KEY_ID=$AWS_ACCESS_KEY_ID \
|
||||
|
|
|
|||
75
.githooks/commit-msg
Executable file
75
.githooks/commit-msg
Executable file
|
|
@ -0,0 +1,75 @@
|
|||
#!/usr/bin/env bash
|
||||
#
|
||||
# commit-msg — enforce Conventional Commits 1.0.0
|
||||
# https://www.conventionalcommits.org/en/v1.0.0/
|
||||
#
|
||||
# Subject format: <type>(<scope>)!: <description>
|
||||
# - <type> must be one of the angular types (feat, fix, ...)
|
||||
# - (<scope>) is optional
|
||||
# - ! is optional and marks a breaking change
|
||||
# - <description> is mandatory and must be non-empty
|
||||
#
|
||||
# Bypass: commit with --no-verify.
|
||||
# Merge, revert, fixup!, squash!, and amend! messages are passed through.
|
||||
|
||||
set -eu
|
||||
|
||||
COMMIT_MSG_FILE="${1:-}"
|
||||
if [ -z "$COMMIT_MSG_FILE" ] || [ ! -f "$COMMIT_MSG_FILE" ]; then
|
||||
echo "commit-msg: missing commit message file" >&2
|
||||
exit 1
|
||||
fi
|
||||
|
||||
# First non-comment, non-empty line is the subject.
|
||||
subject=""
|
||||
while IFS= read -r line || [ -n "$line" ]; do
|
||||
case "$line" in
|
||||
''|'#'*) continue ;;
|
||||
esac
|
||||
subject="$line"
|
||||
break
|
||||
done < "$COMMIT_MSG_FILE"
|
||||
|
||||
if [ -z "$subject" ]; then
|
||||
echo "commit-msg: empty commit message" >&2
|
||||
exit 1
|
||||
fi
|
||||
|
||||
# Pass-through commits generated by git itself.
|
||||
case "$subject" in
|
||||
"Merge "*|"Revert \""*|"fixup! "*|"squash! "*|"amend! "*)
|
||||
exit 0
|
||||
;;
|
||||
esac
|
||||
|
||||
ALLOWED_TYPES="feat|fix|docs|style|refactor|perf|test|build|ci|chore|revert"
|
||||
# Description must not start with an uppercase letter — kept in sync with the
|
||||
# subjectPattern in .github/workflows/conventional-commits.yml so the local
|
||||
# hook is the strictly tighter of the two gates. (Without this guard, a commit
|
||||
# like "feat: Add thing" passes locally but fails the PR-title CI check.)
|
||||
PATTERN="^(${ALLOWED_TYPES})(\([^)]+\))?!?: [^A-Z].*"
|
||||
|
||||
if printf '%s' "$subject" | grep -Eq "$PATTERN"; then
|
||||
exit 0
|
||||
fi
|
||||
|
||||
cat >&2 <<EOF
|
||||
✗ Commit message does not follow Conventional Commits.
|
||||
|
||||
Got: $subject
|
||||
|
||||
Expected: <type>(<scope>)!: <description>
|
||||
(description must start with a lowercase letter)
|
||||
|
||||
Allowed types: feat, fix, docs, style, refactor, perf, test, build, ci, chore, revert
|
||||
Examples:
|
||||
feat(router): add weighted round-robin strategy
|
||||
fix(bedrock): decouple STS region from aws_region_name
|
||||
chore(deps): bump black to 26.3.1
|
||||
refactor!: drop Python 3.8 support
|
||||
|
||||
See https://www.conventionalcommits.org/en/v1.0.0/
|
||||
|
||||
To bypass (use sparingly): git commit --no-verify
|
||||
EOF
|
||||
exit 1
|
||||
92
.githooks/pre-push
Executable file
92
.githooks/pre-push
Executable file
|
|
@ -0,0 +1,92 @@
|
|||
#!/usr/bin/env bash
|
||||
#
|
||||
# pre-push — enforce Conventional Branches
|
||||
# https://conventional-branch.github.io/
|
||||
#
|
||||
# Branch format: <type>/<description>
|
||||
# <type> must be one of: feature, bugfix, hotfix, release, chore
|
||||
#
|
||||
# Protected branches (always allowed):
|
||||
# - main
|
||||
# - litellm_internal_staging
|
||||
# - dependabot/*
|
||||
# - gh-readonly-queue/*
|
||||
#
|
||||
# Tag pushes and branch deletions are skipped.
|
||||
# Bypass: git push --no-verify.
|
||||
|
||||
set -eu
|
||||
|
||||
ZERO_OID="0000000000000000000000000000000000000000"
|
||||
ZERO_OID_SHA256="0000000000000000000000000000000000000000000000000000000000000000"
|
||||
ALLOWED_TYPES="feature|bugfix|hotfix|release|chore"
|
||||
BRANCH_PATTERN="^(${ALLOWED_TYPES})/.+"
|
||||
|
||||
PROTECTED_NAMES="main litellm_internal_staging"
|
||||
PROTECTED_PREFIXES="dependabot/ gh-readonly-queue/"
|
||||
|
||||
is_protected() {
|
||||
branch="$1"
|
||||
for name in $PROTECTED_NAMES; do
|
||||
if [ "$branch" = "$name" ]; then
|
||||
return 0
|
||||
fi
|
||||
done
|
||||
for prefix in $PROTECTED_PREFIXES; do
|
||||
case "$branch" in "$prefix"*) return 0 ;; esac
|
||||
done
|
||||
return 1
|
||||
}
|
||||
|
||||
invalid=""
|
||||
|
||||
while read -r local_ref local_oid remote_ref remote_oid; do
|
||||
# Branch deletion (no local commit being pushed).
|
||||
if [ "$local_oid" = "$ZERO_OID" ] || [ "$local_oid" = "$ZERO_OID_SHA256" ]; then
|
||||
continue
|
||||
fi
|
||||
|
||||
# Only validate branch pushes; ignore tags and other ref namespaces.
|
||||
case "$remote_ref" in
|
||||
refs/heads/*) ;;
|
||||
*) continue ;;
|
||||
esac
|
||||
|
||||
branch="${remote_ref#refs/heads/}"
|
||||
|
||||
if is_protected "$branch"; then
|
||||
continue
|
||||
fi
|
||||
|
||||
if ! printf '%s' "$branch" | grep -Eq "$BRANCH_PATTERN"; then
|
||||
invalid="$invalid $branch"
|
||||
fi
|
||||
done
|
||||
|
||||
if [ -n "$invalid" ]; then
|
||||
cat >&2 <<EOF
|
||||
✗ Branch name does not follow Conventional Branches.
|
||||
|
||||
Invalid:$invalid
|
||||
|
||||
Expected: <type>/<description>
|
||||
|
||||
Allowed types: feature, bugfix, hotfix, release, chore
|
||||
Examples:
|
||||
feature/weighted-round-robin
|
||||
bugfix/streaming-empty-chunks
|
||||
chore/bump-deps
|
||||
hotfix/auth-bypass
|
||||
|
||||
Protected (always allowed): main, litellm_internal_staging,
|
||||
dependabot/*, gh-readonly-queue/*.
|
||||
|
||||
See https://conventional-branch.github.io/
|
||||
|
||||
Rename with: git branch -m <new-name>
|
||||
To bypass (use sparingly): git push --no-verify
|
||||
EOF
|
||||
exit 1
|
||||
fi
|
||||
|
||||
exit 0
|
||||
BIN
.github/deploy-on-aws.png
vendored
Normal file
BIN
.github/deploy-on-aws.png
vendored
Normal file
Binary file not shown.
|
After Width: | Height: | Size: 4 KiB |
BIN
.github/deploy-on-gcp.png
vendored
Normal file
BIN
.github/deploy-on-gcp.png
vendored
Normal file
Binary file not shown.
|
After Width: | Height: | Size: 4.8 KiB |
2
.github/pull_request_template.md
vendored
2
.github/pull_request_template.md
vendored
|
|
@ -4,7 +4,7 @@
|
|||
|
||||
## Linear ticket
|
||||
|
||||
<!-- if you are an internal contributor, add the Linear ticket e.g. "Resolves LIT-1234" to magically link the Linear ticket to the GitHub PR -->
|
||||
<!-- if you are an internal contributor (e.g., your username is postfixed with -berri or -berriai), add "Resolves " followed by the Linear ticket e.g. "Resolves LIT-1234" to magically link the Linear ticket to the GitHub PR -->
|
||||
|
||||
## Pre-Submission checklist
|
||||
|
||||
|
|
|
|||
50
.github/scripts/_agent_shin_actions.py
vendored
Normal file
50
.github/scripts/_agent_shin_actions.py
vendored
Normal file
|
|
@ -0,0 +1,50 @@
|
|||
"""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
Normal file
211
.github/scripts/agent_shin_shared.py
vendored
Normal file
|
|
@ -0,0 +1,211 @@
|
|||
"""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()
|
||||
573
.github/scripts/close_low_quality_prs.py
vendored
Normal file
573
.github/scripts/close_low_quality_prs.py
vendored
Normal file
|
|
@ -0,0 +1,573 @@
|
|||
#!/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
Normal file
282
.github/scripts/triage-requirements.txt
vendored
Normal file
|
|
@ -0,0 +1,282 @@
|
|||
# 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
|
||||
557
.github/scripts/triage_rollout_heads_up.py
vendored
Normal file
557
.github/scripts/triage_rollout_heads_up.py
vendored
Normal file
|
|
@ -0,0 +1,557 @@
|
|||
#!/usr/bin/env python3
|
||||
"""One-shot 7-day heads-up sweep for the Agent Shin rollout.
|
||||
|
||||
Posts a friendly "the OSS triage bot kicks in next Monday" comment on every
|
||||
open external PR/issue that currently *would* fail the new rubric — i.e.,
|
||||
every PR/issue Agent Shin would close once the rollout completes. The point
|
||||
is to give contributors a full week to fix their description before the bot
|
||||
ever takes a destructive action, so nobody is surprised by an auto-close.
|
||||
|
||||
The script is designed to run **exactly once** at rollout, fired by a manual
|
||||
``workflow_dispatch`` (``dry_run=false``) on the heads-up workflow. Re-runs
|
||||
are safe: every comment is stamped with the hidden ``HEADS_UP_MARKER`` and
|
||||
PRs/issues that already carry the marker are skipped.
|
||||
|
||||
Dry-run vs. real run
|
||||
--------------------
|
||||
Defaults to dry-run. Passing ``--close`` flips into real mode. Every GitHub
|
||||
mutation goes through ``_agent_shin_actions``, which has a one-line
|
||||
``if dry_run: log else: do_it`` per call, so the only difference between a
|
||||
dry-run preview and the real run is the call site that actually hits the
|
||||
GitHub API.
|
||||
|
||||
Local preview::
|
||||
|
||||
python3 .github/scripts/triage_rollout_heads_up.py --repo BerriAI/litellm
|
||||
|
||||
Real run (the manual rollout dispatch uses this)::
|
||||
|
||||
python3 .github/scripts/triage_rollout_heads_up.py --repo BerriAI/litellm --close
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import datetime as dt
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
# Make the sibling triage_with_llm + _agent_shin_actions importable when this
|
||||
# script is invoked directly (the GitHub workflow does `python3 .github/scripts/...`).
|
||||
_SCRIPTS_DIR = Path(__file__).resolve().parent
|
||||
if str(_SCRIPTS_DIR) not in sys.path:
|
||||
sys.path.insert(0, str(_SCRIPTS_DIR))
|
||||
|
||||
from _agent_shin_actions import maybe_post_comment # noqa: E402
|
||||
from agent_shin_shared import ( # noqa: E402
|
||||
AGENT_SHIN_DEFAULT_BOT_LOGIN,
|
||||
ALLOWLIST_LOGINS,
|
||||
list_open_items,
|
||||
)
|
||||
from triage_with_llm import ( # noqa: E402
|
||||
DEFAULT_MODEL,
|
||||
call_llm_judge,
|
||||
fetch_issue,
|
||||
fetch_pr,
|
||||
gh,
|
||||
is_internal_contributor,
|
||||
review_gate,
|
||||
triage,
|
||||
)
|
||||
|
||||
# Hidden marker so re-runs skip PRs/issues we've already notified. Distinct from
|
||||
# the within-grace / ready / regressed markers so it can't be confused with the
|
||||
# steady-state lifecycle comments.
|
||||
HEADS_UP_MARKER = "<!-- agent-shin:rollout-heads-up -->"
|
||||
|
||||
# Placeholder until the litellm-docs PR ships. The rollout blog post explains
|
||||
# the new rubric, the 7-day grace, and how to recover after an auto-close.
|
||||
# TODO(docs): replace with the canonical URL once the litellm-docs PR merges.
|
||||
ROLLOUT_BLOG_URL = "https://docs.litellm.ai/docs/agent_shin_triage_rollout"
|
||||
|
||||
# Default cutoff is one week from "now". Computed at runtime so the wording
|
||||
# stays correct even if the rollout is merged later than planned. The user can
|
||||
# override with --close-on YYYY-MM-DD when running the script manually.
|
||||
DEFAULT_GRACE_DAYS = 7
|
||||
|
||||
# The daily auto-close sweeps (close_low_quality_prs.yml at 09:00 UTC and
|
||||
# review_gate.yml at 09:30 UTC) are what actually close a still-failing item,
|
||||
# so the deadline we promise contributors has to name that wall-clock moment.
|
||||
ACTIVATION_TIME_UTC = "09:00 UTC"
|
||||
|
||||
|
||||
def _format_cutoff(cutoff: dt.date) -> str:
|
||||
"""Human-readable, timezone-explicit cutoff, e.g. ``Monday, June 1, 2026
|
||||
(09:00 UTC)`` — the moment a still-failing PR/issue gets closed."""
|
||||
return (
|
||||
f"{cutoff.strftime('%A, %B')} {cutoff.day}, {cutoff.year} "
|
||||
f"({ACTIVATION_TIME_UTC})"
|
||||
)
|
||||
|
||||
|
||||
def _rubric_section_pr() -> str:
|
||||
return (
|
||||
"**Going forward, every external PR needs ONE of:**\n"
|
||||
"\n"
|
||||
"- A linked GitHub issue using a closing keyword: "
|
||||
"`Fixes #1234`, `Closes #1234`, or `Resolves #1234`, OR\n"
|
||||
"- All three of: a clear **problem description**, **expected vs. "
|
||||
"actual behavior**, and **end-to-end QA proof** (at least one of a "
|
||||
"short screen recording / video, before/after screenshots, or the "
|
||||
"exact commands you ran with their real output; mocked or stubbed "
|
||||
"runs don't count).\n"
|
||||
"\n"
|
||||
"PRs also need a **Greptile confidence score of 4/5 or higher** before "
|
||||
"the bot will tag them `ready for review`. You can `@greptileai` to "
|
||||
"request a fresh review at any time, including after the PR is closed."
|
||||
)
|
||||
|
||||
|
||||
def _rubric_section_issue() -> str:
|
||||
return (
|
||||
"**Going forward, every external issue needs:**\n"
|
||||
"\n"
|
||||
"- For **bug reports**: end-to-end evidence of the bug (at least one "
|
||||
"of a screen recording / video, a screenshot, or the exact commands "
|
||||
"you ran with their real output / traceback) plus expected vs. actual "
|
||||
"behavior. Written steps with no run output don't count, and mocked "
|
||||
"or stubbed runs don't count.\n"
|
||||
"- For **feature requests**: a clear description of the proposed "
|
||||
"feature plus a use case + concrete example (config, API call, UI "
|
||||
"flow, or scenario showing what's blocked today)."
|
||||
)
|
||||
|
||||
|
||||
def _description_only_note(kind: str) -> str:
|
||||
noun = "PR" if kind == "pr" else "issue"
|
||||
return (
|
||||
f"⚠️ **The requirements must live in the {noun} *description*, not in "
|
||||
"comments.** Some PRs/issues collect 100+ comments from humans and "
|
||||
"bots; reading the entire thread on every triage run would balloon "
|
||||
"GitHub API usage (we'd start getting 429'd) and blow out the LLM "
|
||||
"judge's context. The bot only reads the description, so anything "
|
||||
"you add as a comment will be invisible to it."
|
||||
)
|
||||
|
||||
|
||||
def _missing_section(verdict: dict, greptile_score: int | None) -> str:
|
||||
"""Bullet list of what's currently missing on this PR/issue.
|
||||
|
||||
Combines the LLM judge's `missing` list (rubric items) with a Greptile
|
||||
shortfall (for PRs) so the contributor sees one list of things to fix.
|
||||
"""
|
||||
missing = list(verdict.get("missing") or [])
|
||||
if greptile_score is not None and greptile_score < 4:
|
||||
missing.insert(
|
||||
0,
|
||||
f"Greptile's most recent review scored this PR {greptile_score}/5 "
|
||||
"(below the 4/5 bar Agent Shin will require).",
|
||||
)
|
||||
if not missing:
|
||||
return (
|
||||
"_The bot couldn't articulate a specific missing piece; see the "
|
||||
"rubric link above and double-check the description includes all "
|
||||
"of it before the rollout._"
|
||||
)
|
||||
bullets = "\n".join(f"- {m}" for m in missing)
|
||||
return f"**What this one is currently missing:**\n\n{bullets}"
|
||||
|
||||
|
||||
def _recovery_section(kind: str) -> str:
|
||||
if kind == "pr":
|
||||
return (
|
||||
"**If the bot closes this PR after the rollout:** update the "
|
||||
"description with the missing pieces, then either open a fresh "
|
||||
"PR or comment `@agent-shin reconsider` on the closed PR. If "
|
||||
"Greptile re-scores you at 4/5 or higher I'll reopen and tag "
|
||||
"the PR `ready for review`. (`@greptileai` works on closed PRs "
|
||||
"too; a fresh review is one of the signals that lifts you back "
|
||||
"into the queue.) This is **not** us losing interest in your "
|
||||
"change; far from it. We just need open PRs to be a list of "
|
||||
"things a maintainer can act on, so we can get to yours faster."
|
||||
)
|
||||
return (
|
||||
"**If the bot closes this issue after the rollout:** edit the issue "
|
||||
"description to add the missing pieces, then comment `@agent-shin "
|
||||
"reconsider` on the closed issue. I'll re-evaluate and, if the rubric "
|
||||
"is met, reopen it. (GitHub doesn't let external authors reopen an "
|
||||
"issue a maintainer or bot closed, so the comment is the reliable "
|
||||
"path.) This is **not** us saying the bug isn't real or the request "
|
||||
"isn't useful; it's so the remaining open issues are a list of things "
|
||||
"a maintainer can act on."
|
||||
)
|
||||
|
||||
|
||||
def format_heads_up_comment(
|
||||
*, kind: str, verdict: dict, greptile_score: int | None, cutoff: dt.date
|
||||
) -> str:
|
||||
"""Compose the friendly 7-day heads-up comment posted on a failing PR/issue."""
|
||||
noun = "PR" if kind == "pr" else "issue"
|
||||
rubric = _rubric_section_pr() if kind == "pr" else _rubric_section_issue()
|
||||
cutoff_str = _format_cutoff(cutoff)
|
||||
explanation = (verdict.get("explanation") or "").strip()
|
||||
explanation_block = (
|
||||
f"> _(The judge's note for this one: {explanation})_\n\n" if explanation else ""
|
||||
)
|
||||
|
||||
return (
|
||||
"🚅 **Heads-up: we're turning on the OSS triage bot in "
|
||||
f"{DEFAULT_GRACE_DAYS} days, on {cutoff_str}.**\n"
|
||||
"\n"
|
||||
"We're rolling out **Agent Shin**, an LLM-as-judge triage bot for "
|
||||
f"external {noun}s. Once it's live, the bot reads each open "
|
||||
f"{noun}'s description, scores it against a small rubric, and "
|
||||
f"auto-closes any {noun} that's missing the basics, with a single "
|
||||
f"comment explaining what's missing and how to recover. Full "
|
||||
f"context: [Agent Shin rollout blog post]({ROLLOUT_BLOG_URL}).\n"
|
||||
"\n"
|
||||
f"{rubric}\n"
|
||||
"\n"
|
||||
f"{_description_only_note(kind)}\n"
|
||||
"\n"
|
||||
f"{_missing_section(verdict, greptile_score)}\n"
|
||||
"\n"
|
||||
f"{explanation_block}"
|
||||
"**Timeline (you have a week):**\n"
|
||||
"\n"
|
||||
f"- We turn the bot on in {DEFAULT_GRACE_DAYS} days, on "
|
||||
f"**{cutoff_str}**. You have until then to update this {noun}'s "
|
||||
"description with the missing pieces above.\n"
|
||||
f"- If this {noun} still fails the rubric at **{cutoff_str}**, "
|
||||
"we'll close it.\n"
|
||||
f"- From then on the bot runs daily, and every {noun} that fails "
|
||||
"the rubric gets a **2-hour lifetime**: one warning comment, then "
|
||||
"auto-close 2 hours later.\n"
|
||||
"\n"
|
||||
f"{_recovery_section(kind)}\n"
|
||||
"\n"
|
||||
f"{HEADS_UP_MARKER}"
|
||||
)
|
||||
|
||||
|
||||
def _list_open_numbers(repo: str, kind: str) -> list[int]:
|
||||
"""Return every open PR or issue number in ``repo``.
|
||||
|
||||
Delegates to ``list_open_items`` so the full backlog is fetched (no cap)
|
||||
and the `gh {pr,issue} list` invocation stays in one shared place. ``gh
|
||||
issue list`` would include PRs, but ``list_open_items`` uses the dedicated
|
||||
command per kind, so the two never mix.
|
||||
"""
|
||||
return [
|
||||
item["number"] for item in list_open_items(kind, repo=repo, fields="number")
|
||||
]
|
||||
|
||||
|
||||
def _has_heads_up_marker(item: dict) -> bool:
|
||||
"""Cheap fast-path: check the PR/issue body itself for the marker.
|
||||
|
||||
The marker is appended to the *comment* we post, not the body, so this
|
||||
will only fire if the body literally contains the marker text. We still
|
||||
do the comment-marker check separately below; this body check just lets
|
||||
us short-circuit for PRs/issues that quote the marker for any reason.
|
||||
"""
|
||||
body = item.get("body") or ""
|
||||
return HEADS_UP_MARKER in body
|
||||
|
||||
|
||||
def _comments_have_marker(repo: str, number: int) -> bool:
|
||||
"""True if the bot already posted a comment carrying the marker.
|
||||
|
||||
Used for idempotency: a re-run skips items the previous run notified.
|
||||
Filters by author (matching the sibling marker-checks in
|
||||
``triage_with_llm._has_marker`` and
|
||||
``agent_shin_shared.seconds_since_latest_marker_comment``) so a
|
||||
contributor who quotes the heads-up via GitHub's "Quote reply" — which
|
||||
preserves HTML comments in the raw markdown — can't trick the
|
||||
idempotency check into silently skipping a real heads-up.
|
||||
|
||||
Comments live on the unified issues endpoint regardless of whether the
|
||||
item is a PR or an issue, so no ``kind`` argument is required here.
|
||||
"""
|
||||
expected_login = (
|
||||
os.environ.get("AGENT_SHIN_BOT_LOGIN") or AGENT_SHIN_DEFAULT_BOT_LOGIN
|
||||
).lower()
|
||||
raw = gh(
|
||||
"api",
|
||||
"--paginate",
|
||||
f"repos/{repo}/issues/{number}/comments?per_page=100",
|
||||
)
|
||||
for line in raw.splitlines():
|
||||
line = line.strip()
|
||||
if not line:
|
||||
continue
|
||||
try:
|
||||
payload = json.loads(line)
|
||||
except json.JSONDecodeError:
|
||||
continue
|
||||
comments = payload if isinstance(payload, list) else [payload]
|
||||
for comment in comments:
|
||||
author = ((comment.get("user") or {}).get("login") or "").lower()
|
||||
if author != expected_login:
|
||||
continue
|
||||
if HEADS_UP_MARKER in (comment.get("body") or ""):
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def _evaluate_pr(*, repo: str, number: int, model: str, judge: Any = None) -> dict:
|
||||
"""Run the future PR rubric (review_gate) in dry-run and return the result."""
|
||||
return review_gate(
|
||||
repo=repo,
|
||||
number=number,
|
||||
close=False, # we only want the verdict, never act here
|
||||
model=model,
|
||||
judge=judge,
|
||||
)
|
||||
|
||||
|
||||
def _evaluate_issue(*, repo: str, number: int, model: str, judge: Any = None) -> dict:
|
||||
"""Run the future issue rubric (triage kind='issue') in dry-run."""
|
||||
return triage(
|
||||
repo=repo,
|
||||
kind="issue",
|
||||
number=number,
|
||||
close=False,
|
||||
model=model,
|
||||
judge=judge,
|
||||
)
|
||||
|
||||
|
||||
def _would_be_closed(kind: str, result: dict) -> bool:
|
||||
"""True if the future triage would auto-close this PR/issue based on the
|
||||
rubric (regardless of grace-period gating).
|
||||
|
||||
For PRs we trust ``review_gate``'s ``passing`` field — it combines the LLM
|
||||
verdict and the Greptile score. For issues we read the LLM verdict
|
||||
directly. Both fields are ``None``/missing on skip paths
|
||||
(skip-internal-author, skip-llm-error, etc.) where the future bot would
|
||||
NOT close the item — those return False.
|
||||
"""
|
||||
if kind == "pr":
|
||||
passing = result.get("passing")
|
||||
if passing is None:
|
||||
return False # skipped — nothing for the heads-up to warn about
|
||||
return passing is False
|
||||
verdict = result.get("verdict") or {}
|
||||
return (verdict.get("verdict") or "").lower() == "fail"
|
||||
|
||||
|
||||
def _process_one(
|
||||
*,
|
||||
repo: str,
|
||||
kind: str,
|
||||
number: int,
|
||||
model: str,
|
||||
cutoff: dt.date,
|
||||
dry_run: bool,
|
||||
judge: Any = None,
|
||||
skip_marker_check: bool = False,
|
||||
allowlist: frozenset[str] = ALLOWLIST_LOGINS,
|
||||
) -> dict:
|
||||
"""Evaluate one PR/issue and post a heads-up if it would be auto-closed.
|
||||
|
||||
Returns a per-item dict for the summary table.
|
||||
"""
|
||||
base = {"kind": kind, "number": number}
|
||||
fetcher = fetch_pr if kind == "pr" else fetch_issue
|
||||
item = fetcher(repo, number)
|
||||
|
||||
if (item.get("state") or "") != "open":
|
||||
return {**base, "action": "skip-not-open"}
|
||||
if allowlist:
|
||||
login = (item.get("user") or {}).get("login") or ""
|
||||
if login.lower() not in allowlist:
|
||||
return {**base, "action": "skip-not-allowlisted"}
|
||||
elif is_internal_contributor(item):
|
||||
return {**base, "action": "skip-internal-author"}
|
||||
if not skip_marker_check and _has_heads_up_marker(item):
|
||||
return {**base, "action": "skip-already-marked-in-body"}
|
||||
if not skip_marker_check and _comments_have_marker(repo, number):
|
||||
return {**base, "action": "skip-already-notified"}
|
||||
|
||||
if kind == "pr":
|
||||
result = _evaluate_pr(repo=repo, number=number, model=model, judge=judge)
|
||||
else:
|
||||
result = _evaluate_issue(repo=repo, number=number, model=model, judge=judge)
|
||||
|
||||
if not _would_be_closed(kind, result):
|
||||
return {**base, "action": "skip-passing", "evaluator": result.get("action")}
|
||||
|
||||
verdict = result.get("verdict") or {}
|
||||
greptile_score = result.get("greptile_score") if kind == "pr" else None
|
||||
comment = format_heads_up_comment(
|
||||
kind=kind, verdict=verdict, greptile_score=greptile_score, cutoff=cutoff
|
||||
)
|
||||
maybe_post_comment(repo, number, comment, dry_run=dry_run)
|
||||
return {
|
||||
**base,
|
||||
"action": "heads-up-posted" if not dry_run else "would-post-heads-up",
|
||||
"verdict": (verdict.get("verdict") or "").lower(),
|
||||
"greptile_score": greptile_score,
|
||||
}
|
||||
|
||||
|
||||
def _print_summary(results: list[dict]) -> None:
|
||||
"""Tally per-action counts so a dry-run preview tells you at a glance how
|
||||
many comments the real run would post."""
|
||||
counts: dict[str, int] = {}
|
||||
for r in results:
|
||||
counts[r["action"]] = counts.get(r["action"], 0) + 1
|
||||
print("\n=== rollout heads-up summary ===")
|
||||
for action in sorted(counts):
|
||||
print(f" {action:35s} {counts[action]}")
|
||||
print(f" total {len(results)}")
|
||||
|
||||
|
||||
def run(
|
||||
*,
|
||||
repo: str,
|
||||
close: bool,
|
||||
cutoff: dt.date,
|
||||
model: str,
|
||||
kinds: tuple[str, ...] = ("pr", "issue"),
|
||||
judge: Any = None,
|
||||
only_numbers: dict[str, list[int]] | None = None,
|
||||
skip_marker_check: bool = False,
|
||||
) -> list[dict]:
|
||||
"""Sweep ``repo`` and post heads-up comments. Returns the per-item results."""
|
||||
dry_run = not close
|
||||
if dry_run:
|
||||
print(
|
||||
f"[DRY RUN] sweeping {repo}; --close not passed, no comments will be posted."
|
||||
)
|
||||
else:
|
||||
print(f"[REAL RUN] sweeping {repo}; comments WILL be posted.")
|
||||
print(f"Cutoff date in comment body: {cutoff.isoformat()}")
|
||||
|
||||
results: list[dict] = []
|
||||
for kind in kinds:
|
||||
if only_numbers and kind in only_numbers:
|
||||
numbers = list(only_numbers[kind])
|
||||
else:
|
||||
numbers = _list_open_numbers(repo, kind)
|
||||
print(f"\n--- {kind}s: {len(numbers)} open ---")
|
||||
for n in numbers:
|
||||
try:
|
||||
result = _process_one(
|
||||
repo=repo,
|
||||
kind=kind,
|
||||
number=n,
|
||||
model=model,
|
||||
cutoff=cutoff,
|
||||
dry_run=dry_run,
|
||||
judge=judge,
|
||||
skip_marker_check=skip_marker_check,
|
||||
)
|
||||
except (
|
||||
Exception
|
||||
) as exc: # noqa: BLE001 - per-item errors don't abort the sweep
|
||||
result = {
|
||||
"kind": kind,
|
||||
"number": n,
|
||||
"action": "error",
|
||||
"error": str(exc),
|
||||
}
|
||||
print(f"!! {kind}#{n}: {exc}", file=sys.stderr)
|
||||
print(f" {kind}#{n}: {result['action']}")
|
||||
results.append(result)
|
||||
_print_summary(results)
|
||||
return results
|
||||
|
||||
|
||||
def main() -> int:
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("--repo", required=True, help="owner/repo")
|
||||
parser.add_argument(
|
||||
"--close",
|
||||
action="store_true",
|
||||
help=(
|
||||
"Actually post comments. Without this flag the script is in "
|
||||
"dry-run mode and only logs what it would do."
|
||||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--close-on",
|
||||
type=dt.date.fromisoformat,
|
||||
default=None,
|
||||
help=(
|
||||
"Cutoff date shown in the heads-up comment as the rollout date "
|
||||
f"(default: today + {DEFAULT_GRACE_DAYS} days)."
|
||||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--model",
|
||||
default=os.environ.get("TRIAGE_MODEL") or DEFAULT_MODEL,
|
||||
help=f"Model for the rubric LLM judge (default: {DEFAULT_MODEL}).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--kind",
|
||||
choices=("pr", "issue", "both"),
|
||||
default="both",
|
||||
help="Restrict the sweep to PRs or issues only (default: both).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--only-pr",
|
||||
type=int,
|
||||
action="append",
|
||||
default=[],
|
||||
help="Limit the PR sweep to these PR numbers (repeat for several).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--only-issue",
|
||||
type=int,
|
||||
action="append",
|
||||
default=[],
|
||||
help="Limit the issue sweep to these issue numbers (repeat for several).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--ignore-existing-marker",
|
||||
action="store_true",
|
||||
help=(
|
||||
"Re-post on PRs/issues that already carry the heads-up marker. "
|
||||
"Useful for testing the comment wording on a known PR."
|
||||
),
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
cutoff = args.close_on or (
|
||||
dt.datetime.now(dt.timezone.utc).date() + dt.timedelta(days=DEFAULT_GRACE_DAYS)
|
||||
)
|
||||
|
||||
kinds: tuple[str, ...]
|
||||
if args.kind == "pr":
|
||||
kinds = ("pr",)
|
||||
elif args.kind == "issue":
|
||||
kinds = ("issue",)
|
||||
else:
|
||||
kinds = ("pr", "issue")
|
||||
|
||||
only: dict[str, list[int]] = {}
|
||||
if args.only_pr:
|
||||
only["pr"] = args.only_pr
|
||||
if args.only_issue:
|
||||
only["issue"] = args.only_issue
|
||||
|
||||
# The script must NOT hit the LLM in dry-run if no key is set — we still
|
||||
# want a useful preview that says "skip-no-llm-key" for items that would
|
||||
# have been judged. Production runs require OPENAI_API_KEY.
|
||||
if args.close and not os.environ.get("OPENAI_API_KEY"):
|
||||
parser.error("OPENAI_API_KEY must be set for --close (real-run) mode.")
|
||||
|
||||
run(
|
||||
repo=args.repo,
|
||||
close=args.close,
|
||||
cutoff=cutoff,
|
||||
model=args.model,
|
||||
kinds=kinds,
|
||||
only_numbers=only or None,
|
||||
skip_marker_check=args.ignore_existing_marker,
|
||||
)
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
1778
.github/scripts/triage_with_llm.py
vendored
Normal file
1778
.github/scripts/triage_with_llm.py
vendored
Normal file
File diff suppressed because it is too large
Load diff
2
.github/workflows/check-ui-api-types.yml
vendored
2
.github/workflows/check-ui-api-types.yml
vendored
|
|
@ -54,7 +54,7 @@ jobs:
|
|||
run: uv run --no-sync prisma generate --schema litellm/proxy/schema.prisma
|
||||
|
||||
- name: Set up Node.js
|
||||
uses: actions/setup-node@a0853c24544627f65ddf259abe73b1d18a591444 # v5.0
|
||||
uses: actions/setup-node@a0853c24544627f65ddf259abe73b1d18a591444 # v5.0.0
|
||||
with:
|
||||
node-version: "20"
|
||||
cache: "npm"
|
||||
|
|
|
|||
92
.github/workflows/close_low_quality_prs.yml
vendored
Normal file
92
.github/workflows/close_low_quality_prs.yml
vendored
Normal file
|
|
@ -0,0 +1,92 @@
|
|||
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[@]}"
|
||||
6
.github/workflows/codeql.yml
vendored
6
.github/workflows/codeql.yml
vendored
|
|
@ -43,14 +43,14 @@ jobs:
|
|||
persist-credentials: false
|
||||
|
||||
- name: Initialize CodeQL
|
||||
uses: github/codeql-action/init@ebcb5b36ded6beda4ceefea6a8bc4cc885255bb3 # v3
|
||||
uses: github/codeql-action/init@ebcb5b36ded6beda4ceefea6a8bc4cc885255bb3 # v3.34.1
|
||||
with:
|
||||
languages: ${{ matrix.language }}
|
||||
build-mode: ${{ matrix.build-mode }}
|
||||
config-file: ./.github/codeql/codeql-config.yml
|
||||
|
||||
- name: Perform CodeQL Analysis
|
||||
uses: github/codeql-action/analyze@ebcb5b36ded6beda4ceefea6a8bc4cc885255bb3 # v3
|
||||
uses: github/codeql-action/analyze@ebcb5b36ded6beda4ceefea6a8bc4cc885255bb3 # v3.34.1
|
||||
with:
|
||||
category: "/language:${{ matrix.language }}"
|
||||
output: sarif-results
|
||||
|
|
@ -77,7 +77,7 @@ jobs:
|
|||
output: sarif-results/python.sarif
|
||||
|
||||
- name: Upload SARIF
|
||||
uses: github/codeql-action/upload-sarif@ebcb5b36ded6beda4ceefea6a8bc4cc885255bb3 # v3
|
||||
uses: github/codeql-action/upload-sarif@ebcb5b36ded6beda4ceefea6a8bc4cc885255bb3 # v3.34.1
|
||||
with:
|
||||
sarif_file: sarif-results
|
||||
category: "/language:${{ matrix.language }}"
|
||||
|
|
|
|||
46
.github/workflows/conventional-commits.yml
vendored
Normal file
46
.github/workflows/conventional-commits.yml
vendored
Normal file
|
|
@ -0,0 +1,46 @@
|
|||
name: Conventional PR Title
|
||||
|
||||
# Squash-merge replaces the merge commit subject with the PR title, so
|
||||
# enforcing Conventional Commits at the PR-title level is what actually gates
|
||||
# the commits that land on the default branch. The local commit-msg hook
|
||||
# (.githooks/commit-msg) is a best-effort assist; this workflow is the gate.
|
||||
#
|
||||
# See https://www.conventionalcommits.org/en/v1.0.0/
|
||||
|
||||
on:
|
||||
pull_request:
|
||||
types: [opened, edited, reopened, synchronize, labeled, unlabeled]
|
||||
|
||||
permissions:
|
||||
pull-requests: read
|
||||
|
||||
jobs:
|
||||
lint-pr-title:
|
||||
name: Validate PR title
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- name: Check title against Conventional Commits
|
||||
uses: amannn/action-semantic-pull-request@48f256284bd46cdaab1048c3721360e808335d50 # v6.1.1
|
||||
env:
|
||||
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
|
||||
with:
|
||||
# Must mirror the type list in .githooks/commit-msg.
|
||||
types: |
|
||||
feat
|
||||
fix
|
||||
docs
|
||||
style
|
||||
refactor
|
||||
perf
|
||||
test
|
||||
build
|
||||
ci
|
||||
chore
|
||||
revert
|
||||
requireScope: false
|
||||
subjectPattern: ^(?![A-Z]).+$
|
||||
subjectPatternError: |
|
||||
The subject "{subject}" must start with a lowercase character.
|
||||
# Allow merges/reverts that GitHub generates automatically.
|
||||
ignoreLabels: |
|
||||
ignore-semantic-pull-request
|
||||
33
.github/workflows/create-release.yml
vendored
33
.github/workflows/create-release.yml
vendored
|
|
@ -52,6 +52,22 @@ jobs:
|
|||
// are stable maintenance releases, not pre-releases.
|
||||
const isPrerelease = /(?:rc|nightly|alpha|beta|[-.]dev)/i.test(tag);
|
||||
|
||||
// A stable release should only claim the repo "latest" badge when its
|
||||
// version is >= the current latest. Otherwise a backport (e.g. 1.84.6)
|
||||
// would steal "latest" from a newer line (e.g. 1.88.1).
|
||||
const versionKey = (rawTag) => {
|
||||
const m = String(rawTag).match(/^v?(\d+)\.(\d+)\.(\d+)/);
|
||||
if (!m) return null;
|
||||
const maintenance = String(rawTag).match(/(?:\.post|\.patch\.)(\d+)/i);
|
||||
return [Number(m[1]), Number(m[2]), Number(m[3]), maintenance ? Number(maintenance[1]) : 0];
|
||||
};
|
||||
const isAtLeast = (a, b) => {
|
||||
for (let i = 0; i < a.length; i++) {
|
||||
if (a[i] !== b[i]) return a[i] > b[i];
|
||||
}
|
||||
return true;
|
||||
};
|
||||
|
||||
const cosignSection = [
|
||||
`## Verify Docker Image Signature`,
|
||||
``,
|
||||
|
|
@ -90,6 +106,22 @@ jobs:
|
|||
].join('\n');
|
||||
|
||||
try {
|
||||
let makeLatest = "false";
|
||||
const newVersion = versionKey(tag);
|
||||
if (!isPrerelease && newVersion) {
|
||||
let latestVersion = null;
|
||||
try {
|
||||
const latest = await github.rest.repos.getLatestRelease({
|
||||
owner: context.repo.owner,
|
||||
repo: context.repo.repo,
|
||||
});
|
||||
latestVersion = versionKey(latest.data.tag_name);
|
||||
} catch (error) {
|
||||
if (error.status !== 404) throw error;
|
||||
}
|
||||
makeLatest = (!latestVersion || isAtLeast(newVersion, latestVersion)) ? "true" : "false";
|
||||
}
|
||||
|
||||
const response = await github.rest.repos.createRelease({
|
||||
draft: true,
|
||||
generate_release_notes: true,
|
||||
|
|
@ -108,6 +140,7 @@ jobs:
|
|||
release_id: response.data.id,
|
||||
body: updatedBody,
|
||||
draft: false,
|
||||
make_latest: makeLatest,
|
||||
});
|
||||
|
||||
} catch (error) {
|
||||
|
|
|
|||
44
.github/workflows/osv-scan.yml
vendored
Normal file
44
.github/workflows/osv-scan.yml
vendored
Normal file
|
|
@ -0,0 +1,44 @@
|
|||
name: OSV Scan
|
||||
|
||||
on:
|
||||
pull_request:
|
||||
branches:
|
||||
- main
|
||||
- litellm_internal_staging
|
||||
- litellm_oss_branch
|
||||
- "litellm_**"
|
||||
schedule:
|
||||
- cron: "23 6 * * *"
|
||||
workflow_dispatch:
|
||||
|
||||
permissions: {}
|
||||
|
||||
concurrency:
|
||||
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
|
||||
cancel-in-progress: true
|
||||
|
||||
jobs:
|
||||
osv-scan:
|
||||
name: osv-scan
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 10
|
||||
permissions:
|
||||
contents: read
|
||||
steps:
|
||||
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
|
||||
with:
|
||||
persist-credentials: false
|
||||
|
||||
- name: Download osv-scanner v2.3.8
|
||||
run: |
|
||||
curl -fsSL --retry 3 -o "$RUNNER_TEMP/osv-scanner" \
|
||||
https://github.com/google/osv-scanner/releases/download/v2.3.8/osv-scanner_linux_amd64
|
||||
echo "bc98e15319ed0d515e3f9235287ba53cdc5535d576d24fd573978ecfe9ab92dc $RUNNER_TEMP/osv-scanner" | sha256sum -c -
|
||||
chmod +x "$RUNNER_TEMP/osv-scanner"
|
||||
|
||||
- name: Scan lockfiles
|
||||
run: |
|
||||
"$RUNNER_TEMP/osv-scanner" scan source \
|
||||
--config osv-scanner.toml \
|
||||
-L uv.lock \
|
||||
-L ui/litellm-dashboard/package-lock.json
|
||||
53
.github/workflows/test-linting.yml
vendored
53
.github/workflows/test-linting.yml
vendored
|
|
@ -14,11 +14,15 @@ permissions:
|
|||
jobs:
|
||||
lint:
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 5
|
||||
timeout-minutes: 15
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
|
||||
# Check out the PR head, not the default refs/pull/N/merge: the merge ref
|
||||
# folds in newer base commits, which the diff-based gates (ruff delta,
|
||||
# Any-discipline) would otherwise blame on this branch.
|
||||
with:
|
||||
ref: ${{ github.event.pull_request.head.sha }}
|
||||
fetch-depth: 0
|
||||
clean: true
|
||||
persist-credentials: false
|
||||
|
|
@ -67,15 +71,27 @@ jobs:
|
|||
uv run --no-sync ruff check .
|
||||
cd ..
|
||||
|
||||
- name: Check strict-rule budget (delta vs base)
|
||||
env:
|
||||
BASE_SHA: ${{ github.event.pull_request.base.sha }}
|
||||
run: |
|
||||
uv run --no-sync python scripts/ruff_strict_gate.py --base "$BASE_SHA"
|
||||
|
||||
- name: Check type-discipline budget (mutable collections / casts / type guards / kwargs / unexplained suppressions, delta vs base)
|
||||
env:
|
||||
BASE_SHA: ${{ github.event.pull_request.base.sha }}
|
||||
run: |
|
||||
uv run --no-sync python scripts/type_discipline_gate.py --base "$BASE_SHA"
|
||||
|
||||
- name: Print OpenAI version
|
||||
run: |
|
||||
uv run --no-sync python -c "import openai; print(f'OpenAI version: {openai.__version__}')"
|
||||
|
||||
- name: Run MyPy type checking
|
||||
- name: Check basedpyright budget (delta vs base)
|
||||
env:
|
||||
BASE_SHA: ${{ github.event.pull_request.base.sha }}
|
||||
run: |
|
||||
cd litellm
|
||||
uv run --no-sync mypy .
|
||||
cd ..
|
||||
(uv run --no-sync basedpyright --outputjson || true) | uv run --no-sync python scripts/type_check_gate.py --base "$BASE_SHA"
|
||||
|
||||
- name: Check for circular imports
|
||||
run: |
|
||||
|
|
@ -87,6 +103,33 @@ jobs:
|
|||
run: |
|
||||
uv run --no-sync python -c "from litellm import *" || (echo '🚨 import failed, this means you introduced unprotected imports! 🚨'; exit 1)
|
||||
|
||||
# Intentionally NON-GATING. This job turns red when a *-budget.json ceiling is
|
||||
# raised (or a rule/budget is dropped) so a loosening is obvious in review, but it
|
||||
# must be kept OUT of the branch-protection required-checks list so a justified
|
||||
# bump can still be merged by a human who has seen and accepted the red.
|
||||
budget-ratchet:
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 5
|
||||
permissions:
|
||||
contents: read
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
|
||||
with:
|
||||
fetch-depth: 0
|
||||
persist-credentials: false
|
||||
|
||||
- name: Set up Python
|
||||
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
|
||||
with:
|
||||
python-version: "3.12"
|
||||
|
||||
- name: Ratchet check (budgets may only decrease; non-gating)
|
||||
env:
|
||||
BASE_SHA: ${{ github.event.pull_request.base.sha }}
|
||||
run: |
|
||||
python scripts/budget_ratchet_check.py --base "$BASE_SHA"
|
||||
|
||||
secret-scan:
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 5
|
||||
|
|
|
|||
6
.github/workflows/test-litellm-ui-build.yml
vendored
6
.github/workflows/test-litellm-ui-build.yml
vendored
|
|
@ -25,7 +25,7 @@ jobs:
|
|||
persist-credentials: false
|
||||
|
||||
- name: Setup Node.js
|
||||
uses: actions/setup-node@a0853c24544627f65ddf259abe73b1d18a591444 # v5.0
|
||||
uses: actions/setup-node@a0853c24544627f65ddf259abe73b1d18a591444 # v5.0.0
|
||||
with:
|
||||
node-version: "20"
|
||||
cache: "npm"
|
||||
|
|
@ -77,7 +77,7 @@ jobs:
|
|||
|
||||
- name: Setup Node.js
|
||||
if: steps.changed.outputs.has_files == 'true'
|
||||
uses: actions/setup-node@a0853c24544627f65ddf259abe73b1d18a591444 # v5.0
|
||||
uses: actions/setup-node@a0853c24544627f65ddf259abe73b1d18a591444 # v5.0.0
|
||||
with:
|
||||
node-version: "20"
|
||||
cache: "npm"
|
||||
|
|
@ -111,4 +111,4 @@ jobs:
|
|||
if: ${{ !cancelled() && steps.changed.outputs.has_files == 'true' }}
|
||||
run: |
|
||||
npx eslint . -f json -o "$RUNNER_TEMP/lint-report.json" || true
|
||||
node scripts/check-lint-budgets.mjs "$RUNNER_TEMP/lint-report.json" eslint-budgets.json
|
||||
node scripts/check-lint-budgets.mjs "$RUNNER_TEMP/lint-report.json" eslint-budgets.json --check eslint-metrics.json
|
||||
|
|
|
|||
65
.github/workflows/test-rust.yml
vendored
Normal file
65
.github/workflows/test-rust.yml
vendored
Normal file
|
|
@ -0,0 +1,65 @@
|
|||
name: LiteLLM Rust
|
||||
|
||||
on:
|
||||
push:
|
||||
paths:
|
||||
- "litellm-rust/**"
|
||||
- ".github/workflows/test-rust.yml"
|
||||
pull_request:
|
||||
branches:
|
||||
- main
|
||||
- litellm_internal_staging
|
||||
- litellm_oss_branch
|
||||
- "litellm_**"
|
||||
paths:
|
||||
- "litellm-rust/**"
|
||||
- ".github/workflows/test-rust.yml"
|
||||
|
||||
permissions:
|
||||
contents: read
|
||||
|
||||
concurrency:
|
||||
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
|
||||
cancel-in-progress: true
|
||||
|
||||
jobs:
|
||||
rust-checks:
|
||||
name: rustfmt, clippy, test
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 10
|
||||
defaults:
|
||||
run:
|
||||
working-directory: litellm-rust
|
||||
env:
|
||||
CARGO_TERM_COLOR: always
|
||||
|
||||
steps:
|
||||
- name: Checkout repository
|
||||
uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
|
||||
with:
|
||||
persist-credentials: false
|
||||
|
||||
- name: Set up Rust
|
||||
run: |
|
||||
rustup toolchain install stable --profile minimal --component clippy,rustfmt
|
||||
rustup default stable
|
||||
|
||||
- name: Cache Cargo registry and target
|
||||
uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0
|
||||
with:
|
||||
path: |
|
||||
~/.cargo/registry
|
||||
~/.cargo/git
|
||||
litellm-rust/target
|
||||
key: ${{ runner.os }}-cargo-${{ hashFiles('litellm-rust/Cargo.lock') }}
|
||||
restore-keys: |
|
||||
${{ runner.os }}-cargo-
|
||||
|
||||
- name: Check Rust formatting
|
||||
run: cargo fmt --check
|
||||
|
||||
- name: Run Clippy
|
||||
run: cargo clippy --workspace --all-targets --locked -- -D warnings
|
||||
|
||||
- name: Run Rust tests
|
||||
run: cargo test --workspace --locked
|
||||
2
.github/workflows/test-unit-misc.yml
vendored
2
.github/workflows/test-unit-misc.yml
vendored
|
|
@ -32,7 +32,9 @@ jobs:
|
|||
tests/test_litellm/repositories
|
||||
tests/test_litellm/images
|
||||
tests/test_litellm/interactions
|
||||
tests/test_litellm/ocr
|
||||
tests/test_litellm/passthrough
|
||||
tests/test_litellm/sandbox
|
||||
tests/test_litellm/vector_stores
|
||||
tests/test_litellm/test_*.py
|
||||
workers: 2
|
||||
|
|
|
|||
10
.github/workflows/test-unit-proxy-endpoints.yml
vendored
10
.github/workflows/test-unit-proxy-endpoints.yml
vendored
|
|
@ -11,8 +11,6 @@ on:
|
|||
|
||||
permissions:
|
||||
contents: read
|
||||
id-token: write
|
||||
pull-requests: write
|
||||
|
||||
concurrency:
|
||||
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
|
||||
|
|
@ -20,6 +18,10 @@ concurrency:
|
|||
|
||||
jobs:
|
||||
proxy-endpoints:
|
||||
permissions:
|
||||
contents: read
|
||||
id-token: write
|
||||
pull-requests: write
|
||||
uses: ./.github/workflows/_test-unit-base.yml
|
||||
with:
|
||||
test-path: >-
|
||||
|
|
@ -52,6 +54,10 @@ jobs:
|
|||
# is independent and its coverage artifact is uploaded separately.
|
||||
# See: https://www.notion.so/36c43b8acdab81ee845fd5365128a2fc
|
||||
proxy-server:
|
||||
permissions:
|
||||
contents: read
|
||||
id-token: write
|
||||
pull-requests: write
|
||||
uses: ./.github/workflows/_test-unit-base.yml
|
||||
with:
|
||||
test-path: tests/test_litellm/proxy/proxy_server
|
||||
|
|
|
|||
7
.github/workflows/test_server_root_path.yml
vendored
7
.github/workflows/test_server_root_path.yml
vendored
|
|
@ -32,17 +32,16 @@ jobs:
|
|||
df -h /
|
||||
|
||||
- name: Set up Docker Buildx
|
||||
uses: docker/setup-buildx-action@8d2750c68a42422c14e847fe6c8ac0403b4cbd6f # v3.12
|
||||
uses: docker/setup-buildx-action@8d2750c68a42422c14e847fe6c8ac0403b4cbd6f # v3.12.0
|
||||
|
||||
- name: Build Docker image
|
||||
uses: docker/build-push-action@0adf9959216b96bec444f325f1e493d4aa344497 #v6.14
|
||||
uses: docker/build-push-action@0adf9959216b96bec444f325f1e493d4aa344497 # v6.14.0
|
||||
with:
|
||||
context: .
|
||||
file: ./docker/Dockerfile.non_root
|
||||
tags: litellm-test:${{ github.sha }}
|
||||
load: true
|
||||
cache-from: type=gha
|
||||
cache-to: type=gha,mode=max
|
||||
push: false
|
||||
|
||||
- name: Start LiteLLM container with SERVER_ROOT_PATH
|
||||
run: |
|
||||
|
|
|
|||
96
.github/workflows/triage_issue_with_llm.yml
vendored
Normal file
96
.github/workflows/triage_issue_with_llm.yml
vendored
Normal file
|
|
@ -0,0 +1,96 @@
|
|||
name: Agent Shin — Issue triage
|
||||
|
||||
# LLM-as-judge triage for external GitHub issues.
|
||||
#
|
||||
# DRY-RUN BY DEFAULT. See .github/workflows/triage_pr_with_llm.yml for the
|
||||
# enablement procedure — same repo variable (`AGENT_SHIN_ENABLED=true`)
|
||||
# unlocks the PR and issue triage flows together.
|
||||
|
||||
on:
|
||||
issues:
|
||||
types: [opened, reopened]
|
||||
workflow_dispatch:
|
||||
inputs:
|
||||
issue_number:
|
||||
description: "Issue number to triage manually."
|
||||
required: true
|
||||
close:
|
||||
description: "If true and AGENT_SHIN_ENABLED=true, actually close on fail."
|
||||
required: false
|
||||
default: "false"
|
||||
type: choice
|
||||
options:
|
||||
- "true"
|
||||
- "false"
|
||||
|
||||
permissions:
|
||||
contents: read
|
||||
issues: write
|
||||
|
||||
jobs:
|
||||
triage:
|
||||
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: Install LLM client
|
||||
run: pip install --no-cache-dir --require-hashes -r .github/scripts/triage-requirements.txt
|
||||
|
||||
- name: Run Agent Shin
|
||||
env:
|
||||
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
|
||||
# Only expose the LLM key when the bot is enabled or a collaborator
|
||||
# triggers it manually, so an external user can't force paid LLM
|
||||
# calls by churning issues 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.
|
||||
OPENAI_API_KEY: ${{ (vars.AGENT_SHIN_ENABLED == 'true' || github.event_name == 'workflow_dispatch') && secrets.OPENAI_API_KEY || '' }}
|
||||
OPENAI_BASE_URL: ${{ vars.OPENAI_BASE_URL }}
|
||||
TRIAGE_MODEL: ${{ vars.TRIAGE_MODEL }}
|
||||
AGENT_SHIN_ENABLED: ${{ vars.AGENT_SHIN_ENABLED }}
|
||||
DISPATCH_CLOSE: ${{ github.event.inputs.close }}
|
||||
ISSUE_NUMBER: ${{ github.event.issue.number || github.event.inputs.issue_number }}
|
||||
run: |
|
||||
set -euo pipefail
|
||||
ARGS=(--repo "${{ github.repository }}" --issue "${ISSUE_NUMBER}")
|
||||
# Fail-safe gating: only the EXACT string "true" enables the
|
||||
# destructive --close path. The workflow_dispatch input is a
|
||||
# `choice` dropdown of "true"/"false" so the UI is constrained,
|
||||
# but the API (`gh workflow run -f close=...`) accepts any
|
||||
# string, and a `!= "false"` check would treat "True", "yes",
|
||||
# "1", "TRUE", typos, and accidental whitespace as enabling
|
||||
# closure. Mirror the Greptile closer's `= "true"` pattern.
|
||||
if [ "${AGENT_SHIN_ENABLED:-false}" = "true" ] && [ "${DISPATCH_CLOSE:-false}" = "true" ]; then
|
||||
ARGS+=(--close)
|
||||
echo "::notice::Agent Shin is ENABLED and running in close-on-fail mode."
|
||||
elif [ "${AGENT_SHIN_ENABLED:-false}" = "true" ]; then
|
||||
echo "::notice::Agent Shin is ENABLED but this trigger is dry-run (workflow_dispatch close != 'true')."
|
||||
else
|
||||
echo "::notice::Agent Shin is in DRY-RUN mode (AGENT_SHIN_ENABLED is not 'true'). No comments will be posted; no issues will be closed."
|
||||
fi
|
||||
# Automatic `issues` events stay dry-run regardless until the team
|
||||
# explicitly invokes workflow_dispatch with close=true.
|
||||
if [ "${GITHUB_EVENT_NAME:-}" = "issues" ]; then
|
||||
# filter out --close rather than substituting to "" (which would
|
||||
# leave an empty positional arg that argparse rejects)
|
||||
FILTERED=()
|
||||
for arg in "${ARGS[@]}"; do
|
||||
if [ "${arg}" != "--close" ]; then
|
||||
FILTERED+=("${arg}")
|
||||
fi
|
||||
done
|
||||
ARGS=("${FILTERED[@]}")
|
||||
echo "::notice::issues trigger -> forcing dry-run."
|
||||
fi
|
||||
python3 .github/scripts/triage_with_llm.py "${ARGS[@]}"
|
||||
172
.github/workflows/triage_reconsider.yml
vendored
Normal file
172
.github/workflows/triage_reconsider.yml
vendored
Normal file
|
|
@ -0,0 +1,172 @@
|
|||
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)"
|
||||
92
.github/workflows/triage_rollout_heads_up.yml
vendored
Normal file
92
.github/workflows/triage_rollout_heads_up.yml
vendored
Normal file
|
|
@ -0,0 +1,92 @@
|
|||
name: Agent Shin — rollout heads-up (one-shot)
|
||||
|
||||
# Fires the 7-day heads-up comment on every open external PR/issue that the
|
||||
# new triage bot would auto-close. The real sweep is a deliberate one-shot:
|
||||
# trigger it at rollout via a manual `workflow_dispatch` with `dry_run=false`.
|
||||
# The script is idempotent (skips items that already carry the
|
||||
# `<!-- agent-shin:rollout-heads-up -->` marker), so a re-run is harmless.
|
||||
#
|
||||
# The automatic push trigger runs DRY-RUN only, so merging the script to
|
||||
# `litellm_internal_staging` never posts a comment; it just confirms the
|
||||
# workflow is wired up. Posting real comments requires the manual dispatch,
|
||||
# which is also the only trigger that exposes `OPENAI_API_KEY`. The heads-up
|
||||
# is intentionally NOT gated on `AGENT_SHIN_ENABLED`: it has to warn
|
||||
# contributors while that flag is still off, ahead of the flip that turns on
|
||||
# auto-closing.
|
||||
#
|
||||
# The workflow is a thin shell over `.github/scripts/triage_rollout_heads_up.py`.
|
||||
# Dry-run vs. real run differ in EXACTLY one CLI flag (`--close`), added only
|
||||
# on a manual dispatch with `dry_run=false`.
|
||||
|
||||
on:
|
||||
push:
|
||||
branches:
|
||||
- litellm_internal_staging
|
||||
paths:
|
||||
# The presence of this script on staging IS the rollout merge marker.
|
||||
# Editing the file later would re-fire the workflow; that's safe because
|
||||
# the script skips PRs/issues that already have the heads-up marker.
|
||||
- ".github/scripts/triage_rollout_heads_up.py"
|
||||
workflow_dispatch:
|
||||
inputs:
|
||||
dry_run:
|
||||
description: "Dry run (true = preview only, false = actually post comments)."
|
||||
required: false
|
||||
default: "true"
|
||||
type: choice
|
||||
options:
|
||||
- "true"
|
||||
- "false"
|
||||
|
||||
permissions:
|
||||
contents: read
|
||||
issues: write
|
||||
pull-requests: write
|
||||
|
||||
jobs:
|
||||
heads-up:
|
||||
if: github.repository == 'BerriAI/litellm'
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- name: Checkout triage scripts
|
||||
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: Install LLM client
|
||||
run: pip install --no-cache-dir --require-hashes -r .github/scripts/triage-requirements.txt
|
||||
|
||||
- name: Run heads-up sweep
|
||||
env:
|
||||
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
|
||||
# Only the manual dispatch (the real-run trigger) needs the LLM key.
|
||||
# The automatic push trigger runs dry-run and never posts, so it gets
|
||||
# no key. Mirrors the sibling triage workflows, which expose the key
|
||||
# only on an enabled/dispatched run rather than unconditionally.
|
||||
OPENAI_API_KEY: ${{ github.event_name == 'workflow_dispatch' && secrets.OPENAI_API_KEY || '' }}
|
||||
OPENAI_BASE_URL: ${{ vars.OPENAI_BASE_URL }}
|
||||
TRIAGE_MODEL: ${{ vars.TRIAGE_MODEL }}
|
||||
# The real run is a deliberate manual dispatch with dry_run=false.
|
||||
# Use the EXACT "false" comparison so any unexpected input value
|
||||
# fail-closes to dry-run (mirrors the AGENT_SHIN_ENABLED pattern in
|
||||
# the sibling workflows). The automatic push trigger always stays
|
||||
# dry-run, so merging the script never posts.
|
||||
DRY_RUN_INPUT: ${{ github.event.inputs.dry_run }}
|
||||
run: |
|
||||
set -euo pipefail
|
||||
ARGS=(--repo "${{ github.repository }}")
|
||||
if [ "${GITHUB_EVENT_NAME:-}" = "workflow_dispatch" ] && [ "${DRY_RUN_INPUT:-true}" = "false" ]; then
|
||||
ARGS+=(--close)
|
||||
echo "::notice::Manual rollout dispatch with dry_run=false -> heads-up comments WILL be posted."
|
||||
elif [ "${GITHUB_EVENT_NAME:-}" = "workflow_dispatch" ]; then
|
||||
echo "::notice::Manual dispatch in dry-run mode -> previewing only, no comments will be posted."
|
||||
else
|
||||
echo "::notice::Automatic push trigger -> dry-run preview only. Fire the real rollout sweep with a manual workflow_dispatch (dry_run=false)."
|
||||
fi
|
||||
python3 .github/scripts/triage_rollout_heads_up.py "${ARGS[@]}"
|
||||
13
.github/workflows/zizmor.yml
vendored
13
.github/workflows/zizmor.yml
vendored
|
|
@ -2,9 +2,9 @@ name: GitHub Actions Security Analysis
|
|||
|
||||
on:
|
||||
push:
|
||||
branches: [main]
|
||||
branches: [main, litellm_internal_staging]
|
||||
pull_request:
|
||||
branches: [main]
|
||||
branches: [main, litellm_internal_staging]
|
||||
|
||||
concurrency:
|
||||
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
|
||||
|
|
@ -18,9 +18,7 @@ jobs:
|
|||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 5
|
||||
permissions:
|
||||
security-events: write
|
||||
contents: read
|
||||
actions: read
|
||||
steps:
|
||||
- name: Checkout repository
|
||||
uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
|
||||
|
|
@ -28,4 +26,9 @@ jobs:
|
|||
persist-credentials: false
|
||||
|
||||
- name: Run zizmor
|
||||
uses: zizmorcore/zizmor-action@71321a20a9ded102f6e9ce5718a2fcec2c4f70d8 # v0.5.2
|
||||
uses: zizmorcore/zizmor-action@5f14fd08f7cf1cb1609c1e344975f152c7ee938d # v0.5.6
|
||||
with:
|
||||
version: "1.24.1"
|
||||
min-severity: medium
|
||||
advanced-security: false
|
||||
annotations: true
|
||||
|
|
|
|||
1
.gitignore
vendored
1
.gitignore
vendored
|
|
@ -74,7 +74,6 @@ tests/local_testing/log.txt
|
|||
.codegpt
|
||||
litellm/proxy/_new_new_secret_config.yaml
|
||||
litellm/proxy/custom_guardrail.py
|
||||
**/.mypy_cache/
|
||||
litellm/proxy/application.log
|
||||
tests/llm_translation/vertex_test_account.json
|
||||
tests/llm_translation/test_vertex_key.json
|
||||
|
|
|
|||
23
CLAUDE.md
23
CLAUDE.md
|
|
@ -29,13 +29,19 @@ If you ever make public-facing PR descriptions, comments, issues, commit message
|
|||
- don't use "—". Instead, reach for ";", ".", etc.
|
||||
- don't use the pattern "It's not X, it's Y", "You're not X, you're Y", etc.
|
||||
- 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)
|
||||
- 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: unless there's a sentence immediately after, don't add a "."
|
||||
- don't use →. Instead, prefer not to use arrows, and if need be, use -> instead
|
||||
|
||||
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
|
||||
|
||||
Run tests, format your code, and lint your code before each commit
|
||||
|
||||
When you fix violations gated by `ruff-strict-budget.json` or `basedpyright-code-budget.json`, run `make lint-budget-update` and commit the lowered baselines so the ceilings ratchet down instead of leaving stale headroom
|
||||
|
||||
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()` / `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
|
||||
|
||||
Ask to commit and push your work when you're done (or if you're confident that your code is good and works, just do it)
|
||||
|
||||
When you must use real LLM models to, for example, write e2e tests, write a QA runbook, etc., make sure to use the latest models (doesn't have to be smartest, can also be a modern small, fast one. No strong preference for smart vs fast here, just use something modern) as of the year and month of the current date. Do a web search as necessary to figure that out
|
||||
|
|
@ -52,6 +58,21 @@ Do not put names of customers or customer company names in code, PRs, and issues
|
|||
|
||||
CI supply-chain safety: Never pipe a remote script into a shell (`curl ... | bash`, `wget ... | sh`); download the artifact to a file, verify its SHA-256 checksum, then install. Pin every external tool to a specific version with a full URL (not `latest` or `stable`). Verify checksums for all downloaded binaries, using the provider's official `.sha256` / `.sha256sum` sidecar when available. These rules apply to every download in CI
|
||||
|
||||
Follow these coding conventions for new/updated code (a three-line fix in a legacy file shouldn't trigger huge drive-by refactors):
|
||||
|
||||
- Composition over inheritance
|
||||
- Never-nester: early returns over deep nesting
|
||||
- Don't throw; model failures as values (One function (e.g., raise_public) maps error union to existing public exception contracts via exhaustive match + assert_never)
|
||||
- No mutation; don't reassign variables, global or local. Instead of mutable lists and dicts, prefer tuples, frozen dataclasses (with slots=True), etc.
|
||||
- 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
|
||||
|
||||
Follow conventional commits for commit names and PR titles
|
||||
|
||||
## Think Before Coding
|
||||
|
||||
**Don't assume. Don't hide confusion. Surface tradeoffs**
|
||||
|
|
|
|||
|
|
@ -38,18 +38,25 @@ Before contributing code to LiteLLM, you must sign our [Contributor License Agre
|
|||
git clone https://github.com/YOUR_USERNAME/litellm.git
|
||||
cd litellm
|
||||
|
||||
# Create a new branch for your feature
|
||||
git checkout -b your-feature-branch
|
||||
# Create a new branch for your feature (see "Commit and Branch Conventions" below)
|
||||
git checkout -b feature/your-feature
|
||||
|
||||
# Install development dependencies
|
||||
make install-dev
|
||||
|
||||
# Install git hooks that enforce commit + branch conventions (one-time, opt-in)
|
||||
make install-hooks
|
||||
|
||||
# Verify your setup works
|
||||
make help
|
||||
```
|
||||
|
||||
That's it! Your local development environment is ready.
|
||||
|
||||
## Commit and Branch Conventions
|
||||
|
||||
Commits follow [Conventional Commits](https://www.conventionalcommits.org/en/v1.0.0/) and branches follow [Conventional Branches](https://conventional-branch.github.io/). Run `make install-hooks` once per clone to enable the local git hooks that enforce these — see the [contributor docs](https://docs.litellm.ai/docs/extras/contributing_code#commit-and-branch-conventions) for the full type list, examples, the protected-branch bypass list, and how to opt out.
|
||||
|
||||
### 2. Development Workflow
|
||||
|
||||
Here's the recommended workflow for making changes:
|
||||
|
|
@ -67,12 +74,12 @@ make lint
|
|||
# Run unit tests to ensure nothing is broken
|
||||
make test-unit
|
||||
|
||||
# Commit your changes
|
||||
# Commit your changes (must follow Conventional Commits — see above)
|
||||
git add .
|
||||
git commit -m "Your descriptive commit message"
|
||||
git commit -m "feat(scope): your descriptive commit message"
|
||||
|
||||
# Push and create a PR
|
||||
git push origin your-feature-branch
|
||||
# Push and create a PR (branch must follow Conventional Branches — see above)
|
||||
git push origin feature/your-feature
|
||||
```
|
||||
|
||||
## Adding Testing
|
||||
|
|
@ -147,7 +154,7 @@ Individual linting commands:
|
|||
```bash
|
||||
make format-check # Check Black formatting
|
||||
make lint-ruff # Run Ruff linting
|
||||
make lint-mypy # Run MyPy type checking
|
||||
make lint-basedpyright # Run basedpyright type checking
|
||||
make check-circular-imports # Check for circular imports
|
||||
make check-import-safety # Check import safety
|
||||
```
|
||||
|
|
@ -209,7 +216,7 @@ LiteLLM follows the [Google Python Style Guide](https://google.github.io/stylegu
|
|||
Our automated quality checks include:
|
||||
- **Black** for consistent code formatting
|
||||
- **Ruff** for linting and code quality
|
||||
- **MyPy** for static type checking
|
||||
- **basedpyright** for static type checking
|
||||
- **Circular import detection**
|
||||
- **Import safety validation**
|
||||
|
||||
|
|
@ -223,7 +230,7 @@ If `make lint` fails:
|
|||
|
||||
1. **Formatting issues**: Run `make format` to auto-fix
|
||||
2. **Ruff issues**: Check the output and fix manually
|
||||
3. **MyPy issues**: Add proper type hints
|
||||
3. **basedpyright issues**: Add proper type hints
|
||||
4. **Circular imports**: Refactor import dependencies
|
||||
5. **Import safety**: Fix any unprotected imports
|
||||
|
||||
|
|
@ -238,7 +245,7 @@ If `make test-unit` fails:
|
|||
|
||||
### 3. Common Development Tips
|
||||
|
||||
- **Use type hints**: MyPy requires proper type annotations
|
||||
- **Use type hints**: basedpyright requires proper type annotations
|
||||
- **Write descriptive commit messages**: Help reviewers understand your changes
|
||||
- **Keep PRs focused**: One feature/fix per PR
|
||||
- **Test edge cases**: Don't just test the happy path
|
||||
|
|
|
|||
30
Dockerfile
30
Dockerfile
|
|
@ -1,8 +1,8 @@
|
|||
# Base image for building
|
||||
ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:31da6565f35af6401031c1d7aa91dc84ac76c5c48edd17fb90f0ed9e3173c7a9
|
||||
ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:c61ac6919b811ea53c4782d69f1fe05218ba3c25d53f01b6ab7892e621bd4370
|
||||
|
||||
# Runtime image
|
||||
ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:31da6565f35af6401031c1d7aa91dc84ac76c5c48edd17fb90f0ed9e3173c7a9
|
||||
ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:c61ac6919b811ea53c4782d69f1fe05218ba3c25d53f01b6ab7892e621bd4370
|
||||
ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a
|
||||
|
||||
FROM $UV_IMAGE AS uvbin
|
||||
|
|
@ -68,22 +68,24 @@ FROM $LITELLM_RUNTIME_IMAGE AS runtime
|
|||
|
||||
USER root
|
||||
|
||||
RUN apk add --no-cache bash openssl tzdata nodejs npm python3 libsndfile && \
|
||||
npm install -g npm@11.14.0 tar@7.5.11 glob@13.0.6 @isaacs/brace-expansion@5.0.1 brace-expansion@5.0.5 minimatch@10.2.4 diff@8.0.3 picomatch@4.0.4 && \
|
||||
GLOBAL="$(npm root -g)" && \
|
||||
for pkg in tar glob @isaacs/brace-expansion brace-expansion minimatch diff picomatch; do \
|
||||
name="${pkg##*/}"; \
|
||||
find "$GLOBAL/npm" -type d -name "$name" -path "*/node_modules/$pkg" | while read d; do \
|
||||
rm -rf "$d" && cp -rL "$GLOBAL/$pkg" "$d"; \
|
||||
done; \
|
||||
done && \
|
||||
npm cache clean --force && \
|
||||
{ apk del --no-cache npm 2>/dev/null || true; }
|
||||
# node (without npm) is required by the prisma CLI at runtime
|
||||
RUN apk add --no-cache bash openssl tzdata nodejs python3 libsndfile
|
||||
|
||||
WORKDIR /app
|
||||
ENV PATH="/app/.venv/bin:${PATH}"
|
||||
|
||||
COPY --from=builder /app /app
|
||||
# Copy only what runtime needs. The application is installed inside the venv;
|
||||
# the rest of the builder's /app is source and build metadata that must not
|
||||
# ship (manifest-scanning tools attribute everything in it to this image).
|
||||
# entrypoint.sh invokes litellm/proxy/prisma_migration.py by source path.
|
||||
COPY --from=builder /app/.venv /app/.venv
|
||||
COPY --from=builder /app/docker /app/docker
|
||||
COPY --from=builder /app/schema.prisma /app/schema.prisma
|
||||
COPY --from=builder /app/litellm/proxy/prisma_migration.py /app/litellm/proxy/prisma_migration.py
|
||||
# enterprise/ is imported by source path at runtime (proxy_cli puts the
|
||||
# working directory on sys.path; litellm/proxy/hooks resolves
|
||||
# enterprise.enterprise_hooks from it)
|
||||
COPY --from=builder /app/enterprise /app/enterprise
|
||||
# Prisma binaries live in $HOME/.cache (default prisma-python location),
|
||||
# which is /root/.cache here. Copy only the Prisma subdirs — copying the
|
||||
# whole /root/.cache drags in the uv build cache (~660 MB, includes a
|
||||
|
|
|
|||
46
Makefile
46
Makefile
|
|
@ -5,7 +5,9 @@
|
|||
test-unit-integrations test-unit-core-utils test-unit-other test-unit-root \
|
||||
test-proxy-unit-a test-proxy-unit-b test-integration test-unit-helm \
|
||||
info lint lint-dev format \
|
||||
install-dev install-proxy-dev install-test-deps \
|
||||
lint-basedpyright lint-basedpyright-budget-update \
|
||||
lint-ruff-budget lint-ruff-budget-update lint-budget-update lint-gate \
|
||||
install-dev install-proxy-dev install-test-deps install-hooks \
|
||||
install-helm-unittest check-circular-imports check-import-safety
|
||||
|
||||
# Default target
|
||||
|
|
@ -17,12 +19,18 @@ help:
|
|||
@echo " make install-proxy-dev-ci - Install proxy dev dependencies (CI-compatible)"
|
||||
@echo " make install-test-deps - Install the full local test environment"
|
||||
@echo " make install-helm-unittest - Install helm unittest plugin"
|
||||
@echo " make install-hooks - Install git hooks (Conventional Commits + Branches)"
|
||||
@echo " make format - Apply Black code formatting"
|
||||
@echo " make format-check - Check Black code formatting (matches CI)"
|
||||
@echo " make lint - Run all linting (Ruff, MyPy, Black check, circular imports, import safety)"
|
||||
@echo " make lint - Run all linting (Ruff, basedpyright, Black check, circular imports, import safety)"
|
||||
@echo " make lint-ruff - Run Ruff linting only"
|
||||
@echo " make lint-mypy - Run MyPy type checking only"
|
||||
@echo " make lint-basedpyright - Run basedpyright strict, gated by per-rule error counts"
|
||||
@echo " make lint-basedpyright-budget-update - Re-capture the basedpyright per-rule budget (ratchet)"
|
||||
@echo " make lint-black - Check Black formatting (matches CI)"
|
||||
@echo " make lint-ruff-budget - Gate the codebase total of each strict ruff rule against its ceiling"
|
||||
@echo " make lint-gate - Strict ruff gate in CI-parity mode (fetches staging, simulates the merge)"
|
||||
@echo " make lint-ruff-budget-update - Re-capture per-rule baselines in ruff-strict-budget.json (ratchet)"
|
||||
@echo " make lint-budget-update - Re-capture all ratchet budgets (ruff + basedpyright)"
|
||||
@echo " make check-circular-imports - Check for circular imports"
|
||||
@echo " make check-import-safety - Check import safety"
|
||||
@echo " make test - Run all tests"
|
||||
|
|
@ -68,6 +76,11 @@ install-test-deps: install-proxy-dev
|
|||
install-helm-unittest:
|
||||
helm plugin install https://github.com/helm-unittest/helm-unittest --version v0.4.4 || echo "ignore error if plugin exists"
|
||||
|
||||
# Install git hooks that enforce Conventional Commits and Conventional Branches.
|
||||
# Opt-in: not chained into install-dev.
|
||||
install-hooks:
|
||||
./scripts/install_git_hooks.sh
|
||||
|
||||
# Formatting
|
||||
format: install-dev
|
||||
cd litellm && $(UV_RUN) black . && cd ..
|
||||
|
|
@ -111,11 +124,30 @@ lint-ruff-FULL-dev: install-dev
|
|||
if [ -n "$$files" ]; then echo "$$files" | xargs $(UV_RUN) ruff check; \
|
||||
else echo "No changed .py files to check."; fi
|
||||
|
||||
lint-mypy: install-dev
|
||||
cd litellm && $(UV_RUN) mypy . --ignore-missing-imports && cd ..
|
||||
lint-basedpyright: install-dev
|
||||
git fetch origin litellm_internal_staging
|
||||
($(UV_RUN) basedpyright --outputjson || true) | $(UV_RUN) python scripts/type_check_gate.py --base origin/litellm_internal_staging
|
||||
|
||||
lint-basedpyright-budget-update: install-dev
|
||||
($(UV_RUN) basedpyright --outputjson || true) | $(UV_RUN) python scripts/type_check_gate.py --update
|
||||
|
||||
lint-black: format-check
|
||||
|
||||
lint-ruff-budget: install-dev
|
||||
$(UV_RUN) python scripts/ruff_strict_gate.py
|
||||
|
||||
# Strict gate, invoked the same way CI does in test-linting.yml so a local pass
|
||||
# means the CI check will pass too.
|
||||
lint-gate: install-dev
|
||||
git fetch origin litellm_internal_staging
|
||||
$(UV_RUN) python scripts/ruff_strict_gate.py --base origin/litellm_internal_staging
|
||||
|
||||
lint-ruff-budget-update: install-dev
|
||||
$(UV_RUN) python scripts/ruff_strict_gate.py --update
|
||||
|
||||
# Ratchet all budgets in one shot (ruff strict + basedpyright)
|
||||
lint-budget-update: lint-ruff-budget-update lint-basedpyright-budget-update
|
||||
|
||||
check-circular-imports: install-dev
|
||||
cd litellm && $(UV_RUN) python ../tests/documentation_tests/test_circular_imports.py && cd ..
|
||||
|
||||
|
|
@ -123,10 +155,10 @@ check-import-safety: install-dev
|
|||
@$(UV_RUN) python -c "from litellm import *; print('[from litellm import *] OK! no issues!');" || (echo '🚨 import failed, this means you introduced unprotected imports! 🚨'; exit 1)
|
||||
|
||||
# Combined linting (matches test-linting.yml workflow)
|
||||
lint: format-check lint-ruff lint-mypy check-circular-imports check-import-safety
|
||||
lint: format-check lint-ruff lint-basedpyright check-circular-imports check-import-safety lint-ruff-budget
|
||||
|
||||
# Faster linting for local development (only checks changed code)
|
||||
lint-dev: lint-format-changed lint-mypy check-circular-imports check-import-safety
|
||||
lint-dev: lint-format-changed check-circular-imports check-import-safety
|
||||
|
||||
# Testing targets
|
||||
test: install-test-deps
|
||||
|
|
|
|||
144
README.md
144
README.md
|
|
@ -6,10 +6,10 @@
|
|||
</p>
|
||||
<p align="center">Open Source AI Gateway for 100+ LLMs. Self-hosted. Enterprise-ready. Call any LLM in OpenAI format.</p>
|
||||
<p align="center">
|
||||
<a href="https://render.com/deploy?repo=https://github.com/BerriAI/litellm" target="_blank" rel="nofollow"><img src="https://render.com/images/deploy-to-render-button.svg" alt="Deploy to Render"></a>
|
||||
<a href="https://railway.com/deploy/RhvhdC?referralCode=7mRv9K&utm_medium=integration&utm_source=template&utm_campaign=generic">
|
||||
<img src="https://railway.com/button.svg" alt="Deploy on Railway">
|
||||
</a>
|
||||
<a href="https://render.com/deploy?repo=https://github.com/BerriAI/litellm" target="_blank" rel="nofollow"><img src="https://render.com/images/deploy-to-render-button.svg" alt="Deploy to Render" height="40"></a>
|
||||
<a href="https://railway.com/deploy/RhvhdC?referralCode=7mRv9K&utm_medium=integration&utm_source=template&utm_campaign=generic"><img src="https://railway.com/button.svg" alt="Deploy on Railway" height="40"></a>
|
||||
<a href="https://console.aws.amazon.com/cloudshell/home" target="_blank" rel="nofollow"><img src="./.github/deploy-on-aws.png" alt="Deploy on AWS" height="40"></a>
|
||||
<a href="https://ssh.cloud.google.com/cloudshell/editor?cloudshell_git_repo=https%3A%2F%2Fgithub.com%2FBerriAI%2Flitellm&cloudshell_workspace=terraform%2Flitellm%2Fgcp%2Fexamples%2Fdefault&cloudshell_tutorial=TUTORIAL.md&cloudshell_image=gcr.io/ds-artifacts-cloudshell/deploystack_custom_image&shellonly=true" target="_blank" rel="nofollow"><img src="./.github/deploy-on-gcp.png" alt="Deploy on GCP" height="40"></a>
|
||||
</p>
|
||||
</p>
|
||||
<h4 align="center"><a href="https://docs.litellm.ai/docs/simple_proxy" target="_blank">LiteLLM Proxy Server (AI Gateway)</a> | <a href="https://docs.litellm.ai/docs/enterprise#hosted-litellm-proxy" target="_blank"> Hosted Proxy</a> | <a href="https://litellm.ai/enterprise"target="_blank">Enterprise Tier</a> | <a href="https://www.litellm.ai/ai-gateway" target="_blank">Website</a></h4>
|
||||
|
|
@ -327,6 +327,7 @@ curl -X POST 'http://0.0.0.0:4000/v1/chat/completions' \
|
|||
| [Maritalk (`maritalk`)](https://docs.litellm.ai/docs/providers/maritalk) | ✅ | ✅ | ✅ | | | | | | | |
|
||||
| [Meta - Llama API (`meta_llama`)](https://docs.litellm.ai/docs/providers/meta_llama) | ✅ | ✅ | ✅ | | | | | | | |
|
||||
| [Mistral AI API (`mistral`)](https://docs.litellm.ai/docs/providers/mistral) | ✅ | ✅ | ✅ | ✅ | | | | | | |
|
||||
| [ModelScope (`modelscope`)](https://docs.litellm.ai/docs/providers/modelscope) | ✅ | ✅ | ✅ | | ✅ | | | | | |
|
||||
| [Moonshot (`moonshot`)](https://docs.litellm.ai/docs/providers/moonshot) | ✅ | ✅ | ✅ | | | | | | | |
|
||||
| [Morph (`morph`)](https://docs.litellm.ai/docs/providers/morph) | ✅ | ✅ | ✅ | | | | | | | |
|
||||
| [Nebius AI Studio (`nebius`)](https://docs.litellm.ai/docs/providers/nebius) | ✅ | ✅ | ✅ | ✅ | | | | | | |
|
||||
|
|
@ -344,6 +345,7 @@ curl -X POST 'http://0.0.0.0:4000/v1/chat/completions' \
|
|||
| [OVHCloud AI Endpoints (`ovhcloud`)](https://docs.litellm.ai/docs/providers/ovhcloud) | ✅ | ✅ | ✅ | | | | | | | |
|
||||
| [Perplexity AI (`perplexity`)](https://docs.litellm.ai/docs/providers/perplexity) | ✅ | ✅ | ✅ | | | | | | | |
|
||||
| [Petals (`petals`)](https://docs.litellm.ai/docs/providers/petals) | ✅ | ✅ | ✅ | | | | | | | |
|
||||
| [Pinstripes (`pinstripes`)](https://docs.litellm.ai/docs/providers/pinstripes) | ✅ | ✅ | ✅ | | | | | | | |
|
||||
| [Predibase (`predibase`)](https://docs.litellm.ai/docs/providers/predibase) | ✅ | ✅ | ✅ | | | | | | | |
|
||||
| [Recraft (`recraft`)](https://docs.litellm.ai/docs/providers/recraft) | | | | | ✅ | | | | | |
|
||||
| [Replicate (`replicate`)](https://docs.litellm.ai/docs/providers/replicate) | ✅ | ✅ | ✅ | | | | | | | |
|
||||
|
|
@ -404,6 +406,140 @@ You can use LiteLLM through either the Proxy Server or Python SDK. Both give you
|
|||
|
||||
Support for more providers. Missing a provider or LLM Platform, raise a [feature request](https://github.com/BerriAI/litellm/issues/new?assignees=&labels=enhancement&projects=&template=feature_request.yml&title=%5BFeature%5D%3A+).
|
||||
|
||||
### Deploy on AWS or GCP with Terraform
|
||||
|
||||
Run the LiteLLM proxy as a production-ready componentized stack (gateway, backend, UI on separate services; managed Postgres + Redis + object store) using the published Terraform modules. Both modules are on the [public Terraform Registry](https://registry.terraform.io/namespaces/BerriAI) — no auth needed.
|
||||
|
||||
#### AWS — ECS Fargate + Aurora + ElastiCache + ALB
|
||||
|
||||
[](https://console.aws.amazon.com/cloudshell/home) — opens an in-browser shell, already authenticated to your AWS account. Once inside, run:
|
||||
|
||||
```bash
|
||||
git clone https://github.com/BerriAI/litellm.git
|
||||
cd litellm/terraform/litellm/aws/examples/default
|
||||
cp terraform.tfvars.example terraform.tfvars # edit region/tenant/env
|
||||
terraform init && terraform apply
|
||||
```
|
||||
|
||||
[Module page →](https://registry.terraform.io/modules/BerriAI/litellm/aws/latest)
|
||||
|
||||
Or call the module from your own root config:
|
||||
|
||||
```hcl
|
||||
# main.tf
|
||||
terraform {
|
||||
required_version = ">= 1.6.0"
|
||||
required_providers {
|
||||
aws = { source = "hashicorp/aws", version = "~> 5.60" }
|
||||
}
|
||||
}
|
||||
|
||||
provider "aws" {
|
||||
region = "us-west-2"
|
||||
}
|
||||
|
||||
module "litellm" {
|
||||
source = "BerriAI/litellm/aws"
|
||||
version = "~> 1.89"
|
||||
|
||||
region = "us-west-2"
|
||||
azs = ["us-west-2a", "us-west-2b"]
|
||||
tenant = "acme"
|
||||
env = "prod"
|
||||
|
||||
# Production: provide an ACM cert. Without one, set allow_plaintext_alb = true
|
||||
# (dev/trial only).
|
||||
# acm_certificate_arn = "arn:aws:acm:us-west-2:111122223333:certificate/..."
|
||||
allow_plaintext_alb = true
|
||||
}
|
||||
|
||||
output "litellm_url" {
|
||||
value = module.litellm.alb_dns_name
|
||||
}
|
||||
```
|
||||
|
||||
```bash
|
||||
terraform init
|
||||
terraform apply
|
||||
```
|
||||
|
||||
Provider API keys live in AWS Secrets Manager; reference ARNs via `gateway_extra_secrets`. Full input list and architecture diagram on the [registry page](https://registry.terraform.io/modules/BerriAI/litellm/aws/latest?tab=inputs).
|
||||
|
||||
#### GCP — Cloud Run + Cloud SQL + Memorystore + HTTPS LB
|
||||
|
||||
[](https://ssh.cloud.google.com/cloudshell/editor?cloudshell_git_repo=https%3A%2F%2Fgithub.com%2FBerriAI%2Flitellm&cloudshell_workspace=terraform%2Flitellm%2Fgcp%2Fexamples%2Fdefault&cloudshell_tutorial=TUTORIAL.md&cloudshell_image=gcr.io/ds-artifacts-cloudshell/deploystack_custom_image&shellonly=true)
|
||||
|
||||
Real 1-click. Opens Cloud Shell, clones this repo, and walks you through `terraform apply` via a built-in [DeployStack tutorial](./terraform/litellm/gcp/examples/default/TUTORIAL.md) — pick the project, the tutorial sets up the Artifact Registry remote repo, writes `terraform.tfvars` from your answers, and runs apply.
|
||||
|
||||
[Module page →](https://registry.terraform.io/modules/BerriAI/litellm/google/latest)
|
||||
|
||||
To call the module from your own config instead, Cloud Run can't pull from `ghcr.io` directly, so first set up a one-time Artifact Registry remote repo backed by GHCR:
|
||||
|
||||
```bash
|
||||
gcloud artifacts repositories create litellm \
|
||||
--location=us-central1 \
|
||||
--repository-format=docker \
|
||||
--mode=remote-repository \
|
||||
--remote-docker-repo=https://ghcr.io \
|
||||
--project=my-gcp-project
|
||||
```
|
||||
|
||||
Then:
|
||||
|
||||
```hcl
|
||||
# main.tf
|
||||
terraform {
|
||||
required_version = ">= 1.6.0"
|
||||
required_providers {
|
||||
google = { source = "hashicorp/google", version = "~> 6.10" }
|
||||
google-beta = { source = "hashicorp/google-beta", version = "~> 6.10" }
|
||||
}
|
||||
}
|
||||
|
||||
provider "google" { project = "my-gcp-project"; region = "us-central1" }
|
||||
provider "google-beta" { project = "my-gcp-project"; region = "us-central1" }
|
||||
|
||||
module "litellm" {
|
||||
source = "BerriAI/litellm/google"
|
||||
version = "~> 1.89"
|
||||
|
||||
project_id = "my-gcp-project"
|
||||
region = "us-central1"
|
||||
tenant = "acme"
|
||||
env = "prod"
|
||||
|
||||
# Replace my-gcp-project with your GCP project ID (same value as project_id above).
|
||||
image_registry = "us-central1-docker.pkg.dev/my-gcp-project/litellm/berriai"
|
||||
|
||||
# Production: provide DNS already pointing at the LB IP for Google-managed certs.
|
||||
# Without one, set allow_plaintext_lb = true (dev/trial only).
|
||||
# lb_domains = ["proxy.example.com"]
|
||||
allow_plaintext_lb = true
|
||||
}
|
||||
|
||||
output "litellm_url" {
|
||||
value = module.litellm.load_balancer_url
|
||||
}
|
||||
```
|
||||
|
||||
```bash
|
||||
terraform init
|
||||
terraform apply
|
||||
```
|
||||
|
||||
Provider API keys live in Secret Manager; reference resource IDs (e.g. `projects/my-gcp-project/secrets/openai-api-key`) via `gateway_extra_secrets`. Full input list and architecture diagram on the [registry page](https://registry.terraform.io/modules/BerriAI/litellm/google/latest?tab=inputs).
|
||||
|
||||
#### Both stacks include
|
||||
|
||||
- The full componentized split (gateway / backend / UI as independent services)
|
||||
- Managed Postgres (writer + reader) and Redis
|
||||
- Versioned object store for proxy state + file uploads
|
||||
- An auto-generated `LITELLM_MASTER_KEY` in your cloud's secret manager
|
||||
- A one-off migration job that runs `prisma migrate deploy` before the proxy starts
|
||||
- The same `proxy_config` surface as the [Helm chart](./helm/litellm/) — pass YAML as a typed map
|
||||
|
||||
The Terraform modules live at [`terraform/litellm/aws/`](./terraform/litellm/aws/) and [`terraform/litellm/gcp/`](./terraform/litellm/gcp/) in this repo; the registry entries are read-only mirrors updated on each release.
|
||||
|
||||
### Run in Developer Mode
|
||||
#### Services
|
||||
1. Setup .env file in root
|
||||
|
|
|
|||
|
|
@ -1,5 +1,5 @@
|
|||
ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:31da6565f35af6401031c1d7aa91dc84ac76c5c48edd17fb90f0ed9e3173c7a9
|
||||
ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:31da6565f35af6401031c1d7aa91dc84ac76c5c48edd17fb90f0ed9e3173c7a9
|
||||
ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:c61ac6919b811ea53c4782d69f1fe05218ba3c25d53f01b6ab7892e621bd4370
|
||||
ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:c61ac6919b811ea53c4782d69f1fe05218ba3c25d53f01b6ab7892e621bd4370
|
||||
ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a
|
||||
|
||||
FROM $UV_IMAGE AS uvbin
|
||||
|
|
|
|||
|
|
@ -20,7 +20,11 @@ DatabaseURLSettings.from_env().apply_to_env()
|
|||
|
||||
from litellm.proxy.proxy_server import app
|
||||
|
||||
from backend.routes.allowlist import BACKEND_EXACT_PATHS, BACKEND_PATH_PREFIXES
|
||||
from backend.routes.allowlist import (
|
||||
BACKEND_EXACT_PATHS,
|
||||
BACKEND_MOUNT_PATHS,
|
||||
BACKEND_PATH_PREFIXES,
|
||||
)
|
||||
|
||||
|
||||
def _is_backend_route(route) -> bool:
|
||||
|
|
@ -29,8 +33,9 @@ def _is_backend_route(route) -> bool:
|
|||
if path is None:
|
||||
return False
|
||||
if isinstance(route, Mount):
|
||||
# Static UI mounts are served by the dedicated UI container, not here.
|
||||
return False
|
||||
# The dashboard UI static mounts are served by the dedicated UI container.
|
||||
# Only Mounts in the backend allowlist (e.g. swagger docs) remain on backend.
|
||||
return path in BACKEND_MOUNT_PATHS
|
||||
if path in BACKEND_EXACT_PATHS:
|
||||
return True
|
||||
return any(path.startswith(prefix) for prefix in BACKEND_PATH_PREFIXES)
|
||||
|
|
|
|||
|
|
@ -120,6 +120,9 @@ BACKEND_PATH_PREFIXES: tuple[str, ...] = (
|
|||
"/robots.txt",
|
||||
# Health (k8s probes)
|
||||
"/health",
|
||||
# Plugin system
|
||||
"/api/plugins",
|
||||
"/plugin-proxy/",
|
||||
)
|
||||
|
||||
BACKEND_EXACT_PATHS: frozenset[str] = frozenset(
|
||||
|
|
@ -133,3 +136,9 @@ BACKEND_EXACT_PATHS: frozenset[str] = frozenset(
|
|||
"/fallback/login",
|
||||
}
|
||||
)
|
||||
|
||||
BACKEND_MOUNT_PATHS: frozenset[str] = frozenset(
|
||||
{
|
||||
"/swagger", # API documentation static assets belong to the backend
|
||||
}
|
||||
)
|
||||
|
|
|
|||
194
basedpyright-code-budget.json
Normal file
194
basedpyright-code-budget.json
Normal file
|
|
@ -0,0 +1,194 @@
|
|||
{
|
||||
"reportAny": {
|
||||
"baseline": 24989,
|
||||
"slack": 2500
|
||||
},
|
||||
"reportArgumentType": {
|
||||
"baseline": 1934,
|
||||
"slack": 180
|
||||
},
|
||||
"reportAssignmentType": {
|
||||
"baseline": 220,
|
||||
"slack": 22
|
||||
},
|
||||
"reportAttributeAccessIssue": {
|
||||
"baseline": 346,
|
||||
"slack": 35
|
||||
},
|
||||
"reportCallIssue": {
|
||||
"baseline": 87,
|
||||
"slack": 10
|
||||
},
|
||||
"reportConstantRedefinition": {
|
||||
"baseline": 39,
|
||||
"slack": 4
|
||||
},
|
||||
"reportDeprecated": {
|
||||
"baseline": 217,
|
||||
"slack": 22
|
||||
},
|
||||
"reportDuplicateImport": {
|
||||
"baseline": 28,
|
||||
"slack": 3
|
||||
},
|
||||
"reportExplicitAny": {
|
||||
"baseline": 6931,
|
||||
"slack": 700
|
||||
},
|
||||
"reportFunctionMemberAccess": {
|
||||
"baseline": 7,
|
||||
"slack": 3
|
||||
},
|
||||
"reportGeneralTypeIssues": {
|
||||
"baseline": 151,
|
||||
"slack": 15
|
||||
},
|
||||
"reportIncompatibleMethodOverride": {
|
||||
"baseline": 52,
|
||||
"slack": 5
|
||||
},
|
||||
"reportIncompatibleVariableOverride": {
|
||||
"baseline": 8,
|
||||
"slack": 3
|
||||
},
|
||||
"reportInconsistentOverload": {
|
||||
"baseline": 12,
|
||||
"slack": 3
|
||||
},
|
||||
"reportIndexIssue": {
|
||||
"baseline": 26,
|
||||
"slack": 3
|
||||
},
|
||||
"reportInvalidTypeForm": {
|
||||
"baseline": 23,
|
||||
"slack": 3
|
||||
},
|
||||
"reportInvalidTypeVarUse": {
|
||||
"baseline": 2,
|
||||
"slack": 3
|
||||
},
|
||||
"reportMatchNotExhaustive": {
|
||||
"baseline": 1,
|
||||
"slack": 3
|
||||
},
|
||||
"reportMissingParameterType": {
|
||||
"baseline": 3933,
|
||||
"slack": 390
|
||||
},
|
||||
"reportMissingTypeArgument": {
|
||||
"baseline": 10612,
|
||||
"slack": 1000
|
||||
},
|
||||
"reportMissingTypeStubs": {
|
||||
"baseline": 27,
|
||||
"slack": 10
|
||||
},
|
||||
"reportOperatorIssue": {
|
||||
"baseline": 6,
|
||||
"slack": 3
|
||||
},
|
||||
"reportOptionalCall": {
|
||||
"baseline": 4,
|
||||
"slack": 3
|
||||
},
|
||||
"reportOptionalIterable": {
|
||||
"baseline": 3,
|
||||
"slack": 3
|
||||
},
|
||||
"reportOptionalMemberAccess": {
|
||||
"baseline": 724,
|
||||
"slack": 72
|
||||
},
|
||||
"reportOptionalOperand": {
|
||||
"baseline": 3,
|
||||
"slack": 3
|
||||
},
|
||||
"reportOptionalSubscript": {
|
||||
"baseline": 11,
|
||||
"slack": 3
|
||||
},
|
||||
"reportPossiblyUnboundVariable": {
|
||||
"baseline": 52,
|
||||
"slack": 10
|
||||
},
|
||||
"reportPrivateUsage": {
|
||||
"baseline": 1625,
|
||||
"slack": 160
|
||||
},
|
||||
"reportRedeclaration": {
|
||||
"baseline": 8,
|
||||
"slack": 3
|
||||
},
|
||||
"reportReturnType": {
|
||||
"baseline": 126,
|
||||
"slack": 100
|
||||
},
|
||||
"reportTypedDictNotRequiredAccess": {
|
||||
"baseline": 20,
|
||||
"slack": 3
|
||||
},
|
||||
"reportUndefinedVariable": {
|
||||
"baseline": 2,
|
||||
"slack": 3
|
||||
},
|
||||
"reportUnknownArgumentType": {
|
||||
"baseline": 30603,
|
||||
"slack": 3000
|
||||
},
|
||||
"reportUnknownLambdaType": {
|
||||
"baseline": 75,
|
||||
"slack": 10
|
||||
},
|
||||
"reportUnknownMemberType": {
|
||||
"baseline": 27037,
|
||||
"slack": 2500
|
||||
},
|
||||
"reportUnknownParameterType": {
|
||||
"baseline": 13612,
|
||||
"slack": 1000
|
||||
},
|
||||
"reportUnknownVariableType": {
|
||||
"baseline": 21445,
|
||||
"slack": 2000
|
||||
},
|
||||
"reportUnnecessaryCast": {
|
||||
"baseline": 118,
|
||||
"slack": 10
|
||||
},
|
||||
"reportUnnecessaryComparison": {
|
||||
"baseline": 683,
|
||||
"slack": 100
|
||||
},
|
||||
"reportUnnecessaryContains": {
|
||||
"baseline": 4,
|
||||
"slack": 3
|
||||
},
|
||||
"reportUnnecessaryIsInstance": {
|
||||
"baseline": 808,
|
||||
"slack": 80
|
||||
},
|
||||
"reportUntypedBaseClass": {
|
||||
"baseline": 110,
|
||||
"slack": 11
|
||||
},
|
||||
"reportUntypedFunctionDecorator": {
|
||||
"baseline": 22,
|
||||
"slack": 3
|
||||
},
|
||||
"reportUnusedClass": {
|
||||
"baseline": 22,
|
||||
"slack": 3
|
||||
},
|
||||
"reportUnusedFunction": {
|
||||
"baseline": 137,
|
||||
"slack": 10
|
||||
},
|
||||
"reportUnusedImport": {
|
||||
"baseline": 670,
|
||||
"slack": 50
|
||||
},
|
||||
"reportUnusedVariable": {
|
||||
"baseline": 865,
|
||||
"slack": 50
|
||||
}
|
||||
}
|
||||
16
codecov.yaml
16
codecov.yaml
|
|
@ -35,6 +35,22 @@ component_management:
|
|||
- component_id: "Enterprise"
|
||||
paths:
|
||||
- "enterprise/**"
|
||||
- component_id: "Batches"
|
||||
paths:
|
||||
- "*/proxy/batches_endpoints/**"
|
||||
- "litellm/batches/**"
|
||||
- "*/llms/*/batches/**"
|
||||
- component_id: "Videos"
|
||||
paths:
|
||||
- "litellm/videos/**"
|
||||
- "*/proxy/video_endpoints/**"
|
||||
- "*/llms/*/videos/**"
|
||||
- component_id: "Realtime"
|
||||
paths:
|
||||
- "litellm/realtime_api/**"
|
||||
- "*/proxy/realtime_endpoints/**"
|
||||
- "*/llms/*/realtime/**"
|
||||
- "litellm/litellm_core_utils/realtime_streaming.py"
|
||||
comment:
|
||||
layout: "header, diff, flags, components" # show component info in the PR comment
|
||||
|
||||
|
|
|
|||
|
|
@ -15,7 +15,7 @@ db = Prisma(
|
|||
)
|
||||
|
||||
|
||||
async def check_view_exists(): # noqa: PLR0915
|
||||
async def check_view_exists():
|
||||
"""
|
||||
Checks if the LiteLLM_VerificationTokenView and MonthlyGlobalSpend exists in the user's db.
|
||||
|
||||
|
|
@ -34,8 +34,7 @@ async def check_view_exists(): # noqa: PLR0915
|
|||
print("LiteLLM_VerificationTokenView Exists!") # noqa
|
||||
except Exception:
|
||||
# If an error occurs, the view does not exist, so create it
|
||||
await db.execute_raw(
|
||||
"""
|
||||
await db.execute_raw("""
|
||||
CREATE VIEW "LiteLLM_VerificationTokenView" AS
|
||||
SELECT
|
||||
v.*,
|
||||
|
|
@ -45,8 +44,7 @@ async def check_view_exists(): # noqa: PLR0915
|
|||
t.rpm_limit AS team_rpm_limit
|
||||
FROM "LiteLLM_VerificationToken" v
|
||||
LEFT JOIN "LiteLLM_TeamTable" t ON v.team_id = t.team_id;
|
||||
"""
|
||||
)
|
||||
""")
|
||||
|
||||
print("LiteLLM_VerificationTokenView Created!") # noqa
|
||||
|
||||
|
|
|
|||
99
db_scripts/partition_spend_logs.sql
Normal file
99
db_scripts/partition_spend_logs.sql
Normal file
|
|
@ -0,0 +1,99 @@
|
|||
-- Converts an existing LiteLLM_SpendLogs table into a native Postgres
|
||||
-- range-partitioned table keyed on "startTime".
|
||||
--
|
||||
-- Why: at high request volume, retention via DELETE leaves dead tuples that
|
||||
-- autovacuum cannot reclaim quickly enough, so the table keeps growing on disk
|
||||
-- (seen at 450GB+ after ~1 month). With partitioning, retention drops whole
|
||||
-- partitions, which is instant and returns disk to the OS immediately.
|
||||
--
|
||||
-- This is an opt-in, manual operation. The default LiteLLM schema is NOT
|
||||
-- partitioned, so existing installs are unaffected until you run this.
|
||||
--
|
||||
-- IMPORTANT
|
||||
-- * Test on a staging copy first and take a backup.
|
||||
-- * Postgres cannot convert a populated table to partitioned in place, so this
|
||||
-- renames the old table aside and creates a fresh partitioned table.
|
||||
-- * The partition key ("startTime") must be part of the primary key, so the
|
||||
-- PK becomes composite ("request_id", "startTime"). LiteLLM's write path uses
|
||||
-- INSERT ... ON CONFLICT DO NOTHING, which is compatible with this.
|
||||
-- * Choose a partition granularity ("day" is the recommended default for
|
||||
-- high-volume tables) and keep it consistent with SPEND_LOG_PARTITION_INTERVAL.
|
||||
--
|
||||
-- After running this, enable the feature and set a retention period in
|
||||
-- proxy_config.yaml:
|
||||
-- general_settings:
|
||||
-- use_spend_logs_partitioning: true
|
||||
-- maximum_spend_logs_retention_period: "30d"
|
||||
-- The spend-log cleanup job then verifies the table is partitioned and reclaims
|
||||
-- disk by dropping expired partitions instead of deleting rows. It also
|
||||
-- pre-creates upcoming partitions on each run. To roll back, see
|
||||
-- db_scripts/unpartition_spend_logs.sql.
|
||||
|
||||
BEGIN;
|
||||
|
||||
ALTER TABLE "LiteLLM_SpendLogs" RENAME TO "LiteLLM_SpendLogs_legacy";
|
||||
|
||||
-- Renaming a table does NOT rename its indexes, and index names are unique per
|
||||
-- schema. Move the legacy table's indexes aside so the CREATE INDEX statements
|
||||
-- below actually create indexes on the new partitioned table instead of being
|
||||
-- silently skipped by IF NOT EXISTS, and so the new PK keeps the canonical
|
||||
-- name instead of getting a "_pkey1" suffix.
|
||||
ALTER INDEX IF EXISTS "LiteLLM_SpendLogs_pkey"
|
||||
RENAME TO "LiteLLM_SpendLogs_legacy_pkey";
|
||||
ALTER INDEX IF EXISTS "LiteLLM_SpendLogs_startTime_idx"
|
||||
RENAME TO "LiteLLM_SpendLogs_legacy_startTime_idx";
|
||||
ALTER INDEX IF EXISTS "LiteLLM_SpendLogs_startTime_request_id_idx"
|
||||
RENAME TO "LiteLLM_SpendLogs_legacy_startTime_request_id_idx";
|
||||
ALTER INDEX IF EXISTS "LiteLLM_SpendLogs_end_user_idx"
|
||||
RENAME TO "LiteLLM_SpendLogs_legacy_end_user_idx";
|
||||
ALTER INDEX IF EXISTS "LiteLLM_SpendLogs_session_id_idx"
|
||||
RENAME TO "LiteLLM_SpendLogs_legacy_session_id_idx";
|
||||
|
||||
CREATE TABLE "LiteLLM_SpendLogs" (
|
||||
LIKE "LiteLLM_SpendLogs_legacy" INCLUDING DEFAULTS INCLUDING GENERATED
|
||||
) PARTITION BY RANGE ("startTime");
|
||||
|
||||
ALTER TABLE "LiteLLM_SpendLogs"
|
||||
ADD PRIMARY KEY ("request_id", "startTime");
|
||||
|
||||
-- Recreate every index Prisma defines on the table. LIKE ... INCLUDING DEFAULTS
|
||||
-- INCLUDING GENERATED copies columns and defaults but NOT indexes, so without
|
||||
-- these the admin-UI cost-reporting queries that filter by end_user/session_id
|
||||
-- fall back to sequential scans. On a partitioned parent these propagate to
|
||||
-- every current and future partition automatically.
|
||||
CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogs_startTime_idx"
|
||||
ON "LiteLLM_SpendLogs" ("startTime");
|
||||
|
||||
CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogs_startTime_request_id_idx"
|
||||
ON "LiteLLM_SpendLogs" ("startTime", "request_id");
|
||||
|
||||
CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogs_end_user_idx"
|
||||
ON "LiteLLM_SpendLogs" ("end_user");
|
||||
|
||||
CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogs_session_id_idx"
|
||||
ON "LiteLLM_SpendLogs" ("session_id");
|
||||
|
||||
-- Safety net: any row whose startTime has no explicit partition lands here so
|
||||
-- writes never fail. The cleanup job never drops the DEFAULT partition.
|
||||
CREATE TABLE IF NOT EXISTS "LiteLLM_SpendLogs_pdefault"
|
||||
PARTITION OF "LiteLLM_SpendLogs" DEFAULT;
|
||||
|
||||
COMMIT;
|
||||
|
||||
-- Backfill (optional). Rows route to the correct partition automatically.
|
||||
-- For large legacy tables, copy in time-bounded batches during a low-traffic
|
||||
-- window instead of one statement, or simply keep "LiteLLM_SpendLogs_legacy"
|
||||
-- read-only until its data ages past your retention, then DROP it.
|
||||
--
|
||||
-- Backfilled rows land in the DEFAULT partition until explicit partitions
|
||||
-- cover their dates. Postgres refuses to create a partition whose range
|
||||
-- overlaps rows already in DEFAULT, so the cleanup job may log a warning when
|
||||
-- pre-creating today's partition right after a backfill; it recovers on its
|
||||
-- own once those dates age out, and future partitions are unaffected because
|
||||
-- they are always created ahead of writes.
|
||||
--
|
||||
-- INSERT INTO "LiteLLM_SpendLogs"
|
||||
-- SELECT * FROM "LiteLLM_SpendLogs_legacy"
|
||||
-- WHERE "startTime" >= now() - interval '30 days';
|
||||
--
|
||||
-- DROP TABLE "LiteLLM_SpendLogs_legacy";
|
||||
69
db_scripts/unpartition_spend_logs.sql
Normal file
69
db_scripts/unpartition_spend_logs.sql
Normal file
|
|
@ -0,0 +1,69 @@
|
|||
-- Rolls back db_scripts/partition_spend_logs.sql: converts the native
|
||||
-- range-partitioned "LiteLLM_SpendLogs" table back into a plain,
|
||||
-- non-partitioned table matching the default LiteLLM schema.
|
||||
--
|
||||
-- When/why: run this if you want to stop using partition-based retention and
|
||||
-- return to DELETE-based cleanup, or to restore the original single-column
|
||||
-- primary key ("request_id") that the partitioned layout had to widen to a
|
||||
-- composite ("request_id", "startTime").
|
||||
--
|
||||
-- IMPORTANT
|
||||
-- * Test on a staging copy first and take a backup.
|
||||
-- * Postgres cannot convert a partitioned table back in place, so this
|
||||
-- renames the partitioned table aside and creates a fresh plain table.
|
||||
-- * The composite PK could in principle hold the same "request_id" in more
|
||||
-- than one partition, so rows are copied with ON CONFLICT DO NOTHING to
|
||||
-- restore the single-column PK without failing on such duplicates.
|
||||
-- * For large tables the INSERT ... SELECT copies every surviving row and may
|
||||
-- run long; do it during a low-traffic window.
|
||||
-- * Also remove use_spend_logs_partitioning from proxy_config.yaml (or set it
|
||||
-- to false) so the cleanup job returns to DELETE-based retention.
|
||||
|
||||
BEGIN;
|
||||
|
||||
ALTER TABLE "LiteLLM_SpendLogs" RENAME TO "LiteLLM_SpendLogs_partitioned";
|
||||
|
||||
-- Renaming a table does NOT rename its indexes, and index names are unique per
|
||||
-- schema. Move the partitioned table's indexes aside so the CREATE INDEX
|
||||
-- statements below actually create indexes on the new plain table instead of
|
||||
-- being silently skipped by IF NOT EXISTS, and so the new PK keeps the
|
||||
-- canonical name.
|
||||
ALTER INDEX IF EXISTS "LiteLLM_SpendLogs_pkey"
|
||||
RENAME TO "LiteLLM_SpendLogs_partitioned_pkey";
|
||||
ALTER INDEX IF EXISTS "LiteLLM_SpendLogs_pkey1"
|
||||
RENAME TO "LiteLLM_SpendLogs_partitioned_pkey1";
|
||||
ALTER INDEX IF EXISTS "LiteLLM_SpendLogs_startTime_idx"
|
||||
RENAME TO "LiteLLM_SpendLogs_partitioned_startTime_idx";
|
||||
ALTER INDEX IF EXISTS "LiteLLM_SpendLogs_startTime_request_id_idx"
|
||||
RENAME TO "LiteLLM_SpendLogs_partitioned_startTime_request_id_idx";
|
||||
ALTER INDEX IF EXISTS "LiteLLM_SpendLogs_end_user_idx"
|
||||
RENAME TO "LiteLLM_SpendLogs_partitioned_end_user_idx";
|
||||
ALTER INDEX IF EXISTS "LiteLLM_SpendLogs_session_id_idx"
|
||||
RENAME TO "LiteLLM_SpendLogs_partitioned_session_id_idx";
|
||||
|
||||
CREATE TABLE "LiteLLM_SpendLogs" (
|
||||
LIKE "LiteLLM_SpendLogs_partitioned" INCLUDING DEFAULTS INCLUDING GENERATED
|
||||
);
|
||||
|
||||
ALTER TABLE "LiteLLM_SpendLogs"
|
||||
ADD PRIMARY KEY ("request_id");
|
||||
|
||||
CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogs_startTime_idx"
|
||||
ON "LiteLLM_SpendLogs" ("startTime");
|
||||
|
||||
CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogs_startTime_request_id_idx"
|
||||
ON "LiteLLM_SpendLogs" ("startTime", "request_id");
|
||||
|
||||
CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogs_end_user_idx"
|
||||
ON "LiteLLM_SpendLogs" ("end_user");
|
||||
|
||||
CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogs_session_id_idx"
|
||||
ON "LiteLLM_SpendLogs" ("session_id");
|
||||
|
||||
INSERT INTO "LiteLLM_SpendLogs"
|
||||
SELECT * FROM "LiteLLM_SpendLogs_partitioned"
|
||||
ON CONFLICT ("request_id") DO NOTHING;
|
||||
|
||||
DROP TABLE "LiteLLM_SpendLogs_partitioned";
|
||||
|
||||
COMMIT;
|
||||
|
|
@ -1,8 +1,8 @@
|
|||
# Base image for building
|
||||
ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:3258be472764337fd13095bcbb3182da170243b5819fd67ad4c0754590588b31
|
||||
ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:c61ac6919b811ea53c4782d69f1fe05218ba3c25d53f01b6ab7892e621bd4370
|
||||
|
||||
# Runtime image
|
||||
ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:3258be472764337fd13095bcbb3182da170243b5819fd67ad4c0754590588b31
|
||||
ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:c61ac6919b811ea53c4782d69f1fe05218ba3c25d53f01b6ab7892e621bd4370
|
||||
ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a
|
||||
|
||||
FROM $UV_IMAGE AS uvbin
|
||||
|
|
@ -66,36 +66,31 @@ FROM $LITELLM_RUNTIME_IMAGE AS runtime
|
|||
|
||||
USER root
|
||||
|
||||
RUN apk add --no-cache bash openssl tzdata nodejs npm python3 libsndfile && \
|
||||
npm install -g npm@11.12.1 tar@7.5.11 glob@11.1.0 @isaacs/brace-expansion@5.0.1 minimatch@10.2.4 diff@8.0.3 && \
|
||||
GLOBAL="$(npm root -g)" && \
|
||||
find "$GLOBAL/npm" -type d -name "tar" -path "*/node_modules/tar" | while read d; do \
|
||||
rm -rf "$d" && cp -rL "$GLOBAL/tar" "$d"; \
|
||||
done && \
|
||||
find "$GLOBAL/npm" -type d -name "glob" -path "*/node_modules/glob" | while read d; do \
|
||||
rm -rf "$d" && cp -rL "$GLOBAL/glob" "$d"; \
|
||||
done && \
|
||||
find "$GLOBAL/npm" -type d -name "brace-expansion" -path "*/node_modules/@isaacs/brace-expansion" | while read d; do \
|
||||
rm -rf "$d" && cp -rL "$GLOBAL/@isaacs/brace-expansion" "$d"; \
|
||||
done && \
|
||||
find "$GLOBAL/npm" -type d -name "minimatch" -path "*/node_modules/minimatch" | while read d; do \
|
||||
rm -rf "$d" && cp -rL "$GLOBAL/minimatch" "$d"; \
|
||||
done && \
|
||||
find "$GLOBAL/npm" -type d -name "diff" -path "*/node_modules/diff" | while read d; do \
|
||||
rm -rf "$d" && cp -rL "$GLOBAL/diff" "$d"; \
|
||||
done && \
|
||||
npm cache clean --force && \
|
||||
{ apk del --no-cache npm 2>/dev/null || true; }
|
||||
# node (without npm) is required by the prisma CLI at runtime
|
||||
RUN apk add --no-cache bash openssl tzdata nodejs python3 libsndfile
|
||||
|
||||
WORKDIR /app
|
||||
ENV PATH="/app/.venv/bin:${PATH}"
|
||||
|
||||
COPY --from=builder /app /app
|
||||
# Copy only what runtime needs. The application is installed inside the venv;
|
||||
# the rest of the builder's /app is source and build metadata that must not
|
||||
# ship (manifest-scanning tools attribute everything in it to this image).
|
||||
# entrypoint.sh invokes litellm/proxy/prisma_migration.py by source path.
|
||||
COPY --from=builder /app/.venv /app/.venv
|
||||
COPY --from=builder /app/docker /app/docker
|
||||
COPY --from=builder /app/schema.prisma /app/schema.prisma
|
||||
COPY --from=builder /app/litellm/proxy/prisma_migration.py /app/litellm/proxy/prisma_migration.py
|
||||
# enterprise/ is imported by source path at runtime (proxy_cli puts the
|
||||
# working directory on sys.path; litellm/proxy/hooks resolves
|
||||
# enterprise.enterprise_hooks from it)
|
||||
COPY --from=builder /app/enterprise /app/enterprise
|
||||
# Prisma binaries live in $HOME/.cache (default prisma-python location),
|
||||
# which is /root/.cache here. Copy them from the builder so they survive
|
||||
# deployments that volume-mount /app/.cache (e.g. readOnlyRootFilesystem
|
||||
# + emptyDir) — otherwise the mount would shadow the baked-in query engine.
|
||||
COPY --from=builder /root/.cache /root/.cache
|
||||
# Only the Prisma subdirs: the whole /root/.cache drags in the uv build cache.
|
||||
COPY --from=builder /root/.cache/prisma /root/.cache/prisma
|
||||
COPY --from=builder /root/.cache/prisma-python /root/.cache/prisma-python
|
||||
|
||||
RUN find /app/.venv -type f -path "*/tornado/test/*" -delete && \
|
||||
find /app/.venv -type d -path "*/tornado/test" -delete
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
# Base images
|
||||
ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:3258be472764337fd13095bcbb3182da170243b5819fd67ad4c0754590588b31
|
||||
ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:3258be472764337fd13095bcbb3182da170243b5819fd67ad4c0754590588b31
|
||||
ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:c61ac6919b811ea53c4782d69f1fe05218ba3c25d53f01b6ab7892e621bd4370
|
||||
ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:c61ac6919b811ea53c4782d69f1fe05218ba3c25d53f01b6ab7892e621bd4370
|
||||
ARG PROXY_EXTRAS_SOURCE=published
|
||||
ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a
|
||||
|
||||
|
|
@ -95,7 +95,21 @@ RUN for i in 1 2 3; do \
|
|||
apk add --no-cache python3 bash openssl tzdata libsndfile nodejs && break || sleep 5; \
|
||||
done
|
||||
|
||||
COPY --from=builder /app /app
|
||||
# Copy only what runtime needs. The application is installed inside the venv;
|
||||
# the rest of the builder's /app is source and build metadata that must not
|
||||
# ship (manifest-scanning tools attribute everything in it to this image).
|
||||
# entrypoint.sh invokes litellm/proxy/prisma_migration.py by source path.
|
||||
# Prisma caches live under /app/.cache here (XDG_CACHE_HOME /
|
||||
# PRISMA_BINARY_CACHE_DIR) so the runtime prisma generate finds them.
|
||||
COPY --from=builder /app/.venv /app/.venv
|
||||
COPY --from=builder /app/docker /app/docker
|
||||
COPY --from=builder /app/schema.prisma /app/schema.prisma
|
||||
COPY --from=builder /app/litellm/proxy/prisma_migration.py /app/litellm/proxy/prisma_migration.py
|
||||
# enterprise/ is imported by source path at runtime (proxy_cli puts the
|
||||
# working directory on sys.path; litellm/proxy/hooks resolves
|
||||
# enterprise.enterprise_hooks from it)
|
||||
COPY --from=builder /app/enterprise /app/enterprise
|
||||
COPY --from=builder /app/.cache /app/.cache
|
||||
COPY --from=builder /var/lib/litellm/ui /var/lib/litellm/ui
|
||||
COPY --from=builder /var/lib/litellm/assets /var/lib/litellm/assets
|
||||
|
||||
|
|
|
|||
|
|
@ -3,7 +3,7 @@ model_list:
|
|||
litellm_params:
|
||||
model: openai/fake
|
||||
api_key: fake-key
|
||||
api_base: https://exampleopenaiendpoint-production.up.railway.app/
|
||||
api_base: os.environ/FAKE_OPENAI_API_BASE
|
||||
|
||||
general_settings:
|
||||
alerting: ["slack"]
|
||||
141
docs/plugin_architecture.md
Normal file
141
docs/plugin_architecture.md
Normal file
|
|
@ -0,0 +1,141 @@
|
|||
# LiteLLM Plugin Architecture
|
||||
|
||||
Plugins let external services appear as selectable modes in the litellm UI sidebar alongside the AI Gateway.
|
||||
|
||||
---
|
||||
|
||||
## Quick start
|
||||
|
||||
### 1. Configure the plugin
|
||||
|
||||
Add a `plugins` block to your litellm `config.yaml`:
|
||||
|
||||
```yaml
|
||||
general_settings:
|
||||
master_key: sk-...
|
||||
plugins:
|
||||
- name: my-plugin # unique identifier (no spaces)
|
||||
display_name: My Plugin # shown in the UI dropdown
|
||||
url: "https://my-plugin.example.com"
|
||||
plugin_key: "sk-..." # plugin's own auth credential
|
||||
```
|
||||
|
||||
`plugin_key` is injected as `Authorization: Bearer <plugin_key>` on every
|
||||
request proxied through `/plugin-proxy/my-plugin/*`. The caller's litellm
|
||||
credential is stripped before forwarding so the plugin never receives a live
|
||||
litellm API key.
|
||||
|
||||
### 2. Implement two endpoints on your service
|
||||
|
||||
| Endpoint | Method | Purpose |
|
||||
|---|---|---|
|
||||
| `GET /api/plugin-manifest` | public | Returns plugin metadata for the UI |
|
||||
| `POST /api/plugin-auth` | public | Decrypts the identity claim for seamless sign-in |
|
||||
|
||||
#### `GET /api/plugin-manifest`
|
||||
|
||||
```json
|
||||
{
|
||||
"name": "my-plugin",
|
||||
"display_name": "My Plugin",
|
||||
"version": "1.0.0",
|
||||
"nav_items": [
|
||||
{ "key": "home", "label": "Home", "icon": "HomeOutlined", "path": "/" },
|
||||
{ "key": "reports", "label": "Reports", "icon": "BarChartOutlined", "path": "/reports" }
|
||||
],
|
||||
"capabilities": ["reports", "data"]
|
||||
}
|
||||
```
|
||||
|
||||
#### `POST /api/plugin-auth`
|
||||
|
||||
Receives `{ "session_claim": "<fernet-ciphertext>" }`.
|
||||
|
||||
The proxy never shares `LITELLM_SALT_KEY` with your plugin. Each plugin is
|
||||
provisioned with its own dedicated key, derived as
|
||||
`HMAC-SHA256(LITELLM_SALT_KEY, plugin_name)`. Compute it once on the proxy
|
||||
host and hand the result to your plugin as a secret (e.g. `PLUGIN_AUTH_KEY`):
|
||||
|
||||
```bash
|
||||
python -c 'import base64,hmac,hashlib,os; \
|
||||
print(base64.urlsafe_b64encode(hmac.new(os.environ["LITELLM_SALT_KEY"].encode(), b"my-plugin", hashlib.sha256).digest()).decode())'
|
||||
```
|
||||
|
||||
A compromised plugin holding only this scoped key cannot recover
|
||||
`LITELLM_SALT_KEY` or decrypt any other litellm secret.
|
||||
|
||||
Decrypt and validate the claim with that key:
|
||||
|
||||
```python
|
||||
import json, os, time
|
||||
from cryptography.fernet import Fernet
|
||||
|
||||
_CLAIM_TTL_SECONDS = 30
|
||||
|
||||
def plugin_auth(session_claim: str) -> dict:
|
||||
cipher = Fernet(os.environ["PLUGIN_AUTH_KEY"].encode())
|
||||
claim = json.loads(cipher.decrypt(session_claim.encode(), ttl=_CLAIM_TTL_SECONDS))
|
||||
if claim.get("plugin") != "my-plugin":
|
||||
raise ValueError("claim audience mismatch")
|
||||
if int(claim.get("exp", 0)) < int(time.time()):
|
||||
raise ValueError("claim expired")
|
||||
return claim
|
||||
```
|
||||
|
||||
The claim is `{ "plugin", "user_id", "user_role", "exp" }`; it carries no
|
||||
litellm bearer token. Establish the plugin's own session from `user_id` /
|
||||
`user_role` and authenticate API calls back to litellm through the
|
||||
`/plugin-proxy/my-plugin/*` reverse proxy, which injects `plugin_key` for you.
|
||||
|
||||
---
|
||||
|
||||
## How iframe auth works
|
||||
|
||||
```
|
||||
litellm UI
|
||||
├─ GET /api/plugins/auth-token -> { session_claim }
|
||||
└─ postMessage({ type:"litellm-auth", session_claim }, pluginOrigin)
|
||||
│
|
||||
▼
|
||||
Plugin iframe browser
|
||||
└─ POST /api/plugin-auth { session_claim }
|
||||
│
|
||||
▼
|
||||
Plugin server
|
||||
├─ decrypt(session_claim, PLUGIN_AUTH_KEY) -> { user_id, user_role, exp }
|
||||
└─ establish plugin session -> stored in sessionStorage
|
||||
```
|
||||
|
||||
No litellm bearer token ever leaves the proxy; the claim only conveys the
|
||||
caller's identity and expires after 30 seconds. A postMessage intercept
|
||||
yields ciphertext that is useless without the plugin's scoped key.
|
||||
|
||||
---
|
||||
|
||||
## Proxy routes
|
||||
|
||||
- `GET /api/plugins` — list registered plugins (`name`, `display_name`, `url`). `plugin_key` is **never** returned; it stays server-side. Requires an authenticated caller.
|
||||
- `GET /api/plugins/auth-token?plugin_name=<name>` — short-lived encrypted identity claim for the named plugin. Requires `LITELLM_SALT_KEY` to be set (503 otherwise) and the plugin to be registered (404 otherwise).
|
||||
- `ANY /plugin-proxy/{name}/{path}` — authenticated reverse proxy to the plugin backend. Restricted to `proxy_admin`.
|
||||
|
||||
---
|
||||
|
||||
## Reverse proxy behaviour
|
||||
|
||||
When an admin (or server-to-server caller) hits `/plugin-proxy/<name>/<path>`, the proxy authenticates the caller locally, then rewrites the request before forwarding it to the plugin's `url`:
|
||||
|
||||
- **Every litellm credential header is stripped** — `Authorization`, `x-api-key`, `API-Key`, `x-goog-api-key`, `Ocp-Apim-Subscription-Key`, `x-litellm-api-key`, any configured `litellm_key_header_name`, plus `Cookie`. The plugin can never be handed the caller's live litellm key.
|
||||
- **`plugin_key` is injected** as `Authorization: Bearer <plugin_key>` — the only credential the plugin receives.
|
||||
- **Caller identity is forwarded** as `x-litellm-user-id` and `x-litellm-user-role` so the plugin can run its own authorization. These are informational, not credentials.
|
||||
- **Responses are sandboxed** — `Content-Security-Policy: sandbox` and `X-Content-Type-Options: nosniff` are set so plugin-controlled bytes served from the litellm origin cannot execute against the dashboard.
|
||||
|
||||
---
|
||||
|
||||
## Security checklist
|
||||
|
||||
- [ ] `LITELLM_SALT_KEY` is set on the proxy and never shared with the plugin
|
||||
- [ ] The plugin holds only its derived `HMAC(LITELLM_SALT_KEY, plugin_name)` key, provisioned as a dedicated secret
|
||||
- [ ] `plugin_key` is a dedicated credential scoped to the plugin (not your litellm master key)
|
||||
- [ ] Plugin's `POST /api/plugin-auth` enforces the claim's `plugin` audience and `exp` (30s TTL)
|
||||
- [ ] Plugin treats `x-litellm-user-id` / `x-litellm-user-role` as identity hints, not as proof of authentication
|
||||
- [ ] Plugin service URL uses HTTPS in production
|
||||
|
|
@ -1,9 +1,9 @@
|
|||
# What is this?
|
||||
## This hook is used to check for LiteLLM managed files in the request body, and replace them with model-specific file id
|
||||
|
||||
import asyncio
|
||||
import base64
|
||||
import json
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Union, cast
|
||||
|
||||
from fastapi import HTTPException
|
||||
|
|
@ -412,7 +412,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
detail=f"User {user_api_key_dict.user_id} does not have access to the file {file_id}",
|
||||
)
|
||||
|
||||
async def async_pre_call_hook( # noqa: PLR0915
|
||||
async def async_pre_call_hook(
|
||||
self,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
cache: DualCache,
|
||||
|
|
@ -504,7 +504,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
if retrieve_file_id
|
||||
else False
|
||||
)
|
||||
if potential_file_id:
|
||||
if potential_file_id and "llm_output_file_id," in potential_file_id:
|
||||
model_id = self.get_model_id_from_unified_file_id(potential_file_id)
|
||||
if model_id:
|
||||
data["model"] = model_id
|
||||
|
|
@ -1058,7 +1058,12 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
return file_id.split("llm_output_file_model_id,")[1].split(";")[0]
|
||||
|
||||
def get_output_file_id_from_unified_file_id(self, file_id: str) -> str:
|
||||
return file_id.split("llm_output_file_id,")[1].split(";")[0]
|
||||
marker = "llm_output_file_id,"
|
||||
if marker not in file_id:
|
||||
raise ValueError(
|
||||
f"Unified id does not contain {marker!r}: {file_id[:80]!r}"
|
||||
)
|
||||
return file_id.split(marker, 1)[1].split(";")[0]
|
||||
|
||||
async def async_post_call_success_hook(
|
||||
self, data: Dict, user_api_key_dict: UserAPIKeyAuth, response: LLMResponseTypes
|
||||
|
|
@ -1099,13 +1104,33 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
for file_attr in ["output_file_id", "error_file_id"]:
|
||||
file_id_value = getattr(response, file_attr, None)
|
||||
if file_id_value and model_id:
|
||||
original_file_id = file_id_value
|
||||
unified_file_id = self.get_unified_output_file_id(
|
||||
output_file_id=original_file_id,
|
||||
model_id=model_id,
|
||||
model_name=resolved_model_name,
|
||||
decoded_output_file_id = _is_base64_encoded_unified_file_id(
|
||||
file_id_value
|
||||
)
|
||||
setattr(response, file_attr, unified_file_id)
|
||||
if (
|
||||
decoded_output_file_id
|
||||
and "llm_output_file_id," in decoded_output_file_id
|
||||
):
|
||||
provider_file_id = (
|
||||
self.get_output_file_id_from_unified_file_id(
|
||||
decoded_output_file_id
|
||||
)
|
||||
)
|
||||
unified_file_id = file_id_value
|
||||
elif decoded_output_file_id:
|
||||
verbose_logger.warning(
|
||||
f"Skipping {file_attr}={file_id_value!r}: "
|
||||
"unified id is not a managed file output id"
|
||||
)
|
||||
continue
|
||||
else:
|
||||
provider_file_id = file_id_value
|
||||
unified_file_id = self.get_unified_output_file_id(
|
||||
output_file_id=provider_file_id,
|
||||
model_id=model_id,
|
||||
model_name=resolved_model_name,
|
||||
)
|
||||
setattr(response, file_attr, unified_file_id)
|
||||
|
||||
# Use llm_router credentials when available. Without credentials,
|
||||
# Azure and other auth-required providers return 500/401.
|
||||
|
|
@ -1125,27 +1150,27 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
or {}
|
||||
)
|
||||
file_object = await litellm.afile_retrieve(
|
||||
file_id=original_file_id,
|
||||
file_id=provider_file_id,
|
||||
**_creds,
|
||||
)
|
||||
else:
|
||||
file_object = await litellm.afile_retrieve(
|
||||
custom_llm_provider=model_name.split("/")[0] if model_name and "/" in model_name else "openai", # type: ignore[arg-type]
|
||||
file_id=original_file_id,
|
||||
file_id=provider_file_id,
|
||||
)
|
||||
verbose_logger.debug(
|
||||
f"Successfully retrieved file object for {file_attr}={original_file_id}"
|
||||
f"Successfully retrieved file object for {file_attr}={provider_file_id}"
|
||||
)
|
||||
except Exception as e:
|
||||
verbose_logger.warning(
|
||||
f"Failed to retrieve file object for {file_attr}={original_file_id}: {str(e)}. Storing with None and will fetch on-demand."
|
||||
f"Failed to retrieve file object for {file_attr}={provider_file_id}: {str(e)}. Storing with None and will fetch on-demand."
|
||||
)
|
||||
|
||||
await self.store_unified_file_id(
|
||||
file_id=unified_file_id,
|
||||
file_object=file_object,
|
||||
litellm_parent_otel_span=user_api_key_dict.parent_otel_span,
|
||||
model_mappings={model_id: original_file_id},
|
||||
model_mappings={model_id: provider_file_id},
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
)
|
||||
await self.store_unified_object_id(
|
||||
|
|
@ -1447,8 +1472,8 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
error_message += f" (showing {MAX_BATCHES_IN_ERROR} most recent): {', '.join(batch_statuses)}. "
|
||||
|
||||
error_message += (
|
||||
f"To delete this file before complete cost tracking, please delete or cancel the referencing batch(es) first. "
|
||||
f"Alternatively, wait for all batches to complete and for cost to be computed (batch_processed=true)."
|
||||
"To delete this file before complete cost tracking, please delete or cancel the referencing batch(es) first. "
|
||||
"Alternatively, wait for all batches to complete and for cost to be computed (batch_processed=true)."
|
||||
)
|
||||
|
||||
# Record blocked deletion metric
|
||||
|
|
@ -1525,9 +1550,22 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
|
||||
if specific_model_file_id_mapping:
|
||||
exception_dict = {}
|
||||
for model_id, file_id in specific_model_file_id_mapping.items():
|
||||
for model_id, provider_file_id in specific_model_file_id_mapping.items():
|
||||
try:
|
||||
return await llm_router.afile_content(model=model_id, file_id=file_id, **data) # type: ignore
|
||||
# Cloud-storage providers (e.g. Bedrock S3) validate file ids
|
||||
# against the deployment's configured bucket, which they only
|
||||
# trust from this immutable server-side snapshot, never from
|
||||
# request params.
|
||||
credentials = llm_router.get_deployment_credentials_with_provider(
|
||||
model_id=model_id
|
||||
)
|
||||
if credentials is not None:
|
||||
data["_litellm_internal_model_credentials"] = cast(
|
||||
Dict, MappingProxyType(dict(credentials))
|
||||
)
|
||||
else:
|
||||
data.pop("_litellm_internal_model_credentials", None)
|
||||
return await llm_router.afile_content(model=model_id, file_id=provider_file_id, **data) # type: ignore
|
||||
except Exception as e:
|
||||
exception_dict[model_id] = str(e)
|
||||
raise Exception(
|
||||
|
|
|
|||
|
|
@ -483,7 +483,7 @@ async def new_project(
|
|||
response_model=LiteLLM_ProjectTable,
|
||||
)
|
||||
@management_endpoint_wrapper
|
||||
async def update_project( # noqa: PLR0915
|
||||
async def update_project(
|
||||
data: UpdateProjectRequest,
|
||||
http_request: Request,
|
||||
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
[project]
|
||||
name = "litellm-enterprise"
|
||||
version = "0.1.42"
|
||||
version = "0.1.43"
|
||||
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.42"
|
||||
version = "0.1.43"
|
||||
version_files = [
|
||||
"pyproject.toml:^version",
|
||||
"../pyproject.toml:litellm-enterprise==",
|
||||
|
|
|
|||
|
|
@ -1,5 +1,5 @@
|
|||
ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:31da6565f35af6401031c1d7aa91dc84ac76c5c48edd17fb90f0ed9e3173c7a9
|
||||
ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:31da6565f35af6401031c1d7aa91dc84ac76c5c48edd17fb90f0ed9e3173c7a9
|
||||
ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:c61ac6919b811ea53c4782d69f1fe05218ba3c25d53f01b6ab7892e621bd4370
|
||||
ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:c61ac6919b811ea53c4782d69f1fe05218ba3c25d53f01b6ab7892e621bd4370
|
||||
ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a
|
||||
|
||||
FROM $UV_IMAGE AS uvbin
|
||||
|
|
|
|||
10
litellm-rust/.cargo/config.toml
Normal file
10
litellm-rust/.cargo/config.toml
Normal file
|
|
@ -0,0 +1,10 @@
|
|||
# PyO3 cdylib (`litellm-python-bridge`) links against the host interpreter's
|
||||
# symbols, which are not present at link time when building an extension module.
|
||||
# On macOS, tell the linker to resolve undefined `_Py*` symbols dynamically at
|
||||
# load time (the standard pyo3 extension-module flag) so the cdylib links without
|
||||
# a libpython on the link line.
|
||||
[target.x86_64-apple-darwin]
|
||||
rustflags = ["-C", "link-arg=-undefined", "-C", "link-arg=dynamic_lookup"]
|
||||
|
||||
[target.aarch64-apple-darwin]
|
||||
rustflags = ["-C", "link-arg=-undefined", "-C", "link-arg=dynamic_lookup"]
|
||||
1
litellm-rust/.gitignore
vendored
Normal file
1
litellm-rust/.gitignore
vendored
Normal file
|
|
@ -0,0 +1 @@
|
|||
/target/
|
||||
9
litellm-rust/ADDING_A_PROVIDER.md
Normal file
9
litellm-rust/ADDING_A_PROVIDER.md
Normal file
|
|
@ -0,0 +1,9 @@
|
|||
# Adding a provider / route to litellm-rust
|
||||
|
||||
Three layers, same for every route (see `ocr` and `realtime` as references):
|
||||
|
||||
1. **Transform contract (pure)** — `crates/core/src/<route>/transformation.rs`: a `…ProviderConfig` trait (URL build + request/response transforms) + types in `types.rs`. No network, env, or auth.
|
||||
2. **Provider config (pure)** — `crates/providers/src/<provider>/<route>/transformation.rs`: implement that trait as a `const <PROVIDER>_<ROUTE>_CONFIG`, mirroring the Python provider tree. Add parity unit tests.
|
||||
3. **HTTP / transport (the host)** — `crates/providers/src/<route>.rs` (e.g. `ocr.rs`, `realtime.rs`): the callable fn (`run_ocr`, `realtime`). It resolves the key, builds the auth header, builds URL + transforms via the config, then does the network call. This is the only layer allowed to do I/O.
|
||||
|
||||
**Calling:** the host invokes the route fn — the Python bridge calls `run_ocr`; the `ai-gateway` server calls `realtime`. Register new modules in `lib.rs` / `mod.rs`, then run `cargo fmt && cargo clippy --workspace -- -D warnings && cargo test --workspace`.
|
||||
17
litellm-rust/AGENTS.md
Normal file
17
litellm-rust/AGENTS.md
Normal file
|
|
@ -0,0 +1,17 @@
|
|||
# AGENTS.md
|
||||
|
||||
litellm-rust has exactly THREE crates. A crate is a LAYER, not a route. Routes (ocr, realtime, chat) and providers (mistral, openai) are MODULES inside the layers.
|
||||
|
||||
## Crates
|
||||
|
||||
| Crate | Role | Pure / I/O |
|
||||
|-------|------|------------|
|
||||
| litellm-core | Translation layer — types, route contracts (traits), provider transforms (modules under providers/), and the router. Builds requests/responses; no network. | Pure |
|
||||
| litellm-ai-gateway | Routes + host — the only crate that touches the network. HTTP/WebSocket I/O (modules under io/) plus the axum server binary (behind the `server` feature). | I/O |
|
||||
| litellm-python-bridge | PyO3 cdylib exposing Rust to the litellm Python SDK — a thin adapter over litellm-ai-gateway's I/O. | Binding |
|
||||
|
||||
Dependency direction (acyclic): litellm-core ← litellm-ai-gateway ← litellm-python-bridge.
|
||||
|
||||
Adding a crate: default to a MODULE. New crate ONLY on a real trigger — separate artifact (binary/cdylib), proc-macro, shared foundation, or publishable standalone. A new provider or route is none of these.
|
||||
|
||||
Adding a crate fails crates/core/tests/workspace_crate_allowlist.rs until you update its allowlist and this file — intentional.
|
||||
95
litellm-rust/CLAUDE.md
Normal file
95
litellm-rust/CLAUDE.md
Normal file
|
|
@ -0,0 +1,95 @@
|
|||
# CLAUDE.md
|
||||
|
||||
This file defines the rules for Rust work in LiteLLM.
|
||||
|
||||
## Crates (exactly three — see AGENTS.md)
|
||||
|
||||
`litellm-core` describes work; `litellm-ai-gateway` executes it; `litellm-python-bridge`
|
||||
exposes it to the Python SDK. A crate is a **layer**, not a route — add modules, not crates.
|
||||
|
||||
## Core Boundary
|
||||
|
||||
`litellm-core` is the pure translation layer; the `litellm-ai-gateway` host executes work.
|
||||
|
||||
Route-level Rust structure mirrors LiteLLM's Python responsibilities:
|
||||
- `core/src/<route>/` owns the route contract, shared types, and provider
|
||||
template traits. For OCR, this means `core/src/ocr`.
|
||||
- `core/src/providers/<provider>/<route>/transformation.rs` owns the
|
||||
provider-specific transform. For Mistral OCR, this means
|
||||
`core/src/providers/mistral/ocr/transformation.rs`.
|
||||
- Network execution lives in the host crate `ai-gateway` (`ai-gateway/src/io/`),
|
||||
never inside `core`.
|
||||
|
||||
Allowed in `core`:
|
||||
- Pure request transforms
|
||||
- Pure response transforms
|
||||
- Pure stream chunk normalization
|
||||
- Shared data types and validation errors
|
||||
- Deterministic token/cost helper logic
|
||||
|
||||
Not allowed in `core`:
|
||||
- Network calls
|
||||
- Environment variable or secret reads
|
||||
- Filesystem access
|
||||
- Database or cache access
|
||||
- Provider SDK signing or auth flows
|
||||
- Logging callbacks, spend writes, or custom callbacks
|
||||
- Global mutable runtime state
|
||||
|
||||
Python owns rollout state and fallback while Rust is being introduced. Rust
|
||||
paths must be off by default until parity tests prove equivalence with Python.
|
||||
|
||||
## Production Bar
|
||||
|
||||
Rust code in this workspace is held to a strict parity and robustness bar from
|
||||
the first PR:
|
||||
|
||||
- Correctness parity is proven with tests. Do not rely on README claims or
|
||||
manual inspection for a port that mirrors Python behavior.
|
||||
- Every provider transform must have unit tests for supported-parameter
|
||||
filtering, request body shape, response normalization, missing/null fields,
|
||||
and bad-input errors.
|
||||
- When Rust is exposed through Python, add Python tests that prove disabled,
|
||||
enabled, and unavailable-bridge fallback behavior.
|
||||
- Avoid panics on user/provider input. Return typed errors and let the host map
|
||||
them to Python exceptions or HTTP responses.
|
||||
- OCR handles documents that often contain personal data. Do not log document
|
||||
contents, base64 payloads, provider response bodies, or secrets.
|
||||
- Error messages must be useful but data-minimized. Truncate or sanitize any
|
||||
upstream body before it crosses a host boundary.
|
||||
- Treat empty or whitespace-only credentials, URLs, and config values as absent
|
||||
at the host/config resolution layer.
|
||||
- Preserve Python output shape intentionally. If a field is always serialized as
|
||||
`null` for Python parity, leave a short comment explaining that parity choice.
|
||||
|
||||
## Host I/O Rules
|
||||
|
||||
These rules apply when adding future crates or modules that execute network I/O,
|
||||
such as `ai-gateway`, router hosts, or standalone servers:
|
||||
|
||||
- Set connect and full-request timeouts. No unbounded waits.
|
||||
- Reuse HTTP clients; do not construct clients per request.
|
||||
- Prefer rustls TLS for portable Python wheels and Linux images unless there is
|
||||
a documented reason not to.
|
||||
- Add request IDs and structured tracing at the host layer, without logging OCR
|
||||
document contents or secrets.
|
||||
- Do not echo raw upstream response bodies to callers. Sanitize and bound them.
|
||||
- Avoid `expect`/`unwrap` in server startup and request paths unless the panic is
|
||||
impossible by construction and documented.
|
||||
|
||||
## Checks
|
||||
|
||||
Run these before pushing Rust changes. The same checks run in GitHub Actions
|
||||
for changes under `litellm-rust/`.
|
||||
|
||||
```bash
|
||||
cd litellm-rust
|
||||
cargo fmt --check
|
||||
# the ai-gateway binary + server code is behind the `server` feature
|
||||
cargo clippy -p litellm-ai-gateway --all-targets --features server -- -D warnings
|
||||
cargo clippy -p litellm-core -p litellm-python-bridge --all-targets -- -D warnings
|
||||
cargo test --workspace
|
||||
```
|
||||
|
||||
When a Rust path is exposed through Python, add Python parity tests that compare
|
||||
the existing Python output with the Rust-backed output.
|
||||
1861
litellm-rust/Cargo.lock
generated
Normal file
1861
litellm-rust/Cargo.lock
generated
Normal file
File diff suppressed because it is too large
Load diff
27
litellm-rust/Cargo.toml
Normal file
27
litellm-rust/Cargo.toml
Normal file
|
|
@ -0,0 +1,27 @@
|
|||
[workspace]
|
||||
members = [
|
||||
"crates/core",
|
||||
"crates/ai-gateway",
|
||||
"crates/python-bridge",
|
||||
]
|
||||
resolver = "2"
|
||||
|
||||
[workspace.package]
|
||||
edition = "2021"
|
||||
license = "MIT"
|
||||
repository = "https://github.com/BerriAI/litellm"
|
||||
|
||||
[workspace.dependencies]
|
||||
litellm-core = { path = "crates/core" }
|
||||
litellm-ai-gateway = { path = "crates/ai-gateway", default-features = false }
|
||||
axum = "0.7"
|
||||
pyo3 = "0.23.5"
|
||||
rand = "0.8"
|
||||
reqwest = { version = "0.12", default-features = false, features = ["blocking", "json", "rustls-tls"] }
|
||||
serde = { version = "1.0", features = ["derive"] }
|
||||
serde_json = "1.0"
|
||||
subtle = "2"
|
||||
thiserror = "2.0"
|
||||
tokio = { version = "1", features = ["rt-multi-thread", "macros", "time"] }
|
||||
tokio-tungstenite = { version = "0.24", default-features = false, features = ["connect", "rustls-tls-native-roots"] }
|
||||
futures-util = { version = "0.3", default-features = false, features = ["sink", "std"] }
|
||||
44
litellm-rust/README.md
Normal file
44
litellm-rust/README.md
Normal file
|
|
@ -0,0 +1,44 @@
|
|||
# LiteLLM Rust
|
||||
|
||||
This workspace contains the staged Rust implementation for LiteLLM.
|
||||
|
||||
Rust starts as a pure transform core used by the existing Python host. Python
|
||||
continues to own auth, configuration, network I/O, retries, routing, logging,
|
||||
callbacks, spend tracking, and customer plugins until each Rust path has parity
|
||||
coverage and production evidence.
|
||||
|
||||
## Crates
|
||||
|
||||
| Crate | Role | Pure / I/O |
|
||||
|-------|------|------------|
|
||||
| litellm-core | Translation layer — types, route contracts (traits), provider transforms (modules under providers/), and the router. Builds requests/responses; no network. | Pure |
|
||||
| litellm-ai-gateway | Routes + host — the only crate that touches the network. HTTP/WebSocket I/O (modules under io/) plus the axum server binary (behind the `server` feature). | I/O |
|
||||
| litellm-python-bridge | PyO3 cdylib exposing Rust to the litellm Python SDK — a thin adapter over litellm-ai-gateway's I/O. | Binding |
|
||||
|
||||
Dependency direction (acyclic): litellm-core ← litellm-ai-gateway ← litellm-python-bridge.
|
||||
|
||||
## Layout
|
||||
|
||||
```text
|
||||
crates/
|
||||
core/ Route contracts, shared pure types, errors, and templates.
|
||||
src/ocr/
|
||||
providers/ Provider-specific pure transforms.
|
||||
src/mistral/ocr/transformation.rs
|
||||
python-bridge/ PyO3 bridge for Python LiteLLM.
|
||||
```
|
||||
|
||||
The folder shape should follow the Python provider tree:
|
||||
`providers/src/<provider>/<route>/transformation.rs`. The bridge should expose
|
||||
one function per top-level route, starting with `ocr(payload)`.
|
||||
|
||||
## Checks
|
||||
|
||||
Run these before pushing Rust changes. GitHub Actions runs the same checks for
|
||||
changes under `litellm-rust/`.
|
||||
|
||||
```bash
|
||||
cargo fmt --check
|
||||
cargo clippy --workspace --all-targets -- -D warnings
|
||||
cargo test --workspace
|
||||
```
|
||||
50
litellm-rust/crates/ai-gateway/AGENTS.md
Normal file
50
litellm-rust/crates/ai-gateway/AGENTS.md
Normal file
|
|
@ -0,0 +1,50 @@
|
|||
# ai-gateway — folder architecture
|
||||
|
||||
The Axum server that fronts the Rust gateway. It owns transport + config + auth
|
||||
only; deployment selection lives in `core::router`, transforms in `core`/`providers`.
|
||||
|
||||
```
|
||||
src/
|
||||
main.rs # entrypoint: build AppState (router + master key), bind, serve
|
||||
state.rs # AppState — shared Arc<Router> + master_key
|
||||
gil.rs # GIL-activity tracker (records Python acquisitions)
|
||||
auth/ # authentication as an axum extractor — added to handler args
|
||||
mod.rs # RequireMasterKey: FromRequestParts, single master key (LITELLM_MASTER_KEY)
|
||||
routes/ # one module per route, all matching the same template
|
||||
AGENTS.md # ← the route template (read this before adding a route)
|
||||
mod.rs # app(): merges every module's router()
|
||||
health.rs # simple route (one file): router() + liveness/readiness
|
||||
gil.rs # simple route (one file): router() + GET /health/gil
|
||||
realtime/ # route with logic → axum surface + a no-axum service:
|
||||
mod.rs # router() + handler + WS<->events adapter (the axum surface)
|
||||
service.rs # business logic (select deployment, call provider) — no axum, testable
|
||||
python/ # Python interop (feature: python-config) — load-time only
|
||||
mod.rs, config.rs, AGENTS.md
|
||||
```
|
||||
|
||||
## Rules
|
||||
|
||||
- **Routes follow one template.** Each route module exposes
|
||||
`pub fn router() -> Router<AppState>`; `routes/mod.rs` only merges them. Simple
|
||||
routes are one file; non-trivial routes are a folder (`handler`/`service`/
|
||||
`transport`). See `routes/AGENTS.md`.
|
||||
- **Auth is an extractor.** Add `crate::auth::RequireMasterKey` to a handler's
|
||||
args; it runs during extraction. Never re-implement the check per route.
|
||||
- **Handlers are thin.** A handler validates and delegates to its `service`. No
|
||||
business logic, no provider calls, no transforms in handlers.
|
||||
- **State is shared and cheap to clone.** Long-lived handles live behind `Arc` in
|
||||
`state.rs`; read env/config only in `main.rs` when building state.
|
||||
|
||||
## Auth (interim)
|
||||
|
||||
A single **master key** (`LITELLM_MASTER_KEY`), enforced by the
|
||||
`auth::RequireMasterKey` extractor: any caller presenting it as
|
||||
`Authorization: Bearer <key>` may invoke the gateway. Fails closed (500) when
|
||||
unset; constant-time compare. The server binds `127.0.0.1` by default (`HOST` to
|
||||
override). Full per-key auth + budgets/rate-limits are delegated to the Python
|
||||
proxy in a later phase. Health routes don't add the extractor (unauthenticated).
|
||||
|
||||
## Python interop
|
||||
|
||||
Anything that calls into Python lives in `python/` and is **load-time only** — see
|
||||
`python/AGENTS.md`. The realtime data path never takes the GIL.
|
||||
36
litellm-rust/crates/ai-gateway/Cargo.toml
Normal file
36
litellm-rust/crates/ai-gateway/Cargo.toml
Normal file
|
|
@ -0,0 +1,36 @@
|
|||
[package]
|
||||
name = "litellm-ai-gateway"
|
||||
version = "0.1.0"
|
||||
edition.workspace = true
|
||||
license.workspace = true
|
||||
repository.workspace = true
|
||||
|
||||
[lib]
|
||||
name = "litellm_ai_gateway"
|
||||
|
||||
[[bin]]
|
||||
name = "litellm-ai-gateway"
|
||||
path = "src/main.rs"
|
||||
required-features = ["server"]
|
||||
|
||||
[dependencies]
|
||||
litellm-core.workspace = true
|
||||
reqwest.workspace = true
|
||||
tokio = { workspace = true, features = ["rt-multi-thread", "macros", "net", "time"] }
|
||||
tokio-tungstenite.workspace = true
|
||||
futures-util.workspace = true
|
||||
serde_json.workspace = true
|
||||
axum = { workspace = true, features = ["ws"], optional = true }
|
||||
serde = { workspace = true, optional = true }
|
||||
subtle = { workspace = true, optional = true }
|
||||
pyo3 = { workspace = true, features = ["auto-initialize"], optional = true }
|
||||
|
||||
[features]
|
||||
default = []
|
||||
server = ["dep:axum", "dep:subtle", "dep:serde"]
|
||||
# Build the gateway's config from the proxy YAML via an embedded Python
|
||||
# interpreter (links libpython; requires `litellm` importable at runtime).
|
||||
python-config = ["dep:pyo3"]
|
||||
|
||||
[dev-dependencies]
|
||||
futures-channel = "0.3"
|
||||
86
litellm-rust/crates/ai-gateway/Dockerfile
Normal file
86
litellm-rust/crates/ai-gateway/Dockerfile
Normal file
|
|
@ -0,0 +1,86 @@
|
|||
# Multi-stage build for the LiteLLM Rust AI Gateway (realtime WebSocket proxy).
|
||||
#
|
||||
# Build context is the **repo root** so we can install `litellm` from this repo's
|
||||
# source (the gateway loads its model_list via litellm.proxy.read_model_list,
|
||||
# which is not in any PyPI release yet) AND build the rust workspace under
|
||||
# litellm-rust/.
|
||||
#
|
||||
# docker build -f litellm-rust/crates/ai-gateway/Dockerfile -t litellm-ai-gateway .
|
||||
#
|
||||
# No secrets live in this file. Runtime config (LITELLM_MASTER_KEY,
|
||||
# OPENAI_API_KEY referenced by config.yaml, etc.) is injected as environment
|
||||
# variables at deploy time.
|
||||
|
||||
# ---- Chef -------------------------------------------------------------------
|
||||
# cargo-chef caches the dependency build so only the gateway crate recompiles on
|
||||
# a source-only change. python3-dev is present in every rust stage because the
|
||||
# `python-config` feature links libpython via pyo3 (even in the cook step).
|
||||
FROM rust:1.90-slim-bookworm AS chef
|
||||
ENV PYO3_PYTHON=python3.11
|
||||
RUN apt-get update \
|
||||
&& apt-get install -y --no-install-recommends \
|
||||
python3 python3-dev pkg-config libssl-dev clang \
|
||||
&& rm -rf /var/lib/apt/lists/* \
|
||||
&& cargo install cargo-chef --locked --version 0.1.77
|
||||
WORKDIR /build/litellm-rust
|
||||
|
||||
# ---- Planner ----------------------------------------------------------------
|
||||
# Produce the dependency recipe from the rust workspace manifests + Cargo.lock.
|
||||
FROM chef AS planner
|
||||
COPY litellm-rust/ .
|
||||
RUN cargo chef prepare --recipe-path recipe.json
|
||||
|
||||
# ---- Builder ----------------------------------------------------------------
|
||||
FROM chef AS builder
|
||||
# Cook (compile) just the dependencies first — this layer is cached and reused
|
||||
# whenever only gateway source changes.
|
||||
COPY --from=planner /build/litellm-rust/recipe.json recipe.json
|
||||
RUN cargo chef cook --locked --release \
|
||||
-p litellm-ai-gateway --features python-config \
|
||||
--recipe-path recipe.json
|
||||
# Now copy the real sources and build the gateway binary. Deps are already cooked
|
||||
# above, so this step only recompiles the gateway crate.
|
||||
COPY litellm-rust/ .
|
||||
RUN cargo build --locked --release -p litellm-ai-gateway --features python-config
|
||||
|
||||
# ---- Runtime ----------------------------------------------------------------
|
||||
# python:3.11-slim-bookworm ships libpython3.11, matching the builder's PyO3
|
||||
# 3.11 ABI so the embedded interpreter links and imports cleanly.
|
||||
FROM python:3.11-slim-bookworm AS runtime
|
||||
|
||||
# CA certificates for outbound TLS to the OpenAI realtime endpoint.
|
||||
RUN apt-get update \
|
||||
&& apt-get install -y --no-install-recommends ca-certificates \
|
||||
&& rm -rf /var/lib/apt/lists/*
|
||||
|
||||
WORKDIR /app
|
||||
|
||||
# Install litellm (with proxy extras) FROM THIS REPO'S SOURCE so
|
||||
# `import litellm.proxy.read_model_list` works — it is not on PyPI yet. Copy the
|
||||
# package + packaging metadata, then pip install the proxy extra.
|
||||
COPY pyproject.toml README.md LICENSE ./
|
||||
COPY litellm/ ./litellm/
|
||||
RUN pip install --no-cache-dir ".[proxy]"
|
||||
|
||||
# The compiled gateway binary (pure-Rust realtime hot path; Python is load-time
|
||||
# only).
|
||||
COPY --from=builder /build/litellm-rust/target/release/litellm-ai-gateway /usr/local/bin/litellm-ai-gateway
|
||||
|
||||
# Default config.yaml. A real deploy can override this (e.g. mount a Render
|
||||
# secret file at the same path) — never bake secrets into the image.
|
||||
COPY litellm-rust/crates/ai-gateway/config.yaml /app/config.yaml
|
||||
|
||||
# Bind to all interfaces (Render routes to 0.0.0.0:$PORT) and load the model_list
|
||||
# from config.yaml via the embedded python config reader.
|
||||
ENV HOST=0.0.0.0 \
|
||||
LITELLM_CONFIG_PATH=/app/config.yaml
|
||||
|
||||
# Drop to a non-root user. The realtime hot path needs no root privileges, so
|
||||
# running unprivileged limits blast radius if the process is ever compromised.
|
||||
# The binary in /usr/local/bin is world-executable (COPY default mode 755); we
|
||||
# only need /app (and the config.yaml it reads) owned by the unprivileged user.
|
||||
RUN useradd --system --no-create-home --uid 10001 appuser \
|
||||
&& chown -R appuser:appuser /app
|
||||
USER appuser
|
||||
|
||||
ENTRYPOINT ["/usr/local/bin/litellm-ai-gateway"]
|
||||
45
litellm-rust/crates/ai-gateway/Dockerfile.dockerignore
Normal file
45
litellm-rust/crates/ai-gateway/Dockerfile.dockerignore
Normal file
|
|
@ -0,0 +1,45 @@
|
|||
# Dockerfile-specific ignore-file for the Rust AI Gateway build.
|
||||
#
|
||||
# The build context is the repo root (so the image can pip install litellm from
|
||||
# source AND build the rust workspace). BuildKit honors `<Dockerfile>.dockerignore`
|
||||
# next to the Dockerfile and it takes precedence over the repo-root `.dockerignore`,
|
||||
# so this file shrinks the (large) repo-root context for THIS build only without
|
||||
# touching the root `.dockerignore` used by the main litellm images.
|
||||
#
|
||||
# Strategy: ignore everything, then re-include only what the build needs:
|
||||
# - litellm/ (pip install . needs the full package + proxy reader)
|
||||
# - litellm-rust/ (the rust workspace; Cargo.lock + crate sources)
|
||||
# - pyproject.toml / README.md / LICENSE (packaging metadata for pip install)
|
||||
*
|
||||
|
||||
# --- re-include the build inputs ---
|
||||
!litellm/
|
||||
!litellm-rust/
|
||||
!pyproject.toml
|
||||
!README.md
|
||||
!LICENSE
|
||||
|
||||
# --- prune heavy / irrelevant subpaths back out of the re-included trees ---
|
||||
# Rust build artifacts (huge; regenerated in the builder).
|
||||
**/target/
|
||||
# Python caches and compiled bytecode.
|
||||
**/__pycache__/
|
||||
**/*.pyc
|
||||
**/*.pyo
|
||||
**/.pytest_cache/
|
||||
**/.ruff_cache/
|
||||
**/.mypy_cache/
|
||||
# Node / UI build output bundled under the python package (not needed to import
|
||||
# litellm.proxy.read_model_list).
|
||||
**/node_modules/
|
||||
litellm/proxy/_experimental/out/
|
||||
# Tests, logs, and local scratch.
|
||||
**/tests/
|
||||
**/test/
|
||||
*.log
|
||||
log.txt
|
||||
*.tgz
|
||||
# VCS / editor / CI metadata that may live under re-included trees.
|
||||
**/.git/
|
||||
.git/
|
||||
**/.DS_Store
|
||||
184
litellm-rust/crates/ai-gateway/README.md
Normal file
184
litellm-rust/crates/ai-gateway/README.md
Normal file
|
|
@ -0,0 +1,184 @@
|
|||
# LiteLLM Rust AI Gateway
|
||||
|
||||
A minimal Axum service that fronts OpenAI's realtime API. Clients open a
|
||||
WebSocket to `GET /v1/realtime`; the gateway authenticates, selects a deployment,
|
||||
dials OpenAI upstream, and splices the two sockets frame-by-frame.
|
||||
|
||||
## Crates
|
||||
|
||||
`litellm-rust` is exactly three crates (a crate is a **layer**, not a route):
|
||||
|
||||
| Crate | Role | Pure / I/O |
|
||||
|-------|------|------------|
|
||||
| litellm-core | Translation layer — types, route contracts (traits), provider transforms (modules under `providers/`), and the router. Builds requests/responses; no network. | Pure |
|
||||
| litellm-ai-gateway | Routes + host — the only crate that touches the network. HTTP/WebSocket I/O (modules under `io/`) plus the Axum server binary (behind the `server` feature). | I/O |
|
||||
| litellm-python-bridge | PyO3 cdylib exposing Rust to the litellm Python SDK — a thin adapter over litellm-ai-gateway's I/O. | Binding |
|
||||
|
||||
Dependency direction (acyclic): litellm-core ← litellm-ai-gateway ← litellm-python-bridge.
|
||||
|
||||
- **Client endpoint:** `wss://<host>/v1/realtime?model=<model>` (WebSocket)
|
||||
- **Auth:** `Authorization: Bearer $LITELLM_MASTER_KEY` (fails closed if unset)
|
||||
- **Health:** `GET /health/readiness`, `GET /health/liveness`, `GET /health/gil`
|
||||
|
||||
> **Realtime serving is pure Rust.** Python is used at **load time only** — to
|
||||
> read the config once at boot. The realtime hot path never touches Python.
|
||||
|
||||
## Configuration (config.yaml)
|
||||
|
||||
The gateway loads its `model_list` from a **config.yaml**, the same as the
|
||||
LiteLLM proxy. Point `LITELLM_CONFIG_PATH` at the file:
|
||||
|
||||
```yaml
|
||||
# config.yaml
|
||||
model_list:
|
||||
- model_name: gpt-realtime
|
||||
litellm_params:
|
||||
model: openai/gpt-realtime
|
||||
api_key: os.environ/OPENAI_API_KEY
|
||||
```
|
||||
|
||||
```bash
|
||||
LITELLM_CONFIG_PATH=./config.yaml ./litellm-ai-gateway
|
||||
```
|
||||
|
||||
At boot the gateway calls into `litellm.proxy.read_model_list`, which reuses the
|
||||
**real proxy config reader** (`ProxyConfig.get_config`). That means everything
|
||||
the proxy supports in config.yaml works here too:
|
||||
|
||||
- `include:` to merge in other config files,
|
||||
- `os.environ/VAR` secret references (resolved via the secret manager, never
|
||||
inlined),
|
||||
- DB-stored models (when a database is configured).
|
||||
|
||||
Secrets stay out of the config — reference them with `os.environ/...` and set
|
||||
the env var at deploy time. The shipped Docker image is built with the
|
||||
`python-config` feature and **bundles litellm**, so config loading works out of
|
||||
the box; the default baked config lives at `/app/config.yaml` and can be
|
||||
overridden at deploy time (e.g. a Render secret file mounted at the same path).
|
||||
|
||||
### Environment variables
|
||||
|
||||
| Var | Required | Default | Purpose |
|
||||
|---|---|---|---|
|
||||
| `LITELLM_CONFIG_PATH` | yes (config mode) | — | Path to the config.yaml the gateway loads its `model_list` from. The Docker image defaults this to `/app/config.yaml`. |
|
||||
| `LITELLM_MASTER_KEY` | yes | — | Bearer token clients must send. Unset ⇒ all `/v1/realtime` requests are rejected (fail closed). |
|
||||
| `OPENAI_API_KEY` | yes | — | Upstream OpenAI key. Referenced by config.yaml as `os.environ/OPENAI_API_KEY` for the gateway→OpenAI dial. |
|
||||
| `HOST` | no | `127.0.0.1` | **Set to `0.0.0.0` in any container/deploy** or external traffic is refused. |
|
||||
| `PORT` | no | `4001` | Listen port. Render and most PaaS inject this automatically. |
|
||||
|
||||
> Secrets (`LITELLM_MASTER_KEY`, `OPENAI_API_KEY`) are never baked into the image
|
||||
> or `render.yaml` — inject them at deploy time only.
|
||||
|
||||
### Lean env stand-in (fallback)
|
||||
|
||||
If the binary is built **without** `python-config` (default features), or
|
||||
`LITELLM_CONFIG_PATH` is unset, the gateway falls back to a single-deployment
|
||||
stand-in built from the environment:
|
||||
|
||||
| Var | Default | Purpose |
|
||||
|---|---|---|
|
||||
| `OPENAI_REALTIME_MODEL` | `gpt-realtime` | The single deployment's model name (also the `?model=` clients pass). |
|
||||
|
||||
This mode links no libpython and needs no config file, but it only supports one
|
||||
hard-coded OpenAI deployment. **config.yaml is the recommended path** — use the
|
||||
stand-in only for the leanest possible build.
|
||||
|
||||
## Build & run with Docker
|
||||
|
||||
The image is built `--features python-config` and installs litellm **from this
|
||||
repo's source** (the config reader is newer than any PyPI release), so the build
|
||||
**context is the repo root**:
|
||||
|
||||
```bash
|
||||
# from the repo root
|
||||
docker build -f litellm-rust/crates/ai-gateway/Dockerfile -t litellm-ai-gateway .
|
||||
|
||||
docker run --rm -p 4001:4001 \
|
||||
-e HOST=0.0.0.0 -e PORT=4001 \
|
||||
-e LITELLM_MASTER_KEY=sk-local \
|
||||
-e OPENAI_API_KEY=$OPENAI_API_KEY \
|
||||
litellm-ai-gateway # LITELLM_CONFIG_PATH defaults to /app/config.yaml
|
||||
|
||||
# smoke test
|
||||
curl -s -o /dev/null -w '%{http_code}\n' localhost:4001/health/readiness # -> 200
|
||||
curl -s -o /dev/null -w '%{http_code}\n' localhost:4001/v1/realtime # -> 401 (auth fails closed)
|
||||
```
|
||||
|
||||
On boot you should see `loaded model_list from /app/config.yaml via python
|
||||
config reader` — that confirms the config path (not the env stand-in fallback).
|
||||
To use your own config, mount it over the default:
|
||||
|
||||
```bash
|
||||
docker run --rm -p 4001:4001 \
|
||||
-e HOST=0.0.0.0 -e LITELLM_MASTER_KEY=sk-local -e OPENAI_API_KEY=$OPENAI_API_KEY \
|
||||
-v $(pwd)/my-config.yaml:/app/config.yaml:ro \
|
||||
litellm-ai-gateway
|
||||
```
|
||||
|
||||
### Cargo-only (no Docker)
|
||||
|
||||
```bash
|
||||
# config.yaml mode — needs litellm importable in the active python env
|
||||
LITELLM_CONFIG_PATH=./crates/ai-gateway/config.yaml \
|
||||
cargo run --release -p litellm-ai-gateway --features python-config
|
||||
|
||||
# env stand-in mode — no python, no config
|
||||
cargo run --release -p litellm-ai-gateway
|
||||
```
|
||||
|
||||
## Deploy on Render
|
||||
|
||||
The service is a Docker **web service**; Render terminates TLS and supports
|
||||
WebSockets, so the public endpoint is `wss://<service>.onrender.com/v1/realtime`.
|
||||
|
||||
### Option A — Blueprint (`render.yaml`)
|
||||
|
||||
`crates/ai-gateway/render.yaml` describes the service (Docker runtime,
|
||||
`healthCheckPath: /health/readiness`, repo-root `dockerContext: .`,
|
||||
`dockerfilePath: ./litellm-rust/crates/ai-gateway/Dockerfile`,
|
||||
`LITELLM_CONFIG_PATH: /app/config.yaml`). `LITELLM_MASTER_KEY` and
|
||||
`OPENAI_API_KEY` are `sync: false` — set them in the dashboard after the first
|
||||
deploy. To use a non-default model_list, mount a **Render Secret File** at
|
||||
`/app/config.yaml`. Point a Render Blueprint at this repo/branch and apply.
|
||||
|
||||
### Option B — Render API
|
||||
|
||||
```bash
|
||||
# create a Docker web service from this repo+branch, then set env vars:
|
||||
curl -X POST https://api.render.com/v1/services \
|
||||
-H "Authorization: Bearer $RENDER_API_KEY" -H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"type": "web_service", "name": "litellm-rust-ai-gateway",
|
||||
"ownerId": "<owner-id>", "repo": "https://github.com/BerriAI/litellm",
|
||||
"branch": "<branch-with-this-dockerfile>",
|
||||
"serviceDetails": {
|
||||
"env": "docker",
|
||||
"envSpecificDetails": {
|
||||
"dockerfilePath": "./litellm-rust/crates/ai-gateway/Dockerfile",
|
||||
"dockerContext": "."
|
||||
},
|
||||
"healthCheckPath": "/health/readiness"
|
||||
}
|
||||
}'
|
||||
# then set env vars LITELLM_MASTER_KEY, OPENAI_API_KEY, HOST=0.0.0.0,
|
||||
# LITELLM_CONFIG_PATH=/app/config.yaml
|
||||
```
|
||||
|
||||
Health check path **must** be `/health/readiness`. `autoDeploy` is off by default
|
||||
in the blueprint — trigger deploys manually (or flip it on) to pick up new commits.
|
||||
|
||||
## Scaling
|
||||
|
||||
Concurrency is what matters, not total connections: each in-flight session holds
|
||||
one client socket + one upstream socket. To scale, raise the instance count /
|
||||
enable autoscaling on the Render service (e.g. baseline 10, max 100). Each
|
||||
instance needs file descriptors for `2 × peak_concurrent_sessions` — raise
|
||||
`ulimit -n` if you push very high concurrency.
|
||||
|
||||
## Latency note
|
||||
|
||||
The gateway adds the cost of one extra hop: client→gateway, then a fresh
|
||||
gateway→OpenAI realtime handshake (TLS + WS upgrade + `session.created`). In
|
||||
benchmarks this is ~100–150 ms of added session-establishment time; first-audio
|
||||
and steady-state streaming add no measurable overhead. To minimize it, deploy the
|
||||
gateway in the Render region with the lowest RTT to OpenAI's realtime endpoint.
|
||||
55
litellm-rust/crates/ai-gateway/benchmarks/realtime/README.md
Normal file
55
litellm-rust/crates/ai-gateway/benchmarks/realtime/README.md
Normal file
|
|
@ -0,0 +1,55 @@
|
|||
# Realtime gateway benchmark — pool on/off
|
||||
|
||||
Measures what the gateway adds over talking to OpenAI's realtime WebSocket
|
||||
directly, and what the pre-warmed connection pool removes. See
|
||||
`../../src/routes/realtime/README.md` for how the pool works.
|
||||
|
||||
## Results
|
||||
|
||||
5000 calls / 500 concurrency, gateway at 10 instances, pool ON
|
||||
(`REALTIME_POOL_SIZE=64`), upstream OpenAI `gpt-realtime`. Each leg run twice.
|
||||
Times in **ms**. Phases per connection: **dial** = TCP+TLS+WS upgrade,
|
||||
**session** = upgrade → `session.created` (the phase the pool removes),
|
||||
**1st-audio** = `response.create` → first audio delta (OpenAI inference),
|
||||
**total** = full wall-clock.
|
||||
|
||||
| metric | Direct OpenAI | Gateway (pool ON) | Overhead (ms) | vs OpenAI |
|
||||
| ------------------ | ------------- | ----------------- | ------------- | ---------- |
|
||||
| success rate (%) | 99.8 | 99.8 | — | — |
|
||||
| dial p50 (ms) | 276 | 158 | −118 | **faster** |
|
||||
| session p50 (ms) | 7 | 0 | −7 | **faster** |
|
||||
| 1st-audio p50 (ms) | 440 | 664 | +224 | slower¹ |
|
||||
| total p50 (ms) | 816 | 1010 | +194 | slower¹ |
|
||||
| total p95 (ms) | 2152 | 1970 | −182 | **faster** |
|
||||
| total p99 (ms) | 2692 | 2610 | −82 | **faster** |
|
||||
|
||||
The gateway is **faster than direct on 4 of 6 metrics**. The warm pool makes the
|
||||
**session phase sub-millisecond** at the median — ~76% of connects hit the pool,
|
||||
~70% had session < 1 ms. ¹ The two "slower" rows are not gateway overhead:
|
||||
`1st-audio` is OpenAI's own inference time (the gateway only relays it), which ran
|
||||
slower during the gateway legs and drags `total p50` with it.
|
||||
|
||||
**Pool OFF** (control, `REALTIME_POOL_SIZE=0`): session p50 was **367 ms** — the
|
||||
fresh-dial overhead the pool removes.
|
||||
|
||||
## Reproduce
|
||||
|
||||
The load generator lives in a separate repo:
|
||||
**https://github.com/ishaan-berri/litellm-realtime-bench**
|
||||
|
||||
```bash
|
||||
git clone https://github.com/ishaan-berri/litellm-realtime-bench
|
||||
cd litellm-realtime-bench && go build -o wsbench .
|
||||
|
||||
# Direct to OpenAI (baseline)
|
||||
./wsbench -host api.openai.com -key "$OPENAI_API_KEY" -m gpt-realtime -n 5000 -c 500 -t 60
|
||||
|
||||
# Through the gateway — run once with pool ON, once with REALTIME_POOL_SIZE=0
|
||||
./wsbench -host <gateway-host> -key "$LITELLM_MASTER_KEY" -m gpt-realtime -n 5000 -c 500 -t 60
|
||||
```
|
||||
|
||||
Run the gateway with the env stand-in (`OPENAI_REALTIME_MODEL=gpt-realtime`,
|
||||
`OPENAI_API_KEY`, `LITELLM_MASTER_KEY`, `REALTIME_POOL_SIZE`, `HOST=0.0.0.0`). At
|
||||
500 concurrency over N instances, size the pool to `≈ 500 / N` per instance (64 was
|
||||
used here for 10 instances). The bench repo's README covers running 500-concurrency
|
||||
legs from a hosted multi-vCPU runner. **Never commit keys — pass them via `-key`.**
|
||||
13
litellm-rust/crates/ai-gateway/config.yaml
Normal file
13
litellm-rust/crates/ai-gateway/config.yaml
Normal file
|
|
@ -0,0 +1,13 @@
|
|||
# Sample realtime config for the LiteLLM Rust AI Gateway.
|
||||
#
|
||||
# The gateway loads this model_list at boot via the embedded python config
|
||||
# reader (litellm.proxy.read_model_list), which reuses the proxy's own reader —
|
||||
# so include:, os.environ/ secrets, and DB-stored models all work here too.
|
||||
#
|
||||
# Secrets are referenced (never inlined) via os.environ/. A real deploy can
|
||||
# override this file (e.g. mount a Render secret file at LITELLM_CONFIG_PATH).
|
||||
model_list:
|
||||
- model_name: gpt-realtime
|
||||
litellm_params:
|
||||
model: openai/gpt-realtime
|
||||
api_key: os.environ/OPENAI_API_KEY
|
||||
35
litellm-rust/crates/ai-gateway/render.yaml
Normal file
35
litellm-rust/crates/ai-gateway/render.yaml
Normal file
|
|
@ -0,0 +1,35 @@
|
|||
# Render blueprint for the LiteLLM Rust AI Gateway (realtime WebSocket proxy).
|
||||
#
|
||||
# Single instance for now (no autoscaling). The public endpoint is a
|
||||
# WebSocket served over TLS: wss://<service>.onrender.com/v1/realtime
|
||||
#
|
||||
# Paths are relative to the **repo root** (Render's convention). The build
|
||||
# context is the repo root so the image can install litellm from source — the
|
||||
# gateway loads its model_list via litellm.proxy.read_model_list at boot.
|
||||
#
|
||||
# Secrets (LITELLM_MASTER_KEY, OPENAI_API_KEY) are marked sync: false — set
|
||||
# them in the Render dashboard or via the API, never inline here.
|
||||
services:
|
||||
- type: web
|
||||
name: litellm-rust-ai-gateway
|
||||
runtime: docker
|
||||
plan: standard
|
||||
dockerfilePath: ./litellm-rust/crates/ai-gateway/Dockerfile
|
||||
dockerContext: .
|
||||
healthCheckPath: /health/readiness
|
||||
numInstances: 1
|
||||
envVars:
|
||||
# The gateway loads its model_list from this config.yaml via the embedded
|
||||
# python config reader. The image bakes a default config at /app/config.yaml;
|
||||
# a real deploy can override it by mounting a Render secret file at this
|
||||
# same path (Dashboard → Environment → Secret Files) — never inline secrets.
|
||||
- key: LITELLM_CONFIG_PATH
|
||||
value: /app/config.yaml
|
||||
- key: HOST
|
||||
value: 0.0.0.0
|
||||
# Bearer token clients must send on /v1/realtime (fail closed if unset).
|
||||
- key: LITELLM_MASTER_KEY
|
||||
sync: false
|
||||
# Referenced by config.yaml as os.environ/OPENAI_API_KEY for the upstream dial.
|
||||
- key: OPENAI_API_KEY
|
||||
sync: false
|
||||
54
litellm-rust/crates/ai-gateway/src/auth/mod.rs
Normal file
54
litellm-rust/crates/ai-gateway/src/auth/mod.rs
Normal file
|
|
@ -0,0 +1,54 @@
|
|||
//! Gateway authentication, as an axum **extractor** (the idiomatic pattern —
|
||||
//! keeps handlers clean and auth testable).
|
||||
//!
|
||||
//! For now this is a single **master key**: any caller presenting it as
|
||||
//! `Authorization: Bearer <key>` may invoke the gateway. Per-key auth, budgets,
|
||||
//! and rate limits are delegated to the Python proxy in a later phase.
|
||||
//!
|
||||
//! A handler opts in by adding [`RequireMasterKey`] to its arguments; auth then
|
||||
//! runs during extraction, before the handler body. Routes never re-implement it.
|
||||
|
||||
use axum::extract::FromRequestParts;
|
||||
use axum::http::header::AUTHORIZATION;
|
||||
use axum::http::request::Parts;
|
||||
use axum::http::StatusCode;
|
||||
use subtle::ConstantTimeEq;
|
||||
|
||||
use crate::state::AppState;
|
||||
|
||||
/// Extractor that requires the configured master key as a bearer token.
|
||||
///
|
||||
/// Rejections: `500` when no master key is configured (permanent
|
||||
/// misconfiguration, not a transient outage); `401` on a missing/incorrect
|
||||
/// token. The comparison is constant-time.
|
||||
pub struct RequireMasterKey;
|
||||
|
||||
#[axum::async_trait]
|
||||
impl FromRequestParts<AppState> for RequireMasterKey {
|
||||
type Rejection = (StatusCode, String);
|
||||
|
||||
async fn from_request_parts(
|
||||
parts: &mut Parts,
|
||||
state: &AppState,
|
||||
) -> Result<Self, Self::Rejection> {
|
||||
let Some(expected) = state.master_key.as_deref() else {
|
||||
return Err((
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
"gateway auth not configured (set LITELLM_MASTER_KEY)".to_string(),
|
||||
));
|
||||
};
|
||||
let provided = parts
|
||||
.headers
|
||||
.get(AUTHORIZATION)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.and_then(|value| value.strip_prefix("Bearer "))
|
||||
.map(str::trim);
|
||||
match provided {
|
||||
Some(token) if bool::from(token.as_bytes().ct_eq(expected.as_bytes())) => Ok(Self),
|
||||
_ => Err((
|
||||
StatusCode::UNAUTHORIZED,
|
||||
"missing or invalid bearer token".to_string(),
|
||||
)),
|
||||
}
|
||||
}
|
||||
}
|
||||
58
litellm-rust/crates/ai-gateway/src/gil.rs
Normal file
58
litellm-rust/crates/ai-gateway/src/gil.rs
Normal file
|
|
@ -0,0 +1,58 @@
|
|||
//! GIL-activity tracking.
|
||||
//!
|
||||
//! Every acquisition of the Python GIL is recorded here so the `/health/gil`
|
||||
//! endpoint can report whether Python was touched recently. The design goal is
|
||||
//! that the GIL is acquired **only at load time** (config read) and never on the
|
||||
//! realtime hot path — polling this endpoint during traffic should show the
|
||||
//! count holding steady and `acquired_last_30s` falling to `false`.
|
||||
|
||||
use std::sync::atomic::{AtomicU64, Ordering};
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
|
||||
/// Window (seconds) for the "recently acquired" signal.
|
||||
pub const RECENT_WINDOW_SECS: u64 = 30;
|
||||
|
||||
static GIL_ACQUISITIONS: AtomicU64 = AtomicU64::new(0);
|
||||
/// Unix seconds of the last acquisition; `0` means "never".
|
||||
static LAST_GIL_UNIX_SECS: AtomicU64 = AtomicU64::new(0);
|
||||
|
||||
fn now_unix_secs() -> u64 {
|
||||
SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.map(|d| d.as_secs())
|
||||
.unwrap_or(0)
|
||||
}
|
||||
|
||||
/// Record that the GIL was just acquired. Call immediately before taking the GIL.
|
||||
///
|
||||
/// Only invoked under the `python-config` feature; without it the gateway never
|
||||
/// touches Python, so the recorder is unused (and the endpoint reports zero).
|
||||
#[cfg_attr(not(feature = "python-config"), allow(dead_code))]
|
||||
pub fn record_acquisition() {
|
||||
GIL_ACQUISITIONS.fetch_add(1, Ordering::Relaxed);
|
||||
LAST_GIL_UNIX_SECS.store(now_unix_secs(), Ordering::Relaxed);
|
||||
}
|
||||
|
||||
/// Point-in-time view of GIL activity.
|
||||
pub struct GilSnapshot {
|
||||
pub total_acquisitions: u64,
|
||||
pub seconds_since_last: Option<u64>,
|
||||
pub acquired_last_30s: bool,
|
||||
}
|
||||
|
||||
/// Read the current GIL-activity snapshot.
|
||||
pub fn snapshot() -> GilSnapshot {
|
||||
let total = GIL_ACQUISITIONS.load(Ordering::Relaxed);
|
||||
let last = LAST_GIL_UNIX_SECS.load(Ordering::Relaxed);
|
||||
let seconds_since_last = if last == 0 {
|
||||
None
|
||||
} else {
|
||||
Some(now_unix_secs().saturating_sub(last))
|
||||
};
|
||||
let acquired_last_30s = seconds_since_last.is_some_and(|secs| secs <= RECENT_WINDOW_SECS);
|
||||
GilSnapshot {
|
||||
total_acquisitions: total,
|
||||
seconds_since_last,
|
||||
acquired_last_30s,
|
||||
}
|
||||
}
|
||||
3
litellm-rust/crates/ai-gateway/src/io/mod.rs
Normal file
3
litellm-rust/crates/ai-gateway/src/io/mod.rs
Normal file
|
|
@ -0,0 +1,3 @@
|
|||
pub mod ocr;
|
||||
pub mod realtime;
|
||||
pub mod realtime_pool;
|
||||
127
litellm-rust/crates/ai-gateway/src/io/ocr.rs
Normal file
127
litellm-rust/crates/ai-gateway/src/io/ocr.rs
Normal file
|
|
@ -0,0 +1,127 @@
|
|||
//! End-to-end OCR orchestration.
|
||||
//!
|
||||
//! Owns the whole Mistral OCR call so the Python side stays a thin bridge:
|
||||
//! resolve the API key, build the URL + body via the pure transforms, POST it,
|
||||
//! and normalize the response. The HTTP client is built once and reused.
|
||||
|
||||
use std::sync::OnceLock;
|
||||
use std::time::Duration;
|
||||
|
||||
use litellm_core::error::CoreError;
|
||||
use litellm_core::ocr::transformation::OcrProviderConfig;
|
||||
use litellm_core::CoreResult;
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
use litellm_core::providers::mistral::ocr::transformation as mistral;
|
||||
use litellm_core::providers::mistral::ocr::transformation::MISTRAL_OCR_CONFIG;
|
||||
|
||||
/// OCR over large documents can take a while; bound it generously rather than
|
||||
/// hanging forever on an unresponsive upstream. The client-level limit is the
|
||||
/// outer ceiling; callers can tighten it per request via ``run_ocr``'s ``timeout``.
|
||||
const OCR_TIMEOUT_SECS: u64 = 600;
|
||||
|
||||
/// Maximum upstream body characters retained in error messages. OCR responses
|
||||
/// can echo document contents and prompts; keep enough for debugging without
|
||||
/// forwarding sensitive payloads across the host boundary.
|
||||
const ERROR_BODY_MAX_CHARS: usize = 256;
|
||||
|
||||
/// Process-wide blocking HTTP client (connection pool + TLS reused across calls).
|
||||
fn http_client() -> &'static reqwest::blocking::Client {
|
||||
static CLIENT: OnceLock<reqwest::blocking::Client> = OnceLock::new();
|
||||
CLIENT.get_or_init(|| {
|
||||
reqwest::blocking::Client::builder()
|
||||
.timeout(Duration::from_secs(OCR_TIMEOUT_SECS))
|
||||
.build()
|
||||
.expect("failed to build reqwest client")
|
||||
})
|
||||
}
|
||||
|
||||
fn truncate_error_body(body: &str) -> String {
|
||||
if body.chars().count() <= ERROR_BODY_MAX_CHARS {
|
||||
return body.to_string();
|
||||
}
|
||||
let truncated: String = body.chars().take(ERROR_BODY_MAX_CHARS).collect();
|
||||
format!("{truncated}... (truncated)")
|
||||
}
|
||||
|
||||
/// Perform a Mistral OCR call end to end and return the normalized response as
|
||||
/// JSON (the shape the Python `OCRResponse` model expects).
|
||||
///
|
||||
/// Blocking: intended to be called with the GIL released from the Python bridge.
|
||||
pub fn run_ocr(
|
||||
model: &str,
|
||||
document: Value,
|
||||
api_key: Option<&str>,
|
||||
api_base: Option<&str>,
|
||||
optional_params: Map<String, Value>,
|
||||
timeout: Option<Duration>,
|
||||
) -> CoreResult<Value> {
|
||||
let config = &MISTRAL_OCR_CONFIG;
|
||||
|
||||
let api_key = mistral::resolve_api_key(api_key, &|key| std::env::var(key).ok())?;
|
||||
let url = mistral::complete_url(api_base);
|
||||
let filtered_params = config.map_ocr_params(&optional_params);
|
||||
let body = config
|
||||
.transform_ocr_request(model, document, filtered_params)?
|
||||
.data;
|
||||
|
||||
let mut request = http_client().post(&url).bearer_auth(&api_key).json(&body);
|
||||
if let Some(duration) = timeout {
|
||||
request = request.timeout(duration);
|
||||
}
|
||||
|
||||
let response = request
|
||||
.send()
|
||||
.map_err(|err| CoreError::Network(err.to_string()))?;
|
||||
|
||||
let status = response.status();
|
||||
let text = response
|
||||
.text()
|
||||
.map_err(|err| CoreError::Network(err.to_string()))?;
|
||||
|
||||
if !status.is_success() {
|
||||
return Err(CoreError::Http {
|
||||
status: status.as_u16(),
|
||||
body: truncate_error_body(&text),
|
||||
});
|
||||
}
|
||||
|
||||
let response_json: Value = serde_json::from_str(&text)
|
||||
.map_err(|err| CoreError::InvalidResponse(format!("invalid OCR response JSON: {err}")))?;
|
||||
|
||||
Ok(config
|
||||
.transform_ocr_response(model, response_json)?
|
||||
.into_json())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn truncate_error_body_passes_short_strings_through() {
|
||||
let body = "Unauthorized";
|
||||
assert_eq!(truncate_error_body(body), "Unauthorized");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn truncate_error_body_caps_long_payloads() {
|
||||
let body = "x".repeat(ERROR_BODY_MAX_CHARS + 50);
|
||||
let truncated = truncate_error_body(&body);
|
||||
|
||||
assert!(truncated.ends_with("... (truncated)"));
|
||||
let prefix_chars = truncated
|
||||
.strip_suffix("... (truncated)")
|
||||
.expect("truncated marker present")
|
||||
.chars()
|
||||
.count();
|
||||
assert_eq!(prefix_chars, ERROR_BODY_MAX_CHARS);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn truncate_error_body_does_not_split_multibyte_chars() {
|
||||
let body = "é".repeat(ERROR_BODY_MAX_CHARS + 10);
|
||||
let truncated = truncate_error_body(&body);
|
||||
assert!(truncated.is_char_boundary(truncated.len()));
|
||||
}
|
||||
}
|
||||
374
litellm-rust/crates/ai-gateway/src/io/realtime.rs
Normal file
374
litellm-rust/crates/ai-gateway/src/io/realtime.rs
Normal file
|
|
@ -0,0 +1,374 @@
|
|||
//! End-to-end OpenAI realtime invocation.
|
||||
//!
|
||||
//! The host-facing entry point, mirroring `crate::io::ocr::run_ocr`: open the
|
||||
//! WebSocket to OpenAI, then splice a client realtime stream to the upstream,
|
||||
//! driving typed events through the pure `OPENAI_REALTIME_CONFIG` transforms.
|
||||
//! Network, auth header, key resolution, and wire (de)serialization live here so
|
||||
//! the `transformation` module stays pure and typed.
|
||||
//!
|
||||
//! The dial and splice steps are factored out ([`dial_upstream`], [`splice`]) so
|
||||
//! the connection pool ([`crate::io::realtime_pool`]) can pre-establish an upstream,
|
||||
//! buffer its `session.created`, and later hand the live socket to the same
|
||||
//! splice loop a fresh dial uses.
|
||||
|
||||
use std::time::Duration;
|
||||
|
||||
use futures_util::stream::{SplitSink, SplitStream};
|
||||
use futures_util::{Sink, SinkExt, Stream, StreamExt};
|
||||
use litellm_core::error::CoreError;
|
||||
use litellm_core::realtime::transformation::RealtimeProviderConfig;
|
||||
use litellm_core::realtime::types::RealtimeEvent;
|
||||
use litellm_core::CoreResult;
|
||||
use tokio::net::TcpStream;
|
||||
use tokio_tungstenite::tungstenite::client::IntoClientRequest;
|
||||
use tokio_tungstenite::tungstenite::http::header::AUTHORIZATION;
|
||||
use tokio_tungstenite::tungstenite::http::HeaderValue;
|
||||
use tokio_tungstenite::tungstenite::Message;
|
||||
use tokio_tungstenite::{connect_async, MaybeTlsStream, WebSocketStream};
|
||||
|
||||
use litellm_core::providers::openai::realtime::transformation::OPENAI_REALTIME_CONFIG;
|
||||
|
||||
/// Environment variable holding the OpenAI API key (last-resort fallback).
|
||||
const OPENAI_API_KEY_ENV: &str = "OPENAI_API_KEY";
|
||||
|
||||
const MISSING_KEY_MESSAGE: &str = "Missing OpenAI API Key - a realtime call is being made but no key was passed via params or the OPENAI_API_KEY environment variable";
|
||||
|
||||
/// Default **idle** timeout: if neither side sends a frame for this long, the
|
||||
/// session is reaped. It resets on any activity, so it does not cap a healthy
|
||||
/// (continuously streaming) session — it only frees a stalled one (e.g. a
|
||||
/// half-open upstream that keeps the socket open but stops sending).
|
||||
const DEFAULT_IDLE_TIMEOUT_SECS: u64 = 300;
|
||||
|
||||
/// The concrete upstream WebSocket type (TLS or plain). Shared by the dial path
|
||||
/// and the pool so warm sockets and fresh sockets are the exact same type.
|
||||
pub type UpstreamWs = WebSocketStream<MaybeTlsStream<TcpStream>>;
|
||||
pub(crate) type UpstreamTx = SplitSink<UpstreamWs, Message>;
|
||||
pub(crate) type UpstreamRx = SplitStream<UpstreamWs>;
|
||||
|
||||
/// Resolve the OpenAI API key from the explicit param or the environment.
|
||||
///
|
||||
/// Blank/whitespace values are treated as absent (guard at resolution time).
|
||||
pub(crate) fn resolve_api_key(api_key: Option<&str>) -> CoreResult<String> {
|
||||
api_key
|
||||
.map(str::trim)
|
||||
.filter(|key| !key.is_empty())
|
||||
.map(str::to_string)
|
||||
.or_else(|| {
|
||||
std::env::var(OPENAI_API_KEY_ENV)
|
||||
.ok()
|
||||
.filter(|key| !key.trim().is_empty())
|
||||
})
|
||||
.ok_or_else(|| CoreError::Auth(MISSING_KEY_MESSAGE.to_string()))
|
||||
}
|
||||
|
||||
/// Open the upstream WebSocket to OpenAI for `(model, api_key, api_base)`.
|
||||
///
|
||||
/// This is the dial half of [`realtime`], factored out so the pool can
|
||||
/// pre-establish sockets ahead of any client. `api_key` here is already resolved
|
||||
/// (non-blank) — the pool resolves it once when it is created.
|
||||
pub(crate) async fn dial_upstream(
|
||||
model: &str,
|
||||
api_key: &str,
|
||||
api_base: Option<&str>,
|
||||
) -> CoreResult<UpstreamWs> {
|
||||
let url = OPENAI_REALTIME_CONFIG.complete_url(api_base, model);
|
||||
|
||||
let mut request = url
|
||||
.as_str()
|
||||
.into_client_request()
|
||||
.map_err(|err| CoreError::Network(err.to_string()))?;
|
||||
// GA realtime: only Authorization. The legacy OpenAI-Beta header triggers
|
||||
// beta_api_shape_disabled, so we do not send it.
|
||||
request.headers_mut().insert(
|
||||
AUTHORIZATION,
|
||||
HeaderValue::from_str(&format!("Bearer {api_key}"))
|
||||
.map_err(|err| CoreError::Auth(err.to_string()))?,
|
||||
);
|
||||
|
||||
let (upstream, _response) = connect_async(request)
|
||||
.await
|
||||
.map_err(|err| CoreError::Network(err.to_string()))?;
|
||||
Ok(upstream)
|
||||
}
|
||||
|
||||
/// Read the next text frame from the upstream and decode it as a typed event.
|
||||
///
|
||||
/// Used by the pool to pre-read OpenAI's unprompted `session.created`. Returns an
|
||||
/// error on a non-text frame, a closed socket, or undecodable JSON so the pool can
|
||||
/// discard a misbehaving socket rather than warm it.
|
||||
pub(crate) async fn read_event(upstream_rx: &mut UpstreamRx) -> CoreResult<RealtimeEvent> {
|
||||
loop {
|
||||
let message = upstream_rx
|
||||
.next()
|
||||
.await
|
||||
.ok_or_else(|| CoreError::Network("upstream closed before first event".to_string()))?
|
||||
.map_err(|err| CoreError::Network(err.to_string()))?;
|
||||
match message {
|
||||
Message::Text(text) => {
|
||||
return serde_json::from_str(&text)
|
||||
.map_err(|err| CoreError::InvalidResponse(err.to_string()));
|
||||
}
|
||||
// Ignore protocol frames (ping/pong) while waiting for the first event.
|
||||
Message::Ping(_) | Message::Pong(_) => continue,
|
||||
Message::Close(_) => {
|
||||
return Err(CoreError::Network(
|
||||
"upstream closed before first event".to_string(),
|
||||
))
|
||||
}
|
||||
_ => continue,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Splice an already-connected upstream to the client streams.
|
||||
///
|
||||
/// `prelude` is relayed to the client first (the pool passes the buffered
|
||||
/// `session.created` here; the fresh-dial path passes `None` and lets the upstream
|
||||
/// deliver it). Then a single select loop forwards both directions through the
|
||||
/// transforms until either side closes or the idle timeout fires.
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub(crate) async fn splice<In, Out>(
|
||||
model: &str,
|
||||
mut upstream_tx: UpstreamTx,
|
||||
mut upstream_rx: UpstreamRx,
|
||||
prelude: Option<RealtimeEvent>,
|
||||
idle_timeout: Option<Duration>,
|
||||
mut client_in: In,
|
||||
mut client_out: Out,
|
||||
) -> CoreResult<()>
|
||||
where
|
||||
In: Stream<Item = RealtimeEvent> + Unpin + Send,
|
||||
Out: Sink<RealtimeEvent> + Unpin + Send,
|
||||
<Out as Sink<RealtimeEvent>>::Error: std::fmt::Display,
|
||||
{
|
||||
let config = &OPENAI_REALTIME_CONFIG;
|
||||
|
||||
// Relay a buffered backend event (warm handoff's session.created) first, so a
|
||||
// warm session looks identical to a fresh one from the client's view.
|
||||
if let Some(event) = prelude {
|
||||
for outbound in config.transform_realtime_response(&event, model)?.events {
|
||||
client_out
|
||||
.send(outbound)
|
||||
.await
|
||||
.map_err(|err| CoreError::Network(err.to_string()))?;
|
||||
}
|
||||
}
|
||||
|
||||
let idle = idle_timeout.unwrap_or(Duration::from_secs(DEFAULT_IDLE_TIMEOUT_SECS));
|
||||
|
||||
// One loop forwarding both directions. The `sleep(idle)` arm is rebuilt every
|
||||
// iteration, so any frame (either way) resets it — it fires only when the
|
||||
// session has been fully idle for `idle`, reaping a stalled connection
|
||||
// (task + upstream TCP socket) instead of leaking it.
|
||||
loop {
|
||||
tokio::select! {
|
||||
// client -> upstream
|
||||
client_event = client_in.next() => {
|
||||
let Some(event) = client_event else { break }; // client disconnected
|
||||
for outbound in config.transform_realtime_request(&event, model)?.events {
|
||||
let payload = serde_json::to_string(&outbound)
|
||||
.map_err(|err| CoreError::InvalidResponse(err.to_string()))?;
|
||||
upstream_tx
|
||||
.send(Message::Text(payload))
|
||||
.await
|
||||
.map_err(|err| CoreError::Network(err.to_string()))?;
|
||||
}
|
||||
}
|
||||
// upstream -> client
|
||||
upstream_message = upstream_rx.next() => {
|
||||
let Some(message) = upstream_message else { break }; // upstream closed
|
||||
match message.map_err(|err| CoreError::Network(err.to_string()))? {
|
||||
Message::Text(text) => {
|
||||
let event: RealtimeEvent = serde_json::from_str(&text)
|
||||
.map_err(|err| CoreError::InvalidResponse(err.to_string()))?;
|
||||
for outbound in config.transform_realtime_response(&event, model)?.events {
|
||||
client_out
|
||||
.send(outbound)
|
||||
.await
|
||||
.map_err(|err| CoreError::Network(err.to_string()))?;
|
||||
}
|
||||
}
|
||||
Message::Close(_) => break,
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
// idle timeout: no activity from either side within `idle`
|
||||
_ = tokio::time::sleep(idle) => break,
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Splice a client realtime stream to OpenAI: forward client events upstream
|
||||
/// (via `transform_realtime_request`) and backend events downstream (via
|
||||
/// `transform_realtime_response`). Returns when either side closes.
|
||||
///
|
||||
/// Generic over the client transport (typed events) so this crate stays
|
||||
/// framework-agnostic; the gateway adapts its axum socket to these. This is the
|
||||
/// fresh-dial path: dial, then splice. The pool's warm-handoff path skips the dial
|
||||
/// and calls [`splice`] directly with a buffered `session.created`.
|
||||
pub async fn realtime<In, Out>(
|
||||
model: &str,
|
||||
api_key: Option<&str>,
|
||||
api_base: Option<&str>,
|
||||
idle_timeout: Option<Duration>,
|
||||
client_in: In,
|
||||
client_out: Out,
|
||||
) -> CoreResult<()>
|
||||
where
|
||||
In: Stream<Item = RealtimeEvent> + Unpin + Send,
|
||||
Out: Sink<RealtimeEvent> + Unpin + Send,
|
||||
<Out as Sink<RealtimeEvent>>::Error: std::fmt::Display,
|
||||
{
|
||||
let api_key = resolve_api_key(api_key)?;
|
||||
let upstream = dial_upstream(model, &api_key, api_base).await?;
|
||||
let (upstream_tx, upstream_rx) = upstream.split();
|
||||
splice(
|
||||
model,
|
||||
upstream_tx,
|
||||
upstream_rx,
|
||||
None,
|
||||
idle_timeout,
|
||||
client_in,
|
||||
client_out,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
/// Splice a pre-warmed upstream (taken from [`crate::io::realtime_pool`]) to the
|
||||
/// client. Relays the buffered `session.created` first, then splices exactly like
|
||||
/// the fresh-dial path — so a warm session is indistinguishable from a fresh one.
|
||||
pub async fn realtime_warm<In, Out>(
|
||||
model: &str,
|
||||
handoff: crate::io::realtime_pool::WarmHandoff,
|
||||
idle_timeout: Option<Duration>,
|
||||
client_in: In,
|
||||
client_out: Out,
|
||||
) -> CoreResult<()>
|
||||
where
|
||||
In: Stream<Item = RealtimeEvent> + Unpin + Send,
|
||||
Out: Sink<RealtimeEvent> + Unpin + Send,
|
||||
<Out as Sink<RealtimeEvent>>::Error: std::fmt::Display,
|
||||
{
|
||||
splice(
|
||||
model,
|
||||
handoff.tx,
|
||||
handoff.rx,
|
||||
Some(handoff.session_created),
|
||||
idle_timeout,
|
||||
client_in,
|
||||
client_out,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn event(raw: &str) -> RealtimeEvent {
|
||||
serde_json::from_str(raw).expect("valid event json")
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolve_api_key_prefers_param_then_blank_falls_through() {
|
||||
assert_eq!(resolve_api_key(Some("sk-test")).unwrap(), "sk-test");
|
||||
// A blank param with no env set should error.
|
||||
if std::env::var(OPENAI_API_KEY_ENV).is_err() {
|
||||
assert!(resolve_api_key(Some(" ")).is_err());
|
||||
}
|
||||
}
|
||||
|
||||
/// Live end-to-end check against OpenAI. Ignored by default (CI never runs
|
||||
/// it); run explicitly with `OPENAI_API_KEY` set:
|
||||
/// `cargo test -p litellm-ai-gateway --features server realtime_invokes_openai -- --ignored --nocapture`
|
||||
#[tokio::test]
|
||||
#[ignore = "hits the live OpenAI realtime API; needs OPENAI_API_KEY"]
|
||||
async fn realtime_invokes_openai_and_responds() {
|
||||
use futures_channel::mpsc;
|
||||
|
||||
let key =
|
||||
std::env::var(OPENAI_API_KEY_ENV).expect("set OPENAI_API_KEY to run this ignored test");
|
||||
|
||||
// client -> provider (we hold `client_tx` to push events upstream)
|
||||
let (mut client_tx, client_in) = mpsc::unbounded::<RealtimeEvent>();
|
||||
// provider -> client (we hold `backend_rx` to read backend events)
|
||||
let (client_out, mut backend_rx) = mpsc::unbounded::<RealtimeEvent>();
|
||||
|
||||
// Clone the key so the spawned task owns its `String` (no borrow across await).
|
||||
let key_owned = key.clone();
|
||||
let call = tokio::spawn(async move {
|
||||
realtime(
|
||||
"gpt-realtime",
|
||||
Some(&key_owned),
|
||||
None,
|
||||
None,
|
||||
client_in,
|
||||
client_out,
|
||||
)
|
||||
.await
|
||||
});
|
||||
|
||||
// 1. First backend event should be session.created.
|
||||
let first = tokio::time::timeout(Duration::from_secs(30), backend_rx.next())
|
||||
.await
|
||||
.expect("timed out waiting for session.created")
|
||||
.expect("backend stream closed before session.created");
|
||||
assert_eq!(
|
||||
first.event_type, "session.created",
|
||||
"expected session.created, got: {}",
|
||||
first.event_type
|
||||
);
|
||||
|
||||
// 2. Ask for a short audio response.
|
||||
client_tx
|
||||
.send(event(
|
||||
r#"{"type":"conversation.item.create","item":{"type":"message","role":"user","content":[{"type":"input_text","text":"Say hi."}]}}"#,
|
||||
))
|
||||
.await
|
||||
.expect("send conversation.item.create");
|
||||
client_tx
|
||||
.send(event(r#"{"type":"response.create"}"#))
|
||||
.await
|
||||
.expect("send response.create");
|
||||
|
||||
// 3. Read backend events; require a non-empty audio delta, then response.done.
|
||||
let mut saw_audio_delta = false;
|
||||
let mut saw_done = false;
|
||||
for _ in 0..500 {
|
||||
let next = tokio::time::timeout(Duration::from_secs(30), backend_rx.next()).await;
|
||||
let event = match next {
|
||||
Ok(Some(event)) => event,
|
||||
Ok(None) => break,
|
||||
Err(_) => panic!("timed out waiting for backend events"),
|
||||
};
|
||||
match event.event_type.as_str() {
|
||||
"response.output_audio.delta" => {
|
||||
let delta = event
|
||||
.data
|
||||
.get("delta")
|
||||
.and_then(|value| value.as_str())
|
||||
.unwrap_or("");
|
||||
if !delta.is_empty() {
|
||||
saw_audio_delta = true;
|
||||
}
|
||||
}
|
||||
"response.done" => {
|
||||
saw_done = true;
|
||||
break;
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
|
||||
assert!(
|
||||
saw_audio_delta,
|
||||
"expected a response.output_audio.delta with non-empty delta"
|
||||
);
|
||||
assert!(saw_done, "expected a response.done event");
|
||||
|
||||
// Drop the client sender so the provider's to_upstream side finishes.
|
||||
drop(client_tx);
|
||||
let _ = call.await;
|
||||
}
|
||||
}
|
||||
712
litellm-rust/crates/ai-gateway/src/io/realtime_pool.rs
Normal file
712
litellm-rust/crates/ai-gateway/src/io/realtime_pool.rs
Normal file
|
|
@ -0,0 +1,712 @@
|
|||
//! Pre-warmed upstream realtime connection pool.
|
||||
//!
|
||||
//! The gateway's realtime overhead lives entirely in session establishment: on
|
||||
//! every client connect it dials a fresh upstream WS to OpenAI and waits for
|
||||
//! `session.created` before it can serve. This pool keeps a small set of upstream
|
||||
//! sockets **already connected and already past `session.created`** so a connect
|
||||
//! can be served from a warm socket and the handshake is off the critical path.
|
||||
//!
|
||||
//! Layering: this lives in the gateway's `io` module next to the dial/splice it
|
||||
//! reuses. The gateway holds an `Arc<RealtimePool>` in its state and asks for a
|
||||
//! warm socket per connect; on a miss it fresh-dials exactly as before. The pool
|
||||
//! is a latency optimization, never a correctness dependency — see the gateway's
|
||||
//! `src/routes/realtime/README.md`.
|
||||
//!
|
||||
//! ## Caveats (enforced here)
|
||||
//! - One warm socket serves exactly one session (realtime isn't multiplexed), so
|
||||
//! the pool is sized to the connect *rate*, not concurrent connections.
|
||||
//! - `session.created` is pre-read once and buffered; nothing else is read from a
|
||||
//! warm socket before handoff, so a warm session starts at OpenAI defaults just
|
||||
//! like a fresh one (`session.update` semantics unchanged).
|
||||
//! - Warm sockets are short-lived (`max_idle`) and liveness-checked at handoff to
|
||||
//! bound idle billing / dodge OpenAI's idle timeout.
|
||||
//! - On miss or dead socket the caller fresh-dials; the pool never blocks or fails
|
||||
//! a connect because it is empty.
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::sync::{Arc, Mutex};
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
use futures_util::StreamExt;
|
||||
use litellm_core::realtime::types::RealtimeEvent;
|
||||
use litellm_core::CoreResult;
|
||||
|
||||
use crate::io::realtime::{
|
||||
dial_upstream, read_event, resolve_api_key, UpstreamRx, UpstreamTx, UpstreamWs,
|
||||
};
|
||||
|
||||
/// Default target warm sockets per key when pooling is enabled.
|
||||
pub const DEFAULT_POOL_SIZE: usize = 4;
|
||||
|
||||
/// Default max time a warm socket may sit before it is closed and replaced.
|
||||
pub const DEFAULT_MAX_IDLE: Duration = Duration::from_secs(30);
|
||||
|
||||
/// Env var: target warm sockets per key. `0` disables pooling (fresh-dial only).
|
||||
pub const POOL_SIZE_ENV: &str = "REALTIME_POOL_SIZE";
|
||||
|
||||
/// Env var: max warm-socket idle lifetime, in seconds.
|
||||
pub const MAX_IDLE_ENV: &str = "REALTIME_POOL_MAX_IDLE_SECS";
|
||||
|
||||
/// How often the background replenisher wakes to top up and reap stale sockets.
|
||||
const REPLENISH_TICK: Duration = Duration::from_millis(250);
|
||||
|
||||
/// Backoff floor after a key's warm-up dials all fail. The first failed pass
|
||||
/// waits this long before retrying that key.
|
||||
const BACKOFF_BASE: Duration = Duration::from_millis(500);
|
||||
|
||||
/// Backoff ceiling. A key that keeps failing (invalid credentials, an
|
||||
/// unreachable upstream) is retried at most once per this interval — instead of
|
||||
/// firing `needed` concurrent TLS dials every 250 ms tick, which would hammer
|
||||
/// the upstream and risk rate-limit exhaustion that degrades valid cold-path
|
||||
/// traffic. Backoff resets the moment a dial for the key succeeds.
|
||||
const BACKOFF_MAX: Duration = Duration::from_secs(30);
|
||||
|
||||
/// Identifies an upstream connection: the tuple that fully determines the dial.
|
||||
/// `api_key` is included so a warm socket is only ever reused for the same key
|
||||
/// (no cross-tenant reuse).
|
||||
#[derive(Clone, PartialEq, Eq, Hash)]
|
||||
pub struct UpstreamKey {
|
||||
pub model: String,
|
||||
pub api_key: String,
|
||||
pub api_base: Option<String>,
|
||||
}
|
||||
|
||||
impl std::fmt::Debug for UpstreamKey {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
f.debug_struct("UpstreamKey")
|
||||
.field("model", &self.model)
|
||||
.field("api_key", &"[REDACTED]")
|
||||
.field("api_base", &self.api_base)
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
|
||||
/// A warm upstream: split halves + the buffered `session.created` + when it was
|
||||
/// warmed (for `max_idle` expiry).
|
||||
struct WarmConnection {
|
||||
tx: UpstreamTx,
|
||||
rx: UpstreamRx,
|
||||
session_created: RealtimeEvent,
|
||||
warmed_at: Instant,
|
||||
}
|
||||
|
||||
/// A live upstream taken from the pool, ready to splice. The caller relays
|
||||
/// `session_created` to the client first, then splices `(tx, rx)` as usual.
|
||||
pub struct WarmHandoff {
|
||||
pub tx: UpstreamTx,
|
||||
pub rx: UpstreamRx,
|
||||
pub session_created: RealtimeEvent,
|
||||
}
|
||||
|
||||
/// Pool configuration, resolved once at startup from the environment.
|
||||
#[derive(Clone, Copy, Debug)]
|
||||
pub struct PoolConfig {
|
||||
/// Target warm sockets per key. `0` disables pooling.
|
||||
pub target_size: usize,
|
||||
/// Max time a warm socket may sit before it is closed and replaced.
|
||||
pub max_idle: Duration,
|
||||
}
|
||||
|
||||
impl Default for PoolConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
target_size: DEFAULT_POOL_SIZE,
|
||||
max_idle: DEFAULT_MAX_IDLE,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl PoolConfig {
|
||||
/// Read config from the environment, falling back to defaults. An invalid
|
||||
/// value warns and uses the default rather than failing startup.
|
||||
pub fn from_env() -> Self {
|
||||
let target_size = match std::env::var(POOL_SIZE_ENV) {
|
||||
Ok(raw) => raw.trim().parse().unwrap_or_else(|_| {
|
||||
eprintln!("warning: {POOL_SIZE_ENV}={raw:?} is not a valid size; using {DEFAULT_POOL_SIZE}");
|
||||
DEFAULT_POOL_SIZE
|
||||
}),
|
||||
Err(_) => DEFAULT_POOL_SIZE,
|
||||
};
|
||||
let max_idle = match std::env::var(MAX_IDLE_ENV) {
|
||||
Ok(raw) => raw
|
||||
.trim()
|
||||
.parse()
|
||||
.map(Duration::from_secs)
|
||||
.unwrap_or_else(|_| {
|
||||
eprintln!(
|
||||
"warning: {MAX_IDLE_ENV}={raw:?} is not a valid number of seconds; using {}s",
|
||||
DEFAULT_MAX_IDLE.as_secs()
|
||||
);
|
||||
DEFAULT_MAX_IDLE
|
||||
}),
|
||||
Err(_) => DEFAULT_MAX_IDLE,
|
||||
};
|
||||
Self {
|
||||
target_size,
|
||||
max_idle,
|
||||
}
|
||||
}
|
||||
|
||||
/// Whether pooling is on (`target_size > 0`).
|
||||
pub fn enabled(&self) -> bool {
|
||||
self.target_size > 0
|
||||
}
|
||||
}
|
||||
|
||||
/// Per-key warm sockets, behind a single `Mutex`. Realtime warm sockets are few
|
||||
/// (the pool is small), so a plain mutex over a `VecDeque`-ish `Vec` is simpler
|
||||
/// and faster than sharding; contention is negligible at this scale.
|
||||
type Warm = HashMap<UpstreamKey, Vec<WarmConnection>>;
|
||||
|
||||
/// Per-key replenish backoff. Absent (or `consecutive_failures == 0`) means the
|
||||
/// key is healthy and replenished every tick. After a pass whose dials all fail,
|
||||
/// `retry_after` is pushed out with exponential backoff so a broken key (invalid
|
||||
/// credentials, unreachable upstream) is not re-dialed on every 250 ms tick.
|
||||
#[derive(Default)]
|
||||
struct Backoff {
|
||||
/// Don't attempt warm-up dials for this key until this instant. `None` =
|
||||
/// eligible now.
|
||||
retry_after: Option<Instant>,
|
||||
consecutive_failures: u32,
|
||||
}
|
||||
|
||||
type Backoffs = HashMap<UpstreamKey, Backoff>;
|
||||
|
||||
/// Pre-warmed upstream realtime connection pool.
|
||||
///
|
||||
/// Cheap to clone-via-`Arc`. The background replenisher is spawned by
|
||||
/// [`RealtimePool::spawn`]; a pool built with [`RealtimePool::disabled`] never
|
||||
/// warms anything and every `take` misses (callers fresh-dial).
|
||||
pub struct RealtimePool {
|
||||
config: PoolConfig,
|
||||
warm: Mutex<Warm>,
|
||||
/// Per-key replenish backoff so a broken key doesn't trigger unbounded
|
||||
/// concurrent dials every tick. Separate lock from `warm` so the request
|
||||
/// hot path (`take`) never contends on it.
|
||||
backoff: Mutex<Backoffs>,
|
||||
}
|
||||
|
||||
impl RealtimePool {
|
||||
/// A disabled pool: no background task, every `take` returns `None`.
|
||||
pub fn disabled() -> Arc<Self> {
|
||||
Arc::new(Self {
|
||||
config: PoolConfig {
|
||||
target_size: 0,
|
||||
..PoolConfig::default()
|
||||
},
|
||||
warm: Mutex::new(HashMap::new()),
|
||||
backoff: Mutex::new(HashMap::new()),
|
||||
})
|
||||
}
|
||||
|
||||
/// Build a pool from config **without** the background replenisher. The pool
|
||||
/// only warms when [`RealtimePool::warm_now`] is called. Used by deterministic
|
||||
/// unit tests; production uses [`RealtimePool::spawn`].
|
||||
#[cfg(test)]
|
||||
fn new_unspawned(config: PoolConfig) -> Arc<Self> {
|
||||
Arc::new(Self {
|
||||
config,
|
||||
warm: Mutex::new(HashMap::new()),
|
||||
backoff: Mutex::new(HashMap::new()),
|
||||
})
|
||||
}
|
||||
|
||||
/// Build a pool from config and, if enabled, spawn the background replenisher.
|
||||
/// Returns the shared handle the gateway stores in its state.
|
||||
pub fn spawn(config: PoolConfig) -> Arc<Self> {
|
||||
let pool = Arc::new(Self {
|
||||
config,
|
||||
warm: Mutex::new(HashMap::new()),
|
||||
backoff: Mutex::new(HashMap::new()),
|
||||
});
|
||||
if config.enabled() {
|
||||
let weak = Arc::downgrade(&pool);
|
||||
tokio::spawn(async move {
|
||||
let mut tick = tokio::time::interval(REPLENISH_TICK);
|
||||
tick.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
|
||||
loop {
|
||||
tick.tick().await;
|
||||
// Stop once the gateway has dropped its handle.
|
||||
let Some(pool) = weak.upgrade() else { break };
|
||||
pool.replenish_all().await;
|
||||
}
|
||||
});
|
||||
}
|
||||
pool
|
||||
}
|
||||
|
||||
/// Resolved config (test/inspection).
|
||||
pub fn config(&self) -> PoolConfig {
|
||||
self.config
|
||||
}
|
||||
|
||||
/// Register a key so the replenisher starts warming it. Idempotent. The
|
||||
/// gateway calls this once per known deployment at startup; the pool only
|
||||
/// warms keys it has seen, so it never dials a model nobody asked for.
|
||||
pub fn register(&self, key: UpstreamKey) {
|
||||
if !self.config.enabled() {
|
||||
return;
|
||||
}
|
||||
self.warm.lock().unwrap().entry(key).or_default();
|
||||
}
|
||||
|
||||
/// Take a warm, live socket for `key`, or `None` on miss / dead socket.
|
||||
///
|
||||
/// Pops the freshest non-expired socket and liveness-checks it; a socket that
|
||||
/// is too old or already dead is dropped (closing it) and the next candidate
|
||||
/// tried. Never blocks: if nothing warm is live, returns `None` so the caller
|
||||
/// fresh-dials.
|
||||
pub fn take(&self, key: &UpstreamKey) -> Option<WarmHandoff> {
|
||||
if !self.config.enabled() {
|
||||
return None;
|
||||
}
|
||||
loop {
|
||||
let mut candidate = {
|
||||
let mut warm = self.warm.lock().unwrap();
|
||||
let bucket = warm.get_mut(key)?;
|
||||
bucket.pop()?
|
||||
};
|
||||
// Discard sockets past their warm lifetime (idle-billing guard).
|
||||
if candidate.warmed_at.elapsed() > self.config.max_idle {
|
||||
continue; // drops `candidate`, closing the socket
|
||||
}
|
||||
// Liveness: a non-blocking check that the socket hasn't already
|
||||
// delivered a Close/Err. A warm socket should be silent after
|
||||
// session.created, so anything pending means it is unhealthy.
|
||||
if is_dead(&mut candidate.rx) {
|
||||
continue;
|
||||
}
|
||||
return Some(WarmHandoff {
|
||||
tx: candidate.tx,
|
||||
rx: candidate.rx,
|
||||
session_created: candidate.session_created,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
/// One replenish pass over every registered key: reap stale sockets, then
|
||||
/// dial up to `target_size`. Dials run concurrently; failures are swallowed
|
||||
/// (a key that can't be warmed just keeps fresh-dialing on the request path)
|
||||
/// and put the key into exponential backoff so a broken key isn't re-dialed
|
||||
/// on every tick.
|
||||
async fn replenish_all(&self) {
|
||||
let keys: Vec<UpstreamKey> = { self.warm.lock().unwrap().keys().cloned().collect() };
|
||||
for key in keys {
|
||||
self.reap_stale(&key);
|
||||
// Skip keys still in backoff from a prior all-failed pass — this is
|
||||
// what bounds dials against an invalid/unreachable key to once per
|
||||
// `BACKOFF_MAX` instead of `needed` dials every 250 ms tick.
|
||||
if self.in_backoff(&key) {
|
||||
continue;
|
||||
}
|
||||
let needed = {
|
||||
let warm = self.warm.lock().unwrap();
|
||||
let have = warm.get(&key).map(Vec::len).unwrap_or(0);
|
||||
self.config.target_size.saturating_sub(have)
|
||||
};
|
||||
if needed == 0 {
|
||||
continue;
|
||||
}
|
||||
// Dial the missing sockets CONCURRENTLY. A sequential loop here makes
|
||||
// a full refill cost `needed × handshake` (~needed × 350 ms), which
|
||||
// can't keep up with a high connect rate — the pool drains faster
|
||||
// than it refills and most connects miss. Firing the dials together
|
||||
// refills in ~one handshake window, keeping warm supply ≈ peak
|
||||
// concurrent connects so the sub-ms warm handoff becomes the median,
|
||||
// not the lucky-hit tail.
|
||||
let dials = (0..needed).map(|_| warm_one(&key));
|
||||
let results = futures_util::future::join_all(dials).await;
|
||||
let mut any_ok = false;
|
||||
// `.flatten()` keeps only the successful dials; a key that can't be
|
||||
// warmed just keeps fresh-dialing on the request path.
|
||||
for conn in results.into_iter().flatten() {
|
||||
any_ok = true;
|
||||
self.warm
|
||||
.lock()
|
||||
.unwrap()
|
||||
.entry(key.clone())
|
||||
.or_default()
|
||||
.push(conn);
|
||||
}
|
||||
// Reset backoff on any success; otherwise grow it. We only ever enter
|
||||
// backoff when a pass that *attempted* dials produced none — a `needed
|
||||
// == 0` pass is handled by the `continue` above and never touches it.
|
||||
self.record_replenish_outcome(&key, any_ok);
|
||||
}
|
||||
}
|
||||
|
||||
/// Whether `key` is currently in a backoff window (a prior pass failed and
|
||||
/// the retry time hasn't arrived). Eligible keys are pruned from the backoff
|
||||
/// map so it doesn't grow unbounded for healthy keys.
|
||||
fn in_backoff(&self, key: &UpstreamKey) -> bool {
|
||||
let mut backoff = self.backoff.lock().unwrap();
|
||||
match backoff.get(key).and_then(|b| b.retry_after) {
|
||||
Some(retry_after) if Instant::now() < retry_after => true,
|
||||
Some(_) => {
|
||||
// Window elapsed — allow the attempt. Keep the failure count so a
|
||||
// still-broken key backs off further, but clear the gate so this
|
||||
// tick proceeds.
|
||||
if let Some(b) = backoff.get_mut(key) {
|
||||
b.retry_after = None;
|
||||
}
|
||||
false
|
||||
}
|
||||
None => false,
|
||||
}
|
||||
}
|
||||
|
||||
/// Update a key's backoff after a replenish attempt. Success clears it;
|
||||
/// failure grows the retry delay exponentially up to `BACKOFF_MAX`.
|
||||
fn record_replenish_outcome(&self, key: &UpstreamKey, any_ok: bool) {
|
||||
let mut backoff = self.backoff.lock().unwrap();
|
||||
if any_ok {
|
||||
backoff.remove(key);
|
||||
return;
|
||||
}
|
||||
let entry = backoff.entry(key.clone()).or_default();
|
||||
entry.consecutive_failures = entry.consecutive_failures.saturating_add(1);
|
||||
// Exponential: BASE * 2^(failures-1), saturating at MAX. `min` of the
|
||||
// shift exponent keeps the doubling from overflowing.
|
||||
let shift = (entry.consecutive_failures - 1).min(16);
|
||||
let delay = BACKOFF_BASE.saturating_mul(1u32 << shift).min(BACKOFF_MAX);
|
||||
entry.retry_after = Some(Instant::now() + delay);
|
||||
}
|
||||
|
||||
/// Drop sockets past `max_idle` or already dead for a key.
|
||||
fn reap_stale(&self, key: &UpstreamKey) {
|
||||
let mut warm = self.warm.lock().unwrap();
|
||||
if let Some(bucket) = warm.get_mut(key) {
|
||||
bucket.retain_mut(|conn| {
|
||||
conn.warmed_at.elapsed() <= self.config.max_idle && !is_dead(&mut conn.rx)
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
/// Test/inspection: number of warm sockets currently held for `key`.
|
||||
#[cfg(test)]
|
||||
pub fn warm_len(&self, key: &UpstreamKey) -> usize {
|
||||
self.warm
|
||||
.lock()
|
||||
.unwrap()
|
||||
.get(key)
|
||||
.map(Vec::len)
|
||||
.unwrap_or(0)
|
||||
}
|
||||
|
||||
/// Test/inspection: consecutive replenish failures recorded for `key` (0 if
|
||||
/// the key is healthy / has no backoff entry).
|
||||
#[cfg(test)]
|
||||
pub fn backoff_failures(&self, key: &UpstreamKey) -> u32 {
|
||||
self.backoff
|
||||
.lock()
|
||||
.unwrap()
|
||||
.get(key)
|
||||
.map(|b| b.consecutive_failures)
|
||||
.unwrap_or(0)
|
||||
}
|
||||
|
||||
/// Test helper: synchronously warm `target_size` sockets for `key` (no
|
||||
/// background task). Lets tests assert handoff behavior deterministically.
|
||||
#[cfg(test)]
|
||||
pub async fn warm_now(&self, key: &UpstreamKey) {
|
||||
let needed = {
|
||||
let warm = self.warm.lock().unwrap();
|
||||
let have = warm.get(key).map(Vec::len).unwrap_or(0);
|
||||
self.config.target_size.saturating_sub(have)
|
||||
};
|
||||
for _ in 0..needed {
|
||||
if let Ok(conn) = warm_one(key).await {
|
||||
self.warm
|
||||
.lock()
|
||||
.unwrap()
|
||||
.entry(key.clone())
|
||||
.or_default()
|
||||
.push(conn);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Test helper: insert an already-built warm connection (used to inject a
|
||||
/// dead socket and assert it is discarded at handoff).
|
||||
#[cfg(test)]
|
||||
fn insert_warm(&self, key: UpstreamKey, conn: WarmConnection) {
|
||||
self.warm.lock().unwrap().entry(key).or_default().push(conn);
|
||||
}
|
||||
}
|
||||
|
||||
/// Dial one upstream and pre-read its `session.created` into a [`WarmConnection`].
|
||||
///
|
||||
/// `key.api_key` is already resolved (non-blank). The first frame OpenAI sends
|
||||
/// unprompted is `session.created`; we buffer exactly that and read nothing more.
|
||||
async fn warm_one(key: &UpstreamKey) -> CoreResult<WarmConnection> {
|
||||
let upstream: UpstreamWs =
|
||||
dial_upstream(&key.model, &key.api_key, key.api_base.as_deref()).await?;
|
||||
let (tx, mut rx) = upstream.split();
|
||||
let session_created = read_event(&mut rx).await?;
|
||||
Ok(WarmConnection {
|
||||
tx,
|
||||
rx,
|
||||
session_created,
|
||||
warmed_at: Instant::now(),
|
||||
})
|
||||
}
|
||||
|
||||
/// Resolve a deployment's API key into the pool key, returning `None` when no key
|
||||
/// can be resolved (those deployments simply aren't pooled — the request path
|
||||
/// still fresh-dials and surfaces the auth error there).
|
||||
pub fn upstream_key(
|
||||
model: &str,
|
||||
api_key: Option<&str>,
|
||||
api_base: Option<&str>,
|
||||
) -> Option<UpstreamKey> {
|
||||
let api_key = resolve_api_key(api_key).ok()?;
|
||||
Some(UpstreamKey {
|
||||
model: model.to_string(),
|
||||
api_key,
|
||||
api_base: api_base.map(str::to_string),
|
||||
})
|
||||
}
|
||||
|
||||
/// Non-blocking liveness check: poll the upstream once. A warm socket is silent
|
||||
/// after `session.created`, so a pending `Close`/`Err`/`None` means it is dead.
|
||||
/// A pending data frame (shouldn't happen pre-handoff) is also treated as
|
||||
/// unhealthy — we'd rather discard and fresh-dial than hand over a socket in an
|
||||
/// unexpected state. `Pending` (the healthy case) returns `false`.
|
||||
fn is_dead(rx: &mut UpstreamRx) -> bool {
|
||||
use futures_util::task::noop_waker_ref;
|
||||
use futures_util::Stream;
|
||||
use std::pin::Pin;
|
||||
use std::task::{Context, Poll};
|
||||
|
||||
let mut cx = Context::from_waker(noop_waker_ref());
|
||||
match Pin::new(rx).poll_next(&mut cx) {
|
||||
Poll::Pending => false,
|
||||
Poll::Ready(None) => true,
|
||||
Poll::Ready(Some(Err(_))) => true,
|
||||
// Any frame arriving before handoff is unexpected for a silent warm
|
||||
// socket; treat it as unhealthy.
|
||||
Poll::Ready(Some(Ok(_))) => true,
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use futures_util::SinkExt;
|
||||
use std::net::SocketAddr;
|
||||
use tokio::net::TcpListener;
|
||||
use tokio_tungstenite::tungstenite::Message;
|
||||
|
||||
/// An in-process fake OpenAI realtime WS server. On connect it sends
|
||||
/// `session.created`; on `response.create` it sends `response.created` +
|
||||
/// `response.output_audio.delta` + `response.done`. Returns its `ws://` base.
|
||||
async fn spawn_fake_openai() -> String {
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let addr: SocketAddr = listener.local_addr().unwrap();
|
||||
tokio::spawn(async move {
|
||||
while let Ok((stream, _)) = listener.accept().await {
|
||||
tokio::spawn(handle_fake_conn(stream));
|
||||
}
|
||||
});
|
||||
format!("ws://{addr}")
|
||||
}
|
||||
|
||||
async fn handle_fake_conn(stream: tokio::net::TcpStream) {
|
||||
let mut ws = match tokio_tungstenite::accept_async(stream).await {
|
||||
Ok(ws) => ws,
|
||||
Err(_) => return,
|
||||
};
|
||||
// Unprompted session.created, exactly like OpenAI.
|
||||
let _ = ws
|
||||
.send(Message::Text(
|
||||
r#"{"type":"session.created","session":{"id":"sess_fake"}}"#.to_string(),
|
||||
))
|
||||
.await;
|
||||
while let Some(Ok(msg)) = ws.next().await {
|
||||
if let Message::Text(text) = msg {
|
||||
if text.contains("response.create") {
|
||||
for frame in [
|
||||
r#"{"type":"response.created"}"#,
|
||||
r#"{"type":"response.output_audio.delta","delta":"AAAA"}"#,
|
||||
r#"{"type":"response.done"}"#,
|
||||
] {
|
||||
let _ = ws.send(Message::Text(frame.to_string())).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn test_config() -> PoolConfig {
|
||||
PoolConfig {
|
||||
target_size: 2,
|
||||
max_idle: Duration::from_secs(30),
|
||||
}
|
||||
}
|
||||
|
||||
fn key_for(base: &str) -> UpstreamKey {
|
||||
UpstreamKey {
|
||||
model: "gpt-realtime".to_string(),
|
||||
api_key: "sk-test".to_string(),
|
||||
api_base: Some(base.to_string()),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn warm_handoff_relays_buffered_session_created() {
|
||||
let base = spawn_fake_openai().await;
|
||||
let pool = RealtimePool::new_unspawned(test_config());
|
||||
let key = key_for(&base);
|
||||
pool.register(key.clone());
|
||||
pool.warm_now(&key).await;
|
||||
assert_eq!(pool.warm_len(&key), 2);
|
||||
|
||||
let handoff = pool.take(&key).expect("a warm socket should be available");
|
||||
assert_eq!(handoff.session_created.event_type, "session.created");
|
||||
assert_eq!(
|
||||
handoff
|
||||
.session_created
|
||||
.data
|
||||
.get("session")
|
||||
.and_then(|s| s.get("id"))
|
||||
.and_then(|v| v.as_str()),
|
||||
Some("sess_fake")
|
||||
);
|
||||
// Taking one leaves one.
|
||||
assert_eq!(pool.warm_len(&key), 1);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn pool_miss_returns_none_for_fresh_dial_fallback() {
|
||||
let base = spawn_fake_openai().await;
|
||||
let pool = RealtimePool::new_unspawned(test_config());
|
||||
let key = key_for(&base);
|
||||
// Registered but never warmed → empty bucket → miss.
|
||||
pool.register(key.clone());
|
||||
assert!(pool.take(&key).is_none());
|
||||
|
||||
// Unknown key → miss.
|
||||
let other = key_for("ws://127.0.0.1:1");
|
||||
assert!(pool.take(&other).is_none());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn disabled_pool_never_hands_off() {
|
||||
let pool = RealtimePool::disabled();
|
||||
let key = key_for("ws://127.0.0.1:1");
|
||||
pool.register(key.clone());
|
||||
assert_eq!(pool.warm_len(&key), 0);
|
||||
assert!(pool.take(&key).is_none());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn dead_warm_socket_is_discarded() {
|
||||
let base = spawn_fake_openai().await;
|
||||
let pool = RealtimePool::new_unspawned(test_config());
|
||||
let key = key_for(&base);
|
||||
pool.register(key.clone());
|
||||
|
||||
// Build one real warm connection, then kill the upstream by dropping the
|
||||
// server side: easiest is to dial, read session.created, then close our
|
||||
// own rx's peer. Instead we forge "dead" via an already-closed socket:
|
||||
// dial a connection and immediately send a Close from the client side so
|
||||
// the server closes back, then warm it. Simpler: warm normally, then
|
||||
// mark it stale by backdating warmed_at past max_idle and confirm it's
|
||||
// dropped — that exercises the same discard path.
|
||||
let mut conn = warm_one(&key).await.expect("warm one");
|
||||
conn.warmed_at = Instant::now() - Duration::from_secs(3600); // past max_idle
|
||||
pool.insert_warm(key.clone(), conn);
|
||||
assert_eq!(pool.warm_len(&key), 1);
|
||||
|
||||
// take() must discard the stale socket and report a miss.
|
||||
assert!(pool.take(&key).is_none());
|
||||
assert_eq!(pool.warm_len(&key), 0);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn background_replenisher_tops_up_registered_key() {
|
||||
let base = spawn_fake_openai().await;
|
||||
let pool = RealtimePool::spawn(test_config());
|
||||
let key = key_for(&base);
|
||||
pool.register(key.clone());
|
||||
|
||||
// Wait (bounded) for the background task to reach the target size.
|
||||
let mut warmed = 0;
|
||||
for _ in 0..40 {
|
||||
tokio::time::sleep(Duration::from_millis(50)).await;
|
||||
warmed = pool.warm_len(&key);
|
||||
if warmed >= test_config().target_size {
|
||||
break;
|
||||
}
|
||||
}
|
||||
assert_eq!(
|
||||
warmed,
|
||||
test_config().target_size,
|
||||
"background replenisher should warm up to target_size"
|
||||
);
|
||||
let handoff = pool.take(&key).expect("a warm socket should be available");
|
||||
assert_eq!(handoff.session_created.event_type, "session.created");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn closed_upstream_socket_is_detected_dead() {
|
||||
// A genuinely dead socket: dial the fake, read session.created, then drop
|
||||
// the server by closing from our side and waiting for the close to land.
|
||||
let base = spawn_fake_openai().await;
|
||||
let pool = RealtimePool::new_unspawned(test_config());
|
||||
let key = key_for(&base);
|
||||
pool.register(key.clone());
|
||||
|
||||
let mut conn = warm_one(&key).await.expect("warm one");
|
||||
// Close the upstream from the client side; the server echoes a close.
|
||||
let _ = conn.tx.send(Message::Close(None)).await;
|
||||
// Give the close a moment to arrive on rx.
|
||||
tokio::time::sleep(Duration::from_millis(50)).await;
|
||||
pool.insert_warm(key.clone(), conn);
|
||||
|
||||
// Liveness check at take() should detect the close and discard it.
|
||||
assert!(pool.take(&key).is_none());
|
||||
assert_eq!(pool.warm_len(&key), 0);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn broken_key_backs_off_instead_of_dialing_every_tick() {
|
||||
// A key whose upstream is unreachable: every warm-up dial fails.
|
||||
let pool = RealtimePool::new_unspawned(test_config());
|
||||
let key = key_for("ws://127.0.0.1:1"); // nothing listens here
|
||||
pool.register(key.clone());
|
||||
|
||||
// First pass attempts dials, they all fail → key enters backoff, no warm
|
||||
// sockets, one recorded failure.
|
||||
pool.replenish_all().await;
|
||||
assert_eq!(pool.warm_len(&key), 0);
|
||||
assert_eq!(pool.backoff_failures(&key), 1);
|
||||
assert!(
|
||||
pool.in_backoff(&key),
|
||||
"a key whose dials all failed must be in backoff"
|
||||
);
|
||||
|
||||
// An immediate next pass must be SKIPPED (still in the backoff window), so
|
||||
// it does NOT fire another round of dials — the failure count is unchanged.
|
||||
pool.replenish_all().await;
|
||||
assert_eq!(
|
||||
pool.backoff_failures(&key),
|
||||
1,
|
||||
"replenish during the backoff window must not re-dial the broken key"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn healthy_key_never_enters_backoff_and_clears_after_recovery() {
|
||||
let base = spawn_fake_openai().await;
|
||||
let pool = RealtimePool::new_unspawned(test_config());
|
||||
let key = key_for(&base);
|
||||
pool.register(key.clone());
|
||||
|
||||
// A reachable upstream: the pass succeeds, so the key is never backed off.
|
||||
pool.replenish_all().await;
|
||||
assert_eq!(pool.warm_len(&key), test_config().target_size);
|
||||
assert_eq!(pool.backoff_failures(&key), 0);
|
||||
assert!(!pool.in_backoff(&key));
|
||||
}
|
||||
}
|
||||
28
litellm-rust/crates/ai-gateway/src/lib.rs
Normal file
28
litellm-rust/crates/ai-gateway/src/lib.rs
Normal file
|
|
@ -0,0 +1,28 @@
|
|||
//! LiteLLM AI Gateway library.
|
||||
//!
|
||||
//! Two layers, split by feature so the Python `cdylib` can depend on the I/O
|
||||
//! without pulling in the HTTP server:
|
||||
//!
|
||||
//! - [`io`]: all network I/O (OCR HTTP call, realtime WebSocket splice, the
|
||||
//! pre-warmed realtime pool). Always available — no feature required. The
|
||||
//! Python bridge links this for `run_ocr`.
|
||||
//! - The server modules ([`auth`], [`routes`], [`state`]) and anything pulling
|
||||
//! `axum` are gated behind the `server` feature, which the `litellm-ai-gateway`
|
||||
//! binary turns on. The `python-config` feature additionally pulls in [`python`]
|
||||
//! for the load-time config reader.
|
||||
|
||||
pub mod io;
|
||||
|
||||
/// GIL-activity tracking. Pure (atomics only); shared by the `server` routes and
|
||||
/// the `python-config` reader, so it is available without either feature.
|
||||
pub mod gil;
|
||||
|
||||
#[cfg(feature = "server")]
|
||||
pub mod auth;
|
||||
#[cfg(feature = "server")]
|
||||
pub mod routes;
|
||||
#[cfg(feature = "server")]
|
||||
pub mod state;
|
||||
|
||||
#[cfg(feature = "python-config")]
|
||||
pub mod python;
|
||||
153
litellm-rust/crates/ai-gateway/src/main.rs
Normal file
153
litellm-rust/crates/ai-gateway/src/main.rs
Normal file
|
|
@ -0,0 +1,153 @@
|
|||
//! LiteLLM AI Gateway — a minimal Axum server fronting the Rust router.
|
||||
//!
|
||||
//! Flow: client → `POST /v1/realtime` → `router.realtime()` selects a deployment
|
||||
//! (simple-shuffle) → `io::realtime::realtime()` invokes OpenAI. The
|
||||
//! server owns transport + config; routing lives in the `router` crate.
|
||||
//!
|
||||
//! The binary requires the `server` feature (declared in `Cargo.toml` via
|
||||
//! `required-features`), so cargo skips it unless that feature is on. Everything
|
||||
//! the binary needs lives in the library (`litellm_ai_gateway`); `main` just
|
||||
//! wires startup.
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use litellm_ai_gateway::io::realtime_pool::{upstream_key, PoolConfig, RealtimePool};
|
||||
use litellm_ai_gateway::routes;
|
||||
use litellm_ai_gateway::state::AppState;
|
||||
use litellm_core::router::{Deployment, LiteLLMParams, Router};
|
||||
|
||||
#[cfg(feature = "python-config")]
|
||||
use litellm_ai_gateway::python;
|
||||
|
||||
/// Bind to localhost by default so the gateway is not a public, unauthenticated
|
||||
/// provider proxy out of the box. Override with `HOST` (e.g. `0.0.0.0`).
|
||||
const DEFAULT_HOST: &str = "127.0.0.1";
|
||||
const DEFAULT_PORT: u16 = 4001;
|
||||
|
||||
#[tokio::main]
|
||||
async fn main() {
|
||||
// Trim before storing so it matches the trimmed bearer token in `auth`
|
||||
// (avoids a silent auth failure when the env var has surrounding whitespace).
|
||||
let master_key: Option<Arc<str>> = std::env::var("LITELLM_MASTER_KEY")
|
||||
.ok()
|
||||
.map(|key| key.trim().to_string())
|
||||
.filter(|key| !key.is_empty())
|
||||
.map(Arc::from);
|
||||
if master_key.is_none() {
|
||||
eprintln!(
|
||||
"warning: LITELLM_MASTER_KEY is not set; /v1/realtime will reject all requests (fail closed)"
|
||||
);
|
||||
}
|
||||
|
||||
let router = Arc::new(build_router());
|
||||
|
||||
// Build the pre-warmed realtime pool and register each deployment's upstream
|
||||
// so the background replenisher starts warming it. `REALTIME_POOL_SIZE=0`
|
||||
// yields a disabled pool → every connect fresh-dials (original behavior).
|
||||
let pool_config = PoolConfig::from_env();
|
||||
let realtime_pool = RealtimePool::spawn(pool_config);
|
||||
if pool_config.enabled() {
|
||||
register_deployments(&router, &realtime_pool);
|
||||
eprintln!(
|
||||
"realtime connection pool enabled: target {} warm sockets/key, max idle {}s",
|
||||
pool_config.target_size,
|
||||
pool_config.max_idle.as_secs()
|
||||
);
|
||||
} else {
|
||||
eprintln!(
|
||||
"realtime connection pool disabled (REALTIME_POOL_SIZE=0); fresh-dialing each connect"
|
||||
);
|
||||
}
|
||||
|
||||
let state = AppState {
|
||||
router,
|
||||
master_key,
|
||||
realtime_pool,
|
||||
};
|
||||
|
||||
let host = std::env::var("HOST").unwrap_or_else(|_| DEFAULT_HOST.to_string());
|
||||
let port = resolve_port();
|
||||
|
||||
let listener = tokio::net::TcpListener::bind((host.as_str(), port))
|
||||
.await
|
||||
.expect("failed to bind listener");
|
||||
eprintln!("litellm-ai-gateway listening on {host}:{port}");
|
||||
axum::serve(listener, routes::app(state))
|
||||
.await
|
||||
.expect("server error");
|
||||
}
|
||||
|
||||
/// Register every deployment's upstream key with the pool so the replenisher
|
||||
/// pre-warms it. Mirrors `service::run`'s key derivation (strip `openai/`, resolve
|
||||
/// api_key); deployments whose key can't be resolved are skipped (they fresh-dial
|
||||
/// and surface the auth error on the request path, as before).
|
||||
fn register_deployments(router: &Router, pool: &RealtimePool) {
|
||||
for deployment in router.deployments() {
|
||||
let params = &deployment.litellm_params;
|
||||
let provider_model = params
|
||||
.model
|
||||
.strip_prefix("openai/")
|
||||
.unwrap_or(¶ms.model);
|
||||
if let Some(key) = upstream_key(
|
||||
provider_model,
|
||||
params.api_key.as_deref(),
|
||||
params.api_base.as_deref(),
|
||||
) {
|
||||
pool.register(key);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Resolve `PORT`, warning (rather than silently defaulting) on an invalid value.
|
||||
fn resolve_port() -> u16 {
|
||||
match std::env::var("PORT") {
|
||||
Ok(raw) => raw.parse().unwrap_or_else(|_| {
|
||||
eprintln!("warning: PORT={raw:?} is not a valid port; using {DEFAULT_PORT}");
|
||||
DEFAULT_PORT
|
||||
}),
|
||||
Err(_) => DEFAULT_PORT,
|
||||
}
|
||||
}
|
||||
|
||||
/// Build the router. With the `python-config` feature and `LITELLM_CONFIG_PATH`
|
||||
/// set, load the resolved `model_list` from the proxy config via the embedded
|
||||
/// Python reader (load time only). Otherwise fall back to the env stand-in.
|
||||
fn build_router() -> Router {
|
||||
#[cfg(feature = "python-config")]
|
||||
if let Ok(config_path) = std::env::var("LITELLM_CONFIG_PATH") {
|
||||
match python::config::load_router_from_config(&config_path) {
|
||||
Ok(router) => {
|
||||
eprintln!("loaded model_list from {config_path} via python config reader");
|
||||
return router;
|
||||
}
|
||||
Err(err) => {
|
||||
eprintln!("config load failed ({err}); falling back to env deployment");
|
||||
}
|
||||
}
|
||||
}
|
||||
build_router_from_env()
|
||||
}
|
||||
|
||||
/// Build a minimal single-deployment `model_list` from the environment.
|
||||
///
|
||||
/// A real deployment loads `model_list` from config; this is the minimal stand-in
|
||||
/// so the gateway has one OpenAI deployment to route to.
|
||||
fn build_router_from_env() -> Router {
|
||||
let model =
|
||||
std::env::var("OPENAI_REALTIME_MODEL").unwrap_or_else(|_| "gpt-realtime".to_string());
|
||||
let api_key = std::env::var("OPENAI_API_KEY").ok();
|
||||
if api_key.is_none() {
|
||||
eprintln!(
|
||||
"warning: OPENAI_API_KEY is not set; realtime requests will fail with auth errors"
|
||||
);
|
||||
}
|
||||
let deployment = Deployment {
|
||||
model_name: model.clone(),
|
||||
litellm_params: LiteLLMParams {
|
||||
model,
|
||||
api_key,
|
||||
api_base: None,
|
||||
},
|
||||
};
|
||||
Router::new(vec![deployment])
|
||||
}
|
||||
27
litellm-rust/crates/ai-gateway/src/python/AGENTS.md
Normal file
27
litellm-rust/crates/ai-gateway/src/python/AGENTS.md
Normal file
|
|
@ -0,0 +1,27 @@
|
|||
# ai-gateway/src/python — Python interop (load-time only)
|
||||
|
||||
Functions here embed the Python interpreter (pyo3) and take the GIL to call into
|
||||
`litellm` (e.g. read the proxy `model_list`). Compiled only under the
|
||||
`python-config` feature.
|
||||
|
||||
## Hard rule: non-hot-path functions only
|
||||
|
||||
Everything in this folder MUST run **at most once per process lifetime — at
|
||||
startup / load time** (config read, warm-up). NEVER call into Python on the
|
||||
request path:
|
||||
|
||||
- No GIL acquisition per request, per connection, or per realtime event.
|
||||
- No Python call inside a route handler, the router's hot path, or any loop that
|
||||
scales with traffic.
|
||||
|
||||
**Why:** the GIL serializes execution and would cap throughput; the realtime data
|
||||
path must stay pure Rust. Every acquisition is recorded by `crate::gil` — poll
|
||||
`GET /health/gil`, and `total_acquisitions` MUST stay flat under load.
|
||||
|
||||
## How to add one
|
||||
|
||||
Resolve whatever Python-derived data you need **once at boot** and hand the rest
|
||||
of the gateway an owned, plain-Rust value (e.g. build a `Router` from the
|
||||
resolved `model_list`). Record the acquisition via `crate::gil::record_acquisition()`
|
||||
immediately before taking the GIL. If a function would need to run per request,
|
||||
it does not belong here — move the work to Rust, or pre-resolve it at startup.
|
||||
39
litellm-rust/crates/ai-gateway/src/python/config.rs
Normal file
39
litellm-rust/crates/ai-gateway/src/python/config.rs
Normal file
|
|
@ -0,0 +1,39 @@
|
|||
//! Build the router by calling the Python proxy config reader (load time only).
|
||||
//!
|
||||
//! Embeds the interpreter via pyo3 and calls
|
||||
//! `litellm.proxy.read_model_list.read_model_list`, which reuses the proxy's
|
||||
//! `os.environ/` + secret-manager resolution. The GIL is taken **once at boot**
|
||||
//! (and recorded in [`crate::gil`]); the realtime hot path never touches Python.
|
||||
//!
|
||||
//! Compiled only under the `python-config` feature.
|
||||
|
||||
use litellm_core::error::CoreError;
|
||||
use litellm_core::router::{Deployment, Router};
|
||||
use litellm_core::CoreResult;
|
||||
use pyo3::prelude::*;
|
||||
|
||||
use crate::gil;
|
||||
|
||||
/// Load the router's `model_list` from `config_path` via the Python reader.
|
||||
pub fn load_router_from_config(config_path: &str) -> CoreResult<Router> {
|
||||
gil::record_acquisition();
|
||||
Python::with_gil(|py| {
|
||||
let model_list = py
|
||||
.import("litellm.proxy.read_model_list")
|
||||
.and_then(|module| module.getattr("read_model_list"))
|
||||
.and_then(|reader| reader.call1((config_path,)))
|
||||
.map_err(|err| CoreError::Routing(format!("read_model_list failed: {err}")))?;
|
||||
|
||||
let model_list_json: String = py
|
||||
.import("json")
|
||||
.and_then(|json| json.getattr("dumps"))
|
||||
.and_then(|dumps| dumps.call1((model_list,)))
|
||||
.and_then(|encoded| encoded.extract())
|
||||
.map_err(|err| CoreError::Routing(format!("serializing model_list failed: {err}")))?;
|
||||
|
||||
let deployments: Vec<Deployment> = serde_json::from_str(&model_list_json)
|
||||
.map_err(|err| CoreError::Routing(format!("parsing model_list failed: {err}")))?;
|
||||
|
||||
Ok(Router::new(deployments))
|
||||
})
|
||||
}
|
||||
4
litellm-rust/crates/ai-gateway/src/python/mod.rs
Normal file
4
litellm-rust/crates/ai-gateway/src/python/mod.rs
Normal file
|
|
@ -0,0 +1,4 @@
|
|||
//! Python interop for the gateway. See `AGENTS.md`: **load-time / non-hot-path
|
||||
//! only.** Compiled only under the `python-config` feature.
|
||||
|
||||
pub mod config;
|
||||
38
litellm-rust/crates/ai-gateway/src/routes/AGENTS.md
Normal file
38
litellm-rust/crates/ai-gateway/src/routes/AGENTS.md
Normal file
|
|
@ -0,0 +1,38 @@
|
|||
# routes/ — the route template
|
||||
|
||||
Every route follows the **same shape** so the layout is predictable. The rule:
|
||||
|
||||
> **Each route module exposes `pub fn router() -> Router<AppState>`.**
|
||||
> `routes/mod.rs::app` merges them all and applies state once. Adding a route is:
|
||||
> create the module, then add one `.merge(<name>::router())` line.
|
||||
|
||||
## Default: one file
|
||||
A route is a single file containing `router()` + its handler(s) (handlers stay
|
||||
private). This is the norm — don't split until it hurts.
|
||||
```
|
||||
pub fn router() -> Router<AppState> { Router::new().route(PATH, get(handle)) }
|
||||
async fn handle(...) -> impl IntoResponse { ... }
|
||||
```
|
||||
`health.rs` and `gil.rs` are examples.
|
||||
|
||||
## Split out `service` when there's real logic
|
||||
When a route has business logic worth testing without axum, put it in a sibling
|
||||
`service` (a file, or a folder if the route grows). The route file stays the
|
||||
**axum surface** (router + handler + any socket/SSE adapter); `service` is plain
|
||||
Rust with **no axum types**. `realtime/` is the example:
|
||||
```
|
||||
realtime/
|
||||
mod.rs # axum surface: router() + handler + the WS<->events adapter
|
||||
service.rs # pure logic: select deployment + call provider (no axum) — testable
|
||||
```
|
||||
Split `service` further (or add `transport`, `repo`, …) only once a single file
|
||||
genuinely gets hard to read.
|
||||
|
||||
## Invariants
|
||||
- **Auth is an extractor, not a manual call.** A handler requires auth by adding
|
||||
`crate::auth::RequireMasterKey` to its arguments; it runs during extraction.
|
||||
Never re-implement the check per route.
|
||||
- **Handlers contain no business logic; `service` contains no axum types.**
|
||||
- A route owns its paths in its own `router()`; `mod.rs` only merges.
|
||||
- Cross-cutting concerns (logging, CORS, timeouts) → Tower layers in `mod.rs`,
|
||||
not duplicated in handlers.
|
||||
30
litellm-rust/crates/ai-gateway/src/routes/gil.rs
Normal file
30
litellm-rust/crates/ai-gateway/src/routes/gil.rs
Normal file
|
|
@ -0,0 +1,30 @@
|
|||
//! `GET /health/gil` — poll to confirm Python is only touched at load time.
|
||||
//! Simple-route template: a `router()` plus its handler, in one file.
|
||||
|
||||
use axum::routing::get;
|
||||
use axum::{Json, Router};
|
||||
use serde::Serialize;
|
||||
|
||||
use crate::gil;
|
||||
use crate::state::AppState;
|
||||
|
||||
/// This route's contribution to the app router.
|
||||
pub fn router() -> Router<AppState> {
|
||||
Router::new().route("/health/gil", get(status))
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
struct GilStatusResponse {
|
||||
gil_acquired_last_30s: bool,
|
||||
total_acquisitions: u64,
|
||||
seconds_since_last: Option<u64>,
|
||||
}
|
||||
|
||||
async fn status() -> Json<GilStatusResponse> {
|
||||
let snapshot = gil::snapshot();
|
||||
Json(GilStatusResponse {
|
||||
gil_acquired_last_30s: snapshot.acquired_last_30s,
|
||||
total_acquisitions: snapshot.total_acquisitions,
|
||||
seconds_since_last: snapshot.seconds_since_last,
|
||||
})
|
||||
}
|
||||
24
litellm-rust/crates/ai-gateway/src/routes/health.rs
Normal file
24
litellm-rust/crates/ai-gateway/src/routes/health.rs
Normal file
|
|
@ -0,0 +1,24 @@
|
|||
//! Health probes. Simple-route template: a `router()` plus its handlers, in one file.
|
||||
|
||||
use axum::http::StatusCode;
|
||||
use axum::routing::get;
|
||||
use axum::Router;
|
||||
|
||||
use crate::state::AppState;
|
||||
|
||||
/// This route's contribution to the app router.
|
||||
pub fn router() -> Router<AppState> {
|
||||
Router::new()
|
||||
.route("/health/liveness", get(liveness))
|
||||
.route("/health/readiness", get(readiness))
|
||||
}
|
||||
|
||||
/// The process is up.
|
||||
async fn liveness() -> StatusCode {
|
||||
StatusCode::OK
|
||||
}
|
||||
|
||||
/// The server is ready to accept traffic.
|
||||
async fn readiness() -> StatusCode {
|
||||
StatusCode::OK
|
||||
}
|
||||
23
litellm-rust/crates/ai-gateway/src/routes/mod.rs
Normal file
23
litellm-rust/crates/ai-gateway/src/routes/mod.rs
Normal file
|
|
@ -0,0 +1,23 @@
|
|||
//! HTTP routes.
|
||||
//!
|
||||
//! **Template:** every route module exposes `pub fn router() -> Router<AppState>`
|
||||
//! that mounts its own paths; [`app`] merges them. A trivial route is a single
|
||||
//! file (`health.rs`, `gil.rs`); a non-trivial one is a folder (`realtime/`) with
|
||||
//! `handler` (entry) + `service` (logic) + `transport` (adapters). See AGENTS.md.
|
||||
|
||||
pub mod gil;
|
||||
pub mod health;
|
||||
pub mod realtime;
|
||||
|
||||
use axum::Router;
|
||||
|
||||
use crate::state::AppState;
|
||||
|
||||
/// Assemble the application router by merging every route module's `router()`.
|
||||
pub fn app(state: AppState) -> Router {
|
||||
Router::new()
|
||||
.merge(health::router())
|
||||
.merge(gil::router())
|
||||
.merge(realtime::router())
|
||||
.with_state(state)
|
||||
}
|
||||
87
litellm-rust/crates/ai-gateway/src/routes/realtime/README.md
Normal file
87
litellm-rust/crates/ai-gateway/src/routes/realtime/README.md
Normal file
|
|
@ -0,0 +1,87 @@
|
|||
# Realtime route (`GET /v1/realtime`)
|
||||
|
||||
Proxies OpenAI's realtime WebSocket. `mod.rs` is the axum surface (handler +
|
||||
socket↔events adapter); `service.rs` is the pure logic (select a deployment, then
|
||||
splice client ↔ upstream). The pool itself lives in
|
||||
`crates/providers/src/realtime_pool.rs`.
|
||||
|
||||
## Connection pooling
|
||||
|
||||
### The problem
|
||||
|
||||
The gateway's realtime overhead lives **entirely in session establishment**. On each
|
||||
client connect it dials a *fresh* upstream WS to OpenAI and waits for
|
||||
`session.created` before it can serve. Measured at 5000 calls / 500 concurrency, the
|
||||
fresh-dial session phase is **~360 ms** vs **~7 ms** direct; dial, first-audio, and
|
||||
streaming add ~0. So the one lever is removing that per-connect handshake from the
|
||||
critical path.
|
||||
|
||||
### The idea
|
||||
|
||||
Keep a few upstream OpenAI sockets **already connected and already past
|
||||
`session.created`** (buffered). On a client connect, hand off a warm socket — relay
|
||||
its buffered `session.created` instantly (a local `Vec::pop`, sub-millisecond) and
|
||||
splice exactly as a fresh dial would. A background task keeps the pool topped up. On
|
||||
a miss or dead socket we fall back to fresh-dial: the pool is a latency optimization,
|
||||
never a correctness dependency.
|
||||
|
||||
```
|
||||
┌───────────────────────────────────────┐
|
||||
client connect ──────► │ routes/realtime → service::run │
|
||||
│ pool.take(key) │
|
||||
│ hit → relay buffered │
|
||||
│ session.created, then splice │
|
||||
│ miss → fresh dial (original path) │
|
||||
└───────────────┬───────────────────────┘
|
||||
│ replenish (async, concurrent)
|
||||
┌───────────────▼───────────────────────┐
|
||||
background task ─────► │ RealtimePool: per-key warm sockets │
|
||||
│ each = { ws, buffered session.created}│
|
||||
│ liveness-checked before handoff │
|
||||
└─────────────────────────────────────────┘
|
||||
```
|
||||
|
||||
A warm session is indistinguishable from a fresh one: OpenAI sends `session.created`
|
||||
unprompted on connect, we pre-read exactly that one frame and relay it on handoff,
|
||||
and we send nothing else on the socket before a client exists — so the client's first
|
||||
`session.update` behaves identically either way.
|
||||
|
||||
### Sizing
|
||||
|
||||
Each warm socket serves **exactly one** session (realtime isn't multiplexed), so the
|
||||
pool is sized to the **peak concurrent connects per instance**, not total live
|
||||
connections:
|
||||
|
||||
```
|
||||
REALTIME_POOL_SIZE ≈ peak_concurrency / instance_count
|
||||
```
|
||||
|
||||
e.g. 500 concurrency over 10 instances → ~50–64 per instance. The replenisher dials
|
||||
the missing sockets **concurrently**, so a drained pool refills in ~one handshake
|
||||
window and keeps supply close to the connect rate. Over-provisioning just burns idle
|
||||
upstream sockets, which is why warm sockets are short-lived
|
||||
(`REALTIME_POOL_MAX_IDLE_SECS`).
|
||||
|
||||
### Config
|
||||
|
||||
| env | default | meaning |
|
||||
| ----------------------------- | ------- | --------------------------------------------------------------- |
|
||||
| `REALTIME_POOL_SIZE` | `4` | target warm sockets per key. `0` disables pooling (fresh-dial). |
|
||||
| `REALTIME_POOL_MAX_IDLE_SECS` | `30` | max time a warm socket sits before it's closed and replaced. |
|
||||
|
||||
### Notes
|
||||
|
||||
- **Miss / dead socket → fresh dial.** Burst beyond warm supply, or a socket that
|
||||
died, never blocks or fails — it falls back to the original path. The pool can only
|
||||
make a connect faster, never slower or more fragile.
|
||||
- **Auth scope.** The pool key includes `api_key`, so a warm socket is only handed to
|
||||
a request resolving to the same key — no cross-tenant reuse.
|
||||
- **Idle billing.** Warm sockets are liveness-checked at handoff and capped at
|
||||
`REALTIME_POOL_MAX_IDLE_SECS` to bound idle billing and dodge OpenAI's idle timeout.
|
||||
- **Replenish backoff.** If a key's warm-up dials all fail (invalid credentials, an
|
||||
unreachable upstream), the replenisher puts that key into exponential backoff
|
||||
(500 ms → 30 s cap) instead of re-dialing it every tick. This bounds connection
|
||||
attempts against a broken key so it can't exhaust upstream rate limits and degrade
|
||||
valid cold-path traffic; the backoff resets the moment a dial succeeds.
|
||||
|
||||
Benchmarks and repro: `../../benchmarks/realtime/README.md`.
|
||||
88
litellm-rust/crates/ai-gateway/src/routes/realtime/mod.rs
Normal file
88
litellm-rust/crates/ai-gateway/src/routes/realtime/mod.rs
Normal file
|
|
@ -0,0 +1,88 @@
|
|||
//! `GET /v1/realtime` (WebSocket).
|
||||
//!
|
||||
//! This file is the **axum surface**: `router()`, the handler, and the small
|
||||
//! socket↔events adapter. The pure logic (no axum) lives in [`service`]. Auth is
|
||||
//! the `RequireMasterKey` extractor, so the handler stays thin.
|
||||
|
||||
mod service;
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use crate::io::realtime_pool::RealtimePool;
|
||||
use axum::extract::ws::{Message, WebSocket, WebSocketUpgrade};
|
||||
use axum::extract::{Query, State};
|
||||
use axum::http::StatusCode;
|
||||
use axum::response::Response;
|
||||
use axum::routing::get;
|
||||
use axum::Router;
|
||||
use futures_util::{SinkExt, StreamExt};
|
||||
use litellm_core::realtime::types::RealtimeEvent;
|
||||
use litellm_core::router::Router as ModelRouter;
|
||||
use serde::Deserialize;
|
||||
|
||||
use crate::auth::RequireMasterKey;
|
||||
use crate::state::AppState;
|
||||
|
||||
/// This route's contribution to the app router.
|
||||
pub fn router() -> Router<AppState> {
|
||||
Router::new().route("/v1/realtime", get(handle))
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct RealtimeQuery {
|
||||
model: String,
|
||||
}
|
||||
|
||||
/// Auth runs via the `RequireMasterKey` extractor. We validate the model BEFORE
|
||||
/// the upgrade so failures are clean HTTP (400/404), not a socket that opens then
|
||||
/// closes, then hand the socket to `bridge`.
|
||||
async fn handle(
|
||||
_auth: RequireMasterKey,
|
||||
ws: WebSocketUpgrade,
|
||||
State(state): State<AppState>,
|
||||
Query(query): Query<RealtimeQuery>,
|
||||
) -> Result<Response, (StatusCode, String)> {
|
||||
if query.model.trim().is_empty() {
|
||||
return Err((
|
||||
StatusCode::BAD_REQUEST,
|
||||
"missing 'model' query param".to_string(),
|
||||
));
|
||||
}
|
||||
if !state.router.has_deployment(&query.model) {
|
||||
return Err((
|
||||
StatusCode::NOT_FOUND,
|
||||
format!("no deployment for model '{}'", query.model),
|
||||
));
|
||||
}
|
||||
|
||||
let router = state.router.clone();
|
||||
let pool = state.realtime_pool.clone();
|
||||
let model = query.model;
|
||||
Ok(ws.on_upgrade(move |socket| bridge(socket, router, pool, model)))
|
||||
}
|
||||
|
||||
/// Adapt the axum socket (text frames) to the typed-event `Stream`/`Sink` the
|
||||
/// service wants, keeping axum types out of `service`.
|
||||
async fn bridge(
|
||||
socket: WebSocket,
|
||||
router: Arc<ModelRouter>,
|
||||
pool: Arc<RealtimePool>,
|
||||
model: String,
|
||||
) {
|
||||
let (ws_sink, ws_stream) = socket.split();
|
||||
|
||||
let client_in = ws_stream.filter_map(|message| async move {
|
||||
match message {
|
||||
Ok(Message::Text(text)) => serde_json::from_str::<RealtimeEvent>(&text).ok(),
|
||||
_ => None,
|
||||
}
|
||||
});
|
||||
let client_out = ws_sink.with(|event: RealtimeEvent| async move {
|
||||
Ok::<Message, axum::Error>(Message::Text(
|
||||
serde_json::to_string(&event).unwrap_or_default(),
|
||||
))
|
||||
});
|
||||
|
||||
futures_util::pin_mut!(client_in, client_out);
|
||||
let _ = service::run(&router, &pool, &model, None, client_in, client_out).await;
|
||||
}
|
||||
|
|
@ -0,0 +1,76 @@
|
|||
//! Business logic: select a deployment with the (pure) core router, then call the
|
||||
//! provider splice. The seam between `core::router` (selection only) and
|
||||
//! `io` (the actual WebSocket I/O).
|
||||
//!
|
||||
//! On connect we try a pre-warmed upstream from the pool (handshake already paid,
|
||||
//! `session.created` buffered) and relay it instantly. On a pool miss or dead warm
|
||||
//! socket we fresh-dial exactly as before — the pool is never on the critical path
|
||||
//! for correctness, only latency.
|
||||
|
||||
use std::time::Duration;
|
||||
|
||||
use crate::io::realtime_pool::{upstream_key, RealtimePool};
|
||||
use futures_util::{Sink, Stream};
|
||||
use litellm_core::error::CoreError;
|
||||
use litellm_core::realtime::types::RealtimeEvent;
|
||||
use litellm_core::router::Router;
|
||||
use litellm_core::CoreResult;
|
||||
|
||||
/// Select a deployment for `model` and splice the client stream to the provider.
|
||||
///
|
||||
/// `pool` supplies a pre-warmed upstream when one is available; otherwise we
|
||||
/// fresh-dial. A disabled pool always misses, so this collapses to the original
|
||||
/// fresh-dial behavior.
|
||||
pub async fn run<In, Out>(
|
||||
router: &Router,
|
||||
pool: &RealtimePool,
|
||||
model: &str,
|
||||
idle_timeout: Option<Duration>,
|
||||
client_in: In,
|
||||
client_out: Out,
|
||||
) -> CoreResult<()>
|
||||
where
|
||||
In: Stream<Item = RealtimeEvent> + Unpin + Send,
|
||||
Out: Sink<RealtimeEvent> + Unpin + Send,
|
||||
<Out as Sink<RealtimeEvent>>::Error: std::fmt::Display,
|
||||
{
|
||||
let deployment = router.get_available_deployment(model).ok_or_else(|| {
|
||||
CoreError::Routing(format!("no deployment available for model '{model}'"))
|
||||
})?;
|
||||
let params = &deployment.litellm_params;
|
||||
// Strip a leading `openai/` so the OpenAI-only realtime fn gets the bare model.
|
||||
let provider_model = params
|
||||
.model
|
||||
.strip_prefix("openai/")
|
||||
.unwrap_or(¶ms.model);
|
||||
|
||||
// Warm path: take a pooled upstream (handshake already paid) and relay its
|
||||
// buffered session.created immediately. On miss/dead socket fall through.
|
||||
if let Some(key) = upstream_key(
|
||||
provider_model,
|
||||
params.api_key.as_deref(),
|
||||
params.api_base.as_deref(),
|
||||
) {
|
||||
if let Some(handoff) = pool.take(&key) {
|
||||
return crate::io::realtime::realtime_warm(
|
||||
provider_model,
|
||||
handoff,
|
||||
idle_timeout,
|
||||
client_in,
|
||||
client_out,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
}
|
||||
|
||||
// Cold path: fresh dial (the original behavior).
|
||||
crate::io::realtime::realtime(
|
||||
provider_model,
|
||||
params.api_key.as_deref(),
|
||||
params.api_base.as_deref(),
|
||||
idle_timeout,
|
||||
client_in,
|
||||
client_out,
|
||||
)
|
||||
.await
|
||||
}
|
||||
17
litellm-rust/crates/ai-gateway/src/state.rs
Normal file
17
litellm-rust/crates/ai-gateway/src/state.rs
Normal file
|
|
@ -0,0 +1,17 @@
|
|||
use std::sync::Arc;
|
||||
|
||||
use crate::io::realtime_pool::RealtimePool;
|
||||
use litellm_core::router::Router;
|
||||
|
||||
/// Shared application state handed to every route handler.
|
||||
#[derive(Clone)]
|
||||
pub struct AppState {
|
||||
pub router: Arc<Router>,
|
||||
/// The gateway master key. Any caller presenting it as a bearer token may
|
||||
/// invoke the gateway. `None` → auth not configured (routes fail closed).
|
||||
pub master_key: Option<Arc<str>>,
|
||||
/// Pre-warmed upstream realtime connection pool. Disabled
|
||||
/// (`RealtimePool::disabled()`) when `REALTIME_POOL_SIZE=0`, in which case
|
||||
/// every realtime connect fresh-dials exactly as before.
|
||||
pub realtime_pool: Arc<RealtimePool>,
|
||||
}
|
||||
3
litellm-rust/crates/core/AGENTS.md
Normal file
3
litellm-rust/crates/core/AGENTS.md
Normal file
|
|
@ -0,0 +1,3 @@
|
|||
litellm-core is the PURE translation layer — types, route contracts (traits), provider transforms (modules under `providers/`), and the router. No network, no I/O, no env reads.
|
||||
|
||||
Routes (ocr, realtime) and providers (mistral, openai) are modules, not crates.
|
||||
47
litellm-rust/crates/core/CLAUDE.md
Normal file
47
litellm-rust/crates/core/CLAUDE.md
Normal file
|
|
@ -0,0 +1,47 @@
|
|||
# CLAUDE.md
|
||||
|
||||
Rules for `litellm-rust/crates/core`.
|
||||
|
||||
## Responsibility
|
||||
|
||||
`core` owns shared data types, typed errors, and deterministic helper contracts.
|
||||
It must stay pure and host-independent.
|
||||
|
||||
Allowed:
|
||||
- Shared request/response structs.
|
||||
- Typed errors with stable, non-sensitive messages.
|
||||
- Deterministic validation helpers.
|
||||
- Serialization helpers that intentionally mirror Python output shape.
|
||||
- Route templates that match Python base config responsibilities, such as
|
||||
`ocr::transformation::OcrProviderConfig`.
|
||||
|
||||
Not allowed:
|
||||
- Network, filesystem, database, cache, or environment access.
|
||||
- Secret reads or auth/header construction.
|
||||
- Logging callbacks, tracing spans, spend writes, or customer callbacks.
|
||||
- Provider-specific branching that belongs in `providers`.
|
||||
- Panics for user/provider-controlled input.
|
||||
|
||||
## Typed Contracts (core rule)
|
||||
|
||||
Trait and function boundaries MUST be strongly typed. No stringly-typed JSON
|
||||
(`&str` / `String` / `Vec<String>` / bare `serde_json::Value`) as a transform
|
||||
input or output. Parse wire bytes into typed structs/enums at the host edge;
|
||||
`core` and `providers` operate only on those types (e.g. `RealtimeEvent`,
|
||||
`RealtimeTransformResult`, `OcrRequestData`). A `type`-style discriminator is a
|
||||
typed field on a struct, not a raw string threaded through the API.
|
||||
|
||||
## Structure
|
||||
|
||||
Use route names directly under `src/`: `ocr`, future `messages`,
|
||||
`chat_completions`, `embeddings`, and similar top-level LiteLLM calls. Do not
|
||||
invent broad names like `engine` for route contracts.
|
||||
|
||||
## Parity Rules
|
||||
|
||||
- Every shared type used by a provider transform needs unit tests for
|
||||
serialization shape.
|
||||
- If Python parity requires always emitting a `null` field instead of omitting
|
||||
it, document that in code and pin it with a test.
|
||||
- Error enums should preserve enough detail for Python/HTTP hosts to map errors
|
||||
consistently without exposing document contents or upstream bodies.
|
||||
12
litellm-rust/crates/core/Cargo.toml
Normal file
12
litellm-rust/crates/core/Cargo.toml
Normal file
|
|
@ -0,0 +1,12 @@
|
|||
[package]
|
||||
name = "litellm-core"
|
||||
version = "0.1.0"
|
||||
edition.workspace = true
|
||||
license.workspace = true
|
||||
repository.workspace = true
|
||||
|
||||
[dependencies]
|
||||
rand.workspace = true
|
||||
serde.workspace = true
|
||||
serde_json.workspace = true
|
||||
thiserror.workspace = true
|
||||
35
litellm-rust/crates/core/src/error.rs
Normal file
35
litellm-rust/crates/core/src/error.rs
Normal file
|
|
@ -0,0 +1,35 @@
|
|||
use thiserror::Error;
|
||||
|
||||
pub type CoreResult<T> = Result<T, CoreError>;
|
||||
|
||||
#[derive(Debug, Error, PartialEq, Eq)]
|
||||
pub enum CoreError {
|
||||
#[error("expected {expected}, got {actual}")]
|
||||
InvalidType {
|
||||
expected: &'static str,
|
||||
actual: &'static str,
|
||||
},
|
||||
#[error("missing required field: {0}")]
|
||||
MissingField(&'static str),
|
||||
#[error("invalid response: {0}")]
|
||||
InvalidResponse(String),
|
||||
#[error("{0}")]
|
||||
Auth(String),
|
||||
#[error("OCR request failed with status {status}: {body}")]
|
||||
Http { status: u16, body: String },
|
||||
#[error("OCR network error: {0}")]
|
||||
Network(String),
|
||||
#[error("routing error: {0}")]
|
||||
Routing(String),
|
||||
}
|
||||
|
||||
pub fn json_type_name(value: &serde_json::Value) -> &'static str {
|
||||
match value {
|
||||
serde_json::Value::Null => "null",
|
||||
serde_json::Value::Bool(_) => "bool",
|
||||
serde_json::Value::Number(_) => "number",
|
||||
serde_json::Value::String(_) => "string",
|
||||
serde_json::Value::Array(_) => "array",
|
||||
serde_json::Value::Object(_) => "object",
|
||||
}
|
||||
}
|
||||
7
litellm-rust/crates/core/src/lib.rs
Normal file
7
litellm-rust/crates/core/src/lib.rs
Normal file
|
|
@ -0,0 +1,7 @@
|
|||
pub mod error;
|
||||
pub mod ocr;
|
||||
pub mod providers;
|
||||
pub mod realtime;
|
||||
pub mod router;
|
||||
|
||||
pub use error::{CoreError, CoreResult};
|
||||
2
litellm-rust/crates/core/src/ocr/mod.rs
Normal file
2
litellm-rust/crates/core/src/ocr/mod.rs
Normal file
|
|
@ -0,0 +1,2 @@
|
|||
pub mod transformation;
|
||||
pub mod types;
|
||||
32
litellm-rust/crates/core/src/ocr/transformation.rs
Normal file
32
litellm-rust/crates/core/src/ocr/transformation.rs
Normal file
|
|
@ -0,0 +1,32 @@
|
|||
use serde_json::{Map, Value};
|
||||
|
||||
use crate::CoreResult;
|
||||
|
||||
use super::types::{OcrRequestData, OcrResponseData};
|
||||
|
||||
pub trait OcrProviderConfig {
|
||||
fn supported_ocr_params(&self) -> &'static [&'static str];
|
||||
|
||||
fn map_ocr_params(&self, non_default_params: &Map<String, Value>) -> Map<String, Value> {
|
||||
let mut mapped_params = Map::new();
|
||||
for (param, value) in non_default_params {
|
||||
if self.supported_ocr_params().contains(¶m.as_str()) {
|
||||
mapped_params.insert(param.clone(), value.clone());
|
||||
}
|
||||
}
|
||||
mapped_params
|
||||
}
|
||||
|
||||
fn transform_ocr_request(
|
||||
&self,
|
||||
model: &str,
|
||||
document: Value,
|
||||
optional_params: Map<String, Value>,
|
||||
) -> CoreResult<OcrRequestData>;
|
||||
|
||||
fn transform_ocr_response(
|
||||
&self,
|
||||
model: &str,
|
||||
response_json: Value,
|
||||
) -> CoreResult<OcrResponseData>;
|
||||
}
|
||||
29
litellm-rust/crates/core/src/ocr/types.rs
Normal file
29
litellm-rust/crates/core/src/ocr/types.rs
Normal file
|
|
@ -0,0 +1,29 @@
|
|||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::Value;
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
|
||||
pub struct OcrRequestData {
|
||||
pub data: Value,
|
||||
pub files: Option<Value>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
|
||||
pub struct OcrResponseData {
|
||||
pub pages: Vec<Value>,
|
||||
pub model: String,
|
||||
pub document_annotation: Option<Value>,
|
||||
pub usage_info: Option<Value>,
|
||||
pub object: String,
|
||||
}
|
||||
|
||||
impl OcrResponseData {
|
||||
pub fn into_json(self) -> Value {
|
||||
serde_json::json!({
|
||||
"pages": self.pages,
|
||||
"model": self.model,
|
||||
"document_annotation": self.document_annotation,
|
||||
"usage_info": self.usage_info,
|
||||
"object": self.object,
|
||||
})
|
||||
}
|
||||
}
|
||||
1
litellm-rust/crates/core/src/providers/mistral/mod.rs
Normal file
1
litellm-rust/crates/core/src/providers/mistral/mod.rs
Normal file
|
|
@ -0,0 +1 @@
|
|||
pub mod ocr;
|
||||
|
|
@ -0,0 +1 @@
|
|||
pub mod transformation;
|
||||
|
|
@ -0,0 +1,292 @@
|
|||
use crate::error::{json_type_name, CoreError, CoreResult};
|
||||
use crate::ocr::transformation::OcrProviderConfig;
|
||||
use crate::ocr::types::{OcrRequestData, OcrResponseData};
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
const SUPPORTED_OCR_PARAMS: &[&str] = &[
|
||||
"pages",
|
||||
"include_image_base64",
|
||||
"image_limit",
|
||||
"image_min_size",
|
||||
"bbox_annotation_format",
|
||||
"document_annotation_format",
|
||||
"document_annotation_prompt",
|
||||
"extract_header",
|
||||
"extract_footer",
|
||||
"table_format",
|
||||
"confidence_scores_granularity",
|
||||
"id",
|
||||
];
|
||||
|
||||
/// Default Mistral API base, used when the caller does not override `api_base`.
|
||||
pub const MISTRAL_DEFAULT_API_BASE: &str = "https://api.mistral.ai/v1";
|
||||
|
||||
/// Environment variable holding the Mistral API key.
|
||||
pub const MISTRAL_API_KEY_ENV: &str = "MISTRAL_API_KEY";
|
||||
|
||||
/// Error message raised when no Mistral API key can be resolved.
|
||||
pub const MISSING_KEY_MESSAGE: &str = "Missing Mistral API Key - A call is being made to Mistral but no key is set either in the environment variables or via params";
|
||||
|
||||
/// Build the complete OCR endpoint URL, de-duplicating a trailing `/v1`.
|
||||
///
|
||||
/// Blank/whitespace `api_base` is treated as absent (guard at resolution time).
|
||||
pub fn complete_url(api_base: Option<&str>) -> String {
|
||||
let base = api_base
|
||||
.map(str::trim)
|
||||
.filter(|base| !base.is_empty())
|
||||
.unwrap_or(MISTRAL_DEFAULT_API_BASE)
|
||||
.trim_end_matches('/');
|
||||
|
||||
if base.ends_with("/v1") {
|
||||
format!("{base}/ocr")
|
||||
} else {
|
||||
format!("{base}/v1/ocr")
|
||||
}
|
||||
}
|
||||
|
||||
/// Resolve the Mistral API key from the explicit param or the environment.
|
||||
///
|
||||
/// Blank/whitespace values are treated as absent. Returns `CoreError::Auth`
|
||||
/// when no usable key is available.
|
||||
///
|
||||
/// Note: the env fallback only reads the process environment. Secret-manager
|
||||
/// backends (AWS/Azure/GCP/Vault) are resolved on the Python side and passed in
|
||||
/// via `api_key`; this fallback is a last resort for direct/standalone use.
|
||||
pub fn resolve_api_key(
|
||||
api_key: Option<&str>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> CoreResult<String> {
|
||||
api_key
|
||||
.map(str::trim)
|
||||
.filter(|key| !key.is_empty())
|
||||
.map(str::to_string)
|
||||
.or_else(|| env_lookup(MISTRAL_API_KEY_ENV).filter(|key| !key.trim().is_empty()))
|
||||
.ok_or_else(|| CoreError::Auth(MISSING_KEY_MESSAGE.to_string()))
|
||||
}
|
||||
|
||||
pub struct MistralOcrConfig;
|
||||
|
||||
pub const MISTRAL_OCR_CONFIG: MistralOcrConfig = MistralOcrConfig;
|
||||
|
||||
impl OcrProviderConfig for MistralOcrConfig {
|
||||
fn supported_ocr_params(&self) -> &'static [&'static str] {
|
||||
SUPPORTED_OCR_PARAMS
|
||||
}
|
||||
|
||||
fn transform_ocr_request(
|
||||
&self,
|
||||
model: &str,
|
||||
document: Value,
|
||||
optional_params: Map<String, Value>,
|
||||
) -> CoreResult<OcrRequestData> {
|
||||
if !document.is_object() {
|
||||
return Err(CoreError::InvalidType {
|
||||
expected: "object",
|
||||
actual: json_type_name(&document),
|
||||
});
|
||||
}
|
||||
|
||||
let mut data = Map::new();
|
||||
data.insert("model".to_string(), Value::String(model.to_string()));
|
||||
data.insert("document".to_string(), document);
|
||||
for (param, value) in optional_params {
|
||||
data.insert(param, value);
|
||||
}
|
||||
|
||||
Ok(OcrRequestData {
|
||||
data: Value::Object(data),
|
||||
files: None,
|
||||
})
|
||||
}
|
||||
|
||||
fn transform_ocr_response(
|
||||
&self,
|
||||
model: &str,
|
||||
response_json: Value,
|
||||
) -> CoreResult<OcrResponseData> {
|
||||
let response_object = response_json
|
||||
.as_object()
|
||||
.ok_or_else(|| CoreError::InvalidType {
|
||||
expected: "object",
|
||||
actual: json_type_name(&response_json),
|
||||
})?;
|
||||
|
||||
let pages = response_object
|
||||
.get("pages")
|
||||
.and_then(Value::as_array)
|
||||
.cloned()
|
||||
.unwrap_or_default();
|
||||
let model = response_object
|
||||
.get("model")
|
||||
.and_then(Value::as_str)
|
||||
.unwrap_or(model)
|
||||
.to_string();
|
||||
let document_annotation = response_object.get("document_annotation").cloned();
|
||||
let usage_info = response_object.get("usage_info").cloned();
|
||||
|
||||
Ok(OcrResponseData {
|
||||
pages,
|
||||
model,
|
||||
document_annotation,
|
||||
usage_info,
|
||||
object: "ocr".to_string(),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
pub fn supported_ocr_params() -> &'static [&'static str] {
|
||||
MISTRAL_OCR_CONFIG.supported_ocr_params()
|
||||
}
|
||||
|
||||
pub fn map_ocr_params(non_default_params: &Map<String, Value>) -> Map<String, Value> {
|
||||
MISTRAL_OCR_CONFIG.map_ocr_params(non_default_params)
|
||||
}
|
||||
|
||||
pub fn transform_ocr_request(
|
||||
model: &str,
|
||||
document: Value,
|
||||
optional_params: Map<String, Value>,
|
||||
) -> CoreResult<OcrRequestData> {
|
||||
MISTRAL_OCR_CONFIG.transform_ocr_request(model, document, optional_params)
|
||||
}
|
||||
|
||||
pub fn transform_ocr_response(model: &str, response_json: Value) -> CoreResult<OcrResponseData> {
|
||||
MISTRAL_OCR_CONFIG.transform_ocr_response(model, response_json)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use serde_json::json;
|
||||
|
||||
#[test]
|
||||
fn supported_params_match_python_mistral_ocr_config() {
|
||||
assert_eq!(
|
||||
supported_ocr_params(),
|
||||
&[
|
||||
"pages",
|
||||
"include_image_base64",
|
||||
"image_limit",
|
||||
"image_min_size",
|
||||
"bbox_annotation_format",
|
||||
"document_annotation_format",
|
||||
"document_annotation_prompt",
|
||||
"extract_header",
|
||||
"extract_footer",
|
||||
"table_format",
|
||||
"confidence_scores_granularity",
|
||||
"id",
|
||||
]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn map_ocr_params_drops_unknown_params() {
|
||||
let params = json!({
|
||||
"extract_header": true,
|
||||
"unsupported_param": "value",
|
||||
"pages": [0, 1]
|
||||
});
|
||||
let mapped = map_ocr_params(params.as_object().unwrap());
|
||||
|
||||
assert_eq!(mapped.get("extract_header"), Some(&json!(true)));
|
||||
assert_eq!(mapped.get("pages"), Some(&json!([0, 1])));
|
||||
assert!(!mapped.contains_key("unsupported_param"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn transform_ocr_request_builds_mistral_body() {
|
||||
let document = json!({
|
||||
"type": "document_url",
|
||||
"document_url": "https://example.com/doc.pdf"
|
||||
});
|
||||
let optional_params = json!({
|
||||
"include_image_base64": true,
|
||||
"table_format": "html"
|
||||
})
|
||||
.as_object()
|
||||
.unwrap()
|
||||
.clone();
|
||||
|
||||
let result = transform_ocr_request("mistral-ocr-latest", document.clone(), optional_params)
|
||||
.expect("request should transform");
|
||||
|
||||
assert_eq!(
|
||||
result.data,
|
||||
json!({
|
||||
"model": "mistral-ocr-latest",
|
||||
"document": document,
|
||||
"include_image_base64": true,
|
||||
"table_format": "html"
|
||||
})
|
||||
);
|
||||
assert_eq!(result.files, None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn transform_ocr_request_rejects_non_object_document() {
|
||||
let err = transform_ocr_request("mistral-ocr-latest", json!("bad"), Map::new())
|
||||
.expect_err("string document should be rejected");
|
||||
|
||||
assert_eq!(
|
||||
err,
|
||||
CoreError::InvalidType {
|
||||
expected: "object",
|
||||
actual: "string",
|
||||
}
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn transform_ocr_response_normalizes_mistral_json() {
|
||||
let response = json!({
|
||||
"pages": [{"index": 0, "markdown": "hello"}],
|
||||
"model": "mistral-ocr-2505-completion",
|
||||
"document_annotation": null,
|
||||
"usage_info": {"pages_processed": 1}
|
||||
});
|
||||
|
||||
let result = transform_ocr_response("mistral-ocr-latest", response)
|
||||
.expect("response should transform");
|
||||
|
||||
assert_eq!(result.pages, vec![json!({"index": 0, "markdown": "hello"})]);
|
||||
assert_eq!(result.model, "mistral-ocr-2505-completion");
|
||||
assert_eq!(result.document_annotation, Some(Value::Null));
|
||||
assert_eq!(result.usage_info, Some(json!({"pages_processed": 1})));
|
||||
assert_eq!(result.object, "ocr");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn complete_url_defaults_and_dedupes_v1() {
|
||||
assert_eq!(complete_url(None), "https://api.mistral.ai/v1/ocr");
|
||||
assert_eq!(complete_url(Some(" ")), "https://api.mistral.ai/v1/ocr");
|
||||
assert_eq!(
|
||||
complete_url(Some("https://proxy.internal")),
|
||||
"https://proxy.internal/v1/ocr"
|
||||
);
|
||||
assert_eq!(
|
||||
complete_url(Some("https://proxy.internal/v1/")),
|
||||
"https://proxy.internal/v1/ocr"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolve_api_key_prefers_param_then_env() {
|
||||
let no_env = |_: &str| None;
|
||||
assert_eq!(
|
||||
resolve_api_key(Some("sk-param"), &no_env).unwrap(),
|
||||
"sk-param"
|
||||
);
|
||||
|
||||
let with_env = |key: &str| (key == MISTRAL_API_KEY_ENV).then(|| "sk-env".to_string());
|
||||
assert_eq!(resolve_api_key(None, &with_env).unwrap(), "sk-env");
|
||||
// Blank param falls through to the environment.
|
||||
assert_eq!(resolve_api_key(Some(" "), &with_env).unwrap(), "sk-env");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolve_api_key_errors_when_absent() {
|
||||
let err = resolve_api_key(None, &|_| None).expect_err("missing key should error");
|
||||
assert_eq!(err, CoreError::Auth(MISSING_KEY_MESSAGE.to_string()));
|
||||
}
|
||||
}
|
||||
2
litellm-rust/crates/core/src/providers/mod.rs
Normal file
2
litellm-rust/crates/core/src/providers/mod.rs
Normal file
|
|
@ -0,0 +1,2 @@
|
|||
pub mod mistral;
|
||||
pub mod openai;
|
||||
1
litellm-rust/crates/core/src/providers/openai/mod.rs
Normal file
1
litellm-rust/crates/core/src/providers/openai/mod.rs
Normal file
|
|
@ -0,0 +1 @@
|
|||
pub mod realtime;
|
||||
|
|
@ -0,0 +1 @@
|
|||
pub mod transformation;
|
||||
|
|
@ -0,0 +1,189 @@
|
|||
use crate::realtime::transformation::RealtimeProviderConfig;
|
||||
use crate::realtime::types::{RealtimeEvent, RealtimeTransformResult};
|
||||
use crate::CoreResult;
|
||||
|
||||
/// Default OpenAI API base, used when the caller does not override `api_base`.
|
||||
pub const OPENAI_REALTIME_DEFAULT_API_BASE: &str = "https://api.openai.com";
|
||||
|
||||
/// Path appended to the resolved host base to reach the realtime endpoint.
|
||||
pub const OPENAI_REALTIME_PATH: &str = "/v1/realtime";
|
||||
|
||||
/// Percent-encode a query value, escaping any char outside the RFC 3986
|
||||
/// unreserved set (`A-Za-z0-9-._~`). Keeps us dependency-free; common realtime
|
||||
/// model slugs have no special chars, but this stays correct for the rest.
|
||||
fn percent_encode(value: &str) -> String {
|
||||
let mut encoded = String::with_capacity(value.len());
|
||||
for byte in value.bytes() {
|
||||
let unreserved = byte.is_ascii_alphanumeric() || matches!(byte, b'-' | b'.' | b'_' | b'~');
|
||||
if unreserved {
|
||||
encoded.push(byte as char);
|
||||
} else {
|
||||
encoded.push('%');
|
||||
encoded.push_str(&format!("{byte:02X}"));
|
||||
}
|
||||
}
|
||||
encoded
|
||||
}
|
||||
|
||||
/// Build the realtime WebSocket URL, porting Python's `OpenAIRealtime._construct_url`.
|
||||
///
|
||||
/// Blank/whitespace `api_base` is treated as absent (guard at resolution time),
|
||||
/// falling back to the default. The scheme is swapped to its WebSocket
|
||||
/// equivalent (`https://`→`wss://`, `http://`→`ws://`); bases already using
|
||||
/// `ws`/`wss` are left untouched. A bare host or unrecognized scheme defaults to
|
||||
/// secure `wss://` so we never hand a scheme-less URL to the connector (this is
|
||||
/// a deliberate hardening over Python's `_construct_url`, which would emit a
|
||||
/// scheme-less URL here). A trailing `/` is trimmed before the path and
|
||||
/// `?model=<encoded>` are appended.
|
||||
pub fn complete_url(api_base: Option<&str>, model: &str) -> String {
|
||||
let base = api_base
|
||||
.map(str::trim)
|
||||
.filter(|base| !base.is_empty())
|
||||
.unwrap_or(OPENAI_REALTIME_DEFAULT_API_BASE);
|
||||
|
||||
let base = if let Some(rest) = base.strip_prefix("https://") {
|
||||
format!("wss://{rest}")
|
||||
} else if let Some(rest) = base.strip_prefix("http://") {
|
||||
format!("ws://{rest}")
|
||||
} else if base.starts_with("wss://") || base.starts_with("ws://") {
|
||||
base.to_string()
|
||||
} else {
|
||||
format!("wss://{base}")
|
||||
};
|
||||
|
||||
let base = base.trim_end_matches('/');
|
||||
|
||||
format!(
|
||||
"{base}{OPENAI_REALTIME_PATH}?model={}",
|
||||
percent_encode(model)
|
||||
)
|
||||
}
|
||||
|
||||
pub struct OpenAiRealtimeConfig;
|
||||
|
||||
pub const OPENAI_REALTIME_CONFIG: OpenAiRealtimeConfig = OpenAiRealtimeConfig;
|
||||
|
||||
impl RealtimeProviderConfig for OpenAiRealtimeConfig {
|
||||
fn complete_url(&self, api_base: Option<&str>, model: &str) -> String {
|
||||
complete_url(api_base, model)
|
||||
}
|
||||
|
||||
fn transform_realtime_request(
|
||||
&self,
|
||||
event: &RealtimeEvent,
|
||||
_model: &str,
|
||||
) -> CoreResult<RealtimeTransformResult> {
|
||||
Ok(RealtimeTransformResult::passthrough(event.clone()))
|
||||
}
|
||||
|
||||
fn transform_realtime_response(
|
||||
&self,
|
||||
event: &RealtimeEvent,
|
||||
_model: &str,
|
||||
) -> CoreResult<RealtimeTransformResult> {
|
||||
Ok(RealtimeTransformResult::passthrough(event.clone()))
|
||||
}
|
||||
}
|
||||
|
||||
pub fn transform_realtime_request(
|
||||
event: &RealtimeEvent,
|
||||
model: &str,
|
||||
) -> CoreResult<RealtimeTransformResult> {
|
||||
OPENAI_REALTIME_CONFIG.transform_realtime_request(event, model)
|
||||
}
|
||||
|
||||
pub fn transform_realtime_response(
|
||||
event: &RealtimeEvent,
|
||||
model: &str,
|
||||
) -> CoreResult<RealtimeTransformResult> {
|
||||
OPENAI_REALTIME_CONFIG.transform_realtime_response(event, model)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn complete_url_defaults_to_openai_wss() {
|
||||
assert_eq!(
|
||||
complete_url(None, "gpt-4o-realtime-preview"),
|
||||
"wss://api.openai.com/v1/realtime?model=gpt-4o-realtime-preview"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn complete_url_blank_base_uses_default() {
|
||||
assert_eq!(
|
||||
complete_url(Some(" "), "gpt-4o-realtime-preview"),
|
||||
"wss://api.openai.com/v1/realtime?model=gpt-4o-realtime-preview"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn complete_url_swaps_http_to_ws() {
|
||||
assert_eq!(
|
||||
complete_url(Some("http://localhost:8080"), "gpt-4o-realtime-preview"),
|
||||
"ws://localhost:8080/v1/realtime?model=gpt-4o-realtime-preview"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn complete_url_dedupes_trailing_slash() {
|
||||
assert_eq!(
|
||||
complete_url(Some("https://api.openai.com/"), "gpt-4o-realtime-preview"),
|
||||
"wss://api.openai.com/v1/realtime?model=gpt-4o-realtime-preview"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn complete_url_custom_base() {
|
||||
assert_eq!(
|
||||
complete_url(Some("https://oai.azure.example"), "gpt-4o-realtime-preview"),
|
||||
"wss://oai.azure.example/v1/realtime?model=gpt-4o-realtime-preview"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn complete_url_preserves_existing_wss_scheme() {
|
||||
assert_eq!(
|
||||
complete_url(Some("wss://api.openai.com"), "gpt-realtime"),
|
||||
"wss://api.openai.com/v1/realtime?model=gpt-realtime"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn complete_url_bare_host_defaults_to_wss() {
|
||||
assert_eq!(
|
||||
complete_url(Some("api.openai.com"), "gpt-realtime"),
|
||||
"wss://api.openai.com/v1/realtime?model=gpt-realtime"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn complete_url_percent_encodes_model_space() {
|
||||
assert_eq!(
|
||||
complete_url(None, "gpt 4o"),
|
||||
"wss://api.openai.com/v1/realtime?model=gpt%204o"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn transform_realtime_request_passthrough_preserves_event() {
|
||||
let event: RealtimeEvent =
|
||||
serde_json::from_str(r#"{"type":"session.update","session":{"voice":"alloy"}}"#)
|
||||
.expect("valid event");
|
||||
let result =
|
||||
transform_realtime_request(&event, "gpt-realtime").expect("passthrough is infallible");
|
||||
assert_eq!(result.events, vec![event]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn transform_realtime_response_passthrough_preserves_event() {
|
||||
let event: RealtimeEvent =
|
||||
serde_json::from_str(r#"{"type":"response.output_audio.delta","delta":"abc=="}"#)
|
||||
.expect("valid event");
|
||||
let result =
|
||||
transform_realtime_response(&event, "gpt-realtime").expect("passthrough is infallible");
|
||||
assert_eq!(result.events, vec![event]);
|
||||
}
|
||||
}
|
||||
Some files were not shown because too many files have changed in this diff Show more
Loading…
Add table
Reference in a new issue