scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-07 02:59:05 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 103

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

Browse source
This commit is contained in:
kerry 2026-09-17 20:52:32 +00:00
commit f5940bcee4
760 changed files with 53978 additions and 15071 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

32
.circleci/config.yml
View file

@ -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:

13
.circleci/scripts/classify_changes.sh
View file

@ -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
;;

29
.circleci/scripts/run_integration.sh
View file

@ -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
View 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()

4
.github/CODEOWNERS vendored
View file

@ -4,7 +4,7 @@
/ui/nginx.conf
/ui/litellm-dashboard/src/lib/http/schema.d.ts
/ui/litellm-dashboard/tsconfig.tsbuildinfo
/model_prices_and_context_window.json @mateo-berri
/litellm/model_prices_and_context_window_backup.json @mateo-berri
/model_prices_and_context_window.json @mateo-berri @ryan-crabbe-berri @kerry-berri
/litellm/model_prices_and_context_window_backup.json @mateo-berri @ryan-crabbe-berri @kerry-berri
/litellm-proxy-extras/litellm_proxy_extras/migrations/ @yuneng-berri @ryan-crabbe-berri
/.github/CODEOWNERS @yuneng-berri

39
.github/scripts/assert_ci_coverage.py vendored
View file

@ -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
View 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
View 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

65
.github/workflows/test-rust.yml vendored
View file

@ -70,7 +70,7 @@ env:
jobs:
rust-lint:
runs-on: ubuntu-latest
timeout-minutes: 10
timeout-minutes: 15
defaults:
run:
working-directory: litellm-rust
@ -81,26 +81,48 @@ jobs:
- run: rustup toolchain install --no-self-update
- run: cargo fmt --check
- run: cargo fmt --all --check
- uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0
- uses: Swatinem/rust-cache@6323deb102c322ba6fcbdcafc7e3dddab59af2b6 # v2.9.2
with:
path: |
~/.cargo/registry
~/.cargo/git
litellm-rust/target
key: ${{ runner.os }}-cargo-${{ github.job }}-${{ hashFiles('rust-toolchain.toml', '.cargo/**', 'litellm-rust/Cargo.lock') }}
restore-keys: |
${{ runner.os }}-cargo-${{ github.job }}-
workspaces: litellm-rust
cache-on-failure: true
- run: cargo clippy --workspace --all-targets --locked -- -D warnings
- run: cargo clippy -p litellm-core --all-targets --features bedrock-auth --locked -- -D warnings
rust-test:
runs-on: ubuntu-latest
timeout-minutes: 30
timeout-minutes: 20
defaults:
run:
working-directory: litellm-rust
steps:
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
persist-credentials: false
- uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
with:
python-version: "3.12"
- run: rustup toolchain install --no-self-update
- uses: taiki-e/install-action@d438492cf8a250514fa2d34b30bc3c0dc37c65ff # v2.87.8
with:
tool: cargo-nextest@0.9.143
- uses: Swatinem/rust-cache@6323deb102c322ba6fcbdcafc7e3dddab59af2b6 # v2.9.2
with:
workspaces: litellm-rust
cache-on-failure: true
- run: cargo nextest run --workspace --locked
- run: cargo test --workspace --doc --locked
rust-wheel:
runs-on: ubuntu-latest
timeout-minutes: 30
steps:
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
@ -116,21 +138,10 @@ jobs:
- run: rustup toolchain install --no-self-update
- uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0
- uses: Swatinem/rust-cache@6323deb102c322ba6fcbdcafc7e3dddab59af2b6 # v2.9.2
with:
path: |
~/.cargo/registry
~/.cargo/git
litellm-rust/target
key: ${{ runner.os }}-cargo-${{ github.job }}-${{ hashFiles('rust-toolchain.toml', '.cargo/**', 'litellm-rust/Cargo.lock') }}
restore-keys: |
${{ runner.os }}-cargo-${{ github.job }}-
- run: cargo test --workspace --locked
working-directory: litellm-rust
- run: cargo test -p litellm-core --features bedrock-auth --locked
working-directory: litellm-rust
workspaces: litellm-rust
cache-on-failure: true
- run: uv build --wheel --out-dir dist

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

@ -100,6 +100,7 @@ jobs:
tests/test_litellm/secret_managers
tests/test_litellm/a2a_protocol
tests/test_litellm/anthropic_interface
tests/test_litellm/chat_completions
tests/test_litellm/completion_extras
tests/test_litellm/compression
tests/test_litellm/containers
@ -109,6 +110,7 @@ jobs:
tests/test_litellm/repositories
tests/test_litellm/images
tests/test_litellm/interactions
tests/test_litellm/messages
tests/test_litellm/ocr
tests/test_litellm/passthrough
tests/test_litellm/rag

3
ARCHITECTURE.md
View file

@ -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 |

5
db_scripts/partition_spend_logs.sql
View file

