scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-28 06:23:07 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 94

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

Browse source
This commit is contained in:
mateo-berri 2026-10-20 10:12:23 -07:00
commit 77d2c84a4a
189 changed files with 12033 additions and 2603 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

11
.github/PULL_REQUEST_TEMPLATE/general.md vendored
View file

@ -4,6 +4,15 @@
Drop every section you have nothing to put in, heading included: a bare "## Relevant issues" or
"## Affected release" with nothing under it must not appear in the final description -->
## Description
<!-- Two or three plain sentences a teammate can read in 10 seconds, written the way you'd explain it in Slack
Say what the PR does for users and why, then anything that keeps it working over time (a sync job, a config someone owns)
No headings, bullets, tables, or code names here. Everything else goes inside the Details block below -->
<details>
<summary>Details</summary>
## TLDR
<!-- Fill in the bullets below and keep each one short and concrete: one line per bullet, roughly 10 words max
@ -164,3 +173,5 @@ Example checklists:
## Final Attestation
- [ ] The tests check the right things, including the edge cases, and regressions in the respective real-world customer use-cases are not possible after this PR
</details>

128
enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py
View file

@ -4,6 +4,7 @@ Polls LiteLLM_ManagedObjectTable to check if the batch job is complete, and if t
import asyncio
import json
import re
from collections.abc import Mapping, Sequence
from dataclasses import replace as dataclasses_replace
from datetime import datetime, timedelta, timezone
@ -39,7 +40,12 @@ if TYPE_CHECKING:
from litellm.types.utils import LiteLLMBatch
def _error_names_resource(error: Exception, resource_id: str) -> bool:
return resource_id in re.findall(r"""[^\s"'`]+""", str(error))
CHECK_BATCH_COST_USER_AGENT = "LiteLLM Proxy/CheckBatchCost"
RECENT_BATCH_POLL_WINDOW: Final = timedelta(hours=48)
PROVIDER_TERMINAL_BATCH_STATUSES: Final[Tuple[str, ...]] = (
"completed",
@ -109,6 +115,11 @@ class CheckBatchCost:
# the guaranteed-failing primary query on every subsequent cycle.
self._has_batch_processed_column: bool = True
self.batch_processed_support_confirmed: bool = False
self._poll_cursor: tuple[datetime, str] | None = None
self._poll_before: datetime = datetime.now(timezone.utc)
self._arrival_cursor: tuple[datetime, str] | None = None
self._arrival_before: datetime = self._poll_before
self._poll_arrivals: bool = False
self._line_item_stores: Final[set[asyncio.Task[None]]] = set()
self._line_item_store_turn: Final = asyncio.Lock()
self._poll_idle: Final = asyncio.Event()
@ -161,29 +172,9 @@ class CheckBatchCost:
f"(older than {MANAGED_OBJECT_STALENESS_CUTOFF_DAYS} days) as stale_expired"
)
if not self._has_batch_processed_column:
return
# A row already in a terminal status is never rewritten by the sweep above, so
# without this it keeps a poll-page slot forever and starves newer batches.
retired: Final = await _managed_object_table(self.prisma_client).update_many(
where={
"file_purpose": "batch",
"batch_processed": False,
"status": {"in": ["complete", "completed"]},
"created_at": {"lt": cutoff},
},
data={"batch_processed": True},
)
if retired > 0:
verbose_proxy_logger.warning(
f"CheckBatchCost: gave up on {retired} completed managed objects older than "
f"{MANAGED_OBJECT_STALENESS_CUTOFF_DAYS} days that were never costed"
)
async def _fallback_find_jobs(self) -> "Sequence[ManagedObjectRow]":
"""Query batch jobs without the batch_processed filter (for older schemas)."""
return await _managed_object_table(self.prisma_client).find_many(
return await self._find_jobs(
where={
"file_purpose": "batch",
"status": {
@ -197,10 +188,59 @@ class CheckBatchCost:
]
},
},
take=MAX_OBJECTS_PER_POLL_CYCLE,
order={"created_at": "asc"},
)
async def _find_jobs(self, where: Mapping[str, object]) -> Sequence[ManagedObjectRow]:
arrivals: Final = self._poll_arrivals
page: Final = await self._find_job_page(where, arrivals=arrivals)
self._poll_arrivals = not arrivals if page else arrivals
if page:
return page
return await self._find_job_page(where, arrivals=not arrivals)
async def _find_job_page(self, where: Mapping[str, object], *, arrivals: bool) -> Sequence[ManagedObjectRow]:
cursor: Final = self._arrival_cursor if arrivals else self._poll_cursor
if cursor is None:
if arrivals:
self._arrival_before = datetime.now(timezone.utc)
else:
self._poll_before = datetime.now(timezone.utc) - RECENT_BATCH_POLL_WINDOW
before: Final = self._arrival_before if arrivals else self._poll_before
page: Final = await _managed_object_table(self.prisma_client).find_many(
where={
**where,
"created_at": {
"lte": before,
**({"gt": before - RECENT_BATCH_POLL_WINDOW} if arrivals else {}),
},
**(
{
"OR": [
{"created_at": {"gt": cursor[0]}},
{"created_at": cursor[0], "id": {"gt": cursor[1]}},
]
}
if cursor is not None
else {}
),
},
take=MAX_OBJECTS_PER_POLL_CYCLE,
order=[{"created_at": "asc"}, {"id": "asc"}],
)
next_cursor: Final = (page[-1].created_at, page[-1].id) if len(page) == MAX_OBJECTS_PER_POLL_CYCLE else None
if arrivals:
self._arrival_cursor = next_cursor
else:
self._poll_cursor = next_cursor
return page
def _unprocessed_snapshot(self, job: ManagedObjectRow) -> Mapping[str, object]:
return {
"id": job.id,
"updated_at": job.updated_at,
**({"batch_processed": False} if self._has_batch_processed_column else {}),
}
async def _retire_job(self, job: "ManagedObjectRow", reason: str) -> None:
"""
Take a row that can never be costed out of the poll page. Leaving it selectable
@ -208,24 +248,19 @@ class CheckBatchCost:
once enough such rows accumulate no newer batch is ever reached. Older schemas
without batch_processed can only be excluded through the status filter.
"""
data: Final = (
{"batch_processed": True}
if self._has_batch_processed_column
else {"status": "stale_expired"}
)
data: Final = {"batch_processed": True} if self._has_batch_processed_column else {"status": "stale_expired"}
try:
await _managed_object_table(self.prisma_client).update(
where={"id": job.id},
retired: Final = await _managed_object_table(self.prisma_client).update_many(
where=self._unprocessed_snapshot(job),
data=data,
)
except Exception as db_err:
verbose_proxy_logger.error(
f"CheckBatchCost: failed to retire uncostable job {job.id} ({reason}): {db_err}"
)
verbose_proxy_logger.error(f"CheckBatchCost: failed to retire uncostable job {job.id} ({reason}): {db_err}")
return
if retired == 0:
return
verbose_proxy_logger.warning(
f"CheckBatchCost: job {job.id} can never be costed ({reason}), "
"so it will no longer be polled"
f"CheckBatchCost: job {job.id} can never be costed ({reason}), so it will no longer be polled"
)
async def _claim_job_for_costing(self, job: "ManagedObjectRow", response: "LiteLLMBatch") -> bool:
@ -312,13 +347,13 @@ class CheckBatchCost:
from litellm.exceptions import NotFoundError
return isinstance(error, (NotFoundError, openai.NotFoundError)) and batch_id in str(error)
return isinstance(error, (NotFoundError, openai.NotFoundError)) and _error_names_resource(error, batch_id)
def _batch_deployment_exists(self, model_id: str) -> bool:
"""A 404 only proves the batch is gone when it came from the batch's own
deployment. Once that deployment leaves the router, default fallbacks can
silently send the retrieve to a provider that never saw the batch, so its
404 must not retire the row; the staleness sweep bounds it instead."""
404 must not retire a potentially recoverable row."""
return self.llm_router.get_deployment(model_id=model_id) is not None
@staticmethod
@ -333,7 +368,7 @@ class CheckBatchCost:
if not output_file_id:
return False
return isinstance(error, (NotFoundError, openai.NotFoundError)) and output_file_id in str(error)
return isinstance(error, (NotFoundError, openai.NotFoundError)) and _error_names_resource(error, output_file_id)
async def _finalize_unbilled_terminal_job(
self, job: "ManagedObjectRow", response: "LiteLLMBatch"
@ -355,18 +390,18 @@ class CheckBatchCost:
db_batch_object=job,
unified_batch_id=_is_base64_encoded_unified_file_id(job.unified_object_id),
)
update_data: Final[dict] = {
update_data: Final[dict[str, object]] = {
"status": response.status,
"file_object": response.model_dump_json(),
**({"batch_processed": True} if self._has_batch_processed_column else {}),
}
await _managed_object_table(self.prisma_client).update(
where={"id": job.id},
finalized: Final = await _managed_object_table(self.prisma_client).update_many(
where=self._unprocessed_snapshot(job),
data=update_data,
)
verbose_proxy_logger.info(
f"CheckBatchCost: marked job {job.id} as {response.status} in DB"
)
if finalized == 0:
return False
verbose_proxy_logger.info(f"CheckBatchCost: marked job {job.id} as {response.status} in DB")
return True
except Exception as db_err:
verbose_proxy_logger.error(
@ -1047,6 +1082,7 @@ class CheckBatchCost:
"""
try:
from litellm.integrations.prometheus import PrometheusLogger
prom_logger = PrometheusLogger.get_instance()
except Exception as e:
verbose_proxy_logger.error(f"CheckBatchCost: could not get Prometheus logger: {e}")
@ -1071,7 +1107,7 @@ class CheckBatchCost:
# endpoint may transition a batch to "complete" before
# CheckBatchCost runs. The batch_processed=False filter
# already prevents reprocessing finished batches.
jobs = await _managed_object_table(self.prisma_client).find_many(
jobs = await self._find_jobs(
where={
"file_purpose": "batch",
"batch_processed": False,
@ -1084,8 +1120,6 @@ class CheckBatchCost:
]
},
},
take=MAX_OBJECTS_PER_POLL_CYCLE,
order={"created_at": "asc"},
)
self.batch_processed_support_confirmed = True
except Exception as query_err:

26
enterprise/litellm_enterprise/proxy/common_utils/check_responses_cost.py
View file

@ -59,6 +59,19 @@ if TYPE_CHECKING:
from litellm.router import Router
_STORED_RESPONSE_ADAPTER: Final[TypeAdapter[dict[str, object]]] = TypeAdapter(dict[str, object])
class _RoutedModelInfo(BaseModel):
model_config = ConfigDict(frozen=True)
id: str | None = None
class _RoutedDeployment(BaseModel):
model_config = ConfigDict(frozen=True)
model_info: _RoutedModelInfo = _RoutedModelInfo()
_ROUTED_DEPLOYMENTS: Final = TypeAdapter(tuple[_RoutedDeployment, ...])
_RETRIEVE_RESPONSE_CALL_TYPE: Final = "aget_responses"
@ -314,10 +327,23 @@ class CheckResponsesCost:
INTERNAL_CALL_ORIGIN_METADATA_KEY: BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN,
**target.model_metadata,
**(_deployment_metadata(deployment) if deployment is not None else {}),
**self._requested_model_group(target, job.team_id),
}
),
)
def _requested_model_group(self, target: _PollTarget, team_id: str | None) -> Mapping[str, object]:
requested: Final = target.model_metadata.get("model_group")
if target.deployment is None or not isinstance(requested, str):
return MappingProxyType({})
routed: Final = _ROUTED_DEPLOYMENTS.validate_python(
self.llm_router.get_model_list(model_name=requested, team_id=team_id) or ()
)
serving_id: Final = target.deployment.model_info.id
return MappingProxyType(
{"model_group": requested} if any(row.model_info.id == serving_id for row in routed) else {}
)
async def _bill_once(
self,
job: ManagedObjectRow,

62
enterprise/litellm_enterprise/proxy/common_utils/managed_object_attribution.py
View file

@ -9,6 +9,7 @@ from __future__ import annotations
from collections.abc import Mapping
from dataclasses import dataclass
from datetime import datetime
from typing import TYPE_CHECKING, Final, Protocol, cast
from pydantic import TypeAdapter
@ -27,6 +28,12 @@ if TYPE_CHECKING:
class ManagedObjectRow(Protocol):
@property
def created_at(self) -> datetime: ...
@property
def updated_at(self) -> datetime: ...
@property
def id(self) -> str: ...
@ -107,8 +114,10 @@ class ManagedObjectCreatorAttribution:
user's alias, because a resolvable name is more useful on the spend row than a null.
user_api_key_org_id is resolved too: the spend update writer reads it off this metadata
to increment organization spend, so leaving it out drops the cost from org accounting.
A background response also carries its end user and the key, team, and user budgets, the
way the read that billed it before did, and never names a key-scoped creator as its user.
A background response also carries its end user, the creating key's project and its alias,
and the key, team, and user budgets and per-model budgets, the way the read that billed it
before did, and never names a key-scoped creator as its user. A key that sets no budget of
its own takes its budget tier's, the same rule auth applies.
"""
def __init__(self, prisma_client: PrismaClient, poller_name: str) -> None:
@ -126,9 +135,23 @@ class ManagedObjectCreatorAttribution:
"user_api_key_user_id": _creator_user_id(job.created_by),
**_shared_metadata(job, rows),
**({"user_api_key_end_user_id": end_user_id} if end_user_id else {}),
**_budget_fields(rows.key, "user_api_key_spend", "user_api_key_max_budget"),
**_budget_fields(rows.team, "user_api_key_team_spend", "user_api_key_team_max_budget"),
**_budget_fields(rows.user, "user_api_key_user_spend", "user_api_key_user_max_budget"),
"user_api_key_project_id": getattr(rows.key, "project_id", None),
"user_api_key_project_alias": getattr(
getattr(rows.key, "litellm_project_table", None), "project_alias", None
),
**_key_budget_fields(rows.key),
**_budget_fields(
rows.team,
"user_api_key_team_spend",
"user_api_key_team_max_budget",
"user_api_key_team_model_max_budget",
),
**_budget_fields(
rows.user,
"user_api_key_user_spend",
"user_api_key_user_max_budget",
"user_api_key_user_model_max_budget",
),
}
async def _rows(self, job: ManagedObjectRow, object_id: str) -> _CreatorRows:
@ -155,7 +178,9 @@ class ManagedObjectCreatorAttribution:
if not api_key or _is_cli_session_key(api_key, created_by):
return None
try:
return await _token_table(self._prisma_client).find_unique(where={"token": api_key})
return await _token_table(self._prisma_client).find_unique(
where={"token": api_key}, include={"litellm_budget_table": True, "litellm_project_table": True}
)
except Exception as e:
verbose_proxy_logger.error(f"{self._poller_name}: could not look up the creating key for {object_id}: {e}")
return None
@ -212,7 +237,28 @@ def _organization_id(row: object) -> str | None:
return cast("str | None", getattr(row, "organization_id", None)) if row is not None else None
def _budget_fields(row: object, spend_key: str, max_budget_key: str) -> Mapping[str, object]:
def _budget_fields(row: object, spend_key: str, max_budget_key: str, model_max_budget_key: str) -> Mapping[str, object]:
if row is None:
return {}
return {spend_key: getattr(row, "spend", None), max_budget_key: getattr(row, "max_budget", None)}
return {
spend_key: getattr(row, "spend", None),
max_budget_key: getattr(row, "max_budget", None),
model_max_budget_key: getattr(row, "model_max_budget", None),
}
def _key_budget_fields(key: object) -> Mapping[str, object]:
if key is None:
return {}
tier: Final = getattr(key, "litellm_budget_table", None)
return {
"user_api_key_spend": getattr(key, "spend", None),
"user_api_key_max_budget": _own_or_tier_budget(key, tier, "max_budget"),
"user_api_key_model_max_budget": _own_or_tier_budget(key, tier, "model_max_budget"),
}
def _own_or_tier_budget(key: object, tier: object, field: str) -> object:
own: Final = getattr(key, field, None)
tier_value: Final = getattr(tier, field, None)
return own if tier_value is None or own not in (None, {}) else tier_value

4
litellm/a2a_protocol/main.py
View file

@ -602,9 +602,9 @@ def _build_streaming_logging_obj(
litellm_params: dict[str, object] | None,
metadata: dict[str, object] | None,
proxy_server_request: dict[str, object] | None,
start_time: datetime.datetime,
) -> Logging:
"""Build logging object for streaming A2A requests."""
start_time: Final = datetime.datetime.now()
model: Final = f"a2a_agent/{agent_name}"
logging_obj: Final = Logging(
@ -718,6 +718,7 @@ async def asend_message_streaming(
if request is None:
raise ValueError("request is required")
called_at: Final = datetime.datetime.now()
_raw_logging_obj: Final = kwargs.get("litellm_logging_obj")
logging_obj: Logging | None = _raw_logging_obj if isinstance(_raw_logging_obj, Logging) else None
@ -751,6 +752,7 @@ async def asend_message_streaming(
litellm_params=litellm_params,
metadata=metadata,
proxy_server_request=proxy_server_request,
start_time=called_at,
)
verbose_logger.info("A2A send_message_streaming request_id=%s, agent=%s", request.id, agent_name)

4
litellm/a2a_protocol/streaming_iterator.py
View file

@ -41,7 +41,7 @@ class A2AStreamingIterator:
self.request = request
self.logging_obj = logging_obj
self.agent_name = agent_name
self.start_time = datetime.now()
self.start_time = logging_obj.start_time
# Collect chunks for token counting
self.chunks: list[SendStreamingMessageResponse] = []
@ -54,6 +54,8 @@ class A2AStreamingIterator:
async def __anext__(self) -> "SendStreamingMessageResponse":
try:
chunk: Final = await self.stream.__anext__()
if self.logging_obj.completion_start_time is None:
self.logging_obj.update_completion_start_time(completion_start_time=datetime.now())
# Store chunk
self.chunks.append(chunk)

26
litellm/caching/dual_cache.py
View file

@ -59,16 +59,6 @@ class PendingBatchRead:
previous_access_times: dict[str, float | None]
@dataclass(frozen=True, slots=True)
class DeclaredBatchRead:
"""A ``async_batch_get_cache`` split in two: the memory half done, the Redis half declared on a ``RedisBatch``
so it rides that batch's next round trip, resolved later with ``async_resolve_batch_get``."""
keys: tuple[str, ...]
pending: PendingBatchRead
result: BatchResult[Mapping[str, object]] | None
def _log_deferred_increment_failure(future: asyncio.Future[float]) -> None:
if future.cancelled():
return
@ -382,22 +372,6 @@ class DualCache(BaseCache):
await self.in_memory_cache.async_set_cache(key, value, **self._backfill_kwargs(kwargs))
return merged
async def declare_batch_get(self, keys: Sequence[str], batch: RedisBatch) -> DeclaredBatchRead:
pending: Final = await self._prepare_batch_get(
list(keys),
local_only=False,
throttle_redis=False,
)
return DeclaredBatchRead(
keys=tuple(keys),
pending=pending,
result=batch.mget(pending.redis_keys) if pending.redis_keys else None,
)
async def async_resolve_batch_get(self, declared: DeclaredBatchRead) -> list[object | None]:
redis_result: Final = None if declared.result is None else await declared.result
return await self._apply_batch_get(declared.pending, redis_result)
async def async_batch_get_cache(
self,
keys: list,

47
litellm/caching/redis_cache.py
View file

@ -134,7 +134,12 @@ _INCREMENT_WITH_FLOOR_LUA: Final = (
"return count"
)
_INCREMENT_IF_PRESENT_LUA: Final = (
"if redis.call('EXISTS', KEYS[1]) == 0 then return false end return redis.call('INCRBYFLOAT', KEYS[1], ARGV[1])"
)
_LUA_COUNT: Final = TypeAdapter(int)
_LUA_OPTIONAL_TOTAL: Final[TypeAdapter[float | None]] = TypeAdapter(float | None)
_OPTIONAL_COUNTS: Final = TypeAdapter(tuple[int | None, ...])
@ -675,7 +680,7 @@ class RedisCache(BaseCache):
# redis namespaces
self.namespace = namespace
# for high traffic, we store the redis results in memory and then batch write to redis
self.redis_batch_writing_buffer: list = []
self.redis_batch_writing_buffer: list[tuple[str, object]] = []
if redis_flush_size is None:
self.redis_flush_size: int = 100
else:
@ -1568,6 +1573,46 @@ class RedisCache(BaseCache):
count: Final = await _redis_client.eval(_INCREMENT_WITH_FLOOR_LUA, 1, namespaced_key, value, ttl)
return _LUA_COUNT.validate_python(count)
@_redis_circuit_breaker_guard
async def async_increment_if_present(
self, key: str, value: float, parent_otel_span: Span | None = None
) -> float | None:
_redis_client: Final = self._async_commands()
namespaced_key: Final = self.check_and_fix_namespace(key=key)
start_time: Final = time.time()
try:
total: Final = _LUA_OPTIONAL_TOTAL.validate_python(
await _redis_client.eval(_INCREMENT_IF_PRESENT_LUA, 1, namespaced_key, value)
)
except Exception as e:
failed_at: Final = time.time()
asyncio.create_task(
self.service_logger_obj.async_service_failure_hook(
service=ServiceTypes.REDIS,
duration=failed_at - start_time,
error=e,
call_type="async_increment_if_present",
caller=_get_call_stack_info(),
start_time=start_time,
end_time=failed_at,
parent_otel_span=parent_otel_span,
)
)
raise
end_time: Final = time.time()
asyncio.create_task(
self.service_logger_obj.async_service_success_hook(
service=ServiceTypes.REDIS,
duration=end_time - start_time,
call_type="async_increment_if_present",
caller=_get_call_stack_info(),
start_time=start_time,
end_time=end_time,
parent_otel_span=parent_otel_span,
)
)
return total
async def flush_cache_buffer(self):
print_verbose(f"flushing to redis....reached size of buffer {len(self.redis_batch_writing_buffer)}")
await self.async_set_cache_pipeline(self.redis_batch_writing_buffer)

2
litellm/cost_calculator.py
View file

@ -1642,7 +1642,7 @@ def completion_cost(
*(((model, request_pricing_provider),) if model is not None else ()),
)
potential_model_names: Final = (
raw_candidates if explicit_pricing else _priced_candidates_before_unpriced_rules(raw_candidates)
raw_candidates if custom_pricing is True else _priced_candidates_before_unpriced_rules(raw_candidates)
)
for idx, (model, pricing_provider) in enumerate(potential_model_names):

2
litellm/decisions/call.py
View file

@ -427,7 +427,7 @@ def _jsonable_fallback(value: object) -> object:
def _logged_request_message(caller_request: object) -> ChatCompletionUserMessage:
caller_fields: Final = _LOGGED_FIELDS_ADAPTER.validate_python(
to_jsonable_python(caller_request, fallback=_jsonable_fallback)
to_jsonable_python(caller_request, fallback=_jsonable_fallback, bytes_mode="base64")
)
sent_fields: Final = json.dumps(
dict((name, value) for name, value in caller_fields.items() if value not in (None, [], {}))

4
litellm/google_genai/streaming_iterator.py
View file

@ -71,7 +71,7 @@ class BaseGoogleGenAIGenerateContentStreamingIterator:
):
self.litellm_logging_obj = litellm_logging_obj
self.request_body = request_body
self.start_time = datetime.now()
self.start_time = litellm_logging_obj.start_time
self.collected_chunks: list[bytes] = []
self.model = model
self.custom_llm_provider = custom_llm_provider
@ -192,6 +192,8 @@ class AsyncGoogleGenAIGenerateContentStreamingIterator(BaseGoogleGenAIGenerateCo
async def __anext__(self):
try:
chunk: Final = await _anext_google_genai_sse_chunk(self.stream_iterator)
if self.litellm_logging_obj.completion_start_time is None:
self.litellm_logging_obj.update_completion_start_time(completion_start_time=datetime.now())
self.collected_chunks.append(chunk)
return chunk
except StopAsyncIteration:

21
litellm/integrations/anthropic_cache_control_hook.py
View file

@ -342,6 +342,20 @@ EXTERNAL_BREAKPOINTS_STAMP: Final = "_litellm_external_breakpoints"
CUSTOM_LLM_PROVIDER_STAMP: Final = "_litellm_custom_llm_provider"
def _accepts_cache_control(part: object) -> bool:
match part:
case {"type": "thinking" | "redacted_thinking"}:
return False
case dict():
return True
case _:
return False
def _last_index_accepting_cache_control(parts: Sequence[object]) -> int | None:
return next((index for index in reversed(range(len(parts))) if _accepts_cache_control(parts[index])), None)
class AnthropicCacheControlHook(CustomPromptManagement):
@staticmethod
def _request_value(request_kwargs: object, key: str) -> object:
@ -740,10 +754,11 @@ class AnthropicCacheControlHook(CustomPromptManagement):
# 1. if string, insert cache control in the message
if isinstance(message_content, str):
message["cache_control"] = control
# 2. list of objects - only apply to last item per Anthropic spec
# 2. list of objects - only apply to the last item Claude accepts a breakpoint on
elif isinstance(message_content, list):
if len(message_content) > 0 and isinstance(message_content[-1], dict):
message_content[-1]["cache_control"] = control # pyright: ignore[reportGeneralTypeIssues] # loose runtime dict
markable_index: Final = _last_index_accepting_cache_control(message_content)
if markable_index is not None:
message_content[markable_index]["cache_control"] = control # pyright: ignore[reportGeneralTypeIssues] # loose runtime dict
return message
@staticmethod

24
litellm/integrations/datadog/datadog_llm_obs.py
View file

@ -12,7 +12,7 @@ import os
from collections.abc import Mapping, Sequence
from datetime import datetime
from types import MappingProxyType
from typing import Any, Final, Literal
from typing import Final, Literal
import httpx
from pydantic import ConfigDict, TypeAdapter, ValidationError
@ -54,6 +54,7 @@ from litellm.types.utils import (
PROMPT_CARRYING_GUARDRAIL_FIELDS,
PROMPT_QUOTING_ROUTING_DECISION_FIELDS,
CallTypes,
StandardLoggingMetadata,
StandardLoggingPayload,
StandardLoggingPayloadErrorInformation,
)
@ -93,8 +94,12 @@ _ROUTER_DIMENSIONS: Final[tuple[str, ...]] = ("router_tier", "router_cause", "ro
_COST_DIMENSIONS: Final[tuple[str, ...]] = ("team", "user", "key_alias", "model_group", *_ROUTER_DIMENSIONS)
def _metadata_of(standard_logging_payload: StandardLoggingPayload) -> Mapping[str, Any]:
metadata: Final = standard_logging_payload.get("metadata")
def _logging_metadata_of(standard_logging_payload: StandardLoggingPayload) -> StandardLoggingMetadata | None:
return standard_logging_payload.get("metadata")
def _metadata_of(standard_logging_payload: StandardLoggingPayload) -> Mapping[str, object]:
metadata: Final = _logging_metadata_of(standard_logging_payload)
return metadata or _EMPTY_MAPPING
@ -1057,14 +1062,23 @@ class DataDogLLMObsLogger(CustomBatchLogger):
# Get budget information from metadata
metadata: Final = _metadata_of(standard_logging_payload)
logging_metadata: Final = _logging_metadata_of(standard_logging_payload)
# API key max budget
user_api_key_max_budget: Final = metadata.get("user_api_key_max_budget")
user_api_key_max_budget: Final = (
logging_metadata["user_api_key_max_budget"]
if logging_metadata and "user_api_key_max_budget" in logging_metadata
else None
)
if user_api_key_max_budget is not None:
spend_metrics["user_api_key_max_budget"] = float(user_api_key_max_budget)
# API key spend
user_api_key_spend: Final = metadata.get("user_api_key_spend")
user_api_key_spend: Final = (
logging_metadata["user_api_key_spend"]
if logging_metadata and "user_api_key_spend" in logging_metadata
else None
)
if user_api_key_spend is not None:
try:
spend_metrics["user_api_key_spend"] = float(user_api_key_spend)

15
litellm/integrations/langfuse/langfuse_otel.py
View file

@ -1,7 +1,7 @@
import base64
import json
import os
from collections.abc import Sequence
from collections.abc import Mapping, Sequence
from datetime import datetime
from typing import TYPE_CHECKING, Any, Final, Optional
@ -9,6 +9,7 @@ from litellm._logging import verbose_logger
from litellm.integrations.arize import _utils
from litellm.integrations.langfuse.langfuse_otel_attributes import (
LangfuseLLMObsOTELAttributes,
serialize_langfuse_datetime,
)
from litellm.integrations.opentelemetry import OpenTelemetry, OpenTelemetryConfig
from litellm.litellm_core_utils.safe_json_loads import safe_json_loads
@ -55,7 +56,7 @@ class LangfuseOtelLogger(OpenTelemetry):
LangfuseOtelLogger._set_langfuse_specific_attributes(span=span, kwargs=kwargs, response_obj=response_obj)
@staticmethod
def _extract_langfuse_metadata(kwargs: dict) -> dict:
def _extract_langfuse_metadata(kwargs: Mapping[str, Any]) -> dict:
"""
Extracts Langfuse metadata from the standard LiteLLM kwargs structure.
@ -221,7 +222,7 @@ class LangfuseOtelLogger(OpenTelemetry):
)
@staticmethod
def _set_langfuse_specific_attributes(span: Span, kwargs, response_obj):
def _set_langfuse_specific_attributes(span: Span, kwargs: Mapping[str, Any], response_obj):
"""
Sets Langfuse specific metadata attributes onto the OTEL span.
@ -255,6 +256,14 @@ class LangfuseOtelLogger(OpenTelemetry):
safe_dumps(messages),
)
completion_start_time: Final = kwargs.get("completion_start_time")
if isinstance(completion_start_time, datetime):
safe_set_attribute(
span,
LangfuseSpanAttributes.OBSERVATION_COMPLETION_START_TIME.value,
serialize_langfuse_datetime(completion_start_time),
)
LangfuseOtelLogger._set_observation_output(span=span, response_obj=response_obj)
@staticmethod

5
litellm/integrations/langfuse/langfuse_otel_attributes.py
View file

@ -6,6 +6,7 @@ Relevant Issue: https://github.com/BerriAI/litellm/issues/13764
import json
from collections.abc import Mapping
from datetime import datetime
from typing import TYPE_CHECKING, Any, Final
from pydantic import BaseModel
@ -29,6 +30,10 @@ if TYPE_CHECKING:
from opentelemetry.trace import Span
def serialize_langfuse_datetime(value: datetime) -> str:
return json.dumps(value.astimezone().isoformat())
def get_output_content_by_type(
response_obj: None
| dict

3
litellm/integrations/langfuse/langfuse_sdk.py
View file

@ -39,6 +39,7 @@ from pydantic import ConfigDict
import litellm
from litellm._logging import verbose_logger
from litellm.integrations.langfuse.langfuse import PROMPT_CACHE_TTL_ENV, parse_langfuse_debug, whole_number
from litellm.integrations.langfuse.langfuse_otel_attributes import serialize_langfuse_datetime
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
from litellm.llms.custom_httpx.http_handler import HTTPHandler, get_httpx_client
from litellm.types.llms.base import LiteLLMBaseModel
@ -168,7 +169,7 @@ def _string_or_none(value: object) -> str | None:
def _serialize_datetime(value: object) -> str | None:
"""A datetime the way the SDK's ``EventSerializer`` sends one: a JSON string, naive values read as local time."""
if isinstance(value, datetime):
return safe_dumps(value.astimezone().isoformat())
return serialize_langfuse_datetime(value)
return _serialize(value)

2
litellm/integrations/otel/README.md
View file

@ -91,7 +91,7 @@ way (`pod_lock`, `budget_reset`, `spend_queue`, `health_check`, `scheduler_queue
rather than a bare `redis.set`; `tests/unit/test_internal_context.py` scans every
module under `litellm/` and `enterprise/` that calls a shared cache or declares a
read or write on the request batch (`reserve_redis_batch_reads`,
`declare_batch_get`, `batch.mget`, `batch.set`, `batch.script`) and fails when
`batch.mget`, `batch.set`, `batch.script`) and fails when
one has no declared family, with the process-local `InMemoryCache` callers listed
as the only exemptions. A batch op carries the family that was active when it was
declared, so the routing prefetch armed before deployment selection

1
litellm/integrations/otel/logger.py
View file

@ -594,6 +594,7 @@ class OpenTelemetryV2(CustomLogger):
payload,
capture_content=self._capture_span_content(),
time_to_first_chunk_seconds=call.time_to_first_chunk_seconds,
completion_start_seconds=call.completion_start_seconds,
request_route=request_root_http_route(),
request_purpose=call.purpose,
trace=call.trace,

7
litellm/integrations/otel/mappers/langfuse.py
View file

@ -12,8 +12,10 @@ the JSON-serialized payloads. ``trace_attributes`` maps the caller's trace contr
import json
from collections.abc import Callable
from datetime import datetime, timezone
from typing import Final
from litellm.integrations.langfuse.langfuse_otel_attributes import serialize_langfuse_datetime
from litellm.integrations.otel.mappers.base import AttributeMap, AttrValue, SpanData
from litellm.integrations.otel.mappers.utils import (
collect,
@ -45,6 +47,11 @@ class LangfuseMapper:
"langfuse.observation.id": lambda d: d.identity.call_id or None,
"langfuse.trace.metadata.team_id": lambda d: d.identity.team_id or None,
"langfuse.trace.metadata.team_alias": lambda d: d.identity.team_alias or None,
"langfuse.observation.completion_start_time": lambda d: (
serialize_langfuse_datetime(datetime.fromtimestamp(d.completion_start_seconds, tz=timezone.utc))
if d.completion_start_seconds is not None
else None
),
}
# Sub-tables folded into their respective JSON blobs.

2
litellm/integrations/otel/model/metadata.py
View file

@ -258,6 +258,7 @@ class LLMCallEvent:
# needs to be reasonable for a span that never gets closed (a leak).
provisional_span_name: str
time_to_first_chunk_seconds: float | None
completion_start_seconds: float | None
trace: TraceControls
session_id: str | None
@ -280,6 +281,7 @@ class LLMCallEvent:
cache_hit=bool(payload.get("cache_hit")) if payload else False, # pyright: ignore[reportUnknownMemberType] # TypedDict .get on a str key is partially unknown
provisional_span_name=f"{operation.value} {model}".strip(),
time_to_first_chunk_seconds=time_to_first_chunk_seconds(kwargs),
completion_start_seconds=_epoch_seconds(kwargs.get("completion_start_time")),
trace=trace,
session_id=caller_session_id(kwargs, trace),
)

3
litellm/integrations/otel/model/payloads.py
View file

@ -429,6 +429,7 @@ class LLMCallSpanData:
choices_out: tuple[Mapping[str, object], ...] = ()
system_fingerprint: str | None = None
time_to_first_chunk_seconds: float | None = None
completion_start_seconds: float | None = None
# The requested output modality, set only on the routes that pin one (image
# generation, speech, transcription, OCR), and the litellm route itself, which
# keeps routes the convention folds into one operation distinguishable.
@ -450,6 +451,7 @@ class LLMCallSpanData:
payload: StandardLoggingPayload,
capture_content: bool = False,
time_to_first_chunk_seconds: float | None = None,
completion_start_seconds: float | None = None,
request_route: str | None = None,
request_purpose: str | None = None,
trace: TraceControls | None = None,
@ -503,6 +505,7 @@ class LLMCallSpanData:
choices_out=choices_out if capture_content else (),
system_fingerprint=as_str(response.get("system_fingerprint")),
time_to_first_chunk_seconds=time_to_first_chunk_seconds,
completion_start_seconds=completion_start_seconds,
output_type=resolve_output_type(call_type),
call_type=call_type or None,
request_route=request_route or context.identity.request_route,

1
litellm/integrations/otel/model/spans.py
View file

@ -214,6 +214,7 @@ _SERVICE_VERB_BY_CALL_TYPE: Final[dict[str, str]] = {
"increment_cache": "incr",
"async_increment": "incr",
"async_increment_pipeline": "incr",
"async_increment_if_present": "incr",
"delete_cache": "delete",
"async_delete_cache": "delete",
"async_rpush": "rpush",

4
litellm/interactions/streaming_iterator.py
View file

@ -48,7 +48,7 @@ class BaseInteractionsAPIStreamingIterator:
self.finished = False
self.interactions_api_config = interactions_api_config
self.completed_response: InteractionsAPIStreamingResponse | None = None
self.start_time = datetime.now()
self.start_time = logging_obj.start_time
# set request kwargs
self.litellm_metadata = litellm_metadata
@ -87,6 +87,8 @@ class BaseInteractionsAPIStreamingIterator:
# Format as InteractionsAPIStreamingResponse
if isinstance(parsed_chunk, dict):
if self.logging_obj.completion_start_time is None:
self.logging_obj.update_completion_start_time(completion_start_time=datetime.now())
streaming_response: Final = self.interactions_api_config.transform_streaming_response(
model=self.model,
parsed_chunk=parsed_chunk,

23
litellm/litellm_core_utils/cloud_storage_security.py
View file

@ -1,6 +1,7 @@
import posixpath
import re
from collections.abc import Mapping, Sequence
from itertools import accumulate, repeat
from types import MappingProxyType
from typing import Any, Final, cast
from urllib.parse import quote, unquote
@ -20,15 +21,33 @@ MANAGED_CLOUD_STORAGE_SCHEMES: Final = ("s3://", "gs://")
_MAPPING_PROXY_TYPE: Final[type] = type(MappingProxyType({}))
MAX_FILE_ID_DECODE_PASSES: Final = 8
def _unquote_once(current: str, _: str) -> str:
return unquote(current)
def _decodings(value: str) -> tuple[str, ...]:
return tuple(accumulate(repeat(value, MAX_FILE_ID_DECODE_PASSES + 2), _unquote_once))
def is_managed_cloud_storage_uri(file_id: str) -> bool:
"""
True if file_id is a raw cloud-storage object URI (e.g. ``s3://bucket/key``).
True if file_id is a raw cloud-storage object URI (e.g. ``s3://bucket/key``),
up to ``MAX_FILE_ID_DECODE_PASSES`` layers of percent-encoding deep, or an id
encoded deeper than that, which no client produces and which is treated as raw
rather than decoded any further.
These are internal provider artifacts. On the multi-tenant proxy they must be
retrieved through their managed unified file id so owner/team access is enforced;
a raw URI supplied by a caller bypasses that check.
"""
return isinstance(file_id, str) and file_id.startswith(MANAGED_CLOUD_STORAGE_SCHEMES)
if not isinstance(file_id, str):
return False
decodings: Final = _decodings(file_id)
settled: Final = decodings[-1] == decodings[-2]
return decodings[-1].startswith(MANAGED_CLOUD_STORAGE_SCHEMES) or not settled
_SAFE_OBJECT_COMPONENT_PATTERN: Final = re.compile(r"[^A-Za-z0-9._-]+")

8
litellm/litellm_core_utils/litellm_logging.py
View file

@ -2015,7 +2015,7 @@ class Logging(LiteLLMLoggingBaseClass):
router_model_id = self.get_router_model_id()
## RESPONSE COST ##
custom_pricing: Final = self._custom_pricing_for(priced_result)
custom_pricing: Final = self.custom_pricing_for(priced_result)
prompt: Final = self._prompt_for_cost_calculation()
@ -2190,7 +2190,7 @@ class Logging(LiteLLMLoggingBaseClass):
model=model,
completion_response=result,
custom_llm_provider=custom_llm_provider,
custom_pricing=self._custom_pricing_for(result),
custom_pricing=self.custom_pricing_for(result),
base_model=get_base_model_from_metadata(model_call_details=self.model_call_details),
router_model_id=router_model_id or self.get_router_model_id(),
region_name=_resolve_mantle_region_for_cost(
@ -2215,7 +2215,7 @@ class Logging(LiteLLMLoggingBaseClass):
usage=usage,
)
def _custom_pricing_for(self, result: object) -> bool:
def custom_pricing_for(self, result: object) -> bool:
litellm_params: Final = getattr(self, "litellm_params", None)
result_hidden_params: Final = getattr(result, HIDDEN_PARAMS_ATTR, None) or MappingProxyType({})
additional_headers: Final = (
@ -7235,7 +7235,7 @@ def get_standard_logging_object_payload(
# Fall back to the actual deployment model set by the router in metadata.
if base_model is None:
base_model = metadata.get("deployment")
custom_pricing: Final = use_custom_pricing_for_model(litellm_params=litellm_params)
custom_pricing: Final = logging_obj.custom_pricing_for(init_response_obj)
raw_response_cost: Final = kwargs.get("response_cost")
llm_response_cost: Final[float] = raw_response_cost or 0.0
guardrail_cost: Final = guardrail_information_cost(metadata.get("standard_logging_guardrail_information"))

16
litellm/litellm_core_utils/prompt_templates/common_utils.py
View file

@ -624,7 +624,7 @@ def update_messages_with_model_file_ids(
def update_responses_input_with_model_file_ids(
input: object,
model_id: str | None = None,
model_file_id_mapping: dict[str, dict[str, str]] | None = None,
model_file_id_mapping: Mapping[str, Mapping[str, str]] | None = None,
) -> object:
"""
Updates responses API input with provider-specific file IDs.
@ -665,7 +665,7 @@ def update_responses_input_with_model_file_ids(
for content_item in content:
if isinstance(content_item, dict) and content_item.get("type") == "input_file":
file_id = content_item.get("file_id")
if file_id:
if isinstance(file_id, str) and file_id:
provider_file_id = file_id # Default to original
# Check if we have a mapping for this file ID
@ -761,7 +761,7 @@ def _decode_vector_store_ids_in_tools(
def update_responses_tools_with_model_file_ids(
tools: list[dict[str, object]] | None,
model_id: str | None = None,
model_file_id_mapping: dict[str, dict[str, str]] | None = None,
model_file_id_mapping: Mapping[str, Mapping[str, str]] | None = None,
) -> list[dict[str, object]] | None:
"""
Updates responses API tools with provider-specific file IDs.
@ -2325,6 +2325,16 @@ def is_unsignable_thinking_block(block: object) -> bool:
return not (isinstance(signature, str) and len(signature) > 0)
def has_blank_redacted_data(block: object) -> bool:
match block:
case {"type": "redacted_thinking", "data": str() as data}:
return not data.strip()
case {"type": "redacted_thinking"}:
return True
case _:
return False
def native_redacted_thinking_data(data: str) -> str:
"""Claude's own `data` string for a redacted_thinking block.

256
litellm/litellm_core_utils/prompt_templates/factory.py
View file

@ -6,10 +6,10 @@ import json
import mimetypes
import re
import xml.etree.ElementTree as ET
from collections.abc import Container, Iterable, Iterator, Mapping, Sequence
from collections.abc import Callable, Container, Iterable, Iterator, Mapping, Sequence
from enum import Enum
from types import MappingProxyType
from typing import Annotated, Any, Final, TypeAlias, TypedDict, cast, overload
from typing import Annotated, Any, Final, TypeAlias, TypedDict, TypeVar, cast, overload
from jinja2.sandbox import ImmutableSandboxedEnvironment
from pydantic import BaseModel, BeforeValidator, ConfigDict, JsonValue, TypeAdapter, ValidationError, model_validator
@ -61,6 +61,7 @@ from litellm.types.utils import GenericImageParsingChunk
from .common_utils import (
converse_redacted_content,
convert_content_list_to_str,
has_blank_redacted_data,
infer_content_type_from_url_and_content,
is_encrypted_reasoning_block,
is_non_content_values_set,
@ -2550,22 +2551,125 @@ def _pair_tool_calls_with_results(messages: Sequence[AllMessageValues]) -> tuple
def _replayable_thinking_blocks(
thinking_blocks: list[ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock],
*,
drop_params: bool = False,
) -> list[ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock]:
keys_to_drop: Final = _UNSUPPORTED_THINKING_KEYS if drop_params else _STREAM_ONLY_THINKING_KEYS
return [
_without_stream_index(_with_native_redacted_data(block))
_without_keys(_with_native_redacted_data(block), keys_to_drop)
for block in thinking_blocks
if not is_unsignable_thinking_block(block)
if not (is_unsignable_thinking_block(block) or (drop_params and has_blank_redacted_data(block)))
]
def _without_stream_index(
def _as_thinking_block(item: object) -> ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock | None:
match item:
case {"type": "thinking" | "redacted_thinking"}:
return cast( # cast-ok: the mapping pattern matched a thinking block's type tag
ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock, item
)
case _:
return None
def _thinking_identity(item: object) -> tuple[str | None, ...] | None:
match item:
case {"type": "redacted_thinking", "data": str() as data}:
return ("redacted_thinking", native_redacted_thinking_data(data))
case {"type": "thinking", "signature": str() as signature} if signature:
return ("thinking", signature)
case {"type": str() as kind} if kind in ("thinking", "redacted_thinking"):
fields: Final = cast(Mapping[str, object], item) # cast-ok: the mapping pattern matched
return _unsigned_thinking_identity(kind, fields.get("thinking"), fields.get("signature"))
case _:
return None
def _unsigned_thinking_identity(kind: str, thinking: object, signature: object) -> tuple[str | None, ...] | None:
match thinking, signature:
case (str() | None as text, str() | None as signed):
return (kind, text, signed)
case _:
return None
_Replayed: Final = TypeVar("_Replayed")
def _with_missing_thinking_blocks(
content: Iterable[_Replayed], thinking_blocks: Sequence[_Replayed]
) -> tuple[_Replayed, ...]:
content_items: Final = tuple(content)
block_identities: Final = tuple(map(_thinking_identity, thinking_blocks))
content_positions: Final[Mapping[tuple[object, ...] | None, int]] = MappingProxyType(
{
identity: index
for index, identity in reversed(tuple(enumerate(map(_thinking_identity, content_items))))
if identity is not None
}
)
anchors: Final = tuple(index for index, identity in enumerate(block_identities) if identity in content_positions)
if not anchors:
return (*thinking_blocks, *content_items)
missing_before: Final = MappingProxyType(
{
content_positions[block_identities[end]]: thinking_blocks[start + 1 : end]
for start, end in zip((-1, *anchors), anchors)
}
)
missing_after: Final = MappingProxyType(
{content_positions[block_identities[anchors[-1]]]: thinking_blocks[anchors[-1] + 1 :]}
)
return tuple(_interleave(content_items, missing_before, missing_after))
def _without_repeated_thinking(
items: Sequence[_Replayed], identity: Callable[[_Replayed], tuple[object, ...] | None]
) -> tuple[_Replayed, ...]:
identities: Final = tuple(map(identity, items))
first_positions: Final[Mapping[tuple[object, ...], int]] = MappingProxyType(
{key: index for index, key in reversed(tuple(enumerate(identities))) if key is not None}
)
return tuple(
item for index, (item, key) in enumerate(zip(items, identities)) if key is None or first_positions[key] == index
)
def _interleave(
content: Sequence[_Replayed],
before: Mapping[int, Sequence[_Replayed]],
after: Mapping[int, Sequence[_Replayed]],
) -> Iterator[_Replayed]:
for index, item in enumerate(content):
yield from before.get(index, ())
yield item
yield from after.get(index, ())
def _thinking_in_replay_order(
content: str | Iterable[object] | None, thinking_blocks: Sequence[object] | None
) -> tuple[ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock, ...]:
content_items: Final = content if isinstance(content, list) else ()
return tuple(
block
for item in _with_missing_thinking_blocks(content_items, thinking_blocks or ())
if (block := _as_thinking_block(item)) is not None
)
_STREAM_ONLY_THINKING_KEYS: Final = frozenset({"index"})
_UNSUPPORTED_THINKING_KEYS: Final = frozenset({"index", "cache_control"})
def _without_keys(
block: ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock,
keys: frozenset[str],
) -> ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock:
if "index" not in block:
if keys.isdisjoint(block):
return block
return cast( # cast-ok: dropping one key keeps the block's TypedDict shape
return cast( # cast-ok: dropping keys keeps the block's TypedDict shape
ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock,
{key: value for key, value in block.items() if key != "index"},
{key: value for key, value in block.items() if key not in keys},
)
@ -2629,6 +2733,7 @@ def anthropic_messages_pt(
messages: list[AllMessageValues],
model: str,
llm_provider: str,
drop_params: bool = False,
) -> _AnthropicMessageList:
"""
format messages for anthropic
@ -2816,7 +2921,9 @@ def anthropic_messages_pt(
_raw_thinking_blocks = assistant_content_block.get("thinking_blocks", None)
thinking_blocks = (
_replayable_thinking_blocks(_raw_thinking_blocks) if _raw_thinking_blocks is not None else None
_replayable_thinking_blocks(_raw_thinking_blocks, drop_params=drop_params)
if _raw_thinking_blocks is not None
else None
)
# Check if tool_calls contain server tool calls (web search, etc.)
@ -2951,42 +3058,28 @@ def anthropic_messages_pt(
# SEQUENTIAL MODE: No server tool calls, or no thinking blocks,
# or content is a list. Use the original sequential approach.
# When content is a list, check if it already contains thinking
# blocks inline. If so, skip prepending thinking_blocks to avoid
# duplication and preserve the original interleaved order.
# Fixes the gap where list-content messages bypass INTERLEAVED
# MODE and still get thinking blocks prepended out of order.
_content_is_list = "content" in assistant_content_block and isinstance(
assistant_content_block["content"], list
)
_content_list = assistant_content_block.get("content") if _content_is_list else None
_list_has_thinking = False
if _content_is_list and _content_list is not None:
for _item in _content_list:
if isinstance(_item, dict) and _item.get("type") in (
"thinking",
"redacted_thinking",
):
_list_has_thinking = True
break
if (
thinking_blocks is not None and not _list_has_thinking
thinking_blocks is not None and not _content_is_list
): # IMPORTANT: ADD THIS FIRST, ELSE ANTHROPIC WILL RAISE AN ERROR
assistant_content.extend(thinking_blocks)
if _content_is_list and _content_list is not None:
for m in _content_list:
for m in _with_missing_thinking_blocks(_content_list, thinking_blocks or ()):
if not isinstance(m, dict):
continue
# handle thinking blocks
text_block = cast(str, m.get("text", ""))
if m.get("type", "") == "thinking" and not is_unsignable_thinking_block(m):
anthropic_message: ChatCompletionThinkingBlock | AnthropicMessagesTextParam = cast(
ChatCompletionThinkingBlock, m
if (thinking_block := _as_thinking_block(m)) is not None:
assistant_content.extend(
_replayable_thinking_blocks([thinking_block], drop_params=drop_params)
)
assistant_content.append(anthropic_message)
continue
text_block = cast(str, m.get("text", ""))
# handle text
elif (
if (
m.get("type", "") == "text" and len(text_block) > 0
): # don't pass empty text blocks. anthropic api raises errors.
anthropic_message = AnthropicMessagesTextParam(type="text", text=text_block)
@ -3058,7 +3151,12 @@ def anthropic_messages_pt(
msg_i += 1
if assistant_content:
new_messages.append({"role": "assistant", "content": assistant_content})
new_messages.append(
{
"role": "assistant",
"content": list(_without_repeated_thinking(assistant_content, _thinking_identity)),
}
)
## MID-CONVERSATION SYSTEM MESSAGES (placement is the caller's job) ##
while msg_i < len(messages) and messages[msg_i]["role"] == "system":
@ -4190,6 +4288,16 @@ def _convert_to_bedrock_tool_call_result(
return content_block
def _reasoning_identity(block: BedrockContentBlock) -> tuple[object, ...] | None:
match block:
case {"reasoningContent": {"redactedContent": redacted}}:
return ("redacted_thinking", redacted)
case {"reasoningContent": {"reasoningText": {"signature": str() as signature}}} if signature:
return ("thinking", signature)
case _:
return None
def _deduplicate_bedrock_content_blocks(
blocks: list[BedrockContentBlock],
block_key: str,
@ -4673,6 +4781,7 @@ class BedrockConverseMessagesProcessor:
assistant_continue_message: str | ChatCompletionAssistantMessage | None = None,
deployment_model: str | None = None,
base_model: str | None = None,
drop_params: bool = False,
) -> list[BedrockMessageBlock]:
from litellm.llms.bedrock.common_utils import bedrock_converse_accepts_redacted_thinking
@ -4826,33 +4935,21 @@ class BedrockConverseMessagesProcessor:
)
_assistant_content = assistant_message_block.get("content", None)
thinking_blocks = assistant_message_block.get("thinking_blocks")
if thinking_blocks is not None:
converted_thinking_blocks = (
BedrockConverseMessagesProcessor.translate_thinking_blocks_to_reasoning_content_blocks(
thinking_blocks, to_claude=to_claude
)
)
assistant_content = BedrockConverseMessagesProcessor.add_thinking_blocks_to_assistant_content(
thinking_blocks=converted_thinking_blocks,
assistant_parts=assistant_content,
converted_thinking_blocks = (
BedrockConverseMessagesProcessor.translate_thinking_blocks_to_reasoning_content_blocks(
_thinking_in_replay_order(_assistant_content, thinking_blocks), to_claude=to_claude
)
)
assistant_content = BedrockConverseMessagesProcessor.add_thinking_blocks_to_assistant_content(
thinking_blocks=converted_thinking_blocks,
assistant_parts=assistant_content,
)
if _assistant_content is not None and isinstance(_assistant_content, list):
assistants_parts: list[BedrockContentBlock] = []
for element in _assistant_content:
if isinstance(element, dict):
if element["type"] == "thinking" or element["type"] == "redacted_thinking":
thinking_block = BedrockConverseMessagesProcessor.translate_thinking_blocks_to_reasoning_content_blocks(
thinking_blocks=[element], to_claude=to_claude
)
assistants_parts = (
BedrockConverseMessagesProcessor.add_thinking_blocks_to_assistant_content(
thinking_blocks=thinking_block,
assistant_parts=assistants_parts,
)
)
elif element["type"] == "text":
if element["type"] == "text":
# Skip completely empty strings to avoid blank content blocks
if element.get("text", "").strip():
assistants_part = BedrockContentBlock(text=element["text"])
@ -4871,7 +4968,9 @@ class BedrockConverseMessagesProcessor:
model=model,
deployment_model=deployment_model,
)
if _cache_point_block is not None:
if _cache_point_block is not None and not (
drop_params and _as_thinking_block(element) is not None
):
assistants_parts.append(_cache_point_block)
assistant_content.extend(assistants_parts)
elif _assistant_content is not None and isinstance(_assistant_content, str):
@ -4899,7 +4998,11 @@ class BedrockConverseMessagesProcessor:
msg_i += 1
assistant_content = _deduplicate_bedrock_content_blocks(assistant_content, "toolUse")
assistant_content = list(
_without_repeated_thinking(
_deduplicate_bedrock_content_blocks(assistant_content, "toolUse"), _reasoning_identity
)
)
assistant_content = _sort_bedrock_assistant_content_blocks(assistant_content)
if assistant_content:
@ -5083,6 +5186,7 @@ def _bedrock_converse_messages_pt(
assistant_continue_message: str | ChatCompletionAssistantMessage | None = None,
deployment_model: str | None = None,
base_model: str | None = None,
drop_params: bool = False,
) -> list[BedrockMessageBlock]:
"""
Converts given messages from OpenAI format to Bedrock format
@ -5244,35 +5348,21 @@ def _bedrock_converse_messages_pt(
)
_assistant_content = assistant_message_block.get("content", None)
thinking_blocks = assistant_message_block.get("thinking_blocks")
if thinking_blocks is not None:
converted_thinking_blocks = (
BedrockConverseMessagesProcessor.translate_thinking_blocks_to_reasoning_content_blocks(
thinking_blocks, to_claude=to_claude
)
)
assistant_content = BedrockConverseMessagesProcessor.add_thinking_blocks_to_assistant_content(
thinking_blocks=converted_thinking_blocks,
assistant_parts=assistant_content,
converted_thinking_blocks = (
BedrockConverseMessagesProcessor.translate_thinking_blocks_to_reasoning_content_blocks(
_thinking_in_replay_order(_assistant_content, thinking_blocks), to_claude=to_claude
)
)
assistant_content = BedrockConverseMessagesProcessor.add_thinking_blocks_to_assistant_content(
thinking_blocks=converted_thinking_blocks,
assistant_parts=assistant_content,
)
if _assistant_content is not None and isinstance(_assistant_content, list):
assistants_parts: list[BedrockContentBlock] = []
for element in _assistant_content:
if isinstance(element, dict):
if element["type"] == "thinking" or element["type"] == "redacted_thinking":
thinking_block = (
BedrockConverseMessagesProcessor.translate_thinking_blocks_to_reasoning_content_blocks(
thinking_blocks=[element], to_claude=to_claude
)
)
assistants_parts = (
BedrockConverseMessagesProcessor.add_thinking_blocks_to_assistant_content(
thinking_blocks=thinking_block,
assistant_parts=assistants_parts,
)
)
elif element["type"] == "text":
if element["type"] == "text":
# AWS Bedrock doesn't allow empty or whitespace-only text content
# Skip completely empty strings to avoid blank content blocks
if element.get("text", "").strip():
@ -5292,7 +5382,9 @@ def _bedrock_converse_messages_pt(
model=model,
deployment_model=deployment_model,
)
if _cache_point_block is not None:
if _cache_point_block is not None and not (
drop_params and _as_thinking_block(element) is not None
):
assistants_parts.append(_cache_point_block)
assistant_content.extend(assistants_parts)
elif _assistant_content is not None and isinstance(_assistant_content, str):
@ -5313,7 +5405,11 @@ def _bedrock_converse_messages_pt(
msg_i += 1
assistant_content = _deduplicate_bedrock_content_blocks(assistant_content, "toolUse")
assistant_content = list(
_without_repeated_thinking(
_deduplicate_bedrock_content_blocks(assistant_content, "toolUse"), _reasoning_identity
)
)
assistant_content = _sort_bedrock_assistant_content_blocks(assistant_content)
if assistant_content:

19
litellm/litellm_core_utils/prompt_templates/mid_conversation_system.py
View file

@ -49,7 +49,7 @@ from litellm.types.llms.openai import (
ChatCompletionUserMessage,
)
from .common_utils import is_unsignable_thinking_block
from .common_utils import has_blank_redacted_data, is_unsignable_thinking_block
CONVERTED_SYSTEM_NOTE: Final = (
"Operator note (not from the user): the following was originally a mid-conversation system-role reminder."
@ -267,27 +267,22 @@ def _block_index_by_message(blocks: Sequence[tuple[bool, tuple[int, ...]]]) -> M
def _thinking_block_renders(block: object) -> bool:
"""A thinking block the converter keeps: one Anthropic can verify, so never bridged encrypted reasoning."""
return message_field(block, "type") in _THINKING_BLOCK_TYPES and not is_unsignable_thinking_block(block)
return message_field(block, "type") in _THINKING_BLOCK_TYPES and not (
is_unsignable_thinking_block(block) or has_blank_redacted_data(block)
)
def _assistant_part_renders(part: object) -> bool:
"""A text part always renders: the converter pads empty text with a placeholder."""
part_type: Final = message_field(part, "type")
if part_type == "thinking":
if part_type in _THINKING_BLOCK_TYPES:
return _thinking_block_renders(part)
return part_type in _RENDERED_ASSISTANT_PART_TYPES or (
isinstance(part_type, str) and part_type.endswith("_tool_result")
)
def _separate_thinking_blocks_render(message: object, parts: Sequence[object]) -> bool:
"""``thinking_blocks`` reach the wire only when no inline thinking part claims the slot.
The converter skips the separate blocks as soon as the content list carries a
``thinking`` or ``redacted_thinking`` part, whether or not that part itself renders.
"""
if any(message_field(part, "type") in _THINKING_BLOCK_TYPES for part in parts):
return False
def _separate_thinking_blocks_render(message: object) -> bool:
return any(_thinking_block_renders(block) for block in parts_of(message_field(message, "thinking_blocks")))
@ -404,7 +399,7 @@ def _assistant_renders(message: object, answered: frozenset[str]) -> bool:
parts: Final = parts_of(content)
return (
any(_assistant_part_renders(part) for part in parts)
or _separate_thinking_blocks_render(message, parts)
or _separate_thinking_blocks_render(message)
or _tool_calls_render(message, answered)
or bool(message_field(message, "function_call"))
or bool(message_field(message_field(message, "provider_specific_fields"), "compaction_blocks"))

5
litellm/llms/anthropic/chat/transformation.py
View file

@ -2267,6 +2267,8 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
anthropic_messages_pt,
)
drop_params: Final = litellm.drop_params is True or litellm_params.get("drop_params") is True
if "tools" not in optional_params and messages is not None and has_tool_call_blocks(messages):
optional_params["tools"], _ = self._map_tools(add_dummy_tool(custom_llm_provider="anthropic"))
@ -2293,7 +2295,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
AnthropicConfig._maybe_drop_speed_param(
model=model,
optional_params=optional_params,
drop_params=litellm.drop_params or litellm_params.get("drop_params") is True,
drop_params=drop_params,
custom_llm_provider=self.custom_llm_provider,
)
@ -2351,6 +2353,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
model=model,
messages=list(conversation),
llm_provider=self._resolved_provider,
drop_params=drop_params,
)
except Exception as e:
raise AnthropicError(

34
litellm/llms/anthropic/count_tokens/token_counter.py
View file

@ -2,6 +2,8 @@
Anthropic Token Counter implementation using the CountTokens API.
"""
from collections.abc import Mapping
from types import MappingProxyType
from typing import Any, Final
from pydantic import JsonValue
@ -10,6 +12,11 @@ import litellm
from litellm._logging import verbose_logger
from litellm.exceptions import AuthenticationError
from litellm.llms.anthropic.count_tokens.handler import AnthropicCountTokensHandler
from litellm.llms.anthropic.count_tokens.transformation import (
AnthropicCountFields,
ChatCountTranslationError,
anthropic_count_fields,
)
from litellm.llms.base_llm.base_utils import BaseTokenCounter
from litellm.types.utils import CountContextManagement, LlmProviders, TokenCountMessagesFormat, TokenCountResponse
@ -38,6 +45,8 @@ class AnthropicTokenCounter(BaseTokenCounter):
messages_format: TokenCountMessagesFormat = "anthropic",
output_format: JsonValue = None,
context_management: CountContextManagement | None = None,
chat_params: Mapping[str, JsonValue] = MappingProxyType({}),
count_options: Mapping[str, JsonValue] = MappingProxyType({}),
) -> TokenCountResponse | None:
"""
Count tokens using Anthropic's CountTokens API.
@ -57,6 +66,30 @@ class AnthropicTokenCounter(BaseTokenCounter):
if not messages:
return None
from litellm.llms.anthropic.chat.transformation import AnthropicConfig
match translated := await anthropic_count_fields(
model_to_use, messages, system, tools, messages_format, AnthropicConfig(), chat_params
):
case ChatCountTranslationError():
return translated.response(request_model, model_to_use, "anthropic_api")
case AnthropicCountFields():
return await self.count_tokens(
model_to_use=model_to_use,
messages=[dict(message) for message in translated.messages],
contents=contents,
deployment=deployment,
request_model=request_model,
tools=[dict(tool) for tool in translated.tools] if translated.tools is not None else None,
system=translated.system,
messages_format="anthropic",
output_format=output_format,
context_management=context_management,
count_options=translated.options,
)
case None:
pass
deployment = deployment or {}
litellm_params: Final = deployment.get("litellm_params", {})
api_base: Final = litellm_params.get("api_base")
@ -89,6 +122,7 @@ class AnthropicTokenCounter(BaseTokenCounter):
api_base=api_base,
tools=tools,
system=system,
optional_params=count_options,
)
if result is not None:

225
litellm/llms/anthropic/count_tokens/transformation.py
View file

@ -4,25 +4,37 @@ Anthropic CountTokens API transformation logic.
This module handles the transformation of requests to Anthropic's CountTokens API format.
"""
import copy
from collections.abc import Mapping, Sequence
from dataclasses import dataclass
from itertools import chain
from types import MappingProxyType
from typing import Final
from typing import (
Final,
Protocol,
TypeAlias,
cast, # noqa: TID251 # caller chat messages go to chat inference translation as-is
)
from pydantic import JsonValue, TypeAdapter
from pydantic import ConfigDict, JsonValue, TypeAdapter, with_config
from typing_extensions import NotRequired, ReadOnly, TypedDict
from litellm.constants import ANTHROPIC_TOKEN_COUNTING_BETA_VERSION
from litellm.litellm_core_utils.asyncify import asyncify
from litellm.litellm_core_utils.prompt_templates.mid_conversation_system import (
anthropic_system_blocks,
split_leading_system_run,
)
from litellm.llms.anthropic.common_utils import merge_anthropic_beta_headers
from litellm.llms.anthropic.wif import resolve_anthropic_base
from litellm.types.llms.openai import ChatCompletionImageObject
from litellm.types.llms.openai import AllMessageValues, ChatCompletionImageObject, ValidUserMessageContentTypes
from litellm.types.utils import TokenCountMessagesFormat, TokenCountResponse
_COUNT_REQUEST: Final = TypeAdapter(dict[str, JsonValue])
_SYSTEM_BLOCKS: Final = TypeAdapter(list[JsonValue])
_IMAGE_BLOCK: Final = TypeAdapter(ChatCompletionImageObject)
COUNT_TOKEN_OPTION_NAMES: Final = ("thinking", "tool_choice", "output_config")
_PLAIN_MESSAGE_KEYS: Final = frozenset({"role", "content"})
def _count_image(block: JsonValue) -> JsonValue:
@ -74,6 +86,213 @@ def _lift_leading_system(
return conversation, lifted or system
@dataclass(frozen=True, slots=True)
class AnthropicCountFields:
messages: tuple[Mapping[str, JsonValue], ...]
system: JsonValue
tools: tuple[Mapping[str, JsonValue], ...] | None
options: Mapping[str, JsonValue]
@dataclass(frozen=True, slots=True)
class ChatCountTranslationError:
message: str
def response(self, request_model: str, model_used: str, tokenizer_type: str) -> TokenCountResponse:
return TokenCountResponse(
total_tokens=0,
request_model=request_model,
model_used=model_used,
tokenizer_type=tokenizer_type,
error=True,
error_message=self.message,
status_code=400,
rejected_before_provider_call=True,
)
@with_config(ConfigDict(extra="ignore"))
class _AnthropicChatRequest(TypedDict):
messages: ReadOnly[tuple[Mapping[str, JsonValue], ...]]
system: NotRequired[ReadOnly[JsonValue]]
tools: NotRequired[ReadOnly[tuple[Mapping[str, JsonValue], ...]]]
thinking: NotRequired[ReadOnly[JsonValue]]
tool_choice: NotRequired[ReadOnly[JsonValue]]
output_config: NotRequired[ReadOnly[JsonValue]]
output_format: NotRequired[ReadOnly[JsonValue]]
_ANTHROPIC_CHAT_REQUEST: Final = TypeAdapter(_AnthropicChatRequest)
_ChatParams: TypeAlias = dict[str, object]
_ChatMessages: TypeAlias = list[AllMessageValues]
class _ChatRequestTranslator(Protocol):
def map_openai_params(
self, non_default_params: _ChatParams, optional_params: _ChatParams, model: str, drop_params: bool
) -> _ChatParams: ...
def transform_request(
self,
model: str,
messages: _ChatMessages,
optional_params: _ChatParams,
litellm_params: _ChatParams,
headers: _ChatParams,
) -> _ChatParams: ...
def _needs_chat_translation(message: Mapping[str, object]) -> bool:
return message.get("role") == "tool" or not _PLAIN_MESSAGE_KEYS.issuperset(message.keys())
def _is_chat_tool(tool: Mapping[str, object]) -> bool:
from litellm.litellm_core_utils.prompt_templates.factory import is_flat_responses_tool
return is_flat_responses_tool(tool) or (tool.get("type") == "function" and "function" in tool)
def _has_non_chat_user_block(message: Mapping[str, object]) -> bool:
content: Final = message.get("content")
if message.get("role") != "user" or not isinstance(content, list):
return False
blocks: Final = _SYSTEM_BLOCKS.validate_python(content)
return any(isinstance(block, dict) and block.get("type") not in ValidUserMessageContentTypes for block in blocks)
def _gets_placeholder_name(tool: Mapping[str, object]) -> bool:
from litellm.llms.anthropic.pass_through.adapters.transformation import is_provider_native_tool_dict
name: Final = tool.get("name")
return not (_is_chat_tool(tool) or is_provider_native_tool_dict(tool) or (isinstance(name, str) and name.strip()))
def _untranslatable_input(
messages: Sequence[Mapping[str, object]], tools: Sequence[Mapping[str, object]]
) -> ChatCountTranslationError | None:
invalid_index: Final = next(
(index for index, message in enumerate(messages) if _has_non_chat_user_block(message)), None
)
if invalid_index is not None:
return ChatCountTranslationError(
message=f"Invalid user message at index {invalid_index}. "
"Please ensure all user messages are valid OpenAI chat completion messages."
)
unnamed_tool: Final = next(filter(_gets_placeholder_name, tools), None)
if unnamed_tool is not None:
return ChatCountTranslationError(
message=f"Invalid request for token counting: tool type {unnamed_tool.get('type')!r} "
"has no Chat Completions equivalent"
)
return None
def _system_messages(system: object) -> tuple[Mapping[str, object], ...]:
message: Final[Mapping[str, object]] = {"role": "system", "content": system}
return (message,) if isinstance(system, (str, list)) and system else ()
def _count_options(request: _AnthropicChatRequest) -> Mapping[str, JsonValue]:
"""Structured output reaches chat inference as the beta ``output_format``, which count_tokens
refuses; its GA ``output_config.format`` form counts the same schema the same way."""
output_config: Final = request.get("output_config")
output_format: Final = request.get("output_format")
counted_config: Final = {
**(output_config if isinstance(output_config, dict) else {}),
**({} if output_format is None else {"format": output_format}),
}
return MappingProxyType(
{
key: value
for key, value in (
("thinking", request.get("thinking")),
("tool_choice", request.get("tool_choice")),
("output_config", counted_config or None),
)
if value is not None
}
)
def _translated_chat_request(
model: str,
messages: Sequence[Mapping[str, object]],
system: object,
tools: Sequence[Mapping[str, object]],
chat_params: Mapping[str, object],
translator: _ChatRequestTranslator,
) -> _ChatParams:
from litellm.litellm_core_utils.prompt_templates.factory import chat_completion_tools_for
chat_tools: Final = [dict(chat_tool) for chat_tool in chain.from_iterable(map(chat_completion_tools_for, tools))]
openai_params: Final = {name: value for name, value in chat_params.items() if name != "drop_params"}
mapped_params: Final = {**openai_params, **({"tools": chat_tools} if chat_tools else {})}
chat_messages: Final = cast( # cast-ok: chat-shaped request messages, read the way chat inference reads them
_ChatMessages, copy.deepcopy([*_system_messages(system), *messages])
)
return translator.transform_request(
model=model,
messages=chat_messages,
optional_params=translator.map_openai_params(
mapped_params, {}, model, drop_params=chat_params.get("drop_params") is True
)
if mapped_params
else {},
litellm_params={},
headers={},
)
def anthropic_fields_for_chat_tool_use(
model: str,
messages: Sequence[Mapping[str, object]],
system: object,
tools: Sequence[Mapping[str, object]] | None,
translator: _ChatRequestTranslator,
chat_params: Mapping[str, object] = MappingProxyType({}),
) -> AnthropicCountFields | ChatCountTranslationError | None:
counted_tools: Final = tools or ()
if (
not chat_params
and not any(map(_is_chat_tool, counted_tools))
and not any(map(_needs_chat_translation, messages))
):
return None
untranslatable: Final = _untranslatable_input(messages, counted_tools)
if untranslatable is not None:
return untranslatable
from litellm.exceptions import BadRequestError
from litellm.llms.anthropic.common_utils import AnthropicError
try:
request: Final = _translated_chat_request(model, messages, system, counted_tools, chat_params, translator)
except (AnthropicError, BadRequestError) as e:
return ChatCountTranslationError(message=str(e.message))
except (ValueError, KeyError) as e:
return ChatCountTranslationError(message=f"Invalid request for token counting: {e}")
translated: Final = _ANTHROPIC_CHAT_REQUEST.validate_python(request)
return AnthropicCountFields(
messages=translated["messages"],
system=translated.get("system"),
tools=translated.get("tools"),
options=_count_options(translated),
)
async def anthropic_count_fields(
model: str,
messages: Sequence[Mapping[str, object]],
system: object,
tools: Sequence[Mapping[str, object]] | None,
messages_format: TokenCountMessagesFormat,
translator: _ChatRequestTranslator,
chat_params: Mapping[str, object] = MappingProxyType({}),
) -> AnthropicCountFields | ChatCountTranslationError | None:
if messages_format != "openai":
return None
return await asyncify(anthropic_fields_for_chat_tool_use)(model, messages, system, tools, translator, chat_params)
class AnthropicCountTokensConfig:
"""
Configuration and transformation logic for Anthropic CountTokens API.

3
litellm/llms/anthropic/pass_through/messages/response_cache.py
View file

@ -1,5 +1,6 @@
import re
from collections.abc import AsyncIterator, Mapping, Sequence
from datetime import datetime
from types import MappingProxyType
from typing import TYPE_CHECKING, Final, cast
@ -153,6 +154,8 @@ class CachedAnthropicMessagesStreamIterator(BaseAnthropicMessagesStreamingIterat
await self._handle_streaming_logging(chunks)
raise StopAsyncIteration
chunk: Final = self.chunks[self.current_index]
if self.completion_start_time is None:
self.completion_start_time = datetime.now()
self.current_index += 1
return chunk

2
litellm/llms/anthropic/pass_through/messages/streaming_iterator.py
View file

@ -506,7 +506,7 @@ class BaseAnthropicMessagesStreamingIterator:
):
self.litellm_logging_obj = litellm_logging_obj
self.request_body = request_body
self.start_time = datetime.now()
self.start_time = litellm_logging_obj.start_time
self.completion_start_time: datetime | None = None
async def _handle_streaming_logging(self, collected_chunks: list[bytes], *, stream_teardown: bool = False):

3
litellm/llms/azure_ai/anthropic/count_tokens/handler.py
View file

@ -4,6 +4,7 @@ Azure AI Anthropic CountTokens API handler.
Uses httpx for HTTP requests with Azure authentication.
"""
from collections.abc import Mapping
from typing import Final
import httpx
@ -36,6 +37,7 @@ class AzureAIAnthropicCountTokensHandler(AzureAIAnthropicCountTokensConfig):
timeout: float | httpx.Timeout | None = None,
tools: list[dict[str, JsonValue]] | None = None,
system: object = None,
optional_params: Mapping[str, JsonValue] | None = None,
) -> dict[str, object]:
"""
Handle a CountTokens request using httpx with Azure authentication.
@ -66,6 +68,7 @@ class AzureAIAnthropicCountTokensHandler(AzureAIAnthropicCountTokensConfig):
messages=messages,
tools=tools,
system=system,
optional_params=optional_params,
)
verbose_logger.debug("Transformed request: %s", request_body)

34
litellm/llms/azure_ai/anthropic/count_tokens/token_counter.py
View file

@ -3,12 +3,19 @@ Azure AI Anthropic Token Counter implementation using the CountTokens API.
"""
import os
from collections.abc import Mapping
from types import MappingProxyType
from typing import Any, Final
from pydantic import JsonValue
import litellm
from litellm._logging import verbose_logger
from litellm.llms.anthropic.count_tokens.transformation import (
AnthropicCountFields,
ChatCountTranslationError,
anthropic_count_fields,
)
from litellm.llms.azure_ai.anthropic.count_tokens.handler import (
AzureAIAnthropicCountTokensHandler,
)
@ -55,6 +62,8 @@ class AzureAIAnthropicTokenCounter(BaseTokenCounter):
messages_format: TokenCountMessagesFormat = "anthropic",
output_format: JsonValue = None,
context_management: CountContextManagement | None = None,
chat_params: Mapping[str, JsonValue] = MappingProxyType({}),
count_options: Mapping[str, JsonValue] = MappingProxyType({}),
) -> TokenCountResponse | None:
"""
Count tokens using Azure AI Anthropic's CountTokens API.
@ -74,6 +83,30 @@ class AzureAIAnthropicTokenCounter(BaseTokenCounter):
if not messages:
return None
from litellm.llms.azure_ai.anthropic.transformation import AzureAnthropicConfig
match translated := await anthropic_count_fields(
model_to_use, messages, system, tools, messages_format, AzureAnthropicConfig(), chat_params
):
case ChatCountTranslationError():
return translated.response(request_model, model_to_use, "azure_ai_anthropic_api")
case AnthropicCountFields():
return await self.count_tokens(
model_to_use=model_to_use,
messages=[dict(message) for message in translated.messages],
contents=contents,
deployment=deployment,
request_model=request_model,
tools=[dict(tool) for tool in translated.tools] if translated.tools is not None else None,
system=translated.system,
messages_format="anthropic",
output_format=output_format,
context_management=context_management,
count_options=translated.options,
)
case None:
pass
deployment = deployment or {}
litellm_params: Final = deployment.get("litellm_params", {})
@ -101,6 +134,7 @@ class AzureAIAnthropicTokenCounter(BaseTokenCounter):
litellm_params=litellm_params,
tools=tools,
system=system,
optional_params=count_options,
)
if result is not None:

2
litellm/llms/base_llm/base_utils.py
View file

@ -9,6 +9,7 @@ import re
from abc import ABC, abstractmethod
from collections.abc import Iterator, Mapping, Sequence
from dataclasses import dataclass
from types import MappingProxyType
from typing import Any, Final
from openai.lib import _parsing, _pydantic # pyright: ignore[reportPrivateUsage] # SDK parser internals
@ -40,6 +41,7 @@ class BaseTokenCounter(ABC):
messages_format: TokenCountMessagesFormat = "openai",
output_format: JsonValue = None,
context_management: CountContextManagement | None = None,
chat_params: Mapping[str, JsonValue] = MappingProxyType({}),
) -> TokenCountResponse | None:
pass

2
litellm/llms/base_llm/chat/transformation.py
View file

@ -428,7 +428,7 @@ class BaseConfig(ABC):
streaming_response: Union[Iterator[str], AsyncIterator[str], "ModelResponse"],
sync_stream: bool,
json_mode: bool | None = False,
) -> Any:
) -> object:
pass
def get_model_response_iterator_for_request(

2
litellm/llms/base_llm/files/batch_records.py
View file

@ -18,6 +18,7 @@ def _responses_request_keys() -> frozenset[str]:
def responses_batch_body_to_chat_body(
openai_request_body: Mapping[str, object],
custom_llm_provider: str | None = None,
drop_params: bool | None = None,
) -> dict[str, object]: # mutable-ok: provider transforms take the bridged chat body as a plain dict
"""
Rewrite the body of an OpenAI `/v1/responses` batch record as a Chat Completions body.
@ -50,5 +51,6 @@ def responses_batch_body_to_chat_body(
input=chat_input,
responses_api_request=responses_request,
custom_llm_provider=custom_llm_provider,
drop_params=drop_params,
metadata=openai_request_body.get("metadata"),
)

37
litellm/llms/bedrock/base_aws_llm.py
View file

@ -80,6 +80,12 @@ class _AssumeRoleParams(TypedDict):
Tags: ReadOnly[NotRequired[tuple[AwsSessionTag, ...]]]
class _StsClientKwargs(TypedDict):
verify: ReadOnly[bool | str]
endpoint_url: ReadOnly[NotRequired[str]]
region_name: ReadOnly[NotRequired[str]]
def _assume_role_params(
aws_role_name: str,
aws_session_name: str,
@ -275,7 +281,7 @@ class BaseAWSLLM(SignsRequestsWithAWS):
"aws_session_tags",
]
def _get_ssl_verify(self, ssl_verify: bool | str | None = None):
def _get_ssl_verify(self, ssl_verify: bool | str | None = None) -> bool | str:
"""
Get SSL verification setting for boto3 clients.
@ -875,15 +881,18 @@ class BaseAWSLLM(SignsRequestsWithAWS):
aws_sts_endpoint: str | None = None,
ssl_verify: bool | str | None = None,
aws_region_name: str | None = None,
) -> dict:
) -> _StsClientKwargs:
"""STS client kwargs with aligned endpoint_url and region_name (SigV4)."""
kwargs: Final[dict] = {"verify": self._get_ssl_verify(ssl_verify)}
if aws_sts_endpoint is not None:
kwargs["endpoint_url"] = aws_sts_endpoint
sts_region: Final = self._resolve_sts_region(aws_sts_endpoint, aws_region_name)
if sts_region is not None:
kwargs["region_name"] = sts_region
return kwargs
verify: Final = self._get_ssl_verify(ssl_verify)
match (aws_sts_endpoint, self._resolve_sts_region(aws_sts_endpoint, aws_region_name)):
case (None, None):
return _StsClientKwargs(verify=verify)
case (None, region):
return _StsClientKwargs(verify=verify, region_name=region)
case (endpoint, None):
return _StsClientKwargs(verify=verify, endpoint_url=endpoint)
case (endpoint, region):
return _StsClientKwargs(verify=verify, endpoint_url=endpoint, region_name=region)
def get_aws_region_name_for_non_llm_api_calls(
self,
@ -1813,13 +1822,13 @@ async def run_aws_signing(
)
async def sign_request_off_loop_if_aws(
async def run_off_loop_if_aws(
provider_config: object,
sign_request: Callable[_SignParams, _SignedRequest],
call: Callable[_SignParams, _SignedRequest],
/,
*args: _SignParams.args,
**kwargs: _SignParams.kwargs, # kwargs-ok: ParamSpec forwarding keeps the wrapped sign_request signature
**kwargs: _SignParams.kwargs, # kwargs-ok: ParamSpec forwarding keeps the wrapped call's signature
) -> _SignedRequest:
if isinstance(provider_config, SignsRequestsWithAWS):
return await run_aws_signing(sign_request, *args, **kwargs)
return sign_request(*args, **kwargs)
return await run_aws_signing(call, *args, **kwargs)
return call(*args, **kwargs)

4
litellm/llms/bedrock/chat/converse_transformation.py
View file

@ -385,7 +385,7 @@ class AmazonConverseConfig(BaseConfig):
@staticmethod
def _requires_min_max_tokens(model: str) -> bool:
return re.search(r"openai\.gpt-\d|xai\.grok-", model) is not None
return re.search(r"openai\.gpt-\d|xai\.grok-|moonshotai\.kimi-k3", model) is not None
def _is_nova_2_model(self, model: str) -> bool:
"""
@ -2127,6 +2127,7 @@ class AmazonConverseConfig(BaseConfig):
llm_provider="bedrock_converse",
user_continue_message=litellm_params.pop("user_continue_message", None),
base_model=self._deployment_base_model(_DEPLOYMENT_PARAMS.validate_python(litellm_params)),
drop_params=litellm.drop_params is True or litellm_params.get("drop_params") is True,
),
capability_model,
)
@ -2207,6 +2208,7 @@ class AmazonConverseConfig(BaseConfig):
llm_provider="bedrock_converse",
user_continue_message=litellm_params.pop("user_continue_message", None),
base_model=self._deployment_base_model(_DEPLOYMENT_PARAMS.validate_python(litellm_params)),
drop_params=litellm.drop_params is True or litellm_params.get("drop_params") is True,
),
capability_model,
)

34
litellm/llms/bedrock/claude_platform/messages_transformation.py
View file

@ -1,10 +1,11 @@
from typing import Any, Final
from typing import Final
import litellm
from litellm.llms.anthropic.pass_through.messages.transformation import (
DEFAULT_ANTHROPIC_API_VERSION,
AnthropicMessagesConfig,
)
from litellm.llms.bedrock.base_aws_llm import run_aws_signing
from litellm.secret_managers.main import get_secret_str
from litellm.types.router import GenericLiteLLMParams
@ -22,14 +23,14 @@ class BedrockClaudePlatformMessagesConfig(BedrockClaudePlatformMixin, AnthropicM
def validate_anthropic_messages_environment(
self,
headers: dict,
headers: dict[str, str], # mutable-ok: overrides the base signature
model: str,
messages: list[Any],
optional_params: dict,
litellm_params: dict,
messages: list[object], # mutable-ok: overrides the base signature
optional_params: dict[str, object], # mutable-ok: overrides the base signature
litellm_params: dict[str, object], # mutable-ok: overrides the base signature
api_key: str | None = None,
api_base: str | None = None,
) -> tuple[dict, str | None]:
) -> tuple[dict[str, str], str | None]: # mutable-ok: overrides the base signature
workspace_id: Final = self._get_workspace_id(optional_params, litellm_params)
if workspace_id is None:
raise litellm.AuthenticationError(
@ -66,6 +67,27 @@ class BedrockClaudePlatformMessagesConfig(BedrockClaudePlatformMixin, AnthropicM
api_base,
)
async def avalidate_anthropic_messages_environment(
self,
headers: dict[str, str], # mutable-ok: mirrors validate_anthropic_messages_environment
model: str,
messages: list[object], # mutable-ok: mirrors validate_anthropic_messages_environment
optional_params: dict[str, object], # mutable-ok: mirrors validate_anthropic_messages_environment
litellm_params: dict[str, object], # mutable-ok: mirrors validate_anthropic_messages_environment
api_key: str | None = None,
api_base: str | None = None,
) -> tuple[dict[str, str], str | None]: # mutable-ok: mirrors validate_anthropic_messages_environment
return await run_aws_signing(
self.validate_anthropic_messages_environment,
headers=headers,
model=model,
messages=messages,
optional_params=optional_params,
litellm_params=litellm_params,
api_key=api_key,
api_base=api_base,
)
def transform_anthropic_messages_request(
self,
model: str,

62
litellm/llms/bedrock/count_tokens/bedrock_token_counter.py
View file

@ -4,12 +4,15 @@ Bedrock Token Counter implementation using the CountTokens API.
from collections.abc import Mapping, Sequence
from dataclasses import dataclass
from types import MappingProxyType
from typing import Any, Final
from pydantic import JsonValue
from typing_extensions import assert_never
from litellm._logging import verbose_logger
from litellm.llms.anthropic.chat.transformation import AnthropicConfig
from litellm.llms.anthropic.count_tokens.transformation import ChatCountTranslationError, anthropic_count_fields
from litellm.llms.base_llm.base_utils import BaseTokenCounter
from litellm.llms.bedrock.common_utils import BedrockError, get_bedrock_base_model
from litellm.llms.bedrock.count_tokens.handler import BedrockCountTokensHandler
@ -58,6 +61,17 @@ def _runtime_rejected_claude_model(outcome: CountTokensOutcome, resolved_model:
)
def _request_data(
model: str, messages: Sequence[Mapping[str, object]], tools: Sequence[Mapping[str, object]] | None, system: object
) -> Mapping[str, object]:
return {
"model": model,
"messages": [dict(message) for message in messages],
**({"tools": [dict(tool) for tool in tools]} if tools else {}),
**({"system": system} if system else {}),
}
class BedrockTokenCounter(BaseTokenCounter):
"""Counts on the configured Bedrock endpoint, with Mantle retry when Runtime rejects a Claude model"""
@ -81,13 +95,13 @@ class BedrockTokenCounter(BaseTokenCounter):
self,
handler: BedrockCountTokensHandler | BedrockMantleCountTokensHandler,
tokenizer_type: str,
request_data: dict[str, object],
request_data: Mapping[str, object],
litellm_params: dict[str, object],
resolved_model: str,
) -> CountTokensOutcome:
try:
result: Final = await handler.handle_count_tokens_request(
request_data=request_data,
request_data=dict(request_data),
litellm_params=litellm_params,
resolved_model=resolved_model,
client=self._client,
@ -115,30 +129,44 @@ class BedrockTokenCounter(BaseTokenCounter):
messages_format: TokenCountMessagesFormat = "anthropic",
output_format: JsonValue = None,
context_management: CountContextManagement | None = None,
chat_params: Mapping[str, JsonValue] = MappingProxyType({}),
) -> TokenCountResponse | None:
if not messages:
return None
litellm_params: Final = (deployment or {}).get("litellm_params", {})
request_data: Final[dict[str, object]] = {
"model": model_to_use,
"messages": messages,
**({"tools": tools} if tools else {}),
**({"system": system} if system else {}),
}
use_mantle: Final = model_to_use.startswith("bedrock_mantle/")
resolved_model: Final = get_bedrock_base_model(model_to_use.removeprefix("bedrock_mantle/"))
counted_model: Final = model_to_use.removeprefix("bedrock_mantle/")
resolved_model: Final = get_bedrock_base_model(counted_model)
primary_tokenizer: Final = MANTLE_TOKENIZER_TYPE if use_mantle else RUNTIME_TOKENIZER_TYPE
translated: Final = (
await anthropic_count_fields(
counted_model, messages, system, tools, messages_format, AnthropicConfig(), chat_params
)
if is_mantle_claude_model(resolved_model)
else None
)
if isinstance(translated, ChatCountTranslationError):
return translated.response(request_model, model_to_use, primary_tokenizer)
litellm_params: Final = (deployment or {}).get("litellm_params", {})
runtime_request: Final = (
_request_data(model_to_use, messages, tools, system)
if translated is None
else _request_data(model_to_use, translated.messages, translated.tools, translated.system)
)
mantle_request: Final = runtime_request if translated is None else {**runtime_request, **translated.options}
primary: Final = await self._count_with(
self._mantle_handler if use_mantle else self._runtime_handler,
MANTLE_TOKENIZER_TYPE if use_mantle else RUNTIME_TOKENIZER_TYPE,
request_data,
litellm_params,
resolved_model,
primary: Final = (
await self._count_with(
self._mantle_handler, primary_tokenizer, mantle_request, litellm_params, resolved_model
)
if use_mantle
else await self._count_with(
self._runtime_handler, primary_tokenizer, runtime_request, litellm_params, resolved_model
)
)
outcome: Final = (
await self._count_with(
self._mantle_handler, MANTLE_TOKENIZER_TYPE, request_data, litellm_params, resolved_model
self._mantle_handler, MANTLE_TOKENIZER_TYPE, mantle_request, litellm_params, resolved_model
)
if not use_mantle and _runtime_rejected_claude_model(primary, resolved_model)
else primary

6
litellm/llms/bedrock/count_tokens/handler.py
View file

@ -48,7 +48,8 @@ class BedrockCountTokensHandler(BedrockCountTokensConfig):
verbose_logger.debug("Processing CountTokens request for resolved model: %s", resolved_model)
# Get AWS region using existing LiteLLM function
aws_region_name: Final = self.get_aws_region_name(
aws_region_name: Final = await run_aws_signing(
self.get_aws_region_name,
optional_params=litellm_params,
model=resolved_model,
model_id=None,
@ -64,7 +65,8 @@ class BedrockCountTokensHandler(BedrockCountTokensConfig):
# Get endpoint URL using simplified function
api_base: Final = litellm_params.get("api_base", None)
aws_bedrock_runtime_endpoint: Final = litellm_params.get("aws_bedrock_runtime_endpoint", None)
endpoint_url: Final = self.get_bedrock_count_tokens_endpoint(
endpoint_url: Final = await run_aws_signing(
self.get_bedrock_count_tokens_endpoint,
model=resolved_model,
aws_region_name=aws_region_name,
api_base=api_base,

13
litellm/llms/bedrock/count_tokens/mantle_handler.py
View file

@ -1,3 +1,4 @@
from collections.abc import Mapping
from typing import Final
import httpx
@ -6,7 +7,7 @@ from typing_extensions import NotRequired, ReadOnly, TypedDict
import litellm
from litellm._logging import verbose_logger
from litellm.llms.anthropic.count_tokens.transformation import AnthropicCountTokensConfig
from litellm.llms.anthropic.count_tokens.transformation import COUNT_TOKEN_OPTION_NAMES, AnthropicCountTokensConfig
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, run_aws_signing
from litellm.llms.bedrock.common_utils import BedrockError, build_mantle_messages_url
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, get_async_httpx_client
@ -23,6 +24,7 @@ class MantleCountTokensRequest(TypedDict):
_COUNT_REQUEST: Final = TypeAdapter(MantleCountTokensRequest)
_COUNT_RESPONSE: Final = TypeAdapter(dict[str, JsonValue])
_COUNT_OPTIONS: Final[TypeAdapter[Mapping[str, JsonValue]]] = TypeAdapter(Mapping[str, JsonValue])
class BedrockMantleCountTokensHandler(BaseAWSLLM):
@ -52,15 +54,18 @@ class BedrockMantleCountTokensHandler(BaseAWSLLM):
) -> dict[str, JsonValue]:
try:
request: Final = _COUNT_REQUEST.validate_python(request_data)
aws_region_name: Final = self.get_aws_region_name(
optional_params=litellm_params, model=resolved_model, model_id=None
aws_region_name: Final = await run_aws_signing(
self.get_aws_region_name, optional_params=litellm_params, model=resolved_model, model_id=None
)
endpoint_url: Final = self.get_mantle_count_tokens_endpoint(aws_region_name)
endpoint_url: Final = await run_aws_signing(self.get_mantle_count_tokens_endpoint, aws_region_name)
body: Final = self._anthropic_config.transform_request_to_count_tokens(
model=resolved_model,
messages=request["messages"],
tools=request.get("tools"),
system=request.get("system"),
optional_params=_COUNT_OPTIONS.validate_python(
{name: value for name, value in request_data.items() if name in COUNT_TOKEN_OPTION_NAMES}
),
)
verbose_logger.debug("Making bedrock-mantle count_tokens request to: %s", endpoint_url)
api_key: Final = litellm_params.get("api_key")

9
litellm/llms/bedrock/count_tokens/transformation.py
View file

@ -52,12 +52,15 @@ def _sanitized_bedrock_tool_name(raw_name: JsonValue) -> str:
def _to_bedrock_tool_spec(tool: JsonValue) -> dict[str, JsonValue]:
fields: Final = _json_dict(tool)
name: Final = _sanitized_bedrock_tool_name(fields.get("name", ""))
chat_function: Final = _json_dict(fields.get("function"))
spec: Final = chat_function or fields
name: Final = _sanitized_bedrock_tool_name(spec.get("name", ""))
schema_field: Final = "parameters" if chat_function else "input_schema"
return {
"toolSpec": {
"name": name,
"description": fields.get("description") or name,
"inputSchema": {"json": fields.get("input_schema", {"type": "object", "properties": {}})},
"description": spec.get("description") or name,
"inputSchema": {"json": spec.get(schema_field, {"type": "object", "properties": {}})},
}
}

5
litellm/llms/bedrock/realtime/handler.py
View file

@ -337,10 +337,7 @@ class BedrockRealtime(BaseAWSLLM):
# Get AWS region
if aws_region_name is None:
optional_params: Final = {
"aws_region_name": aws_region_name,
}
aws_region_name = self._get_aws_region_name(optional_params, model)
aws_region_name = await run_aws_signing(self.get_aws_region_name, {}, model)
# Get endpoint URL
if api_base is not None:

74
litellm/llms/custom_httpx/llm_http_handler.py
View file

@ -104,7 +104,7 @@ from litellm.llms.base_llm.vector_store_files.transformation import (
BaseVectorStoreFilesConfig,
)
from litellm.llms.base_llm.videos.transformation import BaseVideoConfig
from litellm.llms.bedrock.base_aws_llm import SignsRequestsWithAWS, run_aws_signing, sign_request_off_loop_if_aws
from litellm.llms.bedrock.base_aws_llm import SignsRequestsWithAWS, run_aws_signing, run_off_loop_if_aws
from litellm.llms.custom_httpx.container_handler import raise_for_error_status
from litellm.llms.custom_httpx.http_handler import (
AsyncHTTPHandler,
@ -2085,14 +2085,18 @@ class BaseLLMHTTPHandler:
raise ValueError(f"No provider config found for provider: {custom_llm_provider}")
# Validate environment and get headers
headers = provider_config.validate_environment(
headers = await run_off_loop_if_aws(
provider_config,
provider_config.validate_environment,
api_key=api_key,
api_base=api_base,
headers=headers or {},
)
# Transform the request first
data: Final = provider_config.transform_search_request(
data: Final = await run_off_loop_if_aws(
provider_config,
provider_config.transform_search_request,
query=query,
optional_params=optional_params,
api_key=api_key,
@ -2108,7 +2112,7 @@ class BaseLLMHTTPHandler:
api_key=api_key,
)
signed_headers, signed_json_body = await sign_request_off_loop_if_aws(
signed_headers, signed_json_body = await run_off_loop_if_aws(
provider_config,
provider_config.sign_request,
headers=headers,
@ -2236,7 +2240,7 @@ class BaseLLMHTTPHandler:
)
provider_config.transform_anthropic_messages_request_on_http_error(e=e, request_data=request_body)
recovered_errors = (*recovered_errors, e)
headers, signed_json_body = await sign_request_off_loop_if_aws(
headers, signed_json_body = await run_off_loop_if_aws(
provider_config,
provider_config.sign_request,
headers=headers,
@ -2405,7 +2409,9 @@ class BaseLLMHTTPHandler:
logging_obj.model_call_details.update(request_body)
# Make the request
request_url: Final = anthropic_messages_provider_config.get_complete_url(
request_url: Final = await run_off_loop_if_aws(
anthropic_messages_provider_config,
anthropic_messages_provider_config.get_complete_url,
api_base=api_base,
api_key=api_key,
model=model,
@ -2416,7 +2422,7 @@ class BaseLLMHTTPHandler:
stream=stream,
)
headers, signed_json_body = await sign_request_off_loop_if_aws(
headers, signed_json_body = await run_off_loop_if_aws(
anthropic_messages_provider_config,
anthropic_messages_provider_config.sign_request,
headers=headers,
@ -2967,7 +2973,9 @@ class BaseLLMHTTPHandler:
stream = response_api_optional_request_params.get("stream", False)
caller_requested_stream: Final = bool(stream or (extra_body or {}).get("stream"))
api_base: Final = responses_api_provider_config.get_complete_url(
api_base: Final = await run_off_loop_if_aws(
responses_api_provider_config,
responses_api_provider_config.get_complete_url,
api_base=litellm_params.api_base,
litellm_params=dict(litellm_params),
)
@ -3009,7 +3017,7 @@ class BaseLLMHTTPHandler:
fake_stream=fake_stream,
)
headers, signed_body = await sign_request_off_loop_if_aws(
headers, signed_body = await run_off_loop_if_aws(
responses_api_provider_config,
responses_api_provider_config.sign_request,
headers=validated_headers,
@ -3856,7 +3864,9 @@ class BaseLLMHTTPHandler:
litellm_params=litellm_params,
api_key=api_key,
)
complete_api_base: Final = provider_config.get_complete_file_url(
complete_api_base: Final = await run_off_loop_if_aws(
provider_config,
provider_config.get_complete_file_url,
api_base=api_base,
api_key=api_key,
model="",
@ -3867,7 +3877,9 @@ class BaseLLMHTTPHandler:
if not complete_api_base:
raise ValueError("api_base is required for create_file")
return await self.async_create_file(
transformed_request=provider_config.transform_create_file_request(
transformed_request=await run_off_loop_if_aws(
provider_config,
provider_config.transform_create_file_request,
model="",
create_file_data=create_file_data,
litellm_params=litellm_params,
@ -4334,7 +4346,9 @@ class BaseLLMHTTPHandler:
litellm_params=litellm_params,
api_key=api_key,
)
complete_api_base: Final = provider_config.get_complete_batch_url(
complete_api_base: Final = await run_off_loop_if_aws(
provider_config,
provider_config.get_complete_batch_url,
api_base=api_base,
api_key=api_key,
model=model,
@ -4345,7 +4359,9 @@ class BaseLLMHTTPHandler:
if not complete_api_base:
raise ValueError("api_base is required for create_batch")
return await self.async_create_batch(
transformed_request=provider_config.transform_create_batch_request(
transformed_request=await run_off_loop_if_aws(
provider_config,
provider_config.transform_create_batch_request,
model=model,
create_batch_data=create_batch_data,
litellm_params=litellm_params,
@ -4818,7 +4834,9 @@ class BaseLLMHTTPHandler:
**(extra_headers or {}),
}
api_base: Final = responses_api_provider_config.get_complete_url(
api_base: Final = await run_off_loop_if_aws(
responses_api_provider_config,
responses_api_provider_config.get_complete_url,
api_base=litellm_params.api_base,
litellm_params=dict(litellm_params),
)
@ -4836,7 +4854,7 @@ class BaseLLMHTTPHandler:
)
data = BaseResponsesAPIConfig.normalize_responses_api_request_dict(data)
headers, signed_body = await sign_request_off_loop_if_aws(
headers, signed_body = await run_off_loop_if_aws(
responses_api_provider_config,
responses_api_provider_config.sign_request,
headers=validated_headers,
@ -4962,7 +4980,9 @@ class BaseLLMHTTPHandler:
async_httpx_client = client
# Get URL and params from provider config
url, params = provider_config.transform_retrieve_file_request(
url, params = await run_off_loop_if_aws(
provider_config,
provider_config.transform_retrieve_file_request,
file_id=file_id,
optional_params={},
litellm_params=litellm_params,
@ -5089,7 +5109,9 @@ class BaseLLMHTTPHandler:
async_httpx_client = client
# Get URL and params from provider config
url, params = provider_config.transform_delete_file_request(
url, params = await run_off_loop_if_aws(
provider_config,
provider_config.transform_delete_file_request,
file_id=file_id,
optional_params={},
litellm_params=litellm_params,
@ -5214,7 +5236,9 @@ class BaseLLMHTTPHandler:
async_httpx_client = client
# Get URL and params from provider config
url, params = provider_config.transform_list_files_request(
url, params = await run_off_loop_if_aws(
provider_config,
provider_config.transform_list_files_request,
purpose=purpose,
optional_params={},
litellm_params=litellm_params,
@ -5298,7 +5322,9 @@ class BaseLLMHTTPHandler:
)
yield page_files[: MAX_FILE_LIST_LIMIT - listed_count]
listed_count += len(page_files)
next_request = self._next_listing_request(latest_page, provider_config, litellm_params, listed_count)
next_request = await run_off_loop_if_aws(
provider_config, self._next_listing_request, latest_page, provider_config, litellm_params, listed_count
)
if next_request is None:
return
url, params = next_request
@ -5524,7 +5550,9 @@ class BaseLLMHTTPHandler:
headers: dict,
logging_obj: LiteLLMLoggingObj,
) -> "_PreparedFileContentRequest":
url, params = provider_config.transform_file_content_request(
url, params = await run_off_loop_if_aws(
provider_config,
provider_config.transform_file_content_request,
file_content_request=file_content_request,
optional_params={},
litellm_params=litellm_params,
@ -10332,7 +10360,9 @@ class BaseLLMHTTPHandler:
if extra_headers:
headers.update(extra_headers)
api_base: Final = vector_store_provider_config.get_complete_url(
api_base: Final = await run_off_loop_if_aws(
vector_store_provider_config,
vector_store_provider_config.get_complete_url,
api_base=litellm_params.api_base,
litellm_params=dict(litellm_params),
)
@ -10366,7 +10396,7 @@ class BaseLLMHTTPHandler:
)
all_optional_params: Final[dict[str, object]] = dict(litellm_params)
all_optional_params.update(vector_store_search_optional_params or {})
headers, signed_json_body = await sign_request_off_loop_if_aws(
headers, signed_json_body = await run_off_loop_if_aws(
vector_store_provider_config,
vector_store_provider_config.sign_request,
headers=headers,

2
litellm/llms/gemini/common_utils.py
View file

@ -3,6 +3,7 @@ import datetime
import json
import math
from collections.abc import Mapping, Sequence
from types import MappingProxyType
from typing import Any, Final
import httpx
@ -498,6 +499,7 @@ class GoogleAIStudioTokenCounter(BaseTokenCounter):
messages_format: TokenCountMessagesFormat = "openai",
output_format: JsonValue = None,
context_management: CountContextManagement | None = None,
chat_params: Mapping[str, JsonValue] = MappingProxyType({}),
) -> TokenCountResponse | None:
import copy

16
litellm/llms/openai/responses/count_tokens/handler.py
View file

@ -34,7 +34,7 @@ class OpenAICountTokensHandler(OpenAICountTokensConfig):
tools: Sequence[Mapping[str, object]] | None = None,
instructions: str | None = None,
text: Mapping[str, object] | None = None,
) -> dict[str, object]:
) -> Mapping[str, object]:
"""
Handle a token counting request to OpenAI's Responses API.
@ -56,7 +56,19 @@ class OpenAICountTokensHandler(OpenAICountTokensConfig):
instructions=instructions,
text=text,
)
except ValueError as e:
verbose_logger.error("Error in CountTokens handler: %s", e)
raise OpenAIError(status_code=500, message=f"CountTokens processing error: {e}")
return await self.post_count_request(request_body, api_key, api_base, timeout)
async def post_count_request(
self,
request_body: Mapping[str, object],
api_key: str,
api_base: str | None = None,
timeout: float | httpx.Timeout | None = None,
) -> Mapping[str, object]:
try:
endpoint_url: Final = self.get_openai_count_tokens_endpoint(api_base)
verbose_logger.debug("Making request to: %s", endpoint_url)
@ -70,7 +82,7 @@ class OpenAICountTokensHandler(OpenAICountTokensConfig):
response: Final = await async_client.post(
endpoint_url,
headers=headers,
json=request_body,
json=dict(request_body),
timeout=request_timeout,
)

163
litellm/llms/openai/responses/count_tokens/token_counter.py
View file

@ -2,7 +2,9 @@
OpenAI Token Counter implementation using the Responses API /input_tokens endpoint.
"""
from collections.abc import Mapping, Sequence
from collections.abc import Awaitable, Callable, Mapping, Sequence
from dataclasses import dataclass
from types import MappingProxyType
from typing import Any, Final
import openai
@ -11,6 +13,8 @@ from typing_extensions import assert_never
from litellm._logging import verbose_logger
from litellm.litellm_core_utils.asyncify import asyncify
from litellm.litellm_core_utils.core_helpers import as_str_mapping
from litellm.litellm_core_utils.tool_fields import DropPolicy
from litellm.llms.anthropic.pass_through.convertible_blocks import deployment_drops_params, unconvertible_blocks_error
from litellm.llms.anthropic.pass_through.responses_adapters.transformation import (
LiteLLMAnthropicToResponsesAPIAdapter,
@ -23,6 +27,7 @@ from litellm.llms.openai.responses.count_tokens.handler import (
from litellm.llms.openai.responses.count_tokens.transformation import (
OpenAICountTokensConfig,
ResponsesCountRequestParts,
native_count_request,
)
from litellm.llms.openai.workload_identity import (
OpenAIWorkloadIdentityConfig,
@ -50,6 +55,46 @@ def _workload_identity_config(
return e
@dataclass(frozen=True, slots=True)
class _CountCredential:
api_base: str | None
credential: str | OpenAIWorkloadIdentityConfig
def _count_credential(deployment: Mapping[str, object] | None, model_label: str) -> _CountCredential | OpenAIError:
litellm_params: Final = as_str_mapping(None if deployment is None else deployment.get("litellm_params")) or {}
api_base: Final = litellm_params.get("api_base")
api_key: Final = litellm_params.get("api_key")
credentials: Final = get_openai_credentials(
api_base=api_base if isinstance(api_base, str) else None,
api_key=api_key if isinstance(api_key, str) else None,
)
credential: Final = openai_static_api_key(credentials.api_key) or _workload_identity_config(litellm_params)
if isinstance(credential, OpenAIError):
return credential
if credential is None:
return OpenAIError(
status_code=401,
message=(
f"No OpenAI API key is set for model {model_label}, so tokens can't be counted. Set api_key or "
"workload identity on the deployment, or OPENAI_API_KEY in the proxy environment."
),
)
return _CountCredential(api_base=credentials.api_base, credential=credential)
def _count_error(model_to_use: str, request_model: str, message: str, status_code: int) -> TokenCountResponse:
return TokenCountResponse(
total_tokens=0,
request_model=request_model,
model_used=model_to_use,
tokenizer_type="openai_api",
error=True,
error_message=message,
status_code=status_code,
)
class OpenAITokenCounter(BaseTokenCounter):
"""Token counter implementation for OpenAI provider using the Responses API."""
@ -102,6 +147,7 @@ class OpenAITokenCounter(BaseTokenCounter):
messages_format: TokenCountMessagesFormat = "openai",
output_format: JsonValue = None,
context_management: CountContextManagement | None = None,
chat_params: Mapping[str, JsonValue] = MappingProxyType({}),
) -> TokenCountResponse | None:
"""
Count tokens using OpenAI's Responses API /input_tokens endpoint.
@ -129,41 +175,9 @@ class OpenAITokenCounter(BaseTokenCounter):
rejected_before_provider_call=True,
)
deployment = deployment or {}
litellm_params: Final = deployment.get("litellm_params", {})
credentials: Final = get_openai_credentials(
api_base=litellm_params.get("api_base"), api_key=litellm_params.get("api_key")
)
credential: Final = openai_static_api_key(credentials.api_key) or _workload_identity_config(litellm_params)
if isinstance(credential, OpenAIError):
return TokenCountResponse(
total_tokens=0,
request_model=request_model,
model_used=model_to_use,
tokenizer_type="openai_api",
error=True,
error_message=credential.message,
status_code=credential.status_code,
)
if credential is None:
return TokenCountResponse(
total_tokens=0,
request_model=request_model,
model_used=model_to_use,
tokenizer_type="openai_api",
error=True,
error_message=(
f"No OpenAI API key is set for model {request_model or model_to_use}, so tokens can't be "
"counted. Set api_key or workload identity on the deployment, or OPENAI_API_KEY in the proxy "
"environment."
),
status_code=401,
)
api_base: Final = credentials.api_base
resolved: Final = _count_credential(deployment, request_model or model_to_use)
if isinstance(resolved, OpenAIError):
return _count_error(model_to_use, request_model, resolved.message, resolved.status_code)
input_items, instructions, counted_tools, text = self._responses_request_parts(
model_to_use=model_to_use,
@ -178,13 +192,9 @@ class OpenAITokenCounter(BaseTokenCounter):
if not input_items:
return None
try:
api_key: Final = (
credential
if isinstance(credential, str)
else await asyncify(get_workload_identity_bearer_token)(credential)
)
result: Final = await openai_count_tokens_handler.handle_count_tokens_request(
return await self._counted(
resolved,
lambda api_key, api_base: openai_count_tokens_handler.handle_count_tokens_request(
model=model_to_use,
input=input_items,
api_key=api_key,
@ -192,7 +202,44 @@ class OpenAITokenCounter(BaseTokenCounter):
tools=counted_tools,
instructions=instructions,
text=text,
),
model_to_use,
request_model,
)
async def count_responses_request(
self,
model_to_use: str,
request: Mapping[str, object],
policy: DropPolicy,
deployment: Mapping[str, object] | None,
request_model: str,
) -> TokenCountResponse | None:
resolved: Final = _count_credential(deployment, request_model or model_to_use)
if isinstance(resolved, OpenAIError):
return _count_error(model_to_use, request_model, resolved.message, resolved.status_code)
body: Final = native_count_request(request, model_to_use, policy)
return await self._counted(
resolved,
lambda api_key, api_base: openai_count_tokens_handler.post_count_request(body, api_key, api_base),
model_to_use,
request_model,
)
@staticmethod
async def _counted(
resolved: _CountCredential,
send: Callable[[str, str | None], Awaitable[Mapping[str, object]]],
model_to_use: str,
request_model: str,
) -> TokenCountResponse | None:
try:
api_key: Final = (
resolved.credential
if isinstance(resolved.credential, str)
else await asyncify(get_workload_identity_bearer_token)(resolved.credential)
)
result: Final = await send(api_key, resolved.api_base)
if result is not None:
return TokenCountResponse.model_validate(
@ -206,36 +253,12 @@ class OpenAITokenCounter(BaseTokenCounter):
)
except OpenAIError as e:
verbose_logger.warning("OpenAI CountTokens API error: status=%s, message=%s", e.status_code, e.message)
return TokenCountResponse(
total_tokens=0,
request_model=request_model,
model_used=model_to_use,
tokenizer_type="openai_api",
error=True,
error_message=e.message,
status_code=e.status_code,
)
return _count_error(model_to_use, request_model, e.message, e.status_code)
except openai.AuthenticationError as e:
verbose_logger.warning("OpenAI workload identity token exchange failed: %s", e.message)
return TokenCountResponse(
total_tokens=0,
request_model=request_model,
model_used=model_to_use,
tokenizer_type="openai_api",
error=True,
error_message=e.message,
status_code=401,
)
return _count_error(model_to_use, request_model, e.message, 401)
except Exception as e:
verbose_logger.warning("Error calling OpenAI CountTokens API: %s", e)
return TokenCountResponse(
total_tokens=0,
request_model=request_model,
model_used=model_to_use,
tokenizer_type="openai_api",
error=True,
error_message=str(e),
status_code=500,
)
return _count_error(model_to_use, request_model, str(e), 500)
return None

36
litellm/llms/openai/responses/count_tokens/transformation.py
View file

@ -4,12 +4,15 @@ OpenAI Responses API token counting transformation logic.
This module handles the transformation of requests to OpenAI's /v1/responses/input_tokens endpoint.
"""
import copy
from collections.abc import Mapping, Sequence
from typing import Any, Final, Literal, TypeAlias
from openai.types.responses.input_token_count_params import InputTokenCountParams
from pydantic import JsonValue, TypeAdapter, ValidationError
from typing_extensions import ReadOnly, TypedDict
from litellm.litellm_core_utils.tool_fields import DropPolicy
from litellm.llms.anthropic.pass_through.convertible_blocks import malformed_tool_result_text_error
from litellm.llms.anthropic.pass_through.messages.utils import normalize_native_messages
from litellm.llms.anthropic.pass_through.responses_adapters.transformation import (
@ -17,6 +20,8 @@ from litellm.llms.anthropic.pass_through.responses_adapters.transformation impor
)
from litellm.responses.utils import ResponsesAPIRequestUtils
from litellm.types.llms.anthropic import AnthropicMessagesRequest, AnthropicOutputSchema
from litellm.types.llms.openai import ResponsesAPIRequestParams
from litellm.types.utils import all_litellm_params
class ResponsesInputTextPart(TypedDict):
@ -340,3 +345,34 @@ class OpenAICountTokensConfig:
input_items: object, instructions: object, tools: object, text: object = None
) -> ResponsesCountRequestParts:
return _RESPONSES_COUNT_REQUEST_PARTS.validate_python((input_items, instructions, tools, text))
_COUNT_FIELDS_THE_SDK_TYPE_OMITS: Final = frozenset({"prompt"})
_COUNT_REQUEST_FIELDS: Final = (
InputTokenCountParams.__required_keys__ | InputTokenCountParams.__optional_keys__ | _COUNT_FIELDS_THE_SDK_TYPE_OMITS
)
_INFERENCE_REQUEST_FIELDS: Final = (
ResponsesAPIRequestParams.__required_keys__ | ResponsesAPIRequestParams.__optional_keys__
)
_LITELLM_REQUEST_FIELDS: Final = frozenset((*all_litellm_params, "drop_params", "additional_drop_params"))
def _forwards(field: str, policy: DropPolicy) -> bool:
if field in _LITELLM_REQUEST_FIELDS or field in policy.additional_drop_params:
return False
return field in _COUNT_REQUEST_FIELDS or (not policy.drop_params and field in _INFERENCE_REQUEST_FIELDS)
def _restored(field: str, value: object) -> object:
match field, value:
case "input", _:
return ResponsesAPIRequestUtils.restore_encrypted_content_item_ids_in_input(copy.deepcopy(value))
case "previous_response_id", str():
return ResponsesAPIRequestUtils.decode_previous_response_id_to_original_previous_response_id(value)
case _:
return value
def native_count_request(request: Mapping[str, object], model: str, policy: DropPolicy) -> Mapping[str, object]:
forwarded: Final = {field: _restored(field, value) for field, value in request.items() if _forwards(field, policy)}
return {**forwarded, "model": model}

66
litellm/llms/vertex_ai/common_utils.py
View file

@ -1,5 +1,5 @@
import re
from collections.abc import Mapping
from collections.abc import Mapping, Sequence
from copy import deepcopy
from datetime import datetime, timezone
from enum import Enum
@ -1263,6 +1263,52 @@ class VertexAITokenCounter(BaseTokenCounter):
return custom_llm_provider == LlmProviders.VERTEX_AI.value
async def _count_translated_claude_request(
self,
model_to_use: str,
messages: Sequence[Mapping[str, object]] | None,
deployment: Mapping[str, object],
request_model: str,
tools: Sequence[Mapping[str, object]] | None,
system: object | None,
messages_format: TokenCountMessagesFormat,
output_format: JsonValue,
context_management: CountContextManagement | None,
chat_params: Mapping[str, JsonValue],
) -> TokenCountResponse | None:
from litellm.llms.anthropic.count_tokens.transformation import (
AnthropicCountFields,
ChatCountTranslationError,
anthropic_count_fields,
)
from litellm.llms.vertex_ai.vertex_ai_partner_models.count_tokens.handler import count_tokens_publisher
if count_tokens_publisher(model_to_use) != "anthropic":
return None
from litellm.llms.vertex_ai.vertex_ai_partner_models.anthropic.transformation import VertexAIAnthropicConfig
match translated := await anthropic_count_fields(
model_to_use, messages or [], system, tools, messages_format, VertexAIAnthropicConfig(), chat_params
):
case ChatCountTranslationError():
return translated.response(request_model, model_to_use, "vertex_ai_partner_models")
case AnthropicCountFields():
return await self.count_tokens(
model_to_use=model_to_use,
messages=[dict(message) for message in translated.messages],
contents=None,
deployment=dict(deployment),
request_model=request_model,
tools=[dict(tool) for tool in translated.tools] if translated.tools is not None else None,
system=translated.system,
messages_format="anthropic",
output_format=output_format,
context_management=context_management,
count_options=translated.options,
)
case None:
return None
async def count_tokens(
self,
model_to_use: str,
@ -1275,6 +1321,8 @@ class VertexAITokenCounter(BaseTokenCounter):
messages_format: TokenCountMessagesFormat = "openai",
output_format: JsonValue = None,
context_management: CountContextManagement | None = None,
chat_params: Mapping[str, JsonValue] = MappingProxyType({}),
count_options: Mapping[str, JsonValue] = MappingProxyType({}),
) -> TokenCountResponse | None:
import copy
@ -1313,6 +1361,21 @@ class VertexAITokenCounter(BaseTokenCounter):
# Check if this is a partner model (Claude, Mistral, etc.)
if VertexAIPartnerModels.is_vertex_partner_model(model_to_use):
translated_count: Final = await self._count_translated_claude_request(
model_to_use=model_to_use,
messages=messages,
deployment=deployment,
request_model=request_model,
tools=tools,
system=system,
messages_format=messages_format,
output_format=output_format,
context_management=context_management,
chat_params=chat_params,
)
if translated_count is not None:
return translated_count
# Use partner models token counter
partner_models_handler: Final = VertexAIPartnerModels()
@ -1343,6 +1406,7 @@ class VertexAITokenCounter(BaseTokenCounter):
vertex_credentials=vertex_credentials,
system=system,
tools=tools,
count_options=count_options,
)
except (PartnerVertexAIError, httpx.HTTPStatusError) as e:
status_code: Final = e.response.status_code

156
litellm/llms/vertex_ai/context_caching/known_cached_contents.py Normal file
View file

@ -0,0 +1,156 @@
import asyncio
import contextlib
import re
import threading
import time
import weakref
from collections.abc import Awaitable, Callable, Mapping
from datetime import datetime, timezone
from typing import Final, Literal
KNOWN_CACHE_EXPIRY_MARGIN_SECONDS: Final = 10.0
MAX_KNOWN_CACHED_CONTENTS: Final = 10_000
_CACHE_REJECTED_STATUS_CODES: Final = frozenset({400, 403, 404})
_NAMES_CACHED_CONTENT: Final = re.compile(r"cached?\s*content", re.IGNORECASE)
_FORBIDS_FIELDS_BESIDE_A_CACHE: Final = re.compile(r"system[_\s]instruction", re.IGNORECASE)
_RFC3339_FRACTION: Final = re.compile(r"\.(\d{1,6})\d*")
def names_rejected_cache(status_code: int, message: str) -> bool:
return (
status_code in _CACHE_REJECTED_STATUS_CODES
and _NAMES_CACHED_CONTENT.search(message) is not None
and _FORBIDS_FIELDS_BESIDE_A_CACHE.search(message) is None
)
def parse_expire_time(expire_time: object) -> float | None:
"""Parse a Vertex RFC 3339 ``expireTime`` (nanosecond precision, ``Z`` suffix) to epoch seconds."""
if not isinstance(expire_time, str) or not expire_time:
return None
normalized: Final = _RFC3339_FRACTION.sub(lambda m: "." + m.group(1), expire_time.replace("Z", "+00:00"), count=1)
try:
parsed: Final = datetime.fromisoformat(normalized)
except ValueError:
return None
if parsed.tzinfo is None:
return parsed.replace(tzinfo=timezone.utc).timestamp()
return parsed.timestamp()
class KnownCachedContents:
"""Process-local record of the cachedContents this process created or found.
Vertex meters every cachedContents list page against a per-minute, per-location quota,
so listing on every cache-marked request exhausts it under load. A cache is remembered
until ``KNOWN_CACHE_EXPIRY_MARGIN_SECONDS`` before the ``expireTime`` Google reported,
and after that requests list again, as they did before the record existed.
"""
def __init__(self, clock: Callable[[], float] = time.time) -> None:
self._clock: Final = clock
self._entries: dict[tuple[tuple[str, str], str], tuple[str, float]] = {} # mutable-ok: guarded by _lock
self._lock: Final = threading.Lock()
self._inflight: weakref.WeakKeyDictionary[ # mutable-ok: private, guarded by _lock
asyncio.AbstractEventLoop, dict[tuple[tuple[str, str], str], asyncio.Future[str | None]]
] = weakref.WeakKeyDictionary()
@staticmethod
def scope(
base_url: str,
custom_llm_provider: Literal["vertex_ai", "vertex_ai_beta", "gemini"],
headers: Mapping[str, str],
vertex_project: str | None,
vertex_location: str | None,
) -> tuple[str, str]:
"""Where a cache name is valid: a Vertex project and location, or the project a Gemini key belongs to.
A custom ``api_base`` can replace the whole Vertex path, so the project and location are kept beside the URL.
The Gemini key is read from the headers actually sent, so an ``extra_headers`` override scopes the
cache to the caller's project. It is kept as-is in this in-memory scope and never logged or persisted.
"""
if custom_llm_provider != "gemini":
return (base_url, f"{vertex_project}/{vertex_location}")
sent_keys: Final = [value for name, value in headers.items() if name.lower() == "x-goog-api-key"]
return (base_url, sent_keys[-1] if sent_keys else "")
def has_not_expired(self, expire_time: object) -> bool:
expires_at: Final = parse_expire_time(expire_time)
return expires_at is None or expires_at > self._clock()
def get(self, scope: tuple[str, str], cache_key: str) -> str | None:
with self._lock:
entry: Final = self._entries.get((scope, cache_key))
if entry is None:
return None
name, reuse_until = entry
if reuse_until > self._clock():
return name
del self._entries[(scope, cache_key)]
return None
def remember(self, scope: tuple[str, str], cache_key: str, name: str | None, expire_time: object) -> None:
expires_at: Final = parse_expire_time(expire_time)
if not name or expires_at is None:
return
reuse_until: Final = expires_at - KNOWN_CACHE_EXPIRY_MARGIN_SECONDS
now: Final = self._clock()
if reuse_until <= now:
return
with self._lock:
if len(self._entries) >= MAX_KNOWN_CACHED_CONTENTS:
self._entries = {k: v for k, v in self._entries.items() if v[1] > now}
while len(self._entries) >= MAX_KNOWN_CACHED_CONTENTS:
del self._entries[next(iter(self._entries))]
self._entries[(scope, cache_key)] = (name, reuse_until)
def forget_rejected(self, cached_content: str | None, status_code: int) -> None:
if cached_content is None or status_code not in _CACHE_REJECTED_STATUS_CODES:
return
with self._lock:
self._entries = {key: entry for key, entry in self._entries.items() if entry[0] != cached_content}
async def single_flight(
self,
scope: tuple[str, str],
cache_key: str,
resolve: Callable[[], Awaitable[str | None]],
wait_seconds: float | None,
) -> str | None:
"""Run ``resolve`` once per (scope, cache_key) at a time within this event loop.
Concurrent requests for the same prefix wait for the first one's lookup-or-create
instead of each listing (and creating a duplicate cache). A waiter resolves on its
own when that first attempt fails or has not finished within ``wait_seconds``, and
never cancels it.
"""
loop: Final = asyncio.get_running_loop()
own: Final[asyncio.Future[str | None]] = loop.create_future()
with self._lock:
inflight: Final = self._inflight.setdefault(loop, {})
leader: Final = inflight.setdefault((scope, cache_key), own)
if leader is not own:
with contextlib.suppress(asyncio.TimeoutError):
if name := await asyncio.wait_for(asyncio.shield(leader), wait_seconds):
return name
return await resolve()
try:
resolved: Final = await resolve()
except BaseException:
own.set_result(None)
raise
else:
own.set_result(resolved)
return resolved
finally:
with self._lock:
inflight.pop((scope, cache_key), None)
def clear(self) -> None:
with self._lock:
self._entries.clear()
self._inflight.clear()
known_cached_contents: Final = KnownCachedContents()

74
litellm/llms/vertex_ai/context_caching/rejected_cache_retry.py Normal file
View file

@ -0,0 +1,74 @@
from collections.abc import Awaitable, Callable
from copy import deepcopy
from functools import partial
from typing import Final, TypeVar
import httpx
from litellm.litellm_core_utils.core_helpers import independent_snapshot
from litellm.llms.base_llm.chat.transformation import BaseLLMException
from litellm.types.llms.openai import AllMessageValues
from litellm.types.llms.vertex_ai import RequestBody
from litellm.utils import is_cached_message
from .known_cached_contents import names_rejected_cache
Built: Final = TypeVar("Built")
Sent: Final = TypeVar("Sent")
def cache_rebuild(
build: Callable[..., Built], messages: list[AllMessageValues], optional_params: dict
) -> Callable[[], Built] | None:
"""Copies the inputs before the first transform, which pops keys from ``optional_params`` and ``messages``."""
if not any(map(is_cached_message, messages)):
return None
return partial(build, messages=deepcopy(messages), optional_params=independent_snapshot(optional_params))
def _status_and_message(error: httpx.HTTPStatusError | BaseLLMException) -> tuple[int, str]:
match error:
case httpx.HTTPStatusError():
return error.response.status_code, error.response.text
case BaseLLMException():
return error.status_code, error.message
def _cache_rejected(body: RequestBody, error: httpx.HTTPStatusError | BaseLLMException) -> bool:
return body.get("cachedContent") is not None and names_rejected_cache(*_status_and_message(error))
async def post_with_fresh_cache_async(
body: RequestBody,
rebuild: Callable[[], Awaitable[RequestBody]] | None,
post: Callable[[RequestBody], Awaitable[Sent]],
log_request: Callable[[RequestBody], None],
) -> tuple[RequestBody, Sent]:
try:
return body, await post(body)
except (httpx.HTTPStatusError, BaseLLMException) as error:
if rebuild is None or not _cache_rejected(body, error):
raise
fresh_body: Final = await rebuild()
log_request(fresh_body)
if fresh_body.get("cachedContent") == body.get("cachedContent"):
raise
return fresh_body, await post(fresh_body)
def post_with_fresh_cache_sync(
body: RequestBody,
rebuild: Callable[[], RequestBody] | None,
post: Callable[[RequestBody], Sent],
log_request: Callable[[RequestBody], None],
) -> tuple[RequestBody, Sent]:
try:
return body, post(body)
except (httpx.HTTPStatusError, BaseLLMException) as error:
if rebuild is None or not _cache_rejected(body, error):
raise
fresh_body: Final = rebuild()
log_request(fresh_body)
if fresh_body.get("cachedContent") == body.get("cachedContent"):
raise
return fresh_body, post(fresh_body)

179
litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py
View file

@ -20,6 +20,7 @@ from litellm.utils import is_prompt_caching_valid_prompt
from ..common_utils import VertexAIError, get_vertex_base_url
from ..vertex_llm_base import VertexBase
from .known_cached_contents import KnownCachedContents, known_cached_contents
from .transformation import (
cached_messages_end_on_supported_turn,
separate_cached_messages,
@ -35,6 +36,7 @@ MAX_PAGINATION_PAGES: Final = 100 # Reasonable upper bound for pagination
class _CachedContentItem(TypedDict):
displayName: ReadOnly[NotRequired[object]]
name: ReadOnly[NotRequired[str | None]]
expireTime: ReadOnly[NotRequired[str | None]]
@with_config(ConfigDict(extra="allow", strict=True, hide_input_in_errors=True))
@ -51,6 +53,7 @@ class _CachedContentUsageMetadata(TypedDict):
@with_config(ConfigDict(extra="allow", strict=True, hide_input_in_errors=True))
class _CreatedCachedContent(TypedDict):
name: ReadOnly[NotRequired[str | None]]
expireTime: ReadOnly[NotRequired[str | None]]
usageMetadata: ReadOnly[NotRequired[_CachedContentUsageMetadata | None]]
@ -67,6 +70,14 @@ def _record_cache_creation_tokens(logging_obj: Logging, created: _CreatedCachedC
logging_obj.vertex_ai_cache_creation_input_tokens = max(tokens, 0) # rebind-ok: set provider count
def _wait_seconds(timeout: float | httpx.Timeout | None) -> float | None:
match timeout:
case httpx.Timeout():
return timeout.read
case _:
return timeout
class ContextCachingEndpoints(VertexBase):
"""
Covers context caching endpoints for Vertex AI + Google AI Studio
@ -74,8 +85,9 @@ class ContextCachingEndpoints(VertexBase):
v0: covers Google AI Studio
"""
def __init__(self) -> None:
def __init__(self, known_caches: KnownCachedContents = known_cached_contents) -> None:
super().__init__()
self._known_caches: Final = known_caches
def _get_token_and_url_context_caching(
self,
@ -129,7 +141,7 @@ class ContextCachingEndpoints(VertexBase):
self,
cache_key: str,
client: HTTPHandler,
headers: dict,
headers: dict[str, str], # mutable-ok: passed through to the HTTP client, which takes a dict
api_key: str,
api_base: str | None,
logging_obj: Logging,
@ -159,6 +171,12 @@ class ContextCachingEndpoints(VertexBase):
vertex_auth_header=vertex_auth_header,
model=model,
)
scope: Final = KnownCachedContents.scope(
base_url, custom_llm_provider, headers, vertex_project=vertex_project, vertex_location=vertex_location
)
known_name: Final = self._known_caches.get(scope, cache_key)
if known_name is not None:
return known_name
page_token: str | None = None
@ -201,8 +219,14 @@ class ContextCachingEndpoints(VertexBase):
# Check current page for matching cache_key
for cached_item in raw_response["cachedContents"]:
display_name = cached_item.get("displayName")
if display_name is not None and display_name == cache_key:
return cached_item.get("name")
if display_name is None or display_name != cache_key:
continue
expire_time = cached_item.get("expireTime")
if not self._known_caches.has_not_expired(expire_time):
continue
name = cached_item.get("name")
self._known_caches.remember(scope, cache_key, name, expire_time)
return name
# Check if there are more pages
page_token = raw_response.get("nextPageToken")
@ -216,7 +240,7 @@ class ContextCachingEndpoints(VertexBase):
self,
cache_key: str,
client: AsyncHTTPHandler,
headers: dict,
headers: dict[str, str], # mutable-ok: passed through to the HTTP client, which takes a dict
api_key: str,
api_base: str | None,
logging_obj: Logging,
@ -246,6 +270,12 @@ class ContextCachingEndpoints(VertexBase):
vertex_auth_header=vertex_auth_header,
model=model,
)
scope: Final = KnownCachedContents.scope(
base_url, custom_llm_provider, headers, vertex_project=vertex_project, vertex_location=vertex_location
)
known_name: Final = self._known_caches.get(scope, cache_key)
if known_name is not None:
return known_name
page_token: str | None = None
@ -288,8 +318,14 @@ class ContextCachingEndpoints(VertexBase):
# Check current page for matching cache_key
for cached_item in raw_response["cachedContents"]:
display_name = cached_item.get("displayName")
if display_name is not None and display_name == cache_key:
return cached_item.get("name")
if display_name is None or display_name != cache_key:
continue
expire_time = cached_item.get("expireTime")
if not self._known_caches.has_not_expired(expire_time):
continue
name = cached_item.get("name")
self._known_caches.remember(scope, cache_key, name, expire_time)
return name
# Check if there are more pages
page_token = raw_response.get("nextPageToken")
@ -451,6 +487,14 @@ class ContextCachingEndpoints(VertexBase):
raw_response_cached: Final = _CreatedCachedContentRoot.model_validate(response.json()).root
_record_cache_creation_tokens(logging_obj, raw_response_cached)
self._known_caches.remember(
KnownCachedContents.scope(
url, custom_llm_provider, headers, vertex_project=vertex_project, vertex_location=vertex_location
),
generated_cache_key,
raw_response_cached.get("name"),
raw_response_cached.get("expireTime"),
)
return (
non_cached_messages,
optional_params,
@ -548,68 +592,77 @@ class ContextCachingEndpoints(VertexBase):
generated_cache_key: Final = local_cache_obj.get_cache_key(
messages=cached_messages, tools=tools, tool_choice=tool_choice, model=model
)
google_cache_name: Final = await self.async_check_cache(
cache_key=generated_cache_key,
client=client,
headers=headers,
api_key=api_key,
api_base=api_base,
logging_obj=logging_obj,
custom_llm_provider=custom_llm_provider,
vertex_project=vertex_project,
vertex_location=vertex_location,
vertex_auth_header=vertex_auth_header,
model=model,
scope: Final = KnownCachedContents.scope(
url, custom_llm_provider, headers, vertex_project=vertex_project, vertex_location=vertex_location
)
if google_cache_name:
return non_cached_messages, optional_params, google_cache_name
## TRANSFORM REQUEST
cached_content_request_body: Final = transform_openai_messages_to_gemini_context_caching(
model=model,
messages=cached_messages,
cache_key=generated_cache_key,
custom_llm_provider=custom_llm_provider,
vertex_project=vertex_project,
vertex_location=vertex_location,
)
cached_content_request_body["tools"] = tools
if tool_choice is not None:
cached_content_request_body["toolConfig"] = tool_choice
## LOGGING
logging_obj.pre_call(
input=messages,
api_key="",
additional_args={
"complete_input_dict": cached_content_request_body,
"api_base": url,
"headers": headers,
},
)
try:
response: Final = await client.post(
url=url,
async def lookup_or_create() -> str | None:
google_cache_name: Final = await self.async_check_cache(
cache_key=generated_cache_key,
client=client,
headers=headers,
json=cached_content_request_body,
api_key=api_key,
api_base=api_base,
logging_obj=logging_obj,
custom_llm_provider=custom_llm_provider,
vertex_project=vertex_project,
vertex_location=vertex_location,
vertex_auth_header=vertex_auth_header,
model=model,
)
response.raise_for_status()
except httpx.HTTPStatusError as err:
error_code: Final = err.response.status_code
raise VertexAIError(status_code=error_code, message=err.response.text)
except httpx.TimeoutException:
raise VertexAIError(status_code=408, message="Timeout error occurred.")
raw_response_cached: Final = _CreatedCachedContentRoot.model_validate(response.json()).root
_record_cache_creation_tokens(logging_obj, raw_response_cached)
return (
non_cached_messages,
optional_params,
raw_response_cached.get("name"),
if google_cache_name:
return google_cache_name
## TRANSFORM REQUEST
cached_content_request_body: Final = transform_openai_messages_to_gemini_context_caching(
model=model,
messages=cached_messages,
cache_key=generated_cache_key,
custom_llm_provider=custom_llm_provider,
vertex_project=vertex_project,
vertex_location=vertex_location,
)
cached_content_request_body["tools"] = tools
if tool_choice is not None:
cached_content_request_body["toolConfig"] = tool_choice
## LOGGING
logging_obj.pre_call(
input=messages,
api_key="",
additional_args={
"complete_input_dict": cached_content_request_body,
"api_base": url,
"headers": headers,
},
)
try:
response: Final = await client.post(
url=url,
headers=headers,
json=cached_content_request_body,
)
response.raise_for_status()
except httpx.HTTPStatusError as err:
error_code: Final = err.response.status_code
raise VertexAIError(status_code=error_code, message=err.response.text)
except httpx.TimeoutException:
raise VertexAIError(status_code=408, message="Timeout error occurred.")
raw_response_cached: Final = _CreatedCachedContentRoot.model_validate(response.json()).root
_record_cache_creation_tokens(logging_obj, raw_response_cached)
self._known_caches.remember(
scope, generated_cache_key, raw_response_cached.get("name"), raw_response_cached.get("expireTime")
)
return raw_response_cached.get("name")
cache_name: Final = await self._known_caches.single_flight(
scope, generated_cache_key, lookup_or_create, _wait_seconds(timeout)
)
return non_cached_messages, optional_params, cache_name
def get_cache(self):
pass

251
litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py
View file

@ -96,6 +96,12 @@ from ..common_utils import (
build_vertex_schema,
supports_response_json_schema,
)
from ..context_caching.known_cached_contents import known_cached_contents
from ..context_caching.rejected_cache_retry import (
cache_rebuild,
post_with_fresh_cache_async,
post_with_fresh_cache_sync,
)
from ..vertex_llm_base import VertexBase
from .grounding_requests import calculate_grounding_requests
from .long_decode import (
@ -2786,7 +2792,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
litellm_params: dict,
api_key: str | dict | None = None,
api_base: str | None = None,
) -> dict:
) -> dict[str, str]:
default_headers: Final = {
"Content-Type": "application/json",
}
@ -2800,6 +2806,20 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
return default_headers
def _log_generate_request(
logging_obj: LoggingClass,
messages: Sequence[AllMessageValues],
api_base: str,
headers: Mapping[str, str],
body: RequestBody,
) -> None:
logging_obj.pre_call(
input=messages,
api_key="",
additional_args={"complete_input_dict": body, "api_base": api_base, "headers": headers},
)
async def make_call(
client: AsyncHTTPHandler | None, # module-level client
gemini_client: AsyncHTTPHandler | None, # if passed by user
@ -2832,6 +2852,7 @@ async def make_call(
response.raise_for_status()
except httpx.HTTPStatusError as e:
exception_string: Final = str(await e.response.aread())
known_cached_contents.forget_rejected(request_body.get("cachedContent"), e.response.status_code)
raise VertexAIError(
status_code=e.response.status_code,
message=VertexGeminiConfig().translate_exception_str(exception_string),
@ -2845,14 +2866,18 @@ async def make_call(
)
async def open_hop(body: RequestBody) -> httpx.Response:
return await client.post(
url_for(api_base, body),
headers=dict(await hop_headers()),
data=json.dumps(body),
stream=True,
logging_obj=logging_obj,
timeout=timeout,
)
try:
return await client.post(
url_for(api_base, body),
headers=dict(await hop_headers()),
data=json.dumps(body),
stream=True,
logging_obj=logging_obj,
timeout=timeout,
)
except httpx.HTTPStatusError as e:
known_cached_contents.forget_rejected(body.get("cachedContent"), e.response.status_code)
raise
completion_stream: Final = ModelResponseIterator(
streaming_response=response.aiter_lines(),
@ -2891,14 +2916,18 @@ def make_sync_call(
if client is None:
client = HTTPHandler() # Create a new client if none provided
response: Final = client.post(
url_for(api_base, request_body),
headers=headers,
data=data,
stream=True,
logging_obj=logging_obj,
timeout=timeout,
)
try:
response: Final = client.post(
url_for(api_base, request_body),
headers=headers,
data=data,
stream=True,
logging_obj=logging_obj,
timeout=timeout,
)
except httpx.HTTPStatusError as e:
known_cached_contents.forget_rejected(request_body.get("cachedContent"), e.response.status_code)
raise
if response.status_code != 200 and response.status_code != 201:
raise VertexAIError(
@ -2908,14 +2937,18 @@ def make_sync_call(
)
def open_hop(body: RequestBody) -> httpx.Response:
return client.post(
url_for(api_base, body),
headers=dict(hop_headers()),
data=json.dumps(body),
stream=True,
logging_obj=logging_obj,
timeout=timeout,
)
try:
return client.post(
url_for(api_base, body),
headers=dict(hop_headers()),
data=json.dumps(body),
stream=True,
logging_obj=logging_obj,
timeout=timeout,
)
except httpx.HTTPStatusError as e:
known_cached_contents.forget_rejected(body.get("cachedContent"), e.response.status_code)
raise
completion_stream: Final = ModelResponseIterator(
streaming_response=response.iter_lines(),
@ -2983,14 +3016,14 @@ class VertexLLM(VertexBase):
custom_llm_provider: Literal[
"vertex_ai", "vertex_ai_beta", "gemini"
], # if it's vertex_ai or gemini (google ai studio)
messages: list,
messages: list[AllMessageValues],
model_response: ModelResponse,
print_verbose: Callable,
data: dict,
timeout: float | httpx.Timeout | None,
encoding: object,
logging_obj,
stream,
logging_obj: LoggingClass,
stream: bool | None,
optional_params: dict,
litellm_params: dict,
logger_fn=None,
@ -3033,6 +3066,17 @@ class VertexLLM(VertexBase):
litellm_params=litellm_params,
)
rebuild: Final = cache_rebuild(
partial(
async_transform_request_body,
**data,
vertex_project=vertex_project,
vertex_location=vertex_location,
vertex_auth_header=auth_header,
),
messages,
optional_params,
)
request_body: Final = await async_transform_request_body(
**data,
vertex_project=vertex_project,
@ -3041,39 +3085,37 @@ class VertexLLM(VertexBase):
)
## LOGGING
logging_obj.pre_call(
input=messages,
api_key="",
additional_args={
"complete_input_dict": request_body,
"api_base": api_base,
"headers": headers,
},
log_request: Final = partial(_log_generate_request, logging_obj, messages, api_base, headers)
log_request(request_body)
open_stream: Final = partial(
make_call,
gemini_client=(client if client is not None and isinstance(client, AsyncHTTPHandler) else None),
api_base=api_base,
headers=headers,
model=model,
messages=messages,
logging_obj=logging_obj,
hop_headers=partial(
self._fresh_headers_async,
headers,
extra_headers,
vertex_credentials,
vertex_project,
custom_llm_provider,
),
timeout=timeout,
)
request_body_str: Final = json.dumps(request_body)
async def open_with_fresh_cache(client: AsyncHTTPHandler | None) -> ModelResponseIterator:
async def post(body: RequestBody) -> ModelResponseIterator:
return await open_stream(client=client, data=json.dumps(body), request_body=body)
return (await post_with_fresh_cache_async(request_body, rebuild, post, log_request))[1]
streaming_response: Final = CustomStreamWrapper(
completion_stream=None,
make_call=partial(
make_call,
gemini_client=(client if client is not None and isinstance(client, AsyncHTTPHandler) else None),
api_base=api_base,
headers=headers,
data=request_body_str,
model=model,
messages=messages,
logging_obj=logging_obj,
request_body=request_body,
hop_headers=partial(
self._fresh_headers_async,
headers,
extra_headers,
vertex_credentials,
vertex_project,
custom_llm_provider,
),
timeout=timeout,
),
make_call=open_with_fresh_cache,
model=model,
custom_llm_provider="vertex_ai_beta",
logging_obj=logging_obj,
@ -3083,7 +3125,7 @@ class VertexLLM(VertexBase):
async def async_completion(
self,
model: str,
messages: list,
messages: list[AllMessageValues],
model_response: ModelResponse,
print_verbose: Callable,
data: dict,
@ -3093,7 +3135,7 @@ class VertexLLM(VertexBase):
timeout: float | httpx.Timeout | None,
encoding: object,
logging_obj: LoggingClass,
stream,
stream: bool | None,
optional_params: dict,
litellm_params: dict,
logger_fn=None,
@ -3136,6 +3178,17 @@ class VertexLLM(VertexBase):
litellm_params=litellm_params,
)
rebuild: Final = cache_rebuild(
partial(
async_transform_request_body,
**data,
vertex_project=vertex_project,
vertex_location=vertex_location,
vertex_auth_header=auth_header,
),
messages,
optional_params,
)
request_body: Final = await async_transform_request_body(
**data,
vertex_project=vertex_project,
@ -3151,15 +3204,8 @@ class VertexLLM(VertexBase):
else:
client = client
## LOGGING
logging_obj.pre_call(
input=messages,
api_key="",
additional_args={
"complete_input_dict": request_body,
"api_base": api_base,
"headers": headers,
},
)
log_request: Final = partial(_log_generate_request, logging_obj, messages, api_base, headers)
log_request(request_body)
fresh_headers: Final = partial(
self._fresh_headers_async, headers, extra_headers, vertex_credentials, vertex_project, custom_llm_provider
@ -3175,6 +3221,7 @@ class VertexLLM(VertexBase):
)
response.raise_for_status()
except httpx.HTTPStatusError as err:
known_cached_contents.forget_rejected(body.get("cachedContent"), err.response.status_code)
raise VertexAIError(
status_code=err.response.status_code,
message=err.response.text,
@ -3188,14 +3235,14 @@ class VertexLLM(VertexBase):
)
return response
first_response: Final = await post(request_body)
sent_body, first_response = await post_with_fresh_cache_async(request_body, rebuild, post, log_request)
return _transform_long_decode_hops(
model=model,
hops=await collect_hops_async(parse_body(first_response), request_body, post),
hops=await collect_hops_async(parse_body(first_response), sent_body, post),
first_response=first_response,
model_response=model_response,
logging_obj=logging_obj,
request_body=request_body,
request_body=sent_body,
messages=messages,
)
@ -3320,6 +3367,17 @@ class VertexLLM(VertexBase):
)
## TRANSFORMATION ##
rebuild: Final = cache_rebuild(
partial(
sync_transform_request_body,
**transform_request_params,
vertex_project=vertex_project,
vertex_location=vertex_location,
vertex_auth_header=auth_header,
),
messages,
optional_params,
)
data: Final = sync_transform_request_body(
**transform_request_params,
vertex_project=vertex_project,
@ -3328,15 +3386,8 @@ class VertexLLM(VertexBase):
)
## LOGGING
logging_obj.pre_call(
input=messages,
api_key="",
additional_args={
"complete_input_dict": data,
"api_base": url,
"headers": headers,
},
)
log_request: Final = partial(_log_generate_request, logging_obj, messages, url, headers)
log_request(data)
fresh_headers: Final = partial(
self._fresh_headers, headers, extra_headers, vertex_credentials, vertex_project, custom_llm_provider
@ -3344,22 +3395,27 @@ class VertexLLM(VertexBase):
## SYNC STREAMING CALL ##
if stream is True:
request_data_str: Final = json.dumps(data)
open_stream: Final = partial(
make_sync_call,
gemini_client=(client if client is not None and isinstance(client, HTTPHandler) else None),
api_base=url,
model=model,
messages=messages,
logging_obj=logging_obj,
headers=headers,
hop_headers=fresh_headers,
timeout=timeout,
)
def open_with_fresh_cache(client: HTTPHandler | None) -> ModelResponseIterator:
def post(body: RequestBody) -> ModelResponseIterator:
return open_stream(client=client, data=json.dumps(body), request_body=body)
return post_with_fresh_cache_sync(data, rebuild, post, log_request)[1]
streaming_response: Final = CustomStreamWrapper(
completion_stream=None,
make_call=partial(
make_sync_call,
gemini_client=(client if client is not None and isinstance(client, HTTPHandler) else None),
api_base=url,
data=request_data_str,
model=model,
messages=messages,
logging_obj=logging_obj,
headers=headers,
request_body=data,
hop_headers=fresh_headers,
timeout=timeout,
),
make_call=open_with_fresh_cache,
model=model,
custom_llm_provider="vertex_ai_beta",
logging_obj=logging_obj,
@ -3385,6 +3441,7 @@ class VertexLLM(VertexBase):
)
response.raise_for_status()
except httpx.HTTPStatusError as err:
known_cached_contents.forget_rejected(body.get("cachedContent"), err.response.status_code)
raise VertexAIError(
status_code=err.response.status_code,
message=err.response.text,
@ -3398,14 +3455,14 @@ class VertexLLM(VertexBase):
)
return response
first_response: Final = post(data)
sent_body, first_response = post_with_fresh_cache_sync(data, rebuild, post, log_request)
return _transform_long_decode_hops(
model=model,
hops=collect_hops_sync(parse_body(first_response), data, post),
hops=collect_hops_sync(parse_body(first_response), sent_body, post),
first_response=first_response,
model_response=model_response,
logging_obj=logging_obj,
request_body=data,
request_body=sent_body,
messages=messages,
)

4
litellm/llms/vertex_ai/vertex_ai_partner_models/main.py
View file

@ -1,6 +1,6 @@
# What is this?
## API Handler for calling Vertex AI Partner Models
from collections.abc import Callable
from collections.abc import Callable, Mapping
from enum import Enum
from types import MappingProxyType
from typing import Final
@ -266,6 +266,7 @@ class VertexAIPartnerModels(VertexBase):
vertex_credentials=None,
system: object | None = None,
tools: list[dict[str, object]] | None = None,
count_options: Mapping[str, object] = MappingProxyType({}),
):
"""
Count tokens for Vertex AI partner models (Anthropic Claude, Mistral, etc.)
@ -302,6 +303,7 @@ class VertexAIPartnerModels(VertexBase):
**MappingProxyType(
{key: value for key, value in (("system", system), ("tools", tools)) if value is not None}
),
**count_options,
}
# Prepare litellm_params with credentials

14
litellm/main.py
View file

@ -400,7 +400,11 @@ class LiteLLM:
class Chat:
def __init__(self, params, router_obj: "Router | None"):
def __init__(
self,
params: dict[str, object], # mutable-ok: pops the acompletion flag in place
router_obj: "Router | None",
):
self.params = params
if self.params.get("acompletion", False) is True:
self.params.pop("acompletion")
@ -480,6 +484,7 @@ async def acompletion(
service_tier: str | None = None,
store: bool | None = None,
prompt_cache_key: str | None = None,
prompt_cache_retention: str | None = None,
# set api_base, api_version, api_key
base_url: str | None = None,
api_version: str | None = None,
@ -657,6 +662,7 @@ async def acompletion(
"service_tier": service_tier,
"store": store,
"prompt_cache_key": prompt_cache_key,
"prompt_cache_retention": prompt_cache_retention,
"extra_headers": extra_headers,
"acompletion": True, # assuming this is a required parameter
"thinking": thinking,
@ -5356,6 +5362,7 @@ def completion(
service_tier: str | None = None,
store: bool | None = None,
prompt_cache_key: str | None = None,
prompt_cache_retention: str | None = None,
# soon to be deprecated params by OpenAI
functions: list | None = None,
function_call: str | None = None,
@ -5478,6 +5485,7 @@ def completion(
service_tier=service_tier,
store=store,
prompt_cache_key=prompt_cache_key,
prompt_cache_retention=prompt_cache_retention,
base_url=base_url,
api_version=api_version,
api_key=api_key,
@ -5843,6 +5851,7 @@ def completion(
"service_tier": service_tier,
"store": store,
"prompt_cache_key": prompt_cache_key,
"prompt_cache_retention": prompt_cache_retention,
"allowed_openai_params": allowed_openai_params,
"base_model": base_model,
}
@ -9143,7 +9152,7 @@ async def ahealth_check(
####### HELPER FUNCTIONS ################
## Set verbose to true -> ```litellm.set_verbose = True```
def print_verbose(print_statement):
def print_verbose(print_statement: object):
try:
verbose_logger.debug(print_statement)
if litellm.set_verbose:
@ -9247,6 +9256,7 @@ async def acount_tokens(
request_model=model,
tools=tools,
system=system,
messages_format="openai",
)
if result is not None and not result.error:
return result

269
litellm/model_prices_and_context_window_backup.json
View file

@ -4916,6 +4916,26 @@
"supports_tool_choice": true,
"supports_vision": true
},
"azure/eu/gpt-4o": {
"deprecation_date": "2027-04-14",
"cache_creation_input_token_cost": 1.38e-06,
"cache_read_input_token_cost": 1.375e-06,
"input_cost_per_token": 2.75e-06,
"input_cost_per_token_batches": 1.375e-06,
"litellm_provider": "azure",
"max_input_tokens": 128000,
"max_output_tokens": 16384,
"max_tokens": 16384,
"mode": "chat",
"output_cost_per_token": 1.1e-05,
"output_cost_per_token_batches": 5.5e-06,
"source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'",
"supports_function_calling": true,
"supports_parallel_function_calling": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true
},
"azure/eu/gpt-4o-2024-11-20": {
"deprecation_date": "2027-04-14",
"cache_creation_input_token_cost": 1.38e-06,
@ -4936,6 +4956,25 @@
"supports_tool_choice": true,
"supports_vision": true
},
"azure/eu/gpt-4o-mini": {
"cache_read_input_token_cost": 8.3e-08,
"deprecation_date": "2027-04-14",
"input_cost_per_token": 1.65e-07,
"input_cost_per_token_batches": 8.3e-08,
"litellm_provider": "azure",
"max_input_tokens": 128000,
"max_output_tokens": 16384,
"max_tokens": 16384,
"mode": "chat",
"output_cost_per_token": 6.6e-07,
"source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'",
"supports_function_calling": true,
"supports_parallel_function_calling": true,
"supports_prompt_caching": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true
},
"azure/eu/gpt-4o-mini-2024-07-18": {
"cache_read_input_token_cost": 8.3e-08,
"deprecation_date": "2027-04-14",
@ -5247,6 +5286,25 @@
"supports_vision": true,
"cache_read_input_token_cost_batches": 2.75e-09
},
"azure/eu/o1": {
"cache_read_input_token_cost": 8.25e-06,
"deprecation_date": "2026-11-19",
"input_cost_per_token": 1.65e-05,
"input_cost_per_token_batches": 8.25e-06,
"litellm_provider": "azure",
"max_input_tokens": 200000,
"max_output_tokens": 100000,
"max_tokens": 100000,
"mode": "chat",
"output_cost_per_token": 6.6e-05,
"output_cost_per_token_batches": 3.3e-05,
"source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'",
"supports_function_calling": true,
"supports_parallel_function_calling": true,
"supports_prompt_caching": true,
"supports_tool_choice": true,
"supports_vision": true
},
"azure/eu/o1-2024-12-17": {
"cache_read_input_token_cost": 8.25e-06,
"deprecation_date": "2026-11-19",
@ -5283,6 +5341,24 @@
"supports_prompt_caching": true,
"supports_vision": false
},
"azure/eu/o3-mini": {
"cache_read_input_token_cost": 6.05e-07,
"deprecation_date": "2026-11-19",
"input_cost_per_token": 1.21e-06,
"input_cost_per_token_batches": 6.05e-07,
"litellm_provider": "azure",
"max_input_tokens": 200000,
"max_output_tokens": 100000,
"max_tokens": 100000,
"mode": "chat",
"output_cost_per_token": 4.84e-06,
"output_cost_per_token_batches": 2.42e-06,
"source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'",
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_tool_choice": true,
"supports_vision": false
},
"azure/eu/o3-mini-2025-01-31": {
"cache_read_input_token_cost": 6.05e-07,
"deprecation_date": "2026-11-19",
@ -11540,6 +11616,26 @@
"supports_tool_choice": true,
"supports_vision": true
},
"azure/us/gpt-4o": {
"deprecation_date": "2027-04-14",
"cache_creation_input_token_cost": 1.38e-06,
"cache_read_input_token_cost": 1.375e-06,
"input_cost_per_token": 2.75e-06,
"input_cost_per_token_batches": 1.375e-06,
"litellm_provider": "azure",
"max_input_tokens": 128000,
"max_output_tokens": 16384,
"max_tokens": 16384,
"mode": "chat",
"output_cost_per_token": 1.1e-05,
"output_cost_per_token_batches": 5.5e-06,
"source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'",
"supports_function_calling": true,
"supports_parallel_function_calling": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true
},
"azure/us/gpt-4o-2024-11-20": {
"deprecation_date": "2027-04-14",
"cache_creation_input_token_cost": 1.38e-06,
@ -11560,6 +11656,25 @@
"supports_tool_choice": true,
"supports_vision": true
},
"azure/us/gpt-4o-mini": {
"cache_read_input_token_cost": 8.3e-08,
"deprecation_date": "2027-04-14",
"input_cost_per_token": 1.65e-07,
"input_cost_per_token_batches": 8.3e-08,
"litellm_provider": "azure",
"max_input_tokens": 128000,
"max_output_tokens": 16384,
"max_tokens": 16384,
"mode": "chat",
"output_cost_per_token": 6.6e-07,
"source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'",
"supports_function_calling": true,
"supports_parallel_function_calling": true,
"supports_prompt_caching": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true
},
"azure/us/gpt-4o-mini-2024-07-18": {
"cache_read_input_token_cost": 8.3e-08,
"deprecation_date": "2027-04-14",
@ -11871,6 +11986,25 @@
"supports_tool_choice": true,
"supports_vision": true
},
"azure/us/o1": {
"cache_read_input_token_cost": 8.25e-06,
"deprecation_date": "2026-11-19",
"input_cost_per_token": 1.65e-05,
"input_cost_per_token_batches": 8.25e-06,
"litellm_provider": "azure",
"max_input_tokens": 200000,
"max_output_tokens": 100000,
"max_tokens": 100000,
"mode": "chat",
"output_cost_per_token": 6.6e-05,
"output_cost_per_token_batches": 3.3e-05,
"source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'",
"supports_function_calling": true,
"supports_parallel_function_calling": true,
"supports_prompt_caching": true,
"supports_tool_choice": true,
"supports_vision": true
},
"azure/us/o1-2024-12-17": {
"cache_read_input_token_cost": 8.25e-06,
"deprecation_date": "2026-11-19",
@ -11907,6 +12041,39 @@
"supports_prompt_caching": true,
"supports_vision": false
},
"azure/us/o3": {
"deprecation_date": "2026-11-19",
"cache_read_input_token_cost": 5.5e-07,
"input_cost_per_token": 2.2e-06,
"input_cost_per_token_batches": 1.1e-06,
"litellm_provider": "azure",
"max_input_tokens": 200000,
"max_output_tokens": 100000,
"max_tokens": 100000,
"mode": "chat",
"output_cost_per_token": 8.8e-06,
"output_cost_per_token_batches": 4.4e-06,
"source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'",
"supported_endpoints": [
"/v1/chat/completions",
"/v1/batch",
"/v1/responses"
],
"supported_modalities": [
"text",
"image"
],
"supported_output_modalities": [
"text"
],
"supports_function_calling": true,
"supports_parallel_function_calling": false,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true
},
"azure/us/o3-2025-04-16": {
"deprecation_date": "2026-11-19",
"cache_read_input_token_cost": 5.5e-07,
@ -11940,6 +12107,24 @@
"supports_tool_choice": true,
"supports_vision": true
},
"azure/us/o3-mini": {
"cache_read_input_token_cost": 6.05e-07,
"deprecation_date": "2026-11-19",
"input_cost_per_token": 1.21e-06,
"input_cost_per_token_batches": 6.05e-07,
"litellm_provider": "azure",
"max_input_tokens": 200000,
"max_output_tokens": 100000,
"max_tokens": 100000,
"mode": "chat",
"output_cost_per_token": 4.84e-06,
"output_cost_per_token_batches": 2.42e-06,
"source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'",
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_tool_choice": true,
"supports_vision": false
},
"azure/us/o3-mini-2025-01-31": {
"cache_read_input_token_cost": 6.05e-07,
"deprecation_date": "2026-11-19",
@ -11958,6 +12143,27 @@
"supports_tool_choice": true,
"supports_vision": false
},
"azure/us/o4-mini": {
"cache_read_input_token_cost": 3.03e-07,
"deprecation_date": "2026-11-19",
"input_cost_per_token": 1.21e-06,
"input_cost_per_token_batches": 6.05e-07,
"litellm_provider": "azure",
"max_input_tokens": 200000,
"max_output_tokens": 100000,
"max_tokens": 100000,
"mode": "chat",
"output_cost_per_token": 4.84e-06,
"output_cost_per_token_batches": 2.42e-06,
"source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'",
"supports_function_calling": true,
"supports_parallel_function_calling": false,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true
},
"azure/us/o4-mini-2025-04-16": {
"cache_read_input_token_cost": 3.03e-07,
"deprecation_date": "2026-11-19",
@ -61148,6 +61354,41 @@
"supports_tool_choice": false,
"supports_vision": true
},
"volcengine/doubao-seed-2-1-pro-260915": {
"cache_read_input_token_cost": 1.725e-07,
"input_cost_per_token": 8.625e-07,
"litellm_provider": "volcengine",
"max_input_tokens": 1024000,
"max_output_tokens": 256000,
"max_tokens": 256000,
"mode": "chat",
"output_cost_per_token": 4.3125e-06,
"source": "https://www.volcengine.com/docs/82379/1544106",
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_tool_choice": false,
"supports_vision": true
},
"volcengine/doubao-seed-2-1-lite-260915": {
"cache_read_input_audio_token_cost": 3.45e-07,
"cache_read_input_token_cost": 2.3e-08,
"input_cost_per_audio_token": 1.725e-06,
"input_cost_per_token": 1.15e-07,
"litellm_provider": "volcengine",
"max_input_tokens": 1024000,
"max_output_tokens": 256000,
"max_tokens": 256000,
"mode": "chat",
"output_cost_per_token": 3.88126e-07,
"source": "https://www.volcengine.com/docs/82379/1544106",
"supports_audio_input": true,
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_tool_choice": false,
"supports_vision": true
},
"volcengine/doubao-seed-2-1-turbo-260628": {
"cache_read_input_token_cost": 8.625e-08,
"input_cost_per_token": 4.3125e-07,
@ -70141,6 +70382,20 @@
"output_cost_per_token_batches": 2.42e-06,
"source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'"
},
"azure/eu/o3": {
"deprecation_date": "2026-11-19",
"cache_read_input_token_cost": 5.5e-07,
"input_cost_per_token": 2.2e-06,
"input_cost_per_token_batches": 1.1e-06,
"litellm_provider": "azure",
"max_input_tokens": 200000,
"max_output_tokens": 100000,
"max_tokens": 100000,
"mode": "chat",
"output_cost_per_token": 8.8e-06,
"output_cost_per_token_batches": 4.4e-06,
"source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'"
},
"azure/eu/o3-2025-04-16": {
"deprecation_date": "2026-11-19",
"cache_read_input_token_cost": 5.5e-07,
@ -70164,6 +70419,20 @@
"output_cost_per_token": 4.4e-05,
"source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'"
},
"azure/eu/o4-mini": {
"deprecation_date": "2026-11-19",
"cache_read_input_token_cost": 3.03e-07,
"input_cost_per_token": 1.21e-06,
"input_cost_per_token_batches": 6.05e-07,
"litellm_provider": "azure",
"max_input_tokens": 200000,
"max_output_tokens": 100000,
"max_tokens": 100000,
"mode": "chat",
"output_cost_per_token": 4.84e-06,
"output_cost_per_token_batches": 2.42e-06,
"source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'"
},
"azure/eu/o4-mini-2025-04-16": {
"deprecation_date": "2026-11-19",
"cache_read_input_token_cost": 3.03e-07,

3
litellm/models/organization_membership.py
View file

@ -6,7 +6,6 @@ Canonical definition for ``litellm_organizationmembership``. Re-exported from
"""
from datetime import datetime
from typing import Any
from pydantic import ConfigDict, model_validator
@ -24,7 +23,7 @@ class LiteLLM_OrganizationMembershipTable(LiteLLMPydanticObjectBase):
budget_id: str | None = None
created_at: datetime
updated_at: datetime
user: Any | None = None
user: object | None = None
litellm_budget_table: LiteLLM_BudgetTable | None = None
user_email: str | None = None

2
litellm/proxy/auth/auth_checks.py
View file

@ -2663,7 +2663,7 @@ async def _load_team_membership_on_cache_miss(
parent_otel_span: Span | None,
proxy_logging_obj: ProxyLogging | None,
) -> LiteLLM_TeamMembership | None:
redis_cached: Final[object] = await user_api_key_cache.async_get_cache(key=cache_key)
redis_cached: Final = await _raw_cache(user_api_key_cache).async_get_cache(key=cache_key)
redis_membership: Final = _membership_from_cached_payload(redis_cached)
if not isinstance(redis_membership, _TeamMembershipCacheMiss):
return redis_membership

12
litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py
View file

@ -796,7 +796,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
# logic becomes shared across providers.
#### CALL HOOKS - proxy only ####
def _load_credentials(self, bearer_token: str | None = None):
def _load_credentials(self, api_key: str | None = None):
try:
from botocore.credentials import Credentials
except ImportError:
@ -820,7 +820,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
credentials: Final[Credentials | None] = (
None
if bearer_token is not None
if bedrock_bearer_token(api_key) is not None
else self.get_credentials(
aws_access_key_id=aws_access_key_id,
aws_secret_access_key=aws_secret_access_key,
@ -1017,9 +1017,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
start_time: "datetime",
allow_chunking: bool,
) -> BedrockGuardrailResponse:
credentials, aws_region_name = await run_aws_signing(
self._load_credentials, bearer_token=bedrock_bearer_token(api_key)
)
credentials, aws_region_name = await run_aws_signing(self._load_credentials, api_key=api_key)
completed_chunk_usages: Final[list[BedrockGuardrailUsage]] = [] # mutable-ok: billed-chunk usage accumulator
try:
@ -1960,9 +1958,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
return BedrockGuardrailResponse()
api_key: Final[str | None] = request_data.get("api_key") if request_data else None
credentials, aws_region_name = await run_aws_signing(
self._load_credentials, bearer_token=bedrock_bearer_token(api_key)
)
credentials, aws_region_name = await run_aws_signing(self._load_credentials, api_key=api_key)
body: Final[dict[str, object]] = {"messages": checks_messages, "checks": self.checks}
prepared_request: Final = await run_aws_signing(

13
litellm/proxy/hooks/proxy_track_cost_callback.py
View file

@ -359,7 +359,6 @@ class _ProxyDBLogger(CustomLogger):
increment_spend_counters,
proxy_logging_obj,
update_cache,
update_cache_read_keys,
)
verbose_proxy_logger.debug("INSIDE _PROXY_track_cost_callback")
@ -454,12 +453,6 @@ class _ProxyDBLogger(CustomLogger):
request_tags=tags,
model_access_groups=model_access_groups,
project_id=project_id,
update_cache_read_keys=update_cache_read_keys(
user_id=user_id,
end_user_id=end_user_id,
tags=tags,
response_cost=response_cost,
),
)
if not charged:
return
@ -783,7 +776,6 @@ async def _update_database_and_spend_counters(
request_tags: list[str] | None = None,
model_access_groups: Sequence[str] | None = None,
project_id: str | None = None,
update_cache_read_keys: Sequence[str] = (),
) -> bool:
"""The reservation is reconciled before the spend is persisted, from its own read. One spend counter batch then
spans the database write and the counter update, so the post-call counters are read with a single MGET after the
@ -825,7 +817,6 @@ async def _update_database_and_spend_counters(
request_tags=request_tags,
model_access_groups=model_access_groups,
project_id=project_id,
update_cache_read_keys=update_cache_read_keys,
)
@ -846,10 +837,7 @@ async def _update_database_and_spend_counters_in_batch(
request_tags: list[str] | None,
model_access_groups: Sequence[str] | None,
project_id: str | None,
update_cache_read_keys: Sequence[str],
) -> bool:
from litellm.proxy.proxy_server import arm_update_cache_read
try:
charged: Final = await proxy_logging_obj.db_spend_update_writer.update_database(
token=user_api_key,
@ -881,7 +869,6 @@ async def _update_database_and_spend_counters_in_batch(
await _release_budget_reservation(budget_reservation=budget_reservation)
return False
await arm_update_cache_read(update_cache_read_keys)
try:
await increment_spend_counters(
token=user_api_key,

1
litellm/proxy/litellm_pre_call_utils.py
View file

@ -467,6 +467,7 @@ _ROUTER_RESERVED_METADATA_FIELDS: Final = frozenset(
{
"attempted_fallbacks",
"original_model_group",
"previous_models",
"request_retry_count",
CLIENT_OUTPUT_CEILING_METADATA_KEY,
ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY,

35
litellm/proxy/management_endpoints/workflow_management_endpoints.py
View file

@ -147,6 +147,34 @@ class _MessageCreateData(TypedDict, total=False):
session_id: str
class _EventTable(Protocol):
async def create(self, *, data: _EventCreateData) -> object: ...
class _RunTable(Protocol):
async def update(self, *, where: Mapping[str, str], data: Mapping[str, str]) -> object: ...
class _WorkflowTx(Protocol):
@property
def litellm_workflowevent(self) -> _EventTable: ...
@property
def litellm_workflowrun(self) -> _RunTable: ...
async def _append_event_in_tx(
tx: _WorkflowTx, run_id: str, event_data: _EventCreateData, new_status: str | None
) -> object:
event: Final = await tx.litellm_workflowevent.create(data=event_data)
if new_status:
await tx.litellm_workflowrun.update(
where={"run_id": run_id},
data={"status": new_status},
)
return event
# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
@ -387,12 +415,7 @@ async def append_workflow_event(
event_data["data"] = _json(data.data)
async with prisma_client.db.tx() as tx:
event: object = await tx.litellm_workflowevent.create(data=event_data)
if new_status:
await tx.litellm_workflowrun.update(
where={"run_id": run_id},
data={"status": new_status},
)
event: object = await _append_event_in_tx(tx, run_id, event_data, new_status)
return event

10
litellm/proxy/openai_files_endpoints/files_endpoints.py
View file

@ -1102,14 +1102,10 @@ async def get_file_content(
}
)
else:
# A raw cloud-storage URI (s3://, gs://) supplied here would skip the
# managed-file owner/team check that only runs for unified ids, letting
# a caller read another tenant's object by its key. Such objects are only
# reachable through their managed unified id.
if is_managed_cloud_storage_uri(file_id):
if is_managed_cloud_storage_uri(file_id) and user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN:
raise HTTPException(
status_code=400,
detail="Raw cloud storage file ids cannot be retrieved directly. Use the LiteLLM managed file id returned when the file was created.",
status_code=403,
detail="Raw cloud storage file ids can only be retrieved by a proxy admin key. Use the LiteLLM managed file id returned when the file was created.",
)
# Check for model-based credential routing
(

9
litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py
View file

@ -1459,7 +1459,6 @@ async def bedrock_proxy_route(
"""
create_request_copy(request)
aws_region_name: Final = get_secret_str(secret_name="AWS_REGION_NAME")
if not _is_bedrock_agent_runtime_route(endpoint=endpoint):
return await bedrock_llm_proxy_route(
endpoint=endpoint,
@ -1474,6 +1473,10 @@ async def bedrock_proxy_route(
detail="bedrock-agent-runtime pass-through is disabled on this proxy.",
)
from litellm.llms.bedrock.base_aws_llm import run_aws_signing, sign_aws_json_post
from litellm.llms.bedrock.chat import BedrockConverseLLM
aws_region_name: Final = await run_aws_signing(get_secret_str, secret_name="AWS_REGION_NAME")
base_target_url: Final = f"https://bedrock-agent-runtime.{aws_region_name}.{get_aws_dns_suffix(aws_region_name)}"
encoded_endpoint = httpx.URL(endpoint).path
@ -1488,10 +1491,6 @@ async def bedrock_proxy_route(
path=HttpPassThroughEndpointHelpers.join_base_and_endpoint_path(base_url, encoded_endpoint)
)
# Add or update query parameters
from litellm.llms.bedrock.base_aws_llm import run_aws_signing, sign_aws_json_post
from litellm.llms.bedrock.chat import BedrockConverseLLM
bedrock_llm: Final = BedrockConverseLLM()
# Assuming the body contains JSON data, parse it
try:

489
litellm/proxy/proxy_server.py
View file

@ -146,7 +146,6 @@ from litellm.proxy._types import (
TransformRequestBody,
UserAPIKeyAuth,
)
from litellm.proxy.common_utils.cache_pydantic_utils import CacheCodec
from litellm.proxy.common_utils.callback_utils import (
is_sensitive_callback_key,
normalize_callback_names,
@ -308,10 +307,8 @@ from litellm._logging import (
verbose_router_logger,
)
from litellm.caching.caching import DualCache, RedisCache
from litellm.caching.dual_cache import DeclaredBatchRead
from litellm.caching.redis_batch import (
active_post_call_redis_batch,
active_request_redis_batches,
drain_post_call_redis_batches,
)
from litellm.caching.redis_cache import RedisCircuitBreakerOpenError, is_redis_timeout_failure
@ -3118,8 +3115,9 @@ def cost_tracking():
from litellm.integrations.shadow_eval_logger import ShadowEvalLogger
spend_event_producer = build_spend_event_producer(CollectorSettings(), fallback=run_spend_event)
litellm.logging_callback_manager.add_litellm_callback(ProxyDBLogger(spend_event_producer))
litellm.logging_callback_manager.add_litellm_async_success_callback(ProxyDBLogger(spend_event_producer))
db_logger: Final = ProxyDBLogger(spend_event_producer)
litellm.logging_callback_manager.add_litellm_callback(db_logger)
litellm.logging_callback_manager.add_litellm_async_success_callback(db_logger)
litellm.logging_callback_manager.add_litellm_callback(ShadowEvalLogger())
@ -4235,73 +4233,61 @@ async def run_spend_counter_pipeline(pending: Sequence[PendingSpendIncrement]) -
return tuple(float(current_value) for current_value in results or ())
def update_cache_read_keys(
user_id: str | None,
end_user_id: str | None,
tags: Sequence[object] | None,
response_cost: float | None,
) -> tuple[str, ...]:
if response_cost is None:
return ()
user_keys: tuple[str, ...] = (user_id, GLOBAL_PROXY_SPEND_CACHE_KEY) if user_id is not None else ()
end_user_keys: tuple[str, ...] = (end_user_cache_key(end_user_id),) if end_user_id is not None else ()
tag_keys: tuple[str, ...] = tuple(tag_cache_key(tag) for tag in tags or () if isinstance(tag, str) and tag)
return user_keys + end_user_keys + tag_keys
_CachedSpendObjectT = TypeVar(
"_CachedSpendObjectT", LiteLLM_UserTable, LiteLLM_EndUserTable, LiteLLM_TeamTableCachedObj, LiteLLM_TagTable
)
_UPDATE_CACHE_PREFETCH_SLOT: Final = "update_cache_read"
async def arm_update_cache_read(keys: Sequence[str], cache: DualCache | None = None) -> None:
"""Declares the ``update_cache`` read on the request pipeline once the spend is persisted, so it rides the same
round trip as the post-call spend counter read instead of its own."""
request: Final = active_request_redis_batches()
target: Final = user_api_key_cache if cache is None else cache
if request is None or target.redis_cache is None or not keys:
def _add_spend_to_locally_cached_object(
cache: UserApiKeyCache, key: str, model_type: type[_CachedSpendObjectT], response_cost: float
) -> None:
current: Final = cache.get_cache(key=key, local_only=True, model_type=model_type)
if current is None:
return
with service_target(AUTH_OBJECTS_TARGET):
request.prefetched[_UPDATE_CACHE_PREFETCH_SLOT] = await target.declare_batch_get(
keys, request.batch(target.redis_cache)
)
cache.set_cache(
key=key,
value=current.model_copy(update={"spend": (current.spend or 0.0) + response_cost}),
local_only=True,
model_type=model_type,
ttl=get_management_object_ttl(cache),
)
async def _take_armed_update_cache_read(keys: Sequence[str], cache: DualCache) -> Mapping[str, object] | None:
request: Final = active_request_redis_batches()
if request is None:
async def _add_spend_to_shared_global_proxy_spend(
cache: UserApiKeyCache, response_cost: float, parent_otel_span: Span | None
) -> float | None:
if cache.redis_cache is None:
return None
armed: Final = request.prefetched.pop(_UPDATE_CACHE_PREFETCH_SLOT, None)
if not isinstance(armed, DeclaredBatchRead) or armed.keys != tuple(keys):
return None
values: Final = await cache.async_resolve_batch_get(armed)
return MappingProxyType({key: value for key, value in zip(keys, values) if value is not None})
async def _read_update_cache_values(
keys: Sequence[str], parent_otel_span: Span | None, cache: DualCache | None = None
) -> Mapping[str, object]:
"""One batched read for every object ``update_cache`` refreshes; a failed read leaves them all untouched,
exactly as a failed per-object GET left that object untouched."""
if not keys:
return MappingProxyType({})
target: Final = user_api_key_cache if cache is None else cache
try:
armed: Final = await _take_armed_update_cache_read(keys, target)
if armed is not None:
return armed
values: Final = await target.async_batch_get_cache(
keys=list(keys), parent_otel_span=parent_otel_span, throttle_redis=False
return await cache.redis_cache.async_increment_if_present(
key=GLOBAL_PROXY_SPEND_CACHE_KEY, value=response_cost, parent_otel_span=parent_otel_span
)
except Exception as e:
verbose_proxy_logger.warning(
"Spend tracking - failed to read cached spend objects. Budget enforcement may use stale spend values. "
"keys=%s - %s",
keys,
"Spend tracking - failed to add spend to the global proxy spend in Redis. "
"Budget enforcement may use stale spend values. response_cost=%s - %s",
response_cost,
str(e),
)
return MappingProxyType({})
if values is None:
return MappingProxyType({})
return MappingProxyType({key: value for key, value in zip(keys, values) if value is not None})
return None
async def _add_spend_to_global_proxy_spend(
cache: UserApiKeyCache, response_cost: float, parent_otel_span: Span | None
) -> None:
local_before: Final = cache.get_cache(key=GLOBAL_PROXY_SPEND_CACHE_KEY, local_only=True)
shared_total: Final = await _add_spend_to_shared_global_proxy_spend(cache, response_cost, parent_otel_span)
local_total: Final = cache.get_cache(key=GLOBAL_PROXY_SPEND_CACHE_KEY, local_only=True)
if not isinstance(local_total, (int, float)):
return
if shared_total is not None and local_total != local_before and shared_total < local_total:
return
cache.set_cache(
key=GLOBAL_PROXY_SPEND_CACHE_KEY,
value=local_total + response_cost if shared_total is None else shared_total,
local_only=True,
ttl=get_management_object_ttl(cache),
)
async def update_cache(
@ -4312,6 +4298,7 @@ async def update_cache(
response_cost: float | None,
parent_otel_span: Span | None,
tags: list[str] | None = None,
cache: UserApiKeyCache | None = None,
):
"""
Use this to update the cache with new user spend.
@ -4319,14 +4306,7 @@ async def update_cache(
Put any alerting logic in here.
"""
values_to_update_in_cache: Final[list[tuple[str, object]]] = []
with service_target(AUTH_OBJECTS_TARGET):
cached_values: Final = await _read_update_cache_values(
keys=update_cache_read_keys(
user_id=user_id, end_user_id=end_user_id, tags=tags, response_cost=response_cost
),
parent_otel_span=parent_otel_span,
)
target: Final = user_api_key_cache if cache is None else cache
### UPDATE KEY SPEND ###
async def _update_key_cache(token: str, response_cost: float):
@ -4336,7 +4316,7 @@ async def update_cache(
else:
hashed_token = token
verbose_proxy_logger.debug("_update_key_cache: hashed_token=%s", hashed_token)
existing_spend_obj = await user_api_key_cache.async_get_cache(key=hashed_token, model_type=UserAPIKeyAuth)
existing_spend_obj = await target.async_get_cache(key=hashed_token, model_type=UserAPIKeyAuth)
verbose_proxy_logger.debug("_update_key_cache: existing_spend_obj=%s", existing_spend_obj)
if existing_spend_obj is None:
return
@ -4378,184 +4358,17 @@ async def update_cache(
)
# set cooldown on alert
### UPDATE USER SPEND ###
async def _update_user_cache():
## UPDATE CACHE FOR USER ID + GLOBAL PROXY
if response_cost is None:
return
user_ids: Final = [user_id]
def _add_spend(key: str, model_type: type[_CachedSpendObjectT], cost: float) -> None:
try:
for _id in user_ids:
# Fetch the existing cost for the given user
if _id is None:
continue
cached_user = cached_values.get(_id)
if cached_user is None:
# do nothing if there is no cache value
return
existing_spend_obj = CacheCodec.deserialize(cached_user, LiteLLM_UserTable)
if existing_spend_obj is None:
return
verbose_proxy_logger.debug(
"_update_user_db: existing spend: %s; response_cost: %s", existing_spend_obj, response_cost
)
existing_spend = existing_spend_obj.spend or 0.0
# Calculate the new cost by adding the existing cost and response_cost
new_spend = existing_spend + response_cost
existing_spend_obj.spend = new_spend
values_to_update_in_cache.append(
(
_id,
CacheCodec.serialize(existing_spend_obj, model_type=LiteLLM_UserTable),
)
)
## UPDATE GLOBAL PROXY ##
global_proxy_spend: Final = cached_values.get(GLOBAL_PROXY_SPEND_CACHE_KEY)
if not isinstance(global_proxy_spend, (int, float)):
# do nothing if not in cache
return
elif response_cost is not None:
increment: Final = global_proxy_spend + response_cost
values_to_update_in_cache.append((GLOBAL_PROXY_SPEND_CACHE_KEY, increment))
_add_spend_to_locally_cached_object(target, key, model_type, cost)
except Exception as e:
verbose_proxy_logger.warning(
"Spend tracking - failed to update user spend in cache. "
"Spend tracking - failed to update %s spend in cache. "
"Budget enforcement may use stale spend values. "
"user_id=%s, response_cost=%s - %s\n%s",
user_id,
response_cost,
str(e),
traceback.format_exc(),
)
### UPDATE END-USER SPEND ###
async def _update_end_user_cache():
if end_user_id is None or response_cost is None:
return
_id: Final = end_user_cache_key(end_user_id)
try:
# Fetch the existing cost for the given user
cached_end_user: Final = cached_values.get(_id)
if cached_end_user is None:
# if user does not exist in LiteLLM_UserTable, create a new user
# do nothing if end-user not in api key cache
return
existing_spend_obj: Final = CacheCodec.deserialize(cached_end_user, LiteLLM_EndUserTable)
if existing_spend_obj is None:
return
verbose_proxy_logger.debug(
"_update_end_user_db: existing spend: %s; response_cost: %s", existing_spend_obj, response_cost
)
existing_spend: Final = existing_spend_obj.spend or 0.0
# Calculate the new cost by adding the existing cost and response_cost
new_spend: Final = existing_spend + response_cost
existing_spend_obj.spend = new_spend
values_to_update_in_cache.append(
(
_id,
CacheCodec.serialize(existing_spend_obj, model_type=LiteLLM_EndUserTable),
)
)
except Exception as e:
verbose_proxy_logger.warning(
"Spend tracking - failed to update end user spend in cache. "
"Budget enforcement may use stale spend values. "
"end_user_id=%s, response_cost=%s - %s\n%s",
end_user_id,
response_cost,
str(e),
traceback.format_exc(),
)
### UPDATE TEAM SPEND ###
def _update_team_cache():
if team_id is None or response_cost is None:
return
_id: Final = f"team_id:{team_id}"
try:
current_team: Final = CacheCodec.deserialize(
user_api_key_cache.get_cache(key=_id, local_only=True), LiteLLM_TeamTableCachedObj
)
if current_team is None:
return
verbose_proxy_logger.debug(
"_update_team_db: existing spend: %s; response_cost: %s", current_team, response_cost
)
user_api_key_cache.set_cache(
key=_id,
value=CacheCodec.serialize(
current_team.model_copy(update={"spend": (current_team.spend or 0.0) + response_cost}),
model_type=LiteLLM_TeamTableCachedObj,
),
local_only=True,
ttl=get_management_object_ttl(user_api_key_cache),
)
except Exception as e:
verbose_proxy_logger.warning(
"Spend tracking - failed to update team spend in cache. "
"Budget enforcement may use stale spend values. "
"team_id=%s, response_cost=%s - %s\n%s",
team_id,
response_cost,
str(e),
traceback.format_exc(),
)
### UPDATE TAG SPEND ###
async def _update_tag_cache():
"""
Update the tag cache with the new spend.
"""
if tags is None or response_cost is None:
return
try:
for tag_name in tags:
if not tag_name or not isinstance(tag_name, str):
continue
cache_key = tag_cache_key(tag_name)
# Fetch the existing tag object from cache
cached_tag = cached_values.get(cache_key)
if cached_tag is None:
# do nothing if tag not in api key cache
continue
existing_tag_obj = CacheCodec.deserialize(cached_tag, LiteLLM_TagTable)
if existing_tag_obj is None:
continue
verbose_proxy_logger.debug(
"_update_tag_cache: existing spend for tag=%s: %s; response_cost: %s",
tag_name,
existing_tag_obj,
response_cost,
)
existing_spend = existing_tag_obj.spend or 0.0
# Calculate the new cost by adding the existing cost and response_cost
new_spend = existing_spend + response_cost
existing_tag_obj.spend = new_spend
values_to_update_in_cache.append(
(
cache_key,
CacheCodec.serialize(existing_tag_obj, model_type=LiteLLM_TagTable),
)
)
except Exception as e:
verbose_proxy_logger.warning(
"Spend tracking - failed to update tag spend in cache. "
"Budget enforcement may use stale spend values. "
"tags=%s, response_cost=%s - %s\n%s",
tags,
response_cost,
"key=%s, response_cost=%s - %s\n%s",
model_type.__name__,
key,
cost,
str(e),
traceback.format_exc(),
)
@ -4563,42 +4376,23 @@ async def update_cache(
with service_target(AUTH_OBJECTS_TARGET):
if token is not None and response_cost is not None:
await _update_key_cache(token=token, response_cost=response_cost)
if response_cost is None:
return
if user_id is not None:
await _update_user_cache()
if user_id is not None:
_add_spend(user_id, LiteLLM_UserTable, response_cost)
if end_user_id is not None:
_add_spend(end_user_cache_key(end_user_id), LiteLLM_EndUserTable, response_cost)
if team_id is not None:
_add_spend(f"team_id:{team_id}", LiteLLM_TeamTableCachedObj, response_cost)
for tag_name in tags or ():
if tag_name and isinstance(tag_name, str):
_add_spend(tag_cache_key(tag_name), LiteLLM_TagTable, response_cost)
if end_user_id is not None:
await _update_end_user_cache()
if team_id is not None:
_update_team_cache()
if tags is not None:
await _update_tag_cache()
global_proxy_spend_key: Final = GLOBAL_PROXY_SPEND_CACHE_KEY
local_object_updates: Final = tuple((k, v) for k, v in values_to_update_in_cache if k != global_proxy_spend_key)
shared_scalar_updates: Final = tuple((k, v) for k, v in values_to_update_in_cache if k == global_proxy_spend_key)
if local_object_updates:
with service_target(AUTH_OBJECTS_TARGET):
asyncio.create_task(
user_api_key_cache.async_set_cache_pipeline(
cache_list=list(local_object_updates),
ttl=get_management_object_ttl(user_api_key_cache),
litellm_parent_otel_span=parent_otel_span,
local_only=True,
)
)
if shared_scalar_updates:
with service_target(SPEND_COUNTERS_TARGET):
asyncio.create_task(
user_api_key_cache.async_set_cache_pipeline(
cache_list=list(shared_scalar_updates),
ttl=get_management_object_ttl(user_api_key_cache),
litellm_parent_otel_span=parent_otel_span,
)
)
if litellm.max_budget <= 0:
return
with service_target(SPEND_COUNTERS_TARGET):
await _add_spend_to_global_proxy_spend(target, response_cost, parent_otel_span)
def run_ollama_serve():
@ -14701,7 +14495,9 @@ async def run_thread(
# dependencies=[Depends(user_api_key_auth)],
# )
# async def get_available_routes(user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth)):
from litellm.litellm_core_utils.tool_fields import DropPolicy, drop_policy
from litellm.llms.base_llm.base_utils import BaseTokenCounter
from litellm.llms.openai.responses.count_tokens.token_counter import OpenAITokenCounter
from litellm.proxy.db.routing_prisma_wrapper import WriterPinnedClient
from litellm.repositories.config_repository import ConfigRepository
from litellm.repositories.model_repository import ModelRepository
@ -14714,6 +14510,11 @@ from litellm.repositories.table_repositories import (
)
from litellm.repositories.team_repository import TeamRepository
from litellm.repositories.user_repository import UserRepository
from litellm.responses.litellm_completion_transformation.count_request import (
ChatCountBridgeError,
chat_count_request,
with_provider_file_ids,
)
def _get_provider_token_counter(
@ -14771,6 +14572,18 @@ def _get_provider_token_counter(
return None, None, None
def _provider_count_outcome(
result: Union["TokenCountResponse", None],
) -> Union["TokenCountResponse", ProviderCountFailure, None]:
if result is not None and result.error is True:
return ProviderCountFailure(
status_code=result.status_code or 500,
message=result.error_message or "Token counting failed",
rejected_before_provider_call=result.rejected_before_provider_call,
)
return result
def _with_resolved_api_key(litellm_params: object) -> Mapping[str, object]:
call_params: Final = _LITELLM_PARAMS_COPY.validate_python(litellm_params)
try:
@ -14785,6 +14598,12 @@ def _with_resolved_api_key(litellm_params: object) -> Mapping[str, object]:
return call_params
def _with_resolved_credentials(deployment: Mapping[str, object] | None) -> dict[str, object] | None:
if deployment is None:
return None
return {**deployment, "litellm_params": _with_resolved_api_key(deployment.get("litellm_params", {}))}
_DEPLOYMENT_LITELLM_PARAMS: Final = TypeAdapter(Mapping[str, object])
@ -14825,6 +14644,7 @@ async def _try_provider_token_count(
messages_format: TokenCountMessagesFormat = "openai",
output_format: JsonValue = None,
context_management: CountContextManagement | None = None,
chat_params: Mapping[str, JsonValue] = MappingProxyType({}),
) -> Union["TokenCountResponse", ProviderCountFailure, None]:
"""Attempt provider-specific token counting. Returns None to fall through to local counting."""
if not provider_counter.should_use_token_counting_api(custom_llm_provider=custom_llm_provider):
@ -14844,31 +14664,44 @@ async def _try_provider_token_count(
model_to_use=model_to_use or "",
messages=messages,
contents=contents,
deployment=None
if deployment_in_region is None
else {
**deployment_in_region,
"litellm_params": _with_resolved_api_key(deployment_in_region.get("litellm_params", {})),
},
deployment=_with_resolved_credentials(deployment_in_region),
request_model=request_model,
tools=tools,
system=system,
messages_format=messages_format,
output_format=output_format,
context_management=context_management,
chat_params=chat_params,
)
except httpx.HTTPStatusError as e:
return ProviderCountFailure(
status_code=getattr(e, "status_code", None) or e.response.status_code,
message=getattr(e, "message", None) or str(e),
)
if result is not None and result.error is True:
return ProviderCountFailure(
status_code=result.status_code or 500,
message=result.error_message or "Token counting failed",
rejected_before_provider_call=result.rejected_before_provider_call,
return _provider_count_outcome(result)
async def _try_native_responses_count(
provider_counter: OpenAITokenCounter,
custom_llm_provider: str | None,
model_to_use: str,
_region: str | None,
responses_request: Mapping[str, object],
policy: DropPolicy,
deployment: Mapping[str, object] | None,
request_model: str,
) -> Union["TokenCountResponse", ProviderCountFailure, None]:
if not provider_counter.should_use_token_counting_api(custom_llm_provider=custom_llm_provider):
return None
return _provider_count_outcome(
await provider_counter.count_responses_request(
model_to_use=model_to_use,
request=responses_request,
policy=policy,
deployment=_with_resolved_credentials(deployment),
request_model=request_model,
)
return result
)
async def _provider_token_count(
@ -15150,6 +14983,7 @@ class _TokenCountDeployment(TypedDict, total=False):
_TOKEN_COUNT_DEPLOYMENT: Final = TypeAdapter(_TokenCountDeployment)
_ROUTED_DEPLOYMENT: Final[TypeAdapter[Mapping[str, object] | None]] = TypeAdapter(Mapping[str, object] | None)
def _model_group_deployment(router: Router | None, model_group: str) -> Mapping[str, object] | None:
@ -15228,6 +15062,43 @@ def _count_context_management(
)
def _responses_routing_kwargs(responses_request: Mapping[str, object] | None) -> Mapping[str, object]:
previous_response_id: Final = None if responses_request is None else responses_request.get("previous_response_id")
return {"previous_response_id": previous_response_id} if isinstance(previous_response_id, str) else {}
def _count_drop_policy(deployment: Mapping[str, object] | None, responses_request: Mapping[str, object]) -> DropPolicy:
litellm_params: Final = as_str_mapping(None if deployment is None else deployment.get("litellm_params")) or {}
router_defaults: Final[Mapping[str, object]] = (
{} if llm_router is None else _router_default_litellm_params(llm_router)
)
return drop_policy({**litellm_params, **router_defaults, **responses_request})
def _native_responses_count(
provider_counter: BaseTokenCounter,
custom_llm_provider: str | None,
responses_request: Mapping[str, object] | None,
deployment: Mapping[str, object] | None,
request_model: str,
) -> Callable[[str, str | None], Awaitable[Union["TokenCountResponse", ProviderCountFailure, None]]] | None:
if responses_request is None or not isinstance(provider_counter, OpenAITokenCounter):
return None
return partial(
_try_native_responses_count,
provider_counter,
custom_llm_provider,
responses_request=responses_request,
policy=_count_drop_policy(deployment, responses_request),
deployment=deployment,
request_model=request_model,
)
def _deployment_api_base(deployment: _TokenCountDeployment) -> str | None:
return _non_empty_str(deployment.get("litellm_params", {}).get("api_base"))
def _system_message(system: object) -> ChatCompletionSystemMessage | None:
if not isinstance(system, (str, list)) or not system:
return None
@ -15260,6 +15131,7 @@ async def count_request_tokens(
output_format: JsonValue = None,
context_management: JsonValue = None,
request_litellm_params: Mapping[str, JsonValue] = MappingProxyType({}),
responses_request: Mapping[str, object] | None = None,
) -> TokenCountResponse:
prompt: Final = request.prompt
messages: Final = request.messages
@ -15281,7 +15153,7 @@ async def count_request_tokens(
try:
deployment = await llm_router.async_get_available_deployment(
model=request.model,
request_kwargs={},
request_kwargs=dict(_responses_routing_kwargs(responses_request)),
)
except Exception:
verbose_proxy_logger.exception(
@ -15335,23 +15207,54 @@ async def count_request_tokens(
request_messages: Final = _counted_request_messages(messages, messages_format, count_deployment)
contents_estimable: Final = not contents or (bool(messages) and _estimable_contents(contents))
prompt_as_message: Final = prompt_only and custom_llm_provider in _PROVIDERS_COUNTING_WITHOUT_MESSAGES
routed_deployment: Final = _ROUTED_DEPLOYMENT.validate_python(deployment)
routed_request: Final = (
None if responses_request is None else with_provider_file_ids(responses_request, routed_deployment)
)
chat_count: Final = (
None
if routed_request is None or provider_counter is None or isinstance(provider_counter, OpenAITokenCounter)
else await chat_count_request(
model_to_use,
custom_llm_provider,
_deployment_api_base(deployment_view),
routed_request,
_count_drop_policy(routed_deployment, routed_request),
)
)
if isinstance(chat_count, ChatCountBridgeError):
raise ProxyException(message=chat_count.message, type="invalid_request_error", param=None, code=400)
ladder: Final = (
None
if provider_counter is None
else await _provider_token_count(
count_with=partial(
count_with=_native_responses_count(
provider_counter,
custom_llm_provider,
routed_request,
routed_deployment,
request.model,
)
or partial(
_try_provider_token_count,
provider_counter,
custom_llm_provider,
messages=[{"role": "user", "content": prompt}] if prompt_as_message else request_messages,
messages=(
[dict(message) for message in chat_count.messages]
if chat_count is not None
else [{"role": "user", "content": prompt}]
if prompt_as_message
else request_messages
),
contents=None if prompt_as_message else contents,
deployment=count_deployment,
request_model=request.model,
tools=tools,
tools=tools if chat_count is None else [dict(tool) for tool in chat_count.tools or ()] or None,
system=system,
messages_format=messages_format,
output_format=output_format,
context_management=_count_context_management(context_management, request.model, request_litellm_params),
chat_params=MappingProxyType({}) if chat_count is None else chat_count.params,
),
model_to_use=model_to_use,
settings=settings,

93
litellm/proxy/response_api_endpoints/endpoints.py
View file

@ -32,12 +32,13 @@ from typing_extensions import NotRequired, ReadOnly, TypedDict
import litellm
from litellm._logging import verbose_proxy_logger
from litellm.anthropic_interface.exceptions.exception_mapping_utils import AnthropicExceptionMapping
from litellm.caching.dual_cache import DualCache
from litellm.constants import EMPTY_MAPPING, MANAGED_BACKGROUND_RESPONSE_METADATA_KEY, TERMINAL_RESPONSE_STATUSES
from litellm.integrations.custom_guardrail import ModifyResponseException
from litellm.litellm_core_utils.core_helpers import as_str_mapping, get_metadata_variable_name_from_kwargs
from litellm.litellm_core_utils.internal_call_metadata import is_background_response, requests_background
from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup
from litellm.litellm_core_utils.tool_fields import first_validation_error
from litellm.litellm_core_utils.tool_fields import TOOL_FIELDS, first_validation_error
from litellm.llms.openai.responses.guardrail_translation.handler import (
OpenAIResponsesHandler,
build_blocked_response,
@ -79,6 +80,7 @@ from litellm.proxy.pass_through_endpoints.llm_provider_handlers.batch_attributio
)
from litellm.proxy.route_llm_request import raise_if_required_body_param_missing
from litellm.proxy.utils import handle_exception_on_proxy
from litellm.responses.litellm_completion_transformation.count_request import MODEL_FILE_ID_MAPPING_FIELD
from litellm.responses.streaming_iterator import GuardedResponsesWsFrame, ResponsesWsFrameGuard, spawn_logging_task
from litellm.router import Router
from litellm.types.llms.base import LiteLLMBaseModel
@ -89,7 +91,7 @@ from litellm.types.llms.openai import (
)
from litellm.types.responses.main import DeleteResponseResult
from litellm.types.token_count import TokenCountSourceFields
from litellm.types.utils import TokenCountResponse
from litellm.types.utils import ADDRESSED_RESPONSE_ID_FIELD, TokenCountResponse
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
@ -1388,13 +1390,15 @@ class _TokenCountPayload(TypedDict):
class _TokenCounter(Protocol):
def __call__(self, request: TokenCountRequest, call_endpoint: bool) -> Awaitable[TokenCountResponse]: ...
def __call__(
self, *, request: TokenCountRequest, call_endpoint: bool, responses_request: Mapping[str, object] | None
) -> Awaitable[TokenCountResponse]: ...
def _proxy_token_counter() -> _TokenCounter:
from litellm.proxy.proxy_server import token_counter
from litellm.proxy.proxy_server import count_request_tokens
return token_counter
return count_request_tokens
_token_counter_dep: Final = Depends(_proxy_token_counter)
@ -1409,6 +1413,67 @@ def _proxy_llm_router() -> Router | None:
_llm_router_dep: Final = Depends(_proxy_llm_router)
def _proxy_response_id_security() -> ResponsesIDSecurity | None:
from litellm.proxy.proxy_server import proxy_logging_obj
security_hook: Final = proxy_logging_obj.get_proxy_hook("responses_id_security")
return security_hook if isinstance(security_hook, ResponsesIDSecurity) else None
_response_id_security_dep: Final = Depends(_proxy_response_id_security)
class _ManagedFileAccess(Protocol):
async def async_pre_call_hook(
self,
user_api_key_dict: UserAPIKeyAuth,
cache: DualCache,
data: dict[str, object],
call_type: Literal["aresponses"],
) -> object: ...
def _proxy_managed_files_hook() -> _ManagedFileAccess | None:
from litellm.proxy.proxy_server import proxy_logging_obj
return proxy_logging_obj.get_proxy_hook("managed_files")
_managed_files_hook_dep: Final = Depends(_proxy_managed_files_hook)
async def _with_managed_file_id_mapping(
data: Mapping[str, object], managed_files: _ManagedFileAccess | None, user_api_key_dict: UserAPIKeyAuth
) -> Mapping[str, object]:
caller_fields: Final = {name: value for name, value in data.items() if name != MODEL_FILE_ID_MAPPING_FIELD}
if managed_files is None:
return caller_fields
from litellm.proxy.proxy_server import user_api_key_cache
hooked: Final = dict(caller_fields)
try:
await managed_files.async_pre_call_hook(
user_api_key_dict=user_api_key_dict, cache=user_api_key_cache, data=hooked, call_type="aresponses"
)
except HTTPException as e:
raise ProxyException(message=str(e.detail), type="permission_error", param=None, code=e.status_code) from e
mapping: Final = hooked.get(MODEL_FILE_ID_MAPPING_FIELD)
return caller_fields if mapping is None else {**caller_fields, MODEL_FILE_ID_MAPPING_FIELD: mapping}
def _with_authorized_previous_response(
data: Mapping[str, object], security: ResponsesIDSecurity | None, user_api_key_dict: UserAPIKeyAuth
) -> Mapping[str, object]:
previous_response_id: Final = data.get("previous_response_id")
if security is None or not isinstance(previous_response_id, str) or not previous_response_id:
return data
try:
authorized_id: Final = security.authorize_response_id(previous_response_id, user_api_key_dict)
except HTTPException as e:
raise ProxyException(message=str(e.detail), type="permission_error", param=None, code=e.status_code) from e
return {**data, "previous_response_id": authorized_id, ADDRESSED_RESPONSE_ID_FIELD: previous_response_id}
def _responses_invalid_request_response(message: str, param: str | None, code: str | None) -> JSONResponse:
body: Final[_ResponsesApiErrorBody] = {
"error": {
@ -1546,6 +1611,8 @@ async def responses_input_tokens(
token_counter: _TokenCounter = _token_counter_dep,
llm_router: Router | None = _llm_router_dep,
user_api_key_dict: UserAPIKeyAuth = _user_api_key_auth_dep,
response_id_security: ResponsesIDSecurity | None = _response_id_security_dep,
managed_files: _ManagedFileAccess | None = _managed_files_hook_dep,
):
"""
Count the input tokens of a Responses API request without calling the model.
@ -1579,7 +1646,7 @@ async def responses_input_tokens(
)
counted_fields: Final = {
name: data[name] for name in ("model", "tools", "drop_params", "additional_drop_params") if name in data
name: data[name] for name in ("model", "drop_params", "additional_drop_params", *TOOL_FIELDS) if name in data
}
token_fields: Final = resolve_counted_tool_fields(counted_fields, llm_router, user_api_key_dict.team_id)
if isinstance(token_fields, ProxyException):
@ -1604,8 +1671,20 @@ async def responses_input_tokens(
message=f"Invalid request for token counting: {e}", param=None, code=None
)
request_fields: Final = _PROXY_SERVER_REQUEST_ADAPTER.validate_python(data)
responses_request: Final = await _with_managed_file_id_mapping(
_with_authorized_previous_response(
{**{name: value for name, value in request_fields.items() if name not in counted_fields}, **token_fields},
response_id_security,
user_api_key_dict,
),
managed_files,
user_api_key_dict,
)
try:
token_response: Final = await token_counter(request=token_request, call_endpoint=True)
token_response: Final = await token_counter(
request=token_request, call_endpoint=True, responses_request=responses_request
)
except ProxyException as e:
if e.type != _PROVIDER_COUNT_ERROR_TYPE:
raise

208
litellm/proxy/utils.py
View file

@ -1031,9 +1031,7 @@ def _failure_usage_to_lift(
at the rate the router registered under ``deployment_id`` (``Logging.get_router_model_id``)
when the deployment declares one, else the cost map's rate for ``base_model`` (an Azure
deployment's ``model_info.base_model``) or the deployment model; any other post-dispatch
failure stays at zero cost. The raw request body backfills
the system prompt when the SDK bridges an endpoint (e.g. /v1/messages on a
chat-completions provider) without filling optional_params. Returns the
failure stays at zero cost. Returns the
(combined_usage_object, response_cost) pair to lift, or None."""
recovered_usage: Final = model_call_details.get("combined_usage_object")
failed_attempts: Final = _FAILED_ATTEMPTS_USAGE.validate_python(model_call_details.get("failed_attempts_usage", ()))
@ -1047,6 +1045,22 @@ def _failure_usage_to_lift(
# The proxy-only error handler stamps this flag on a client disconnect too, even though the
# request had already reached the provider, so a billed failure keeps its estimate.
return None
estimated_usage: Final = _estimated_prompt_usage(model_call_details, request_body)
if estimated_usage is None:
return None
if not provider_billed:
return estimated_usage, 0.0
return estimated_usage, _estimated_input_cost(
model_call_details, estimated_usage.prompt_tokens, deployment_id=deployment_id, base_model=base_model
)
def _estimated_prompt_usage(
model_call_details: Mapping[str, object], request_body: Mapping[str, object]
) -> Usage | None:
"""The input-only Usage estimate for the prompt this request sent, or None when the call type has no
estimable input. The raw request body backfills the system prompt when the SDK bridges an endpoint
(e.g. /v1/messages on a chat-completions provider) without filling optional_params."""
if str(model_call_details.get("call_type")) not in _INPUT_ESTIMABLE_CALL_TYPES:
return None
optional_params: Final = model_call_details.get("optional_params")
@ -1056,17 +1070,118 @@ def _failure_usage_to_lift(
else None
)
system_input: Final = dispatched_system or request_body.get("system") or request_body.get("instructions")
estimated_usage: Final = _estimate_dispatched_failure_usage(
return _estimate_dispatched_failure_usage(
model=str(model_call_details.get("model") or ""),
request_input=model_call_details.get("messages"),
system_input=system_input,
)
if estimated_usage is None:
return None
if not provider_billed:
return estimated_usage, 0.0
return estimated_usage, _estimated_input_cost(
model_call_details, estimated_usage.prompt_tokens, deployment_id=deployment_id, base_model=base_model
@dataclass(frozen=True)
class BilledFailedAttempt:
"""An earlier retry attempt the provider also billed, written as its own spend row."""
attempted_retries: int
deployment_id: str | None
usage: Usage
response_cost: float
_NO_RETRY_METADATA: Final[Mapping[str, object]] = MappingProxyType({})
def _retry_metadata_bucket(litellm_params: Mapping[str, object] | None) -> Mapping[str, object]:
"""The request metadata bucket the router writes retry state into, ``litellm_metadata`` when the
route carries one (/messages, /responses) and ``metadata`` otherwise, same order as the deployment id."""
if litellm_params is None:
return _NO_RETRY_METADATA
for metadata_key in ("litellm_metadata", "metadata"):
metadata = litellm_params.get(metadata_key)
if isinstance(metadata, dict):
return cast("Mapping[str, object]", metadata) # cast-ok: untyped router dict, keys read with .get
return _NO_RETRY_METADATA
def _attempt_number(record: Mapping[str, object], index: int) -> int:
"""The attempt a retry record describes: its own ``attempted_retries`` count, else its position."""
recorded: Final = record.get("attempted_retries")
return recorded if isinstance(recorded, int) else index
def _retry_attempt_records(litellm_params: Mapping[str, object] | None) -> tuple[Mapping[str, object], ...]:
"""The router's ``RetryAttemptRecord`` breadcrumbs, one per failed attempt, read from the same bucket."""
if litellm_params is None:
return ()
for metadata_key in ("litellm_metadata", "metadata"):
metadata = litellm_params.get(metadata_key)
if not isinstance(metadata, dict):
continue
records = cast("Mapping[str, object]", metadata).get("previous_models") # cast-ok: untyped router dict
if isinstance(records, (list, tuple)):
return tuple(
cast("Mapping[str, object]", record) # cast-ok: RetryAttemptRecord TypedDict, read with .get
for record in cast("Sequence[object]", records) # cast-ok: untyped router tuple
if isinstance(record, dict)
)
return ()
def _record_deployment_id(record: Mapping[str, object]) -> str | None:
deployment_id: Final = record.get("deployment_id")
return deployment_id if isinstance(deployment_id, str) else None
def _attempt_identity(record: Mapping[str, object]) -> tuple[str | None, int | None]:
"""What tells one provider call from another across the request: the deployment that served it
and the retry count within its model group. Two records with the same identity describe the
same call (the fallback handler re-logs the attempt that exhausted the previous group)."""
attempted_retries: Final = record.get("attempted_retries")
return _record_deployment_id(record), attempted_retries if isinstance(attempted_retries, int) else None
def _timed_out_earlier_records(
litellm_params: Mapping[str, object] | None, final_identity: tuple[str | None, int | None]
) -> tuple[Mapping[str, object], ...]:
"""The ``Timeout`` breadcrumbs for calls other than the final attempt, one per provider call."""
by_identity: Final = {
_attempt_identity(record): record
for record in _retry_attempt_records(litellm_params)
if record.get("exception_type") == litellm.Timeout.__name__
}
return tuple(record for identity, record in by_identity.items() if identity != final_identity)
def _billed_earlier_attempts(
model_call_details: Mapping[str, object],
litellm_params: Mapping[str, object] | None,
estimated_usage: Usage,
base_model: str | None,
deployment_id: str | None,
) -> tuple[BilledFailedAttempt, ...]:
"""Every earlier retry attempt the provider also billed. The router records one
``RetryAttemptRecord`` per failed attempt in ``metadata.previous_models``, the final attempt
included, which the failure row itself bills. The final attempt is the record with the
request's own deployment and retry count (``attempted_retries`` restarts at 0 after a fallback
hop, so the count alone cannot tell the groups apart). An earlier attempt that ended in a
``Timeout`` had the same prompt accepted by its deployment, so it is billed like the final
attempt, at that deployment's input rate (else the cost map's rate for ``base_model`` or the
model). Only the exception class survives on the record, so an earlier connect timeout is
billed too; any other failure class is not."""
final_attempt: Final = _retry_metadata_bucket(litellm_params).get("attempted_retries")
final_identity: Final = (deployment_id, final_attempt if isinstance(final_attempt, int) else None)
return tuple(
BilledFailedAttempt(
attempted_retries=_attempt_number(record, index),
deployment_id=_record_deployment_id(record),
usage=estimated_usage,
response_cost=_estimated_input_cost(
model_call_details,
estimated_usage.prompt_tokens,
deployment_id=_record_deployment_id(record),
base_model=base_model,
),
)
for index, record in enumerate(_timed_out_earlier_records(litellm_params, final_identity))
)
@ -1085,6 +1200,25 @@ def _deployment_cost_per_token(deployment_id: str | None, model: str) -> CostPer
)
def _priced_input_cost(
pricing_model: str,
custom_llm_provider: str | None,
prompt_tokens: int,
custom_cost_per_token: CostPerToken | None,
) -> float | None:
try:
estimated_input_cost, _ = litellm.cost_per_token(
model=pricing_model,
custom_llm_provider=custom_llm_provider,
prompt_tokens=prompt_tokens,
custom_cost_per_token=custom_cost_per_token,
)
except Exception as exc: # noqa: BLE001 # cost_per_token raises BadRequestError for an unmapped model and the row must still be written
verbose_proxy_logger.debug("failure spend estimate: could not price %s: %s", pricing_model, exc)
return None
return estimated_input_cost
def _estimated_input_cost(
model_call_details: Mapping[str, object],
prompt_tokens: int,
@ -1092,23 +1226,19 @@ def _estimated_input_cost(
base_model: str | None,
) -> float:
"""Input price for the estimated prompt: the deployment's registered rate when it has one, else
the cost map's rate for ``base_model`` or the deployment model (an Azure deployment name is not
in the cost map, its ``base_model`` is). An unpriceable model keeps the row at $0 rather than
dropping the failure row."""
try:
model: Final = str(model_call_details.get("model") or "")
estimated_input_cost, _ = litellm.cost_per_token(
model=base_model or model,
custom_llm_provider=str(model_call_details.get("custom_llm_provider") or "") or None,
prompt_tokens=prompt_tokens,
custom_cost_per_token=_deployment_cost_per_token(deployment_id, model),
)
except Exception as exc: # noqa: BLE001 # cost_per_token raises BadRequestError for an unmapped model and the row must still be written
verbose_proxy_logger.debug(
"failure spend estimate: could not price %s: %s", model_call_details.get("model"), exc
)
return 0.0
return estimated_input_cost
the cost map's rate for ``base_model``, else the deployment model, the same candidate order
``completion_cost`` walks (an Azure deployment name is not in the cost map, its ``base_model``
is; a ``base_model`` the cost map does not know falls through to the deployment model). A model
no candidate can price keeps the row at $0 rather than dropping the failure row."""
model: Final = str(model_call_details.get("model") or "")
custom_llm_provider: Final = str(model_call_details.get("custom_llm_provider") or "") or None
custom_cost_per_token: Final = _deployment_cost_per_token(deployment_id, model)
pricing_models: Final = (base_model, model) if base_model is not None else (model,)
candidate_costs: Final = (
_priced_input_cost(pricing_model, custom_llm_provider, prompt_tokens, custom_cost_per_token)
for pricing_model in pricing_models
)
return next((cost for cost in candidate_costs if cost is not None), 0.0)
_EMPTY_LIFT: Final = MappingProxyType({})
@ -1265,13 +1395,26 @@ def _failure_fields_to_lift(request_data: Mapping[str, object], provider_billed:
return _EMPTY_LIFT
_model_call_details: Final = getattr(_logging_obj, "model_call_details", {})
_first_handoff: Final = _model_call_details.get("first_api_call_start_time")
_litellm_params: Final = getattr(_logging_obj, "litellm_params", None)
_base_model: Final = get_base_model_from_metadata(model_call_details=_model_call_details)
_deployment_id: Final = _logging_obj.get_router_model_id() if isinstance(_logging_obj, Logging) else None
_usage_to_lift: Final = _failure_usage_to_lift(
model_call_details=_model_call_details,
request_body=request_data,
dispatched=_first_handoff is not None,
provider_billed=provider_billed,
deployment_id=_logging_obj.get_router_model_id() if isinstance(_logging_obj, Logging) else None,
base_model=get_base_model_from_metadata(model_call_details=_model_call_details),
deployment_id=_deployment_id,
base_model=_base_model,
)
_prompt_estimate: Final = (
_estimated_prompt_usage(_model_call_details, request_data)
if provider_billed and _first_handoff is not None and not _model_call_details.get("failed_attempts_usage")
else None
)
_earlier_attempts: Final = (
_billed_earlier_attempts(_model_call_details, _litellm_params, _prompt_estimate, _base_model, _deployment_id)
if _prompt_estimate is not None
else ()
)
_entries: Final = (
("start_time", _model_call_details.get("start_time")),
@ -1280,6 +1423,7 @@ def _failure_fields_to_lift(request_data: Mapping[str, object], provider_billed:
("combined_usage_object", None if _usage_to_lift is None else _usage_to_lift[0]),
("response_cost", None if _usage_to_lift is None else (_usage_to_lift[1] or 0.0)),
("standard_logging_object", _model_call_details.get("standard_logging_object")),
("billed_failed_attempts", _earlier_attempts or None),
)
return MappingProxyType({key: value for key, value in _entries if value is not None})
@ -8375,7 +8519,11 @@ MAX_SPEND_LOG_ISOLATION_FAILURES_PER_BATCH: Final = 256
def _is_transient_spend_log_write_error(e: Exception) -> bool:
return PrismaDBExceptionHandler.is_database_transport_error(e) or PrismaDBExceptionHandler.is_deadlock_error(e)
return (
PrismaDBExceptionHandler.is_database_transport_error(e)
or PrismaDBExceptionHandler.is_prisma_engine_internal_error(e)
or PrismaDBExceptionHandler.is_deadlock_error(e)
)
async def _create_spend_logs_with_poison_isolation(

6
litellm/responses/file_search/emulated_handler.py
View file

@ -45,7 +45,7 @@ class FileSearchToolCallArgs(TypedDict):
def should_use_emulated_file_search(
tools: Iterable[ToolParam] | None,
tools: Iterable[object] | None,
provider_config: "BaseResponsesAPIConfig | None",
) -> bool:
"""Return True when there is a file_search tool and the provider can't handle it natively."""
@ -105,7 +105,7 @@ def _build_function_tool(vector_store_ids: list[str]) -> dict[str, object]:
}
def _replace_file_search_tools(
def replace_file_search_tools(
tools: Iterable[ToolParam] | None,
) -> tuple[list[object], list[str]]:
"""
@ -529,7 +529,7 @@ async def aresponses_with_emulated_file_search(
_include_search_results, call_kwargs = _prepare_emulated_file_search_call(kwargs=kwargs)
# 1. Replace file_search tools with function tool
transformed_tools, all_vs_ids = _replace_file_search_tools(tools)
transformed_tools, all_vs_ids = replace_file_search_tools(tools)
# 2. First provider call — provider will call the file_search function.
# Mark as an internal sub-call so wrapper_async skips billing callbacks;

207
litellm/responses/litellm_completion_transformation/count_request.py Normal file
View file

@ -0,0 +1,207 @@
from collections.abc import Mapping
from dataclasses import dataclass
from typing import Final, cast # noqa: TID251 # input items are forwarded as sent, as inference does
from pydantic import BaseModel, JsonValue, TypeAdapter, ValidationError
from litellm.exceptions import BadRequestError
from litellm.litellm_core_utils.core_helpers import as_str_mapping
from litellm.litellm_core_utils.prompt_templates.common_utils import (
update_responses_input_with_model_file_ids,
update_responses_tools_with_model_file_ids,
)
from litellm.litellm_core_utils.tool_fields import DropPolicy
from litellm.llms.base_llm.files.batch_records import responses_batch_body_to_chat_body
from litellm.responses.additional_tools import hoist_additional_tools
from litellm.responses.file_search.emulated_handler import replace_file_search_tools
from litellm.responses.litellm_completion_transformation.transformation import LiteLLMCompletionResponsesConfig
from litellm.responses.main import previous_response_client_ids
from litellm.types.llms.openai import ResponseInputParam
from litellm.utils import ProviderConfigManager
_COUNTED_CHAT_PARAMS: Final = (
"tool_choice",
"parallel_tool_calls",
"response_format",
"web_search_options",
"reasoning_effort",
)
_ITEMS: Final[TypeAdapter[tuple[object, ...]]] = TypeAdapter(tuple[object, ...])
_MAPPINGS: Final[TypeAdapter[tuple[Mapping[str, object], ...]]] = TypeAdapter(tuple[Mapping[str, object], ...])
_JSON_PARAMS: Final[TypeAdapter[Mapping[str, JsonValue]]] = TypeAdapter(Mapping[str, JsonValue])
_ENTRY: Final[TypeAdapter[Mapping[str, object]]] = TypeAdapter(Mapping[str, object])
_MARKED_ENTRY_LISTS: Final = frozenset({"content", "tool_calls"})
_FILE_ID_MAPPING: Final[TypeAdapter[Mapping[str, Mapping[str, str]] | None]] = TypeAdapter(
Mapping[str, Mapping[str, str]] | None
)
MODEL_FILE_ID_MAPPING_FIELD: Final = "model_file_id_mapping"
@dataclass(frozen=True, slots=True)
class ChatCountRequest:
messages: tuple[Mapping[str, object], ...]
tools: tuple[Mapping[str, object], ...] | None
params: Mapping[str, JsonValue]
@dataclass(frozen=True, slots=True)
class ChatCountBridgeError:
message: str
def _plain(item: object) -> object:
return item.model_dump(exclude_none=True) if isinstance(item, BaseModel) else item
def _mappings(value: object) -> tuple[Mapping[str, object], ...]:
return _MAPPINGS.validate_python(tuple(_plain(item) for item in _ITEMS.validate_python(value or ())))
def _unmarked(entry: Mapping[str, object]) -> Mapping[str, object]:
return {
key: _unmarked_entries(value) if key in _MARKED_ENTRY_LISTS else value
for key, value in entry.items()
if key != "cache_control"
}
def _unmarked_entries(value: object) -> object:
return (
[
_unmarked(_ENTRY.validate_python(item)) if isinstance(item, Mapping) else item
for item in _ITEMS.validate_python(value)
]
if isinstance(value, list)
else value
)
def _routed_model_id(deployment: Mapping[str, object] | None) -> str | None:
model_info: Final = as_str_mapping(None if deployment is None else deployment.get("model_info")) or {}
model_id: Final = model_info.get("id")
return model_id if isinstance(model_id, str) else None
def with_provider_file_ids(
responses_request: Mapping[str, object], deployment: Mapping[str, object] | None
) -> Mapping[str, object]:
mapping: Final = _FILE_ID_MAPPING.validate_python(responses_request.get(MODEL_FILE_ID_MAPPING_FIELD))
model_id: Final = _routed_model_id(deployment)
fields: Final = {name: value for name, value in responses_request.items() if name != MODEL_FILE_ID_MAPPING_FIELD}
tools: Final = fields.get("tools")
return {
**fields,
**(
{"input": update_responses_input_with_model_file_ids(fields["input"], model_id, mapping)}
if "input" in fields
else {}
),
**(
{
"tools": update_responses_tools_with_model_file_ids(
cast(list[dict[str, object]], tools), # cast-ok: entries that are not tools pass through as sent
model_id,
mapping,
)
}
if isinstance(tools, list)
else {}
),
}
def _with_hoisted_additional_tools(responses_request: Mapping[str, object]) -> Mapping[str, object]:
responses_input: Final = responses_request.get("input")
if not isinstance(responses_input, list):
return responses_request
input_items: Final = cast(ResponseInputParam, responses_input) # cast-ok: forwarded as sent, as inference does
hoisted: Final = hoist_additional_tools(input_items, None)
if not hoisted.hoisted:
return responses_request
return {
**responses_request,
"input": hoisted.input,
"tools": [*_ITEMS.validate_python(responses_request.get("tools") or ()), *hoisted.hoisted],
}
def _with_emulated_file_search(responses_request: Mapping[str, object]) -> Mapping[str, object]:
tools: Final = responses_request.get("tools")
if tools is None:
return responses_request
return {**responses_request, "tools": replace_file_search_tools(_ITEMS.validate_python(tools))[0]}
async def _with_previous_turns(
chat_body: Mapping[str, object],
model: str,
custom_llm_provider: str,
responses_request: Mapping[str, object],
policy: DropPolicy,
) -> Mapping[str, object]:
previous_response_id: Final = responses_request.get("previous_response_id")
if not isinstance(previous_response_id, str) or not previous_response_id:
return chat_body
instructions: Final = responses_request.get("instructions")
return await LiteLLMCompletionResponsesConfig.async_responses_api_session_handler( # pyright: ignore[reportUnknownMemberType] # the session handler takes a bare dict request
previous_response_id=previous_response_id,
litellm_completion_request={**chat_body, "model": model, "custom_llm_provider": custom_llm_provider},
instructions=instructions if isinstance(instructions, str) else None,
request_kwargs={"drop_params": policy.drop_params},
previous_response_client_ids=previous_response_client_ids(previous_response_id, responses_request),
)
def _counted_params(
chat_body: Mapping[str, object], tools: tuple[Mapping[str, object], ...], policy: DropPolicy
) -> Mapping[str, JsonValue]:
sent: Final = {
name: chat_body[name]
for name in _COUNTED_CHAT_PARAMS
if chat_body.get(name) is not None and name not in policy.additional_drop_params
}
return _JSON_PARAMS.validate_python(
{**sent, **({"drop_params": True} if policy.drop_params and (sent or tools) else {})}
)
async def chat_count_request(
model: str,
custom_llm_provider: str | None,
api_base: str | None,
responses_request: Mapping[str, object],
policy: DropPolicy,
) -> ChatCountRequest | ChatCountBridgeError | None:
"""The chat request Responses inference sends a provider with no native Responses API, or None
when inference sends this provider the Responses body itself."""
if custom_llm_provider is None or (
ProviderConfigManager.get_provider_responses_api_config(
provider=custom_llm_provider, model=model, api_base=api_base
)
is not None
):
return None
try:
chat_body: Final = await _with_previous_turns(
responses_batch_body_to_chat_body(
_with_emulated_file_search(_with_hoisted_additional_tools(responses_request)),
custom_llm_provider,
policy.drop_params,
),
model,
custom_llm_provider,
responses_request,
policy,
)
tools: Final = (
() if "tools" in policy.additional_drop_params else tuple(map(_unmarked, _mappings(chat_body.get("tools"))))
)
return ChatCountRequest(
messages=tuple(map(_unmarked, _mappings(chat_body.get("messages")))),
tools=tools or None,
params=_counted_params(chat_body, tools, policy),
)
except BadRequestError as e:
return ChatCountBridgeError(message=str(e.message))
except (ValidationError, ValueError, KeyError) as e:
return ChatCountBridgeError(message=f"Invalid request for token counting: {e}")

2
litellm/responses/litellm_completion_transformation/transformation.py
View file

@ -1005,7 +1005,7 @@ class LiteLLMCompletionResponsesConfig:
session_lookup: Callable[[str], Awaitable[ChatCompletionSession]] | None = None,
request_kwargs: Mapping[str, object] = MappingProxyType({}),
previous_response_client_ids: tuple[str, ...] = (),
) -> dict:
) -> dict[str, object]:
"""
Async hook to get the chain of previous input and output pairs and return a list of Chat Completion messages
"""

4
litellm/responses/main.py
View file

@ -548,7 +548,7 @@ def _bridges_to_chat_completions(
_RESPONSES_ONLY_REQUEST_FIELDS_NEVER_BRIDGED: Final = frozenset({"client_metadata"})
def _previous_response_client_ids(previous_response_id: str | None, kwargs: Mapping[str, object]) -> tuple[str, ...]:
def previous_response_client_ids(previous_response_id: str | None, kwargs: Mapping[str, object]) -> tuple[str, ...]:
"""The forms of the previous response's id the client held, which its spend log row can be stored under."""
if not previous_response_id:
return ()
@ -1622,7 +1622,7 @@ def responses(
if _bridges_to_chat_completions(responses_api_provider_config, use_chat_completions_api):
bridge_kwargs: Final = _bridge_kwargs(kwargs, responses_api_provider_config, allowed_openai_params)
return litellm_completion_transformation_handler.response_api_handler(
_previous_response_client_ids(previous_response_id, kwargs),
previous_response_client_ids(previous_response_id, kwargs),
model=model,
input=input,
responses_api_request=response_api_optional_params,

116
litellm/router.py
View file

@ -243,6 +243,15 @@ from litellm.router_utils.cooldown_handlers import (
is_video_lookup_not_found,
set_cooldown_deployments,
)
from litellm.router_utils.embedding_batch_size_check import (
async_log_rejection_before_provider_call,
clear_rejected_before_provider_call,
deployments_admitting_embedding_batch,
deployments_when_none_admits_embedding_batch,
embedding_batch_too_large_error,
embedding_deployment_selection,
log_rejection_before_provider_call,
)
from litellm.router_utils.fallback_event_handlers import (
MID_STREAM_FALLBACK_CONTROLS_KEY,
AttemptedFallbackTargets,
@ -6962,19 +6971,33 @@ class Router:
except Exception as e:
raise e
def _embedding(self, input: str | list, model: str, **kwargs):
def _embedding(self, input: str | list[object], model: str, **kwargs):
model_name = None
try:
verbose_router_logger.debug("Inside embedding()- model: %s; kwargs: %s", model, kwargs)
deployment: Final = self.get_available_deployment(
model=model,
input=input,
specific_deployment=kwargs.pop("specific_deployment", None),
request_kwargs=kwargs,
)
with embedding_deployment_selection():
deployment: Final = self.get_available_deployment(
model=model,
input=input,
specific_deployment=kwargs.pop("specific_deployment", None),
request_kwargs=kwargs,
)
self._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs)
data: Final = deployment["litellm_params"].copy()
model_name = data["model"]
batch_error: Final = embedding_batch_too_large_error(
deployment=_OPAQUE_VALUE_ADAPTER.validate_python(deployment),
input=input,
request_kwargs=_OPAQUE_VALUE_ADAPTER.validate_python(kwargs),
model=model,
)
logging_obj: Final = _REQUEST_KWARGS_ADAPTER.validate_python(kwargs).get("litellm_logging_obj")
if batch_error is not None:
if isinstance(logging_obj, LiteLLMLogging):
log_rejection_before_provider_call(logging_obj, batch_error)
raise batch_error
if isinstance(logging_obj, LiteLLMLogging):
clear_rejected_before_provider_call(logging_obj)
potential_model_client: Final = self._get_client(deployment=deployment, kwargs=kwargs, client_type="sync")
# check if provided keys == client keys #
@ -7037,20 +7060,34 @@ class Router:
)
raise e
async def _aembedding(self, input: str | list, model: str, **kwargs):
async def _aembedding(self, input: str | list[object], model: str, **kwargs):
model_name = None
try:
verbose_router_logger.debug("Inside _aembedding()- model: %s; kwargs: %s", model, kwargs)
parent_otel_span: Final = get_parent_otel_span_from_kwargs(kwargs)
deployment: Final = await self.async_get_available_deployment(
model=model,
input=input,
specific_deployment=kwargs.pop("specific_deployment", None),
request_kwargs=kwargs,
)
with embedding_deployment_selection():
deployment: Final = await self.async_get_available_deployment(
model=model,
input=input,
specific_deployment=kwargs.pop("specific_deployment", None),
request_kwargs=kwargs,
)
self._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs)
data: Final = deployment["litellm_params"].copy()
model_name = data["model"]
batch_error: Final = embedding_batch_too_large_error(
deployment=_OPAQUE_VALUE_ADAPTER.validate_python(deployment),
input=input,
request_kwargs=_OPAQUE_VALUE_ADAPTER.validate_python(kwargs),
model=model,
)
logging_obj: Final = _REQUEST_KWARGS_ADAPTER.validate_python(kwargs).get("litellm_logging_obj")
if batch_error is not None:
if isinstance(logging_obj, LiteLLMLogging):
await async_log_rejection_before_provider_call(logging_obj, batch_error)
raise batch_error
if isinstance(logging_obj, LiteLLMLogging):
clear_rejected_before_provider_call(logging_obj)
model_client: Final = self._get_async_openai_model_client(
deployment=deployment,
kwargs=kwargs,
@ -14784,10 +14821,34 @@ class Router:
batch_candidates: Final = cast( # cast-ok: router deployments are plain dicts
list[dict[str, object]], healthy_deployments
)
embedding_batch_admitting: Final = deployments_admitting_embedding_batch(
batch_candidates,
_OPAQUE_VALUE_ADAPTER.validate_python(input),
_OPAQUE_VALUE_ADAPTER.validate_python(request_kwargs),
)
embedding_batch_routed: Final = (
embedding_batch_admitting
if embedding_batch_admitting is not None
else deployments_when_none_admits_embedding_batch(
batch_candidates,
cast(list[dict[str, object]], _pre_cooldown_deployments), # cast-ok: router deployments are plain dicts
await self._async_routable_cooled_deployment_ids(
model=model,
request_deployments=request_deployments,
cooldown_deployments=cooldown_deployments,
request_kwargs=cast(Mapping[str, object], request_kwargs), # cast-ok: str-keyed request kwargs
has_request_input=True,
),
_OPAQUE_VALUE_ADAPTER.validate_python(input),
_OPAQUE_VALUE_ADAPTER.validate_python(request_kwargs),
)
)
healthy_deployments = litellm.utils.get_order_filtered_deployments(
list(
self._narrow_to_video_deployment(
without_deployments_that_cannot_batch(batch_candidates, caller_deployments=lambda: batch_candidates)
without_deployments_that_cannot_batch(
embedding_batch_routed, caller_deployments=lambda: batch_candidates
)
)
),
target_order=_target_order,
@ -15994,11 +16055,34 @@ class Router:
batch_candidates: Final = cast( # cast-ok: router deployments are plain dicts
list[dict[str, object]], healthy_deployments
)
embedding_batch_admitting: Final = deployments_admitting_embedding_batch(
batch_candidates,
_OPAQUE_VALUE_ADAPTER.validate_python(input),
routing_kwargs,
)
embedding_batch_routed: Final = (
embedding_batch_admitting
if embedding_batch_admitting is not None
else deployments_when_none_admits_embedding_batch(
batch_candidates,
cast(list[dict[str, object]], _pre_cooldown_deployments), # cast-ok: router deployments are plain dicts
self._deployment_ids(
self._region_routable_deployments(
request_deployments,
cast(Mapping[str, object] | None, request_kwargs), # cast-ok: str-keyed request kwargs
True,
)
)
& frozenset(cooldown_deployments),
_OPAQUE_VALUE_ADAPTER.validate_python(input),
routing_kwargs,
)
)
healthy_deployments = litellm.utils.get_order_filtered_deployments(
list(
self._narrow_to_video_deployment(
without_deployments_that_cannot_batch(
batch_candidates,
embedding_batch_routed,
caller_deployments=lambda: deployments_the_caller_may_use(batch_candidates, routing_kwargs),
)
)

19
litellm/router_strategy/complexity_router/complexity_router.py
View file

@ -22,7 +22,7 @@ import json
import random
import re
import time
from collections.abc import Callable, Iterator, Mapping, Sequence
from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence
from itertools import accumulate, chain, islice, takewhile
from threading import Lock
from types import MappingProxyType
@ -30,6 +30,7 @@ from typing import TYPE_CHECKING, Any, Final, Literal, NamedTuple, cast
from pydantic import BaseModel, TypeAdapter, ValidationError, create_model
from pydantic_core import ErrorDetails
from typing_extensions import ReadOnly, TypedDict
from litellm._internal_context import with_service_target
from litellm._logging import verbose_router_logger
@ -411,9 +412,19 @@ def _append_custom_keywords(base_keywords: list[str], custom_keywords: list[str]
return [*base_keywords, *deduped_custom.values()]
def _parent_session_kwargs(request_kwargs: Mapping[str, object] | None) -> Mapping[str, Any]:
class _ParentSessionKwargs(TypedDict, total=False):
litellm_session_id: ReadOnly[object]
litellm_trace_id: ReadOnly[object]
def _parent_session_kwargs(request_kwargs: Mapping[str, object] | None) -> _ParentSessionKwargs:
kwargs: Final = request_kwargs or {}
return {k: kwargs[k] for k in ("litellm_session_id", "litellm_trace_id") if kwargs.get(k) is not None}
session_id: Final = kwargs.get("litellm_session_id")
trace_id: Final = kwargs.get("litellm_trace_id")
return _ParentSessionKwargs(
**({"litellm_session_id": session_id} if session_id is not None else {}),
**({"litellm_trace_id": trace_id} if trace_id is not None else {}),
)
def _response_cost_or_none(response: object) -> float | None:
@ -2828,7 +2839,7 @@ class ComplexityRouter(CustomLogger):
"originating_request_masked": masked_originating_request(request_kwargs),
"body": {"model": llm_config.model, **payload},
}
classify: Final = (
classify: Final[Callable[..., Awaitable[ModelResponse | ResponsesAPIResponse]]] = (
self.litellm_router_instance.aresponses
if encrypted_task is not None
else self.litellm_router_instance.acompletion

199
litellm/router_utils/embedding_batch_size_check.py Normal file
View file

@ -0,0 +1,199 @@
import traceback
from collections.abc import Generator, Iterable, Mapping, Sequence
from contextlib import contextmanager
from contextvars import ContextVar
from typing import TYPE_CHECKING, Final
from typing_extensions import TypeIs # noqa: TID251 # TypeIs reaches typing only on 3.13
import litellm
from litellm.constants import PROXY_REJECTED_BEFORE_ROUTING_KEY
from litellm.types.utils import CallTypes
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
_SELECTING_FOR_EMBEDDING: Final[ContextVar[bool]] = ContextVar("selecting_for_embedding", default=False)
def _is_str_keyed_mapping(
value: object,
) -> TypeIs[Mapping[str, object]]: # guard-ok: model_info, litellm_params and extra_body are str-keyed config objects
return isinstance(value, Mapping)
def _is_embedding_input(
value: object,
) -> TypeIs[str | Sequence[object]]: # guard-ok: only the length and first item are read
return isinstance(value, Sequence)
def count_embedding_inputs(input: str | Sequence[object]) -> int:
if isinstance(input, str):
return 1
if input and isinstance(input[0], int):
return 1
return len(input)
def effective_embedding_input(
input: object,
extra_body: object,
) -> str | Sequence[object] | None:
"""Return the input shape used for batch-size checks, or None when neither is an embedding input.
Some providers merge ``extra_body`` over the request body (OpenAI-compatible paths);
others ignore ``extra_body.input`` and use the top-level input (e.g. Bedrock Titan).
Use whichever of the two yields the larger input count so limits cannot be bypassed.
"""
override: Final = extra_body.get("input") if _is_str_keyed_mapping(extra_body) else None
candidates: Final = tuple(value for value in (input, override) if _is_embedding_input(value))
return max(candidates, key=count_embedding_inputs, default=None)
def _litellm_param(deployment: Mapping[str, object], key: str) -> object:
litellm_params: Final = deployment.get("litellm_params")
return litellm_params.get(key) if _is_str_keyed_mapping(litellm_params) else None
def _embedding_batch_limit(deployment: Mapping[str, object]) -> int | None:
model_info: Final = deployment.get("model_info")
limit: Final = model_info.get("max_embedding_batch_size") if _is_str_keyed_mapping(model_info) else None
return limit if isinstance(limit, int) else None
def _requested_embedding_count(
deployment: Mapping[str, object],
input: object,
request_kwargs: Mapping[str, object],
) -> int:
extra_body: Final = request_kwargs.get("extra_body", _litellm_param(deployment, "extra_body"))
effective_input: Final = effective_embedding_input(input, extra_body)
return 0 if effective_input is None else count_embedding_inputs(effective_input)
def _deployment_admits_embedding_batch(
deployment: Mapping[str, object],
input: object,
request_kwargs: Mapping[str, object],
) -> bool:
limit: Final = _embedding_batch_limit(deployment)
return limit is None or _requested_embedding_count(deployment, input, request_kwargs) <= limit
def _str_keyed_or_empty(value: object) -> Mapping[str, object]:
return value if _is_str_keyed_mapping(value) else {}
def _deployment_id(deployment: Mapping[str, object]) -> str | None:
model_info: Final = deployment.get("model_info")
deployment_id: Final = model_info.get("id") if _is_str_keyed_mapping(model_info) else None
return None if deployment_id is None else str(deployment_id)
def _is_paused(deployment: Mapping[str, object]) -> bool:
model_info: Final = deployment.get("model_info")
return _is_str_keyed_mapping(model_info) and model_info.get("blocked") is True
@contextmanager
def embedding_deployment_selection() -> Generator[None]:
"""Apply the input limit to deployment selection inside this block, which only embedding calls enter."""
token: Final = _SELECTING_FOR_EMBEDDING.set(True)
try:
yield
finally:
_SELECTING_FOR_EMBEDDING.reset(token)
def deployments_admitting_embedding_batch(
deployments: list[dict[str, object]], # mutable-ok: router deployment-list contract
input: object,
request_kwargs: object,
) -> list[dict[str, object]] | None: # mutable-ok: returns a mutable deployment list
"""Keep the deployments whose input limit admits the request, all of them outside embedding selection, else None"""
if not _SELECTING_FOR_EMBEDDING.get():
return deployments
kwargs: Final = _str_keyed_or_empty(request_kwargs)
admitting: Final = [
deployment for deployment in deployments if _deployment_admits_embedding_batch(deployment, input, kwargs)
]
return admitting or None
def deployments_when_none_admits_embedding_batch(
deployments: list[dict[str, object]], # mutable-ok: router deployment-list contract
before_cooldown: Iterable[Mapping[str, object]],
routable_cooling_ids: frozenset[str],
input: object,
request_kwargs: object,
) -> list[dict[str, object]]: # mutable-ok: returns a mutable deployment list
"""None while an unpaused cooling deployment the request routes to admits it, so it waits out the cooldown,
else the deployments with the largest limit, so the 400 names the most the group accepts"""
kwargs: Final = _str_keyed_or_empty(request_kwargs)
routable_cooling: Final = (
deployment
for deployment in before_cooldown
if _deployment_id(deployment) in routable_cooling_ids and not _is_paused(deployment)
)
if any(_deployment_admits_embedding_batch(deployment, input, kwargs) for deployment in routable_cooling):
return []
largest_limit: Final = max((_embedding_batch_limit(deployment) or 0 for deployment in deployments), default=0)
return [deployment for deployment in deployments if _embedding_batch_limit(deployment) == largest_limit]
def _mark_litellm_params(logging_obj: "LiteLLMLoggingObj", *keys: str) -> None:
litellm_params: Final = logging_obj.model_call_details.get("litellm_params")
if not isinstance(litellm_params, dict):
return
for key in keys:
litellm_params[key] = True
def log_rejection_before_provider_call(logging_obj: "LiteLLMLoggingObj", error: Exception) -> None:
"""Mark the request as never reaching a provider, so deployment health metrics skip its failure, then log it"""
_mark_litellm_params(logging_obj, PROXY_REJECTED_BEFORE_ROUTING_KEY)
logging_obj.failure_handler(error, "".join(traceback.format_exception(error)))
async def async_log_rejection_before_provider_call(logging_obj: "LiteLLMLoggingObj", error: Exception) -> None:
"""Log it like an async provider failure: sync handlers inline, before a fallback swaps the deployment they read,
and the call marked async, so a logger with both hooks records it once"""
_mark_litellm_params(logging_obj, PROXY_REJECTED_BEFORE_ROUTING_KEY, CallTypes.aembedding.value)
traceback_text: Final = "".join(traceback.format_exception(error))
logging_obj.failure_handler(error, traceback_text)
await logging_obj.async_failure_handler(error, traceback_text)
def clear_rejected_before_provider_call(logging_obj: "LiteLLMLoggingObj") -> None:
"""Drop the mark an earlier rejected attempt left, since a fallback or retry shares the logging object"""
litellm_params: Final = logging_obj.model_call_details.get("litellm_params")
if isinstance(litellm_params, dict) and PROXY_REJECTED_BEFORE_ROUTING_KEY in litellm_params:
del litellm_params[PROXY_REJECTED_BEFORE_ROUTING_KEY]
def embedding_batch_too_large_error(
deployment: object,
input: object,
request_kwargs: object,
model: str,
) -> litellm.BadRequestError | None:
if not _is_str_keyed_mapping(deployment):
return None
limit: Final = _embedding_batch_limit(deployment)
if limit is None:
return None
input_count: Final = _requested_embedding_count(deployment, input, _str_keyed_or_empty(request_kwargs))
if input_count <= limit:
return None
return litellm.BadRequestError(
message=(
f"Embedding request has {input_count} inputs, which exceeds "
f"max_embedding_batch_size={limit} for model={model}. "
"Split the input into smaller batches."
),
model=model,
llm_provider="",
)

6
litellm/router_utils/prompt_caching_cache.py
View file

@ -94,8 +94,10 @@ def _is_ephemeral_mark(mark: object) -> bool:
return mark_map.get("type") == "ephemeral"
def _tool_carries_ephemeral_mark(tool: ChatCompletionToolParam) -> bool:
tool_map: Final = cast( # cast-ok: TypedDict.get rejects keys the type does not declare
def _tool_carries_ephemeral_mark(tool: object) -> bool:
if not isinstance(tool, dict):
return False
tool_map: Final = cast( # cast-ok: tools are plain JSON dicts
Mapping[str, object], tool
)
if _is_ephemeral_mark(tool_map.get("cache_control")):

1
litellm/types/integrations/langfuse_otel.py
View file

@ -29,6 +29,7 @@ class LangfuseSpanAttributes(str, Enum):
# ---- Observation input/output ----
OBSERVATION_INPUT = "langfuse.observation.input"
OBSERVATION_OUTPUT = "langfuse.observation.output"
OBSERVATION_COMPLETION_START_TIME = "langfuse.observation.completion_start_time"
# ---- Trace-level metadata ----
TRACE_USER_ID = "user.id"

2
litellm/types/router.py
View file

@ -323,6 +323,8 @@ class ModelInfo(MirroredPricingParams):
# in the spend log row's metadata. Set it on every deployment of the group.
internal_router_model: bool | None = None
max_embedding_batch_size: Annotated[int, Field(gt=0)] | None = None
def __init__(self, id: str | int | None = None, **params) -> None:
if id is None:
id = str(uuid.uuid4()) # Generate a UUID if id is None or not provided

3
litellm/utils.py
View file

@ -5552,8 +5552,9 @@ def add_provider_specific_params_to_optional_params(
# for openai, azure we should pass the extra/passed params within `extra_body` https://github.com/openai/openai-python/blob/ac33853ba10d13ac149b1fa3ca6dba7d613065c9/src/openai/resources/models.py#L46
if _should_drop_param(k="extra_body", additional_drop_params=additional_drop_params) is False:
extra_body: Final = dict(passed_params.pop("extra_body", None) or {})
withheld_params: Final = frozenset[str]() if custom_llm_provider == "litellm_proxy" else LITELLM_ONLY_PARAMS
for k in passed_params:
if k not in openai_params and k not in LITELLM_ONLY_PARAMS and passed_params[k] is not None:
if k not in openai_params and k not in withheld_params and passed_params[k] is not None:
extra_body[k] = passed_params[k]
if not isinstance(optional_params.get("extra_body"), dict):
optional_params["extra_body"] = {}

269
model_prices_and_context_window.json
View file

@ -4916,6 +4916,26 @@
"supports_tool_choice": true,
"supports_vision": true
},
"azure/eu/gpt-4o": {
"deprecation_date": "2027-04-14",
"cache_creation_input_token_cost": 1.38e-06,
"cache_read_input_token_cost": 1.375e-06,
"input_cost_per_token": 2.75e-06,
"input_cost_per_token_batches": 1.375e-06,
"litellm_provider": "azure",
"max_input_tokens": 128000,
"max_output_tokens": 16384,
"max_tokens": 16384,
"mode": "chat",
"output_cost_per_token": 1.1e-05,
"output_cost_per_token_batches": 5.5e-06,
"source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'",
"supports_function_calling": true,
"supports_parallel_function_calling": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true
},
"azure/eu/gpt-4o-2024-11-20": {
"deprecation_date": "2027-04-14",
"cache_creation_input_token_cost": 1.38e-06,
@ -4936,6 +4956,25 @@
"supports_tool_choice": true,
"supports_vision": true
},
"azure/eu/gpt-4o-mini": {
"cache_read_input_token_cost": 8.3e-08,
"deprecation_date": "2027-04-14",
"input_cost_per_token": 1.65e-07,
"input_cost_per_token_batches": 8.3e-08,
"litellm_provider": "azure",
"max_input_tokens": 128000,
"max_output_tokens": 16384,
"max_tokens": 16384,
"mode": "chat",
"output_cost_per_token": 6.6e-07,
"source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'",
"supports_function_calling": true,
"supports_parallel_function_calling": true,
"supports_prompt_caching": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true
},
"azure/eu/gpt-4o-mini-2024-07-18": {
"cache_read_input_token_cost": 8.3e-08,
"deprecation_date": "2027-04-14",
@ -5247,6 +5286,25 @@
"supports_vision": true,
"cache_read_input_token_cost_batches": 2.75e-09
},
"azure/eu/o1": {
"cache_read_input_token_cost": 8.25e-06,
"deprecation_date": "2026-11-19",
"input_cost_per_token": 1.65e-05,
"input_cost_per_token_batches": 8.25e-06,
"litellm_provider": "azure",
"max_input_tokens": 200000,
"max_output_tokens": 100000,
"max_tokens": 100000,
"mode": "chat",
"output_cost_per_token": 6.6e-05,
"output_cost_per_token_batches": 3.3e-05,
"source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'",
"supports_function_calling": true,
"supports_parallel_function_calling": true,
"supports_prompt_caching": true,
"supports_tool_choice": true,
"supports_vision": true
},
"azure/eu/o1-2024-12-17": {
"cache_read_input_token_cost": 8.25e-06,
"deprecation_date": "2026-11-19",
@ -5283,6 +5341,24 @@
"supports_prompt_caching": true,
"supports_vision": false
},
"azure/eu/o3-mini": {
"cache_read_input_token_cost": 6.05e-07,
"deprecation_date": "2026-11-19",
"input_cost_per_token": 1.21e-06,
"input_cost_per_token_batches": 6.05e-07,
"litellm_provider": "azure",
"max_input_tokens": 200000,
"max_output_tokens": 100000,
"max_tokens": 100000,
"mode": "chat",
"output_cost_per_token": 4.84e-06,
"output_cost_per_token_batches": 2.42e-06,
"source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'",
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_tool_choice": true,
"supports_vision": false
},
"azure/eu/o3-mini-2025-01-31": {
"cache_read_input_token_cost": 6.05e-07,
"deprecation_date": "2026-11-19",
@ -11540,6 +11616,26 @@
"supports_tool_choice": true,
"supports_vision": true
},
"azure/us/gpt-4o": {
"deprecation_date": "2027-04-14",
"cache_creation_input_token_cost": 1.38e-06,
"cache_read_input_token_cost": 1.375e-06,
"input_cost_per_token": 2.75e-06,
"input_cost_per_token_batches": 1.375e-06,
"litellm_provider": "azure",
"max_input_tokens": 128000,
"max_output_tokens": 16384,
"max_tokens": 16384,
"mode": "chat",
"output_cost_per_token": 1.1e-05,
"output_cost_per_token_batches": 5.5e-06,
"source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'",
"supports_function_calling": true,
"supports_parallel_function_calling": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true
},
"azure/us/gpt-4o-2024-11-20": {
"deprecation_date": "2027-04-14",
"cache_creation_input_token_cost": 1.38e-06,
@ -11560,6 +11656,25 @@
"supports_tool_choice": true,
"supports_vision": true
},
"azure/us/gpt-4o-mini": {
"cache_read_input_token_cost": 8.3e-08,
"deprecation_date": "2027-04-14",
"input_cost_per_token": 1.65e-07,
"input_cost_per_token_batches": 8.3e-08,
"litellm_provider": "azure",
"max_input_tokens": 128000,
"max_output_tokens": 16384,
"max_tokens": 16384,
"mode": "chat",
"output_cost_per_token": 6.6e-07,
"source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'",
"supports_function_calling": true,
"supports_parallel_function_calling": true,
"supports_prompt_caching": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true
},
"azure/us/gpt-4o-mini-2024-07-18": {
"cache_read_input_token_cost": 8.3e-08,
"deprecation_date": "2027-04-14",
@ -11871,6 +11986,25 @@
"supports_tool_choice": true,
"supports_vision": true
},
"azure/us/o1": {
"cache_read_input_token_cost": 8.25e-06,
"deprecation_date": "2026-11-19",
"input_cost_per_token": 1.65e-05,
"input_cost_per_token_batches": 8.25e-06,
"litellm_provider": "azure",
"max_input_tokens": 200000,
"max_output_tokens": 100000,
"max_tokens": 100000,
"mode": "chat",
"output_cost_per_token": 6.6e-05,
"output_cost_per_token_batches": 3.3e-05,
"source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'",
"supports_function_calling": true,
"supports_parallel_function_calling": true,
"supports_prompt_caching": true,
"supports_tool_choice": true,
"supports_vision": true
},
"azure/us/o1-2024-12-17": {
"cache_read_input_token_cost": 8.25e-06,
"deprecation_date": "2026-11-19",
@ -11907,6 +12041,39 @@
"supports_prompt_caching": true,
"supports_vision": false
},
"azure/us/o3": {
"deprecation_date": "2026-11-19",
"cache_read_input_token_cost": 5.5e-07,
"input_cost_per_token": 2.2e-06,
"input_cost_per_token_batches": 1.1e-06,
"litellm_provider": "azure",
"max_input_tokens": 200000,
"max_output_tokens": 100000,
"max_tokens": 100000,
"mode": "chat",
"output_cost_per_token": 8.8e-06,
"output_cost_per_token_batches": 4.4e-06,
"source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'",
"supported_endpoints": [
"/v1/chat/completions",
"/v1/batch",
"/v1/responses"
],
"supported_modalities": [
"text",
"image"
],
"supported_output_modalities": [
"text"
],
"supports_function_calling": true,
"supports_parallel_function_calling": false,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true
},
"azure/us/o3-2025-04-16": {
"deprecation_date": "2026-11-19",
"cache_read_input_token_cost": 5.5e-07,
@ -11940,6 +12107,24 @@
"supports_tool_choice": true,
"supports_vision": true
},
"azure/us/o3-mini": {
"cache_read_input_token_cost": 6.05e-07,
"deprecation_date": "2026-11-19",
"input_cost_per_token": 1.21e-06,
"input_cost_per_token_batches": 6.05e-07,
"litellm_provider": "azure",
"max_input_tokens": 200000,
"max_output_tokens": 100000,
"max_tokens": 100000,
"mode": "chat",
"output_cost_per_token": 4.84e-06,
"output_cost_per_token_batches": 2.42e-06,
"source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'",
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_tool_choice": true,
"supports_vision": false
},
"azure/us/o3-mini-2025-01-31": {
"cache_read_input_token_cost": 6.05e-07,
"deprecation_date": "2026-11-19",
@ -11958,6 +12143,27 @@
"supports_tool_choice": true,
"supports_vision": false
},
"azure/us/o4-mini": {
"cache_read_input_token_cost": 3.03e-07,
"deprecation_date": "2026-11-19",
"input_cost_per_token": 1.21e-06,
"input_cost_per_token_batches": 6.05e-07,
"litellm_provider": "azure",
"max_input_tokens": 200000,
"max_output_tokens": 100000,
"max_tokens": 100000,
"mode": "chat",
"output_cost_per_token": 4.84e-06,
"output_cost_per_token_batches": 2.42e-06,
"source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'",
"supports_function_calling": true,
"supports_parallel_function_calling": false,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true
},
"azure/us/o4-mini-2025-04-16": {
"cache_read_input_token_cost": 3.03e-07,
"deprecation_date": "2026-11-19",
@ -61148,6 +61354,41 @@
"supports_tool_choice": false,
"supports_vision": true
},
"volcengine/doubao-seed-2-1-pro-260915": {
"cache_read_input_token_cost": 1.725e-07,
"input_cost_per_token": 8.625e-07,
"litellm_provider": "volcengine",
"max_input_tokens": 1024000,
"max_output_tokens": 256000,
"max_tokens": 256000,
"mode": "chat",
"output_cost_per_token": 4.3125e-06,
"source": "https://www.volcengine.com/docs/82379/1544106",
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_tool_choice": false,
"supports_vision": true
},
"volcengine/doubao-seed-2-1-lite-260915": {
"cache_read_input_audio_token_cost": 3.45e-07,
"cache_read_input_token_cost": 2.3e-08,
"input_cost_per_audio_token": 1.725e-06,
"input_cost_per_token": 1.15e-07,
"litellm_provider": "volcengine",
"max_input_tokens": 1024000,
"max_output_tokens": 256000,
"max_tokens": 256000,
"mode": "chat",
"output_cost_per_token": 3.88126e-07,
"source": "https://www.volcengine.com/docs/82379/1544106",
"supports_audio_input": true,
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_tool_choice": false,
"supports_vision": true
},
"volcengine/doubao-seed-2-1-turbo-260628": {
"cache_read_input_token_cost": 8.625e-08,
"input_cost_per_token": 4.3125e-07,
@ -70141,6 +70382,20 @@
"output_cost_per_token_batches": 2.42e-06,
"source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'"
},
"azure/eu/o3": {
"deprecation_date": "2026-11-19",
"cache_read_input_token_cost": 5.5e-07,
"input_cost_per_token": 2.2e-06,
"input_cost_per_token_batches": 1.1e-06,
"litellm_provider": "azure",
"max_input_tokens": 200000,
"max_output_tokens": 100000,
"max_tokens": 100000,
"mode": "chat",
"output_cost_per_token": 8.8e-06,
"output_cost_per_token_batches": 4.4e-06,
"source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'"
},
"azure/eu/o3-2025-04-16": {
"deprecation_date": "2026-11-19",
"cache_read_input_token_cost": 5.5e-07,
@ -70164,6 +70419,20 @@
"output_cost_per_token": 4.4e-05,
"source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'"
},
"azure/eu/o4-mini": {
"deprecation_date": "2026-11-19",
"cache_read_input_token_cost": 3.03e-07,
"input_cost_per_token": 1.21e-06,
"input_cost_per_token_batches": 6.05e-07,
"litellm_provider": "azure",
"max_input_tokens": 200000,
"max_output_tokens": 100000,
"max_tokens": 100000,
"mode": "chat",
"output_cost_per_token": 4.84e-06,
"output_cost_per_token_batches": 2.42e-06,
"source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'"
},
"azure/eu/o4-mini-2025-04-16": {
"deprecation_date": "2026-11-19",
"cache_read_input_token_cost": 3.03e-07,

3
tests/e2e/coverage_registry/llm_conversational.yaml
View file

@ -13,6 +13,8 @@
- {id: llm.chat_completions.openai.vision.nonstream.works, module: llm, tier: P0, subject_endpoint: chat_completions, route: openai, capability: vision, streaming: nonstream, assertions: [works], source: "model_prices json", rationale: "gpt-4o vision; high usage"}
- {id: llm.chat_completions.openai.prompt_cache_5m.nonstream.works, module: llm, tier: P0, subject_endpoint: chat_completions, route: openai, capability: prompt_cache_5m, streaming: nonstream, assertions: [works], source: "model_prices json", rationale: "Prompt caching cost optimization"}
- {id: llm.chat_completions.openai.service_tier.nonstream.works, module: llm, tier: P1, subject_endpoint: chat_completions, route: openai, capability: service_tier, streaming: nonstream, assertions: [works], source: "OpenAI service_tier param", rationale: "OpenAI scale-tier request option is forwarded and echoed"}
- {id: llm.chat_completions.openai.prompt_cache_24h.nonstream.forwards_retention, module: llm, tier: P1, subject_endpoint: chat_completions, route: openai, capability: prompt_cache_24h, streaming: nonstream, assertions: [forwards_retention], source: "main.py completion()", fail_before_fix: proven, rationale: "prompt_cache_retention sent through the proxy, litellm.completion or Router.completion reaches the body sent to OpenAI /v1/chat/completions, so a caller asking for 24h prompt caching gets it instead of a silent drop"}
- {id: llm.chat_completions.openai.prompt_cache_24h.nonstream.bridge_forwards_retention, module: llm, tier: P1, subject_endpoint: chat_completions, route: openai, capability: prompt_cache_24h, streaming: nonstream, assertions: [forwards_retention], source: "main.py completion()", fail_before_fix: proven, rationale: "prompt_cache_retention on an openai/responses/ model (chat completions in, Responses out) reaches the body sent to OpenAI /v1/responses"}
- {id: llm.chat_completions.openai.service_tier.stream.echoes_served_tier, module: llm, tier: P1, subject_endpoint: chat_completions, route: openai, capability: service_tier, streaming: stream, assertions: [works], source: "litellm_core_utils/streaming_handler.py", fail_before_fix: proven, rationale: "Every relayed stream chunk carries the service_tier OpenAI stamped on it, so a streaming caller can see which tier served the request"}
- {id: llm.chat_completions.openai.thinking.nonstream.works, module: llm, tier: P1, subject_endpoint: chat_completions, route: openai, capability: thinking, streaming: nonstream, assertions: [works], source: "model_prices json", rationale: "o-series reasoning; emerging"}
- {id: llm.chat_completions.openai.structured_output.nonstream.works, module: llm, tier: P1, subject_endpoint: chat_completions, route: openai, capability: structured_output, streaming: nonstream, assertions: [works], source: "model_prices json", rationale: "response_schema extraction"}
@ -68,6 +70,7 @@
- {id: llm.messages.anthropic.vision.nonstream.works, module: llm, tier: P0, subject_endpoint: messages, route: anthropic, capability: vision, streaming: nonstream, assertions: [works], source: "model_prices json", rationale: "Vision via Messages API"}
- {id: llm.messages.anthropic.prompt_cache_5m.nonstream.works, module: llm, tier: P0, subject_endpoint: messages, route: anthropic, capability: prompt_cache_5m, streaming: nonstream, assertions: [works], source: "model_prices json", rationale: "Prompt caching via Messages API"}
- {id: llm.messages.anthropic.thinking.nonstream.works, module: llm, tier: P1, subject_endpoint: messages, route: anthropic, capability: thinking, streaming: nonstream, assertions: [works], source: "model_prices json", rationale: "Extended thinking via Messages API"}
- {id: llm.messages.bedrock_converse.basic.nonstream.works_for_kimi_k3_max_tokens_1, module: llm, tier: P0, subject_endpoint: messages, route: bedrock_converse, capability: basic, streaming: nonstream, assertions: [works_for_kimi_k3_max_tokens_1], source: "llm_translation/test_bedrock_kimi_k3_messages_e2e.py", rationale: "Bedrock Converse accepts max_tokens=1 on the Anthropic Messages surface after LiteLLM applies the model minimum", fail_before_fix: proven}
- {id: llm.messages.bedrock_invoke.mid_conversation_system.nonstream.cache_hit, module: llm, tier: P0, subject_endpoint: messages, route: bedrock_invoke, capability: mid_conversation_system, streaming: nonstream, assertions: [works, cache_hit], source: "llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py", rationale: "Flagged Claude 4.8+/5 must keep mid-conversation system reminders in messages; hoisting mutates the system prefix and collapses the prompt cache (#32578/#32831/#32882)", fail_before_fix: proven}
- {id: llm.messages.bedrock_invoke.mid_conversation_system.nonstream.works, module: llm, tier: P0, subject_endpoint: messages, route: bedrock_invoke, capability: mid_conversation_system, streaming: nonstream, assertions: [works], source: "llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py", rationale: "Claude <= 4.7 rejects role system inside messages; unflagged models must convert reminders to user turns in place (hoisting collapses the prompt cache) or every Claude Code session 400s (#32831)", fail_before_fix: proven}
- {id: llm.messages.bedrock_invoke.web_search_server_tool.nonstream.works, module: llm, tier: P0, subject_endpoint: messages, route: bedrock_invoke, capability: web_search_server_tool, streaming: nonstream, assertions: [works], source: "llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py", rationale: "Bedrock hosts no web_search server tool, so this only works because interception rewrites it before the upstream call and the agentic loop feeds the results back in native shape; a regression that short-circuits or forwards it instead yields raw text or AWS's 400", fail_before_fix: unproven}

1
tests/e2e/coverage_registry/schema.py
View file

@ -81,6 +81,7 @@ LlmCapability = Literal[
"pdf_input",
"logprobs",
"prompt_cache_1h",
"prompt_cache_24h",
"prompt_cache_5m",
"response_headers",
"service_tier",

71
tests/e2e/llm_translation/test_bedrock_kimi_k3_messages_e2e.py Normal file
View file

@ -0,0 +1,71 @@
from __future__ import annotations
from typing import Final
import pytest
from anthropic.types import Message, MessageParam
from e2e_metadata import Domain, Mode, Provider, Route, Subject, meta
from lifecycle import ResourceManager
from models import LiteLLMParamsBody
from proxy_client import ProxyClient
from sdk_clients import NO_PROXY_CACHE, SdkClients
pytestmark = pytest.mark.e2e
BEDROCK_MODEL: Final = "bedrock/us.moonshotai.kimi-k3"
MODEL_NAME: Final = "claude-kimi-k3"
MODEL_SWITCH_BETA: Final = (
"claude-code-20250219,context-1m-2025-08-07,interleaved-thinking-2025-05-14,"
"redact-thinking-2026-02-12,context-management-2025-06-27,"
"prompt-caching-scope-2026-01-05,mid-conversation-system-2026-04-07"
)
def _register_model(proxy: ProxyClient, resources: ResourceManager) -> tuple[str, str]:
model_id: Final = proxy.create_model(
MODEL_NAME,
LiteLLMParamsBody(
model=BEDROCK_MODEL,
aws_access_key_id="os.environ/AWS_ACCESS_KEY_ID",
aws_secret_access_key="os.environ/AWS_SECRET_ACCESS_KEY",
aws_region_name="os.environ/AWS_REGION",
drop_params=True,
),
)
resources.defer(lambda: proxy.delete_model(model_id))
return MODEL_NAME, resources.key(models=[MODEL_NAME])
class TestBedrockKimiK3Messages:
@pytest.mark.covers("llm.messages.bedrock_converse.basic.nonstream.works_for_kimi_k3_max_tokens_1")
@meta(
Subject(
domain=Domain.LLM_TRANSLATION,
route=Route.MESSAGES,
providers=(Provider.BEDROCK,),
models=(BEDROCK_MODEL,),
mode=Mode.NONSTREAM,
)
)
def test_model_switch_probe_max_tokens_1_returns_message(
self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients
) -> None:
model, key = _register_model(proxy, resources)
client: Final = sdk.anthropic(key)
message_param: Final[MessageParam] = {
"role": "user",
"content": [{"type": "text", "text": "Hi", "cache_control": {"type": "ephemeral"}}],
}
message: Final = client.messages.create(
model=model,
max_tokens=1,
stream=False,
system="You are Claude Code, an assistant for software development.",
messages=[message_param],
extra_headers={"anthropic-beta": MODEL_SWITCH_BETA},
extra_query={"beta": "true"},
extra_body=NO_PROXY_CACHE,
)
assert isinstance(message, Message)
assert message.role == "assistant"

168
tests/e2e/llm_translation/test_provider_features_e2e.py
View file

@ -1,4 +1,4 @@
"""Live e2e for model-specific request features: service_tier.
"""Live e2e for model-specific request features: service_tier and prompt_cache_retention.
Each case asserts the feature took effect, not just a 200.
@ -11,26 +11,38 @@ avoided here because it is capacity-constrained and returns a transient 429 when
flex resources are unavailable. Bedrock and Vertex do not accept service_tier, so
that cell is OpenAI-only by design.
Prompt caching lives in test_cache_control.py.
prompt_cache_retention is judged on the body OpenAI receives, captured by a live
provider edge, not on the reply: OpenAI echoes "24h" on gpt-5.5 even when the
request never carried the param, so an echo would pass with the param dropped.
Prompt caching hits and misses live in test_cache_control.py.
"""
from __future__ import annotations
import queue
from collections.abc import Mapping
from dataclasses import dataclass, field
from typing import Final
import pytest
from pydantic import BaseModel
from e2e_config import unique_marker
from e2e_config import PROVIDER_EDGE_ADVERTISE_HOST, PROVIDER_EDGE_BIND_HOST, unique_marker
from e2e_http import unwrap
from e2e_metadata import Domain, Mode, Provider, Route, Subject, meta
from lifecycle import ResourceManager
from models import ChatBody, ChatMessage, LiteLLMParamsBody
from passthrough_client import PassthroughClient
from provider_edge import LiveEdge, start_provider_edge
from proxy_client import ProxyClient
pytestmark = pytest.mark.e2e
SERVICE_TIER = "priority"
OPENAI_BACKEND: Final = "openai/gpt-5.5"
OPENAI_RESPONSES_BRIDGE_BACKEND: Final = "openai/responses/gpt-5.5"
PROMPT_CACHE_RETENTION: Final = "24h"
class TestServiceTier:
@ -74,3 +86,153 @@ class TestServiceTier:
f"service_tier not honored: sent {SERVICE_TIER!r}, response reported "
f"{response.service_tier!r} ({response})"
)
class _OutboundRetention(BaseModel, frozen=True):
prompt_cache_retention: str | None = None
@dataclass(frozen=True, slots=True)
class _OpenAIBodyCapture:
"""The raw bodies sent to one OpenAI path, seen by a live edge in front of
api.openai.com. Pass ``observe`` as the edge's ``observe_request``."""
path: str
_bodies: queue.SimpleQueue[bytes] = field(default_factory=queue.SimpleQueue)
def observe(self, url: str, headers: Mapping[str, str], body: bytes | None) -> None:
if body is not None and url.split("?")[0].endswith(self.path):
self._bodies.put(body)
def retentions(self) -> tuple[str | None, ...]:
bodies: Final = tuple(self._bodies.get_nowait() for _ in range(self._bodies.qsize()))
return tuple(_OutboundRetention.model_validate_json(body).prompt_cache_retention for body in bodies)
def _openai_edge_api_base(resources: ResourceManager, capture: _OpenAIBodyCapture, *, reachable_from_proxy: bool) -> str:
"""The proxy may run in a container, so its edge binds where the proxy can
reach it; an in-process SDK call reaches the edge on loopback."""
edge: Final = (
start_provider_edge(
LiveEdge(observe_request=capture.observe),
bind_host=PROVIDER_EDGE_BIND_HOST,
advertise_host=PROVIDER_EDGE_ADVERTISE_HOST,
)
if reachable_from_proxy
else start_provider_edge(LiveEdge(observe_request=capture.observe))
)
resources.defer(edge.shutdown)
return f"{edge.edge.api_base('openai')}/v1"
def _retention_backends() -> tuple[object, ...]:
"""gpt-5.5 straight to /v1/chat/completions, and the openai/responses/ bridge
(chat completions in, Responses out) to /v1/responses."""
return (
pytest.param(
OPENAI_BACKEND,
"/chat/completions",
marks=(
pytest.mark.covers("llm.chat_completions.openai.prompt_cache_24h.nonstream.forwards_retention"),
meta(
Subject(
domain=Domain.LLM_TRANSLATION,
route=Route.CHAT_COMPLETIONS,
providers=(Provider.OPENAI,),
models=(OPENAI_BACKEND,),
mode=Mode.NONSTREAM,
)
),
),
id="chat_completions",
),
pytest.param(
OPENAI_RESPONSES_BRIDGE_BACKEND,
"/responses",
marks=(
pytest.mark.covers("llm.chat_completions.openai.prompt_cache_24h.nonstream.bridge_forwards_retention"),
meta(
Subject(
domain=Domain.LLM_TRANSLATION,
route=Route.CHAT_COMPLETIONS,
providers=(Provider.OPENAI,),
models=(OPENAI_RESPONSES_BRIDGE_BACKEND,),
mode=Mode.NONSTREAM,
)
),
),
id="responses_bridge",
),
)
class TestPromptCacheRetention:
@pytest.mark.provider_edge_host
@pytest.mark.parametrize(("backend", "openai_path"), _retention_backends())
def test_proxy_chat_completions_forwards_prompt_cache_retention_to_openai(
self, proxy: ProxyClient, resources: ResourceManager, backend: str, openai_path: str
) -> None:
capture: Final = _OpenAIBodyCapture(openai_path)
api_base: Final = _openai_edge_api_base(resources, capture, reachable_from_proxy=True)
model: Final = f"e2e-cache-retention-{unique_marker()}"
model_id: Final = proxy.create_model(
model, LiteLLMParamsBody(model=backend, api_key="os.environ/OPENAI_API_KEY", api_base=api_base)
)
resources.defer(lambda: proxy.delete_model(model_id))
unwrap(
proxy.chat(
resources.key(),
ChatBody(
model=model,
messages=[ChatMessage(role="user", content="reply with one word")],
prompt_cache_retention=PROMPT_CACHE_RETENTION,
),
)
)
assert capture.retentions() == (PROMPT_CACHE_RETENTION,), (
f"one {openai_path} call to OpenAI must carry prompt_cache_retention={PROMPT_CACHE_RETENTION!r}"
)
@pytest.mark.parametrize(("backend", "openai_path"), _retention_backends())
def test_sdk_completion_forwards_prompt_cache_retention_to_openai(
self, resources: ResourceManager, backend: str, openai_path: str
) -> None:
import litellm
capture: Final = _OpenAIBodyCapture(openai_path)
api_base: Final = _openai_edge_api_base(resources, capture, reachable_from_proxy=False)
litellm.completion(
model=backend,
messages=[{"role": "user", "content": "reply with one word"}],
api_base=api_base,
prompt_cache_retention=PROMPT_CACHE_RETENTION,
)
assert capture.retentions() == (PROMPT_CACHE_RETENTION,), (
f"one {openai_path} call to OpenAI must carry prompt_cache_retention={PROMPT_CACHE_RETENTION!r}"
)
@pytest.mark.parametrize(("backend", "openai_path"), _retention_backends())
def test_router_completion_forwards_prompt_cache_retention_to_openai(
self, resources: ResourceManager, backend: str, openai_path: str
) -> None:
import litellm
capture: Final = _OpenAIBodyCapture(openai_path)
api_base: Final = _openai_edge_api_base(resources, capture, reachable_from_proxy=False)
router: Final = litellm.Router(
model_list=[{"model_name": "cache-retention", "litellm_params": {"model": backend, "api_base": api_base}}]
)
router.completion( # pyright: ignore[reportUnknownMemberType] # Router.completion forwards an untyped **kwargs to litellm.completion
model="cache-retention",
messages=[{"role": "user", "content": "reply with one word"}],
prompt_cache_retention=PROMPT_CACHE_RETENTION,
)
assert capture.retentions() == (PROMPT_CACHE_RETENTION,), (
f"one {openai_path} call to OpenAI must carry prompt_cache_retention={PROMPT_CACHE_RETENTION!r}"
)

1
tests/e2e/models.py
View file

@ -364,6 +364,7 @@ class ChatBody(BaseModel):
thinking: ThinkingParam | None = None
service_tier: str | None = None
prompt_cache_key: str | None = None
prompt_cache_retention: str | None = None
tools: Sequence[ChatTool | McpChatTool | GeminiUrlContextTool] | None = None
tool_choice: str | None = None
guardrails: list[str] | None = None

111
tests/e2e/perf/UI_QUERIES.md
View file

@ -1,92 +1,55 @@
# SpendLog UI query performance
# Admin UI journey latency on the 1TB SpendLogs restore
Run `test_ui_queries.py` against a dedicated proxy and a disposable database restored from the metadata-only performance seed. These tests do not call an LLM or measure browser rendering
Run `test_ui_journeys.py` against a dedicated proxy, its Admin UI and a disposable database restored from the metadata-only performance seed. A headless Chromium logs into the dashboard and performs the journeys below, timing each one from the user's click or navigation until the data is on screen. No LLM is called. The test does not look at which routes, SQL or tables served the page, so it stays valid whether the backend gets faster queries or new aggregate tables
The suite measures 42 request shapes across 20 routes used by current dashboard screens. It replays direct HTTP requests against the existing bulk seed, with the seed anchor replacing the wall clock and a UTC browser timezone. It does not run a browser or claim to cover every dashboard request
## Journeys
## UI callers
| name | what the user does | timer stops when |
| --- | --- | --- |
| `logs_open_24h` | opens Logs (default Last 24 Hours) | request rows are rendered and the refresh control is idle |
| `logs_last_7_days` | switches the time range to Last 7 Days | same |
| `logs_next_page` | clicks the next page control | same |
| `logs_filter_key_hash` | filters by the seeded key hash and applies | same |
| `logs_filter_key_alias` | picks the first key alias in the filter drawer and applies | same |
| `logs_filter_team` | picks the first team in the filter drawer and applies | same |
| `logs_filter_model` | picks the first model in the filter drawer and applies | same |
| `usage_open` | opens Usage | summary tiles show dollar values with no skeletons left |
| `usage_last_30_days` | picks Last 30 days in the Usage date picker | same |
| `usage_team_view` | switches the Usage view to Team Usage | same |
| `usage_customer_view` | switches the Usage view to Customer Usage | same |
| `usage_model_activity_30_days` | on the Model Activity tab, picks Last 30 days | same |
| `cache_open` | opens Caching | the Cache Hits tile is visible and the Cache Hits vs API Requests chart has bars |
| `models_open` | opens Models + Endpoints | model rows are rendered and the loading state is gone |
| `keys_open` | opens Virtual Keys | key rows are rendered and the loading state is gone |
All paths below are relative to `ui/litellm-dashboard/src/`. Each case requires an authenticated admin session with access to that screen, and selected resources must exist in the fixture
The mix follows what customers report as slow: the Logs page on `LiteLLM_SpendLogs`, the Usage page on the `LiteLLM_Daily*Spend` rollups (30 day ranges, Customer and Model views), and the Models and Keys pages.
| Cases | UI action and source |
| --- | --- |
| `logs_*` | Logs > Request Logs, change date range, filter, sort or next page. `components/logs/request/useLogFilterLogic.ts` and `RequestLogsFilters.tsx` |
| `usage_*` | Usage overview or entity tab, open key activity, search keys, expand a model group. `app/(dashboard)/usage/_components/components/UsagePageView.tsx`, `EntityUsage/EntityUsage.tsx`, `EntityUsage/entityFetchFns.ts`, `components/UsagePage/components/KeyActivityPanel.tsx` |
| `team_users_*` | Team Usage. `EntityUsage/EntityUsage.tsx` passes selected teams or all available non-dashboard teams to `TeamUserSpendCard.tsx` |
| `cache_activity_*` | Caching > Cache Analytics, optionally filter aliases/models. `app/(dashboard)/hooks/caching/useCacheActivity.ts` |
| `cache_leakage_keys` | Cost Optimization, switch cache leakage to Keys. `app/(dashboard)/cost-optimization/_components/useCacheLeakageKeys.ts` |
## Seeded rows
Request serialization lives in `components/networking.tsx`. Logs use page size 25 and session grouping, including the next-session cursor returned by page one. Individual Logs filters use a 30-day window to exercise larger matching datasets. Their case names end in `_30d` so the measurements do not share history with the old 24-hour filters. The synthetic cursorless page-three request is excluded
The snapshot carries the SpendLogs rows and their Daily* rollups. `setup/ui_query_fixture.py` checks that `LiteLLM_DailyUserSpend` and `LiteLLM_DailyTeamSpend` have rows in the week before the anchor, then fills the tables the snapshot does not have so the Keys, Models and filter journeys see realistic counts: 100 teams (`perf-team-N`), 2,000 users (`perf-user-N`), 5,000 keys whose token hashes and `perf-key-N` aliases match the seeded log rows, and 200 DB models pointed at the perf mock provider. Every insert is `ON CONFLICT DO NOTHING`, so reruns on the same restore are no-ops. Nothing here runs outside the disposable restore
Usage list, search and model drilldown requests do not inject an `api_key` filter. Team activity excludes `litellm-dashboard`. Cache leakage supplies `include_current_utc_day=true`. Dates and filter values are deterministic fixture values, not captured production traffic
Each journey is measured once cold, then `E2E_UI_JOURNEY_REPEATS` more times (default 9) from the same starting state, inside one logged-in tab whose clock is pinned to the seed anchor so relative windows land on seeded rows. The browser timezone is UTC to match the seed
Old Usage is excluded because its expensive-query gate normally hides those calls at this database size. LiteAdmin's agent-selected reports are excluded from this deterministic dashboard set. Missing-token alias recovery is backend functional coverage, not a separate UI action, and its artificial requests and orphan fixtures are removed
## Safety
This mapping is source-verified, not browser-recorded. Recheck the named callers when changing cases. Neither browser capture nor an automatic UI-drift checker is installed by this suite
The seeded source is never the target. `E2E_UI_PERF_DISPOSABLE=1`, `E2E_RESET_SPEND_LOGS=0` and a hostname matching `db-perf-perf-<id>.litellm.svc` are all required, and in-flight SQL is cancelled only on that disposable restore between journeys (see `setup/ui_query_isolation.py`)
## Dataset scope
## Running
Logs, team-user spend and cache activity query the bulk SpendLogs data. User, team, tag and customer usage, plus cache leakage, query daily tables aggregated from that dataset. These are sequential large-dataset queries, not concurrent traffic or saturation tests
A read-only restore of `db-perf-v1-105-1tb` gave exact request totals from daily aggregates: 9,442,606 across 24 hours, 66,077,435 across seven days and 283,189,684 across 30 days. The retained 30-day key filter matches 1,134,559 requests, user filter 11,316,454, team filter 113,298,119 and model filter 94,395,473. These describe matching requests, not rows scanned or returned by the API, and are documentation rather than performance assertions
PostgreSQL table-wide estimates (`pg_class.reltuples`) are approximately 2.70 million daily user rows, 2.70 million daily team rows, 5.40 million daily tag rows and 78.27 million daily end-user rows. These are not exact counts within each Usage case's seven-day window
The individual end-user filter is excluded because it matched only 104 requests in 24 hours and 8,616 across all dates. User/end-user dropdown requests are also excluded because their queries cap the input at 10,000 matching SpendLogs before extracting distinct values. Neither is a broad matching-data workload in this seed
Excluded because they only had tiny matching fixtures or individual-record lookups: organization/agent usage, guardrail/tool logs, request-ID lookup, payload/session drawers, session/search filters for `perf-ui-*`, the OAuth-token filter, and all prompt-caching request cases. The seed has response-cache hits but no prompt-cache usage or injection metadata, so the prompt-caching cases only matched the 320 added rows. No UI-specific rows are inserted anymore
| Remaining cases | Count | Variations |
| --- | ---: | --- |
| Logs | 21 | 24h/7d/30d, page 2 with session cursor, 12 individual filters over 30 days, five additional descending sorts |
| User/team/tag/customer usage | 16 | Aggregated totals, keys, key search, model top keys for each |
| Team-user spend | 2 | All 100 seeded teams or `perf-team-0` |
| Cache activity | 2 | All or selected aliases/model |
| Cache leakage keys | 1 | 30-day window |
## Deployment prerequisites
Use a fresh proxy against the disposable restore, never the active seed source or a shared test database. Apply the target migrations and wait for readiness first. The source seed's `litellm_perf_seed.checkpoint` must exist and contain rows
Configure `litellm_settings.service_callback: [prometheus_system]` and expose metrics for every proxy worker/pod. Set `E2E_DATABASE_METRICS_URLS` to their comma-separated URLs, not a load-balanced endpoint that can miss a pod. Absent failure counters count as zero only when the endpoint exposes a positive PostgreSQL request counter. Missing instrumentation stops the run
No expensive-query override or legacy view creation is needed for these cases. Keep the normal large-database UI restrictions
The cases use explicit dates from the seed anchor where the UI supplies dates. A missing nonempty fixture result records a contract failure rather than reporting a fast empty query as success
## Run
Set `DATABASE_URL` and `LITELLM_MASTER_KEY` securely for this disposable deployment. The URL must name the same database used by the proxy, with hostname `db-perf-perf-<run>.litellm.svc`. The explicit disposable confirmation permits cancellation of leftover queries in this database. Setup reads the existing seed checkpoint without adding UI-specific rows. Destroy the disposable database after the run
The componentized perf stack serves the dashboard behind the same router as the proxy, so `LITELLM_PROXY_URL` covers `/ui`. Set `E2E_UI_BASE_URL` only if the UI lives elsewhere. The admin login uses `E2E_UI_USERNAME` (default `admin`) and `E2E_UI_PASSWORD` (default the master key)
```bash
export LITELLM_PROXY_URL=http://localhost:4000
export E2E_DATABASE_METRICS_URLS=http://localhost:4000/metrics
export E2E_PERF_TESTS=1
export E2E_UI_PERF_DISPOSABLE=1
export E2E_RESET_SPEND_LOGS=0
export E2E_ARTIFACT_DIR="$PWD/ui-query-results"
mkdir -p "$E2E_ARTIFACT_DIR"
uv run pytest tests/e2e/perf/test_ui_queries.py --reruns 0 --junitxml="$E2E_ARTIFACT_DIR/junit.xml" -q
export E2E_PERF_TESTS=1 E2E_UI_PERF_DISPOSABLE=1 E2E_RESET_SPEND_LOGS=0
export DATABASE_URL=postgresql://...@db-perf-perf-<id>.litellm.svc:5432/litellm
export E2E_DATABASE_METRICS_URLS=http://<backend-pod>:4000/metrics
export E2E_ARTIFACT_DIR=/tmp/e2e-artifacts
uv run pytest tests/e2e/perf/test_ui_journeys.py --reruns 0 --junitxml="$E2E_ARTIFACT_DIR/junit.xml" -q
```
Do not enable SpendLog retention or run other traffic on this stack. The `quiet_stack` marker coordinates local pytest callers but cannot stop external traffic. Use one pytest worker for comparable timings
`E2E_UI_JOURNEY_TIMEOUT_SECONDS` (default 120) bounds one journey. A journey that never settles is recorded as a `timeout` sample and a journey whose page rendered empty or threw in the browser as a `ui` sample. The test fails when any sample failed, or when the p50 or p90 across all samples is over the journey's `Ceiling` in `UI_JOURNEY_LIMITS` (`setup/ui_journeys.py`, 10s/20s for Logs journeys, 20s/40s for Usage and Caching, 5s/10s for the Models and Keys tables, generous until the first 1TB run sets real numbers). The artifacts are written before the assertion so a failed run still ships its samples
`E2E_UI_QUERY_REPEATS` defaults to 9, giving ten measured requests including the first, and `E2E_UI_QUERY_TIMEOUT_SECONDS` defaults to 120. Every request is issued once without transport retries. A failed first request stops that case's repeated samples. Cursor prerequisites are recorded separately, and failed prerequisites stop that case's measurements
## Artifacts
## Reading results
Each JUnit testcase `test_ui_journey_latency[<name>]` carries `perf_ui_journeys`, `perf_ui_failures`, `perf_ui_seconds_p50`, `perf_ui_seconds_p90`, `perf_ui_seconds_max`, `perf_ui_first_seconds`, `perf_ui_repeated_seconds_p50`, `perf_ui_repeated_seconds_max` and `perf_ui_database_failures`. `ui-journey-<name>.json` holds every sample with its on-screen row count and `ui-isolation-<name>.jsonl` the SQL isolation evidence
Each case writes `ui-query-<name>.json` with its exact method, path, parameters and body, response status counts, elapsed times, row counts, errors and before/after PostgreSQL failure counters. Response bodies and authentication headers are not saved
The dashboard's `perf_ui_seconds_p50`, `_p90` and `_max` include all successful measured requests, including the first. `perf_ui_samples` is their sample count. p50 is the median and p90 uses nearest rank, the ninth sorted value with ten successes. Failures and cursor prerequisites contribute no latency. Pass percentage is `(perf_ui_requests - perf_ui_failures) / perf_ui_requests * 100`, including cursor prerequisites in the request counts
The first-request and repeated-request metrics are retained for existing consumers. This is not a cold-cache claim: previous cases and previous runs can warm database buffers or application caches. Restart the disposable proxy between independent runs when comparing independent runs
Repeated p50 is reported for successful samples, p95 only with at least 20 and p99 only with at least 100. First-request latency is omitted when that request fails. HTTP, JSON, response-shape, nonempty-result, cursor failures are recorded, never asserted. Endpoint failures, database counter increases and latency do not fail CI
Each JUnit testcase `test_ui_query_performance[<case>]` carries `perf_ui_requests`, `perf_ui_failures`, `perf_ui_repeated_samples`, and nonzero `perf_ui_error_<code>` counts. Codes distinguish `http_400`, `http_500`, `http_503`, other HTTP statuses, `timeout`, `network`, `invalid_json` and `contract`. Successful latency properties are `perf_ui_first_seconds` and `perf_ui_repeated_seconds_p50`, `_p95`, `_p99`, `_max`. A failed cursor prerequisite is counted once and contributes no latency
Before every case the runner requires three idle database observations and stable failure counters. After each case it cancels leftover SQL, terminates idle transactions and waits for idle and stable counters again. `ui-isolation-<case>.jsonl` retains active SQL, waits, elapsed times, raw samples and counter observations. `perf_ui_database_failures` measures the increase before cleanup, while `perf_ui_cleanup_database_failures` measures the increase observed afterward, including cancellations and potentially delayed request failures. Isolation, instrumentation or setup failures still stop the run rather than publish contaminated measurements. Disk and process caches are not reset between cases
The Buildkite `ui_queries` lane restores the pinned snapshot independently, waits for migrations and indexes to complete, runs these cases with no retries, uploads JUnit and JSON artifacts, and removes the temporary resources. It follows the nightly regular perf and database-upgrade builds using the same LiteLLM SHA. Shinkiro ingestion is a separate change
Validate a small restore first. Small-database numbers validate the runner only, not full-size performance. Run the same cases against the completed full-size snapshot before publishing large-database results
The Buildkite `ui_queries` lane restores the pinned snapshot, waits for migrations and indexes, runs these journeys with no retries and uploads the JUnit and JSON artifacts. Validate selectors on a small restore first, then take numbers from the completed 1TB snapshot

14
tests/e2e/perf/cases/mixed_size_requests_2rps.py
View file

@ -64,14 +64,14 @@ CASE: Final = PerformanceTestCase(
Limit("unexpected_status_count", max=3),
Limit("gateway_added_seconds", percentile=50, max=0.58),
Limit("gateway_added_seconds", percentile=99, max=6.1),
Limit("litellm_overhead_seconds", percentile=50, max=0.05),
Limit("litellm_overhead_seconds", percentile=99, max=2.2),
Limit("litellm_response_seconds", percentile=50, max=1.2),
Limit("litellm_overhead_seconds", percentile=50, max=0.1),
Limit("litellm_overhead_seconds", percentile=99, max=3.4),
Limit("litellm_response_seconds", percentile=50, max=1.0),
Limit("litellm_response_seconds", percentile=99, max=3.8),
Limit("prometheus_redis_latency_mean_seconds", max=0.003),
Limit("postgres_active_sessions", max=0.004),
Limit("postgres_slowest_statement_seconds", max=0.045),
Limit("postgres_peak_connections", max=30),
Limit("prometheus_redis_latency_mean_seconds", max=0.005),
Limit("postgres_active_sessions", max=0.016),
Limit("postgres_slowest_statement_seconds", max=0.25),
Limit("postgres_peak_connections", max=45),
Limit("time_to_first_token_seconds", percentile=50, max=2.7),
Limit("time_to_first_token_seconds", percentile=99, max=8.8),
Limit("memory_growth_megabytes", max=600),

14
tests/e2e/perf/cases/small_bodies_5rps.py
View file

@ -65,14 +65,14 @@ CASE: Final = PerformanceTestCase(
limits=(
Limit("unexpected_status_count", max=3),
Limit("gateway_added_seconds", percentile=50, max=0.29),
Limit("litellm_overhead_seconds", percentile=50, max=0.02),
Limit("litellm_overhead_seconds", percentile=99, max=0.8),
Limit("litellm_response_seconds", percentile=50, max=0.8),
Limit("litellm_response_seconds", percentile=99, max=2.8),
Limit("litellm_overhead_seconds", percentile=50, max=0.012),
Limit("litellm_overhead_seconds", percentile=99, max=0.35),
Limit("litellm_response_seconds", percentile=50, max=1.0),
Limit("litellm_response_seconds", percentile=99, max=2.0),
Limit("prometheus_redis_latency_mean_seconds", max=0.006),
Limit("postgres_active_sessions", max=0.006),
Limit("postgres_slowest_statement_seconds", max=0.035),
Limit("postgres_peak_connections", max=30),
Limit("postgres_active_sessions", max=0.01),
Limit("postgres_slowest_statement_seconds", max=0.025),
Limit("postgres_peak_connections", max=45),
Limit("time_to_first_token_seconds", percentile=99, max=3.7),
Limit("memory_growth_megabytes", max=100),
Limit("postgres_calls_per_request", max=1.8),

397
tests/e2e/perf/setup/ui_journeys.py Normal file
View file

@ -0,0 +1,397 @@
from __future__ import annotations
import math
import re
import statistics
import time
from collections.abc import Callable, Generator
from contextlib import contextmanager
from dataclasses import dataclass
from datetime import datetime
from itertools import chain
from types import MappingProxyType
from typing import Final
from playwright.sync_api import Error as PlaywrightError
from playwright.sync_api import Locator, Page, Response, expect, sync_playwright
from playwright.sync_api import TimeoutError as PlaywrightTimeoutError
from prometheus_client.parser import text_string_to_metric_families
from pydantic import BaseModel, ConfigDict
MONEY: Final = re.compile(r"^\$[\d,]+\.\d{2}$")
LOADING_LOGS: Final = "Loading request logs…"
LOADING_MODELS: Final = "Loading models…"
LOADING_KEYS: Final = "Loading keys..."
class Sample(BaseModel):
model_config = ConfigDict(frozen=True)
seconds: float
rows: int = 0
error: str = ""
class Timings(BaseModel):
model_config = ConfigDict(frozen=True)
first_seconds: float | None
repeated_samples: int
repeated_p50_seconds: float | None
repeated_p95_seconds: float | None
repeated_p99_seconds: float | None
repeated_max_seconds: float | None
@dataclass(frozen=True, slots=True)
class Journey:
"""One dashboard interaction. ``start`` brings the page to the state the user acts from (untimed),
``act`` is the click or navigation being measured and ``settled`` blocks until the data the user
was waiting for is on screen, returning the number of rows or tiles that rendered."""
name: str
surface: str
start: Callable[[Page, float], None]
act: Callable[[Page], None]
settled: Callable[[Page, float], int]
@dataclass(frozen=True, slots=True)
class Dashboard:
base_url: str
username: str
password: str
anchor: datetime
def summarize(samples: tuple[Sample, ...]) -> Timings:
assert samples, "No UI journeys were measured"
repeated: Final = tuple(sorted(sample.seconds for sample in samples[1:] if not sample.error))
return Timings(
first_seconds=samples[0].seconds if not samples[0].error else None,
repeated_samples=len(repeated),
repeated_p50_seconds=statistics.median(repeated) if repeated else None,
repeated_p95_seconds=repeated[math.ceil(len(repeated) * 0.95) - 1] if len(repeated) >= 20 else None,
repeated_p99_seconds=repeated[math.ceil(len(repeated) * 0.99) - 1] if len(repeated) >= 100 else None,
repeated_max_seconds=max(repeated) if repeated else None,
)
def error_code(sample: Sample) -> str:
if not sample.error:
return ""
return "timeout" if sample.error == "timeout" else "ui"
def ui_junit_properties(samples: tuple[Sample, ...]) -> tuple[tuple[str, float], ...]:
codes: Final = tuple(error_code(sample) for sample in samples if sample.error)
timings: Final = summarize(samples)
successful: Final = tuple(sorted(sample.seconds for sample in samples if not sample.error))
latency: Final = (
("seconds_p50", statistics.median(successful) if successful else None),
("seconds_p90", successful[math.ceil(len(successful) * 0.90) - 1] if successful else None),
("seconds_max", max(successful) if successful else None),
("first_seconds", timings.first_seconds),
("repeated_seconds_p50", timings.repeated_p50_seconds),
("repeated_seconds_p95", timings.repeated_p95_seconds),
("repeated_seconds_p99", timings.repeated_p99_seconds),
("repeated_seconds_max", timings.repeated_max_seconds),
)
return (
("perf_ui_journeys", len(samples)),
("perf_ui_failures", len(codes)),
("perf_ui_samples", len(successful)),
("perf_ui_repeated_samples", timings.repeated_samples),
*((f"perf_ui_error_{code}", codes.count(code)) for code in sorted(frozenset(codes))),
*((f"perf_ui_{name}", value) for name, value in latency if value is not None),
)
@dataclass(frozen=True, slots=True)
class Ceiling:
p50_seconds: float
p90_seconds: float
LOGS_CEILING: Final = Ceiling(p50_seconds=10, p90_seconds=20)
CHART_CEILING: Final = Ceiling(p50_seconds=20, p90_seconds=40)
TABLE_CEILING: Final = Ceiling(p50_seconds=5, p90_seconds=10)
UI_JOURNEY_LIMITS: Final = MappingProxyType(
{
"logs_open_24h": LOGS_CEILING,
"logs_last_7_days": LOGS_CEILING,
"logs_next_page": LOGS_CEILING,
"logs_filter_key_hash": LOGS_CEILING,
"logs_filter_key_alias": LOGS_CEILING,
"logs_filter_team": LOGS_CEILING,
"logs_filter_model": LOGS_CEILING,
"usage_open": CHART_CEILING,
"usage_last_30_days": CHART_CEILING,
"usage_team_view": CHART_CEILING,
"usage_customer_view": CHART_CEILING,
"usage_model_activity_30_days": CHART_CEILING,
"cache_open": CHART_CEILING,
"models_open": TABLE_CEILING,
"keys_open": TABLE_CEILING,
}
)
def assert_within_ceiling(journey_name: str, samples: tuple[Sample, ...]) -> None:
failed: Final = tuple(sample.error for sample in samples if sample.error)
assert not failed, f"{len(failed)} of {len(samples)} {journey_name} samples failed: {failed[0]}"
metrics: Final = dict(ui_junit_properties(samples))
ceiling: Final = UI_JOURNEY_LIMITS[journey_name]
assert metrics["perf_ui_seconds_p50"] <= ceiling.p50_seconds, (
f"{journey_name} p50 {metrics['perf_ui_seconds_p50']:.1f}s over {ceiling.p50_seconds}s"
)
assert metrics["perf_ui_seconds_p90"] <= ceiling.p90_seconds, (
f"{journey_name} p90 {metrics['perf_ui_seconds_p90']:.1f}s over {ceiling.p90_seconds}s"
)
@contextmanager
def dashboard_page(dashboard: Dashboard, timeout: float) -> Generator[Page]:
"""A logged-in admin tab whose clock reads the seed anchor, so relative windows such as
Last 24 Hours land on seeded rows. Live Tail is off so its 15 second refetch cannot land
inside a sample. Playwright keeps a UTC timezone to match the seed."""
with sync_playwright() as playwright:
browser: Final = playwright.chromium.launch()
context: Final = browser.new_context(
base_url=dashboard.base_url, timezone_id="UTC", viewport={"width": 1600, "height": 1000}
)
context.set_default_timeout(timeout * 1000)
context.add_init_script('sessionStorage.setItem("isLiveTail", "false")')
page: Final = context.new_page()
page.clock.set_fixed_time(dashboard.anchor)
_login(page, dashboard)
try:
yield page
finally:
context.close()
browser.close()
def _login(page: Page, dashboard: Dashboard) -> None:
_ = page.goto("/ui/login")
page.get_by_placeholder("Enter your username").fill(dashboard.username)
page.get_by_placeholder("Enter your password").fill(dashboard.password)
page.get_by_role("button", name="Login", exact=True).click()
page.wait_for_url(lambda url: "/ui" in url and "/login" not in url)
expect(page.get_by_label("LiteLLM home")).to_be_visible()
def _is_data_fetch(response: Response) -> bool:
return response.request.resource_type in ("fetch", "xhr") and "/_next/" not in response.url
def measure(page: Page, journey: Journey, timeout: float) -> Sample:
"""Wall time from the user's action to the data being on screen. Every sample starts from the
same pre-action state so repeats measure the same journey, not a warmed client cache."""
try:
journey.start(page, timeout)
except (PlaywrightError, AssertionError) as error:
return Sample(seconds=0.0, error=f"start {type(error).__name__}: {error}"[:300])
started: Final = time.monotonic()
try:
with page.expect_response(_is_data_fetch, timeout=timeout * 1000):
journey.act(page)
rows: Final = journey.settled(page, timeout)
except PlaywrightTimeoutError:
return Sample(seconds=time.monotonic() - started, error="timeout")
except (PlaywrightError, AssertionError) as error:
return Sample(seconds=time.monotonic() - started, error=f"{type(error).__name__}: {error}"[:300])
return Sample(seconds=time.monotonic() - started, rows=rows)
def _visible(page: Page, test_id: str) -> Locator:
return page.get_by_test_id(test_id).filter(visible=True).first
def _logs_rows(page: Page) -> Locator:
return page.locator("table").filter(visible=True).first.locator("tbody tr")
def _logs_settled(page: Page, timeout: float) -> int:
millis: Final = timeout * 1000
rows: Final = _logs_rows(page).filter(has_not_text="No results")
expect(rows.first).to_be_visible(timeout=millis)
expect(_visible(page, "datatable-refresh")).to_be_enabled(timeout=millis)
expect(page.get_by_text(LOADING_LOGS)).to_have_count(0, timeout=millis)
return rows.count()
def _no_skeletons(page: Page, timeout: float) -> None:
expect(page.locator('[data-slot="skeleton"]')).to_have_count(0, timeout=timeout * 1000)
def _usage_settled(page: Page, timeout: float) -> int:
millis: Final = timeout * 1000
expect(page.get_by_text("Total Spend", exact=True).first).to_be_visible(timeout=millis)
_no_skeletons(page, timeout)
money: Final = page.get_by_text(MONEY)
expect(money.first).to_be_visible(timeout=millis)
return money.count()
def _cache_settled(page: Page, timeout: float) -> int:
millis: Final = timeout * 1000
expect(page.get_by_text("Cache Hits", exact=True)).to_be_visible(timeout=millis)
bars: Final = page.locator(".recharts-bar-rectangle")
expect(bars.first).to_be_visible(timeout=millis)
return bars.count()
def _table_settled(loading: str) -> Callable[[Page, float], int]:
def settled(page: Page, timeout: float) -> int:
millis: Final = timeout * 1000
rows: Final = page.locator("[data-row-id]")
expect(rows.first).to_be_visible(timeout=millis)
expect(page.get_by_text(loading)).to_have_count(0, timeout=millis)
return rows.count()
return settled
def _open_logs(page: Page, timeout: float) -> None:
_ = page.goto("/ui/logs")
_ = _logs_settled(page, timeout)
def _open_usage(page: Page, timeout: float) -> None:
_ = page.goto("/ui/usage")
_ = _usage_settled(page, timeout)
def _filters_drawer(page: Page) -> Locator:
_visible(page, "datatable-filters-trigger").click()
drawer: Final = page.get_by_role("dialog", name="Filters")
expect(drawer).to_be_visible()
return drawer
def _apply_filters(page: Page) -> None:
page.get_by_role("dialog", name="Filters").get_by_role("button", name="Apply Filters").click()
def _pick_first_option(page: Page, combobox: str) -> None:
_filters_drawer(page).get_by_role("combobox", name=combobox).click()
page.get_by_role("option").first.click()
def _open_time_range(page: Page, timeout: float) -> None:
_open_logs(page, timeout)
page.get_by_role("button", name="Last 24 Hours").click()
def _open_usage_range(page: Page, timeout: float) -> None:
_open_usage(page, timeout)
page.locator('[data-slot="advanced-date-picker-trigger"]').click()
def _model_activity_settled(page: Page, timeout: float) -> int:
millis: Final = timeout * 1000
expect(page.locator('[role="tab"][data-active]', has_text="Model Activity")).to_be_visible(timeout=millis)
_no_skeletons(page, timeout)
money: Final = page.get_by_text(MONEY).filter(visible=True)
expect(money.first).to_be_visible(timeout=millis)
return money.count()
def _open_model_activity_range(page: Page, timeout: float) -> None:
_open_usage(page, timeout)
page.get_by_role("tab", name="Model Activity").click()
_ = _model_activity_settled(page, timeout)
page.locator('[data-slot="advanced-date-picker-trigger"]').click()
def _pick_usage_preset(preset: str) -> Callable[[Page], None]:
def act(page: Page) -> None:
page.locator('[data-slot="advanced-date-picker-preset"]', has_text=preset).click()
return act
def _select_usage_view(page: Page, view: str) -> None:
page.locator('[data-slot="select-trigger"]').filter(has_text="Usage").first.click()
page.get_by_role("option", name=re.compile(rf"^{re.escape(view)}")).click()
def _nothing(page: Page, timeout: float) -> None:
return None
def _goto(path: str) -> Callable[[Page], None]:
def act(page: Page) -> None:
_ = page.goto(path)
return act
def _logs_filter(name: str, prepare: Callable[[Page], None]) -> Journey:
def start(page: Page, timeout: float) -> None:
_open_logs(page, timeout)
prepare(page)
return Journey(name, "logs", start, _apply_filters, _logs_settled)
def journeys(api_key: str) -> tuple[Journey, ...]:
return (
Journey("logs_open_24h", "logs", _nothing, _goto("/ui/logs"), _logs_settled),
Journey(
"logs_last_7_days",
"logs",
_open_time_range,
lambda page: page.get_by_role("button", name="Last 7 Days", exact=True).click(),
_logs_settled,
),
Journey(
"logs_next_page", "logs", _open_logs, lambda page: _visible(page, "pagination-next").click(), _logs_settled
),
_logs_filter(
"logs_filter_key_hash",
lambda page: _filters_drawer(page).get_by_placeholder("Enter key hash…").fill(api_key),
),
_logs_filter("logs_filter_key_alias", lambda page: _pick_first_option(page, "Search a key alias")),
_logs_filter("logs_filter_team", lambda page: _pick_first_option(page, "Search or select a team")),
_logs_filter("logs_filter_model", lambda page: _pick_first_option(page, "Search a model")),
Journey("usage_open", "usage", _nothing, _goto("/ui/usage"), _usage_settled),
Journey("usage_last_30_days", "usage", _open_usage_range, _pick_usage_preset("Last 30 days"), _usage_settled),
Journey(
"usage_team_view", "usage", _open_usage, lambda page: _select_usage_view(page, "Team Usage"), _usage_settled
),
Journey(
"usage_customer_view",
"usage",
_open_usage,
lambda page: _select_usage_view(page, "Customer Usage"),
_usage_settled,
),
Journey(
"usage_model_activity_30_days",
"usage",
_open_model_activity_range,
_pick_usage_preset("Last 30 days"),
_model_activity_settled,
),
Journey("cache_open", "caching", _nothing, _goto("/ui/caching"), _cache_settled),
Journey("models_open", "models", _nothing, _goto("/ui/models-and-endpoints"), _table_settled(LOADING_MODELS)),
Journey("keys_open", "keys", _nothing, _goto("/ui/api-keys"), _table_settled(LOADING_KEYS)),
)
def database_failures(text: str) -> float:
families: Final = tuple(text_string_to_metric_families(text))
required: Final = frozenset(("litellm_postgres_failed_requests", "litellm_batch_write_to_db_failed_requests"))
missing: Final = required - frozenset(family.name for family in families)
all_samples: Final = tuple(chain.from_iterable(family.samples for family in families))
has_activity: Final = any(
sample.name == "litellm_postgres_total_requests_total" and math.isfinite(sample.value) and sample.value > 0
for sample in all_samples
)
assert not missing or has_activity, f"Missing PostgreSQL failure counters {missing}"
names: Final = frozenset(f"{name}_total" for name in required)
values: Final = tuple(sample.value for sample in all_samples if sample.name in names)
assert all(math.isfinite(value) and value >= 0 for value in values), "Invalid PostgreSQL failure counter"
return sum(values)

227
tests/e2e/perf/setup/ui_queries.py
View file

@ -1,227 +0,0 @@
from __future__ import annotations
import math
import statistics
import time
from collections.abc import Callable
from dataclasses import dataclass
from itertools import chain
from typing import Final, Literal
from urllib.parse import urlencode
from e2e_http import (
NetworkError,
RateLimitedError,
Result,
Success,
UnauthorizedError,
UnknownApiError,
ValidationError,
without_retries,
)
from e2e_metadata import step
from prometheus_client.parser import text_string_to_metric_families
from pydantic import BaseModel, ConfigDict, JsonValue, RootModel
from transport import HttpTransport
class Query(BaseModel):
model_config = ConfigDict(frozen=True)
class Document(RootModel[JsonValue]):
model_config = ConfigDict(frozen=True)
@dataclass(frozen=True, slots=True)
class Contract:
collection: tuple[str, ...] = ()
row_fields: tuple[str, ...] = ()
positive: tuple[str, ...] = ()
object_fields: tuple[str, ...] = ()
excluded_ids: frozenset[str] = frozenset()
def count(self, document: Document) -> int:
value: Final = descend(document.root, self.collection)
if self.object_fields:
assert isinstance(value, dict), "Expected a JSON object"
assert set(self.object_fields) <= value.keys(), f"Missing fields {self.object_fields}"
return 1
assert isinstance(value, list) and value, f"Empty/missing collection {self.collection}, check the fixture"
for row in value:
if self.excluded_ids:
assert isinstance(row, dict) and isinstance(row_id := row.get("request_id"), str), "Missing request_id"
assert row_id not in self.excluded_ids, "Repeated page-one row"
if self.row_fields or self.positive:
assert isinstance(row, dict), "Expected object rows"
assert set(self.row_fields) <= row.keys(), f"Missing row fields {self.row_fields}"
for name in self.positive:
assert any(tuple(_positive(row, name) for row in value)), f"Expected positive {name}"
return len(value)
def _positive(row: JsonValue, name: str) -> bool:
assert isinstance(row, dict), "Expected object row"
number: Final = row.get(name)
assert isinstance(number, (int, float)) and not isinstance(number, bool), f"Expected numeric {name}"
assert math.isfinite(number), f"Expected finite {name}"
return number > 0
def descend(value: JsonValue, path: tuple[str, ...]) -> JsonValue:
if not path:
return value
assert isinstance(value, dict) and path[0] in value, f"Missing response path {path}"
return descend(value[path[0]], path[1:])
@dataclass(frozen=True, slots=True)
class UiQueryCase:
name: str
path: str
params: BaseModel
contract: Contract
surface: Literal["current"] = "current"
body: BaseModel | None = None
cursor: Literal["none", "session"] = "none"
class Sample(BaseModel):
model_config = ConfigDict(frozen=True)
seconds: float
status: int
rows: int = 0
error: str = ""
class Timings(BaseModel):
model_config = ConfigDict(frozen=True)
first_seconds: float | None
repeated_samples: int
repeated_p50_seconds: float | None
repeated_p95_seconds: float | None
repeated_p99_seconds: float | None
repeated_max_seconds: float | None
def summarize(samples: tuple[Sample, ...]) -> Timings:
assert samples, "No UI requests were measured"
repeated: Final = tuple(sorted(sample.seconds for sample in samples[1:] if not sample.error))
return Timings(
first_seconds=samples[0].seconds if not samples[0].error else None,
repeated_samples=len(repeated),
repeated_p50_seconds=statistics.median(repeated) if repeated else None,
repeated_p95_seconds=repeated[math.ceil(len(repeated) * 0.95) - 1] if len(repeated) >= 20 else None,
repeated_p99_seconds=repeated[math.ceil(len(repeated) * 0.99) - 1] if len(repeated) >= 100 else None,
repeated_max_seconds=max(repeated) if repeated else None,
)
def error_code(sample: Sample) -> str:
if not sample.error:
return ""
if sample.status > 0 and sample.status != 200:
return f"http_{sample.status}"
if sample.error in ("timeout", "network"):
return sample.error
if sample.error == "Response is not valid JSON":
return "invalid_json"
return "contract"
def ui_junit_properties(
samples: tuple[Sample, ...], prerequisite: Sample | None = None
) -> tuple[tuple[str, float], ...]:
requests: Final = (prerequisite, *samples) if prerequisite is not None and not prerequisite.error else samples
codes: Final = tuple(error_code(sample) for sample in requests if sample.error)
timings: Final = summarize(samples)
successful: Final = tuple(sorted(sample.seconds for sample in samples if not sample.error))
latency: Final = (
("seconds_p50", statistics.median(successful) if successful else None),
("seconds_p90", successful[math.ceil(len(successful) * 0.90) - 1] if successful else None),
("seconds_max", max(successful) if successful else None),
("first_seconds", timings.first_seconds),
("repeated_seconds_p50", timings.repeated_p50_seconds),
("repeated_seconds_p95", timings.repeated_p95_seconds),
("repeated_seconds_p99", timings.repeated_p99_seconds),
("repeated_seconds_max", timings.repeated_max_seconds),
)
return (
("perf_ui_requests", len(requests)),
("perf_ui_failures", len(codes)),
("perf_ui_samples", len(successful)),
("perf_ui_repeated_samples", timings.repeated_samples),
*((f"perf_ui_error_{code}", codes.count(code)) for code in sorted(frozenset(codes))),
*((f"perf_ui_{name}", value) for name, value in latency if value is not None),
)
def _failure(result: Result[Document], seconds: float) -> Sample:
match result:
case Success(status_code=status):
return Sample(seconds=seconds, status=status, error=f"Unexpected HTTP {status}")
case UnknownApiError(status_code=status):
return Sample(seconds=seconds, status=status, error=f"HTTP {status}")
case UnauthorizedError():
return Sample(seconds=seconds, status=401, error="Unauthorized")
case RateLimitedError():
return Sample(seconds=seconds, status=429, error="Rate limited")
case NetworkError(message=message):
return Sample(seconds=seconds, status=-1, error="timeout" if "timed out" in message.lower() else "network")
case ValidationError():
return Sample(seconds=seconds, status=-1, error="Response is not valid JSON")
def measure(
issue: Callable[[], Result[Document]],
contract: Contract,
*,
clock: Callable[[], float] = time.monotonic,
) -> tuple[Sample, Document | None]:
started: Final = clock()
with without_retries():
result: Final = issue()
seconds: Final = clock() - started
match result:
case Success(status_code=200, data=document):
try:
count: Final = contract.count(document)
except AssertionError as error:
return Sample(seconds=seconds, status=200, error=str(error)), document
return Sample(seconds=seconds, status=200, rows=count), document
case _:
return _failure(result, seconds), None
def query_path(case: UiQueryCase) -> str:
query: Final = urlencode(case.params.model_dump(by_alias=True, exclude_none=True), doseq=True)
return f"{case.path}?{query}" if query else case.path
@step("Measure UI query {case.name} at {case.path}")
def issue_query(transport: HttpTransport, case: UiQueryCase, *, timeout: float) -> Result[Document]:
if case.body is not None:
return transport.post(
case.path, headers=transport.master, json=case.body, response_type=Document, timeout=timeout
)
return transport.get(
query_path(case), headers=transport.master, params=Query(), response_type=Document, timeout=timeout
)
def database_failures(text: str) -> float:
families: Final = tuple(text_string_to_metric_families(text))
required: Final = frozenset(("litellm_postgres_failed_requests", "litellm_batch_write_to_db_failed_requests"))
missing: Final = required - frozenset(family.name for family in families)
all_samples: Final = tuple(chain.from_iterable(family.samples for family in families))
has_activity: Final = any(
sample.name == "litellm_postgres_total_requests_total" and math.isfinite(sample.value) and sample.value > 0
for sample in all_samples
)
assert not missing or has_activity, f"Missing PostgreSQL failure counters {missing}"
names: Final = frozenset(f"{name}_total" for name in required)
values: Final = tuple(sample.value for sample in all_samples if sample.name in names)
assert all(math.isfinite(value) and value >= 0 for value in values), "Invalid PostgreSQL failure counter"
return sum(values)

206
tests/e2e/perf/setup/ui_query_cases.py
View file

@ -1,206 +0,0 @@
from __future__ import annotations
from dataclasses import replace
from datetime import datetime, timedelta
from itertools import chain
from typing import Final, Literal
from pydantic import ConfigDict
from setup.ui_queries import Contract, Document, Query, UiQueryCase, descend
class Params(Query):
model_config = ConfigDict(frozen=True, extra="forbid", populate_by_name=True)
start_date: str | None = None
timezone: str | None = None
offset: int | None = None
model_group: str | None = None
by_model_group: str | None = None
end_date: str | None = None
page: int | None = None
page_size: int | None = None
sort_by: str | None = None
sort_order: str | None = None
group_by_session: Literal["true"] | None = None
session_cursor: str | None = None
api_key: str | None = None
team_id: str | None = None
team_ids: str | None = None
user_id: str | None = None
model: str | None = None
model_id: str | None = None
key_alias: str | None = None
search: str | None = None
status_filter: str | None = None
cache_hit_filter: str | None = None
error_code: str | None = None
error_message: str | None = None
span_type: str | None = None
exclude_internal_health_checks: Literal["true"] | None = None
key_aliases: tuple[str, ...] | None = None
models: tuple[str, ...] | None = None
limit: int | None = None
tags: str | None = None
exclude_team_ids: str | None = None
include_current_utc_day: Literal["true"] | None = None
def _timestamp(value: datetime) -> str:
return value.strftime("%Y-%m-%d %H:%M:%S")
def next_page(case: UiQueryCase, document: Document) -> UiQueryCase:
params: Final = Params.model_validate(case.params.model_dump())
rows: Final = descend(document.root, case.contract.collection)
assert isinstance(rows, list), "Missing page-one rows"
ids: Final = frozenset(str(row["request_id"]) for row in rows if isinstance(row, dict))
contract: Final = replace(case.contract, excluded_ids=ids)
cursor: Final = descend(document.root, ("next_session_cursor",))
assert isinstance(cursor, str) and cursor, "Missing session cursor"
return replace(case, params=params.model_copy(update={"page": 2, "session_cursor": cursor}), contract=contract)
_LOGS: Final = Contract(("data",), ("request_id", "startTime", "metadata"))
def log_cases(anchor: datetime, api_key: str) -> tuple[UiQueryCase, ...]:
params: Final = Params(
start_date=_timestamp(anchor - timedelta(days=1)),
end_date=_timestamp(anchor),
page=1,
page_size=25,
sort_by="startTime",
sort_order="desc",
group_by_session="true",
)
base: Final = UiQueryCase("logs_default_24h", "/spend/logs/ui", params, _LOGS)
filtered: Final = params.model_copy(update={"start_date": _timestamp(anchor - timedelta(days=30))})
filters: Final = (
("api_key", api_key),
("team_id", "perf-team-0"),
("user_id", "perf-user-0"),
("model", "perf-chat"),
("model_id", "perf-model-0"),
("key_alias", "perf-key-0"),
("status_filter", "failure"),
("cache_hit_filter", "hit"),
("error_code", "429"),
("error_message", "Synthetic"),
("span_type", "llm"),
("exclude_internal_health_checks", "true"),
)
return (
base,
replace(base, name="logs_cursor_page_2", params=params.model_copy(update={"page": 2}), cursor="session"),
replace(
base,
name="logs_7d",
params=params.model_copy(update={"start_date": _timestamp(anchor - timedelta(days=7))}),
),
replace(
base,
name="logs_30d",
params=params.model_copy(update={"start_date": _timestamp(anchor - timedelta(days=30))}),
),
*(
replace(base, name=f"logs_{name}_30d", params=filtered.model_copy(update={name: value}))
for name, value in filters
),
*(
replace(base, name=f"logs_sort_{name}", params=params.model_copy(update={"sort_by": name}))
for name in ("spend", "total_tokens", "request_duration_ms", "model", "ttft_ms")
),
)
def analytics_cases(anchor: datetime) -> tuple[UiQueryCase, ...]:
dates: Final = Params(
start_date=(anchor - timedelta(days=7)).date().isoformat(), end_date=anchor.date().isoformat()
)
cache: Final = UiQueryCase(
"cache_activity_7d", "/global/activity/cache_hits", dates, Contract(("groups",), ("call_type",))
)
return (
UiQueryCase(
"team_users_all_teams",
"/team/spend/by_user",
dates.model_copy(update={"team_ids": ",".join(f"perf-team-{i}" for i in range(100))}),
Contract(("results",), ("team_id", "user_id", "spend"), ("api_requests",)),
),
UiQueryCase(
"team_users_filtered",
"/team/spend/by_user",
dates.model_copy(update={"team_ids": "perf-team-0"}),
Contract(("results",), ("team_id", "user_id", "spend"), ("api_requests",)),
),
cache,
replace(
cache,
name="cache_activity_filtered",
params=dates.model_copy(update={"key_aliases": ("perf-key-0", "perf-key-1"), "models": ("perf-chat",)}),
),
)
def _entity_activity_cases(anchor: datetime, entity: str) -> tuple[UiQueryCase, ...]:
dates: Final = Params(
start_date=(anchor - timedelta(days=7)).date().isoformat(),
end_date=anchor.date().isoformat(),
timezone="0",
exclude_team_ids="litellm-dashboard" if entity == "team" else None,
)
routes: Final = (
("aggregated", dates, Contract(("results",), ("date", "breakdown"))),
(
"aggregated/keys",
dates.model_copy(update={"offset": 0, "limit": 50}),
Contract(("api_keys",), ("api_key", "metrics", "metadata")),
),
(
"aggregated/search",
dates.model_copy(update={"search": "perf-key"}),
Contract(("api_keys",), ("api_key", "metrics", "metadata")),
),
(
"aggregated/model_top_keys",
dates.model_copy(update={"model_group": "perf-chat", "by_model_group": "true"}),
Contract(("api_keys",), ("api_key", "metrics", "metadata")),
),
)
return tuple(
UiQueryCase(
f"usage_{entity}_{route.replace('/', '_')}",
f"/{entity}/daily/activity/{route}",
params,
contract,
)
for route, params, contract in routes
)
def activity_cases(anchor: datetime) -> tuple[UiQueryCase, ...]:
groups: Final = tuple(_entity_activity_cases(anchor, entity) for entity in ("user", "team", "tag", "customer"))
return (
*chain.from_iterable(groups),
UiQueryCase(
"cache_leakage_keys",
"/user/daily/activity/aggregated/cache_leakage_keys",
Params(
start_date=(anchor - timedelta(days=30)).date().isoformat(),
end_date=anchor.date().isoformat(),
timezone="0",
include_current_utc_day="true",
),
Contract(("api_keys",), ("api_key", "metrics", "metadata")),
),
)
def cases(anchor: datetime, api_key: str) -> tuple[UiQueryCase, ...]:
return (
*log_cases(anchor, api_key),
*analytics_cases(anchor),
*activity_cases(anchor),
)

72
tests/e2e/perf/setup/ui_query_fixture.py
View file

@ -3,7 +3,7 @@ from __future__ import annotations
import hashlib
import os
import time
from datetime import datetime
from datetime import datetime, timedelta
from typing import Final
import psycopg
@ -20,6 +20,45 @@ class UiFixture(BaseModel):
api_key: str
SEEDED_TEAMS: Final = 100
SEEDED_USERS: Final = 2000
SEEDED_KEYS: Final = 5000
SEEDED_MODELS: Final = 200
DAILY_TABLES: Final = ("LiteLLM_DailyUserSpend", "LiteLLM_DailyTeamSpend")
_SEED_SQL: Final = (
"""
INSERT INTO "LiteLLM_TeamTable" (team_id, team_alias, members_with_roles, metadata, model_spend, model_max_budget)
SELECT 'perf-team-' || n, 'perf-team-' || n, '[]', '{}', '{}', '{}' FROM generate_series(0, %(teams)s - 1) AS n
ON CONFLICT (team_id) DO NOTHING
""",
"""
INSERT INTO "LiteLLM_UserTable" (user_id, user_alias, user_email, user_role, team_id, metadata, model_spend, model_max_budget)
SELECT 'perf-user-' || n, 'perf-user-' || n, 'perf-user-' || n || '@example.com', 'internal_user',
'perf-team-' || (n %% %(teams)s), '{}', '{}', '{}'
FROM generate_series(0, %(users)s - 1) AS n
ON CONFLICT (user_id) DO NOTHING
""",
"""
INSERT INTO "LiteLLM_VerificationToken" (token, key_name, key_alias, team_id, user_id, aliases, config, permissions,
metadata, model_spend, model_max_budget, created_at, updated_at)
SELECT md5('perf-key:' || n) || md5('perf-key-tail:' || n), 'sk-...perf' || n, 'perf-key-' || n,
'perf-team-' || (n %% %(teams)s), 'perf-user-' || (n %% %(users)s), '{}', '{}', '{}', '{}', '{}', '{}',
%(anchor)s - make_interval(days => n %% 365), %(anchor)s
FROM generate_series(0, %(keys)s - 1) AS n
ON CONFLICT (token) DO NOTHING
""",
"""
INSERT INTO "LiteLLM_ProxyModelTable" (model_id, model_name, litellm_params, model_info, created_by, updated_by)
SELECT 'perf-model-' || n, 'perf-chat-' || (n %% 20),
jsonb_build_object('model', 'openai/perf-chat', 'api_key', 'perf-seed',
'api_base', 'os.environ/E2E_PERF_MOCK_OPENAI_URL'),
jsonb_build_object('id', 'perf-model-' || n, 'db_model', true), 'perf-seed', 'perf-seed'
FROM generate_series(0, %(models)s - 1) AS n
ON CONFLICT (model_id) DO NOTHING
""",
)
def prepare_fixture() -> UiFixture:
validate_target(os.environ["DATABASE_URL"])
_finish_upgrade()
@ -27,6 +66,7 @@ def prepare_fixture() -> UiFixture:
_ = connection.execute("SET statement_timeout='30s'")
row: Final = connection.execute("SELECT spec->>'anchor', rows FROM litellm_perf_seed.checkpoint").fetchone()
assert row is not None and TypeAdapter(int).validate_python(row[1]) > 0, "Missing completed seed rows"
seed_ui_tables(connection, anchor=TypeAdapter(datetime).validate_python(row[0]))
fixture: Final = UiFixture(
anchor=TypeAdapter(datetime).validate_python(row[0]),
api_key=hashlib.md5(b"perf-key:0", usedforsecurity=False).hexdigest()
@ -35,6 +75,36 @@ def prepare_fixture() -> UiFixture:
return fixture
def seed_ui_tables(connection: psycopg.Connection, anchor: datetime) -> None:
"""Give the Keys, Teams, Users and Models pages realistic row counts on the disposable restore.
Names follow the SpendLogs seed (perf-key-N aliases, perf-team-N, perf-user-N) so the Logs filters and
Team Usage line up with real spend. The Daily* rollups come from the snapshot itself and are only checked.
"""
window: Final = {
"since": (anchor - timedelta(days=7)).strftime("%Y-%m-%d"),
"until": anchor.strftime("%Y-%m-%d"),
}
for table in DAILY_TABLES:
row = connection.execute(f'SELECT count(*) FROM "{table}" WHERE date BETWEEN %(since)s AND %(until)s', window)
assert _count(row.fetchone()) > 0, f"{table} has no rows in the week before the anchor"
params: Final = {
"teams": SEEDED_TEAMS,
"users": SEEDED_USERS,
"keys": SEEDED_KEYS,
"models": SEEDED_MODELS,
"anchor": anchor,
}
for statement in _SEED_SQL:
_ = connection.execute(statement, params)
connection.commit()
def _count(row: tuple[object, ...] | None) -> int:
assert row is not None
return TypeAdapter(int).validate_python(row[0])
def _finish_upgrade() -> None:
if os.environ.get("E2E_DATABASE_UPGRADE") != "1":
return

4
tests/e2e/perf/setup/ui_query_isolation.py
View file

@ -62,7 +62,9 @@ def wait_idle(connection: psycopg.Connection[dict[str, object]], evidence: Path,
for query in queries:
if cancel and query.leader_pid is None:
_ = connection.execute(
"SELECT pg_cancel_backend(%s)" if query.state == "active" else "SELECT pg_terminate_backend(%s)",
"SELECT pg_cancel_backend(%s)"
if query.state == "active"
else "SELECT pg_terminate_backend(%s)",
(query.pid,),
)
else:

108
tests/e2e/perf/test_ui_journeys.py Normal file
View file

@ -0,0 +1,108 @@
from __future__ import annotations
import json
import os
from datetime import UTC, datetime
from pathlib import Path
from typing import Final
import pytest
from e2e_config import MASTER_KEY, UI_BASE_URL, UI_PASSWORD, UI_USERNAME
from e2e_http import send_text_external, without_retries
from setup.ui_journeys import (
assert_within_ceiling,
Dashboard,
Journey,
Sample,
dashboard_page,
database_failures,
error_code,
journeys,
measure,
summarize,
ui_junit_properties,
)
from setup.ui_query_fixture import UiFixture, prepare_fixture
from setup.ui_query_isolation import isolated_database, settled_counter
from transport import HttpTransport
pytestmark = [
pytest.mark.perf,
pytest.mark.quiet_stack,
pytest.mark.skipif(os.environ.get("E2E_UI_PERF_DISPOSABLE") != "1", reason="Dedicated disposable database only"),
]
def pytest_generate_tests(metafunc: pytest.Metafunc) -> None:
if "journey_name" in metafunc.fixturenames:
metafunc.parametrize("journey_name", [journey.name for journey in journeys("")])
@pytest.fixture(scope="module")
def ui_fixture() -> UiFixture:
return prepare_fixture()
def _metrics(transport: HttpTransport) -> float:
urls: Final = tuple(url.strip() for url in os.environ["E2E_DATABASE_METRICS_URLS"].split(",") if url.strip())
try:
assert urls, "Provide a metrics URL for every proxy worker/pod"
readings: Final = tuple(send_text_external("GET", url, headers=transport.master, timeout=10) for url in urls)
assert all(reading.status_code == 200 for reading in readings), "Cannot read database failure metrics"
return sum(database_failures(reading.body) for reading in readings)
except Exception as error:
pytest.exit(f"Database metrics unavailable ({type(error).__name__}). Stopping UI measurements", returncode=2)
def _samples(dashboard: Dashboard, journey: Journey, timeout: float, repeats: int) -> tuple[Sample, ...]:
with dashboard_page(dashboard, timeout) as page:
return tuple(measure(page, journey, timeout) for _ in range(repeats + 1))
def test_ui_journey_latency(journey_name: str, ui_fixture: UiFixture, request: pytest.FixtureRequest) -> None:
journey: Final = next(journey for journey in journeys(ui_fixture.api_key) if journey.name == journey_name)
transport: Final = HttpTransport(base_url=UI_BASE_URL, master_key=MASTER_KEY)
timeout: Final = float(os.environ.get("E2E_UI_JOURNEY_TIMEOUT_SECONDS", "120"))
repeats: Final = int(os.environ.get("E2E_UI_JOURNEY_REPEATS", "9"))
assert 0 < timeout <= 600 and 1 <= repeats <= 1000, "Invalid UI performance sample settings"
target: Final = Path(os.environ["E2E_ARTIFACT_DIR"])
assert target.is_dir(), "E2E_ARTIFACT_DIR must exist"
evidence: Final = target / f"ui-isolation-{journey.name}.jsonl"
dashboard: Final = Dashboard(
base_url=UI_BASE_URL, username=UI_USERNAME, password=UI_PASSWORD, anchor=ui_fixture.anchor
)
with without_retries():
with isolated_database(journey.name):
before: Final = settled_counter(lambda: _metrics(transport), evidence)
samples: Final = _samples(dashboard, journey, timeout, repeats)
with evidence.open("a", encoding="utf-8") as sample_output:
_ = sample_output.write(
json.dumps({"samples": tuple(sample.model_dump() for sample in samples)}) + "\n"
)
after_requests: Final = _metrics(transport)
item: Final = request.node # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # pytest: untyped
assert isinstance(item, pytest.Item)
item.user_properties.extend(ui_junit_properties(samples))
after: Final = settled_counter(lambda: _metrics(transport), evidence)
item.user_properties.extend(
(
("perf_ui_database_failures", after_requests - before),
("perf_ui_cleanup_database_failures", after - after_requests),
)
)
artifact: Final = {
"name": journey.name,
"surface": journey.surface,
"measured_at": datetime.now(UTC).isoformat(),
"samples": [sample.model_dump() for sample in samples],
"error_codes": [error_code(sample) for sample in samples],
"timings": summarize(samples).model_dump(),
"timeouts": sum(sample.error == "timeout" for sample in samples),
"failures": sum(bool(sample.error) for sample in samples),
"database_failures_before": before,
"database_failures_after_requests": after_requests,
"database_failures_after_cleanup": after,
}
output: Final = target / f"ui-journey-{journey.name}.json"
_ = output.write_text(json.dumps(artifact, indent=2), encoding="utf-8")
assert_within_ceiling(journey.name, samples)

159
tests/e2e/perf/test_ui_queries.py
View file

@ -1,159 +0,0 @@
from __future__ import annotations
import json
import os
from dataclasses import dataclass, replace
from datetime import UTC, datetime
from pathlib import Path
from typing import Final
import pytest
from e2e_config import CONTROL_PLANE_BASE_URL, MASTER_KEY
from e2e_http import send_text_external, without_retries
from setup.ui_queries import (
Sample,
UiQueryCase,
database_failures,
error_code,
issue_query,
measure,
summarize,
ui_junit_properties,
)
from setup.ui_query_cases import Params, cases, next_page
from setup.ui_query_fixture import UiFixture, prepare_fixture
from setup.ui_query_isolation import isolated_database, settled_counter
from transport import HttpTransport
pytestmark = [
pytest.mark.perf,
pytest.mark.quiet_stack,
pytest.mark.skipif(os.environ.get("E2E_UI_PERF_DISPOSABLE") != "1", reason="Dedicated disposable database only"),
]
def pytest_generate_tests(metafunc: pytest.Metafunc) -> None:
if "ui_case_name" not in metafunc.fixturenames:
return
if os.environ.get("E2E_UI_PERF_DISPOSABLE") == "1" and os.environ.get("E2E_RESET_SPEND_LOGS") == "1":
raise pytest.UsageError("UI performance tests cannot run with E2E_RESET_SPEND_LOGS=1")
names: Final = tuple(case.name for case in cases(datetime(2000, 1, 1, tzinfo=UTC), "collection-only"))
metafunc.parametrize("ui_case_name", names, ids=names)
@pytest.fixture(scope="module")
def ui_fixture() -> UiFixture:
return prepare_fixture()
def _metrics(transport: HttpTransport) -> float:
urls: Final = tuple(url.strip() for url in os.environ["E2E_DATABASE_METRICS_URLS"].split(",") if url.strip())
try:
assert urls, "Provide a metrics URL for every proxy worker/pod"
readings: Final = tuple(send_text_external("GET", url, headers=transport.master, timeout=10) for url in urls)
assert all(reading.status_code == 200 for reading in readings), "Cannot read database failure metrics"
return sum(database_failures(reading.body) for reading in readings)
except Exception as error:
pytest.exit(f"Database metrics unavailable ({type(error).__name__}). Stopping UI measurements", returncode=2)
@dataclass(frozen=True, slots=True)
class Prepared:
case: UiQueryCase
prerequisite: Sample | None = None
def _page_two(transport: HttpTransport, case: UiQueryCase, timeout: float) -> Prepared:
if case.cursor == "none":
return Prepared(case)
params: Final = Params.model_validate(case.params.model_dump()).model_copy(update={"page": 1})
first: Final = replace(case, params=params)
sample, document = measure(lambda: issue_query(transport, first, timeout=timeout), case.contract)
if sample.error or document is None:
return Prepared(case, sample)
try:
return Prepared(next_page(case, document), sample)
except AssertionError as error:
return Prepared(case, sample.model_copy(update={"error": str(error)}))
def _samples(transport: HttpTransport, case: UiQueryCase, timeout: float, repeats: int) -> tuple[Sample, ...]:
measured, _ = measure(lambda: issue_query(transport, case, timeout=timeout), case.contract)
if measured.error:
return (measured,)
repeated: Final = tuple(
measure(lambda: issue_query(transport, case, timeout=timeout), case.contract)[0] for _ in range(repeats)
)
return (measured, *repeated)
def test_ui_query_performance(ui_case_name: str, ui_fixture: UiFixture, request: pytest.FixtureRequest) -> None:
case: Final = next(case for case in cases(ui_fixture.anchor, ui_fixture.api_key) if case.name == ui_case_name)
transport: Final = HttpTransport(base_url=CONTROL_PLANE_BASE_URL, master_key=MASTER_KEY)
timeout: Final = float(os.environ.get("E2E_UI_QUERY_TIMEOUT_SECONDS", "120"))
repeats: Final = int(os.environ.get("E2E_UI_QUERY_REPEATS", "9"))
assert 0 < timeout <= 600 and 1 <= repeats <= 1000, "Invalid UI performance sample settings"
target: Final = Path(os.environ["E2E_ARTIFACT_DIR"])
assert target.is_dir(), "E2E_ARTIFACT_DIR must exist"
evidence: Final = target / f"ui-isolation-{case.name}.jsonl"
with without_retries():
with isolated_database(case.name):
before: Final = settled_counter(lambda: _metrics(transport), evidence)
paginated: Final = _page_two(transport, case, timeout)
samples: Final = (
(paginated.prerequisite,)
if paginated.prerequisite is not None and paginated.prerequisite.error
else _samples(transport, paginated.case, timeout, repeats)
)
with evidence.open("a", encoding="utf-8") as sample_output:
_ = sample_output.write(
json.dumps(
{
"samples": tuple(sample.model_dump() for sample in samples),
"prerequisite": paginated.prerequisite.model_dump()
if paginated.prerequisite is not None
else None,
}
)
+ "\n"
)
after_requests: Final = _metrics(transport)
item: Final = request.node # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # pytest: untyped
assert isinstance(item, pytest.Item)
item.user_properties.extend(ui_junit_properties(samples, paginated.prerequisite))
after: Final = settled_counter(lambda: _metrics(transport), evidence)
item.user_properties.extend(
(
("perf_ui_database_failures", after_requests - before),
("perf_ui_cleanup_database_failures", after - after_requests),
)
)
timings: Final = summarize(samples)
artifact: Final = {
"name": case.name,
"surface": case.surface,
"path": case.path,
"method": "POST" if case.body is not None else "GET",
"params": paginated.case.params.model_dump(mode="json", by_alias=True, exclude_none=True),
"prerequisite": paginated.prerequisite.model_dump() if paginated.prerequisite is not None else None,
"prerequisite_failed": paginated.prerequisite is not None and bool(paginated.prerequisite.error),
"body": case.body.model_dump(mode="json", exclude_none=True) if case.body is not None else None,
"samples": [sample.model_dump() for sample in samples],
"error_codes": [error_code(sample) for sample in samples],
"timings": timings.model_dump(),
"timeouts": sum(sample.error == "timeout" for sample in samples),
"failures": sum(bool(sample.error) for sample in samples),
"status_counts": {
str(status): sum(sample.status == status for sample in samples)
for status in frozenset(sample.status for sample in samples)
},
"validation_failures": sum(
bool(sample.error) and (sample.status == 200 or sample.error == "Response is not valid JSON")
for sample in samples
),
"database_failures_before": before,
"database_failures_after_requests": after_requests,
"database_failures_after_cleanup": after,
}
output: Final = target / f"ui-query-{case.name}.json"
_ = output.write_text(json.dumps(artifact, indent=2), encoding="utf-8")

144
tests/e2e_harness/perf/setup/test_ui_journeys.py Normal file
View file

@ -0,0 +1,144 @@
from __future__ import annotations
from typing import Final
import pytest
from setup.ui_journeys import (
UI_JOURNEY_LIMITS,
Sample,
assert_within_ceiling,
database_failures,
journeys,
summarize,
ui_junit_properties,
)
def test_first_journey_is_separate_from_repeated_latency() -> None:
samples: Final = (Sample(seconds=80), *(Sample(seconds=i) for i in range(1, 101)))
result: Final = summarize(samples)
assert result.first_seconds == 80
assert result.repeated_samples == 100
assert result.repeated_p50_seconds == 50.5
assert result.repeated_p95_seconds == 95
assert result.repeated_p99_seconds == 99
assert result.repeated_max_seconds == 100
def test_failed_repeats_do_not_look_like_fast_successes() -> None:
result: Final = summarize((Sample(seconds=10), Sample(seconds=0.01, error="AssertionError: no rows")))
assert result.repeated_samples == 0
assert result.repeated_p50_seconds is None
assert result.repeated_p95_seconds is None
assert result.repeated_p99_seconds is None
assert result.repeated_max_seconds is None
def test_small_sample_does_not_report_tail_percentiles() -> None:
result: Final = summarize(tuple(Sample(seconds=1) for _ in range(6)))
assert result.repeated_samples == 5
assert result.repeated_p50_seconds == 1
assert result.repeated_p95_seconds is None
assert result.repeated_p99_seconds is None
def test_failed_first_journey_emits_error_code_without_latency() -> None:
sample: Final = Sample(seconds=120, error="timeout")
assert summarize((sample,)).first_seconds is None
assert dict(ui_junit_properties((sample,))) == {
"perf_ui_journeys": 1,
"perf_ui_failures": 1,
"perf_ui_samples": 0,
"perf_ui_repeated_samples": 0,
"perf_ui_error_timeout": 1,
}
def test_junit_counts_failures_but_only_measures_successful_latency() -> None:
samples: Final = (
Sample(seconds=2, rows=25),
Sample(seconds=3, rows=25),
Sample(seconds=120, error="timeout"),
Sample(seconds=1, error="AssertionError: Logs table rendered empty"),
Sample(seconds=4, rows=25),
)
assert dict(ui_junit_properties(samples)) == {
"perf_ui_journeys": 5,
"perf_ui_failures": 2,
"perf_ui_samples": 3,
"perf_ui_repeated_samples": 2,
"perf_ui_error_timeout": 1,
"perf_ui_error_ui": 1,
"perf_ui_seconds_p50": 3,
"perf_ui_seconds_p90": 4,
"perf_ui_seconds_max": 4,
"perf_ui_first_seconds": 2,
"perf_ui_repeated_seconds_p50": 3.5,
"perf_ui_repeated_seconds_max": 4,
}
def test_all_journey_percentiles_include_first_and_use_nearest_rank_p90() -> None:
samples: Final = tuple(Sample(seconds=seconds) for seconds in (10, 1, 9, 2, 8, 3, 7, 4, 6, 5))
metrics: Final = dict(ui_junit_properties(samples))
assert {key: value for key, value in metrics.items() if key.startswith("perf_ui_seconds_")} == {
"perf_ui_seconds_p50": 5.5,
"perf_ui_seconds_p90": 9,
"perf_ui_seconds_max": 10,
}
assert metrics["perf_ui_samples"] == 10
def test_journeys_have_unique_names_and_cover_every_dashboard_surface() -> None:
names: Final = tuple(journey.name for journey in journeys("key"))
assert len(names) == len(frozenset(names))
assert frozenset(journey.surface for journey in journeys("key")) == {"logs", "usage", "caching", "models", "keys"}
def _metrics(postgres: str, batch: str) -> str:
return (
"# TYPE litellm_postgres_failed_requests_total counter\n"
f'litellm_postgres_failed_requests_total{{call_type="query_raw"}} {postgres}\n'
"# TYPE litellm_batch_write_to_db_failed_requests_total counter\n"
f'litellm_batch_write_to_db_failed_requests_total{{call_type="batch"}} {batch}\n'
)
def test_postgres_failures_include_background_batch_writes() -> None:
assert database_failures(_metrics("0", "0")) == 0
assert database_failures(_metrics("2", "3")) == 5
def test_lazy_failure_counters_require_evidence_of_database_activity() -> None:
assert (
database_failures(
"# TYPE litellm_postgres_total_requests_total counter\nlitellm_postgres_total_requests_total 3\n"
)
== 0
)
@pytest.mark.parametrize(
"text", ("", "# TYPE unrelated counter\nunrelated_total 0\n", _metrics("NaN", "0"), _metrics("-1", "0"))
)
def test_missing_or_invalid_metrics_cannot_pass_as_zero(text: str) -> None:
with pytest.raises(AssertionError):
_ = database_failures(text)
def test_every_journey_has_a_latency_ceiling() -> None:
assert frozenset(journey.name for journey in journeys("key")) == frozenset(UI_JOURNEY_LIMITS)
def test_failed_sample_fails_the_journey_before_latency_is_judged() -> None:
with pytest.raises(AssertionError, match="1 of 2 logs_open_24h samples failed: timeout"):
assert_within_ceiling("logs_open_24h", (Sample(seconds=0.2, rows=25), Sample(seconds=120, error="timeout")))
def test_slow_journey_fails_its_ceiling() -> None:
fast: Final = tuple(Sample(seconds=1, rows=25) for _ in range(9))
assert_within_ceiling("logs_open_24h", fast)
with pytest.raises(AssertionError, match=r"logs_open_24h p90 25\.0s over 20s"):
assert_within_ceiling("logs_open_24h", (*fast[:8], Sample(seconds=25, rows=25), Sample(seconds=25, rows=25)))
with pytest.raises(AssertionError, match=r"logs_open_24h p50 11\.0s over 10s"):
assert_within_ceiling("logs_open_24h", tuple(Sample(seconds=11, rows=25) for _ in range(3)))

257
tests/e2e_harness/perf/setup/test_ui_queries.py
View file

@ -1,257 +0,0 @@
from __future__ import annotations
from datetime import UTC, datetime
from typing import Final
from urllib.parse import parse_qs, urlsplit
import pytest
from e2e_http import NetworkError, Result, Success, UnknownApiError
from setup.ui_queries import (
Contract,
Document,
Sample,
database_failures,
measure,
query_path,
summarize,
ui_junit_properties,
)
from setup.ui_query_cases import Params, cases, next_page
def test_query_string_preserves_ui_arrays_and_boolean_strings() -> None:
requests: Final = cases(datetime(2026, 1, 1, tzinfo=UTC), "key")
cache: Final = next(case for case in requests if case.name == "cache_activity_filtered")
assert parse_qs(urlsplit(query_path(cache)).query) == {
"start_date": ["2025-12-25"],
"end_date": ["2026-01-01"],
"key_aliases": ["perf-key-0", "perf-key-1"],
"models": ["perf-chat"],
}
health: Final = next(case for case in requests if case.name == "logs_exclude_internal_health_checks_30d")
query: Final = parse_qs(urlsplit(query_path(health)).query)
assert query["exclude_internal_health_checks"] == ["true"]
assert query["group_by_session"] == ["true"]
@pytest.mark.parametrize("entity", ("user", "team", "tag", "customer"))
@pytest.mark.parametrize(
("route", "extra"),
(
("aggregated", {}),
("aggregated/keys", {"offset": 0, "limit": 50}),
("aggregated/search", {"search": "perf-key"}),
("aggregated/model_top_keys", {"model_group": "perf-chat", "by_model_group": "true"}),
),
)
def test_usage_requests_list_keys_without_a_single_key_filter(
entity: str, route: str, extra: dict[str, str | int]
) -> None:
case: Final = next(
case
for case in cases(datetime(2026, 1, 1, tzinfo=UTC), "key")
if case.path == f"/{entity}/daily/activity/{route}"
)
assert case.params.model_dump(exclude_none=True) == {
"start_date": "2025-12-25",
"end_date": "2026-01-01",
"timezone": "0",
**({"exclude_team_ids": "litellm-dashboard"} if entity == "team" else {}),
**extra,
}
@pytest.mark.parametrize(("name", "value"), (("api_key", "key"), ("key_alias", "perf-key-0")))
def test_key_filters_request_30_days_of_logs(name: str, value: str) -> None:
case: Final = next(
case for case in cases(datetime(2026, 1, 1, tzinfo=UTC), "key") if case.name == f"logs_{name}_30d"
)
assert parse_qs(urlsplit(query_path(case)).query) == {
"start_date": ["2025-12-02 00:00:00"],
"end_date": ["2026-01-01 00:00:00"],
"page_size": ["25"],
"page": ["1"],
"sort_by": ["startTime"],
"sort_order": ["desc"],
"group_by_session": ["true"],
name: [value],
}
@pytest.mark.parametrize("invalid", ("not-a-number", True, float("nan"), float("inf")))
def test_positive_row_does_not_hide_a_malformed_later_row(invalid: str | bool | float) -> None:
with pytest.raises(AssertionError, match=r"Expected (numeric|finite) spend"):
Contract(positive=("spend",)).count(Document([{"spend": 1}, {"spend": invalid}]))
def test_session_cursor_is_reused_and_repeated_rows_rejected() -> None:
case: Final = next(case for case in cases(datetime(2026, 1, 1, tzinfo=UTC), "key") if case.cursor == "session")
first: Final = Document.model_validate_json('{"data":[{"request_id":"one"}],"next_session_cursor":"opaque"}')
second: Final = next_page(case, first)
assert Params.model_validate(second.params).session_cursor == "opaque"
assert Params.model_validate(second.params).page == 2
with pytest.raises(AssertionError, match="Repeated page-one row"):
_ = second.contract.count(first)
def test_first_query_is_separate_from_repeated_latency() -> None:
samples: Final = (Sample(seconds=80, status=200), *(Sample(seconds=i, status=200) for i in range(1, 101)))
result: Final = summarize(samples)
assert result.first_seconds == 80
assert result.repeated_samples == 100
assert result.repeated_p50_seconds == 50.5
assert result.repeated_p95_seconds == 95
assert result.repeated_p99_seconds == 99
assert result.repeated_max_seconds == 100
def test_failed_repeats_do_not_look_like_fast_successes() -> None:
result: Final = summarize((Sample(seconds=10, status=200), Sample(seconds=0.01, status=503, error="HTTP 503")))
assert result.repeated_samples == 0
assert result.repeated_p50_seconds is None
assert result.repeated_p95_seconds is None
assert result.repeated_p99_seconds is None
assert result.repeated_max_seconds is None
def test_small_sample_does_not_report_tail_percentiles() -> None:
result: Final = summarize(tuple(Sample(seconds=1, status=200) for _ in range(6)))
assert result.repeated_samples == 5
assert result.repeated_p50_seconds == 1
assert result.repeated_p95_seconds is None
assert result.repeated_p99_seconds is None
def test_failed_first_request_emits_error_code_without_latency() -> None:
sample: Final = Sample(seconds=30, status=500, error="HTTP 500")
assert summarize((sample,)).first_seconds is None
assert dict(ui_junit_properties((sample,))) == {
"perf_ui_requests": 1,
"perf_ui_failures": 1,
"perf_ui_samples": 0,
"perf_ui_repeated_samples": 0,
"perf_ui_error_http_500": 1,
}
def test_junit_counts_failures_but_only_measures_successful_latency() -> None:
samples: Final = (
Sample(seconds=2, status=200),
Sample(seconds=3, status=200),
Sample(seconds=30, status=-1, error="timeout"),
Sample(seconds=30, status=503, error="HTTP 503"),
Sample(seconds=1, status=200, error="Empty collection"),
)
assert dict(ui_junit_properties(samples, Sample(seconds=4, status=200))) == {
"perf_ui_requests": 6,
"perf_ui_failures": 3,
"perf_ui_samples": 2,
"perf_ui_seconds_p50": 2.5,
"perf_ui_seconds_p90": 3,
"perf_ui_seconds_max": 3,
"perf_ui_repeated_samples": 1,
"perf_ui_error_timeout": 1,
"perf_ui_error_http_503": 1,
"perf_ui_error_contract": 1,
"perf_ui_first_seconds": 2,
"perf_ui_repeated_seconds_p50": 3,
"perf_ui_repeated_seconds_max": 3,
}
def test_failed_cursor_prerequisite_is_counted_once() -> None:
sample: Final = Sample(seconds=30, status=503, error="HTTP 503")
assert dict(ui_junit_properties((sample,), sample)) == dict(ui_junit_properties((sample,)))
def test_all_request_percentiles_include_first_and_use_nearest_rank_p90() -> None:
samples: Final = tuple(Sample(seconds=seconds, status=200) for seconds in (10, 1, 9, 2, 8, 3, 7, 4, 6, 5))
metrics: Final = dict(ui_junit_properties(samples))
assert {key: value for key, value in metrics.items() if key.startswith("perf_ui_seconds_")} == {
"perf_ui_seconds_p50": 5.5,
"perf_ui_seconds_p90": 9,
"perf_ui_seconds_max": 10,
}
assert metrics["perf_ui_samples"] == 10
@pytest.mark.parametrize("body", ("{}", '{"data": []}', '{"data": [{"spend": 1}]}', '{"data": [null]}'))
def test_empty_or_malformed_rows_fail_even_with_http_200(body: str) -> None:
ticks: Final = iter((10.0, 12.5))
result, _ = measure(
lambda: Success(status_code=200, data=Document.model_validate_json(body)),
Contract(("data",), ("request_id",)),
clock=lambda: next(ticks),
)
assert result.status == 200
assert result.seconds == 2.5
assert result.rows == 0
assert result.error
def test_valid_response_records_elapsed_seconds_and_rows() -> None:
ticks: Final = iter((5.0, 5.125))
result, document = measure(
lambda: Success(status_code=200, data=Document.model_validate_json('{"data":[{"request_id":"one"}]}')),
Contract(("data",), ("request_id",)),
clock=lambda: next(ticks),
)
assert result == Sample(seconds=0.125, status=200, rows=1)
assert document is not None
def test_zero_spend_rows_are_allowed_alongside_nonzero_fixture_data() -> None:
contract: Final = Contract(positive=("spend",))
assert contract.count(Document.model_validate_json('[{"spend":0},{"spend":2}]')) == 2
with pytest.raises(AssertionError, match="positive spend"):
_ = contract.count(Document.model_validate_json('[{"spend":0}]'))
@pytest.mark.parametrize(
"failure,error,status",
(
(UnknownApiError(status_code=500, body="sensitive error detail"), "HTTP 500", 500),
(NetworkError(message="Read timed out"), "timeout", -1),
(NetworkError(message="connection refused"), "network", -1),
),
)
def test_failed_requests_are_recorded_without_response_bodies(
failure: Result[Document],
error: str,
status: int,
) -> None:
ticks: Final = iter((1.0, 2.0))
result, document = measure(lambda: failure, Contract(), clock=lambda: next(ticks))
assert result == Sample(seconds=1, status=status, error=error)
assert document is None
def _metrics(postgres: str, batch: str) -> str:
return (
"# TYPE litellm_postgres_failed_requests_total counter\n"
f'litellm_postgres_failed_requests_total{{call_type="query_raw"}} {postgres}\n'
"# TYPE litellm_batch_write_to_db_failed_requests_total counter\n"
f'litellm_batch_write_to_db_failed_requests_total{{call_type="batch"}} {batch}\n'
)
def test_postgres_failures_include_background_batch_writes() -> None:
assert database_failures(_metrics("0", "0")) == 0
assert database_failures(_metrics("2", "3")) == 5
def test_lazy_failure_counters_require_evidence_of_database_activity() -> None:
assert (
database_failures(
"# TYPE litellm_postgres_total_requests_total counter\nlitellm_postgres_total_requests_total 3\n"
)
== 0
)
@pytest.mark.parametrize(
"text", ("", "# TYPE unrelated counter\nunrelated_total 0\n", _metrics("NaN", "0"), _metrics("-1", "0"))
)
def test_missing_or_invalid_metrics_cannot_pass_as_zero(text: str) -> None:
with pytest.raises(AssertionError):
_ = database_failures(text)

4
tests/e2e_harness/perf/test_upgrade.py
View file

@ -308,9 +308,7 @@ def test_invalid_200_is_not_an_allowed_non_200(failed: WireRecord) -> None:
def test_upgrade_verdict_uses_regular_limits_without_cpu_or_datastore_samples(
failures: int, overhead: float, ttft: float, growth: float, expected: tuple[str, ...]
) -> None:
nonstream: Final = tuple(
record(i).model_copy(update={"total_s": overhead, "mock_delay_s": 0}) for i in range(100)
)
nonstream: Final = tuple(record(i).model_copy(update={"total_s": overhead, "mock_delay_s": 0}) for i in range(100))
stream: Final = record(100).model_copy(update={"stream": True, "first_token_s": ttft})
observed: Final = Observation(
records=(*nonstream, stream, *(failed_record(i) for i in range(101, 101 + failures))),

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