mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-09 03:18:44 +00:00
Merge remote-tracking branch 'berri/litellm_internal_staging' into litellm_e2e_testing
This commit is contained in:
commit
c5922e87bc
1128 changed files with 61751 additions and 14135 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 \
|
||||
|
|
|
|||
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 }}"
|
||||
|
|
|
|||
8
.github/workflows/test-linting.yml
vendored
8
.github/workflows/test-linting.yml
vendored
|
|
@ -14,7 +14,7 @@ permissions:
|
|||
jobs:
|
||||
lint:
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 10
|
||||
timeout-minutes: 15
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
|
||||
|
|
@ -87,9 +87,11 @@ jobs:
|
|||
run: |
|
||||
uv run --no-sync python -c "import openai; print(f'OpenAI version: {openai.__version__}')"
|
||||
|
||||
- name: Run basedpyright type checking
|
||||
- name: Check basedpyright budget (delta vs base)
|
||||
env:
|
||||
BASE_SHA: ${{ github.event.pull_request.base.sha }}
|
||||
run: |
|
||||
(uv run --no-sync basedpyright --outputjson || true) | uv run --no-sync python scripts/type_check_gate.py
|
||||
(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: |
|
||||
|
|
|
|||
4
.github/workflows/test-litellm-ui-build.yml
vendored
4
.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"
|
||||
|
|
|
|||
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
|
||||
|
|
|
|||
12
Makefile
12
Makefile
|
|
@ -6,7 +6,7 @@
|
|||
test-proxy-unit-a test-proxy-unit-b test-integration test-unit-helm \
|
||||
info lint lint-dev format \
|
||||
lint-basedpyright lint-basedpyright-budget-update \
|
||||
lint-ruff-budget lint-ruff-budget-update lint-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
|
||||
|
||||
|
|
@ -28,6 +28,7 @@ help:
|
|||
@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"
|
||||
|
|
@ -124,7 +125,8 @@ lint-ruff-FULL-dev: install-dev
|
|||
else echo "No changed .py files to check."; fi
|
||||
|
||||
lint-basedpyright: install-dev
|
||||
($(UV_RUN) basedpyright --outputjson || true) | $(UV_RUN) python scripts/type_check_gate.py
|
||||
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
|
||||
|
|
@ -134,6 +136,12 @@ 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
|
||||
|
||||
|
|
|
|||
143
README.md
143
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>
|
||||
|
|
@ -345,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) | ✅ | ✅ | ✅ | | | | | | | |
|
||||
|
|
@ -405,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
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -1,31 +1,31 @@
|
|||
{
|
||||
"reportAny": {
|
||||
"baseline": 24954,
|
||||
"baseline": 24989,
|
||||
"slack": 2500
|
||||
},
|
||||
"reportArgumentType": {
|
||||
"baseline": 1863,
|
||||
"baseline": 1934,
|
||||
"slack": 180
|
||||
},
|
||||
"reportAssignmentType": {
|
||||
"baseline": 220,
|
||||
"slack": 3
|
||||
"slack": 22
|
||||
},
|
||||
"reportAttributeAccessIssue": {
|
||||
"baseline": 335,
|
||||
"slack": 3
|
||||
"baseline": 346,
|
||||
"slack": 35
|
||||
},
|
||||
"reportCallIssue": {
|
||||
"baseline": 77,
|
||||
"baseline": 87,
|
||||
"slack": 10
|
||||
},
|
||||
"reportConstantRedefinition": {
|
||||
"baseline": 39,
|
||||
"slack": 3
|
||||
"slack": 4
|
||||
},
|
||||
"reportDeprecated": {
|
||||
"baseline": 217,
|
||||
"slack": 10
|
||||
"slack": 22
|
||||
},
|
||||
"reportDuplicateImport": {
|
||||
"baseline": 28,
|
||||
|
|
@ -41,11 +41,11 @@
|
|||
},
|
||||
"reportGeneralTypeIssues": {
|
||||
"baseline": 151,
|
||||
"slack": 3
|
||||
"slack": 15
|
||||
},
|
||||
"reportIncompatibleMethodOverride": {
|
||||
"baseline": 52,
|
||||
"slack": 10
|
||||
"slack": 5
|
||||
},
|
||||
"reportIncompatibleVariableOverride": {
|
||||
"baseline": 8,
|
||||
|
|
@ -73,7 +73,7 @@
|
|||
},
|
||||
"reportMissingParameterType": {
|
||||
"baseline": 3933,
|
||||
"slack": 10
|
||||
"slack": 390
|
||||
},
|
||||
"reportMissingTypeArgument": {
|
||||
"baseline": 10612,
|
||||
|
|
@ -97,7 +97,7 @@
|
|||
},
|
||||
"reportOptionalMemberAccess": {
|
||||
"baseline": 724,
|
||||
"slack": 10
|
||||
"slack": 72
|
||||
},
|
||||
"reportOptionalOperand": {
|
||||
"baseline": 3,
|
||||
|
|
@ -120,8 +120,8 @@
|
|||
"slack": 3
|
||||
},
|
||||
"reportReturnType": {
|
||||
"baseline": 118,
|
||||
"slack": 10
|
||||
"baseline": 126,
|
||||
"slack": 100
|
||||
},
|
||||
"reportTypedDictNotRequiredAccess": {
|
||||
"baseline": 20,
|
||||
|
|
@ -136,19 +136,19 @@
|
|||
"slack": 3000
|
||||
},
|
||||
"reportUnknownLambdaType": {
|
||||
"baseline": 76,
|
||||
"baseline": 75,
|
||||
"slack": 10
|
||||
},
|
||||
"reportUnknownMemberType": {
|
||||
"baseline": 27322,
|
||||
"baseline": 27037,
|
||||
"slack": 2500
|
||||
},
|
||||
"reportUnknownParameterType": {
|
||||
"baseline": 13636,
|
||||
"baseline": 13612,
|
||||
"slack": 1000
|
||||
},
|
||||
"reportUnknownVariableType": {
|
||||
"baseline": 21776,
|
||||
"baseline": 21445,
|
||||
"slack": 2000
|
||||
},
|
||||
"reportUnnecessaryCast": {
|
||||
|
|
@ -156,20 +156,20 @@
|
|||
"slack": 10
|
||||
},
|
||||
"reportUnnecessaryComparison": {
|
||||
"baseline": 680,
|
||||
"slack": 10
|
||||
"baseline": 683,
|
||||
"slack": 100
|
||||
},
|
||||
"reportUnnecessaryContains": {
|
||||
"baseline": 4,
|
||||
"slack": 3
|
||||
},
|
||||
"reportUnnecessaryIsInstance": {
|
||||
"baseline": 807,
|
||||
"slack": 10
|
||||
"baseline": 808,
|
||||
"slack": 80
|
||||
},
|
||||
"reportUntypedBaseClass": {
|
||||
"baseline": 110,
|
||||
"slack": 3
|
||||
"slack": 11
|
||||
},
|
||||
"reportUntypedFunctionDecorator": {
|
||||
"baseline": 22,
|
||||
|
|
@ -185,10 +185,10 @@
|
|||
},
|
||||
"reportUnusedImport": {
|
||||
"baseline": 670,
|
||||
"slack": 10
|
||||
"slack": 50
|
||||
},
|
||||
"reportUnusedVariable": {
|
||||
"baseline": 865,
|
||||
"slack": 10
|
||||
"slack": 50
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
@ -1472,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
|
||||
|
|
@ -1550,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(
|
||||
|
|
|
|||
|
|
@ -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
litellm-rust/.gitignore
vendored
Normal file
1
litellm-rust/.gitignore
vendored
Normal file
|
|
@ -0,0 +1 @@
|
|||
/target/
|
||||
88
litellm-rust/CLAUDE.md
Normal file
88
litellm-rust/CLAUDE.md
Normal file
|
|
@ -0,0 +1,88 @@
|
|||
# CLAUDE.md
|
||||
|
||||
This file defines the rules for Rust work in LiteLLM.
|
||||
|
||||
## Core Boundary
|
||||
|
||||
The `core` and `providers` crates describe work; hosts execute 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`.
|
||||
- `providers/src/<provider>/<route>/transformation.rs` owns the
|
||||
provider-specific transform. For Mistral OCR, this means
|
||||
`providers/src/mistral/ocr/transformation.rs`.
|
||||
- Future network execution belongs in a host/transport layer such as
|
||||
`llm_http_handler`, not inside `core` or `providers`.
|
||||
|
||||
Allowed in `core` and `providers`:
|
||||
- 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` or `providers`:
|
||||
- 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
|
||||
cargo clippy --workspace --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.
|
||||
1498
litellm-rust/Cargo.lock
generated
Normal file
1498
litellm-rust/Cargo.lock
generated
Normal file
File diff suppressed because it is too large
Load diff
21
litellm-rust/Cargo.toml
Normal file
21
litellm-rust/Cargo.toml
Normal file
|
|
@ -0,0 +1,21 @@
|
|||
[workspace]
|
||||
members = [
|
||||
"crates/core",
|
||||
"crates/providers",
|
||||
"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-providers = { path = "crates/providers" }
|
||||
pyo3 = "0.23.5"
|
||||
reqwest = { version = "0.12", default-features = false, features = ["blocking", "json", "rustls-tls"] }
|
||||
serde = { version = "1.0", features = ["derive"] }
|
||||
serde_json = "1.0"
|
||||
thiserror = "2.0"
|
||||
34
litellm-rust/README.md
Normal file
34
litellm-rust/README.md
Normal file
|
|
@ -0,0 +1,34 @@
|
|||
# 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.
|
||||
|
||||
## 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
|
||||
```
|
||||
38
litellm-rust/crates/core/CLAUDE.md
Normal file
38
litellm-rust/crates/core/CLAUDE.md
Normal file
|
|
@ -0,0 +1,38 @@
|
|||
# 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.
|
||||
|
||||
## 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.
|
||||
11
litellm-rust/crates/core/Cargo.toml
Normal file
11
litellm-rust/crates/core/Cargo.toml
Normal file
|
|
@ -0,0 +1,11 @@
|
|||
[package]
|
||||
name = "litellm-core"
|
||||
version = "0.1.0"
|
||||
edition.workspace = true
|
||||
license.workspace = true
|
||||
repository.workspace = true
|
||||
|
||||
[dependencies]
|
||||
serde.workspace = true
|
||||
serde_json.workspace = true
|
||||
thiserror.workspace = true
|
||||
33
litellm-rust/crates/core/src/error.rs
Normal file
33
litellm-rust/crates/core/src/error.rs
Normal file
|
|
@ -0,0 +1,33 @@
|
|||
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),
|
||||
}
|
||||
|
||||
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",
|
||||
}
|
||||
}
|
||||
4
litellm-rust/crates/core/src/lib.rs
Normal file
4
litellm-rust/crates/core/src/lib.rs
Normal file
|
|
@ -0,0 +1,4 @@
|
|||
pub mod error;
|
||||
pub mod ocr;
|
||||
|
||||
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,
|
||||
})
|
||||
}
|
||||
}
|
||||
53
litellm-rust/crates/providers/CLAUDE.md
Normal file
53
litellm-rust/crates/providers/CLAUDE.md
Normal file
|
|
@ -0,0 +1,53 @@
|
|||
# CLAUDE.md
|
||||
|
||||
Rules for `litellm-rust/crates/providers`.
|
||||
|
||||
## Responsibility
|
||||
|
||||
`providers` owns provider-specific pure transforms. It mirrors the existing
|
||||
Python provider modules closely enough that parity review is mechanical.
|
||||
|
||||
Provider files should map to the Python provider tree:
|
||||
|
||||
```text
|
||||
providers/src/<provider>/<route>/transformation.rs
|
||||
```
|
||||
|
||||
For example, Mistral OCR lives at
|
||||
`providers/src/mistral/ocr/transformation.rs`, matching
|
||||
`litellm/llms/mistral/ocr/transformation.py`.
|
||||
|
||||
Allowed:
|
||||
- Provider request transforms.
|
||||
- Provider response normalization.
|
||||
- Supported-parameter filtering.
|
||||
- Provider-specific validation that does not require I/O or secrets.
|
||||
|
||||
Not allowed:
|
||||
- HTTP clients or provider SDK calls.
|
||||
- Environment variable reads.
|
||||
- API key resolution or auth header construction.
|
||||
- Logging, callbacks, spend tracking, retries, routing, cooldowns, or fallbacks.
|
||||
- Panics on bad user/provider input.
|
||||
|
||||
## Required Tests
|
||||
|
||||
Every provider transform must include focused unit tests for:
|
||||
- Supported params matching the Python provider config.
|
||||
- Unknown params being dropped or transformed the same way as Python.
|
||||
- Request body shape matching Python output.
|
||||
- Response normalization with complete, missing, null, and extra fields.
|
||||
- Bad input returning typed errors.
|
||||
|
||||
For OCR specifically, assume documents can contain personal data. Tests should
|
||||
prove transforms do not copy document contents into error messages.
|
||||
|
||||
## Implementation Rules
|
||||
|
||||
- Prefer static supported-parameter lists over allocating strings on every call.
|
||||
- Keep transforms deterministic and allocation-conscious, but choose clarity over
|
||||
premature micro-optimization for tiny parameter lists.
|
||||
- Use typed errors from `core`; avoid stringly-typed error plumbing.
|
||||
- Add comments only when they explain Python-parity decisions or provider quirks.
|
||||
- Put route-level provider dispatch in a route file such as `providers/src/ocr.rs`.
|
||||
Do not move provider-specific transform logic into the Python bridge.
|
||||
14
litellm-rust/crates/providers/Cargo.toml
Normal file
14
litellm-rust/crates/providers/Cargo.toml
Normal file
|
|
@ -0,0 +1,14 @@
|
|||
[package]
|
||||
name = "litellm-providers"
|
||||
version = "0.1.0"
|
||||
edition.workspace = true
|
||||
license.workspace = true
|
||||
repository.workspace = true
|
||||
|
||||
[dependencies]
|
||||
litellm-core.workspace = true
|
||||
reqwest.workspace = true
|
||||
serde_json.workspace = true
|
||||
|
||||
[dev-dependencies]
|
||||
serde_json.workspace = true
|
||||
2
litellm-rust/crates/providers/src/lib.rs
Normal file
2
litellm-rust/crates/providers/src/lib.rs
Normal file
|
|
@ -0,0 +1,2 @@
|
|||
pub mod mistral;
|
||||
pub mod ocr;
|
||||
1
litellm-rust/crates/providers/src/mistral/mod.rs
Normal file
1
litellm-rust/crates/providers/src/mistral/mod.rs
Normal file
|
|
@ -0,0 +1 @@
|
|||
pub mod ocr;
|
||||
1
litellm-rust/crates/providers/src/mistral/ocr/mod.rs
Normal file
1
litellm-rust/crates/providers/src/mistral/ocr/mod.rs
Normal file
|
|
@ -0,0 +1 @@
|
|||
pub mod transformation;
|
||||
292
litellm-rust/crates/providers/src/mistral/ocr/transformation.rs
Normal file
292
litellm-rust/crates/providers/src/mistral/ocr/transformation.rs
Normal file
|
|
@ -0,0 +1,292 @@
|
|||
use litellm_core::error::{json_type_name, CoreError, CoreResult};
|
||||
use litellm_core::ocr::transformation::OcrProviderConfig;
|
||||
use litellm_core::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()));
|
||||
}
|
||||
}
|
||||
127
litellm-rust/crates/providers/src/ocr.rs
Normal file
127
litellm-rust/crates/providers/src/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 crate::mistral::ocr::transformation as mistral;
|
||||
use crate::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()));
|
||||
}
|
||||
}
|
||||
36
litellm-rust/crates/python-bridge/CLAUDE.md
Normal file
36
litellm-rust/crates/python-bridge/CLAUDE.md
Normal file
|
|
@ -0,0 +1,36 @@
|
|||
# CLAUDE.md
|
||||
|
||||
Rules for `litellm-rust/crates/python-bridge`.
|
||||
|
||||
## Responsibility
|
||||
|
||||
`python-bridge` is the PyO3 boundary between Python LiteLLM and Rust transforms.
|
||||
Keep this crate thin. It adapts Python objects to Rust payloads and returns
|
||||
Python-compatible dictionaries.
|
||||
|
||||
## Bridge Shape
|
||||
|
||||
- Prefer one stable method per top-level LiteLLM route, for example
|
||||
`ocr(payload)`.
|
||||
- Do not add one exported PyO3 function per provider helper unless there is a
|
||||
measured reason.
|
||||
- Provider dispatch belongs in Rust route modules such as
|
||||
`litellm_providers::ocr`, not in this PyO3 crate.
|
||||
- Python owns rollout state and fallback. Rust should return errors; Python
|
||||
decides whether to raise or fall back.
|
||||
|
||||
## Data Handling
|
||||
|
||||
- OCR payloads can contain personal data and large base64 images. Do not log
|
||||
payloads or provider responses.
|
||||
- Avoid copying large payloads more than needed. The current JSON round-trip is
|
||||
acceptable for the first scaffold, but future performance work should evaluate
|
||||
direct PyO3 conversion before expanding Rust coverage to image-heavy paths.
|
||||
- Do not expose raw Rust errors that include document contents or upstream
|
||||
bodies.
|
||||
|
||||
## Tests
|
||||
|
||||
- `cargo test --workspace` must compile this crate.
|
||||
- Python tests must cover bridge disabled, bridge enabled, and module-missing
|
||||
fallback behavior for every exposed route.
|
||||
16
litellm-rust/crates/python-bridge/Cargo.toml
Normal file
16
litellm-rust/crates/python-bridge/Cargo.toml
Normal file
|
|
@ -0,0 +1,16 @@
|
|||
[package]
|
||||
name = "litellm-python-bridge"
|
||||
version = "0.1.0"
|
||||
edition.workspace = true
|
||||
license.workspace = true
|
||||
repository.workspace = true
|
||||
|
||||
[lib]
|
||||
name = "litellm_python_bridge"
|
||||
crate-type = ["cdylib"]
|
||||
|
||||
[dependencies]
|
||||
litellm-core.workspace = true
|
||||
litellm-providers.workspace = true
|
||||
pyo3 = { workspace = true, features = ["extension-module"] }
|
||||
serde_json.workspace = true
|
||||
32
litellm-rust/crates/python-bridge/src/gil.rs
Normal file
32
litellm-rust/crates/python-bridge/src/gil.rs
Normal file
|
|
@ -0,0 +1,32 @@
|
|||
//! GIL accounting.
|
||||
//!
|
||||
//! A single chokepoint for releasing the GIL around blocking work. Every
|
||||
//! blocking call in the bridge goes through [`release_gil`] instead of calling
|
||||
//! `Python::allow_threads` directly, so the release count stays accurate and we
|
||||
//! have one place to extend later (timing histograms, per-call labels, etc.).
|
||||
|
||||
use std::sync::atomic::{AtomicU64, Ordering};
|
||||
|
||||
use pyo3::prelude::*;
|
||||
|
||||
/// Number of times the bridge has released the GIL since process start.
|
||||
static GIL_RELEASES: AtomicU64 = AtomicU64::new(0);
|
||||
|
||||
/// Release the GIL around `f`, recording the release.
|
||||
///
|
||||
/// `f` must not touch any Python state — that is what makes releasing the GIL
|
||||
/// safe. Returning the value back to Python re-acquires the GIL at the call
|
||||
/// site, after `f` has finished.
|
||||
pub fn release_gil<T, F>(py: Python<'_>, f: F) -> T
|
||||
where
|
||||
F: FnOnce() -> T + Send,
|
||||
T: Send,
|
||||
{
|
||||
GIL_RELEASES.fetch_add(1, Ordering::Relaxed);
|
||||
py.allow_threads(f)
|
||||
}
|
||||
|
||||
/// Total GIL releases performed by the bridge so far.
|
||||
pub fn release_count() -> u64 {
|
||||
GIL_RELEASES.load(Ordering::Relaxed)
|
||||
}
|
||||
100
litellm-rust/crates/python-bridge/src/lib.rs
Normal file
100
litellm-rust/crates/python-bridge/src/lib.rs
Normal file
|
|
@ -0,0 +1,100 @@
|
|||
use std::time::Duration;
|
||||
|
||||
use litellm_core::error::CoreError;
|
||||
use litellm_providers::ocr::run_ocr;
|
||||
use pyo3::exceptions::{PyRuntimeError, PyValueError};
|
||||
use pyo3::prelude::*;
|
||||
use pyo3::types::{PyAny, PyDict};
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
mod gil;
|
||||
|
||||
fn py_to_json(py: Python<'_>, value: &Bound<'_, PyAny>) -> PyResult<Value> {
|
||||
let json = py.import("json")?;
|
||||
let encoded: String = json.call_method1("dumps", (value,))?.extract()?;
|
||||
serde_json::from_str(&encoded).map_err(|err| PyValueError::new_err(err.to_string()))
|
||||
}
|
||||
|
||||
fn json_to_py(py: Python<'_>, value: Value) -> PyResult<Py<PyAny>> {
|
||||
let json = py.import("json")?;
|
||||
let encoded =
|
||||
serde_json::to_string(&value).map_err(|err| PyValueError::new_err(err.to_string()))?;
|
||||
Ok(json.call_method1("loads", (encoded,))?.unbind())
|
||||
}
|
||||
|
||||
/// Map a core error to the closest Python exception. Caller-input problems
|
||||
/// (auth, bad types, missing fields) -> `ValueError`; everything else
|
||||
/// (network, upstream status, parse failures) -> `RuntimeError`.
|
||||
fn core_error_to_pyerr(err: CoreError) -> PyErr {
|
||||
match err {
|
||||
CoreError::Auth(message) => PyValueError::new_err(message),
|
||||
CoreError::InvalidType { .. } | CoreError::MissingField(_) => {
|
||||
PyValueError::new_err(err.to_string())
|
||||
}
|
||||
other => PyRuntimeError::new_err(other.to_string()),
|
||||
}
|
||||
}
|
||||
|
||||
/// Perform a Mistral OCR call end to end and return the response as a dict.
|
||||
#[pyfunction]
|
||||
#[pyo3(signature = (model, document, api_key=None, api_base=None, optional_params=None, timeout_seconds=None))]
|
||||
fn ocr(
|
||||
py: Python<'_>,
|
||||
model: String,
|
||||
document: Py<PyAny>,
|
||||
api_key: Option<String>,
|
||||
api_base: Option<String>,
|
||||
optional_params: Option<Py<PyAny>>,
|
||||
timeout_seconds: Option<f64>,
|
||||
) -> PyResult<Py<PyAny>> {
|
||||
let document = py_to_json(py, document.bind(py))?;
|
||||
|
||||
let optional_params = match optional_params {
|
||||
Some(params) => match py_to_json(py, params.bind(py))? {
|
||||
Value::Object(map) => map,
|
||||
_ => return Err(PyValueError::new_err("optional_params must be a dict")),
|
||||
},
|
||||
None => Map::new(),
|
||||
};
|
||||
|
||||
let timeout = timeout_seconds.and_then(|secs| {
|
||||
if secs.is_finite() && secs > 0.0 {
|
||||
Some(Duration::from_secs_f64(secs))
|
||||
} else {
|
||||
None
|
||||
}
|
||||
});
|
||||
|
||||
// Release the GIL during the blocking HTTP call (counted for observability).
|
||||
let result = gil::release_gil(py, || {
|
||||
run_ocr(
|
||||
&model,
|
||||
document,
|
||||
api_key.as_deref(),
|
||||
api_base.as_deref(),
|
||||
optional_params,
|
||||
timeout,
|
||||
)
|
||||
});
|
||||
|
||||
match result {
|
||||
Ok(value) => json_to_py(py, value),
|
||||
Err(err) => Err(core_error_to_pyerr(err)),
|
||||
}
|
||||
}
|
||||
|
||||
/// Bridge GIL accounting, e.g. `{"releases": 12}`. Lets the Python side observe
|
||||
/// how often the bridge has dropped the GIL for blocking work.
|
||||
#[pyfunction]
|
||||
fn gil_stats(py: Python<'_>) -> PyResult<Py<PyAny>> {
|
||||
let stats = PyDict::new(py);
|
||||
stats.set_item("releases", gil::release_count())?;
|
||||
Ok(stats.into_any().unbind())
|
||||
}
|
||||
|
||||
#[pymodule]
|
||||
fn litellm_python_bridge(module: &Bound<'_, PyModule>) -> PyResult<()> {
|
||||
module.add_function(wrap_pyfunction!(ocr, module)?)?;
|
||||
module.add_function(wrap_pyfunction!(gil_stats, module)?)?;
|
||||
Ok(())
|
||||
}
|
||||
|
|
@ -80,6 +80,7 @@ from litellm.constants import (
|
|||
WANDB_MODELS,
|
||||
REPEATED_STREAMING_CHUNK_LIMIT,
|
||||
request_timeout,
|
||||
request_timeout_explicitly_set as request_timeout_explicitly_set,
|
||||
open_ai_embedding_models,
|
||||
cohere_embedding_models,
|
||||
bedrock_embedding_models,
|
||||
|
|
@ -213,6 +214,15 @@ standard_logging_payload_excluded_fields: Optional[List[str]] = (
|
|||
log_raw_request_response: bool = False
|
||||
redact_messages_in_exceptions: Optional[bool] = False
|
||||
redact_user_api_key_info: Optional[bool] = False
|
||||
# When True (default — preserves historical behavior), the Router appends
|
||||
# internal config names (model_group, fallback model groups, deployment
|
||||
# timeouts, fallback failure details) onto exception messages and surfaces
|
||||
# them to clients via ProxyException.message. Set to False if you do NOT
|
||||
# want the proxy's internal model_name / fallback wiring visible to clients.
|
||||
# Deprecation: planned to flip to False (redact by default) in a future
|
||||
# major release; opt in early with `litellm.expose_router_debug_in_errors
|
||||
# = False`.
|
||||
expose_router_debug_in_errors: bool = True
|
||||
filter_invalid_headers: Optional[bool] = False
|
||||
add_user_information_to_llm_headers: Optional[bool] = (
|
||||
None # adds user_id, team_id, token hash (params from StandardLoggingMetadata) to request headers
|
||||
|
|
@ -235,6 +245,17 @@ modify_params = bool(os.getenv("LITELLM_MODIFY_PARAMS", False))
|
|||
use_chat_completions_url_for_anthropic_messages: bool = bool(
|
||||
os.getenv("LITELLM_USE_CHAT_COMPLETIONS_URL_FOR_ANTHROPIC_MESSAGES", False)
|
||||
) # When True, routes OpenAI /v1/messages requests to chat/completions instead of the Responses API
|
||||
# When True, strip the OpenAI-flavored `usage.total_tokens` field that
|
||||
# LiteLLM injects into non-streaming /v1/messages responses, bringing the
|
||||
# wire response into line with the Anthropic spec (matches the streaming
|
||||
# SSE path, which already omits total_tokens). Default False to preserve
|
||||
# backward compatibility for clients that read the LiteLLM-shaped
|
||||
# `usage.total_tokens` today. Planned to flip to True in a future major
|
||||
# release; opt in early via Python:
|
||||
# `litellm.strip_anthropic_total_tokens = True`
|
||||
# Or via `litellm_settings.strip_anthropic_total_tokens: true` in
|
||||
# config.yaml.
|
||||
strip_anthropic_total_tokens: bool = False
|
||||
route_all_chat_openai_to_responses: bool = (
|
||||
os.getenv("LITELLM_ROUTE_ALL_CHAT_OPENAI_TO_RESPONSES", "false").lower() == "true"
|
||||
) # When True, routes all OpenAI /chat/completions requests through the Responses API bridge
|
||||
|
|
@ -413,7 +434,7 @@ anthropic_beta_headers_url: str = os.getenv(
|
|||
"LITELLM_ANTHROPIC_BETA_HEADERS_URL",
|
||||
"https://raw.githubusercontent.com/BerriAI/litellm/main/litellm/anthropic_beta_headers_config.json",
|
||||
)
|
||||
suppress_debug_info = False
|
||||
suppress_debug_info: bool = False
|
||||
dynamodb_table_name: Optional[str] = None
|
||||
s3_callback_params: Optional[Dict] = None
|
||||
s3_audit_callback_params: Optional[Dict] = None
|
||||
|
|
@ -653,6 +674,7 @@ elevenlabs_models: Set = set()
|
|||
dashscope_models: Set = set()
|
||||
moonshot_models: Set = set()
|
||||
publicai_models: Set = set()
|
||||
darkbloom_models: Set = set()
|
||||
v0_models: Set = set()
|
||||
morph_models: Set = set()
|
||||
lambda_ai_models: Set = set()
|
||||
|
|
@ -907,6 +929,8 @@ def add_known_models(model_cost_map: Optional[Dict] = None):
|
|||
moonshot_models.add(key)
|
||||
elif value.get("litellm_provider") == "publicai":
|
||||
publicai_models.add(key)
|
||||
elif value.get("litellm_provider") == "darkbloom":
|
||||
darkbloom_models.add(key)
|
||||
elif value.get("litellm_provider") == "v0":
|
||||
v0_models.add(key)
|
||||
elif value.get("litellm_provider") == "morph":
|
||||
|
|
@ -1055,6 +1079,7 @@ model_list = list(
|
|||
| dashscope_models
|
||||
| moonshot_models
|
||||
| publicai_models
|
||||
| darkbloom_models
|
||||
| v0_models
|
||||
| morph_models
|
||||
| lambda_ai_models
|
||||
|
|
@ -1159,6 +1184,7 @@ models_by_provider: dict = {
|
|||
"modelscope": modelscope_models,
|
||||
"moonshot": moonshot_models,
|
||||
"publicai": publicai_models,
|
||||
"darkbloom": darkbloom_models,
|
||||
"v0": v0_models,
|
||||
"morph": morph_models,
|
||||
"lambda_ai": lambda_ai_models,
|
||||
|
|
@ -1380,7 +1406,9 @@ from .skills.main import (
|
|||
)
|
||||
from .containers.main import *
|
||||
from .ocr.main import *
|
||||
from .ocr.rust_bridge import use_litellm_rust
|
||||
from .rag.main import *
|
||||
from .sandbox.main import *
|
||||
from .search.main import *
|
||||
from .realtime_api.main import (
|
||||
_arealtime,
|
||||
|
|
@ -1901,9 +1929,6 @@ if TYPE_CHECKING:
|
|||
from .llms.fireworks_ai.completion.transformation import (
|
||||
FireworksAITextCompletionConfig as FireworksAITextCompletionConfig,
|
||||
)
|
||||
from .llms.fireworks_ai.audio_transcription.transformation import (
|
||||
FireworksAIAudioTranscriptionConfig as FireworksAIAudioTranscriptionConfig,
|
||||
)
|
||||
from .llms.fireworks_ai.embed.fireworks_ai_transformation import (
|
||||
FireworksAIEmbeddingConfig as FireworksAIEmbeddingConfig,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -260,7 +260,6 @@ LLM_CONFIG_NAMES = (
|
|||
"SambaNovaEmbeddingConfig",
|
||||
"FireworksAIConfig",
|
||||
"FireworksAITextCompletionConfig",
|
||||
"FireworksAIAudioTranscriptionConfig",
|
||||
"FireworksAIEmbeddingConfig",
|
||||
"FriendliaiChatConfig",
|
||||
"JinaAIEmbeddingConfig",
|
||||
|
|
@ -1027,10 +1026,6 @@ _LLM_CONFIGS_IMPORT_MAP = {
|
|||
".llms.fireworks_ai.completion.transformation",
|
||||
"FireworksAITextCompletionConfig",
|
||||
),
|
||||
"FireworksAIAudioTranscriptionConfig": (
|
||||
".llms.fireworks_ai.audio_transcription.transformation",
|
||||
"FireworksAIAudioTranscriptionConfig",
|
||||
),
|
||||
"FireworksAIEmbeddingConfig": (
|
||||
".llms.fireworks_ai.embed.fireworks_ai_transformation",
|
||||
"FireworksAIEmbeddingConfig",
|
||||
|
|
|
|||
|
|
@ -100,6 +100,8 @@ class Cache:
|
|||
gcs_path: Optional[str] = None,
|
||||
redis_semantic_cache_embedding_model: str = "text-embedding-ada-002",
|
||||
redis_semantic_cache_index_name: Optional[str] = None,
|
||||
valkey_semantic_cache_embedding_model: str = "text-embedding-ada-002",
|
||||
valkey_semantic_cache_index_name: str | None = None,
|
||||
redis_flush_size: Optional[int] = None,
|
||||
redis_startup_nodes: Optional[List] = None,
|
||||
disk_cache_dir: Optional[str] = None,
|
||||
|
|
@ -208,6 +210,21 @@ class Cache:
|
|||
index_name=redis_semantic_cache_index_name,
|
||||
**kwargs,
|
||||
)
|
||||
elif type == LiteLLMCacheType.VALKEY_SEMANTIC:
|
||||
# Imported here, not at module top, so the optional redis dependency
|
||||
# is only required when this backend is actually selected.
|
||||
from .valkey_semantic_cache import ValkeySemanticCache
|
||||
|
||||
self.cache = ValkeySemanticCache(
|
||||
host=host,
|
||||
port=port,
|
||||
password=password,
|
||||
similarity_threshold=similarity_threshold,
|
||||
embedding_model=valkey_semantic_cache_embedding_model,
|
||||
index_name=valkey_semantic_cache_index_name,
|
||||
startup_nodes=redis_startup_nodes,
|
||||
**kwargs,
|
||||
)
|
||||
elif type == LiteLLMCacheType.QDRANT_SEMANTIC:
|
||||
self.cache = QdrantSemanticCache(
|
||||
qdrant_api_base=qdrant_api_base,
|
||||
|
|
@ -267,12 +284,50 @@ class Cache:
|
|||
if (
|
||||
self.type == LiteLLMCacheType.REDIS
|
||||
or self.type == LiteLLMCacheType.REDIS_SEMANTIC
|
||||
or self.type == LiteLLMCacheType.VALKEY_SEMANTIC
|
||||
) and default_in_redis_ttl is not None:
|
||||
self.ttl = default_in_redis_ttl
|
||||
|
||||
if self.namespace is not None and isinstance(self.cache, RedisCache):
|
||||
self.cache.namespace = self.namespace
|
||||
|
||||
# Params whose values carry prompt content. Excluded from semantic-cache
|
||||
# scope keys so differently worded prompts share a bucket and match via
|
||||
# vector similarity rather than being split into per-wording buckets.
|
||||
_SEMANTIC_CACHE_SCOPE_EXCLUDED_PARAMS: frozenset = frozenset(
|
||||
{"messages", "prompt", "input"}
|
||||
)
|
||||
|
||||
# Server-set identity (from proxy auth) used to isolate semantic-cache
|
||||
# buckets per tenant. Required once the prompt is out of the scope key, so a
|
||||
# similar prompt from another key/team/org stays in a separate bucket.
|
||||
_SEMANTIC_CACHE_TENANT_SCOPE_FIELDS: tuple[str, ...] = (
|
||||
"user_api_key",
|
||||
"user_api_key_team_id",
|
||||
"user_api_key_org_id",
|
||||
)
|
||||
|
||||
def _is_semantic_cache(self) -> bool:
|
||||
return self.type in (
|
||||
LiteLLMCacheType.REDIS_SEMANTIC,
|
||||
LiteLLMCacheType.QDRANT_SEMANTIC,
|
||||
LiteLLMCacheType.VALKEY_SEMANTIC,
|
||||
)
|
||||
|
||||
def _get_semantic_cache_tenant_scope(self, kwargs: dict) -> str:
|
||||
metadata: dict = kwargs.get("metadata") or {}
|
||||
litellm_params: dict = kwargs.get("litellm_params") or {}
|
||||
metadata_in_litellm_params: dict = litellm_params.get("metadata") or {}
|
||||
|
||||
scope = ""
|
||||
for field in self._SEMANTIC_CACHE_TENANT_SCOPE_FIELDS:
|
||||
value = metadata.get(field)
|
||||
if value is None:
|
||||
value = metadata_in_litellm_params.get(field)
|
||||
if value is not None:
|
||||
scope += f"{field}: {value}"
|
||||
return scope
|
||||
|
||||
def get_cache_key(self, **kwargs) -> str:
|
||||
"""
|
||||
Get the cache key for the given arguments.
|
||||
|
|
@ -293,7 +348,15 @@ class Cache:
|
|||
|
||||
combined_kwargs = ModelParamHelper._get_all_llm_api_params()
|
||||
litellm_param_kwargs = all_litellm_params
|
||||
is_semantic_cache = self._is_semantic_cache()
|
||||
scope_excluded_params = (
|
||||
self._SEMANTIC_CACHE_SCOPE_EXCLUDED_PARAMS
|
||||
if is_semantic_cache
|
||||
else frozenset()
|
||||
)
|
||||
for param in kwargs:
|
||||
if param in scope_excluded_params:
|
||||
continue
|
||||
if param in combined_kwargs:
|
||||
param_value: Optional[str] = self._get_param_value(param, kwargs)
|
||||
if param_value is not None:
|
||||
|
|
@ -309,6 +372,9 @@ class Cache:
|
|||
param_value = kwargs[param]
|
||||
cache_key += f"{str(param)}: {str(param_value)}"
|
||||
|
||||
if is_semantic_cache:
|
||||
cache_key += self._get_semantic_cache_tenant_scope(kwargs)
|
||||
|
||||
hashed_cache_key = Cache._get_hashed_cache_key(cache_key)
|
||||
hashed_cache_key = self._add_namespace_to_cache_key(hashed_cache_key, **kwargs)
|
||||
verbose_logger.debug(
|
||||
|
|
|
|||
|
|
@ -5,6 +5,7 @@ Supports syncing responses to Google Cloud Storage Buckets using HTTP requests.
|
|||
import json
|
||||
import asyncio
|
||||
from typing import Optional
|
||||
from urllib.parse import quote
|
||||
|
||||
from litellm._logging import print_verbose, verbose_logger
|
||||
from litellm.integrations.gcs_bucket.gcs_bucket_base import GCSBucketBase
|
||||
|
|
@ -48,7 +49,7 @@ class GCSCache(BaseCache):
|
|||
headers = self._construct_headers()
|
||||
object_name = self.key_prefix + key
|
||||
bucket_name = self.bucket_name
|
||||
url = f"https://storage.googleapis.com/upload/storage/v1/b/{bucket_name}/o?uploadType=media&name={object_name}"
|
||||
url = f"https://storage.googleapis.com/upload/storage/v1/b/{bucket_name}/o?uploadType=media&name={quote(object_name, safe='')}"
|
||||
data = json.dumps(value)
|
||||
self.sync_client.post(url=url, data=data, headers=headers)
|
||||
except Exception as e:
|
||||
|
|
@ -59,7 +60,7 @@ class GCSCache(BaseCache):
|
|||
headers = self._construct_headers()
|
||||
object_name = self.key_prefix + key
|
||||
bucket_name = self.bucket_name
|
||||
url = f"https://storage.googleapis.com/upload/storage/v1/b/{bucket_name}/o?uploadType=media&name={object_name}"
|
||||
url = f"https://storage.googleapis.com/upload/storage/v1/b/{bucket_name}/o?uploadType=media&name={quote(object_name, safe='')}"
|
||||
data = json.dumps(value)
|
||||
await self.async_client.post(url=url, data=data, headers=headers)
|
||||
except Exception as e:
|
||||
|
|
@ -72,7 +73,7 @@ class GCSCache(BaseCache):
|
|||
headers = self._construct_headers()
|
||||
object_name = self.key_prefix + key
|
||||
bucket_name = self.bucket_name
|
||||
url = f"https://storage.googleapis.com/storage/v1/b/{bucket_name}/o/{object_name}?alt=media"
|
||||
url = f"https://storage.googleapis.com/storage/v1/b/{bucket_name}/o/{quote(object_name, safe='')}?alt=media"
|
||||
response = self.sync_client.get(url=url, headers=headers)
|
||||
if response.status_code == 200:
|
||||
cached_response = json.loads(response.text)
|
||||
|
|
@ -91,7 +92,7 @@ class GCSCache(BaseCache):
|
|||
headers = self._construct_headers()
|
||||
object_name = self.key_prefix + key
|
||||
bucket_name = self.bucket_name
|
||||
url = f"https://storage.googleapis.com/storage/v1/b/{bucket_name}/o/{object_name}?alt=media"
|
||||
url = f"https://storage.googleapis.com/storage/v1/b/{bucket_name}/o/{quote(object_name, safe='')}?alt=media"
|
||||
response = await self.async_client.get(url=url, headers=headers)
|
||||
if response.status_code == 200:
|
||||
return json.loads(response.text)
|
||||
|
|
|
|||
|
|
@ -903,6 +903,43 @@ class RedisCache(BaseCache):
|
|||
)
|
||||
raise e
|
||||
|
||||
@_redis_circuit_breaker_guard
|
||||
async def async_set_max(
|
||||
self,
|
||||
key: str,
|
||||
value: float,
|
||||
ttl: int | None = None,
|
||||
) -> float | None:
|
||||
"""Atomically set ``key`` to ``value`` only when ``value`` is greater
|
||||
than the stored value (or the key is unset), refreshing the TTL.
|
||||
|
||||
Monotonic by construction: it never lowers the stored value, so a repair
|
||||
that writes an authoritative-but-slightly-stale total cannot clobber a
|
||||
concurrent increment that has already pushed the counter higher. The
|
||||
GET/compare/SET runs in a single Lua call, so it is also atomic across
|
||||
racing callers and pods. Returns the resulting value.
|
||||
"""
|
||||
_redis_client = self.init_async_client()
|
||||
_used_ttl = self.get_ttl(ttl=ttl)
|
||||
key = self.check_and_fix_namespace(key=key)
|
||||
lua = (
|
||||
"local cur = redis.call('GET', KEYS[1]) "
|
||||
"if cur == false or tonumber(cur) < tonumber(ARGV[1]) then "
|
||||
"redis.call('SET', KEYS[1], ARGV[1]) "
|
||||
"if tonumber(ARGV[2]) > 0 then redis.call('EXPIRE', KEYS[1], ARGV[2]) end "
|
||||
"return ARGV[1] end "
|
||||
"return cur"
|
||||
)
|
||||
result = cast(
|
||||
"str | bytes | int | float | None",
|
||||
await _redis_client.eval(lua, 1, key, str(value), str(int(_used_ttl or 0))),
|
||||
)
|
||||
if result is None:
|
||||
return None
|
||||
if isinstance(result, bytes):
|
||||
result = result.decode()
|
||||
return float(result)
|
||||
|
||||
async def flush_cache_buffer(self):
|
||||
print_verbose(
|
||||
f"flushing to redis....reached size of buffer {len(self.redis_batch_writing_buffer)}"
|
||||
|
|
|
|||
353
litellm/caching/valkey_semantic_cache.py
Normal file
353
litellm/caching/valkey_semantic_cache.py
Normal file
|
|
@ -0,0 +1,353 @@
|
|||
"""
|
||||
Valkey Semantic Cache implementation for LiteLLM
|
||||
|
||||
Backs semantic caching with Valkey (for example AWS ElastiCache for Valkey)
|
||||
running the valkey-search module.
|
||||
|
||||
RedisVL cannot drive valkey-search: it gates on a RediSearch module version
|
||||
that valkey-search does not report, and its SemanticCache index uses a TEXT
|
||||
field that valkey-search does not implement. This backend therefore talks to
|
||||
valkey-search directly over redis-py, building a vector index from the field
|
||||
types valkey-search does support (TAG for cache-key isolation and VECTOR for
|
||||
the prompt embedding) and running KNN queries for retrieval. Prompt extraction,
|
||||
embedding generation, and cached-response parsing are reused from
|
||||
RedisSemanticCache since those are backend agnostic.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import hashlib
|
||||
import os
|
||||
import struct
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
|
||||
from redis import Redis
|
||||
from redis.asyncio import Redis as AsyncRedis
|
||||
from redis.commands.search.field import TagField, VectorField
|
||||
from redis.commands.search.indexDefinition import IndexDefinition, IndexType
|
||||
from redis.commands.search.query import Query
|
||||
|
||||
from litellm._logging import print_verbose
|
||||
from litellm._uuid import uuid
|
||||
|
||||
from .redis_semantic_cache import RedisSemanticCache
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _ValkeyCacheHit:
|
||||
response: str
|
||||
distance: float
|
||||
|
||||
|
||||
class ValkeySemanticCache(RedisSemanticCache):
|
||||
"""Valkey-backed semantic cache for LLM responses."""
|
||||
|
||||
DEFAULT_VALKEY_INDEX_NAME: str = "litellm_semantic_cache_index"
|
||||
EMBEDDING_FIELD_NAME: str = "embedding"
|
||||
PROMPT_FIELD_NAME: str = "prompt"
|
||||
RESPONSE_FIELD_NAME: str = "response"
|
||||
DISTANCE_FIELD_NAME: str = "vector_distance"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
host: str | None = None,
|
||||
port: str | None = None,
|
||||
password: str | None = None,
|
||||
redis_url: str | None = None,
|
||||
similarity_threshold: float | None = None,
|
||||
embedding_model: str = "text-embedding-ada-002",
|
||||
index_name: str | None = None,
|
||||
ssl: bool = False,
|
||||
startup_nodes: list | None = None,
|
||||
sync_client: Redis | None = None,
|
||||
async_client: AsyncRedis | None = None,
|
||||
**kwargs: Any,
|
||||
):
|
||||
if similarity_threshold is None:
|
||||
raise ValueError("similarity_threshold must be provided, passed None")
|
||||
|
||||
if startup_nodes:
|
||||
raise ValueError(
|
||||
"valkey-semantic does not support cluster-mode-enabled (multi-shard) "
|
||||
"endpoints. The async cluster client cannot route the FT.* search "
|
||||
"commands reliably. Point it at a cluster-mode-disabled endpoint "
|
||||
"instead (a primary with replicas is fine; only horizontal sharding "
|
||||
"is unsupported), or pass a single redis_url. On AWS, vector search "
|
||||
"needs ElastiCache for Valkey 8.2+ on a node-based cluster."
|
||||
)
|
||||
|
||||
self.similarity_threshold = similarity_threshold
|
||||
self.embedding_model = embedding_model
|
||||
self.index_name = index_name or self.DEFAULT_VALKEY_INDEX_NAME
|
||||
self.key_prefix = f"{self.index_name}:"
|
||||
self._index_dim: int | None = None
|
||||
|
||||
resolved_url = None
|
||||
if sync_client is None or async_client is None:
|
||||
resolved_url = redis_url or self._build_valkey_url(
|
||||
host, port, password, ssl
|
||||
)
|
||||
self.sync_client = (
|
||||
sync_client if sync_client is not None else Redis.from_url(resolved_url) # type: ignore[arg-type]
|
||||
)
|
||||
self.async_client = (
|
||||
async_client
|
||||
if async_client is not None
|
||||
else AsyncRedis.from_url(resolved_url) # type: ignore[arg-type]
|
||||
)
|
||||
|
||||
print_verbose(f"Valkey semantic-cache initializing index - {self.index_name}")
|
||||
|
||||
@staticmethod
|
||||
def _build_valkey_url(
|
||||
host: str | None, port: str | None, password: str | None, ssl: bool = False
|
||||
) -> str:
|
||||
host = host or os.environ.get("VALKEY_HOST") or os.environ.get("REDIS_HOST")
|
||||
port = port or os.environ.get("VALKEY_PORT") or os.environ.get("REDIS_PORT")
|
||||
password = (
|
||||
password
|
||||
or os.environ.get("VALKEY_PASSWORD")
|
||||
or os.environ.get("REDIS_PASSWORD")
|
||||
)
|
||||
|
||||
if not host or not port:
|
||||
raise ValueError(
|
||||
"Missing required Valkey configuration. Provide host and port "
|
||||
"(or VALKEY_HOST/VALKEY_PORT), or pass redis_url."
|
||||
)
|
||||
|
||||
credentials = f":{password}@" if password else ""
|
||||
scheme = "rediss" if ssl else "redis"
|
||||
return f"{scheme}://{credentials}{host}:{port}"
|
||||
|
||||
@classmethod
|
||||
def _scope_tag(cls, key: str) -> str:
|
||||
# valkey-search TAG fields tokenize on punctuation and do not honour
|
||||
# backslash escaping, so an arbitrary cache key cannot be matched
|
||||
# verbatim. Hashing to hex yields a token that is always exact-match
|
||||
# safe and still uniquely isolates a caller's scope.
|
||||
return hashlib.sha256(str(key).encode("utf-8")).hexdigest()
|
||||
|
||||
@staticmethod
|
||||
def _embedding_to_bytes(embedding: list[float]) -> bytes:
|
||||
return struct.pack(f"<{len(embedding)}f", *embedding)
|
||||
|
||||
def _index_schema(self, dim: int) -> tuple[TagField, VectorField]:
|
||||
return (
|
||||
TagField(self.CACHE_KEY_FIELD_NAME),
|
||||
VectorField(
|
||||
self.EMBEDDING_FIELD_NAME,
|
||||
"HNSW",
|
||||
{"TYPE": "FLOAT32", "DIM": dim, "DISTANCE_METRIC": "COSINE"},
|
||||
),
|
||||
)
|
||||
|
||||
def _index_definition(self) -> IndexDefinition:
|
||||
return IndexDefinition(prefix=[self.key_prefix], index_type=IndexType.HASH)
|
||||
|
||||
@staticmethod
|
||||
def _is_index_exists_error(exc: Exception) -> bool:
|
||||
return "already exists" in str(exc).lower()
|
||||
|
||||
@staticmethod
|
||||
def _extract_index_dim(info: dict) -> int | None:
|
||||
# FT.INFO nests the vector field's "dimensions" one level inside its
|
||||
# "index" block, so flatten each field descriptor a single level and
|
||||
# scan for the dimensions marker.
|
||||
for field in info.get("attributes") or []:
|
||||
if not isinstance(field, (list, tuple)):
|
||||
continue
|
||||
flat = [
|
||||
sub
|
||||
for item in field
|
||||
for sub in (item if isinstance(item, (list, tuple)) else [item])
|
||||
]
|
||||
for i, marker in enumerate(flat):
|
||||
if marker in (b"dimensions", "dimensions") and i + 1 < len(flat):
|
||||
return int(flat[i + 1])
|
||||
return None
|
||||
|
||||
def _assert_dim_matches(self, info: dict, dim: int) -> None:
|
||||
existing_dim = self._extract_index_dim(info)
|
||||
if existing_dim is not None and existing_dim != dim:
|
||||
raise ValueError(
|
||||
f"Valkey semantic-cache index '{self.index_name}' already exists with "
|
||||
f"embedding dimension {existing_dim}, but the configured embedding "
|
||||
f"model produced dimension {dim}. Use a different "
|
||||
f"valkey_semantic_cache_index_name or drop the existing index."
|
||||
)
|
||||
|
||||
def _ensure_index_sync(self, dim: int) -> None:
|
||||
if self._index_dim == dim:
|
||||
return
|
||||
try:
|
||||
self.sync_client.ft(self.index_name).create_index(
|
||||
self._index_schema(dim), definition=self._index_definition()
|
||||
)
|
||||
except Exception as exc:
|
||||
if not self._is_index_exists_error(exc):
|
||||
raise
|
||||
self._assert_dim_matches(self.sync_client.ft(self.index_name).info(), dim)
|
||||
self._index_dim = dim
|
||||
|
||||
async def _ensure_index_async(self, dim: int) -> None:
|
||||
if self._index_dim == dim:
|
||||
return
|
||||
try:
|
||||
await self.async_client.ft(self.index_name).create_index(
|
||||
self._index_schema(dim), definition=self._index_definition()
|
||||
)
|
||||
except Exception as exc:
|
||||
if not self._is_index_exists_error(exc):
|
||||
raise
|
||||
info = await self.async_client.ft(self.index_name).info()
|
||||
self._assert_dim_matches(info, dim)
|
||||
self._index_dim = dim
|
||||
|
||||
def _doc_key(self, key: str) -> str:
|
||||
return f"{self.key_prefix}{self._scope_tag(key)}:{uuid.uuid4()}"
|
||||
|
||||
def _doc_mapping(
|
||||
self, key: str, prompt: str, value_str: str, embedding: list[float]
|
||||
) -> dict:
|
||||
return {
|
||||
self.CACHE_KEY_FIELD_NAME: self._scope_tag(key),
|
||||
self.PROMPT_FIELD_NAME: prompt,
|
||||
self.RESPONSE_FIELD_NAME: value_str,
|
||||
self.EMBEDDING_FIELD_NAME: self._embedding_to_bytes(embedding),
|
||||
}
|
||||
|
||||
def _knn_query(self, key: str) -> Query:
|
||||
scope = self._scope_tag(key)
|
||||
query_string = (
|
||||
f"(@{self.CACHE_KEY_FIELD_NAME}:{{{scope}}})"
|
||||
f"=>[KNN 1 @{self.EMBEDDING_FIELD_NAME} $vec AS {self.DISTANCE_FIELD_NAME}]"
|
||||
)
|
||||
return (
|
||||
Query(query_string)
|
||||
.return_fields(self.RESPONSE_FIELD_NAME, self.DISTANCE_FIELD_NAME)
|
||||
.dialect(2)
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def _first_hit(cls, search_result: Any) -> _ValkeyCacheHit | None:
|
||||
docs = getattr(search_result, "docs", [])
|
||||
if not docs:
|
||||
return None
|
||||
doc = docs[0]
|
||||
return _ValkeyCacheHit(
|
||||
response=str(getattr(doc, cls.RESPONSE_FIELD_NAME)),
|
||||
distance=float(getattr(doc, cls.DISTANCE_FIELD_NAME)),
|
||||
)
|
||||
|
||||
def _resolve_hit(self, hit: _ValkeyCacheHit | None, key: str, **kwargs: Any) -> Any:
|
||||
if hit is None:
|
||||
kwargs.setdefault("metadata", {})["semantic-similarity"] = 0.0
|
||||
return None
|
||||
|
||||
similarity = 1 - hit.distance
|
||||
kwargs.setdefault("metadata", {})["semantic-similarity"] = similarity
|
||||
|
||||
if similarity < self.similarity_threshold:
|
||||
return None
|
||||
return self._get_cache_logic(cached_response=hit.response)
|
||||
|
||||
def set_cache(self, key: str, value: Any, **kwargs: Any) -> None:
|
||||
print_verbose(f"Valkey semantic-cache set_cache, kwargs: {kwargs}")
|
||||
try:
|
||||
prompt = self._get_prompt_from_kwargs(**kwargs)
|
||||
if prompt is None:
|
||||
print_verbose("No prompt provided for semantic caching")
|
||||
return
|
||||
|
||||
embedding = self._get_embedding(prompt)
|
||||
self._ensure_index_sync(len(embedding))
|
||||
|
||||
doc_key = self._doc_key(key)
|
||||
self.sync_client.hset(
|
||||
doc_key, mapping=self._doc_mapping(key, prompt, str(value), embedding)
|
||||
)
|
||||
ttl = self._get_ttl(**kwargs)
|
||||
if ttl is not None:
|
||||
self.sync_client.expire(doc_key, ttl)
|
||||
except Exception as e:
|
||||
print_verbose(f"Error in Valkey semantic-cache set_cache: {str(e)}")
|
||||
|
||||
def get_cache(self, key: str, **kwargs: Any) -> Any:
|
||||
print_verbose(f"Valkey semantic-cache get_cache, kwargs: {kwargs}")
|
||||
try:
|
||||
prompt = self._get_prompt_from_kwargs(**kwargs)
|
||||
if prompt is None:
|
||||
kwargs.setdefault("metadata", {})["semantic-similarity"] = 0.0
|
||||
return None
|
||||
|
||||
embedding = self._get_embedding(prompt)
|
||||
self._ensure_index_sync(len(embedding))
|
||||
|
||||
search_result = self.sync_client.ft(self.index_name).search(
|
||||
self._knn_query(key),
|
||||
query_params={"vec": self._embedding_to_bytes(embedding)},
|
||||
)
|
||||
return self._resolve_hit(self._first_hit(search_result), key, **kwargs)
|
||||
except Exception as e:
|
||||
print_verbose(f"Error in Valkey semantic-cache get_cache: {str(e)}")
|
||||
kwargs.setdefault("metadata", {})["semantic-similarity"] = 0.0
|
||||
|
||||
async def async_set_cache(self, key: str, value: Any, **kwargs: Any) -> None:
|
||||
print_verbose(f"Async Valkey semantic-cache set_cache, kwargs: {kwargs}")
|
||||
try:
|
||||
prompt = self._get_prompt_from_kwargs(**kwargs)
|
||||
if prompt is None:
|
||||
print_verbose("No prompt provided for semantic caching")
|
||||
return
|
||||
|
||||
embedding = await self._get_async_embedding(prompt, **kwargs)
|
||||
await self._ensure_index_async(len(embedding))
|
||||
|
||||
doc_key = self._doc_key(key)
|
||||
await self.async_client.hset(
|
||||
doc_key, mapping=self._doc_mapping(key, prompt, str(value), embedding)
|
||||
)
|
||||
ttl = self._get_ttl(**kwargs)
|
||||
if ttl is not None:
|
||||
await self.async_client.expire(doc_key, ttl)
|
||||
except Exception as e:
|
||||
print_verbose(f"Error in async Valkey semantic-cache set_cache: {str(e)}")
|
||||
|
||||
async def async_get_cache(self, key: str, **kwargs: Any) -> Any:
|
||||
print_verbose(f"Async Valkey semantic-cache get_cache, kwargs: {kwargs}")
|
||||
try:
|
||||
prompt = self._get_prompt_from_kwargs(**kwargs)
|
||||
if prompt is None:
|
||||
kwargs.setdefault("metadata", {})["semantic-similarity"] = 0.0
|
||||
return None
|
||||
|
||||
embedding = await self._get_async_embedding(prompt, **kwargs)
|
||||
await self._ensure_index_async(len(embedding))
|
||||
|
||||
search_result = await self.async_client.ft(self.index_name).search(
|
||||
self._knn_query(key),
|
||||
query_params={"vec": self._embedding_to_bytes(embedding)},
|
||||
)
|
||||
return self._resolve_hit(self._first_hit(search_result), key, **kwargs)
|
||||
except Exception as e:
|
||||
print_verbose(f"Error in async Valkey semantic-cache get_cache: {str(e)}")
|
||||
kwargs.setdefault("metadata", {})["semantic-similarity"] = 0.0
|
||||
|
||||
async def async_set_cache_pipeline(
|
||||
self, cache_list: list[tuple[str, Any]], **kwargs: Any
|
||||
) -> None:
|
||||
try:
|
||||
await asyncio.gather(
|
||||
*[
|
||||
self.async_set_cache(key, value, **kwargs)
|
||||
for key, value in cache_list
|
||||
]
|
||||
)
|
||||
except Exception as e:
|
||||
print_verbose(
|
||||
f"Error in Valkey semantic-cache async_set_cache_pipeline: {str(e)}"
|
||||
)
|
||||
|
||||
async def _index_info(self) -> dict:
|
||||
return await self.async_client.ft(self.index_name).info()
|
||||
|
|
@ -201,6 +201,18 @@ DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET = int(
|
|||
|
||||
# Provider-specific API base URLs
|
||||
XAI_API_BASE = "https://api.x.ai/v1"
|
||||
OPEN_SANDBOX_API_BASE_ENV_VAR = "OPEN_SANDBOX_API_BASE"
|
||||
OPEN_SANDBOX_API_KEY_ENV_VAR = "OPEN_SANDBOX_API_KEY"
|
||||
OPEN_SANDBOX_DEFAULT_TEMPLATE = "opensandbox/code-interpreter:v1.1.0"
|
||||
_OPEN_SANDBOX_FALLBACK_ENTRYPOINT = "/opt/code-interpreter/code-interpreter.sh"
|
||||
OPEN_SANDBOX_DEFAULT_ENTRYPOINT = (_OPEN_SANDBOX_FALLBACK_ENTRYPOINT,)
|
||||
OPEN_SANDBOX_DEFAULT_LANGUAGE = "python"
|
||||
OPEN_SANDBOX_DEFAULT_CPU_LIMIT = "1"
|
||||
OPEN_SANDBOX_DEFAULT_MEMORY_LIMIT = "2Gi"
|
||||
OPEN_SANDBOX_EXECD_PORT = 44772
|
||||
OPEN_SANDBOX_DEFAULT_TIMEOUT = 300
|
||||
OPEN_SANDBOX_READY_TIMEOUT = 30.0
|
||||
OPEN_SANDBOX_POLL_INTERVAL = 0.2
|
||||
|
||||
DEFAULT_REASONING_EFFORT_LOW_THINKING_BUDGET = int(
|
||||
os.getenv("DEFAULT_REASONING_EFFORT_LOW_THINKING_BUDGET", 1024)
|
||||
|
|
@ -456,6 +468,7 @@ HTTP_HANDLER_CONNECT_TIMEOUT_SECONDS: float = 5.0
|
|||
request_timeout: float = float(
|
||||
os.getenv("REQUEST_TIMEOUT", str(int(DEFAULT_REQUEST_TIMEOUT_SECONDS)))
|
||||
)
|
||||
request_timeout_explicitly_set: bool = "REQUEST_TIMEOUT" in os.environ
|
||||
DEFAULT_A2A_AGENT_TIMEOUT: float = float(
|
||||
os.getenv("DEFAULT_A2A_AGENT_TIMEOUT", 6000)
|
||||
) # 10 minutes
|
||||
|
|
@ -802,6 +815,7 @@ openai_compatible_endpoints: List = [
|
|||
"https://api.inference.wandb.ai/v1",
|
||||
"https://api.clarifai.com/v2/ext/openai/v1",
|
||||
"https://api.libertai.io/v1",
|
||||
"https://pinstripes.io/v1",
|
||||
]
|
||||
|
||||
|
||||
|
|
@ -865,6 +879,8 @@ openai_compatible_providers: List = [
|
|||
"clarifai",
|
||||
"docker_model_runner",
|
||||
"ragflow",
|
||||
"pinstripes", # Pinstripes - JSON-configured provider
|
||||
"darkbloom",
|
||||
]
|
||||
openai_text_completion_compatible_providers: List = (
|
||||
[ # providers that support `/v1/completions`
|
||||
|
|
|
|||
|
|
@ -888,6 +888,23 @@ def _map_traffic_type_to_service_tier(traffic_type: Optional[str]) -> Optional[s
|
|||
return service_tier
|
||||
|
||||
|
||||
def _normalize_service_tier(service_tier: object) -> str | None:
|
||||
"""
|
||||
Reduce a service_tier value to a concrete billable tier string or None.
|
||||
|
||||
"auto" is a routing preference and any non-string value is not a billable
|
||||
tier, so both defer to standard pricing (or to the tier the provider reports
|
||||
on the response usage) instead of crashing the downstream cost-key lookup,
|
||||
which calls service_tier.lower()
|
||||
"""
|
||||
if (
|
||||
not isinstance(service_tier, str)
|
||||
or service_tier.lower() == ServiceTier.AUTO.value
|
||||
):
|
||||
return None
|
||||
return service_tier
|
||||
|
||||
|
||||
def _get_usage_object(
|
||||
completion_response: Any,
|
||||
) -> Optional[Usage]:
|
||||
|
|
@ -1227,15 +1244,7 @@ def completion_cost(
|
|||
if service_tier is None and optional_params is not None:
|
||||
service_tier = optional_params.get("service_tier")
|
||||
|
||||
# A request-level service_tier only prices the request when it is a
|
||||
# concrete billable tier string. "auto" is a routing preference and any
|
||||
# non-string value is not a billable tier, so defer to the tier the
|
||||
# provider reports on the response/usage instead of crashing or mispricing
|
||||
if (
|
||||
not isinstance(service_tier, str)
|
||||
or service_tier.lower() == ServiceTier.AUTO.value
|
||||
):
|
||||
service_tier = None
|
||||
service_tier = _normalize_service_tier(service_tier)
|
||||
|
||||
# Extract service_tier from completion_response if not provided
|
||||
if service_tier is None and completion_response is not None:
|
||||
|
|
@ -1244,6 +1253,8 @@ def completion_cost(
|
|||
elif isinstance(completion_response, dict):
|
||||
service_tier = completion_response.get("service_tier")
|
||||
|
||||
service_tier = _normalize_service_tier(service_tier)
|
||||
|
||||
# Extract service_tier from usage object if not provided
|
||||
if service_tier is None and cost_per_token_usage_object is not None:
|
||||
if isinstance(cost_per_token_usage_object, BaseModel):
|
||||
|
|
@ -1253,6 +1264,8 @@ def completion_cost(
|
|||
elif isinstance(cost_per_token_usage_object, dict):
|
||||
service_tier = cost_per_token_usage_object.get("service_tier")
|
||||
|
||||
service_tier = _normalize_service_tier(service_tier)
|
||||
|
||||
selected_model = _select_model_name_for_cost_calc(
|
||||
model=model,
|
||||
completion_response=completion_response,
|
||||
|
|
|
|||
|
|
@ -0,0 +1,15 @@
|
|||
"""
|
||||
Code Interpreter Interception Module
|
||||
|
||||
Converts the native OpenAI Responses ``code_interpreter`` tool into a function
|
||||
tool, runs the model-emitted code in a sandbox, and feeds the result back into
|
||||
the agentic loop.
|
||||
"""
|
||||
|
||||
from litellm.integrations.code_interpreter_interception.handler import (
|
||||
CodeInterpreterInterceptionLogger,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"CodeInterpreterInterceptionLogger",
|
||||
]
|
||||
839
litellm/integrations/code_interpreter_interception/handler.py
Normal file
839
litellm/integrations/code_interpreter_interception/handler.py
Normal file
|
|
@ -0,0 +1,839 @@
|
|||
"""
|
||||
Code Interpreter Interception Handler
|
||||
|
||||
CustomLogger that swaps the native OpenAI Responses ``code_interpreter`` tool for
|
||||
a function tool, executes the code the model emits inside a sandbox, and feeds the
|
||||
captured stdout back through the typed agentic loop plan.
|
||||
"""
|
||||
|
||||
import json
|
||||
import time
|
||||
import uuid
|
||||
from typing import Any, Literal, TypedDict, cast
|
||||
|
||||
import litellm
|
||||
from pydantic import ValidationError
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.types.integrations.code_interpreter_interception import (
|
||||
CodeInterpreterInterceptionConfig,
|
||||
)
|
||||
from litellm.types.integrations.custom_logger import (
|
||||
AgenticLoopPlan,
|
||||
AgenticLoopRequestPatch,
|
||||
CHAT_COMPLETION_AGENTIC_SURFACE,
|
||||
NON_CODE_INTERPRETER_INTERCEPTION_INTERNAL_PREFIXES,
|
||||
is_interception_internal_key,
|
||||
)
|
||||
from litellm.types.llms.openai import (
|
||||
ChatCompletionAssistantMessage,
|
||||
ChatCompletionAssistantToolCall,
|
||||
ChatCompletionToolMessage,
|
||||
)
|
||||
from litellm.types.utils import (
|
||||
CallTypes,
|
||||
ChatCompletionMessageToolCall,
|
||||
ModelResponse,
|
||||
)
|
||||
|
||||
LITELLM_CODE_EXECUTION_TOOL_NAME = "litellm_code_execution"
|
||||
_INTERCEPTION_ACTIVE_KEY = "_code_interpreter_interception_active"
|
||||
_SANDBOX_KEY = "_code_interpreter_interception_sandbox_key"
|
||||
_CONVERTED_STREAM_KEY = "_code_interpreter_interception_converted_stream"
|
||||
_LITELLM_METADATA_KEY = "litellm_metadata"
|
||||
_CACHE_TTL_SECONDS = 15 * 60
|
||||
|
||||
|
||||
class CodeExecutionToolCall(TypedDict, total=False):
|
||||
id: str | None
|
||||
call_id: str | None
|
||||
type: Literal["function"]
|
||||
name: str
|
||||
arguments: str
|
||||
|
||||
|
||||
class CodeInterpreterLogOutput(TypedDict):
|
||||
type: Literal["logs"]
|
||||
logs: str
|
||||
|
||||
|
||||
class CodeInterpreterCall(TypedDict):
|
||||
id: str
|
||||
type: Literal["code_interpreter_call"]
|
||||
status: Literal["completed"]
|
||||
code: str
|
||||
container_id: str | None
|
||||
outputs: list[CodeInterpreterLogOutput]
|
||||
|
||||
|
||||
class CodeExecutionFunctionParameters(TypedDict):
|
||||
type: Literal["object"]
|
||||
properties: dict[str, dict[str, str]]
|
||||
required: list[str]
|
||||
|
||||
|
||||
class ResponsesFunctionTool(TypedDict):
|
||||
type: Literal["function"]
|
||||
name: str
|
||||
description: str
|
||||
parameters: CodeExecutionFunctionParameters
|
||||
|
||||
|
||||
class ChatCompletionFunctionDefinition(TypedDict):
|
||||
name: str
|
||||
description: str
|
||||
parameters: CodeExecutionFunctionParameters
|
||||
|
||||
|
||||
class ChatCompletionFunctionTool(TypedDict):
|
||||
type: Literal["function"]
|
||||
function: ChatCompletionFunctionDefinition
|
||||
|
||||
|
||||
CodeExecutionFunctionTool = ResponsesFunctionTool | ChatCompletionFunctionTool
|
||||
|
||||
|
||||
class ResponsesFunctionToolChoice(TypedDict):
|
||||
type: Literal["function"]
|
||||
name: str
|
||||
|
||||
|
||||
class ChatCompletionFunctionToolChoice(TypedDict):
|
||||
type: Literal["function"]
|
||||
function: dict[str, str]
|
||||
|
||||
|
||||
CodeExecutionFunctionToolChoice = (
|
||||
ResponsesFunctionToolChoice | ChatCompletionFunctionToolChoice
|
||||
)
|
||||
|
||||
|
||||
def _resolve_sandbox_tool(sandbox_tool_name: str | None) -> dict[str, Any] | None:
|
||||
try:
|
||||
from litellm.sandbox.sandbox_tools import resolve_sandbox_tool
|
||||
except ImportError:
|
||||
return None
|
||||
return resolve_sandbox_tool(sandbox_tool_name)
|
||||
|
||||
|
||||
class CodeInterpreterInterceptionLogger(CustomLogger):
|
||||
"""
|
||||
CustomLogger that implements transparent code-interpreter execution loops.
|
||||
|
||||
Flow:
|
||||
1. Replace the native ``code_interpreter`` tool with a function tool in the
|
||||
pre-call hook so the model emits code as function-call arguments.
|
||||
2. Detect ``litellm_code_execution`` function calls in the model response.
|
||||
3. Run the emitted code in a sandbox (reused per request via a server-minted
|
||||
sandbox key) and build a typed rerun plan that appends the
|
||||
function_call_output.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
enabled: bool = True,
|
||||
enabled_providers: list[str] | None = None,
|
||||
sandbox_tool_name: str | None = None,
|
||||
sandbox_config: Any | None = None,
|
||||
):
|
||||
super().__init__()
|
||||
self.enabled = enabled
|
||||
self.enabled_providers = enabled_providers
|
||||
self.sandbox_tool_name = sandbox_tool_name
|
||||
self.sandbox_config = sandbox_config
|
||||
self._container_cache: dict[str, tuple[Any, dict[str, Any] | None, float]] = {}
|
||||
|
||||
@classmethod
|
||||
def from_config_yaml(
|
||||
cls, config: CodeInterpreterInterceptionConfig
|
||||
) -> "CodeInterpreterInterceptionLogger":
|
||||
return cls(
|
||||
enabled=bool(config.get("enabled", True)),
|
||||
enabled_providers=config.get("enabled_providers"),
|
||||
sandbox_tool_name=config.get("sandbox_tool_name"),
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def initialize_from_proxy_config(
|
||||
litellm_settings: dict[str, Any],
|
||||
callback_specific_params: dict[str, Any],
|
||||
) -> "CodeInterpreterInterceptionLogger":
|
||||
params: CodeInterpreterInterceptionConfig = {}
|
||||
if "code_interpreter_interception_params" in litellm_settings:
|
||||
params = litellm_settings["code_interpreter_interception_params"]
|
||||
elif "code_interpreter_interception" in callback_specific_params and isinstance(
|
||||
callback_specific_params["code_interpreter_interception"], dict
|
||||
):
|
||||
params = cast(
|
||||
CodeInterpreterInterceptionConfig,
|
||||
callback_specific_params["code_interpreter_interception"],
|
||||
)
|
||||
return CodeInterpreterInterceptionLogger.from_config_yaml(params)
|
||||
|
||||
async def async_pre_call_deployment_hook(
|
||||
self, kwargs: dict[str, Any], call_type: CallTypes | None
|
||||
) -> dict | None:
|
||||
if not kwargs.get("_agentic_loop_depth"):
|
||||
kwargs.pop(_INTERCEPTION_ACTIVE_KEY, None)
|
||||
kwargs.pop(_SANDBOX_KEY, None)
|
||||
self._strip_interception_metadata(kwargs)
|
||||
if not self.enabled:
|
||||
return None
|
||||
if call_type not in (
|
||||
CallTypes.responses,
|
||||
CallTypes.aresponses,
|
||||
CallTypes.completion,
|
||||
CallTypes.acompletion,
|
||||
):
|
||||
return None
|
||||
if (
|
||||
self.enabled_providers is not None
|
||||
and self._resolve_provider(kwargs) not in self.enabled_providers
|
||||
):
|
||||
return None
|
||||
|
||||
tools = kwargs.get("tools")
|
||||
if not isinstance(tools, list):
|
||||
return None
|
||||
if not any(
|
||||
isinstance(tool, dict) and tool.get("type") == "code_interpreter"
|
||||
for tool in tools
|
||||
):
|
||||
return None
|
||||
|
||||
kwargs[_INTERCEPTION_ACTIVE_KEY] = True
|
||||
kwargs[_SANDBOX_KEY] = uuid.uuid4().hex
|
||||
if kwargs.get("stream"):
|
||||
kwargs["stream"] = False
|
||||
kwargs[_CONVERTED_STREAM_KEY] = True
|
||||
self._write_interception_metadata(kwargs)
|
||||
|
||||
function_tool = self._get_function_tool(call_type=call_type)
|
||||
kwargs["tools"] = [
|
||||
(
|
||||
function_tool
|
||||
if isinstance(tool, dict) and tool.get("type") == "code_interpreter"
|
||||
else tool
|
||||
)
|
||||
for tool in tools
|
||||
]
|
||||
if self._tool_choice_targets_code_interpreter(kwargs.get("tool_choice")):
|
||||
kwargs["tool_choice"] = self._get_function_tool_choice(call_type=call_type)
|
||||
return kwargs
|
||||
|
||||
@staticmethod
|
||||
def _strip_interception_metadata(kwargs: dict[str, Any]) -> None:
|
||||
metadata = kwargs.get(_LITELLM_METADATA_KEY)
|
||||
if not isinstance(metadata, dict):
|
||||
return
|
||||
filtered_metadata = {
|
||||
key: value
|
||||
for key, value in metadata.items()
|
||||
if not is_interception_internal_key(key)
|
||||
and not key.startswith("_agentic_loop")
|
||||
and key != "max_agentic_loops"
|
||||
}
|
||||
if filtered_metadata:
|
||||
kwargs[_LITELLM_METADATA_KEY] = filtered_metadata
|
||||
else:
|
||||
kwargs.pop(_LITELLM_METADATA_KEY, None)
|
||||
|
||||
@staticmethod
|
||||
def _write_interception_metadata(kwargs: dict[str, Any]) -> None:
|
||||
metadata = kwargs.get(_LITELLM_METADATA_KEY)
|
||||
metadata = dict(metadata) if isinstance(metadata, dict) else {}
|
||||
for key in (_INTERCEPTION_ACTIVE_KEY, _SANDBOX_KEY, _CONVERTED_STREAM_KEY):
|
||||
if key in kwargs:
|
||||
metadata[key] = kwargs[key]
|
||||
kwargs[_LITELLM_METADATA_KEY] = metadata
|
||||
|
||||
@staticmethod
|
||||
def _get_function_parameters() -> CodeExecutionFunctionParameters:
|
||||
return {
|
||||
"type": "object",
|
||||
"properties": {"code": {"type": "string"}},
|
||||
"required": ["code"],
|
||||
}
|
||||
|
||||
def _get_function_tool(
|
||||
self, call_type: CallTypes | None
|
||||
) -> CodeExecutionFunctionTool:
|
||||
description = "Execute python code in a sandbox and return stdout."
|
||||
if call_type in (CallTypes.completion, CallTypes.acompletion):
|
||||
return {
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": LITELLM_CODE_EXECUTION_TOOL_NAME,
|
||||
"description": description,
|
||||
"parameters": self._get_function_parameters(),
|
||||
},
|
||||
}
|
||||
return {
|
||||
"type": "function",
|
||||
"name": LITELLM_CODE_EXECUTION_TOOL_NAME,
|
||||
"description": description,
|
||||
"parameters": self._get_function_parameters(),
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def _get_function_tool_choice(
|
||||
call_type: CallTypes | None,
|
||||
) -> CodeExecutionFunctionToolChoice:
|
||||
if call_type in (CallTypes.completion, CallTypes.acompletion):
|
||||
return {
|
||||
"type": "function",
|
||||
"function": {"name": LITELLM_CODE_EXECUTION_TOOL_NAME},
|
||||
}
|
||||
return {
|
||||
"type": "function",
|
||||
"name": LITELLM_CODE_EXECUTION_TOOL_NAME,
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def _tool_choice_targets_code_interpreter(tool_choice: Any) -> bool:
|
||||
if not isinstance(tool_choice, dict):
|
||||
return False
|
||||
function = tool_choice.get("function")
|
||||
return (
|
||||
tool_choice.get("type") == "code_interpreter"
|
||||
or tool_choice.get("name") == "code_interpreter"
|
||||
or tool_choice.get("name") == LITELLM_CODE_EXECUTION_TOOL_NAME
|
||||
or (
|
||||
isinstance(function, dict)
|
||||
and function.get("name") == LITELLM_CODE_EXECUTION_TOOL_NAME
|
||||
)
|
||||
)
|
||||
|
||||
def _resolve_provider(self, kwargs: dict[str, Any]) -> str | None:
|
||||
provider = kwargs.get("custom_llm_provider")
|
||||
if provider:
|
||||
return provider
|
||||
model = kwargs.get("model")
|
||||
if not isinstance(model, str):
|
||||
return None
|
||||
try:
|
||||
return litellm.get_llm_provider(model=model)[1]
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
async def async_should_run_agentic_loop(
|
||||
self,
|
||||
response: Any,
|
||||
model: str,
|
||||
messages: list[dict],
|
||||
tools: list[dict] | None,
|
||||
stream: bool,
|
||||
custom_llm_provider: str,
|
||||
kwargs: dict,
|
||||
) -> tuple[bool, dict]:
|
||||
if not self.enabled:
|
||||
return False, {}
|
||||
if not kwargs.get(_INTERCEPTION_ACTIVE_KEY):
|
||||
return False, {}
|
||||
if (
|
||||
self.enabled_providers is not None
|
||||
and custom_llm_provider not in self.enabled_providers
|
||||
):
|
||||
return False, {}
|
||||
|
||||
tool_calls = (
|
||||
self._extract_chat_completion_code_execution_tool_calls(response=response)
|
||||
if kwargs.get("_agentic_loop_api_surface")
|
||||
== CHAT_COMPLETION_AGENTIC_SURFACE
|
||||
else self._extract_code_execution_tool_calls(response=response)
|
||||
)
|
||||
if not tool_calls:
|
||||
return False, {}
|
||||
|
||||
return True, {"tool_calls": tool_calls}
|
||||
|
||||
async def async_build_agentic_loop_plan(
|
||||
self,
|
||||
tools: dict,
|
||||
model: str,
|
||||
messages: list[dict],
|
||||
response: Any,
|
||||
anthropic_messages_provider_config: Any,
|
||||
anthropic_messages_optional_request_params: dict,
|
||||
logging_obj: Any,
|
||||
stream: bool,
|
||||
kwargs: dict,
|
||||
) -> AgenticLoopPlan:
|
||||
if kwargs.get("_agentic_loop_api_surface") == CHAT_COMPLETION_AGENTIC_SURFACE:
|
||||
return await self._build_chat_completion_agentic_loop_plan(
|
||||
tools=tools,
|
||||
model=model,
|
||||
messages=messages,
|
||||
optional_params=anthropic_messages_optional_request_params,
|
||||
kwargs=kwargs,
|
||||
)
|
||||
|
||||
await self._prune_expired_cache()
|
||||
tool_calls = cast(list[CodeExecutionToolCall], tools.get("tool_calls", []))
|
||||
sandbox_key = kwargs.get(_SANDBOX_KEY)
|
||||
container, params = await self._get_or_create_container(cache_key=sandbox_key)
|
||||
|
||||
try:
|
||||
container_id = cast(str | None, getattr(container, "id", None))
|
||||
input_list = self._normalize_messages(messages)
|
||||
code_interpreter_calls: list[CodeInterpreterCall] = []
|
||||
for tool_call in tool_calls:
|
||||
arguments = tool_call.get("arguments", "")
|
||||
code = self._parse_code(arguments)
|
||||
stdout = await self._run_tool_call(
|
||||
container=container, params=params, arguments=arguments
|
||||
)
|
||||
input_list.append(
|
||||
{
|
||||
"type": "function_call",
|
||||
"call_id": tool_call.get("call_id"),
|
||||
"name": LITELLM_CODE_EXECUTION_TOOL_NAME,
|
||||
"arguments": arguments,
|
||||
}
|
||||
)
|
||||
input_list.append(
|
||||
{
|
||||
"type": "function_call_output",
|
||||
"call_id": tool_call.get("call_id"),
|
||||
"output": stdout,
|
||||
}
|
||||
)
|
||||
code_interpreter_calls.append(
|
||||
{
|
||||
"id": f"ci_{uuid.uuid4().hex}",
|
||||
"type": "code_interpreter_call",
|
||||
"status": "completed",
|
||||
"code": code,
|
||||
"container_id": container_id,
|
||||
"outputs": (
|
||||
[{"type": "logs", "logs": stdout}] if stdout else []
|
||||
),
|
||||
}
|
||||
)
|
||||
except Exception:
|
||||
await self._delete_container_for_cache_key(sandbox_key)
|
||||
raise
|
||||
|
||||
optional_params = anthropic_messages_optional_request_params
|
||||
request_patch = AgenticLoopRequestPatch(
|
||||
model=model,
|
||||
messages=input_list,
|
||||
tools=self._get_followup_tools(
|
||||
tools=optional_params.get("tools"),
|
||||
call_type=CallTypes.responses,
|
||||
),
|
||||
optional_params=self._get_followup_optional_params(optional_params),
|
||||
kwargs=self._filter_agentic_loop_kwargs(kwargs),
|
||||
)
|
||||
|
||||
return AgenticLoopPlan(
|
||||
run_agentic_loop=True,
|
||||
request_patch=request_patch,
|
||||
metadata={
|
||||
"tool_type": "code_interpreter",
|
||||
"sandbox_key": sandbox_key or "",
|
||||
"code_interpreter_calls": code_interpreter_calls,
|
||||
},
|
||||
)
|
||||
|
||||
async def _build_chat_completion_agentic_loop_plan(
|
||||
self,
|
||||
tools: dict[str, object],
|
||||
model: str,
|
||||
messages: list[dict],
|
||||
optional_params: dict[str, object],
|
||||
kwargs: dict[str, object],
|
||||
) -> AgenticLoopPlan:
|
||||
await self._prune_expired_cache()
|
||||
tool_calls = cast(list[CodeExecutionToolCall], tools.get("tool_calls", []))
|
||||
sandbox_key = cast(str | None, kwargs.get(_SANDBOX_KEY))
|
||||
container, params = await self._get_or_create_container(cache_key=sandbox_key)
|
||||
|
||||
try:
|
||||
container_id = cast(str | None, getattr(container, "id", None))
|
||||
tool_results = [
|
||||
await self._build_chat_completion_tool_result(
|
||||
container=container,
|
||||
params=params,
|
||||
tool_call=tool_call,
|
||||
container_id=container_id,
|
||||
)
|
||||
for tool_call in tool_calls
|
||||
]
|
||||
except Exception:
|
||||
await self._delete_container_for_cache_key(sandbox_key)
|
||||
raise
|
||||
tool_messages = [result[0] for result in tool_results]
|
||||
code_interpreter_calls = [result[1] for result in tool_results]
|
||||
|
||||
request_patch = AgenticLoopRequestPatch(
|
||||
model=model,
|
||||
messages=list(messages)
|
||||
+ [self._build_chat_completion_assistant_message(tool_calls)]
|
||||
+ tool_messages,
|
||||
tools=self._get_followup_tools(
|
||||
tools=optional_params.get("tools"),
|
||||
call_type=CallTypes.completion,
|
||||
),
|
||||
optional_params=self._get_followup_optional_params(optional_params),
|
||||
kwargs=self._filter_agentic_loop_kwargs(kwargs),
|
||||
)
|
||||
|
||||
return AgenticLoopPlan(
|
||||
run_agentic_loop=True,
|
||||
request_patch=request_patch,
|
||||
metadata={
|
||||
"tool_type": "code_interpreter",
|
||||
"sandbox_key": sandbox_key or "",
|
||||
"code_interpreter_calls": code_interpreter_calls,
|
||||
"response_format": "openai",
|
||||
},
|
||||
)
|
||||
|
||||
async def _build_chat_completion_tool_result(
|
||||
self,
|
||||
container: object,
|
||||
params: dict[str, Any] | None,
|
||||
tool_call: CodeExecutionToolCall,
|
||||
container_id: str | None,
|
||||
) -> tuple[ChatCompletionToolMessage, CodeInterpreterCall]:
|
||||
arguments = tool_call.get("arguments", "")
|
||||
code = self._parse_code(arguments)
|
||||
stdout = await self._run_tool_call(
|
||||
container=container, params=params, arguments=arguments
|
||||
)
|
||||
tool_call_id = (
|
||||
tool_call.get("id") or tool_call.get("call_id") or uuid.uuid4().hex
|
||||
)
|
||||
return (
|
||||
{
|
||||
"role": "tool",
|
||||
"tool_call_id": tool_call_id,
|
||||
"content": stdout,
|
||||
},
|
||||
{
|
||||
"id": f"ci_{uuid.uuid4().hex}",
|
||||
"type": "code_interpreter_call",
|
||||
"status": "completed",
|
||||
"code": code,
|
||||
"container_id": container_id,
|
||||
"outputs": [{"type": "logs", "logs": stdout}] if stdout else [],
|
||||
},
|
||||
)
|
||||
|
||||
async def async_agentic_loop_cleanup_hook(
|
||||
self, plan: AgenticLoopPlan, kwargs: dict
|
||||
) -> None:
|
||||
metadata = plan.metadata or {} if plan else {}
|
||||
await self._delete_container_for_cache_key(metadata.get("sandbox_key"))
|
||||
|
||||
@staticmethod
|
||||
def _filter_agentic_loop_kwargs(kwargs: dict[str, object]) -> dict[str, object]:
|
||||
return {
|
||||
k: v
|
||||
for k, v in kwargs.items()
|
||||
if k not in {"litellm_logging_obj", "acompletion"}
|
||||
and not is_interception_internal_key(
|
||||
k, prefixes=NON_CODE_INTERPRETER_INTERCEPTION_INTERNAL_PREFIXES
|
||||
)
|
||||
}
|
||||
|
||||
def _get_followup_tools(
|
||||
self, tools: object, call_type: CallTypes | None
|
||||
) -> list[dict[str, Any]] | None:
|
||||
if not isinstance(tools, list):
|
||||
return None
|
||||
return [
|
||||
(
|
||||
self._get_function_tool(call_type=call_type)
|
||||
if isinstance(tool, dict) and tool.get("type") == "code_interpreter"
|
||||
else tool
|
||||
)
|
||||
for tool in tools
|
||||
]
|
||||
|
||||
def _get_followup_optional_params(
|
||||
self, optional_params: dict[str, object]
|
||||
) -> dict[str, object]:
|
||||
drop_tool_choice = self._tool_choice_targets_code_interpreter(
|
||||
optional_params.get("tool_choice")
|
||||
)
|
||||
return {
|
||||
k: v
|
||||
for k, v in optional_params.items()
|
||||
if k != "tools" and not (k == "tool_choice" and drop_tool_choice)
|
||||
}
|
||||
|
||||
async def async_post_agentic_loop_response_hook(
|
||||
self, response: Any, plan: AgenticLoopPlan, kwargs: dict
|
||||
) -> Any:
|
||||
metadata = plan.metadata or {} if plan else {}
|
||||
await self._delete_container_for_cache_key(metadata.get("sandbox_key"))
|
||||
|
||||
calls = metadata.get("code_interpreter_calls")
|
||||
if not calls:
|
||||
return response
|
||||
|
||||
is_dict = isinstance(response, dict)
|
||||
output = (
|
||||
response.get("output") if is_dict else getattr(response, "output", None)
|
||||
)
|
||||
if not isinstance(output, list):
|
||||
return response
|
||||
|
||||
def _item_type(item: Any) -> Any:
|
||||
return (
|
||||
item.get("type")
|
||||
if isinstance(item, dict)
|
||||
else getattr(item, "type", None)
|
||||
)
|
||||
|
||||
insert_at = next(
|
||||
(i for i, item in enumerate(output) if _item_type(item) == "message"),
|
||||
len(output),
|
||||
)
|
||||
new_output = output[:insert_at] + list(calls) + output[insert_at:]
|
||||
if is_dict:
|
||||
response["output"] = new_output
|
||||
else:
|
||||
response.output = new_output
|
||||
return response
|
||||
|
||||
@staticmethod
|
||||
def _parse_code(arguments: str) -> str:
|
||||
try:
|
||||
return json.loads(arguments).get("code", "") if arguments else ""
|
||||
except (json.JSONDecodeError, TypeError, AttributeError):
|
||||
return ""
|
||||
|
||||
async def _run_tool_call(
|
||||
self, container: Any, params: dict[str, Any] | None, arguments: str
|
||||
) -> str:
|
||||
try:
|
||||
code = json.loads(arguments).get("code", "") if arguments else ""
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
return "[invalid tool arguments: could not parse code]"
|
||||
|
||||
result = await self._run_code(container=container, params=params, code=code)
|
||||
if getattr(result, "error", None):
|
||||
error = result.error
|
||||
message = (
|
||||
error.get("value") or error.get("name")
|
||||
if isinstance(error, dict)
|
||||
else str(error)
|
||||
)
|
||||
return f"[execution error] {message}"
|
||||
return getattr(result, "stdout", "") or ""
|
||||
|
||||
async def _get_or_create_container(
|
||||
self, cache_key: str | None
|
||||
) -> tuple[Any, dict[str, Any] | None]:
|
||||
if cache_key:
|
||||
cached = self._container_cache.get(cache_key)
|
||||
if cached is not None:
|
||||
return cached[0], cached[1]
|
||||
|
||||
container, params = await self._create_container()
|
||||
if cache_key:
|
||||
self._container_cache[cache_key] = (container, params, time.time())
|
||||
return container, params
|
||||
|
||||
async def _create_container(self) -> tuple[Any, dict[str, Any] | None]:
|
||||
if self.sandbox_config is not None:
|
||||
return await self.sandbox_config.acreate_sandbox(), None
|
||||
|
||||
params = _resolve_sandbox_tool(self.sandbox_tool_name)
|
||||
if params is None:
|
||||
raise ValueError(
|
||||
"CodeInterpreterInterception: no sandbox available. Provide a "
|
||||
"sandbox_config or configure a sandbox tool resolvable via "
|
||||
"sandbox_tool_name."
|
||||
)
|
||||
container = await litellm.acreate_sandbox(
|
||||
provider=params["sandbox_provider"],
|
||||
api_key=params.get("api_key"),
|
||||
api_base=params.get("api_base"),
|
||||
)
|
||||
return container, params
|
||||
|
||||
async def _run_code(
|
||||
self, container: Any, params: dict[str, Any] | None, code: str
|
||||
) -> Any:
|
||||
if self.sandbox_config is not None:
|
||||
return await self.sandbox_config.arun_code(container=container, code=code)
|
||||
if params is None:
|
||||
raise ValueError(
|
||||
"CodeInterpreterInterception: no sandbox available to run code."
|
||||
)
|
||||
return await litellm.arun_code(
|
||||
provider=params["sandbox_provider"],
|
||||
container=container,
|
||||
code=code,
|
||||
api_key=params.get("api_key"),
|
||||
)
|
||||
|
||||
async def _delete_container(
|
||||
self, container: Any, params: dict[str, Any] | None
|
||||
) -> None:
|
||||
try:
|
||||
if self.sandbox_config is not None:
|
||||
await self.sandbox_config.adelete_sandbox(container=container)
|
||||
return
|
||||
if params is None:
|
||||
return
|
||||
await litellm.adelete_sandbox(
|
||||
provider=params["sandbox_provider"],
|
||||
container=container,
|
||||
api_key=params.get("api_key"),
|
||||
api_base=params.get("api_base"),
|
||||
)
|
||||
except Exception:
|
||||
verbose_logger.exception(
|
||||
"CodeInterpreterInterception: failed to delete sandbox container"
|
||||
)
|
||||
|
||||
async def _delete_container_for_cache_key(self, cache_key: str | None) -> None:
|
||||
if not cache_key:
|
||||
return
|
||||
cached = self._container_cache.pop(cache_key, None)
|
||||
if cached is None:
|
||||
return
|
||||
await self._delete_container(container=cached[0], params=cached[1])
|
||||
|
||||
def _normalize_messages(self, messages: Any) -> list[dict[str, Any]]:
|
||||
if isinstance(messages, str):
|
||||
return [{"role": "user", "content": messages}]
|
||||
if isinstance(messages, list):
|
||||
return list(messages)
|
||||
return []
|
||||
|
||||
def _extract_code_execution_tool_calls(
|
||||
self, response: object
|
||||
) -> list[CodeExecutionToolCall]:
|
||||
if isinstance(response, dict):
|
||||
output = response.get("output", [])
|
||||
else:
|
||||
output = getattr(response, "output", []) or []
|
||||
if not isinstance(output, list):
|
||||
return []
|
||||
|
||||
return [
|
||||
{
|
||||
"call_id": (
|
||||
item.get("call_id")
|
||||
if isinstance(item, dict)
|
||||
else getattr(item, "call_id", None)
|
||||
),
|
||||
"name": LITELLM_CODE_EXECUTION_TOOL_NAME,
|
||||
"arguments": (
|
||||
item.get("arguments")
|
||||
if isinstance(item, dict)
|
||||
else getattr(item, "arguments", "")
|
||||
),
|
||||
}
|
||||
for item in output
|
||||
if self._is_code_execution_call(item)
|
||||
]
|
||||
|
||||
def _extract_chat_completion_code_execution_tool_calls(
|
||||
self, response: ModelResponse | dict[str, Any]
|
||||
) -> list[CodeExecutionToolCall]:
|
||||
model_response = self._to_model_response(response)
|
||||
if model_response is None:
|
||||
return []
|
||||
choices = model_response.choices or []
|
||||
if not choices:
|
||||
return []
|
||||
message = choices[0].message
|
||||
tool_calls = message.tool_calls or []
|
||||
|
||||
return [
|
||||
normalized
|
||||
for tool_call in tool_calls
|
||||
if (normalized := self._normalize_chat_completion_tool_call(tool_call))
|
||||
is not None
|
||||
]
|
||||
|
||||
@staticmethod
|
||||
def _normalize_chat_completion_tool_call(
|
||||
tool_call: ChatCompletionMessageToolCall,
|
||||
) -> CodeExecutionToolCall | None:
|
||||
if (
|
||||
tool_call.type != "function"
|
||||
or tool_call.function.name != LITELLM_CODE_EXECUTION_TOOL_NAME
|
||||
):
|
||||
return None
|
||||
|
||||
arguments = tool_call.function.arguments
|
||||
if isinstance(arguments, dict):
|
||||
arguments = json.dumps(arguments)
|
||||
elif not isinstance(arguments, str):
|
||||
arguments = "" if arguments is None else str(arguments)
|
||||
|
||||
return {
|
||||
"id": tool_call.id,
|
||||
"call_id": tool_call.id,
|
||||
"type": "function",
|
||||
"name": LITELLM_CODE_EXECUTION_TOOL_NAME,
|
||||
"arguments": arguments,
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def _build_chat_completion_assistant_message(
|
||||
tool_calls: list[CodeExecutionToolCall],
|
||||
) -> ChatCompletionAssistantMessage:
|
||||
return {
|
||||
"role": "assistant",
|
||||
"tool_calls": [
|
||||
cast(
|
||||
ChatCompletionAssistantToolCall,
|
||||
{
|
||||
"id": tool_call.get("id"),
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": LITELLM_CODE_EXECUTION_TOOL_NAME,
|
||||
"arguments": tool_call.get("arguments", ""),
|
||||
},
|
||||
},
|
||||
)
|
||||
for tool_call in tool_calls
|
||||
],
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def _to_model_response(
|
||||
response: ModelResponse | dict[str, Any],
|
||||
) -> ModelResponse | None:
|
||||
if isinstance(response, ModelResponse):
|
||||
return response
|
||||
try:
|
||||
return ModelResponse(**response)
|
||||
except (TypeError, ValidationError):
|
||||
return None
|
||||
|
||||
def _is_code_execution_call(self, item: Any) -> bool:
|
||||
if isinstance(item, dict):
|
||||
return (
|
||||
item.get("type") == "function_call"
|
||||
and item.get("name") == LITELLM_CODE_EXECUTION_TOOL_NAME
|
||||
)
|
||||
return (
|
||||
getattr(item, "type", None) == "function_call"
|
||||
and getattr(item, "name", None) == LITELLM_CODE_EXECUTION_TOOL_NAME
|
||||
)
|
||||
|
||||
async def _prune_expired_cache(self) -> None:
|
||||
now = time.time()
|
||||
expired = [
|
||||
(cache_key, container, params)
|
||||
for cache_key, (
|
||||
container,
|
||||
params,
|
||||
created_at,
|
||||
) in self._container_cache.items()
|
||||
if now - created_at > _CACHE_TTL_SECONDS
|
||||
]
|
||||
for cache_key, container, params in expired:
|
||||
self._container_cache.pop(cache_key, None)
|
||||
await self._delete_container(container=container, params=params)
|
||||
|
|
@ -718,6 +718,24 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
|
|||
"""
|
||||
return response
|
||||
|
||||
async def async_agentic_loop_cleanup_hook(
|
||||
self,
|
||||
plan: AgenticLoopPlan,
|
||||
kwargs: dict,
|
||||
) -> None:
|
||||
"""
|
||||
Release resources held for an agentic-loop iteration.
|
||||
|
||||
Runs in a ``finally`` around the follow-up provider call, so it fires
|
||||
whether the rerun returns normally, hits a loop safety abort, or raises
|
||||
an upstream error. Implementations must be idempotent because the
|
||||
post-response hook may already have released the same resource on the
|
||||
success path. Use ``plan.metadata`` to locate what to clean up.
|
||||
|
||||
Default does nothing.
|
||||
"""
|
||||
return None
|
||||
|
||||
async def async_should_run_chat_completion_agentic_loop(
|
||||
self,
|
||||
response: Any,
|
||||
|
|
|
|||
|
|
@ -3,7 +3,7 @@
|
|||
from collections import OrderedDict
|
||||
from contextlib import contextmanager
|
||||
from datetime import datetime
|
||||
from typing import TYPE_CHECKING, Any, Iterator, Mapping, cast
|
||||
from typing import TYPE_CHECKING, Any, Callable, Iterator, Mapping, Sequence, cast
|
||||
|
||||
from opentelemetry.context import attach, get_current
|
||||
from opentelemetry.sdk.trace import TracerProvider
|
||||
|
|
@ -546,6 +546,58 @@ class OpenTelemetryV2(CustomLogger):
|
|||
return span
|
||||
|
||||
|
||||
def select_global_otel_v2_logger(
|
||||
in_memory_loggers: Sequence[object],
|
||||
registered: "OpenTelemetryV2 | None" = None,
|
||||
) -> "OpenTelemetryV2":
|
||||
"""The single ``OpenTelemetryV2`` whose provider should become the OTel global.
|
||||
|
||||
The callback factory designates one logger as canonical the moment it builds
|
||||
the first one (``_init_otel_logger_on_litellm_proxy`` sets
|
||||
``proxy_server.open_telemetry_logger``), and every other v2 entry point —
|
||||
guardrail, identity seeding, phase spans — already routes through that same
|
||||
``registered`` owner. Reuse it here too so the global provider has one source
|
||||
of truth instead of a second, independently-derived guess; this is the logger
|
||||
a preset (arize, langfuse, …) folds the ``OTEL_*`` base exporter and its own
|
||||
exporter into, so the FastAPI server span and the gen-ai spans share one
|
||||
provider and one trace.
|
||||
|
||||
Fall back to ``in_memory_loggers`` for the SDK path, where no proxy global is
|
||||
set (selecting from there, not ``service_callback``, which a preset logger does
|
||||
not always reach), and build a generic logger from ``OTEL_*`` only when none was
|
||||
configured at all. Each fallback still avoids the second generic logger that
|
||||
orphaned the gen-ai spans onto a different backend than the server span.
|
||||
"""
|
||||
if registered is not None:
|
||||
return registered
|
||||
existing = next(
|
||||
(cb for cb in in_memory_loggers if isinstance(cb, OpenTelemetryV2)), None
|
||||
)
|
||||
return existing if existing is not None else OpenTelemetryV2()
|
||||
|
||||
|
||||
def publish_global_otel_v2_provider(
|
||||
in_memory_loggers: Sequence[object],
|
||||
set_global_provider: Callable[[TracerProvider], None],
|
||||
registered: "OpenTelemetryV2 | None" = None,
|
||||
) -> "OpenTelemetryV2":
|
||||
"""Select the single v2 logger and publish its provider as the OTel global.
|
||||
|
||||
The proxy calls this once at startup, after callbacks are initialized, so the
|
||||
preset logger already exists; it passes ``registered`` (the canonical owner the
|
||||
factory designated as ``proxy_server.open_telemetry_logger``) so the global
|
||||
provider reuses the same logger the rest of the v2 code emits through (see
|
||||
:func:`select_global_otel_v2_logger`). Both ``registered`` and
|
||||
``set_global_provider`` (the proxy passes
|
||||
``opentelemetry.trace.set_tracer_provider``) are injected so the publish step is
|
||||
unit-testable without reading or mutating real global OTel state. Returns the
|
||||
logger whose provider was published.
|
||||
"""
|
||||
logger = select_global_otel_v2_logger(in_memory_loggers, registered=registered)
|
||||
set_global_provider(logger._tracer_provider)
|
||||
return logger
|
||||
|
||||
|
||||
def _registered_v2_logger() -> "OpenTelemetryV2 | None":
|
||||
try:
|
||||
from litellm.proxy import proxy_server
|
||||
|
|
|
|||
|
|
@ -1,5 +1,7 @@
|
|||
"""Typed configuration for the OpenTelemetry instrumentation."""
|
||||
|
||||
from enum import Enum
|
||||
from functools import lru_cache
|
||||
from typing import Any, List
|
||||
|
||||
from pydantic import AliasChoices, BaseModel, Field, field_validator, model_validator
|
||||
|
|
@ -23,13 +25,35 @@ class CaptureMessageContent(str):
|
|||
SPAN_AND_EVENT = "span_and_event"
|
||||
|
||||
|
||||
class ExporterOwner(str, Enum):
|
||||
"""The preset that contributed an exporter. Values match the callback names
|
||||
in ``presets.PRESET_BY_CALLBACK`` so per-request dynamic-credential routing
|
||||
can match an exporter's owner against the credential source's callback name.
|
||||
A ``str`` enum so the value compares equal to the bare callback-name string."""
|
||||
|
||||
# Arize AX (the hosted platform) and Arize Phoenix (the open-source / Phoenix
|
||||
# Cloud tracer) are distinct backends with separate config and auth, so they
|
||||
# are separate owners. The member value stays the public callback name.
|
||||
ARIZE_AX = "arize"
|
||||
ARIZE_PHOENIX = "arize_phoenix"
|
||||
LANGFUSE_OTEL = "langfuse_otel"
|
||||
WEAVE_OTEL = "weave_otel"
|
||||
LEVO = "levo"
|
||||
AGENTOPS = "agentops"
|
||||
|
||||
|
||||
class _OTelV2Flag(BaseSettings):
|
||||
model_config = SettingsConfigDict(extra="ignore")
|
||||
|
||||
enabled: bool = Field(default=False, validation_alias=AliasChoices(OTEL_V2_ENV))
|
||||
|
||||
|
||||
@lru_cache(maxsize=1)
|
||||
def is_otel_v2_enabled() -> bool:
|
||||
# Resolved once at startup and cached: constructing the pydantic-settings
|
||||
# model re-scans the environment and cost ~28us, which on the proxy hot path
|
||||
# (auth, logging-callback setup) compounded into a measurable throughput
|
||||
# regression. Tests that toggle the env must call ``is_otel_v2_enabled.cache_clear()``.
|
||||
return _OTelV2Flag().enabled
|
||||
|
||||
|
||||
|
|
@ -49,6 +73,15 @@ class ExporterSpec(BaseModel):
|
|||
)
|
||||
endpoint: str | None = None
|
||||
headers: str | None = None
|
||||
owner: ExporterOwner | None = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"The preset that contributed this exporter. Per-request dynamic OTLP "
|
||||
"credentials are applied only to the exporter whose owner matches the "
|
||||
"credential source, so one tenant's vendor key never lands on a "
|
||||
"different backend's exporter."
|
||||
),
|
||||
)
|
||||
options: dict[str, str] | None = Field(
|
||||
default=None,
|
||||
description=(
|
||||
|
|
|
|||
|
|
@ -88,13 +88,23 @@ class TenantTracerCache:
|
|||
return get_tracer(provider, self._tracer_name)
|
||||
|
||||
def _config_with_headers(self, headers: Mapping[str, str]) -> OpenTelemetryV2Config:
|
||||
"""Clone the config, replacing OTLP exporter headers with ``headers``."""
|
||||
"""Clone the config, stamping ``headers`` onto the credential's own exporter.
|
||||
|
||||
``headers`` are the per-request credentials of ``self._callback_name`` (the
|
||||
integration that built this cache), so they apply only to the exporter that
|
||||
integration contributed (``spec.owner``). A request that carries one
|
||||
tenant's Arize key must never rewrite the headers of a co-configured
|
||||
Langfuse or self-hosted collector exporter, which would leak that key to a
|
||||
different backend.
|
||||
"""
|
||||
header_str = ",".join(f"{key}={value}" for key, value in headers.items())
|
||||
header_update: dict[str, str] = {"headers": header_str}
|
||||
exporters = [
|
||||
(
|
||||
spec
|
||||
if spec.kind.lower() in _NON_OTLP_KINDS
|
||||
else spec.model_copy(update={"headers": header_str})
|
||||
spec.model_copy(update=header_update)
|
||||
if spec.owner == self._callback_name
|
||||
and spec.kind.lower() not in _NON_OTLP_KINDS
|
||||
else spec
|
||||
)
|
||||
for spec in self._config.exporters
|
||||
]
|
||||
|
|
|
|||
|
|
@ -16,7 +16,11 @@ from pydantic import Field
|
|||
from pydantic_settings import BaseSettings, SettingsConfigDict
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.integrations.otel.model.config import ExporterSpec, OpenTelemetryV2Config
|
||||
from litellm.integrations.otel.model.config import (
|
||||
ExporterOwner,
|
||||
ExporterSpec,
|
||||
OpenTelemetryV2Config,
|
||||
)
|
||||
from litellm.integrations.otel.plumbing.providers import register_exporter_factory
|
||||
|
||||
_AGENTOPS_ENDPOINT = "https://otlp.agentops.cloud/v1/traces"
|
||||
|
|
@ -59,6 +63,7 @@ def agentops_preset(
|
|||
options=(
|
||||
{"api_key": settings.api_key} if settings.api_key else None
|
||||
),
|
||||
owner=ExporterOwner.AGENTOPS,
|
||||
),
|
||||
],
|
||||
"resource_attributes": {
|
||||
|
|
|
|||
|
|
@ -4,7 +4,11 @@ from pydantic import Field
|
|||
from pydantic_settings import BaseSettings, SettingsConfigDict
|
||||
|
||||
from litellm.integrations.arize.arize import ArizeLogger as _V1ArizeLogger
|
||||
from litellm.integrations.otel.model.config import ExporterSpec, OpenTelemetryV2Config
|
||||
from litellm.integrations.otel.model.config import (
|
||||
ExporterOwner,
|
||||
ExporterSpec,
|
||||
OpenTelemetryV2Config,
|
||||
)
|
||||
from litellm.integrations.otel.presets.utils import ensure_mappers
|
||||
from litellm.types.utils import StandardCallbackDynamicParams
|
||||
|
||||
|
|
@ -34,6 +38,7 @@ def arize_preset(
|
|||
kind=arize_cfg.protocol or "otlp_grpc",
|
||||
endpoint=arize_cfg.endpoint or "https://otlp.arize.com/v1",
|
||||
headers=headers,
|
||||
owner=ExporterOwner.ARIZE_AX,
|
||||
),
|
||||
],
|
||||
"mapper_names": ensure_mappers(base.mapper_names, "openinference"),
|
||||
|
|
|
|||
|
|
@ -3,7 +3,11 @@
|
|||
from litellm.integrations.langfuse.langfuse_otel import (
|
||||
LangfuseOtelLogger as _V1Langfuse,
|
||||
)
|
||||
from litellm.integrations.otel.model.config import ExporterSpec, OpenTelemetryV2Config
|
||||
from litellm.integrations.otel.model.config import (
|
||||
ExporterOwner,
|
||||
ExporterSpec,
|
||||
OpenTelemetryV2Config,
|
||||
)
|
||||
from litellm.integrations.otel.presets.utils import ensure_mappers
|
||||
from litellm.types.utils import StandardCallbackDynamicParams
|
||||
|
||||
|
|
@ -23,6 +27,7 @@ def langfuse_preset(
|
|||
kind=kind,
|
||||
endpoint=cfg.endpoint,
|
||||
headers=cfg.headers,
|
||||
owner=ExporterOwner.LANGFUSE_OTEL,
|
||||
),
|
||||
],
|
||||
"mapper_names": ensure_mappers(base.mapper_names, "langfuse"),
|
||||
|
|
|
|||
|
|
@ -1,7 +1,11 @@
|
|||
"""Levo preset — OTLP/HTTP to a Levo collector with org+workspace headers."""
|
||||
|
||||
from litellm.integrations.levo.levo import LevoLogger as _V1Levo
|
||||
from litellm.integrations.otel.model.config import ExporterSpec, OpenTelemetryV2Config
|
||||
from litellm.integrations.otel.model.config import (
|
||||
ExporterOwner,
|
||||
ExporterSpec,
|
||||
OpenTelemetryV2Config,
|
||||
)
|
||||
|
||||
|
||||
def levo_preset(
|
||||
|
|
@ -18,6 +22,7 @@ def levo_preset(
|
|||
kind="otlp_http",
|
||||
endpoint=cfg.endpoint,
|
||||
headers=cfg.otlp_auth_headers,
|
||||
owner=ExporterOwner.LEVO,
|
||||
),
|
||||
],
|
||||
}
|
||||
|
|
|
|||
|
|
@ -6,7 +6,11 @@ from pydantic_settings import BaseSettings, SettingsConfigDict
|
|||
from litellm.integrations.arize.arize_phoenix import (
|
||||
ArizePhoenixLogger as _V1Phoenix,
|
||||
)
|
||||
from litellm.integrations.otel.model.config import ExporterSpec, OpenTelemetryV2Config
|
||||
from litellm.integrations.otel.model.config import (
|
||||
ExporterOwner,
|
||||
ExporterSpec,
|
||||
OpenTelemetryV2Config,
|
||||
)
|
||||
from litellm.integrations.otel.presets.utils import ensure_mappers
|
||||
|
||||
|
||||
|
|
@ -37,6 +41,7 @@ def phoenix_preset(
|
|||
kind=cfg.protocol if hasattr(cfg, "protocol") else "otlp_http",
|
||||
endpoint=cfg.endpoint,
|
||||
headers=headers,
|
||||
owner=ExporterOwner.ARIZE_PHOENIX,
|
||||
),
|
||||
],
|
||||
"mapper_names": ensure_mappers(base.mapper_names, "openinference"),
|
||||
|
|
|
|||
|
|
@ -1,6 +1,10 @@
|
|||
"""Weave (W&B) preset."""
|
||||
|
||||
from litellm.integrations.otel.model.config import ExporterSpec, OpenTelemetryV2Config
|
||||
from litellm.integrations.otel.model.config import (
|
||||
ExporterOwner,
|
||||
ExporterSpec,
|
||||
OpenTelemetryV2Config,
|
||||
)
|
||||
from litellm.integrations.otel.presets.utils import ensure_mappers
|
||||
from litellm.integrations.weave.weave_otel import (
|
||||
_get_weave_authorization_header,
|
||||
|
|
@ -23,6 +27,7 @@ def weave_preset(
|
|||
kind=weave_cfg.protocol or "otlp_http",
|
||||
endpoint=weave_cfg.endpoint,
|
||||
headers=weave_cfg.otlp_auth_headers,
|
||||
owner=ExporterOwner.WEAVE_OTEL,
|
||||
),
|
||||
],
|
||||
# Weave consumes OpenInference + a small Weave-specific overlay.
|
||||
|
|
|
|||
|
|
@ -300,9 +300,6 @@ class LiteLLMResponsesInteractionsConfig:
|
|||
"total_output_tokens": getattr(usage, "output_tokens", 0),
|
||||
}
|
||||
|
||||
# Add role
|
||||
interactions_response_dict["role"] = "model"
|
||||
|
||||
# Add updated (same as created for now)
|
||||
interactions_response_dict["updated"] = created
|
||||
|
||||
|
|
|
|||
332
litellm/litellm_core_utils/chat_completion_agentic_loop.py
Normal file
332
litellm/litellm_core_utils/chat_completion_agentic_loop.py
Normal file
|
|
@ -0,0 +1,332 @@
|
|||
# this is a patch to allow for agentic loops covering llm_http_handler.py and openai sdk based calling flows for the .completion() api
|
||||
|
||||
import json
|
||||
from typing import cast
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.types.integrations.custom_logger import (
|
||||
CHAT_COMPLETION_AGENTIC_SURFACE,
|
||||
NON_CODE_INTERPRETER_INTERCEPTION_INTERNAL_PREFIXES,
|
||||
AgenticLoopPlan,
|
||||
AgenticLoopRequestPatch,
|
||||
is_interception_internal_key,
|
||||
)
|
||||
from litellm.types.utils import ModelResponse
|
||||
from litellm.utils import CustomStreamWrapper
|
||||
|
||||
_FOLLOWUP_INTERNAL_PARAMS = frozenset(
|
||||
(
|
||||
"acompletion",
|
||||
"litellm_logging_obj",
|
||||
"custom_llm_provider",
|
||||
"model_alias_map",
|
||||
"stream_response",
|
||||
"custom_prompt_dict",
|
||||
"_agentic_loop_api_surface",
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def _gate_overridden(callback: CustomLogger) -> bool:
|
||||
base = CustomLogger.async_should_run_agentic_loop
|
||||
func = type(callback).async_should_run_agentic_loop
|
||||
return getattr(func, "__func__", func) is not getattr(base, "__func__", base)
|
||||
|
||||
|
||||
def _build_plan_overridden(callback: CustomLogger) -> bool:
|
||||
base = CustomLogger.async_build_agentic_loop_plan
|
||||
func = type(callback).async_build_agentic_loop_plan
|
||||
return getattr(func, "__func__", func) is not getattr(base, "__func__", base)
|
||||
|
||||
|
||||
def _post_hook_overridden(callback: CustomLogger) -> bool:
|
||||
base = CustomLogger.async_post_agentic_loop_response_hook
|
||||
func = type(callback).async_post_agentic_loop_response_hook
|
||||
return getattr(func, "__func__", func) is not getattr(base, "__func__", base)
|
||||
|
||||
|
||||
def _coerce_int(value: object, default: int) -> int:
|
||||
return int(value) if isinstance(value, (int, str)) else default
|
||||
|
||||
|
||||
def _agentic_loop_settings(kwargs: dict[str, object]) -> tuple[int, int, list[str]]:
|
||||
depth = _coerce_int(kwargs.get("_agentic_loop_depth"), 0)
|
||||
max_loops = max(_coerce_int(kwargs.get("max_agentic_loops"), 3), 1)
|
||||
raw_fingerprints = kwargs.get("_agentic_loop_fingerprints")
|
||||
fingerprints = (
|
||||
[str(fp) for fp in raw_fingerprints]
|
||||
if isinstance(raw_fingerprints, list)
|
||||
else []
|
||||
)
|
||||
return depth, max_loops, fingerprints
|
||||
|
||||
|
||||
def _fingerprint_tools(tool_calls: object) -> str:
|
||||
try:
|
||||
return json.dumps(tool_calls, sort_keys=True, default=str)
|
||||
except Exception:
|
||||
return str(tool_calls)
|
||||
|
||||
|
||||
def _check_agentic_loop_safety(
|
||||
tool_calls: object,
|
||||
fingerprints: list[str],
|
||||
depth: int,
|
||||
max_loops: int,
|
||||
model: str,
|
||||
) -> str:
|
||||
fingerprint = _fingerprint_tools(tool_calls)
|
||||
if fingerprint in fingerprints:
|
||||
raise ValueError(
|
||||
"Agentic loop detected repeated tool-call fingerprint; aborting rerun"
|
||||
)
|
||||
if depth >= max_loops:
|
||||
raise ValueError(f"Exceeded max_agentic_loops={max_loops} for model={model}")
|
||||
return fingerprint
|
||||
|
||||
|
||||
def _wrap_response_as_fake_stream(response: object) -> object:
|
||||
if getattr(response, "object", None) == "chat.completion.chunk":
|
||||
return response
|
||||
if not hasattr(response, "choices"):
|
||||
return response
|
||||
from litellm.llms.base_llm.base_model_iterator import (
|
||||
convert_model_response_to_streaming,
|
||||
)
|
||||
|
||||
return convert_model_response_to_streaming(cast(ModelResponse, response))
|
||||
|
||||
|
||||
def _add_agentic_loop_metadata(kwargs_for_followup: dict[str, object]) -> None:
|
||||
metadata = kwargs_for_followup.get("litellm_metadata")
|
||||
metadata = dict(metadata) if isinstance(metadata, dict) else {}
|
||||
for key, value in kwargs_for_followup.items():
|
||||
if (
|
||||
key.startswith("_agentic_loop")
|
||||
or key == "max_agentic_loops"
|
||||
or is_interception_internal_key(key)
|
||||
):
|
||||
metadata[key] = value
|
||||
kwargs_for_followup["litellm_metadata"] = metadata
|
||||
|
||||
|
||||
def _filter_followup_kwargs(source: dict[str, object]) -> dict[str, object]:
|
||||
return {
|
||||
k: v
|
||||
for k, v in source.items()
|
||||
if not is_interception_internal_key(
|
||||
k, prefixes=NON_CODE_INTERPRETER_INTERCEPTION_INTERNAL_PREFIXES
|
||||
)
|
||||
and k not in _FOLLOWUP_INTERNAL_PARAMS
|
||||
}
|
||||
|
||||
|
||||
async def _execute_chat_completion_agentic_plan(
|
||||
*,
|
||||
plan: AgenticLoopPlan,
|
||||
callback: CustomLogger,
|
||||
model: str,
|
||||
optional_params: dict[str, object],
|
||||
kwargs: dict[str, object],
|
||||
logging_obj: object,
|
||||
custom_llm_provider: str,
|
||||
depth: int,
|
||||
max_loops: int,
|
||||
fingerprints: list[str],
|
||||
fingerprint: str,
|
||||
) -> object:
|
||||
import litellm
|
||||
|
||||
patch = plan.request_patch or AgenticLoopRequestPatch()
|
||||
if patch.messages is None:
|
||||
raise ValueError("Agentic loop plan missing patched messages")
|
||||
|
||||
full_model_name = patch.model or model
|
||||
if "/" not in full_model_name:
|
||||
full_model_name = f"{custom_llm_provider}/{full_model_name}"
|
||||
|
||||
optional_params_for_followup = {**optional_params, **patch.optional_params}
|
||||
if patch.tools is not None:
|
||||
optional_params_for_followup["tools"] = patch.tools
|
||||
if "tool_choice" not in patch.optional_params:
|
||||
optional_params_for_followup.pop("tool_choice", None)
|
||||
|
||||
kwargs_for_followup = _filter_followup_kwargs(kwargs)
|
||||
kwargs_for_followup.update(
|
||||
{
|
||||
k: v
|
||||
for k, v in _filter_followup_kwargs(patch.kwargs).items()
|
||||
if k not in optional_params_for_followup
|
||||
}
|
||||
)
|
||||
kwargs_for_followup["_agentic_loop_depth"] = depth + 1
|
||||
kwargs_for_followup["max_agentic_loops"] = max_loops
|
||||
kwargs_for_followup["_agentic_loop_fingerprints"] = fingerprints + [fingerprint]
|
||||
_add_agentic_loop_metadata(kwargs_for_followup)
|
||||
|
||||
try:
|
||||
response_followup = await litellm.acompletion(
|
||||
model=full_model_name,
|
||||
messages=patch.messages,
|
||||
**optional_params_for_followup,
|
||||
**kwargs_for_followup,
|
||||
)
|
||||
if _post_hook_overridden(callback):
|
||||
try:
|
||||
response_followup = (
|
||||
await callback.async_post_agentic_loop_response_hook(
|
||||
response=response_followup, plan=plan, kwargs=kwargs
|
||||
)
|
||||
)
|
||||
except Exception as e:
|
||||
_call_id = getattr(logging_obj, "litellm_call_id", "unknown")
|
||||
verbose_logger.exception(
|
||||
"LiteLLM.AgenticHookError: Exception in "
|
||||
"async_post_agentic_loop_response_hook [call_id=%s model=%s]: %s",
|
||||
_call_id,
|
||||
model,
|
||||
str(e),
|
||||
)
|
||||
if kwargs.get("_code_interpreter_interception_converted_stream") and not depth:
|
||||
return _wrap_response_as_fake_stream(response_followup)
|
||||
return response_followup
|
||||
finally:
|
||||
try:
|
||||
await callback.async_agentic_loop_cleanup_hook(plan=plan, kwargs=kwargs)
|
||||
except Exception as e:
|
||||
_call_id = getattr(logging_obj, "litellm_call_id", "unknown")
|
||||
verbose_logger.exception(
|
||||
"LiteLLM.AgenticHookError: Exception in "
|
||||
"async_agentic_loop_cleanup_hook [call_id=%s model=%s]: %s",
|
||||
_call_id,
|
||||
model,
|
||||
str(e),
|
||||
)
|
||||
|
||||
|
||||
async def maybe_run_chat_completion_agentic_loop(
|
||||
*,
|
||||
response: ModelResponse,
|
||||
model: str,
|
||||
messages: list,
|
||||
optional_params: dict,
|
||||
kwargs: dict,
|
||||
logging_obj: object,
|
||||
custom_llm_provider: str,
|
||||
stream: bool,
|
||||
) -> ModelResponse | CustomStreamWrapper | None:
|
||||
import litellm
|
||||
|
||||
callbacks = litellm.callbacks + (
|
||||
getattr(logging_obj, "dynamic_success_callbacks", None) or []
|
||||
)
|
||||
depth, max_loops, fingerprints = _agentic_loop_settings(kwargs)
|
||||
tools = optional_params.get("tools", [])
|
||||
|
||||
for callback in callbacks:
|
||||
if not isinstance(callback, CustomLogger):
|
||||
continue
|
||||
if not _gate_overridden(callback):
|
||||
continue
|
||||
|
||||
gate_kwargs = {
|
||||
**kwargs,
|
||||
"_agentic_loop_api_surface": CHAT_COMPLETION_AGENTIC_SURFACE,
|
||||
"custom_llm_provider": custom_llm_provider,
|
||||
}
|
||||
try:
|
||||
should_run, tool_calls = await callback.async_should_run_agentic_loop(
|
||||
response=response,
|
||||
model=model,
|
||||
messages=messages,
|
||||
tools=tools,
|
||||
stream=stream,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
kwargs=gate_kwargs,
|
||||
)
|
||||
except Exception as e:
|
||||
verbose_logger.exception(
|
||||
"LiteLLM.AgenticHookError: Exception in chat completion agentic gate: %s",
|
||||
str(e),
|
||||
)
|
||||
continue
|
||||
|
||||
if not should_run:
|
||||
continue
|
||||
|
||||
fingerprint = _check_agentic_loop_safety(
|
||||
tool_calls=tool_calls,
|
||||
fingerprints=fingerprints,
|
||||
depth=depth,
|
||||
max_loops=max_loops,
|
||||
model=model,
|
||||
)
|
||||
|
||||
try:
|
||||
plan_kwargs = {
|
||||
**kwargs,
|
||||
"_agentic_loop_api_surface": CHAT_COMPLETION_AGENTIC_SURFACE,
|
||||
"custom_llm_provider": custom_llm_provider,
|
||||
}
|
||||
if not _build_plan_overridden(callback):
|
||||
return await callback.async_run_agentic_loop(
|
||||
tools=tool_calls,
|
||||
model=model,
|
||||
messages=messages,
|
||||
response=response,
|
||||
anthropic_messages_provider_config=None,
|
||||
anthropic_messages_optional_request_params=optional_params,
|
||||
logging_obj=logging_obj,
|
||||
stream=stream,
|
||||
kwargs=plan_kwargs,
|
||||
)
|
||||
|
||||
plan = await callback.async_build_agentic_loop_plan(
|
||||
tools=tool_calls,
|
||||
model=model,
|
||||
messages=messages,
|
||||
response=response,
|
||||
anthropic_messages_provider_config=None,
|
||||
anthropic_messages_optional_request_params=optional_params,
|
||||
logging_obj=logging_obj,
|
||||
stream=stream,
|
||||
kwargs=plan_kwargs,
|
||||
)
|
||||
|
||||
if plan.response_override is not None:
|
||||
return plan.response_override
|
||||
if plan.terminate:
|
||||
return response
|
||||
if not plan.run_agentic_loop:
|
||||
continue
|
||||
|
||||
return await _execute_chat_completion_agentic_plan(
|
||||
plan=plan,
|
||||
callback=callback,
|
||||
model=model,
|
||||
optional_params=optional_params,
|
||||
kwargs=kwargs,
|
||||
logging_obj=logging_obj,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
depth=depth,
|
||||
max_loops=max_loops,
|
||||
fingerprints=fingerprints,
|
||||
fingerprint=fingerprint,
|
||||
)
|
||||
except Exception as e:
|
||||
verbose_logger.exception(
|
||||
"LiteLLM.AgenticHookError: Exception in chat completion agentic hooks: %s",
|
||||
str(e),
|
||||
)
|
||||
|
||||
if (
|
||||
kwargs.get("_code_interpreter_interception_converted_stream")
|
||||
and not depth
|
||||
and hasattr(response, "choices")
|
||||
):
|
||||
return cast(
|
||||
"ModelResponse | CustomStreamWrapper",
|
||||
_wrap_response_as_fake_stream(response),
|
||||
)
|
||||
return None
|
||||
|
|
@ -15,8 +15,23 @@ BEDROCK_MANAGED_S3_PREFIXES = (
|
|||
BEDROCK_MANAGED_S3_UPLOAD_PREFIX,
|
||||
BEDROCK_MANAGED_S3_OUTPUT_PREFIX,
|
||||
)
|
||||
MANAGED_CLOUD_STORAGE_SCHEMES = ("s3://", "gs://")
|
||||
_MAPPING_PROXY_TYPE: type = type(MappingProxyType({}))
|
||||
|
||||
|
||||
def is_managed_cloud_storage_uri(file_id: str) -> bool:
|
||||
"""
|
||||
True if file_id is a raw cloud-storage object URI (e.g. ``s3://bucket/key``).
|
||||
|
||||
These are internal provider artifacts. On the multi-tenant proxy they must be
|
||||
retrieved through their managed unified file id so owner/team access is enforced;
|
||||
a raw URI supplied by a caller bypasses that check.
|
||||
"""
|
||||
return isinstance(file_id, str) and file_id.startswith(
|
||||
MANAGED_CLOUD_STORAGE_SCHEMES
|
||||
)
|
||||
|
||||
|
||||
_SAFE_OBJECT_COMPONENT_PATTERN = re.compile(r"[^A-Za-z0-9._-]+")
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -6,10 +6,7 @@ from typing import Callable, Optional, Union
|
|||
|
||||
import httpx
|
||||
|
||||
from litellm.constants import (
|
||||
COMPLETION_HTTP_FALLBACK_SECONDS,
|
||||
DEFAULT_REQUEST_TIMEOUT_SECONDS,
|
||||
)
|
||||
from litellm.constants import COMPLETION_HTTP_FALLBACK_SECONDS
|
||||
|
||||
|
||||
class CompletionTimeout:
|
||||
|
|
@ -22,17 +19,13 @@ class CompletionTimeout:
|
|||
"""
|
||||
Used when ``model_timeout`` and kwargs timeouts are all unset.
|
||||
|
||||
``global_timeout`` is :attr:`litellm.request_timeout` (numeric / string), not
|
||||
:class:`httpx.Timeout`.
|
||||
|
||||
If it equals :data:`~litellm.constants.DEFAULT_REQUEST_TIMEOUT_SECONDS` (6000),
|
||||
return :data:`~litellm.constants.COMPLETION_HTTP_FALLBACK_SECONDS`. Same if
|
||||
``None``. Otherwise return ``float(global_timeout)``.
|
||||
``global_timeout`` is the explicitly-configured ``litellm.request_timeout``
|
||||
(numeric / string) or ``None`` when it was never set. ``None`` falls back to
|
||||
:data:`~litellm.constants.COMPLETION_HTTP_FALLBACK_SECONDS`; any explicit value
|
||||
(including ``6000``) is honored.
|
||||
"""
|
||||
if global_timeout is None:
|
||||
return COMPLETION_HTTP_FALLBACK_SECONDS
|
||||
if float(global_timeout) == float(DEFAULT_REQUEST_TIMEOUT_SECONDS):
|
||||
return COMPLETION_HTTP_FALLBACK_SECONDS
|
||||
return float(global_timeout)
|
||||
|
||||
@staticmethod
|
||||
|
|
@ -50,11 +43,10 @@ class CompletionTimeout:
|
|||
1. ``model_timeout`` (call argument / merged ``litellm_params``)
|
||||
2. ``kwargs["timeout"]``
|
||||
3. ``kwargs["request_timeout"]``
|
||||
4. Fallback from ``global_timeout`` (:attr:`litellm.request_timeout`) — if it is
|
||||
the package default (6000), use 600 instead.
|
||||
4. ``global_timeout`` (the explicitly-configured ``litellm.request_timeout``),
|
||||
or 600 when nothing was configured.
|
||||
|
||||
Coerce :class:`httpx.Timeout` when the provider does not support it.
|
||||
Explicit ``6000`` on the model or in kwargs is kept as ``6000``.
|
||||
"""
|
||||
resolved: Union[float, str, httpx.Timeout]
|
||||
if model_timeout is not None:
|
||||
|
|
|
|||
File diff suppressed because it is too large
Load diff
|
|
@ -388,6 +388,9 @@ def get_llm_provider(
|
|||
elif endpoint == "https://api.inference.wandb.ai/v1":
|
||||
custom_llm_provider = "wandb"
|
||||
dynamic_api_key = get_secret_str("WANDB_API_KEY")
|
||||
elif endpoint == "https://pinstripes.io/v1":
|
||||
custom_llm_provider = "pinstripes"
|
||||
dynamic_api_key = get_secret_str("PINSTRIPES_API_KEY")
|
||||
|
||||
if api_base is not None and not isinstance(api_base, str):
|
||||
raise Exception(
|
||||
|
|
@ -641,7 +644,7 @@ def _get_openai_compatible_provider_info(
|
|||
api_base,
|
||||
dynamic_api_key,
|
||||
) = litellm.BedrockMantleChatConfig()._get_openai_compatible_provider_info(
|
||||
api_base, api_key, litellm_params=litellm_params
|
||||
api_base, api_key, litellm_params=litellm_params, model=model
|
||||
)
|
||||
elif custom_llm_provider == "nvidia_nim":
|
||||
# nvidia_nim is openai compatible, we just need to set this to custom_openai and have the api_base be https://api.endpoints.anyscale.com/v1
|
||||
|
|
|
|||
|
|
@ -86,9 +86,7 @@ def get_supported_openai_params(
|
|||
model=model
|
||||
)
|
||||
elif request_type == "transcription":
|
||||
return litellm.FireworksAIAudioTranscriptionConfig().get_supported_openai_params(
|
||||
model=model
|
||||
)
|
||||
return None
|
||||
else:
|
||||
return litellm.FireworksAIConfig().get_supported_openai_params(model=model)
|
||||
elif custom_llm_provider == "nvidia_nim":
|
||||
|
|
@ -191,7 +189,9 @@ def get_supported_openai_params(
|
|||
)
|
||||
elif custom_llm_provider == "sambanova":
|
||||
if request_type == "embeddings":
|
||||
litellm.SambaNovaEmbeddingConfig().get_supported_openai_params(model=model)
|
||||
return litellm.SambaNovaEmbeddingConfig().get_supported_openai_params(
|
||||
model=model
|
||||
)
|
||||
else:
|
||||
return litellm.SambanovaConfig().get_supported_openai_params(model=model)
|
||||
elif custom_llm_provider == "nebius":
|
||||
|
|
|
|||
|
|
@ -44,8 +44,8 @@ class HealthCheckHelpers:
|
|||
model_params["litellm_logging_obj"] = litellm_logging_obj
|
||||
model_params["fallbacks"] = fallback_models
|
||||
model_params["max_tokens"] = model_params.get(
|
||||
"max_tokens", 10
|
||||
) # gpt-5-nano throws errors for max_tokens=1
|
||||
"max_tokens", 16
|
||||
) # GPT-5 models require max_output_tokens >= 16
|
||||
await acompletion(**model_params)
|
||||
return {}
|
||||
|
||||
|
|
|
|||
|
|
@ -2975,7 +2975,12 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
)
|
||||
self.model_call_details["end_time"] = end_time
|
||||
self.model_call_details.setdefault("original_response", None)
|
||||
self.model_call_details["response_cost"] = 0
|
||||
# A stream interrupted mid-flight still billed the provider for the
|
||||
# chunks already delivered; the router stashes that recovered usage as
|
||||
# ``combined_usage_object`` and pre-computes its cost, so preserve it
|
||||
# here instead of zeroing the spend on an otherwise-failed request.
|
||||
if self.model_call_details.get("combined_usage_object") is None:
|
||||
self.model_call_details["response_cost"] = 0
|
||||
|
||||
if hasattr(exception, "headers") and isinstance(exception.headers, dict):
|
||||
self.model_call_details.setdefault("litellm_params", {})
|
||||
|
|
@ -5416,19 +5421,20 @@ class StandardLoggingPayloadSetup:
|
|||
tb_lines[:MAXIMUM_TRACEBACK_LINES_TO_LOG]
|
||||
) # Limit to first 100 lines
|
||||
|
||||
# Prefer the `.message` attribute (set by ProxyException and every
|
||||
# litellm.exceptions.* class) over str(exc); ProxyException does not
|
||||
# call super().__init__() nor define __str__, so str() on it returns
|
||||
# an empty string, which used to silently strip the human-readable
|
||||
# message from spend_logs.metadata.error_information.
|
||||
# Use isinstance, not truthiness: an explicit empty string on
|
||||
# `.message` is a deliberate value and must not be replaced by
|
||||
# `str(exc)`.
|
||||
explicit_message = getattr(original_exception, "message", None)
|
||||
error_message = (
|
||||
explicit_message
|
||||
if isinstance(explicit_message, str) and explicit_message
|
||||
else str(original_exception)
|
||||
)
|
||||
if isinstance(explicit_message, str):
|
||||
error_message = explicit_message
|
||||
else:
|
||||
error_message = str(original_exception) if original_exception else ""
|
||||
|
||||
# Duck-typed read so bare-Exception subclasses like
|
||||
# `litellm.BudgetExceededError` can participate without joining the
|
||||
# RateLimitError hierarchy (which would break `except BudgetExceededError`).
|
||||
# Validated against the enum value sets so a third-party exception that
|
||||
# happens to declare a `.category` or `.rate_limit_type` string attribute
|
||||
# can't leak garbage into the payload or Prometheus label cardinality.
|
||||
rate_limit_category = validate_rate_limit_category(
|
||||
getattr(original_exception, "category", None)
|
||||
)
|
||||
|
|
@ -5441,7 +5447,7 @@ class StandardLoggingPayloadSetup:
|
|||
error_class=error_class,
|
||||
llm_provider=_llm_provider_in_exception,
|
||||
traceback=traceback_info,
|
||||
error_message=error_message if original_exception else "",
|
||||
error_message=error_message,
|
||||
error_rate_limit_category=rate_limit_category,
|
||||
error_rate_limit_type=rate_limit_type,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -191,6 +191,11 @@ def _get_service_tier_cost_key(base_key: str, service_tier: Optional[str]) -> st
|
|||
return base_key
|
||||
|
||||
|
||||
def _parse_above_token_threshold(key: str) -> float:
|
||||
threshold_str = key.split("_above_")[1].split("_tokens")[0]
|
||||
return float(threshold_str.replace("k", "")) * (1000 if "k" in threshold_str else 1)
|
||||
|
||||
|
||||
def _get_token_base_cost(
|
||||
model_info: ModelInfo, usage: Usage, service_tier: Optional[str] = None
|
||||
) -> Tuple[float, float, float, float, float]:
|
||||
|
|
@ -256,15 +261,13 @@ def _get_token_base_cost(
|
|||
|
||||
# Only sort the threshold keys (typically 1-2 keys instead of 66+)
|
||||
threshold: Optional[float] = None
|
||||
for key in sorted(threshold_keys, reverse=True):
|
||||
for key in sorted(threshold_keys, key=_parse_above_token_threshold, reverse=True):
|
||||
value = model_info.get(key)
|
||||
if value is not None:
|
||||
try:
|
||||
# Handle both formats: _above_128k_tokens and _above_128_tokens
|
||||
threshold_str = key.split("_above_")[1].split("_tokens")[0]
|
||||
threshold = float(threshold_str.replace("k", "")) * (
|
||||
1000 if "k" in threshold_str else 1
|
||||
)
|
||||
threshold = _parse_above_token_threshold(key)
|
||||
if usage.prompt_tokens > threshold:
|
||||
# Prefer a service_tier-specific above-threshold key when available,
|
||||
# e.g. input_cost_per_token_priority_above_200k_tokens for Gemini
|
||||
|
|
|
|||
29
litellm/litellm_core_utils/request_timeout_resolver.py
Normal file
29
litellm/litellm_core_utils/request_timeout_resolver.py
Normal file
|
|
@ -0,0 +1,29 @@
|
|||
"""Single source of truth for whether ``litellm.request_timeout`` was configured.
|
||||
|
||||
``litellm.request_timeout`` always holds a value (the package default,
|
||||
:data:`~litellm.constants.DEFAULT_REQUEST_TIMEOUT_SECONDS`), so a bare read can't
|
||||
tell "user asked for this" from "nobody set it". This resolver answers that:
|
||||
|
||||
* ``request_timeout_explicitly_set`` is the authoritative signal, set when the
|
||||
value comes from the ``REQUEST_TIMEOUT`` env var or ``litellm_settings``.
|
||||
* A runtime value that differs from the package default (e.g. ``litellm.request_timeout
|
||||
= 300`` in SDK code) is also treated as explicit, for backwards compatibility.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Optional
|
||||
|
||||
from litellm.constants import DEFAULT_REQUEST_TIMEOUT_SECONDS
|
||||
|
||||
|
||||
def get_configured_request_timeout() -> Optional[float]:
|
||||
"""Return the explicitly-configured ``litellm.request_timeout``, else ``None``."""
|
||||
import litellm
|
||||
|
||||
timeout = float(litellm.request_timeout)
|
||||
if litellm.request_timeout_explicitly_set:
|
||||
return timeout
|
||||
if timeout != float(DEFAULT_REQUEST_TIMEOUT_SECONDS):
|
||||
return timeout
|
||||
return None
|
||||
|
|
@ -12,6 +12,7 @@ class SensitiveDataMasker:
|
|||
visible_prefix: int = 4,
|
||||
visible_suffix: int = 4,
|
||||
mask_char: str = "*",
|
||||
mask_short_values: bool = True,
|
||||
):
|
||||
self.sensitive_patterns = sensitive_patterns or {
|
||||
"password",
|
||||
|
|
@ -38,12 +39,17 @@ class SensitiveDataMasker:
|
|||
self.visible_prefix = visible_prefix
|
||||
self.visible_suffix = visible_suffix
|
||||
self.mask_char = mask_char
|
||||
self.mask_short_values = mask_short_values
|
||||
|
||||
def _mask_value(self, value: str) -> str:
|
||||
if not value or len(str(value)) < (self.visible_prefix + self.visible_suffix):
|
||||
return value
|
||||
|
||||
value_str = str(value)
|
||||
if not value_str:
|
||||
return value
|
||||
if len(value_str) <= (self.visible_prefix + self.visible_suffix):
|
||||
return (
|
||||
self.mask_char * len(value_str) if self.mask_short_values else value_str
|
||||
)
|
||||
|
||||
masked_length = len(value_str) - (self.visible_prefix + self.visible_suffix)
|
||||
|
||||
# Handle the case where visible_suffix is 0 to avoid showing the entire string
|
||||
|
|
|
|||
|
|
@ -6,6 +6,7 @@ import logging
|
|||
import threading
|
||||
import time
|
||||
import traceback
|
||||
from dataclasses import dataclass
|
||||
from typing import (
|
||||
Any,
|
||||
AsyncIterator,
|
||||
|
|
@ -97,6 +98,19 @@ def print_verbose(print_statement):
|
|||
pass
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _ProviderChunkParsed:
|
||||
response_obj: dict[str, Any]
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _ProviderChunkEarlyReturn:
|
||||
value: Any
|
||||
|
||||
|
||||
_ProviderChunkResult = Union[_ProviderChunkParsed, _ProviderChunkEarlyReturn]
|
||||
|
||||
|
||||
class CustomStreamWrapper:
|
||||
def __init__(
|
||||
self,
|
||||
|
|
@ -1145,381 +1159,392 @@ class CustomStreamWrapper:
|
|||
del model_response.choices[0].delta.reasoning_content
|
||||
return
|
||||
|
||||
def _dispatch_provider_chunk(
|
||||
self,
|
||||
chunk: Any,
|
||||
model_response: ModelResponseStream,
|
||||
completion_obj: dict[str, Any],
|
||||
) -> _ProviderChunkResult:
|
||||
response_obj: dict[str, Any] = {}
|
||||
if (
|
||||
isinstance(chunk, ModelResponseStream)
|
||||
and self.custom_llm_provider is not None
|
||||
and self.custom_llm_provider in litellm._custom_providers
|
||||
):
|
||||
_has_content = bool(
|
||||
chunk.choices
|
||||
and chunk.choices[0].delta is not None
|
||||
and (
|
||||
chunk.choices[0].delta.content or chunk.choices[0].delta.tool_calls
|
||||
)
|
||||
)
|
||||
if self.received_finish_reason is not None:
|
||||
if not _has_content:
|
||||
raise StopIteration
|
||||
if chunk.choices and chunk.choices[0].finish_reason:
|
||||
self.received_finish_reason = chunk.choices[0].finish_reason
|
||||
if not _has_content:
|
||||
return _ProviderChunkEarlyReturn(None)
|
||||
# Strip finish_reason from the content chunk so it appears
|
||||
# only on the trailing empty-delta chunk (OpenAI spec).
|
||||
# finish_reason_handler() will emit the proper terminal chunk.
|
||||
chunk.choices[0].finish_reason = None # type: ignore[assignment]
|
||||
return _ProviderChunkEarlyReturn(chunk)
|
||||
|
||||
if (
|
||||
isinstance(chunk, dict)
|
||||
and generic_chunk_has_all_required_fields(
|
||||
chunk=chunk
|
||||
) # check if chunk is a generic streaming chunk
|
||||
) or (
|
||||
self.custom_llm_provider
|
||||
and self.custom_llm_provider in litellm._custom_providers
|
||||
):
|
||||
if self.received_finish_reason is not None:
|
||||
_chunk_has_content = isinstance(chunk, dict) and (
|
||||
bool(chunk.get("text", ""))
|
||||
or chunk.get("tool_use") is not None
|
||||
# Usage-only final chunks are valid and needed to surface
|
||||
# finish_reason/usage to downstream translators.
|
||||
or chunk.get("usage") is not None
|
||||
)
|
||||
if not _chunk_has_content and (
|
||||
not isinstance(chunk, dict)
|
||||
or "provider_specific_fields" not in chunk
|
||||
):
|
||||
raise StopIteration
|
||||
anthropic_response_obj: GChunk = cast(GChunk, chunk)
|
||||
completion_obj["content"] = anthropic_response_obj["text"]
|
||||
if anthropic_response_obj["is_finished"]:
|
||||
self.received_finish_reason = anthropic_response_obj["finish_reason"]
|
||||
|
||||
if anthropic_response_obj["finish_reason"]:
|
||||
self.intermittent_finish_reason = anthropic_response_obj[
|
||||
"finish_reason"
|
||||
]
|
||||
|
||||
if anthropic_response_obj["usage"] is not None:
|
||||
setattr(
|
||||
model_response,
|
||||
"usage",
|
||||
litellm.Usage(**anthropic_response_obj["usage"]),
|
||||
)
|
||||
|
||||
if (
|
||||
"tool_use" in anthropic_response_obj
|
||||
and anthropic_response_obj["tool_use"] is not None
|
||||
):
|
||||
completion_obj["tool_calls"] = [anthropic_response_obj["tool_use"]]
|
||||
|
||||
if (
|
||||
"provider_specific_fields" in anthropic_response_obj
|
||||
and anthropic_response_obj["provider_specific_fields"] is not None
|
||||
):
|
||||
for key, value in anthropic_response_obj[
|
||||
"provider_specific_fields"
|
||||
].items():
|
||||
setattr(model_response, key, value)
|
||||
|
||||
response_obj = cast(dict[str, Any], anthropic_response_obj)
|
||||
elif self.model == "replicate" or self.custom_llm_provider == "replicate":
|
||||
response_obj = self.handle_replicate_chunk(chunk)
|
||||
completion_obj["content"] = response_obj["text"]
|
||||
if response_obj["is_finished"]:
|
||||
self.received_finish_reason = response_obj["finish_reason"]
|
||||
elif self.custom_llm_provider and self.custom_llm_provider == "predibase":
|
||||
response_obj = self.handle_predibase_chunk(chunk)
|
||||
completion_obj["content"] = response_obj["text"]
|
||||
if response_obj["is_finished"]:
|
||||
self.received_finish_reason = response_obj["finish_reason"]
|
||||
elif (
|
||||
self.custom_llm_provider and self.custom_llm_provider == "baseten"
|
||||
): # baseten doesn't provide streaming
|
||||
completion_obj["content"] = self.handle_baseten_chunk(chunk)
|
||||
elif (
|
||||
self.custom_llm_provider and self.custom_llm_provider == "ai21"
|
||||
): # ai21 doesn't provide streaming
|
||||
response_obj = self.handle_ai21_chunk(chunk)
|
||||
completion_obj["content"] = response_obj["text"]
|
||||
if response_obj["is_finished"]:
|
||||
self.received_finish_reason = response_obj["finish_reason"]
|
||||
elif self.custom_llm_provider and self.custom_llm_provider == "maritalk":
|
||||
response_obj = self.handle_maritalk_chunk(chunk)
|
||||
completion_obj["content"] = response_obj["text"]
|
||||
if response_obj["is_finished"]:
|
||||
self.received_finish_reason = response_obj["finish_reason"]
|
||||
elif self.custom_llm_provider and self.custom_llm_provider == "vllm":
|
||||
completion_obj["content"] = chunk[0].outputs[0].text
|
||||
elif (
|
||||
self.custom_llm_provider and self.custom_llm_provider == "aleph_alpha"
|
||||
): # aleph alpha doesn't provide streaming
|
||||
response_obj = self.handle_aleph_alpha_chunk(chunk)
|
||||
completion_obj["content"] = response_obj["text"]
|
||||
if response_obj["is_finished"]:
|
||||
self.received_finish_reason = response_obj["finish_reason"]
|
||||
elif self.custom_llm_provider == "nlp_cloud":
|
||||
try:
|
||||
response_obj = self.handle_nlp_cloud_chunk(chunk)
|
||||
completion_obj["content"] = response_obj["text"]
|
||||
if response_obj["is_finished"]:
|
||||
self.received_finish_reason = response_obj["finish_reason"]
|
||||
except Exception as e:
|
||||
if self.received_finish_reason:
|
||||
raise e
|
||||
else:
|
||||
if self.sent_first_chunk is False:
|
||||
raise Exception("An unknown error occurred with the stream")
|
||||
self.received_finish_reason = "stop"
|
||||
elif self.custom_llm_provider == "vertex_ai" and not isinstance(
|
||||
chunk, ModelResponseStream
|
||||
):
|
||||
chunk = cast(Any, chunk)
|
||||
import proto # type: ignore
|
||||
|
||||
if hasattr(chunk, "candidates") is True:
|
||||
try:
|
||||
try:
|
||||
completion_obj["content"] = chunk.text # type: ignore
|
||||
except Exception as e:
|
||||
original_exception = e
|
||||
if "Part has no text." in str(e):
|
||||
## check for function calling
|
||||
function_call = (
|
||||
chunk.candidates[0].content.parts[0].function_call # type: ignore
|
||||
)
|
||||
|
||||
args_dict = {}
|
||||
|
||||
# Check if it's a RepeatedComposite instance
|
||||
for key, val in function_call.args.items():
|
||||
if isinstance(
|
||||
val,
|
||||
proto.marshal.collections.repeated.RepeatedComposite, # type: ignore
|
||||
):
|
||||
# If so, convert to list
|
||||
args_dict[key] = [v for v in val]
|
||||
else:
|
||||
args_dict[key] = val
|
||||
|
||||
try:
|
||||
args_str = json.dumps(args_dict)
|
||||
except Exception as e:
|
||||
raise e
|
||||
_delta_obj = litellm.utils.Delta(
|
||||
content=None,
|
||||
tool_calls=[
|
||||
{
|
||||
"id": f"call_{str(uuid.uuid4())}",
|
||||
"function": {
|
||||
"arguments": args_str,
|
||||
"name": function_call.name,
|
||||
},
|
||||
"type": "function",
|
||||
}
|
||||
],
|
||||
)
|
||||
_streaming_response = StreamingChoices(delta=_delta_obj)
|
||||
_model_response = ModelResponseStream()
|
||||
_model_response.choices = [_streaming_response]
|
||||
response_obj = {"original_chunk": _model_response}
|
||||
else:
|
||||
raise original_exception
|
||||
if (
|
||||
hasattr(chunk.candidates[0], "finish_reason") # type: ignore
|
||||
and chunk.candidates[0].finish_reason.name # type: ignore
|
||||
!= "FINISH_REASON_UNSPECIFIED"
|
||||
): # every non-final chunk in vertex ai has this
|
||||
self.received_finish_reason = map_finish_reason( # type: ignore
|
||||
chunk.candidates[0].finish_reason.name
|
||||
)
|
||||
except Exception:
|
||||
if chunk.candidates[0].finish_reason.name == "SAFETY": # type: ignore
|
||||
raise Exception(
|
||||
f"The response was blocked by VertexAI. {str(chunk)}"
|
||||
)
|
||||
else:
|
||||
completion_obj["content"] = str(chunk)
|
||||
elif self.custom_llm_provider == "petals":
|
||||
if self.completion_stream is None or len(self.completion_stream) == 0:
|
||||
if self.received_finish_reason is not None:
|
||||
raise StopIteration
|
||||
else:
|
||||
self.received_finish_reason = "stop"
|
||||
chunk_size = 30
|
||||
stream = cast(Any, self.completion_stream)
|
||||
new_chunk = stream[:chunk_size]
|
||||
completion_obj["content"] = new_chunk
|
||||
self.completion_stream = stream[chunk_size:]
|
||||
elif self.custom_llm_provider == "palm":
|
||||
# fake streaming
|
||||
response_obj = {}
|
||||
if self.completion_stream is None or len(self.completion_stream) == 0:
|
||||
if self.received_finish_reason is not None:
|
||||
raise StopIteration
|
||||
else:
|
||||
self.received_finish_reason = "stop"
|
||||
chunk_size = 30
|
||||
stream = cast(Any, self.completion_stream)
|
||||
new_chunk = stream[:chunk_size]
|
||||
completion_obj["content"] = new_chunk
|
||||
self.completion_stream = stream[chunk_size:]
|
||||
elif self.custom_llm_provider == "triton":
|
||||
response_obj = self.handle_triton_stream(chunk)
|
||||
completion_obj["content"] = response_obj["text"]
|
||||
print_verbose(f"completion obj content: {completion_obj['content']}")
|
||||
if response_obj["is_finished"]:
|
||||
self.received_finish_reason = response_obj["finish_reason"]
|
||||
elif self.custom_llm_provider == "text-completion-openai":
|
||||
response_obj = self.handle_openai_text_completion_chunk(chunk)
|
||||
completion_obj["content"] = response_obj["text"]
|
||||
print_verbose(f"completion obj content: {completion_obj['content']}")
|
||||
if response_obj["is_finished"]:
|
||||
self.received_finish_reason = response_obj["finish_reason"]
|
||||
if response_obj["usage"] is not None:
|
||||
setattr(
|
||||
model_response,
|
||||
"usage",
|
||||
litellm.Usage(
|
||||
prompt_tokens=response_obj["usage"].prompt_tokens,
|
||||
completion_tokens=response_obj["usage"].completion_tokens,
|
||||
total_tokens=response_obj["usage"].total_tokens,
|
||||
),
|
||||
)
|
||||
elif self.custom_llm_provider == "text-completion-codestral":
|
||||
if not isinstance(chunk, str):
|
||||
raise ValueError(f"chunk is not a string: {chunk}")
|
||||
response_obj = cast(
|
||||
dict[str, Any],
|
||||
litellm.CodestralTextCompletionConfig()._chunk_parser(chunk),
|
||||
)
|
||||
completion_obj["content"] = response_obj["text"]
|
||||
print_verbose(f"completion obj content: {completion_obj['content']}")
|
||||
if response_obj["is_finished"]:
|
||||
self.received_finish_reason = response_obj["finish_reason"]
|
||||
if "usage" in response_obj is not None:
|
||||
setattr(
|
||||
model_response,
|
||||
"usage",
|
||||
litellm.Usage(
|
||||
prompt_tokens=response_obj["usage"].prompt_tokens,
|
||||
completion_tokens=response_obj["usage"].completion_tokens,
|
||||
total_tokens=response_obj["usage"].total_tokens,
|
||||
),
|
||||
)
|
||||
elif self.custom_llm_provider == "azure_text":
|
||||
response_obj = self.handle_azure_text_completion_chunk(chunk)
|
||||
completion_obj["content"] = response_obj["text"]
|
||||
print_verbose(f"completion obj content: {completion_obj['content']}")
|
||||
if response_obj["is_finished"]:
|
||||
self.received_finish_reason = response_obj["finish_reason"]
|
||||
elif self.custom_llm_provider == "cached_response":
|
||||
chunk = cast(ModelResponseStream, chunk)
|
||||
response_obj = {
|
||||
"text": chunk.choices[0].delta.content,
|
||||
"is_finished": True,
|
||||
"finish_reason": chunk.choices[0].finish_reason,
|
||||
"original_chunk": chunk,
|
||||
"tool_calls": (
|
||||
chunk.choices[0].delta.tool_calls
|
||||
if hasattr(chunk.choices[0].delta, "tool_calls")
|
||||
else None
|
||||
),
|
||||
}
|
||||
|
||||
completion_obj["content"] = response_obj["text"]
|
||||
if response_obj["tool_calls"] is not None:
|
||||
completion_obj["tool_calls"] = response_obj["tool_calls"]
|
||||
print_verbose(f"completion obj content: {completion_obj['content']}")
|
||||
if hasattr(chunk, "id"):
|
||||
model_response.id = chunk.id
|
||||
self.response_id = chunk.id
|
||||
if hasattr(chunk, "system_fingerprint"):
|
||||
self.system_fingerprint = chunk.system_fingerprint
|
||||
if response_obj["is_finished"]:
|
||||
self.received_finish_reason = response_obj["finish_reason"]
|
||||
else: # openai / azure chat model
|
||||
if self.custom_llm_provider in [
|
||||
LlmProviders.AZURE.value,
|
||||
LlmProviders.AZURE_AI.value,
|
||||
]:
|
||||
if isinstance(chunk, BaseModel) and hasattr(chunk, "model"):
|
||||
# for azure, we need to pass the model from the original chunk
|
||||
self.model = getattr(chunk, "model", self.model)
|
||||
response_obj = self.handle_openai_chat_completion_chunk(chunk)
|
||||
if response_obj is None:
|
||||
return _ProviderChunkEarlyReturn(None)
|
||||
completion_obj["content"] = response_obj["text"]
|
||||
self.intermittent_finish_reason = response_obj.get("finish_reason", None)
|
||||
if response_obj["is_finished"]:
|
||||
if response_obj["finish_reason"] == "error":
|
||||
raise Exception(
|
||||
"{} raised a streaming error - finish_reason: error, no content string given. Received Chunk={}".format(
|
||||
self.custom_llm_provider, response_obj
|
||||
)
|
||||
)
|
||||
self.received_finish_reason = response_obj["finish_reason"]
|
||||
if response_obj.get("original_chunk", None) is not None:
|
||||
if hasattr(response_obj["original_chunk"], "id"):
|
||||
model_response = self.set_model_id(
|
||||
response_obj["original_chunk"].id, model_response
|
||||
)
|
||||
if hasattr(response_obj["original_chunk"], "system_fingerprint"):
|
||||
model_response.system_fingerprint = response_obj[
|
||||
"original_chunk"
|
||||
].system_fingerprint
|
||||
self.system_fingerprint = response_obj[
|
||||
"original_chunk"
|
||||
].system_fingerprint
|
||||
if response_obj["logprobs"] is not None:
|
||||
model_response.choices[0].logprobs = response_obj["logprobs"]
|
||||
|
||||
if response_obj["usage"] is not None:
|
||||
if isinstance(response_obj["usage"], dict):
|
||||
setattr(
|
||||
model_response,
|
||||
"usage",
|
||||
litellm.Usage(
|
||||
prompt_tokens=response_obj["usage"].get(
|
||||
"prompt_tokens", None
|
||||
)
|
||||
or None,
|
||||
completion_tokens=response_obj["usage"].get(
|
||||
"completion_tokens", None
|
||||
)
|
||||
or None,
|
||||
total_tokens=response_obj["usage"].get("total_tokens", None)
|
||||
or None,
|
||||
),
|
||||
)
|
||||
elif isinstance(response_obj["usage"], Usage):
|
||||
setattr(
|
||||
model_response,
|
||||
"usage",
|
||||
response_obj["usage"],
|
||||
)
|
||||
elif isinstance(response_obj["usage"], BaseModel):
|
||||
setattr(
|
||||
model_response,
|
||||
"usage",
|
||||
litellm.Usage(**response_obj["usage"].model_dump()),
|
||||
)
|
||||
return _ProviderChunkParsed(response_obj)
|
||||
|
||||
def chunk_creator(self, chunk: Any): # type: ignore
|
||||
if hasattr(chunk, "id"):
|
||||
self.response_id = chunk.id
|
||||
model_response = self.model_response_creator()
|
||||
response_obj: Dict[str, Any] = {}
|
||||
response_obj: dict[str, Any] = {}
|
||||
try:
|
||||
# return this for all models
|
||||
completion_obj: Dict[str, Any] = {"content": ""}
|
||||
from litellm.types.utils import GenericStreamingChunk as GChunk
|
||||
|
||||
if (
|
||||
isinstance(chunk, ModelResponseStream)
|
||||
and self.custom_llm_provider is not None
|
||||
and self.custom_llm_provider in litellm._custom_providers
|
||||
):
|
||||
_has_content = bool(
|
||||
chunk.choices
|
||||
and chunk.choices[0].delta is not None
|
||||
and (
|
||||
chunk.choices[0].delta.content
|
||||
or chunk.choices[0].delta.tool_calls
|
||||
)
|
||||
)
|
||||
if self.received_finish_reason is not None:
|
||||
if not _has_content:
|
||||
raise StopIteration
|
||||
if chunk.choices and chunk.choices[0].finish_reason:
|
||||
self.received_finish_reason = chunk.choices[0].finish_reason
|
||||
if not _has_content:
|
||||
return None
|
||||
# Strip finish_reason from the content chunk so it appears
|
||||
# only on the trailing empty-delta chunk (OpenAI spec).
|
||||
# finish_reason_handler() will emit the proper terminal chunk.
|
||||
chunk.choices[0].finish_reason = None # type: ignore[assignment]
|
||||
return chunk
|
||||
|
||||
if (
|
||||
isinstance(chunk, dict)
|
||||
and generic_chunk_has_all_required_fields(
|
||||
chunk=chunk
|
||||
) # check if chunk is a generic streaming chunk
|
||||
) or (
|
||||
self.custom_llm_provider
|
||||
and self.custom_llm_provider in litellm._custom_providers
|
||||
):
|
||||
if self.received_finish_reason is not None:
|
||||
_chunk_has_content = isinstance(chunk, dict) and (
|
||||
bool(chunk.get("text", ""))
|
||||
or chunk.get("tool_use") is not None
|
||||
# Usage-only final chunks are valid and needed to surface
|
||||
# finish_reason/usage to downstream translators.
|
||||
or chunk.get("usage") is not None
|
||||
)
|
||||
if not _chunk_has_content and (
|
||||
not isinstance(chunk, dict)
|
||||
or "provider_specific_fields" not in chunk
|
||||
):
|
||||
raise StopIteration
|
||||
anthropic_response_obj: GChunk = cast(GChunk, chunk)
|
||||
completion_obj["content"] = anthropic_response_obj["text"]
|
||||
if anthropic_response_obj["is_finished"]:
|
||||
self.received_finish_reason = anthropic_response_obj[
|
||||
"finish_reason"
|
||||
]
|
||||
|
||||
if anthropic_response_obj["finish_reason"]:
|
||||
self.intermittent_finish_reason = anthropic_response_obj[
|
||||
"finish_reason"
|
||||
]
|
||||
|
||||
if anthropic_response_obj["usage"] is not None:
|
||||
setattr(
|
||||
model_response,
|
||||
"usage",
|
||||
litellm.Usage(**anthropic_response_obj["usage"]),
|
||||
)
|
||||
|
||||
if (
|
||||
"tool_use" in anthropic_response_obj
|
||||
and anthropic_response_obj["tool_use"] is not None
|
||||
):
|
||||
completion_obj["tool_calls"] = [anthropic_response_obj["tool_use"]]
|
||||
|
||||
if (
|
||||
"provider_specific_fields" in anthropic_response_obj
|
||||
and anthropic_response_obj["provider_specific_fields"] is not None
|
||||
):
|
||||
for key, value in anthropic_response_obj[
|
||||
"provider_specific_fields"
|
||||
].items():
|
||||
setattr(model_response, key, value)
|
||||
|
||||
response_obj = cast(Dict[str, Any], anthropic_response_obj)
|
||||
elif self.model == "replicate" or self.custom_llm_provider == "replicate":
|
||||
response_obj = self.handle_replicate_chunk(chunk)
|
||||
completion_obj["content"] = response_obj["text"]
|
||||
if response_obj["is_finished"]:
|
||||
self.received_finish_reason = response_obj["finish_reason"]
|
||||
elif self.custom_llm_provider and self.custom_llm_provider == "predibase":
|
||||
response_obj = self.handle_predibase_chunk(chunk)
|
||||
completion_obj["content"] = response_obj["text"]
|
||||
if response_obj["is_finished"]:
|
||||
self.received_finish_reason = response_obj["finish_reason"]
|
||||
elif (
|
||||
self.custom_llm_provider and self.custom_llm_provider == "baseten"
|
||||
): # baseten doesn't provide streaming
|
||||
completion_obj["content"] = self.handle_baseten_chunk(chunk)
|
||||
elif (
|
||||
self.custom_llm_provider and self.custom_llm_provider == "ai21"
|
||||
): # ai21 doesn't provide streaming
|
||||
response_obj = self.handle_ai21_chunk(chunk)
|
||||
completion_obj["content"] = response_obj["text"]
|
||||
if response_obj["is_finished"]:
|
||||
self.received_finish_reason = response_obj["finish_reason"]
|
||||
elif self.custom_llm_provider and self.custom_llm_provider == "maritalk":
|
||||
response_obj = self.handle_maritalk_chunk(chunk)
|
||||
completion_obj["content"] = response_obj["text"]
|
||||
if response_obj["is_finished"]:
|
||||
self.received_finish_reason = response_obj["finish_reason"]
|
||||
elif self.custom_llm_provider and self.custom_llm_provider == "vllm":
|
||||
completion_obj["content"] = chunk[0].outputs[0].text
|
||||
elif (
|
||||
self.custom_llm_provider and self.custom_llm_provider == "aleph_alpha"
|
||||
): # aleph alpha doesn't provide streaming
|
||||
response_obj = self.handle_aleph_alpha_chunk(chunk)
|
||||
completion_obj["content"] = response_obj["text"]
|
||||
if response_obj["is_finished"]:
|
||||
self.received_finish_reason = response_obj["finish_reason"]
|
||||
elif self.custom_llm_provider == "nlp_cloud":
|
||||
try:
|
||||
response_obj = self.handle_nlp_cloud_chunk(chunk)
|
||||
completion_obj["content"] = response_obj["text"]
|
||||
if response_obj["is_finished"]:
|
||||
self.received_finish_reason = response_obj["finish_reason"]
|
||||
except Exception as e:
|
||||
if self.received_finish_reason:
|
||||
raise e
|
||||
else:
|
||||
if self.sent_first_chunk is False:
|
||||
raise Exception("An unknown error occurred with the stream")
|
||||
self.received_finish_reason = "stop"
|
||||
elif self.custom_llm_provider == "vertex_ai" and not isinstance(
|
||||
chunk, ModelResponseStream
|
||||
):
|
||||
import proto # type: ignore
|
||||
|
||||
if hasattr(chunk, "candidates") is True:
|
||||
try:
|
||||
try:
|
||||
completion_obj["content"] = chunk.text # type: ignore
|
||||
except Exception as e:
|
||||
original_exception = e
|
||||
if "Part has no text." in str(e):
|
||||
## check for function calling
|
||||
function_call = (
|
||||
chunk.candidates[0].content.parts[0].function_call # type: ignore
|
||||
)
|
||||
|
||||
args_dict = {}
|
||||
|
||||
# Check if it's a RepeatedComposite instance
|
||||
for key, val in function_call.args.items():
|
||||
if isinstance(
|
||||
val,
|
||||
proto.marshal.collections.repeated.RepeatedComposite, # type: ignore
|
||||
):
|
||||
# If so, convert to list
|
||||
args_dict[key] = [v for v in val]
|
||||
else:
|
||||
args_dict[key] = val
|
||||
|
||||
try:
|
||||
args_str = json.dumps(args_dict)
|
||||
except Exception as e:
|
||||
raise e
|
||||
_delta_obj = litellm.utils.Delta(
|
||||
content=None,
|
||||
tool_calls=[
|
||||
{
|
||||
"id": f"call_{str(uuid.uuid4())}",
|
||||
"function": {
|
||||
"arguments": args_str,
|
||||
"name": function_call.name,
|
||||
},
|
||||
"type": "function",
|
||||
}
|
||||
],
|
||||
)
|
||||
_streaming_response = StreamingChoices(delta=_delta_obj)
|
||||
_model_response = ModelResponseStream()
|
||||
_model_response.choices = [_streaming_response]
|
||||
response_obj = {"original_chunk": _model_response}
|
||||
else:
|
||||
raise original_exception
|
||||
if (
|
||||
hasattr(chunk.candidates[0], "finish_reason") # type: ignore
|
||||
and chunk.candidates[0].finish_reason.name # type: ignore
|
||||
!= "FINISH_REASON_UNSPECIFIED"
|
||||
): # every non-final chunk in vertex ai has this
|
||||
self.received_finish_reason = map_finish_reason( # type: ignore
|
||||
chunk.candidates[0].finish_reason.name
|
||||
)
|
||||
except Exception:
|
||||
if chunk.candidates[0].finish_reason.name == "SAFETY": # type: ignore
|
||||
raise Exception(
|
||||
f"The response was blocked by VertexAI. {str(chunk)}"
|
||||
)
|
||||
else:
|
||||
completion_obj["content"] = str(chunk)
|
||||
elif self.custom_llm_provider == "petals":
|
||||
if self.completion_stream is None or len(self.completion_stream) == 0:
|
||||
if self.received_finish_reason is not None:
|
||||
raise StopIteration
|
||||
else:
|
||||
self.received_finish_reason = "stop"
|
||||
chunk_size = 30
|
||||
new_chunk = self.completion_stream[:chunk_size] # type: ignore[index]
|
||||
completion_obj["content"] = new_chunk
|
||||
self.completion_stream = self.completion_stream[chunk_size:] # type: ignore[index]
|
||||
elif self.custom_llm_provider == "palm":
|
||||
# fake streaming
|
||||
response_obj = {}
|
||||
if self.completion_stream is None or len(self.completion_stream) == 0:
|
||||
if self.received_finish_reason is not None:
|
||||
raise StopIteration
|
||||
else:
|
||||
self.received_finish_reason = "stop"
|
||||
chunk_size = 30
|
||||
new_chunk = self.completion_stream[:chunk_size] # type: ignore[index]
|
||||
completion_obj["content"] = new_chunk
|
||||
self.completion_stream = self.completion_stream[chunk_size:] # type: ignore[index]
|
||||
elif self.custom_llm_provider == "triton":
|
||||
response_obj = self.handle_triton_stream(chunk)
|
||||
completion_obj["content"] = response_obj["text"]
|
||||
print_verbose(f"completion obj content: {completion_obj['content']}")
|
||||
if response_obj["is_finished"]:
|
||||
self.received_finish_reason = response_obj["finish_reason"]
|
||||
elif self.custom_llm_provider == "text-completion-openai":
|
||||
response_obj = self.handle_openai_text_completion_chunk(chunk)
|
||||
completion_obj["content"] = response_obj["text"]
|
||||
print_verbose(f"completion obj content: {completion_obj['content']}")
|
||||
if response_obj["is_finished"]:
|
||||
self.received_finish_reason = response_obj["finish_reason"]
|
||||
if response_obj["usage"] is not None:
|
||||
setattr(
|
||||
model_response,
|
||||
"usage",
|
||||
litellm.Usage(
|
||||
prompt_tokens=response_obj["usage"].prompt_tokens,
|
||||
completion_tokens=response_obj["usage"].completion_tokens,
|
||||
total_tokens=response_obj["usage"].total_tokens,
|
||||
),
|
||||
)
|
||||
elif self.custom_llm_provider == "text-completion-codestral":
|
||||
if not isinstance(chunk, str):
|
||||
raise ValueError(f"chunk is not a string: {chunk}")
|
||||
response_obj = cast(
|
||||
Dict[str, Any],
|
||||
litellm.CodestralTextCompletionConfig()._chunk_parser(chunk),
|
||||
)
|
||||
completion_obj["content"] = response_obj["text"]
|
||||
print_verbose(f"completion obj content: {completion_obj['content']}")
|
||||
if response_obj["is_finished"]:
|
||||
self.received_finish_reason = response_obj["finish_reason"]
|
||||
if "usage" in response_obj is not None:
|
||||
setattr(
|
||||
model_response,
|
||||
"usage",
|
||||
litellm.Usage(
|
||||
prompt_tokens=response_obj["usage"].prompt_tokens,
|
||||
completion_tokens=response_obj["usage"].completion_tokens,
|
||||
total_tokens=response_obj["usage"].total_tokens,
|
||||
),
|
||||
)
|
||||
elif self.custom_llm_provider == "azure_text":
|
||||
response_obj = self.handle_azure_text_completion_chunk(chunk)
|
||||
completion_obj["content"] = response_obj["text"]
|
||||
print_verbose(f"completion obj content: {completion_obj['content']}")
|
||||
if response_obj["is_finished"]:
|
||||
self.received_finish_reason = response_obj["finish_reason"]
|
||||
elif self.custom_llm_provider == "cached_response":
|
||||
chunk = cast(ModelResponseStream, chunk)
|
||||
response_obj = {
|
||||
"text": chunk.choices[0].delta.content,
|
||||
"is_finished": True,
|
||||
"finish_reason": chunk.choices[0].finish_reason,
|
||||
"original_chunk": chunk,
|
||||
"tool_calls": (
|
||||
chunk.choices[0].delta.tool_calls
|
||||
if hasattr(chunk.choices[0].delta, "tool_calls")
|
||||
else None
|
||||
),
|
||||
}
|
||||
|
||||
completion_obj["content"] = response_obj["text"]
|
||||
if response_obj["tool_calls"] is not None:
|
||||
completion_obj["tool_calls"] = response_obj["tool_calls"]
|
||||
print_verbose(f"completion obj content: {completion_obj['content']}")
|
||||
if hasattr(chunk, "id"):
|
||||
model_response.id = chunk.id
|
||||
self.response_id = chunk.id
|
||||
if hasattr(chunk, "system_fingerprint"):
|
||||
self.system_fingerprint = chunk.system_fingerprint
|
||||
if response_obj["is_finished"]:
|
||||
self.received_finish_reason = response_obj["finish_reason"]
|
||||
else: # openai / azure chat model
|
||||
if self.custom_llm_provider in [
|
||||
LlmProviders.AZURE.value,
|
||||
LlmProviders.AZURE_AI.value,
|
||||
]:
|
||||
if isinstance(chunk, BaseModel) and hasattr(chunk, "model"):
|
||||
# for azure, we need to pass the model from the original chunk
|
||||
self.model = getattr(chunk, "model", self.model)
|
||||
response_obj = self.handle_openai_chat_completion_chunk(chunk)
|
||||
if response_obj is None:
|
||||
return
|
||||
completion_obj["content"] = response_obj["text"]
|
||||
self.intermittent_finish_reason = response_obj.get(
|
||||
"finish_reason", None
|
||||
)
|
||||
if response_obj["is_finished"]:
|
||||
if response_obj["finish_reason"] == "error":
|
||||
raise Exception(
|
||||
"{} raised a streaming error - finish_reason: error, no content string given. Received Chunk={}".format(
|
||||
self.custom_llm_provider, response_obj
|
||||
)
|
||||
)
|
||||
self.received_finish_reason = response_obj["finish_reason"]
|
||||
if response_obj.get("original_chunk", None) is not None:
|
||||
if hasattr(response_obj["original_chunk"], "id"):
|
||||
model_response = self.set_model_id(
|
||||
response_obj["original_chunk"].id, model_response
|
||||
)
|
||||
if hasattr(response_obj["original_chunk"], "system_fingerprint"):
|
||||
model_response.system_fingerprint = response_obj[
|
||||
"original_chunk"
|
||||
].system_fingerprint
|
||||
self.system_fingerprint = response_obj[
|
||||
"original_chunk"
|
||||
].system_fingerprint
|
||||
if response_obj["logprobs"] is not None:
|
||||
model_response.choices[0].logprobs = response_obj["logprobs"]
|
||||
|
||||
if response_obj["usage"] is not None:
|
||||
if isinstance(response_obj["usage"], dict):
|
||||
setattr(
|
||||
model_response,
|
||||
"usage",
|
||||
litellm.Usage(
|
||||
prompt_tokens=response_obj["usage"].get(
|
||||
"prompt_tokens", None
|
||||
)
|
||||
or None,
|
||||
completion_tokens=response_obj["usage"].get(
|
||||
"completion_tokens", None
|
||||
)
|
||||
or None,
|
||||
total_tokens=response_obj["usage"].get(
|
||||
"total_tokens", None
|
||||
)
|
||||
or None,
|
||||
),
|
||||
)
|
||||
elif isinstance(response_obj["usage"], Usage):
|
||||
setattr(
|
||||
model_response,
|
||||
"usage",
|
||||
response_obj["usage"],
|
||||
)
|
||||
elif isinstance(response_obj["usage"], BaseModel):
|
||||
setattr(
|
||||
model_response,
|
||||
"usage",
|
||||
litellm.Usage(**response_obj["usage"].model_dump()),
|
||||
)
|
||||
completion_obj: dict[str, Any] = {"content": ""}
|
||||
dispatch_result = self._dispatch_provider_chunk(
|
||||
chunk=chunk,
|
||||
model_response=model_response,
|
||||
completion_obj=completion_obj,
|
||||
)
|
||||
if isinstance(dispatch_result, _ProviderChunkEarlyReturn):
|
||||
return dispatch_result.value
|
||||
response_obj = dispatch_result.response_obj
|
||||
|
||||
model_response.model = self.model
|
||||
## FUNCTION CALL PARSING
|
||||
|
|
@ -1980,11 +2005,29 @@ class CustomStreamWrapper:
|
|||
|
||||
except StopIteration:
|
||||
if self.sent_last_chunk is True:
|
||||
complete_streaming_response = litellm.stream_chunk_builder(
|
||||
chunks=self.chunks,
|
||||
messages=self.messages,
|
||||
logging_obj=self.logging_obj,
|
||||
)
|
||||
try:
|
||||
complete_streaming_response = litellm.stream_chunk_builder(
|
||||
chunks=self.chunks,
|
||||
messages=self.messages,
|
||||
logging_obj=self.logging_obj,
|
||||
)
|
||||
except Exception as e:
|
||||
# stream_chunk_builder can re-raise (as APIError) on large agentic
|
||||
# streams. The raise originates inside this except-StopIteration block,
|
||||
# so the sibling `except Exception` below does not catch it; it would
|
||||
# escape __next__ and drop the request from SpendLogs. Recover
|
||||
# best-effort usage from the raw chunks so cost is still tracked
|
||||
verbose_logger.warning(
|
||||
"stream_chunk_builder raised at end-of-stream (%s); logging "
|
||||
"best-effort usage from chunks.",
|
||||
str(e),
|
||||
)
|
||||
try:
|
||||
complete_streaming_response = self.model_response_creator(
|
||||
chunk={"usage": calculate_total_usage(chunks=self.chunks)}
|
||||
)
|
||||
except Exception:
|
||||
complete_streaming_response = None
|
||||
|
||||
response = self.model_response_creator()
|
||||
if complete_streaming_response is not None:
|
||||
|
|
@ -2209,11 +2252,27 @@ class CustomStreamWrapper:
|
|||
except (StopAsyncIteration, StopIteration):
|
||||
if self.sent_last_chunk is True:
|
||||
# log the final chunk with accurate streaming values
|
||||
complete_streaming_response = litellm.stream_chunk_builder(
|
||||
chunks=self.chunks,
|
||||
messages=self.messages,
|
||||
logging_obj=self.logging_obj,
|
||||
)
|
||||
try:
|
||||
complete_streaming_response = litellm.stream_chunk_builder(
|
||||
chunks=self.chunks,
|
||||
messages=self.messages,
|
||||
logging_obj=self.logging_obj,
|
||||
)
|
||||
except Exception as e:
|
||||
# see sync __next__: a raise from stream_chunk_builder inside this
|
||||
# except handler escapes __anext__ and drops the request from SpendLogs.
|
||||
# Recover best-effort usage from the raw chunks so cost is still tracked
|
||||
verbose_logger.warning(
|
||||
"stream_chunk_builder raised at end-of-stream (%s); logging "
|
||||
"best-effort usage from chunks.",
|
||||
str(e),
|
||||
)
|
||||
try:
|
||||
complete_streaming_response = self.model_response_creator(
|
||||
chunk={"usage": calculate_total_usage(chunks=self.chunks)}
|
||||
)
|
||||
except Exception:
|
||||
complete_streaming_response = None
|
||||
|
||||
response = self.model_response_creator()
|
||||
if complete_streaming_response is not None:
|
||||
|
|
@ -2290,6 +2349,7 @@ class CustomStreamWrapper:
|
|||
litellm.request_timeout
|
||||
)
|
||||
if self.logging_obj is not None:
|
||||
self._record_partial_usage_for_failure()
|
||||
## LOGGING
|
||||
threading.Thread(
|
||||
target=self.logging_obj.failure_handler,
|
||||
|
|
@ -2303,6 +2363,7 @@ class CustomStreamWrapper:
|
|||
except Exception as e:
|
||||
traceback_exception = traceback.format_exc()
|
||||
if self.logging_obj is not None:
|
||||
self._record_partial_usage_for_failure()
|
||||
## LOGGING
|
||||
threading.Thread(
|
||||
target=self.logging_obj.failure_handler,
|
||||
|
|
@ -2314,6 +2375,33 @@ class CustomStreamWrapper:
|
|||
)
|
||||
self._handle_stream_fallback_error(e)
|
||||
|
||||
def _record_partial_usage_for_failure(self) -> None:
|
||||
"""
|
||||
A stream that breaks mid-flight still billed the provider for the chunks
|
||||
already delivered. Recover that partial usage from the chunks seen so
|
||||
far and stash it, with its cost, on the logging object so the failure
|
||||
handler records the real partial spend instead of zero. A request that
|
||||
later recovers via a router fallback overwrites this with the combined
|
||||
success log on the same request id, so this never double counts.
|
||||
"""
|
||||
if self.logging_obj is None or not self.chunks:
|
||||
return
|
||||
try:
|
||||
partial_response = litellm.stream_chunk_builder(chunks=self.chunks)
|
||||
usage = cast(Optional[Usage], getattr(partial_response, "usage", None))
|
||||
if usage is None:
|
||||
return
|
||||
self.logging_obj.model_call_details["combined_usage_object"] = usage
|
||||
self.logging_obj.model_call_details["response_cost"] = (
|
||||
self.logging_obj._response_cost_calculator(result=partial_response)
|
||||
or 0.0
|
||||
)
|
||||
except Exception as recover_error:
|
||||
verbose_logger.debug(
|
||||
"could not recover partial usage for interrupted stream: %s",
|
||||
recover_error,
|
||||
)
|
||||
|
||||
def _handle_stream_fallback_error(self, e: Exception) -> "NoReturn":
|
||||
"""
|
||||
Common error handling for both __next__ and __anext__.
|
||||
|
|
|
|||
|
|
@ -859,7 +859,17 @@ class LiteLLMAnthropicMessagesAdapter:
|
|||
"""
|
||||
new_tools: List[ChatCompletionToolParam] = []
|
||||
tool_name_mapping: Dict[str, str] = {}
|
||||
mapped_tool_params = ["name", "input_schema", "description", "cache_control"]
|
||||
# "type" is the Anthropic tool type (e.g. "custom"); it must not be
|
||||
# merged into the OpenAI function `parameters` schema below, or it
|
||||
# overwrites the real parameters.type ("object") and the provider
|
||||
# rejects the request. See #30557.
|
||||
mapped_tool_params = [
|
||||
"name",
|
||||
"input_schema",
|
||||
"description",
|
||||
"cache_control",
|
||||
"type",
|
||||
]
|
||||
|
||||
for idx, tool in enumerate(tools):
|
||||
# Check if this is an Anthropic-native tool that should be kept as-is
|
||||
|
|
|
|||
|
|
@ -84,8 +84,14 @@ class AdvisorOrchestrationHandler(MessagesInterceptor):
|
|||
)
|
||||
# Optional routing overrides for the advisor sub-call (e.g. proxy routing).
|
||||
# If not set in the tool definition, litellm resolves from env vars.
|
||||
advisor_api_key: Optional[str] = advisor_tool.get("api_key")
|
||||
advisor_api_base: Optional[str] = advisor_tool.get("api_base")
|
||||
# The advisor tool is caller-controlled; only honor a client-supplied
|
||||
# api_base/api_key when the proxy has enabled clientside credentials,
|
||||
# otherwise let litellm resolve from server config.
|
||||
advisor_api_key: Optional[str] = None
|
||||
advisor_api_base: Optional[str] = None
|
||||
if _allow_client_side_advisor_credentials():
|
||||
advisor_api_key = advisor_tool.get("api_key")
|
||||
advisor_api_base = advisor_tool.get("api_base")
|
||||
|
||||
# Build the synthetic tool definition the provider will receive.
|
||||
synthetic_advisor_tool = _make_synthetic_advisor_tool()
|
||||
|
|
@ -181,6 +187,20 @@ class AdvisorOrchestrationHandler(MessagesInterceptor):
|
|||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _allow_client_side_advisor_credentials() -> bool:
|
||||
"""Whether a caller-supplied advisor api_base/api_key may be honored.
|
||||
|
||||
Gated on the proxy's ``allow_client_side_credentials`` opt-in. When the
|
||||
interceptor runs outside the proxy (SDK use), there is no admin boundary
|
||||
to protect, so client-supplied routing is allowed.
|
||||
"""
|
||||
try:
|
||||
from litellm.proxy.proxy_server import general_settings
|
||||
except (ImportError, ModuleNotFoundError):
|
||||
return True
|
||||
return general_settings.get("allow_client_side_credentials") is True
|
||||
|
||||
|
||||
def _make_synthetic_advisor_tool() -> Dict:
|
||||
"""Build a regular tool definition the executor provider can understand."""
|
||||
return {
|
||||
|
|
|
|||
|
|
@ -53,7 +53,13 @@ class APISerpentSearchConfig(BaseSearchConfig):
|
|||
api_base: Optional[str] = None,
|
||||
**kwargs,
|
||||
) -> Dict:
|
||||
api_key = api_key or get_secret_str("APISERPENT_API_KEY")
|
||||
api_key = self.resolve_server_api_key(
|
||||
caller_api_key=api_key,
|
||||
caller_api_base=api_base,
|
||||
key_env_vars=("APISERPENT_API_KEY",),
|
||||
base_env_var="APISERPENT_API_BASE",
|
||||
default_api_base=APISERPENT_BASE,
|
||||
)
|
||||
if not api_key:
|
||||
raise ValueError(
|
||||
"APISERPENT_API_KEY is not set. Set `APISERPENT_API_KEY` environment variable."
|
||||
|
|
|
|||
0
litellm/llms/base_llm/sandbox/__init__.py
Normal file
0
litellm/llms/base_llm/sandbox/__init__.py
Normal file
96
litellm/llms/base_llm/sandbox/transformation.py
Normal file
96
litellm/llms/base_llm/sandbox/transformation.py
Normal file
|
|
@ -0,0 +1,96 @@
|
|||
"""
|
||||
Base Sandbox transformation configuration.
|
||||
|
||||
A sandbox provider runs an executable string inside an isolated container and
|
||||
returns whatever the sandbox produced. The lifecycle is create container ->
|
||||
run code -> delete container; `code_interpreter_tool` combines all three.
|
||||
"""
|
||||
|
||||
from typing import Any, Union
|
||||
|
||||
import httpx
|
||||
|
||||
from pydantic import Field, PrivateAttr
|
||||
|
||||
from litellm.types.llms.base import LiteLLMPydanticObjectBase
|
||||
|
||||
SANDBOX_MAX_OUTPUT_BYTES = 10 * 1024 * 1024
|
||||
|
||||
|
||||
class ContainerHandle(LiteLLMPydanticObjectBase):
|
||||
"""A live sandbox container. Carries everything needed to reach it again."""
|
||||
|
||||
id: str
|
||||
provider: str
|
||||
domain: str | None = None
|
||||
|
||||
model_config = {"extra": "allow"}
|
||||
|
||||
_hidden_params: dict = PrivateAttr(default_factory=dict)
|
||||
|
||||
|
||||
class CodeExecutionResult(LiteLLMPydanticObjectBase):
|
||||
"""Passthrough of the sandbox's own execution output."""
|
||||
|
||||
stdout: str = ""
|
||||
stderr: str = ""
|
||||
results: list[dict[str, Any]] = Field(default_factory=list)
|
||||
error: dict[str, Any] | None = None
|
||||
execution_count: int | None = None
|
||||
object: str = "code_execution"
|
||||
|
||||
model_config = {"extra": "allow"}
|
||||
|
||||
_hidden_params: dict = PrivateAttr(default_factory=dict)
|
||||
|
||||
|
||||
class BaseSandboxConfig:
|
||||
"""Provider-agnostic sandbox operations."""
|
||||
|
||||
def validate_environment(self, api_key: str | None = None, **kwargs) -> str:
|
||||
raise NotImplementedError(
|
||||
"validate_environment must be implemented by provider"
|
||||
)
|
||||
|
||||
async def acreate_sandbox(
|
||||
self,
|
||||
*,
|
||||
template: str | None = None,
|
||||
timeout: int | None = None,
|
||||
allow_internet_access: bool | None = None,
|
||||
api_key: str | None = None,
|
||||
**kwargs,
|
||||
) -> ContainerHandle:
|
||||
raise NotImplementedError("acreate_sandbox must be implemented by provider")
|
||||
|
||||
async def arun_code(
|
||||
self,
|
||||
*,
|
||||
container: Union[ContainerHandle, str],
|
||||
code: str,
|
||||
api_key: str | None = None,
|
||||
**kwargs,
|
||||
) -> CodeExecutionResult:
|
||||
raise NotImplementedError("arun_code must be implemented by provider")
|
||||
|
||||
async def adelete_sandbox(
|
||||
self,
|
||||
*,
|
||||
container: Union[ContainerHandle, str],
|
||||
api_key: str | None = None,
|
||||
**kwargs,
|
||||
) -> bool:
|
||||
raise NotImplementedError("adelete_sandbox must be implemented by provider")
|
||||
|
||||
async def _read_capped_lines(self, response: httpx.Response) -> list[str]:
|
||||
lines: list[str] = []
|
||||
total = 0
|
||||
async for line in response.aiter_lines():
|
||||
total += len(line.encode("utf-8"))
|
||||
if total > SANDBOX_MAX_OUTPUT_BYTES:
|
||||
raise ValueError(
|
||||
f"Sandbox output exceeded {SANDBOX_MAX_OUTPUT_BYTES} bytes; aborting "
|
||||
"to avoid unbounded memory use."
|
||||
)
|
||||
lines.append(line)
|
||||
return lines
|
||||
|
|
@ -3,11 +3,13 @@ Base Search transformation configuration.
|
|||
"""
|
||||
|
||||
from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Union
|
||||
from urllib.parse import urlsplit
|
||||
|
||||
import httpx
|
||||
from pydantic import PrivateAttr
|
||||
|
||||
from litellm.llms.base_llm.chat.transformation import BaseLLMException
|
||||
from litellm.secret_managers.main import get_secret_str
|
||||
from litellm.types.llms.base import LiteLLMPydanticObjectBase
|
||||
|
||||
if TYPE_CHECKING:
|
||||
|
|
@ -16,6 +18,29 @@ else:
|
|||
LiteLLMLoggingObj = Any
|
||||
|
||||
|
||||
def _search_host(url: str) -> str:
|
||||
return urlsplit(url).netloc.lower()
|
||||
|
||||
|
||||
def _is_trusted_search_api_base(
|
||||
caller_api_base: str,
|
||||
default_api_base: str | None,
|
||||
base_env_var: str | None,
|
||||
) -> bool:
|
||||
candidate = _search_host(caller_api_base)
|
||||
if not candidate:
|
||||
return False
|
||||
trusted = {
|
||||
_search_host(base)
|
||||
for base in (
|
||||
default_api_base,
|
||||
get_secret_str(base_env_var) if base_env_var else None,
|
||||
)
|
||||
if base
|
||||
}
|
||||
return candidate in trusted
|
||||
|
||||
|
||||
class SearchResult(LiteLLMPydanticObjectBase):
|
||||
"""Single search result."""
|
||||
|
||||
|
|
@ -86,6 +111,60 @@ class BaseSearchConfig:
|
|||
"max_tokens_per_page",
|
||||
}
|
||||
|
||||
def _assert_trusted_api_base_for_server_credential(
|
||||
self,
|
||||
caller_api_base: str | None,
|
||||
default_api_base: str | None,
|
||||
base_env_var: str | None,
|
||||
credential_name: str,
|
||||
) -> None:
|
||||
"""
|
||||
Block sending a server-managed credential to a caller-chosen host.
|
||||
|
||||
A caller-supplied api_base is honored when constructing the request URL, so
|
||||
falling back to a server-configured secret while the caller controls the host
|
||||
leaks that secret. The provider default and the operator's own api_base
|
||||
override are the only trusted destinations for a server-managed credential.
|
||||
"""
|
||||
if not caller_api_base:
|
||||
return
|
||||
if _is_trusted_search_api_base(caller_api_base, default_api_base, base_env_var):
|
||||
return
|
||||
raise ValueError(
|
||||
f"Refusing to send the server-configured {credential_name} to the "
|
||||
f"caller-supplied api_base '{caller_api_base}'. Pass an explicit api_key "
|
||||
f"when overriding api_base for this search provider."
|
||||
)
|
||||
|
||||
def resolve_server_api_key(
|
||||
self,
|
||||
*,
|
||||
caller_api_key: str | None,
|
||||
caller_api_base: str | None,
|
||||
key_env_vars: tuple[str, ...],
|
||||
base_env_var: str | None,
|
||||
default_api_base: str | None,
|
||||
) -> str | None:
|
||||
"""
|
||||
Resolve a single-secret search API key, falling back to a server-managed
|
||||
secret only when the request targets a trusted host.
|
||||
|
||||
Returns the caller's key when provided, otherwise the first set
|
||||
server-managed secret (or None when none is set, for keyless providers).
|
||||
"""
|
||||
if caller_api_key:
|
||||
return caller_api_key
|
||||
server_key = next(
|
||||
(key for key in (get_secret_str(var) for var in key_env_vars) if key),
|
||||
None,
|
||||
)
|
||||
if server_key is None:
|
||||
return None
|
||||
self._assert_trusted_api_base_for_server_credential(
|
||||
caller_api_base, default_api_base, base_env_var, key_env_vars[0]
|
||||
)
|
||||
return server_key
|
||||
|
||||
def validate_environment(
|
||||
self,
|
||||
headers: Dict,
|
||||
|
|
|
|||
|
|
@ -10,7 +10,6 @@ from typing import (
|
|||
Callable,
|
||||
ClassVar,
|
||||
Dict,
|
||||
List,
|
||||
Literal,
|
||||
Optional,
|
||||
Tuple,
|
||||
|
|
@ -210,32 +209,11 @@ class BaseAWSLLM:
|
|||
"""
|
||||
Return a boto3.Credentials object
|
||||
"""
|
||||
## CHECK IS 'os.environ/' passed in
|
||||
params_to_check: List[Optional[str]] = [
|
||||
aws_access_key_id,
|
||||
aws_secret_access_key,
|
||||
aws_session_token,
|
||||
aws_region_name,
|
||||
aws_session_name,
|
||||
aws_profile_name,
|
||||
aws_role_name,
|
||||
aws_web_identity_token,
|
||||
aws_sts_endpoint,
|
||||
aws_external_id,
|
||||
]
|
||||
|
||||
# Iterate over parameters and update if needed
|
||||
for i, param in enumerate(params_to_check):
|
||||
if param and param.startswith("os.environ/"):
|
||||
_v = get_secret(param)
|
||||
if _v is not None and isinstance(_v, str):
|
||||
params_to_check[i] = _v
|
||||
elif param is None: # check if uppercase value in env
|
||||
key = self.aws_authentication_params[i]
|
||||
if key.upper() in os.environ:
|
||||
params_to_check[i] = os.getenv(key.upper())
|
||||
|
||||
# Assign updated values back to parameters
|
||||
# Only config-sourced credentials are expanded against the environment.
|
||||
# os.environ/<VAR> references in the model config are resolved at load time,
|
||||
# so any reference still present at this point is caller-supplied input and is
|
||||
# left as-is rather than expanded into a process environment variable. Each
|
||||
# unset param falls back to its matching fixed AWS_* ambient env var.
|
||||
(
|
||||
aws_access_key_id,
|
||||
aws_secret_access_key,
|
||||
|
|
@ -247,7 +225,21 @@ class BaseAWSLLM:
|
|||
aws_web_identity_token,
|
||||
aws_sts_endpoint,
|
||||
aws_external_id,
|
||||
) = params_to_check
|
||||
) = tuple(
|
||||
value if value is not None else os.getenv(env_var)
|
||||
for value, env_var in (
|
||||
(aws_access_key_id, "AWS_ACCESS_KEY_ID"),
|
||||
(aws_secret_access_key, "AWS_SECRET_ACCESS_KEY"),
|
||||
(aws_session_token, "AWS_SESSION_TOKEN"),
|
||||
(aws_region_name, "AWS_REGION_NAME"),
|
||||
(aws_session_name, "AWS_SESSION_NAME"),
|
||||
(aws_profile_name, "AWS_PROFILE_NAME"),
|
||||
(aws_role_name, "AWS_ROLE_NAME"),
|
||||
(aws_web_identity_token, "AWS_WEB_IDENTITY_TOKEN"),
|
||||
(aws_sts_endpoint, "AWS_STS_ENDPOINT"),
|
||||
(aws_external_id, "AWS_EXTERNAL_ID"),
|
||||
)
|
||||
)
|
||||
|
||||
verbose_logger.debug(
|
||||
"in get credentials\n"
|
||||
|
|
@ -845,6 +837,20 @@ class BaseAWSLLM:
|
|||
f"IN Web Identity Token: {aws_web_identity_token} | Role Name: {aws_role_name} | Session Name: {aws_session_name}"
|
||||
)
|
||||
|
||||
# get_secret() expands environment-variable references (an os.environ/<VAR>
|
||||
# prefix, or a bare name matching an environment variable). Config-sourced
|
||||
# references are expanded at load time, so such a reference reaching here is
|
||||
# caller-supplied input; reject it rather than expanding a process-environment
|
||||
# value for use as the token.
|
||||
if (
|
||||
aws_web_identity_token.startswith("os.environ/")
|
||||
or aws_web_identity_token in os.environ
|
||||
):
|
||||
raise AwsAuthError(
|
||||
message="Invalid web identity token reference.",
|
||||
status_code=400,
|
||||
)
|
||||
|
||||
oidc_token = get_secret(aws_web_identity_token)
|
||||
|
||||
if oidc_token is None:
|
||||
|
|
|
|||
|
|
@ -218,8 +218,20 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM):
|
|||
- Qualifier goes as query parameter
|
||||
- Only the payload goes in the request body
|
||||
|
||||
Payload shape:
|
||||
- ``prompt`` is always present and contains the text-only flatten of the
|
||||
last message's content (existing behavior).
|
||||
- ``content`` is added ONLY when the ``forward_multimodal_content`` litellm
|
||||
param is truthy AND the last message's ``content`` is a list containing a
|
||||
non-text block (e.g. ``image_url``, ``file``, ``input_audio``). The list is
|
||||
forwarded verbatim so the agent's ``@app.entrypoint`` handler can parse the
|
||||
OpenAI-shaped multimodal blocks. This is opt-in because an AgentCore agent
|
||||
must be explicitly written to read ``payload["content"]``; by default the
|
||||
payload stays byte-identical to the legacy ``{"prompt": "..."}`` shape.
|
||||
|
||||
Returns:
|
||||
dict: Payload dict containing the prompt
|
||||
dict: Payload dict containing the prompt and (optionally) the OpenAI
|
||||
content list.
|
||||
"""
|
||||
verbose_logger.debug(
|
||||
f"AgentCore transform_request - optional_params keys: {list(optional_params.keys())}"
|
||||
|
|
@ -231,6 +243,20 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM):
|
|||
# Create the payload - this is what goes in the body (raw JSON)
|
||||
payload: dict = {"prompt": prompt}
|
||||
|
||||
# Opt-in: when forward_multimodal_content is set, forward the OpenAI content
|
||||
# list verbatim under "content" so an attachment-aware agent can read the raw
|
||||
# blocks (image_url, file, etc.). Default off keeps the payload byte-identical
|
||||
# to the legacy {"prompt": "..."} shape for agents that only read the prompt.
|
||||
if self._should_forward_multimodal_content(optional_params, litellm_params):
|
||||
last_content = messages[-1].get("content")
|
||||
if isinstance(last_content, list) and any(
|
||||
isinstance(block, dict) and block.get("type") not in (None, "text")
|
||||
for block in last_content
|
||||
):
|
||||
# Copy so the payload never aliases messages[-1]["content"]; shallow,
|
||||
# not deep, to avoid cloning large base64 media on the request path.
|
||||
payload["content"] = list(last_content)
|
||||
|
||||
# Get or generate session ID - this goes in the header
|
||||
runtime_session_id = self._get_runtime_session_id(optional_params)
|
||||
headers["X-Amzn-Bedrock-AgentCore-Runtime-Session-Id"] = runtime_session_id
|
||||
|
|
@ -246,6 +272,29 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM):
|
|||
verbose_logger.debug(f"PAYLOAD: {payload}")
|
||||
return payload
|
||||
|
||||
@staticmethod
|
||||
def _should_forward_multimodal_content(
|
||||
optional_params: dict, litellm_params: dict
|
||||
) -> bool:
|
||||
"""Whether to forward raw OpenAI content blocks under ``payload["content"]``.
|
||||
|
||||
Opt-in via the ``forward_multimodal_content`` litellm param (default ``False``)
|
||||
because AgentCore agents must be explicitly written to read the field. The
|
||||
value may arrive as a bool or a config/env string ("true", "1", ...). Checks
|
||||
``optional_params`` first (where other AgentCore params land), then
|
||||
``litellm_params``.
|
||||
"""
|
||||
for source in (optional_params, litellm_params):
|
||||
if not isinstance(source, dict):
|
||||
continue
|
||||
value = source.get("forward_multimodal_content")
|
||||
if value is None:
|
||||
continue
|
||||
if isinstance(value, str):
|
||||
return value.strip().lower() in ("1", "true", "yes", "on")
|
||||
return bool(value)
|
||||
return False
|
||||
|
||||
def _extract_sse_json(self, line: str) -> Optional[Dict]:
|
||||
"""Extract and parse JSON from an SSE data line."""
|
||||
if not line.startswith("data:"):
|
||||
|
|
|
|||
|
|
@ -70,6 +70,7 @@ from ..base_aws_llm import BaseAWSLLM
|
|||
from ..common_utils import (
|
||||
BedrockError,
|
||||
ModelResponseIterator,
|
||||
build_bedrock_stream_error,
|
||||
get_bedrock_response_stream_shape,
|
||||
get_bedrock_tool_name,
|
||||
)
|
||||
|
|
@ -1841,23 +1842,7 @@ class AWSEventStreamDecoder:
|
|||
parsed_response = self.parser.parse(response_dict, response_stream_shape)
|
||||
|
||||
if response_dict["status_code"] != 200:
|
||||
decoded_body = response_dict["body"].decode()
|
||||
if isinstance(decoded_body, dict):
|
||||
error_message = decoded_body.get("message")
|
||||
elif isinstance(decoded_body, str):
|
||||
error_message = decoded_body
|
||||
else:
|
||||
error_message = ""
|
||||
exception_status = response_dict["headers"].get(":exception-type")
|
||||
error_message = exception_status + " " + error_message
|
||||
raise BedrockError(
|
||||
status_code=response_dict["status_code"],
|
||||
message=(
|
||||
json.dumps(error_message)
|
||||
if isinstance(error_message, dict)
|
||||
else error_message
|
||||
),
|
||||
)
|
||||
raise build_bedrock_stream_error(response_dict, response_stream_shape)
|
||||
if "chunk" in parsed_response:
|
||||
chunk = parsed_response.get("chunk")
|
||||
if not chunk:
|
||||
|
|
|
|||
|
|
@ -7,9 +7,21 @@ Common utilities used across bedrock chat/embedding/image generation
|
|||
import functools
|
||||
import json
|
||||
import os
|
||||
from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Union
|
||||
from typing import (
|
||||
TYPE_CHECKING,
|
||||
Any,
|
||||
Dict,
|
||||
List,
|
||||
Literal,
|
||||
Mapping,
|
||||
Optional,
|
||||
TypedDict,
|
||||
Union,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from botocore.model import Shape
|
||||
|
||||
from litellm.types.llms.bedrock import BedrockCreateBatchRequest
|
||||
|
||||
import httpx
|
||||
|
|
@ -1132,6 +1144,39 @@ def get_bedrock_response_stream_shape():
|
|||
return _load_bedrock_response_stream_shape()
|
||||
|
||||
|
||||
class BedrockEventStreamResponseDict(TypedDict):
|
||||
status_code: int
|
||||
headers: Mapping[str, str]
|
||||
body: bytes
|
||||
|
||||
|
||||
def build_bedrock_stream_error(
|
||||
response_dict: BedrockEventStreamResponseDict,
|
||||
response_stream_shape: Shape | None,
|
||||
) -> BedrockError:
|
||||
"""Build a BedrockError for a non-200 event-stream error event.
|
||||
|
||||
botocore hard-codes HTTP 400 on every mid-stream error event, so the modeled
|
||||
ResponseStream member's httpStatusCode is the real status. Resolve it from the
|
||||
shape and fall back to the raw status when the type is not modeled.
|
||||
"""
|
||||
exception_type = response_dict["headers"].get(":exception-type")
|
||||
decoded_body = response_dict["body"].decode()
|
||||
message = f"{exception_type} {decoded_body}" if exception_type else decoded_body
|
||||
|
||||
status_code = response_dict["status_code"]
|
||||
if exception_type is not None and response_stream_shape is not None:
|
||||
member = response_stream_shape.members.get(exception_type)
|
||||
if member is not None:
|
||||
modeled_status = (
|
||||
(member.metadata or {}).get("error", {}).get("httpStatusCode")
|
||||
)
|
||||
if modeled_status is not None:
|
||||
status_code = int(modeled_status)
|
||||
|
||||
return BedrockError(status_code=status_code, message=message)
|
||||
|
||||
|
||||
class BedrockEventStreamDecoderBase:
|
||||
"""
|
||||
Base class for event stream decoding for Bedrock
|
||||
|
|
@ -1156,23 +1201,7 @@ class BedrockEventStreamDecoderBase:
|
|||
parsed_response = self.parser.parse(response_dict, response_stream_shape)
|
||||
|
||||
if response_dict["status_code"] != 200:
|
||||
decoded_body = response_dict["body"].decode()
|
||||
if isinstance(decoded_body, dict):
|
||||
error_message = decoded_body.get("message")
|
||||
elif isinstance(decoded_body, str):
|
||||
error_message = decoded_body
|
||||
else:
|
||||
error_message = ""
|
||||
exception_status = response_dict["headers"].get(":exception-type")
|
||||
error_message = exception_status + " " + error_message
|
||||
raise BedrockError(
|
||||
status_code=response_dict["status_code"],
|
||||
message=(
|
||||
json.dumps(error_message)
|
||||
if isinstance(error_message, dict)
|
||||
else error_message
|
||||
),
|
||||
)
|
||||
raise build_bedrock_stream_error(response_dict, response_stream_shape)
|
||||
if "chunk" in parsed_response:
|
||||
chunk = parsed_response.get("chunk")
|
||||
if not chunk:
|
||||
|
|
|
|||
|
|
@ -1,8 +1,6 @@
|
|||
import asyncio
|
||||
import base64
|
||||
import os
|
||||
from types import MappingProxyType
|
||||
from typing import Any, Coroutine, Mapping, Optional, Tuple, Union, cast
|
||||
from collections.abc import Mapping
|
||||
from typing import Any, Coroutine, Optional, Tuple, Union
|
||||
|
||||
import httpx
|
||||
|
||||
|
|
@ -17,7 +15,6 @@ from litellm.types.llms.openai import (
|
|||
FileContentRequest,
|
||||
HttpxBinaryResponseContent,
|
||||
)
|
||||
from litellm.types.utils import SpecialEnums
|
||||
|
||||
from ..base_aws_llm import BaseAWSLLM
|
||||
|
||||
|
|
@ -37,40 +34,9 @@ class BedrockFilesHandler(BaseAWSLLM):
|
|||
)
|
||||
|
||||
def _extract_s3_uri_from_file_id(self, file_id: str) -> str:
|
||||
"""
|
||||
Extract S3 URI from encoded file ID.
|
||||
from .transformation import extract_s3_uri_from_file_id
|
||||
|
||||
The file ID can be in two formats:
|
||||
1. Base64-encoded unified file ID containing: llm_output_file_id,s3://bucket/path
|
||||
2. Direct S3 URI: s3://bucket/litellm-managed-prefix/path
|
||||
|
||||
Args:
|
||||
file_id: Encoded file ID or direct S3 URI
|
||||
|
||||
Returns:
|
||||
S3 URI (e.g., "s3://bucket-name/path/to/file")
|
||||
"""
|
||||
# First, try to decode if it's a base64-encoded unified file ID
|
||||
try:
|
||||
# Add padding if needed
|
||||
padded = file_id + "=" * (-len(file_id) % 4)
|
||||
decoded = base64.urlsafe_b64decode(padded).decode()
|
||||
|
||||
# Check if it's a unified file ID format
|
||||
if decoded.startswith(SpecialEnums.LITELM_MANAGED_FILE_ID_PREFIX.value):
|
||||
# Extract llm_output_file_id from the decoded string
|
||||
if "llm_output_file_id," in decoded:
|
||||
s3_uri = decoded.split("llm_output_file_id,")[1].split(";")[0]
|
||||
return s3_uri
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# If not base64 encoded or doesn't contain llm_output_file_id, accept only
|
||||
# explicit S3 URIs. Bucket and key validation happens before any S3 call.
|
||||
if file_id.startswith("s3://"):
|
||||
return file_id
|
||||
|
||||
raise ValueError("file_id must be a managed LiteLLM S3 file id")
|
||||
return extract_s3_uri_from_file_id(file_id)
|
||||
|
||||
def _parse_s3_uri(
|
||||
self,
|
||||
|
|
@ -95,26 +61,12 @@ class BedrockFilesHandler(BaseAWSLLM):
|
|||
allow_legacy_cloud_file_ids=allow_legacy_cloud_file_ids,
|
||||
)
|
||||
|
||||
def _get_configured_s3_bucket_name(self, litellm_params: dict) -> str:
|
||||
trusted_model_credentials = litellm_params.get(
|
||||
"_litellm_internal_model_credentials"
|
||||
)
|
||||
bucket_name = None
|
||||
if isinstance(trusted_model_credentials, type(MappingProxyType({}))):
|
||||
trusted_model_credentials_mapping = cast(
|
||||
Mapping[str, Any], trusted_model_credentials
|
||||
)
|
||||
candidate_bucket_name = trusted_model_credentials_mapping.get(
|
||||
"s3_bucket_name"
|
||||
)
|
||||
if isinstance(candidate_bucket_name, str):
|
||||
bucket_name = candidate_bucket_name
|
||||
bucket_name = bucket_name or os.getenv("AWS_S3_BUCKET_NAME")
|
||||
if not bucket_name:
|
||||
raise ValueError(
|
||||
"S3 bucket_name is required. Set 's3_bucket_name' in proxy config or AWS_S3_BUCKET_NAME for Bedrock file content retrieval."
|
||||
)
|
||||
return bucket_name
|
||||
def _get_configured_s3_bucket_name(
|
||||
self, litellm_params: Mapping[str, object]
|
||||
) -> str:
|
||||
from .transformation import get_configured_s3_bucket_name
|
||||
|
||||
return get_configured_s3_bucket_name(litellm_params)
|
||||
|
||||
async def afile_content(
|
||||
self,
|
||||
|
|
|
|||
|
|
@ -1,23 +1,37 @@
|
|||
import base64
|
||||
import json
|
||||
import os
|
||||
import time
|
||||
from typing import Any, Dict, List, Optional, Tuple, Union
|
||||
from collections.abc import Mapping, MutableMapping
|
||||
from types import MappingProxyType
|
||||
from typing import (
|
||||
Any,
|
||||
Dict,
|
||||
List,
|
||||
Optional,
|
||||
Tuple,
|
||||
Union,
|
||||
)
|
||||
from urllib.parse import unquote
|
||||
|
||||
import httpx
|
||||
from httpx import Headers, Response
|
||||
from openai.types.file_deleted import FileDeleted
|
||||
from pydantic import BaseModel, ConfigDict
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm._uuid import uuid
|
||||
from litellm.files.utils import FilesAPIUtils
|
||||
from litellm.litellm_core_utils.cloud_storage_security import (
|
||||
BEDROCK_MANAGED_S3_BATCH_PREFIX,
|
||||
BEDROCK_MANAGED_S3_PREFIXES,
|
||||
BEDROCK_MANAGED_S3_UPLOAD_PREFIX,
|
||||
build_managed_cloud_object_name,
|
||||
encode_s3_object_key_for_url,
|
||||
sanitize_cloud_object_component,
|
||||
should_allow_legacy_cloud_file_ids,
|
||||
split_configured_cloud_bucket_name,
|
||||
validate_managed_cloud_file_id,
|
||||
)
|
||||
from litellm.litellm_core_utils.prompt_templates.common_utils import extract_file_data
|
||||
from litellm.llms.base_llm.chat.transformation import BaseLLMException
|
||||
|
|
@ -28,18 +42,98 @@ from litellm.llms.base_llm.files.transformation import (
|
|||
from litellm.types.llms.openai import (
|
||||
AllMessageValues,
|
||||
CreateFileRequest,
|
||||
FileContentRequest,
|
||||
FileTypes,
|
||||
HttpxBinaryResponseContent,
|
||||
OpenAICreateFileRequestOptionalParams,
|
||||
OpenAIFileObject,
|
||||
PathLike,
|
||||
)
|
||||
from litellm.types.utils import ExtractedFileData, LlmProviders
|
||||
from litellm.types.utils import ExtractedFileData, LlmProviders, SpecialEnums
|
||||
from litellm.utils import get_llm_provider
|
||||
|
||||
from ..base_aws_llm import BaseAWSLLM
|
||||
from ..common_utils import BedrockError
|
||||
|
||||
# litellm_params key used to hand the SigV4-signed GET headers from
|
||||
# `transform_file_content_request` to `validate_environment` (the only hook
|
||||
# the shared file-content HTTP handler exposes for setting request headers).
|
||||
# Same pattern as the `upload_url` handoff in `transform_create_file_request`.
|
||||
S3_SIGNED_GET_HEADERS_PARAM = "_s3_signed_get_headers"
|
||||
|
||||
|
||||
class _BedrockS3RequestParams(BaseModel):
|
||||
"""Typed view of the credential/region params the S3 GetObject path reads."""
|
||||
|
||||
model_config = ConfigDict(extra="ignore")
|
||||
|
||||
aws_access_key_id: str | None = None
|
||||
aws_secret_access_key: str | None = None
|
||||
aws_session_token: str | None = None
|
||||
aws_region_name: str | None = None
|
||||
aws_session_name: str | None = None
|
||||
aws_profile_name: str | None = None
|
||||
aws_role_name: str | None = None
|
||||
aws_web_identity_token: str | None = None
|
||||
aws_sts_endpoint: str | None = None
|
||||
s3_region_name: str | None = None
|
||||
s3_endpoint_url: str | None = None
|
||||
|
||||
|
||||
class _TrustedS3ModelCredentials(BaseModel):
|
||||
"""The S3 bucket the server trusts file ids against, from the deployment snapshot."""
|
||||
|
||||
model_config = ConfigDict(extra="ignore")
|
||||
|
||||
s3_bucket_name: str | None = None
|
||||
|
||||
|
||||
def extract_s3_uri_from_file_id(file_id: str) -> str:
|
||||
"""
|
||||
Resolve a Bedrock file id to its S3 URI.
|
||||
|
||||
Accepts either a base64-encoded LiteLLM unified file id (whose decoded
|
||||
form carries `llm_output_file_id,s3://...`) or a direct `s3://` URI.
|
||||
"""
|
||||
try:
|
||||
padded = file_id + "=" * (-len(file_id) % 4)
|
||||
decoded = base64.urlsafe_b64decode(padded).decode()
|
||||
|
||||
if decoded.startswith(SpecialEnums.LITELM_MANAGED_FILE_ID_PREFIX.value):
|
||||
if "llm_output_file_id," in decoded:
|
||||
return decoded.split("llm_output_file_id,")[1].split(";")[0]
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
if file_id.startswith("s3://"):
|
||||
return file_id
|
||||
|
||||
raise ValueError("file_id must be a managed LiteLLM S3 file id")
|
||||
|
||||
|
||||
def get_configured_s3_bucket_name(litellm_params: Mapping[str, object]) -> str:
|
||||
"""
|
||||
Resolve the server-configured S3 bucket for Bedrock file operations.
|
||||
|
||||
Only trusts the immutable server-side credential snapshot or the
|
||||
environment; never a request-supplied param, since the bucket is what
|
||||
`validate_managed_cloud_file_id` checks file ids against.
|
||||
"""
|
||||
trusted_model_credentials = litellm_params.get(
|
||||
"_litellm_internal_model_credentials"
|
||||
)
|
||||
bucket_name: str | None = None
|
||||
if isinstance(trusted_model_credentials, MappingProxyType):
|
||||
snapshot: dict[str, object] = {}
|
||||
snapshot.update(trusted_model_credentials) # any-ok: untyped snapshot
|
||||
bucket_name = _TrustedS3ModelCredentials.model_validate(snapshot).s3_bucket_name
|
||||
bucket_name = bucket_name or os.getenv("AWS_S3_BUCKET_NAME")
|
||||
if not bucket_name:
|
||||
raise ValueError(
|
||||
"S3 bucket_name is required. Set 's3_bucket_name' in proxy config or AWS_S3_BUCKET_NAME for Bedrock file content retrieval."
|
||||
)
|
||||
return bucket_name
|
||||
|
||||
|
||||
class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig):
|
||||
"""
|
||||
|
|
@ -63,16 +157,21 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig):
|
|||
|
||||
def validate_environment(
|
||||
self,
|
||||
headers: dict,
|
||||
headers: MutableMapping[str, object],
|
||||
model: str,
|
||||
messages: List[AllMessageValues],
|
||||
optional_params: dict,
|
||||
litellm_params: dict,
|
||||
api_key: Optional[str] = None,
|
||||
api_base: Optional[str] = None,
|
||||
litellm_params: MutableMapping[str, object],
|
||||
api_key: str | None = None,
|
||||
api_base: str | None = None,
|
||||
) -> dict:
|
||||
# No additional headers needed for S3 uploads - AWS credentials handled by BaseAWSLLM
|
||||
return headers
|
||||
result: dict[str, object] = {}
|
||||
result.update(headers)
|
||||
signed_headers = litellm_params.pop(S3_SIGNED_GET_HEADERS_PARAM, None)
|
||||
if isinstance(signed_headers, Mapping):
|
||||
result.update(signed_headers) # any-ok: untyped handoff headers
|
||||
# otherwise no extra headers - AWS credentials are handled by BaseAWSLLM
|
||||
return result
|
||||
|
||||
def _get_content_from_openai_file(self, openai_file_content: FileTypes) -> str:
|
||||
"""
|
||||
|
|
@ -927,23 +1026,114 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig):
|
|||
|
||||
def transform_file_content_request(
|
||||
self,
|
||||
file_content_request,
|
||||
optional_params: dict,
|
||||
litellm_params: dict,
|
||||
) -> tuple[str, dict]:
|
||||
raise NotImplementedError(
|
||||
"BedrockFilesConfig does not support file content retrieval"
|
||||
file_content_request: FileContentRequest,
|
||||
optional_params: Mapping[str, object],
|
||||
litellm_params: MutableMapping[str, object],
|
||||
) -> tuple[str, dict[str, str]]:
|
||||
"""
|
||||
Build a SigV4-signed S3 GetObject request for a Bedrock batch file.
|
||||
|
||||
Bedrock batch file ids are `s3://bucket/key` URIs (or unified ids
|
||||
that decode to one); the bucket and key are validated against the
|
||||
server-configured bucket before any request is signed.
|
||||
"""
|
||||
file_id = file_content_request.get("file_id")
|
||||
if not file_id:
|
||||
raise ValueError("file_id is required for Bedrock file content retrieval")
|
||||
|
||||
s3_uri = extract_s3_uri_from_file_id(file_id)
|
||||
bucket_name, object_key = validate_managed_cloud_file_id(
|
||||
file_id=s3_uri,
|
||||
scheme="s3://",
|
||||
configured_bucket_name=get_configured_s3_bucket_name(litellm_params),
|
||||
allowed_object_prefixes=BEDROCK_MANAGED_S3_PREFIXES,
|
||||
allow_legacy_cloud_file_ids=should_allow_legacy_cloud_file_ids(
|
||||
litellm_params
|
||||
),
|
||||
)
|
||||
|
||||
# The shared file-content handler passes optional_params={}, so AWS
|
||||
# credentials/region arrive via litellm_params here (unlike the upload
|
||||
# path). s3_region_name wins over aws_region_name, same priority as
|
||||
# get_complete_file_url above.
|
||||
merged_params: dict[str, object] = {}
|
||||
merged_params.update(litellm_params)
|
||||
merged_params.update(optional_params)
|
||||
request_params = _BedrockS3RequestParams.model_validate(merged_params)
|
||||
|
||||
region_preference = (
|
||||
request_params.s3_region_name or request_params.aws_region_name
|
||||
)
|
||||
region_params: dict[str, str | None] = {"aws_region_name": region_preference}
|
||||
aws_region_name = self._get_aws_region_name(
|
||||
optional_params=region_params, model=""
|
||||
)
|
||||
|
||||
s3_endpoint_url = (
|
||||
request_params.s3_endpoint_url
|
||||
or f"https://s3.{aws_region_name}.amazonaws.com"
|
||||
).rstrip("/")
|
||||
url = f"{s3_endpoint_url}/{bucket_name}/{encode_s3_object_key_for_url(object_key)}"
|
||||
|
||||
litellm_params[S3_SIGNED_GET_HEADERS_PARAM] = self._sign_s3_get_request(
|
||||
api_base=url,
|
||||
aws_region_name=aws_region_name,
|
||||
request_params=request_params,
|
||||
)
|
||||
return url, {}
|
||||
|
||||
def _sign_s3_get_request(
|
||||
self,
|
||||
api_base: str,
|
||||
aws_region_name: str,
|
||||
request_params: _BedrockS3RequestParams,
|
||||
) -> dict[str, str]:
|
||||
"""
|
||||
SigV4-sign an S3 GetObject request, mirroring `_sign_s3_request` (PUT).
|
||||
"""
|
||||
try:
|
||||
import hashlib
|
||||
|
||||
from botocore.auth import SigV4Auth
|
||||
from botocore.awsrequest import AWSRequest
|
||||
except ImportError:
|
||||
raise ImportError("Missing boto3 to call bedrock. Run 'pip install boto3'.")
|
||||
|
||||
credentials = self.get_credentials( # any-ok: boto3 Credentials is untyped
|
||||
aws_access_key_id=request_params.aws_access_key_id,
|
||||
aws_secret_access_key=request_params.aws_secret_access_key,
|
||||
aws_session_token=request_params.aws_session_token,
|
||||
aws_region_name=aws_region_name,
|
||||
aws_session_name=request_params.aws_session_name,
|
||||
aws_profile_name=request_params.aws_profile_name,
|
||||
aws_role_name=request_params.aws_role_name,
|
||||
aws_web_identity_token=request_params.aws_web_identity_token,
|
||||
aws_sts_endpoint=request_params.aws_sts_endpoint,
|
||||
)
|
||||
|
||||
empty_body_hash = hashlib.sha256(b"").hexdigest()
|
||||
aws_request = AWSRequest( # any-ok: botocore AWSRequest is untyped
|
||||
method="GET",
|
||||
url=api_base,
|
||||
headers={"x-amz-content-sha256": empty_body_hash},
|
||||
)
|
||||
auth = SigV4Auth(credentials, "s3", aws_region_name) # any-ok: botocore untyped
|
||||
auth.add_auth(aws_request) # any-ok: botocore request mutation is untyped
|
||||
return dict(aws_request.headers) # any-ok: botocore headers are untyped
|
||||
|
||||
def transform_file_content_response(
|
||||
self,
|
||||
raw_response: httpx.Response,
|
||||
logging_obj: LiteLLMLoggingObj,
|
||||
litellm_params: dict,
|
||||
) -> HttpxBinaryResponseContent:
|
||||
raise NotImplementedError(
|
||||
"BedrockFilesConfig does not support file content retrieval"
|
||||
)
|
||||
if raw_response.status_code >= 400:
|
||||
raise BedrockError(
|
||||
status_code=raw_response.status_code,
|
||||
message=raw_response.text,
|
||||
headers=raw_response.headers,
|
||||
)
|
||||
return HttpxBinaryResponseContent(response=raw_response)
|
||||
|
||||
|
||||
class BedrockJsonlFilesTransformation:
|
||||
|
|
|
|||
|
|
@ -4,8 +4,10 @@ Amazon Bedrock Mantle - OpenAI-compatible inference engine in Amazon Bedrock.
|
|||
API docs: https://docs.aws.amazon.com/bedrock/latest/userguide/bedrock-mantle.html
|
||||
|
||||
Base URL: https://bedrock-mantle.{region}.api.aws/v1
|
||||
Auth: AWS Bedrock API key as Bearer token (set via BEDROCK_MANTLE_API_KEY env var)
|
||||
or region-aware key via BEDROCK_MANTLE_{REGION}_API_KEY.
|
||||
Auth: Bearer token (litellm_params.api_key, BEDROCK_MANTLE_API_KEY, or the
|
||||
standard AWS_BEARER_TOKEN_BEDROCK) when present; otherwise AWS SigV4
|
||||
(service "bedrock") over the standard credential chain. See
|
||||
BedrockMantleAuthMixin in common_utils.
|
||||
"""
|
||||
|
||||
from typing import Iterator, AsyncIterator, Any, List, Optional, Tuple, Union
|
||||
|
|
@ -13,20 +15,27 @@ from typing import Iterator, AsyncIterator, Any, List, Optional, Tuple, Union
|
|||
import litellm
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
|
||||
from litellm.llms.bedrock_mantle.common_utils import (
|
||||
BEDROCK_MANTLE_DEFAULT_REGION,
|
||||
BedrockMantleAuthMixin,
|
||||
)
|
||||
from litellm.secret_managers.main import get_secret_str
|
||||
from litellm.types.llms.openai import AllMessageValues
|
||||
from litellm.types.router import GenericLiteLLMParams
|
||||
|
||||
from ..common_utils import mantle_base_segment
|
||||
from ...openai_like.chat.transformation import OpenAILikeChatConfig
|
||||
|
||||
BEDROCK_MANTLE_DEFAULT_REGION = "us-east-1"
|
||||
|
||||
|
||||
class BedrockMantleChatConfig(OpenAILikeChatConfig):
|
||||
class BedrockMantleChatConfig(BedrockMantleAuthMixin, OpenAILikeChatConfig):
|
||||
"""
|
||||
Transformation config for Amazon Bedrock Mantle OpenAI-compatible API.
|
||||
"""
|
||||
|
||||
def __init__(self, aws_signer: BaseAWSLLM | None = None):
|
||||
super().__init__()
|
||||
self._aws_signer = aws_signer or BaseAWSLLM()
|
||||
|
||||
@property
|
||||
def custom_llm_provider(self) -> Optional[str]:
|
||||
return "bedrock_mantle"
|
||||
|
|
@ -40,6 +49,7 @@ class BedrockMantleChatConfig(OpenAILikeChatConfig):
|
|||
api_base: Optional[str],
|
||||
api_key: Optional[str],
|
||||
litellm_params: Optional[GenericLiteLLMParams] = None,
|
||||
model: str | None = None,
|
||||
) -> Tuple[Optional[str], Optional[str]]:
|
||||
region = (
|
||||
(litellm_params.aws_region_name if litellm_params else None)
|
||||
|
|
@ -49,12 +59,15 @@ class BedrockMantleChatConfig(OpenAILikeChatConfig):
|
|||
or BEDROCK_MANTLE_DEFAULT_REGION
|
||||
)
|
||||
BaseAWSLLM._validate_aws_region_name(region)
|
||||
# The base path segment is data-driven per model (use_openai_responses_path
|
||||
# flag): gemma-4-* and gpt-5.x are served on /openai/v1, everything else on
|
||||
# /v1. An explicit api_base still wins over the derived default.
|
||||
api_base = (
|
||||
api_base
|
||||
or get_secret_str("BEDROCK_MANTLE_API_BASE")
|
||||
or f"https://bedrock-mantle.{region}.api.aws/v1"
|
||||
or f"https://bedrock-mantle.{region}.api.aws/{mantle_base_segment(model, litellm.model_cost)}"
|
||||
)
|
||||
dynamic_api_key = api_key or get_secret_str("BEDROCK_MANTLE_API_KEY")
|
||||
dynamic_api_key = self._resolve_bearer_token(api_key)
|
||||
return api_base, dynamic_api_key
|
||||
|
||||
def validate_environment(
|
||||
|
|
|
|||
Some files were not shown because too many files have changed in this diff Show more
Loading…
Add table
Reference in a new issue