@ -53,6 +53,8 @@ ALTER INDEX IF EXISTS "LiteLLM_SpendLogs_end_user_idx"
RENAME TO "LiteLLM_SpendLogs_legacy_end_user_idx";
ALTER INDEX IF EXISTS "LiteLLM_SpendLogs_session_id_idx"
RENAME TO "LiteLLM_SpendLogs_legacy_session_id_idx";
ALTER INDEX IF EXISTS "LiteLLM_SpendLogs_api_key_startTime_idx"
RENAME TO "LiteLLM_SpendLogs_legacy_api_key_startTime_idx";
CREATE TABLE "LiteLLM_SpendLogs" (
LIKE "LiteLLM_SpendLogs_legacy" INCLUDING DEFAULTS INCLUDING GENERATED
@ -78,6 +80,9 @@ CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogs_end_user_idx"
CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogs_session_id_idx"
ON "LiteLLM_SpendLogs" ("session_id");
CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogs_api_key_startTime_idx"
ON "LiteLLM_SpendLogs" ("api_key", "startTime");
-- Safety net: any row whose startTime has no explicit partition lands here so
-- writes never fail. The cleanup job never drops the DEFAULT partition.
CREATE TABLE IF NOT EXISTS "LiteLLM_SpendLogs_pdefault"

5
db_scripts/unpartition_spend_logs.sql
View file

@ -40,6 +40,8 @@ ALTER INDEX IF EXISTS "LiteLLM_SpendLogs_end_user_idx"
RENAME TO "LiteLLM_SpendLogs_partitioned_end_user_idx";
ALTER INDEX IF EXISTS "LiteLLM_SpendLogs_session_id_idx"
RENAME TO "LiteLLM_SpendLogs_partitioned_session_id_idx";
ALTER INDEX IF EXISTS "LiteLLM_SpendLogs_api_key_startTime_idx"
RENAME TO "LiteLLM_SpendLogs_partitioned_api_key_startTime_idx";
CREATE TABLE "LiteLLM_SpendLogs" (
LIKE "LiteLLM_SpendLogs_partitioned" INCLUDING DEFAULTS INCLUDING GENERATED
@ -60,6 +62,9 @@ CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogs_end_user_idx"
CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogs_session_id_idx"
ON "LiteLLM_SpendLogs" ("session_id");
CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogs_api_key_startTime_idx"
ON "LiteLLM_SpendLogs" ("api_key", "startTime");
INSERT INTO "LiteLLM_SpendLogs"
SELECT * FROM "LiteLLM_SpendLogs_partitioned"
ON CONFLICT ("request_id") DO NOTHING;

2
litellm-proxy-extras/litellm_proxy_extras/migrations/20260823000000_add_spend_logs_api_key_starttime_index/migration.sql Normal file
View file

@ -0,0 +1,2 @@
-- CreateIndex
CREATE INDEX IF NOT EXISTS "LiteLLM_SpendLogs_api_key_startTime_idx" ON "LiteLLM_SpendLogs"("api_key", "startTime");

5
litellm-proxy-extras/litellm_proxy_extras/migrations/20260916000000_add_key_total_spend/migration.sql Normal file
View file

@ -0,0 +1,5 @@
-- AlterTable
ALTER TABLE "LiteLLM_VerificationToken" ADD COLUMN IF NOT EXISTS "total_spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0;
-- AlterTable
ALTER TABLE "LiteLLM_DeletedVerificationToken" ADD COLUMN IF NOT EXISTS "total_spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0;

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

@ -426,6 +426,7 @@ model LiteLLM_VerificationToken {
key_alias String?
soft_budget_cooldown Boolean @default(false) // key-level state on if budget alerts need to be cooled down
spend Float @default(0.0)
total_spend Float @default(0.0)
expires DateTime?
models String[]
aliases Json @default("{}")
@ -528,6 +529,7 @@ model LiteLLM_DeletedVerificationToken {
key_alias String?
soft_budget_cooldown Boolean @default(false)
spend Float @default(0.0)
total_spend Float @default(0.0)
expires DateTime?
models String[]
aliases Json @default("{}")
@ -676,6 +678,7 @@ model LiteLLM_SpendLogs {
@@index([end_user])
@@index([session_id])
@@index([litellm_call_id])
@@index([api_key, startTime])
}
model LiteLLM_BudgetWindowSpend {

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

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

369
litellm-rust/Cargo.lock generated
View file

@ -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"
@ -931,8 +948,18 @@ version = "0.20.11"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "fc7f46116c46ff9ab3eb1597a45688b6715c6e628b5c133e288e709a29bcb4ee"
dependencies = [
"darling_core",
"darling_macro",
"darling_core 0.20.11",
"darling_macro 0.20.11",
]
[[package]]
name = "darling"
version = "0.21.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "9cdf337090841a411e2a7f3deb9187445851f91b309c0c0a29e05f74a00a48c0"
dependencies = [
"darling_core 0.21.3",
"darling_macro 0.21.3",
]
[[package]]
@ -949,13 +976,38 @@ dependencies = [
"syn 2.0.119",
]
[[package]]
name = "darling_core"
version = "0.21.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "1247195ecd7e3c85f83c8d2a366e4210d588e802133e1e355180a9870b517ea4"
dependencies = [
"fnv",
"ident_case",
"proc-macro2",
"quote",
"strsim",
"syn 2.0.119",
]
[[package]]
name = "darling_macro"
version = "0.20.11"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "fc34b93ccb385b40dc71c6fceac4b2ad23662c7eeb248cf10d529b7e055b6ead"
dependencies = [
"darling_core",
"darling_core 0.20.11",
"quote",
"syn 2.0.119",
]
[[package]]
name = "darling_macro"
version = "0.21.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "d38308df82d1080de0afee5d069fa14b0326a88c14f15c5ccda35b4a6c414c81"
dependencies = [
"darling_core 0.21.3",
"quote",
"syn 2.0.119",
]
@ -1005,7 +1057,7 @@ version = "0.20.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "2d5bcf7b024d6835cfb3d473887cd966994907effbe9227e8c8219824d06c4e8"
dependencies = [
"darling",
"darling 0.20.11",
"proc-macro2",
"quote",
"syn 2.0.119",
@ -1346,7 +1398,7 @@ dependencies = [
"futures-sink",
"futures-util",
"http 0.2.12",
"indexmap",
"indexmap 2.14.0",
"slab",
"tokio",
"tokio-util",
@ -1365,7 +1417,7 @@ dependencies = [
"futures-core",
"futures-sink",
"http 1.4.2",
"indexmap",
"indexmap 2.14.0",
"slab",
"tokio",
"tokio-util",
@ -1383,6 +1435,12 @@ dependencies = [
"zerocopy",
]
[[package]]
name = "hashbrown"
version = "0.12.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "8a9ee70c43aaf417c914396645a0fa852624801b24ebb7ae78fe8272889ac888"
[[package]]
name = "hashbrown"
version = "0.17.1"
@ -1719,6 +1777,17 @@ dependencies = [
"icu_properties",
]
[[package]]
name = "indexmap"
version = "1.9.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "bd070e393353796e801d209ad339e89596eb4c8d430d18ede6a1cced8fafbd99"
dependencies = [
"autocfg",
"hashbrown 0.12.3",
"serde",
]
[[package]]
name = "indexmap"
version = "2.14.0"
@ -1726,7 +1795,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "d466e9454f08e4a911e14806c24e16fba1b4c121d1ea474396f396069cf949d9"
dependencies = [
"equivalent",
"hashbrown",
"hashbrown 0.17.1",
"serde",
"serde_core",
]
@ -1838,7 +1907,24 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "68ab91017fe16c622486840e4c83c9a37afeff978bd239b5293d61ece587de66"
[[package]]
name = "litellm-core"
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"
dependencies = [
"serde",
"subtle",
"thiserror 2.0.19",
"tokio",
"veil",
]
[[package]]
name = "litellm-auth-aws"
version = "0.1.0"
dependencies = [
"aws-config",
@ -1847,13 +1933,89 @@ dependencies = [
"aws-sigv4",
"aws-smithy-runtime-api",
"aws-types",
"litellm-auth",
"moka",
"reqwest 0.12.28",
"serde_json",
"sha2 0.10.9",
"thiserror 2.0.19",
"tokio",
]
[[package]]
name = "litellm-auth-azure"
version = "0.1.0"
dependencies = [
"azure_core",
"azure_identity",
"litellm-auth",
"moka",
"serde_json",
"sha2 0.10.9",
"strum",
"tokio",
"url",
]
[[package]]
name = "litellm-auth-gcp"
version = "0.1.0"
dependencies = [
"gcp_auth",
"litellm-auth",
"moka",
"serde_json",
"sha2 0.10.9",
"tokio",
]
[[package]]
name = "litellm-cache"
version = "0.1.0"
dependencies = [
"rstest",
"serde",
"serde_json",
"sha2 0.10.9",
"thiserror 2.0.19",
]
[[package]]
name = "litellm-cache-memory"
version = "0.1.0"
dependencies = [
"litellm-cache",
"rstest",
"serde_json",
"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",
"futures-util",
"gcp_auth",
"litellm-auth",
"litellm-auth-aws",
"litellm-auth-azure",
"litellm-auth-gcp",
"litellm-framing",
"mime_guess",
"moka",
"rand 0.8.7",
@ -1864,22 +2026,40 @@ dependencies = [
"serde",
"serde_json",
"serde_path_to_error",
"serde_with",
"sha2 0.10.9",
"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"
dependencies = [
"bytes",
"criterion",
"futures-util",
"litellm-auth",
"litellm-core",
"litellm-python-interop",
"litellm-token-counter",
@ -1909,7 +2089,7 @@ version = "0.1.0"
dependencies = [
"base64 0.22.1",
"criterion",
"indexmap",
"indexmap 2.14.0",
"itoa",
"rand 0.8.7",
"rstest",
@ -2065,6 +2245,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"
@ -2581,6 +2771,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"
@ -2590,6 +2810,26 @@ dependencies = [
"bitflags",
]
[[package]]
name = "ref-cast"
version = "1.0.27"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "7e440fb4e4b4147295338efb76001ab9e4efc0e5839df2c47fc5ac2381d365c3"
dependencies = [
"ref-cast-impl",
]
[[package]]
name = "ref-cast-impl"
version = "1.0.27"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "92ecd8964f8453721699a1ed72037b0db49ce2f5a5138486ee89bed6f67cdf3a"
dependencies = [
"proc-macro2",
"quote",
"syn 3.0.0",
]
[[package]]
name = "regex"
version = "1.13.1"
@ -2771,6 +3011,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"
@ -2899,6 +3152,30 @@ dependencies = [
"windows-sys 0.61.2",
]
[[package]]
name = "schemars"
version = "0.9.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "4cd191f9397d57d581cddd31014772520aa448f65ef991055d7f61582c65165f"
dependencies = [
"dyn-clone",
"ref-cast",
"serde",
"serde_json",
]
[[package]]
name = "schemars"
version = "1.2.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "687274d293b6cdc6e73e0fee520bf2049650090d7164f87672d212a3c530cf4a"
dependencies = [
"dyn-clone",
"ref-cast",
"serde",
"serde_json",
]
[[package]]
name = "scopeguard"
version = "1.2.0"
@ -2980,6 +3257,7 @@ version = "1.0.150"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e8014e44b4736ed0538adeecded0fce2a272f22dc9578a7eb6b2d9993c74cfb9"
dependencies = [
"indexmap 2.14.0",
"itoa",
"memchr",
"serde",
@ -3010,6 +3288,37 @@ dependencies = [
"serde",
]
[[package]]
name = "serde_with"
version = "3.16.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "4fa237f2807440d238e0364a218270b98f767a00d3dada77b1c53ae88940e2e7"
dependencies = [
"base64 0.22.1",
"chrono",
"hex",
"indexmap 1.9.3",
"indexmap 2.14.0",
"schemars 0.9.0",
"schemars 1.2.2",
"serde_core",
"serde_json",
"serde_with_macros",
"time",
]
[[package]]
name = "serde_with_macros"
version = "3.16.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "52a8e3ca0ca629121f70ab50f95249e5a6f925cc0f6ffe8256c45b728875706c"
dependencies = [
"darling 0.21.3",
"proc-macro2",
"quote",
"syn 2.0.119",
]
[[package]]
name = "sha1"
version = "0.10.7"
@ -3021,6 +3330,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"
@ -3125,6 +3440,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"
@ -3224,6 +3552,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"
@ -3453,7 +3794,7 @@ version = "0.25.13+spec-1.1.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "6975367e4d2ef766d86af01ffad14b622fecc8d4357a998fbc4deb6e9bacaf9b"
dependencies = [
"indexmap",
"indexmap 2.14.0",
"toml_datetime",
"toml_parser",
"winnow",
@ -4107,6 +4448,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"

12
litellm-rust/Cargo.toml
View file

@ -11,6 +11,13 @@ 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" }
litellm-auth-gcp = { path = "crates/auth-gcp" }
litellm-cache = { path = "crates/cache" }
litellm-cache-memory = { path = "crates/cache-memory" }
litellm-token-counter = { path = "crates/token-counter" }
litellm-python-interop = { path = "crates/python-interop" }
pyo3 = "0.29.2"
@ -23,6 +30,7 @@ rustls = { version = "0.23", default-features = false, features = ["ring", "std"
rustls-native-certs = "0.8"
serde = { version = "1.0", features = ["derive"] }
serde_json = { version = "1.0", features = ["float_roundtrip"] }
serde_with = { version = "=3.16.1", default-features = false, features = ["std", "macros"] }
sha2 = "0.10"
subtle = "2"
thiserror = "2.0"
@ -30,12 +38,10 @@ tokio = { version = "1", features = ["rt-multi-thread", "macros", "time", "net"]
tokio-tungstenite = { version = "0.24", default-features = false, features = ["connect", "rustls-tls-native-roots"] }
futures-util = { version = "0.3", default-features = false, features = ["sink", "std"] }
base64 = "0.22"
gcp_auth = "0.12.7"
azure_core = "1.0.0"
azure_identity = { version = "1.0.0", features = ["tokio"] }
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"

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

@ -0,0 +1,25 @@
[package]
name = "litellm-auth-aws"
version = "0.1.0"
edition.workspace = true
license.workspace = true
repository.workspace = true
[dependencies]
litellm-auth.workspace = true
moka = { workspace = true, features = ["sync"] }
serde_json.workspace = true
sha2.workspace = true
thiserror.workspace = true
aws-config = { version = "1.9.0", default-features = false, features = ["rustls", "rt-tokio"] }
aws-credential-types = { version = "1.3.0", features = ["hardcoded-credentials"] }
aws-sdk-sts = { version = "1.108.0", default-features = false, features = ["rustls", "rt-tokio"] }
aws-sigv4 = "1.5.1"
aws-types = "1.4.0"
aws-smithy-runtime-api = "1.13.0"
[dev-dependencies]
reqwest.workspace = true
tokio.workspace = true

79
litellm-rust/crates/core/src/providers/bedrock/aws_base.rs → litellm-rust/crates/auth-aws/src/aws.rs
View file

@ -1,10 +1,12 @@
use std::collections::BTreeMap;
use std::sync::{Mutex, OnceLock};
use std::sync::OnceLock;
use std::time::Duration;
use std::time::{SystemTime, UNIX_EPOCH};
use crate::caching::in_memory_cache::InMemoryCache;
use crate::error::Error;
use moka::sync::Cache;
use serde_json::{Map, Value};
use sha2::{Digest, Sha256};
use aws_credential_types::Credentials;
use aws_credential_types::provider::ProvideCredentials;
use aws_sigv4::http_request::{
@ -12,9 +14,8 @@ use aws_sigv4::http_request::{
};
use aws_sigv4::sign::v4;
use aws_smithy_runtime_api::client::identity::Identity;
use serde_json::{Map, Value};
use sha2::{Digest, Sha256};
use super::Error;
use super::constants::{
AWS_ACCESS_KEY_ID, AWS_EXTERNAL_ID, AWS_PROFILE_NAME, AWS_REGION, AWS_REGION_NAME,
AWS_ROLE_ARN, AWS_ROLE_NAME, AWS_SECRET_ACCESS_KEY, AWS_SESSION_NAME, AWS_SESSION_TOKEN,
@ -26,7 +27,8 @@ use super::constants::{
const STATIC_CREDENTIALS_TTL: Duration = Duration::from_secs(3600 - 60);
const AMBIENT_CREDENTIALS_TTL: Duration = Duration::from_secs(600);
static IAM_CREDENTIALS_CACHE: OnceLock<Mutex<InMemoryCache<Credentials>>> = OnceLock::new();
static STATIC_CREDENTIALS_CACHE: OnceLock<Cache<String, Credentials>> = OnceLock::new();
static AMBIENT_CREDENTIALS_CACHE: OnceLock<Cache<String, Credentials>> = OnceLock::new();
fn credential_cache_ttl(flow: &AwsAuthFlow) -> Option<Duration> {
match flow {
@ -107,16 +109,35 @@ fn cache_key(config: &AwsAuthConfig, flow: &AwsAuthFlow) -> String {
format!("{:x}", hasher.finalize())
}
fn static_credentials_cache() -> &'static Cache<String, Credentials> {
STATIC_CREDENTIALS_CACHE.get_or_init(|| {
Cache::builder()
.max_capacity(200)
.time_to_live(STATIC_CREDENTIALS_TTL)
.build()
})
}
fn ambient_credentials_cache() -> &'static Cache<String, Credentials> {
AMBIENT_CREDENTIALS_CACHE.get_or_init(|| {
Cache::builder()
.max_capacity(200)
.time_to_live(AMBIENT_CREDENTIALS_TTL)
.build()
})
}
fn get_cached_credentials(key: &str) -> Option<Credentials> {
let cache = IAM_CREDENTIALS_CACHE.get_or_init(|| Mutex::new(InMemoryCache::default()));
let mut entries = cache.lock().ok()?;
entries.get_cache(key)
static_credentials_cache()
.get(key)
.or_else(|| ambient_credentials_cache().get(key))
}
fn set_cached_credentials(key: String, credentials: Credentials, ttl: Duration) {
let cache = IAM_CREDENTIALS_CACHE.get_or_init(|| Mutex::new(InMemoryCache::default()));
if let Ok(mut entries) = cache.lock() {
entries.set_cache(key, credentials, Some(ttl));
if ttl == STATIC_CREDENTIALS_TTL {
static_credentials_cache().insert(key, credentials);
} else {
ambient_credentials_cache().insert(key, credentials);
}
}
@ -247,7 +268,7 @@ pub async fn resolve_credentials(
provider
.provide_credentials()
.await
.map_err(|error| Error::Auth(format!("AWS profile credentials failed: {error}")))
.map_err(|error| Error::AwsProfile(error.to_string()))
}
AwsAuthFlow::AssumeRole { role, session_name } => {
if is_already_running_as_role(&role, &resolved).await? {
@ -260,9 +281,10 @@ pub async fn resolve_credentials(
aws_config::default_provider::credentials::DefaultCredentialsChain::builder()
.build()
.await;
let credentials = provider.provide_credentials().await.map_err(|error| {
Error::Auth(format!("AWS default credentials failed: {error}"))
})?;
let credentials = provider
.provide_credentials()
.await
.map_err(|error| Error::AwsDefaultChain(error.to_string()))?;
set_cached_credentials(
key,
credentials.clone(),
@ -302,7 +324,7 @@ pub async fn resolve_credentials(
provider
.provide_credentials()
.await
.map_err(|error| Error::Auth(format!("AWS role credentials failed: {error}")))
.map_err(|error| Error::AwsAssumeRole(error.to_string()))
}
AwsAuthFlow::WebIdentity {
token,
@ -325,15 +347,12 @@ pub async fn resolve_credentials(
.web_identity_token(token)
.send()
.await
.map_err(|error| {
Error::Auth(format!("AWS web identity credentials failed: {error}"))
})?;
let credentials = response.credentials().ok_or_else(|| {
Error::Auth("AWS web identity response had no credentials".to_string())
})?;
let expiration = SystemTime::try_from(*credentials.expiration()).map_err(|error| {
Error::Auth(format!("AWS web identity expiration was invalid: {error}"))
})?;
.map_err(|error| Error::AwsWebIdentity(error.to_string()))?;
let credentials = response
.credentials()
.ok_or(Error::AwsMissingWebIdentityCredentials)?;
let expiration = SystemTime::try_from(*credentials.expiration())
.map_err(|error| Error::AwsWebIdentityExpiration(error.to_string()))?;
Ok(Credentials::new(
credentials.access_key_id(),
credentials.secret_access_key(),
@ -354,7 +373,7 @@ pub async fn resolve_credentials(
let credentials = provider
.provide_credentials()
.await
.map_err(|error| Error::Auth(format!("AWS default credentials failed: {error}")))?;
.map_err(|error| Error::AwsDefaultChain(error.to_string()))?;
set_cached_credentials(
key,
credentials.clone(),
@ -449,14 +468,14 @@ pub fn sign_bedrock_post(
.settings(SigningSettings::default())
.build()
.map(SigningParams::from)
.map_err(|error| Error::Auth(format!("AWS signing parameters failed: {error}")))?;
.map_err(|error| Error::AwsSigningParameters(error.to_string()))?;
let header_refs = headers
.iter()
.map(|(name, value)| (name.as_str(), value.as_str()));
let request = SignableRequest::new("POST", url, header_refs, SignableBody::Bytes(body))
.map_err(|error| Error::Auth(format!("AWS signable request failed: {error}")))?;
.map_err(|error| Error::AwsSignableRequest(error.to_string()))?;
let (instructions, _) = sign(request, &params)
.map_err(|error| Error::Auth(format!("AWS request signing failed: {error}")))?
.map_err(|error| Error::AwsSigning(error.to_string()))?
.into_parts();
Ok(instructions
.headers()

0
litellm-rust/crates/core/src/providers/bedrock/constants.rs → litellm-rust/crates/auth-aws/src/constants.rs
View file

46
litellm-rust/crates/auth-aws/src/error.rs Normal file
View file

@ -0,0 +1,46 @@
use thiserror::Error as ThisError;
#[derive(Clone, Debug, ThisError, PartialEq, Eq)]
pub enum Error {
#[error("AWS profile credentials failed: {0}")]
AwsProfile(String),
#[error("AWS default credentials failed: {0}")]
AwsDefaultChain(String),
#[error("AWS role credentials failed: {0}")]
AwsAssumeRole(String),
#[error("AWS web identity credentials failed: {0}")]
AwsWebIdentity(String),
#[error("AWS web identity expiration was invalid: {0}")]
AwsWebIdentityExpiration(String),
#[error("AWS signing parameters failed: {0}")]
AwsSigningParameters(String),
#[error("AWS signable request failed: {0}")]
AwsSignableRequest(String),
#[error("AWS request signing failed: {0}")]
AwsSigning(String),
#[error("AWS web identity response had no credentials")]
AwsMissingWebIdentityCredentials,
}
impl From<Error> for litellm_auth::Error {
fn from(error: Error) -> Self {
Self::ProviderAuthentication(error.to_string())
}
}
#[cfg(test)]
mod tests {
use super::Error;
#[test]
fn converts_to_shared_auth_error_without_losing_context() {
let error = litellm_auth::Error::from(Error::AwsProfile("profile not found".into()));
assert_eq!(
error,
litellm_auth::Error::ProviderAuthentication(
"AWS profile credentials failed: profile not found".into()
)
);
}
}

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

@ -0,0 +1,6 @@
mod aws;
pub mod constants;
mod error;
pub use aws::*;
pub use error::Error;

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

@ -0,0 +1,21 @@
[package]
name = "litellm-auth-azure"
version = "0.1.0"
edition.workspace = true
license.workspace = true
repository.workspace = true
[dependencies]
litellm-auth.workspace = true
moka.workspace = true
serde_json.workspace = true
sha2.workspace = true
strum.workspace = true
url.workspace = true
azure_core = "1.0.0"
azure_identity = { version = "1.0.0", features = ["tokio"] }
[dev-dependencies]
tokio.workspace = true

6
litellm-rust/crates/core/src/providers/azure_ai/auth/credential_provider_cache.rs → litellm-rust/crates/auth-azure/src/credential_provider_cache.rs
View file

@ -4,7 +4,7 @@ use std::sync::Arc;
use azure_core::credentials::TokenCredential;
use moka::future::Cache;
use crate::AuthError;
use litellm_auth::Error;
#[derive(Clone, Debug, PartialEq, Eq, Hash)]
pub(crate) struct AzureCredentialProviderCacheKey {
@ -31,9 +31,9 @@ impl AzureCredentialProviderCache {
&self,
key: AzureCredentialProviderCacheKey,
create: F,
) -> Result<Arc<dyn TokenCredential>, AuthError>
) -> Result<Arc<dyn TokenCredential>, Error>
where
F: Future<Output = Result<Arc<dyn TokenCredential>, AuthError>>,
F: Future<Output = Result<Arc<dyn TokenCredential>, Error>>,
{
self.entries
.try_get_with(key, create)

7
litellm-rust/crates/auth-azure/src/lib.rs Normal file
View file

@ -0,0 +1,7 @@
mod credential_provider_cache;
mod native;
mod resolve;
mod types;
pub use resolve::AzureAuthService;
pub use types::AzureAuthInputs;

51
litellm-rust/crates/core/src/providers/azure_ai/auth/native.rs → litellm-rust/crates/auth-azure/src/native.rs
View file

@ -1,4 +1,3 @@
use crate::auth::error::AuthConfigurationError;
use std::sync::Arc;
use std::time::{Duration, UNIX_EPOCH};
@ -13,8 +12,8 @@ use azure_identity::{
};
use sha2::{Digest, Sha256};
use crate::AuthError;
use crate::auth::{InputSource, ResolvedCredential, SecretValue, Sourced};
use litellm_auth::Error;
use litellm_auth::{InputSource, ResolvedCredential, SecretValue, Sourced};
use super::credential_provider_cache::{
AzureCredentialProviderCache, AzureCredentialProviderCacheKey,
@ -62,7 +61,7 @@ pub(crate) struct ValidatedAzureRequest {
}
impl ValidatedAzureRequest {
pub(crate) fn new(request: NativeAzureRequest) -> Result<Self, AuthError> {
pub(crate) fn new(request: NativeAzureRequest) -> Result<Self, Error> {
validate_authority(&request)?;
let credential_source = validate_sources(&request)?;
Ok(Self {
@ -120,7 +119,7 @@ impl NativeAzureTokenAcquirer {
pub(crate) async fn acquire(
&self,
request: ValidatedAzureRequest,
) -> Result<ResolvedCredential, AuthError> {
) -> Result<ResolvedCredential, Error> {
let scope = request.request.scope().to_string();
let key = request.request.cache_key();
let transport = self.transport.clone();
@ -134,7 +133,7 @@ impl NativeAzureTokenAcquirer {
let token = credential
.get_token(&[scope.as_str()], None)
.await
.map_err(|error| AuthError::AzureTokenAcquisition(error.to_string()))?;
.map_err(|error| Error::AzureTokenAcquisition(error.to_string()))?;
let expires_on = u64::try_from(token.expires_on.unix_timestamp())
.ok()
.map(|seconds| UNIX_EPOCH + Duration::from_secs(seconds));
@ -239,7 +238,7 @@ impl NativeAzureRequest {
}
}
fn validate_authority(request: &NativeAzureRequest) -> Result<(), AuthError> {
fn validate_authority(request: &NativeAzureRequest) -> Result<(), Error> {
let authority = match request {
NativeAzureRequest::ClientSecret { authority, .. }
| NativeAzureRequest::ClientAssertion { authority, .. }
@ -251,8 +250,7 @@ fn validate_authority(request: &NativeAzureRequest) -> Result<(), AuthError> {
let Some(authority) = authority else {
return Ok(());
};
let url = url::Url::parse(authority.value())
.map_err(|_| AuthError::Configuration(AuthConfigurationError::InvalidAzureAuthority))?;
let url = url::Url::parse(authority.value()).map_err(|_| Error::InvalidAzureAuthority)?;
if url.scheme() != "https"
|| url.host_str().is_none()
|| !url.username().is_empty()
@ -261,14 +259,12 @@ fn validate_authority(request: &NativeAzureRequest) -> Result<(), AuthError> {
|| url.fragment().is_some()
|| !matches!(url.path(), "" | "/")
{
return Err(AuthError::Configuration(
AuthConfigurationError::InvalidAzureAuthority,
));
return Err(Error::InvalidAzureAuthority);
}
Ok(())
}
fn validate_sources(request: &NativeAzureRequest) -> Result<InputSource, AuthError> {
fn validate_sources(request: &NativeAzureRequest) -> Result<InputSource, Error> {
match request {
NativeAzureRequest::ClientSecret {
tenant_id,
@ -356,7 +352,7 @@ fn is_request_controlled<T>(value: &Sourced<T>, optional: Option<&Sourced<String
|| optional.is_some_and(|value| value.source() == InputSource::Request)
}
fn trusted_only(sources: &[InputSource]) -> Result<InputSource, AuthError> {
fn trusted_only(sources: &[InputSource]) -> Result<InputSource, Error> {
if sources.contains(&InputSource::Request) {
return mixed_sources();
}
@ -371,16 +367,14 @@ fn trusted_source(sources: &[InputSource]) -> InputSource {
}
}
fn mixed_sources<T>() -> Result<T, AuthError> {
Err(AuthError::Configuration(
AuthConfigurationError::MixedAzureCredentialSources,
))
fn mixed_sources<T>() -> Result<T, Error> {
Err(Error::MixedAzureCredentialSources)
}
fn build_credential(
request: NativeAzureRequest,
transport: Option<azure_core::http::Transport>,
) -> Result<Arc<dyn TokenCredential>, AuthError> {
) -> Result<Arc<dyn TokenCredential>, Error> {
match request {
NativeAzureRequest::ClientSecret {
tenant_id,
@ -439,11 +433,7 @@ fn build_credential(
NativeAzureRequest::DeveloperTools { .. } => DeveloperToolsCredential::new(None)
.map(|credential| credential as Arc<dyn TokenCredential>),
}
.map_err(|error| {
AuthError::Configuration(AuthConfigurationError::AzureCredentialInitialization(
error.to_string(),
))
})
.map_err(|error| Error::AzureCredentialInitialization(error.to_string()))
}
fn client_options(
@ -494,7 +484,7 @@ mod tests {
use azure_core::{Bytes, Result};
use super::{NativeAzureRequest, NativeAzureTokenAcquirer, ValidatedAzureRequest};
use crate::auth::{InputSource, SecretValue, Sourced};
use litellm_auth::{InputSource, SecretValue, Sourced};
fn deployment<T>(value: T) -> Sourced<T> {
Sourced::new(value, InputSource::Deployment)
@ -659,9 +649,7 @@ mod tests {
assert!(matches!(
error,
crate::AuthError::Configuration(
crate::auth::error::AuthConfigurationError::MixedAzureCredentialSources
)
litellm_auth::Error::MixedAzureCredentialSources
));
}
@ -691,12 +679,7 @@ mod tests {
authority,
))
.unwrap_err();
assert!(matches!(
error,
crate::AuthError::Configuration(
crate::auth::error::AuthConfigurationError::InvalidAzureAuthority
)
));
assert!(matches!(error, litellm_auth::Error::InvalidAzureAuthority));
}
}
}

93
litellm-rust/crates/core/src/providers/azure_ai/auth/resolve.rs → litellm-rust/crates/auth-azure/src/resolve.rs
View file

@ -1,6 +1,5 @@
use crate::AuthError;
use crate::auth::error::AuthConfigurationError;
use crate::auth::{
use litellm_auth::Error;
use litellm_auth::{
CredentialFileRef, CredentialLookup, CredentialRef, InputSource, ResolvedCredential,
SecretValue, Sourced, TokenProviderHandle,
};
@ -37,7 +36,7 @@ pub(crate) enum AzureCredentialPlan {
}
/// Rust counterpart to Python's `get_azure_ad_token`, not `BaseAzureLLM`.
pub(crate) struct AzureAuthService {
pub struct AzureAuthService {
native: Arc<dyn AzureTokenAcquirer>,
}
@ -45,14 +44,14 @@ trait AzureTokenAcquirer: Send + Sync {
fn acquire(
&self,
request: ValidatedAzureRequest,
) -> Pin<Box<dyn Future<Output = Result<ResolvedCredential, AuthError>> + Send + '_>>;
) -> Pin<Box<dyn Future<Output = Result<ResolvedCredential, Error>> + Send + '_>>;
}
impl AzureTokenAcquirer for NativeAzureTokenAcquirer {
fn acquire(
&self,
request: ValidatedAzureRequest,
) -> Pin<Box<dyn Future<Output = Result<ResolvedCredential, AuthError>> + Send + '_>> {
) -> Pin<Box<dyn Future<Output = Result<ResolvedCredential, Error>> + Send + '_>> {
Box::pin(NativeAzureTokenAcquirer::acquire(self, request))
}
}
@ -71,17 +70,17 @@ impl AzureAuthService {
Self { native }
}
pub(crate) async fn get_azure_ad_token(
pub async fn get_azure_ad_token(
&self,
inputs: &AzureAuthInputs,
env_lookup: &(dyn Fn(&str) -> Option<String> + Sync),
) -> Result<Option<Sourced<ResolvedCredential>>, AuthError> {
) -> Result<Option<Sourced<ResolvedCredential>>, Error> {
match select_auth_plan(inputs, env_lookup)? {
AzureCredentialPlan::Supplied(credential) => Ok(Some(credential)),
AzureCredentialPlan::Caller(caller) => {
let credential = caller.acquire().await?;
if credential.secret().expose().is_empty() {
return Err(AuthError::EmptyAzureToken);
return Err(Error::EmptyAzureToken);
}
Ok(Some(Sourced::new(credential, InputSource::Deployment)))
}
@ -94,7 +93,7 @@ impl AzureAuthService {
} => {
let assertion = resolve_reference(inputs, env_lookup, reference.value())
.await?
.ok_or(AuthError::UnresolvedOidcReference)?;
.ok_or(Error::UnresolvedOidcReference)?;
let request = ValidatedAzureRequest::new(NativeAzureRequest::ClientAssertion {
tenant_id,
client_id,
@ -126,7 +125,7 @@ impl AzureAuthService {
Err(error) => failures.push(error),
}
}
Err(AuthError::CredentialChain(failures))
Err(Error::CredentialChain(failures))
}
AzureCredentialPlan::Missing => Ok(None),
}
@ -136,7 +135,7 @@ impl AzureAuthService {
pub(crate) fn select_auth_plan(
inputs: &AzureAuthInputs,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Result<AzureCredentialPlan, AuthError> {
) -> Result<AzureCredentialPlan, Error> {
let token = configured_secret(&inputs.azure_ad_token, AZURE_AD_TOKEN_ENV, env_lookup);
let tenant_id = configured_string(&inputs.tenant_id, AZURE_TENANT_ID_ENV, env_lookup);
let client_id = configured_string(&inputs.client_id, AZURE_CLIENT_ID_ENV, env_lookup);
@ -157,7 +156,7 @@ pub(crate) fn select_auth_plan(
.map(|selector| Sourced::new(selector, value.source()))
})
.transpose()
.map_err(|_| AuthError::Configuration(AuthConfigurationError::InvalidAzureSelector))?;
.map_err(|_| Error::InvalidAzureSelector)?;
let federated_token_file = configured_string(
&inputs.federated_token_file,
AZURE_FEDERATED_TOKEN_FILE_ENV,
@ -229,7 +228,7 @@ fn select_native_plan(
scope: Sourced<String>,
authority: Option<Sourced<String>>,
refresh_source: InputSource,
) -> Result<AzureCredentialPlan, AuthError> {
) -> Result<AzureCredentialPlan, Error> {
let selected = selector.unwrap_or_else(|| {
Sourced::new(
{
@ -247,9 +246,7 @@ fn select_native_plan(
let selection_source = selected.source();
match selected.into_value() {
AzureCredentialType::ClientSecretCredential => Err(AuthError::Configuration(
AuthConfigurationError::MissingClientSecretFields,
)),
AzureCredentialType::ClientSecretCredential => Err(Error::MissingClientSecretFields),
AzureCredentialType::WorkloadIdentityCredential => {
Ok(AzureCredentialPlan::Native(ValidatedAzureRequest::new(
workload_request(tenant_id, client_id, federated_token_file, scope, authority)?,
@ -331,17 +328,11 @@ fn workload_request(
token_file_path: Option<Sourced<String>>,
scope: Sourced<String>,
authority: Option<Sourced<String>>,
) -> Result<NativeAzureRequest, AuthError> {
) -> Result<NativeAzureRequest, Error> {
Ok(NativeAzureRequest::WorkloadIdentity {
tenant_id: tenant_id.ok_or(AuthError::Configuration(
AuthConfigurationError::MissingWorkloadTenant,
))?,
client_id: client_id.ok_or(AuthError::Configuration(
AuthConfigurationError::MissingWorkloadClient,
))?,
token_file_path: token_file_path.ok_or(AuthError::Configuration(
AuthConfigurationError::MissingWorkloadTokenFile,
))?,
tenant_id: tenant_id.ok_or(Error::MissingWorkloadTenant)?,
client_id: client_id.ok_or(Error::MissingWorkloadClient)?,
token_file_path: token_file_path.ok_or(Error::MissingWorkloadTokenFile)?,
scope,
authority,
})
@ -383,7 +374,7 @@ async fn resolve_reference(
inputs: &AzureAuthInputs,
env_lookup: &(dyn Fn(&str) -> Option<String> + Sync),
reference: &CredentialRef,
) -> Result<Option<SecretValue>, AuthError> {
) -> Result<Option<SecretValue>, Error> {
let lookup = match reference {
CredentialRef::Explicit(secret) => return Ok(Some(secret.clone())),
CredentialRef::Env(name) => env_lookup(name)
@ -395,9 +386,7 @@ async fn resolve_reference(
let resolver = inputs
.credential_resolver
.as_ref()
.ok_or(AuthError::Configuration(
AuthConfigurationError::MissingHostResolver,
))?;
.ok_or(Error::MissingHostResolver)?;
resolver.resolve(reference).await?
}
};
@ -409,15 +398,13 @@ async fn resolve_reference(
fn oidc_reference(
token: &Option<Sourced<SecretValue>>,
) -> Result<Option<Sourced<CredentialRef>>, AuthError> {
) -> Result<Option<Sourced<CredentialRef>>, Error> {
let Some(token) = token.as_ref() else {
return Ok(None);
};
let value = token.value().expose();
if token.source() == InputSource::Request && value.starts_with("oidc/") {
return Err(AuthError::Configuration(
AuthConfigurationError::RequestAzureCredentialReference,
));
return Err(Error::RequestAzureCredentialReference);
}
if let Some(name) = value.strip_prefix("oidc/env/") {
return non_empty_reference(name, "OIDC environment reference")
@ -439,18 +426,14 @@ fn oidc_reference(
)));
}
if value.starts_with("oidc/") {
return Err(AuthError::Configuration(
AuthConfigurationError::UnsupportedOidcReference,
));
return Err(Error::UnsupportedOidcReference);
}
Ok(None)
}
fn non_empty_reference(value: &str, kind: &str) -> Result<String, AuthError> {
fn non_empty_reference(value: &str, kind: &str) -> Result<String, Error> {
if value.is_empty() {
return Err(AuthError::Configuration(
AuthConfigurationError::EmptyReference(kind.to_string()),
));
return Err(Error::EmptyReference(kind.to_string()));
}
Ok(value.to_string())
}
@ -466,14 +449,14 @@ mod tests {
AzureAuthService, AzureCredentialPlan, AzureTokenAcquirer, oidc_reference,
resolve_reference, select_auth_plan,
};
use crate::AuthError;
use crate::auth::ResolvedCredential;
use crate::auth::{
use crate::native::ValidatedAzureRequest;
use crate::types::AzureAuthInputs;
use litellm_auth::Error;
use litellm_auth::ResolvedCredential;
use litellm_auth::{
CredentialFileRef, CredentialLookup, CredentialLookupFuture, CredentialRef,
CredentialResolver, CredentialResolverHandle, InputSource, SecretValue, Sourced,
};
use crate::providers::azure_ai::auth::native::ValidatedAzureRequest;
use crate::providers::azure_ai::auth::types::AzureAuthInputs;
#[derive(Debug)]
struct FileResolver;
@ -487,9 +470,8 @@ mod tests {
fn acquire(
&self,
request: ValidatedAzureRequest,
) -> std::pin::Pin<
Box<dyn Future<Output = Result<ResolvedCredential, AuthError>> + Send + '_>,
> {
) -> std::pin::Pin<Box<dyn Future<Output = Result<ResolvedCredential, Error>> + Send + '_>>
{
let kind = request.kind();
self.requests.lock().unwrap().push(kind);
Box::pin(async move {
@ -499,7 +481,7 @@ mod tests {
expires_on: None,
})
} else {
Err(AuthError::AzureTokenAcquisition(format!("{kind} failed")))
Err(Error::AzureTokenAcquisition(format!("{kind} failed")))
}
})
}
@ -612,12 +594,7 @@ mod tests {
})
.unwrap_err();
assert!(matches!(
error,
AuthError::Configuration(
crate::auth::error::AuthConfigurationError::RequestAzureCredentialReference
)
));
assert!(matches!(error, Error::RequestAzureCredentialReference));
}
#[tokio::test]
@ -678,6 +655,6 @@ mod tests {
.await
.unwrap_err();
assert!(matches!(error, AuthError::CredentialChain(errors) if errors.len() == 2));
assert!(matches!(error, Error::CredentialChain(errors) if errors.len() == 2));
}
}

19
litellm-rust/crates/core/src/providers/azure_ai/auth/types.rs → litellm-rust/crates/auth-azure/src/types.rs
View file

@ -1,10 +1,9 @@
use crate::auth::error::AuthConfigurationError;
use serde_json::{Map, Value};
use std::collections::BTreeMap;
use strum::EnumString;
use crate::AuthError;
use crate::auth::{
use litellm_auth::Error;
use litellm_auth::{
CredentialResolverHandle, InputSource, SecretValue, Sourced, TokenProviderHandle,
};
@ -54,14 +53,14 @@ pub struct AzureAuthInputs {
impl AzureAuthInputs {
#[cfg(test)]
pub fn from_optional_params(params: &Map<String, Value>) -> Result<Self, AuthError> {
pub fn from_optional_params(params: &Map<String, Value>) -> Result<Self, Error> {
Self::from_sourced_optional_params(params, &BTreeMap::new())
}
pub fn from_sourced_optional_params(
params: &Map<String, Value>,
sources: &BTreeMap<String, InputSource>,
) -> Result<Self, AuthError> {
) -> Result<Self, Error> {
Ok(Self {
azure_ad_token: secret_config(params, sources, "azure_ad_token")?,
azure_ad_token_provider: None,
@ -88,15 +87,13 @@ fn string_config(
params: &Map<String, Value>,
sources: &BTreeMap<String, InputSource>,
name: &str,
) -> Result<ConfigValue<String>, AuthError> {
) -> Result<ConfigValue<String>, Error> {
let source = source_for(sources, name);
match params.get(name) {
None => Ok(ConfigValue::Absent),
Some(Value::Null) => Ok(ConfigValue::ExplicitNone(source)),
Some(Value::String(value)) => Ok(ConfigValue::Value(Sourced::new(value.clone(), source))),
Some(_) => Err(AuthError::Configuration(
AuthConfigurationError::InvalidFieldType(name.to_string()),
)),
Some(_) => Err(Error::InvalidFieldType(name.to_string())),
}
}
@ -104,7 +101,7 @@ fn secret_config(
params: &Map<String, Value>,
sources: &BTreeMap<String, InputSource>,
name: &str,
) -> Result<ConfigValue<SecretValue>, AuthError> {
) -> Result<ConfigValue<SecretValue>, Error> {
Ok(match string_config(params, sources, name)? {
ConfigValue::Absent => ConfigValue::Absent,
ConfigValue::ExplicitNone(source) => ConfigValue::ExplicitNone(source),
@ -123,7 +120,7 @@ mod tests {
use std::collections::BTreeMap;
use super::{AzureAuthInputs, AzureCredentialType, ConfigValue};
use crate::auth::{InputSource, Sourced};
use litellm_auth::{InputSource, Sourced};
#[test]
fn selector_parsing_is_exact() {

16
litellm-rust/crates/auth-gcp/Cargo.toml Normal file
View file

@ -0,0 +1,16 @@
[package]
name = "litellm-auth-gcp"
version = "0.1.0"
edition.workspace = true
license.workspace = true
repository.workspace = true
[dependencies]
litellm-auth.workspace = true
moka.workspace = true
serde_json.workspace = true
sha2.workspace = true
tokio.workspace = true
gcp_auth = "0.12.7"

71
litellm-rust/crates/core/src/auth/vertex.rs → litellm-rust/crates/auth-gcp/src/lib.rs
View file

@ -9,9 +9,8 @@ use moka::future::Cache;
use serde_json::{Map, Value};
use sha2::{Digest, Sha256};
use crate::auth::error::AuthConfigurationError;
use crate::auth::http::apply_credential;
use crate::auth::{AuthError, CredentialPlacement, InputSource, SecretValue, Sourced};
use litellm_auth::http::apply_credential;
use litellm_auth::{CredentialPlacement, Error, InputSource, SecretValue, Sourced};
const CLOUD_PLATFORM_SCOPE: &str = "https://www.googleapis.com/auth/cloud-platform";
const GOOGLE_OAUTH_TOKEN_ENDPOINT: &str = "https://oauth2.googleapis.com/token";
@ -24,17 +23,17 @@ const VERTEXAI_LOCATION_ENV: &str = "VERTEXAI_LOCATION";
const VERTEX_LOCATION_ENV: &str = "VERTEX_LOCATION";
#[derive(Clone, Debug, Default)]
pub(crate) struct VertexConfig {
pub struct VertexConfig {
credentials: Option<Sourced<SecretValue>>,
project_id: Option<String>,
location: Option<String>,
}
impl VertexConfig {
pub(crate) fn from_sourced_optional_params(
pub fn from_sourced_optional_params(
params: &Map<String, Value>,
sources: &BTreeMap<String, InputSource>,
) -> Result<Self, AuthError> {
) -> Result<Self, Error> {
Ok(Self {
credentials: optional_credentials(
params,
@ -46,16 +45,16 @@ impl VertexConfig {
})
}
pub(crate) fn project_id(&self) -> Option<&str> {
pub fn project_id(&self) -> Option<&str> {
self.project_id.as_deref()
}
pub(crate) fn location(&self) -> Option<&str> {
pub fn location(&self) -> Option<&str> {
self.location.as_deref()
}
}
pub(crate) struct VertexEnvironment {
pub struct VertexEnvironment {
pub headers: Vec<(String, String)>,
pub project_id: String,
}
@ -65,7 +64,7 @@ struct VertexAccessToken {
project_id: String,
}
pub(crate) fn get_vertex_ai_project(
pub fn get_vertex_ai_project(
config: &VertexConfig,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Option<String> {
@ -75,7 +74,7 @@ pub(crate) fn get_vertex_ai_project(
.or_else(|| non_empty_env(env_lookup, VERTEXAI_PROJECT_ENV))
}
pub(crate) fn get_vertex_ai_location(
pub fn get_vertex_ai_location(
config: &VertexConfig,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Option<String> {
@ -87,7 +86,7 @@ pub(crate) fn get_vertex_ai_location(
}
#[derive(Clone)]
pub(crate) struct VertexAuth {
pub struct VertexAuth {
providers: Cache<CredentialCacheKey, Arc<dyn VertexTokenSource>>,
loader: Arc<dyn VertexProviderLoader>,
}
@ -106,13 +105,13 @@ impl VertexAuth {
}
}
pub(crate) async fn validate_environment(
pub async fn validate_environment(
&self,
headers: Vec<(String, String)>,
api_key: Option<&str>,
config: &VertexConfig,
env_lookup: &(dyn Fn(&str) -> Option<String> + Sync),
) -> Result<VertexEnvironment, AuthError> {
) -> Result<VertexEnvironment, Error> {
let has_authorization = headers
.iter()
.any(|(name, _)| name.eq_ignore_ascii_case("Authorization"));
@ -160,7 +159,7 @@ impl VertexAuth {
&self,
config: &VertexConfig,
env_lookup: &(dyn Fn(&str) -> Option<String> + Sync),
) -> Result<VertexAccessToken, AuthError> {
) -> Result<VertexAccessToken, Error> {
let provider = self.load_provider(config, env_lookup).await?;
let (token, project_id) = tokio::try_join!(provider.token(), provider.project_id())?;
Ok(VertexAccessToken { token, project_id })
@ -170,7 +169,7 @@ impl VertexAuth {
&self,
config: &VertexConfig,
env_lookup: &(dyn Fn(&str) -> Option<String> + Sync),
) -> Result<Arc<dyn VertexTokenSource>, AuthError> {
) -> Result<Arc<dyn VertexTokenSource>, Error> {
let source = credential_source(config, env_lookup);
let key = source.cache_key();
self.providers
@ -189,7 +188,7 @@ trait VertexProviderLoader: Send + Sync {
fn load(&self, source: CredentialSource) -> VertexAuthFuture<'_, Arc<dyn VertexTokenSource>>;
}
type VertexAuthFuture<'a, T> = Pin<Box<dyn Future<Output = Result<T, AuthError>> + Send + 'a>>;
type VertexAuthFuture<'a, T> = Pin<Box<dyn Future<Output = Result<T, Error>> + Send + 'a>>;
struct GcpTokenSource(Arc<dyn TokenProvider>);
@ -249,7 +248,7 @@ impl VertexProviderLoader for GcpProviderLoader {
}
}
fn validate_request_credentials(configured: &str) -> Result<&str, AuthError> {
fn validate_request_credentials(configured: &str) -> Result<&str, Error> {
let token_uri = serde_json::from_str::<Value>(configured)
.ok()
.and_then(|credentials| {
@ -259,7 +258,7 @@ fn validate_request_credentials(configured: &str) -> Result<&str, AuthError> {
.map(str::to_string)
});
if token_uri.as_deref() != Some(GOOGLE_OAUTH_TOKEN_ENDPOINT) {
return Err(AuthConfigurationError::RequestVertexTokenEndpoint.into());
return Err(Error::RequestVertexTokenEndpoint);
}
Ok(configured)
}
@ -321,7 +320,7 @@ fn optional_credentials(
params: &Map<String, Value>,
sources: &BTreeMap<String, InputSource>,
names: &[&str],
) -> Result<Option<Sourced<SecretValue>>, AuthError> {
) -> Result<Option<Sourced<SecretValue>>, Error> {
for name in names {
let source = source_for(sources, name);
match params.get(*name) {
@ -336,17 +335,10 @@ fn optional_credentials(
.map(SecretValue::new)
.map(|value| Sourced::new(value, source))
.map(Some)
.map_err(|error| {
AuthError::Configuration(AuthConfigurationError::InvalidFieldType(format!(
"{}: {error}",
names[0]
)))
});
.map_err(|error| Error::InvalidFieldType(format!("{}: {error}", names[0])));
}
Some(_) => {
return Err(AuthError::Configuration(
AuthConfigurationError::InvalidFieldType(names[0].to_string()),
));
return Err(Error::InvalidFieldType(names[0].to_string()));
}
}
}
@ -357,19 +349,14 @@ fn source_for(sources: &BTreeMap<String, InputSource>, name: &str) -> InputSourc
sources.get(name).copied().unwrap_or_default()
}
fn optional_string(
params: &Map<String, Value>,
names: &[&str],
) -> Result<Option<String>, AuthError> {
fn optional_string(params: &Map<String, Value>, names: &[&str]) -> Result<Option<String>, Error> {
for name in names {
match params.get(*name) {
None | Some(Value::Null) => continue,
Some(Value::String(value)) if value.trim().is_empty() => continue,
Some(Value::String(value)) => return Ok(Some(value.clone())),
Some(_) => {
return Err(AuthError::Configuration(
AuthConfigurationError::InvalidFieldType(names[0].to_string()),
));
return Err(Error::InvalidFieldType(names[0].to_string()));
}
}
}
@ -382,8 +369,8 @@ fn non_empty_env(env_lookup: &dyn Fn(&str) -> Option<String>, name: &str) -> Opt
.filter(|value| !value.is_empty())
}
fn auth_acquisition_error(error: gcp_auth::Error) -> AuthError {
AuthError::VertexTokenAcquisition(error.to_string())
fn auth_acquisition_error(error: gcp_auth::Error) -> Error {
Error::VertexTokenAcquisition(error.to_string())
}
#[cfg(test)]
@ -537,15 +524,11 @@ mod tests {
);
assert!(matches!(
validate_request_credentials(r#"{"token_uri":"http://127.0.0.1/token"}"#),
Err(AuthError::Configuration(
AuthConfigurationError::RequestVertexTokenEndpoint
))
Err(Error::RequestVertexTokenEndpoint)
));
assert!(matches!(
validate_request_credentials("{}"),
Err(AuthError::Configuration(
AuthConfigurationError::RequestVertexTokenEndpoint
))
Err(Error::RequestVertexTokenEndpoint)
));
}

15
litellm-rust/crates/auth/Cargo.toml Normal file
View file

@ -0,0 +1,15 @@
[package]
name = "litellm-auth"
version = "0.1.0"
edition.workspace = true
license.workspace = true
repository.workspace = true
[dependencies]
serde.workspace = true
subtle.workspace = true
thiserror.workspace = true
veil.workspace = true
[dev-dependencies]
tokio.workspace = true

18
litellm-rust/crates/core/src/auth/credential.rs → litellm-rust/crates/auth/src/credential.rs
View file

@ -5,7 +5,7 @@ use std::sync::Arc;
use veil::Redact;
use crate::AuthError;
use crate::Error;
use super::{ResolvedCredential, SecretValue, TokenProviderHandle};
@ -48,7 +48,7 @@ pub enum CredentialLookup {
}
pub type CredentialLookupFuture<'a> =
Pin<Box<dyn Future<Output = Result<CredentialLookup, AuthError>> + Send + 'a>>;
Pin<Box<dyn Future<Output = Result<CredentialLookup, Error>> + Send + 'a>>;
pub trait CredentialResolver: std::fmt::Debug + Send + Sync {
fn resolve<'a>(&'a self, reference: &'a CredentialRef) -> CredentialLookupFuture<'a>;
@ -62,7 +62,7 @@ impl CredentialResolverHandle {
Self(resolver)
}
pub async fn resolve(&self, reference: &CredentialRef) -> Result<CredentialLookup, AuthError> {
pub async fn resolve(&self, reference: &CredentialRef) -> Result<CredentialLookup, Error> {
self.0.resolve(reference).await
}
}
@ -84,7 +84,7 @@ impl CredentialPlan {
pub async fn resolve(
&self,
resolver: &CredentialResolverHandle,
) -> Result<CredentialPlanResolution, AuthError> {
) -> Result<CredentialPlanResolution, Error> {
match self {
Self::Static(CredentialRef::Explicit(secret)) => Ok(
CredentialPlanResolution::Resolved(ResolvedCredential::Static(secret.clone())),
@ -103,7 +103,7 @@ impl CredentialPlan {
Self::Caller(caller) => {
let credential = caller.acquire().await?;
if credential.secret().expose().is_empty() {
return Err(AuthError::EmptyCallerCredential);
return Err(Error::EmptyCallerCredential);
}
Ok(CredentialPlanResolution::Resolved(credential))
}
@ -119,8 +119,8 @@ mod tests {
CredentialLookup, CredentialLookupFuture, CredentialPlan, CredentialPlanResolution,
CredentialRef, CredentialResolver, CredentialResolverHandle,
};
use crate::AuthError;
use crate::auth::SecretValue;
use crate::Error;
use crate::SecretValue;
#[derive(Debug)]
struct HostResolver;
@ -164,7 +164,7 @@ mod tests {
impl CredentialResolver for FailingResolver {
fn resolve<'a>(&'a self, _reference: &'a CredentialRef) -> CredentialLookupFuture<'a> {
Box::pin(async { Err(AuthError::UnresolvedOidcReference) })
Box::pin(async { Err(Error::UnresolvedOidcReference) })
}
}
@ -178,6 +178,6 @@ mod tests {
.await
.expect_err("acquisition errors cannot become fallback");
assert_eq!(error, AuthError::UnresolvedOidcReference);
assert_eq!(error, Error::UnresolvedOidcReference);
}
}

120
litellm-rust/crates/auth/src/error.rs Normal file
View file

@ -0,0 +1,120 @@
use thiserror::Error as ThisError;
#[derive(Clone, Debug, ThisError, PartialEq, Eq)]
pub enum Error {
#[error("invalid authentication configuration: credential header already exists")]
ExistingCredentialHeader,
#[error(
"invalid authentication configuration: credential plan is not allowed by the provider auth policy"
)]
DisallowedCredentialPlan,
#[error("invalid authentication configuration: credential cannot be empty")]
EmptyCredential,
#[error("invalid authentication configuration: invalid Azure credential selector")]
InvalidAzureSelector,
#[error(
"invalid authentication configuration: ClientSecretCredential requires tenant_id, client_id, and client_secret"
)]
MissingClientSecretFields,
#[error("invalid authentication configuration: WorkloadIdentityCredential requires tenant_id")]
MissingWorkloadTenant,
#[error("invalid authentication configuration: WorkloadIdentityCredential requires client_id")]
MissingWorkloadClient,
#[error(
"invalid authentication configuration: WorkloadIdentityCredential requires azure_federated_token_file"
)]
MissingWorkloadTokenFile,
#[error(
"invalid authentication configuration: credential reference requires a host credential resolver"
)]
MissingHostResolver,
#[error(
"invalid authentication configuration: caller credential plan requires provider-specific inputs"
)]
MissingCallerInputs,
#[error("invalid authentication configuration: credential header {0} already exists")]
DuplicateHeader(&'static str),
#[error("invalid authentication configuration: {0} must be a string or null")]
InvalidFieldType(String),
#[error("invalid authentication configuration: unsupported OIDC reference")]
UnsupportedOidcReference,
#[error("invalid authentication configuration: {0} cannot be empty")]
EmptyReference(String),
#[error("invalid authentication configuration: Azure credential initialization failed: {0}")]
AzureCredentialInitialization(String),
#[error(
"invalid authentication configuration: Azure authority must be an HTTPS origin without credentials, query, or fragment"
)]
InvalidAzureAuthority,
#[error(
"invalid authentication configuration: request-controlled Azure auth inputs cannot be combined with host credentials"
)]
MixedAzureCredentialSources,
#[error(
"invalid authentication configuration: request-controlled Azure credential references are not allowed"
)]
RequestAzureCredentialReference,
#[error(
"invalid authentication configuration: host credentials cannot be sent to a request-controlled Azure endpoint"
)]
RequestAzureCredentialDestination,
#[error(
"invalid authentication configuration: credentials cannot be sent to a request-controlled Vertex AI endpoint"
)]
RequestVertexCredentialDestination,
#[error(
"invalid authentication configuration: request-controlled Vertex credentials must use the canonical Google OAuth token endpoint"
)]
RequestVertexTokenEndpoint,
#[error("credential acquisition failed: {0}")]
AzureTokenAcquisition(String),
#[error("credential acquisition failed: Vertex AI credentials: {0}")]
VertexTokenAcquisition(String),
#[error("{0}")]
ProviderAuthentication(String),
#[error("credential acquisition failed: {}", .0.iter().map(ToString::to_string).collect::<Vec<_>>().join("; "))]
CredentialChain(Vec<Error>),
#[error("credential caller failed: credential caller returned an empty credential")]
EmptyCallerCredential,
#[error("credential caller failed: Azure AD token provider returned an empty token")]
EmptyAzureToken,
#[error("credential acquisition failed: Azure OIDC reference did not resolve to a value")]
UnresolvedOidcReference,
#[error(
"Missing {provider} API Key - Set `api_key` or the {environment_variable} environment variable"
)]
MissingApiKey {
provider: &'static str,
environment_variable: &'static str,
},
#[error(
"Missing {provider} API Base - Set {environment_variable} environment variable or pass api_base parameter"
)]
MissingApiBase {
provider: &'static str,
environment_variable: &'static str,
},
#[error(
"Missing Azure API Base - Set `api_base` or the AZURE_API_BASE environment variable. Expected format: https://<resource-name>.services.ai.azure.com/anthropic"
)]
MissingAzureApiBase,
#[error("invalid authentication header")]
InvalidHeader,
}
#[cfg(test)]
mod tests {
use super::Error;
#[test]
fn missing_api_key_names_provider_and_environment_variable() {
assert_eq!(
Error::MissingApiKey {
provider: "Anthropic",
environment_variable: "ANTHROPIC_API_KEY",
}
.to_string(),
"Missing Anthropic API Key - Set `api_key` or the ANTHROPIC_API_KEY environment variable"
);
}
}

15
litellm-rust/crates/core/src/auth/http.rs → litellm-rust/crates/auth/src/http.rs
View file

@ -1,5 +1,4 @@
use crate::AuthError;
use crate::auth::error::AuthConfigurationError;
use crate::Error;
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum CredentialPlacement {
@ -16,23 +15,19 @@ impl CredentialPlacement {
}
}
pub(crate) fn apply_credential(
pub fn apply_credential(
headers: Vec<(String, String)>,
credential: &str,
placement: CredentialPlacement,
) -> Result<Vec<(String, String)>, AuthError> {
) -> Result<Vec<(String, String)>, Error> {
if credential.trim().is_empty() {
return Err(AuthError::Configuration(
AuthConfigurationError::EmptyCredential,
));
return Err(Error::EmptyCredential);
}
if headers
.iter()
.any(|(name, _)| name.eq_ignore_ascii_case(placement.header_name()))
{
return Err(AuthError::Configuration(
AuthConfigurationError::DuplicateHeader(placement.header_name()),
));
return Err(Error::DuplicateHeader(placement.header_name()));
}
let value = match placement {
CredentialPlacement::Bearer => format!("Bearer {credential}"),

7
litellm-rust/crates/core/src/auth/mod.rs → litellm-rust/crates/auth/src/lib.rs
View file

@ -1,8 +1,6 @@
mod credential;
pub mod error;
pub(crate) mod vertex;
pub use error::AuthError;
pub(crate) mod http;
mod error;
pub mod http;
mod policy;
mod secret;
mod token;
@ -51,6 +49,7 @@ pub use credential::{
CredentialPlanResolution, CredentialRef, CredentialResolver, CredentialResolverHandle,
credential_default_fields, credential_index,
};
pub use error::Error;
pub use http::{CredentialPlacement, RequestAuth};
pub use policy::{CredentialPlanKind, CredentialRule, ExistingHeaderBehavior, ProviderAuthPolicy};
pub use secret::SecretValue;

23
litellm-rust/crates/core/src/auth/policy.rs → litellm-rust/crates/auth/src/policy.rs
View file

@ -1,5 +1,4 @@
use crate::AuthError;
use crate::auth::error::AuthConfigurationError;
use crate::Error;
use super::http::apply_credential;
use super::{CredentialPlacement, ResolvedCredential};
@ -46,22 +45,18 @@ impl ProviderAuthPolicy {
headers: Vec<(String, String)>,
kind: CredentialPlanKind,
credential: &ResolvedCredential,
) -> Result<Vec<(String, String)>, AuthError> {
) -> Result<Vec<(String, String)>, Error> {
if self.has_existing_credential(&headers) {
return match self.existing_header_behavior {
ExistingHeaderBehavior::Preserve => Ok(headers),
ExistingHeaderBehavior::Reject => Err(AuthError::Configuration(
AuthConfigurationError::ExistingCredentialHeader,
)),
ExistingHeaderBehavior::Reject => Err(Error::ExistingCredentialHeader),
};
}
let rule =
self.rules
.iter()
.find(|rule| rule.kind == kind)
.ok_or(AuthError::Configuration(
AuthConfigurationError::DisallowedCredentialPlan,
))?;
let rule = self
.rules
.iter()
.find(|rule| rule.kind == kind)
.ok_or(Error::DisallowedCredentialPlan)?;
apply_credential(headers, credential.secret().expose(), rule.placement)
}
}
@ -69,7 +64,7 @@ impl ProviderAuthPolicy {
#[cfg(test)]
mod tests {
use super::{CredentialPlanKind, CredentialRule, ExistingHeaderBehavior, ProviderAuthPolicy};
use crate::auth::{CredentialPlacement, ResolvedCredential, SecretValue};
use crate::{CredentialPlacement, ResolvedCredential, SecretValue};
const RULES: &[CredentialRule] = &[CredentialRule {
kind: CredentialPlanKind::Static,

0
litellm-rust/crates/core/src/auth/secret.rs → litellm-rust/crates/auth/src/secret.rs
View file

6
litellm-rust/crates/core/src/auth/token.rs → litellm-rust/crates/auth/src/token.rs
View file

@ -5,7 +5,7 @@ use std::time::SystemTime;
use veil::Redact;
use crate::AuthError;
use crate::Error;
use super::secret::SecretValue;
@ -27,7 +27,7 @@ impl ResolvedCredential {
}
pub type TokenFuture<'a> =
Pin<Box<dyn Future<Output = Result<ResolvedCredential, AuthError>> + Send + 'a>>;
Pin<Box<dyn Future<Output = Result<ResolvedCredential, Error>> + Send + 'a>>;
pub trait TokenProvider: std::fmt::Debug + Send + Sync {
fn acquire(&self) -> TokenFuture<'_>;
@ -41,7 +41,7 @@ impl TokenProviderHandle {
Self(caller)
}
pub async fn acquire(&self) -> Result<ResolvedCredential, AuthError> {
pub async fn acquire(&self) -> Result<ResolvedCredential, Error> {
self.0.acquire().await
}
}

14
litellm-rust/crates/cache-memory/Cargo.toml Normal file
View file

@ -0,0 +1,14 @@
[package]
name = "litellm-cache-memory"
version = "0.1.0"
edition.workspace = true
license.workspace = true
repository.workspace = true
[dependencies]
litellm-cache.workspace = true
serde_json.workspace = true
[dev-dependencies]
rstest.workspace = true
tokio.workspace = true

254
litellm-rust/crates/cache-memory/src/cache.rs Normal file
View file

@ -0,0 +1,254 @@
use std::cmp::Reverse;
use std::collections::{BinaryHeap, HashMap};
use std::sync::{Arc, Mutex};
use std::time::{Duration, SystemTime, UNIX_EPOCH};
use litellm_cache::{
BaseCache, CacheConnectionResult, CacheConnectionStatus, CacheEntry, CacheFuture, CacheKwargs,
Error,
};
const DEFAULT_MAX_SIZE_IN_MEMORY: usize = 200;
const DEFAULT_TTL: Duration = Duration::from_secs(600);
type ValueMeasure<V> = Arc<dyn Fn(&V) -> Result<usize, Error> + Send + Sync>;
type ValueValidator<V> = Arc<dyn Fn(&V) -> Result<(), Error> + Send + Sync>;
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum CacheWrite {
Stored,
Disabled,
TooLarge,
}
struct CacheState<V> {
values: HashMap<String, V>,
expirations: HashMap<String, Duration>,
expiration_heap: BinaryHeap<Reverse<(Duration, String)>>,
}
pub struct InMemoryCache<V: Clone> {
state: Mutex<CacheState<V>>,
max_size_in_memory: usize,
default_ttl: Duration,
max_entry_bytes: Option<usize>,
measure_value: Option<ValueMeasure<V>>,
validate_value: Option<ValueValidator<V>>,
now: Arc<dyn Fn() -> Duration + Send + Sync>,
}
impl<V: Clone> Default for InMemoryCache<V> {
fn default() -> Self {
Self::new(None, None)
}
}
impl<V: Clone> InMemoryCache<V> {
pub fn new(max_size_in_memory: Option<usize>, default_ttl: Option<Duration>) -> Self {
Self::with_clock(max_size_in_memory, default_ttl, || {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap_or_default()
})
}
pub fn with_clock(
max_size_in_memory: Option<usize>,
default_ttl: Option<Duration>,
now: impl Fn() -> Duration + Send + Sync + 'static,
) -> Self {
Self::with_clock_and_size_measurement(max_size_in_memory, default_ttl, None, None, now)
}
pub fn with_clock_and_size_measurement(
max_size_in_memory: Option<usize>,
default_ttl: Option<Duration>,
max_entry_bytes: Option<usize>,
measure_value: Option<ValueMeasure<V>>,
now: impl Fn() -> Duration + Send + Sync + 'static,
) -> Self {
Self {
state: Mutex::new(CacheState {
values: HashMap::new(),
expirations: HashMap::new(),
expiration_heap: BinaryHeap::new(),
}),
max_size_in_memory: max_size_in_memory.unwrap_or(DEFAULT_MAX_SIZE_IN_MEMORY),
default_ttl: default_ttl.unwrap_or(DEFAULT_TTL),
max_entry_bytes,
measure_value,
validate_value: None,
now: Arc::new(now),
}
}
pub fn set_cache(
&self,
key: impl Into<String>,
value: V,
ttl: Option<Duration>,
) -> Result<CacheWrite, Error> {
if self.max_size_in_memory == 0 {
return Ok(CacheWrite::Disabled);
}
if let Some(validate) = &self.validate_value {
validate(&value)?;
}
if let (Some(limit), Some(measure)) = (self.max_entry_bytes, &self.measure_value)
&& measure(&value)? > limit
{
return Ok(CacheWrite::TooLarge);
}
let now = (self.now)();
let mut state = self.state.lock().map_err(|_| Error::Unavailable)?;
Self::evict(&mut state, self.max_size_in_memory, now);
let key = key.into();
state.values.insert(key.clone(), value);
let expiration = state.expirations.get(&key).copied();
if expiration.is_none_or(|expiration| expiration < now) {
let expiration = now + ttl.unwrap_or(self.default_ttl);
state.expirations.insert(key.clone(), expiration);
state.expiration_heap.push(Reverse((expiration, key)));
}
Ok(CacheWrite::Stored)
}
pub fn get_cache(&self, key: &str) -> Result<Option<V>, Error> {
let now = (self.now)();
let mut state = self.state.lock().map_err(|_| Error::Unavailable)?;
if state
.expirations
.get(key)
.is_some_and(|expiration| *expiration < now)
{
Self::remove(&mut state, key);
}
Ok(state.values.get(key).cloned())
}
pub fn expires_at(&self, key: &str) -> Result<Option<Duration>, Error> {
Ok(self
.state
.lock()
.map_err(|_| Error::Unavailable)?
.expirations
.get(key)
.copied())
}
pub fn delete_cache(&self, key: &str) -> Result<(), Error> {
let mut state = self.state.lock().map_err(|_| Error::Unavailable)?;
Self::remove(&mut state, key);
Ok(())
}
pub fn flush_cache(&self) -> Result<(), Error> {
let mut state = self.state.lock().map_err(|_| Error::Unavailable)?;
state.values.clear();
state.expirations.clear();
state.expiration_heap.clear();
Ok(())
}
fn evict(state: &mut CacheState<V>, capacity: usize, now: Duration) {
while let Some(Reverse((expiration, key))) = state.expiration_heap.peek().cloned() {
if state.expirations.get(&key).copied() != Some(expiration) {
state.expiration_heap.pop();
} else if expiration <= now {
state.expiration_heap.pop();
Self::remove(state, &key);
} else {
break;
}
}
while state.values.len() >= capacity {
let Some(Reverse((expiration, key))) = state.expiration_heap.pop() else {
break;
};
if state.expirations.get(&key).copied() == Some(expiration) {
Self::remove(state, &key);
}
}
}
fn remove(state: &mut CacheState<V>, key: &str) {
state.values.remove(key);
state.expirations.remove(key);
}
}
impl InMemoryCache<CacheEntry> {
pub fn response_cache(capacity: usize, ttl: Duration, max_entry_bytes: usize) -> Self {
Self::response_cache_with_clock(capacity, ttl, max_entry_bytes, || {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap_or_default()
})
}
pub fn response_cache_with_clock(
capacity: usize,
ttl: Duration,
max_entry_bytes: usize,
now: impl Fn() -> Duration + Send + Sync + 'static,
) -> Self {
let mut cache = Self::with_clock_and_size_measurement(
Some(capacity),
Some(ttl),
Some(max_entry_bytes),
Some(Arc::new(|entry: &CacheEntry| {
serde_json::to_vec(entry)
.map(|bytes| bytes.len())
.map_err(|_| Error::InvalidEntry)
})),
now,
);
cache.validate_value = Some(Arc::new(|entry: &CacheEntry| {
entry
.timestamp
.is_finite()
.then_some(())
.ok_or(Error::InvalidEntry)
}));
cache
}
}
impl BaseCache for InMemoryCache<CacheEntry> {
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 ttl = self.get_ttl(&kwargs);
self.set_cache(key, value, Some(ttl)).map(|_| ())
}
fn get_cache(&self, key: &str, _: &CacheKwargs) -> Result<Option<Self::Value>, Error> {
self.get_cache(key)
}
fn delete_cache(&self, key: &str) -> Result<(), Error> {
self.delete_cache(key)
}
fn flush_cache(&self) -> Result<(), Error> {
self.flush_cache()
}
fn disconnect(&self) -> CacheFuture<'_, ()> {
Box::pin(async { Ok(()) })
}
fn test_connection(&self) -> CacheFuture<'_, CacheConnectionResult> {
Box::pin(async {
Ok(CacheConnectionResult {
status: CacheConnectionStatus::Success,
message: "In-memory cache connection test successful".into(),
error: None,
})
})
}
}

3
litellm-rust/crates/cache-memory/src/lib.rs Normal file
View file

@ -0,0 +1,3 @@
mod cache;
pub use cache::{CacheWrite, InMemoryCache};

158
litellm-rust/crates/cache-memory/tests/cache.rs Normal file
View file

@ -0,0 +1,158 @@
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::Duration;
use litellm_cache::{BaseCache, CacheConnectionStatus, CacheEntry, Error};
use litellm_cache_memory::{CacheWrite, InMemoryCache};
use rstest::{fixture, rstest};
#[fixture]
fn clock() -> Arc<AtomicU64> {
Arc::new(AtomicU64::new(100))
}
fn cache(clock: Arc<AtomicU64>, capacity: usize) -> InMemoryCache<String> {
InMemoryCache::with_clock(Some(capacity), Some(Duration::from_secs(60)), move || {
Duration::from_secs(clock.load(Ordering::SeqCst))
})
}
#[rstest]
fn default_explicit_and_override_ttls_follow_python_rules(clock: Arc<AtomicU64>) {
let cache = cache(clock.clone(), 4);
cache.set_cache("key", "first".into(), None).unwrap();
assert_eq!(
cache.expires_at("key").unwrap(),
Some(Duration::from_secs(160))
);
cache
.set_cache("key", "second".into(), Some(Duration::from_secs(10)))
.unwrap();
assert_eq!(
cache.expires_at("key").unwrap(),
Some(Duration::from_secs(160))
);
clock.store(160, Ordering::SeqCst);
assert_eq!(cache.get_cache("key").unwrap(), Some("second".into()));
clock.store(161, Ordering::SeqCst);
assert_eq!(cache.get_cache("key").unwrap(), None);
cache
.set_cache("key", "third".into(), Some(Duration::from_secs(10)))
.unwrap();
assert_eq!(
cache.expires_at("key").unwrap(),
Some(Duration::from_secs(171))
);
}
#[rstest]
fn write_at_expiry_boundary_refreshes_ttl(clock: Arc<AtomicU64>) {
let cache = cache(clock.clone(), 4);
cache
.set_cache("key", "first".into(), Some(Duration::from_secs(10)))
.unwrap();
clock.store(110, Ordering::SeqCst);
cache
.set_cache("key", "second".into(), Some(Duration::from_secs(10)))
.unwrap();
assert_eq!(
cache.expires_at("key").unwrap(),
Some(Duration::from_secs(120))
);
clock.store(115, Ordering::SeqCst);
assert_eq!(cache.get_cache("key").unwrap(), Some("second".into()));
}
#[rstest]
fn capacity_evicts_earliest_and_ignores_stale_heap_entries(clock: Arc<AtomicU64>) {
let cache = cache(clock, 2);
cache
.set_cache("early", "a".into(), Some(Duration::from_secs(10)))
.unwrap();
cache
.set_cache("late", "b".into(), Some(Duration::from_secs(20)))
.unwrap();
cache.delete_cache("early").unwrap();
cache
.set_cache("new", "c".into(), Some(Duration::from_secs(30)))
.unwrap();
assert_eq!(cache.get_cache("late").unwrap(), Some("b".into()));
cache
.set_cache("last", "d".into(), Some(Duration::from_secs(40)))
.unwrap();
assert_eq!(cache.get_cache("late").unwrap(), None);
}
#[test]
fn disabled_size_limited_and_synchronized_response_writes_are_observable() {
let disabled = InMemoryCache::<CacheEntry>::response_cache(0, Duration::from_secs(60), 80);
assert_eq!(
disabled
.set_cache(
"a",
CacheEntry {
timestamp: 1.0,
response: serde_json::json!("x")
},
None
)
.unwrap(),
CacheWrite::Disabled
);
let cache = InMemoryCache::<CacheEntry>::response_cache(2, Duration::from_secs(60), 80);
assert_eq!(
cache
.set_cache(
"large",
CacheEntry {
timestamp: 1.0,
response: serde_json::json!("x".repeat(100))
},
None
)
.unwrap(),
CacheWrite::TooLarge
);
cache
.set_cache(
"small",
CacheEntry {
timestamp: 1.0,
response: serde_json::json!("ok"),
},
None,
)
.unwrap();
assert!(cache.get_cache("small").unwrap().is_some());
assert_eq!(
cache
.set_cache(
"invalid",
CacheEntry {
timestamp: f64::NAN,
response: serde_json::json!("bad"),
},
None,
)
.unwrap_err(),
Error::InvalidEntry
);
cache.delete_cache("small").unwrap();
cache.flush_cache().unwrap();
}
#[tokio::test]
async fn connection_test_matches_python_result_contract() {
let cache = InMemoryCache::<CacheEntry>::default();
let result = BaseCache::test_connection(&cache).await.unwrap();
assert_eq!(result.status, CacheConnectionStatus::Success);
assert_eq!(result.message, "In-memory cache connection test successful");
assert_eq!(result.error, None);
assert_eq!(
serde_json::to_value(result).unwrap(),
serde_json::json!({
"status": "success",
"message": "In-memory cache connection test successful"
})
);
}

15
litellm-rust/crates/cache-redis/Cargo.toml Normal file
View 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
View 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
View file

@ -0,0 +1,3 @@
mod cache;
pub use cache::RedisCache;

6
litellm-rust/crates/cache-redis/tests/cache.rs Normal file
View 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
litellm-rust/crates/cache/Cargo.toml vendored Normal file
View file

@ -0,0 +1,15 @@
[package]
name = "litellm-cache"
version = "0.1.0"
edition.workspace = true
license.workspace = true
repository.workspace = true
[dependencies]
serde.workspace = true
serde_json.workspace = true
sha2.workspace = true
thiserror.workspace = true
[dev-dependencies]
rstest.workspace = true

98
litellm-rust/crates/cache/src/base_cache.rs vendored Normal file
View file

@ -0,0 +1,98 @@
use std::future::Future;
use std::pin::Pin;
use std::time::Duration;
use serde::{Deserialize, Serialize};
use serde_json::{Map, Value};
use crate::Error;
pub type CacheFuture<'a, T> = Pin<Box<dyn Future<Output = Result<T, Error>> + Send + 'a>>;
#[derive(Clone, Debug, Default, PartialEq)]
pub struct CacheKwargs {
pub ttl: Option<Duration>,
pub extras: Map<String, Value>,
}
#[derive(Clone, Copy, Debug, Deserialize, Serialize, PartialEq, Eq)]
#[serde(rename_all = "lowercase")]
pub enum CacheConnectionStatus {
Success,
Failed,
}
#[derive(Clone, Debug, Deserialize, Serialize, PartialEq, Eq)]
pub struct CacheConnectionResult {
pub status: CacheConnectionStatus,
pub message: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub error: Option<String>,
}
pub trait BaseCache: Send + Sync {
type Value: Clone + Send + Sync + 'static;
fn default_ttl(&self) -> Duration {
Duration::from_secs(60)
}
fn get_ttl(&self, kwargs: &CacheKwargs) -> Duration {
kwargs.ttl.unwrap_or_else(|| self.default_ttl())
}
fn set_cache(&self, key: &str, value: Self::Value, kwargs: CacheKwargs) -> Result<(), Error>;
fn get_cache(&self, key: &str, kwargs: &CacheKwargs) -> Result<Option<Self::Value>, Error>;
fn async_set_cache<'a>(
&'a self,
key: &'a str,
value: Self::Value,
kwargs: CacheKwargs,
) -> CacheFuture<'a, ()> {
Box::pin(async move { self.set_cache(key, value, kwargs) })
}
fn async_get_cache<'a>(
&'a self,
key: &'a str,
kwargs: &'a CacheKwargs,
) -> CacheFuture<'a, Option<Self::Value>> {
Box::pin(async move { self.get_cache(key, kwargs) })
}
fn async_set_cache_pipeline<'a>(
&'a self,
cache_list: Vec<(String, Self::Value)>,
kwargs: CacheKwargs,
) -> CacheFuture<'a, ()> {
Box::pin(async move {
for (key, value) in cache_list {
self.set_cache(&key, value, kwargs.clone())?;
}
Ok(())
})
}
fn batch_cache_write<'a>(
&'a self,
key: &'a str,
value: Self::Value,
kwargs: CacheKwargs,
) -> CacheFuture<'a, ()> {
self.async_set_cache(key, value, kwargs)
}
fn delete_cache(&self, key: &str) -> Result<(), Error>;
fn async_delete_cache<'a>(&'a self, key: &'a str) -> CacheFuture<'a, ()> {
Box::pin(async move { self.delete_cache(key) })
}
fn flush_cache(&self) -> Result<(), Error>;
fn disconnect(&self) -> CacheFuture<'_, ()>;
fn test_connection(&self) -> CacheFuture<'_, CacheConnectionResult>;
}

166
litellm-rust/crates/cache/src/caching.rs vendored Normal file
View file

@ -0,0 +1,166 @@
use std::sync::Arc;
use std::time::Duration;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use sha2::{Digest, Sha256};
use crate::{BaseCache, CacheKwargs, Error};
pub use crate::BaseCache as Cache;
#[derive(Clone, Copy, Debug, Default, Deserialize, Serialize, PartialEq, Eq)]
pub enum CacheMode {
#[default]
#[serde(rename = "default_on")]
DefaultOn,
#[serde(rename = "default_off")]
DefaultOff,
}
#[derive(Clone, Debug, Deserialize, Serialize)]
pub struct CacheKeyField {
pub name: String,
pub value: Option<String>,
pub api_parameter: bool,
pub internal_parameter: bool,
}
#[derive(Clone, Debug, Default, Deserialize, Serialize)]
pub struct CacheKeyInput {
pub fields: Vec<CacheKeyField>,
pub preset: Option<String>,
pub namespace: Option<String>,
pub include_provider_parameters: bool,
}
#[derive(Default)]
pub struct CacheKeyContext {
pub model_group: Option<String>,
pub caching_groups: Vec<(Vec<String>, String)>,
pub file_checksum: Option<String>,
pub file_object_name: Option<String>,
pub metadata_file_name: Option<String>,
pub parameters_file_name: Option<String>,
}
impl CacheKeyContext {
pub fn apply(self, input: &mut CacheKeyInput) {
let group = self.model_group.as_ref().and_then(|model| {
self.caching_groups
.iter()
.find(|(models, _)| models.contains(model))
});
for field in &mut input.fields {
match field.name.as_str() {
"model" => {
field.value = group
.map(|(_, formatted)| formatted.clone())
.or_else(|| self.model_group.clone())
.or_else(|| field.value.take())
}
"file" => {
field.value = self
.file_checksum
.clone()
.or_else(|| self.file_object_name.clone())
.or_else(|| self.metadata_file_name.clone())
.or_else(|| self.parameters_file_name.clone())
}
_ => {}
}
}
}
}
pub fn get_cache_key(input: &CacheKeyInput) -> String {
cache_key(input)
}
pub fn cache_key(input: &CacheKeyInput) -> String {
if let Some(preset) = &input.preset {
return preset.clone();
}
let mut digest = Sha256::new();
for field in &input.fields {
if (field.api_parameter || (input.include_provider_parameters && !field.internal_parameter))
&& let Some(value) = &field.value
{
digest.update(field.name.as_bytes());
digest.update(b": ");
digest.update(value.as_bytes());
}
}
let hash = format!("{:x}", digest.finalize());
input
.namespace
.as_deref()
.filter(|namespace| !namespace.is_empty())
.map_or(hash.clone(), |namespace| format!("{namespace}:{hash}"))
}
#[derive(Clone, Copy, Debug, Default, Deserialize, Serialize)]
pub struct CacheControls {
pub supported_call_type: bool,
pub configured: bool,
pub native_backend: bool,
pub default_on: bool,
pub caching: Option<bool>,
pub no_cache: bool,
pub no_store: bool,
#[serde(default)]
pub use_cache: bool,
}
impl CacheControls {
pub fn reads(self) -> bool {
self.supported_call_type
&& self.configured
&& self.caching.unwrap_or(true)
&& !self.no_cache
&& (self.default_on || self.use_cache)
}
pub fn writes(self) -> bool {
self.supported_call_type
&& self.configured
&& !self.no_store
&& (self.default_on || self.use_cache)
}
}
pub fn should_use_cache(controls: CacheControls) -> bool {
controls.reads() || controls.writes()
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct CacheEntry {
pub timestamp: f64,
pub response: Value,
}
impl CacheEntry {
pub fn fresh(&self, now: Duration, max_age: Option<Duration>) -> bool {
self.timestamp.is_finite()
&& max_age.is_none_or(|age| now.as_secs_f64() - self.timestamp <= age.as_secs_f64())
}
}
pub fn get_cache(
cache: &dyn BaseCache<Value = CacheEntry>,
key: &str,
kwargs: &CacheKwargs,
) -> Result<Option<CacheEntry>, Error> {
cache.get_cache(key, kwargs)
}
pub fn set_cache(
cache: &dyn BaseCache<Value = CacheEntry>,
key: &str,
entry: CacheEntry,
kwargs: CacheKwargs,
) -> Result<(), Error> {
cache.set_cache(key, entry, kwargs)
}
pub type CacheBackend = Arc<dyn BaseCache<Value = CacheEntry>>;

7
litellm-rust/crates/cache/src/error.rs vendored Normal file
View file

@ -0,0 +1,7 @@
#[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)]
pub enum Error {
#[error("cache is unavailable")]
Unavailable,
#[error("invalid cache entry")]
InvalidEntry,
}

12
litellm-rust/crates/cache/src/lib.rs vendored Normal file
View file

@ -0,0 +1,12 @@
mod base_cache;
mod caching;
mod error;
pub use base_cache::{
BaseCache, CacheConnectionResult, CacheConnectionStatus, CacheFuture, CacheKwargs,
};
pub use caching::{
Cache, CacheBackend, CacheControls, CacheEntry, CacheKeyContext, CacheKeyField, CacheKeyInput,
CacheMode, cache_key, get_cache, get_cache_key, set_cache, should_use_cache,
};
pub use error::Error;

139
litellm-rust/crates/cache/tests/caching.rs vendored Normal file
View file

@ -0,0 +1,139 @@
use litellm_cache::{
BaseCache, CacheConnectionResult, CacheControls, CacheEntry, CacheFuture, CacheKeyContext,
CacheKeyField, CacheKeyInput, CacheKwargs, Error, cache_key, get_cache_key,
};
use sha2::{Digest, Sha256};
use std::time::Duration;
struct TestCache {
default_ttl: Duration,
}
impl BaseCache for TestCache {
type Value = CacheEntry;
fn default_ttl(&self) -> Duration {
self.default_ttl
}
fn set_cache(&self, _: &str, _: Self::Value, _: CacheKwargs) -> Result<(), Error> {
Ok(())
}
fn get_cache(&self, _: &str, _: &CacheKwargs) -> Result<Option<Self::Value>, Error> {
Ok(None)
}
fn delete_cache(&self, _: &str) -> Result<(), Error> {
Ok(())
}
fn flush_cache(&self) -> Result<(), Error> {
Ok(())
}
fn disconnect(&self) -> CacheFuture<'_, ()> {
Box::pin(async { Ok(()) })
}
fn test_connection(&self) -> CacheFuture<'_, CacheConnectionResult> {
unreachable!()
}
}
#[test]
fn ttl_uses_default_and_allows_per_call_override() {
let cache = TestCache {
default_ttl: Duration::from_secs(60),
};
assert_eq!(
cache.get_ttl(&CacheKwargs::default()),
Duration::from_secs(60)
);
assert_eq!(
cache.get_ttl(&CacheKwargs {
ttl: Some(Duration::from_secs(5)),
..Default::default()
}),
Duration::from_secs(5)
);
}
#[test]
fn keys_match_python_order_groups_files_presets_and_namespaces() {
let mut input = CacheKeyInput {
fields: vec![
CacheKeyField {
name: "model".into(),
value: Some("deployment".into()),
api_parameter: true,
internal_parameter: false,
},
CacheKeyField {
name: "file".into(),
value: None,
api_parameter: true,
internal_parameter: false,
},
],
namespace: Some("team".into()),
..Default::default()
};
CacheKeyContext {
model_group: Some("group".into()),
caching_groups: vec![(vec!["group".into()], "['group']".into())],
file_checksum: Some("checksum".into()),
..Default::default()
}
.apply(&mut input);
assert_eq!(
cache_key(&input),
format!(
"team:{:x}",
Sha256::digest(b"model: ['group']file: checksum")
)
);
input.preset = Some("preset".into());
assert_eq!(get_cache_key(&input), "preset");
}
#[test]
fn cache_controls_honor_default_modes_and_directives() {
let enabled = CacheControls {
supported_call_type: true,
configured: true,
default_on: true,
..Default::default()
};
assert!(enabled.reads());
assert!(enabled.writes());
assert!(
!CacheControls {
default_on: false,
..enabled
}
.reads()
);
assert!(
CacheControls {
default_on: false,
use_cache: true,
..enabled
}
.reads()
);
assert!(
!CacheControls {
no_cache: true,
..enabled
}
.reads()
);
assert!(
!CacheControls {
no_store: true,
..enabled
}
.writes()
);
}

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

@ -1,7 +1,27 @@
litellm-core is the LiteLLM SDK in Rust — it makes the LLM call. Each top-level call is a module under `src/<route>/` exposing a public entrypoint named after the route (`messages::messages()`, the Rust equivalent of `litellm.messages()`): you call it and get a typed non-streaming response back.
A route module owns everything the call needs: types, the provider template trait, provider transforms (under `providers/`), provider/auth/URL resolution, and the handler that performs the HTTP call. Handlers belong here, never in a host crate.
A route module owns the call entrypoint, runtime types, provider/auth/URL resolution, and the handler that performs the HTTP call. Provider code and base config traits live under `src/llms/`, mirroring their Python source paths. This applies to every API surface: shared orchestration stays in its route module (`ocr/`, `chat_completions/`, `messages/`, `audio_transcription/`, or `responses/`), while provider transformations live under the corresponding Python-mirrored `llms/<provider>/` path. Import implementations directly from their canonical paths; do not add a `src/providers/` layer or compatibility re-exports. Shared provider resolution lives under `src/litellm_core_utils/get_llm_provider_logic.rs`. Handlers belong in core, never in a host crate
Not here: serving HTTP (axum routes, extractors), config file reading, rollout state, databases, or host-specific callback execution. Core owns lifecycle sequencing and callback payload construction; hosts execute the selected integrations. Env reads are limited to credential fallback in a route's `prepare.rs`.
Routes (messages, ocr, realtime) and providers (anthropic, mistral, openai) are modules, not crates.
## Python/Rust transformation pairs
Use the base OCR and Mistral OCR pairs as the reference when aligning transformations. Derive `src/<relative_path>.rs` from `litellm/<relative_path>.py`, preserving meaningful basenames such as `messages_transformation`
Keep corresponding operation names and parameter names when their responsibilities match. Rust types retain the Python semantic name with Rust acronym casing (`BaseOCRConfig` / `BaseOcrConfig`, `MistralOCRConfig` / `MistralOcrConfig`). Private Python helpers can drop their leading underscore. Give Rust adapter helpers distinct responsibility names rather than duplicating trait method names
Order OCR config methods as supported parameters, credential metadata and connection resolution, health-check input, parameter mapping, environment validation, URL construction, request transformation, async request transformation, response transformation, async response transformation, and error conversion. Put constants and data types before the config, private helpers after it in operation order, and tests last. Rust-only trait hooks follow the corresponding Python methods
Use trait defaults for unchanged inherited behavior and explicit delegation for shared provider behavior. Keep typed inputs, ownership, `Result`, and async I/O idiomatic. A matching path or symbol identifies the counterpart, not a claim of full behavioral parity
Use named `#[rstest]` cases for independent input/output scenarios instead of loops or repeated calls in one test. Inject reusable setup with `#[fixture]` arguments and use `#[with(...)]` for fixture overrides. Keep assertions about the same result together
For base OCR, Python response models correspond to `src/ocr/types.rs`; Rust context/environment types support the runtime. `BaseOcrConfig::prepare_request` corresponds to Python's HTTP-handler preparation rather than a `BaseOCRConfig` method, and `validate_request_body` is a Rust-only hook
For Mistral, `async_transform_ocr_request` uses the base default in both languages. `resolve_headers` and `build_ocr_url` implement the respective environment and URL operations, and `normalize_response` implements the typed part of response transformation. Existing auth key/header handling and top-level response-extra preservation differ between languages; layout refactors must preserve those behaviors and verify them with the existing tests
For non-OCR pairs, order corresponding methods as parameter support/mapping, environment validation, URL construction, request transformation, and response transformation, followed by Rust-only runtime hooks. Auth resolution remains split between configs and route preparation. Chat `supported_openai_param_mappings` describes accepted OpenAI/provider name pairs, unlike Python's `get_supported_openai_params` name list. Audio `map_transcription_params` remains a Rust filtering helper
Azure Messages maps to `llms/azure_ai/anthropic/messages_transformation.py`; Bedrock Converse maps to `llms/bedrock/chat/converse_transformation.py`. `AnthropicConfig`, `AmazonConverseConfig`, and the non-OCR base traits are partial ports. `OpenAiResponsesApiConfig` currently implements only the WebSocket surface. Preserve their acceptance gates, passthrough behavior, and host fallback contracts when aligning layout

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

@ -10,10 +10,12 @@ autotests = false
bytes.workspace = true
futures-util.workspace = true
base64.workspace = true
azure_core.workspace = true
azure_identity.workspace = true
data-url = "0.3.2"
gcp_auth.workspace = true
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
@ -21,33 +23,20 @@ reqwest.workspace = true
rustls.workspace = true
rustls-native-certs.workspace = true
serde.workspace = true
serde_json.workspace = true
serde_json = { workspace = true, features = ["preserve_order"] }
serde_with.workspace = true
serde_path_to_error = "0.1"
strum.workspace = true
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
aws-config = { version = "1.9.0", default-features = false, features = ["rustls", "rt-tokio"], optional = true }
aws-credential-types = { version = "1.3.0", features = ["hardcoded-credentials"], optional = true }
aws-sdk-sts = { version = "1.108.0", default-features = false, features = ["rustls", "rt-tokio"], optional = true }
aws-sigv4 = { version = "1.5.1", optional = true }
aws-types = { version = "1.4.0", optional = true }
aws-smithy-runtime-api = { version = "1.13.0", optional = true }
[features]
default = []
bedrock-auth = [
"dep:aws-config",
"dep:aws-credential-types",
"dep:aws-sdk-sts",
"dep:aws-sigv4",
"dep:aws-types",
"dep:aws-smithy-runtime-api",
]
[dev-dependencies]
aws-smithy-eventstream = "=0.61.1"
aws-smithy-types = "1.6.1"
rstest.workspace = true

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

@ -0,0 +1,26 @@
#[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)]
pub enum Error {
#[error("expected {expected}, got {actual}")]
InvalidType {
expected: &'static str,
actual: &'static str,
},
#[error("missing required field: {0}")]
MissingField(&'static str),
#[error("invalid provider: {0}")]
InvalidProvider(String),
#[error("invalid request: {0}")]
InvalidRequest(String),
#[error("invalid response: {0}")]
InvalidResponse(String),
#[error("unsupported by the rust path: {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(transparent)]
Aws(#[from] litellm_auth_aws::Error),
}

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

@ -1,10 +1,9 @@
use serde_json::Value;
use crate::error::Error;
use crate::http_utils::{http_request, truncate_error_body};
use super::Error;
use super::client::http_client;
use super::types::ProviderAudioTranscriptionRequest;
use crate::http_utils::{http_request, truncate_error_body};
pub async fn execute_audio_transcription_provider_call(
request: ProviderAudioTranscriptionRequest,
@ -21,27 +20,26 @@ pub async fn execute_audio_transcription_provider_call(
}
let response = http_request(request_builder)
.await
.map_err(|error| Error::Network(error.to_string()))?;
.map_err(|error| Error::Transport(crate::transport::Error::Network(error.to_string())))?;
let status = response.status();
let text = response
.text()
.await
.map_err(|error| Error::Network(error.to_string()))?;
.map_err(|error| Error::Transport(crate::transport::Error::Network(error.to_string())))?;
if !status.is_success() {
return Err(Error::Http {
return Err(Error::Transport(crate::transport::Error::Http {
status: status.as_u16(),
body: truncate_error_body(&text),
});
}));
}
let response_json = serde_json::from_str(&text)
.map_err(|error| Error::InvalidResponse(format!("invalid audio response JSON: {error}")))?;
Ok(request
.config
.transform_transcription_response(&request.model, response_json)?
.transform_audio_transcription_response(&request.model, response_json)?
.into_json())
}
#[cfg(feature = "bedrock-auth")]
async fn signed_headers(
request: &ProviderAudioTranscriptionRequest,
body: &[u8],
@ -49,9 +47,8 @@ async fn signed_headers(
use std::collections::BTreeMap;
use std::time::SystemTime;
use crate::audio_transcription::transformation::AudioTranscriptionAuth;
use crate::providers::bedrock::audio_transcription::aws_auth_config;
use crate::providers::bedrock::aws_base::{resolve_credentials, sign_bedrock_post};
use crate::llms::base_llm::audio_transcription::transformation::AudioTranscriptionAuth;
use litellm_auth_aws::{aws_auth_config, resolve_credentials, sign_bedrock_post};
let AudioTranscriptionAuth::AwsSigV4 { region, .. } = &request.auth else {
return Ok(request.upstream_headers.clone());
@ -73,18 +70,3 @@ async fn signed_headers(
)?;
Ok(unsigned.into_iter().chain(signature).collect())
}
#[cfg(not(feature = "bedrock-auth"))]
async fn signed_headers(
request: &ProviderAudioTranscriptionRequest,
_body: &[u8],
) -> Result<Vec<(String, String)>, Error> {
use crate::audio_transcription::transformation::AudioTranscriptionAuth;
match request.auth {
AudioTranscriptionAuth::AwsSigV4 { .. } => Err(Error::Unsupported(
"AWS SigV4 requires the bedrock-auth feature",
)),
AudioTranscriptionAuth::Bearer => Ok(request.upstream_headers.clone()),
}
}

7
litellm-rust/crates/core/src/audio_transcription/mod.rs
View file

@ -1,14 +1,13 @@
use crate::Error;
mod error;
pub use error::Error;
mod client;
mod handler;
mod prepare;
pub mod transformation;
pub mod types;
use serde_json::Value;
pub use handler::execute_audio_transcription_provider_call;
pub use prepare::prepare_audio_transcription_provider_call;
use serde_json::Value;
pub use types::{AudioTranscriptionRequest, ProviderAudioTranscriptionRequest};
pub async fn audio_transcription(request: AudioTranscriptionRequest<'_>) -> Result<Value, Error> {

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

@ -1,14 +1,15 @@
use crate::error::Error;
use crate::http_utils::{has_header, string_headers};
#[cfg(feature = "bedrock-auth")]
use crate::providers::bedrock::audio_transcription::BEDROCK_AUDIO_TRANSCRIPTION_CONFIG;
use crate::providers::custom_llm_provider::{CustomLlmProvider, get_custom_llm_provider};
use super::transformation::{AudioTranscriptionAuth, AudioTranscriptionProviderConfig};
use super::Error;
use super::types::{AudioTranscriptionRequest, ProviderAudioTranscriptionRequest};
use crate::http_utils::{has_header, string_headers};
use crate::litellm_core_utils::get_llm_provider_logic::{
CustomLlmProvider, get_custom_llm_provider,
};
use crate::llms::base_llm::audio_transcription::transformation::{
AudioTranscriptionAuth, BaseAudioTranscriptionConfig,
};
use crate::llms::bedrock::audio_transcription::BEDROCK_AUDIO_TRANSCRIPTION_CONFIG;
fn provider_config(provider: &str) -> Option<&'static dyn AudioTranscriptionProviderConfig> {
#[cfg(feature = "bedrock-auth")]
fn provider_config(provider: &str) -> Option<&'static dyn BaseAudioTranscriptionConfig> {
if provider == "bedrock" {
return Some(&BEDROCK_AUDIO_TRANSCRIPTION_CONFIG);
}
@ -48,7 +49,7 @@ pub fn prepare_audio_transcription_provider_call(
if !has_header(&headers, "content-type") {
headers.push(("Content-Type".to_string(), "application/json".to_string()));
}
let url = config.complete_url(
let url = config.get_complete_url(
request.api_base,
&model,
&request.optional_params,
@ -56,7 +57,7 @@ pub fn prepare_audio_transcription_provider_call(
)?;
let filtered_params = config.map_transcription_params(&request.optional_params);
let transformed =
config.transform_transcription_request(&model, request.audio, filtered_params)?;
config.transform_audio_transcription_request(&model, request.audio, filtered_params)?;
Ok(ProviderAudioTranscriptionRequest {
model,
custom_llm_provider: provider_info.custom_llm_provider.to_string(),
@ -65,7 +66,6 @@ pub fn prepare_audio_transcription_provider_call(
body: transformed.body,
upstream_headers: headers,
auth,
#[cfg(feature = "bedrock-auth")]
optional_params: request.optional_params,
timeout: request.timeout,
})

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

@ -3,7 +3,9 @@ use std::time::Duration;
use serde::{Deserialize, Serialize};
use serde_json::{Map, Value};
use super::transformation::{AudioTranscriptionAuth, AudioTranscriptionProviderConfig};
use crate::llms::base_llm::audio_transcription::transformation::{
AudioTranscriptionAuth, BaseAudioTranscriptionConfig,
};
pub struct AudioTranscriptionRequest<'a> {
pub model: &'a str,
@ -20,12 +22,11 @@ pub struct AudioTranscriptionRequest<'a> {
pub struct ProviderAudioTranscriptionRequest {
pub(super) model: String,
pub(super) custom_llm_provider: String,
pub(super) config: &'static dyn AudioTranscriptionProviderConfig,
pub(super) config: &'static dyn BaseAudioTranscriptionConfig,
pub(super) url: String,
pub(super) body: Value,
pub(super) upstream_headers: Vec<(String, String)>,
pub(super) auth: AudioTranscriptionAuth,
#[cfg(feature = "bedrock-auth")]
pub(super) optional_params: Map<String, Value>,
pub(super) timeout: Option<Duration>,
}

128
litellm-rust/crates/core/src/auth/error.rs
View file

@ -1,128 +0,0 @@
use thiserror::Error;
#[derive(Clone, Debug, Error, PartialEq, Eq)]
pub enum AuthError {
#[error("invalid authentication configuration: {0}")]
Configuration(#[from] AuthConfigurationError),
#[error("credential acquisition failed: {0}")]
AzureTokenAcquisition(String),
#[error("credential acquisition failed: Vertex AI credentials: {0}")]
VertexTokenAcquisition(String),
#[error("credential acquisition failed: {}", .0.iter().map(ToString::to_string).collect::<Vec<_>>().join("; "))]
CredentialChain(Vec<AuthError>),
#[error("credential caller failed: credential caller returned an empty credential")]
EmptyCallerCredential,
#[error("credential caller failed: Azure AD token provider returned an empty token")]
EmptyAzureToken,
#[error("credential acquisition failed: Azure OIDC reference did not resolve to a value")]
UnresolvedOidcReference,
#[error(
"Missing {provider} API Key - A call is being made to {provider} but no key is set either in the environment variables or via params"
)]
MissingApiKey { provider: &'static str },
#[error(
"Missing {provider} API Base - Set {environment_variable} environment variable or pass api_base parameter"
)]
MissingApiBase {
provider: &'static str,
environment_variable: &'static str,
},
#[error("{0}")]
MissingCredential(#[from] MissingCredential),
#[error("{0}")]
Aws(#[from] AwsAuthError),
#[error("invalid authentication header")]
InvalidHeader,
}
#[derive(Clone, Debug, Error, PartialEq, Eq)]
pub enum AuthConfigurationError {
#[error("credential header already exists")]
ExistingCredentialHeader,
#[error("credential plan is not allowed by the provider auth policy")]
DisallowedCredentialPlan,
#[error("credential cannot be empty")]
EmptyCredential,
#[error("invalid Azure credential selector")]
InvalidAzureSelector,
#[error("ClientSecretCredential requires tenant_id, client_id, and client_secret")]
MissingClientSecretFields,
#[error("WorkloadIdentityCredential requires tenant_id")]
MissingWorkloadTenant,
#[error("WorkloadIdentityCredential requires client_id")]
MissingWorkloadClient,
#[error("WorkloadIdentityCredential requires azure_federated_token_file")]
MissingWorkloadTokenFile,
#[error("credential reference requires a host credential resolver")]
MissingHostResolver,
#[error("caller credential plan requires provider-specific inputs")]
MissingCallerInputs,
#[error("credential header {0} already exists")]
DuplicateHeader(&'static str),
#[error("{0} must be a string or null")]
InvalidFieldType(String),
#[error("unsupported OIDC reference")]
UnsupportedOidcReference,
#[error("{0} cannot be empty")]
EmptyReference(String),
#[error("Azure credential initialization failed: {0}")]
AzureCredentialInitialization(String),
#[error("Azure authority must be an HTTPS origin without credentials, query, or fragment")]
InvalidAzureAuthority,
#[error("request-controlled Azure auth inputs cannot be combined with host credentials")]
MixedAzureCredentialSources,
#[error("request-controlled Azure credential references are not allowed")]
RequestAzureCredentialReference,
#[error("host credentials cannot be sent to a request-controlled Azure endpoint")]
RequestAzureCredentialDestination,
#[error("credentials cannot be sent to a request-controlled Vertex AI endpoint")]
RequestVertexCredentialDestination,
#[error(
"request-controlled Vertex credentials must use the canonical Google OAuth token endpoint"
)]
RequestVertexTokenEndpoint,
}
#[derive(Clone, Debug, Error, PartialEq, Eq)]
pub enum MissingCredential {
#[error(
"Missing Anthropic API Key - Set `api_key` or the ANTHROPIC_API_KEY environment variable"
)]
AnthropicApiKey,
#[error("Missing Azure API Key - Set `api_key` or the AZURE_API_KEY environment variable")]
AzureApiKey,
#[error(
"Missing Azure API Base - Set `api_base` or the AZURE_API_BASE environment variable. Expected format: https://<resource-name>.services.ai.azure.com/anthropic"
)]
AzureApiBase,
#[error(
"Missing OpenAI API Key - a realtime call is being made but no key was passed via params or the OPENAI_API_KEY environment variable"
)]
OpenAiRealtimeApiKey,
#[error(
"Missing OpenAI API Key - a Responses WebSocket call is being made but no key was passed via params or the OPENAI_API_KEY environment variable"
)]
OpenAiResponsesApiKey,
}
#[derive(Clone, Debug, Error, PartialEq, Eq)]
pub enum AwsAuthError {
#[error("AWS profile credentials failed: {0}")]
Profile(String),
#[error("AWS default credentials failed: {0}")]
DefaultChain(String),
#[error("AWS role credentials failed: {0}")]
AssumeRole(String),
#[error("AWS web identity credentials failed: {0}")]
WebIdentity(String),
#[error("AWS web identity expiration was invalid: {0}")]
WebIdentityExpiration(String),
#[error("AWS signing parameters failed: {0}")]
SigningParameters(String),
#[error("AWS signable request failed: {0}")]
SignableRequest(String),
#[error("AWS request signing failed: {0}")]
Signing(String),
#[error("AWS web identity response had no credentials")]
MissingWebIdentityCredentials,
}

258
litellm-rust/crates/core/src/caching/in_memory_cache.rs
View file

@ -1,258 +0,0 @@
use std::cmp::Reverse;
use std::collections::{BinaryHeap, HashMap};
use std::time::{Duration, SystemTime, UNIX_EPOCH};
const DEFAULT_MAX_SIZE_IN_MEMORY: usize = 200;
const DEFAULT_TTL: Duration = Duration::from_secs(600);
pub struct InMemoryCache<V: Clone> {
pub cache_dict: HashMap<String, V>,
pub ttl_dict: HashMap<String, Duration>,
pub expiration_heap: BinaryHeap<Reverse<(Duration, String)>>,
pub max_size_in_memory: usize,
pub default_ttl: Duration,
now: Box<dyn Fn() -> Duration + Send + Sync>,
}
impl<V: Clone> Default for InMemoryCache<V> {
fn default() -> Self {
Self::new(None, None)
}
}
impl<V: Clone> InMemoryCache<V> {
pub fn new(max_size_in_memory: Option<usize>, default_ttl: Option<Duration>) -> Self {
Self::with_clock(max_size_in_memory, default_ttl, || {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap_or_default()
})
}
pub fn with_clock(
max_size_in_memory: Option<usize>,
default_ttl: Option<Duration>,
now: impl Fn() -> Duration + Send + Sync + 'static,
) -> Self {
Self {
cache_dict: HashMap::new(),
ttl_dict: HashMap::new(),
expiration_heap: BinaryHeap::new(),
max_size_in_memory: max_size_in_memory.unwrap_or(DEFAULT_MAX_SIZE_IN_MEMORY),
default_ttl: default_ttl.unwrap_or(DEFAULT_TTL),
now: Box::new(now),
}
}
pub fn evict_cache(&mut self) {
if self.max_size_in_memory == 0 {
return;
}
let current_time = (self.now)();
while let Some(Reverse((expiration_time, key))) = self.expiration_heap.peek().cloned() {
if self.ttl_dict.get(&key).copied() != Some(expiration_time) {
self.expiration_heap.pop();
} else if expiration_time <= current_time {
self.expiration_heap.pop();
self.remove_key(&key);
} else {
break;
}
}
while self.cache_dict.len() >= self.max_size_in_memory {
let Some(Reverse((expiration_time, key))) = self.expiration_heap.pop() else {
break;
};
if self.ttl_dict.get(&key).copied() == Some(expiration_time) {
self.remove_key(&key);
}
}
}
pub fn allow_ttl_override(&self, key: &str) -> bool {
match self.ttl_dict.get(key).copied() {
None => true,
Some(expiration_time) => expiration_time < (self.now)(),
}
}
pub fn set_cache(&mut self, key: impl Into<String>, value: V, ttl: Option<Duration>) {
if self.max_size_in_memory == 0 {
return;
}
self.evict_cache();
let key = key.into();
self.cache_dict.insert(key.clone(), value);
if self.allow_ttl_override(&key) {
let expiration_time = (self.now)() + ttl.unwrap_or(self.default_ttl);
self.ttl_dict.insert(key.clone(), expiration_time);
self.expiration_heap.push(Reverse((expiration_time, key)));
}
}
// Generic values intentionally omit Python's per-item size check.
pub fn get_cache(&mut self, key: &str) -> Option<V> {
if self.cache_dict.contains_key(key) {
if self.is_key_expired(key) {
self.remove_key(key);
return None;
}
return self.cache_dict.get(key).cloned();
}
None
}
pub fn get_ttl(&self, key: &str) -> Option<Duration> {
self.ttl_dict.get(key).copied()
}
pub fn delete_cache(&mut self, key: &str) {
self.remove_key(key);
}
pub fn flush_cache(&mut self) {
self.cache_dict.clear();
self.ttl_dict.clear();
self.expiration_heap.clear();
}
fn is_key_expired(&self, key: &str) -> bool {
self.ttl_dict
.get(key)
.is_some_and(|expiration_time| *expiration_time < (self.now)())
}
fn remove_key(&mut self, key: &str) {
self.cache_dict.remove(key);
self.ttl_dict.remove(key);
}
}
#[cfg(test)]
mod tests {
use std::sync::{
Arc,
atomic::{AtomicU64, Ordering},
};
use super::InMemoryCache;
use std::time::Duration;
fn cache(now: Arc<AtomicU64>, max_size: usize, default_ttl: Duration) -> InMemoryCache<String> {
InMemoryCache::with_clock(Some(max_size), Some(default_ttl), move || {
Duration::from_secs(now.load(Ordering::Relaxed))
})
}
#[test]
fn ttl_expiry_is_deterministic() {
let now = Arc::new(AtomicU64::new(100));
let mut cache = cache(now.clone(), 10, Duration::from_secs(60));
cache.set_cache("key", "value".to_string(), None);
assert_eq!(cache.get_cache("key"), Some("value".to_string()));
now.store(161, Ordering::Relaxed);
assert_eq!(cache.get_cache("key"), None);
assert_eq!(cache.get_ttl("key"), None);
}
#[test]
fn default_and_per_set_ttl_are_applied() {
let now = Arc::new(AtomicU64::new(100));
let mut cache = cache(now.clone(), 10, Duration::from_secs(60));
cache.set_cache("default", "value".to_string(), None);
cache.set_cache("custom", "value".to_string(), Some(Duration::from_secs(20)));
assert_eq!(cache.get_ttl("default"), Some(Duration::from_secs(160)));
assert_eq!(cache.get_ttl("custom"), Some(Duration::from_secs(120)));
}
#[test]
fn unexpired_entries_do_not_allow_ttl_override() {
let now = Arc::new(AtomicU64::new(100));
let mut cache = cache(now.clone(), 10, Duration::from_secs(60));
cache.set_cache("key", "first".to_string(), Some(Duration::from_secs(20)));
cache.set_cache("key", "second".to_string(), Some(Duration::from_secs(80)));
assert_eq!(cache.get_cache("key"), Some("second".to_string()));
assert_eq!(cache.get_ttl("key"), Some(Duration::from_secs(120)));
now.store(121, Ordering::Relaxed);
cache.set_cache("key", "third".to_string(), Some(Duration::from_secs(80)));
assert_eq!(cache.get_ttl("key"), Some(Duration::from_secs(201)));
}
#[test]
fn max_size_evicts_earliest_expiration() {
let now = Arc::new(AtomicU64::new(100));
let mut cache = cache(now, 2, Duration::from_secs(60));
cache.set_cache("early", "value".to_string(), Some(Duration::from_secs(10)));
cache.set_cache("late", "value".to_string(), Some(Duration::from_secs(20)));
cache.set_cache("new", "value".to_string(), Some(Duration::from_secs(30)));
assert_eq!(cache.get_cache("early"), None);
assert!(cache.get_cache("late").is_some());
assert!(cache.get_cache("new").is_some());
}
#[test]
fn expired_entries_are_evicted_before_live_entries() {
let now = Arc::new(AtomicU64::new(100));
let mut cache = cache(now.clone(), 3, Duration::from_secs(60));
cache.set_cache(
"expired-one",
"value".to_string(),
Some(Duration::from_secs(10)),
);
cache.set_cache(
"expired-two",
"value".to_string(),
Some(Duration::from_secs(20)),
);
cache.set_cache("live", "value".to_string(), Some(Duration::from_secs(100)));
now.store(121, Ordering::Relaxed);
cache.set_cache("new", "value".to_string(), Some(Duration::from_secs(100)));
assert_eq!(cache.get_cache("expired-one"), None);
assert_eq!(cache.get_cache("expired-two"), None);
assert!(cache.get_cache("live").is_some());
assert!(cache.get_cache("new").is_some());
}
#[test]
fn stale_heap_entries_are_skipped() {
let now = Arc::new(AtomicU64::new(100));
let mut cache = cache(now, 1, Duration::from_secs(60));
cache.set_cache(
"removed",
"value".to_string(),
Some(Duration::from_secs(10)),
);
cache.delete_cache("removed");
cache.set_cache("kept", "value".to_string(), Some(Duration::from_secs(20)));
cache.set_cache("new", "value".to_string(), Some(Duration::from_secs(30)));
assert_eq!(cache.get_cache("removed"), None);
assert_eq!(cache.get_cache("kept"), None);
assert!(cache.get_cache("new").is_some());
}
#[test]
fn delete_and_flush_remove_values_and_ttls() {
let now = Arc::new(AtomicU64::new(100));
let mut cache = cache(now, 10, Duration::from_secs(60));
cache.set_cache("one", "value".to_string(), None);
cache.set_cache("two", "value".to_string(), None);
cache.delete_cache("one");
assert_eq!(cache.get_cache("one"), None);
cache.flush_cache();
assert!(cache.cache_dict.is_empty());
assert!(cache.ttl_dict.is_empty());
assert!(cache.expiration_heap.is_empty());
}
#[test]
fn zero_max_size_does_not_cache() {
let now = Arc::new(AtomicU64::new(100));
let mut cache = cache(now, 0, Duration::from_secs(60));
cache.set_cache("key", "value".to_string(), None);
assert_eq!(cache.get_cache("key"), None);
assert!(cache.cache_dict.is_empty());
}
}

1
litellm-rust/crates/core/src/caching/mod.rs
View file

@ -1 +0,0 @@
pub mod in_memory_cache;

468
litellm-rust/crates/core/src/call_arguments.rs Normal file
View file

@ -0,0 +1,468 @@
use std::ops::Deref;
use serde::{Deserialize, Serialize, de::DeserializeOwned};
use serde_json::{Map, Value};
#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
#[serde(transparent)]
pub struct CallArguments(Map<String, Value>);
impl CallArguments {
pub(crate) fn select(&self, names: &[&str]) -> Map<String, Value> {
self.iter()
.filter(|(name, _)| names.contains(&name.as_str()))
.map(|(name, value)| (name.clone(), value.clone()))
.collect()
}
}
#[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)]
#[error("invalid argument: {path}")]
pub struct ArgumentError {
pub path: String,
}
pub fn parse_options<T: DeserializeOwned>(arguments: &CallArguments) -> Result<T, ArgumentError> {
let deserializer = serde::de::value::MapDeserializer::new(
arguments.iter().map(|(name, value)| (name.as_str(), value)),
);
serde_path_to_error::deserialize(deserializer).map_err(|error| ArgumentError {
path: error.path().to_string(),
})
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct ArgumentSpec {
pub name: &'static str,
pub secret: bool,
}
pub fn should_project(name: &str, consumed: &[ArgumentSpec], bound_fields: &[&str]) -> bool {
consumed.iter().any(|field| field.name == name)
|| (!bound_fields.contains(&name) && !is_control(name))
}
pub fn is_control(name: &str) -> bool {
crate::params::is_control_param(name) || HOST_CONTROLS.contains(&name)
}
const HOST_CONTROLS: &[&str] = &[
"_agentic_loop_api_surface",
"_agentic_loop_depth",
"_agentic_loop_fingerprints",
"_code_interpreter_interception_active",
"_code_interpreter_interception_converted_stream",
"_code_interpreter_interception_sandbox_key",
"_code_interpreter_interception_session_scoped",
"_headroom_interception_converted_stream",
"_litellm_strip_stream_usage",
"_router_weights",
"_websearch_interception_converted_stream",
"_websearch_interception_emit_native_blocks",
"acompletion",
"adaptive_router_config",
"adaptive_router_default_model",
"aembedding",
"aimg_generation",
"allm_passthrough_route",
"allow_client_keepalive_override",
"allowed_model_region",
"allowed_openai_params",
"annotation_cost_per_page",
"api_version",
"arize_api_key",
"arize_space_id",
"arize_space_key",
"assistant_continue_message",
"async_call",
"atext_completion",
"attempted_targets",
"auto_router_config",
"auto_router_config_path",
"auto_router_default_model",
"auto_router_embedding_model",
"auto_router_max_input_chars",
"auto_router_model_compression",
"auto_router_routing_compression",
"aws_batch_role_arn",
"azure",
"azure_password",
"azure_username",
"base_model",
"bedrock_tags",
"bos_token",
"budget_duration",
"cache",
"cache_creation_input_audio_token_cost",
"cache_creation_input_token_cost",
"cache_creation_input_token_cost_above_1hr",
"cache_creation_input_token_cost_above_200k_tokens",
"cache_creation_input_token_cost_above_272k_tokens",
"cache_creation_input_token_cost_above_272k_tokens_flex",
"cache_creation_input_token_cost_above_272k_tokens_priority",
"cache_creation_input_token_cost_flex",
"cache_creation_input_token_cost_priority",
"cache_creation_input_token_cost_ultrafast",
"cache_key",
"cache_read_input_audio_token_cost",
"cache_read_input_token_cost",
"cache_read_input_token_cost_above_200k_tokens",
"cache_read_input_token_cost_above_200k_tokens_priority",
"cache_read_input_token_cost_above_272k_tokens",
"cache_read_input_token_cost_above_272k_tokens_flex",
"cache_read_input_token_cost_above_272k_tokens_priority",
"cache_read_input_token_cost_above_512k_tokens",
"cache_read_input_token_cost_flex",
"cache_read_input_token_cost_priority",
"cache_read_input_token_cost_ultrafast",
"caching",
"caching_groups",
"citation_cost_per_token",
"client",
"client_side_timeout",
"complete_response",
"completion_call_id",
"complexity_router_config",
"complexity_router_default_model",
"configurable_clientside_auth_params",
"context_window_fallback_dict",
"cooldown_time",
"cost_per_query",
"custom_prompt_dict",
"data_residency",
"dd_agent_host",
"dd_agent_port",
"dd_api_key",
"dd_site",
"default_api_key_rpm_limit",
"default_api_key_tpm_limit",
"disable_add_transform_inline_image_block",
"enable_json_schema_validation",
"enable_prompt_caching",
"enable_tag_filtering",
"ensure_alternating_roles",
"eos_token",
"fallback_depth",
"fallbacks",
"fastest_response",
"final_prompt_value",
"force_timeout",
"gcs_bucket_name",
"gcs_path_service_account",
"google_maps_grounding_cost_per_query",
"headers",
"hf_model_name",
"humanloop_api_key",
"id",
"input_cost_per_audio_per_second",
"input_cost_per_audio_per_second_above_128k_tokens",
"input_cost_per_audio_token",
"input_cost_per_audio_token_batches",
"input_cost_per_character",
"input_cost_per_character_above_128k_tokens",
"input_cost_per_image",
"input_cost_per_image_above_128k_tokens",
"input_cost_per_image_token",
"input_cost_per_image_token_batches",
"input_cost_per_pixel",
"input_cost_per_query",
"input_cost_per_second",
"input_cost_per_token",
"input_cost_per_token_above_128k_tokens",
"input_cost_per_token_above_200k_tokens",
"input_cost_per_token_above_200k_tokens_priority",
"input_cost_per_token_above_272k_tokens",
"input_cost_per_token_above_272k_tokens_flex",
"input_cost_per_token_above_272k_tokens_priority",
"input_cost_per_token_above_512k_tokens",
"input_cost_per_token_batches",
"input_cost_per_token_cache_hit",
"input_cost_per_token_flex",
"input_cost_per_token_priority",
"input_cost_per_token_ultrafast",
"input_cost_per_video_per_second",
"input_cost_per_video_per_second_above_128k_tokens",
"input_cost_per_video_per_second_above_15s_interval",
"input_cost_per_video_per_second_above_8s_interval",
"input_cost_per_video_token",
"input_cost_per_video_token_batches",
"itpm",
"keepalive_seconds",
"langfuse_environment",
"langfuse_host",
"langfuse_prompt_version",
"langfuse_public_key",
"langfuse_secret",
"langfuse_secret_key",
"langsmith_api_key",
"langsmith_base_url",
"langsmith_project",
"langsmith_sampling_rate",
"langsmith_tenant_id",
"litellm_credential_name",
"litellm_disabled_callbacks",
"litellm_request_debug",
"litellm_session_id",
"litellm_system_prompt",
"litellm_trace_id",
"litellm_trusted_callback_vars",
"logger_fn",
"max_agentic_loops",
"max_budget",
"max_fallbacks",
"max_parallel_requests",
"merge_reasoning_content_in_choices",
"metadata",
"mock_response",
"mock_timeout",
"model_alias_map",
"model_config",
"model_file_id_mapping",
"model_info",
"model_list",
"newrelic_api_key",
"newrelic_region",
"no-log",
"num_retries",
"ocr_cost_per_credit",
"ocr_cost_per_page",
"order",
"otpm",
"output_cost_per_audio_per_second",
"output_cost_per_audio_token",
"output_cost_per_character",
"output_cost_per_character_above_128k_tokens",
"output_cost_per_image",
"output_cost_per_image_token",
"output_cost_per_pixel",
"output_cost_per_reasoning_token",
"output_cost_per_reasoning_token_flex",
"output_cost_per_reasoning_token_priority",
"output_cost_per_second",
"output_cost_per_second_1080p",
"output_cost_per_second_480p",
"output_cost_per_second_4k",
"output_cost_per_second_720p",
"output_cost_per_token",
"output_cost_per_token_above_128k_tokens",
"output_cost_per_token_above_200k_tokens",
"output_cost_per_token_above_200k_tokens_priority",
"output_cost_per_token_above_272k_tokens",
"output_cost_per_token_above_272k_tokens_flex",
"output_cost_per_token_above_272k_tokens_priority",
"output_cost_per_token_above_512k_tokens",
"output_cost_per_token_batches",
"output_cost_per_token_flex",
"output_cost_per_token_priority",
"output_cost_per_token_ultrafast",
"output_cost_per_video_per_second",
"output_cost_per_video_token",
"output_vector_size",
"posthog_api_key",
"posthog_api_url",
"preset_cache_key",
"prompt_environment",
"prompt_id",
"prompt_label",
"prompt_variables",
"prompt_version",
"provider_specific_header",
"quality_router_config",
"quality_router_default_model",
"region_name",
"regional_endpoint_uplift_multiplier",
"regional_processing_uplift_multiplier_eu",
"regional_processing_uplift_multiplier_us",
"retry_policy",
"retry_strategy",
"roles",
"routing_strategy",
"rpm",
"rust",
"s3_bucket_name",
"s3_output_bucket_name",
"s3_region_name",
"search_context_cost_per_query",
"search_tool_name",
"secret_fields",
"self",
"shared_session",
"ssl_verify",
"stream_response",
"stream_timeout",
"supports_system_message",
"tags",
"text_completion",
"tiered_pricing",
"tpm",
"ttl",
"turn_off_message_logging",
"use_chat_completions_api",
"use_client",
"use_in_pass_through",
"use_litellm_proxy",
"use_xai_oauth",
"user_continue_message",
"verbose",
"wandb_api_key",
"weave_project_id",
"weight",
];
pub fn compose_body<B: Serialize>(
arguments: &CallArguments,
body: &B,
consumed: &[&str],
) -> Result<Value, crate::params::Error> {
let Value::Object(fields) =
serde_json::to_value(body).map_err(|_| crate::params::Error::Body)?
else {
return Err(crate::params::Error::Body);
};
let overrides = match arguments.get("extra_body") {
None | Some(Value::Null) => None,
Some(Value::Object(fields)) => Some(fields),
Some(_) => return Err(crate::params::Error::ExtraBody),
};
let extensions = arguments.iter().filter(|(name, _)| {
!consumed.contains(&name.as_str()) && name.as_str() != "extra_body" && !is_control(name)
});
Ok(Value::Object(
fields
.into_iter()
.chain(
extensions
.chain(overrides.into_iter().flatten())
.filter(|(name, _)| {
name.as_str() != "model"
&& name.as_str() != "extra_body"
&& !crate::params::is_control_param(name)
})
.map(|(name, value)| (name.clone(), value.clone())),
)
.collect(),
))
}
impl Deref for CallArguments {
type Target = Map<String, Value>;
fn deref(&self) -> &Self::Target {
&self.0
}
}
impl From<Map<String, Value>> for CallArguments {
fn from(values: Map<String, Value>) -> Self {
Self(values)
}
}
impl From<CallArguments> for Map<String, Value> {
fn from(arguments: CallArguments) -> Self {
arguments.0
}
}
impl FromIterator<(String, Value)> for CallArguments {
fn from_iter<T: IntoIterator<Item = (String, Value)>>(iter: T) -> Self {
Self(iter.into_iter().collect())
}
}
impl IntoIterator for CallArguments {
type Item = (String, Value);
type IntoIter = serde_json::map::IntoIter;
fn into_iter(self) -> Self::IntoIter {
self.0.into_iter()
}
}
#[cfg(test)]
mod tests {
use serde_json::json;
use super::*;
#[test]
fn composition_preserves_extensions_and_applies_shallow_explicit_overrides() {
let original = json!({
"known": false, "future": {"old": 1}, "null": null, "zero": 0,
"metadata": {"host": true}, "shared_session": "host", "api_key": "secret",
"extra_body": {
"known": null, "future": {"new": [false, 0, null]},
"metadata": {"provider": true}, "model": "ignored", "api_key": "ignored"
}
});
let arguments = serde_json::from_value(original.clone()).unwrap();
let body = compose_body(
&arguments,
&json!({"model":"resolved", "known":false}),
&["known"],
)
.unwrap();
assert_eq!(
body,
json!({
"model":"resolved", "known":null, "future":{"new":[false,0,null]},
"null":null, "zero":0, "metadata":{"provider":true}
})
);
assert_eq!(serde_json::to_value(arguments).unwrap(), original);
}
#[test]
fn projection_prioritizes_consumed_fields_and_keeps_unknown_names() {
let fields = [ArgumentSpec {
name: "id",
secret: false,
}];
assert!(should_project("id", &fields, &[]));
assert!(!should_project("id", &[], &[]));
assert!(should_project("future_option", &[], &[]));
assert!(!should_project("document", &fields, &["document"]));
assert!(!should_project("metadata", &fields, &[]));
assert!(!should_project("callbacks", &fields, &[]));
assert!(!should_project("ocr_cost_per_page", &fields, &[]));
}
#[test]
fn invalid_extra_body_is_rejected_without_coercing_it_to_empty() {
for value in [json!(false), json!(0), json!([]), json!("")] {
let arguments = serde_json::from_value(json!({"extra_body":value})).unwrap();
assert_eq!(
compose_body(&arguments, &json!({}), &[]),
Err(crate::params::Error::ExtraBody)
);
}
let arguments = serde_json::from_value(json!({"extra_body":null})).unwrap();
assert_eq!(
compose_body(&arguments, &json!({}), &[]).unwrap(),
json!({})
);
}
#[test]
fn typed_views_preserve_missing_and_explicit_null_in_the_source() {
#[derive(Deserialize)]
struct Options {
enabled: Option<bool>,
}
let arguments: CallArguments =
serde_json::from_value(json!({"enabled":null,"future":0})).unwrap();
assert!(
parse_options::<Options>(&arguments)
.unwrap()
.enabled
.is_none()
);
assert_eq!(arguments.get("enabled"), Some(&Value::Null));
assert_eq!(arguments.get("missing"), None);
let invalid = serde_json::from_value(json!({"enabled":0})).unwrap();
assert_eq!(
parse_options::<Options>(&invalid).err().unwrap().path,
"enabled"
);
}
}

19
litellm-rust/crates/core/src/call_lifecycle/host.rs
View file

@ -6,10 +6,11 @@ pub enum HostCallStep<O, C> {
Complete(C),
}
pub type HostCallFuture<'a, O, C> =
Pin<Box<dyn Future<Output = Result<HostCallStep<O, C>, crate::Error>> + Send + 'a>>;
pub type HostCallFuture<'a, O, C, E> =
Pin<Box<dyn Future<Output = Result<HostCallStep<O, C>, E>> + Send + 'a>>;
pub trait HostCall: Send + Sync {
type Error: Send + Sync + 'static;
type Operation: Send + 'static;
type Result: Send + 'static;
type Complete: Send + 'static;
@ -17,12 +18,12 @@ pub trait HostCall: Send + Sync {
fn resume(
&mut self,
result: Option<Self::Result>,
) -> HostCallFuture<'_, Self::Operation, Self::Complete>;
) -> HostCallFuture<'_, Self::Operation, Self::Complete, Self::Error>;
fn interrupt(
&mut self,
failure: HostFailure,
) -> HostCallFuture<'_, Self::Operation, Self::Complete>;
failure: HostFailure<Self::Error>,
) -> HostCallFuture<'_, Self::Operation, Self::Complete, Self::Error>;
}
pub enum HostStep<V, S> {
@ -48,9 +49,9 @@ pub enum HostPhase {
}
#[derive(Clone, Debug)]
pub enum HostFailure {
Error(crate::Error),
Cancelled(crate::Error),
pub enum HostFailure<E> {
Error(E),
Cancelled(E),
}
pub struct HostLifecycle {
@ -70,7 +71,7 @@ impl HostLifecycle {
self.phase
}
pub fn accept(&mut self, result: Result<(), HostFailure>) -> Option<crate::Error> {
pub fn accept<E>(&mut self, result: Result<(), HostFailure<E>>) -> Option<E> {
if let Err(failure) = result {
if self.phase == HostPhase::DeploymentFailure {
self.phase = HostPhase::Failure;

47
litellm-rust/crates/core/src/call_lifecycle/mod.rs
View file

@ -1,8 +1,6 @@
use std::future::Future;
use std::time::{Instant, SystemTime, UNIX_EPOCH};
use crate::Error;
pub mod host;
#[cfg(test)]
#[path = "../../tests/host_lifecycle.rs"]
@ -15,14 +13,15 @@ pub use types::{
};
pub trait CallLifecycleHooks<InitialReq, ProviderReq, Resp>: Send + Sync {
type PreCallFuture<'a>: Future<Output = Result<InitialReq, Error>> + Send + 'a
type Error: Send + Sync;
type PreCallFuture<'a>: Future<Output = Result<InitialReq, Self::Error>> + Send + 'a
where
Self: 'a,
InitialReq: 'a,
ProviderReq: 'a,
Resp: 'a;
type DuringCallFuture<'a>: Future<Output = Result<ProviderReq, Error>> + Send + 'a
type DuringCallFuture<'a>: Future<Output = Result<ProviderReq, Self::Error>> + Send + 'a
where
Self: 'a,
InitialReq: 'a,
@ -60,7 +59,7 @@ pub trait CallLifecycleHooks<InitialReq, ProviderReq, Resp>: Send + Sync {
fn async_log_failure_event<'a>(
&'a self,
context: &'a CallLifecycleContext,
error: &'a Error,
error: &'a Self::Error,
timing: &'a CallLifecycleTiming,
) -> Self::FailureFuture<'a>;
}
@ -90,12 +89,12 @@ impl<'a> CallLifecycle<'a> {
request: InitialReq,
hooks: &Hooks,
provider_call: ProviderCall,
) -> Result<Resp, Error>
) -> Result<Resp, Hooks::Error>
where
InitialReq: CallLifecycleRequest,
Hooks: CallLifecycleHooks<InitialReq, ProviderReq, Resp>,
ProviderCall: FnOnce(ProviderReq) -> ProviderFuture,
ProviderFuture: Future<Output = Result<Resp, Error>>,
ProviderFuture: Future<Output = Result<Resp, Hooks::Error>>,
{
let context = request.lifecycle_context();
self.run(context, request, hooks, provider_call).await
@ -107,11 +106,11 @@ impl<'a> CallLifecycle<'a> {
request: InitialReq,
hooks: &Hooks,
provider_call: ProviderCall,
) -> Result<Resp, Error>
) -> Result<Resp, Hooks::Error>
where
Hooks: CallLifecycleHooks<InitialReq, ProviderReq, Resp>,
ProviderCall: FnOnce(ProviderReq) -> ProviderFuture,
ProviderFuture: Future<Output = Result<Resp, Error>>,
ProviderFuture: Future<Output = Result<Resp, Hooks::Error>>,
{
let call_start = epoch_seconds();
let mut phases = Vec::new();
@ -170,7 +169,7 @@ impl<'a> CallLifecycle<'a> {
&self,
context: &CallLifecycleContext,
hooks: &Hooks,
error: &Error,
error: &Hooks::Error,
call_start: f64,
phases: &mut Vec<CallLifecyclePhaseTiming>,
) where
@ -229,10 +228,11 @@ fn epoch_seconds() -> f64 {
#[cfg(test)]
mod tests {
use super::*;
use std::pin::Pin;
use std::sync::Mutex;
use super::*;
type BoxFuture<'a, T> = Pin<Box<dyn Future<Output = T> + Send + 'a>>;
#[derive(Default)]
@ -255,8 +255,9 @@ mod tests {
}
impl CallLifecycleHooks<String, String, String> for RecordingHooks {
type PreCallFuture<'a> = BoxFuture<'a, Result<String, Error>>;
type DuringCallFuture<'a> = BoxFuture<'a, Result<String, Error>>;
type Error = crate::messages::Error;
type PreCallFuture<'a> = BoxFuture<'a, Result<String, crate::messages::Error>>;
type DuringCallFuture<'a> = BoxFuture<'a, Result<String, crate::messages::Error>>;
type SuccessFuture<'a> = BoxFuture<'a, ()>;
type FailureFuture<'a> = BoxFuture<'a, ()>;
@ -298,7 +299,7 @@ mod tests {
fn async_log_failure_event<'a>(
&'a self,
_context: &'a CallLifecycleContext,
_error: &'a Error,
_error: &'a crate::messages::Error,
_timing: &'a CallLifecycleTiming,
) -> Self::FailureFuture<'a> {
Box::pin(async move {
@ -308,8 +309,9 @@ mod tests {
}
impl CallLifecycleHooks<RecordingRequest, String, String> for RecordingHooks {
type PreCallFuture<'a> = BoxFuture<'a, Result<RecordingRequest, Error>>;
type DuringCallFuture<'a> = BoxFuture<'a, Result<String, Error>>;
type Error = crate::messages::Error;
type PreCallFuture<'a> = BoxFuture<'a, Result<RecordingRequest, crate::messages::Error>>;
type DuringCallFuture<'a> = BoxFuture<'a, Result<String, crate::messages::Error>>;
type SuccessFuture<'a> = BoxFuture<'a, ()>;
type FailureFuture<'a> = BoxFuture<'a, ()>;
@ -349,7 +351,7 @@ mod tests {
fn async_log_failure_event<'a>(
&'a self,
_context: &'a CallLifecycleContext,
_error: &'a Error,
_error: &'a crate::messages::Error,
_timing: &'a CallLifecycleTiming,
) -> Self::FailureFuture<'a> {
Box::pin(async move {
@ -387,13 +389,20 @@ mod tests {
"request".to_string(),
&hooks,
|_request| async move {
Err::<String, Error>(Error::Network("provider down".to_string()))
Err::<String, crate::messages::Error>(crate::messages::Error::Transport(
crate::transport::Error::Network("provider down".to_string()),
))
},
)
.await
.expect_err("call fails");
assert_eq!(error, Error::Network("provider down".to_string()));
assert_eq!(
error,
crate::messages::Error::Transport(crate::transport::Error::Network(
"provider down".to_string()
))
);
assert_eq!(hooks.events(), vec!["pre_call", "during_call", "failure"]);
}

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

@ -1,20 +1,17 @@
use crate::Error;
use crate::http_utils::string_headers as shared_string_headers;
use crate::providers::anthropic::chat_completions::transformation::ANTHROPIC_CHAT_COMPLETIONS_CONFIG;
use serde_json::{Map, Value};
use super::transformation::ChatCompletionsProviderConfig;
use super::Error;
use crate::http_utils::string_headers as shared_string_headers;
use crate::llms::anthropic::chat::transformation::ANTHROPIC_CHAT_COMPLETIONS_CONFIG;
use crate::llms::base_llm::chat::transformation::BaseConfig;
const HEADER_CONTEXT: &str = "chat completions";
pub(super) fn chat_completions_provider_config(
provider: &str,
) -> Option<&'static dyn ChatCompletionsProviderConfig> {
pub(super) fn chat_completions_provider_config(provider: &str) -> Option<&'static dyn BaseConfig> {
match provider {
"anthropic" => Some(&ANTHROPIC_CHAT_COMPLETIONS_CONFIG),
#[cfg(feature = "bedrock-auth")]
"bedrock" => Some(
&crate::providers::bedrock::chat_completions::transformation::BEDROCK_CHAT_COMPLETIONS_CONFIG,
&crate::llms::bedrock::chat::converse_transformation::BEDROCK_CHAT_COMPLETIONS_CONFIG,
),
_ => None,
}
@ -23,5 +20,5 @@ pub(super) fn chat_completions_provider_config(
pub(super) fn string_headers(
extra_headers: Option<Map<String, Value>>,
) -> Result<Vec<(String, String)>, Error> {
shared_string_headers(HEADER_CONTEXT, extra_headers)
shared_string_headers(HEADER_CONTEXT, extra_headers).map_err(Error::from)
}

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

@ -10,9 +10,8 @@
//! `_bedrock_converse_messages_pt` for the text-only surface this route
//! accepts; anything richer is declined upstream by the capability gate.
use crate::constants::EMPTY_TEXT_PLACEHOLDER;
use super::types::{ChatMessage, ChatMessageContent};
use crate::constants::EMPTY_TEXT_PLACEHOLDER;
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum TurnRole {
@ -132,9 +131,10 @@ pub fn build_conversation(messages: &[ChatMessage]) -> Conversation {
#[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
use super::*;
fn messages(value: serde_json::Value) -> Vec<ChatMessage> {
serde_json::from_value(value).expect("valid messages")
}

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

@ -0,0 +1,26 @@
#[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)]
pub enum Error {
#[error("expected {expected}, got {actual}")]
InvalidType {
expected: &'static str,
actual: &'static str,
},
#[error("missing required field: {0}")]
MissingField(&'static str),
#[error("invalid provider: {0}")]
InvalidProvider(String),
#[error("invalid request: {0}")]
InvalidRequest(String),
#[error("invalid response: {0}")]
InvalidResponse(String),
#[error("unsupported by the rust path: {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(transparent)]
Aws(#[from] litellm_auth_aws::Error),
}

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

@ -1,15 +1,14 @@
use serde_json::Value;
use crate::error::Error;
use crate::http_utils::{http_request, truncate_error_body};
use super::Error;
use super::client::http_client;
use super::prepare::prepare_provider_request;
use super::transformation::ChatCompletionsAuth;
use super::types::{
ChatCompletionsResponse, ProviderChatCompletionsRequest, ProviderChatResponseData,
ResolvedChatCompletionsRequest,
};
use crate::http_utils::{http_request, truncate_error_body};
use crate::llms::base_llm::chat::transformation::ChatCompletionsAuth;
pub(super) async fn execute_chat_completions_provider_call(
request: ResolvedChatCompletionsRequest<'_>,
@ -35,9 +34,9 @@ pub(super) async fn execute_chat_completions_provider_call(
// so the host can still serve it. Everything else here, a timeout
// above all, may have reached the provider and been answered.
if err.is_connect() || err.is_builder() {
Error::Connect(err.to_string())
Error::Transport(crate::transport::Error::Connect(err.to_string()))
} else {
Error::Network(err.to_string())
Error::Transport(crate::transport::Error::Network(err.to_string()))
}
})?;
@ -45,13 +44,13 @@ pub(super) async fn execute_chat_completions_provider_call(
let text = response
.text()
.await
.map_err(|err| Error::Network(err.to_string()))?;
.map_err(|err| Error::Transport(crate::transport::Error::Network(err.to_string())))?;
if !status.is_success() {
return Err(Error::Http {
return Err(Error::Transport(crate::transport::Error::Http {
status: status.as_u16(),
body: truncate_error_body(&text),
});
}));
}
let body: Value = serde_json::from_str(&text).map_err(|err| {
@ -74,12 +73,12 @@ pub(super) async fn execute_chat_completions_provider_call(
/// can only mean the provider was already called.
pub(super) fn as_response_error(err: Error) -> Error {
match err {
already @ (Error::InvalidResponse(_) | Error::Http { .. }) => already,
already @ (Error::InvalidResponse(_)
| Error::Transport(crate::transport::Error::Http { .. })) => already,
other => Error::InvalidResponse(other.to_string()),
}
}
#[cfg(feature = "bedrock-auth")]
pub(super) async fn signed_headers(
request: &ProviderChatCompletionsRequest,
body: &[u8],
@ -87,7 +86,7 @@ pub(super) async fn signed_headers(
use std::collections::BTreeMap;
use std::time::SystemTime;
use crate::providers::bedrock::aws_base::{
use litellm_auth_aws::{
aws_auth_config, aws_signature_headers, host_supplied_credentials,
is_sigv4_computed_header, resolve_credentials, sign_bedrock_post,
};
@ -135,16 +134,3 @@ pub(super) async fn signed_headers(
// that would collide, so no name appears twice.
Ok(unsigned.into_iter().chain(signature).collect())
}
#[cfg(not(feature = "bedrock-auth"))]
pub(super) async fn signed_headers(
request: &ProviderChatCompletionsRequest,
_body: &[u8],
) -> Result<Vec<(String, String)>, Error> {
match &request.auth {
ChatCompletionsAuth::AwsSigV4 { .. } => Err(Error::Unsupported(
"AWS SigV4 requires the bedrock-auth feature",
)),
_ => Ok(request.upstream_headers.clone()),
}
}

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

@ -6,20 +6,20 @@
//! credentials, and it resolves the provider, translates the conversation,
//! calls the provider, and returns a typed OpenAI-shaped response.
use crate::Error;
mod error;
pub use error::Error;
mod client;
mod common_utils;
pub mod conversation;
pub(crate) mod handler;
mod prepare;
pub mod response_utils;
pub mod transformation;
pub mod streaming;
pub mod types;
use serde_json::{Map, Value};
use handler::execute_chat_completions_provider_call;
use prepare::{parse_messages, resolve_provider_config, resolve_request};
use serde_json::{Map, Value};
use types::{ChatCompletionsRequest, ChatCompletionsResponse};
pub async fn chat_completions(

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

@ -1,20 +1,21 @@
use serde_json::Value;
use crate::error::Error;
use crate::http_utils::has_header;
use crate::providers::custom_llm_provider::{CustomLlmProvider, get_custom_llm_provider};
use super::Error;
use super::common_utils::{chat_completions_provider_config, string_headers};
use super::transformation::{ChatCompletionsAuth, ChatCompletionsProviderConfig};
use super::types::{
ChatCompletionsRequest, ChatMessage, ProviderChatCompletionsRequest,
ResolvedChatCompletionsRequest,
};
use crate::http_utils::has_header;
use crate::litellm_core_utils::get_llm_provider_logic::{
CustomLlmProvider, get_custom_llm_provider,
};
use crate::llms::base_llm::chat::transformation::{BaseConfig, ChatCompletionsAuth};
pub(super) fn resolve_provider_config<'a>(
model: &'a str,
custom_llm_provider: Option<&'a str>,
) -> Result<(String, &'static dyn ChatCompletionsProviderConfig), Error> {
) -> Result<(String, &'static dyn BaseConfig), Error> {
let provider_info = get_custom_llm_provider(model, custom_llm_provider)
.or_else(|| {
custom_llm_provider.map(|provider| CustomLlmProvider {
@ -65,7 +66,7 @@ pub(super) fn resolve_request(
fn validate_environment(
request: &ResolvedChatCompletionsRequest<'_>,
model: &str,
config: &dyn ChatCompletionsProviderConfig,
config: &dyn BaseConfig,
) -> Result<(Vec<(String, String)>, ChatCompletionsAuth), Error> {
let env_lookup = |key: &str| std::env::var(key).ok();
let mut headers = string_headers(request.extra_headers.clone())?;
@ -122,7 +123,7 @@ pub(super) fn prepare_provider_request(
let model = request.model;
let config = request.config;
let env_lookup = |key: &str| std::env::var(key).ok();
let url = config.complete_url(
let url = config.get_complete_url(
request.api_base,
&model,
&request.optional_params,

9
litellm-rust/crates/core/src/chat_completions/streaming.rs Normal file
View file

@ -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>;
}

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

@ -1,10 +1,9 @@
use serde_json::{Map, Value, json};
use crate::error::Error;
use super::Error;
use super::prepare::{prepare_provider_request, resolve_request};
use super::transformation::ChatCompletionsAuth;
use super::types::{ChatCompletionsRequest, ProviderChatCompletionsRequest};
use crate::llms::base_llm::chat::transformation::ChatCompletionsAuth;
fn prepare_chat_completions_call(
request: ChatCompletionsRequest<'_>,
@ -264,13 +263,14 @@ fn rejects_non_string_extra_headers() {
call.extra_headers = Some(Map::from_iter([("x-trace".to_string(), json!(7))]));
assert_eq!(
decline(call),
Error::InvalidRequest(
"chat completions extra_headers.x-trace must be a string, got number".to_string()
)
Error::Headers(crate::http_utils::HeaderError {
context: "chat completions",
name: "x-trace".to_string(),
actual: "number",
})
);
}
#[cfg(feature = "bedrock-auth")]
#[test]
fn prepares_a_bedrock_call_without_resolving_credentials() {
let mut call = request(
@ -302,7 +302,6 @@ fn prepares_a_bedrock_call_without_resolving_credentials() {
assert_eq!(prepared.body["inferenceConfig"], json!({"maxTokens": 16}));
}
#[cfg(feature = "bedrock-auth")]
#[tokio::test]
async fn a_forwarded_client_header_does_not_enter_the_bedrock_signature() {
// Python signs only the AWS header set and reattaches the rest, so a header
@ -351,7 +350,6 @@ async fn a_forwarded_client_header_does_not_enter_the_bedrock_signature() {
);
}
#[cfg(feature = "bedrock-auth")]
#[tokio::test]
async fn a_forwarded_header_the_signer_computes_declines_to_python() {
// Reattaching the caller's copy next to the computed one puts the name on
@ -386,7 +384,6 @@ async fn a_forwarded_header_the_signer_computes_declines_to_python() {
}
}
#[cfg(feature = "bedrock-auth")]
#[test]
fn a_bedrock_deployment_bearer_outranks_a_forwarded_authorization() {
// `get_request_headers` assigns `headers["Authorization"]` unconditionally
@ -453,7 +450,6 @@ fn an_anthropic_forwarded_oauth_bearer_still_outranks_the_resolved_key() {
);
}
#[cfg(feature = "bedrock-auth")]
#[test]
fn a_bedrock_api_key_is_sent_as_a_bearer_token_instead_of_being_signed() {
// The configured bearer identity has its own account and quota boundary,
@ -591,10 +587,10 @@ fn the_gate_agrees_with_prepare_on_every_case_it_accepts() {
}
mod round_trip {
use super::*;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::{TcpListener, TcpStream};
use super::*;
use crate::chat_completions::chat_completions;
async fn read_http_request(socket: &mut TcpStream) -> String {
@ -769,7 +765,10 @@ mod round_trip {
.expect_err("upstream rejects");
handle.await.expect("server task");
assert!(
matches!(err, Error::Http { status: 429, .. }),
matches!(
err,
Error::Transport(crate::transport::Error::Http { status: 429, .. })
),
"expected a 429, got {err:?}"
);
}
@ -793,7 +792,7 @@ mod round_trip {
.await
.expect_err("nothing is listening");
assert!(
matches!(err, Error::Connect(_)),
matches!(err, Error::Transport(crate::transport::Error::Connect(_))),
"expected a pre-send connect failure, got {err:?}"
);
}
@ -806,7 +805,7 @@ mod round_trip {
Error::MissingField("usage"),
Error::Unsupported("non-text response content block"),
Error::InvalidRequest("whatever".to_string()),
Error::Auth("whatever".to_string()),
Error::Auth(litellm_auth::Error::InvalidHeader),
] {
let label = format!("{original:?}");
assert!(
@ -816,11 +815,11 @@ mod round_trip {
}
// An upstream status is already unambiguous, so it survives intact.
assert!(matches!(
as_response_error(Error::Http {
as_response_error(Error::Transport(crate::transport::Error::Http {
status: 500,
body: "boom".to_string()
}),
Error::Http { status: 500, .. }
})),
Error::Transport(crate::transport::Error::Http { status: 500, .. })
));
}
}

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

@ -3,7 +3,7 @@ use std::time::Duration;
use serde::{Deserialize, Serialize};
use serde_json::{Map, Value};
use super::transformation::{ChatCompletionsAuth, ChatCompletionsProviderConfig};
use crate::llms::base_llm::chat::transformation::{BaseConfig, ChatCompletionsAuth};
/// A `/chat/completions` call as it crosses into the core.
///
@ -24,7 +24,7 @@ pub struct ChatCompletionsRequest<'a> {
pub(super) struct ResolvedChatCompletionsRequest<'a> {
pub(super) model: String,
pub(super) config: &'static dyn ChatCompletionsProviderConfig,
pub(super) config: &'static dyn BaseConfig,
pub(super) messages: Vec<ChatMessage>,
pub(super) optional_params: Map<String, Value>,
pub(super) api_key: Option<&'a str>,
@ -35,12 +35,11 @@ pub(super) struct ResolvedChatCompletionsRequest<'a> {
pub(super) struct ProviderChatCompletionsRequest {
pub(super) model: String,
pub(super) config: &'static dyn ChatCompletionsProviderConfig,
pub(super) config: &'static dyn BaseConfig,
pub(super) url: String,
pub(super) body: Value,
pub(super) upstream_headers: Vec<(String, String)>,
pub(super) auth: ChatCompletionsAuth,
#[cfg_attr(not(feature = "bedrock-auth"), allow(dead_code))]
pub(super) optional_params: Map<String, Value>,
pub(super) timeout: Option<Duration>,
}
@ -121,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>>,
}

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

@ -1,220 +1,13 @@
use thiserror::Error as ThisError;
#[derive(Clone, Debug, ThisError, PartialEq, Eq)]
#[derive(Debug, thiserror::Error)]
pub enum Error {
#[error("expected {expected}, got {actual}")]
InvalidType {
expected: &'static str,
actual: &'static str,
},
#[error("missing required field: {0}")]
MissingField(&'static str),
#[error("Document URL is required")]
MissingDocumentUrl,
#[error("invalid response: {0}")]
InvalidResponse(String),
#[error("invalid provider: {0}")]
InvalidProvider(String),
#[error("invalid request: {0}")]
InvalidRequest(String),
#[error("{0}")]
Auth(String),
#[error(
"Missing {provider} API Key - A call is being made to {provider} but no key is set either in the environment variables or via params"
)]
MissingApiKey { provider: &'static str },
#[error(
"invalid authentication configuration: Missing Azure AI credentials - set AZURE_AI_API_KEY or configure Entra ID"
)]
MissingAzureAiCredentials,
#[error(
"invalid authentication configuration: Missing Azure Document Intelligence credentials - set AZURE_DOCUMENT_INTELLIGENCE_API_KEY or configure Entra ID"
)]
MissingAzureDocumentIntelligenceCredentials,
#[error(
"Missing REDUCTO_API_KEY - set it in the environment or pass api_key to litellm.ocr()/litellm.aocr()"
)]
MissingReductoApiKey,
#[error("upstream request failed with status {status}: {body}")]
Http { status: u16, body: String },
#[error("upstream network error: {0}")]
Network(String),
/// The provider was never reached: DNS, TCP, TLS or proxy setup failed
/// before any byte of the request went out. Nothing was billed, so a host
/// that keeps a reference implementation can serve the request itself.
/// A timeout is deliberately not this, since the provider may have received
/// and answered the request already.
#[error("could not reach the provider: {0}")]
Connect(String),
#[error("routing error: {0}")]
Routing(String),
/// The request is outside the surface this route covers in Rust. Hosts that
/// keep a reference implementation treat this as "fall back", not "fail".
#[error("unsupported by the rust path: {0}")]
Unsupported(&'static str),
}
impl Error {
pub const fn http_status_code(&self) -> Option<u16> {
match self {
Self::InvalidRequest(_) => Some(400),
Self::MissingDocumentUrl => Some(500),
Self::Http { status, .. } => Some(*status),
_ => None,
}
}
}
#[derive(Debug, ThisError)]
pub(crate) enum MediaError {
#[error("media URL rejected by network policy")]
BlockedUrl,
#[error("media download is disabled")]
DownloadDisabled,
#[error("media download exceeds the maximum size")]
DownloadTooLarge,
#[error("too many redirects while fetching media")]
TooManyRedirects,
#[error("media redirect is missing a Location header")]
MissingRedirectLocation,
#[error("invalid media redirect")]
InvalidRedirect,
#[error("media download failed with status {0}")]
Http(u16),
#[error("media download timed out")]
Timeout,
#[error("{0}")]
Transport(#[from] TransportError),
}
#[derive(Clone, Debug, ThisError, PartialEq, Eq)]
pub enum TransportError {
#[error("upstream request failed with status {status}: {body}")]
Http { status: u16, body: String },
#[error("upstream network error: {0}")]
Network(String),
#[error("could not reach the provider: {0}")]
Connect(String),
}
impl TransportError {
pub fn from_reqwest_before_dispatch(error: reqwest::Error) -> Self {
let before_dispatch = !error.is_timeout() && (error.is_connect() || error.is_builder());
let message = error.without_url().to_string();
if before_dispatch {
Self::Connect(message)
} else {
Self::Network(message)
}
}
}
impl From<reqwest::Error> for TransportError {
fn from(error: reqwest::Error) -> Self {
Self::Network(error.without_url().to_string())
}
}
impl From<crate::ocr::error::OcrRequestError> for Error {
fn from(error: crate::ocr::error::OcrRequestError) -> Self {
match error {
crate::ocr::error::OcrRequestError::MissingField(field) => Self::MissingField(field),
crate::ocr::error::OcrRequestError::MissingDocumentUrl => Self::MissingDocumentUrl,
error => Self::InvalidRequest(error.to_string()),
}
}
}
impl From<crate::ocr::error::OcrResponseError> for Error {
fn from(error: crate::ocr::error::OcrResponseError) -> Self {
Self::InvalidResponse(error.to_string())
}
}
impl From<TransportError> for Error {
fn from(error: TransportError) -> Self {
match error {
TransportError::Http { status, body } => Self::Http { status, body },
TransportError::Network(message) => Self::Network(message),
TransportError::Connect(message) => Self::Connect(message),
}
}
}
impl From<crate::AuthError> for Error {
fn from(error: crate::AuthError) -> Self {
match error {
crate::AuthError::MissingApiKey { provider } => Self::MissingApiKey { provider },
error => Self::Auth(error.to_string()),
}
}
}
pub fn json_type_name(value: &serde_json::Value) -> &'static str {
match value {
serde_json::Value::Null => "null",
serde_json::Value::Bool(_) => "bool",
serde_json::Value::Number(_) => "number",
serde_json::Value::String(_) => "string",
serde_json::Value::Array(_) => "array",
serde_json::Value::Object(_) => "object",
}
}
#[cfg(test)]
mod transport_tests {
use super::*;
#[test]
fn missing_auth_key_preserves_provider_in_public_error() {
assert_eq!(
Error::from(crate::AuthError::MissingApiKey { provider: "Vertex" }),
Error::MissingApiKey { provider: "Vertex" }
);
}
#[tokio::test]
async fn transport_errors_remove_urls_and_keep_dispatch_context() {
let error = reqwest::Client::builder()
.no_proxy()
.build()
.expect("client")
.get("http://localhost:invalid/private?api_key=secret")
.send()
.await
.expect_err("invalid port");
let error = TransportError::from_reqwest_before_dispatch(error);
assert!(matches!(error, TransportError::Connect(_)));
assert!(!error.to_string().contains("secret"));
assert!(!error.to_string().contains("private"));
}
#[tokio::test]
async fn request_timeout_is_not_safe_to_retry_as_a_connect_failure() {
use std::time::Duration;
let listener = tokio::net::TcpListener::bind("127.0.0.1:0")
.await
.expect("bind");
let address = listener.local_addr().expect("address");
let request = reqwest::Client::builder()
.no_proxy()
.build()
.expect("client")
.get(format!("http://{address}"))
.timeout(Duration::from_millis(200))
.send();
let (response, accepted) = tokio::join!(
request,
tokio::time::timeout(Duration::from_secs(2), listener.accept())
);
let _connection = accepted
.expect("accept deadline")
.expect("accepted connection");
let error = response.expect_err("server does not respond");
assert!(error.is_timeout());
assert!(matches!(
TransportError::from_reqwest_before_dispatch(error),
TransportError::Network(_)
));
}
#[error(transparent)]
Ocr(#[from] crate::ocr::Error),
#[error(transparent)]
Messages(#[from] crate::messages::Error),
#[error(transparent)]
ChatCompletions(#[from] crate::chat_completions::Error),
#[error(transparent)]
AudioTranscription(#[from] crate::audio_transcription::Error),
#[error(transparent)]
Responses(#[from] crate::responses::Error),
}

49
litellm-rust/crates/core/src/http_utils.rs
View file

@ -1,7 +1,14 @@
#[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)]
#[error("invalid request: {context} extra_headers.{name} must be a string, got {actual}")]
pub struct HeaderError {
pub context: &'static str,
pub name: String,
pub actual: &'static str,
}
use serde_json::{Map, Value};
use crate::constants::UPSTREAM_ERROR_BODY_MAX_CHARS;
use crate::error::{Error, json_type_name};
#[allow(
dead_code,
@ -44,6 +51,13 @@ pub async fn http_request(
request.send().await
}
pub async fn execute_http_request(
client: &reqwest::Client,
request: reqwest::Request,
) -> Result<reqwest::Response, reqwest::Error> {
client.execute(request).await
}
pub fn truncate_error_body(body: &str) -> String {
if body.chars().count() <= UPSTREAM_ERROR_BODY_MAX_CHARS {
return body.to_string();
@ -55,7 +69,7 @@ pub fn truncate_error_body(body: &str) -> String {
pub fn string_headers(
context: &'static str,
extra_headers: Option<Map<String, Value>>,
) -> Result<Vec<(String, String)>, Error> {
) -> Result<Vec<(String, String)>, HeaderError> {
extra_headers
.unwrap_or_default()
.into_iter()
@ -63,11 +77,10 @@ pub fn string_headers(
value
.as_str()
.map(|value| (key.clone(), value.to_string()))
.ok_or_else(|| {
Error::InvalidRequest(format!(
"{context} extra_headers.{key} must be a string, got {}",
json_type_name(&value)
))
.ok_or_else(|| HeaderError {
context,
name: key,
actual: json_type_name(&value),
})
})
.collect()
@ -105,11 +118,23 @@ where
<Option<T> as serde::Deserialize>::deserialize(deserializer).map(Some)
}
pub fn json_type_name(value: &serde_json::Value) -> &'static str {
match value {
serde_json::Value::Null => "null",
serde_json::Value::Bool(_) => "bool",
serde_json::Value::Number(_) => "number",
serde_json::Value::String(_) => "string",
serde_json::Value::Array(_) => "array",
serde_json::Value::Object(_) => "object",
}
}
#[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
use super::*;
#[rstest::rstest]
#[case(HeaderPolicy::All, true, true)]
#[case(HeaderPolicy::Only(&["authorization"]), true, false)]
@ -184,9 +209,11 @@ mod tests {
let err = string_headers("chat completions", Some(headers)).expect_err("non-string value");
assert_eq!(
err,
Error::InvalidRequest(
"chat completions extra_headers.x-trace must be a string, got number".to_string()
)
HeaderError {
context: "chat completions",
name: "x-trace".into(),
actual: "number"
}
);
}

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

@ -1,17 +1,19 @@
pub mod audio_transcription;
pub mod auth;
pub mod caching;
pub mod call_arguments;
pub mod call_lifecycle;
pub mod chat_completions;
pub mod constants;
pub mod error;
pub mod http_utils;
pub mod litellm_core_utils;
pub mod llms;
mod media;
pub mod messages;
pub mod ocr;
pub mod providers;
pub mod params;
pub mod responses;
mod serde_compat;
pub mod transport;
mod url_utils;
pub use auth::AuthError;
pub use error::Error;

0
litellm-rust/crates/core/src/providers/custom_llm_provider.rs → litellm-rust/crates/core/src/litellm_core_utils/get_llm_provider_logic.rs
View file

1
litellm-rust/crates/core/src/litellm_core_utils/mod.rs Normal file
View file

@ -0,0 +1 @@
pub mod get_llm_provider_logic;

1
litellm-rust/crates/core/src/providers/openai/responses/mod.rs → litellm-rust/crates/core/src/llms/anthropic/chat/mod.rs
View file

@ -1 +1,2 @@
pub mod streaming;
pub mod transformation;

164
litellm-rust/crates/core/src/llms/anthropic/chat/streaming.rs Normal file
View file

@ -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::llms::anthropic::experimental_pass_through::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!()
}
}

7
litellm-rust/crates/core/src/providers/anthropic/chat_completions/tests.rs → litellm-rust/crates/core/src/llms/anthropic/chat/tests.rs
View file

@ -1,7 +1,8 @@
use super::*;
use crate::Error;
use serde_json::json;
use super::*;
use crate::chat_completions::Error;
fn messages(value: Value) -> Vec<ChatMessage> {
serde_json::from_value(value).expect("valid messages")
}
@ -419,7 +420,7 @@ fn resolves_the_messages_url_and_x_api_key_auth() {
let config = &ANTHROPIC_CHAT_COMPLETIONS_CONFIG;
assert_eq!(
config
.complete_url(None, "claude-sonnet-4-5", &Map::new(), &|_| None)
.get_complete_url(None, "claude-sonnet-4-5", &Map::new(), &|_| None)
.expect("url builds"),
"https://api.anthropic.com/v1/messages"
);

199
litellm-rust/crates/core/src/providers/anthropic/chat_completions/transformation.rs → litellm-rust/crates/core/src/llms/anthropic/chat/transformation.rs
View file

@ -1,21 +1,19 @@
use serde_json::{Map, Value, json};
use crate::chat_completions::Error;
use crate::chat_completions::conversation::{Conversation, build_conversation};
use crate::chat_completions::transformation::{
ChatCompletionsAuth, ChatCompletionsProviderConfig, Unsupported, unsupported_message,
unsupported_param,
};
use crate::chat_completions::response_utils::{finish_reason_for, unix_now, usage_from_parts};
use crate::chat_completions::types::{
ChatCompletionsChoice, ChatCompletionsChoiceMessage, ChatCompletionsResponse, ChatMessage,
ProviderChatRequestData, ProviderChatResponseData,
};
use crate::constants::ANTHROPIC_OAUTH_TOKEN_PREFIX;
use crate::error::Error;
use crate::providers::anthropic::messages::transformation::{
use crate::llms::anthropic::experimental_pass_through::messages::transformation::{
complete_anthropic_url, resolve_anthropic_api_key,
};
use crate::chat_completions::response_utils::{finish_reason_for, unix_now, usage_from_parts};
use crate::llms::base_llm::chat::transformation::{
BaseConfig, ChatCompletionsAuth, Unsupported, unsupported_message, unsupported_param,
};
/// Anthropic parameter names, post `map_openai_params`, that the Rust path can
/// place verbatim in the Messages body.
@ -34,46 +32,16 @@ const SUPPORTED_PARAMS: &[(&str, &str)] = &[
("stop", "stop_sequences"),
];
pub struct AnthropicChatCompletionsConfig;
pub struct AnthropicConfig;
pub const ANTHROPIC_CHAT_COMPLETIONS_CONFIG: AnthropicChatCompletionsConfig =
AnthropicChatCompletionsConfig;
pub const ANTHROPIC_CHAT_COMPLETIONS_CONFIG: AnthropicConfig = AnthropicConfig;
fn text_block(text: &str) -> Value {
json!({"type": "text", "text": text})
}
impl BaseConfig for AnthropicConfig {
fn supported_openai_param_mappings(&self) -> &'static [(&'static str, &'static str)] {
SUPPORTED_PARAMS
}
fn anthropic_body(model: &str, conversation: &Conversation, params: Map<String, Value>) -> Value {
let messages: Vec<Value> = conversation
.turns
.iter()
.map(|turn| {
json!({
"role": turn.role.as_str(),
"content": turn.texts.iter().map(|text| text_block(text)).collect::<Vec<_>>(),
})
})
.collect();
let system: Vec<Value> = conversation.system.iter().map(|s| text_block(s)).collect();
let body = Map::from_iter(
[
("model".to_string(), json!(model)),
("messages".to_string(), json!(messages)),
]
.into_iter()
// Python builds `{"model", "messages", **optional_params}` with
// `system` already folded into optional_params, so a caller-supplied
// key of the same name wins here too.
.chain((!system.is_empty()).then(|| ("system".to_string(), json!(system))))
.chain(params),
);
Value::Object(body)
}
impl ChatCompletionsProviderConfig for AnthropicChatCompletionsConfig {
fn complete_url(
fn get_complete_url(
&self,
api_base: Option<&str>,
_model: &str,
@ -83,60 +51,6 @@ impl ChatCompletionsProviderConfig for AnthropicChatCompletionsConfig {
Ok(complete_anthropic_url(api_base, env_lookup))
}
fn auth(
&self,
api_key: Option<&str>,
_model: &str,
_optional_params: &Map<String, Value>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Result<ChatCompletionsAuth, Error> {
Ok(ChatCompletionsAuth::Header {
name: "x-api-key",
value: resolve_anthropic_api_key(api_key, env_lookup)?,
})
}
fn default_headers(&self) -> &'static [(&'static str, &'static str)] {
&[
("anthropic-version", "2023-06-01"),
("content-type", "application/json"),
]
}
/// An OAuth bearer is the whole credential: Python's `validate_environment`
/// authenticates with it and drops `x-api-key` rather than resolving one, so
/// the resolved key must not be applied over the top. Any other forwarded
/// `authorization` is unrelated to this header and does not defer, which is
/// also what Python does: it sends the deployment's `x-api-key` alongside.
fn defers_to_forwarded_auth(&self, headers: &[(String, String)]) -> bool {
headers.iter().any(|(name, value)| {
name.eq_ignore_ascii_case("authorization")
&& value
.strip_prefix("Bearer ")
.is_some_and(|token| token.starts_with(ANTHROPIC_OAUTH_TOKEN_PREFIX))
})
}
fn supported_openai_params(&self) -> &'static [(&'static str, &'static str)] {
SUPPORTED_PARAMS
}
fn unsupported_reason(
&self,
messages: &[ChatMessage],
optional_params: &Map<String, Value>,
) -> Option<Unsupported> {
unsupported_param(self.supported_openai_params(), &[], optional_params)
.or_else(|| messages.iter().find_map(unsupported_message))
// Anthropic rejects a request whose first turn is not a user turn.
// Python only repairs that under `litellm.modify_params`, which the
// core cannot observe, so decline instead of guessing.
.or_else(|| {
(!build_conversation(messages).opens_on_user_turn())
.then_some(Unsupported("conversation does not open on a user turn"))
})
}
fn transform_request(
&self,
model: &str,
@ -210,6 +124,93 @@ impl ChatCompletionsProviderConfig for AnthropicChatCompletionsConfig {
),
})
}
fn auth(
&self,
api_key: Option<&str>,
_model: &str,
_optional_params: &Map<String, Value>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Result<ChatCompletionsAuth, Error> {
Ok(ChatCompletionsAuth::Header {
name: "x-api-key",
value: resolve_anthropic_api_key(api_key, env_lookup)?,
})
}
fn default_headers(&self) -> &'static [(&'static str, &'static str)] {
&[
("anthropic-version", "2023-06-01"),
("content-type", "application/json"),
]
}
/// An OAuth bearer is the whole credential: Python's `validate_environment`
/// authenticates with it and drops `x-api-key` rather than resolving one, so
/// the resolved key must not be applied over the top. Any other forwarded
/// `authorization` is unrelated to this header and does not defer, which is
/// also what Python does: it sends the deployment's `x-api-key` alongside.
fn defers_to_forwarded_auth(&self, headers: &[(String, String)]) -> bool {
headers.iter().any(|(name, value)| {
name.eq_ignore_ascii_case("authorization")
&& value
.strip_prefix("Bearer ")
.is_some_and(|token| token.starts_with(ANTHROPIC_OAUTH_TOKEN_PREFIX))
})
}
fn unsupported_reason(
&self,
messages: &[ChatMessage],
optional_params: &Map<String, Value>,
) -> Option<Unsupported> {
unsupported_param(self.supported_openai_param_mappings(), &[], optional_params)
.or_else(|| messages.iter().find_map(unsupported_message))
// Anthropic rejects a request whose first turn is not a user turn.
// Python only repairs that under `litellm.modify_params`, which the
// core cannot observe, so decline instead of guessing.
.or_else(|| {
(!build_conversation(messages).opens_on_user_turn())
.then_some(Unsupported("conversation does not open on a user turn"))
})
}
}
fn text_block(text: &str) -> Value {
json!({"type": "text", "text": text})
}
fn anthropic_body(
model: &str,
conversation: &Conversation,
optional_params: Map<String, Value>,
) -> Value {
let messages: Vec<Value> = conversation
.turns
.iter()
.map(|turn| {
json!({
"role": turn.role.as_str(),
"content": turn.texts.iter().map(|text| text_block(text)).collect::<Vec<_>>(),
})
})
.collect();
let system: Vec<Value> = conversation.system.iter().map(|s| text_block(s)).collect();
let body = Map::from_iter(
[
("model".to_string(), json!(model)),
("messages".to_string(), json!(messages)),
]
.into_iter()
// Python builds `{"model", "messages", **optional_params}` with
// `system` already folded into optional_params, so a caller-supplied
// key of the same name wins here too.
.chain((!system.is_empty()).then(|| ("system".to_string(), json!(system))))
.chain(optional_params),
);
Value::Object(body)
}
#[cfg(test)]

338
litellm-rust/crates/core/src/llms/anthropic/experimental_pass_through/messages/batches.rs Normal file
View file

@ -0,0 +1,338 @@
use serde::{Deserialize, Serialize};
use serde_json::Value;
use time::OffsetDateTime;
use url::Url;
use crate::llms::anthropic::experimental_pass_through::messages::transformation::resolve_anthropic_api_base;
use crate::messages::Error;
use crate::messages::types::AnthropicMessagesResponse;
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"))
));
}
}

168
litellm-rust/crates/core/src/llms/anthropic/experimental_pass_through/messages/count_tokens.rs Normal file
View file

@ -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())));
}
}

4
litellm-rust/crates/core/src/llms/anthropic/experimental_pass_through/messages/mod.rs Normal file
View file

@ -0,0 +1,4 @@
pub mod batches;
pub mod count_tokens;
pub mod streaming;
pub mod transformation;

282
litellm-rust/crates/core/src/llms/anthropic/experimental_pass_through/messages/streaming.rs Normal file
View file

@ -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(),
},
}]
);
}
}

91
litellm-rust/crates/core/src/providers/anthropic/messages/transformation.rs → litellm-rust/crates/core/src/llms/anthropic/experimental_pass_through/messages/transformation.rs
View file

@ -1,6 +1,5 @@
use crate::auth::error::MissingCredential;
use crate::error::Error;
use crate::messages::transformation::{AnthropicMessagesProviderConfig, MessagesAuthStrategy};
use crate::llms::base_llm::anthropic_messages::transformation::BaseAnthropicMessagesConfig;
use crate::messages::Error;
const ANTHROPIC_API_KEY_ENV: &str = "ANTHROPIC_API_KEY";
const ANTHROPIC_API_BASE_ENV: &str = "ANTHROPIC_API_BASE";
@ -11,38 +10,8 @@ pub struct AnthropicMessagesConfig;
pub const ANTHROPIC_MESSAGES_CONFIG: AnthropicMessagesConfig = AnthropicMessagesConfig;
pub fn non_empty(value: Option<&str>) -> Option<&str> {
value.map(str::trim).filter(|value| !value.is_empty())
}
pub fn resolve_anthropic_api_key(
api_key: Option<&str>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Result<String, Error> {
non_empty(api_key)
.map(str::to_string)
.or_else(|| env_lookup(ANTHROPIC_API_KEY_ENV).filter(|value| !value.trim().is_empty()))
.ok_or_else(|| Error::from(crate::AuthError::from(MissingCredential::AnthropicApiKey)))
}
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 = api_base.trim_end_matches('/');
if api_base.ends_with(MESSAGES_PATH_SUFFIX) {
return api_base.to_string();
}
format!("{api_base}{MESSAGES_PATH_SUFFIX}")
}
impl AnthropicMessagesProviderConfig for AnthropicMessagesConfig {
fn complete_url(
impl BaseAnthropicMessagesConfig for AnthropicMessagesConfig {
fn get_complete_url(
&self,
api_base: Option<&str>,
_model: &str,
@ -56,12 +25,48 @@ impl AnthropicMessagesProviderConfig for AnthropicMessagesConfig {
api_key: Option<&str>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Result<String, Error> {
resolve_anthropic_api_key(api_key, env_lookup)
resolve_anthropic_api_key(api_key, env_lookup).map_err(Error::from)
}
}
fn auth_strategy(&self) -> MessagesAuthStrategy {
MessagesAuthStrategy::Header("x-api-key")
pub fn non_empty(value: Option<&str>) -> Option<&str> {
value.map(str::trim).filter(|value| !value.is_empty())
}
pub fn resolve_anthropic_api_key(
api_key: Option<&str>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Result<String, litellm_auth::Error> {
non_empty(api_key)
.map(str::to_string)
.or_else(|| env_lookup(ANTHROPIC_API_KEY_ENV).filter(|value| !value.trim().is_empty()))
.ok_or(litellm_auth::Error::MissingApiKey {
provider: "Anthropic",
environment_variable: ANTHROPIC_API_KEY_ENV,
})
}
pub fn complete_anthropic_url(
api_base: Option<&str>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> 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) {
return api_base.to_string();
}
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())
}
#[cfg(test)]
@ -114,10 +119,12 @@ mod tests {
resolve_anthropic_api_key(Some(" "), &with_env).unwrap(),
"sk-env"
);
assert!(matches!(
resolve_anthropic_api_key(None, &|_| None).expect_err("missing key"),
Error::Auth(_)
));
assert_eq!(
resolve_anthropic_api_key(None, &|_| None)
.expect_err("missing key")
.to_string(),
"Missing Anthropic API Key - Set `api_key` or the ANTHROPIC_API_KEY environment variable"
);
}
#[test]

1
litellm-rust/crates/core/src/llms/anthropic/experimental_pass_through/mod.rs Normal file
View file

@ -0,0 +1 @@
pub mod messages;

2
litellm-rust/crates/core/src/llms/anthropic/mod.rs Normal file
View file

@ -0,0 +1,2 @@
pub mod chat;
pub mod experimental_pass_through;

153
litellm-rust/crates/core/src/providers/azure_ai/messages/transformation.rs → litellm-rust/crates/core/src/llms/azure_ai/anthropic/messages_transformation.rs
View file

@ -1,14 +1,16 @@
use crate::auth::error::MissingCredential;
use crate::error::Error;
use crate::messages::transformation::{AnthropicMessagesProviderConfig, MessagesAuthStrategy};
use serde_json::{Map, Value};
use crate::llms::anthropic::experimental_pass_through::messages::transformation::{
ANTHROPIC_MESSAGES_CONFIG, AnthropicMessagesConfig, non_empty,
};
use crate::llms::base_llm::anthropic_messages::transformation::{
BaseAnthropicMessagesConfig, MessagesAuthStrategy,
};
use crate::messages::Error;
use crate::messages::types::{
AnthropicMessage, AnthropicMessagesRequest, AnthropicMessagesResponse, ContentBlock,
MessageContent, SystemPrompt,
};
use crate::providers::anthropic::messages::transformation::{
ANTHROPIC_MESSAGES_CONFIG, AnthropicMessagesConfig, non_empty,
};
use serde_json::{Map, Value};
const AZURE_API_KEY_ENV: &str = "AZURE_API_KEY";
const AZURE_API_BASE_ENV: &str = "AZURE_API_BASE";
@ -26,6 +28,61 @@ pub const AZURE_ANTHROPIC_MESSAGES_CONFIG: AzureAnthropicMessagesConfig =
anthropic: ANTHROPIC_MESSAGES_CONFIG,
};
impl BaseAnthropicMessagesConfig for AzureAnthropicMessagesConfig {
fn get_complete_url(
&self,
api_base: Option<&str>,
_model: &str,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Result<String, Error> {
complete_azure_anthropic_url(api_base, env_lookup)
}
fn transform_anthropic_messages_request(
&self,
request: AnthropicMessagesRequest,
) -> Result<AnthropicMessagesRequest, Error> {
let mut request = fold_system_role_messages(request);
if let Some(system) = request.system.as_mut() {
strip_scope_from_system(system);
}
request
.messages
.iter_mut()
.for_each(strip_scope_from_message);
self.anthropic.transform_anthropic_messages_request(request)
}
fn transform_anthropic_messages_response(
&self,
model: &str,
response: AnthropicMessagesResponse,
) -> Result<AnthropicMessagesResponse, Error> {
self.anthropic
.transform_anthropic_messages_response(model, response)
}
fn resolve_api_key(
&self,
api_key: Option<&str>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Result<String, Error> {
resolve_azure_api_key(api_key, env_lookup)
}
fn auth_strategy(&self) -> MessagesAuthStrategy {
self.anthropic.auth_strategy()
}
fn accepts_bearer_auth(&self) -> bool {
true
}
fn default_headers(&self) -> &'static [(&'static str, &'static str)] {
self.anthropic.default_headers()
}
}
pub fn resolve_azure_api_key(
api_key: Option<&str>,
env_lookup: &dyn Fn(&str) -> Option<String>,
@ -33,7 +90,12 @@ pub fn resolve_azure_api_key(
non_empty(api_key)
.map(str::to_string)
.or_else(|| env_lookup(AZURE_API_KEY_ENV).filter(|value| !value.trim().is_empty()))
.ok_or_else(|| Error::from(crate::AuthError::from(MissingCredential::AzureApiKey)))
.ok_or_else(|| {
Error::from(litellm_auth::Error::MissingApiKey {
provider: "Azure",
environment_variable: AZURE_API_KEY_ENV,
})
})
}
pub fn complete_azure_anthropic_url(
@ -43,7 +105,7 @@ pub fn complete_azure_anthropic_url(
let api_base = non_empty(api_base)
.map(str::to_string)
.or_else(|| env_lookup(AZURE_API_BASE_ENV).filter(|value| !value.trim().is_empty()))
.ok_or_else(|| Error::from(crate::AuthError::from(MissingCredential::AzureApiBase)))?;
.ok_or_else(|| Error::from(litellm_auth::Error::MissingAzureApiBase))?;
let api_base = api_base.trim_end_matches('/');
@ -131,65 +193,12 @@ fn fold_system_role_messages(request: AnthropicMessagesRequest) -> AnthropicMess
}
}
impl AnthropicMessagesProviderConfig for AzureAnthropicMessagesConfig {
fn complete_url(
&self,
api_base: Option<&str>,
_model: &str,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Result<String, Error> {
complete_azure_anthropic_url(api_base, env_lookup)
}
fn resolve_api_key(
&self,
api_key: Option<&str>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> Result<String, Error> {
resolve_azure_api_key(api_key, env_lookup)
}
fn auth_strategy(&self) -> MessagesAuthStrategy {
self.anthropic.auth_strategy()
}
fn accepts_bearer_auth(&self) -> bool {
true
}
fn default_headers(&self) -> &'static [(&'static str, &'static str)] {
self.anthropic.default_headers()
}
fn transform_request(
&self,
request: AnthropicMessagesRequest,
) -> Result<AnthropicMessagesRequest, Error> {
let mut request = fold_system_role_messages(request);
if let Some(system) = request.system.as_mut() {
strip_scope_from_system(system);
}
request
.messages
.iter_mut()
.for_each(strip_scope_from_message);
self.anthropic.transform_request(request)
}
fn transform_response(
&self,
model: &str,
response: AnthropicMessagesResponse,
) -> Result<AnthropicMessagesResponse, Error> {
self.anthropic.transform_response(model, response)
}
}
#[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
use super::*;
fn request_from(value: serde_json::Value) -> AnthropicMessagesRequest {
serde_json::from_value(value).expect("valid request")
}
@ -333,7 +342,7 @@ mod tests {
let transformed = to_value(
AZURE_ANTHROPIC_MESSAGES_CONFIG
.transform_request(request)
.transform_anthropic_messages_request(request)
.expect("request transforms"),
);
@ -360,10 +369,10 @@ mod tests {
"messages": [{"role": "user", "content": "hi"}]
}));
let once = AZURE_ANTHROPIC_MESSAGES_CONFIG
.transform_request(request)
.transform_anthropic_messages_request(request)
.expect("request transforms");
let twice = AZURE_ANTHROPIC_MESSAGES_CONFIG
.transform_request(once.clone())
.transform_anthropic_messages_request(once.clone())
.expect("request transforms");
assert_eq!(once, twice);
assert_eq!(to_value(once)["system"], json!("plain string system"));
@ -397,7 +406,7 @@ mod tests {
});
let transformed = to_value(
AZURE_ANTHROPIC_MESSAGES_CONFIG
.transform_request(request_from(body.clone()))
.transform_anthropic_messages_request(request_from(body.clone()))
.expect("request transforms"),
);
assert_eq!(transformed, body);
@ -417,7 +426,7 @@ mod tests {
let transformed = to_value(
AZURE_ANTHROPIC_MESSAGES_CONFIG
.transform_request(request)
.transform_anthropic_messages_request(request)
.expect("request transforms"),
);
@ -447,7 +456,7 @@ mod tests {
let transformed = to_value(
AZURE_ANTHROPIC_MESSAGES_CONFIG
.transform_request(request)
.transform_anthropic_messages_request(request)
.expect("request transforms"),
);
@ -474,7 +483,7 @@ mod tests {
});
let transformed = to_value(
AZURE_ANTHROPIC_MESSAGES_CONFIG
.transform_request(request_from(body.clone()))
.transform_anthropic_messages_request(request_from(body.clone()))
.expect("request transforms"),
);
assert_eq!(transformed, body);
@ -501,7 +510,7 @@ mod tests {
}))
.expect("valid response");
let transformed = AZURE_ANTHROPIC_MESSAGES_CONFIG
.transform_response("claude-sonnet-4-5", response)
.transform_anthropic_messages_response("claude-sonnet-4-5", response)
.expect("response transforms");
let value = serde_json::to_value(transformed).expect("serializable");
assert_eq!(value["stop_reason"], json!("end_turn"));

1
litellm-rust/crates/core/src/llms/azure_ai/anthropic/mod.rs Normal file
View file

@ -0,0 +1 @@
pub mod messages_transformation;

2
litellm-rust/crates/core/src/llms/azure_ai/mod.rs Normal file
View file

@ -0,0 +1,2 @@
pub mod anthropic;
pub(crate) mod ocr;

166
litellm-rust/crates/core/src/llms/azure_ai/ocr/cohere_parse_transformation.rs Normal file
View file

@ -0,0 +1,166 @@
use serde_json::Value;
use crate::call_arguments::CallArguments;
use crate::llms::base_llm::ocr::transformation::{BaseOcrConfig, OcrRequestContext};
use crate::llms::cohere::ocr::transformation::{CohereParseConfig, CohereRequest};
use crate::llms::cohere::ocr::{CohereOptions, validate_document};
use crate::ocr::OcrClient;
use crate::ocr::document::{inline_remote_document, validate_inline_document};
use crate::ocr::types::{LiteLLMOcrResponse, OcrDocument, PreparedOcrRequest};
use crate::url_utils::ApiUrl;
#[derive(Default)]
pub(crate) struct AzureAICohereParseConfig;
impl BaseOcrConfig for AzureAICohereParseConfig {
type OcrParams = CohereOptions;
type ProviderRequest = CohereRequest;
type Environment = Vec<(String, String)>;
fn get_api_key_env_var(&self) -> Option<&'static str> {
super::transformation::AzureAiOcrConfig.get_api_key_env_var()
}
fn get_health_check_document(&self) -> OcrDocument {
CohereParseConfig.get_health_check_document()
}
async fn validate_environment(
&self,
request: &PreparedOcrRequest,
client: &OcrClient,
) -> Result<Self::Environment, crate::ocr::Error> {
BaseOcrConfig::validate_environment(
&super::transformation::AzureAiOcrConfig,
request,
client,
)
.await
}
fn get_complete_url(
&self,
request: &PreparedOcrRequest,
_params: &Self::OcrParams,
_environment: &Self::Environment,
) -> Result<String, crate::ocr::Error> {
let base = super::transformation::AzureAiOcrConfig::resolve_api_base(
request.connection.api_base.as_deref(),
&crate::ocr::prepare::credential_env,
)?;
self.get_complete_url(&base)
}
fn transform_ocr_request(
&self,
model: &str,
document: OcrDocument,
params: &CohereOptions,
headers: &[(String, String)],
) -> Result<CohereRequest, crate::ocr::Error> {
CohereParseConfig.transform_ocr_request(model, document, params, headers)
}
fn get_supported_ocr_params(&self, model: &str) -> &'static [&'static str] {
CohereParseConfig.get_supported_ocr_params(model)
}
fn map_ocr_params(
&self,
arguments: &CallArguments,
model: &str,
) -> Result<CohereOptions, crate::ocr::Error> {
CohereParseConfig.map_ocr_params(arguments, model)
}
async fn async_transform_ocr_request(
&self,
model: &str,
document: OcrDocument,
optional_params: &CohereOptions,
headers: &[(String, String)],
context: OcrRequestContext<'_>,
) -> Result<CohereRequest, crate::ocr::Error> {
validate_document(&document)?;
let document = inline_remote_document(
context.client.document_fetcher(),
document,
context.connection,
)
.await?;
self.transform_ocr_request(model, document, optional_params, headers)
}
fn transform_ocr_response(
&self,
model: &str,
raw_response: &[u8],
request_format: crate::ocr::types::OcrResponseFormat,
) -> Result<LiteLLMOcrResponse, crate::ocr::Error> {
CohereParseConfig.transform_ocr_response(model, raw_response, request_format)
}
fn validate_request_body(&self, body: &Value) -> Result<(), crate::ocr::Error> {
let document = crate::ocr::prepare::body_document(body)?;
validate_document(&document)?;
validate_inline_document(&document)
}
}
impl AzureAICohereParseConfig {
fn get_complete_url(&self, base: &str) -> Result<String, crate::ocr::Error> {
let mut url = reqwest::Url::parse(base).map_err(|_| invalid_api_base())?;
if !matches!(url.scheme(), "http" | "https") {
return Err(invalid_api_base());
}
let path = url.path().trim_end_matches('/').to_string();
if path.ends_with("/v2/parse") {
url.set_path(&path);
return Ok(url.into());
}
url.set_path(path.strip_suffix("/models").unwrap_or(&path));
ApiUrl::parse(url.as_str())
.and_then(|url| url.complete_path(&["providers", "cohere", "v2", "parse"]))
.map(|url| url.into_string())
.map_err(|_| invalid_api_base())
}
}
fn invalid_api_base() -> crate::ocr::Error {
crate::ocr::Error::RequestField {
path: "api_base".into(),
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn completes_foundry_urls_without_duplicate_paths_and_preserves_queries() {
for suffix in [
"",
"/models",
"/providers/cohere/v2",
"/providers/cohere/v2/parse",
] {
assert_eq!(
AzureAICohereParseConfig
.get_complete_url(&format!("https://example.com{suffix}?tenant=a"))
.unwrap(),
"https://example.com/providers/cohere/v2/parse?tenant=a"
);
}
assert_eq!(
AzureAICohereParseConfig
.get_complete_url("https://example.com/v2/parse?tenant=a")
.unwrap(),
"https://example.com/v2/parse?tenant=a"
);
assert!(
AzureAICohereParseConfig
.get_complete_url("relative/path")
.is_err()
);
}
}

34
litellm-rust/crates/core/src/ocr/adapters/azure/mod.rs → litellm-rust/crates/core/src/llms/azure_ai/ocr/common_utils.rs
View file

@ -1,32 +1,21 @@
mod cohere;
mod document_intelligence;
mod mistral;
use std::sync::OnceLock;
use crate::Error;
use crate::auth::error::AuthConfigurationError;
use crate::auth::{InputSource, Sourced};
use crate::ocr::error::OcrError;
use litellm_auth::{InputSource, Sourced};
use litellm_auth_azure::{AzureAuthInputs, AzureAuthService};
use crate::ocr::types::OcrConnection;
use crate::providers::azure_ai::auth::{AzureAuthInputs, AzureAuthService};
pub(crate) use cohere::AzureCohereAdapter;
pub(crate) use document_intelligence::AzureDocumentIntelligenceAdapter;
pub(crate) use mistral::AzureMistralAdapter;
pub(super) use mistral::validate_environment as validate_ai_environment;
async fn resolve_entra(
pub(super) async fn resolve_entra(
config: &AzureAuthInputs,
env_lookup: &(dyn Fn(&str) -> Option<String> + Sync),
) -> Result<Option<Sourced<String>>, Error> {
) -> Result<Option<Sourced<String>>, crate::ocr::Error> {
static SERVICE: OnceLock<AzureAuthService> = OnceLock::new();
SERVICE
.get_or_init(AzureAuthService::default)
.get_azure_ad_token(config, env_lookup)
.await
.or_else(|error| match error {
crate::AuthError::EmptyAzureToken => Ok(None),
litellm_auth::Error::EmptyAzureToken => Ok(None),
other => Err(other),
})
.map(|credential| {
@ -36,21 +25,18 @@ async fn resolve_entra(
Sourced::new(value, source)
})
})
.map_err(Error::from)
.map_err(crate::ocr::Error::from)
}
fn validate_destination(
pub(super) fn validate_destination(
connection: &OcrConnection,
credential_source: InputSource,
) -> Result<(), OcrError> {
) -> Result<(), crate::ocr::Error> {
if connection.api_base.is_some()
&& connection.api_base_source == InputSource::Request
&& credential_source != InputSource::Request
{
return Err(Error::from(crate::AuthError::Configuration(
AuthConfigurationError::RequestAzureCredentialDestination,
))
.into());
return Err(litellm_auth::Error::RequestAzureCredentialDestination.into());
}
Ok(())
}

1
litellm-rust/crates/core/src/llms/azure_ai/ocr/document_intelligence/mod.rs Normal file
View file

@ -0,0 +1 @@
pub(crate) mod transformation;

1261
litellm-rust/crates/core/src/llms/azure_ai/ocr/document_intelligence/transformation.rs Normal file
View file

File diff suppressed because it is too large Load diff

4
litellm-rust/crates/core/src/llms/azure_ai/ocr/mod.rs Normal file
View file

@ -0,0 +1,4 @@
pub(crate) mod cohere_parse_transformation;
pub(crate) mod common_utils;
pub(crate) mod document_intelligence;
pub(crate) mod transformation;

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