mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-19 00:01:29 +00:00
Merge remote-tracking branch 'github/main' into litellm_rust_bridge_declarative_route_catalog
# Conflicts: # tests/e2e/access_control/test_model_access_group_e2e.py
This commit is contained in:
commit
b26935416a
327 changed files with 23015 additions and 2667 deletions
|
|
@ -257,7 +257,7 @@ commands:
|
|||
- install_rust
|
||||
- restore_cache:
|
||||
keys:
|
||||
- v1-uv-cache-{{ checksum "uv.lock" }}
|
||||
- v3-integration-uv-cache-{{ checksum "uv.lock" }}
|
||||
- run:
|
||||
name: Install Dependencies
|
||||
command: |
|
||||
|
|
@ -266,7 +266,7 @@ commands:
|
|||
- save_cache:
|
||||
paths:
|
||||
- ~/.cache/uv
|
||||
key: v1-uv-cache-{{ checksum "uv.lock" }}
|
||||
key: v3-integration-uv-cache-{{ checksum "uv.lock" }}
|
||||
|
||||
jobs:
|
||||
# Add Windows testing job
|
||||
|
|
@ -2955,6 +2955,32 @@ jobs:
|
|||
working_directory: ~/project
|
||||
steps:
|
||||
- setup_litellm_test_deps
|
||||
- when:
|
||||
condition:
|
||||
equal: [browser, << parameters.suite >>]
|
||||
steps:
|
||||
- install_node
|
||||
- restore_cache:
|
||||
keys:
|
||||
- integration-ui-v3-{{ checksum "ui/litellm-dashboard/package-lock.json" }}-{{ checksum "tests/e2e/ui/package-lock.json" }}
|
||||
- run:
|
||||
name: Install locked browser dependencies
|
||||
command: |
|
||||
cd ui/litellm-dashboard
|
||||
npm ci
|
||||
cd ../../tests/e2e/ui
|
||||
npm ci
|
||||
sudo env PATH="$PATH" DEBIAN_FRONTEND=noninteractive NEEDRESTART_MODE=l \
|
||||
timeout --signal=TERM --kill-after=20s 6m node node_modules/@playwright/test/cli.js install-deps chromium
|
||||
timeout --signal=TERM --kill-after=20s 3m node node_modules/@playwright/test/cli.js install chromium
|
||||
- save_cache:
|
||||
key: integration-ui-v3-{{ checksum "ui/litellm-dashboard/package-lock.json" }}-{{ checksum "tests/e2e/ui/package-lock.json" }}
|
||||
paths:
|
||||
- ~/.npm
|
||||
- ~/.cache/ms-playwright
|
||||
- run:
|
||||
name: Build the candidate dashboard
|
||||
command: cd ui/litellm-dashboard && NEXT_TELEMETRY_DISABLED=1 npm run build
|
||||
- start_postgres:
|
||||
image: postgres:16@sha256:e17e86066e5ef83e0952a9347f5c792b7ece00972e2aa787a6986f471b3dd3d5
|
||||
- start_redis
|
||||
|
|
@ -2983,7 +3009,7 @@ workflows:
|
|||
name: integration-<< matrix.suite >>
|
||||
matrix:
|
||||
parameters:
|
||||
suite: [management, accounting, providers]
|
||||
suite: [management, accounting, database, providers, extensions, sdk, browser]
|
||||
filters:
|
||||
branches:
|
||||
only:
|
||||
|
|
|
|||
|
|
@ -1,12 +1,14 @@
|
|||
#!/usr/bin/env bash
|
||||
set -uo pipefail
|
||||
|
||||
category="${1:?usage: classify_changes.sh <backend|client|ui|provider-harness>}"
|
||||
category="${1:?usage: classify_changes.sh <backend|client|ui|provider-harness|cost-map-only>}"
|
||||
|
||||
has_client=false
|
||||
has_backend=false
|
||||
has_ci=false
|
||||
has_provider_harness=false
|
||||
has_cost_map=false
|
||||
outside_cost_map_set=false
|
||||
while IFS= read -r file || [ -n "$file" ]; do
|
||||
[ -n "$file" ] || continue
|
||||
case "$file" in
|
||||
|
|
@ -20,9 +22,18 @@ while IFS= read -r file || [ -n "$file" ]; do
|
|||
.github/* | .circleci/*) has_ci=true; has_backend=true ;;
|
||||
*) has_backend=true ;;
|
||||
esac
|
||||
case "$file" in
|
||||
model_prices_and_context_window.json | litellm/model_prices_and_context_window_backup.json | model_prices_and_context_window.schema.json)
|
||||
has_cost_map=true ;;
|
||||
tests/test_litellm/* | tests/proxy_unit_tests/*) : ;;
|
||||
*) outside_cost_map_set=true ;;
|
||||
esac
|
||||
done
|
||||
|
||||
case "$category" in
|
||||
cost-map-only)
|
||||
{ [ "$has_cost_map" = true ] && [ "$outside_cost_map_set" = false ]; } && echo run || echo skip
|
||||
;;
|
||||
provider-harness)
|
||||
[ "$has_provider_harness" = true ] && echo run || echo skip
|
||||
;;
|
||||
|
|
|
|||
|
|
@ -1,6 +1,11 @@
|
|||
#!/usr/bin/env bash
|
||||
set -euo pipefail
|
||||
|
||||
if [ "${GITHUB_ACTIONS:-}" = true ]; then
|
||||
echo "Integration contracts are owned by CircleCI" >&2
|
||||
exit 1
|
||||
fi
|
||||
|
||||
suite="${1:?integration suite required}"
|
||||
results="test-results/integration-${suite}"
|
||||
mkdir -p "$results"
|
||||
|
|
@ -65,7 +70,13 @@ export INTEGRATION_PROXY_URL=http://127.0.0.1:4000
|
|||
export INTEGRATION_PEER_URL=""
|
||||
export INTEGRATION_UPSTREAM_URL=http://127.0.0.1:8190
|
||||
export INTEGRATION_MASTER_KEY="$LITELLM_MASTER_KEY"
|
||||
export INTEGRATION_SEED="$((16#$(git rev-parse --short=8 HEAD)))"
|
||||
export LITELLM_UI_PATH="$PWD/litellm/proxy/_experimental/out"
|
||||
if [ "$suite" = browser ]; then
|
||||
export LITELLM_UI_PATH="$PWD/ui/litellm-dashboard/out"
|
||||
test -f "$LITELLM_UI_PATH/index.html"
|
||||
fi
|
||||
export INTEGRATION_SEED="$(.venv/bin/python -c 'import hashlib,os; print(int(hashlib.sha256((os.environ.get("CIRCLE_SHA1", "local") + os.environ.get("CIRCLE_WORKFLOW_ID", "local")).encode()).hexdigest()[:8],16))')"
|
||||
export INTEGRATION_ORDER_SEED="$INTEGRATION_SEED"
|
||||
|
||||
uv run --no-sync prisma generate --schema litellm/proxy/schema.prisma > "$results/prisma-generate.log" 2>&1
|
||||
|
||||
|
|
@ -102,7 +113,7 @@ start_proxy() {
|
|||
local log_name="$2"
|
||||
setsid env -i PATH="$PATH" HOME="$HOME" PYTHONPATH="$PYTHONPATH" INTEGRATION_RUN_ID="$integration_identity" \
|
||||
DATABASE_URL="$DATABASE_URL" REDIS_HOST="$REDIS_HOST" REDIS_PORT="$REDIS_PORT" \
|
||||
LITELLM_MASTER_KEY="$LITELLM_MASTER_KEY" LITELLM_SALT_KEY="$LITELLM_SALT_KEY" \
|
||||
LITELLM_MASTER_KEY="$LITELLM_MASTER_KEY" LITELLM_SALT_KEY="$LITELLM_SALT_KEY" LITELLM_UI_PATH="$LITELLM_UI_PATH" PROXY_BASE_URL="http://127.0.0.1:$port" \
|
||||
LITELLM_MODE=PRODUCTION LITELLM_LOCAL_MODEL_COST_MAP=True STORE_MODEL_IN_DB=True \
|
||||
AWS_EC2_METADATA_DISABLED=true DO_NOT_TRACK=1 \
|
||||
.venv/bin/python -m integration._support.proxy --config tests/integration/proxy_config.yaml \
|
||||
|
|
@ -131,6 +142,19 @@ if [ "$suite" = providers ]; then
|
|||
--junitxml="$results/replay-controls.xml"
|
||||
fi
|
||||
|
||||
if [ "$suite" = browser ]; then
|
||||
export E2E_UI_BASE_URL="$INTEGRATION_PROXY_URL" E2E_UI_ARTIFACT_DIR="$PWD/$results"
|
||||
export INTEGRATION_PYTHON="$PWD/.venv/bin/python"
|
||||
timeout --signal=TERM --kill-after=20s 3m env -i PATH="$PATH" HOME="$HOME" PYTHONPATH="$PYTHONPATH" \
|
||||
INTEGRATION_RUN_ID="$integration_identity" DATABASE_URL="$DATABASE_URL" \
|
||||
INTEGRATION_UPSTREAM_URL="$INTEGRATION_UPSTREAM_URL" INTEGRATION_PYTHON="$INTEGRATION_PYTHON" \
|
||||
E2E_UI_BASE_URL="$E2E_UI_BASE_URL" E2E_UI_ARTIFACT_DIR="$E2E_UI_ARTIFACT_DIR" \
|
||||
LITELLM_MASTER_KEY="$LITELLM_MASTER_KEY" CI=true \
|
||||
node tests/e2e/ui/node_modules/@playwright/test/cli.js test --config tests/e2e/ui/integration.config.ts
|
||||
.venv/bin/python .circleci/scripts/verify_integration_browser.py "$results/browser-results.json"
|
||||
exit 0
|
||||
fi
|
||||
|
||||
timeout --signal=TERM --kill-after=20s 11m env -i PATH="$PATH" HOME="$HOME" PYTHONPATH="$PYTHONPATH" \
|
||||
INTEGRATION_RUN_ID="$integration_identity" \
|
||||
DATABASE_URL="$DATABASE_URL" REDIS_HOST="$REDIS_HOST" REDIS_PORT="$REDIS_PORT" \
|
||||
|
|
@ -138,5 +162,6 @@ timeout --signal=TERM --kill-after=20s 11m env -i PATH="$PATH" HOME="$HOME" PYTH
|
|||
INTEGRATION_UPSTREAM_URL="$INTEGRATION_UPSTREAM_URL" \
|
||||
INTEGRATION_MASTER_KEY="$INTEGRATION_MASTER_KEY" LITELLM_MODE=PRODUCTION \
|
||||
INTEGRATION_SEED="$INTEGRATION_SEED" \
|
||||
INTEGRATION_ORDER_SEED="$INTEGRATION_ORDER_SEED" \
|
||||
LITELLM_LOCAL_MODEL_COST_MAP=True AWS_EC2_METADATA_DISABLED=true DO_NOT_TRACK=1 \
|
||||
.venv/bin/python tests/integration/run.py "$suite" --results "$results"
|
||||
|
|
|
|||
60
.circleci/scripts/verify_integration_browser.py
Normal file
60
.circleci/scripts/verify_integration_browser.py
Normal file
|
|
@ -0,0 +1,60 @@
|
|||
import json
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
|
||||
from pydantic import TypeAdapter
|
||||
from typing_extensions import NotRequired, ReadOnly, TypedDict
|
||||
|
||||
|
||||
class BrowserAttempt(TypedDict):
|
||||
status: ReadOnly[str]
|
||||
retry: ReadOnly[int]
|
||||
|
||||
|
||||
class BrowserTest(TypedDict):
|
||||
results: ReadOnly[list[BrowserAttempt]]
|
||||
|
||||
|
||||
class BrowserSpec(TypedDict):
|
||||
file: ReadOnly[str]
|
||||
title: ReadOnly[str]
|
||||
tests: ReadOnly[list[BrowserTest]]
|
||||
|
||||
|
||||
class BrowserSuite(TypedDict):
|
||||
specs: NotRequired[ReadOnly[list[BrowserSpec]]]
|
||||
suites: NotRequired[ReadOnly[list["BrowserSuite"]]]
|
||||
|
||||
|
||||
def main() -> None:
|
||||
result: Final = json.loads(Path(sys.argv[1]).read_text())
|
||||
assert not result.get("errors"), result.get("errors")
|
||||
expected: Final = json.loads(
|
||||
(Path(__file__).resolve().parents[2] / "tests/integration/contracts.json").read_text()
|
||||
)["browser"]
|
||||
assert expected and result["stats"]["expected"] == len(expected)
|
||||
assert all(result["stats"][name] == 0 for name in ("unexpected", "flaky", "skipped"))
|
||||
|
||||
def cases(suite: BrowserSuite) -> tuple[BrowserSpec, ...]:
|
||||
return tuple(suite.get("specs", ())) + tuple(spec for child in suite.get("suites", ()) for spec in cases(child))
|
||||
|
||||
suites: Final = TypeAdapter(list[BrowserSuite]).validate_python(result["suites"], strict=True)
|
||||
specs: Final = tuple(spec for suite in suites for spec in cases(suite))
|
||||
repository: Final = Path(__file__).resolve().parents[2]
|
||||
report_root: Final = Path(result["config"]["rootDir"])
|
||||
assert report_root.is_absolute(), "Playwright rootDir must be explicit"
|
||||
observed: Final = tuple(
|
||||
str((report_root / spec["file"]).resolve().relative_to(repository)) + "::" + spec["title"] for spec in specs
|
||||
)
|
||||
assert sorted(observed) == sorted(expected)
|
||||
for spec in specs:
|
||||
tests: Final = spec["tests"]
|
||||
assert len(tests) == 1 and len(tests[0]["results"]) == 1
|
||||
assert tests[0]["results"][0]["status"] == "passed" and tests[0]["results"][0]["retry"] == 0
|
||||
|
||||
sys.stdout.write("One canonical browser contract passed once without skips or retries\n")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
39
.github/scripts/assert_ci_coverage.py
vendored
39
.github/scripts/assert_ci_coverage.py
vendored
|
|
@ -505,6 +505,7 @@ def _integration_ownership(repo_root: pathlib.Path = REPO_ROOT) -> tuple[frozens
|
|||
return frozenset(), ()
|
||||
entries: Final = json.loads(manifest.read_text())
|
||||
paths: Final = frozenset(node.split("::", 1)[0] for node in entries["tests"])
|
||||
browser_paths: Final = frozenset(node.split("::", 1)[0] for node in entries.get("browser", {}))
|
||||
circle_path: Final = repo_root / ".circleci/config.yml"
|
||||
circle: Final = yaml.safe_load(circle_path.read_text()) if circle_path.exists() else {}
|
||||
steps: Final = circle.get("jobs", {}).get("integration_contracts", {}).get("steps", ())
|
||||
|
|
@ -523,7 +524,7 @@ def _integration_ownership(repo_root: pathlib.Path = REPO_ROOT) -> tuple[frozens
|
|||
.get("suite", (job["integration_contracts"].get("suite"),))
|
||||
if isinstance(suite, str)
|
||||
)
|
||||
required: Final = frozenset(
|
||||
required: Final = (frozenset({"browser"}) if browser_paths else frozenset()) | frozenset(
|
||||
group
|
||||
for group, folders in entries["groups"].items()
|
||||
if any(any(path.startswith(f"tests/integration/{folder}/") for folder in folders) for path in paths)
|
||||
|
|
@ -551,6 +552,40 @@ def _integration_ownership(repo_root: pathlib.Path = REPO_ROOT) -> tuple[frozens
|
|||
for path in paths
|
||||
if not (repo_root / path).is_file()
|
||||
)
|
||||
browser_commands: Final = tuple(
|
||||
scalar.value
|
||||
for path in (repo_root / ".github/workflows").glob("*.y*ml")
|
||||
for scalar in _scalars(yaml.safe_load(path.read_text()), path.name)
|
||||
if scalar.key in {"run", "command"}
|
||||
)
|
||||
browser_findings: Final = tuple(
|
||||
Finding(path, "browser integration contract is explicitly selected by GitHub Actions")
|
||||
for path in browser_paths
|
||||
if any(
|
||||
path in command
|
||||
or pathlib.Path(path).name in command
|
||||
or "integrationCritical" in command
|
||||
or "integration.config.ts" in command
|
||||
or ("run_integration.sh" in command and "browser" in command)
|
||||
for command in browser_commands
|
||||
)
|
||||
) + tuple(
|
||||
Finding(path, "canonical browser integration file is missing")
|
||||
for path in browser_paths
|
||||
if not (repo_root / path).is_file()
|
||||
)
|
||||
default_browser: Final = repo_root / "tests/e2e/ui/playwright.config.ts"
|
||||
exclusion_findings: Final = (
|
||||
(
|
||||
Finding(
|
||||
str(default_browser.relative_to(repo_root)),
|
||||
"default Playwright selection must exclude integrationCritical",
|
||||
),
|
||||
)
|
||||
if browser_paths
|
||||
and (not default_browser.exists() or "**/integrationCritical/**" not in default_browser.read_text())
|
||||
else ()
|
||||
)
|
||||
group_findings: Final = tuple(
|
||||
Finding(group, "canonical integration group is not scheduled by CircleCI")
|
||||
for group in sorted(required - scheduled)
|
||||
|
|
@ -559,7 +594,7 @@ def _integration_ownership(repo_root: pathlib.Path = REPO_ROOT) -> tuple[frozens
|
|||
return frozenset(), findings + (
|
||||
Finding(str(manifest.relative_to(repo_root)), "dedicated CircleCI runner is missing"),
|
||||
)
|
||||
return paths, findings + group_findings
|
||||
return paths | browser_paths, findings + group_findings + browser_findings + exclusion_findings
|
||||
|
||||
|
||||
def main() -> int:
|
||||
|
|
|
|||
465
.github/scripts/auto_merge_price_sync.py
vendored
Normal file
465
.github/scripts/auto_merge_price_sync.py
vendored
Normal file
|
|
@ -0,0 +1,465 @@
|
|||
"""Auto-merge the provider-info-sync bot's cost-map pull requests.
|
||||
|
||||
Evaluates every gate (author allowlist, cost-map-only diff, required and
|
||||
non-required checks, Greptile confidence, Bugbot review, human reviews) and
|
||||
merges with a merge commit when all of them hold. Every hold reason is
|
||||
logged; the process exits 0 on hold and 1 only on API or programming errors.
|
||||
``DRY_RUN=1`` prints the verdict without calling the merge endpoint.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import subprocess
|
||||
import sys
|
||||
import time
|
||||
import urllib.error
|
||||
import urllib.request
|
||||
from collections.abc import Callable, Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime, timezone
|
||||
from typing import Final
|
||||
|
||||
REPO_ROOT: Final = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||||
CLASSIFY_SCRIPT: Final = os.path.join(REPO_ROOT, ".circleci", "scripts", "classify_changes.sh")
|
||||
API_ROOT: Final = "https://api.github.com"
|
||||
CHANGED_FILE_CEILING: Final = 3000
|
||||
OK_CHECK_CONCLUSIONS: Final = frozenset({"success", "skipped", "neutral"})
|
||||
GREPTILE_LOGIN: Final = "greptile-apps[bot]"
|
||||
BUGBOT_LOGIN: Final = "cursor[bot]"
|
||||
GREPTILE_SCORE_RE: Final = re.compile(r"Confidence Score:\s*(\d)/5")
|
||||
BUGBOT_REVIEW_MARKER: Final = "<!-- BUGBOT_REVIEW -->"
|
||||
BUGBOT_STALE_MARKER: Final = "<!-- BUGBOT_REVIEW_STALE -->"
|
||||
BUGBOT_CLEAN: Final = "found no new issues"
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class PullRequest:
|
||||
number: int
|
||||
title: str
|
||||
author_login: str
|
||||
state: str
|
||||
draft: bool
|
||||
mergeable: bool | None
|
||||
mergeable_state: str
|
||||
head_sha: str
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class CheckRun:
|
||||
name: str
|
||||
status: str
|
||||
conclusion: str | None
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class CommitStatus:
|
||||
context: str
|
||||
state: str
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class IssueComment:
|
||||
author_login: str
|
||||
body: str
|
||||
updated_at: datetime
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Review:
|
||||
author_login: str
|
||||
state: str
|
||||
body: str
|
||||
commit_id: str
|
||||
submitted_at: datetime
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Verdict:
|
||||
merge: bool
|
||||
reasons: tuple[str, ...]
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class EvaluationInputs:
|
||||
pr: PullRequest
|
||||
changed_files: tuple[str, ...]
|
||||
required_contexts: frozenset[str]
|
||||
check_runs: tuple[CheckRun, ...]
|
||||
statuses: tuple[CommitStatus, ...]
|
||||
comments: tuple[IssueComment, ...]
|
||||
reviews: tuple[Review, ...]
|
||||
head_commit_date: datetime
|
||||
self_check_name: str
|
||||
author_allowlist: frozenset[str]
|
||||
|
||||
|
||||
def _is_bot_login(login: str) -> bool:
|
||||
return login.lower().endswith("[bot]")
|
||||
|
||||
|
||||
def _classify(changed_files: Sequence[str]) -> str:
|
||||
result: Final = subprocess.run(
|
||||
["bash", CLASSIFY_SCRIPT, "cost-map-only"],
|
||||
input="\n".join(changed_files),
|
||||
capture_output=True,
|
||||
text=True,
|
||||
check=False,
|
||||
)
|
||||
if result.returncode != 0:
|
||||
return "error"
|
||||
return result.stdout.strip()
|
||||
|
||||
|
||||
def evaluate(
|
||||
inputs: EvaluationInputs,
|
||||
*,
|
||||
classify: Callable[[Sequence[str]], str] = _classify,
|
||||
) -> Verdict:
|
||||
pr: Final = inputs.pr
|
||||
reasons: list[str] = []
|
||||
|
||||
if pr.author_login.lower() not in {login.lower() for login in inputs.author_allowlist}:
|
||||
reasons.append(f"author {pr.author_login!r} not in allowlist")
|
||||
if pr.state != "open":
|
||||
reasons.append("pr not open")
|
||||
if pr.draft:
|
||||
reasons.append("pr is a draft")
|
||||
if pr.mergeable is None:
|
||||
reasons.append("mergeability unknown")
|
||||
elif not pr.mergeable:
|
||||
reasons.append("pr not mergeable")
|
||||
if pr.mergeable_state == "dirty":
|
||||
reasons.append("pr has merge conflicts")
|
||||
|
||||
if len(inputs.changed_files) > CHANGED_FILE_CEILING:
|
||||
reasons.append(f"changed file count {len(inputs.changed_files)} over {CHANGED_FILE_CEILING} ceiling")
|
||||
else:
|
||||
decision: Final = classify(inputs.changed_files)
|
||||
if decision != "run":
|
||||
reasons.append("changed files outside the cost-map-only set")
|
||||
|
||||
green_runs: Final = frozenset(run.name for run in inputs.check_runs if run.conclusion in OK_CHECK_CONCLUSIONS)
|
||||
green_statuses: Final = frozenset(status.context for status in inputs.statuses if status.state == "success")
|
||||
for context in sorted(inputs.required_contexts):
|
||||
if context not in green_runs and context not in green_statuses:
|
||||
reasons.append(f"required check {context!r} not green")
|
||||
for run in inputs.check_runs:
|
||||
if run.name == inputs.self_check_name:
|
||||
continue
|
||||
if run.status != "completed" or run.conclusion not in OK_CHECK_CONCLUSIONS:
|
||||
reasons.append(f"check run {run.name!r} is {run.status}/{run.conclusion}")
|
||||
for status in inputs.statuses:
|
||||
if status.state != "success":
|
||||
reasons.append(f"commit status {status.context!r} is {status.state}")
|
||||
|
||||
greptile: Final = tuple(
|
||||
comment
|
||||
for comment in inputs.comments
|
||||
if comment.author_login == GREPTILE_LOGIN and GREPTILE_SCORE_RE.search(comment.body)
|
||||
)
|
||||
if not greptile:
|
||||
reasons.append("greptile score not available")
|
||||
else:
|
||||
latest: Final = max(greptile, key=lambda comment: comment.updated_at)
|
||||
match: Final = GREPTILE_SCORE_RE.search(latest.body)
|
||||
score: Final = int(match.group(1)) if match else 0
|
||||
if latest.updated_at < inputs.head_commit_date:
|
||||
reasons.append("greptile score older than head commit")
|
||||
elif score != 5:
|
||||
reasons.append(f"greptile score {score}/5 below 5")
|
||||
|
||||
bugbot: Final = tuple(
|
||||
review
|
||||
for review in inputs.reviews
|
||||
if review.author_login == BUGBOT_LOGIN
|
||||
and BUGBOT_REVIEW_MARKER in review.body
|
||||
and BUGBOT_STALE_MARKER not in review.body
|
||||
and review.commit_id == pr.head_sha
|
||||
)
|
||||
if not bugbot:
|
||||
reasons.append("bugbot review not available")
|
||||
else:
|
||||
latest_review: Final = max(bugbot, key=lambda review: review.submitted_at)
|
||||
if BUGBOT_CLEAN not in latest_review.body:
|
||||
reasons.append("bugbot reported issues")
|
||||
|
||||
latest_state_by_reviewer: Final[dict[str, str]] = {}
|
||||
for review in sorted(inputs.reviews, key=lambda review: review.submitted_at):
|
||||
if _is_bot_login(review.author_login):
|
||||
continue
|
||||
latest_state_by_reviewer[review.author_login] = review.state
|
||||
for reviewer, state in latest_state_by_reviewer.items():
|
||||
if state == "CHANGES_REQUESTED":
|
||||
reasons.append(f"changes requested by {reviewer}")
|
||||
|
||||
return Verdict(merge=not reasons, reasons=tuple(reasons))
|
||||
|
||||
|
||||
def _request(token: str, method: str, path: str, body: Mapping[str, object] | None = None) -> object:
|
||||
url: Final = path if path.startswith("http") else f"{API_ROOT}{path}"
|
||||
data: Final = None if body is None else json.dumps(body).encode("utf-8")
|
||||
request: Final = urllib.request.Request(
|
||||
url,
|
||||
data=data,
|
||||
method=method,
|
||||
headers={
|
||||
"Accept": "application/vnd.github+json",
|
||||
"Authorization": f"Bearer {token}",
|
||||
"X-GitHub-Api-Version": "2022-11-28",
|
||||
},
|
||||
)
|
||||
with urllib.request.urlopen(request) as response:
|
||||
return json.loads(response.read().decode("utf-8"))
|
||||
|
||||
|
||||
def _request_allow_fail(
|
||||
token: str, method: str, path: str, body: Mapping[str, object] | None = None
|
||||
) -> tuple[int, object | None]:
|
||||
url: Final = path if path.startswith("http") else f"{API_ROOT}{path}"
|
||||
data: Final = None if body is None else json.dumps(body).encode("utf-8")
|
||||
request: Final = urllib.request.Request(
|
||||
url,
|
||||
data=data,
|
||||
method=method,
|
||||
headers={
|
||||
"Accept": "application/vnd.github+json",
|
||||
"Authorization": f"Bearer {token}",
|
||||
"X-GitHub-Api-Version": "2022-11-28",
|
||||
},
|
||||
)
|
||||
try:
|
||||
with urllib.request.urlopen(request) as response:
|
||||
return response.status, json.loads(response.read().decode("utf-8"))
|
||||
except urllib.error.HTTPError as exc:
|
||||
return exc.code, None
|
||||
|
||||
|
||||
def _items(payload: object, key: str | None = None) -> tuple[object, ...]:
|
||||
source: Final = payload.get(key) if key and isinstance(payload, Mapping) else payload
|
||||
if not isinstance(source, list):
|
||||
return ()
|
||||
return tuple(source)
|
||||
|
||||
|
||||
def _paginate(token: str, path: str, key: str | None = None) -> list[object]:
|
||||
separator: Final = "&" if "?" in path else "?"
|
||||
results: list[object] = []
|
||||
for page in range(1, 10_000):
|
||||
batch: Final = _items(_request(token, "GET", f"{path}{separator}per_page=100&page={page}"), key)
|
||||
results.extend(batch)
|
||||
if len(batch) < 100:
|
||||
return results
|
||||
return results
|
||||
|
||||
|
||||
def _text(value: object) -> str:
|
||||
return value if isinstance(value, str) else ""
|
||||
|
||||
|
||||
def _int(value: object) -> int:
|
||||
return value if isinstance(value, int) else 0
|
||||
|
||||
|
||||
def _bool(value: object) -> bool:
|
||||
return value is True
|
||||
|
||||
|
||||
def _nested(value: object, *keys: str) -> object:
|
||||
current: object = value
|
||||
for key in keys:
|
||||
if not isinstance(current, Mapping):
|
||||
return None
|
||||
current = current.get(key)
|
||||
return current
|
||||
|
||||
|
||||
def _parse_time(value: object) -> datetime:
|
||||
text: Final = _text(value)
|
||||
if not text:
|
||||
return datetime.min.replace(tzinfo=timezone.utc)
|
||||
return datetime.fromisoformat(text.replace("Z", "+00:00"))
|
||||
|
||||
|
||||
def _load_pr(token: str, repo: str, number: int) -> PullRequest:
|
||||
data: Final = _request(token, "GET", f"/repos/{repo}/pulls/{number}")
|
||||
if not isinstance(data, Mapping):
|
||||
raise RuntimeError(f"unexpected pull payload for #{number}")
|
||||
return PullRequest(
|
||||
number=number,
|
||||
title=_text(data.get("title")),
|
||||
author_login=_text(_nested(data, "user", "login")),
|
||||
state=_text(data.get("state")),
|
||||
draft=_bool(data.get("draft")),
|
||||
mergeable=data.get("mergeable") if isinstance(data.get("mergeable"), bool) else None,
|
||||
mergeable_state=_text(data.get("mergeable_state")),
|
||||
head_sha=_text(_nested(data, "head", "sha")),
|
||||
)
|
||||
|
||||
|
||||
def _list_candidate_prs(token: str, repo: str, base: str, allowlist: frozenset[str]) -> list[int]:
|
||||
candidates: Final = _paginate(token, f"/repos/{repo}/pulls?state=open&base={base}")
|
||||
return [
|
||||
_int(item.get("number"))
|
||||
for item in candidates
|
||||
if isinstance(item, Mapping) and _text(_nested(item, "user", "login")).lower() in allowlist
|
||||
]
|
||||
|
||||
|
||||
def _changed_files(token: str, repo: str, number: int) -> tuple[str, ...]:
|
||||
files: Final = _paginate(token, f"/repos/{repo}/pulls/{number}/files")
|
||||
return tuple(_text(item.get("filename")) for item in files if isinstance(item, Mapping))
|
||||
|
||||
|
||||
def _required_contexts(token: str, repo: str, base: str) -> frozenset[str]:
|
||||
payload: Final = _request(token, "GET", f"/repos/{repo}/rules/branches/{base}")
|
||||
contexts: set[str] = set()
|
||||
for rule in _items(payload):
|
||||
if not isinstance(rule, Mapping) or rule.get("type") != "required_status_checks":
|
||||
continue
|
||||
checks: Final = _nested(rule, "parameters", "required_status_checks")
|
||||
for check in _items(checks):
|
||||
if isinstance(check, Mapping):
|
||||
context: Final = _text(check.get("context"))
|
||||
if context:
|
||||
contexts.add(context)
|
||||
return frozenset(contexts)
|
||||
|
||||
|
||||
def _check_runs(token: str, repo: str, sha: str) -> tuple[CheckRun, ...]:
|
||||
runs: Final = _paginate(token, f"/repos/{repo}/commits/{sha}/check-runs", key="check_runs")
|
||||
return tuple(
|
||||
CheckRun(
|
||||
name=_text(item.get("name")),
|
||||
status=_text(item.get("status")),
|
||||
conclusion=item.get("conclusion") if isinstance(item.get("conclusion"), str) else None,
|
||||
)
|
||||
for item in runs
|
||||
if isinstance(item, Mapping)
|
||||
)
|
||||
|
||||
|
||||
def _statuses(token: str, repo: str, sha: str) -> tuple[CommitStatus, ...]:
|
||||
payload: Final = _request(token, "GET", f"/repos/{repo}/commits/{sha}/status")
|
||||
return tuple(
|
||||
CommitStatus(context=_text(item.get("context")), state=_text(item.get("state")))
|
||||
for item in _items(payload, "statuses")
|
||||
if isinstance(item, Mapping)
|
||||
)
|
||||
|
||||
|
||||
def _comments(token: str, repo: str, number: int) -> tuple[IssueComment, ...]:
|
||||
comments: Final = _paginate(token, f"/repos/{repo}/issues/{number}/comments")
|
||||
return tuple(
|
||||
IssueComment(
|
||||
author_login=_text(_nested(item, "user", "login")),
|
||||
body=_text(item.get("body")),
|
||||
updated_at=_parse_time(item.get("updated_at")),
|
||||
)
|
||||
for item in comments
|
||||
if isinstance(item, Mapping)
|
||||
)
|
||||
|
||||
|
||||
def _reviews(token: str, repo: str, number: int) -> tuple[Review, ...]:
|
||||
reviews: Final = _paginate(token, f"/repos/{repo}/pulls/{number}/reviews")
|
||||
return tuple(
|
||||
Review(
|
||||
author_login=_text(_nested(item, "user", "login")),
|
||||
state=_text(item.get("state")),
|
||||
body=_text(item.get("body")),
|
||||
commit_id=_text(item.get("commit_id")),
|
||||
submitted_at=_parse_time(item.get("submitted_at")),
|
||||
)
|
||||
for item in reviews
|
||||
if isinstance(item, Mapping)
|
||||
)
|
||||
|
||||
|
||||
def _head_commit_date(token: str, repo: str, number: int) -> datetime:
|
||||
commits: Final = _paginate(token, f"/repos/{repo}/pulls/{number}/commits")
|
||||
if not commits:
|
||||
return datetime.min.replace(tzinfo=timezone.utc)
|
||||
last: Final = commits[-1]
|
||||
if not isinstance(last, Mapping):
|
||||
return datetime.min.replace(tzinfo=timezone.utc)
|
||||
return _parse_time(_nested(last, "commit", "committer", "date"))
|
||||
|
||||
|
||||
def _mergeable_or_refetch(token: str, repo: str, pr: PullRequest) -> PullRequest:
|
||||
if pr.mergeable is not None:
|
||||
return pr
|
||||
time.sleep(5)
|
||||
return _load_pr(token, repo, pr.number)
|
||||
|
||||
|
||||
def _gather_inputs(
|
||||
token: str,
|
||||
repo: str,
|
||||
number: int,
|
||||
base: str,
|
||||
self_check_name: str,
|
||||
allowlist: frozenset[str],
|
||||
) -> EvaluationInputs:
|
||||
pr: Final = _mergeable_or_refetch(token, repo, _load_pr(token, repo, number))
|
||||
return EvaluationInputs(
|
||||
pr=pr,
|
||||
changed_files=_changed_files(token, repo, number),
|
||||
required_contexts=_required_contexts(token, repo, base),
|
||||
check_runs=_check_runs(token, repo, pr.head_sha),
|
||||
statuses=_statuses(token, repo, pr.head_sha),
|
||||
comments=_comments(token, repo, number),
|
||||
reviews=_reviews(token, repo, number),
|
||||
head_commit_date=_head_commit_date(token, repo, number),
|
||||
self_check_name=self_check_name,
|
||||
author_allowlist=allowlist,
|
||||
)
|
||||
|
||||
|
||||
def merge_request_body(pr: PullRequest) -> dict[str, str]:
|
||||
return {"merge_method": "merge", "commit_title": f"{pr.title} (#{pr.number})", "sha": pr.head_sha}
|
||||
|
||||
|
||||
def _merge(token: str, repo: str, pr: PullRequest) -> None:
|
||||
status, _ = _request_allow_fail(token, "PUT", f"/repos/{repo}/pulls/{pr.number}/merge", merge_request_body(pr))
|
||||
if status in (200, 405, 409):
|
||||
print(f"auto-merge-price-sync: PR #{pr.number} merge call returned {status}")
|
||||
return
|
||||
raise RuntimeError(f"merge call for PR #{pr.number} returned {status}")
|
||||
|
||||
|
||||
def main() -> int:
|
||||
token: Final = os.environ.get("GH_TOKEN", "")
|
||||
repo: Final = os.environ.get("REPO", "")
|
||||
base: Final = os.environ.get("BASE_BRANCH", "main")
|
||||
dry_run: Final = os.environ.get("DRY_RUN", "") != ""
|
||||
self_check_name: Final = os.environ.get("SELF_CHECK_NAME", "auto-merge-price-sync")
|
||||
allowlist: Final = frozenset(login.lower() for login in os.environ.get("PR_AUTHOR_ALLOWLIST", "").split() if login)
|
||||
if not token:
|
||||
print("auto-merge-price-sync: app credentials not configured")
|
||||
return 0
|
||||
if not repo:
|
||||
print("auto-merge-price-sync: REPO not set", file=sys.stderr)
|
||||
return 1
|
||||
|
||||
pr_number_env: Final = os.environ.get("PR_NUMBER", "")
|
||||
candidates: Final = [int(pr_number_env)] if pr_number_env else _list_candidate_prs(token, repo, base, allowlist)
|
||||
for number in candidates:
|
||||
inputs: Final = _gather_inputs(token, repo, number, base, self_check_name, allowlist)
|
||||
verdict: Final = evaluate(inputs)
|
||||
for reason in verdict.reasons:
|
||||
print(f"auto-merge-price-sync: PR #{number} hold: {reason}")
|
||||
if not verdict.merge:
|
||||
continue
|
||||
print(f"auto-merge-price-sync: PR #{number} all gates green")
|
||||
if dry_run:
|
||||
print(f"auto-merge-price-sync: DRY_RUN merge suppressed for PR #{number}")
|
||||
continue
|
||||
_merge(token, repo, inputs.pr)
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
61
.github/workflows/auto-merge-price-sync.yml
vendored
Normal file
61
.github/workflows/auto-merge-price-sync.yml
vendored
Normal file
|
|
@ -0,0 +1,61 @@
|
|||
name: auto-merge-price-sync
|
||||
|
||||
on:
|
||||
issue_comment:
|
||||
types: [created, edited]
|
||||
check_suite:
|
||||
types: [completed]
|
||||
status: {}
|
||||
schedule:
|
||||
- cron: "*/30 * * * *"
|
||||
workflow_dispatch:
|
||||
inputs:
|
||||
pr-number:
|
||||
description: "Evaluate only this PR number (empty = scan all open sync-bot PRs)"
|
||||
required: false
|
||||
default: ""
|
||||
|
||||
permissions:
|
||||
contents: read
|
||||
pull-requests: read
|
||||
checks: read
|
||||
statuses: read
|
||||
|
||||
concurrency:
|
||||
group: auto-merge-price-sync
|
||||
cancel-in-progress: false
|
||||
|
||||
jobs:
|
||||
auto-merge-price-sync:
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 15
|
||||
env:
|
||||
PROVIDER_INFO_SYNC_APP_ID: ${{ secrets.PROVIDER_INFO_SYNC_APP_ID }}
|
||||
PROVIDER_INFO_SYNC_APP_PRIVATE_KEY: ${{ secrets.PROVIDER_INFO_SYNC_APP_PRIVATE_KEY }}
|
||||
steps:
|
||||
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
|
||||
with:
|
||||
persist-credentials: false
|
||||
|
||||
- name: Set up Python
|
||||
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
|
||||
with:
|
||||
python-version: "3.12"
|
||||
|
||||
- name: Mint app token
|
||||
id: app-token
|
||||
if: ${{ env.PROVIDER_INFO_SYNC_APP_ID != '' && env.PROVIDER_INFO_SYNC_APP_PRIVATE_KEY != '' }}
|
||||
uses: actions/create-github-app-token@bcd2ba49218906704ab6c1aa796996da409d3eb1 # v3.2.0
|
||||
with:
|
||||
app-id: ${{ secrets.PROVIDER_INFO_SYNC_APP_ID }}
|
||||
private-key: ${{ secrets.PROVIDER_INFO_SYNC_APP_PRIVATE_KEY }}
|
||||
|
||||
- name: Auto-merge eligible sync PRs
|
||||
env:
|
||||
GH_TOKEN: ${{ steps.app-token.outputs.token }}
|
||||
REPO: ${{ github.repository }}
|
||||
PR_NUMBER: ${{ (github.event.issue.pull_request && github.event.issue.number) || github.event.inputs.pr-number || '' }}
|
||||
BASE_BRANCH: main
|
||||
PR_AUTHOR_ALLOWLIST: "berriai-litellm-provider-info-sync[bot]"
|
||||
SELF_CHECK_NAME: auto-merge-price-sync
|
||||
run: python3 .github/scripts/auto_merge_price_sync.py
|
||||
|
|
@ -45,7 +45,7 @@ sequenceDiagram
|
|||
ProxyServer->>Auth: user_api_key_auth()
|
||||
Auth->>Redis: Check API key cache
|
||||
Redis-->>Auth: Key info + spend limits
|
||||
ProxyServer->>Hooks: max_budget_limiter, parallel_request_limiter
|
||||
ProxyServer->>Hooks: parallel_request_limiter, cache_control_check
|
||||
Hooks->>Redis: Check/increment rate limit counters
|
||||
ProxyServer->>Router: route_request()
|
||||
Router->>Main: litellm.acompletion()
|
||||
|
|
@ -145,7 +145,6 @@ graph TD
|
|||
|
||||
| Hook | File | Purpose |
|
||||
|------|------|---------|
|
||||
| `max_budget_limiter` | `proxy/hooks/max_budget_limiter.py` | Enforce budget limits |
|
||||
| `parallel_request_limiter` | `proxy/hooks/parallel_request_limiter_v3.py` | Rate limiting per key/user |
|
||||
| `cache_control_check` | `proxy/hooks/cache_control_check.py` | Cache validation |
|
||||
| `responses_id_security` | `proxy/hooks/responses_id_security.py` | Response ID validation |
|
||||
|
|
|
|||
143
litellm-rust/Cargo.lock
generated
143
litellm-rust/Cargo.lock
generated
|
|
@ -70,6 +70,12 @@ dependencies = [
|
|||
"rustversion",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "arcstr"
|
||||
version = "1.2.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "03918c3dbd7701a85c6b9887732e2921175f26c350b4563841d0958c21d57e6d"
|
||||
|
||||
[[package]]
|
||||
name = "async-compression"
|
||||
version = "0.4.46"
|
||||
|
|
@ -262,6 +268,17 @@ dependencies = [
|
|||
"tokio",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "aws-smithy-eventstream"
|
||||
version = "0.61.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "5a9381123ab62d20c13082b151f30f962a3b112b727345394536dfa39a482944"
|
||||
dependencies = [
|
||||
"aws-smithy-types",
|
||||
"bytes",
|
||||
"crc32fast",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "aws-smithy-http"
|
||||
version = "0.64.0"
|
||||
|
|
@ -1837,6 +1854,12 @@ version = "0.2.186"
|
|||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "68ab91017fe16c622486840e4c83c9a37afeff978bd239b5293d61ece587de66"
|
||||
|
||||
[[package]]
|
||||
name = "linux-raw-sys"
|
||||
version = "0.12.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "32a66949e030da00e8c7d4434b251670a91556f4144941d37452769c25d58a53"
|
||||
|
||||
[[package]]
|
||||
name = "litellm-auth"
|
||||
version = "0.1.0"
|
||||
|
|
@ -1915,10 +1938,23 @@ dependencies = [
|
|||
"tokio",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "litellm-cache-redis"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"litellm-cache",
|
||||
"redis",
|
||||
"redis-test",
|
||||
"serde_json",
|
||||
"tokio",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "litellm-core"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"aws-smithy-eventstream",
|
||||
"aws-smithy-types",
|
||||
"base64 0.22.1",
|
||||
"bytes",
|
||||
"data-url",
|
||||
|
|
@ -1927,6 +1963,7 @@ dependencies = [
|
|||
"litellm-auth-aws",
|
||||
"litellm-auth-azure",
|
||||
"litellm-auth-gcp",
|
||||
"litellm-framing",
|
||||
"mime_guess",
|
||||
"moka",
|
||||
"rand 0.8.7",
|
||||
|
|
@ -1941,12 +1978,27 @@ dependencies = [
|
|||
"strum",
|
||||
"subtle",
|
||||
"thiserror 2.0.19",
|
||||
"time",
|
||||
"tokio",
|
||||
"tokio-tungstenite",
|
||||
"url",
|
||||
"veil",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "litellm-framing"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"aws-smithy-eventstream",
|
||||
"aws-smithy-types",
|
||||
"bytes",
|
||||
"futures-util",
|
||||
"rstest",
|
||||
"sse-stream",
|
||||
"thiserror 2.0.19",
|
||||
"tokio",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "litellm-python-bridge"
|
||||
version = "0.1.0"
|
||||
|
|
@ -2140,6 +2192,16 @@ dependencies = [
|
|||
"minimal-lexical",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "num-bigint"
|
||||
version = "0.5.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "93e7820bc0a80a0238e650327316f929ba18d5be054b647490a3a6a339f3e7c0"
|
||||
dependencies = [
|
||||
"num-integer",
|
||||
"num-traits",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "num-conv"
|
||||
version = "0.2.2"
|
||||
|
|
@ -2656,6 +2718,36 @@ dependencies = [
|
|||
"crossbeam-utils",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "redis"
|
||||
version = "1.7.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "2acbc41a996f7652b2ddd9dfd98cc4ff602cfd742ae35382f07f608405ab50ed"
|
||||
dependencies = [
|
||||
"arcstr",
|
||||
"combine",
|
||||
"itoa",
|
||||
"num-bigint",
|
||||
"percent-encoding",
|
||||
"ryu",
|
||||
"sha1_smol",
|
||||
"socket2 0.6.5",
|
||||
"url",
|
||||
"xxhash-rust",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "redis-test"
|
||||
version = "1.0.4"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "804d36862e4323b69f96440cbb13c9894fc90176abdeaf91264e21d5d77f6aca"
|
||||
dependencies = [
|
||||
"rand 0.9.5",
|
||||
"redis",
|
||||
"socket2 0.6.5",
|
||||
"tempfile",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "redox_syscall"
|
||||
version = "0.5.18"
|
||||
|
|
@ -2846,6 +2938,19 @@ dependencies = [
|
|||
"semver",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "rustix"
|
||||
version = "1.1.5"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "891efababe418670775f199f0d233d84843c227a0949a883ce15b37c78d6629d"
|
||||
dependencies = [
|
||||
"bitflags",
|
||||
"errno",
|
||||
"libc",
|
||||
"linux-raw-sys",
|
||||
"windows-sys 0.61.2",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "rustls"
|
||||
version = "0.21.12"
|
||||
|
|
@ -3096,6 +3201,12 @@ dependencies = [
|
|||
"digest 0.10.7",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "sha1_smol"
|
||||
version = "1.0.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "bbfa15b3dddfee50a0fff136974b3e1bde555604ba463834a7eb7deb6417705d"
|
||||
|
||||
[[package]]
|
||||
name = "sha2"
|
||||
version = "0.10.9"
|
||||
|
|
@ -3200,6 +3311,19 @@ dependencies = [
|
|||
"unicode-segmentation",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "sse-stream"
|
||||
version = "0.2.6"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "c25ac7aff0abd1dbc474536e40416e1102c7dd9bfba0b9861c6d357f835dcfb4"
|
||||
dependencies = [
|
||||
"bytes",
|
||||
"futures-util",
|
||||
"http-body 1.1.0",
|
||||
"http-body-util",
|
||||
"pin-project-lite",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "stable_deref_trait"
|
||||
version = "1.2.1"
|
||||
|
|
@ -3299,6 +3423,19 @@ version = "0.13.5"
|
|||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "adb6935a6f5c20170eeceb1a3835a49e12e19d792f6dd344ccc76a985ca5a6ca"
|
||||
|
||||
[[package]]
|
||||
name = "tempfile"
|
||||
version = "3.27.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "32497e9a4c7b38532efcdebeef879707aa9f794296a4f0244f6f69e9bc8574bd"
|
||||
dependencies = [
|
||||
"fastrand",
|
||||
"getrandom 0.4.3",
|
||||
"once_cell",
|
||||
"rustix",
|
||||
"windows-sys 0.61.2",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "thiserror"
|
||||
version = "1.0.69"
|
||||
|
|
@ -4182,6 +4319,12 @@ version = "0.13.6"
|
|||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "66fee0b777b0f5ac1c69bb06d361268faafa61cd4682ae064a171c16c433e9e4"
|
||||
|
||||
[[package]]
|
||||
name = "xxhash-rust"
|
||||
version = "0.8.18"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "aee1b19627c7c60102ab80d3a9cbe18de90bfe03bfa6c3715447681f0e8c8af6"
|
||||
|
||||
[[package]]
|
||||
name = "yoke"
|
||||
version = "0.8.3"
|
||||
|
|
|
|||
|
|
@ -11,6 +11,7 @@ repository = "https://github.com/BerriAI/litellm"
|
|||
[workspace.dependencies]
|
||||
bytes = "1"
|
||||
litellm-core = { path = "crates/core" }
|
||||
litellm-framing = { path = "crates/framer" }
|
||||
litellm-auth = { path = "crates/auth" }
|
||||
litellm-auth-aws = { path = "crates/auth-aws" }
|
||||
litellm-auth-azure = { path = "crates/auth-azure" }
|
||||
|
|
@ -39,6 +40,7 @@ base64 = "0.22"
|
|||
moka = { version = "0.12.16", features = ["future"] }
|
||||
strum = { version = "0.28.0", features = ["derive"] }
|
||||
url = "2.5.8"
|
||||
time = { version = "0.3.53", features = ["parsing"] }
|
||||
criterion = "0.8.2"
|
||||
veil = "0.3.0"
|
||||
|
||||
|
|
|
|||
15
litellm-rust/crates/cache-redis/Cargo.toml
Normal file
15
litellm-rust/crates/cache-redis/Cargo.toml
Normal file
|
|
@ -0,0 +1,15 @@
|
|||
[package]
|
||||
name = "litellm-cache-redis"
|
||||
version = "0.1.0"
|
||||
edition.workspace = true
|
||||
license.workspace = true
|
||||
repository.workspace = true
|
||||
|
||||
[dependencies]
|
||||
litellm-cache.workspace = true
|
||||
redis = "1.7.0"
|
||||
serde_json.workspace = true
|
||||
tokio.workspace = true
|
||||
|
||||
[dev-dependencies]
|
||||
redis-test = "1.0.4"
|
||||
315
litellm-rust/crates/cache-redis/src/cache.rs
Normal file
315
litellm-rust/crates/cache-redis/src/cache.rs
Normal file
|
|
@ -0,0 +1,315 @@
|
|||
use std::sync::{Arc, Mutex, MutexGuard};
|
||||
use std::time::Duration;
|
||||
|
||||
use litellm_cache::{
|
||||
BaseCache, CacheConnectionResult, CacheConnectionStatus, CacheEntry, CacheFuture, CacheKwargs,
|
||||
Error,
|
||||
};
|
||||
use redis::Commands;
|
||||
|
||||
const DEFAULT_TTL: Duration = Duration::from_secs(600);
|
||||
const KEY_PREFIX: &str = "litellm-cache:";
|
||||
|
||||
pub struct RedisCache<C = redis::Connection> {
|
||||
connection: Arc<Mutex<C>>,
|
||||
default_ttl: Duration,
|
||||
}
|
||||
|
||||
impl RedisCache<redis::Connection> {
|
||||
pub fn new(url: &str, default_ttl: Option<Duration>) -> Result<Self, Error> {
|
||||
let client = redis::Client::open(url).map_err(|_| Error::Unavailable)?;
|
||||
let connection = client.get_connection().map_err(|_| Error::Unavailable)?;
|
||||
Ok(Self::with_connection(connection, default_ttl))
|
||||
}
|
||||
}
|
||||
|
||||
impl<C> RedisCache<C>
|
||||
where
|
||||
C: redis::ConnectionLike + Send + 'static,
|
||||
{
|
||||
fn with_connection(connection: C, default_ttl: Option<Duration>) -> Self {
|
||||
Self {
|
||||
connection: Arc::new(Mutex::new(connection)),
|
||||
default_ttl: default_ttl.unwrap_or(DEFAULT_TTL),
|
||||
}
|
||||
}
|
||||
|
||||
fn connection(&self) -> Result<MutexGuard<'_, C>, Error> {
|
||||
self.connection.lock().map_err(|_| Error::Unavailable)
|
||||
}
|
||||
|
||||
fn namespaced_key(key: &str) -> String {
|
||||
format!("{KEY_PREFIX}{key}")
|
||||
}
|
||||
|
||||
fn namespaced_pattern() -> &'static str {
|
||||
const PATTERN: &str = "litellm-cache:*";
|
||||
PATTERN
|
||||
}
|
||||
|
||||
fn encode(value: &CacheEntry) -> Result<Vec<u8>, Error> {
|
||||
serde_json::to_vec(value).map_err(|_| Error::InvalidEntry)
|
||||
}
|
||||
|
||||
fn decode(value: Vec<u8>) -> Result<CacheEntry, Error> {
|
||||
serde_json::from_slice(&value).map_err(|_| Error::InvalidEntry)
|
||||
}
|
||||
|
||||
fn ttl_seconds(ttl: Duration) -> u64 {
|
||||
ttl.as_secs()
|
||||
.saturating_add(u64::from(ttl.subsec_nanos() > 0))
|
||||
.max(1)
|
||||
}
|
||||
|
||||
fn run_blocking<T, F>(connection: Arc<Mutex<C>>, operation: F) -> CacheFuture<'static, T>
|
||||
where
|
||||
T: Send + 'static,
|
||||
F: FnOnce(&mut C) -> Result<T, Error> + Send + 'static,
|
||||
{
|
||||
Box::pin(async move {
|
||||
tokio::task::spawn_blocking(move || {
|
||||
let mut connection = connection.lock().map_err(|_| Error::Unavailable)?;
|
||||
operation(&mut connection)
|
||||
})
|
||||
.await
|
||||
.map_err(|_| Error::Unavailable)?
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
impl<C> BaseCache for RedisCache<C>
|
||||
where
|
||||
C: redis::ConnectionLike + Send + 'static,
|
||||
{
|
||||
type Value = CacheEntry;
|
||||
|
||||
fn default_ttl(&self) -> Duration {
|
||||
self.default_ttl
|
||||
}
|
||||
|
||||
fn set_cache(&self, key: &str, value: Self::Value, kwargs: CacheKwargs) -> Result<(), Error> {
|
||||
let payload = Self::encode(&value)?;
|
||||
let ttl = Self::ttl_seconds(self.get_ttl(&kwargs));
|
||||
self.connection()?
|
||||
.set_ex::<_, _, ()>(Self::namespaced_key(key), payload, ttl)
|
||||
.map_err(|_| Error::Unavailable)
|
||||
}
|
||||
|
||||
fn get_cache(&self, key: &str, _: &CacheKwargs) -> Result<Option<Self::Value>, Error> {
|
||||
self.connection()?
|
||||
.get::<_, Option<Vec<u8>>>(Self::namespaced_key(key))
|
||||
.map_err(|_| Error::Unavailable)?
|
||||
.map(Self::decode)
|
||||
.transpose()
|
||||
}
|
||||
|
||||
fn delete_cache(&self, key: &str) -> Result<(), Error> {
|
||||
self.connection()?
|
||||
.del::<_, ()>(Self::namespaced_key(key))
|
||||
.map_err(|_| Error::Unavailable)
|
||||
}
|
||||
|
||||
fn flush_cache(&self) -> Result<(), Error> {
|
||||
let mut connection = self.connection()?;
|
||||
let keys = connection
|
||||
.scan_match(Self::namespaced_pattern())
|
||||
.map_err(|_| Error::Unavailable)?
|
||||
.collect::<redis::RedisResult<Vec<String>>>()
|
||||
.map_err(|_| Error::Unavailable)?;
|
||||
if keys.is_empty() {
|
||||
return Ok(());
|
||||
}
|
||||
connection
|
||||
.del::<_, usize>(keys)
|
||||
.map(|_| ())
|
||||
.map_err(|_| Error::Unavailable)
|
||||
}
|
||||
|
||||
fn async_set_cache<'a>(
|
||||
&'a self,
|
||||
key: &'a str,
|
||||
value: Self::Value,
|
||||
kwargs: CacheKwargs,
|
||||
) -> CacheFuture<'a, ()> {
|
||||
let payload = Self::encode(&value);
|
||||
let key = Self::namespaced_key(key);
|
||||
let ttl = Self::ttl_seconds(self.get_ttl(&kwargs));
|
||||
Self::run_blocking(Arc::clone(&self.connection), move |connection| {
|
||||
connection
|
||||
.set_ex::<_, _, ()>(key, payload?, ttl)
|
||||
.map_err(|_| Error::Unavailable)
|
||||
})
|
||||
}
|
||||
|
||||
fn async_get_cache<'a>(
|
||||
&'a self,
|
||||
key: &'a str,
|
||||
_: &'a CacheKwargs,
|
||||
) -> CacheFuture<'a, Option<Self::Value>> {
|
||||
let key = Self::namespaced_key(key);
|
||||
Box::pin(async move {
|
||||
Self::run_blocking(Arc::clone(&self.connection), move |connection| {
|
||||
connection
|
||||
.get::<_, Option<Vec<u8>>>(key)
|
||||
.map_err(|_| Error::Unavailable)
|
||||
})
|
||||
.await?
|
||||
.map(Self::decode)
|
||||
.transpose()
|
||||
})
|
||||
}
|
||||
|
||||
fn async_set_cache_pipeline<'a>(
|
||||
&'a self,
|
||||
cache_list: Vec<(String, Self::Value)>,
|
||||
kwargs: CacheKwargs,
|
||||
) -> CacheFuture<'a, ()> {
|
||||
let entries = cache_list
|
||||
.into_iter()
|
||||
.map(|(key, value)| {
|
||||
Self::encode(&value).map(|payload| (Self::namespaced_key(&key), payload))
|
||||
})
|
||||
.collect::<Result<Vec<_>, _>>();
|
||||
let ttl = Self::ttl_seconds(self.get_ttl(&kwargs));
|
||||
Self::run_blocking(Arc::clone(&self.connection), move |connection| {
|
||||
for (key, payload) in entries? {
|
||||
connection
|
||||
.set_ex::<_, _, ()>(key, payload, ttl)
|
||||
.map_err(|_| Error::Unavailable)?;
|
||||
}
|
||||
Ok(())
|
||||
})
|
||||
}
|
||||
|
||||
fn async_delete_cache<'a>(&'a self, key: &'a str) -> CacheFuture<'a, ()> {
|
||||
let key = Self::namespaced_key(key);
|
||||
Self::run_blocking(Arc::clone(&self.connection), move |connection| {
|
||||
connection.del::<_, ()>(key).map_err(|_| Error::Unavailable)
|
||||
})
|
||||
}
|
||||
|
||||
fn disconnect(&self) -> CacheFuture<'_, ()> {
|
||||
Box::pin(async { Ok(()) })
|
||||
}
|
||||
|
||||
fn test_connection(&self) -> CacheFuture<'_, CacheConnectionResult> {
|
||||
Box::pin(async move {
|
||||
Self::run_blocking(Arc::clone(&self.connection), |connection| {
|
||||
redis::cmd("PING")
|
||||
.query::<String>(connection)
|
||||
.map_err(|_| Error::Unavailable)
|
||||
})
|
||||
.await?;
|
||||
Ok(CacheConnectionResult {
|
||||
status: CacheConnectionStatus::Success,
|
||||
message: "Redis cache connection test successful".into(),
|
||||
error: None,
|
||||
})
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::RedisCache;
|
||||
use litellm_cache::{BaseCache, CacheEntry, CacheKwargs};
|
||||
use redis_test::{MockCmd, MockRedisConnection};
|
||||
use serde_json::json;
|
||||
use std::time::Duration;
|
||||
|
||||
fn entry() -> CacheEntry {
|
||||
CacheEntry {
|
||||
timestamp: 123.0,
|
||||
response: json!({"choices": [{"text": "cached"}]}),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cache_entries_round_trip_through_json() {
|
||||
let entry = entry();
|
||||
let encoded = RedisCache::<redis::Connection>::encode(&entry).unwrap();
|
||||
assert_eq!(
|
||||
RedisCache::<redis::Connection>::decode(encoded).unwrap(),
|
||||
entry
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn invalid_json_is_rejected() {
|
||||
assert!(RedisCache::<redis::Connection>::decode(b"not json".to_vec()).is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn ttl_seconds_rounds_up_and_keeps_expiration_positive() {
|
||||
assert_eq!(
|
||||
RedisCache::<redis::Connection>::ttl_seconds(Duration::ZERO),
|
||||
1
|
||||
);
|
||||
assert_eq!(
|
||||
RedisCache::<redis::Connection>::ttl_seconds(Duration::from_millis(1500)),
|
||||
2
|
||||
);
|
||||
assert_eq!(
|
||||
RedisCache::<redis::Connection>::ttl_seconds(Duration::from_secs(15)),
|
||||
15
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn redis_commands_round_trip_entries_and_delete_only_namespaced_keys() {
|
||||
let value = entry();
|
||||
let payload = RedisCache::<redis::Connection>::encode(&value).unwrap();
|
||||
let connection = MockRedisConnection::new([
|
||||
MockCmd::new(
|
||||
redis::cmd("SETEX")
|
||||
.arg("litellm-cache:key")
|
||||
.arg(600)
|
||||
.arg(payload.clone()),
|
||||
Ok("OK"),
|
||||
),
|
||||
MockCmd::new(redis::cmd("GET").arg("litellm-cache:key"), Ok(payload)),
|
||||
MockCmd::new(redis::cmd("DEL").arg("litellm-cache:key"), Ok(1u32)),
|
||||
])
|
||||
.assert_all_commands_consumed();
|
||||
let cache = RedisCache::with_connection(connection, None);
|
||||
|
||||
cache
|
||||
.set_cache("key", value.clone(), CacheKwargs::default())
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
cache.get_cache("key", &CacheKwargs::default()).unwrap(),
|
||||
Some(value)
|
||||
);
|
||||
cache.delete_cache("key").unwrap();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn flush_scans_and_deletes_only_cache_keys() {
|
||||
let connection = MockRedisConnection::new([
|
||||
MockCmd::new(
|
||||
redis::cmd("SCAN")
|
||||
.cursor_arg(0)
|
||||
.arg("MATCH")
|
||||
.arg("litellm-cache:*"),
|
||||
Ok(redis_test::redis_value!(["0", ["litellm-cache:key"]])),
|
||||
),
|
||||
MockCmd::new(redis::cmd("DEL").arg("litellm-cache:key"), Ok(1u32)),
|
||||
])
|
||||
.assert_all_commands_consumed();
|
||||
let cache = RedisCache::with_connection(connection, None);
|
||||
|
||||
cache.flush_cache().unwrap();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_connection_runs_ping_off_executor() {
|
||||
let connection = MockRedisConnection::new([MockCmd::new(redis::cmd("PING"), Ok("PONG"))])
|
||||
.assert_all_commands_consumed();
|
||||
let cache = RedisCache::with_connection(connection, None);
|
||||
|
||||
assert_eq!(
|
||||
cache.test_connection().await.unwrap().status,
|
||||
litellm_cache::CacheConnectionStatus::Success
|
||||
);
|
||||
}
|
||||
}
|
||||
3
litellm-rust/crates/cache-redis/src/lib.rs
Normal file
3
litellm-rust/crates/cache-redis/src/lib.rs
Normal file
|
|
@ -0,0 +1,3 @@
|
|||
mod cache;
|
||||
|
||||
pub use cache::RedisCache;
|
||||
6
litellm-rust/crates/cache-redis/tests/cache.rs
Normal file
6
litellm-rust/crates/cache-redis/tests/cache.rs
Normal file
|
|
@ -0,0 +1,6 @@
|
|||
use litellm_cache_redis::RedisCache;
|
||||
|
||||
#[test]
|
||||
fn constructor_rejects_invalid_urls() {
|
||||
assert!(RedisCache::new("not a redis url", None).is_err());
|
||||
}
|
||||
|
|
@ -15,6 +15,7 @@ litellm-auth.workspace = true
|
|||
litellm-auth-aws.workspace = true
|
||||
litellm-auth-azure.workspace = true
|
||||
litellm-auth-gcp.workspace = true
|
||||
litellm-framing.workspace = true
|
||||
moka.workspace = true
|
||||
mime_guess = "2.0.5"
|
||||
rand.workspace = true
|
||||
|
|
@ -29,9 +30,12 @@ subtle.workspace = true
|
|||
tokio = { workspace = true, features = ["sync"] }
|
||||
tokio-tungstenite.workspace = true
|
||||
thiserror.workspace = true
|
||||
time.workspace = true
|
||||
sha2.workspace = true
|
||||
url.workspace = true
|
||||
veil.workspace = true
|
||||
|
||||
[dev-dependencies]
|
||||
aws-smithy-eventstream = "=0.61.1"
|
||||
aws-smithy-types = "1.6.1"
|
||||
rstest.workspace = true
|
||||
|
|
|
|||
|
|
@ -14,6 +14,7 @@ pub mod conversation;
|
|||
pub(crate) mod handler;
|
||||
mod prepare;
|
||||
pub mod response_utils;
|
||||
pub mod streaming;
|
||||
pub mod transformation;
|
||||
pub mod types;
|
||||
|
||||
|
|
|
|||
|
|
@ -0,0 +1,9 @@
|
|||
pub trait StreamTransformer {
|
||||
type Input;
|
||||
type Output;
|
||||
type Error;
|
||||
|
||||
fn transform(&mut self, input: Self::Input) -> Result<Vec<Self::Output>, Self::Error>;
|
||||
|
||||
fn finish(&mut self) -> Result<Vec<Self::Output>, Self::Error>;
|
||||
}
|
||||
|
|
@ -120,3 +120,83 @@ pub struct ChatCompletionsResponse {
|
|||
pub choices: Vec<ChatCompletionsChoice>,
|
||||
pub usage: ChatCompletionsUsage,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
|
||||
pub struct ChatCompletionToolCallFunctionChunk {
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub name: Option<String>,
|
||||
pub arguments: String,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub provider_specific_fields: Option<Map<String, Value>>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
|
||||
pub struct ChatCompletionToolCallChunk {
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub id: Option<String>,
|
||||
#[serde(rename = "type")]
|
||||
pub tool_type: String,
|
||||
pub function: ChatCompletionToolCallFunctionChunk,
|
||||
pub index: i64,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
|
||||
#[serde(tag = "type", rename_all = "snake_case")]
|
||||
pub enum ChatCompletionThinkingBlock {
|
||||
Thinking {
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
thinking: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
signature: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
cache_control: Option<Value>,
|
||||
},
|
||||
RedactedThinking {
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
data: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
cache_control: Option<Value>,
|
||||
},
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
|
||||
pub struct ChatCompletionDelta {
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub content: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub role: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub tool_calls: Option<Vec<ChatCompletionToolCallChunk>>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub reasoning_content: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub thinking_blocks: Option<Vec<ChatCompletionThinkingBlock>>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub provider_specific_fields: Option<Map<String, Value>>,
|
||||
#[serde(flatten)]
|
||||
pub extra: Map<String, Value>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
|
||||
pub struct ChatCompletionStreamingChoice {
|
||||
pub index: u64,
|
||||
pub delta: ChatCompletionDelta,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub finish_reason: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub logprobs: Option<Value>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
|
||||
pub struct ChatCompletionChunk {
|
||||
pub id: String,
|
||||
pub created: u64,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub model: Option<String>,
|
||||
pub object: String,
|
||||
pub choices: Vec<ChatCompletionStreamingChoice>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub usage: Option<ChatCompletionsUsage>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub provider_specific_fields: Option<Map<String, Value>>,
|
||||
}
|
||||
|
|
|
|||
|
|
@ -2,16 +2,54 @@
|
|||
pub enum Error {
|
||||
#[error("invalid provider: {0}")]
|
||||
InvalidProvider(String),
|
||||
#[error("missing required field: {0}")]
|
||||
MissingField(&'static str),
|
||||
#[error("invalid request: {0}")]
|
||||
InvalidRequest(String),
|
||||
#[error("invalid response: {0}")]
|
||||
InvalidResponse(String),
|
||||
#[error("routing error: {0}")]
|
||||
Routing(String),
|
||||
#[error("unsupported by the Rust messages route: {0}")]
|
||||
Unsupported(&'static str),
|
||||
#[error(transparent)]
|
||||
Auth(#[from] litellm_auth::Error),
|
||||
#[error(transparent)]
|
||||
Transport(#[from] crate::transport::Error),
|
||||
#[error(transparent)]
|
||||
Headers(#[from] crate::http_utils::HeaderError),
|
||||
#[error("stream framing failed: {0}")]
|
||||
StreamFraming(String),
|
||||
#[error("Anthropic SSE frame has no data")]
|
||||
MissingStreamData,
|
||||
#[error("Anthropic stream event is invalid: {0}")]
|
||||
InvalidStreamEvent(String),
|
||||
#[error("Bedrock event payload is invalid: {0}")]
|
||||
InvalidBedrockPayload(String),
|
||||
#[error("Bedrock event payload has invalid base64: {0}")]
|
||||
InvalidBedrockBase64(String),
|
||||
}
|
||||
|
||||
impl Error {
|
||||
pub fn is_request(&self) -> bool {
|
||||
match self {
|
||||
Self::InvalidProvider(_)
|
||||
| Self::MissingField(_)
|
||||
| Self::InvalidRequest(_)
|
||||
| Self::Unsupported(_)
|
||||
| Self::Headers(_) => true,
|
||||
Self::Auth(error) => !matches!(error, litellm_auth::Error::MissingApiKey { .. }),
|
||||
_ => false,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn is_response(&self) -> bool {
|
||||
matches!(
|
||||
self,
|
||||
Self::InvalidResponse(_)
|
||||
| Self::StreamFraming(_)
|
||||
| Self::MissingStreamData
|
||||
| Self::InvalidStreamEvent(_)
|
||||
| Self::InvalidBedrockPayload(_)
|
||||
| Self::InvalidBedrockBase64(_)
|
||||
)
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -46,9 +46,7 @@ pub(super) async fn execute_messages_provider_stream(
|
|||
) -> Result<reqwest::Response, Error> {
|
||||
let request = prepare_provider_request(request)?;
|
||||
if request.provider != ANTHROPIC_MESSAGES_PROVIDER {
|
||||
return Err(Error::InvalidRequest(
|
||||
"streaming messages is not supported for this provider".to_string(),
|
||||
));
|
||||
return Err(Error::Unsupported("streaming messages for this provider"));
|
||||
}
|
||||
|
||||
let mut request_builder = http_client().post(&request.url).json(&request.body);
|
||||
|
|
|
|||
|
|
@ -1 +1,2 @@
|
|||
pub mod streaming;
|
||||
pub mod transformation;
|
||||
|
|
|
|||
|
|
@ -0,0 +1,164 @@
|
|||
use std::collections::HashMap;
|
||||
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::chat_completions::Error;
|
||||
use crate::chat_completions::streaming::StreamTransformer;
|
||||
use crate::chat_completions::types::{
|
||||
ChatCompletionChunk, ChatCompletionThinkingBlock, ChatCompletionToolCallChunk,
|
||||
ChatCompletionsUsage,
|
||||
};
|
||||
use crate::providers::anthropic::messages::streaming::{
|
||||
AnthropicContentBlock, AnthropicContentBlockDelta, AnthropicMessagesStreamEvent,
|
||||
AnthropicStreamUsage,
|
||||
};
|
||||
|
||||
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
|
||||
pub enum AnthropicJsonChunkType {
|
||||
ValidJson,
|
||||
AccumulatedJson,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Eq, PartialEq)]
|
||||
pub enum AnthropicContentBlockType {
|
||||
Text,
|
||||
ToolUse,
|
||||
ServerToolUse,
|
||||
Thinking,
|
||||
RedactedThinking,
|
||||
Compaction,
|
||||
ToolResult(String),
|
||||
Other(String),
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq)]
|
||||
pub struct AnthropicContentBlockDeltaEvent {
|
||||
pub index: u64,
|
||||
pub delta: AnthropicContentBlockDelta,
|
||||
}
|
||||
|
||||
pub struct AnthropicChatCompletionsStreamTransformer {
|
||||
pub content_blocks: Vec<AnthropicContentBlockDeltaEvent>,
|
||||
pub tool_index: i64,
|
||||
pub json_mode: bool,
|
||||
pub speed: Option<String>,
|
||||
pub tool_name_reverse_map: HashMap<String, String>,
|
||||
pub response_id: String,
|
||||
pub served_model: Option<String>,
|
||||
pub is_response_format_tool: bool,
|
||||
pub converted_response_format_tool: bool,
|
||||
pub accumulated_json: String,
|
||||
pub chunk_type: AnthropicJsonChunkType,
|
||||
pub current_content_block_type: Option<AnthropicContentBlockType>,
|
||||
pub web_search_results: Vec<Value>,
|
||||
pub web_search_calls: HashMap<String, Value>,
|
||||
pub compaction_blocks: Vec<Value>,
|
||||
pub reasoning_content_chunks: Vec<String>,
|
||||
pub server_tool_inputs: HashMap<String, Value>,
|
||||
pub tool_results: Vec<Value>,
|
||||
pub current_server_tool_id: Option<String>,
|
||||
pub container_id: Option<String>,
|
||||
}
|
||||
|
||||
impl AnthropicChatCompletionsStreamTransformer {
|
||||
pub fn new(
|
||||
_json_mode: bool,
|
||||
_speed: Option<String>,
|
||||
_tool_name_reverse_map: HashMap<String, String>,
|
||||
) -> Self {
|
||||
todo!()
|
||||
}
|
||||
|
||||
pub fn check_empty_tool_call_args(&self) -> bool {
|
||||
todo!()
|
||||
}
|
||||
|
||||
pub fn handle_usage(&mut self, _usage: AnthropicStreamUsage) -> ChatCompletionsUsage {
|
||||
todo!()
|
||||
}
|
||||
|
||||
pub fn handle_content_block_delta(
|
||||
&mut self,
|
||||
_index: u64,
|
||||
_delta: AnthropicContentBlockDelta,
|
||||
) -> (
|
||||
String,
|
||||
Option<ChatCompletionToolCallChunk>,
|
||||
Vec<ChatCompletionThinkingBlock>,
|
||||
Option<Value>,
|
||||
Option<String>,
|
||||
) {
|
||||
todo!()
|
||||
}
|
||||
|
||||
pub fn handle_content_block_start(
|
||||
&mut self,
|
||||
_index: u64,
|
||||
_content_block: AnthropicContentBlock,
|
||||
) -> Result<ChatCompletionChunk, Error> {
|
||||
todo!()
|
||||
}
|
||||
|
||||
pub fn handle_json_mode_chunk(
|
||||
&mut self,
|
||||
_text: String,
|
||||
_tool_use: Option<ChatCompletionToolCallChunk>,
|
||||
) -> (String, Option<ChatCompletionToolCallChunk>) {
|
||||
todo!()
|
||||
}
|
||||
|
||||
pub fn handle_accumulated_json_chunk(
|
||||
&mut self,
|
||||
_data: &str,
|
||||
_is_final: bool,
|
||||
) -> Result<Option<ChatCompletionChunk>, Error> {
|
||||
todo!()
|
||||
}
|
||||
|
||||
pub fn handle_redacted_thinking_content(
|
||||
&mut self,
|
||||
_content_block: &AnthropicContentBlock,
|
||||
) -> Vec<ChatCompletionThinkingBlock> {
|
||||
todo!()
|
||||
}
|
||||
|
||||
pub fn web_search_call_snapshot(&self) -> HashMap<String, Value> {
|
||||
todo!()
|
||||
}
|
||||
|
||||
pub fn complete_web_search_call(&mut self, _result: Value) {
|
||||
todo!()
|
||||
}
|
||||
|
||||
pub fn build_code_interpreter_results(&self) -> Vec<Value> {
|
||||
todo!()
|
||||
}
|
||||
|
||||
pub fn handle_message_delta(
|
||||
&mut self,
|
||||
_event: AnthropicMessagesStreamEvent,
|
||||
) -> (Option<String>, Option<ChatCompletionsUsage>, Option<Value>) {
|
||||
todo!()
|
||||
}
|
||||
|
||||
pub fn chunk_parser(
|
||||
&mut self,
|
||||
_event: AnthropicMessagesStreamEvent,
|
||||
) -> Result<ChatCompletionChunk, Error> {
|
||||
todo!()
|
||||
}
|
||||
}
|
||||
|
||||
impl StreamTransformer for AnthropicChatCompletionsStreamTransformer {
|
||||
type Input = AnthropicMessagesStreamEvent;
|
||||
type Output = ChatCompletionChunk;
|
||||
type Error = Error;
|
||||
|
||||
fn transform(&mut self, _input: Self::Input) -> Result<Vec<Self::Output>, Self::Error> {
|
||||
todo!()
|
||||
}
|
||||
|
||||
fn finish(&mut self) -> Result<Vec<Self::Output>, Self::Error> {
|
||||
todo!()
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,338 @@
|
|||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::Value;
|
||||
use time::OffsetDateTime;
|
||||
use url::Url;
|
||||
|
||||
use crate::messages::Error;
|
||||
use crate::messages::types::AnthropicMessagesResponse;
|
||||
use crate::providers::anthropic::messages::transformation::resolve_anthropic_api_base;
|
||||
|
||||
const BATCHES_PATH_SUFFIX: &str = "/v1/messages/batches";
|
||||
|
||||
#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct AnthropicBatchRequestCounts {
|
||||
#[serde(default)]
|
||||
pub processing: u64,
|
||||
#[serde(default)]
|
||||
pub succeeded: u64,
|
||||
#[serde(default)]
|
||||
pub errored: u64,
|
||||
#[serde(default)]
|
||||
pub canceled: u64,
|
||||
#[serde(default)]
|
||||
pub expired: u64,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct AnthropicMessageBatch {
|
||||
#[serde(default)]
|
||||
pub id: String,
|
||||
#[serde(default = "default_processing_status")]
|
||||
pub processing_status: String,
|
||||
pub created_at: Option<String>,
|
||||
pub ended_at: Option<String>,
|
||||
pub expires_at: Option<String>,
|
||||
pub cancel_initiated_at: Option<String>,
|
||||
pub archived_at: Option<String>,
|
||||
#[serde(default)]
|
||||
pub request_counts: AnthropicBatchRequestCounts,
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum BatchStatus {
|
||||
InProgress,
|
||||
Cancelling,
|
||||
Completed,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct BatchRequestCounts {
|
||||
pub total: u64,
|
||||
pub completed: u64,
|
||||
pub failed: u64,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct LiteLlmMessageBatch {
|
||||
pub id: String,
|
||||
pub object: String,
|
||||
pub endpoint: String,
|
||||
pub input_file_id: String,
|
||||
pub completion_window: String,
|
||||
pub status: BatchStatus,
|
||||
pub output_file_id: String,
|
||||
pub created_at: i64,
|
||||
pub in_progress_at: Option<i64>,
|
||||
pub expires_at: Option<i64>,
|
||||
pub completed_at: Option<i64>,
|
||||
pub expired_at: Option<i64>,
|
||||
pub cancelling_at: Option<i64>,
|
||||
pub cancelled_at: Option<i64>,
|
||||
pub request_counts: BatchRequestCounts,
|
||||
}
|
||||
|
||||
pub trait AnthropicBatchesConfig {
|
||||
fn create_batch_url(
|
||||
&self,
|
||||
api_base: Option<&str>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> Result<String, Error>;
|
||||
|
||||
fn transform_create_batch_request(&self) -> Result<Value, Error>;
|
||||
|
||||
fn transform_create_batch_response(
|
||||
&self,
|
||||
response: AnthropicMessageBatch,
|
||||
now: i64,
|
||||
) -> Result<LiteLlmMessageBatch, Error>;
|
||||
|
||||
fn retrieve_batch_url(
|
||||
&self,
|
||||
api_base: Option<&str>,
|
||||
batch_id: &str,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> Result<String, Error>;
|
||||
|
||||
fn transform_retrieve_batch_request(&self) -> Value;
|
||||
|
||||
fn transform_retrieve_batch_response(
|
||||
&self,
|
||||
response: AnthropicMessageBatch,
|
||||
now: i64,
|
||||
) -> LiteLlmMessageBatch;
|
||||
|
||||
fn transform_batch_results(&self, body: &str) -> Result<Vec<AnthropicMessagesResponse>, Error>;
|
||||
}
|
||||
|
||||
pub struct AnthropicBatchesTransformation;
|
||||
|
||||
pub const ANTHROPIC_BATCHES_TRANSFORMATION: AnthropicBatchesTransformation =
|
||||
AnthropicBatchesTransformation;
|
||||
|
||||
fn default_processing_status() -> String {
|
||||
"in_progress".into()
|
||||
}
|
||||
|
||||
fn timestamp(value: Option<&str>) -> Option<i64> {
|
||||
value
|
||||
.and_then(|value| {
|
||||
OffsetDateTime::parse(value, &time::format_description::well_known::Rfc3339).ok()
|
||||
})
|
||||
.map(OffsetDateTime::unix_timestamp)
|
||||
}
|
||||
|
||||
fn batches_base_url(
|
||||
api_base: Option<&str>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> Result<Url, Error> {
|
||||
let api_base = resolve_anthropic_api_base(api_base, env_lookup);
|
||||
let api_base = api_base.trim_end_matches('/');
|
||||
let complete_url = if api_base.ends_with(BATCHES_PATH_SUFFIX) {
|
||||
api_base.to_string()
|
||||
} else if let Some(base) = api_base.strip_suffix("/v1/messages") {
|
||||
format!("{base}{BATCHES_PATH_SUFFIX}")
|
||||
} else {
|
||||
format!("{api_base}{BATCHES_PATH_SUFFIX}")
|
||||
};
|
||||
Url::parse(&complete_url)
|
||||
.map_err(|error| Error::InvalidRequest(format!("invalid Anthropic API base: {error}")))
|
||||
}
|
||||
|
||||
impl AnthropicBatchesConfig for AnthropicBatchesTransformation {
|
||||
fn create_batch_url(
|
||||
&self,
|
||||
api_base: Option<&str>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> Result<String, Error> {
|
||||
Ok(batches_base_url(api_base, env_lookup)?.into())
|
||||
}
|
||||
|
||||
fn transform_create_batch_request(&self) -> Result<Value, Error> {
|
||||
Err(Error::Unsupported("Anthropic message batch creation"))
|
||||
}
|
||||
|
||||
fn transform_create_batch_response(
|
||||
&self,
|
||||
_response: AnthropicMessageBatch,
|
||||
_now: i64,
|
||||
) -> Result<LiteLlmMessageBatch, Error> {
|
||||
Err(Error::Unsupported("Anthropic message batch creation"))
|
||||
}
|
||||
|
||||
fn retrieve_batch_url(
|
||||
&self,
|
||||
api_base: Option<&str>,
|
||||
batch_id: &str,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> Result<String, Error> {
|
||||
if batch_id.is_empty() {
|
||||
return Err(Error::MissingField("batch_id"));
|
||||
}
|
||||
let mut url = batches_base_url(api_base, env_lookup)?;
|
||||
url.path_segments_mut()
|
||||
.map_err(|_| Error::InvalidRequest("Anthropic API base cannot be a base URL".into()))?
|
||||
.push(batch_id);
|
||||
Ok(url.into())
|
||||
}
|
||||
|
||||
fn transform_retrieve_batch_request(&self) -> Value {
|
||||
Value::Object(Default::default())
|
||||
}
|
||||
|
||||
fn transform_retrieve_batch_response(
|
||||
&self,
|
||||
response: AnthropicMessageBatch,
|
||||
now: i64,
|
||||
) -> LiteLlmMessageBatch {
|
||||
let created_at = timestamp(response.created_at.as_deref());
|
||||
let ended_at = timestamp(response.ended_at.as_deref());
|
||||
let expires_at = timestamp(response.expires_at.as_deref());
|
||||
let cancel_initiated_at = timestamp(response.cancel_initiated_at.as_deref());
|
||||
let archived_at = timestamp(response.archived_at.as_deref());
|
||||
let status = match response.processing_status.as_str() {
|
||||
"canceling" => BatchStatus::Cancelling,
|
||||
"ended" => BatchStatus::Completed,
|
||||
_ => BatchStatus::InProgress,
|
||||
};
|
||||
let request_counts = BatchRequestCounts {
|
||||
total: response.request_counts.processing
|
||||
+ response.request_counts.succeeded
|
||||
+ response.request_counts.errored
|
||||
+ response.request_counts.canceled
|
||||
+ response.request_counts.expired,
|
||||
completed: response.request_counts.succeeded,
|
||||
failed: response.request_counts.errored,
|
||||
};
|
||||
|
||||
LiteLlmMessageBatch {
|
||||
id: response.id.clone(),
|
||||
object: "batch".into(),
|
||||
endpoint: "/v1/messages".into(),
|
||||
input_file_id: "None".into(),
|
||||
completion_window: "24h".into(),
|
||||
status,
|
||||
output_file_id: response.id,
|
||||
created_at: created_at.unwrap_or(now),
|
||||
in_progress_at: (response.processing_status == "in_progress")
|
||||
.then_some(created_at)
|
||||
.flatten(),
|
||||
expires_at,
|
||||
completed_at: (response.processing_status == "ended")
|
||||
.then_some(ended_at)
|
||||
.flatten(),
|
||||
expired_at: archived_at,
|
||||
cancelling_at: (response.processing_status == "canceling")
|
||||
.then_some(cancel_initiated_at)
|
||||
.flatten(),
|
||||
cancelled_at: (response.processing_status == "canceling")
|
||||
.then_some(ended_at)
|
||||
.flatten(),
|
||||
request_counts,
|
||||
}
|
||||
}
|
||||
|
||||
fn transform_batch_results(&self, body: &str) -> Result<Vec<AnthropicMessagesResponse>, Error> {
|
||||
body.lines()
|
||||
.filter(|line| !line.trim().is_empty())
|
||||
.filter_map(|line| serde_json::from_str::<Value>(line.trim()).ok())
|
||||
.map(|record| {
|
||||
serde_json::from_value(record["result"]["message"].clone()).map_err(|error| {
|
||||
Error::InvalidResponse(format!("invalid Anthropic batch result: {error}"))
|
||||
})
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use serde_json::json;
|
||||
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn builds_and_encodes_message_batch_urls() {
|
||||
assert_eq!(
|
||||
ANTHROPIC_BATCHES_TRANSFORMATION
|
||||
.create_batch_url(None, &|_| None)
|
||||
.unwrap(),
|
||||
"https://api.anthropic.com/v1/messages/batches"
|
||||
);
|
||||
assert_eq!(
|
||||
ANTHROPIC_BATCHES_TRANSFORMATION
|
||||
.create_batch_url(Some("https://proxy.test/v1/messages/batches"), &|_| None)
|
||||
.unwrap(),
|
||||
"https://proxy.test/v1/messages/batches"
|
||||
);
|
||||
assert_eq!(
|
||||
ANTHROPIC_BATCHES_TRANSFORMATION
|
||||
.retrieve_batch_url(Some("https://proxy.test"), "batch/id ?", &|_| None)
|
||||
.unwrap(),
|
||||
"https://proxy.test/v1/messages/batches/batch%2Fid%20%3F"
|
||||
);
|
||||
assert_eq!(
|
||||
ANTHROPIC_BATCHES_TRANSFORMATION.transform_retrieve_batch_request(),
|
||||
json!({})
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn maps_retrieved_batch_status_counts_and_timestamps_like_python() {
|
||||
let response: AnthropicMessageBatch = serde_json::from_value(json!({
|
||||
"id": "msgbatch_1",
|
||||
"processing_status": "ended",
|
||||
"created_at": "2025-01-01T00:00:00Z",
|
||||
"ended_at": "2025-01-01T00:01:00Z",
|
||||
"expires_at": "not-a-timestamp",
|
||||
"request_counts": {
|
||||
"processing": 1,
|
||||
"succeeded": 2,
|
||||
"errored": 3,
|
||||
"canceled": 4,
|
||||
"expired": 5
|
||||
}
|
||||
}))
|
||||
.unwrap();
|
||||
|
||||
let batch = ANTHROPIC_BATCHES_TRANSFORMATION.transform_retrieve_batch_response(response, 7);
|
||||
assert_eq!(batch.status, BatchStatus::Completed);
|
||||
assert_eq!(batch.created_at, 1_735_689_600);
|
||||
assert_eq!(batch.completed_at, Some(1_735_689_660));
|
||||
assert_eq!(batch.expires_at, None);
|
||||
assert_eq!(
|
||||
batch.request_counts,
|
||||
BatchRequestCounts {
|
||||
total: 15,
|
||||
completed: 2,
|
||||
failed: 3
|
||||
}
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn extracts_message_responses_from_ndjson_and_skips_non_json_lines() {
|
||||
let body = r#"not-json
|
||||
{"result":{"message":{"id":"msg_1","type":"message","role":"assistant","model":"claude-test","content":[],"stop_reason":"end_turn","stop_sequence":null}}}
|
||||
"#;
|
||||
let messages = ANTHROPIC_BATCHES_TRANSFORMATION
|
||||
.transform_batch_results(body)
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(messages.len(), 1);
|
||||
assert_eq!(messages[0].id, "msg_1");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn preserves_python_placeholder_for_batch_creation() {
|
||||
assert!(matches!(
|
||||
ANTHROPIC_BATCHES_TRANSFORMATION.transform_create_batch_request(),
|
||||
Err(Error::Unsupported("Anthropic message batch creation"))
|
||||
));
|
||||
let response: AnthropicMessageBatch = serde_json::from_value(json!({})).unwrap();
|
||||
assert!(matches!(
|
||||
ANTHROPIC_BATCHES_TRANSFORMATION.transform_create_batch_response(response, 0),
|
||||
Err(Error::Unsupported("Anthropic message batch creation"))
|
||||
));
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,168 @@
|
|||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::constants::ANTHROPIC_OAUTH_TOKEN_PREFIX;
|
||||
use crate::messages::Error;
|
||||
use crate::messages::types::{AnthropicMessage, SystemPrompt};
|
||||
|
||||
const COUNT_TOKENS_ENDPOINT: &str = "https://api.anthropic.com/v1/messages/count_tokens";
|
||||
const TOKEN_COUNTING_BETA: &str = "token-counting-2024-11-01";
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
|
||||
pub struct AnthropicCountTokensRequest {
|
||||
pub model: String,
|
||||
pub messages: Vec<AnthropicMessage>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub tools: Option<Vec<Value>>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub system: Option<SystemPrompt>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct AnthropicCountTokensResponse {
|
||||
pub input_tokens: u64,
|
||||
}
|
||||
|
||||
pub trait AnthropicCountTokensConfig {
|
||||
fn endpoint(&self) -> &'static str;
|
||||
|
||||
fn validate_request(&self, model: &str, messages: &[AnthropicMessage]) -> Result<(), Error>;
|
||||
|
||||
fn transform_request(
|
||||
&self,
|
||||
model: &str,
|
||||
messages: Vec<AnthropicMessage>,
|
||||
tools: Option<Vec<Value>>,
|
||||
system: Option<SystemPrompt>,
|
||||
) -> Result<AnthropicCountTokensRequest, Error>;
|
||||
|
||||
fn required_headers(&self, api_key: &str) -> Vec<(&'static str, String)>;
|
||||
}
|
||||
|
||||
pub struct AnthropicCountTokensTransformation;
|
||||
|
||||
pub const ANTHROPIC_COUNT_TOKENS_TRANSFORMATION: AnthropicCountTokensTransformation =
|
||||
AnthropicCountTokensTransformation;
|
||||
|
||||
impl AnthropicCountTokensConfig for AnthropicCountTokensTransformation {
|
||||
fn endpoint(&self) -> &'static str {
|
||||
COUNT_TOKENS_ENDPOINT
|
||||
}
|
||||
|
||||
fn transform_request(
|
||||
&self,
|
||||
model: &str,
|
||||
messages: Vec<AnthropicMessage>,
|
||||
tools: Option<Vec<Value>>,
|
||||
system: Option<SystemPrompt>,
|
||||
) -> Result<AnthropicCountTokensRequest, Error> {
|
||||
self.validate_request(model, &messages)?;
|
||||
|
||||
Ok(AnthropicCountTokensRequest {
|
||||
model: model.to_string(),
|
||||
messages,
|
||||
tools,
|
||||
system,
|
||||
})
|
||||
}
|
||||
|
||||
fn validate_request(&self, model: &str, messages: &[AnthropicMessage]) -> Result<(), Error> {
|
||||
if model.is_empty() {
|
||||
return Err(Error::MissingField("model"));
|
||||
}
|
||||
if messages.is_empty() {
|
||||
return Err(Error::MissingField("messages"));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn required_headers(&self, api_key: &str) -> Vec<(&'static str, String)> {
|
||||
let auth = if api_key.starts_with(ANTHROPIC_OAUTH_TOKEN_PREFIX) {
|
||||
("authorization", format!("Bearer {api_key}"))
|
||||
} else {
|
||||
("x-api-key", api_key.to_string())
|
||||
};
|
||||
vec![
|
||||
("content-type", "application/json".to_string()),
|
||||
auth,
|
||||
("anthropic-version", "2023-06-01".to_string()),
|
||||
("anthropic-beta", TOKEN_COUNTING_BETA.to_string()),
|
||||
]
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use serde_json::{Map, json};
|
||||
|
||||
use super::*;
|
||||
use crate::messages::types::MessageContent;
|
||||
|
||||
fn message() -> AnthropicMessage {
|
||||
AnthropicMessage {
|
||||
role: "user".into(),
|
||||
content: MessageContent::Text("hello".into()),
|
||||
extra: Map::new(),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn maps_the_python_count_tokens_contract() {
|
||||
let request = ANTHROPIC_COUNT_TOKENS_TRANSFORMATION
|
||||
.transform_request(
|
||||
"claude-test",
|
||||
vec![message()],
|
||||
Some(vec![json!({"name": "lookup"})]),
|
||||
Some(SystemPrompt::Text("system".into())),
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(
|
||||
serde_json::to_value(request).unwrap(),
|
||||
json!({
|
||||
"model": "claude-test",
|
||||
"messages": [{"role": "user", "content": "hello"}],
|
||||
"tools": [{"name": "lookup"}],
|
||||
"system": "system"
|
||||
})
|
||||
);
|
||||
assert_eq!(
|
||||
ANTHROPIC_COUNT_TOKENS_TRANSFORMATION.endpoint(),
|
||||
COUNT_TOKENS_ENDPOINT
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_the_invalid_requests_python_rejects() {
|
||||
assert!(matches!(
|
||||
ANTHROPIC_COUNT_TOKENS_TRANSFORMATION.transform_request(
|
||||
"",
|
||||
vec![message()],
|
||||
None,
|
||||
None
|
||||
),
|
||||
Err(Error::MissingField("model"))
|
||||
));
|
||||
assert!(matches!(
|
||||
ANTHROPIC_COUNT_TOKENS_TRANSFORMATION.transform_request(
|
||||
"claude-test",
|
||||
vec![],
|
||||
None,
|
||||
None
|
||||
),
|
||||
Err(Error::MissingField("messages"))
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn uses_api_key_or_oauth_headers_without_combining_credentials() {
|
||||
let api_key = ANTHROPIC_COUNT_TOKENS_TRANSFORMATION.required_headers("sk-ant-api");
|
||||
assert!(api_key.contains(&("x-api-key", "sk-ant-api".into())));
|
||||
assert!(!api_key.iter().any(|(name, _)| *name == "authorization"));
|
||||
|
||||
let oauth = ANTHROPIC_COUNT_TOKENS_TRANSFORMATION.required_headers("sk-ant-oat-test");
|
||||
assert!(oauth.contains(&("authorization", "Bearer sk-ant-oat-test".into())));
|
||||
assert!(!oauth.iter().any(|(name, _)| *name == "x-api-key"));
|
||||
assert!(oauth.contains(&("anthropic-beta", TOKEN_COUNTING_BETA.into())));
|
||||
}
|
||||
}
|
||||
|
|
@ -1 +1,4 @@
|
|||
pub mod batches;
|
||||
pub mod count_tokens;
|
||||
pub mod streaming;
|
||||
pub mod transformation;
|
||||
|
|
|
|||
|
|
@ -0,0 +1,282 @@
|
|||
use base64::Engine;
|
||||
use bytes::Buf;
|
||||
use futures_util::{Stream, StreamExt};
|
||||
use litellm_framing::Framer;
|
||||
use litellm_framing::aws_event_stream::{AwsEventStreamFrame, AwsEventStreamFramer};
|
||||
use litellm_framing::sse::{SseFrame, SseFramer};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
use crate::messages::Error;
|
||||
|
||||
#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
|
||||
pub struct AnthropicStreamUsage {
|
||||
#[serde(default)]
|
||||
pub input_tokens: u64,
|
||||
#[serde(default)]
|
||||
pub output_tokens: u64,
|
||||
#[serde(default)]
|
||||
pub cache_creation_input_tokens: u64,
|
||||
#[serde(default)]
|
||||
pub cache_read_input_tokens: u64,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub server_tool_use: Option<Value>,
|
||||
#[serde(flatten)]
|
||||
pub extra: Map<String, Value>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
|
||||
pub struct AnthropicStreamMessage {
|
||||
pub id: String,
|
||||
#[serde(rename = "type")]
|
||||
pub message_type: String,
|
||||
pub role: String,
|
||||
pub model: String,
|
||||
pub content: Vec<Value>,
|
||||
pub stop_reason: Option<String>,
|
||||
pub stop_sequence: Option<String>,
|
||||
pub usage: AnthropicStreamUsage,
|
||||
#[serde(flatten)]
|
||||
pub extra: Map<String, Value>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
|
||||
#[serde(tag = "type", rename_all = "snake_case")]
|
||||
pub enum AnthropicContentBlockDelta {
|
||||
TextDelta {
|
||||
text: String,
|
||||
},
|
||||
InputJsonDelta {
|
||||
partial_json: String,
|
||||
},
|
||||
#[serde(rename = "citations_delta")]
|
||||
Citations {
|
||||
citation: Value,
|
||||
},
|
||||
ThinkingDelta {
|
||||
thinking: String,
|
||||
},
|
||||
SignatureDelta {
|
||||
signature: String,
|
||||
},
|
||||
CompactionDelta {
|
||||
content: String,
|
||||
},
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
|
||||
pub struct AnthropicContentBlock {
|
||||
#[serde(rename = "type")]
|
||||
pub block_type: String,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub id: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub name: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub text: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub input: Option<Value>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub thinking: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub signature: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub data: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub content: Option<Value>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub caller: Option<Value>,
|
||||
#[serde(flatten)]
|
||||
pub extra: Map<String, Value>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
|
||||
pub struct AnthropicMessageDelta {
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub stop_reason: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub stop_sequence: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub stop_details: Option<Value>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub container: Option<Value>,
|
||||
#[serde(flatten)]
|
||||
pub extra: Map<String, Value>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
|
||||
pub struct AnthropicStreamError {
|
||||
#[serde(rename = "type")]
|
||||
pub error_type: String,
|
||||
pub message: String,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub details: Option<Value>,
|
||||
#[serde(flatten)]
|
||||
pub extra: Map<String, Value>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
|
||||
#[serde(tag = "type", rename_all = "snake_case")]
|
||||
pub enum AnthropicMessagesStreamEvent {
|
||||
MessageStart {
|
||||
message: AnthropicStreamMessage,
|
||||
},
|
||||
ContentBlockStart {
|
||||
index: u64,
|
||||
content_block: AnthropicContentBlock,
|
||||
},
|
||||
ContentBlockDelta {
|
||||
index: u64,
|
||||
delta: AnthropicContentBlockDelta,
|
||||
},
|
||||
ContentBlockStop {
|
||||
index: u64,
|
||||
},
|
||||
MessageDelta {
|
||||
delta: AnthropicMessageDelta,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
usage: Option<AnthropicStreamUsage>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
context_management: Option<Value>,
|
||||
},
|
||||
MessageStop,
|
||||
Ping,
|
||||
Error {
|
||||
error: AnthropicStreamError,
|
||||
},
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct BedrockChunkPayload {
|
||||
bytes: String,
|
||||
}
|
||||
|
||||
pub fn decode_anthropic_sse_frame(frame: SseFrame) -> Result<AnthropicMessagesStreamEvent, Error> {
|
||||
let data = frame.data.ok_or(Error::MissingStreamData)?;
|
||||
serde_json::from_str(&data).map_err(|error| Error::InvalidStreamEvent(error.to_string()))
|
||||
}
|
||||
|
||||
pub fn decode_bedrock_anthropic_frame(
|
||||
frame: AwsEventStreamFrame,
|
||||
) -> Result<AnthropicMessagesStreamEvent, Error> {
|
||||
let payload: BedrockChunkPayload = serde_json::from_slice(&frame.payload)
|
||||
.map_err(|error| Error::InvalidBedrockPayload(error.to_string()))?;
|
||||
let event = base64::engine::general_purpose::STANDARD
|
||||
.decode(payload.bytes)
|
||||
.map_err(|error| Error::InvalidBedrockBase64(error.to_string()))?;
|
||||
serde_json::from_slice(&event).map_err(|error| Error::InvalidStreamEvent(error.to_string()))
|
||||
}
|
||||
|
||||
pub fn direct_anthropic_event_stream<S, B, E>(
|
||||
input: S,
|
||||
) -> impl Stream<Item = Result<AnthropicMessagesStreamEvent, Error>> + Send
|
||||
where
|
||||
S: Stream<Item = Result<B, E>> + Send,
|
||||
B: Buf + Send,
|
||||
E: std::error::Error + Send + Sync + 'static,
|
||||
{
|
||||
SseFramer.frame(input).map(|frame| {
|
||||
let frame = frame.map_err(|error| Error::StreamFraming(error.to_string()))?;
|
||||
decode_anthropic_sse_frame(frame)
|
||||
})
|
||||
}
|
||||
|
||||
pub fn bedrock_anthropic_event_stream<S, B, E>(
|
||||
input: S,
|
||||
) -> impl Stream<Item = Result<AnthropicMessagesStreamEvent, Error>> + Send
|
||||
where
|
||||
S: Stream<Item = Result<B, E>> + Send,
|
||||
B: Buf + Send,
|
||||
E: std::error::Error + Send + Sync + 'static,
|
||||
{
|
||||
AwsEventStreamFramer.frame(input).map(|frame| {
|
||||
let frame = frame.map_err(|error| Error::StreamFraming(error.to_string()))?;
|
||||
decode_bedrock_anthropic_frame(frame)
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::io;
|
||||
|
||||
use aws_smithy_eventstream::frame::write_message_to;
|
||||
use aws_smithy_types::event_stream::{Header, HeaderValue, Message};
|
||||
use base64::engine::general_purpose::STANDARD;
|
||||
use bytes::Bytes;
|
||||
use futures_util::TryStreamExt;
|
||||
|
||||
use super::*;
|
||||
|
||||
const TEXT_DELTA: &str =
|
||||
r#"{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"hello"}}"#;
|
||||
|
||||
#[tokio::test]
|
||||
async fn direct_anthropic_sse_frames_into_typed_events() {
|
||||
let wire = format!("event: content_block_delta\ndata: {TEXT_DELTA}\n\n");
|
||||
let events = direct_anthropic_event_stream(futures_util::stream::iter(
|
||||
wire.as_bytes().chunks(3).map(Ok::<_, io::Error>),
|
||||
))
|
||||
.try_collect::<Vec<_>>()
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(
|
||||
events,
|
||||
vec![AnthropicMessagesStreamEvent::ContentBlockDelta {
|
||||
index: 0,
|
||||
delta: AnthropicContentBlockDelta::TextDelta {
|
||||
text: "hello".into(),
|
||||
},
|
||||
}]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn decodes_citations_delta_events() {
|
||||
let event = decode_anthropic_sse_frame(SseFrame {
|
||||
event: Some("content_block_delta".into()),
|
||||
data: Some(
|
||||
r#"{"type":"content_block_delta","index":0,"delta":{"type":"citations_delta","citation":{"type":"char_location"}}}"#
|
||||
.into(),
|
||||
),
|
||||
id: None,
|
||||
retry: None,
|
||||
})
|
||||
.unwrap();
|
||||
|
||||
assert!(matches!(
|
||||
event,
|
||||
AnthropicMessagesStreamEvent::ContentBlockDelta {
|
||||
delta: AnthropicContentBlockDelta::Citations { .. },
|
||||
..
|
||||
}
|
||||
));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn bedrock_aws_frames_into_the_same_typed_events() {
|
||||
let payload = serde_json::json!({"bytes": STANDARD.encode(TEXT_DELTA)});
|
||||
let message = Message::new(Bytes::from(serde_json::to_vec(&payload).unwrap())).add_header(
|
||||
Header::new(":event-type", HeaderValue::String("chunk".into())),
|
||||
);
|
||||
let mut wire = Vec::new();
|
||||
write_message_to(&message, &mut wire).unwrap();
|
||||
|
||||
let events = bedrock_anthropic_event_stream(futures_util::stream::iter(
|
||||
wire.chunks(3).map(Ok::<_, io::Error>),
|
||||
))
|
||||
.try_collect::<Vec<_>>()
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(
|
||||
events,
|
||||
vec![AnthropicMessagesStreamEvent::ContentBlockDelta {
|
||||
index: 0,
|
||||
delta: AnthropicContentBlockDelta::TextDelta {
|
||||
text: "hello".into(),
|
||||
},
|
||||
}]
|
||||
);
|
||||
}
|
||||
}
|
||||
|
|
@ -31,10 +31,7 @@ pub fn complete_anthropic_url(
|
|||
api_base: Option<&str>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> String {
|
||||
let api_base = non_empty(api_base)
|
||||
.map(str::to_string)
|
||||
.or_else(|| env_lookup(ANTHROPIC_API_BASE_ENV).filter(|value| !value.trim().is_empty()))
|
||||
.unwrap_or_else(|| DEFAULT_ANTHROPIC_API_BASE.to_string());
|
||||
let api_base = resolve_anthropic_api_base(api_base, env_lookup);
|
||||
|
||||
let api_base = api_base.trim_end_matches('/');
|
||||
if api_base.ends_with(MESSAGES_PATH_SUFFIX) {
|
||||
|
|
@ -43,6 +40,16 @@ pub fn complete_anthropic_url(
|
|||
format!("{api_base}{MESSAGES_PATH_SUFFIX}")
|
||||
}
|
||||
|
||||
pub fn resolve_anthropic_api_base(
|
||||
api_base: Option<&str>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> String {
|
||||
non_empty(api_base)
|
||||
.map(str::to_string)
|
||||
.or_else(|| env_lookup(ANTHROPIC_API_BASE_ENV).filter(|value| !value.trim().is_empty()))
|
||||
.unwrap_or_else(|| DEFAULT_ANTHROPIC_API_BASE.to_string())
|
||||
}
|
||||
|
||||
impl AnthropicMessagesProviderConfig for AnthropicMessagesConfig {
|
||||
fn complete_url(
|
||||
&self,
|
||||
|
|
|
|||
23
litellm-rust/crates/framer/Cargo.toml
Normal file
23
litellm-rust/crates/framer/Cargo.toml
Normal file
|
|
@ -0,0 +1,23 @@
|
|||
[package]
|
||||
name = "litellm-framing"
|
||||
version = "0.1.0"
|
||||
edition.workspace = true
|
||||
license.workspace = true
|
||||
repository.workspace = true
|
||||
|
||||
[features]
|
||||
default = ["aws", "sse"]
|
||||
aws = ["dep:aws-smithy-eventstream", "dep:aws-smithy-types"]
|
||||
sse = ["dep:sse-stream"]
|
||||
|
||||
[dependencies]
|
||||
aws-smithy-eventstream = { version = "=0.61.1", optional = true }
|
||||
aws-smithy-types = { version = "1.6.1", optional = true }
|
||||
bytes = "1"
|
||||
futures-util.workspace = true
|
||||
sse-stream = { version = "=0.2.6", optional = true }
|
||||
thiserror.workspace = true
|
||||
|
||||
[dev-dependencies]
|
||||
rstest.workspace = true
|
||||
tokio.workspace = true
|
||||
66
litellm-rust/crates/framer/src/aws_event_stream.rs
Normal file
66
litellm-rust/crates/framer/src/aws_event_stream.rs
Normal file
|
|
@ -0,0 +1,66 @@
|
|||
use bytes::{Buf, Bytes, BytesMut};
|
||||
use futures_util::{Stream, StreamExt};
|
||||
|
||||
use aws_smithy_eventstream::frame::read_message_from;
|
||||
use aws_smithy_types::event_stream::Header;
|
||||
|
||||
use crate::{Error, Framer};
|
||||
|
||||
const MAX_FRAME_BYTES: usize = 16 * 1024 * 1024;
|
||||
|
||||
#[derive(Clone, Debug, PartialEq)]
|
||||
pub struct AwsEventStreamFrame {
|
||||
pub headers: Vec<Header>,
|
||||
pub payload: Bytes,
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, Default)]
|
||||
pub struct AwsEventStreamFramer;
|
||||
|
||||
impl Framer for AwsEventStreamFramer {
|
||||
type Frame = AwsEventStreamFrame;
|
||||
|
||||
fn frame<S, B, E>(self, input: S) -> impl Stream<Item = Result<Self::Frame, Error>> + Send
|
||||
where
|
||||
S: Stream<Item = Result<B, E>> + Send,
|
||||
B: Buf + Send,
|
||||
E: std::error::Error + Send + Sync + 'static,
|
||||
{
|
||||
futures_util::stream::try_unfold(
|
||||
(Box::pin(input), BytesMut::new()),
|
||||
|(mut input, mut buffer)| async move {
|
||||
loop {
|
||||
if buffer.len() >= 4 {
|
||||
let length = (&buffer[..4]).get_u32() as usize;
|
||||
if !(16..=MAX_FRAME_BYTES).contains(&length) {
|
||||
return Err(Error::InvalidLength(length));
|
||||
}
|
||||
if buffer.len() >= length {
|
||||
let raw = buffer.split_to(length).freeze();
|
||||
let message = read_message_from(raw)?;
|
||||
let frame = AwsEventStreamFrame {
|
||||
headers: message.headers().to_vec(),
|
||||
payload: message.payload().clone(),
|
||||
};
|
||||
return Ok(Some((frame, (input, buffer))));
|
||||
}
|
||||
}
|
||||
match input.next().await {
|
||||
Some(Ok(mut chunk)) => {
|
||||
while chunk.has_remaining() {
|
||||
let bytes = chunk.chunk();
|
||||
buffer.extend_from_slice(bytes);
|
||||
let length = bytes.len();
|
||||
chunk.advance(length);
|
||||
}
|
||||
}
|
||||
Some(Err(error)) => return Err(Error::Body(Box::new(error))),
|
||||
None if buffer.is_empty() => return Ok(None),
|
||||
None => return Err(Error::Truncated),
|
||||
}
|
||||
}
|
||||
},
|
||||
)
|
||||
.fuse()
|
||||
}
|
||||
}
|
||||
17
litellm-rust/crates/framer/src/error.rs
Normal file
17
litellm-rust/crates/framer/src/error.rs
Normal file
|
|
@ -0,0 +1,17 @@
|
|||
#[derive(Debug, thiserror::Error)]
|
||||
pub enum Error {
|
||||
#[cfg(feature = "sse")]
|
||||
#[error("SSE framing failed: {0}")]
|
||||
Sse(#[from] sse_stream::Error),
|
||||
#[cfg(feature = "aws")]
|
||||
#[error("AWS EventStream framing failed: {0}")]
|
||||
Aws(#[from] aws_smithy_eventstream::error::Error),
|
||||
#[error("body stream failed: {0}")]
|
||||
Body(#[source] Box<dyn std::error::Error + Send + Sync>),
|
||||
#[cfg(feature = "aws")]
|
||||
#[error("invalid AWS EventStream frame length: {0}")]
|
||||
InvalidLength(usize),
|
||||
#[cfg(feature = "aws")]
|
||||
#[error("truncated AWS EventStream frame")]
|
||||
Truncated,
|
||||
}
|
||||
13
litellm-rust/crates/framer/src/framer.rs
Normal file
13
litellm-rust/crates/framer/src/framer.rs
Normal file
|
|
@ -0,0 +1,13 @@
|
|||
use futures_util::Stream;
|
||||
|
||||
use crate::Error;
|
||||
|
||||
pub trait Framer: Send {
|
||||
type Frame: Send;
|
||||
|
||||
fn frame<S, B, E>(self, input: S) -> impl Stream<Item = Result<Self::Frame, Error>> + Send
|
||||
where
|
||||
S: Stream<Item = Result<B, E>> + Send,
|
||||
B: bytes::Buf + Send,
|
||||
E: std::error::Error + Send + Sync + 'static;
|
||||
}
|
||||
10
litellm-rust/crates/framer/src/lib.rs
Normal file
10
litellm-rust/crates/framer/src/lib.rs
Normal file
|
|
@ -0,0 +1,10 @@
|
|||
mod error;
|
||||
mod framer;
|
||||
|
||||
pub use error::*;
|
||||
pub use framer::*;
|
||||
|
||||
#[cfg(feature = "aws")]
|
||||
pub mod aws_event_stream;
|
||||
#[cfg(feature = "sse")]
|
||||
pub mod sse;
|
||||
43
litellm-rust/crates/framer/src/sse.rs
Normal file
43
litellm-rust/crates/framer/src/sse.rs
Normal file
|
|
@ -0,0 +1,43 @@
|
|||
use futures_util::{Stream, StreamExt};
|
||||
|
||||
use crate::{Error, Framer};
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub struct SseFrame {
|
||||
pub event: Option<String>,
|
||||
pub data: Option<String>,
|
||||
pub id: Option<String>,
|
||||
pub retry: Option<u64>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, Default)]
|
||||
pub struct SseFramer;
|
||||
|
||||
impl Framer for SseFramer {
|
||||
type Frame = SseFrame;
|
||||
|
||||
fn frame<S, B, E>(self, input: S) -> impl Stream<Item = Result<SseFrame, Error>> + Send
|
||||
where
|
||||
S: Stream<Item = Result<B, E>> + Send,
|
||||
B: bytes::Buf + Send,
|
||||
E: std::error::Error + Send + Sync + 'static,
|
||||
{
|
||||
let frames = Box::pin(sse_stream::SseStream::from_bytes_stream(input));
|
||||
futures_util::stream::try_unfold(frames, |mut frames| async move {
|
||||
let Some(frame) = frames.next().await else {
|
||||
return Ok(None);
|
||||
};
|
||||
let frame = frame?;
|
||||
Ok(Some((
|
||||
SseFrame {
|
||||
event: frame.event,
|
||||
data: frame.data,
|
||||
id: frame.id,
|
||||
retry: frame.retry,
|
||||
},
|
||||
frames,
|
||||
)))
|
||||
})
|
||||
.fuse()
|
||||
}
|
||||
}
|
||||
92
litellm-rust/crates/framer/tests/aws_event_stream.rs
Normal file
92
litellm-rust/crates/framer/tests/aws_event_stream.rs
Normal file
|
|
@ -0,0 +1,92 @@
|
|||
#![cfg(feature = "aws")]
|
||||
|
||||
mod support;
|
||||
|
||||
use std::io;
|
||||
|
||||
use futures_util::TryStreamExt;
|
||||
use litellm_framing::aws_event_stream::{AwsEventStreamFrame, AwsEventStreamFramer};
|
||||
use litellm_framing::{Error, Framer};
|
||||
use rstest::{fixture, rstest};
|
||||
|
||||
use support::encode;
|
||||
|
||||
async fn collect_aws(bytes: &[u8], chunk_size: usize) -> Result<Vec<AwsEventStreamFrame>, Error> {
|
||||
AwsEventStreamFramer
|
||||
.frame(futures_util::stream::iter(
|
||||
bytes.chunks(chunk_size).map(Ok::<_, io::Error>),
|
||||
))
|
||||
.try_collect()
|
||||
.await
|
||||
}
|
||||
|
||||
#[fixture]
|
||||
fn two_frames() -> Vec<u8> {
|
||||
[encode(b"\xff\x00"), encode(b"second")].concat()
|
||||
}
|
||||
|
||||
#[fixture]
|
||||
fn payload_frame() -> Vec<u8> {
|
||||
encode(b"payload")
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case(1)]
|
||||
#[case(3)]
|
||||
#[case(12)]
|
||||
#[case(usize::MAX)]
|
||||
#[tokio::test]
|
||||
async fn fragmented_and_coalesced_frames_preserve_typed_headers_and_binary_payloads(
|
||||
two_frames: Vec<u8>,
|
||||
#[case] chunk_size: usize,
|
||||
) {
|
||||
let chunk_size = chunk_size.min(two_frames.len());
|
||||
let frames = collect_aws(&two_frames, chunk_size).await.unwrap();
|
||||
assert_eq!(frames.len(), 2);
|
||||
assert_eq!(frames[0].payload, &b"\xff\x00"[..]);
|
||||
assert_eq!(frames[1].payload, "second");
|
||||
assert_eq!(
|
||||
frames[0].headers[0].value().as_string().unwrap().as_str(),
|
||||
"payload"
|
||||
);
|
||||
assert_eq!(frames[0].headers[1].value().as_int32(), Ok(7));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case(8)]
|
||||
#[case(0)]
|
||||
#[tokio::test]
|
||||
async fn rejects_corrupt_crcs(payload_frame: Vec<u8>, #[case] index: usize) {
|
||||
let corrupt_index = if index == 0 {
|
||||
payload_frame.len() - 1
|
||||
} else {
|
||||
index
|
||||
};
|
||||
let mut corrupt = payload_frame;
|
||||
corrupt[corrupt_index] ^= 1;
|
||||
assert!(matches!(collect_aws(&corrupt, 3).await, Err(Error::Aws(_))));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case(0_u32)]
|
||||
#[case(15)]
|
||||
#[case(u32::MAX)]
|
||||
#[tokio::test]
|
||||
async fn rejects_invalid_lengths(#[case] length: u32) {
|
||||
assert!(matches!(
|
||||
collect_aws(&length.to_be_bytes(), 1).await,
|
||||
Err(Error::InvalidLength(_))
|
||||
));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case(1)]
|
||||
#[case(3)]
|
||||
#[case(5)]
|
||||
#[tokio::test]
|
||||
async fn rejects_truncation(payload_frame: Vec<u8>, #[case] end: usize) {
|
||||
assert!(matches!(
|
||||
collect_aws(&payload_frame[..end], 1).await,
|
||||
Err(Error::Truncated)
|
||||
));
|
||||
}
|
||||
29
litellm-rust/crates/framer/tests/chaining.rs
Normal file
29
litellm-rust/crates/framer/tests/chaining.rs
Normal file
|
|
@ -0,0 +1,29 @@
|
|||
#![cfg(all(feature = "aws", feature = "sse"))]
|
||||
|
||||
mod support;
|
||||
|
||||
use std::io;
|
||||
|
||||
use futures_util::TryStreamExt;
|
||||
use litellm_framing::Framer;
|
||||
use litellm_framing::aws_event_stream::{AwsEventStreamFrame, AwsEventStreamFramer};
|
||||
use litellm_framing::sse::SseFramer;
|
||||
|
||||
use support::encode;
|
||||
|
||||
#[tokio::test]
|
||||
async fn hosting_payloads_feed_the_same_sse_framer_across_envelope_boundaries() {
|
||||
let bytes = [encode(b"event: delta\ndata: hel"), encode(b"lo\nid: 7\n\n")].concat();
|
||||
let envelopes = AwsEventStreamFramer.frame(futures_util::stream::iter(
|
||||
bytes.chunks(3).map(Ok::<_, io::Error>),
|
||||
));
|
||||
let frames = SseFramer
|
||||
.frame(envelopes.map_ok(|frame: AwsEventStreamFrame| frame.payload))
|
||||
.try_collect::<Vec<_>>()
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(frames.len(), 1);
|
||||
assert_eq!(frames[0].event.as_deref(), Some("delta"));
|
||||
assert_eq!(frames[0].data.as_deref(), Some("hello"));
|
||||
assert_eq!(frames[0].id.as_deref(), Some("7"));
|
||||
}
|
||||
67
litellm-rust/crates/framer/tests/sse.rs
Normal file
67
litellm-rust/crates/framer/tests/sse.rs
Normal file
|
|
@ -0,0 +1,67 @@
|
|||
#![cfg(feature = "sse")]
|
||||
|
||||
use std::io;
|
||||
|
||||
use futures_util::{StreamExt, TryStreamExt};
|
||||
use litellm_framing::sse::{SseFrame, SseFramer};
|
||||
use litellm_framing::{Error, Framer};
|
||||
use rstest::rstest;
|
||||
|
||||
async fn collect_sse(chunks: &[&[u8]]) -> Result<Vec<SseFrame>, Error> {
|
||||
SseFramer
|
||||
.frame(futures_util::stream::iter(
|
||||
chunks.iter().copied().map(Ok::<_, io::Error>),
|
||||
))
|
||||
.try_collect()
|
||||
.await
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case(
|
||||
&[&b":ping\r\nevent: delta\r\nid: 7\r\nretry: 10\r\ndata: \xe2"[..], &b"\x82"[..], &b"\xac\r"[..], &b"\ndata: next\r\n\r"[..], &b"\ndata: [DONE]\n\n"[..]],
|
||||
vec![
|
||||
SseFrame {
|
||||
event: Some("delta".into()),
|
||||
data: Some("€\nnext".into()),
|
||||
id: Some("7".into()),
|
||||
retry: Some(10),
|
||||
},
|
||||
SseFrame {
|
||||
event: None,
|
||||
data: Some("[DONE]".into()),
|
||||
id: None,
|
||||
retry: None,
|
||||
},
|
||||
]
|
||||
)]
|
||||
#[tokio::test]
|
||||
async fn fragmented_utf8_crlf_and_multiline_data_retain_metadata_and_sentinel(
|
||||
#[case] chunks: &[&[u8]],
|
||||
#[case] expected: Vec<SseFrame>,
|
||||
) {
|
||||
assert_eq!(collect_sse(chunks).await.unwrap(), expected);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn eof_does_not_dispatch_an_unterminated_frame() {
|
||||
assert!(collect_sse(&[b"data: partial\n"]).await.unwrap().is_empty());
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case(io::ErrorKind::ConnectionReset)]
|
||||
#[case(io::ErrorKind::UnexpectedEof)]
|
||||
#[tokio::test]
|
||||
async fn framing_errors_terminate_and_preserve_input_error_causes(#[case] kind: io::ErrorKind) {
|
||||
let mut frames = Box::pin(SseFramer.frame(futures_util::stream::iter([
|
||||
Err(io::Error::new(kind, "reset")),
|
||||
Ok(&b"data: later\n\n"[..]),
|
||||
])));
|
||||
let error = frames.next().await.unwrap().unwrap_err();
|
||||
assert!(matches!(
|
||||
error,
|
||||
Error::Sse(sse_stream::Error::Body(ref cause))
|
||||
if cause.downcast_ref::<io::Error>().unwrap().kind() == kind
|
||||
));
|
||||
assert!(frames.next().await.is_none());
|
||||
assert!(frames.next().await.is_none());
|
||||
}
|
||||
15
litellm-rust/crates/framer/tests/support/mod.rs
Normal file
15
litellm-rust/crates/framer/tests/support/mod.rs
Normal file
|
|
@ -0,0 +1,15 @@
|
|||
use aws_smithy_eventstream::frame::write_message_to;
|
||||
use aws_smithy_types::event_stream::{Header, HeaderValue, Message};
|
||||
use bytes::Bytes;
|
||||
|
||||
pub fn encode(payload: &'static [u8]) -> Vec<u8> {
|
||||
let message = Message::new(Bytes::from_static(payload))
|
||||
.add_header(Header::new(
|
||||
":event-type",
|
||||
HeaderValue::String("payload".into()),
|
||||
))
|
||||
.add_header(Header::new("sequence", HeaderValue::Int32(7)));
|
||||
let mut bytes = Vec::new();
|
||||
write_message_to(&message, &mut bytes).unwrap();
|
||||
bytes
|
||||
}
|
||||
|
|
@ -46,10 +46,7 @@ pub(crate) fn core_error_to_pyerr(error: Error) -> PyErr {
|
|||
),
|
||||
Error::Messages(error) => match error {
|
||||
messages::Error::Auth(source) => auth_is_value_error(source),
|
||||
messages::Error::InvalidProvider(_)
|
||||
| messages::Error::InvalidRequest(_)
|
||||
| messages::Error::Headers(_) => true,
|
||||
_ => false,
|
||||
_ => error.is_request(),
|
||||
},
|
||||
Error::AudioTranscription(error) => match error {
|
||||
audio_transcription::Error::Auth(source) => auth_is_value_error(source),
|
||||
|
|
|
|||
|
|
@ -244,6 +244,7 @@ telemetry = True
|
|||
max_tokens: int = DEFAULT_MAX_TOKENS # OpenAI Defaults
|
||||
drop_params = drop_params_env_flag(os.environ, verbose_logger)
|
||||
modify_params = bool(os.getenv("LITELLM_MODIFY_PARAMS", False))
|
||||
bedrock_neutralize_orphaned_tool_blocks: bool = True
|
||||
use_chat_completions_url_for_anthropic_messages: bool = bool(
|
||||
os.getenv("LITELLM_USE_CHAT_COMPLETIONS_URL_FOR_ANTHROPIC_MESSAGES", False)
|
||||
) # When True, routes OpenAI /v1/messages requests to chat/completions instead of the Responses API
|
||||
|
|
@ -1833,6 +1834,9 @@ if TYPE_CHECKING:
|
|||
from .llms.azure.responses.o_series_transformation import (
|
||||
AzureOpenAIOSeriesResponsesAPIConfig as AzureOpenAIOSeriesResponsesAPIConfig,
|
||||
)
|
||||
from .llms.azure_ai.responses.transformation import (
|
||||
AzureAIResponsesAPIConfig as AzureAIResponsesAPIConfig,
|
||||
)
|
||||
from .llms.xai.responses.transformation import (
|
||||
XAIResponsesAPIConfig as XAIResponsesAPIConfig,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -234,6 +234,7 @@ LLM_CONFIG_NAMES: Final = (
|
|||
"OpenAIResponsesAPIConfig",
|
||||
"AzureOpenAIResponsesAPIConfig",
|
||||
"AzureOpenAIOSeriesResponsesAPIConfig",
|
||||
"AzureAIResponsesAPIConfig",
|
||||
"XAIResponsesAPIConfig",
|
||||
"LiteLLMProxyResponsesAPIConfig",
|
||||
"HostedVLLMResponsesAPIConfig",
|
||||
|
|
@ -946,6 +947,10 @@ _LLM_CONFIGS_IMPORT_MAP: Final = {
|
|||
".llms.azure.responses.o_series_transformation",
|
||||
"AzureOpenAIOSeriesResponsesAPIConfig",
|
||||
),
|
||||
"AzureAIResponsesAPIConfig": (
|
||||
".llms.azure_ai.responses.transformation",
|
||||
"AzureAIResponsesAPIConfig",
|
||||
),
|
||||
"XAIResponsesAPIConfig": (
|
||||
".llms.xai.responses.transformation",
|
||||
"XAIResponsesAPIConfig",
|
||||
|
|
|
|||
|
|
@ -3,6 +3,7 @@ import contextvars
|
|||
import functools
|
||||
import logging
|
||||
import os
|
||||
import re
|
||||
import sys
|
||||
from datetime import datetime
|
||||
from logging import Formatter
|
||||
|
|
@ -13,10 +14,11 @@ import litellm
|
|||
from litellm.constants import (
|
||||
LITELLM_TRUNCATED_PAYLOAD_FIELD,
|
||||
LITELLM_TRUNCATION_STDOUT_SAFEGUARD_NOTE,
|
||||
MAX_BASE64_LENGTH_STDOUT_LOG,
|
||||
MAX_STRING_LENGTH_STDOUT_LOG,
|
||||
)
|
||||
from litellm.litellm_core_utils.env_utils import get_env_int
|
||||
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
|
||||
from litellm.litellm_core_utils.safe_json_dumps import UNSERIALIZABLE_OBJECT, safe_dumps, safe_json_structure
|
||||
from litellm.litellm_core_utils.safe_json_loads import safe_json_loads
|
||||
from litellm.litellm_core_utils.secret_redaction import (
|
||||
redact_internal_details,
|
||||
|
|
@ -77,6 +79,37 @@ def _redact_structured_value(key: str | None, value: str) -> str:
|
|||
return redact_structured_value(key, value)
|
||||
|
||||
|
||||
_REDACTED_RECORD_ATTR: Final = "litellm_redacted"
|
||||
_REDACTED_STAMP: Final = object()
|
||||
_UNREDACTED_SCALAR_TYPES: Final = (bool, int, float, type(None))
|
||||
|
||||
|
||||
def _is_redacted(record: logging.LogRecord) -> bool:
|
||||
return getattr(record, _REDACTED_RECORD_ATTR, None) is _REDACTED_STAMP
|
||||
|
||||
|
||||
def _scrubbing_changed_nothing(scrubbed: object, original: object) -> bool:
|
||||
try:
|
||||
return bool(scrubbed == original)
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
|
||||
def _plain_text(value: object) -> str:
|
||||
try:
|
||||
return str(value)
|
||||
except Exception:
|
||||
return UNSERIALIZABLE_OBJECT
|
||||
|
||||
|
||||
def _redact_extra_value(key: str, value: object) -> object:
|
||||
try:
|
||||
scrubbed: Final = safe_json_structure(value, value_transform=_redact_structured_value, key=key)
|
||||
except Exception:
|
||||
return _redact_string(_plain_text(value))
|
||||
return value if _scrubbing_changed_nothing(scrubbed, value) else scrubbed
|
||||
|
||||
|
||||
def redact_secrets(value: str) -> str:
|
||||
"""Public API: redact known secret/credential patterns from an arbitrary string.
|
||||
|
||||
|
|
@ -126,7 +159,7 @@ class SecretRedactionFilter(logging.Filter):
|
|||
_formatter = logging.Formatter()
|
||||
|
||||
def filter(self, record: logging.LogRecord) -> bool:
|
||||
if not _ENABLE_SECRET_REDACTION:
|
||||
if not _ENABLE_SECRET_REDACTION or _is_redacted(record):
|
||||
return True
|
||||
|
||||
# Runs before args are cleared, and before the extra-field loop below
|
||||
|
|
@ -149,11 +182,19 @@ class SecretRedactionFilter(logging.Filter):
|
|||
except Exception:
|
||||
pass
|
||||
|
||||
if isinstance(record.stack_info, str):
|
||||
record.stack_info = _redact_string(record.stack_info) # rebind-ok: a Filter scrubs records in place
|
||||
|
||||
# Redact extra fields passed via logger.debug("msg", extra={...})
|
||||
for key, value in list(record.__dict__.items()):
|
||||
if key not in _STANDARD_RECORD_ATTRS and isinstance(value, str):
|
||||
setattr(record, key, _redact_string(value))
|
||||
if key in _STANDARD_RECORD_ATTRS:
|
||||
continue
|
||||
if isinstance(value, str):
|
||||
setattr(record, key, _redact_structured_value(key, value))
|
||||
elif not isinstance(value, _UNREDACTED_SCALAR_TYPES):
|
||||
setattr(record, key, _redact_extra_value(key, value))
|
||||
|
||||
setattr(record, _REDACTED_RECORD_ATTR, _REDACTED_STAMP)
|
||||
return True
|
||||
|
||||
|
||||
|
|
@ -277,6 +318,51 @@ def _truncate_for_stdout_log(text: str, limit: int) -> str:
|
|||
return f"{text[:head_chars]}{_stdout_truncation_marker(len(text) - kept_chars)}{text[-tail_chars:]}"
|
||||
|
||||
|
||||
_BYTES_PER_KIB: Final = 1024
|
||||
_BYTES_PER_MIB: Final = 1024 * 1024
|
||||
|
||||
|
||||
def format_base64_size(num_chars: int) -> str:
|
||||
"""Return a human-readable byte-size estimate from a base64 character count."""
|
||||
num_bytes: Final = num_chars * 3 / 4
|
||||
if num_bytes >= _BYTES_PER_MIB:
|
||||
return f"{num_bytes / _BYTES_PER_MIB:.2f}MB"
|
||||
if num_bytes >= _BYTES_PER_KIB:
|
||||
return f"{num_bytes / _BYTES_PER_KIB:.1f}KB"
|
||||
return f"{int(num_bytes)}B"
|
||||
|
||||
|
||||
def _get_max_base64_length_stdout_log() -> int:
|
||||
return get_env_int("MAX_BASE64_LENGTH_STDOUT_LOG", MAX_BASE64_LENGTH_STDOUT_LOG)
|
||||
|
||||
|
||||
@functools.lru_cache(maxsize=8)
|
||||
def _base64_run_pattern(min_chars: int) -> "re.Pattern[str]":
|
||||
return re.compile(rf"(?<![A-Za-z0-9+/])[A-Za-z0-9+/]{{{min_chars},}}={{0,2}}")
|
||||
|
||||
|
||||
_LOWER_HEX_DIGITS: Final = "0123456789abcdef"
|
||||
_UPPER_HEX_DIGITS: Final = "0123456789ABCDEF"
|
||||
|
||||
|
||||
def _looks_like_base64(run: str) -> bool:
|
||||
unpadded: Final = run.rstrip("=")
|
||||
is_hex_or_decimal: Final = not unpadded.strip(_LOWER_HEX_DIGITS) or not unpadded.strip(_UPPER_HEX_DIGITS)
|
||||
is_one_repeated_char: Final = not unpadded.strip(unpadded[0])
|
||||
return not is_hex_or_decimal or is_one_repeated_char
|
||||
|
||||
|
||||
def _replace_base64_run(match: "re.Match[str]") -> str:
|
||||
run: Final = match.group(0)
|
||||
if not _looks_like_base64(run):
|
||||
return run
|
||||
return f"[base64_data truncated: {format_base64_size(len(run))}]"
|
||||
|
||||
|
||||
def _collapse_base64_runs(text: str, limit: int) -> str:
|
||||
return _base64_run_pattern(limit + 1).sub(_replace_base64_run, text)
|
||||
|
||||
|
||||
class StdoutLogTruncationFilter(logging.Filter):
|
||||
"""Bounds how much of an oversized log line reaches stdout.
|
||||
|
||||
|
|
@ -284,36 +370,42 @@ class StdoutLogTruncationFilter(logging.Filter):
|
|||
request writes hundreds of KB to stdout, repeatedly as the exception propagates from
|
||||
the router to the proxy handler and into its traceback, all inline on the event loop.
|
||||
|
||||
DEBUG records pass through untouched, since dumping full payloads is the point of
|
||||
At every level, in the message and in the traceback alike, a base64 run longer than
|
||||
MAX_BASE64_LENGTH_STDOUT_LOG collapses to a size placeholder first: a multi-megabyte
|
||||
document upload otherwise costs seconds of event-loop time per DEBUG line in the
|
||||
secret regex alone. Hex and decimal runs (digests, numeric ids) are left alone unless
|
||||
they are one repeated character, which is what a zero-filled payload encodes to.
|
||||
The text around a run stays, since dumping payloads is the point of
|
||||
`--detailed_debug`, and logging callbacks (OTEL, Datadog, etc.) don't run through
|
||||
logging filters at all, so they still get the untruncated error.
|
||||
logging filters at all, so they still get the untouched record.
|
||||
"""
|
||||
|
||||
_formatter = logging.Formatter()
|
||||
|
||||
def filter(self, record: logging.LogRecord) -> bool:
|
||||
if record.levelno < logging.INFO:
|
||||
return True
|
||||
|
||||
limit: Final = _get_max_string_length_stdout_log()
|
||||
if limit <= 0:
|
||||
return True
|
||||
|
||||
try:
|
||||
message: Final = record.getMessage()
|
||||
except (TypeError, ValueError):
|
||||
return True
|
||||
|
||||
if len(message) > limit:
|
||||
record.msg = _truncate_for_stdout_log(message, limit) # rebind-ok: the Filter interface mutates the record
|
||||
record.args = None # rebind-ok: args are consumed by the truncated message above
|
||||
base64_limit: Final = _get_max_base64_length_stdout_log()
|
||||
collapsed: Final = _collapse_base64_runs(message, base64_limit) if base64_limit > 0 else message
|
||||
limit: Final = _get_max_string_length_stdout_log() if record.levelno >= logging.INFO else 0
|
||||
bounded: Final = _truncate_for_stdout_log(collapsed, limit) if 0 < limit < len(collapsed) else collapsed
|
||||
if bounded != message:
|
||||
record.msg = bounded # rebind-ok: the Filter interface mutates the record
|
||||
record.args = None # rebind-ok: args are consumed by the rewritten message above
|
||||
|
||||
if isinstance(record.exc_info, tuple):
|
||||
exc_text: Final = record.exc_text or self._formatter.formatException(record.exc_info)
|
||||
if len(exc_text) > limit:
|
||||
record.exc_text = _truncate_for_stdout_log( # rebind-ok: the Filter interface mutates the record
|
||||
exc_text, limit
|
||||
)
|
||||
if not isinstance(record.exc_info, tuple):
|
||||
return True
|
||||
|
||||
exc_text: Final = record.exc_text or self._formatter.formatException(record.exc_info)
|
||||
collapsed_exc: Final = _collapse_base64_runs(exc_text, base64_limit) if base64_limit > 0 else exc_text
|
||||
bounded_exc: Final = (
|
||||
_truncate_for_stdout_log(collapsed_exc, limit) if 0 < limit < len(collapsed_exc) else collapsed_exc
|
||||
)
|
||||
if bounded_exc != exc_text:
|
||||
record.exc_text = bounded_exc # rebind-ok: the Filter interface mutates the record
|
||||
|
||||
return True
|
||||
|
||||
|
|
@ -474,6 +566,7 @@ def _get_standard_record_attrs() -> frozenset:
|
|||
|
||||
|
||||
_STANDARD_RECORD_ATTRS: Final = _get_standard_record_attrs()
|
||||
_NON_EXTRA_RECORD_ATTRS: Final = _STANDARD_RECORD_ATTRS | {_REDACTED_RECORD_ATTR}
|
||||
|
||||
# CorrelationContextFilter is the only legitimate source for these two JSON fields;
|
||||
# see JsonFormatter.format() for why they're excluded from the generic message-content
|
||||
|
|
@ -514,7 +607,7 @@ class JsonFormatter(Formatter):
|
|||
|
||||
# Include extra attributes passed via logger.debug("msg", extra={...})
|
||||
for key, value in record.__dict__.items():
|
||||
if key not in _STANDARD_RECORD_ATTRS and key not in json_record:
|
||||
if key not in _NON_EXTRA_RECORD_ATTRS and key not in json_record:
|
||||
json_record[key] = value
|
||||
|
||||
# trace_id/session_id are reserved: CorrelationContextFilter is the only
|
||||
|
|
@ -538,7 +631,7 @@ class JsonFormatter(Formatter):
|
|||
if record.exc_info:
|
||||
json_record["stacktrace"] = record.exc_text or self.formatException(record.exc_info)
|
||||
|
||||
return safe_dumps(json_record, value_transform=_redact_structured_value)
|
||||
return safe_dumps(json_record, value_transform=None if _is_redacted(record) else _redact_structured_value)
|
||||
|
||||
|
||||
class CorrelationPlainFormatter(logging.Formatter):
|
||||
|
|
@ -549,7 +642,8 @@ class CorrelationPlainFormatter(logging.Formatter):
|
|||
"""
|
||||
|
||||
def format(self, record: logging.LogRecord) -> str:
|
||||
formatted: Final = _redact_string(super().format(record))
|
||||
rendered: Final = super().format(record)
|
||||
formatted: Final = rendered if _is_redacted(record) else _redact_string(rendered)
|
||||
trace_id: Final = getattr(record, "trace_id", None)
|
||||
session_id: Final = getattr(record, "session_id", None)
|
||||
if not trace_id and not session_id:
|
||||
|
|
@ -567,8 +661,8 @@ def _setup_json_exception_handlers(formatter):
|
|||
# Create a handler with JSON formatting for exceptions
|
||||
error_handler: Final = logging.StreamHandler()
|
||||
error_handler.setFormatter(formatter)
|
||||
error_handler.addFilter(_secret_filter)
|
||||
error_handler.addFilter(_stdout_truncation_filter)
|
||||
error_handler.addFilter(_secret_filter)
|
||||
error_handler.addFilter(_correlation_filter)
|
||||
|
||||
# Setup excepthook for uncaught exceptions
|
||||
|
|
|
|||
|
|
@ -502,7 +502,12 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
|
|||
elif key == "response_format":
|
||||
text_format = self._transform_response_format_to_text_format(value)
|
||||
if text_format:
|
||||
responses_api_request["text"] = text_format
|
||||
responses_api_request["text"] = self._merge_text(responses_api_request, text_format)
|
||||
elif key == "verbosity":
|
||||
responses_api_request["text"] = self._merge_text(
|
||||
responses_api_request,
|
||||
MappingProxyType({"verbosity": value}), # pyright: ignore[reportUnknownArgumentType] # untyped value
|
||||
)
|
||||
elif key == "tool_choice":
|
||||
responses_api_request["tool_choice"] = self._normalize_tool_choice_for_responses_api(value)
|
||||
elif key == "stream_options":
|
||||
|
|
@ -518,6 +523,19 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
|
|||
elif key == "web_search_options":
|
||||
self._add_web_search_tool(responses_api_request, value)
|
||||
|
||||
@staticmethod
|
||||
def _merge_text(
|
||||
responses_api_request: "ResponsesAPIOptionalRequestParams", update: Mapping[str, object]
|
||||
) -> "ResponseText":
|
||||
existing: Final = cast( # cast-ok: text field is a ResponseText | dict[str, Any] | None union
|
||||
"dict[str, object]",
|
||||
dict(responses_api_request).get("text") or {}, # mutable-ok: one-shot merge seed
|
||||
)
|
||||
return cast( # cast-ok: merged mapping is a valid ResponseText shape
|
||||
"ResponseText",
|
||||
{**existing, **update}, # mutable-ok: one-shot merged payload
|
||||
)
|
||||
|
||||
def _build_sanitized_litellm_params(self, litellm_params: dict) -> dict[str, object]:
|
||||
"""Build sanitized litellm_params with merged metadata."""
|
||||
responses_optional_param_keys: Final = set(ResponsesAPIOptionalRequestParams.__annotations__.keys())
|
||||
|
|
|
|||
|
|
@ -102,6 +102,7 @@ REDACTED_BY_LITELLM: Final = "redacted-by-litellm"
|
|||
REDACTED_TOOL_CALL_ARGUMENTS_PLACEHOLDER: Final = "{}"
|
||||
|
||||
MAX_STRING_LENGTH_STDOUT_LOG: Final = get_env_int("MAX_STRING_LENGTH_STDOUT_LOG", 4096)
|
||||
MAX_BASE64_LENGTH_STDOUT_LOG: Final = get_env_int("MAX_BASE64_LENGTH_STDOUT_LOG", 4096)
|
||||
|
||||
# When true, adds detailed per-phase timing breakdown headers to responses.
|
||||
# Headers: x-litellm-timing-{pre-processing,llm-api,post-processing,message-copy}-ms
|
||||
|
|
|
|||
|
|
@ -346,13 +346,17 @@ class MCPClient:
|
|||
self.update_auth_value(auth_value)
|
||||
|
||||
async def discovery_auth_fingerprint(self) -> str:
|
||||
return self._hash_discovery_auth(await self.prepare_request_auth())
|
||||
|
||||
async def prepare_request_auth(self) -> httpx.Request:
|
||||
"""Preview the authenticated request without sending it, closing the auth flow afterwards."""
|
||||
request: Final = httpx.Request("POST", self.server_url or "http://localhost/", headers=self._get_auth_headers())
|
||||
if self._resolved_auth is None:
|
||||
return self._hash_discovery_auth(request)
|
||||
return request
|
||||
flow: Final = self._resolved_auth.async_auth_flow(request)
|
||||
try:
|
||||
authenticated: Final = await flow.__anext__()
|
||||
return self._hash_discovery_auth(authenticated)
|
||||
return authenticated
|
||||
finally:
|
||||
await flow.aclose()
|
||||
|
||||
|
|
|
|||
|
|
@ -16,6 +16,7 @@ from litellm.litellm_core_utils.core_helpers import (
|
|||
get_or_create_metadata_bucket,
|
||||
redact_nested_match_and_regex_keys,
|
||||
)
|
||||
from litellm.llms.base_llm.guardrail_translation.base_translation import REQUEST_SCAN_CONTEXT_KEY
|
||||
from litellm.secret_managers.main import str_to_bool
|
||||
from litellm.types.guardrails import (
|
||||
DynamicGuardrailParams,
|
||||
|
|
@ -949,9 +950,28 @@ class CustomGuardrail(CustomLogger):
|
|||
await translation.process_input_messages(data=scratch_request, guardrail_to_apply=self)
|
||||
if response is None:
|
||||
return
|
||||
await output_translation.process_output_response(
|
||||
response=copy.deepcopy(response), guardrail_to_apply=self, request_data=scratch_request
|
||||
output_request: Final = (
|
||||
scratch_request
|
||||
if type(output_translation) is type(translation)
|
||||
else self._chat_shaped_request(scratch_request, translation)
|
||||
)
|
||||
await output_translation.process_output_response(
|
||||
response=copy.deepcopy(response), guardrail_to_apply=self, request_data=output_request
|
||||
)
|
||||
|
||||
def _chat_shaped_request(
|
||||
self,
|
||||
scratch_request: Mapping[str, object],
|
||||
translation: "BaseTranslation",
|
||||
) -> dict[str, object]: # mutable-ok: BaseTranslation.process_output_response contract
|
||||
"""The logged request in OpenAI chat shape, for an output scan whose translation differs from the input's."""
|
||||
context: Final = translation.request_scan_context(scratch_request, self)
|
||||
return {
|
||||
**scratch_request,
|
||||
"messages": list(context.structured_messages),
|
||||
"tools": list(context.tools),
|
||||
REQUEST_SCAN_CONTEXT_KEY: context,
|
||||
}
|
||||
|
||||
def supports_scan_only_tool_results(self) -> bool:
|
||||
"""Whether this guardrail can scan tool-result content.
|
||||
|
|
@ -1379,8 +1399,9 @@ class CustomGuardrail(CustomLogger):
|
|||
raise e
|
||||
|
||||
def _inputs_were_modified(self, original_inputs: Mapping[str, object], response: Mapping[str, object]) -> bool:
|
||||
"""True when any key of either mapping differs between them (mask), False otherwise (allow)."""
|
||||
return any(original_inputs.get(key) != response.get(key) for key in original_inputs.keys() | response.keys())
|
||||
"""True when any content key of either mapping differs between them (mask), False otherwise (allow)."""
|
||||
compared_keys: Final = (original_inputs.keys() | response.keys()) - _STREAM_CONTROL_KEYS
|
||||
return any(original_inputs.get(key) != response.get(key) for key in compared_keys)
|
||||
|
||||
def mask_content_in_string(
|
||||
self,
|
||||
|
|
@ -1490,6 +1511,7 @@ def _sync_guardrail_info_to_logging_obj(request_data: dict, logging_obj: object)
|
|||
_PRE_CALL_CONTENT_KEYS: Final = frozenset(
|
||||
{"messages", "input", "prompt", "system", "instructions", "tools", "functions", "function_call", "tool_choice"}
|
||||
)
|
||||
_STREAM_CONTROL_KEYS: Final = frozenset({"stream_holdback_chars"})
|
||||
|
||||
|
||||
def _original_inputs_for(
|
||||
|
|
|
|||
|
|
@ -2965,16 +2965,13 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
|
|||
)
|
||||
|
||||
propagator: Final = TraceContextTextMapPropagator()
|
||||
carrier: Final = {"traceparent": _traceparent}
|
||||
carrier: Final = {key: headers[key] for key in ("traceparent", "tracestate") if headers.get(key) is not None}
|
||||
_parent_context: Final = propagator.extract(carrier=carrier)
|
||||
|
||||
return _parent_context
|
||||
|
||||
def _get_span_context(self, kwargs, default_span: Span | None = None):
|
||||
from opentelemetry import context, trace
|
||||
from opentelemetry.trace.propagation.tracecontext import (
|
||||
TraceContextTextMapPropagator,
|
||||
)
|
||||
|
||||
litellm_params: Final = kwargs.get("litellm_params", {}) or {}
|
||||
proxy_server_request: Final = litellm_params.get("proxy_server_request", {}) or {}
|
||||
|
|
@ -2998,11 +2995,7 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
|
|||
# Priority 2: HTTP traceparent header
|
||||
if traceparent is not None:
|
||||
verbose_logger.debug("OpenTelemetry: Using traceparent header for context propagation")
|
||||
carrier: Final = {"traceparent": traceparent}
|
||||
return (
|
||||
TraceContextTextMapPropagator().extract(carrier=carrier),
|
||||
None,
|
||||
)
|
||||
return self.get_traceparent_from_header(headers=headers), None
|
||||
|
||||
# Priority 3: Active span from global context (auto-detection)
|
||||
try:
|
||||
|
|
|
|||
|
|
@ -1,15 +1,19 @@
|
|||
"""The span engine: dedup, start, run the mapper chain, set status, end."""
|
||||
|
||||
from collections import OrderedDict
|
||||
from collections.abc import Callable, Sequence
|
||||
from collections.abc import Callable, Mapping, Sequence
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
|
||||
from opentelemetry.context import Context
|
||||
from opentelemetry.sdk.trace import ReadableSpan, SpanLimits
|
||||
from opentelemetry.sdk.trace import Span as SdkSpan
|
||||
from opentelemetry.trace import Link, Span, Tracer
|
||||
from opentelemetry.trace.status import Status, StatusCode
|
||||
|
||||
from litellm.integrations.otel.mappers import resolve_mappers
|
||||
from litellm.integrations.otel.mappers.base import AttributeMapper, SpanData
|
||||
from litellm.integrations.otel.mappers.base import AttributeMapper, AttrValue, SpanData
|
||||
from litellm.integrations.otel.mappers.openinference import fit_indexed_messages
|
||||
from litellm.integrations.otel.model.config import OpenTelemetryV2Config
|
||||
from litellm.integrations.otel.model.payloads import (
|
||||
GuardrailSpanData,
|
||||
|
|
@ -52,25 +56,48 @@ _NAME_BUILDERS: Final[dict[SpanRole, Callable[..., str]]] = {
|
|||
_DEDUP_CACHE_MAX: Final = 10_000
|
||||
|
||||
|
||||
def _stamp_otel_error_attributes(span: Span, error_type: str, resolved_message: str) -> None:
|
||||
"""Stamp the OTel-semconv error attributes (``error.type`` + ``error.message``).
|
||||
``error_type`` and ``resolved_message`` are ``finish_span``'s already-computed
|
||||
fallback chains, so the pair on the status, event, and attributes stays in
|
||||
lockstep."""
|
||||
span.set_attribute(Error.TYPE, error_type)
|
||||
span.set_attribute(Error.MESSAGE, resolved_message)
|
||||
def _resolve_error(error: SpanError) -> tuple[str, str] | None:
|
||||
"""The ``(error_type, message)`` fallback chain shared by the status, the event and the attributes, or
|
||||
``None`` when ``error`` carries neither a type nor a message."""
|
||||
if not (error.error_type or error.message):
|
||||
return None
|
||||
return error.error_type or "error", error.message or error.error_type or "error"
|
||||
|
||||
|
||||
def _stamp_litellm_error_attributes(span: Span, error: SpanError) -> None:
|
||||
"""Stamp litellm-specific error detail attributes. Emitted only when the
|
||||
corresponding field is populated so guardrail-shape errors carrying only a
|
||||
message aren't polluted with empty detail keys."""
|
||||
if error.code:
|
||||
span.set_attribute(LiteLLMError.CODE, error.code)
|
||||
if error.stack_trace:
|
||||
span.set_attribute(LiteLLMError.STACK_TRACE, error.stack_trace)
|
||||
if error.llm_provider:
|
||||
span.set_attribute(LiteLLMError.LLM_PROVIDER, error.llm_provider)
|
||||
_NO_ATTRIBUTES: Final[Mapping[str, AttrValue]] = MappingProxyType({})
|
||||
|
||||
|
||||
def error_attributes(error: SpanError) -> Mapping[str, AttrValue]:
|
||||
"""The v2 error attribute set: the OTel-semconv ``error.*`` pair plus the litellm detail keys that are
|
||||
populated, so guardrail-shape errors carrying only a message aren't polluted with empty detail keys."""
|
||||
resolved: Final = _resolve_error(error)
|
||||
if resolved is None:
|
||||
return _NO_ATTRIBUTES
|
||||
error_type, message = resolved
|
||||
pairs: Final = (
|
||||
(Error.TYPE, error_type),
|
||||
(Error.MESSAGE, message),
|
||||
(LiteLLMError.CODE, error.code),
|
||||
(LiteLLMError.STACK_TRACE, error.stack_trace),
|
||||
(LiteLLMError.LLM_PROVIDER, error.llm_provider),
|
||||
)
|
||||
return MappingProxyType({key: value for key, value in pairs if value})
|
||||
|
||||
|
||||
def span_attribute_limit(span: Span) -> int | None:
|
||||
"""The attribute count limit ``span`` was built with, ``None`` when unbounded."""
|
||||
if not isinstance(span, SdkSpan):
|
||||
return SpanLimits().max_span_attributes
|
||||
return span._limits.max_span_attributes # pyright: ignore[reportPrivateUsage] # SDK has no public getter
|
||||
|
||||
|
||||
def attribute_budget(span: Span, reserved: int) -> int | None:
|
||||
"""How many mapped attributes fit on ``span`` next to what it already carries and ``reserved`` more."""
|
||||
limit: Final = span_attribute_limit(span)
|
||||
if limit is None:
|
||||
return None
|
||||
on_span: Final = len(span.attributes or ()) if isinstance(span, ReadableSpan) else 0
|
||||
return limit - on_span - reserved
|
||||
|
||||
|
||||
def stamp_error(
|
||||
|
|
@ -93,12 +120,12 @@ def stamp_error(
|
|||
``set_status`` are opt-outs for callers whose span lifecycle (``use_span``) or
|
||||
owner (the FastAPI instrumentor) already records the event or the status.
|
||||
"""
|
||||
if not (error.error_type or error.message):
|
||||
resolved: Final = _resolve_error(error)
|
||||
if resolved is None:
|
||||
return None
|
||||
error_type: Final = error.error_type or "error"
|
||||
message: Final = error.message or error.error_type or "error"
|
||||
_stamp_otel_error_attributes(span, error_type, message)
|
||||
_stamp_litellm_error_attributes(span, error)
|
||||
error_type, message = resolved
|
||||
for key, value in error_attributes(error).items():
|
||||
span.set_attribute(key, value)
|
||||
if set_status:
|
||||
span.set_status(Status(StatusCode.ERROR, message))
|
||||
if record_event:
|
||||
|
|
@ -238,9 +265,6 @@ class SpanEmitter:
|
|||
data, since the boundary opener only has a provisional name.
|
||||
"""
|
||||
span.update_name(_NAME_BUILDERS[role](data))
|
||||
for mapper in self._mappers:
|
||||
for key, value in mapper.map(data).items():
|
||||
span.set_attribute(key, value)
|
||||
error: Final = (
|
||||
data.error
|
||||
if isinstance(
|
||||
|
|
@ -255,6 +279,13 @@ class SpanEmitter:
|
|||
)
|
||||
else None
|
||||
)
|
||||
mapped: Final = MappingProxyType(
|
||||
{key: value for mapper in self._mappers for key, value in mapper.map(data).items()}
|
||||
)
|
||||
stamped_later: Final = error_attributes(error) if error else _NO_ATTRIBUTES
|
||||
reserved: Final = len(stamped_later.keys() - mapped.keys())
|
||||
for key, value in fit_indexed_messages(mapped, attribute_budget(span, reserved)).items():
|
||||
span.set_attribute(key, value)
|
||||
if error:
|
||||
stamped: Final = stamp_error(span, error)
|
||||
if stamped is not None and self._event_recorder is not None and role is SpanRole.LLM_CALL:
|
||||
|
|
|
|||
|
|
@ -6,10 +6,10 @@ from litellm.integrations.otel.logger import OpenTelemetryV2
|
|||
from litellm.integrations.otel.mappers.langfuse import (
|
||||
LANGFUSE_OBSERVATION_INPUT,
|
||||
LANGFUSE_OBSERVATION_OUTPUT,
|
||||
LANGFUSE_TRACE_NAME,
|
||||
LangfuseMapper,
|
||||
)
|
||||
from litellm.integrations.otel.model.metadata import caller_trace_name
|
||||
from litellm.integrations.otel.model.request_io import request_input, response_output, stream_output
|
||||
from litellm.integrations.otel.model.trace_controls import caller_trace_controls
|
||||
from litellm.integrations.otel.plumbing.context import request_root_span
|
||||
|
||||
if TYPE_CHECKING:
|
||||
|
|
@ -18,14 +18,13 @@ if TYPE_CHECKING:
|
|||
|
||||
|
||||
class LangfuseOpenTelemetryV2(OpenTelemetryV2):
|
||||
"""Names the trace from the request. Langfuse reads ``langfuse.trace.name`` off the root observation,
|
||||
and the proxy's root span is still recording when the LLM call starts."""
|
||||
"""Stamps the caller's trace controls (name, user, session, tags) on the request. Langfuse reads them off
|
||||
the root observation, and the proxy's root span is still recording when the LLM call starts."""
|
||||
|
||||
def log_pre_api_call(self, model: str, messages: object, kwargs: Mapping[str, object]) -> None:
|
||||
root: Final = request_root_span()
|
||||
name: Final = caller_trace_name(kwargs)
|
||||
if root is not None and root.is_recording() and name is not None:
|
||||
root.set_attribute(LANGFUSE_TRACE_NAME, name)
|
||||
if root is not None and root.is_recording():
|
||||
root.set_attributes(LangfuseMapper.trace_attributes(caller_trace_controls(kwargs)))
|
||||
super().log_pre_api_call(model, messages, kwargs)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -555,7 +555,7 @@ class OpenTelemetryV2(CustomLogger):
|
|||
capture_content=self.config.capture_span_content,
|
||||
time_to_first_chunk_seconds=call.time_to_first_chunk_seconds,
|
||||
request_route=request_root_http_route(),
|
||||
trace_name=call.trace_name,
|
||||
trace=call.trace,
|
||||
)
|
||||
end_time_ns: Final = to_ns(end_time)
|
||||
if carrier is not None and carrier.span is not None:
|
||||
|
|
|
|||
|
|
@ -6,7 +6,8 @@ Langfuse ingests OTLP spans and reads from its own vendor namespace
|
|||
|
||||
Every attribute is declared as a ``key -> extractor`` table entry (one callable
|
||||
per mapping operation): ``_LLM_CALL_ATTRS`` for scalars and ``_BLOB_ATTRS`` for
|
||||
the JSON-serialized payloads. ``_llm_call`` just applies both tables.
|
||||
the JSON-serialized payloads. ``trace_attributes`` maps the caller's trace controls
|
||||
(shared with the root observation); ``_llm_call`` applies both tables plus it.
|
||||
"""
|
||||
|
||||
import json
|
||||
|
|
@ -16,6 +17,7 @@ from typing import Final
|
|||
from litellm.integrations.otel.mappers.base import AttributeMap, AttrValue, SpanData
|
||||
from litellm.integrations.otel.mappers.utils import (
|
||||
collect,
|
||||
drop_none_pairs,
|
||||
json_if,
|
||||
output_messages,
|
||||
serialize_messages,
|
||||
|
|
@ -25,10 +27,14 @@ from litellm.integrations.otel.model.payloads import (
|
|||
LLMRequestParams,
|
||||
LLMUsage,
|
||||
)
|
||||
from litellm.integrations.otel.model.trace_controls import TraceControls
|
||||
|
||||
LANGFUSE_OBSERVATION_INPUT: Final = "langfuse.observation.input"
|
||||
LANGFUSE_OBSERVATION_OUTPUT: Final = "langfuse.observation.output"
|
||||
LANGFUSE_TRACE_NAME: Final = "langfuse.trace.name"
|
||||
LANGFUSE_TRACE_USER_ID: Final = "user.id"
|
||||
LANGFUSE_TRACE_SESSION_ID: Final = "session.id"
|
||||
LANGFUSE_TRACE_TAGS: Final = "langfuse.trace.tags"
|
||||
|
||||
|
||||
class LangfuseMapper:
|
||||
|
|
@ -37,7 +43,6 @@ class LangfuseMapper:
|
|||
"langfuse.observation.model.name": lambda d: d.request_model or None,
|
||||
"langfuse.observation.metadata.provider": lambda d: d.provider or None,
|
||||
"langfuse.observation.id": lambda d: d.identity.call_id or None,
|
||||
LANGFUSE_TRACE_NAME: lambda d: d.trace_name or None,
|
||||
"langfuse.trace.metadata.team_id": lambda d: d.identity.team_id or None,
|
||||
"langfuse.trace.metadata.team_alias": lambda d: d.identity.team_alias or None,
|
||||
}
|
||||
|
|
@ -77,9 +82,21 @@ class LangfuseMapper:
|
|||
case _:
|
||||
return {}
|
||||
|
||||
@staticmethod
|
||||
def trace_attributes(trace: TraceControls) -> AttributeMap:
|
||||
return drop_none_pairs(
|
||||
(
|
||||
(LANGFUSE_TRACE_NAME, trace.name or None),
|
||||
(LANGFUSE_TRACE_USER_ID, trace.user_id or None),
|
||||
(LANGFUSE_TRACE_SESSION_ID, trace.session_id or None),
|
||||
(LANGFUSE_TRACE_TAGS, trace.tags or None),
|
||||
)
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def _llm_call(cls, data: LLMCallSpanData) -> AttributeMap:
|
||||
return {
|
||||
**collect(cls._LLM_CALL_ATTRS, data),
|
||||
**cls.trace_attributes(data.trace),
|
||||
**collect(cls._BLOB_ATTRS, data),
|
||||
}
|
||||
|
|
|
|||
|
|
@ -7,12 +7,13 @@ Phoenix + any other OpenInference-aware backend simultaneously.
|
|||
"""
|
||||
|
||||
import json
|
||||
from collections.abc import Callable, Sequence
|
||||
from collections.abc import Callable, Mapping, Sequence
|
||||
from itertools import accumulate, chain, groupby
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
|
||||
from litellm.integrations.otel.mappers.base import AttributeMap, AttrValue, SpanData
|
||||
from litellm.integrations.otel.mappers.utils import (
|
||||
MAX_MESSAGE_ATTRS_PER_SPAN,
|
||||
MAX_TOOL_DEFINITION_ATTRS_PER_SPAN,
|
||||
collect,
|
||||
drop_none,
|
||||
|
|
@ -27,7 +28,53 @@ from litellm.integrations.otel.model.payloads import (
|
|||
ToolDefinition,
|
||||
)
|
||||
|
||||
_MAX_INDEXED_MESSAGES: Final = MAX_MESSAGE_ATTRS_PER_SPAN // 2
|
||||
_INPUT_MESSAGES: Final = "llm.input_messages"
|
||||
_OUTPUT_MESSAGES: Final = "llm.output_messages"
|
||||
_MESSAGE_FAMILIES: Final = (_INPUT_MESSAGES, _OUTPUT_MESSAGES)
|
||||
|
||||
|
||||
def _message_key_groups(attrs: Mapping[str, AttrValue]) -> Mapping[tuple[str, int], tuple[str, ...]]:
|
||||
"""Per-index message keys in ``attrs`` grouped by ``(family, index)``."""
|
||||
tagged: Final = sorted(
|
||||
(family, int(key.split(".")[2]), key)
|
||||
for key in attrs
|
||||
for family in _MESSAGE_FAMILIES
|
||||
if key.startswith(f"{family}.")
|
||||
)
|
||||
return MappingProxyType(
|
||||
{group: tuple(key for _, _, key in keys) for group, keys in groupby(tagged, key=lambda tag: tag[:2])}
|
||||
)
|
||||
|
||||
|
||||
def _shed_order(groups: Mapping[tuple[str, int], tuple[str, ...]]) -> tuple[tuple[str, int], ...]:
|
||||
"""Message groups least valuable first: middle prompt turns, extra choices, then the opener, the newest turn
|
||||
and the first choice."""
|
||||
inputs: Final = sorted(idx for family, idx in groups if family == _INPUT_MESSAGES)
|
||||
outputs: Final = sorted(idx for family, idx in groups if family == _OUTPUT_MESSAGES)
|
||||
pinned_inputs: Final = tuple(dict.fromkeys((*inputs[:1], *inputs[-1:])))
|
||||
return (
|
||||
*((_INPUT_MESSAGES, idx) for idx in inputs[1:-1]),
|
||||
*((_OUTPUT_MESSAGES, idx) for idx in reversed(outputs[1:])),
|
||||
*((_INPUT_MESSAGES, idx) for idx in pinned_inputs),
|
||||
*((_OUTPUT_MESSAGES, idx) for idx in outputs[:1]),
|
||||
)
|
||||
|
||||
|
||||
def fit_indexed_messages(attrs: Mapping[str, AttrValue], budget: int | None) -> Mapping[str, AttrValue]:
|
||||
"""``attrs`` with whole per-index messages shed, least valuable first, until at most ``budget`` keys remain.
|
||||
|
||||
``None`` means the span has no attribute count limit. Every message still rides the ``input.value`` and
|
||||
``output.value`` blobs, so shedding a per-index pair loses no content.
|
||||
"""
|
||||
if budget is None or len(attrs) <= budget:
|
||||
return attrs
|
||||
groups: Final = _message_key_groups(attrs)
|
||||
order: Final = _shed_order(groups)
|
||||
running: Final = tuple(accumulate(len(groups[group]) for group in order))
|
||||
excess: Final = len(attrs) - budget
|
||||
shed_count: Final = next((n + 1 for n, total in enumerate(running) if total >= excess), len(order))
|
||||
shed: Final = frozenset(chain.from_iterable(groups[group] for group in order[:shed_count]))
|
||||
return MappingProxyType({key: value for key, value in attrs.items() if key not in shed})
|
||||
|
||||
|
||||
class OpenInferenceMapper:
|
||||
|
|
@ -87,42 +134,22 @@ class OpenInferenceMapper:
|
|||
return {}
|
||||
|
||||
def _llm_call(self, data: LLMCallSpanData) -> AttributeMap:
|
||||
outputs: Final = output_messages(data)
|
||||
indexed_in, indexed_out = self._indexed_split(len(data.messages_in), len(outputs))
|
||||
return {
|
||||
**collect(self._LLM_CALL_ATTRS, data),
|
||||
**collect(self._BLOB_ATTRS, data),
|
||||
**self._messages(
|
||||
"llm.input_messages",
|
||||
"input.value",
|
||||
data.messages_in,
|
||||
self._prompt_positions(len(data.messages_in), indexed_in),
|
||||
),
|
||||
**self._messages("llm.output_messages", "output.value", outputs, range(indexed_out)),
|
||||
**self._messages(_INPUT_MESSAGES, "input.value", data.messages_in),
|
||||
**self._messages(_OUTPUT_MESSAGES, "output.value", output_messages(data)),
|
||||
**self._tools(data),
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def _indexed_split(inputs: int, outputs: int) -> tuple[int, int]:
|
||||
"""Prompt and response share one allowance; the response is reserved at least half of it."""
|
||||
indexed_out: Final = min(outputs, max(_MAX_INDEXED_MESSAGES // 2, _MAX_INDEXED_MESSAGES - inputs))
|
||||
return _MAX_INDEXED_MESSAGES - indexed_out, indexed_out
|
||||
|
||||
@staticmethod
|
||||
def _prompt_positions(total: int, indexed: int) -> tuple[int, ...]:
|
||||
"""Prompt messages that get per-index attributes: message 0 and the most recent turns."""
|
||||
if total <= indexed:
|
||||
return tuple(range(total))
|
||||
return (0, *range(total - indexed + 1, total))
|
||||
|
||||
@staticmethod
|
||||
def _messages(prefix: str, value_key: str, messages: Sequence[object], positions: Sequence[int]) -> AttributeMap:
|
||||
"""``{prefix}.{idx}.message.*`` keys for the messages at ``positions`` + the ``value_key`` blob of all."""
|
||||
def _messages(prefix: str, value_key: str, messages: Sequence[object]) -> AttributeMap:
|
||||
"""``{prefix}.{idx}.message.*`` keys for every message + the ``value_key`` blob of all of them."""
|
||||
parsed: Final = [(m.get("role") if isinstance(m, dict) else None, message_content(m)) for m in messages]
|
||||
attrs: Final = drop_none(
|
||||
{
|
||||
key: value
|
||||
for idx, (role, content) in ((idx, parsed[idx]) for idx in positions)
|
||||
for idx, (role, content) in enumerate(parsed)
|
||||
for key, value in (
|
||||
(f"{prefix}.{idx}.message.role", role if isinstance(role, str) else None),
|
||||
(f"{prefix}.{idx}.message.content", content),
|
||||
|
|
|
|||
|
|
@ -6,7 +6,7 @@ they live in one place.
|
|||
"""
|
||||
|
||||
import json
|
||||
from collections.abc import Callable, Mapping, Sequence
|
||||
from collections.abc import Callable, Iterable, Mapping, Sequence
|
||||
from typing import Final
|
||||
|
||||
from litellm.integrations.otel.mappers.base import AttributeMap, AttrValue
|
||||
|
|
@ -32,14 +32,6 @@ core telemetry no matter how many vocabularies are configured.
|
|||
"""
|
||||
|
||||
|
||||
MAX_MESSAGE_ATTRS_PER_SPAN: Final = DEFAULT_SPAN_ATTRIBUTE_LIMIT // 8
|
||||
"""Span-wide ceiling on per-index chat message attributes, prompt and response together.
|
||||
|
||||
An eighth is the largest share that still fits beside the tool ceiling and the core
|
||||
of every vocabulary at once. The complete conversation still rides the JSON blobs.
|
||||
"""
|
||||
|
||||
|
||||
def tool_attr_budget(vocabularies: int) -> int:
|
||||
"""Split the span-wide tool-definition ceiling across active vocabularies."""
|
||||
return MAX_TOOL_DEFINITION_ATTRS_PER_SPAN // max(vocabularies, 1)
|
||||
|
|
@ -47,7 +39,12 @@ def tool_attr_budget(vocabularies: int) -> int:
|
|||
|
||||
def drop_none(values: Mapping[str, AttrValue | None]) -> AttributeMap:
|
||||
"""Return ``values`` with ``None``-valued entries removed."""
|
||||
return {k: v for k, v in values.items() if v is not None}
|
||||
return drop_none_pairs(values.items())
|
||||
|
||||
|
||||
def drop_none_pairs(pairs: Iterable[tuple[str, AttrValue | None]]) -> AttributeMap:
|
||||
"""Return ``pairs`` as a map with ``None``-valued entries removed."""
|
||||
return {k: v for k, v in pairs if v is not None}
|
||||
|
||||
|
||||
def tool_definition_attrs(
|
||||
|
|
|
|||
|
|
@ -43,12 +43,12 @@ from typing import TYPE_CHECKING, Any, Final, cast
|
|||
|
||||
from litellm.constants import LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL
|
||||
from litellm.integrations.otel.model.semconv import resolve_operation
|
||||
from litellm.integrations.otel.model.utils import as_str, to_seconds
|
||||
from litellm.integrations.otel.model.trace_controls import TraceControls, caller_trace_controls
|
||||
from litellm.integrations.otel.model.utils import as_str, as_str_mapping, to_seconds
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.types.utils import StandardLoggingPayload
|
||||
|
||||
LANGFUSE_TRACE_NAME_HEADER: Final = "langfuse_trace_name"
|
||||
REQUESTER_METADATA_KEY: Final = "requester_metadata"
|
||||
REQUESTER_METADATA_PATH: Final = f"{REQUESTER_METADATA_KEY}."
|
||||
|
||||
|
|
@ -225,7 +225,7 @@ class LLMCallEvent:
|
|||
# needs to be reasonable for a span that never gets closed (a leak).
|
||||
provisional_span_name: str
|
||||
time_to_first_chunk_seconds: float | None
|
||||
trace_name: str | None
|
||||
trace: TraceControls
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, kwargs: Mapping[str, Any]) -> LLMCallEvent:
|
||||
|
|
@ -242,30 +242,10 @@ class LLMCallEvent:
|
|||
upstream_started=kwargs.get("api_call_start_time") is not None,
|
||||
provisional_span_name=f"{operation.value} {model}".strip(),
|
||||
time_to_first_chunk_seconds=time_to_first_chunk_seconds(kwargs),
|
||||
trace_name=caller_trace_name(kwargs),
|
||||
trace=caller_trace_controls(kwargs),
|
||||
)
|
||||
|
||||
|
||||
def caller_trace_name(kwargs: Mapping[str, object]) -> str | None:
|
||||
request: Final = _as_str_mapping(kwargs.get("litellm_params"))
|
||||
if request is None:
|
||||
return None
|
||||
proxy_request: Final = _as_str_mapping(request.get("proxy_server_request"))
|
||||
headers: Final = _as_str_mapping(proxy_request.get("headers")) if proxy_request is not None else None
|
||||
from_header: Final = as_str(headers.get(LANGFUSE_TRACE_NAME_HEADER)) if headers is not None else None
|
||||
if from_header:
|
||||
return from_header
|
||||
return next(
|
||||
(
|
||||
name
|
||||
for key in ("metadata", "litellm_metadata")
|
||||
if (metadata := _as_str_mapping(request.get(key))) is not None
|
||||
and (name := as_str(metadata.get("trace_name")))
|
||||
),
|
||||
None,
|
||||
)
|
||||
|
||||
|
||||
def time_to_first_chunk_seconds(kwargs: Mapping[str, Any]) -> float | None:
|
||||
"""Seconds from the upstream request being issued (``api_call_start_time``)
|
||||
to the first streamed chunk (``completion_start_time``); ``None`` for
|
||||
|
|
@ -300,15 +280,8 @@ def auth_metadata(payload: StandardLoggingPayload | None, kwargs: Mapping[str, o
|
|||
)
|
||||
|
||||
|
||||
def _as_str_mapping(value: object) -> Mapping[str, object] | None:
|
||||
"""A read-only view of ``value`` when it is a mapping, else ``None``."""
|
||||
if not isinstance(value, Mapping):
|
||||
return None
|
||||
return cast("Mapping[str, object]", value) # cast-ok: isinstance-guarded, JSON metadata has str keys
|
||||
|
||||
|
||||
def _string_entries(value: object) -> Mapping[str, str] | None:
|
||||
entries: Final = _as_str_mapping(value)
|
||||
entries: Final = as_str_mapping(value)
|
||||
if entries is None:
|
||||
return None
|
||||
typed: Final = MappingProxyType({key: item for key, item in entries.items() if isinstance(item, str)})
|
||||
|
|
@ -324,18 +297,18 @@ def _metadata_dicts(
|
|||
litellm copies it onto ``metadata``, but both are yielded so a route that
|
||||
populates only one is still covered.
|
||||
"""
|
||||
payload_view: Final = _as_str_mapping(payload)
|
||||
payload_view: Final = as_str_mapping(payload)
|
||||
if payload_view is not None:
|
||||
payload_metadata: Final = _as_str_mapping(payload_view.get("metadata"))
|
||||
payload_metadata: Final = as_str_mapping(payload_view.get("metadata"))
|
||||
if payload_metadata is not None:
|
||||
yield payload_metadata
|
||||
params: Final = _as_str_mapping(kwargs.get("litellm_params"))
|
||||
params: Final = as_str_mapping(kwargs.get("litellm_params"))
|
||||
if params is None:
|
||||
return
|
||||
yield from (
|
||||
metadata
|
||||
for key in ("metadata", "litellm_metadata")
|
||||
if (metadata := _as_str_mapping(params.get(key))) is not None
|
||||
if (metadata := as_str_mapping(params.get(key))) is not None
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -365,14 +338,14 @@ def metadata_from_request_data(data: object) -> Mapping[str, object] | None:
|
|||
The proxy stores it under ``metadata`` or ``litellm_metadata`` depending on the route;
|
||||
the proxy-owned siblings (``user_api_key_*``, ``requester_ip_address``) are not read.
|
||||
"""
|
||||
top: Final = _as_str_mapping(data)
|
||||
top: Final = as_str_mapping(data)
|
||||
if top is None:
|
||||
return None
|
||||
snapshots: Final = tuple(
|
||||
snapshot
|
||||
for name in ("metadata", "litellm_metadata")
|
||||
if (nested := _as_str_mapping(top.get(name))) is not None
|
||||
and (snapshot := _as_str_mapping(nested.get(REQUESTER_METADATA_KEY))) is not None
|
||||
if (nested := as_str_mapping(top.get(name))) is not None
|
||||
and (snapshot := as_str_mapping(nested.get(REQUESTER_METADATA_KEY))) is not None
|
||||
)
|
||||
return MappingProxyType({REQUESTER_METADATA_KEY: snapshots[0]}) if snapshots else None
|
||||
|
||||
|
|
@ -382,7 +355,7 @@ def flatten_metadata(raw: Mapping[str, object]) -> Iterator[tuple[str, str]]:
|
|||
stack: Final = list(tuple(raw.items())[::-1]) # mutable-ok: iterative worklist keeps the walk off the call stack
|
||||
while stack:
|
||||
key, value = stack.pop()
|
||||
if (nested := _as_str_mapping(value)) is not None:
|
||||
if (nested := as_str_mapping(value)) is not None:
|
||||
stack.extend(tuple((f"{key}.{sub_key}", sub_value) for sub_key, sub_value in nested.items())[::-1])
|
||||
elif isinstance(value, (str, bool, int, float)):
|
||||
yield key, str(value)
|
||||
|
|
|
|||
|
|
@ -10,10 +10,7 @@ from types import MappingProxyType
|
|||
from typing import TYPE_CHECKING, ClassVar, Final, cast
|
||||
from urllib.parse import urlsplit
|
||||
|
||||
from litellm.integrations.otel.model.metadata import (
|
||||
RequestContext,
|
||||
RequestIdentity,
|
||||
)
|
||||
from litellm.integrations.otel.model.metadata import RequestContext, RequestIdentity
|
||||
from litellm.integrations.otel.model.semconv import (
|
||||
GenAIOperation,
|
||||
GenAIOutputType,
|
||||
|
|
@ -22,6 +19,7 @@ from litellm.integrations.otel.model.semconv import (
|
|||
resolve_output_type,
|
||||
resolve_provider,
|
||||
)
|
||||
from litellm.integrations.otel.model.trace_controls import TraceControls
|
||||
from litellm.integrations.otel.model.utils import (
|
||||
as_bool,
|
||||
as_float,
|
||||
|
|
@ -387,7 +385,7 @@ class LLMCallSpanData:
|
|||
output_type: GenAIOutputType | None = None
|
||||
call_type: str | None = None
|
||||
request_route: str | None = None
|
||||
trace_name: str | None = None
|
||||
trace: TraceControls = field(default_factory=TraceControls)
|
||||
|
||||
@classmethod
|
||||
def from_standard_logging_payload(
|
||||
|
|
@ -396,7 +394,7 @@ class LLMCallSpanData:
|
|||
capture_content: bool = False,
|
||||
time_to_first_chunk_seconds: float | None = None,
|
||||
request_route: str | None = None,
|
||||
trace_name: str | None = None,
|
||||
trace: TraceControls | None = None,
|
||||
) -> LLMCallSpanData:
|
||||
params: Final = cast(Mapping[str, object], payload.get("model_parameters") or {})
|
||||
# The single parse of the request's metadata — the request-vs-provider
|
||||
|
|
@ -438,7 +436,7 @@ class LLMCallSpanData:
|
|||
output_type=resolve_output_type(call_type),
|
||||
call_type=call_type or None,
|
||||
request_route=request_route or context.identity.request_route,
|
||||
trace_name=trace_name,
|
||||
trace=trace or TraceControls(),
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
61
litellm/integrations/otel/model/trace_controls.py
Normal file
61
litellm/integrations/otel/model/trace_controls.py
Normal file
|
|
@ -0,0 +1,61 @@
|
|||
"""The caller's Langfuse trace controls, parsed from the live callback kwargs."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Mapping
|
||||
from dataclasses import dataclass
|
||||
from typing import Final
|
||||
|
||||
from pydantic import TypeAdapter, ValidationError
|
||||
|
||||
from litellm.integrations.otel.model.utils import as_str, as_str_mapping
|
||||
|
||||
LANGFUSE_HEADER_PREFIX: Final = "langfuse_"
|
||||
_ITEMS: Final = TypeAdapter(tuple[object, ...])
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class TraceControls:
|
||||
"""The caller's trace-level Langfuse controls: ``metadata.trace_name`` / ``trace_user_id`` / ``session_id`` /
|
||||
``tags`` on the request (SDK or proxy body), with the proxy's ``langfuse_<control>`` headers winning over the
|
||||
body for the scalar ones. Mutation controls (``trace_id``, ``existing_trace_id``, ``update_trace_keys``) are
|
||||
deliberately not carried."""
|
||||
|
||||
name: str | None = None
|
||||
user_id: str | None = None
|
||||
session_id: str | None = None
|
||||
tags: tuple[str, ...] = ()
|
||||
|
||||
|
||||
def caller_trace_controls(kwargs: Mapping[str, object]) -> TraceControls:
|
||||
request: Final = as_str_mapping(kwargs.get("litellm_params"))
|
||||
if request is None:
|
||||
return TraceControls()
|
||||
proxy_request: Final = as_str_mapping(request.get("proxy_server_request"))
|
||||
headers: Final = as_str_mapping(proxy_request.get("headers")) if proxy_request is not None else None
|
||||
bodies: Final = tuple(
|
||||
metadata
|
||||
for key in ("metadata", "litellm_metadata")
|
||||
if (metadata := as_str_mapping(request.get(key))) is not None
|
||||
)
|
||||
|
||||
def scalar(control: str) -> str | None:
|
||||
from_header: Final = as_str(headers.get(f"{LANGFUSE_HEADER_PREFIX}{control}")) if headers is not None else None
|
||||
if from_header:
|
||||
return from_header
|
||||
return next((value for body in bodies if (value := as_str(body.get(control)))), None)
|
||||
|
||||
return TraceControls(
|
||||
name=scalar("trace_name"),
|
||||
user_id=scalar("trace_user_id"),
|
||||
session_id=scalar("session_id"),
|
||||
tags=next((tags for body in bodies if (tags := _str_items(body.get("tags")))), ()),
|
||||
)
|
||||
|
||||
|
||||
def _str_items(value: object) -> tuple[str, ...]:
|
||||
try:
|
||||
items: Final = _ITEMS.validate_python(value)
|
||||
except ValidationError:
|
||||
return ()
|
||||
return tuple(item for item in items if isinstance(item, str) and item)
|
||||
|
|
@ -8,7 +8,13 @@ parsing lives in :mod:`litellm.integrations.otel.plumbing.providers` instead,
|
|||
because it delegates to the OTel SDK's own W3C Baggage parser.
|
||||
"""
|
||||
|
||||
from collections.abc import Mapping
|
||||
from datetime import datetime
|
||||
from typing import Final
|
||||
|
||||
from pydantic import TypeAdapter, ValidationError
|
||||
|
||||
_STR_MAPPING: Final = TypeAdapter(Mapping[str, object])
|
||||
|
||||
|
||||
def as_str(value: object) -> str | None:
|
||||
|
|
@ -55,6 +61,13 @@ def as_bool(value: object) -> bool | None:
|
|||
return bool(value)
|
||||
|
||||
|
||||
def as_str_mapping(value: object) -> Mapping[str, object] | None:
|
||||
try:
|
||||
return _STR_MAPPING.validate_python(value)
|
||||
except ValidationError:
|
||||
return None
|
||||
|
||||
|
||||
def as_str_tuple(value: object) -> tuple[str, ...] | None:
|
||||
if value is None:
|
||||
return None
|
||||
|
|
|
|||
|
|
@ -26,6 +26,7 @@ if TYPE_CHECKING:
|
|||
from litellm.integrations.otel.model.destination import OtelDestination
|
||||
|
||||
_PROPAGATOR: Final = TraceContextTextMapPropagator()
|
||||
_W3C_TRACE_HEADERS: Final = frozenset(("traceparent", "tracestate"))
|
||||
|
||||
# The request's root span — the FastAPI-owned SERVER span — captured ONCE when the
|
||||
# proxy first resolves it, so request-level spans (the LLM call, guardrails) can
|
||||
|
|
@ -310,6 +311,37 @@ def extract_traceparent(headers: Mapping[str, str]) -> Context | None:
|
|||
return _PROPAGATOR.extract(carrier)
|
||||
|
||||
|
||||
def _outgoing_trace_context(parent_span: object) -> Context | None:
|
||||
if isinstance(parent_span, Span) and is_recordable_span(parent_span):
|
||||
return context_from_span(parent_span)
|
||||
|
||||
root: Final = request_root_span()
|
||||
if root is not None:
|
||||
return context_from_span(root)
|
||||
|
||||
current: Final = get_current()
|
||||
if is_recordable_span(get_current_span(current)):
|
||||
return current
|
||||
return None
|
||||
|
||||
|
||||
def inject_trace_context(headers: Mapping[str, str], parent_span: object = None) -> dict[str, str]:
|
||||
"""``headers`` plus W3C ``traceparent``/``tracestate`` for this request's span.
|
||||
|
||||
Parent preference: ``parent_span`` (the request span auth stashed on the key), then
|
||||
the anchored request root span, then the ambient active span. Only trace context is
|
||||
injected, never Baggage. Unchanged when no valid span exists anywhere.
|
||||
"""
|
||||
context: Final = _outgoing_trace_context(parent_span)
|
||||
if context is None:
|
||||
return dict(headers) # mutable-ok: OpenTelemetry propagator requires a mutable carrier
|
||||
carrier: Final = { # mutable-ok: OpenTelemetry propagator requires a mutable carrier
|
||||
key: value for key, value in headers.items() if key.lower() not in _W3C_TRACE_HEADERS
|
||||
}
|
||||
_PROPAGATOR.inject(carrier, context=context)
|
||||
return carrier
|
||||
|
||||
|
||||
# The OTLP destinations this request's key or team pointed its traces at, resolved
|
||||
# once during auth. A ``ContextVar`` for the same reason the root span above is one:
|
||||
# it rides the request task's context into the ``asyncio.create_task`` children that
|
||||
|
|
|
|||
|
|
@ -160,7 +160,7 @@ def get_llm_provider(
|
|||
if model is None:
|
||||
raise ValueError("model parameter is required but was None. Please provide a valid model name.")
|
||||
|
||||
if litellm.LiteLLMProxyChatConfig._should_use_litellm_proxy_by_default(
|
||||
if litellm.LiteLLMProxyChatConfig.should_use_litellm_proxy_by_default(
|
||||
litellm_params=cast(LiteLLM_Params | None, litellm_params)
|
||||
):
|
||||
return litellm.LiteLLMProxyChatConfig.litellm_proxy_get_custom_llm_provider_info(
|
||||
|
|
|
|||
|
|
@ -7,7 +7,7 @@ from collections.abc import Iterator, Mapping, Sequence
|
|||
from datetime import datetime
|
||||
from typing import TYPE_CHECKING, Any, Final
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm._logging import format_base64_size, verbose_logger
|
||||
from litellm.constants import (
|
||||
BASE64_TRUNCATION_OFFLOAD_THRESHOLD_CHARS,
|
||||
MAX_BASE64_LENGTH_FOR_LOGGING,
|
||||
|
|
@ -40,9 +40,6 @@ import litellm
|
|||
Helper utils used for logging callbacks
|
||||
"""
|
||||
|
||||
_BYTES_PER_KIB: Final = 1024
|
||||
_BYTES_PER_MIB: Final = 1024 * 1024
|
||||
|
||||
# Regex matching data-URI base64 content: "data:<mime>;base64,<payload>"
|
||||
# Captures: group(1)=mime_type, group(2)=base64_payload
|
||||
_DATA_URI_RE: Final = re.compile(r"data:([^;]+);base64,([A-Za-z0-9+/=]+)")
|
||||
|
|
@ -52,23 +49,13 @@ _DATA_URI_RE: Final = re.compile(r"data:([^;]+);base64,([A-Za-z0-9+/=]+)")
|
|||
_MAX_TRUNCATION_DEPTH: Final = 20
|
||||
|
||||
|
||||
def _format_base64_size(num_chars: int) -> str:
|
||||
"""Return a human-readable byte-size estimate from a base64 character count."""
|
||||
num_bytes: Final = num_chars * 3 / 4
|
||||
if num_bytes >= _BYTES_PER_MIB:
|
||||
return f"{num_bytes / _BYTES_PER_MIB:.2f}MB"
|
||||
if num_bytes >= _BYTES_PER_KIB:
|
||||
return f"{num_bytes / _BYTES_PER_KIB:.1f}KB"
|
||||
return f"{int(num_bytes)}B"
|
||||
|
||||
|
||||
def _base64_data_uri_replacer(match: re.Match) -> str:
|
||||
"""Replace a single base64 data-URI match with a size placeholder if too long."""
|
||||
mime_type: Final = match.group(1)
|
||||
payload: Final = match.group(2)
|
||||
if len(payload) <= MAX_BASE64_LENGTH_FOR_LOGGING:
|
||||
return match.group(0)
|
||||
size_str: Final = _format_base64_size(len(payload))
|
||||
size_str: Final = format_base64_size(len(payload))
|
||||
return f"data:{mime_type};base64,[base64_data truncated: {size_str}]"
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -6,25 +6,29 @@ from pydantic import BaseModel
|
|||
|
||||
from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH
|
||||
|
||||
UNSERIALIZABLE_OBJECT: Final = "Unserializable Object"
|
||||
|
||||
|
||||
def strip_null_bytes(value: str) -> str:
|
||||
"""Strip NUL bytes, which PostgreSQL text/jsonb columns reject (error 22P05)."""
|
||||
return value.replace("\x00", "")
|
||||
|
||||
|
||||
def safe_dumps(
|
||||
data: Any,
|
||||
def safe_json_structure(
|
||||
data: object,
|
||||
max_depth: int = DEFAULT_MAX_RECURSE_DEPTH,
|
||||
value_transform: Callable[[str | None, str], str] | None = None,
|
||||
) -> str:
|
||||
key: str | None = None,
|
||||
) -> object:
|
||||
"""
|
||||
Recursively serialize data while detecting circular references.
|
||||
Rebuild data out of JSON-native pieces while detecting circular references.
|
||||
If a circular reference is detected then a marker string is returned.
|
||||
NUL bytes are stripped from strings to prevent PostgreSQL 22P05 errors.
|
||||
|
||||
value_transform, when given, is applied to every string leaf (and to the
|
||||
str() fallback for non-serializable objects) with the mapping key the leaf
|
||||
was reached under, so callers can rewrite values without touching structure.
|
||||
key is the mapping key data itself was reached under, when the caller has one.
|
||||
"""
|
||||
|
||||
def _transform(key: str | None, value: str) -> str:
|
||||
|
|
@ -75,7 +79,15 @@ def safe_dumps(
|
|||
try:
|
||||
return _transform(key, strip_null_bytes(str(obj)))
|
||||
except Exception:
|
||||
return "Unserializable Object"
|
||||
return UNSERIALIZABLE_OBJECT
|
||||
|
||||
safe_data: Final = _serialize(data, set(), 0)
|
||||
return json.dumps(safe_data, default=str)
|
||||
return _serialize(data, set(), 0, key)
|
||||
|
||||
|
||||
def safe_dumps(
|
||||
data: Any,
|
||||
max_depth: int = DEFAULT_MAX_RECURSE_DEPTH,
|
||||
value_transform: Callable[[str | None, str], str] | None = None,
|
||||
) -> str:
|
||||
"""Serialize data to JSON text through safe_json_structure."""
|
||||
return json.dumps(safe_json_structure(data, max_depth, value_transform), default=str)
|
||||
|
|
|
|||
|
|
@ -148,6 +148,7 @@ class _ToolCallChunk(TypedDict):
|
|||
class _UsageBearingChunk(TypedDict, total=False):
|
||||
usage: Usage | None
|
||||
_hidden_params: Mapping[str, str]
|
||||
choices: ReadOnly[Sequence[StreamingChoices | Mapping[str, object]]]
|
||||
|
||||
|
||||
class _UsageSummary(TypedDict):
|
||||
|
|
@ -921,21 +922,22 @@ class ChunkProcessor:
|
|||
|
||||
prompt_tokens_details = attach_cache_creation_token_details(prompt_tokens_details, cache_creation_token_details)
|
||||
|
||||
completion_tokens = self._reset_anthropic_cursor_completion_tokens(
|
||||
recovered_completion_tokens: Final = self._reset_anthropic_cursor_completion_tokens(
|
||||
chunks=chunks,
|
||||
completion_tokens=completion_tokens,
|
||||
completion_usage_updates=completion_usage_updates,
|
||||
)
|
||||
cursor_was_reset: Final = recovered_completion_tokens != completion_tokens
|
||||
|
||||
return UsagePerChunk(
|
||||
prompt_tokens=prompt_tokens,
|
||||
completion_tokens=completion_tokens,
|
||||
completion_tokens=recovered_completion_tokens,
|
||||
cache_creation_input_tokens=cache_creation_input_tokens,
|
||||
cache_read_input_tokens=cache_read_input_tokens,
|
||||
server_tool_use=server_tool_use,
|
||||
web_search_requests=web_search_requests,
|
||||
google_maps_grounding_requests=google_maps_grounding_requests,
|
||||
completion_tokens_details=completion_tokens_details,
|
||||
completion_tokens_details=None if cursor_was_reset else completion_tokens_details,
|
||||
prompt_tokens_details=prompt_tokens_details,
|
||||
cost=cost,
|
||||
inference_geo=self._last_provider_pricing_field(chunks, "inference_geo"),
|
||||
|
|
@ -960,6 +962,30 @@ class ChunkProcessor:
|
|||
]
|
||||
return values[-1] if values else None
|
||||
|
||||
@staticmethod
|
||||
def _finish_reason_of_choice(choice: object) -> str | None:
|
||||
match choice:
|
||||
case StreamingChoices(finish_reason=reason) | Choices(finish_reason=reason):
|
||||
return reason
|
||||
case {"finish_reason": str() as reason}:
|
||||
return reason
|
||||
case _:
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def _chunk_choices(chunk: "_UsageBearingChunk | ModelResponse | ModelResponseStream") -> Sequence[object]:
|
||||
if isinstance(chunk, dict):
|
||||
return chunk.get("choices", ())
|
||||
return getattr(chunk, "choices", ())
|
||||
|
||||
@staticmethod
|
||||
def _saw_finish_reason(chunks: Sequence["_UsageBearingChunk | ModelResponse"]) -> bool:
|
||||
return any(
|
||||
ChunkProcessor._finish_reason_of_choice(choice) is not None
|
||||
for chunk in chunks
|
||||
for choice in ChunkProcessor._chunk_choices(chunk)
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _reset_anthropic_cursor_completion_tokens(
|
||||
chunks: Sequence["_UsageBearingChunk | ModelResponse"],
|
||||
|
|
@ -970,18 +996,18 @@ class ChunkProcessor:
|
|||
|
||||
See the ``completion_usage_updates`` comment in
|
||||
``_calculate_usage_per_chunk``. The accumulated value is NOT a stale
|
||||
cursor when either it is > 1 (definitely not a placeholder) or we saw
|
||||
>= 2 completion-bearing usage events (positive evidence ``message_delta``
|
||||
arrived). Otherwise — the only completion update we ever saw was the
|
||||
Anthropic ``message_start`` cursor (=1) — reset to 0 so
|
||||
``calculate_usage()``'s ``or token_counter(text=...)`` fallback estimates
|
||||
from the actually-received completion text instead of trusting the
|
||||
placeholder. Gated on ``custom_llm_provider == "anthropic"`` so the
|
||||
heuristic (which encodes Anthropic's specific message_start SSE shape)
|
||||
does not silently affect other providers that may legitimately report
|
||||
``completion_tokens=1`` from a single usage event.
|
||||
cursor when we saw >= 2 completion-bearing usage events or any chunk
|
||||
carried a ``finish_reason`` (positive evidence ``message_delta``
|
||||
arrived). Otherwise the only completion update we ever saw was the
|
||||
Anthropic ``message_start`` cursor, a small placeholder whose magnitude
|
||||
varies per request (1 and 8 both observed live), so reset to 0 and let
|
||||
``calculate_usage()``'s ``or token_counter(...)`` fallback estimate from
|
||||
the actually-received text and reasoning instead. Gated on
|
||||
``custom_llm_provider == "anthropic"`` so the heuristic (which encodes
|
||||
Anthropic's specific message_start SSE shape) does not silently affect
|
||||
other providers that legitimately report usage from a single event.
|
||||
"""
|
||||
saw_non_cursor_completion: Final = completion_tokens > 1 or completion_usage_updates >= 2
|
||||
saw_non_cursor_completion: Final = completion_usage_updates >= 2 or ChunkProcessor._saw_finish_reason(chunks)
|
||||
if saw_non_cursor_completion:
|
||||
return completion_tokens
|
||||
|
||||
|
|
@ -995,7 +1021,7 @@ class ChunkProcessor:
|
|||
if isinstance(hp, dict):
|
||||
custom_llm_provider = hp.get("custom_llm_provider")
|
||||
|
||||
if custom_llm_provider == "anthropic" and completion_tokens == 1:
|
||||
if custom_llm_provider == "anthropic":
|
||||
return 0
|
||||
return completion_tokens
|
||||
|
||||
|
|
@ -1039,10 +1065,13 @@ class ChunkProcessor:
|
|||
returned_usage.prompt_tokens = 0
|
||||
returned_usage.completion_tokens = (
|
||||
completion_tokens
|
||||
or token_counter(
|
||||
model=model,
|
||||
text=completion_output,
|
||||
count_response_tokens=True, # count_response_tokens is a Flag to tell token counter this is a response, No need to add extra tokens we do for input messages
|
||||
or (
|
||||
token_counter(
|
||||
model=model,
|
||||
text=completion_output,
|
||||
count_response_tokens=True, # count_response_tokens is a Flag to tell token counter this is a response, No need to add extra tokens we do for input messages
|
||||
)
|
||||
+ (reasoning_tokens or 0)
|
||||
)
|
||||
)
|
||||
returned_usage.total_tokens = returned_usage.prompt_tokens + returned_usage.completion_tokens
|
||||
|
|
@ -1066,15 +1095,16 @@ class ChunkProcessor:
|
|||
returned_usage.completion_tokens_details = completion_tokens_details
|
||||
|
||||
if reasoning_tokens is not None:
|
||||
capped_reasoning_tokens: Final = min(max(0, reasoning_tokens), returned_usage.completion_tokens)
|
||||
if returned_usage.completion_tokens_details is None:
|
||||
returned_usage.completion_tokens_details = CompletionTokensDetailsWrapper(
|
||||
reasoning_tokens=reasoning_tokens
|
||||
reasoning_tokens=capped_reasoning_tokens,
|
||||
text_tokens=returned_usage.completion_tokens - capped_reasoning_tokens,
|
||||
)
|
||||
elif (
|
||||
returned_usage.completion_tokens_details is not None
|
||||
and returned_usage.completion_tokens_details.reasoning_tokens is None
|
||||
):
|
||||
capped_reasoning_tokens: Final = min(max(0, reasoning_tokens), returned_usage.completion_tokens)
|
||||
returned_usage.completion_tokens_details.reasoning_tokens = capped_reasoning_tokens
|
||||
if returned_usage.completion_tokens_details.text_tokens is None:
|
||||
returned_usage.completion_tokens_details.text_tokens = (
|
||||
|
|
|
|||
|
|
@ -31,6 +31,7 @@ from litellm.llms.anthropic.experimental_pass_through.adapters.transformation im
|
|||
)
|
||||
from litellm.llms.base_llm.guardrail_translation.base_translation import (
|
||||
BaseTranslation,
|
||||
RequestScanContext,
|
||||
StreamingScanKey,
|
||||
StreamTransformSink,
|
||||
)
|
||||
|
|
@ -527,6 +528,26 @@ class AnthropicMessagesHandler(BaseTranslation):
|
|||
)
|
||||
return result if result else None
|
||||
|
||||
def request_scan_context(
|
||||
self, data: Mapping[str, object], guardrail_to_apply: "CustomGuardrail"
|
||||
) -> RequestScanContext:
|
||||
if data.get("messages") is None:
|
||||
return RequestScanContext()
|
||||
translated: Final = self._translate_to_openai(
|
||||
{key: value for key, value in data.items() if key != "system"} # mutable-ok: API message payload
|
||||
)
|
||||
hoisted_system_message: Final = (
|
||||
None
|
||||
if effective_skip_system_message_for_guardrail(guardrail_to_apply)
|
||||
else self._hoisted_top_level_system_message(data)
|
||||
)
|
||||
return RequestScanContext.scoped(
|
||||
(*(() if hoisted_system_message is None else (hoisted_system_message,)), *translated["messages"]),
|
||||
tuple(tool for tool in translated.get("tools") or () if not is_provider_native_tool_dict(tool)),
|
||||
guardrail_to_apply,
|
||||
skip_system=False,
|
||||
)
|
||||
|
||||
async def process_input_messages(
|
||||
self,
|
||||
data: dict,
|
||||
|
|
@ -696,9 +717,7 @@ class AnthropicMessagesHandler(BaseTranslation):
|
|||
|
||||
return data
|
||||
|
||||
def _hoisted_top_level_system_message(
|
||||
self, data: dict
|
||||
) -> AllMessageValues | None: # mutable-ok: API message payload
|
||||
def _hoisted_top_level_system_message(self, data: Mapping[str, object]) -> AllMessageValues | None:
|
||||
"""Return the system message produced by translating the top-level prompt."""
|
||||
system: Final = data.get("system")
|
||||
if not system:
|
||||
|
|
@ -1200,7 +1219,7 @@ class AnthropicMessagesHandler(BaseTranslation):
|
|||
)
|
||||
|
||||
guardrailed_inputs: Final = await guardrail_to_apply.apply_guardrail(
|
||||
inputs=inputs,
|
||||
inputs=self.with_response_context(inputs, request_data, guardrail_to_apply),
|
||||
request_data=request_data,
|
||||
input_type="response",
|
||||
logging_obj=litellm_logging_obj,
|
||||
|
|
@ -1273,7 +1292,7 @@ class AnthropicMessagesHandler(BaseTranslation):
|
|||
key="response",
|
||||
)
|
||||
_guardrailed_inputs = await guardrail_to_apply.apply_guardrail(
|
||||
inputs=guardrail_inputs,
|
||||
inputs=self.with_response_context(guardrail_inputs, prepared_request_data, guardrail_to_apply),
|
||||
request_data=prepared_request_data,
|
||||
input_type="response",
|
||||
logging_obj=litellm_logging_obj,
|
||||
|
|
@ -1319,7 +1338,11 @@ class AnthropicMessagesHandler(BaseTranslation):
|
|||
key="responses",
|
||||
)
|
||||
_guardrailed_inputs = await guardrail_to_apply.apply_guardrail(
|
||||
inputs={"texts": [string_so_far]},
|
||||
inputs=self.with_response_context(
|
||||
GenericGuardrailAPIInputs(texts=[string_so_far]), # mutable-ok: guardrail inputs want a list
|
||||
prepared_request_data,
|
||||
guardrail_to_apply,
|
||||
),
|
||||
request_data=prepared_request_data,
|
||||
input_type="response",
|
||||
logging_obj=litellm_logging_obj,
|
||||
|
|
|
|||
|
|
@ -1180,7 +1180,7 @@ class LiteLLMAnthropicMessagesAdapter:
|
|||
self._add_system_message_to_messages(new_messages, anthropic_message_request)
|
||||
|
||||
new_kwargs: Final[ChatCompletionRequest] = {
|
||||
"model": anthropic_message_request["model"],
|
||||
"model": anthropic_message_request.get("model", ""),
|
||||
"messages": new_messages,
|
||||
}
|
||||
## CONVERT METADATA (user_id + litellm metadata)
|
||||
|
|
|
|||
|
|
@ -49,12 +49,7 @@ class AzureOpenAIResponsesAPIConfig(OpenAIResponsesAPIConfig):
|
|||
return BaseAzureLLM._base_validate_azure_environment(headers=headers, litellm_params=litellm_params)
|
||||
|
||||
def get_stripped_model_name(self, model: str) -> str:
|
||||
# if "responses/" is in the model name, remove it
|
||||
if "responses/" in model:
|
||||
model = model.replace("responses/", "")
|
||||
if "o_series" in model:
|
||||
model = model.replace("o_series/", "")
|
||||
return model
|
||||
return model.replace("responses/", "").replace("o_series/", "").replace("azure_ai/", "")
|
||||
|
||||
def _handle_reasoning_item(self, item: dict[str, Any]) -> dict[str, Any]:
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -9,6 +9,7 @@ from litellm.types.llms.openai import AllMessageValues
|
|||
from litellm.types.router import GenericLiteLLMParams
|
||||
|
||||
AzureAIApiKeyHeader = Literal["Authorization", "api-key", "Api-Key", "Ocp-Apim-Subscription-Key"]
|
||||
AZURE_OPENAI_V1_HOST_SUFFIXES: Final = (".services.ai.azure.com", ".openai.azure.com")
|
||||
|
||||
|
||||
def is_foundry_model_inference_base(api_base: str) -> bool:
|
||||
|
|
@ -19,11 +20,13 @@ def is_foundry_model_inference_base(api_base: str) -> bool:
|
|||
return "/openai/deployments" not in parsed.path
|
||||
|
||||
|
||||
def api_key_header_for_base(api_base: str | None) -> AzureAIApiKeyHeader:
|
||||
def is_azure_openai_v1_host(api_base: str | None) -> bool:
|
||||
host: Final = urlparse(api_base).hostname if api_base else None
|
||||
if host and (host.endswith(".services.ai.azure.com") or host.endswith(".openai.azure.com")):
|
||||
return "api-key"
|
||||
return "Authorization"
|
||||
return host is not None and host.endswith(AZURE_OPENAI_V1_HOST_SUFFIXES)
|
||||
|
||||
|
||||
def api_key_header_for_base(api_base: str | None) -> AzureAIApiKeyHeader:
|
||||
return "api-key" if is_azure_openai_v1_host(api_base) else "Authorization"
|
||||
|
||||
|
||||
def get_azure_ai_entra_token(litellm_params: Mapping[str, object] | None = None) -> str | None:
|
||||
|
|
@ -70,6 +73,17 @@ def get_azure_ai_auth_headers(
|
|||
AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY: Final = "azure_model_router_selected_model"
|
||||
|
||||
|
||||
def azure_ai_supports_native_responses(model: str | None, api_base: str | None) -> bool:
|
||||
resolved_base: Final = AzureFoundryModelInfo.get_api_base(api_base)
|
||||
if resolved_base is not None and not is_azure_openai_v1_host(resolved_base):
|
||||
return False
|
||||
if model is None:
|
||||
return True
|
||||
if "claude" in model.lower():
|
||||
return False
|
||||
return AzureFoundryModelInfo.get_azure_ai_route(model) == "default"
|
||||
|
||||
|
||||
class AzureFoundryModelInfo(BaseLLMModelInfo):
|
||||
"""Model info for Azure AI / Azure Foundry models."""
|
||||
|
||||
|
|
|
|||
0
litellm/llms/azure_ai/responses/__init__.py
Normal file
0
litellm/llms/azure_ai/responses/__init__.py
Normal file
53
litellm/llms/azure_ai/responses/transformation.py
Normal file
53
litellm/llms/azure_ai/responses/transformation.py
Normal file
|
|
@ -0,0 +1,53 @@
|
|||
from typing import Final
|
||||
|
||||
import httpx
|
||||
|
||||
from litellm.llms.azure.responses.transformation import AzureOpenAIResponsesAPIConfig
|
||||
from litellm.llms.azure_ai.common_utils import (
|
||||
AzureFoundryModelInfo,
|
||||
api_key_header_for_base,
|
||||
get_azure_ai_auth_headers,
|
||||
)
|
||||
from litellm.types.router import GenericLiteLLMParams
|
||||
from litellm.types.utils import LlmProviders
|
||||
|
||||
_PROJECT_PATH_PREFIX: Final = ("api", "projects")
|
||||
_RESPONSES_PATH: Final = ("openai", "v1", "responses")
|
||||
|
||||
|
||||
def _responses_url(api_base: str) -> str:
|
||||
base_url: Final = httpx.URL(api_base)
|
||||
segments: Final = tuple(segment for segment in base_url.path.split("/") if segment)
|
||||
project_root: Final = segments[:3] if segments[:2] == _PROJECT_PATH_PREFIX else ()
|
||||
return str(base_url.copy_with(path="/" + "/".join((*project_root, *_RESPONSES_PATH)), query=None))
|
||||
|
||||
|
||||
class AzureAIResponsesAPIConfig(AzureOpenAIResponsesAPIConfig):
|
||||
@property
|
||||
def custom_llm_provider(self) -> LlmProviders:
|
||||
return LlmProviders.AZURE_AI
|
||||
|
||||
def validate_environment(self, headers: dict, model: str, litellm_params: GenericLiteLLMParams | None) -> dict:
|
||||
params: Final = litellm_params or GenericLiteLLMParams()
|
||||
auth_headers: Final = get_azure_ai_auth_headers(
|
||||
api_key=AzureFoundryModelInfo.get_api_key(params.api_key),
|
||||
litellm_params=params.model_dump(),
|
||||
api_key_header=api_key_header_for_base(AzureFoundryModelInfo.get_api_base(params.api_base)),
|
||||
)
|
||||
return { # mutable-ok: the handler updates the returned headers in place per the dict contract
|
||||
**headers,
|
||||
**auth_headers,
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
|
||||
def supports_native_websocket(self) -> bool:
|
||||
return False
|
||||
|
||||
def get_complete_url(self, api_base: str | None, litellm_params: dict) -> str:
|
||||
resolved_base: Final = AzureFoundryModelInfo.get_api_base(api_base)
|
||||
if resolved_base is None:
|
||||
raise ValueError(
|
||||
"api_base is required for the Azure AI Foundry Responses API. "
|
||||
"Set the api_base parameter or the AZURE_AI_API_BASE environment variable."
|
||||
)
|
||||
return _responses_url(resolved_base)
|
||||
|
|
@ -1,8 +1,17 @@
|
|||
from abc import ABC, abstractmethod
|
||||
from collections.abc import Sequence
|
||||
from collections.abc import Mapping, Sequence
|
||||
from dataclasses import dataclass, field
|
||||
from typing import TYPE_CHECKING, Any, ClassVar, Final, Optional
|
||||
|
||||
from litellm.llms.base_llm.guardrail_translation.utils import (
|
||||
effective_scan_only_tool_results_for_guardrail,
|
||||
effective_skip_system_message_for_guardrail,
|
||||
effective_skip_tool_message_for_guardrail,
|
||||
request_tools,
|
||||
response_assistant_turn,
|
||||
scoped_structured_message_indices,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from fastapi import HTTPException
|
||||
|
||||
|
|
@ -12,7 +21,43 @@ if TYPE_CHECKING:
|
|||
)
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.types.llms.openai import AllMessageValues
|
||||
from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolParam
|
||||
from litellm.types.utils import GenericGuardrailAPIInputs
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class RequestScanContext:
|
||||
"""The scoped request turns and tool definitions a guardrail's request scan sees, in OpenAI chat shape."""
|
||||
|
||||
structured_messages: tuple["AllMessageValues", ...] = ()
|
||||
tools: tuple["ChatCompletionToolParam", ...] = ()
|
||||
conversation_supplied: bool = False
|
||||
|
||||
@staticmethod
|
||||
def scoped(
|
||||
structured_messages: Sequence["AllMessageValues"],
|
||||
tools: Sequence["ChatCompletionToolParam"],
|
||||
guardrail_to_apply: "CustomGuardrail",
|
||||
*,
|
||||
skip_system: bool | None = None,
|
||||
) -> "RequestScanContext":
|
||||
scan_only_tool_results: Final = effective_scan_only_tool_results_for_guardrail(guardrail_to_apply)
|
||||
scoped_indices: Final = scoped_structured_message_indices(
|
||||
structured_messages,
|
||||
scan_only_tool_results=scan_only_tool_results,
|
||||
skip_system=(
|
||||
effective_skip_system_message_for_guardrail(guardrail_to_apply) if skip_system is None else skip_system
|
||||
),
|
||||
skip_tool=effective_skip_tool_message_for_guardrail(guardrail_to_apply),
|
||||
)
|
||||
return RequestScanContext(
|
||||
structured_messages=tuple(structured_messages[index] for index in scoped_indices),
|
||||
tools=() if scan_only_tool_results else tuple(tools),
|
||||
conversation_supplied=bool(structured_messages),
|
||||
)
|
||||
|
||||
|
||||
REQUEST_SCAN_CONTEXT_KEY: Final = "litellm_request_scan_context"
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
|
|
@ -257,6 +302,50 @@ class BaseTranslation(ABC):
|
|||
"""
|
||||
return None
|
||||
|
||||
def request_scan_context(
|
||||
self, data: Mapping[str, object], guardrail_to_apply: "CustomGuardrail"
|
||||
) -> RequestScanContext:
|
||||
"""Override wherever ``process_input_messages`` scopes or translates the request differently."""
|
||||
structured_messages: Final = self.get_structured_messages(
|
||||
dict(data) # mutable-ok: get_structured_messages takes the request as a dict
|
||||
)
|
||||
return RequestScanContext.scoped(
|
||||
structured_messages or (), request_tools(data.get("tools")), guardrail_to_apply
|
||||
)
|
||||
|
||||
def with_response_context(
|
||||
self,
|
||||
inputs: "GenericGuardrailAPIInputs",
|
||||
request_data: Mapping[str, object] | None,
|
||||
guardrail_to_apply: "CustomGuardrail",
|
||||
) -> "GenericGuardrailAPIInputs":
|
||||
"""``inputs`` plus the scoped request conversation, closed by the scanned reply, and the request tools."""
|
||||
if request_data is None:
|
||||
return inputs
|
||||
precomputed: Final = request_data.get(REQUEST_SCAN_CONTEXT_KEY)
|
||||
context: Final = (
|
||||
precomputed
|
||||
if isinstance(precomputed, RequestScanContext)
|
||||
else self.request_scan_context(request_data, guardrail_to_apply)
|
||||
)
|
||||
if not context.conversation_supplied:
|
||||
return inputs
|
||||
assistant_turn: Final = response_assistant_turn(inputs.get("texts") or (), inputs.get("tool_calls") or ())
|
||||
contextual_inputs: Final[GenericGuardrailAPIInputs] = {
|
||||
**inputs,
|
||||
"structured_messages": [ # mutable-ok: GenericGuardrailAPIInputs fields are lists
|
||||
*context.structured_messages,
|
||||
*(() if assistant_turn is None else (assistant_turn,)),
|
||||
],
|
||||
}
|
||||
if not context.tools:
|
||||
return contextual_inputs
|
||||
with_tools: Final[GenericGuardrailAPIInputs] = {
|
||||
**contextual_inputs,
|
||||
"tools": list(context.tools), # mutable-ok: GenericGuardrailAPIInputs fields are lists
|
||||
}
|
||||
return with_tools
|
||||
|
||||
def extract_request_tool_names(self, data: dict) -> list[str]:
|
||||
"""
|
||||
Extract tool names from the request body for allowlist/policy checks.
|
||||
|
|
|
|||
|
|
@ -2,12 +2,24 @@ from __future__ import annotations
|
|||
|
||||
import json
|
||||
from collections.abc import Callable, Iterator, Mapping, Sequence
|
||||
from typing import Final, TypeVar, cast # noqa: TID251 # a rebuilt chat row has no typed constructor across roles
|
||||
from typing import TYPE_CHECKING, Final, TypeVar, cast # noqa: TID251 # a rebuilt chat row has no typed constructor
|
||||
|
||||
from pydantic import BaseModel
|
||||
|
||||
from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicUsage
|
||||
from litellm.types.llms.openai import AllMessageValues, ResponseAPIUsage
|
||||
from litellm.types.llms.openai import (
|
||||
AllMessageValues,
|
||||
ChatCompletionAssistantMessage,
|
||||
ChatCompletionAssistantToolCall,
|
||||
ChatCompletionTextObject,
|
||||
ChatCompletionToolCallChunk,
|
||||
ChatCompletionToolCallFunctionChunk,
|
||||
ChatCompletionToolParam,
|
||||
ResponseAPIUsage,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.types.utils import ChatCompletionMessageToolCall
|
||||
|
||||
|
||||
def _anthropic_stream_chunk_events(item: object) -> list[dict]:
|
||||
|
|
@ -278,9 +290,57 @@ def scoped_structured_message_indices(
|
|||
)
|
||||
|
||||
|
||||
def _assistant_tool_call(
|
||||
tool_call: ChatCompletionToolCallChunk | ChatCompletionMessageToolCall,
|
||||
) -> ChatCompletionAssistantToolCall:
|
||||
function: Final = stream_item_field(tool_call, "function")
|
||||
tool_call_id: Final = stream_item_field(tool_call, "id")
|
||||
name: Final = stream_item_field(function, "name")
|
||||
arguments: Final = stream_item_field(function, "arguments")
|
||||
return ChatCompletionAssistantToolCall(
|
||||
id=tool_call_id if isinstance(tool_call_id, str) else None,
|
||||
type="function",
|
||||
function=ChatCompletionToolCallFunctionChunk(
|
||||
name=name if isinstance(name, str) else None,
|
||||
arguments=arguments if isinstance(arguments, str) else "",
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def response_assistant_turn(
|
||||
texts: Sequence[str],
|
||||
tool_calls: Sequence[ChatCompletionToolCallChunk] | Sequence[ChatCompletionMessageToolCall],
|
||||
) -> ChatCompletionAssistantMessage | None:
|
||||
"""The scanned reply as the assistant turn closing the request conversation."""
|
||||
assistant_tool_calls: Final = tuple(_assistant_tool_call(tool_call) for tool_call in tool_calls)
|
||||
if not texts and not assistant_tool_calls:
|
||||
return None
|
||||
content: Final = (
|
||||
texts[0]
|
||||
if len(texts) == 1
|
||||
else tuple(ChatCompletionTextObject(type="text", text=text) for text in texts) or None
|
||||
)
|
||||
if not assistant_tool_calls:
|
||||
return ChatCompletionAssistantMessage(role="assistant", content=content)
|
||||
return ChatCompletionAssistantMessage(
|
||||
role="assistant",
|
||||
content=content,
|
||||
tool_calls=list(assistant_tool_calls), # mutable-ok: the assistant message type takes a list
|
||||
)
|
||||
|
||||
|
||||
ToolT = TypeVar("ToolT")
|
||||
|
||||
|
||||
def request_tools(raw_tools: object) -> tuple[ChatCompletionToolParam, ...]:
|
||||
"""The request's ``tools`` list, as the chat completion request model already validated it upstream."""
|
||||
if not isinstance(raw_tools, list):
|
||||
return ()
|
||||
return tuple(
|
||||
cast(Sequence[ChatCompletionToolParam], raw_tools) # cast-ok: the request model validated tools upstream
|
||||
)
|
||||
|
||||
|
||||
def openai_tool_name(tool: object) -> str | None:
|
||||
if not isinstance(tool, dict):
|
||||
return None
|
||||
|
|
|
|||
|
|
@ -53,6 +53,7 @@ from litellm.types.llms.openai import (
|
|||
AllMessageValues,
|
||||
ChatCompletionAnnotation,
|
||||
ChatCompletionAssistantMessage,
|
||||
ChatCompletionAssistantToolCall,
|
||||
ChatCompletionRedactedThinkingBlock,
|
||||
ChatCompletionResponseMessage,
|
||||
ChatCompletionSystemMessage,
|
||||
|
|
@ -205,6 +206,84 @@ class AmazonConverseConfig(BaseConfig):
|
|||
|
||||
return messages_copy
|
||||
|
||||
@staticmethod
|
||||
def _has_orphaned_tool_blocks(messages: list[AllMessageValues]) -> bool:
|
||||
return any(
|
||||
(m.get("role") == "assistant" and m.get("tool_calls")) or m.get("role") in ("tool", "function")
|
||||
for m in messages
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _neutralize_orphaned_tool_blocks(
|
||||
messages: list[AllMessageValues], optional_params: dict
|
||||
) -> list[AllMessageValues]:
|
||||
if optional_params.get("tools") or not AmazonConverseConfig._has_orphaned_tool_blocks(messages):
|
||||
return messages
|
||||
|
||||
from litellm.litellm_core_utils.prompt_templates.common_utils import (
|
||||
convert_content_list_to_str,
|
||||
)
|
||||
|
||||
def _tool_call_text(tool_call: ChatCompletionAssistantToolCall) -> str:
|
||||
function = tool_call.get("function") or {}
|
||||
name = function.get("name") or "unknown_tool"
|
||||
arguments = function.get("arguments") or ""
|
||||
call_id = tool_call.get("id")
|
||||
label = f"tool call {call_id}" if call_id else "tool call"
|
||||
return f"[{label}: {name}({arguments})]"
|
||||
|
||||
def _result_text(message: AllMessageValues) -> str:
|
||||
rendered = convert_content_list_to_str(message).strip()
|
||||
return rendered or "<non-text tool result omitted>"
|
||||
|
||||
guardrail_active: Final = "guardrailConfig" in optional_params
|
||||
|
||||
def _rewrite(message: AllMessageValues) -> AllMessageValues:
|
||||
role = message.get("role")
|
||||
tool_calls = message.get("tool_calls")
|
||||
if role == "assistant" and tool_calls:
|
||||
base_text: Final = convert_content_list_to_str(message)
|
||||
call_texts: Final = tuple(_tool_call_text(call) for call in tool_calls)
|
||||
text: Final = "\n".join(part for part in (base_text, *call_texts) if part)
|
||||
return ChatCompletionAssistantMessage(role="assistant", content=text)
|
||||
if role in ("tool", "function"):
|
||||
tool_call_id = message.get("tool_call_id")
|
||||
name = message.get("name")
|
||||
label = f"tool result for {tool_call_id or name or 'unknown'}"
|
||||
result_text: Final = f"[{label}: {_result_text(message)}]"
|
||||
# Tool results are externally controlled, so guard them wherever they
|
||||
# land in history; _convert_consecutive_user_messages_to_guarded_text
|
||||
# only covers the trailing user turn.
|
||||
content: Final = [{"type": "guarded_text", "text": result_text}] if guardrail_active else result_text
|
||||
return ChatCompletionUserMessage(role="user", content=content)
|
||||
return message
|
||||
|
||||
verbose_logger.warning(
|
||||
"litellm.bedrock: request has tool blocks in message history but no "
|
||||
"`tools=` param; neutralizing orphaned tool blocks to text so Bedrock "
|
||||
"accepts the request without a toolConfig. Non-text tool-result "
|
||||
"payloads are dropped. Pass `tools=` to preserve structured tool calling."
|
||||
)
|
||||
return [_rewrite(message) for message in messages]
|
||||
|
||||
@staticmethod
|
||||
def _handle_orphaned_tool_blocks(messages: list[AllMessageValues], optional_params: dict) -> list[AllMessageValues]:
|
||||
if litellm.bedrock_neutralize_orphaned_tool_blocks:
|
||||
return AmazonConverseConfig._neutralize_orphaned_tool_blocks(messages, optional_params)
|
||||
|
||||
if "tools" in optional_params or not has_tool_call_blocks(messages):
|
||||
return messages
|
||||
|
||||
if litellm.modify_params:
|
||||
optional_params["tools"] = add_dummy_tool(custom_llm_provider="bedrock_converse")
|
||||
return messages
|
||||
|
||||
raise litellm.utils.UnsupportedParamsError(
|
||||
message="Bedrock doesn't support tool calling without `tools=` param specified. Pass `tools=` param OR set `litellm.modify_params = True` // `litellm_settings::modify_params: True` to add dummy tool to the request.",
|
||||
model="",
|
||||
llm_provider="bedrock",
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def get_config(cls):
|
||||
return {
|
||||
|
|
@ -1609,20 +1688,6 @@ class AmazonConverseConfig(BaseConfig):
|
|||
drop_params: bool = False,
|
||||
litellm_params: Mapping[str, object] | None = None,
|
||||
) -> CommonRequestObject:
|
||||
## VALIDATE REQUEST
|
||||
"""
|
||||
Bedrock doesn't support tool calling without `tools=` param specified.
|
||||
"""
|
||||
if "tools" not in optional_params and messages is not None and has_tool_call_blocks(messages):
|
||||
if litellm.modify_params:
|
||||
optional_params["tools"] = add_dummy_tool(custom_llm_provider="bedrock_converse")
|
||||
else:
|
||||
raise litellm.UnsupportedParamsError(
|
||||
message="Bedrock doesn't support tool calling without `tools=` param specified. Pass `tools=` param OR set `litellm.modify_params = True` // `litellm_settings::modify_params: True` to add dummy tool to the request.",
|
||||
model="",
|
||||
llm_provider="bedrock",
|
||||
)
|
||||
|
||||
# Drop thinking param if thinking is enabled but thinking_blocks are missing
|
||||
# This prevents the error: "Expected thinking or redacted_thinking, but found tool_use"
|
||||
#
|
||||
|
|
@ -1735,7 +1800,9 @@ class AmazonConverseConfig(BaseConfig):
|
|||
messages, system_content_blocks = self._transform_system_message(messages, model=model)
|
||||
|
||||
# Convert last user message to guarded_text if guardrailConfig is present
|
||||
messages = self._convert_consecutive_user_messages_to_guarded_text(messages, optional_params)
|
||||
messages = self._convert_consecutive_user_messages_to_guarded_text(
|
||||
self._handle_orphaned_tool_blocks(messages, optional_params), optional_params
|
||||
)
|
||||
## TRANSFORMATION ##
|
||||
|
||||
_data: Final[CommonRequestObject] = self._transform_request_helper(
|
||||
|
|
@ -1796,7 +1863,9 @@ class AmazonConverseConfig(BaseConfig):
|
|||
messages, system_content_blocks = self._transform_system_message(messages, model=model)
|
||||
|
||||
# Convert last user message to guarded_text if guardrailConfig is present
|
||||
messages = self._convert_consecutive_user_messages_to_guarded_text(messages, optional_params)
|
||||
messages = self._convert_consecutive_user_messages_to_guarded_text(
|
||||
self._handle_orphaned_tool_blocks(messages, optional_params), optional_params
|
||||
)
|
||||
|
||||
_data: Final[CommonRequestObject] = self._transform_request_helper(
|
||||
model=model,
|
||||
|
|
|
|||
|
|
@ -22,6 +22,7 @@ from litellm.llms.bedrock_mantle.common_utils import (
|
|||
BEDROCK_MANTLE_DEFAULT_REGION,
|
||||
BedrockMantleAuthMixin,
|
||||
)
|
||||
from litellm.llms.openai.chat.gpt_5_transformation import is_gpt_reasoning_series_name
|
||||
from litellm.secret_managers.main import get_secret_str
|
||||
from litellm.types.llms.openai import AllMessageValues
|
||||
from litellm.types.router import GenericLiteLLMParams
|
||||
|
|
@ -108,13 +109,22 @@ class BedrockMantleChatConfig(BedrockMantleAuthMixin, OpenAILikeChatConfig):
|
|||
|
||||
def get_supported_openai_params(self, model: str) -> list:
|
||||
base_params: Final = super().get_supported_openai_params(model)
|
||||
extra_params: Final = tuple(
|
||||
param
|
||||
for param, supported in (
|
||||
("verbosity", is_gpt_reasoning_series_name(model)),
|
||||
("reasoning_effort", self._supports_reasoning(model)),
|
||||
)
|
||||
if supported and param not in base_params
|
||||
)
|
||||
return [*base_params, *extra_params] # mutable-ok: fresh list required by the inherited signature
|
||||
|
||||
def _supports_reasoning(self, model: str) -> bool:
|
||||
try:
|
||||
if litellm.supports_reasoning(model=model, custom_llm_provider=self.custom_llm_provider):
|
||||
if "reasoning_effort" not in base_params:
|
||||
base_params.append("reasoning_effort")
|
||||
return litellm.supports_reasoning(model=model, custom_llm_provider=self.custom_llm_provider)
|
||||
except Exception as e:
|
||||
verbose_logger.debug("BedrockMantleChatConfig: error checking reasoning support: %s", e)
|
||||
return base_params
|
||||
return False
|
||||
|
||||
def get_model_response_iterator(
|
||||
self,
|
||||
|
|
|
|||
|
|
@ -12,6 +12,12 @@ from ...openai.chat.gpt_transformation import OpenAIGPTConfig
|
|||
|
||||
|
||||
class DashScopeChatConfig(OpenAIGPTConfig):
|
||||
def get_supported_openai_params(self, model: str) -> list[str]: # mutable-ok: base class contract returns a list
|
||||
return [ # mutable-ok: base class contract returns a list
|
||||
*super().get_supported_openai_params(model=model),
|
||||
"reasoning_effort",
|
||||
]
|
||||
|
||||
def remove_cache_control_flag_from_messages_and_tools(
|
||||
self,
|
||||
model: str,
|
||||
|
|
|
|||
|
|
@ -54,7 +54,7 @@ class LiteLLMProxyChatConfig(OpenAIGPTConfig):
|
|||
return api_key or get_secret_str("LITELLM_PROXY_API_KEY")
|
||||
|
||||
@staticmethod
|
||||
def _should_use_litellm_proxy_by_default(
|
||||
def should_use_litellm_proxy_by_default(
|
||||
litellm_params: LiteLLM_Params | None = None,
|
||||
):
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -4,9 +4,9 @@ from typing import Final
|
|||
|
||||
import litellm
|
||||
from litellm.utils import (
|
||||
_is_explicitly_disabled_factory,
|
||||
_supports_factory,
|
||||
declared_value_factory,
|
||||
is_explicitly_disabled_factory,
|
||||
)
|
||||
|
||||
from .gpt_transformation import OpenAIGPTConfig
|
||||
|
|
@ -192,7 +192,7 @@ class OpenAIGPT5Config(OpenAIGPTConfig):
|
|||
|
||||
Use this for opt-out checks where unknown models should be allowed through.
|
||||
"""
|
||||
return _is_explicitly_disabled_factory(
|
||||
return is_explicitly_disabled_factory(
|
||||
model=cls._model_map_lookup_name(model),
|
||||
custom_llm_provider=None,
|
||||
key=f"supports_{level}_reasoning_effort",
|
||||
|
|
|
|||
|
|
@ -452,7 +452,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
|
|||
inputs["model"] = response.model
|
||||
|
||||
guardrailed_inputs: Final = await guardrail_to_apply.apply_guardrail(
|
||||
inputs=inputs,
|
||||
inputs=self.with_response_context(inputs, request_data, guardrail_to_apply),
|
||||
request_data=request_data,
|
||||
input_type="response",
|
||||
logging_obj=litellm_logging_obj,
|
||||
|
|
@ -615,7 +615,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
|
|||
if responses_so_far and hasattr(responses_so_far[0], "model") and responses_so_far[0].model:
|
||||
inputs["model"] = responses_so_far[0].model
|
||||
guardrailed_inputs: Final = await guardrail_to_apply.apply_guardrail(
|
||||
inputs=inputs,
|
||||
inputs=self.with_response_context(inputs, request_data, guardrail_to_apply),
|
||||
request_data=request_data,
|
||||
input_type="response",
|
||||
logging_obj=litellm_logging_obj,
|
||||
|
|
@ -760,7 +760,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
|
|||
if responses_so_far and getattr(responses_so_far[0], "model", None):
|
||||
inputs["model"] = responses_so_far[0].model
|
||||
guardrailed_inputs: Final = await guardrail_to_apply.apply_guardrail(
|
||||
inputs=inputs,
|
||||
inputs=self.with_response_context(inputs, request_data, guardrail_to_apply),
|
||||
request_data=request_data,
|
||||
input_type="response",
|
||||
logging_obj=litellm_logging_obj,
|
||||
|
|
|
|||
|
|
@ -48,6 +48,7 @@ from litellm.completion_extras.litellm_responses_transformation.transformation i
|
|||
)
|
||||
from litellm.llms.base_llm.guardrail_translation.base_translation import (
|
||||
BaseTranslation,
|
||||
RequestScanContext,
|
||||
StreamingScanKey,
|
||||
StreamTransformSink,
|
||||
)
|
||||
|
|
@ -451,6 +452,28 @@ class OpenAIResponsesHandler(BaseTranslation):
|
|||
)
|
||||
return cast(list[AllMessageValues], messages) if messages else None
|
||||
|
||||
def request_scan_context(
|
||||
self, data: Mapping[str, object], guardrail_to_apply: "CustomGuardrail"
|
||||
) -> RequestScanContext:
|
||||
raw_tools: Final = data.get("tools")
|
||||
structured_messages: Final = tuple(
|
||||
self.get_structured_messages(
|
||||
dict(data) # mutable-ok: get_structured_messages takes the request as a dict
|
||||
)
|
||||
or ()
|
||||
)
|
||||
return RequestScanContext(
|
||||
structured_messages=structured_messages,
|
||||
tools=tuple(
|
||||
cast(ChatCompletionToolParam, tool) # cast-ok: mcp tools ride along in the guardrail's tool list
|
||||
for form in LiteLLMCompletionResponsesConfig.responses_tools_to_chat_forms(
|
||||
tuple(raw_tools) if isinstance(raw_tools, list) else ()
|
||||
)
|
||||
for tool in form.chat_tools
|
||||
),
|
||||
conversation_supplied=bool(structured_messages),
|
||||
)
|
||||
|
||||
async def process_input_messages(
|
||||
self,
|
||||
data: dict,
|
||||
|
|
@ -754,7 +777,7 @@ class OpenAIResponsesHandler(BaseTranslation):
|
|||
|
||||
pre_guardrail_tool_calls: Final = _tool_call_shapes(tool_calls_to_check)
|
||||
guardrailed_inputs: Final = await guardrail_to_apply.apply_guardrail(
|
||||
inputs=inputs,
|
||||
inputs=self.with_response_context(inputs, request_data, guardrail_to_apply),
|
||||
request_data=request_data,
|
||||
input_type="response",
|
||||
logging_obj=litellm_logging_obj,
|
||||
|
|
@ -867,7 +890,7 @@ class OpenAIResponsesHandler(BaseTranslation):
|
|||
|
||||
pre_guardrail_tool_calls: Final = _tool_call_shapes(tool_calls_to_check)
|
||||
guardrailed_inputs: Final = await guardrail_to_apply.apply_guardrail(
|
||||
inputs=inputs,
|
||||
inputs=self.with_response_context(inputs, request_data, guardrail_to_apply),
|
||||
request_data=request_data,
|
||||
input_type="response",
|
||||
logging_obj=litellm_logging_obj,
|
||||
|
|
@ -926,7 +949,7 @@ class OpenAIResponsesHandler(BaseTranslation):
|
|||
if hasattr(model_response_stream, "model") and model_response_stream.model:
|
||||
inputs["model"] = model_response_stream.model
|
||||
await guardrail_to_apply.apply_guardrail(
|
||||
inputs=inputs,
|
||||
inputs=self.with_response_context(inputs, request_data, guardrail_to_apply),
|
||||
request_data=request_data if request_data is not None else {},
|
||||
input_type="response",
|
||||
logging_obj=litellm_logging_obj,
|
||||
|
|
@ -949,7 +972,7 @@ class OpenAIResponsesHandler(BaseTranslation):
|
|||
if response_model:
|
||||
fallback_inputs["model"] = response_model
|
||||
fallback_outputs: Final = await guardrail_to_apply.apply_guardrail(
|
||||
inputs=fallback_inputs,
|
||||
inputs=self.with_response_context(fallback_inputs, request_data, guardrail_to_apply),
|
||||
request_data=request_data if request_data is not None else {},
|
||||
input_type="response",
|
||||
logging_obj=litellm_logging_obj,
|
||||
|
|
|
|||
92
litellm/llms/openai_like/model_info.py
Normal file
92
litellm/llms/openai_like/model_info.py
Normal file
|
|
@ -0,0 +1,92 @@
|
|||
import hashlib
|
||||
import json
|
||||
from collections.abc import Mapping
|
||||
from types import MappingProxyType
|
||||
from typing import Annotated, Final, TypeAlias
|
||||
|
||||
import httpx
|
||||
from pydantic import BaseModel, BeforeValidator, ConfigDict
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.caching.in_memory_cache import InMemoryCache
|
||||
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
|
||||
from litellm.utils import _add_path_to_api_base # pyright: ignore[reportPrivateUsage] # shared provider URL helper
|
||||
|
||||
MODEL_INFO_REFRESH_SECONDS: Final = 300
|
||||
MODEL_INFO_REFRESH_CONCURRENCY: Final = 8
|
||||
MODEL_INFO_DISCOVERY_PROVIDERS: Final = frozenset({"hosted_vllm", "openai", "text-completion-openai", "openai_like"})
|
||||
_EMPTY_LIMITS: Final[Mapping[str, int]] = MappingProxyType({})
|
||||
|
||||
|
||||
def _positive_limit(value: object) -> int | None:
|
||||
return value if isinstance(value, int) and not isinstance(value, bool) and value > 0 else None
|
||||
|
||||
|
||||
_TokenLimit: TypeAlias = Annotated[int | None, BeforeValidator(_positive_limit)]
|
||||
|
||||
|
||||
class _ModelCard(BaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
id: str
|
||||
max_model_len: _TokenLimit = None
|
||||
context_length: _TokenLimit = None
|
||||
max_input_tokens: _TokenLimit = None
|
||||
max_output_tokens: _TokenLimit = None
|
||||
|
||||
def token_limits(self) -> Mapping[str, int]:
|
||||
context: Final = self.max_model_len or self.context_length
|
||||
input_limit: Final = self.max_input_tokens or context
|
||||
output_limit: Final = self.max_output_tokens or context
|
||||
return MappingProxyType(
|
||||
{
|
||||
key: value
|
||||
for key, value in (
|
||||
("max_tokens", context),
|
||||
("max_input_tokens", min(input_limit, context) if input_limit and context else input_limit),
|
||||
("max_output_tokens", min(output_limit, context) if output_limit and context else output_limit),
|
||||
)
|
||||
if value is not None
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
class _ModelList(BaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
data: tuple[_ModelCard, ...] = ()
|
||||
|
||||
|
||||
async def get_openai_compatible_model_info(
|
||||
*,
|
||||
model: str,
|
||||
api_base: str,
|
||||
headers: Mapping[str, str],
|
||||
client: AsyncHTTPHandler,
|
||||
cache: InMemoryCache,
|
||||
) -> Mapping[str, int]:
|
||||
url: Final = _add_path_to_api_base(api_base, "/v1/models")
|
||||
cache_key: Final = (
|
||||
"upstream_model_info:" + hashlib.sha256(json.dumps((url, sorted(headers.items()))).encode()).hexdigest()
|
||||
)
|
||||
cached: Final[object] = cache.get_cache(cache_key)
|
||||
if isinstance(cached, _ModelList):
|
||||
return next((card.token_limits() for card in cached.data if card.id == model), _EMPTY_LIMITS)
|
||||
|
||||
try:
|
||||
response: Final = await client.get(
|
||||
url=url,
|
||||
headers=dict(headers), # mutable-ok: AsyncHTTPHandler requires a concrete dict
|
||||
timeout=httpx.Timeout(5.0),
|
||||
follow_redirects=False,
|
||||
max_response_bytes=2 * 1024 * 1024,
|
||||
)
|
||||
response.raise_for_status()
|
||||
models: Final = _ModelList.model_validate_json(response.content)
|
||||
except Exception: # noqa: BLE001 # optional upstream metadata must not interrupt proxy refresh
|
||||
verbose_logger.debug("Could not discover upstream model token limits")
|
||||
cache.set_cache(cache_key, _ModelList(), ttl=60)
|
||||
return _EMPTY_LIMITS
|
||||
|
||||
cache.set_cache(cache_key, models, ttl=MODEL_INFO_REFRESH_SECONDS)
|
||||
return next((card.token_limits() for card in models.data if card.id == model), _EMPTY_LIMITS)
|
||||
|
|
@ -79,6 +79,7 @@ from litellm.utils import (
|
|||
CustomStreamWrapper,
|
||||
ModelResponse,
|
||||
is_base64_encoded,
|
||||
is_explicitly_disabled_factory,
|
||||
supports_reasoning,
|
||||
)
|
||||
|
||||
|
|
@ -866,6 +867,14 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
|
|||
else:
|
||||
raise _unsupported_reasoning_effort(reasoning_effort)
|
||||
|
||||
@staticmethod
|
||||
def _supports_minimal_thinking_level(model: str) -> bool:
|
||||
lowered: Final = model.lower()
|
||||
is_gemini3flash: Final = "gemini-3" in lowered and "flash" in lowered
|
||||
return is_gemini3flash and not is_explicitly_disabled_factory(
|
||||
model=model, custom_llm_provider=None, key="supports_minimal_reasoning_effort"
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _map_reasoning_effort_to_thinking_level(
|
||||
reasoning_effort: str,
|
||||
|
|
@ -880,13 +889,11 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
|
|||
Returns:
|
||||
GeminiThinkingConfig with thinkingLevel and includeThoughts
|
||||
"""
|
||||
# Check if this is gemini-3-flash which supports MINIMAL thinking level
|
||||
# Covers gemini-3-flash, gemini-3-flash-preview, gemini-3.1-flash, gemini-3.1-flash-lite-preview,
|
||||
# gemini-3.5-flash, and any future 3.x-flash variants.
|
||||
is_gemini3flash: Final = model and ("flash" in model.lower() and "gemini-3" in model.lower())
|
||||
supports_minimal: Final = bool(model) and VertexGeminiConfig._supports_minimal_thinking_level(model)
|
||||
is_gemini31pro: Final = model and ("gemini-3.1-pro-preview" in model.lower())
|
||||
if reasoning_effort == "minimal":
|
||||
if is_gemini3flash:
|
||||
if supports_minimal:
|
||||
return {"thinkingLevel": "minimal", "includeThoughts": True}
|
||||
else:
|
||||
return {"thinkingLevel": "low", "includeThoughts": True}
|
||||
|
|
@ -899,18 +906,11 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
|
|||
return {"thinkingLevel": "high", "includeThoughts": True}
|
||||
elif reasoning_effort == "high":
|
||||
return {"thinkingLevel": "high", "includeThoughts": True}
|
||||
elif reasoning_effort == "disable":
|
||||
# Gemini 3 cannot fully disable thinking, so we use "minimal" for gemini-3-flash-preview, "low" for others
|
||||
if is_gemini3flash:
|
||||
return {"thinkingLevel": "minimal", "includeThoughts": False}
|
||||
else:
|
||||
return {"thinkingLevel": "low", "includeThoughts": False}
|
||||
elif reasoning_effort == "none":
|
||||
# For gemini-3-flash-preview, use "minimal" instead of "low"
|
||||
if is_gemini3flash:
|
||||
return {"thinkingLevel": "minimal", "includeThoughts": False}
|
||||
else:
|
||||
return {"thinkingLevel": "low", "includeThoughts": False}
|
||||
elif reasoning_effort in ("disable", "none"):
|
||||
return {
|
||||
"thinkingLevel": "minimal" if supports_minimal else "low",
|
||||
"includeThoughts": False,
|
||||
}
|
||||
else:
|
||||
raise _unsupported_reasoning_effort(reasoning_effort)
|
||||
|
||||
|
|
@ -977,8 +977,9 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
|
|||
params["includeThoughts"] = True
|
||||
# Follow provider defaults unless explicitly opted into legacy behavior.
|
||||
if litellm.enable_gemini_default_thinking_level_low is True:
|
||||
is_gemini3flash: Final = "gemini-3" in model.lower() and "flash" in model.lower()
|
||||
params["thinkingLevel"] = "minimal" if is_gemini3flash else "low"
|
||||
params["thinkingLevel"] = (
|
||||
"minimal" if VertexGeminiConfig._supports_minimal_thinking_level(model) else "low"
|
||||
)
|
||||
else:
|
||||
# Thinking disabled
|
||||
params["includeThoughts"] = False
|
||||
|
|
|
|||
|
|
@ -377,8 +377,7 @@
|
|||
"supports_prompt_caching": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_vision": true,
|
||||
"supports_tool_choice": true,
|
||||
"cache_read_input_token_cost": 1.5e-08
|
||||
"supports_tool_choice": true
|
||||
},
|
||||
"amazon.nova-2-lite-v1:0": {
|
||||
"cache_read_input_token_cost": 7.5e-08,
|
||||
|
|
@ -561,8 +560,7 @@
|
|||
"supports_function_calling": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_tool_choice": true,
|
||||
"cache_read_input_token_cost": 8.75e-09
|
||||
"supports_tool_choice": true
|
||||
},
|
||||
"amazon.nova-pro-v1:0": {
|
||||
"cache_read_input_token_cost": 2e-07,
|
||||
|
|
@ -578,8 +576,7 @@
|
|||
"supports_prompt_caching": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_vision": true,
|
||||
"supports_tool_choice": true,
|
||||
"cache_read_input_token_cost": 2e-07
|
||||
"supports_tool_choice": true
|
||||
},
|
||||
"amazon.nova-sonic-v1:0": {
|
||||
"deprecation_date": "2026-09-14",
|
||||
|
|
@ -26108,6 +26105,7 @@
|
|||
"supports_pdf_input": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_minimal_reasoning_effort": false,
|
||||
"supports_response_schema": true,
|
||||
"supports_system_messages": true,
|
||||
"supports_tool_choice": true,
|
||||
|
|
@ -26165,6 +26163,7 @@
|
|||
"supports_pdf_input": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_minimal_reasoning_effort": false,
|
||||
"supports_response_schema": true,
|
||||
"supports_system_messages": true,
|
||||
"supports_tool_choice": true,
|
||||
|
|
@ -28114,6 +28113,7 @@
|
|||
"supports_pdf_input": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_minimal_reasoning_effort": false,
|
||||
"supports_response_schema": true,
|
||||
"supports_system_messages": true,
|
||||
"supports_tool_choice": true,
|
||||
|
|
@ -28173,6 +28173,7 @@
|
|||
"supports_pdf_input": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_minimal_reasoning_effort": false,
|
||||
"supports_response_schema": true,
|
||||
"supports_system_messages": true,
|
||||
"supports_tool_choice": true,
|
||||
|
|
@ -28595,6 +28596,7 @@
|
|||
"supports_pdf_input": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_minimal_reasoning_effort": false,
|
||||
"supports_response_schema": true,
|
||||
"supports_system_messages": true,
|
||||
"supports_tool_choice": true,
|
||||
|
|
@ -28652,6 +28654,7 @@
|
|||
"supports_pdf_input": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_minimal_reasoning_effort": false,
|
||||
"supports_response_schema": true,
|
||||
"supports_system_messages": true,
|
||||
"supports_tool_choice": true,
|
||||
|
|
@ -42310,6 +42313,20 @@
|
|||
"max_tokens": 128000,
|
||||
"mode": "chat"
|
||||
},
|
||||
"openrouter/stealth/union-alpha": {
|
||||
"input_cost_per_token": 0,
|
||||
"output_cost_per_token": 0,
|
||||
"litellm_provider": "openrouter",
|
||||
"max_input_tokens": 262144,
|
||||
"max_output_tokens": 131072,
|
||||
"max_tokens": 131072,
|
||||
"mode": "chat",
|
||||
"source": "https://openrouter.ai/stealth/union-alpha",
|
||||
"supports_function_calling": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_vision": true
|
||||
},
|
||||
"ovhcloud/DeepSeek-R1-Distill-Llama-70B": {
|
||||
"input_cost_per_token": 6.7e-07,
|
||||
"litellm_provider": "ovhcloud",
|
||||
|
|
@ -45794,8 +45811,7 @@
|
|||
"supports_prompt_caching": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_vision": true,
|
||||
"supports_tool_choice": true,
|
||||
"cache_read_input_token_cost": 1.5e-08
|
||||
"supports_tool_choice": true
|
||||
},
|
||||
"us.amazon.nova-micro-v1:0": {
|
||||
"cache_read_input_token_cost": 8.75e-09,
|
||||
|
|
@ -45809,8 +45825,7 @@
|
|||
"supports_function_calling": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_tool_choice": true,
|
||||
"cache_read_input_token_cost": 8.75e-09
|
||||
"supports_tool_choice": true
|
||||
},
|
||||
"us.amazon.nova-premier-v1:0": {
|
||||
"deprecation_date": "2026-09-14",
|
||||
|
|
@ -45842,8 +45857,7 @@
|
|||
"supports_prompt_caching": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_vision": true,
|
||||
"supports_tool_choice": true,
|
||||
"cache_read_input_token_cost": 2e-07
|
||||
"supports_tool_choice": true
|
||||
},
|
||||
"us.anthropic.claude-3-5-haiku-20241022-v1:0": {
|
||||
"cache_creation_input_token_cost": 1e-06,
|
||||
|
|
|
|||
|
|
@ -102,6 +102,7 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials import (
|
|||
UpstreamCredentialProvider,
|
||||
)
|
||||
from litellm.proxy._experimental.mcp_server.outbound_credentials.adapter import (
|
||||
prepare_mcp_client,
|
||||
raise_public,
|
||||
raise_token_exchange_challenge,
|
||||
raise_user_oauth_challenge,
|
||||
|
|
@ -2804,6 +2805,8 @@ class MCPServerManager:
|
|||
headers=headers,
|
||||
server_label=server.name or server.server_name or server.alias or server.server_id,
|
||||
relays_upstream_auth=server.is_client_forwarded_token,
|
||||
auth_type=server.auth_type,
|
||||
upstream_token_header=server.upstream_token_header,
|
||||
)
|
||||
tool_func.__name__ = prefixed_tool_name
|
||||
tool_func.__doc__ = description
|
||||
|
|
@ -4259,15 +4262,20 @@ class MCPServerManager:
|
|||
user_api_key_auth=user_api_key_auth,
|
||||
extra_headers=extra_headers,
|
||||
)
|
||||
return MCPClient(
|
||||
server_url=server_url,
|
||||
transport_type=transport,
|
||||
auth_type=resolved_server.auth_type,
|
||||
timeout=(resolved_server.timeout if resolved_server.timeout is not None else MCP_CLIENT_TIMEOUT),
|
||||
extra_headers=extra_headers,
|
||||
resolved_auth=resolved_auth,
|
||||
sampling_callback=sampling_cb,
|
||||
elicitation_callback=elicitation_cb,
|
||||
return await prepare_mcp_client(
|
||||
resolved_server,
|
||||
MCPClient(
|
||||
server_url=server_url,
|
||||
transport_type=transport,
|
||||
auth_type=resolved_server.auth_type,
|
||||
timeout=(
|
||||
resolved_server.timeout if resolved_server.timeout is not None else MCP_CLIENT_TIMEOUT
|
||||
),
|
||||
extra_headers=extra_headers,
|
||||
resolved_auth=resolved_auth,
|
||||
sampling_callback=sampling_cb,
|
||||
elicitation_callback=elicitation_cb,
|
||||
),
|
||||
)
|
||||
|
||||
# Create SigV4 auth if configured
|
||||
|
|
@ -4297,17 +4305,20 @@ class MCPServerManager:
|
|||
else AuthResolution.no_auth
|
||||
)
|
||||
record_auth_resolution(server.server_id, legacy_source)
|
||||
return MCPClient(
|
||||
server_url=server_url,
|
||||
transport_type=transport,
|
||||
auth_type=resolved_server.auth_type,
|
||||
auth_value=auth_value,
|
||||
auth_header_name=auth_header_name,
|
||||
timeout=(resolved_server.timeout if resolved_server.timeout is not None else MCP_CLIENT_TIMEOUT),
|
||||
extra_headers=extra_headers,
|
||||
aws_auth=aws_auth,
|
||||
sampling_callback=sampling_cb,
|
||||
elicitation_callback=elicitation_cb,
|
||||
return await prepare_mcp_client(
|
||||
resolved_server,
|
||||
MCPClient(
|
||||
server_url=server_url,
|
||||
transport_type=transport,
|
||||
auth_type=resolved_server.auth_type,
|
||||
auth_value=auth_value,
|
||||
auth_header_name=auth_header_name,
|
||||
timeout=(resolved_server.timeout if resolved_server.timeout is not None else MCP_CLIENT_TIMEOUT),
|
||||
extra_headers=extra_headers,
|
||||
aws_auth=aws_auth,
|
||||
sampling_callback=sampling_cb,
|
||||
elicitation_callback=elicitation_cb,
|
||||
),
|
||||
)
|
||||
|
||||
async def _get_tools_from_server(
|
||||
|
|
|
|||
|
|
@ -54,7 +54,7 @@ from litellm.llms.custom_httpx.http_handler import (
|
|||
from litellm.proxy._experimental.mcp_server.tool_registry import (
|
||||
global_mcp_tool_registry,
|
||||
)
|
||||
from litellm.types.mcp import credential_redirect_hook, custom_credential_slot
|
||||
from litellm.types.mcp import MCPAuthType, credential_redirect_hook, custom_credential_slot
|
||||
|
||||
|
||||
class _OpenAPIJSONSchema(TypedDict, total=False):
|
||||
|
|
@ -471,6 +471,8 @@ def create_tool_function(
|
|||
headers: dict[str, str] | None = None,
|
||||
server_label: str | None = None,
|
||||
relays_upstream_auth: bool = False,
|
||||
auth_type: MCPAuthType = None,
|
||||
upstream_token_header: str | None = None,
|
||||
):
|
||||
"""Create a tool function for an OpenAPI operation.
|
||||
|
||||
|
|
@ -503,6 +505,18 @@ def create_tool_function(
|
|||
by using **kwargs instead of named parameters.
|
||||
"""
|
||||
effective_headers: Final = _merge_openapi_tool_request_headers(headers)
|
||||
if auth_type is not None:
|
||||
from litellm.proxy._experimental.mcp_server.outbound_credentials.adapter import (
|
||||
raise_public,
|
||||
validate_static_credential,
|
||||
)
|
||||
from litellm.proxy._experimental.mcp_server.outbound_credentials.result import Error, Ok
|
||||
|
||||
match validate_static_credential(auth_type, effective_headers, upstream_token_header, headers or ()):
|
||||
case Error(error):
|
||||
raise_public(error)
|
||||
case Ok():
|
||||
pass
|
||||
|
||||
# Build URL from base_url and path
|
||||
url = base_url + path
|
||||
|
|
|
|||
|
|
@ -13,15 +13,17 @@ from __future__ import annotations
|
|||
|
||||
import base64
|
||||
import os
|
||||
from collections.abc import Iterable, Mapping
|
||||
from typing import TYPE_CHECKING, Final, Literal, NoReturn
|
||||
|
||||
from fastapi import HTTPException
|
||||
from pydantic import SecretStr
|
||||
from typing_extensions import assert_never
|
||||
|
||||
from litellm.experimental_mcp_client.client import strip_auth_scheme, to_basic_credentials
|
||||
from litellm.experimental_mcp_client.client import MCPClient, strip_auth_scheme, to_basic_credentials
|
||||
from litellm.proxy._experimental.mcp_server.exceptions import MCPServerURLCredentialsError
|
||||
from litellm.proxy._experimental.mcp_server.oauth_utils import resolve_upstream_resource
|
||||
from litellm.proxy._experimental.mcp_server.outbound_credentials.result import Error, Ok, Result
|
||||
from litellm.proxy._experimental.mcp_server.outbound_credentials.types import (
|
||||
DEFAULT_CREDENTIAL_HEADER,
|
||||
ApiKeyConfig,
|
||||
|
|
@ -39,7 +41,7 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.types import (
|
|||
Subject,
|
||||
TokenExchangeConfig,
|
||||
)
|
||||
from litellm.types.mcp import DEFAULT_SUBJECT_TOKEN_TYPE, MCPAuth
|
||||
from litellm.types.mcp import DEFAULT_SUBJECT_TOKEN_TYPE, MCPAuth, MCPAuthType, MCPTransport
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
|
|
@ -79,7 +81,7 @@ def to_server_spec(server: MCPServer) -> ServerSpec | None:
|
|||
|
||||
BYOK is the per-user source of the ``api_key`` mode; its scheme rides on ``auth_type`` just
|
||||
like a shared key, but the value is per-user and not migrated yet, so a BYOK server defers
|
||||
to v1 regardless of ``auth_type`` (this guard is the seam the BYOK arm replaces later).
|
||||
to v1 for its static schemes. Declared OBO always stays with the exchange arm.
|
||||
|
||||
Dispatches on the declared ``auth_type``. The match is exhaustive over ``MCPAuthType`` with
|
||||
an ``assert_never`` tail, so a newly added auth mode fails the type gate here until it is
|
||||
|
|
@ -90,8 +92,8 @@ def to_server_spec(server: MCPServer) -> ServerSpec | None:
|
|||
modes ``true_passthrough`` / ``oauth_delegate`` (``PassthroughConfig``); delegated/passthrough
|
||||
oauth2 and SigV4 return None and stay on v1.
|
||||
"""
|
||||
if server.is_byok:
|
||||
return None # per-user BYOK source not migrated yet -> defer to v1 (any auth_type)
|
||||
if server.is_byok and server.auth_type != MCPAuth.oauth2_token_exchange:
|
||||
return None # per-user BYOK source not migrated yet -> defer to v1
|
||||
resource: Final = server.url or server.server_id
|
||||
auth_type: Final = server.auth_type
|
||||
match auth_type:
|
||||
|
|
@ -165,21 +167,9 @@ def _client_credentials_spec(server: MCPServer, resource: str) -> ServerSpec:
|
|||
)
|
||||
|
||||
|
||||
def _token_exchange_spec(server: MCPServer, resource: str) -> ServerSpec | None:
|
||||
"""Build a token_exchange (OBO) spec, or defer (None) when it is not OBO-configured.
|
||||
|
||||
An OBO server with ``client_id``/``client_secret`` is owned by the v2 arm even if the
|
||||
``token_exchange_endpoint``/``token_url`` is absent: a missing endpoint then fails closed (412) at
|
||||
the exchanger rather than silently deferring to v1 and connecting unauthenticated, since the
|
||||
gateway must not guess the IdP or fall back to a weaker source. Without client credentials there is
|
||||
nothing to own, so the server stays on v1 (parity-safe). ``profile`` selects the wire dialect
|
||||
(``rfc8693`` default, ``entra_obo`` for Microsoft Entra On-Behalf-Of); an unrecognized value
|
||||
normalizes to ``rfc8693`` so a bad config value cannot crash spec-building. ``audience`` is
|
||||
forwarded only when the operator set it; a missing one is omitted, not derived.
|
||||
"""
|
||||
def _token_exchange_spec(server: MCPServer, resource: str) -> ServerSpec:
|
||||
"""Keep declared OBO owned by the resolver, including incomplete client configuration."""
|
||||
endpoint: Final = server.token_exchange_endpoint or server.effective_token_url
|
||||
if not server.client_id or not server.client_secret:
|
||||
return None
|
||||
profile: Final[Literal["rfc8693", "entra_obo"]] = (
|
||||
"entra_obo" if server.token_exchange_profile == "entra_obo" else "rfc8693"
|
||||
)
|
||||
|
|
@ -193,7 +183,7 @@ def _token_exchange_spec(server: MCPServer, resource: str) -> ServerSpec | None:
|
|||
token_exchange_endpoint=endpoint,
|
||||
audience=server.audience,
|
||||
client_id=server.client_id,
|
||||
client_secret=SecretStr(server.client_secret),
|
||||
client_secret=SecretStr(server.client_secret) if server.client_secret else None,
|
||||
token_endpoint_auth_method=server.token_endpoint_auth_method,
|
||||
scopes=tuple(server.scopes or ()),
|
||||
),
|
||||
|
|
@ -397,3 +387,74 @@ def raise_token_exchange_challenge(
|
|||
detail="Unauthorized",
|
||||
headers={"WWW-Authenticate": www_authenticate},
|
||||
)
|
||||
|
||||
|
||||
_STATIC_MODES: Final = frozenset(
|
||||
(MCPAuth.api_key, MCPAuth.bearer_token, MCPAuth.basic, MCPAuth.token, MCPAuth.authorization)
|
||||
)
|
||||
|
||||
|
||||
def _usable_credential_value(auth_type: MCPAuthType, name: str, value: str) -> bool:
|
||||
if not value:
|
||||
return False
|
||||
if auth_type == MCPAuth.api_key and name != "authorization":
|
||||
return True
|
||||
if value.lower() in ("bearer", "basic", "token", "apikey"):
|
||||
return False
|
||||
if auth_type == MCPAuth.api_key:
|
||||
api_scheme: Final = value.split(None, 1)[0]
|
||||
if api_scheme.lower() in ("bearer", "token", "apikey"):
|
||||
api_credential: Final = strip_auth_scheme(value, api_scheme).strip()
|
||||
return api_credential.lower() != api_scheme.lower()
|
||||
if auth_type in (MCPAuth.bearer_token, MCPAuth.token):
|
||||
scheme: Final = "Bearer" if auth_type == MCPAuth.bearer_token else "token"
|
||||
credential: Final = strip_auth_scheme(value, scheme).strip()
|
||||
return bool(credential) and credential.lower() != scheme.lower()
|
||||
if auth_type == MCPAuth.basic:
|
||||
parts: Final = value.split(None, 1)
|
||||
if len(parts) != 2 or parts[0].lower() != "basic":
|
||||
return False
|
||||
try:
|
||||
decoded: Final = base64.b64decode(parts[1], validate=True).strip()
|
||||
return b":" in decoded
|
||||
except ValueError:
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
def validate_static_credential(
|
||||
auth_type: MCPAuthType,
|
||||
headers: Mapping[str, str],
|
||||
upstream_token_header: str | None = None,
|
||||
static_header_names: Iterable[str] = (),
|
||||
) -> Result[None, CredError]:
|
||||
if auth_type not in _STATIC_MODES:
|
||||
return Ok(None)
|
||||
default_slot: Final = "X-API-Key" if auth_type == MCPAuth.api_key else "Authorization"
|
||||
admin_chosen_slots: Final = tuple(static_header_names) if auth_type == MCPAuth.api_key else ()
|
||||
slots: Final = frozenset(
|
||||
name.lower()
|
||||
for name in (
|
||||
upstream_token_header or default_slot,
|
||||
default_slot,
|
||||
"Authorization",
|
||||
*admin_chosen_slots,
|
||||
)
|
||||
)
|
||||
values: Final = tuple((name.lower(), value.strip()) for name, value in headers.items() if name.lower() in slots)
|
||||
if any(_usable_credential_value(auth_type, name, value) for name, value in values):
|
||||
return Ok(None)
|
||||
return Error(CredError.of_misconfigured(f"{auth_type} requires a usable upstream credential"))
|
||||
|
||||
|
||||
async def prepare_mcp_client(server: MCPServer, client: MCPClient) -> MCPClient:
|
||||
if server.auth_type not in _STATIC_MODES or client.transport_type == MCPTransport.stdio:
|
||||
return client
|
||||
request: Final = await client.prepare_request_auth()
|
||||
match validate_static_credential(
|
||||
server.auth_type, request.headers, server.upstream_token_header, server.static_headers or ()
|
||||
):
|
||||
case Error(error):
|
||||
raise_public(error)
|
||||
case Ok():
|
||||
return client
|
||||
|
|
|
|||
|
|
@ -844,6 +844,9 @@ class LiteLLMRoutes(enum.Enum):
|
|||
)
|
||||
|
||||
self_managed_routes = [
|
||||
# update_team resolves proxy/org/team admin itself and filters team admins
|
||||
# through the team_admin_editable_team_fields setting
|
||||
"/team/update",
|
||||
"/team/member_add",
|
||||
"/team/member_delete",
|
||||
"/management/v1/teams/{team_id}/members/bulk_delete",
|
||||
|
|
@ -4467,6 +4470,29 @@ class TeamInfoMember(Member):
|
|||
user_alias: str | None = None
|
||||
|
||||
|
||||
class TeamEditUnrestricted(BaseModel):
|
||||
kind: Literal["unrestricted"] = "unrestricted"
|
||||
|
||||
|
||||
class TeamEditAsTeamAdmin(BaseModel):
|
||||
kind: Literal["team_admin"] = "team_admin"
|
||||
editable_fields: tuple[str, ...]
|
||||
|
||||
|
||||
class TeamEditAsTeamAdminDisabled(BaseModel):
|
||||
kind: Literal["team_admin_disabled"] = "team_admin_disabled"
|
||||
|
||||
|
||||
class TeamEditNone(BaseModel):
|
||||
kind: Literal["none"] = "none"
|
||||
|
||||
|
||||
TeamEditAccess = Annotated[
|
||||
TeamEditUnrestricted | TeamEditAsTeamAdmin | TeamEditAsTeamAdminDisabled | TeamEditNone,
|
||||
Field(discriminator="kind"),
|
||||
]
|
||||
|
||||
|
||||
class TeamInfoResponseObjectTeamTable(LiteLLM_TeamTable):
|
||||
members_with_roles: tuple[TeamInfoMember, ...] = ()
|
||||
team_member_budget_table: LiteLLM_BudgetTableFull | None = None
|
||||
|
|
@ -4479,6 +4505,7 @@ class TeamInfoResponseObjectTeamTable(LiteLLM_TeamTable):
|
|||
# None = no org or not a manager; [] or ["all-proxy-models"] = no ceiling.
|
||||
organization_models: list[str] | None = None
|
||||
model_max_budget_usage: Mapping[str, Mapping[str, object]] | None = None
|
||||
caller_edit_access: TeamEditAccess = Field(default_factory=TeamEditNone)
|
||||
|
||||
|
||||
class TeamInfoResponseObject(TypedDict):
|
||||
|
|
|
|||
|
|
@ -8,7 +8,6 @@ from fastapi import APIRouter, Depends, HTTPException, Request, Response
|
|||
from fastapi.responses import JSONResponse
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.anthropic_interface.exceptions import AnthropicErrorResponse, AnthropicExceptionMapping
|
||||
from litellm.integrations.custom_guardrail import ModifyResponseException
|
||||
from litellm.llms.anthropic.experimental_pass_through.context_management import (
|
||||
|
|
@ -22,13 +21,16 @@ from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
|||
from litellm.proxy.common_request_processing import (
|
||||
ProxyBaseLLMRequestProcessing,
|
||||
create_response,
|
||||
log_llm_api_exception,
|
||||
proxy_exception_from_http_exception,
|
||||
resolve_litellm_call_id,
|
||||
)
|
||||
from litellm.proxy.common_utils.http_parsing_utils import _read_request_body
|
||||
from litellm.proxy.common_utils.openai_error_payload import (
|
||||
error_status_code,
|
||||
openai_error_param,
|
||||
openai_error_type,
|
||||
with_litellm_call_id,
|
||||
)
|
||||
from litellm.types.utils import TokenCountResponse
|
||||
|
||||
|
|
@ -218,10 +220,12 @@ async def anthropic_response(
|
|||
await proxy_logging_obj.post_call_failure_hook(
|
||||
user_api_key_dict=user_api_key_dict, original_exception=e, request_data=base_llm_response_processor.data
|
||||
)
|
||||
verbose_proxy_logger.exception("litellm.proxy.proxy_server.anthropic_response(): Exception occured - %s", e)
|
||||
log_llm_api_exception(e, base_llm_response_processor.litellm_call_id)
|
||||
|
||||
if isinstance(e, ProxyException):
|
||||
return _anthropic_error_json_response(e, request)
|
||||
return _anthropic_error_json_response(
|
||||
with_litellm_call_id(e, base_llm_response_processor.litellm_call_id), request
|
||||
)
|
||||
|
||||
# Extract model_id from request metadata (same as success path)
|
||||
litellm_metadata: Final = data.get("litellm_metadata", {}) or {}
|
||||
|
|
@ -231,7 +235,7 @@ async def anthropic_response(
|
|||
# Get headers
|
||||
headers: Final = ProxyBaseLLMRequestProcessing.get_custom_headers(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
call_id=data.get("litellm_call_id", ""),
|
||||
call_id=base_llm_response_processor.litellm_call_id,
|
||||
model_id=model_id,
|
||||
version=version,
|
||||
response_cost=0,
|
||||
|
|
@ -288,6 +292,7 @@ async def count_tokens(
|
|||
"""
|
||||
from litellm.proxy.proxy_server import token_counter as internal_token_counter
|
||||
|
||||
litellm_call_id: Final = resolve_litellm_call_id(request.headers.get("x-litellm-call-id"))
|
||||
try:
|
||||
request_data: Final = await _read_request_body(request=request)
|
||||
data: Final[dict] = {**request_data}
|
||||
|
|
@ -339,7 +344,7 @@ async def count_tokens(
|
|||
detail=detail,
|
||||
)
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.exception("litellm.proxy.anthropic_endpoints.count_tokens(): Exception occurred - %s", e)
|
||||
log_llm_api_exception(e, litellm_call_id)
|
||||
raise HTTPException(status_code=500, detail={"error": f"Internal server error: {e}"})
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -856,6 +856,16 @@ BUDGET_ENFORCED_SIDE_EFFECT_ROUTES: Final = frozenset(
|
|||
)
|
||||
|
||||
|
||||
def route_skips_budget_checks(route: str) -> bool:
|
||||
return route not in BUDGET_ENFORCED_SIDE_EFFECT_ROUTES and (
|
||||
route in MODEL_DISCOVERY_ROUTES or not RouteChecks.is_llm_api_route(route=route)
|
||||
)
|
||||
|
||||
|
||||
def request_skips_budget_checks(route: str, model: str | list[str] | None, llm_router: Router | None) -> bool:
|
||||
return route_skips_budget_checks(route=route) or _is_model_cost_zero(model=model, llm_router=llm_router)
|
||||
|
||||
|
||||
async def common_checks(
|
||||
request_body: dict,
|
||||
team_object: LiteLLM_TeamTable | None,
|
||||
|
|
@ -903,10 +913,7 @@ async def common_checks(
|
|||
team_id=valid_token.team_id if valid_token is not None else None,
|
||||
)
|
||||
|
||||
skip_all_budget_checks: Final = skip_budget_checks or (
|
||||
route not in BUDGET_ENFORCED_SIDE_EFFECT_ROUTES
|
||||
and (route in MODEL_DISCOVERY_ROUTES or not RouteChecks.is_llm_api_route(route=route))
|
||||
)
|
||||
skip_all_budget_checks: Final = skip_budget_checks or route_skips_budget_checks(route=route)
|
||||
|
||||
membership_user_id: Final = (
|
||||
valid_token.user_id if valid_token is not None and (bool(_model) or not skip_all_budget_checks) else None
|
||||
|
|
@ -2104,7 +2111,7 @@ async def _fetch_uncached_tags(
|
|||
|
||||
@log_db_metrics
|
||||
async def get_tag_objects_batch(
|
||||
tag_names: list[str],
|
||||
tag_names: Sequence[str],
|
||||
prisma_client: PrismaClient | None,
|
||||
user_api_key_cache: UserApiKeyCache,
|
||||
parent_otel_span: Span | None = None,
|
||||
|
|
@ -5863,15 +5870,25 @@ async def _tag_max_budget_check(
|
|||
"""
|
||||
from litellm.proxy.common_utils.http_parsing_utils import get_tags_from_request_body
|
||||
|
||||
if prisma_client is None:
|
||||
await tag_max_budget_check_for_tags(
|
||||
tags=get_tags_from_request_body(request_body=request_body),
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
valid_token=valid_token,
|
||||
)
|
||||
|
||||
|
||||
async def tag_max_budget_check_for_tags(
|
||||
tags: Sequence[str],
|
||||
prisma_client: PrismaClient | None,
|
||||
user_api_key_cache: UserApiKeyCache,
|
||||
proxy_logging_obj: ProxyLogging,
|
||||
valid_token: UserAPIKeyAuth | None,
|
||||
) -> None:
|
||||
if prisma_client is None or not tags:
|
||||
return
|
||||
|
||||
# Get tags from request metadata
|
||||
tags: Final = get_tags_from_request_body(request_body=request_body)
|
||||
if not tags:
|
||||
return
|
||||
|
||||
# Batch fetch all tags in one go
|
||||
tag_objects: Final = await get_tag_objects_batch(
|
||||
tag_names=tags,
|
||||
prisma_client=prisma_client,
|
||||
|
|
|
|||
|
|
@ -7,6 +7,7 @@
|
|||
import asyncio
|
||||
import os
|
||||
from collections.abc import Mapping
|
||||
from types import MappingProxyType
|
||||
from typing import Any, Final, cast
|
||||
|
||||
from fastapi import APIRouter, Depends, HTTPException, Path, Request, Response
|
||||
|
|
@ -17,7 +18,11 @@ from litellm.batches.main import CancelBatchRequest, RetrieveBatchRequest
|
|||
from litellm.proxy._types import *
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
from litellm.proxy.batches_endpoints.common_utils import validate_batch_list_limit
|
||||
from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
|
||||
from litellm.proxy.common_request_processing import (
|
||||
ProxyBaseLLMRequestProcessing,
|
||||
log_llm_api_exception,
|
||||
request_litellm_call_id,
|
||||
)
|
||||
from litellm.proxy.common_utils.callback_utils import sanitize_openai_provider_metadata
|
||||
from litellm.proxy.common_utils.http_parsing_utils import _read_request_body
|
||||
from litellm.proxy.common_utils.openai_endpoint_utils import (
|
||||
|
|
@ -383,8 +388,9 @@ async def create_batch(
|
|||
await proxy_logging_obj.post_call_failure_hook(
|
||||
user_api_key_dict=user_api_key_dict, original_exception=e, request_data=data
|
||||
)
|
||||
verbose_proxy_logger.exception("litellm.proxy.proxy_server.create_batch(): Exception occured - %s", e)
|
||||
raise handle_exception_on_proxy(e)
|
||||
litellm_call_id: Final = request_litellm_call_id(data)
|
||||
log_llm_api_exception(e, litellm_call_id)
|
||||
raise handle_exception_on_proxy(e, litellm_call_id)
|
||||
|
||||
|
||||
@router.get(
|
||||
|
|
@ -674,8 +680,9 @@ async def retrieve_batch(
|
|||
await proxy_logging_obj.post_call_failure_hook(
|
||||
user_api_key_dict=user_api_key_dict, original_exception=e, request_data=data
|
||||
)
|
||||
verbose_proxy_logger.exception("litellm.proxy.proxy_server.retrieve_batch(): Exception occured - %s", e)
|
||||
raise handle_exception_on_proxy(e)
|
||||
litellm_call_id: Final = request_litellm_call_id(data)
|
||||
log_llm_api_exception(e, litellm_call_id)
|
||||
raise handle_exception_on_proxy(e, litellm_call_id)
|
||||
|
||||
|
||||
@router.get(
|
||||
|
|
@ -725,6 +732,7 @@ async def list_batches(
|
|||
)
|
||||
|
||||
verbose_proxy_logger.debug("GET /v1/batches after=%s limit=%s", after, limit)
|
||||
data: Mapping[str, object] = MappingProxyType({})
|
||||
try:
|
||||
if llm_router is None:
|
||||
raise HTTPException(
|
||||
|
|
@ -854,10 +862,11 @@ async def list_batches(
|
|||
await proxy_logging_obj.post_call_failure_hook(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
original_exception=e,
|
||||
request_data={"after": after, "limit": limit},
|
||||
request_data={**data, "after": after, "limit": limit},
|
||||
)
|
||||
verbose_proxy_logger.error("litellm.proxy.proxy_server.retrieve_batch(): Exception occured - %s", e)
|
||||
raise handle_exception_on_proxy(e)
|
||||
litellm_call_id: Final = request_litellm_call_id(data)
|
||||
log_llm_api_exception(e, litellm_call_id)
|
||||
raise handle_exception_on_proxy(e, litellm_call_id)
|
||||
|
||||
|
||||
@router.post(
|
||||
|
|
@ -1079,8 +1088,9 @@ async def cancel_batch(
|
|||
await proxy_logging_obj.post_call_failure_hook(
|
||||
user_api_key_dict=user_api_key_dict, original_exception=e, request_data=data
|
||||
)
|
||||
verbose_proxy_logger.exception("litellm.proxy.proxy_server.create_batch(): Exception occured - %s", e)
|
||||
raise handle_exception_on_proxy(e)
|
||||
litellm_call_id: Final = request_litellm_call_id(data)
|
||||
log_llm_api_exception(e, litellm_call_id)
|
||||
raise handle_exception_on_proxy(e, litellm_call_id)
|
||||
|
||||
|
||||
######################################################################
|
||||
|
|
|
|||
|
|
@ -7,14 +7,25 @@ from collections.abc import AsyncGenerator, Awaitable, Callable, Mapping, Sequen
|
|||
from datetime import datetime
|
||||
from functools import lru_cache
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Any, Final, Literal, NamedTuple, Protocol, TypeAlias, TypeVar, overload
|
||||
from typing import (
|
||||
TYPE_CHECKING,
|
||||
Any,
|
||||
Final,
|
||||
Literal,
|
||||
NamedTuple,
|
||||
Protocol,
|
||||
TypeAlias,
|
||||
TypeVar,
|
||||
overload,
|
||||
runtime_checkable,
|
||||
)
|
||||
|
||||
import anyio
|
||||
import httpx
|
||||
import orjson
|
||||
from fastapi import HTTPException, Request, status
|
||||
from fastapi.responses import JSONResponse, Response, StreamingResponse
|
||||
from pydantic import ValidationError
|
||||
from pydantic import TypeAdapter, ValidationError
|
||||
from starlette.types import Receive, Scope, Send
|
||||
|
||||
import litellm
|
||||
|
|
@ -34,7 +45,11 @@ from litellm.constants import (
|
|||
UNSAFE_PROXY_RESPONSE_HEADERS,
|
||||
)
|
||||
from litellm.integrations.custom_guardrail import CustomGuardrail
|
||||
from litellm.litellm_core_utils.core_helpers import get_or_create_metadata_bucket, is_expected_client_error
|
||||
from litellm.litellm_core_utils.core_helpers import (
|
||||
get_or_create_metadata_bucket,
|
||||
independent_snapshot,
|
||||
is_expected_client_error,
|
||||
)
|
||||
from litellm.litellm_core_utils.dd_tracing import NullTracer, tracer
|
||||
from litellm.litellm_core_utils.get_supported_openai_params import (
|
||||
get_supported_openai_params,
|
||||
|
|
@ -49,14 +64,21 @@ from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
|
|||
from litellm.litellm_core_utils.streaming_handler import (
|
||||
backfill_missing_cache_usage_fields,
|
||||
)
|
||||
from litellm.proxy._types import ProxyException, UserAPIKeyAuth
|
||||
from litellm.proxy.auth.auth_checks import can_key_call_resolved_model
|
||||
from litellm.proxy.auth.auth_utils import check_response_size_is_safe
|
||||
from litellm.proxy._types import ProxyErrorTypes, ProxyException, UserAPIKeyAuth
|
||||
from litellm.proxy.auth.auth_checks import (
|
||||
can_key_call_resolved_model,
|
||||
request_skips_budget_checks,
|
||||
tag_max_budget_check_for_tags,
|
||||
)
|
||||
from litellm.proxy.auth.auth_utils import check_response_size_is_safe, get_request_route
|
||||
from litellm.proxy.common_utils.callback_utils import (
|
||||
get_logging_caching_headers,
|
||||
get_remaining_tokens_and_requests_from_request_data,
|
||||
)
|
||||
from litellm.proxy.common_utils.http_parsing_utils import get_client_requested_model
|
||||
from litellm.proxy.common_utils.http_parsing_utils import (
|
||||
get_client_requested_model,
|
||||
get_tags_from_request_body,
|
||||
)
|
||||
from litellm.proxy.common_utils.openai_error_payload import (
|
||||
attribute_of,
|
||||
error_status_code,
|
||||
|
|
@ -643,6 +665,48 @@ async def _resolve_per_request_model_group_alias(
|
|||
return target
|
||||
|
||||
|
||||
_REQUEST_MODEL: Final[TypeAdapter[str | list[str] | None]] = TypeAdapter(str | list[str] | None)
|
||||
|
||||
|
||||
def _request_model(data: Mapping[str, object]) -> str | list[str] | None:
|
||||
try:
|
||||
return _REQUEST_MODEL.validate_python(data.get("model"), strict=True)
|
||||
except ValidationError:
|
||||
return None
|
||||
|
||||
|
||||
async def _enforce_guardrail_added_tag_budgets(
|
||||
data: Mapping[str, object],
|
||||
tags_before_guardrails: frozenset[str],
|
||||
route: str,
|
||||
llm_router: Router | None,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
proxy_logging_obj: ProxyLogging,
|
||||
) -> None:
|
||||
added_tags: Final = tuple(
|
||||
tag for tag in get_tags_from_request_body(request_body=data) if tag not in tags_before_guardrails
|
||||
)
|
||||
if not added_tags or request_skips_budget_checks(route=route, model=_request_model(data), llm_router=llm_router):
|
||||
return
|
||||
from litellm.proxy.proxy_server import prisma_client, user_api_key_cache
|
||||
|
||||
try:
|
||||
await tag_max_budget_check_for_tags(
|
||||
tags=added_tags,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
valid_token=user_api_key_dict,
|
||||
)
|
||||
except litellm.BudgetExceededError as e:
|
||||
raise ProxyException(
|
||||
message=e.message,
|
||||
type=ProxyErrorTypes.budget_exceeded,
|
||||
param=None,
|
||||
code=e.status_code,
|
||||
) from e
|
||||
|
||||
|
||||
async def _parse_event_data_for_error(event_line: str | bytes) -> int | None:
|
||||
"""Parses an event line and returns an error code if present, else None."""
|
||||
event_line = event_line.decode("utf-8") if isinstance(event_line, bytes) else event_line
|
||||
|
|
@ -1452,7 +1516,19 @@ def _has_attribute_error_in_chain(exc: Exception) -> bool:
|
|||
_CLIENT_DISCONNECT_DETAIL: Final = "Client disconnected the request"
|
||||
|
||||
|
||||
def _log_llm_api_exception(e: Exception, litellm_call_id: str | None) -> None:
|
||||
@runtime_checkable
|
||||
class _CarriesLitellmCallId(Protocol):
|
||||
litellm_call_id: str | None
|
||||
|
||||
|
||||
def request_litellm_call_id(data: Mapping[str, object]) -> str | None:
|
||||
logging_obj: Final = data.get("litellm_logging_obj")
|
||||
logged_id: Final = logging_obj.litellm_call_id if isinstance(logging_obj, _CarriesLitellmCallId) else None
|
||||
call_id: Final = logged_id or data.get("litellm_call_id")
|
||||
return call_id if isinstance(call_id, str) else None
|
||||
|
||||
|
||||
def log_llm_api_exception(e: Exception, litellm_call_id: str | None) -> None:
|
||||
if getattr(e, "status_code", None) == 499 and getattr(e, "detail", None) == _CLIENT_DISCONNECT_DETAIL:
|
||||
verbose_proxy_logger.info(
|
||||
"litellm.proxy.proxy_server._handle_llm_api_exception(): client disconnected, "
|
||||
|
|
@ -1531,6 +1607,11 @@ def _timing_values(
|
|||
class ProxyBaseLLMRequestProcessing:
|
||||
def __init__(self, data: dict):
|
||||
self.data = data
|
||||
self._tags_before_guardrails: frozenset[str] | None = None
|
||||
|
||||
@property
|
||||
def litellm_call_id(self) -> str | None:
|
||||
return request_litellm_call_id(self.data)
|
||||
|
||||
@staticmethod
|
||||
def _merge_passthrough_streaming_headers(
|
||||
|
|
@ -2020,11 +2101,21 @@ class ProxyBaseLLMRequestProcessing:
|
|||
# to run below.
|
||||
await _arm_auto_router_compression(data=self.data, llm_router=llm_router)
|
||||
|
||||
if self._tags_before_guardrails is None:
|
||||
self._tags_before_guardrails = frozenset(get_tags_from_request_body(request_body=self.data))
|
||||
self.data = await proxy_logging_obj.pre_call_hook(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
data=self.data,
|
||||
call_type=route_type,
|
||||
)
|
||||
await _enforce_guardrail_added_tag_budgets(
|
||||
data=self.data,
|
||||
tags_before_guardrails=self._tags_before_guardrails,
|
||||
route=get_request_route(request=request),
|
||||
llm_router=llm_router,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
if route_type == "aget_responses":
|
||||
attach_post_call_pipelines_to_retrieval(
|
||||
data=self.data,
|
||||
|
|
@ -2062,6 +2153,13 @@ class ProxyBaseLLMRequestProcessing:
|
|||
) -> tuple[dict, LiteLLMLoggingObj]:
|
||||
from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError
|
||||
|
||||
configured_fallbacks: Final = (
|
||||
self._configured_fallbacks(llm_router=llm_router, user_api_key_dict=user_api_key_dict)
|
||||
if llm_router is not None and not self.data.get("disable_fallbacks")
|
||||
else None
|
||||
)
|
||||
pristine: Final = independent_snapshot(self.data) if configured_fallbacks else None
|
||||
|
||||
try:
|
||||
return await self.common_processing_pre_call_logic(
|
||||
request=request,
|
||||
|
|
@ -2080,14 +2178,19 @@ class ProxyBaseLLMRequestProcessing:
|
|||
llm_router=llm_router,
|
||||
)
|
||||
except ProxyRateLimitError as original_exc:
|
||||
original_model: Final = self.data.get("model")
|
||||
if not original_model or not llm_router or self.data.get("disable_fallbacks"):
|
||||
rate_limited_data: Final = self.data
|
||||
original_model: Final = rate_limited_data.get("model")
|
||||
if (
|
||||
pristine is None
|
||||
or not configured_fallbacks
|
||||
or rate_limited_data.get("disable_fallbacks")
|
||||
or not isinstance(original_model, str)
|
||||
):
|
||||
raise
|
||||
|
||||
fallback_models: Final = self._resolve_fallback_models(
|
||||
model=original_model,
|
||||
llm_router=llm_router,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
fallbacks=configured_fallbacks,
|
||||
)
|
||||
if not fallback_models:
|
||||
raise
|
||||
|
|
@ -2102,6 +2205,7 @@ class ProxyBaseLLMRequestProcessing:
|
|||
for fallback_model in fallback_models:
|
||||
if fallback_model == original_model:
|
||||
continue
|
||||
self.data = independent_snapshot(pristine)
|
||||
self.data["model"] = fallback_model
|
||||
try:
|
||||
return await self.common_processing_pre_call_logic(
|
||||
|
|
@ -2123,39 +2227,30 @@ class ProxyBaseLLMRequestProcessing:
|
|||
except ProxyRateLimitError:
|
||||
continue
|
||||
except BaseException:
|
||||
self.data["model"] = original_model
|
||||
self.data = rate_limited_data
|
||||
raise
|
||||
|
||||
self.data["model"] = original_model
|
||||
self.data = rate_limited_data
|
||||
raise original_exc
|
||||
|
||||
def _resolve_fallback_models(
|
||||
self,
|
||||
model: str,
|
||||
llm_router: Router,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
) -> list | None:
|
||||
from litellm.router_utils.fallback_event_handlers import get_fallback_model_group
|
||||
|
||||
fallbacks = None
|
||||
|
||||
@staticmethod
|
||||
def _configured_fallbacks(llm_router: Router, user_api_key_dict: UserAPIKeyAuth) -> list | None:
|
||||
key_router_settings: Final = user_api_key_dict.router_settings
|
||||
if isinstance(key_router_settings, dict) and "fallbacks" in key_router_settings:
|
||||
fallbacks = key_router_settings["fallbacks"]
|
||||
key_fallbacks: Final = key_router_settings.get("fallbacks") if isinstance(key_router_settings, dict) else None
|
||||
fallbacks: Final = key_fallbacks if key_fallbacks is not None else llm_router.fallbacks
|
||||
return fallbacks if isinstance(fallbacks, list) and fallbacks else None
|
||||
|
||||
if fallbacks is None:
|
||||
fallbacks = llm_router.fallbacks
|
||||
|
||||
if not fallbacks:
|
||||
return None
|
||||
@staticmethod
|
||||
def _resolve_fallback_models(model: str, fallbacks: list) -> list | None:
|
||||
from litellm.router_utils.fallback_event_handlers import get_fallback_model_group
|
||||
|
||||
fallback_model_group, generic_fallback_idx = get_fallback_model_group(
|
||||
fallbacks=fallbacks,
|
||||
model_group=model,
|
||||
)
|
||||
if fallback_model_group is None and generic_fallback_idx is not None:
|
||||
fallback_model_group = fallbacks[generic_fallback_idx]["*"]
|
||||
return fallback_model_group
|
||||
if fallback_model_group is not None:
|
||||
return fallback_model_group
|
||||
return fallbacks[generic_fallback_idx]["*"] if generic_fallback_idx is not None else None
|
||||
|
||||
@staticmethod
|
||||
def _get_model_id_from_response(hidden_params: Mapping[str, object], data: Mapping[str, object]) -> str:
|
||||
|
|
@ -3429,11 +3524,7 @@ class ProxyBaseLLMRequestProcessing:
|
|||
version: str | None = None,
|
||||
):
|
||||
"""Raises ProxyException (OpenAI API compatible) if an exception is raised"""
|
||||
logging_obj: Final[LiteLLMLoggingObj | None] = self.data.get("litellm_logging_obj", None)
|
||||
_log_llm_api_exception(
|
||||
e,
|
||||
(logging_obj.litellm_call_id if logging_obj is not None else None) or self.data.get("litellm_call_id"),
|
||||
)
|
||||
log_llm_api_exception(e, self.litellm_call_id)
|
||||
# Allow callbacks to transform the error response
|
||||
transformed_exception: Final = await proxy_logging_obj.post_call_failure_hook(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
|
|
@ -3463,9 +3554,7 @@ class ProxyBaseLLMRequestProcessing:
|
|||
|
||||
custom_headers: Final = ProxyBaseLLMRequestProcessing.get_custom_headers(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
call_id=(
|
||||
_litellm_logging_obj.litellm_call_id if _litellm_logging_obj else self.data.get("litellm_call_id")
|
||||
),
|
||||
call_id=self.litellm_call_id,
|
||||
model_id=model_id,
|
||||
version=version,
|
||||
response_cost=0,
|
||||
|
|
|
|||
|
|
@ -9,6 +9,9 @@ from typing import Final
|
|||
from fastapi import status
|
||||
|
||||
from litellm.constants import STRINGIFIED_NONE
|
||||
from litellm.proxy._types import ProxyException
|
||||
|
||||
LITELLM_CALL_ID_HEADER: Final = "x-litellm-call-id"
|
||||
|
||||
_OPENAI_ERROR_TYPE_BY_STATUS: Final[Mapping[int, str]] = MappingProxyType(
|
||||
{
|
||||
|
|
@ -52,3 +55,23 @@ def openai_error_param(exc: object) -> str | None:
|
|||
serializes as JSON ``null``."""
|
||||
carried: Final = attribute_of(exc, "param")
|
||||
return carried if isinstance(carried, str) and carried != STRINGIFIED_NONE else None
|
||||
|
||||
|
||||
def litellm_call_id_headers(litellm_call_id: str | None) -> dict[str, str] | None: # mutable-ok: ProxyException.headers
|
||||
if litellm_call_id is None:
|
||||
return None
|
||||
return {LITELLM_CALL_ID_HEADER: litellm_call_id} # mutable-ok: ProxyException mutates its headers dict
|
||||
|
||||
|
||||
def with_litellm_call_id(exc: ProxyException, litellm_call_id: str | None) -> ProxyException:
|
||||
"""The same error object, answering with ``x-litellm-call-id`` when it was raised without one."""
|
||||
if litellm_call_id is not None:
|
||||
exc.headers.setdefault(LITELLM_CALL_ID_HEADER, litellm_call_id)
|
||||
return exc
|
||||
|
||||
|
||||
def headers_with_litellm_call_id(headers: Mapping[str, str] | None, litellm_call_id: str) -> Mapping[str, str]:
|
||||
"""``headers`` plus ``x-litellm-call-id``, keeping the value they already carry under that name."""
|
||||
if headers is None:
|
||||
return MappingProxyType({LITELLM_CALL_ID_HEADER: litellm_call_id})
|
||||
return MappingProxyType({LITELLM_CALL_ID_HEADER: litellm_call_id, **headers})
|
||||
|
|
|
|||
|
|
@ -11,7 +11,7 @@ exception types:
|
|||
an upstream LLM provider returns 429.
|
||||
* :class:`fastapi.HTTPException` (status 429) — raised directly by proxy hooks
|
||||
such as ``parallel_request_limiter``, ``dynamic_rate_limiter``,
|
||||
``batch_rate_limiter``, ``max_budget_limiter``, ``max_iterations_limiter``,
|
||||
``batch_rate_limiter``, ``max_iterations_limiter``,
|
||||
etc.
|
||||
* :class:`litellm.llms.base_llm.chat.transformation.BaseLLMException` (status
|
||||
429) — raised by some provider transports.
|
||||
|
|
|
|||
|
|
@ -257,8 +257,8 @@ class DBSpendUpdateWriter:
|
|||
# Completion object fields
|
||||
kwargs: dict | None,
|
||||
completion_response: object,
|
||||
start_time: datetime | None,
|
||||
end_time: datetime | None,
|
||||
start_time: datetime,
|
||||
end_time: datetime,
|
||||
response_cost: float | None,
|
||||
) -> bool:
|
||||
"""Record the request's spend, answering whether its cost still needs charging.
|
||||
|
|
@ -299,6 +299,7 @@ class DBSpendUpdateWriter:
|
|||
response_obj=completion_response,
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
llm_router=get_llm_router(),
|
||||
)
|
||||
payload["spend"] = response_cost or 0.0
|
||||
if isinstance(payload["startTime"], datetime):
|
||||
|
|
|
|||
|
|
@ -232,7 +232,8 @@ class AktoGuardrail(CustomGuardrail):
|
|||
"""
|
||||
request_path: Final = self.extract_request_path(request_data)
|
||||
request_headers: Final = self.build_request_headers(request_data)
|
||||
request_body: Final = self.build_request_body(inputs, request_data)
|
||||
request_inputs: Final = GenericGuardrailAPIInputs(model=inputs.get("model")) if include_response else inputs
|
||||
request_body: Final = self.build_request_body(request_inputs, request_data)
|
||||
tag: Final = self.build_tag_metadata(request_data)
|
||||
|
||||
response_payload = json.dumps({}) # Empty body wrapper when no response yet
|
||||
|
|
|
|||
|
|
@ -425,10 +425,7 @@ class CrowdStrikeAIDRHandler(CustomGuardrail):
|
|||
|
||||
def _build_guard_input_for_response(self, inputs: GenericGuardrailAPIInputs) -> _GuardInput:
|
||||
output_texts: Final[list[str]] = inputs.get("texts", [])
|
||||
return _GuardInput(
|
||||
messages=[_Message(role="assistant", content=text) for text in output_texts],
|
||||
tools=inputs.get("tools", []),
|
||||
)
|
||||
return _GuardInput(messages=[_Message(role="assistant", content=text) for text in output_texts], tools=[])
|
||||
|
||||
def _extract_transformed_texts(self, guard_output: _GuardInput, num_assistant_messages: int) -> list[str]:
|
||||
tail: Final = guard_output.messages[-num_assistant_messages:] if num_assistant_messages > 0 else []
|
||||
|
|
|
|||
|
|
@ -286,7 +286,7 @@ class HiddenlayerGuardrail(CustomGuardrail):
|
|||
hl_request_metadata["requester_id"] = headers.get("hl-requester-id") or "LiteLLM"
|
||||
project_id: Final = headers.get("hl-project-id")
|
||||
|
||||
if scan_params := inputs.get("structured_messages"):
|
||||
if input_type == "request" and (scan_params := inputs.get("structured_messages")):
|
||||
last_msg: Final = scan_params[-1]
|
||||
result: _HiddenlayerResponse = await self._call_hiddenlayer(
|
||||
project_id,
|
||||
|
|
|
|||
|
|
@ -197,7 +197,7 @@ class OpenAIModerationGuardrail(OpenAIGuardrailBase, CustomGuardrail):
|
|||
text_to_moderate: str | None = None
|
||||
|
||||
# Prefer structured_messages if available (has role context)
|
||||
if structured_messages := inputs.get("structured_messages"):
|
||||
if input_type == "request" and (structured_messages := inputs.get("structured_messages")):
|
||||
text_to_moderate = self.get_user_prompt(structured_messages)
|
||||
|
||||
# Fall back to texts
|
||||
|
|
|
|||
|
|
@ -20,6 +20,7 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail"
|
|||
guardrail_name=guardrail.get("guardrail_name", ""),
|
||||
event_hook=litellm_params.mode,
|
||||
default_on=litellm_params.default_on,
|
||||
streaming_transform_mode=getattr(litellm_params, "streaming_transform_mode", None),
|
||||
file_sanitization_fail_open=getattr(litellm_params, "file_sanitization_fail_open", None),
|
||||
block_on_file_modify=getattr(litellm_params, "block_on_file_modify", None),
|
||||
)
|
||||
|
|
|
|||
|
|
@ -38,6 +38,11 @@ class PromptSecurityGuardrailMissingSecrets(Exception):
|
|||
pass
|
||||
|
||||
|
||||
def _modified_or_original(text: str, verdict: "_ProtectVerdict") -> str:
|
||||
modified_text: Final = verdict.get("modified_text") if verdict.get("action") == "modify" else None
|
||||
return text if modified_text is None else modified_text
|
||||
|
||||
|
||||
def _inputs_with_structured_messages(
|
||||
inputs: GenericGuardrailAPIInputs, rewritten_messages: Sequence[AllMessageValues] | None
|
||||
) -> GenericGuardrailAPIInputs:
|
||||
|
|
@ -119,6 +124,7 @@ class PromptSecurityGuardrail(CustomGuardrail):
|
|||
user: str | None = None,
|
||||
system_prompt: str | None = None,
|
||||
check_tool_results: bool | None = None,
|
||||
streaming_transform_mode: Literal["block_only", "incremental_diff"] | None = None,
|
||||
file_sanitization_timeout: float = _SANITIZE_FILE_FAIL_OPEN_TIMEOUT_SECONDS,
|
||||
file_sanitization_fail_open: bool | None = None,
|
||||
block_on_file_modify: bool | None = None,
|
||||
|
|
@ -148,6 +154,10 @@ class PromptSecurityGuardrail(CustomGuardrail):
|
|||
)
|
||||
raise PromptSecurityGuardrailMissingSecrets(msg)
|
||||
|
||||
self.streaming_transform_mode: Literal["block_only", "incremental_diff"] = (
|
||||
"block_only" if streaming_transform_mode is None else streaming_transform_mode
|
||||
)
|
||||
|
||||
# Configuration for file sanitization
|
||||
self.max_poll_attempts = 30 # Maximum number of polling attempts
|
||||
self.poll_interval = 2 # Seconds between polling attempts
|
||||
|
|
@ -342,16 +352,46 @@ class PromptSecurityGuardrail(CustomGuardrail):
|
|||
texts: list[str],
|
||||
user_api_key_alias: str | None,
|
||||
) -> GenericGuardrailAPIInputs:
|
||||
"""Handle response-side guardrail checks."""
|
||||
"""Handle response-side guardrail checks, one protect verdict per text.
|
||||
|
||||
Prompt Security rewrites a single string, so texts from several choices must be scanned separately
|
||||
or one ``modified_text`` cannot be mapped back onto the choice it came from. It also returns no span
|
||||
offsets, so on a stream every text is held back in full until the final verdict: a value the vendor
|
||||
redacts later may start anywhere in text that looked clean so far, and streamed bytes cannot be recalled.
|
||||
"""
|
||||
if not texts:
|
||||
return inputs
|
||||
|
||||
# Combine all texts for response checking
|
||||
combined_text: Final = "\n".join(texts)
|
||||
verdicts: Final = await asyncio.gather(
|
||||
*(self._protect_response_text(text, user_api_key_alias) for text in texts)
|
||||
)
|
||||
violations: Final = tuple(
|
||||
violation
|
||||
for verdict in verdicts
|
||||
if verdict.get("action") == "block"
|
||||
for violation in verdict.get("violations", ())
|
||||
)
|
||||
if any(verdict.get("action") == "block" for verdict in verdicts):
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail="Blocked by Prompt Security, Violations: " + ", ".join(violations),
|
||||
)
|
||||
returned_texts: Final = [ # mutable-ok: GenericGuardrailAPIInputs.texts is list[str]
|
||||
_modified_or_original(text, verdict) for text, verdict in zip(texts, verdicts, strict=True)
|
||||
]
|
||||
patched: Final[GenericGuardrailAPIInputs] = {
|
||||
**inputs,
|
||||
"texts": returned_texts,
|
||||
"stream_holdback_chars": [ # mutable-ok: GenericGuardrailAPIInputs.stream_holdback_chars is list[int]
|
||||
len(text) for text in returned_texts
|
||||
],
|
||||
}
|
||||
return patched
|
||||
|
||||
async def _protect_response_text(self, text: str, user_api_key_alias: str | None) -> _ProtectVerdict:
|
||||
headers: Final = self._build_headers(user_api_key_alias)
|
||||
payload: Final = {
|
||||
"response": combined_text,
|
||||
"response": text,
|
||||
"user": user_api_key_alias or self.user,
|
||||
"system_prompt": self.system_prompt,
|
||||
}
|
||||
|
|
@ -360,7 +400,7 @@ class PromptSecurityGuardrail(CustomGuardrail):
|
|||
method="POST",
|
||||
url=f"{self.api_base}/api/protect",
|
||||
headers=headers,
|
||||
payload={"response_length": len(combined_text)},
|
||||
payload={"response_length": len(text)},
|
||||
)
|
||||
|
||||
response: Final = await self.async_handler.post(
|
||||
|
|
@ -377,26 +417,8 @@ class PromptSecurityGuardrail(CustomGuardrail):
|
|||
payload={"result": res.get("result")},
|
||||
)
|
||||
|
||||
result: Final = res.get("result", {}).get("response", {})
|
||||
if result is None:
|
||||
return inputs
|
||||
|
||||
action: Final = result.get("action")
|
||||
violations: Final = result.get("violations", [])
|
||||
|
||||
if action == "block":
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail="Blocked by Prompt Security, Violations: " + ", ".join(violations),
|
||||
)
|
||||
elif action == "modify":
|
||||
modified_text: Final = result.get("modified_text")
|
||||
if modified_text is not None:
|
||||
# If we combined multiple texts, return the modified version as single text
|
||||
# The framework will handle distributing it back
|
||||
inputs["texts"] = [modified_text]
|
||||
|
||||
return inputs
|
||||
verdict: Final = res.get("result", {}).get("response", {})
|
||||
return {} if verdict is None else verdict
|
||||
|
||||
def _extract_texts_from_messages(self, messages: Sequence[Mapping[str, object]]) -> list[str]:
|
||||
return [text for message in messages for text in message_slot_texts(message)]
|
||||
|
|
|
|||
|
|
@ -121,7 +121,7 @@ class PromptGuardGuardrail(CustomGuardrail):
|
|||
) -> GenericGuardrailAPIInputs:
|
||||
texts: Final = inputs.get("texts", [])
|
||||
images: Final = inputs.get("images", [])
|
||||
structured_messages: Final = inputs.get("structured_messages", [])
|
||||
structured_messages: Final = inputs.get("structured_messages") if input_type == "request" else None
|
||||
model: Final = inputs.get("model")
|
||||
|
||||
if structured_messages:
|
||||
|
|
|
|||
|
|
@ -452,7 +452,7 @@ class QualifireGuardrail(CustomGuardrail):
|
|||
dynamic_params: Final = self.get_guardrail_dynamic_request_body_params(request_data=request_data)
|
||||
|
||||
# Extract messages from structured_messages or request_data
|
||||
messages: list[AllMessageValues] | None = inputs.get("structured_messages")
|
||||
messages: list[AllMessageValues] | None = inputs.get("structured_messages") if input_type == "request" else None
|
||||
if not messages:
|
||||
messages = request_data.get("messages")
|
||||
|
||||
|
|
|
|||
|
|
@ -380,11 +380,12 @@ class StraikerGuardrail(CustomGuardrail):
|
|||
call_id: Final = getattr(logging_obj, "litellm_call_id", None) if logging_obj else None
|
||||
event_id: Final = f"{call_id or 'litellm'}:{input_type}"
|
||||
|
||||
is_request: Final = input_type == "request"
|
||||
content: Final = StraikerWebhookContent(
|
||||
texts=list(inputs.get("texts") or []),
|
||||
images=list(inputs.get("images") or []),
|
||||
structured_messages=_opaque_dict_list(inputs.get("structured_messages")),
|
||||
tools=_opaque_dict_list(inputs.get("tools")),
|
||||
structured_messages=_opaque_dict_list(inputs.get("structured_messages")) if is_request else None,
|
||||
tools=_opaque_dict_list(inputs.get("tools")) if is_request else None,
|
||||
tool_calls=_opaque_dict_list(inputs.get("tool_calls")),
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -104,6 +104,10 @@ def _chunk_choices(item: object) -> Sequence[object]:
|
|||
return choices
|
||||
|
||||
|
||||
def _held_choices(held_chars_per_choice: Mapping[int, int]) -> frozenset[int]:
|
||||
return frozenset(idx for idx, held in held_chars_per_choice.items() if held > 0)
|
||||
|
||||
|
||||
def _is_redundant_scan(scan_key: "StreamingScanKey | None", last_scan_key: "StreamingScanKey | None") -> bool:
|
||||
if scan_key is None:
|
||||
return False
|
||||
|
|
@ -472,6 +476,7 @@ class UnifiedLLMGuardrails(CustomLogger):
|
|||
emitted_text_per_choice: dict[int, str],
|
||||
holdback_per_choice: dict[int, int],
|
||||
finish_reason_per_choice: dict[int, str | None],
|
||||
held_chars_per_choice: dict[int, int],
|
||||
is_final: bool,
|
||||
) -> ModelResponseStream | None:
|
||||
"""Build the synthetic chunk carrying the newly-guardrailed deltas.
|
||||
|
|
@ -479,7 +484,9 @@ class UnifiedLLMGuardrails(CustomLogger):
|
|||
For each choice, the new delta is the mutated accumulated text past what
|
||||
has already been emitted, minus a trailing holdback (forced to 0 on the
|
||||
final flush). ``emitted_text_per_choice`` holds the exact bytes already
|
||||
sent per choice and is extended in place. Returns None when there is no
|
||||
sent per choice and is extended in place; ``held_chars_per_choice`` is
|
||||
updated in place with how many mutated chars per choice are still withheld
|
||||
after this round. Returns None when there is no
|
||||
text to emit (e.g. a tool-call-only turn) or nothing new and this is not
|
||||
the final chunk.
|
||||
|
||||
|
|
@ -536,6 +543,7 @@ class UnifiedLLMGuardrails(CustomLogger):
|
|||
holdback = 0 if is_final else max(0, holdback_per_choice.get(choice_idx, 0))
|
||||
end = max(len(already), len(text) - holdback)
|
||||
deltas[choice_idx] = text[len(already) : end]
|
||||
held_chars_per_choice[choice_idx] = len(text) - end
|
||||
|
||||
# Iterate the mutated choices (not just those in reference_chunk) so a
|
||||
# choice with pending text is never dropped for n > 1. finish_reason is
|
||||
|
|
@ -590,6 +598,7 @@ class UnifiedLLMGuardrails(CustomLogger):
|
|||
responses_yielded: list[object],
|
||||
emitted_text_per_choice: dict[int, str],
|
||||
finish_reason_per_choice: dict[int, str | None],
|
||||
held_chars_per_choice: dict[int, int],
|
||||
is_final: bool,
|
||||
) -> AsyncGenerator[object, None]:
|
||||
"""Run one guardrail processing round and emit the resulting diff chunk.
|
||||
|
|
@ -618,6 +627,7 @@ class UnifiedLLMGuardrails(CustomLogger):
|
|||
emitted_text_per_choice=emitted_text_per_choice,
|
||||
holdback_per_choice=sink.holdback_per_choice,
|
||||
finish_reason_per_choice=finish_reason_per_choice,
|
||||
held_chars_per_choice=held_chars_per_choice,
|
||||
is_final=is_final,
|
||||
)
|
||||
except ModifyResponseException as e:
|
||||
|
|
@ -673,6 +683,7 @@ class UnifiedLLMGuardrails(CustomLogger):
|
|||
responses_yielded: Final[list[object]] = []
|
||||
emitted_text_per_choice: Final[dict[int, str]] = {}
|
||||
finish_reason_per_choice: Final[dict[int, str | None]] = {}
|
||||
held_chars_per_choice: Final[dict[int, int]] = {}
|
||||
chunk_counter = 0
|
||||
last_chunk: object | None = None
|
||||
|
||||
|
|
@ -688,6 +699,7 @@ class UnifiedLLMGuardrails(CustomLogger):
|
|||
responses_yielded=responses_yielded,
|
||||
emitted_text_per_choice=emitted_text_per_choice,
|
||||
finish_reason_per_choice=finish_reason_per_choice,
|
||||
held_chars_per_choice=held_chars_per_choice,
|
||||
is_final=is_final,
|
||||
)
|
||||
|
||||
|
|
@ -724,12 +736,18 @@ class UnifiedLLMGuardrails(CustomLogger):
|
|||
# finish_reason to the final text terminator (see the
|
||||
# _tool_call_passthrough_chunk docstring).
|
||||
tool_only = self._tool_call_passthrough_chunk(
|
||||
item, finish_reason_per_choice=finish_reason_per_choice
|
||||
item,
|
||||
finish_reason_per_choice=finish_reason_per_choice,
|
||||
held_choices=_held_choices(held_chars_per_choice),
|
||||
)
|
||||
responses_yielded.append(tool_only)
|
||||
yield tool_only
|
||||
continue
|
||||
|
||||
if self._is_trailing_metadata_chunk(item):
|
||||
responses_so_far.append(item)
|
||||
continue
|
||||
|
||||
chunk_counter += 1
|
||||
responses_so_far.append(item)
|
||||
last_chunk = item
|
||||
|
|
@ -773,12 +791,33 @@ class UnifiedLLMGuardrails(CustomLogger):
|
|||
):
|
||||
yield out
|
||||
|
||||
if last_chunk is not None:
|
||||
async for out in _round(last_chunk, is_final=True):
|
||||
yield out
|
||||
async for out in self._emit_stream_tail(
|
||||
last_chunk=last_chunk,
|
||||
final_round=_round,
|
||||
responses_so_far=responses_so_far,
|
||||
responses_yielded=responses_yielded,
|
||||
):
|
||||
yield out
|
||||
except _StreamTerminated:
|
||||
return
|
||||
|
||||
async def _emit_stream_tail(
|
||||
self,
|
||||
*,
|
||||
last_chunk: object | None,
|
||||
final_round: Callable[[object, bool], AsyncGenerator[object, None]],
|
||||
responses_so_far: Sequence[object],
|
||||
responses_yielded: list[object],
|
||||
) -> AsyncGenerator[object, None]:
|
||||
"""Flush the held text with holdback 0, then replay metadata-only chunks
|
||||
(usage) so they land after the text and its finish_reason, as upstream sent them."""
|
||||
if last_chunk is not None:
|
||||
async for out in final_round(last_chunk, True):
|
||||
yield out
|
||||
for trailing in self._trailing_metadata_chunks(responses_so_far):
|
||||
responses_yielded.append(trailing)
|
||||
yield trailing
|
||||
|
||||
async def _inspect_full_response_for_block(
|
||||
self,
|
||||
*,
|
||||
|
|
@ -829,6 +868,23 @@ class UnifiedLLMGuardrails(CustomLogger):
|
|||
return True
|
||||
return False
|
||||
|
||||
@classmethod
|
||||
def _is_trailing_metadata_chunk(cls, item: object) -> bool:
|
||||
"""True for a chunk that carries only stream metadata (no choices, or a
|
||||
``usage`` chunk whose deltas are empty); such chunks are replayed after
|
||||
the final text flush instead of being folded into the transform."""
|
||||
if not _chunk_choices(item):
|
||||
return True
|
||||
return (
|
||||
getattr(item, "usage", None) is not None
|
||||
and not cls._chunk_carries_text(item)
|
||||
and not cls._chunk_has_finish_reason(item)
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def _trailing_metadata_chunks(cls, items: Sequence[object]) -> tuple[object, ...]:
|
||||
return tuple(item for item in items if cls._is_trailing_metadata_chunk(item))
|
||||
|
||||
@staticmethod
|
||||
def _chunk_carries_text(item: object) -> bool:
|
||||
"""True if any choice in this chunk has non-empty string ``delta.content``."""
|
||||
|
|
@ -843,6 +899,7 @@ class UnifiedLLMGuardrails(CustomLogger):
|
|||
def _tool_call_passthrough_chunk(
|
||||
item: object,
|
||||
finish_reason_per_choice: "dict[int, str | None] | None" = None,
|
||||
held_choices: frozenset[int] = frozenset(),
|
||||
) -> ModelResponseStream:
|
||||
"""Copy of a chunk carrying tool calls with all text content stripped.
|
||||
|
||||
|
|
@ -851,8 +908,9 @@ class UnifiedLLMGuardrails(CustomLogger):
|
|||
transform instead). Applies per choice so an n>1 chunk mixing a text
|
||||
choice and a tool-call choice does not leak the text choice.
|
||||
|
||||
For a choice that carries BOTH text content AND tool_calls, ``finish_reason``
|
||||
is suppressed on the passthrough and recorded on
|
||||
For a choice that carries BOTH text content AND tool_calls, or whose earlier
|
||||
text is still withheld (``held_choices``), ``finish_reason`` is suppressed on
|
||||
the passthrough and recorded on
|
||||
``finish_reason_per_choice`` (when provided) so the final synthetic text
|
||||
chunk delivers it. Emitting the passthrough's ``finish_reason`` before the
|
||||
text flush would let a spec-compliant SSE client stop reading at
|
||||
|
|
@ -865,7 +923,8 @@ class UnifiedLLMGuardrails(CustomLogger):
|
|||
idx = getattr(choice, "index", 0) or 0
|
||||
original_finish = getattr(choice, "finish_reason", None)
|
||||
has_text = isinstance(getattr(delta, "content", None), str) and getattr(delta, "content", "") != ""
|
||||
if has_text and original_finish is not None and finish_reason_per_choice is not None:
|
||||
text_pending = has_text or idx in held_choices
|
||||
if text_pending and original_finish is not None and finish_reason_per_choice is not None:
|
||||
finish_reason_per_choice[idx] = original_finish
|
||||
passthrough_finish: str | None = None
|
||||
else:
|
||||
|
|
|
|||
Some files were not shown because too many files have changed in this diff Show more
Loading…
Add table
Reference in a new issue