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

This commit is contained in:
yassin 2026-09-15 20:46:14 +00:00
commit 4ba136946c
422 changed files with 30527 additions and 4928 deletions

View file

@ -2915,6 +2915,25 @@ jobs:
exit 1
fi
provider_replay_harness:
docker:
- *python312_image
working_directory: ~/project
resource_class: medium
steps:
- setup_litellm_test_deps
- run:
name: Test provider replay harness
command: |
mkdir -p test-results/provider-replay-harness
uv run --no-sync pytest -q --noconftest -o addopts= -o pythonpath=tests/e2e -p no:rerunfailures \
--junitxml=test-results/provider-replay-harness/junit.xml \
tests/e2e/test_provider_edge.py tests/e2e/test_fixture_bundle.py \
tests/e2e/test_fixture_canonical.py tests/e2e/test_fixture_mode.py \
tests/code_coverage_tests/test_provider_replay_harness.py
- store_test_results:
path: test-results/provider-replay-harness
integration_contracts:
parameters:
suite:
@ -2967,6 +2986,7 @@ workflows:
only:
- main
- /litellm_.*/
- provider_replay_harness
- base_sdk_install:
filters: *main_branches
- local_testing_part1:

View file

@ -14,6 +14,17 @@ query-filters:
id: py/clear-text-logging-sensitive-data # CWE-312
- exclude:
id: py/polynomial-redos # CWE-730
# Import resolution confuses stdlib types with management_endpoints/types.py.
# The generic cycle query also reports intentional deferred imports.
- exclude:
id: py/cyclic-import
- exclude:
id: py/unsafe-cyclic-import
# Known false positives on live settings and Protocol placeholders.
- exclude:
id: py/unused-global-variable
- exclude:
id: py/ineffectual-statement
paths-ignore:
- tests

View file

@ -3,6 +3,9 @@ import xml.etree.ElementTree as ET
from pathlib import Path
from typing import Final
sys.path.insert(0, str(Path(__file__).resolve().parents[2] / "tests/e2e"))
from coverage_registry.management_cases import MANAGEMENT_CASES
def main() -> int:
selected: Final = tuple(sys.argv[2:])
@ -16,6 +19,17 @@ def main() -> int:
case.get("file") for case in cases if all(case.find(tag) is None for tag in ("skipped", "failure", "error"))
)
missing: Final = tuple(path for path in selected if path not in passed)
required_nodes: Final = frozenset(case.node for case in MANAGEMENT_CASES if case.node.split("::", 1)[0] in selected)
passed_nodes: Final = frozenset(
prop.get("value")
for case in cases
if all(case.find(tag) is None for tag in ("skipped", "failure", "error"))
for prop in case.findall("./properties/property")
if prop.get("name") == "management_node"
)
missing_nodes: Final = required_nodes - passed_nodes
for node in sorted(missing_nodes):
_ = sys.stdout.write(f"::error::required management case did not pass: {node}\n")
for path in selected:
collected: Final = sum(case.get("file") == path for case in cases)
skipped: Final = sum(case.get("file") == path and case.find("skipped") is not None for case in cases)
@ -27,6 +41,7 @@ def main() -> int:
if (
selected
and not missing
and not missing_nodes
and not any(case.find(tag) is not None for case in cases for tag in ("failure", "error"))
):
return 0

5
.github/e2e-stack/oidc-profile.sh vendored Executable file
View file

@ -0,0 +1,5 @@
#!/usr/bin/env bash
set -euo pipefail
REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)"
cd "${REPO_ROOT}"
exec uv run --no-sync python tests/e2e/idp.py "$@"

View file

@ -12,6 +12,8 @@ UNSUPPORTED: Final = re.compile(
HARNESS: Final = re.compile(
r"^tests/e2e/[A-Za-z0-9_.-]+\.(py|ini)$"
r"|^tests/e2e/idp_realm\.json$"
r"|^tests/e2e/management/(management_client|jwt_actors|conftest)\.py$"
r"|^tests/e2e/coverage_registry/management_cases\.py$"
r"|^tests/e2e/gateway/"
r"|^\.github/e2e-stack/"
r"|^\.github/workflows/test-e2e-changed\.yml$"

View file

@ -101,7 +101,8 @@ If you're seeing a delay in your PR being merged, ping the LiteLLM Team on [Slac
For bug fixes: Before shows the reproduction, After shows the same steps passing
For new features: Before shows the capability missing, After shows it working end-to-end
If the change applies to all three LLM endpoints (/v1/responses, /v1/chat/completions, /v1/messages), make each endpoint its own case, not just one
For UI changes: before/after screenshots under the same headings -->
For UI changes: before/after screenshots under the same headings
If the main use case runs through a coding tool like Claude Code or Codex, drive that tool interactively the way the user does (never `claude -p`, `codex exec`, or curl on its own) and embed before/after screenshots of its pane under the same headings; curl replays and headless runs can follow as extra cases, never as the only proof -->
## Type

View file

@ -57,6 +57,7 @@ permissions:
env:
UV_PYTHON: "3.12"
LITELLM_LOCAL_MODEL_COST_MAP: "True"
jobs:
run:
@ -113,6 +114,7 @@ jobs:
if: steps.changes.outputs.decision != 'skip'
timeout-minutes: 8
run: |
diff -u model_prices_and_context_window.json litellm/model_prices_and_context_window_backup.json
.github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra saml
uv run --no-sync python -c 'import os, sys; print(sys.version); assert f"{sys.version_info.major}.{sys.version_info.minor}" == os.environ["UV_PYTHON"]'

View file

@ -178,7 +178,7 @@ jobs:
version: "0.10.9"
- name: Install dependencies
run: uv sync --frozen --extra proxy --python 3.10
run: uv sync --frozen --extra proxy --extra cli --python 3.10
- run: uv run --no-sync python --version
@ -187,3 +187,6 @@ jobs:
- name: Check litellm CLI
run: uv run --no-sync litellm --version
- name: Check lite CLI
run: uv run --no-sync lite version

View file

@ -183,7 +183,7 @@ jobs:
log="${RUNNER_TEMP}/e2e-pass-${pass}.log"
echo "::group::pass ${pass} of 3"
set +e
uv run --no-sync pytest "${test_files[@]}" --rootdir=. -v -p no:cacheprovider \
uv run --no-sync pytest "${test_files[@]}" --rootdir=. -v --reruns 0 -p no:cacheprovider \
-o junit_family=xunit1 --junitxml="${report}" > "${log}" 2>&1
status=$?
uv run --no-sync python .github/e2e-stack/assert_tests_ran.py "${report}" "${test_files[@]}"

View file

@ -25,6 +25,8 @@ Same thing for bug fixes. The tests should make it so that this specific bug can
Never test structure of code only function of it
A test must only fail when litellm code changes. Never pin facts we don't own (a vendor's price, a third party's field, an upstream default, today's date) as literals or as "X must be absent"; assert the invariant our code guarantees instead, e.g. two rows agree, a value is within range, a field is derived from another. If an outside fact is truly load-bearing, cite its source and date next to the assertion so a reader can tell stale from broken
`tests/test_litellm/` mirrors `litellm/` in a parallel path (see `tests/test_litellm/readme.md`). Name tests `test_<filename>.py`, but always match the existing test file in the directory you touch — many provider dirs use longer descriptive names (e.g. `test_anthropic_chat_transformation.py`) to avoid ambiguity across sibling folders. For bug fixes, extend the existing mapped test file rather than creating a new one. Only create a new test file for a new feature (provider, endpoint, or transformation module) that has no mapped test yet, following that directory's naming convention (or `test_<filename>.py` if you're the first test there). One focused regression test beats many shallow ones
End-to-end tests belong in `tests/e2e/` and must follow the harness conventions documented in that directory's `CLAUDE.md`

View file

@ -299,6 +299,9 @@ test-rust-extension:
[ "$$#" -eq 1 ] && \
UV_PROJECT_ENVIRONMENT="$$temporary/venv" $(UV) sync --python 3.12 --frozen --no-install-project --all-groups --all-extras && \
$(UV) pip install --python "$$temporary/venv/bin/python" --no-deps "$$1" && \
"$$temporary/venv/bin/python" -I -m mypy.stubtest \
--mypy-config-file tests/test_litellm/rust_bridge/stubtest.ini \
litellm.rust_bridge._native && \
LITELLM_RUST=1 LITELLM_LOCAL_MODEL_COST_MAP=True \
"$$temporary/venv/bin/python" -I -m pytest --import-mode=importlib -m requires_rust_extension tests/test_litellm_rust

View file

@ -0,0 +1,14 @@
-- AlterTable
ALTER TABLE "LiteLLM_BudgetTable" ADD COLUMN IF NOT EXISTS "tpd_limit" BIGINT;
-- AlterTable
ALTER TABLE "LiteLLM_TeamTable" ADD COLUMN IF NOT EXISTS "tpd_limit" BIGINT;
-- AlterTable
ALTER TABLE "LiteLLM_DeletedTeamTable" ADD COLUMN IF NOT EXISTS "tpd_limit" BIGINT;
-- AlterTable
ALTER TABLE "LiteLLM_VerificationToken" ADD COLUMN IF NOT EXISTS "tpd_limit" BIGINT;
-- AlterTable
ALTER TABLE "LiteLLM_DeletedVerificationToken" ADD COLUMN IF NOT EXISTS "tpd_limit" BIGINT;

View file

