diff --git a/.github/scripts/_agent_shin_actions.py b/.github/scripts/_agent_shin_actions.py deleted file mode 100644 index b3d1ff055b3..00000000000 --- a/.github/scripts/_agent_shin_actions.py +++ /dev/null @@ -1,50 +0,0 @@ -"""Dry-run wrapper(s) around Agent Shin GitHub mutations. - -The rollout scripts currently need only one mutation wrapped, so this module -exposes a single ``maybe_post_comment`` helper. It takes a ``dry_run: bool`` -keyword argument and the body is intentionally trivial: - - if dry_run: - print(...) # log what we would do, return - return - real_mutation(...) # otherwise, actually do it - -That shape means a dry-run preview differs from the real run in exactly one -line per side effect: the call site. So when you `python3 script.py` locally -without ``--close``, you can be confident the actions printed are the ones the -GitHub Action would have performed (modulo ordering on retry/error paths, -which are deliberately simple). Any further mutation a rollout script needs -should get the same ``maybe_*`` treatment instead of calling the raw -``triage_with_llm`` mutation directly. - -Importing from this module pulls in the real mutation from ``triage_with_llm`` -— call sites in the rollout scripts should NEVER import ``post_comment`` -directly; that would skip the dry-run gate and is the bug class this module -exists to prevent. -""" - -from __future__ import annotations - -import sys -import textwrap - -# Import the module itself rather than the bare names so monkeypatching -# `triage_with_llm.post_comment` (or any of the other mutations) in tests is -# reflected here — `from triage_with_llm import post_comment` would bind the -# original function to a local name and bypass the patch, defeating the whole -# point of these wrappers. -import triage_with_llm - - -def _log(line: str) -> None: - """Print a single dry-run line to stdout (one log statement per side effect).""" - print(line, file=sys.stdout, flush=True) - - -def maybe_post_comment(repo: str, number: int, body: str, *, dry_run: bool) -> None: - """Post a comment on ``repo#number`` — or, in dry-run, log what we would post.""" - if dry_run: - _log(f"[DRY RUN] comment {repo}#{number}:") - _log(textwrap.indent(body, " ")) - return - triage_with_llm.post_comment(repo, number, body) diff --git a/.github/scripts/agent_shin_shared.py b/.github/scripts/agent_shin_shared.py deleted file mode 100644 index 8f3dc3c2322..00000000000 --- a/.github/scripts/agent_shin_shared.py +++ /dev/null @@ -1,211 +0,0 @@ -"""Constants and helpers shared by Agent Shin's triage scripts. - -Both `triage_with_llm.py` (the LLM-judge entrypoint) and -`close_low_quality_prs.py` (the daily Greptile-score sweep) need to -agree on the same notions of: - - * What counts as a Greptile-authored review comment - (``GREPTILE_BOT_LOGINS``) and how to extract a confidence score from - its body (``SCORE_PATTERN`` / :func:`extract_greptile_score`). - * How long the 2-hour grace window is (``GRACE_PERIOD_SECONDS``) and - the HTML marker stamped into a grace-warning comment so the *other* - script can see "Agent Shin already warned" and behave accordingly - (``GRACE_COMMENT_MARKER``). - * Who Agent Shin is on GitHub (``AGENT_SHIN_DEFAULT_BOT_LOGIN``). - * How GitHub-style ISO-8601 timestamps round-trip into timezone-aware - :class:`datetime.datetime` (:func:`parse_iso8601`). - -Keeping these in one module means a future change (new Greptile output -format, a longer grace window, a new allowlisted account) is a single edit -instead of two — the original split version had to call out in comments -that the two copies "must stay in sync" precisely because nothing -enforced it. -""" - -from __future__ import annotations - -import datetime as dt -import json -import os -import re -import subprocess -from typing import Iterable - -GREPTILE_BOT_LOGINS = frozenset({"greptile-apps", "greptile-apps[bot]"}) - -SCORE_PATTERN = re.compile( - r"confidence\s*score\s*[:\-]?\s*(\d+)\s*/\s*5", - re.IGNORECASE, -) - -GRACE_COMMENT_MARKER = "" - -# 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 = "" - -# 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 ``, 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() diff --git a/.github/scripts/close_low_quality_prs.py b/.github/scripts/close_low_quality_prs.py deleted file mode 100644 index 7b9bbb579e3..00000000000 --- a/.github/scripts/close_low_quality_prs.py +++ /dev/null @@ -1,573 +0,0 @@ -#!/usr/bin/env python3 -""" -Auto-close low-quality pull requests. - -Closes open PRs (including drafts, regardless of age) that satisfy ALL of: - 1. Have a Greptile (`greptile-apps`) review comment whose latest - "Confidence Score: X/5" is below the configured threshold (default: 4). - 2. Are authored by an external OSS contributor (internal BerriAI - contributors are exempt). - 3. Do not carry an opt-out label (default: "do not close"). - -`--min-age-days` is retained as an opt-in safety net for one-off backfill -runs (default: 0). The team's intent is that the count of open PRs equals -the count of PRs internal collaborators need to action on, so neither age -nor draft status acts as a free pass. - -For each match, the script posts an explanatory comment and closes the PR. -Because OSS contributors *cannot* reopen a PR closed by the bot/maintainer -(GitHub limitation), the close-comment instructs them to push their fixes -and **open a fresh PR**, or to comment `@agent-shin reconsider` on the -closed PR to have the LLM judge re-evaluate (and reopen on pass). - -Requires the `gh` CLI to be authenticated. - -Usage examples: - # Dry run (default) - prints what would be closed - python3 close_low_quality_prs.py - - # Actually close matching PRs - python3 close_low_quality_prs.py --close - - # Restrict to PRs at least N days old (one-off backfill safety net) - python3 close_low_quality_prs.py --min-age-days 7 --min-score 4 --close -""" - -from __future__ import annotations - -import argparse -import datetime as dt -import json -import os -import subprocess -import sys -from typing import Iterable - -# Add this script's directory to `sys.path` so the sibling -# `agent_shin_shared` module is importable when the script is invoked -# directly (e.g. `python3 .github/scripts/close_low_quality_prs.py ...`). -sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) - -from agent_shin_shared import ( # noqa: E402 -- sys.path adjusted above - AGENT_SHIN_CLOSE_MARKER, - ALLOWLIST_LOGINS, - GRACE_COMMENT_MARKER, - GRACE_PERIOD_SECONDS, - GREPTILE_BOT_LOGINS, - SCORE_PATTERN, - extract_greptile_score, - gh, - list_open_items, - parse_iso8601, - seconds_since_latest_marker_comment, -) - -# `GREPTILE_BOT_LOGINS` and `SCORE_PATTERN` (Greptile's GitHub App login -# variants and the "Confidence Score: X/5" regex) are imported from -# `agent_shin_shared` so the LLM judge in `triage_with_llm.py` and this -# daily Greptile sweep read the score through the same set of logins -# and the same regex. - -# `author_association` values for internal BerriAI contributors who should be -# exempt from auto-triage. -INTERNAL_AUTHOR_ASSOCIATIONS = frozenset({"OWNER", "MEMBER", "COLLABORATOR"}) - -# Default labels that exempt a PR from auto-close. Defined at module scope (not -# as a mutable argparse default) so that `--optout-label foo` REPLACES the -# defaults instead of appending to them — the argparse `action="append"` + -# `default=[...]` combination silently mutates the shared default list. -DEFAULT_OPTOUT_LABELS = ("do not close", "keep open", "wip") - -# `GRACE_COMMENT_MARKER` (HTML marker appended to grace-period warning -# comments — used by either script to recognize that a warning was -# already posted) and `GRACE_PERIOD_SECONDS` (length of the grace -# period between the warning and the actual auto-close, 2 hours) are -# imported from `agent_shin_shared` so the Agent Shin LLM judge and -# this daily Greptile sweep agree on the same marker and duration. - - -def fetch_open_prs(repo: str | None) -> list[dict]: - """Fetch all open PRs (number, createdAt, isDraft, labels, author). - - Includes drafts: `gh pr list --state open` returns both ready-for-review - and draft PRs by default. This is the desired behavior — drafts are not - a free pass; the internal-collaborator open-PR queue should reflect every - PR that needs human attention regardless of draft status. - """ - fields = "number,title,createdAt,isDraft,labels,author,url" - return list_open_items("pr", repo=repo, fields=fields) - - -def fetch_pr_author_association(pr_number: int, repo: str | None) -> str: - """Return the GitHub `author_association` for a PR, uppercase. - - Values: OWNER, MEMBER, COLLABORATOR, CONTRIBUTOR, FIRST_TIME_CONTRIBUTOR, - FIRST_TIMER, MANNEQUIN, NONE. Returns "" on lookup failure. - """ - endpoint = ( - f"repos/{repo}/pulls/{pr_number}" - if repo - else f"repos/{{owner}}/{{repo}}/pulls/{pr_number}" - ) - try: - data = json.loads(gh("api", endpoint)) - except subprocess.CalledProcessError: - return "" - return (data.get("author_association") or "").upper() - - -def is_external_pr_author(pr: dict, repo: str | None) -> bool: - """Return True if the PR author is an external OSS contributor. - - Internal = `OWNER` / `MEMBER` / `COLLABORATOR` association, or a bot login. - """ - login = ((pr.get("author") or {}).get("login") or "").lower() - if login.endswith("[bot]") or login in {"dependabot", "github-actions"}: - return False - association = fetch_pr_author_association(pr["number"], repo) - # Fail-safe: if the API lookup failed (empty string), treat the author as - # internal so we don't auto-close their PR. Auto-close is destructive, so - # an unknown association should never make a PR eligible for closing. - if not association or association in INTERNAL_AUTHOR_ASSOCIATIONS: - return False - return True - - -def fetch_pr_comments(pr_number: int, repo: str | None) -> list[dict]: - """Fetch issue-level comments on a PR (where Greptile posts its summary).""" - endpoint = ( - f"repos/{repo}/issues/{pr_number}/comments?per_page=100" - if repo - else f"repos/{{owner}}/{{repo}}/issues/{pr_number}/comments?per_page=100" - ) - raw = gh("api", "--paginate", endpoint) - comments: list[dict] = [] - for line in raw.strip().splitlines(): - line = line.strip() - if not line: - continue - try: - parsed = json.loads(line) - except json.JSONDecodeError: - # A malformed line should not blow up the whole sweep. Skip and - # carry on so the remaining PRs in this run still get evaluated. - continue - if isinstance(parsed, list): - comments.extend(parsed) - else: - comments.append(parsed) - return comments - - -def has_optout_label(pr: dict, optout_labels: set[str]) -> bool: - labels = {label.get("name", "").lower() for label in pr.get("labels", [])} - return bool(labels & {lbl.lower() for lbl in optout_labels}) - - -def seconds_since_last_grace_warning( - comments: Iterable[dict], - *, - bot_login: str | None = None, - now: dt.datetime | None = None, -) -> float | None: - """Return seconds since the bot's most recent grace-period warning, or - None if no such warning has ever been posted on this PR. - - Thin wrapper over - `agent_shin_shared.seconds_since_latest_marker_comment` — the - centralized helper handles the bot-author filter, marker match, - timestamp parsing, and `now` injection. Keeping this wrapper - preserves the closer's "already-fetched comments + injectable now" - interface so callers (and tests) don't need to change. - """ - return seconds_since_latest_marker_comment( - comments, - marker=GRACE_COMMENT_MARKER, - bot_login=bot_login, - now=now, - ) - - -def format_grace_warning_comment(score: int, threshold: int) -> str: - """Comment posted on the FIRST low-Greptile-score detection — gives - the contributor a 2-hour grace window before the auto-close fires on - the next daily cron run. - - Mirrors `format_grace_warning_pr_comment` in - `triage_with_llm.py` in spirit (2-hour grace + escape hatches), but - framed around Greptile's confidence score instead of the LLM judge's - rubric since the close trigger here is the Greptile signal. - """ - return ( - "🚅 Hi, thanks for the PR! I'm **Agent Shin**, the automated triage bot for this " - "repository.\n" - "\n" - "Heads up: Greptile's most recent review scored this PR " - f"**{score}/5**, below our merge bar of **{threshold}/5**.\n" - "\n" - "If the score isn't lifted in the next **2 hours**, I'll auto-close this PR. That's " - "**not** us saying the change isn't worthwhile. We want the open-PR list to mirror " - "what a maintainer can act on *right now*, so contributors like you don't get lost in " - "a backlog. Take your time; everything below still works after the close.\n" - "\n" - "**During the grace period:** push fixes that address Greptile's feedback, then comment " - "`@greptileai` to request a fresh review. If " - f"the new score is **{threshold}/5 or higher**, the PR stays open and no further " - "action is needed on your side.\n" - "\n" - "**If the PR does get auto-closed in 2 hours, you still have an easy recovery path:**\n" - "\n" - "- Comment `@greptileai` to request a fresh review. **This still works even after " - f"the PR is closed**, and a score of {threshold}/5 or higher is one of the signals " - "that lifts the PR back into the review queue. A low Greptile score isn't a blocker.\n" - "- Comment `@agent-shin reconsider` after pushing fixes; I'll re-run the rubric and " - "reopen the PR if both gates (description rubric + Greptile score) now pass.\n" - "\n" - f"{GRACE_COMMENT_MARKER}" - ) - - -def post_grace_warning( - pr: dict, - score: int, - threshold: int, - repo: str | None, - dry_run: bool, -) -> None: - """Post the 2-hour grace-period warning comment on `pr`. - - The warning carries `GRACE_COMMENT_MARKER` so subsequent runs can - detect that the contributor has already been told about the - pending close. Does NOT close the PR — the close happens on the - next eligible run after `GRACE_PERIOD_SECONDS` elapses (handled - by `close_pr`). - """ - pr_number = pr["number"] - repo_args = ["--repo", repo] if repo else [] - - if dry_run: - print( - f" [DRY RUN] Would post grace warning to PR #{pr_number} " - f"(greptile={score}/5): {pr['title']}" - ) - return - - comment_body = format_grace_warning_comment(score, threshold) - gh("pr", "comment", str(pr_number), "--body", comment_body, *repo_args) - print(f" Posted grace warning on PR #{pr_number} (greptile={score}/5)") - - -def format_close_comment(score: int, threshold: int) -> str: - """Comment posted when a low-Greptile-score PR is auto-closed. - - Carries `AGENT_SHIN_CLOSE_MARKER` so the `@agent-shin reconsider` path - (guarded by `was_closed_by_agent_shin`) recognizes this as an Agent Shin - close and is allowed to reopen the PR once it passes again; without the - marker that recovery path the comment advertises silently rejects the - contributor. - """ - score_sentence = ( - f"Greptile's most recent review scored this PR **{score}/5**, below " - f"our merge bar of **{threshold}/5**, and the 2-hour grace period since " - "the warning has elapsed.\n\n" - ) - return ( - f"Closing as part of automated PR triage.\n\n" - f"{score_sentence}" - "We close low-confidence PRs aggressively to keep the review queue " - "manageable for maintainers and contributors alike. **This is not a " - "rejection of the idea.** To bring this back:\n\n" - "1. Push the fixes that address Greptile's feedback (continue using " - "your existing branch is fine).\n" - "2. **Open a new PR** with the updated branch. Greptile will review " - "it again, and if it scores " - f"**{threshold}/5 or higher** a maintainer will take another look.\n\n" - "_Why open a new PR instead of reopening this one?_ GitHub does not " - "let external contributors reopen a PR that was closed by a bot or " - "maintainer, so a fresh PR is the most reliable path forward. If you " - "would prefer this exact PR re-evaluated, comment " - "`@agent-shin reconsider` once you've pushed the fixes; Agent Shin " - "will re-run triage and reopen this PR if it now meets the bar. " - "You can also comment `@greptileai` to request a fresh Greptile " - "review; that works **even after the PR is closed**.\n\n" - "Thanks for contributing to LiteLLM. We know auto-closures can sting; " - "the goal is to keep the project healthy, not to dismiss your work." - f"\n\n{AGENT_SHIN_CLOSE_MARKER}" - ) - - -def close_pr( - pr: dict, - score: int, - threshold: int, - age_days: int, - repo: str | None, - dry_run: bool, - label: str | None, -) -> None: - """Post the explanatory comment and close the PR.""" - pr_number = pr["number"] - repo_args = ["--repo", repo] if repo else [] - - if dry_run: - print( - f" [DRY RUN] Would close PR #{pr_number} " - f"(age={age_days}d, greptile={score}/5): {pr['title']}" - ) - return - - comment_body = format_close_comment(score, threshold) - gh("pr", "comment", str(pr_number), "--body", comment_body, *repo_args) - - if label: - try: - gh("pr", "edit", str(pr_number), "--add-label", label, *repo_args) - except subprocess.CalledProcessError as exc: - stderr = (exc.stderr or "").strip() - print(f" warn: failed to add label '{label}' to #{pr_number}: {stderr}") - - gh("pr", "close", str(pr_number), *repo_args) - print(f" Closed PR #{pr_number} (greptile={score}/5, age={age_days}d)") - - -def evaluate_pr( - pr: dict, - now: dt.datetime, - min_age_days: int, - min_score: int, - repo: str | None, - optout_labels: set[str], - allowlist: frozenset[str] = ALLOWLIST_LOGINS, -) -> tuple[str, int | None, int | None]: - """Decide what to do with `pr` on this triage run. - - Returns (action, score_or_none, age_days_or_none) where action is one of: - "skip-too-young", "skip-optout-label", "skip-not-allowlisted", - "skip-internal", "skip-no-greptile-score", "skip-score-ok", - "warn-grace", "skip-in-grace-period", or "close". - - Drafts are NOT skipped — the goal is "open PR count == PRs internal - collaborators need to action on", and a draft that Greptile scored <4/5 - is still in that queue. Authors can opt out via the `wip` label (see - `DEFAULT_OPTOUT_LABELS`) if they need to keep a long-lived draft open. - - Grace-period semantics: the first time a PR fails the rubric, the - action is `warn-grace` — the caller should post a warning comment but - NOT close the PR. On a subsequent run, if the warning is still less - than `GRACE_PERIOD_SECONDS` old AND the PR still fails, the action is - `skip-in-grace-period`. Once the warning ages out and the rubric is - still failing, the action is `close`. - """ - if has_optout_label(pr, optout_labels): - return ("skip-optout-label", None, None) - - created = parse_iso8601(pr["createdAt"]) - age_days = (now - created).days - # `min_age_days` defaults to 0 (close as soon as Greptile scores low). - # Set a positive value via --min-age-days for one-off backfill runs that - # want to skip very-young PRs. - if min_age_days > 0 and age_days < min_age_days: - return ("skip-too-young", None, age_days) - - # While the allowlist is active it is the sole author gate: only those - # logins are acted on and the external-only restriction is bypassed for - # them. Otherwise auto-close only external OSS contributors — internal - # contributors (BerriAI org members) handle their own backlog. - login = ((pr.get("author") or {}).get("login") or "").lower() - if allowlist: - if login not in allowlist: - return ("skip-not-allowlisted", None, age_days) - elif not is_external_pr_author(pr, repo): - return ("skip-internal", None, age_days) - - comments = fetch_pr_comments(pr["number"], repo) - extraction = extract_greptile_score(comments) - if extraction is None: - return ("skip-no-greptile-score", None, age_days) - - score, _ = extraction - if score >= min_score: - return ("skip-score-ok", score, age_days) - - grace_age = seconds_since_last_grace_warning(comments, now=now) - if grace_age is None: - return ("warn-grace", score, age_days) - if grace_age < GRACE_PERIOD_SECONDS: - return ("skip-in-grace-period", score, age_days) - - return ("close", score, age_days) - - -def main() -> int: - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument( - "--repo", - type=str, - default=None, - help="Repository (owner/repo). Auto-detected if omitted.", - ) - parser.add_argument( - "--min-age-days", - type=int, - default=0, - help=( - "Minimum age (in days) before a PR is eligible. Default 0 = " - "close as soon as Greptile flags it. Set a positive value for " - "one-off backfill runs that want to spare very-young PRs." - ), - ) - parser.add_argument( - "--min-score", - type=int, - default=4, - choices=range(1, 6), - help="Greptile score below which a PR is closed (default: 4 -> closes <4/5).", - ) - parser.add_argument( - "--optout-label", - action="append", - default=None, - help=( - "Label(s) that exempt a PR from auto-close. Repeat to add more. " - "Case-insensitive. When omitted, defaults to " - f"{list(DEFAULT_OPTOUT_LABELS)!r}; passing this flag REPLACES the " - "defaults (argparse `append` with a mutable default would append " - "instead, which we explicitly avoid)." - ), - ) - parser.add_argument( - "--close-label", - type=str, - default=None, - help=( - "Optional label to add to PRs that get auto-closed " - "(e.g. 'auto-closed-low-quality'). Must already exist on the repo." - ), - ) - parser.add_argument( - "--close", - action="store_true", - help="Actually close matching PRs (default is dry-run).", - ) - parser.add_argument( - "--limit", - type=int, - default=None, - help="Maximum number of PRs to close in one run (safety net).", - ) - args = parser.parse_args() - - dry_run = not args.close - if dry_run: - print("=== DRY RUN MODE (pass --close to actually close PRs) ===\n") - - print("Fetching open PRs...") - prs = fetch_open_prs(args.repo) - print(f"Found {len(prs)} open PRs.\n") - - now = dt.datetime.now(dt.timezone.utc) - optout_labels = set(args.optout_label or DEFAULT_OPTOUT_LABELS) - - closed = 0 - summary = { - "close": 0, - "warn-grace": 0, - "skip-in-grace-period": 0, - "skip-too-young": 0, - "skip-optout-label": 0, - "skip-not-allowlisted": 0, - "skip-internal": 0, - "skip-no-greptile-score": 0, - "skip-score-ok": 0, - } - - # `warned` tracks grace-warning comments posted in this run so the - # `--limit` safety net bounds *all* destructive write actions, not - # just closures. Without this cap, a backlog of PRs failing the - # threshold simultaneously could flood contributors with comments. - warned = 0 - for pr in sorted(prs, key=lambda p: p["createdAt"]): - try: - action, score, age_days = evaluate_pr( - pr, - now, - args.min_age_days, - args.min_score, - args.repo, - optout_labels, - ) - summary[action] = summary.get(action, 0) + 1 - - if action == "warn-grace": - assert score is not None - print( - f"#{pr['number']}: \"{pr['title']}\" " - f"(age={age_days}d, greptile={score}/5) -> warn-grace" - ) - post_grace_warning( - pr, - score=score, - threshold=args.min_score, - repo=args.repo, - dry_run=dry_run, - ) - if not dry_run: - warned += 1 - if args.limit is not None and (warned + closed) >= args.limit: - print( - f"\nReached --limit={args.limit} " - f"(closed={closed}, warned={warned}); stopping." - ) - break - continue - - if action != "close": - continue - - assert score is not None and age_days is not None - print( - f"#{pr['number']}: \"{pr['title']}\" " - f"(age={age_days}d, greptile={score}/5) -> close" - ) - close_pr( - pr, - score=score, - threshold=args.min_score, - age_days=age_days, - repo=args.repo, - dry_run=dry_run, - label=args.close_label, - ) - - if not dry_run: - closed += 1 - if args.limit is not None and (warned + closed) >= args.limit: - print( - f"\nReached --limit={args.limit} " - f"(closed={closed}, warned={warned}); stopping." - ) - break - except Exception as exc: # noqa: BLE001 - per-PR errors don't abort the sweep - summary["error"] = summary.get("error", 0) + 1 - print( - f"!! PR #{pr.get('number')}: {exc}", - file=sys.stderr, - ) - continue - - print("\n=== Summary ===") - for key, value in summary.items(): - print(f" {key:28s} {value}") - if dry_run: - print(f"\nTotal would close: {summary['close']}") - else: - print(f"\nTotal closed: {closed}") - print( - f"Total {'would warn (grace)' if dry_run else 'warned (grace)'}: " - f"{summary['warn-grace']}" - ) - return 0 - - -if __name__ == "__main__": - sys.exit(main()) diff --git a/.github/scripts/triage-requirements.txt b/.github/scripts/triage-requirements.txt deleted file mode 100644 index a18f05fbb95..00000000000 --- a/.github/scripts/triage-requirements.txt +++ /dev/null @@ -1,282 +0,0 @@ -# Hash-pinned dependency set for the Agent Shin triage scripts. -# Installed in privileged triage workflows, so every package is pinned to an -# exact version with SHA-256 hashes and installed with pip --require-hashes. -# -# Regenerate after bumping openai: -# echo 'openai==' \ -# | 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 diff --git a/.github/scripts/triage_with_llm.py b/.github/scripts/triage_with_llm.py deleted file mode 100644 index e23a012425a..00000000000 --- a/.github/scripts/triage_with_llm.py +++ /dev/null @@ -1,1797 +0,0 @@ -#!/usr/bin/env python3 -""" -Agent Shin — LLM-as-judge triage for external OSS pull requests and issues. - -Evaluates a single PR or issue against the contribution rubric and, when the -LLM judge marks it as failing, posts an explanatory comment + closes the -PR/issue. Re-triggers on `reopened` so contributors can iterate back in by -filling in the missing pieces and reopening. - -Internal BerriAI contributors (`author_association` in {OWNER, MEMBER, -COLLABORATOR}) and bot accounts are skipped entirely. - -Usage: - triage_with_llm.py --repo owner/repo --pr 1234 - triage_with_llm.py --repo owner/repo --issue 5678 - triage_with_llm.py --repo owner/repo --pr 1234 --close # actually close - triage_with_llm.py --repo owner/repo --pr 1234 --print-prompt # show prompt - -Defaults are SAFE: without `--close` the script writes a verdict to stdout (and, -when running in GitHub Actions, to $GITHUB_STEP_SUMMARY) but takes no GitHub -write actions. - -Environment: - GH_TOKEN / GITHUB_TOKEN - for `gh` CLI auth (auto-set in Actions) - OPENAI_API_KEY - required when --close is passed - OPENAI_BASE_URL - optional (route to any OpenAI-compatible API) - TRIAGE_MODEL - optional model override (default: gpt-5.4-mini) -""" - -from __future__ import annotations - -import argparse -import datetime as dt -import json -import os -import re -import subprocess -import sys -import textwrap -import urllib.parse -from typing import Any, 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/triage_with_llm.py ...`) and -# also when the tests load this script via -# `importlib.util.spec_from_file_location`. -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, - AGENT_SHIN_DEFAULT_BOT_LOGIN, - ALLOWLIST_LOGINS, - GRACE_COMMENT_MARKER, - GRACE_PERIOD_SECONDS, - GREPTILE_BOT_LOGINS, - SCORE_PATTERN, - extract_greptile_score, - gh, - parse_iso8601, - seconds_since_latest_marker_comment, -) - -DEFAULT_MODEL = "gpt-5.4-mini" - -INTERNAL_ASSOCIATIONS = frozenset({"OWNER", "MEMBER", "COLLABORATOR"}) - -# `AGENT_SHIN_DEFAULT_BOT_LOGIN` is imported from `agent_shin_shared`. -# When the workflow uses the default `secrets.GITHUB_TOKEN`, the -# closure / reopen event's `actor.login` is `github-actions[bot]`. The -# env override `AGENT_SHIN_BOT_LOGIN` exists for local debugging and for -# repos that wire Agent Shin to a PAT. - -# HTML marker appended to every reconsider verdict comment. We grep for this -# on subsequent reconsider triggers to enforce a short cooldown so that -# repeated `@agent-shin reconsider` comments don't burn CI/LLM budget. -# Using a unique HTML comment keeps the marker invisible to humans while -# being trivially greppable from a comments-list API response. -RECONSIDER_COMMENT_MARKER = "" - -# Minimum gap between two reconsider verdicts on the same PR/issue. Set to -# 10 minutes — long enough that a contributor can't trivially spam the -# trigger, short enough that a genuine "I just pushed a fix and reupdated -# the body" iteration loop isn't punished. -RECONSIDER_RATE_LIMIT_SECONDS = 600 - -# `GRACE_COMMENT_MARKER` (HTML marker on the grace-period warning comment -# posted on the first low-quality detection — used on subsequent triage -# runs to detect that a warning was already posted and measure how long -# ago it was 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 daily Greptile sweep and the -# LLM judge agree on the same marker and duration. - -# --- Review-gate ("ready for review" label lifecycle) configuration ---------- -# The review gate keeps a single label in sync with whether a PR currently -# clears BOTH quality bars: the LLM rubric (clear problem + expected/actual + -# QA proof, or a linked issue) AND Greptile's most recent confidence score. -READY_FOR_REVIEW_LABEL = "ready for review" -DEFAULT_GRACE_DAYS = 1 # 24h before an un-passing, un-tagged PR is auto-closed -DEFAULT_MIN_GREPTILE_SCORE = 4 # Greptile < 4/5 counts as "not passing" - -# Hidden HTML-comment markers stamped into review-gate comments. They never -# render in the GitHub UI but let the gate detect its own prior actions so it -# (a) posts the within-grace "what's missing" notice at most once and (b) can -# tell a first-time pass ("ready for review") from a recovery after a -# regression ("all clear again"). -READY_MARKER = "" -REGRESSED_MARKER = "" -WITHIN_GRACE_MARKER = "" - -# `GREPTILE_BOT_LOGINS` (Greptile's GitHub App login variants — -# `greptile-apps[bot]` in REST API comments, `greptile-apps` in -# `gh pr view --json` output) and `SCORE_PATTERN` (regex matching lines -# like `Confidence Score: 3/5`) are imported from `agent_shin_shared` -# so the daily sweep and the review gate read the score through the -# same set of logins / patterns. - -# `AGENT_SHIN_CLOSE_MARKER` is imported from `agent_shin_shared` so this LLM -# judge and the daily Greptile sweep stamp the same marker on their close -# comments — `was_closed_by_agent_shin` keys the reconsider reopen path off it. - -# Model families that require `reasoning_effort` to be set, and that reject -# `temperature != 1` unless `reasoning_effort` is "none". For these models we -# pass `reasoning_effort="none"` so a `temperature=0` deterministic judgment -# is still accepted. See litellm/llms/openai/chat/gpt_5_transformation.py for -# the full set of constraints LiteLLM applies to these models. -GPT5_FAMILY_PREFIX = "gpt-5" - -# Regexes for picking off "obvious passes" without burning LLM tokens. -# -# Keep this list to GitHub's documented PR-closing keywords only -# (https://docs.github.com/issues/tracking-your-work-with-issues/linking-a-pull-request-to-an-issue). -# Casual mentions like "see #1234" or "ref #1234" are intentionally NOT -# auto-passed — they should fall through to the LLM judge, which has the -# stricter rubric "a bare issue number without a closing keyword counts only -# if it's clearly the related issue (not a passing mention)". -LINKED_ISSUE_PATTERN = re.compile( - r"\b(?:fixes|fix|fixed|closes|close|closed|resolves|resolve|resolved)\s+" - r"(?:#\d+|https?://github\.com/[\w.-]+/[\w.-]+/issues/\d+)", - re.IGNORECASE, -) -HTML_COMMENT_PATTERN = re.compile(r"", re.DOTALL) - - -# --------------------------------------------------------------------------- -# gh helpers -# -# `gh` is imported from `agent_shin_shared` so a future change (timeout, -# logging, retry) only needs to be made once. - - -def fetch_pr(repo: str, number: int) -> dict: - """Return the full GitHub REST representation of a PR.""" - return json.loads(gh("api", f"repos/{repo}/pulls/{number}")) - - -def fetch_issue(repo: str, number: int) -> dict: - """Return the full GitHub REST representation of an issue.""" - return json.loads(gh("api", f"repos/{repo}/issues/{number}")) - - -def post_comment(repo: str, number: int, body: str) -> None: - """Post an issue-style comment (works for both issues and PRs).""" - gh( - "api", - f"repos/{repo}/issues/{number}/comments", - "-X", - "POST", - "-f", - f"body={body}", - ) - - -def close_pr(repo: str, number: int) -> None: - """Close a pull request (state=closed).""" - gh( - "api", - f"repos/{repo}/pulls/{number}", - "-X", - "PATCH", - "-f", - "state=closed", - ) - - -def reopen_pr(repo: str, number: int) -> None: - """Reopen a previously-closed pull request (state=open). - - Used by the `@agent-shin reconsider` comment-trigger flow: the bot has - write access via GH_TOKEN, so it can reopen on the contributor's behalf - even though GitHub doesn't let the OSS author do it themselves. - """ - gh( - "api", - f"repos/{repo}/pulls/{number}", - "-X", - "PATCH", - "-f", - "state=open", - ) - - -def close_issue(repo: str, number: int, *, not_planned: bool = True) -> None: - """Close an issue, marking state_reason=not_planned by default.""" - args = [ - "api", - f"repos/{repo}/issues/{number}", - "-X", - "PATCH", - "-f", - "state=closed", - ] - if not_planned: - args.extend(["-f", "state_reason=not_planned"]) - gh(*args) - - -def reopen_issue(repo: str, number: int) -> None: - """Reopen a previously-closed issue (state=open, state_reason=reopened).""" - gh( - "api", - f"repos/{repo}/issues/{number}", - "-X", - "PATCH", - "-f", - "state=open", - "-f", - "state_reason=reopened", - ) - - -def add_label(repo: str, number: int, label: str) -> None: - """Add a label to a PR/issue (GitHub creates the label if it's missing).""" - gh( - "api", - f"repos/{repo}/issues/{number}/labels", - "-X", - "POST", - "-f", - f"labels[]={label}", - ) - - -def remove_label(repo: str, number: int, label: str) -> None: - """Remove a label from a PR/issue. A missing label (404) is not an error.""" - encoded = urllib.parse.quote(label, safe="") - try: - gh( - "api", - f"repos/{repo}/issues/{number}/labels/{encoded}", - "-X", - "DELETE", - ) - except subprocess.CalledProcessError as exc: - stderr = (exc.stderr or "").lower() - if "404" in stderr or "not found" in stderr: - return - raise - - -def _iter_paginated_json(*api_args: str) -> Any: - """Yield JSON objects from `gh api --paginate ... -q '.[]'`. - - `gh api --paginate` on a JSON-array endpoint concatenates pages into - one stream; `-q '.[]'` flattens that stream into newline-delimited - objects (jq-style). This keeps memory bounded for chatty endpoints - like issue events/comments on long-lived PRs. - """ - raw = gh("api", "--paginate", *api_args, "-q", ".[]") - for line in raw.splitlines(): - line = line.strip() - if not line: - continue - try: - yield json.loads(line) - except json.JSONDecodeError: - # A malformed line should not blow up the whole guard. Skip and - # carry on — at worst the guard fail-closes (returns False / - # None) and the caller treats it as "unknown". - continue - - -def fetch_last_close_event( - repo: str, number: int -) -> tuple[str | None, dt.datetime | None]: - """Return the actor login and timestamp of the most recent `closed` event. - - Either field may be None: actor when the events API returns nothing - (unusual for a closed item, but possible on transient errors), and - timestamp when the event lacks `created_at` or the value can't be - parsed. `was_closed_by_agent_shin` fail-closes on either. - """ - actor: str | None = None - closed_at: dt.datetime | None = None - for event in _iter_paginated_json(f"repos/{repo}/issues/{number}/events"): - if event.get("event") != "closed": - continue - actor = (event.get("actor") or {}).get("login") - created = event.get("created_at") - if not created: - closed_at = None - continue - try: - closed_at = parse_iso8601(created) - except ValueError: - closed_at = None - return actor, closed_at - - -# How much older than the latest `closed` event the Agent Shin marker -# comment is allowed to be while still counting as "this close was Agent -# Shin's". Agent Shin posts the close comment immediately before closing, -# so the marker timestamp is normally at most a few seconds before the -# close event; the buffer just absorbs clock skew between the comments -# API and the events API. -AGENT_SHIN_CLOSE_MARKER_SKEW_SECONDS = 300 - - -def was_closed_by_agent_shin( - repo: str, number: int, *, bot_login: str | None = None -) -> bool: - """Return True iff Agent Shin itself most-recently closed this PR/issue. - - This is the guard that stops `@agent-shin reconsider` from reopening an - item Agent Shin did not close — a maintainer closing for non-rubric - reasons (security, duplicate, design rejection), or a different workflow - (stale/duplicate sweeps) closing under the shared `github-actions[bot]` - identity. Three independent signals must all hold, because that identity - is not unique to Agent Shin and a marker comment from a prior - closed/reopened cycle would otherwise vouch for an unrelated close: - - 1. The most recent `closed` event's actor is the bot identity. - 2. Agent Shin left one of its auto-close comments, detected via - `AGENT_SHIN_CLOSE_MARKER`. The actor check alone can't tell an - Agent Shin close from any other `github-actions[bot]` close. - 3. That marker comment was posted at (or just before) the latest - close event, not on a previous close in an - Agent-Shin-close -> reconsider-reopen -> other-bot-reclose cycle. - - The check is intentionally fail-closed: any uncertainty about who closed - the item is treated as "not Agent Shin" so the destructive reopen path - stays gated. - """ - expected = ( - bot_login - or os.environ.get("AGENT_SHIN_BOT_LOGIN") - or AGENT_SHIN_DEFAULT_BOT_LOGIN - ).lower() - actor, closed_at = fetch_last_close_event(repo, number) - if not actor or actor.lower() != expected or closed_at is None: - return False - marker_seconds = seconds_since_last_agent_shin_close( - repo, number, bot_login=bot_login - ) - if marker_seconds is None: - return False - close_age_seconds = (dt.datetime.now(dt.timezone.utc) - closed_at).total_seconds() - return marker_seconds <= close_age_seconds + AGENT_SHIN_CLOSE_MARKER_SKEW_SECONDS - - -def _seconds_since_latest_marker_comment( - repo: str, - number: int, - *, - marker: str, - bot_login: str | None = None, -) -> float | None: - """Return seconds since the bot's most recent comment with ``marker``. - - Fetches comments via `_iter_paginated_json` and delegates the - iteration / author-filter / timestamp logic to - `agent_shin_shared.seconds_since_latest_marker_comment` so the daily - Greptile sweep and the LLM judge use one source of truth for the - "bot already posted X" detection. The wall-clock `now` is resolved - against this module's `dt` so tests that freeze time via - `monkeypatch.setattr(triage_module, "dt", ...)` still apply. - """ - return seconds_since_latest_marker_comment( - _iter_paginated_json(f"repos/{repo}/issues/{number}/comments"), - marker=marker, - bot_login=bot_login, - now=dt.datetime.now(dt.timezone.utc), - ) - - -def seconds_since_last_reconsider_verdict( - repo: str, number: int, *, bot_login: str | None = None -) -> float | None: - """Return seconds since the bot's most recent reconsider verdict comment. - - Detects comments by matching the HTML marker `RECONSIDER_COMMENT_MARKER` - appended by `format_reopen_comment` and - `format_reconsider_still_failing_comment`. Returns None when the bot - has never posted a reconsider verdict on this PR/issue (or when the - only matching comments are missing a `created_at` timestamp, which - shouldn't happen on a real GitHub response). - """ - return _seconds_since_latest_marker_comment( - repo, number, marker=RECONSIDER_COMMENT_MARKER, bot_login=bot_login - ) - - -def seconds_since_last_grace_warning( - repo: str, number: int, *, bot_login: str | None = None -) -> float | None: - """Return seconds since the bot's most recent grace-period warning. - - Detects warning comments by matching the HTML marker - `GRACE_COMMENT_MARKER` appended by `format_grace_warning_pr_comment` - and `format_grace_warning_issue_comment`. Returns None when no - grace warning has ever been posted on this PR/issue — that's the - "first low-quality detection" signal that drives the warning path. - """ - return _seconds_since_latest_marker_comment( - repo, number, marker=GRACE_COMMENT_MARKER, bot_login=bot_login - ) - - -def seconds_since_last_agent_shin_close( - repo: str, number: int, *, bot_login: str | None = None -) -> float | None: - """Return seconds since Agent Shin's most recent auto-close comment. - - Detects close comments by matching `AGENT_SHIN_CLOSE_MARKER` (stamped by - `format_pr_close_comment` / `format_issue_close_comment`). Returns None - when Agent Shin has never closed this PR/issue — the signal - `was_closed_by_agent_shin` uses to keep the reconsider reopen path gated - against closures performed by other workflows sharing the bot identity. - """ - return _seconds_since_latest_marker_comment( - repo, number, marker=AGENT_SHIN_CLOSE_MARKER, bot_login=bot_login - ) - - -# --------------------------------------------------------------------------- -# Author classification - - -def is_internal_contributor(item: dict) -> bool: - """Return True if the PR/issue author should be exempted from triage. - - Fail-safe: if `author_association` is missing or empty (which should never - happen on a successful GitHub REST response but is possible on schema - changes or partial responses), treat the author as INTERNAL so the - destructive close path never fires on an unknown contributor. This matches - the sibling `is_external_pr_author` in `close_low_quality_prs.py`. - """ - login = ((item.get("user") or {}).get("login") or "").lower() - if login.endswith("[bot]") or login in {"dependabot", "github-actions"}: - return True - association = (item.get("author_association") or "").upper() - if not association or association in INTERNAL_ASSOCIATIONS: - return True - return False - - -# --------------------------------------------------------------------------- -# Greptile score + age helpers (`extract_greptile_score`, `parse_iso8601`) -# live in `agent_shin_shared` — they're imported at the top of this module -# so both `triage_with_llm.py` and `close_low_quality_prs.py` share a -# single source of truth for the Confidence-Score regex and ISO-8601 -# parsing. - - -# --------------------------------------------------------------------------- -# Prompt construction - - -def strip_html_comments(text: str) -> str: - """Remove HTML comments — template placeholder text shouldn't fool the judge.""" - return HTML_COMMENT_PATTERN.sub("", text or "") - - -def has_linked_issue(text: str) -> bool: - """Heuristic: does this body link to an open issue (Fixes #123 etc.)?""" - return bool(LINKED_ISSUE_PATTERN.search(strip_html_comments(text or ""))) - - -def build_pr_prompt(*, title: str, body: str) -> str: - cleaned_body = strip_html_comments(body or "").strip() or "(empty)" - # Dedent the static template *before* interpolating dynamic fields so that - # multi-line bodies (whose 2nd+ lines start at column 0) don't defeat the - # common-indent computation in textwrap.dedent. - template = textwrap.dedent(""" - You are "Agent Shin", the OSS triage bot for the LiteLLM open-source - repository (BerriAI/litellm). Decide whether this external pull request - meets the project's contribution standards. - - A PR PASSES triage only if BOTH (1) AND (2) are satisfied. A linked - issue alone is NOT enough — it covers context, not proof. - - (1) CONTEXT — the PR provides AT LEAST ONE of: - (a) A link to a related GitHub issue. Acceptable forms: - "Fixes #1234", "Closes #1234", "Resolves #1234", - "Refs https://github.com/BerriAI/litellm/issues/1234". A - bare "#1234" without a closing keyword counts only if it - is clearly the related issue (not a passing mention). - (b) A clear problem description in the body (what bug or - missing feature this addresses, beyond the title) AND - expected vs. actual behavior (or, for features, "what's - possible now vs. with this PR"). - - (2) END-TO-END QA PROOF: the PR body contains AT LEAST ONE of: - (a) A screen recording / video showing the behavior before - and after the change (the bug reproducing, then the fix - working). For a brand-new feature with no meaningful - "before", a recording of it working end-to-end is fine. - (b) A screenshot (or before/after screenshots) showing the - fix or feature working. - (c) Specific commands that were actually run (curl, python, - a CLI invocation, etc.) PAIRED WITH their real - output, demonstrating the change works end-to-end against - the real system. Commands whose external dependencies - (LLM provider, DB, network) are mocked or stubbed do NOT - satisfy (2c); they are not end-to-end. - - `has_qa_proof` must be set to `true` only when (2a), (2b), - or a non-mocked (2c) is actually present in the body. If the - only "proof" is mocked tests, `has_qa_proof` is `false` and - the verdict is "fail". - - The following do NOT count as QA proof: - - Generic claims like "I tested it", "works locally", "all - tests pass", or a checked "I added tests" checkbox with no - output shown. - - A description of what tests exist or were added, without - their actual output in the PR body. - - `pytest` (or any test runner) executed against the - repository's own unit tests. Those mock the LLM provider, - DB, and network, so they are NOT end-to-end and never - satisfy (2), no matter how much passing output is pasted. - - A linked issue. The linked issue is context (1a), never - proof (2). - - FAIL the PR if EITHER (1) or (2) is missing. Do not bias toward PASS: - if QA proof is absent, the verdict is "fail" even when the rest of - the PR is well-written. - - Respond with a single JSON object, no prose: - - {{ - "verdict": "pass" | "fail", - "linked_issue": boolean, - "has_problem_description": boolean, - "has_expected_vs_actual": boolean, - "has_qa_proof": boolean, - "qa_proof_type": "video" | "screenshot" | "commands_with_output" | "none", - "missing": ["plain-english strings naming what is missing"], - "explanation": "1-2 sentence reasoning for the team to skim" - }} - - --- - PR title: {title} - - PR body: - --- - {cleaned_body} - --- - """).strip() - return template.format(title=title, cleaned_body=cleaned_body) - - -def build_issue_prompt(*, title: str, body: str) -> str: - cleaned_body = strip_html_comments(body or "").strip() or "(empty)" - # Dedent the static template *before* interpolating dynamic fields so that - # multi-line bodies (whose 2nd+ lines start at column 0) don't defeat the - # common-indent computation in textwrap.dedent. - template = textwrap.dedent(""" - You are "Agent Shin", the OSS triage bot for the LiteLLM open-source - repository (BerriAI/litellm). Decide whether this GitHub issue meets - the project's reporting standards. - - For a BUG REPORT the issue PASSES triage only when it contains BOTH: - (1) END-TO-END EVIDENCE OF THE BUG (the "before"; set - `has_repro=true` only when this is present): AT LEAST ONE of: - (a) A screen recording / video of the bug happening. - (b) A screenshot of the bug. - (c) The exact command(s) actually run (curl, python, a CLI - invocation, etc.) PAIRED WITH their real output, traceback, - or logs showing the failure against the real system. - Commands whose external dependencies (LLM provider, DB, - network) are mocked or stubbed do NOT count. - Prose-only "steps to reproduce" with no run output, video, or - screenshot do NOT satisfy (1). An unfilled template scaffold - (bare headings such as "Version or commit:" with nothing under - them, empty numbered lists) counts as absent, not as evidence. - (2) Expected vs. actual behavior (`has_expected_vs_actual`). - - FAIL the bug report if either (1) or (2) is missing. Do not bias - toward PASS: if the bug isn't demonstrated end-to-end, the verdict is - "fail" even when the report is well-written. - - For a FEATURE REQUEST the issue PASSES triage only when it contains - ALL of: - - A clear description of the proposed feature (what should LiteLLM do - that it does not today). - - Motivation / use case with a concrete example (config, API call, - UI flow, or scenario showing what's blocked today). - - END-TO-END EVIDENCE OF THE DEAD-END (set - `has_dead_end_evidence=true` only when this is present): a video, - a screenshot, or the exact command(s) actually run paired with - their real output, showing the point where the flow stops today. - Mocked or stubbed dependencies do NOT count, and an unfilled - template scaffold (bare headings, empty numbered lists) counts as - absent. - - For an issue that is neither a bug report nor a feature request (a - question, support request, or discussion), PASS as long as it has a - clear, specific ask and is not empty or template placeholder text. - - Respond with a single JSON object, no prose: - - {{ - "verdict": "pass" | "fail", - "kind": "bug" | "feature" | "other", - "has_repro": boolean, - "has_expected_vs_actual": boolean, - "has_motivation_example": boolean, - "has_dead_end_evidence": boolean, - "missing": ["plain-english strings naming what is missing"], - "explanation": "1-2 sentence reasoning for the team to skim" - }} - - --- - Issue title: {title} - - Issue body: - --- - {cleaned_body} - --- - """).strip() - return template.format(title=title, cleaned_body=cleaned_body) - - -# --------------------------------------------------------------------------- -# LLM call + verdict parsing - - -def call_llm_judge( - prompt: str, *, model: str, api_key: str, base_url: str | None -) -> str: - """Call an OpenAI-compatible chat completions endpoint. Returns raw text.""" - # Import inside the function so unit tests that monkey-patch this never - # need the openai package installed. - from openai import OpenAI - - client = ( - OpenAI(api_key=api_key, base_url=base_url) - if base_url - else OpenAI(api_key=api_key) - ) - kwargs: dict[str, Any] = { - "model": model, - "messages": [{"role": "user", "content": prompt}], - "temperature": 0, - "response_format": {"type": "json_object"}, - } - # gpt-5.x reasoning models reject `temperature != 1` unless - # `reasoning_effort` is explicitly "none". Set it via `extra_body` so this - # works across openai SDK versions regardless of whether the SDK natively - # types `reasoning_effort` as a top-level chat-completions param yet. - if model.lower().startswith(GPT5_FAMILY_PREFIX): - kwargs["extra_body"] = {"reasoning_effort": "none"} - response = client.chat.completions.create(**kwargs) - return response.choices[0].message.content or "" - - -def parse_verdict(raw: str) -> dict: - """Parse the LLM's JSON response. Tolerates ```json fences and stray text.""" - if not raw: - raise ValueError("empty LLM response") - text = raw.strip() - if text.startswith("```"): - text = re.sub(r"^```(?:json)?\s*", "", text) - text = re.sub(r"\s*```$", "", text) - try: - return json.loads(text) - except json.JSONDecodeError: - match = re.search(r"\{.*\}", text, re.DOTALL) - if not match: - raise ValueError(f"could not extract JSON from LLM response: {raw[:200]}") - return json.loads(match.group(0)) - - -# --------------------------------------------------------------------------- -# Comment composition - - -def _format_missing(missing: list[str]) -> str: - if not missing: - return "- (see explanation below)" - return "\n".join(f"- {m}" for m in missing) - - -# Rubric items the judge can mark present. The first element of each tuple is -# the verdict-JSON boolean field, the second is the human-readable label we -# render in the "what you got right" section of close / grace-warning comments. -_PR_PRESENT_LABELS: tuple[tuple[str, str], ...] = ( - ("linked_issue", "Linked a related GitHub issue"), - ("has_problem_description", "Clear problem description"), - ("has_expected_vs_actual", "Expected vs. actual behavior"), - ("has_qa_proof", "End-to-end QA proof"), -) - -# Issue rubric labels grouped by `kind`. The judge sets `kind` to one of -# {"bug", "feature", "other"}; when "other" we render both groups so we don't -# silently drop a present-flag the judge actually set to True. -_ISSUE_BUG_LABELS: tuple[tuple[str, str], ...] = ( - ( - "has_repro", - "End-to-end evidence of the bug (video, screenshot, or command + real output)", - ), - ("has_expected_vs_actual", "Expected vs. actual behavior"), -) -_ISSUE_FEATURE_LABELS: tuple[tuple[str, str], ...] = ( - ("has_motivation_example", "Motivation and concrete example"), - ( - "has_dead_end_evidence", - "End-to-end evidence of the dead-end (video, screenshot, or command + real output)", - ), -) - - -def _format_present_for_pr(verdict: dict) -> list[str]: - """Human-readable rubric items the judge confirmed are present on a PR. - - Drives the "what you got right" section in close / grace-warning comments. - The user gave explicit feedback: contributors should see what they nailed - *before* the list of gaps, so the comment doesn't read as pure rejection. - """ - return [label for field, label in _PR_PRESENT_LABELS if verdict.get(field)] - - -def _format_present_for_issue(verdict: dict) -> list[str]: - """Human-readable rubric items the judge confirmed are present on an issue. - - Branches on the judge's `kind` field. For `"other"` (or missing kind) we - render the union so a present-flag isn't dropped just because the judge - couldn't classify the issue cleanly. - """ - kind = (verdict.get("kind") or "").lower() - groups: list[tuple[tuple[str, str], ...]] = [] - if kind in ("bug", "other", ""): - groups.append(_ISSUE_BUG_LABELS) - if kind in ("feature", "other", ""): - groups.append(_ISSUE_FEATURE_LABELS) - out: list[str] = [] - for group in groups: - for field, label in group: - if verdict.get(field) and label not in out: - out.append(label) - return out - - -def _format_present_block(items: list[str]) -> str: - """Render the optional "what you got right" block. Empty string when the - judge didn't confirm anything as present — better to omit the section - entirely than to show "What you got right: (nothing)". - """ - if not items: - return "" - bullets = "\n".join(f"- ✅ {item}" for item in items) - return f"**What you got right:**\n\n{bullets}\n\n" - - -def format_pr_close_comment(verdict: dict) -> str: - missing_lines = _format_missing(verdict.get("missing") or []) - present_block = _format_present_block(_format_present_for_pr(verdict)) - explanation = verdict.get("explanation") or "" - return ( - "🚅 Hi, thanks for the PR! I'm **Agent Shin**, the automated triage bot for this " - "repository. " - "[What's this and why am I getting it?](https://docs.litellm.ai/blog/agent-shin-triage)\n" - "\n" - "I read the description against our " - "[contribution rubric](https://github.com/BerriAI/litellm/blob/main/.github/pull_request_template.md). " - "Here's how it lined up:\n" - "\n" - f"{present_block}" - "**What's still missing:**\n" - "\n" - f"{missing_lines}\n" - "\n" - f"> {explanation}\n" - "\n" - "**Closing this PR isn't a rejection of the change.** We want the open-PR list to " - "mirror what a maintainer can act on *right now*, so contributors don't get lost in a " - 'backlog. A closed PR is a soft "park this for later"; your work is still here, ' - "the diff is still here, and getting it reopened is one comment away. Take your time.\n" - "\n" - "**To bring this PR back:**\n" - "\n" - "- Update the description with the missing pieces, then comment `@agent-shin reconsider` " - "on this PR. I'll re-evaluate and reopen if it now passes.\n" - "- Or **Open a new PR** with the same fix and the updated description. GitHub doesn't " - "always let external contributors reopen a bot-closed PR, so a fresh PR is the most " - "reliable path back into the review queue.\n" - "- If Greptile's most recent score on this PR was below 4/5, comment `@greptileai` to " - "request a fresh review; that **still works even after the PR is closed**, and a " - "stronger score is one of the signals that lifts the PR back into the queue. A low " - "Greptile score isn't a blocker.\n" - "\n" - '**What "end-to-end QA proof" means**, since it\'s the most common gap: at least one ' - "of a short before/after screen recording / video (the bug reproducing, then the fix " - "working; for a brand-new feature, a recording of it working end-to-end), a screenshot " - "(or before/after screenshots) of it working, or the exact commands you ran paired " - "with their **real output** against the real system. Running `pytest` on the repo's " - "unit tests doesn't count; those mock the LLM provider, DB, and network, so they " - "aren't end-to-end. Output from a real, no-mocks integration run is what we look " - "for. A linked issue alone isn't enough either: it covers context, not proof. See " - "[the full rubric](https://docs.litellm.ai/blog/agent-shin-triage#the-rubric-for-pull-requests).\n" - "\n" - "Internal BerriAI contributors: this rubric doesn't apply to you; ping a maintainer.\n" - "\n" - "_(I'm an LLM, so I'm not infallible. If you think I got this wrong, comment " - "`@agent-shin reconsider` or ping a maintainer; they'll override me.)_" - f"\n\n{AGENT_SHIN_CLOSE_MARKER}" - ) - - -def format_issue_close_comment(verdict: dict) -> str: - missing_lines = _format_missing(verdict.get("missing") or []) - present_block = _format_present_block(_format_present_for_issue(verdict)) - explanation = verdict.get("explanation") or "" - return ( - "🚅 Hi, thanks for filing this! I'm **Agent Shin**, the automated triage bot for this " - "repository. " - "[What's this and why am I getting it?](https://docs.litellm.ai/blog/agent-shin-triage)\n" - "\n" - "I read the issue against our reporting checklist. Here's how it lined up:\n" - "\n" - f"{present_block}" - "**What's still missing:**\n" - "\n" - f"{missing_lines}\n" - "\n" - f"> {explanation}\n" - "\n" - "**Closing this isn't us saying the bug isn't real or the request isn't useful.** We " - "want the open-issue list to mirror what a maintainer can act on *right now*, so " - "reports like yours don't get buried in a backlog. A closed issue is a soft \"park " - 'this for later"; your report is still here, and getting it reopened is one comment ' - "away. Take your time.\n" - "\n" - "**To bring this issue back:**\n" - "\n" - "1. Edit the issue description to add the missing pieces:\n" - " - For **bug reports**: end-to-end evidence of the bug (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, " - "video, or screenshot don't count, and mocked or stubbed runs don't count.\n" - " - For **feature requests**: a concrete description of what should change, a " - "use case and example (config / API call / UI flow), plus end-to-end evidence of " - "the dead-end (a video, a screenshot, or the exact commands you ran with their " - "real output showing where the flow stops today). Mocked or stubbed runs don't " - "count.\n" - "2. Comment `@agent-shin reconsider`. I'll re-run triage and reopen the issue if it " - "now meets the bar. (GitHub doesn't let external authors reopen an issue a maintainer " - "or bot closed, so the comment-based reconsider is the reliable path.)\n" - "\n" - "Internal BerriAI contributors: this rubric doesn't apply to you; ping a maintainer.\n" - "\n" - "_(I'm an LLM, so I'm not infallible. If you think I got this wrong, comment " - "`@agent-shin reconsider` or ping a maintainer; they'll override me.)_" - f"\n\n{AGENT_SHIN_CLOSE_MARKER}" - ) - - -def format_grace_warning_pr_comment(verdict: dict) -> str: - """Comment posted on the FIRST low-quality detection — gives the - contributor a 2-hour grace window to fix the PR before the next - triage run actually closes it. - - This is the "before-close" warning. On the second triage run, if the - grace marker is older than `GRACE_PERIOD_SECONDS` AND the PR still - fails the rubric, the close path runs (which posts - `format_pr_close_comment` and closes the PR). - """ - missing_lines = _format_missing(verdict.get("missing") or []) - present_block = _format_present_block(_format_present_for_pr(verdict)) - explanation = verdict.get("explanation") or "" - return ( - "🚅 Hi, thanks for the PR! I'm **Agent Shin**, the automated triage bot for this " - "repository. " - "[What's this and why am I getting it?](https://docs.litellm.ai/blog/agent-shin-triage)\n" - "\n" - "I read the description against our " - "[contribution rubric](https://github.com/BerriAI/litellm/blob/main/.github/pull_request_template.md). " - "Here's how it lined up:\n" - "\n" - f"{present_block}" - "**What's still missing:**\n" - "\n" - f"{missing_lines}\n" - "\n" - f"> {explanation}\n" - "\n" - "If the description isn't updated in the next **2 hours**, I'll auto-close this PR. " - "That's **not** us saying we don't care about the change; we want the open-PR list to " - "mirror what a maintainer can act on *right now*, so contributors don't get lost in a " - 'backlog. A closed PR is a soft "park this for later," not a rejection. Take your ' - "time; everything below still works after the close.\n" - "\n" - "**During the grace period:** just update the PR description with the missing pieces. " - "No need to ping me; I'll re-check on the next sweep and skip the auto-close if it " - "now passes. See " - "[what counts as QA proof](https://docs.litellm.ai/blog/agent-shin-triage#the-rubric-for-pull-requests) " - "for the full rubric (a linked issue alone isn't enough; it covers context, not proof).\n" - "\n" - "**If the PR does get auto-closed in 2 hours, you still have easy recovery paths:**\n" - "\n" - "- Comment `@agent-shin reconsider` after updating the description. I'll re-evaluate " - "and reopen the PR if it now passes.\n" - "- Comment `@greptileai` to request a fresh Greptile review; that **still works even " - "after the PR is closed**, and a stronger score is one of the signals that lifts the " - "PR back into the queue. So a low Greptile score isn't a blocker either.\n" - "\n" - "Internal BerriAI contributors: this rubric doesn't apply to you; ping a maintainer.\n" - "\n" - "_(I'm an LLM, so I'm not infallible. If you think I got this wrong, ping a " - "maintainer; they'll override me.)_\n" - "\n" - f"{GRACE_COMMENT_MARKER}" - ) - - -def format_grace_warning_issue_comment(verdict: dict) -> str: - """Issue analogue of `format_grace_warning_pr_comment`.""" - missing_lines = _format_missing(verdict.get("missing") or []) - present_block = _format_present_block(_format_present_for_issue(verdict)) - explanation = verdict.get("explanation") or "" - return ( - "🚅 Hi, thanks for filing this! I'm **Agent Shin**, the automated triage bot for this " - "repository. " - "[What's this and why am I getting it?](https://docs.litellm.ai/blog/agent-shin-triage)\n" - "\n" - "I read the issue against our reporting checklist. Here's how it lined up:\n" - "\n" - f"{present_block}" - "**What's still missing:**\n" - "\n" - f"{missing_lines}\n" - "\n" - f"> {explanation}\n" - "\n" - "If the issue isn't updated in the next **2 hours**, I'll auto-close it. That's **not** us " - "saying the bug isn't real or the request isn't useful; we want the open-issue list " - "to mirror what a maintainer can act on *right now*, so reports like yours don't get " - 'buried in a backlog. A closed issue is a soft "park this for later," not a ' - "rejection. Take your time; reopening is one comment away.\n" - "\n" - "**During the grace period:** just edit the issue description with the missing " - "pieces. No need to ping me; I'll re-check on the next sweep and skip the auto-close " - "if it now passes.\n" - "\n" - "Missing pieces, depending on what this is:\n" - "\n" - "- For **bug reports**: end-to-end evidence of the bug (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 concrete description of what should change, a use " - "case and example (config / API call / UI flow), plus end-to-end evidence of the " - "dead-end (a video, a screenshot, or the exact commands you ran with their real " - "output showing where the flow stops today). Mocked or stubbed runs don't count.\n" - "\n" - "**If the issue does get auto-closed in 2 hours**, comment `@agent-shin reconsider` " - "and I'll re-evaluate. If it now meets the bar, I'll reopen the issue.\n" - "\n" - "Internal BerriAI contributors: this rubric doesn't apply to you; ping a maintainer.\n" - "\n" - "_(I'm an LLM, so I'm not infallible. If you think I got this wrong, ping a " - "maintainer; they'll override me.)_\n" - "\n" - f"{GRACE_COMMENT_MARKER}" - ) - - -# --------------------------------------------------------------------------- -# Step-summary helpers - - -def write_step_summary(content: str) -> None: - """When running inside GitHub Actions, append to the step summary file.""" - path = os.environ.get("GITHUB_STEP_SUMMARY") - if not path: - return - try: - with open(path, "a", encoding="utf-8") as handle: - handle.write(content) - if not content.endswith("\n"): - handle.write("\n") - except OSError as exc: - print(f"warn: failed to write step summary: {exc}", file=sys.stderr) - - -# --------------------------------------------------------------------------- -# Core orchestration - - -def format_reopen_comment(kind: str) -> str: - """Comment posted when Agent Shin reopens after a successful reconsider.""" - noun = "PR" if kind == "pr" else "issue" - # The trailing HTML marker is used by `seconds_since_last_reconsider_verdict` - # to enforce a cooldown between repeated `@agent-shin reconsider` triggers. - # Keep the marker on its own line so it doesn't disturb the rendered text. - return ( - f"♻️ **Re-evaluated and reopened.** Thanks for updating the {noun}!\n" - "\n" - "Agent Shin re-ran triage on the latest description and it now meets " - "the bar. A maintainer will take another look soon; please don't " - f"close this {noun} again unless asked to.\n" - "\n" - "_(If a maintainer ends up closing this for non-rubric reasons, that " - "decision stands; comment `@agent-shin reconsider` again only if you " - "have substantively new information.)_\n" - "\n" - f"{RECONSIDER_COMMENT_MARKER}" - ) - - -def format_reconsider_still_failing_comment(kind: str, verdict: dict) -> str: - """Comment posted when reconsider re-runs triage but the verdict is still fail.""" - missing_lines = _format_missing(verdict.get("missing") or []) - explanation = verdict.get("explanation") or "" - noun = "PR" if kind == "pr" else "issue" - # The trailing HTML marker is used by `seconds_since_last_reconsider_verdict` - # to enforce a cooldown between repeated `@agent-shin reconsider` triggers. - return ( - f"⏸️ **Re-evaluated; this {noun} still doesn't meet the rubric.**\n" - "\n" - "Agent Shin re-ran triage on the current description but is still " - "missing:\n" - "\n" - f"{missing_lines}\n" - "\n" - f"> {explanation}\n" - "\n" - "Update the description with the missing pieces and comment " - "`@agent-shin reconsider` again, or ping a maintainer if you think " - "I got this wrong.\n" - "\n" - "_(I'm an LLM and I'm not infallible.)_\n" - "\n" - f"{RECONSIDER_COMMENT_MARKER}" - ) - - -# --------------------------------------------------------------------------- -# Review gate — "ready for review" label lifecycle - -_UNSET = object() - - -def _combine_missing( - verdict: dict, greptile_score: int | None, min_score: int -) -> list[str]: - """Merge the LLM rubric's `missing` list with a Greptile-score shortfall.""" - missing = list(verdict.get("missing") or []) - if greptile_score is not None and greptile_score < min_score: - missing.insert( - 0, - f"Greptile's most recent review scored this PR {greptile_score}/5 " - f"(below the {min_score}/5 bar)", - ) - return missing or ["(see explanation below)"] - - -def _has_marker( - comments: Iterable[dict], marker: str, *, bot_login: str | None = None -) -> bool: - """Return True iff the bot itself posted a comment containing ``marker``. - - Filters by author so a contributor who quotes the marker (e.g. via - GitHub's "Quote reply" feature, which preserves HTML comments in - raw markdown) is not mistaken for a bot action — that would - silently suppress notifications or change which "recovered" wording - is selected. Matches the author-filter pattern used by the sibling - `_seconds_since_latest_marker_comment` helper. - """ - expected_login = ( - bot_login - or os.environ.get("AGENT_SHIN_BOT_LOGIN") - or AGENT_SHIN_DEFAULT_BOT_LOGIN - ).lower() - for comment in comments: - author = ((comment.get("user") or {}).get("login") or "").lower() - if author != expected_login: - continue - if marker in (comment.get("body") or ""): - return True - return False - - -def format_ready_for_review_comment( - verdict: dict, - greptile_score: int | None, - min_greptile_score: int = DEFAULT_MIN_GREPTILE_SCORE, -) -> str: - """Posted the first time a PR clears the bar (label added).""" - score_line = ( - f" Greptile scored it **{greptile_score}/5**." - if greptile_score is not None - else "" - ) - explanation = verdict.get("explanation") or "" - return ( - "✅ **Triage passed, tagging `ready for review`.**\n" - "\n" - "Agent Shin checked this PR against the " - "[contribution rubric](https://github.com/BerriAI/litellm/blob/main/.github/pull_request_template.md) " - "and it clears the bar (a linked issue, or a clear problem description " - f"+ expected vs. actual + QA proof).{score_line}\n" - "\n" - f"> {explanation}\n" - "\n" - "A maintainer will take it from here. If a later re-check finds the PR " - f"has regressed (Greptile drops below {min_greptile_score}/5, " - "the QA proof is removed, etc.) I'll pull the tag and comment with " - "what's missing; fix it and the tag comes back automatically.\n" - f"{READY_MARKER}" - ) - - -def format_all_clear_comment(verdict: dict, greptile_score: int | None) -> str: - """Posted when a PR recovers after a regression (label re-added).""" - score_line = ( - f" Greptile is back to **{greptile_score}/5**." - if greptile_score is not None - else "" - ) - explanation = verdict.get("explanation") or "" - return ( - "✅ **All clear again, re-adding `ready for review`.**\n" - "\n" - "Thanks for addressing the earlier feedback. On re-check this PR meets " - f"the contribution bar once more.{score_line}\n" - "\n" - f"> {explanation}\n" - "\n" - "A maintainer will take another look.\n" - f"{READY_MARKER}" - ) - - -def format_regression_comment( - missing: list[str], explanation: str, grace_days: int -) -> str: - """Posted when a previously-tagged PR regresses (label removed, PR stays open). - - Discloses the same ``grace_days`` deadline the state machine enforces: - once that window elapses with the PR still failing, the close path fires. - Hiding the deadline behind a bare "stays open" would surprise contributors - with an auto-close they were never warned about. - """ - window = "24 hours" if grace_days == 1 else f"{grace_days} days" - return ( - "⚠️ **Removing the `ready for review` tag.**\n" - "\n" - "On a re-check this PR no longer meets the contribution bar. What's " - "missing now:\n" - "\n" - f"{_format_missing(missing)}\n" - "\n" - f"> {explanation}\n" - "\n" - f"The PR stays open for ~{window}; address the points above and Agent " - 'Shin will post an "all clear" comment and re-add the tag ' - "automatically. If the points still aren't addressed after that " - "window, the PR is auto-closed; that's not a rejection, and you can " - "comment `@agent-shin reconsider` to have it re-evaluated and reopened " - "once it passes.\n" - f"{REGRESSED_MARKER}" - ) - - -def format_within_grace_comment( - missing: list[str], explanation: str, grace_days: int -) -> str: - """Posted once while a failing PR is still inside its grace window.""" - window = "24 hours" if grace_days == 1 else f"{grace_days} days" - return ( - "🚅 Hi, thanks for the PR! This is **Agent Shin**, the automated triage " - "bot. This PR doesn't quite meet the contribution bar yet:\n" - "\n" - f"{_format_missing(missing)}\n" - "\n" - f"> {explanation}\n" - "\n" - f"You have ~{window} from when this PR was opened to add the missing " - "pieces; just update the description and I'll re-check on the next " - "sweep. Once it passes I'll tag it `ready for review`. If it does get " - "auto-closed, that's not a rejection; comment `@agent-shin reconsider` " - "and I'll re-evaluate and reopen if it now passes.\n" - f"{WITHIN_GRACE_MARKER}" - ) - - -def review_gate( - *, - repo: str, - number: int, - close: bool, - model: str, - judge: Any = None, - greptile_score: Any = _UNSET, - comments: Any = _UNSET, - now: dt.datetime | None = None, - grace_days: int = DEFAULT_GRACE_DAYS, - min_greptile_score: int = DEFAULT_MIN_GREPTILE_SCORE, - label: str = READY_FOR_REVIEW_LABEL, - allowlist: frozenset[str] = ALLOWLIST_LOGINS, -) -> dict: - """Reconcile the `ready for review` label with a PR's current quality. - - A PR is *passing* when it clears BOTH gates: the LLM rubric (linked issue, - or problem description + expected/actual + QA proof) AND Greptile's most - recent confidence score (>= ``min_greptile_score``; absence of a score is - not held against the PR). The gate then drives a small state machine, using - the label itself as the persisted state so comments fire only on - transitions (never on every scheduled run): - - passing, untagged -> add label + "ready for review" / "all clear" - passing, tagged -> noop-passing - not passing, tagged -> remove label + regression comment (stays open) - not passing, untagged, old -> close + comment (past the grace window) - not passing, untagged, new -> one-time "what's missing" notice (within grace) - - ``close`` gates every destructive side effect: with ``close=False`` the - function returns a ``would-*`` preview and touches nothing, mirroring the - dry-run contract of :func:`triage`. ``judge``/``greptile_score``/ - ``comments``/``now`` are injectable for tests; in production they are - resolved from the OpenAI judge, the PR's Greptile comment, the live comment - list, and the wall clock respectively. - """ - item = fetch_pr(repo, number) - - title = item.get("title") or "" - body = item.get("body") or "" - login = (item.get("user") or {}).get("login") or "" - association = item.get("author_association") or "" - state = item.get("state") or "" - # GitHub label names are case-insensitive; compare lowercased so a repo - # that already has e.g. "Ready for Review" is recognized as the same - # label as our READY_FOR_REVIEW_LABEL constant ("ready for review"). - labels_now = {(lbl.get("name") or "").lower() for lbl in (item.get("labels") or [])} - label_key = label.lower() - created_raw = item.get("created_at") or "" - - base_result = { - "kind": "pr", - "number": number, - "title": title, - "author": login, - "author_association": association, - "state": state, - "labeled": label_key in labels_now, - "review_gate": True, - } - - if state != "open": - return {**base_result, "action": "skip-not-open"} - - if allowlist: - if login.lower() not in allowlist: - return {**base_result, "action": "skip-not-allowlisted"} - elif is_internal_contributor(item): - return {**base_result, "action": "skip-internal-author"} - - # Resolve the comment list once — used for both the Greptile score and the - # marker-based dedup below. - if comments is _UNSET: - comments = list(_iter_paginated_json(f"repos/{repo}/issues/{number}/comments")) - - # --- rubric verdict: linked-issue short-circuit, else the LLM judge ------- - if has_linked_issue(body): - verdict = { - "verdict": "pass", - "linked_issue": True, - "missing": [], - "explanation": "Linked-issue regex matched; LLM was not called.", - } - rubric_pass = True - else: - prompt = build_pr_prompt(title=title, body=body) - if judge is None: - api_key = os.environ.get("OPENAI_API_KEY") - if not api_key: - return {**base_result, "action": "skip-no-llm-key"} - base_url = os.environ.get("OPENAI_BASE_URL") or None - - def judge(p: str) -> str: - return call_llm_judge( - p, model=model, api_key=api_key, base_url=base_url - ) - - try: - verdict = parse_verdict(judge(prompt)) - except Exception as exc: # noqa: BLE001 - judge errors must never act - return {**base_result, "action": "skip-llm-error", "error": str(exc)} - rubric_pass = (verdict.get("verdict") or "").lower() == "pass" - - # --- Greptile score ------------------------------------------------------- - if greptile_score is _UNSET: - extraction = extract_greptile_score(comments) - greptile_score = extraction[0] if extraction else None - greptile_ok = greptile_score is None or greptile_score >= min_greptile_score - passing = rubric_pass and greptile_ok - - # --- age ------------------------------------------------------------------ - age_days = None - if created_raw: - reference = now or dt.datetime.now(dt.timezone.utc) - age_days = (reference - parse_iso8601(created_raw)).days - - label_present = label_key in labels_now - explanation = verdict.get("explanation") or "" - # When the rubric short-circuited to pass (linked-issue regex) but - # Greptile dragged the PR below the bar, the synthetic verdict's - # explanation ("LLM was not called") would mislead a contributor reading - # the regression / close comment. Surface the real reason instead. - if rubric_pass and not greptile_ok: - explanation = ( - f"Greptile's most recent review scored this PR " - f"{greptile_score}/5 (below the {min_greptile_score}/5 bar)." - ) - verdict = {**verdict, "explanation": explanation} - base_result = { - **base_result, - "verdict": verdict, - "greptile_score": greptile_score, - "passing": passing, - "age_days": age_days, - } - - if passing: - if label_present: - return {**base_result, "action": "noop-passing"} - recovered = _has_marker(comments, REGRESSED_MARKER) - comment = ( - format_all_clear_comment(verdict, greptile_score) - if recovered - else format_ready_for_review_comment( - verdict, greptile_score, min_greptile_score - ) - ) - if not close: - return {**base_result, "action": "would-label-ready", "comment": comment} - post_comment(repo, number, comment) - add_label(repo, number, label) - return {**base_result, "action": "labeled-ready", "comment": comment} - - missing = _combine_missing(verdict, greptile_score, min_greptile_score) - - if label_present: - comment = format_regression_comment(missing, explanation, grace_days) - if not close: - return {**base_result, "action": "would-remove-label", "comment": comment} - remove_label(repo, number, label) - post_comment(repo, number, comment) - return {**base_result, "action": "label-removed-regressed", "comment": comment} - - # Not passing and not tagged. If the PR was previously tagged and then - # regressed (we removed the label and posted REGRESSED_MARKER), honor the - # "PR stays open — fix it and the tag comes back" promise from - # `format_regression_comment` and skip the close path. Without this guard, - # any PR older than `grace_days` would be closed on the next evaluation, - # giving the contributor no realistic window to address the regression. - # - # The promise has a deliberate expiration: once `grace_days` have elapsed - # since the regression notice, fall through to the close path so a PR that - # was abandoned post-regression doesn't sit open forever. - if _has_marker(comments, REGRESSED_MARKER): - reference = now or dt.datetime.now(dt.timezone.utc) - seconds_since_regression = seconds_since_latest_marker_comment( - comments, marker=REGRESSED_MARKER, now=reference - ) - grace_seconds = grace_days * 86400 - if seconds_since_regression is None or seconds_since_regression < grace_seconds: - return {**base_result, "action": "regressed-already-notified"} - - # Not passing and not tagged: close if past the grace window, else notify once. - if age_days is not None and age_days >= grace_days: - comment = format_pr_close_comment({**verdict, "missing": missing}) - if not close: - return {**base_result, "action": "would-close", "comment": comment} - post_comment(repo, number, comment) - close_pr(repo, number) - return {**base_result, "action": "closed", "comment": comment} - - if _has_marker(comments, WITHIN_GRACE_MARKER): - return {**base_result, "action": "within-grace-already-notified"} - comment = format_within_grace_comment(missing, explanation, grace_days) - if not close: - return { - **base_result, - "action": "would-notify-within-grace", - "comment": comment, - } - post_comment(repo, number, comment) - return {**base_result, "action": "within-grace-notified", "comment": comment} - - -def triage( - *, - repo: str, - kind: str, - number: int, - close: bool, - model: str, - judge: Any = None, - print_prompt: bool = False, - reconsider: bool = False, - allowlist: frozenset[str] = ALLOWLIST_LOGINS, -) -> dict: - """Triage a single PR or issue. Returns a result dict for logging/tests. - - `judge` is an optional callable `(prompt) -> str` for tests / dry-run with - a stub. In production, leave it None and the script uses `call_llm_judge`. - - When `reconsider=True`, the closed-state guard is skipped and a - fail-but-no-comment is replaced with a "still failing" comment + leave - closed; a pass triggers `reopen_pr`/`reopen_issue` plus a reopen comment. - Reconsider mode is intended for the `@agent-shin reconsider` comment - trigger. Like regular triage, `close=False` keeps reconsider in dry-run - (returns `would-reopen` / `would-reconsider-still-failing` so a local - operator can preview without write side effects); the workflow only - passes `--close` when `AGENT_SHIN_ENABLED=true`. - - Reconsider mode adds two extra safety guards on top of the regular - triage skip-internal-author check: - - 1. **Bot-closed guard.** Only reopens if the most recent close was - performed by the bot identity (default `github-actions[bot]`). - This stops a contributor from using `@agent-shin reconsider` to - override a maintainer's close for non-rubric reasons. - 2. **Rate-limit guard.** If the bot has already posted a reconsider - verdict on this PR/issue within `RECONSIDER_RATE_LIMIT_SECONDS`, - skip — repeated triggers from the same contributor shouldn't burn - CI minutes or LLM budget. - """ - fetcher = {"pr": fetch_pr, "issue": fetch_issue}[kind] - item = fetcher(repo, number) - - title = item.get("title") or "" - body = item.get("body") or "" - login = (item.get("user") or {}).get("login") or "" - association = item.get("author_association") or "" - state = item.get("state") or "" - - base_result = { - "kind": kind, - "number": number, - "title": title, - "author": login, - "author_association": association, - "state": state, - "reconsider": reconsider, - } - - # Reconsider only makes sense on a closed PR/issue. A "reconsider on an - # open PR" is a no-op (the regular triage flow already evaluates open - # PRs); return a clear skip so the workflow can short-circuit. - if reconsider: - if state != "closed": - return {**base_result, "action": "skip-not-closed"} - else: - if state != "open": - return {**base_result, "action": "skip-not-open"} - - if allowlist: - if login.lower() not in allowlist: - return {**base_result, "action": "skip-not-allowlisted"} - elif is_internal_contributor(item): - return {**base_result, "action": "skip-internal-author"} - - # Reconsider-only guards — these run BEFORE the LLM call so a - # maintainer-closed PR / rate-limited trigger never spends LLM budget. - if reconsider: - if not was_closed_by_agent_shin(repo, number): - return {**base_result, "action": "skip-not-bot-closed"} - age = seconds_since_last_reconsider_verdict(repo, number) - if age is not None and age < RECONSIDER_RATE_LIMIT_SECONDS: - return { - **base_result, - "action": "skip-rate-limited", - "rate_limit_age_seconds": age, - "rate_limit_window_seconds": RECONSIDER_RATE_LIMIT_SECONDS, - } - - if kind == "pr": - # Short-circuit: if body very clearly links a related issue, just pass. - if has_linked_issue(body): - base = { - **base_result, - "action": "pass-linked-issue", - "verdict": { - "verdict": "pass", - "linked_issue": True, - "explanation": "Linked-issue regex matched; LLM was not called.", - }, - } - if reconsider: - # Pass-on-reconsider -> reopen the PR with a friendly comment. - reopen_body = format_reopen_comment(kind) - if not close: - return { - **base, - "action": "would-reopen", - "comment": reopen_body, - } - post_comment(repo, number, reopen_body) - reopen_pr(repo, number) - return { - **base, - "action": "reopened", - "comment": reopen_body, - } - return base - prompt = build_pr_prompt(title=title, body=body) - else: - prompt = build_issue_prompt(title=title, body=body) - - if print_prompt: - return {**base_result, "action": "print-prompt", "prompt": prompt} - - if judge is None: - api_key = os.environ.get("OPENAI_API_KEY") - if not api_key: - # No key configured — never take a destructive action. Report skip. - return { - **base_result, - "action": "skip-no-llm-key", - "prompt_preview": prompt[:200], - } - base_url = os.environ.get("OPENAI_BASE_URL") or None - - def judge(p: str) -> str: - return call_llm_judge(p, model=model, api_key=api_key, base_url=base_url) - - try: - raw = judge(prompt) - verdict = parse_verdict(raw) - except Exception as exc: # noqa: BLE001 - judge errors must never close PRs - return {**base_result, "action": "skip-llm-error", "error": str(exc)} - - decision = (verdict.get("verdict") or "").lower() - - if reconsider: - # Reconsider: an explicit `pass` -> reopen + post reopen comment; - # anything else (fail, missing/malformed verdict, typo) -> leave - # closed + post a "still failing" comment so the contributor can - # iterate again. Reopen is destructive, so a flaky/empty verdict - # must not satisfy the gate. - # In dry-run (`close=False`) we return `would-*` actions instead - # of touching GitHub state, mirroring the regular triage flow's - # `would-close`. This lets a local operator preview the outcome - # of `python triage_with_llm.py --reconsider --pr N` without - # risking accidental comments or reopens. - if decision == "pass": - reopen_body = format_reopen_comment(kind) - if not close: - return { - **base_result, - "action": "would-reopen", - "verdict": verdict, - "comment": reopen_body, - } - post_comment(repo, number, reopen_body) - if kind == "pr": - reopen_pr(repo, number) - else: - reopen_issue(repo, number) - return { - **base_result, - "action": "reopened", - "verdict": verdict, - "comment": reopen_body, - } - still_failing = format_reconsider_still_failing_comment(kind, verdict) - if not close: - return { - **base_result, - "action": "would-reconsider-still-failing", - "verdict": verdict, - "comment": still_failing, - } - post_comment(repo, number, still_failing) - return { - **base_result, - "action": "reconsider-still-failing", - "verdict": verdict, - "comment": still_failing, - } - - if decision != "fail": - return {**base_result, "action": "pass-llm", "verdict": verdict} - - # Grace-period flow: on the first low-quality detection, post a warning - # comment instead of closing immediately. On a subsequent triage run - # (manual re-trigger, or the daily `close_low_quality_prs.py` cron - # finding the same PR in its own pass), if `GRACE_PERIOD_SECONDS` has - # elapsed since the warning AND the PR still fails the rubric, close. - grace_age = seconds_since_last_grace_warning(repo, number) - if grace_age is None: - warning_body = ( - format_grace_warning_pr_comment(verdict) - if kind == "pr" - else format_grace_warning_issue_comment(verdict) - ) - if not close: - return { - **base_result, - "action": "would-warn-grace", - "verdict": verdict, - "comment": warning_body, - } - post_comment(repo, number, warning_body) - return { - **base_result, - "action": "warned-grace", - "verdict": verdict, - "comment": warning_body, - } - if grace_age < GRACE_PERIOD_SECONDS: - return { - **base_result, - "action": "skip-in-grace-period", - "verdict": verdict, - "grace_age_seconds": grace_age, - "grace_period_seconds": GRACE_PERIOD_SECONDS, - } - - # The grace window has elapsed. `--close` still gates the destructive - # write so a dry-run preview never posts or closes — the workflow only - # passes `--close` when `AGENT_SHIN_ENABLED=true`, which keeps the bot - # inert by default. - if not close: - return {**base_result, "action": "would-close", "verdict": verdict} - - comment_body = ( - format_pr_close_comment(verdict) - if kind == "pr" - else format_issue_close_comment(verdict) - ) - post_comment(repo, number, comment_body) - if kind == "pr": - close_pr(repo, number) - else: - close_issue(repo, number) - - return { - **base_result, - "action": "closed", - "verdict": verdict, - "comment": comment_body, - } - - -# --------------------------------------------------------------------------- -# CLI - - -def render_summary(result: dict) -> str: - """Render a human-readable summary block (used for stdout + step summary).""" - lines = ["## Agent Shin verdict", ""] - lines.append( - f"- **{result['kind'].upper()} #{result['number']}**: {result.get('title', '')}" - ) - lines.append( - f"- **Author**: `{result.get('author', '')}` ({result.get('author_association', '')})" - ) - lines.append(f"- **State**: {result.get('state', '')}") - lines.append(f"- **Action**: `{result['action']}`") - verdict = result.get("verdict") - if verdict: - lines.append("") - lines.append("```json") - lines.append(json.dumps(verdict, indent=2)) - lines.append("```") - error = result.get("error") - if error: - lines.append("") - lines.append(f"_LLM error: {error}_") - comment = result.get("comment") - if comment: - lines.append("") - lines.append("### Posted comment:") - lines.append("") - lines.append("> " + comment.replace("\n", "\n> ")) - return "\n".join(lines) - - -def main() -> int: - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--repo", required=True, help="Repository (owner/repo).") - target = parser.add_mutually_exclusive_group(required=True) - target.add_argument("--pr", type=int, help="Pull request number to triage.") - target.add_argument("--issue", type=int, help="Issue number to triage.") - parser.add_argument( - "--close", - action="store_true", - help="Actually post comment + close on fail (default: dry run).", - ) - parser.add_argument( - "--model", - # `os.environ.get("TRIAGE_MODEL", DEFAULT_MODEL)` would return "" when - # GitHub Actions exposes an unset repo variable as an empty-string env - # var, silently bypassing DEFAULT_MODEL and causing every call to fail - # as `skip-llm-error`. The `or` guard collapses empty -> default. - default=os.environ.get("TRIAGE_MODEL") or DEFAULT_MODEL, - help=f"OpenAI-compatible model name (default: {DEFAULT_MODEL}).", - ) - parser.add_argument( - "--print-prompt", - action="store_true", - help="Print the prompt that would be sent to the judge and exit.", - ) - parser.add_argument( - "--reconsider", - action="store_true", - help=( - "Re-run triage on a CLOSED PR/issue and reopen it on pass. " - "Used by the `@agent-shin reconsider` comment-trigger workflow. " - "Only invoke this from a workflow that has already gated on " - "AGENT_SHIN_ENABLED=true and verified the commenter is the " - "PR/issue author or an internal collaborator." - ), - ) - parser.add_argument( - "--review-gate", - action="store_true", - help=( - "Reconcile the `ready for review` label for an OPEN PR: tag on " - "pass, remove the tag + comment on regression, close after the " - "grace window if it never passed. PR-only." - ), - ) - parser.add_argument( - "--grace-days", - type=int, - default=DEFAULT_GRACE_DAYS, - help=( - "Review-gate only: hours/24 a failing, un-tagged PR may stay open " - f"before auto-close (default: {DEFAULT_GRACE_DAYS} = 24h)." - ), - ) - parser.add_argument( - "--min-greptile-score", - type=int, - default=DEFAULT_MIN_GREPTILE_SCORE, - choices=range(1, 6), - help=( - "Review-gate only: Greptile score below which a PR counts as not " - f"passing (default: {DEFAULT_MIN_GREPTILE_SCORE} -> <4/5 regresses)." - ), - ) - args = parser.parse_args() - - kind = "pr" if args.pr is not None else "issue" - number = args.pr if args.pr is not None else args.issue - - if args.review_gate: - if kind != "pr": - parser.error("--review-gate applies to pull requests only (use --pr).") - result = review_gate( - repo=args.repo, - number=number, - close=args.close, - model=args.model, - grace_days=args.grace_days, - min_greptile_score=args.min_greptile_score, - ) - else: - result = triage( - repo=args.repo, - kind=kind, - number=number, - close=args.close, - model=args.model, - print_prompt=args.print_prompt, - reconsider=args.reconsider, - ) - - if result.get("action") == "print-prompt": - print(result["prompt"]) - return 0 - - summary = render_summary(result) - print(summary) - write_step_summary(summary + "\n") - return 0 - - -if __name__ == "__main__": - sys.exit(main()) diff --git a/.github/workflows/close_low_quality_prs.yml b/.github/workflows/close_low_quality_prs.yml deleted file mode 100644 index 2401be84000..00000000000 --- a/.github/workflows/close_low_quality_prs.yml +++ /dev/null @@ -1,92 +0,0 @@ -name: Close Low-Quality PRs - -# Auto-close any open PR (including drafts, regardless of age) authored by an -# external OSS contributor that Greptile reviewed with a confidence score -# below 4/5. Closures are explained in a comment that tells the contributor -# to push fixes and open a fresh PR (since OSS authors cannot reopen a PR -# closed by a bot/maintainer) or comment `@agent-shin reconsider` to have -# Agent Shin re-evaluate. -# -# Manual one-off run: -# gh workflow run "Close Low-Quality PRs" -f close=true -# -# Dry-run preview (no PRs are touched): -# gh workflow run "Close Low-Quality PRs" -f close=false - -on: - schedule: - # Daily at 09:00 UTC. Pairs well with the stale-issue workflow at midnight. - - cron: "0 9 * * *" - workflow_dispatch: - inputs: - close: - description: "Actually close matching PRs (false = dry run)." - required: false - default: "false" - type: choice - options: - - "true" - - "false" - min_age_days: - description: "Minimum PR age in days (default 0 = no age filter)." - required: false - default: "0" - min_score: - description: "Greptile score below which a PR is closed (1-5)." - required: false - default: "4" - limit: - description: "Maximum number of PRs to close in a single run." - required: false - default: "25" - -permissions: - contents: read - pull-requests: write - issues: write - -jobs: - close-low-quality-prs: - if: github.repository == 'BerriAI/litellm' - runs-on: ubuntu-latest - steps: - - name: Checkout triage script - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 - with: - sparse-checkout: .github/scripts - persist-credentials: false - - - name: Set up Python - uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0 - with: - python-version: "3.12" - - - name: Run low-quality PR closer - env: - GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} - # Scheduled runs are ALWAYS dry-run, even when AGENT_SHIN_ENABLED is - # "true", so the team can QA the closer's verdicts in step summaries - # before any contributor sees a PR closed. Real closures only happen - # on manual workflow_dispatch with close=true (and the variable set). - CLOSE_FLAG: ${{ github.event.inputs.close || 'false' }} - AGENT_SHIN_ENABLED: ${{ vars.AGENT_SHIN_ENABLED }} - MIN_AGE_DAYS: ${{ github.event.inputs.min_age_days || '0' }} - MIN_SCORE: ${{ github.event.inputs.min_score || '4' }} - LIMIT: ${{ github.event.inputs.limit || '25' }} - run: | - set -euo pipefail - ARGS=( - --repo "${{ github.repository }}" - --min-age-days "${MIN_AGE_DAYS}" - --min-score "${MIN_SCORE}" - --limit "${LIMIT}" - ) - if [ "${AGENT_SHIN_ENABLED:-false}" != "true" ]; then - echo "::notice::AGENT_SHIN_ENABLED is not 'true' -> forcing dry-run regardless of close input." - elif [ "${GITHUB_EVENT_NAME:-}" = "workflow_dispatch" ] && [ "${CLOSE_FLAG}" = "true" ]; then - ARGS+=(--close) - echo "::notice::Running in close-on-fail mode." - else - echo "::notice::AGENT_SHIN_ENABLED is true but this trigger is dry-run (scheduled event or close=false)." - fi - python3 .github/scripts/close_low_quality_prs.py "${ARGS[@]}" diff --git a/.github/workflows/create_daily_oss_agent_shin_branch.yml b/.github/workflows/create_daily_oss_agent_shin_branch.yml deleted file mode 100644 index 9baf9f142f6..00000000000 --- a/.github/workflows/create_daily_oss_agent_shin_branch.yml +++ /dev/null @@ -1,28 +0,0 @@ -name: Create Daily oss-agent-shin Branch - -on: - schedule: - - cron: "0 0 * * *" # Runs every day at midnight UTC - workflow_dispatch: # Allow manual trigger - -jobs: - create-oss-agent-shin-branch: - if: github.repository == 'BerriAI/litellm' - runs-on: ubuntu-latest - permissions: - contents: write - - steps: - - name: Create daily oss-agent-shin branch - env: - GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} - run: | - BRANCH_NAME="litellm_oss_agent_shin_$(date +'%m_%d_%Y')" - echo "Creating branch: $BRANCH_NAME" - if gh api "repos/${{ github.repository }}/git/ref/heads/$BRANCH_NAME" --silent 2>/dev/null; then - echo "Branch $BRANCH_NAME already exists. Skipping creation." - exit 0 - fi - MAIN_SHA=$(gh api "repos/${{ github.repository }}/git/ref/heads/main" --jq '.object.sha') - gh api "repos/${{ github.repository }}/git/refs" -f ref="refs/heads/$BRANCH_NAME" -f sha="$MAIN_SHA" --silent - echo "Successfully created branch: $BRANCH_NAME at $MAIN_SHA" diff --git a/.github/workflows/triage_reconsider.yml b/.github/workflows/triage_reconsider.yml deleted file mode 100644 index f35f681d09a..00000000000 --- a/.github/workflows/triage_reconsider.yml +++ /dev/null @@ -1,172 +0,0 @@ -name: Agent Shin — reconsider - -# Comment-trigger workflow: when the PR/issue author (or an internal -# collaborator) comments `@agent-shin reconsider` on a CLOSED PR/issue, -# Agent Shin re-runs LLM-judge triage on the current title+body and: -# -# - on PASS: posts a "re-evaluated and reopened" comment + reopens. -# - on FAIL: posts a "still missing X" comment and leaves it closed, -# so the contributor can iterate again. -# -# This exists because GitHub does NOT let an external (non-write-access) -# OSS contributor reopen a PR/issue closed by a bot or maintainer. Without -# this comment trigger, a contributor whose PR Agent Shin auto-closed -# would have no path back into the review queue except opening a fresh PR -# (which loses the original PR's history). The bot, on the other hand, -# has write access via GH_TOKEN and can reopen on their behalf. -# -# DRY-RUN BY DEFAULT — gated on `vars.AGENT_SHIN_ENABLED == 'true'` just -# like the other Agent Shin workflows. The workflow also gates on the -# commenter being either the PR/issue author or an internal collaborator -# (OWNER/MEMBER/COLLABORATOR) so random commenters cannot DOS the LLM -# judge or force a reopen. - -on: - issue_comment: - types: [created] - -permissions: - contents: read - issues: write - pull-requests: write - -jobs: - reconsider: - if: | - github.repository == 'BerriAI/litellm' - && contains(github.event.comment.body, '@agent-shin reconsider') - runs-on: ubuntu-latest - steps: - - name: Authorize commenter - # Only the PR/issue author OR an internal collaborator may trigger - # a reconsider. Outside random commenters could otherwise spam the - # phrase to burn LLM budget or, if a fail-open bug were ever - # introduced, force a reopen on someone else's behalf. - # - # We expose the authorization decision as a step output and gate - # every subsequent (potentially destructive) step on it. A `run:` - # step with `exit 0` would NOT stop the job — only `if:` gating - # on a known-true output is safe here. - id: auth - env: - COMMENTER: ${{ github.event.comment.user.login }} - AUTHOR: ${{ github.event.issue.user.login }} - ASSOCIATION: ${{ github.event.comment.author_association }} - run: | - set -euo pipefail - if [ "${COMMENTER}" = "${AUTHOR}" ]; then - echo "::notice::Authorized: commenter is the PR/issue author." - echo "authorized=true" >> "$GITHUB_OUTPUT" - exit 0 - fi - case "${ASSOCIATION}" in - OWNER|MEMBER|COLLABORATOR) - echo "::notice::Authorized: commenter is an internal collaborator (${ASSOCIATION})." - echo "authorized=true" >> "$GITHUB_OUTPUT" - ;; - *) - echo "::notice::Commenter '${COMMENTER}' (${ASSOCIATION}) is not authorized to trigger reconsider; skipping subsequent steps." - echo "authorized=false" >> "$GITHUB_OUTPUT" - ;; - esac - - - name: React 👀 to acknowledge the reconsider - # Add an eyes reaction to the triggering comment the moment we accept - # it, so the contributor gets instant feedback that the bot saw their - # `@agent-shin reconsider` before the slower triage steps run. Gated on - # AGENT_SHIN_ENABLED so dry-run leaves no visible trace. Best-effort: - # a reactions API hiccup must never fail the actual reconsider. - if: steps.auth.outputs.authorized == 'true' && vars.AGENT_SHIN_ENABLED == 'true' - env: - GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} - COMMENT_ID: ${{ github.event.comment.id }} - run: | - set -euo pipefail - gh api --method POST \ - -H "Accept: application/vnd.github+json" \ - "repos/${{ github.repository }}/issues/comments/${COMMENT_ID}/reactions" \ - -f content=eyes \ - || echo "::warning::failed to add 👀 reaction (non-fatal)" - - - name: Checkout triage script - if: steps.auth.outputs.authorized == 'true' - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 - with: - sparse-checkout: .github/scripts - persist-credentials: false - - - name: Set up Python - if: steps.auth.outputs.authorized == 'true' - uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0 - with: - python-version: "3.12" - - - name: Install LLM client - if: steps.auth.outputs.authorized == 'true' - run: pip install --no-cache-dir --require-hashes -r .github/scripts/triage-requirements.txt - - - name: Run Agent Shin reconsider - if: steps.auth.outputs.authorized == 'true' - env: - GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} - # Only expose the LLM key when the bot is enabled, so a PR/issue - # author can't force paid LLM calls by spamming `@agent-shin - # reconsider` while the bot is still in dry-run. The Python script - # calls the LLM whenever this var is set (regardless of `--close`); - # stripping `--close` doesn't suppress the API call, only the - # destructive side effects. Mirror the gating used by every other - # Agent Shin workflow (triage_pr_with_llm.yml, review_gate.yml, ...). - OPENAI_API_KEY: ${{ vars.AGENT_SHIN_ENABLED == 'true' && secrets.OPENAI_API_KEY || '' }} - OPENAI_BASE_URL: ${{ vars.OPENAI_BASE_URL }} - TRIAGE_MODEL: ${{ vars.TRIAGE_MODEL }} - AGENT_SHIN_ENABLED: ${{ vars.AGENT_SHIN_ENABLED }} - # `issue_comment` events fire for both issues and PR comments. - # `issue.pull_request` is set iff this is a PR comment, so we use - # its presence to decide whether to invoke `--pr N` or `--issue N`. - IS_PR: ${{ github.event.issue.pull_request != null }} - NUMBER: ${{ github.event.issue.number }} - run: | - set -euo pipefail - if [ "${IS_PR}" = "true" ]; then - ARGS=(--repo "${{ github.repository }}" --pr "${NUMBER}" --reconsider) - else - ARGS=(--repo "${{ github.repository }}" --issue "${NUMBER}" --reconsider) - fi - # Reconsider's destructive actions (post comment + reopen) are - # gated on `--close`, mirroring the regular triage workflows. - # When AGENT_SHIN_ENABLED is not the EXACT string "true", we - # still run the script so its verdict + would-X action lands in - # the step summary for QA — but without `--close`, the script - # returns `would-reopen` / `would-reconsider-still-failing` - # instead of touching GitHub state. - # - # Use the positive `= "true"` gate (not `!= "true" -> exit`) so - # the workflow guardrails in - # tests/test_litellm/test_github_triage_workflows.py see the - # canonical fail-safe enable pattern. Unknown values like - # "True", "yes", "1", or typos fall through to the dry-run - # branch, which is the safe default. - if [ "${AGENT_SHIN_ENABLED:-false}" = "true" ]; then - ARGS+=(--close) - echo "::notice::Agent Shin reconsider ENABLED — running real triage (close=true)." - else - echo "::notice::AGENT_SHIN_ENABLED is not 'true' -> reconsider stays in dry-run (no comment, no reopen)." - fi - python3 .github/scripts/triage_with_llm.py "${ARGS[@]}" - - - name: React 👍 when the reconsider finishes - # Once the reconsider run has completed successfully, add a thumbs-up so - # the contributor sees the bot is done (the 👀 stays, signalling - # seen -> handled). `success()` keeps this from firing if the run - # errored, and the AGENT_SHIN_ENABLED gate keeps dry-run inert. - if: success() && steps.auth.outputs.authorized == 'true' && vars.AGENT_SHIN_ENABLED == 'true' - env: - GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} - COMMENT_ID: ${{ github.event.comment.id }} - run: | - set -euo pipefail - gh api --method POST \ - -H "Accept: application/vnd.github+json" \ - "repos/${{ github.repository }}/issues/comments/${COMMENT_ID}/reactions" \ - -f content=+1 \ - || echo "::warning::failed to add 👍 reaction (non-fatal)" diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index c359ca19986..5a8a613204f 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -2027,24 +2027,18 @@ dependencies = [ "tokio", ] -[[package]] -name = "litellm-callbacks" -version = "0.1.0" -dependencies = [ - "rstest", - "serde_json", - "tokio", -] - [[package]] name = "litellm-callbacks-legacy" version = "0.1.0" dependencies = [ - "litellm-callbacks", + "litellm-auth", + "litellm-host", "litellm-host-python", + "proptest", "pyo3", "rstest", "serde_json", + "strum", ] [[package]] @@ -2056,8 +2050,8 @@ dependencies = [ "futures-util", "litellm-auth", "litellm-auth-aws", - "litellm-callbacks", "litellm-core-utils", + "litellm-host", "litellm-llms", "litellm-types", "mime_guess", @@ -2110,12 +2104,22 @@ dependencies = [ "tokio", ] +[[package]] +name = "litellm-host" +version = "0.1.0" +dependencies = [ + "litellm-auth", + "rstest", + "serde_json", + "tokio", +] + [[package]] name = "litellm-host-python" version = "0.1.0" dependencies = [ "futures-util", - "litellm-callbacks", + "litellm-host", "pyo3", "pyo3-async-runtimes", "pythonize", @@ -2139,9 +2143,9 @@ dependencies = [ "litellm-auth-aws", "litellm-auth-azure", "litellm-auth-gcp", - "litellm-callbacks", "litellm-core-utils", "litellm-framing", + "litellm-host", "litellm-types", "reqwest 0.12.28", "rstest", @@ -2598,6 +2602,25 @@ dependencies = [ "unicode-ident", ] +[[package]] +name = "proptest" +version = "1.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4b45fcc2344c680f5025fe57779faef368840d0bd1f42f216291f0dc4ace4744" +dependencies = [ + "bit-set", + "bit-vec", + "bitflags", + "num-traits", + "rand 0.9.5", + "rand_chacha 0.9.0", + "rand_xorshift", + "regex-syntax", + "rusty-fork", + "tempfile", + "unarray", +] + [[package]] name = "pyo3" version = "0.29.2" @@ -2679,6 +2702,12 @@ dependencies = [ "serde", ] +[[package]] +name = "quick-error" +version = "1.2.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a1d01941d82fa2ab50be1e79e6714289dd7cde78eba4c074bc5a4374f650dfe0" + [[package]] name = "quinn" version = "0.11.11" @@ -2842,6 +2871,15 @@ dependencies = [ "rand_core 0.10.1", ] +[[package]] +name = "rand_xorshift" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "513962919efc330f829edb2535844d1b912b0fbe2ca165d613e4e8788bb05a5a" +dependencies = [ + "rand_core 0.9.5", +] + [[package]] name = "rayon" version = "1.12.0" @@ -3241,6 +3279,18 @@ version = "1.0.23" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "cf54715a573b99ac80df0bc206da022bcd442c974952c7b9720069370852e21f" +[[package]] +name = "rusty-fork" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cc6bf79ff24e648f6da1f8d1f011e9cac26491b619e6b9280f2b47f1774e6ee2" +dependencies = [ + "fnv", + "quick-error", + "tempfile", + "wait-timeout", +] + [[package]] name = "ryu" version = "1.0.23" @@ -4096,6 +4146,12 @@ dependencies = [ "syn 2.0.119", ] +[[package]] +name = "unarray" +version = "0.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "eaea85b334db583fe3274d12b4cd1880032beab409c0d774be044d4480ab9a94" + [[package]] name = "unicase" version = "2.9.0" @@ -4209,6 +4265,15 @@ version = "0.8.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5c3082ca00d5a5ef149bb8b555a72ae84c9c59f7250f013ac822ac2e49b19c64" +[[package]] +name = "wait-timeout" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09ac3b126d3914f9849036f826e054cbabdc8519970b8998ddaf3b5bd3c65f11" +dependencies = [ + "libc", +] + [[package]] name = "walkdir" version = "2.5.0" diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index ffdbf64bb49..de6eacc62ee 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -10,7 +10,7 @@ repository = "https://github.com/BerriAI/litellm" [workspace.dependencies] litellm-core = { path = "crates/core" } -litellm-callbacks = { path = "crates/callbacks" } +litellm-host = { path = "crates/host" } litellm-callbacks-legacy = { path = "crates/callbacks-legacy" } litellm-framing = { path = "crates/framer" } litellm-auth = { path = "crates/auth" } @@ -26,6 +26,7 @@ litellm-token-counter = { path = "crates/token-counter" } litellm-host-python = { path = "crates/host-python" } bytes = "1" +proptest = "1.7.0" pyo3 = "0.29.2" pyo3-async-runtimes = { version = "0.29.0", features = ["tokio-runtime"] } pythonize = "0.29.0" diff --git a/litellm-rust/crates/auth/src/secret.rs b/litellm-rust/crates/auth/src/secret.rs index 3ecb0a835ee..a07fe3eaad9 100644 --- a/litellm-rust/crates/auth/src/secret.rs +++ b/litellm-rust/crates/auth/src/secret.rs @@ -1,6 +1,8 @@ +use serde::Deserialize; use veil::Redact; -#[derive(Redact, Clone)] +#[derive(Redact, Clone, Deserialize)] +#[serde(transparent)] pub struct SecretValue(#[redact(with = "[REDACTED]")] String); impl SecretValue { diff --git a/litellm-rust/crates/callbacks-legacy/AGENTS.md b/litellm-rust/crates/callbacks-legacy/AGENTS.md index e4762d3037a..8b2e1c15f6e 100644 --- a/litellm-rust/crates/callbacks-legacy/AGENTS.md +++ b/litellm-rust/crates/callbacks-legacy/AGENTS.md @@ -1,15 +1,17 @@ - Target invariants, not completion claims - Keep this crate the legacy `@client` wrapper as the native call sees it, and nothing else: the `Logging` contract (`function_setup`, the deployment hooks, `pre_call`/`post_call`, the sync and async success and failure fan-out, the deferred proxy release, the argument sharing those callbacks rely on) plus the kwargs rewrites the wrapper makes on the way in (credential-name inheritance, the budget and retry-count limits) - - The driver in `litellm-host-python`, the routes and core see one `CallbackAdapter`; they never learn which Python objects consume a call + - The driver in `litellm-host-python`, the routes and core see one `PythonLifecycle`; they never learn which Python objects consume a call +- Rust drives the call; every litellm Python internal it still borrows is a variant of `LegacyPython`, grouped by subsystem (`Wrapper`, `Logging`, `DeploymentHooks`) + - The enum only shrinks: when Rust owns a subsystem, delete its group rather than adding a Rust path beside it + - Calling a user's own callback directly is permanent Python surface and gets its own type outside `LegacyPython` - `PublicCall` is the caller's call as `Logging` sees it: the positional arguments, the keyword view as the legacy path rewrites it (setup, deployment hook, prepare) and the bound request object whose attributes back keywords the caller omitted; routes hand it over through `run_legacy_call` and keep no copy -- `setup` decides once who owns the `Logging` instance and returns it as `CallSetup.bridge_owned`; `PythonLogger` carries it and nothing on the instance records it - - A logger the caller passed as `litellm_logging_obj` is caller-owned and observed in full, because the caller reads it after the call; the proxy is the live case - - A logger `function_setup` built for this call is bridge-owned, so each fan-out phase is skipped when `callbacks_needed` finds no registry, dynamic callback, `logger_fn` or debug switch for it; cost, timing and response metadata still run +- `setup` reuses a `Logging` the caller passed as `litellm_logging_obj` (the proxy and Router are the live cases) and otherwise builds one through `function_setup`, as `@client` does + - Either way every phase calls the same `Logging` method the Python path calls; which callbacks run is `Logging`'s decision, never this crate's - Callbacks receive the caller's own objects and may mutate them; this crate alone carries that obligation - Retain complete boundary arguments, opaque unknown values, aliases, omitted/default distinctions and deliberate copies; preserve the established deployment-hook kwargs view - - Re-alias every `passthrough_fields` body key to the caller's object before `pre_call`; a keyword wins over the request attribute even when it is an explicit `None` + - Before `pre_call`, re-alias every body key whose value equals the caller's argument to the caller's own object; this crate compares the two itself, and the argument is resolved by `litellm_host_python::lookup` - Retain independently captured body/header roots from `pre_call` to `post_call`; in-place mutation reaches the wire, envelope field replacement is visible to later callbacks only - - A later kind of callback host (WASM, in-process Rust) has none of these obligations, so they stay out of `litellm-callbacks`, `litellm-host-python` and the bridge; the only facts that cross from the route are the prepared keyword view and `RequestContext.passthrough_fields` + - A later kind of callback host (WASM, in-process Rust) has none of these obligations, so they stay out of `litellm-host`, `litellm-host-python` and the bridge; the only fact that crosses from the route is the prepared keyword view - Success and failure handlers receive the exact selected public response or exception; logging projections, redaction and snapshots keep their own copy contracts - Ordinary failure-handler errors cannot suppress the other eligible family or replace the mapped provider error; a cancellation ends the call with no further dispatch - Dispatch errors never replay provider work or trigger the opposite outcome; the proxy's acceptance or rejection releases deferred success at most once diff --git a/litellm-rust/crates/callbacks-legacy/Cargo.toml b/litellm-rust/crates/callbacks-legacy/Cargo.toml index 96c9c9ed560..023c13d912b 100644 --- a/litellm-rust/crates/callbacks-legacy/Cargo.toml +++ b/litellm-rust/crates/callbacks-legacy/Cargo.toml @@ -7,10 +7,15 @@ repository.workspace = true autotests = false [dependencies] -litellm-callbacks.workspace = true +litellm-host.workspace = true litellm-host-python.workspace = true + pyo3.workspace = true +strum.workspace = true +serde_json.workspace = true [dev-dependencies] +litellm-auth.workspace = true +proptest.workspace = true rstest.workspace = true serde_json.workspace = true diff --git a/litellm-rust/crates/callbacks-legacy/python_contract.json b/litellm-rust/crates/callbacks-legacy/python_contract.json new file mode 100644 index 00000000000..8a7f3b98f47 --- /dev/null +++ b/litellm-rust/crates/callbacks-legacy/python_contract.json @@ -0,0 +1,118 @@ +{ + "setup": [ + "call_type", + "args", + "kwargs", + "start_time", + "asynchronous" + ], + "check_limits": [ + "kwargs" + ], + "finalize": [ + "response", + "logger", + "kwargs", + "start_time", + "end_time" + ], + "update_logging": [ + "logger", + "kwargs", + "model", + "optional_params", + "litellm_params", + "custom_llm_provider" + ], + "pre_call": [ + "logger", + "input", + "api_key", + "additional_args" + ], + "post_call": [ + "logger", + "original_response", + "api_key", + "additional_args" + ], + "defers_async_logging": [ + "logger" + ], + "defer_success": [ + "logger", + "pending" + ], + "sync_success_for_async_call": [ + "logger", + "response", + "start", + "end" + ], + "failure_handler": [ + "logger", + "error", + "start", + "end", + "asynchronous" + ], + "submit_success": [ + "logger", + "response", + "start", + "end" + ], + "async_success_handler": [ + "logger", + "response", + "start", + "end" + ], + "enqueue_logging": [ + "coroutine" + ], + "restore_context": [ + "logger" + ], + "custom_pricing_fields": [], + "is_internal_call": [], + "credential_list": [], + "warn_unknown_credential": [ + "name", + "loaded" + ], + "before_deployment_call": [ + "kwargs", + "call_type" + ], + "after_deployment_success": [ + "kwargs", + "response", + "call_type" + ], + "after_deployment_failure": [ + "kwargs", + "error", + "call_type" + ], + "stream_opened": [ + "logger" + ], + "stream_success": [ + "logger", + "url_route", + "endpoint_type", + "request_body", + "chunks", + "start", + "end", + "first_chunk" + ], + "stream_failure": [ + "logger", + "endpoint_type", + "request_body", + "chunks", + "error" + ] +} diff --git a/litellm-rust/crates/callbacks-legacy/src/adapter.rs b/litellm-rust/crates/callbacks-legacy/src/adapter.rs index df346506094..6c013cd1ea5 100644 --- a/litellm-rust/crates/callbacks-legacy/src/adapter.rs +++ b/litellm-rust/crates/callbacks-legacy/src/adapter.rs @@ -2,21 +2,26 @@ //! raises is answered with the same `Logging` calls, in the same order, as the Python //! `@client` path makes them. -use litellm_callbacks::event::{CallEvent, FailureOrigin, RequestContext, Timing, WireRequest}; +use litellm_host::event::{ + FailureOrigin, MachineEvent, RequestContext, Timing, WireRequest, epoch_seconds, +}; use litellm_host_python::{ - AdapterStep, CallbackAdapter, PublicValue, from_py, missing_state, to_py, + LifecycleEvent, LifecycleStep, PythonLifecycle, from_py, missing_state, to_py, }; use pyo3::{ exceptions::{PyBaseException, PyException}, gc::{PyTraverseError, PyVisit}, prelude::*, - types::PyDict, + types::{PyDict, PyList}, }; +use serde_json::Value; use crate::{ DeploymentHooks, LegacyCallbacks, PublicCall, PythonLogger, deferred::{PendingLogging, PendingSuccess}, - finalize, is_internal_call, prepare, setup, + finalize, is_internal_call, + legacy_python::Streaming, + prepare, setup, }; /// What the legacy contract needs to know about the route it is logging. @@ -25,6 +30,22 @@ pub struct LegacySurface { pub call_type: &'static str, /// What `Logging.pre_call` is told the input was. pub input_description: &'static str, + /// How a streamed response is billed; `None` for a route that never streams. + pub stream: Option, +} + +/// The pass-through billing a streamed response goes through once its chunks are in. +#[derive(Clone, Copy, Debug)] +pub struct PassThroughStream { + pub url_route: &'static str, + /// A value of Python's `EndpointType`. + pub endpoint_type: &'static str, +} + +/// What the Messages stream iterator keeps for its end-of-stream billing. +struct DeliveredStream { + chunks: Py, + first_chunk: Option>, } enum Pending { @@ -44,6 +65,8 @@ pub struct LegacyLogging { error: Option>, body: Option>, headers: Option>, + context: Option, + stream: Option, asynchronous: bool, internal: bool, pending: Option, @@ -77,6 +100,8 @@ impl LegacyLogging { error: None, body: None, headers: None, + context: None, + stream: None, asynchronous, internal: false, pending: None, @@ -85,8 +110,8 @@ impl LegacyLogging { /// Deployment hooks are awaited, and Python's synchronous `@client` wrapper never /// runs them. - fn deployment_hooks(&self, py: Python<'_>) -> PyResult { - Ok(self.asynchronous && DeploymentHooks::needed(py)?) + fn runs_deployment_hooks(&self) -> bool { + self.asynchronous } fn logger(&self) -> PyResult<&PythonLogger> { @@ -95,13 +120,13 @@ impl LegacyLogging { }) } - fn prepare(&mut self, py: Python<'_>) -> PyResult { + fn prepare(&mut self, py: Python<'_>) -> PyResult { let prepared = prepare(py, self.call.kwargs().bind(py), self.logger()?)?.unbind(); self.call.set_kwargs(prepared); - Ok(AdapterStep::Arguments(self.call.kwargs().clone_ref(py))) + Ok(LifecycleStep::Arguments(self.call.kwargs().clone_ref(py))) } - fn finalize(&mut self, py: Python<'_>) -> PyResult { + fn finalize(&mut self, py: Python<'_>) -> PyResult { finalize( py, &self.response, @@ -112,7 +137,7 @@ impl LegacyLogging { )?; self.response .as_ref() - .map(|response| AdapterStep::Response(response.clone_ref(py))) + .map(|response| LifecycleStep::Response(response.clone_ref(py))) .ok_or_else(missing_state) } @@ -145,9 +170,7 @@ impl LegacyLogging { .get_item("fallbacks")? .is_none_or(|value| value.is_none()) { - if !logger.callbacks_needed(py, "async_success")? { - logger.success_bookkeeping(py, &self.response, &self.start, &self.end, true)?; - } else if logger.defers_async_logging(py) { + if logger.defers_async_logging(py) { let pending = Py::new( py, PendingLogging { @@ -162,15 +185,72 @@ impl LegacyLogging { logger.sync_success_for_async_call(py, &self.response, &self.start, &self.end) } + fn stream_success(&self, py: Python<'_>, stream: &DeliveredStream) -> PyResult<()> { + let logger = self.logger()?; + let billing = self.surface.stream.ok_or_else(missing_state)?; + let billed = Streaming::Success.call( + py, + ( + logger.object(py), + billing.url_route, + billing.endpoint_type, + &self.body, + &stream.chunks, + &self.start, + &self.end, + &stream.first_chunk, + ), + ); + match billed { + Err(error) if error.is_instance_of::(py) => { + error.write_unraisable(py, Some(logger.object(py))); + Ok(()) + } + result => result.map(|_| ()), + } + } + + /// A failure after the stream reached the caller bills the delivered chunks as + /// partial usage. The sync path has no loop to schedule that on, so it falls back to + /// the plain failure handler. + fn stream_failure(&mut self, py: Python<'_>) -> PyResult { + let (Some(logger), Some(error), Some(stream), Some(billing)) = + (&self.logger, &self.error, &self.stream, self.surface.stream) + else { + return Ok(LifecycleStep::Done); + }; + if !self.asynchronous { + return self.dispatch_failure(py); + } + let scheduled = Streaming::Failure.call( + py, + ( + logger.object(py), + billing.endpoint_type, + &self.body, + &stream.chunks, + error, + ), + ); + match scheduled { + Ok(awaitable) => { + self.pending = Some(Pending::AsyncFailure); + Ok(LifecycleStep::Await(awaitable.unbind())) + } + Err(failure) if is_cancellation(py, &failure) => Err(failure), + Err(_) => Ok(LifecycleStep::Done), + } + } + /// The sync failure handler, then the async one for async calls. Ordinary handler /// errors never replace the selected failure or suppress the other family; a /// cancellation does end the call. - fn dispatch_failure(&mut self, py: Python<'_>) -> PyResult { + fn dispatch_failure(&mut self, py: Python<'_>) -> PyResult { let (Some(logger), Some(error)) = (&self.logger, &self.error) else { - return Ok(AdapterStep::Done); + return Ok(LifecycleStep::Done); }; if self.asynchronous && self.internal { - return Ok(AdapterStep::Done); + return Ok(LifecycleStep::Done); } if let Err(failure) = logger.failure(py, error, &self.start, &self.end, false) && is_cancellation(py, &failure) @@ -178,27 +258,27 @@ impl LegacyLogging { return Err(failure); } if !self.asynchronous { - return Ok(AdapterStep::Done); + return Ok(LifecycleStep::Done); } match logger.failure(py, error, &self.start, &self.end, true) { Ok(Some(awaitable)) => { self.pending = Some(Pending::AsyncFailure); - Ok(AdapterStep::Await(awaitable)) + Ok(LifecycleStep::Await(awaitable)) } - Ok(None) => Ok(AdapterStep::Done), + Ok(None) => Ok(LifecycleStep::Done), Err(failure) if is_cancellation(py, &failure) => Err(failure), - Err(_) => Ok(AdapterStep::Done), + Err(_) => Ok(LifecycleStep::Done), } } } -impl CallbackAdapter for LegacyLogging { +impl PythonLifecycle for LegacyLogging { fn begin( &mut self, py: Python<'_>, arguments: Py, started_at: f64, - ) -> PyResult { + ) -> PyResult { self.call.set_kwargs(arguments); self.start = datetime(py, started_at)?; self.internal = is_internal_call(py)?; @@ -212,9 +292,9 @@ impl CallbackAdapter for LegacyLogging { )?; self.logger = Some(result.logger()?); self.call.set_kwargs(result.kwargs()?); - if self.deployment_hooks(py)? { + if self.runs_deployment_hooks() { self.pending = Some(Pending::DeploymentPreCall); - return Ok(AdapterStep::Await(DeploymentHooks::before_call( + return Ok(LifecycleStep::Await(DeploymentHooks::before_call( py, self.call.kwargs(), self.surface.call_type, @@ -228,18 +308,16 @@ impl CallbackAdapter for LegacyLogging { py: Python<'_>, wire: Box, context: &RequestContext, - ) -> PyResult { + ) -> PyResult { let logger = self.logger()?; logger.update_from_kwargs(py, self.call.kwargs(), &wire, context)?; - if !logger.callbacks_needed(py, "payload")? { - logger.record_api_call_start(py)?; - return Ok(AdapterStep::Wire(wire)); - } let body = to_py(py, &wire.body)? .into_bound(py) .cast_into::()?; - for name in context.passthrough_fields.iter() { - if let Some(value) = self.call.lookup(py, name)? { + for (name, sent) in wire.body.as_object().into_iter().flatten() { + if let Some(value) = self.call.lookup(py, name)? + && from_py::(&value).is_ok_and(|caller| caller == *sent) + { body.set_item(name, value)?; } } @@ -249,11 +327,11 @@ impl CallbackAdapter for LegacyLogging { } self.body = Some(body.clone().unbind()); self.headers = Some(headers.clone().unbind()); - let api_key = self.call.lookup(py, "api_key")?; + self.context = Some(context.clone()); self.logger()?.pre_call( py, self.surface.input_description, - api_key.as_ref(), + context.api_key.as_ref().map(|api_key| api_key.expose()), &body, &headers, &wire.url, @@ -262,7 +340,7 @@ impl CallbackAdapter for LegacyLogging { .iter() .map(|(name, value)| Ok((name.extract::()?, value.extract::()?))) .collect::>>()?; - Ok(AdapterStep::Wire(Box::new(WireRequest { + Ok(LifecycleStep::Wire(Box::new(WireRequest { body: from_py(&body)?, headers, ..*wire @@ -274,12 +352,12 @@ impl CallbackAdapter for LegacyLogging { py: Python<'_>, response: Py, timing: Timing, - ) -> PyResult { + ) -> PyResult { self.end = Some(datetime(py, timing.end_time)?); self.response = Some(response); - if self.deployment_hooks(py)? { + if self.runs_deployment_hooks() { self.pending = Some(Pending::DeploymentPostCall); - return Ok(AdapterStep::Await(DeploymentHooks::after_success( + return Ok(LifecycleStep::Await(DeploymentHooks::after_success( py, self.call.kwargs(), &self.response, @@ -289,36 +367,50 @@ impl CallbackAdapter for LegacyLogging { self.finalize(py) } - fn emit( - &mut self, - py: Python<'_>, - event: &CallEvent, - public: Option>, - ) -> PyResult { - match (event, public) { - (CallEvent::ResponseReceived { raw }, _) => { - let logger = self.logger()?; - if logger.callbacks_needed(py, "payload")? { - logger.post_call(py, &raw.body, self.body.as_ref(), self.headers.as_ref())?; - } - Ok(AdapterStep::Done) + fn emit(&mut self, py: Python<'_>, event: LifecycleEvent<'_>) -> PyResult { + match event { + LifecycleEvent::Started { .. } => Ok(LifecycleStep::Done), + LifecycleEvent::Machine(MachineEvent::ResponseReceived { raw }) => { + let api_key = self + .context + .as_ref() + .and_then(|context| context.api_key.as_ref()) + .map(|api_key| api_key.expose()); + self.logger()?.post_call( + py, + &raw.body, + api_key, + self.body.as_ref(), + self.headers.as_ref(), + )?; + Ok(LifecycleStep::Done) } - (CallEvent::Succeeded { timing }, Some(PublicValue::Response(response))) => { + LifecycleEvent::Succeeded { timing, response } => { self.end = Some(datetime(py, timing.end_time)?); self.response = Some(response.clone_ref(py)); - self.dispatch_success(py)?; - Ok(AdapterStep::Done) + match &self.stream { + Some(stream) => self.stream_success(py, stream)?, + None => self.dispatch_success(py)?, + } + Ok(LifecycleStep::Done) } - (CallEvent::Failed { timing, origin }, Some(PublicValue::Error(error))) => { + LifecycleEvent::Failed { + timing, + origin, + error, + } => { self.end = Some(datetime(py, timing.end_time)?); self.error = Some(error.clone_ref(py).into_value(py)); - if *origin == FailureOrigin::Call + if self.stream.is_some() { + return self.stream_failure(py); + } + if origin == FailureOrigin::Call && self.logger.is_some() - && self.deployment_hooks(py)? + && self.runs_deployment_hooks() { let error = self.error.as_ref().ok_or_else(missing_state)?; self.pending = Some(Pending::DeploymentFailure); - return Ok(AdapterStep::Await(DeploymentHooks::after_failure( + return Ok(LifecycleStep::Await(DeploymentHooks::after_failure( py, self.call.kwargs(), error, @@ -327,11 +419,30 @@ impl CallbackAdapter for LegacyLogging { } self.dispatch_failure(py) } - _ => Err(missing_state()), } } - fn resume(&mut self, py: Python<'_>, result: PyResult>) -> PyResult { + fn opened(&mut self, py: Python<'_>) -> PyResult<()> { + if self.surface.stream.is_none() { + return Err(missing_state()); + } + Streaming::Opened.call(py, (self.logger()?.object(py),))?; + self.stream = Some(DeliveredStream { + chunks: PyList::empty(py).unbind(), + first_chunk: None, + }); + Ok(()) + } + + fn delivered(&mut self, py: Python<'_>, chunk: &Py) -> PyResult<()> { + let stream = self.stream.as_mut().ok_or_else(missing_state)?; + if stream.first_chunk.is_none() { + stream.first_chunk = Some(datetime(py, epoch_seconds())?); + } + stream.chunks.bind(py).append(chunk) + } + + fn resume(&mut self, py: Python<'_>, result: PyResult>) -> PyResult { match self.pending.take().ok_or_else(missing_state)? { Pending::DeploymentPreCall => { self.call @@ -345,7 +456,7 @@ impl CallbackAdapter for LegacyLogging { Pending::DeploymentFailure => self.dispatch_failure(py), Pending::AsyncFailure => match result { Err(failure) if is_cancellation(py, &failure) => Err(failure), - _ => Ok(AdapterStep::Done), + _ => Ok(LifecycleStep::Done), }, } } @@ -357,7 +468,8 @@ impl CallbackAdapter for LegacyLogging { error.write_unraisable(py, None); } self.body = None; - self.headers = None; + self.context = None; + self.stream = None; } fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { @@ -369,8 +481,11 @@ impl CallbackAdapter for LegacyLogging { visit.call(&self.end)?; visit.call(&self.response)?; visit.call(&self.error)?; - visit.call(&self.body)?; - visit.call(&self.headers) + if let Some(stream) = &self.stream { + visit.call(&stream.chunks)?; + visit.call(&stream.first_chunk)?; + } + visit.call(&self.body) } } diff --git a/litellm-rust/crates/callbacks-legacy/src/call.rs b/litellm-rust/crates/callbacks-legacy/src/call.rs index 59090ee8d60..b37790f60a8 100644 --- a/litellm-rust/crates/callbacks-legacy/src/call.rs +++ b/litellm-rust/crates/callbacks-legacy/src/call.rs @@ -3,8 +3,8 @@ //! lifetime. No other callback host has that obligation, which is why nothing outside //! this crate holds them. -use litellm_callbacks::{machine::Machine, route::Route}; -use litellm_host_python::{RouteHost, run_call}; +use litellm_host::{machine::Machine, route::Route}; +use litellm_host_python::{RouteHost, lookup, run_call}; use pyo3::{ gc::{PyTraverseError, PyVisit}, prelude::*, @@ -63,21 +63,6 @@ impl PublicCall { } } -/// The caller's own object for a public argument, as every legacy reader resolves it: the -/// keyword if given, even an explicit `None`, else the bound request's attribute. A route -/// host projecting from the prepared keyword view uses the same rule, so the callbacks -/// and the provider see one object per argument. -pub fn lookup<'py>( - kwargs: &Bound<'py, PyDict>, - request: &Bound<'py, PyAny>, - name: &str, -) -> PyResult>> { - if let Some(value) = kwargs.get_item(name)? { - return Ok(Some(value)); - } - request.getattr_opt(name) -} - /// Runs one native call under the legacy `Logging` contract: the route host projects from /// the keyword view the contract prepares, and the contract observes the call. pub fn run_legacy_call( @@ -121,32 +106,6 @@ mod tests { (call, locals) } - #[test] - fn lookup_prefers_the_keyword_even_when_none_and_falls_back_to_the_request() { - Python::initialize(); - Python::attach(|py| { - let (call, locals) = capture( - py, - c" -key = object() -document = {'type': 'document_url'} -class Request: - api_key = 'from-request' - api_base = 'from-request' - document = document -request = Request() -kwargs = {'api_key': key, 'api_base': None} -", - ); - let key = locals.get_item("key").unwrap().unwrap(); - let document = locals.get_item("document").unwrap().unwrap(); - assert!(call.lookup(py, "api_key").unwrap().unwrap().is(&key)); - assert!(call.lookup(py, "api_base").unwrap().unwrap().is_none()); - assert!(call.lookup(py, "document").unwrap().unwrap().is(&document)); - assert!(call.lookup(py, "model").unwrap().is_none()); - }); - } - #[test] fn capture_copies_the_keyword_dict_without_copying_its_values() { Python::initialize(); diff --git a/litellm-rust/crates/callbacks-legacy/src/callbacks.rs b/litellm-rust/crates/callbacks-legacy/src/callbacks.rs index aa586013e75..5f04224e6d7 100644 --- a/litellm-rust/crates/callbacks-legacy/src/callbacks.rs +++ b/litellm-rust/crates/callbacks-legacy/src/callbacks.rs @@ -2,15 +2,14 @@ //! the deferred and worker-submitted success paths, and the sync-callbacks-for-async-calls //! duplication. All of it expires with the legacy callback contract. -use litellm_callbacks::event::{RequestContext, WireRequest}; +use litellm_host::event::{RequestContext, WireRequest}; use litellm_host_python::to_py; use pyo3::{exceptions::PyBaseException, prelude::*, types::PyDict}; +use crate::legacy_python::{Logging, Wrapper}; use crate::logger::PythonLogger; pub trait LegacyCallbacks { - fn callbacks_needed(&self, py: Python<'_>, phase: &str) -> PyResult; - /// `Logging.update_from_kwargs`: what the logger is told about the request it is /// about to see, with consumed credentials redacted. fn update_from_kwargs( @@ -21,24 +20,23 @@ pub trait LegacyCallbacks { context: &RequestContext, ) -> PyResult<()>; - fn record_api_call_start(&self, py: Python<'_>) -> PyResult<()>; - - /// `Logging.pre_call`, or its payload-free shortcut when no input callback listens. + /// `Logging.pre_call`. fn pre_call( &self, py: Python<'_>, input: &str, - api_key: Option<&Bound<'_, PyAny>>, + api_key: Option<&str>, body: &Bound<'_, PyDict>, headers: &Bound<'_, PyDict>, url: &str, ) -> PyResult<()>; - /// `Logging.post_call`, or its payload-free shortcut when no input callback listens. + /// `Logging.post_call`. fn post_call( &self, py: Python<'_>, original_response: &str, + api_key: Option<&str>, body: Option<&Py>, headers: Option<&Py>, ) -> PyResult<()>; @@ -82,16 +80,6 @@ pub trait LegacyCallbacks { } impl LegacyCallbacks for PythonLogger { - fn callbacks_needed(&self, py: Python<'_>, phase: &str) -> PyResult { - if !self.bridge_owned() { - return Ok(true); - } - py.import("litellm.rust_bridge.legacy_callbacks")? - .getattr("callbacks_needed")? - .call1((self.object(py), phase))? - .extract() - } - fn update_from_kwargs( &self, py: Python<'_>, @@ -100,18 +88,13 @@ impl LegacyCallbacks for PythonLogger { context: &RequestContext, ) -> PyResult<()> { let secret_fields: Vec<&str> = context.secret_fields.iter().map(String::as_str).collect(); - let update = PyDict::new(py); - update.set_item("kwargs", redact(py, kwargs.bind(py), &secret_fields)?)?; - update.set_item("model", &context.model)?; - update.set_item( - "optional_params", - redact( - py, - &to_py(py, &context.optional_params)? - .into_bound(py) - .cast_into::()?, - &secret_fields, - )?, + let redacted_kwargs = redact(py, kwargs.bind(py), &secret_fields)?; + let optional_params = redact( + py, + &to_py(py, &context.optional_params)? + .into_bound(py) + .cast_into::()?, + &secret_fields, )?; let params = PyDict::new(py); params.set_item( @@ -131,15 +114,17 @@ impl LegacyCallbacks for PythonLogger { params.set_item(name, value)?; } } - update.set_item("litellm_params", params)?; - update.set_item("custom_llm_provider", &context.custom_llm_provider)?; - self.object(py) - .call_method("update_from_kwargs", (), Some(&update))?; - Ok(()) - } - - fn record_api_call_start(&self, py: Python<'_>) -> PyResult<()> { - self.object(py).call_method0("record_api_call_start_time")?; + Logging::Update.call( + py, + ( + self.object(py), + redacted_kwargs, + &context.model, + optional_params, + params, + &context.custom_llm_provider, + ), + )?; Ok(()) } @@ -147,7 +132,7 @@ impl LegacyCallbacks for PythonLogger { &self, py: Python<'_>, input: &str, - api_key: Option<&Bound<'_, PyAny>>, + api_key: Option<&str>, body: &Bound<'_, PyDict>, headers: &Bound<'_, PyDict>, url: &str, @@ -156,17 +141,7 @@ impl LegacyCallbacks for PythonLogger { additional.set_item("complete_input_dict", body)?; additional.set_item("headers", headers)?; additional.set_item("api_base", url)?; - let kwargs = PyDict::new(py); - kwargs.set_item("input", input)?; - kwargs.set_item("api_key", api_key)?; - kwargs.set_item("additional_args", &additional)?; - if self.callbacks_needed(py, "input")? { - self.object(py).call_method("pre_call", (), Some(&kwargs))?; - } else { - self.object(py) - .call_method("_pre_call", (), Some(&kwargs))?; - self.record_api_call_start(py)?; - } + Logging::PreCall.call(py, (self.object(py), input, api_key, &additional))?; Ok(()) } @@ -174,37 +149,30 @@ impl LegacyCallbacks for PythonLogger { &self, py: Python<'_>, original_response: &str, + api_key: Option<&str>, body: Option<&Py>, headers: Option<&Py>, ) -> PyResult<()> { let additional = PyDict::new(py); additional.set_item("complete_input_dict", body)?; additional.set_item("headers", headers)?; - if self.callbacks_needed(py, "input")? { - let kwargs = PyDict::new(py); - kwargs.set_item("original_response", original_response)?; - kwargs.set_item("additional_args", &additional)?; - self.object(py) - .call_method("post_call", (), Some(&kwargs))?; - } else { - let response = py - .import("json")? - .call_method1("dumps", (original_response,))?; - self.object(py).call_method1( - "record_post_call", - (response, py.None(), py.None(), additional), - )?; - } + Logging::PostCall.call( + py, + (self.object(py), original_response, api_key, &additional), + )?; Ok(()) } + fn defers_async_logging(&self, py: Python<'_>) -> bool { - self.object(py) - .getattr("_defer_async_logging") - .is_ok_and(|value| value.is_truthy().unwrap_or(false)) + Logging::DefersAsync + .call(py, (self.object(py),)) + .and_then(|value| value.extract()) + .unwrap_or(false) } fn defer_success(&self, py: Python<'_>, pending: &Bound<'_, PyAny>) -> PyResult<()> { - self.object(py).setattr("_native_pending_logging", pending) + Logging::DeferSuccess.call(py, (self.object(py), pending))?; + Ok(()) } fn sync_success_for_async_call( @@ -214,13 +182,7 @@ impl LegacyCallbacks for PythonLogger { start: &Py, end: &Option>, ) -> PyResult<()> { - if !self.callbacks_needed(py, "sync_success_async")? { - return Ok(()); - } - self.object(py).call_method1( - "handle_sync_success_callbacks_for_async_calls", - (response, start, end), - )?; + Logging::SyncSuccessForAsyncCall.call(py, (self.object(py), response, start, end))?; Ok(()) } @@ -232,34 +194,11 @@ impl LegacyCallbacks for PythonLogger { end: &Option>, asynchronous: bool, ) -> PyResult>> { - if !self.callbacks_needed( - py, - if asynchronous { - "async_failure" - } else { - "sync_failure" - }, - )? { - py.import("litellm.rust_bridge.legacy_callbacks")? - .getattr("failure_bookkeeping")? - .call1((self.object(py), error, start, end, asynchronous))?; - return Ok(None); - } - let trace = py - .import("traceback")? - .getattr("format_exception")? - .call1((error,))?; - let trace = pyo3::types::PyString::new(py, "").call_method1("join", (trace,))?; - let value = self.object(py).call_method1( - if asynchronous { - "async_failure_handler" - } else { - "failure_handler" - }, - (error, trace, start, end), - )?; + let value = + Logging::FailureHandler.call(py, (self.object(py), error, start, end, asynchronous))?; Ok(asynchronous.then(|| value.unbind())) } + fn submit_success( &self, py: Python<'_>, @@ -267,22 +206,7 @@ impl LegacyCallbacks for PythonLogger { start: &Py, end: &Option>, ) -> PyResult<()> { - if !self.callbacks_needed(py, "sync_success")? { - return self.success_bookkeeping(py, response, start, end, false); - } - let context = py.import("contextvars")?.call_method0("copy_context")?; - py.import("litellm.litellm_core_utils.litellm_logging")? - .getattr("executor")? - .call_method1( - "submit", - ( - context.getattr("run")?, - self.object(py).getattr("success_handler")?, - response, - start, - end, - ), - )?; + Logging::SubmitSuccess.call(py, (self.object(py), response, start, end))?; Ok(()) } @@ -293,18 +217,9 @@ impl LegacyCallbacks for PythonLogger { start: &Py, end: &Option>, ) -> PyResult<()> { - if !self.callbacks_needed(py, "async_success")? { - return self.success_bookkeeping(py, response, start, end, true); - } - let context = py.import("contextvars")?.call_method0("copy_context")?; - let worker = py - .import("litellm.litellm_core_utils.logging_worker")? - .getattr("GLOBAL_LOGGING_WORKER")? - .getattr("ensure_initialized_and_enqueue")?; - let coroutine = self - .object(py) - .call_method1("async_success_handler", (response, start, end))?; - let enqueue = context.call_method1("run", (worker, &coroutine)); + let coroutine = + Logging::AsyncSuccessHandler.call(py, (self.object(py), response, start, end))?; + let enqueue = Logging::Enqueue.call(py, (&coroutine,)); if enqueue.is_err() && let Err(error) = coroutine.call_method0("close") { @@ -315,14 +230,7 @@ impl LegacyCallbacks for PythonLogger { } fn custom_pricing_fields(py: Python<'_>) -> PyResult> { - py.import("litellm.types.utils")? - .getattr("CustomPricingLiteLLMParams")? - .getattr("model_fields")? - .cast_into::()? - .keys() - .iter() - .map(|name| name.extract::()) - .collect() + Logging::CustomPricingFields.call(py, ())?.extract() } fn redact( @@ -347,58 +255,5 @@ fn redact( /// Proxy-internal calls skip the legacy success fan-out. pub fn is_internal_call(py: Python<'_>) -> PyResult { - py.import("litellm._internal_context")? - .getattr("is_internal_call")? - .call_method0("get")? - .extract() -} - -#[cfg(test)] -mod tests { - use pyo3::types::PyDict; - - use super::*; - - fn logger_whose_registries_need_no_input(py: Python<'_>, bridge_owned: bool) -> PythonLogger { - let locals = PyDict::new(py); - py.run( - c" -import sys -import types -for name in ('litellm', 'litellm.rust_bridge', 'litellm.rust_bridge.legacy_callbacks'): - sys.modules.setdefault(name, types.ModuleType(name)) -legacy = sys.modules['litellm.rust_bridge.legacy_callbacks'] -legacy.callbacks_needed = lambda logger, phase: logger.needed.get(phase, True) -class Logger: - needed = {'input': False} -logger = Logger() -", - Some(&locals), - Some(&locals), - ) - .unwrap(); - PythonLogger::new( - locals.get_item("logger").unwrap().unwrap().unbind(), - bridge_owned, - ) - } - - #[test] - fn a_caller_owned_logger_is_observed_in_full() { - Python::initialize(); - Python::attach(|py| { - let logger = logger_whose_registries_need_no_input(py, false); - assert!(logger.callbacks_needed(py, "input").unwrap()); - }); - } - - #[test] - fn a_bridge_owned_logger_is_elided_where_no_registry_needs_it() { - Python::initialize(); - Python::attach(|py| { - let logger = logger_whose_registries_need_no_input(py, true); - assert!(!logger.callbacks_needed(py, "input").unwrap()); - assert!(logger.callbacks_needed(py, "payload").unwrap()); - }); - } + Wrapper::IsInternalCall.call(py, ())?.extract() } diff --git a/litellm-rust/crates/callbacks-legacy/src/legacy_python.rs b/litellm-rust/crates/callbacks-legacy/src/legacy_python.rs new file mode 100644 index 00000000000..7f5c77c1735 --- /dev/null +++ b/litellm-rust/crates/callbacks-legacy/src/legacy_python.rs @@ -0,0 +1,183 @@ +use pyo3::prelude::*; +use strum::{IntoStaticStr, VariantArray}; + +const MODULE: &str = "litellm.rust_bridge.legacy_callbacks"; + +/// Every litellm Python internal the native call still borrows, grouped by the subsystem it +/// belongs to. Rust drives the call; these exist only so behaviour that Python owns today +/// (span tracking, the standard logging payload, spend, callback fan-out) keeps working. +/// A group is deleted once Rust owns that subsystem, so this enum only shrinks. Calling a +/// user's own callback is not borrowing and does not belong here. +/// +/// `litellm/rust_bridge/legacy_callbacks.py` is the only Python module behind it, and +/// `python_contract.json` pins each function's parameters on both sides. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(crate) enum LegacyPython { + Wrapper(Wrapper), + Logging(Logging), + DeploymentHooks(DeploymentHooks), + Streaming(Streaming), +} + +/// The `@client` wrapper around the call: `function_setup`, limits, credentials, +/// response metadata and the correlation context. +#[derive(Clone, Copy, Debug, IntoStaticStr, PartialEq, Eq, VariantArray)] +pub(crate) enum Wrapper { + #[strum(serialize = "setup")] + Setup, + #[strum(serialize = "check_limits")] + CheckLimits, + #[strum(serialize = "credential_list")] + CredentialList, + #[strum(serialize = "warn_unknown_credential")] + WarnUnknownCredential, + #[strum(serialize = "is_internal_call")] + IsInternalCall, + #[strum(serialize = "finalize")] + Finalize, + #[strum(serialize = "restore_context")] + RestoreContext, +} + +/// litellm's `Logging` object and the sync and async callback fan-out behind it. +#[derive(Clone, Copy, Debug, IntoStaticStr, PartialEq, Eq, VariantArray)] +pub(crate) enum Logging { + #[strum(serialize = "custom_pricing_fields")] + CustomPricingFields, + #[strum(serialize = "update_logging")] + Update, + #[strum(serialize = "pre_call")] + PreCall, + #[strum(serialize = "post_call")] + PostCall, + #[strum(serialize = "defers_async_logging")] + DefersAsync, + #[strum(serialize = "defer_success")] + DeferSuccess, + #[strum(serialize = "sync_success_for_async_call")] + SyncSuccessForAsyncCall, + #[strum(serialize = "submit_success")] + SubmitSuccess, + #[strum(serialize = "async_success_handler")] + AsyncSuccessHandler, + #[strum(serialize = "enqueue_logging")] + Enqueue, + #[strum(serialize = "failure_handler")] + FailureHandler, +} + +/// The `litellm.utils` fan-outs that run every callback's deployment hook. +#[derive(Clone, Copy, Debug, IntoStaticStr, PartialEq, Eq, VariantArray)] +pub(crate) enum DeploymentHooks { + #[strum(serialize = "before_deployment_call")] + BeforeDeploymentCall, + #[strum(serialize = "after_deployment_success")] + AfterDeploymentSuccess, + #[strum(serialize = "after_deployment_failure")] + AfterDeploymentFailure, +} + +/// The Messages stream iterator's logging: the stream flag, the end-of-stream billing +/// from the delivered chunks, and the partial-usage failure path. +#[derive(Clone, Copy, Debug, IntoStaticStr, PartialEq, Eq, VariantArray)] +pub(crate) enum Streaming { + #[strum(serialize = "stream_opened")] + Opened, + #[strum(serialize = "stream_success")] + Success, + #[strum(serialize = "stream_failure")] + Failure, +} + +impl LegacyPython { + fn name(self) -> &'static str { + match self { + Self::Wrapper(function) => function.into(), + Self::Logging(function) => function.into(), + Self::DeploymentHooks(function) => function.into(), + Self::Streaming(function) => function.into(), + } + } + + pub(crate) fn call<'py, A>(self, py: Python<'py>, args: A) -> PyResult> + where + A: pyo3::call::PyCallArgs<'py>, + { + py.import(MODULE)?.getattr(self.name())?.call1(args) + } +} + +impl Wrapper { + pub(crate) fn call<'py, A>(self, py: Python<'py>, args: A) -> PyResult> + where + A: pyo3::call::PyCallArgs<'py>, + { + LegacyPython::Wrapper(self).call(py, args) + } +} + +impl Logging { + pub(crate) fn call<'py, A>(self, py: Python<'py>, args: A) -> PyResult> + where + A: pyo3::call::PyCallArgs<'py>, + { + LegacyPython::Logging(self).call(py, args) + } +} + +impl Streaming { + pub(crate) fn call<'py, A>(self, py: Python<'py>, args: A) -> PyResult> + where + A: pyo3::call::PyCallArgs<'py>, + { + LegacyPython::Streaming(self).call(py, args) + } +} + +impl DeploymentHooks { + pub(crate) fn call<'py, A>(self, py: Python<'py>, args: A) -> PyResult> + where + A: pyo3::call::PyCallArgs<'py>, + { + LegacyPython::DeploymentHooks(self).call(py, args) + } +} + +#[cfg(test)] +mod tests { + use std::collections::BTreeSet; + + use strum::VariantArray; + + use super::{DeploymentHooks, LegacyPython, Logging, Streaming, Wrapper}; + use crate::test_support::PYTHON_CONTRACT; + + #[test] + fn every_borrowed_function_is_in_the_python_contract() { + let contract: serde_json::Map = + serde_json::from_str(PYTHON_CONTRACT).unwrap(); + let declared: BTreeSet<&str> = contract.keys().map(String::as_str).collect(); + let called: Vec<&str> = Wrapper::VARIANTS + .iter() + .map(|&function| LegacyPython::Wrapper(function)) + .chain( + Logging::VARIANTS + .iter() + .map(|&function| LegacyPython::Logging(function)), + ) + .chain( + DeploymentHooks::VARIANTS + .iter() + .map(|&function| LegacyPython::DeploymentHooks(function)), + ) + .chain( + Streaming::VARIANTS + .iter() + .map(|&function| LegacyPython::Streaming(function)), + ) + .map(LegacyPython::name) + .collect(); + assert_eq!(called.len(), declared.len(), "a function is borrowed twice"); + assert_eq!(called.into_iter().collect::>(), declared); + } +} diff --git a/litellm-rust/crates/callbacks-legacy/src/lib.rs b/litellm-rust/crates/callbacks-legacy/src/lib.rs index 06783ac255d..eaa1a8b714e 100644 --- a/litellm-rust/crates/callbacks-legacy/src/lib.rs +++ b/litellm-rust/crates/callbacks-legacy/src/lib.rs @@ -2,7 +2,7 @@ //! sync and async callback registries it fans out to, the deployment hooks, the deferred //! proxy release, and the kwargs rewrites the wrapper makes on the way in (credential-name //! inheritance, budget and retry-count limits). All of it sits behind one -//! [`CallbackAdapter`](litellm_host_python::CallbackAdapter), so the driver, the routes and +//! [`PythonLifecycle`](litellm_host_python::PythonLifecycle), so the driver, the routes and //! core never learn which Python object is on the other end. //! //! Legacy callbacks receive the caller's own objects and may mutate them. [`PublicCall`] @@ -13,6 +13,7 @@ mod adapter; mod call; mod callbacks; mod deferred; +mod legacy_python; mod logger; mod preparation; #[cfg(test)] @@ -20,8 +21,8 @@ mod preparation; mod test_support; pub(crate) use adapter::LegacyLogging; -pub use adapter::LegacySurface; -pub use call::{PublicCall, lookup, run_legacy_call}; +pub use adapter::{LegacySurface, PassThroughStream}; +pub use call::{PublicCall, run_legacy_call}; pub(crate) use callbacks::{LegacyCallbacks, is_internal_call}; pub(crate) use logger::{DeploymentHooks, PythonLogger, finalize, setup}; pub(crate) use preparation::prepare; diff --git a/litellm-rust/crates/callbacks-legacy/src/logger.rs b/litellm-rust/crates/callbacks-legacy/src/logger.rs index a0e525000b8..061941f05b9 100644 --- a/litellm-rust/crates/callbacks-legacy/src/logger.rs +++ b/litellm-rust/crates/callbacks-legacy/src/logger.rs @@ -5,34 +5,25 @@ use pyo3::{ types::{PyDict, PyTuple}, }; -/// The `Logging` instance one call fans out through, and who owns it. A logger the caller -/// handed in is observed in full, because the caller reads it after the call; one this -/// crate built through `function_setup` is elided wherever no registry needs it. +use crate::legacy_python::{self, Wrapper}; + +/// The `Logging` instance one call fans out through. pub struct PythonLogger { object: Py, - bridge_owned: bool, } impl PythonLogger { - pub(crate) fn new(object: Py, bridge_owned: bool) -> Self { - Self { - object, - bridge_owned, - } + pub(crate) fn new(object: Py) -> Self { + Self { object } } pub(crate) fn object<'py>(&self, py: Python<'py>) -> &Bound<'py, PyAny> { self.object.bind(py) } - pub(crate) fn bridge_owned(&self) -> bool { - self.bridge_owned - } - pub fn clone_ref(&self, py: Python<'_>) -> Self { Self { object: self.object.clone_ref(py), - bridge_owned: self.bridge_owned, } } @@ -40,34 +31,17 @@ impl PythonLogger { visit.call(&self.object) } - pub fn success_bookkeeping( - &self, - py: Python<'_>, - response: &Option>, - start: &Py, - end: &Option>, - asynchronous: bool, - ) -> PyResult<()> { - py.import("litellm.rust_bridge.legacy_callbacks")? - .getattr("success_bookkeeping")? - .call1((self.object(py), response, start, end, asynchronous))?; - Ok(()) - } - pub fn restore_context(&self, py: Python<'_>) -> PyResult<()> { - py.import("litellm.utils")? - .getattr("_restore_correlation_context_if_supported")? - .call1((self.object(py),))?; + Wrapper::RestoreContext.call(py, (self.object(py),))?; Ok(()) } } -/// A bare Python object was not obtained from `setup`, so it is caller-owned. impl FromPyObject<'_, '_> for PythonLogger { type Error = PyErr; fn extract(object: Borrowed<'_, '_, PyAny>) -> PyResult { - Ok(Self::new(object.to_owned().unbind(), false)) + Ok(Self::new(object.to_owned().unbind())) } } @@ -75,9 +49,7 @@ pub struct SetupResult<'py>(Bound<'py, PyAny>); impl SetupResult<'_> { pub fn logger(&self) -> PyResult { - let object = self.0.getattr("logger")?.unbind(); - let bridge_owned = self.0.getattr("bridge_owned")?.extract()?; - Ok(PythonLogger::new(object, bridge_owned)) + Ok(PythonLogger::new(self.0.getattr("logger")?.unbind())) } pub fn kwargs(&self) -> PyResult> { @@ -93,9 +65,8 @@ pub fn setup<'py>( start: &Py, asynchronous: bool, ) -> PyResult> { - py.import("litellm.rust_bridge.legacy_callbacks")? - .getattr("setup")? - .call1((call_type, args, kwargs, start, asynchronous)) + Wrapper::Setup + .call(py, (call_type, args, kwargs, start, asynchronous)) .map(SetupResult) } @@ -107,30 +78,20 @@ pub fn finalize( start: &Py, end: &Option>, ) -> PyResult<()> { - py.import("litellm.rust_bridge.legacy_callbacks")? - .getattr("finalize")? - .call1((response, logger.object(py), kwargs, start, end))?; + Wrapper::Finalize.call(py, (response, logger.object(py), kwargs, start, end))?; Ok(()) } pub struct DeploymentHooks; impl DeploymentHooks { - pub fn needed(py: Python<'_>) -> PyResult { - py.import("litellm.rust_bridge.legacy_callbacks")? - .getattr("deployment_callbacks_needed")? - .call0()? - .extract() - } - pub fn before_call( py: Python<'_>, kwargs: &Py, call_type: &str, ) -> PyResult> { - py.import("litellm.utils")? - .getattr("async_pre_call_deployment_hook")? - .call1((kwargs, call_type)) + legacy_python::DeploymentHooks::BeforeDeploymentCall + .call(py, (kwargs, call_type)) .map(Bound::unbind) } @@ -140,9 +101,8 @@ impl DeploymentHooks { response: &Option>, call_type: &str, ) -> PyResult> { - py.import("litellm.utils")? - .getattr("async_post_call_success_deployment_hook")? - .call1((kwargs, response, call_type)) + legacy_python::DeploymentHooks::AfterDeploymentSuccess + .call(py, (kwargs, response, call_type)) .map(Bound::unbind) } @@ -152,9 +112,8 @@ impl DeploymentHooks { error: &Py, call_type: &str, ) -> PyResult> { - py.import("litellm.utils")? - .getattr("async_post_call_failure_deployment_hook")? - .call1((kwargs, error, call_type)) + legacy_python::DeploymentHooks::AfterDeploymentFailure + .call(py, (kwargs, error, call_type)) .map(Bound::unbind) } } @@ -185,10 +144,6 @@ class Setup: reads.append('logger') return logger @property - def bridge_owned(self): - reads.append('bridge_owned') - return True - @property def kwargs(self): reads.append('kwargs') return [] @@ -206,7 +161,6 @@ result = Setup() .object(py) .is(locals.get_item("logger").unwrap().unwrap()) ); - assert!(logger.bridge_owned()); assert!( result .kwargs() @@ -220,17 +174,8 @@ result = Setup() .unwrap() .extract::>() .unwrap(), - ["logger", "bridge_owned", "kwargs"] + ["logger", "kwargs"] ); }); } - - #[test] - fn a_logger_extracted_from_a_bare_object_is_caller_owned() { - Python::initialize(); - Python::attach(|py| { - let logger: PythonLogger = py.None().into_bound(py).extract().unwrap(); - assert!(!logger.bridge_owned()); - }); - } } diff --git a/litellm-rust/crates/callbacks-legacy/src/preparation.rs b/litellm-rust/crates/callbacks-legacy/src/preparation.rs index 981b1702f2e..fa1ff9acd4d 100644 --- a/litellm-rust/crates/callbacks-legacy/src/preparation.rs +++ b/litellm-rust/crates/callbacks-legacy/src/preparation.rs @@ -3,6 +3,8 @@ use pyo3::{ types::{PyDict, PyList}, }; +use crate::legacy_python::Wrapper; + struct CredentialEntry<'py>(Bound<'py, PyAny>); impl<'py> CredentialEntry<'py> { @@ -22,18 +24,19 @@ pub fn prepare<'py>( ) -> PyResult> { let arguments = kwargs.copy()?; arguments.set_item("litellm_logging_obj", logger.object(py))?; - let litellm = py.import("litellm")?; - inherit_credentials(py, &litellm, &arguments)?; - py.import("litellm.rust_bridge.legacy_callbacks")? - .getattr("check_limits")? - .call1((&arguments,))?; + inherit_credentials(py, &arguments, || { + Ok(Wrapper::CredentialList + .call(py, ())? + .cast_into::()?) + })?; + Wrapper::CheckLimits.call(py, (&arguments,))?; Ok(arguments) } -fn inherit_credentials( - py: Python<'_>, - litellm: &Bound<'_, PyModule>, - arguments: &Bound<'_, PyDict>, +fn inherit_credentials<'py>( + py: Python<'py>, + arguments: &Bound<'py, PyDict>, + credential_list: impl FnOnce() -> PyResult>, ) -> PyResult<()> { let Some(requested) = arguments .get_item("litellm_credential_name")? @@ -45,16 +48,13 @@ fn inherit_credentials( return Ok(()); } let requested: String = requested.extract()?; - let credentials = litellm.getattr("credential_list")?.cast_into::()?; + let credentials = credential_list()?; let names = credentials .iter() .map(|credential| CredentialEntry(credential).name()) .collect::>>()?; let Some(index) = names.iter().position(|name| *name == requested) else { - py.import("litellm._logging")?.getattr("verbose_logger")?.call_method1( - "warning", - ("litellm_credential_name=%s matched none of the %d loaded credentials; the request runs without it", requested, names.len()), - )?; + Wrapper::WarnUnknownCredential.call(py, (requested, names.len()))?; return Ok(()); }; let selected = CredentialEntry(credentials.get_item(index)?); @@ -80,19 +80,19 @@ mod tests { } fn inherit(py: Python<'_>, locals: &Bound<'_, PyDict>) -> PyResult<()> { - let litellm = PyModule::new(py, "credential_host")?; - litellm.setattr( - "credential_list", - locals.get_item("credentials").unwrap().unwrap(), - )?; inherit_credentials( py, - &litellm, &locals .get_item("arguments") .unwrap() .unwrap() .cast_into::()?, + || { + Ok(locals + .get_item("credentials")? + .unwrap() + .cast_into::()?) + }, ) } @@ -304,11 +304,11 @@ arguments = {'litellm_credential_name': 'ocr-test'} fn falsy_credential_names_return_before_loading_credentials() { Python::initialize(); Python::attach(|py| { - let litellm = PyModule::new(py, "credential_host").unwrap(); for name in [py.None(), py.eval(c"''", None, None).unwrap().unbind()] { let arguments = PyDict::new(py); arguments.set_item("litellm_credential_name", name).unwrap(); - inherit_credentials(py, &litellm, &arguments).unwrap(); + inherit_credentials(py, &arguments, || panic!("credentials must not be loaded")) + .unwrap(); } }); } diff --git a/litellm-rust/crates/callbacks-legacy/tests/deferred.rs b/litellm-rust/crates/callbacks-legacy/tests/deferred.rs index 3daea8840d8..289ea1b2e7f 100644 --- a/litellm-rust/crates/callbacks-legacy/tests/deferred.rs +++ b/litellm-rust/crates/callbacks-legacy/tests/deferred.rs @@ -16,7 +16,7 @@ fn defer<'py>(py: Python<'py>, script: &CStr) -> Bound<'py, PyDict> { py, PendingLogging { pending: Some(PendingSuccess { - logger: PythonLogger::new(local(&locals, "logger").unbind(), true), + logger: PythonLogger::new(local(&locals, "logger").unbind()), response: Some(local(&locals, "response").unbind()), start: py.None(), end: Some(py.None()), @@ -79,22 +79,6 @@ assert logger.calls == [], logger.calls }); } -#[test] -fn a_release_after_the_async_callbacks_went_away_only_keeps_the_books() { - Python::initialize(); - Python::attach(|py| { - let locals = defer(py, c"logger.needed = {'async_success': False}"); - run( - py, - &locals, - c" -pending.release(True) -assert logger.calls == [('success_bookkeeping', True)], logger.calls -", - ); - }); -} - #[rstest] #[case::ordinary_error(c"RuntimeError('queue full')", false)] #[case::cancellation(c"asyncio.CancelledError()", true)] diff --git a/litellm-rust/crates/callbacks-legacy/tests/deployment_hooks.rs b/litellm-rust/crates/callbacks-legacy/tests/deployment_hooks.rs index 3ceda4441a7..52c5e47f83f 100644 --- a/litellm-rust/crates/callbacks-legacy/tests/deployment_hooks.rs +++ b/litellm-rust/crates/callbacks-legacy/tests/deployment_hooks.rs @@ -1,7 +1,7 @@ use std::ffi::CStr; -use litellm_callbacks::event::{CallEvent, FailureOrigin, Timing}; -use litellm_host_python::{AdapterStep, CallbackAdapter, PublicValue}; +use litellm_host::event::{FailureOrigin, Timing}; +use litellm_host_python::{LifecycleEvent, LifecycleStep, PythonLifecycle}; use pyo3::exceptions::asyncio::CancelledError; use pyo3::prelude::*; use pyo3::types::PyDict; @@ -24,7 +24,7 @@ fn begin<'py>( py: Python<'py>, locals: &Bound<'py, PyDict>, asynchronous: bool, -) -> (LegacyLogging, AdapterStep) { +) -> (LegacyLogging, LifecycleStep) { let mut logging = legacy_call(py, locals, asynchronous); let kwargs = local(locals, "kwargs") .cast_into::() @@ -34,15 +34,15 @@ fn begin<'py>( (logging, step) } -fn arguments<'py>(py: Python<'py>, step: AdapterStep) -> Bound<'py, PyDict> { - let AdapterStep::Arguments(arguments) = step else { +fn arguments<'py>(py: Python<'py>, step: LifecycleStep) -> Bound<'py, PyDict> { + let LifecycleStep::Arguments(arguments) = step else { panic!("expected the prepared arguments"); }; arguments.into_bound(py) } -fn awaits_deployment_hook(step: &AdapterStep) -> bool { - matches!(step, AdapterStep::Await(_)) +fn awaits_deployment_hook(step: &LifecycleStep) -> bool { + matches!(step, LifecycleStep::Await(_)) } #[rstest] @@ -97,6 +97,43 @@ assert checked is prepared }); } +#[rstest] +#[case::synchronous(false)] +#[case::asynchronous(true)] +fn a_keyword_the_bridge_never_reads_reaches_every_reader_as_the_callers_object( + #[case] asynchronous: bool, +) { + Python::initialize(); + Python::attach(|py| { + let locals = namespace( + py, + c" +opaque = object() +hooked = [] +logger.hooks = {'pre': lambda kwargs: hooked.append(kwargs['vendor_extension']) or kwargs} +kwargs = {'logger': logger, 'vendor_extension': opaque} +", + ); + let (mut logging, step) = begin(py, &locals, asynchronous); + let step = match step { + LifecycleStep::Await(hook_result) => logging.resume(py, Ok(hook_result)).unwrap(), + step => step, + }; + locals.set_item("prepared", arguments(py, step)).unwrap(); + locals.set_item("asynchronous", asynchronous).unwrap(); + run( + py, + &locals, + c" +assert prepared['vendor_extension'] is opaque +[checked] = [value for name, value in logger.calls if name == 'check_limits'] +assert checked['vendor_extension'] is opaque +assert hooked == ([opaque] if asynchronous else []), hooked +", + ); + }); +} + #[test] fn response_returned_by_the_post_call_hook_is_finalized_and_returned() { Python::initialize(); @@ -121,7 +158,7 @@ logger.hooks = {'pre': lambda kwargs: kwargs} let step = logging .resume(py, Ok(local(&locals, "replacement").unbind())) .unwrap(); - let AdapterStep::Response(returned) = step else { + let LifecycleStep::Response(returned) = step else { panic!("expected the finalized response"); }; assert!(returned.bind(py).is(local(&locals, "replacement"))); @@ -180,13 +217,12 @@ fn failure_callbacks_run_after_the_failure_hook_however_it_ends(#[case] cancelle .resume(py, Ok(local(&locals, "kwargs").unbind())) .unwrap(); let failure = PyErr::from_value(local(&locals, "failure")); - let failed = CallEvent::Failed { + let failed = LifecycleEvent::Failed { timing: TIMING, origin: FailureOrigin::Call, + error: &failure, }; - let step = logging - .emit(py, &failed, Some(PublicValue::Error(&failure))) - .unwrap(); + let step = logging.emit(py, failed).unwrap(); assert!(awaits_deployment_hook(&step)); let hook_result = if cancelled { Err(CancelledError::new_err("cancelled")) @@ -195,7 +231,7 @@ fn failure_callbacks_run_after_the_failure_hook_however_it_ends(#[case] cancelle }; assert!(matches!( logging.resume(py, hook_result).unwrap(), - AdapterStep::Await(_) + LifecycleStep::Await(_) )); run( py, @@ -237,7 +273,7 @@ kwargs = {'logger': logger} .unwrap() .unbind(); let result = logging.begin(py, kwargs, 0.0).and_then(|step| match step { - AdapterStep::Await(_) => logging.resume(py, Ok(local(&locals, "kwargs").unbind())), + LifecycleStep::Await(_) => logging.resume(py, Ok(local(&locals, "kwargs").unbind())), step => Ok(step), }); let error = result.err().unwrap(); diff --git a/litellm-rust/crates/callbacks-legacy/tests/payload.rs b/litellm-rust/crates/callbacks-legacy/tests/payload.rs index 480bedf8548..5459b36af27 100644 --- a/litellm-rust/crates/callbacks-legacy/tests/payload.rs +++ b/litellm-rust/crates/callbacks-legacy/tests/payload.rs @@ -1,10 +1,12 @@ use std::ffi::CStr; -use litellm_callbacks::event::{CallEvent, Passthrough, RawResponse, RequestContext, WireRequest}; -use litellm_host_python::{AdapterStep, CallbackAdapter}; +use litellm_auth::SecretValue; +use litellm_host::event::{MachineEvent, RawResponse, RequestContext, WireRequest}; +use litellm_host_python::{LifecycleEvent, LifecycleStep, PythonLifecycle, to_py}; +use proptest::prelude::*; use pyo3::prelude::*; use rstest::rstest; -use serde_json::{Value, json}; +use serde_json::{Map, Value, json}; use super::LegacyLogging; use crate::PythonLogger; @@ -23,20 +25,12 @@ class PayloadLogger(StubLogger): def pre_call(self, input, api_key, additional_args): self.record('pre_call', None) self.pre = additional_args + self.pre_api_key = api_key on_pre_call(additional_args) - def _pre_call(self, input, api_key, additional_args): - self.record('_pre_call', None) - - def record_api_call_start_time(self): - self.record('record_api_call_start_time', None) - - def post_call(self, original_response, additional_args): + def post_call(self, original_response, api_key, additional_args): self.record('post_call', None) - self.post = (original_response, additional_args) - - def record_post_call(self, response, *rest): - self.record('record_post_call', response) + self.post = (original_response, api_key, additional_args) request = Request() kwargs = {} @@ -52,33 +46,47 @@ fn document(source: &str) -> Value { json!({"type": "document_url", "document_url": source}) } -fn before_send(script: &CStr, caller: Value, body: Value) -> WireRequest { - before_send_with_secrets(script, caller, body, &[]) +fn before_send(script: &CStr, body: Value) -> WireRequest { + before_send_with_secrets(script, json!({}), body, &[]) } -/// Runs `before_send` over `body` for a caller whose route-side view is `caller`, with the -/// Python objects `script` binds, then delivers the provider's raw response the way the +/// Runs `before_send` over `body` for a route whose parameters are `optional_params`, with +/// the Python objects `script` binds, then delivers the provider's raw response the way the /// driver does and runs the script's `check()`. fn before_send_with_secrets( script: &CStr, - caller: Value, + optional_params: Value, + body: Value, + secret_fields: &[&str], +) -> WireRequest { + before_send_bound(&[], script, optional_params, body, secret_fields) +} + +/// [`before_send_with_secrets`] with `bindings` placed in the namespace before `script` runs. +fn before_send_bound( + bindings: &[(&str, &Value)], + script: &CStr, + optional_params: Value, body: Value, secret_fields: &[&str], ) -> WireRequest { Python::initialize(); Python::attach(|py| { let locals = namespace(py, PAYLOAD_LOGGER); + for &(name, value) in bindings { + locals.set_item(name, to_py(py, value).unwrap()).unwrap(); + } run(py, &locals, script); let mut logging = LegacyLogging { - logger: Some(PythonLogger::new(local(&locals, "logger").unbind(), true)), + logger: Some(PythonLogger::new(local(&locals, "logger").unbind())), ..legacy_call(py, &locals, false) }; let context = RequestContext { model: "model".into(), custom_llm_provider: "provider".into(), - optional_params: caller.clone(), - passthrough_fields: Passthrough::unchanged(caller.as_object().unwrap(), &body), + optional_params, secret_fields: secret_fields.iter().map(|name| name.to_string()).collect(), + api_key: Some(SecretValue::new("route-key")), }; let wire = WireRequest { url: "https://provider.invalid/ocr".into(), @@ -86,17 +94,17 @@ fn before_send_with_secrets( body, }; let step = logging.before_send(py, Box::new(wire), &context).unwrap(); - let raw = CallEvent::ResponseReceived { + let raw = MachineEvent::ResponseReceived { raw: RawResponse { body: "raw response".into(), }, }; assert!(matches!( - logging.emit(py, &raw, None).unwrap(), - AdapterStep::Done + logging.emit(py, LifecycleEvent::Machine(&raw)).unwrap(), + LifecycleStep::Done )); run(py, &locals, c"check()"); - let AdapterStep::Wire(wire) = step else { + let LifecycleStep::Wire(wire) = step else { panic!("before_send did not hand back the wire request"); }; *wire @@ -129,11 +137,7 @@ def check(): ")] fn passthrough_keys_reach_pre_call_as_the_callers_own_objects(#[case] script: &CStr) { let body = json!({"model": "model", "document": document(DOCUMENT), "pages": [0]}); - let wire = before_send( - script, - json!({"document": document(DOCUMENT), "pages": [0]}), - body.clone(), - ); + let wire = before_send(script, body.clone()); assert_eq!(wire.body, body); } @@ -149,7 +153,6 @@ def check(): assert document['document_url'] == 'data:application/pdf;base64,ZWRpdGVk' ", json!({"document": document(DOCUMENT)}), - json!({"document": document(DOCUMENT)}), ); assert_eq!(wire.body["document"], document(EDITED)); } @@ -168,7 +171,6 @@ def check(): assert observed == [False], observed assert document == {'type': 'document_url', 'document_url': 'https://example.invalid/scan.pdf'} ", - json!({"document": document("https://example.invalid/scan.pdf")}), json!({"document": document(DOCUMENT)}), ); assert_eq!( @@ -177,6 +179,23 @@ def check(): ); } +#[test] +fn a_caller_value_with_no_json_form_is_left_out_of_realiasing() { + let body = json!({"pages": [0]}); + let wire = before_send( + c" +opaque = object() +kwargs = {'pages': opaque} +observed = [] +on_pre_call = lambda args: observed.append(args['complete_input_dict']['pages']) +def check(): + assert observed == [[0]], observed +", + body.clone(), + ); + assert_eq!(wire.body, body); +} + #[rstest] #[case::body( c" @@ -192,7 +211,7 @@ def on_pre_call(args): )] fn rebinding_the_payload_envelope_does_not_reach_the_wire(#[case] script: &CStr) { let body = json!({"document": document(DOCUMENT)}); - let wire = before_send(script, json!({}), body.clone()); + let wire = before_send(script, body.clone()); assert_eq!(wire.body, body); assert_eq!(wire.headers, [("x-route".to_string(), "route".to_string())]); } @@ -205,7 +224,6 @@ def on_pre_call(args): args['headers']['x-callback'] = 'edited' ", json!({}), - json!({}), ); assert_eq!( wire.headers, @@ -288,7 +306,7 @@ def on_pre_call(args): )] fn pre_call_body_edits_reach_the_wire(#[case] script: &CStr, #[case] expected: Value) { let body = json!({"document": document(DOCUMENT)}); - let wire = before_send(script, json!({"document": document(DOCUMENT)}), body); + let wire = before_send(script, body); assert_eq!(wire.body, expected); } @@ -302,7 +320,6 @@ def on_pre_call(args): retained['x-retained'] = 'sent' ", json!({}), - json!({}), ); assert_eq!( wire.headers, @@ -314,52 +331,193 @@ def on_pre_call(args): } #[test] -fn post_call_receives_the_raw_response_and_the_payload_dicts_pre_call_saw() { +fn post_call_receives_the_raw_response_the_route_key_and_the_body_and_headers_pre_call_saw() { before_send( c" def check(): - original_response, additional_args = logger.post + original_response, api_key, additional_args = logger.post assert original_response == 'raw response', original_response + assert api_key == logger.pre_api_key == 'route-key', (api_key, logger.pre_api_key) + assert additional_args == { + 'complete_input_dict': logger.pre['complete_input_dict'], + 'headers': logger.pre['headers'], + }, additional_args assert additional_args['complete_input_dict'] is logger.pre['complete_input_dict'] assert additional_args['headers'] is logger.pre['headers'] ", - json!({}), json!({"document": document(DOCUMENT)}), ); } -#[rstest] -#[case::every_phase_listens(c"{}", &["pre_call", "post_call"])] -#[case::no_input_callback( - c"{'input': False}", - &["_pre_call", "record_api_call_start_time", "record_post_call"] -)] -#[case::no_payload_consumer(c"{'payload': False}", &["record_api_call_start_time"])] -fn payload_callbacks_run_only_for_the_phases_someone_listens_to( - #[case] needed: &CStr, - #[case] expected_calls: &[&str], -) { - let script = std::ffi::CString::new(format!( - " -logger.needed = {needed} +#[test] +fn every_request_runs_the_full_pre_call_and_post_call() { + let wire = before_send( + c" def on_pre_call(args): args['complete_input_dict']['include_image_base64'] = True def check(): - assert logger.names() == {expected_calls:?}, logger.calls + assert logger.names() == ['pre_call', 'post_call'], logger.calls ", - needed = needed.to_str().unwrap(), - expected_calls = expected_calls, - )) - .unwrap(); - let body = json!({"document": document(DOCUMENT)}); - let wire = before_send(&script, json!({}), body.clone()); - let edited = json!({"document": document(DOCUMENT), "include_image_base64": true}); + json!({"document": document(DOCUMENT)}), + ); assert_eq!( wire.body, - if expected_calls.contains(&"pre_call") { - edited - } else { - body - } + json!({"document": document(DOCUMENT), "include_image_base64": true}) ); } + +/// What one pre-call callback does to the payload it is handed. +#[derive(Clone, Debug)] +enum Edit { + Nothing, + Set(String, Value), + Remove(String), + Rebind(Value), + RebindThenSetRetained(String, Value), +} + +impl Edit { + fn script(&self) -> Value { + match self { + Self::Nothing => json!({"kind": "nothing"}), + Self::Set(key, value) => json!({"kind": "set", "key": key, "value": value}), + Self::Remove(key) => json!({"kind": "remove", "key": key}), + Self::Rebind(value) => json!({"kind": "rebind", "value": value}), + Self::RebindThenSetRetained(key, value) => { + json!({"kind": "rebind_then_set_retained", "key": key, "value": value}) + } + } + } + + /// The legacy contract: the provider is sent the body object `pre_call` received, as + /// the callback left it. Rebinding the envelope's key points the envelope elsewhere and + /// leaves that object alone. + fn sent(&self, body: &Map) -> Value { + let mut sent = body.clone(); + match self { + Self::Nothing | Self::Rebind(_) => {} + Self::Set(key, value) | Self::RebindThenSetRetained(key, value) => { + sent.insert(key.clone(), value.clone()); + } + Self::Remove(key) => { + sent.remove(key); + } + } + Value::Object(sent) + } +} + +/// How the caller's keyword for a body key relates to what the route sends under it. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum Caller { + PassedUnchanged, + RewrittenByTheRoute, + NotPassed, +} + +const MODEL: &CStr = c" +aliased = {} +def on_pre_call(args): + body = args['complete_input_dict'] + aliased.update({name: body[name] is kwargs[name] for name in unchanged}) + kind = edit['kind'] + if kind == 'set': + body[edit['key']] = edit['value'] + elif kind == 'remove': + body.pop(edit['key'], None) + elif kind == 'rebind': + args['complete_input_dict'] = edit['value'] + elif kind == 'rebind_then_set_retained': + args['complete_input_dict'] = {} + body[edit['key']] = edit['value'] +def check(): + assert aliased == {name: True for name in unchanged}, aliased + assert logger.names() == ['pre_call', 'post_call'], logger.calls +"; + +fn json_value() -> impl Strategy { + let leaf = prop_oneof![ + Just(Value::Null), + any::().prop_map(Value::from), + any::().prop_map(Value::from), + any::() + .prop_filter("JSON has no NaN or infinity", |number| number.is_finite()) + .prop_map(Value::from), + ".{0,8}".prop_map(Value::from), + ]; + leaf.prop_recursive(3, 24, 4, |inner| { + prop_oneof![ + prop::collection::vec(inner.clone(), 0..4).prop_map(Value::from), + prop::collection::btree_map(key(), inner, 0..4) + .prop_map(|fields| Value::Object(fields.into_iter().collect())), + ] + }) +} + +fn key() -> impl Strategy { + "[a-z]{1,6}" +} + +fn caller() -> impl Strategy { + prop_oneof![ + Just(Caller::PassedUnchanged), + Just(Caller::RewrittenByTheRoute), + Just(Caller::NotPassed), + ] +} + +fn edit() -> impl Strategy { + prop_oneof![ + Just(Edit::Nothing), + (key(), json_value()).prop_map(|(key, value)| Edit::Set(key, value)), + key().prop_map(Edit::Remove), + json_value().prop_map(Edit::Rebind), + (key(), json_value()).prop_map(|(key, value)| Edit::RebindThenSetRetained(key, value)), + ] +} + +proptest! { + #![proptest_config(ProptestConfig::with_cases(128))] + + /// For any body, any caller keywords and any callback edit: every keyword the route + /// sends unchanged reaches `pre_call` as the caller's own object, and the provider is + /// sent exactly what the model says, so a callback that edits nothing changes nothing. + #[test] + fn the_wire_is_the_body_pre_call_received_as_the_callback_left_it( + fields in prop::collection::btree_map(key(), (json_value(), caller()), 0..5), + edit in edit(), + ) { + let body: Map = fields + .iter() + .map(|(name, (value, _))| (name.clone(), value.clone())) + .collect(); + let kwargs: Map = fields + .iter() + .filter_map(|(name, (value, caller))| match caller { + Caller::PassedUnchanged => Some((name.clone(), value.clone())), + Caller::RewrittenByTheRoute => Some((name.clone(), json!([value]))), + Caller::NotPassed => None, + }) + .collect(); + let unchanged: Value = fields + .iter() + .filter(|(_, (_, caller))| *caller == Caller::PassedUnchanged) + .map(|(name, _)| Value::from(name.clone())) + .collect(); + + let wire = before_send_bound( + &[ + ("kwargs", &Value::Object(kwargs)), + ("unchanged", &unchanged), + ("edit", &edit.script()), + ], + MODEL, + json!({}), + Value::Object(body.clone()), + &[], + ); + + prop_assert_eq!(wire.body, edit.sent(&body)); + prop_assert_eq!(wire.headers, [("x-route".to_string(), "route".to_string())]); + } +} diff --git a/litellm-rust/crates/callbacks-legacy/tests/support.rs b/litellm-rust/crates/callbacks-legacy/tests/support.rs index 1663e11963e..d3cc32e301f 100644 --- a/litellm-rust/crates/callbacks-legacy/tests/support.rs +++ b/litellm-rust/crates/callbacks-legacy/tests/support.rs @@ -5,64 +5,94 @@ use pyo3::types::{PyDict, PyTuple}; use crate::{LegacyLogging, LegacySurface, PublicCall}; -/// Stand-ins for every litellm function the legacy contract calls. Tests share one -/// interpreter and run concurrently, so each stub is installed idempotently and forwards to -/// the per-test `StubLogger` it is handed (directly, or as `kwargs['logger']`). +/// The parameters of every `legacy_callbacks` function, as the real module declares them. +/// `tests/test_litellm/rust_bridge/test_legacy_callbacks.py` pins this file to the Python +/// signatures, and [`namespace`] binds every fake call against it. +pub(crate) const PYTHON_CONTRACT: &str = include_str!("../python_contract.json"); + +/// Stand-ins for `legacy_callbacks`, the only Python module the crate calls. Tests +/// share one interpreter and run concurrently, so each fake is installed idempotently and +/// forwards to the per-test `StubLogger` it is handed (directly, or as `kwargs['logger']`). +/// Every fake is bound against the contract first, so a call the real module would reject +/// fails here too. const STUBS: &CStr = c" import contextvars +import inspect +import json import sys +import traceback import types -for name in ( - 'litellm', - 'litellm.utils', - 'litellm.types', - 'litellm.types.utils', - 'litellm._internal_context', - 'litellm.litellm_core_utils', - 'litellm.litellm_core_utils.logging_worker', - 'litellm.litellm_core_utils.litellm_logging', - 'litellm.rust_bridge', - 'litellm.rust_bridge.legacy_callbacks', -): +for name in ('litellm', 'litellm.rust_bridge', 'litellm.rust_bridge.legacy_callbacks'): sys.modules.setdefault(name, types.ModuleType(name)) legacy = sys.modules['litellm.rust_bridge.legacy_callbacks'] -legacy.setup = lambda call_type, args, kwargs, start, asynchronous: types.SimpleNamespace( - logger=kwargs['logger_factory'](kwargs) if 'logger_factory' in kwargs else kwargs['logger'], - kwargs=kwargs, - bridge_owned=True, -) -legacy.deployment_callbacks_needed = lambda: True -legacy.check_limits = lambda arguments: arguments['logger'].check_limits(arguments) -legacy.callbacks_needed = lambda logger, phase: logger.needed.get(phase, True) -legacy.success_bookkeeping = lambda logger, response, start, end, asynchronous: logger.record( - 'success_bookkeeping', asynchronous -) -legacy.failure_bookkeeping = lambda logger, error, start, end, asynchronous: logger.record( - 'failure_bookkeeping', asynchronous -) -legacy.finalize = lambda response, logger, kwargs, start, end: logger.record('finalize', response) +CONTRACT = json.loads(python_contract) -utils = sys.modules['litellm.utils'] -utils.async_pre_call_deployment_hook = lambda kwargs, call_type: kwargs['logger'].hook( - 'pre', kwargs, call_type -) -utils.async_post_call_success_deployment_hook = lambda kwargs, response, call_type: kwargs[ - 'logger' -].hook('success', response, call_type) -utils.async_post_call_failure_deployment_hook = lambda kwargs, error, call_type: kwargs[ - 'logger' -].hook('failure', error, call_type) -utils._restore_correlation_context_if_supported = lambda logger: logger.record('restore', None) -internal = sys.modules['litellm._internal_context'] -if not hasattr(internal, 'is_internal_call'): - internal.is_internal_call = contextvars.ContextVar('is_internal_call', default=False) +def contracted(name, fake): + signature = inspect.Signature( + [inspect.Parameter(parameter, inspect.Parameter.POSITIONAL_OR_KEYWORD) for parameter in CONTRACT[name]] + ) -sys.modules['litellm.types.utils'].CustomPricingLiteLLMParams = type( - 'CustomPricingLiteLLMParams', (), {'model_fields': {'ocr_cost_per_page': None}} -) + def checked(*args, **kwargs): + signature.bind(*args, **kwargs) + return fake(*args, **kwargs) + + return checked + + +if not hasattr(legacy, 'is_internal'): + legacy.is_internal = contextvars.ContextVar('is_internal_call', default=False) + +FAKES = { + 'setup': lambda call_type, args, kwargs, start, asynchronous: types.SimpleNamespace( + logger=kwargs['logger_factory'](kwargs) if 'logger_factory' in kwargs else kwargs['logger'], + kwargs=kwargs, + ), + 'check_limits': lambda arguments: arguments['logger'].check_limits(arguments), + 'finalize': lambda response, logger, kwargs, start, end: logger.record('finalize', response), + 'update_logging': lambda logger, kwargs, model, optional_params, litellm_params, provider: logger.update_from_kwargs( + kwargs=kwargs, + model=model, + optional_params=optional_params, + litellm_params=litellm_params, + custom_llm_provider=provider, + ), + 'pre_call': lambda logger, input, api_key, additional_args: logger.pre_call(input, api_key, additional_args), + 'post_call': lambda logger, original_response, api_key, additional_args: logger.post_call( + original_response, api_key, additional_args + ), + 'defers_async_logging': lambda logger: bool(getattr(logger, '_defer_async_logging', False)), + 'defer_success': lambda logger, pending: setattr(logger, '_native_pending_logging', pending), + 'sync_success_for_async_call': lambda logger, response, start, end: logger.handle_sync_success_callbacks_for_async_calls( + response, start, end + ), + 'failure_handler': lambda logger, error, start, end, asynchronous: ( + logger.async_failure_handler if asynchronous else logger.failure_handler + )(error, ''.join(traceback.format_exception(error)), start, end), + 'submit_success': lambda logger, response, start, end: logger.record('submit', (response, start, end)), + 'async_success_handler': lambda logger, response, start, end: logger.async_success_handler(response, start, end), + 'enqueue_logging': lambda coroutine: coroutine.enqueue(), + 'restore_context': lambda logger: logger.record('restore', None), + 'custom_pricing_fields': lambda: ('ocr_cost_per_page',), + 'is_internal_call': lambda: legacy.is_internal.get(), + 'credential_list': lambda: [], + 'warn_unknown_credential': lambda name, loaded: None, + 'before_deployment_call': lambda kwargs, call_type: kwargs['logger'].hook('pre', kwargs, call_type), + 'after_deployment_success': lambda kwargs, response, call_type: kwargs['logger'].hook( + 'success', response, call_type + ), + 'after_deployment_failure': lambda kwargs, error, call_type: kwargs['logger'].hook('failure', error, call_type), + 'stream_opened': lambda logger: logger.record('stream_opened', None), + 'stream_success': lambda logger, request_body, chunks, start, end, first_chunk: logger.record( + 'stream_success', list(chunks) + ), + 'stream_failure': lambda logger, request_body, chunks, error: logger.record('stream_failure', error), +} +assert FAKES.keys() == CONTRACT.keys(), sorted(FAKES.keys() ^ CONTRACT.keys()) +for name, fake in FAKES.items(): + setattr(legacy, name, contracted(name, fake)) unraisable = sys.modules.setdefault( @@ -77,20 +107,6 @@ def unraisable_from(owner): return [error for source, error in unraisable.events if source is owner] -class Worker: - def ensure_initialized_and_enqueue(self, coroutine): - return coroutine.enqueue() - - -class Executor: - def submit(self, run, handler, *args): - handler.__self__.record('submit', args) - - -sys.modules['litellm.litellm_core_utils.logging_worker'].GLOBAL_LOGGING_WORKER = Worker() -sys.modules['litellm.litellm_core_utils.litellm_logging'].executor = Executor() - - class StubCoroutine: def __init__(self, logger): self.logger = logger @@ -106,7 +122,6 @@ class StubCoroutine: class StubLogger: def __init__(self): self.calls = [] - self.needed = {} self.hooks = {} self.on_enqueue = lambda coroutine: None @@ -147,6 +162,7 @@ logger = StubLogger() /// A namespace with the stubs, `StubLogger` and a fresh `logger`, after `script` ran in it. pub(crate) fn namespace<'py>(py: Python<'py>, script: &CStr) -> Bound<'py, PyDict> { let locals = PyDict::new(py); + locals.set_item("python_contract", PYTHON_CONTRACT).unwrap(); py.run(STUBS, Some(&locals), Some(&locals)).unwrap(); py.run(script, Some(&locals), Some(&locals)).unwrap(); locals @@ -181,6 +197,7 @@ pub(crate) fn legacy_call( LegacySurface { call_type: "test", input_description: "test input", + stream: None, }, call, asynchronous, diff --git a/litellm-rust/crates/callbacks-legacy/tests/terminal.rs b/litellm-rust/crates/callbacks-legacy/tests/terminal.rs index 9b9d29108f6..f68209233f2 100644 --- a/litellm-rust/crates/callbacks-legacy/tests/terminal.rs +++ b/litellm-rust/crates/callbacks-legacy/tests/terminal.rs @@ -1,7 +1,7 @@ use std::ffi::CStr; -use litellm_callbacks::event::{CallEvent, FailureOrigin, Timing}; -use litellm_host_python::{AdapterStep, CallbackAdapter, PublicValue}; +use litellm_host::event::{FailureOrigin, Timing}; +use litellm_host_python::{LifecycleEvent, LifecycleStep, PythonLifecycle}; use pyo3::exceptions::PyRuntimeError; use pyo3::exceptions::asyncio::CancelledError; use pyo3::prelude::*; @@ -19,52 +19,52 @@ const TIMING: Timing = Timing { fn logged(py: Python<'_>, locals: &Bound<'_, PyDict>, asynchronous: bool) -> LegacyLogging { LegacyLogging { - logger: Some(PythonLogger::new(local(locals, "logger").unbind(), true)), + logger: Some(PythonLogger::new(local(locals, "logger").unbind())), ..legacy_call(py, locals, asynchronous) } } -fn succeed(py: Python<'_>, locals: &Bound<'_, PyDict>, logging: &mut LegacyLogging) -> AdapterStep { +fn succeed( + py: Python<'_>, + locals: &Bound<'_, PyDict>, + logging: &mut LegacyLogging, +) -> LifecycleStep { let response = local(locals, "response").unbind(); logging .emit( py, - &CallEvent::Succeeded { timing: TIMING }, - Some(PublicValue::Response(&response)), + LifecycleEvent::Succeeded { + timing: TIMING, + response: &response, + }, ) .unwrap() } -fn fail(py: Python<'_>, locals: &Bound<'_, PyDict>, logging: &mut LegacyLogging) -> AdapterStep { +fn fail(py: Python<'_>, locals: &Bound<'_, PyDict>, logging: &mut LegacyLogging) -> LifecycleStep { let failure = PyErr::from_value(local(locals, "failure")); logging .emit( py, - &CallEvent::Failed { + LifecycleEvent::Failed { timing: TIMING, origin: FailureOrigin::Host, + error: &failure, }, - Some(PublicValue::Error(&failure)), ) .unwrap() } #[rstest] #[case::sync_listened(false, c"", &["submit"])] -#[case::sync_unlistened(false, c"logger.needed = {'sync_success': False}", &["success_bookkeeping"])] #[case::async_listened( true, c"", &["async_success_handler", "enqueued", "sync_success_for_async_call"] )] -#[case::async_unlistened( - true, - c"logger.needed = {'async_success': False, 'sync_success_async': False}", - &["success_bookkeeping"] -)] #[case::async_deferred(true, c"logger._defer_async_logging = True", &["sync_success_for_async_call"])] #[case::async_with_fallbacks(true, c"kwargs = {'fallbacks': ['other']}", &["sync_success_for_async_call"])] -fn success_reaches_only_the_callbacks_that_listen( +fn success_reaches_the_logging_handlers( #[case] asynchronous: bool, #[case] script: &CStr, #[case] expected: &[&str], @@ -76,7 +76,7 @@ fn success_reaches_only_the_callbacks_that_listen( let mut logging = logged(py, &locals, asynchronous); assert!(matches!( succeed(py, &locals, &mut logging), - AdapterStep::Done + LifecycleStep::Done )); let names: Vec = local(&locals, "logger") .call_method0("names") @@ -109,7 +109,10 @@ fn internal_calls_skip_failure_callbacks_only_when_asynchronous( internal: true, ..logged(py, &locals, asynchronous) }; - assert!(matches!(fail(py, &locals, &mut logging), AdapterStep::Done)); + assert!(matches!( + fail(py, &locals, &mut logging), + LifecycleStep::Done + )); let names: Vec = local(&locals, "logger") .call_method0("names") .unwrap() @@ -157,7 +160,7 @@ logger = FailingLogger() let mut logging = logged(py, &locals, true); assert!(matches!( succeed(py, &locals, &mut logging), - AdapterStep::Done + LifecycleStep::Done )); assert!( logging @@ -173,14 +176,8 @@ logger = FailingLogger() #[rstest] #[case::sync_listened(false, c"", &["failure_handler"])] -#[case::sync_unlistened(false, c"logger.needed = {'sync_failure': False}", &["failure_bookkeeping"])] #[case::async_listened(true, c"", &["failure_handler", "async_failure_handler"])] -#[case::async_unlistened( - true, - c"logger.needed = {'sync_failure': False, 'async_failure': False}", - &["failure_bookkeeping", "failure_bookkeeping"] -)] -fn failure_reaches_only_the_callbacks_that_listen( +fn failure_reaches_the_logging_handlers( #[case] asynchronous: bool, #[case] script: &CStr, #[case] expected: &[&str], @@ -192,7 +189,10 @@ fn failure_reaches_only_the_callbacks_that_listen( let mut logging = logged(py, &locals, asynchronous); let step = fail(py, &locals, &mut logging); let awaits_async_handler = expected.contains(&"async_failure_handler"); - assert_eq!(matches!(step, AdapterStep::Await(_)), awaits_async_handler); + assert_eq!( + matches!(step, LifecycleStep::Await(_)), + awaits_async_handler + ); let names: Vec = local(&locals, "logger") .call_method0("names") .unwrap() @@ -227,7 +227,7 @@ logger = FailingLogger() let mut logging = logged(py, &locals, true); assert!(matches!( fail(py, &locals, &mut logging), - AdapterStep::Await(_) + LifecycleStep::Await(_) )); assert!( logging @@ -265,7 +265,7 @@ fn the_async_failure_handler_ends_the_call_unless_it_was_cancelled( }; let expected = result.as_ref().err().map(|error| error.value(py).clone()); match logging.resume(py, result) { - Ok(step) => assert!(done && matches!(step, AdapterStep::Done)), + Ok(step) => assert!(done && matches!(step, LifecycleStep::Done)), Err(propagated) => { assert!(!done); assert!(propagated.value(py).is(expected.unwrap())); diff --git a/litellm-rust/crates/callbacks/src/event.rs b/litellm-rust/crates/callbacks/src/event.rs deleted file mode 100644 index e6f88fd9709..00000000000 --- a/litellm-rust/crates/callbacks/src/event.rs +++ /dev/null @@ -1,135 +0,0 @@ -use std::time::{SystemTime, UNIX_EPOCH}; - -use serde_json::{Map, Value}; - -/// Seconds since the Unix epoch, on one clock for every host. -pub fn epoch_seconds() -> f64 { - SystemTime::now() - .duration_since(UNIX_EPOCH) - .map(|duration| duration.as_secs_f64()) - .unwrap_or(0.0) -} - -#[derive(Clone, Copy, Debug, PartialEq)] -pub struct Timing { - pub start_time: f64, - pub end_time: f64, -} - -/// The provider request as it is about to leave, offered to the host for rewriting. -#[derive(Clone, Debug, PartialEq)] -pub struct WireRequest { - pub url: String, - pub headers: Vec<(String, String)>, - pub body: Value, -} - -/// What the route knows about the request it is sending, for a host that logs it. The -/// route owns these facts; a host reads them beside the wire request and never rewrites -/// them. -#[derive(Clone, Debug, PartialEq)] -pub struct RequestContext { - pub model: String, - pub custom_llm_provider: String, - /// The route's parameters before the provider transformation. - pub optional_params: Value, - pub passthrough_fields: Passthrough, - /// Optional-param names that carry credentials and must be redacted when logged. - pub secret_fields: Vec, -} - -/// Body keys whose values are the caller's inputs, unchanged by the route. The only way to -/// build one is to compare the two, so a route cannot name a key it rewrote. -#[derive(Clone, Debug, Default, PartialEq, Eq)] -pub struct Passthrough(Vec); - -impl Passthrough { - pub fn unchanged(caller: &Map, body: &Value) -> Self { - Self( - caller - .iter() - .filter(|(name, value)| body.get(name.as_str()) == Some(*value)) - .map(|(name, _)| name.clone()) - .collect(), - ) - } - - pub fn iter(&self) -> impl Iterator { - self.0.iter().map(String::as_str) - } - - pub fn contains(&self, name: &str) -> bool { - self.0.iter().any(|field| field == name) - } -} - -#[derive(Clone, Debug, PartialEq, Eq)] -pub struct RawResponse { - pub body: String, -} - -/// Whether a failure surfaced inside the call, including a host op the call asked for, -/// or in a host step around it (preparing the arguments, finalizing the response). -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -pub enum FailureOrigin { - Call, - Host, -} - -#[derive(Clone, Debug, PartialEq)] -pub enum CallEvent { - ResponseReceived { - raw: RawResponse, - }, - Succeeded { - timing: Timing, - }, - Failed { - timing: Timing, - origin: FailureOrigin, - }, -} - -#[cfg(test)] -mod tests { - use rstest::rstest; - use serde_json::json; - - use super::*; - - #[rstest] - #[case::unchanged_scalar(json!({"pages": [0]}), json!({"pages": [0]}), &["pages"])] - #[case::unchanged_explicit_null(json!({"pages": null}), json!({"pages": null}), &["pages"])] - #[case::unchanged_nested_object( - json!({"document": {"type": "document_url", "document_url": "https://a/b.pdf"}}), - json!({"document": {"type": "document_url", "document_url": "https://a/b.pdf"}, "model": "m"}), - &["document"] - )] - #[case::rewritten_value( - json!({"document": {"type": "document_url", "document_url": "https://a/b.pdf"}}), - json!({"document": {"type": "document_url", "document_url": "data:application/pdf;base64,YWJj"}}), - &[] - )] - #[case::dropped_nested_field( - json!({"document": {"type": "image_url", "image_url": "https://a/b.png", "document_name": "b.png"}}), - json!({"document": {"type": "image_url", "image_url": "https://a/b.png"}}), - &[] - )] - #[case::added_nested_field( - json!({"document": {"type": "image_url", "image_url": "https://a/b.png"}}), - json!({"document": {"type": "image_url", "image_url": "https://a/b.png", "detail": "high"}}), - &[] - )] - #[case::reordered_array(json!({"pages": [0, 1]}), json!({"pages": [1, 0]}), &[])] - #[case::consumed_by_the_route(json!({"api_key": "k", "pages": [0]}), json!({"pages": [0]}), &["pages"])] - #[case::added_by_the_route(json!({}), json!({"model": "m"}), &[])] - #[case::non_object_body(json!({"pages": [0]}), json!([{"pages": [0]}]), &[])] - fn passthrough_is_exactly_the_callers_unchanged_keys( - #[case] caller: Value, - #[case] body: Value, - #[case] expected: &[&str], - ) { - let passthrough = Passthrough::unchanged(caller.as_object().unwrap(), &body); - assert_eq!(passthrough.iter().collect::>(), expected); - } -} diff --git a/litellm-rust/crates/core/Cargo.toml b/litellm-rust/crates/core/Cargo.toml index db6cfc4b340..3995a235778 100644 --- a/litellm-rust/crates/core/Cargo.toml +++ b/litellm-rust/crates/core/Cargo.toml @@ -9,7 +9,7 @@ autotests = false [dependencies] litellm-types.workspace = true litellm-core-utils.workspace = true -litellm-callbacks.workspace = true +litellm-host.workspace = true bytes.workspace = true futures-util.workspace = true base64.workspace = true diff --git a/litellm-rust/crates/core/src/lib.rs b/litellm-rust/crates/core/src/lib.rs index 58aef6cd629..e3e2fb48721 100644 --- a/litellm-rust/crates/core/src/lib.rs +++ b/litellm-rust/crates/core/src/lib.rs @@ -2,7 +2,6 @@ pub mod audio_transcription; pub mod chat_completions; pub mod constants; pub mod error; -pub mod machine; pub mod messages; pub mod ocr; pub mod responses; diff --git a/litellm-rust/crates/core/src/messages/handler.rs b/litellm-rust/crates/core/src/messages/handler.rs index b95402b1a7a..22e2c398ff7 100644 --- a/litellm-rust/crates/core/src/messages/handler.rs +++ b/litellm-rust/crates/core/src/messages/handler.rs @@ -1,88 +1,54 @@ -use litellm_llms::custom_httpx::http_handler::http_request; -use litellm_types::llms::anthropic_messages::anthropic_response::AnthropicMessagesResponse; +use std::time::Duration; -use super::{ - Error, client::http_client, common_utils::truncate_error_body, - prepare::prepare_provider_request, +use litellm_llms::{ + base_llm::anthropic_messages::transformation::BaseAnthropicMessagesConfig, + custom_httpx::{http_handler::http_request, transport::Error as TransportError}, }; -use crate::{constants::ANTHROPIC_MESSAGES_PROVIDER, messages::types::MessagesRequest}; +use litellm_types::llms::anthropic_messages::anthropic_response::AnthropicMessagesResponse; +use serde_json::Value; -pub(super) async fn execute_messages_provider_call( - request: MessagesRequest<'_>, +use super::{Error, client::http_client, common_utils::truncate_error_body}; + +pub(super) fn network(error: reqwest::Error) -> Error { + Error::Transport(TransportError::Network(error.to_string())) +} + +pub(super) async fn send( + url: &str, + headers: &[(String, String)], + body: &Value, + timeout: Option, +) -> Result { + let builder = headers.iter().fold( + http_client().post(url).json(body), + |builder, (key, value)| builder.header(key, value), + ); + let builder = match timeout { + Some(duration) => builder.timeout(duration), + None => builder, + }; + http_request(builder).await.map_err(network) +} + +pub(super) async fn provider_error(response: reqwest::Response) -> Error { + let status = response.status().as_u16(); + match response.text().await { + Ok(text) => Error::Transport(TransportError::Http { + status, + body: truncate_error_body(&text), + }), + Err(error) => network(error), + } +} + +pub(super) fn decode_response( + config: &dyn BaseAnthropicMessagesConfig, + model: &str, + text: &str, ) -> Result { - let request = prepare_provider_request(request)?; - let mut request_builder = http_client().post(&request.url).json(&request.body); - for (key, value) in &request.upstream_headers { - request_builder = request_builder.header(key, value); - } - if let Some(duration) = request.timeout { - request_builder = request_builder.timeout(duration); - } - - let response = http_request(request_builder).await.map_err(|err| { - Error::Transport(litellm_llms::custom_httpx::transport::Error::Network( - err.to_string(), - )) - })?; - - let status = response.status(); - let text = response.text().await.map_err(|err| { - Error::Transport(litellm_llms::custom_httpx::transport::Error::Network( - err.to_string(), - )) - })?; - - if !status.is_success() { - return Err(Error::Transport( - litellm_llms::custom_httpx::transport::Error::Http { - status: status.as_u16(), - body: truncate_error_body(&text), - }, - )); - } - - let response = serde_json::from_str(&text) + let response = serde_json::from_str(text) .map_err(|err| Error::InvalidResponse(format!("invalid messages response JSON: {err}")))?; - request - .config - .transform_anthropic_messages_response(&request.model, response) + config + .transform_anthropic_messages_response(model, response) .map_err(Error::from) } - -pub(super) async fn execute_messages_provider_stream( - request: MessagesRequest<'_>, -) -> Result { - let request = prepare_provider_request(request)?; - if request.provider != ANTHROPIC_MESSAGES_PROVIDER { - return Err(Error::Unsupported("streaming messages for this provider")); - } - - let mut request_builder = http_client().post(&request.url).json(&request.body); - for (key, value) in &request.upstream_headers { - request_builder = request_builder.header(key, value); - } - if let Some(duration) = request.timeout { - request_builder = request_builder.timeout(duration); - } - - let response = http_request(request_builder).await.map_err(|err| { - Error::Transport(litellm_llms::custom_httpx::transport::Error::Network( - err.to_string(), - )) - })?; - let status = response.status(); - if !status.is_success() { - let text = response.text().await.map_err(|err| { - Error::Transport(litellm_llms::custom_httpx::transport::Error::Network( - err.to_string(), - )) - })?; - return Err(Error::Transport( - litellm_llms::custom_httpx::transport::Error::Http { - status: status.as_u16(), - body: truncate_error_body(&text), - }, - )); - } - Ok(response) -} diff --git a/litellm-rust/crates/core/src/messages/mod.rs b/litellm-rust/crates/core/src/messages/mod.rs index c3d7bea48ff..289f79109dd 100644 --- a/litellm-rust/crates/core/src/messages/mod.rs +++ b/litellm-rust/crates/core/src/messages/mod.rs @@ -1,11 +1,8 @@ //! The Anthropic Messages call, the Rust equivalent of Python's //! `litellm.messages()`. //! -//! [`messages`] is the top-level entrypoint: give it a model, a body, and -//! credentials, and it resolves the provider, transforms the request, calls the -//! provider, and returns a typed non-streaming response. [`messages_stream`] -//! is the streaming variant; it hands the raw upstream response back so a host -//! can splice the event stream to its own caller. +//! [`route`] is the call as a machine a host drives, streaming or not. [`messages`] runs +//! it in process for a caller that already holds the request and wants the message. mod error; pub mod types; @@ -14,17 +11,34 @@ mod client; mod common_utils; mod handler; mod prepare; -use handler::{execute_messages_provider_call, execute_messages_provider_stream}; +pub mod route; use litellm_types::llms::anthropic_messages::anthropic_response::AnthropicMessagesResponse; +use route::{LocalMessagesHost, MessagesCall, MessagesOutput, messages_machine}; +use serde_json::Value; use crate::messages::types::MessagesRequest; pub async fn messages(request: MessagesRequest<'_>) -> Result { - execute_messages_provider_call(request).await -} - -pub async fn messages_stream(request: MessagesRequest<'_>) -> Result { - execute_messages_provider_stream(request).await + let Value::Object(body) = request.body else { + return Err(Error::InvalidRequest( + "messages body must be an object".into(), + )); + }; + let call = MessagesCall { + model: request.model.into(), + body, + api_key: request.api_key.map(Into::into), + api_base: request.api_base.map(Into::into), + custom_llm_provider: request.custom_llm_provider.map(Into::into), + extra_headers: request.extra_headers, + timeout: request.timeout, + }; + match litellm_host::run::run(messages_machine(), &LocalMessagesHost::new(call)).await? { + MessagesOutput::Message(message) => Ok(*message), + MessagesOutput::Streamed => Err(Error::Unsupported( + "streamed responses need a streaming host", + )), + } } #[cfg(test)] diff --git a/litellm-rust/crates/core/src/messages/prepare.rs b/litellm-rust/crates/core/src/messages/prepare.rs index 8b676803871..850f9108869 100644 --- a/litellm-rust/crates/core/src/messages/prepare.rs +++ b/litellm-rust/crates/core/src/messages/prepare.rs @@ -2,6 +2,7 @@ use litellm_core_utils::get_llm_provider_logic::{CustomLlmProvider, get_custom_l use litellm_llms::base_llm::anthropic_messages::transformation::{ BaseAnthropicMessagesConfig, MessagesAuthStrategy, }; +use litellm_types::llms::anthropic_messages::anthropic_request::AnthropicMessagesRequest; use serde_json::{Map, Value}; use super::{ @@ -37,10 +38,14 @@ pub(super) fn prepare_provider_request( let headers = validate_environment(config, request.extra_headers, request.api_key, &env_lookup)?; - let typed_request = serde_json::from_value(request.body).map_err(|err| { - Error::InvalidRequest(format!("invalid Anthropic messages request: {err}")) + let typed_request: AnthropicMessagesRequest = + serde_json::from_value(request.body).map_err(|err| { + Error::InvalidRequest(format!("invalid Anthropic messages request: {err}")) + })?; + let transformed = config.transform_anthropic_messages_request(AnthropicMessagesRequest { + model: model.clone(), + ..typed_request })?; - let transformed = config.transform_anthropic_messages_request(typed_request)?; let body = serde_json::to_value(transformed).map_err(|err| { Error::InvalidRequest(format!( "failed to serialize Anthropic messages request: {err}" diff --git a/litellm-rust/crates/core/src/messages/route.rs b/litellm-rust/crates/core/src/messages/route.rs new file mode 100644 index 00000000000..838b56fcb4b --- /dev/null +++ b/litellm-rust/crates/core/src/messages/route.rs @@ -0,0 +1,196 @@ +use std::{sync::Mutex, time::Duration}; + +use bytes::Bytes; +use litellm_auth::SecretValue; +use litellm_core_utils::get_llm_provider_logic::get_custom_llm_provider; +use litellm_host::{ + event::{MachineEvent, RawResponse, RequestContext, WireRequest}, + host::{Demand, Host}, + machine::{HostChannel, MachineFault, RouteMachine}, + route::Route, +}; +use litellm_types::llms::anthropic_messages::anthropic_response::AnthropicMessagesResponse; +use serde_json::{Map, Value}; + +use super::{ + Error, + common_utils::messages_provider_config, + handler::{decode_response, network, provider_error, send}, + prepare::prepare_provider_request, + types::MessagesRequest, +}; +use crate::constants::ANTHROPIC_MESSAGES_PROVIDER; + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum MessagesOp { + ProjectRequest, +} + +pub enum MessagesOpResult { + Request(Box), +} + +/// The caller's request as the host projects it. +pub struct MessagesCall { + pub model: String, + pub body: Map, + pub api_key: Option, + pub api_base: Option, + pub custom_llm_provider: Option, + pub extra_headers: Option>, + pub timeout: Option, +} + +impl MessagesCall { + fn streams(&self) -> bool { + self.body.get("stream").and_then(Value::as_bool) == Some(true) + } +} + +pub enum MessagesOutput { + Message(Box), + /// Every chunk already reached the host through `Deliver`. + Streamed, +} + +pub struct Messages; + +impl Route for Messages { + type Response = MessagesOutput; + type Error = Error; + type Op = MessagesOp; + type OpResult = MessagesOpResult; + type Chunk = Bytes; + type StreamHead = (); +} + +impl From for Error { + fn from(fault: MachineFault) -> Self { + Self::InvalidRequest(match fault { + MachineFault::Abandoned => "messages host driver was abandoned".into(), + MachineFault::Protocol(message) => format!("messages {message}"), + MachineFault::Mismatch => "invalid messages host operation result".into(), + }) + } +} + +pub type MessagesHost = HostChannel; +pub type MessagesMachine = RouteMachine; + +/// Whether this route serves the request, decided before any callback runs so a host +/// can still run its own path. +pub fn supports(model: &str, custom_llm_provider: Option<&str>, stream: bool) -> bool { + let provider = get_custom_llm_provider(model, custom_llm_provider) + .map(|resolved| resolved.custom_llm_provider) + .or(custom_llm_provider); + match provider { + Some(ANTHROPIC_MESSAGES_PROVIDER) => true, + Some(provider) => !stream && messages_provider_config(provider).is_some(), + None => false, + } +} + +/// The in-process host for a request already in hand. It answers projection once and +/// observes nothing. +pub struct LocalMessagesHost { + call: Mutex>, +} + +impl LocalMessagesHost { + pub fn new(call: MessagesCall) -> Self { + Self { + call: Mutex::new(Some(call)), + } + } +} + +impl Host for LocalMessagesHost { + async fn route(&self, op: MessagesOp) -> Result { + match op { + MessagesOp::ProjectRequest => self + .call + .lock() + .unwrap_or_else(|error| error.into_inner()) + .take() + .map(|call| MessagesOpResult::Request(Box::new(call))) + .ok_or_else(|| { + Error::InvalidRequest("messages request was already projected".into()) + }), + } + } +} + +pub fn messages_machine() -> MessagesMachine { + RouteMachine::new(|host| Box::pin(execute(host))) +} + +async fn execute(host: MessagesHost) -> Result { + let MessagesOpResult::Request(call) = host.route(MessagesOp::ProjectRequest).await?; + let stream = call.streams(); + let request = prepare_provider_request(MessagesRequest { + model: &call.model, + body: Value::Object(call.body.clone()), + api_key: call.api_key.as_deref(), + api_base: call.api_base.as_deref(), + custom_llm_provider: call.custom_llm_provider.as_deref(), + extra_headers: call.extra_headers.clone(), + timeout: call.timeout, + })?; + if stream && request.provider != ANTHROPIC_MESSAGES_PROVIDER { + return Err(Error::Unsupported("streaming messages for this provider")); + } + let context = RequestContext { + model: request.model.clone(), + custom_llm_provider: request.provider.clone(), + optional_params: Value::Object( + call.body + .iter() + .filter(|(name, _)| !matches!(name.as_str(), "model" | "messages")) + .map(|(name, value)| (name.clone(), value.clone())) + .collect(), + ), + secret_fields: Vec::new(), + api_key: call.api_key.clone().map(SecretValue::new), + }; + let wire = host + .before_send( + WireRequest { + url: request.url, + headers: request.upstream_headers, + body: request.body, + }, + context, + ) + .await?; + let response = send(&wire.url, &wire.headers, &wire.body, request.timeout).await?; + if !response.status().is_success() { + return Err(provider_error(response).await); + } + if stream { + return relay(&host, response).await; + } + let text = response.text().await.map_err(network)?; + host.emit(MachineEvent::ResponseReceived { + raw: RawResponse { body: text.clone() }, + }) + .await?; + decode_response(request.config, &request.model, &text) + .map(|message| MessagesOutput::Message(Box::new(message))) +} + +/// Hands each upstream chunk to the caller as it arrives. A caller that stops reading +/// ends the upstream read, and the call completes with what it delivered. +async fn relay( + host: &MessagesHost, + mut response: reqwest::Response, +) -> Result { + if host.open(()).await? == Demand::Detached { + return Ok(MessagesOutput::Streamed); + } + while let Some(chunk) = response.chunk().await.map_err(network)? { + if host.deliver(chunk).await? == Demand::Detached { + break; + } + } + Ok(MessagesOutput::Streamed) +} diff --git a/litellm-rust/crates/core/src/ocr/client.rs b/litellm-rust/crates/core/src/ocr/client.rs index 03782d91f24..c05622932b1 100644 --- a/litellm-rust/crates/core/src/ocr/client.rs +++ b/litellm-rust/crates/core/src/ocr/client.rs @@ -12,7 +12,7 @@ pub async fn perform( client: &OcrClient, request: LiteLLMOcrRequest, ) -> Result { - litellm_callbacks::run::run(ocr_machine(client.clone()), &LocalOcrHost::new(request)).await + litellm_host::run::run(ocr_machine(client.clone()), &LocalOcrHost::new(request)).await } pub async fn ocr(request: LiteLLMOcrRequest) -> Result { diff --git a/litellm-rust/crates/core/src/ocr/handler.rs b/litellm-rust/crates/core/src/ocr/handler.rs index 33cb8a8d32a..bbf9cfa0e02 100644 --- a/litellm-rust/crates/core/src/ocr/handler.rs +++ b/litellm-rust/crates/core/src/ocr/handler.rs @@ -1,5 +1,6 @@ use futures_util::future::BoxFuture; -use litellm_callbacks::event::{CallEvent, Passthrough, RawResponse, RequestContext, WireRequest}; +use litellm_auth::SecretValue; +use litellm_host::event::{MachineEvent, RawResponse, RequestContext, WireRequest}; use litellm_llms::{ base_llm::ocr::{ error::Error, @@ -36,6 +37,7 @@ pub(crate) struct OcrCallHooks { custom_llm_provider: &'static str, optional_params: Value, secret_fields: Vec, + api_key: Option, } impl OcrCallHooks { @@ -51,28 +53,25 @@ impl OcrCallHooks { .filter(|name| is_secret_param(name)) .cloned() .collect(), + api_key: request.connection.api_key.clone(), } } } impl CallHooks for OcrCallHooks { - fn before_send( - &self, - wire: WireRequest, - passthrough_fields: Passthrough, - ) -> BoxFuture<'_, Result> { + fn before_send(&self, wire: WireRequest) -> BoxFuture<'_, Result> { let context = RequestContext { model: self.model.clone(), custom_llm_provider: self.custom_llm_provider.into(), optional_params: self.optional_params.clone(), - passthrough_fields, secret_fields: self.secret_fields.clone(), + api_key: self.api_key.clone(), }; Box::pin(self.host.before_send(wire, context)) } fn response_received<'a>(&'a self, body: &'a [u8]) -> BoxFuture<'a, Result<(), Error>> { - Box::pin(self.host.emit(CallEvent::ResponseReceived { + Box::pin(self.host.emit(MachineEvent::ResponseReceived { raw: RawResponse { body: String::from_utf8_lossy(body).into_owned(), }, diff --git a/litellm-rust/crates/core/src/ocr/mod.rs b/litellm-rust/crates/core/src/ocr/mod.rs index e7f77acc3f8..c977f721a70 100644 --- a/litellm-rust/crates/core/src/ocr/mod.rs +++ b/litellm-rust/crates/core/src/ocr/mod.rs @@ -21,8 +21,8 @@ mod cohere_tests; #[path = "../../tests/deepseek_ocr.rs"] mod deepseek_tests; #[cfg(test)] -#[path = "../../tests/ocr/passthrough.rs"] -mod passthrough_tests; +#[path = "../../tests/ocr/document.rs"] +mod document_tests; #[cfg(test)] #[path = "../../tests/reducto_ocr.rs"] mod reducto_tests; diff --git a/litellm-rust/crates/core/src/ocr/prepare.rs b/litellm-rust/crates/core/src/ocr/prepare.rs index 24c3f43e2b4..8ac038290b7 100644 --- a/litellm-rust/crates/core/src/ocr/prepare.rs +++ b/litellm-rust/crates/core/src/ocr/prepare.rs @@ -1,4 +1,4 @@ -use litellm_auth::{InputSource, Sourced}; +use litellm_auth::{InputSource, SecretValue, Sourced}; use litellm_llms::base_llm::ocr::transformation::{ OcrConnection, OcrCredentialInputs, PreparedOcrRequest, credential_env, }; @@ -22,7 +22,7 @@ pub(crate) fn prepare_request( .config .get_api_key_env_var() .and_then(credential_env) - .map(|value| Sourced::new(value, InputSource::Environment)) + .map(|value| Sourced::new(SecretValue::new(value), InputSource::Environment)) }) }); let dynamic_api_base = credentials.dynamic_api_base.or_else(|| { diff --git a/litellm-rust/crates/core/src/ocr/provider_config.rs b/litellm-rust/crates/core/src/ocr/provider_config.rs index d12b8cfee95..14b34ea4564 100644 --- a/litellm-rust/crates/core/src/ocr/provider_config.rs +++ b/litellm-rust/crates/core/src/ocr/provider_config.rs @@ -277,12 +277,18 @@ mod tests { #[test] fn connection_resolution_preserves_dynamic_precedence_and_input_sources() { let connection = OcrConfigKind::Mistral.resolve_connection_params(OcrCredentialInputs { - api_key: Some(Sourced::new("explicit-key".into(), InputSource::Deployment)), + api_key: Some(Sourced::new( + litellm_auth::SecretValue::new("explicit-key"), + InputSource::Deployment, + )), api_base: Some(Sourced::new( "https://explicit.test".into(), InputSource::Deployment, )), - dynamic_api_key: Some(Sourced::new("dynamic-key".into(), InputSource::Environment)), + dynamic_api_key: Some(Sourced::new( + litellm_auth::SecretValue::new("dynamic-key"), + InputSource::Environment, + )), dynamic_api_base: Some(Sourced::new( "https://dynamic.test".into(), InputSource::Request, @@ -292,7 +298,7 @@ mod tests { connection .api_key .as_ref() - .map(|value| value.value().as_str()), + .map(|value| value.value().expose()), Some("dynamic-key") ); assert_eq!( @@ -318,22 +324,31 @@ mod tests { fn empty_or_missing_dynamic_credentials_preserve_explicit_values( #[case] dynamic_value: Option<&str>, ) { - let dynamic = + let dynamic_key = dynamic_value.map(|value| { + Sourced::new( + litellm_auth::SecretValue::new(value), + InputSource::Environment, + ) + }); + let dynamic_base = dynamic_value.map(|value| Sourced::new(value.into(), InputSource::Environment)); let connection = OcrConfigKind::Mistral.resolve_connection_params(OcrCredentialInputs { - api_key: Some(Sourced::new("explicit-key".into(), InputSource::Deployment)), + api_key: Some(Sourced::new( + litellm_auth::SecretValue::new("explicit-key"), + InputSource::Deployment, + )), api_base: Some(Sourced::new( "https://explicit.test".into(), InputSource::Deployment, )), - dynamic_api_key: dynamic.clone(), - dynamic_api_base: dynamic, + dynamic_api_key: dynamic_key, + dynamic_api_base: dynamic_base, }); assert_eq!( connection .api_key .as_ref() - .map(|value| value.value().as_str()), + .map(|value| value.value().expose()), Some("explicit-key") ); assert_eq!( @@ -356,11 +371,18 @@ mod tests { ) { let connection = OcrConfigKind::AzureDocumentIntelligence.resolve_connection_params( OcrCredentialInputs { - api_key: explicit_key - .map(|value| Sourced::new(value.into(), InputSource::Deployment)), + api_key: explicit_key.map(|value| { + Sourced::new( + litellm_auth::SecretValue::new(value), + InputSource::Deployment, + ) + }), api_base: explicit_base .map(|value| Sourced::new(value.into(), InputSource::Deployment)), - dynamic_api_key: Some(Sourced::new("dynamic-key".into(), InputSource::Environment)), + dynamic_api_key: Some(Sourced::new( + litellm_auth::SecretValue::new("dynamic-key"), + InputSource::Environment, + )), dynamic_api_base: Some(Sourced::new( "https://dynamic.test".into(), InputSource::Deployment, @@ -371,7 +393,7 @@ mod tests { connection .api_key .as_ref() - .map(|value| value.value().as_str()), + .map(|value| value.value().expose()), explicit_key.map(|_| "dynamic-key") ); assert_eq!( diff --git a/litellm-rust/crates/core/src/ocr/route.rs b/litellm-rust/crates/core/src/ocr/route.rs index ac4237651da..bfc8c5ca965 100644 --- a/litellm-rust/crates/core/src/ocr/route.rs +++ b/litellm-rust/crates/core/src/ocr/route.rs @@ -1,8 +1,9 @@ use std::sync::{Arc, Mutex}; use litellm_auth::ResolvedCredential; -use litellm_callbacks::{ +use litellm_host::{ event::{CallEvent, RequestContext, WireRequest}, + machine::{HostChannel, HostTokenProvider, MachineFault, RouteMachine, TokenRoute}, route::Route, }; use litellm_llms::{ @@ -11,10 +12,7 @@ use litellm_llms::{ }; use super::handler::perform_ocr_request; -use crate::{ - machine::{HostChannel, HostTokenProvider, MachineFault, RouteMachine, TokenRoute}, - ocr::types::{LiteLLMOcrRequest, OcrDocumentInput, OcrFileContent, ResolvedOcrRequest}, -}; +use crate::ocr::types::{LiteLLMOcrRequest, OcrDocumentInput, OcrFileContent, ResolvedOcrRequest}; #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub enum OcrOp { @@ -39,6 +37,8 @@ impl Route for Ocr { type Error = Error; type Op = OcrOp; type OpResult = OcrOpResult; + type Chunk = std::convert::Infallible; + type StreamHead = std::convert::Infallible; } impl TokenRoute for Ocr { @@ -54,16 +54,6 @@ impl TokenRoute for Ocr { } } -impl From for Error { - fn from(fault: MachineFault) -> Self { - Self::InvalidRequest(match fault { - MachineFault::Abandoned => "OCR host driver was abandoned".into(), - MachineFault::Protocol(message) => format!("OCR {message}"), - MachineFault::Mismatch => "invalid OCR host operation result".into(), - }) - } -} - pub type OcrHost = HostChannel; pub type OcrMachine = RouteMachine; @@ -173,7 +163,7 @@ impl LocalOcrHost { } } -impl litellm_callbacks::host::Host for LocalOcrHost { +impl litellm_host::host::Host for LocalOcrHost { async fn route(&self, op: OcrOp) -> Result { match op { OcrOp::ProjectRequest => self diff --git a/litellm-rust/crates/core/src/ocr/types.rs b/litellm-rust/crates/core/src/ocr/types.rs index 75202ed52a5..6316088dec8 100644 --- a/litellm-rust/crates/core/src/ocr/types.rs +++ b/litellm-rust/crates/core/src/ocr/types.rs @@ -1,7 +1,7 @@ use std::{collections::BTreeMap, path::PathBuf, time::Duration}; use bytes::Bytes; -use litellm_auth::{InputSource, TokenProviderHandle}; +use litellm_auth::{InputSource, SecretValue, TokenProviderHandle}; use litellm_core_utils::call_arguments::CallArguments; use litellm_llms::base_llm::ocr::{ error::Error, @@ -56,7 +56,7 @@ pub struct OcrFileContent { /// credentials, and per-field provenance in `input_sources`. #[derive(Clone, Debug, Default)] pub struct OcrConnectionInputs { - pub api_key: Option, + pub api_key: Option, pub api_base: Option, pub extra_headers: Map, pub timeout: Option, @@ -237,6 +237,16 @@ mod tests { .unwrap() } + #[test] + fn connection_inputs_debug_hides_the_api_key() { + let inputs = OcrConnectionInputs { + api_key: Some(SecretValue::new("caller-api-key")), + ..OcrConnectionInputs::default() + }; + + assert!(!format!("{inputs:?}").contains("caller-api-key")); + } + #[test] fn from_inputs_applies_connection_overrides_with_field_sources() { let request = LiteLLMOcrRequest::from_inputs( @@ -245,7 +255,7 @@ mod tests { None, Default::default(), OcrConnectionInputs { - api_key: Some(" key ".into()), + api_key: Some(SecretValue::new(" key ")), api_base: Some("".into()), extra_headers: json!({"x-a": "1"}).as_object().unwrap().clone(), timeout: Some(Duration::from_secs(7)), @@ -259,7 +269,7 @@ mod tests { .unwrap(); let api_key = request.credentials.api_key.as_ref().unwrap(); - assert_eq!(api_key.clone().into_value(), "key"); + assert_eq!(api_key.value().expose(), "key"); assert_eq!(api_key.source(), InputSource::Request); assert!(request.credentials.api_base.is_none()); assert_eq!( diff --git a/litellm-rust/crates/core/src/ocr/wire.rs b/litellm-rust/crates/core/src/ocr/wire.rs index 29345e38885..b9c60f57e3c 100644 --- a/litellm-rust/crates/core/src/ocr/wire.rs +++ b/litellm-rust/crates/core/src/ocr/wire.rs @@ -1,6 +1,6 @@ use std::{collections::BTreeMap, time::Duration}; -use litellm_auth::InputSource; +use litellm_auth::{InputSource, SecretValue}; use litellm_llms::base_llm::ocr::{ error::Error, transformation::{OcrDocument, decode_request_value}, @@ -44,7 +44,7 @@ pub fn consumed_optional_param_names( pub struct OcrWireRequest { pub model: String, pub document: D, - pub api_key: Option, + pub api_key: Option, pub api_base: Option, pub custom_llm_provider: Option, pub extra_headers: Option>, diff --git a/litellm-rust/crates/core/tests/azure_document_intelligence_ocr.rs b/litellm-rust/crates/core/tests/azure_document_intelligence_ocr.rs index 01a4e5efb3b..3cbe6fe3159 100644 --- a/litellm-rust/crates/core/tests/azure_document_intelligence_ocr.rs +++ b/litellm-rust/crates/core/tests/azure_document_intelligence_ocr.rs @@ -1,4 +1,4 @@ -use litellm_callbacks::event::CallEvent; +use litellm_host::event::{CallEvent, MachineEvent}; use litellm_llms::base_llm::ocr::error::Error; use rstest::rstest; use serde_json::{Value, json}; @@ -69,7 +69,7 @@ async fn rejects_invalid_pages_features_and_format( let result = decode_request(OcrWireRequest { model: "azure_ai/doc-intelligence/prebuilt-read".into(), document: json!({"type":"document_url","document_url":"https://example.com/a.pdf"}), - api_key: Some("key".into()), + api_key: Some(litellm_auth::SecretValue::new("key")), api_base: Some(base), custom_llm_provider: None, extra_headers: None, @@ -263,7 +263,7 @@ async fn accepted_response_emits_response_received_before_polling() { json!({}), )) .with_observer(move |event| { - let CallEvent::ResponseReceived { raw } = event else { + let CallEvent::Machine(MachineEvent::ResponseReceived { raw }) = event else { return; }; match request_count.lock().unwrap().len() { @@ -466,7 +466,7 @@ async fn model_id_is_encoded_and_dot_segments_are_rejected() { mod transformation { use std::sync::{Arc, Mutex}; - use litellm_callbacks::event::CallEvent; + use litellm_host::event::{CallEvent, MachineEvent}; use litellm_llms::base_llm::ocr::transformation::OcrDocument; use serde_json::{Value, json}; @@ -646,7 +646,7 @@ mod transformation { json!({}), )) .with_observer(move |event| { - if let CallEvent::ResponseReceived { raw } = event { + if let CallEvent::Machine(MachineEvent::ResponseReceived { raw }) = event { observed .lock() .unwrap() diff --git a/litellm-rust/crates/core/tests/ocr.rs b/litellm-rust/crates/core/tests/ocr.rs index 1f591d74d5d..41a650945bc 100644 --- a/litellm-rust/crates/core/tests/ocr.rs +++ b/litellm-rust/crates/core/tests/ocr.rs @@ -1,7 +1,7 @@ use std::sync::{Arc, Mutex}; -use litellm_callbacks::{ - event::{CallEvent, WireRequest}, +use litellm_host::{ + event::{CallEvent, MachineEvent, WireRequest}, host::{Host, HostOp, HostResult}, machine::{HostFailure, Machine, MachineStep}, }; @@ -81,7 +81,7 @@ fn request_boundary_selects_mistral_and_rejects_unknown_providers() { let request = OcrWireRequest { model: "mistral/model".into(), document: json!({"type":"document_url","document_url":"https://example.com/doc.pdf"}), - api_key: Some("key".into()), + api_key: Some(litellm_auth::SecretValue::new("key")), api_base: None, custom_llm_provider: None, extra_headers: None, @@ -97,7 +97,7 @@ fn request_boundary_selects_mistral_and_rejects_unknown_providers() { decode_request(OcrWireRequest { model: "model".into(), document: json!({"type":"document_url","document_url":"https://example.com/doc.pdf"}), - api_key: Some("key".into()), + api_key: Some(litellm_auth::SecretValue::new("key")), api_base: None, custom_llm_provider: Some("unknown".into()), extra_headers: None, @@ -194,7 +194,8 @@ async fn facade_uses_the_injected_http_client() { fn event_name(event: &CallEvent) -> &'static str { match event { - CallEvent::ResponseReceived { .. } => "response", + CallEvent::Started { .. } => "started", + CallEvent::Machine(MachineEvent::ResponseReceived { .. }) => "response", CallEvent::Succeeded { .. } => "success", CallEvent::Failed { .. } => "failure", } @@ -235,7 +236,7 @@ async fn lifecycle_sends_headers_returned_by_the_before_send_operation() { } #[tokio::test] -async fn before_send_context_names_passthrough_fields_and_secrets() { +async fn before_send_context_names_the_route_and_its_secrets() { let (base, _, server) = mock_server(vec![MockResponse::json(json!({"pages":[]}))]).await; let observed = Arc::new(Mutex::new(None)); let captured = observed.clone(); @@ -254,8 +255,6 @@ async fn before_send_context_names_passthrough_fields_and_secrets() { assert_eq!(context.custom_llm_provider, "mistral"); assert_eq!(context.model, "model"); assert_eq!(wire.body["pages"], json!([0])); - assert!(context.passthrough_fields.contains("pages")); - assert!(context.passthrough_fields.contains("document")); assert!(context.secret_fields.is_empty()); assert_eq!(context.optional_params["req_format"], "native"); @@ -279,7 +278,6 @@ async fn before_send_context_names_passthrough_fields_and_secrets() { perform_ocr_with(host).await.unwrap(); server.await.unwrap(); let context = observed.lock().unwrap().take().unwrap(); - assert!(!context.passthrough_fields.contains("document")); assert_eq!(context.secret_fields, ["client_secret"]); } @@ -296,7 +294,7 @@ async fn lifecycle_orders_hooks_and_emits_one_success() { server.await.unwrap(); assert_eq!( *events.lock().unwrap(), - ["before_send", "response", "success"] + ["started", "before_send", "response", "success"] ); assert_eq!(seen.lock().unwrap().len(), 1); } @@ -311,7 +309,10 @@ async fn lifecycle_blocking_prevents_execution_and_emits_one_failure() { ); let error = perform_ocr_with(host).await.unwrap_err(); assert!(matches!(error, OcrError::InvalidRequest(message) if message == "blocked")); - assert_eq!(*events.lock().unwrap(), ["before_send", "failure"]); + assert_eq!( + *events.lock().unwrap(), + ["started", "before_send", "failure"] + ); } #[tokio::test] @@ -330,7 +331,10 @@ async fn upstream_failure_emits_one_terminal_failure() { ); assert!(perform_ocr_with(host).await.is_err()); server.await.unwrap(); - assert_eq!(*events.lock().unwrap(), ["before_send", "failure"]); + assert_eq!( + *events.lock().unwrap(), + ["started", "before_send", "failure"] + ); assert_eq!(seen.lock().unwrap().len(), 1); } @@ -371,6 +375,7 @@ async fn drive_until( intercept(*wire).map(|wire| HostResult::BeforeSend(Box::new(wire))) } HostOp::Emit(event) => { + let event = CallEvent::Machine(event); ops.push(event_name(&event)); host.emit(&event) .await @@ -414,7 +419,7 @@ async fn invalid_provider_response_emits_response_received_before_normalization_ let observed = responses_received.clone(); let host = LocalOcrHost::new(wire_request("mistral/model", &base, json!({}))).with_observer( move |event| { - if let CallEvent::ResponseReceived { raw } = event { + if let CallEvent::Machine(MachineEvent::ResponseReceived { raw }) = event { observed.lock().unwrap().push(raw.body.clone()); } }, @@ -815,7 +820,7 @@ impl Host for CallerTokenHost { async fn before_send( &self, wire: WireRequest, - _: &litellm_callbacks::event::RequestContext, + _: &litellm_host::event::RequestContext, ) -> Result { let is_authorization = |name: &str| name.eq_ignore_ascii_case("authorization"); let authorization = wire @@ -850,7 +855,7 @@ async fn the_callers_azure_token_is_acquired_before_before_send_which_can_still_ trace: Mutex::new(Vec::new()), }; - litellm_callbacks::run::run(ocr_machine(ocr_client()), &host) + litellm_host::run::run(ocr_machine(ocr_client()), &host) .await .unwrap(); server.await.unwrap(); diff --git a/litellm-rust/crates/core/tests/ocr/document.rs b/litellm-rust/crates/core/tests/ocr/document.rs new file mode 100644 index 00000000000..855548dc6bf --- /dev/null +++ b/litellm-rust/crates/core/tests/ocr/document.rs @@ -0,0 +1,152 @@ +use litellm_host::event::WireRequest; +use litellm_llms::base_llm::ocr::error::Error; +use rstest::rstest; +use serde_json::{Value, json}; + +use super::test_support::{ + MockResponse, SERVED_DOCUMENT, document_server, mock_server, perform_ocr_with, request_body, + wire_request_with_document, +}; +use crate::ocr::route::LocalOcrHost; + +#[derive(Clone, Copy, Debug)] +enum Route { + Mistral, + AzureAi, + VertexMistral, + AzureCohereParse, + Cohere, +} + +impl Route { + fn model(self) -> &'static str { + match self { + Self::Mistral => "mistral/model", + Self::AzureAi => "azure_ai/model", + Self::VertexMistral => "vertex_ai/mistral-ocr-maas", + Self::AzureCohereParse => "azure_ai/cohere-parse", + Self::Cohere => "cohere/model", + } + } + + fn document_type(self) -> &'static str { + match self { + Self::Mistral | Self::AzureAi | Self::VertexMistral => "document_url", + Self::AzureCohereParse | Self::Cohere => "image_url", + } + } + + fn options(self) -> Value { + match self { + Self::Mistral | Self::AzureAi => json!({"pages": [0]}), + Self::VertexMistral => json!({"pages": [0], "vertex_project": "project-1"}), + Self::AzureCohereParse | Self::Cohere => json!({"output_format": "markdown"}), + } + } +} + +/// What the host does to the wire request in `before_send`. +#[derive(Clone, Copy, Debug)] +enum Host { + Detached, + ReplacesDocument, +} + +const REPLACED_DOCUMENT: &str = "data:image/png;base64,cmVwbGFjZWQ="; + +impl Host { + fn before_send(self, wire: WireRequest) -> WireRequest { + let Value::Object(fields) = wire.body else { + return wire; + }; + let body = fields + .into_iter() + .map(|(name, value)| match self { + Self::Detached => (name, value), + Self::ReplacesDocument if name == "document" => { + let document_type = value["type"].clone(); + let key = document_type.as_str().unwrap_or_default().to_string(); + (name, json!({"type": document_type, key: REPLACED_DOCUMENT})) + } + Self::ReplacesDocument => (name, value), + }) + .collect(); + WireRequest { + body: Value::Object(body), + ..wire + } + } +} + +struct Sent { + result: Result<(), Error>, + provider_body: Option, +} + +async fn send(route: Route, host: Host, document_base: &str) -> Sent { + let (base, seen, provider) = mock_server(vec![MockResponse::json(json!({"pages": []}))]).await; + let document_type = route.document_type(); + let document = + json!({"type": document_type, document_type: format!("{document_base}/scan.png")}); + let request = wire_request_with_document(route.model(), &base, document, route.options()); + let local = + LocalOcrHost::new(request).with_before_send(move |wire, _| Ok(host.before_send(wire))); + let result = perform_ocr_with(local).await.map(|_| ()); + match result { + Ok(()) => provider.await.unwrap(), + Err(_) => provider.abort(), + } + let provider_body = seen + .lock() + .unwrap() + .first() + .map(|request| request_body(request)); + Sent { + result, + provider_body, + } +} + +fn served_document_uri() -> String { + use base64::Engine; + format!( + "data:image/png;base64,{}", + base64::engine::general_purpose::STANDARD.encode(SERVED_DOCUMENT) + ) +} + +#[rstest] +#[case::azure_ai(Route::AzureAi)] +#[case::vertex_mistral(Route::VertexMistral)] +#[case::azure_cohere_parse(Route::AzureCohereParse)] +#[tokio::test] +async fn inlining_routes_send_the_downloaded_document(#[case] route: Route) { + let (document_base, _documents) = document_server().await; + let sent = send(route, Host::Detached, &document_base).await; + sent.result.unwrap(); + assert_eq!( + sent.provider_body.unwrap()["document"][route.document_type()], + json!(served_document_uri()) + ); +} + +#[rstest] +#[tokio::test] +async fn document_replaced_by_the_host_reaches_the_provider( + #[values( + Route::Mistral, + Route::AzureAi, + Route::VertexMistral, + Route::AzureCohereParse, + Route::Cohere + )] + route: Route, +) { + let (document_base, _documents) = document_server().await; + let sent = send(route, Host::ReplacesDocument, &document_base).await; + sent.result.unwrap(); + assert_eq!( + sent.provider_body.unwrap()["document"][route.document_type()], + json!(REPLACED_DOCUMENT) + ); +} diff --git a/litellm-rust/crates/core/tests/ocr/passthrough.rs b/litellm-rust/crates/core/tests/ocr/passthrough.rs deleted file mode 100644 index 0273b48664d..00000000000 --- a/litellm-rust/crates/core/tests/ocr/passthrough.rs +++ /dev/null @@ -1,282 +0,0 @@ -use std::{ - collections::BTreeSet, - sync::{Arc, Mutex}, -}; - -use litellm_callbacks::event::{RequestContext, WireRequest}; -use litellm_llms::base_llm::ocr::error::Error; -use rstest::rstest; -use rstest_reuse::{self, apply, template}; -use serde_json::{Map, Value, json}; - -use super::test_support::{ - MockResponse, SERVED_DOCUMENT, document_server, mock_server, perform_ocr_with, request_body, - wire_request_with_document, -}; -use crate::ocr::route::LocalOcrHost; - -#[derive(Clone, Copy, Debug)] -enum Route { - Mistral, - AzureAi, - VertexMistral, - AzureCohereParse, - Cohere, -} - -impl Route { - fn model(self) -> &'static str { - match self { - Self::Mistral => "mistral/model", - Self::AzureAi => "azure_ai/model", - Self::VertexMistral => "vertex_ai/mistral-ocr-maas", - Self::AzureCohereParse => "azure_ai/cohere-parse", - Self::Cohere => "cohere/model", - } - } - - fn document_type(self) -> &'static str { - match self { - Self::Mistral | Self::AzureAi | Self::VertexMistral => "document_url", - Self::AzureCohereParse | Self::Cohere => "image_url", - } - } - - fn options(self) -> Value { - match self { - Self::Mistral | Self::AzureAi => json!({"pages": [0]}), - Self::VertexMistral => json!({"pages": [0], "vertex_project": "project-1"}), - Self::AzureCohereParse | Self::Cohere => json!({"output_format": "markdown"}), - } - } -} - -#[derive(Clone, Copy, Debug)] -enum Source { - Inline, - Remote, - RemoteWithExtraField, -} - -/// What the host does to the wire request in `before_send`. -#[derive(Clone, Copy, Debug)] -enum Host { - Detached, - /// What `litellm-callbacks-legacy` does before `pre_call`: every passthrough body key - /// is replaced by the caller's own value. - Realiasing, - ReplacesDocument, -} - -const REPLACED_DOCUMENT: &str = "data:image/png;base64,cmVwbGFjZWQ="; - -impl Host { - fn before_send( - self, - caller: &Map, - wire: WireRequest, - context: &RequestContext, - ) -> WireRequest { - let Value::Object(fields) = wire.body else { - return wire; - }; - let body = fields - .into_iter() - .map(|(name, value)| match self { - Self::Detached => (name, value), - Self::Realiasing => { - let aliased = context - .passthrough_fields - .contains(&name) - .then(|| caller.get(&name).cloned()) - .flatten() - .unwrap_or(value); - (name, aliased) - } - Self::ReplacesDocument if name == "document" => { - let document_type = value["type"].clone(); - let key = document_type.as_str().unwrap_or_default().to_string(); - (name, json!({"type": document_type, key: REPLACED_DOCUMENT})) - } - Self::ReplacesDocument => (name, value), - }) - .collect(); - WireRequest { - body: Value::Object(body), - ..wire - } - } -} - -struct Sent { - caller: Map, - result: Result<(), Error>, - before_send: Option<(WireRequest, RequestContext)>, - provider_body: Option, -} - -fn caller_document(route: Route, source: Source, document_base: &str) -> Value { - let document_type = route.document_type(); - let remote = format!("{document_base}/scan.png"); - match source { - Source::Inline => { - json!({"type": document_type, document_type: "data:image/png;base64,YWJj"}) - } - Source::Remote => json!({"type": document_type, document_type: remote}), - Source::RemoteWithExtraField => { - json!({"type": document_type, document_type: remote, "document_name": "scan.png"}) - } - } -} - -async fn send(route: Route, source: Source, host: Host, document_base: &str) -> Sent { - let (base, seen, provider) = mock_server(vec![MockResponse::json(json!({"pages": []}))]).await; - let document = caller_document(route, source, document_base); - let caller: Map = route - .options() - .as_object() - .unwrap() - .clone() - .into_iter() - .chain([("document".to_string(), document.clone())]) - .collect(); - let observed = Arc::new(Mutex::new(None)); - let captured = observed.clone(); - let host_caller = caller.clone(); - let request = wire_request_with_document(route.model(), &base, document, route.options()); - let local = LocalOcrHost::new(request).with_before_send(move |wire, context| { - *captured.lock().unwrap() = Some((wire.clone(), context.clone())); - Ok(host.before_send(&host_caller, wire, context)) - }); - let result = perform_ocr_with(local).await.map(|_| ()); - match result { - Ok(()) => provider.await.unwrap(), - Err(_) => provider.abort(), - } - let provider_body = seen - .lock() - .unwrap() - .first() - .map(|request| request_body(request)); - let before_send = observed.lock().unwrap().take(); - Sent { - caller, - result, - before_send, - provider_body, - } -} - -fn served_document_uri() -> String { - use base64::Engine; - format!( - "data:image/png;base64,{}", - base64::engine::general_purpose::STANDARD.encode(SERVED_DOCUMENT) - ) -} - -#[template] -#[rstest] -fn every_route_and_source( - #[values( - Route::Mistral, - Route::AzureAi, - Route::VertexMistral, - Route::AzureCohereParse, - Route::Cohere - )] - route: Route, - #[values(Source::Inline, Source::Remote, Source::RemoteWithExtraField)] source: Source, -) { -} - -#[template] -#[rstest] -fn every_route( - #[values( - Route::Mistral, - Route::AzureAi, - Route::VertexMistral, - Route::AzureCohereParse, - Route::Cohere - )] - route: Route, -) { -} - -#[template] -#[rstest] -#[case::azure_ai(Route::AzureAi)] -#[case::vertex_mistral(Route::VertexMistral)] -#[case::azure_cohere_parse(Route::AzureCohereParse)] -fn inlining_routes(#[case] route: Route) {} - -#[apply(every_route_and_source)] -#[tokio::test] -async fn passthrough_fields_are_exactly_the_caller_values_sent_unchanged( - route: Route, - source: Source, -) { - let (document_base, _documents) = document_server().await; - let sent = send(route, source, Host::Detached, &document_base).await; - sent.result.unwrap(); - let (wire, context) = sent.before_send.unwrap(); - let passthrough: BTreeSet<&str> = context.passthrough_fields.iter().collect(); - let unchanged: BTreeSet<&str> = sent - .caller - .iter() - .filter(|(name, value)| wire.body.get(name.as_str()) == Some(*value)) - .map(|(name, _)| name.as_str()) - .collect(); - assert_eq!( - passthrough, - unchanged, - "body: {:#}\ncaller: {:#}", - wire.body, - Value::Object(sent.caller.clone()) - ); -} - -#[apply(every_route_and_source)] -#[tokio::test] -async fn realiasing_leaves_the_provider_request_unchanged(route: Route, source: Source) { - let (document_base, _documents) = document_server().await; - let detached = send(route, source, Host::Detached, &document_base).await; - let realiased = send(route, source, Host::Realiasing, &document_base).await; - detached.result.unwrap(); - realiased.result.unwrap(); - assert_eq!(realiased.provider_body, detached.provider_body); -} - -#[apply(inlining_routes)] -#[tokio::test] -async fn inlining_routes_send_the_downloaded_document( - route: Route, - #[values(Host::Detached, Host::Realiasing)] host: Host, -) { - let (document_base, _documents) = document_server().await; - let sent = send(route, Source::Remote, host, &document_base).await; - sent.result.unwrap(); - assert_eq!( - sent.provider_body.unwrap()["document"][route.document_type()], - json!(served_document_uri()) - ); -} - -#[apply(every_route)] -#[tokio::test] -async fn document_replaced_by_the_host_reaches_the_provider(route: Route) { - let (document_base, _documents) = document_server().await; - let sent = send( - route, - Source::Remote, - Host::ReplacesDocument, - &document_base, - ) - .await; - sent.result.unwrap(); - assert_eq!( - sent.provider_body.unwrap()["document"][route.document_type()], - json!(REPLACED_DOCUMENT) - ); -} diff --git a/litellm-rust/crates/core/tests/ocr/support.rs b/litellm-rust/crates/core/tests/ocr/support.rs index f3adf27cfa6..b368a754656 100644 --- a/litellm-rust/crates/core/tests/ocr/support.rs +++ b/litellm-rust/crates/core/tests/ocr/support.rs @@ -1,7 +1,7 @@ use std::sync::{Arc, Mutex}; use futures_util::future::BoxFuture; -use litellm_callbacks::event::{Passthrough, WireRequest}; +use litellm_host::event::WireRequest; use litellm_llms::{ base_llm::ocr::{error::Error, transformation::LiteLLMOcrResponse}, custom_httpx::llm_http_handler::{CallHooks, OcrClient}, @@ -23,11 +23,7 @@ use crate::ocr::{ pub(crate) struct NoHooks; impl CallHooks for NoHooks { - fn before_send( - &self, - wire: WireRequest, - _passthrough_fields: Passthrough, - ) -> BoxFuture<'_, Result> { + fn before_send(&self, wire: WireRequest) -> BoxFuture<'_, Result> { Box::pin(async move { Ok(wire) }) } @@ -49,7 +45,7 @@ pub(crate) async fn perform_ocr(request: LiteLLMOcrRequest) -> Result Result { - litellm_callbacks::run::run(ocr_machine(ocr_client()), &host).await + litellm_host::run::run(ocr_machine(ocr_client()), &host).await } pub(crate) fn wire_request(model: &str, base: &str, options: Value) -> LiteLLMOcrRequest { @@ -70,7 +66,7 @@ pub(crate) fn wire_request_with_document( decode_request(OcrWireRequest { model: model.into(), document, - api_key: Some("test-key".into()), + api_key: Some(litellm_auth::SecretValue::new("test-key")), api_base: Some(base.into()), custom_llm_provider: None, extra_headers: None, diff --git a/litellm-rust/crates/core/tests/reducto_ocr.rs b/litellm-rust/crates/core/tests/reducto_ocr.rs index 59891b16e90..83e7754122b 100644 --- a/litellm-rust/crates/core/tests/reducto_ocr.rs +++ b/litellm-rust/crates/core/tests/reducto_ocr.rs @@ -1,4 +1,4 @@ -use litellm_callbacks::event::{CallEvent, WireRequest}; +use litellm_host::event::{CallEvent, MachineEvent, WireRequest}; use litellm_llms::base_llm::ocr::{error::Error, transformation::OcrDocument}; use rstest::rstest; use serde_json::{Value, json}; @@ -139,7 +139,7 @@ async fn response_received_stays_after_reducto_upload_and_parse() { let request_count = seen.clone(); let host = LocalOcrHost::new(wire_request("reducto/parse-v3", &base, json!({}))).with_observer( move |event| { - if let CallEvent::ResponseReceived { raw } = event { + if let CallEvent::Machine(MachineEvent::ResponseReceived { raw }) = event { assert_eq!(request_count.lock().unwrap().len(), 2); assert_eq!(raw.body, r#"{"result":{"chunks":[]}}"#); } @@ -351,7 +351,7 @@ async fn guardrail_rewrites_document_before_upload() { } mod transformation { - use litellm_callbacks::event::{CallEvent, WireRequest}; + use litellm_host::event::{CallEvent, MachineEvent, WireRequest}; use litellm_llms::{ base_llm::ocr::transformation::{BaseOcrConfig, OcrConnection, OcrRequestContext}, reducto::ocr::transformation::*, @@ -506,7 +506,7 @@ mod transformation { let request_count = seen.clone(); let host = LocalOcrHost::new(wire_request("reducto/parse-v3", &base, json!({}))) .with_observer(move |event| { - if let CallEvent::ResponseReceived { raw } = event { + if let CallEvent::Machine(MachineEvent::ResponseReceived { raw }) = event { assert_eq!(request_count.lock().unwrap().len(), 2); assert_eq!(raw.body, r#"{"result":{"chunks":[]}}"#); } diff --git a/litellm-rust/crates/host-python/AGENTS.md b/litellm-rust/crates/host-python/AGENTS.md index a3fdd2340b3..5aca13eeb18 100644 --- a/litellm-rust/crates/host-python/AGENTS.md +++ b/litellm-rust/crates/host-python/AGENTS.md @@ -1,9 +1,10 @@ - Target invariants; implementation and runtime validation may lag these rules -- Keep this crate the CPython runtime adapter and nothing more: Serde marshalling, interpreter detachment, tokio/asyncio glue, the `Execution` handle, the call driver and the `CallbackAdapter`/`RouteHost` traits - - No LiteLLM domain dependencies beyond `litellm-callbacks`: no route types, no `Logging` policy, no public API registration, no cdylib build features +- Keep this crate the CPython runtime adapter and nothing more: Serde marshalling, interpreter detachment, tokio/asyncio glue, the `Execution` handle, the call driver and the `PythonLifecycle`/`RouteHost` traits + - No LiteLLM domain dependencies beyond `litellm-host`: no route types, no `Logging` policy, no public API registration, no cdylib build features - The driver emits `Succeeded` or `Failed` exactly once and never dispatches after a cancellation; which Python objects consume those events is the adapter's business - `RouteHost::invoke` receives the keyword view the adapter's `begin` returned, not the caller's dict; a route host that projects from it inherits that adapter's rewrites (for the legacy adapter: setup, deployment hooks, credential inheritance) - - A failure that surfaces inside the call, including a host op the call asked for, is mapped through the route's `map_failure`; a failure in `begin` or `after_success` is raised as is + - A native failure, including one a host op returns as `InvokeError::Native`, is classified exactly once through the route's `classify`; a Python exception raised inside the call, and a failure in `begin` or `after_success`, is raised as is + - A failing `classify` is raised with the native error's text as its `__context__`, never swallowed - Use standard PyO3 ownership and conversion APIs - Prefer `Bound<'py, T>` for attached operations/results, `Py` for retention; binding/unbinding does not copy payloads - Use `pythonize` for selected Serde data, never a JSON-text round trip; share conversion with `Pythonized` diff --git a/litellm-rust/crates/host-python/Cargo.toml b/litellm-rust/crates/host-python/Cargo.toml index ae0cebada59..e2c83fe1081 100644 --- a/litellm-rust/crates/host-python/Cargo.toml +++ b/litellm-rust/crates/host-python/Cargo.toml @@ -7,7 +7,7 @@ repository.workspace = true [dependencies] futures-util.workspace = true -litellm-callbacks.workspace = true +litellm-host.workspace = true pyo3.workspace = true pyo3-async-runtimes.workspace = true pythonize.workspace = true diff --git a/litellm-rust/crates/host-python/src/adapter.rs b/litellm-rust/crates/host-python/src/adapter.rs index f1bc3142a25..3a4cb49be4d 100644 --- a/litellm-rust/crates/host-python/src/adapter.rs +++ b/litellm-rust/crates/host-python/src/adapter.rs @@ -1,5 +1,5 @@ -use litellm_callbacks::event::{CallEvent, RequestContext, Timing, WireRequest}; -use litellm_callbacks::route::Route; +use litellm_host::event::{FailureOrigin, MachineEvent, RequestContext, Timing, WireRequest}; +use litellm_host::route::Route; use pyo3::exceptions::PyRuntimeError; use pyo3::gc::{PyTraverseError, PyVisit}; use pyo3::prelude::*; @@ -11,7 +11,7 @@ pub fn missing_state() -> PyErr { /// What an adapter step produced: either the value the driver asked for, or a Python /// awaitable the driver hands back to the caller's task before asking again. -pub enum AdapterStep { +pub enum LifecycleStep { Await(Py), Arguments(Py), Wire(Box), @@ -19,64 +19,97 @@ pub enum AdapterStep { Done, } -/// The host-typed value the driver attaches to a terminal event. -pub enum PublicValue<'a> { - Response(&'a Py), - Error(&'a PyErr), +/// What a lifecycle observes: the driver's start, the machine's own events, and one +/// terminal event carrying the public value the caller receives. +pub enum LifecycleEvent<'a> { + Started { + start_time: f64, + }, + Machine(&'a MachineEvent), + Succeeded { + timing: Timing, + response: &'a Py, + }, + Failed { + timing: Timing, + origin: FailureOrigin, + error: &'a PyErr, + }, } /// One consumer of a call's lifecycle on the Python side. The driver calls the steps in /// order: `begin` before the machine starts, `before_send` and `emit` while it runs, /// `after_success` and one terminal `emit` after it completes. Whenever a step returns -/// [`AdapterStep::Await`], the driver awaits it in the caller's task and continues the +/// [`LifecycleStep::Await`], the driver awaits it in the caller's task and continues the /// same step through `resume`. /// /// A step that fails with an ordinary exception fails the call with that exception, /// except on a terminal event, where the adapter is expected to report and swallow its /// own errors. An exception that is not a `PyException`, such as a cancellation, ends /// the call without further dispatch. -pub trait CallbackAdapter: Send + Sync { +pub trait PythonLifecycle: Send + Sync { fn begin( &mut self, py: Python<'_>, arguments: Py, started_at: f64, - ) -> PyResult; + ) -> PyResult; fn before_send( &mut self, py: Python<'_>, wire: Box, context: &RequestContext, - ) -> PyResult; + ) -> PyResult; fn after_success( &mut self, py: Python<'_>, response: Py, timing: Timing, - ) -> PyResult; + ) -> PyResult; - fn emit( - &mut self, - py: Python<'_>, - event: &CallEvent, - public: Option>, - ) -> PyResult; + fn emit(&mut self, py: Python<'_>, event: LifecycleEvent<'_>) -> PyResult; - fn resume(&mut self, py: Python<'_>, result: PyResult>) -> PyResult; + /// The call streams and its stream was handed to the caller. The caller is not + /// inside an await here, so this step and `delivered` cannot suspend. + fn opened(&mut self, py: Python<'_>) -> PyResult<()>; + + /// One chunk of an open stream is about to reach the caller. + fn delivered(&mut self, py: Python<'_>, chunk: &Py) -> PyResult<()>; + + fn resume(&mut self, py: Python<'_>, result: PyResult>) -> PyResult; fn close(&mut self, py: Python<'_>); fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError>; } -/// The Python side of one route: answers the route's own operations, builds the public -/// response and maps failures to public exceptions. -pub trait RouteHost: Send + Sync { - type Route: Route; +/// Why a route operation the host answered did not produce a result: the route's own code +/// rejected it, which the route classifies like any other native failure, or Python code +/// raised, which reaches the caller as it was raised. +#[derive(Debug)] +pub enum InvokeError { + Native(E), + Python(PyErr), +} - /// `arguments` is the keyword view the callback adapter's `begin` produced, not the +impl From for InvokeError { + fn from(error: PyErr) -> Self { + Self::Python(error) + } +} + +/// The Python side of one route: answers the route's own operations, builds the public +/// response and classifies native failures into public exceptions. +pub trait RouteHost: Send + Sync { + type Route: Route; + + /// The public exception a native failure maps to, kept as a value until the driver + /// raises it. + type Failure: Into; + + /// `arguments` is the keyword view the lifecycle's `begin` produced, not the /// caller's own dict. A route host that projects from it inherits whatever that /// adapter rewrote. fn invoke( @@ -84,7 +117,7 @@ pub trait RouteHost: Send + Sync { py: Python<'_>, arguments: &Bound<'_, PyDict>, op: ::Op, - ) -> PyResult<::OpResult>; + ) -> Result<::OpResult, InvokeError<::Error>>; fn complete( &mut self, @@ -92,12 +125,21 @@ pub trait RouteHost: Send + Sync { response: ::Response, ) -> PyResult>; - fn native_error(error: ::Error) -> PyErr; + /// One streamed chunk as the caller receives it. + fn chunk( + &mut self, + py: Python<'_>, + chunk: ::Chunk, + ) -> PyResult>; + + fn classify( + &self, + py: Python<'_>, + error: ::Error, + ) -> PyResult; fn host_error(error: &PyErr) -> ::Error; - fn map_failure(&self, py: Python<'_>, error: &PyErr) -> PyResult; - fn close(&mut self, py: Python<'_>); fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError>; diff --git a/litellm-rust/crates/host-python/src/argument.rs b/litellm-rust/crates/host-python/src/argument.rs new file mode 100644 index 00000000000..34e07cdfbd5 --- /dev/null +++ b/litellm-rust/crates/host-python/src/argument.rs @@ -0,0 +1,51 @@ +use pyo3::{prelude::*, types::PyDict}; + +/// The caller's own object for a public argument: the keyword if given, even an explicit +/// `None`, else the bound request's attribute. Every reader of a public Python call uses +/// this rule, so the callbacks and the provider see one object per argument. +pub fn lookup<'py>( + kwargs: &Bound<'py, PyDict>, + request: &Bound<'py, PyAny>, + name: &str, +) -> PyResult>> { + if let Some(value) = kwargs.get_item(name)? { + return Ok(Some(value)); + } + request.getattr_opt(name) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn lookup_prefers_the_keyword_even_when_none_and_falls_back_to_the_request() { + crate::initialize_python(); + Python::attach(|py| { + let locals = PyDict::new(py); + py.run( + c" +key = object() +document = {'type': 'document_url'} +class Request: + api_key = 'from-request' + api_base = 'from-request' + document = document +request = Request() +kwargs = {'api_key': key, 'api_base': None} +", + Some(&locals), + Some(&locals), + ) + .unwrap(); + let item = |name: &str| locals.get_item(name).unwrap().unwrap(); + let kwargs = item("kwargs").cast_into::().unwrap(); + let request = item("request"); + let find = |name: &str| lookup(&kwargs, &request, name).unwrap(); + assert!(find("api_key").unwrap().is(item("key"))); + assert!(find("api_base").unwrap().is_none()); + assert!(find("document").unwrap().is(item("document"))); + assert!(find("model").is_none()); + }); + } +} diff --git a/litellm-rust/crates/host-python/src/driver.rs b/litellm-rust/crates/host-python/src/driver.rs index 8bda13b44d0..392d36e10f4 100644 --- a/litellm-rust/crates/host-python/src/driver.rs +++ b/litellm-rust/crates/host-python/src/driver.rs @@ -2,17 +2,19 @@ use std::sync::Arc; use std::task::Poll; use futures_util::future::{AbortHandle, Abortable}; -use litellm_callbacks::event::{CallEvent, FailureOrigin, Timing, epoch_seconds}; -use litellm_callbacks::host::{HostOp, HostResult, HostStep}; -use litellm_callbacks::machine::{HostFailure, Machine, MachineStep}; -use litellm_callbacks::route::Route; +use litellm_host::event::{FailureOrigin, Timing, epoch_seconds}; +use litellm_host::host::{Demand, HostOp, HostResult, HostStep}; +use litellm_host::machine::{HostFailure, Machine, MachineStep}; +use litellm_host::route::Route; use pyo3::exceptions::{PyBaseException, PyException, PyRuntimeError}; use pyo3::gc::{PyTraverseError, PyVisit}; use pyo3::prelude::*; use pyo3::types::PyDict; use tokio::sync::Mutex; -use crate::adapter::{AdapterStep, CallbackAdapter, PublicValue, RouteHost, missing_state}; +use crate::adapter::{ + InvokeError, LifecycleEvent, LifecycleStep, PythonLifecycle, RouteHost, missing_state, +}; use crate::execution::{poll_async_value, run_async_value, run_sync_value}; use crate::handle::{Execution, ExecutionBody, ExecutionStep}; @@ -36,6 +38,7 @@ struct MachineState { enum Stage { Begin, Call, + Streaming, AfterSuccess, Succeeded(Py), Failed(Py), @@ -43,6 +46,7 @@ enum Stage { #[derive(Clone, Copy)] enum Expect { + Started, Arguments, Wire, Emitted, @@ -53,6 +57,8 @@ enum Expect { enum Pending { Native, Adapter(Expect), + /// The stream handed to the caller waits for its next read or its close. + Consumer, } enum Next { @@ -66,7 +72,7 @@ where M: Machine> + 'static, { route: H, - adapter: Box, + adapter: Box, machine: Option>>>, arguments: Option>, started_at: f64, @@ -84,7 +90,7 @@ pub fn run_call( py: Python<'_>, machine: M, route: H, - adapter: Box, + adapter: Box, arguments: Py, asynchronous: bool, ) -> PyResult> @@ -118,7 +124,14 @@ where } match driver.resume(None)? { ExecutionStep::Return(value) => Ok(value), - ExecutionStep::Await(_) => Err(PyRuntimeError::new_err("sync call suspended")), + ExecutionStep::Open => py + .import("litellm.rust_bridge.lifecycle")? + .getattr("SyncStream")? + .call1((Py::new(py, Execution::suspended(driver))?,)) + .map(Bound::unbind), + ExecutionStep::Await(_) | ExecutionStep::Yield(_) => { + Err(PyRuntimeError::new_err("sync call suspended")) + } } } @@ -146,9 +159,11 @@ where match (self.pending.take(), result) { (None, None) => { self.started_at = epoch_seconds(); - let arguments = self.arguments.take().ok_or_else(missing_state)?; - match self.adapter.begin(py, arguments, self.started_at) { - Ok(step) => self.on_adapter(py, step, Expect::Arguments), + let started = LifecycleEvent::Started { + start_time: self.started_at, + }; + match self.adapter.emit(py, started) { + Ok(step) => self.on_adapter(py, step, Expect::Started), Err(error) => self.adapter_failed(py, error), } } @@ -157,6 +172,14 @@ where self.run_steps(py, HostStep::Ready(result)) } (Some(Pending::Native), Some(Err(error))) => self.interrupt(py, error), + (Some(Pending::Consumer), Some(read)) => { + let demand = if read.is_ok() { + Demand::More + } else { + Demand::Detached + }; + self.resume_machine(py, Some(Ok(HostResult::Demand(demand)))) + } (Some(Pending::Adapter(expect)), Some(result)) => { match self.adapter.resume(py, result) { Ok(step) => self.on_adapter(py, step, expect), @@ -170,27 +193,28 @@ where fn on_adapter( &mut self, py: Python<'_>, - step: AdapterStep, + step: LifecycleStep, expect: Expect, ) -> PyResult { match (expect, step) { - (_, AdapterStep::Await(awaitable)) => { + (_, LifecycleStep::Await(awaitable)) => { self.pending = Some(Pending::Adapter(expect)); Ok(ExecutionStep::Await(awaitable)) } - (Expect::Arguments, AdapterStep::Arguments(arguments)) => { + (Expect::Started, LifecycleStep::Done) => self.begin(py), + (Expect::Arguments, LifecycleStep::Arguments(arguments)) => { self.arguments = Some(arguments); self.stage = Stage::Call; self.resume_machine(py, None) } - (Expect::Wire, AdapterStep::Wire(wire)) => { + (Expect::Wire, LifecycleStep::Wire(wire)) => { self.resume_machine(py, Some(Ok(HostResult::BeforeSend(wire)))) } - (Expect::Emitted, AdapterStep::Done) => { + (Expect::Emitted, LifecycleStep::Done) => { self.resume_machine(py, Some(Ok(HostResult::Emitted))) } - (Expect::Response, AdapterStep::Response(response)) => self.succeeded(py, response), - (Expect::Terminal, AdapterStep::Done) => match &self.stage { + (Expect::Response, LifecycleStep::Response(response)) => self.succeeded(py, response), + (Expect::Terminal, LifecycleStep::Done) => match &self.stage { Stage::Succeeded(response) => Ok(ExecutionStep::Return(response.clone_ref(py))), Stage::Failed(error) => Err(PyErr::from_value(error.bind(py).clone().into_any())), _ => Err(missing_state()), @@ -199,10 +223,18 @@ where } } + fn begin(&mut self, py: Python<'_>) -> PyResult { + let arguments = self.arguments.take().ok_or_else(missing_state)?; + match self.adapter.begin(py, arguments, self.started_at) { + Ok(step) => self.on_adapter(py, step, Expect::Arguments), + Err(error) => self.adapter_failed(py, error), + } + } + fn adapter_failed(&mut self, py: Python<'_>, error: PyErr) -> PyResult { match self.stage { Stage::Begin | Stage::AfterSuccess => self.failure(py, error, FailureOrigin::Host), - Stage::Call => self.interrupt(py, error), + Stage::Call | Stage::Streaming => self.interrupt(py, error), Stage::Succeeded(_) | Stage::Failed(_) => Err(error), } } @@ -248,14 +280,20 @@ where let answer = match op { HostOp::Route(op) => { let arguments = self.arguments.as_ref().ok_or_else(missing_state)?; - self.route - .invoke(py, arguments.bind(py), op) - .map(HostResult::Route) + match self.route.invoke(py, arguments.bind(py), op) { + Ok(result) => Ok(HostResult::Route(result)), + Err(InvokeError::Native(error)) => { + return self + .resume_core(py, Some(Err(HostFailure::Error(error)))) + .map(Next::Continue); + } + Err(InvokeError::Python(error)) => Err(error), + } } HostOp::BeforeSend { wire, context } => { match self.adapter.before_send(py, wire, &context) { - Ok(AdapterStep::Wire(wire)) => Ok(HostResult::BeforeSend(wire)), - Ok(AdapterStep::Await(awaitable)) => { + Ok(LifecycleStep::Wire(wire)) => Ok(HostResult::BeforeSend(wire)), + Ok(LifecycleStep::Await(awaitable)) => { self.pending = Some(Pending::Adapter(Expect::Wire)); return Ok(Next::Return(ExecutionStep::Await(awaitable))); } @@ -263,9 +301,11 @@ where Err(error) => Err(error), } } - HostOp::Emit(event) => match self.adapter.emit(py, &event, None) { - Ok(AdapterStep::Done) => Ok(HostResult::Emitted), - Ok(AdapterStep::Await(awaitable)) => { + HostOp::Open(_) => return self.opened(py).map(Next::Return), + HostOp::Deliver(chunk) => return self.delivered(py, chunk).map(Next::Return), + HostOp::Emit(event) => match self.adapter.emit(py, LifecycleEvent::Machine(&event)) { + Ok(LifecycleStep::Done) => Ok(HostResult::Emitted), + Ok(LifecycleStep::Await(awaitable)) => { self.pending = Some(Pending::Adapter(Expect::Emitted)); return Ok(Next::Return(ExecutionStep::Await(awaitable))); } @@ -279,6 +319,35 @@ where } } + fn opened(&mut self, py: Python<'_>) -> PyResult { + self.stage = Stage::Streaming; + match self.adapter.opened(py) { + Ok(()) => { + self.pending = Some(Pending::Consumer); + Ok(ExecutionStep::Open) + } + Err(error) => self.interrupt(py, error), + } + } + + fn delivered( + &mut self, + py: Python<'_>, + chunk: as Route>::Chunk, + ) -> PyResult { + let chunk = match self.route.chunk(py, chunk) { + Ok(chunk) => chunk, + Err(error) => return self.interrupt(py, error), + }; + match self.adapter.delivered(py, &chunk) { + Ok(()) => { + self.pending = Some(Pending::Consumer); + Ok(ExecutionStep::Yield(chunk)) + } + Err(error) => self.interrupt(py, error), + } + } + fn interrupt(&mut self, py: Python<'_>, error: PyErr) -> PyResult { let cancelled = is_cancellation(py, &error); let native = H::host_error(&error); @@ -349,6 +418,9 @@ where Ok(public) => public, Err(error) => return self.failure(py, error, FailureOrigin::Call), }; + if let Stage::Streaming = self.stage { + return self.succeeded(py, public); + } self.stage = Stage::AfterSuccess; match self.adapter.after_success(py, public, self.timing()) { Ok(step) => self.on_adapter(py, step, Expect::Response), @@ -360,18 +432,35 @@ where self.ended_at.get_or_insert_with(epoch_seconds); let error = match self.interrupted.take() { Some(retained) => PyErr::from_value(retained.into_bound(py).into_any()), - None => H::native_error(error), + None => self.classified(py, error), }; self.failure(py, error, FailureOrigin::Call) } - fn succeeded(&mut self, py: Python<'_>, response: Py) -> PyResult { - let event = CallEvent::Succeeded { - timing: self.timing(), + /// The route's public exception for a native failure. When classification itself + /// fails, that failure is raised with the native error's text as its `__context__`. + fn classified(&self, py: Python<'_>, error: ErrorOf) -> PyErr { + let native = error.to_string(); + let classifier_error = match self.route.classify(py, error) { + Ok(failure) => return failure.into(), + Err(classifier_error) => classifier_error, }; - let step = self - .adapter - .emit(py, &event, Some(PublicValue::Response(&response)))?; + let attached = classifier_error.value(py).setattr( + "__context__", + PyRuntimeError::new_err(native).into_value(py), + ); + match attached { + Ok(()) => classifier_error, + Err(error) => error, + } + } + + fn succeeded(&mut self, py: Python<'_>, response: Py) -> PyResult { + let event = LifecycleEvent::Succeeded { + timing: self.timing(), + response: &response, + }; + let step = self.adapter.emit(py, event)?; self.stage = Stage::Succeeded(response); self.on_adapter(py, step, Expect::Terminal) } @@ -386,18 +475,13 @@ where if is_cancellation(py, &error) { return Err(error); } - let public = match origin { - FailureOrigin::Call => self.route.map_failure(py, &error).unwrap_or(error), - FailureOrigin::Host => error, - }; - let event = CallEvent::Failed { + let event = LifecycleEvent::Failed { timing: self.timing(), origin, + error: &error, }; - let step = self - .adapter - .emit(py, &event, Some(PublicValue::Error(&public)))?; - self.stage = Stage::Failed(public.into_value(py)); + let step = self.adapter.emit(py, event)?; + self.stage = Stage::Failed(error.into_value(py)); self.on_adapter(py, step, Expect::Terminal) } @@ -450,8 +534,8 @@ where mod tests { use std::sync::{Arc, Mutex}; - use litellm_callbacks::event::{RequestContext, WireRequest}; - use litellm_callbacks::machine::{Interrupted, Step}; + use litellm_host::event::{MachineEvent, RequestContext, WireRequest}; + use litellm_host::machine::{Interrupted, Step}; use pyo3::exceptions::{PyBaseException, PyValueError}; use pyo3::types::PyDict; @@ -489,6 +573,12 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri #[derive(Clone, Debug, PartialEq, Eq)] struct Error(String); + impl std::fmt::Display for Error { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter.write_str(&self.0) + } + } + struct Synthetic; impl Route for Synthetic { @@ -496,6 +586,8 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri type Error = Error; type Op = &'static str; type OpResult = String; + type Chunk = std::convert::Infallible; + type StreamHead = std::convert::Infallible; } /// Yields the scripted ops in order, then completes or fails as scripted. @@ -518,8 +610,8 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri model: "model".into(), custom_llm_provider: "provider".into(), optional_params: serde_json::json!({}), - passthrough_fields: Default::default(), secret_fields: Vec::new(), + api_key: None, } } @@ -534,6 +626,7 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri HostResult::Route(value) => value, HostResult::BeforeSend(wire) => wire.url, HostResult::Emitted => "emitted".into(), + HostResult::Demand(demand) => format!("{demand:?}"), }); } if !self.ops.is_empty() { @@ -566,25 +659,50 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri } } + #[derive(Clone, Copy)] + enum OpScript { + Answer, + RaisePython, + RejectNatively, + } + struct SyntheticHost { log: Log, - fail_op: bool, + op: OpScript, + classifier_fails: bool, + } + + /// The fake route's public exception, kept as a value so a test sees what `classify` + /// produced before the driver raises it. + #[derive(Debug, PartialEq, Eq)] + struct Classified(String); + + impl From for PyErr { + fn from(classified: Classified) -> Self { + PyValueError::new_err(format!("classified: {}", classified.0)) + } } impl RouteHost for SyntheticHost { type Route = Synthetic; + type Failure = Classified; fn invoke( &mut self, _: Python<'_>, arguments: &Bound<'_, PyDict>, op: &'static str, - ) -> PyResult { + ) -> Result> { self.log.push(format!("route:{op}")); - if self.fail_op { - return Err(PyValueError::new_err("op failed")); + match self.op { + OpScript::Answer => Ok(format!("{op}:{}", arguments.len())), + OpScript::RaisePython => Err(PyValueError::new_err("op failed").into()), + OpScript::RejectNatively => Err(InvokeError::Native(Error("op rejected".into()))), } - Ok(format!("{op}:{}", arguments.len())) + } + + fn chunk(&mut self, _: Python<'_>, chunk: std::convert::Infallible) -> PyResult> { + match chunk {} } fn complete(&mut self, py: Python<'_>, response: String) -> PyResult> { @@ -594,22 +712,18 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri .unbind()) } - fn native_error(error: Error) -> PyErr { - PyValueError::new_err(error.0) + fn classify(&self, _: Python<'_>, error: Error) -> PyResult { + self.log.push(format!("classify:{error}")); + if self.classifier_fails { + return Err(pyo3::exceptions::PyTypeError::new_err("classifier failed")); + } + Ok(Classified(error.0)) } fn host_error(error: &PyErr) -> Error { Error(error.to_string()) } - fn map_failure(&self, py: Python<'_>, error: &PyErr) -> PyResult { - self.log.push("map_failure"); - Ok(PyValueError::new_err(format!( - "mapped: {}", - error.value(py) - ))) - } - fn close(&mut self, _: Python<'_>) { self.log.push("route.close"); } @@ -632,13 +746,18 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri script: AdapterScript, } - impl CallbackAdapter for SyntheticAdapter { - fn begin(&mut self, _: Python<'_>, arguments: Py, _: f64) -> PyResult { + impl PythonLifecycle for SyntheticAdapter { + fn begin( + &mut self, + _: Python<'_>, + arguments: Py, + _: f64, + ) -> PyResult { self.log.push("begin"); if matches!(self.script, AdapterScript::FailBegin) { return Err(PyValueError::new_err("begin failed")); } - Ok(AdapterStep::Arguments(arguments)) + Ok(LifecycleStep::Arguments(arguments)) } fn before_send( @@ -646,9 +765,9 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri _: Python<'_>, wire: Box, _: &RequestContext, - ) -> PyResult { + ) -> PyResult { self.log.push("before_send"); - Ok(AdapterStep::Wire(Box::new(WireRequest { + Ok(LifecycleStep::Wire(Box::new(WireRequest { url: "rewritten".into(), ..*wire }))) @@ -659,41 +778,48 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri py: Python<'_>, response: Py, _: Timing, - ) -> PyResult { + ) -> PyResult { self.log.push("after_success"); match self.script { - AdapterScript::ReplaceResponse => Ok(AdapterStep::Response( + AdapterScript::ReplaceResponse => Ok(LifecycleStep::Response( "replaced".into_pyobject(py)?.into_any().unbind(), )), AdapterScript::FailAfterSuccess => { Err(PyValueError::new_err("after_success failed")) } AdapterScript::Plain | AdapterScript::FailBegin => { - Ok(AdapterStep::Response(response)) + Ok(LifecycleStep::Response(response)) } } } - fn emit( - &mut self, - py: Python<'_>, - event: &CallEvent, - public: Option>, - ) -> PyResult { - self.log.push(match (event, public) { - (CallEvent::ResponseReceived { raw }, None) => format!("response:{}", raw.body), - (CallEvent::Succeeded { .. }, Some(PublicValue::Response(value))) => { - format!("succeeded:{}", value.bind(py)) + fn emit(&mut self, py: Python<'_>, event: LifecycleEvent<'_>) -> PyResult { + self.log.push(match event { + LifecycleEvent::Started { .. } => "started".into(), + LifecycleEvent::Machine(MachineEvent::ResponseReceived { raw }) => { + format!("response:{}", raw.body) } - (CallEvent::Failed { origin, .. }, Some(PublicValue::Error(error))) => { + LifecycleEvent::Succeeded { response, .. } => { + format!("succeeded:{}", response.bind(py)) + } + LifecycleEvent::Failed { origin, error, .. } => { format!("failed:{origin:?}:{}", error.value(py)) } - _ => "unexpected".into(), }); - Ok(AdapterStep::Done) + Ok(LifecycleStep::Done) } - fn resume(&mut self, _: Python<'_>, _: PyResult>) -> PyResult { + fn opened(&mut self, _: Python<'_>) -> PyResult<()> { + self.log.push("opened"); + Ok(()) + } + + fn delivered(&mut self, _: Python<'_>, _: &Py) -> PyResult<()> { + self.log.push("delivered"); + Ok(()) + } + + fn resume(&mut self, _: Python<'_>, _: PyResult>) -> PyResult { Err(missing_state()) } @@ -709,15 +835,31 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri fn run_scripted( py: Python<'_>, machine: ScriptedMachine, - fail_op: bool, + op: OpScript, script: AdapterScript, asynchronous: bool, ) -> (PyResult>, Vec) { - let log = Log::default(); - let route = SyntheticHost { - log: Log(log.0.clone()), - fail_op, - }; + run_hosted( + py, + machine, + SyntheticHost { + log: Log::default(), + op, + classifier_fails: false, + }, + script, + asynchronous, + ) + } + + fn run_hosted( + py: Python<'_>, + machine: ScriptedMachine, + route: SyntheticHost, + script: AdapterScript, + asynchronous: bool, + ) -> (PyResult>, Vec) { + let log = Log(route.log.0.clone()); let adapter = SyntheticAdapter { log: Log(log.0.clone()), script, @@ -756,8 +898,8 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri wire: Box::new(wire()), context: Box::new(context()), }, - HostOp::Emit(CallEvent::ResponseReceived { - raw: litellm_callbacks::event::RawResponse { body: "raw".into() }, + HostOp::Emit(MachineEvent::ResponseReceived { + raw: litellm_host::event::RawResponse { body: "raw".into() }, }), ], outcome: Some(Ok("done".into())), @@ -777,7 +919,7 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri let (result, log) = run_scripted( py, success_machine(), - false, + OpScript::Answer, AdapterScript::Plain, asynchronous, ); @@ -785,6 +927,7 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri assert_eq!( log, [ + "started", "begin", "route:project", "before_send", @@ -800,28 +943,75 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri }); } + fn failing_machine() -> ScriptedMachine { + ScriptedMachine { + ops: vec![HostOp::Route("project")], + outcome: Some(Err(Error("provider exploded".into()))), + answers: Vec::new(), + } + } + #[test] - fn machine_failures_are_mapped_and_dispatched_once_as_call_failures() { + fn a_native_failure_is_classified_once_and_reported_classified() { let _guard = PYTHON_GLOBALS .lock() .unwrap_or_else(|error| error.into_inner()); crate::initialize_python(); Python::attach(|py| { - let machine = ScriptedMachine { - ops: vec![HostOp::Route("project")], - outcome: Some(Err(Error("provider exploded".into()))), - answers: Vec::new(), - }; - let (result, log) = run_scripted(py, machine, false, AdapterScript::Plain, false); - let error = result.unwrap_err(); - assert_eq!(error.value(py).to_string(), "mapped: provider exploded"); + install_lifecycle_module(py); + for asynchronous in [false, true] { + let (result, log) = run_scripted( + py, + failing_machine(), + OpScript::Answer, + AdapterScript::Plain, + asynchronous, + ); + let error = result.unwrap_err(); + assert!(error.is_instance_of::(py)); + assert_eq!(error.value(py).to_string(), "classified: provider exploded"); + assert_eq!( + log, + [ + "started", + "begin", + "route:project", + "classify:provider exploded", + "failed:Call:classified: provider exploded", + "adapter.close", + "route.close", + ] + ); + } + }); + } + + #[test] + fn a_native_rejection_from_a_host_operation_is_classified_once() { + let _guard = PYTHON_GLOBALS + .lock() + .unwrap_or_else(|error| error.into_inner()); + crate::initialize_python(); + Python::attach(|py| { + let (result, log) = run_scripted( + py, + success_machine(), + OpScript::RejectNatively, + AdapterScript::Plain, + false, + ); + assert_eq!( + result.unwrap_err().value(py).to_string(), + "classified: op rejected" + ); assert_eq!( log, [ + "started", "begin", "route:project", - "map_failure", - "failed:Call:mapped: provider exploded", + "classify:op rejected", + "failed:Call:classified: op rejected", "adapter.close", "route.close", ] @@ -830,18 +1020,72 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri } #[test] - fn host_operation_failures_interrupt_the_call_and_keep_the_python_exception() { + fn a_python_exception_from_a_host_operation_is_reported_as_raised() { let _guard = PYTHON_GLOBALS .lock() .unwrap_or_else(|error| error.into_inner()); crate::initialize_python(); Python::attach(|py| { - let (result, log) = - run_scripted(py, success_machine(), true, AdapterScript::Plain, false); + let (result, log) = run_scripted( + py, + success_machine(), + OpScript::RaisePython, + AdapterScript::Plain, + false, + ); let error = result.unwrap_err(); - assert_eq!(error.value(py).to_string(), "mapped: op failed"); - assert!(!log.contains(&"before_send".to_string())); - assert!(log.contains(&"failed:Call:mapped: op failed".to_string())); + assert!(error.is_instance_of::(py)); + assert_eq!(error.value(py).to_string(), "op failed"); + assert_eq!( + log, + [ + "started", + "begin", + "route:project", + "failed:Call:op failed", + "adapter.close", + "route.close", + ] + ); + }); + } + + #[test] + fn a_failing_classifier_surfaces_with_the_native_error_as_context() { + let _guard = PYTHON_GLOBALS + .lock() + .unwrap_or_else(|error| error.into_inner()); + crate::initialize_python(); + Python::attach(|py| { + let (result, log) = run_hosted( + py, + failing_machine(), + SyntheticHost { + log: Log::default(), + op: OpScript::Answer, + classifier_fails: true, + }, + AdapterScript::Plain, + false, + ); + let error = result.unwrap_err(); + assert!(error.is_instance_of::(py)); + assert_eq!(error.value(py).to_string(), "classifier failed"); + let context = error.value(py).getattr("__context__").unwrap(); + assert!(context.is_instance_of::()); + assert_eq!(context.str().unwrap().to_string(), "provider exploded"); + assert_eq!( + log, + [ + "started", + "begin", + "route:project", + "classify:provider exploded", + "failed:Call:classifier failed", + "adapter.close", + "route.close", + ] + ); }); } @@ -855,7 +1099,7 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri let (result, log) = run_scripted( py, success_machine(), - false, + OpScript::Answer, AdapterScript::FailBegin, false, ); @@ -864,6 +1108,7 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri assert_eq!( log, [ + "started", "begin", "failed:Host:begin failed", "adapter.close", @@ -885,7 +1130,7 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri let (result, log) = run_scripted( py, success_machine(), - false, + OpScript::Answer, AdapterScript::ReplaceResponse, asynchronous, ); @@ -908,7 +1153,7 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri let (result, log) = run_scripted( py, success_machine(), - false, + OpScript::Answer, AdapterScript::FailAfterSuccess, asynchronous, ); @@ -938,12 +1183,13 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri struct Cancelling(Log); impl RouteHost for Cancelling { type Route = Synthetic; + type Failure = Classified; fn invoke( &mut self, py: Python<'_>, _: &Bound<'_, PyDict>, _: &'static str, - ) -> PyResult { + ) -> Result> { self.0.push("route"); Err(PyErr::from_value( py.import("asyncio") @@ -952,21 +1198,26 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri .unwrap() .call0() .unwrap(), - )) + ) + .into()) + } + fn chunk( + &mut self, + _: Python<'_>, + chunk: std::convert::Infallible, + ) -> PyResult> { + match chunk {} } fn complete(&mut self, _: Python<'_>, _: String) -> PyResult> { Err(missing_state()) } - fn native_error(error: Error) -> PyErr { - PyValueError::new_err(error.0) + fn classify(&self, _: Python<'_>, error: Error) -> PyResult { + self.0.push("classify"); + Ok(Classified(error.0)) } fn host_error(error: &PyErr) -> Error { Error(error.to_string()) } - fn map_failure(&self, _: Python<'_>, _: &PyErr) -> PyResult { - self.0.push("map_failure"); - Err(missing_state()) - } fn close(&mut self, _: Python<'_>) {} fn traverse(&self, _: &PyVisit<'_>) -> Result<(), PyTraverseError> { Ok(()) @@ -988,7 +1239,10 @@ sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bri ) .unwrap_err(); assert!(!error.is_instance_of::(py)); - assert_eq!(log.entries(), ["begin", "route", "adapter.close"]); + assert_eq!( + log.entries(), + ["started", "begin", "route", "adapter.close"] + ); }); } diff --git a/litellm-rust/crates/host-python/src/handle.rs b/litellm-rust/crates/host-python/src/handle.rs index d8cd6c92130..10abbadbda5 100644 --- a/litellm-rust/crates/host-python/src/handle.rs +++ b/litellm-rust/crates/host-python/src/handle.rs @@ -8,6 +8,10 @@ use pyo3::prelude::*; pub enum ExecutionStep { Return(Py), Await(Py), + /// The call streams: the caller gets a stream over this execution, which stays + /// suspended until the stream asks for a chunk. + Open, + Yield(Py), } pub trait ExecutionBody: Send + Sync { @@ -34,6 +38,13 @@ impl Execution { } } + /// An execution already started elsewhere and now waiting for its next input. + pub fn suspended(body: impl ExecutionBody + 'static) -> Self { + Self { + state: ExecutionState::Suspended(Box::new(body)), + } + } + fn advance( slf: &Bound<'_, Self>, py: Python<'_>, @@ -64,6 +75,8 @@ impl Execution { let step = body.resume(result)?; let (tag, value, suspended) = match step { ExecutionStep::Await(value) => ("Await", value, true), + ExecutionStep::Open => ("Open", py.None(), true), + ExecutionStep::Yield(value) => ("Yield", value, true), ExecutionStep::Return(value) => ("Complete", value, false), }; let step = py diff --git a/litellm-rust/crates/host-python/src/lib.rs b/litellm-rust/crates/host-python/src/lib.rs index bb0b5b1c3b1..583a4eb91b6 100644 --- a/litellm-rust/crates/host-python/src/lib.rs +++ b/litellm-rust/crates/host-python/src/lib.rs @@ -1,9 +1,10 @@ //! The CPython runtime adapter: value marshalling, interpreter detachment, the tokio and -//! asyncio glue, and the driver that runs a native [`Machine`](litellm_callbacks::machine::Machine) -//! against a Python route host and a callback adapter. Everything here is Python-specific by +//! asyncio glue, and the driver that runs a native [`Machine`](litellm_host::machine::Machine) +//! against a Python route host and a Python lifecycle. Everything here is Python-specific by //! construction; another host language gets its own crate of the same shape. mod adapter; +mod argument; mod callable; mod driver; mod execution; @@ -11,7 +12,10 @@ mod gil; mod handle; mod marshal; -pub use adapter::{AdapterStep, CallbackAdapter, PublicValue, RouteHost, missing_state}; +pub use adapter::{ + InvokeError, LifecycleEvent, LifecycleStep, PythonLifecycle, RouteHost, missing_state, +}; +pub use argument::lookup; pub use callable::wrap_failure; pub use driver::run_call; pub use execution::{poll_async_value, run_async, run_async_value, run_sync, run_sync_value}; diff --git a/litellm-rust/crates/callbacks/Cargo.toml b/litellm-rust/crates/host/Cargo.toml similarity index 65% rename from litellm-rust/crates/callbacks/Cargo.toml rename to litellm-rust/crates/host/Cargo.toml index 4b966271478..0c7c46192b5 100644 --- a/litellm-rust/crates/callbacks/Cargo.toml +++ b/litellm-rust/crates/host/Cargo.toml @@ -1,13 +1,14 @@ [package] -name = "litellm-callbacks" +name = "litellm-host" version = "0.1.0" edition.workspace = true license.workspace = true repository.workspace = true [dependencies] +litellm-auth.workspace = true serde_json.workspace = true +tokio = { workspace = true, features = ["sync"] } [dev-dependencies] rstest.workspace = true -tokio = { workspace = true, features = ["macros"] } diff --git a/litellm-rust/crates/host/src/event.rs b/litellm-rust/crates/host/src/event.rs new file mode 100644 index 00000000000..182dab657d3 --- /dev/null +++ b/litellm-rust/crates/host/src/event.rs @@ -0,0 +1,76 @@ +use std::time::{SystemTime, UNIX_EPOCH}; + +use serde_json::Value; + +/// Seconds since the Unix epoch, on one clock for every host. +pub fn epoch_seconds() -> f64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .map(|duration| duration.as_secs_f64()) + .unwrap_or(0.0) +} + +#[derive(Clone, Copy, Debug, PartialEq)] +pub struct Timing { + pub start_time: f64, + pub end_time: f64, +} + +/// The provider request as it is about to leave, offered to the host for rewriting. +#[derive(Clone, Debug, PartialEq)] +pub struct WireRequest { + pub url: String, + pub headers: Vec<(String, String)>, + pub body: Value, +} + +/// What the route knows about the request it is sending, for a host that logs it. The +/// route owns these facts; a host reads them beside the wire request and never rewrites +/// them. +#[derive(Clone, Debug, PartialEq)] +pub struct RequestContext { + pub model: String, + pub custom_llm_provider: String, + /// The route's parameters before the provider transformation. + pub optional_params: Value, + /// Optional-param names that carry credentials and must be redacted when logged. + pub secret_fields: Vec, + /// The credential the route resolved for the provider call. + pub api_key: Option, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct RawResponse { + pub body: String, +} + +/// Whether a failure surfaced inside the call, including a host op the call asked for, +/// or in a host step around it (preparing the arguments, finalizing the response). +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum FailureOrigin { + Call, + Host, +} + +/// What a machine reports while it runs. +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum MachineEvent { + ResponseReceived { raw: RawResponse }, +} + +/// What an in-process host observes: the machine's own events between the driver's +/// start and terminal ones. +#[derive(Clone, Debug, PartialEq)] +pub enum CallEvent { + Started { + start_time: f64, + }, + Machine(MachineEvent), + Succeeded { + timing: Timing, + }, + Failed { + timing: Timing, + origin: FailureOrigin, + }, +} diff --git a/litellm-rust/crates/callbacks/src/host.rs b/litellm-rust/crates/host/src/host.rs similarity index 58% rename from litellm-rust/crates/callbacks/src/host.rs rename to litellm-rust/crates/host/src/host.rs index 2392718a18d..aba35185a18 100644 --- a/litellm-rust/crates/callbacks/src/host.rs +++ b/litellm-rust/crates/host/src/host.rs @@ -1,6 +1,6 @@ use std::future::Future; -use crate::event::{CallEvent, RequestContext, WireRequest}; +use crate::event::{CallEvent, MachineEvent, RequestContext, WireRequest}; use crate::route::Route; /// One suspension point of a native call, performed by the host. @@ -10,13 +10,26 @@ pub enum HostOp { wire: Box, context: Box, }, - Emit(CallEvent), + Emit(MachineEvent), + /// The response streams: the host hands the caller a stream and answers once the + /// caller asks for the first chunk or goes away. + Open(R::StreamHead), + /// The next chunk of an open stream, answered once the caller asks for the one after. + Deliver(R::Chunk), } pub enum HostResult { Route(R::OpResult), BeforeSend(Box), Emitted, + Demand(Demand), +} + +/// Whether the caller of a streamed call still reads it. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum Demand { + More, + Detached, } /// A host answer that is either available now or arrives once the host's own @@ -42,4 +55,12 @@ pub trait Host: Send + Sync { fn emit(&self, _event: &CallEvent) -> impl Future> + Send { async { Ok(()) } } + + fn open(&self, _head: R::StreamHead) -> impl Future> + Send { + async { Ok(Demand::More) } + } + + fn deliver(&self, _chunk: R::Chunk) -> impl Future> + Send { + async { Ok(Demand::More) } + } } diff --git a/litellm-rust/crates/callbacks/src/lib.rs b/litellm-rust/crates/host/src/lib.rs similarity index 81% rename from litellm-rust/crates/callbacks/src/lib.rs rename to litellm-rust/crates/host/src/lib.rs index 41b0983f0ce..65479c2380f 100644 --- a/litellm-rust/crates/callbacks/src/lib.rs +++ b/litellm-rust/crates/host/src/lib.rs @@ -1,7 +1,7 @@ //! The contract between a native call and the host runtime that drives it. //! //! A host is whatever sits on the far side of the language boundary: CPython today, -//! another runtime later. Core implements [`machine::Machine`] per route and never learns +//! another runtime later. Core runs each route on a [`machine::RouteMachine`] and never learns //! which host is on the other end. The machine yields [`host::HostOp`]s; a driver answers //! them, observes [`event::CallEvent`]s and may rewrite the wire request before it is sent. diff --git a/litellm-rust/crates/core/src/machine/auth.rs b/litellm-rust/crates/host/src/machine/auth.rs similarity index 97% rename from litellm-rust/crates/core/src/machine/auth.rs rename to litellm-rust/crates/host/src/machine/auth.rs index 6a3e4daf6ee..ba7e242e766 100644 --- a/litellm-rust/crates/core/src/machine/auth.rs +++ b/litellm-rust/crates/host/src/machine/auth.rs @@ -1,9 +1,8 @@ use std::sync::Arc; -use litellm_auth::{Error, ResolvedCredential, TokenFuture, TokenProvider, TokenProviderHandle}; -use litellm_callbacks::route::Route; - use super::{HostChannel, MachineFault}; +use crate::route::Route; +use litellm_auth::{Error, ResolvedCredential, TokenFuture, TokenProvider, TokenProviderHandle}; /// A route whose host can mint credentials on the call's behalf. pub trait TokenRoute: Route { diff --git a/litellm-rust/crates/callbacks/src/machine.rs b/litellm-rust/crates/host/src/machine/mod.rs similarity index 91% rename from litellm-rust/crates/callbacks/src/machine.rs rename to litellm-rust/crates/host/src/machine/mod.rs index 2942913f095..2c26db61582 100644 --- a/litellm-rust/crates/callbacks/src/machine.rs +++ b/litellm-rust/crates/host/src/machine/mod.rs @@ -1,6 +1,12 @@ +mod auth; +mod route_machine; + use std::future::Future; use std::pin::Pin; +pub use auth::{HostTokenProvider, TokenRoute}; +pub use route_machine::{ExecuteFuture, HostChannel, MachineFault, RouteMachine}; + use crate::host::{HostOp, HostResult}; use crate::route::Route; diff --git a/litellm-rust/crates/core/src/machine/mod.rs b/litellm-rust/crates/host/src/machine/route_machine.rs similarity index 88% rename from litellm-rust/crates/core/src/machine/mod.rs rename to litellm-rust/crates/host/src/machine/route_machine.rs index 279a2d65c97..38a0b8bc16a 100644 --- a/litellm-rust/crates/core/src/machine/mod.rs +++ b/litellm-rust/crates/host/src/machine/route_machine.rs @@ -2,18 +2,16 @@ //! place, and turns the host operations that future requests into [`Machine`] steps. No //! task is spawned; dropping the machine drops the in-flight call. -mod auth; - use std::{future::Future, pin::Pin}; -pub use auth::{HostTokenProvider, TokenRoute}; -use litellm_callbacks::{ - event::{CallEvent, RequestContext, WireRequest}, - host::{HostOp, HostResult}, - machine::{HostFailure, Interrupted, Machine, MachineStep, Step}, +use tokio::sync::{mpsc, oneshot}; + +use super::{HostFailure, Interrupted, Machine, MachineStep, Step}; +use crate::{ + event::{MachineEvent, RequestContext, WireRequest}, + host::{Demand, HostOp, HostResult}, route::Route, }; -use tokio::sync::{mpsc, oneshot}; /// The machine's own failures, distinct from anything the provider call reports. #[derive(Clone, Copy, Debug, PartialEq, Eq)] @@ -82,12 +80,27 @@ where } } - pub async fn emit(&self, event: CallEvent) -> Result<(), R::Error> { + pub async fn emit(&self, event: MachineEvent) -> Result<(), R::Error> { match self.invoke(HostOp::Emit(event)).await? { HostResult::Emitted => Ok(()), _ => Err(MachineFault::Mismatch.into()), } } + + pub async fn open(&self, head: R::StreamHead) -> Result { + self.demand(HostOp::Open(head)).await + } + + pub async fn deliver(&self, chunk: R::Chunk) -> Result { + self.demand(HostOp::Deliver(chunk)).await + } + + async fn demand(&self, op: HostOp) -> Result { + match self.invoke(op).await? { + HostResult::Demand(demand) => Ok(demand), + _ => Err(MachineFault::Mismatch.into()), + } + } } enum Execution { diff --git a/litellm-rust/crates/callbacks/src/route.rs b/litellm-rust/crates/host/src/route.rs similarity index 57% rename from litellm-rust/crates/callbacks/src/route.rs rename to litellm-rust/crates/host/src/route.rs index 97738c8da8b..8ab2b125760 100644 --- a/litellm-rust/crates/callbacks/src/route.rs +++ b/litellm-rust/crates/host/src/route.rs @@ -6,4 +6,9 @@ pub trait Route: Send + Sync + 'static { type Error: Clone + Send + Sync + 'static; type Op: Send + 'static; type OpResult: Send + 'static; + /// One piece of a streamed response, handed to the caller as it arrives. A route + /// that never streams uses `Infallible`. + type Chunk: Send + 'static; + /// What the route knows once a streamed response starts, before its first chunk. + type StreamHead: Send + 'static; } diff --git a/litellm-rust/crates/callbacks/src/run.rs b/litellm-rust/crates/host/src/run.rs similarity index 72% rename from litellm-rust/crates/callbacks/src/run.rs rename to litellm-rust/crates/host/src/run.rs index 57bf134f345..6a0c08fba68 100644 --- a/litellm-rust/crates/callbacks/src/run.rs +++ b/litellm-rust/crates/host/src/run.rs @@ -11,6 +11,7 @@ where H: Host, { let start_time = epoch_seconds(); + let _ = host.emit(&CallEvent::Started { start_time }).await; let mut result = None; let outcome = loop { let step = match machine.resume(result.take()).await { @@ -24,7 +25,12 @@ where .before_send(*wire, &context) .await .map(|wire| HostResult::BeforeSend(Box::new(wire))), - HostOp::Emit(event) => host.emit(&event).await.map(|()| HostResult::Emitted), + HostOp::Emit(event) => host + .emit(&CallEvent::Machine(event)) + .await + .map(|()| HostResult::Emitted), + HostOp::Open(head) => host.open(head).await.map(HostResult::Demand), + HostOp::Deliver(chunk) => host.deliver(chunk).await.map(HostResult::Demand), }; match answer { Ok(answer) => result = Some(answer), @@ -60,6 +66,8 @@ mod tests { type Error = &'static str; type Op = &'static str; type OpResult = (); + type Chunk = std::convert::Infallible; + type StreamHead = std::convert::Infallible; } struct Scripted { @@ -102,6 +110,7 @@ mod tests { async fn emit(&self, event: &CallEvent) -> Result<(), &'static str> { self.seen.lock().unwrap().push(match event { + CallEvent::Started { .. } => "started".into(), CallEvent::Succeeded { .. } => "succeeded".into(), CallEvent::Failed { .. } => "failed".into(), other => format!("{other:?}"), @@ -124,7 +133,7 @@ mod tests { assert_eq!(outcome, Ok(())); assert_eq!( *host.seen.lock().unwrap(), - ["route:project", "route:send", "succeeded"] + ["started", "route:project", "route:send", "succeeded"] ); } @@ -133,7 +142,7 @@ mod tests { let host = Recording::default(); let outcome = run(scripted(&[], Err("boom")), &host).await; assert_eq!(outcome, Err("boom")); - assert_eq!(*host.seen.lock().unwrap(), ["failed"]); + assert_eq!(*host.seen.lock().unwrap(), ["started", "failed"]); let host = Recording { fail: Some("send"), @@ -143,7 +152,35 @@ mod tests { assert_eq!(outcome, Err("host failed")); assert_eq!( *host.seen.lock().unwrap(), - ["route:project", "route:send", "failed"] + ["started", "route:project", "route:send", "failed"] ); } + + struct StartTimes(Mutex>); + + impl Host for StartTimes { + async fn route(&self, _: &'static str) -> Result<(), &'static str> { + Ok(()) + } + + async fn emit(&self, event: &CallEvent) -> Result<(), &'static str> { + if let CallEvent::Started { start_time } + | CallEvent::Succeeded { + timing: Timing { start_time, .. }, + } = event + { + self.0.lock().unwrap().push(*start_time); + } + Err("observer failed") + } + } + + #[tokio::test] + async fn started_opens_the_call_at_the_terminal_start_time_and_cannot_fail_it() { + let host = StartTimes(Mutex::default()); + assert_eq!(run(scripted(&["project"], Ok(())), &host).await, Ok(())); + let times = host.0.lock().unwrap(); + assert_eq!(times.len(), 2); + assert_eq!(times[0], times[1]); + } } diff --git a/litellm-rust/crates/llms/Cargo.toml b/litellm-rust/crates/llms/Cargo.toml index 4ca6c7cb2a5..d295e4407ba 100644 --- a/litellm-rust/crates/llms/Cargo.toml +++ b/litellm-rust/crates/llms/Cargo.toml @@ -15,7 +15,7 @@ litellm-auth.workspace = true litellm-auth-aws.workspace = true litellm-auth-azure.workspace = true litellm-auth-gcp.workspace = true -litellm-callbacks.workspace = true +litellm-host.workspace = true litellm-framing.workspace = true base64.workspace = true bytes.workspace = true diff --git a/litellm-rust/crates/llms/src/azure_ai/ocr/document_intelligence/transformation.rs b/litellm-rust/crates/llms/src/azure_ai/ocr/document_intelligence/transformation.rs index 87945bf8785..2e398d0287e 100644 --- a/litellm-rust/crates/llms/src/azure_ai/ocr/document_intelligence/transformation.rs +++ b/litellm-rust/crates/llms/src/azure_ai/ocr/document_intelligence/transformation.rs @@ -150,7 +150,7 @@ impl BaseOcrConfig for AzureDocumentIntelligenceOcrConfig { api_key: inputs.api_key.and_then(|key| { inputs .dynamic_api_key - .filter(|value| !value.value().is_empty()) + .filter(|value| !value.value().expose().is_empty()) .or(Some(key)) }), api_base: inputs.api_base.and_then(|base| { @@ -592,12 +592,17 @@ impl AzureDocumentIntelligenceOcrConfig { )?; return Ok(connection.extra_headers.clone()); } - let key = nonblank(connection.api_key.clone()) - .map(|value| Sourced::new(value, connection.api_key_source)) - .or_else(|| { - nonblank(self.get_api_key_env_var().and_then(env_lookup)) - .map(|value| Sourced::new(value, InputSource::Environment)) - }); + let key = nonblank( + connection + .api_key + .as_ref() + .map(|key| key.expose().to_string()), + ) + .map(|value| Sourced::new(value, connection.api_key_source)) + .or_else(|| { + nonblank(self.get_api_key_env_var().and_then(env_lookup)) + .map(|value| Sourced::new(value, InputSource::Environment)) + }); if let Some(key) = key { super::super::common_utils::validate_destination(connection, key.source())?; return Ok( @@ -796,7 +801,7 @@ mod tests { #[tokio::test] async fn request_endpoint_accepts_request_owned_key() { let connection = OcrConnection { - api_key: Some("request-key".into()), + api_key: Some(litellm_auth::SecretValue::new("request-key")), api_key_source: InputSource::Request, api_base: Some("https://request.example".into()), api_base_source: InputSource::Request, diff --git a/litellm-rust/crates/llms/src/azure_ai/ocr/transformation.rs b/litellm-rust/crates/llms/src/azure_ai/ocr/transformation.rs index 2012f740173..7ef051e8986 100644 --- a/litellm-rust/crates/llms/src/azure_ai/ocr/transformation.rs +++ b/litellm-rust/crates/llms/src/azure_ai/ocr/transformation.rs @@ -142,12 +142,17 @@ impl AzureAiOcrConfig { super::common_utils::validate_destination(connection, connection.extra_headers_source)?; return Ok(connection.extra_headers.clone()); } - let key = nonblank(connection.api_key.clone()) - .map(|value| Sourced::new(value, connection.api_key_source)) - .or_else(|| { - nonblank(self.get_api_key_env_var().and_then(env_lookup)) - .map(|value| Sourced::new(value, InputSource::Environment)) - }); + let key = nonblank( + connection + .api_key + .as_ref() + .map(|key| key.expose().to_string()), + ) + .map(|value| Sourced::new(value, connection.api_key_source)) + .or_else(|| { + nonblank(self.get_api_key_env_var().and_then(env_lookup)) + .map(|value| Sourced::new(value, InputSource::Environment)) + }); if let Some(key) = key { super::common_utils::validate_destination(connection, key.source())?; return Ok(bearer_headers(connection, key.value())); @@ -196,7 +201,7 @@ mod tests { #[fixture] fn connection() -> OcrConnection { OcrConnection { - api_key: Some("request-key".into()), + api_key: Some(litellm_auth::SecretValue::new("request-key")), api_base: Some("https://example.com".into()), ..Default::default() } @@ -288,7 +293,7 @@ mod tests { #[tokio::test] async fn request_endpoint_accepts_request_owned_key() { let connection = OcrConnection { - api_key: Some("request-key".into()), + api_key: Some(litellm_auth::SecretValue::new("request-key")), api_key_source: InputSource::Request, api_base: Some("https://request.example".into()), api_base_source: InputSource::Request, diff --git a/litellm-rust/crates/llms/src/base_llm/ocr/error.rs b/litellm-rust/crates/llms/src/base_llm/ocr/error.rs index c3f481d7d44..3061a9fe2b2 100644 --- a/litellm-rust/crates/llms/src/base_llm/ocr/error.rs +++ b/litellm-rust/crates/llms/src/base_llm/ocr/error.rs @@ -102,6 +102,17 @@ pub enum Error { Headers(#[from] crate::custom_httpx::http_handler::HeaderError), } +impl From for Error { + fn from(fault: litellm_host::machine::MachineFault) -> Self { + use litellm_host::machine::MachineFault; + Self::InvalidRequest(match fault { + MachineFault::Abandoned => "OCR host driver was abandoned".into(), + MachineFault::Protocol(message) => format!("OCR {message}"), + MachineFault::Mismatch => "invalid OCR host operation result".into(), + }) + } +} + impl From for Error { fn from(error: litellm_core_utils::call_arguments::ArgumentError) -> Self { Self::RequestField { diff --git a/litellm-rust/crates/llms/src/base_llm/ocr/transformation.rs b/litellm-rust/crates/llms/src/base_llm/ocr/transformation.rs index e6fe5d9556d..8321dcfb4ce 100644 --- a/litellm-rust/crates/llms/src/base_llm/ocr/transformation.rs +++ b/litellm-rust/crates/llms/src/base_llm/ocr/transformation.rs @@ -1,6 +1,6 @@ use std::{collections::BTreeMap, future::Future, time::Duration}; -use litellm_auth::{InputSource, Sourced, TokenProviderHandle}; +use litellm_auth::{InputSource, SecretValue, Sourced, TokenProviderHandle}; use litellm_core_utils::{ call_arguments::CallArguments, serde_compat::{FiniteF64, LaxI64}, @@ -90,21 +90,22 @@ pub enum OcrResponseFormat { #[derive(Clone, Default)] pub struct OcrCredentialInputs { - pub api_key: Option>, - pub dynamic_api_key: Option>, + pub api_key: Option>, + pub dynamic_api_key: Option>, pub api_base: Option>, pub dynamic_api_base: Option>, } impl OcrCredentialInputs { pub fn new( - api_key: Option, + api_key: Option, api_key_source: InputSource, api_base: Option, api_base_source: InputSource, ) -> Self { Self { - api_key: nonblank(api_key).map(|value| Sourced::new(value, api_key_source)), + api_key: nonblank(api_key.as_ref().map(|key| key.expose().to_string())) + .map(|value| Sourced::new(SecretValue::new(value), api_key_source)), dynamic_api_key: None, api_base: nonblank(api_base).map(|value| Sourced::new(value, api_base_source)), dynamic_api_base: None, @@ -159,7 +160,7 @@ fn nonblank(value: Option) -> Option { #[derive(Clone)] pub struct OcrConnection { - pub api_key: Option, + pub api_key: Option, pub api_key_source: InputSource, pub api_base: Option, pub api_base_source: InputSource, @@ -209,7 +210,7 @@ impl Default for OcrConnection { #[derive(Clone, Default)] pub struct ResolvedOcrCredentials { - pub api_key: Option>, + pub api_key: Option>, pub api_base: Option>, } @@ -428,7 +429,7 @@ pub trait BaseOcrConfig: Send + Sync + Sized + 'static { ResolvedOcrCredentials { api_key: inputs .dynamic_api_key - .filter(|value| !value.value().is_empty()) + .filter(|value| !value.value().expose().is_empty()) .or(inputs.api_key), api_base: inputs .dynamic_api_base diff --git a/litellm-rust/crates/llms/src/cohere/ocr/transformation.rs b/litellm-rust/crates/llms/src/cohere/ocr/transformation.rs index f353c22d8c4..2528c967f41 100644 --- a/litellm-rust/crates/llms/src/cohere/ocr/transformation.rs +++ b/litellm-rust/crates/llms/src/cohere/ocr/transformation.rs @@ -179,8 +179,8 @@ impl CohereParseConfig { } let key = connection .api_key - .as_deref() - .map(str::trim) + .as_ref() + .map(|key| key.expose().trim()) .filter(|key| !key.is_empty()) .map(str::to_string) .or_else(|| { @@ -718,7 +718,7 @@ mod tests { assert!(matches!( CohereParseConfig.resolve_headers( &OcrConnection { - api_key: Some(" ".into()), + api_key: Some(litellm_auth::SecretValue::new(" ")), ..Default::default() }, &|_| None, diff --git a/litellm-rust/crates/llms/src/custom_httpx/llm_http_handler.rs b/litellm-rust/crates/llms/src/custom_httpx/llm_http_handler.rs index e93ddee3c50..fdd568d83fd 100644 --- a/litellm-rust/crates/llms/src/custom_httpx/llm_http_handler.rs +++ b/litellm-rust/crates/llms/src/custom_httpx/llm_http_handler.rs @@ -3,9 +3,9 @@ use std::{sync::OnceLock, time::Duration}; use bytes::{Bytes, BytesMut}; use futures_util::future::BoxFuture; use litellm_auth_gcp::VertexAuth; -use litellm_callbacks::event::{Passthrough, WireRequest}; +use litellm_host::event::WireRequest; use serde::{Serialize, de::DeserializeOwned}; -use serde_json::{Map, Value}; +use serde_json::Value; use crate::{ base_llm::ocr::{ @@ -26,11 +26,7 @@ use crate::{ /// The route's view of one call, handed to provider code that has to reach the /// caller's hooks mid-flight (guardrails on the outgoing body, raw response events). pub trait CallHooks: Send + Sync { - fn before_send( - &self, - wire: WireRequest, - passthrough_fields: Passthrough, - ) -> BoxFuture<'_, Result>; + fn before_send(&self, wire: WireRequest) -> BoxFuture<'_, Result>; fn response_received<'a>(&'a self, body: &'a [u8]) -> BoxFuture<'a, Result<(), E>>; } @@ -231,9 +227,8 @@ pub async fn transform_request_body( config.get_supported_ocr_params(&request.model), )?; config.validate_request_body(&composed)?; - let passthrough_fields = Passthrough::unchanged(&caller_inputs(request)?, &composed); let changed = hooks - .before_send(wire_request(url, headers, composed), passthrough_fields) + .before_send(wire_request(url, headers, composed)) .await?; if !changed.body.is_object() { return Err(Error::RequestField { @@ -252,21 +247,6 @@ fn wire_request(url: &str, headers: &[(String, String)], body: Value) -> WireReq } } -fn caller_inputs(request: &PreparedOcrRequest) -> Result, Error> { - let document = request - .caller_document - .then(|| serde_json::to_value(&request.document)) - .transpose() - .map_err(|_| Error::RequestField { - path: "document".into(), - })?; - let params: Map = request.optional_params.clone().into(); - Ok(params - .into_iter() - .chain(document.map(|document| ("document".to_string(), document))) - .collect()) -} - pub fn build_http_request( client: &OcrClient, request: &PreparedOcrRequest, @@ -294,9 +274,7 @@ pub async fn guardrail_document( let body = serde_json::to_value(&request.document).map_err(|_| Error::RequestField { path: "document".into(), })?; - let changed = hooks - .before_send(wire_request(url, headers, body), Passthrough::default()) - .await?; + let changed = hooks.before_send(wire_request(url, headers, body)).await?; let document = decode_request_value(changed.body, "guardrail.document")?; Ok((document, changed.headers)) } diff --git a/litellm-rust/crates/llms/src/mistral/ocr/transformation.rs b/litellm-rust/crates/llms/src/mistral/ocr/transformation.rs index c2038d0552d..9028f09c5ab 100644 --- a/litellm-rust/crates/llms/src/mistral/ocr/transformation.rs +++ b/litellm-rust/crates/llms/src/mistral/ocr/transformation.rs @@ -135,8 +135,8 @@ impl MistralOcrConfig { } let api_key = connection .api_key - .as_deref() - .map(str::trim) + .as_ref() + .map(|key| key.expose().trim()) .filter(|key| !key.is_empty()) .map(str::to_string) .or_else(|| { @@ -212,7 +212,7 @@ mod tests { #[default(vec![])] extra_headers: Vec<(String, String)>, ) -> OcrConnection { OcrConnection { - api_key: api_key.map(str::to_string), + api_key: api_key.map(litellm_auth::SecretValue::new), extra_headers, ..OcrConnection::default() } diff --git a/litellm-rust/crates/llms/src/reducto/ocr/transformation.rs b/litellm-rust/crates/llms/src/reducto/ocr/transformation.rs index ca2bae9c3bb..ec876fafb8f 100644 --- a/litellm-rust/crates/llms/src/reducto/ocr/transformation.rs +++ b/litellm-rust/crates/llms/src/reducto/ocr/transformation.rs @@ -442,8 +442,8 @@ fn resolve_headers( } let api_key = connection .api_key - .as_deref() - .map(str::trim) + .as_ref() + .map(|key| key.expose().trim()) .filter(|key| !key.is_empty()) .map(str::to_string) .or_else(|| { @@ -629,7 +629,7 @@ mod tests { #[test] fn explicit_key_precedes_environment_key() { let connection = OcrConnection { - api_key: Some("passed-key".into()), + api_key: Some(litellm_auth::SecretValue::new("passed-key")), ..Default::default() }; let headers = resolve_headers(&connection, &|_| Some("env-key".into())).unwrap(); @@ -639,7 +639,7 @@ mod tests { #[test] fn blank_explicit_key_uses_environment_key() { let connection = OcrConnection { - api_key: Some(" ".into()), + api_key: Some(litellm_auth::SecretValue::new(" ")), ..Default::default() }; let headers = resolve_headers(&connection, &|_| Some(" env-key ".into())).unwrap(); diff --git a/litellm-rust/crates/llms/src/vertex_ai/ocr/transformation.rs b/litellm-rust/crates/llms/src/vertex_ai/ocr/transformation.rs index ea0bcf3d08c..c2cb23d0010 100644 --- a/litellm-rust/crates/llms/src/vertex_ai/ocr/transformation.rs +++ b/litellm-rust/crates/llms/src/vertex_ai/ocr/transformation.rs @@ -134,7 +134,10 @@ impl VertexAiOcrConfig { .vertex_auth() .validate_environment( connection.extra_headers.clone(), - connection.api_key.as_deref(), + connection + .api_key + .as_ref() + .map(litellm_auth::SecretValue::expose), config, &credential_env, ) diff --git a/litellm-rust/crates/python-bridge/AGENTS.md b/litellm-rust/crates/python-bridge/AGENTS.md index 9932594e2f5..5dccfb4aca8 100644 --- a/litellm-rust/crates/python-bridge/AGENTS.md +++ b/litellm-rust/crates/python-bridge/AGENTS.md @@ -1,7 +1,7 @@ - Target invariants, not completion claims; these supersede older conflicting bridge guidance - Keep this crate the product-specific PyO3 consumer of `litellm-host-python` - Own registration, input projection, the route host and the caller callables it answers operations with (file readers, token providers), public response/error construction and the per-call composition of machine, route host and callback contract - - Legacy callback sharing (the caller's args, kwargs and request object, body/header roots, `passthrough_fields` re-aliasing) lives in `litellm-callbacks-legacy` behind `PublicCall` and `run_legacy_call`; the bridge hands the public call over and keeps no copy + - Legacy callback sharing (the caller's args, kwargs and request object, body/header roots, re-aliasing unchanged body keys) lives in `litellm-callbacks-legacy` behind `PublicCall` and `run_legacy_call`; the bridge hands the public call over and keeps no copy - Value-oriented execution, sync waiting, nested-runtime checks, signal polling and panic containment live in `litellm-host-python`; native async work uses `pyo3-async-runtimes`, Serde output uses `Pythonized` - Core owns typed native state, the route machine, provider preparation/I/O and normalization; the host driver owns terminal events; the legacy adapter in `litellm-callbacks-legacy` owns `Logging` dispatch policy - Python, Rust SDK and gateway use one lifecycle-bearing core route entrypoint; provider helpers stay private, never bridge-accessible transport drivers diff --git a/litellm-rust/crates/python-bridge/src/marshal.rs b/litellm-rust/crates/python-bridge/src/marshal.rs index ea4077b102f..2aba51cc4ff 100644 --- a/litellm-rust/crates/python-bridge/src/marshal.rs +++ b/litellm-rust/crates/python-bridge/src/marshal.rs @@ -18,10 +18,6 @@ pub(crate) struct RouteOptions { pub(crate) timeout: Option, } -pub(crate) fn body_argument(value: &Bound<'_, PyAny>) -> PyResult> { - required_object("body", from_py_argument(value)?) -} - pub(crate) fn messages_argument(value: &Bound<'_, PyAny>) -> PyResult> { match from_py_argument(value)? { Value::Array(values) => Ok(values), @@ -192,18 +188,6 @@ mod tests { json!([{"role": "user", "content": [{"type": "text", "text": "hi"}]}]) ); - let body = py - .eval( - c"{'model': 'claude', 'metadata': {'user': '1'}}", - None, - None, - ) - .unwrap(); - assert_eq!( - Value::Object(body_argument(&body).unwrap()), - json!({"model": "claude", "metadata": {"user": "1"}}) - ); - let params = py.eval(c"{'temperature': 0.2}", None, None).unwrap(); assert_eq!( optional_params_argument(¶ms).unwrap(), diff --git a/litellm-rust/crates/python-bridge/src/routes/messages.rs b/litellm-rust/crates/python-bridge/src/routes/messages.rs deleted file mode 100644 index daec931c92e..00000000000 --- a/litellm-rust/crates/python-bridge/src/routes/messages.rs +++ /dev/null @@ -1,88 +0,0 @@ -use litellm_core::messages::{Error, messages as run_messages, types::MessagesRequest}; -use litellm_host_python::{run_async, run_sync}; -use litellm_types::llms::anthropic_messages::anthropic_response::AnthropicMessagesResponse; -use pyo3::prelude::*; -use serde_json::{Map, Value}; - -use crate::{ - errors::messages_error_to_pyerr, - marshal::{RouteOptions, body_argument, extra_headers_argument, optional_timeout}, -}; - -async fn execute( - body: Map, - options: RouteOptions, -) -> Result { - let RouteOptions { - model, - api_key, - api_base, - custom_llm_provider, - extra_headers, - timeout, - } = options; - run_messages(MessagesRequest { - model: &model, - body: Value::Object(body), - api_key: api_key.as_deref(), - api_base: api_base.as_deref(), - custom_llm_provider: custom_llm_provider.as_deref(), - extra_headers, - timeout, - }) - .await -} - -#[pyfunction] -#[pyo3(signature = (model, body, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, timeout_seconds=None))] -#[expect( - clippy::too_many_arguments, - reason = "one parameter per Python keyword" -)] -pub(crate) fn messages( - py: Python<'_>, - model: String, - #[pyo3(from_py_with = body_argument)] body: Map, - api_key: Option, - api_base: Option, - custom_llm_provider: Option, - #[pyo3(from_py_with = extra_headers_argument)] extra_headers: Option>, - timeout_seconds: Option, -) -> PyResult> { - let options = RouteOptions { - model, - api_key, - api_base, - custom_llm_provider, - extra_headers, - timeout: optional_timeout(timeout_seconds), - }; - run_sync(py, execute(body, options), messages_error_to_pyerr) -} - -#[pyfunction] -#[pyo3(signature = (model, body, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, timeout_seconds=None))] -#[expect( - clippy::too_many_arguments, - reason = "one parameter per Python keyword" -)] -pub(crate) fn amessages<'py>( - py: Python<'py>, - model: String, - #[pyo3(from_py_with = body_argument)] body: Map, - api_key: Option, - api_base: Option, - custom_llm_provider: Option, - #[pyo3(from_py_with = extra_headers_argument)] extra_headers: Option>, - timeout_seconds: Option, -) -> PyResult> { - let options = RouteOptions { - model, - api_key, - api_base, - custom_llm_provider, - extra_headers, - timeout: optional_timeout(timeout_seconds), - }; - run_async(py, execute(body, options), messages_error_to_pyerr) -} diff --git a/litellm-rust/crates/python-bridge/src/routes/messages/host.rs b/litellm-rust/crates/python-bridge/src/routes/messages/host.rs new file mode 100644 index 00000000000..c1b3f59df58 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/routes/messages/host.rs @@ -0,0 +1,186 @@ +use bytes::Bytes; +use litellm_core::messages::{ + Error, + route::{Messages, MessagesCall, MessagesOp, MessagesOpResult, MessagesOutput}, +}; +use litellm_host_python::{InvokeError, RouteHost, from_py, lookup, to_py}; +use litellm_llms::custom_httpx::transport::Error as TransportError; +use pyo3::{ + exceptions::{PyException, PyValueError}, + gc::{PyTraverseError, PyVisit}, + prelude::*, + types::{PyBytes, PyDict}, +}; +use serde_json::{Map, Value}; + +use crate::{ + errors::{RustUpstreamError, messages_error_to_pyerr}, + marshal::{optional_timeout, python_timeout_seconds}, +}; + +/// The Anthropic Messages body fields a caller may pass besides `model` and `messages`, +/// as `AnthropicMessagesRequestOptionalParams` declares them. +const BODY_FIELDS: [&str; 20] = [ + "max_tokens", + "metadata", + "stop_sequences", + "stream", + "system", + "temperature", + "thinking", + "tool_choice", + "tools", + "top_k", + "inference_geo", + "top_p", + "mcp_servers", + "context_management", + "container", + "output_format", + "speed", + "output_config", + "cache_control", + "reasoning_effort", +]; + +/// The Python side of the Messages route: projects the prepared arguments and builds the +/// public response, chunks and exceptions. +pub(super) struct MessagesRouteHost { + request: Py, +} + +impl MessagesRouteHost { + pub(super) fn new(request: Py) -> Self { + Self { request } + } + + fn project(&self, py: Python<'_>, arguments: &Bound<'_, PyDict>) -> PyResult { + let request = self.request.bind(py); + let argument = |name: &str| -> PyResult>> { + Ok(lookup(arguments, request, name)?.filter(|value| !value.is_none())) + }; + let string = |name: &str| -> PyResult> { + argument(name)?.map(|value| value.extract()).transpose() + }; + let model = string("model")?.ok_or_else(|| PyValueError::new_err("model is required"))?; + let messages = + argument("messages")?.ok_or_else(|| PyValueError::new_err("messages is required"))?; + let fields = BODY_FIELDS + .iter() + .filter_map(|name| match argument(name) { + Ok(Some(value)) => Some(from_py(&value).map(|value| ((*name).to_string(), value))), + Ok(None) => None, + Err(error) => Some(Err(error)), + }) + .collect::>>()?; + let body = [ + ("model".to_string(), Value::String(model.clone())), + ("messages".to_string(), from_py(&messages)?), + ] + .into_iter() + .chain(fields) + .collect::>(); + let timeout = argument("timeout")? + .map(|value| python_timeout_seconds(py, value.unbind())) + .transpose()? + .flatten(); + Ok(MessagesCall { + model, + body, + api_key: string("api_key")?, + api_base: string("api_base")?, + custom_llm_provider: string("custom_llm_provider")?, + extra_headers: argument("extra_headers")? + .map(|value| from_py(&value)) + .transpose()?, + timeout: optional_timeout(timeout), + }) + } + + fn provider(&self, py: Python<'_>) -> String { + self.request + .bind(py) + .getattr("custom_llm_provider") + .and_then(|value| value.extract::>()) + .ok() + .flatten() + .unwrap_or_else(|| "anthropic".into()) + } + + fn map_failure(&self, py: Python<'_>, error: PyErr) -> PyErr { + if !error.is_instance_of::(py) { + return error; + } + let mapped = py + .import("litellm.rust_bridge.messages.route_host") + .and_then(|module| module.getattr("map_failure")) + .and_then(|map| map.call1((error.value(py), self.request.bind(py), self.provider(py)))) + .and_then(|mapped| { + mapped + .extract::>() + .map_err(PyErr::from) + }); + match mapped { + Ok(mapped) => PyErr::from_value(mapped.into_bound(py).into_any()), + Err(_) => error, + } + } +} + +impl RouteHost for MessagesRouteHost { + type Route = Messages; + type Failure = PyErr; + + fn invoke( + &mut self, + py: Python<'_>, + arguments: &Bound<'_, PyDict>, + op: MessagesOp, + ) -> Result> { + match op { + MessagesOp::ProjectRequest => self + .project(py, arguments) + .map(|call| MessagesOpResult::Request(Box::new(call))) + .map_err(|error| InvokeError::Python(self.map_failure(py, error))), + } + } + + fn complete(&mut self, py: Python<'_>, response: MessagesOutput) -> PyResult> { + match response { + MessagesOutput::Message(message) => py + .import("litellm.rust_bridge.messages.route_host")? + .getattr("response")? + .call1((to_py(py, message.as_ref())?,)) + .map(Bound::unbind), + MessagesOutput::Streamed => Ok(py.None()), + } + } + + fn chunk(&mut self, py: Python<'_>, chunk: Bytes) -> PyResult> { + Ok(PyBytes::new(py, &chunk).into_any().unbind()) + } + + fn classify(&self, py: Python<'_>, error: Error) -> PyResult { + let native = match error { + Error::Transport(TransportError::Http { status, body }) => { + let error = RustUpstreamError::new_err((status, body)); + error + .value(py) + .setattr("headers", Vec::<(String, String)>::new())?; + error + } + other => messages_error_to_pyerr(other), + }; + Ok(self.map_failure(py, native)) + } + + fn host_error(error: &PyErr) -> Error { + Error::InvalidRequest(error.to_string()) + } + + fn close(&mut self, _: Python<'_>) {} + + fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { + visit.call(&self.request) + } +} diff --git a/litellm-rust/crates/python-bridge/src/routes/messages/mod.rs b/litellm-rust/crates/python-bridge/src/routes/messages/mod.rs new file mode 100644 index 00000000000..8c42315ac59 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/routes/messages/mod.rs @@ -0,0 +1,68 @@ +mod host; + +use host::MessagesRouteHost; +use litellm_callbacks_legacy::{LegacySurface, PassThroughStream, PublicCall, run_legacy_call}; +use litellm_core::messages::route::{messages_machine, supports}; +use pyo3::{ + prelude::*, + types::{PyDict, PyTuple}, +}; + +use crate::errors::RustBridgeDeclined; + +const SURFACE: LegacySurface = LegacySurface { + call_type: "anthropic_messages", + input_description: "Messages", + stream: Some(PassThroughStream { + url_route: "/v1/messages", + endpoint_type: "anthropic", + }), +}; + +fn run_messages( + py: Python<'_>, + request: Bound<'_, PyAny>, + args: Bound<'_, PyTuple>, + kwargs: Bound<'_, PyDict>, + asynchronous: bool, +) -> PyResult> { + let model: String = request.getattr("model")?.extract()?; + let provider: Option = request.getattr("custom_llm_provider")?.extract()?; + let stream = request + .getattr("stream")? + .extract::>()? + .unwrap_or(false); + if !supports(&model, provider.as_deref(), stream) { + return Err(RustBridgeDeclined::new_err( + "the Rust Messages route does not serve this provider", + )); + } + run_legacy_call( + py, + SURFACE, + PublicCall::capture(&request, &args, &kwargs)?, + messages_machine(), + MessagesRouteHost::new(request.unbind()), + asynchronous, + ) +} + +#[pyfunction] +pub(crate) fn messages( + py: Python<'_>, + request: Bound<'_, PyAny>, + args: Bound<'_, PyTuple>, + kwargs: Bound<'_, PyDict>, +) -> PyResult> { + run_messages(py, request, args, kwargs, false) +} + +#[pyfunction] +pub(crate) fn amessages( + py: Python<'_>, + request: Bound<'_, PyAny>, + args: Bound<'_, PyTuple>, + kwargs: Bound<'_, PyDict>, +) -> PyResult> { + run_messages(py, request, args, kwargs, true) +} diff --git a/litellm-rust/crates/python-bridge/src/routes/mod.rs b/litellm-rust/crates/python-bridge/src/routes/mod.rs index f59e32a28e2..2d6b849a6b1 100644 --- a/litellm-rust/crates/python-bridge/src/routes/mod.rs +++ b/litellm-rust/crates/python-bridge/src/routes/mod.rs @@ -22,11 +22,6 @@ mod tests { "atranscription", "(model, audio, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, optional_params=None, timeout_seconds=None)", ), - ( - "messages", - "amessages", - "(model, body, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, timeout_seconds=None)", - ), ( "chat_completions", "achat_completions", @@ -113,25 +108,6 @@ value = Broken() ); assert_eq!(async_chat_error.to_string(), sync_chat_error.to_string()); - let invalid_body = PyList::empty(py); - let sync_messages_error = module - .getattr("messages") - .and_then(|function| function.call1(("model", &invalid_body))) - .expect_err("sync Messages should reject a non-dict body"); - let async_messages_error = module - .getattr("amessages") - .and_then(|function| function.call1(("model", &invalid_body))) - .expect_err("async Messages should reject a non-dict body"); - - assert_eq!( - sync_messages_error.to_string(), - "ValueError: body must be a dict" - ); - assert_eq!( - async_messages_error.to_string(), - sync_messages_error.to_string() - ); - let invalid_headers = PyList::empty(py); let kwargs = PyDict::new(py); kwargs @@ -193,13 +169,6 @@ value = Broken() headers_kwargs .set_item("extra_headers", &invalid) .expect("kwargs should accept extra_headers"); - let invalid_body = PyList::empty(py); - let error = module - .getattr("messages") - .and_then(|function| function.call(("model", &invalid_body), Some(&headers_kwargs))) - .expect_err("body should be validated before headers"); - assert_eq!(error.to_string(), "ValueError: body must be a dict"); - let invalid_payload = PyModule::new(py, "invalid_payload").expect("invalid payload should be created"); let error = module diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr/host.rs b/litellm-rust/crates/python-bridge/src/routes/ocr/host.rs index 9dc891a91d6..77c8d5d6641 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr/host.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/host.rs @@ -1,9 +1,9 @@ use litellm_auth::ResolvedCredential; use litellm_core::ocr::route::{Ocr, OcrOp, OcrOpResult}; -use litellm_host_python::{RouteHost, missing_state, to_py}; +use litellm_host_python::{InvokeError, RouteHost, missing_state, to_py}; use litellm_llms::base_llm::ocr::{error::Error, transformation::LiteLLMOcrResponse}; use pyo3::{ - exceptions::PyBaseException, + exceptions::{PyBaseException, PyException}, gc::{PyTraverseError, PyVisit}, prelude::*, types::PyDict, @@ -57,12 +57,8 @@ impl OcrRouteHost { .ok_or_else(missing_state)? .acquire(py) } -} -impl RouteHost for OcrRouteHost { - type Route = Ocr; - - fn invoke( + fn answer( &mut self, py: Python<'_>, arguments: &Bound<'_, PyDict>, @@ -88,6 +84,40 @@ impl RouteHost for OcrRouteHost { } } + fn map_failure(&self, py: Python<'_>, error: PyErr) -> PyErr { + if !error.is_instance_of::(py) { + return error; + } + let provider = match &self.data { + OcrHostData::Projected(handles) => handles.provider, + _ => "", + }; + let mapped = py + .import("litellm.rust_bridge.ocr.route_host") + .and_then(|module| module.getattr("map_failure")) + .and_then(|map| map.call1((error.value(py), self.request.bind(py), provider))) + .and_then(|mapped| mapped.extract::>().map_err(PyErr::from)); + match mapped { + Ok(mapped) => PyErr::from_value(mapped.into_bound(py).into_any()), + Err(_) => error, + } + } +} + +impl RouteHost for OcrRouteHost { + type Route = Ocr; + type Failure = PyErr; + + fn invoke( + &mut self, + py: Python<'_>, + arguments: &Bound<'_, PyDict>, + op: OcrOp, + ) -> Result> { + self.answer(py, arguments, op) + .map_err(|error| InvokeError::Python(self.map_failure(py, error))) + } + fn complete(&mut self, py: Python<'_>, response: LiteLLMOcrResponse) -> PyResult> { py.import("litellm.rust_bridge.ocr.route_host")? .getattr("response")? @@ -95,27 +125,18 @@ impl RouteHost for OcrRouteHost { .map(Bound::unbind) } - fn native_error(error: Error) -> PyErr { - ocr_error_to_pyerr(error) + fn chunk(&mut self, _: Python<'_>, chunk: std::convert::Infallible) -> PyResult> { + match chunk {} + } + + fn classify(&self, py: Python<'_>, error: Error) -> PyResult { + Ok(self.map_failure(py, ocr_error_to_pyerr(error))) } fn host_error(error: &PyErr) -> Error { Error::InvalidRequest(error.to_string()) } - fn map_failure(&self, py: Python<'_>, error: &PyErr) -> PyResult { - let provider = match &self.data { - OcrHostData::Projected(handles) => handles.provider, - _ => "", - }; - let mapped: Py = py - .import("litellm.rust_bridge.ocr.route_host")? - .getattr("map_failure")? - .call1((error.value(py), self.request.bind(py), provider))? - .extract()?; - Ok(PyErr::from_value(mapped.into_bound(py).into_any())) - } - fn close(&mut self, _: Python<'_>) { self.data = OcrHostData::Released; } diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs b/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs index b5bb941708d..8afa1e2a906 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs @@ -15,6 +15,7 @@ use pyo3::{ const SURFACE: LegacySurface = LegacySurface { call_type: "ocr", input_description: "OCR document processing", + stream: None, }; const ASYNC_SURFACE: LegacySurface = LegacySurface { diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr/project.rs b/litellm-rust/crates/python-bridge/src/routes/ocr/project.rs index 7ffa129f85c..5dd2aa804b8 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr/project.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/project.rs @@ -1,3 +1,4 @@ +use litellm_auth::SecretValue; use litellm_core::ocr::{ types::{LiteLLMOcrRequest, OcrDocumentInput}, wire::{OcrWireRequest, consumed_optional_params, decode_document, decode_request_input}, @@ -31,7 +32,7 @@ struct OcrArguments<'a, 'py> { impl<'py> OcrArguments<'_, 'py> { fn lookup(&self, name: &str) -> PyResult> { - litellm_callbacks_legacy::lookup(self.kwargs, self.request, name)? + litellm_host_python::lookup(self.kwargs, self.request, name)? .ok_or_else(|| PyValueError::new_err(format!("missing argument: {name}"))) } @@ -47,8 +48,11 @@ impl<'py> OcrArguments<'_, 'py> { self.lookup("document") } - fn api_key(&self) -> PyResult> { - self.lookup("api_key")?.extract() + fn api_key(&self) -> PyResult> { + Ok(self + .lookup("api_key")? + .extract::>()? + .map(SecretValue::new)) } fn api_base(&self) -> PyResult> { diff --git a/litellm/integrations/otel/plumbing/context.py b/litellm/integrations/otel/plumbing/context.py index 1282e654365..19243d64c64 100644 --- a/litellm/integrations/otel/plumbing/context.py +++ b/litellm/integrations/otel/plumbing/context.py @@ -325,12 +325,27 @@ def _outgoing_trace_context(parent_span: object) -> Context | None: return None +def _propagated_context(headers: Mapping[str, str], request_context: Context) -> Context: + """``request_context`` when it continues the trace ``headers`` already name, else the + caller's own context, so an explicit upstream ``traceparent`` (``x-pass-traceparent``) + is never swapped for an unrelated trace and its ``tracestate`` survives.""" + caller: Final = extract_traceparent(headers) + if caller is None: + return request_context + caller_span: Final = get_current_span(caller).get_span_context() + request_span: Final = get_current_span(request_context).get_span_context() + if not caller_span.is_valid or caller_span.trace_id == request_span.trace_id: + return request_context + return caller + + def inject_trace_context(headers: Mapping[str, str], parent_span: object = None) -> dict[str, str]: """``headers`` plus W3C ``traceparent``/``tracestate`` for this request's span. Parent preference: ``parent_span`` (the request span auth stashed on the key), then the anchored request root span, then the ambient active span. Only trace context is - injected, never Baggage. Unchanged when no valid span exists anywhere. + injected, never Baggage. Unchanged when no valid span exists anywhere. A ``traceparent`` + already in ``headers`` from a different trace is forwarded as-is instead of replaced. """ context: Final = _outgoing_trace_context(parent_span) if context is None: @@ -338,7 +353,7 @@ def inject_trace_context(headers: Mapping[str, str], parent_span: object = None) carrier: Final = { # mutable-ok: OpenTelemetry propagator requires a mutable carrier key: value for key, value in headers.items() if key.lower() not in _W3C_TRACE_HEADERS } - _PROPAGATOR.inject(carrier, context=context) + _PROPAGATOR.inject(carrier, context=_propagated_context(headers, context)) return carrier diff --git a/litellm/litellm_core_utils/core_helpers.py b/litellm/litellm_core_utils/core_helpers.py index 15380bc5d57..d29b1fc74ef 100644 --- a/litellm/litellm_core_utils/core_helpers.py +++ b/litellm/litellm_core_utils/core_helpers.py @@ -224,6 +224,12 @@ _FINISH_REASON_MAP: Final[dict[str, OpenAIChatCompletionFinishReason]] = { "IMAGE_PROHIBITED_CONTENT": "content_filter", "TOO_MANY_TOOL_CALLS": "stop", "MALFORMED_RESPONSE": "stop", + "NO_IMAGE": "content_filter", + "IMAGE_RECITATION": "content_filter", + "IMAGE_OTHER": "content_filter", + "ESCALATION": "content_filter", + "UNEXPECTED_TOOL_CALL": "stop", + "MISSING_THOUGHT_SIGNATURE": "stop", # Zhipu GLM "network_error": "stop", "sensitive": "content_filter", diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 4a9a65b1485..7248c2f3590 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -1265,11 +1265,6 @@ class Logging(LiteLLMLoggingBaseClass): additional_args.get("api_base", "") ) - def record_api_call_start_time(self) -> None: - self.model_call_details["api_call_start_time"] = datetime.datetime.now() - if self.model_call_details.get("first_api_call_start_time") is None: - self.model_call_details["first_api_call_start_time"] = self.model_call_details["api_call_start_time"] - def pre_call(self, input, api_key, model=None, additional_args={}): # Log the exact input to the LLM API try: @@ -1334,7 +1329,15 @@ class Logging(LiteLLMLoggingBaseClass): "LiteLLM.LoggingError: [Non-Blocking] Exception occurred while logging %s", e ) - self.record_api_call_start_time() + self.model_call_details["api_call_start_time"] = datetime.datetime.now() + # Set-once first provider-handoff instant. api_call_start_time + # is overwritten on every retry, so it can't measure one-time + # preprocessing; pinning the first attempt excludes retry loops + # + backoff. Logging object only — must NOT go into + # litellm_params["metadata"] (caller request metadata, typed + # Dict[str, str], echoed downstream; a datetime breaks it). + if self.model_call_details.get("first_api_call_start_time") is None: + self.model_call_details["first_api_call_start_time"] = self.model_call_details["api_call_start_time"] # Input Integration Logging -> If you want to log the fact that an attempt to call the model was made callbacks: Final = litellm.input_callback + (self.dynamic_input_callbacks or []) for callback in callbacks: @@ -1468,21 +1471,16 @@ class Logging(LiteLLMLoggingBaseClass): """ return _get_masked_values(headers, ignore_sensitive_values=ignore_sensitive_headers) - def record_post_call( - self, original_response: object, input: object, api_key: object, additional_args: dict[str, object] - ) -> None: - self.model_call_details["input"] = input - self.model_call_details["api_key"] = api_key - self.model_call_details["original_response"] = original_response - self.model_call_details["additional_args"] = additional_args - self.model_call_details["log_event_type"] = "post_api_call" - def post_call(self, original_response, input=None, api_key=None, additional_args={}): # Log the exact result from the LLM API, for streaming - log the type of response received if isinstance(original_response, dict): original_response = json.dumps(original_response, default=str) try: - self.record_post_call(original_response, input, api_key, additional_args) + self.model_call_details["input"] = input + self.model_call_details["api_key"] = api_key + self.model_call_details["original_response"] = original_response + self.model_call_details["additional_args"] = additional_args + self.model_call_details["log_event_type"] = "post_api_call" attr: Literal["warning", "debug"] if self.litellm_request_debug: @@ -2177,7 +2175,6 @@ class Logging(LiteLLMLoggingBaseClass): logging_result, start_time, end_time, - build_logging_payload: bool = True, ): """Resolve hidden params, compute response cost, and emit the standard logging payload.""" hidden_params: Final = getattr(logging_result, "_hidden_params", {}) @@ -2202,9 +2199,6 @@ class Logging(LiteLLMLoggingBaseClass): else: self.model_call_details["response_cost"] = self._response_cost_calculator(result=logging_result) - if not build_logging_payload: - return - self.model_call_details["standard_logging_object"] = self._build_standard_logging_payload( logging_result, start_time, end_time ) @@ -2266,7 +2260,6 @@ class Logging(LiteLLMLoggingBaseClass): end_time=None, cache_hit=None, standard_logging_object: StandardLoggingPayload | None = None, - build_logging_payload: bool = True, ): try: if start_time is None: @@ -2304,7 +2297,6 @@ class Logging(LiteLLMLoggingBaseClass): logging_result=logging_result, start_time=start_time, end_time=end_time, - build_logging_payload=build_logging_payload, ) elif standard_logging_object is not None: self.model_call_details["standard_logging_object"] = standard_logging_object @@ -3328,9 +3320,7 @@ class Logging(LiteLLMLoggingBaseClass): except Exception as e: verbose_logger.debug("Error in _handle_callback_failure: %s", e) - def _failure_handler_helper_fn( - self, exception, traceback_exception, start_time=None, end_time=None, build_logging_payload: bool = True - ): + def _failure_handler_helper_fn(self, exception, traceback_exception, start_time=None, end_time=None): if start_time is None: start_time = self.start_time if end_time is None: @@ -3365,9 +3355,6 @@ class Logging(LiteLLMLoggingBaseClass): metadata: Final = self.model_call_details["litellm_params"].get("metadata", {}) or {} metadata.update(exception.headers) - if not build_logging_payload: - return start_time, end_time - ## STANDARDIZED LOGGING PAYLOAD self.model_call_details["standard_logging_object"] = get_standard_logging_object_payload( diff --git a/litellm/llms/anthropic/chat/transformation.py b/litellm/llms/anthropic/chat/transformation.py index 0f99441a115..1f90d375bc2 100644 --- a/litellm/llms/anthropic/chat/transformation.py +++ b/litellm/llms/anthropic/chat/transformation.py @@ -94,6 +94,7 @@ from litellm.utils import ( from ..common_utils import ( AnthropicError, AnthropicModelInfo, + eager_input_streaming_flag, process_anthropic_headers, strip_advisor_blocks_from_messages, ) @@ -732,10 +733,20 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): input_anthropic_schema: Final = sanitize_input_schema_for_anthropic(_input_schema) - _tool: Final = AnthropicMessagesTool( - name=tool["function"]["name"], - input_schema=input_anthropic_schema, - type="custom", + _eager_input_streaming: Final = eager_input_streaming_flag(tool) + _tool: Final = ( + AnthropicMessagesTool( + name=tool["function"]["name"], + input_schema=input_anthropic_schema, + type="custom", + ) + if _eager_input_streaming is None + else AnthropicMessagesTool( + name=tool["function"]["name"], + input_schema=input_anthropic_schema, + type="custom", + eager_input_streaming=_eager_input_streaming, + ) ) _description: Final = tool["function"].get("description") diff --git a/litellm/llms/anthropic/common_utils.py b/litellm/llms/anthropic/common_utils.py index d35a9372058..2de9ab41d95 100644 --- a/litellm/llms/anthropic/common_utils.py +++ b/litellm/llms/anthropic/common_utils.py @@ -10,7 +10,7 @@ from types import MappingProxyType from typing import Any, Final, Literal import httpx -from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError +from pydantic import BaseModel, ConfigDict, StrictBool, TypeAdapter, ValidationError import litellm from litellm.constants import ( @@ -19,6 +19,7 @@ from litellm.constants import ( DEFAULT_REASONING_EFFORT_MEDIUM_THINKING_BUDGET, DEFAULT_REASONING_EFFORT_XHIGH_THINKING_BUDGET, ) +from litellm.exceptions import UnsupportedParamsError from litellm.litellm_core_utils.prompt_templates.common_utils import ( get_file_ids_from_messages, is_encrypted_reasoning_block, @@ -231,6 +232,27 @@ def optionally_handle_anthropic_oauth(headers: dict, api_key: str | None) -> tup return headers, api_key +class _EagerInputStreamingFunction(BaseModel): + eager_input_streaming: StrictBool | None = None + + +class _EagerInputStreamingTool(BaseModel): + eager_input_streaming: StrictBool | None = None + function: _EagerInputStreamingFunction | None = None + + +def eager_input_streaming_flag(tool: object) -> bool | None: + try: + parsed: Final = _EagerInputStreamingTool.model_validate(tool) + except ValidationError as error: + if isinstance(tool, Mapping): + raise UnsupportedParamsError(message="eager_input_streaming must be a boolean") from error + return None + if parsed.eager_input_streaming is not None: + return parsed.eager_input_streaming + return parsed.function.eager_input_streaming if parsed.function is not None else None + + class AnthropicError(BaseLLMException): def __init__( self, @@ -373,6 +395,9 @@ class AnthropicModelInfo(BaseLLMModelInfo): return False + def is_eager_input_streaming_used(self, tools: Sequence[object] | None) -> bool: + return any(eager_input_streaming_flag(tool) is True for tool in tools or ()) + @staticmethod def _supports_sampling_params(model: str) -> bool: """Claude 4.7+ (Opus 4.7/4.8, Fable 5) removed sampling params: the API diff --git a/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py b/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py index 7eb56ae55d3..e9235bc80a7 100644 --- a/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py +++ b/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py @@ -111,6 +111,7 @@ from litellm.litellm_core_utils.reasoning_effort_utils import ( reasoning_effort_from_thinking_budget, ) from litellm.llms.anthropic.common_utils import ( + eager_input_streaming_flag, is_empty_unsigned_thinking_block, normalize_anthropic_tool_use_id, strip_encrypted_reasoning_blocks_from_anthropic_messages, @@ -197,6 +198,15 @@ def target_supports_mid_conversation_system(model: str | None, custom_llm_provid return supports_mid_conversation_system(model=model, custom_llm_provider=custom_llm_provider) +def _chat_tool_param(function_chunk: ChatCompletionToolParamFunctionChunk, tool: object) -> ChatCompletionToolParam: + eager_input_streaming: Final = eager_input_streaming_flag(tool) + if eager_input_streaming is None: + return ChatCompletionToolParam(type="function", function=function_chunk) + return ChatCompletionToolParam( + type="function", function=function_chunk, eager_input_streaming=eager_input_streaming + ) + + class AnthropicAdapter: def __init__(self) -> None: pass @@ -770,6 +780,7 @@ class LiteLLMAnthropicMessagesAdapter: "cache_control", "strict", "type", + "eager_input_streaming", ] for idx, tool in enumerate(tools): @@ -808,7 +819,7 @@ class LiteLLMAnthropicMessagesAdapter: for k, v in tool.items(): if k not in mapped_tool_params: # pass additional computer kwargs function_chunk.setdefault("parameters", {}).update({k: v}) - tool_param = ChatCompletionToolParam(type="function", function=function_chunk) + tool_param = _chat_tool_param(function_chunk, tool) self._add_cache_control_if_applicable(tool, tool_param, model) new_tools.append(tool_param) @@ -1399,6 +1410,8 @@ class LiteLLMAnthropicMessagesAdapter: return "max_tokens" elif openai_finish_reason == "tool_calls": return "tool_use" + elif openai_finish_reason in ["content_filter", "refusal"]: + return "refusal" return "end_turn" @staticmethod diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index 1a4f27e2ddd..1a32fec45e3 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/litellm/llms/bedrock/chat/converse_transformation.py @@ -48,6 +48,7 @@ from litellm.llms.bedrock.request_metadata import ( merge_bedrock_invoke_headers, resolve_bedrock_request_metadata, ) +from litellm.types.llms.anthropic import ANTHROPIC_FINE_GRAINED_TOOL_STREAMING_BETA_HEADER from litellm.types.llms.bedrock import * from litellm.types.llms.openai import ( AllMessageValues, @@ -1517,12 +1518,6 @@ class AmazonConverseConfig(BaseConfig): """Process tools and collect anthropic_beta values.""" bedrock_tools: list[ToolBlock] = [] - # Collect anthropic_beta values from user headers - anthropic_beta_list: Final = [] - if headers: - user_betas: Final = get_anthropic_beta_from_headers(headers) - anthropic_beta_list.extend(user_betas) - # Separate pre-formatted Bedrock tools (e.g. systemTool from web_search_options) # from OpenAI-format tools that need transformation via _bedrock_tools_pt filtered_tools: Final = [] @@ -1542,6 +1537,17 @@ class AmazonConverseConfig(BaseConfig): continue filtered_tools.append(tool) + base_model: Final = BedrockModelInfo.get_base_model(model) + client_beta_list: Final = get_anthropic_beta_from_headers(headers or {}) + eager_beta: Final = ( + (ANTHROPIC_FINE_GRAINED_TOOL_STREAMING_BETA_HEADER,) + if base_model.startswith("anthropic") + and AnthropicModelInfo().is_eager_input_streaming_used(filtered_tools) + and ANTHROPIC_FINE_GRAINED_TOOL_STREAMING_BETA_HEADER not in client_beta_list + else () + ) + anthropic_beta_list: Final = [*client_beta_list, *eager_beta] + # Only separate tools if computer use tools are actually present if filtered_tools and self.is_computer_use_tool_used(filtered_tools, model): # Separate computer use tools from regular function tools @@ -1619,7 +1625,6 @@ class AmazonConverseConfig(BaseConfig): # Opus 4.5 gates ``output_config.effort`` behind a beta header; # Claude 4.6/4.7 accept it without one. - base_model: Final = BedrockModelInfo.get_base_model(model) if base_model.startswith("anthropic"): output_config: Final = additional_request_params.get("output_config") if ( diff --git a/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py index 72bc43ba938..1326dc22ca0 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py @@ -24,8 +24,12 @@ from litellm.llms.bedrock.common_utils import ( normalize_custom_field_on_tools, normalize_tool_input_schema_types_for_bedrock_invoke, strip_unsupported_bedrock_invoke_output_config_keys, + tools_without_eager_input_streaming, +) +from litellm.types.llms.anthropic import ( + ANTHROPIC_FINE_GRAINED_TOOL_STREAMING_BETA_HEADER, + ANTHROPIC_TOOL_SEARCH_BETA_HEADER, ) -from litellm.types.llms.anthropic import ANTHROPIC_TOOL_SEARCH_BETA_HEADER from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import ModelResponse @@ -237,6 +241,9 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig): # Hoist `custom.defer_loading` then drop `custom` (Bedrock doesn't support it) normalize_custom_field_on_tools(anthropic_request) normalize_tool_input_schema_types_for_bedrock_invoke(anthropic_request) + outbound_tools: Final = tools_without_eager_input_streaming(anthropic_request) + if outbound_tools is not None: + anthropic_request["tools"] = outbound_tools return anthropic_request def _compute_bedrock_invoke_beta_headers( @@ -269,6 +276,9 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig): if bedrock_supports_tool_search(model): beta_set.add("tool-search-tool-2025-10-19") + if self.is_eager_input_streaming_used(tools): + beta_set.add(ANTHROPIC_FINE_GRAINED_TOOL_STREAMING_BETA_HEADER) + auto_beta_list: Final = filter_and_transform_beta_headers( beta_headers=list(beta_set - user_beta_set), provider="bedrock", diff --git a/litellm/llms/bedrock/common_utils.py b/litellm/llms/bedrock/common_utils.py index 62b485588ac..7e24292a87e 100644 --- a/litellm/llms/bedrock/common_utils.py +++ b/litellm/llms/bedrock/common_utils.py @@ -9,7 +9,7 @@ import functools import json import os import re -from collections.abc import Mapping +from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Any, Final, Literal, TypedDict if TYPE_CHECKING: @@ -18,6 +18,7 @@ if TYPE_CHECKING: from litellm.types.llms.bedrock import BedrockCreateBatchRequest import httpx +from pydantic import TypeAdapter, ValidationError import litellm from litellm import verbose_logger @@ -330,6 +331,17 @@ def normalize_custom_field_on_tools(request_body: dict) -> None: tool["defer_loading"] = deferred +_TOOL_DICTS_ADAPTER: Final = TypeAdapter(tuple[Mapping[str, object], ...]) + + +def tools_without_eager_input_streaming(request_body: Mapping[str, object]) -> Sequence[object] | None: + try: + tools: Final = _TOOL_DICTS_ADAPTER.validate_python(request_body.get("tools")) + except ValidationError: + return None + return [{key: value for key, value in tool.items() if key != "eager_input_streaming"} for tool in tools] + + def normalize_json_schema_custom_types_to_object(schema: dict) -> None: """ In-place: replace JSON Schema ``type: \"custom\"`` with ``\"object\"`` (iterative walk). diff --git a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py index 4aa2afdbc78..d2be1ad9156 100644 --- a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py +++ b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py @@ -39,6 +39,7 @@ from litellm.llms.bedrock.common_utils import ( normalize_custom_field_on_tools, normalize_tool_input_schema_types_for_bedrock_invoke, strip_unsupported_bedrock_invoke_output_config_keys, + tools_without_eager_input_streaming, ) from litellm.llms.bedrock.request_metadata import ( bedrock_request_metadata_headers, @@ -46,6 +47,7 @@ from litellm.llms.bedrock.request_metadata import ( ) from litellm.types.llms.anthropic import ( ANTHROPIC_BETA_HEADER_VALUES, + ANTHROPIC_FINE_GRAINED_TOOL_STREAMING_BETA_HEADER, ANTHROPIC_TOOL_SEARCH_BETA_HEADER, ) from litellm.types.llms.bedrock import BedrockInvokeAnthropicMessagesRequest @@ -525,6 +527,9 @@ class AmazonAnthropicClaudeMessagesConfig( if injected_thinking_for_clear_thinking: beta_set.add("interleaved-thinking-2025-05-14") + if anthropic_model_info.is_eager_input_streaming_used(tools): + beta_set.add(ANTHROPIC_FINE_GRAINED_TOOL_STREAMING_BETA_HEADER) + self._filter_context_management_for_bedrock_invoke( anthropic_messages_request=anthropic_messages_request, beta_set=beta_set, @@ -719,6 +724,10 @@ class AmazonAnthropicClaudeMessagesConfig( if filtered_betas: anthropic_messages_request["anthropic_beta"] = filtered_betas + outbound_tools: Final = tools_without_eager_input_streaming(anthropic_messages_request) + if outbound_tools is not None: + anthropic_messages_request["tools"] = outbound_tools + remaining_output_config: Final = anthropic_messages_request.get("output_config") if ( litellm.drop_params is True diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 228d170a937..b87d0deb9bc 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -144,6 +144,7 @@ from litellm.types.llms.openai import ( from litellm.types.realtime import RealtimeQueryParams from litellm.types.rerank import RerankResponse from litellm.types.responses.main import DeleteResponseResult +from litellm.types.responses.streaming_websocket import ResponsesWebSocketRequestDefaults from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import ( CallTypes, @@ -6593,6 +6594,7 @@ class BaseLLMHTTPHandler: litellm_metadata: dict[str, object] | None = None, custom_llm_provider: str | None = None, first_message: str | None = None, + request_defaults: ResponsesWebSocketRequestDefaults | None = None, **kwargs: Any, ): """ @@ -6747,6 +6749,7 @@ class BaseLLMHTTPHandler: output_guardrail_callbacks=_ws_output_guardrail_callbacks, quota_callbacks=_ws_quota_callbacks, authorized_model=model, + request_defaults=request_defaults, ) await streaming.bidirectional_forward() diff --git a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py index e8b316b5902..46f1b948026 100644 --- a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py +++ b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py @@ -6,7 +6,7 @@ import time from collections.abc import Callable, Mapping, Sequence from copy import deepcopy from functools import partial -from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Union, cast +from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Union, cast, get_args import httpx @@ -57,6 +57,7 @@ from litellm.types.llms.vertex_ai import ( ContentType, FunctionCallingConfig, FunctionDeclaration, + GeminiFinishReason, GeminiThinkingConfig, GenerateContentResponseBody, HttpxPartType, @@ -1330,25 +1331,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): "IMAGE_PROHIBITED_CONTENT": "The token generation was stopped as the response was flagged for prohibited image content.", } - _GEMINI_FINISH_REASON_KEYS = frozenset( - { - "STOP", - "MAX_TOKENS", - "SAFETY", - "RECITATION", - "FINISH_REASON_UNSPECIFIED", - "MALFORMED_FUNCTION_CALL", - "LANGUAGE", - "OTHER", - "BLOCKLIST", - "PROHIBITED_CONTENT", - "SPII", - "IMAGE_SAFETY", - "IMAGE_PROHIBITED_CONTENT", - "TOO_MANY_TOOL_CALLS", - "MALFORMED_RESPONSE", - } - ) + _GEMINI_FINISH_REASON_KEYS: Final[frozenset[str]] = frozenset(get_args(GeminiFinishReason)) @staticmethod def get_finish_reason_mapping() -> dict[str, OpenAIChatCompletionFinishReason]: @@ -2232,22 +2215,23 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): grounding_metadata: Final[list[dict]] = [] url_context_metadata: Final[list[dict]] = [] - image_response: list[ImageURLListItem] | None = None safety_ratings: Final[list] = [] citation_metadata: Final[list] = [] - chat_completion_message: Final[ChatCompletionResponseMessage] = {"role": "assistant"} - chat_completion_logprobs: ChoiceLogprobs | None = None - tools: list[ChatCompletionToolCallChunk] | None = [] - functions: ChatCompletionToolCallFunctionChunk | None = None - thinking_blocks: list[ChatCompletionThinkingBlock] | None = None - reasoning_content: str | None = None - thought_signatures: Sequence[str] | None = None - server_side_tool_invocations: list[dict[str, object]] | None = None for idx, candidate in enumerate(_candidates): - if "content" not in candidate: + if "content" not in candidate and "finishReason" not in candidate: continue + image_response: list[ImageURLListItem] | None = None + chat_completion_message: ChatCompletionResponseMessage = {"role": "assistant"} + chat_completion_logprobs: ChoiceLogprobs | None = None + tools: list[ChatCompletionToolCallChunk] | None = None + functions: ChatCompletionToolCallFunctionChunk | None = None + thinking_blocks: list[ChatCompletionThinkingBlock] | None = None + reasoning_content: str | None = None + thought_signatures: Sequence[str] | None = None + server_side_tool_invocations: list[dict[str, object]] | None = None + # Extract metadata using helper function ( candidate_grounding_metadata, @@ -2261,7 +2245,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): safety_ratings.extend(candidate_safety_ratings) citation_metadata.extend(candidate_citation_metadata) - if "parts" in candidate["content"]: + if "content" in candidate and candidate["content"] and "parts" in candidate["content"]: ( content, reasoning_content, @@ -2368,14 +2352,18 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): ) model_response.choices.append(choice) elif isinstance(model_response, ModelResponse): + native_finish_reason = candidate.get("finishReason") choice = litellm.Choices( finish_reason=VertexGeminiConfig._check_finish_reason( - chat_completion_message, candidate.get("finishReason") + chat_completion_message, native_finish_reason ), index=candidate.get("index", idx), message=chat_completion_message, logprobs=chat_completion_logprobs, enhancements=None, + provider_specific_fields=( + {"native_finish_reason": native_finish_reason} if native_finish_reason is not None else None + ), ) model_response.choices.append(choice) @@ -3173,12 +3161,10 @@ class ModelResponseIterator: self.has_seen_tool_calls = True break - # _process_candidates skips candidates without a "content" part, so a - # content-less chunk leaves choices empty and the downstream streaming - # handler hits IndexError on choices[0]. This covers the final chunk - # (finishReason, no content) and mid-stream metadata-only chunks - # (grounding/web-search/thought, no content and no finishReason — seen - # with web_search + reasoning) by emitting an empty-delta choice. + # _process_candidates skips candidates with neither "content" nor + # "finishReason", so a metadata-only chunk (grounding/web-search/thought, + # seen with web_search + reasoning) leaves choices empty and the downstream + # streaming handler hits IndexError on choices[0]. Emit an empty-delta choice. if not model_response.choices and _candidates: from litellm.types.utils import Delta, StreamingChoices diff --git a/litellm/main.py b/litellm/main.py index 859e0df142c..ac8fa507728 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -1146,37 +1146,35 @@ def responses_api_bridge_check( return model_info, model -def _should_allow_input_examples(custom_llm_provider: str | None, model: str) -> bool: +_ANTHROPIC_ONLY_TOOL_KEYS: Final = frozenset({"input_examples", "eager_input_streaming"}) + + +def _is_claude_tool_target(custom_llm_provider: str | None, model: str) -> bool: if custom_llm_provider == "anthropic": return True - if custom_llm_provider == "azure_ai" or custom_llm_provider == "bedrock" or custom_llm_provider == "vertex_ai": - return "claude" in model.lower() + model_lower: Final = model.lower() + if custom_llm_provider == "bedrock": + return "claude" in model_lower or ("arn:" in model_lower and ":bedrock:" in model_lower) + if custom_llm_provider == "azure_ai" or custom_llm_provider == "vertex_ai": + return "claude" in model_lower return False -def _drop_input_examples_from_tool(tool: dict) -> dict: - tool_copy: Final = tool.copy() - tool_copy.pop("input_examples", None) - function = tool_copy.get("function") - if isinstance(function, dict): - function = function.copy() - function.pop("input_examples", None) - tool_copy["function"] = function - return tool_copy +def _without_anthropic_only_tool_keys(tool: dict) -> dict: + kept: Final = {key: value for key, value in tool.items() if key not in _ANTHROPIC_ONLY_TOOL_KEYS} + function: Final = tool.get("function") + if not isinstance(function, dict): + return kept + return { + **kept, + "function": {key: value for key, value in function.items() if key not in _ANTHROPIC_ONLY_TOOL_KEYS}, + } -def _drop_input_examples_from_tools( - tools: list[dict] | None, -) -> list[dict] | None: +def _drop_anthropic_only_tool_keys(tools: list[dict] | None) -> list[dict] | None: if tools is None: return None - cleaned_tools: Final[list[dict]] = [] - for tool in tools: - if isinstance(tool, dict): - cleaned_tools.append(_drop_input_examples_from_tool(tool)) - else: - cleaned_tools.append(tool) - return cleaned_tools + return [_without_anthropic_only_tool_keys(tool) if isinstance(tool, dict) else tool for tool in tools] class _ProxyAuthHeadersProvider(Protocol): @@ -5360,8 +5358,8 @@ def completion( api_base=api_base, ) - if not _should_allow_input_examples(custom_llm_provider=custom_llm_provider, model=model): - tools = _drop_input_examples_from_tools(tools=tools) + if not _is_claude_tool_target(custom_llm_provider=custom_llm_provider, model=model): + tools = _drop_anthropic_only_tool_keys(tools=tools) if provider_specific_header is not None: headers.update( diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 20fbb4ed956..954754a6f14 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -50869,7 +50869,7 @@ "vertex_ai/google/gemma-4-26b-a4b-it-maas": { "input_cost_per_token": 1.5e-07, "litellm_provider": "vertex_ai-openai_models", - "max_input_tokens": 256000, + "max_input_tokens": 262144, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 213cd88b6ce..584b1e05b89 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -33378,6 +33378,10 @@ "cache_control": { "$ref": "#/components/schemas/ChatCompletionCachedContent" }, + "eager_input_streaming": { + "title": "Eager Input Streaming", + "type": "boolean" + }, "function": { "$ref": "#/components/schemas/ChatCompletionToolParamFunctionChunk" }, @@ -33407,6 +33411,10 @@ "title": "Description", "type": "string" }, + "eager_input_streaming": { + "title": "Eager Input Streaming", + "type": "boolean" + }, "name": { "title": "Name", "type": "string" diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 174b9ceff93..6f769e6971a 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -2593,10 +2593,12 @@ class ProxyBaseLLMRequestProcessing: async def refresh_stream_headers() -> Mapping[str, str]: """`custom_headers` rebuilt for whichever deployment served the stream.""" - if not getattr(response, "fallback_headers_adopted", False): - return custom_headers return self._stream_response_headers( - hidden_params=get_hidden_params_dict(response), + hidden_params=( + get_hidden_params_dict(response) + if getattr(response, "fallback_headers_adopted", False) + else hidden_params + ), user_api_key_dict=user_api_key_dict, logging_obj=logging_obj, version=version, diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 1c078c0bfa2..a3f9924ee55 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -12,13 +12,36 @@ import sys import threading import time import traceback -from collections.abc import AsyncGenerator, Awaitable, Callable, Coroutine, Mapping, Sequence +from collections.abc import ( + AsyncGenerator, + AsyncIterable, + AsyncIterator, + Awaitable, + Callable, + Coroutine, + Mapping, + Sequence, +) from dataclasses import dataclass, field from datetime import date, datetime, timedelta, timezone from email.mime.multipart import MIMEMultipart from email.mime.text import MIMEText +from functools import partial from types import MappingProxyType -from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, Optional, Protocol, TypeVar, Union, cast, overload +from typing import ( + TYPE_CHECKING, + Any, + ClassVar, + Final, + Generic, + Literal, + Optional, + Protocol, + TypeVar, + Union, + cast, + overload, +) from typing_extensions import ReadOnly, TypedDict @@ -123,6 +146,7 @@ from litellm.proxy._types import ( UserAPIKeyAuth, ) from litellm.proxy.auth.route_checks import RouteChecks +from litellm.proxy.common_utils.callback_utils import add_guardrail_to_applied_guardrails_header from litellm.proxy.common_utils.config_sync_pubsub import publish_config_param_change from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.db.create_views import ( @@ -438,6 +462,36 @@ def _enrich_http_exception_with_guardrail_context(exc: BaseException, callback: detail.setdefault("guardrail_mode", event_hook) +def _record_raising_guardrail(request_data: Mapping[str, object], callback: object) -> None: + guardrail_name: Final[object] = getattr(callback, "guardrail_name", None) + if isinstance(request_data, dict) and isinstance(guardrail_name, str): + add_guardrail_to_applied_guardrails_header(request_data=request_data, guardrail_name=guardrail_name) + + +class _UpstreamStreamBoundary(Generic[_T]): + __slots__ = ("_upstream", "failure") + + def __init__(self, upstream: AsyncIterable[_T]) -> None: + self._upstream: Final = upstream.__aiter__() + self.failure: BaseException | None = None + + def __aiter__(self) -> "_UpstreamStreamBoundary[_T]": + return self + + async def __anext__(self) -> _T: + try: + return await self._upstream.__anext__() + except StopAsyncIteration: + raise + except Exception as e: + self.failure = e + raise + + +class _StreamIteratorHook(Protocol[_T]): + def __call__(self, *, response: AsyncIterator[_T]) -> AsyncGenerator[_T, None]: ... + + def _is_client_error_exception(exc: Exception) -> bool: if isinstance(exc, HTTPException): return exc.status_code < 500 @@ -1816,13 +1870,19 @@ class ProxyLogging: ) if expected_if_unmutated is not None: callback.mark_pre_call_hook_ran(expected_if_unmutated) - result: Final = await self._process_guardrail_callback( - callback=callback, - data=input_data, - user_api_key_dict=user_api_key_dict, - call_type=call_type, - event_type=GuardrailEventHooks.pre_call, - ) + try: + result: Final = await self._process_guardrail_callback( + callback=callback, + data=input_data, + user_api_key_dict=user_api_key_dict, + call_type=call_type, + event_type=GuardrailEventHooks.pre_call, + ) + except SensitiveDataRouteException: + raise + except Exception: + _record_raising_guardrail(data, callback) + raise if ( scans_raw_request and expected_if_unmutated is not None @@ -2045,13 +2105,18 @@ class ProxyLogging: _merge_pipeline_metadata_writes(data, result.modified_data) if result.terminal_action == "block": + blocking_step: Final = result.step_results[-1] if result.step_results else None + callback: Final = ( + PipelineExecutor.find_guardrail_callback(blocking_step.guardrail_name) + if blocking_step is not None + else None + ) + if callback is not None: + _record_raising_guardrail(data, callback) original_exception: Final = result.original_exception if original_exception is not None and not _exception_changes_request_flow(original_exception): - blocking_step: Final = result.step_results[-1] if result.step_results else None - if blocking_step is not None: - callback: Final = PipelineExecutor.find_guardrail_callback(blocking_step.guardrail_name) - if callback is not None: - _enrich_http_exception_with_guardrail_context(original_exception, callback) + if callback is not None: + _enrich_http_exception_with_guardrail_context(original_exception, callback) raise original_exception step_results_serializable: Final = [ @@ -2317,8 +2382,10 @@ class ProxyLogging: if data is not None: self._process_guardrail_metadata(data) return data - except Exception as e: - raise e + except Exception: + if data is not None: + self._process_guardrail_metadata(data) + raise async def _run_parallel_pre_call_guardrails( self, @@ -2376,6 +2443,8 @@ class ProxyLogging: # live kwargs. if callback.scan_raw_request and not isinstance(result, BaseException) and result is not None: callback.mark_pre_call_hook_ran(data) + if isinstance(result, BaseException) and not isinstance(result, SensitiveDataRouteException): + _record_raising_guardrail(data, callback) raised: Final = tuple(result for result in results if isinstance(result, BaseException)) blocking: Final = next((exc for exc in raised if not _exception_changes_request_flow(exc)), None) if blocking is not None: @@ -2454,7 +2523,12 @@ class ProxyLogging: break @staticmethod - async def _run_guardrail_with_metrics(callback: object, coro: Awaitable[_T], hook_type: str) -> _T: + async def _run_guardrail_with_metrics( + callback: object, + coro: Awaitable[_T], + hook_type: str, + request_data: Mapping[str, object], + ) -> _T: """ Await `coro`, recording its latency and status to the `litellm_guardrail_latency_seconds` metric under `hook_type`, and @@ -2474,6 +2548,7 @@ class ProxyLogging: status = "error" error_type = type(e).__name__ _enrich_http_exception_with_guardrail_context(e, callback) + _record_raising_guardrail(request_data, callback) raise finally: ProxyLogging._emit_guardrail_metrics( @@ -2486,21 +2561,19 @@ class ProxyLogging: @staticmethod async def _wrap_streaming_iterator_with_enrichment( - callback: object, gen: AsyncGenerator[_T, None] + callback: object, + response: AsyncIterable[_T], + hook: _StreamIteratorHook[_T], + request_data: Mapping[str, object], ) -> AsyncGenerator[_T, None]: - """ - Yield from `gen`; if iteration raises an HTTPException with dict detail, - enrich the detail with the originating callback's `guardrail_name` and - `guardrail_mode` before re-raising. Used to wrap each layer of the - async_post_call_streaming_iterator_hook chain so the enrichment is - attributed to the callback that produced the chunk pipeline at that - point in the chain. - """ + upstream: Final = _UpstreamStreamBoundary(response) try: - async for chunk in gen: + async for chunk in hook(response=upstream): yield chunk except Exception as e: - _enrich_http_exception_with_guardrail_context(e, callback) + if e is not upstream.failure: + _enrich_http_exception_with_guardrail_context(e, callback) + _record_raising_guardrail(request_data, callback) raise # Cache for callback-capability detection. Keyed on a signature of @@ -2735,6 +2808,7 @@ class ProxyLogging: call_type=call_type, ), "during_call", + request_data=data, ) return await self._run_guardrail_with_metrics( @@ -2745,6 +2819,7 @@ class ProxyLogging: call_type=call_type, ), "during_call", + request_data=data, ) async def failed_tracking_alert( @@ -3263,6 +3338,7 @@ class ProxyLogging: response=response, ), "post_call", + request_data=data, ) else: guardrail_response = await self._run_guardrail_with_metrics( @@ -3273,6 +3349,7 @@ class ProxyLogging: response=response, ), "post_call", + request_data=data, ) if guardrail_response is not None: @@ -3336,6 +3413,7 @@ class ProxyLogging: response=response, ), "post_call", + request_data=data, ) else: await self._run_guardrail_with_metrics( @@ -3346,6 +3424,7 @@ class ProxyLogging: response=response, ), "post_call", + request_data=data, ) results: Final = await asyncio.gather( @@ -3409,6 +3488,7 @@ class ProxyLogging: request_data=request_data, ), "post_mcp_call", + request_data=request_data, ) return response @@ -3650,27 +3730,27 @@ class ProxyLogging: ) else kind ) - if effective_kind == "override": - current_response = self._wrap_streaming_iterator_with_enrichment( - resolved_callback, - resolved_callback.async_post_call_streaming_iterator_hook( - user_api_key_dict=user_api_key_dict, - response=current_response, - request_data=request_data, - ), + hook: _StreamIteratorHook[object] = ( + partial( + resolved_callback.async_post_call_streaming_iterator_hook, + user_api_key_dict=user_api_key_dict, + request_data=request_data, ) - else: - # kind == "apply_guardrail": route through unified_guardrail - current_response = self._wrap_streaming_iterator_with_enrichment( - resolved_callback, - unified_guardrail.async_post_call_streaming_iterator_hook( - user_api_key_dict=user_api_key_dict, - request_data=request_data, - response=current_response, - guardrail_to_apply=resolved_callback, - buffer_until_moderated_default=(kind == "override"), - ), + if effective_kind == "override" + else partial( + unified_guardrail.async_post_call_streaming_iterator_hook, + user_api_key_dict=user_api_key_dict, + request_data=request_data, + guardrail_to_apply=resolved_callback, + buffer_until_moderated_default=(kind == "override"), ) + ) + current_response = self._wrap_streaming_iterator_with_enrichment( + resolved_callback, + current_response, + hook, + request_data=request_data, + ) pipeline_translation: Final = ( resolve_endpoint_translation(user_api_key_dict, None) if post_call_pipelines else None diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index 01fb6cb483d..044596676dd 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -61,6 +61,7 @@ from litellm.types.llms.openai import ( ChatCompletionToolParamFunctionChunk, ChatCompletionUserMessage, GenericChatCompletionMessage, + IncompleteDetails, InputTokensDetails, OpenAIChatCompletionTextObject, OpenAIMcpServerTool, @@ -111,6 +112,9 @@ ResponseTools: TypeAlias = Sequence[Mapping[str, object]] | None ChatToolParam: TypeAlias = ChatCompletionToolParam | OpenAIMcpServerTool NAMESPACE_DESCRIPTION_SEPARATOR: Final = "\n\n" NAMESPACE_MEMBER_TYPES_WITH_CHAT_TOOLS: Final = frozenset({"function", "custom"}) +_INCOMPLETE_REASON_BY_FINISH_REASON: Final[Mapping[str, Literal["max_output_tokens", "content_filter"]]] = ( + MappingProxyType({"length": "max_output_tokens", "content_filter": "content_filter", "refusal": "content_filter"}) +) @dataclass(frozen=True, slots=True) @@ -2020,6 +2024,8 @@ class LiteLLMCompletionResponsesConfig: chat_completion_tool["allowed_callers"] = tool.get("allowed_callers") if tool.get("input_examples"): chat_completion_tool["input_examples"] = tool.get("input_examples") + if tool.get("eager_input_streaming") is not None: + chat_completion_tool["eager_input_streaming"] = tool.get("eager_input_streaming") return ResponsesToolChatForm( chat_tools=(cast(ChatCompletionToolParam, chat_completion_tool),), web_search_options=None ) @@ -2096,6 +2102,8 @@ class LiteLLMCompletionResponsesConfig: responses_tool["allowed_callers"] = tool.get("allowed_callers") if tool.get("input_examples") is not None: responses_tool["input_examples"] = tool.get("input_examples") + if tool.get("eager_input_streaming") is not None: + responses_tool["eager_input_streaming"] = tool.get("eager_input_streaming") result.append(responses_tool) else: # mcp or other: pass through unchanged @@ -2295,6 +2303,18 @@ class LiteLLMCompletionResponsesConfig: # Default to completed for unknown finish reasons return "completed" + @staticmethod + def _incomplete_details_for_finish_reason( + finish_reason: str | None, + existing: IncompleteDetails | None, + ) -> IncompleteDetails | None: + if existing is not None: + return existing + if finish_reason is None: + return None + reason: Final = _INCOMPLETE_REASON_BY_FINISH_REASON.get(finish_reason) + return IncompleteDetails(reason=reason) if reason is not None else None + @staticmethod def _tool_call_id_from_responses_item(item_id: str | None, call_id: str | None) -> str: """Bedrock Mantle returns a non-unique, index-based ``call_id`` (``call_0``, @@ -2411,13 +2431,18 @@ class LiteLLMCompletionResponsesConfig: if choices and len(choices) > 0: finish_reason = choices[0].finish_reason + incomplete_details: Final = LiteLLMCompletionResponsesConfig._incomplete_details_for_finish_reason( + finish_reason=finish_reason, + existing=getattr(chat_completion_response, "incomplete_details", None), + ) + responses_api_response: Final[ResponsesAPIResponse] = ResponsesAPIResponse( id=chat_completion_response.id, created_at=chat_completion_response.created, model=chat_completion_response.model, object="response", error=getattr(chat_completion_response, "error", None), - incomplete_details=getattr(chat_completion_response, "incomplete_details", None), + incomplete_details=incomplete_details, instructions=getattr(chat_completion_response, "instructions", None), metadata=getattr(chat_completion_response, "metadata", {}), output=LiteLLMCompletionResponsesConfig._transform_chat_completion_choices_to_responses_output( diff --git a/litellm/responses/main.py b/litellm/responses/main.py index 6dc34bb93ef..a5912bb42b1 100644 --- a/litellm/responses/main.py +++ b/litellm/responses/main.py @@ -8,7 +8,7 @@ from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, NoReturn, Optional, TypeAlias, cast import httpx -from pydantic import BaseModel +from pydantic import BaseModel, TypeAdapter from typing_extensions import assert_never import litellm @@ -53,6 +53,7 @@ from litellm.litellm_core_utils.get_litellm_params import get_litellm_params from litellm.llms.openai.data_residency import infer_openai_data_residency from litellm.secret_managers.main import get_secret_str from litellm.types.responses.main import * +from litellm.types.responses.streaming_websocket import ResponsesWebSocketRequestDefaults from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import all_litellm_params from litellm.utils import ( @@ -2261,6 +2262,31 @@ def _build_litellm_metadata_for_ws(kwargs: dict) -> dict: return metadata +_JSON_OBJECT_ADAPTER: Final = TypeAdapter(dict[str, object] | None) + + +def _deployment_reasoning_default(kwargs: Mapping[str, object]) -> Reasoning | dict[str, object] | None: + if kwargs.get("reasoning") is not None: + return None + reasoning_effort: Final = kwargs.get("reasoning_effort") + if isinstance(reasoning_effort, str): + return LiteLLMResponsesTransformationHandler()._map_reasoning_effort(reasoning_effort) + return _JSON_OBJECT_ADAPTER.validate_python(reasoning_effort) if isinstance(reasoning_effort, Mapping) else None + + +def _build_responses_websocket_request_defaults(kwargs: Mapping[str, object]) -> ResponsesWebSocketRequestDefaults: + default_reasoning: Final = _deployment_reasoning_default(kwargs) + candidate_params: Final[dict[str, object]] = { + **kwargs, + **({"reasoning": default_reasoning} if default_reasoning is not None else {}), + } + fill_missing: Final = ResponsesAPIRequestUtils.get_requested_response_api_optional_param(candidate_params) + return ResponsesWebSocketRequestDefaults( + fill_missing=MappingProxyType(dict(fill_missing)), + overrides=MappingProxyType(_JSON_OBJECT_ADAPTER.validate_python(kwargs.get("extra_body")) or {}), + ) + + @client async def _aresponses_websocket( model: str, @@ -2352,5 +2378,6 @@ async def _aresponses_websocket( user_api_key_dict=kwargs.get("user_api_key_dict"), litellm_metadata=_build_litellm_metadata_for_ws(kwargs), custom_llm_provider=_custom_llm_provider, + request_defaults=_build_responses_websocket_request_defaults(kwargs), **remaining_kwargs, ) diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 61c93063826..e1ec1f00f0c 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -54,6 +54,7 @@ if TYPE_CHECKING: PresidioGuardrailCallback, ResponsesBackendWebSocket, ResponsesClientWebSocket, + ResponsesWebSocketRequestDefaults, ) from litellm.types.router import LiteLLM_Params @@ -1717,6 +1718,7 @@ class ResponsesWebSocketStreaming: output_guardrail_callbacks: list[PresidioGuardrailCallback] | None = None, quota_callbacks: Sequence[ProjectQuotaCallback] | None = None, authorized_model: str | None = None, + request_defaults: ResponsesWebSocketRequestDefaults | None = None, ): self.websocket = websocket self.backend_ws = backend_ws @@ -1732,6 +1734,7 @@ class ResponsesWebSocketStreaming: # Model name authorized at connection time; enforced on every # response.create frame to prevent deployment-substitution attacks. self.authorized_model: str | None = authorized_model + self.request_defaults: ResponsesWebSocketRequestDefaults | None = request_defaults def _should_store_event(self, event_obj: _MutableJsonObject) -> bool: return event_obj.get("type") in RESPONSES_WS_LOGGED_EVENT_TYPES @@ -1874,12 +1877,23 @@ class ResponsesWebSocketStreaming: modified = True return modified + def _with_request_defaults(self, msg_obj: dict[str, object]) -> dict[str, object]: + if self.request_defaults is None: + return msg_obj + nested: Final = msg_obj.get("response") + if _is_json_object(nested): + return {**msg_obj, "response": self.request_defaults.merged_into(nested)} + return {**self.request_defaults.merged_into(msg_obj), "type": msg_obj["type"]} + async def _mask_response_create(self, message: str) -> str: """ - Enforce the authorized model and apply Presidio PII masking to a - ``response.create`` message before it is forwarded to the upstream - provider. + Merge deployment defaults, enforce the authorized model, and apply + Presidio PII masking to a ``response.create`` message before it is + forwarded to the upstream provider. + - Fills the deployment's ``litellm_params`` request defaults into the + frame the way the HTTP ``/v1/responses`` path does: client-set keys + win, ``extra_body`` entries override. - Overwrites any ``model`` field with the connection-authorized model to prevent deployment-substitution attacks (always applied). - Walks the ``input`` and ``instructions`` fields, calls ``check_pii`` @@ -1889,23 +1903,26 @@ class ResponsesWebSocketStreaming: Non-``response.create`` messages are returned unchanged. """ try: - msg_obj: Final = _load_json_object(message) + parsed: Final = _load_json_object(message) except (json.JSONDecodeError, TypeError): return message - if msg_obj.get("type") != "response.create": + if parsed.get("type") != "response.create": return message + msg_obj: Final = self._with_request_defaults(parsed) + defaults_applied: Final = msg_obj != parsed + # Always enforce the authorized model, even when PII masking is off. model_modified: Final = self._enforce_authorized_model(msg_obj) if not self.guardrail_callbacks: - return json.dumps(msg_obj) if model_modified else message + return json.dumps(msg_obj) if model_modified or defaults_applied else message if "metadata" not in self.request_data: self.request_data["metadata"] = {} - modified = model_modified + modified = model_modified or defaults_applied guardrail_cbs: Final[tuple[PresidioGuardrailCallback, ...]] = tuple(self.guardrail_callbacks) for cb in guardrail_cbs: presidio_config = cb.get_presidio_settings_from_request_data(self.request_data) diff --git a/litellm/rust_bridge/_native.pyi b/litellm/rust_bridge/_native.pyi index 488e278cca7..9f959c056de 100644 --- a/litellm/rust_bridge/_native.pyi +++ b/litellm/rust_bridge/_native.pyi @@ -1,9 +1,11 @@ from asyncio import Future -from collections.abc import Coroutine, Mapping, Sequence +from collections.abc import AsyncIterator, Coroutine, Iterator, Mapping, Sequence from typing import Never, final from litellm.llms.base_llm.ocr.transformation import OCRResponse +from litellm.rust_bridge.messages.entrypoints import LiteLLMMessagesRequest from litellm.rust_bridge.ocr.entrypoints import LiteLLMOcrRequest +from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicMessagesResponse class RustBridgeDeclined(Exception): ... class RustUpstreamError(Exception): ... @@ -39,23 +41,15 @@ def atranscription( timeout_seconds: float | None = None, ) -> Future[dict[str, object]]: ... def messages( - model: str, - body: Mapping[str, object], - api_key: str | None = None, - api_base: str | None = None, - custom_llm_provider: str | None = None, - extra_headers: Mapping[str, object] | None = None, - timeout_seconds: float | None = None, -) -> dict[str, object]: ... + request: LiteLLMMessagesRequest, + args: tuple[object, ...], + kwargs: dict[str, object], +) -> AnthropicMessagesResponse | Iterator[bytes]: ... def amessages( - model: str, - body: Mapping[str, object], - api_key: str | None = None, - api_base: str | None = None, - custom_llm_provider: str | None = None, - extra_headers: Mapping[str, object] | None = None, - timeout_seconds: float | None = None, -) -> Future[dict[str, object]]: ... + request: LiteLLMMessagesRequest, + args: tuple[object, ...], + kwargs: dict[str, object], +) -> Coroutine[object, object, AnthropicMessagesResponse | AsyncIterator[bytes]]: ... def chat_completions_decline( model: str, messages: Sequence[object], diff --git a/litellm/rust_bridge/catalog.py b/litellm/rust_bridge/catalog.py index 9efbbfa2e9e..d843a874fe3 100644 --- a/litellm/rust_bridge/catalog.py +++ b/litellm/rust_bridge/catalog.py @@ -59,6 +59,7 @@ Rules: TypeAlias = tuple[Rule, ...] RULES: Final[Rules] = ( Rule(Route.OCR, Rollout.RUST_OPT_OUT), + Rule(Route.MESSAGES, Rollout.RUST_OPT_IN), Rule(Route.TRANSCRIPTION, Rollout.RUST_REQUIRED, providers=frozenset({"bedrock"})), ) diff --git a/litellm/rust_bridge/failures.py b/litellm/rust_bridge/failures.py index b714341fe43..80805b7ff69 100644 --- a/litellm/rust_bridge/failures.py +++ b/litellm/rust_bridge/failures.py @@ -5,8 +5,37 @@ from __future__ import annotations from collections.abc import Mapping from typing import Final, Protocol, cast # noqa: TID251 # adapts the public exception mapper +import httpx +import openai +from pydantic import TypeAdapter, ValidationError + import litellm +_UPSTREAM_ARGS: Final = TypeAdapter(tuple[int, str]) +_UPSTREAM_HEADERS: Final = TypeAdapter(list[tuple[str, str]]) + + +class UpstreamFailure(Exception): + def __init__(self, response: httpx.Response, cause: Exception) -> None: + super().__init__(str(cause)) + self.message: Final = str(cause) + self.response: Final = response + self.status_code: Final = response.status_code + self.__cause__ = cause + + +def _upstream_failure(error: Exception, api_base: str | None) -> Exception: + try: + status, body = _UPSTREAM_ARGS.validate_python(error.args) + headers: Final = _UPSTREAM_HEADERS.validate_python(getattr(error, "headers", None)) + except ValidationError: + return error + http_request: Final = httpx.Request("POST", api_base or "https://docs.litellm.ai/docs") + return UpstreamFailure( + httpx.Response(status, content=body.encode(), headers=headers, request=http_request), + error, + ) + class ExceptionMapper(Protocol): def __call__( @@ -35,3 +64,17 @@ def map_failure(error: Exception, model: str, request_provider: str, kwargs: Map except Exception as public_error: public_error.__context__ = error return public_error + + +def map_native_failure( + error: Exception, model: str, request_provider: str, kwargs: Mapping[str, object], api_base: str | None = None +) -> Exception: + """`map_failure`, reading a native `(status, body)` provider failure as the HTTP response it was.""" + original: Final = _upstream_failure(error, api_base) + public_error: Final = map_failure(original, model, request_provider, kwargs) + if isinstance(original, UpstreamFailure) and public_error.__context__ is original: + public_error.__context__ = error + if isinstance(public_error, openai.APIStatusError): + public_error.response = original.response + public_error.status_code = original.status_code + return public_error diff --git a/litellm/rust_bridge/legacy_callbacks.py b/litellm/rust_bridge/legacy_callbacks.py index e05d9368fa8..30aa1d97bfc 100644 --- a/litellm/rust_bridge/legacy_callbacks.py +++ b/litellm/rust_bridge/legacy_callbacks.py @@ -6,24 +6,23 @@ registries it fans out to. It expires with that contract. from __future__ import annotations +import asyncio +import contextvars import datetime -import os +import traceback import uuid -from collections.abc import Mapping +from collections.abc import Awaitable, Coroutine, Mapping from dataclasses import dataclass from typing import ( TYPE_CHECKING, Final, - Literal, Protocol, - TypeAlias, cast, # noqa: TID251 # bounded compatibility calls into legacy Python integrations ) -from typing_extensions import assert_never - if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.types.utils import CredentialItem class MetadataUpdater(Protocol): @@ -42,7 +41,6 @@ class MetadataUpdater(Protocol): class CallSetup: logger: Logging kwargs: dict[str, object] - bridge_owned: bool def setup( @@ -61,19 +59,23 @@ def setup( } supplied: Final = arguments.get("litellm_logging_obj") if isinstance(supplied, Logging): - return CallSetup(supplied, arguments, bridge_owned=False) + return CallSetup(supplied, arguments) logger, prepared = function_setup(call_type, Rules(), start_time, *args, is_async_call=asynchronous, **arguments) - return CallSetup(logger, prepared, bridge_owned=True) + return CallSetup(logger, prepared) def check_limits(kwargs: Mapping[str, object]) -> None: - import litellm + from litellm import ( + BudgetExceededError, + _current_cost, # pyright: ignore[reportPrivateUsage] # shared SDK budget counter has no public accessor + max_budget, + num_retries_per_request, + ) from litellm.litellm_core_utils.core_helpers import max_retries_per_request_hit - current_cost: Final = litellm._current_cost # pyright: ignore[reportPrivateUsage] # shared SDK budget counter has no public accessor - if litellm.max_budget and current_cost > litellm.max_budget: - raise litellm.BudgetExceededError(current_cost=current_cost, max_budget=litellm.max_budget) - if max_retries_per_request_hit(kwargs, litellm.num_retries_per_request): + if max_budget and _current_cost > max_budget: + raise BudgetExceededError(current_cost=_current_cost, max_budget=max_budget) + if max_retries_per_request_hit(kwargs, num_retries_per_request): raise RuntimeError("Max retries per request hit!") @@ -93,87 +95,304 @@ def finalize( update(response, logger, model if isinstance(model, str) else None, kwargs, start_time, end_time) -def deployment_callbacks_needed() -> bool: - import litellm - from litellm.integrations.custom_logger import CustomLogger +class LoggingSurface(Protocol): + def update_from_kwargs( + self, + kwargs: dict[str, object], + litellm_params: dict[str, object] | None = None, + optional_params: dict[str, object] | None = None, + model: str | None = None, + user: str | None = None, + **additional_params: object, + ) -> None: ... - return any(isinstance(callback, CustomLogger) for callback in litellm.callbacks) + def pre_call( + self, input: object, api_key: object, model: object = None, additional_args: dict[str, object] = ... + ) -> object: ... + + def post_call( + self, + original_response: object, + input: object = None, + api_key: object = None, + additional_args: dict[str, object] = ..., + ) -> object: ... + + def handle_sync_success_callbacks_for_async_calls( + self, result: object, start_time: datetime.datetime, end_time: datetime.datetime, cache_hit: object = None + ) -> None: ... + + def failure_handler( + self, + exception: Exception, + traceback_exception: str, + start_time: datetime.datetime | None = None, + end_time: datetime.datetime | None = None, + ) -> None: ... + + def async_failure_handler( + self, + exception: Exception, + traceback_exception: str, + start_time: datetime.datetime | None = None, + end_time: datetime.datetime | None = None, + ) -> Coroutine[object, object, None]: ... + + def success_handler( + self, + result: object = None, + start_time: datetime.datetime | None = None, + end_time: datetime.datetime | None = None, + cache_hit: bool | None = None, + **kwargs: object, + ) -> None: ... + + def async_success_handler( + self, + result: object = None, + start_time: datetime.datetime | None = None, + end_time: datetime.datetime | None = None, + cache_hit: bool | None = None, + **kwargs: object, + ) -> Coroutine[object, object, None]: ... -Phase: TypeAlias = Literal[ - "input", "sync_success", "sync_success_async", "async_success", "sync_failure", "async_failure", "payload" -] +if TYPE_CHECKING: + _LOGGING_CONFORMS: type[LoggingSurface] = Logging -def callbacks_needed(logger: Logging, phase: Phase) -> bool: - import litellm - from litellm._logging import ( - _is_debugging_on, # pyright: ignore[reportPrivateUsage] # use the same debug gate as Logging +class LoggingWorker(Protocol): + def ensure_initialized_and_enqueue(self, async_coroutine: Coroutine[object, object, None]) -> None: ... + + +class StreamingLogBuilder(Protocol): + def __call__( + self, + *, + litellm_logging_obj: Logging, + passthrough_success_handler_obj: object, + url_route: str, + request_body: dict[str, object], + endpoint_type: object, + start_time: datetime.datetime, + raw_bytes: list[bytes], + end_time: datetime.datetime, + ) -> Coroutine[object, object, None]: ... + + +class DeploymentHook(Protocol): + def __call__(self, kwargs: dict[str, object], call_type: str) -> Awaitable[object]: ... + + +class DeploymentSuccessHook(Protocol): + def __call__(self, request_data: dict[str, object], response: object, call_type: object) -> Awaitable[object]: ... + + +class DeploymentFailureHook(Protocol): + def __call__(self, request_data: Mapping[str, object], exception: Exception, call_type: str) -> Awaitable[None]: ... + + +def update_logging( + logger: LoggingSurface, + kwargs: dict[str, object], + model: str, + optional_params: dict[str, object], + litellm_params: dict[str, object], + custom_llm_provider: str, +) -> None: + logger.update_from_kwargs( + kwargs=kwargs, + model=model, + optional_params=optional_params, + litellm_params=litellm_params, + custom_llm_provider=custom_llm_provider, ) - if ( - _is_debugging_on() - or getattr(logger, "litellm_request_debug", False) - or os.getenv("LITELLM_PRINT_STANDARD_LOGGING_PAYLOAD") - ): - return True - input_needed: Final = bool( - litellm.input_callback - or litellm._async_input_callback # pyright: ignore[reportPrivateUsage] # live async registries have no public accessor - or logger.dynamic_input_callbacks - or callable(getattr(logger, "logger_fn", None)) - or logger.log_raw_request_response - or litellm.log_raw_request_response + +def pre_call(logger: LoggingSurface, input: str, api_key: str | None, additional_args: dict[str, object]) -> None: + logger.pre_call(input=input, api_key=api_key, additional_args=additional_args) + + +def post_call( + logger: LoggingSurface, original_response: str, api_key: str | None, additional_args: dict[str, object] +) -> None: + logger.post_call(original_response=original_response, api_key=api_key, additional_args=additional_args) + + +def defers_async_logging(logger: LoggingSurface) -> bool: + return bool(getattr(logger, "_defer_async_logging", False)) + + +def defer_success(logger: LoggingSurface, pending: object) -> None: + setattr(logger, "_native_pending_logging", pending) + + +def sync_success_for_async_call( + logger: LoggingSurface, response: object, start: datetime.datetime, end: datetime.datetime +) -> None: + logger.handle_sync_success_callbacks_for_async_calls(result=response, start_time=start, end_time=end) + + +def failure_handler( + logger: LoggingSurface, error: Exception, start: datetime.datetime, end: datetime.datetime, asynchronous: bool +) -> Coroutine[object, object, None] | None: + trace: Final = "".join(traceback.format_exception(error)) + if asynchronous: + return logger.async_failure_handler(error, trace, start, end) + logger.failure_handler(error, trace, start, end) + return None + + +def submit_success(logger: LoggingSurface, response: object, start: datetime.datetime, end: datetime.datetime) -> None: + from litellm.litellm_core_utils.litellm_logging import executor + + executor.submit(contextvars.copy_context().run, logger.success_handler, response, start, end) + + +def async_success_handler( + logger: LoggingSurface, response: object, start: datetime.datetime, end: datetime.datetime +) -> Coroutine[object, object, None]: + return logger.async_success_handler(response, start, end) + + +def enqueue_logging(coroutine: Coroutine[object, object, None]) -> None: + from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER + + worker: Final = cast( # cast-ok: bounded adapter for the untyped logging worker + LoggingWorker, GLOBAL_LOGGING_WORKER ) - match phase: - case "input": - return input_needed - case "sync_success": - return bool(litellm.success_callback or logger.dynamic_success_callbacks) - case "sync_success_async": - return bool( - (litellm.success_callback or logger.dynamic_success_callbacks) - and logger._should_run_sync_callbacks_for_async_calls() # pyright: ignore[reportPrivateUsage] # preserve async call filtering of sync callbacks - ) - case "async_success": - return bool(litellm._async_success_callback or logger.dynamic_async_success_callbacks) # pyright: ignore[reportPrivateUsage] # live async registries have no public accessor - case "sync_failure": - return bool(litellm.failure_callback or logger.dynamic_failure_callbacks) - case "async_failure": - return bool(litellm._async_failure_callback or logger.dynamic_async_failure_callbacks) # pyright: ignore[reportPrivateUsage] # live async registries have no public accessor - case "payload": - return bool( - input_needed - or litellm.success_callback - or litellm.failure_callback - or litellm._async_success_callback # pyright: ignore[reportPrivateUsage] # live async registries have no public accessor - or litellm._async_failure_callback # pyright: ignore[reportPrivateUsage] # live async registries have no public accessor - or logger.dynamic_success_callbacks - or logger.dynamic_async_success_callbacks - or logger.dynamic_failure_callbacks - or logger.dynamic_async_failure_callbacks - ) - case _: - assert_never(phase) + contextvars.copy_context().run(worker.ensure_initialized_and_enqueue, coroutine) -def success_bookkeeping( - logger: Logging, response: object, start: datetime.datetime, end: datetime.datetime, asynchronous: bool +def restore_context(logger: LoggingSurface) -> None: + from litellm.utils import ( + _restore_correlation_context_if_supported, # pyright: ignore[reportPrivateUsage] # the @client wrapper restores the same correlation context + ) + + _restore_correlation_context_if_supported(logger) + + +def custom_pricing_fields() -> tuple[str, ...]: + from litellm.types.utils import CustomPricingLiteLLMParams + + return tuple(CustomPricingLiteLLMParams.model_fields) + + +def is_internal_call() -> bool: + from litellm._internal_context import is_internal_call as internal + + return internal.get() + + +def credential_list() -> list[CredentialItem]: + from litellm import credential_list as credentials + + return credentials + + +def warn_unknown_credential(name: str, loaded: int) -> None: + from litellm._logging import verbose_logger + + verbose_logger.warning( + "litellm_credential_name=%s matched none of the %d loaded credentials; the request runs without it", + name, + loaded, + ) + + +def before_deployment_call(kwargs: dict[str, object], call_type: str) -> Awaitable[object]: + from litellm import utils + + hook: Final = cast( # cast-ok: bounded adapter for the untyped deployment hook + DeploymentHook, utils.async_pre_call_deployment_hook + ) + return hook(kwargs, call_type) + + +def after_deployment_success(kwargs: dict[str, object], response: object, call_type: str) -> Awaitable[object]: + from litellm import utils + from litellm.types.utils import CallTypes + + hook: Final = cast( # cast-ok: bounded adapter for the untyped deployment hook + DeploymentSuccessHook, utils.async_post_call_success_deployment_hook + ) + return hook(kwargs, response, CallTypes(call_type)) + + +def after_deployment_failure(kwargs: dict[str, object], error: Exception, call_type: str) -> Awaitable[None]: + from litellm import utils + + hook: Final = cast( # cast-ok: bounded adapter for the untyped deployment hook + DeploymentFailureHook, utils.async_post_call_failure_deployment_hook + ) + return hook(kwargs, error, call_type) + + +def stream_opened(logger: Logging) -> None: + logger.stream = True + logger.model_call_details["stream"] = True + + +def stream_success( + logger: Logging, + url_route: str, + endpoint_type: str, + request_body: dict[str, object], + chunks: list[bytes], + start: datetime.datetime, + end: datetime.datetime, + first_chunk: datetime.datetime | None, ) -> None: - phase: Final = "async_success" if asynchronous else "sync_success" - if logger.should_run_logging(phase): - logger._success_handler_helper_fn( # pyright: ignore[reportPrivateUsage] # retain success bookkeeping without constructing a callback payload - result=response, start_time=start, end_time=end, build_logging_payload=False - ) - logger.has_run_logging(phase) + from litellm.llms.anthropic.experimental_pass_through.messages.streaming_iterator import ( + GLOBAL_PASS_THROUGH_SUCCESS_HANDLER_OBJ, + ) + from litellm.proxy.pass_through_endpoints.streaming_handler import PassThroughStreamingHandler + from litellm.types.passthrough_endpoints.pass_through_endpoints import EndpointType + + if first_chunk is not None: + logger.completion_start_time = first_chunk + logger.model_call_details["completion_start_time"] = first_chunk + build: Final = cast( # cast-ok: bounded adapter for the untyped pass-through logging builder + StreamingLogBuilder, + PassThroughStreamingHandler._route_streaming_logging_to_handler, # pyright: ignore[reportPrivateUsage] # the Messages stream iterator bills through the same builder + ) + coroutine: Final = build( + litellm_logging_obj=logger, + passthrough_success_handler_obj=GLOBAL_PASS_THROUGH_SUCCESS_HANDLER_OBJ, + url_route=url_route, + request_body=request_body, + endpoint_type=EndpointType(endpoint_type), + start_time=start, + raw_bytes=chunks, + end_time=end, + ) + if getattr(logger, "_on_deferred_stream_complete", None) is not None: + logger._deferred_stream_complete_args = (coroutine,) # pyright: ignore[reportAttributeAccessIssue] # the proxy's deferred stream release reads this slot + return + try: + asyncio.get_running_loop() + except RuntimeError: + from litellm.litellm_core_utils.litellm_logging import executor + + executor.submit(contextvars.copy_context().run, asyncio.run, coroutine) + return + enqueue_logging(coroutine) -def failure_bookkeeping( - logger: Logging, error: BaseException, start: datetime.datetime, end: datetime.datetime, asynchronous: bool -) -> None: - phase: Final = "async_failure" if asynchronous else "sync_failure" - if logger.should_run_logging(phase): - logger._failure_handler_helper_fn( # pyright: ignore[reportPrivateUsage] # retain failure accounting without formatting an unused traceback or payload - error, "", start, end, build_logging_payload=False - ) - logger.has_run_logging(phase) +def stream_failure( + logger: Logging, + endpoint_type: str, + request_body: dict[str, object], + chunks: list[bytes], + error: Exception, +) -> Coroutine[object, object, None]: + from litellm.proxy.pass_through_endpoints.streaming_handler import PassThroughStreamingHandler + from litellm.types.passthrough_endpoints.pass_through_endpoints import EndpointType + + return PassThroughStreamingHandler.schedule_stream_failure_logging( + litellm_logging_obj=logger, + endpoint_type=EndpointType(endpoint_type), + request_body=request_body, + raw_bytes=chunks, + exception=error, + ) diff --git a/litellm/rust_bridge/lifecycle.py b/litellm/rust_bridge/lifecycle.py index d903021b6f3..4096d386964 100644 --- a/litellm/rust_bridge/lifecycle.py +++ b/litellm/rust_bridge/lifecycle.py @@ -1,8 +1,8 @@ from __future__ import annotations -from collections.abc import Awaitable +from collections.abc import AsyncIterator, Awaitable, Iterator from dataclasses import dataclass -from typing import Protocol +from typing import Final, Protocol @dataclass(frozen=True, slots=True) @@ -15,28 +15,133 @@ class Complete: value: object +@dataclass(frozen=True, slots=True) +class Open: + value: None + + +@dataclass(frozen=True, slots=True) +class Yield: + value: object + + +Settled = Complete | Open | Yield +Step = Await | Settled + + class Execution(Protocol): - def start(self) -> Await | Complete: ... + def start(self) -> Step: ... - def resume_value(self, value: object) -> Await | Complete: ... + def resume_value(self, value: object) -> Step: ... - def resume_error(self, error: BaseException) -> Await | Complete: ... + def resume_error(self, error: BaseException) -> Step: ... def close(self) -> None: ... +class StreamClosed(Exception): + """Tells a streaming execution that its caller stopped reading.""" + + +async def _settle(execution: Execution, step: Step) -> Settled: + while isinstance(step, Await): + try: + value = await step.awaitable # rebind-ok: each selected await produces the next protocol input + except GeneratorExit: + raise + except BaseException as error: + step = execution.resume_error(error) # rebind-ok: advance the execution protocol + else: + step = execution.resume_value(value) # rebind-ok: advance the execution protocol + return step + + +def _settled(step: Step) -> Settled: + if isinstance(step, Await): + raise RuntimeError("sync call suspended") + return step + + async def drive(execution: Execution) -> object: + handed_off = False # rebind-ok: set once the execution belongs to the returned stream try: - step = execution.start() # rebind-ok: the execution protocol advances after each selected await - while isinstance(step, Await): - try: - value = await step.awaitable # rebind-ok: each selected await produces the next protocol input - except GeneratorExit: - raise - except BaseException as error: - step = execution.resume_error(error) # rebind-ok: advance the execution protocol - else: - step = execution.resume_value(value) # rebind-ok: advance the execution protocol + step: Final = await _settle(execution, execution.start()) + if isinstance(step, Open): + handed_off = True + return Stream(execution) return step.value finally: - execution.close() + if not handed_off: + execution.close() + + +class Stream(AsyncIterator[object]): + """A streamed native call: each read resumes the execution until its next chunk.""" + + def __init__(self, execution: Execution) -> None: + self._execution: Final = execution + self._done = False + + def __aiter__(self) -> Stream: + return self + + async def __anext__(self) -> object: + if self._done: + raise StopAsyncIteration + try: + step: Final = await _settle(self._execution, self._execution.resume_value(None)) + except BaseException: + self._finish() + raise + if isinstance(step, Yield): + return step.value + self._finish() + raise StopAsyncIteration + + async def aclose(self) -> None: + if self._done: + return + try: + await _settle(self._execution, self._execution.resume_error(StreamClosed())) + finally: + self._finish() + + def _finish(self) -> None: + self._done = True + self._execution.close() + + +class SyncStream(Iterator[object]): + """The sync form of `Stream`; its execution never suspends on an awaitable.""" + + def __init__(self, execution: Execution) -> None: + self._execution: Final = execution + self._done = False + + def __iter__(self) -> SyncStream: + return self + + def __next__(self) -> object: + if self._done: + raise StopIteration + try: + step: Final = _settled(self._execution.resume_value(None)) + except BaseException: + self._finish() + raise + if isinstance(step, Yield): + return step.value + self._finish() + raise StopIteration + + def close(self) -> None: + if self._done: + return + try: + _settled(self._execution.resume_error(StreamClosed())) + finally: + self._finish() + + def _finish(self) -> None: + self._done = True + self._execution.close() diff --git a/litellm/rust_bridge/messages/entrypoints.py b/litellm/rust_bridge/messages/entrypoints.py index 46565bfd46a..d25c906c4c1 100644 --- a/litellm/rust_bridge/messages/entrypoints.py +++ b/litellm/rust_bridge/messages/entrypoints.py @@ -1,6 +1,6 @@ from __future__ import annotations -from collections.abc import Awaitable, Mapping, Sequence +from collections.abc import AsyncIterator, Awaitable, Iterator, Mapping, Sequence from dataclasses import dataclass from typing import Final, Protocol, cast # noqa: TID251 # validates dynamically loaded native callables @@ -26,7 +26,7 @@ class NativeMessages(Protocol): request: LiteLLMMessagesRequest, args: tuple[object, ...], kwargs: Mapping[str, object], - ) -> AnthropicMessagesResponse: ... + ) -> AnthropicMessagesResponse | Iterator[bytes]: ... class NativeAmessages(Protocol): @@ -35,7 +35,7 @@ class NativeAmessages(Protocol): request: LiteLLMMessagesRequest, args: tuple[object, ...], kwargs: Mapping[str, object], - ) -> Awaitable[AnthropicMessagesResponse]: ... + ) -> Awaitable[AnthropicMessagesResponse | AsyncIterator[bytes]]: ... def _messages_binding(value: object) -> NativeMessages | None: @@ -50,5 +50,5 @@ def _amessages_binding(value: object) -> NativeAmessages | None: return cast("NativeAmessages", value) # cast-ok: callable validated at the native binding boundary -NATIVE_MESSAGES: Final = NativeBinding("anthropic_messages_handler", validate=_messages_binding) -NATIVE_AMESSAGES: Final = NativeBinding("anthropic_messages", validate=_amessages_binding) +NATIVE_MESSAGES: Final = NativeBinding("messages", validate=_messages_binding) +NATIVE_AMESSAGES: Final = NativeBinding("amessages", validate=_amessages_binding) diff --git a/litellm/rust_bridge/messages/route_host.py b/litellm/rust_bridge/messages/route_host.py index 1aff6c7f75d..beef0f81eca 100644 --- a/litellm/rust_bridge/messages/route_host.py +++ b/litellm/rust_bridge/messages/route_host.py @@ -20,4 +20,4 @@ def arguments(request: LiteLLMMessagesRequest) -> Mapping[str, object]: def map_failure(error: Exception, request: LiteLLMMessagesRequest, request_provider: str) -> Exception: - return failures.map_failure(error, request.model, request_provider, arguments(request)) + return failures.map_native_failure(error, request.model, request_provider, arguments(request), request.api_base) diff --git a/litellm/rust_bridge/ocr/route_host.py b/litellm/rust_bridge/ocr/route_host.py index 277fdceb734..bfbd5c11d4e 100644 --- a/litellm/rust_bridge/ocr/route_host.py +++ b/litellm/rust_bridge/ocr/route_host.py @@ -4,40 +4,17 @@ from collections.abc import Mapping from types import MappingProxyType from typing import Final -import httpx -import openai -from pydantic import TypeAdapter, ValidationError +from pydantic import TypeAdapter import litellm from litellm.llms.base_llm.ocr.transformation import PROVIDER_NATIVE_RESPONSE_KEY, OCRResponse from litellm.rust_bridge import failures +from litellm.rust_bridge.failures import UpstreamFailure from litellm.rust_bridge.ocr.entrypoints import LiteLLMOcrRequest +__all__ = ("UpstreamFailure", "arguments", "map_failure", "response") + _RESPONSE_ADAPTER: Final = TypeAdapter(dict[str, object]) -_UPSTREAM_ARGS: Final = TypeAdapter(tuple[int, str]) -_UPSTREAM_HEADERS: Final = TypeAdapter(list[tuple[str, str]]) - - -class UpstreamFailure(Exception): - def __init__(self, response: httpx.Response, cause: Exception) -> None: - super().__init__(str(cause)) - self.message: Final = str(cause) - self.response: Final = response - self.status_code: Final = response.status_code - self.__cause__ = cause - - -def _upstream_failure(error: Exception, request: LiteLLMOcrRequest) -> Exception: - try: - status, body = _UPSTREAM_ARGS.validate_python(error.args) - headers: Final = _UPSTREAM_HEADERS.validate_python(getattr(error, "headers", None)) - except ValidationError: - return error - http_request: Final = httpx.Request("POST", request.api_base or "https://docs.litellm.ai/docs") - return UpstreamFailure( - httpx.Response(status, content=body.encode(), headers=headers, request=http_request), - error, - ) def response(value: Mapping[str, object]) -> OCRResponse: @@ -61,11 +38,4 @@ def map_failure(error: Exception, request: LiteLLMOcrRequest, request_provider: model=request.model.removeprefix(f"{request_provider}/"), llm_provider=request_provider, ) - original: Final = _upstream_failure(error, request) - public_error: Final = failures.map_failure(original, request.model, request_provider, arguments(request)) - if isinstance(original, UpstreamFailure) and public_error.__context__ is original: - public_error.__context__ = error - if isinstance(public_error, openai.APIStatusError): - public_error.response = original.response - public_error.status_code = original.status_code - return public_error + return failures.map_native_failure(error, request.model, request_provider, arguments(request), request.api_base) diff --git a/litellm/types/llms/anthropic.py b/litellm/types/llms/anthropic.py index bcdee86360e..bcd24695f25 100644 --- a/litellm/types/llms/anthropic.py +++ b/litellm/types/llms/anthropic.py @@ -56,6 +56,7 @@ class AnthropicMessagesTool(TypedDict, total=False): defer_loading: bool allowed_callers: list[str] | None input_examples: list[dict[str, Any]] | None + eager_input_streaming: ReadOnly[bool] class AnthropicComputerTool(TypedDict, total=False): @@ -755,6 +756,8 @@ ANTHROPIC_TOOL_SEARCH_BETA_HEADER: Final = "advanced-tool-use-2025-11-20" # Effort beta header constant ANTHROPIC_EFFORT_BETA_HEADER: Final = "effort-2025-11-24" +ANTHROPIC_FINE_GRAINED_TOOL_STREAMING_BETA_HEADER: Final = "fine-grained-tool-streaming-2025-05-14" + # OAuth constants ANTHROPIC_OAUTH_TOKEN_PREFIX: Final = "sk-ant-oat" ANTHROPIC_OAUTH_BETA_HEADER: Final = "oauth-2025-04-20" diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py index 710b34116e5..632efcc3c4f 100644 --- a/litellm/types/llms/openai.py +++ b/litellm/types/llms/openai.py @@ -992,6 +992,7 @@ class ChatCompletionToolParamFunctionChunk(TypedDict, total=False): description: str parameters: dict strict: bool + eager_input_streaming: ReadOnly[bool] class OpenAIChatCompletionToolParam(TypedDict): @@ -1002,6 +1003,7 @@ class OpenAIChatCompletionToolParam(TypedDict): class ChatCompletionToolParam(OpenAIChatCompletionToolParam, total=False): cache_control: ChatCompletionCachedContent allowed_callers: list[str] + eager_input_streaming: ReadOnly[bool] class Function(TypedDict, total=False): diff --git a/litellm/types/llms/vertex_ai.py b/litellm/types/llms/vertex_ai.py index 3b95b786631..ce51e46ef15 100644 --- a/litellm/types/llms/vertex_ai.py +++ b/litellm/types/llms/vertex_ai.py @@ -425,22 +425,35 @@ class UrlContextMetadata(TypedDict, total=False): urlMetadata: list[UrlMetadata] +GeminiFinishReason = Literal[ + "FINISH_REASON_UNSPECIFIED", + "STOP", + "MAX_TOKENS", + "SAFETY", + "RECITATION", + "LANGUAGE", + "OTHER", + "BLOCKLIST", + "PROHIBITED_CONTENT", + "SPII", + "MALFORMED_FUNCTION_CALL", + "IMAGE_SAFETY", + "IMAGE_PROHIBITED_CONTENT", + "TOO_MANY_TOOL_CALLS", + "MALFORMED_RESPONSE", + "NO_IMAGE", + "IMAGE_RECITATION", + "IMAGE_OTHER", + "ESCALATION", + "UNEXPECTED_TOOL_CALL", + "MISSING_THOUGHT_SIGNATURE", +] + + class Candidates(TypedDict, total=False): index: int content: HttpxContentType - finishReason: Literal[ - "FINISH_REASON_UNSPECIFIED", - "STOP", - "MAX_TOKENS", - "SAFETY", - "RECITATION", - "OTHER", - "BLOCKLIST", - "PROHIBITED_CONTENT", - "SPII", - "MALFORMED_FUNCTION_CALL", - "IMAGE_SAFETY", - ] + finishReason: GeminiFinishReason safetyRatings: list[SafetyRatings] citationMetadata: CitationMetadata groundingMetadata: GroundingMetadata diff --git a/litellm/types/responses/streaming_websocket.py b/litellm/types/responses/streaming_websocket.py index 2aa71647955..f369cbcebf8 100644 --- a/litellm/types/responses/streaming_websocket.py +++ b/litellm/types/responses/streaming_websocket.py @@ -1,5 +1,7 @@ from __future__ import annotations +from collections.abc import Mapping +from dataclasses import dataclass from typing import Protocol from litellm.types.guardrails import PresidioPerRequestConfig @@ -39,3 +41,14 @@ class PresidioGuardrailCallback(Protocol): presidio_config: PresidioPerRequestConfig | None, request_data: dict[str, object], ) -> str: ... + + +@dataclass(frozen=True, slots=True) +class ResponsesWebSocketRequestDefaults: + """Deployment-level request parameters merged into every ``response.create`` frame relayed over a native websocket.""" + + fill_missing: Mapping[str, object] + overrides: Mapping[str, object] + + def merged_into(self, request: Mapping[str, object]) -> dict[str, object]: + return {**self.fill_missing, **request, **self.overrides} diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 20fbb4ed956..954754a6f14 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -50869,7 +50869,7 @@ "vertex_ai/google/gemma-4-26b-a4b-it-maas": { "input_cost_per_token": 1.5e-07, "litellm_provider": "vertex_ai-openai_models", - "max_input_tokens": 256000, + "max_input_tokens": 262144, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", diff --git a/tests/test_litellm/integrations/otel/test_otel_v2_components.py b/tests/test_litellm/integrations/otel/test_otel_v2_components.py index 72f6213e880..0c95049ce05 100644 --- a/tests/test_litellm/integrations/otel/test_otel_v2_components.py +++ b/tests/test_litellm/integrations/otel/test_otel_v2_components.py @@ -477,19 +477,21 @@ def _test_tracer(): return provider.get_tracer("test") +_CALLER_TRACEPARENT = "00-11111111111111111111111111111111-2222222222222222-01" + + def test_inject_trace_context_prefers_request_root_span(): def run(): tracer = _test_tracer() - with tracer.start_as_current_span("root") as root: + inbound = TraceContextTextMapPropagator().extract({"traceparent": _CALLER_TRACEPARENT}) + with tracer.start_as_current_span("root", context=inbound) as root: ctx_mod.set_request_root_span(root) - result = ctx_mod.inject_trace_context( - {"traceparent": "00-11111111111111111111111111111111-2222222222222222-01"} - ) + result = ctx_mod.inject_trace_context({"traceparent": _CALLER_TRACEPARENT}) propagated = get_current_span(TraceContextTextMapPropagator().extract(result)) return result, root, propagated result, root, propagated = ContextVarContext().run(run) - assert result["traceparent"] != "00-11111111111111111111111111111111-2222222222222222-01" + assert result["traceparent"] != _CALLER_TRACEPARENT assert propagated.get_span_context().trace_id == root.get_span_context().trace_id assert propagated.get_span_context().span_id == root.get_span_context().span_id @@ -507,24 +509,57 @@ def test_inject_trace_context_uses_ambient_span_without_request_root(): assert propagated.get_span_context().span_id == ambient.get_span_context().span_id -def test_inject_trace_context_replaces_stale_trace_headers(): +def test_inject_trace_context_replaces_same_trace_headers_with_request_span(): def run(): tracer = _test_tracer() - with tracer.start_as_current_span("ambient") as ambient: - headers = { - "Traceparent": "00-" + "a" * 32 + "-" + "b" * 16 + "-01", - "Tracestate": "vendor=old", - "x-keep": "1", - } + headers = {"Traceparent": _CALLER_TRACEPARENT, "Tracestate": "vendor=caller", "x-keep": "1"} + inbound = TraceContextTextMapPropagator().extract({key.lower(): value for key, value in headers.items()}) + with tracer.start_as_current_span("ambient", context=inbound) as ambient: result = ctx_mod.inject_trace_context(headers) propagated = get_current_span(TraceContextTextMapPropagator().extract(result)) return result, ambient, propagated result, ambient, propagated = ContextVarContext().run(run) assert sum(key.lower() == "traceparent" for key in result) == 1 - assert not any(key.lower() == "tracestate" for key in result) + assert sum(key.lower() == "tracestate" for key in result) == 1 assert result["x-keep"] == "1" - assert propagated.get_span_context().trace_id == ambient.get_span_context().trace_id + assert result["tracestate"] == "vendor=caller" + assert propagated.get_span_context().span_id == ambient.get_span_context().span_id + + +def test_inject_trace_context_keeps_caller_traceparent_from_another_trace(): + def run(): + tracer = _test_tracer() + parent = tracer.start_span("litellm_request") + with tracer.start_as_current_span("ambient") as ambient: + ctx_mod.set_request_root_span(ambient) + headers = {"Traceparent": _CALLER_TRACEPARENT, "Tracestate": "vendor=caller", "x-keep": "1"} + result = ctx_mod.inject_trace_context(headers, parent_span=parent) + propagated = get_current_span(TraceContextTextMapPropagator().extract(result)) + return result, parent, propagated + + result, parent, propagated = ContextVarContext().run(run) + assert result["traceparent"] == _CALLER_TRACEPARENT + assert result["tracestate"] == "vendor=caller" + assert result["x-keep"] == "1" + assert sum(key.lower() == "traceparent" for key in result) == 1 + assert sum(key.lower() == "tracestate" for key in result) == 1 + assert propagated.get_span_context().trace_id != parent.get_span_context().trace_id + + +def test_inject_trace_context_replaces_malformed_caller_traceparent(): + def run(): + tracer = _test_tracer() + parent = tracer.start_span("litellm_request") + with tracer.start_as_current_span("ambient"): + headers = {"traceparent": "not-a-traceparent", "tracestate": "vendor=caller"} + result = ctx_mod.inject_trace_context(headers, parent_span=parent) + propagated = get_current_span(TraceContextTextMapPropagator().extract(result)) + return result, parent, propagated + + result, parent, propagated = ContextVarContext().run(run) + assert propagated.get_span_context().span_id == parent.get_span_context().span_id + assert "tracestate" not in result def test_inject_trace_context_prefers_explicit_parent_span_over_root_and_ambient(): diff --git a/tests/test_litellm/litellm_core_utils/test_core_helpers.py b/tests/test_litellm/litellm_core_utils/test_core_helpers.py index e937be47441..b2ad13c205e 100644 --- a/tests/test_litellm/litellm_core_utils/test_core_helpers.py +++ b/tests/test_litellm/litellm_core_utils/test_core_helpers.py @@ -151,6 +151,12 @@ class TestMapFinishReasonGemini: ("IMAGE_PROHIBITED_CONTENT", "content_filter"), ("TOO_MANY_TOOL_CALLS", "stop"), ("MALFORMED_RESPONSE", "stop"), + ("NO_IMAGE", "content_filter"), + ("IMAGE_RECITATION", "content_filter"), + ("IMAGE_OTHER", "content_filter"), + ("ESCALATION", "content_filter"), + ("UNEXPECTED_TOOL_CALL", "stop"), + ("MISSING_THOUGHT_SIGNATURE", "stop"), ], ) def test_gemini_finish_reasons(self, gemini_reason, expected): diff --git a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py index 269c351f866..5e179f950a0 100644 --- a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py +++ b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py @@ -6370,3 +6370,74 @@ def test_response_format_tool_path_skips_forced_tool_choice_when_unsupported(loc assert "tools" in result assert "tool_choice" not in result + + +def _eager_chat_function(**extra: object) -> dict[str, object]: + return { + "name": "write_file", + "description": "Write a file", + "parameters": {"type": "object", "properties": {"path": {"type": "string"}}, "required": ["path"]}, + **extra, + } + + +def _eager_chat_tool(**extra: object) -> dict[str, object]: + return {"type": "function", "function": _eager_chat_function(), **extra} + + +@pytest.mark.parametrize("flag", [True, False]) +def test_eager_input_streaming_passed_through_from_tool_top_level(flag): + mapped_tool, _ = AnthropicConfig()._map_tool_helper(_eager_chat_tool(eager_input_streaming=flag)) + + assert mapped_tool == { + "name": "write_file", + "description": "Write a file", + "input_schema": {"type": "object", "properties": {"path": {"type": "string"}}, "required": ["path"]}, + "type": "custom", + "eager_input_streaming": flag, + } + + +def test_eager_input_streaming_passed_through_from_function(): + mapped_tool, _ = AnthropicConfig()._map_tool_helper( + {"type": "function", "function": _eager_chat_function(eager_input_streaming=True)} + ) + + assert mapped_tool["eager_input_streaming"] is True + assert "eager_input_streaming" not in mapped_tool["input_schema"] + + +def test_eager_input_streaming_absent_stays_absent(): + mapped_tool, _ = AnthropicConfig()._map_tool_helper(_eager_chat_tool()) + + assert "eager_input_streaming" not in mapped_tool + + +def test_eager_input_streaming_rejects_non_boolean(): + with pytest.raises(litellm.BadRequestError, match="eager_input_streaming must be a boolean"): + AnthropicConfig()._map_tool_helper(_eager_chat_tool(eager_input_streaming="true")) + + +def test_eager_input_streaming_not_set_on_computer_use_tool(): + computer_tool = { + "type": "computer_20250124", + "function": {"name": "computer", "parameters": {"display_width_px": 1024, "display_height_px": 768}}, + "eager_input_streaming": True, + } + + mapped_tool, _ = AnthropicConfig()._map_tool_helper(computer_tool) + + assert mapped_tool["type"] == "computer_20250124" + assert "eager_input_streaming" not in mapped_tool + + +def test_eager_input_streaming_reaches_anthropic_request_tools(): + result = AnthropicConfig().map_openai_params( + non_default_params={"tools": [_eager_chat_tool(eager_input_streaming=True)], "stream": True}, + optional_params={}, + model="claude-sonnet-5", + drop_params=False, + ) + + assert result["tools"][0]["eager_input_streaming"] is True + assert result["tools"][0]["name"] == "write_file" diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py index e6782b70d3e..a5fb19b236f 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py @@ -105,6 +105,46 @@ def test_translate_chat_length_takes_precedence_over_refusal(): assert result.get("stop_details") is None +def test_translate_chat_content_filter_to_anthropic_response(): + response = ModelResponse( + id="chatcmpl-content-filter", + model="openai-model", + choices=[ + Choices( + index=0, + finish_reason="content_filter", + message=Message(content=None, role="assistant"), + ) + ], + usage=Usage(prompt_tokens=1, completion_tokens=0, total_tokens=1), + ) + + result = LiteLLMAnthropicMessagesAdapter().translate_openai_response_to_anthropic(response) + + assert result["content"] == [] + assert result["stop_reason"] == "refusal" + + +def test_translate_chat_refusal_finish_reason_to_anthropic_response(): + response = ModelResponse( + id="chatcmpl-refusal-reason", + model="openai-model", + choices=[ + Choices( + index=0, + finish_reason="refusal", + message=Message(content=None, role="assistant"), + ) + ], + usage=Usage(prompt_tokens=1, completion_tokens=0, total_tokens=1), + ) + + result = LiteLLMAnthropicMessagesAdapter().translate_openai_response_to_anthropic(response) + + assert result["content"] == [] + assert result["stop_reason"] == "refusal" + + def test_translate_streaming_openai_chunk_to_anthropic_content_block(): choices = [ StreamingChoices( @@ -5017,3 +5057,45 @@ def test_redacted_thinking_blocks_never_carry_cache_control(): replayed: Final = outbound["messages"][1]["content"][0] assert replayed["type"] == "redacted_thinking" assert "cache_control" not in replayed + + +EAGER_INPUT_SCHEMA: Final = {"type": "object", "properties": {"path": {"type": "string"}}, "required": ["path"]} + + +@pytest.mark.parametrize("flag", [True, False]) +def test_translate_anthropic_tools_to_openai_carries_eager_input_streaming_onto_tool(flag): + """The per-tool flag lands on the OpenAI tool object, never inside the JSON schema Bedrock sends as inputSchema.""" + tools: Final = [{"name": "write_file", "input_schema": EAGER_INPUT_SCHEMA, "eager_input_streaming": flag}] + + new_tools, _ = LiteLLMAnthropicMessagesAdapter().translate_anthropic_tools_to_openai(tools=tools) + + assert new_tools[0]["eager_input_streaming"] is flag + assert new_tools[0]["function"]["parameters"] == EAGER_INPUT_SCHEMA + assert "eager_input_streaming" not in new_tools[0]["function"] + + +def test_translate_anthropic_tools_to_openai_omits_unset_eager_input_streaming(): + tools: Final = [{"name": "write_file", "input_schema": EAGER_INPUT_SCHEMA}] + + new_tools, _ = LiteLLMAnthropicMessagesAdapter().translate_anthropic_tools_to_openai(tools=tools) + + assert "eager_input_streaming" not in new_tools[0] + assert "eager_input_streaming" not in new_tools[0]["function"]["parameters"] + + +def test_eager_input_streaming_tool_reaches_bedrock_converse_as_beta(): + """An Anthropic Messages request routed to bedrock/converse/ turns the flag into the fine-grained streaming beta.""" + from litellm.llms.bedrock.chat.converse_transformation import AmazonConverseConfig + + tools: Final = [{"name": "write_file", "input_schema": EAGER_INPUT_SCHEMA, "eager_input_streaming": True}] + new_tools, _ = LiteLLMAnthropicMessagesAdapter().translate_anthropic_tools_to_openai(tools=tools) + + data: Final = AmazonConverseConfig()._transform_request_helper( + model="us.anthropic.claude-sonnet-4-5-20250929-v1:0", + system_content_blocks=[], + optional_params={"tools": new_tools}, + messages=[{"role": "user", "content": "write a big file"}], + ) + + assert data["additionalModelRequestFields"]["anthropic_beta"] == ["fine-grained-tool-streaming-2025-05-14"] + assert data["toolConfig"]["tools"][0]["toolSpec"]["inputSchema"]["json"] == EAGER_INPUT_SCHEMA diff --git a/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py b/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py index ec0bf6b842a..84db0733227 100644 --- a/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py @@ -1,6 +1,7 @@ import asyncio import json import uuid +from typing import Final from unittest.mock import patch import httpx @@ -859,3 +860,58 @@ def test_bedrock_chat_invoke_tool_search_beta_follows_model_map( ) assert result.get("anthropic_beta") == expected_betas + + +FINE_GRAINED_TOOL_STREAMING_BETA: Final = "fine-grained-tool-streaming-2025-05-14" +EAGER_TOOL_SCHEMA: Final = {"type": "object", "properties": {"path": {"type": "string"}}, "required": ["path"]} + + +def _chat_invoke_request_with_tools( + tools: list[dict[str, object]], headers: dict[str, str] | None = None +) -> dict[str, object]: + config: Final = AmazonAnthropicClaudeConfig() + model: Final = "us.anthropic.claude-sonnet-4-5-20250929-v1:0" + optional_params: Final = config.map_openai_params( + non_default_params={"max_tokens": 64, "stream": True, "tools": tools}, + optional_params={}, + model=model, + drop_params=False, + ) + return config.transform_request( + model=model, + messages=[{"role": "user", "content": "write a big file"}], + optional_params=optional_params, + litellm_params={}, + headers=headers or {}, + ) + + +def _eager_openai_tool(name: str, **extra: object) -> dict[str, object]: + return {"type": "function", "function": {"name": name, "parameters": EAGER_TOOL_SCHEMA}, **extra} + + +def test_bedrock_chat_invoke_eager_input_streaming_tool_adds_beta_and_strips_key(): + result = _chat_invoke_request_with_tools( + [_eager_openai_tool("write_file", eager_input_streaming=True), _eager_openai_tool("read_file")] + ) + + assert result["anthropic_beta"] == [FINE_GRAINED_TOOL_STREAMING_BETA] + assert [tool["name"] for tool in result["tools"]] == ["write_file", "read_file"] + assert all("eager_input_streaming" not in tool for tool in result["tools"]) + assert result["tools"][0]["input_schema"] == EAGER_TOOL_SCHEMA + + +def test_bedrock_chat_invoke_eager_input_streaming_false_strips_key_without_beta(): + result = _chat_invoke_request_with_tools([_eager_openai_tool("write_file", eager_input_streaming=False)]) + + assert "anthropic_beta" not in result + assert "eager_input_streaming" not in result["tools"][0] + + +def test_bedrock_chat_invoke_eager_input_streaming_beta_not_duplicated_with_client_header(): + result = _chat_invoke_request_with_tools( + [_eager_openai_tool("write_file", eager_input_streaming=True)], + headers={"anthropic-beta": FINE_GRAINED_TOOL_STREAMING_BETA}, + ) + + assert result["anthropic_beta"] == [FINE_GRAINED_TOOL_STREAMING_BETA] diff --git a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py index 96c78c1cf75..cd33e8d34d8 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py @@ -4,6 +4,7 @@ import os import httpx import pytest +from typing import Final from unittest.mock import MagicMock, patch import litellm @@ -7400,3 +7401,111 @@ def test_transform_response_honors_json_mode_kwarg_when_optional_params_lack_it( ) assert result.choices[0].message.tool_calls is None assert json.loads(result.choices[0].message.content) == {"city": "Paris", "population": 2100000} + + +FINE_GRAINED_TOOL_STREAMING_BETA: Final = "fine-grained-tool-streaming-2025-05-14" +EAGER_TOOL_SCHEMA: Final = {"type": "object", "properties": {"path": {"type": "string"}}, "required": ["path"]} + + +def _eager_openai_tool(**extra: object) -> dict[str, object]: + return {"type": "function", "function": {"name": "write_file", "parameters": EAGER_TOOL_SCHEMA}, **extra} + + +def _eager_openai_function_tool(**extra: object) -> dict[str, object]: + return {"type": "function", "function": {"name": "write_file", "parameters": EAGER_TOOL_SCHEMA, **extra}} + + +def _eager_anthropic_tool(**extra: object) -> dict[str, object]: + return {"name": "write_file", "input_schema": EAGER_TOOL_SCHEMA, **extra} + + +def _converse_request( + model: str, tools: list[dict[str, object]], headers: dict[str, object] | None = None +) -> dict[str, object]: + return AmazonConverseConfig()._transform_request_helper( + model=model, + system_content_blocks=[], + optional_params={"tools": tools}, + messages=[{"role": "user", "content": "write a big file"}], + headers=headers, + ) + + +@pytest.mark.parametrize( + "tool", + [ + _eager_openai_tool(eager_input_streaming=True), + _eager_openai_function_tool(eager_input_streaming=True), + _eager_anthropic_tool(eager_input_streaming=True), + ], + ids=["openai_top_level", "openai_under_function", "anthropic_shape"], +) +def test_eager_input_streaming_tool_adds_fine_grained_tool_streaming_beta(tool): + data = _converse_request("us.anthropic.claude-sonnet-4-5-20250929-v1:0", [tool]) + + assert data["additionalModelRequestFields"]["anthropic_beta"] == [FINE_GRAINED_TOOL_STREAMING_BETA] + tool_spec = data["toolConfig"]["tools"][0]["toolSpec"] + assert tool_spec["name"] == "write_file" + assert "eager_input_streaming" not in tool_spec + assert "eager_input_streaming" not in tool_spec["inputSchema"]["json"] + + +@pytest.mark.parametrize( + "tool", + [ + _eager_openai_tool(eager_input_streaming=False), + _eager_openai_function_tool(eager_input_streaming=False), + _eager_anthropic_tool(eager_input_streaming=False), + _eager_openai_tool(), + ], + ids=["openai_false", "function_false", "anthropic_false", "absent"], +) +def test_eager_input_streaming_false_or_absent_adds_no_beta(tool): + data = _converse_request("us.anthropic.claude-sonnet-4-5-20250929-v1:0", [tool]) + + assert "anthropic_beta" not in data.get("additionalModelRequestFields", {}) + assert "eager_input_streaming" not in data["toolConfig"]["tools"][0]["toolSpec"] + + +def test_eager_input_streaming_beta_only_on_anthropic_models(): + data = _converse_request("amazon.nova-pro-v1:0", [_eager_openai_tool(eager_input_streaming=True)]) + + assert "anthropic_beta" not in data.get("additionalModelRequestFields", {}) + assert data["toolConfig"]["tools"][0]["toolSpec"]["name"] == "write_file" + + +def test_eager_input_streaming_beta_not_duplicated_with_client_header(): + data = _converse_request( + "us.anthropic.claude-sonnet-4-5-20250929-v1:0", + [_eager_openai_tool(eager_input_streaming=True)], + headers={"anthropic-beta": f"{FINE_GRAINED_TOOL_STREAMING_BETA},interleaved-thinking-2025-05-14"}, + ) + + assert data["additionalModelRequestFields"]["anthropic_beta"] == [ + FINE_GRAINED_TOOL_STREAMING_BETA, + "interleaved-thinking-2025-05-14", + ] + + +def test_eager_input_streaming_beta_never_written_back_into_client_header_list(): + headers = {"anthropic-beta": ["interleaved-thinking-2025-05-14"]} + + data = _converse_request( + "us.anthropic.claude-sonnet-4-5-20250929-v1:0", + [_eager_openai_tool(eager_input_streaming=True)], + headers=headers, + ) + + assert data["additionalModelRequestFields"]["anthropic_beta"] == [ + "interleaved-thinking-2025-05-14", + FINE_GRAINED_TOOL_STREAMING_BETA, + ] + assert headers == {"anthropic-beta": ["interleaved-thinking-2025-05-14"]} + + +def test_eager_input_streaming_non_boolean_is_a_bad_request(): + with pytest.raises(litellm.BadRequestError, match="eager_input_streaming must be a boolean"): + _converse_request( + "us.anthropic.claude-sonnet-4-5-20250929-v1:0", + [_eager_openai_tool(eager_input_streaming="true")], + ) diff --git a/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py b/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py index e43accdb835..7be005c0efe 100644 --- a/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py +++ b/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py @@ -4,6 +4,7 @@ import json import os from datetime import datetime from types import SimpleNamespace +from typing import Final from unittest.mock import Mock import pytest @@ -3244,3 +3245,67 @@ def test_bedrock_messages_strips_effort_but_keeps_format_for_sonnet_4_5(local_mo ) assert result.get("output_config") == {"format": schema_format} + + +FINE_GRAINED_TOOL_STREAMING_BETA: Final = "fine-grained-tool-streaming-2025-05-14" + + +def _invoke_request_with_tools( + tools: list[dict[str, object]], headers: dict[str, str] | None = None +) -> dict[str, object]: + from litellm.types.router import GenericLiteLLMParams + + return AmazonAnthropicClaudeMessagesConfig().transform_anthropic_messages_request( + model="us.anthropic.claude-sonnet-4-5-20250929-v1:0", + messages=[{"role": "user", "content": "write a big file"}], + anthropic_messages_optional_request_params={"max_tokens": 4096, "tools": copy.deepcopy(tools), "stream": True}, + litellm_params=GenericLiteLLMParams(), + headers=headers or {}, + ) + + +def _eager_invoke_tool(name: str, eager_input_streaming: bool) -> dict[str, object]: + return { + "name": name, + "description": f"{name} tool", + "input_schema": {"type": "object", "properties": {"path": {"type": "string"}}}, + "eager_input_streaming": eager_input_streaming, + } + + +def test_bedrock_invoke_eager_input_streaming_tool_adds_beta_and_strips_key(): + result = _invoke_request_with_tools( + [ + _eager_invoke_tool("write_file", True), + _eager_invoke_tool("read_file", False), + {"name": "list_files", "input_schema": {"type": "object", "properties": {}}}, + ] + ) + + assert result["anthropic_beta"] == [FINE_GRAINED_TOOL_STREAMING_BETA] + assert [tool["name"] for tool in result["tools"]] == ["write_file", "read_file", "list_files"] + assert all("eager_input_streaming" not in tool for tool in result["tools"]) + assert result["tools"][0]["description"] == "write_file tool" + assert result["tools"][0]["input_schema"] == {"type": "object", "properties": {"path": {"type": "string"}}} + + +def test_bedrock_invoke_eager_input_streaming_false_strips_key_without_beta(): + result = _invoke_request_with_tools([_eager_invoke_tool("write_file", False)]) + + assert "anthropic_beta" not in result + assert result["tools"] == [ + { + "name": "write_file", + "description": "write_file tool", + "input_schema": {"type": "object", "properties": {"path": {"type": "string"}}}, + } + ] + + +def test_bedrock_invoke_eager_input_streaming_beta_not_duplicated_with_client_header(): + result = _invoke_request_with_tools( + [_eager_invoke_tool("write_file", True)], + headers={"anthropic-beta": FINE_GRAINED_TOOL_STREAMING_BETA}, + ) + + assert result["anthropic_beta"] == [FINE_GRAINED_TOOL_STREAMING_BETA] diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py index 84295666fbf..6c818016c87 100644 --- a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py +++ b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py @@ -2,7 +2,7 @@ import asyncio import json import re from copy import deepcopy -from typing import Final, List, cast +from typing import Final, List, cast, get_args from unittest.mock import MagicMock, patch import httpx @@ -18,7 +18,7 @@ from litellm.llms.vertex_ai.common_utils import VertexAIError from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( VertexGeminiConfig, ) -from litellm.types.llms.vertex_ai import UsageMetadata +from litellm.types.llms.vertex_ai import GeminiFinishReason, UsageMetadata from litellm.types.utils import ChoiceLogprobs, Usage from litellm.utils import CustomStreamWrapper @@ -940,6 +940,11 @@ def test_check_finish_reason(): ) +def test_every_documented_gemini_finish_reason_has_an_explicit_mapping(): + documented: Final = frozenset(get_args(GeminiFinishReason)) + assert set(VertexGeminiConfig.get_finish_reason_mapping()) == documented + + def test_finish_reason_unspecified_and_malformed_function_call(): """ Test that FINISH_REASON_UNSPECIFIED and MALFORMED_FUNCTION_CALL @@ -968,6 +973,12 @@ def test_finish_reason_unspecified_and_malformed_function_call(): # Test new Gemini finish reasons assert finish_reason_mappings["TOO_MANY_TOOL_CALLS"] == "stop" assert finish_reason_mappings["MALFORMED_RESPONSE"] == "stop" + assert finish_reason_mappings["NO_IMAGE"] == "content_filter" + assert finish_reason_mappings["IMAGE_RECITATION"] == "content_filter" + assert finish_reason_mappings["IMAGE_OTHER"] == "content_filter" + assert finish_reason_mappings["ESCALATION"] == "content_filter" + assert finish_reason_mappings["UNEXPECTED_TOOL_CALL"] == "stop" + assert finish_reason_mappings["MISSING_THOUGHT_SIGNATURE"] == "stop" def test_vertex_ai_usage_metadata_response_token_count(): @@ -6074,3 +6085,210 @@ def test_prompt_blocked_chunk_keeps_served_model_version(): assert streaming_chunk.model == "gemini-3.8-flash-001" assert streaming_chunk.choices[0].finish_reason == "content_filter" + + +def test_gemini_candidate_with_finish_reason_no_content_chat_completion(): + config = VertexGeminiConfig() + completion_response = { + "candidates": [ + { + "finishReason": "NO_IMAGE", + "index": 0, + } + ], + "usageMetadata": { + "promptTokenCount": 19, + "candidatesTokenCount": 0, + "totalTokenCount": 19, + }, + } + model_response = ModelResponse() + logging_obj = MagicMock() + raw_response = MagicMock() + raw_response.headers = {} + + resp = config._transform_google_generate_content_to_openai_model_response( + completion_response=completion_response, + model_response=model_response, + model="gemini-2.5-flash-image", + logging_obj=logging_obj, + raw_response=raw_response, + ) + assert len(resp.choices) == 1 + assert resp.choices[0].finish_reason == "content_filter" + assert resp.choices[0].message.content is None + assert resp.choices[0].provider_specific_fields["native_finish_reason"] == "NO_IMAGE" + + +def test_gemini_candidate_with_finish_reason_no_content_anthropic_messages(): + from litellm.llms.anthropic.experimental_pass_through.adapters.transformation import ( + LiteLLMAnthropicMessagesAdapter, + ) + + config = VertexGeminiConfig() + completion_response = { + "candidates": [ + { + "finishReason": "NO_IMAGE", + "index": 0, + } + ], + "usageMetadata": { + "promptTokenCount": 19, + "candidatesTokenCount": 0, + "totalTokenCount": 19, + }, + } + resp = config._transform_google_generate_content_to_openai_model_response( + completion_response=completion_response, + model_response=ModelResponse(), + model="gemini-2.5-flash-image", + logging_obj=MagicMock(), + raw_response=MagicMock(headers={}), + ) + + adapter = LiteLLMAnthropicMessagesAdapter() + anthropic_resp = adapter.translate_openai_response_to_anthropic( + response=resp, + tool_name_mapping={}, + ) + assert anthropic_resp["stop_reason"] == "refusal" + assert anthropic_resp["content"] == [] + + +def test_gemini_candidate_with_finish_reason_no_content_responses_api(): + from litellm.responses.litellm_completion_transformation.transformation import ( + LiteLLMCompletionResponsesConfig, + ) + + config = VertexGeminiConfig() + completion_response = { + "candidates": [ + { + "finishReason": "NO_IMAGE", + "index": 0, + } + ], + "usageMetadata": { + "promptTokenCount": 19, + "candidatesTokenCount": 0, + "totalTokenCount": 19, + }, + } + resp = config._transform_google_generate_content_to_openai_model_response( + completion_response=completion_response, + model_response=ModelResponse(), + model="gemini-2.5-flash-image", + logging_obj=MagicMock(), + raw_response=MagicMock(headers={}), + ) + + responses_resp = LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( + request_input="Generate picture", + responses_api_request={}, + chat_completion_response=resp, + ) + assert responses_resp.status == "incomplete" + assert responses_resp.incomplete_details is not None + assert responses_resp.incomplete_details.reason == "content_filter" + + +def test_gemini_candidate_other_finish_reasons_no_content(): + from litellm.llms.anthropic.experimental_pass_through.adapters.transformation import ( + LiteLLMAnthropicMessagesAdapter, + ) + from litellm.responses.litellm_completion_transformation.transformation import ( + LiteLLMCompletionResponsesConfig, + ) + + config = VertexGeminiConfig() + max_tokens_response = { + "candidates": [{"finishReason": "MAX_TOKENS", "index": 0}], + "usageMetadata": {"promptTokenCount": 10, "candidatesTokenCount": 50, "totalTokenCount": 60}, + } + resp_length = config._transform_google_generate_content_to_openai_model_response( + completion_response=max_tokens_response, + model_response=ModelResponse(), + model="gemini-2.5-flash", + logging_obj=MagicMock(), + raw_response=MagicMock(headers={}), + ) + assert len(resp_length.choices) == 1 + assert resp_length.choices[0].finish_reason == "length" + assert resp_length.choices[0].provider_specific_fields["native_finish_reason"] == "MAX_TOKENS" + + anthropic_length = LiteLLMAnthropicMessagesAdapter().translate_openai_response_to_anthropic( + response=resp_length, + tool_name_mapping={}, + ) + assert anthropic_length["stop_reason"] == "max_tokens" + + responses_length = LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( + request_input="thinking request", + responses_api_request={}, + chat_completion_response=resp_length, + ) + assert responses_length.status == "incomplete" + assert responses_length.incomplete_details.reason == "max_output_tokens" + + +def test_gemini_candidate_with_finish_reason_no_content_streaming_chunk(): + from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( + ModelResponseIterator, + ) + + chunk: Final = { + "candidates": [{"finishReason": "NO_IMAGE", "index": 0}], + "usageMetadata": {"promptTokenCount": 19, "candidatesTokenCount": 0, "totalTokenCount": 19}, + } + iterator: Final = ModelResponseIterator(streaming_response=[], sync_stream=True, logging_obj=MagicMock()) + + streaming_chunk: Final = iterator.chunk_parser(chunk) + + assert len(streaming_chunk.choices) == 1 + assert streaming_chunk.choices[0].finish_reason == "content_filter" + assert streaming_chunk.choices[0].delta.content is None + assert streaming_chunk.choices[0].delta.tool_calls is None + + +def test_gemini_multi_candidate_messages_do_not_share_state(): + config: Final = VertexGeminiConfig() + completion_response: Final = { + "candidates": [ + { + "content": { + "role": "model", + "parts": [ + {"text": "Let me check the weather.", "thought": True}, + {"functionCall": {"name": "get_weather", "args": {"city": "Paris"}}}, + ], + }, + "finishReason": "STOP", + "index": 0, + }, + { + "content": {"role": "model", "parts": [{"text": "It is sunny in Paris."}]}, + "finishReason": "STOP", + "index": 1, + }, + ], + "usageMetadata": {"promptTokenCount": 10, "candidatesTokenCount": 20, "totalTokenCount": 30}, + } + + resp: Final = config._transform_google_generate_content_to_openai_model_response( + completion_response=completion_response, + model_response=ModelResponse(), + model="gemini-2.5-flash", + logging_obj=MagicMock(), + raw_response=MagicMock(headers={}), + ) + + assert len(resp.choices) == 2 + assert resp.choices[0].finish_reason == "tool_calls" + assert resp.choices[0].message.tool_calls[0].function.name == "get_weather" + assert resp.choices[0].message.reasoning_content == "Let me check the weather." + assert resp.choices[1].finish_reason == "stop" + assert resp.choices[1].message.content == "It is sunny in Paris." + assert resp.choices[1].message.tool_calls is None + assert getattr(resp.choices[1].message, "reasoning_content", None) is None + assert resp.choices[1].provider_specific_fields["native_finish_reason"] == "STOP" diff --git a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/gemma/test_vertex_ai_gemma_global_endpoint.py b/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/gemma/test_vertex_ai_gemma_global_endpoint.py index e9b58622a4b..85a2124ab02 100644 --- a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/gemma/test_vertex_ai_gemma_global_endpoint.py +++ b/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/gemma/test_vertex_ai_gemma_global_endpoint.py @@ -30,7 +30,7 @@ from litellm.types.llms.vertex_ai import VertexPartnerProvider _GEMMA_MODEL_COST_ENTRY = { "vertex_ai/google/gemma-4-26b-a4b-it-maas": { "litellm_provider": "vertex_ai-openai_models", - "max_input_tokens": 256000, + "max_input_tokens": 262144, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", @@ -180,6 +180,15 @@ class TestCreateVertexURLGemma: # --------------------------------------------------------------------------- +def test_gemma_maas_context_window_matches_google(local_model_cost_map): + info = litellm.get_model_info("vertex_ai/google/gemma-4-26b-a4b-it-maas") + + # 262,144 context length and 128,000 maximum output per Google's model page, checked 2026-09-18: + # https://docs.cloud.google.com/gemini-enterprise-agent-platform/models/maas/google/gemma-4-26b-a4b-it + assert info["max_input_tokens"] == 262144 + assert info["max_output_tokens"] == 128000 + + # --------------------------------------------------------------------------- # Integration tests: verify payloads reach the global OpenAI endpoint # diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py index be090a309fc..bf8ef920bdc 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -4418,6 +4418,65 @@ async def test_pass_through_request_propagates_active_trace_context(span_source: assert propagated.get_span_context().span_id == span.get_span_context().span_id +async def _relay_with_trace_headers(inbound_headers: dict[str, str], forward_headers: bool): + from opentelemetry.sdk.trace import TracerProvider + + captured: dict[str, httpx.Headers] = {} + + def transport_handler(upstream_request: httpx.Request) -> httpx.Response: + captured["headers"] = upstream_request.headers + return httpx.Response(200, json={"ok": True}, request=upstream_request) + + fake_client, cleanup = _inject_fake_passthrough_client(httpx.MockTransport(transport_handler), timeout=None) + tracer = TracerProvider().get_tracer("test") + try: + with ExitStack() as stack: + _enter_relay_logging_mocks(stack, {}) + span = tracer.start_span("litellm_request") + stack.callback(span.end) + request = _relay_client_request(method="POST") + request.headers = Headers(inbound_headers) + response = await pass_through_request( + request=request, + target="http://internal-api.test/v1/generate", + custom_headers={}, + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test", parent_otel_span=span), + forward_headers=forward_headers, + ) + finally: + cleanup() + await fake_client.aclose() + assert response.status_code == 200 + return captured["headers"], span + + +@pytest.mark.asyncio +@pytest.mark.parametrize("forward_headers", [False, True]) +async def test_pass_through_request_keeps_x_pass_trace_headers_when_otel_span_is_active(forward_headers: bool): + caller_traceparent = "00-11111111111111111111111111111111-2222222222222222-01" + + upstream_headers, span = await _relay_with_trace_headers( + {"x-pass-traceparent": caller_traceparent, "x-pass-tracestate": "vendor=caller"}, + forward_headers=forward_headers, + ) + + assert upstream_headers["traceparent"] == caller_traceparent + assert upstream_headers["tracestate"] == "vendor=caller" + assert format(span.get_span_context().trace_id, "032x") not in upstream_headers["traceparent"] + + +@pytest.mark.asyncio +async def test_pass_through_request_without_caller_trace_headers_still_propagates_proxy_span(): + from opentelemetry.trace import get_current_span + from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator + + upstream_headers, span = await _relay_with_trace_headers({"x-pass-anthropic-beta": "beta-1"}, forward_headers=False) + + propagated = get_current_span(TraceContextTextMapPropagator().extract(upstream_headers)) + assert propagated.get_span_context().span_id == span.get_span_context().span_id + assert upstream_headers["anthropic-beta"] == "beta-1" + + @pytest.mark.asyncio async def test_pass_through_request_relays_non_json_body_without_buffering(): """ diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index 135ce019813..e4ca0b03d59 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -40,9 +40,12 @@ from litellm.proxy.common_request_processing import ( _parse_event_data_for_error, _resolve_per_request_model_group_alias, _should_return_raw_model_name, + _sse_error_frames, _UpstreamClosingStreamingResponse, create_response, + sse_error_payload, ) +from litellm.proxy.common_utils.callback_utils import add_guardrail_to_applied_guardrails_header from litellm.proxy.dd_span_tagger import DDSpanTagger from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.proxy._types import ProxyErrorTypes, ProxyException @@ -9065,6 +9068,60 @@ class TestStreamingResponseHeadersFollowFallback: assert "llm_provider-stale-marker" not in result.headers assert result.headers["x-callback-header"] == "kept" + @pytest.mark.asyncio + async def test_streaming_block_headers_name_the_blocking_guardrail(self, monkeypatch): + processor_data: dict[str, object] = {"model": "oa", "stream": True, "metadata": {}} + + def select_data_generator(**kwargs): + async def generator(): + add_guardrail_to_applied_guardrails_header(processor_data, "stream-blocker") + _, error_obj = sse_error_payload(HTTPException(status_code=400, detail="blocked")) + for frame in _sse_error_frames(error_obj): + yield frame + + return generator() + + logging_obj = MagicMock() + logging_obj.litellm_call_id = "lit-7144-call" + logging_obj._defer_async_logging = False + logging_obj._on_deferred_stream_complete = None + logging_obj.cost_breakdown = None + processor_data["litellm_logging_obj"] = logging_obj + processor = ProxyBaseLLMRequestProcessing(data=processor_data) + + proxy_logging_obj = MagicMock(spec=ProxyLogging) + proxy_logging_obj.during_call_hook = AsyncMock(return_value=None) + proxy_logging_obj.update_request_status = AsyncMock(return_value=None) + proxy_logging_obj.post_call_success_hook = AsyncMock( + side_effect=lambda data, user_api_key_dict, response: response + ) + proxy_logging_obj.post_call_response_headers_hook = AsyncMock(return_value={}) + + async def fake_route_request(**kwargs): + async def call(): + return SimpleNamespace(_hidden_params={}, fallback_headers_adopted=False) + + return call() + + monkeypatch.setattr(litellm.proxy.common_request_processing, "route_request", fake_route_request) + + result = await processor.base_process_llm_request( + request=Request(scope={"type": "http", "headers": []}), + fastapi_response=Response(), + user_api_key_dict=ProxyUserAPIKeyAuth(api_key="sk-test"), + route_type="acompletion", + proxy_logging_obj=proxy_logging_obj, + general_settings={}, + proxy_config=MagicMock(spec=ProxyConfig), + select_data_generator=select_data_generator, + is_streaming_request=True, + skip_pre_call_logic=True, + ) + + assert isinstance(result, JSONResponse) + assert result.status_code == 400 + assert result.headers["x-litellm-applied-guardrails"] == "stream-blocker" + class _MessagesFallbackStream: def __init__(self) -> None: diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_during_call_hook.py b/tests/test_litellm/proxy/utils/proxy_logging/test_during_call_hook.py index 3c5d879c2dc..46f39ef6fb7 100644 --- a/tests/test_litellm/proxy/utils/proxy_logging/test_during_call_hook.py +++ b/tests/test_litellm/proxy/utils/proxy_logging/test_during_call_hook.py @@ -6,6 +6,7 @@ from typing import Any, Dict from unittest.mock import AsyncMock, MagicMock import pytest +from fastapi import HTTPException import litellm from litellm.integrations.custom_guardrail import CustomGuardrail @@ -84,3 +85,20 @@ async def test_during_call_hook_guardrail_error_raises(proxy_logging, make_user_ user_api_key_dict=make_user_api_key_auth(), call_type="completion", ) + + +@pytest.mark.asyncio +async def test_during_call_block_names_the_blocking_guardrail_in_applied_guardrails( + proxy_logging, make_user_api_key_auth, monkeypatch +): + g = _make_guardrail("blocker") + g.async_moderation_hook = AsyncMock(side_effect=HTTPException(status_code=400, detail="blocked")) + monkeypatch.setattr(litellm, "callbacks", [_make_guardrail("passer"), g]) + data = {"model": "m", "metadata": {}} + with pytest.raises(HTTPException): + await proxy_logging.during_call_hook( + data=data, + user_api_key_dict=make_user_api_key_auth(), + call_type="completion", + ) + assert "blocker" in data["metadata"]["applied_guardrails"] diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py b/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py index 077bf5a313e..5f3c09d9195 100644 --- a/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py +++ b/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py @@ -527,17 +527,19 @@ def test_handle_pipeline_result_block_enriches_with_guardrail_name_and_mode(): result.step_results = [MagicMock(guardrail_name="g")] result.original_exception = original + data: dict[str, object] = {"model": "m"} saved = litellm.callbacks litellm.callbacks = [cb] try: with pytest.raises(HTTPException) as info: - ProxyLogging._handle_pipeline_result(result=result, data={"model": "m"}, policy_name="p") + ProxyLogging._handle_pipeline_result(result=result, data=data, policy_name="p") finally: litellm.callbacks = saved assert info.value is original assert info.value.detail["guardrail_name"] == "g" assert info.value.detail["guardrail_mode"] == GuardrailEventHooks.pre_call + assert data["metadata"] == {"applied_guardrails": ["g"]} def test_handle_pipeline_result_block_does_not_reraise_sensitive_data_route(): @@ -549,14 +551,23 @@ def test_handle_pipeline_result_block_does_not_reraise_sensitive_data_route(): session_id="sess-1", guardrail_name="pii-router", ) + cb = _make_guardrail() + cb.guardrail_name = "pii-router" result = MagicMock() result.terminal_action = "block" result.step_results = [MagicMock(guardrail_name="pii-router")] result.original_exception = original - with pytest.raises(HTTPException) as info: - ProxyLogging._handle_pipeline_result(result=result, data={"model": "m"}, policy_name="p") + data: dict[str, object] = {"model": "m"} + saved = litellm.callbacks + litellm.callbacks = [cb] + try: + with pytest.raises(HTTPException) as info: + ProxyLogging._handle_pipeline_result(result=result, data=data, policy_name="p") + finally: + litellm.callbacks = saved assert info.value.status_code == 400 assert info.value.detail["error"]["type"] == "guardrail_pipeline_error" + assert data["metadata"] == {"applied_guardrails": ["pii-router"]} def test_handle_pipeline_result_block_does_not_reraise_modify_response(): @@ -569,14 +580,23 @@ def test_handle_pipeline_result_block_does_not_reraise_modify_response(): request_data={"model": "m"}, guardrail_name="masker", ) + cb = _make_guardrail() + cb.guardrail_name = "masker" result = MagicMock() result.terminal_action = "block" result.step_results = [MagicMock(guardrail_name="masker")] result.original_exception = original - with pytest.raises(HTTPException) as info: - ProxyLogging._handle_pipeline_result(result=result, data={"model": "m"}, policy_name="p") + data: dict[str, object] = {"model": "m"} + saved = litellm.callbacks + litellm.callbacks = [cb] + try: + with pytest.raises(HTTPException) as info: + ProxyLogging._handle_pipeline_result(result=result, data=data, policy_name="p") + finally: + litellm.callbacks = saved assert info.value.status_code == 400 assert info.value.detail["error"]["type"] == "guardrail_pipeline_error" + assert data["metadata"] == {"applied_guardrails": ["masker"]} def test_handle_pipeline_result_modify_response_raises_modify_exception(): @@ -617,7 +637,7 @@ async def test_run_guardrail_with_metrics_passes_result_and_records_success(monk monkeypatch.setattr(litellm, "callbacks", [prom]) out = await ProxyLogging._run_guardrail_with_metrics( - callback=MagicMock(guardrail_name="g"), coro=task(), hook_type="during_call" + callback=MagicMock(guardrail_name="g"), coro=task(), hook_type="during_call", request_data={} ) assert out == {"a": 1, "b": 2, "c": 3} @@ -643,7 +663,7 @@ async def test_run_guardrail_with_metrics_records_error_and_enriches(monkeypatch monkeypatch.setattr(litellm, "callbacks", [prom]) with pytest.raises(HTTPException): - await ProxyLogging._run_guardrail_with_metrics(callback=cb, coro=task(), hook_type="post_call") + await ProxyLogging._run_guardrail_with_metrics(callback=cb, coro=task(), hook_type="post_call", request_data={}) assert detail["guardrail_name"] == "presidio" recorded = prom._record_guardrail_metrics.call_args.kwargs diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_post_call_success_hook.py b/tests/test_litellm/proxy/utils/proxy_logging/test_post_call_success_hook.py index 715d66db181..53d8948869f 100644 --- a/tests/test_litellm/proxy/utils/proxy_logging/test_post_call_success_hook.py +++ b/tests/test_litellm/proxy/utils/proxy_logging/test_post_call_success_hook.py @@ -6,6 +6,7 @@ from typing import Any from unittest.mock import AsyncMock, MagicMock import pytest +from fastapi import HTTPException import litellm from litellm.integrations.custom_guardrail import CustomGuardrail @@ -96,3 +97,26 @@ async def test_post_call_success_hook_guardrail_returns_modified_response( data={}, response={"orig": True}, user_api_key_dict=make_user_api_key_auth() ) assert out == modified + + +@pytest.mark.asyncio +@pytest.mark.parametrize("run_in_parallel", [False, True], ids=["sequential", "parallel"]) +async def test_post_call_block_names_the_blocking_guardrail_in_applied_guardrails( + proxy_logging, make_user_api_key_auth, monkeypatch, run_in_parallel +): + def _passer_that_records(data, user_api_key_dict, response): + data["metadata"]["applied_guardrails"] = ["passer"] + + passer = _make_guardrail("passer") + passer.async_post_call_success_hook = AsyncMock(side_effect=_passer_that_records) + passer.run_in_parallel = run_in_parallel + blocker = _make_guardrail("blocker") + blocker.async_post_call_success_hook = AsyncMock(side_effect=HTTPException(status_code=400, detail="blocked")) + blocker.run_in_parallel = run_in_parallel + monkeypatch.setattr(litellm, "callbacks", [passer, blocker]) + data = {"model": "m", "metadata": {}} + with pytest.raises(HTTPException): + await proxy_logging.post_call_success_hook( + data=data, response=MagicMock(), user_api_key_dict=make_user_api_key_auth() + ) + assert data["metadata"]["applied_guardrails"] == ["passer", "blocker"] diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_pre_call_hook.py b/tests/test_litellm/proxy/utils/proxy_logging/test_pre_call_hook.py index 6e5cb7fcae3..dbc6fba4ab1 100644 --- a/tests/test_litellm/proxy/utils/proxy_logging/test_pre_call_hook.py +++ b/tests/test_litellm/proxy/utils/proxy_logging/test_pre_call_hook.py @@ -905,3 +905,43 @@ async def test_scan_raw_request_warns_on_in_place_mutation_returning_none( ) mock_logger.warning.assert_called_once() assert "scan_raw_request" in str(mock_logger.warning.call_args) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "blocker_kwargs", + [ + pytest.param({}, id="sequential"), + pytest.param({"scan_raw_request": True}, id="scan_raw_request"), + pytest.param({"run_in_parallel": True}, id="parallel"), + ], +) +async def test_pre_call_block_names_the_blocking_guardrail_in_applied_guardrails( + proxy_logging, make_user_api_key_auth, monkeypatch, blocker_kwargs +): + monkeypatch.setattr(litellm, "callbacks", [_BlockOnSecretGuardrail(**blocker_kwargs)]) + proxy_logging.slack_alerting_instance = MagicMock(alerting=None) + data = _secret_request() + with pytest.raises(HTTPException, match="blocked"): + await proxy_logging.pre_call_hook( + user_api_key_dict=make_user_api_key_auth(), + data=data, + call_type="completion", + ) + assert data["metadata"]["applied_guardrails"] == ["blocker"] + + +@pytest.mark.asyncio +async def test_pre_call_block_keeps_request_declared_guardrail_in_applied_guardrails( + proxy_logging, make_user_api_key_auth, monkeypatch +): + monkeypatch.setattr(litellm, "callbacks", [_BlockOnSecretGuardrail(default_on=False)]) + proxy_logging.slack_alerting_instance = MagicMock(alerting=None) + data = {**_secret_request(), "metadata": {"guardrails": ["blocker", "declared-post-call"]}} + with pytest.raises(HTTPException, match="blocked"): + await proxy_logging.pre_call_hook( + user_api_key_dict=make_user_api_key_auth(), + data=data, + call_type="completion", + ) + assert data["metadata"]["applied_guardrails"] == ["blocker", "declared-post-call"] diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_streaming_hooks.py b/tests/test_litellm/proxy/utils/proxy_logging/test_streaming_hooks.py index ebc831b4102..6fb000b4fa7 100644 --- a/tests/test_litellm/proxy/utils/proxy_logging/test_streaming_hooks.py +++ b/tests/test_litellm/proxy/utils/proxy_logging/test_streaming_hooks.py @@ -20,6 +20,7 @@ from fastapi import HTTPException import litellm from litellm.exceptions import GuardrailRaisedException +from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.anthropic.experimental_pass_through.messages.streaming_iterator import ( @@ -27,6 +28,7 @@ from litellm.llms.anthropic.experimental_pass_through.messages.streaming_iterato ) from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.utils import ProxyLogging +from litellm.types.guardrails import GuardrailEventHooks from litellm.types.utils import Usage @@ -168,6 +170,15 @@ def test_init_response_taking_too_long_task_no_slack_instance_no_error_raises(pr # --------------------------------------------------------------------------- +async def _passthrough_hook(*, response: AsyncIterator[object]) -> AsyncGenerator[object, None]: + async for chunk in response: + yield chunk + + +async def _one_chunk() -> AsyncGenerator[object, None]: + yield "chunk" + + @pytest.mark.asyncio async def test_wrap_streaming_iterator_with_enrichment_passes_through_chunks(proxy_logging): async def gen(): @@ -175,7 +186,9 @@ async def test_wrap_streaming_iterator_with_enrichment_passes_through_chunks(pro yield ch cb = MagicMock(guardrail_name="g", event_hook="pre_call") - wrapped = proxy_logging._wrap_streaming_iterator_with_enrichment(callback=cb, gen=gen()) + wrapped = proxy_logging._wrap_streaming_iterator_with_enrichment( + callback=cb, response=gen(), hook=_passthrough_hook, request_data={} + ) out = [ch async for ch in wrapped] snapshot = { "chunks": out, @@ -195,18 +208,43 @@ async def test_wrap_streaming_iterator_with_enrichment_passes_through_chunks(pro async def test_wrap_streaming_iterator_with_enrichment_enriches_http_exception_raises(proxy_logging): detail = {"error": "blocked"} - async def boom_gen(): + async def boom_hook(*, response: AsyncIterator[object]) -> AsyncGenerator[object, None]: if False: yield # pragma: no cover raise HTTPException(status_code=400, detail=detail) cb = MagicMock(guardrail_name="presidio", event_hook="post_call") - wrapped = proxy_logging._wrap_streaming_iterator_with_enrichment(callback=cb, gen=boom_gen()) + request_data: dict[str, object] = {} + wrapped = proxy_logging._wrap_streaming_iterator_with_enrichment( + callback=cb, response=_one_chunk(), hook=boom_hook, request_data=request_data + ) with pytest.raises(HTTPException): async for _ in wrapped: pass assert detail["guardrail_name"] == "presidio" assert detail["guardrail_mode"] == "post_call" + assert request_data["metadata"]["applied_guardrails"] == ["presidio"] + + +@pytest.mark.asyncio +async def test_wrap_streaming_iterator_leaves_upstream_http_exception_unattributed(proxy_logging): + detail = {"error": "upstream rejected the stream"} + + async def failing_upstream() -> AsyncGenerator[object, None]: + if False: + yield # pragma: no cover + raise HTTPException(status_code=502, detail=detail) + + cb = MagicMock(guardrail_name="presidio", event_hook="post_call") + request_data: dict[str, object] = {} + wrapped = proxy_logging._wrap_streaming_iterator_with_enrichment( + callback=cb, response=failing_upstream(), hook=_passthrough_hook, request_data=request_data + ) + with pytest.raises(HTTPException): + async for _ in wrapped: + pass + assert detail == {"error": "upstream rejected the stream"} + assert request_data == {} # --------------------------------------------------------------------------- @@ -696,3 +734,85 @@ async def test_post_call_response_headers_hook_swallows_callback_error(proxy_log data={}, user_api_key_dict=make_user_api_key_auth(), response=response ) assert out == {} + + +class _StreamBlocker(CustomGuardrail): + def __init__(self, guardrail_name: str = "stream-blocker") -> None: + super().__init__(guardrail_name=guardrail_name, event_hook=GuardrailEventHooks.post_call, default_on=True) + + async def async_post_call_streaming_iterator_hook( + self, user_api_key_dict: UserAPIKeyAuth, response: AsyncIterator[object], request_data: dict[str, object] + ) -> AsyncGenerator[object, None]: + async for _ in response: + raise HTTPException(status_code=400, detail={"error": "blocked"}) + yield # pragma: no cover + + +class _StreamPasser(CustomGuardrail): + def __init__(self, guardrail_name: str = "stream-passer") -> None: + super().__init__(guardrail_name=guardrail_name, event_hook=GuardrailEventHooks.post_call, default_on=True) + + async def async_post_call_streaming_iterator_hook( + self, user_api_key_dict: UserAPIKeyAuth, response: AsyncIterator[object], request_data: dict[str, object] + ) -> AsyncGenerator[object, None]: + async for chunk in response: + yield chunk + + +async def _drain_stream_chain( + proxy_logging: ProxyLogging, + user_api_key_dict: UserAPIKeyAuth, + upstream: AsyncIterator[object], + request_data: dict[str, object], +) -> None: + async for _ in proxy_logging.async_post_call_streaming_iterator_hook( + response=upstream, + user_api_key_dict=user_api_key_dict, + request_data=request_data, + ): + pass + + +async def _failing_provider_stream() -> AsyncGenerator[object, None]: + yield "chunk" + raise RuntimeError("provider connection dropped") + + +@pytest.mark.asyncio +async def test_stream_guardrail_block_names_the_blocking_guardrail_in_applied_guardrails( + proxy_logging, make_user_api_key_auth, monkeypatch +): + monkeypatch.setattr(litellm, "callbacks", [_StreamBlocker()]) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + + request_data: dict[str, object] = {"metadata": {}} + with pytest.raises(HTTPException): + await _drain_stream_chain(proxy_logging, make_user_api_key_auth(), _one_chunk(), request_data) + assert request_data["metadata"]["applied_guardrails"] == ["stream-blocker"] + + +@pytest.mark.asyncio +async def test_stream_block_by_inner_guardrail_does_not_name_the_outer_layers( + proxy_logging, make_user_api_key_auth, monkeypatch +): + monkeypatch.setattr(litellm, "callbacks", [_StreamBlocker(), _StreamPasser("outer-a"), _StreamPasser("outer-b")]) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + + request_data: dict[str, object] = {"metadata": {}} + with pytest.raises(HTTPException) as info: + await _drain_stream_chain(proxy_logging, make_user_api_key_auth(), _one_chunk(), request_data) + assert info.value.detail["guardrail_name"] == "stream-blocker" + assert request_data["metadata"]["applied_guardrails"] == ["stream-blocker"] + + +@pytest.mark.asyncio +async def test_stream_provider_failure_is_not_attributed_to_any_guardrail( + proxy_logging, make_user_api_key_auth, monkeypatch +): + monkeypatch.setattr(litellm, "callbacks", [_StreamPasser("outer-a"), _StreamPasser("outer-b")]) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + + request_data: dict[str, object] = {"metadata": {}} + with pytest.raises(RuntimeError, match="provider connection dropped"): + await _drain_stream_chain(proxy_logging, make_user_api_key_auth(), _failing_provider_stream(), request_data) + assert request_data["metadata"] == {} diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py b/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py index 0ed101952be..c66133ed5f1 100644 --- a/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py @@ -1900,6 +1900,38 @@ class TestToolTransformation: assert "defer_loading" not in result_tool assert "allowed_callers" not in result_tool assert "input_examples" not in result_tool + assert "eager_input_streaming" not in result_tool + + @pytest.mark.parametrize("eager_input_streaming", [True, False]) + def test_transform_function_tools_forwards_eager_input_streaming(self, eager_input_streaming: bool) -> None: + function_tool: Final = { + "type": "function", + "name": "write_file", + "parameters": {"type": "object", "properties": {"path": {"type": "string"}}}, + "eager_input_streaming": eager_input_streaming, + } + + result_tools, _ = LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools( + tools=[function_tool] + ) + + assert result_tools[0]["eager_input_streaming"] is eager_input_streaming + + @pytest.mark.parametrize("eager_input_streaming", [True, False]) + def test_chat_completion_tools_to_responses_tools_keeps_eager_input_streaming( + self, eager_input_streaming: bool + ) -> None: + chat_tool: Final = { + "type": "function", + "function": {"name": "write_file", "parameters": {"type": "object"}}, + "eager_input_streaming": eager_input_streaming, + } + + result_tools: Final = LiteLLMCompletionResponsesConfig.transform_chat_completion_tool_params_to_responses_api_tools( + [chat_tool] + ) + + assert result_tools[0]["eager_input_streaming"] is eager_input_streaming def test_transform_code_execution_tools(self): """Test that code_execution tools are passed through as-is""" @@ -4906,3 +4938,65 @@ class TestStreamingSnapshotItemIds: reasoning_items = _bridged_output_items(completed_event.response, "reasoning") assert len(reasoning_items) == 1 assert reasoning_items[0].id == streamed_event.item_id + + +def test_transform_chat_completion_response_incomplete_details(): + from litellm.types.llms.openai import IncompleteDetails + + resp_length = ModelResponse( + id="resp-length", + choices=[Choices(index=0, finish_reason="length", message=Message(content="cutoff", role="assistant"))], + model="gpt-4o", + ) + result_length = LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( + request_input="test prompt", + responses_api_request={}, + chat_completion_response=resp_length, + ) + assert result_length.status == "incomplete" + assert result_length.incomplete_details is not None + assert result_length.incomplete_details.reason == "max_output_tokens" + + resp_filter = ModelResponse( + id="resp-filter", + choices=[Choices(index=0, finish_reason="content_filter", message=Message(content=None, role="assistant"))], + model="gpt-4o", + ) + result_filter = LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( + request_input="test prompt", + responses_api_request={}, + chat_completion_response=resp_filter, + ) + assert result_filter.status == "incomplete" + assert result_filter.incomplete_details is not None + assert result_filter.incomplete_details.reason == "content_filter" + + resp_refusal = ModelResponse( + id="resp-refusal", + choices=[Choices(index=0, finish_reason="refusal", message=Message(content=None, role="assistant"))], + model="gpt-4o", + ) + result_refusal = LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( + request_input="test prompt", + responses_api_request={}, + chat_completion_response=resp_refusal, + ) + assert result_refusal.status == "incomplete" + assert result_refusal.incomplete_details is not None + assert result_refusal.incomplete_details.reason == "content_filter" + + existing_details = IncompleteDetails(reason="content_filter") + resp_existing = ModelResponse( + id="resp-existing", + choices=[Choices(index=0, finish_reason="length", message=Message(content="cutoff", role="assistant"))], + model="gpt-4o", + ) + resp_existing.incomplete_details = existing_details + result_existing = LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( + request_input="test prompt", + responses_api_request={}, + chat_completion_response=resp_existing, + ) + assert result_existing.status == "incomplete" + assert result_existing.incomplete_details == existing_details + diff --git a/tests/test_litellm/responses/test_responses_websocket_all_providers.py b/tests/test_litellm/responses/test_responses_websocket_all_providers.py index 0394e260606..50fbfb592a5 100644 --- a/tests/test_litellm/responses/test_responses_websocket_all_providers.py +++ b/tests/test_litellm/responses/test_responses_websocket_all_providers.py @@ -1257,6 +1257,252 @@ class TestWebSocketProjectQuotaEnforcement: quota_callback.enforce_project_io_token_quota_for_frame.assert_awaited_once() +def _deployment_defaults(): + from types import MappingProxyType + + from litellm.types.responses.streaming_websocket import ResponsesWebSocketRequestDefaults + + return ResponsesWebSocketRequestDefaults( + fill_missing=MappingProxyType({"reasoning": {"effort": "high"}, "service_tier": "priority"}), + overrides=MappingProxyType({"provider_default": "configured"}), + ) + + +class TestNativeWebSocketDeploymentDefaults: + """The native relay merges deployment litellm_params into every response.create like HTTP does.""" + + def test_builder_maps_router_kwargs_like_the_http_path(self): + from litellm.responses.main import _build_responses_websocket_request_defaults + + defaults = _build_responses_websocket_request_defaults( + { + "model": "gpt-5-pro", + "reasoning_effort": "high", + "service_tier": "priority", + "extra_body": {"provider_default": "configured"}, + "temperature": None, + "timeout": 600, + "max_retries": 2, + "caching": False, + "custom_llm_provider": "openai", + "litellm_metadata": {"user_api_key": "hashed"}, + "user_api_key_dict": MagicMock(), + "litellm_logging_obj": MagicMock(), + "websocket": MagicMock(), + } + ) + + assert dict(defaults.fill_missing) == {"reasoning": {"effort": "high"}, "service_tier": "priority"} + assert dict(defaults.overrides) == {"provider_default": "configured"} + + def test_builder_keeps_explicit_reasoning_over_reasoning_effort(self): + from litellm.responses.main import _build_responses_websocket_request_defaults + + defaults = _build_responses_websocket_request_defaults( + {"model": "gpt-5-pro", "reasoning": {"effort": "low"}, "reasoning_effort": "high"} + ) + + assert dict(defaults.fill_missing) == {"reasoning": {"effort": "low"}} + assert dict(defaults.overrides) == {} + + def test_builder_copies_dict_valued_reasoning_effort_like_the_http_path(self): + from litellm.responses.main import _build_responses_websocket_request_defaults + + defaults = _build_responses_websocket_request_defaults( + {"model": "gpt-5-pro", "reasoning_effort": {"effort": "xhigh", "summary": "auto"}} + ) + + assert dict(defaults.fill_missing) == {"reasoning": {"effort": "xhigh", "summary": "auto"}} + + @pytest.mark.asyncio + async def test_extra_body_type_key_never_replaces_the_frame_type(self): + from types import MappingProxyType + + from litellm.types.responses.streaming_websocket import ResponsesWebSocketRequestDefaults + + handler = _make_streaming( + authorized_model="gpt-5-pro", + request_defaults=ResponsesWebSocketRequestDefaults( + fill_missing=MappingProxyType({}), + overrides=MappingProxyType({"type": "session.update", "provider_default": "configured"}), + ), + ) + + forwarded = json.loads( + await handler._mask_response_create( + json.dumps({"type": "response.create", "model": "gpt-5-pro", "input": "hi"}) + ) + ) + + assert forwarded == { + "type": "response.create", + "model": "gpt-5-pro", + "input": "hi", + "provider_default": "configured", + } + + @pytest.mark.asyncio + async def test_flat_frame_gets_defaults_client_keys_win_extra_body_overrides(self): + handler = _make_streaming(authorized_model="gpt-5-pro", request_defaults=_deployment_defaults()) + + forwarded = json.loads( + await handler._mask_response_create( + json.dumps( + { + "type": "response.create", + "model": "gpt-5-pro", + "input": "Say hello", + "service_tier": "default", + "provider_default": "client", + } + ) + ) + ) + + assert forwarded == { + "type": "response.create", + "model": "gpt-5-pro", + "input": "Say hello", + "service_tier": "default", + "provider_default": "configured", + "reasoning": {"effort": "high"}, + } + + @pytest.mark.asyncio + async def test_nested_response_frame_gets_defaults_inside_response(self): + handler = _make_streaming(authorized_model="gpt-5-pro", request_defaults=_deployment_defaults()) + + forwarded = json.loads( + await handler._mask_response_create( + json.dumps({"type": "response.create", "response": {"model": "gpt-5-pro", "input": "hi"}}) + ) + ) + + assert forwarded == { + "type": "response.create", + "response": { + "model": "gpt-5-pro", + "input": "hi", + "reasoning": {"effort": "high"}, + "service_tier": "priority", + "provider_default": "configured", + }, + } + + @pytest.mark.asyncio + async def test_frames_that_need_nothing_pass_through_untouched(self): + handler = _make_streaming(authorized_model="gpt-5-pro", request_defaults=_deployment_defaults()) + cancel_frame = json.dumps({"type": "response.cancel"}) + complete_frame = json.dumps( + { + "type": "response.create", + "model": "gpt-5-pro", + "input": "hi", + "reasoning": {"effort": "high"}, + "service_tier": "priority", + "provider_default": "configured", + } + ) + + assert await handler._mask_response_create(cancel_frame) is cancel_frame + assert await handler._mask_response_create(complete_frame) is complete_frame + + @pytest.mark.asyncio + async def test_handler_applies_defaults_to_the_first_frame_sent_upstream(self): + import asyncio + from unittest.mock import AsyncMock, patch + + from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler + + class FakeBackend: + def __init__(self): + self.sent = [] + + async def send(self, message): + self.sent.append(message) + + async def recv(self, decode=False): + raise RuntimeError("backend closed") + + async def close(self): + pass + + backend = FakeBackend() + + class FakeConnect: + def __init__(self, url, **kwargs): + pass + + async def __aenter__(self): + return backend + + async def __aexit__(self, *args): + pass + + mock_config = MagicMock(spec=OpenAIResponsesAPIConfig) + mock_config.supports_native_websocket.return_value = True + mock_config.model_in_websocket_url.return_value = True + mock_config.get_websocket_url.return_value = "wss://api.openai.com/v1/responses" + mock_config.validate_environment.return_value = {} + + mock_logging = MagicMock() + mock_logging.pre_call = MagicMock() + mock_logging.dispatch_success_handlers = AsyncMock() + + client_ws = MagicMock() + client_ws.receive_text = AsyncMock(side_effect=RuntimeError("client closed")) + client_ws.send_text = AsyncMock() + client_ws.close = AsyncMock() + + with patch("websockets.connect", FakeConnect): + await BaseLLMHTTPHandler().async_responses_websocket( + model="gpt-5-pro", + websocket=client_ws, + logging_obj=mock_logging, + responses_api_provider_config=mock_config, + api_key="sk-test", + first_message=json.dumps({"type": "response.create", "model": "gpt-5-pro", "input": "Say hello"}), + request_defaults=_deployment_defaults(), + ) + await asyncio.sleep(0) + + assert [json.loads(frame) for frame in backend.sent] == [ + { + "type": "response.create", + "model": "gpt-5-pro", + "input": "Say hello", + "reasoning": {"effort": "high"}, + "service_tier": "priority", + "provider_default": "configured", + } + ] + + @pytest.mark.asyncio + async def test_aresponses_websocket_builds_defaults_from_deployment_kwargs(self, monkeypatch): + import importlib + from unittest.mock import AsyncMock + + responses_main = importlib.import_module("litellm.responses.main") + + stub = MagicMock() + stub.async_responses_websocket = AsyncMock() + monkeypatch.setattr(responses_main, "base_llm_http_handler", stub) + + await responses_main._aresponses_websocket.__wrapped__( + model="openai/gpt-5-pro", + websocket=MagicMock(), + api_key="sk-test", + litellm_logging_obj=MagicMock(), + reasoning_effort="high", + service_tier="priority", + extra_body={"provider_default": "configured"}, + ) + + request_defaults = stub.async_responses_websocket.call_args.kwargs["request_defaults"] + assert dict(request_defaults.fill_missing) == {"reasoning": {"effort": "high"}, "service_tier": "priority"} + assert dict(request_defaults.overrides) == {"provider_default": "configured"} + + class TestNativeWebSocketGuardrails: @pytest.mark.asyncio async def test_response_create_injects_authorized_model(self): diff --git a/tests/test_litellm/rust_bridge/native_route_wheel_test.py b/tests/test_litellm/rust_bridge/native_route_wheel_test.py index 4fa4c0b95ec..0b442f1f269 100644 --- a/tests/test_litellm/rust_bridge/native_route_wheel_test.py +++ b/tests/test_litellm/rust_bridge/native_route_wheel_test.py @@ -73,7 +73,7 @@ def assert_native_request( headers: HTTPMessage, body: object, ) -> None: - if route not in {"transcription", "messages", "chat_completions"}: + if route not in {"transcription", "chat_completions"}: raise AssertionError(f"unexpected route marker: {route!r}") if outcome not in {"success", "429", "hang"}: raise AssertionError(f"unexpected outcome marker: {outcome!r}") @@ -89,10 +89,6 @@ def assert_native_request( assert path == "/v1/messages" assert headers.get("x-api-key") == "sk-native" assert body["model"] == "claude-sonnet-4-5" - if route == "messages": - assert body["max_tokens"] == 16 - assert body["messages"][0]["content"] == "hello-from-messages" - return assert body["max_tokens"] == 17 assert body["messages"][0]["content"] == [{"type": "text", "text": "hello-from-chat"}] @@ -132,17 +128,6 @@ def route_kwargs(route: str, api_base: str, outcome: str) -> dict[str, object]: "language": "en", }, } - if route == "messages": - return common | { - "model": "claude-sonnet-4-5", - "body": { - "model": "claude-sonnet-4-5", - "max_tokens": 16, - "messages": [{"role": "user", "content": "hello-from-messages"}], - }, - "api_key": "sk-native", - "custom_llm_provider": "anthropic", - } if route == "chat_completions": return common | { "model": "anthropic/claude-sonnet-4-5", @@ -165,8 +150,6 @@ def assert_success(route: str, response: object) -> None: def success_value(route: str, response: dict[object, object]) -> object: if route == "transcription": return response["text"] - if route == "messages": - return response["content"][0]["text"] return response["choices"][0]["message"]["content"] @@ -181,7 +164,7 @@ def assert_rate_limit(native: object, route: str, error: BaseException) -> None: def exercise_sync(native: object, api_base: str) -> None: - for route in ("transcription", "messages", "chat_completions"): + for route in ("transcription", "chat_completions"): function: Final = getattr(native, route) assert_success(route, function(**route_kwargs(route, api_base, "success"))) try: @@ -193,7 +176,7 @@ def exercise_sync(native: object, api_base: str) -> None: async def exercise_async(native: object, api_base: str) -> None: - for route in ("transcription", "messages", "chat_completions"): + for route in ("transcription", "chat_completions"): function: Final = getattr(native, f"a{route}") assert_success(route, await function(**route_kwargs(route, api_base, "success"))) try: @@ -206,11 +189,11 @@ async def exercise_async(native: object, api_base: str) -> None: async def exercise_async_concurrency(native: object, api_base: str) -> None: responses: Final = await asyncio.wait_for( - asyncio.gather(*(native.amessages(**route_kwargs("messages", api_base, "success")) for _ in range(32))), + asyncio.gather(*(native.achat_completions(**route_kwargs("chat_completions", api_base, "success")) for _ in range(32))), timeout=15, ) for response in responses: - assert_success("messages", response) + assert_success("chat_completions", response) def exercise_routes(native_path: Path, api_base: str) -> object: @@ -223,8 +206,8 @@ def exercise_routes(native_path: Path, api_base: str) -> object: def exercise_signal(native: object, api_base: str) -> int: try: - native.messages( - **route_kwargs("messages", api_base, "hang"), + native.chat_completions( + **route_kwargs("chat_completions", api_base, "hang"), ) except KeyboardInterrupt: sys.stdout.write("KeyboardInterrupt\n") diff --git a/tests/test_litellm/rust_bridge/test_bindings.py b/tests/test_litellm/rust_bridge/test_bindings.py index 72390b79141..b882a1bb8c2 100644 --- a/tests/test_litellm/rust_bridge/test_bindings.py +++ b/tests/test_litellm/rust_bridge/test_bindings.py @@ -43,8 +43,8 @@ def test_binding_validates_native_attribute( ROUTE_BINDINGS: Final = ( ("completion", chat_completions.NATIVE_COMPLETION), ("acompletion", chat_completions.NATIVE_ACOMPLETION), - ("anthropic_messages_handler", messages.NATIVE_MESSAGES), - ("anthropic_messages", messages.NATIVE_AMESSAGES), + ("messages", messages.NATIVE_MESSAGES), + ("amessages", messages.NATIVE_AMESSAGES), ("responses", responses.NATIVE_RESPONSES), ("aresponses", responses.NATIVE_ARESPONSES), ("ocr", ocr.NATIVE_OCR), diff --git a/tests/test_litellm/rust_bridge/test_catalog.py b/tests/test_litellm/rust_bridge/test_catalog.py index 2c737b0160e..e9fdbf859f4 100644 --- a/tests/test_litellm/rust_bridge/test_catalog.py +++ b/tests/test_litellm/rust_bridge/test_catalog.py @@ -40,6 +40,10 @@ def test_shipped_decisions( enabled: Final = environment == "1" if environment is not None else process is not False assert catalog.rollout(context) is Rollout.RUST_OPT_OUT assert catalog.decision(context) is (Decision.RUST_WITH_FALLBACK if enabled else Decision.PYTHON) + elif route is Route.MESSAGES: + enabled: Final = environment == "1" if environment is not None else process is True + assert catalog.rollout(context) is Rollout.RUST_OPT_IN + assert catalog.decision(context) is (Decision.RUST_WITH_FALLBACK if enabled else Decision.PYTHON) elif route is Route.TRANSCRIPTION and provider == "bedrock": assert catalog.rollout(context) is Rollout.RUST_REQUIRED assert catalog.decision(context) is Decision.RUST_REQUIRED diff --git a/tests/test_litellm/rust_bridge/test_legacy_callbacks.py b/tests/test_litellm/rust_bridge/test_legacy_callbacks.py index a4474c85230..a0906c7c5be 100644 --- a/tests/test_litellm/rust_bridge/test_legacy_callbacks.py +++ b/tests/test_litellm/rust_bridge/test_legacy_callbacks.py @@ -1,12 +1,16 @@ import datetime +import inspect from collections.abc import Mapping +from pathlib import Path from types import MappingProxyType from typing import Final import pytest +from pydantic import TypeAdapter import litellm from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.rust_bridge import legacy_callbacks as legacy from litellm.rust_bridge.legacy_callbacks import check_limits, setup _OCR_KWARGS: Final = MappingProxyType( @@ -56,13 +60,12 @@ def _supplied_logger() -> Logging: ) -def test_setup_adopts_a_supplied_logger_as_caller_owned() -> None: +def test_setup_reuses_a_supplied_logger() -> None: supplied: Final = _supplied_logger() result: Final = setup( "aocr", (), {**_OCR_KWARGS, "litellm_logging_obj": supplied}, datetime.datetime.now(), asynchronous=True ) assert result.logger is supplied - assert result.bridge_owned is False @pytest.mark.parametrize( @@ -73,7 +76,15 @@ def test_setup_adopts_a_supplied_logger_as_caller_owned() -> None: ], ids=["ocr", "embedding"], ) -def test_setup_owns_every_logger_it_builds(call_type: str, kwargs: Mapping[str, object]) -> None: +def test_setup_builds_a_logger_when_none_is_supplied(call_type: str, kwargs: Mapping[str, object]) -> None: result: Final = setup(call_type, (), kwargs, datetime.datetime.now(), asynchronous=True) - assert result.bridge_owned is True assert result.logger.litellm_call_id == result.kwargs["litellm_call_id"] + + +CONTRACT_PATH: Final = Path(__file__).parents[3] / "litellm-rust/crates/callbacks-legacy/python_contract.json" + + +def test_the_rust_contract_matches_the_shim_signatures() -> None: + contract: Final = TypeAdapter(dict[str, list[str]]).validate_json(CONTRACT_PATH.read_text()) + + assert contract == {name: list(inspect.signature(getattr(legacy, name)).parameters) for name in contract} diff --git a/tests/test_litellm/rust_bridge/test_runtime.py b/tests/test_litellm/rust_bridge/test_runtime.py index ade0ae549fb..fa6c0b30413 100644 --- a/tests/test_litellm/rust_bridge/test_runtime.py +++ b/tests/test_litellm/rust_bridge/test_runtime.py @@ -157,7 +157,6 @@ def test_context_outside_rule_stays_on_python() -> None: ( Context(Route.CHAT_COMPLETIONS, provider="anthropic"), Context(Route.CHAT_COMPLETIONS, provider="bedrock"), - Context(Route.MESSAGES, provider="anthropic"), Context(Route.RESPONSES, provider="openai"), Context(Route.TRANSCRIPTION, provider="openai"), ), diff --git a/tests/test_litellm/test_github_close_low_quality_prs.py b/tests/test_litellm/test_github_close_low_quality_prs.py deleted file mode 100644 index 2a891ca72f5..00000000000 --- a/tests/test_litellm/test_github_close_low_quality_prs.py +++ /dev/null @@ -1,856 +0,0 @@ -"""Unit tests for `.github/scripts/close_low_quality_prs.py`. - -These exercise the pure logic (score extraction and per-PR evaluation) without -hitting GitHub. Network/CLI calls are stubbed via monkeypatch. -""" - -from __future__ import annotations - -import datetime as dt -import importlib.util -import sys -from pathlib import Path - -import pytest - -SCRIPT_PATH = ( - Path(__file__).resolve().parents[2] - / ".github" - / "scripts" - / "close_low_quality_prs.py" -) - - -@pytest.fixture(scope="module") -def closer_module(): - """Load the script as a module via its file path (it lives outside the package).""" - spec = importlib.util.spec_from_file_location("close_low_quality_prs", SCRIPT_PATH) - assert spec and spec.loader, f"Could not load spec for {SCRIPT_PATH}" - module = importlib.util.module_from_spec(spec) - sys.modules["close_low_quality_prs"] = module - spec.loader.exec_module(module) - return module - - -def _greptile_comment( - body: str, - updated_at: str = "2026-05-10T00:00:00Z", - login: str = "greptile-apps[bot]", -) -> dict: - return { - "user": {"login": login}, - "body": body, - "created_at": updated_at, - "updated_at": updated_at, - } - - -class TestExtractGreptileScore: - def test_should_extract_score_from_html_header(self, closer_module): - comments = [ - _greptile_comment("

Confidence Score: 3/5

\nSome body text.") - ] - result = closer_module.extract_greptile_score(comments) - assert result is not None - score, _ = result - assert score == 3 - - def test_should_accept_both_greptile_login_variants(self, closer_module): - # REST API form ("greptile-apps[bot]") and GraphQL form ("greptile-apps") - for login in ("greptile-apps", "greptile-apps[bot]"): - comments = [ - _greptile_comment("

Confidence Score: 2/5

", login=login) - ] - result = closer_module.extract_greptile_score(comments) - assert result is not None, f"failed to detect score for login={login}" - score, _ = result - assert score == 2 - - def test_should_extract_score_from_plain_text(self, closer_module): - comments = [_greptile_comment("Confidence Score: 5/5 — looks good!")] - result = closer_module.extract_greptile_score(comments) - assert result is not None - score, _ = result - assert score == 5 - - def test_should_tolerate_whitespace_and_case(self, closer_module): - comments = [_greptile_comment("**confidence score : 2 / 5**")] - result = closer_module.extract_greptile_score(comments) - assert result is not None - score, _ = result - assert score == 2 - - def test_should_pick_most_recent_comment_when_rereview_happens(self, closer_module): - comments = [ - _greptile_comment( - "Confidence Score: 2/5", updated_at="2026-05-01T00:00:00Z" - ), - _greptile_comment( - "Confidence Score: 5/5", updated_at="2026-05-12T00:00:00Z" - ), - ] - result = closer_module.extract_greptile_score(comments) - assert result is not None - score, _ = result - assert score == 5 - - def test_should_ignore_non_greptile_authors(self, closer_module): - comments = [ - { - "user": {"login": "some-human"}, - "body": "Confidence Score: 1/5", - "created_at": "2026-05-12T00:00:00Z", - "updated_at": "2026-05-12T00:00:00Z", - } - ] - assert closer_module.extract_greptile_score(comments) is None - - def test_should_return_none_when_no_score_present(self, closer_module): - comments = [_greptile_comment("Greptile summary without a score.")] - assert closer_module.extract_greptile_score(comments) is None - - def test_should_return_none_for_empty_comments(self, closer_module): - assert closer_module.extract_greptile_score([]) is None - - -class TestEvaluatePr: - @pytest.fixture(autouse=True) - def _now(self): - return dt.datetime(2026, 5, 17, tzinfo=dt.timezone.utc) - - def _make_pr( - self, - *, - number: int = 1, - created_days_ago: int = 10, - is_draft: bool = False, - labels: list[str] | None = None, - author_login: str = "mateo-berri", - ) -> dict: - created = dt.datetime(2026, 5, 17, tzinfo=dt.timezone.utc) - dt.timedelta( - days=created_days_ago - ) - return { - "number": number, - "title": f"PR #{number}", - "createdAt": created.isoformat().replace("+00:00", "Z"), - "isDraft": is_draft, - "labels": [{"name": lbl} for lbl in (labels or [])], - "author": {"login": author_login}, - "url": f"https://example.com/pr/{number}", - } - - @pytest.fixture(autouse=True) - def _external_author(self, closer_module, monkeypatch): - """Treat every test PR as external unless overridden.""" - monkeypatch.setattr( - closer_module, "is_external_pr_author", lambda pr, repo: True - ) - - def test_should_warn_drafts_when_score_low_first_time( - self, closer_module, _now, monkeypatch - ): - # Drafts are NOT a free pass — the open-PR queue should reflect any - # PR that needs human attention regardless of draft status. Authors - # who need a long-lived draft can use the `wip` opt-out label. - # First run: warn the contributor (1-day grace), don't close yet. - monkeypatch.setattr( - closer_module, - "fetch_pr_comments", - lambda *a, **kw: [_greptile_comment("Confidence Score: 2/5")], - ) - action, score, age = closer_module.evaluate_pr( - self._make_pr(is_draft=True, created_days_ago=0), - now=_now, - min_age_days=0, - min_score=4, - repo=None, - optout_labels=set(), - ) - assert action == "warn-grace" - assert score == 2 and age == 0 - - def test_should_warn_brand_new_pr_when_min_age_zero( - self, closer_module, _now, monkeypatch - ): - # `min_age_days=0` means no age filter — a freshly-opened PR is - # eligible the moment Greptile scores it below threshold. The - # first detection still goes through the warn-grace step rather - # than closing immediately, giving the contributor 2 hours to - # respond before the next run actually closes the PR. - monkeypatch.setattr( - closer_module, - "fetch_pr_comments", - lambda *a, **kw: [_greptile_comment("Confidence Score: 1/5")], - ) - action, score, age = closer_module.evaluate_pr( - self._make_pr(created_days_ago=0), - now=_now, - min_age_days=0, - min_score=4, - repo=None, - optout_labels=set(), - ) - assert action == "warn-grace" - assert score == 1 and age == 0 - - def test_should_skip_optout_label_case_insensitive( - self, closer_module, _now, monkeypatch - ): - monkeypatch.setattr( - closer_module, - "fetch_pr_comments", - lambda *a, **kw: pytest.fail("should not fetch comments for opt-outs"), - ) - action, _, _ = closer_module.evaluate_pr( - self._make_pr(labels=["WIP"]), - now=_now, - min_age_days=7, - min_score=4, - repo=None, - optout_labels={"wip"}, - ) - assert action == "skip-optout-label" - - def test_should_skip_too_young_when_min_age_set( - self, closer_module, _now, monkeypatch - ): - # The min-age-days flag is now opt-in (default 0). When a maintainer - # explicitly passes a positive value (e.g. for a backfill run that - # wants to spare brand-new PRs), the skip-too-young path still works. - monkeypatch.setattr( - closer_module, - "fetch_pr_comments", - lambda *a, **kw: pytest.fail("should not fetch comments for young PRs"), - ) - action, _, age = closer_module.evaluate_pr( - self._make_pr(created_days_ago=2), - now=_now, - min_age_days=7, - min_score=4, - repo=None, - optout_labels=set(), - ) - assert action == "skip-too-young" - assert age == 2 - - def test_should_not_skip_when_min_age_is_zero( - self, closer_module, _now, monkeypatch - ): - # With the new default min_age_days=0, even a 0-day-old PR is - # evaluated. This test pins that behavior so future refactors don't - # silently restore an age filter. - monkeypatch.setattr( - closer_module, - "fetch_pr_comments", - lambda *a, **kw: [_greptile_comment("Confidence Score: 5/5")], - ) - action, score, age = closer_module.evaluate_pr( - self._make_pr(created_days_ago=0), - now=_now, - min_age_days=0, - min_score=4, - repo=None, - optout_labels=set(), - ) - assert action == "skip-score-ok" - assert score == 5 and age == 0 - - def test_should_skip_when_greptile_has_not_reviewed( - self, closer_module, _now, monkeypatch - ): - monkeypatch.setattr(closer_module, "fetch_pr_comments", lambda *a, **kw: []) - action, score, age = closer_module.evaluate_pr( - self._make_pr(created_days_ago=10), - now=_now, - min_age_days=7, - min_score=4, - repo=None, - optout_labels=set(), - ) - assert action == "skip-no-greptile-score" - assert score is None and age == 10 - - def test_should_skip_when_score_meets_threshold( - self, closer_module, _now, monkeypatch - ): - monkeypatch.setattr( - closer_module, - "fetch_pr_comments", - lambda *a, **kw: [_greptile_comment("Confidence Score: 4/5")], - ) - action, score, age = closer_module.evaluate_pr( - self._make_pr(created_days_ago=10), - now=_now, - min_age_days=7, - min_score=4, - repo=None, - optout_labels=set(), - ) - assert action == "skip-score-ok" - assert score == 4 and age == 10 - - def test_should_warn_when_old_and_low_score_no_prior_warning( - self, closer_module, _now, monkeypatch - ): - # Even an old PR that still has no grace warning gets one on the - # first eligible run — the daily cron is the natural cadence, so - # an existing-but-never-warned PR enters the grace flow normally. - monkeypatch.setattr( - closer_module, - "fetch_pr_comments", - lambda *a, **kw: [_greptile_comment("Confidence Score: 3/5")], - ) - action, score, age = closer_module.evaluate_pr( - self._make_pr(created_days_ago=10), - now=_now, - min_age_days=7, - min_score=4, - repo=None, - optout_labels=set(), - ) - assert action == "warn-grace" - assert score == 3 and age == 10 - - def test_should_close_when_grace_warning_aged_out_and_score_still_low( - self, closer_module, _now, monkeypatch - ): - # Day-1 the closer posted a warning. Day-2 the PR still scores <4 - # AND the warning is older than `GRACE_PERIOD_SECONDS`, so the - # action flips to `close`. This is the "grace expired" path. - old_warning = { - "user": {"login": "github-actions[bot]"}, - "body": ( - "you have 2 hours to fix this\n\n" + closer_module.GRACE_COMMENT_MARKER - ), - "created_at": ( - _now - dt.timedelta(seconds=closer_module.GRACE_PERIOD_SECONDS + 60) - ) - .isoformat() - .replace("+00:00", "Z"), - "updated_at": "2026-05-15T00:00:00Z", - } - monkeypatch.setattr( - closer_module, - "fetch_pr_comments", - lambda *a, **kw: [ - _greptile_comment( - "

Confidence Score: 1/5

", - updated_at="2026-05-15T00:00:00Z", - ), - old_warning, - ], - ) - action, score, _ = closer_module.evaluate_pr( - self._make_pr(created_days_ago=14), - now=_now, - min_age_days=7, - min_score=4, - repo=None, - optout_labels=set(), - ) - assert action == "close" - assert score == 1 - - def test_should_skip_when_grace_warning_within_window( - self, closer_module, _now, monkeypatch - ): - # Within the 2-hour grace window the closer must NOT close the - # PR even if the score is still low. The warning is only an hour - # old; give the contributor time to push fixes before destruction. - recent_warning = { - "user": {"login": "github-actions[bot]"}, - "body": "warning text\n\n" + closer_module.GRACE_COMMENT_MARKER, - "created_at": (_now - dt.timedelta(hours=1)) - .isoformat() - .replace("+00:00", "Z"), - } - monkeypatch.setattr( - closer_module, - "fetch_pr_comments", - lambda *a, **kw: [ - _greptile_comment("Confidence Score: 2/5"), - recent_warning, - ], - ) - action, score, _ = closer_module.evaluate_pr( - self._make_pr(created_days_ago=10), - now=_now, - min_age_days=0, - min_score=4, - repo=None, - optout_labels=set(), - ) - assert action == "skip-in-grace-period" - assert score == 2 - - def test_should_warn_grace_for_swiftwinds_not_close_immediately( - self, closer_module, _now, monkeypatch - ): - # Regression: SwiftWinds (the dogfood account) used to be in a - # now-removed `IMMEDIATE_CLOSE_LOGINS` bypass that closed on first - # detection. It must now follow the SAME grace path as every other - # external author: warn first, close only after the window elapses. - monkeypatch.setattr( - closer_module, - "fetch_pr_comments", - lambda *a, **kw: [_greptile_comment("Confidence Score: 1/5")], - ) - action, score, _ = closer_module.evaluate_pr( - self._make_pr(created_days_ago=0, author_login="SwiftWinds"), - now=_now, - min_age_days=0, - min_score=4, - repo=None, - optout_labels=set(), - ) - assert action == "warn-grace" - assert score == 1 - - def test_should_skip_internal_authors(self, closer_module, _now, monkeypatch): - # Override the fixture for this one test. - monkeypatch.setattr( - closer_module, "is_external_pr_author", lambda pr, repo: False - ) - monkeypatch.setattr( - closer_module, - "fetch_pr_comments", - lambda *a, **kw: pytest.fail("should not fetch comments for internal"), - ) - action, score, _ = closer_module.evaluate_pr( - self._make_pr(created_days_ago=14, author_login="krrishdholakia"), - now=_now, - min_age_days=7, - min_score=4, - repo=None, - optout_labels=set(), - allowlist=frozenset(), - ) - assert action == "skip-internal" - assert score is None - - -class TestMainOptoutLabelDefault: - """`--optout-label` must REPLACE the canonical defaults, not append.""" - - def _patch_no_op(self, closer_module, monkeypatch): - monkeypatch.setattr(closer_module, "fetch_open_prs", lambda repo: []) - # `optout_labels` is captured indirectly via evaluate_pr; sniff the - # set passed in by stubbing evaluate_pr. - captured: dict = {} - - def fake_evaluate(pr, now, min_age_days, min_score, repo, optout_labels): - captured["optout_labels"] = set(optout_labels) - return ("skip-internal", None, None) - - monkeypatch.setattr(closer_module, "evaluate_pr", fake_evaluate) - return captured - - def test_should_use_canonical_defaults_when_flag_omitted( - self, closer_module, monkeypatch - ): - captured = self._patch_no_op(closer_module, monkeypatch) - # No PRs -> capture won't fire; instead inject one synthetic PR via - # fetch_open_prs so evaluate_pr is invoked at least once. - monkeypatch.setattr( - closer_module, - "fetch_open_prs", - lambda repo: [ - { - "number": 1, - "title": "p", - "createdAt": "2026-05-10T00:00:00Z", - "isDraft": True, - "labels": [], - "author": {"login": "x"}, - } - ], - ) - monkeypatch.setattr(sys, "argv", ["close_low_quality_prs.py"]) - rc = closer_module.main() - assert rc == 0 - assert captured["optout_labels"] == set(closer_module.DEFAULT_OPTOUT_LABELS) - - def test_should_replace_defaults_when_flag_provided( - self, closer_module, monkeypatch - ): - captured = self._patch_no_op(closer_module, monkeypatch) - monkeypatch.setattr( - closer_module, - "fetch_open_prs", - lambda repo: [ - { - "number": 1, - "title": "p", - "createdAt": "2026-05-10T00:00:00Z", - "isDraft": True, - "labels": [], - "author": {"login": "x"}, - } - ], - ) - monkeypatch.setattr( - sys, - "argv", - [ - "close_low_quality_prs.py", - "--optout-label", - "hold", - "--optout-label", - "needs-discussion", - ], - ) - rc = closer_module.main() - assert rc == 0 - # Crucially, none of the canonical defaults leak in. - assert captured["optout_labels"] == {"hold", "needs-discussion"} - for default in closer_module.DEFAULT_OPTOUT_LABELS: - assert default not in captured["optout_labels"], default - - -class TestSecondsSinceLastGraceWarning: - """Grace-period detection: only counts comments by the bot identity - that contain the shared `GRACE_COMMENT_MARKER`.""" - - def _make_marker_comment( - self, - closer_module, - *, - login: str = "github-actions[bot]", - created_at: str = "2026-05-16T00:00:00Z", - include_marker: bool = True, - ) -> dict: - body = "warning text" - if include_marker: - body += "\n\n" + closer_module.GRACE_COMMENT_MARKER - return { - "user": {"login": login}, - "body": body, - "created_at": created_at, - } - - def test_should_return_none_when_no_marker_comment(self, closer_module): - comments = [ - { - "user": {"login": "github-actions[bot]"}, - "body": "Some other bot comment", - "created_at": "2026-05-16T00:00:00Z", - } - ] - assert closer_module.seconds_since_last_grace_warning(comments) is None - - def test_should_return_none_for_empty(self, closer_module): - assert closer_module.seconds_since_last_grace_warning([]) is None - - def test_should_ignore_non_bot_comments_with_marker(self, closer_module): - # If a curious user quotes the marker in a comment, we must NOT - # treat it as a bot warning. The grace timer would then never fire. - comments = [ - self._make_marker_comment(closer_module, login="random-user"), - ] - assert closer_module.seconds_since_last_grace_warning(comments) is None - - def test_should_pick_latest_marker_comment(self, closer_module): - # When multiple grace warnings exist (e.g. a re-open cycle), use - # the most recent one to compute the age. - comments = [ - self._make_marker_comment(closer_module, created_at="2026-05-15T00:00:00Z"), - self._make_marker_comment(closer_module, created_at="2026-05-16T23:00:00Z"), - ] - now = dt.datetime(2026, 5, 17, 0, 0, 0, tzinfo=dt.timezone.utc) - age = closer_module.seconds_since_last_grace_warning(comments, now=now) - # 1h = 3600s - assert age == 3600.0 - - -class TestGraceWarningCommentText: - """Pin the user-facing language in the grace warning comment so the - grace-window and `@greptileai still works after close` promises - don't get accidentally dropped in a future refactor. - """ - - def test_should_state_grace_window(self, closer_module): - body = closer_module.format_grace_warning_comment(score=2, threshold=4) - # The user's PR explicitly said "specify in the comment" — pin - # that the grace window appears in the comment. - assert "2 hours" in body - - def test_should_mention_agent_shin_reconsider(self, closer_module): - body = closer_module.format_grace_warning_comment(score=2, threshold=4) - assert "@agent-shin reconsider" in body - - def test_should_promise_greptileai_works_after_close(self, closer_module): - body = closer_module.format_grace_warning_comment(score=2, threshold=4) - assert "@greptileai" in body - assert "even after the PR is closed" in body - - def test_should_carry_grace_marker(self, closer_module): - # The marker is what `seconds_since_last_grace_warning` greps for - # to detect a prior warning — dropping it would silently break - # the cooldown. - body = closer_module.format_grace_warning_comment(score=2, threshold=4) - assert closer_module.GRACE_COMMENT_MARKER in body - - def test_close_comment_should_mention_greptileai_post_close(self, closer_module): - # The close comment should ALSO point at the @greptileai post-close - # re-review path so contributors see the same options whether they - # read the warning or only catch the close comment. - body = closer_module.format_close_comment(score=2, threshold=4) - assert "@greptileai" in body - assert "even after the PR is closed" in body - - def test_close_comment_should_advertise_reconsider(self, closer_module): - body = closer_module.format_close_comment(score=2, threshold=4) - assert "@agent-shin reconsider" in body - - def test_close_comment_should_carry_agent_shin_close_marker(self, closer_module): - # The close comment advertises `@agent-shin reconsider`, and the - # reconsider reopen guard (`was_closed_by_agent_shin`) only treats a - # PR as Agent-Shin-closed when the close comment carries this marker. - # Dropping it silently breaks the advertised recovery path for every - # PR closed by this daily sweep. - body = closer_module.format_close_comment(score=2, threshold=4) - assert closer_module.AGENT_SHIN_CLOSE_MARKER in body - - def test_close_comment_should_state_score_and_threshold(self, closer_module): - body = closer_module.format_close_comment(score=1, threshold=4) - assert "1/5" in body - assert "4/5" in body - - -class TestHasOptoutLabel: - def test_should_match_label_case_insensitively(self, closer_module): - pr = {"labels": [{"name": "Do Not Close"}, {"name": "bug"}]} - assert closer_module.has_optout_label(pr, {"do not close"}) is True - - def test_should_return_false_when_no_match(self, closer_module): - pr = {"labels": [{"name": "bug"}, {"name": "enhancement"}]} - assert closer_module.has_optout_label(pr, {"wip", "keep open"}) is False - - def test_should_handle_missing_labels(self, closer_module): - assert closer_module.has_optout_label({}, {"wip"}) is False - - -class TestListOpenItemsNoCap: - """The bulk sweeps must fetch the ENTIRE open backlog. - - Regression guard for the old hard-coded ``--limit 1000``: gh lists - newest-first, so a low cap silently dropped the *oldest* PRs/issues — - exactly the stale ones a low-quality sweep exists to catch. - """ - - @staticmethod - def _shared(closer_module): - # `closer_module` loading puts `.github/scripts` on sys.path and - # imports agent_shin_shared, so it's already in sys.modules. - import agent_shin_shared - - return agent_shin_shared - - def _capture_gh_args(self, closer_module, monkeypatch, *, returns="[]"): - shared = self._shared(closer_module) - captured: dict = {} - - def fake_gh(*args): - captured["args"] = args - return returns - - # `list_open_items` looks up `gh` in agent_shin_shared's namespace. - monkeypatch.setattr(shared, "gh", fake_gh) - return shared, captured - - def test_list_open_items_passes_no_cap_limit_not_1000( - self, closer_module, monkeypatch - ): - shared, captured = self._capture_gh_args(closer_module, monkeypatch) - shared.list_open_items("pr", repo="o/r", fields="number,title") - args = captured["args"] - assert "--limit" in args - limit_value = args[args.index("--limit") + 1] - assert limit_value == str(shared.GH_LIST_ALL_LIMIT) - assert limit_value != "1000" - # A meaningful ceiling: comfortably above any realistic open backlog. - assert shared.GH_LIST_ALL_LIMIT >= 100_000 - - def test_list_open_items_uses_dedicated_command_state_and_fields( - self, closer_module, monkeypatch - ): - shared, captured = self._capture_gh_args(closer_module, monkeypatch) - shared.list_open_items("issue", repo="o/r", fields="number") - args = captured["args"] - assert args[0] == "issue" and args[1] == "list" - assert args[args.index("--state") + 1] == "open" - assert args[args.index("--json") + 1] == "number" - assert tuple(args[-2:]) == ("--repo", "o/r") - - def test_list_open_items_omits_repo_when_none(self, closer_module, monkeypatch): - shared, captured = self._capture_gh_args(closer_module, monkeypatch) - shared.list_open_items("pr", repo=None, fields="number") - assert "--repo" not in captured["args"] - - def test_list_open_items_parses_json_array(self, closer_module, monkeypatch): - shared, _ = self._capture_gh_args( - closer_module, monkeypatch, returns='[{"number": 1}, {"number": 2}]' - ) - items = shared.list_open_items("pr", repo=None, fields="number") - assert [i["number"] for i in items] == [1, 2] - - def test_list_open_items_rejects_unknown_kind(self, closer_module): - shared = self._shared(closer_module) - with pytest.raises(ValueError, match="kind must be 'pr' or 'issue', got 'both"): - shared.list_open_items("both", repo="o/r", fields="number") - - def test_fetch_open_prs_delegates_with_no_cap(self, closer_module, monkeypatch): - shared, captured = self._capture_gh_args(closer_module, monkeypatch) - closer_module.fetch_open_prs("o/r") - args = captured["args"] - assert args[0] == "pr" - assert args[args.index("--limit") + 1] == str(shared.GH_LIST_ALL_LIMIT) - # Still requests every field downstream evaluate_pr / labels logic needs. - assert "createdAt" in args[args.index("--json") + 1] - - -class TestEvaluatePrAllowlist: - """While the dogfood allowlist is active `evaluate_pr` only acts on the - named accounts and bypasses the external-only restriction for them. - Emptying it restores the internal-author skip.""" - - @pytest.fixture(autouse=True) - def _now(self): - return dt.datetime(2026, 5, 17, tzinfo=dt.timezone.utc) - - def _make_pr(self, *, author_login: str, created_days_ago: int = 10) -> dict: - created = dt.datetime(2026, 5, 17, tzinfo=dt.timezone.utc) - dt.timedelta( - days=created_days_ago - ) - return { - "number": 1, - "title": "PR #1", - "createdAt": created.isoformat().replace("+00:00", "Z"), - "isDraft": False, - "labels": [], - "author": {"login": author_login}, - "url": "https://example.com/pr/1", - } - - def test_should_skip_author_not_on_allowlist( - self, closer_module, _now, monkeypatch - ): - monkeypatch.setattr( - closer_module, - "fetch_pr_comments", - lambda *a, **kw: pytest.fail("must not fetch comments for non-allowlisted"), - ) - action, score, _ = closer_module.evaluate_pr( - self._make_pr(author_login="random-oss-dev"), - now=_now, - min_age_days=0, - min_score=4, - repo=None, - optout_labels=set(), - ) - assert action == "skip-not-allowlisted" - assert score is None - - def test_should_act_on_allowlisted_internal_author( - self, closer_module, _now, monkeypatch - ): - monkeypatch.setattr( - closer_module, "is_external_pr_author", lambda pr, repo: False - ) - monkeypatch.setattr( - closer_module, - "fetch_pr_comments", - lambda *a, **kw: [_greptile_comment("Confidence Score: 2/5")], - ) - action, score, _ = closer_module.evaluate_pr( - self._make_pr(author_login="mateo-berri", created_days_ago=0), - now=_now, - min_age_days=0, - min_score=4, - repo=None, - optout_labels=set(), - ) - assert action == "warn-grace" - assert score == 2 - - def test_empty_allowlist_restores_internal_skip( - self, closer_module, _now, monkeypatch - ): - monkeypatch.setattr( - closer_module, "is_external_pr_author", lambda pr, repo: False - ) - monkeypatch.setattr( - closer_module, - "fetch_pr_comments", - lambda *a, **kw: pytest.fail("must not fetch comments for internal"), - ) - action, score, _ = closer_module.evaluate_pr( - self._make_pr(author_login="krrishdholakia"), - now=_now, - min_age_days=0, - min_score=4, - repo=None, - optout_labels=set(), - allowlist=frozenset(), - ) - assert action == "skip-internal" - - def test_allowlist_constant_is_the_two_dogfood_accounts(self, closer_module): - assert closer_module.ALLOWLIST_LOGINS == frozenset( - {"mateo-berri", "swiftwinds"} - ) - - -class TestDryRunGateOnClose: - """Regression: the daily sweep is dry-run unless `--close` is passed - (the workflow only adds it when `AGENT_SHIN_ENABLED=true`). A closeable - PR (low score, grace window elapsed) must be DETECTED and reported as - "would close", but the dry run must never make a real GitHub mutation, - so merging Agent Shin stays inert by default.""" - - def _closeable_pr(self) -> dict: - return { - "number": 7, - "title": "thin PR", - "createdAt": "2026-05-10T00:00:00Z", - "isDraft": False, - "labels": [], - "author": {"login": "SwiftWinds"}, - "url": "https://example.com/pr/7", - } - - def test_dry_run_sweep_detects_but_does_not_close( - self, closer_module, monkeypatch, capsys - ): - aged_out_warning = { - "user": {"login": "github-actions[bot]"}, - "body": "warned\n\n" + closer_module.GRACE_COMMENT_MARKER, - # Far enough in the past that it's aged out regardless of - # GRACE_PERIOD_SECONDS, since main() pins `now` to real time. - "created_at": "2020-01-01T00:00:00Z", - } - monkeypatch.setattr( - closer_module, "fetch_open_prs", lambda repo: [self._closeable_pr()] - ) - monkeypatch.setattr( - closer_module, - "fetch_pr_comments", - lambda *a, **kw: [ - _greptile_comment("Confidence Score: 1/5"), - aged_out_warning, - ], - ) - # Any real GitHub mutation during a dry run is the bug under test. - monkeypatch.setattr( - closer_module, - "gh", - lambda *a, **kw: pytest.fail(f"dry run must not call gh: {a}"), - ) - monkeypatch.setattr(sys, "argv", ["close_low_quality_prs.py"]) - - rc = closer_module.main() - - assert rc == 0 - # The PR is detected as closeable, just not acted on. - assert "Total would close: 1" in capsys.readouterr().out diff --git a/tests/test_litellm/test_github_review_gate.py b/tests/test_litellm/test_github_review_gate.py deleted file mode 100644 index 001fa8f43f5..00000000000 --- a/tests/test_litellm/test_github_review_gate.py +++ /dev/null @@ -1,524 +0,0 @@ -"""Unit tests for the `ready for review` label lifecycle (Agent Shin review gate). - -Exercises `triage_with_llm.review_gate`, the state machine that keeps the -`ready for review` label in sync with whether a PR clears both the LLM rubric -and Greptile's confidence score: - - * pass (untagged) -> add label + "ready for review" comment - * pass (untagged, recovered) -> add label + "all clear again" comment - * pass (already tagged) -> noop - * regress (tagged) -> remove label + "what's missing" comment, stays open - * fail (untagged, within 24h)-> one-time "what's missing" notice - * fail (untagged, >24h) -> close + comment - * dry run (close=False) -> would-* previews, no side effects -""" - -from __future__ import annotations - -import datetime as dt -import importlib.util -import sys -from pathlib import Path - -import pytest - -SCRIPT_PATH = ( - Path(__file__).resolve().parents[2] / ".github" / "scripts" / "triage_with_llm.py" -) - -NOW = dt.datetime(2026, 5, 24, 12, 0, 0, tzinfo=dt.timezone.utc) -JUST_NOW = "2026-05-24T11:00:00Z" # 1h old -> within 24h grace -TWO_DAYS_AGO = "2026-05-22T11:00:00Z" # >24h old -> past grace - - -@pytest.fixture(scope="module") -def triage_module(): - spec = importlib.util.spec_from_file_location("triage_with_llm", SCRIPT_PATH) - assert spec and spec.loader - module = importlib.util.module_from_spec(spec) - sys.modules["triage_with_llm"] = module - spec.loader.exec_module(module) - return module - - -class _Recorder: - """Captures every gh mutation review_gate could fire, and fails loudly - on the ones a given scenario forbids.""" - - def __init__(self, triage_module, monkeypatch): - self.comments: list[str] = [] - self.added: list[str] = [] - self.removed: list[str] = [] - self.closed: list[int] = [] - monkeypatch.setattr( - triage_module, - "post_comment", - lambda repo, n, body: self.comments.append(body), - ) - monkeypatch.setattr( - triage_module, - "add_label", - lambda repo, n, label: self.added.append(label), - ) - monkeypatch.setattr( - triage_module, - "remove_label", - lambda repo, n, label: self.removed.append(label), - ) - monkeypatch.setattr( - triage_module, - "close_pr", - lambda repo, n: self.closed.append(n), - ) - - -def _make_pr(**overrides): - base = { - "number": 7, - "title": "feat: do a thing", - "body": "some body without a linked issue or QA proof", - "state": "open", - "author_association": "NONE", - "user": {"login": "mateo-berri"}, - "labels": [], - "created_at": JUST_NOW, - } - base.update(overrides) - return base - - -def _pass(prompt): - return '{"verdict": "pass", "missing": [], "explanation": "looks good"}' - - -def _fail(prompt): - return ( - '{"verdict": "fail", "missing": ["QA proof", "expected vs. actual"],' - ' "explanation": "thin description"}' - ) - - -def _gate(triage_module, **kwargs): - """Call review_gate with safe defaults for the injectable hooks.""" - params = dict( - repo="o/r", - number=7, - close=True, - model="m", - judge=_pass, - greptile_score=None, - comments=[], - now=NOW, - ) - params.update(kwargs) - return triage_module.review_gate(**params) - - -class TestReviewGatePass: - def test_pass_untagged_adds_label_and_ready_comment( - self, triage_module, monkeypatch - ): - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: _make_pr()) - rec = _Recorder(triage_module, monkeypatch) - - result = _gate(triage_module, judge=_pass, greptile_score=5) - - assert result["action"] == "labeled-ready" - assert rec.added == [triage_module.READY_FOR_REVIEW_LABEL] - assert rec.removed == [] and rec.closed == [] - assert len(rec.comments) == 1 - assert "ready for review" in rec.comments[0].lower() - assert triage_module.READY_MARKER in rec.comments[0] - assert "5/5" in rec.comments[0] - - def test_pass_already_tagged_is_noop(self, triage_module, monkeypatch): - pr = _make_pr(labels=[{"name": "ready for review"}]) - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: pr) - rec = _Recorder(triage_module, monkeypatch) - - result = _gate(triage_module, judge=_pass, greptile_score=5) - - assert result["action"] == "noop-passing" - assert rec.added == [] and rec.removed == [] and rec.comments == [] - - def test_pass_after_prior_regression_uses_all_clear_wording( - self, triage_module, monkeypatch - ): - # A regression marker in history -> this is a recovery, not a first pass. - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: _make_pr()) - rec = _Recorder(triage_module, monkeypatch) - prior = [ - { - "user": {"login": "github-actions[bot]"}, - "body": triage_module.REGRESSED_MARKER, - } - ] - - result = _gate(triage_module, judge=_pass, greptile_score=5, comments=prior) - - assert result["action"] == "labeled-ready" - assert "all clear" in rec.comments[0].lower() - - def test_linked_issue_passes_without_calling_judge( - self, triage_module, monkeypatch - ): - pr = _make_pr(body="Fixes #4321\n\nbody") - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: pr) - rec = _Recorder(triage_module, monkeypatch) - - result = _gate( - triage_module, - judge=lambda p: pytest.fail("LLM must not be called for linked issue"), - greptile_score=5, - ) - assert result["action"] == "labeled-ready" - assert rec.added == [triage_module.READY_FOR_REVIEW_LABEL] - - -class TestReviewGateRegression: - def test_regression_removes_label_and_keeps_pr_open( - self, triage_module, monkeypatch - ): - pr = _make_pr(labels=[{"name": "ready for review"}]) - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: pr) - rec = _Recorder(triage_module, monkeypatch) - - result = _gate(triage_module, judge=_fail, greptile_score=5) - - assert result["action"] == "label-removed-regressed" - assert rec.removed == [triage_module.READY_FOR_REVIEW_LABEL] - assert rec.closed == [] # regression NEVER closes the PR - assert triage_module.REGRESSED_MARKER in rec.comments[0] - assert "QA proof" in rec.comments[0] - # The state machine closes a still-failing PR `grace_days` after this - # notice (default 24h); the comment must disclose that deadline rather - # than implying the PR stays open indefinitely. - assert "24 hours" in rec.comments[0] - assert "auto-closed" in rec.comments[0] - - def test_regression_comment_discloses_grace_deadline(self, triage_module): - one_day = triage_module.format_regression_comment( - ["QA proof"], "needs work", grace_days=1 - ) - assert "24 hours" in one_day - assert "auto-closed" in one_day - - three_days = triage_module.format_regression_comment( - ["QA proof"], "needs work", grace_days=3 - ) - assert "3 days" in three_days - assert "auto-closed" in three_days - - def test_greptile_drop_alone_triggers_regression(self, triage_module, monkeypatch): - # Rubric still passes, but Greptile fell to 2/5 -> not passing. - pr = _make_pr(labels=[{"name": "ready for review"}]) - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: pr) - rec = _Recorder(triage_module, monkeypatch) - - result = _gate(triage_module, judge=_pass, greptile_score=2) - - assert result["action"] == "label-removed-regressed" - assert rec.removed == [triage_module.READY_FOR_REVIEW_LABEL] - assert "2/5" in rec.comments[0] - - def test_greptile_score_read_from_comments_when_not_injected( - self, triage_module, monkeypatch - ): - pr = _make_pr(labels=[{"name": "ready for review"}]) - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: pr) - rec = _Recorder(triage_module, monkeypatch) - greptile = [ - { - "user": {"login": "greptile-apps[bot]"}, - "body": "Confidence Score: 2/5", - "created_at": "2026-05-24T10:00:00Z", - } - ] - - result = _gate( - triage_module, - judge=_pass, - greptile_score=triage_module._UNSET, - comments=greptile, - ) - assert result["action"] == "label-removed-regressed" - assert "2/5" in rec.comments[0] - - -class TestReviewGateGraceAndClose: - def test_within_grace_posts_one_time_notice(self, triage_module, monkeypatch): - monkeypatch.setattr( - triage_module, "fetch_pr", lambda repo, n: _make_pr(created_at=JUST_NOW) - ) - rec = _Recorder(triage_module, monkeypatch) - - result = _gate(triage_module, judge=_fail, greptile_score=None) - - assert result["action"] == "within-grace-notified" - assert rec.closed == [] and rec.added == [] and rec.removed == [] - assert triage_module.WITHIN_GRACE_MARKER in rec.comments[0] - assert "QA proof" in rec.comments[0] - - def test_within_grace_does_not_double_notify(self, triage_module, monkeypatch): - monkeypatch.setattr( - triage_module, "fetch_pr", lambda repo, n: _make_pr(created_at=JUST_NOW) - ) - rec = _Recorder(triage_module, monkeypatch) - prior = [ - { - "user": {"login": "github-actions[bot]"}, - "body": triage_module.WITHIN_GRACE_MARKER, - } - ] - - result = _gate(triage_module, judge=_fail, greptile_score=None, comments=prior) - - assert result["action"] == "within-grace-already-notified" - assert rec.comments == [] - - def test_past_grace_closes_with_comment(self, triage_module, monkeypatch): - monkeypatch.setattr( - triage_module, - "fetch_pr", - lambda repo, n: _make_pr(created_at=TWO_DAYS_AGO), - ) - rec = _Recorder(triage_module, monkeypatch) - - result = _gate(triage_module, judge=_fail, greptile_score=None) - - assert result["action"] == "closed" - assert rec.closed == [7] - assert len(rec.comments) == 1 - # The close comment must carry the reconsider provenance marker so - # `was_closed_by_agent_shin` can later recognize this as an Agent Shin - # close (and not some other workflow's `github-actions[bot]` close). - assert triage_module.AGENT_SHIN_CLOSE_MARKER in rec.comments[0] - - def test_recent_regression_marker_blocks_close(self, triage_module, monkeypatch): - """A failing PR with a fresh regression notice must NOT be closed — - the contributor needs a window to address the regression.""" - monkeypatch.setattr( - triage_module, - "fetch_pr", - lambda repo, n: _make_pr(created_at=TWO_DAYS_AGO), - ) - rec = _Recorder(triage_module, monkeypatch) - prior = [ - { - "user": {"login": "github-actions[bot]"}, - "body": triage_module.REGRESSED_MARKER, - # Posted just an hour before NOW -> well inside grace_days. - "created_at": "2026-05-24T11:00:00Z", - } - ] - - result = _gate(triage_module, judge=_fail, greptile_score=None, comments=prior) - - assert result["action"] == "regressed-already-notified" - assert rec.closed == [] and rec.comments == [] - - def test_stale_regression_marker_allows_close(self, triage_module, monkeypatch): - """Once grace_days have elapsed since the regression notice, the - review gate must let the close path fire — otherwise PRs that were - regressed and then abandoned stay open forever.""" - monkeypatch.setattr( - triage_module, - "fetch_pr", - lambda repo, n: _make_pr(created_at=TWO_DAYS_AGO), - ) - rec = _Recorder(triage_module, monkeypatch) - prior = [ - { - "user": {"login": "github-actions[bot]"}, - "body": triage_module.REGRESSED_MARKER, - # Posted 30 days before NOW -> well past the default 1-day grace. - "created_at": "2026-04-24T11:00:00Z", - } - ] - - result = _gate(triage_module, judge=_fail, greptile_score=None, comments=prior) - - assert result["action"] == "closed" - assert rec.closed == [7] - assert len(rec.comments) == 1 - - def test_linked_issue_with_greptile_fail_uses_greptile_explanation( - self, triage_module, monkeypatch - ): - """When the rubric short-circuits to pass (linked-issue regex) but - Greptile dragged the PR under the bar, the close comment's - explanation must describe the Greptile shortfall, not the - misleading "LLM was not called" rubric placeholder.""" - pr = _make_pr(body="Fixes #4321\n\nbody", created_at=TWO_DAYS_AGO) - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: pr) - rec = _Recorder(triage_module, monkeypatch) - - result = _gate( - triage_module, - judge=lambda p: pytest.fail("LLM must not be called for linked issue"), - greptile_score=2, - ) - - assert result["action"] == "closed" - assert len(rec.comments) == 1 - body = rec.comments[0] - assert "LLM was not called" not in body - assert "Greptile" in body and "2/5" in body - - -class TestReviewGateDryRun: - @pytest.mark.parametrize( - "scenario,labels,judge,score,created,expected", - [ - ("pass", [], _pass, 5, JUST_NOW, "would-label-ready"), - ( - "regress", - [{"name": "ready for review"}], - _fail, - 5, - JUST_NOW, - "would-remove-label", - ), - ("within-grace", [], _fail, None, JUST_NOW, "would-notify-within-grace"), - ("past-grace", [], _fail, None, TWO_DAYS_AGO, "would-close"), - ], - ) - def test_dry_run_previews_without_side_effects( - self, - triage_module, - monkeypatch, - scenario, - labels, - judge, - score, - created, - expected, - ): - pr = _make_pr(labels=labels, created_at=created) - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: pr) - rec = _Recorder(triage_module, monkeypatch) - - result = _gate(triage_module, close=False, judge=judge, greptile_score=score) - - assert result["action"] == expected - # Dry run touches nothing. - assert rec.added == [] and rec.removed == [] and rec.closed == [] - assert rec.comments == [] - assert "comment" in result # preview body still surfaced - - -class TestReviewGateGuards: - def test_skips_internal_author(self, triage_module, monkeypatch): - pr = _make_pr(author_association="MEMBER", user={"login": "krrish"}) - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: pr) - result = _gate( - triage_module, - judge=lambda p: pytest.fail("no LLM for internal"), - allowlist=frozenset(), - ) - assert result["action"] == "skip-internal-author" - - def test_skips_closed_pr(self, triage_module, monkeypatch): - pr = _make_pr(state="closed") - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: pr) - result = _gate(triage_module, judge=lambda p: pytest.fail("no LLM for closed")) - assert result["action"] == "skip-not-open" - - def test_llm_error_is_non_destructive(self, triage_module, monkeypatch): - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: _make_pr()) - rec = _Recorder(triage_module, monkeypatch) - - def boom(prompt): - raise RuntimeError("api down") - - result = _gate(triage_module, judge=boom, greptile_score=None) - - assert result["action"] == "skip-llm-error" - assert rec.closed == [] and rec.added == [] and rec.removed == [] - - def test_full_recovery_cycle(self, triage_module, monkeypatch): - """pass -> regress -> recover, threading labels/comments like GitHub would.""" - state = {"labels": [], "comments": []} - - def fake_fetch(repo, n): - return _make_pr(labels=list(state["labels"]), created_at=JUST_NOW) - - monkeypatch.setattr(triage_module, "fetch_pr", fake_fetch) - monkeypatch.setattr( - triage_module, - "post_comment", - lambda repo, n, body: state["comments"].append( - {"user": {"login": "github-actions[bot]"}, "body": body} - ), - ) - monkeypatch.setattr( - triage_module, - "add_label", - lambda repo, n, label: state["labels"].append({"name": label}), - ) - monkeypatch.setattr( - triage_module, - "remove_label", - lambda repo, n, label: state["labels"].clear(), - ) - monkeypatch.setattr( - triage_module, "close_pr", lambda repo, n: pytest.fail("must not close") - ) - - # 1) passes -> tagged - r1 = _gate( - triage_module, judge=_pass, greptile_score=5, comments=state["comments"] - ) - assert r1["action"] == "labeled-ready" - assert any(lbl["name"] == "ready for review" for lbl in state["labels"]) - - # 2) regresses -> tag removed, comment posted, PR still open - r2 = _gate( - triage_module, judge=_fail, greptile_score=2, comments=state["comments"] - ) - assert r2["action"] == "label-removed-regressed" - assert state["labels"] == [] - - # 3) fixed again -> "all clear" + tag back - r3 = _gate( - triage_module, judge=_pass, greptile_score=5, comments=state["comments"] - ) - assert r3["action"] == "labeled-ready" - assert any(lbl["name"] == "ready for review" for lbl in state["labels"]) - assert "all clear" in state["comments"][-1]["body"].lower() - - -class TestReviewGateAllowlist: - """While the dogfood allowlist is active it is the sole author gate: - only the named accounts pass, and for them the internal-author exemption - is bypassed. Emptying it restores the normal internal-author skip.""" - - def test_should_skip_author_not_on_allowlist(self, triage_module, monkeypatch): - pr = _make_pr(user={"login": "random-oss-dev"}) - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: pr) - rec = _Recorder(triage_module, monkeypatch) - result = _gate( - triage_module, judge=lambda p: pytest.fail("no LLM for non-allowlisted") - ) - assert result["action"] == "skip-not-allowlisted" - assert rec.added == [] and rec.comments == [] and rec.closed == [] - - def test_should_act_on_allowlisted_internal_author( - self, triage_module, monkeypatch - ): - pr = _make_pr(author_association="MEMBER", user={"login": "mateo-berri"}) - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: pr) - rec = _Recorder(triage_module, monkeypatch) - result = _gate(triage_module, judge=_pass, greptile_score=5) - assert result["action"] == "labeled-ready" - assert rec.added == [triage_module.READY_FOR_REVIEW_LABEL] - - def test_empty_allowlist_restores_internal_skip(self, triage_module, monkeypatch): - pr = _make_pr(author_association="MEMBER", user={"login": "krrish"}) - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: pr) - result = _gate( - triage_module, - judge=lambda p: pytest.fail("no LLM for internal"), - allowlist=frozenset(), - ) - assert result["action"] == "skip-internal-author" diff --git a/tests/test_litellm/test_github_triage_with_llm.py b/tests/test_litellm/test_github_triage_with_llm.py deleted file mode 100644 index ddffb978b48..00000000000 --- a/tests/test_litellm/test_github_triage_with_llm.py +++ /dev/null @@ -1,2134 +0,0 @@ -"""Unit tests for `.github/scripts/triage_with_llm.py` (Agent Shin).""" - -from __future__ import annotations - -import importlib.util -import json -import sys -from pathlib import Path - -import pytest - -SCRIPT_PATH = ( - Path(__file__).resolve().parents[2] / ".github" / "scripts" / "triage_with_llm.py" -) - - -@pytest.fixture(scope="module") -def triage_module(): - spec = importlib.util.spec_from_file_location("triage_with_llm", SCRIPT_PATH) - assert spec and spec.loader - module = importlib.util.module_from_spec(spec) - sys.modules["triage_with_llm"] = module - spec.loader.exec_module(module) - return module - - -class TestIsInternalContributor: - @pytest.mark.parametrize("association", ["OWNER", "MEMBER", "COLLABORATOR"]) - def test_should_mark_org_associations_as_internal(self, triage_module, association): - item = { - "author_association": association, - "user": {"login": "krrishdholakia"}, - } - assert triage_module.is_internal_contributor(item) is True - - @pytest.mark.parametrize( - "association", - ["CONTRIBUTOR", "FIRST_TIME_CONTRIBUTOR", "FIRST_TIMER", "NONE"], - ) - def test_should_mark_outside_associations_as_external( - self, triage_module, association - ): - item = { - "author_association": association, - "user": {"login": "random-oss-dev"}, - } - assert triage_module.is_internal_contributor(item) is False - - @pytest.mark.parametrize( - "item", - [ - {"author_association": "", "user": {"login": "random-oss-dev"}}, - {"user": {"login": "random-oss-dev"}}, # association field absent - ], - ) - def test_should_fail_safe_when_author_association_is_missing( - self, triage_module, item - ): - # Fail-safe: an empty/missing association must never make a PR - # eligible for the destructive close path. Treat as internal (skip). - assert triage_module.is_internal_contributor(item) is True - - @pytest.mark.parametrize( - "login", - ["dependabot[bot]", "greptile-apps[bot]", "dependabot", "github-actions"], - ) - def test_should_skip_bot_accounts_regardless_of_association( - self, triage_module, login - ): - item = {"author_association": "NONE", "user": {"login": login}} - assert triage_module.is_internal_contributor(item) is True - - -class TestHasLinkedIssue: - @pytest.mark.parametrize( - "body", - [ - "Fixes #1234", - "closes #1", - "Resolves #99", - "fix #42 — this addresses the regression", - "Closes https://github.com/BerriAI/litellm/issues/27000", - "Resolved https://github.com/BerriAI/litellm/issues/27001", - ], - ) - def test_should_detect_common_link_phrases(self, triage_module, body): - assert triage_module.has_linked_issue(body) is True - - @pytest.mark.parametrize( - "body", - [ - "", - "Some change", - # Casual mentions must NOT auto-pass — they should fall through to - # the LLM judge so the stricter "not a passing mention" rule fires. - "See #1234", - "see #1234 for context", - "ref #1234", - "Refs https://github.com/BerriAI/litellm/issues/27000", - "this addresses #1234", - ], - ) - def test_should_not_auto_pass_casual_mentions(self, triage_module, body): - assert triage_module.has_linked_issue(body) is False - - def test_should_not_detect_when_only_html_comment_template(self, triage_module): - body = "" - assert triage_module.has_linked_issue(body) is False - - -class TestStripHtmlComments: - def test_should_remove_single_line_comments(self, triage_module): - text = "before after" - assert "placeholder" not in triage_module.strip_html_comments(text) - - def test_should_remove_multiline_comments(self, triage_module): - text = "kept\n\nkept2" - cleaned = triage_module.strip_html_comments(text) - assert "Fixes #1" not in cleaned - assert "kept" in cleaned and "kept2" in cleaned - - def test_should_handle_none(self, triage_module): - assert triage_module.strip_html_comments(None) == "" - - -class TestCloseCommentText: - """Pin the user-facing language in close comments so changes are intentional.""" - - def test_pr_close_comment_should_recommend_new_pr_primarily(self, triage_module): - body = triage_module.format_pr_close_comment( - {"verdict": "fail", "missing": ["QA proof"], "explanation": "thin"} - ) - # Primary path: open a new PR (because OSS authors can't reopen a - # bot-closed PR). Secondary path: `@agent-shin reconsider`. - assert "Open a new PR" in body - assert "@agent-shin reconsider" in body - # Old advice that no longer works for OSS contributors must NOT - # appear (they can't reopen a PR closed by a bot/maintainer). - assert "Reopen the PR" not in body - - def test_reopen_comment_should_carry_reconsider_marker(self, triage_module): - # The marker is what the rate-limit guard greps for to detect a - # prior reconsider verdict on the same PR. If the marker ever - # gets dropped from this comment, the cooldown silently breaks - # and a contributor can spam `@agent-shin reconsider` to burn - # LLM budget. - body = triage_module.format_reopen_comment("pr") - assert triage_module.RECONSIDER_COMMENT_MARKER in body - - def test_still_failing_comment_should_carry_reconsider_marker(self, triage_module): - body = triage_module.format_reconsider_still_failing_comment( - "pr", - {"verdict": "fail", "missing": ["QA proof"], "explanation": "thin"}, - ) - assert triage_module.RECONSIDER_COMMENT_MARKER in body - - def test_pr_close_comment_should_not_promise_automatic_reopen_on_open( - self, triage_module - ): - # The previous comment said "I'll re-evaluate automatically" — that - # only worked because the author could reopen, which they often - # can't. The new wording must point them at the comment trigger or - # a new PR instead. - body = triage_module.format_pr_close_comment( - {"verdict": "fail", "missing": [], "explanation": ""} - ) - assert "I'll re-evaluate automatically" not in body - - def test_issue_close_comment_should_use_reconsider_trigger(self, triage_module): - # OSS authors have read access, which only lets them reopen issues - # they closed themselves; they CANNOT reopen an issue a maintainer or - # bot closed. So the recovery path is `@agent-shin reconsider` (the - # bot reopens), exactly like the PR path. If this regresses to "reopen - # it yourself", contributors hit a dead end on bot-closed issues. - body = triage_module.format_issue_close_comment( - {"verdict": "fail", "missing": ["repro"], "explanation": "thin"} - ) - assert "@agent-shin reconsider" in body - - def test_pr_close_comment_should_link_blog_explainer(self, triage_module): - # The blog post is the canonical public explanation of what the bot - # checks and why. Every action-required bot comment must link to it - # so contributors landing on a bot-closed PR can self-serve context - # without pinging a maintainer. - body = triage_module.format_pr_close_comment( - {"verdict": "fail", "missing": [], "explanation": ""} - ) - assert "https://docs.litellm.ai/blog/agent-shin-triage" in body - - def test_issue_close_comment_should_link_blog_explainer(self, triage_module): - body = triage_module.format_issue_close_comment( - {"verdict": "fail", "missing": [], "explanation": ""} - ) - assert "https://docs.litellm.ai/blog/agent-shin-triage" in body - - def test_pr_close_comment_should_flag_mocked_tests_as_insufficient_proof( - self, triage_module - ): - # The PR rubric was tightened to require end-to-end QA proof and - # explicitly exclude mocked-dependency unit tests. The user-facing - # close comment must say so — otherwise contributors will keep - # re-submitting "pytest passed (mocks)" runs and getting closed - # again with no explanation of why. - body = triage_module.format_pr_close_comment( - {"verdict": "fail", "missing": [], "explanation": ""} - ) - assert "end-to-end qa proof" in body.lower() - assert "mock" in body.lower() - - def test_issue_recovery_comments_should_name_feature_dead_end_evidence( - self, triage_module - ): - # The feature-request pass bar demands end-to-end evidence of the - # dead-end, so the close and grace-warning recovery bullets must ask - # for it too — otherwise a requester follows those exact instructions - # (description + use case only) and fails `reconsider` again with no - # hint of what else was needed. - verdict = {"verdict": "fail", "missing": [], "explanation": ""} - for body in ( - triage_module.format_issue_close_comment(verdict), - triage_module.format_grace_warning_issue_comment(verdict), - ): - normalized = " ".join(body.split()) - assert "end-to-end evidence of the dead-end" in normalized - assert "showing where the flow stops today" in normalized - - def test_all_agent_shin_comments_should_use_bullet_train_emoji(self, triage_module): - # The bullet train (🚅) is Agent Shin's symbol, matching the LiteLLM - # logo; the previous wave (👋) was generic and didn't match the bot's - # identity. Every action-required comment the bot can post must use the - # bullet train so the contributor recognizes who's writing without - # reading the signoff. - verdict = {"verdict": "fail", "missing": [], "explanation": ""} - comments = { - "pr_close": triage_module.format_pr_close_comment(verdict), - "issue_close": triage_module.format_issue_close_comment(verdict), - "pr_grace": triage_module.format_grace_warning_pr_comment(verdict), - "issue_grace": triage_module.format_grace_warning_issue_comment(verdict), - "within_grace": triage_module.format_within_grace_comment( - [], "", grace_days=1 - ), - } - for name, body in comments.items(): - assert "🚅" in body, f"{name} comment is missing the bullet train emoji" - assert "👋" not in body, f"{name} comment still uses the old wave emoji" - - def test_pr_close_comment_should_show_what_pr_got_right(self, triage_module): - # The user explicitly asked for a "things you got right" section so - # the comment doesn't read as pure rejection. When the judge confirms - # a field is present (e.g. linked_issue), the bullet for it MUST - # appear in the close comment. - body = triage_module.format_pr_close_comment( - { - "verdict": "fail", - "linked_issue": True, - "has_problem_description": True, - "has_expected_vs_actual": False, - "has_qa_proof": False, - "missing": ["QA proof"], - "explanation": "no proof", - } - ) - assert "What you got right" in body - # The two present fields surface as ✅ bullets; the two absent - # fields do not get a ✅ bullet (the QA-proof rubric block still - # mentions the concept, but only the affirmed fields get checkmarks). - assert "- ✅ Linked a related GitHub issue" in body - assert "- ✅ Clear problem description" in body - assert "- ✅ Expected vs. actual behavior" not in body - assert "- ✅ End-to-end QA proof" not in body - - def test_pr_close_comment_should_omit_present_section_when_nothing_present( - self, triage_module - ): - # If the judge says nothing is present (every flag False), the - # "what you got right" block is skipped entirely — better to omit - # than to render "What you got right: (nothing)". - body = triage_module.format_pr_close_comment( - { - "verdict": "fail", - "linked_issue": False, - "has_problem_description": False, - "has_expected_vs_actual": False, - "has_qa_proof": False, - "missing": [], - "explanation": "", - } - ) - assert "What you got right" not in body - - def test_issue_close_comment_should_show_what_issue_got_right(self, triage_module): - # `has_expected_vs_actual` is present, the end-to-end bug evidence is - # not: the "what you got right" block must surface the former and omit - # the latter (no "✅ (nothing)"-style noise for absent items). - body = triage_module.format_issue_close_comment( - { - "verdict": "fail", - "kind": "bug", - "has_repro": False, - "has_expected_vs_actual": True, - "missing": ["end-to-end evidence of the bug"], - "explanation": "no repro shown", - } - ) - assert "What you got right" in body - assert "Expected vs. actual behavior" in body - assert "- ✅ End-to-end evidence of the bug" not in body - - def test_issue_close_comment_should_credit_feature_dead_end_evidence( - self, triage_module - ): - # A feature requester who pasted their dead-end run but skipped the - # motivation must see the evidence credited and only the motivation - # listed as a gap — without a dedicated verdict field the praise - # block could never acknowledge the work they did do. - body = triage_module.format_issue_close_comment( - { - "verdict": "fail", - "kind": "feature", - "has_motivation_example": False, - "has_dead_end_evidence": True, - "missing": ["motivation / use case"], - "explanation": "no use case given", - } - ) - assert "What you got right" in body - assert "- ✅ End-to-end evidence of the dead-end" in body - assert "- ✅ Motivation and concrete example" not in body - - def test_close_comments_should_use_softer_park_for_later_framing( - self, triage_module - ): - # User feedback: the messaging shouldn't feel like punishment. The - # comment must explicitly frame close as a "park this for later," not - # a rejection, and ground that in the queue-hygiene reason. - for body in ( - triage_module.format_pr_close_comment( - {"verdict": "fail", "missing": [], "explanation": ""} - ), - triage_module.format_issue_close_comment( - {"verdict": "fail", "missing": [], "explanation": ""} - ), - ): - assert "park this for later" in body - assert ( - "not a rejection" in body - or "isn't a rejection" in body - or ("isn't us saying" in body) - ) - - def test_only_close_comments_carry_the_agent_shin_close_marker(self, triage_module): - # The reconsider reopen guard keys off AGENT_SHIN_CLOSE_MARKER to tell - # an Agent Shin close from a same-identity close by another workflow. - # That only works if the marker is stamped on the close comments and - # NOT on the grace warnings (which don't close anything). - verdict = {"verdict": "fail", "missing": [], "explanation": ""} - marker = triage_module.AGENT_SHIN_CLOSE_MARKER - assert marker in triage_module.format_pr_close_comment(verdict) - assert marker in triage_module.format_issue_close_comment(verdict) - assert marker not in triage_module.format_grace_warning_pr_comment(verdict) - assert marker not in triage_module.format_grace_warning_issue_comment(verdict) - - -class TestWasClosedByAgentShin: - """Bot-closed guard: only Agent Shin's own closures are reopen candidates.""" - - @staticmethod - def _stub_close_event( - triage_module, - monkeypatch, - *, - actor: str | None, - closed_at: object = "now", - ): - """Stub the most recent `closed` event used by the guard. - - `actor` is the login that closed the item. `closed_at` defaults - to "now" so the marker comment (stubbed at 42s ago) reads as - recent enough relative to the close; tests can pass a concrete - ``datetime`` to simulate older closes (e.g. the stale-marker - regression scenario). - """ - import datetime as real_dt - - if closed_at == "now": - closed_at = real_dt.datetime.now(real_dt.timezone.utc) - monkeypatch.setattr( - triage_module, - "fetch_last_close_event", - lambda repo, n: (actor, closed_at), - ) - - @staticmethod - def _stub_close_marker_present( - triage_module, monkeypatch, *, present: bool, age_seconds: float = 42.0 - ): - """Stub the Agent Shin close-comment marker lookup. - - `was_closed_by_agent_shin` requires the closing actor AND a - recent Agent Shin close comment; these tests pin the latter so - they exercise the actor half in isolation. - """ - monkeypatch.setattr( - triage_module, - "seconds_since_last_agent_shin_close", - lambda *a, **kw: age_seconds if present else None, - ) - - def test_should_return_true_when_bot_closed_and_close_comment_present( - self, triage_module, monkeypatch - ): - self._stub_close_event(triage_module, monkeypatch, actor="github-actions[bot]") - self._stub_close_marker_present(triage_module, monkeypatch, present=True) - assert triage_module.was_closed_by_agent_shin("o/r", 1) is True - - def test_should_return_false_when_bot_closed_but_no_agent_shin_comment( - self, triage_module, monkeypatch - ): - # The `github-actions[bot]` identity is shared across workflows. A - # stale/duplicate sweep closing under that identity must NOT let - # @agent-shin reconsider reopen the item: without an Agent Shin close - # comment the guard fails closed. - self._stub_close_event(triage_module, monkeypatch, actor="github-actions[bot]") - self._stub_close_marker_present(triage_module, monkeypatch, present=False) - assert triage_module.was_closed_by_agent_shin("o/r", 1) is False - - def test_should_return_false_when_last_close_actor_is_maintainer( - self, triage_module, monkeypatch - ): - # A maintainer closed it (e.g. duplicate, security, design). The - # bot must refuse to reopen on @agent-shin reconsider even if an - # earlier Agent Shin close comment is still on the thread. - self._stub_close_event(triage_module, monkeypatch, actor="krrishdholakia") - self._stub_close_marker_present(triage_module, monkeypatch, present=True) - assert triage_module.was_closed_by_agent_shin("o/r", 1) is False - - def test_should_fail_closed_when_no_close_event(self, triage_module, monkeypatch): - # If the events API returns nothing (network blip, repo permission - # quirk), the guard must fail-closed: refuse to reopen rather than - # assume the bot did it. - self._stub_close_event(triage_module, monkeypatch, actor=None, closed_at=None) - self._stub_close_marker_present(triage_module, monkeypatch, present=True) - assert triage_module.was_closed_by_agent_shin("o/r", 1) is False - - def test_should_fail_closed_when_close_event_has_no_timestamp( - self, triage_module, monkeypatch - ): - # Without a usable close timestamp the guard cannot prove the - # marker comment belongs to the latest close; fail-closed. - self._stub_close_event( - triage_module, monkeypatch, actor="github-actions[bot]", closed_at=None - ) - self._stub_close_marker_present(triage_module, monkeypatch, present=True) - assert triage_module.was_closed_by_agent_shin("o/r", 1) is False - - def test_should_return_false_when_marker_predates_latest_close( - self, triage_module, monkeypatch - ): - # Regression for the stale-marker bug: Agent Shin closed once - # (marker stamped), reconsider reopened, and a different workflow - # later closed under the same bot identity without stamping the - # marker. The old marker is still on the thread but does NOT - # belong to the latest close, so reconsider must not reopen. - import datetime as real_dt - - now = real_dt.datetime.now(real_dt.timezone.utc) - # Latest close happened a minute ago. - self._stub_close_event( - triage_module, - monkeypatch, - actor="github-actions[bot]", - closed_at=now - real_dt.timedelta(seconds=60), - ) - # The most recent Agent Shin marker is from an hour ago (a prior - # closed/reopened cycle), which is well outside the skew window. - self._stub_close_marker_present( - triage_module, monkeypatch, present=True, age_seconds=3600.0 - ) - assert triage_module.was_closed_by_agent_shin("o/r", 1) is False - - def test_should_respect_bot_login_override_via_env( - self, triage_module, monkeypatch - ): - # Operators wiring Agent Shin to a PAT (instead of GITHUB_TOKEN) - # can override the expected bot login via env. The guard must - # respect the override so non-default deployments still work. - monkeypatch.setenv("AGENT_SHIN_BOT_LOGIN", "my-bot") - self._stub_close_marker_present(triage_module, monkeypatch, present=True) - self._stub_close_event(triage_module, monkeypatch, actor="my-bot") - assert triage_module.was_closed_by_agent_shin("o/r", 1) is True - # Default "github-actions[bot]" should NOT match when env is set. - self._stub_close_event(triage_module, monkeypatch, actor="github-actions[bot]") - assert triage_module.was_closed_by_agent_shin("o/r", 1) is False - - -class TestSecondsSinceLastAgentShinClose: - """Close-provenance lookup: detects the bot's own auto-close marker.""" - - def _make_comment(self, *, login: str, body: str) -> dict: - return { - "user": {"login": login}, - "body": body, - "created_at": "2026-05-18T05:00:00Z", - } - - def test_should_return_none_when_bot_never_closed(self, triage_module, monkeypatch): - # Comments exist, but none is an Agent Shin close — e.g. only a grace - # warning, or a close by another workflow with no Agent Shin comment. - comments = [ - self._make_comment(login="outside-dev", body="any update?"), - self._make_comment( - login="github-actions[bot]", - body=triage_module.format_grace_warning_pr_comment( - {"verdict": "fail", "missing": [], "explanation": ""} - ), - ), - ] - monkeypatch.setattr( - triage_module, "_iter_paginated_json", lambda *a, **kw: iter(comments) - ) - assert triage_module.seconds_since_last_agent_shin_close("o/r", 1) is None - - def test_should_detect_bot_close_comment(self, triage_module, monkeypatch): - comments = [ - self._make_comment( - login="github-actions[bot]", - body=triage_module.format_pr_close_comment( - {"verdict": "fail", "missing": [], "explanation": ""} - ), - ), - ] - monkeypatch.setattr( - triage_module, "_iter_paginated_json", lambda *a, **kw: iter(comments) - ) - assert triage_module.seconds_since_last_agent_shin_close("o/r", 1) is not None - - def test_should_ignore_non_bot_comment_quoting_marker( - self, triage_module, monkeypatch - ): - # A contributor quoting the hidden marker (GitHub "Quote reply" - # preserves HTML comments) must not be mistaken for a bot close. - comments = [ - self._make_comment( - login="curious-user", - body=f"what is this? {triage_module.AGENT_SHIN_CLOSE_MARKER}", - ), - ] - monkeypatch.setattr( - triage_module, "_iter_paginated_json", lambda *a, **kw: iter(comments) - ) - assert triage_module.seconds_since_last_agent_shin_close("o/r", 1) is None - - -class TestSecondsSinceLastReconsiderVerdict: - """Rate-limit guard: detects the bot's own reconsider verdict marker.""" - - def _make_comment( - self, *, login: str, body: str, created_at: str | None = "2026-05-18T05:00:00Z" - ) -> dict: - comment: dict = {"user": {"login": login}, "body": body} - if created_at is not None: - comment["created_at"] = created_at - return comment - - def test_should_return_none_when_no_bot_reconsider_comments( - self, triage_module, monkeypatch - ): - # An issue with chatter from other users but no bot reconsider - # verdict must not be rate-limited. - comments = [ - self._make_comment(login="outside-dev", body="ping?"), - self._make_comment( - login="github-actions[bot]", body="some other bot message" - ), - ] - monkeypatch.setattr( - triage_module, "_iter_paginated_json", lambda *a, **kw: iter(comments) - ) - assert triage_module.seconds_since_last_reconsider_verdict("o/r", 1) is None - - def test_should_pick_latest_bot_reconsider_marker(self, triage_module, monkeypatch): - # When multiple reconsider verdicts exist, return the AGE of the - # most recent one. Using a frozen reference helps pin the math. - comments = [ - self._make_comment( - login="github-actions[bot]", - body="old verdict " + triage_module.RECONSIDER_COMMENT_MARKER, - created_at="2026-05-18T04:00:00Z", - ), - self._make_comment( - login="github-actions[bot]", - body="newer verdict " + triage_module.RECONSIDER_COMMENT_MARKER, - created_at="2026-05-18T04:55:00Z", - ), - ] - monkeypatch.setattr( - triage_module, "_iter_paginated_json", lambda *a, **kw: iter(comments) - ) - - # Freeze "now" via a tiny shim on the module's `dt` import. - import datetime as real_dt - - class FrozenDateTime(real_dt.datetime): - @classmethod - def now(cls, tz=None): - return real_dt.datetime(2026, 5, 18, 5, 0, 0, tzinfo=tz) - - frozen_module = type(triage_module.dt)("datetime") - frozen_module.datetime = FrozenDateTime - frozen_module.timezone = real_dt.timezone - monkeypatch.setattr(triage_module, "dt", frozen_module) - - age = triage_module.seconds_since_last_reconsider_verdict("o/r", 1) - # newer verdict is 5 minutes (300 seconds) before "now" - assert age == 300.0 - - def test_should_ignore_non_bot_comments_with_marker( - self, triage_module, monkeypatch - ): - # A user comment that happens to quote the marker (e.g. in - # a "what does this hidden marker do?" question) must NOT count. - # The rate-limit guard only trusts comments authored by the bot. - comments = [ - self._make_comment( - login="curious-user", - body=f"Saw this marker: {triage_module.RECONSIDER_COMMENT_MARKER}", - ), - ] - monkeypatch.setattr( - triage_module, "_iter_paginated_json", lambda *a, **kw: iter(comments) - ) - assert triage_module.seconds_since_last_reconsider_verdict("o/r", 1) is None - - def test_should_ignore_bot_comments_without_marker( - self, triage_module, monkeypatch - ): - # The bot posts other things too (Agent Shin close comments, - # CI status, etc.) — only the reconsider-verdict marker should - # arm the cooldown. - comments = [ - self._make_comment( - login="github-actions[bot]", - body="Agent Shin closed this PR (no marker)", - ), - ] - monkeypatch.setattr( - triage_module, "_iter_paginated_json", lambda *a, **kw: iter(comments) - ) - assert triage_module.seconds_since_last_reconsider_verdict("o/r", 1) is None - - -class TestParseVerdict: - def test_should_parse_plain_json(self, triage_module): - raw = '{"verdict": "pass", "missing": []}' - assert triage_module.parse_verdict(raw)["verdict"] == "pass" - - def test_should_strip_markdown_fence(self, triage_module): - raw = '```json\n{"verdict": "fail", "missing": ["foo"]}\n```' - result = triage_module.parse_verdict(raw) - assert result["verdict"] == "fail" - assert result["missing"] == ["foo"] - - def test_should_extract_embedded_json_from_prose(self, triage_module): - raw = 'Here you go: {"verdict": "pass", "missing": []}\nThanks.' - assert triage_module.parse_verdict(raw)["verdict"] == "pass" - - def test_should_raise_for_unparseable_text(self, triage_module): - with pytest.raises(ValueError, match='could not extract JSON from LLM response: not even close to'): - triage_module.parse_verdict("not even close to json") - - def test_should_raise_for_empty(self, triage_module): - with pytest.raises(ValueError, match='empty LLM response'): - triage_module.parse_verdict("") - - -class TestBuildPrompts: - def test_should_include_pr_title_and_body(self, triage_module): - prompt = triage_module.build_pr_prompt( - title="Add foo", body=" Real body" - ) - assert "Add foo" in prompt - assert "Real body" in prompt - assert "comment" not in prompt # HTML comments are stripped - - def test_should_show_empty_marker_for_empty_pr_body(self, triage_module): - prompt = triage_module.build_pr_prompt(title="t", body="") - assert "(empty)" in prompt - - def test_should_include_issue_title_and_body(self, triage_module): - prompt = triage_module.build_issue_prompt(title="Bug", body="repro here") - assert "Bug" in prompt - assert "repro here" in prompt - - def test_issue_bug_rubric_requires_end_to_end_evidence_and_drops_pass_bias( - self, triage_module - ): - # The bug bar was tightened: a report needs the "before" half shown - # end-to-end (video / screenshot / real command output), prose-only - # repro steps no longer pass, and the old "bias toward PASS" leniency - # is gone. If any of these regress, the judge silently goes soft on - # undemonstrated bug reports again. - prompt = triage_module.build_issue_prompt(title="t", body="x") - normalized = " ".join(prompt.split()) - assert "Bias toward PASS when the issue has structure" not in normalized - assert "END-TO-END EVIDENCE OF THE BUG" in normalized - assert "Do not bias toward PASS" in normalized - # The three accepted forms of the "before" demonstration must be named. - assert "screen recording / video" in normalized - assert "screenshot of the bug" in normalized - assert "mocked or stubbed" in normalized - # Prose-only steps are explicitly insufficient now. - assert "steps to reproduce" in normalized - # An unedited issue-form scaffold must not read as evidence: the proof - # field ships with visible headings, so the judge has to be told that - # bare headings with nothing under them count as absent. - assert "unfilled template scaffold" in normalized - assert "counts as absent, not as evidence" in normalized - - def test_issue_feature_rubric_requires_evidence_of_the_dead_end( - self, triage_module - ): - # The feature form asks the requester to walk the ideal flow against a - # live proxy and paste output up to the step that dead-ends, so the - # judge has to demand that evidence, and must not accept an unedited - # scaffold of bare headings as if it were a real attempt. - prompt = triage_module.build_issue_prompt(title="t", body="x") - normalized = " ".join(prompt.split()) - assert "END-TO-END EVIDENCE OF THE DEAD-END" in normalized - assert "showing the point where the flow stops today" in normalized - assert "unfilled template scaffold" in normalized - # The evidence has its own verdict field so feature requesters who - # provided it get credited in "What you got right", exactly like - # `has_repro` credits bug evidence. - assert "`has_dead_end_evidence=true` only when this is present" in normalized - assert '"has_dead_end_evidence": boolean' in normalized - - def test_should_not_crash_when_pr_body_contains_curly_braces(self, triage_module): - """User-supplied content with `{` / `}` must NOT be re-parsed by - `str.format()`. `format` only scans the template literal for - replacement fields; values being substituted in are inserted as - plain strings, so a body like `{"foo": "bar"}` or `{unmatched` - cannot blow up the script. Pinning this here so a future - "improvement" to the templating doesn't reintroduce a crash on - every PR that quotes JSON. - """ - for body in ( - 'Here is some JSON: {"foo": "bar", "n": 1}', - "Half a brace { left dangling, and a stray }", - "Format-spec-looking thing: {0}, {name:>10}, {!r}", - "Nested {a: {b: c}} braces", - ): - pr_prompt = triage_module.build_pr_prompt(title="t", body=body) - issue_prompt = triage_module.build_issue_prompt(title="t", body=body) - assert body in pr_prompt - assert body in issue_prompt - - def test_should_not_crash_when_pr_title_contains_curly_braces(self, triage_module): - title = "Fix bug in {0:>10} format-spec handling" - pr_prompt = triage_module.build_pr_prompt(title=title, body="x") - issue_prompt = triage_module.build_issue_prompt(title=title, body="x") - assert title in pr_prompt - assert title in issue_prompt - - def test_should_preserve_template_indentation_with_multiline_body( - self, triage_module - ): - """`textwrap.dedent` runs on the static template *before* user - content is interpolated, so a multi-line body (whose 2nd+ lines - start at column 0) cannot defeat the common-indent computation - and leave 8-space indentation on every template line. Pin the - dedented shape so the rendered prompt stays consistent for the - LLM judge. - """ - body = "first line\nsecond line at column 0\nthird line at column 0" - for builder in ( - triage_module.build_pr_prompt, - triage_module.build_issue_prompt, - ): - prompt = builder(title="t", body=body) - # Template lines should NOT carry the 8 leading spaces from - # the source-file indentation of the triple-quoted string. - assert " You are " not in prompt - assert 'You are "Agent Shin"' in prompt - assert body in prompt - - -class TestMainModelDefault: - """`--model` falls back to DEFAULT_MODEL even when TRIAGE_MODEL is empty.""" - - def _stub_triage(self, triage_module, monkeypatch): - captured: dict = {} - - def fake_triage(**kwargs): - captured.update(kwargs) - return { - "kind": kwargs["kind"], - "number": kwargs["number"], - "title": "", - "author": "x", - "author_association": "NONE", - "state": "open", - "action": "skip-no-llm-key", - } - - monkeypatch.setattr(triage_module, "triage", fake_triage) - return captured - - def test_should_fall_back_to_default_when_triage_model_env_empty( - self, triage_module, monkeypatch - ): - captured = self._stub_triage(triage_module, monkeypatch) - monkeypatch.setenv("TRIAGE_MODEL", "") - monkeypatch.setattr( - sys, - "argv", - ["triage_with_llm.py", "--repo", "o/r", "--pr", "1"], - ) - rc = triage_module.main() - assert rc == 0 - assert captured["model"] == triage_module.DEFAULT_MODEL - - def test_should_respect_explicit_triage_model_env(self, triage_module, monkeypatch): - captured = self._stub_triage(triage_module, monkeypatch) - monkeypatch.setenv("TRIAGE_MODEL", "gpt-4o-mini") - monkeypatch.setattr( - sys, - "argv", - ["triage_with_llm.py", "--repo", "o/r", "--pr", "1"], - ) - rc = triage_module.main() - assert rc == 0 - assert captured["model"] == "gpt-4o-mini" - - -class TestCallLlmJudge: - """call_llm_judge sets gpt-5 specific kwargs correctly.""" - - def _stub_openai(self, monkeypatch, captured: dict): - """Install a fake `openai.OpenAI` client into sys.modules. - - The fake client records the kwargs passed to chat.completions.create - and returns a minimal response object whose .choices[0].message.content - is "ok". - """ - import types - - class FakeMessage: - content = '{"verdict": "pass"}' - - class FakeChoice: - message = FakeMessage() - - class FakeResponse: - choices = [FakeChoice()] - - class FakeCompletions: - def create(self, **kwargs): - captured.update(kwargs) - return FakeResponse() - - class FakeChat: - completions = FakeCompletions() - - class FakeClient: - def __init__(self, api_key, base_url=None): - captured["__client_kwargs__"] = { - "api_key": api_key, - "base_url": base_url, - } - self.chat = FakeChat() - - fake_module = types.ModuleType("openai") - fake_module.OpenAI = FakeClient - monkeypatch.setitem(sys.modules, "openai", fake_module) - - def test_should_set_reasoning_effort_none_for_gpt5_family( - self, triage_module, monkeypatch - ): - captured: dict = {} - self._stub_openai(monkeypatch, captured) - triage_module.call_llm_judge( - "prompt", model="gpt-5.4-mini", api_key="sk-test", base_url=None - ) - assert captured["model"] == "gpt-5.4-mini" - assert captured["temperature"] == 0 - assert captured["extra_body"] == {"reasoning_effort": "none"} - - def test_should_set_reasoning_effort_for_capitalized_or_dated_gpt5( - self, triage_module, monkeypatch - ): - for model in ("GPT-5.4-mini", "gpt-5.4-mini-2026-03-17", "gpt-5"): - captured: dict = {} - self._stub_openai(monkeypatch, captured) - triage_module.call_llm_judge( - "prompt", model=model, api_key="sk-test", base_url=None - ) - assert captured["extra_body"] == {"reasoning_effort": "none"}, model - - def test_should_omit_reasoning_effort_for_non_gpt5( - self, triage_module, monkeypatch - ): - captured: dict = {} - self._stub_openai(monkeypatch, captured) - triage_module.call_llm_judge( - "prompt", model="gpt-4o-mini", api_key="sk-test", base_url=None - ) - assert "extra_body" not in captured - - def test_should_pass_base_url_when_provided(self, triage_module, monkeypatch): - captured: dict = {} - self._stub_openai(monkeypatch, captured) - triage_module.call_llm_judge( - "p", - model="gpt-5.4-mini", - api_key="sk-test", - base_url="https://proxy.example.com/v1", - ) - assert ( - captured["__client_kwargs__"]["base_url"] == "https://proxy.example.com/v1" - ) - - -class TestTriageOrchestration: - """End-to-end-ish tests that mock both gh fetchers and the LLM.""" - - def _make_pr(self, **overrides): - base = { - "number": 1, - "title": "PR title", - "body": "PR body", - "state": "open", - "author_association": "NONE", - "user": {"login": "mateo-berri"}, - } - base.update(overrides) - return base - - def test_should_skip_internal_author(self, triage_module, monkeypatch): - pr = self._make_pr( - author_association="MEMBER", user={"login": "krrishdholakia"} - ) - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: pr) - - def boom(*a, **kw): - pytest.fail("LLM should not be called for internal authors") - - result = triage_module.triage( - repo="o/r", - kind="pr", - number=1, - close=True, - model="m", - judge=boom, - allowlist=frozenset(), - ) - assert result["action"] == "skip-internal-author" - - def test_should_skip_closed_pr(self, triage_module, monkeypatch): - pr = self._make_pr(state="closed") - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: pr) - result = triage_module.triage( - repo="o/r", - kind="pr", - number=1, - close=True, - model="m", - judge=lambda p: pytest.fail("should not run on closed PRs"), - ) - assert result["action"] == "skip-not-open" - - def test_should_short_circuit_on_linked_issue(self, triage_module, monkeypatch): - pr = self._make_pr(body="Fixes #1234\n\nFoo bar") - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: pr) - result = triage_module.triage( - repo="o/r", - kind="pr", - number=1, - close=True, - model="m", - judge=lambda p: pytest.fail("LLM should not be called"), - ) - assert result["action"] == "pass-linked-issue" - assert result["verdict"]["verdict"] == "pass" - - def test_should_not_short_circuit_on_casual_mention( - self, triage_module, monkeypatch - ): - # "See #1234" is a passing mention, not a closing keyword. The LLM - # must get a chance to apply the stricter rubric. With no prior - # grace warning, the first failing verdict triggers the warning - # path (`would-warn-grace` in dry-run). - pr = self._make_pr(body="See #1234 for context. No QA proof here.") - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: pr) - self._stub_grace_no_warning(triage_module, monkeypatch) - called = {"judge": False} - - def judge(prompt): - called["judge"] = True - return json.dumps( - {"verdict": "fail", "missing": ["QA proof"], "explanation": "thin."} - ) - - result = triage_module.triage( - repo="o/r", - kind="pr", - number=1, - close=False, - model="m", - judge=judge, - ) - assert called["judge"] is True - assert result["action"] == "would-warn-grace" - - def test_should_return_pass_llm_when_judge_passes(self, triage_module, monkeypatch): - pr = self._make_pr(body="Long body, no linked issue.") - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: pr) - captured = {} - - def judge(prompt): - captured["prompt"] = prompt - return json.dumps({"verdict": "pass", "missing": [], "explanation": "ok"}) - - result = triage_module.triage( - repo="o/r", kind="pr", number=1, close=True, model="m", judge=judge - ) - assert result["action"] == "pass-llm" - assert "Long body" in captured["prompt"] - - def test_should_return_would_close_in_dry_run_after_grace_aged_out( - self, triage_module, monkeypatch - ): - # When the grace warning has already aged out (>= GRACE_PERIOD_SECONDS) - # AND the rubric still fails, the dry-run preview returns - # `would-close` so a step-summary writer can render the close - # comment without touching GitHub state. - pr = self._make_pr(body="just a sentence.") - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: pr) - self._stub_grace_aged_out(triage_module, monkeypatch) - - def fake_post(*a, **kw): - pytest.fail("should not post comments in dry-run") - - def fake_close(*a, **kw): - pytest.fail("should not close in dry-run") - - monkeypatch.setattr(triage_module, "post_comment", fake_post) - monkeypatch.setattr(triage_module, "close_pr", fake_close) - - verdict = { - "verdict": "fail", - "missing": ["problem description", "QA proof"], - "explanation": "Body is one sentence.", - } - result = triage_module.triage( - repo="o/r", - kind="pr", - number=1, - close=False, - model="m", - judge=lambda p: json.dumps(verdict), - ) - assert result["action"] == "would-close" - assert result["verdict"]["missing"] == ["problem description", "QA proof"] - - def test_should_post_comment_and_close_after_grace_window( - self, triage_module, monkeypatch - ): - # The "real close" path: --close passed AND the grace warning has - # aged out AND the rubric still fails. The bot posts the close - # comment and closes the PR. - pr = self._make_pr(body="just a sentence.") - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: pr) - self._stub_grace_aged_out(triage_module, monkeypatch) - posted = {} - closed = {} - monkeypatch.setattr( - triage_module, - "post_comment", - lambda repo, n, body: posted.update({"repo": repo, "n": n, "body": body}), - ) - monkeypatch.setattr( - triage_module, - "close_pr", - lambda repo, n: closed.update({"repo": repo, "n": n}), - ) - - verdict = { - "verdict": "fail", - "missing": ["QA proof"], - "explanation": "Body too thin.", - } - result = triage_module.triage( - repo="o/r", - kind="pr", - number=42, - close=True, - model="m", - judge=lambda p: json.dumps(verdict), - ) - assert result["action"] == "closed" - assert posted["n"] == 42 and closed["n"] == 42 - assert "Agent Shin" in posted["body"] - assert "QA proof" in posted["body"] - - def test_should_skip_on_llm_error_in_close_mode(self, triage_module, monkeypatch): - pr = self._make_pr(body="something.") - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: pr) - monkeypatch.setattr( - triage_module, - "post_comment", - lambda *a, **kw: pytest.fail("must not comment on LLM error"), - ) - monkeypatch.setattr( - triage_module, - "close_pr", - lambda *a, **kw: pytest.fail("must not close on LLM error"), - ) - - def broken_judge(prompt): - raise RuntimeError("upstream 500") - - result = triage_module.triage( - repo="o/r", - kind="pr", - number=1, - close=True, - model="m", - judge=broken_judge, - ) - assert result["action"] == "skip-llm-error" - assert "upstream 500" in result["error"] - - def test_should_skip_open_pr_in_reconsider_mode(self, triage_module, monkeypatch): - # Reconsider only makes sense on a CLOSED PR — running it on an open - # one is a no-op (the regular triage flow already evaluated it). - pr = self._make_pr(state="open") - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: pr) - result = triage_module.triage( - repo="o/r", - kind="pr", - number=1, - close=False, - model="m", - judge=lambda p: pytest.fail("should not run on open PR in reconsider"), - reconsider=True, - ) - assert result["action"] == "skip-not-closed" - - @staticmethod - def _stub_reconsider_guards(triage_module, monkeypatch): - """Default reconsider-guard stubs: pretend bot closed + no cooldown. - - The new safety guards (`was_closed_by_agent_shin`, - `seconds_since_last_reconsider_verdict`) hit the GitHub API in - production. Tests that exercise the reconsider happy path stub - them to "yes the bot closed it, no recent reconsider comment" - so the test stays focused on its actual assertion. - """ - monkeypatch.setattr( - triage_module, "was_closed_by_agent_shin", lambda *a, **kw: True - ) - monkeypatch.setattr( - triage_module, - "seconds_since_last_reconsider_verdict", - lambda *a, **kw: None, - ) - - @staticmethod - def _stub_grace_aged_out(triage_module, monkeypatch): - """Pretend the grace warning has aged out. - - For tests that exercise the post-grace close path. Set the age - to twice the grace window so a future tweak to - `GRACE_PERIOD_SECONDS` doesn't accidentally make the stub fall - back inside the window. - """ - monkeypatch.setattr( - triage_module, - "seconds_since_last_grace_warning", - lambda *a, **kw: triage_module.GRACE_PERIOD_SECONDS * 2, - ) - - @staticmethod - def _stub_grace_no_warning(triage_module, monkeypatch): - """Pretend no grace warning has been posted yet (first detection).""" - monkeypatch.setattr( - triage_module, - "seconds_since_last_grace_warning", - lambda *a, **kw: None, - ) - - def test_should_reopen_on_reconsider_pass(self, triage_module, monkeypatch): - # Reconsider on a closed PR with a passing verdict -> reopen + post a - # friendly "re-evaluated" comment. close=True is the production path - # (the workflow only adds --close when AGENT_SHIN_ENABLED=true). - pr = self._make_pr( - state="closed", body="Updated body with QA proof + screenshots." - ) - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: pr) - self._stub_reconsider_guards(triage_module, monkeypatch) - posted = {} - reopened = {} - monkeypatch.setattr( - triage_module, - "post_comment", - lambda repo, n, body: posted.update({"n": n, "body": body}), - ) - monkeypatch.setattr( - triage_module, - "reopen_pr", - lambda repo, n: reopened.update({"n": n}), - ) - # close_pr / close_issue MUST NOT fire in reconsider mode. - monkeypatch.setattr( - triage_module, - "close_pr", - lambda *a, **kw: pytest.fail("must not close on reconsider pass"), - ) - - result = triage_module.triage( - repo="o/r", - kind="pr", - number=42, - close=True, - model="m", - judge=lambda p: json.dumps( - {"verdict": "pass", "missing": [], "explanation": "ok now"} - ), - reconsider=True, - ) - assert result["action"] == "reopened" - assert reopened["n"] == 42 - assert posted["n"] == 42 - assert "reopened" in posted["body"].lower() - - def test_should_dry_run_reconsider_pass_when_close_false( - self, triage_module, monkeypatch - ): - # Reconsider must honor `close=False` (dry-run) just like the - # regular triage flow. A local invocation of - # `python triage_with_llm.py --reconsider --pr N` (no --close) - # must NOT post a comment or reopen the PR — it should return - # `would-reopen` so the operator can preview the outcome. - pr = self._make_pr( - state="closed", body="Updated body with QA proof + screenshots." - ) - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: pr) - self._stub_reconsider_guards(triage_module, monkeypatch) - monkeypatch.setattr( - triage_module, - "post_comment", - lambda *a, **kw: pytest.fail("must not post comment in dry-run reconsider"), - ) - monkeypatch.setattr( - triage_module, - "reopen_pr", - lambda *a, **kw: pytest.fail("must not reopen PR in dry-run reconsider"), - ) - - result = triage_module.triage( - repo="o/r", - kind="pr", - number=42, - close=False, - model="m", - judge=lambda p: json.dumps( - {"verdict": "pass", "missing": [], "explanation": "ok now"} - ), - reconsider=True, - ) - assert result["action"] == "would-reopen" - # The previewed comment body is still returned so a step-summary - # writer can render exactly what would have been posted. - assert "reopened" in result["comment"].lower() - - def test_should_post_still_failing_on_reconsider_fail( - self, triage_module, monkeypatch - ): - pr = self._make_pr(state="closed", body="still empty") - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: pr) - self._stub_reconsider_guards(triage_module, monkeypatch) - posted = {} - monkeypatch.setattr( - triage_module, - "post_comment", - lambda repo, n, body: posted.update({"n": n, "body": body}), - ) - # Neither reopen nor close should fire when reconsider verdict is fail. - monkeypatch.setattr( - triage_module, - "reopen_pr", - lambda *a, **kw: pytest.fail("must not reopen on fail"), - ) - monkeypatch.setattr( - triage_module, - "close_pr", - lambda *a, **kw: pytest.fail("must not close again on reconsider fail"), - ) - - verdict = { - "verdict": "fail", - "missing": ["QA proof"], - "explanation": "Still no QA proof.", - } - result = triage_module.triage( - repo="o/r", - kind="pr", - number=42, - close=True, - model="m", - judge=lambda p: json.dumps(verdict), - reconsider=True, - ) - assert result["action"] == "reconsider-still-failing" - assert posted["n"] == 42 - assert "QA proof" in posted["body"] - - def test_should_not_reopen_on_reconsider_with_ambiguous_verdict( - self, triage_module, monkeypatch - ): - # Regression: only an explicit `pass` verdict reopens. Missing, - # empty, or unexpected verdict strings ("failed", "", garbage) - # must fall through to the still-failing branch rather than - # reopen a PR the rubric did not actually clear. - pr = self._make_pr(state="closed", body="still empty") - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: pr) - self._stub_reconsider_guards(triage_module, monkeypatch) - posted = {} - monkeypatch.setattr( - triage_module, - "post_comment", - lambda repo, n, body: posted.update({"body": body}), - ) - monkeypatch.setattr( - triage_module, - "reopen_pr", - lambda *a, **kw: pytest.fail("must not reopen on ambiguous verdict"), - ) - - for ambiguous in ("", "failed", "needs-info", "unknown"): - posted.clear() - result = triage_module.triage( - repo="o/r", - kind="pr", - number=42, - close=True, - model="m", - judge=lambda p, v=ambiguous: json.dumps( - {"verdict": v, "missing": [], "explanation": "weird"} - ), - reconsider=True, - ) - assert result["action"] == "reconsider-still-failing", ambiguous - assert "body" in posted, ambiguous - - def test_should_dry_run_reconsider_fail_when_close_false( - self, triage_module, monkeypatch - ): - # Mirror dry-run behavior for the FAIL branch — `close=False` - # must NOT post the "still failing" comment. - pr = self._make_pr(state="closed", body="still empty") - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: pr) - self._stub_reconsider_guards(triage_module, monkeypatch) - monkeypatch.setattr( - triage_module, - "post_comment", - lambda *a, **kw: pytest.fail( - "must not post still-failing comment in dry-run" - ), - ) - - verdict = { - "verdict": "fail", - "missing": ["QA proof"], - "explanation": "Still no QA proof.", - } - result = triage_module.triage( - repo="o/r", - kind="pr", - number=42, - close=False, - model="m", - judge=lambda p: json.dumps(verdict), - reconsider=True, - ) - assert result["action"] == "would-reconsider-still-failing" - assert "QA proof" in result["comment"] - - def test_should_reopen_on_reconsider_with_linked_issue_short_circuit( - self, triage_module, monkeypatch - ): - # The linked-issue short-circuit also has to honor reconsider mode: - # if the contributor edited the body to add `Fixes #1234`, the regex - # path should reopen the PR without calling the LLM. - pr = self._make_pr(state="closed", body="Fixes #1234\n\nAddresses the bug.") - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: pr) - self._stub_reconsider_guards(triage_module, monkeypatch) - posted = {} - reopened = {} - monkeypatch.setattr( - triage_module, - "post_comment", - lambda repo, n, body: posted.update({"body": body}), - ) - monkeypatch.setattr( - triage_module, - "reopen_pr", - lambda repo, n: reopened.update({"n": n}), - ) - - result = triage_module.triage( - repo="o/r", - kind="pr", - number=55, - close=True, - model="m", - judge=lambda p: pytest.fail("LLM must not run when linked-issue matches"), - reconsider=True, - ) - assert result["action"] == "reopened" - assert reopened["n"] == 55 - assert "reopened" in posted["body"].lower() - - def test_should_dry_run_reconsider_with_linked_issue_when_close_false( - self, triage_module, monkeypatch - ): - # Linked-issue short-circuit must ALSO honor dry-run. - pr = self._make_pr(state="closed", body="Fixes #1234\n\nAddresses the bug.") - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: pr) - self._stub_reconsider_guards(triage_module, monkeypatch) - monkeypatch.setattr( - triage_module, - "post_comment", - lambda *a, **kw: pytest.fail("must not post in dry-run"), - ) - monkeypatch.setattr( - triage_module, - "reopen_pr", - lambda *a, **kw: pytest.fail("must not reopen in dry-run"), - ) - - result = triage_module.triage( - repo="o/r", - kind="pr", - number=55, - close=False, - model="m", - judge=lambda p: pytest.fail("LLM must not run when linked-issue matches"), - reconsider=True, - ) - assert result["action"] == "would-reopen" - - def test_should_skip_internal_in_reconsider_mode(self, triage_module, monkeypatch): - # Internal authors are exempt from triage in both regular and - # reconsider mode — Agent Shin should never reopen one of their PRs - # automatically, in case a maintainer closed it intentionally. - pr = self._make_pr( - state="closed", - author_association="MEMBER", - user={"login": "krrishdholakia"}, - ) - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: pr) - monkeypatch.setattr( - triage_module, - "reopen_pr", - lambda *a, **kw: pytest.fail("must not reopen for internal author"), - ) - result = triage_module.triage( - repo="o/r", - kind="pr", - number=1, - close=False, - model="m", - judge=lambda p: pytest.fail("LLM must not run for internal author"), - reconsider=True, - allowlist=frozenset(), - ) - assert result["action"] == "skip-internal-author" - - def test_should_skip_reconsider_when_not_bot_closed( - self, triage_module, monkeypatch - ): - # SECURITY: `@agent-shin reconsider` must NOT reopen a PR/issue - # that a MAINTAINER closed for non-rubric reasons (e.g. duplicate, - # design rejection, security report). Only PRs closed by the bot - # itself should ever be candidates for the reconsider reopen path. - pr = self._make_pr(state="closed", body="something.") - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: pr) - monkeypatch.setattr( - triage_module, "was_closed_by_agent_shin", lambda *a, **kw: False - ) - # Even though there's no rate-limit conflict, the bot-closed guard - # alone is sufficient to block. The LLM judge must never run on a - # maintainer-closed PR. - monkeypatch.setattr( - triage_module, - "seconds_since_last_reconsider_verdict", - lambda *a, **kw: None, - ) - monkeypatch.setattr( - triage_module, - "post_comment", - lambda *a, **kw: pytest.fail("must not comment on maintainer-closed PR"), - ) - monkeypatch.setattr( - triage_module, - "reopen_pr", - lambda *a, **kw: pytest.fail("must not reopen maintainer-closed PR"), - ) - - result = triage_module.triage( - repo="o/r", - kind="pr", - number=1, - close=True, - model="m", - judge=lambda p: pytest.fail("LLM must not run before bot-closed guard"), - reconsider=True, - ) - assert result["action"] == "skip-not-bot-closed" - - def test_should_rate_limit_repeated_reconsider_triggers( - self, triage_module, monkeypatch - ): - # COST CONTROL: each `@agent-shin reconsider` event burns CI - # minutes + an OpenAI API call. If the bot already posted a - # reconsider verdict within the cooldown window - # (RECONSIDER_RATE_LIMIT_SECONDS), refuse to run again. This - # bounds the damage from a contributor spamming the trigger. - pr = self._make_pr(state="closed", body="something with new edits.") - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: pr) - monkeypatch.setattr( - triage_module, "was_closed_by_agent_shin", lambda *a, **kw: True - ) - # Pretend the bot posted a reconsider verdict 1 second ago. - monkeypatch.setattr( - triage_module, - "seconds_since_last_reconsider_verdict", - lambda *a, **kw: 1.0, - ) - monkeypatch.setattr( - triage_module, - "post_comment", - lambda *a, **kw: pytest.fail("must not comment during cooldown"), - ) - monkeypatch.setattr( - triage_module, - "reopen_pr", - lambda *a, **kw: pytest.fail("must not reopen during cooldown"), - ) - - result = triage_module.triage( - repo="o/r", - kind="pr", - number=1, - close=True, - model="m", - judge=lambda p: pytest.fail("LLM must not run during cooldown"), - reconsider=True, - ) - assert result["action"] == "skip-rate-limited" - assert result["rate_limit_age_seconds"] == 1.0 - assert ( - result["rate_limit_window_seconds"] - == triage_module.RECONSIDER_RATE_LIMIT_SECONDS - ) - - def test_should_allow_reconsider_after_cooldown_window( - self, triage_module, monkeypatch - ): - # The cooldown is a window, not a one-shot lock — once - # RECONSIDER_RATE_LIMIT_SECONDS has elapsed since the last bot - # verdict, a fresh `@agent-shin reconsider` is allowed through. - pr = self._make_pr(state="closed", body="updated with screenshots now.") - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: pr) - monkeypatch.setattr( - triage_module, "was_closed_by_agent_shin", lambda *a, **kw: True - ) - # Last reconsider was 1 hour ago — well outside the 10-min window. - monkeypatch.setattr( - triage_module, - "seconds_since_last_reconsider_verdict", - lambda *a, **kw: 3600.0, - ) - posted = {} - reopened = {} - monkeypatch.setattr( - triage_module, - "post_comment", - lambda repo, n, body: posted.update({"n": n, "body": body}), - ) - monkeypatch.setattr( - triage_module, - "reopen_pr", - lambda repo, n: reopened.update({"n": n}), - ) - - result = triage_module.triage( - repo="o/r", - kind="pr", - number=1, - close=True, - model="m", - judge=lambda p: json.dumps( - {"verdict": "pass", "missing": [], "explanation": "ok"} - ), - reconsider=True, - ) - assert result["action"] == "reopened" - assert reopened["n"] == 1 - - def test_should_reopen_issue_on_reconsider_pass(self, triage_module, monkeypatch): - issue = { - "number": 7, - "title": "Bug: now with repro", - "body": "## Repro\n```bash\ncurl ...\n```\n\nExpected X, got Y.", - "state": "closed", - "author_association": "NONE", - "user": {"login": "mateo-berri"}, - } - monkeypatch.setattr(triage_module, "fetch_issue", lambda repo, n: issue) - self._stub_reconsider_guards(triage_module, monkeypatch) - posted = {} - reopened = {} - monkeypatch.setattr( - triage_module, - "post_comment", - lambda repo, n, body: posted.update({"body": body}), - ) - monkeypatch.setattr( - triage_module, - "reopen_issue", - lambda repo, n: reopened.update({"n": n}), - ) - - result = triage_module.triage( - repo="o/r", - kind="issue", - number=7, - close=True, - model="m", - judge=lambda p: json.dumps( - {"verdict": "pass", "missing": [], "explanation": "now reproducible"} - ), - reconsider=True, - ) - assert result["action"] == "reopened" - assert reopened["n"] == 7 - assert "reopened" in posted["body"].lower() - - def test_should_triage_issues_kind(self, triage_module, monkeypatch): - issue = { - "number": 7, - "title": "Bug: X is broken", - "body": "no detail", - "state": "open", - "author_association": "NONE", - "user": {"login": "mateo-berri"}, - } - monkeypatch.setattr(triage_module, "fetch_issue", lambda repo, n: issue) - # Grace already aged out -> close path. (Issues use the same - # GRACE_COMMENT_MARKER detection as PRs.) - self._stub_grace_aged_out(triage_module, monkeypatch) - closed = {} - posted = {} - monkeypatch.setattr( - triage_module, - "post_comment", - lambda repo, n, body: posted.update(body=body), - ) - monkeypatch.setattr( - triage_module, "close_issue", lambda repo, n: closed.update(n=n) - ) - - verdict = { - "verdict": "fail", - "kind": "bug", - "has_repro": False, - "missing": ["reproduction", "expected vs. actual"], - "explanation": "No repro provided.", - } - result = triage_module.triage( - repo="o/r", - kind="issue", - number=7, - close=True, - model="m", - judge=lambda p: json.dumps(verdict), - ) - assert result["action"] == "closed" - assert closed["n"] == 7 - assert "reproduction" in posted["body"] - - # ---- Grace-period flow ------------------------------------------------ - - def test_should_post_grace_warning_on_first_failing_run_in_close_mode( - self, triage_module, monkeypatch - ): - # First low-quality detection -> bot posts a warning comment with - # the GRACE_COMMENT_MARKER. The PR must NOT be closed yet. - pr = self._make_pr(body="just a sentence.") - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: pr) - self._stub_grace_no_warning(triage_module, monkeypatch) - posted = {} - monkeypatch.setattr( - triage_module, - "post_comment", - lambda repo, n, body: posted.update({"n": n, "body": body}), - ) - monkeypatch.setattr( - triage_module, - "close_pr", - lambda *a, **kw: pytest.fail("must not close on first detection"), - ) - - verdict = { - "verdict": "fail", - "missing": ["QA proof"], - "explanation": "Body too thin.", - } - result = triage_module.triage( - repo="o/r", - kind="pr", - number=42, - close=True, - model="m", - judge=lambda p: json.dumps(verdict), - ) - assert result["action"] == "warned-grace" - assert posted["n"] == 42 - # Pin the user-facing language pieces the user explicitly asked for. - assert "2 hours" in posted["body"] - assert "@agent-shin reconsider" in posted["body"] - assert "@greptileai" in posted["body"] - assert "even after the PR is closed" in posted["body"] - assert triage_module.GRACE_COMMENT_MARKER in posted["body"] - - def test_should_skip_close_inside_grace_window(self, triage_module, monkeypatch): - # A warning was posted recently; do nothing on this run regardless - # of close=True. The next run after `GRACE_PERIOD_SECONDS` elapses - # is the one that flips to actual close. - pr = self._make_pr(body="just a sentence.") - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: pr) - monkeypatch.setattr( - triage_module, - "seconds_since_last_grace_warning", - lambda *a, **kw: 60.0, - ) - monkeypatch.setattr( - triage_module, - "post_comment", - lambda *a, **kw: pytest.fail("must not comment during grace window"), - ) - monkeypatch.setattr( - triage_module, - "close_pr", - lambda *a, **kw: pytest.fail("must not close during grace window"), - ) - - verdict = { - "verdict": "fail", - "missing": ["QA proof"], - "explanation": "Body too thin.", - } - result = triage_module.triage( - repo="o/r", - kind="pr", - number=42, - close=True, - model="m", - judge=lambda p: json.dumps(verdict), - ) - assert result["action"] == "skip-in-grace-period" - assert result["grace_age_seconds"] == 60.0 - assert result["grace_period_seconds"] == triage_module.GRACE_PERIOD_SECONDS - - def test_should_dry_run_grace_warning_when_close_false( - self, triage_module, monkeypatch - ): - # In dry-run mode the FIRST failing detection returns - # `would-warn-grace` (with the previewed comment body) and never - # touches GitHub state. Lets a local operator preview the - # warning before flipping --close on. - pr = self._make_pr(body="thin") - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: pr) - self._stub_grace_no_warning(triage_module, monkeypatch) - monkeypatch.setattr( - triage_module, - "post_comment", - lambda *a, **kw: pytest.fail("must not post in dry-run grace warn"), - ) - - verdict = { - "verdict": "fail", - "missing": ["QA proof"], - "explanation": "thin", - } - result = triage_module.triage( - repo="o/r", - kind="pr", - number=1, - close=False, - model="m", - judge=lambda p: json.dumps(verdict), - ) - assert result["action"] == "would-warn-grace" - assert "2 hours" in result["comment"] - - def test_should_warn_grace_for_swiftwinds_not_close_instantly( - self, triage_module, monkeypatch - ): - # Regression: SwiftWinds (the dogfood account) used to be in a - # now-removed `IMMEDIATE_CLOSE_LOGINS` bypass that skipped the grace - # window and closed on first detection. It must follow the SAME - # grace path as every other author: warn first, close only after the - # window elapses. A re-added instant-close bypass would call - # close_pr here and fail the test. - pr = self._make_pr(body="just a sentence.", user={"login": "SwiftWinds"}) - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: pr) - self._stub_grace_no_warning(triage_module, monkeypatch) - posted = {} - monkeypatch.setattr( - triage_module, - "post_comment", - lambda repo, n, body: posted.update({"n": n, "body": body}), - ) - monkeypatch.setattr( - triage_module, - "close_pr", - lambda *a, **kw: pytest.fail( - "SwiftWinds must not close on first detection; it gets the grace window" - ), - ) - - verdict = { - "verdict": "fail", - "missing": ["QA proof"], - "explanation": "Body too thin.", - } - result = triage_module.triage( - repo="o/r", - kind="pr", - number=99, - close=True, - model="m", - judge=lambda p: json.dumps(verdict), - ) - assert result["action"] == "warned-grace" - assert "2 hours" in posted["body"] - - -class TestGraceWarningCommentText: - """Pin the user-facing promises in the grace warning so a future - refactor can't silently drop them.""" - - def test_pr_grace_warning_should_state_grace_window(self, triage_module): - body = triage_module.format_grace_warning_pr_comment( - {"verdict": "fail", "missing": ["QA proof"], "explanation": "thin"} - ) - # The user explicitly asked: "specify in the comment" the grace window. - assert "2 hours" in body - - def test_pr_grace_warning_should_mention_reconsider_during_grace( - self, triage_module - ): - body = triage_module.format_grace_warning_pr_comment( - {"verdict": "fail", "missing": [], "explanation": ""} - ) - assert "@agent-shin reconsider" in body - - def test_pr_grace_warning_should_promise_greptileai_works_post_close( - self, triage_module - ): - body = triage_module.format_grace_warning_pr_comment( - {"verdict": "fail", "missing": [], "explanation": ""} - ) - # Per user: comment should state @greptileai works even after close. - assert "@greptileai" in body - assert "even after the PR is closed" in body - - def test_pr_grace_warning_should_carry_grace_marker(self, triage_module): - # The marker is what `seconds_since_last_grace_warning` greps for - # on subsequent runs to detect that a warning has been posted. - # Dropping it would silently break the close-after-grace path. - body = triage_module.format_grace_warning_pr_comment( - {"verdict": "fail", "missing": [], "explanation": ""} - ) - assert triage_module.GRACE_COMMENT_MARKER in body - - def test_issue_grace_warning_should_carry_grace_marker(self, triage_module): - body = triage_module.format_grace_warning_issue_comment( - {"verdict": "fail", "missing": [], "explanation": ""} - ) - assert triage_module.GRACE_COMMENT_MARKER in body - assert "2 hours" in body - # OSS authors can't reopen a bot-closed issue, so recovery is - # `@agent-shin reconsider` (the bot reopens), like the PR path. - assert "@agent-shin reconsider" in body - - def test_pr_close_comment_should_promise_greptileai_works_post_close( - self, triage_module - ): - # The standard close comment must ALSO point at @greptileai so - # contributors see the same options whether they read the warning - # or only catch the close comment. - body = triage_module.format_pr_close_comment( - {"verdict": "fail", "missing": [], "explanation": ""} - ) - assert "@greptileai" in body - assert "even after the PR is closed" in body - - def test_pr_grace_warning_should_not_prompt_reconsider_during_grace_window( - self, triage_module - ): - # Per user feedback: during the 24h grace window, the contributor - # should just update the PR description. Asking them to also comment - # "@agent-shin reconsider" right away adds a step they don't need — - # the bot re-checks automatically on the next sweep. The reconsider - # trigger is reserved for the post-close recovery path. - # - # We pin this by checking that the grace section explicitly tells - # the contributor they don't need to ping the bot during the grace - # window. The presence of "@agent-shin reconsider" elsewhere in the - # comment (as the post-close path) is fine and required by other - # tests. - body = triage_module.format_grace_warning_pr_comment( - {"verdict": "fail", "missing": [], "explanation": ""} - ) - assert "No need to ping" in body or "no need to ping" in body - - def test_grace_warnings_should_show_what_got_right(self, triage_module): - # The "What you got right" section must appear in the grace warning - # too, not only the close comment — the contributor sees the warning - # first and that's their best chance to know what to keep. - pr_body = triage_module.format_grace_warning_pr_comment( - { - "verdict": "fail", - "linked_issue": True, - "has_problem_description": True, - "has_expected_vs_actual": True, - "has_qa_proof": False, - "missing": ["QA proof"], - "explanation": "thin", - } - ) - assert "What you got right" in pr_body - assert "Linked a related GitHub issue" in pr_body - - issue_body = triage_module.format_grace_warning_issue_comment( - { - "verdict": "fail", - "kind": "feature", - "has_motivation_example": True, - "missing": ["concrete description"], - "explanation": "vague", - } - ) - assert "What you got right" in issue_body - assert "Motivation and concrete example" in issue_body - - def test_grace_warnings_should_use_softer_park_for_later_framing( - self, triage_module - ): - # Same softer-framing pin as the close comment, but for the warning - # — the contributor's first contact with the bot must not read as a - # hard deadline / ultimatum. - for body in ( - triage_module.format_grace_warning_pr_comment( - {"verdict": "fail", "missing": [], "explanation": ""} - ), - triage_module.format_grace_warning_issue_comment( - {"verdict": "fail", "missing": [], "explanation": ""} - ), - ): - assert "park this for later" in body - assert ( - "not a rejection" in body - or "isn't a rejection" in body - or ("isn't us saying" in body) - ) - - -class TestSecondsSinceLastGraceWarning: - """Mirror of TestSecondsSinceLastReconsiderVerdict for the new helper. - Both helpers share `_seconds_since_latest_marker_comment` underneath - so the parsing logic is exercised either way; these tests pin the - grace-marker-specific behavior.""" - - def _make_comment( - self, - *, - login: str, - body: str, - created_at: str | None = "2026-05-18T05:00:00Z", - ) -> dict: - comment: dict = {"user": {"login": login}, "body": body} - if created_at is not None: - comment["created_at"] = created_at - return comment - - def test_should_return_none_when_no_grace_marker(self, triage_module, monkeypatch): - comments = [ - self._make_comment( - login="github-actions[bot]", - body="Some other bot message", - ), - self._make_comment(login="random-user", body="ping?"), - ] - monkeypatch.setattr( - triage_module, "_iter_paginated_json", lambda *a, **kw: iter(comments) - ) - assert triage_module.seconds_since_last_grace_warning("o/r", 1) is None - - def test_should_ignore_non_bot_comments_with_marker( - self, triage_module, monkeypatch - ): - # A user who quotes the marker in a question must NOT be treated - # as the bot warning; otherwise the close-after-grace path would - # never fire because the timer keeps resetting. - comments = [ - self._make_comment( - login="random-user", - body=f"What is {triage_module.GRACE_COMMENT_MARKER}?", - ) - ] - monkeypatch.setattr( - triage_module, "_iter_paginated_json", lambda *a, **kw: iter(comments) - ) - assert triage_module.seconds_since_last_grace_warning("o/r", 1) is None - - def test_should_pick_latest_grace_marker(self, triage_module, monkeypatch): - comments = [ - self._make_comment( - login="github-actions[bot]", - body="old warning " + triage_module.GRACE_COMMENT_MARKER, - created_at="2026-05-18T03:00:00Z", - ), - self._make_comment( - login="github-actions[bot]", - body="newer warning " + triage_module.GRACE_COMMENT_MARKER, - created_at="2026-05-18T04:55:00Z", - ), - ] - monkeypatch.setattr( - triage_module, "_iter_paginated_json", lambda *a, **kw: iter(comments) - ) - - import datetime as real_dt - - class FrozenDateTime(real_dt.datetime): - @classmethod - def now(cls, tz=None): - return real_dt.datetime(2026, 5, 18, 5, 0, 0, tzinfo=tz) - - frozen_module = type(triage_module.dt)("datetime") - frozen_module.datetime = FrozenDateTime - frozen_module.timezone = real_dt.timezone - monkeypatch.setattr(triage_module, "dt", frozen_module) - - age = triage_module.seconds_since_last_grace_warning("o/r", 1) - # Newer warning is 5 minutes (300s) before "now". - assert age == 300.0 - - -class TestTriageAllowlist: - """The dogfood allowlist gates `triage`: while non-empty it is the sole - author filter (only the named accounts are acted on) and it bypasses the - internal-author exemption for them, so a maintainer can dogfood on their - own org account. Emptying it restores the internal-author skip.""" - - def _make_pr(self, **overrides): - base = { - "number": 1, - "title": "PR title", - "body": "Body with no linked issue and no QA proof.", - "state": "open", - "author_association": "NONE", - "user": {"login": "mateo-berri"}, - } - base.update(overrides) - return base - - def test_should_skip_author_not_on_allowlist(self, triage_module, monkeypatch): - pr = self._make_pr(user={"login": "random-oss-dev"}) - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: pr) - result = triage_module.triage( - repo="o/r", - kind="pr", - number=1, - close=True, - model="m", - judge=lambda p: pytest.fail("LLM must not run for non-allowlisted author"), - ) - assert result["action"] == "skip-not-allowlisted" - - def test_should_act_on_allowlisted_internal_author( - self, triage_module, monkeypatch - ): - pr = self._make_pr(author_association="MEMBER", user={"login": "mateo-berri"}) - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: pr) - result = triage_module.triage( - repo="o/r", - kind="pr", - number=1, - close=True, - model="m", - judge=lambda p: json.dumps( - {"verdict": "pass", "missing": [], "explanation": "ok"} - ), - ) - assert result["action"] == "pass-llm" - - def test_empty_allowlist_restores_internal_skip(self, triage_module, monkeypatch): - pr = self._make_pr( - author_association="MEMBER", user={"login": "krrishdholakia"} - ) - monkeypatch.setattr(triage_module, "fetch_pr", lambda repo, n: pr) - result = triage_module.triage( - repo="o/r", - kind="pr", - number=1, - close=True, - model="m", - judge=lambda p: pytest.fail("LLM must not run for internal author"), - allowlist=frozenset(), - ) - assert result["action"] == "skip-internal-author" - - def test_allowlist_constant_is_the_two_dogfood_accounts(self, triage_module): - assert triage_module.ALLOWLIST_LOGINS == frozenset( - {"mateo-berri", "swiftwinds"} - ) - for login in triage_module.ALLOWLIST_LOGINS: - assert login == login.lower(), login diff --git a/tests/test_litellm/test_github_triage_workflows.py b/tests/test_litellm/test_github_triage_workflows.py deleted file mode 100644 index f96c9b7e974..00000000000 --- a/tests/test_litellm/test_github_triage_workflows.py +++ /dev/null @@ -1,264 +0,0 @@ -"""Static guardrails for the Agent Shin + Greptile workflow YAML files. - -These workflows can post comments and close PRs/issues on -BerriAI/litellm, so the gating logic that decides "is this a real -close-on-fail run?" must fail-safe on any unexpected input. The risk -is mostly maintenance: someone edits the bash gate, drops a quote, -inverts a comparison, or uses `!= "false"` (which treats "True", -"yes", "1", and typos as enabling closure) and the regression isn't -caught until a real OSS contributor's PR gets auto-closed. - -The tests below pin a set of invariants. The first two apply to every -workflow that gates a destructive `--close`: - - 1. The gate uses the fail-safe `= "true"` comparison — not `!= "false"`, - not `!= ""`. Only the literal string "true" should ever enable - closure. - 2. The gate also requires `AGENT_SHIN_ENABLED = "true"` (or the - scheduled-job equivalent) — disabling the variable must always - force dry-run. - -A third invariant covers every workflow that installs the OpenAI client. -These run with a write-scoped `GITHUB_TOKEN`, so a compromised package -release would execute in that context; the install must therefore come -from the hash-pinned `.github/scripts/triage-requirements.txt` via -`pip --require-hashes`, never a floating `pip install openai>=...`. - -Static parsing of the YAML + bash text is the right level of test here: -the gating logic lives in a `run:` block, not in a Python module we can -import, and end-to-end testing a GitHub Actions workflow from CI is -infeasible. A YAML-level guardrail is exactly what would have caught -the original `!= "false"` regression at PR time. -""" - -from __future__ import annotations - -from pathlib import Path - -import pytest -import yaml - -REPO_ROOT = Path(__file__).resolve().parents[2] -WORKFLOWS_DIR = REPO_ROOT / ".github" / "workflows" - -# Map of workflow file -> the env var name that drives the destructive -# gate inside that workflow's `run:` block. Keeping this table explicit -# (rather than scraping every workflow file) means a new workflow file -# that bypasses the dry-run gating doesn't silently slip past this test. -DESTRUCTIVE_GATE_ENV: dict[str, str] = { - "close_low_quality_prs.yml": "CLOSE_FLAG", - # The reconsider workflow has no per-run "really do it?" knob — its - # only kill switch is `AGENT_SHIN_ENABLED`, which already serves as - # both the destructive gate and the global enablement gate. - "triage_reconsider.yml": "AGENT_SHIN_ENABLED", -} - - -# Privileged workflows that install the OpenAI client. They run with a -# write-scoped GITHUB_TOKEN, so the install must be hash-pinned: a poisoned -# release would otherwise execute in that context. A new workflow that -# installs the client must be added here and use the same pinned file. -LLM_CLIENT_INSTALLER_WORKFLOWS = ( - "triage_reconsider.yml", -) - -PINNED_INSTALL = "--require-hashes -r .github/scripts/triage-requirements.txt" -REQUIREMENTS_FILE = REPO_ROOT / ".github" / "scripts" / "triage-requirements.txt" - - -def _load_workflow(name: str) -> dict: - return yaml.safe_load((WORKFLOWS_DIR / name).read_text()) - - -def _all_run_blocks(workflow: dict) -> list[str]: - """Return every `run:` step's command text, joined.""" - commands: list[str] = [] - jobs = workflow.get("jobs") or {} - for job in jobs.values(): - for step in job.get("steps", []) or []: - if not isinstance(step, dict): - continue - run = step.get("run") - if isinstance(run, str): - commands.append(run) - return commands - - -@pytest.mark.parametrize("workflow_file,env_var", sorted(DESTRUCTIVE_GATE_ENV.items())) -def test_should_use_failsafe_equals_true_comparison(workflow_file: str, env_var: str) -> None: - """The destructive `--close` gate must use `= "true"` (fail-safe), not - `!= "false"` (which would treat "True", "yes", "1", or any typo as - enabling closure). - - Both bare `${ENV_VAR}` and `${ENV_VAR:-false}` (with a default) are - accepted forms — what matters is the comparison operator. The - Greptile closer relies on an outer `AGENT_SHIN_ENABLED` gate so it - can use the bare form; the Agent Shin workflows include `:-false` - for defense in depth. Either is fine. - """ - workflow = _load_workflow(workflow_file) - text = "\n".join(_all_run_blocks(workflow)) - assert env_var in text, ( - f"{workflow_file} no longer references {env_var}; was the gating env var renamed without updating this test?" - ) - accepted_patterns = ( - f'"${{{env_var}}}" = "true"', - f'"${{{env_var}:-false}}" = "true"', - ) - assert any(p in text for p in accepted_patterns), ( - f"{workflow_file} must gate the destructive --close flag on the " - f'EXACT string "true" (one of: {accepted_patterns!r}). Mirror ' - 'the Greptile closer pattern; do NOT use `!= "false"` which ' - 'fail-opens on unknown values like "True", "yes", "1", or typos.' - ) - forbidden_patterns = ( - f'"${{{env_var}}}" != "false"', - f'"${{{env_var}:-false}}" != "false"', - f'"${{{env_var}:-true}}" != "false"', - ) - for forbidden in forbidden_patterns: - assert forbidden not in text, ( - f"{workflow_file} uses the fail-open pattern {forbidden!r}. " - 'Switch to `= "true"` so unknown values stay dry-run.' - ) - - -@pytest.mark.parametrize("workflow_file", sorted(DESTRUCTIVE_GATE_ENV)) -def test_should_require_agent_shin_enabled_for_close(workflow_file: str) -> None: - """Every destructive gate must also gate on the global enablement - variable, so flipping `AGENT_SHIN_ENABLED` off is a kill switch - regardless of any per-run input. - - Two patterns are equally fine: - - Positive: `[ "${AGENT_SHIN_ENABLED:-false}" = "true" ]` to enter - the close branch (Agent Shin workflows). - - Negative: `[ "${AGENT_SHIN_ENABLED:-false}" != "true" ]` then - bail out / force dry-run (Greptile closer). - - What matters is that the comparison value is the literal "true"; - `!= "false"` or `= "1"` etc. would not be a true kill switch. - """ - workflow = _load_workflow(workflow_file) - text = "\n".join(_all_run_blocks(workflow)) - accepted_patterns = ( - '"${AGENT_SHIN_ENABLED:-false}" = "true"', - '"${AGENT_SHIN_ENABLED:-false}" != "true"', - ) - assert any(p in text for p in accepted_patterns), ( - f"{workflow_file} must gate destructive actions on " - '`AGENT_SHIN_ENABLED = "true"` (or the inverted `!= "true"` ' - "guard that forces dry-run). Without this, an unset repo " - "variable would not be treated as a kill switch." - ) - - -@pytest.mark.parametrize("workflow_file", LLM_CLIENT_INSTALLER_WORKFLOWS) -def test_llm_client_install_is_hash_pinned(workflow_file: str) -> None: - """Every privileged workflow installs the OpenAI client from the - hash-pinned requirements file, never by floating version. - - A bare `pip install "openai>=1.40.0"` resolves to whatever PyPI serves - at run time and executes during install/import while a write-scoped - `GITHUB_TOKEN` is in scope, so a compromised release runs in a - privileged context. This test fails if that floating form comes back or - if the `--require-hashes` install is loosened. - """ - blocks = _all_run_blocks(_load_workflow(workflow_file)) - assert PINNED_INSTALL in "\n".join(blocks), ( - f"{workflow_file} must install the client via `pip install " - f"{PINNED_INSTALL}`; a floating install runs unverified code with a " - "write-scoped token." - ) - offenders = [b for b in blocks if "pip install" in b and "openai" in b] - assert not offenders, ( - f"{workflow_file} installs openai by name ({offenders!r}); pin it " - "through the hash-locked requirements file so the version and " - "checksum are fixed." - ) - - -def test_triage_requirements_are_fully_hash_pinned() -> None: - """The shared requirements file pins every package to an exact version - with a sha256 hash, which is what `pip --require-hashes` enforces at - install time. A loosened pin or a missing hash here would silently widen - the supply-chain surface for all the installer workflows. - """ - assert REQUIREMENTS_FILE.exists(), ( - f"the hash-pinned requirements file the triage workflows install from is missing at {REQUIREMENTS_FILE}" - ) - joined = REQUIREMENTS_FILE.read_text().replace("\\\n", " ") - entries = [line.strip() for line in joined.splitlines() if line.strip() and not line.strip().startswith("#")] - assert any(e.split()[0].startswith("openai==") for e in entries), ( - "openai must be pinned to an exact version in the triage requirements" - ) - for entry in entries: - spec = entry.split()[0] - assert "==" in spec, ( - f"requirement {spec!r} is not pinned to an exact version; " - "--require-hashes needs every package pinned with ==" - ) - assert "--hash=sha256:" in entry, ( - f"requirement {spec!r} has no sha256 hash; every pin must carry " - "checksums so --require-hashes can verify the download" - ) - - -def _reconsider_steps() -> list[dict]: - workflow = _load_workflow("triage_reconsider.yml") - return workflow["jobs"]["reconsider"]["steps"] - - -def _index_of_run_step(steps: list[dict], needle: str) -> int: - for i, step in enumerate(steps): - run = step.get("run") - if isinstance(run, str) and needle in run: - return i - raise AssertionError(f"no run step contains {needle!r}") - - -def _reaction_steps(steps: list[dict], content: str) -> list[tuple[int, dict]]: - return [ - (i, s) - for i, s in enumerate(steps) - if isinstance(s.get("run"), str) and f"content={content}" in s["run"] and "/reactions" in s["run"] - ] - - -class TestReconsiderReactions: - """The reconsider workflow acknowledges the triggering comment with a 👀 - reaction the moment it accepts the trigger, and a 👍 once the run finishes, - so the contributor gets feedback immediately instead of waiting on a cron. - - Both reactions are gated on `AGENT_SHIN_ENABLED == 'true'` so a dry-run - leaves no visible trace, and both target the comment that fired the event - (`github.event.comment.id`). The ordering (👀 before the triage run, 👍 - after) is the whole point — these tests fail if a refactor reorders the - steps, drops a reaction, or stops gating them. - """ - - def test_eyes_reaction_is_posted_before_the_triage_run(self) -> None: - steps = _reconsider_steps() - run_idx = _index_of_run_step(steps, "triage_with_llm.py") - eyes = _reaction_steps(steps, "eyes") - assert len(eyes) == 1, "expected exactly one 👀 (eyes) reaction step" - idx, step = eyes[0] - assert idx < run_idx, "👀 must be posted BEFORE the slow triage run, not after" - assert "github.event.comment.id" in (step.get("env") or {}).get("COMMENT_ID", ""), ( - "👀 must react to the comment that triggered the workflow" - ) - assert "${COMMENT_ID}" in step["run"], "👀 must react to the triggering comment, not a hardcoded id" - assert "vars.AGENT_SHIN_ENABLED == 'true'" in step["if"], ( - "👀 must be gated on AGENT_SHIN_ENABLED so dry-run stays inert" - ) - - def test_thumbs_up_reaction_is_posted_after_a_successful_run(self) -> None: - steps = _reconsider_steps() - run_idx = _index_of_run_step(steps, "triage_with_llm.py") - thumbs = _reaction_steps(steps, "+1") - assert len(thumbs) == 1, "expected exactly one 👍 (+1) reaction step" - idx, step = thumbs[0] - assert idx > run_idx, "👍 must come AFTER the triage run" - assert "success()" in step["if"], "👍 must only fire when the reconsider run succeeded" - assert "vars.AGENT_SHIN_ENABLED == 'true'" in step["if"], ( - "👍 must be gated on AGENT_SHIN_ENABLED so dry-run stays inert" - ) diff --git a/tests/test_litellm/test_main.py b/tests/test_litellm/test_main.py index cbbac3d247f..bd115c699d5 100644 --- a/tests/test_litellm/test_main.py +++ b/tests/test_litellm/test_main.py @@ -349,28 +349,66 @@ def test_bedrock_latency_optimized_inference(): assert json_data["performanceConfig"]["latency"] == "optimized" -def test_strip_input_examples_for_non_anthropic_providers(): +@pytest.mark.parametrize( + ("custom_llm_provider", "model", "expected"), + [ + ("anthropic", "claude-sonnet-5", True), + ("bedrock", "us.anthropic.claude-sonnet-5-20260501-v1:0", True), + ("bedrock", "arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/abc123", True), + ("bedrock", "us.amazon.nova-2-lite-v1:0", False), + ("vertex_ai", "claude-sonnet-5", True), + ("vertex_ai", "gemini-3.8-flash", False), + ("azure_ai", "claude-sonnet-4-6", True), + ("azure_ai", "gpt-5.6", False), + ("openai", "gpt-5.6", False), + ("gemini", "gemini-3.8-flash", False), + ], +) +def test_is_claude_tool_target(custom_llm_provider: str, model: str, expected: bool): + assert litellm_main._is_claude_tool_target(custom_llm_provider=custom_llm_provider, model=model) is expected + + +@pytest.mark.parametrize("key", ["input_examples", "eager_input_streaming"]) +def test_drop_anthropic_only_tool_keys_strips_tool_and_function_levels(key: str): tools = [ - { - "type": "function", - "name": "example_tool", - "input_examples": [{"foo": "bar"}], - "function": { - "name": "example_tool", - "input_examples": [{"foo": "bar"}], - }, - } + {"type": "function", "name": "example_tool", key: True, "function": {"name": "example_tool", key: True}}, + "opaque_tool", ] - assert not litellm_main._should_allow_input_examples( - custom_llm_provider="openai", model="gpt-4o-mini" + cleaned = litellm_main._drop_anthropic_only_tool_keys(tools=tools) + + assert cleaned == [ + {"type": "function", "name": "example_tool", "function": {"name": "example_tool"}}, + "opaque_tool", + ] + assert tools[0][key] is True + assert tools[0]["function"][key] is True + + +def test_completion_strips_eager_input_streaming_before_openai(respx_mock: respx.MockRouter, openai_api_response): + api_base: Final = "http://localhost:12346/v1" + mock_route: Final = respx_mock.post(url__regex=rf"{api_base}/chat/completions.*").mock( + return_value=httpx.Response(status_code=200, json=openai_api_response) ) - cleaned = litellm_main._drop_input_examples_from_tools(tools=tools) + litellm.completion( + model="openai/gpt-5.6", + messages=[{"role": "user", "content": "Write the file"}], + tools=[ + { + "type": "function", + "function": {"name": "write_file", "parameters": {"type": "object", "properties": {}}}, + "eager_input_streaming": True, + } + ], + api_base=api_base, + api_key="fake_openai_api_key", + ) - assert isinstance(cleaned, list) - assert "input_examples" not in cleaned[0] - assert "input_examples" not in cleaned[0]["function"] + assert mock_route.called + sent_tool: Final = json.loads(respx_mock.calls[0].request.content)["tools"][0] + assert "eager_input_streaming" not in sent_tool + assert sent_tool["function"]["name"] == "write_file" def test_custom_provider_with_extra_headers(): diff --git a/tests/test_litellm_rust/conftest.py b/tests/test_litellm_rust/conftest.py index 4387ea2e2fd..1b6fcfa00db 100644 --- a/tests/test_litellm_rust/conftest.py +++ b/tests/test_litellm_rust/conftest.py @@ -1,10 +1,9 @@ import asyncio import os -from collections.abc import AsyncIterator, Generator, Iterator +from collections.abc import AsyncIterator, Generator from concurrent.futures import ThreadPoolExecutor -from contextlib import ExitStack, contextmanager -from types import ModuleType -from typing import Final, cast +from contextlib import ExitStack +from typing import Final import pytest import pytest_asyncio @@ -18,62 +17,20 @@ from litellm.rust_bridge.configuration import ( # pyright: ignore[reportPrivate _parse_env_bool, ) from tests.test_litellm_rust.support.callback_recorder import drain_logging +from tests.test_litellm_rust.support.isolation import isolated_callback_registries, rebound from tests.test_litellm_rust.support.recording_server import RecordingServer, recording_service -CALLBACK_ATTRIBUTES: Final = ( - "callbacks", - "input_callback", - "success_callback", - "failure_callback", - "_async_input_callback", - "_async_success_callback", - "_async_failure_callback", -) - - -def _list_attribute(container: ModuleType, attribute: str) -> list[object]: - value: Final = getattr(container, attribute) - if not isinstance(value, list): - raise AssertionError(f"{container.__name__}.{attribute} is not a list") - return cast(list[object], value) - - -@contextmanager -def _isolated_list(container: ModuleType, attribute: str) -> Iterator[None]: - source: Final = _list_attribute(container, attribute) - original: Final = list(source) - source.clear() # mutable-ok: test isolation mutates global registries by design - try: - yield - finally: - source.clear() - source.extend(original) - setattr(container, attribute, source) - - -@contextmanager -def _rebound(container: object, attribute: str, value: object) -> Iterator[None]: - original: Final[object] = getattr(container, attribute) - setattr(container, attribute, value) - try: - yield - finally: - setattr(container, attribute, original) - @pytest_asyncio.fixture(autouse=True, loop_scope="function") async def isolate_ocr_test_state() -> AsyncIterator[None]: with ExitStack() as stack: - for attribute in CALLBACK_ATTRIBUTES: - stack.enter_context(_isolated_list(litellm, attribute)) - stack.enter_context(_isolated_list(litellm_logging, "_in_memory_loggers")) # pyright: ignore[reportPrivateUsage] # no public callback-cache accessor - stack.enter_context(_rebound(utils, "callback_list", [])) # rebind-ok: isolate legacy callback registry - stack.enter_context(_rebound(litellm, "cache", None)) # test-quality-ok: isolate process-global cache - stack.enter_context(_rebound(_CONFIGURATION, "override", None)) + stack.enter_context(isolated_callback_registries()) + stack.enter_context(rebound(litellm, "cache", None)) # test-quality-ok: isolate process-global cache + stack.enter_context(rebound(_CONFIGURATION, "override", None)) executor: Final = ThreadPoolExecutor(thread_name_prefix="rust-ocr-test-logging") - stack.enter_context(_rebound(litellm_logging, "executor", executor)) - stack.enter_context(_rebound(utils, "executor", executor)) - stack.enter_context(_rebound(thread_pool_executor, "executor", executor)) + stack.enter_context(rebound(litellm_logging, "executor", executor)) + stack.enter_context(rebound(utils, "executor", executor)) + stack.enter_context(rebound(thread_pool_executor, "executor", executor)) try: yield finally: diff --git a/tests/test_litellm_rust/messages/__init__.py b/tests/test_litellm_rust/messages/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm_rust/messages/test_callbacks.py b/tests/test_litellm_rust/messages/test_callbacks.py new file mode 100644 index 00000000000..b55bc47d640 --- /dev/null +++ b/tests/test_litellm_rust/messages/test_callbacks.py @@ -0,0 +1,175 @@ +from collections.abc import AsyncIterator, Iterator +from typing import Final + +import pytest + +import litellm +from litellm.integrations.custom_logger import CustomLogger +from tests.test_litellm_rust.support.callback_recorder import RecordingLogger, drain_logging +from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec +from tests.test_litellm_rust.support.requests import ( + MESSAGES, + MESSAGES_EVENTS, + MESSAGES_MODEL, + MESSAGES_RESPONSE, + request_body, +) + +pytestmark = pytest.mark.requires_rust_extension + +STREAM: Final = ResponseSpec(body=None, events=MESSAGES_EVENTS) + + +@pytest.fixture +def messages_server(recording_server: RecordingServer) -> RecordingServer: + recording_server.default_response = ResponseSpec(body=MESSAGES_RESPONSE) + return recording_server + + +def arguments(server: RecordingServer, **kwargs: object) -> dict[str, object]: + return { + "model": MESSAGES_MODEL, + "messages": [dict(message) for message in MESSAGES], + "max_tokens": 64, + "api_key": "test-key", + "api_base": server.base_url, + **kwargs, + } + + +def assert_served_natively(server: RecordingServer) -> None: + assert len(server.requests) == 1 + assert not server.requests[0].headers.get("user-agent", "").startswith("python-httpx") + + +@pytest.mark.asyncio +async def test_native_messages_callbacks_see_the_provider_request_and_the_public_response( + messages_server: RecordingServer, +) -> None: + recorder: Final = RecordingLogger() + + response: Final = await litellm.anthropic.messages.acreate( + **arguments(messages_server, callbacks=[recorder], litellm_call_id="messages-success") + ) + + assert_served_natively(messages_server) + assert response["content"] == MESSAGES_RESPONSE["content"] + sent: Final = messages_server.requests[0] + assert sent.path == "/v1/messages" + assert sent.body == {"model": "claude-sonnet-5", "messages": list(MESSAGES), "max_tokens": 64, "stream": False} + pre_call: Final = recorder.wait_for("log_pre_api_call") + assert request_body(pre_call[0].kwargs) == sent.body + success: Final = await recorder.wait_for_async("async_log_success_event") + assert len(success) == 1 + assert success[0].call_type == "anthropic_messages" + assert success[0].kwargs["litellm_call_id"] == "messages-success" + assert success[0].response.choices[0].message.content == "Hello from native Messages" + + +@pytest.mark.asyncio +async def test_native_messages_pre_call_body_edit_reaches_the_provider(messages_server: RecordingServer) -> None: + class Edit(CustomLogger): + def log_pre_api_call(self, model, messages, kwargs): + request_body(kwargs)["temperature"] = 0.25 + + await litellm.anthropic.messages.acreate(**arguments(messages_server, callbacks=[Edit()])) + + assert messages_server.requests[0].body["temperature"] == 0.25 + + +@pytest.mark.asyncio +async def test_native_messages_provider_error_reaches_caller_and_failure_callbacks_as_one_public_error( + messages_server: RecordingServer, +) -> None: + messages_server.enqueue( + ResponseSpec(body={"type": "error", "error": {"type": "invalid_request_error", "message": "bad"}}, status=400) + ) + observed: Final = [] + + class Observe(CustomLogger): + def log_failure_event(self, kwargs, response_obj, start_time, end_time): + observed.append(("sync", kwargs["exception"])) + + async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): + observed.append(("async", kwargs["exception"])) + + with pytest.raises(litellm.BadRequestError) as raised: + await litellm.anthropic.messages.acreate(**arguments(messages_server, callbacks=[Observe()])) + + assert_served_natively(messages_server) + assert [phase for phase, _ in observed] == ["sync", "async"] + assert all(error is raised.value for _, error in observed) + + +def sse_payload() -> bytes: + return b"".join(STREAM.payloads()) + + +@pytest.mark.asyncio +async def test_native_messages_stream_relays_provider_events_and_logs_success_once_after_the_last_chunk( + messages_server: RecordingServer, +) -> None: + messages_server.enqueue(STREAM) + recorder: Final = RecordingLogger() + + stream: Final = await litellm.anthropic.messages.acreate( + **arguments(messages_server, stream=True, callbacks=[recorder]) + ) + assert isinstance(stream, AsyncIterator) + first: Final = await anext(stream) + await drain_logging() + assert "async_log_success_event" not in recorder.names + rest: Final = [chunk async for chunk in stream] + + assert first + b"".join(rest) == sse_payload() + assert_served_natively(messages_server) + assert messages_server.requests[0].body["stream"] is True + success: Final = await recorder.wait_for_async("async_log_success_event") + assert len(success) == 1 + assert success[0].kwargs["stream"] is True + assert success[0].kwargs["completion_start_time"] is not None + assert "log_failure_event" not in recorder.names + + +@pytest.mark.asyncio +async def test_native_messages_stream_closed_early_logs_success_once_for_what_was_delivered( + messages_server: RecordingServer, +) -> None: + messages_server.enqueue(STREAM) + recorder: Final = RecordingLogger() + + stream: Final = await litellm.anthropic.messages.acreate( + **arguments(messages_server, stream=True, callbacks=[recorder]) + ) + assert isinstance(stream, AsyncIterator) + await anext(stream) + await stream.aclose() + + success: Final = await recorder.wait_for_async("async_log_success_event") + assert len(success) == 1 + with pytest.raises(StopAsyncIteration): + await anext(stream) + + +def test_native_sync_messages_stream_relays_provider_events_and_logs_success_once( + messages_server: RecordingServer, +) -> None: + messages_server.enqueue(STREAM) + recorder: Final = RecordingLogger() + + stream: Final = litellm.anthropic.messages.create(**arguments(messages_server, stream=True, callbacks=[recorder])) + assert isinstance(stream, Iterator) + + assert b"".join(stream) == sse_payload() + assert_served_natively(messages_server) + assert len(recorder.wait_for("async_log_success_event")) == 1 + + +def test_native_sync_messages_returns_the_provider_message(messages_server: RecordingServer) -> None: + recorder: Final = RecordingLogger() + + response: Final = litellm.anthropic.messages.create(**arguments(messages_server, callbacks=[recorder])) + + assert_served_natively(messages_server) + assert response["content"] == MESSAGES_RESPONSE["content"] + assert len(recorder.wait_for("log_success_event")) == 1 diff --git a/tests/test_litellm_rust/ocr/test_callbacks.py b/tests/test_litellm_rust/ocr/test_callbacks.py index 27cdcc4d997..ac4a1a11a80 100644 --- a/tests/test_litellm_rust/ocr/test_callbacks.py +++ b/tests/test_litellm_rust/ocr/test_callbacks.py @@ -1,24 +1,31 @@ import asyncio import copy +import gc import queue import threading +from collections.abc import Mapping +from types import MappingProxyType from typing import Final import pytest +from hypothesis import HealthCheck, given, settings +from hypothesis import strategies as st import litellm from litellm.integrations.custom_logger import CustomLogger from litellm.llms.base_llm.ocr.transformation import OCRResponse -from tests.test_litellm_rust.support.callback_recorder import RecordingLogger +from tests.test_litellm_rust.support.callback_recorder import RecordingLogger, drain_logging +from tests.test_litellm_rust.support.isolation import isolated_callback_registries +from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec from tests.test_litellm_rust.support.requests import ( OCR_DOCUMENT, OCR_RESPONSE, + call_native, call_native_aocr, call_native_ocr, request_body, request_headers, ) -from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec pytestmark = pytest.mark.requires_rust_extension @@ -123,9 +130,7 @@ async def test_native_ocr_pre_call_nested_document_edit_updates_caller_callback_ "callbacks": [Retain(), Edit()], } response: Final = ( - await call_native_aocr(ocr_server, **arguments) - if asynchronous - else call_native_ocr(ocr_server, **arguments) + await call_native_aocr(ocr_server, **arguments) if asynchronous else call_native_ocr(ocr_server, **arguments) ) assert aliases == [True] @@ -291,6 +296,153 @@ def test_native_ocr_dispatches_each_callback_phase_once_when_logger_is_registere assert "log_failure_event" not in recorder.names +JSON_SCALARS: Final = ( + st.none() + | st.booleans() + | st.integers(min_value=-(2**63), max_value=2**63 - 1) + | st.floats(allow_nan=False, allow_infinity=False) + | st.text(max_size=8) +) +JSON_VALUES: Final = st.recursive( + JSON_SCALARS, + lambda children: st.lists(children, max_size=3) | st.dictionaries(st.text(max_size=6), children, max_size=3), + max_leaves=8, +) + + +class ApplyEdits(CustomLogger): + def __init__(self, edits: Mapping[str, object]) -> None: + super().__init__() + self.edits: Final = edits + + def log_pre_api_call(self, model, messages, kwargs): + request_body(kwargs).update(copy.deepcopy(dict(self.edits))) + + +@settings(max_examples=25, deadline=None, suppress_health_check=[HealthCheck.function_scoped_fixture]) +@given(edits=st.dictionaries(st.from_regex(r"x_[a-z]{1,6}", fullmatch=True), JSON_VALUES, max_size=3)) +def test_native_ocr_provider_receives_the_body_exactly_as_pre_call_callbacks_left_it( + ocr_server: RecordingServer, edits: dict[str, object] +) -> None: + ocr_server.expected_requests = None + + with isolated_callback_registries(): + call_native_ocr_with_callbacks(ocr_server, [ApplyEdits(MappingProxyType(edits))]) + + assert ocr_server.requests[-1].body == {"model": "mistral-ocr-latest", "document": OCR_DOCUMENT, **edits} + + +@pytest.mark.parametrize("hook", ["log_pre_api_call", "logging_hook", "log_success_event"]) +def test_native_ocr_sync_hooks_see_no_running_event_loop(ocr_server: RecordingServer, hook: str) -> None: + recorder: Final = RecordingLogger() + + call_native_ocr_with_callbacks(ocr_server, [recorder]) + + [event] = recorder.wait_for(hook) + assert event.loop is None + assert (event.thread is threading.current_thread()) == (hook == "log_pre_api_call") + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"]) +async def test_native_ocr_payload_a_callback_retains_outlives_the_call_intact( + ocr_server: RecordingServer, asynchronous: bool +) -> None: + retained: Final = [] + + class Retain(CustomLogger): + def log_pre_api_call(self, model, messages, kwargs): + retained.append((kwargs, request_body(kwargs), request_headers(kwargs))) + + await call_native(ocr_server, asynchronous, callbacks=[Retain()]) + await drain_logging() + gc.collect() + + [(details, body, headers)] = retained + assert body == ocr_server.requests[0].body + assert headers + assert all(ocr_server.requests[0].headers[name.lower()] == value for name, value in headers.items()) + assert details["additional_args"]["complete_input_dict"] is body + assert details["additional_args"]["headers"] is headers + + +@pytest.mark.asyncio +@pytest.mark.parametrize("family", ["sync", "async"]) +async def test_native_ocr_success_callbacks_share_one_logging_payload(ocr_server: RecordingServer, family: str) -> None: + queued: Final = [] + finished: Final = threading.Event() + + def queue_payload(kwargs: dict[str, object]) -> None: + queued.append(kwargs["standard_logging_object"]) + + def strip_payload(kwargs: dict[str, object]) -> None: + payload: Final = kwargs["standard_logging_object"] + assert isinstance(payload, dict) + payload["stripped-by-a-later-callback"] = True + finished.set() + + class QueuePayload(CustomLogger): + if family == "sync": + + def log_success_event(self, kwargs, response_obj, start_time, end_time): + queue_payload(kwargs) + + else: + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + queue_payload(kwargs) + + class StripPayload(CustomLogger): + if family == "sync": + + def log_success_event(self, kwargs, response_obj, start_time, end_time): + strip_payload(kwargs) + + else: + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + strip_payload(kwargs) + + await call_native(ocr_server, family == "async", callbacks=[QueuePayload(), StripPayload()]) + await drain_logging() + + assert await asyncio.to_thread(finished.wait, 10) + assert [payload["stripped-by-a-later-callback"] for payload in queued] == [True] + + +@pytest.mark.asyncio +async def test_native_aocr_state_stashed_before_a_blocking_hook_raises_reaches_failure_callbacks( + ocr_server: RecordingServer, +) -> None: + token: Final = object() + observed: Final = [] + + class Blocked(Exception): + pass + + class Block(CustomLogger): + async def async_post_call_success_deployment_hook(self, request_data, response, call_type): + request_data["litellm_logging_obj"].model_call_details["blocked-by"] = token + raise Blocked("blocked after the provider answered") + + def log_success_event(self, kwargs, response_obj, start_time, end_time): + observed.append(("success", None, None)) + + def log_failure_event(self, kwargs, response_obj, start_time, end_time): + observed.append(("sync", kwargs.get("blocked-by"), kwargs["exception"])) + + async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): + observed.append(("async", kwargs.get("blocked-by"), kwargs["exception"])) + + litellm.callbacks.append(Block()) + + with pytest.raises(Blocked) as raised: + await call_native_aocr(ocr_server) + await drain_logging() + + assert observed == [("sync", token, raised.value), ("async", token, raised.value)] + + @pytest.mark.asyncio @pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"]) async def test_native_azure_ocr_resolves_token_before_pre_call_on_caller_context( diff --git a/tests/test_litellm_rust/ocr/test_lifecycle.py b/tests/test_litellm_rust/ocr/test_lifecycle.py index 085ea4a14c0..5fca927bea3 100644 --- a/tests/test_litellm_rust/ocr/test_lifecycle.py +++ b/tests/test_litellm_rust/ocr/test_lifecycle.py @@ -4,7 +4,7 @@ import gc import json import threading import weakref -from collections.abc import Coroutine +from collections.abc import Awaitable, Callable, Coroutine from contextvars import ContextVar from typing import Final @@ -400,7 +400,7 @@ async def test_response_limit_is_enforced_at_the_public_boundary(ocr_server: Rec @pytest.mark.asyncio @pytest.mark.parametrize("failure", [False, True]) -async def test_empty_callbacks_keep_bookkeeping_without_optional_dispatch( +async def test_empty_callbacks_run_deployment_hooks_and_defer_like_the_python_client_wrapper( ocr_server: RecordingServer, monkeypatch: pytest.MonkeyPatch, failure: bool, @@ -414,8 +414,12 @@ async def test_empty_callbacks_keep_bookkeeping_without_optional_dispatch( submissions = 0 enqueues = 0 - def deployment(self, *args: object, **kwargs: object) -> None: - self.deployments += 1 + def counting(self, hook: Callable[..., Awaitable[object]]) -> Callable[..., Awaitable[object]]: + async def counted(*args: object, **kwargs: object) -> object: + self.deployments += 1 + return await hook(*args, **kwargs) + + return counted def submit(self, *args: object, **kwargs: object) -> None: self.submissions += 1 @@ -430,7 +434,7 @@ async def test_empty_callbacks_keep_bookkeeping_without_optional_dispatch( "async_post_call_success_deployment_hook", "async_post_call_failure_deployment_hook", ): - monkeypatch.setattr(utils, name, probe.deployment) + monkeypatch.setattr(utils, name, probe.counting(getattr(utils, name))) monkeypatch.setattr(litellm_logging, "executor", probe) monkeypatch.setattr(logging_worker, "GLOBAL_LOGGING_WORKER", probe) if failure: @@ -447,17 +451,16 @@ async def test_empty_callbacks_keep_bookkeeping_without_optional_dispatch( assert response._hidden_params["response_cost"] is not None assert response._hidden_params["_response_ms"] > 0 assert trace_id_var.get() == "callback-free-parent" - assert probe.deployments == probe.submissions == probe.enqueues == 0 + assert probe.deployments == 2 + assert probe.submissions == probe.enqueues == 0 assert len(created_loggers) == 1 logger: Final = created_loggers[0] - assert not hasattr(logger, "_native_pending_logging") - assert logger.model_call_details["first_api_call_start_time"] <= logger.model_call_details["end_time"] - assert "standard_logging_object" not in logger.model_call_details - assert ( - "original_response" not in logger.model_call_details or logger.model_call_details["original_response"] is None - ) - assert "complete_input_dict" not in logger.model_call_details.get("additional_args", {}) - assert logger.model_call_details["response_cost"] == (0 if failure else response._hidden_params["response_cost"]) + if failure: + assert logger.model_call_details["first_api_call_start_time"] <= logger.model_call_details["end_time"] + assert logger.model_call_details["response_cost"] == 0 + else: + assert getattr(logger, "_native_pending_logging", None) is not None + assert "end_time" not in logger.model_call_details @pytest.mark.asyncio diff --git a/tests/test_litellm_rust/support/isolation.py b/tests/test_litellm_rust/support/isolation.py new file mode 100644 index 00000000000..f98ce4843a8 --- /dev/null +++ b/tests/test_litellm_rust/support/isolation.py @@ -0,0 +1,58 @@ +from collections.abc import Generator +from contextlib import ExitStack, contextmanager +from types import ModuleType +from typing import Final, cast + +import litellm +from litellm import utils +from litellm.litellm_core_utils import litellm_logging + +CALLBACK_ATTRIBUTES: Final = ( + "callbacks", + "input_callback", + "success_callback", + "failure_callback", + "_async_input_callback", + "_async_success_callback", + "_async_failure_callback", +) + + +def _list_attribute(container: ModuleType, attribute: str) -> list[object]: + value: Final = getattr(container, attribute) + if not isinstance(value, list): + raise AssertionError(f"{container.__name__}.{attribute} is not a list") + return cast(list[object], value) + + +@contextmanager +def _isolated_list(container: ModuleType, attribute: str) -> Generator[None]: + source: Final = _list_attribute(container, attribute) + original: Final = list(source) + source.clear() # mutable-ok: test isolation mutates global registries by design + try: + yield + finally: + source.clear() + source.extend(original) + setattr(container, attribute, source) + + +@contextmanager +def rebound(container: object, attribute: str, value: object) -> Generator[None]: + original: Final[object] = getattr(container, attribute) + setattr(container, attribute, value) + try: + yield + finally: + setattr(container, attribute, original) + + +@contextmanager +def isolated_callback_registries() -> Generator[None]: + with ExitStack() as stack: + for attribute in CALLBACK_ATTRIBUTES: + stack.enter_context(_isolated_list(litellm, attribute)) + stack.enter_context(_isolated_list(litellm_logging, "_in_memory_loggers")) # pyright: ignore[reportPrivateUsage] # no public callback-cache accessor + stack.enter_context(rebound(utils, "callback_list", [])) # rebind-ok: isolate legacy callback registry + yield diff --git a/tests/test_litellm_rust/support/recording_server.py b/tests/test_litellm_rust/support/recording_server.py index 228ed2cc454..3eea47751d3 100644 --- a/tests/test_litellm_rust/support/recording_server.py +++ b/tests/test_litellm_rust/support/recording_server.py @@ -25,6 +25,12 @@ class ResponseSpec: status: int = 200 headers: dict[str, str] = field(default_factory=dict) delay: float = 0 + events: tuple[tuple[str, object], ...] = () + + def payloads(self) -> tuple[bytes, ...]: + if not self.events: + return (json.dumps(self.body).encode(),) + return tuple(f"event: {event}\ndata: {json.dumps(data)}\n\n".encode() for event, data in self.events) @dataclass @@ -73,15 +79,17 @@ def recording_service() -> Iterator[RecordingServer]: response: Final = responses.pop(0) if responses else copy.deepcopy(recording_server.default_response) if response.delay: time.sleep(response.delay) - payload: Final = json.dumps(response.body).encode() + payloads: Final = response.payloads() self.send_response(response.status) - self.send_header("Content-Type", "application/json") - self.send_header("Content-Length", str(len(payload))) + self.send_header("Content-Type", "text/event-stream" if response.events else "application/json") + self.send_header("Content-Length", str(sum(len(payload) for payload in payloads))) for name, value in response.headers.items(): self.send_header(name, value) self.end_headers() try: - self.wfile.write(payload) + for payload in payloads: + self.wfile.write(payload) + self.wfile.flush() except (BrokenPipeError, ConnectionResetError): pass diff --git a/tests/test_litellm_rust/support/requests.py b/tests/test_litellm_rust/support/requests.py index b60cf5eac02..c9cf81b83ca 100644 --- a/tests/test_litellm_rust/support/requests.py +++ b/tests/test_litellm_rust/support/requests.py @@ -12,6 +12,41 @@ OCR_RESPONSE: Final = { "usage_info": {"pages_processed": 1, "doc_size_bytes": 3}, } +MESSAGES_MODEL: Final = "anthropic/claude-sonnet-5" +MESSAGES: Final = ({"role": "user", "content": "Hello"},) +MESSAGES_RESPONSE: Final = { + "id": "msg_native", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-5", + "content": [{"type": "text", "text": "Hello from native Messages"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 5, "output_tokens": 4}, +} +MESSAGES_EVENTS: Final = ( + ("message_start", {"type": "message_start", "message": {**MESSAGES_RESPONSE, "content": [], "stop_reason": None}}), + ("content_block_start", {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}), + ( + "content_block_delta", + { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "text_delta", "text": "Hello from native Messages"}, + }, + ), + ("content_block_stop", {"type": "content_block_stop", "index": 0}), + ( + "message_delta", + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn", "stop_sequence": None}, + "usage": {"output_tokens": 4}, + }, + ), + ("message_stop", {"type": "message_stop"}), +) + def ocr_arguments(server: RecordingServer, **kwargs: object) -> dict[str, object]: return { diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 55b3a5fef11..136b8f26784 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -26070,6 +26070,8 @@ export interface components { /** Allowed Callers */ allowed_callers?: string[]; cache_control?: components["schemas"]["ChatCompletionCachedContent"]; + /** Eager Input Streaming */ + eager_input_streaming?: boolean; function: components["schemas"]["ChatCompletionToolParamFunctionChunk"]; /** Type */ type: "function" | string; @@ -26078,6 +26080,8 @@ export interface components { ChatCompletionToolParamFunctionChunk: { /** Description */ description?: string; + /** Eager Input Streaming */ + eager_input_streaming?: boolean; /** Name */ name: string; /** Parameters */