iMerge upstream/litellm_internal_stagingitellm_internal_staging' into fix/vertex-ai-image-models-supports-reasoning

This commit is contained in:
bennedictbett 2026-08-21 08:50:25 +03:00
commit 5776980c05
874 changed files with 51549 additions and 11801 deletions

View file

@ -1025,7 +1025,7 @@ jobs:
name: Run tests
command: |
mkdir -p test-results
TEST_FILES=$(circleci tests glob "tests/agent_tests/**/test_*.py" | grep -v "^tests/agent_tests/local_only_agent_tests/")
TEST_FILES=$(circleci tests glob "tests/agent_tests/test_*.py")
echo "$TEST_FILES" | circleci tests run \
--verbose \
--command="tr ' ' '\\n' | awk '/\\.py/ {print; next} {sub(/\\.[A-Z][^.]*$/, \"\"); gsub(/\\./, \"/\"); print \$0 \".py\"}' | xargs uv run --no-sync python -m pytest \

View file

@ -4,6 +4,25 @@ description: >-
by a job nor listed here, so every entry below is a decision on the record.
test_paths:
- reason: >-
The caching suite in tests/local_testing, which runs nowhere. Every job that globs that
directory either deselects it (local_testing_part1 and part2 carry `-k "... and not caching
and not cache"`) or keeps only another keyword (langfuse, router, assistants), and no job
names these files the way redis_caching_unit_tests names test_dual_cache.py. Measured
2026-08-20 by collecting the directory under each job's own selector: 118 tests across
these eight files are selected by none of them. Listed so the gap is a decision rather
than an accident, and so the --slices guard has a baseline to ratchet down from. Revisit
when tests/local_testing is ported off CircleCI, where the keyless part of this suite
belongs in a real job
paths:
- tests/local_testing/test_cache_preset_key.py
- tests/local_testing/test_caching.py
- tests/local_testing/test_caching_handler.py
- tests/local_testing/test_disk_cache_unit_tests.py
- tests/local_testing/test_gcs_cache_unit_tests.py
- tests/local_testing/test_prompt_caching.py
- tests/local_testing/test_responses_stream_cache_keys.py
- tests/local_testing/test_unit_test_caching.py
- reason: >-
The end-to-end suite runs against a deployed proxy from its own in-cluster rig rather than
from a pull request; it needs a live gateway and provider credentials no PR job holds
@ -21,72 +40,24 @@ test_paths:
- tests/documentation_tests/test_requests_lib_usage.py
- tests/documentation_tests/test_standard_logging_payload.py
- reason: >-
Sibling files here are executed by name from the code-quality workflow; this one is referenced
by no job
Named like a test but shaped like a benchmark: it fetches live image URLs, times aiohttp
against httpx, prints the ratio, and asserts nothing, so pytest cannot collect it (its
functions take arguments, not fixtures) and running it beside its siblings in the
code-quality workflow would add a network dependency for a number nothing reads. Exempt
as a script rather than as an unresolved gap; revisit by deleting it once the aiohttp
choice it informed is settled
paths:
- tests/code_coverage_tests/test_aio_http_image_conversion.py
- reason: >-
A second mirror of the package tree living beside tests/test_litellm, which is the mirror the
repo convention names; only test_no_hardcoded_secrets.py is invoked, from the linting
workflow, and whether this directory should exist at all is unresolved
The last file of a second mirror that sat beside tests/test_litellm and ran nowhere. Its
other 33 files landed in the real mirror during August 2026, 30 as moves and 3 by merging
their bodies into the live file of the same name. This one cannot follow either route yet:
its live twin was rewritten from 1268 lines to 9434, and of the 19 tests here 5 have no
counterpart while 25 assertions fail against today's code, so what survives that rewrite
is a judgement about the endpoints, not a merge. Revisit by deciding which of the five
behaviours still hold
paths:
- tests/litellm/a2a_protocol/providers/pydantic_ai_agents/test_pydantic_ai_agent_headers.py
- tests/litellm/a2a_protocol/providers/pydantic_ai_agents/test_pydantic_ai_agent_transformation.py
- tests/litellm/integrations/helicone/test_helicone_gemini.py
- tests/litellm/litellm_core_utils/test_json_schema_validation.py
- tests/litellm/llms/anthropic/test_anthropic_reasoning_effort.py
- tests/litellm/llms/anthropic/test_anthropic_schema_filter.py
- tests/litellm/llms/azure/test_azure_embedding.py
- tests/litellm/llms/bedrock/embed/test_embedding.py
- tests/litellm/llms/bedrock/test_nova_imported_models.py
- tests/litellm/llms/deepseek/chat/test_deepseek_chat_transformation.py
- tests/litellm/llms/gradient_ai/chat/test_gradient_ai_chat_transformation.py
- tests/litellm/llms/oci/chat/test_oci_chat_transformation.py
- tests/litellm/llms/openai_like/test_abliteration_provider.py
- tests/litellm/llms/openai_like/test_assemblyai_provider.py
- tests/litellm/llms/openai_like/test_empiriolabs_provider.py
- tests/litellm/llms/vertex_ai/agent_engine/test_transformation.py
- tests/litellm/llms/vertex_ai/gemini/test_transformation.py
- tests/litellm/llms/vertex_ai/text_to_speech/test_transformation.py
- tests/litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py
- tests/litellm/proxy/agent_endpoints/test_agent_rbac.py
- tests/litellm/proxy/common_utils/test_rbac_utils.py
- tests/litellm/proxy/management_endpoints/test_common_utils.py
- tests/litellm/proxy/management_endpoints/test_cost_estimate_endpoint.py
- tests/litellm/proxy/test_claude_code_marketplace.py
- tests/litellm/proxy/test_init_litellm_callbacks.py
- tests/litellm/proxy/test_prisma_engine_watchdog.py
- tests/litellm/proxy/vector_store_endpoints/test_vector_store_rbac.py
- tests/litellm/test_bedrock_extended_beta_models.py
- tests/litellm/test_bedrock_nemotron_super.py
- tests/litellm/test_proxy_auth.py
- tests/litellm/test_router_retry_backoff_headers.py
- tests/litellm/test_sambanova_model_metadata.py
- tests/litellm/test_stream_chunk_builder_images.py
- reason: >-
Legacy proxy suite superseded by the proxy shards; no job invokes it and whether it still
describes supported behaviour is unresolved
paths:
- tests/old_proxy_tests/tests/test_anthropic_context_caching.py
- tests/old_proxy_tests/tests/test_anthropic_sdk.py
- tests/old_proxy_tests/tests/test_async.py
- tests/old_proxy_tests/tests/test_gemini_context_caching.py
- tests/old_proxy_tests/tests/test_langchain_embedding.py
- tests/old_proxy_tests/tests/test_langchain_request.py
- tests/old_proxy_tests/tests/test_llamaindex.py
- tests/old_proxy_tests/tests/test_mistral_sdk.py
- tests/old_proxy_tests/tests/test_openai_embedding.py
- tests/old_proxy_tests/tests/test_openai_exception_request.py
- tests/old_proxy_tests/tests/test_openai_request.py
- tests/old_proxy_tests/tests/test_openai_request_with_traceparent.py
- tests/old_proxy_tests/tests/test_openai_simple_embedding.py
- tests/old_proxy_tests/tests/test_openai_tts_request.py
- tests/old_proxy_tests/tests/test_pass_through_langfuse.py
- tests/old_proxy_tests/tests/test_q.py
- tests/old_proxy_tests/tests/test_simple_traceparent_openai.py
- tests/old_proxy_tests/tests/test_vertex_sdk_forward_headers.py
- tests/old_proxy_tests/tests/test_vtx_embedding.py
- tests/old_proxy_tests/tests/test_vtx_sdk_embedding.py
- reason: >-
No job invokes this suite and its files mix pure transformation tests with ones driving live
vendor vector stores, so assigning them needs a per-file decision
@ -116,6 +87,14 @@ test_paths:
- tests/load_tests/test_otel_load_test.py
- tests/load_tests/test_vertex_embeddings_load_test.py
- tests/load_tests/test_vertex_load_tests.py
- reason: >-
A local-only agent rig: test_a2a_completion_bridge.py needs a LangGraph server on
localhost:2024 and test_a2a.py drives a live A2A endpoint, so neither can run in a
pull request job. Until 2026-08-20 the CircleCI agent job hid them behind a grep -v
that this census could not see; the glob now excludes them structurally and this entry
is the decision on the record. Revisit when the A2A bridge gets a recorded-wire fixture
paths:
- tests/agent_tests/local_only_agent_tests
- reason: >-
Third-party integration tests that skip themselves without OCI configuration or sandbox
credentials, neither of which a pull request job holds
@ -124,14 +103,14 @@ test_paths:
- tests/integration/test_oci_integration.py
- tests/integration/test_oci_proxy_integration.py
- reason: >-
Two prompt-factory tests sitting at the top level of tests/ instead of under the
tests/test_litellm mirror the shards enumerate; they need moving rather than a shard entry
paths:
- tests/litellm_core_utils/test_anthropic_dedup_factory.py
- tests/litellm_core_utils/test_bedrock_converse_dedup_factory.py
- reason: >-
A unit test for the proxy-extras package that no job invokes, while the package's other tests
live under tests/proxy_migration_tests
A unit test for the proxy-extras package that no job invokes, while the package's other
tests live under tests/proxy_migration_tests. Measured 2026-08-20: 24 of its 28 tests pass
and the 4 in TestMigrationSQLIdempotency fail, because 13 migrations from 2026-03 onward use
bare CREATE TABLE, ADD COLUMN, CREATE INDEX and ADD CONSTRAINT rather than the guarded forms
this file requires. It also matches those keywords inside SQL comments, so two further
migrations are reported that are in fact fine. Wiring it up means deciding what to do about
the 13 first, and they cannot simply be edited: Prisma checksums an applied migration, so a
changed one breaks migrate deploy for existing installs
paths:
- tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py

View file

@ -1,10 +1,14 @@
from __future__ import annotations
import ast
import operator
import pathlib
import re
import sys
from collections.abc import Iterable, Mapping, Sequence
import warnings
from collections.abc import Callable, Iterable, Mapping, Sequence
from dataclasses import dataclass
from typing import Final
import yaml
@ -25,6 +29,15 @@ DOCKERFILE_TOKEN_RE = re.compile(r"[A-Za-z0-9_./-]*Dockerfile[A-Za-z0-9_.-]*")
COMMENT_RE = re.compile(r"^\s*#.*$", re.MULTILINE)
GLOB_CHARS = frozenset("*?")
# Trees whose jobs are sharded with no catch-all bucket, so every child that holds
# tests has to be named by some shard or it runs nowhere. A child listed here is
# itself decomposed one level deeper and is checked through its own entry.
SHARDED_ROOTS: tuple[str, ...] = (
"tests/proxy_unit_tests",
"tests/test_litellm",
"tests/test_litellm/proxy",
)
@dataclass(frozen=True, slots=True)
class AllowEntry:
@ -44,6 +57,14 @@ class Allowlist:
return any(relative_path == path for entry in self.dockerfiles for path in entry.paths)
@dataclass(frozen=True, slots=True)
class Section:
name: str
entries: tuple[AllowEntry, ...]
candidates: tuple[str, ...]
matches: Callable[[str, str], bool]
@dataclass(frozen=True, slots=True)
class Scalar:
key: str
@ -107,20 +128,34 @@ def _built_dockerfile_tokens(scalars: Iterable[Scalar]) -> frozenset[str]:
)
def _glob_to_regex(token: str) -> re.Pattern[str]:
parts = re.split(r"(\*\*/|\*\*|\*|\?)", token)
def _glob_to_regex(token: str, *, subtree: bool) -> re.Pattern[str]:
parts = re.split(r"(\*\*/|\*\*|\*|\?|\[[^\]]*\])", token)
translated = "".join(
{"**/": r"(?:.*/)?", "**": r".*", "*": r"[^/]*", "?": r"[^/]"}.get(part, re.escape(part)) for part in parts
{"**/": r"(?:.*/)?", "**": r".*", "*": r"[^/]*", "?": r"[^/]"}.get(part)
or (part if part.startswith("[") and part.endswith("]") else re.escape(part))
for part in parts
)
return re.compile(rf"{translated}(?:/.*)?$")
return re.compile(rf"{translated}(?:/.*)?$" if subtree else rf"{translated}$")
def _token_covers(token: str, relative_path: str) -> bool:
if GLOB_CHARS & set(token):
return _glob_to_regex(token).match(relative_path) is not None
return _glob_to_regex(token, subtree=True).match(relative_path) is not None
return relative_path == token or relative_path.startswith(f"{token}/")
def _token_names(token: str, relative_path: str) -> bool:
"""Whether the token names this path itself, rather than merely containing it.
A sharded tree has no catch-all bucket, so the ancestor token the census is happy
with (`tests/x` standing in for everything below it) is exactly what would let a
newly added child ride along without a shard.
"""
if GLOB_CHARS & set(token):
return _glob_to_regex(token, subtree=False).match(relative_path) is not None
return token == relative_path
def _test_files() -> tuple[str, ...]:
return tuple(
sorted(
@ -166,6 +201,174 @@ def _describe(paths: tuple[str, ...]) -> str:
return f"{len(paths)} test file(s) invoked by no job: {names}{suffix}"
GLOB_CALL_RE = re.compile(r'circleci tests glob "([^"]+)"')
KEYWORD_RE = re.compile(r"-k\s+\\?[\"']([^\"'\\]+)")
@dataclass(frozen=True, slots=True)
class Slice:
"""One job's selection: the files it globs, narrowed by its `-k` expression."""
job: str
globs: tuple[str, ...]
named: frozenset[str]
required: tuple[str, ...]
excluded: tuple[str, ...]
understood: bool
def claims(self, relative_path: str, inner_names: frozenset[str]) -> bool:
"""Whether this job runs any test in the file.
The question is deliberately per-file, not per-test. An excluded term is only
honoured when it appears in the path, because that is the case where it takes
the whole module with it; a term matching one function inside drops that test
and leaves the file claimed. Losing a whole file is the failure worth a gate,
and answering per-test would mean a baseline of test ids that churns on every
rename.
"""
if relative_path in self.named:
return True
if not any(_token_covers(glob, relative_path) for glob in self.globs):
return False
if not self.understood:
return True # a `-k` this parser cannot model is assumed to claim everything
if any(term.lower() in relative_path.lower() for term in self.excluded):
return False
return not self.required or any(
term.lower() in name.lower() for term in self.required for name in inner_names
)
def _strings(node: object) -> Iterable[str]:
if isinstance(node, str):
yield node
elif isinstance(node, dict):
for value in node.values():
yield from _strings(value)
elif isinstance(node, list):
for value in node:
yield from _strings(value)
def _keyword_terms(
expressions: Sequence[str], *, attributable: bool = True
) -> tuple[tuple[str, ...], tuple[str, ...], bool]:
"""A `-k` expression as (required, excluded, understood).
Only flat `and` chains of bare terms are modelled. Anything with `or`, parentheses
or negation of a group is left unmodelled, and its job is then treated as claiming
every file it globs, so an unparsed selector can never raise a false alarm.
`attributable` is False when a job runs several pytest commands, since a selector
read out of the job's text cannot then be tied to the glob it belongs to, and
pairing one command's exclusion with another's glob would invent a gap.
"""
terms: Final = tuple(part.strip() for expression in expressions for part in expression.split(" and "))
if not attributable and terms:
return (), (), False
if any(("or " in term) or ("(" in term) or (term.startswith("not ") and " " in term[4:]) for term in terms):
return (), (), False
return (
tuple(term for term in terms if term and not term.startswith("not ")),
tuple(term[4:].strip() for term in terms if term.startswith("not ")),
True,
)
def _slices() -> tuple[Slice, ...]:
if not CIRCLECI_CONFIG.exists():
return ()
jobs: Final = yaml.safe_load(CIRCLECI_CONFIG.read_text()).get("jobs", {})
return tuple(
Slice(job=job, globs=globs, named=named, required=required, excluded=excluded, understood=understood)
for job, body in jobs.items()
for text in ("\n".join(_strings(body)),)
if "pytest" in text
for globs in (tuple(GLOB_CALL_RE.findall(text)),)
for named in (frozenset(TEST_TOKEN_RE.findall(text)) & frozenset(_test_files()),)
for required, excluded, understood in (
_keyword_terms(tuple(KEYWORD_RE.findall(text)), attributable=len(globs) < 2),
)
if globs or named
)
def _matchable_names(relative_path: str) -> frozenset[str]:
"""Every name a `-k` term can match for this file: its path, plus the names inside it.
pytest matches a keyword against an item's own name and each of its parents', so a
positive term hits a file when it appears in the path or in a class or function name.
"""
try:
with warnings.catch_warnings():
warnings.simplefilter("ignore") # test files carry stray escapes; their names still parse
tree: Final = ast.parse((REPO_ROOT / relative_path).read_text())
except (OSError, SyntaxError):
return frozenset({relative_path})
return frozenset({relative_path}) | frozenset(
node.name
for node in ast.walk(tree)
if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef))
)
def _deselected_everywhere(allowlist: Allowlist) -> tuple[Finding, ...]:
slices: Final = _slices()
globbed: Final = tuple(
path
for path in _test_files()
if any(_token_covers(glob, path) for slice_ in slices for glob in slice_.globs)
)
return tuple(
Finding(
subject=path,
detail="globbed by a job, then deselected by every one of their -k expressions",
)
for path in globbed
if not allowlist.covers_test(path)
and not any(slice_.claims(path, _matchable_names(path)) for slice_ in slices)
)
def _holds_tests(directory: pathlib.Path) -> bool:
return any(directory.rglob("test_*.py"))
def _shard_children(root: str, repo_root: pathlib.Path = REPO_ROOT) -> tuple[str, ...]:
"""Children of a sharded root that carry tests, so each one needs its own shard.
A directory earns an entry by containing a test file rather than by being named
`test_*`, which is what keeps fixture directories (`test_configs`, `expected_*`)
out without a hand-maintained list of exceptions.
"""
return tuple(
sorted(
child.relative_to(repo_root).as_posix()
for child in (repo_root / root).iterdir()
if not child.name.startswith(".")
and (
_holds_tests(child)
if child.is_dir()
else child.name.startswith("test_") and child.suffix == ".py"
)
)
)
def _unassigned_shard_children(
tokens: frozenset[str],
roots: tuple[str, ...] = SHARDED_ROOTS,
repo_root: pathlib.Path = REPO_ROOT,
) -> tuple[Finding, ...]:
return tuple(
Finding(subject=child, detail=f"holds tests but no shard of {root} names it")
for root in roots
if (repo_root / root).is_dir()
for child in _shard_children(root, repo_root)
if child not in roots and not any(_token_names(token, child) for token in tokens)
)
def _uncovered_dockerfiles(allowlist: Allowlist, tokens: frozenset[str]) -> tuple[Finding, ...]:
return tuple(
Finding(subject=relative_path, detail="built by no job")
@ -174,6 +377,25 @@ def _uncovered_dockerfiles(allowlist: Allowlist, tokens: frozenset[str]) -> tupl
)
def _stale_allowlist_paths(
allowlist: Allowlist,
*,
test_files: tuple[str, ...],
dockerfiles: tuple[str, ...],
) -> tuple[Finding, ...]:
sections: Final[tuple[Section, ...]] = (
Section("test_paths", allowlist.test_paths, test_files, _token_covers),
Section("dockerfiles", allowlist.dockerfiles, dockerfiles, operator.eq),
)
return tuple(
Finding(subject=path, detail=f"listed under '{section.name}' but matches no file the census looks at")
for section in sections
for entry in section.entries
for path in entry.paths
if not any(section.matches(path, candidate) for candidate in section.candidates)
)
def _parse_entry(item: object, section: str) -> AllowEntry:
if not isinstance(item, dict):
raise SystemExit(f"{ALLOWLIST_FILE.name}: '{section}' entries must be mappings")
@ -229,13 +451,56 @@ def _report(title: str, findings: tuple[Finding, ...], remedy: str) -> None:
_write("")
def _check_slices() -> int:
findings: Final = _deselected_everywhere(_load_allowlist())
if findings:
_report(
"test files a -k expression removes from every job that globs them",
findings,
"Give each one a job whose -k keeps it, or list it in "
".github/ci-coverage-allowlist.yml with the reason it may stay unrun.",
)
return 1
_write(f"OK: no test file is globbed by a job and then deselected by every -k across {len(_slices())} slices.")
return 0
def _check_shards() -> int:
findings = _unassigned_shard_children(_invoked_test_tokens(_all_scalars()))
if findings:
_report(
"test directories and files that no shard claims",
findings,
"Add each to the shard it belongs to. A directory that is itself split across "
"several shards belongs in SHARDED_ROOTS instead, so its own children get checked.",
)
return 1
counted = sum(len(_shard_children(root)) for root in SHARDED_ROOTS if (REPO_ROOT / root).is_dir())
_write(f"OK: all {counted} test children across {len(SHARDED_ROOTS)} sharded trees are assigned to a shard.")
return 0
def main() -> int:
if "--shards" in sys.argv[1:]:
return _check_shards()
if "--slices" in sys.argv[1:]:
return _check_slices()
allowlist = _load_allowlist()
scalars = _all_scalars()
test_findings = _uncovered_tests(allowlist, _invoked_test_tokens(scalars))
dockerfile_findings = _uncovered_dockerfiles(allowlist, _built_dockerfile_tokens(scalars))
stale_findings = _stale_allowlist_paths(allowlist, test_files=_test_files(), dockerfiles=_dockerfiles())
if stale_findings:
_report(
"allowlist entries that exempt nothing",
stale_findings,
"Delete each from .github/ci-coverage-allowlist.yml; the file it named is gone or was renamed.",
)
if test_findings:
_report(
"test files that no CI job invokes",
@ -248,7 +513,7 @@ def main() -> int:
dockerfile_findings,
"Build each in a workflow, or list it in .github/ci-coverage-allowlist.yml with a reason.",
)
if test_findings or dockerfile_findings:
if stale_findings or test_findings or dockerfile_findings:
return 1
_write(

View file

@ -0,0 +1,149 @@
#!/usr/bin/env python3
"""Three invariants about what lives in .github/workflows/ and what its names mean.
`.github/workflows/` is a directory GitHub reads, not a place to keep things. Every
file at its top level is parsed as a workflow, so a script or a data file parked there
is either an invalid workflow or an orphan nobody can find. A subdirectory is not read
at all, so helper files may live in one. GitHub accepts both `.yml` and `.yaml`, and
this repo spells them `.yml`, which is a naming rule rather than a validity one and is
reported separately. And the `_` prefix is the repo's only signal that a workflow is a
reusable building block rather than something that runs on its own, which is worth
nothing unless it is true both ways.
WF001 a top-level file in .github/workflows/ that is not a workflow at all
WF002 a workflow whose only trigger is `workflow_call` but is not `_`-prefixed
WF003 a `_`-prefixed workflow that no other workflow can call
WF004 a real workflow spelled `.yaml` where this directory spells them `.yml`
A workflow with `workflow_call` alongside a human trigger is deliberately dual-mode
and belongs under its plain name, so only the call-only ones are held to WF002.
Usage
-----
python assert_workflow_dir_hygiene.py
Exit code 1 if any violation is found.
"""
from __future__ import annotations
import pathlib
import sys
from dataclasses import dataclass
from typing import Final
import yaml
REPO_ROOT: Final = pathlib.Path(__file__).resolve().parents[2]
WORKFLOW_DIR: Final = REPO_ROOT / ".github" / "workflows"
SCRIPT_HOME: Final = ".github/scripts/"
REUSABLE_PREFIX: Final = "_"
CALL_TRIGGER: Final = "workflow_call"
CANONICAL_SUFFIX: Final = ".yml"
WORKFLOW_SUFFIXES: Final = frozenset((CANONICAL_SUFFIX, ".yaml"))
@dataclass(frozen=True, slots=True)
class Finding:
subject: str
code: str
detail: str
def render(self) -> str:
return f" - {self.subject}: {self.code} {self.detail}"
def _triggers(document: object) -> frozenset[str]:
if not isinstance(document, dict):
return frozenset()
raw: Final = document.get("on", document.get(True))
if isinstance(raw, str):
return frozenset({raw})
if isinstance(raw, dict):
return frozenset(str(key) for key in raw)
if isinstance(raw, list):
return frozenset(str(item) for item in raw)
return frozenset()
def _workflows(directory: pathlib.Path) -> tuple[pathlib.Path, ...]:
return tuple(
path
for path in sorted(directory.iterdir())
if path.is_file() and path.suffix in WORKFLOW_SUFFIXES
)
def _strays(directory: pathlib.Path) -> tuple[Finding, ...]:
return tuple(
Finding(
path.name,
"WF001",
f"is not a workflow, and GitHub parses every top-level file here as one; "
f"move it to {SCRIPT_HOME} or into a subdirectory, which GitHub does not read",
)
for path in sorted(directory.iterdir())
if path.is_file() and path.suffix not in WORKFLOW_SUFFIXES
)
def _misspelled(directory: pathlib.Path) -> tuple[Finding, ...]:
return tuple(
Finding(
path.name,
"WF004",
f"is a real workflow and GitHub reads it, but this directory spells them "
f"{CANONICAL_SUFFIX}; rename it to {path.stem}{CANONICAL_SUFFIX}",
)
for path in _workflows(directory)
if path.suffix != CANONICAL_SUFFIX
)
def _misnamed(directory: pathlib.Path) -> tuple[Finding, ...]:
return tuple(
finding
for path in _workflows(directory)
for finding in _naming_findings(path, _triggers(yaml.safe_load(path.read_text(encoding="utf-8"))))
)
def _naming_findings(path: pathlib.Path, triggers: frozenset[str]) -> tuple[Finding, ...]:
underscored: Final = path.name.startswith(REUSABLE_PREFIX)
if triggers == frozenset({CALL_TRIGGER}) and not underscored:
return (
Finding(
path.name,
"WF002",
f"is only callable by another workflow, so name it {REUSABLE_PREFIX}{path.name}",
),
)
if underscored and CALL_TRIGGER not in triggers:
return (
Finding(
path.name,
"WF003",
f"is named as a reusable workflow but has no {CALL_TRIGGER} trigger; "
"add one or drop the prefix",
),
)
return ()
def main() -> int:
findings: Final = _strays(WORKFLOW_DIR) + _misspelled(WORKFLOW_DIR) + _misnamed(WORKFLOW_DIR)
if not findings:
total: Final = len(_workflows(WORKFLOW_DIR))
sys.stdout.write(
f"OK: {total} workflows, every file in .github/workflows/ is one, and the "
f"{REUSABLE_PREFIX} prefix means callable in both directions.\n"
)
return 0
sys.stdout.write("ERROR: .github/workflows/ holds files that break its own conventions\n")
for finding in findings:
sys.stdout.write(f"{finding.render()}\n")
return 1
if __name__ == "__main__":
sys.exit(main())

View file

@ -1,557 +0,0 @@
#!/usr/bin/env python3
"""One-shot 7-day heads-up sweep for the Agent Shin rollout.
Posts a friendly "the OSS triage bot kicks in next Monday" comment on every
open external PR/issue that currently *would* fail the new rubric i.e.,
every PR/issue Agent Shin would close once the rollout completes. The point
is to give contributors a full week to fix their description before the bot
ever takes a destructive action, so nobody is surprised by an auto-close.
The script is designed to run **exactly once** at rollout, fired by a manual
``workflow_dispatch`` (``dry_run=false``) on the heads-up workflow. Re-runs
are safe: every comment is stamped with the hidden ``HEADS_UP_MARKER`` and
PRs/issues that already carry the marker are skipped.
Dry-run vs. real run
--------------------
Defaults to dry-run. Passing ``--close`` flips into real mode. Every GitHub
mutation goes through ``_agent_shin_actions``, which has a one-line
``if dry_run: log else: do_it`` per call, so the only difference between a
dry-run preview and the real run is the call site that actually hits the
GitHub API.
Local preview::
python3 .github/scripts/triage_rollout_heads_up.py --repo BerriAI/litellm
Real run (the manual rollout dispatch uses this)::
python3 .github/scripts/triage_rollout_heads_up.py --repo BerriAI/litellm --close
"""
from __future__ import annotations
import argparse
import datetime as dt
import json
import os
import sys
from pathlib import Path
from typing import Any
# Make the sibling triage_with_llm + _agent_shin_actions importable when this
# script is invoked directly (the GitHub workflow does `python3 .github/scripts/...`).
_SCRIPTS_DIR = Path(__file__).resolve().parent
if str(_SCRIPTS_DIR) not in sys.path:
sys.path.insert(0, str(_SCRIPTS_DIR))
from _agent_shin_actions import maybe_post_comment # noqa: E402
from agent_shin_shared import ( # noqa: E402
AGENT_SHIN_DEFAULT_BOT_LOGIN,
ALLOWLIST_LOGINS,
list_open_items,
)
from triage_with_llm import ( # noqa: E402
DEFAULT_MODEL,
call_llm_judge,
fetch_issue,
fetch_pr,
gh,
is_internal_contributor,
review_gate,
triage,
)
# Hidden marker so re-runs skip PRs/issues we've already notified. Distinct from
# the within-grace / ready / regressed markers so it can't be confused with the
# steady-state lifecycle comments.
HEADS_UP_MARKER = "<!-- agent-shin:rollout-heads-up -->"
# Placeholder until the litellm-docs PR ships. The rollout blog post explains
# the new rubric, the 7-day grace, and how to recover after an auto-close.
# TODO(docs): replace with the canonical URL once the litellm-docs PR merges.
ROLLOUT_BLOG_URL = "https://docs.litellm.ai/docs/agent_shin_triage_rollout"
# Default cutoff is one week from "now". Computed at runtime so the wording
# stays correct even if the rollout is merged later than planned. The user can
# override with --close-on YYYY-MM-DD when running the script manually.
DEFAULT_GRACE_DAYS = 7
# The daily auto-close sweeps (close_low_quality_prs.yml at 09:00 UTC and
# review_gate.yml at 09:30 UTC) are what actually close a still-failing item,
# so the deadline we promise contributors has to name that wall-clock moment.
ACTIVATION_TIME_UTC = "09:00 UTC"
def _format_cutoff(cutoff: dt.date) -> str:
"""Human-readable, timezone-explicit cutoff, e.g. ``Monday, June 1, 2026
(09:00 UTC)`` the moment a still-failing PR/issue gets closed."""
return (
f"{cutoff.strftime('%A, %B')} {cutoff.day}, {cutoff.year} "
f"({ACTIVATION_TIME_UTC})"
)
def _rubric_section_pr() -> str:
return (
"**Going forward, every external PR needs ONE of:**\n"
"\n"
"- A linked GitHub issue using a closing keyword: "
"`Fixes #1234`, `Closes #1234`, or `Resolves #1234`, OR\n"
"- All three of: a clear **problem description**, **expected vs. "
"actual behavior**, and **end-to-end QA proof** (at least one of a "
"short screen recording / video, before/after screenshots, or the "
"exact commands you ran with their real output; mocked or stubbed "
"runs don't count).\n"
"\n"
"PRs also need a **Greptile confidence score of 4/5 or higher** before "
"the bot will tag them `ready for review`. You can `@greptileai` to "
"request a fresh review at any time, including after the PR is closed."
)
def _rubric_section_issue() -> str:
return (
"**Going forward, every external issue needs:**\n"
"\n"
"- For **bug reports**: end-to-end evidence of the bug (at least one "
"of a screen recording / video, a screenshot, or the exact commands "
"you ran with their real output / traceback) plus expected vs. actual "
"behavior. Written steps with no run output don't count, and mocked "
"or stubbed runs don't count.\n"
"- For **feature requests**: a clear description of the proposed "
"feature plus a use case + concrete example (config, API call, UI "
"flow, or scenario showing what's blocked today)."
)
def _description_only_note(kind: str) -> str:
noun = "PR" if kind == "pr" else "issue"
return (
f"⚠️ **The requirements must live in the {noun} *description*, not in "
"comments.** Some PRs/issues collect 100+ comments from humans and "
"bots; reading the entire thread on every triage run would balloon "
"GitHub API usage (we'd start getting 429'd) and blow out the LLM "
"judge's context. The bot only reads the description, so anything "
"you add as a comment will be invisible to it."
)
def _missing_section(verdict: dict, greptile_score: int | None) -> str:
"""Bullet list of what's currently missing on this PR/issue.
Combines the LLM judge's `missing` list (rubric items) with a Greptile
shortfall (for PRs) so the contributor sees one list of things to fix.
"""
missing = list(verdict.get("missing") or [])
if greptile_score is not None and greptile_score < 4:
missing.insert(
0,
f"Greptile's most recent review scored this PR {greptile_score}/5 "
"(below the 4/5 bar Agent Shin will require).",
)
if not missing:
return (
"_The bot couldn't articulate a specific missing piece; see the "
"rubric link above and double-check the description includes all "
"of it before the rollout._"
)
bullets = "\n".join(f"- {m}" for m in missing)
return f"**What this one is currently missing:**\n\n{bullets}"
def _recovery_section(kind: str) -> str:
if kind == "pr":
return (
"**If the bot closes this PR after the rollout:** update the "
"description with the missing pieces, then either open a fresh "
"PR or comment `@agent-shin reconsider` on the closed PR. If "
"Greptile re-scores you at 4/5 or higher I'll reopen and tag "
"the PR `ready for review`. (`@greptileai` works on closed PRs "
"too; a fresh review is one of the signals that lifts you back "
"into the queue.) This is **not** us losing interest in your "
"change; far from it. We just need open PRs to be a list of "
"things a maintainer can act on, so we can get to yours faster."
)
return (
"**If the bot closes this issue after the rollout:** edit the issue "
"description to add the missing pieces, then comment `@agent-shin "
"reconsider` on the closed issue. I'll re-evaluate and, if the rubric "
"is met, reopen it. (GitHub doesn't let external authors reopen an "
"issue a maintainer or bot closed, so the comment is the reliable "
"path.) This is **not** us saying the bug isn't real or the request "
"isn't useful; it's so the remaining open issues are a list of things "
"a maintainer can act on."
)
def format_heads_up_comment(
*, kind: str, verdict: dict, greptile_score: int | None, cutoff: dt.date
) -> str:
"""Compose the friendly 7-day heads-up comment posted on a failing PR/issue."""
noun = "PR" if kind == "pr" else "issue"
rubric = _rubric_section_pr() if kind == "pr" else _rubric_section_issue()
cutoff_str = _format_cutoff(cutoff)
explanation = (verdict.get("explanation") or "").strip()
explanation_block = (
f"> _(The judge's note for this one: {explanation})_\n\n" if explanation else ""
)
return (
"🚅 **Heads-up: we're turning on the OSS triage bot in "
f"{DEFAULT_GRACE_DAYS} days, on {cutoff_str}.**\n"
"\n"
"We're rolling out **Agent Shin**, an LLM-as-judge triage bot for "
f"external {noun}s. Once it's live, the bot reads each open "
f"{noun}'s description, scores it against a small rubric, and "
f"auto-closes any {noun} that's missing the basics, with a single "
f"comment explaining what's missing and how to recover. Full "
f"context: [Agent Shin rollout blog post]({ROLLOUT_BLOG_URL}).\n"
"\n"
f"{rubric}\n"
"\n"
f"{_description_only_note(kind)}\n"
"\n"
f"{_missing_section(verdict, greptile_score)}\n"
"\n"
f"{explanation_block}"
"**Timeline (you have a week):**\n"
"\n"
f"- We turn the bot on in {DEFAULT_GRACE_DAYS} days, on "
f"**{cutoff_str}**. You have until then to update this {noun}'s "
"description with the missing pieces above.\n"
f"- If this {noun} still fails the rubric at **{cutoff_str}**, "
"we'll close it.\n"
f"- From then on the bot runs daily, and every {noun} that fails "
"the rubric gets a **2-hour lifetime**: one warning comment, then "
"auto-close 2 hours later.\n"
"\n"
f"{_recovery_section(kind)}\n"
"\n"
f"{HEADS_UP_MARKER}"
)
def _list_open_numbers(repo: str, kind: str) -> list[int]:
"""Return every open PR or issue number in ``repo``.
Delegates to ``list_open_items`` so the full backlog is fetched (no cap)
and the `gh {pr,issue} list` invocation stays in one shared place. ``gh
issue list`` would include PRs, but ``list_open_items`` uses the dedicated
command per kind, so the two never mix.
"""
return [
item["number"] for item in list_open_items(kind, repo=repo, fields="number")
]
def _has_heads_up_marker(item: dict) -> bool:
"""Cheap fast-path: check the PR/issue body itself for the marker.
The marker is appended to the *comment* we post, not the body, so this
will only fire if the body literally contains the marker text. We still
do the comment-marker check separately below; this body check just lets
us short-circuit for PRs/issues that quote the marker for any reason.
"""
body = item.get("body") or ""
return HEADS_UP_MARKER in body
def _comments_have_marker(repo: str, number: int) -> bool:
"""True if the bot already posted a comment carrying the marker.
Used for idempotency: a re-run skips items the previous run notified.
Filters by author (matching the sibling marker-checks in
``triage_with_llm._has_marker`` and
``agent_shin_shared.seconds_since_latest_marker_comment``) so a
contributor who quotes the heads-up via GitHub's "Quote reply" — which
preserves HTML comments in the raw markdown can't trick the
idempotency check into silently skipping a real heads-up.
Comments live on the unified issues endpoint regardless of whether the
item is a PR or an issue, so no ``kind`` argument is required here.
"""
expected_login = (
os.environ.get("AGENT_SHIN_BOT_LOGIN") or AGENT_SHIN_DEFAULT_BOT_LOGIN
).lower()
raw = gh(
"api",
"--paginate",
f"repos/{repo}/issues/{number}/comments?per_page=100",
)
for line in raw.splitlines():
line = line.strip()
if not line:
continue
try:
payload = json.loads(line)
except json.JSONDecodeError:
continue
comments = payload if isinstance(payload, list) else [payload]
for comment in comments:
author = ((comment.get("user") or {}).get("login") or "").lower()
if author != expected_login:
continue
if HEADS_UP_MARKER in (comment.get("body") or ""):
return True
return False
def _evaluate_pr(*, repo: str, number: int, model: str, judge: Any = None) -> dict:
"""Run the future PR rubric (review_gate) in dry-run and return the result."""
return review_gate(
repo=repo,
number=number,
close=False, # we only want the verdict, never act here
model=model,
judge=judge,
)
def _evaluate_issue(*, repo: str, number: int, model: str, judge: Any = None) -> dict:
"""Run the future issue rubric (triage kind='issue') in dry-run."""
return triage(
repo=repo,
kind="issue",
number=number,
close=False,
model=model,
judge=judge,
)
def _would_be_closed(kind: str, result: dict) -> bool:
"""True if the future triage would auto-close this PR/issue based on the
rubric (regardless of grace-period gating).
For PRs we trust ``review_gate``'s ``passing`` field — it combines the LLM
verdict and the Greptile score. For issues we read the LLM verdict
directly. Both fields are ``None``/missing on skip paths
(skip-internal-author, skip-llm-error, etc.) where the future bot would
NOT close the item those return False.
"""
if kind == "pr":
passing = result.get("passing")
if passing is None:
return False # skipped — nothing for the heads-up to warn about
return passing is False
verdict = result.get("verdict") or {}
return (verdict.get("verdict") or "").lower() == "fail"
def _process_one(
*,
repo: str,
kind: str,
number: int,
model: str,
cutoff: dt.date,
dry_run: bool,
judge: Any = None,
skip_marker_check: bool = False,
allowlist: frozenset[str] = ALLOWLIST_LOGINS,
) -> dict:
"""Evaluate one PR/issue and post a heads-up if it would be auto-closed.
Returns a per-item dict for the summary table.
"""
base = {"kind": kind, "number": number}
fetcher = fetch_pr if kind == "pr" else fetch_issue
item = fetcher(repo, number)
if (item.get("state") or "") != "open":
return {**base, "action": "skip-not-open"}
if allowlist:
login = (item.get("user") or {}).get("login") or ""
if login.lower() not in allowlist:
return {**base, "action": "skip-not-allowlisted"}
elif is_internal_contributor(item):
return {**base, "action": "skip-internal-author"}
if not skip_marker_check and _has_heads_up_marker(item):
return {**base, "action": "skip-already-marked-in-body"}
if not skip_marker_check and _comments_have_marker(repo, number):
return {**base, "action": "skip-already-notified"}
if kind == "pr":
result = _evaluate_pr(repo=repo, number=number, model=model, judge=judge)
else:
result = _evaluate_issue(repo=repo, number=number, model=model, judge=judge)
if not _would_be_closed(kind, result):
return {**base, "action": "skip-passing", "evaluator": result.get("action")}
verdict = result.get("verdict") or {}
greptile_score = result.get("greptile_score") if kind == "pr" else None
comment = format_heads_up_comment(
kind=kind, verdict=verdict, greptile_score=greptile_score, cutoff=cutoff
)
maybe_post_comment(repo, number, comment, dry_run=dry_run)
return {
**base,
"action": "heads-up-posted" if not dry_run else "would-post-heads-up",
"verdict": (verdict.get("verdict") or "").lower(),
"greptile_score": greptile_score,
}
def _print_summary(results: list[dict]) -> None:
"""Tally per-action counts so a dry-run preview tells you at a glance how
many comments the real run would post."""
counts: dict[str, int] = {}
for r in results:
counts[r["action"]] = counts.get(r["action"], 0) + 1
print("\n=== rollout heads-up summary ===")
for action in sorted(counts):
print(f" {action:35s} {counts[action]}")
print(f" total {len(results)}")
def run(
*,
repo: str,
close: bool,
cutoff: dt.date,
model: str,
kinds: tuple[str, ...] = ("pr", "issue"),
judge: Any = None,
only_numbers: dict[str, list[int]] | None = None,
skip_marker_check: bool = False,
) -> list[dict]:
"""Sweep ``repo`` and post heads-up comments. Returns the per-item results."""
dry_run = not close
if dry_run:
print(
f"[DRY RUN] sweeping {repo}; --close not passed, no comments will be posted."
)
else:
print(f"[REAL RUN] sweeping {repo}; comments WILL be posted.")
print(f"Cutoff date in comment body: {cutoff.isoformat()}")
results: list[dict] = []
for kind in kinds:
if only_numbers and kind in only_numbers:
numbers = list(only_numbers[kind])
else:
numbers = _list_open_numbers(repo, kind)
print(f"\n--- {kind}s: {len(numbers)} open ---")
for n in numbers:
try:
result = _process_one(
repo=repo,
kind=kind,
number=n,
model=model,
cutoff=cutoff,
dry_run=dry_run,
judge=judge,
skip_marker_check=skip_marker_check,
)
except (
Exception
) as exc: # noqa: BLE001 - per-item errors don't abort the sweep
result = {
"kind": kind,
"number": n,
"action": "error",
"error": str(exc),
}
print(f"!! {kind}#{n}: {exc}", file=sys.stderr)
print(f" {kind}#{n}: {result['action']}")
results.append(result)
_print_summary(results)
return results
def main() -> int:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--repo", required=True, help="owner/repo")
parser.add_argument(
"--close",
action="store_true",
help=(
"Actually post comments. Without this flag the script is in "
"dry-run mode and only logs what it would do."
),
)
parser.add_argument(
"--close-on",
type=dt.date.fromisoformat,
default=None,
help=(
"Cutoff date shown in the heads-up comment as the rollout date "
f"(default: today + {DEFAULT_GRACE_DAYS} days)."
),
)
parser.add_argument(
"--model",
default=os.environ.get("TRIAGE_MODEL") or DEFAULT_MODEL,
help=f"Model for the rubric LLM judge (default: {DEFAULT_MODEL}).",
)
parser.add_argument(
"--kind",
choices=("pr", "issue", "both"),
default="both",
help="Restrict the sweep to PRs or issues only (default: both).",
)
parser.add_argument(
"--only-pr",
type=int,
action="append",
default=[],
help="Limit the PR sweep to these PR numbers (repeat for several).",
)
parser.add_argument(
"--only-issue",
type=int,
action="append",
default=[],
help="Limit the issue sweep to these issue numbers (repeat for several).",
)
parser.add_argument(
"--ignore-existing-marker",
action="store_true",
help=(
"Re-post on PRs/issues that already carry the heads-up marker. "
"Useful for testing the comment wording on a known PR."
),
)
args = parser.parse_args()
cutoff = args.close_on or (
dt.datetime.now(dt.timezone.utc).date() + dt.timedelta(days=DEFAULT_GRACE_DAYS)
)
kinds: tuple[str, ...]
if args.kind == "pr":
kinds = ("pr",)
elif args.kind == "issue":
kinds = ("issue",)
else:
kinds = ("pr", "issue")
only: dict[str, list[int]] = {}
if args.only_pr:
only["pr"] = args.only_pr
if args.only_issue:
only["issue"] = args.only_issue
# The script must NOT hit the LLM in dry-run if no key is set — we still
# want a useful preview that says "skip-no-llm-key" for items that would
# have been judged. Production runs require OPENAI_API_KEY.
if args.close and not os.environ.get("OPENAI_API_KEY"):
parser.error("OPENAI_API_KEY must be set for --close (real-run) mode.")
run(
repo=args.repo,
close=args.close,
cutoff=cutoff,
model=args.model,
kinds=kinds,
only_numbers=only or None,
skip_marker_check=args.ignore_existing_marker,
)
return 0
if __name__ == "__main__":
sys.exit(main())

View file

@ -129,6 +129,13 @@ jobs:
WORKERS: ${{ inputs.workers }}
RERUNS: ${{ inputs.reruns }}
DIST: ${{ inputs.dist }}
# coverage.py's sys.monitoring backend (PEP 669), the cheapest core it has.
# It is only the default from Python 3.14, and these shards run 3.12, so it
# has to be asked for. Coverage refuses it when branch measurement is on
# (`branch_right_left` needs > 3.14.0a5) and falls back to the slow core with
# a `no-sysmon` warning, so turning on `branch = true` here means giving this
# back until the runners move to 3.14.
COVERAGE_CORE: sysmon
run: |
if [ "${WORKERS}" = "0" ]; then
uv run --no-sync pytest ${TEST_PATH:?} \

View file

@ -23,7 +23,7 @@ jobs:
version: "0.10.9"
- name: Update JSON Data
run: |
uv run --frozen --with 'aiohttp==3.13.3' python ".github/workflows/auto_update_price_and_context_window_file.py"
uv run --frozen --with 'aiohttp==3.13.3' python ".github/scripts/auto_update_price_and_context_window_file.py"
- name: Regenerate JSON Schema
run: |
uv run --frozen python ci_cd/generate_model_prices_schema.py

View file

@ -40,3 +40,12 @@ jobs:
run: |
python -m pip install "pyyaml==6.0.3"
python .github/scripts/assert_ci_coverage.py
# The census asks whether a job names a file; this asks whether that job's -k
# then throws it back out. A file both globbed and deselected everywhere runs
# nowhere while counting as covered, which is how the caching suite went unrun.
- name: Assert no -k expression deselects a file from every job that globs it
run: python .github/scripts/assert_ci_coverage.py --slices
- name: Assert .github/workflows/ holds only workflows, correctly named
run: python .github/scripts/assert_workflow_dir_hygiene.py

View file

@ -122,6 +122,11 @@ jobs:
uv run --no-sync ruff check .
cd ..
- name: Run Ruff linting (test tree)
if: steps.changes.outputs.decision != 'skip'
run: |
uv run --no-sync ruff check --config ruff-tests.toml tests
- name: Check strict-rule budget (delta vs base)
if: steps.changes.outputs.decision != 'skip'
run: |
@ -132,6 +137,11 @@ jobs:
run: |
uv run --no-sync python scripts/type_discipline_gate.py --base "$GATE_BASE_SHA"
- name: Check test-quality budget (zero-assert / mock-echo tests, sys.path.insert, raw env writes, litellm global mutation, credential-gated skips, conftest snapshot inventory, delta vs base)
if: steps.changes.outputs.decision != 'skip'
run: |
uv run --no-sync python scripts/test_quality_gate.py --base "$GATE_BASE_SHA"
- name: Print OpenAI version
if: steps.changes.outputs.decision != 'skip'
run: |
@ -223,7 +233,7 @@ jobs:
- name: Run secret scan test
run: |
uv run --no-project --with 'pytest==9.0.2' pytest tests/litellm/test_no_hardcoded_secrets.py -v
uv run --no-project --with 'pytest==9.0.2' pytest tests/code_coverage_tests/test_no_hardcoded_secrets.py -v
- name: Run ggshield secret scan
env:

View file

@ -1,31 +0,0 @@
name: "Unit Tests: Core Utilities"
on:
pull_request:
branches:
- main
- litellm_internal_staging
- litellm_oss_staging
- "litellm_**"
push:
branches:
- main
- litellm_internal_staging
permissions:
contents: read
id-token: write
pull-requests: write
concurrency:
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}
cancel-in-progress: ${{ github.event_name == 'pull_request' }}
jobs:
core-utils:
uses: ./.github/workflows/_test-unit-base.yml
with:
test-path: "tests/test_litellm/litellm_core_utils"
workers: 2
reruns: 1
artifact-name: core-utils

View file

@ -1,35 +0,0 @@
name: "Unit Tests: Enterprise, Google GenAI & Routing"
on:
pull_request:
branches:
- main
- litellm_internal_staging
- litellm_oss_staging
- "litellm_**"
push:
branches:
- main
- litellm_internal_staging
permissions:
contents: read
id-token: write
pull-requests: write
concurrency:
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}
cancel-in-progress: ${{ github.event_name == 'pull_request' }}
jobs:
enterprise-routing:
uses: ./.github/workflows/_test-unit-base.yml
with:
test-path: >-
tests/test_litellm/enterprise
tests/test_litellm/google_genai
tests/test_litellm/router_utils
tests/test_litellm/router_strategy
workers: 2
reruns: 2
artifact-name: enterprise-routing

View file

@ -1,31 +0,0 @@
name: "Unit Tests: Integrations (Callbacks & Logging)"
on:
pull_request:
branches:
- main
- litellm_internal_staging
- litellm_oss_staging
- "litellm_**"
push:
branches:
- main
- litellm_internal_staging
permissions:
contents: read
id-token: write
pull-requests: write
concurrency:
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}
cancel-in-progress: ${{ github.event_name == 'pull_request' }}
jobs:
integrations:
uses: ./.github/workflows/_test-unit-base.yml
with:
test-path: "tests/test_litellm/integrations"
workers: 2
reruns: 3
artifact-name: integrations

View file

@ -1,47 +0,0 @@
name: "Unit Tests: LLM Provider Transformations"
on:
pull_request:
branches:
- main
- litellm_internal_staging
- litellm_oss_staging
- "litellm_**"
push:
branches:
- main
- litellm_internal_staging
permissions:
contents: read
concurrency:
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}
cancel-in-progress: ${{ github.event_name == 'pull_request' }}
jobs:
vertex-ai:
name: Vertex AI
permissions:
contents: read
id-token: write
pull-requests: write
uses: ./.github/workflows/_test-unit-base.yml
with:
test-path: "tests/test_litellm/llms/vertex_ai"
workers: 1
reruns: 2
artifact-name: llm-vertex-ai
other-providers:
name: All Other Providers
permissions:
contents: read
id-token: write
pull-requests: write
uses: ./.github/workflows/_test-unit-base.yml
with:
test-path: "tests/test_litellm/llms --ignore=tests/test_litellm/llms/vertex_ai"
workers: 2
reruns: 2
artifact-name: llm-other-providers

View file

@ -1,53 +0,0 @@
name: "Unit Tests: MCP, Secrets, Containers & Misc"
on:
pull_request:
branches:
- main
- litellm_internal_staging
- litellm_oss_staging
- "litellm_**"
push:
branches:
- main
- litellm_internal_staging
permissions:
contents: read
id-token: write
pull-requests: write
concurrency:
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}
cancel-in-progress: ${{ github.event_name == 'pull_request' }}
jobs:
misc:
uses: ./.github/workflows/_test-unit-base.yml
with:
test-path: >-
tests/test_litellm/batches
tests/test_litellm/secret_managers
tests/test_litellm/a2a_protocol
tests/test_litellm/anthropic_interface
tests/test_litellm/completion_extras
tests/test_litellm/compression
tests/test_litellm/containers
tests/test_litellm/experimental_mcp_client
tests/test_litellm/models
tests/test_litellm/repositories
tests/test_litellm/images
tests/test_litellm/interactions
tests/test_litellm/ocr
tests/test_litellm/passthrough
tests/test_litellm/rag
tests/test_litellm/realtime_api
tests/test_litellm/rerank_api
tests/test_litellm/sandbox
tests/test_litellm/test_router
tests/test_litellm/vector_stores
tests/test_litellm/videos
tests/test_litellm/test_*.py
workers: 2
reruns: 2
artifact-name: misc

View file

@ -1,31 +0,0 @@
name: "Unit Tests: Proxy Auth & Key Management"
on:
pull_request:
branches:
- main
- litellm_internal_staging
- litellm_oss_staging
- "litellm_**"
push:
branches:
- main
- litellm_internal_staging
permissions:
contents: read
id-token: write
pull-requests: write
concurrency:
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}
cancel-in-progress: ${{ github.event_name == 'pull_request' }}
jobs:
proxy-auth:
uses: ./.github/workflows/_test-unit-base.yml
with:
test-path: "tests/test_litellm/proxy/auth tests/test_litellm/proxy/hooks tests/test_litellm/proxy/policy_engine tests/test_litellm/proxy/client"
workers: 2
reruns: 2
artifact-name: proxy-auth

View file

@ -42,11 +42,10 @@ concurrency:
# pinning the whole file to one worker (the default --dist=loadscope
# behavior for single-file targets).
jobs:
# Fast guard — fails the workflow if a test_*.py file under
# tests/proxy_unit_tests/ is not referenced by any matrix entry below.
# The semantic-shard design (no catch-all "remaining" bucket) relies on
# every test file being explicitly assigned; this guard prevents a new
# file from silently dropping out of CI.
# Fast guard — fails the workflow when a test directory or file inside a sharded
# tree is claimed by no shard. The semantic-shard design has no catch-all bucket,
# so an unassigned child runs nowhere; assert_ci_coverage.py holds the tree list
# and reads the same test-path keys the coverage census does.
assert-shard-coverage:
runs-on: ubuntu-latest
timeout-minutes: 2
@ -56,31 +55,8 @@ jobs:
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
persist-credentials: false
- name: Assert every test_*.py is in a matrix shard
run: |
python3 - <<'PY'
import pathlib, sys, yaml
wf = yaml.safe_load(open(".github/workflows/test-unit-proxy-db.yml"))
matrix = wf["jobs"]["proxy-db"]["strategy"]["matrix"]["include"]
referenced = set()
for entry in matrix:
for token in entry["test-path"].split():
if token.startswith("tests/proxy_unit_tests/"):
referenced.add(pathlib.PurePosixPath(token).name)
actual = {p.name for p in pathlib.Path("tests/proxy_unit_tests").iterdir()
if p.name.startswith("test_") and (p.suffix == ".py" or p.is_dir())
and p.name != "test_configs"}
orphans = sorted(actual - referenced)
if orphans:
print("ERROR: the following files/dirs under tests/proxy_unit_tests/")
print(" are not assigned to any shard in test-unit-proxy-db.yml:")
for o in orphans:
print(f" - {o}")
print()
print("Add each to whichever semantic shard it belongs to.")
sys.exit(1)
print(f"OK: all {len(actual)} files assigned to a shard.")
PY
- name: Assert every test directory and file is claimed by a shard
run: python3 .github/scripts/assert_ci_coverage.py --shards
proxy-db:
needs: assert-shard-coverage

View file

@ -1,81 +0,0 @@
name: "Unit Tests: Proxy API Endpoints"
on:
pull_request:
branches:
- main
- litellm_internal_staging
- litellm_oss_staging
- "litellm_**"
push:
branches:
- main
- litellm_internal_staging
workflow_dispatch:
permissions:
contents: read
concurrency:
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}
cancel-in-progress: ${{ github.event_name == 'pull_request' }}
jobs:
proxy-endpoints:
permissions:
contents: read
id-token: write
pull-requests: write
uses: ./.github/workflows/_test-unit-base.yml
with:
test-path: >-
tests/test_litellm/proxy/analytics_endpoints
tests/test_litellm/proxy/management_endpoints
tests/test_litellm/proxy/memory
tests/test_litellm/proxy/guardrails
tests/test_litellm/proxy/management_helpers
tests/test_litellm/proxy/anthropic_endpoints
tests/test_litellm/proxy/google_endpoints
tests/test_litellm/proxy/openai_files_endpoint
tests/test_litellm/proxy/batches_endpoints
tests/test_litellm/proxy/fine_tuning_endpoints
tests/test_litellm/proxy/vector_store_files_endpoints
tests/test_litellm/proxy/video_endpoints
tests/test_litellm/proxy/response_api_endpoints
tests/test_litellm/proxy/image_endpoints
tests/test_litellm/proxy/ocr_endpoints
tests/test_litellm/proxy/vector_store_endpoints
tests/test_litellm/proxy/agent_endpoints
tests/test_litellm/proxy/a2a
tests/test_litellm/proxy/credential_endpoints
tests/test_litellm/proxy/discovery_endpoints
tests/test_litellm/proxy/health_endpoints
tests/test_litellm/proxy/shutdown
tests/test_litellm/proxy/public_endpoints
tests/test_litellm/proxy/prompts
tests/test_litellm/proxy/rag_endpoints
tests/test_litellm/proxy/realtime_endpoints
tests/test_litellm/proxy/ui_crud_endpoints
tests/test_litellm/proxy/config_resolvers
tests/test_litellm/proxy/utils
workers: 2
reruns: 2
artifact-name: proxy-endpoints
# Behavior-pinning tests for litellm/proxy/proxy_server.py. Owns its
# own job (not a path on the proxy-endpoints job above) so its budget
# is independent and its coverage artifact is uploaded separately.
# See: https://www.notion.so/36c43b8acdab81ee845fd5365128a2fc
proxy-server:
permissions:
contents: read
id-token: write
pull-requests: write
uses: ./.github/workflows/_test-unit-base.yml
with:
test-path: tests/test_litellm/proxy/proxy_server
workers: 4
reruns: 2
timeout-minutes: 60
job-timeout-minutes: 95
artifact-name: proxy-server

View file

@ -1,42 +0,0 @@
name: "Unit Tests: Proxy Infrastructure"
on:
pull_request:
branches:
- main
- litellm_internal_staging
- litellm_oss_staging
- "litellm_**"
push:
branches:
- main
- litellm_internal_staging
permissions:
contents: read
id-token: write
pull-requests: write
concurrency:
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}
cancel-in-progress: ${{ github.event_name == 'pull_request' }}
jobs:
proxy-infra:
uses: ./.github/workflows/_test-unit-base.yml
with:
test-path: >-
tests/test_litellm/proxy/db
tests/test_litellm/proxy/middleware
tests/test_litellm/proxy/spend_tracking
tests/test_litellm/proxy/pass_through_endpoints
tests/test_litellm/proxy/_experimental
tests/test_litellm/proxy/experimental
tests/test_litellm/proxy/common_utils
tests/test_litellm/proxy/enterprise_billing
tests/test_litellm/proxy/types_utils
tests/test_litellm/proxy/logging_endpoints
tests/test_litellm/proxy/test_*.py
workers: 2
reruns: 2
artifact-name: proxy-infra

View file

@ -1,31 +0,0 @@
name: "Unit Tests: Responses, Caching & Types"
on:
pull_request:
branches:
- main
- litellm_internal_staging
- litellm_oss_staging
- "litellm_**"
push:
branches:
- main
- litellm_internal_staging
permissions:
contents: read
id-token: write
pull-requests: write
concurrency:
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}
cancel-in-progress: ${{ github.event_name == 'pull_request' }}
jobs:
responses-caching-types:
uses: ./.github/workflows/_test-unit-base.yml
with:
test-path: "tests/test_litellm/responses tests/test_litellm/caching tests/test_litellm/types"
workers: 2
reruns: 2
artifact-name: responses-caching-types

220
.github/workflows/test-unit.yml vendored Normal file
View file

@ -0,0 +1,220 @@
name: "Unit Tests"
on:
pull_request:
branches:
- main
- litellm_internal_staging
- litellm_oss_staging
- "litellm_**"
push:
branches:
- main
- litellm_internal_staging
workflow_dispatch:
permissions:
contents: read
concurrency:
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}
cancel-in-progress: ${{ github.event_name == 'pull_request' }}
# One caller for every tests/test_litellm shard, replacing the nine thin workflow
# files that each wrapped a single call to _test-unit-base.yml. Adding a shard is
# now one matrix entry rather than a new file.
#
# `name` is the shard id and nothing else, so each check reports as
# "<shard> / Run tests" exactly as it did when the shard had its own file. Those
# strings are the branch ruleset's required contexts, so they are load-bearing:
# renaming an entry renames a required check and the ruleset stops matching it.
#
# Every entry states its timeouts even when they equal the base workflow's
# defaults. An absent matrix key renders as an empty string, which is not a
# number, so a partially-specified entry would fail the call rather than fall
# back to the default.
#
# tests/proxy_unit_tests keeps its own caller (test-unit-proxy-db.yml): it is
# already a matrix and carries a shard-coverage guard that reads that file by
# name. Folding it in here is a follow-up, together with generalising that guard
# into assert_ci_coverage.py.
jobs:
unit:
name: ${{ matrix.shard }}
permissions:
contents: read
id-token: write
pull-requests: write
strategy:
fail-fast: false
matrix:
include:
- shard: core-utils
artifact-name: core-utils
test-path: "tests/test_litellm/litellm_core_utils"
workers: 2
reruns: 1
timeout-minutes: 20
job-timeout-minutes: 55
- shard: enterprise-routing
artifact-name: enterprise-routing
test-path: >-
tests/test_litellm/enterprise
tests/test_litellm/google_genai
tests/test_litellm/router_utils
tests/test_litellm/router_strategy
workers: 2
reruns: 2
timeout-minutes: 20
job-timeout-minutes: 55
- shard: integrations
artifact-name: integrations
test-path: "tests/test_litellm/integrations"
workers: 2
reruns: 3
timeout-minutes: 20
job-timeout-minutes: 55
- shard: Vertex AI
artifact-name: llm-vertex-ai
test-path: "tests/test_litellm/llms/vertex_ai"
workers: 1
reruns: 2
timeout-minutes: 20
job-timeout-minutes: 55
- shard: All Other Providers
artifact-name: llm-other-providers
test-path: "tests/test_litellm/llms --ignore=tests/test_litellm/llms/vertex_ai"
workers: 2
reruns: 2
timeout-minutes: 20
job-timeout-minutes: 55
- shard: misc
artifact-name: misc
test-path: >-
tests/test_litellm/batches
tests/test_litellm/secret_managers
tests/test_litellm/a2a_protocol
tests/test_litellm/anthropic_interface
tests/test_litellm/completion_extras
tests/test_litellm/compression
tests/test_litellm/containers
tests/test_litellm/experimental_mcp_client
tests/test_litellm/models
tests/test_litellm/repositories
tests/test_litellm/images
tests/test_litellm/interactions
tests/test_litellm/ocr
tests/test_litellm/passthrough
tests/test_litellm/rag
tests/test_litellm/realtime_api
tests/test_litellm/rerank_api
tests/test_litellm/rust_bridge
tests/test_litellm/sandbox
tests/test_litellm/test_router
tests/test_litellm/vector_stores
tests/test_litellm/videos
tests/test_litellm/test_*.py
workers: 2
reruns: 2
timeout-minutes: 20
job-timeout-minutes: 55
- shard: proxy-auth
artifact-name: proxy-auth
test-path: >-
tests/test_litellm/proxy/auth
tests/test_litellm/proxy/hooks
tests/test_litellm/proxy/policy_engine
tests/test_litellm/proxy/client
workers: 2
reruns: 2
timeout-minutes: 20
job-timeout-minutes: 55
- shard: proxy-endpoints
artifact-name: proxy-endpoints
test-path: >-
tests/test_litellm/proxy/analytics_endpoints
tests/test_litellm/proxy/management_endpoints
tests/test_litellm/proxy/memory
tests/test_litellm/proxy/guardrails
tests/test_litellm/proxy/management_helpers
tests/test_litellm/proxy/anthropic_endpoints
tests/test_litellm/proxy/google_endpoints
tests/test_litellm/proxy/openai_files_endpoint
tests/test_litellm/proxy/batches_endpoints
tests/test_litellm/proxy/fine_tuning_endpoints
tests/test_litellm/proxy/vector_store_files_endpoints
tests/test_litellm/proxy/video_endpoints
tests/test_litellm/proxy/response_api_endpoints
tests/test_litellm/proxy/image_endpoints
tests/test_litellm/proxy/ocr_endpoints
tests/test_litellm/proxy/vector_store_endpoints
tests/test_litellm/proxy/agent_endpoints
tests/test_litellm/proxy/a2a
tests/test_litellm/proxy/credential_endpoints
tests/test_litellm/proxy/discovery_endpoints
tests/test_litellm/proxy/health_endpoints
tests/test_litellm/proxy/shutdown
tests/test_litellm/proxy/public_endpoints
tests/test_litellm/proxy/prompts
tests/test_litellm/proxy/rag_endpoints
tests/test_litellm/proxy/realtime_endpoints
tests/test_litellm/proxy/ui_crud_endpoints
tests/test_litellm/proxy/config_resolvers
tests/test_litellm/proxy/utils
workers: 2
reruns: 2
timeout-minutes: 20
job-timeout-minutes: 55
- shard: proxy-server
artifact-name: proxy-server
test-path: "tests/test_litellm/proxy/proxy_server"
workers: 4
reruns: 2
timeout-minutes: 60
job-timeout-minutes: 95
- shard: proxy-infra
artifact-name: proxy-infra
test-path: >-
tests/test_litellm/proxy/db
tests/test_litellm/proxy/middleware
tests/test_litellm/proxy/spend_tracking
tests/test_litellm/proxy/pass_through_endpoints
tests/test_litellm/proxy/_experimental
tests/test_litellm/proxy/experimental
tests/test_litellm/proxy/common_utils
tests/test_litellm/proxy/enterprise_billing
tests/test_litellm/proxy/types_utils
tests/test_litellm/proxy/logging_endpoints
tests/test_litellm/proxy/test_*.py
workers: 2
reruns: 2
timeout-minutes: 20
job-timeout-minutes: 55
- shard: responses-caching-types
artifact-name: responses-caching-types
test-path: >-
tests/test_litellm/responses
tests/test_litellm/caching
tests/test_litellm/types
workers: 2
reruns: 2
timeout-minutes: 20
job-timeout-minutes: 55
uses: ./.github/workflows/_test-unit-base.yml
with:
test-path: ${{ matrix.test-path }}
workers: ${{ matrix.workers }}
reruns: ${{ matrix.reruns }}
timeout-minutes: ${{ matrix.timeout-minutes }}
job-timeout-minutes: ${{ matrix.job-timeout-minutes }}
artifact-name: ${{ matrix.artifact-name }}

View file

@ -1,92 +0,0 @@
name: Agent Shin — rollout heads-up (one-shot)
# Fires the 7-day heads-up comment on every open external PR/issue that the
# new triage bot would auto-close. The real sweep is a deliberate one-shot:
# trigger it at rollout via a manual `workflow_dispatch` with `dry_run=false`.
# The script is idempotent (skips items that already carry the
# `<!-- agent-shin:rollout-heads-up -->` marker), so a re-run is harmless.
#
# The automatic push trigger runs DRY-RUN only, so merging the script to
# `litellm_internal_staging` never posts a comment; it just confirms the
# workflow is wired up. Posting real comments requires the manual dispatch,
# which is also the only trigger that exposes `OPENAI_API_KEY`. The heads-up
# is intentionally NOT gated on `AGENT_SHIN_ENABLED`: it has to warn
# contributors while that flag is still off, ahead of the flip that turns on
# auto-closing.
#
# The workflow is a thin shell over `.github/scripts/triage_rollout_heads_up.py`.
# Dry-run vs. real run differ in EXACTLY one CLI flag (`--close`), added only
# on a manual dispatch with `dry_run=false`.
on:
push:
branches:
- litellm_internal_staging
paths:
# The presence of this script on staging IS the rollout merge marker.
# Editing the file later would re-fire the workflow; that's safe because
# the script skips PRs/issues that already have the heads-up marker.
- ".github/scripts/triage_rollout_heads_up.py"
workflow_dispatch:
inputs:
dry_run:
description: "Dry run (true = preview only, false = actually post comments)."
required: false
default: "true"
type: choice
options:
- "true"
- "false"
permissions:
contents: read
issues: write
pull-requests: write
jobs:
heads-up:
if: github.repository == 'BerriAI/litellm'
runs-on: ubuntu-latest
steps:
- name: Checkout triage scripts
uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
sparse-checkout: .github/scripts
persist-credentials: false
- name: Set up Python
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
with:
python-version: "3.12"
- name: Install LLM client
run: pip install --no-cache-dir --require-hashes -r .github/scripts/triage-requirements.txt
- name: Run heads-up sweep
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
# Only the manual dispatch (the real-run trigger) needs the LLM key.
# The automatic push trigger runs dry-run and never posts, so it gets
# no key. Mirrors the sibling triage workflows, which expose the key
# only on an enabled/dispatched run rather than unconditionally.
OPENAI_API_KEY: ${{ github.event_name == 'workflow_dispatch' && secrets.OPENAI_API_KEY || '' }}
OPENAI_BASE_URL: ${{ vars.OPENAI_BASE_URL }}
TRIAGE_MODEL: ${{ vars.TRIAGE_MODEL }}
# The real run is a deliberate manual dispatch with dry_run=false.
# Use the EXACT "false" comparison so any unexpected input value
# fail-closes to dry-run (mirrors the AGENT_SHIN_ENABLED pattern in
# the sibling workflows). The automatic push trigger always stays
# dry-run, so merging the script never posts.
DRY_RUN_INPUT: ${{ github.event.inputs.dry_run }}
run: |
set -euo pipefail
ARGS=(--repo "${{ github.repository }}")
if [ "${GITHUB_EVENT_NAME:-}" = "workflow_dispatch" ] && [ "${DRY_RUN_INPUT:-true}" = "false" ]; then
ARGS+=(--close)
echo "::notice::Manual rollout dispatch with dry_run=false -> heads-up comments WILL be posted."
elif [ "${GITHUB_EVENT_NAME:-}" = "workflow_dispatch" ]; then
echo "::notice::Manual dispatch in dry-run mode -> previewing only, no comments will be posted."
else
echo "::notice::Automatic push trigger -> dry-run preview only. Fire the real rollout sweep with a manual workflow_dispatch (dry_run=false)."
fi
python3 .github/scripts/triage_rollout_heads_up.py "${ARGS[@]}"

View file

@ -7,6 +7,7 @@
info lint lint-inner lint-dev lint-checks format \
lint-basedpyright lint-e2e-basedpyright lint-basedpyright-budget-update lint-type-discipline lint-type-discipline-budget-update \
lint-ruff-budget lint-ruff-budget-update lint-budget-update lint-gate \
lint-test-quality lint-test-quality-budget-update \
install-dev install-proxy-dev install-test-deps install-hooks \
install-helm-unittest check-circular-imports check-import-safety check check-inner pre-commit \
lint-install lint-fetch-base bootstrap
@ -35,7 +36,8 @@ help:
@echo " make lint-ruff-budget - Gate the codebase total of each strict ruff rule against its limit"
@echo " make lint-gate - Strict ruff gate in CI-parity mode (fetches staging, simulates the merge)"
@echo " make lint-ruff-budget-update - Ratchet ruff-strict-budget.json limits down by what this branch fixed"
@echo " make lint-budget-update - Ratchet all budgets down (ruff + type-discipline + basedpyright)"
@echo " make lint-test-quality - Gate the test suite against test-quality-budget.json"
@echo " make lint-budget-update - Ratchet all budgets down (ruff + type-discipline + test quality + basedpyright)"
@echo " make check-circular-imports - Check for circular imports"
@echo " make check-import-safety - Check import safety"
@echo " make test - Run all tests"
@ -142,11 +144,13 @@ lint-install:
$(UV) sync --inexact --frozen --group proxy-dev --group e2e-dev
$(UV_RUN) python scripts/prisma_generate_if_needed.py
# Diff-scoped format check, identical to test-linting.yml's "Check ruff format" step:
# Diff-scoped format check, mirroring test-linting.yml's "Check ruff format" step:
# only the litellm Python files changed vs the base are checked, so a pre-existing
# format issue elsewhere doesn't block an unrelated commit.
# format issue elsewhere doesn't block an unrelated commit. Git pathspecs match
# recursively, so 'litellm/*.py' covers nested modules and the top-level files that
# CI's 'litellm/**/*.py' skips, which makes this target a superset of the CI step.
lint-format-check-changed: $(LINT_DEP_INSTALL) $(LINT_DEP_BASE)
@files=$$(git diff --name-only origin/litellm_internal_staging...HEAD -- 'litellm/**/*.py' | grep -v '^litellm/enterprise/' || true); \
@files=$$(git diff --name-only --diff-filter=ACMR origin/litellm_internal_staging...HEAD -- 'litellm/*.py' | grep -v '^litellm/enterprise/' || true); \
if [ -z "$$files" ]; then \
echo "No changed litellm Python files to format-check."; \
else \
@ -156,6 +160,7 @@ lint-format-check-changed: $(LINT_DEP_INSTALL) $(LINT_DEP_BASE)
# Linting targets
lint-ruff: $(LINT_DEP_INSTALL)
cd litellm && $(UV_RUN) ruff check . && cd ..
$(UV_RUN) ruff check --config ruff-tests.toml tests
# faster linter for developing ...
# inspiration from:
@ -200,6 +205,12 @@ lint-e2e-basedpyright: $(LINT_E2E_DEP_INSTALL)
lint-type-discipline: $(LINT_DEP_INSTALL) $(LINT_DEP_BASE)
$(UV_RUN) python scripts/type_discipline_gate.py --base origin/litellm_internal_staging
# Test-quality budget (zero-assert / mock-echo tests, sys.path.insert, raw env writes,
# litellm module-global mutation, credential-gated skips, conftest snapshot
# inventory), counted across tests/ the same delta-vs-base way.
lint-test-quality: $(LINT_DEP_INSTALL) $(LINT_DEP_BASE)
$(UV_RUN) python scripts/test_quality_gate.py --base origin/litellm_internal_staging
# --update lowers each limit by what this branch fixed since its branch point, so
# it needs the base ref fetched to resolve the merge-base.
lint-basedpyright-budget-update: install-dev lint-fetch-base
@ -221,8 +232,11 @@ lint-ruff-budget-update: install-dev lint-fetch-base
lint-type-discipline-budget-update: install-dev lint-fetch-base
$(UV_RUN) python scripts/type_discipline_gate.py --update
# Ratchet all budgets in one shot (ruff strict + type-discipline + basedpyright)
lint-budget-update: lint-ruff-budget-update lint-type-discipline-budget-update lint-basedpyright-budget-update
lint-test-quality-budget-update: install-dev lint-fetch-base
$(UV_RUN) python scripts/test_quality_gate.py --update
# Ratchet all budgets in one shot (ruff strict + type-discipline + test quality + basedpyright)
lint-budget-update: lint-ruff-budget-update lint-type-discipline-budget-update lint-test-quality-budget-update lint-basedpyright-budget-update
check-circular-imports: $(LINT_DEP_INSTALL)
cd litellm && $(UV_RUN) python ../tests/documentation_tests/test_circular_imports.py && cd ..
@ -244,7 +258,7 @@ lint:
lint-inner: lint-install lint-fetch-base
$(MAKE) -j $(LINT_JOBS) $(LINT_OUTPUT_SYNC) LINT_DEP_INSTALL= LINT_E2E_DEP_INSTALL= LINT_DEP_BASE= lint-checks
lint-checks: lint-format-check-changed lint-ruff lint-gate lint-type-discipline lint-basedpyright lint-e2e-basedpyright check-circular-imports check-import-safety
lint-checks: lint-format-check-changed lint-ruff lint-gate lint-type-discipline lint-test-quality lint-basedpyright lint-e2e-basedpyright check-circular-imports check-import-safety
# Faster linting for local development (only checks changed code)
lint-dev: lint-format-changed check-circular-imports check-import-safety
@ -314,7 +328,7 @@ test-unit-helm: install-helm-unittest
# LLM Translation testing targets
test-llm-translation: install-test-deps
@echo "Running LLM translation tests..."
@python .github/workflows/run_llm_translation_tests.py
@python .github/scripts/run_llm_translation_tests.py
test-llm-translation-single: install-test-deps
@echo "Running single LLM translation test file..."

View file

@ -292,6 +292,7 @@ curl -X POST 'http://0.0.0.0:4000/v1/chat/completions' \
| [Clarifai (`clarifai`)](https://docs.litellm.ai/docs/providers/clarifai) | ✅ | ✅ | ✅ | | | | | | | |
| [Cloudflare AI Workers (`cloudflare`)](https://docs.litellm.ai/docs/providers/cloudflare_workers) | ✅ | ✅ | ✅ | | | | | | | |
| [Codestral (`codestral`)](https://docs.litellm.ai/docs/providers/codestral) | ✅ | ✅ | ✅ | | | | | | | |
| [Cognition (`cognition`)](https://docs.litellm.ai/docs/providers/cognition) | ✅ | ✅ | ✅ | | | | | | | |
| [Cohere (`cohere`)](https://docs.litellm.ai/docs/providers/cohere) | ✅ | ✅ | ✅ | ✅ | | | | | | ✅ |
| [Cohere Chat (`cohere_chat`)](https://docs.litellm.ai/docs/providers/cohere) | ✅ | ✅ | ✅ | | | | | | | |
| [CometAPI (`cometapi`)](https://docs.litellm.ai/docs/providers/cometapi) | ✅ | ✅ | ✅ | ✅ | | | | | | |

View file

@ -105,13 +105,13 @@
"limit": 109
},
"reportUnknownMemberType": {
"limit": 39017
"limit": 39011
},
"reportUnknownParameterType": {
"limit": 19885
},
"reportUnknownVariableType": {
"limit": 30572
"limit": 30569
},
"reportUnnecessaryCast": {
"limit": 117

View file

@ -60,4 +60,4 @@ if __name__ == "__main__":
print("\n💡 Tips:")
print("1. Run 'litellm-proxy login' to authenticate first")
print("2. Replace 'https://your-proxy.com' with your actual proxy URL")
print("3. The token is stored locally at ~/.litellm/token.json")
print("3. The token is stored in your OS keychain, or in ~/.litellm/token.json when there is none")

View file

@ -255,6 +255,52 @@ class CheckBatchCost:
"so it will no longer be polled"
)
async def _claim_job_for_costing(self, job: "LiteLLM_ManagedObjectTable") -> bool:
"""
Atomically flip batch_processed from false to true, returning whether this pod won
the row. Every pod and uvicorn worker schedules its own poller against the shared
table, so without this compare-and-swap two of them can select the same completed
batch in one window and both emit an aretrieve_batch spend log for it. Schemas
without the column can't be claimed, so they keep the pre-existing behavior.
Called immediately before the spend log is written rather than before the results
fetch, because batch_processed is also what holds off deletion of the files that
fetch reads and what keeps an unbilled row selectable by the next poll cycle.
"""
if not self._has_batch_processed_column:
return True
try:
claimed: Final = await self.prisma_client.db.litellm_managedobjecttable.update_many(
where={"id": job.id, "batch_processed": False},
data={"batch_processed": True},
)
except Exception as db_err:
verbose_proxy_logger.error(
f"CheckBatchCost: failed to claim job {job.id} for cost tracking: {db_err}"
)
return False
return claimed > 0
async def _release_job_claim(self, job: "LiteLLM_ManagedObjectTable") -> None:
"""Give a claimed row back once billing it failed, so a later poll cycle retries it.
Safe to match on batch_processed=True: while this poller is active the retrieve
path leaves the column alone (batch_cost_poller_is_active), so a true value here
is always this pod's own claim.
"""
if not self._has_batch_processed_column:
return
try:
await self.prisma_client.db.litellm_managedobjecttable.update_many(
where={"id": job.id, "batch_processed": True},
data={"batch_processed": False},
)
except Exception as db_err:
verbose_proxy_logger.error(
f"CheckBatchCost: failed to release the claim on job {job.id}, "
f"so its cost will not be retried: {db_err}"
)
@staticmethod
def _has_unified_id_without_model(job: "LiteLLM_ManagedObjectTable") -> bool:
"""A unified id that decodes but carries no model_id can never be routed."""
@ -572,9 +618,10 @@ class CheckBatchCost:
"""
Fetch a completed batch's results, compute cost/usage, and emit the
aretrieve_batch spend log. Returns (model_name, llm_provider) on
success, None when the job can't be routed to a deployment. Raises on
results-fetch or cost-computation failures so the caller can leave the
job unprocessed and retry it on a later poll.
success, None when the job can't be routed to a deployment or when
another pod claimed it. Raises on results-fetch or cost-computation
failures so the caller can leave the job unprocessed and retry it on a
later poll.
"""
from litellm.batches.batch_utils import (
_get_file_content_as_dictionary,
@ -743,12 +790,23 @@ class CheckBatchCost:
optional_params={},
)
await logging_obj.async_success_handler(
result=response,
batch_cost=batch_cost,
batch_usage=batch_usage,
batch_models=batch_models,
)
if not await self._claim_job_for_costing(job):
verbose_proxy_logger.info(
f"CheckBatchCost: batch {batch_id} (job {job.id}) was claimed by another pod "
"in this window, so its cost is already being tracked there"
)
return None
try:
await logging_obj.async_success_handler(
result=response,
batch_cost=batch_cost,
batch_usage=batch_usage,
batch_models=batch_models,
)
except Exception:
await self._release_job_claim(job)
raise
# Record batch duration (completed_at - created_at)
if prom_logger and response.completed_at and response.created_at:

View file

@ -1,6 +1,6 @@
[project]
name = "litellm-enterprise"
version = "0.1.57"
version = "0.1.58"
description = "Package for LiteLLM Enterprise features"
readme = "README.md"
requires-python = ">=3.9"
@ -26,7 +26,7 @@ required-version = ">=0.10.9"
module-root = ""
[tool.commitizen]
version = "0.1.57"
version = "0.1.58"
version_files = [
"pyproject.toml:^version",
"../pyproject.toml:litellm-enterprise==",

View file

@ -100,6 +100,13 @@ spec:
- name: DATABASE_URL
value: {{ .Values.db.url | quote }}
{{- end }}
{{- if and .Values.db.useExisting .Values.db.readReplicaUrl .Values.db.secret.readReplicaEndpointKey (not .Values.db.secret.readReplicaUrlKey) }}
- name: DATABASE_READER_HOST
valueFrom:
secretKeyRef:
name: {{ .Values.db.secret.name }}
key: {{ .Values.db.secret.readReplicaEndpointKey }}
{{- end }}
{{- if and .Values.db.useExisting .Values.db.secret.readReplicaUrlKey }}
- name: DATABASE_URL_READ_REPLICA
valueFrom:

View file

@ -80,6 +80,96 @@ tests:
secretKeyRef:
name: my-secret
key: my-key
- it: should inject DATABASE_READER_HOST from readReplicaEndpointKey before DATABASE_URL_READ_REPLICA
template: deployment.yaml
set:
db:
deployStandalone: false
useExisting: true
secret:
name: postgres
usernameKey: username
passwordKey: password
readReplicaEndpointKey: reader-host
readReplicaUrl: postgresql://$(DATABASE_USERNAME):$(DATABASE_PASSWORD)@$(DATABASE_READER_HOST):5432/$(DATABASE_NAME)?sslmode=require
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_READER_HOST
valueFrom:
secretKeyRef:
name: postgres
key: reader-host
- contains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_URL_READ_REPLICA
value: postgresql://$(DATABASE_USERNAME):$(DATABASE_PASSWORD)@$(DATABASE_READER_HOST):5432/$(DATABASE_NAME)?sslmode=require
# $(VAR) interpolation only resolves vars defined EARLIER in the env
# array, so the reader host must precede the composed URL
- equal:
path: spec.template.spec.containers[0].env[7].name
value: DATABASE_READER_HOST
- equal:
path: spec.template.spec.containers[0].env[8].name
value: DATABASE_URL_READ_REPLICA
- it: should omit reader host when readReplicaUrl is unset
template: deployment.yaml
set:
db:
deployStandalone: false
useExisting: true
secret:
name: postgres
usernameKey: username
passwordKey: password
readReplicaEndpointKey: reader-host
asserts:
- notContains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_READER_HOST
valueFrom:
secretKeyRef:
name: postgres
key: reader-host
- it: should prefer readReplicaUrlKey over readReplicaEndpointKey composition
template: deployment.yaml
set:
db:
useExisting: true
secret:
name: postgres
usernameKey: username
passwordKey: password
readReplicaUrlKey: reader-url
readReplicaEndpointKey: reader-host
readReplicaUrl: postgresql://ignored
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_URL_READ_REPLICA
valueFrom:
secretKeyRef:
name: postgres
key: reader-url
- notContains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_URL_READ_REPLICA
value: postgresql://ignored
# the unused reader-host secret ref must be suppressed so a missing
# key can't fail pod creation
- notContains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_READER_HOST
valueFrom:
secretKeyRef:
name: postgres
key: reader-host
- it: should work with extraEnvVars
template: deployment.yaml
set:

View file

@ -277,6 +277,14 @@ db:
# written to db.readReplicaUrl ends up visible in the rendered pod spec
# and the Helm release secret.
readReplicaUrlKey: ""
# Optional: when set, a DATABASE_READER_HOST env var is sourced from this
# secret key, so db.readReplicaUrl can compose the reader URL from
# individual secret components, e.g.
# postgresql://$(DATABASE_USERNAME):$(DATABASE_PASSWORD)@$(DATABASE_READER_HOST):5432/$(DATABASE_NAME)
# Use this when your secret store holds the bare reader hostname rather
# than a full connection URL. Only takes effect when readReplicaUrl is
# set; ignored when readReplicaUrlKey is set.
readReplicaEndpointKey: ""
# Optional read-replica routing. When set, the proxy sends read-only
# queries (find_*, count, group_by, query_raw/_first) to this URL while

View file

@ -213,18 +213,21 @@ whenever the password contains a URL-reserved character (@, /, ?, %, +,
When `database.writer.useIAMAuth: true`, the chart injects
IAM_TOKEN_DB_AUTH=true and omits DATABASE_PASSWORD — the entrypoint mints
the URL from DATABASE_HOST/PORT/USER/NAME plus a short-lived IAM token
instead of a static password.
the URL from DATABASE_HOST/PORT/USER/NAME plus a short-lived AWS RDS IAM
token instead of a static password. `database.writer.useAzureEntraAuth: true`
does the same with AZURE_POSTGRESQL_AUTH=true and a Microsoft Entra ID token,
for Azure Database for PostgreSQL. The two are mutually exclusive.
The read replica is opt-in via `database.reader.host`. The chart emits
DATABASE_HOST_READ_REPLICA / DATABASE_PORT_READ_REPLICA /
DATABASE_NAME_READ_REPLICA (+ DATABASE_SCHEMA_READ_REPLICA) for both auth
modes, plus DATABASE_USER_READ_REPLICA / DATABASE_PASSWORD_READ_REPLICA for
password auth. When `database.reader.useIAMAuth: true` it omits
password auth. When `database.reader.useIAMAuth: true` (or
`database.reader.useAzureEntraAuth: true`) it omits
DATABASE_PASSWORD_READ_REPLICA and the entrypoint mints the reader URL the
same way. Reader IAM only takes effect when the writer also uses IAM auth
(the proxy gates URL minting on IAM_TOKEN_DB_AUTH, which only the writer
sets).
same way. Reader token auth only takes effect when the writer uses the same
token source, since the proxy gates URL minting on the single global
IAM_TOKEN_DB_AUTH / AZURE_POSTGRESQL_AUTH toggle that only the writer sets.
*/}}
{{- define "litellm.serverEnv" -}}
{{- $root := .root -}}
@ -254,9 +257,15 @@ sets).
- name: DATABASE_SCHEMA
value: {{ .schema | quote }}
{{- end }}
{{- if and .useIAMAuth .useAzureEntraAuth }}
{{- fail "database.writer.useIAMAuth and database.writer.useAzureEntraAuth are mutually exclusive: the database password can only come from one token source" }}
{{- end }}
{{- if .useIAMAuth }}
- name: IAM_TOKEN_DB_AUTH
value: "true"
{{- else if .useAzureEntraAuth }}
- name: AZURE_POSTGRESQL_AUTH
value: "true"
{{- else }}
- name: DATABASE_PASSWORD
valueFrom:
@ -270,6 +279,9 @@ sets).
{{- if and .useIAMAuth (not $root.Values.database.writer.useIAMAuth) }}
{{- fail "database.reader.useIAMAuth requires database.writer.useIAMAuth: true (the proxy gates IAM URL minting on IAM_TOKEN_DB_AUTH, which is only set by the writer)" }}
{{- end }}
{{- if and .useAzureEntraAuth (not $root.Values.database.writer.useAzureEntraAuth) }}
{{- fail "database.reader.useAzureEntraAuth requires database.writer.useAzureEntraAuth: true (the proxy gates Entra URL minting on AZURE_POSTGRESQL_AUTH, which is only set by the writer)" }}
{{- end }}
- name: DATABASE_HOST_READ_REPLICA
value: {{ .host | quote }}
- name: DATABASE_PORT_READ_REPLICA
@ -280,7 +292,7 @@ sets).
- name: DATABASE_SCHEMA_READ_REPLICA
value: {{ .schema | quote }}
{{- end }}
{{- if .useIAMAuth }}
{{- if or .useIAMAuth .useAzureEntraAuth }}
{{- if .passwordSecret.name }}
- name: DATABASE_USER_READ_REPLICA
valueFrom:

View file

@ -0,0 +1,116 @@
suite: test database token auth env vars
templates:
- gateway/deployment.yaml
- gateway/configmap.yaml
- backend/deployment.yaml
- backend/configmap.yaml
values:
- ./values/required.yaml
tests:
- it: writer emits DATABASE_PASSWORD and no token toggle by default
template: gateway/deployment.yaml
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_PASSWORD
valueFrom:
secretKeyRef:
name: litellm-writer-secret
key: password
any: true
- notContains:
path: spec.template.spec.containers[0].env
content:
name: IAM_TOKEN_DB_AUTH
value: "true"
any: true
- notContains:
path: spec.template.spec.containers[0].env
content:
name: AZURE_POSTGRESQL_AUTH
value: "true"
any: true
- it: writer emits AZURE_POSTGRESQL_AUTH and omits DATABASE_PASSWORD under Entra auth
template: gateway/deployment.yaml
set:
database.writer.useAzureEntraAuth: true
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: AZURE_POSTGRESQL_AUTH
value: "true"
any: true
- notContains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_PASSWORD
any: true
- notContains:
path: spec.template.spec.containers[0].env
content:
name: IAM_TOKEN_DB_AUTH
value: "true"
any: true
- it: backend gets the same Entra toggle as the gateway
template: backend/deployment.yaml
set:
database.writer.useAzureEntraAuth: true
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: AZURE_POSTGRESQL_AUTH
value: "true"
any: true
- it: writer rejects both token sources at once
template: gateway/deployment.yaml
set:
database.writer.useIAMAuth: true
database.writer.useAzureEntraAuth: true
asserts:
- failedTemplate:
errorMessage: "database.writer.useIAMAuth and database.writer.useAzureEntraAuth are mutually exclusive: the database password can only come from one token source"
- it: reader Entra auth without writer Entra auth is rejected
template: gateway/deployment.yaml
set:
database.reader.host: reader.example.com
database.reader.dbname: litellm
database.reader.useAzureEntraAuth: true
asserts:
- failedTemplate:
errorMessage: "database.reader.useAzureEntraAuth requires database.writer.useAzureEntraAuth: true (the proxy gates Entra URL minting on AZURE_POSTGRESQL_AUTH, which is only set by the writer)"
- it: reader under Entra auth omits DATABASE_PASSWORD_READ_REPLICA
template: gateway/deployment.yaml
set:
database.writer.useAzureEntraAuth: true
database.reader.host: reader.example.com
database.reader.dbname: litellm
database.reader.useAzureEntraAuth: true
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_HOST_READ_REPLICA
value: reader.example.com
any: true
- contains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_USER_READ_REPLICA
valueFrom:
secretKeyRef:
name: litellm-reader-secret
key: username
any: true
- notContains:
path: spec.template.spec.containers[0].env
content:
name: DATABASE_PASSWORD_READ_REPLICA
any: true

View file

@ -145,6 +145,8 @@ database:
dbname: ""
schema: ""
useIAMAuth: false
# Azure Database for PostgreSQL with a Microsoft Entra ID token; mutually exclusive with useIAMAuth
useAzureEntraAuth: false
passwordSecret:
name: litellm-writer-secret
usernameKey: username
@ -159,6 +161,8 @@ database:
dbname: ""
schema: ""
useIAMAuth: false
# Azure Database for PostgreSQL with a Microsoft Entra ID token; mutually exclusive with useIAMAuth
useAzureEntraAuth: false
passwordSecret:
name: litellm-reader-secret
usernameKey: username

View file

@ -0,0 +1,5 @@
-- AlterTable
ALTER TABLE "LiteLLM_ShadowEvalJob" ADD COLUMN "max_budget" DOUBLE PRECISION;
-- AlterTable
ALTER TABLE "LiteLLM_ShadowEvalAttempt" ADD COLUMN "shadow_cost" DOUBLE PRECISION NOT NULL DEFAULT 0;

View file

@ -1502,7 +1502,8 @@ model LiteLLM_ShadowEvalJob {
baseline_model String? // reverse only: the fixed model the router is judged against
judge_model String
shadow_percentage Float
max_turns Int // this key's sample budget: judge at most this many turns
max_turns Int // sample-count ceiling: the whole budget on pre-max_budget jobs, the error-loop valve otherwise
max_budget Float? // per-key USD cap on the eval's own shadow + judge spend; null on jobs from before spend budgets
created_at DateTime @default(now())
created_by String?
ends_at DateTime
@ -1525,6 +1526,7 @@ model LiteLLM_ShadowEvalAttempt {
shadow_model String?
confidence Float?
judge_cost Float @default(0)
shadow_cost Float @default(0)
error String?
created_at DateTime @default(now())

View file

@ -1,6 +1,6 @@
[project]
name = "litellm-proxy-extras"
version = "0.4.87"
version = "0.4.88"
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.87"
version = "0.4.88"
version_files = [
"pyproject.toml:^version",
"../pyproject.toml:litellm-proxy-extras==",

View file

@ -295,6 +295,8 @@ fn core_error_kind(error: &CoreError) -> &'static str {
CoreError::Http { .. } => "HttpError",
CoreError::InvalidResponse(_) => "InvalidResponse",
CoreError::Network(_) => "NetworkError",
CoreError::Connect(_) => "ConnectError",
CoreError::Routing(_) => "RoutingError",
CoreError::Unsupported(_) => "UnsupportedRequest",
}
}

View file

@ -324,6 +324,8 @@ fn core_error_kind(error: &CoreError) -> &'static str {
CoreError::Http { .. } => "HttpError",
CoreError::InvalidResponse(_) => "InvalidResponse",
CoreError::Network(_) => "NetworkError",
CoreError::Connect(_) => "ConnectError",
CoreError::Routing(_) => "RoutingError",
CoreError::Unsupported(_) => "UnsupportedRequest",
}
}

View file

@ -105,12 +105,20 @@ impl IntoResponse for MessagesRouteError {
),
CoreError::Http { .. }
| CoreError::Network(_)
| CoreError::Connect(_)
| CoreError::InvalidResponse(_)
| CoreError::InvalidType { .. }
| CoreError::MissingField(_) => (
StatusCode::BAD_GATEWAY,
"messages provider request failed".to_string(),
),
// The gateway has no Python implementation to decline to, so a
// request the core cannot serve is reported to the caller. The
// reason is a fixed internal string, never provider content.
CoreError::Unsupported(reason) => (
StatusCode::BAD_REQUEST,
format!("messages request is not supported: {reason}"),
),
};
(
status,

View file

@ -0,0 +1,15 @@
use std::sync::OnceLock;
use std::time::Duration;
use crate::constants::{CHAT_COMPLETIONS_CONNECT_TIMEOUT_SECS, CHAT_COMPLETIONS_TIMEOUT_SECS};
pub(super) fn http_client() -> &'static reqwest::Client {
static CLIENT: OnceLock<reqwest::Client> = OnceLock::new();
CLIENT.get_or_init(|| {
reqwest::Client::builder()
.timeout(Duration::from_secs(CHAT_COMPLETIONS_TIMEOUT_SECS))
.connect_timeout(Duration::from_secs(CHAT_COMPLETIONS_CONNECT_TIMEOUT_SECS))
.build()
.unwrap_or_else(|_| reqwest::Client::new())
})
}

View file

@ -0,0 +1,28 @@
use serde_json::{Map, Value};
use crate::error::CoreResult;
use crate::http_utils::string_headers as shared_string_headers;
use crate::providers::anthropic::chat_completions::transformation::ANTHROPIC_CHAT_COMPLETIONS_CONFIG;
use super::transformation::ChatCompletionsProviderConfig;
const HEADER_CONTEXT: &str = "chat completions";
pub(super) fn chat_completions_provider_config(
provider: &str,
) -> Option<&'static dyn ChatCompletionsProviderConfig> {
match provider {
"anthropic" => Some(&ANTHROPIC_CHAT_COMPLETIONS_CONFIG),
#[cfg(feature = "bedrock-auth")]
"bedrock" => Some(
&crate::providers::bedrock::chat_completions::transformation::BEDROCK_CHAT_COMPLETIONS_CONFIG,
),
_ => None,
}
}
pub(super) fn string_headers(
extra_headers: Option<Map<String, Value>>,
) -> CoreResult<Vec<(String, String)>> {
shared_string_headers(HEADER_CONTEXT, extra_headers)
}

View file

@ -0,0 +1,254 @@
//! Provider-neutral conversation shape.
//!
//! Both Anthropic Messages and Bedrock Converse want the same thing out of an
//! OpenAI message list: the system prompt lifted out, consecutive same-role
//! turns merged, and text blocks that are never empty. That normalization is
//! shared here so a provider config only renders the result into its own wire
//! shape.
//!
//! Mirrors Python's `anthropic_messages_pt` /
//! `_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};
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum TurnRole {
User,
Assistant,
}
impl TurnRole {
pub fn as_str(self) -> &'static str {
match self {
Self::User => "user",
Self::Assistant => "assistant",
}
}
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Turn {
pub role: TurnRole,
pub texts: Vec<String>,
}
#[derive(Clone, Debug, Default, PartialEq, Eq)]
pub struct Conversation {
pub system: Vec<String>,
pub turns: Vec<Turn>,
}
/// True when the conversation can be sent as-is.
///
/// Python inserts a placeholder first user turn only under
/// `litellm.modify_params`, which the core cannot see, so a conversation that
/// does not open on a user turn is declined rather than guessed at.
impl Conversation {
pub fn opens_on_user_turn(&self) -> bool {
self.turns
.first()
.is_some_and(|turn| turn.role == TurnRole::User)
}
}
fn message_texts(content: &ChatMessageContent) -> Vec<String> {
match content {
ChatMessageContent::Text(text) => vec![text.clone()],
ChatMessageContent::Parts(parts) => parts
.iter()
.filter_map(|part| part.get("text").and_then(|text| text.as_str()))
.map(str::to_string)
.collect(),
}
}
/// Python rewrites empty or whitespace-only text rather than dropping it, so an
/// entirely empty content list never reaches a provider that rejects one.
fn sanitize(text: String) -> String {
if text.trim().is_empty() {
return EMPTY_TEXT_PLACEHOLDER.to_string();
}
text
}
pub fn build_conversation(messages: &[ChatMessage]) -> Conversation {
let system = messages
.iter()
.filter(|message| message.role == "system")
.filter_map(|message| message.content.as_ref())
.flat_map(message_texts)
.filter(|text| !text.is_empty())
.collect();
let turns = messages
.iter()
.filter(|message| message.role != "system")
.fold(Vec::<Turn>::new(), |mut turns, message| {
let role = if message.role == "assistant" {
TurnRole::Assistant
} else {
TurnRole::User
};
let texts = message
.content
.as_ref()
.map(message_texts)
.unwrap_or_default()
.into_iter()
.map(sanitize);
match turns.last_mut() {
Some(last) if last.role == role => last.texts.extend(texts),
_ => turns.push(Turn {
role,
texts: texts.collect(),
}),
}
turns
});
// Anthropic and Bedrock both reject trailing whitespace on the final
// assistant turn, so Python right-strips it there; mirror that exactly.
let turns = match turns.split_last() {
Some((last, rest)) if last.role == TurnRole::Assistant => rest
.iter()
.cloned()
.chain([Turn {
role: last.role,
texts: last
.texts
.iter()
.map(|text| text.trim_end().to_string())
.collect(),
}])
.collect(),
_ => turns,
};
Conversation { system, turns }
}
#[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
fn messages(value: serde_json::Value) -> Vec<ChatMessage> {
serde_json::from_value(value).expect("valid messages")
}
#[test]
fn lifts_system_messages_out_of_the_turn_list() {
let conversation = build_conversation(&messages(json!([
{"role": "system", "content": "be terse"},
{"role": "user", "content": "hi"}
])));
assert_eq!(conversation.system, vec!["be terse".to_string()]);
assert_eq!(
conversation.turns,
vec![Turn {
role: TurnRole::User,
texts: vec!["hi".to_string()]
}]
);
}
#[test]
fn merges_consecutive_same_role_turns() {
let conversation = build_conversation(&messages(json!([
{"role": "user", "content": "one"},
{"role": "user", "content": "two"},
{"role": "assistant", "content": "ack"},
{"role": "user", "content": "three"}
])));
assert_eq!(
conversation.turns,
vec![
Turn {
role: TurnRole::User,
texts: vec!["one".to_string(), "two".to_string()]
},
Turn {
role: TurnRole::Assistant,
texts: vec!["ack".to_string()]
},
Turn {
role: TurnRole::User,
texts: vec!["three".to_string()]
},
]
);
}
#[test]
fn flattens_text_parts_in_order() {
let conversation = build_conversation(&messages(json!([
{"role": "user", "content": [
{"type": "text", "text": "first"},
{"type": "text", "text": "second"}
]}
])));
assert_eq!(
conversation.turns[0].texts,
vec!["first".to_string(), "second".to_string()]
);
}
#[test]
fn rewrites_empty_and_whitespace_only_text_to_the_python_placeholder() {
let conversation = build_conversation(&messages(json!([
{"role": "user", "content": ""},
{"role": "assistant", "content": " "},
{"role": "user", "content": "real"}
])));
assert_eq!(conversation.turns[0].texts, vec![EMPTY_TEXT_PLACEHOLDER]);
assert_eq!(conversation.turns[1].texts, vec![EMPTY_TEXT_PLACEHOLDER]);
}
#[test]
fn right_strips_only_the_final_assistant_turn() {
let conversation = build_conversation(&messages(json!([
{"role": "user", "content": "hi"},
{"role": "assistant", "content": "kept "},
{"role": "user", "content": "more"},
{"role": "assistant", "content": "stripped "}
])));
assert_eq!(conversation.turns[1].texts, vec!["kept ".to_string()]);
assert_eq!(conversation.turns[3].texts, vec!["stripped".to_string()]);
}
#[test]
fn does_not_strip_when_the_last_turn_is_a_user_turn() {
let conversation = build_conversation(&messages(json!([
{"role": "assistant", "content": "kept "},
{"role": "user", "content": "hi "}
])));
assert_eq!(conversation.turns[0].texts, vec!["kept ".to_string()]);
assert_eq!(conversation.turns[1].texts, vec!["hi ".to_string()]);
}
#[test]
fn reports_whether_the_conversation_opens_on_a_user_turn() {
assert!(
build_conversation(&messages(json!([{"role": "user", "content": "hi"}])))
.opens_on_user_turn()
);
assert!(
!build_conversation(&messages(json!([{"role": "assistant", "content": "hi"}])))
.opens_on_user_turn()
);
assert!(!Conversation::default().opens_on_user_turn());
}
#[test]
fn drops_empty_system_text_the_way_python_skips_empty_system_blocks() {
let conversation = build_conversation(&messages(json!([
{"role": "system", "content": ""},
{"role": "system", "content": "kept"},
{"role": "user", "content": "hi"}
])));
assert_eq!(conversation.system, vec!["kept".to_string()]);
}
}

View file

@ -0,0 +1,147 @@
use serde_json::Value;
use crate::error::{CoreError, CoreResult};
use crate::http_utils::truncate_error_body;
use super::client::http_client;
use super::transformation::ChatCompletionsAuth;
use super::types::{
ChatCompletionsResponse, ProviderChatCompletionsRequest, ProviderChatResponseData,
};
pub(super) async fn execute_chat_completions_provider_call(
request: ProviderChatCompletionsRequest,
) -> CoreResult<ChatCompletionsResponse> {
let body = serde_json::to_vec(&request.body).map_err(|err| {
CoreError::InvalidRequest(format!(
"failed to serialize chat completions request: {err}"
))
})?;
let headers = signed_headers(&request, &body).await?;
let mut request_builder = http_client().post(&request.url).body(body);
for (key, value) in &headers {
request_builder = request_builder.header(key, value);
}
if let Some(duration) = request.timeout {
request_builder = request_builder.timeout(duration);
}
let response = request_builder.send().await.map_err(|err| {
// Failing to establish the connection means the request never went out,
// 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() {
CoreError::Connect(err.to_string())
} else {
CoreError::Network(err.to_string())
}
})?;
let status = response.status();
let text = response
.text()
.await
.map_err(|err| CoreError::Network(err.to_string()))?;
if !status.is_success() {
return Err(CoreError::Http {
status: status.as_u16(),
body: truncate_error_body(&text),
});
}
let body: Value = serde_json::from_str(&text).map_err(|err| {
CoreError::InvalidResponse(format!("invalid chat completions response JSON: {err}"))
})?;
request
.config
.transform_response(&request.model, ProviderChatResponseData { body })
.map_err(as_response_error)
}
/// Re-tag an error raised while normalizing a response the provider already
/// returned.
///
/// A config reports the same variants on either side of the call: a missing
/// field or an unsupported block can mean "this request cannot be translated"
/// during prepare and "this response cannot be normalized" here. Only the
/// second kind has already been billed, and a host that keeps a reference
/// implementation must not retry those, so collapse them to one variant that
/// can only mean the provider was already called.
pub(super) fn as_response_error(err: CoreError) -> CoreError {
match err {
already @ (CoreError::InvalidResponse(_) | CoreError::Http { .. }) => already,
other => CoreError::InvalidResponse(other.to_string()),
}
}
#[cfg(feature = "bedrock-auth")]
pub(super) async fn signed_headers(
request: &ProviderChatCompletionsRequest,
body: &[u8],
) -> CoreResult<Vec<(String, String)>> {
use std::collections::BTreeMap;
use std::time::SystemTime;
use crate::providers::bedrock::aws_base::{
aws_auth_config, aws_signature_headers, host_supplied_credentials,
is_sigv4_computed_header, resolve_credentials, sign_bedrock_post,
};
let ChatCompletionsAuth::AwsSigV4 { region } = &request.auth else {
return Ok(request.upstream_headers.clone());
};
// Reattaching a header the signer also emits would put both copies on the
// wire, and Bedrock rejects that pair. Python instead drops the caller's
// copy and prefers a forwarded Authorization over the signature, so leave
// the request to Python rather than serving it a different way here.
if request
.upstream_headers
.iter()
.any(|(name, _)| is_sigv4_computed_header(name))
{
return Err(CoreError::Unsupported(
"request forwards a header AWS SigV4 computes",
));
}
let env_lookup = |key: &str| std::env::var(key).ok();
let unsigned: BTreeMap<String, String> = request.upstream_headers.iter().cloned().collect();
// A host with its own resolution chain hands the result down; only fall
// back to deriving credentials here when it supplied none.
let credentials = match host_supplied_credentials(&request.optional_params) {
Some(credentials) => credentials,
None => {
resolve_credentials(
aws_auth_config(&request.optional_params, &env_lookup),
&env_lookup,
)
.await?
}
};
let signature = sign_bedrock_post(
&request.url,
body,
&aws_signature_headers(&unsigned),
region,
&credentials,
SystemTime::now(),
)?;
// Every original header goes back on the wire alongside the computed ones,
// as Python reattaches them. The guard above already rejected the names
// 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],
) -> CoreResult<Vec<(String, String)>> {
match &request.auth {
ChatCompletionsAuth::AwsSigV4 { .. } => Err(CoreError::Unsupported(
"AWS SigV4 requires the bedrock-auth feature",
)),
_ => Ok(request.upstream_headers.clone()),
}
}

View file

@ -0,0 +1,59 @@
//! The `/chat/completions` call, the Rust equivalent of Python's
//! `litellm.completion()`.
//!
//! [`chat_completions`] is the top-level entrypoint: give it a model, the
//! OpenAI-shaped message list, the provider-mapped optional params, and
//! credentials, and it resolves the provider, translates the conversation,
//! calls the provider, and returns a typed OpenAI-shaped response.
mod client;
mod common_utils;
pub mod conversation;
pub(crate) mod handler;
mod prepare;
pub mod response_utils;
pub mod transformation;
pub mod types;
use serde_json::{Map, Value};
use crate::error::CoreResult;
use handler::execute_chat_completions_provider_call;
use prepare::{parse_messages, prepare_chat_completions_call, resolve_provider_config};
use types::{ChatCompletionsRequest, ChatCompletionsResponse};
pub async fn chat_completions(
request: ChatCompletionsRequest<'_>,
) -> CoreResult<ChatCompletionsResponse> {
execute_chat_completions_provider_call(prepare_chat_completions_call(request)?).await
}
/// Whether the core would accept this request, without resolving credentials or
/// touching the network.
///
/// A host that keeps the Python implementation asks this first so it can emit
/// its pre-call logging exactly once, on whichever path is about to run.
/// Returns the decline reason, or `None` when the request is accepted.
pub fn chat_completions_decline_reason(
model: &str,
custom_llm_provider: Option<&str>,
messages: Value,
optional_params: &Map<String, Value>,
) -> Option<&'static str> {
let Ok((_, config)) = resolve_provider_config(model, custom_llm_provider) else {
return Some("provider is not on the rust chat completions path");
};
let Ok(messages) = parse_messages(messages) else {
return Some("unreadable message list");
};
if messages.is_empty() {
return Some("empty message list");
}
config
.unsupported_reason(&messages, optional_params)
.map(|reason| reason.0)
}
#[cfg(test)]
mod tests;

View file

@ -0,0 +1,118 @@
use serde_json::Value;
use crate::error::{CoreError, CoreResult};
use crate::http_utils::has_header;
use crate::routing_utils::provider::{CustomLlmProvider, get_custom_llm_provider};
use super::common_utils::{chat_completions_provider_config, string_headers};
use super::transformation::{ChatCompletionsAuth, ChatCompletionsProviderConfig};
use super::types::{ChatCompletionsRequest, ChatMessage, ProviderChatCompletionsRequest};
pub(super) fn resolve_provider_config<'a>(
model: &'a str,
custom_llm_provider: Option<&'a str>,
) -> CoreResult<(String, &'static dyn ChatCompletionsProviderConfig)> {
let provider_info = get_custom_llm_provider(model, custom_llm_provider)
.or_else(|| {
custom_llm_provider.map(|provider| CustomLlmProvider {
model,
custom_llm_provider: provider,
})
})
.ok_or_else(|| {
CoreError::InvalidProvider(
"unable to resolve custom_llm_provider for chat completions request".to_string(),
)
})?;
let config = chat_completions_provider_config(provider_info.custom_llm_provider)
.ok_or_else(|| CoreError::InvalidProvider(provider_info.custom_llm_provider.to_string()))?;
Ok((provider_info.model.to_string(), config))
}
pub(super) fn parse_messages(messages: Value) -> CoreResult<Vec<ChatMessage>> {
serde_json::from_value(messages).map_err(|err| {
CoreError::InvalidRequest(format!("invalid chat completions messages: {err}"))
})
}
pub(super) fn prepare_chat_completions_call(
request: ChatCompletionsRequest<'_>,
) -> CoreResult<ProviderChatCompletionsRequest> {
let (model, config) = resolve_provider_config(request.model, request.custom_llm_provider)?;
let env_lookup = |key: &str| std::env::var(key).ok();
let messages = parse_messages(request.messages)?;
if messages.is_empty() {
return Err(CoreError::InvalidRequest(
"chat completions requires at least one message".to_string(),
));
}
if let Some(reason) = config.unsupported_reason(&messages, &request.optional_params) {
return Err(CoreError::Unsupported(reason.0));
}
let mut headers = string_headers(request.extra_headers)?;
let auth = config.auth(
request.api_key,
&model,
&request.optional_params,
&env_lookup,
)?;
match &auth {
ChatCompletionsAuth::Header { name, value } => {
// The deployment's credential replaces whatever the caller forwarded
// under the same name, mirroring Python's
// `{**headers, **anthropic_headers}`: letting a request header win
// would let its sender choose the principal the call bills to.
//
// The exception is a scheme the provider hands off to entirely, such
// as an Anthropic OAuth bearer, where Python drops `x-api-key`
// instead of resolving one. Re-adding it there would put the
// credential into a header the host removed on purpose.
if !config.defers_to_forwarded_auth(&headers) {
headers.retain(|(header, _)| !header.eq_ignore_ascii_case(name));
headers.push(((*name).to_string(), value.clone()));
}
}
ChatCompletionsAuth::Bearer { token } => {
// Bedrock's `get_request_headers` assigns `headers["Authorization"]`
// unconditionally once a bearer token resolves, so the deployment's
// identity outranks whatever the caller forwarded. Keeping the
// caller's would bill and authorize the call as a different
// principal than the same deployment uses on Python.
//
// The `Header` arm below keeps the opposite precedence on purpose:
// Anthropic's transform honours a forwarded OAuth bearer.
headers.retain(|(name, _)| !name.eq_ignore_ascii_case("authorization"));
headers.push(("authorization".to_string(), format!("Bearer {token}")));
}
// SigV4 signs the serialized body, so the handler adds its headers.
ChatCompletionsAuth::AwsSigV4 { .. } => {}
}
for (name, value) in config.default_headers() {
if !has_header(&headers, name) {
headers.push(((*name).to_string(), (*value).to_string()));
}
}
let url = config.complete_url(
request.api_base,
&model,
&request.optional_params,
&env_lookup,
)?;
let transformed =
config.transform_request(&model, messages, request.optional_params.clone())?;
Ok(ProviderChatCompletionsRequest {
model,
config,
url,
body: transformed.body,
upstream_headers: headers,
auth,
optional_params: request.optional_params,
timeout: request.timeout,
})
}

View file

@ -0,0 +1,101 @@
//! Response normalization shared by every chat completions provider config.
use std::time::{SystemTime, UNIX_EPOCH};
use super::types::{ChatCompletionsUsage, PromptTokensDetails};
/// OpenAI finish reasons, mirroring Python's `_FINISH_REASON_MAP` for the
/// reasons the providers on this route can emit. Python warns and falls back to
/// `stop` for anything unmapped, so do the same.
const FINISH_REASONS: &[(&str, &str)] = &[
("end_turn", "stop"),
("stop_sequence", "stop"),
("max_tokens", "length"),
("refusal", "content_filter"),
("compaction", "length"),
("guardrail_intervened", "content_filter"),
("content_filtered", "content_filter"),
("content_filter", "content_filter"),
("stop", "stop"),
("length", "length"),
];
pub fn finish_reason_for(provider_reason: &str) -> &'static str {
FINISH_REASONS
.iter()
.find(|(reason, _)| *reason == provider_reason)
.map_or("stop", |(_, mapped)| *mapped)
}
/// Python folds cache tokens into `prompt_tokens` and reports the split under
/// `prompt_tokens_details`; mirror that so cost tracking agrees on both paths.
pub fn usage_from_parts(
input_tokens: u64,
output_tokens: u64,
cache_read_tokens: u64,
cache_creation_tokens: u64,
) -> ChatCompletionsUsage {
let prompt_tokens = input_tokens + cache_read_tokens + cache_creation_tokens;
ChatCompletionsUsage {
prompt_tokens,
completion_tokens: output_tokens,
total_tokens: prompt_tokens + output_tokens,
prompt_tokens_details: PromptTokensDetails {
cached_tokens: cache_read_tokens,
cache_creation_tokens,
text_tokens: input_tokens,
},
}
}
pub fn unix_now() -> u64 {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.map_or(0, |elapsed| elapsed.as_secs())
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn maps_every_reason_the_route_can_observe() {
assert_eq!(finish_reason_for("end_turn"), "stop");
assert_eq!(finish_reason_for("stop_sequence"), "stop");
assert_eq!(finish_reason_for("max_tokens"), "length");
assert_eq!(finish_reason_for("refusal"), "content_filter");
assert_eq!(finish_reason_for("guardrail_intervened"), "content_filter");
// Converse emits these two, and folding them into `stop` would report a
// filtered completion as a normal one.
assert_eq!(finish_reason_for("content_filtered"), "content_filter");
assert_eq!(finish_reason_for("content_filter"), "content_filter");
}
#[test]
fn defaults_an_unmapped_reason_to_stop_like_python() {
// Python warns and falls back to `stop` for a reason its own map does
// not carry, so only a reason absent from `_FINISH_REASON_MAP` belongs
// here.
assert_eq!(finish_reason_for("something_new"), "stop");
assert_eq!(finish_reason_for(""), "stop");
}
#[test]
fn folds_cache_tokens_into_prompt_tokens() {
let usage = usage_from_parts(10, 4, 7, 3);
assert_eq!(usage.prompt_tokens, 20);
assert_eq!(usage.completion_tokens, 4);
assert_eq!(usage.total_tokens, 24);
assert_eq!(usage.prompt_tokens_details.cached_tokens, 7);
assert_eq!(usage.prompt_tokens_details.cache_creation_tokens, 3);
assert_eq!(usage.prompt_tokens_details.text_tokens, 10);
}
#[test]
fn reports_raw_input_tokens_when_no_cache_is_involved() {
let usage = usage_from_parts(12, 5, 0, 0);
assert_eq!(usage.prompt_tokens, 12);
assert_eq!(usage.total_tokens, 17);
assert_eq!(usage.prompt_tokens_details.text_tokens, 12);
}
}

View file

@ -0,0 +1,820 @@
use serde_json::{Map, Value, json};
use crate::error::CoreError;
use super::prepare::prepare_chat_completions_call;
use super::transformation::ChatCompletionsAuth;
use super::types::ChatCompletionsRequest;
fn request<'a>(
model: &'a str,
provider: Option<&'a str>,
messages: Value,
optional_params: Value,
) -> ChatCompletionsRequest<'a> {
ChatCompletionsRequest {
model,
messages,
optional_params: match optional_params {
Value::Object(map) => map,
other => panic!("params must be an object, got {other}"),
},
api_key: Some("sk-test"),
api_base: None,
custom_llm_provider: provider,
extra_headers: None,
timeout: None,
}
}
/// `ProviderChatCompletionsRequest` deliberately has no `Debug` (its headers
/// carry resolved credentials), so unwrap the failure case by hand.
fn decline(request: ChatCompletionsRequest<'_>) -> CoreError {
match prepare_chat_completions_call(request) {
Err(error) => error,
Ok(prepared) => panic!("expected a decline, prepared a call to {}", prepared.url),
}
}
#[test]
fn resolves_the_provider_from_the_model_prefix() {
let prepared = prepare_chat_completions_call(request(
"anthropic/claude-sonnet-4-5",
None,
json!([{"role": "user", "content": "hi"}]),
json!({"max_tokens": 16}),
))
.expect("prepares");
assert_eq!(prepared.model, "claude-sonnet-4-5");
assert_eq!(prepared.url, "https://api.anthropic.com/v1/messages");
assert_eq!(prepared.body["model"], json!("claude-sonnet-4-5"));
}
#[test]
fn strips_an_explicit_provider_prefix_from_the_model() {
let prepared = prepare_chat_completions_call(request(
"anthropic/claude-sonnet-4-5",
Some("anthropic"),
json!([{"role": "user", "content": "hi"}]),
json!({}),
))
.expect("prepares");
assert_eq!(prepared.model, "claude-sonnet-4-5");
}
#[test]
fn adds_the_auth_and_default_headers() {
let prepared = prepare_chat_completions_call(request(
"claude-sonnet-4-5",
Some("anthropic"),
json!([{"role": "user", "content": "hi"}]),
json!({}),
))
.expect("prepares");
assert!(
prepared
.upstream_headers
.contains(&("x-api-key".to_string(), "sk-test".to_string()))
);
assert!(
prepared
.upstream_headers
.contains(&("anthropic-version".to_string(), "2023-06-01".to_string()))
);
assert!(matches!(
prepared.auth,
ChatCompletionsAuth::Header {
name: "x-api-key",
..
}
));
}
#[test]
fn the_deployment_credential_replaces_a_caller_supplied_auth_header() {
// Python builds `{**headers, **anthropic_headers}`, so the deployment's key
// overwrites a forwarded one. Honouring the caller's would let whoever sends
// the request choose the Anthropic principal it bills to.
let mut call = request(
"claude-sonnet-4-5",
Some("anthropic"),
json!([{"role": "user", "content": "hi"}]),
json!({}),
);
call.extra_headers = Some(Map::from_iter([(
"X-Api-Key".to_string(),
json!("sk-caller"),
)]));
let prepared = prepare_chat_completions_call(call).expect("prepares");
let keys: Vec<_> = prepared
.upstream_headers
.iter()
.filter(|(name, _)| name.eq_ignore_ascii_case("x-api-key"))
.collect();
assert_eq!(keys.len(), 1, "got {:?}", prepared.upstream_headers);
assert_eq!(keys[0].1, "sk-test");
}
#[test]
fn a_forwarded_authorization_header_suppresses_the_resolved_api_key_header() {
// Anthropic's `validate_environment` pops `x-api-key` and sets `authorization`
// for an OAuth token, so re-adding the key here would put the credential into
// a header the host removed on purpose.
let mut call = request(
"claude-sonnet-4-5",
Some("anthropic"),
json!([{"role": "user", "content": "hi"}]),
json!({}),
);
call.extra_headers = Some(Map::from_iter([
(
"Authorization".to_string(),
json!("Bearer sk-ant-oat01-token"),
),
("X-Api-Key".to_string(), json!("sk-caller")),
]));
let prepared = prepare_chat_completions_call(call).expect("prepares");
assert!(
!prepared
.upstream_headers
.iter()
.any(|(name, value)| name.eq_ignore_ascii_case("x-api-key") && value == "sk-test"),
"the resolved key must not be applied over an OAuth bearer, got {:?}",
prepared.upstream_headers
);
assert!(
prepared
.upstream_headers
.iter()
.any(|(name, value)| name.eq_ignore_ascii_case("authorization")
&& value == "Bearer sk-ant-oat01-token")
);
}
#[test]
fn an_unrelated_forwarded_authorization_does_not_defer_the_resolved_key() {
// Only an OAuth bearer replaces the credential. Python sends the deployment's
// `x-api-key` alongside any other forwarded `authorization`, so deferring on
// the mere presence of that header would drop the deployment's auth.
let mut call = request(
"claude-sonnet-4-5",
Some("anthropic"),
json!([{"role": "user", "content": "hi"}]),
json!({}),
);
call.extra_headers = Some(Map::from_iter([
("Authorization".to_string(), json!("Bearer unrelated")),
("X-Api-Key".to_string(), json!("sk-caller")),
]));
let prepared = prepare_chat_completions_call(call).expect("prepares");
let keys: Vec<_> = prepared
.upstream_headers
.iter()
.filter(|(name, _)| name.eq_ignore_ascii_case("x-api-key"))
.collect();
assert_eq!(keys.len(), 1, "got {:?}", prepared.upstream_headers);
assert_eq!(keys[0].1, "sk-test");
assert!(
prepared
.upstream_headers
.iter()
.any(|(name, value)| name.eq_ignore_ascii_case("authorization")
&& value == "Bearer unrelated"),
"the unrelated authorization must survive, got {:?}",
prepared.upstream_headers
);
}
#[test]
fn declines_an_unsupported_request_before_resolving_credentials() {
let mut call = request(
"claude-sonnet-4-5",
Some("anthropic"),
json!([{"role": "user", "content": "hi"}]),
json!({"stream": true}),
);
call.api_key = None;
// No api_key is set and no env is consulted: the gate must run first, so the
// error is the decline rather than a missing-credential error.
assert_eq!(decline(call), CoreError::Unsupported("streaming"));
}
#[test]
fn rejects_an_unknown_provider() {
assert_eq!(
decline(request(
"openai/gpt-4o",
None,
json!([{"role": "user", "content": "hi"}]),
json!({}),
)),
CoreError::InvalidProvider("openai".to_string())
);
}
#[test]
fn rejects_a_model_with_no_resolvable_provider() {
assert!(matches!(
decline(request(
"claude-sonnet-4-5",
None,
json!([{"role": "user", "content": "hi"}]),
json!({}),
)),
CoreError::InvalidProvider(_)
));
}
#[test]
fn rejects_an_empty_or_malformed_message_list() {
assert_eq!(
decline(request(
"anthropic/claude-sonnet-4-5",
None,
json!([]),
json!({}),
)),
CoreError::InvalidRequest("chat completions requires at least one message".to_string())
);
assert!(matches!(
decline(request(
"anthropic/claude-sonnet-4-5",
None,
json!("not a list"),
json!({}),
)),
CoreError::InvalidRequest(_)
));
}
#[test]
fn rejects_non_string_extra_headers() {
let mut call = request(
"anthropic/claude-sonnet-4-5",
None,
json!([{"role": "user", "content": "hi"}]),
json!({}),
);
call.extra_headers = Some(Map::from_iter([("x-trace".to_string(), json!(7))]));
assert_eq!(
decline(call),
CoreError::InvalidRequest(
"chat completions extra_headers.x-trace must be a string, got number".to_string()
)
);
}
#[cfg(feature = "bedrock-auth")]
#[test]
fn prepares_a_bedrock_call_without_resolving_credentials() {
let mut call = request(
"bedrock/us-east-1/anthropic.claude-v2",
None,
json!([{"role": "user", "content": "hi"}]),
json!({"maxTokens": 16}),
);
call.api_key = None;
let prepared = prepare_chat_completions_call(call).expect("prepares");
assert_eq!(
prepared.url,
"https://bedrock-runtime.us-east-1.amazonaws.com/model/anthropic.claude-v2/converse"
);
assert_eq!(
prepared.auth,
ChatCompletionsAuth::AwsSigV4 {
region: "us-east-1".to_string()
}
);
// SigV4 signs the serialized body, so prepare must not have added an
// Authorization header; the handler does it.
assert!(
!prepared
.upstream_headers
.iter()
.any(|(name, _)| name.eq_ignore_ascii_case("authorization"))
);
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
// the caller forwarded rides along without joining the canonical request.
// Signing it makes Converse 403 on a deployment that works on Python.
let mut call = request(
"bedrock/us-east-1/anthropic.claude-v2",
None,
json!([{"role": "user", "content": "hi"}]),
json!({
"maxTokens": 16,
"aws_access_key_id": "AKIDEXAMPLE",
"aws_secret_access_key": "wJalrXUtnFEMI/K7MDENG+bPxRfiCYEXAMPLEKEY"
}),
);
// A key would resolve to a bearer token and never reach the signer.
call.api_key = None;
call.extra_headers = Some(Map::from_iter([(
"x-request-id".to_string(),
json!("abc-123"),
)]));
let prepared = prepare_chat_completions_call(call).expect("prepares");
let signed = super::handler::signed_headers(&prepared, br#"{"a":1}"#)
.await
.expect("signs");
let authorization = signed
.iter()
.find(|(name, _)| name.eq_ignore_ascii_case("authorization"))
.map(|(_, value)| value.clone())
.expect("carries an authorization header");
assert!(
authorization.starts_with("AWS4-HMAC-SHA256"),
"expected a SigV4 signature, got {authorization}"
);
assert!(
!authorization.contains("x-request-id"),
"forwarded header reached SignedHeaders: {authorization}"
);
// It still goes on the wire, it is just not part of the signature.
assert!(
signed
.iter()
.any(|(name, value)| name == "x-request-id" && value == "abc-123"),
"forwarded header was dropped instead of reattached"
);
}
#[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
// the wire twice and Bedrock rejects the pair, so a request carrying one
// has to go to Python instead of being signed here.
for forwarded in [
"Authorization",
"x-amz-date",
"x-amz-security-token",
"Date",
] {
let mut call = request(
"bedrock/us-east-1/anthropic.claude-v2",
None,
json!([{"role": "user", "content": "hi"}]),
json!({
"maxTokens": 16,
"aws_access_key_id": "AKIDEXAMPLE",
"aws_secret_access_key": "wJalrXUtnFEMI/K7MDENG+bPxRfiCYEXAMPLEKEY"
}),
);
call.api_key = None;
call.extra_headers = Some(Map::from_iter([(forwarded.to_string(), json!("forged"))]));
let prepared = prepare_chat_completions_call(call).expect("prepares");
let error = super::handler::signed_headers(&prepared, br#"{"a":1}"#)
.await
.expect_err("{forwarded} should decline instead of being signed");
assert!(
matches!(error, CoreError::Unsupported(_)),
"{forwarded} declined as {error:?}, which the host would not fall back on"
);
}
}
#[cfg(feature = "bedrock-auth")]
#[test]
fn a_bedrock_deployment_bearer_outranks_a_forwarded_authorization() {
// `get_request_headers` assigns `headers["Authorization"]` unconditionally
// once a bearer token resolves, so the deployment's identity wins on
// Python. Keeping the caller's would authorize and bill the call as a
// different principal, and only when the deployment carries `rust: true`.
let mut call = request(
"bedrock/us-east-1/anthropic.claude-v2",
None,
json!([{"role": "user", "content": "hi"}]),
json!({"maxTokens": 16}),
);
call.extra_headers = Some(Map::from_iter([(
"Authorization".to_string(),
json!("Bearer caller-supplied"),
)]));
let prepared = prepare_chat_completions_call(call).expect("prepares");
let authorizations: Vec<_> = prepared
.upstream_headers
.iter()
.filter(|(name, _)| name.eq_ignore_ascii_case("authorization"))
.map(|(_, value)| value.as_str())
.collect();
assert_eq!(
authorizations,
vec!["Bearer sk-test"],
"the deployment token must be the only authorization on the wire"
);
}
#[test]
fn an_anthropic_forwarded_oauth_bearer_still_outranks_the_resolved_key() {
// The opposite precedence, and deliberate: Anthropic's own transform
// honours a forwarded OAuth bearer, so the Bedrock fix above must not be
// generalized into a rule that the configured key always wins.
//
// An OAuth bearer is the whole of that exception. This forwarded a plain
// `x-api-key` until round 17, which read as the same claim and was not:
// Python overwrites a forwarded `x-api-key` with the deployment's.
let mut call = request(
"claude-sonnet-4-5",
Some("anthropic"),
json!([{"role": "user", "content": "hi"}]),
json!({}),
);
call.extra_headers = Some(Map::from_iter([(
"authorization".to_string(),
json!("Bearer sk-ant-oat01-forwarded"),
)]));
let prepared = prepare_chat_completions_call(call).expect("prepares");
let keys: Vec<_> = prepared
.upstream_headers
.iter()
.filter(|(name, _)| name.eq_ignore_ascii_case("x-api-key"))
.map(|(_, value)| value.as_str())
.collect();
assert!(keys.is_empty(), "got {:?}", prepared.upstream_headers);
assert!(
prepared
.upstream_headers
.iter()
.any(|(name, value)| name.eq_ignore_ascii_case("authorization")
&& value == "Bearer sk-ant-oat01-forwarded")
);
}
#[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,
// so a request carrying one must not be signed as whatever principal the
// host's AWS credentials resolve to.
let prepared = prepare_chat_completions_call(request(
"bedrock/us-east-1/anthropic.claude-v2",
None,
json!([{"role": "user", "content": "hi"}]),
json!({"maxTokens": 16}),
))
.expect("prepares");
assert_eq!(
prepared.auth,
ChatCompletionsAuth::Bearer {
token: "sk-test".to_string()
}
);
assert!(
prepared
.upstream_headers
.iter()
.any(|(name, value)| name.eq_ignore_ascii_case("authorization")
&& value == "Bearer sk-test"),
"prepare did not carry the bearer token"
);
}
fn decline_reason(
model: &str,
provider: Option<&str>,
messages: Value,
params: Value,
) -> Option<&'static str> {
let params = match params {
Value::Object(map) => map,
other => panic!("params must be an object, got {other}"),
};
super::chat_completions_decline_reason(model, provider, messages, &params)
}
#[test]
fn the_gate_accepts_what_prepare_accepts() {
assert_eq!(
decline_reason(
"anthropic/claude-sonnet-4-5",
None,
json!([{"role": "user", "content": "hi"}]),
json!({"max_tokens": 16}),
),
None
);
}
#[test]
fn the_gate_declines_without_resolving_credentials_or_calling_out() {
assert_eq!(
decline_reason(
"anthropic/claude-sonnet-4-5",
None,
json!([{"role": "user", "content": "hi"}]),
json!({"stream": true}),
),
Some("streaming")
);
assert_eq!(
decline_reason(
"openai/gpt-4o",
None,
json!([{"role": "user", "content": "hi"}]),
json!({}),
),
Some("provider is not on the rust chat completions path")
);
assert_eq!(
decline_reason(
"claude-sonnet-4-5",
None,
json!([{"role": "user", "content": "hi"}]),
json!({}),
),
Some("provider is not on the rust chat completions path")
);
assert_eq!(
decline_reason(
"anthropic/claude-sonnet-4-5",
None,
json!("nope"),
json!({})
),
Some("unreadable message list")
);
assert_eq!(
decline_reason("anthropic/claude-sonnet-4-5", None, json!([]), json!({})),
Some("empty message list")
);
}
#[test]
fn the_gate_agrees_with_prepare_on_every_case_it_accepts() {
// A gate that accepts what prepare then declines would make the host emit
// its pre-call logging on a path that falls back, so pin the agreement.
for (messages, params) in [
(
json!([{"role": "user", "content": "hi"}]),
json!({"max_tokens": 8}),
),
(
json!([{"role": "system", "content": "s"}, {"role": "user", "content": "hi"}]),
json!({"temperature": 0.1}),
),
(
json!([{"role": "user", "content": "hi"}, {"role": "assistant", "content": "yo"}]),
json!({}),
),
] {
assert_eq!(
decline_reason(
"anthropic/claude-sonnet-4-5",
None,
messages.clone(),
params.clone()
),
None,
"gate declined {messages}"
);
prepare_chat_completions_call(request(
"anthropic/claude-sonnet-4-5",
None,
messages.clone(),
params,
))
.unwrap_or_else(|error| panic!("prepare declined {messages}: {error}"));
}
}
mod round_trip {
use super::*;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::{TcpListener, TcpStream};
use crate::chat_completions::chat_completions;
async fn read_http_request(socket: &mut TcpStream) -> String {
let mut request = Vec::new();
let mut buffer = [0_u8; 1024];
let header_end = loop {
let n = socket.read(&mut buffer).await.expect("reads request");
if n == 0 {
break request.len();
}
request.extend_from_slice(&buffer[..n]);
if let Some(position) = request.windows(4).position(|window| window == b"\r\n\r\n") {
break position + 4;
}
};
let headers = String::from_utf8_lossy(&request[..header_end]);
let content_length = headers
.lines()
.find_map(|line| {
let (name, value) = line.split_once(':')?;
name.eq_ignore_ascii_case("content-length")
.then(|| value.trim().parse::<usize>().ok())
.flatten()
})
.unwrap_or(0);
while request.len().saturating_sub(header_end) < content_length {
let n = socket.read(&mut buffer).await.expect("reads body");
if n == 0 {
break;
}
request.extend_from_slice(&buffer[..n]);
}
String::from_utf8(request).expect("request is utf8")
}
fn http_response(status: &str, body: &str) -> String {
format!(
"HTTP/1.1 {status}\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}",
body.len(),
body
)
}
/// Serve one request from a stub upstream and hand back what it received.
async fn serve_once(
status: &'static str,
body: &'static str,
) -> (String, tokio::task::JoinHandle<String>) {
let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds");
let port = listener.local_addr().expect("addr").port();
let handle = tokio::spawn(async move {
let (mut socket, _) = listener.accept().await.expect("accepts");
let received = read_http_request(&mut socket).await;
socket
.write_all(http_response(status, body).as_bytes())
.await
.expect("writes response");
socket.flush().await.expect("flushes");
received
});
(format!("http://127.0.0.1:{port}/v1/messages"), handle)
}
fn call(api_base: &str, messages: Value, params: Value) -> ChatCompletionsRequest<'_> {
ChatCompletionsRequest {
model: "anthropic/claude-sonnet-4-5",
messages,
optional_params: match params {
Value::Object(map) => map,
other => panic!("params must be an object, got {other}"),
},
api_key: Some("sk-test"),
api_base: Some(api_base),
custom_llm_provider: None,
extra_headers: None,
timeout: Some(std::time::Duration::from_secs(10)),
}
}
const GOOD_BODY: &str = r#"{"id":"msg_1","type":"message","role":"assistant","model":"claude-sonnet-4-5-20260101","content":[{"type":"text","text":"hello"}],"stop_reason":"end_turn","stop_sequence":null,"usage":{"input_tokens":11,"output_tokens":4}}"#;
#[tokio::test]
async fn round_trip_sends_the_translated_body_and_normalizes_the_response() {
let (api_base, handle) = serve_once("200 OK", GOOD_BODY).await;
let response = chat_completions(call(
&api_base,
json!([
{"role": "system", "content": "be terse"},
{"role": "user", "content": "hi"}
]),
json!({"max_tokens": 16}),
))
.await
.expect("call succeeds");
let received = handle.await.expect("server task");
let sent: Value = serde_json::from_str(
received
.split_once("\r\n\r\n")
.expect("request has a body")
.1,
)
.expect("body is json");
assert_eq!(
sent["messages"],
json!([{"role": "user", "content": [{"type": "text", "text": "hi"}]}])
);
assert_eq!(
sent["system"],
json!([{"type": "text", "text": "be terse"}])
);
assert_eq!(sent["max_tokens"], json!(16));
assert!(received.to_lowercase().contains("x-api-key: sk-test"));
assert_eq!(
response.choices[0].message.content.as_deref(),
Some("hello")
);
assert_eq!(response.usage.total_tokens, 15);
}
#[tokio::test]
async fn a_response_it_cannot_normalize_is_reported_as_already_sent() {
// The provider was called and billed, so the host must not retry this
// on its own path. `MissingField` here would read as a pre-send
// decline and be retried; `InvalidResponse` cannot.
const NO_USAGE: &str =
r#"{"model":"m","content":[{"type":"text","text":"hi"}],"stop_reason":"end_turn"}"#;
let (api_base, handle) = serve_once("200 OK", NO_USAGE).await;
let err = chat_completions(call(
&api_base,
json!([{"role": "user", "content": "hi"}]),
json!({"max_tokens": 16}),
))
.await
.expect_err("response cannot be normalized");
handle.await.expect("server task");
assert!(
matches!(err, CoreError::InvalidResponse(_)),
"expected a post-send error, got {err:?}"
);
}
#[tokio::test]
async fn a_tool_use_block_in_the_response_is_also_reported_as_already_sent() {
const TOOL_USE: &str = r#"{"model":"m","content":[{"type":"tool_use","id":"t","name":"f","input":{}}],"stop_reason":"tool_use","usage":{"input_tokens":1,"output_tokens":1}}"#;
let (api_base, handle) = serve_once("200 OK", TOOL_USE).await;
let err = chat_completions(call(
&api_base,
json!([{"role": "user", "content": "hi"}]),
json!({"max_tokens": 16}),
))
.await
.expect_err("response cannot be normalized");
handle.await.expect("server task");
assert!(
matches!(err, CoreError::InvalidResponse(_)),
"expected a post-send error, got {err:?}"
);
}
#[tokio::test]
async fn an_upstream_error_status_keeps_its_code() {
let (api_base, handle) =
serve_once("429 Too Many Requests", r#"{"error":"slow down"}"#).await;
let err = chat_completions(call(
&api_base,
json!([{"role": "user", "content": "hi"}]),
json!({"max_tokens": 16}),
))
.await
.expect_err("upstream rejects");
handle.await.expect("server task");
assert!(
matches!(err, CoreError::Http { status: 429, .. }),
"expected a 429, got {err:?}"
);
}
#[tokio::test]
async fn a_connection_that_is_never_established_declines_instead_of_failing() {
// Nothing was sent, so nothing was billed and the host can still serve
// the request. Classing this with the post-send failures would turn a
// recoverable fallback into a user-facing error on exactly the
// deployments whose transport is configured only on the Python client.
let port = {
let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds");
listener.local_addr().expect("has an address").port()
// Dropped here, so the port is closed and the connect is refused.
};
let err = chat_completions(call(
&format!("http://127.0.0.1:{port}/v1/messages"),
json!([{"role": "user", "content": "hi"}]),
json!({"max_tokens": 16}),
))
.await
.expect_err("nothing is listening");
assert!(
matches!(err, CoreError::Connect(_)),
"expected a pre-send connect failure, got {err:?}"
);
}
#[test]
fn response_errors_collapse_to_one_variant_that_can_only_mean_already_sent() {
use crate::chat_completions::handler::as_response_error;
for original in [
CoreError::MissingField("usage"),
CoreError::Unsupported("non-text response content block"),
CoreError::InvalidRequest("whatever".to_string()),
CoreError::Auth("whatever".to_string()),
] {
let label = format!("{original:?}");
assert!(
matches!(as_response_error(original), CoreError::InvalidResponse(_)),
"{label} must not stay retryable once the provider has answered"
);
}
// An upstream status is already unambiguous, so it survives intact.
assert!(matches!(
as_response_error(CoreError::Http {
status: 500,
body: "boom".to_string()
}),
CoreError::Http { status: 500, .. }
));
}
}

View file

@ -0,0 +1,155 @@
use serde_json::{Map, Value};
use crate::error::CoreResult;
use super::types::{
ChatCompletionsResponse, ChatMessage, ChatMessageContent, ProviderChatRequestData,
ProviderChatResponseData,
};
/// How the upstream call is authenticated. API-key strategies are resolved in
/// `prepare`; SigV4 needs the serialized body, so the handler signs it.
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum ChatCompletionsAuth {
Header { name: &'static str, value: String },
Bearer { token: String },
AwsSigV4 { region: String },
}
/// Why a request cannot be served by the Rust path.
///
/// The core declines rather than guessing: the host turns this into a
/// transparent fallback to the Python implementation, which covers the full
/// surface. Acceptance is an allowlist, so a parameter or message shape the
/// core has never seen declines by construction instead of being translated
/// wrong.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct Unsupported(pub &'static str);
pub const STREAM_PARAM: &str = "stream";
/// Message fields that carry no meaning for the upstream body, so their
/// presence does not make a request untranslatable.
const IGNORABLE_MESSAGE_FIELDS: &[&str] = &["name"];
pub trait ChatCompletionsProviderConfig: Sync {
fn complete_url(
&self,
api_base: Option<&str>,
model: &str,
optional_params: &Map<String, Value>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> CoreResult<String>;
fn auth(
&self,
api_key: Option<&str>,
model: &str,
optional_params: &Map<String, Value>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> CoreResult<ChatCompletionsAuth>;
fn default_headers(&self) -> &'static [(&'static str, &'static str)] {
&[("content-type", "application/json")]
}
/// Whether an auth header the caller already supplied is the credential this
/// request should authenticate with, so the resolved one is not applied.
///
/// Defaults to false: the deployment's credential outranks anything
/// forwarded, which is what every provider wants for its own auth header.
/// A provider overrides this only for a scheme it hands off to entirely.
fn defers_to_forwarded_auth(&self, _headers: &[(String, String)]) -> bool {
false
}
/// Provider parameter names (post-mapping) the Rust path knows how to place
/// in the upstream body. Anything outside this set declines the request.
fn supported_params(&self) -> &'static [&'static str];
/// Parameters consumed as call configuration (credentials, endpoints)
/// rather than placed in the body. Accepted, never serialized.
fn config_params(&self) -> &'static [&'static str] {
&[]
}
fn unsupported_reason(
&self,
messages: &[ChatMessage],
optional_params: &Map<String, Value>,
) -> Option<Unsupported> {
unsupported_param(
self.supported_params(),
self.config_params(),
optional_params,
)
.or_else(|| messages.iter().find_map(unsupported_message))
}
fn transform_request(
&self,
model: &str,
messages: Vec<ChatMessage>,
optional_params: Map<String, Value>,
) -> CoreResult<ProviderChatRequestData>;
fn transform_response(
&self,
model: &str,
response: ProviderChatResponseData,
) -> CoreResult<ChatCompletionsResponse>;
}
pub fn unsupported_param(
supported: &'static [&'static str],
config: &'static [&'static str],
optional_params: &Map<String, Value>,
) -> Option<Unsupported> {
if optional_params
.get(STREAM_PARAM)
.and_then(Value::as_bool)
.unwrap_or(false)
{
return Some(Unsupported("streaming"));
}
optional_params
.keys()
.any(|key| {
key != STREAM_PARAM
&& !supported.contains(&key.as_str())
&& !config.contains(&key.as_str())
})
.then_some(Unsupported("unrecognized request parameter"))
}
/// Message shapes the core can translate faithfully: text content, either a
/// plain string or a non-empty list of parts that are all
/// `{"type": "text", "text": ...}`. Tool calls, tool results, and multimodal
/// parts decline so Python's fuller translation handles them.
pub fn unsupported_message(message: &ChatMessage) -> Option<Unsupported> {
if message
.extra
.keys()
.any(|key| !IGNORABLE_MESSAGE_FIELDS.contains(&key.as_str()))
{
return Some(Unsupported("unrecognized message field"));
}
if !matches!(message.role.as_str(), "system" | "user" | "assistant") {
return Some(Unsupported("unrecognized message role"));
}
match &message.content {
None => Some(Unsupported("message without content")),
Some(ChatMessageContent::Text(_)) => None,
Some(ChatMessageContent::Parts(parts)) if parts.is_empty() => {
Some(Unsupported("message without content"))
}
Some(ChatMessageContent::Parts(parts)) => parts
.iter()
.any(|part| {
part.get("type").and_then(Value::as_str) != Some("text")
|| part.get("text").and_then(Value::as_str).is_none()
|| part.as_object().is_some_and(|object| object.len() != 2)
})
.then_some(Unsupported("non-text message content")),
}
}

View file

@ -0,0 +1,112 @@
use std::time::Duration;
use serde::{Deserialize, Serialize};
use serde_json::{Map, Value};
use super::transformation::{ChatCompletionsAuth, ChatCompletionsProviderConfig};
/// A `/chat/completions` call as it crosses into the core.
///
/// `optional_params` arrives already mapped to the provider's own parameter
/// names by the host, exactly as the messages route receives an already
/// Anthropic-shaped body. The core owns the conversation translation, the
/// provider call, and the response normalization.
pub struct ChatCompletionsRequest<'a> {
pub model: &'a str,
pub messages: Value,
pub optional_params: Map<String, Value>,
pub api_key: Option<&'a str>,
pub api_base: Option<&'a str>,
pub custom_llm_provider: Option<&'a str>,
pub extra_headers: Option<Map<String, Value>>,
pub timeout: Option<Duration>,
}
pub(super) struct ProviderChatCompletionsRequest {
pub(super) model: String,
pub(super) config: &'static dyn ChatCompletionsProviderConfig,
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>,
}
/// The provider-shaped request body a config produces. Named rather than a bare
/// `Value` so the transform contract stays a typed one, mirroring
/// [`crate::audio_transcription::types::AudioTranscriptionRequestData`].
pub struct ProviderChatRequestData {
pub body: Value,
}
/// The raw provider response body handed back to a config for normalization.
pub struct ProviderChatResponseData {
pub body: Value,
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
#[serde(untagged)]
pub enum ChatMessageContent {
Text(String),
Parts(Vec<Value>),
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct ChatMessage {
pub role: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub content: Option<ChatMessageContent>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub name: Option<String>,
#[serde(flatten)]
pub extra: Map<String, Value>,
}
/// OpenAI `usage`, including the `prompt_tokens_details` split LiteLLM's Python
/// path reports so cost tracking sees the same numbers on either path.
#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
pub struct PromptTokensDetails {
pub cached_tokens: u64,
pub cache_creation_tokens: u64,
pub text_tokens: u64,
}
#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
pub struct ChatCompletionsUsage {
pub prompt_tokens: u64,
pub completion_tokens: u64,
pub total_tokens: u64,
pub prompt_tokens_details: PromptTokensDetails,
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct ChatCompletionsChoiceMessage {
pub role: String,
// Whether an empty turn is `None` or `""` is the provider's choice, not a
// shared invariant: Anthropic's transform ends on `merged_text or None`
// while Converse assigns the joined string unconditionally. Each config
// mirrors its own, so keep this optional and serialize it even when None.
pub content: Option<String>,
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct ChatCompletionsChoice {
pub index: u64,
pub message: ChatCompletionsChoiceMessage,
pub finish_reason: String,
}
/// The normalized response handed back to the host.
///
/// There is deliberately no `id`: Python mints the `chatcmpl-…` id on the
/// `ModelResponse` it already created, and echoing the provider's own id here
/// would change it. Pinned by `response_carries_no_id` in `tests.rs`.
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct ChatCompletionsResponse {
pub created: u64,
pub model: String,
pub choices: Vec<ChatCompletionsChoice>,
pub usage: ChatCompletionsUsage,
}

View file

@ -12,8 +12,30 @@ pub(crate) const MESSAGES_CONNECT_TIMEOUT_SECS: u64 = 10;
/// Max characters of an upstream error body echoed across the call boundary
/// before truncation, so provider bodies are bounded and data-minimized.
pub(crate) const MESSAGES_ERROR_BODY_MAX_CHARS: usize = 256;
pub(crate) const UPSTREAM_ERROR_BODY_MAX_CHARS: usize = 256;
/// Provider name used for Anthropic Messages when a deployment's provider model
/// does not carry an explicit provider prefix.
pub const ANTHROPIC_MESSAGES_PROVIDER: &str = "anthropic";
/// Prefix identifying an Anthropic OAuth token. Mirrors Python's
/// `ANTHROPIC_OAUTH_TOKEN_PREFIX`, which is what makes `validate_environment`
/// authenticate with `authorization` and drop `x-api-key` entirely.
pub(crate) const ANTHROPIC_OAUTH_TOKEN_PREFIX: &str = "sk-ant-oat";
/// Full-request timeout ceiling for chat completions provider calls, in
/// seconds. Mirrors the Python chat completions default.
pub(crate) const CHAT_COMPLETIONS_TIMEOUT_SECS: u64 = 600;
/// Connect timeout for chat completions provider calls, in seconds.
pub(crate) const CHAT_COMPLETIONS_CONNECT_TIMEOUT_SECS: u64 = 10;
/// `object` field every non-streaming chat completion response carries.
pub const CHAT_COMPLETION_OBJECT: &str = "chat.completion";
/// Placeholder Python substitutes for empty or whitespace-only message text,
/// which Anthropic and Bedrock both reject. Must match
/// `_EMPTY_TEXT_PLACEHOLDER` in
/// `litellm/litellm_core_utils/prompt_templates/factory.py`.
pub const EMPTY_TEXT_PLACEHOLDER: &str =
"[System: Empty message content sanitised to satisfy protocol]";

View file

@ -23,8 +23,19 @@ pub enum CoreError {
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),
}
pub fn json_type_name(value: &serde_json::Value) -> &'static str {

View file

@ -0,0 +1,112 @@
//! Header and upstream-body helpers shared by every route module.
use serde_json::{Map, Value};
use crate::constants::UPSTREAM_ERROR_BODY_MAX_CHARS;
use crate::error::{CoreError, CoreResult, json_type_name};
/// Bound an upstream error body before it crosses a host boundary, so provider
/// bodies stay data-minimized.
pub fn truncate_error_body(body: &str) -> String {
if body.chars().count() <= UPSTREAM_ERROR_BODY_MAX_CHARS {
return body.to_string();
}
let truncated: String = body.chars().take(UPSTREAM_ERROR_BODY_MAX_CHARS).collect();
format!("{truncated}... (truncated)")
}
pub fn string_headers(
context: &'static str,
extra_headers: Option<Map<String, Value>>,
) -> CoreResult<Vec<(String, String)>> {
extra_headers
.unwrap_or_default()
.into_iter()
.map(|(key, value)| {
value
.as_str()
.map(|value| (key.clone(), value.to_string()))
.ok_or_else(|| {
CoreError::InvalidRequest(format!(
"{context} extra_headers.{key} must be a string, got {}",
json_type_name(&value)
))
})
})
.collect()
}
pub fn has_header(headers: &[(String, String)], name: &str) -> bool {
headers
.iter()
.any(|(key, _)| key.eq_ignore_ascii_case(name))
}
pub fn has_bearer_auth(headers: &[(String, String)]) -> bool {
headers.iter().any(|(name, value)| {
if !name.eq_ignore_ascii_case("authorization") {
return false;
}
let value = value.trim();
value.len() > 7
&& value[..7].eq_ignore_ascii_case("bearer ")
&& !value[7..].trim().is_empty()
})
}
#[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
#[test]
fn truncate_leaves_short_bodies_untouched() {
assert_eq!(truncate_error_body("short"), "short");
}
#[test]
fn truncate_bounds_long_bodies_by_characters() {
let body = "\u{00e9}".repeat(UPSTREAM_ERROR_BODY_MAX_CHARS + 10);
let truncated = truncate_error_body(&body);
assert!(truncated.ends_with("... (truncated)"));
assert_eq!(
truncated.chars().count(),
UPSTREAM_ERROR_BODY_MAX_CHARS + "... (truncated)".chars().count()
);
}
#[test]
fn string_headers_rejects_non_string_values() {
let headers = Map::from_iter([("x-trace".to_string(), json!(7))]);
let err = string_headers("chat completions", Some(headers)).expect_err("non-string value");
assert_eq!(
err,
CoreError::InvalidRequest(
"chat completions extra_headers.x-trace must be a string, got number".to_string()
)
);
}
#[test]
fn header_lookup_is_case_insensitive() {
let headers = vec![("X-Api-Key".to_string(), "k".to_string())];
assert!(has_header(&headers, "x-api-key"));
assert!(!has_header(&headers, "authorization"));
}
#[test]
fn bearer_detection_requires_a_non_empty_token() {
assert!(has_bearer_auth(&[(
"Authorization".to_string(),
"Bearer abc".to_string()
)]));
assert!(!has_bearer_auth(&[(
"Authorization".to_string(),
"Bearer ".to_string()
)]));
assert!(!has_bearer_auth(&[(
"Authorization".to_string(),
"Basic abc".to_string()
)]));
}
}

View file

@ -1,8 +1,10 @@
pub mod audio_transcription;
pub mod caching;
pub mod call_lifecycle;
pub mod chat_completions;
pub mod constants;
pub mod error;
pub mod http_utils;
pub mod messages;
pub mod ocr;
pub mod providers;

View file

@ -1,19 +1,15 @@
use serde_json::{Map, Value};
use crate::constants::MESSAGES_ERROR_BODY_MAX_CHARS;
use crate::error::{CoreError, CoreResult, json_type_name};
use crate::error::CoreResult;
use crate::http_utils::string_headers as shared_string_headers;
use crate::providers::anthropic::messages::transformation::ANTHROPIC_MESSAGES_CONFIG;
use crate::providers::azure_ai::messages::transformation::AZURE_ANTHROPIC_MESSAGES_CONFIG;
use super::transformation::AnthropicMessagesProviderConfig;
pub(super) fn truncate_error_body(body: &str) -> String {
if body.chars().count() <= MESSAGES_ERROR_BODY_MAX_CHARS {
return body.to_string();
}
let truncated: String = body.chars().take(MESSAGES_ERROR_BODY_MAX_CHARS).collect();
format!("{truncated}... (truncated)")
}
pub(super) use crate::http_utils::{has_bearer_auth, has_header, truncate_error_body};
const HEADER_CONTEXT: &str = "messages";
pub(super) fn messages_provider_config(
provider: &str,
@ -28,37 +24,5 @@ pub(super) fn messages_provider_config(
pub(super) fn string_headers(
extra_headers: Option<Map<String, Value>>,
) -> CoreResult<Vec<(String, String)>> {
extra_headers
.unwrap_or_default()
.into_iter()
.map(|(key, value)| {
value
.as_str()
.map(|value| (key.clone(), value.to_string()))
.ok_or_else(|| {
CoreError::InvalidRequest(format!(
"messages extra_headers.{key} must be a string, got {}",
json_type_name(&value)
))
})
})
.collect()
}
pub(super) fn has_header(headers: &[(String, String)], name: &str) -> bool {
headers
.iter()
.any(|(key, _)| key.eq_ignore_ascii_case(name))
}
pub(super) fn has_bearer_auth(headers: &[(String, String)]) -> bool {
headers.iter().any(|(name, value)| {
if !name.eq_ignore_ascii_case("authorization") {
return false;
}
let value = value.trim();
value.len() > 7
&& value[..7].eq_ignore_ascii_case("bearer ")
&& !value[7..].trim().is_empty()
})
shared_string_headers(HEADER_CONTEXT, extra_headers)
}

View file

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

View file

@ -0,0 +1,444 @@
use super::*;
use serde_json::json;
fn messages(value: Value) -> Vec<ChatMessage> {
serde_json::from_value(value).expect("valid messages")
}
fn params(value: Value) -> Map<String, Value> {
match value {
Value::Object(map) => map,
other => panic!("params must be an object, got {other}"),
}
}
fn transform(model: &str, msgs: Value, opts: Value) -> Value {
ANTHROPIC_CHAT_COMPLETIONS_CONFIG
.transform_request(model, messages(msgs), params(opts))
.expect("request transforms")
.body
}
fn transform_response(body: Value) -> CoreResult<ChatCompletionsResponse> {
ANTHROPIC_CHAT_COMPLETIONS_CONFIG
.transform_response("claude-sonnet-4-5", ProviderChatResponseData { body })
}
fn reason(msgs: Value, opts: Value) -> Option<Unsupported> {
ANTHROPIC_CHAT_COMPLETIONS_CONFIG.unsupported_reason(&messages(msgs), &params(opts))
}
#[test]
fn builds_the_messages_body_python_builds() {
let body = transform(
"claude-sonnet-4-5",
json!([
{"role": "system", "content": "be terse"},
{"role": "user", "content": "hi"}
]),
json!({"max_tokens": 128, "temperature": 0.2}),
);
assert_eq!(
body,
json!({
"model": "claude-sonnet-4-5",
"messages": [
{"role": "user", "content": [{"type": "text", "text": "hi"}]}
],
"system": [{"type": "text", "text": "be terse"}],
"max_tokens": 128,
"temperature": 0.2
})
);
}
#[test]
fn omits_system_when_no_system_message_is_present() {
let body = transform(
"claude-sonnet-4-5",
json!([{"role": "user", "content": "hi"}]),
json!({"max_tokens": 16}),
);
assert!(body.get("system").is_none());
}
#[test]
fn merges_consecutive_turns_and_wraps_every_text_in_a_block() {
let body = transform(
"claude-sonnet-4-5",
json!([
{"role": "user", "content": "one"},
{"role": "user", "content": [{"type": "text", "text": "two"}]},
{"role": "assistant", "content": "ack"}
]),
json!({"max_tokens": 16}),
);
assert_eq!(
body["messages"],
json!([
{"role": "user", "content": [
{"type": "text", "text": "one"},
{"type": "text", "text": "two"}
]},
{"role": "assistant", "content": [{"type": "text", "text": "ack"}]}
])
);
}
#[test]
fn right_strips_a_trailing_assistant_prefill_like_python() {
let body = transform(
"claude-sonnet-4-5",
json!([
{"role": "user", "content": "hi"},
{"role": "assistant", "content": "Argentina "}
]),
json!({"max_tokens": 16}),
);
assert_eq!(
body["messages"][1]["content"][0]["text"],
json!("Argentina")
);
}
#[test]
fn passes_every_supported_param_through_untouched() {
let body = transform(
"claude-sonnet-4-5",
json!([{"role": "user", "content": "hi"}]),
json!({
"max_tokens": 64,
"temperature": 0.1,
"top_p": 0.9,
"stop_sequences": ["STOP"]
}),
);
assert_eq!(body["max_tokens"], json!(64));
assert_eq!(body["temperature"], json!(0.1));
assert_eq!(body["top_p"], json!(0.9));
assert_eq!(body["stop_sequences"], json!(["STOP"]));
}
#[test]
fn declines_top_k_because_python_gates_it_by_model_below_this_point() {
// `temperature` and `top_p` arrive already resolved, because
// `map_openai_params` applies `_apply_sampling_param` to them before the
// gate runs. `top_k` bypasses that and is gated inside `transform_request`,
// the function this route replaces, so forwarding it would send `top_k` to
// a model that removed sampling params and take a 400 after the call, where
// Python drops it and succeeds.
assert_eq!(
reason(
json!([{"role": "user", "content": "hi"}]),
json!({"top_k": 40})
),
Some(Unsupported("unrecognized request parameter"))
);
}
#[test]
fn declines_streaming_before_anything_else() {
assert_eq!(
reason(
json!([{"role": "user", "content": "hi"}]),
json!({"stream": true, "max_tokens": 16})
),
Some(Unsupported("streaming"))
);
}
#[test]
fn accepts_an_explicit_stream_false() {
assert_eq!(
reason(
json!([{"role": "user", "content": "hi"}]),
json!({"stream": false, "max_tokens": 16})
),
None
);
}
#[test]
fn declines_any_param_outside_the_allowlist() {
for param in [
json!({"tools": []}),
json!({"tool_choice": {"type": "auto"}}),
json!({"thinking": {"type": "enabled"}}),
json!({"system": "injected"}),
json!({"metadata": {"user_id": "u1"}}),
json!({"output_config": {"effort": "high"}}),
] {
assert_eq!(
reason(json!([{"role": "user", "content": "hi"}]), param.clone()),
Some(Unsupported("unrecognized request parameter")),
"expected {param} to decline"
);
}
}
#[test]
fn declines_tool_calls_tool_results_and_multimodal_content() {
assert_eq!(
reason(
json!([
{"role": "user", "content": "hi"},
{"role": "assistant", "content": null, "tool_calls": [
{"id": "c1", "type": "function",
"function": {"name": "f", "arguments": "{}"}}
]}
]),
json!({})
),
Some(Unsupported("unrecognized message field"))
);
assert_eq!(
reason(
json!([
{"role": "user", "content": "hi"},
{"role": "tool", "tool_call_id": "c1", "content": "ok"}
]),
json!({})
),
Some(Unsupported("unrecognized message field"))
);
assert_eq!(
reason(
json!([{"role": "user", "content": [
{"type": "image_url", "image_url": {"url": "https://x/y.png"}}
]}]),
json!({})
),
Some(Unsupported("non-text message content"))
);
assert_eq!(
reason(
json!([{"role": "user", "content": [
{"type": "text", "text": "hi", "cache_control": {"type": "ephemeral"}}
]}]),
json!({})
),
Some(Unsupported("non-text message content"))
);
}
#[test]
fn declines_a_message_whose_content_list_is_empty() {
// An empty list passes every per-part check, so without this it would reach
// the provider as an empty `content` array and fail after the call rather
// than declining to Python before it.
assert_eq!(
reason(json!([{"role": "user", "content": []}]), json!({})),
Some(Unsupported("message without content"))
);
assert_eq!(
reason(
json!([{"role": "user", "content": [{"type": "text", "text": "hi"}]}]),
json!({})
),
None
);
}
#[test]
fn declines_a_conversation_that_does_not_open_on_a_user_turn() {
assert_eq!(
reason(
json!([
{"role": "system", "content": "be terse"},
{"role": "assistant", "content": "prefill"}
]),
json!({})
),
Some(Unsupported("conversation does not open on a user turn"))
);
}
#[test]
fn accepts_a_plain_text_conversation() {
assert_eq!(
reason(
json!([
{"role": "system", "content": "be terse"},
{"role": "user", "content": "hi"},
{"role": "assistant", "content": "hello"},
{"role": "user", "content": [{"type": "text", "text": "again"}]}
]),
json!({"max_tokens": 16, "temperature": 0.5})
),
None
);
}
#[test]
fn normalizes_a_text_response_into_openai_shape() {
let response = transform_response(json!({
"id": "msg_123",
"type": "message",
"role": "assistant",
"model": "claude-sonnet-4-5-20260101",
"content": [{"type": "text", "text": "hello"}, {"type": "text", "text": " there"}],
"stop_reason": "end_turn",
"stop_sequence": null,
"usage": {"input_tokens": 11, "output_tokens": 4}
}))
.expect("response transforms");
assert_eq!(response.model, "claude-sonnet-4-5-20260101");
assert_eq!(response.choices.len(), 1);
assert_eq!(response.choices[0].index, 0);
assert_eq!(response.choices[0].message.role, "assistant");
assert_eq!(
response.choices[0].message.content.as_deref(),
Some("hello there")
);
assert_eq!(response.choices[0].finish_reason, "stop");
assert_eq!(response.usage.prompt_tokens, 11);
assert_eq!(response.usage.completion_tokens, 4);
assert_eq!(response.usage.total_tokens, 15);
}
#[test]
fn folds_cache_tokens_into_prompt_tokens_like_python() {
let response = transform_response(json!({
"model": "claude-sonnet-4-5",
"content": [{"type": "text", "text": "hi"}],
"stop_reason": "end_turn",
"usage": {
"input_tokens": 10,
"output_tokens": 2,
"cache_read_input_tokens": 5,
"cache_creation_input_tokens": 3
}
}))
.expect("response transforms");
assert_eq!(response.usage.prompt_tokens, 18);
assert_eq!(response.usage.total_tokens, 20);
assert_eq!(response.usage.prompt_tokens_details.cached_tokens, 5);
assert_eq!(
response.usage.prompt_tokens_details.cache_creation_tokens,
3
);
assert_eq!(response.usage.prompt_tokens_details.text_tokens, 10);
}
#[test]
fn maps_max_tokens_stop_reason_to_length() {
let response = transform_response(json!({
"model": "claude-sonnet-4-5",
"content": [{"type": "text", "text": "hi"}],
"stop_reason": "max_tokens",
"usage": {"input_tokens": 1, "output_tokens": 1}
}))
.expect("response transforms");
assert_eq!(response.choices[0].finish_reason, "length");
}
#[test]
fn a_refusal_returns_the_completion_python_returns() {
// `refusal` is a stop_reason, not a content block type, so the content is
// ordinary text and this normalizes rather than declining. Python maps it
// to content_filter in _FINISH_REASON_MAP and returns the completion.
let response = transform_response(json!({
"model": "claude-sonnet-4-5",
"content": [{"type": "text", "text": "I can't help with that."}],
"stop_reason": "refusal",
"usage": {"input_tokens": 9, "output_tokens": 6}
}))
.expect("a refusal still transforms");
assert_eq!(response.choices[0].finish_reason, "content_filter");
assert_eq!(
response.choices[0].message.content.as_deref(),
Some("I can't help with that.")
);
}
#[test]
fn reports_no_content_rather_than_an_empty_string() {
let response = transform_response(json!({
"model": "claude-sonnet-4-5",
"content": [],
"stop_reason": "end_turn",
"usage": {"input_tokens": 1, "output_tokens": 0}
}))
.expect("response transforms");
assert_eq!(response.choices[0].message.content, None);
}
#[test]
fn response_carries_no_id_so_python_keeps_its_chatcmpl_id() {
let response = transform_response(json!({
"id": "msg_should_not_leak",
"model": "claude-sonnet-4-5",
"content": [{"type": "text", "text": "hi"}],
"stop_reason": "end_turn",
"usage": {"input_tokens": 1, "output_tokens": 1}
}))
.expect("response transforms");
let value = serde_json::to_value(response).expect("serializable");
assert!(
value.get("id").is_none(),
"the rust response must not carry an id, got {value}"
);
}
#[test]
fn declines_a_response_carrying_a_non_text_block() {
let err = transform_response(json!({
"model": "claude-sonnet-4-5",
"content": [{"type": "tool_use", "id": "t1", "name": "f", "input": {}}],
"stop_reason": "tool_use",
"usage": {"input_tokens": 1, "output_tokens": 1}
}))
.expect_err("non-text block");
assert_eq!(
err,
CoreError::Unsupported("non-text response content block")
);
}
#[test]
fn errors_on_a_response_missing_required_fields() {
assert_eq!(
transform_response(json!("nope")).expect_err("not an object"),
CoreError::InvalidResponse("messages response is not an object".to_string())
);
assert_eq!(
transform_response(json!({"model": "m", "usage": {}})).expect_err("no content"),
CoreError::MissingField("content")
);
assert_eq!(
transform_response(json!({"model": "m", "content": []})).expect_err("no usage"),
CoreError::MissingField("usage")
);
assert_eq!(
transform_response(json!({"content": [], "usage": {}})).expect_err("no model"),
CoreError::MissingField("model")
);
}
#[test]
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)
.expect("url builds"),
"https://api.anthropic.com/v1/messages"
);
assert_eq!(
config
.auth(Some("sk-x"), "claude-sonnet-4-5", &Map::new(), &|_| None)
.expect("auth resolves"),
ChatCompletionsAuth::Header {
name: "x-api-key",
value: "sk-x".to_string()
}
);
assert_eq!(
config.default_headers(),
&[
("anthropic-version", "2023-06-01"),
("content-type", "application/json"),
]
);
}

View file

@ -0,0 +1,211 @@
use serde_json::{Map, Value, json};
use crate::chat_completions::conversation::{Conversation, build_conversation};
use crate::chat_completions::transformation::{
ChatCompletionsAuth, ChatCompletionsProviderConfig, Unsupported, unsupported_message,
unsupported_param,
};
use crate::chat_completions::types::{
ChatCompletionsChoice, ChatCompletionsChoiceMessage, ChatCompletionsResponse, ChatMessage,
ProviderChatRequestData, ProviderChatResponseData,
};
use crate::constants::ANTHROPIC_OAUTH_TOKEN_PREFIX;
use crate::error::{CoreError, CoreResult};
use crate::providers::anthropic::messages::transformation::{
complete_anthropic_url, resolve_anthropic_api_key,
};
use crate::chat_completions::response_utils::{finish_reason_for, unix_now, usage_from_parts};
/// Anthropic parameter names, post `map_openai_params`, that the Rust path can
/// place verbatim in the Messages body.
///
/// `top_k` is deliberately absent even though the Messages API takes it.
/// `temperature` and `top_p` reach this gate already resolved, because
/// `map_openai_params` runs first and applies `_apply_sampling_param` to them.
/// `top_k` bypasses `map_openai_params` entirely, so Python applies that same
/// per-model gate inside `transform_request`, the function this route replaces.
/// Forwarding it would send `top_k` to a model that removed sampling params and
/// take a 400 after the call, where Python drops it and succeeds.
const SUPPORTED_PARAMS: &[&str] = &["max_tokens", "temperature", "top_p", "stop_sequences"];
pub struct AnthropicChatCompletionsConfig;
pub const ANTHROPIC_CHAT_COMPLETIONS_CONFIG: AnthropicChatCompletionsConfig =
AnthropicChatCompletionsConfig;
fn text_block(text: &str) -> Value {
json!({"type": "text", "text": text})
}
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(
&self,
api_base: Option<&str>,
_model: &str,
_optional_params: &Map<String, Value>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> CoreResult<String> {
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>,
) -> CoreResult<ChatCompletionsAuth> {
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_params(&self) -> &'static [&'static str] {
SUPPORTED_PARAMS
}
fn unsupported_reason(
&self,
messages: &[ChatMessage],
optional_params: &Map<String, Value>,
) -> Option<Unsupported> {
unsupported_param(SUPPORTED_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,
messages: Vec<ChatMessage>,
optional_params: Map<String, Value>,
) -> CoreResult<ProviderChatRequestData> {
Ok(ProviderChatRequestData {
body: anthropic_body(model, &build_conversation(&messages), optional_params),
})
}
fn transform_response(
&self,
_model: &str,
response: ProviderChatResponseData,
) -> CoreResult<ChatCompletionsResponse> {
let body = response.body.as_object().ok_or_else(|| {
CoreError::InvalidResponse("messages response is not an object".into())
})?;
let content = body
.get("content")
.and_then(Value::as_array)
.ok_or(CoreError::MissingField("content"))?;
// The route declines tool and thinking requests, so a non-text block
// means the response carries something this path never asked for.
// Decline rather than silently dropping it; the host falls back.
if content
.iter()
.any(|block| block.get("type").and_then(Value::as_str) != Some("text"))
{
return Err(CoreError::Unsupported("non-text response content block"));
}
let text: String = content
.iter()
.filter_map(|block| block.get("text").and_then(Value::as_str))
.collect();
let usage = body
.get("usage")
.and_then(Value::as_object)
.ok_or(CoreError::MissingField("usage"))?;
let field = |name: &str| usage.get(name).and_then(Value::as_u64).unwrap_or(0);
Ok(ChatCompletionsResponse {
created: unix_now(),
model: body
.get("model")
.and_then(Value::as_str)
.ok_or(CoreError::MissingField("model"))?
.to_string(),
choices: vec![ChatCompletionsChoice {
index: 0,
message: ChatCompletionsChoiceMessage {
role: "assistant".to_string(),
content: (!text.is_empty()).then_some(text),
},
finish_reason: finish_reason_for(
body.get("stop_reason")
.and_then(Value::as_str)
.unwrap_or(""),
)
.to_string(),
}],
usage: usage_from_parts(
field("input_tokens"),
field("output_tokens"),
field("cache_read_input_tokens"),
field("cache_creation_input_tokens"),
),
})
}
}
#[cfg(test)]
#[path = "tests.rs"]
mod tests;

View file

@ -1 +1,2 @@
pub mod chat_completions;
pub mod messages;

View file

@ -8,11 +8,8 @@ use crate::audio_transcription::types::{
};
use crate::error::{CoreError, CoreResult, json_type_name};
use super::aws_base::AwsAuthConfig;
use super::constants::{
AWS_REGION, AWS_REGION_NAME, BEDROCK_RUNTIME_ENDPOINT_TEMPLATE, BEDROCK_SERVICE,
DEFAULT_BEDROCK_REGION,
};
pub use super::aws_base::{aws_auth_config, bedrock_model_id_and_region, resolve_bedrock_region};
use super::constants::{BEDROCK_RUNTIME_ENDPOINT_TEMPLATE, BEDROCK_SERVICE};
const SUPPORTED_PARAMS: &[&str] = &["language", "prompt", "temperature", "response_format"];
@ -21,64 +18,6 @@ pub static BEDROCK_AUDIO_TRANSCRIPTION_CONFIG: BedrockAudioTranscriptionConfig =
pub struct BedrockAudioTranscriptionConfig;
pub fn bedrock_model_id_and_region(model: &str) -> (String, Option<String>) {
let mut stripped = model;
for prefix in ["bedrock/converse/", "bedrock/", "converse/"] {
if let Some(value) = stripped.strip_prefix(prefix) {
stripped = value;
break;
}
}
let mut region = None;
if let Some((candidate, remainder)) = stripped.split_once('/')
&& is_bedrock_region(candidate)
{
region = Some(candidate.to_string());
stripped = remainder;
}
for prefix in ["nova-2/", "nova/"] {
if let Some(value) = stripped.strip_prefix(prefix) {
stripped = value;
break;
}
}
if region.is_none() {
region = stripped
.strip_prefix("arn:")
.and_then(|value| value.split(':').nth(3))
.filter(|value| !value.is_empty())
.map(str::to_string);
}
(stripped.to_string(), region)
}
fn is_bedrock_region(value: &str) -> bool {
value.len() > 3
&& value.contains('-')
&& value
.chars()
.all(|char| char.is_ascii_alphanumeric() || char == '-')
}
pub fn resolve_bedrock_region(
model_region: Option<&str>,
optional_params: &Map<String, Value>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> String {
if let Some(region) = optional_params
.get("aws_region_name")
.and_then(Value::as_str)
{
return region.to_string();
}
if let Some(region) = model_region {
return region.to_string();
}
env_lookup(AWS_REGION_NAME)
.or_else(|| env_lookup(AWS_REGION))
.unwrap_or_else(|| DEFAULT_BEDROCK_REGION.to_string())
}
fn audio_fields(audio: Value) -> CoreResult<(String, String)> {
let object = audio.as_object().ok_or_else(|| CoreError::InvalidType {
expected: "object",
@ -203,32 +142,6 @@ impl AudioTranscriptionProviderConfig for BedrockAudioTranscriptionConfig {
}
}
pub fn aws_auth_config(
optional_params: &Map<String, Value>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> AwsAuthConfig {
let value = |key: &str| {
optional_params
.get(key)
.and_then(Value::as_str)
.map(str::to_string)
};
let env = |key: &str| env_lookup(key);
AwsAuthConfig {
access_key_id: value("aws_access_key_id").or_else(|| env("AWS_ACCESS_KEY_ID")),
secret_access_key: value("aws_secret_access_key").or_else(|| env("AWS_SECRET_ACCESS_KEY")),
session_token: value("aws_session_token").or_else(|| env("AWS_SESSION_TOKEN")),
region_name: value("aws_region_name").or_else(|| env(AWS_REGION_NAME)),
session_name: value("aws_session_name").or_else(|| env("AWS_SESSION_NAME")),
profile_name: value("aws_profile_name").or_else(|| env("AWS_PROFILE_NAME")),
role_name: value("aws_role_name").or_else(|| env("AWS_ROLE_NAME")),
web_identity_token: value("aws_web_identity_token")
.or_else(|| env("AWS_WEB_IDENTITY_TOKEN")),
sts_endpoint: value("aws_sts_endpoint").or_else(|| env("AWS_STS_ENDPOINT")),
external_id: value("aws_external_id").or_else(|| env("AWS_EXTERNAL_ID")),
}
}
#[cfg(test)]
mod tests {
use super::*;

View file

@ -12,13 +12,15 @@ 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::constants::{
AWS_ACCESS_KEY_ID, AWS_EXTERNAL_ID, AWS_PROFILE_NAME, AWS_REGION_NAME, AWS_ROLE_ARN,
AWS_ROLE_NAME, AWS_SECRET_ACCESS_KEY, AWS_SESSION_NAME, AWS_SESSION_TOKEN, AWS_STS_ENDPOINT,
AWS_WEB_IDENTITY_TOKEN, AWS_WEB_IDENTITY_TOKEN_FILE, BEDROCK_SERVICE,
DEFAULT_SESSION_NAME_PREFIX,
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,
AWS_SIGNED_HEADER_NAMES, AWS_STS_ENDPOINT, AWS_WEB_IDENTITY_TOKEN, AWS_WEB_IDENTITY_TOKEN_FILE,
BEDROCK_SERVICE, DEFAULT_BEDROCK_REGION, DEFAULT_SESSION_NAME_PREFIX,
SIGV4_COMPUTED_HEADER_NAMES,
};
const STATIC_CREDENTIALS_TTL: Duration = Duration::from_secs(3600 - 60);
@ -401,6 +403,33 @@ fn default_session_name() -> String {
format!("{DEFAULT_SESSION_NAME_PREFIX}-{seconds}")
}
/// The subset of `headers` SigV4 should cover.
///
/// Python signs only these and reattaches the rest afterwards, so a forwarded
/// client header cannot change the canonical request and invalidate the
/// signature. Signing everything instead makes the request 403 on a header the
/// caller supplied, on a deployment that works on the Python path.
pub fn aws_signature_headers(headers: &BTreeMap<String, String>) -> BTreeMap<String, String> {
headers
.iter()
.filter(|(name, _)| {
let name = name.to_ascii_lowercase();
AWS_SIGNED_HEADER_NAMES.contains(&name.as_str())
|| name.starts_with("x-amz-")
|| name.starts_with("x-amzn-")
})
.map(|(name, value)| (name.clone(), value.clone()))
.collect()
}
/// Whether the signer produces `name` itself.
///
/// Python's reattach loop skips these, so a caller-supplied copy never reaches
/// the wire next to the computed one.
pub fn is_sigv4_computed_header(name: &str) -> bool {
SIGV4_COMPUTED_HEADER_NAMES.contains(&name.to_ascii_lowercase().as_str())
}
pub fn sign_bedrock_post(
url: &str,
body: &[u8],
@ -441,6 +470,121 @@ pub fn sign_bedrock_post(
.collect())
}
/// Model-id and region parsing shared by every Bedrock route.
pub fn bedrock_model_id_and_region(model: &str) -> (String, Option<String>) {
let mut stripped = model;
for prefix in ["bedrock/converse/", "bedrock/", "converse/"] {
if let Some(value) = stripped.strip_prefix(prefix) {
stripped = value;
break;
}
}
let mut region = None;
if let Some((candidate, remainder)) = stripped.split_once('/')
&& is_bedrock_region(candidate)
{
region = Some(candidate.to_string());
stripped = remainder;
}
for prefix in ["nova-2/", "nova/"] {
if let Some(value) = stripped.strip_prefix(prefix) {
stripped = value;
break;
}
}
if region.is_none() {
// Python splits the whole ARN and takes field 3, the region. Stripping
// `arn:` first shifts every field down one, so the region is field 2
// here; field 3 is the account id.
region = stripped
.strip_prefix("arn:")
.and_then(|value| value.split(':').nth(2))
.filter(|value| !value.is_empty())
.map(str::to_string);
}
(stripped.to_string(), region)
}
fn is_bedrock_region(value: &str) -> bool {
value.len() > 3
&& value.contains('-')
&& value
.chars()
.all(|char| char.is_ascii_alphanumeric() || char == '-')
}
pub fn resolve_bedrock_region(
model_region: Option<&str>,
optional_params: &Map<String, Value>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> String {
if let Some(region) = optional_params
.get("aws_region_name")
.and_then(Value::as_str)
{
return region.to_string();
}
if let Some(region) = model_region {
return region.to_string();
}
env_lookup(AWS_REGION_NAME)
.or_else(|| env_lookup(AWS_REGION))
.unwrap_or_else(|| DEFAULT_BEDROCK_REGION.to_string())
}
pub fn aws_auth_config(
optional_params: &Map<String, Value>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> AwsAuthConfig {
let value = |key: &str| {
optional_params
.get(key)
.and_then(Value::as_str)
.map(str::to_string)
};
let env = |key: &str| env_lookup(key);
AwsAuthConfig {
access_key_id: value("aws_access_key_id").or_else(|| env("AWS_ACCESS_KEY_ID")),
secret_access_key: value("aws_secret_access_key").or_else(|| env("AWS_SECRET_ACCESS_KEY")),
session_token: value("aws_session_token").or_else(|| env("AWS_SESSION_TOKEN")),
region_name: value("aws_region_name").or_else(|| env(AWS_REGION_NAME)),
session_name: value("aws_session_name").or_else(|| env("AWS_SESSION_NAME")),
profile_name: value("aws_profile_name").or_else(|| env("AWS_PROFILE_NAME")),
role_name: value("aws_role_name").or_else(|| env("AWS_ROLE_NAME")),
web_identity_token: value("aws_web_identity_token")
.or_else(|| env("AWS_WEB_IDENTITY_TOKEN")),
sts_endpoint: value("aws_sts_endpoint").or_else(|| env("AWS_STS_ENDPOINT")),
external_id: value("aws_external_id").or_else(|| env("AWS_EXTERNAL_ID")),
}
}
/// Credentials a host resolved through its own chain and handed down verbatim.
///
/// A host with its own resolution (LiteLLM's Python `BaseAWSLLM`, which reads
/// profiles, STS and boto sessions) passes the result here so the core signs
/// with exactly those. Without this the core would re-derive from ambient
/// state, where an unrelated `AWS_ROLE_NAME` or `AWS_PROFILE_NAME` in the
/// environment outranks explicit keys in [`classify_auth`] and the two sides
/// would sign as different principals.
pub fn host_supplied_credentials(optional_params: &Map<String, Value>) -> Option<Credentials> {
let value = |key: &str| {
optional_params
.get(key)
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
};
let access_key_id = value("aws_access_key_id")?;
let secret_access_key = value("aws_secret_access_key")?;
Some(Credentials::new(
access_key_id,
secret_access_key,
value("aws_session_token").map(str::to_string),
None,
"litellm-host-supplied",
))
}
#[cfg(test)]
mod tests {
use super::*;
@ -458,6 +602,18 @@ mod tests {
)
}
#[test]
fn reads_the_region_field_of_a_model_arn_not_the_account_id() {
// Python's `_get_aws_region_from_model_arn` splits the whole ARN and
// takes field 3. Stripping `arn:` first shifts every field down one, so
// the region is field 2 here. Taking field 3 after the strip returns
// the account id, which is not a region at all.
let (_, region) = bedrock_model_id_and_region(
"bedrock/arn:aws:bedrock:us-west-2:123456789012:foundation-model/anthropic.claude-v2",
);
assert_eq!(region.as_deref(), Some("us-west-2"));
}
#[test]
fn classification_preserves_python_precedence() {
let config = AwsAuthConfig {
@ -610,6 +766,52 @@ mod tests {
));
}
#[test]
fn a_forwarded_client_header_is_not_folded_into_the_signature() {
// Python signs only the AWS header set, so a header a caller forwarded
// cannot change the canonical request. Signing it instead makes the
// request 403 the moment anything on the wire rewrites or drops it.
let (url, body, mut headers) = parity_inputs();
headers.insert("x-request-id".to_string(), "abc-123".to_string());
headers.insert("Accept-Encoding".to_string(), "gzip".to_string());
headers.insert("x-amzn-trace-id".to_string(), "Root=1-abc".to_string());
let signable = aws_signature_headers(&headers);
assert!(!signable.contains_key("x-request-id"));
assert!(!signable.contains_key("Accept-Encoding"));
// The AWS-prefixed one is genuinely part of the signature.
assert!(signable.contains_key("x-amzn-trace-id"));
assert!(signable.contains_key("Content-Type"));
let credentials = Credentials::new(
"AKIDEXAMPLE",
"wJalrXUtnFEMI/K7MDENG+bPxRfiCYEXAMPLEKEY",
None,
None,
"test",
);
let signed = sign_bedrock_post(
&url,
&body,
&signable,
"us-east-1",
&credentials,
SystemTime::UNIX_EPOCH,
)
.expect("signs");
let authorization = signed
.get("Authorization")
.expect("carries an authorization header");
assert!(
!authorization.contains("x-request-id"),
"forwarded header reached SignedHeaders: {authorization}"
);
assert!(
!authorization.contains("accept-encoding"),
"forwarded header reached SignedHeaders: {authorization}"
);
}
#[test]
fn signing_matches_botocore_golden_vector() {
let (url, body, headers) = parity_inputs();

View file

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

View file

@ -0,0 +1,580 @@
use super::*;
use serde_json::json;
fn messages(value: Value) -> Vec<ChatMessage> {
serde_json::from_value(value).expect("valid messages")
}
fn params(value: Value) -> Map<String, Value> {
match value {
Value::Object(map) => map,
other => panic!("params must be an object, got {other}"),
}
}
fn transform(msgs: Value, opts: Value) -> Value {
BEDROCK_CHAT_COMPLETIONS_CONFIG
.transform_request(
"anthropic.claude-sonnet-4-5-v1:0",
messages(msgs),
params(opts),
)
.expect("request transforms")
.body
}
fn transform_response(body: Value) -> CoreResult<ChatCompletionsResponse> {
BEDROCK_CHAT_COMPLETIONS_CONFIG.transform_response(
"anthropic.claude-sonnet-4-5-v1:0",
ProviderChatResponseData { body },
)
}
fn reason(msgs: Value, opts: Value) -> Option<Unsupported> {
BEDROCK_CHAT_COMPLETIONS_CONFIG.unsupported_reason(&messages(msgs), &params(opts))
}
#[test]
fn builds_the_converse_body_python_builds() {
let body = transform(
json!([
{"role": "system", "content": "be terse"},
{"role": "user", "content": "hi"}
]),
json!({"maxTokens": 128, "temperature": 0.2}),
);
assert_eq!(
body,
json!({
"inferenceConfig": {"maxTokens": 128, "temperature": 0.2},
"messages": [{"role": "user", "content": [{"text": "hi"}]}],
"system": [{"text": "be terse"}]
})
);
}
#[test]
fn always_emits_inference_config_even_when_empty() {
let body = transform(json!([{"role": "user", "content": "hi"}]), json!({}));
assert_eq!(body["inferenceConfig"], json!({}));
assert!(body.get("system").is_none());
}
#[test]
fn places_only_inference_params_in_inference_config() {
let body = transform(
json!([{"role": "user", "content": "hi"}]),
json!({
"maxTokens": 64,
"temperature": 0.1,
"topP": 0.9,
"stopSequences": ["STOP"]
}),
);
assert_eq!(
body["inferenceConfig"],
json!({"maxTokens": 64, "temperature": 0.1, "topP": 0.9, "stopSequences": ["STOP"]})
);
assert!(body.get("additionalModelRequestFields").is_none());
}
#[test]
fn merges_consecutive_user_turns_into_one_message() {
let body = transform(
json!([
{"role": "user", "content": "one"},
{"role": "user", "content": [{"type": "text", "text": "two"}]},
{"role": "assistant", "content": "ack"},
{"role": "user", "content": "three"}
]),
json!({}),
);
assert_eq!(
body["messages"],
json!([
{"role": "user", "content": [{"text": "one"}, {"text": "two"}]},
{"role": "assistant", "content": [{"text": "ack"}]},
{"role": "user", "content": [{"text": "three"}]}
])
);
}
#[test]
fn declines_streaming() {
assert_eq!(
reason(
json!([{"role": "user", "content": "hi"}]),
json!({"stream": true})
),
Some(Unsupported("streaming"))
);
}
#[test]
fn declines_top_k_because_python_routes_it_by_base_model() {
assert_eq!(
reason(
json!([{"role": "user", "content": "hi"}]),
json!({"topK": 40})
),
Some(Unsupported("unrecognized request parameter"))
);
}
#[test]
fn declines_tools_and_other_params_outside_the_allowlist() {
for param in [
json!({"tools": []}),
json!({"tool_choice": {"auto": {}}}),
json!({"thinking": {"type": "enabled"}}),
json!({"requestMetadata": {"k": "v"}}),
json!({"outputConfig": {}}),
json!({"_parallel_tool_use_config": {}}),
] {
assert_eq!(
reason(json!([{"role": "user", "content": "hi"}]), param.clone()),
Some(Unsupported("unrecognized request parameter")),
"expected {param} to decline"
);
}
}
#[test]
fn declines_blank_text_rather_than_substituting_the_anthropic_placeholder() {
for content in [
json!(""),
json!(" "),
json!([{"type": "text", "text": " "}]),
] {
assert_eq!(
reason(
json!([{"role": "user", "content": content}, {"role": "user", "content": "hi"}]),
json!({})
),
Some(Unsupported("blank message text")),
"expected blank content {content} to decline"
);
}
}
#[test]
fn declines_a_message_whose_content_list_is_empty() {
// The blank-text check scans parts, so an empty list clears it; Converse
// rejects an empty `content` array, which is a decline the core owes the
// host before the call rather than an error after it.
assert_eq!(
reason(json!([{"role": "user", "content": []}]), json!({})),
Some(Unsupported("message without content"))
);
assert_eq!(
reason(
json!([{"role": "user", "content": [{"type": "text", "text": "hi"}]}]),
json!({})
),
None
);
}
#[test]
fn declines_a_conversation_that_opens_or_closes_on_an_assistant_turn() {
assert_eq!(
reason(
json!([
{"role": "assistant", "content": "prefill"},
{"role": "user", "content": "hi"}
]),
json!({})
),
Some(Unsupported(
"conversation does not run user turn to user turn"
))
);
assert_eq!(
reason(
json!([
{"role": "user", "content": "hi"},
{"role": "assistant", "content": "prefill"}
]),
json!({})
),
Some(Unsupported(
"conversation does not run user turn to user turn"
))
);
}
#[test]
fn accepts_a_user_to_user_text_conversation() {
assert_eq!(
reason(
json!([
{"role": "system", "content": "be terse"},
{"role": "user", "content": "hi"},
{"role": "assistant", "content": "hello"},
{"role": "user", "content": "again"}
]),
json!({"maxTokens": 16})
),
None
);
}
#[test]
fn builds_the_converse_url_from_the_region_in_the_model_id() {
let config = &BEDROCK_CHAT_COMPLETIONS_CONFIG;
assert_eq!(
config
.complete_url(None, "us-east-1/anthropic.claude-v2", &Map::new(), &|_| {
None
})
.expect("url builds"),
"https://bedrock-runtime.us-east-1.amazonaws.com/model/anthropic.claude-v2/converse"
);
}
#[test]
fn falls_back_to_the_region_env_then_the_default_region() {
let config = &BEDROCK_CHAT_COMPLETIONS_CONFIG;
let with_env = |key: &str| (key == "AWS_REGION_NAME").then(|| "eu-west-1".to_string());
assert_eq!(
config
.complete_url(None, "anthropic.claude-v2", &Map::new(), &with_env)
.expect("url builds"),
"https://bedrock-runtime.eu-west-1.amazonaws.com/model/anthropic.claude-v2/converse"
);
assert_eq!(
config
.complete_url(None, "anthropic.claude-v2", &Map::new(), &|_| None)
.expect("url builds"),
"https://bedrock-runtime.us-west-2.amazonaws.com/model/anthropic.claude-v2/converse"
);
}
#[test]
fn prefers_an_explicit_runtime_endpoint_over_the_api_base() {
let config = &BEDROCK_CHAT_COMPLETIONS_CONFIG;
let overrides = params(json!({"aws_bedrock_runtime_endpoint": "https://vpce.internal/"}));
assert_eq!(
config
.complete_url(
Some("https://ignored.example"),
"anthropic.claude-v2",
&overrides,
&|_| None
)
.expect("url builds"),
"https://vpce.internal/model/anthropic.claude-v2/converse"
);
}
#[test]
fn signs_with_sigv4_in_the_resolved_region() {
let config = &BEDROCK_CHAT_COMPLETIONS_CONFIG;
assert_eq!(
config
.auth(
None,
"eu-central-1/anthropic.claude-v2",
&Map::new(),
&|_| None
)
.expect("auth resolves"),
ChatCompletionsAuth::AwsSigV4 {
region: "eu-central-1".to_string()
}
);
}
#[test]
fn a_bearer_token_outranks_sigv4_the_way_python_resolves_it() {
// Python's get_request_headers reads `api_key` as the Bedrock bearer token
// and only falls back to the env when the caller passed none, so each case
// pins one of its precedence rules. Signing as the host principal when a
// bearer identity is configured would cross an account and quota boundary.
let bedrock_env =
|key: &str| (key == "AWS_BEARER_TOKEN_BEDROCK").then(|| "from-env".to_string());
let no_env = |_: &str| None;
let resolve = |api_key, env: &dyn Fn(&str) -> Option<String>| {
BEDROCK_CHAT_COMPLETIONS_CONFIG
.auth(
api_key,
"eu-central-1/anthropic.claude-v2",
&Map::new(),
env,
)
.expect("auth resolves")
};
let bearer = |token: &str| ChatCompletionsAuth::Bearer {
token: token.to_string(),
};
let sigv4 = ChatCompletionsAuth::AwsSigV4 {
region: "eu-central-1".to_string(),
};
// A caller-supplied key is the bearer token, and outranks the env.
assert_eq!(
resolve(Some("bedrock-api-key"), &bedrock_env),
bearer("bedrock-api-key")
);
// No key, so the env supplies it.
assert_eq!(resolve(None, &bedrock_env), bearer("from-env"));
// An empty key is not a bearer token, and deliberately does NOT reach for
// the env, which is what Python's `is not None` check does.
assert_eq!(resolve(Some(""), &bedrock_env), sigv4);
// Whitespace is truthy in Python, so it stays a bearer token rather than
// silently becoming a host-credentialed SigV4 request.
assert_eq!(resolve(Some(" "), &no_env), bearer(" "));
// Neither present, so SigV4 as before.
assert_eq!(resolve(None, &no_env), sigv4);
}
#[test]
fn normalizes_a_converse_response_into_openai_shape() {
let response = transform_response(json!({
"output": {"message": {"role": "assistant", "content": [
{"text": "hello"}, {"text": " there"}
]}},
"stopReason": "end_turn",
"usage": {"inputTokens": 11, "outputTokens": 4, "totalTokens": 15}
}))
.expect("response transforms");
assert_eq!(response.model, "anthropic.claude-sonnet-4-5-v1:0");
assert_eq!(
response.choices[0].message.content.as_deref(),
Some("hello there")
);
assert_eq!(response.choices[0].finish_reason, "stop");
assert_eq!(response.usage.prompt_tokens, 11);
assert_eq!(response.usage.completion_tokens, 4);
assert_eq!(response.usage.total_tokens, 15);
}
#[test]
fn maps_converse_stop_reasons_python_maps() {
for (provider_reason, expected) in [
("end_turn", "stop"),
("stop_sequence", "stop"),
("max_tokens", "length"),
("guardrail_intervened", "content_filter"),
// Converse emits this one, and Python's `_FINISH_REASON_MAP` carries
// it. Folding it into `stop` reports a filtered completion as a normal
// one to anything keying on the finish reason.
("content_filtered", "content_filter"),
("content_filter", "content_filter"),
] {
let response = transform_response(json!({
"output": {"message": {"content": [{"text": "x"}]}},
"stopReason": provider_reason,
"usage": {"inputTokens": 1, "outputTokens": 1}
}))
.expect("response transforms");
assert_eq!(
response.choices[0].finish_reason, expected,
"stopReason {provider_reason}"
);
}
}
#[test]
fn reports_an_empty_converse_answer_as_an_empty_string_not_null() {
// Converse assigns the joined text unconditionally
// (`chat_completion_message["content"] = content_str`), unlike Anthropic's
// `merged_text or None`, so an empty answer is `""` on both paths. A caller
// calling `.strip()` on it would break on the Rust path alone. Reachable
// through a filtered or guardrail-intervened response.
for content in [json!([]), json!([{"text": ""}])] {
let response = transform_response(json!({
"output": {"message": {"content": content}},
"stopReason": "content_filtered",
"usage": {"inputTokens": 1, "outputTokens": 0}
}))
.expect("response transforms");
assert_eq!(response.choices[0].message.content, Some(String::new()));
}
}
#[test]
fn reports_the_total_tokens_converse_sent_rather_than_recomputing_them() {
// Python reads `usage["totalTokens"]` straight through here, where Anthropic
// has no such field and adds the two counts instead. The two agree while the
// gate declines every cache_control request, so this is what keeps them
// agreeing if that ever widens.
let response = transform_response(json!({
"output": {"message": {"content": [{"text": "x"}]}},
"stopReason": "end_turn",
"usage": {"inputTokens": 10, "outputTokens": 4, "cacheReadInputTokens": 7, "totalTokens": 14}
}))
.expect("response transforms");
assert_eq!(
response.usage.total_tokens, 14,
"provider total was recomputed"
);
assert_eq!(response.usage.prompt_tokens, 17);
assert_eq!(response.usage.completion_tokens, 4);
}
#[test]
fn falls_back_to_the_computed_total_when_converse_omits_it() {
// Python raises a KeyError on a body with no `totalTokens`. Reporting a zero
// instead would be a worse divergence than the one above, so the computed
// total stands in.
let response = transform_response(json!({
"output": {"message": {"content": [{"text": "x"}]}},
"stopReason": "end_turn",
"usage": {"inputTokens": 10, "outputTokens": 4}
}))
.expect("response transforms");
assert_eq!(response.usage.total_tokens, 14);
}
#[test]
fn declines_a_cache_control_message_so_widening_the_gate_is_a_red_test() {
// Converse only reports cache token counts when the request carries a
// cachePoint block, which is why the provider total and the computed one
// cannot disagree today. This is the tripwire: whoever widens the gate to
// admit prompt caching has to come back and re-check the usage mapping
// rather than discovering a silent number change in production.
assert_eq!(
reason(
json!([{"role": "user", "content": [
{"type": "text", "text": "hi", "cache_control": {"type": "ephemeral"}}
]}]),
json!({})
),
Some(Unsupported("non-text message content"))
);
}
#[test]
fn folds_converse_cache_tokens_into_prompt_tokens() {
let response = transform_response(json!({
"output": {"message": {"content": [{"text": "x"}]}},
"stopReason": "end_turn",
"usage": {
"inputTokens": 10,
"outputTokens": 2,
"cacheReadInputTokens": 5,
"cacheWriteInputTokens": 3
}
}))
.expect("response transforms");
assert_eq!(response.usage.prompt_tokens, 18);
assert_eq!(response.usage.prompt_tokens_details.cached_tokens, 5);
assert_eq!(
response.usage.prompt_tokens_details.cache_creation_tokens,
3
);
assert_eq!(response.usage.prompt_tokens_details.text_tokens, 10);
}
#[test]
fn declines_a_response_carrying_a_tool_use_block() {
let err = transform_response(json!({
"output": {"message": {"content": [
{"toolUse": {"toolUseId": "t1", "name": "f", "input": {}}}
]}},
"stopReason": "tool_use",
"usage": {"inputTokens": 1, "outputTokens": 1}
}))
.expect_err("tool use block");
assert_eq!(
err,
CoreError::Unsupported("non-text response content block")
);
}
#[test]
fn errors_on_a_response_missing_required_fields() {
assert_eq!(
transform_response(json!("nope")).expect_err("not an object"),
CoreError::InvalidResponse("converse response is not an object".to_string())
);
assert_eq!(
transform_response(json!({"usage": {}})).expect_err("no output"),
CoreError::MissingField("output.message.content")
);
assert_eq!(
transform_response(json!({"output": {"message": {"content": []}}})).expect_err("no usage"),
CoreError::MissingField("usage")
);
}
#[test]
fn accepts_aws_call_configuration_without_serializing_it() {
let call_config = json!({
"maxTokens": 16,
"aws_access_key_id": "AKIA",
"aws_secret_access_key": "secret",
"aws_session_token": "token",
"aws_region_name": "us-east-1",
"aws_profile_name": "litellm-stage",
"aws_role_name": "role",
"aws_session_name": "session",
"aws_web_identity_token": "wit",
"aws_sts_endpoint": "https://sts.example",
"aws_external_id": "ext",
"aws_bedrock_runtime_endpoint": "https://vpce.internal"
});
assert_eq!(
reason(
json!([{"role": "user", "content": "hi"}]),
call_config.clone()
),
None
);
let body = transform(json!([{"role": "user", "content": "hi"}]), call_config);
assert_eq!(
body,
json!({
"inferenceConfig": {"maxTokens": 16},
"messages": [{"role": "user", "content": [{"text": "hi"}]}]
}),
"aws call configuration must not reach the Converse body"
);
}
#[test]
fn leaves_a_complete_converse_url_untouched() {
let config = &BEDROCK_CHAT_COMPLETIONS_CONFIG;
let already_built =
"https://bedrock-runtime.us-east-1.amazonaws.com/model/us.anthropic.claude-v2%3A0/converse";
assert_eq!(
config
.complete_url(
Some(already_built),
"anthropic.claude-v2",
&Map::new(),
&|_| None
)
.expect("url builds"),
already_built,
"a host that encoded the model id itself must not have it re-derived"
);
}
#[test]
fn host_supplied_credentials_outrank_ambient_profile_and_role_state() {
use crate::providers::bedrock::aws_base::host_supplied_credentials;
let supplied = params(json!({
"aws_access_key_id": "AKIAHOST",
"aws_secret_access_key": "hostsecret",
"aws_session_token": "hosttoken"
}));
let credentials = host_supplied_credentials(&supplied).expect("host credentials");
assert_eq!(credentials.access_key_id(), "AKIAHOST");
assert_eq!(credentials.secret_access_key(), "hostsecret");
assert_eq!(credentials.session_token(), Some("hosttoken"));
// Without a full static pair there is nothing to honor, so the core falls
// back to deriving credentials itself.
assert!(host_supplied_credentials(&params(json!({"aws_access_key_id": "AKIA"}))).is_none());
assert!(
host_supplied_credentials(&params(
json!({"aws_access_key_id": " ", "aws_secret_access_key": "s"})
))
.is_none()
);
assert!(host_supplied_credentials(&Map::new()).is_none());
}

View file

@ -0,0 +1,297 @@
use serde_json::{Map, Value, json};
use crate::chat_completions::conversation::{Conversation, TurnRole, build_conversation};
use crate::chat_completions::response_utils::{finish_reason_for, unix_now, usage_from_parts};
use crate::chat_completions::transformation::{
ChatCompletionsAuth, ChatCompletionsProviderConfig, Unsupported, unsupported_message,
unsupported_param,
};
use crate::chat_completions::types::{
ChatCompletionsChoice, ChatCompletionsChoiceMessage, ChatCompletionsResponse,
ChatCompletionsUsage, ChatMessage, ChatMessageContent, ProviderChatRequestData,
ProviderChatResponseData,
};
use crate::error::{CoreError, CoreResult};
use super::super::aws_base::{bedrock_model_id_and_region, resolve_bedrock_region};
use super::super::constants::{AWS_BEARER_TOKEN_BEDROCK, BEDROCK_RUNTIME_ENDPOINT_TEMPLATE};
/// Converse parameter names, post `map_openai_params`, that the Rust path can
/// place verbatim in `inferenceConfig`.
///
/// `topK` is deliberately absent: Python routes it to
/// `additionalModelRequestFields` for Anthropic base models and to
/// `inferenceConfig` otherwise, and that branch reads the model catalog the
/// core cannot see.
const SUPPORTED_PARAMS: &[&str] = &["maxTokens", "temperature", "topP", "stopSequences"];
/// Params that belong in `inferenceConfig`, in the order Python's
/// `AmazonConverseConfig` declares them, so bodies compare cleanly.
const INFERENCE_CONFIG_PARAMS: &[&str] = SUPPORTED_PARAMS;
const AWS_BEDROCK_RUNTIME_ENDPOINT: &str = "aws_bedrock_runtime_endpoint";
/// AWS call configuration a host passes down: consumed for signing and endpoint
/// resolution, never serialized into the Converse body.
const CONFIG_PARAMS: &[&str] = &[
"aws_access_key_id",
"aws_secret_access_key",
"aws_session_token",
"aws_region_name",
"aws_session_name",
"aws_profile_name",
"aws_role_name",
"aws_web_identity_token",
"aws_sts_endpoint",
"aws_external_id",
AWS_BEDROCK_RUNTIME_ENDPOINT,
];
const CONVERSE_PATH_SUFFIX: &str = "/converse";
pub struct BedrockChatCompletionsConfig;
pub const BEDROCK_CHAT_COMPLETIONS_CONFIG: BedrockChatCompletionsConfig =
BedrockChatCompletionsConfig;
fn converse_body(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| json!({"text": text})).collect::<Vec<_>>(),
})
})
.collect();
let inference_config = Map::from_iter(INFERENCE_CONFIG_PARAMS.iter().filter_map(|name| {
params
.get(*name)
.map(|value| ((*name).to_string(), value.clone()))
}));
let system: Vec<Value> = conversation
.system
.iter()
.map(|text| json!({"text": text}))
.collect();
Value::Object(Map::from_iter(
[
(
"inferenceConfig".to_string(),
Value::Object(inference_config),
),
("messages".to_string(), json!(messages)),
]
.into_iter()
.chain((!system.is_empty()).then(|| ("system".to_string(), json!(system)))),
))
}
fn has_blank_text(message: &ChatMessage) -> bool {
match &message.content {
None => false,
Some(ChatMessageContent::Text(text)) => text.trim().is_empty(),
Some(ChatMessageContent::Parts(parts)) => parts.iter().any(|part| {
part.get("text")
.and_then(Value::as_str)
.is_none_or(|text| text.trim().is_empty())
}),
}
}
impl ChatCompletionsProviderConfig for BedrockChatCompletionsConfig {
fn complete_url(
&self,
api_base: Option<&str>,
model: &str,
optional_params: &Map<String, Value>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> CoreResult<String> {
let (model_id, model_region) = bedrock_model_id_and_region(model);
let region = resolve_bedrock_region(model_region.as_deref(), optional_params, env_lookup);
let endpoint = optional_params
.get(AWS_BEDROCK_RUNTIME_ENDPOINT)
.and_then(Value::as_str)
.or(api_base)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(str::to_string)
.unwrap_or_else(|| BEDROCK_RUNTIME_ENDPOINT_TEMPLATE.replace("{region}", &region));
let endpoint = endpoint.trim_end_matches('/');
// A host that already built the full Converse URL (LiteLLM's Python
// path encodes the model id itself) passes it through untouched, the
// way the Anthropic config leaves a complete `/v1/messages` URL alone.
if endpoint.ends_with(CONVERSE_PATH_SUFFIX) {
return Ok(endpoint.to_string());
}
Ok(format!("{endpoint}/model/{model_id}{CONVERSE_PATH_SUFFIX}"))
}
fn auth(
&self,
api_key: Option<&str>,
model: &str,
optional_params: &Map<String, Value>,
env_lookup: &dyn Fn(&str) -> Option<String>,
) -> CoreResult<ChatCompletionsAuth> {
// Python reads `api_key` as the Bedrock bearer token and consults the
// env only when the caller passed none, so a caller-supplied empty key
// falls through to SigV4 without reaching for the environment. An
// all-whitespace token stays a bearer token here because Python sends
// it too: treating it as absent would sign as the host principal
// instead, which is the identity swap this branch exists to prevent.
let bearer = match api_key {
Some(key) => Some(key.to_string()),
None => env_lookup(AWS_BEARER_TOKEN_BEDROCK),
}
.filter(|token| !token.is_empty());
if let Some(token) = bearer {
return Ok(ChatCompletionsAuth::Bearer { token });
}
let (_, model_region) = bedrock_model_id_and_region(model);
Ok(ChatCompletionsAuth::AwsSigV4 {
region: resolve_bedrock_region(model_region.as_deref(), optional_params, env_lookup),
})
}
fn default_headers(&self) -> &'static [(&'static str, &'static str)] {
&[("Content-Type", "application/json")]
}
fn supported_params(&self) -> &'static [&'static str] {
SUPPORTED_PARAMS
}
fn config_params(&self) -> &'static [&'static str] {
CONFIG_PARAMS
}
fn unsupported_reason(
&self,
messages: &[ChatMessage],
optional_params: &Map<String, Value>,
) -> Option<Unsupported> {
unsupported_param(SUPPORTED_PARAMS, CONFIG_PARAMS, optional_params)
.or_else(|| messages.iter().find_map(unsupported_message))
// Python's Converse translation drops blank text blocks instead of
// substituting the placeholder the shared conversation builder
// applies, so decline blank text rather than diverge.
.or_else(|| {
messages
.iter()
.any(has_blank_text)
.then_some(Unsupported("blank message text"))
})
// Converse has no assistant prefill: Python inserts a continue turn
// when a conversation opens or closes on an assistant message, and
// only under `litellm.modify_params`, which the core cannot see.
// Declining both ends also keeps the shared builder's final
// assistant right-strip (an Anthropic rule) unreachable here.
.or_else(|| {
let conversation = build_conversation(messages);
let ends_on_assistant = conversation
.turns
.last()
.is_some_and(|turn| turn.role == TurnRole::Assistant);
(!conversation.opens_on_user_turn() || ends_on_assistant).then_some(Unsupported(
"conversation does not run user turn to user turn",
))
})
}
fn transform_request(
&self,
_model: &str,
messages: Vec<ChatMessage>,
optional_params: Map<String, Value>,
) -> CoreResult<ProviderChatRequestData> {
Ok(ProviderChatRequestData {
body: converse_body(&build_conversation(&messages), &optional_params),
})
}
fn transform_response(
&self,
model: &str,
response: ProviderChatResponseData,
) -> CoreResult<ChatCompletionsResponse> {
let body = response.body.as_object().ok_or_else(|| {
CoreError::InvalidResponse("converse response is not an object".into())
})?;
let content = body
.get("output")
.and_then(|output| output.get("message"))
.and_then(|message| message.get("content"))
.and_then(Value::as_array)
.ok_or(CoreError::MissingField("output.message.content"))?;
// The route declines tool requests, so anything other than a text block
// is something this path never asked for. Decline; the host falls back.
if content.iter().any(|block| {
block
.as_object()
.is_none_or(|block| block.len() != 1 || !block.contains_key("text"))
}) {
return Err(CoreError::Unsupported("non-text response content block"));
}
let text: String = content
.iter()
.filter_map(|block| block.get("text").and_then(Value::as_str))
.collect();
let usage = body
.get("usage")
.and_then(Value::as_object)
.ok_or(CoreError::MissingField("usage"))?;
let field = |name: &str| usage.get(name).and_then(Value::as_u64).unwrap_or(0);
let computed = usage_from_parts(
field("inputTokens"),
field("outputTokens"),
field("cacheReadInputTokens"),
field("cacheWriteInputTokens"),
);
// Converse reports `totalTokens` and Python passes it straight through,
// where Anthropic has no such field and Python adds the two counts
// instead, so only this provider overrides the computed total. Python
// does a bare `usage["totalTokens"]` lookup, so a body without the key
// raises there rather than reporting a zero; fall back to the computed
// total, which is the closest thing to that without failing the call.
let usage = ChatCompletionsUsage {
total_tokens: usage
.get("totalTokens")
.and_then(Value::as_u64)
.unwrap_or(computed.total_tokens),
..computed
};
Ok(ChatCompletionsResponse {
created: unix_now(),
// Converse echoes no model id, so Python reports the requested one.
model: model.to_string(),
choices: vec![ChatCompletionsChoice {
index: 0,
message: ChatCompletionsChoiceMessage {
role: "assistant".to_string(),
// Converse assigns the joined string unconditionally, so an
// empty response is `""` here and not `None` as it is on
// Anthropic. A caller calling `.strip()` on it would break
// on this path alone.
content: Some(text),
},
finish_reason: finish_reason_for(
body.get("stopReason").and_then(Value::as_str).unwrap_or(""),
)
.to_string(),
}],
usage,
})
}
}
#[cfg(test)]
#[path = "tests.rs"]
mod tests;

View file

@ -11,6 +11,31 @@ pub const AWS_ROLE_ARN: &str = "AWS_ROLE_ARN";
pub const AWS_WEB_IDENTITY_TOKEN_FILE: &str = "AWS_WEB_IDENTITY_TOKEN_FILE";
pub const AWS_STS_ENDPOINT: &str = "AWS_STS_ENDPOINT";
pub const AWS_EXTERNAL_ID: &str = "AWS_EXTERNAL_ID";
pub const AWS_BEARER_TOKEN_BEDROCK: &str = "AWS_BEARER_TOKEN_BEDROCK";
/// Headers SigV4 covers, beyond the `x-amz-` / `x-amzn-` prefixes. Mirrors
/// Python's `_filter_headers_for_aws_signature` allowlist.
pub const AWS_SIGNED_HEADER_NAMES: &[&str] = &[
"host",
"content-type",
"date",
"x-amz-date",
"x-amz-security-token",
"x-amz-content-sha256",
"x-amz-algorithm",
"x-amz-credential",
"x-amz-signedheaders",
"x-amz-signature",
];
/// Headers the signer emits itself. Mirrors Python's `SIGV4_COMPUTED_HEADERS`,
/// which the reattach loop skips so a caller's copy cannot ride alongside the
/// computed one.
pub const SIGV4_COMPUTED_HEADER_NAMES: &[&str] = &[
"authorization",
"x-amz-date",
"x-amz-security-token",
"date",
];
pub const BEDROCK_SERVICE: &str = "bedrock";
pub const DEFAULT_SESSION_NAME_PREFIX: &str = "litellm-session";
pub const DEFAULT_BEDROCK_REGION: &str = "us-west-2";

View file

@ -5,4 +5,5 @@
#[cfg(feature = "bedrock-auth")]
pub mod audio_transcription;
pub mod aws_base;
pub mod chat_completions;
mod constants;

View file

@ -6,6 +6,10 @@ use litellm_ai_gateway::io::audio_transcription::{
};
use litellm_ai_gateway::io::ocr::{OcrRequest, ocr as run_ocr};
use litellm_ai_gateway::io::responses_ws::ResponsesWebSocketConnection as RustResponsesWebSocketConnection;
use litellm_core::chat_completions::types::{ChatCompletionsRequest, ChatCompletionsResponse};
use litellm_core::chat_completions::{
chat_completions as run_chat_completions, chat_completions_decline_reason,
};
use litellm_core::error::CoreError;
use litellm_core::messages::messages as run_messages;
use litellm_core::messages::types::{AnthropicMessagesResponse, MessagesRequest};
@ -16,6 +20,20 @@ use serde_json::{Map, Value};
mod gil;
pyo3::create_exception!(
_native,
RustBridgeDeclined,
pyo3::exceptions::PyException,
"The route declined before calling the provider, so the host may retry on its own path."
);
pyo3::create_exception!(
_native,
RustUpstreamError,
pyo3::exceptions::PyException,
"The provider call was already issued and failed. Args are (status, message); status is 0 when there was no HTTP response."
);
type MarshaledOcrInputs = (
Value,
Option<Map<String, Value>>,
@ -45,6 +63,15 @@ fn messages_response_to_py(
json_to_py(py, value)
}
fn chat_completions_response_to_py(
py: Python<'_>,
response: ChatCompletionsResponse,
) -> PyResult<Py<PyAny>> {
let value =
serde_json::to_value(response).map_err(|err| PyValueError::new_err(err.to_string()))?;
json_to_py(py, value)
}
fn core_error_to_pyerr(err: CoreError) -> PyErr {
match err {
CoreError::Auth(message) => PyValueError::new_err(message),
@ -56,6 +83,33 @@ fn core_error_to_pyerr(err: CoreError) -> PyErr {
}
}
/// Map a core error for a route whose host keeps a Python implementation.
///
/// The distinction the host needs is whether the provider was already called.
/// Everything raised before the request goes out is safe for the host to retry
/// on its own path; anything after it is not, because the provider has already
/// done the work and billed for it.
fn chat_completions_error_to_pyerr(err: CoreError) -> PyErr {
match err {
CoreError::Unsupported(_)
| CoreError::Auth(_)
| CoreError::InvalidProvider(_)
| CoreError::InvalidRequest(_)
| CoreError::InvalidType { .. }
| CoreError::MissingField(_)
| CoreError::Routing(_)
// Nothing reached the provider, so serving it on Python cannot double
// bill and is the only way the caller gets an answer at all.
| CoreError::Connect(_) => RustBridgeDeclined::new_err(err.to_string()),
CoreError::Http { status, body } => {
RustUpstreamError::new_err((status, format!("{status}: {body}")))
}
CoreError::Network(message) | CoreError::InvalidResponse(message) => {
RustUpstreamError::new_err((0u16, message))
}
}
}
fn optional_object_to_map(
py: Python<'_>,
name: &'static str,
@ -430,6 +484,143 @@ fn amessages(
})
}
type MarshaledChatCompletionsInputs = (
Value,
Map<String, Value>,
Option<Map<String, Value>>,
Option<Duration>,
);
fn marshal_chat_completions_inputs(
py: Python<'_>,
messages: Py<PyAny>,
optional_params: Option<Py<PyAny>>,
extra_headers: Option<Py<PyAny>>,
timeout_seconds: Option<f64>,
) -> PyResult<MarshaledChatCompletionsInputs> {
let messages = py_to_json(py, messages.bind(py))?;
if !messages.is_array() {
return Err(PyValueError::new_err("messages must be a list"));
}
let optional_params = optional_object_to_map(py, "optional_params", optional_params)?;
let extra_headers = match extra_headers {
Some(headers) => Some(optional_object_to_map(py, "extra_headers", Some(headers))?),
None => None,
};
Ok((
messages,
optional_params,
extra_headers,
optional_timeout(timeout_seconds),
))
}
/// The decline reason for this request, or `None` when the Rust path accepts
/// it. Resolves no credentials and performs no I/O, so a host can ask before
/// committing to either path.
#[pyfunction]
#[pyo3(signature = (model, messages, optional_params=None, custom_llm_provider=None))]
fn chat_completions_decline(
py: Python<'_>,
model: String,
messages: Py<PyAny>,
optional_params: Option<Py<PyAny>>,
custom_llm_provider: Option<String>,
) -> PyResult<Option<String>> {
let messages = py_to_json(py, messages.bind(py))?;
let optional_params = optional_object_to_map(py, "optional_params", optional_params)?;
Ok(chat_completions_decline_reason(
&model,
custom_llm_provider.as_deref(),
messages,
&optional_params,
)
.map(str::to_string))
}
#[pyfunction]
#[pyo3(signature = (model, messages, optional_params=None, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, timeout_seconds=None))]
#[allow(clippy::too_many_arguments)]
fn chat_completions(
py: Python<'_>,
model: String,
messages: Py<PyAny>,
optional_params: Option<Py<PyAny>>,
api_key: Option<String>,
api_base: Option<String>,
custom_llm_provider: Option<String>,
extra_headers: Option<Py<PyAny>>,
timeout_seconds: Option<f64>,
) -> PyResult<Py<PyAny>> {
let (messages, optional_params, extra_headers, timeout) = marshal_chat_completions_inputs(
py,
messages,
optional_params,
extra_headers,
timeout_seconds,
)?;
let result = gil::release_gil(py, || {
pyo3_async_runtimes::tokio::get_runtime().block_on(run_chat_completions(
ChatCompletionsRequest {
model: &model,
messages,
optional_params,
api_key: api_key.as_deref(),
api_base: api_base.as_deref(),
custom_llm_provider: custom_llm_provider.as_deref(),
extra_headers,
timeout,
},
))
});
match result {
Ok(response) => chat_completions_response_to_py(py, response),
Err(err) => Err(chat_completions_error_to_pyerr(err)),
}
}
#[pyfunction]
#[pyo3(signature = (model, messages, optional_params=None, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, timeout_seconds=None))]
#[allow(clippy::too_many_arguments)]
fn achat_completions(
py: Python<'_>,
model: String,
messages: Py<PyAny>,
optional_params: Option<Py<PyAny>>,
api_key: Option<String>,
api_base: Option<String>,
custom_llm_provider: Option<String>,
extra_headers: Option<Py<PyAny>>,
timeout_seconds: Option<f64>,
) -> PyResult<Bound<'_, PyAny>> {
let (messages, optional_params, extra_headers, timeout) = marshal_chat_completions_inputs(
py,
messages,
optional_params,
extra_headers,
timeout_seconds,
)?;
pyo3_async_runtimes::tokio::future_into_py(py, async move {
let response = run_chat_completions(ChatCompletionsRequest {
model: &model,
messages,
optional_params,
api_key: api_key.as_deref(),
api_base: api_base.as_deref(),
custom_llm_provider: custom_llm_provider.as_deref(),
extra_headers,
timeout,
})
.await
.map_err(chat_completions_error_to_pyerr)?;
Python::attach(|py| chat_completions_response_to_py(py, response))
})
}
#[pyfunction]
fn gil_stats(py: Python<'_>) -> PyResult<Py<PyAny>> {
let stats = PyDict::new(py);
@ -439,12 +630,18 @@ fn gil_stats(py: Python<'_>) -> PyResult<Py<PyAny>> {
#[pymodule]
fn _native(module: &Bound<'_, PyModule>) -> PyResult<()> {
let py = module.py();
module.add_function(wrap_pyfunction!(ocr, module)?)?;
module.add_function(wrap_pyfunction!(aocr, module)?)?;
module.add_function(wrap_pyfunction!(transcription, module)?)?;
module.add_function(wrap_pyfunction!(atranscription, module)?)?;
module.add_function(wrap_pyfunction!(messages, module)?)?;
module.add_function(wrap_pyfunction!(amessages, module)?)?;
module.add("RustBridgeDeclined", py.get_type::<RustBridgeDeclined>())?;
module.add("RustUpstreamError", py.get_type::<RustUpstreamError>())?;
module.add_function(wrap_pyfunction!(chat_completions_decline, module)?)?;
module.add_function(wrap_pyfunction!(chat_completions, module)?)?;
module.add_function(wrap_pyfunction!(achat_completions, module)?)?;
module.add_class::<ResponsesWebSocketConnection>()?;
module.add_function(wrap_pyfunction!(gil_stats, module)?)?;
Ok(())

View file

@ -453,6 +453,7 @@ max_end_user_budget_id: Optional[str] = None
# backwards compatibility — arbitrary client-supplied identifiers still
# pass through unchanged.
validate_end_user_id_in_db: bool = False
block_requests_for_models_without_pricing: bool = False
disable_end_user_cost_tracking: Optional[bool] = None
disable_end_user_cost_tracking_prometheus_only: Optional[bool] = None
enable_end_user_cost_tracking_prometheus_only: Optional[bool] = None

View file

@ -8,6 +8,12 @@ from logging import Formatter
from typing import Any, Final
import litellm
from litellm.constants import (
LITELLM_TRUNCATED_PAYLOAD_FIELD,
LITELLM_TRUNCATION_STDOUT_SAFEGUARD_NOTE,
MAX_STRING_LENGTH_STDOUT_LOG,
)
from litellm.litellm_core_utils.env_utils import get_env_int
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
from litellm.litellm_core_utils.safe_json_loads import safe_json_loads
from litellm.litellm_core_utils.secret_redaction import redact_string, redact_structured_value
@ -101,7 +107,7 @@ class SecretRedactionFilter(logging.Filter):
# Redact exception tracebacks
if record.exc_info and record.exc_info[1] is not None:
try:
record.exc_text = _redact_string(self._formatter.formatException(record.exc_info))
record.exc_text = _redact_string(record.exc_text or self._formatter.formatException(record.exc_info))
except Exception:
pass
@ -116,6 +122,72 @@ class SecretRedactionFilter(logging.Filter):
_secret_filter: Final = SecretRedactionFilter()
def _get_max_string_length_stdout_log() -> int:
"""Read the limit per record so a value loaded later via proxy config
environment_variables is honored."""
return get_env_int("MAX_STRING_LENGTH_STDOUT_LOG", MAX_STRING_LENGTH_STDOUT_LOG)
def _stdout_truncation_marker(skipped_chars: int) -> str:
return (
f"... ({LITELLM_TRUNCATED_PAYLOAD_FIELD} skipped {skipped_chars} chars. "
f"{LITELLM_TRUNCATION_STDOUT_SAFEGUARD_NOTE}) ..."
)
def _truncate_for_stdout_log(text: str, limit: int) -> str:
kept_chars: Final = limit - len(_stdout_truncation_marker(len(text)))
if kept_chars <= 0:
return text[:limit]
head_chars: Final = kept_chars // 2
tail_chars: Final = kept_chars - head_chars
return f"{text[:head_chars]}{_stdout_truncation_marker(len(text) - kept_chars)}{text[-tail_chars:]}"
class StdoutLogTruncationFilter(logging.Filter):
"""Bounds how much of an oversized log line reaches stdout.
A provider error string can echo the whole request payload, so one failed agentic
request writes hundreds of KB to stdout, repeatedly as the exception propagates from
the router to the proxy handler and into its traceback, all inline on the event loop.
DEBUG records pass through untouched, since dumping full payloads is the point of
`--detailed_debug`, and logging callbacks (OTEL, Datadog, etc.) don't run through
logging filters at all, so they still get the untruncated error.
"""
_formatter = logging.Formatter()
def filter(self, record: logging.LogRecord) -> bool:
if record.levelno < logging.INFO:
return True
limit: Final = _get_max_string_length_stdout_log()
if limit <= 0:
return True
try:
message: Final = record.getMessage()
except (TypeError, ValueError):
return True
if len(message) > limit:
record.msg = _truncate_for_stdout_log(message, limit) # rebind-ok: the Filter interface mutates the record
record.args = None # rebind-ok: args are consumed by the truncated message above
if isinstance(record.exc_info, tuple):
exc_text: Final = record.exc_text or self._formatter.formatException(record.exc_info)
if len(exc_text) > limit:
record.exc_text = _truncate_for_stdout_log( # rebind-ok: the Filter interface mutates the record
exc_text, limit
)
return True
_stdout_truncation_filter: Final = StdoutLogTruncationFilter()
class CorrelationContextFilter(logging.Filter):
"""Stamps each log record with the current request's trace_id and session_id from contextvars.
@ -301,6 +373,7 @@ def _setup_json_exception_handlers(formatter):
error_handler: Final = logging.StreamHandler()
error_handler.setFormatter(formatter)
error_handler.addFilter(_secret_filter)
error_handler.addFilter(_stdout_truncation_filter)
error_handler.addFilter(_correlation_filter)
# Setup excepthook for uncaught exceptions
@ -365,6 +438,12 @@ verbose_router_logger.addHandler(handler)
verbose_proxy_logger.addHandler(handler)
verbose_logger.addHandler(handler)
# Filters attached to the logger, not the handler, survive callers swapping in their own
# handlers (JSON mode, uvicorn log config, a host app's root handler).
verbose_router_logger.addFilter(_stdout_truncation_filter)
verbose_proxy_logger.addFilter(_stdout_truncation_filter)
verbose_logger.addFilter(_stdout_truncation_filter)
def _suppress_loggers():
"""Suppress noisy loggers at INFO level"""

View file

@ -17,6 +17,7 @@ from typing import Final
import redis
import redis.asyncio as async_redis
from redis.credentials import CredentialProvider
from litellm import get_secret, get_secret_str
from litellm._redis_credential_provider import (
@ -134,6 +135,7 @@ def _get_redis_cluster_kwargs(client=None):
"ssl_check_hostname",
"ssl_ca_certs",
"redis_connect_func", # Needed for sync clusters and IAM detection
"credential_provider",
"gcp_service_account",
"gcp_ssl_ca_certs",
"azure_redis_ad_token",
@ -549,14 +551,22 @@ def _init_redis_sentinel(redis_kwargs) -> redis.Redis:
return sentinel.master_for(service_name, **connection_kwargs)
def _sentinel_auth_kwargs(connection_kwargs: dict, sentinel_password: str | None) -> dict:
"""The Sentinel monitors are separate servers that authenticate with their own password, so the
data node's credential provider never belongs on them: leaving it there makes redis-py send the
data node's token to a monitor, which fails whether the monitor is unauthenticated or has its
own password."""
kept: Final = ((k, v) for k, v in connection_kwargs.items() if k != "credential_provider")
return dict(kept, password=sentinel_password)
def _init_async_redis_sentinel(redis_kwargs) -> async_redis.Redis:
sentinel_nodes: Final = redis_kwargs.get("sentinel_nodes")
sentinel_password: Final = redis_kwargs.get("sentinel_password")
service_name: Final = redis_kwargs.get("service_name")
connection_kwargs: Final = _get_redis_sentinel_connection_kwargs(redis_kwargs)
connection_kwargs.setdefault("socket_timeout", REDIS_SOCKET_TIMEOUT)
sentinel_kwargs: Final = dict(connection_kwargs)
sentinel_kwargs["password"] = sentinel_password
sentinel_kwargs: Final = _sentinel_auth_kwargs(connection_kwargs, sentinel_password)
if not sentinel_nodes or not service_name:
raise ValueError("Both 'sentinel_nodes' and 'service_name' are required for Redis Sentinel.")
@ -574,6 +584,36 @@ def _init_async_redis_sentinel(redis_kwargs) -> async_redis.Redis:
return sentinel.master_for(service_name, **connection_kwargs)
def _async_credential_provider(redis_connect_func: object | None) -> CredentialProvider | None:
"""The Azure AD and GCP IAM connect funcs run their AUTH exchange with the blocking client
API, so on an async connection their ``send_command``/``read_response`` calls return
coroutines nobody awaits and every connect fails. Async paths authenticate through a
``CredentialProvider`` instead, which redis-py consults per connection so the token stays
fresh. Any other ``redis_connect_func`` is left where it is, since redis-py awaits it
itself when it is a coroutine function."""
gcp_service_account: Final = getattr(redis_connect_func, "_gcp_service_account", None)
if gcp_service_account is not None:
return GCPIAMCredentialProvider(gcp_service_account)
azure_credential: Final = getattr(redis_connect_func, "_azure_credential", None)
if azure_credential is not None:
return AzureADCredentialProvider(azure_credential, username=os.environ.get("REDIS_USERNAME") or None)
return None
def _async_auth_kwargs(redis_kwargs: dict) -> dict:
"""Swaps a connect func an async path cannot run for the equivalent credential provider,
which supersedes any static username or password redis-py would otherwise reject it with."""
credential_provider: Final = _async_credential_provider(redis_kwargs.get("redis_connect_func"))
if credential_provider is None:
return redis_kwargs
superseded: Final = frozenset({"redis_connect_func", "username", "password"})
kept: Final = ((k, v) for k, v in redis_kwargs.items() if k not in superseded)
return dict(kept, credential_provider=credential_provider) # mutable-ok: the branches below mutate these kwargs
def get_redis_client(**env_overrides):
redis_kwargs: Final = _get_redis_client_logic(**env_overrides)
@ -600,7 +640,7 @@ def get_redis_async_client(
connection_pool: async_redis.BlockingConnectionPool | None = None,
**env_overrides,
) -> async_redis.Redis | async_redis.RedisCluster:
redis_kwargs: Final = _get_redis_client_logic(**env_overrides)
redis_kwargs: Final = _async_auth_kwargs(_get_redis_client_logic(**env_overrides))
if "startup_nodes" in redis_kwargs:
from redis.cluster import ClusterNode
@ -611,28 +651,12 @@ def get_redis_async_client(
if arg in args:
cluster_kwargs[arg] = redis_kwargs[arg]
# Handle GCP IAM authentication for async clusters
redis_connect_func = cluster_kwargs.pop("redis_connect_func", None)
# Use a CredentialProvider so the IAM token is regenerated on every new
# connection — mirrors the sync path where redis_connect_func is invoked
# per connection. Without this, the token would expire after ~1 hour.
if redis_connect_func and hasattr(redis_connect_func, "_gcp_service_account"):
cluster_kwargs["credential_provider"] = GCPIAMCredentialProvider(redis_connect_func._gcp_service_account)
# Handle Azure AD authentication for async clusters via CredentialProvider
# so the credential's internal cache + silent refresh runs per connection
# (mirrors GCP IAM above; avoids static-token-baked-in-pool expiry).
elif redis_connect_func and hasattr(redis_connect_func, "_azure_credential"):
cluster_kwargs["credential_provider"] = AzureADCredentialProvider(
redis_connect_func._azure_credential,
username=os.environ.get("REDIS_USERNAME") or None,
)
new_startup_nodes: Final[list[ClusterNode]] = []
for item in redis_kwargs["startup_nodes"]:
new_startup_nodes.append(ClusterNode(**item))
cluster_kwargs.pop("startup_nodes", None)
cluster_kwargs.pop("redis_connect_func", None)
# Default to a periodic health check + TCP keepalive so a connection silently dropped
# by a cluster restart (e.g. ElastiCache Serverless maintenance) is revalidated and
@ -667,19 +691,6 @@ def get_redis_async_client(
if "sentinel_nodes" in redis_kwargs and "service_name" in redis_kwargs:
return _init_async_redis_sentinel(redis_kwargs)
# Wrap GCP / Azure AD auth in a CredentialProvider for the standard async
# Redis client. The async client doesn't support redis_connect_func, but it
# does honour credential_provider — which is called per connection, so the
# underlying SDK can refresh tokens silently before they expire.
redis_connect_func = redis_kwargs.pop("redis_connect_func", None)
if redis_connect_func and hasattr(redis_connect_func, "_azure_credential"):
redis_kwargs["credential_provider"] = AzureADCredentialProvider(
redis_connect_func._azure_credential,
username=os.environ.get("REDIS_USERNAME") or None,
)
elif redis_connect_func and hasattr(redis_connect_func, "_gcp_service_account"):
redis_kwargs["credential_provider"] = GCPIAMCredentialProvider(redis_connect_func._gcp_service_account)
_pretty_print_redis_config(redis_kwargs=redis_kwargs)
if connection_pool is not None:
@ -693,7 +704,7 @@ def get_redis_async_client(
def get_redis_connection_pool(
**env_overrides,
) -> async_redis.BlockingConnectionPool | None:
redis_kwargs: Final = _get_redis_client_logic(**env_overrides)
redis_kwargs: Final = _async_auth_kwargs(_get_redis_client_logic(**env_overrides))
verbose_logger.debug("get_redis_connection_pool: redis_kwargs", redis_kwargs)
if "startup_nodes" in redis_kwargs:
@ -714,18 +725,6 @@ def get_redis_connection_pool(
)
return async_redis.BlockingConnectionPool.from_url(**pool_kwargs)
# Wrap GCP / Azure AD auth in a CredentialProvider so pool-managed
# connections re-fetch tokens via the SDK's internal cache + silent refresh
# rather than reusing a single token captured at pool creation.
redis_connect_func: Final = redis_kwargs.pop("redis_connect_func", None)
if redis_connect_func and hasattr(redis_connect_func, "_azure_credential"):
redis_kwargs["credential_provider"] = AzureADCredentialProvider(
redis_connect_func._azure_credential,
username=os.environ.get("REDIS_USERNAME") or None,
)
elif redis_connect_func and hasattr(redis_connect_func, "_gcp_service_account"):
redis_kwargs["credential_provider"] = GCPIAMCredentialProvider(redis_connect_func._gcp_service_account)
if redis_kwargs.pop("ssl", None):
redis_kwargs["connection_class"] = async_redis.SSLConnection
return async_redis.BlockingConnectionPool(timeout=REDIS_CONNECTION_POOL_TIMEOUT, **redis_kwargs)

View file

@ -16,6 +16,7 @@ from collections.abc import Sequence
from typing import TYPE_CHECKING, Any, Final
import litellm
from litellm.constants import SEMANTIC_CACHE_EMBEDDING_TIMEOUT_SECONDS
if TYPE_CHECKING:
from litellm.router import Router
@ -60,6 +61,13 @@ def resolve_embedding_max_input_tokens(
return deployment_max_input_tokens
def resolve_embedding_timeout(configured_timeout: float | None) -> float:
"""Explicit cache setting first, else the short semantic-cache default."""
if configured_timeout is not None:
return configured_timeout
return SEMANTIC_CACHE_EMBEDDING_TIMEOUT_SECONDS
def truncate_embedding_input(prompt: str, embedding_model: str, max_input_tokens: int | None) -> str:
"""Keep only the first ``max_input_tokens`` tokens of ``prompt`` for the embedding call."""
if max_input_tokens is None:

View file

@ -98,6 +98,7 @@ class Cache:
qdrant_semantic_cache_embedding_model: str = "text-embedding-ada-002",
qdrant_semantic_cache_vector_size: int | None = None,
semantic_cache_embedding_max_input_tokens: int | None = None,
semantic_cache_embedding_timeout: float | None = None,
# GCP IAM authentication parameters
gcp_service_account: str | None = None,
gcp_ssl_ca_certs: str | None = None,
@ -124,6 +125,7 @@ class Cache:
qdrant_collection_name (str, optional): The name for your qdrant collection. Required if type is "qdrant-semantic".
similarity_threshold (float, optional): The similarity threshold for semantic-caching, Required if type is "redis-semantic" or "qdrant-semantic".
semantic_cache_embedding_max_input_tokens (int, optional): Truncate prompts to this many tokens before embedding them for semantic caching. Defaults to the embedding deployment's configured max_input_tokens.
semantic_cache_embedding_timeout (float, optional): Seconds a semantic-cache lookup may spend embedding the prompt before it gives up and lets the request continue to the LLM. Defaults to SEMANTIC_CACHE_EMBEDDING_TIMEOUT_SECONDS.
# Disk Cache Args
disk_cache_dir (str, optional): The directory for the disk cache. Defaults to None.
@ -195,6 +197,7 @@ class Cache:
embedding_model=redis_semantic_cache_embedding_model,
index_name=redis_semantic_cache_index_name,
embedding_max_input_tokens=semantic_cache_embedding_max_input_tokens,
embedding_timeout=semantic_cache_embedding_timeout,
**kwargs,
)
elif type == LiteLLMCacheType.VALKEY_SEMANTIC:
@ -211,6 +214,7 @@ class Cache:
index_name=valkey_semantic_cache_index_name,
startup_nodes=redis_startup_nodes,
embedding_max_input_tokens=semantic_cache_embedding_max_input_tokens,
embedding_timeout=semantic_cache_embedding_timeout,
**kwargs,
)
elif type == LiteLLMCacheType.QDRANT_SEMANTIC:
@ -223,6 +227,7 @@ class Cache:
embedding_model=qdrant_semantic_cache_embedding_model,
vector_size=qdrant_semantic_cache_vector_size,
embedding_max_input_tokens=semantic_cache_embedding_max_input_tokens,
embedding_timeout=semantic_cache_embedding_timeout,
)
elif type == LiteLLMCacheType.LOCAL:
self.cache = InMemoryCache()

View file

@ -16,7 +16,11 @@ from typing import TYPE_CHECKING, Any, Final, cast
import litellm
from litellm._logging import print_verbose
from litellm.constants import QDRANT_SCALAR_QUANTILE, QDRANT_VECTOR_SIZE
from litellm.constants import (
QDRANT_SCALAR_QUANTILE,
QDRANT_VECTOR_SIZE,
SEMANTIC_CACHE_EMBEDDING_TIMEOUT_SECONDS,
)
from litellm.litellm_core_utils.prompt_templates.common_utils import (
get_str_from_messages,
)
@ -26,6 +30,7 @@ from ._embedding_router import (
build_router_embedding_metadata,
resolve_embedding_max_input_tokens,
resolve_embedding_router,
resolve_embedding_timeout,
truncate_embedding_input,
)
from .base_cache import BaseCache
@ -37,6 +42,7 @@ if TYPE_CHECKING:
class QdrantSemanticCache(BaseCache):
CACHE_KEY_FIELD_NAME = "litellm_cache_key"
embedding_max_input_tokens: int | None = None
embedding_timeout: float = SEMANTIC_CACHE_EMBEDDING_TIMEOUT_SECONDS
def __init__(
self,
@ -49,6 +55,7 @@ class QdrantSemanticCache(BaseCache):
host_type=None,
vector_size=None,
embedding_max_input_tokens: int | None = None,
embedding_timeout: float | None = None,
):
from litellm.llms.custom_httpx.http_handler import (
_get_httpx_client,
@ -68,6 +75,7 @@ class QdrantSemanticCache(BaseCache):
self.similarity_threshold = similarity_threshold
self.embedding_model = embedding_model
self.embedding_max_input_tokens = embedding_max_input_tokens
self.embedding_timeout = resolve_embedding_timeout(embedding_timeout)
self.vector_size = vector_size if vector_size is not None else QDRANT_VECTOR_SIZE
headers = {}
@ -222,11 +230,15 @@ class QdrantSemanticCache(BaseCache):
input=embedding_input,
cache={"no-store": True, "no-cache": True},
metadata=build_router_embedding_metadata(metadata),
timeout=self.embedding_timeout,
num_retries=0,
)
return litellm.embedding(
model=self.embedding_model,
input=embedding_input,
cache={"no-store": True, "no-cache": True},
timeout=self.embedding_timeout,
num_retries=0,
)
async def _get_async_embedding(self, prompt: str, metadata: dict[str, Any] | None = None) -> EmbeddingResponse:
@ -238,19 +250,25 @@ class QdrantSemanticCache(BaseCache):
router: Final = resolve_embedding_router(self.embedding_model, llm_router, llm_model_list)
embedding_input: Final = self._embedding_input(prompt, router)
if router is not None:
return await router.aembedding(
embedding_call: Final = (
router.aembedding(
model=self.embedding_model,
input=embedding_input,
cache={"no-store": True, "no-cache": True},
metadata=build_router_embedding_metadata(metadata),
timeout=self.embedding_timeout,
num_retries=0,
)
if router is not None
else litellm.aembedding(
model=self.embedding_model,
input=embedding_input,
cache={"no-store": True, "no-cache": True},
timeout=self.embedding_timeout,
num_retries=0,
)
return await litellm.aembedding(
model=self.embedding_model,
input=embedding_input,
cache={"no-store": True, "no-cache": True},
)
return await asyncio.wait_for(embedding_call, self.embedding_timeout)
def set_cache(self, key, value, **kwargs):
print_verbose(f"qdrant semantic-cache set_cache, kwargs: {kwargs}")

View file

@ -18,6 +18,7 @@ from typing import TYPE_CHECKING, Any, Final, cast
import litellm
from litellm._logging import print_verbose, verbose_logger
from litellm.constants import SEMANTIC_CACHE_EMBEDDING_TIMEOUT_SECONDS
from litellm.litellm_core_utils.prompt_templates.common_utils import (
get_str_from_messages,
)
@ -27,6 +28,7 @@ from ._embedding_router import (
build_router_embedding_metadata,
resolve_embedding_max_input_tokens,
resolve_embedding_router,
resolve_embedding_timeout,
truncate_embedding_input,
)
from .base_cache import BaseCache
@ -47,6 +49,7 @@ class RedisSemanticCache(BaseCache):
DEFAULT_REDIS_INDEX_NAME: str = "litellm_semantic_cache_index"
CACHE_KEY_FIELD_NAME: str = "litellm_cache_key"
embedding_max_input_tokens: int | None = None
embedding_timeout: float = SEMANTIC_CACHE_EMBEDDING_TIMEOUT_SECONDS
def __init__(
self,
@ -58,6 +61,7 @@ class RedisSemanticCache(BaseCache):
embedding_model: str = "text-embedding-ada-002",
index_name: str | None = None,
embedding_max_input_tokens: int | None = None,
embedding_timeout: float | None = None,
**kwargs: object,
):
"""
@ -74,6 +78,8 @@ class RedisSemanticCache(BaseCache):
index_name: Name for the Redis index
embedding_max_input_tokens: Truncate prompts to this many tokens before
embedding; defaults to the Router deployment's configured max_input_tokens
embedding_timeout: Seconds a cache lookup may spend embedding the prompt before it
gives up and lets the request continue to the LLM
ttl: Default time-to-live for cache entries in seconds
**kwargs: Additional arguments passed to the Redis client
@ -99,6 +105,7 @@ class RedisSemanticCache(BaseCache):
self.distance_threshold = 1 - similarity_threshold
self.embedding_model = embedding_model
self.embedding_max_input_tokens = embedding_max_input_tokens
self.embedding_timeout = resolve_embedding_timeout(embedding_timeout)
# Set up Redis connection
if redis_url is None:
@ -349,6 +356,8 @@ class RedisSemanticCache(BaseCache):
input=embedding_input,
cache={"no-store": True, "no-cache": True},
metadata=build_router_embedding_metadata(metadata),
timeout=self.embedding_timeout,
num_retries=0,
),
)
else:
@ -358,6 +367,8 @@ class RedisSemanticCache(BaseCache):
model=self.embedding_model,
input=embedding_input,
cache={"no-store": True, "no-cache": True},
timeout=self.embedding_timeout,
num_retries=0,
),
)
return embedding_response["data"][0]["embedding"]
@ -512,20 +523,26 @@ class RedisSemanticCache(BaseCache):
router: Final = resolve_embedding_router(self.embedding_model, llm_router, llm_model_list)
embedding_input: Final = self._embedding_input(prompt, router)
embedding_call: Final = (
router.aembedding(
model=self.embedding_model,
input=embedding_input,
cache={"no-store": True, "no-cache": True},
metadata=build_router_embedding_metadata(metadata),
timeout=self.embedding_timeout,
num_retries=0,
)
if router is not None
else litellm.aembedding(
model=self.embedding_model,
input=embedding_input,
cache={"no-store": True, "no-cache": True},
timeout=self.embedding_timeout,
num_retries=0,
)
)
try:
if router is not None:
embedding_response = await router.aembedding(
model=self.embedding_model,
input=embedding_input,
cache={"no-store": True, "no-cache": True},
metadata=build_router_embedding_metadata(metadata),
)
else:
embedding_response = await litellm.aembedding(
model=self.embedding_model,
input=embedding_input,
cache={"no-store": True, "no-cache": True},
)
embedding_response: Final = await asyncio.wait_for(embedding_call, self.embedding_timeout)
return embedding_response["data"][0]["embedding"]
except Exception as e:
print_verbose(f"Error generating async embedding: {e}")

View file

@ -30,6 +30,7 @@ from litellm._logging import print_verbose
from litellm._uuid import uuid
from litellm.llms.valkey.common_utils import build_valkey_url, pack_vector
from ._embedding_router import resolve_embedding_timeout
from .redis_semantic_cache import RedisSemanticCache
@ -62,6 +63,7 @@ class ValkeySemanticCache(RedisSemanticCache):
sync_client: Redis | None = None,
async_client: AsyncRedis | None = None,
embedding_max_input_tokens: int | None = None,
embedding_timeout: float | None = None,
**kwargs: Any,
):
if similarity_threshold is None:
@ -80,6 +82,7 @@ class ValkeySemanticCache(RedisSemanticCache):
self.similarity_threshold = similarity_threshold
self.embedding_model = embedding_model
self.embedding_max_input_tokens = embedding_max_input_tokens
self.embedding_timeout = resolve_embedding_timeout(embedding_timeout)
self.index_name = index_name or self.DEFAULT_VALKEY_INDEX_NAME
self.key_prefix = f"{self.index_name}:"
self._index_dim: int | None = None

View file

@ -25,6 +25,11 @@ class ResponsesToCompletionBridgeHandlerInputKwargs(TypedDict):
encoding: object
def _restore_routing_prefix(model: str, custom_llm_provider: str) -> str:
"""`responses()` runs `get_llm_provider()` itself, so hand back the prefixed model `completion()` started from."""
return f"{custom_llm_provider}/{model}"
class ResponsesToCompletionBridgeHandler:
def __init__(self):
from .transformation import LiteLLMResponsesTransformationHandler
@ -184,14 +189,11 @@ class ResponsesToCompletionBridgeHandler:
client=kwargs.get("client"),
)
# Pin the resolved provider so `responses()` doesn't re-run
# `get_llm_provider()` on the model string and strip a second
# provider prefix (see GitHub issue #28505). request_data already
# carries `custom_llm_provider` via the spread of
# `sanitized_litellm_params`; overwriting it on the dict (rather
# than adding an explicit kwarg) avoids the duplicate-keyword
# TypeError that would otherwise fire on the real bridge path.
# Set on request_data rather than passed as explicit kwargs: the spread of
# `sanitized_litellm_params` already carries both, so passing them again
# would raise a duplicate-keyword TypeError.
request_data["custom_llm_provider"] = custom_llm_provider
request_data["model"] = _restore_routing_prefix(model, custom_llm_provider)
result: Final = responses(
**request_data,
)
@ -282,13 +284,11 @@ class ResponsesToCompletionBridgeHandler:
except Exception as e:
raise e
# Pin the resolved provider so `aresponses()` doesn't re-run
# `get_llm_provider()` on the model string and strip a second
# provider prefix (see GitHub issue #28505). Set on request_data
# rather than passed as a separate kwarg to avoid the duplicate-
# keyword TypeError when `sanitized_litellm_params` already
# carries `custom_llm_provider`.
# Set on request_data rather than passed as explicit kwargs: the spread of
# `sanitized_litellm_params` already carries both, so passing them again
# would raise a duplicate-keyword TypeError.
request_data["custom_llm_provider"] = custom_llm_provider
request_data["model"] = _restore_routing_prefix(model, custom_llm_provider)
result: Final = await aresponses(
**request_data,
aresponses=True,

View file

@ -113,6 +113,58 @@ def _build_reasoning_item(
}
def _reasoning_item_from_output_item(item: object) -> _BuiltReasoningItem | None:
from openai.types.responses import ResponseReasoningItem
if isinstance(item, ResponseReasoningItem):
return _build_reasoning_item(
item_id=item.id,
encrypted_content=getattr(item, "encrypted_content", None),
summary_raw=item.summary,
)
if isinstance(item, dict) and item.get("type") == "reasoning":
return _build_reasoning_item(
item_id=item.get("id", ""),
encrypted_content=item.get("encrypted_content"),
summary_raw=item.get("summary"),
)
return None
def _reasoning_items_from_output_items(output_items: Sequence[object]) -> tuple[_BuiltReasoningItem, ...]:
return tuple(
reasoning_item
for reasoning_item in (_reasoning_item_from_output_item(item) for item in output_items)
if reasoning_item is not None
)
def _as_chat_reasoning_items(
reasoning_items: Sequence[_BuiltReasoningItem],
) -> list[ChatCompletionReasoningItem] | None:
if not reasoning_items:
return None
# cast-ok: _BuiltReasoningItem is the structural shape ChatCompletionReasoningItem
# describes, and TypedDict invariance is what stops the two from unifying here.
return cast(list[ChatCompletionReasoningItem], list(reasoning_items))
def _map_incomplete_reason_to_finish_reason(incomplete_reason: str | None) -> Literal["length", "content_filter"]:
if incomplete_reason == "content_filter":
return "content_filter"
return "length"
def _incomplete_reason_from_response_payload(response_payload: object) -> str | None:
if not isinstance(response_payload, Mapping):
return None
incomplete_details: Final = response_payload.get("incomplete_details")
if not isinstance(incomplete_details, Mapping):
return None
reason: Final = incomplete_details.get("reason")
return reason if isinstance(reason, str) else None
class _ChatToolCallDict(ChatCompletionToolCallChunk, total=False):
provider_specific_fields: Mapping[str, object]
@ -657,6 +709,27 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
return choices
@staticmethod
def _build_empty_incomplete_choice(
output_items: Sequence[object],
finish_reason: Literal["length", "content_filter"],
) -> "Choices":
from litellm.types.utils import Choices, Message
reasoning_items: Final = _reasoning_items_from_output_items(output_items)
reasoning_content: Final = " ".join(
summary_block["text"]
for reasoning_item in reasoning_items
for summary_block in reasoning_item["summary"]
if summary_block.get("text")
)
message: Final = Message(
content="",
reasoning_content=reasoning_content if reasoning_content else None,
reasoning_items=_as_chat_reasoning_items(reasoning_items),
)
return Choices(message=message, finish_reason=finish_reason, index=0)
@classmethod
def _extract_output_from_completed_event(cls, parsed_chunk: Mapping[str, object]) -> list[dict[str, object]] | None:
response_payload: Final = parsed_chunk.get("response")
@ -763,11 +836,22 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
handle_raw_dict_callback=self._handle_raw_dict_response_item,
)
if len(choices) == 0:
if raw_response.incomplete_details is not None and raw_response.incomplete_details.reason is not None:
raise ValueError(f"{model} unable to complete request: {raw_response.incomplete_details.reason}")
response_is_incomplete: Final = raw_response.status == "incomplete" or (
raw_response.incomplete_details is not None and raw_response.incomplete_details.reason is not None
)
if len(choices) == 0 and not response_is_incomplete:
raise ValueError(f"Unknown items in responses API response: {output_items}")
if response_is_incomplete:
incomplete_finish_reason: Final = _map_incomplete_reason_to_finish_reason(
raw_response.incomplete_details.reason if raw_response.incomplete_details is not None else None
)
if len(choices) == 0:
choices.append(self._build_empty_incomplete_choice(output_items, incomplete_finish_reason))
else:
raise ValueError(f"Unknown items in responses API response: {output_items}")
for choice in choices:
choice.finish_reason = incomplete_finish_reason
setattr(model_response, "choices", choices)
@ -1392,12 +1476,7 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator):
)
]
)
elif event_type == "response.completed":
# Response is fully complete - now we can signal is_finished=True
# This ensures we don't prematurely end the stream before tool_calls arrive
# Check if response contains function_call items in output
# to determine correct finish_reason
elif event_type in ("response.completed", "response.incomplete"):
response_data: Final = parsed_chunk.get("response", {})
output_items: Final = response_data.get("output", []) if response_data else []
@ -1407,25 +1486,14 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator):
if isinstance(item, dict)
)
finish_reason: Final = "tool_calls" if has_function_calls else "stop"
finish_reason: Final = (
_map_incomplete_reason_to_finish_reason(_incomplete_reason_from_response_payload(response_data))
if event_type == "response.incomplete"
else ("tool_calls" if has_function_calls else "stop")
)
# Extract reasoning items with encrypted_content for round-tripping
completed_reasoning_items: list[_BuiltReasoningItem] | None = None
for item in output_items:
if not isinstance(item, dict) or item.get("type") != "reasoning":
continue
if completed_reasoning_items is None:
completed_reasoning_items = []
completed_reasoning_items.append(
_build_reasoning_item(
item_id=item.get("id", ""),
encrypted_content=item.get("encrypted_content"),
summary_raw=item.get("summary"),
)
)
completed_reasoning_items_typed: Final = cast(
list[ChatCompletionReasoningItem] | None,
completed_reasoning_items,
terminal_reasoning_items_typed: Final = _as_chat_reasoning_items(
_reasoning_items_from_output_items(output_items)
)
usage = None
@ -1439,7 +1507,7 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator):
index=0,
delta=Delta(
content="",
reasoning_items=completed_reasoning_items_typed,
reasoning_items=terminal_reasoning_items_typed,
),
finish_reason=finish_reason,
)

View file

@ -3,7 +3,7 @@ import sys
from types import MappingProxyType
from typing import Final, Literal
from litellm.litellm_core_utils.env_utils import get_env_int, get_env_int_or_none
from litellm.litellm_core_utils.env_utils import get_env_int, get_env_int_in_range, get_env_int_or_none
DEFAULT_HEALTH_CHECK_PROMPT: Final = str(os.getenv("DEFAULT_HEALTH_CHECK_PROMPT", "test from litellm"))
AZURE_DEFAULT_RESPONSES_API_VERSION: Final = str(os.getenv("AZURE_DEFAULT_RESPONSES_API_VERSION", "preview"))
@ -49,6 +49,8 @@ LITELLM_MAX_STREAMING_DURATION_SECONDS: Final = (
# Set to 0 to disable truncation.
MAX_BASE64_LENGTH_FOR_LOGGING: Final = int(os.getenv("MAX_BASE64_LENGTH_FOR_LOGGING", 64))
MAX_STRING_LENGTH_STDOUT_LOG: Final = get_env_int("MAX_STRING_LENGTH_STDOUT_LOG", 4096)
# When true, adds detailed per-phase timing breakdown headers to responses.
# Headers: x-litellm-timing-{pre-processing,llm-api,post-processing,message-copy}-ms
LITELLM_DETAILED_TIMING: Final = os.getenv("LITELLM_DETAILED_TIMING", "false").lower() == "true"
@ -243,6 +245,12 @@ AIOHTTP_NEEDS_CLEANUP_CLOSED: Final = (3, 13, 0) <= sys.version_info < (
# https://github.com/openai/openai-agents-python/blob/cf1b933660e44fd37b4350c41febab8221801409/src/agents/realtime/openai_realtime.py#L235
_max_size_env: Final = os.getenv("REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES")
REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES: Final = int(_max_size_env) if _max_size_env is not None else None
REALTIME_CREDENTIAL_RESOLUTION_TIMEOUT_SECONDS: Final = float(
os.getenv("REALTIME_CREDENTIAL_RESOLUTION_TIMEOUT_SECONDS", "20.0")
)
# RFC 6455 caps the close frame payload at 125 bytes, 2 of which carry the status code
WEBSOCKET_CLOSE_REASON_MAX_BYTES: Final = 123
# SSL/TLS cipher configuration for faster handshakes
# Strategy: Strongly prefer fast modern ciphers, but allow fallback to commonly supported ones
@ -317,6 +325,17 @@ DEFAULT_MOCK_RESPONSE_PROMPT_TOKEN_COUNT: Final = int(os.getenv("DEFAULT_MOCK_RE
DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT: Final = int(os.getenv("DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT", 20))
MAX_SHORT_SIDE_FOR_IMAGE_HIGH_RES: Final = int(os.getenv("MAX_SHORT_SIDE_FOR_IMAGE_HIGH_RES", 768))
MAX_LONG_SIDE_FOR_IMAGE_HIGH_RES: Final = int(os.getenv("MAX_LONG_SIDE_FOR_IMAGE_HIGH_RES", 2000))
# tiktoken's BPE merge loop is quadratic in the length of a single regex piece, so a long run of one
# repeated character (dot leaders, whitespace, zero-padded base64) can take minutes on a multi-MB payload.
# Encoding in chunks makes the cost linear, at a drift of at most ~1 token per chunk boundary. The upper
# bound keeps a misconfigured chunk size from restoring the quadratic cost this exists to remove.
TIKTOKEN_ENCODE_MAX_CHUNK_SIZE_CHARS: Final = 4096
TIKTOKEN_ENCODE_CHUNK_SIZE_CHARS: Final = get_env_int_in_range(
"TIKTOKEN_ENCODE_CHUNK_SIZE_CHARS",
default=1024,
minimum=1,
maximum=TIKTOKEN_ENCODE_MAX_CHUNK_SIZE_CHARS,
)
MAX_TILE_WIDTH: Final = int(os.getenv("MAX_TILE_WIDTH", 512))
MAX_TILE_HEIGHT: Final = int(os.getenv("MAX_TILE_HEIGHT", 512))
OPENAI_FILE_SEARCH_COST_PER_1K_CALLS: Final = float(os.getenv("OPENAI_FILE_SEARCH_COST_PER_1K_CALLS", 2.5 / 1000))
@ -417,6 +436,9 @@ DEFAULT_REQUEST_TIMEOUT_SECONDS: Final[float] = 6000.0
# deadline and connect handshake (see ``http_handler`` cached handler paths).
COMPLETION_HTTP_FALLBACK_SECONDS: Final[float] = 600.0
HTTP_HANDLER_CONNECT_TIMEOUT_SECONDS: Final[float] = 5.0
SEMANTIC_CACHE_EMBEDDING_TIMEOUT_SECONDS: Final[float] = float(
os.getenv("SEMANTIC_CACHE_EMBEDDING_TIMEOUT_SECONDS", "5.0")
)
request_timeout: float = float(os.getenv("REQUEST_TIMEOUT", str(int(DEFAULT_REQUEST_TIMEOUT_SECONDS))))
request_timeout_explicitly_set: bool = "REQUEST_TIMEOUT" in os.environ
DEFAULT_A2A_AGENT_TIMEOUT: Final[float] = float(os.getenv("DEFAULT_A2A_AGENT_TIMEOUT", 6000)) # 10 minutes
@ -461,6 +483,9 @@ MAX_TIME_TO_CLEAR_QUEUE: Final = float(os.getenv("MAX_TIME_TO_CLEAR_QUEUE", 5.0)
LOGGING_WORKER_AGGRESSIVE_CLEAR_COOLDOWN_SECONDS: Final = float(
os.getenv("LOGGING_WORKER_AGGRESSIVE_CLEAR_COOLDOWN_SECONDS", 0.5)
) # Cooldown time in seconds before allowing another aggressive clear (default: 0.5s)
LOGGING_EXECUTOR_MAX_THREADS: Final = get_env_int("LOGGING_EXECUTOR_MAX_THREADS", 100)
LOGGING_EXECUTOR_MAX_PENDING_TASKS: Final = get_env_int("LOGGING_EXECUTOR_MAX_PENDING_TASKS", 10_000)
LOGGING_EXECUTOR_DROPPED_TASK_LOG_INTERVAL_SECONDS: Final = 30.0
DD_TRACER_STREAMING_CHUNK_YIELD_RESOURCE: Final = os.getenv(
"DD_TRACER_STREAMING_CHUNK_YIELD_RESOURCE", "streaming.chunk.yield"
)
@ -757,6 +782,7 @@ openai_compatible_endpoints: Final[list] = [
"https://api.libertai.io/v1",
"https://pinstripes.io/v1",
"https://api.meta.ai/v1",
"https://api.cognition.ai/v1",
]
@ -824,6 +850,7 @@ openai_compatible_providers: Final[list] = [
"pinstripes", # Pinstripes - JSON-configured provider
"darkbloom",
"meta", # Meta Model API (Muse Spark) - JSON-configured provider
"cognition",
]
openai_text_completion_compatible_providers: Final[list] = [ # providers that support `/v1/completions`
"together_ai",
@ -1336,6 +1363,11 @@ LITELLM_TRUNCATION_DB_SAFEGUARD_NOTE: Final = (
"Full, untruncated data is logged to logging callbacks (OTEL, Datadog, etc.). "
"To increase the truncation limit, set `MAX_STRING_LENGTH_PROMPT_IN_DB` in your env."
)
LITELLM_TRUNCATION_STDOUT_SAFEGUARD_NOTE: Final = (
"Truncation is a stdout logging safeguard. "
"Full, untruncated data is logged to logging callbacks (OTEL, Datadog, etc.) and at DEBUG level. "
"To increase the truncation limit, set `MAX_STRING_LENGTH_STDOUT_LOG` in your env."
)
########################### LiteLLM Proxy Specific Constants ###########################
########################################################################################
@ -1510,6 +1542,11 @@ DEFAULT_CRON_JOB_LOCK_TTL_SECONDS: Final = int(os.getenv("DEFAULT_CRON_JOB_LOCK_
PROXY_BUDGET_RESCHEDULER_MIN_TIME: Final = int(os.getenv("PROXY_BUDGET_RESCHEDULER_MIN_TIME", 597))
RESET_BUDGET_JOB_BATCH_SIZE: Final = max(1, int(os.getenv("RESET_BUDGET_JOB_BATCH_SIZE", "500")))
RESET_BUDGET_JOB_MAX_CHUNKS_PER_RUN: Final = max(1, int(os.getenv("RESET_BUDGET_JOB_MAX_CHUNKS_PER_RUN", "100")))
RESET_BUDGET_JOB_NAME: Final = "reset_budget_job"
# Comfortably longer than one PROXY_BUDGET_RESCHEDULER_MIN_TIME tick, so a healthy
# leader keeps the lease across its own run, and a crashed one strands the sweep for
# at most a single tick.
RESET_BUDGET_JOB_LOCK_TTL_SECONDS: Final[int] = 900
PROXY_BATCH_POLLING_INTERVAL: Final = int(os.getenv("PROXY_BATCH_POLLING_INTERVAL", 3600))
MAX_OBJECTS_PER_POLL_CYCLE: Final = max(1, int(os.getenv("MAX_OBJECTS_PER_POLL_CYCLE", 50)))
MANAGED_OBJECT_STALENESS_CUTOFF_DAYS: Final = max(1, int(os.getenv("MANAGED_OBJECT_STALENESS_CUTOFF_DAYS", 7)))
@ -1574,6 +1611,7 @@ LITELLM_SETTINGS_SAFE_DB_OVERRIDES: Final = [
"public_model_groups_links",
"cost_discount_config",
"cost_margin_config",
"block_requests_for_models_without_pricing",
"budget_exceeded_throttle_percentage",
# Every field editable from the Admin UI (proxy_server._GENERAL_SETTINGS_UI_LITELLM_FIELDS)
# must be listed here so a DB write from one worker overrides the live litellm attribute on

View file

@ -1039,16 +1039,29 @@ class LiteLLMUnknownProvider(BadRequestError):
class GuardrailRaisedException(Exception):
"""
Raised both when a guardrail judged content and when it could not judge it at all, since a
guardrail that fails closed refuses the request the same way a policy violation does.
``blocked_content`` separates the two. Set it only where the guardrail actually reached a
verdict on the payload; leave it alone for an unreachable backend, a timeout, or a response
the integration could not parse. Callers that treat a block as something other than a plain
failure, such as the batch path dropping one record and submitting the rest, must gate on it,
because dropping a record no guardrail ever inspected is a silent loss of enforcement.
"""
def __init__(
self,
guardrail_name: str | None = None,
message: str = "",
should_wrap_with_default_message: bool = True,
status_code: int = 400,
blocked_content: bool = False,
):
default_message: Final = f"Guardrail raised an exception, Guardrail: {guardrail_name}, Message: {message}"
self.guardrail_name = guardrail_name
self.status_code = status_code
self.blocked_content = blocked_content
self.message = default_message if should_wrap_with_default_message else message
super().__init__(self.message)

View file

@ -51,6 +51,42 @@ def to_basic_auth(auth_value: str) -> str:
return base64.b64encode(auth_value.encode("utf-8")).decode()
def strip_auth_scheme(auth_value: str, scheme: str) -> str:
"""Return ``auth_value`` with a leading ``<scheme> `` removed, or unchanged when absent.
Callers supply both a bare credential and a complete header value, so prefixing
unconditionally yields ``Bearer Bearer <jwt>``. Scheme names are case-insensitive per
RFC 7235. A credential is required after the scheme, so both a token that merely begins
with the scheme text and a scheme with nothing behind it are returned untouched.
Surrounding whitespace is left to ``_strip_header_whitespace`` at header-build time.
"""
scheme_name, _, remainder = auth_value.lstrip().partition(" ")
credential: Final = remainder.lstrip()
if credential and scheme_name.lower() == scheme.lower():
return credential
return auth_value
def to_basic_credentials(auth_value: str) -> str:
"""Return the base64 credentials for a ``Basic`` header, encoding only when needed.
``Basic <credentials>`` carries credentials that are already encoded, so encoding the whole
value again would bury the scheme inside the payload. This has to run before
:func:`to_basic_auth` rather than at header-build time, where no prefix is left to find.
A schemed value whose remainder does not decode is the bare ``username:password`` shape with
the scheme written in front of it, and is encoded rather than forwarded as an invalid header;
a pair always contains ``:``, which is outside the base64 alphabet, so the two never collide.
"""
credentials: Final = strip_auth_scheme(auth_value, "Basic")
if credentials == auth_value:
return to_basic_auth(auth_value)
try:
base64.b64decode(credentials, validate=True)
except ValueError:
return to_basic_auth(credentials)
return credentials
def _strip_header_whitespace(headers: dict[str, str]) -> dict[str, str]:
return {
(key.strip() if isinstance(key, str) else key): (value.strip() if isinstance(value, str) else value)
@ -441,16 +477,15 @@ class MCPClient:
except BaseException as e:
verbose_logger.debug("Error during http_client cleanup: %s", e)
def update_auth_value(self, mcp_auth_value: str | dict[str, str]):
def update_auth_value(self, mcp_auth_value: str | dict[str, str]) -> None:
"""
Set the authentication header for the MCP client.
"""
if isinstance(mcp_auth_value, dict):
self._mcp_auth_value = mcp_auth_value
elif self.auth_type == MCPAuth.basic:
self._mcp_auth_value = to_basic_credentials(mcp_auth_value)
else:
if self.auth_type == MCPAuth.basic:
# Assuming mcp_auth_value is in format "username:password", convert it when updating
mcp_auth_value = to_basic_auth(mcp_auth_value)
self._mcp_auth_value = mcp_auth_value
def _get_auth_headers(self) -> dict:
@ -459,19 +494,20 @@ class MCPClient:
if self._mcp_auth_value:
if isinstance(self._mcp_auth_value, str):
if self.auth_type == MCPAuth.bearer_token:
headers["Authorization"] = f"Bearer {self._mcp_auth_value}"
headers["Authorization"] = f"Bearer {strip_auth_scheme(self._mcp_auth_value, 'Bearer')}"
elif self.auth_type == MCPAuth.basic:
headers["Authorization"] = f"Basic {self._mcp_auth_value}"
elif self.auth_type == MCPAuth.api_key:
headers["X-API-Key"] = self._mcp_auth_value
elif self.auth_type == MCPAuth.authorization:
# This auth type means the caller owns the whole header value.
headers["Authorization"] = self._mcp_auth_value
elif self.auth_type == MCPAuth.oauth2:
headers["Authorization"] = f"Bearer {self._mcp_auth_value}"
headers["Authorization"] = f"Bearer {strip_auth_scheme(self._mcp_auth_value, 'Bearer')}"
elif self.auth_type == MCPAuth.token:
headers["Authorization"] = f"token {self._mcp_auth_value}"
headers["Authorization"] = f"token {strip_auth_scheme(self._mcp_auth_value, 'token')}"
elif self.auth_type == MCPAuth.oauth2_token_exchange:
headers["Authorization"] = f"Bearer {self._mcp_auth_value}"
headers["Authorization"] = f"Bearer {strip_auth_scheme(self._mcp_auth_value, 'Bearer')}"
elif isinstance(self._mcp_auth_value, dict):
headers.update(self._mcp_auth_value)
# Note: aws_sigv4 auth is not handled here — SigV4 requires per-request

View file

@ -10,17 +10,35 @@ Supported for both `v1/chat/completions` (via the prompt-management hook) and
"""
import copy
import os
import re
from collections.abc import Iterable, Mapping, Sequence
from typing import TYPE_CHECKING, Any, Final, cast
from urllib.parse import urlparse
from litellm._logging import verbose_logger
from litellm.integrations.custom_logger import CustomLogger
from litellm.integrations.custom_prompt_management import CustomPromptManagement
from litellm.integrations.prompt_management_base import PromptManagementClient
from litellm.litellm_core_utils.prompt_templates.common_utils import (
with_prompt_cache_breakpoint,
)
from litellm.types.integrations.anthropic_cache_control_hook import (
CacheControlInjectionPoint,
CacheControlMessageInjectionPoint,
)
from litellm.types.llms.openai import AllMessageValues, ChatCompletionCachedContent
from litellm.types.llms.anthropic import (
AllAnthropicToolsValues,
AnthropicSystemMessageContent,
)
from litellm.types.llms.openai import (
AllMessageValues,
ChatCompletionCachedContent,
ChatCompletionTextObject,
ChatCompletionToolParam,
PromptCacheBreakpoint,
PromptCacheOptions,
)
from litellm.types.prompts.init_prompts import PromptSpec
from litellm.types.utils import StandardCallbackDynamicParams
@ -34,6 +52,57 @@ else:
# breakpoints: "A maximum of 4 blocks with cache_control may be provided."
MAX_CACHE_CONTROL_BLOCKS: Final = 4
CACHE_BREAKPOINT_KEYS: Final = ("cache_control", "prompt_cache_breakpoint")
OPENAI_PROMPT_CACHE_BREAKPOINT_MIN_GPT_VERSION: Final = (5, 6)
_GPT_VERSION_PATTERN: Final = re.compile(r"^gpt-(\d+)(?:\.(\d+))?")
OPENAI_PROMPT_CACHE_BREAKPOINT_BLOCK_TYPES: Final = frozenset(
{"text", "image", "image_url", "file", "input_audio", "input_text", "input_image", "input_file"}
)
OPENAI_API_HOST: Final = "api.openai.com"
OPENAI_API_BASE_ENV_VARS: Final = ("OPENAI_BASE_URL", "OPENAI_API_BASE")
AllToolParamValues = ChatCompletionToolParam | AllAnthropicToolsValues
def supports_openai_prompt_cache_breakpoint(model: str) -> bool:
model_map_flag: Final = _model_map_prompt_cache_breakpoint_flag(model)
if model_map_flag is not None:
return model_map_flag
version_match: Final = _GPT_VERSION_PATTERN.match(model.rsplit("/", 1)[-1].lower())
if version_match is None:
return False
version: Final = (int(version_match.group(1)), int(version_match.group(2) or 0))
return version >= OPENAI_PROMPT_CACHE_BREAKPOINT_MIN_GPT_VERSION
def _model_map_prompt_cache_breakpoint_flag(model: str) -> bool | None:
import litellm
entries: Final = (litellm.model_cost.get(key) for key in (model, model.rsplit("/", 1)[-1]))
flags: Final = (entry.get("supports_prompt_cache_breakpoint") for entry in entries if isinstance(entry, dict))
return next((bool(flag) for flag in flags if flag is not None), None)
def targets_openai_api(api_base: object) -> bool:
import litellm
resolved: Final = next(
(value for value in (api_base, litellm.api_base, *map(os.getenv, OPENAI_API_BASE_ENV_VARS)) if value),
None,
)
if not isinstance(resolved, str):
return True
host: Final = urlparse(resolved).hostname
return host is not None and (host == OPENAI_API_HOST or host.endswith(f".{OPENAI_API_HOST}"))
def _carries_cache_breakpoint(block: object) -> bool:
return isinstance(block, dict) and any(block.get(key) is not None for key in CACHE_BREAKPOINT_KEYS)
def _accepts_prompt_cache_breakpoint(block: object) -> bool:
return isinstance(block, dict) and block.get("type") in OPENAI_PROMPT_CACHE_BREAKPOINT_BLOCK_TYPES
class AnthropicCacheControlHook(CustomPromptManagement):
def get_chat_completion_prompt(
@ -81,13 +150,32 @@ class AnthropicCacheControlHook(CustomPromptManagement):
# provider transform, where each tool_config point appends at most one
# cachePoint to the tools. That block also counts toward Anthropic's
# limit, so reserve a slot for it here to leave room.
reserved_blocks: Final = 1 if any(p.get("location") == "tool_config" for p in remaining_points) else 0
stamped_dialect: Final = injection_points[0].get("_litellm_openai_dialect")
openai_dialect: Final = (
stamped_dialect
if isinstance(stamped_dialect, bool)
else AnthropicCacheControlHook._targets_openai_prompt_cache_breakpoint(
model,
non_default_params.get("custom_llm_provider"),
non_default_params.get("api_base") or non_default_params.get("base_url"),
non_default_params.get("prompt_cache_options"),
)
)
reserved_blocks: Final = (
1 if not openai_dialect and any(p.get("location") == "tool_config" for p in remaining_points) else 0
)
breakpoints_before: Final = AnthropicCacheControlHook._count_request_cache_breakpoints(processed_messages)
processed_messages = self._apply_message_injections(
points=message_points,
messages=processed_messages,
max_blocks=MAX_CACHE_CONTROL_BLOCKS - reserved_blocks,
openai_dialect=openai_dialect,
)
if (
openai_dialect
and AnthropicCacheControlHook._count_request_cache_breakpoints(processed_messages) > breakpoints_before
):
non_default_params.setdefault("prompt_cache_options", PromptCacheOptions(mode="explicit"))
# Pass through non-message injection points for provider-specific handling
if remaining_points:
@ -97,11 +185,43 @@ class AnthropicCacheControlHook(CustomPromptManagement):
return model, processed_messages, non_default_params
@staticmethod
def _targets_openai_prompt_cache_breakpoint(
model: str | None,
custom_llm_provider: str | None,
api_base: object = None,
prompt_cache_options: object = None,
) -> bool:
if model is None or not supports_openai_prompt_cache_breakpoint(model):
return False
if (custom_llm_provider or AnthropicCacheControlHook._resolve_provider(model)) != "openai":
return False
return prompt_cache_options is not None or targets_openai_api(api_base)
@staticmethod
def _resolve_provider(model: str) -> str | None:
from litellm.exceptions import BadRequestError
from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider
try:
_, provider, _, _ = get_llm_provider(model=model)
except BadRequestError:
return None
return provider
@staticmethod
def _count_request_cache_breakpoints(messages: Iterable[object], system: object = None) -> int:
system_blocks: Final = (
sum(1 for block in system if _carries_cache_breakpoint(block)) if isinstance(system, list) else 0
)
return system_blocks + sum(AnthropicCacheControlHook._count_cache_control_blocks(msg) for msg in messages)
@staticmethod
def _apply_message_injections(
points: list[CacheControlMessageInjectionPoint],
messages: list[AllMessageValues],
max_blocks: int,
openai_dialect: bool = False,
) -> list[AllMessageValues]:
"""Apply message-level cache control injection points in order.
@ -112,7 +232,7 @@ class AnthropicCacheControlHook(CustomPromptManagement):
``max_blocks`` is reached. Injection points are honored in config order,
so earlier points win when slots are scarce.
"""
used_blocks = sum(AnthropicCacheControlHook._count_cache_control_blocks(msg) for msg in messages)
used_blocks = AnthropicCacheControlHook._count_request_cache_breakpoints(messages)
limit_reached = False
for point in points:
@ -134,16 +254,17 @@ class AnthropicCacheControlHook(CustomPromptManagement):
continue
messages[target_index] = AnthropicCacheControlHook._safe_insert_cache_control_in_message(
messages[target_index], control
messages[target_index], control, openai_dialect
)
used_blocks += 1
if AnthropicCacheControlHook._message_has_cache_control(messages[target_index]):
used_blocks += 1
if limit_reached:
break
if limit_reached:
verbose_logger.warning(
"AnthropicCacheControlHook: Reached the Anthropic limit of %s cache_control blocks. Skipping further injection.",
"AnthropicCacheControlHook: Reached the provider limit of %s cache breakpoints. Skipping further injection.",
MAX_CACHE_CONTROL_BLOCKS,
)
@ -189,16 +310,13 @@ class AnthropicCacheControlHook(CustomPromptManagement):
return []
@staticmethod
def _count_cache_control_blocks(message: AllMessageValues) -> int:
"""Count cache_control breakpoints on a message (message + content level)."""
count = 0
if message.get("cache_control") is not None:
count += 1
def _count_cache_control_blocks(message: object) -> int:
if not isinstance(message, dict):
return 0
count = 1 if _carries_cache_breakpoint(message) else 0
content: Final = message.get("content")
if isinstance(content, list):
for block in content:
if isinstance(block, dict) and block.get("cache_control") is not None:
count += 1
count += sum(1 for block in content if _carries_cache_breakpoint(block))
return count
@staticmethod
@ -208,7 +326,7 @@ class AnthropicCacheControlHook(CustomPromptManagement):
@staticmethod
def _safe_insert_cache_control_in_message(
message: AllMessageValues, control: ChatCompletionCachedContent
message: AllMessageValues, control: ChatCompletionCachedContent, openai_dialect: bool = False
) -> AllMessageValues:
"""
Safe way to insert cache control in a message
@ -221,6 +339,9 @@ class AnthropicCacheControlHook(CustomPromptManagement):
Per Anthropic's API specification, when using multiple content blocks,
only the last content block can have cache_control.
"""
if openai_dialect:
return AnthropicCacheControlHook._insert_prompt_cache_breakpoint_in_message(message)
message_content: Final = message.get("content", None)
# 1. if string, insert cache control in the message
@ -232,11 +353,51 @@ class AnthropicCacheControlHook(CustomPromptManagement):
message_content[-1]["cache_control"] = control
return message
@staticmethod
def _insert_prompt_cache_breakpoint_in_message(message: AllMessageValues) -> AllMessageValues:
if message.get("role") == "assistant":
return message
message_content: Final = message.get("content", None)
if isinstance(message_content, str):
marked: Final = copy.copy(message)
marked["content"] = [
with_prompt_cache_breakpoint(
ChatCompletionTextObject(type="text", text=message_content), PromptCacheBreakpoint(mode="explicit")
)
]
return marked
if isinstance(message_content, list):
target_index: Final = next(
(
index
for index in range(len(message_content) - 1, -1, -1)
if _accepts_prompt_cache_breakpoint(message_content[index])
),
None,
)
if target_index is not None:
message_content[target_index] = with_prompt_cache_breakpoint(
message_content[target_index], PromptCacheBreakpoint(mode="explicit")
)
return message
@staticmethod
def _system_block_with_breakpoint(
block: Mapping[str, object], control: ChatCompletionCachedContent, openai_dialect: bool
) -> Mapping[str, object]:
marker: Final = (
("prompt_cache_breakpoint", PromptCacheBreakpoint(mode="explicit"))
if openai_dialect
else ("cache_control", control)
)
return {**block, marker[0]: marker[1]}
@staticmethod
def apply_to_anthropic_messages_request(
messages: list[dict],
system: str | list | None,
injection_points: list[CacheControlInjectionPoint],
openai_dialect: bool = False,
) -> tuple[list[dict], str | list | None, list[CacheControlInjectionPoint]]:
"""Apply cache control injection for the Anthropic-native v1/messages endpoint.
@ -262,30 +423,32 @@ class AnthropicCacheControlHook(CustomPromptManagement):
else:
remaining_points.append(point)
reserved_blocks: Final = 1 if any(p.get("location") == "tool_config" for p in remaining_points) else 0
reserved_blocks: Final = (
1 if not openai_dialect and any(p.get("location") == "tool_config" for p in remaining_points) else 0
)
max_blocks: Final = MAX_CACHE_CONTROL_BLOCKS - reserved_blocks
used_blocks = sum(
AnthropicCacheControlHook._count_cache_control_blocks(cast(AllMessageValues, msg))
for msg in processed_messages
)
if isinstance(processed_system, list):
used_blocks += sum(
1 for b in processed_system if isinstance(b, dict) and b.get("cache_control") is not None
)
message_blocks: Final = AnthropicCacheControlHook._count_request_cache_breakpoints(processed_messages)
system_blocks = AnthropicCacheControlHook._count_request_cache_breakpoints((), processed_system)
if system_points and processed_system is not None and used_blocks < max_blocks:
if system_points and processed_system is not None and message_blocks + system_blocks < max_blocks:
system_already_has_cc: Final = isinstance(processed_system, list) and any(
isinstance(b, dict) and b.get("cache_control") is not None for b in processed_system
_carries_cache_breakpoint(b) for b in processed_system
)
if not system_already_has_cc:
control: Final = system_points[0].get("control") or ChatCompletionCachedContent(type="ephemeral")
if isinstance(processed_system, str):
processed_system = [{"type": "text", "text": processed_system, "cache_control": control}]
used_blocks += 1
processed_system = [
AnthropicCacheControlHook._system_block_with_breakpoint(
AnthropicSystemMessageContent(type="text", text=processed_system), control, openai_dialect
)
]
system_blocks += 1
elif len(processed_system) > 0 and isinstance(processed_system[-1], dict):
processed_system[-1] = {**processed_system[-1], "cache_control": control}
used_blocks += 1
processed_system[-1] = AnthropicCacheControlHook._system_block_with_breakpoint(
processed_system[-1], control, openai_dialect
)
system_blocks += 1
for i, msg in enumerate(processed_messages):
content = msg.get("content")
@ -295,7 +458,8 @@ class AnthropicCacheControlHook(CustomPromptManagement):
processed_messages = AnthropicCacheControlHook._apply_message_injections(
points=message_points,
messages=cast(list[AllMessageValues], processed_messages),
max_blocks=max_blocks - used_blocks,
max_blocks=max_blocks - system_blocks,
openai_dialect=openai_dialect,
)
return processed_messages, processed_system, remaining_points
@ -315,17 +479,57 @@ class AnthropicCacheControlHook(CustomPromptManagement):
return ChatCompletionCachedContent(type="ephemeral")
@staticmethod
def _stamped_as_judged(points: list[CacheControlInjectionPoint]) -> list[dict[str, object]]:
def _stamped_as_judged(points: Sequence[CacheControlInjectionPoint]) -> Sequence[Mapping[str, object]]:
"""Mark written-back points as having passed the client cache_control judgment.
Builds copies because config-owned point dicts are shared across
requests; mutating them would leak the stamp into future requests.
"""
return [{**point, "_litellm_judged": True} for point in points]
return AnthropicCacheControlHook._stamped(points, "_litellm_judged", True)
@staticmethod
def _judged_configured_points(
points: Sequence[CacheControlInjectionPoint],
messages: list[AllMessageValues],
tools: list[object] | None,
model: str,
custom_llm_provider: str | None,
api_base: object,
prompt_cache_options: object,
) -> Sequence[Mapping[str, object]] | None:
if AnthropicCacheControlHook._should_stand_down(points, messages, None, tools):
return None
return AnthropicCacheControlHook._stamped_with_dialect(
points, model, custom_llm_provider, api_base, prompt_cache_options
)
@staticmethod
def _stamped_with_dialect(
points: Sequence[CacheControlInjectionPoint],
model: str,
custom_llm_provider: str | None,
api_base: object,
prompt_cache_options: object,
) -> Sequence[Mapping[str, object]]:
if not supports_openai_prompt_cache_breakpoint(model):
return points
return AnthropicCacheControlHook._stamped(
points,
"_litellm_openai_dialect",
AnthropicCacheControlHook._targets_openai_prompt_cache_breakpoint(
model, custom_llm_provider, api_base, prompt_cache_options
),
)
@staticmethod
def _stamped(
points: Sequence[CacheControlInjectionPoint], key: str, value: object
) -> Sequence[Mapping[str, object]]:
return [{**point, key: value} for point in points]
@staticmethod
def _should_stand_down(
points: list[CacheControlInjectionPoint],
points: Sequence[CacheControlInjectionPoint],
messages: list[AllMessageValues],
system: str | list | None,
tools: list | None,
@ -359,11 +563,8 @@ class AnthropicCacheControlHook(CustomPromptManagement):
carry the mark either at the top level (Anthropic shape) or nested under
``function`` (OpenAI shape); the Anthropic chat transform accepts both.
"""
if any(AnthropicCacheControlHook._count_cache_control_blocks(msg) for msg in messages):
if AnthropicCacheControlHook._count_request_cache_breakpoints(messages, system) > 0:
return True
if isinstance(system, list):
if any(isinstance(block, dict) and block.get("cache_control") is not None for block in system):
return True
if tools is not None:
return any(
isinstance(tool, dict)
@ -430,6 +631,50 @@ class AnthropicCacheControlHook(CustomPromptManagement):
]
return points
@staticmethod
def messages_with_default_injections(
messages: list[AllMessageValues],
models: Iterable[str],
tools: list[AllToolParamValues] | None = None,
enable_prompt_caching: bool | None = None,
) -> list[AllMessageValues]:
"""Return the messages auto prompt caching will send, default breakpoints included.
Router cache affinity depends on this. Deployment selection runs before the injection in
`litellm.acompletion`, so it has to reproduce the markers to derive the same cache key the
success event later writes from the sent messages. `models` is every candidate model of the
group: the first that would auto-inject decides, since the default breakpoints (system
prompt and trailing turn) do not depend on which deployment serves the call. Returns the
input list itself when auto-injection would not apply
"""
points: Final = next(
(
candidate
for candidate in (
AnthropicCacheControlHook.get_default_injection_points(
messages=messages,
system=None,
model=model,
custom_llm_provider=None,
tools=tools,
enable_prompt_caching=enable_prompt_caching,
)
for model in models
)
if candidate
),
None,
)
if not points:
return messages
return AnthropicCacheControlHook._apply_message_injections(
points=cast( # cast-ok: the default points are all message-location points
list[CacheControlMessageInjectionPoint], points
),
messages=copy.deepcopy(messages),
max_blocks=MAX_CACHE_CONTROL_BLOCKS,
)
@staticmethod
def maybe_seed_default_injection_points(
non_default_params: dict[str, Any],
@ -438,6 +683,7 @@ class AnthropicCacheControlHook(CustomPromptManagement):
custom_llm_provider: str | None,
tools: list | None = None,
enable_prompt_caching: bool | None = None,
api_base: object = None,
) -> None:
"""For /chat/completions: resolve the injection points the request should carry.
@ -452,10 +698,19 @@ class AnthropicCacheControlHook(CustomPromptManagement):
unchanged.
"""
if non_default_params.get("cache_control_injection_points"):
if AnthropicCacheControlHook._should_stand_down(
non_default_params["cache_control_injection_points"], messages, None, tools
):
judged: Final = AnthropicCacheControlHook._judged_configured_points(
non_default_params["cache_control_injection_points"],
messages,
tools,
model,
custom_llm_provider,
api_base,
non_default_params.get("prompt_cache_options"),
)
if judged is None:
non_default_params.pop("cache_control_injection_points")
else:
non_default_params["cache_control_injection_points"] = judged
return
points: Final = AnthropicCacheControlHook.get_default_injection_points(
messages=messages,
@ -476,6 +731,7 @@ class AnthropicCacheControlHook(CustomPromptManagement):
model: str | None = None,
custom_llm_provider: str | None = None,
tools: list[dict] | None = None,
api_base: str | None = None,
) -> tuple[list[dict], str | list | None]:
"""Extract cache_control_injection_points from kwargs and apply if present.
@ -513,11 +769,21 @@ class AnthropicCacheControlHook(CustomPromptManagement):
if not injection_points:
return messages, system
openai_dialect: Final = AnthropicCacheControlHook._targets_openai_prompt_cache_breakpoint(
model, custom_llm_provider, api_base, kwargs.get("prompt_cache_options")
)
breakpoints_before: Final = AnthropicCacheControlHook._count_request_cache_breakpoints(messages, system)
messages, system, remaining = AnthropicCacheControlHook.apply_to_anthropic_messages_request(
messages=messages,
system=system,
injection_points=injection_points,
openai_dialect=openai_dialect,
)
if (
openai_dialect
and AnthropicCacheControlHook._count_request_cache_breakpoints(messages, system) > breakpoints_before
):
kwargs.setdefault("prompt_cache_options", PromptCacheOptions(mode="explicit"))
if remaining:
kwargs["cache_control_injection_points"] = AnthropicCacheControlHook._stamped_as_judged(remaining)
return messages, system

View file

@ -359,10 +359,10 @@ class ArizePhoenixPromptManager(CustomPromptManagement):
"""
Determine if prompt management should run based on the prompt_id.
For Arize Phoenix, we always return True and handle the prompt loading
in the _compile_prompt_helper method.
Arize Phoenix needs a prompt_id to compile, so it declines requests without one;
prompt loading itself happens in the _compile_prompt_helper method.
"""
return True
return prompt_id is not None
def _compile_prompt_helper(
self,

View file

@ -65,6 +65,41 @@ _guardrail_self_recorded: Final[contextvars.ContextVar[bool]] = contextvars.Cont
)
def is_guardrail_intervention(e: Exception) -> bool:
"""
Returns True if the exception represents an intentional guardrail block
(this was logged previously as an API failure - guardrail_failed_to_respond).
Guardrails signal intentional blocks by raising:
- GuardrailRaisedException (generic guardrail API, tool permission)
- BlockedPiiEntityError (Presidio PII detection)
- SensitiveDataRouteException (sensitive-data reroute to on-premise model)
- HTTPException with a block-signalling status (400, 403, 422)
- ModifyResponseException (passthrough mode violation)
Only the statuses guardrails use in-tree to signal a deliberate rejection
count as an intervention: 400 (content policy), 403 (e.g. akto) and 422
(e.g. llm_as_a_judge). Other 4xx codes are commonly propagated from an
upstream guardrail provider response (401 bad key, 408 timeout, 429 rate
limit, or a raw upstream status), which are technical failures, not
blocks, so they stay guardrail_failed_to_respond.
"""
if isinstance(e, ModifyResponseException):
return True
if isinstance(
e,
(
GuardrailRaisedException,
BlockedPiiEntityError,
SensitiveDataRouteException,
),
):
return True
if HTTPException is not None and isinstance(e, HTTPException) and e.status_code in _GUARDRAIL_BLOCK_STATUS_CODES:
return True
return False
def _strict_guardrail_modes_enabled() -> bool:
"""Whether guardrail-mode validation raises (default) or logs a warning.
@ -429,11 +464,13 @@ class CustomGuardrail(CustomLogger):
f"Sensitive data detected by {self.guardrail_name} (routing skipped: request has no session_id)"
),
guardrail_name=self.guardrail_name,
blocked_content=True,
)
else:
raise GuardrailRaisedException(
message=f"Sensitive data detected by {self.guardrail_name}",
guardrail_name=self.guardrail_name,
blocked_content=True,
)
@staticmethod
@ -1068,42 +1105,8 @@ class CustomGuardrail(CustomLogger):
@staticmethod
def _is_guardrail_intervention(e: Exception) -> bool:
"""
Returns True if the exception represents an intentional guardrail block
(this was logged previously as an API failure - guardrail_failed_to_respond).
Guardrails signal intentional blocks by raising:
- GuardrailRaisedException (generic guardrail API, tool permission)
- BlockedPiiEntityError (Presidio PII detection)
- SensitiveDataRouteException (sensitive-data reroute to on-premise model)
- HTTPException with a block-signalling status (400, 403, 422)
- ModifyResponseException (passthrough mode violation)
Only the statuses guardrails use in-tree to signal a deliberate rejection
count as an intervention: 400 (content policy), 403 (e.g. akto) and 422
(e.g. llm_as_a_judge). Other 4xx codes are commonly propagated from an
upstream guardrail provider response (401 bad key, 408 timeout, 429 rate
limit, or a raw upstream status), which are technical failures, not
blocks, so they stay guardrail_failed_to_respond.
"""
if isinstance(e, ModifyResponseException):
return True
if isinstance(
e,
(
GuardrailRaisedException,
BlockedPiiEntityError,
SensitiveDataRouteException,
),
):
return True
if (
HTTPException is not None
and isinstance(e, HTTPException)
and e.status_code in _GUARDRAIL_BLOCK_STATUS_CODES
):
return True
return False
"""Retained spelling for existing callers; prefer ``is_guardrail_intervention``."""
return is_guardrail_intervention(e)
def _process_error(
self,

View file

@ -11,6 +11,7 @@ from litellm.integrations.datadog.datadog_handler import (
get_datadog_hostname,
get_datadog_pod_name,
get_datadog_service,
normalize_datadog_tag_value,
)
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
from litellm.llms.custom_httpx.http_handler import (
@ -184,7 +185,7 @@ class DatadogCostManagementLogger(CustomBatchLogger):
# Backwards-compat: team/user/model_group preserved regardless of allowlist.
if metadata.get("user_api_key_alias"):
tags["user"] = str(metadata["user_api_key_alias"])
tags["user"] = normalize_datadog_tag_value(metadata["user_api_key_alias"])
team_tag: Final = (
metadata.get("user_api_key_team_alias")
or metadata.get("team_alias")
@ -192,7 +193,7 @@ class DatadogCostManagementLogger(CustomBatchLogger):
or metadata.get("team_id")
)
if team_tag:
tags["team"] = str(team_tag)
tags["team"] = normalize_datadog_tag_value(team_tag)
if metadata.get("model_group"):
tags["model_group"] = str(metadata["model_group"])
@ -229,7 +230,7 @@ class DatadogCostManagementLogger(CustomBatchLogger):
value,
)
return
tags[key] = value
tags[key] = normalize_datadog_tag_value(value)
@staticmethod
def _add_tag(tags: dict[str, str], key: str, value: Any) -> None:

View file

@ -3,6 +3,7 @@
from __future__ import annotations
import os
import re
from typing import Final
from litellm.types.utils import StandardLoggingPayload
@ -36,6 +37,13 @@ def get_datadog_pod_name() -> str:
return os.getenv("POD_NAME", "unknown")
def normalize_datadog_tag_value(value: object) -> str:
normalized_value: Final = "".join(
character if character.isalnum() or character in "_-:./" else "_" for character in str(value).lower()
)
return re.sub(r"_+", "_", normalized_value).strip("_")
def get_datadog_tags(
standard_logging_object: StandardLoggingPayload | None = None,
) -> list[str]:
@ -58,7 +66,7 @@ def get_datadog_tags(
if standard_logging_object:
request_tags: Final = standard_logging_object.get("request_tags", []) or []
tags.extend(f"request_tag:{tag}" for tag in request_tags)
tags.extend(f"request_tag:{normalize_datadog_tag_value(tag)}" for tag in request_tags)
# Add Team Tag
metadata: Final = standard_logging_object.get("metadata", {}) or {}
@ -69,6 +77,6 @@ def get_datadog_tags(
or metadata.get("team_id")
)
if team_tag:
tags.append(f"team:{team_tag}")
tags.append(f"team:{normalize_datadog_tag_value(team_tag)}")
return tags

View file

@ -12,6 +12,7 @@ from litellm.integrations.datadog.datadog_handler import (
get_datadog_hostname,
get_datadog_pod_name,
get_datadog_service,
normalize_datadog_tag_value,
)
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
from litellm.llms.custom_httpx.http_handler import (
@ -97,7 +98,7 @@ class DatadogMetricsLogger(CustomBatchLogger):
)
if team_tag:
tags.append(f"team:{team_tag}")
tags.append(f"team:{normalize_datadog_tag_value(team_tag)}")
return tags

View file

@ -9,7 +9,15 @@ from typing import TYPE_CHECKING, Any, Final, cast
from opentelemetry.context import Context, attach, get_current
from opentelemetry.sdk._logs import LoggerProvider
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.trace import Span, Tracer, get_current_span, use_span
from opentelemetry.trace import (
INVALID_SPAN,
Link,
Span,
Tracer,
get_current_span,
set_span_in_context,
use_span,
)
import litellm
from litellm._logging import verbose_logger
@ -21,6 +29,7 @@ from litellm.integrations.otel.model.config import OpenTelemetryV2Config
from litellm.integrations.otel.model.metadata import (
LLMCallEvent,
RequestIdentity,
auth_metadata,
model_from_request_data,
)
from litellm.integrations.otel.model.payloads import (
@ -118,6 +127,12 @@ _OTEL_MODULES: Final = (
_OPEN_CALLS_MAX: Final = 10_000
def _request_trace_links(context: Context | None) -> tuple[Link, ...] | None:
"""A link back to the request trace, for a span detached into its own trace."""
anchor: Final = get_current_span(context).get_span_context()
return (Link(anchor),) if anchor.is_valid else None
class _LLMCallSpan:
"""The state carried from the ``pre_call`` boundary to span close.
@ -127,13 +142,24 @@ class _LLMCallSpan:
own (worker-copied) ambient context using ``start_time_ns``. The presence of
a carrier for a call at all is the proof that ``pre_call`` ran, i.e. that an
upstream call was actually attempted.
``provider`` is the routed provider the live span was opened on (``None`` on
the default route or when creation was deferred). It is held in the tenant
cache while the span is open so LRU eviction can't shut the provider down
under it, and must be released exactly once when the carrier is removed.
"""
__slots__ = ("span", "start_time_ns")
__slots__ = ("provider", "span", "start_time_ns")
def __init__(self, span: "Span | None", start_time_ns: int | None) -> None:
def __init__(
self,
span: "Span | None",
start_time_ns: int | None,
provider: "TracerProvider | None" = None,
) -> None:
self.span = span
self.start_time_ns = start_time_ns
self.provider = provider
class OpenTelemetryV2(CustomLogger):
@ -258,26 +284,37 @@ class OpenTelemetryV2(CustomLogger):
if call_id in self._open_llm_calls:
return
start_time_ns: Final = to_ns(datetime.now())
span: Span | None = None
# Parent to the request's anchored root span (stable across the request),
# falling back to ambient on the SDK path. Open the span live only when
# that resolves to a recordable parent; otherwise defer to the close
# callback (the thread-pool case, where the anchor isn't visible here).
# Do not route on the deferred path: creating or LRU-touching a tenant
# provider here would evict idle ones even though close re-routes.
parent_context: Final = resolve_request_span_context()
if is_recordable_span(get_current_span(parent_context)):
span = self._emitter.start_span(
if not is_recordable_span(get_current_span(parent_context)):
self._store_open_call(call_id, _LLMCallSpan(span=None, start_time_ns=start_time_ns))
return
# A detached route roots its own trace instead (linked to the request
# trace) — see ``TenantRoute.detached``.
route: Final = self._tenant_tracers.route_for(self.tracer, call.dynamic_params, call.auth_metadata)
try:
span: Final = self._emitter.start_span(
SpanRole.LLM_CALL,
call.provisional_span_name,
parent_context=parent_context,
parent_context=(
set_span_in_context(INVALID_SPAN, parent_context) if route.detached else parent_context
),
start_time_ns=start_time_ns,
tracer=self._tenant_tracers.tracer_for(self.tracer, call.dynamic_params),
tracer=route.tracer,
links=_request_trace_links(parent_context) if route.detached else None,
)
self._open_llm_calls[call_id] = _LLMCallSpan(span=span, start_time_ns=start_time_ns)
# Evict the oldest open call if the map is over budget. A call that opens
# but never closes (a stream that only fires stream events) would linger
# otherwise; the evicted span is simply dropped (never exported).
if len(self._open_llm_calls) > _OPEN_CALLS_MAX:
self._open_llm_calls.popitem(last=False)
except BaseException:
self._tenant_tracers.release(route.provider)
raise
self._store_open_call(
call_id,
_LLMCallSpan(span=span, start_time_ns=start_time_ns, provider=route.provider),
)
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
if self._emit_mcp_tool_call(kwargs, start_time, end_time):
@ -371,17 +408,22 @@ class OpenTelemetryV2(CustomLogger):
# otherwise linger until evicted; drop it so it's neither leaked nor closed
# as a phantom LLM span.
if data.identity.call_id:
self._open_llm_calls.pop(data.identity.call_id, None)
parent_context, links = resolve_mcp_span_context()
parent_context = self._seed_identity_baggage(data.identity, None, parent_context)
self._emitter.emit(
SpanRole.MCP_TOOL_CALL,
data,
parent_context=parent_context,
start_time_ns=to_ns(start_time),
end_time_ns=to_ns(end_time),
links=links,
)
self._release_carrier(self._open_llm_calls.pop(data.identity.call_id, None))
route: Final = self._tenant_tracers.route_for(self.tracer, None, auth_metadata(payload, kwargs))
try:
parent_context, links = resolve_mcp_span_context()
seeded: Final = self._seed_identity_baggage(data.identity, None, parent_context)
self._emitter.emit(
SpanRole.MCP_TOOL_CALL,
data,
parent_context=(set_span_in_context(INVALID_SPAN, seeded) if route.detached else seeded),
start_time_ns=to_ns(start_time),
end_time_ns=to_ns(end_time),
links=((*(links or ()), *(_request_trace_links(seeded) or ())) if route.detached else links),
tracer=route.tracer,
)
finally:
self._tenant_tracers.release(route.provider)
return True
def _emit_mcp_list_tools(
@ -407,17 +449,22 @@ class OpenTelemetryV2(CustomLogger):
payload, capture_content=self.config.capture_span_content
)
if data.identity.call_id:
self._open_llm_calls.pop(data.identity.call_id, None)
parent_context, links = resolve_mcp_span_context()
parent_context = self._seed_identity_baggage(data.identity, None, parent_context)
self._emitter.emit(
SpanRole.MCP_LIST_TOOLS,
data,
parent_context=parent_context,
start_time_ns=to_ns(start_time),
end_time_ns=to_ns(end_time),
links=links,
)
self._release_carrier(self._open_llm_calls.pop(data.identity.call_id, None))
route: Final = self._tenant_tracers.route_for(self.tracer, None, auth_metadata(payload, kwargs))
try:
parent_context, links = resolve_mcp_span_context()
seeded: Final = self._seed_identity_baggage(data.identity, None, parent_context)
self._emitter.emit(
SpanRole.MCP_LIST_TOOLS,
data,
parent_context=(set_span_in_context(INVALID_SPAN, seeded) if route.detached else seeded),
start_time_ns=to_ns(start_time),
end_time_ns=to_ns(end_time),
links=((*(links or ()), *(_request_trace_links(seeded) or ())) if route.detached else links),
tracer=route.tracer,
)
finally:
self._tenant_tracers.release(route.provider)
return True
def _close_llm_call(
@ -439,6 +486,36 @@ class OpenTelemetryV2(CustomLogger):
carrier: Final = self._open_llm_calls.pop(call_id, None) if call_id else None
if carrier is None:
return None
try:
return self._finish_carrier(carrier, call, end_time)
finally:
# After the span has ended, so a release-triggered provider shutdown
# force-flushes it out rather than racing its enqueue.
self._release_carrier(carrier)
def _store_open_call(self, call_id: str, carrier: _LLMCallSpan) -> None:
"""Remember an in-flight LLM call, evicting the oldest if over budget.
A call that opens but never closes (a stream that only fires stream
events) would linger otherwise; the evicted span is simply dropped
(never exported).
"""
self._open_llm_calls[call_id] = carrier
if len(self._open_llm_calls) > _OPEN_CALLS_MAX:
_, evicted = self._open_llm_calls.popitem(last=False)
self._release_carrier(evicted)
def _release_carrier(self, carrier: "_LLMCallSpan | None") -> None:
"""Release the routed provider a removed carrier was holding open."""
if carrier is not None:
self._tenant_tracers.release(carrier.provider)
def _finish_carrier(
self,
carrier: _LLMCallSpan,
call: LLMCallEvent,
end_time: datetime | float | None,
) -> Span | None:
payload: Final = call.payload
if payload is None:
if carrier.span is not None:
@ -462,16 +539,23 @@ class OpenTelemetryV2(CustomLogger):
# The worker copied the request task's context, which carries the anchored
# root span — parent to it (ambient fallback on the SDK path). Seed identity
# Baggage so the span — and the SDK path, which has none — is labeled
# consistently.
parent_ctx = self._seed_identity_baggage(data.identity, data.request_model, resolve_request_span_context())
return self._emitter.emit(
SpanRole.LLM_CALL,
data,
parent_context=parent_ctx,
start_time_ns=carrier.start_time_ns,
end_time_ns=end_time_ns,
tracer=self._tenant_tracers.tracer_for(self.tracer, call.dynamic_params),
)
# consistently. A detached route roots its own trace instead, linked back.
route: Final = self._tenant_tracers.route_for(self.tracer, call.dynamic_params, call.auth_metadata)
try:
parent_ctx: Final = self._seed_identity_baggage(
data.identity, data.request_model, resolve_request_span_context()
)
return self._emitter.emit(
SpanRole.LLM_CALL,
data,
parent_context=(set_span_in_context(INVALID_SPAN, parent_ctx) if route.detached else parent_ctx),
start_time_ns=carrier.start_time_ns,
end_time_ns=end_time_ns,
tracer=route.tracer,
links=_request_trace_links(parent_ctx) if route.detached else None,
)
finally:
self._tenant_tracers.release(route.provider)
# ====================================================================== #
# Service hooks

View file

@ -36,8 +36,9 @@ model. They coincide on the SDK path, which is correct.
from __future__ import annotations
from collections.abc import Mapping
from collections.abc import Iterator, Mapping
from dataclasses import dataclass, field
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, cast
from litellm.constants import LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL
@ -195,6 +196,10 @@ class LLMCallEvent:
# The ``standard_callback_dynamic_params`` routing the call to a per-tenant
# tracer (its own exporter/endpoint), or ``None`` when the call isn't scoped.
dynamic_params: Any
# The key/team config the proxy resolved at auth (``user_api_key_auth_metadata``),
# routing the call to that tenant's telemetry project. Server-set and so
# trusted, unlike ``dynamic_params``, which carries client-supplied metadata.
auth_metadata: Mapping[str, str] | None
# True for synthetic proxy-gate logs (auth / rate-limit rejections): they fire
# the ``pre_call`` hook but never made an upstream call, so they get no span.
is_no_upstream_call: bool
@ -214,6 +219,7 @@ class LLMCallEvent:
call_id=_call_id(payload, kwargs),
payload=payload,
dynamic_params=kwargs.get("standard_callback_dynamic_params"),
auth_metadata=auth_metadata(payload, kwargs),
is_no_upstream_call=bool(kwargs.get(LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL)),
provisional_span_name=f"{operation.value} {model}".strip(),
time_to_first_chunk_seconds=time_to_first_chunk_seconds(kwargs),
@ -235,6 +241,64 @@ def time_to_first_chunk_seconds(kwargs: Mapping[str, Any]) -> float | None:
return completion_start - api_call_start
def auth_metadata(payload: StandardLoggingPayload | None, kwargs: Mapping[str, object]) -> Mapping[str, str] | None:
"""The key/team config the proxy resolved at auth, or ``None`` off the proxy.
Read from the payload once the call closes and from ``litellm_params`` at
``pre_call``, where no payload exists yet the LLM-call span is *created* at
``pre_call``, so the tracer (and therefore the destination) must be
resolvable there. Values arrive untyped, so non-string entries are dropped
rather than passed on to header builders.
"""
return next(
(
typed
for metadata in _metadata_dicts(payload, kwargs)
if (typed := _string_entries(metadata.get("user_api_key_auth_metadata")))
),
None,
)
def _as_str_mapping(value: object) -> Mapping[str, object] | None:
"""A read-only view of ``value`` when it is a mapping, else ``None``."""
if not isinstance(value, Mapping):
return None
return cast("Mapping[str, object]", value) # cast-ok: isinstance-guarded, JSON metadata has str keys
def _string_entries(value: object) -> Mapping[str, str] | None:
entries: Final = _as_str_mapping(value)
if entries is None:
return None
typed: Final = MappingProxyType({key: item for key, item in entries.items() if isinstance(item, str)})
return typed or None
def _metadata_dicts(
payload: StandardLoggingPayload | None, kwargs: Mapping[str, object]
) -> Iterator[Mapping[str, object]]:
"""Request metadata dicts, closed-call payload first then the live kwargs.
``litellm_metadata`` is the metadata field on the Anthropic-shaped routes;
litellm copies it onto ``metadata``, but both are yielded so a route that
populates only one is still covered.
"""
payload_view: Final = _as_str_mapping(payload)
if payload_view is not None:
payload_metadata: Final = _as_str_mapping(payload_view.get("metadata"))
if payload_metadata is not None:
yield payload_metadata
params: Final = _as_str_mapping(kwargs.get("litellm_params"))
if params is None:
return
yield from (
metadata
for key in ("metadata", "litellm_metadata")
if (metadata := _as_str_mapping(params.get(key))) is not None
)
def _call_id(payload: StandardLoggingPayload | None, kwargs: Mapping[str, Any]) -> str | None:
"""The call id from the payload (when closed) or the bare kwargs (at pre_call)."""
if payload is not None:

View file

@ -1,31 +1,45 @@
"""Per-request multi-tenant tracer routing.
When a request carries team/key vendor credentials in
``standard_callback_dynamic_params``, its spans must export through a
``TracerProvider`` whose OTLP headers carry those credentials.
``TenantTracerCache`` builds and caches one provider per distinct credential
set, and otherwise hands back the logger's default tracer. This lets a single
logger fan requests out to many tenants without needing a logger per tenant.
``standard_callback_dynamic_params``, or the key/team config resolved at auth
names a destination project, its spans must export through a
``TracerProvider`` whose OTLP headers carry those credentials / that project.
``TenantTracerCache`` builds and caches one provider per distinct
(credentials, project) pair, and otherwise hands back the logger's default
tracer. This lets a single logger fan requests out to many tenants without
needing a logger per tenant.
"""
import threading
from collections import OrderedDict
from collections.abc import Mapping
from typing import Any, Final
from dataclasses import dataclass
from types import MappingProxyType
from typing import Any, Final, TypeAlias
from urllib.parse import quote
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.trace import Tracer
from litellm._logging import verbose_logger
from litellm.integrations.otel.model.config import OpenTelemetryV2Config
from litellm.integrations.otel.model.config import ExporterSpec, OpenTelemetryV2Config
from litellm.integrations.otel.plumbing.providers import (
build_tracer_provider,
get_tracer,
)
from litellm.integrations.otel.presets import dynamic_otlp_headers
from litellm.integrations.otel.presets import (
dynamic_otlp_headers,
project_routing_headers,
)
# Exporter kinds that ignore headers — never rewritten with dynamic credentials.
_NON_OTLP_KINDS: Final = ("console", "in_memory", "inmemory", "memory")
# gRPC exporters still take dynamic credentials (as gRPC metadata) but not
# project headers: the routing headers backends read (Phoenix's
# ``x-project-name``) are only honored on the OTLP/HTTP endpoint.
_GRPC_KINDS: Final = ("otlp_grpc", "grpc")
# Cap on distinct credential-scoped providers held at once. ``dynamic_params``
# can be populated from request metadata, so an unbounded cache lets a caller
# spawn one ``TracerProvider`` (plus its ``BatchSpanProcessor`` background
@ -34,6 +48,23 @@ _NON_OTLP_KINDS: Final = ("console", "in_memory", "inmemory", "memory")
# evicted providers so their threads are reclaimed.
_MAX_CACHED_PROVIDERS: Final = 256
# Cap on providers evicted from the cache while still holding open spans, which
# are kept alive to drain instead of being shut down under them. Their only
# other bound is the logger's open-call map (10k), so without this a caller
# cycling unique credential sets across long-lived calls could pin far more
# live providers, and exporter threads, than the cache cap allows. Past this
# many, the stalest retiree is shut down and whatever it was draining is
# dropped (a shut-down ``BatchSpanProcessor`` discards spans handed to it after
# the fact), which by then means a span on a route evicted long ago. A quarter
# of the cache cap: enough that a burst of tenant churn during long-lived calls
# still drains normally, small enough that the worst case is a bounded 320
# providers rather than one per concurrent call.
_MAX_RETIRED_PROVIDERS: Final = 64
_HeaderItems: TypeAlias = tuple[tuple[str, str], ...]
_NO_HEADERS: Final[Mapping[str, str]] = MappingProxyType({})
def _shutdown_provider(provider: TracerProvider) -> None:
"""Flush + stop an evicted provider's processors (reclaims their threads).
@ -49,8 +80,41 @@ def _shutdown_provider(provider: TracerProvider) -> None:
verbose_logger.debug("OTel V2: error shutting down evicted provider: %s", e)
def _plain_header_string(headers: Mapping[str, str]) -> str:
return ",".join(f"{key}={value}" for key, value in headers.items())
def _encoded_header_string(headers: Mapping[str, str]) -> str:
"""Percent-encode values so one containing the ``k=v,k=v`` separators (e.g.
a project name with a comma) survives; ``parse_env_headers`` decodes it back.
"""
return ",".join(f"{key}={quote(value, safe='')}" for key, value in headers.items())
@dataclass(frozen=True, slots=True)
class TenantRoute:
"""The tracer to create a span on, plus whether it must root its own trace.
``detached`` is True when project routing engaged. Phoenix assigns a whole
trace to one project by whichever of its spans arrives first, so a
project-routed span parented into the request trace gets dragged into the
project of the default-exported request spans and the header does nothing.
The span must therefore start a fresh trace (with a link back to the
request trace for correlation) which is also how the v1 Phoenix logger
behaved, exporting each request under its own Phoenix-local parent span.
"""
tracer: Tracer
detached: bool
#: The provider ``tracer`` came from, or ``None`` on the default route. It
#: is returned already held (counted as an open span, atomically with the
#: cache update), so LRU eviction can't shut it down before the caller's
#: span lands; the caller must ``release`` it exactly once when done.
provider: TracerProvider | None = None
class TenantTracerCache:
"""Credential-scoped ``TracerProvider`` cache keyed by the dynamic headers."""
"""Credential/project-scoped ``TracerProvider`` cache keyed by the routing headers."""
def __init__(
self,
@ -61,49 +125,170 @@ class TenantTracerCache:
self._config = config
self._callback_name = callback_name
self._tracer_name = tracer_name
self._providers: OrderedDict[tuple[tuple[str, str], ...], TracerProvider] = OrderedDict()
# Guards the three mutable structures below: ``pre_call`` can run on
# thread-pool workers concurrently with the event loop, so cache
# updates, span counts, and retirement must be atomic.
self._lock: Final = threading.Lock()
self._providers: OrderedDict[tuple[_HeaderItems, _HeaderItems], TracerProvider] = OrderedDict()
self._open_span_counts: dict[TracerProvider, int] = {} # mutable-ok: live refcount state
# Oldest-first so an overflow of draining providers sheds the stalest.
self._retired: OrderedDict[TracerProvider, None] = OrderedDict() # mutable-ok: draining evicted providers
self._project_routable = any(
spec.owner == callback_name and spec.kind.lower() not in (*_NON_OTLP_KINDS, *_GRPC_KINDS)
for spec in config.exporters
)
self._warned_project_unroutable = False
def tracer_for(self, default: Tracer, dynamic_params: Any) -> Tracer:
"""Return the tracer for this request.
def release(self, provider: TracerProvider | None) -> None:
"""Drop one open-span count; shut a retired provider down once drained.
Use ``default`` unless the request's dynamic credentials require a
credential-scoped tracer, in which case build (or reuse) one. The cache
is a bounded LRU: the least-recently-used provider is flushed and shut
down on overflow so its exporter threads don't accumulate.
``None`` (the default route) is a no-op so callers can release a
``TenantRoute.provider`` unconditionally. The shutdown itself runs
outside the lock: it force-flushes over the network and must not stall
every concurrently routing request.
"""
headers: Final = dynamic_otlp_headers(self._callback_name, dynamic_params)
if not headers:
return default
cache_key: Final = tuple(sorted(headers.items()))
provider = self._providers.get(cache_key)
if provider is not None:
if provider is None:
return
with self._lock:
remaining: Final = self._open_span_counts.get(provider, 0) - 1
if remaining > 0:
self._open_span_counts[provider] = remaining
return
self._open_span_counts.pop(provider, None)
drained: Final = provider in self._retired
self._retired.pop(provider, None)
if drained:
_shutdown_provider(provider)
def route_for(
self,
default: Tracer,
dynamic_params: Any,
auth_metadata: Mapping[str, str] | None = None,
) -> TenantRoute:
"""Return the tracer (and trace-detachment flag) for this request.
Use ``default`` unless the request's dynamic credentials or its key/team
project require a scoped tracer, in which case build (or reuse) one. The
cache is a bounded LRU: the least-recently-used provider is flushed and
shut down on overflow so its exporter threads don't accumulate.
A routed provider is returned already held its open-span count is
incremented in the same critical section as the cache update so a
concurrent overflow eviction can't shut it down between selection and
the caller's span start. The caller must ``release`` it exactly once.
"""
credential_headers: Final = dynamic_otlp_headers(self._callback_name, dynamic_params) or _NO_HEADERS
project_headers: Final = self._project_headers(auth_metadata)
if not credential_headers and not project_headers:
return TenantRoute(tracer=default, detached=False)
cache_key: Final = (
tuple(sorted(credential_headers.items())),
tuple(sorted(project_headers.items())),
)
with self._lock:
provider: Final = self._cached_provider_locked(cache_key, credential_headers, project_headers)
self._open_span_counts[provider] = self._open_span_counts.get(provider, 0) + 1
evicted: Final = self._evicted_on_overflow_locked()
if evicted is not None:
_shutdown_provider(evicted)
return TenantRoute(
tracer=get_tracer(provider, self._tracer_name),
detached=bool(project_headers),
provider=provider,
)
def _cached_provider_locked(
self,
cache_key: tuple[_HeaderItems, _HeaderItems],
credential_headers: Mapping[str, str],
project_headers: Mapping[str, str],
) -> TracerProvider:
cached: Final = self._providers.get(cache_key)
if cached is not None:
self._providers.move_to_end(cache_key)
else:
provider = build_tracer_provider(self._config_with_headers(headers))
self._providers[cache_key] = provider
if len(self._providers) > _MAX_CACHED_PROVIDERS:
_, evicted = self._providers.popitem(last=False)
_shutdown_provider(evicted)
return get_tracer(provider, self._tracer_name)
return cached
built: Final = build_tracer_provider(self._routed_config(credential_headers, project_headers))
self._providers[cache_key] = built
return built
def _config_with_headers(self, headers: Mapping[str, str]) -> OpenTelemetryV2Config:
"""Clone the config, stamping ``headers`` onto the credential's own exporter.
def _evicted_on_overflow_locked(self) -> TracerProvider | None:
"""Pop the LRU provider past the cap; return it if the caller must shut it down.
``headers`` are the per-request credentials of ``self._callback_name`` (the
integration that built this cache), so they apply only to the exporter that
integration contributed (``spec.owner``). A request that carries one
tenant's Arize key must never rewrite the headers of a co-configured
Langfuse or self-hosted collector exporter, which would leak that key to a
different backend.
A provider with open spans is retired to drain instead: stopping its
processors while a span opened at ``pre_call`` is still live would
silently drop that span at end instead of exporting it. Retirees are
themselves capped, so the stalest one is shut down (and its open-span
count dropped, making its eventual ``release`` a no-op) once too many
pile up rather than letting them accumulate a thread each.
"""
header_str: Final = ",".join(f"{key}={value}" for key, value in headers.items())
header_update: Final[dict[str, str]] = {"headers": header_str}
exporters: Final = [
(
spec.model_copy(update=header_update)
if spec.owner == self._callback_name and spec.kind.lower() not in _NON_OTLP_KINDS
else spec
if len(self._providers) <= _MAX_CACHED_PROVIDERS:
return None
_, evicted = self._providers.popitem(last=False)
if self._open_span_counts.get(evicted, 0) == 0:
return evicted
self._retired[evicted] = None
if len(self._retired) <= _MAX_RETIRED_PROVIDERS:
return None
overflowed, _ = self._retired.popitem(last=False)
self._open_span_counts.pop(overflowed, None)
return overflowed
def _project_headers(self, auth_metadata: Mapping[str, str] | None) -> Mapping[str, str]:
"""The per-request project-routing headers, if this cache can apply them.
A gRPC-only exporter can't (the project header route is HTTP-only), so
the request warns once and stays on the env-configured default project.
"""
requested: Final = project_routing_headers(self._callback_name, auth_metadata)
if not requested or self._project_routable:
return requested
if not self._warned_project_unroutable:
self._warned_project_unroutable = True
verbose_logger.warning(
"OTel V2: %s key/team config names a per-request project, but its exporter "
"is not OTLP/HTTP and the project header is HTTP-only; spans stay in the "
"default project.",
self._callback_name,
)
for spec in self._config.exporters
return _NO_HEADERS
def _routed_config(
self,
credential_headers: Mapping[str, str],
project_headers: Mapping[str, str],
) -> OpenTelemetryV2Config:
"""Clone the config, rewriting headers on the callback's own exporter.
Both header sets apply only to the exporter ``self._callback_name``
contributed (``spec.owner``). A request that carries one tenant's Arize
key must never rewrite the headers of a co-configured Langfuse or
self-hosted collector exporter, which would leak that key to a
different backend.
Dynamic credentials REPLACE the exporter's headers — they are the
tenant's complete credential set. Project headers APPEND instead: the
preset's static headers carry the backend auth (Phoenix's
``Authorization``), which must survive routing to a project.
"""
exporters: Final = [
self._routed_exporter(spec, credential_headers, project_headers) for spec in self._config.exporters
]
return self._config.model_copy(update={"exporters": exporters})
def _routed_exporter(
self,
spec: ExporterSpec,
credential_headers: Mapping[str, str],
project_headers: Mapping[str, str],
) -> ExporterSpec:
kind: Final = spec.kind.lower()
if spec.owner != self._callback_name or kind in _NON_OTLP_KINDS:
return spec
base: Final = _plain_header_string(credential_headers) if credential_headers else spec.headers
routed: Final = (
",".join(part for part in (base, _encoded_header_string(project_headers)) if part)
if project_headers and kind not in _GRPC_KINDS
else base
)
return spec if routed == spec.headers else spec.model_copy(update={"headers": routed})

View file

@ -8,7 +8,8 @@ the factory in ``litellm_logging`` can resolve a name and build a single
``OpenTelemetryV2`` instance from the result.
"""
from collections.abc import Callable
from collections.abc import Callable, Mapping
from types import MappingProxyType
from typing import Final
from litellm.integrations.otel.presets.agentops import agentops_preset
@ -20,7 +21,10 @@ from litellm.integrations.otel.presets.langfuse import (
)
from litellm.integrations.otel.presets.langtrace import langtrace_preset
from litellm.integrations.otel.presets.levo import levo_preset
from litellm.integrations.otel.presets.phoenix import phoenix_preset
from litellm.integrations.otel.presets.phoenix import (
phoenix_preset,
phoenix_project_headers,
)
from litellm.integrations.otel.presets.weave import weave_dynamic_headers, weave_preset
from litellm.types.utils import StandardCallbackDynamicParams
@ -47,6 +51,23 @@ DYNAMIC_HEADERS_BY_CALLBACK: Final[dict[str, Callable[[StandardCallbackDynamicPa
}
#: Callback name → per-request *routing* header builder, sourced from the key/team
#: config the proxy resolved at auth. Deliberately separate from
#: ``DYNAMIC_HEADERS_BY_CALLBACK``: that one is fed
#: ``StandardCallbackDynamicParams``, which is populated from client-supplied
#: request metadata. Naming a destination project is a data-exfiltration
#: primitive, so it must only ever come from server-set key/team config.
PROJECT_HEADERS_BY_CALLBACK: Final[Mapping[str, Callable[[Mapping[str, str] | None], Mapping[str, str]]]] = (
MappingProxyType(
{
"arize_phoenix": phoenix_project_headers,
}
)
)
_NO_PROJECT_HEADERS: Final[Mapping[str, str]] = MappingProxyType({})
def dynamic_otlp_headers(
callback_name: str | None,
dynamic_params: StandardCallbackDynamicParams | None,
@ -62,9 +83,25 @@ def dynamic_otlp_headers(
return headers or None
def project_routing_headers(
callback_name: str | None,
auth_metadata: Mapping[str, str] | None,
) -> Mapping[str, str]:
"""Per-request project-routing headers from trusted key/team config.
Empty means "no per-request project" the caller keeps its default tracer,
whose resource attributes carry the env-configured project.
"""
builder: Final = PROJECT_HEADERS_BY_CALLBACK.get(callback_name or "")
if builder is None:
return _NO_PROJECT_HEADERS
return builder(auth_metadata)
__all__ = [
"DYNAMIC_HEADERS_BY_CALLBACK",
"PRESET_BY_CALLBACK",
"PROJECT_HEADERS_BY_CALLBACK",
"Preset",
"agentops_preset",
"arize_preset",
@ -73,5 +110,6 @@ __all__ = [
"langtrace_preset",
"levo_preset",
"phoenix_preset",
"project_routing_headers",
"weave_preset",
]

View file

@ -1,5 +1,7 @@
"""Arize-Phoenix preset."""
from collections.abc import Mapping
from types import MappingProxyType
from typing import Final
from pydantic import AliasChoices, Field
@ -25,6 +27,36 @@ class _PhoenixSettings(BaseSettings):
)
#: Phoenix routes an OTLP/HTTP export to a project by this header, which takes
#: precedence over the ``openinference.project.name`` resource attribute the env
#: var sets. Requires arize-phoenix 15.5.0+; older collectors ignore it and the
#: spans land in the resource attribute's project.
PHOENIX_PROJECT_HEADER: Final = "x-project-name"
#: Key/team config fields naming the target project, highest precedence first.
_PROJECT_KEYS: Final = ("phoenix_project_name_override", "phoenix_project_name")
_NO_PROJECT: Final[Mapping[str, str]] = MappingProxyType({})
def phoenix_project_headers(auth_metadata: Mapping[str, str] | None) -> Mapping[str, str]:
"""The per-request Phoenix project header for this key/team, if any.
``auth_metadata`` must be the key/team config the proxy resolved at auth
(``user_api_key_auth_metadata``), never client-supplied request metadata:
choosing the destination project is a data-exfiltration primitive, so a
caller must not be able to name one. Returns an empty mapping when the key
and team name no project, leaving the request on the env-configured default.
"""
if not auth_metadata:
return _NO_PROJECT
project: Final = next(
(stripped for key in _PROJECT_KEYS if (stripped := (auth_metadata.get(key) or "").strip())),
"",
)
return MappingProxyType({PHOENIX_PROJECT_HEADER: project}) if project else _NO_PROJECT
def phoenix_preset(
*,
config_overrides: OpenTelemetryV2Config | None = None,

View file

@ -4067,9 +4067,10 @@ class PrometheusLogger(CustomLogger):
require_auth (bool, optional): Whether to require authentication for the metrics endpoint.
Defaults to False.
"""
from prometheus_client import make_asgi_app
from prometheus_client import REGISTRY
from litellm._logging import verbose_proxy_logger
from litellm.integrations.prometheus_metrics_endpoint import make_metrics_asgi_app
from litellm.proxy.proxy_server import app
# Create metrics ASGI app
@ -4078,9 +4079,9 @@ class PrometheusLogger(CustomLogger):
registry: Final = CollectorRegistry()
multiprocess.MultiProcessCollector(registry)
metrics_app = make_asgi_app(registry)
metrics_app = make_metrics_asgi_app(registry)
else:
metrics_app = make_asgi_app()
metrics_app = make_metrics_asgi_app(REGISTRY)
# Mount the metrics app to the app
app.mount("/metrics", metrics_app)

View file

@ -0,0 +1,100 @@
"""ASGI app for `/metrics` that keeps registry rendering off the event loop.
``prometheus_client.make_asgi_app`` collects and serializes the whole registry
inline in the coroutine, so a large scrape (tens of MB on high cardinality
deployments) blocks every other request on the loop for its whole duration. This
app renders in a worker thread instead, shares one render across concurrent
scrapes that want the same output, and streams the payload back in chunks.
"""
from __future__ import annotations
import asyncio
import gzip
from collections.abc import Callable, Iterator, Mapping
from dataclasses import dataclass
from types import MappingProxyType
from typing import Final
from prometheus_client import CollectorRegistry
from prometheus_client.exposition import choose_encoder, gzip_accepted
from starlette.requests import Request
from starlette.responses import StreamingResponse
from starlette.types import ASGIApp, Receive, Scope, Send
RESPONSE_CHUNK_SIZE_BYTES: Final = 64 * 1024
_GZIP_HEADERS: Final = MappingProxyType({"Content-Encoding": "gzip"})
@dataclass(frozen=True, slots=True)
class ScrapeRequest:
"""What a scrape asks for, normalized so that header spellings sharing an output share a render."""
encoder: Callable[[CollectorRegistry], bytes]
content_type: str
gzipped: bool
metric_names: tuple[str, ...]
def parse_scrape_request(accept: str, accept_encoding: str, metric_names: tuple[str, ...]) -> ScrapeRequest:
encoder, content_type = choose_encoder(accept)
return ScrapeRequest(
encoder=encoder,
content_type=content_type,
gzipped=gzip_accepted(accept_encoding),
metric_names=metric_names,
)
def render_scrape(registry: CollectorRegistry, request: ScrapeRequest) -> bytes:
rendered: Final = request.encoder(
registry.restricted_registry(request.metric_names) if request.metric_names else registry # pyright: ignore[reportArgumentType] # RestrictedRegistry is registry-shaped but not a subclass
)
return gzip.compress(rendered) if request.gzipped else rendered
class CoalescedScrapeRenderer:
"""Renders the registry in a worker thread, sharing one render per distinct output across concurrent scrapes."""
def __init__(self, registry: CollectorRegistry) -> None:
self._registry = registry
self._inflight: Mapping[ScrapeRequest, asyncio.Task[bytes]] = MappingProxyType({})
def _forget(self, finished: asyncio.Task[bytes]) -> None:
self._inflight = MappingProxyType({key: task for key, task in self._inflight.items() if task is not finished})
async def render(self, request: ScrapeRequest) -> bytes:
inflight: Final = self._inflight.get(request)
if inflight is not None:
return await asyncio.shield(inflight)
task: Final = asyncio.create_task(asyncio.to_thread(render_scrape, self._registry, request))
self._inflight = MappingProxyType({**self._inflight, request: task})
task.add_done_callback(self._forget)
return await asyncio.shield(task)
def _chunks(body: bytes) -> Iterator[bytes]:
return (body[start : start + RESPONSE_CHUNK_SIZE_BYTES] for start in range(0, len(body), RESPONSE_CHUNK_SIZE_BYTES))
def make_metrics_asgi_app(registry: CollectorRegistry) -> ASGIApp:
renderer: Final = CoalescedScrapeRenderer(registry)
async def metrics_app(scope: Scope, receive: Receive, send: Send) -> None:
request: Final = Request(scope, receive)
scrape: Final = parse_scrape_request(
accept=request.headers.get("accept", ""),
accept_encoding=request.headers.get("accept-encoding", ""),
metric_names=tuple(request.query_params.getlist("name[]")),
)
body: Final = await renderer.render(scrape)
response: Final = StreamingResponse(
_chunks(body),
media_type=scrape.content_type,
headers=_GZIP_HEADERS if scrape.gzipped else None,
)
await response(scope, receive, send)
return metrics_app

View file

@ -165,7 +165,7 @@ class PromptManagementBase(ABC):
ignore_prompt_manager_optional_params: bool | None = False,
) -> tuple[str, list[AllMessageValues], dict]:
if prompt_id is None:
raise ValueError("prompt_id is required for Prompt Management Base class")
return model, messages, non_default_params
if not self.should_run_prompt_management(
prompt_id=prompt_id,
prompt_spec=prompt_spec,

View file

@ -10,7 +10,7 @@ import asyncio
import hashlib
import random
import traceback
from collections.abc import Callable, Mapping, Sequence
from collections.abc import Awaitable, Callable, Mapping, Sequence
from dataclasses import dataclass
from datetime import datetime, timezone
from itertools import groupby
@ -42,8 +42,9 @@ if TYPE_CHECKING:
from litellm.router import Router
from litellm.types.utils import StandardLoggingPayload
# A job starting, stopping, or hitting its turn budget propagates to sampling within one
# TTL; the turn budget can overshoot by at most one TTL of in-flight samples per pod.
# A job starting, stopping, or hitting a budget propagates to sampling within one TTL;
# the spend gate re-checks the cross-pod counter at pipeline entry, so it overshoots
# only by the samples already in flight when the cap is crossed.
_JOBS_CACHE_TTL_SECONDS: Final = 10
# Concurrent shadow+judge pipelines per pod: a traffic spike turns into skipped samples
@ -340,13 +341,24 @@ def _failure_detail(e: BaseException) -> str:
return f"{type(e).__name__}{location}: {e}"
def _judge_call_cost(response: object) -> float:
"""Price a judge call, treating an unmapped judge model as free rather than fatal."""
def _call_cost(response: object) -> float:
"""Price one eval-arm call with the figure the spend pipeline bills: the router client
stamps _hidden_params.response_cost from the deployment's own pricing, which the public
price map lookup below cannot see (it reads 0 for deployment-priced models)."""
getter: Final = getattr(getattr(response, "_hidden_params", None), "get", None)
stamped: Final = getter("response_cost") if callable(getter) else None
if isinstance(stamped, (int, float)):
return float(stamped)
return _price_map_cost(response)
def _price_map_cost(response: object) -> float:
"""Public price map fallback, treating an unmapped model as free rather than fatal."""
import litellm
try:
return litellm.completion_cost(completion_response=response) or 0.0
except Exception: # noqa: BLE001 # unmapped judge model: the verdict still counts, cost stays 0
except Exception: # noqa: BLE001 # unmapped model: the attempt still counts, cost stays 0
return 0.0
@ -374,6 +386,32 @@ def _judge_user_prompt(conversation: str, response_a: str, response_b: str) -> s
)
def _job_spend_counter_key(job_id: str) -> str:
return f"spend:shadow_eval:{job_id}"
async def _job_spend_from_counter(counter_key: str, fallback_spend: float, max_budget: float) -> float:
"""The leg's spend through the cross-pod counter the key budget gates read. The owner
degrades internally to the fill-time DB floor and raises only under fail-closed
enforcement, which the caller honors by skipping the sample."""
from litellm.proxy.proxy_server import get_current_spend
return await get_current_spend(counter_key=counter_key, fallback_spend=fallback_spend, max_budget=max_budget)
async def _add_job_spend_to_counter(counter_key: str, cost: float) -> None:
"""Advance the counter the moment a cost is known, so even a lost row closes the gate.
Known failure mode: a Redis outage freezes the counter (the owner invalidates it), the
gate degrades to the fill floor, and overshoot grows to in-flight plus one TTL of
samples, the same degradation the key budget counters accept."""
try:
from litellm.proxy.proxy_server import increment_spend_counter
await increment_spend_counter(counter_key=counter_key, increment=cost)
except Exception as e: # noqa: BLE001 # attempt recording must proceed; the row stays truth and the fill floor gates
verbose_logger.warning("shadow_eval: spend counter increment failed for %s: %s", counter_key, e)
async def _key_or_team_is_over_budget(metadata: Mapping[str, object]) -> bool:
"""Whether the shadowed key or its team is over budget, decided by the same owners
the request path uses, so counter keys and thresholds can never drift from auth's.
@ -438,8 +476,8 @@ def _request_was_routed_by(request_metadata: Mapping[str, object], router_name:
@dataclass(frozen=True, slots=True)
class _CallFailure:
"""A shadow or judge call that produced no usable response. cost carries any judge
spend the failed attempt still billed, so job-level judge_spend never undercounts."""
"""A shadow or judge call that produced no usable response. cost carries any spend
the failed call still billed, so job-level spend figures never undercount."""
error: str
cost: float = 0.0
@ -452,6 +490,7 @@ class _ShadowResponse:
text: str
model: str
tier: str | None
cost: float
@dataclass(frozen=True, slots=True)
@ -478,8 +517,10 @@ class ActiveShadowEvalJob(BaseModel):
shadow_percentage: float
judge_model: str
max_turns: int
max_budget: float | None = None
ends_at: datetime
attempts: int = 0
spend: float = 0.0
@field_validator("ends_at")
@classmethod
@ -500,7 +541,7 @@ class ActiveShadowEvalJob(BaseModel):
return self.baseline_model or self.router_name
def _as_active_job(record: object, attempts: int) -> ActiveShadowEvalJob | None:
def _as_active_job(record: object, attempts: int, spend: float) -> ActiveShadowEvalJob | None:
"""The sampling path's view of one job row, or None for a row it cannot sample: an
unknown direction, or a reverse job with no baseline model to duplicate against.
Failing closed here is what keeps the dispatch path total."""
@ -509,7 +550,7 @@ def _as_active_job(record: object, attempts: int) -> ActiveShadowEvalJob | None:
except ValidationError as e:
verbose_logger.debug("shadow_eval: skipping unsamplable job row: %s", e)
return None
return job.model_copy(update={"attempts": attempts})
return job.model_copy(update={"attempts": attempts, "spend": spend}) # mutable-ok: pydantic update payload
_jobs_cache: Final = InMemoryCache(max_size_in_memory=4, default_ttl=_JOBS_CACHE_TTL_SECONDS)
@ -524,12 +565,17 @@ class ShadowEvalLogger(CustomLogger):
router_provider: Callable[[], "Router | None"] | None = None,
prisma_provider: Callable[[], "PrismaClient | None"] | None = None,
jobs_cache: InMemoryCache | None = None,
job_spend_reader: Callable[[str, float, float], Awaitable[float]] | None = None,
job_spend_writer: Callable[[str, float], Awaitable[None]] | None = None,
) -> None:
"""Providers are callables so the proxy's lazily-initialized globals are resolved
at call time, not at logger construction."""
at call time, not at logger construction. The spend reader and writer wrap the
proxy's cross-pod spend counter; tests inject a plain in-memory pair."""
self._router_provider = router_provider or default_router_provider
self._prisma_provider = prisma_provider or _default_prisma_provider
self._jobs_cache = jobs_cache or _jobs_cache
self._read_job_spend = job_spend_reader or _job_spend_from_counter
self._write_job_spend = job_spend_writer or _add_job_spend_to_counter
self._inflight_shadow_tasks: int = 0
# Starts per job since the last cache fill, never decremented within a
# generation; the refill absorbs written rows and resets.
@ -556,18 +602,26 @@ class ShadowEvalLogger(CustomLogger):
await prisma.db.litellm_shadowevalattempt.group_by(
by=["job_id"],
count=True,
sum={"judge_cost": True, "shadow_cost": True}, # mutable-ok: Prisma aggregate spec
where={"job_id": {"in": [str(record.id) for record in records]}}, # mutable-ok: Prisma filter
)
if records
else ()
)
attempt_counts: Final = {str(row["job_id"]): int(row["_count"]["_all"]) for row in grouped or []}
attempt_stats: Final = { # mutable-ok: frozen snapshot of the grouped read
str(row["job_id"]): (
int(row["_count"]["_all"]),
float((row["_sum"] or {}).get("judge_cost") or 0.0)
+ float((row["_sum"] or {}).get("shadow_cost") or 0.0),
)
for row in grouped or []
}
by_key: Final = tuple(
sorted(
(
(str(record.api_key_id), job)
for record in records or []
if (job := _as_active_job(record, attempt_counts.get(str(record.id), 0))) is not None
if (job := _as_active_job(record, *attempt_stats.get(str(record.id), (0, 0.0)))) is not None
),
key=itemgetter(0),
)
@ -624,6 +678,7 @@ class ShadowEvalLogger(CustomLogger):
for job in (await self._active_jobs()).get(str(api_key_hash), ())
if datetime.now(timezone.utc) < job.ends_at
and job.attempts + self._job_starts.get(job.id, 0) < job.max_turns
and (job.max_budget is None or job.spend < job.max_budget)
and _sample_hits(request_id, job.id, job.shadow_percentage)
and _request_was_routed_by(request_metadata, job.router_name) == (job.direction == "reverse")
)
@ -684,12 +739,28 @@ class ShadowEvalLogger(CustomLogger):
return
if await _key_or_team_is_over_budget(parent_metadata):
return
if job.max_budget is not None:
try:
spend: Final = await self._read_job_spend(_job_spend_counter_key(job.id), job.spend, job.max_budget)
except Exception as e: # noqa: BLE001 # unverifiable budget: skip the sample rather than spend on it
verbose_logger.warning("shadow_eval: budget unverifiable for %s, sample skipped: %s", job.id, e)
return
if spend >= job.max_budget:
return
shadow: Final = await self._call_router_shadow(job.shadow_target, messages, shadow_params, parent_metadata)
if isinstance(shadow, _CallFailure):
await self._record_attempt(prisma, job, request_id, control_tier, outcome="error", error=shadow.error)
return
except Exception as e: # noqa: BLE001 # detached task: nothing billed yet, record and never raise
verbose_logger.debug("shadow_eval: pipeline failed for %s: %s", request_id, e)
await self._record_attempt(
prisma, job, request_id, control_tier, outcome="error", error=f"pipeline error: {e}"
)
return
if isinstance(shadow, _CallFailure):
await self._record_attempt(
prisma, job, request_id, control_tier, outcome="error", error=shadow.error, shadow_cost=shadow.cost
)
return
# From here the shadow call has billed, so every exit records its cost.
try:
verdict: Final = await self._call_judge(
judge_model=job.judge_model,
messages=messages,
@ -707,6 +778,7 @@ class ShadowEvalLogger(CustomLogger):
error=verdict.error,
shadow=shadow,
judge_cost=verdict.cost,
shadow_cost=shadow.cost,
)
return
await self._record_attempt(
@ -719,15 +791,23 @@ class ShadowEvalLogger(CustomLogger):
real_model=real_model,
confidence=verdict.confidence,
judge_cost=verdict.cost,
shadow_cost=shadow.cost,
)
except Exception as e: # noqa: BLE001 # detached task: record what happened, never raise
except Exception as e: # noqa: BLE001 # detached task: the shadow call billed, record its cost, never raise
verbose_logger.debug("shadow_eval: pipeline failed for %s: %s", request_id, e)
await self._record_attempt(
prisma, job, request_id, control_tier, outcome="error", error=f"pipeline error: {e}"
prisma,
job,
request_id,
control_tier,
outcome="error",
error=f"pipeline error: {e}",
shadow=shadow,
shadow_cost=shadow.cost,
)
@staticmethod
async def _record_attempt(
self,
prisma: "PrismaClient | None",
job: ActiveShadowEvalJob,
request_id: str,
@ -738,8 +818,11 @@ class ShadowEvalLogger(CustomLogger):
real_model: str = "",
confidence: float | None = None,
judge_cost: float = 0.0,
shadow_cost: float = 0.0,
error: str | None = None,
) -> None:
if judge_cost + shadow_cost > 0:
await self._write_job_spend(_job_spend_counter_key(job.id), judge_cost + shadow_cost)
if prisma is None:
return
try:
@ -753,6 +836,7 @@ class ShadowEvalLogger(CustomLogger):
"shadow_model": shadow.model if shadow else None,
"confidence": confidence,
"judge_cost": judge_cost,
"shadow_cost": shadow_cost,
"error": error[:_MAX_ERROR_CHARS] if error else None,
}
)
@ -792,11 +876,12 @@ class ShadowEvalLogger(CustomLogger):
return _CallFailure(f"shadow router call failed: {_failure_detail(e)}")
text: Final = _chat_final_text(response)
if not text:
return _CallFailure("shadow router returned an empty response")
return _CallFailure("shadow router returned an empty response", cost=_call_cost(response))
return _ShadowResponse(
text=text,
model=str(getattr(response, "model", None) or _routing_decision(shadow_metadata).get("routed_model") or ""),
tier=_routed_tier(shadow_metadata),
cost=_call_cost(response),
)
async def _call_judge(
@ -843,11 +928,11 @@ class ShadowEvalLogger(CustomLogger):
verdict: Final = PairwiseVerdict.model_validate(parse_json_verdict(raw))
except Exception as e: # noqa: BLE001 # malformed verdicts become error rows
verbose_logger.debug("shadow_eval: unparseable judge verdict: %s", e)
return _CallFailure(f"unparseable judge verdict: {e}", cost=_judge_call_cost(response))
return _CallFailure(f"unparseable judge verdict: {e}", cost=_call_cost(response))
return _JudgeVerdict(
preference=_unmask_preference(verdict.preference, real_is_a),
confidence=max(0.0, min(1.0, verdict.confidence)),
cost=_judge_call_cost(response),
cost=_call_cost(response),
)

View file

@ -0,0 +1,231 @@
"""
CLI Keyring Access
SDK-level access to the OS keychain (macOS Keychain, Windows Credential Manager,
Linux Secret Service) that holds the credential minted by `lite login`.
The `keyring` package is optional and imported lazily, so importing this module
never pulls it in. Every failure is returned as a value, naming which of the
ways the keychain can be out of reach applies, so callers can degrade to the
token file and tell the user what to do about it.
A write is only reported as stored once it has been read back, because keyring's
null backend, which `keyring --disable` and headless CI images both select,
accepts every write and keeps nothing. Writes are also pre-flighted with a
throwaway value, because a keychain can answer neither way and block forever.
"""
import os
import threading
from contextlib import suppress
from dataclasses import dataclass, field
from typing import Final, Protocol, TypeAlias
KEYRING_SERVICE: Final = "litellm-cli"
KEYRING_ACCOUNT: Final = "credential"
KEYRING_PREFLIGHT_ACCOUNT: Final = "credential-preflight"
DISABLE_KEYRING_ENV_VAR: Final = "LITELLM_CLI_DISABLE_KEYRING"
_DISABLED_VALUES: Final = frozenset(("1", "true", "yes", "on"))
_PREFLIGHT_VALUE: Final = "preflight"
_PREFLIGHT_TIMEOUT_SECONDS: Final = 5.0
@dataclass(frozen=True, slots=True)
class SecretFound:
blob: str
@dataclass(frozen=True, slots=True)
class SecretMissing:
pass
@dataclass(frozen=True, slots=True)
class SecretStored:
pass
@dataclass(frozen=True, slots=True)
class SecretErased:
pass
@dataclass(frozen=True, slots=True)
class SecretStranded:
pass
@dataclass(frozen=True, slots=True)
class KeyringNotInstalled:
pass
@dataclass(frozen=True, slots=True)
class KeyringDisabled:
pass
@dataclass(frozen=True, slots=True)
class KeyringUnreachable:
pass
@dataclass(frozen=True, slots=True)
class KeyringDiscardsWrites:
pass
KeyringUnusable: TypeAlias = KeyringNotInstalled | KeyringDisabled | KeyringUnreachable
SecretRead: TypeAlias = SecretFound | SecretMissing | KeyringUnusable
SecretWrite: TypeAlias = SecretStored | KeyringUnusable | KeyringDiscardsWrites
SecretErase: TypeAlias = SecretErased | SecretStranded | KeyringUnusable
class SecretVault(Protocol):
"""The single slot holding the CLI credential's secret material."""
def read(self) -> SecretRead: ...
def write(self, blob: str) -> SecretWrite: ...
def erase(self) -> SecretErase: ...
class KeyringApi(Protocol):
def get_password(self, service_name: str, username: str) -> str | None: ...
def set_password(self, service_name: str, username: str, password: str) -> None: ...
def delete_password(self, service_name: str, username: str) -> None: ...
def _keyring_disabled() -> bool:
return os.getenv(DISABLE_KEYRING_ENV_VAR, "").strip().lower() in _DISABLED_VALUES
def _import_keyring() -> KeyringApi | None:
try:
import keyring
except ImportError:
return None
return keyring
def _keyring_api() -> KeyringApi | KeyringNotInstalled | KeyringDisabled:
if _keyring_disabled():
return KeyringDisabled()
api: Final = _import_keyring()
return KeyringNotInstalled() if api is None else api
def _answers_a_write(api: KeyringApi, timeout_seconds: float) -> bool:
"""Whether the keychain answers a write at all, asked with a value worth nothing.
macOS derives the login keychain from `$HOME`, and `set_password` against a HOME with no usable
one blocks forever with no timeout of its own. Containers, CI images, `sudo -H`, and service
accounts all run there, and reads answer normally, so nothing cheaper tells them apart. Asking
with a throwaway value keeps a keychain that never answers from taking `lite login` down with
it, and keeps the real credential out of a store that might accept it long after we gave up.
A keychain that refuses the probe outright still answered it, so only silence counts against it.
"""
answered: Final = threading.Event()
def ask() -> None:
with suppress(Exception):
api.set_password(KEYRING_SERVICE, KEYRING_PREFLIGHT_ACCOUNT, _PREFLIGHT_VALUE)
answered.set()
threading.Thread(target=ask, daemon=True, name="litellm-cli-keyring-preflight").start()
return answered.wait(timeout_seconds)
def _forget_the_preflight(api: KeyringApi) -> None:
"""Take the throwaway probe back out.
A backend that kept nothing has nothing to remove, and the probe is worth nothing either way,
so a keychain that refuses to give it up costs the caller nothing.
"""
with suppress(Exception):
api.delete_password(KEYRING_SERVICE, KEYRING_PREFLIGHT_ACCOUNT)
@dataclass(frozen=True, slots=True)
class KeyringVault:
"""The OS keychain, reached through the optional `keyring` package.
A keychain that let the pre-flight time out is not asked anything else for the rest of the
process. The probe that timed out is still sitting in the keychain on a thread of its own, and
it holds the keychain against every later call, so the read after it would block on the main
thread with no timeout to save it. One silence is answer enough.
"""
preflight_timeout_seconds: float = _PREFLIGHT_TIMEOUT_SECONDS
stopped_answering: threading.Event = field(default_factory=threading.Event, compare=False, repr=False)
def read(self) -> SecretRead:
if self.stopped_answering.is_set():
return KeyringUnreachable()
api: Final = _keyring_api()
if isinstance(api, (KeyringNotInstalled, KeyringDisabled)):
return api
try:
blob: Final = api.get_password(KEYRING_SERVICE, KEYRING_ACCOUNT)
except Exception: # noqa: BLE001 # backends raise outside keyring.errors; never break the SDK
return KeyringUnreachable()
return SecretMissing() if blob is None else SecretFound(blob)
def write(self, blob: str) -> SecretWrite:
"""Store the secret, reporting stored only once the keychain hands the same bytes back.
A backend that accepts writes and keeps nothing, which is exactly what `keyring --disable`
and `PYTHON_KEYRING_BACKEND=keyring.backends.null.Keyring` select, raises nothing to
distinguish itself. Reading the value back is the only way to tell it apart from a keychain
that really stored the credential, and the caller is about to drop its own copy on our word.
The keychain is pre-flighted first, because one that blocks rather than answering would
otherwise hang `lite login` outright.
"""
if self.stopped_answering.is_set():
return KeyringUnreachable()
api: Final = _keyring_api()
if isinstance(api, (KeyringNotInstalled, KeyringDisabled)):
return api
if not _answers_a_write(api, self.preflight_timeout_seconds):
self.stopped_answering.set()
return KeyringUnreachable()
_forget_the_preflight(api)
try:
api.set_password(KEYRING_SERVICE, KEYRING_ACCOUNT, blob)
except Exception: # noqa: BLE001 # a keychain that refuses the write falls back to the token file
return KeyringUnreachable()
return SecretStored() if self.read() == SecretFound(blob) else KeyringDiscardsWrites()
def erase(self) -> SecretErase:
"""Remove our entry, reporting whether the keychain is guaranteed to be free of it.
A keychain out of reach is never an erasure: the entry belongs to the OS, not to this
install, so it outlives an uninstalled `keyring` package and a kill switch set after login.
Those cases are reported apart from a confirmed entry that would not delete, because only
the caller knows whether this machine ever put a secret in a keychain.
"""
match self.read():
case KeyringNotInstalled() | KeyringDisabled() | KeyringUnreachable() as unusable:
return unusable
case SecretMissing():
return SecretErased()
case SecretFound():
return self._delete()
def _delete(self) -> SecretErase:
api: Final = _keyring_api()
if isinstance(api, (KeyringNotInstalled, KeyringDisabled)):
return api
try:
api.delete_password(KEYRING_SERVICE, KEYRING_ACCOUNT)
except Exception: # noqa: BLE001 # report the failure as a value so `lite logout` can warn
return SecretStranded()
return SecretErased()
SYSTEM_KEYRING: Final[SecretVault] = KeyringVault()

View file

@ -1,16 +1,134 @@
"""
CLI Token Utilities
SDK-level utilities for reading CLI authentication tokens.
SDK-level utilities for reading the credential minted by `lite login`.
Non-secret metadata lives in ~/.litellm/token.json. The secret material (the
bearer key, the refresh token that renews it, and a JWT when one is issued)
lives in the OS keychain when the machine has one, and in that same 0600 file
otherwise. This module hides the split from callers, and migrates a plaintext
file into the keychain the first time it reads one.
This module has no dependencies on proxy code and can be safely imported at the SDK level.
"""
import json
import os
import math
import time
from collections.abc import Mapping
from dataclasses import dataclass
from pathlib import Path
from typing import Final
from types import MappingProxyType
from typing import Final, TypeAlias
from pydantic import BaseModel, ConfigDict, ValidationError
from litellm.litellm_core_utils.cli_keyring import (
SYSTEM_KEYRING,
KeyringDisabled,
KeyringNotInstalled,
KeyringUnreachable,
SecretErase,
SecretErased,
SecretFound,
SecretMissing,
SecretStored,
SecretStranded,
SecretVault,
SecretWrite,
)
from litellm.litellm_core_utils.private_json import (
commit_staged_json,
discard_staged_json,
ensure_private_dir,
overwrite_private_json,
stage_private_json,
write_private_json,
)
@dataclass(frozen=True, slots=True)
class CredentialNotSaved:
"""The credential was minted but no store would keep it, so this machine has none.
Nothing was touched on the way to this, so a login that already worked still does.
"""
detail: str
@dataclass(frozen=True, slots=True)
class CredentialNotRecorded:
"""The keychain took the credential, but the file that names it could not be replaced.
The keychain holds one entry, so the secret that was there is already gone and no rollback
brings it back. Removing the new one as well would only turn a login this machine may still
be able to use into no login at all, so it stays, and the user is told what is where.
"""
@dataclass(frozen=True, slots=True)
class CredentialNotCleared:
"""The token file still holds the secret, because it could not be removed or rewritten.
Logging out of the keychain is only half of it. A `~/.litellm` that refuses both the scrubbed
rewrite and the removal leaves the credential readable on disk, which is the one thing a logout
is for, so it is reported instead of being counted as a clean sweep.
"""
detail: str
SecretSave: TypeAlias = SecretWrite | CredentialNotSaved | CredentialNotRecorded
SecretClear: TypeAlias = SecretErase | CredentialNotCleared
class CliTokenRecord(BaseModel):
"""A stored CLI credential.
`key is None` means the metadata was found but the secret could not be
produced: the keychain holds nothing for us, or we could not reach it.
"""
model_config = ConfigDict(frozen=True, extra="allow")
base_url: str = ""
key: str | None = None
user_id: str = ""
user_email: str = ""
user_role: str = ""
auth_header_name: str = "Authorization"
jwt_token: str = ""
timestamp: float = 0.0
expires_at: float | None = None
refresh_token: str | None = None
class CliTokenSecret(BaseModel):
"""The secret material as stored in the OS keychain.
`base_url` is duplicated from the metadata file purely as a pairing tag: a
secret minted for one server is never handed to another, even if the
metadata file is edited underneath us. `timestamp` is the sign-in this
secret came from, which is what decides it against a secret still on disk.
Every field a thief could sign in with belongs here, which is why the
refresh token is one of them: it buys a fresh key from the proxy on demand,
so leaving it on disk would leave the login readable there. `key` is
optional because the file can hold a refresh token without one, and moving
that into the keychain must not invent a key to go with it.
"""
model_config = ConfigDict(frozen=True)
base_url: str
key: str | None = None
jwt_token: str = ""
refresh_token: str | None = None
timestamp: float = 0.0
CLI_TOKEN_FRESHNESS_BUFFER_SECONDS: Final = 360
def get_cli_token_file_path() -> str:
@ -20,26 +138,183 @@ def get_cli_token_file_path() -> str:
return str(config_dir / "token.json")
def load_cli_token() -> dict | None:
"""Load CLI token data from file"""
token_file: Final = get_cli_token_file_path()
if not os.path.exists(token_file):
def load_cli_token(*, vault: SecretVault = SYSTEM_KEYRING) -> CliTokenRecord | None:
"""Load the stored CLI credential, or None when this machine has none"""
record: Final = _read_token_file()
if record is None:
return None
return _resolve_secret(record, vault)
def save_cli_token(record: CliTokenRecord, *, vault: SecretVault = SYSTEM_KEYRING) -> SecretSave:
"""Store a freshly minted credential. Reports where its secret material ended up, and why.
The token file is what makes a keychain-backed credential findable again, and it is also the
half that a read-only or full directory refuses, so it is staged before the keychain is handed
anything. A save that cannot land then leaves both stores exactly as it found them, which
matters most when the login it failed to replace is still perfectly good.
Staging can still succeed and the replacement fail afterwards. That is the one case where the
keychain has already taken the new secret, and it reports itself as such rather than claiming
the previous login survived.
"""
stamped: Final = _stamped_past_every_stored_login(record, vault)
staged: Final = _stage_token_file(_without_secret(stamped))
if isinstance(staged, CredentialNotSaved):
return staged
outcome: Final = vault.write(_encode_secret(stamped)) if _holds_a_secret(stamped) else SecretStored()
if isinstance(outcome, SecretStored):
return outcome if _commit_token_file(staged) else CredentialNotRecorded()
discard_staged_json(staged)
return _keep_the_secret_in_the_file(stamped, outcome)
def _stamped_past_every_stored_login(record: CliTokenRecord, vault: SecretVault) -> CliTokenRecord:
"""Keep a sign-in's stamp ahead of every login already stored, whatever the clock did in between.
The stamp is what decides a keychain secret against one still on disk, so a clock that stepped
backwards between two logins would hand the older of them the win and put a superseded
credential back in use. Pinning the new stamp just past the highest one either store holds costs
one read each and changes nothing on a clock that only moves forwards.
"""
highest: Final = _highest_stamp_already_stored(record.base_url, vault)
if highest < record.timestamp:
return record
return record.model_copy(update=MappingProxyType({"timestamp": math.nextafter(highest, math.inf)}))
def _highest_stamp_already_stored(base_url: str, vault: SecretVault) -> float:
"""When the latest login either store still holds was made, or minus infinity when neither has one.
Both are asked because the file names the login being replaced only while the two agree. A login
the keychain took but the file could not record afterwards leaves the keychain holding the later
of the two, and reading only the file would stamp the next sign-in below it.
"""
previous: Final = _read_token_file()
secret: Final = _stored_secret(base_url, vault)
return max(
-math.inf if previous is None else previous.timestamp,
-math.inf if secret is None else secret.timestamp,
)
def _stored_secret(base_url: str, vault: SecretVault) -> CliTokenSecret | None:
"""The keychain's secret for this server, when it holds one this login may be compared against"""
match vault.read():
case SecretFound(blob=blob):
return _decode_secret(blob, base_url)
case SecretMissing() | KeyringNotInstalled() | KeyringDisabled() | KeyringUnreachable():
return None
def _keep_the_secret_in_the_file(record: CliTokenRecord, outcome: SecretWrite) -> SecretSave:
"""Fall back to the owner-only file, which is all that is left when no keychain took the secret"""
try:
with open(token_file, "r") as f:
return json.load(f)
except (OSError, json.JSONDecodeError):
return None
_write_token_file(record)
except OSError as error:
return CredentialNotSaved(str(error))
return outcome
def clear_cli_token(*, vault: SecretVault = SYSTEM_KEYRING) -> SecretClear:
"""Remove the credential from both stores. Reports whether the keychain is now free of it.
A logout the keychain never answered keeps the token file, with its secret taken out, because
that file is the only remaining record that something may still be in there to remove. It is
what lets a later run tell a machine with a credential it cannot reach apart from one that never
had a login at all, and taking it away would leave the next logout answering the warning this
one just issued with a false all-clear. The secret goes either way, and a file that will give up
neither its copy nor itself is removed rather than kept, with the note written again afterwards
so the warning still outlives this run.
"""
outcome: Final = vault.erase()
record: Final = _read_token_file()
settled: Final = _nothing_left_behind(outcome, record)
if not settled and _keep_the_unchecked_keychain_on_record(outcome, record):
return outcome
removal: Final = _remove_token_file()
if removal is not None and record is not None and not _scrub_file_secret(record):
return removal
if removal is None and record is not None and _the_keychain_went_unchecked(outcome):
_write_the_note_the_removal_took_with_it(record)
return SecretErased() if settled else outcome
def _remove_token_file() -> CredentialNotCleared | None:
try:
Path(get_cli_token_file_path()).unlink(missing_ok=True)
except OSError as error:
return CredentialNotCleared(str(error))
return None
def _write_the_note_the_removal_took_with_it(record: CliTokenRecord) -> None:
"""Put the secret-free note back after the file carrying it had to go to get the secret off disk.
Reaching here means neither rewrite would take, so the file went instead, and its absence is
what the next logout would read as a keychain already known to be clean. Removing it is also
what frees the room the rewrite was refused for, so the note usually lands on this second try.
When it does not, the warning this logout printed is the only one the user gets.
"""
staged: Final = _stage_scrubbed_file(record)
if staged is not None:
_commit_token_file(staged)
def _keep_the_unchecked_keychain_on_record(outcome: SecretErase, record: CliTokenRecord | None) -> bool:
"""Whether the token file, stripped of its secret, is worth keeping as the note that says so.
Only a keychain that could not be reached leaves the question open. One that answered for itself
is remembered without any help from the file, and a file it can still pair a live entry with
would leave the machine signed in to the login that was just ended. A copy that will give up
its secret neither to a staged replacement nor to an overwrite is not kept either, because the
secret goes first.
"""
if record is None or not _the_keychain_went_unchecked(outcome):
return False
return _scrub_file_secret(record)
def _the_keychain_went_unchecked(outcome: SecretErase) -> bool:
"""Whether the keychain neither confirmed the erase nor answered that it still holds the secret"""
match outcome:
case SecretErased() | SecretStranded():
return False
case KeyringDisabled() | KeyringNotInstalled() | KeyringUnreachable():
return True
def _nothing_left_behind(outcome: SecretErase, record: CliTokenRecord | None) -> bool:
"""Whether the keychain can be trusted to hold no credential of ours once the file is gone.
A machine with no token file has no stored login to end, and `clear_cli_token` keeps one behind
whenever the keychain is left unconfirmed, taking the secret out in place when it cannot stage a
replacement and writing the note again when the file holding it had to go, so a missing file is
real evidence rather than the absence of it. Past that, a
keychain that could not be reached is never trusted, whatever the file looks like. Even a file
holding its own secret says only that the login which wrote it had no keychain to write to, and
the login before it may well have had one: the entry that login left outlives both the
uninstalled package and the file that replaced it. `SecretStranded` is
the keychain answering for itself and outranks the file.
"""
match outcome:
case SecretErased():
return True
case SecretStranded():
return False
case KeyringDisabled() | KeyringNotInstalled() | KeyringUnreachable():
return record is None
def get_litellm_gateway_api_key(
expected_base_url: str | None = None,
*,
vault: SecretVault = SYSTEM_KEYRING,
) -> str | None:
"""
Get the stored CLI API key for use with LiteLLM SDK.
This function reads the token file created by `lite login`
This function reads the credential created by `lite login`
and returns the API key for use in Python scripts.
Args:
@ -47,6 +322,7 @@ def get_litellm_gateway_api_key(
originally issued for this URL. Pass the target server URL to
prevent credential leakage when the client is pointed at a
different (possibly malicious) server.
vault: Where the secret material is stored. Defaults to the OS keychain.
Returns:
str: The API key if found (and origin matches), None otherwise
@ -62,25 +338,222 @@ def get_litellm_gateway_api_key(
>>> base_url="https://your-proxy.com/v1"
>>> )
"""
token_data: Final = load_cli_token()
if not token_data or "key" not in token_data:
record: Final = _read_token_file()
if record is None:
return None
if expected_base_url is not None:
stored_url: Final = token_data.get("base_url")
if stored_url != expected_base_url.rstrip("/"):
return None
return token_data["key"]
if expected_base_url is not None and record.base_url != expected_base_url.rstrip("/"):
return None
resolved: Final = _resolve_secret(record, vault)
return None if resolved is None else resolved.key
def is_cli_token_fresh(token_data: Mapping[str, object], buffer_hours: float = 0.1) -> bool:
"""Check whether a cached CLI token (as stored in token.json) is still
within its expiration window. Used by `lite auth print-token` to fail
fast, without a network round trip, once the cached token is past
`LITELLM_CLI_JWT_EXPIRATION_HOURS`."""
def is_cli_token_fresh(
token_data: CliTokenRecord | Mapping[str, object],
buffer_hours: float = CLI_TOKEN_FRESHNESS_BUFFER_SECONDS / 3600,
) -> bool:
"""Check whether a cached CLI token is still within its expiration window.
Used by `lite auth print-token` to fail fast, without a network round trip,
once the cached token is past `LITELLM_CLI_JWT_EXPIRATION_HOURS`. A `--pkce`
credential carries its own `expires_at`, which is authoritative when present."""
from litellm.constants import CLI_JWT_EXPIRATION_HOURS
timestamp: Final = token_data.get("timestamp")
expires_at: Final = (
token_data.expires_at if isinstance(token_data, CliTokenRecord) else token_data.get("expires_at")
)
if isinstance(expires_at, (int, float)):
return time.time() < expires_at - buffer_hours * 3600
timestamp: Final = token_data.timestamp if isinstance(token_data, CliTokenRecord) else token_data.get("timestamp")
if not isinstance(timestamp, (int, float)):
return False
age_hours: Final = (time.time() - timestamp) / 3600
return age_hours < (CLI_JWT_EXPIRATION_HOURS - buffer_hours)
def _read_token_file() -> CliTokenRecord | None:
try:
raw: Final = Path(get_cli_token_file_path()).read_text()
except (OSError, ValueError):
return None
try:
return CliTokenRecord.model_validate_json(raw)
except ValidationError:
return None
def _resolve_secret(record: CliTokenRecord, vault: SecretVault) -> CliTokenRecord | None:
match vault.read():
case SecretFound(blob=blob):
return _apply_vault_secret(record, blob, vault)
case SecretMissing():
return _migrate_file_secret(record, vault)
case KeyringNotInstalled() | KeyringDisabled() | KeyringUnreachable():
return record
def _apply_vault_secret(record: CliTokenRecord, blob: str, vault: SecretVault) -> CliTokenRecord | None:
"""Resolve the credential when both stores hold one.
The sign-in each secret came from decides it, because either store can be the stale one. A
secret is usually left on disk by a keychain that would not take it, which makes the file the
fresher of the two. It is the older one when a login the keychain did take could not replace
the file afterwards, and serving that one would put a superseded credential back in use. Equal
stamps are one login sitting in both stores, left by a migration whose scrub was refused or by
an upgrade that took the key into the keychain and left the refresh token behind, so that branch
rejoins the halves and retries the migration rather than trading one credential for another.
A scrub the file refuses leaves that superseded secret where it lies, which is the state the
login already named when it could not replace the file, and which `lite logout` reports rather
than counting as a clean sweep. Rolling the vault back the way a migration does is not the
answer here, because the two stores hold different credentials and the rollback would hand the
superseded one back out.
"""
secret: Final = _decode_secret(blob, record.base_url)
if secret is None or (_holds_a_secret(record) and secret.timestamp <= record.timestamp):
return _migrate_file_secret(_rejoined(record, secret), vault, replacing=secret)
_scrub_file_secret(record)
return record.model_copy(
update=MappingProxyType(
{
"key": secret.key,
"jwt_token": secret.jwt_token,
"refresh_token": secret.refresh_token,
"timestamp": max(secret.timestamp, record.timestamp),
}
)
)
def _rejoined(record: CliTokenRecord, secret: CliTokenSecret | None) -> CliTokenRecord:
"""Put one sign-in's secret material back together when each store holds part of it.
Upgrading from the release that kept only the key in the keychain leaves the refresh token
behind in the file, so a single login sits across both stores. Filling in whatever the file is
missing before the migration writes its entry is what stops that write from replacing a live key
with nothing. Only a matching stamp is one login. Two stamps are two logins, and pairing one's
key with the other's refresh token would build a credential neither store ever held.
"""
if secret is None or secret.timestamp != record.timestamp:
return record
return record.model_copy(
update=MappingProxyType(
{
"key": record.key if record.key is not None else secret.key,
"jwt_token": record.jwt_token or secret.jwt_token,
"refresh_token": record.refresh_token if record.refresh_token is not None else secret.refresh_token,
}
)
)
def _migrate_file_secret(
record: CliTokenRecord, vault: SecretVault, *, replacing: CliTokenSecret | None = None
) -> CliTokenRecord | None:
"""Move a file-held secret into the vault, but only once the file's copy can be taken away.
The scrubbed file is staged first so a directory that will not accept it stops the migration
before the keychain is handed anything. Copying the credential into a second store and only
then discovering the first one cannot be cleaned would widen exposure instead of narrowing it,
which is the opposite of what moving it into the keychain is for.
A staged file that will not go into place is overwritten where it lies before the keychain is
asked to take the new entry back, so the migration finishes on a directory that would only ever
have refused it. Rolling back is the last resort, and a rollback the keychain also refuses
leaves the secret in both stores until the next read, which retries this same migration.
Only an entry this migration put there is taken back. `replacing` names one that was already in
the keychain, whose material the new entry carries forward, so erasing it would take away the
half the file never had, and a machine that refuses the scrub is exactly the one with nowhere
else to keep it. The next read finds the same two halves and tries the move again.
"""
if not _holds_a_secret(record):
return None
staged: Final = _stage_scrubbed_file(record)
if staged is None:
return record
if not isinstance(vault.write(_encode_secret(record)), SecretStored):
discard_staged_json(staged)
return record
if not _commit_token_file(staged) and not _overwrite_file_secret(record) and replacing is None:
vault.erase()
return record
def _scrub_file_secret(record: CliTokenRecord) -> bool:
"""Leave no secret material in the token file once the vault holds it"""
if not _holds_a_secret(record):
return True
staged: Final = _stage_scrubbed_file(record)
if staged is not None and _commit_token_file(staged):
return True
return _overwrite_file_secret(record)
def _overwrite_file_secret(record: CliTokenRecord) -> bool:
"""Take the secret out of the token file where it lies, when no replacement can be put in place.
The atomic rewrite wants room for a second file and a directory that will accept it. A full disk
refuses the first and a read-only `~/.litellm` the second, and neither stands in the way of
shortening the file that is already there. It is worth the loss of atomicity because a partial
write reads as no login at all, which is where the refused rewrite left the next run anyway.
"""
try:
overwrite_private_json(get_cli_token_file_path(), _without_secret(record).model_dump(exclude_none=True))
except OSError:
return False
return True
def _stage_scrubbed_file(record: CliTokenRecord) -> str | None:
staged: Final = _stage_token_file(_without_secret(record))
return None if isinstance(staged, CredentialNotSaved) else staged
def _stage_token_file(record: CliTokenRecord) -> str | CredentialNotSaved:
path: Final = Path(get_cli_token_file_path())
try:
ensure_private_dir(path.parent)
return stage_private_json(str(path), record.model_dump(exclude_none=True))
except OSError as error:
return CredentialNotSaved(str(error))
def _commit_token_file(staged: str) -> bool:
try:
commit_staged_json(staged, get_cli_token_file_path())
except OSError:
return False
return True
def _holds_a_secret(record: CliTokenRecord) -> bool:
"""Whether the record carries anything that would sign someone in as this user"""
return record.key is not None or bool(record.jwt_token) or record.refresh_token is not None
def _without_secret(record: CliTokenRecord) -> CliTokenRecord:
return record.model_copy(update=MappingProxyType({"key": None, "jwt_token": "", "refresh_token": None}))
def _encode_secret(record: CliTokenRecord) -> str:
return CliTokenSecret(
base_url=record.base_url,
key=record.key,
jwt_token=record.jwt_token,
refresh_token=record.refresh_token,
timestamp=record.timestamp,
).model_dump_json()
def _decode_secret(blob: str, base_url: str) -> CliTokenSecret | None:
"""The keychain entry, when it is one this metadata file may be paired with"""
try:
secret: Final = CliTokenSecret.model_validate_json(blob)
except ValidationError:
return None
return secret if secret.base_url == base_url else None
def _write_token_file(record: CliTokenRecord) -> None:
path: Final = Path(get_cli_token_file_path())
ensure_private_dir(path.parent)
write_private_json(str(path), record.model_dump(exclude_none=True))

View file

@ -2,6 +2,7 @@
Utility helpers for reading and parsing environment variables.
"""
import logging
import os
from typing import Final
@ -22,6 +23,26 @@ def get_env_int(env_var: str, default: int) -> int:
return default
def get_env_int_in_range(env_var: str, default: int, minimum: int, maximum: int) -> int:
"""Parse an environment variable as an integer constrained to ``[minimum, maximum]``.
Values outside the range fall back to the default and warn, so a misconfigured knob can
neither crash the caller nor silently change the meaning of what it computes.
"""
value: Final = get_env_int(env_var, default)
if minimum <= value <= maximum:
return value
logging.getLogger("LiteLLM").warning(
"%s=%s is outside the supported range [%s, %s]. Falling back to %s.",
env_var,
value,
minimum,
maximum,
default,
)
return default
def get_env_int_or_none(env_var: str) -> int | None:
"""Parse an environment variable as an integer, returning None when it is unset or unusable.

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