@ -17,6 +17,7 @@ model LiteLLM_BudgetTable {
max_parallel_requests Int?
tpm_limit BigInt?
rpm_limit BigInt?
tpd_limit BigInt?
model_max_budget Json?
budget_duration String?
budget_reset_at DateTime?
@ -133,6 +134,7 @@ model LiteLLM_TeamTable {
max_parallel_requests Int?
tpm_limit BigInt?
rpm_limit BigInt?
tpd_limit BigInt?
budget_duration String?
budget_reset_at DateTime?
blocked Boolean @default(false)
@ -203,6 +205,7 @@ model LiteLLM_DeletedTeamTable {
max_parallel_requests Int?
tpm_limit BigInt?
rpm_limit BigInt?
tpd_limit BigInt?
budget_duration String?
budget_reset_at DateTime?
blocked Boolean @default(false)
@ -438,6 +441,7 @@ model LiteLLM_VerificationToken {
blocked Boolean?
tpm_limit BigInt?
rpm_limit BigInt?
tpd_limit BigInt?
max_budget Float?
budget_duration String?
budget_reset_at DateTime?
@ -534,6 +538,7 @@ model LiteLLM_DeletedVerificationToken {
blocked Boolean?
tpm_limit BigInt?
rpm_limit BigInt?
tpd_limit BigInt?
max_budget Float?
budget_duration String?
budget_reset_at DateTime?

View file

@ -538,7 +538,7 @@ context_window_fallbacks: Optional[List] = None
content_policy_fallbacks: Optional[List] = None
allowed_fails: int = 3
allow_dynamic_callback_disabling: bool = True
num_retries_per_request: Optional[int] = None # cap on Router retries of one model group; resets per fallback hop
num_retries_per_request: Optional[int] = None # for the request overall (incl. fallbacks + model retries)
####### SECRET MANAGERS #####################
secret_manager_client: Optional[Any] = (
None # list of instantiated key management clients - e.g. azure kv, infisical, etc.

View file

@ -21,6 +21,7 @@ from litellm._logging import verbose_logger, verbose_proxy_logger
from litellm.a2a_protocol.streaming_iterator import A2AStreamingIterator
from litellm.a2a_protocol.utils import A2ARequestUtils
from litellm.constants import DEFAULT_A2A_AGENT_TIMEOUT
from litellm.litellm_core_utils.asyncify import asyncify
from litellm.litellm_core_utils.litellm_logging import Logging
from litellm.llms.custom_httpx.http_handler import (
get_async_httpx_client,
@ -507,7 +508,7 @@ async def asend_message(
prompt_tokens,
completion_tokens,
_,
) = A2ARequestUtils.calculate_usage_from_request_response(
) = await asyncify(A2ARequestUtils.calculate_usage_from_request_response)(
request=request,
response_dict=response_dict,
)

View file

@ -11,6 +11,7 @@ import litellm
from litellm._logging import verbose_logger
from litellm.a2a_protocol.cost_calculator import A2ACostCalculator
from litellm.a2a_protocol.utils import A2ARequestUtils
from litellm.litellm_core_utils.asyncify import asyncify
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
if TYPE_CHECKING:
@ -99,11 +100,11 @@ class A2AStreamingIterator:
# Calculate tokens from collected text
input_message: Final = A2ARequestUtils.get_input_message_from_request(self.request)
input_text: Final = A2ARequestUtils.extract_text_from_message(input_message)
prompt_tokens: Final = A2ARequestUtils.count_tokens(input_text)
prompt_tokens: Final = await asyncify(A2ARequestUtils.count_tokens)(input_text)
# Use the last (most complete) text from chunks
output_text: Final = self.collected_text_parts[-1] if self.collected_text_parts else ""
completion_tokens: Final = A2ARequestUtils.count_tokens(output_text)
completion_tokens: Final = await asyncify(A2ARequestUtils.count_tokens)(output_text)
total_tokens: Final = prompt_tokens + completion_tokens

View file

@ -22,6 +22,7 @@
"mcp-servers-2025-12-04": null,
"oauth-2025-04-20": "oauth-2025-04-20",
"output-128k-2025-02-19": "output-128k-2025-02-19",
"per-turn-control-2026-07-01": "per-turn-control-2026-07-01",
"prompt-caching-scope-2026-01-05": "prompt-caching-scope-2026-01-05",
"skills-2025-10-02": "skills-2025-10-02",
"structured-outputs-2025-11-13": "structured-outputs-2025-11-13",
@ -52,6 +53,7 @@
"mcp-servers-2025-12-04": null,
"output-128k-2025-02-19": null,
"structured-output-2024-03-01": null,
"per-turn-control-2026-07-01": null,
"prompt-caching-scope-2026-01-05": "prompt-caching-scope-2026-01-05",
"skills-2025-10-02": "skills-2025-10-02",
"structured-outputs-2025-11-13": "structured-outputs-2025-11-13",
@ -82,6 +84,7 @@
"mcp-servers-2025-12-04": null,
"output-128k-2025-02-19": null,
"structured-output-2024-03-01": null,
"per-turn-control-2026-07-01": null,
"prompt-caching-scope-2026-01-05": null,
"skills-2025-10-02": null,
"structured-outputs-2025-11-13": "structured-outputs-2025-11-13",
@ -113,6 +116,7 @@
"mcp-servers-2025-12-04": null,
"output-128k-2025-02-19": null,
"structured-output-2024-03-01": null,
"per-turn-control-2026-07-01": null,
"prompt-caching-scope-2026-01-05": null,
"skills-2025-10-02": null,
"structured-outputs-2025-11-13": null,
@ -144,6 +148,7 @@
"mcp-servers-2025-12-04": null,
"output-128k-2025-02-19": null,
"structured-output-2024-03-01": null,
"per-turn-control-2026-07-01": null,
"prompt-caching-scope-2026-01-05": null,
"skills-2025-10-02": null,
"structured-outputs-2025-11-13": null,
@ -176,6 +181,7 @@
"mcp-servers-2025-12-04": null,
"oauth-2025-04-20": "oauth-2025-04-20",
"output-128k-2025-02-19": "output-128k-2025-02-19",
"per-turn-control-2026-07-01": null,
"prompt-caching-scope-2026-01-05": "prompt-caching-scope-2026-01-05",
"skills-2025-10-02": "skills-2025-10-02",
"structured-outputs-2025-11-13": "structured-outputs-2025-11-13",

View file

@ -9,6 +9,7 @@ import litellm
from litellm._logging import verbose_logger
from litellm.litellm_core_utils.get_litellm_params import AWS_CREDENTIAL_KWARGS_KEYS
from litellm.litellm_core_utils.llm_cost_calc.utils import parse_prompt_tokens_details
from litellm.llms.vertex_ai.batches.transformation import vertex_prompt_tokens_details
from litellm.types.llms.openai import Batch
from litellm.types.utils import ModelInfo, Usage
from litellm.utils import token_counter
@ -356,6 +357,7 @@ def calculate_vertex_ai_batch_cost_and_usage(
prompt_tokens=_prompt,
completion_tokens=_completion,
total_tokens=_total,
prompt_tokens_details=vertex_prompt_tokens_details(usage_metadata),
)
try:

View file

@ -0,0 +1,125 @@
"""Atomic affinity claims shared by deployment and tier-model selection."""
import json
from collections.abc import Mapping
from typing import (
Final,
cast, # noqa: TID251 # Redis script results are narrowed only to object, then validated
)
from pydantic import JsonValue, TypeAdapter, ValidationError
from litellm._logging import verbose_router_logger
from litellm.caching.dual_cache import DualCache
_PIN_JSON_ADAPTER: Final = TypeAdapter[JsonValue](JsonValue)
_CLAIM_PIN_SCRIPT: Final = """
local current = redis.call('GET', KEYS[1])
if current == false then
redis.call('SET', KEYS[1], ARGV[1], 'EX', ARGV[2])
return ARGV[1]
end
if ARGV[3] then
local decoded, stored = pcall(cjson.decode, current)
if decoded and type(stored) == 'table' then
for _, eligible in ipairs(cjson.decode(ARGV[3])) do
local matches = true
for key, value in pairs(eligible) do
if stored[key] ~= value then matches = false; break end
end
for key, _ in pairs(stored) do
if eligible[key] == nil then matches = false; break end
end
if matches then
redis.call('EXPIRE', KEYS[1], ARGV[2])
return current
end
end
end
redis.call('SET', KEYS[1], ARGV[1], 'EX', ARGV[2])
return ARGV[1]
end
if current == ARGV[1] then
redis.call('EXPIRE', KEYS[1], ARGV[2])
end
return current
"""
def set_local_affinity_pin(cache: DualCache, cache_key: str, value: object, ttl_seconds: int) -> None:
"""Replace the entry because InMemoryCache.set_cache preserves a live key's expiry."""
cache.in_memory_cache.delete_cache(cache_key)
cache.in_memory_cache.set_cache(cache_key, value, ttl=ttl_seconds)
def _legacy_pin_matches(stored: object, pin_value: Mapping[str, str]) -> bool:
if isinstance(stored, dict):
return all(stored.get(key) is not None and str(stored[key]) == value for key, value in pin_value.items())
return isinstance(stored, str) and len(pin_value) == 1 and stored in pin_value.values()
def claim_affinity_pin_in_memory(
cache: DualCache,
cache_key: str,
pin_value: Mapping[str, str],
ttl_seconds: int,
*,
eligible_values: tuple[Mapping[str, str], ...] | None = None,
) -> object:
"""No await between read and write, so same-loop claims agree during a Redis outage."""
existing: Final[object] = cache.in_memory_cache.get_cache(cache_key)
if existing is not None and eligible_values is None:
if _legacy_pin_matches(existing, pin_value):
set_local_affinity_pin(cache, cache_key, pin_value, ttl_seconds)
return existing
winner: Final = existing if existing is not None and existing in (eligible_values or ()) else pin_value
set_local_affinity_pin(cache, cache_key, winner, ttl_seconds)
return winner
def _decode_pin(value: str) -> object:
try:
return _PIN_JSON_ADAPTER.validate_json(value)
except ValidationError:
return value
async def claim_affinity_pin(
cache: DualCache,
cache_key: str,
pin_value: Mapping[str, str],
ttl_seconds: int,
*,
eligible_values: tuple[Mapping[str, str], ...] | None = None,
) -> object:
"""Return the authoritative first writer, replacing it only when it becomes ineligible.
Eligible claims refresh the returned winner. Legacy deployment claims only refresh
a matching candidate. Resolve Redis per call because the proxy attaches it lazily.
"""
redis_cache: Final = cache.redis_cache
if redis_cache is not None:
try:
claim_script: Final = redis_cache.async_register_script(_CLAIM_PIN_SCRIPT)
args: Final = (
json.dumps(dict(pin_value)), # mutable-ok: JSON serialization requires dict, not a generic Mapping
int(ttl_seconds),
*(
(json.dumps(tuple(dict(value) for value in eligible_values)),) # mutable-ok: JSON requires dict
if eligible_values is not None
else ()
),
)
raw: Final = cast( # cast-ok: Redis scripts return heterogeneous values; only object is asserted here
object, await claim_script(keys=(cache_key,), args=args)
)
decoded: Final = raw.decode("utf-8") if isinstance(raw, bytes) else raw
if not isinstance(decoded, str):
return pin_value
winner: Final = _decode_pin(decoded)
set_local_affinity_pin(cache, cache_key, winner, ttl_seconds)
return winner
except Exception as error: # noqa: BLE001 # Redis/Lua faults retain same-pod affinity through local claims
verbose_router_logger.debug("Affinity cache: Redis claim failed, using pod-local claim. error=%s", error)
return claim_affinity_pin_in_memory(cache, cache_key, pin_value, ttl_seconds, eligible_values=eligible_values)

View file

@ -21,6 +21,7 @@ from litellm.constants import (
QDRANT_VECTOR_SIZE,
SEMANTIC_CACHE_EMBEDDING_TIMEOUT_SECONDS,
)
from litellm.litellm_core_utils.asyncify import asyncify
from litellm.litellm_core_utils.prompt_templates.common_utils import (
get_str_from_messages,
)
@ -255,7 +256,7 @@ class QdrantSemanticCache(BaseCache):
llm_router = None
router: Final = resolve_embedding_router(self.embedding_model, llm_router, llm_model_list)
embedding_input: Final = self._embedding_input(prompt, router)
embedding_input: Final = await asyncify(self._embedding_input)(prompt, router)
embedding_call: Final = (
router.aembedding(
model=self.embedding_model,

View file

@ -19,6 +19,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.asyncify import asyncify
from litellm.litellm_core_utils.prompt_templates.common_utils import (
get_str_from_messages,
)
@ -522,7 +523,7 @@ class RedisSemanticCache(BaseCache):
llm_router = None
router: Final = resolve_embedding_router(self.embedding_model, llm_router, llm_model_list)
embedding_input: Final = self._embedding_input(prompt, router)
embedding_input: Final = await asyncify(self._embedding_input)(prompt, router)
embedding_call: Final = (
router.aembedding(
model=self.embedding_model,

View file

@ -205,21 +205,42 @@ def _extract_anthropic_tool_exchange_spans(
return spans, None
def _message_has_cache_control(message: Mapping[str, object]) -> bool:
if message.get("cache_control") is not None:
return True
content: Final = message.get("content")
if isinstance(content, list):
return any(isinstance(part, Mapping) and part.get("cache_control") is not None for part in content)
return False
def _cached_prefix_indices(messages: Sequence[Mapping[str, object]]) -> tuple[int, ...]:
last_breakpoint: Final = max(
(index for index, msg in enumerate(messages) if _message_has_cache_control(msg)),
default=-1,
)
return tuple(range(last_breakpoint + 1))
def get_protected_indices(messages: Sequence[Mapping[str, object]]) -> tuple[int, ...]:
"""
Return indices of messages that must never be compressed:
- All system messages
- The last user message
- The last assistant message
- Every message up to and including the last one carrying an Anthropic cache_control breakpoint
The last user message is what the model is being asked to act on right now,
so compressing it replaces the live instruction with a marker. Compression
guardrails share this policy; see the Headroom guardrail.
guardrails share this policy; see the Headroom guardrail. A cache_control
breakpoint pins the provider's prompt-cache prefix to the exact bytes of every
row up to it, so rewriting any row inside that prefix turns the next request's
cache read into a cache write.
"""
system_indices: Final = tuple(index for index, msg in enumerate(messages) if msg.get("role", "") == "system")
last_user: Final = tuple(index for index, msg in enumerate(messages) if msg.get("role", "") == "user")[-1:]
last_assistant = tuple(index for index, msg in enumerate(messages) if msg.get("role", "") == "assistant")[-1:]
return system_indices + last_user + last_assistant
assistant_indices: Final = tuple(index for index, msg in enumerate(messages) if msg.get("role", "") == "assistant")
return tuple(dict.fromkeys(system_indices + last_user + assistant_indices[-1:] + _cached_prefix_indices(messages)))
def _combine_scores(
@ -421,7 +442,7 @@ def compress(
combined_scores = bm25_scores
# Protected messages are never compressed
protected_indices: Final = get_protected_indices(normalized_messages)
protected_indices: Final = get_protected_indices(original_messages)
kept_indices: set[int] = set(protected_indices)
tool_exchange_spans: list[set[int]] = []

View file

@ -1566,6 +1566,8 @@ BASE_MCP_ROUTE: Final = "/mcp"
BATCH_STATUS_POLL_INTERVAL_SECONDS: Final = int(os.getenv("BATCH_STATUS_POLL_INTERVAL_SECONDS", 3600)) # 1 hour
BATCH_STATUS_POLL_MAX_ATTEMPTS: Final = int(os.getenv("BATCH_STATUS_POLL_MAX_ATTEMPTS", 24)) # for 24 hours
BATCH_TPD_WINDOW_SECONDS: Final = 86400
BATCH_TPD_DESCRIPTOR_SUFFIX: Final = "_tpd"
HEALTH_CHECK_TIMEOUT_SECONDS: Final = int(os.getenv("HEALTH_CHECK_TIMEOUT_SECONDS", 60)) # 60 seconds
_background_health_check_max_tokens_env: Final = os.getenv("BACKGROUND_HEALTH_CHECK_MAX_TOKENS")
@ -1976,6 +1978,8 @@ BROWSER_SECURITY_HEADERS: Final[frozenset[str]] = frozenset(
UNSAFE_PROXY_RESPONSE_HEADERS: Final[frozenset[str]] = HTTP_FRAMING_HEADERS | BROWSER_SECURITY_HEADERS
STRINGIFIED_NONE: Final[str] = "None"
# A retrieved response replays the usage of the call that created it, so pricing these
# read/management routes like inference bills the same tokens twice.
NON_INFERENCE_CALL_TYPES: Final[frozenset[str]] = frozenset(

View file

@ -2278,6 +2278,19 @@ def default_video_cost_calculator(
return 0.0
def _batch_rate(
model_info: ModelInfo,
key: Literal[
"input_cost_per_audio_token_batches",
"input_cost_per_image_token_batches",
"input_cost_per_video_token_batches",
],
fallback: float,
) -> float:
rate: Final = model_info.get(key)
return fallback if rate is None else rate
def batch_cost_calculator(
usage: Usage,
model: str,
@ -2337,7 +2350,29 @@ def batch_cost_calculator(
total_prompt_cost = 0.0
total_completion_cost = 0.0
if input_cost_per_token_batches is not None:
total_prompt_cost = usage.prompt_tokens * input_cost_per_token_batches
batch_details: Final = parse_prompt_tokens_details(usage)
audio_tokens, image_tokens, video_tokens = (
batch_details["audio_tokens"],
batch_details["image_tokens"],
batch_details["video_tokens"],
)
modality_rates: Final = (
_batch_rate(model_info, "input_cost_per_audio_token_batches", input_cost_per_token_batches),
_batch_rate(model_info, "input_cost_per_image_token_batches", input_cost_per_token_batches),
_batch_rate(model_info, "input_cost_per_video_token_batches", input_cost_per_token_batches),
)
total_prompt_cost = sum(
tokens * rate
for tokens, rate in zip(
(
max((usage.prompt_tokens or 0) - audio_tokens - image_tokens - video_tokens, 0),
audio_tokens,
image_tokens,
video_tokens,
),
(input_cost_per_token_batches, *modality_rates),
)
)
elif input_cost_per_token:
details: Final = parse_prompt_tokens_details(usage)
cache_read_tokens: Final = details["cache_hit_tokens"]

View file

@ -15,6 +15,7 @@ from typing_extensions import ReadOnly, TypedDict
from litellm._logging import verbose_logger
from litellm.compression import compress
from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.asyncify import asyncify
from litellm.types.integrations.compression_interception import (
CompressionInterceptionConfig,
CompressionSavingsMetadata,
@ -153,7 +154,7 @@ class CompressionInterceptionLogger(CustomLogger):
self._prune_expired_cache()
compressed: Final = compress(
compressed: Final = await asyncify(compress)(
messages=messages,
model=model,
call_type=CallTypes.anthropic_messages,

View file

@ -884,7 +884,9 @@ class CustomGuardrail(CustomLogger):
"""logging_only: run apply_guardrail on copies of the logged request/response and record the verdict."""
from litellm.llms import get_guardrail_translation_mapping
if not self.uses_apply_guardrail_interface() or self.use_native_lifecycle_hooks:
if not self.uses_apply_guardrail_interface():
return kwargs, result
if not self._event_hook_is_event_type(GuardrailEventHooks.logging_only):
return kwargs, result
try:
translation: Final = get_guardrail_translation_mapping(CallTypes(call_type))()
@ -901,8 +903,18 @@ class CustomGuardrail(CustomLogger):
for key, value in (litellm_params.get("metadata") or {}).items()
if key != "standard_logging_guardrail_information"
}
response: Final = (
kwargs.get("async_complete_streaming_response") or kwargs.get("complete_streaming_response") or result
)
from litellm.types.utils import ModelResponse
output_translation: Final = (
get_guardrail_translation_mapping(CallTypes.acompletion)()
if isinstance(response, ModelResponse)
else translation
)
try:
await self._scan_logged_call(kwargs, result, translation, scratch_metadata)
await self._scan_logged_call(kwargs, response, translation, output_translation, scratch_metadata)
except Exception as e:
verbose_logger.warning("Guardrail %s: logging_only scan raised: %s", self.guardrail_name, e)
recorded: Final = scratch_metadata.get("standard_logging_guardrail_information")
@ -919,8 +931,9 @@ class CustomGuardrail(CustomLogger):
async def _scan_logged_call(
self,
kwargs: dict, # mutable-ok: CustomLogger.async_logging_hook contract
result: object,
response: object | None,
translation: "BaseTranslation",
output_translation: "BaseTranslation",
scratch_metadata: dict, # mutable-ok: apply_guardrail records its verdict into request metadata
) -> None:
optional_params: Final = kwargs.get("optional_params") or {}
@ -934,8 +947,10 @@ class CustomGuardrail(CustomLogger):
"metadata": scratch_metadata,
}
await translation.process_input_messages(data=scratch_request, guardrail_to_apply=self)
await translation.process_output_response(
response=copy.deepcopy(result), guardrail_to_apply=self, request_data=scratch_request
if response is None:
return
await output_translation.process_output_response(
response=copy.deepcopy(response), guardrail_to_apply=self, request_data=scratch_request
)
def supports_scan_only_tool_results(self) -> bool:

View file

@ -822,46 +822,33 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
def truncate_standard_logging_payload_content(
self,
standard_logging_object: StandardLoggingPayload,
):
) -> StandardLoggingPayload:
"""
Truncate error strings and message content in logging payload
Return a copy of the logging payload with error_str, messages, and response truncated
Some loggers like DataDog/ GCS Bucket have a limit on the size of the payload. (1MB)
This function truncates the error string and the message content if they exceed a certain length.
Every callback of a request shares one standard logging object, so the payload passed in is left
untouched and the callbacks that run later (the prompt caching router check, spend logs) still see
the original fields.
"""
MAX_STR_LENGTH: Final = 10_000
max_str_length: Final = 10_000
candidates: Final = {
field: self._truncate_field(field_value=standard_logging_object.get(field), max_length=max_str_length)
for field in ("error_str", "messages", "response")
}
truncated_fields: Final = {field: text for field, text in candidates.items() if text is not None}
return {**standard_logging_object, **truncated_fields}
# Truncate fields that might exceed max length
fields_to_truncate: Final = ["error_str", "messages", "response"]
for field in fields_to_truncate:
self._truncate_field(
standard_logging_object=standard_logging_object,
field_name=field,
max_length=MAX_STR_LENGTH,
)
def _truncate_field(
self,
standard_logging_object: StandardLoggingPayload,
field_name: str,
max_length: int,
) -> None:
def _truncate_field(self, field_value: object, max_length: int) -> str | None:
"""
Helper function to truncate a field in the logging payload
Return the truncated text of a field that exceeds max_length, or None when the field fits
This converts the field to a string and then truncates it if it exceeds the max length.
Why convert to string ?
1. User was sending a poorly formatted list for `messages` field, we could not predict where they would send content
- Converting to string and then truncating the logged content catches this
2. We want to avoid modifying the original `messages`, `response`, and `error_str` in the logging payload since these are in kwargs and could be returned to the user
The field is measured as a string because users send poorly formatted lists for `messages`, so there is
no fixed place the content would be.
"""
field_value: Final[object] = standard_logging_object.get(field_name)
if field_value:
str_value: Final = str(field_value)
if len(str_value) > max_length:
standard_logging_object[field_name] = self._truncate_text(text=str_value, max_length=max_length)
text: Final = str(field_value or "")
return self._truncate_text(text=text, max_length=max_length) if len(text) > max_length else None
def _truncate_text(self, text: str, max_length: int) -> str:
"""Truncate text if it exceeds max_length"""

View file

@ -563,11 +563,10 @@ class DataDogLogger(
if standard_logging_object.get("status") == "failure":
status = DataDogStatus.ERROR
# Build the initial payload
self.truncate_standard_logging_payload_content(standard_logging_object)
truncated_payload: Final = self.truncate_standard_logging_payload_content(standard_logging_object)
dd_payload: Final = self._create_datadog_logging_payload_helper(
standard_logging_object=standard_logging_object,
standard_logging_object=truncated_payload,
status=status,
)
return dd_payload

View file

@ -2610,12 +2610,6 @@ class PrometheusLogger(CustomLogger):
StandardLoggingPayloadSetup,
)
if self._should_skip_metrics_for_invalid_key(
user_api_key_dict=user_api_key_dict,
exception=original_exception,
):
return
status_code: Final = self._extract_status_code(exception=original_exception)
try:
@ -2633,7 +2627,7 @@ class PrometheusLogger(CustomLogger):
end_user=user_api_key_dict.end_user_id,
user=user_api_key_dict.user_id,
user_email=user_api_key_dict.user_email,
hashed_api_key=user_api_key_dict.api_key,
hashed_api_key=None if status_code == 401 else user_api_key_dict.api_key,
api_key_alias=user_api_key_dict.key_alias,
team=user_api_key_dict.team_id,
team_alias=user_api_key_dict.team_alias,

View file

@ -309,8 +309,8 @@ def max_retries_per_request_hit(kwargs: Mapping[str, object], num_retries_per_re
metadata: Final = kwargs.get(get_metadata_variable_name_from_kwargs(kwargs))
if not isinstance(metadata, Mapping):
return False
attempted_retries: Final = metadata.get("attempted_retries")
return type(attempted_retries) is int and 0 < attempted_retries and num_retries_per_request <= attempted_retries
retry_count: Final = metadata.get("request_retry_count")
return type(retry_count) is int and 0 < retry_count and num_retries_per_request <= retry_count
def get_or_create_metadata_bucket(

View file

@ -1,6 +1,9 @@
import inspect
import json
import re
import traceback
from collections.abc import Mapping
from types import MappingProxyType
from typing import Any, Final, Protocol, cast
import httpx
@ -202,11 +205,17 @@ def _get_response_headers(original_exception: Exception) -> httpx.Headers | None
return _response_headers
def _accepted_init_kwargs(exception_class: type[Exception], candidates: Mapping[str, object]) -> Mapping[str, object]:
accepted: Final = inspect.signature(exception_class).parameters
return MappingProxyType({name: value for name, value in candidates.items() if name in accepted})
def extract_and_raise_litellm_exception(
response: Any | None,
error_str: str,
model: str,
custom_llm_provider: str,
body: object | None = None,
):
"""
Covers scenario where litellm sdk calling proxy.
@ -216,32 +225,19 @@ def extract_and_raise_litellm_exception(
Relevant Issue: https://github.com/BerriAI/litellm/issues/7259
"""
pattern: Final = r"litellm\.\w+Error"
# Search for the exception in the error string
match: Final = re.search(pattern, error_str)
# Extract the exception if found
if match:
exception_name = match.group(0)
exception_name = exception_name.strip().replace("litellm.", "")
raised_exception_obj: Final = getattr(litellm, exception_name, None)
if raised_exception_obj:
# Try with response parameter first, fall back to without it
# Some exceptions (e.g., APIConnectionError) don't accept response param
try:
raise raised_exception_obj(
message=error_str,
llm_provider=custom_llm_provider,
model=model,
response=response,
)
except TypeError:
# Exception doesn't accept response parameter
raise raised_exception_obj(
message=error_str,
llm_provider=custom_llm_provider,
model=model,
)
if match is None:
return
exception_name: Final = match.group(0).removeprefix("litellm.")
raised_exception_obj: Final = getattr(litellm, exception_name, None)
if not raised_exception_obj:
return
raise raised_exception_obj(
message=error_str,
llm_provider=custom_llm_provider,
model=model,
**_accepted_init_kwargs(raised_exception_obj, MappingProxyType({"response": response, "body": body})),
)
class _ProviderHTTPException(Protocol):
@ -254,6 +250,23 @@ class _ProviderHTTPException(Protocol):
llm_provider: str
def _litellm_proxy_response(
original_exception: _ProviderHTTPException, custom_llm_provider: str
) -> httpx.Response | None:
response: Final = getattr(original_exception, "response", None)
if custom_llm_provider != "litellm_proxy" or not isinstance(response, httpx.Response) or response.headers:
return response
headers: Final = getattr(original_exception, "headers", None)
if not isinstance(headers, Mapping) or not headers:
return response
pairs: Final = headers.multi_items() if isinstance(headers, httpx.Headers) else headers.items()
return httpx.Response(
status_code=response.status_code,
headers=[(str(k), str(v)) for k, v in pairs],
request=getattr(original_exception, "request", None),
)
def _map_openai_exception(
*,
model: str,
@ -264,6 +277,7 @@ def _map_openai_exception(
exception_provider: str,
extra_information: str,
) -> None:
response: Final = _litellm_proxy_response(original_exception, custom_llm_provider)
# custom_llm_provider is openai, make it OpenAI
message = get_error_message(error_obj=original_exception)
if message is None:
@ -292,14 +306,14 @@ def _map_openai_exception(
message=f"RateLimitError: {exception_provider} - {message}",
model=model,
llm_provider=custom_llm_provider,
response=getattr(original_exception, "response", None),
response=response,
)
elif ExceptionCheckers.is_error_str_context_window_exceeded(error_str):
raise ContextWindowExceededError(
message=f"ContextWindowExceededError: {exception_provider} - {message}",
llm_provider=custom_llm_provider,
model=model,
response=getattr(original_exception, "response", None),
response=response,
litellm_debug_info=extra_information,
)
elif "invalid_request_error" in error_str and "model_not_found" in error_str:
@ -307,7 +321,7 @@ def _map_openai_exception(
message=f"{exception_provider} - {message}",
llm_provider=custom_llm_provider,
model=model,
response=getattr(original_exception, "response", None),
response=response,
litellm_debug_info=extra_information,
)
elif "A timeout occurred" in error_str:
@ -326,8 +340,9 @@ def _map_openai_exception(
message=f"ContentPolicyViolationError: {exception_provider} - {message}",
llm_provider=custom_llm_provider,
model=model,
response=getattr(original_exception, "response", None),
response=response,
litellm_debug_info=extra_information,
body=getattr(original_exception, "body", None),
)
elif "invalid_encrypted_content" in error_str or "could not be verified" in error_str:
helpful_message: Final = (
@ -345,7 +360,7 @@ def _map_openai_exception(
message=helpful_message,
llm_provider=custom_llm_provider,
model=model,
response=getattr(original_exception, "response", None),
response=response,
litellm_debug_info=extra_information,
body=getattr(original_exception, "body", None),
)
@ -354,7 +369,7 @@ def _map_openai_exception(
message=f"{exception_provider} - {message}",
llm_provider=custom_llm_provider,
model=model,
response=getattr(original_exception, "response", None),
response=response,
litellm_debug_info=extra_information,
body=getattr(original_exception, "body", None),
)
@ -372,7 +387,7 @@ def _map_openai_exception(
message=f"RateLimitError: {exception_provider} - {message}",
model=model,
llm_provider=custom_llm_provider,
response=getattr(original_exception, "response", None),
response=response,
litellm_debug_info=extra_information,
)
elif (
@ -383,7 +398,7 @@ def _map_openai_exception(
message=f"AuthenticationError: {exception_provider} - {message}",
llm_provider=custom_llm_provider,
model=model,
response=getattr(original_exception, "response", None),
response=response,
litellm_debug_info=extra_information,
)
elif "Mistral API raised a streaming error" in error_str:
@ -402,15 +417,16 @@ def _map_openai_exception(
message=f"{exception_provider} - {message}",
llm_provider=custom_llm_provider,
model=model,
response=getattr(original_exception, "response", None),
response=response,
litellm_debug_info=extra_information,
body=getattr(original_exception, "body", None),
)
elif original_exception.status_code == 401:
raise AuthenticationError(
message=f"AuthenticationError: {exception_provider} - {message}",
llm_provider=custom_llm_provider,
model=model,
response=getattr(original_exception, "response", None),
response=response,
litellm_debug_info=extra_information,
)
elif original_exception.status_code == 404:
@ -418,7 +434,7 @@ def _map_openai_exception(
message=f"NotFoundError: {exception_provider} - {message}",
model=model,
llm_provider=custom_llm_provider,
response=getattr(original_exception, "response", None),
response=response,
litellm_debug_info=extra_information,
)
elif original_exception.status_code == 408:
@ -433,7 +449,7 @@ def _map_openai_exception(
message=f"{exception_provider} - {message}",
model=model,
llm_provider=custom_llm_provider,
response=getattr(original_exception, "response", None),
response=response,
litellm_debug_info=extra_information,
body=getattr(original_exception, "body", None),
)
@ -442,7 +458,7 @@ def _map_openai_exception(
message=f"RateLimitError: {exception_provider} - {message}",
model=model,
llm_provider=custom_llm_provider,
response=getattr(original_exception, "response", None),
response=response,
litellm_debug_info=extra_information,
)
elif original_exception.status_code == 500:
@ -450,7 +466,7 @@ def _map_openai_exception(
message=f"InternalServerError: {exception_provider} - {message}",
model=model,
llm_provider=custom_llm_provider,
response=getattr(original_exception, "response", None),
response=response,
litellm_debug_info=extra_information,
)
elif original_exception.status_code == 502:
@ -458,7 +474,7 @@ def _map_openai_exception(
message=f"BadGatewayError: {exception_provider} - {message}",
model=model,
llm_provider=custom_llm_provider,
response=getattr(original_exception, "response", None),
response=response,
litellm_debug_info=extra_information,
)
elif original_exception.status_code == 503:
@ -466,7 +482,7 @@ def _map_openai_exception(
message=f"ServiceUnavailableError: {exception_provider} - {message}",
model=model,
llm_provider=custom_llm_provider,
response=getattr(original_exception, "response", None),
response=response,
litellm_debug_info=extra_information,
)
elif original_exception.status_code == 504: # gateway timeout error
@ -2423,10 +2439,11 @@ def exception_type(
custom_llm_provider == "litellm_proxy"
): # handle special case where calling litellm proxy + exception str contains error message
extract_and_raise_litellm_exception(
response=getattr(original_exception, "response", None),
response=_litellm_proxy_response(mappable_exception, custom_llm_provider),
error_str=error_str,
model=model,
custom_llm_provider=custom_llm_provider,
body=getattr(original_exception, "body", None),
)
if (
custom_llm_provider == "openai"

View file

@ -644,6 +644,24 @@ class Logging(LiteLLMLoggingBaseClass):
"""Keep ``_response_ms`` / ``litellm_overhead_time_ms`` for a result that has no ``_hidden_params``."""
self.response_timing_metrics = dict(timing_metrics) # mutable-ok: kept deep-copyable
def add_dynamic_callback(self, callback: CustomLogger) -> None:
self.dynamic_input_callbacks = self._with_dynamic_callback(self.dynamic_input_callbacks, callback)
self.dynamic_success_callbacks = self._with_dynamic_callback(self.dynamic_success_callbacks, callback)
self.dynamic_async_success_callbacks = self._with_dynamic_callback(
self.dynamic_async_success_callbacks, callback
)
self.dynamic_failure_callbacks = self._with_dynamic_callback(self.dynamic_failure_callbacks, callback)
self.dynamic_async_failure_callbacks = self._with_dynamic_callback(
self.dynamic_async_failure_callbacks, callback
)
@staticmethod
def _with_dynamic_callback(
callbacks: Sequence[str | Callable | CustomLogger] | None, callback: CustomLogger
) -> list[str | Callable | CustomLogger]:
existing: Final = tuple(callbacks or ())
return [*existing, *(() if callback in existing else (callback,))]
def process_dynamic_callbacks(self):
"""
Initializes CustomLogger compatible callbacks in self.dynamic_* callbacks
@ -1973,6 +1991,12 @@ class Logging(LiteLLMLoggingBaseClass):
self.model_call_details["combined_usage_object"] = usage
self.model_call_details["response_cost"] = response_cost
def record_assembled_response_for_failure(self, assembled: ModelResponse) -> None:
"""Bill a fully streamed response on the failure log when a post-call hook rejects it."""
usage: Final = getattr(assembled, "usage", None)
if isinstance(usage, Usage):
self.record_partial_usage_for_failure(usage, self._response_cost_calculator(result=assembled) or 0.0)
async def dispatch_failure_handlers(
self,
exception: Exception,

View file

@ -484,11 +484,12 @@ def apply_off_peak_pricing(model_info: ModelInfo, current_time: datetime | None,
def _apply_off_peak_to_base_costs(
model_info: ModelInfo,
current_time: datetime | None,
base_costs: tuple[float, float, float, float, float],
base_costs: tuple[float, float, float, float | None, float],
) -> tuple[float, float, float, float, float]:
"""Apply off-peak rates to an already-resolved set of base costs, whichever pricing path
produced them. The one-hour cache-creation rate passes through untouched, since
off_peak_pricing has no field for it, and reasoning is left to _resolve_billed_reasoning_rate.
produced them. off_peak_pricing has no field for the one-hour cache-creation rate, so a
present one passes through untouched and an absent one resolves to the applied
cache-creation rate. Reasoning is left to _resolve_billed_reasoning_rate.
"""
prompt, completion, cache_creation, cache_creation_above_1hr, cache_read = base_costs
rates: Final = apply_off_peak_pricing(
@ -506,7 +507,7 @@ def _apply_off_peak_to_base_costs(
rates.input_rate,
rates.output_rate,
rates.cache_creation_rate,
cache_creation_above_1hr,
rates.cache_creation_rate if cache_creation_above_1hr is None else cache_creation_above_1hr,
rates.cache_read_rate,
)
@ -532,6 +533,11 @@ def _get_token_base_cost(
`missing_cache_read_uses_input` resolves an absent cache-read rate to the resolved
input rate instead of 0.0; an explicit 0.0 rate stays a real price either way.
An absent cache-creation rate always resolves to the resolved input rate, the way the
tiered table and custom deployment pricing already do, since a provider that publishes
no write price bills cache writes as ordinary input. An absent 1h write rate resolves
to the cache-creation rate, off-peak included. An explicit 0.0 stays a real price for both.
Returns:
Tuple[float, float, float, float] - (prompt_cost, completion_cost, cache_creation_cost, cache_read_cost)
"""
@ -554,10 +560,9 @@ def _get_token_base_cost(
output_image_cost: Final = _get_cost_per_unit(model_info, "output_cost_per_image_token", None)
if output_image_cost is not None:
completion_base_cost = cast(float, output_image_cost)
cache_creation_cost = cast(float, _get_cost_per_unit(model_info, cache_creation_cost_key))
cache_creation_cost_above_1hr = cast(
float,
_get_cost_per_unit(model_info, "cache_creation_input_token_cost_above_1hr"),
cache_creation_cost = _get_cost_per_unit(model_info, cache_creation_cost_key, default_value=None)
cache_creation_cost_above_1hr = _get_cost_per_unit(
model_info, "cache_creation_input_token_cost_above_1hr", default_value=None
)
cache_read_cost = _get_cost_per_unit(model_info, cache_read_cost_key, default_value=None)
@ -639,22 +644,10 @@ def _get_token_base_cost(
else f"cache_read_input_token_cost_above_{threshold_str}_tokens"
)
cache_creation_cost = cast(
float,
_get_cost_per_unit(
model_info,
cache_creation_tiered_key,
cache_creation_cost,
),
)
cache_creation_cost = _get_cost_per_unit(model_info, cache_creation_tiered_key, cache_creation_cost)
cache_creation_cost_above_1hr = cast(
float,
_get_cost_per_unit(
model_info,
cache_creation_1hr_tiered_key,
cache_creation_cost_above_1hr,
),
cache_creation_cost_above_1hr = _get_cost_per_unit(
model_info, cache_creation_1hr_tiered_key, cache_creation_cost_above_1hr
)
cache_read_cost = _get_cost_per_unit(model_info, cache_read_tiered_key, cache_read_cost)
@ -665,16 +658,16 @@ def _get_token_base_cost(
except Exception:
continue
input_rate_for_missing_cache_rates: Final = _off_peak_rate(
_open_off_peak_block(model_info, current_time) or MappingProxyType({}),
"input_cost_per_token",
prompt_base_cost,
)
if cache_read_cost is None:
cache_read_cost = (
_off_peak_rate(
_open_off_peak_block(model_info, current_time) or MappingProxyType({}),
"input_cost_per_token",
prompt_base_cost,
)
if missing_cache_read_uses_input
else 0.0
)
cache_read_cost = input_rate_for_missing_cache_rates if missing_cache_read_uses_input else 0.0
resolved_cache_creation_cost: Final = (
input_rate_for_missing_cache_rates if cache_creation_cost is None else cache_creation_cost
)
return _apply_off_peak_to_base_costs(
model_info,
@ -682,7 +675,7 @@ def _get_token_base_cost(
(
prompt_base_cost,
completion_base_cost,
cache_creation_cost,
resolved_cache_creation_cost,
cache_creation_cost_above_1hr,
cache_read_cost,
),
@ -956,12 +949,16 @@ def _calculate_input_cost(
)
### AUDIO COST
if prompt_tokens_details["audio_tokens"]:
if prompt_tokens_details["audio_tokens"] and not (
prompt_tokens_details["audio_length_seconds"] and model_info.get("input_cost_per_audio_per_second") is not None
):
audio_cost_key: Final = _get_service_tier_cost_key("input_cost_per_audio_token", service_tier)
prompt_cost += calculate_cost_component(model_info, audio_cost_key, prompt_tokens_details["audio_tokens"])
### IMAGE TOKEN COST
if prompt_tokens_details["image_tokens"]:
if prompt_tokens_details["image_tokens"] and not (
prompt_tokens_details["image_count"] and model_info.get("input_cost_per_image") is not None
):
# For image token costs:
# First check if input_cost_per_image_token is available. If not, default to generic input_cost_per_token.
image_token_cost_key = "input_cost_per_image_token"
@ -970,7 +967,9 @@ def _calculate_input_cost(
prompt_cost += calculate_cost_component(model_info, image_token_cost_key, prompt_tokens_details["image_tokens"])
### VIDEO TOKEN COST
if prompt_tokens_details["video_tokens"]:
if prompt_tokens_details["video_tokens"] and not (
prompt_tokens_details["video_length_seconds"] and model_info.get("input_cost_per_video_per_second") is not None
):
video_token_cost_key = "input_cost_per_video_token"
if model_info.get(video_token_cost_key) is None:
video_token_cost_key = "input_cost_per_token"

View file

@ -19,6 +19,7 @@ from typing_extensions import NotRequired, TypedDict
import litellm
from litellm import verbose_logger
from litellm._uuid import uuid
from litellm.litellm_core_utils.asyncify import asyncify
from litellm.litellm_core_utils.model_response_utils import (
is_model_response_stream_empty,
)
@ -2247,7 +2248,7 @@ class CustomStreamWrapper:
if self.sent_last_chunk is True:
# log the final chunk with accurate streaming values
try:
complete_streaming_response = litellm.stream_chunk_builder(
complete_streaming_response = await asyncify(litellm.stream_chunk_builder)(
chunks=self.chunks,
messages=self.messages,
logging_obj=self.logging_obj,

View file

@ -20,6 +20,7 @@ from itertools import chain, repeat
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, Protocol, cast, overload, runtime_checkable
from pydantic import TypeAdapter, ValidationError
from typing_extensions import ReadOnly, TypedDict, assert_never
from litellm._logging import verbose_proxy_logger
@ -44,6 +45,7 @@ from litellm.llms.base_llm.guardrail_translation.utils import (
scoped_structured_message_indices,
stream_item_field,
stream_item_fingerprint,
unappliable_request_rewrite,
)
from litellm.proxy.pass_through_endpoints.llm_provider_handlers.anthropic_passthrough_logging_handler import (
AnthropicPassthroughLoggingHandler,
@ -103,9 +105,24 @@ class ToolResultBlockTextTarget:
block_idx: int
InputWriteBackTarget = (
MessageContentTarget | ContentBlockTextTarget | ToolResultStringTarget | ToolResultBlockTextTarget
)
@dataclass(frozen=True, slots=True)
class SystemStringTarget:
pass
@dataclass(frozen=True, slots=True)
class SystemBlockTextTarget:
block_idx: int
@dataclass(frozen=True, slots=True)
class ToolUseInputTarget:
msg_idx: int
content_idx: int
MessageTextTarget = MessageContentTarget | ContentBlockTextTarget | ToolResultStringTarget | ToolResultBlockTextTarget
InputWriteBackTarget = SystemStringTarget | SystemBlockTextTarget | MessageTextTarget
def _as_str_mapping(value: Mapping[str, object]) -> Mapping[str, object]:
@ -146,10 +163,17 @@ class ScannedText:
target: InputWriteBackTarget
@dataclass(frozen=True, slots=True)
class ScannedToolCall:
tool_call: ChatCompletionToolCallChunk
target: ToolUseInputTarget
@dataclass(frozen=True, slots=True)
class ExtractedInput:
scanned: tuple[ScannedText, ...]
images: tuple[str, ...]
tool_calls: tuple[ScannedToolCall, ...] = ()
EMPTY_EXTRACTED_INPUT: Final = ExtractedInput(scanned=(), images=())
@ -161,6 +185,74 @@ class _ToolCallShape:
arguments: str
def _is_client_tool_use(block: Mapping[str, object]) -> bool:
return (
block.get("type") == "tool_use"
and isinstance(block.get("id"), str)
and isinstance(block.get("name"), str)
and isinstance(block.get("input"), dict)
)
def _write_back_system_block(system: object, block_idx: int, response: str) -> None:
if not isinstance(system, list):
return
text_blocks: Final = tuple(block for block in system if isinstance(block, dict) and block.get("type") == "text")
if block_idx < len(text_blocks):
text_blocks[block_idx]["text"] = (
response # mutable-ok: guardrails rewrite the caller's request payload in place
)
def _write_back_message_text(message: _WritableMessage, target: MessageTextTarget, response: str) -> None:
content: Final = message.get("content", None)
if content is None:
return
match target:
case MessageContentTarget():
if isinstance(content, str):
message["content"] = response # mutable-ok: guardrails rewrite the caller's request payload in place
case ContentBlockTextTarget(content_idx=content_idx):
if isinstance(content, list):
content[content_idx]["text"] = (
response # mutable-ok: guardrails rewrite the caller's request payload in place
)
case ToolResultStringTarget(content_idx=content_idx):
if isinstance(content, list):
content[content_idx]["content"] = (
response # mutable-ok: guardrails rewrite the caller's request payload in place
)
case ToolResultBlockTextTarget(content_idx=content_idx, block_idx=block_idx):
if isinstance(content, list):
content[content_idx]["content"][block_idx]["text"] = (
response # mutable-ok: guardrails rewrite the caller's request payload in place
)
case _:
assert_never(target)
_TOOL_USE_INPUT_ADAPTER: Final = TypeAdapter(dict[str, object])
def _rewritten_tool_use_input(arguments: str) -> Mapping[str, object] | None:
try:
return _TOOL_USE_INPUT_ADAPTER.validate_json(arguments)
except ValidationError:
return None
def _write_back_tool_use(
message: _WritableMessage, target: ToolUseInputTarget, shape: _ToolCallShape, rewritten_input: Mapping[str, object]
) -> None:
content: Final = message.get("content", None)
block: Final = content[target.content_idx] if isinstance(content, list) else None
if not isinstance(block, dict):
return
block["input"] = rewritten_input # mutable-ok: guardrails rewrite the caller's request payload in place
if shape.name is not None and shape.name != block.get("name"):
block["name"] = shape.name # mutable-ok: guardrails rewrite the caller's request payload in place
@dataclass(frozen=True, slots=True)
class _SSEFieldRewrite:
"""One field of one nested section of a buffered SSE event, rewritten."""
@ -452,9 +544,8 @@ class AnthropicMessagesHandler(BaseTranslation):
skip_tool: Final = effective_skip_tool_message_for_guardrail(guardrail_to_apply)
scan_only_tool_results: Final = effective_scan_only_tool_results_for_guardrail(guardrail_to_apply)
# Exclude only the trusted top-level prompt. In-sequence system entries are untrusted
# and must stay aligned with texts_to_check for positional masking. When the top-level
# prompt is included, the pre-existing count mismatch disables positional masking.
# The top-level prompt is translated on its own below so it can be hoisted in front of
# any mid-turn system entries and scanned first, aligned with that structured position.
translation_source: Final = { # mutable-ok: API message payload
key: value for key, value in data.items() if key != "system"
}
@ -490,7 +581,12 @@ class AnthropicMessagesHandler(BaseTranslation):
]
)
# Step 1: Extract all text content and images
# Step 1: Extract all text content, images, and tool calls
top_level_system_scanned: Final = (
()
if hoisted_system_message is None or scan_only_tool_results
else self._extract_top_level_system_text(hoisted_system_message)
)
extracted: Final = tuple(
self._extract_input_text_and_images(
message=message,
@ -501,17 +597,27 @@ class AnthropicMessagesHandler(BaseTranslation):
)
for msg_idx, message in enumerate(messages)
)
scanned: Final = tuple(item for one_message in extracted for item in one_message.scanned)
scanned: Final = (
*top_level_system_scanned,
*(item for one_message in extracted for item in one_message.scanned),
)
texts_to_check: Final = [item.text for item in scanned] # mutable-ok: GenericGuardrailAPIInputs takes list[str]
images_to_check: Final = [
image for one_message in extracted for image in one_message.images
] # mutable-ok: GenericGuardrailAPIInputs takes list[str]
scanned_tool_calls: Final = tuple(item for one_message in extracted for item in one_message.tool_calls)
tool_calls_to_check: Final = [
item.tool_call for item in scanned_tool_calls
] # mutable-ok: GenericGuardrailAPIInputs takes list[ChatCompletionToolCallChunk]
pre_guardrail_tool_calls: Final = _tool_call_shapes(tool_calls_to_check)
# Step 2: Apply guardrail to all texts in batch
if texts_to_check:
# Step 2: Apply guardrail to all texts and tool calls in batch
if texts_to_check or tool_calls_to_check:
inputs: Final = GenericGuardrailAPIInputs(texts=texts_to_check)
if images_to_check:
inputs["images"] = images_to_check
if tool_calls_to_check:
inputs["tool_calls"] = tool_calls_to_check
if tools_to_check:
inputs["tools"] = tools_to_check
original_structured_messages: Final = structured_messages
@ -570,9 +676,18 @@ class AnthropicMessagesHandler(BaseTranslation):
preserve_system_messages=has_midturn_system_message,
)
else:
if guardrailed_texts and len(guardrailed_texts) != len(scanned):
raise unappliable_request_rewrite(guardrail_to_apply.guardrail_name)
self._apply_guardrail_tool_calls_to_input(
messages=messages,
scanned_tool_calls=scanned_tool_calls,
pre_guardrail_tool_calls=pre_guardrail_tool_calls,
returned_tool_calls=guardrailed_inputs.get("tool_calls"),
guardrail_name=guardrail_to_apply.guardrail_name,
)
# Step 3: Map guardrail responses back to original message structure
await self._apply_guardrail_responses_to_input(
messages=messages,
data=data,
responses=guardrailed_texts,
scanned=scanned,
)
@ -598,6 +713,19 @@ class AnthropicMessagesHandler(BaseTranslation):
hoisted: Final = probe.get("messages") or [] # mutable-ok: API message payload
return hoisted[0] if hoisted else None
@staticmethod
def _extract_top_level_system_text(hoisted_system_message: AllMessageValues) -> tuple[ScannedText, ...]:
content: Final = hoisted_system_message.get("content")
if isinstance(content, str):
return (ScannedText(content, SystemStringTarget()),)
if not isinstance(content, list):
return ()
return tuple(
ScannedText(text_str, SystemBlockTextTarget(block_idx))
for block_idx, block in enumerate(content)
if isinstance(block, dict) and isinstance(text_str := block.get("text"), str)
)
@staticmethod
def _openai_system_message_to_anthropic(
message: Mapping[str, object],
@ -852,9 +980,25 @@ class AnthropicMessagesHandler(BaseTranslation):
for content_idx, content_item in enumerate(content)
if isinstance(content_item, dict)
)
tool_use_blocks: Final = (
()
if scan_only_tool_results
else tuple(
(content_idx, content_item)
for content_idx, content_item in enumerate(content)
if isinstance(content_item, dict) and _is_client_tool_use(content_item)
)
)
return ExtractedInput(
scanned=tuple(item for block in blocks for item in block.scanned),
images=tuple(image for block in blocks for image in block.images),
tool_calls=tuple(
ScannedToolCall(
tool_call=AnthropicConfig.convert_tool_use_to_openai_format(content_item, tool_call_idx),
target=ToolUseInputTarget(msg_idx, content_idx),
)
for tool_call_idx, (content_idx, content_item) in enumerate(tool_use_blocks)
),
)
@classmethod
@ -940,43 +1084,59 @@ class AnthropicMessagesHandler(BaseTranslation):
async def _apply_guardrail_responses_to_input(
self,
messages: Sequence[_WritableMessage],
responses: list[str],
data: dict[str, object], # mutable-ok: API message payload
responses: Sequence[str],
scanned: tuple[ScannedText, ...],
) -> None:
"""
Apply guardrail responses back to input messages.
Apply guardrail responses back to the top-level system prompt and the input messages.
"""
raw_messages: Final = data.get("messages")
messages: Final[Sequence[_WritableMessage]] = raw_messages if isinstance(raw_messages, list) else ()
for item, guardrail_response in zip(scanned, responses):
target = item.target
message = messages[target.msg_idx]
content = message.get("content", None)
if content is None:
continue
match target:
case MessageContentTarget():
if isinstance(content, str):
message["content"] = (
guardrail_response # mutable-ok: guardrails rewrite the caller's request payload in place
)
case ContentBlockTextTarget(content_idx=content_idx):
if isinstance(content, list):
content[content_idx]["text"] = (
guardrail_response # mutable-ok: guardrails rewrite the caller's request payload in place
)
case ToolResultStringTarget(content_idx=content_idx):
if isinstance(content, list):
content[content_idx]["content"] = (
guardrail_response # mutable-ok: guardrails rewrite the caller's request payload in place
)
case ToolResultBlockTextTarget(content_idx=content_idx, block_idx=block_idx):
if isinstance(content, list):
content[content_idx]["content"][block_idx]["text"] = (
match item.target:
case SystemStringTarget():
if isinstance(data.get("system"), str):
data["system"] = (
guardrail_response # mutable-ok: guardrails rewrite the caller's request payload in place
)
case SystemBlockTextTarget(block_idx=block_idx):
_write_back_system_block(data.get("system"), block_idx, guardrail_response)
case (
MessageContentTarget()
| ContentBlockTextTarget()
| ToolResultStringTarget()
| ToolResultBlockTextTarget() as message_target
):
_write_back_message_text(messages[message_target.msg_idx], message_target, guardrail_response)
case _:
assert_never(target)
assert_never(item.target)
@staticmethod
def _apply_guardrail_tool_calls_to_input(
messages: Sequence[_WritableMessage],
scanned_tool_calls: tuple[ScannedToolCall, ...],
pre_guardrail_tool_calls: tuple[_ToolCallShape, ...],
returned_tool_calls: Sequence[object] | None,
guardrail_name: str | None,
) -> None:
post_guardrail_tool_calls: Final = _tool_call_shapes(
returned_tool_calls
if returned_tool_calls is not None and len(returned_tool_calls) == len(pre_guardrail_tool_calls)
else tuple(item.tool_call for item in scanned_tool_calls)
)
rewritten: Final = tuple(
(item, after, _rewritten_tool_use_input(after.arguments))
for item, before, after in zip(scanned_tool_calls, pre_guardrail_tool_calls, post_guardrail_tool_calls)
if before != after
)
applicable: Final = tuple(
(item, after, rewritten_input) for item, after, rewritten_input in rewritten if rewritten_input is not None
)
if len(applicable) != len(rewritten):
raise unappliable_request_rewrite(guardrail_name)
for item, after, rewritten_input in applicable:
_write_back_tool_use(messages[item.target.msg_idx], item.target, after, rewritten_input)
async def process_output_response(
self,

View file

@ -5,6 +5,7 @@ from collections.abc import Awaitable, Callable
from typing import TYPE_CHECKING, Final, TypeAlias
from litellm._logging import verbose_logger
from litellm.litellm_core_utils.asyncify import asyncify
from litellm.types.llms.anthropic import AppliedEdit
from .constants import CLEAR_TOOL_USES_EDIT_TYPE, COMPACT_EDIT_TYPE
@ -82,9 +83,9 @@ async def apply_context_management(
"""Run edits in order; return a single ``PolyfillResult``.
The dispatcher is async so async editors (``compact_20260112``) can
``await`` the configured summarization model. Sync editors are called
inline ``inspect.iscoroutinefunction`` decides how each editor is
invoked.
``await`` the configured summarization model. Sync editors run in a
worker thread so their token counts stay off the event loop;
``inspect.iscoroutinefunction`` decides how each editor is invoked.
"""
edits: Final = _normalize_spec(context_management_spec)
if not edits:
@ -121,7 +122,7 @@ async def apply_context_management(
user_api_key_auth=user_api_key_auth,
)
if editor_is_async
else editor(
else await asyncify(editor)(
model=model,
messages=current_messages,
tools=tools,

View file

@ -20,6 +20,7 @@ from typing_extensions import NotRequired, ReadOnly, TypedDict, Unpack
import litellm
from litellm._logging import verbose_logger
from litellm.litellm_core_utils.asyncify import asyncify
from litellm.types.llms.anthropic import (
AppliedEdit,
CompactionBlock,
@ -1157,7 +1158,7 @@ async def apply_compact_20260112(
# Phase B: threshold check.
try:
current_tokens = _count_effective_tokens(
current_tokens = await asyncify(_count_effective_tokens)(
model=model,
effective_messages=effective_messages,
# ``augmented_system`` already carries the prior compaction summary

View file

@ -678,7 +678,7 @@ class BaseAnthropicMessagesStreamingIterator:
"""
from litellm.proxy.pass_through_endpoints.streaming_handler import PassThroughStreamingHandler
PassThroughStreamingHandler.schedule_stream_failure_logging(
await PassThroughStreamingHandler.schedule_stream_failure_logging(
litellm_logging_obj=self.litellm_logging_obj,
endpoint_type=EndpointType.ANTHROPIC,
request_body=self.request_body,

View file

@ -42,6 +42,10 @@ DROP_UNFITTING_REASONING_EFFORT_WARNING: Final = (
)
def _messages_carry_output_config(messages: Sequence[object]) -> bool:
return any(isinstance(message, Mapping) and "output_config" in message for message in messages)
class AnthropicMessagesConfig(BaseAnthropicMessagesConfig):
@property
def custom_llm_provider(self) -> str | None:
@ -331,6 +335,7 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig):
headers = self._update_headers_with_anthropic_beta(
headers=headers,
optional_params=optional_params,
messages=messages,
)
return headers, api_base
@ -664,6 +669,7 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig):
headers: dict,
optional_params: dict,
custom_llm_provider: str = "anthropic",
messages: Sequence[object] = (),
) -> dict:
"""
Auto-inject anthropic-beta headers based on features used.
@ -673,24 +679,30 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig):
- tool_search: adds provider-specific tool search header
- output_format: adds 'structured-outputs-2025-11-13'
- speed: adds 'fast-mode-2026-02-01'
- a message carrying output_config: adds 'per-turn-control-2026-07-01'
Args:
headers: Request headers dict
optional_params: Optional parameters including tools, context_management, output_format, speed
custom_llm_provider: Provider name for looking up correct tool search header
messages: Request messages, scanned for per-message output_config
"""
beta_values: Final[set] = set()
# Get existing beta headers if any
existing_beta: Final = headers.get("anthropic-beta")
if existing_beta:
beta_values.update(b.strip() for b in existing_beta.split(","))
existing_beta: Final = tuple(
piece.strip()
for key, value in headers.items()
if key.lower() == "anthropic-beta"
for piece in value.split(",")
if piece.strip()
)
beta_values.update(existing_beta)
# Check for context management
context_management_param: Final = optional_params.get("context_management")
if context_management_param is not None:
# Check edits array for compact_20260112 type
edits: Final = context_management_param.get("edits", [])
edits: Final = context_management_param.get("edits", ())
has_compact = False
has_other = False
@ -722,24 +734,18 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig):
if optional_params.get("speed") == "fast":
beta_values.add(ANTHROPIC_BETA_HEADER_VALUES.FAST_MODE_2026_02_01.value)
# Check for advisor tool
tools = optional_params.get("tools")
if tools:
for tool in tools:
if isinstance(tool, dict) and tool.get("type") == ANTHROPIC_ADVISOR_TOOL_TYPE:
beta_values.add(ANTHROPIC_BETA_HEADER_VALUES.ADVISOR_TOOL_2026_03_01.value)
break
if _messages_carry_output_config(messages):
beta_values.add(ANTHROPIC_BETA_HEADER_VALUES.PER_TURN_CONTROL_2026_07_01.value)
# Check for tool search tools
tools = optional_params.get("tools")
if tools:
anthropic_model_info: Final = AnthropicModelInfo()
if anthropic_model_info.is_tool_search_used(tools):
# Use provider-specific tool search header
tool_search_header: Final = get_tool_search_beta_header(custom_llm_provider)
beta_values.add(tool_search_header)
tools: Final = optional_params.get("tools")
if any(isinstance(tool, dict) and tool.get("type") == ANTHROPIC_ADVISOR_TOOL_TYPE for tool in tools or ()):
beta_values.add(ANTHROPIC_BETA_HEADER_VALUES.ADVISOR_TOOL_2026_03_01.value)
if beta_values:
headers["anthropic-beta"] = ",".join(sorted(beta_values))
if AnthropicModelInfo().is_tool_search_used(tools):
beta_values.add(get_tool_search_beta_header(custom_llm_provider))
return headers
if not beta_values:
return headers
merged: Final = {key: value for key, value in headers.items() if key.lower() != "anthropic-beta"}
merged["anthropic-beta"] = ",".join(sorted(beta_values))
return merged

View file

@ -0,0 +1,388 @@
from __future__ import annotations
import hashlib
import json
from collections.abc import Mapping, Sequence
from dataclasses import dataclass
from itertools import accumulate
from types import MappingProxyType
from typing import Annotated, Final, Literal, Protocol, TypeAlias
import httpx
from pydantic import BaseModel, ConfigDict, Field, JsonValue, StrictInt, TypeAdapter, ValidationError
import litellm
from litellm.llms.anthropic.common_utils import AnthropicModelInfo, is_anthropic_oauth_key
from litellm.llms.anthropic.count_tokens.handler import AnthropicCountTokensHandler
from litellm.llms.anthropic.experimental_pass_through.messages.transformation import DEFAULT_ANTHROPIC_API_VERSION
from litellm.types.router import LiteLLM_Params
from litellm.types.utils import ModelResponse
_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
_HEADERS: Final = TypeAdapter(dict[str, str])
_counter: Final = AnthropicCountTokensHandler()
_NATIVE_HEADERS: Final = frozenset(
(
"host",
"accept",
"accept-encoding",
"connection",
"user-agent",
"content-length",
"content-type",
"x-api-key",
"anthropic-version",
)
)
_DEPLOYMENT_OPTIONS: Final = frozenset(
{
"model",
"api_key",
"api_base",
"custom_llm_provider",
"rpm",
"tpm",
"timeout",
"stream_timeout",
"max_retries",
"num_retries",
"max_parallel_requests",
"input_cost_per_token",
"output_cost_per_token",
"cache_read_input_token_cost",
"cache_creation_input_token_cost",
"cache_creation_input_token_cost_above_1hr",
}
)
class _StrictModel(BaseModel):
model_config = ConfigDict(extra="forbid", frozen=True, strict=True)
class _CacheControl(_StrictModel):
type: Literal["ephemeral"]
ttl: Literal["5m", "1h"] = "5m"
class _Text(_StrictModel):
type: Literal["text"]
text: str = Field(min_length=1, pattern=r"\S")
cache_control: _CacheControl | None = None
class _ToolUse(_StrictModel):
type: Literal["tool_use"]
id: str = Field(min_length=1)
name: str = Field(min_length=1)
input: Mapping[str, JsonValue]
cache_control: _CacheControl | None = None
class _ResultText(_StrictModel):
type: Literal["text"]
text: str
class _ToolResult(_StrictModel):
type: Literal["tool_result"]
tool_use_id: str = Field(min_length=1)
content: str | Annotated[tuple[_ResultText, ...], Field(strict=False)]
is_error: bool | None = None
cache_control: _CacheControl | None = None
_Block: TypeAlias = Annotated[_Text | _ToolUse | _ToolResult, Field(discriminator="type")]
class _Message(_StrictModel):
role: Literal["user", "assistant"]
content: str | Annotated[tuple[_Block, ...], Field(strict=False)]
def blocks(self) -> tuple[_Text | _ToolUse | _ToolResult, ...]:
return (_Text(type="text", text=self.content),) if isinstance(self.content, str) else tuple(self.content)
class _Tool(_StrictModel):
name: str = Field(min_length=1)
description: str | None = None
input_schema: Mapping[str, JsonValue]
type: Literal["custom"] | None = None
class _Request(_StrictModel):
messages: tuple[_Message, ...] = Field(min_length=1, strict=False)
system: str | Annotated[tuple[_ResultText, ...], Field(strict=False)] | None = None
tools: Annotated[tuple[_Tool, ...], Field(strict=False)] | None = None
model: str | None = None
max_tokens: int | None = None
stream: bool | None = None
temperature: float | int | None = None
top_p: float | int | None = None
top_k: int | None = None
stop_sequences: Annotated[tuple[str, ...], Field(strict=False)] | None = None
metadata: Mapping[str, JsonValue] | None = None
@dataclass(frozen=True, slots=True)
class PromptPrefix:
prefix_body: Mapping[str, JsonValue]
fingerprint: str
fingerprints: tuple[str, ...]
ttl_seconds: int
def _digest(value: object) -> str:
return hashlib.sha256(
json.dumps(value, sort_keys=True, separators=(",", ":"), ensure_ascii=False).encode()
).hexdigest()
def _next_digest(previous: str, boundary: tuple[int, str, Mapping[str, JsonValue]]) -> str:
return _digest((previous, boundary))
def parse_prompt(body: Mapping[str, JsonValue]) -> PromptPrefix | None:
try:
request: Final = _Request.model_validate(body)
blocks: Final = tuple(message.blocks() for message in request.messages)
except ValidationError:
return None
markers: Final = tuple(
(message_index, block_index, block.cache_control)
for message_index, message_blocks in enumerate(blocks)
for block_index, block in enumerate(message_blocks)
if block.cache_control is not None
)
if len(markers) != 1:
return None
message_end, block_end, marker = markers[0]
normalized: Final = _JSON_OBJECT.validate_python(request.model_dump(mode="json", exclude_none=True))
context: Final = MappingProxyType({key: normalized[key] for key in ("system", "tools") if key in normalized})
boundaries: Final = tuple(
(
message_index,
request.messages[message_index].role,
_JSON_OBJECT.validate_python(
block.model_dump(mode="json", exclude=MappingProxyType({"cache_control": True}), exclude_none=True)
),
)
for message_index, message_blocks in enumerate(blocks[: message_end + 1])
for block_index, block in enumerate(message_blocks)
if message_index < message_end or block_index <= block_end
)
hashes: Final = tuple(
accumulate(boundaries, _next_digest, initial=_digest((_JSON_OBJECT.validate_python(context), marker.ttl)))
)[1:]
prefix_messages: Final = tuple(
_Message(
role=request.messages[message_index].role,
content=tuple(
block
for block_index, block in enumerate(message_blocks)
if message_index < message_end or block_index <= block_end
),
)
for message_index, message_blocks in enumerate(blocks[: message_end + 1])
)
return PromptPrefix(
prefix_body=MappingProxyType(
_JSON_OBJECT.validate_python(
_Request(messages=prefix_messages, system=request.system, tools=request.tools).model_dump(
mode="json", exclude_none=True
)
)
),
fingerprint=hashes[-1],
fingerprints=tuple(reversed(hashes[-20:])),
ttl_seconds=3600 if marker.ttl == "1h" else 300,
)
def cache_scope(
caller_key_hash: str,
deployment_id: str,
provider_key: str,
model: str,
anthropic_version: str = DEFAULT_ANTHROPIC_API_VERSION,
) -> str:
return _digest((caller_key_hash, deployment_id, provider_key, model, anthropic_version))
class _TTLUsage(BaseModel):
model_config = ConfigDict(strict=True)
ephemeral_5m_input_tokens: int = Field(default=0, ge=0)
ephemeral_1h_input_tokens: int = Field(default=0, ge=0)
class _CacheUsage(BaseModel):
model_config = ConfigDict(strict=True)
cached_tokens: int = Field(default=0, ge=0)
cache_creation_tokens: int = Field(default=0, ge=0)
cache_creation_token_details: _TTLUsage | None = None
class _Usage(BaseModel):
model_config = ConfigDict(strict=True)
prompt_tokens: int = Field(ge=0)
prompt_tokens_details: _CacheUsage
class _Choice(BaseModel):
finish_reason: str = Field(min_length=1)
class _Response(BaseModel):
model_config = ConfigDict(strict=True)
model: str
usage: _Usage
choices: tuple[_Choice, ...] = Field(min_length=1, strict=False)
class _CountBody(BaseModel):
messages: Sequence[Mapping[str, JsonValue]]
tools: Sequence[Mapping[str, JsonValue]] | None = None
system: str | Sequence[Mapping[str, JsonValue]] | None = None
class _CountResult(BaseModel):
input_tokens: Annotated[StrictInt, Field(ge=0)]
class TokenCounter(Protocol):
async def __call__(self, model: str, api_key: str, body: Mapping[str, JsonValue]) -> int | None: ...
def _count_objects(
values: Sequence[Mapping[str, JsonValue]],
) -> list[dict[str, JsonValue]]: # mutable-ok: the existing provider count API requires JSON lists/dicts
return [dict(value) for value in values] # mutable-ok: serialize read-only inputs at the provider API boundary
async def count_prompt_tokens(model: str, api_key: str, body: Mapping[str, JsonValue]) -> int | None:
native: Final = _CountBody.model_validate(body)
try:
result: Final = _CountResult.model_validate(
await _counter.handle_count_tokens_request(
model=model,
messages=_count_objects(native.messages),
tools=_count_objects(native.tools) if native.tools is not None else None,
system=native.system,
api_key=api_key,
timeout=15.0,
)
)
except Exception: # noqa: BLE001 # provider/count validation failures are unavailable estimates, not zero tokens
return None
return result.input_tokens
@dataclass(frozen=True, slots=True)
class NativePredictionTarget:
model: str
api_key: str
@dataclass(frozen=True, slots=True)
class UnsupportedPredictionTarget:
reason: Literal[
"unsupported_deployment_configuration",
"unsupported_provider_endpoint",
"unsupported_provider",
"unsupported_provider_credentials",
]
def resolve_prediction_target(params: LiteLLM_Params) -> NativePredictionTarget | UnsupportedPredictionTarget:
configured_options: Final = frozenset(params.model_dump(exclude_defaults=True, exclude_none=True))
if configured_options - _DEPLOYMENT_OPTIONS:
return UnsupportedPredictionTarget("unsupported_deployment_configuration")
api_base: Final = AnthropicModelInfo.get_api_base(params.api_base)
if api_base not in ("https://api.anthropic.com", "https://api.anthropic.com/v1/messages"):
return UnsupportedPredictionTarget("unsupported_provider_endpoint")
try:
model, provider, _, _ = litellm.get_llm_provider(
model=params.model, custom_llm_provider=params.custom_llm_provider
)
except Exception: # noqa: BLE001 # the shared provider resolver raises for unknown deployments
return UnsupportedPredictionTarget("unsupported_provider")
if provider != "anthropic":
return UnsupportedPredictionTarget("unsupported_provider")
api_key: Final = AnthropicModelInfo.get_api_key(params.api_key)
if api_key is None or not _supported_provider_key(api_key):
return UnsupportedPredictionTarget("unsupported_provider_credentials")
return NativePredictionTarget(model=model, api_key=api_key)
def _supported_provider_key(api_key: str) -> bool:
return bool(api_key) and not is_anthropic_oauth_key(api_key)
def supported_prediction_headers(headers: Mapping[str, str]) -> bool:
return all(
name.lower() != "anthropic-beta"
and (name.lower() != "anthropic-version" or value == DEFAULT_ANTHROPIC_API_VERSION)
for name, value in headers.items()
)
@dataclass(frozen=True, slots=True)
class ObservedCachePrefix:
prefix: PromptPrefix
scope: str
cached_tokens: int
cache_creation_tokens: int
def parse_observed_cache(
wire: httpx.Request, response_obj: ModelResponse, caller_key_hash: str, deployment_id: str
) -> ObservedCachePrefix | None:
try:
response: Final = _Response.model_validate(response_obj, from_attributes=True)
body: Final = _JSON_OBJECT.validate_json(wire.content)
headers: Final = _HEADERS.validate_python(wire.headers)
except (ValidationError, RuntimeError, httpx.RequestNotRead):
return None
if (
wire.url.scheme != "https"
or wire.url.host != "api.anthropic.com"
or wire.url.path != "/v1/messages"
or wire.url.query
or wire.url.port not in (None, 443)
):
return None
if (
frozenset(headers) - _NATIVE_HEADERS
or not supported_prediction_headers(headers)
or headers.get("anthropic-version") != DEFAULT_ANTHROPIC_API_VERSION
):
return None
provider_key: Final = headers.get("x-api-key", "")
model: Final = body.get("model")
if not _supported_provider_key(provider_key) or not isinstance(model, str) or model != response.model:
return None
prefix: Final = parse_prompt(body)
if prefix is None:
return None
usage: Final = response.usage.prompt_tokens_details
cache_tokens: Final = usage.cached_tokens + usage.cache_creation_tokens
if cache_tokens <= 0 or cache_tokens > response.usage.prompt_tokens:
return None
split: Final = usage.cache_creation_token_details
if usage.cache_creation_tokens and split is None:
return None
if split is not None and (
split.ephemeral_5m_input_tokens + split.ephemeral_1h_input_tokens != usage.cache_creation_tokens
or (prefix.ttl_seconds == 300 and split.ephemeral_1h_input_tokens > 0)
or (prefix.ttl_seconds == 3600 and split.ephemeral_5m_input_tokens > 0)
):
return None
return ObservedCachePrefix(
prefix=prefix,
scope=cache_scope(caller_key_hash, deployment_id, provider_key, model),
cached_tokens=cache_tokens,
cache_creation_tokens=usage.cache_creation_tokens,
)

View file

@ -68,6 +68,7 @@ class AzureAnthropicMessagesConfig(AnthropicMessagesConfig):
headers = self._update_headers_with_anthropic_beta(
headers=headers,
optional_params=optional_params,
messages=messages,
)
return headers, api_base

View file

@ -1,8 +1,8 @@
from __future__ import annotations
import json
from collections.abc import Callable, Iterator, Sequence
from typing import Final, TypeVar
from collections.abc import Callable, Iterator, Mapping, Sequence
from typing import Final, TypeVar, cast # noqa: TID251 # a rebuilt chat row has no typed constructor across roles
from pydantic import BaseModel
@ -364,3 +364,67 @@ def merge_guardrailed_scoped_messages(
yield from appended
return list(_merged())
def _content_part_text(part: object) -> str | None:
if not isinstance(part, Mapping):
return None
text: Final = part.get("text")
return text if isinstance(text, str) else None
def message_slot_texts(message: Mapping[str, object]) -> tuple[str, ...]:
content: Final = message.get("content")
if isinstance(content, str):
return (content,)
if isinstance(content, list):
return tuple(text for part in content if (text := _content_part_text(part)) is not None)
return ()
def message_text_slot_count(message: AllMessageValues) -> int:
return len(message_slot_texts(message))
def _part_with_text(part: object, text: str) -> object:
if not isinstance(part, Mapping):
return part
return {**part, "text": text} # mutable-ok: content parts stay JSON-plain dicts
def _content_with_slot_texts(content: Sequence[object], texts: Sequence[str]) -> Sequence[object]:
remaining_texts: Final = iter(texts)
return [ # mutable-ok: message content stays a JSON list
_part_with_text(part, next(remaining_texts)) if _content_part_text(part) is not None else part
for part in content
]
def message_with_slot_texts(message: AllMessageValues, texts: Sequence[str]) -> AllMessageValues | None:
"""Swap one rewritten text into each text slot of a chat row, in order.
A slot is a string ``content`` or one list part carrying a string ``text``;
images and other parts ride along untouched. Returns None unless the counts
line up exactly, so a rewrite never lands on the wrong slot.
"""
if message_text_slot_count(message) != len(texts):
return None
content: Final = message.get("content")
if not isinstance(content, (str, list)):
return message
rewritten_content: Final = texts[0] if isinstance(content, str) else _content_with_slot_texts(content, texts)
rewritten: Final = {**message, "content": rewritten_content} # mutable-ok: chat rows stay JSON-plain dicts
return cast("AllMessageValues", rewritten) # cast-ok: the same row with only its text slots swapped
class UnappliableRequestRewrite(Exception):
def __init__(self, guardrail_name: str) -> None:
super().__init__(
f"Guardrail '{guardrail_name}' rewrote the request in a way this endpoint cannot apply, "
"so the request was rejected rather than sent unrewritten"
)
self.guardrail_name: Final = guardrail_name
def unappliable_request_rewrite(guardrail_name: str | None) -> UnappliableRequestRewrite:
return UnappliableRequestRewrite(guardrail_name or "unknown")

View file

@ -11,6 +11,7 @@ from concurrent.futures import ThreadPoolExecutor
from datetime import datetime
from functools import partial
from threading import Lock
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, ParamSpec, TypeVar, cast, get_args, overload
import httpx
@ -96,6 +97,77 @@ def _assume_role_params(
)
_SecureTransportBool = TypedDict("_SecureTransportBool", {"aws:SecureTransport": ReadOnly[Literal["true"]]})
class _SecureTransportCondition(TypedDict):
Bool: ReadOnly[_SecureTransportBool]
class _SessionPolicyStatement(TypedDict):
Sid: ReadOnly[str]
Effect: ReadOnly[Literal["Allow"]]
Action: ReadOnly[tuple[str, ...]]
Resource: ReadOnly[Literal["*"]]
Condition: ReadOnly[_SecureTransportCondition]
class WebIdentitySessionPolicy(TypedDict):
Version: ReadOnly[Literal["2012-10-17"]]
Statement: ReadOnly[tuple[_SessionPolicyStatement, ...]]
_WEB_IDENTITY_SESSION_POLICY_ACTIONS: Final[Mapping[str, tuple[str, ...]]] = MappingProxyType(
{
"BedrockLiteLLM": (
"bedrock:InvokeModel",
"bedrock:InvokeModelWithResponseStream",
"bedrock:CountTokens",
"bedrock:Rerank",
"bedrock:Retrieve",
"bedrock:ListKnowledgeBases",
"bedrock:InvokeAgent",
"bedrock:ApplyGuardrail",
"bedrock:GetGuardrail",
"bedrock:ListGuardrails",
),
"BedrockAgentCoreLiteLLM": (
"bedrock-agentcore:InvokeAgentRuntime",
"bedrock-agentcore:InvokeAgentRuntimeForUser",
"bedrock-agentcore:InvokeGateway",
),
"ClaudePlatformLiteLLM": (
"aws-external-anthropic:CreateInference",
"aws-external-anthropic:CreateBatchInference",
"aws-external-anthropic:CancelBatchInference",
"aws-external-anthropic:DeleteBatchInference",
"aws-external-anthropic:CountTokens",
"aws-external-anthropic:Get*",
"aws-external-anthropic:List*",
),
"BedrockMantleLiteLLM": ("bedrock-mantle:CreateInference",),
}
)
_SECURE_TRANSPORT_ONLY: Final = _SecureTransportCondition(Bool=_SecureTransportBool({"aws:SecureTransport": "true"}))
def build_web_identity_session_policy() -> WebIdentitySessionPolicy:
return WebIdentitySessionPolicy(
Version="2012-10-17",
Statement=tuple(
_SessionPolicyStatement(
Sid=sid,
Effect="Allow",
Action=actions,
Resource="*",
Condition=_SECURE_TRANSPORT_ONLY,
)
for sid, actions in _WEB_IDENTITY_SESSION_POLICY_ACTIONS.items()
),
)
class BedrockRequestTarget(BaseModel):
aws_region_name: str
aws_bedrock_runtime_endpoint: str | None
@ -940,60 +1012,12 @@ class BaseAWSLLM(SignsRequestsWithAWS):
# auth only (static creds + IRSA take other code paths).
# https://docs.aws.amazon.com/STS/latest/APIReference/API_AssumeRoleWithWebIdentity.html
# https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/sts/client/assume_role_with_web_identity.html
bedrock_session_policy: Final = {
"Version": "2012-10-17",
"Statement": [
{
"Sid": "BedrockLiteLLM",
"Effect": "Allow",
"Action": [
"bedrock:InvokeModel",
"bedrock:InvokeModelWithResponseStream",
"bedrock:CountTokens",
"bedrock:ApplyGuardrail",
"bedrock:GetGuardrail",
"bedrock:ListGuardrails",
],
"Resource": "*",
"Condition": {"Bool": {"aws:SecureTransport": "true"}},
},
# Claude Platform on AWS (added by #27678 for the
# ``bedrock/claude_platform/<model>`` route) lives under
# a separate IAM action namespace; without these entries
# the OIDC path 403s on every claude_platform request
# even with a fully permissive identity policy (#30200).
{
"Sid": "ClaudePlatformLiteLLM",
"Effect": "Allow",
"Action": [
"aws-external-anthropic:CreateInference",
"aws-external-anthropic:CreateBatchInference",
"aws-external-anthropic:CancelBatchInference",
"aws-external-anthropic:DeleteBatchInference",
"aws-external-anthropic:CountTokens",
"aws-external-anthropic:Get*",
"aws-external-anthropic:List*",
],
"Resource": "*",
"Condition": {"Bool": {"aws:SecureTransport": "true"}},
},
{
"Sid": "BedrockMantleLiteLLM",
"Effect": "Allow",
"Action": [
"bedrock-mantle:CreateInference",
],
"Resource": "*",
"Condition": {"Bool": {"aws:SecureTransport": "true"}},
},
],
}
assume_role_params: Final = {
"RoleArn": aws_role_name,
"RoleSessionName": aws_session_name,
"WebIdentityToken": oidc_token,
"DurationSeconds": 3600,
"Policy": json.dumps(bedrock_session_policy, separators=(",", ":")),
"Policy": json.dumps(build_web_identity_session_policy(), separators=(",", ":")),
}
# Add ExternalId parameter if provided

View file

@ -46,6 +46,7 @@ class BedrockClaudePlatformMessagesConfig(BedrockClaudePlatformMixin, AnthropicM
headers = self._update_headers_with_anthropic_beta(
headers=headers,
optional_params=optional_params,
messages=messages,
)
return headers, api_base

View file

@ -17,7 +17,7 @@ BaseAWSLLM._sign_request after the request body is finalized.
import json
from collections.abc import Mapping
from typing import Any, Final
from typing import Any, Final, cast # noqa: TID251 # map_openai_params returns the filtered params as a bare dict
import httpx
from typing_extensions import ReadOnly, TypedDict
@ -32,6 +32,7 @@ from litellm.llms.bedrock_mantle.common_utils import (
BedrockMantleAuthMixin,
)
from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig
from litellm.responses.additional_tools import HoistedAdditionalTools, hoist_additional_tools
from litellm.secret_managers.main import get_secret_str
from litellm.types.llms.openai import (
ResponseInputParam,
@ -58,8 +59,6 @@ _BEDROCK_MANTLE_SUPPORTED_RESPONSE_TOOL_TYPES: Final = frozenset(
_BEDROCK_MANTLE_SUPPORTED_SERVICE_TIERS: Final = frozenset({"auto", "default"})
_BEDROCK_MANTLE_OPENAI_PATH_SUPPORTED_REASONING_SUMMARIES: Final = frozenset({"auto"})
_CODEX_ADDITIONAL_TOOLS_INPUT_ITEM_TYPE: Final = "additional_tools"
_CODEX_AGENT_MESSAGE_INPUT_ITEM_TYPE: Final = "agent_message"
_CODEX_CONTEXT_COMPACTION_INPUT_ITEM_TYPE: Final = "context_compaction"
_CODEX_LOCAL_SHELL_CALL_INPUT_ITEM_TYPE: Final = "local_shell_call"
@ -233,17 +232,14 @@ class BedrockMantleResponsesAPIConfig(BedrockMantleAuthMixin, OpenAIResponsesAPI
litellm_params: GenericLiteLLMParams,
headers: dict,
) -> dict:
remaining_input, hoisted_tools = self._hoist_codex_additional_tools(input)
normalized_input: Final = self._normalize_codex_input_items(remaining_input)
params: Final = cast( # cast-ok: the base signature leaves the params dict untyped
"ResponsesAPIOptionalRequestParams", response_api_optional_request_params
)
hoisted: Final = hoist_additional_tools(input, params.get("tools"))
normalized_input: Final = self._normalize_codex_input_items(hoisted.input)
request_params: Final = (
{
**response_api_optional_request_params,
"tools": [
*(response_api_optional_request_params.get("tools") or []),
*hoisted_tools,
],
}
if hoisted_tools
self._params_with_hoisted_tools(params, hoisted)
if hoisted.hoisted
else response_api_optional_request_params
)
return super().transform_responses_api_request(
@ -254,41 +250,14 @@ class BedrockMantleResponsesAPIConfig(BedrockMantleAuthMixin, OpenAIResponsesAPI
headers=headers,
)
@staticmethod
def _is_codex_additional_tools_item(item: Any) -> bool:
return isinstance(item, dict) and item.get("type") == _CODEX_ADDITIONAL_TOOLS_INPUT_ITEM_TYPE
@staticmethod
def _tools_of_additional_tools_item(item: "dict[str, Any]") -> "list[Any]":
tools: Final = item.get("tools")
return tools if isinstance(tools, list) else []
@classmethod
def _hoist_codex_additional_tools(
cls,
input: "str | ResponseInputParam",
) -> "tuple[str | ResponseInputParam, list[Any]]":
"""Codex's "responses lite" wire mode ships tool definitions inside
`input` as {"type": "additional_tools", "role": "developer",
"tools": [...]} items. api.openai.com accepts that item type; Mantle
rejects the whole request with 400 "Invalid 'input': value did not
match any expected variant" but accepts the same tools at the top
level, so move them there and strip the items from `input`.
"""
if not isinstance(input, list):
return input, []
additional_tools_items: Final = [item for item in input if cls._is_codex_additional_tools_item(item)]
if not additional_tools_items:
return input, []
remaining_input: Final = [item for item in input if not cls._is_codex_additional_tools_item(item)]
hoisted_tools = [tool for item in additional_tools_items for tool in cls._tools_of_additional_tools_item(item)]
verbose_logger.debug(
"Bedrock Mantle Responses API: hoisting %d tool(s) out of %d 'additional_tools' input item(s) "
"into the top-level tools param (Mantle rejects that input item type).",
len(hoisted_tools),
len(additional_tools_items),
)
return remaining_input, cls._filter_unsupported_tools(hoisted_tools)
def _params_with_hoisted_tools(
cls, params: Mapping[str, object], hoisted: HoistedAdditionalTools
) -> dict[str, object]:
supported_tools: Final = cls._filter_unsupported_tools(list(hoisted.tools))
if supported_tools:
return {**params, "tools": supported_tools}
return {key: value for key, value in params.items() if key != "tools"}
@staticmethod
def _agent_message_text(item: "Mapping[str, object]") -> str:

View file

@ -66,6 +66,7 @@ class DeepSeekAnthropicMessagesConfig(AnthropicMessagesConfig):
headers=headers,
optional_params=optional_params,
custom_llm_provider=self.custom_llm_provider or "deepseek",
messages=messages,
)
return headers, api_base

View file

@ -115,6 +115,19 @@ def _parse_setup(session_configuration_request: str) -> BidiGenerateContentSetup
return envelope.get("setup", empty_setup)
def _grounding_metadata_from_frame(frame: Mapping[str, object]) -> tuple[Mapping[str, object], ...]:
"""Read ``serverContent.groundingMetadata`` off the frame that carries the turn's usage.
Live reports grounding in the server frames rather than in ``usageMetadata``, and it emits both
on the same frame, so the per-query charge is countable at the point usage is built.
"""
server_content: Final = frame.get("serverContent")
if not isinstance(server_content, Mapping):
return ()
metadata: Final = server_content.get("groundingMetadata")
return (metadata,) if isinstance(metadata, Mapping) else ()
# Google bills Live transcription at an estimated 25 audio tokens/sec of input and
# 175 text tokens/min of output (ai.google.dev/gemini-api/docs/pricing).
GEMINI_LIVE_TRANSCRIBE_AUDIO_TOKENS_PER_SECOND: Final = 25
@ -323,7 +336,7 @@ class GeminiRealtimeConfig(BaseRealtimeConfig):
)
elif key == "input_audio_transcription" and value is not None:
optional_params["inputAudioTranscription"] = {}
elif key == "turn_detection":
elif key == "turn_detection" and value is not None:
value_typed = cast(OpenAIRealtimeTurnDetection, value)
if (
isinstance(value_typed, dict)
@ -1049,6 +1062,11 @@ class GeminiRealtimeConfig(BaseRealtimeConfig):
{**cast(dict, message), "usageMetadata": resolved_usage_metadata},
),
)
grounding_metadata: Final = _grounding_metadata_from_frame(message)
if grounding_metadata:
VertexGeminiConfig._set_grounding_usage_counters( # pyright: ignore[reportPrivateUsage] # shared with the chat path; no public alias exists yet
_chat_completion_usage, grounding_metadata
)
else:
_chat_completion_usage = get_empty_usage()

View file

@ -92,7 +92,7 @@ class GithubCopilotAnthropicMessagesConfig(AnthropicMessagesConfig):
headers["anthropic-version"] = "2023-06-01"
headers = self._update_headers_with_anthropic_beta(
headers, optional_params, custom_llm_provider="github_copilot"
headers, optional_params, custom_llm_provider="github_copilot", messages=messages
)
return headers, dynamic_api_base

View file

@ -42,6 +42,7 @@ from litellm.llms.base_llm.guardrail_translation.utils import (
stream_item_field,
stream_item_fingerprint,
stream_item_items,
unappliable_request_rewrite,
)
from litellm.main import stream_chunk_builder
from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolParam
@ -196,6 +197,8 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
else:
# Step 3: Map guardrail responses back to original message structure
if guardrailed_texts and texts_to_check:
if len(guardrailed_texts) != len(text_task_mappings):
raise unappliable_request_rewrite(guardrail_to_apply.guardrail_name)
await self._apply_guardrail_responses_to_input_texts(
messages=messages,
responses=guardrailed_texts,
@ -210,6 +213,17 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
task_mappings=tool_call_task_mappings,
)
elif (
not images_to_check
and not guardrail_to_apply.records_own_guardrail_information
and (not_run_reason := self._not_run_reason(messages)) is not None
):
guardrail_to_apply.add_standard_logging_guardrail_information_to_request_data(
guardrail_json_response=not_run_reason,
request_data=data,
guardrail_status="not_run",
)
verbose_proxy_logger.debug(
"OpenAI Chat Completions: Processed input messages: %s",
data.get("messages"),
@ -217,6 +231,28 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
return data
def _not_run_reason(
self,
messages: Sequence[dict[str, Any]], # mutable-ok: raw request messages consumed by _extract_inputs
) -> str | None:
"""Why nothing was scanned, or None when the only unscoped content is images, which this handler never scans."""
texts: Final[list[str]] = [] # mutable-ok: filled by _extract_inputs
images: Final[list[str]] = [] # mutable-ok: filled by _extract_inputs
tool_calls: Final[list[ChatCompletionToolParam]] = [] # mutable-ok: filled by _extract_inputs
for msg_idx, message in enumerate(messages):
self._extract_inputs(
message=message,
msg_idx=msg_idx,
texts_to_check=texts,
images_to_check=images,
tool_calls_to_check=tool_calls,
text_task_mappings=[], # mutable-ok: required by _extract_inputs, unused here
tool_call_task_mappings=[], # mutable-ok: required by _extract_inputs, unused here
)
if texts or tool_calls:
return "no scannable content after message scoping"
return None if images else "no scannable content"
def extract_request_tool_names(self, data: dict) -> list[str]:
"""Extract tool names from OpenAI chat completions request (tools[].function.name, functions[].name)."""
names: Final[list[str]] = []

View file

@ -56,6 +56,7 @@ from litellm.llms.base_llm.guardrail_translation.utils import (
stream_item_field,
stream_item_fingerprint,
stream_item_items,
unappliable_request_rewrite,
)
from litellm.llms.openai.responses.guardrail_translation.tool_merge import merge_guardrailed_tools
from litellm.responses.litellm_completion_transformation.transformation import (
@ -495,13 +496,13 @@ class OpenAIResponsesHandler(BaseTranslation):
data["instructions"] = written_back.instructions # rebind-ok: data is an out-param
elif isinstance(input_data, str):
guardrailed_texts: Final = guardrailed_inputs.get("texts") or ()
if len(guardrailed_texts) > 1:
raise unappliable_request_rewrite(guardrail_to_apply.guardrail_name)
data["input"] = guardrailed_texts[0] if guardrailed_texts else input_data # rebind-ok: data is an out-param
else:
rewritten_texts: Final = guardrailed_inputs.get("texts") or ()
if len(rewritten_texts) != len(extracted.task_mappings):
from litellm.proxy.policy_engine.pipeline_executor import UnappliableRequestRewrite
raise UnappliableRequestRewrite(guardrail_to_apply.guardrail_name or "unknown")
raise unappliable_request_rewrite(guardrail_to_apply.guardrail_name)
await self._apply_guardrail_responses_to_input(
messages=input_data,
responses=rewritten_texts,

View file

@ -6,8 +6,10 @@ from typing import Final, TypeAlias
from pydantic import BaseModel, TypeAdapter, ValidationError
from litellm._logging import verbose_logger
from litellm.responses.litellm_completion_transformation.custom_tools import custom_tool_grammar_suffix
from litellm.responses.litellm_completion_transformation.transformation import (
NAMESPACE_DESCRIPTION_SEPARATOR,
NAMESPACE_MEMBER_TYPES_WITH_CHAT_TOOLS,
LiteLLMCompletionResponsesConfig,
)
@ -34,8 +36,8 @@ def _validated_tools(values: Iterable[object]) -> tuple[Tool, ...]:
return tuple(tool for tool in validated if tool is not None)
def _is_function(tool: Tool) -> bool:
return tool.get("type") == "function"
def _has_chat_tool(member: Tool) -> bool:
return member.get("type") in NAMESPACE_MEMBER_TYPES_WITH_CHAT_TOOLS
def _chat_tool_key(tool: Tool) -> str:
@ -67,18 +69,19 @@ def _function_fields(tool: Tool) -> Tool:
return function if function is not None else MappingProxyType({})
def _without_namespace_prefix(key: str, value: object, prefix: str) -> object:
if key != "description" or not isinstance(value, str) or not value.startswith(prefix):
def _member_description(key: str, value: object, prefix: str, suffix: str) -> object:
if key != "description" or not isinstance(value, str):
return value
return value[len(prefix) :]
return value.replace(prefix, "", 1).replace(suffix, "", 1)
def _rebuilt_member(member: Tool, flattened: Tool, guardrailed: Tool, namespace_description: str) -> Tool:
flattened_function: Final = _function_fields(flattened)
prefix: Final = f"{namespace_description}{NAMESPACE_DESCRIPTION_SEPARATOR}" if namespace_description else ""
suffix: Final = custom_tool_grammar_suffix(member.get("format")) if member.get("type") == "custom" else ""
changed_function: Final = MappingProxyType(
{
key: _without_namespace_prefix(key, value, prefix)
key: _member_description(key, value, prefix, suffix)
for key, value in _function_fields(guardrailed).items()
if flattened_function.get(key) != value
}
@ -93,8 +96,8 @@ def _rebuilt_member(member: Tool, flattened: Tool, guardrailed: Tool, namespace_
return {**member, **changed_extras, **changed_function} # mutable-ok: json.dumps rejects MappingProxyType
def _rebuilt_function_members(
function_members: Sequence[Tool],
def _rebuilt_flattened_members(
flattened_members: Sequence[Tool],
flattened_group: Sequence[Tool],
group_keys: Sequence[IndexedKey],
guardrailed_by_key: Mapping[IndexedKey, Tool],
@ -106,7 +109,7 @@ def _rebuilt_function_members(
else member
if guardrailed_by_key[key] == flattened
else _rebuilt_member(member, flattened, guardrailed_by_key[key], namespace_description)
for member, flattened, key in zip(function_members, flattened_group, group_keys)
for member, flattened, key in zip(flattened_members, flattened_group, group_keys)
)
@ -118,9 +121,9 @@ def _rebuilt_namespace(
guardrailed_by_key: Mapping[IndexedKey, Tool],
) -> tuple[Tool, ...]:
namespace_description: Final = str(original.get("description") or "")
rebuilt_functions: Final = iter(
_rebuilt_function_members(
tuple(member for member in members if _is_function(member)),
rebuilt_flattened: Final = iter(
_rebuilt_flattened_members(
tuple(member for member in members if _has_chat_tool(member)),
flattened_group,
group_keys,
guardrailed_by_key,
@ -129,7 +132,7 @@ def _rebuilt_namespace(
)
rebuilt_members: Final = tuple(
rebuilt
for rebuilt in (next(rebuilt_functions) if _is_function(member) else member for member in members)
for rebuilt in (next(rebuilt_flattened) if _has_chat_tool(member) else member for member in members)
if rebuilt is not None
)
if not rebuilt_members:
@ -149,7 +152,7 @@ def _merged_original(
if guardrailed_group == tuple(flattened_group):
return (original,)
members: Final = _namespace_members(original) if original.get("type") == "namespace" else ()
if members and sum(map(_is_function, members)) == len(flattened_group):
if members and sum(map(_has_chat_tool, members)) == len(flattened_group):
return _rebuilt_namespace(original, members, flattened_group, group_keys, guardrailed_by_key)
if not guardrailed_group:
return ()

View file

@ -56,6 +56,7 @@ class OpenAILikeAnthropicMessagesConfig(AnthropicMessagesConfig):
merged: Final = self._update_headers_with_anthropic_beta(
headers=normalized,
optional_params=optional_params,
messages=messages,
)
return merged, api_base

View file

@ -1,3 +1,4 @@
from collections.abc import Mapping
from typing import Any, Final
from urllib.parse import unquote
@ -8,7 +9,36 @@ from litellm.llms.vertex_ai.common_utils import (
)
from litellm.types.llms.openai import BatchJobStatus, CreateBatchRequest
from litellm.types.llms.vertex_ai import *
from litellm.types.utils import LiteLLMBatch
from litellm.types.utils import LiteLLMBatch, PromptTokensDetailsWrapper
def vertex_prompt_tokens_details(
usage_metadata: Mapping[str, object],
) -> PromptTokensDetailsWrapper | None:
raw_details: Final = usage_metadata.get("promptTokensDetails")
if not isinstance(raw_details, list):
return None
def _normalize(detail: object) -> tuple[str, int] | None:
if not isinstance(detail, Mapping):
return None
modality: Final = detail.get("modality")
token_count: Final = detail.get("tokenCount")
if not isinstance(modality, str) or not isinstance(token_count, int):
return None
return modality.upper(), token_count
parsed_details: Final = tuple(_normalize(detail) for detail in raw_details)
normalized: Final = tuple(detail for detail in parsed_details if detail is not None)
if len(normalized) != len(parsed_details):
return None
return PromptTokensDetailsWrapper(
text_tokens=sum(token_count for modality, token_count in normalized if modality in ("TEXT", "DOCUMENT")),
audio_tokens=sum(token_count for modality, token_count in normalized if modality == "AUDIO"),
image_tokens=sum(token_count for modality, token_count in normalized if modality == "IMAGE"),
video_tokens=sum(token_count for modality, token_count in normalized if modality == "VIDEO"),
)
class VertexAIBatchTransformation:

View file

@ -298,8 +298,6 @@ def transform_openai_input_gemini_embed_content(
_IMAGE_MIME_TYPES: Final = frozenset({"image/png", "image/jpeg"})
_VIDEO_TOKENS_PER_SECOND: Final = 258.0
_AUDIO_TOKENS_PER_SECOND: Final = 32.0
_usage_metadata_adapter: Final = TypeAdapter(UsageMetadata)
@ -339,11 +337,12 @@ def _is_image_element(
return False
def _count_input_images(
def _is_image_only_input(
input: GeminiEmbeddingInput,
resolved_files: Mapping[str, Mapping[str, str]],
) -> int:
return sum(1 for element in _flatten_input(input) if _is_image_element(element, resolved_files))
) -> bool:
elements: Final = _flatten_input(input)
return bool(elements) and all(_is_image_element(element, resolved_files) for element in elements)
def _tokens_for_modality(details: Sequence[PromptTokensDetails], modality: str) -> int:
@ -372,30 +371,29 @@ def _usage_from_embed_content_response(
total_tokens: Final = usage_metadata.get("totalTokenCount") or prompt_tokens
details: Final[Sequence[PromptTokensDetails]] = usage_metadata.get("promptTokensDetails") or ()
if not details:
return Usage(
prompt_tokens=prompt_tokens,
total_tokens=total_tokens,
prompt_tokens_details=PromptTokensDetailsWrapper(
text_tokens=0,
image_tokens=prompt_tokens if _is_image_only_input(input, resolved_files) else 0,
),
)
text_tokens: Final = _tokens_for_modality(details, "TEXT")
audio_tokens: Final = _tokens_for_modality(details, "AUDIO")
image_tokens: Final = _tokens_for_modality(details, "IMAGE")
video_tokens: Final = _tokens_for_modality(details, "VIDEO")
image_count: Final = _count_input_images(input, resolved_files)
video_length_seconds: Final = video_tokens / _VIDEO_TOKENS_PER_SECOND if video_tokens > 0 else 0.0
audio_length_seconds: Final = audio_tokens / _AUDIO_TOKENS_PER_SECOND if audio_tokens > 0 else 0.0
# generic_cost_per_token rewrites text_tokens to the full prompt minus
# other modalities when both text_tokens and image_count are zero. For
# video, that misallocates video tokens to text; a 1-token floor sidesteps
# the rewrite and keeps billing on input_cost_per_video_per_second.
needs_video_text_floor: Final = video_length_seconds > 0 and text_tokens == 0 and image_count == 0
resolved_text_tokens: Final = 1 if needs_video_text_floor else text_tokens
return Usage(
prompt_tokens=prompt_tokens,
total_tokens=total_tokens,
prompt_tokens_details=PromptTokensDetailsWrapper(
text_tokens=resolved_text_tokens,
text_tokens=text_tokens,
audio_tokens=audio_tokens,
image_count=image_count,
video_length_seconds=video_length_seconds,
audio_length_seconds=audio_length_seconds,
image_tokens=image_tokens,
video_tokens=video_tokens,
),
)
@ -415,8 +413,7 @@ def process_embed_content_response(
model_response: EmbeddingResponse to populate
model: Model name
response_json: Raw JSON response from embedContent endpoint
resolved_files: Mapping of file references (files/abc) to {mime_type, uri},
used to bill resolved image references at the per-image rate
resolved_files: Mapping of file references to resolved metadata
Returns:
EmbeddingResponse with single embedding

View file

@ -67,7 +67,7 @@ from litellm.constants import (
)
from litellm.exceptions import LiteLLMUnknownProvider
from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.asyncify import run_async_function
from litellm.litellm_core_utils.asyncify import asyncify, run_async_function
from litellm.litellm_core_utils.audio_utils.utils import (
calculate_request_duration,
get_audio_file_for_health_check,
@ -9127,7 +9127,7 @@ async def acount_tokens(
fallback_messages = messages or []
if system and fallback_messages:
fallback_messages = [{"role": "system", "content": system}] + fallback_messages
local_count: Final = litellm.token_counter(
local_count: Final = await asyncify(litellm.token_counter)(
model=model,
messages=fallback_messages,
tools=tools,

File diff suppressed because it is too large Load diff

View file

@ -26,6 +26,7 @@ class LiteLLM_BudgetTable(LiteLLMPydanticObjectBase):
max_parallel_requests: int | None = None
tpm_limit: int | None = None
rpm_limit: int | None = None
tpd_limit: int | None = None
model_max_budget: dict | None = None
budget_duration: str | None = None
allowed_models: list[str] | None = None # per-member model scope; empty = inherit team models

View file

@ -5,6 +5,8 @@ These are the canonical credential types for the proxy. They live in the model
layer; ``litellm.types.utils`` re-exports them for backwards compatibility.
"""
from collections.abc import Mapping
from pydantic import BaseModel, model_validator
@ -27,3 +29,10 @@ class CreateCredentialItem(CredentialBase):
if not values.get("credential_values") and not values.get("model_id"):
raise ValueError("Either credential_values or model_id must be set")
return values
class UpdateCredentialItem(BaseModel):
credential_name: str
credential_info: Mapping[str, object]
credential_values: Mapping[str, object] | None = None
model_id: str | None = None

View file

@ -71,6 +71,7 @@ class TeamBase(LiteLLMPydanticObjectBase):
metadata: dict | None = None
tpm_limit: int | None = None
rpm_limit: int | None = None
tpd_limit: int | None = None
max_budget: float | None = None
soft_budget: float | None = None
budget_duration: str | None = None

View file

@ -31,6 +31,7 @@ class LiteLLM_VerificationToken(LiteLLMPydanticObjectBase):
metadata: dict = {}
tpm_limit: int | None = None
rpm_limit: int | None = None
tpd_limit: int | None = None
budget_duration: str | None = None
budget_reset_at: datetime | None = None
allowed_cache_controls: list | None = []

View file

@ -11503,6 +11503,12 @@
"description": "Enable content moderation to check for harmful content (harassment, hate speech, etc.).",
"title": "Content Moderation Check"
},
"contextual_grounding_from_messages": {
"default": false,
"description": "ApplyGuardrail: when True, post-call scans of a request with no grounding_source / query content parts send the system and developer messages as the grounding source and the latest user message as the query, so the guardrail's contextual grounding policy can score the response. Bedrock bills contextual grounding units for these scans and rejects queries, sources and responses over its contextual grounding length limits, so leave this off for guardrails without a contextual grounding policy. Default False: plain messages are never sent as grounding context.",
"title": "Contextual Grounding From Messages",
"type": "boolean"
},
"credentials": {
"anyOf": [
{
@ -12962,18 +12968,24 @@
"PHONE_NUMBER",
"MEDICAL_LICENSE",
"URL",
"MAC_ADDRESS",
"UUID",
"US_BANK_NUMBER",
"US_DRIVER_LICENSE",
"US_ITIN",
"US_PASSPORT",
"US_SSN",
"US_MBI",
"US_NPI",
"UK_NHS",
"UK_NINO",
"UK_PASSPORT",
"UK_POSTCODE",
"UK_VEHICLE_REGISTRATION",
"UK_DRIVING_LICENCE",
"ES_NIF",
"ES_NIE",
"ES_PASSPORT",
"IT_FISCAL_CODE",
"IT_DRIVER_LICENSE",
"IT_VAT_CODE",
@ -12991,7 +13003,38 @@
"IN_VEHICLE_REGISTRATION",
"IN_VOTER",
"IN_PASSPORT",
"FI_PERSONAL_IDENTITY_CODE"
"IN_GSTIN",
"FI_PERSONAL_IDENTITY_CODE",
"DE_TAX_ID",
"DE_TAX_NUMBER",
"DE_VAT_ID",
"DE_PASSPORT",
"DE_ID_CARD",
"DE_FUEHRERSCHEIN",
"DE_SOCIAL_SECURITY",
"DE_HEALTH_INSURANCE",
"DE_LANR",
"DE_BSNR",
"DE_KFZ",
"DE_HANDELSREGISTER",
"DE_PLZ",
"KR_RRN",
"KR_FRN",
"KR_PASSPORT",
"KR_DRIVER_LICENSE",
"KR_BRN",
"CA_SIN",
"SE_PERSONNUMMER",
"SE_ORGANISATIONSNUMMER",
"TH_TNIN",
"TR_NATIONAL_ID",
"TR_LICENSE_PLATE",
"NG_NIN",
"NG_VEHICLE_REGISTRATION",
"PH_TIN",
"PH_UMID",
"PH_PASSPORT",
"ZA_ID_NUMBER"
],
"title": "PiiEntityType",
"type": "string"

View file

@ -4,11 +4,12 @@ import os
from collections.abc import Callable, Mapping
from datetime import datetime
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, Literal, NamedTuple
from typing import TYPE_CHECKING, Annotated, Any, Final, Literal, NamedTuple
import httpx
from pydantic import (
BaseModel,
BeforeValidator,
ConfigDict,
Field,
Json,
@ -47,6 +48,7 @@ from litellm.types.proxy.carried_budget_state import (
)
from litellm.types.proxy.control_plane_endpoints import WorkerRegistryEntry
from litellm.types.router import RouterErrors, UpdateRouterConfig
from litellm.types.router_weights import validate_router_settings_dict
from litellm.types.secret_managers.main import KeyManagementSystem
from litellm.types.utils import (
CallTypes,
@ -284,6 +286,7 @@ class KeyManagementRoutes(str, enum.Enum):
# team's `team_member_permissions`, non-admin members of that team may set
# `access_group_ids` on keys they create/update. Default-deny.
KEY_ACCESS_GROUP_ASSIGNMENT = "/key/access_group_assignment"
AUTO_ROUTER_MANAGE = "/auto_router/manage"
# info and health routes
KEY_INFO = "/key/info"
@ -650,15 +653,18 @@ class LiteLLMRoutes(enum.Enum):
KeyManagementRoutes.KEY_RESET_SPEND.value,
KeyManagementRoutes.KEY_ALIASES.value,
KeyManagementRoutes.KEY_ACCESS_GROUP_ASSIGNMENT.value,
KeyManagementRoutes.AUTO_ROUTER_MANAGE.value,
]
management_routes = (
[
# user
"/user/new",
"/management/v1/users/bulk",
"/user/update",
"/user/bulk_update",
"/user/delete",
"/management/v1/users/bulk_delete",
"/user/info",
"/user/list",
"/user/daily/activity",
@ -838,6 +844,7 @@ class LiteLLMRoutes(enum.Enum):
self_managed_routes = [
"/team/member_add",
"/team/member_delete",
"/management/v1/teams/{team_id}/members/bulk_delete",
"/team/member_update",
"/team/{team_id}/member/{user_id}/reset_spend",
"/team/permissions_list",
@ -864,6 +871,7 @@ class LiteLLMRoutes(enum.Enum):
"/organization/daily/activity",
"/user/available_roles", # read-only role metadata; any authenticated user may read
"/user/list", # org admins checked in endpoint; non-admins get 403
"/management/v1/users/bulk_delete", # proxy admins delete anyone, org admins only their orgs' users; others 403
"/model/{model_id}/update",
"/prompt/list",
"/prompt/info",
@ -887,6 +895,7 @@ class LiteLLMRoutes(enum.Enum):
"/auto_router/validate_complexity_router_config",
# Per-session auto-router read - the endpoint scopes the row to the caller's own key hash
"/auto_router/session",
"/cost/predict-cache",
# Agent registry - reads are role-scoped and writes are proxy-admin-gated
# inside agent_endpoints/endpoints.py
*agent_management_routes,
@ -1197,6 +1206,7 @@ class AllowedVectorStoreIndexItem(LiteLLMPydanticObjectBase):
class KeyRequestBase(GenerateRequestBase):
key: str | None = None
tpd_limit: int | None = None
default_estimated_output_tokens: PositiveInt | None = None
default_estimated_output_tokens_per_model: Mapping[str, PositiveInt] | None = None
budget_id: str | None = None
@ -1882,6 +1892,9 @@ class BudgetNewRequest(LiteLLMPydanticObjectBase):
)
tpm_limit: int | None = Field(default=None, description="Max tokens per minute, allowed for this budget id.")
rpm_limit: int | None = Field(default=None, description="Max requests per minute, allowed for this budget id.")
tpd_limit: int | None = Field(
default=None, description="Max tokens per day, charged by batch submissions, allowed for this budget id."
)
budget_duration: str | None = Field(
default=None,
description="Max duration budget should be set for (e.g. '1hr', '1d', '28d')",
@ -1980,8 +1993,14 @@ class OrgMember(MemberBase):
from litellm.models.team import TeamBase as TeamBase # noqa: E402
RouterSettingsDict = Annotated[
dict[str, object],
BeforeValidator(validate_router_settings_dict, json_schema_input_type=UpdateRouterConfig),
]
class NewTeamRequest(TeamBase):
router_settings: RouterSettingsDict | None = None
model_aliases: dict | None = None
tags: list | None = None
guardrails: list[str] | None = None
@ -2052,6 +2071,7 @@ class UpdateTeamRequest(LiteLLMPydanticObjectBase):
metadata: dict | None = None
tpm_limit: int | None = None
rpm_limit: int | None = None
tpd_limit: int | None = None
max_budget: float | None = None
soft_budget: float | None = None
models: list | None = None
@ -2079,7 +2099,7 @@ class UpdateTeamRequest(LiteLLMPydanticObjectBase):
allowed_vector_store_indexes: list[AllowedVectorStoreIndexItem] | None = None
enforced_batch_output_expires_after: dict | None = None
enforced_file_expires_after: dict | None = None
router_settings: dict | None = None
router_settings: RouterSettingsDict | None = None
access_group_ids: list[str] | None = None
budget_limits: list[BudgetLimitEntry] | None = None # multiple concurrent budget windows
default_team_member_models: list[str] | None = None # default allowed_models seeded onto new team members
@ -3007,6 +3027,7 @@ class LiteLLM_VerificationTokenView(LiteLLM_VerificationToken):
team_alias: str | None = None
team_tpm_limit: int | None = None
team_rpm_limit: int | None = None
team_tpd_limit: int | None = None
team_max_budget: float | None = None
team_soft_budget: float | None = None
team_models: list = []
@ -3026,6 +3047,7 @@ class LiteLLM_VerificationTokenView(LiteLLM_VerificationToken):
end_user_id: str | None = None
end_user_tpm_limit: int | None = None
end_user_rpm_limit: int | None = None
end_user_tpd_limit: int | None = None
end_user_max_budget: float | None = None
end_user_model_max_budget: dict | None = None
@ -4701,6 +4723,7 @@ class JWTAuthBuilderResult(TypedDict):
org_id: str | None
team_membership: LiteLLM_TeamMembership | None
jwt_claims: dict # Decoded JWT token claims (avoids re-decoding)
agent_id: ReadOnly[str | None]
class ClientSideFallbackModel(TypedDict, total=False):
@ -4939,6 +4962,14 @@ class LiteLLM_JWTAuth(LiteLLMPydanticObjectBase):
user_allowed_roles: list[str] | None = None
user_id_upsert: bool = Field(default=False, description="If user doesn't exist, upsert them into the db.")
end_user_id_jwt_field: str | None = None
agent_id_jwt_field: str | None = Field(
default=None,
description=(
"The field in the JWT token that identifies the calling agent (e.g. 'azp' for a Microsoft Entra ID "
"app token). Supports dot notation. The value is matched against a registered agent's agent_id, "
"then agent_name, and the request is rejected when it matches neither."
),
)
public_key_ttl: float = 600
public_key_stale_ttl: float = Field(
default=DEFAULT_JWKS_STALE_TTL,

View file

@ -39,6 +39,7 @@ from litellm.constants import (
from litellm.litellm_core_utils.dd_tracing import tracer
from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider
from litellm.litellm_core_utils.safe_json_loads import safe_json_loads
from litellm.models.project import LiteLLM_ProjectTable
from litellm.proxy._types import (
RBAC_ROLES,
CallInfo,
@ -109,7 +110,7 @@ from litellm.proxy.utils import PrismaClient, ProxyLogging, log_db_metrics
from litellm.repositories.budget_repository import BudgetRepository
from litellm.repositories.object_permission_repository import ObjectPermissionRepository
from litellm.repositories.organization_repository import OrganizationRepository
from litellm.repositories.prisma_protocols import RowT_co
from litellm.repositories.prisma_protocols import DatabaseClient, RowT_co
from litellm.repositories.project_repository import ProjectRepository
from litellm.repositories.table_repositories import (
AccessGroupRepository,
@ -847,6 +848,7 @@ BUDGET_ENFORCED_SIDE_EFFECT_ROUTES: Final = frozenset(
"/health",
"/health/services",
"/health/test_connection",
"/auto_router/test_routing",
}
)
@ -895,6 +897,7 @@ async def common_checks(
request_query_params=_safe_get_request_query_params(request=request),
llm_router=llm_router,
request=request,
team_id=valid_token.team_id if valid_token is not None else None,
)
skip_all_budget_checks: Final = skip_budget_checks or (
@ -3171,7 +3174,7 @@ async def _delete_cache_access_object(
@log_db_metrics
async def get_access_object(
access_group_id: str,
prisma_client: PrismaClient | None,
prisma_client: DatabaseClient | None,
user_api_key_cache: UserApiKeyCache,
proxy_logging_obj: ProxyLogging | None = None,
) -> LiteLLM_AccessGroupTable:
@ -3917,7 +3920,7 @@ async def get_org_object(
async def _get_resources_from_access_groups(
access_group_ids: Sequence[str],
resource_field: Literal["access_model_names", "access_mcp_server_ids", "access_agent_ids"],
prisma_client: PrismaClient | None = None,
prisma_client: DatabaseClient | None = None,
user_api_key_cache: UserApiKeyCache | None = None,
proxy_logging_obj: ProxyLogging | None = None,
) -> list[str]:
@ -3975,7 +3978,7 @@ async def _get_resources_from_access_groups(
async def _get_models_from_access_groups(
access_group_ids: Sequence[str],
prisma_client: PrismaClient | None = None,
prisma_client: DatabaseClient | None = None,
user_api_key_cache: UserApiKeyCache | None = None,
proxy_logging_obj: ProxyLogging | None = None,
) -> list[str]:
@ -4471,9 +4474,10 @@ async def stamp_matched_model_access_groups(
async def can_key_call_model(
model: str | list[str],
llm_model_list: list | None,
llm_model_list: Sequence[object] | None,
valid_token: UserAPIKeyAuth,
llm_router: litellm.Router | None,
prisma_client: DatabaseClient | None = None,
) -> Literal[True]:
"""
Checks if token can call a given model
@ -4503,6 +4507,7 @@ async def can_key_call_model(
if key_access_group_ids:
models_from_groups: Final = await _get_models_from_access_groups(
access_group_ids=key_access_group_ids,
prisma_client=prisma_client,
)
if models_from_groups:
return _can_object_call_model(
@ -4518,7 +4523,7 @@ async def can_key_call_model(
async def can_key_call_resolved_model(
model: str,
llm_model_list: list | None,
llm_model_list: Sequence[object] | None,
valid_token: UserAPIKeyAuth,
llm_router: litellm.Router | None,
) -> None:
@ -4631,6 +4636,7 @@ async def can_team_access_model(
team_object: LiteLLM_TeamTable | None,
llm_router: Router | None,
team_model_aliases: dict[str, str] | None = None,
prisma_client: DatabaseClient | None = None,
) -> Literal[True]:
"""
Returns True if the team can access a specific model.
@ -4653,12 +4659,13 @@ async def can_team_access_model(
if team_access_group_ids:
models_from_groups: Final = await _get_models_from_access_groups(
access_group_ids=team_access_group_ids,
prisma_client=prisma_client,
)
if models_from_groups:
return _can_object_call_model(
model=model,
llm_router=llm_router,
models=models_from_groups,
models=list(dict.fromkeys([*(team_object.models if team_object else []), *models_from_groups])),
team_model_aliases=team_model_aliases,
team_id=team_object.team_id if team_object else None,
object_type="team",
@ -4748,7 +4755,7 @@ async def _key_access_group_grants_model(
def can_project_access_model(
model: str | list[str],
project_object: LiteLLM_ProjectTableCachedObj,
project_object: LiteLLM_ProjectTable,
llm_router: Router | None,
) -> Literal[True]:
"""
@ -5766,8 +5773,7 @@ async def _organization_max_budget_check(
if org_table.litellm_budget_table is not None:
org_max_budget = org_table.litellm_budget_table.max_budget
# Only check if organization has a valid max_budget set
if org_max_budget is None or org_max_budget <= 0:
if org_max_budget is None:
return
# Read spend from cross-pod counter (Redis-first) or cached object (fallback)

View file

@ -23,6 +23,7 @@ from litellm.proxy.auth.auth_utils import (
_get_request_ip_address,
is_invalid_virtual_key_error,
mark_invalid_virtual_key_error,
normalize_request_route,
)
from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler
from litellm.types.services import ServiceTypes
@ -172,7 +173,7 @@ class UserAPIKeyAuthExceptionHandler:
# so the handler is side-effect-free for the caller's identity object.
user_api_key_dict = resolved_identity.model_copy() if resolved_identity is not None else UserAPIKeyAuth()
user_api_key_dict.parent_otel_span = parent_otel_span
user_api_key_dict.request_route = route
user_api_key_dict.request_route = normalize_request_route(route)
user_api_key_dict.api_key = user_api_key_dict.api_key or UserAPIKeyAuth(api_key=api_key).api_key
# Stamp identity onto the request's server span now, before the request

View file

@ -33,7 +33,7 @@ from litellm.proxy.common_utils.http_parsing_utils import extract_nested_form_me
from litellm.types.passthrough_endpoints.pass_through_endpoints import (
LITELLM_PASS_THROUGH_ENDPOINT_MARKER,
)
from litellm.types.router import CONFIGURABLE_CLIENTSIDE_AUTH_PARAMS
from litellm.types.router import CONFIGURABLE_CLIENTSIDE_AUTH_PARAMS, Deployment
from litellm.types.utils import CustomPricingLiteLLMParams
@ -1736,7 +1736,7 @@ def _append_model_candidates(candidates: list[str], value: Any) -> None:
candidates.extend(model for model in model_names if model)
def _dedupe_model_candidates(candidates: list[str]) -> list[str]:
def _dedupe_model_candidates(candidates: Collection[str]) -> list[str]:
deduped: Final[list[str]] = []
for model in candidates:
if model not in deduped:
@ -1845,13 +1845,42 @@ def _resolve_model_id_with_router(model_id: str | None, llm_router: Router | Non
return model_id
def get_cache_prediction_deployments(
*, current_deployment_id: str, candidate_deployment_id: str, llm_router: Router, team_id: str | None
) -> tuple[Deployment, Deployment] | None:
current: Final = llm_router.get_deployment(current_deployment_id)
candidate: Final = llm_router.get_deployment(candidate_deployment_id)
if current is None or candidate is None:
return None
if any(deployment.model_info.team_id not in (None, team_id) for deployment in (current, candidate)):
return None
return current, candidate
def _cache_prediction_model_candidates(
request_data: Mapping[str, object], llm_router: Router | None, team_id: str | None
) -> tuple[str, ...]:
current_id: Final = request_data.get("current_deployment_id")
candidate_id: Final = request_data.get("candidate_deployment_id")
if llm_router is None or not isinstance(current_id, str) or not isinstance(candidate_id, str):
return ()
deployments: Final = get_cache_prediction_deployments(
current_deployment_id=current_id, candidate_deployment_id=candidate_id, llm_router=llm_router, team_id=team_id
)
return tuple(deployment.model_name for deployment in deployments) if deployments is not None else ()
def _extract_model_candidates_from_request(
request_data: dict,
route: str,
request_headers: Mapping[str, object] | None = None,
request_query_params: Mapping[str, object] | None = None,
llm_router: Router | None = None,
team_id: str | None = None,
) -> list[str]:
if route == "/cost/predict-cache":
prediction_models: Final = _cache_prediction_model_candidates(request_data, llm_router, team_id) # pyright: ignore[reportUnknownArgumentType] # the typed reader validates each deployment ID from this legacy payload
return _dedupe_model_candidates(prediction_models)
candidates: Final[list[str]] = []
uses_model_routing_sources: Final = _route_uses_model_routing_sources(route=route)
uses_header_or_query_model_sources: Final = _route_matches_any_marker(
@ -1945,6 +1974,7 @@ def get_model_from_request(
request_query_params: Mapping[str, object] | None = None,
llm_router: Router | None = None,
request: Request | None = None,
team_id: str | None = None,
) -> str | list[str] | None:
"""Resolve the model(s) a request targets, for model-access and budget checks.
@ -1967,6 +1997,7 @@ def get_model_from_request(
request_headers=request_headers,
request_query_params=request_query_params,
llm_router=llm_router,
team_id=team_id,
)
model = _format_model_candidates(candidates)

View file

@ -0,0 +1,136 @@
from __future__ import annotations
from collections.abc import Mapping
from typing import TYPE_CHECKING, Final
from pydantic import TypeAdapter, ValidationError
from litellm.litellm_core_utils.core_helpers import get_metadata_variable_name_from_kwargs
if TYPE_CHECKING:
from litellm.router import Router
_MAPPING_ADAPTER: Final = TypeAdapter(Mapping[str, object])
def _mapping(value: object) -> Mapping[str, object] | None:
try:
return _MAPPING_ADAPTER.validate_python(value)
except ValidationError:
return None
async def authorize_member_auto_router_inference(
*,
deployment: Mapping[str, object] | None,
request_kwargs: Mapping[str, object],
llm_router: Router,
) -> None:
if deployment is None:
return
model_info: Final = _mapping(deployment.get("model_info"))
if model_info is None or model_info.get("member_auto_router") is not True:
return
from fastapi import HTTPException
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
from litellm.proxy.auth.auth_checks import (
OrganizationNotFoundError,
TeamNotFoundError,
get_org_object,
get_project_object,
get_team_membership,
get_team_object,
)
from litellm.proxy.management_helpers.auto_router_permissions import (
MemberAutoRouterDependencyObjects,
authorize_member_auto_router_dependencies,
validate_member_auto_router_config,
)
metadata: Final = _mapping(request_kwargs.get(get_metadata_variable_name_from_kwargs(request_kwargs)))
actor: Final = metadata.get("user_api_key_auth") if metadata is not None else None
team_id: Final = model_info.get("team_id")
if not isinstance(actor, UserAPIKeyAuth) or not isinstance(team_id, str) or not team_id:
raise HTTPException(status_code=403, detail="Member auto-routers require authenticated team access")
if actor.team_id != team_id and actor.user_role != LitellmUserRoles.PROXY_ADMIN:
raise HTTPException(status_code=403, detail="This auto-router belongs to a different team")
from litellm.proxy.proxy_server import prisma_client, proxy_logging_obj, user_api_key_cache
if prisma_client is None:
raise HTTPException(status_code=503, detail="Cannot verify auto-router model access without a database")
try:
team: Final = await get_team_object(
team_id=team_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=actor.parent_otel_span,
proxy_logging_obj=proxy_logging_obj,
)
except TeamNotFoundError as error:
raise HTTPException(status_code=403, detail="The auto-router team no longer exists") from error
if (
actor.user_role != LitellmUserRoles.PROXY_ADMIN
and actor.user_id is not None
and (not actor.user_id or not any(member.user_id == actor.user_id for member in team.members_with_roles))
):
raise HTTPException(status_code=403, detail="You are no longer a member of this auto-router's team")
if team.blocked:
raise HTTPException(status_code=403, detail="This auto router's team is blocked.")
params: Final = _mapping(deployment.get("litellm_params"))
if params is None:
raise HTTPException(status_code=403, detail="The member auto-router configuration is invalid")
raw_config: Final = _mapping(params.get("complexity_router_config"))
if raw_config is None:
raise HTTPException(status_code=403, detail="The member auto-router configuration is invalid")
default_model: Final = params.get("complexity_router_default_model")
config: Final = validate_member_auto_router_config(raw_config)
membership: Final = (
await get_team_membership(
user_id=actor.user_id,
team_id=team_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=actor.parent_otel_span,
proxy_logging_obj=proxy_logging_obj,
)
if actor.user_id
else None
)
try:
organization: Final = (
await get_org_object(
org_id=team.organization_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=actor.parent_otel_span,
proxy_logging_obj=proxy_logging_obj,
)
if team.organization_id
else None
)
except OrganizationNotFoundError as error:
raise HTTPException(status_code=403, detail="The auto router's organization is unavailable.") from error
project: Final = (
await get_project_object(
project_id=actor.project_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
)
if actor.project_id
else None
)
await authorize_member_auto_router_dependencies(
config=config,
default_model=default_model if isinstance(default_model, str) else None,
user_api_key_dict=actor,
team=team,
prisma_client=None,
llm_router=llm_router,
dependency_objects=MemberAutoRouterDependencyObjects(
membership=membership, organization=organization, project=project
),
)

View file

@ -14,7 +14,7 @@ import hashlib
import os
import re
import time
from collections.abc import Awaitable, Callable, Sequence
from collections.abc import Awaitable, Callable, Mapping, Sequence
from typing import Any, Final, Literal, NoReturn, Protocol, TypeVar, cast
import httpx
@ -61,6 +61,7 @@ from litellm.proxy.common_utils.user_api_key_cache import (
)
from litellm.proxy.utils import PrismaClient, ProxyLogging
from litellm.repositories.user_repository import UserRepository
from litellm.types.agents import AgentResponse
from .auth_checks import (
_allowed_routes_check,
@ -127,6 +128,26 @@ class _UserInfoResponse(Protocol):
def json(self) -> dict[str, object]: ...
class AgentLookup(Protocol):
"""The registered-agent lookups a JWT agent claim is matched against."""
def get_agent_by_id(self, agent_id: str) -> AgentResponse | None:
"""The agent registered under ``agent_id``, if any."""
def get_agent_by_name(self, agent_name: str) -> AgentResponse | None:
"""The agent registered under ``agent_name``, if any."""
class _NoRegisteredAgents:
"""The lookup in force until the proxy binds its agent registry: no agent is registered, so no claim matches."""
def get_agent_by_id(self, agent_id: str) -> None:
return None
def get_agent_by_name(self, agent_name: str) -> None:
return None
def _discovery_document(response: _OIDCDiscoveryResponse) -> _OIDCDiscoveryBody:
"""Decode an OIDC discovery response body."""
return response.json()
@ -198,6 +219,10 @@ class JWTHandler:
self.leeway = 0
# Per-cache-key locks so a TTL lapse triggers one refresh instead of one per in-flight request.
self._refresh_locks: dict[str, asyncio.Lock] = {} # mutable-ok: lock registry, keyed by JWKS url
self.agent_lookup: AgentLookup = _NoRegisteredAgents()
def bind_agent_lookup(self, agent_lookup: AgentLookup) -> None:
self.agent_lookup = agent_lookup
def update_environment(
self,
@ -623,6 +648,12 @@ class JWTHandler:
object_id = default_value
return object_id
def get_agent_claim(self, token: Mapping[str, object]) -> str | None:
if self.litellm_jwtauth.agent_id_jwt_field is None:
return None
claim: Final[object] = get_nested_value(data=token, key_path=self.litellm_jwtauth.agent_id_jwt_field)
return claim if isinstance(claim, str) and claim else None
def get_org_id(self, token: dict, default_value: str | None) -> str | None:
if self._has_trusted_issuer_normalized_claim(token=token, claim=self.LITELLM_ORG_ID_CLAIM):
return token.get(self.LITELLM_ORG_ID_CLAIM)
@ -1380,6 +1411,7 @@ class JWTAuthManager:
api_key: str,
jwt_valid_token: dict | None = None,
user_email: str | None = None,
agent_id: str | None = None,
) -> JWTAuthBuilderResult | None:
"""Check admin status and route access permissions"""
if not jwt_handler.is_admin(scopes=scopes):
@ -1409,8 +1441,28 @@ class JWTAuthManager:
org_id=org_id,
team_membership=None,
jwt_claims=jwt_valid_token or {},
agent_id=agent_id,
)
@staticmethod
def resolve_agent_id(
jwt_handler: JWTHandler,
jwt_valid_token: Mapping[str, object],
agent_registry: AgentLookup,
) -> str | None:
agent_claim: Final = jwt_handler.get_agent_claim(token=jwt_valid_token)
if agent_claim is None:
return None
agent: Final = agent_registry.get_agent_by_id(agent_id=agent_claim) or agent_registry.get_agent_by_name(
agent_name=agent_claim
)
if agent is None:
raise HTTPException(
status_code=status.HTTP_403_FORBIDDEN,
detail=f"No registered agent matches JWT claim {jwt_handler.litellm_jwtauth.agent_id_jwt_field}={agent_claim}",
)
return agent.agent_id
@staticmethod
async def find_and_validate_specific_team_id(
jwt_handler: JWTHandler,
@ -2268,9 +2320,23 @@ class JWTAuthManager:
elif rbac_role == LitellmUserRoles.INTERNAL_USER:
user_id = object_id
agent_id: Final = JWTAuthManager.resolve_agent_id(
jwt_handler=jwt_handler,
jwt_valid_token=jwt_valid_token,
agent_registry=jwt_handler.agent_lookup,
)
# Check admin access
admin_result: Final = await JWTAuthManager.check_admin_access(
jwt_handler, scopes, route, user_id, org_id, api_key, jwt_valid_token, user_email=user_email
jwt_handler,
scopes,
route,
user_id,
org_id,
api_key,
jwt_valid_token,
user_email=user_email,
agent_id=agent_id,
)
if admin_result:
await JWTAuthManager._attach_team_from_header_for_admin(
@ -2514,4 +2580,5 @@ class JWTAuthManager:
token=api_key,
team_membership=team_membership_object,
jwt_claims=jwt_valid_token,
agent_id=agent_id,
)

View file

@ -249,6 +249,7 @@ async def authenticate_user(
if os.getenv("DATABASE_URL") is not None:
response = await generate_key_helper_fn(
llm_router=None,
request_type="key",
**{
"user_role": LitellmUserRoles.PROXY_ADMIN,
@ -324,6 +325,7 @@ async def authenticate_user(
await _rehash_password_if_needed(_user_row.user_id, password, _password)
if os.getenv("DATABASE_URL") is not None:
response = await generate_key_helper_fn(
llm_router=None,
request_type="key",
**{
"user_role": user_role,

View file

@ -1,5 +1,5 @@
import re
from collections.abc import Sequence
from collections.abc import Collection
from typing import Final
from fastapi import HTTPException, Request, status
@ -24,10 +24,13 @@ _PROXY_ADMIN_VIEW_ONLY_BLOCKED_ROUTES: Final = frozenset(
[
# user
"/user/new",
"/management/v1/users/bulk",
"/user/delete",
"/management/v1/users/bulk_delete",
"/user/bulk_update",
# team
"/team/new",
"/management/v1/teams/{team_id}/members/bulk_delete",
"/team/update",
"/team/delete",
"/team/block",
@ -587,7 +590,7 @@ class RouteChecks:
return False
@staticmethod
def check_route_access(route: str, allowed_routes: Sequence[str]) -> bool:
def check_route_access(route: str, allowed_routes: Collection[str]) -> bool:
"""
Check if a route has access by checking both exact matches and patterns
@ -758,9 +761,12 @@ class RouteChecks:
_ADMIN_VIEWER_BLOCKED_WRITE_ROUTES = frozenset(
[
"/user/new",
"/management/v1/users/bulk",
"/user/delete",
"/management/v1/users/bulk_delete",
"/user/bulk_update",
"/team/new",
"/management/v1/teams/{team_id}/members/bulk_delete",
"/team/update",
"/team/delete",
"/model/new",
@ -824,7 +830,7 @@ class RouteChecks:
status_code=status.HTTP_403_FORBIDDEN,
detail=f"user not allowed to access this route, role= {_user_role}. Trying to access: {route} and updating invalid param: {param}. only user_email and password can be updated",
)
elif route in _PROXY_ADMIN_VIEW_ONLY_BLOCKED_ROUTES or (
elif RouteChecks.check_route_access(route=route, allowed_routes=_PROXY_ADMIN_VIEW_ONLY_BLOCKED_ROUTES) or (
route.startswith("/key/") and route.endswith(_PROXY_ADMIN_VIEW_ONLY_BLOCKED_KEY_SUFFIXES)
):
# Block write operations for PROXY_ADMIN_VIEW_ONLY
@ -859,9 +865,9 @@ class RouteChecks:
# Hard-block known write routes regardless of HTTP method (defensive
# — these are POSTs in practice, but pinning them here protects
# against future GET-shaped writes).
if route in RouteChecks._ADMIN_VIEWER_BLOCKED_WRITE_ROUTES or (
route.startswith("/key/") and route.endswith("/regenerate")
):
if RouteChecks.check_route_access(
route=route, allowed_routes=RouteChecks._ADMIN_VIEWER_BLOCKED_WRITE_ROUTES
) or (route.startswith("/key/") and route.endswith("/regenerate")):
raise HTTPException(
status_code=status.HTTP_403_FORBIDDEN,
detail=f"user not allowed to access this route, role= {_user_role}. Trying to access: {route}",

View file

@ -56,6 +56,7 @@ class TeamGrants(TypedDict, total=False):
team_alias: ReadOnly[str | None]
team_tpm_limit: ReadOnly[int | None]
team_rpm_limit: ReadOnly[int | None]
team_tpd_limit: ReadOnly[int | None]
team_max_budget: ReadOnly[float | None]
team_soft_budget: ReadOnly[float | None]
team_spend: ReadOnly[float | None]
@ -97,6 +98,7 @@ def team_grants(
team_alias=team_object.team_alias,
team_tpm_limit=team_object.tpm_limit,
team_rpm_limit=team_object.rpm_limit,
team_tpd_limit=team_object.tpd_limit,
team_max_budget=team_object.max_budget,
team_soft_budget=team_object.soft_budget,
team_spend=team_object.spend,

View file

@ -191,6 +191,7 @@ def _get_model_from_request_context(
route: str,
request: Request | None,
llm_router: Any | None = None,
team_id: str | None = None,
) -> str | list[str] | None:
return get_model_from_request(
request_data=request_data,
@ -199,6 +200,7 @@ def _get_model_from_request_context(
request_query_params=_safe_get_request_query_params(request=request),
llm_router=llm_router,
request=request,
team_id=team_id,
)
@ -217,7 +219,7 @@ async def _normalize_claude_model(
return
if request is not None and request.scope.get(_CLAUDE_MODEL_NORMALIZED) is True:
return
requested: Final = _get_model_from_request_context(request_data, route, request, llm_router)
requested: Final = _get_model_from_request_context(request_data, route, request, llm_router, valid_token.team_id)
if not isinstance(requested, str) or requested != request_data.get("model"):
return
if not requested.startswith("claude-router-") and not requested.lower().endswith("[1m]"):
@ -535,6 +537,9 @@ def _apply_budget_limits_to_end_user_params(
if budget_info.rpm_limit is not None:
end_user_params["end_user_rpm_limit"] = budget_info.rpm_limit
if budget_info.tpd_limit is not None:
end_user_params["end_user_tpd_limit"] = budget_info.tpd_limit
if budget_info.max_budget is not None:
end_user_params["end_user_max_budget"] = budget_info.max_budget
@ -619,6 +624,8 @@ def update_valid_token_with_end_user_params(valid_token: UserAPIKeyAuth, end_use
valid_token.end_user_tpm_limit = end_user_params["end_user_tpm_limit"]
if end_user_params.get("end_user_rpm_limit") is not None:
valid_token.end_user_rpm_limit = end_user_params["end_user_rpm_limit"]
if end_user_params.get("end_user_tpd_limit") is not None:
valid_token.end_user_tpd_limit = end_user_params["end_user_tpd_limit"]
if end_user_params.get("allowed_model_region") is not None:
valid_token.allowed_model_region = end_user_params["allowed_model_region"]
if end_user_params.get("end_user_model_max_budget") is not None:
@ -850,6 +857,7 @@ async def _auto_register_jwt_mapping(
user_id: str | None = None,
org_id: str | None = None,
end_user_id: str | None = None,
agent_id: str | None = None,
) -> UserAPIKeyAuth | None:
"""
Auto-register: create a new virtual key + mapping for an unrecognised JWT
@ -876,11 +884,13 @@ async def _auto_register_jwt_mapping(
# the NOT NULL @id constraint. Every successful key-creation caller (e.g.
# /key/generate) passes table_name="key" explicitly.
key_data: Final = await generate_key_helper_fn(
llm_router=None,
request_type="key",
table_name="key",
team_id=team_id,
user_id=user_id,
organization_id=org_id,
agent_id=agent_id,
metadata={
"auto_registered": True,
"jwt_claim_field": virtual_key_claim_field,
@ -1564,6 +1574,7 @@ async def _user_api_key_auth_builder(
org_id: Final = result["org_id"]
team_membership: Final[LiteLLM_TeamMembership | None] = result.get("team_membership", None)
jwt_claims = result.get("jwt_claims", None)
agent_id: Final[str | None] = result.get("agent_id")
if is_proxy_admin:
# Proxy admins authenticate via auth_builder (full
@ -1589,6 +1600,7 @@ async def _user_api_key_auth_builder(
end_user_id=end_user_id,
parent_otel_span=parent_otel_span,
jwt_claims=jwt_claims,
agent_id=agent_id,
**team_grants(team_object=team_object, team_membership=team_membership, user_id=user_id),
)
@ -1609,6 +1621,7 @@ async def _user_api_key_auth_builder(
user_rpm_limit=(user_object.rpm_limit if user_object is not None else None),
user_model_max_budget=(user_object.model_max_budget if user_object is not None else None),
jwt_claims=jwt_claims,
agent_id=agent_id,
**team_grants(team_object=team_object, team_membership=team_membership, user_id=user_id),
)
@ -1632,6 +1645,7 @@ async def _user_api_key_auth_builder(
user_id=user_id,
org_id=org_id,
end_user_id=end_user_id,
agent_id=agent_id,
)
if auto_registered is not None:
auto_registered.jwt_claims = jwt_claims
@ -1652,6 +1666,7 @@ async def _user_api_key_auth_builder(
route=route,
request=request,
llm_router=llm_router,
team_id=valid_token.team_id,
)
skip_budget_checks = False
if model is not None and llm_router is not None:
@ -1692,6 +1707,7 @@ async def _user_api_key_auth_builder(
route=route,
request=request,
llm_router=llm_router,
team_id=valid_token.team_id,
)
),
)
@ -2015,6 +2031,7 @@ async def _user_api_key_auth_builder(
valid_token.end_user_id = end_user_params.get("end_user_id")
valid_token.end_user_tpm_limit = end_user_params.get("end_user_tpm_limit")
valid_token.end_user_rpm_limit = end_user_params.get("end_user_rpm_limit")
valid_token.end_user_tpd_limit = end_user_params.get("end_user_tpd_limit")
valid_token.allowed_model_region = end_user_params.get("allowed_model_region")
if valid_token is not None:
@ -2091,6 +2108,7 @@ async def _user_api_key_auth_builder(
route=route,
request=request,
llm_router=llm_router,
team_id=valid_token.team_id,
)
skip_budget_checks = False
if model is not None and llm_router is not None:
@ -2209,6 +2227,7 @@ async def _user_api_key_auth_builder(
route=route,
request=request,
llm_router=llm_router,
team_id=valid_token.team_id,
)
current_models = _get_model_names_for_budget_checks(model=current_model)
@ -2239,6 +2258,7 @@ async def _user_api_key_auth_builder(
route=route,
request=request,
llm_router=llm_router,
team_id=valid_token.team_id,
)
current_models = _get_model_names_for_budget_checks(model=current_model)
@ -2288,6 +2308,7 @@ async def _user_api_key_auth_builder(
spend=valid_token.team_spend,
tpm_limit=valid_token.team_tpm_limit,
rpm_limit=valid_token.team_rpm_limit,
tpd_limit=valid_token.team_tpd_limit,
blocked=valid_token.team_blocked,
models=token_team_models,
metadata=valid_token.team_metadata,
@ -2441,6 +2462,7 @@ def _team_obj_from_token(valid_token: UserAPIKeyAuth) -> LiteLLM_TeamTableCached
spend=valid_token.team_spend,
tpm_limit=valid_token.team_tpm_limit,
rpm_limit=valid_token.team_rpm_limit,
tpd_limit=valid_token.team_tpd_limit,
blocked=valid_token.team_blocked,
models=token_team_models,
metadata=valid_token.team_metadata,
@ -2482,7 +2504,7 @@ def _token_can_vouch_for_team(valid_token: UserAPIKeyAuth, lookup_error: BaseExc
async def _run_centralized_common_checks(
user_api_key_auth_obj: UserAPIKeyAuth,
request: Request,
request_data: dict,
request_data: dict[str, object],
route: str,
) -> None:
"""Run ``common_checks`` once at the ``user_api_key_auth`` wrapper
@ -2734,6 +2756,7 @@ async def _run_centralized_common_checks(
route=route,
request=request,
llm_router=llm_router,
team_id=user_api_key_auth_obj.team_id,
)
# Pin the metadata variable name (litellm_metadata vs metadata) before
@ -2850,12 +2873,14 @@ def _should_skip_budget_checks(
route: str,
request: Request | None,
llm_router: Any | None,
team_id: str | None = None,
) -> bool:
model: Final = _get_model_from_request_context(
request_data=request_data,
route=route,
request=request,
llm_router=llm_router,
team_id=team_id,
)
if model is not None and llm_router is not None:
return _is_model_cost_zero(model=model, llm_router=llm_router)
@ -3301,6 +3326,7 @@ async def _enforce_key_and_fallback_model_access(
route=route,
request=request,
llm_router=llm_router,
team_id=valid_token.team_id,
)
if model is not None:
@ -3408,6 +3434,7 @@ async def _run_post_custom_auth_checks(
route=route,
request=request,
llm_router=llm_router,
team_id=valid_token.team_id,
)
current_models = _get_model_names_for_budget_checks(model=current_model)
@ -3449,6 +3476,7 @@ async def _run_post_custom_auth_checks(
route=route,
request=request,
llm_router=llm_router,
team_id=valid_token.team_id,
)
current_models = _get_model_names_for_budget_checks(model=current_model)

View file

@ -580,8 +580,8 @@ What the command changed is recorded in `~/.litellm/claude_configure_state.json`
`lite configure claude`, `lite login --config-claude`, `lite up` and `lite autoroute up` also install a status line (`~/.litellm/statusline.py`, registered as `statusLine` in `~/.claude/settings.json` unless you already run one) that shows which model the auto-router actually served the last turn and, once the proxy has recorded the session, what the session cost against the router's savings baseline:
```
claude-auto · Routed to: claude-haiku-4-5 -63% vs Claude Opus 5
LiteLLM ████████░░░░░░░░░░░░░░░░ $0.14
Routed to: claude-haiku-4-5 -63% vs Claude Opus 5
claude-auto ████████░░░░░░░░░░░░░░░░ $0.14
Claude Opus 5 ████████████████████████ $0.38
```

View file

@ -10,7 +10,7 @@ import os
import tempfile
from collections.abc import Callable, Mapping
from dataclasses import dataclass
from enum import StrEnum
from enum import Enum
from pathlib import Path
from types import MappingProxyType
from typing import Annotated, Final
@ -25,7 +25,7 @@ LITELLM_PROXY_API_KEY_ENV: Final = "LITELLM_PROXY_API_KEY"
_REJECTED_STATUSES: Final = frozenset((401, 403))
class ListingFailure(StrEnum):
class ListingFailure(str, Enum):
"""Why a proxy could not be listed, decided once where the HTTP outcome is classified.
`unreachable` means no response at all; the other kinds prove the proxy answered, so callers

View file

@ -28,6 +28,7 @@ import os
import sys
import tempfile
import time
import unicodedata
import urllib.error
import urllib.request
from collections.abc import Callable, Mapping
@ -42,7 +43,6 @@ FETCH_TIMEOUT_SECONDS: Final = 3
BAR_WIDTH: Final = 24
BAR_FULL: Final = "\u2588"
BAR_EMPTY: Final = "\u2591"
SEPARATOR: Final = " \u00b7 "
TRANSCRIPT_SCAN_LIMIT_BYTES: Final = 4 * 1024 * 1024
CLAUDE_BASE_URL_ENV_KEYS: Final = ("ANTHROPIC_BASE_URL",)
CLAUDE_API_KEY_ENV_KEYS: Final = ("ANTHROPIC_AUTH_TOKEN", "ANTHROPIC_API_KEY")
@ -50,7 +50,6 @@ CODEX_BASE_URL_ENV_KEYS: Final = ("OPENAI_BASE_URL",)
CODEX_API_KEY_ENV_KEYS: Final = ("OPENAI_API_KEY",)
CODEX_STOP_EVENT: Final = "Stop"
SYNTHETIC_MODEL: Final = "<synthetic>"
LITELLM_LABEL: Final = "LiteLLM"
RESET: Final = "\033[0m"
BOLD: Final = "\033[1m"
DIM: Final = "\033[90m"
@ -302,31 +301,37 @@ def _bar(fraction: float, color: str, width: int, use_color: bool) -> str:
return f"{color}{BAR_FULL * filled}{DIM}{BAR_EMPTY * (width - filled)}{RESET}"
def _display_width(label: str) -> int:
return sum(
2 if unicodedata.east_asian_width(character) in ("W", "F") else 1
for character in label
if unicodedata.category(character) not in ("Mn", "Me")
)
def render(model: str, session: Session | None, config_dir: Path, use_color: bool, bar_width: int = BAR_WIDTH) -> str:
def paint(code: str, text: str) -> str:
return f"{code}{text}{RESET}" if use_color else text
routed: Final = paint(BOLD, f"Routed to: {model}")
if session is None:
if session is None or session.baseline_model is None or session.baseline_spend <= 0:
return routed
header: Final = f"{session.router_name}{SEPARATOR}{routed}"
if session.baseline_model is None or session.baseline_spend <= 0:
return header
reference: Final = baseline_label(session.baseline_model, config_dir)
pct: Final = (session.baseline_spend - session.spend) / session.baseline_spend * 100
delta: Final = paint(LITELLM_COLOR, f"{'-' if pct >= 0 else '+'}{abs(round(pct))}% vs {reference}")
peak: Final = max(session.spend, session.baseline_spend)
label_width: Final = max(len(LITELLM_LABEL), len(reference))
label_width: Final = max(_display_width(session.router_name), _display_width(reference))
rows: Final = (
(LITELLM_LABEL, session.spend, LITELLM_COLOR),
(session.router_name, session.spend, LITELLM_COLOR),
(reference, session.baseline_spend, BASELINE_COLOR),
)
lines: Final = (
f"{paint(DIM, label.ljust(label_width))} {_bar(amount / peak, color, bar_width, use_color)} "
f"{paint(DIM, label + ' ' * (label_width - _display_width(label)))} "
f"{_bar(amount / peak, color, bar_width, use_color)} "
f"{paint(DIM, f'${amount:.2f}')}"
for label, amount, color in rows
)
return "\n".join((f"{header} {delta}", *lines))
return "\n".join((f"{routed} {delta}", *lines))
def color_enabled(env: Mapping[str, str]) -> bool:

View file

@ -14,6 +14,7 @@ import httpx
import orjson
from fastapi import HTTPException, Request, status
from fastapi.responses import JSONResponse, Response, StreamingResponse
from pydantic import ValidationError
from starlette.types import Receive, Scope, Send
import litellm
@ -76,6 +77,7 @@ from litellm.router_utils.add_retry_fallback_headers import get_hidden_params_di
from litellm.router_utils.common_utils import resolve_model_group_alias
from litellm.types.guardrails import GuardrailEventHooks
from litellm.types.router import RouterRateLimitError
from litellm.types.router_weights import validate_router_weights
_LateResponseT = TypeVar("_LateResponseT", bound=Response)
_LlmCallT = TypeVar("_LlmCallT")
@ -1571,6 +1573,9 @@ class ProxyBaseLLMRequestProcessing:
) -> dict:
exclude_values: Final = {"", None, "None"}
hidden_params = hidden_params or {}
resolved_call_id: Final = (
call_id or hidden_params.get("litellm_call_id") or (request_data or {}).get("litellm_call_id")
)
timing_values: Final = _timing_values(
hidden_params=hidden_params,
logging_obj=litellm_logging_obj,
@ -1598,7 +1603,7 @@ class ProxyBaseLLMRequestProcessing:
classifier_cost: Final = _classifier_cost_from_request_data(request_data)
headers: Final = {
"x-litellm-call-id": call_id,
"x-litellm-call-id": resolved_call_id,
"x-litellm-model-id": model_id,
"x-litellm-model-name": model_name,
"x-litellm-cache-key": cache_key,
@ -1936,6 +1941,13 @@ class ProxyBaseLLMRequestProcessing:
# This avoids expensive Router instantiation on each request
if router_settings is not None:
self.data["router_settings_override"] = router_settings
try:
self.data["_router_weights"] = validate_router_weights(router_settings.get("weights"))
except ValidationError:
self.data["_router_weights"] = None
verbose_proxy_logger.warning(
"Ignoring invalid saved router weights; update team/key router_settings"
)
alias_target: Final = await _resolve_per_request_model_group_alias(
requested_model=self.data.get("model"),
router_settings=router_settings,

View file

@ -124,7 +124,7 @@ def decrypt_value_helper(
key: str, # this is just for debug purposes, showing the k,v pair that's invalid. not a signing key.
exception_type: Literal["debug", "error"] = "error",
return_original_value: bool = False,
):
) -> str | None:
signing_key: Final = _get_salt_key()
try:

View file

@ -8,6 +8,8 @@ from typing import Final
from fastapi import status
from litellm.constants import STRINGIFIED_NONE
_OPENAI_ERROR_TYPE_BY_STATUS: Final[Mapping[int, str]] = MappingProxyType(
{
status.HTTP_401_UNAUTHORIZED: "authentication_error",
@ -35,7 +37,7 @@ def openai_error_type(exc: object, status_code: int) -> str:
"""OpenAI types ``error.type`` as a required string, so an exception carrying none
falls back to the type its status code stands for."""
carried: Final = attribute_of(exc, "type")
if isinstance(carried, str):
if isinstance(carried, str) and carried != STRINGIFIED_NONE:
return carried
mapped: Final = _OPENAI_ERROR_TYPE_BY_STATUS.get(status_code)
if mapped is not None:
@ -49,4 +51,4 @@ def openai_error_param(exc: object) -> str | None:
"""OpenAI types ``error.param`` as nullable, so an exception carrying none
serializes as JSON ``null``."""
carried: Final = attribute_of(exc, "param")
return carried if isinstance(carried, str) else None
return carried if isinstance(carried, str) and carried != STRINGIFIED_NONE else None

View file

@ -0,0 +1,91 @@
from collections.abc import Mapping
from math import isfinite
from typing import Final
from pydantic import TypeAdapter
import litellm
from litellm.cost_calculator import (
_select_model_name_for_cost_calc, # pyright: ignore[reportPrivateUsage] # shares completion_cost's deployment tariff selection
completion_cost, # pyright: ignore[reportUnknownVariableType] # legacy optional parameters are untyped
)
from litellm.litellm_core_utils.litellm_logging import Logging
from litellm.types.management_endpoints.prompt_cache_prediction import CacheTokenBuckets
from litellm.types.utils import CacheCreationTokenDetails, ModelResponse, PromptTokensDetailsWrapper, Usage
_PRICE_ENTRY: Final = TypeAdapter(Mapping[str, object])
def _valid_price(value: object) -> bool:
return isinstance(value, (int, float)) and not isinstance(value, bool) and isfinite(value) and value >= 0
def _has_required_prices(prices: Mapping[str, object], tokens: CacheTokenBuckets) -> bool:
required: Final = (
("input_cost_per_token", True),
("cache_read_input_token_cost", tokens.cache_read_input_tokens > 0),
("cache_creation_input_token_cost", tokens.cache_creation_5m_input_tokens > 0),
("cache_creation_input_token_cost_above_1hr", tokens.cache_creation_1h_input_tokens > 0),
)
if any(needed and not _valid_price(prices.get(key)) for key, needed in required):
return False
return all(
_valid_price(value)
for key, value in prices.items()
if value is not None and any(needed and key.startswith(f"{base}_above_") for base, needed in required)
)
def price_cache_tokens(model: str, deployment_id: str, tokens: CacheTokenBuckets) -> float | None:
try:
selected_model: Final = _select_model_name_for_cost_calc(
model=model,
completion_response=None,
custom_pricing=True,
custom_llm_provider="anthropic",
router_model_id=deployment_id,
)
if selected_model is None:
return None
model_info: Final = litellm.get_model_info(model=selected_model, custom_llm_provider="anthropic")
registry: Final = _PRICE_ENTRY.validate_python(litellm.model_cost) # pyright: ignore[reportUnknownMemberType] # legacy registry is validated at this boundary
price_entry: Final = registry.get(model_info["key"])
if price_entry is None:
return None
prices: Final = _PRICE_ENTRY.validate_python(price_entry)
if not _has_required_prices(prices, tokens):
return None
usage: Final = Usage(
prompt_tokens=tokens.total_tokens,
completion_tokens=0,
total_tokens=tokens.total_tokens,
prompt_tokens_details=PromptTokensDetailsWrapper(
cached_tokens=tokens.cache_read_input_tokens,
cache_creation_tokens=tokens.cache_creation_5m_input_tokens + tokens.cache_creation_1h_input_tokens,
cache_creation_token_details=CacheCreationTokenDetails(
ephemeral_5m_input_tokens=tokens.cache_creation_5m_input_tokens,
ephemeral_1h_input_tokens=tokens.cache_creation_1h_input_tokens,
),
),
)
logging_obj: Final = Logging(
model=model,
messages=[], # mutable-ok: Logging requires a list
stream=False,
call_type="completion",
start_time=None,
litellm_call_id="prompt-cache-prediction",
function_id="prompt-cache-prediction",
)
completion_cost(
completion_response=ModelResponse(model=model, usage=usage),
model=model,
custom_llm_provider="anthropic",
custom_pricing=True,
router_model_id=deployment_id,
litellm_logging_obj=logging_obj,
)
cost: Final = logging_obj.cost_breakdown.get("input_cost") if logging_obj.cost_breakdown is not None else None
return cost if cost is not None and _valid_price(cost) else None
except Exception: # noqa: BLE001 # the shared pricing owners raise plain Exception for unpriceable models
return None

View file

@ -26,7 +26,7 @@ class ComplianceChecker:
def __init__(self, data: ComplianceCheckRequest):
self.data = data
self.guardrails = data.guardrail_information or []
self.guardrails = tuple(g for g in data.guardrail_information or () if g.get("guardrail_status") != "not_run")
def _get_guardrails_by_mode(self, mode: str) -> list[dict]:
"""

View file

@ -2,25 +2,31 @@
CRUD endpoints for storing reusable credentials.
"""
from collections.abc import Mapping
from typing import (
Annotated,
Final,
cast, # noqa: TID251 # jsonify_object in proxy/utils.py is annotated with a bare dict
)
from fastapi import APIRouter, Depends, HTTPException, Path, Request, Response
from pydantic import TypeAdapter
import litellm
from litellm._logging import verbose_proxy_logger
from litellm.litellm_core_utils.credential_accessor import CredentialAccessor
from litellm.litellm_core_utils.litellm_logging import _get_masked_values
from litellm.models.credentials import UpdateCredentialItem
from litellm.proxy._types import CommonProxyErrors, UserAPIKeyAuth
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_value_helper
from litellm.proxy.utils import handle_exception_on_proxy, jsonify_object
from litellm.repositories.base_repository import is_unique_violation
from litellm.repositories.credentials_repository import CredentialsRepository
from litellm.types.utils import CreateCredentialItem, CredentialItem
router: Final = APIRouter()
_CREDENTIAL_DICT_ADAPTER: Final = TypeAdapter(dict[str, object])
class CredentialHelperUtils:
@ -40,6 +46,33 @@ class CredentialHelperUtils:
)
def _credential_exists_detail(credential_name: str) -> str:
return (
f"Credential '{credential_name}' already exists. "
f"Update it with PATCH /credentials/{credential_name}, or delete it first."
)
def get_llm_router() -> litellm.Router | None:
from litellm.proxy.proxy_server import llm_router
return llm_router
def _resolve_deployment_credentials(llm_router: litellm.Router | None, model_id: str) -> Mapping[str, object]:
if llm_router is None:
raise HTTPException(
status_code=500,
detail="LLM router not found. Please ensure you have a valid router instance.",
)
if llm_router.get_deployment(model_id) is None:
raise HTTPException(status_code=404, detail="Model not found")
credential_values: Final = llm_router.get_deployment_credentials(model_id)
if credential_values is None:
raise HTTPException(status_code=404, detail="Model not found")
return _CREDENTIAL_DICT_ADAPTER.validate_python(credential_values)
@router.post(
"/credentials",
dependencies=[Depends(user_api_key_auth)],
@ -50,13 +83,14 @@ async def create_credential(
fastapi_response: Response,
credential: CreateCredentialItem,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
llm_router: Annotated[litellm.Router | None, Depends(get_llm_router)] = None,
):
"""
[BETA] endpoint. This might change unexpectedly.
Stores credential in DB.
Reloads credentials in memory.
"""
from litellm.proxy.proxy_server import llm_router, prisma_client
from litellm.proxy.proxy_server import prisma_client
try:
if prisma_client is None:
@ -64,29 +98,19 @@ async def create_credential(
status_code=500,
detail={"error": CommonProxyErrors.db_not_connected_error.value},
)
if credential.model_id:
if llm_router is None:
raise HTTPException(
status_code=500,
detail="LLM router not found. Please ensure you have a valid router instance.",
)
# get model from router
model: Final = llm_router.get_deployment(credential.model_id)
if model is None:
raise HTTPException(status_code=404, detail="Model not found")
credential_values: Final = llm_router.get_deployment_credentials(credential.model_id)
if credential_values is None:
raise HTTPException(status_code=404, detail="Model not found")
credential.credential_values = credential_values
if credential.credential_values is None:
credential_values: Final = (
_resolve_deployment_credentials(llm_router, credential.model_id)
if credential.model_id
else credential.credential_values
)
if credential_values is None:
raise HTTPException(
status_code=400,
detail="Credential values are required. Unable to infer credential values from model ID.",
)
processed_credential: Final = CredentialItem(
credential_name=credential.credential_name,
credential_values=credential.credential_values,
credential_values=_CREDENTIAL_DICT_ADAPTER.validate_python(credential_values),
credential_info=credential.credential_info,
)
encrypted_credential: Final = CredentialHelperUtils.encrypt_credential_values(processed_credential)
@ -94,13 +118,18 @@ async def create_credential(
credentials_dict_jsonified: Final = cast( # cast-ok: deep-copies a model_dump, so keys are str
"dict[str, object]", jsonify_object(credentials_dict)
)
await CredentialsRepository(prisma_client).create(
data={
**credentials_dict_jsonified,
"created_by": user_api_key_dict.user_id,
"updated_by": user_api_key_dict.user_id,
}
)
try:
await CredentialsRepository(prisma_client).create(
data={
**credentials_dict_jsonified,
"created_by": user_api_key_dict.user_id,
"updated_by": user_api_key_dict.user_id,
}
)
except Exception as e:
if not is_unique_violation(e):
raise
raise HTTPException(status_code=409, detail=_credential_exists_detail(credential.credential_name))
## ADD TO LITELLM ##
CredentialAccessor.upsert_credentials([processed_credential])
@ -300,9 +329,10 @@ def update_db_credential(
async def update_credential(
request: Request,
fastapi_response: Response,
credential: CredentialItem,
credential: UpdateCredentialItem,
credential_name: str = Path(..., description="The credential name, percent-decoded; may contain slashes"),
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
llm_router: Annotated[litellm.Router | None, Depends(get_llm_router)] = None,
):
"""
[BETA] endpoint. This might change unexpectedly.
@ -319,7 +349,16 @@ async def update_credential(
db_credential: Final = await credentials_repository.find_by_name(credential_name)
if db_credential is None:
raise HTTPException(status_code=404, detail="Credential not found in DB.")
merged_credential: Final = update_db_credential(db_credential, credential)
patch: Final = CredentialItem(
credential_name=credential.credential_name,
credential_info=_CREDENTIAL_DICT_ADAPTER.validate_python(credential.credential_info),
credential_values=_CREDENTIAL_DICT_ADAPTER.validate_python(
_resolve_deployment_credentials(llm_router, credential.model_id)
if credential.model_id
else credential.credential_values or {}
),
)
merged_credential: Final = update_db_credential(db_credential, patch)
credential_object_jsonified: Final = cast( # cast-ok: deep-copies a model_dump, so keys are str
"dict[str, object]", jsonify_object(merged_credential.model_dump())
)
@ -341,11 +380,11 @@ async def update_credential(
if existing_in_memory is not None:
in_memory_values: Final = dict(existing_in_memory.credential_values or {})
if credential.credential_values:
in_memory_values.update(credential.credential_values)
if patch.credential_values:
in_memory_values.update(patch.credential_values)
in_memory_info: Final = dict(existing_in_memory.credential_info or {})
if credential.credential_info:
in_memory_info.update(credential.credential_info)
if patch.credential_info:
in_memory_info.update(patch.credential_info)
updated_in_memory: Final = CredentialItem(
credential_name=new_name,
credential_values=in_memory_values,

View file

@ -80,6 +80,7 @@ async def create_missing_views(db: SupportsRawQueries) -> None:
t.max_budget AS team_max_budget,
t.tpm_limit AS team_tpm_limit,
t.rpm_limit AS team_rpm_limit,
t.tpd_limit AS team_tpd_limit,
p.project_alias AS project_alias
FROM "LiteLLM_VerificationToken" v
LEFT JOIN "LiteLLM_TeamTable" t ON v.team_id = t.team_id

View file

@ -74,10 +74,14 @@ class LatestHealthCheckRow(BaseModel):
_ROWS_ADAPTER: Final = TypeAdapter(tuple[LatestHealthCheckRow, ...])
async def query_latest_health_checks(prisma_client: PrismaClient) -> tuple[LatestHealthCheckRow, ...]:
rows: Final = await prisma_client.db.query_raw(LATEST_HEALTH_CHECKS_SQL)
return _ROWS_ADAPTER.validate_python(rows)
async def fetch_latest_health_checks(prisma_client: PrismaClient) -> tuple[LatestHealthCheckRow, ...]:
try:
rows: Final = await prisma_client.db.query_raw(LATEST_HEALTH_CHECKS_SQL)
return _ROWS_ADAPTER.validate_python(rows)
return await query_latest_health_checks(prisma_client)
except Exception as query_err: # noqa: BLE001 # health decorates other reads; a driver error must not fail them
verbose_proxy_logger.error("Error getting all latest health checks: %s", query_err)
return ()

View file

@ -81,6 +81,11 @@ class WriterPinnedClient:
self.db: Final = db.writer if isinstance(db, RoutingPrismaWrapper) and not db.writer_unavailable else db
def writer_wrapper(db: "PrismaWrapper | RoutingPrismaWrapper") -> PrismaWrapper:
"""Unlike `WriterPinnedClient`, ignores `writer_unavailable`: a raw SQL write has no replica fallback."""
return db.writer if isinstance(db, RoutingPrismaWrapper) else db
class RoutingPrismaWrapper:
"""
Routes Prisma operations between a writer and a reader Prisma client.

View file

@ -244,6 +244,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
prompt_attack_threshold: float | None = 0.5,
pii_confidence_threshold: float | None = 0.5,
chunk_budget_chars: int = BEDROCK_APPLY_GUARDRAIL_CHUNK_BUDGET_CHARS,
contextual_grounding_from_messages: bool = False,
streaming_buffer_until_moderated: bool | None = None,
streaming_sampling_rate: int | None = None,
streaming_end_of_stream_only: bool | None = None,
@ -265,6 +266,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
self.guardrailVersion = guardrailVersion
self.guardrail_provider = "bedrock"
self.chunk_budget_chars = chunk_budget_chars
self.contextual_grounding_from_messages = contextual_grounding_from_messages
self.experimental_use_latest_role_message_only = bool(kwargs.get("experimental_use_latest_role_message_only"))
# Resource-less, detect-only InvokeGuardrailChecks mode. Present `checks`
@ -459,8 +461,8 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
"""
Flatten a message into text blocks, preserving any contextual-grounding
qualifier carried by the content-block ``type`` (grounding_source / query).
Untagged text keeps ``qualifier=None`` so the payload is unchanged for
callers that do not use grounding.
Untagged text keeps ``qualifier=None``; the OUTPUT scan decides whether to
derive grounding qualifiers from it.
"""
content: Final = message.get("content")
if content is None:
@ -493,6 +495,10 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
result carrying externally-influenced content can supply fake evidence for the
contextual-grounding check to grade the response against. ``query`` is accepted
from any role (it is the user's question).
With ``contextual_grounding_from_messages`` on, a request with no tagged blocks
falls back to the plain messages: system / developer text is the grounding
source and the latest user message is the query.
"""
grounding: Final[list[QualifiedTextBlock]] = []
for message in messages or []:
@ -504,7 +510,33 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
and role in _GROUNDING_SOURCE_TRUSTED_ROLES
):
grounding.append(block)
return grounding
if grounding or not self.contextual_grounding_from_messages:
return grounding
return self._derive_grounding_blocks_from_plain_messages(messages)
def _derive_grounding_blocks_from_plain_messages(
self, messages: list[AllMessageValues] | None
) -> list[QualifiedTextBlock]:
if not messages:
return []
latest_user_index: Final = self._find_latest_message_index(messages, target_role="user")
if latest_user_index is None:
return []
sources: Final = tuple(
QualifiedTextBlock(text=block.text, qualifier="grounding_source")
for message in messages
if message.get("role") in _GROUNDING_SOURCE_TRUSTED_ROLES
for block in self.get_content_items_for_message(message=message) or []
if block.text
)
queries: Final = tuple(
QualifiedTextBlock(text=block.text, qualifier="query")
for block in self.get_content_items_for_message(message=messages[latest_user_index]) or []
if block.text
)
if not sources or not queries:
return []
return [*sources, *queries]
def supports_scan_only_tool_results(self) -> bool:
return self.experimental_use_latest_role_message_only is not True
@ -3210,6 +3242,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
bedrock_response = await self.make_bedrock_api_request(
source="OUTPUT",
response=synthetic_response,
messages=request_data.get("messages"),
request_data=request_data,
logging_event_type=_log_hook,
)

View file

@ -7,7 +7,7 @@
import fnmatch
import os
from collections.abc import Mapping
from collections.abc import Mapping, Sequence
from typing import TYPE_CHECKING, Any, Final, Literal, Optional
import httpx
@ -24,7 +24,7 @@ from litellm.llms.custom_httpx.http_handler import (
httpxSpecialProvider,
)
from litellm.types.guardrails import GuardrailEventHooks
from litellm.types.llms.openai import ChatCompletionToolParam
from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolParam
from litellm.types.proxy.guardrails.guardrail_hooks.generic_guardrail_api import (
GenericGuardrailAPIMetadata,
GenericGuardrailAPIRequest,
@ -150,6 +150,26 @@ def _extract_inbound_headers(
return None
def _structured_rows_to_write_back(
original_rows: Sequence[AllMessageValues] | None,
shown_rows: Sequence[AllMessageValues] | None,
returned_rows: Sequence[AllMessageValues],
) -> tuple[AllMessageValues, ...] | None:
"""The request model drops row keys its message types do not declare, so a
row the server echoes back verbatim is restored to the original row object.
A server that echoes every row back unchanged has not rewritten anything
per row, so its answer is read from texts, as it was before rows could be
returned at all."""
if original_rows is None or shown_rows is None or len(returned_rows) != len(original_rows):
return tuple(returned_rows)
if all(returned == shown for shown, returned in zip(shown_rows, returned_rows)):
return None
return tuple(
original if returned == shown else returned
for original, shown, returned in zip(original_rows, shown_rows, returned_rows)
)
class GenericGuardrailAPI(CustomGuardrail):
"""
Generic Guardrail API integration for LiteLLM.
@ -322,6 +342,8 @@ class GenericGuardrailAPI(CustomGuardrail):
texts: list,
images: list[str] | None,
tools: list[ChatCompletionToolParam] | None,
structured_messages: Sequence[AllMessageValues] | None,
shown_messages: Sequence[AllMessageValues] | None,
guardrail_response: GenericGuardrailAPIResponse,
) -> GenericGuardrailAPIInputs:
# Action is NONE or no modifications needed
@ -336,6 +358,13 @@ class GenericGuardrailAPI(CustomGuardrail):
return_inputs["tools"] = guardrail_response.tools
elif tools:
return_inputs["tools"] = tools
rows_to_write_back: Final = (
_structured_rows_to_write_back(structured_messages, shown_messages, guardrail_response.structured_messages)
if guardrail_response.structured_messages
else None
)
if rows_to_write_back is not None:
return_inputs["structured_messages"] = list(rows_to_write_back) # mutable-ok: guardrail inputs take a list
if guardrail_response.stream_holdback_chars is not None:
return_inputs["stream_holdback_chars"] = guardrail_response.stream_holdback_chars
return return_inputs
@ -473,6 +502,8 @@ class GenericGuardrailAPI(CustomGuardrail):
texts=texts,
images=images,
tools=tools,
structured_messages=structured_messages,
shown_messages=guardrail_request.structured_messages,
guardrail_response=guardrail_response,
)

View file

@ -1,11 +1,13 @@
from collections.abc import AsyncGenerator, Mapping, Sequence
import time
from collections.abc import AsyncGenerator, Awaitable, Callable, Mapping, Sequence
from enum import Enum, auto
from typing import TYPE_CHECKING, Any, Final, Literal
from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal
import httpx
from fastapi import HTTPException
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel
import json
@ -23,6 +25,7 @@ from litellm.litellm_core_utils.core_helpers import (
get_or_create_metadata_bucket,
)
from litellm.llms.custom_httpx.http_handler import (
AsyncHTTPHandler,
get_async_httpx_client,
httpxSpecialProvider,
)
@ -52,6 +55,7 @@ from litellm.types.utils import (
CallTypes,
CallTypesLiteral,
Choices,
GenericGuardrailAPIInputs,
GuardrailStatus,
ModelResponse,
ModelResponseStream,
@ -118,8 +122,12 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
Supports:
- Pre-call sanitization (sanitizeUserPrompt)
- Post-call sanitization (sanitizeModelResponse)
- logging_only: scans the completed response after it reaches the client and
records the verdict in spend logs without blocking
"""
use_native_lifecycle_hooks: ClassVar[bool] = True
@classmethod
def get_supported_event_hooks(cls) -> list[GuardrailEventHooks]:
return [
@ -128,6 +136,7 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
GuardrailEventHooks.post_call,
GuardrailEventHooks.pre_mcp_call,
GuardrailEventHooks.during_mcp_call,
GuardrailEventHooks.logging_only,
]
def __init__(
@ -138,6 +147,8 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
credentials: VERTEX_CREDENTIALS_TYPES | None = None,
api_endpoint: str | None = None,
sanitize_error_detail: "bool | None" = True,
async_handler: AsyncHTTPHandler | None = None,
access_token_provider: Callable[[], Awaitable[tuple[str, str]]] | None = None,
**kwargs,
):
# Set supported event hooks if not already provided
@ -154,7 +165,10 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
VertexBase.__init__(self)
# Then set our attributes (this ensures project_id is not overwritten)
self.async_handler = get_async_httpx_client(llm_provider=httpxSpecialProvider.GuardrailCallback)
self.async_handler = async_handler or get_async_httpx_client(
llm_provider=httpxSpecialProvider.GuardrailCallback
)
self.access_token_provider = access_token_provider
self.template_id = template_id
self.project_id = project_id
self.location = location or "us-central1"
@ -278,11 +292,14 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
If file_bytes and file_type are provided, file prompt sanitization is performed.
"""
# Get access token using VertexBase auth
access_token, resolved_project_id = await self._ensure_access_token_async(
credentials=self.credentials,
project_id=self.project_id,
custom_llm_provider="vertex_ai",
)
if self.access_token_provider is not None:
access_token, resolved_project_id = await self.access_token_provider()
else:
access_token, resolved_project_id = await self._ensure_access_token_async(
credentials=self.credentials,
project_id=self.project_id,
custom_llm_provider="vertex_ai",
)
# Use resolved project ID if not explicitly set
if not self.project_id and resolved_project_id:
@ -1096,6 +1113,11 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
add_guardrail_to_applied_guardrails_header,
)
if self.should_run_guardrail(data=request_data, event_type=GuardrailEventHooks.post_call) is not True:
async for chunk in response:
yield chunk
return
all_chunks: Final[Sequence[object]] = tuple([chunk async for chunk in response])
if not all_chunks or self._is_terminal_error_stream(all_chunks):
@ -1213,6 +1235,60 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
for chunk in all_chunks:
yield chunk
@log_guardrail_information
async def apply_guardrail(
self,
inputs: GenericGuardrailAPIInputs,
request_data: dict,
input_type: Literal["request", "response"],
logging_obj: "LiteLLMLoggingObj | None" = None,
) -> GenericGuardrailAPIInputs:
content: Final = "\n".join(text for text in inputs.get("texts") or () if text)
if not content:
return inputs
source: Final[Literal["user_prompt", "model_response"]] = (
"user_prompt" if input_type == "request" else "model_response"
)
start_time: Final = time.time()
try:
armor_response: Final = await self.make_model_armor_request(
content=content, source=source, request_data=request_data
)
except (ModelArmorAPIError, httpx.HTTPError) as e:
error_end_time: Final = time.time()
self.add_standard_logging_guardrail_information_to_request_data(
guardrail_json_response=str(e),
request_data=request_data,
guardrail_status="guardrail_failed_to_respond",
guardrail_provider="model_armor",
start_time=start_time,
end_time=error_end_time,
duration=error_end_time - start_time,
)
return inputs
flagged: Final = self._should_block_content(armor_response, allow_sanitization=False)
end_time: Final = time.time()
self.add_standard_logging_guardrail_information_to_request_data(
guardrail_json_response=self._build_logging_response(armor_response),
request_data=request_data,
guardrail_status="guardrail_flagged" if flagged else "success",
guardrail_provider="model_armor",
start_time=start_time,
end_time=end_time,
duration=end_time - start_time,
)
if flagged and not self._event_hook_is_event_type(GuardrailEventHooks.logging_only):
raise HTTPException(
status_code=400,
detail=self._build_block_error_detail(
"Response blocked by Model Armor" if input_type == "response" else "Content blocked by Model Armor",
armor_response,
),
)
return inputs
@staticmethod
def get_config_model() -> type["GuardrailConfigModel"] | None:
"""

View file

@ -1600,8 +1600,8 @@ class PanwPrismaAirsHandler(CustomGuardrail):
Args:
texts: Flattened text entries from the framework.
messages: Original request messages (request_data["messages"]),
NOT structured_messages (which may have injected system content).
messages: The structured messages the framework flattened into ``texts``,
hoisted top-level system prompt included, so positions line up.
Returns a set of scannable indices, or None on count mismatch or no user/developer
message (safety fallback to existing role-filter behavior).
@ -1788,15 +1788,10 @@ class PanwPrismaAirsHandler(CustomGuardrail):
structured_messages: Final = inputs.get("structured_messages")
if structured_messages:
# For Anthropic /v1/messages: default to latest-user-only scanning.
# Uses request_data["messages"] (original format), NOT structured_messages
# (which has injected system content from adapter translation).
if self._use_latest_user_only(request_data, logging_obj):
original_messages: Final = request_data.get("messages")
if original_messages:
scannable_indices = self._get_latest_user_text_indices(texts, original_messages)
scannable_indices = self._get_latest_user_text_indices(texts, structured_messages)
# Fall through to existing role filtering if:
# - not Anthropic, OR flag explicitly False, OR
# - no original messages, OR
# - latest-user extraction returned None (no user / count mismatch)
if scannable_indices is None:
scannable_indices = self._get_scannable_text_indices(texts, structured_messages)

View file

@ -2,6 +2,7 @@ import asyncio
import base64
import os
from collections.abc import Mapping, Sequence
from types import MappingProxyType
from typing import TYPE_CHECKING, Final, Literal, Optional
import httpx
@ -14,11 +15,13 @@ from litellm.integrations.custom_guardrail import (
CustomGuardrail,
log_guardrail_information,
)
from litellm.llms.base_llm.guardrail_translation.utils import message_slot_texts, message_with_slot_texts
from litellm.llms.custom_httpx.http_handler import (
get_async_httpx_client,
httpxSpecialProvider,
)
from litellm.types.guardrails import GuardrailEventHooks
from litellm.types.llms.openai import AllMessageValues
from litellm.types.utils import GenericGuardrailAPIInputs
if TYPE_CHECKING:
@ -27,12 +30,37 @@ if TYPE_CHECKING:
_SANITIZE_FILE_FAIL_OPEN_TIMEOUT_SECONDS: Final = 30.0
_SANITIZE_FILE_QUEUED_STATUSES: Final = frozenset({"created", "in progress"})
_PROTECT_ROLES: Final = frozenset({"system", "user", "assistant"})
class PromptSecurityGuardrailMissingSecrets(Exception):
pass
def _inputs_with_structured_messages(
inputs: GenericGuardrailAPIInputs, rewritten_messages: Sequence[AllMessageValues] | None
) -> GenericGuardrailAPIInputs:
if rewritten_messages is None:
return inputs
patched: Final[GenericGuardrailAPIInputs] = {
**inputs,
"structured_messages": list(rewritten_messages), # mutable-ok: the TypedDict field is declared as a list
}
return patched
def _inputs_with_modifications(
inputs: GenericGuardrailAPIInputs,
modified_texts: list[str],
rewritten_messages: Sequence[AllMessageValues] | None,
) -> GenericGuardrailAPIInputs:
if not modified_texts:
return _inputs_with_structured_messages(inputs, rewritten_messages)
with_texts: Final[GenericGuardrailAPIInputs] = {**inputs, "texts": modified_texts}
return _inputs_with_structured_messages(with_texts, rewritten_messages)
class _ProtectVerdict(TypedDict, total=False):
"""One side (``prompt`` or ``response``) of an ``/api/protect`` verdict."""
@ -275,14 +303,39 @@ class PromptSecurityGuardrail(CustomGuardrail):
detail="Blocked by Prompt Security, Violations: " + ", ".join(violations),
)
elif action == "modify":
# Extract modified texts from modified_messages
modified_messages: Final = result.get("modified_messages", [])
modified_texts: Final = self._extract_texts_from_messages(modified_messages)
if modified_texts:
inputs["texts"] = modified_texts
return _inputs_with_modifications(
inputs,
self._extract_texts_from_messages(modified_messages),
self._structured_messages_with_modifications(structured_messages, modified_messages),
)
return inputs
def _is_sent_to_protect(self, message: Mapping[str, object]) -> bool:
return self.check_tool_results or message.get("role") in _PROTECT_ROLES
def _structured_messages_with_modifications(
self,
structured_messages: Sequence[AllMessageValues],
modified_messages: Sequence[Mapping[str, object]],
) -> tuple[AllMessageValues, ...] | None:
sent_indices: Final = tuple(
index for index, message in enumerate(structured_messages) if self._is_sent_to_protect(message)
)
if not sent_indices or len(sent_indices) != len(modified_messages):
return None
rewritten: Final = tuple(
message_with_slot_texts(structured_messages[index], self._extract_texts_from_messages((modified,)))
for index, modified in zip(sent_indices, modified_messages)
)
replacements: Final = MappingProxyType(
{index: message for index, message in zip(sent_indices, rewritten) if message is not None}
)
if len(replacements) != len(sent_indices):
return None
return tuple(replacements.get(index, message) for index, message in enumerate(structured_messages))
async def _apply_guardrail_on_response(
self,
inputs: GenericGuardrailAPIInputs,
@ -346,19 +399,7 @@ class PromptSecurityGuardrail(CustomGuardrail):
return inputs
def _extract_texts_from_messages(self, messages: Sequence[Mapping[str, object]]) -> list[str]:
"""Extract text content from messages."""
texts: Final = []
for message in messages:
content = message.get("content")
if isinstance(content, str):
texts.append(content)
elif isinstance(content, list):
for item in content:
if isinstance(item, dict) and item.get("type") == "text":
text = item.get("text")
if text:
texts.append(text)
return texts
return [text for message in messages for text in message_slot_texts(message)]
async def _process_standalone_images(self, images: list[str], user_api_key_alias: str | None) -> None:
"""Process standalone images from inputs (data URLs)."""
@ -512,16 +553,18 @@ class PromptSecurityGuardrail(CustomGuardrail):
"metadata": result.get("metadata", {}),
"violations": result.get("metadata", {}).get("violations", []),
}
elif status == "in progress":
verbose_proxy_logger.debug(
"Prompt Security Guardrail: File sanitization in progress (attempt %d/%d)",
attempt + 1,
self.max_poll_attempts,
)
continue
else:
if status not in _SANITIZE_FILE_QUEUED_STATUSES:
raise HTTPException(status_code=500, detail=f"Unexpected sanitization status: {status}")
verbose_proxy_logger.debug(
"Prompt Security Guardrail: File sanitization status=%s for jobId=%s (attempt %d/%d)",
status,
job_id,
attempt + 1,
self.max_poll_attempts,
)
raise HTTPException(status_code=408, detail="File sanitization timeout")
def _raise_if_file_blocked(self, sanitization_result: _SanitizeResult, resource_name: str) -> None:
@ -678,14 +721,13 @@ class PromptSecurityGuardrail(CustomGuardrail):
This allows checking tool results for indirect prompt injection when enabled.
"""
supported_roles: Final = ["system", "user", "assistant"]
filtered_messages: Final = []
transformed_count = 0
filtered_count = 0
for message in messages:
role = message.get("role", "")
if role in supported_roles:
if role in _PROTECT_ROLES:
filtered_messages.append(message)
else:
if self.check_tool_results:

View file

@ -23,6 +23,7 @@ def initialize_bedrock(litellm_params: LitellmParams, guardrail: Guardrail):
prompt_attack_threshold=litellm_params.prompt_attack_threshold,
pii_confidence_threshold=litellm_params.pii_confidence_threshold,
chunk_budget_chars=litellm_params.chunk_budget_chars,
contextual_grounding_from_messages=litellm_params.contextual_grounding_from_messages,
default_on=litellm_params.default_on,
disable_exception_on_block=litellm_params.disable_exception_on_block,
mask_request_content=litellm_params.mask_request_content,

View file

@ -42,7 +42,7 @@ if TYPE_CHECKING:
router: Final = APIRouter()
_EMPTY_UNITS: Final[Mapping[str, int]] = MappingProxyType({})
_ACTION_SEVERITY: Final[Mapping[str, int]] = MappingProxyType({"passed": 0, "flagged": 1, "blocked": 2})
_ACTION_SEVERITY: Final[Mapping[str, int]] = MappingProxyType({"not_run": 0, "passed": 1, "flagged": 2, "blocked": 3})
_T = TypeVar("_T")
@ -325,7 +325,7 @@ class UsageDetailResponse(BaseModel):
class UsageLogEntry(BaseModel):
id: str
timestamp: str
action: str # blocked | passed | flagged
action: str # blocked | passed | flagged | not_run
score: float | None
latency_ms: float | None
model: str | None

View file

@ -193,10 +193,12 @@ async def _upsert_rows_with_retry(
def guardrail_status_to_action(status: str | None) -> str:
"""Map StandardLogging guardrail_status to blocked/passed/flagged."""
"""Map StandardLogging guardrail_status to blocked/passed/flagged/not_run."""
if not status:
return "passed"
s: Final = (status or "").lower()
if s == "not_run":
return "not_run"
if "intervened" in s or "block" in s:
return "blocked"
if "flagged" in s or "fail" in s or "error" in s:
@ -354,37 +356,49 @@ async def process_spend_logs_guardrail_usage(
"flagged_count": 0,
}
)
index_rows: Final[list[dict[str, object]]] = []
index_rows_by_key: Final[dict[tuple[str, str], dict[str, object]]] = {}
for payload in logs_to_process:
request_id = payload.get("request_id")
start_time = _parse_payload_start_time(payload)
if not request_id or start_time is None:
if not isinstance(request_id, str) or not request_id or start_time is None:
continue
date_key = _date_str(start_time)
for entry in _parse_guardrail_info_from_payload(payload):
guardrail_id = entry.get("guardrail_id") or entry.get("guardrail_name") or ""
if not guardrail_id:
entries = _parse_guardrail_info_from_payload(payload)
ids_by_name = MappingProxyType(
{
e["guardrail_name"]: e["guardrail_id"]
for e in entries
if e.get("guardrail_id") and isinstance(e.get("guardrail_name"), str) and e["guardrail_name"]
}
)
for entry in entries:
raw_name = entry.get("guardrail_name")
guardrail_name = raw_name if isinstance(raw_name, str) else ""
guardrail_id = entry.get("guardrail_id") or ids_by_name.get(guardrail_name) or guardrail_name
if not isinstance(guardrail_id, str) or not guardrail_id:
continue
key = _MetricsKey(guardrail_id, date_key)
daily_guardrail[key]["requests_evaluated"] += 1
action = guardrail_status_to_action(entry.get("guardrail_status"))
if action == "passed":
daily_guardrail[key]["passed_count"] += 1
elif action == "blocked":
daily_guardrail[key]["blocked_count"] += 1
else:
daily_guardrail[key]["flagged_count"] += 1
if action != "not_run":
key = _MetricsKey(guardrail_id, date_key)
daily_guardrail[key]["requests_evaluated"] += 1
if action == "passed":
daily_guardrail[key]["passed_count"] += 1
elif action == "blocked":
daily_guardrail[key]["blocked_count"] += 1
else:
daily_guardrail[key]["flagged_count"] += 1
policy_id = entry.get("policy_id")
index_rows.append(
{
prior = index_rows_by_key.get((request_id, guardrail_id))
if prior is None or (prior["policy_id"] is None and policy_id is not None):
index_rows_by_key[(request_id, guardrail_id)] = {
"request_id": request_id,
"guardrail_id": guardrail_id,
"policy_id": policy_id,
"start_time": start_time,
}
)
index_rows: Final = tuple(index_rows_by_key.values())
async with pending.lock:
pending_metrics: Final = pending.metrics

View file

@ -45,7 +45,10 @@ from litellm.proxy.auth.auth_utils import (
from litellm.proxy.auth.model_checks import get_key_models
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler
from litellm.proxy.db.health_check_latest import LatestHealthCheckRow
from litellm.proxy.db.health_check_latest import (
LatestHealthCheckRow,
query_latest_health_checks,
)
from litellm.proxy.db.proxy_worker_heartbeat import count_live_proxy_workers
from litellm.proxy.health_check import (
ADMIN_ONLY_HEALTH_DISPLAY_PARAMS,
@ -876,7 +879,7 @@ async def _save_background_health_checks_to_db(
)
# Step 3: Get latest health checks for all models in one query to compare status
latest_checks: Final = await prisma_client.get_all_latest_health_checks()
latest_checks: Final = await query_latest_health_checks(prisma_client)
latest_checks_map: Final = {}
for check in latest_checks:
# Use model_id as primary key, fallback to model_name

View file

@ -9,6 +9,7 @@ from .max_budget_per_session_limiter import _PROXY_MaxBudgetPerSessionHandler
from .max_iterations_limiter import _PROXY_MaxIterationsHandler
from .parallel_request_limiter import _PROXY_MaxParallelRequestsHandler
from .parallel_request_limiter_v3 import _PROXY_MaxParallelRequestsHandler_v3
from .prompt_cache_prediction import PromptCacheObserver
from .responses_id_security import ResponsesIDSecurity
from .sensitive_data_routing import _PROXY_SensitiveDataRoutingHandler
@ -25,6 +26,7 @@ PROXY_HOOKS: Final = {
"max_iterations_limiter": _PROXY_MaxIterationsHandler,
"max_budget_per_session_limiter": _PROXY_MaxBudgetPerSessionHandler,
"sensitive_data_routing": _PROXY_SensitiveDataRoutingHandler,
"prompt_cache_prediction": PromptCacheObserver,
}
## FEATURE FLAG HOOKS ##

View file

@ -18,12 +18,13 @@ Quick summary:
"""
import json
from collections.abc import Iterable, Mapping, Sequence
from collections.abc import Callable, Iterable, Mapping, Sequence
from datetime import datetime
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, Literal, NoReturn, TypeAlias
from fastapi import HTTPException
from pydantic import BaseModel, Field, TypeAdapter
from pydantic import BaseModel, Field, TypeAdapter, ValidationError
import litellm
from litellm._logging import verbose_proxy_logger
@ -33,6 +34,7 @@ from litellm.batches.batch_utils import (
_extract_file_access_credentials,
_iter_batch_input_lines,
)
from litellm.constants import BATCH_TPD_DESCRIPTOR_SUFFIX, BATCH_TPD_WINDOW_SECONDS
from litellm.exceptions import RateLimitErrorCategory
from litellm.integrations.custom_logger import CustomLogger
from litellm.proxy._types import (
@ -55,6 +57,7 @@ from litellm.proxy.hooks.batch_enqueued_tokens import (
from litellm.proxy.hooks.parallel_request_limiter_v3 import (
PROJECT_ITPM_DESCRIPTOR_KEY,
PROJECT_OTPM_DESCRIPTOR_KEY,
ReservationAwareIncrementOperation,
get_or_create_request_stash,
)
from litellm.proxy.hooks.rate_limiter_utils import resolve_llm_provider_for_rate_limit
@ -92,6 +95,7 @@ else:
_BATCH_BODY_ADAPTER: Final = TypeAdapter(dict[str, object])
_WINDOW_START_ADAPTER: Final[TypeAdapter[int | float | str | None]] = TypeAdapter(int | float | str | None)
IncrementAmounts: TypeAlias = dict[Literal["requests", "tokens"], int]
@ -128,6 +132,7 @@ class _PROXY_BatchRateLimiter(CustomLogger):
self,
internal_usage_cache: InternalUsageCache,
parallel_request_limiter: ParallelRequestLimiter,
time_provider: Callable[[], datetime] | None = None,
):
"""
Initialize the batch rate limiter.
@ -138,9 +143,11 @@ class _PROXY_BatchRateLimiter(CustomLogger):
Args:
internal_usage_cache: Cache for storing rate limit data (auto-injected)
parallel_request_limiter: Existing rate limiter to integrate with (needs custom injection)
time_provider: Clock used for rate limit reset times (defaults to ``datetime.now``)
"""
self.internal_usage_cache = internal_usage_cache
self.parallel_request_limiter = parallel_request_limiter
self._time_provider: Final = time_provider or datetime.now
self._warned_unsupported_model_skip = False
def _get_file_bound_batch_model(self, data: dict) -> str | None:
@ -236,14 +243,48 @@ class _PROXY_BatchRateLimiter(CustomLogger):
file-bound/top-level routing model this function resolves. Charging
project quotas here would let a caller bind the file to a model
without a quota while rows execute against a quota-limited model.
Scopes with a ``tpd_limit`` (key, team, end user) are charged against a
daily token descriptor instead of their per-minute RPM/TPM descriptor,
because a batch's rows are scheduled by the provider and never share a
minute with the submission. The daily descriptor uses its own key so
its 24h window never collides with the online limiter's counters.
"""
return self.parallel_request_limiter._create_rate_limit_descriptors(
descriptors: Final = self.parallel_request_limiter._create_rate_limit_descriptors(
user_api_key_dict=user_api_key_dict,
data=data,
rpm_limit_type=None,
tpm_limit_type=None,
model_has_failures=False,
)
tpd_limits: Final[Mapping[str, tuple[str, int]]] = MappingProxyType(
{
key: (value, limit)
for key, value, limit in (
("api_key", user_api_key_dict.api_key, user_api_key_dict.tpd_limit),
("team", user_api_key_dict.team_id, user_api_key_dict.team_tpd_limit),
("end_user", user_api_key_dict.end_user_id, user_api_key_dict.end_user_tpd_limit),
)
if value and limit is not None
}
)
if not tpd_limits:
return descriptors
return [
*(d for d in descriptors if d["key"] not in tpd_limits),
*(
RateLimitDescriptor(
key=f"{key}{BATCH_TPD_DESCRIPTOR_SUFFIX}",
value=value,
rate_limit={
"requests_per_unit": None,
"tokens_per_unit": limit,
"window_size": BATCH_TPD_WINDOW_SECONDS,
},
)
for key, (value, limit) in tpd_limits.items()
),
]
@staticmethod
def _project_has_any_io_token_limits(user_api_key_dict: UserAPIKeyAuth) -> bool:
@ -583,9 +624,14 @@ class _PROXY_BatchRateLimiter(CustomLogger):
batch_usage: BatchFileUsage,
limit_type: str,
requested_model: str | None = None,
window_start: int | None = None,
) -> NoReturn:
"""Raise :class:`ProxyRateLimitError` (a 429) for batch rate limit exceeded."""
from datetime import datetime
"""Raise :class:`ProxyRateLimitError` (a 429) for batch rate limit exceeded.
``window_start`` is the active counter window's start (unix seconds) when
known, so the reset time reflects that window's actual end rather than a
full window from now.
"""
# Find the descriptor for this status. Matching on (key, value) is
# required, not key alone: a batch can carry several project ITPM/OTPM
@ -609,9 +655,12 @@ class _PROXY_BatchRateLimiter(CustomLogger):
descriptors[descriptor_index] if descriptors else {"key": "", "value": "", "rate_limit": None}
)
now: Final = datetime.now().timestamp()
window_size: Final = self.parallel_request_limiter.window_size
reset_time: Final = now + window_size
now: Final = self._time_provider().timestamp()
window_size: Final = (descriptor.get("rate_limit") or {}).get(
"window_size"
) or self.parallel_request_limiter.window_size
reset_time: Final = now + window_size if window_start is None else window_start + window_size
retry_after: Final = max(0, int(reset_time - now))
reset_time_formatted: Final = datetime.fromtimestamp(reset_time).strftime("%Y-%m-%d %H:%M:%S UTC")
remaining_display: Final = max(0, status["limit_remaining"])
@ -643,10 +692,13 @@ class _PROXY_BatchRateLimiter(CustomLogger):
if descriptor.get("key") == PROJECT_ITPM_DESCRIPTOR_KEY
else batch_usage.total_tokens
)
token_limit_label: Final = (
"TPD" if descriptor.get("key", "").endswith(BATCH_TPD_DESCRIPTOR_SUFFIX) else "TPM"
)
detail = (
f"Batch rate limit exceeded for {descriptor.get('key', 'unknown')}: {descriptor.get('value', 'unknown')}. "
f"Batch contains {batch_token_count} tokens but only {remaining_display} tokens remaining "
f"out of {current_limit} TPM limit. "
f"out of {current_limit} {token_limit_label} limit. "
f"Limit resets at: {reset_time_formatted}"
)
@ -654,7 +706,7 @@ class _PROXY_BatchRateLimiter(CustomLogger):
raise ProxyRateLimitError(
detail=detail,
headers={
"retry-after": str(window_size),
"retry-after": str(retry_after),
"rate_limit_type": limit_type,
"reset_at": reset_time_formatted,
},
@ -712,6 +764,8 @@ class _PROXY_BatchRateLimiter(CustomLogger):
parent_otel_span=user_api_key_dict.parent_otel_span,
)
stash: Final = get_or_create_request_stash()
stash.batch_tpd_refund_ops = ()
if rate_limit_response["overall_code"] == "OVER_LIMIT":
requested_model: Final = data.get("model") if data else None
for status in rate_limit_response["statuses"]:
@ -722,8 +776,70 @@ class _PROXY_BatchRateLimiter(CustomLogger):
batch_usage,
status["rate_limit_type"],
requested_model=requested_model,
window_start=await self._read_tpd_window_start(
status=status, parent_otel_span=user_api_key_dict.parent_otel_span
),
)
stash.batch_tpd_refund_ops = self._build_tpd_refund_ops(
descriptors=descriptors,
tokens=batch_usage.total_tokens,
reservation_windows=rate_limit_response.get("reservation_windows", frozenset()),
)
async def _read_tpd_window_start(self, status: "RateLimitStatus", parent_otel_span: "Span | None") -> int | None:
descriptor_key: Final = status.get("descriptor_key") or ""
if not descriptor_key.endswith(BATCH_TPD_DESCRIPTOR_SUFFIX):
return None
try:
window_start: Final = _WINDOW_START_ADAPTER.validate_python(
await self.parallel_request_limiter.internal_usage_cache.async_get_cache(
key=f"{{{descriptor_key}:{status.get('descriptor_value') or ''}}}:window",
litellm_parent_otel_span=parent_otel_span,
),
strict=True,
)
return None if window_start is None else int(float(window_start))
except (ValidationError, ValueError):
return None
def _build_tpd_refund_ops(
self,
descriptors: Sequence["RateLimitDescriptor"],
tokens: int,
reservation_windows: frozenset[tuple[str, str, Literal["redis", "local"]]],
) -> tuple[ReservationAwareIncrementOperation, ...]:
"""Refund operations for the daily token counters this batch charged.
The v3 limiter's failure hook applies them when the submission fails
after the counters were incremented. Each operation carries the window
identity the charge landed in, so the refund is skipped once that
window has rolled over.
"""
if tokens <= 0 or not reservation_windows:
return ()
tpd_descriptors_by_counter: Final[Mapping[str, RateLimitDescriptor]] = MappingProxyType(
{
self.parallel_request_limiter.create_rate_limit_keys(
descriptor["key"], descriptor["value"], "tokens"
): descriptor
for descriptor in descriptors
if descriptor["key"].endswith(BATCH_TPD_DESCRIPTOR_SUFFIX)
}
)
return tuple(
ReservationAwareIncrementOperation(
key=counter_key,
increment_value=-tokens,
ttl=BATCH_TPD_WINDOW_SECONDS,
window_key=f"{{{descriptor['key']}:{descriptor['value']}}}:window",
expected_window_start=window_start,
reservation_backend=backend,
)
for counter_key, window_start, backend in sorted(reservation_windows)
if (descriptor := tpd_descriptors_by_counter.get(counter_key)) is not None
)
async def count_input_file_usage(
self,
file_id: str,

View file

@ -9,10 +9,12 @@ import binascii
import logging
import os
import uuid
from collections.abc import Awaitable, Callable, Mapping, Sequence, Set
from collections.abc import AsyncGenerator, Awaitable, Callable, Mapping, Sequence, Set
from contextlib import asynccontextmanager
from contextvars import ContextVar
from dataclasses import dataclass, field
from datetime import datetime
from types import MappingProxyType
from typing import (
TYPE_CHECKING,
Any,
@ -23,6 +25,7 @@ from typing import (
TypedDict,
)
from pydantic import TypeAdapter
from typing_extensions import NotRequired, ReadOnly
from litellm import DualCache
@ -84,6 +87,9 @@ else:
InternalUsageCache = Any
_REQUEST_RATE_LIMIT_DATA: Final = TypeAdapter(Mapping[str, object])
BATCH_RATE_LIMITER_SCRIPT: Final = """
local results = {}
local now = tonumber(ARGV[1])
@ -390,6 +396,8 @@ CacheCounterValue: TypeAlias = int | float | str | bytes
CacheCounterValues: TypeAlias = Sequence[CacheCounterValue | None]
ReservationWindowIdentity: TypeAlias = tuple[str, str, Literal["redis", "local"]]
ParallelGaugeCacheValue: TypeAlias = dict[str, object] | int | float | str | bytes
@ -536,6 +544,7 @@ class RequestRateLimiterStash:
default_factory=frozenset
)
batch_enqueued_reservation: BatchEnqueuedTokenReservation | None = None
batch_tpd_refund_ops: tuple[ReservationAwareIncrementOperation, ...] = ()
reservation_released: bool = False
@ -677,6 +686,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
self._batch_rate_limiter = _PROXY_BatchRateLimiter(
internal_usage_cache=self.internal_usage_cache,
parallel_request_limiter=self,
time_provider=self._time_provider,
)
except Exception as e:
verbose_proxy_logger.debug("Could not load batch rate limiter: %s", e)
@ -1817,6 +1827,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
)
applied: Final[list[list[AtomicCounterMeta]]] = []
statuses: Final[list[RateLimitStatus]] = []
reservation_windows: Final[set[ReservationWindowIdentity]] = set() # mutable-ok: filled by the group loop
raw: list[CacheCounterValue]
for _idx, (keys, args, meta) in enumerate(descriptor_groups):
@ -1854,11 +1865,12 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
return response
applied.append(meta)
statuses.extend(response["statuses"])
reservation_windows.update(response.get("reservation_windows", frozenset()))
return RateLimitResponse(
overall_code="OK",
statuses=statuses,
reservation_windows=frozenset(),
reservation_windows=frozenset(reservation_windows),
)
async def _refund_applied_descriptor_groups(
@ -2673,12 +2685,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
Returns list of descriptors for API key, user, team, team member, end user,
model-specific, agent, and agent-session limits.
"""
from litellm.proxy.auth.auth_utils import (
get_team_model_rpm_limit,
get_team_model_tpm_limit,
)
descriptors: Final = []
descriptors: Final[list[RateLimitDescriptor]] = [] # mutable-ok: existing descriptor helpers append in place
# API Key rate limits
if user_api_key_dict.api_key and (
@ -2803,34 +2810,11 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
descriptors=descriptors,
)
if (
get_team_model_rpm_limit(user_api_key_dict) is not None
or get_team_model_tpm_limit(user_api_key_dict) is not None
):
_tpm_limit_for_team_model: Final = get_team_model_tpm_limit(user_api_key_dict) or {}
_rpm_limit_for_team_model: Final = get_team_model_rpm_limit(user_api_key_dict) or {}
should_check_rate_limit = False
if requested_model in _tpm_limit_for_team_model or requested_model in _rpm_limit_for_team_model:
should_check_rate_limit = True
if should_check_rate_limit:
model_specific_tpm_limit = None
model_specific_rpm_limit = None
if requested_model in _tpm_limit_for_team_model:
model_specific_tpm_limit = _tpm_limit_for_team_model[requested_model]
if requested_model in _rpm_limit_for_team_model:
model_specific_rpm_limit = _rpm_limit_for_team_model[requested_model]
descriptors.append(
RateLimitDescriptor(
key="model_per_team",
value=f"{user_api_key_dict.team_id}:{requested_model}",
rate_limit={
"requests_per_unit": model_specific_rpm_limit,
"tokens_per_unit": model_specific_tpm_limit,
"window_size": self.window_size,
},
)
)
self._add_team_model_rate_limit_descriptor_from_metadata(
user_api_key_dict=user_api_key_dict,
requested_model=requested_model if isinstance(requested_model, str) else None,
descriptors=descriptors,
)
# Agent-level and session-level rate limits
resolved_agent_id: Final = self._get_resolved_agent_id(user_api_key_dict, data)
@ -3416,6 +3400,108 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
requested_model,
)
async def _build_request_rate_limit_descriptors(
self,
user_api_key_dict: UserAPIKeyAuth,
data: Mapping[str, object],
call_type: str | None,
) -> list[RateLimitDescriptor]: # mutable-ok: the shared generation reservation helpers require a list
metadata: Final = _REQUEST_RATE_LIMIT_DATA.validate_python(
user_api_key_dict.metadata or MappingProxyType({}) # pyright: ignore[reportUnknownMemberType] # validates the legacy auth metadata boundary
)
rpm_value: Final = metadata.get("rpm_limit_type")
tpm_value: Final = metadata.get("tpm_limit_type")
rpm_limit_type: Final = rpm_value if isinstance(rpm_value, str) else None
tpm_limit_type: Final = tpm_value if isinstance(tpm_value, str) else None
model_value: Final = data.get("model")
requested_model: Final = model_value if isinstance(model_value, str) else None
model_has_failures: Final = (
await self._check_model_has_recent_failures(
model=requested_model,
parent_otel_span=user_api_key_dict.parent_otel_span,
)
if requested_model and self._is_dynamic_rate_limiting_enabled(rpm_limit_type, tpm_limit_type)
else False
)
descriptors: Final = self._create_rate_limit_descriptors( # pyright: ignore[reportUnknownMemberType] # legacy helper reads a dictionary with validated keys
user_api_key_dict=user_api_key_dict,
data=dict(data), # mutable-ok: legacy descriptor helpers accept a request dictionary
rpm_limit_type=rpm_limit_type,
tpm_limit_type=tpm_limit_type,
model_has_failures=model_has_failures,
call_type=call_type,
)
self._add_project_model_rate_limit_descriptor_from_metadata(
user_api_key_dict=user_api_key_dict,
requested_model=requested_model,
descriptors=descriptors,
)
self.add_project_io_token_rate_limit_descriptors_from_metadata(
user_api_key_dict=user_api_key_dict,
requested_model=requested_model,
descriptors=descriptors,
)
return [ # mutable-ok: the shared generation reservation helpers require a list
*descriptors,
*self.create_organization_rate_limit_descriptor(user_api_key_dict, requested_model),
]
async def _release_request_capacity_when_admitted(
self,
admission: asyncio.Task[RateLimitResponse],
acquisition: ParallelSlotAcquisition,
user_api_key_dict: UserAPIKeyAuth,
) -> None:
response: Final = await admission
if response["overall_code"] == "OK":
await self._release_parallel_request_slots(acquisition, user_api_key_dict.parent_otel_span)
@asynccontextmanager
async def request_capacity(
self,
user_api_key_dict: UserAPIKeyAuth,
model: str,
*,
request_data: Mapping[str, object] | None = None,
) -> AsyncGenerator[None, None]:
"""Charge one non-generation provider request to RPM and hold its concurrency slot."""
data: Final = MappingProxyType({**(request_data or MappingProxyType({})), "model": model})
descriptors: Final = await self._build_request_rate_limit_descriptors(user_api_key_dict, data, None)
acquisition: Final = ParallelSlotAcquisition(
slot_id=uuid.uuid4().hex,
counter_keys=[ # mutable-ok: the shared slot-release contract requires a list
self.create_rate_limit_keys(d["key"], d["value"], "max_parallel_requests")
for d in descriptors
if d["rate_limit"] is not None and d["rate_limit"].get("max_parallel_requests") is not None
],
)
admission: Final = asyncio.create_task(
self.should_rate_limit(
descriptors=descriptors,
parent_otel_span=user_api_key_dict.parent_otel_span,
skip_tpm_check=True,
parallel_slot_id=acquisition["slot_id"],
)
)
try:
response: Final = await asyncio.shield(admission)
if response["overall_code"] == "OVER_LIMIT":
self._handle_rate_limit_error(response, descriptors, model)
yield
finally:
cleanup: Final = asyncio.create_task(
self._release_request_capacity_when_admitted(admission, acquisition, user_api_key_dict)
)
cancellation: asyncio.CancelledError | None = None # rebind-ok: retain cancellation until cleanup finishes
while not cleanup.done():
try:
await asyncio.shield(cleanup)
except asyncio.CancelledError as exc:
cancellation = exc # rebind-ok: retain the latest cancellation without interrupting slot release
cleanup.result()
if cancellation is not None:
raise cancellation
async def async_pre_call_hook(
self,
user_api_key_dict: UserAPIKeyAuth,
@ -3444,59 +3530,15 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
call_type=call_type,
)
# Get rate limit types from metadata
metadata: Final = user_api_key_dict.metadata or {}
rpm_limit_type: Final = metadata.get("rpm_limit_type")
tpm_limit_type: Final = metadata.get("tpm_limit_type")
# For dynamic mode, check if the model has recent failures
model_has_failures = False
requested_model: Final = data.get("model", None)
if (
self._is_dynamic_rate_limiting_enabled(
rpm_limit_type=rpm_limit_type,
tpm_limit_type=tpm_limit_type,
)
and requested_model
):
model_has_failures = await self._check_model_has_recent_failures(
model=requested_model,
parent_otel_span=user_api_key_dict.parent_otel_span,
)
# Create rate limit descriptors
descriptors: Final = self._create_rate_limit_descriptors(
request_data: Final = _REQUEST_RATE_LIMIT_DATA.validate_python(data)
model_value: Final = request_data.get("model")
requested_model: Final = model_value if isinstance(model_value, str) else None
descriptors: Final = await self._build_request_rate_limit_descriptors(
user_api_key_dict=user_api_key_dict,
data=data,
rpm_limit_type=rpm_limit_type,
tpm_limit_type=tpm_limit_type,
model_has_failures=model_has_failures,
data=request_data,
call_type=call_type,
)
# Add team model rate limits from team_metadata
self._add_team_model_rate_limit_descriptor_from_metadata(
user_api_key_dict=user_api_key_dict,
requested_model=requested_model,
descriptors=descriptors,
)
# Project Level Rate Limits
self._add_project_model_rate_limit_descriptor_from_metadata(
user_api_key_dict=user_api_key_dict,
requested_model=requested_model,
descriptors=descriptors,
)
self.add_project_io_token_rate_limit_descriptors_from_metadata(
user_api_key_dict=user_api_key_dict,
requested_model=requested_model,
descriptors=descriptors,
)
# Org Level Rate Limits
descriptors.extend(self.create_organization_rate_limit_descriptor(user_api_key_dict, requested_model))
# Only check rate limits if we have descriptors with actual limits
if descriptors:
# First pass: RPM and max_parallel_requests sliding-window check.
@ -4788,6 +4830,13 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
)
stash.batch_enqueued_reservation = None
if stash.batch_tpd_refund_ops:
await self.async_increment_reservation_aware_tokens(
pipeline_operations=stash.batch_tpd_refund_ops,
parent_otel_span=user_api_key_dict.parent_otel_span,
)
stash.batch_tpd_refund_ops = ()
if stash.reservation_released:
return
reserved_tokens: Final = stash.reserved_tokens

View file

@ -0,0 +1,142 @@
from __future__ import annotations
import asyncio
import time
from collections.abc import Callable, Mapping
from datetime import datetime
from typing import TYPE_CHECKING, Final, Literal
import httpx
from pydantic import BaseModel, ConfigDict, Field, TypeAdapter, ValidationError
from litellm.caching.dual_cache import DualCache
from litellm.integrations.custom_logger import CustomLogger
from litellm.llms.anthropic.prompt_cache_prediction import PromptPrefix, parse_observed_cache
from litellm.types.utils import ModelResponse
if TYPE_CHECKING:
from litellm.proxy.utils import InternalUsageCache
_RETENTION_SECONDS: Final = 86_400
class CacheObservation(BaseModel):
model_config = ConfigDict(extra="forbid", frozen=True, strict=True)
fingerprint: str = Field(pattern=r"^[0-9a-f]{64}$")
cached_tokens: int = Field(gt=0)
observed_at: float = Field(ge=0, allow_inf_nan=False)
expires_at: float = Field(ge=0, allow_inf_nan=False)
_CACHE_ENTRY: Final[TypeAdapter[CacheObservation | str | None]] = TypeAdapter(CacheObservation | str | None)
def _cache_key(scope: str, fingerprint: str) -> str:
return f"prompt-cache-observation:{scope}:{fingerprint}"
async def lookup(
cache: DualCache, scope: str, prefix: PromptPrefix, now: float | None = None
) -> CacheObservation | None:
checked_at: Final = time.time() if now is None else now
exact: Final = await _read_exact(cache, scope, prefix.fingerprint)
if exact is not None and exact.expires_at > checked_at:
return exact
older: Final = await asyncio.gather(
*(_read_exact(cache, scope, fingerprint) for fingerprint in prefix.fingerprints[1:])
)
observations: Final = tuple(observation for observation in (exact, *older) if observation is not None)
return next(
(observation for observation in observations if observation.expires_at > checked_at),
next(iter(observations), None),
)
async def _read_exact(cache: DualCache, scope: str, fingerprint: str) -> CacheObservation | None:
try:
value: Final = _CACHE_ENTRY.validate_python(await cache.async_get_cache(_cache_key(scope, fingerprint), ttl=1)) # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # validate the legacy cache's untyped result at the I/O boundary
if value is None:
return None
observation: Final = CacheObservation.model_validate_json(value) if isinstance(value, str) else value
except ValidationError:
return None
return observation if observation.fingerprint == fingerprint else None
class _Metadata(BaseModel):
model_config = ConfigDict(strict=True)
user_api_key_hash: str = Field(min_length=1)
class _Logged(BaseModel):
model_config = ConfigDict(strict=True)
status: Literal["success"]
model_id: str = Field(min_length=1)
metadata: _Metadata
class _Event(BaseModel):
model_config = ConfigDict(strict=True, arbitrary_types_allowed=True)
call_type: Literal["anthropic_messages"]
custom_llm_provider: Literal["anthropic"]
cache_hit: bool | None = None
httpx_response: httpx.Response
first_api_call_start_time: datetime
standard_logging_object: _Logged
stream: bool = False
prompt_cache_response_complete: bool = False
class PromptCacheObserver(CustomLogger):
def __init__(self, internal_usage_cache: InternalUsageCache, clock: Callable[[], float] = time.time) -> None:
super().__init__() # pyright: ignore[reportUnknownMemberType] # base callback constructor accepts untyped kwargs
self.cache = internal_usage_cache.dual_cache
self.clock = clock
async def async_log_success_event(
self, kwargs: Mapping[str, object], response_obj: object, start_time: datetime, end_time: datetime
) -> None:
if not isinstance(response_obj, ModelResponse):
return
try:
event: Final = _Event.model_validate(kwargs)
wire: Final = event.httpx_response.request
except (ValidationError, RuntimeError, httpx.RequestNotRead):
return
if (
event.cache_hit
or event.httpx_response.status_code != 200
or (event.stream and not event.prompt_cache_response_complete)
):
return
observed: Final = parse_observed_cache(
wire,
response_obj,
event.standard_logging_object.metadata.user_api_key_hash,
event.standard_logging_object.model_id,
)
if observed is None:
return
prefix: Final = observed.prefix
scope: Final = observed.scope
cache_tokens: Final = observed.cached_tokens
now: Final = self.clock()
started: Final = event.first_api_call_start_time.timestamp()
if started > now:
return
if observed.cache_creation_tokens == 0:
previous: Final = await _read_exact(self.cache, scope, prefix.fingerprint)
if previous is None or previous.fingerprint != prefix.fingerprint or previous.cached_tokens != cache_tokens:
return
observation: Final = CacheObservation(
fingerprint=prefix.fingerprint,
cached_tokens=cache_tokens,
observed_at=now,
expires_at=started + prefix.ttl_seconds,
)
key: Final = _cache_key(scope, prefix.fingerprint)
payload: Final = observation.model_dump_json()
await self.cache.async_set_cache(key, payload, ttl=_RETENTION_SECONDS) # pyright: ignore[reportUnknownMemberType] # legacy cache accepts a serialized validated observation
if self.cache.redis_cache is not None:
await self.cache.async_set_cache(key, payload, local_only=True, ttl=1) # pyright: ignore[reportUnknownMemberType] # keep the local copy short-lived while Redis retains stale evidence

View file

@ -1,5 +1,6 @@
"""Contract machinery shared by every LiteLLM-defined list route, on any surface."""
from collections.abc import Sequence
from typing import Final
from urllib.parse import urlencode
@ -7,6 +8,7 @@ from fastapi import Request
from fastapi.dependencies.utils import get_flat_params
from fastapi.params import ParamTypes
from fastapi.responses import JSONResponse
from typing_extensions import ReadOnly, TypedDict
from litellm.types.proxy.management_endpoints.management_v1 import (
ListLinks,
@ -56,6 +58,40 @@ def escape_like(value: str) -> str:
return value.replace("\\", "\\\\").replace("%", "\\%").replace("_", "\\_")
class ValidationErrorDetail(TypedDict):
"""The keys of a pydantic/FastAPI validation error a problem document needs."""
type: ReadOnly[str]
loc: ReadOnly[tuple[int | str, ...]]
msg: ReadOnly[str]
def _is_length_error_of_rejected_items(error: ValidationErrorDetail, errors: Sequence[ValidationErrorDetail]) -> bool:
"""pydantic counts only items that validated, so a bad item also trips the parent's min_length."""
return error["type"] == "too_short" and any(
len(other["loc"]) > len(error["loc"]) and other["loc"][: len(error["loc"])] == error["loc"] for other in errors
)
def request_validation_problem(raw_errors: Sequence[ValidationErrorDetail]) -> ProblemDetail:
"""A body that fails validation (an unknown field included) is 422; a bad query parameter is 400."""
errors: Final = tuple(error for error in raw_errors if not _is_length_error_of_rejected_items(error, raw_errors))
detail: Final = "; ".join(f"{'.'.join(str(part) for part in error['loc'][1:])}: {error['msg']}" for error in errors)
if any(error["loc"] and error["loc"][0] == "body" for error in errors):
return ProblemDetail(
type=f"{PROBLEM_TYPE_BASE}invalid-request-body",
title="Invalid request body",
status=422,
detail=detail or "The request body is invalid.",
)
return ProblemDetail(
type=f"{PROBLEM_TYPE_BASE}invalid-query-parameter",
title="Invalid query parameter",
status=400,
detail=detail or "The request query parameters are invalid.",
)
def unknown_query_param_problem(unknown: tuple[str, ...], allowed: tuple[str, ...]) -> ProblemDetail:
return ProblemDetail(
type=f"{PROBLEM_TYPE_BASE}unknown-query-parameter",

View file

@ -221,6 +221,8 @@ LITELLM_TRACE_CONTROL_METADATA_FIELDS: Final = frozenset(
)
_UNTRUSTED_ROOT_CONTROL_FIELDS: Final = (
"weights",
"_router_weights",
"proxy_server_request",
"standard_logging_object",
"secret_fields",
@ -334,7 +336,7 @@ _CLIENT_PRICING_METADATA_FIELDS: Final = frozenset({"model_info", "standard_logg
# and read by spend logs as fact; a client value has no legitimate meaning and no
# key or team setting keeps it, so the strip is never gated.
_ROUTER_RESERVED_METADATA_FIELDS: Final = frozenset(
{"attempted_fallbacks", "original_model_group", CLIENT_OUTPUT_CEILING_METADATA_KEY}
{"attempted_fallbacks", "original_model_group", "request_retry_count", CLIENT_OUTPUT_CEILING_METADATA_KEY}
)
_ALLOW_CLIENT_PRICING_OVERRIDE_METADATA_KEY: Final = "allow_client_pricing_override"

View file

@ -40,6 +40,14 @@ from litellm.proxy.litellm_pre_call_utils import (
LiteLLMProxyRequestSetup,
refresh_proxy_server_request_body_snapshot,
)
from litellm.proxy.management_endpoints.common_utils import (
_is_user_team_admin, # pyright: ignore[reportPrivateUsage] # shared owner of team-admin membership
)
from litellm.proxy.management_helpers.auto_router_permissions import (
authorize_member_auto_router_dependencies,
authorize_member_auto_router_team,
validate_member_auto_router_config,
)
from litellm.repositories.autorouter_session_repository import AutoRouterSessionRepository
from litellm.repositories.base_repository import SupportsModelDump
from litellm.repositories.team_repository import TeamRepository
@ -72,13 +80,13 @@ from litellm.types.management_endpoints.auto_router_endpoints import (
)
if TYPE_CHECKING:
from fastapi import APIRouter, Depends, HTTPException, Query, status
from fastapi import APIRouter, Depends, HTTPException, Query, Request, status
from litellm.proxy.utils import PrismaClient
from litellm.router import Router
else:
try:
from fastapi import APIRouter, Depends, HTTPException, Query, status
from fastapi import APIRouter, Depends, HTTPException, Query, Request, status
except ImportError:
# fastapi is only required for proxy, not for SDK usage
pass
@ -201,21 +209,14 @@ async def _query_raw(prisma_client: "PrismaClient", query: str, *args: object) -
return await prisma_client.db.query_raw(query, *args)
async def _authorize_router_dry_run(user_api_key_dict: UserAPIKeyAuth, team_id: str | None) -> None:
"""Allow exactly the callers who could create this router.
Both dry runs are gated like the write they rehearse rather than as reads: a proxy
admin, or a team admin naming their own team, matching /model/new. Routing a test
prompt can also spend money (an `llm` classifier config calls its classifier, a
semantic config embeds the prompt), so a read-level gate would be too loose anyway.
"""
async def _authorize_router_dry_run(user_api_key_dict: UserAPIKeyAuth, team_id: str | None) -> LiteLLM_TeamTable | None:
from litellm.proxy.management_endpoints.model_management_endpoints import (
ModelManagementAuthChecks,
)
from litellm.proxy.proxy_server import premium_user, prisma_client
if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN:
return
return None
if team_id is None:
raise HTTPException(
@ -244,12 +245,47 @@ async def _authorize_router_dry_run(user_api_key_dict: UserAPIKeyAuth, team_id:
},
)
ModelManagementAuthChecks.can_user_make_team_model_call(
team_id=team_id,
team: Final = LiteLLM_TeamTable.model_validate(team_row.model_dump())
if _is_user_team_admin(user_api_key_dict=user_api_key_dict, team_obj=team):
ModelManagementAuthChecks.can_user_make_team_model_call(
team_id=team_id,
user_api_key_dict=user_api_key_dict,
team_obj=team,
premium_user=premium_user,
)
return None
authorize_member_auto_router_team(
user_api_key_dict=user_api_key_dict,
team_obj=LiteLLM_TeamTable.model_validate(team_row.model_dump()),
team=team,
premium_user=premium_user,
)
return team
async def _authorize_member_dry_run_config(
*,
config: Mapping[str, object],
default_model: str | None,
user_api_key_dict: UserAPIKeyAuth,
team: LiteLLM_TeamTable,
) -> UserAPIKeyAuth:
from litellm.proxy.proxy_server import llm_router, prisma_client
if prisma_client is None or llm_router is None:
raise HTTPException(status_code=503, detail="Cannot verify auto-router model access")
validated: Final = validate_member_auto_router_config(config)
scoped_actor: Final = user_api_key_dict.model_copy(
update=MappingProxyType({"team_id": team.team_id, "team_models": team.models, "org_id": team.organization_id})
)
await authorize_member_auto_router_dependencies(
config=validated,
default_model=default_model,
user_api_key_dict=scoped_actor,
team=team,
prisma_client=prisma_client,
llm_router=llm_router,
)
return scoped_actor
def _models_this_test_can_call(config: RequestComplexityRouterConfig) -> tuple[str, ...]:
@ -326,16 +362,23 @@ async def validate_complexity_router_config(
Runs the same check every write path runs (the router's own pydantic model), so a form can
show the backend's exact verdict while the operator is still editing rather than after a
rejected save. Gated exactly like the save it rehearses: a proxy admin, or a team admin
naming their own team. Nothing is created, routed, or billed.
rejected save. Uses the same team opt-in and model-access checks as configuration
writes for members. Nothing is created, routed, or billed.
"""
await _authorize_router_dry_run(user_api_key_dict=user_api_key_dict, team_id=data.team_id)
member_team: Final = await _authorize_router_dry_run(user_api_key_dict=user_api_key_dict, team_id=data.team_id)
from litellm.router_utils.auto_router_model_naming import (
validate_complexity_router_config_write,
)
error: Final = validate_complexity_router_config_write(data.complexity_router_config)
if error is None and member_team is not None:
await _authorize_member_dry_run_config(
config=data.complexity_router_config,
default_model=None,
user_api_key_dict=user_api_key_dict,
team=member_team,
)
return ComplexityRouterConfigValidationResponse(valid=error is None, error=error)
@ -349,6 +392,7 @@ async def validate_complexity_router_config(
async def preview_auto_router_routing(
data: AutoRouterRoutingTestRequest,
user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)],
http_request: Request,
) -> AutoRouterRoutingTestResponse:
"""
Route a single request through a complexity-router config and report where it landed.
@ -392,7 +436,34 @@ async def preview_auto_router_routing(
)
from litellm.proxy.utils import get_available_models_for_user
await _authorize_router_dry_run(user_api_key_dict=user_api_key_dict, team_id=data.team_id)
member_team: Final = await _authorize_router_dry_run(user_api_key_dict=user_api_key_dict, team_id=data.team_id)
actor: Final = (
await _authorize_member_dry_run_config(
config=data.complexity_router_config.model_dump(exclude_none=True),
default_model=data.default_model,
user_api_key_dict=user_api_key_dict,
team=member_team,
)
if member_team is not None
else user_api_key_dict
)
request_data: Final[dict[str, object]] = { # mutable-ok: auth and routing enrich this request in place
**data.wire_body(),
"metadata": {}, # mutable-ok: centralized auth and identity stamping share this metadata bucket
"proxy_server_request": {"body": None}, # mutable-ok: the snapshot owner fills this body in place
}
if member_team is not None and _models_this_test_can_call(data.complexity_router_config):
from litellm.proxy.auth.user_api_key_auth import (
_run_centralized_common_checks, # pyright: ignore[reportPrivateUsage] # reuse the serving admission policy
)
await _run_centralized_common_checks(
user_api_key_auth_obj=actor,
request=http_request,
request_data=request_data,
route="/auto_router/test_routing",
)
if llm_router is None:
raise HTTPException(
@ -404,7 +475,7 @@ async def preview_auto_router_routing(
await _authorize_models_this_test_can_call(
config=data.complexity_router_config,
user_api_key_dict=user_api_key_dict,
user_api_key_dict=actor,
llm_router=llm_router,
)
@ -417,12 +488,8 @@ async def preview_auto_router_routing(
)
request_kwargs: Final = LiteLLMProxyRequestSetup.add_user_api_key_auth_to_request_metadata(
data={ # mutable-ok: the request-metadata helper takes and returns request kwargs as a dict
**data.wire_body(),
"metadata": {}, # mutable-ok: the request-metadata helper writes the auth fields into this dict
"proxy_server_request": {"body": None}, # mutable-ok: the snapshot owner fills body in place
},
user_api_key_dict=user_api_key_dict,
data=request_data,
user_api_key_dict=actor,
_metadata_variable_name="metadata",
)
refresh_proxy_server_request_body_snapshot(request_kwargs)

View file

@ -52,6 +52,7 @@ async def new_budget(
- max_parallel_requests: Optional[int] - The max number of parallel requests for the budget.
- tpm_limit: Optional[int] - The tokens per minute limit for the budget.
- rpm_limit: Optional[int] - The requests per minute limit for the budget.
- tpd_limit: Optional[int] - The tokens per day limit for the budget. Charged by batch submissions instead of tpm_limit/rpm_limit.
- model_max_budget: Optional[dict] - Specify max budget for a given model. Example: {"openai/gpt-4o-mini": {"max_budget": 100.0, "budget_duration": "1d", "tpm_limit": 100000, "rpm_limit": 100000}}
- budget_reset_at: Optional[datetime] - Datetime when the initial budget is reset. Default is now.
"""
@ -135,6 +136,7 @@ async def update_budget(
- max_parallel_requests: Optional[int] - The max number of parallel requests for the budget.
- tpm_limit: Optional[int] - The tokens per minute limit for the budget.
- rpm_limit: Optional[int] - The requests per minute limit for the budget.
- tpd_limit: Optional[int] - The tokens per day limit for the budget. Charged by batch submissions instead of tpm_limit/rpm_limit.
- model_max_budget: Optional[dict] - Specify max budget for a given model. Example: {"openai/gpt-4o-mini": {"max_budget": 100.0, "budget_duration": "1d", "tpm_limit": 100000, "rpm_limit": 100000}}
- budget_reset_at: Optional[datetime] - Update the Datetime when the budget was last reset.
"""
@ -272,6 +274,7 @@ async def budget_settings(
"max_parallel_requests": {"type": "Integer"},
"tpm_limit": {"type": "Integer"},
"rpm_limit": {"type": "Integer"},
"tpd_limit": {"type": "Integer"},
"budget_duration": {"type": "String"},
"max_budget": {"type": "Float"},
"soft_budget": {"type": "Float"},

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