mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-28 06:23:07 +00:00
Merge stable/1.104.x (v1.104.6 prep) into patch_d351e54e
Keeps stable/1.104.x's 1.104.6 version in pyproject.toml and uv.lock, so this branch no longer bumps the version. In litellm_logging.py both async success helpers stay, the base's enqueued flag and the claim that #44508 adds, matching main. The two test files keep both sides' new tests
This commit is contained in:
commit
2b2e960848
144 changed files with 35645 additions and 24214 deletions
|
|
@ -91,6 +91,17 @@ OBJECT_KEYS: dict[str, JsonSchema] = {
|
|||
},
|
||||
"additionalProperties": False,
|
||||
},
|
||||
"search_free_requests": {
|
||||
"type": "object",
|
||||
"description": "Free Google Search grounding requests per period, shared across models in the same shared_group.",
|
||||
"required": ["count", "period", "shared_group"],
|
||||
"additionalProperties": False,
|
||||
"properties": {
|
||||
"count": {"type": "integer", "minimum": 0},
|
||||
"period": {"type": "string", "enum": ["day", "month"]},
|
||||
"shared_group": {"type": "string", "pattern": "\\S"},
|
||||
},
|
||||
},
|
||||
"guardrail_cost_per_unit": {
|
||||
"type": "object",
|
||||
"description": "USD cost per billable guardrail unit, keyed by the provider's usage counter name (e.g. Bedrock's contentPolicyUnits).",
|
||||
|
|
@ -111,8 +122,9 @@ OBJECT_KEYS: dict[str, JsonSchema] = {
|
|||
"properties": {
|
||||
"model": NON_BLANK_STRING,
|
||||
"region": NON_BLANK_STRING,
|
||||
"litellm_provider": {"type": "string", "enum": ["bedrock", "bedrock_mantle"]},
|
||||
},
|
||||
"description": "Model that counts tokens for this one, and the AWS region its count API runs in.",
|
||||
"description": "Model that counts tokens for this one, its counting provider, and the AWS region its count API runs in.",
|
||||
},
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -2,7 +2,9 @@
|
|||
Polls LiteLLM_ManagedObjectTable to check if the batch job is complete, and if the cost has been tracked.
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
import json
|
||||
import re
|
||||
from collections.abc import Mapping, Sequence
|
||||
from dataclasses import replace as dataclasses_replace
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from types import MappingProxyType
|
||||
|
|
@ -27,7 +29,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",
|
||||
|
|
@ -44,6 +51,12 @@ TERMINAL_MANAGED_OBJECT_STATUSES: Final[Tuple[str, ...]] = (
|
|||
|
||||
|
||||
class _ManagedObjectRow(Protocol):
|
||||
@property
|
||||
def created_at(self) -> datetime: ...
|
||||
|
||||
@property
|
||||
def updated_at(self) -> datetime: ...
|
||||
|
||||
@property
|
||||
def id(self) -> str: ...
|
||||
|
||||
|
|
@ -98,6 +111,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
|
||||
|
||||
@staticmethod
|
||||
def _is_missing_batch_processed_column_error(err: Exception) -> bool:
|
||||
|
|
@ -269,29 +287,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": {
|
||||
|
|
@ -305,10 +303,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
|
||||
|
|
@ -316,27 +363,23 @@ 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"
|
||||
)
|
||||
|
||||
async def _claim_job_for_costing(self, job: "_ManagedObjectRow") -> bool:
|
||||
async def _claim_job_for_costing(self, job: "_ManagedObjectRow", response: "LiteLLMBatch") -> bool:
|
||||
"""
|
||||
Atomically flip batch_processed from false to true, returning whether this pod won
|
||||
the row. Every pod and uvicorn worker schedules its own poller against the shared
|
||||
|
|
@ -350,10 +393,16 @@ class CheckBatchCost:
|
|||
"""
|
||||
if not self._has_batch_processed_column:
|
||||
return True
|
||||
from litellm_enterprise.proxy.hooks.managed_files import postgres_storable_json
|
||||
|
||||
try:
|
||||
claimed: Final = await _managed_object_table(self.prisma_client).update_many(
|
||||
where={"id": job.id, "batch_processed": False},
|
||||
data={"batch_processed": True},
|
||||
data={
|
||||
"batch_processed": True,
|
||||
"status": response.status if response.status != "completed" else "complete",
|
||||
"file_object": postgres_storable_json(response),
|
||||
},
|
||||
)
|
||||
except Exception as db_err:
|
||||
verbose_proxy_logger.error(
|
||||
|
|
@ -362,7 +411,9 @@ class CheckBatchCost:
|
|||
return False
|
||||
return claimed > 0
|
||||
|
||||
async def _release_job_claim(self, job: "_ManagedObjectRow") -> None:
|
||||
async def _release_job_claim(
|
||||
self, job: "_ManagedObjectRow", original_status: str | None, original_file_object: object
|
||||
) -> None:
|
||||
"""Give a claimed row back once billing it failed, so a later poll cycle retries it.
|
||||
|
||||
Safe to match on batch_processed=True: while this poller is active the retrieve
|
||||
|
|
@ -374,7 +425,11 @@ class CheckBatchCost:
|
|||
try:
|
||||
await _managed_object_table(self.prisma_client).update_many(
|
||||
where={"id": job.id, "batch_processed": True},
|
||||
data={"batch_processed": False},
|
||||
data={
|
||||
"batch_processed": False,
|
||||
"status": original_status,
|
||||
"file_object": json.dumps(original_file_object),
|
||||
},
|
||||
)
|
||||
except Exception as db_err:
|
||||
verbose_proxy_logger.error(
|
||||
|
|
@ -408,13 +463,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
|
||||
|
|
@ -429,7 +484,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"
|
||||
|
|
@ -451,15 +506,17 @@ 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,
|
||||
)
|
||||
if finalized == 0:
|
||||
return
|
||||
verbose_proxy_logger.info(
|
||||
f"CheckBatchCost: marked job {job.id} as {response.status} in DB"
|
||||
)
|
||||
|
|
@ -671,8 +728,6 @@ class CheckBatchCost:
|
|||
|
||||
@staticmethod
|
||||
def _get_input_file_id(job: "_ManagedObjectRow") -> Optional[str]:
|
||||
import json
|
||||
|
||||
from litellm.types.utils import LiteLLMBatch
|
||||
|
||||
file_object = job.file_object
|
||||
|
|
@ -712,7 +767,11 @@ class CheckBatchCost:
|
|||
from litellm.files.main import afile_content
|
||||
from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
|
||||
from litellm.litellm_core_utils.litellm_logging import deployment_pricing_model_info, mask_api_base_credentials
|
||||
from litellm.litellm_core_utils.litellm_logging import (
|
||||
deployment_pricing_model_info,
|
||||
mask_api_base_credentials,
|
||||
resolve_vertex_location_for_cost,
|
||||
)
|
||||
from litellm.proxy.openai_files_endpoints.common_utils import (
|
||||
_is_base64_encoded_unified_file_id,
|
||||
)
|
||||
|
|
@ -849,6 +908,12 @@ class CheckBatchCost:
|
|||
custom_llm_provider=batch_file_provider,
|
||||
model_name=model_name,
|
||||
model_info=deployment_model_info,
|
||||
vertex_location=resolve_vertex_location_for_cost(
|
||||
custom_llm_provider=llm_provider,
|
||||
litellm_params=deployment_info.litellm_params.model_dump(exclude_none=True),
|
||||
optional_params=None,
|
||||
model=model_name,
|
||||
),
|
||||
)
|
||||
error_file_failed_requests: Final = await count_error_file_failed_requests(
|
||||
response,
|
||||
|
|
@ -866,8 +931,9 @@ class CheckBatchCost:
|
|||
failed_requests=output_file_result.failed_requests + error_file_failed_requests,
|
||||
)
|
||||
)
|
||||
logging_model: Final = batch_result.models[0] if batch_result.models else litellm_model_name or model_id
|
||||
logging_obj = LiteLLMLogging(
|
||||
model=batch_result.models[0],
|
||||
model=logging_model,
|
||||
messages=[{"role": "user", "content": "<retrieve_batch>"}],
|
||||
stream=False,
|
||||
call_type="aretrieve_batch",
|
||||
|
|
@ -898,27 +964,35 @@ class CheckBatchCost:
|
|||
custom_llm_provider=str(llm_provider) if llm_provider else None,
|
||||
)
|
||||
|
||||
if not await self._claim_job_for_costing(job):
|
||||
original_status: Final = job.status
|
||||
original_file_object: Final = job.file_object
|
||||
if not await self._claim_job_for_costing(job, response.model_copy(update={"usage": batch_result.usage})):
|
||||
verbose_proxy_logger.info(
|
||||
f"CheckBatchCost: batch {batch_id} (job {job.id}) was claimed by another pod "
|
||||
"in this window, so its cost is already being tracked there"
|
||||
)
|
||||
return None
|
||||
|
||||
try:
|
||||
await logging_obj.async_success_handler(
|
||||
result=response,
|
||||
batch_cost=batch_result.cost,
|
||||
batch_usage=batch_result.usage,
|
||||
batch_models=batch_result.models,
|
||||
batch_successful_requests=batch_result.successful_requests,
|
||||
batch_failed_requests=batch_result.failed_requests,
|
||||
batch_prompt_cost=batch_result.prompt_cost,
|
||||
batch_completion_cost=batch_result.completion_cost,
|
||||
)
|
||||
except Exception:
|
||||
await self._release_job_claim(job)
|
||||
raise
|
||||
from litellm.proxy.hooks.proxy_track_cost_callback import track_batch_accounting
|
||||
|
||||
with track_batch_accounting(logging_obj.litellm_call_id) as accounting:
|
||||
try:
|
||||
await logging_obj.async_success_handler(
|
||||
result=response,
|
||||
batch_cost=batch_result.cost,
|
||||
batch_usage=batch_result.usage,
|
||||
batch_models=batch_result.models,
|
||||
batch_successful_requests=batch_result.successful_requests,
|
||||
batch_failed_requests=batch_result.failed_requests,
|
||||
batch_prompt_cost=batch_result.prompt_cost,
|
||||
batch_completion_cost=batch_result.completion_cost,
|
||||
)
|
||||
finally:
|
||||
if accounting.retryable_failure:
|
||||
await self._release_job_claim(job, original_status, original_file_object)
|
||||
|
||||
if accounting.retryable_failure:
|
||||
return None
|
||||
|
||||
# Record batch duration (completed_at - created_at)
|
||||
if prom_logger and response.completed_at and response.created_at:
|
||||
|
|
@ -966,7 +1040,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,
|
||||
|
|
@ -979,8 +1053,6 @@ class CheckBatchCost:
|
|||
]
|
||||
},
|
||||
},
|
||||
take=MAX_OBJECTS_PER_POLL_CYCLE,
|
||||
order={"created_at": "asc"},
|
||||
)
|
||||
self.batch_processed_support_confirmed = True
|
||||
except Exception as query_err:
|
||||
|
|
@ -1050,7 +1122,7 @@ class CheckBatchCost:
|
|||
continue
|
||||
verbose_proxy_logger.error(
|
||||
f"CheckBatchCost: failed to track cost for batch {batch_id} "
|
||||
f"(job {job.id}); leaving it unprocessed so the next poll retries: {tracking_err}"
|
||||
f"(job {job.id}): {tracking_err}"
|
||||
)
|
||||
self._record_error(prom_logger, "cost_tracking_error")
|
||||
continue
|
||||
|
|
@ -1060,22 +1132,19 @@ class CheckBatchCost:
|
|||
# Track this job for the final metrics summary
|
||||
processed_models.append(tracked)
|
||||
|
||||
# mark the job as complete
|
||||
try:
|
||||
update_data: dict = {
|
||||
"status": response.status if response.status != "completed" else "complete",
|
||||
"file_object": response.model_dump_json(),
|
||||
}
|
||||
if self._has_batch_processed_column:
|
||||
update_data["batch_processed"] = True
|
||||
await _managed_object_table(self.prisma_client).update(
|
||||
where={"id": job.id},
|
||||
data=update_data,
|
||||
)
|
||||
except Exception as db_err:
|
||||
verbose_proxy_logger.error(
|
||||
f"CheckBatchCost: failed to mark job {job.id} complete in DB: {db_err}"
|
||||
)
|
||||
if not self._has_batch_processed_column:
|
||||
try:
|
||||
await _managed_object_table(self.prisma_client).update(
|
||||
where={"id": job.id},
|
||||
data={
|
||||
"status": response.status if response.status != "completed" else "complete",
|
||||
"file_object": response.model_dump_json(),
|
||||
},
|
||||
)
|
||||
except Exception as db_err:
|
||||
verbose_proxy_logger.error(
|
||||
f"CheckBatchCost: failed to mark job {job.id} complete in DB: {db_err}"
|
||||
)
|
||||
|
||||
elif response.status in PROVIDER_TERMINAL_BATCH_STATUSES:
|
||||
from litellm.proxy.openai_files_endpoints.common_utils import (
|
||||
|
|
|
|||
|
|
@ -22,7 +22,7 @@ from uuid import NAMESPACE_URL, uuid5
|
|||
|
||||
import httpx
|
||||
from fastapi import HTTPException
|
||||
from pydantic import ValidationError
|
||||
from pydantic import BaseModel, JsonValue, TypeAdapter, ValidationError
|
||||
|
||||
import litellm
|
||||
from litellm import Router, verbose_logger
|
||||
|
|
@ -205,6 +205,31 @@ def _managed_object_table(prisma_client: PrismaClient) -> _ManagedObjectTableAct
|
|||
return prisma_client.db.litellm_managedobjecttable
|
||||
|
||||
|
||||
_JSON_VALUE_ADAPTER: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue)
|
||||
_POSTGRES_UNSTORABLE_CHARACTER: Final = "\x00"
|
||||
|
||||
|
||||
def _storable_text(text: str) -> str:
|
||||
return text.replace(_POSTGRES_UNSTORABLE_CHARACTER, "")
|
||||
|
||||
|
||||
def _storable_json_value(value: JsonValue) -> JsonValue:
|
||||
if isinstance(value, str):
|
||||
return _storable_text(value)
|
||||
if isinstance(value, list):
|
||||
return [_storable_json_value(item) for item in value]
|
||||
if isinstance(value, dict):
|
||||
return {_storable_text(key): _storable_json_value(item) for key, item in value.items()}
|
||||
return value
|
||||
|
||||
|
||||
def postgres_storable_json(file_object: BaseModel) -> str:
|
||||
dumped: Final = file_object.model_dump_json()
|
||||
if "\\u0000" not in dumped:
|
||||
return dumped
|
||||
return _JSON_VALUE_ADAPTER.dump_json(_storable_json_value(_JSON_VALUE_ADAPTER.validate_json(dumped))).decode()
|
||||
|
||||
|
||||
def _storage_metadata_of(file_object: OpenAIFileObject | None) -> Mapping[str, str]:
|
||||
hidden_params: Final = cast( # cast-ok: _hidden_params is an untyped attribute the upload path sets
|
||||
"Mapping[str, object]", getattr(file_object, "_hidden_params", None) or {}
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
[project]
|
||||
name = "litellm-enterprise"
|
||||
version = "0.1.71"
|
||||
version = "0.1.71.post1"
|
||||
description = "Package for LiteLLM Enterprise features"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.9"
|
||||
|
|
@ -26,7 +26,7 @@ required-version = ">=0.10.9"
|
|||
module-root = ""
|
||||
|
||||
[tool.commitizen]
|
||||
version = "0.1.71"
|
||||
version = "0.1.71.post1"
|
||||
version_files = [
|
||||
"pyproject.toml:^version",
|
||||
"../pyproject.toml:litellm-enterprise==",
|
||||
|
|
|
|||
|
|
@ -140,11 +140,14 @@ impl BaseAnthropicMessagesConfig for AnthropicMessagesConfig {
|
|||
}
|
||||
|
||||
fn request_headers(&self, headers: Headers, request: &AnthropicMessagesRequest) -> Headers {
|
||||
update_headers_with_anthropic_beta(headers, request)
|
||||
merge_beta_headers(
|
||||
update_headers_with_anthropic_beta(headers, request),
|
||||
safeguards_betas(&request.params),
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
fn update_headers_with_anthropic_beta(
|
||||
pub(crate) fn update_headers_with_anthropic_beta(
|
||||
headers: Headers,
|
||||
request: &AnthropicMessagesRequest,
|
||||
) -> Headers {
|
||||
|
|
@ -199,6 +202,20 @@ fn uses_structured_output(params: &AnthropicMessagesOptionalParams) -> bool {
|
|||
.is_some_and(|config| config.format.is_some())
|
||||
}
|
||||
|
||||
fn safeguards_betas(params: &AnthropicMessagesOptionalParams) -> BetaSet {
|
||||
carries_safeguards(params)
|
||||
.then_some(AnthropicBeta::DangerousToolUse20260903)
|
||||
.into_iter()
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn carries_safeguards(params: &AnthropicMessagesOptionalParams) -> bool {
|
||||
params
|
||||
.extra
|
||||
.get("safeguards")
|
||||
.is_some_and(|value| !value.is_null())
|
||||
}
|
||||
|
||||
fn messages_carry_output_config(messages: &[AnthropicMessage]) -> bool {
|
||||
messages
|
||||
.iter()
|
||||
|
|
@ -1047,6 +1064,12 @@ mod tests {
|
|||
json!({"messages": [{"role": "user", "content": "hi", "output_config": {"effort": "low"}}]}),
|
||||
&[("anthropic-beta", "claude-code-20250219,context-management-2025-06-27,effort-2025-11-24,interleaved-thinking-2025-05-14,per-turn-control-2026-07-01")],
|
||||
)]
|
||||
#[case::safeguards_add_the_dangerous_tool_use_beta(
|
||||
&[("anthropic-beta", "claude-code-20250219")],
|
||||
json!({"safeguards": [{"type": "dangerous_tool_use"}]}),
|
||||
&[("anthropic-beta", "claude-code-20250219,dangerous-tool-use-2026-09-03")],
|
||||
)]
|
||||
#[case::null_safeguards_add_no_beta(&[("x-api-key", "k")], json!({"safeguards": null}), &[("x-api-key", "k")])]
|
||||
fn request_headers_merge_the_feature_betas(
|
||||
#[case] input: &[(&str, &str)],
|
||||
#[case] fields: Value,
|
||||
|
|
|
|||
|
|
@ -12,7 +12,9 @@ use crate::{
|
|||
Error,
|
||||
anthropic::{
|
||||
common_utils::{API_KEY_PLACEMENT, MESSAGES_PATH_SUFFIX, non_empty},
|
||||
messages::transformation::{ANTHROPIC_MESSAGES_CONFIG, AnthropicMessagesConfig},
|
||||
messages::transformation::{
|
||||
ANTHROPIC_MESSAGES_CONFIG, AnthropicMessagesConfig, update_headers_with_anthropic_beta,
|
||||
},
|
||||
},
|
||||
base_llm::{
|
||||
anthropic_messages::transformation::{
|
||||
|
|
@ -103,7 +105,7 @@ impl BaseAnthropicMessagesConfig for AzureAnthropicMessagesConfig {
|
|||
}
|
||||
|
||||
fn request_headers(&self, headers: Headers, request: &AnthropicMessagesRequest) -> Headers {
|
||||
self.anthropic.request_headers(headers, request)
|
||||
update_headers_with_anthropic_beta(headers, request)
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -29,10 +29,16 @@ pub struct ModelInfo {
|
|||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub cache_creation_input_token_cost_above_32k_tokens: Option<f64>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub cache_creation_input_token_cost_above_100k_tokens: Option<f64>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub cache_creation_input_token_cost_above_100k_tokens_batches: Option<f64>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub cache_creation_input_token_cost_above_128k_tokens: Option<f64>,
|
||||
/// Rate applied once the prompt exceeds the token threshold in the field name.
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub cache_creation_input_token_cost_above_1hr: Option<f64>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub cache_creation_input_token_cost_above_1hr_above_100k_tokens: Option<f64>,
|
||||
/// Rate applied once the prompt exceeds the token threshold in the field name.
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub cache_creation_input_token_cost_above_1hr_above_200k_tokens: Option<f64>,
|
||||
|
|
@ -76,6 +82,10 @@ pub struct ModelInfo {
|
|||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub cache_read_input_token_cost_above_32k_tokens: Option<f64>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub cache_read_input_token_cost_above_100k_tokens: Option<f64>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub cache_read_input_token_cost_above_100k_tokens_batches: Option<f64>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub cache_read_input_token_cost_above_128k_tokens: Option<f64>,
|
||||
/// Rate applied once the prompt exceeds the token threshold in the field name.
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
|
|
@ -183,6 +193,10 @@ pub struct ModelInfo {
|
|||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub input_cost_per_token_above_32k_tokens: Option<f64>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub input_cost_per_token_above_100k_tokens: Option<f64>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub input_cost_per_token_above_100k_tokens_batches: Option<f64>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub input_cost_per_token_above_128k_tokens: Option<f64>,
|
||||
/// Rate applied once the prompt exceeds the token threshold in the field name.
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
|
|
@ -329,6 +343,10 @@ pub struct ModelInfo {
|
|||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub output_cost_per_token_above_32k_tokens: Option<f64>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub output_cost_per_token_above_100k_tokens: Option<f64>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub output_cost_per_token_above_100k_tokens_batches: Option<f64>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub output_cost_per_token_above_128k_tokens: Option<f64>,
|
||||
/// Rate applied once the prompt exceeds the token threshold in the field name.
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
|
|
|
|||
|
|
@ -74,3 +74,20 @@ fn checked_in_catalog_and_backup_match() {
|
|||
"invalid registry aliases"
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::input("input_cost_per_token_above_100k_tokens")]
|
||||
#[case::input_batches("input_cost_per_token_above_100k_tokens_batches")]
|
||||
#[case::output("output_cost_per_token_above_100k_tokens")]
|
||||
#[case::output_batches("output_cost_per_token_above_100k_tokens_batches")]
|
||||
#[case::cache_creation("cache_creation_input_token_cost_above_100k_tokens")]
|
||||
#[case::cache_creation_batches("cache_creation_input_token_cost_above_100k_tokens_batches")]
|
||||
#[case::cache_creation_1hr("cache_creation_input_token_cost_above_1hr_above_100k_tokens")]
|
||||
#[case::cache_read("cache_read_input_token_cost_above_100k_tokens")]
|
||||
#[case::cache_read_batches("cache_read_input_token_cost_above_100k_tokens_batches")]
|
||||
fn registry_validation_keeps_above_100k_tier_rates(#[case] field: &str) {
|
||||
let mut entry = Map::new();
|
||||
entry.insert("litellm_provider".into(), "anthropic".into());
|
||||
entry.insert(field.into(), 5e-7.into());
|
||||
validate_model_entry("test", &Value::Object(entry)).unwrap();
|
||||
}
|
||||
|
|
|
|||
|
|
@ -556,6 +556,7 @@ from litellm.litellm_core_utils.get_model_cost_map import get_model_cost_map, ma
|
|||
|
||||
model_cost = get_model_cost_map(url=model_cost_map_url)
|
||||
cost_discount_config: Dict[str, float] = {} # Provider-specific cost discounts {"vertex_ai": 0.05} = 5% discount
|
||||
apply_search_free_tier: bool = False
|
||||
cost_margin_config: Dict[
|
||||
str, Union[float, Dict[str, float]]
|
||||
] = {} # Provider-specific or global cost margins. Examples:
|
||||
|
|
|
|||
|
|
@ -217,7 +217,7 @@
|
|||
"computer-use-2025-11-24": "computer-use-2025-11-24",
|
||||
"context-1m-2025-08-07": "context-1m-2025-08-07",
|
||||
"context-management-2025-06-27": "context-management-2025-06-27",
|
||||
"dangerous-tool-use-2026-09-03": null,
|
||||
"dangerous-tool-use-2026-09-03": "dangerous-tool-use-2026-09-03",
|
||||
"effort-2025-11-24": "effort-2025-11-24",
|
||||
"fast-mode-2026-02-01": "fast-mode-2026-02-01",
|
||||
"files-api-2025-04-14": "files-api-2025-04-14",
|
||||
|
|
|
|||
|
|
@ -22,17 +22,28 @@ Environment Variables:
|
|||
- LITELLM_ANTHROPIC_BETA_HEADERS_URL: Custom URL for remote config (optional)
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import os
|
||||
import threading
|
||||
from collections.abc import Mapping
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from functools import cache
|
||||
from importlib.resources import files
|
||||
from typing import Final
|
||||
from typing import Final, TypeAlias
|
||||
|
||||
import httpx
|
||||
from pydantic import TypeAdapter
|
||||
|
||||
from litellm.litellm_core_utils.litellm_logging import verbose_logger
|
||||
|
||||
# Cache for the loaded configuration
|
||||
_BETA_HEADERS_CONFIG: dict | None = None
|
||||
_BETA_HEADERS_CONFIG_LOCK: Final = threading.Lock()
|
||||
_BETA_HEADERS_CONFIG_LOADER: Final = ThreadPoolExecutor(max_workers=1, thread_name_prefix="litellm-beta-headers-config")
|
||||
|
||||
BetaHeadersConfig: TypeAlias = Mapping[str, str | Mapping[str, str | None]]
|
||||
_BETA_HEADERS_CONFIG_ADAPTER: Final = TypeAdapter(BetaHeadersConfig)
|
||||
|
||||
|
||||
class GetAnthropicBetaHeadersConfig:
|
||||
|
|
@ -174,13 +185,27 @@ def _load_beta_headers_config() -> dict:
|
|||
if _BETA_HEADERS_CONFIG is not None:
|
||||
return _BETA_HEADERS_CONFIG
|
||||
|
||||
# Get the URL from environment or use default
|
||||
from litellm import anthropic_beta_headers_url
|
||||
with _BETA_HEADERS_CONFIG_LOCK:
|
||||
if _BETA_HEADERS_CONFIG is not None:
|
||||
return _BETA_HEADERS_CONFIG
|
||||
|
||||
_BETA_HEADERS_CONFIG = get_beta_headers_config(url=anthropic_beta_headers_url)
|
||||
verbose_logger.debug("Loaded and cached beta headers config")
|
||||
# Get the URL from environment or use default
|
||||
from litellm import anthropic_beta_headers_url
|
||||
|
||||
return _BETA_HEADERS_CONFIG
|
||||
_BETA_HEADERS_CONFIG = get_beta_headers_config(url=anthropic_beta_headers_url)
|
||||
verbose_logger.debug("Loaded and cached beta headers config")
|
||||
|
||||
return _BETA_HEADERS_CONFIG
|
||||
|
||||
|
||||
async def aload_beta_headers_config() -> None:
|
||||
if _BETA_HEADERS_CONFIG is None:
|
||||
await asyncio.get_running_loop().run_in_executor(_BETA_HEADERS_CONFIG_LOADER, _load_beta_headers_config)
|
||||
|
||||
|
||||
@cache
|
||||
def _bundled_beta_headers_config() -> BetaHeadersConfig:
|
||||
return _BETA_HEADERS_CONFIG_ADAPTER.validate_python(GetAnthropicBetaHeadersConfig.load_local_beta_headers_config())
|
||||
|
||||
|
||||
def reload_beta_headers_config() -> dict:
|
||||
|
|
@ -288,6 +313,14 @@ def is_beta_header_supported(
|
|||
return beta_header in provider_mapping and provider_mapping[beta_header] is not None
|
||||
|
||||
|
||||
def is_beta_header_supported_by_bundled_config(beta_header: str, provider: str) -> bool:
|
||||
config: Final = _bundled_beta_headers_config()
|
||||
aliases: Final = config.get("provider_aliases", {})
|
||||
canonical_provider: Final = (aliases.get(provider) if isinstance(aliases, Mapping) else None) or provider
|
||||
provider_mapping: Final = config.get(canonical_provider, {})
|
||||
return isinstance(provider_mapping, Mapping) and provider_mapping.get(beta_header) is not None
|
||||
|
||||
|
||||
def get_provider_beta_header(
|
||||
anthropic_beta_header: str,
|
||||
provider: str,
|
||||
|
|
|
|||
|
|
@ -3,7 +3,7 @@ from collections.abc import Iterable, Iterator, Mapping
|
|||
from dataclasses import dataclass
|
||||
from dataclasses import replace as dataclasses_replace
|
||||
from enum import Enum
|
||||
from typing import Any, Final, Literal
|
||||
from typing import Any, Final, Literal, cast
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_logger
|
||||
|
|
@ -73,6 +73,7 @@ async def calculate_batch_cost_and_usage(
|
|||
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "mistral"],
|
||||
model_name: str | None = None,
|
||||
model_info: ModelInfo | None = None,
|
||||
vertex_location: str | None = None,
|
||||
) -> BatchCostUsageResult:
|
||||
"""
|
||||
Calculate the cost and usage of a batch.
|
||||
|
|
@ -82,16 +83,21 @@ async def calculate_batch_cost_and_usage(
|
|||
pricing. Threaded through to batch_cost_calculator so that
|
||||
deployment-specific pricing (e.g. input_cost_per_token_batches)
|
||||
is used instead of the global cost map.
|
||||
vertex_location: The Vertex AI location the batch ran on; non-global
|
||||
locations apply the model's regional-endpoint uplift multiplier.
|
||||
"""
|
||||
first_row: Final = file_content_dictionary[0] if file_content_dictionary else None
|
||||
if _uses_native_vertex_output(custom_llm_provider, model_name, first_row):
|
||||
return calculate_vertex_ai_batch_cost_and_usage(file_content_dictionary, model_name, model_info=model_info)
|
||||
return calculate_vertex_ai_batch_cost_and_usage(
|
||||
file_content_dictionary, model_name, model_info=model_info, vertex_location=vertex_location
|
||||
)
|
||||
|
||||
return _aggregate_batch_cost_usage_models(
|
||||
entries=file_content_dictionary,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
model_name=model_name,
|
||||
model_info=model_info,
|
||||
vertex_location=vertex_location,
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -101,6 +107,7 @@ async def _handle_completed_batch(
|
|||
model_name: str | None = None,
|
||||
litellm_params: dict | None = None,
|
||||
model_info: ModelInfo | None = None,
|
||||
vertex_location: str | None = None,
|
||||
) -> BatchCostUsageResult:
|
||||
"""Fetch a completed batch's output file and aggregate its cost, usage, and
|
||||
models in a single pass over the JSONL lines, so the parsed file content is
|
||||
|
|
@ -114,6 +121,8 @@ async def _handle_completed_batch(
|
|||
model_info: Optional deployment-level model info with custom pricing,
|
||||
threaded through so a deployment's configured rates win over the
|
||||
global cost map.
|
||||
vertex_location: The Vertex AI location the batch ran on; non-global
|
||||
locations apply the model's regional-endpoint uplift multiplier.
|
||||
"""
|
||||
# A completed batch whose request lines all failed has no output file - the
|
||||
# results are written to a separate error_file_id and output_file_id is None.
|
||||
|
|
@ -141,7 +150,10 @@ async def _handle_completed_batch(
|
|||
|
||||
output_file_result: Final = (
|
||||
calculate_vertex_ai_batch_cost_and_usage(
|
||||
_iter_batch_output_entries(file_content), model_name, model_info=model_info
|
||||
_iter_batch_output_entries(file_content),
|
||||
model_name,
|
||||
model_info=model_info,
|
||||
vertex_location=vertex_location,
|
||||
)
|
||||
if _uses_native_vertex_output(
|
||||
custom_llm_provider, model_name, next(_iter_batch_output_entries(file_content), None)
|
||||
|
|
@ -151,6 +163,7 @@ async def _handle_completed_batch(
|
|||
custom_llm_provider=custom_llm_provider,
|
||||
model_name=model_name,
|
||||
model_info=model_info,
|
||||
vertex_location=vertex_location,
|
||||
)
|
||||
)
|
||||
|
||||
|
|
@ -186,6 +199,7 @@ def _classify_output_line_stats(
|
|||
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock", "mistral"],
|
||||
model_name: str | None,
|
||||
model_info: ModelInfo | None,
|
||||
vertex_location: str | None = None,
|
||||
) -> Iterator[_BatchOutputLineStats | _LineOutcome]:
|
||||
"""Classify every output line in a single pass, so counting failures never needs
|
||||
a second read of a potentially huge output file. A line the provider reported as
|
||||
|
|
@ -196,7 +210,7 @@ def _classify_output_line_stats(
|
|||
if not _batch_response_was_successful(entry, custom_llm_provider):
|
||||
yield _LineOutcome.PROVIDER_FAILED
|
||||
continue
|
||||
stats = _safe_output_line_stats(entry, custom_llm_provider, model_name, model_info)
|
||||
stats = _safe_output_line_stats(entry, custom_llm_provider, model_name, model_info, vertex_location)
|
||||
yield stats if stats is not None else _LineOutcome.UNCOSTABLE
|
||||
|
||||
|
||||
|
|
@ -205,13 +219,14 @@ def _safe_output_line_stats(
|
|||
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock", "mistral"],
|
||||
model_name: str | None,
|
||||
model_info: ModelInfo | None,
|
||||
vertex_location: str | None = None,
|
||||
) -> _BatchOutputLineStats | None:
|
||||
"""Return the stats for one provider-successful batch output line, or None when
|
||||
it cannot be costed, so a single bad line never aborts the whole batch's cost
|
||||
accounting."""
|
||||
custom_id: Final = entry.get("custom_id") if isinstance(entry, dict) else None
|
||||
try:
|
||||
return _compute_output_line_stats(entry, custom_llm_provider, model_name, model_info)
|
||||
return _compute_output_line_stats(entry, custom_llm_provider, model_name, model_info, vertex_location)
|
||||
except Exception as e: # noqa: BLE001 # any single line's costing failure must not abort the whole batch
|
||||
verbose_logger.warning(
|
||||
"batch output line could not be costed, so it is billed at $0 and the rest of the batch "
|
||||
|
|
@ -227,6 +242,7 @@ def _compute_output_line_stats(
|
|||
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock", "mistral"],
|
||||
model_name: str | None,
|
||||
model_info: ModelInfo | None,
|
||||
vertex_location: str | None = None,
|
||||
) -> _BatchOutputLineStats:
|
||||
response_body: Final = _get_response_from_batch_job_output_file(entry, custom_llm_provider)
|
||||
usage: Final = _get_batch_job_usage_from_response_body(response_body, custom_llm_provider)
|
||||
|
|
@ -241,6 +257,7 @@ def _compute_output_line_stats(
|
|||
model_name=model_name,
|
||||
response_model=response_model,
|
||||
model_info=model_info,
|
||||
vertex_location=vertex_location,
|
||||
)
|
||||
return _BatchOutputLineStats(
|
||||
prompt_cost=line_prompt_cost,
|
||||
|
|
@ -270,6 +287,7 @@ def _output_line_cost(
|
|||
model_name: str | None,
|
||||
response_model: str | None,
|
||||
model_info: ModelInfo | None,
|
||||
vertex_location: str | None = None,
|
||||
) -> tuple[float, float]:
|
||||
"""(prompt_cost, completion_cost) for one output line, priced at batch rates."""
|
||||
from litellm.cost_calculator import batch_cost_calculator, ocr_batch_cost
|
||||
|
|
@ -290,6 +308,7 @@ def _output_line_cost(
|
|||
model=cost_model,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
model_info=model_info,
|
||||
vertex_location=vertex_location,
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -298,11 +317,14 @@ def _aggregate_batch_cost_usage_models(
|
|||
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock", "mistral"],
|
||||
model_name: str | None = None,
|
||||
model_info: ModelInfo | None = None,
|
||||
vertex_location: str | None = None,
|
||||
) -> BatchCostUsageResult:
|
||||
"""Aggregate cost, usage, models, and pass/fail counts from batch output
|
||||
entries in a single pass, holding one small stats record per line instead
|
||||
of the parsed file."""
|
||||
all_results: Final = tuple(_classify_output_line_stats(entries, custom_llm_provider, model_name, model_info))
|
||||
all_results: Final = tuple(
|
||||
_classify_output_line_stats(entries, custom_llm_provider, model_name, model_info, vertex_location)
|
||||
)
|
||||
line_stats: Final = tuple(result for result in all_results if isinstance(result, _BatchOutputLineStats))
|
||||
failed_requests: Final = sum(1 for result in all_results if result is _LineOutcome.PROVIDER_FAILED)
|
||||
successful_requests: Final = len(all_results) - failed_requests
|
||||
|
|
@ -349,6 +371,7 @@ def calculate_vertex_ai_batch_cost_and_usage(
|
|||
vertex_ai_batch_responses: Iterable[dict],
|
||||
model_name: str | None = None,
|
||||
model_info: ModelInfo | None = None,
|
||||
vertex_location: str | None = None,
|
||||
) -> BatchCostUsageResult:
|
||||
"""
|
||||
Cost and usage of a native Vertex predictions.jsonl, one
|
||||
|
|
@ -367,6 +390,7 @@ def calculate_vertex_ai_batch_cost_and_usage(
|
|||
model_info=model_info,
|
||||
calculate_usage=VertexGeminiConfig._calculate_usage,
|
||||
cost_calculator=batch_cost_calculator,
|
||||
vertex_location=vertex_location,
|
||||
)
|
||||
for row in vertex_ai_batch_responses
|
||||
)
|
||||
|
|
@ -744,9 +768,17 @@ def _batch_response_was_successful(
|
|||
message batch results lines report ``result.type == "succeeded"``; Bedrock
|
||||
batch output lines report ``modelOutput`` (and no ``error``).
|
||||
"""
|
||||
if custom_llm_provider == "anthropic":
|
||||
return _get_anthropic_result_from_batch_results_line(batch_job_output_file).get("type") == "succeeded"
|
||||
if custom_llm_provider == "bedrock":
|
||||
return batch_job_output_file.get("modelOutput") is not None and batch_job_output_file.get("error") is None
|
||||
_response: Final[dict] = batch_job_output_file.get("response", None) or {}
|
||||
return _response.get("status_code", None) == 200
|
||||
envelope_key: Final = "result" if custom_llm_provider == "anthropic" else "response"
|
||||
envelope: Final = batch_job_output_file.get(envelope_key)
|
||||
if not isinstance(envelope, Mapping):
|
||||
if envelope is not None:
|
||||
verbose_logger.warning("batch output %s must be an object, treating the line as unsuccessful", envelope_key)
|
||||
return False
|
||||
response: Final = cast(Mapping[object, object], envelope) # cast-ok: mapping shape validated above
|
||||
return (
|
||||
response.get("type") == "succeeded"
|
||||
if custom_llm_provider == "anthropic"
|
||||
else response.get("status_code") == 200
|
||||
)
|
||||
|
|
|
|||
|
|
@ -1871,6 +1871,7 @@ LITELLM_SETTINGS_SAFE_DB_OVERRIDES: Final = [
|
|||
"public_model_groups",
|
||||
"public_model_groups_links",
|
||||
"cost_discount_config",
|
||||
"apply_search_free_tier",
|
||||
"cost_margin_config",
|
||||
"block_requests_for_models_without_pricing",
|
||||
"budget_exceeded_throttle_percentage",
|
||||
|
|
|
|||
|
|
@ -8,7 +8,7 @@ from types import MappingProxyType
|
|||
from typing import TYPE_CHECKING, Any, Final, Literal, cast
|
||||
|
||||
from httpx import Response
|
||||
from pydantic import BaseModel
|
||||
from pydantic import BaseModel, TypeAdapter
|
||||
from typing_extensions import ReadOnly, TypedDict, assert_never
|
||||
|
||||
import litellm
|
||||
|
|
@ -36,6 +36,7 @@ from litellm.litellm_core_utils.llm_cost_calc.utils import (
|
|||
get_batch_cost_rates,
|
||||
get_billable_input_tokens,
|
||||
get_token_type_cost_breakdown,
|
||||
get_vertex_regional_endpoint_uplift,
|
||||
parse_prompt_tokens_details,
|
||||
select_cost_metric_for_model,
|
||||
)
|
||||
|
|
@ -125,6 +126,7 @@ from litellm.types.utils import (
|
|||
TranscriptionUsageTokensObject,
|
||||
Usage,
|
||||
VectorStoreSearchResponse,
|
||||
is_server_derived_pricing_key,
|
||||
)
|
||||
from litellm.utils import (
|
||||
CallTypes,
|
||||
|
|
@ -137,6 +139,8 @@ from litellm.utils import (
|
|||
TextCompletionResponse,
|
||||
TranscriptionResponse,
|
||||
_cached_get_model_info_helper,
|
||||
model_info_with_pricing_overrides,
|
||||
runtime_model_cost_registrations,
|
||||
token_counter,
|
||||
)
|
||||
|
||||
|
|
@ -402,6 +406,7 @@ def cost_per_token(
|
|||
### REQUEST MODEL ###
|
||||
request_model: str | None = None, # original request model for router detection
|
||||
custom_model_info: OCRPricing | None = None,
|
||||
bedrock_model_info: ModelInfo | None = None,
|
||||
) -> tuple[float, float]:
|
||||
"""
|
||||
Calculates the cost per token for a given model, prompt tokens, and completion tokens.
|
||||
|
|
@ -627,6 +632,7 @@ def cost_per_token(
|
|||
model=model,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
data_residency=data_residency,
|
||||
vertex_location=vertex_location,
|
||||
)
|
||||
elif call_type == "atranscription" or call_type == "transcription":
|
||||
if _transcription_usage_has_token_details(usage_block):
|
||||
|
|
@ -688,6 +694,15 @@ def cost_per_token(
|
|||
elif custom_llm_provider == "anthropic":
|
||||
return anthropic_cost_per_token(model=model, usage=usage_block, service_tier=service_tier)
|
||||
elif custom_llm_provider == "bedrock":
|
||||
resolved_bedrock_info: Final = bedrock_model_info or _bedrock_pricing_info(model, None, "bedrock", False)
|
||||
if resolved_bedrock_info is not None:
|
||||
return generic_cost_per_token(
|
||||
model=model,
|
||||
usage=usage_block,
|
||||
custom_llm_provider="bedrock",
|
||||
service_tier=service_tier,
|
||||
model_info=resolved_bedrock_info,
|
||||
)
|
||||
return bedrock_cost_per_token(model=model, usage=usage_block, service_tier=service_tier)
|
||||
elif custom_llm_provider == "openai":
|
||||
return openai_cost_per_token(
|
||||
|
|
@ -801,6 +816,55 @@ def _cost_map_entry_prices_anything(entry: Mapping[str, object]) -> bool:
|
|||
)
|
||||
|
||||
|
||||
_LOGGED_LITELLM_PARAMS: Final = TypeAdapter(Mapping[str, object])
|
||||
|
||||
|
||||
def _bedrock_registration_has_prices(entry: Mapping[str, object]) -> bool:
|
||||
return _cost_map_entry_prices_anything(entry) or any(
|
||||
value is not None and field.startswith("cache_") and is_server_derived_pricing_key(field)
|
||||
for field, value in entry.items()
|
||||
)
|
||||
|
||||
|
||||
def _bedrock_pricing_info(
|
||||
model: str | None,
|
||||
invoked_model_id: str | None,
|
||||
custom_llm_provider: str | None,
|
||||
explicit_pricing: bool,
|
||||
) -> ModelInfo | None:
|
||||
if explicit_pricing or custom_llm_provider != "bedrock" or not model:
|
||||
return None
|
||||
registration: Final = next(
|
||||
(
|
||||
(key, entry)
|
||||
for key, entry in runtime_model_cost_registrations(model, "bedrock")
|
||||
if _bedrock_registration_has_prices(entry)
|
||||
),
|
||||
None,
|
||||
)
|
||||
profile: Final = _cost_map_model_info(invoked_model_id, "bedrock") if invoked_model_id else None
|
||||
if profile is None or not _declares_rates(profile):
|
||||
return _cost_map_model_info(registration[0], None) if registration is not None else None
|
||||
if registration is None or registration[0] == profile["key"]:
|
||||
return profile
|
||||
return model_info_with_pricing_overrides(profile, registration[1])
|
||||
|
||||
|
||||
def _declares_rates(model_info: ModelInfo | None) -> bool:
|
||||
entry: Final = _raw_cost_map_entry(model_info["key"]) if model_info is not None else None
|
||||
return entry is not None and _cost_map_entry_prices_anything(entry)
|
||||
|
||||
|
||||
def _invoked_model_id_for_cost(
|
||||
optional_params: object,
|
||||
logging_obj: LitellmLoggingObject | None,
|
||||
) -> str | None:
|
||||
recorded: Final[object] = getattr(logging_obj, "model_call_details", None)
|
||||
return _get_hidden_str_for_cost_calc(optional_params, "model_id") or _get_hidden_str_for_cost_calc(
|
||||
recorded, "model_id"
|
||||
)
|
||||
|
||||
|
||||
def _select_model_name_for_cost_calc(
|
||||
model: str | None,
|
||||
completion_response: object | None,
|
||||
|
|
@ -809,6 +873,7 @@ def _select_model_name_for_cost_calc(
|
|||
custom_llm_provider: str | None = None,
|
||||
router_model_id: str | None = None,
|
||||
region_name: str | None = None,
|
||||
invoked_model_id: str | None = None,
|
||||
) -> str | None:
|
||||
"""
|
||||
1. If custom pricing is true, return received model name
|
||||
|
|
@ -829,6 +894,13 @@ def _select_model_name_for_cost_calc(
|
|||
hidden_params: Final[dict | None] = getattr(completion_response, "_hidden_params", None)
|
||||
provider_response_model: Final = _get_hidden_str_for_cost_calc(hidden_params, "provider_response_model")
|
||||
explicit_pricing: Final = custom_pricing is True or base_model is not None
|
||||
configured_model: Final = model or (
|
||||
completion_response_model if isinstance(completion_response_model, str) else None
|
||||
)
|
||||
bedrock_model_info: Final = _bedrock_pricing_info(
|
||||
configured_model, invoked_model_id, custom_llm_provider, explicit_pricing
|
||||
)
|
||||
bedrock_pricing_model: Final = bedrock_model_info["key"] if bedrock_model_info is not None else None
|
||||
priced_from_response: Final = provider_response_model is not None or completion_response_model is not None
|
||||
priced_region: Final = (
|
||||
_get_hidden_str_for_cost_calc(hidden_params, "region_name") or region_name
|
||||
|
|
@ -846,8 +918,8 @@ def _select_model_name_for_cost_calc(
|
|||
else:
|
||||
return_model = model
|
||||
|
||||
elif base_model is not None or provider_response_model is not None:
|
||||
return_model = base_model if base_model is not None else provider_response_model
|
||||
elif base_model is not None or bedrock_pricing_model is not None or provider_response_model is not None:
|
||||
return_model = base_model if base_model is not None else bedrock_pricing_model or provider_response_model
|
||||
|
||||
elif completion_response_model is None and hidden_params is not None:
|
||||
if hidden_params.get("model", None) is not None and len(hidden_params["model"]) > 0:
|
||||
|
|
@ -1426,6 +1498,10 @@ def completion_cost(
|
|||
service_tier = _normalize_service_tier(service_tier)
|
||||
|
||||
explicit_pricing: Final = custom_pricing is True or base_model is not None
|
||||
invoked_model_id: Final = _invoked_model_id_for_cost(
|
||||
_LOGGED_LITELLM_PARAMS.validate_python(optional_params) if optional_params is not None else None,
|
||||
litellm_logging_obj,
|
||||
)
|
||||
selected_model: Final = _select_model_name_for_cost_calc(
|
||||
model=model,
|
||||
completion_response=completion_response,
|
||||
|
|
@ -1434,6 +1510,14 @@ def completion_cost(
|
|||
base_model=base_model,
|
||||
router_model_id=router_model_id,
|
||||
region_name=region_name,
|
||||
invoked_model_id=invoked_model_id,
|
||||
)
|
||||
|
||||
bedrock_model_info: Final = _bedrock_pricing_info(
|
||||
model or _get_response_model(completion_response),
|
||||
invoked_model_id,
|
||||
_get_provider_for_cost_calc(model, custom_llm_provider),
|
||||
explicit_pricing,
|
||||
)
|
||||
|
||||
potential_model_names: Final = [
|
||||
|
|
@ -1808,6 +1892,7 @@ def completion_cost(
|
|||
vertex_location=vertex_location,
|
||||
response=completion_response,
|
||||
custom_model_info=_ocr_model_info(litellm_logging_obj, custom_pricing, router_model_id),
|
||||
bedrock_model_info=bedrock_model_info,
|
||||
)
|
||||
|
||||
# Get additional costs from provider (e.g., routing fees, infrastructure costs)
|
||||
|
|
@ -1896,6 +1981,7 @@ def completion_cost(
|
|||
data_residency=data_residency,
|
||||
vertex_location=vertex_location,
|
||||
custom_cost_per_token=custom_cost_per_token,
|
||||
model_info=bedrock_model_info,
|
||||
)
|
||||
_reasoning_cost = _token_type_breakdown.reasoning_cost
|
||||
_cache_read_cost = _token_type_breakdown.cache_read_cost
|
||||
|
|
@ -2137,7 +2223,16 @@ def pricing_entry_for_cost_calc(
|
|||
custom_llm_provider=custom_llm_provider,
|
||||
router_model_id=router_model_id,
|
||||
region_name=region_name,
|
||||
invoked_model_id=_invoked_model_id_for_cost(None, litellm_logging_obj),
|
||||
)
|
||||
bedrock_model_info: Final = _bedrock_pricing_info(
|
||||
model or _get_response_model(completion_response),
|
||||
_invoked_model_id_for_cost(None, litellm_logging_obj),
|
||||
_get_provider_for_cost_calc(model, custom_llm_provider),
|
||||
custom_pricing is True or base_model is not None,
|
||||
)
|
||||
if bedrock_model_info is not None:
|
||||
return bedrock_model_info["key"], bedrock_model_info
|
||||
candidates: Final = (selected_model, _get_response_model(completion_response), model)
|
||||
resolved: Final = next(
|
||||
(info for info in (_cost_map_model_info(name, custom_llm_provider) for name in candidates if name) if info),
|
||||
|
|
@ -2571,6 +2666,8 @@ def batch_cost_calculator(
|
|||
custom_llm_provider: str | None = None,
|
||||
model_info: ModelInfo | None = None,
|
||||
data_residency: str | None = None,
|
||||
*,
|
||||
vertex_location: str | None = None,
|
||||
) -> tuple[float, float]:
|
||||
"""
|
||||
Calculate the cost of a batch job.
|
||||
|
|
@ -2656,7 +2753,8 @@ def batch_cost_calculator(
|
|||
total_prompt_cost *= uplift
|
||||
total_completion_cost *= uplift
|
||||
|
||||
return total_prompt_cost, total_completion_cost
|
||||
vertex_uplift: Final = get_vertex_regional_endpoint_uplift(model_info, vertex_location)
|
||||
return total_prompt_cost * vertex_uplift, total_completion_cost * vertex_uplift
|
||||
|
||||
|
||||
def _batch_prompt_cost(
|
||||
|
|
|
|||
|
|
@ -1,9 +1,11 @@
|
|||
import json
|
||||
from collections.abc import Mapping, Sequence
|
||||
from dataclasses import dataclass, field, replace
|
||||
from typing import Final, TypeAlias
|
||||
|
||||
import httpx
|
||||
from pydantic import ConfigDict, TypeAdapter, ValidationError
|
||||
from pydantic_core import to_jsonable_python
|
||||
from typing_extensions import assert_never
|
||||
|
||||
import litellm
|
||||
|
|
@ -30,6 +32,7 @@ from litellm.types.decisions import (
|
|||
OpenAIDecisionResponse,
|
||||
UnsupportedDecisionsRequest,
|
||||
)
|
||||
from litellm.types.llms.openai import ChatCompletionUserMessage
|
||||
from litellm.types.utils import LlmProviders
|
||||
from litellm.utils import ProviderConfigManager
|
||||
|
||||
|
|
@ -148,6 +151,30 @@ def _provider_ir_request(
|
|||
return _UnsupportedSafetyIdentifier()
|
||||
|
||||
|
||||
_LOGGED_FIELDS_ADAPTER: Final[TypeAdapter[Mapping[str, object]]] = TypeAdapter(Mapping[str, object])
|
||||
|
||||
|
||||
def _jsonable_fallback(value: object) -> object:
|
||||
return dict(_LOGGED_FIELDS_ADAPTER.validate_python(value)) if isinstance(value, Mapping) else str(value)
|
||||
|
||||
|
||||
def _logged_request_message(caller_request: object) -> ChatCompletionUserMessage:
|
||||
caller_fields: Final = _LOGGED_FIELDS_ADAPTER.validate_python(
|
||||
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, [], {}))
|
||||
)
|
||||
return ChatCompletionUserMessage(role="user", content=sent_fields)
|
||||
|
||||
|
||||
def log_caller_request(kwargs: Mapping[str, object], caller_request: object) -> None:
|
||||
"""Record every request field the caller sent as the logged `messages`, before any validation."""
|
||||
logging_obj: Final = kwargs.get("litellm_logging_obj")
|
||||
if isinstance(logging_obj, LiteLLMLoggingObj):
|
||||
logging_obj.update_messages([_logged_request_message(caller_request)])
|
||||
|
||||
|
||||
def prepare_call(
|
||||
*,
|
||||
model: str,
|
||||
|
|
|
|||
|
|
@ -13,6 +13,7 @@ from litellm.decisions.call import (
|
|||
SystemOneQuestions,
|
||||
SystemOneRequestFields,
|
||||
asend,
|
||||
log_caller_request,
|
||||
prepare_call,
|
||||
send,
|
||||
)
|
||||
|
|
@ -59,6 +60,14 @@ def _prepare_decisions_call(
|
|||
extra_headers: Mapping[str, str] | None,
|
||||
kwargs: Mapping[str, object],
|
||||
) -> DecisionsCall:
|
||||
log_caller_request(
|
||||
kwargs,
|
||||
{
|
||||
"input": decision_input,
|
||||
"questions": questions,
|
||||
**({"state": kwargs["state"]} if "state" in kwargs else {}),
|
||||
},
|
||||
)
|
||||
return prepare_call(
|
||||
model=model,
|
||||
fields=_request_fields(decision_input, questions, kwargs),
|
||||
|
|
|
|||
|
|
@ -42,6 +42,7 @@ from litellm.constants import (
|
|||
DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT,
|
||||
DEFAULT_MOCK_RESPONSE_PROMPT_TOKEN_COUNT,
|
||||
EMPTY_MAPPING,
|
||||
LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL,
|
||||
PROVIDER_REQUEST_ID_HEADERS,
|
||||
SENTRY_DENYLIST,
|
||||
SENTRY_PII_DENYLIST,
|
||||
|
|
@ -484,7 +485,7 @@ def _published_pricing(deployment_model: str | None) -> ModelInfo | None:
|
|||
return None
|
||||
|
||||
|
||||
def _resolve_vertex_location_for_cost(
|
||||
def resolve_vertex_location_for_cost(
|
||||
custom_llm_provider: str | None,
|
||||
litellm_params: Mapping[str, object] | None,
|
||||
optional_params: Mapping[str, object] | None,
|
||||
|
|
@ -714,6 +715,7 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
self._enqueue_deferred_logging: Callable[[], None] | None = None
|
||||
self._async_success_scheduled: bool = False
|
||||
self._on_detached_stream_failure: Callable[[Exception], Awaitable[None]] | None = None
|
||||
self._async_success_enqueued: bool = False
|
||||
self.shadow_eval_request_snapshot: GuardrailRequestSnapshot | None = None
|
||||
|
||||
def set_response_timing_metrics(self, timing_metrics: Mapping[str, float]) -> None:
|
||||
|
|
@ -1405,15 +1407,20 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
"LiteLLM.LoggingError: [Non-Blocking] Exception occurred while logging %s", e
|
||||
)
|
||||
|
||||
self.model_call_details["api_call_start_time"] = datetime.datetime.now()
|
||||
# Set-once first provider-handoff instant. api_call_start_time
|
||||
# is overwritten on every retry, so it can't measure one-time
|
||||
# preprocessing; pinning the first attempt excludes retry loops
|
||||
# + backoff. Logging object only — must NOT go into
|
||||
# litellm_params["metadata"] (caller request metadata, typed
|
||||
# Dict[str, str], echoed downstream; a datetime breaks it).
|
||||
if self.model_call_details.get("first_api_call_start_time") is None:
|
||||
self.model_call_details["first_api_call_start_time"] = self.model_call_details["api_call_start_time"]
|
||||
# A proxy-gate failure (auth, rate limit, client disconnect) fires pre_call for its callbacks
|
||||
# without any provider handoff, so it must not stamp the handoff instants.
|
||||
if LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL not in self.model_call_details:
|
||||
self.model_call_details["api_call_start_time"] = datetime.datetime.now()
|
||||
# Set-once first provider-handoff instant. api_call_start_time
|
||||
# is overwritten on every retry, so it can't measure one-time
|
||||
# preprocessing; pinning the first attempt excludes retry loops
|
||||
# + backoff. Logging object only — must NOT go into
|
||||
# litellm_params["metadata"] (caller request metadata, typed
|
||||
# Dict[str, str], echoed downstream; a datetime breaks it).
|
||||
if self.model_call_details.get("first_api_call_start_time") is None:
|
||||
self.model_call_details["first_api_call_start_time"] = self.model_call_details[
|
||||
"api_call_start_time"
|
||||
]
|
||||
# Input Integration Logging -> If you want to log the fact that an attempt to call the model was made
|
||||
callbacks: Final = litellm.input_callback + (self.dynamic_input_callbacks or [])
|
||||
for callback in callbacks:
|
||||
|
|
@ -1880,7 +1887,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 = self._prompt_for_cost_calculation()
|
||||
|
||||
|
|
@ -1907,7 +1914,7 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
if hasattr(self, "litellm_params") and self.litellm_params
|
||||
else None
|
||||
),
|
||||
"vertex_location": _resolve_vertex_location_for_cost(
|
||||
"vertex_location": resolve_vertex_location_for_cost(
|
||||
custom_llm_provider=self.model_call_details.get("custom_llm_provider", None),
|
||||
litellm_params=(self.litellm_params if hasattr(self, "litellm_params") else None),
|
||||
optional_params=self.optional_params,
|
||||
|
|
@ -2019,7 +2026,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(
|
||||
|
|
@ -2044,20 +2051,23 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
usage=usage,
|
||||
)
|
||||
|
||||
def _custom_pricing_for(self, result: object) -> bool:
|
||||
litellm_params: Final = getattr(self, "litellm_params", None)
|
||||
def spilled_over_ptu_for(self, result: object) -> bool:
|
||||
result_hidden_params: Final = getattr(result, "_hidden_params", None) or MappingProxyType({})
|
||||
additional_headers: Final = (
|
||||
result_hidden_params.get("additional_headers")
|
||||
if isinstance(result_hidden_params, dict)
|
||||
else getattr(result_hidden_params, "additional_headers", None)
|
||||
)
|
||||
spilled_over: Final = is_spilled_over_ptu_request(
|
||||
model_info=_deployment_model_info(litellm_params),
|
||||
return is_spilled_over_ptu_request(
|
||||
model_info=_deployment_model_info(getattr(self, "litellm_params", None)),
|
||||
response_headers=self.model_call_details.get("response_headers"),
|
||||
additional_headers=additional_headers,
|
||||
)
|
||||
return False if spilled_over else use_custom_pricing_for_model(litellm_params=litellm_params)
|
||||
|
||||
def custom_pricing_for(self, result: object) -> bool:
|
||||
if self.spilled_over_ptu_for(result):
|
||||
return False
|
||||
return use_custom_pricing_for_model(litellm_params=getattr(self, "litellm_params", None))
|
||||
|
||||
def _prompt_for_cost_calculation(self) -> str:
|
||||
"""
|
||||
|
|
@ -2264,6 +2274,16 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
except Exception:
|
||||
return True
|
||||
|
||||
@property
|
||||
def async_success_log_enqueued(self) -> bool:
|
||||
"""True once the success log task was actually created, not merely claimed, so a later failure raised
|
||||
around it (a disconnect cancelling a still-pending guardrail task) must not bill the prompt a second
|
||||
time. A claimed log parked in the deferred slot never runs once a 499 interrupts, so it stays False."""
|
||||
return self._async_success_enqueued
|
||||
|
||||
def mark_async_success_log_enqueued(self) -> None:
|
||||
self._async_success_enqueued = True
|
||||
|
||||
def claim_async_success_log(self) -> bool:
|
||||
"""One async success log per request. The innermost @client wrapper always exits first,
|
||||
since the outer one is awaiting it, so it claims the log here and the outer wrapper gets
|
||||
|
|
@ -3285,6 +3305,12 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
model_name=self.get_deployment_model_for_cost(),
|
||||
litellm_params=self.litellm_params,
|
||||
model_info=self.get_router_deployment_model_info(),
|
||||
vertex_location=resolve_vertex_location_for_cost(
|
||||
custom_llm_provider=self.custom_llm_provider,
|
||||
litellm_params=self.litellm_params if hasattr(self, "litellm_params") else None,
|
||||
optional_params=None,
|
||||
model=self.get_deployment_model_for_cost() or self.model,
|
||||
),
|
||||
)
|
||||
|
||||
result._hidden_params["response_cost"] = batch_result.cost
|
||||
|
|
@ -5890,6 +5916,7 @@ class StandardLoggingPayloadSetup:
|
|||
custom_llm_provider: str | None,
|
||||
init_response_obj: object,
|
||||
api_base: str | None = None,
|
||||
router_model_id: str | None = None,
|
||||
) -> StandardLoggingModelInformation:
|
||||
model_cost_name: Final = _select_model_name_for_cost_calc(
|
||||
model=base_model if custom_pricing else None,
|
||||
|
|
@ -5897,12 +5924,23 @@ class StandardLoggingPayloadSetup:
|
|||
base_model=base_model,
|
||||
custom_pricing=custom_pricing,
|
||||
)
|
||||
pricing_model_name: Final = (
|
||||
_select_model_name_for_cost_calc(
|
||||
model=base_model if custom_pricing else None,
|
||||
completion_response=init_response_obj,
|
||||
base_model=base_model,
|
||||
custom_pricing=custom_pricing,
|
||||
router_model_id=router_model_id,
|
||||
)
|
||||
if router_model_id is not None
|
||||
else model_cost_name
|
||||
)
|
||||
if model_cost_name is None:
|
||||
model_cost_information = StandardLoggingModelInformation(model_map_key="", model_map_value=None)
|
||||
else:
|
||||
try:
|
||||
_model_cost_information: Final = litellm.get_model_info(
|
||||
model=model_cost_name,
|
||||
model=pricing_model_name or model_cost_name,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
api_base=api_base,
|
||||
)
|
||||
|
|
@ -6534,13 +6572,16 @@ def get_standard_logging_object_payload(
|
|||
)
|
||||
|
||||
## Get model cost information ##
|
||||
base_model = _get_base_model_from_metadata(model_call_details=kwargs)
|
||||
# The router overrides completion_response.model to the model-group alias before
|
||||
# this payload is built, so cost-map lookup via that alias always misses.
|
||||
# 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)
|
||||
configured_base_model: Final = _get_base_model_from_metadata(model_call_details=kwargs)
|
||||
spilled_over: Final = logging_obj.spilled_over_ptu_for(init_response_obj)
|
||||
# The router overrides completion_response.model to the model-group alias before this payload is
|
||||
# built, so cost-map lookup via that alias misses and the deployment model stands in for base_model.
|
||||
# A PTU request Azure spilled onto pay-as-you-go capacity is billed at the served model's rates,
|
||||
# which the deployment's zeroed PTU entry would misreport, so that fallback is skipped there.
|
||||
base_model: Final = (
|
||||
metadata.get("deployment") if configured_base_model is None and not spilled_over else configured_base_model
|
||||
)
|
||||
custom_pricing: Final = False if spilled_over else 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"))
|
||||
|
|
@ -6564,6 +6605,7 @@ def get_standard_logging_object_payload(
|
|||
custom_llm_provider=kwargs.get("custom_llm_provider"),
|
||||
init_response_obj=init_response_obj,
|
||||
api_base=litellm_params.get("api_base"),
|
||||
router_model_id=logging_obj.get_router_model_id(),
|
||||
)
|
||||
|
||||
error_information, error_str = StandardLoggingPayloadSetup.get_error_information_for_logging_payload(
|
||||
|
|
|
|||
|
|
@ -75,6 +75,9 @@ _SERVICE_TIER_TO_COST_KEY_SUFFIX: Final[Mapping[str, str]] = MappingProxyType(
|
|||
ServiceTier.ULTRAFAST.value: ServiceTier.ULTRAFAST.value,
|
||||
}
|
||||
)
|
||||
SERVICE_TIER_COST_KEY_SUFFIXES: Final[tuple[str, ...]] = tuple(
|
||||
sorted(frozenset(f"_{suffix}" for suffix in _SERVICE_TIER_TO_COST_KEY_SUFFIX.values()))
|
||||
)
|
||||
|
||||
_INCLUSIVE_THRESHOLD_PROVIDERS: Final = frozenset({"xai"})
|
||||
_BATCH_KEY_SUFFIX: Final = "_batches"
|
||||
|
|
@ -272,7 +275,7 @@ def _parse_token_threshold(threshold: str) -> float:
|
|||
|
||||
|
||||
def _parse_above_token_threshold(key: str) -> float:
|
||||
return _parse_token_threshold(key.split("_above_")[1].split("_tokens")[0])
|
||||
return _parse_token_threshold(key.rsplit("_above_", 1)[1].split("_tokens")[0])
|
||||
|
||||
|
||||
def _prompt_exceeds_threshold(prompt_tokens: int, threshold: float, inclusive: bool) -> bool:
|
||||
|
|
@ -665,20 +668,26 @@ def _get_token_base_cost(
|
|||
# Exclude service_tier-specific variants (e.g. input_cost_per_token_above_200k_tokens_priority)
|
||||
# so that the threshold detection loop only processes standard keys. The
|
||||
# service_tier-specific above-threshold key is resolved later via _get_service_tier_cost_key.
|
||||
tier_key_suffix: Final = _get_service_tier_cost_key("", service_tier)
|
||||
threshold_prefixes: Final = tuple(f"{prefix}_above_" for prefix in _BATCH_RATE_PREFIXES)
|
||||
threshold_keys: Final = [
|
||||
k
|
||||
for k in model_info
|
||||
if k.startswith("input_cost_per_token_above_") and not k.endswith(_NON_STANDARD_THRESHOLD_SUFFIXES)
|
||||
if k.startswith(threshold_prefixes)
|
||||
and "_tokens" in k
|
||||
and (
|
||||
not k.endswith(_NON_STANDARD_THRESHOLD_SUFFIXES) or (tier_key_suffix != "" and k.endswith(tier_key_suffix))
|
||||
)
|
||||
]
|
||||
|
||||
# Only sort the threshold keys (typically 1-2 keys instead of 66+)
|
||||
threshold: float | None = None
|
||||
for key in sorted(threshold_keys, key=_parse_above_token_threshold, reverse=True):
|
||||
for key in sorted(threshold_keys, key=_parse_above_token_threshold):
|
||||
value = model_info.get(key)
|
||||
if value is not None:
|
||||
try:
|
||||
# Handle both formats: _above_128k_tokens and _above_128_tokens
|
||||
threshold_str = key.split("_above_")[1].split("_tokens")[0]
|
||||
threshold_str = key.rsplit("_above_", 1)[1].split("_tokens")[0]
|
||||
threshold = _parse_above_token_threshold(key)
|
||||
if _prompt_exceeds_threshold(usage.prompt_tokens, threshold, threshold_is_inclusive):
|
||||
# Prefer a service_tier-specific above-threshold key when available,
|
||||
|
|
@ -691,7 +700,7 @@ def _get_token_base_cost(
|
|||
service_tier,
|
||||
)
|
||||
if service_tier
|
||||
else key
|
||||
else f"input_cost_per_token_above_{threshold_str}_tokens"
|
||||
)
|
||||
prompt_base_cost = cast(
|
||||
float,
|
||||
|
|
@ -747,8 +756,6 @@ def _get_token_base_cost(
|
|||
)
|
||||
|
||||
cache_read_cost = _get_cost_per_unit(model_info, cache_read_tiered_key, cache_read_cost)
|
||||
|
||||
break
|
||||
except (IndexError, ValueError):
|
||||
continue
|
||||
except Exception:
|
||||
|
|
@ -1064,6 +1071,25 @@ def parse_completion_tokens_details(usage: Usage) -> CompletionTokensDetailsResu
|
|||
)
|
||||
|
||||
|
||||
def _modality_token_cost(
|
||||
model_info: ModelInfo,
|
||||
modality_cost_key: str,
|
||||
tokens: int,
|
||||
prompt_base_cost: float,
|
||||
service_tier: str | None,
|
||||
) -> float:
|
||||
"""
|
||||
Audio, image and video prompt tokens bill at their modality rate when the model has one.
|
||||
Without one they are ordinary prompt tokens, so they bill at `prompt_base_cost`, which
|
||||
`_get_token_base_cost` has already resolved for the service tier and long-context tier.
|
||||
"""
|
||||
tiered_key: Final = _get_service_tier_cost_key(modality_cost_key, service_tier)
|
||||
cost_key: Final = tiered_key if model_info.get(tiered_key) is not None else modality_cost_key
|
||||
if model_info.get(cost_key) is None:
|
||||
return float(tokens) * prompt_base_cost
|
||||
return calculate_cost_component(model_info, cost_key, tokens)
|
||||
|
||||
|
||||
def _calculate_input_cost(
|
||||
prompt_tokens_details: PromptTokensDetailsResult,
|
||||
model_info: ModelInfo,
|
||||
|
|
@ -1075,8 +1101,14 @@ def _calculate_input_cost(
|
|||
) -> float:
|
||||
"""
|
||||
Calculates the input cost for a given model, prompt tokens, and completion tokens.
|
||||
|
||||
A model priced per character (`input_cost_per_character`) bills the reported character count
|
||||
in place of its text tokens, so the same prompt is never charged under both units.
|
||||
"""
|
||||
prompt_cost = float(prompt_tokens_details["text_tokens"]) * prompt_base_cost
|
||||
bills_characters: Final = (
|
||||
prompt_tokens_details["character_count"] > 0 and model_info.get("input_cost_per_character") is not None
|
||||
)
|
||||
prompt_cost = 0.0 if bills_characters else float(prompt_tokens_details["text_tokens"]) * prompt_base_cost
|
||||
|
||||
### CACHE READ COST - Now uses tiered pricing
|
||||
cache_hit_audio_tokens: Final = prompt_tokens_details["cache_hit_audio_tokens"]
|
||||
|
|
@ -1094,28 +1126,37 @@ def _calculate_input_cost(
|
|||
if prompt_tokens_details["audio_tokens"] and not (
|
||||
prompt_tokens_details["audio_length_seconds"] and model_info.get("input_cost_per_audio_per_second") is not None
|
||||
):
|
||||
audio_cost_key: Final = _get_service_tier_cost_key("input_cost_per_audio_token", service_tier)
|
||||
prompt_cost += calculate_cost_component(model_info, audio_cost_key, prompt_tokens_details["audio_tokens"])
|
||||
prompt_cost += _modality_token_cost(
|
||||
model_info,
|
||||
"input_cost_per_audio_token",
|
||||
prompt_tokens_details["audio_tokens"],
|
||||
prompt_base_cost,
|
||||
service_tier,
|
||||
)
|
||||
|
||||
### IMAGE TOKEN COST
|
||||
if prompt_tokens_details["image_tokens"] and not (
|
||||
prompt_tokens_details["image_count"] and model_info.get("input_cost_per_image") is not None
|
||||
):
|
||||
# For image token costs:
|
||||
# First check if input_cost_per_image_token is available. If not, default to generic input_cost_per_token.
|
||||
image_token_cost_key = "input_cost_per_image_token"
|
||||
if model_info.get(image_token_cost_key) is None:
|
||||
image_token_cost_key = "input_cost_per_token"
|
||||
prompt_cost += calculate_cost_component(model_info, image_token_cost_key, prompt_tokens_details["image_tokens"])
|
||||
prompt_cost += _modality_token_cost(
|
||||
model_info,
|
||||
"input_cost_per_image_token",
|
||||
prompt_tokens_details["image_tokens"],
|
||||
prompt_base_cost,
|
||||
service_tier,
|
||||
)
|
||||
|
||||
### VIDEO TOKEN COST
|
||||
if prompt_tokens_details["video_tokens"] and not (
|
||||
prompt_tokens_details["video_length_seconds"] and model_info.get("input_cost_per_video_per_second") is not None
|
||||
):
|
||||
video_token_cost_key = "input_cost_per_video_token"
|
||||
if model_info.get(video_token_cost_key) is None:
|
||||
video_token_cost_key = "input_cost_per_token"
|
||||
prompt_cost += calculate_cost_component(model_info, video_token_cost_key, prompt_tokens_details["video_tokens"])
|
||||
prompt_cost += _modality_token_cost(
|
||||
model_info,
|
||||
"input_cost_per_video_token",
|
||||
prompt_tokens_details["video_tokens"],
|
||||
prompt_base_cost,
|
||||
service_tier,
|
||||
)
|
||||
|
||||
### CACHE WRITING COST - Now uses tiered pricing
|
||||
if (
|
||||
|
|
@ -1629,20 +1670,23 @@ def get_billed_token_rates(
|
|||
vertex_location: str | None = None,
|
||||
current_time: datetime | None = None,
|
||||
custom_cost_per_token: CostPerToken | None = None,
|
||||
model_info: ModelInfo | None = None,
|
||||
) -> BilledTokenRates | None:
|
||||
"""Rates the cost calculator bills ``usage`` at, resolved exactly as the totals and the token-type
|
||||
breakdown resolve them. None when the model's pricing cannot be resolved."""
|
||||
if custom_cost_per_token is not None:
|
||||
return _custom_pricing_rates(custom_cost_per_token)
|
||||
try:
|
||||
model_info: Final = apply_provider_cache_read_default(
|
||||
get_model_info(model=model, custom_llm_provider=custom_llm_provider),
|
||||
resolved_model_info: Final = apply_provider_cache_read_default(
|
||||
model_info
|
||||
if model_info is not None
|
||||
else get_model_info(model=model, custom_llm_provider=custom_llm_provider),
|
||||
custom_llm_provider,
|
||||
)
|
||||
except Exception: # noqa: BLE001 # get_model_info raises a bare Exception for an unmapped model: no rates
|
||||
return None
|
||||
return _cost_map_billed_rates(
|
||||
model_info=model_info,
|
||||
model_info=resolved_model_info,
|
||||
usage=usage,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
service_tier=service_tier,
|
||||
|
|
@ -1661,6 +1705,7 @@ def get_token_type_cost_breakdown(
|
|||
vertex_location: str | None = None,
|
||||
current_time: datetime | None = None,
|
||||
custom_cost_per_token: CostPerToken | None = None,
|
||||
model_info: ModelInfo | None = None,
|
||||
) -> TokenTypeCostBreakdown:
|
||||
"""
|
||||
Provider-agnostic cost of reasoning and cache tokens, derived from the usage
|
||||
|
|
@ -1685,6 +1730,7 @@ def get_token_type_cost_breakdown(
|
|||
vertex_location=vertex_location,
|
||||
current_time=current_time,
|
||||
custom_cost_per_token=custom_cost_per_token,
|
||||
model_info=model_info,
|
||||
)
|
||||
if rates is None:
|
||||
return TokenTypeCostBreakdown(0.0, 0.0, 0.0)
|
||||
|
|
|
|||
|
|
@ -4,6 +4,7 @@ from functools import reduce
|
|||
from typing import Final
|
||||
|
||||
import httpx
|
||||
from pydantic import TypeAdapter, ValidationError
|
||||
|
||||
from litellm.constants import LITELLM_DETAILED_TIMING
|
||||
from litellm.litellm_core_utils.core_helpers import get_litellm_metadata_from_kwargs, process_response_headers
|
||||
|
|
@ -101,6 +102,32 @@ def response_timing_metrics(
|
|||
return {"_response_ms": total_response_time_ms, "litellm_overhead_time_ms": overhead_ms}
|
||||
|
||||
|
||||
_HIDDEN_PARAMS_DICT: Final = TypeAdapter(dict[str, object])
|
||||
|
||||
|
||||
def _hidden_params_mapping(hidden_params: object) -> Mapping[str, object]:
|
||||
try:
|
||||
return _HIDDEN_PARAMS_DICT.validate_python(hidden_params)
|
||||
except ValidationError:
|
||||
return {}
|
||||
|
||||
|
||||
def _existing_provider_response_model(result: object) -> str | None:
|
||||
"""The pricing model a response already carries, like the decisions path sets, wins over the response model."""
|
||||
hidden_params: Final[object] = getattr(result, "_hidden_params", None)
|
||||
carried: Final[object] = _hidden_params_mapping(hidden_params).get("provider_response_model") or getattr(
|
||||
hidden_params, "provider_response_model", None
|
||||
)
|
||||
return carried if isinstance(carried, str) and carried else None
|
||||
|
||||
|
||||
def _provider_response_model(result: object) -> str | None:
|
||||
"""The model the provider answered with, kept because the proxy later rewrites ``result.model`` to the alias the
|
||||
client asked for and cost-map lookups for the standard logging payload would otherwise miss."""
|
||||
model: Final = getattr(result, "model", None) if isinstance(result, ModelResponse) else None
|
||||
return model if isinstance(model, str) and model else None
|
||||
|
||||
|
||||
class ResponseMetadata:
|
||||
"""
|
||||
Handles setting and managing `_hidden_params`, `response_time_ms`, and `litellm_overhead_time_ms` for LiteLLM responses
|
||||
|
|
@ -133,6 +160,8 @@ class ResponseMetadata:
|
|||
preserve_litellm_internal_headers=True,
|
||||
),
|
||||
"litellm_model_name": model,
|
||||
"provider_response_model": _existing_provider_response_model(self.result)
|
||||
or _provider_response_model(self.result),
|
||||
}
|
||||
self._update_hidden_params(new_params)
|
||||
|
||||
|
|
|
|||
|
|
@ -12,10 +12,15 @@ tell "user asked for this" from "nobody set it". This resolver answers that:
|
|||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Final
|
||||
from collections.abc import Mapping
|
||||
from typing import Final, Literal, Protocol
|
||||
|
||||
from pydantic import TypeAdapter
|
||||
|
||||
from litellm.constants import DEFAULT_REQUEST_TIMEOUT_SECONDS
|
||||
|
||||
_TIMEOUTS_BY_SERVICE_TIER: Final[TypeAdapter[Mapping[str, float] | None]] = TypeAdapter(Mapping[str, float] | None)
|
||||
|
||||
|
||||
def get_configured_request_timeout() -> float | None:
|
||||
"""Return the explicitly-configured ``litellm.request_timeout``, else ``None``."""
|
||||
|
|
@ -27,3 +32,19 @@ def get_configured_request_timeout() -> float | None:
|
|||
if timeout != float(DEFAULT_REQUEST_TIMEOUT_SECONDS):
|
||||
return timeout
|
||||
return None
|
||||
|
||||
|
||||
class _DeploymentParams(Protocol):
|
||||
def get(self, key: str, /) -> object: ...
|
||||
|
||||
|
||||
def get_service_tier_timeout(
|
||||
service_tier: object,
|
||||
litellm_params: _DeploymentParams,
|
||||
key: Literal["timeout_by_service_tier", "stream_timeout_by_service_tier"],
|
||||
) -> float | None:
|
||||
"""Return the deployment's ``key`` timeout for the request's ``service_tier``, else ``None``."""
|
||||
if not isinstance(service_tier, str):
|
||||
return None
|
||||
timeouts: Final = _TIMEOUTS_BY_SERVICE_TIER.validate_python(litellm_params.get(key))
|
||||
return None if timeouts is None else timeouts.get(service_tier)
|
||||
|
|
|
|||
|
|
@ -103,6 +103,7 @@ from ..common_utils import (
|
|||
process_anthropic_headers,
|
||||
requires_native_compaction_beta,
|
||||
strip_advisor_blocks_from_messages,
|
||||
tools_marked_for_a_tool_config_point,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
|
|
@ -2056,6 +2057,19 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
|
|||
# Remove internal LiteLLM parameters that should not be sent to Anthropic API
|
||||
optional_params.pop("is_vertex_request", None)
|
||||
optional_params.pop("client_metadata", None)
|
||||
request_params: Final = cast( # cast-ok: optional_params is the untyped dict of the shared transform signature
|
||||
dict[str, object], optional_params
|
||||
)
|
||||
marked_tools: Final = tools_marked_for_a_tool_config_point(
|
||||
request_params.get("tools"), request_params.pop("cache_control_injection_points", None)
|
||||
)
|
||||
if marked_tools is not None:
|
||||
from litellm.integrations.anthropic_cache_control_hook import AnthropicCacheControlHook
|
||||
|
||||
request_litellm_params: Final = cast( # cast-ok: litellm_params is the untyped dict of the shared signature
|
||||
Mapping[str, object], litellm_params
|
||||
)
|
||||
AnthropicCacheControlHook.record_gateway_injection(request_litellm_params, 1)
|
||||
|
||||
# ``top_k`` is a provider-specific kwarg that bypasses
|
||||
# ``map_openai_params``; gate it here, the single boundary shared by
|
||||
|
|
@ -2075,6 +2089,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
|
|||
"model": model,
|
||||
"messages": anthropic_messages,
|
||||
**optional_params,
|
||||
**({} if marked_tools is None else {"tools": list(marked_tools)}),
|
||||
}
|
||||
|
||||
self._apply_output_config(data=data, model=model, optional_params=optional_params)
|
||||
|
|
|
|||
|
|
@ -33,6 +33,7 @@ from litellm.types.llms.anthropic import (
|
|||
ANTHROPIC_HOSTED_TOOLS,
|
||||
ANTHROPIC_OAUTH_BETA_HEADER,
|
||||
ANTHROPIC_OAUTH_TOKEN_PREFIX,
|
||||
ANTHROPIC_TOOL_SEARCH_TOOL_TYPES,
|
||||
AllAnthropicToolsValues,
|
||||
AnthropicMcpServerTool,
|
||||
AnthropicMessagesToolChoice,
|
||||
|
|
@ -73,7 +74,7 @@ _DATED_RELEASE_SUFFIX_RE: Final = re.compile(r"-\d{8}$")
|
|||
_DOTTED_VERSION_RE: Final = re.compile(r"(\d)\.(\d)")
|
||||
_CLAUDE_CODE_BILLING_HEADER_PREFIX: Final = "x-anthropic-billing-header:"
|
||||
_CLAUDE_CODE_OBJECT_MAPPING_ADAPTER: Final = TypeAdapter(dict[object, object])
|
||||
_CLAUDE_CODE_OBJECT_LIST_ADAPTER: Final = TypeAdapter(list[object])
|
||||
OBJECT_LIST_ADAPTER: Final = TypeAdapter(list[object])
|
||||
|
||||
|
||||
_CLAUDE_CODE_USER_AGENT_PREFIXES: Final = ("claude-cli/", "claude-code/")
|
||||
|
|
@ -115,6 +116,40 @@ def supports_anthropic_cache_control(model: str, custom_llm_provider: str | None
|
|||
)
|
||||
|
||||
|
||||
def _tool_config_point_control(points: object) -> Mapping[object, object] | None:
|
||||
point_fields: Final = tuple(_validated_claude_code_mapping(point) for point in validated_object_list(points) or ())
|
||||
tool_config_point: Final = next(
|
||||
(fields for fields in point_fields if fields is not None and fields.get("location") == "tool_config"), None
|
||||
)
|
||||
if tool_config_point is None:
|
||||
return None
|
||||
control: Final = tool_config_point.get("control")
|
||||
return {"type": "ephemeral"} if control is None else _validated_claude_code_mapping(control)
|
||||
|
||||
|
||||
def tools_marked_for_a_tool_config_point(tools: object, points: object) -> tuple[object, ...] | None:
|
||||
"""Put the ``tool_config`` point's breakpoint on the last tool Anthropic caches.
|
||||
|
||||
Anthropic caches tools as one prefix ending at the marked tool, so the last one carries
|
||||
the mark, skipping tool-search tools whose transform drops cache markers. A client that
|
||||
already marked a tool keeps its own breakpoint. Returns ``None`` when nothing is marked.
|
||||
"""
|
||||
control: Final = _tool_config_point_control(points)
|
||||
tool_list: Final = validated_object_list(tools) or ()
|
||||
tool_fields: Final = tuple(_validated_claude_code_mapping(tool) for tool in tool_list)
|
||||
if control is None or any(fields is not None and "cache_control" in fields for fields in tool_fields):
|
||||
return None
|
||||
cacheable: Final = tuple(
|
||||
(position, fields)
|
||||
for position, fields in enumerate(tool_fields)
|
||||
if fields is not None and fields.get("type") not in ANTHROPIC_TOOL_SEARCH_TOOL_TYPES
|
||||
)
|
||||
if not cacheable:
|
||||
return None
|
||||
position, fields = cacheable[-1]
|
||||
return (*tool_list[:position], {**fields, "cache_control": control}, *tool_list[position + 1 :])
|
||||
|
||||
|
||||
def is_claude_code_user_agent(user_agent: str) -> bool:
|
||||
"""Claude Code sends its API calls through the Anthropic SDK as `claude-cli/<version>` and its own
|
||||
fetches, such as gateway model discovery, as `claude-code/<version>`"""
|
||||
|
|
@ -128,9 +163,9 @@ def _validated_claude_code_mapping(value: object) -> dict[object, object] | None
|
|||
return None
|
||||
|
||||
|
||||
def _validated_claude_code_list(value: object) -> list[object] | None:
|
||||
def validated_object_list(value: object) -> Sequence[object] | None:
|
||||
try:
|
||||
return _CLAUDE_CODE_OBJECT_LIST_ADAPTER.validate_python(value)
|
||||
return OBJECT_LIST_ADAPTER.validate_python(value)
|
||||
except ValidationError:
|
||||
return None
|
||||
|
||||
|
|
@ -157,7 +192,7 @@ def _claude_code_billing_fields(text: str) -> tuple[tuple[str, str], ...] | None
|
|||
def _claude_code_billing_texts(system: object) -> tuple[str, ...] | None:
|
||||
if isinstance(system, str):
|
||||
return (system,)
|
||||
blocks: Final = _validated_claude_code_list(system)
|
||||
blocks: Final = validated_object_list(system)
|
||||
if blocks is None:
|
||||
return None
|
||||
block_mappings: Final = tuple(_validated_claude_code_mapping(block) for block in blocks)
|
||||
|
|
|
|||
|
|
@ -12,6 +12,7 @@ from functools import partial
|
|||
from typing import Any, Final, cast
|
||||
|
||||
import litellm
|
||||
from litellm.integrations.anthropic_cache_control_hook import AnthropicCacheControlHook
|
||||
from litellm.litellm_core_utils.exception_mapping_utils import exception_type
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.llms.anthropic.common_utils import (
|
||||
|
|
@ -19,6 +20,7 @@ from litellm.llms.anthropic.common_utils import (
|
|||
sanitize_tool_use_ids_in_anthropic_messages,
|
||||
strip_empty_content_blocks_from_anthropic_messages,
|
||||
strip_provider_specific_fields_from_anthropic_messages,
|
||||
tools_marked_for_a_tool_config_point,
|
||||
)
|
||||
from litellm.llms.base_llm.anthropic_messages.transformation import (
|
||||
BaseAnthropicMessagesConfig,
|
||||
|
|
@ -636,7 +638,13 @@ def anthropic_messages_handler(
|
|||
f"custom_llm_provider is required for Anthropic messages, passed in model={model}, custom_llm_provider={custom_llm_provider}"
|
||||
)
|
||||
|
||||
local_vars.update(kwargs)
|
||||
request_kwargs: Final = cast(dict[str, object], kwargs) # cast-ok: **kwargs of the handler is untyped
|
||||
marked_tools: Final = tools_marked_for_a_tool_config_point(
|
||||
local_vars.get("tools"), request_kwargs.pop("cache_control_injection_points", None)
|
||||
)
|
||||
if marked_tools is not None:
|
||||
AnthropicCacheControlHook.record_gateway_injection(request_kwargs, 1)
|
||||
local_vars.update({**request_kwargs, **({"tools": list(marked_tools)} if marked_tools is not None else {})})
|
||||
anthropic_messages_optional_request_params: Final = (
|
||||
AnthropicMessagesRequestUtils.get_requested_anthropic_messages_optional_param(
|
||||
params=local_vars,
|
||||
|
|
|
|||
|
|
@ -3,6 +3,7 @@ from typing import Any, Final
|
|||
|
||||
import httpx
|
||||
|
||||
from litellm.anthropic_beta_headers_manager import is_beta_header_supported_by_bundled_config
|
||||
from litellm.exceptions import AuthenticationError
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.litellm_core_utils.litellm_logging import verbose_logger
|
||||
|
|
@ -611,6 +612,7 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig):
|
|||
optional_params: dict,
|
||||
custom_llm_provider: str = "anthropic",
|
||||
messages: Sequence[object] = (),
|
||||
beta_map_provider: str | None = None,
|
||||
) -> dict:
|
||||
"""
|
||||
Auto-inject anthropic-beta headers based on features used.
|
||||
|
|
@ -627,8 +629,9 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig):
|
|||
optional_params: Optional parameters including tools, context_management, output_format, speed
|
||||
custom_llm_provider: Provider name for looking up correct tool search header
|
||||
messages: Request messages, scanned for per-message output_config
|
||||
beta_map_provider: Provider whose beta map column gates safeguards, defaults to custom_llm_provider
|
||||
"""
|
||||
beta_values: Final[set] = set()
|
||||
beta_values: Final[set[str]] = set()
|
||||
|
||||
existing_beta: Final = tuple(
|
||||
piece.strip()
|
||||
|
|
@ -678,6 +681,12 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig):
|
|||
if optional_params.get("speed") == "fast":
|
||||
beta_values.add(ANTHROPIC_BETA_HEADER_VALUES.FAST_MODE_2026_02_01.value)
|
||||
|
||||
if optional_params.get("safeguards") is not None and is_beta_header_supported_by_bundled_config(
|
||||
ANTHROPIC_BETA_HEADER_VALUES.DANGEROUS_TOOL_USE_2026_09_03.value,
|
||||
beta_map_provider or custom_llm_provider,
|
||||
):
|
||||
beta_values.add(ANTHROPIC_BETA_HEADER_VALUES.DANGEROUS_TOOL_USE_2026_09_03.value)
|
||||
|
||||
if _messages_carry_output_config(messages):
|
||||
beta_values.add(ANTHROPIC_BETA_HEADER_VALUES.PER_TURN_CONTROL_2026_07_01.value)
|
||||
|
||||
|
|
|
|||
|
|
@ -4,11 +4,12 @@ import asyncio
|
|||
import json
|
||||
import traceback
|
||||
from collections import deque
|
||||
from collections.abc import AsyncIterator, Iterator, Mapping
|
||||
from collections.abc import AsyncGenerator, Iterator, Mapping
|
||||
from typing import TYPE_CHECKING, Any, Final
|
||||
|
||||
from litellm import verbose_logger
|
||||
from litellm._uuid import uuid
|
||||
from litellm.responses.streaming_iterator import partial_responses_incomplete_event_on_disconnect
|
||||
from litellm.litellm_core_utils.prompt_templates.common_utils import (
|
||||
encrypted_reasoning_signature,
|
||||
)
|
||||
|
|
@ -25,6 +26,7 @@ from .transformation import (
|
|||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObject
|
||||
from litellm.types.llms.openai import ResponseIncompleteEvent
|
||||
|
||||
|
||||
class AnthropicResponsesStreamWrapper:
|
||||
|
|
@ -391,8 +393,11 @@ class AnthropicResponsesStreamWrapper:
|
|||
|
||||
raise StopAsyncIteration
|
||||
|
||||
async def async_anthropic_sse_wrapper(self) -> AsyncIterator[bytes]:
|
||||
"""Yield SSE-encoded bytes for each Anthropic event chunk."""
|
||||
def async_anthropic_sse_wrapper(self) -> "AnthropicResponsesSSEByteStream":
|
||||
"""SSE-encoded bytes for each Anthropic event chunk."""
|
||||
return AnthropicResponsesSSEByteStream(self)
|
||||
|
||||
async def _sse_frames(self) -> AsyncGenerator[bytes, None]:
|
||||
async for chunk in self:
|
||||
if isinstance(chunk, dict):
|
||||
event_type: str = str(chunk.get("type", "message"))
|
||||
|
|
@ -400,3 +405,30 @@ class AnthropicResponsesStreamWrapper:
|
|||
yield payload.encode()
|
||||
else:
|
||||
yield chunk
|
||||
|
||||
def partial_response_on_disconnect(self) -> "ResponseIncompleteEvent | None":
|
||||
return partial_responses_incomplete_event_on_disconnect(self.responses_stream)
|
||||
|
||||
|
||||
class AnthropicResponsesSSEByteStream:
|
||||
"""
|
||||
The byte stream the proxy relays for a default-route /v1/messages request.
|
||||
Keeps the Responses API iterator reachable so a client disconnect can bill
|
||||
the partial output the way native /v1/responses streams do.
|
||||
"""
|
||||
|
||||
def __init__(self, wrapper: AnthropicResponsesStreamWrapper) -> None:
|
||||
self._wrapper: Final = wrapper
|
||||
self._frames: Final = wrapper._sse_frames() # pyright: ignore[reportPrivateUsage] # the wrapper's own byte stream
|
||||
|
||||
def __aiter__(self) -> "AnthropicResponsesSSEByteStream":
|
||||
return self
|
||||
|
||||
async def __anext__(self) -> bytes:
|
||||
return await self._frames.__anext__()
|
||||
|
||||
async def aclose(self) -> None:
|
||||
await self._frames.aclose()
|
||||
|
||||
def partial_response_on_disconnect(self) -> "ResponseIncompleteEvent | None":
|
||||
return self._wrapper.partial_response_on_disconnect()
|
||||
|
|
|
|||
|
|
@ -1,12 +1,10 @@
|
|||
"""Support for Azure OpenAI gpt-5 model family."""
|
||||
|
||||
from typing import Final
|
||||
from typing import ClassVar, Final
|
||||
|
||||
import litellm
|
||||
from litellm.exceptions import UnsupportedParamsError
|
||||
from litellm.llms.openai.chat.gpt_5_transformation import (
|
||||
OpenAIGPT5Config,
|
||||
_get_effort_level,
|
||||
is_gpt_reasoning_series_name,
|
||||
)
|
||||
from litellm.types.llms.openai import AllMessageValues
|
||||
|
|
@ -15,9 +13,16 @@ from .gpt_transformation import AzureOpenAIConfig
|
|||
|
||||
|
||||
class AzureOpenAIGPT5Config(AzureOpenAIConfig, OpenAIGPT5Config):
|
||||
"""Azure specific handling for gpt-5 models."""
|
||||
"""Azure specific handling for gpt-5 models.
|
||||
|
||||
reasoning_effort none joins xhigh as opt-in here: Azure rejected it on every gpt-5
|
||||
deployment before gpt-5.1 (https://github.com/BerriAI/litellm/issues/16704), and a deployment
|
||||
name the map never heard of must not forward it. See
|
||||
https://learn.microsoft.com/en-us/azure/ai-foundry/openai/how-to/reasoning
|
||||
"""
|
||||
|
||||
GPT5_SERIES_ROUTE = "gpt5_series/"
|
||||
OPT_IN_REASONING_EFFORTS: ClassVar[frozenset[str]] = frozenset(("xhigh", "none"))
|
||||
|
||||
@classmethod
|
||||
def _model_map_lookup_name(cls, model: str) -> str:
|
||||
|
|
@ -27,11 +32,18 @@ class AzureOpenAIGPT5Config(AzureOpenAIConfig, OpenAIGPT5Config):
|
|||
model_prices_and_context_window.json; ``azure/gpt-5.1`` is. Overriding the shared
|
||||
resolver rather than one lookup means the supports, explicitly-disabled and
|
||||
default-effort answers all read the same entry.
|
||||
|
||||
An ``azure_ai/`` name reaches this config when ``get_llm_provider`` re-routes a Foundry
|
||||
deployment on an OpenAI-v1 host to the azure provider with the prefix intact. It is the
|
||||
same product the azure_ai config reads ``azure_ai/<model>`` for, so that row answers when
|
||||
the map carries it, and the azure/ twin when it does not.
|
||||
"""
|
||||
if model.startswith(cls.GPT5_SERIES_ROUTE):
|
||||
return "azure/" + model[len(cls.GPT5_SERIES_ROUTE) :]
|
||||
if model.startswith("azure/"):
|
||||
return model
|
||||
if model.startswith("azure_ai/"):
|
||||
return model if model in litellm.model_cost else "azure/" + model.removeprefix("azure_ai/")
|
||||
return "azure/" + model
|
||||
|
||||
@classmethod
|
||||
|
|
@ -75,34 +87,7 @@ class AzureOpenAIGPT5Config(AzureOpenAIConfig, OpenAIGPT5Config):
|
|||
drop_params: bool,
|
||||
api_version: str = "",
|
||||
) -> dict:
|
||||
reasoning_effort_value = non_default_params.get("reasoning_effort") or optional_params.get("reasoning_effort")
|
||||
effective_effort: Final = _get_effort_level(reasoning_effort_value)
|
||||
|
||||
# gpt-5.1/5.2/5.4 support reasoning_effort='none', but other gpt-5 models don't
|
||||
# See: https://learn.microsoft.com/en-us/azure/ai-foundry/openai/how-to/reasoning
|
||||
supports_none: Final = self._supports_reasoning_effort_level(model, "none")
|
||||
|
||||
if effective_effort == "none" and not supports_none:
|
||||
if litellm.drop_params is True or (drop_params is not None and drop_params is True):
|
||||
non_default_params = non_default_params.copy()
|
||||
optional_params = optional_params.copy()
|
||||
if _get_effort_level(non_default_params.get("reasoning_effort")) == "none":
|
||||
non_default_params.pop("reasoning_effort")
|
||||
if _get_effort_level(optional_params.get("reasoning_effort")) == "none":
|
||||
optional_params.pop("reasoning_effort")
|
||||
else:
|
||||
raise UnsupportedParamsError(
|
||||
status_code=400,
|
||||
message=(
|
||||
"Azure OpenAI does not support reasoning_effort='none' for this model. "
|
||||
"Supported values are: 'low', 'medium', and 'high'. "
|
||||
"To drop this parameter, set `litellm.drop_params=True` or for proxy:\n\n"
|
||||
"`litellm_settings:\n drop_params: true`\n"
|
||||
"Issue: https://github.com/BerriAI/litellm/issues/16704"
|
||||
),
|
||||
)
|
||||
|
||||
result: Final = OpenAIGPT5Config.map_openai_params(
|
||||
return OpenAIGPT5Config.map_openai_params(
|
||||
self,
|
||||
non_default_params=non_default_params,
|
||||
optional_params=optional_params,
|
||||
|
|
@ -110,17 +95,6 @@ class AzureOpenAIGPT5Config(AzureOpenAIConfig, OpenAIGPT5Config):
|
|||
drop_params=drop_params,
|
||||
)
|
||||
|
||||
# Only drop reasoning_effort='none' for models that don't support it
|
||||
result_effort: Final = _get_effort_level(result.get("reasoning_effort"))
|
||||
if result_effort == "none" and not supports_none:
|
||||
result.pop("reasoning_effort")
|
||||
|
||||
# Azure gpt-5.4+ with tools + reasoning_effort is now routed to the
|
||||
# Responses API bridge (same as OpenAI), so we no longer need to drop
|
||||
# reasoning_effort here. See: responses_api_bridge_check() in main.py.
|
||||
|
||||
return result
|
||||
|
||||
def transform_request(
|
||||
self,
|
||||
model: str,
|
||||
|
|
|
|||
|
|
@ -30,10 +30,18 @@ class AzureOpenAIResponsesAPIConfig(OpenAIResponsesAPIConfig):
|
|||
def custom_llm_provider(self) -> LlmProviders:
|
||||
return LlmProviders.AZURE
|
||||
|
||||
@staticmethod
|
||||
def _is_gpt_5_model(model: str) -> bool:
|
||||
return AzureOpenAIGPT5Config.is_model_gpt_5_model(model)
|
||||
|
||||
@staticmethod
|
||||
def _supports_reasoning_effort_none(model: str) -> bool:
|
||||
return AzureOpenAIGPT5Config._supports_reasoning_effort_level(model, "none")
|
||||
|
||||
@staticmethod
|
||||
def _effort_level_is_disabled(model: str, level: str) -> bool:
|
||||
return AzureOpenAIGPT5Config.row_disables_reasoning_effort_level(model, level)
|
||||
|
||||
@staticmethod
|
||||
def _effort_resolves_to_none(model: str, effort: str | None) -> bool:
|
||||
return AzureOpenAIGPT5Config.effort_resolves_to_none(model, effort)
|
||||
|
|
|
|||
|
|
@ -3,6 +3,7 @@ from typing import Final
|
|||
import httpx
|
||||
|
||||
from litellm.llms.azure.responses.transformation import AzureOpenAIResponsesAPIConfig
|
||||
from litellm.llms.azure_ai.chat.transformation import AzureAIGPT5Config
|
||||
from litellm.llms.azure_ai.common_utils import (
|
||||
AzureFoundryModelInfo,
|
||||
api_key_header_for_base,
|
||||
|
|
@ -27,6 +28,18 @@ class AzureAIResponsesAPIConfig(AzureOpenAIResponsesAPIConfig):
|
|||
def custom_llm_provider(self) -> LlmProviders:
|
||||
return LlmProviders.AZURE_AI
|
||||
|
||||
@staticmethod
|
||||
def _supports_reasoning_effort_none(model: str) -> bool:
|
||||
return AzureAIGPT5Config._supports_reasoning_effort_level(model, "none")
|
||||
|
||||
@staticmethod
|
||||
def _effort_level_is_disabled(model: str, level: str) -> bool:
|
||||
return AzureAIGPT5Config.row_disables_reasoning_effort_level(model, level)
|
||||
|
||||
@staticmethod
|
||||
def _effort_resolves_to_none(model: str, effort: str | None) -> bool:
|
||||
return AzureAIGPT5Config.effort_resolves_to_none(model, effort)
|
||||
|
||||
def validate_environment(self, headers: dict, model: str, litellm_params: GenericLiteLLMParams | None) -> dict:
|
||||
params: Final = litellm_params or GenericLiteLLMParams()
|
||||
auth_headers: Final = get_azure_ai_auth_headers(
|
||||
|
|
|
|||
|
|
@ -2659,13 +2659,13 @@ class AmazonConverseConfig(BaseConfig):
|
|||
|
||||
## HANDLE TOOL CALLS
|
||||
_message: Final = Message(**chat_completion_message)
|
||||
initial_finish_reason = map_finish_reason(completion_response["stopReason"])
|
||||
mapped_finish_reason: Final = map_finish_reason(completion_response["stopReason"])
|
||||
|
||||
# When json_mode filtered out all synthetic tool calls the response
|
||||
# is plain content, not a pending tool invocation. Fix finish_reason
|
||||
# so callers (e.g. OpenAI SDK) don't misinterpret it.
|
||||
if resolved_json_mode and not filtered_tools and tools:
|
||||
initial_finish_reason = "stop"
|
||||
initial_finish_reason: Final = (
|
||||
"stop"
|
||||
if resolved_json_mode and not filtered_tools and tools and mapped_finish_reason == "tool_calls"
|
||||
else mapped_finish_reason
|
||||
)
|
||||
|
||||
(
|
||||
returned_message,
|
||||
|
|
|
|||
|
|
@ -3,29 +3,103 @@ 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 pydantic import JsonValue, TypeAdapter
|
||||
from typing_extensions import assert_never
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
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
|
||||
from litellm.types.utils import LlmProviders, TokenCountMessagesFormat, TokenCountResponse
|
||||
from litellm.llms.bedrock.count_tokens.mantle_handler import BedrockMantleCountTokensHandler
|
||||
from litellm.llms.bedrock_mantle.common_utils import is_mantle_claude_model
|
||||
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
|
||||
from litellm.types.utils import LiteLLMPydanticObjectBase, LlmProviders, TokenCountMessagesFormat, TokenCountResponse
|
||||
|
||||
RUNTIME_TOKENIZER_TYPE: Final = "bedrock_api"
|
||||
MANTLE_TOKENIZER_TYPE: Final = "bedrock_mantle_api"
|
||||
_RUNTIME_FALLBACK_STATUSES: Final = frozenset({400, 404})
|
||||
_COUNT_RESPONSE: Final = TypeAdapter(dict[str, JsonValue])
|
||||
|
||||
|
||||
class _CountTokensReply(LiteLLMPydanticObjectBase):
|
||||
input_tokens: int
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class CountedTokens:
|
||||
input_tokens: int
|
||||
original_response: Mapping[str, JsonValue]
|
||||
tokenizer_type: str
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class CountTokensFailure:
|
||||
status_code: int
|
||||
message: str
|
||||
tokenizer_type: str
|
||||
|
||||
|
||||
CountTokensOutcome = CountedTokens | CountTokensFailure
|
||||
|
||||
|
||||
def _runtime_rejected_claude_model(outcome: CountTokensOutcome, resolved_model: str) -> bool:
|
||||
return (
|
||||
isinstance(outcome, CountTokensFailure)
|
||||
and outcome.status_code in _RUNTIME_FALLBACK_STATUSES
|
||||
and is_mantle_claude_model(resolved_model)
|
||||
)
|
||||
|
||||
|
||||
class BedrockTokenCounter(BaseTokenCounter):
|
||||
"""Token counter implementation for AWS Bedrock provider using the CountTokens API."""
|
||||
"""Token counter for AWS Bedrock: bedrock-runtime CountTokens first, and Anthropic's count_tokens
|
||||
on bedrock-mantle for the Claude models bedrock-runtime answers 400 or 404."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
runtime_handler: BedrockCountTokensHandler | None = None,
|
||||
mantle_handler: BedrockMantleCountTokensHandler | None = None,
|
||||
client: AsyncHTTPHandler | None = None,
|
||||
) -> None:
|
||||
self._runtime_handler: Final = runtime_handler or BedrockCountTokensHandler()
|
||||
self._mantle_handler: Final = mantle_handler or BedrockMantleCountTokensHandler()
|
||||
self._client: Final = client
|
||||
|
||||
def should_use_token_counting_api(
|
||||
self,
|
||||
custom_llm_provider: str | None = None,
|
||||
) -> bool:
|
||||
"""
|
||||
Returns True if we should use the Bedrock CountTokens API for token counting.
|
||||
"""
|
||||
return custom_llm_provider == LlmProviders.BEDROCK.value
|
||||
|
||||
async def _count_with(
|
||||
self,
|
||||
handler: BedrockCountTokensHandler | BedrockMantleCountTokensHandler,
|
||||
tokenizer_type: str,
|
||||
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,
|
||||
litellm_params=litellm_params,
|
||||
resolved_model=resolved_model,
|
||||
client=self._client,
|
||||
)
|
||||
reply: Final = _CountTokensReply.model_validate(result)
|
||||
except BedrockError as e:
|
||||
verbose_logger.debug(
|
||||
"%s CountTokens API error: status=%s, message=%s", tokenizer_type, e.status_code, e.message
|
||||
)
|
||||
return CountTokensFailure(status_code=e.status_code, message=e.message, tokenizer_type=tokenizer_type)
|
||||
except Exception as e:
|
||||
verbose_logger.debug("Error calling %s CountTokens API: %s", tokenizer_type, e)
|
||||
return CountTokensFailure(status_code=500, message=str(e), tokenizer_type=tokenizer_type)
|
||||
return CountedTokens(input_tokens=reply.input_tokens, original_response=result, tokenizer_type=tokenizer_type)
|
||||
|
||||
async def count_tokens(
|
||||
self,
|
||||
model_to_use: str,
|
||||
|
|
@ -38,81 +112,53 @@ class BedrockTokenCounter(BaseTokenCounter):
|
|||
messages_format: TokenCountMessagesFormat = "anthropic",
|
||||
output_format: JsonValue = None,
|
||||
) -> TokenCountResponse | None:
|
||||
"""
|
||||
Count tokens using AWS Bedrock's CountTokens API.
|
||||
|
||||
This method calls the existing BedrockCountTokensHandler to make an API call
|
||||
to Bedrock's token counting endpoint, bypassing the local tiktoken-based counting.
|
||||
|
||||
Args:
|
||||
model_to_use: The model identifier
|
||||
messages: The messages to count tokens for
|
||||
contents: Alternative content format (not used for Bedrock)
|
||||
deployment: Deployment configuration containing litellm_params
|
||||
request_model: The original request model name
|
||||
|
||||
Returns:
|
||||
TokenCountResponse with token count, or None if counting fails
|
||||
"""
|
||||
if not messages:
|
||||
return None
|
||||
|
||||
deployment = deployment or {}
|
||||
litellm_params: Final = deployment.get("litellm_params", {})
|
||||
|
||||
# Build request data in the format expected by BedrockCountTokensHandler
|
||||
request_data: Final[dict[str, object]] = {
|
||||
"model": model_to_use,
|
||||
"messages": messages,
|
||||
}
|
||||
|
||||
if tools:
|
||||
request_data["tools"] = tools
|
||||
|
||||
if system:
|
||||
request_data["system"] = system
|
||||
|
||||
# Get the resolved model (strip prefixes like bedrock/, converse/, etc.)
|
||||
litellm_params: Final = (deployment or {}).get("litellm_params", {})
|
||||
request_data: Final = MappingProxyType(
|
||||
{
|
||||
"model": model_to_use,
|
||||
"messages": messages,
|
||||
**MappingProxyType({name: value for name, value in (("tools", tools), ("system", system)) if value}),
|
||||
}
|
||||
)
|
||||
resolved_model: Final = get_bedrock_base_model(model_to_use)
|
||||
|
||||
try:
|
||||
handler: Final = BedrockCountTokensHandler()
|
||||
result: Final = await handler.handle_count_tokens_request(
|
||||
request_data=request_data,
|
||||
litellm_params=litellm_params,
|
||||
resolved_model=resolved_model,
|
||||
runtime: Final = await self._count_with(
|
||||
self._runtime_handler, RUNTIME_TOKENIZER_TYPE, request_data, litellm_params, resolved_model
|
||||
)
|
||||
outcome: Final = (
|
||||
await self._count_with(
|
||||
self._mantle_handler, MANTLE_TOKENIZER_TYPE, request_data, litellm_params, resolved_model
|
||||
)
|
||||
|
||||
# Transform response to TokenCountResponse
|
||||
if result is not None:
|
||||
if _runtime_rejected_claude_model(runtime, resolved_model)
|
||||
else runtime
|
||||
)
|
||||
match outcome:
|
||||
case CountedTokens():
|
||||
return TokenCountResponse(
|
||||
total_tokens=result.get("input_tokens", 0),
|
||||
total_tokens=outcome.input_tokens,
|
||||
request_model=request_model,
|
||||
model_used=model_to_use,
|
||||
tokenizer_type="bedrock_api",
|
||||
original_response=result,
|
||||
tokenizer_type=outcome.tokenizer_type,
|
||||
original_response=_COUNT_RESPONSE.validate_python(outcome.original_response),
|
||||
)
|
||||
except BedrockError as e:
|
||||
verbose_logger.warning("Bedrock 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="bedrock_api",
|
||||
error=True,
|
||||
error_message=e.message,
|
||||
status_code=e.status_code,
|
||||
)
|
||||
except Exception as e:
|
||||
verbose_logger.warning("Error calling Bedrock CountTokens API: %s", e)
|
||||
return TokenCountResponse(
|
||||
total_tokens=0,
|
||||
request_model=request_model,
|
||||
model_used=model_to_use,
|
||||
tokenizer_type="bedrock_api",
|
||||
error=True,
|
||||
error_message=str(e),
|
||||
status_code=500,
|
||||
)
|
||||
|
||||
return None
|
||||
case CountTokensFailure():
|
||||
verbose_logger.warning(
|
||||
"%s CountTokens API error: status=%s, message=%s",
|
||||
outcome.tokenizer_type,
|
||||
outcome.status_code,
|
||||
outcome.message,
|
||||
)
|
||||
return TokenCountResponse(
|
||||
total_tokens=0,
|
||||
request_model=request_model,
|
||||
model_used=model_to_use,
|
||||
tokenizer_type=outcome.tokenizer_type,
|
||||
error=True,
|
||||
error_message=outcome.message,
|
||||
status_code=outcome.status_code,
|
||||
)
|
||||
case _:
|
||||
assert_never(outcome)
|
||||
|
|
|
|||
|
|
@ -4,6 +4,7 @@ AWS Bedrock CountTokens API handler.
|
|||
Simplified handler leveraging existing LiteLLM Bedrock infrastructure.
|
||||
"""
|
||||
|
||||
from collections.abc import Mapping
|
||||
from typing import Any, Final
|
||||
|
||||
import httpx
|
||||
|
|
@ -25,7 +26,7 @@ class BedrockCountTokensHandler(BedrockCountTokensConfig):
|
|||
|
||||
async def handle_count_tokens_request(
|
||||
self,
|
||||
request_data: dict[str, Any],
|
||||
request_data: Mapping[str, Any],
|
||||
litellm_params: dict[str, Any],
|
||||
resolved_model: str,
|
||||
client: AsyncHTTPHandler | None = None,
|
||||
|
|
@ -125,7 +126,7 @@ class BedrockCountTokensHandler(BedrockCountTokensConfig):
|
|||
raise
|
||||
except httpx.HTTPStatusError as e:
|
||||
# HTTP errors - preserve the actual status code
|
||||
verbose_logger.error("HTTP error in CountTokens handler: %s", e)
|
||||
verbose_logger.debug("HTTP error in CountTokens handler: %s", e)
|
||||
raise BedrockError(
|
||||
status_code=e.response.status_code,
|
||||
message=e.response.text,
|
||||
|
|
|
|||
113
litellm/llms/bedrock/count_tokens/mantle_handler.py
Normal file
113
litellm/llms/bedrock/count_tokens/mantle_handler.py
Normal file
|
|
@ -0,0 +1,113 @@
|
|||
from collections.abc import Mapping
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
from pydantic import JsonValue, TypeAdapter
|
||||
from typing_extensions import NotRequired, ReadOnly, TypedDict
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_logger
|
||||
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
|
||||
|
||||
MANTLE_COUNT_TOKENS_SUFFIX: Final = "/count_tokens"
|
||||
MANTLE_ANTHROPIC_VERSION: Final = "2023-06-01"
|
||||
|
||||
|
||||
class MantleCountTokensRequest(TypedDict):
|
||||
messages: ReadOnly[list[dict[str, JsonValue]]]
|
||||
system: ReadOnly[NotRequired[JsonValue]]
|
||||
tools: ReadOnly[NotRequired[list[dict[str, JsonValue]]]]
|
||||
|
||||
|
||||
_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):
|
||||
"""Counts tokens through Anthropic's count_tokens on the bedrock-mantle endpoint.
|
||||
|
||||
Claude models that Bedrock offers only through cross-region inference answer 400 on
|
||||
bedrock-runtime's CountTokens; AWS documents Mantle's /anthropic/v1/messages/count_tokens
|
||||
as the way to count them, with the base model id and the deployment's AWS credentials
|
||||
"""
|
||||
|
||||
def __init__(self, anthropic_config: AnthropicCountTokensConfig | None = None) -> None:
|
||||
super().__init__()
|
||||
self._anthropic_config: Final = anthropic_config or AnthropicCountTokensConfig()
|
||||
|
||||
def get_mantle_count_tokens_endpoint(self, aws_region_name: str) -> str:
|
||||
messages_url: Final = build_mantle_messages_url(
|
||||
api_base=None, aws_bedrock_runtime_endpoint=None, region=aws_region_name
|
||||
)
|
||||
return f"{messages_url}{MANTLE_COUNT_TOKENS_SUFFIX}"
|
||||
|
||||
async def handle_count_tokens_request(
|
||||
self,
|
||||
request_data: Mapping[str, object],
|
||||
litellm_params: dict[str, object],
|
||||
resolved_model: str,
|
||||
client: AsyncHTTPHandler | None = None,
|
||||
) -> 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
|
||||
)
|
||||
endpoint_url: Final = 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(
|
||||
MappingProxyType(
|
||||
{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")
|
||||
signed_headers, signed_body = await run_aws_signing(
|
||||
self._sign_request,
|
||||
service_name="bedrock",
|
||||
headers={ # mutable-ok: [LIT002] bearer signing may add the Authorization header
|
||||
"Content-Type": "application/json",
|
||||
"anthropic-version": MANTLE_ANTHROPIC_VERSION,
|
||||
},
|
||||
optional_params=litellm_params,
|
||||
request_data=body,
|
||||
api_base=endpoint_url,
|
||||
model=resolved_model,
|
||||
api_key=api_key if isinstance(api_key, str) else None,
|
||||
)
|
||||
async_client: Final = client or get_async_httpx_client(llm_provider=litellm.LlmProviders.BEDROCK)
|
||||
response: Final = await async_client.post(
|
||||
endpoint_url, headers=signed_headers, data=signed_body, timeout=30.0
|
||||
)
|
||||
if response.status_code != 200:
|
||||
raise BedrockError(
|
||||
status_code=response.status_code,
|
||||
message=response.text,
|
||||
headers=response.headers,
|
||||
response=response,
|
||||
)
|
||||
return _COUNT_RESPONSE.validate_json(response.content)
|
||||
except BedrockError:
|
||||
raise
|
||||
except httpx.HTTPStatusError as e:
|
||||
raise BedrockError(
|
||||
status_code=e.response.status_code,
|
||||
message=e.response.text,
|
||||
headers=e.response.headers,
|
||||
response=e.response,
|
||||
)
|
||||
except Exception as e:
|
||||
raise BedrockError(status_code=500, message=f"bedrock-mantle count_tokens processing error: {e}")
|
||||
|
|
@ -45,6 +45,10 @@ def split_mantle_region_prefix(model: str) -> tuple[str | None, str]:
|
|||
return None, model
|
||||
|
||||
|
||||
def is_mantle_claude_model(model: str) -> bool:
|
||||
return "claude" in model.lower()
|
||||
|
||||
|
||||
def resolve_mantle_region(params: Mapping[str, object]) -> str:
|
||||
region: Final = params.get("aws_region_name")
|
||||
if isinstance(region, str) and region:
|
||||
|
|
|
|||
|
|
@ -3,7 +3,7 @@ import json
|
|||
import ssl
|
||||
from collections.abc import AsyncGenerator, AsyncIterator, Coroutine, Iterator, Mapping, Sequence
|
||||
from contextlib import asynccontextmanager
|
||||
from functools import lru_cache, partial
|
||||
from functools import lru_cache, partial, reduce
|
||||
from types import MappingProxyType, ModuleType
|
||||
from typing import (
|
||||
TYPE_CHECKING,
|
||||
|
|
@ -32,7 +32,7 @@ import litellm.litellm_core_utils
|
|||
import litellm.types
|
||||
import litellm.types.utils
|
||||
from litellm._logging import _redact_string, verbose_logger
|
||||
from litellm.anthropic_beta_headers_manager import update_headers_with_filtered_beta
|
||||
from litellm.anthropic_beta_headers_manager import aload_beta_headers_config, update_headers_with_filtered_beta
|
||||
from litellm.constants import MAX_FILE_LIST_LIMIT, REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES
|
||||
from litellm.files.types import FileContentStreamingResult
|
||||
from litellm.litellm_core_utils.agentic_followup_kwargs import build_agentic_followup_kwargs
|
||||
|
|
@ -46,6 +46,7 @@ from litellm.litellm_core_utils.audio_utils.subtitle_utils import (
|
|||
synthesize_subtitle_document,
|
||||
)
|
||||
from litellm.litellm_core_utils.core_helpers import set_provider_response_headers_in_hidden_params
|
||||
from litellm.litellm_core_utils.dot_notation_indexing import delete_nested_value
|
||||
from litellm.litellm_core_utils.get_litellm_params import AWS_CREDENTIAL_KWARGS_KEYS
|
||||
from litellm.litellm_core_utils.llm_request_utils import serialize_multipart_form_fields
|
||||
from litellm.litellm_core_utils.realtime_errors import (
|
||||
|
|
@ -2116,6 +2117,16 @@ class BaseLLMHTTPHandler:
|
|||
merged_headers.update(extra_headers_from_kwargs)
|
||||
if provider_specific_headers:
|
||||
merged_headers.update(provider_specific_headers)
|
||||
request_params: Final = reduce(
|
||||
delete_nested_value,
|
||||
litellm_params.get("additional_drop_params") or (),
|
||||
anthropic_messages_optional_request_params,
|
||||
)
|
||||
if (
|
||||
request_params.get("safeguards") is not None
|
||||
and anthropic_messages_provider_config.should_filter_anthropic_beta_headers()
|
||||
):
|
||||
await aload_beta_headers_config()
|
||||
(
|
||||
headers,
|
||||
api_base,
|
||||
|
|
@ -2123,7 +2134,7 @@ class BaseLLMHTTPHandler:
|
|||
headers=merged_headers or {},
|
||||
model=model,
|
||||
messages=messages,
|
||||
optional_params=anthropic_messages_optional_request_params,
|
||||
optional_params=request_params,
|
||||
litellm_params=dict(litellm_params),
|
||||
api_key=api_key,
|
||||
api_base=api_base,
|
||||
|
|
@ -2154,20 +2165,11 @@ class BaseLLMHTTPHandler:
|
|||
custom_llm_provider=custom_llm_provider,
|
||||
)
|
||||
|
||||
additional_drop_params: Final[list[str]] = litellm_params.get("additional_drop_params") or []
|
||||
if additional_drop_params:
|
||||
from litellm.litellm_core_utils.dot_notation_indexing import delete_nested_value
|
||||
|
||||
for path in additional_drop_params:
|
||||
anthropic_messages_optional_request_params = delete_nested_value(
|
||||
anthropic_messages_optional_request_params, path
|
||||
)
|
||||
|
||||
# Prepare request body
|
||||
request_body: Final = anthropic_messages_provider_config.transform_anthropic_messages_request(
|
||||
model=model,
|
||||
messages=messages,
|
||||
anthropic_messages_optional_request_params=anthropic_messages_optional_request_params,
|
||||
anthropic_messages_optional_request_params=request_params,
|
||||
litellm_params=litellm_params,
|
||||
headers=headers,
|
||||
)
|
||||
|
|
@ -2273,7 +2275,7 @@ class BaseLLMHTTPHandler:
|
|||
|
||||
held_back_tool_names: Final = self._server_fulfilled_tools_in_request(
|
||||
logging_obj=logging_obj,
|
||||
tools=anthropic_messages_optional_request_params.get("tools"),
|
||||
tools=request_params.get("tools"),
|
||||
)
|
||||
initial_response = AgenticAnthropicStreamingIterator(
|
||||
completion_stream=completion_stream,
|
||||
|
|
@ -2281,7 +2283,7 @@ class BaseLLMHTTPHandler:
|
|||
model=model,
|
||||
messages=messages,
|
||||
anthropic_messages_provider_config=anthropic_messages_provider_config,
|
||||
anthropic_messages_optional_request_params=anthropic_messages_optional_request_params,
|
||||
anthropic_messages_optional_request_params=request_params,
|
||||
logging_obj=logging_obj,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
kwargs=kwargs_for_agentic,
|
||||
|
|
@ -2304,7 +2306,7 @@ class BaseLLMHTTPHandler:
|
|||
model=model,
|
||||
messages=messages,
|
||||
anthropic_messages_provider_config=anthropic_messages_provider_config,
|
||||
anthropic_messages_optional_request_params=anthropic_messages_optional_request_params,
|
||||
anthropic_messages_optional_request_params=request_params,
|
||||
logging_obj=logging_obj,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
kwargs=kwargs_for_agentic,
|
||||
|
|
|
|||
|
|
@ -7,7 +7,7 @@ from collections.abc import AsyncIterator, Coroutine, Iterator, Mapping, Sequenc
|
|||
from typing import TYPE_CHECKING, Any, Final, Literal, cast, overload
|
||||
|
||||
import httpx
|
||||
from pydantic import BaseModel
|
||||
from pydantic import BaseModel, TypeAdapter
|
||||
|
||||
from litellm.constants import RESPONSE_FORMAT_TOOL_NAME
|
||||
from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import (
|
||||
|
|
@ -21,6 +21,9 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import (
|
|||
strip_name_from_message,
|
||||
)
|
||||
from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator
|
||||
from litellm.llms.base_llm.base_utils import (
|
||||
type_to_response_format_param, # pyright: ignore[reportUnknownVariableType] # base_utils helper returns an untyped dict
|
||||
)
|
||||
from litellm.types.llms.anthropic import AllAnthropicToolsValues
|
||||
from litellm.types.llms.databricks import (
|
||||
AllDatabricksContentValues,
|
||||
|
|
@ -59,6 +62,10 @@ from ...anthropic.chat.transformation import (
|
|||
from ...openai_like.chat.transformation import OpenAILikeChatConfig
|
||||
from ..common_utils import DatabricksBase, DatabricksException
|
||||
|
||||
_RESPONSE_FORMAT_ADAPTER: ( # mutable-ok: mirrors the dict return contract of get_json_schema_from_pydantic_object
|
||||
Final[TypeAdapter[dict[str, object] | None]]
|
||||
) = TypeAdapter(dict[str, object] | None)
|
||||
|
||||
|
||||
def _is_bare_assistant_message(message_dict: Mapping[str, object]) -> bool:
|
||||
"""Databricks rejects assistant messages with neither content nor tool calls, e.g. a replayed
|
||||
|
|
@ -68,6 +75,15 @@ def _is_bare_assistant_message(message_dict: Mapping[str, object]) -> bool:
|
|||
)
|
||||
|
||||
|
||||
def _collapse_single_text_block(message: Mapping[str, object]) -> Mapping[str, object]:
|
||||
content: Final = message.get("content")
|
||||
match content:
|
||||
case [{"type": "text", "text": str() as text}] if content == [{"type": "text", "text": text}]:
|
||||
return {**message, "content": text}
|
||||
case _:
|
||||
return message
|
||||
|
||||
|
||||
def _sanitize_empty_content(message_dict: dict[str, object]) -> None:
|
||||
"""
|
||||
Remove or filter content so empty text blocks are not sent.
|
||||
|
|
@ -191,6 +207,14 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig):
|
|||
def get_config(cls, *, model: str | None = None):
|
||||
return super().get_config()
|
||||
|
||||
def get_json_schema_from_pydantic_object(
|
||||
self,
|
||||
response_format: ( # mutable-ok: matches BaseConfig override signature
|
||||
type[BaseModel] | dict[str, object] | None
|
||||
),
|
||||
) -> dict[str, object] | None: # mutable-ok: BaseConfig contract returns a dict
|
||||
return _RESPONSE_FORMAT_ADAPTER.validate_python(type_to_response_format_param(response_format=response_format))
|
||||
|
||||
def get_required_params(self) -> list[ProviderField]:
|
||||
"""For a given provider, return it's required fields with a description"""
|
||||
return [
|
||||
|
|
@ -447,6 +471,7 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig):
|
|||
Databricks does not support:
|
||||
- 'name' in user message.
|
||||
- litellm's internal `thinking_blocks` / `reasoning_content` on assistant messages.
|
||||
- non-Claude json_object "json" checks read only string content, so a single text block is sent as a string.
|
||||
"""
|
||||
new_messages = []
|
||||
for idx, message in enumerate(messages):
|
||||
|
|
@ -462,7 +487,7 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig):
|
|||
_sanitize_empty_content(cast(dict[str, Any], _message))
|
||||
if _is_bare_assistant_message(_message):
|
||||
continue
|
||||
new_messages.append(_message)
|
||||
new_messages.append(_collapse_single_text_block(_message) if "claude" not in model else _message)
|
||||
|
||||
if "claude" not in model:
|
||||
new_messages = _split_parallel_tool_calls(
|
||||
|
|
|
|||
|
|
@ -1,7 +1,7 @@
|
|||
"""Support for OpenAI gpt-5 model family."""
|
||||
|
||||
import re
|
||||
from typing import Final
|
||||
from typing import ClassVar, Final
|
||||
|
||||
import litellm
|
||||
from litellm.utils import (
|
||||
|
|
@ -67,6 +67,10 @@ def _get_effort_level(value: str | dict | None) -> str | None:
|
|||
GPT_REASONING_SERIES_MARKERS: Final = ("gpt-5", "gpt-6")
|
||||
|
||||
|
||||
def unsupported_reasoning_effort_message(param: str, level: str, model: str) -> str:
|
||||
return f"{param}={level} is not supported for {model}. To drop unsupported params set `litellm.drop_params = True`"
|
||||
|
||||
|
||||
def is_gpt_reasoning_series_name(model: str) -> bool:
|
||||
normalized: Final = model.split("/")[-1]
|
||||
return any(marker in model for marker in GPT_REASONING_SERIES_MARKERS) and not normalized.startswith("gpt-5-chat")
|
||||
|
|
@ -82,6 +86,9 @@ class OpenAIGPT5Config(OpenAIGPTConfig):
|
|||
- Support for GPT-5-Codex models optimized for code generation.
|
||||
"""
|
||||
|
||||
OPT_IN_REASONING_EFFORTS: ClassVar[frozenset[str]] = frozenset(("xhigh",))
|
||||
ROUTE_SPECIFIC_REASONING_EFFORTS: ClassVar[frozenset[str]] = frozenset(("max",))
|
||||
|
||||
@classmethod
|
||||
def is_model_gpt_5_model(cls, model: str) -> bool:
|
||||
return is_gpt_reasoning_series_name(model)
|
||||
|
|
@ -193,15 +200,10 @@ class OpenAIGPT5Config(OpenAIGPTConfig):
|
|||
return False
|
||||
|
||||
@classmethod
|
||||
def _is_reasoning_effort_level_explicitly_disabled(cls, model: str, level: str) -> bool:
|
||||
def is_reasoning_effort_level_explicitly_disabled(cls, model: str, level: str) -> bool:
|
||||
"""Return True only when the model map explicitly sets the capability to False.
|
||||
|
||||
Unlike ``_supports_reasoning_effort_level`` (which requires an explicit True),
|
||||
this method returns True only when ``supports_{level}_reasoning_effort`` is
|
||||
explicitly set to ``False`` in the model map. A missing key is treated as
|
||||
supported (i.e. this method returns False = not disabled).
|
||||
|
||||
Use this for opt-out checks where unknown models should be allowed through.
|
||||
A missing key is not a disable. Unknown models pass through.
|
||||
"""
|
||||
return is_explicitly_disabled_factory(
|
||||
model=cls._model_map_lookup_name(model),
|
||||
|
|
@ -209,6 +211,39 @@ class OpenAIGPT5Config(OpenAIGPTConfig):
|
|||
key=f"supports_{level}_reasoning_effort",
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def _is_reasoning_effort_level_explicitly_disabled(cls, model: str, level: str) -> bool:
|
||||
return cls.is_reasoning_effort_level_explicitly_disabled(model, level)
|
||||
|
||||
@classmethod
|
||||
def refuses_reasoning_effort_level(cls, model: str, level: str) -> bool:
|
||||
"""The chat gate for any requested effort level, read off the model's map row alone.
|
||||
|
||||
An opt-in level (xhigh here; Azure adds none) is refused unless the row says true, since
|
||||
the providers rejected it on most of the family when it arrived. Every other level is
|
||||
refused only when the row says false: a model the map never heard of, or a row without the
|
||||
flag, keeps forwarding what it was sent, and a level with no flag at all (medium, high) is
|
||||
never refused. So a new model in the family needs its map row and no code change, the
|
||||
row's supports_<level>_reasoning_effort flags being the whole contract.
|
||||
|
||||
A route-specific level (max) is never refused by either gate: the same model honors it
|
||||
on Responses and refuses it on Chat Completions, so one row flag cannot gate it on
|
||||
either route. Chat forwards it for the Responses bridge and a plain chat call gets the
|
||||
provider's own 400, exactly as before the gate read every level.
|
||||
"""
|
||||
if level in cls.ROUTE_SPECIFIC_REASONING_EFFORTS:
|
||||
return False
|
||||
if level in cls.OPT_IN_REASONING_EFFORTS:
|
||||
return not cls._supports_reasoning_effort_level(model, level)
|
||||
return cls.is_reasoning_effort_level_explicitly_disabled(model, level)
|
||||
|
||||
@classmethod
|
||||
def row_disables_reasoning_effort_level(cls, model: str, level: str) -> bool:
|
||||
"""The Responses gate: opt-out for every level except the route-specific ones."""
|
||||
if level in cls.ROUTE_SPECIFIC_REASONING_EFFORTS:
|
||||
return False
|
||||
return cls.is_reasoning_effort_level_explicitly_disabled(model, level)
|
||||
|
||||
def get_supported_openai_params(self, model: str) -> list:
|
||||
if self.is_model_gpt_5_search_model(model):
|
||||
return [
|
||||
|
|
@ -269,12 +304,12 @@ class OpenAIGPT5Config(OpenAIGPTConfig):
|
|||
drop_params=drop_params,
|
||||
)
|
||||
|
||||
# Get raw reasoning_effort and effective effort level for all guards.
|
||||
# Use effective_effort (extracted string) for xhigh validation, "none" checks, and
|
||||
# tool/sampling guards — dict inputs like {"effort": "none", "summary": "detailed"}
|
||||
# must be treated as effort="none" to avoid incorrect tool-drop or sampling errors.
|
||||
raw_reasoning_effort = non_default_params.get("reasoning_effort") or optional_params.get("reasoning_effort")
|
||||
effective_effort: Final = _get_effort_level(raw_reasoning_effort)
|
||||
# Dict inputs like {"effort": "none", "summary": "detailed"} are judged by their effort
|
||||
# string so the gate, the sampling guard, and the temperature guard all see one level.
|
||||
raw_reasoning_effort: Final = non_default_params.get("reasoning_effort") or optional_params.get(
|
||||
"reasoning_effort"
|
||||
)
|
||||
requested_effort: Final = _get_effort_level(raw_reasoning_effort)
|
||||
|
||||
# Normalize dict reasoning_effort to string for Chat Completions API.
|
||||
# Example: {"effort": "high", "summary": "detailed"} -> "high"
|
||||
|
|
@ -286,31 +321,20 @@ class OpenAIGPT5Config(OpenAIGPTConfig):
|
|||
if "reasoning_effort" in optional_params:
|
||||
optional_params["reasoning_effort"] = normalized
|
||||
|
||||
if effective_effort == "xhigh":
|
||||
# xhigh is an opt-in capability: only allow if model explicitly supports it.
|
||||
if not self._supports_reasoning_effort_level(model, effective_effort):
|
||||
if litellm.drop_params or drop_params:
|
||||
non_default_params.pop("reasoning_effort", None)
|
||||
optional_params.pop("reasoning_effort", None)
|
||||
else:
|
||||
raise litellm.utils.UnsupportedParamsError(
|
||||
message=(f"reasoning_effort={effective_effort} is not supported for this model."),
|
||||
status_code=400,
|
||||
)
|
||||
elif effective_effort in ("minimal", "low"):
|
||||
# minimal/low are opt-out: unknown models pass through; only block when
|
||||
# the model map explicitly sets supports_{level}_reasoning_effort=false.
|
||||
# Example: gpt-5.5-pro only accepts {medium, high, xhigh}, so it sets
|
||||
# supports_low_reasoning_effort=false (and supports_minimal=false).
|
||||
if self._is_reasoning_effort_level_explicitly_disabled(model, effective_effort):
|
||||
if litellm.drop_params or drop_params:
|
||||
non_default_params.pop("reasoning_effort", None)
|
||||
optional_params.pop("reasoning_effort", None)
|
||||
else:
|
||||
raise litellm.utils.UnsupportedParamsError(
|
||||
message=(f"reasoning_effort={effective_effort} is not supported for this model."),
|
||||
status_code=400,
|
||||
)
|
||||
if requested_effort is not None and self.refuses_reasoning_effort_level(model, requested_effort):
|
||||
if litellm.drop_params or drop_params:
|
||||
non_default_params.pop("reasoning_effort", None)
|
||||
optional_params.pop("reasoning_effort", None)
|
||||
else:
|
||||
raise litellm.utils.UnsupportedParamsError(
|
||||
message=unsupported_reasoning_effort_message("reasoning_effort", requested_effort, model),
|
||||
status_code=400,
|
||||
)
|
||||
|
||||
# The level the provider will see: None once dropped, so the guards below judge by the model's default effort
|
||||
effective_effort: Final = _get_effort_level(
|
||||
non_default_params.get("reasoning_effort") or optional_params.get("reasoning_effort")
|
||||
)
|
||||
|
||||
################################################################
|
||||
# max_tokens is not supported for gpt-5 models on OpenAI API
|
||||
|
|
|
|||
|
|
@ -22,7 +22,10 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import (
|
|||
from litellm.litellm_core_utils.safe_json_loads import safe_json_loads
|
||||
from litellm.litellm_core_utils.url_utils import encode_url_path_segment
|
||||
from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig
|
||||
from litellm.llms.openai.chat.gpt_5_transformation import is_gpt_reasoning_series_name
|
||||
from litellm.llms.openai.chat.gpt_5_transformation import (
|
||||
is_gpt_reasoning_series_name,
|
||||
unsupported_reasoning_effort_message,
|
||||
)
|
||||
from litellm.responses.litellm_completion_transformation.custom_tools import TOOL_CALL_ITEM_ID_PREFIX_BY_TYPE
|
||||
from litellm.secret_managers.main import get_secret_str
|
||||
from litellm.types.llms.openai import *
|
||||
|
|
@ -136,6 +139,20 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig):
|
|||
|
||||
return supports_none_reasoning_effort(model=model, custom_llm_provider=None)
|
||||
|
||||
@staticmethod
|
||||
def _effort_level_is_disabled(model: str, level: str) -> bool:
|
||||
"""Whether the model's cost-map row turns this effort level off.
|
||||
|
||||
The same row the chat gate reads, so a level the row refuses is refused on both wires
|
||||
and a new model needs its row, not code. This surface stays opt-out for every level,
|
||||
xhigh included, since it never refused a level the row omits; max is the exception, a
|
||||
level the row cannot gate because the same model takes it here and refuses it on chat.
|
||||
Azure overrides this so a bare deployment name reads the azure/ entry.
|
||||
"""
|
||||
from litellm.llms.openai.chat.gpt_5_transformation import OpenAIGPT5Config
|
||||
|
||||
return OpenAIGPT5Config.row_disables_reasoning_effort_level(model, level)
|
||||
|
||||
@staticmethod
|
||||
def _effort_resolves_to_none(model: str, effort: str | None) -> bool:
|
||||
"""Whether this request's reasoning effort ends up as "none", the one condition
|
||||
|
|
@ -242,7 +259,25 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig):
|
|||
lookup_name: Final = self._model_map_lookup_name(model)
|
||||
if self._is_gpt_5_model(model=lookup_name):
|
||||
reasoning: Final = params.get("reasoning") or {}
|
||||
effort: Final = reasoning.get("effort") if isinstance(reasoning, dict) else None
|
||||
requested_effort: Final = reasoning.get("effort") if isinstance(reasoning, dict) else None
|
||||
if (
|
||||
isinstance(reasoning, dict)
|
||||
and isinstance(requested_effort, str)
|
||||
and self._effort_level_is_disabled(lookup_name, requested_effort)
|
||||
):
|
||||
if not (drop_params or litellm.drop_params):
|
||||
raise litellm.UnsupportedParamsError(
|
||||
message=unsupported_reasoning_effort_message("reasoning.effort", requested_effort, model),
|
||||
status_code=400,
|
||||
)
|
||||
remaining: Final = reasoning.copy()
|
||||
remaining.pop("effort", None)
|
||||
if remaining:
|
||||
params["reasoning"] = remaining
|
||||
else:
|
||||
params.pop("reasoning", None)
|
||||
forwarded_reasoning: Final = params.get("reasoning")
|
||||
effort: Final = forwarded_reasoning.get("effort") if isinstance(forwarded_reasoning, dict) else None
|
||||
supports_none: Final = self._supports_reasoning_effort_none(model=lookup_name)
|
||||
effort_is_none: Final = supports_none and self._effort_resolves_to_none(lookup_name, effort)
|
||||
|
||||
|
|
|
|||
|
|
@ -8,6 +8,7 @@ from litellm.llms.anthropic.experimental_pass_through.messages.transformation im
|
|||
from litellm.llms.openai_like.json_loader import SimpleProviderConfig
|
||||
from litellm.secret_managers.main import get_secret_str
|
||||
from litellm.types.router import GenericLiteLLMParams
|
||||
from litellm.types.utils import LlmProviders
|
||||
|
||||
DEFAULT_ANTHROPIC_API_VERSION: Final = "2023-06-01"
|
||||
|
||||
|
|
@ -57,6 +58,7 @@ class OpenAILikeAnthropicMessagesConfig(AnthropicMessagesConfig):
|
|||
headers=normalized,
|
||||
optional_params=optional_params,
|
||||
messages=messages,
|
||||
beta_map_provider=LlmProviders.OPENAI_LIKE.value,
|
||||
)
|
||||
return merged, api_base
|
||||
|
||||
|
|
|
|||
|
|
@ -54,6 +54,8 @@ class NativeVertexBatchCostCalculator(Protocol):
|
|||
model: str,
|
||||
custom_llm_provider: str | None = None,
|
||||
model_info: ModelInfo | None = None,
|
||||
*,
|
||||
vertex_location: str | None = None,
|
||||
) -> tuple[float, float]: ...
|
||||
|
||||
|
||||
|
|
@ -90,6 +92,7 @@ def native_vertex_batch_row_stats(
|
|||
model_info: ModelInfo | None,
|
||||
calculate_usage: Callable[[GenerateContentResponseBody], Usage],
|
||||
cost_calculator: NativeVertexBatchCostCalculator,
|
||||
vertex_location: str | None = None,
|
||||
) -> NativeVertexBatchRowStats | None:
|
||||
"""
|
||||
Usage and cost of one native Vertex predictions.jsonl row, a
|
||||
|
|
@ -121,7 +124,11 @@ def native_vertex_batch_row_stats(
|
|||
)
|
||||
try:
|
||||
prompt_cost, completion_cost = cost_calculator(
|
||||
usage=usage, model=model, custom_llm_provider="vertex_ai", model_info=model_info
|
||||
usage=usage,
|
||||
model=model,
|
||||
custom_llm_provider="vertex_ai",
|
||||
model_info=model_info,
|
||||
vertex_location=vertex_location,
|
||||
)
|
||||
except Exception as e: # noqa: BLE001 # one unpriceable row must not abort the batch's cost accounting
|
||||
verbose_logger.warning(
|
||||
|
|
|
|||
|
|
@ -1,14 +1,34 @@
|
|||
import types
|
||||
from typing import Final, Literal
|
||||
|
||||
from pydantic import BaseModel
|
||||
from pydantic import BaseModel, TypeAdapter
|
||||
from typing_extensions import ReadOnly, TypedDict
|
||||
|
||||
from litellm.llms.vertex_ai.common_utils import pop_vertex_request_labels
|
||||
from litellm.types.utils import EmbeddingResponse, Usage
|
||||
from litellm.types.utils import EmbeddingResponse, PromptTokensDetailsWrapper, Usage
|
||||
|
||||
from .types import *
|
||||
|
||||
|
||||
class _PredictMetadata(TypedDict, total=False):
|
||||
billableCharacterCount: ReadOnly[int]
|
||||
|
||||
|
||||
class _PredictResponseMetadata(TypedDict, total=False):
|
||||
metadata: ReadOnly[_PredictMetadata]
|
||||
|
||||
|
||||
_PREDICT_RESPONSE_METADATA: Final = TypeAdapter(_PredictResponseMetadata)
|
||||
|
||||
|
||||
def billable_character_count(response: _PredictResponseMetadata) -> int | None:
|
||||
"""
|
||||
`metadata.billableCharacterCount` of a Vertex text embedding predict response, the count Google
|
||||
bills per-character embedding models (text-embedding-005, text-multilingual-embedding-002) on.
|
||||
"""
|
||||
return response.get("metadata", {}).get("billableCharacterCount")
|
||||
|
||||
|
||||
class VertexAITextEmbeddingConfig(BaseModel):
|
||||
"""
|
||||
Reference: https://cloud.google.com/vertex-ai/generative-ai/docs/model-reference/text-embeddings-api#TextEmbeddingInput
|
||||
|
|
@ -248,7 +268,15 @@ class VertexAITextEmbeddingConfig(BaseModel):
|
|||
model_response.object = "list"
|
||||
model_response.data = embedding_response
|
||||
model_response.model = model
|
||||
usage: Final = Usage(prompt_tokens=input_tokens, completion_tokens=0, total_tokens=input_tokens)
|
||||
billable_characters: Final = billable_character_count(_PREDICT_RESPONSE_METADATA.validate_python(response))
|
||||
usage: Final = Usage(
|
||||
prompt_tokens=input_tokens,
|
||||
completion_tokens=0,
|
||||
total_tokens=input_tokens,
|
||||
prompt_tokens_details=(
|
||||
None if billable_characters is None else PromptTokensDetailsWrapper(character_count=billable_characters)
|
||||
),
|
||||
)
|
||||
setattr(model_response, "usage", usage)
|
||||
return model_response
|
||||
|
||||
|
|
|
|||
|
|
@ -1057,6 +1057,18 @@ def _resolve_openai_api_base(api_base: str | None) -> str:
|
|||
)
|
||||
|
||||
|
||||
def _chat_gate_refuses_reasoning_effort_none(model: str, custom_llm_provider: str) -> bool:
|
||||
from litellm.llms.azure_ai.chat.transformation import AzureAIGPT5Config
|
||||
|
||||
match custom_llm_provider:
|
||||
case "azure":
|
||||
return litellm.AzureOpenAIGPT5Config.refuses_reasoning_effort_level(model, "none")
|
||||
case "azure_ai":
|
||||
return AzureAIGPT5Config.refuses_reasoning_effort_level(model, "none")
|
||||
case _:
|
||||
return OpenAIGPT5Config.refuses_reasoning_effort_level(model, "none")
|
||||
|
||||
|
||||
def responses_api_bridge_check(
|
||||
model: str,
|
||||
custom_llm_provider: str,
|
||||
|
|
@ -1129,10 +1141,13 @@ def responses_api_bridge_check(
|
|||
)
|
||||
for tool in (tools or ())
|
||||
)
|
||||
if isinstance(reasoning_effort, dict):
|
||||
reasoning_active = reasoning_effort.get("effort") != "none" or reasoning_effort.get("summary") is not None
|
||||
else:
|
||||
reasoning_active = reasoning_effort != "none"
|
||||
requested_effort: Final = reasoning_effort.get("effort") if isinstance(reasoning_effort, dict) else reasoning_effort
|
||||
requested_summary: Final = reasoning_effort.get("summary") if isinstance(reasoning_effort, dict) else None
|
||||
reasoning_active: Final = (
|
||||
requested_effort != "none"
|
||||
or requested_summary is not None
|
||||
or _chat_gate_refuses_reasoning_effort_none(model, custom_llm_provider)
|
||||
)
|
||||
# The reasoning+tools constraint is enforced by the real OpenAI backend behind any api.openai.com
|
||||
# host (the default URL or a PrivateLink hostname such as <region>.privatelink.api.openai.com) and
|
||||
# by Azure OpenAI through the azure provider. Resolve the effective OpenAI base arg>global>env>default
|
||||
|
|
@ -5707,6 +5722,7 @@ def completion(
|
|||
optional_params=processed_non_default_params, # [IMPORTANT] - using processed_non_default_params ensures consistent params logged to langfuse for finetuning / eval datasets.
|
||||
litellm_params=litellm_params,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
model_id=optional_params.get("model_id"),
|
||||
)
|
||||
if mock_response or mock_tool_calls or mock_timeout:
|
||||
kwargs.pop("mock_timeout", None) # remove for any fallbacks triggered
|
||||
|
|
|
|||
File diff suppressed because it is too large
Load diff
|
|
@ -5,6 +5,8 @@ from typing import Final
|
|||
|
||||
from pydantic import TypeAdapter
|
||||
|
||||
from litellm.litellm_core_utils.request_timeout_resolver import get_service_tier_timeout
|
||||
|
||||
DEFAULT_PASS_THROUGH_REQUEST_TIMEOUT_SECONDS: Final = 600.0
|
||||
|
||||
_SECONDS: Final = TypeAdapter(float)
|
||||
|
|
@ -48,21 +50,27 @@ def resolve_llm_passthrough_timeout(
|
|||
Anthropic /v1/messages).
|
||||
|
||||
Non-streaming precedence: kwargs timeout/request_timeout -> litellm_params
|
||||
timeout/request_timeout -> router_timeout -> general_settings.pass_through_request_timeout
|
||||
-> 600s default.
|
||||
timeout_by_service_tier[kwargs service_tier] -> litellm_params timeout/request_timeout -> router_timeout
|
||||
-> general_settings.pass_through_request_timeout -> 600s default.
|
||||
|
||||
Streaming (``kwargs["stream"]`` truthy) resolves ``stream_timeout`` at every level before
|
||||
any generic timeout, matching ``Router._get_stream_timeout`` on the completion route:
|
||||
kwargs stream_timeout -> litellm_params stream_timeout -> router_stream_timeout, then the
|
||||
non-streaming chain above.
|
||||
kwargs stream_timeout -> litellm_params stream_timeout_by_service_tier[kwargs service_tier] ->
|
||||
litellm_params stream_timeout -> router_stream_timeout, then the non-streaming chain above.
|
||||
|
||||
Only the first set value is validated as seconds, so a value in a lower-precedence
|
||||
field never fails the call.
|
||||
"""
|
||||
request: Final = kwargs if kwargs is not None else _NO_PARAMS
|
||||
deployment: Final = litellm_params if litellm_params is not None else _NO_PARAMS
|
||||
service_tier: Final = request.get("service_tier")
|
||||
stream_candidates: Final = (
|
||||
(request.get("stream_timeout"), deployment.get("stream_timeout"), router_stream_timeout)
|
||||
(
|
||||
request.get("stream_timeout"),
|
||||
get_service_tier_timeout(service_tier, deployment, "stream_timeout_by_service_tier"),
|
||||
deployment.get("stream_timeout"),
|
||||
router_stream_timeout,
|
||||
)
|
||||
if request.get("stream")
|
||||
else ()
|
||||
)
|
||||
|
|
@ -70,6 +78,7 @@ def resolve_llm_passthrough_timeout(
|
|||
*stream_candidates,
|
||||
request.get("timeout"),
|
||||
request.get("request_timeout"),
|
||||
get_service_tier_timeout(service_tier, deployment, "timeout_by_service_tier"),
|
||||
deployment.get("timeout"),
|
||||
deployment.get("request_timeout"),
|
||||
router_timeout,
|
||||
|
|
|
|||
|
|
@ -111,6 +111,7 @@ from litellm.proxy.guardrails.auto_router_compression import arm_pre_call as _ar
|
|||
from litellm.proxy.native_compaction import with_proxy_compaction_executor
|
||||
from litellm.proxy.route_llm_request import route_request
|
||||
from litellm.proxy.utils import ProxyLogging, _check_and_merge_model_level_guardrails
|
||||
from litellm.responses.streaming_iterator import partial_responses_incomplete_event_on_disconnect
|
||||
from litellm.router import Router
|
||||
from litellm.router_utils.add_retry_fallback_headers import get_hidden_params_dict
|
||||
from litellm.router_utils.common_utils import resolve_model_group_alias
|
||||
|
|
@ -418,6 +419,13 @@ async def _bill_partial_streamed_spend_on_disconnect(request_data: dict, respons
|
|||
# A natural end-of-stream success event already fired and released the
|
||||
# slot; do not bill again, and let the caller skip the slot release.
|
||||
return True
|
||||
partial_responses_event: Final = partial_responses_incomplete_event_on_disconnect(response)
|
||||
if partial_responses_event is not None:
|
||||
verbose_proxy_logger.debug(
|
||||
"Billing partial streamed Responses API spend after client disconnect, litellm_call_id=%s",
|
||||
logging_obj.litellm_call_id,
|
||||
)
|
||||
return await _dispatch_disconnect_billing_event(logging_obj, partial_responses_event)
|
||||
chunks: Final[object] = getattr(response, "chunks", None)
|
||||
if not isinstance(chunks, list) or not chunks:
|
||||
return False
|
||||
|
|
@ -447,6 +455,10 @@ async def _bill_partial_streamed_spend_on_disconnect(request_data: dict, respons
|
|||
partial_usage: Final = getattr(partial_response, "usage", None)
|
||||
if isinstance(partial_usage, Usage):
|
||||
backfill_missing_cache_usage_fields(partial_usage)
|
||||
return await _dispatch_disconnect_billing_event(logging_obj, partial_response)
|
||||
|
||||
|
||||
async def _dispatch_disconnect_billing_event(logging_obj: LiteLLMLoggingObj, partial_response: object) -> bool:
|
||||
try:
|
||||
await logging_obj.dispatch_success_handlers(
|
||||
partial_response,
|
||||
|
|
|
|||
|
|
@ -119,6 +119,23 @@ def coerce_numeric_form_fields(
|
|||
}
|
||||
|
||||
|
||||
def _parse_form_metadata(raw_metadata: str) -> object:
|
||||
"""The JSON body path rejects a lone surrogate escape (``"\\ud800"``) with a
|
||||
400 because no provider or database can encode it. ``json.loads`` accepts
|
||||
the escape, so the multipart ``metadata`` field gets the same check."""
|
||||
parsed: Final[object] = json.loads(raw_metadata)
|
||||
try:
|
||||
json.dumps(parsed, ensure_ascii=False).encode("utf-8")
|
||||
except UnicodeEncodeError as e:
|
||||
raise ProxyException(
|
||||
message=f"Invalid form payload: metadata {e}",
|
||||
type="invalid_request_error",
|
||||
param="metadata",
|
||||
code=status.HTTP_400_BAD_REQUEST,
|
||||
) from e
|
||||
return parsed
|
||||
|
||||
|
||||
async def _read_request_body(request: Request | None) -> dict:
|
||||
"""
|
||||
Safely read the request body and parse it as JSON.
|
||||
|
|
@ -158,9 +175,13 @@ async def _read_request_body(request: Request | None) -> dict:
|
|||
param="request_body",
|
||||
code=status.HTTP_400_BAD_REQUEST,
|
||||
)
|
||||
parsed_body = dict(form_data)
|
||||
if "metadata" in parsed_body and isinstance(parsed_body["metadata"], str):
|
||||
parsed_body["metadata"] = json.loads(parsed_body["metadata"])
|
||||
form_body: Final = dict(form_data)
|
||||
raw_metadata: Final = form_body.get("metadata")
|
||||
parsed_body = (
|
||||
{**form_body, "metadata": _parse_form_metadata(raw_metadata)}
|
||||
if isinstance(raw_metadata, str)
|
||||
else form_body
|
||||
)
|
||||
else:
|
||||
# Read the request body
|
||||
body: Final = await request.body()
|
||||
|
|
|
|||
|
|
@ -88,6 +88,7 @@ if TYPE_CHECKING:
|
|||
from litellm.proxy.db.autorouter_session_rollup import AutoRouterTurnTransaction
|
||||
from litellm.proxy.db.baseline_accounting import DailyBaselineAttribution
|
||||
from litellm.proxy.utils import PrismaClient, ProxyLogging
|
||||
from litellm.router import Router
|
||||
else:
|
||||
PrismaClient = Any
|
||||
ProxyLogging = Any
|
||||
|
|
@ -97,6 +98,34 @@ RESPONSES_SESSION_CALL_TYPES: Final = frozenset({CallTypes.responses.value, Call
|
|||
_SPEND_METADATA_ADAPTER: Final = TypeAdapter(Mapping[str, object])
|
||||
|
||||
|
||||
def _turn_chains_through_the_db(payload: Mapping[str, object], llm_router: "Router | None") -> bool:
|
||||
"""A Responses turn the chat-completions bridge replays from its spend log on the next `previous_response_id`.
|
||||
|
||||
Native Responses providers hold the conversation themselves, so their rows can wait for the batch write.
|
||||
A row with no router deployment behind it is treated as bridged, which keeps today's immediate flush.
|
||||
"""
|
||||
if payload.get("call_type") not in RESPONSES_SESSION_CALL_TYPES:
|
||||
return False
|
||||
model_id: Final = payload.get("model_id")
|
||||
deployment: Final = (
|
||||
llm_router.get_deployment(model_id=model_id) if llm_router is not None and isinstance(model_id, str) else None
|
||||
)
|
||||
if deployment is None:
|
||||
return True
|
||||
from litellm.responses.main import provider_stores_the_conversation
|
||||
|
||||
litellm_params: Final = deployment.litellm_params
|
||||
return not provider_stores_the_conversation(
|
||||
litellm_params.model,
|
||||
litellm_params.custom_llm_provider,
|
||||
{
|
||||
"use_chat_completions_api": litellm_params.use_chat_completions_api,
|
||||
"model_info": deployment.model_info.model_dump(),
|
||||
"api_base": litellm_params.api_base,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def _org_member_transaction_key(org_id: str, user_id: str) -> str:
|
||||
return f"organization_id::{quote(org_id, safe='')}::user_id::{quote(user_id, safe='')}"
|
||||
|
||||
|
|
@ -527,12 +556,13 @@ class DBSpendUpdateWriter:
|
|||
get_request_model_access_groups,
|
||||
)
|
||||
|
||||
llm_router: Final = get_llm_router()
|
||||
payload: Final = get_logging_payload(
|
||||
kwargs=kwargs,
|
||||
response_obj=completion_response,
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
llm_router=get_llm_router(),
|
||||
llm_router=llm_router,
|
||||
)
|
||||
payload["spend"] = response_cost or 0.0
|
||||
if isinstance(payload["startTime"], datetime):
|
||||
|
|
@ -547,7 +577,10 @@ class DBSpendUpdateWriter:
|
|||
payload["team_id"] = team_id
|
||||
|
||||
if not await self._record_spend_log(
|
||||
payload=payload, prisma_client=prisma_client, disable_spend_logs=disable_spend_logs
|
||||
payload=payload,
|
||||
prisma_client=prisma_client,
|
||||
disable_spend_logs=disable_spend_logs,
|
||||
llm_router=llm_router,
|
||||
):
|
||||
return False
|
||||
|
||||
|
|
@ -608,14 +641,18 @@ class DBSpendUpdateWriter:
|
|||
return True
|
||||
|
||||
async def _record_spend_log(
|
||||
self, payload: SpendLogsPayload, prisma_client: "PrismaClient | None", disable_spend_logs: bool
|
||||
self,
|
||||
payload: SpendLogsPayload,
|
||||
prisma_client: "PrismaClient | None",
|
||||
disable_spend_logs: bool,
|
||||
llm_router: "Router | None" = None,
|
||||
) -> bool:
|
||||
if prisma_client is not None and _is_batch_cost_row(payload):
|
||||
return await self._claim_batch_cost_spend_log(
|
||||
payload=payload, prisma_client=prisma_client, disable_spend_logs=disable_spend_logs
|
||||
)
|
||||
if disable_spend_logs is False:
|
||||
await self._insert_spend_log_to_db(payload=payload, prisma_client=prisma_client)
|
||||
await self._insert_spend_log_to_db(payload=payload, prisma_client=prisma_client, llm_router=llm_router)
|
||||
return True
|
||||
|
||||
async def _claim_batch_cost_spend_log(
|
||||
|
|
@ -1462,6 +1499,7 @@ class DBSpendUpdateWriter:
|
|||
payload: dict | SpendLogsPayload,
|
||||
prisma_client: PrismaClient | None = None,
|
||||
spend_logs_url: str | None = os.getenv("SPEND_LOGS_URL"),
|
||||
llm_router: "Router | None" = None,
|
||||
) -> PrismaClient | None:
|
||||
verbose_proxy_logger.debug(
|
||||
"Writing spend log to db - request_id: {}, spend: {}".format(
|
||||
|
|
@ -1472,7 +1510,7 @@ class DBSpendUpdateWriter:
|
|||
from litellm.proxy.utils import enqueue_spend_logs, request_spend_log_flush
|
||||
|
||||
await enqueue_spend_logs(prisma_client, (payload,))
|
||||
if payload.get("call_type") in RESPONSES_SESSION_CALL_TYPES:
|
||||
if _turn_chains_through_the_db(_SPEND_METADATA_ADAPTER.validate_python(payload), llm_router):
|
||||
request_spend_log_flush(prisma_client)
|
||||
else:
|
||||
verbose_proxy_logger.debug("prisma_client is None. Skipping writing spend logs to db.")
|
||||
|
|
|
|||
|
|
@ -1538,7 +1538,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail):
|
|||
|
||||
def _unmask_responses_api_completed_chunk(self, chunk: object, pii_tokens: dict[str, str]) -> None:
|
||||
"""
|
||||
Unmask PII tokens in-place for a ``response.completed`` Responses API event.
|
||||
Unmask PII tokens in-place for a ``response.completed`` / ``response.incomplete`` Responses API event.
|
||||
|
||||
The chunk carries a ``response`` attribute (ResponsesAPIResponse) whose
|
||||
``output`` list holds message items. Each item has a ``content`` list of
|
||||
|
|
@ -1598,7 +1598,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail):
|
|||
yield buffered_chunk
|
||||
remaining_chunks = []
|
||||
chunk_type = getattr(chunk, "type", None)
|
||||
if chunk_type == "response.completed" and pii_tokens:
|
||||
if chunk_type in ("response.completed", "response.incomplete") and pii_tokens:
|
||||
self._unmask_responses_api_completed_chunk(chunk, pii_tokens)
|
||||
saw_non_chat_chunk = True
|
||||
yield chunk
|
||||
|
|
|
|||
|
|
@ -6,8 +6,6 @@ from dataclasses import dataclass
|
|||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
|
||||
from openai.types import Batch
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.caching.caching import DualCache
|
||||
|
|
@ -16,7 +14,7 @@ from litellm.litellm_core_utils.duration_parser import duration_in_seconds
|
|||
from litellm.llms.bedrock.common_utils import get_bedrock_base_model
|
||||
from litellm.proxy._types import Litellm_EntityType, UserAPIKeyAuth
|
||||
from litellm.router_strategy.budget_limiter import RouterBudgetLimiting
|
||||
from litellm.router_utils.batch_utils import is_batch_retrieve_call_type
|
||||
from litellm.router_utils.batch_utils import batch_charged_once_marker_key, batch_id_to_charge_once
|
||||
from litellm.types.llms.openai import AllMessageValues
|
||||
from litellm.types.utils import BudgetConfig, StandardLoggingPayload
|
||||
|
||||
|
|
@ -121,17 +119,6 @@ def model_budget_start_time_cache_key(
|
|||
return f"{_BUDGET_START_TIME_KEY_PREFIXES[entity_type]}:{entity_id}:{budget_model}:{budget_duration}"
|
||||
|
||||
|
||||
def batch_charged_once_marker_key(spend_key: str, batch_id: str) -> str:
|
||||
return f"{spend_key}:batch:{batch_id}"
|
||||
|
||||
|
||||
def batch_id_to_charge_once(call_type: object, response_obj: object, response_cost: float) -> str | None:
|
||||
"""A finished batch reports its whole cost on every poll, so its id is charged once per counter."""
|
||||
if response_cost <= 0 or not is_batch_retrieve_call_type(call_type):
|
||||
return None
|
||||
return response_obj.id if isinstance(response_obj, Batch) else None
|
||||
|
||||
|
||||
def resolve_model_budget(model: str, model_max_budget: Mapping[str, object]) -> ResolvedModelBudget | None:
|
||||
"""Find the `model_max_budget` entry that governs `model`, or None."""
|
||||
for candidate in _budget_model_candidates(model):
|
||||
|
|
|
|||
|
|
@ -1,9 +1,14 @@
|
|||
import asyncio
|
||||
import traceback
|
||||
from collections.abc import Callable, Mapping, Sequence
|
||||
from datetime import datetime
|
||||
from collections.abc import Callable, Generator, Mapping, Sequence
|
||||
from contextlib import contextmanager
|
||||
from contextvars import ContextVar
|
||||
from datetime import datetime, timezone
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Any, Final, Protocol, cast
|
||||
|
||||
from pydantic import TypeAdapter, ValidationError
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.batches.batch_utils import batch_cost_is_final
|
||||
|
|
@ -29,6 +34,7 @@ from litellm.proxy.db.db_spend_update_writer import (
|
|||
debitable_model_access_groups,
|
||||
get_llm_router,
|
||||
)
|
||||
from litellm.proxy.hooks.search_free_tier import apply_search_free_tier
|
||||
from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup
|
||||
from litellm.proxy.spend_tracking.spend_event import (
|
||||
ObjectMapping,
|
||||
|
|
@ -49,7 +55,7 @@ from litellm.proxy.spend_tracking.spend_tracking_utils import (
|
|||
get_request_model_access_groups,
|
||||
should_store_prompts_and_responses_in_spend_logs,
|
||||
)
|
||||
from litellm.proxy.utils import ProxyUpdateSpend
|
||||
from litellm.proxy.utils import BilledFailedAttempt, ProxyUpdateSpend
|
||||
from litellm.types.utils import (
|
||||
CallTypes,
|
||||
LiteLLMBatch,
|
||||
|
|
@ -61,6 +67,37 @@ from litellm.utils import get_end_user_id_for_cost_tracking
|
|||
if TYPE_CHECKING:
|
||||
from litellm.proxy.utils import ProxyLogging
|
||||
|
||||
|
||||
class BatchAccountingOutcome:
|
||||
def __init__(self, call_id: str) -> None:
|
||||
self.call_id: Final = call_id
|
||||
self.write_started = False
|
||||
self.failed = False
|
||||
|
||||
@property
|
||||
def retryable_failure(self) -> bool:
|
||||
return self.failed and not self.write_started
|
||||
|
||||
|
||||
_BATCH_ACCOUNTING: Final[ContextVar[BatchAccountingOutcome | None]] = ContextVar("batch_accounting", default=None)
|
||||
_BATCH_CALLBACK_FIELDS: Final = TypeAdapter(dict[str, object])
|
||||
|
||||
|
||||
@contextmanager
|
||||
def track_batch_accounting(call_id: str) -> Generator[BatchAccountingOutcome, None, None]:
|
||||
outcome: Final = BatchAccountingOutcome(call_id)
|
||||
token: Final = _BATCH_ACCOUNTING.set(outcome)
|
||||
try:
|
||||
yield outcome
|
||||
finally:
|
||||
_BATCH_ACCOUNTING.reset(token)
|
||||
|
||||
|
||||
def _batch_accounting_for(call_id: object) -> BatchAccountingOutcome | None:
|
||||
outcome: Final = _BATCH_ACCOUNTING.get()
|
||||
return outcome if outcome is not None and call_id == outcome.call_id else None
|
||||
|
||||
|
||||
_UNATTRIBUTED_TRACKABLE_CALL_TYPES: Final[frozenset[str]] = frozenset(
|
||||
{
|
||||
CallTypes.pass_through.value,
|
||||
|
|
@ -91,6 +128,79 @@ def _proxy_spend_writer() -> DBSpendUpdateWriter:
|
|||
return proxy_logging_obj.db_spend_update_writer
|
||||
|
||||
|
||||
def _billed_failed_attempts(request_data: Mapping[str, object]) -> tuple[BilledFailedAttempt, ...]:
|
||||
"""The earlier billed attempts ``_failure_fields_to_lift`` stashed on the request, if any."""
|
||||
attempts: Final = request_data.get("billed_failed_attempts")
|
||||
if not isinstance(attempts, tuple):
|
||||
return ()
|
||||
lifted: Final = cast("tuple[object, ...]", attempts) # cast-ok: lifted by _failure_fields_to_lift, checked below
|
||||
return tuple(attempt for attempt in lifted if isinstance(attempt, BilledFailedAttempt))
|
||||
|
||||
|
||||
def _failure_spend_rows(
|
||||
request_data: Mapping[str, object], final_cost: float, late_guardrails_snapshot: object
|
||||
) -> tuple[tuple[float, Mapping[str, object]], ...]:
|
||||
"""One ``(response_cost, kwargs)`` per spend row a failure writes: the final attempt under the
|
||||
request id, then every earlier attempt the provider also billed under its own row."""
|
||||
final_row: Final[Mapping[str, object]] = (
|
||||
request_data
|
||||
if late_guardrails_snapshot is None
|
||||
else {**request_data, "standard_logging_object": late_guardrails_snapshot}
|
||||
)
|
||||
return (
|
||||
(final_cost, final_row),
|
||||
*(
|
||||
(attempt.response_cost, _attempt_spend_row_kwargs(request_data, attempt, position))
|
||||
for position, attempt in enumerate(_billed_failed_attempts(request_data))
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def _attempt_spend_row_kwargs(
|
||||
request_data: Mapping[str, object], attempt: BilledFailedAttempt, position: int
|
||||
) -> Mapping[str, object]:
|
||||
"""The final attempt's row kwargs re-keyed for one earlier billed attempt: its own
|
||||
``request_id`` (``<id>-attempt-<position>``, since ``request_id`` is the spend-log primary key
|
||||
and the retry count restarts after a fallback hop), its estimated usage and cost, its retry
|
||||
count under ``attempted_retries`` and the deployment that billed it under ``model_info.id``."""
|
||||
suffix: Final = f"-attempt-{position}"
|
||||
raw_params: Final = cast("Mapping[str, object]", request_data.get("litellm_params") or {}) # cast-ok: request dict
|
||||
litellm_params: Final = dict(raw_params)
|
||||
metadata: Final = {
|
||||
**cast("Mapping[str, object]", litellm_params.get("metadata") or {}), # cast-ok: request dict
|
||||
"attempted_retries": attempt.attempted_retries,
|
||||
}
|
||||
metadata.pop("previous_models", None)
|
||||
if attempt.deployment_id is not None:
|
||||
metadata["model_info"] = {
|
||||
**cast("Mapping[str, object]", metadata.get("model_info") or {}), # cast-ok: router dict
|
||||
"id": attempt.deployment_id,
|
||||
}
|
||||
litellm_params["metadata"] = metadata
|
||||
row: Final = {
|
||||
**request_data,
|
||||
"litellm_params": litellm_params,
|
||||
"combined_usage_object": attempt.usage,
|
||||
"response_cost": attempt.response_cost,
|
||||
}
|
||||
row.pop("billed_failed_attempts", None)
|
||||
call_id: Final = request_data.get("litellm_call_id") or litellm_params.get("litellm_call_id")
|
||||
if isinstance(call_id, str):
|
||||
row["litellm_call_id"] = f"{call_id}{suffix}"
|
||||
standard_logging_object: Final = request_data.get("standard_logging_object")
|
||||
if isinstance(standard_logging_object, dict):
|
||||
payload: Final = cast("Mapping[str, object]", standard_logging_object) # cast-ok: StandardLoggingPayload dict
|
||||
if isinstance(payload.get("id"), str):
|
||||
row["standard_logging_object"] = {**payload, "id": f"{payload['id']}{suffix}"}
|
||||
return row
|
||||
|
||||
|
||||
def _str_keyed_fields(value: object) -> Mapping[str, object]:
|
||||
if not isinstance(value, Mapping):
|
||||
return {}
|
||||
return cast("Mapping[str, object]", value) # cast-ok: payload dicts are str-keyed
|
||||
|
||||
|
||||
class _ProxyDBLogger(CustomLogger):
|
||||
def __init__(
|
||||
self,
|
||||
|
|
@ -258,20 +368,40 @@ class _ProxyDBLogger(CustomLogger):
|
|||
existing_metadata.get("standard_logging_guardrail_information")
|
||||
)
|
||||
|
||||
await self._spend_writer().update_database(
|
||||
token=LiteLLMProxyRequestSetup.get_logged_api_key(user_api_key_dict),
|
||||
response_cost=recovered_response_cost,
|
||||
user_id=user_api_key_dict.user_id,
|
||||
end_user_id=user_api_key_dict.end_user_id,
|
||||
team_id=user_api_key_dict.team_id,
|
||||
kwargs=request_data,
|
||||
completion_response=original_exception,
|
||||
await self._write_failure_spend_rows(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
original_exception=original_exception,
|
||||
start_time=actual_start_time,
|
||||
end_time=datetime.now(),
|
||||
org_id=user_api_key_dict.org_id,
|
||||
project_id=user_api_key_dict.project_id,
|
||||
spend_rows=_failure_spend_rows(
|
||||
cast("Mapping[str, object]", request_data), # cast-ok: request dict, copied per row below
|
||||
recovered_response_cost,
|
||||
None,
|
||||
),
|
||||
)
|
||||
|
||||
async def _write_failure_spend_rows(
|
||||
self,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
original_exception: Exception,
|
||||
start_time: datetime,
|
||||
spend_rows: Sequence[tuple[float, Mapping[str, object]]],
|
||||
) -> None:
|
||||
"""One spend row (and one key/team/user spend deduction) per ``(response_cost, kwargs)`` pair."""
|
||||
for row_cost, row_kwargs in spend_rows:
|
||||
await self._spend_writer().update_database(
|
||||
token=LiteLLMProxyRequestSetup.get_logged_api_key(user_api_key_dict),
|
||||
response_cost=row_cost,
|
||||
user_id=user_api_key_dict.user_id,
|
||||
end_user_id=user_api_key_dict.end_user_id,
|
||||
team_id=user_api_key_dict.team_id,
|
||||
kwargs=dict(row_kwargs),
|
||||
completion_response=original_exception,
|
||||
start_time=start_time,
|
||||
end_time=datetime.now(),
|
||||
org_id=user_api_key_dict.org_id,
|
||||
project_id=user_api_key_dict.project_id,
|
||||
)
|
||||
|
||||
@log_db_metrics
|
||||
async def _PROXY_track_cost_callback(
|
||||
self,
|
||||
|
|
@ -286,6 +416,11 @@ class _ProxyDBLogger(CustomLogger):
|
|||
update_cache,
|
||||
)
|
||||
|
||||
accounting: Final = (
|
||||
_batch_accounting_for(_BATCH_CALLBACK_FIELDS.validate_python(kwargs).get("litellm_call_id"))
|
||||
if _BATCH_ACCOUNTING.get() is not None
|
||||
else None
|
||||
)
|
||||
verbose_proxy_logger.debug("INSIDE _PROXY_track_cost_callback")
|
||||
try:
|
||||
verbose_proxy_logger.debug(
|
||||
|
|
@ -306,6 +441,8 @@ class _ProxyDBLogger(CustomLogger):
|
|||
)
|
||||
_write_spend_metadata_to_kwargs(kwargs=kwargs, metadata=metadata)
|
||||
budget_reservation: Final = _get_budget_reservation_from_metadata(metadata=metadata)
|
||||
if accounting is not None and budget_reservation is not None:
|
||||
accounting.write_started = True
|
||||
if (
|
||||
isinstance(completion_response, LiteLLMBatch)
|
||||
and kwargs.get("call_type") == CallTypes.aretrieve_batch.value
|
||||
|
|
@ -329,8 +466,21 @@ class _ProxyDBLogger(CustomLogger):
|
|||
key_alias: Final = cast(str | None, metadata.get("user_api_key_alias", None))
|
||||
end_user_max_budget: Final = metadata.get("user_api_end_user_max_budget", None)
|
||||
sl_object: Final[StandardLoggingPayload | None] = kwargs.get("standard_logging_object", None)
|
||||
raw_cost: Final = _str_keyed_fields(
|
||||
sl_object if sl_object is not None else kwargs # pyright: ignore[reportUnknownArgumentType] # untyped success-callback payload
|
||||
).get("response_cost")
|
||||
priced_cost: Final = raw_cost if isinstance(raw_cost, (int, float)) else None
|
||||
response_cost: Final = (
|
||||
sl_object.get("response_cost", None) if sl_object is not None else kwargs.get("response_cost", None)
|
||||
await apply_search_free_tier(
|
||||
kwargs=kwargs, # pyright: ignore[reportUnknownArgumentType] # untyped success-callback payload
|
||||
sl_object=sl_object, # pyright: ignore[reportUnknownArgumentType] # payload carries untyped lists
|
||||
response_cost=priced_cost,
|
||||
completion_response=completion_response,
|
||||
counter=proxy_logging_obj.internal_usage_cache.dual_cache.redis_cache,
|
||||
now=datetime.now(timezone.utc),
|
||||
)
|
||||
if litellm.apply_search_free_tier and sl_object is not None and priced_cost is not None
|
||||
else priced_cost
|
||||
)
|
||||
tags: Final = _get_request_tags_for_cost_tracking(
|
||||
sl_object=sl_object,
|
||||
|
|
@ -448,6 +598,8 @@ class _ProxyDBLogger(CustomLogger):
|
|||
f"Cost tracking failed for model={model}.\nDebug info - {cost_tracking_failure_debug_info}\nAdd custom pricing - https://docs.litellm.ai/docs/proxy/custom_pricing"
|
||||
)
|
||||
except Exception as e:
|
||||
if accounting is not None:
|
||||
accounting.failed = True
|
||||
failing_model: Final = kwargs.get("model", "")
|
||||
failing_call_type: Final = kwargs.get("call_type", "")
|
||||
error_msg: Final = (
|
||||
|
|
@ -698,7 +850,15 @@ async def _update_database_and_spend_counters(
|
|||
budget_reservation=budget_reservation, response_cost=response_cost
|
||||
)
|
||||
try:
|
||||
charged: Final = await proxy_logging_obj.db_spend_update_writer.update_database(
|
||||
writer: Final = proxy_logging_obj.db_spend_update_writer
|
||||
accounting: Final = (
|
||||
_batch_accounting_for(_BATCH_CALLBACK_FIELDS.validate_python(kwargs).get("litellm_call_id"))
|
||||
if _BATCH_ACCOUNTING.get() is not None
|
||||
else None
|
||||
)
|
||||
if accounting is not None:
|
||||
accounting.write_started = True
|
||||
charged: Final = await writer.update_database(
|
||||
token=user_api_key,
|
||||
response_cost=response_cost,
|
||||
user_id=user_id,
|
||||
|
|
|
|||
123
litellm/proxy/hooks/search_free_tier.py
Normal file
123
litellm/proxy/hooks/search_free_tier.py
Normal file
|
|
@ -0,0 +1,123 @@
|
|||
from collections.abc import MutableMapping
|
||||
from datetime import datetime
|
||||
from types import MappingProxyType
|
||||
from typing import Final, Protocol
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.litellm_core_utils.duration_parser import get_next_standardized_reset_time
|
||||
from litellm.litellm_core_utils.llm_cost_calc.utils import get_web_search_requests_from_usage
|
||||
from litellm.llms.gemini.cost_calculator import cost_per_web_search_request
|
||||
from litellm.types.utils import CostBreakdown, ModelInfo, SearchFreeRequests, StandardLoggingPayload, Usage
|
||||
|
||||
_PERIOD_DURATION: Final = MappingProxyType({"day": "1d", "month": "1mo"})
|
||||
_PERIOD_STAMP: Final = MappingProxyType({"day": "%Y-%m-%d", "month": "%Y-%m"})
|
||||
_TTL_GRACE_SECONDS: Final = 3600
|
||||
_COUNTED_KWARG: Final = "search_free_tier_counted"
|
||||
_COST_BREAKDOWN: Final = TypeAdapter(CostBreakdown)
|
||||
|
||||
|
||||
class SharedCounter(Protocol):
|
||||
async def async_increment(self, key: str, value: float, ttl: int | None = None) -> float: ...
|
||||
|
||||
|
||||
class _UsageCarrier(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="ignore", from_attributes=True)
|
||||
|
||||
usage: Usage | None = None
|
||||
|
||||
|
||||
def freed_units(allowance: int, units: int, total_after: int) -> int:
|
||||
"""Units of this request that landed under the allowance, given the shared total after they were added."""
|
||||
return max(0, min(units, allowance - (total_after - units)))
|
||||
|
||||
|
||||
def search_free_tier_key(provider: str, allowance: SearchFreeRequests, now: datetime) -> str:
|
||||
period_stamp: Final = now.strftime(_PERIOD_STAMP[allowance["period"]])
|
||||
return f"litellm:search_free_tier:{provider}:{allowance['shared_group']}:{period_stamp}"
|
||||
|
||||
|
||||
def search_free_tier_ttl(allowance: SearchFreeRequests, now: datetime) -> int:
|
||||
reset: Final = get_next_standardized_reset_time(_PERIOD_DURATION[allowance["period"]], now, "UTC")
|
||||
return int((reset - now).total_seconds()) + _TTL_GRACE_SECONDS
|
||||
|
||||
|
||||
def _usage_of(completion_response: object) -> Usage | None:
|
||||
try:
|
||||
return _UsageCarrier.model_validate(completion_response, from_attributes=True).usage
|
||||
except ValidationError:
|
||||
return None
|
||||
|
||||
|
||||
def _grounding_queries(usage: Usage) -> int:
|
||||
details: Final = usage.prompt_tokens_details
|
||||
from_details: Final = getattr(details, "web_search_requests", None) if details is not None else None
|
||||
return from_details or get_web_search_requests_from_usage(usage) or 0
|
||||
|
||||
|
||||
def _billable_units(usage: Usage, model_info: ModelInfo) -> int:
|
||||
queries: Final = _grounding_queries(usage)
|
||||
if queries == 0:
|
||||
return 0
|
||||
return queries if model_info.get("web_search_billing_unit") == "per_query" else 1
|
||||
|
||||
|
||||
def _record_freed_cost(
|
||||
kwargs: MutableMapping[str, object], # mutable-ok: writes the reduced cost where the spend row reads it
|
||||
sl_object: StandardLoggingPayload,
|
||||
original_cost: float,
|
||||
freed_cost: float,
|
||||
) -> float:
|
||||
reduced_cost: Final = original_cost - freed_cost
|
||||
breakdown: Final = sl_object["cost_breakdown"] or CostBreakdown(total_cost=original_cost)
|
||||
sl_object["cost_breakdown"] = _COST_BREAKDOWN.validate_python( # rebind-ok: the spend row reads this payload
|
||||
{
|
||||
**breakdown,
|
||||
"total_cost": reduced_cost,
|
||||
"original_cost": breakdown.get("original_cost", original_cost),
|
||||
"discount_amount": breakdown.get("discount_amount", 0.0) + freed_cost,
|
||||
}
|
||||
)
|
||||
sl_object["response_cost"] = reduced_cost # rebind-ok: the spend row reads this payload
|
||||
kwargs["response_cost"] = reduced_cost # rebind-ok: the spend row reads kwargs["response_cost"]
|
||||
return reduced_cost
|
||||
|
||||
|
||||
async def apply_search_free_tier(
|
||||
kwargs: MutableMapping[str, object], # mutable-ok: writes the reduced cost where the spend row reads it
|
||||
sl_object: StandardLoggingPayload,
|
||||
response_cost: float,
|
||||
completion_response: object,
|
||||
counter: SharedCounter | None,
|
||||
now: datetime,
|
||||
) -> float:
|
||||
"""Zero the grounding queries that landed under the model's shared free allowance, everything else is unchanged."""
|
||||
model_info: Final = sl_object["model_map_information"]["model_map_value"]
|
||||
allowance: Final = model_info.get("search_free_requests") if model_info is not None else None
|
||||
usage: Final = _usage_of(completion_response)
|
||||
if model_info is None or allowance is None or usage is None or sl_object["cache_hit"]:
|
||||
return response_cost
|
||||
units: Final = _billable_units(usage, model_info)
|
||||
if units == 0 or kwargs.get(_COUNTED_KWARG):
|
||||
return response_cost
|
||||
if counter is None:
|
||||
verbose_proxy_logger.warning(
|
||||
"apply_search_free_tier is on but the proxy has no Redis, grounding is billed at full price"
|
||||
)
|
||||
return response_cost
|
||||
provider: Final = sl_object["custom_llm_provider"] or model_info["litellm_provider"]
|
||||
key: Final = search_free_tier_key(provider, allowance, now)
|
||||
try:
|
||||
total_after: Final = await counter.async_increment(key, units, ttl=search_free_tier_ttl(allowance, now))
|
||||
except Exception: # noqa: BLE001 # any counter failure means full price, never a lost spend row
|
||||
verbose_proxy_logger.warning(
|
||||
"search free tier counter %s is unavailable, grounding is billed at full price", key
|
||||
)
|
||||
return response_cost
|
||||
kwargs[_COUNTED_KWARG] = True # rebind-ok: the proxy runs this callback twice per request, count it once
|
||||
freed: Final = freed_units(allowance["count"], units, int(total_after))
|
||||
if freed == 0:
|
||||
return response_cost
|
||||
freed_cost: Final = min(response_cost, cost_per_web_search_request(usage, model_info) * freed / units)
|
||||
return _record_freed_cost(kwargs, sl_object, response_cost, freed_cost)
|
||||
|
|
@ -358,8 +358,9 @@ _ALLOW_CLIENT_MESSAGE_REDACTION_OPT_OUT_METADATA_KEY: Final = "allow_client_mess
|
|||
# ``litellm.model_cost`` map. Both effects belong to deployment configuration,
|
||||
# not to user-supplied request bodies, so the proxy strips them before they
|
||||
# reach the call path. Built from the Pydantic model so newly-added pricing
|
||||
# fields are covered automatically.
|
||||
_CLIENT_PRICING_CONTROL_FIELDS: Final = frozenset(CustomPricingLiteLLMParams.model_fields.keys())
|
||||
# fields are covered automatically. ``base_model`` selects which cost-map entry
|
||||
# prices the row, so a caller naming a cheaper model lowers their own spend.
|
||||
_CLIENT_PRICING_CONTROL_FIELDS: Final = frozenset(CustomPricingLiteLLMParams.model_fields.keys()) | {"base_model"}
|
||||
# ``model_info`` carries the same pricing fields when read by
|
||||
# ``use_custom_pricing_for_model``; strip from metadata for the same reason.
|
||||
# ``standard_logging_guardrail_information`` is proxy-written telemetry summed
|
||||
|
|
@ -373,6 +374,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,
|
||||
|
|
|
|||
|
|
@ -46,11 +46,13 @@ from typing import (
|
|||
overload,
|
||||
)
|
||||
|
||||
import httpx
|
||||
from typing_extensions import ReadOnly, TypedDict
|
||||
|
||||
from litellm import _custom_logger_compatible_callbacks_literal
|
||||
from litellm.constants import (
|
||||
DEFAULT_MODEL_CREATED_AT_TIME,
|
||||
LITELLM_HTTP_STATUS_CLIENT_DISCONNECTED,
|
||||
LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL,
|
||||
MAX_TEAM_LIST_LIMIT,
|
||||
PROXY_REJECTED_BEFORE_ROUTING_KEY,
|
||||
|
|
@ -143,7 +145,8 @@ from litellm.litellm_core_utils.core_helpers import (
|
|||
independent_snapshot,
|
||||
is_expected_client_error,
|
||||
)
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging, deployment_pricing_model_info
|
||||
from litellm.litellm_core_utils.llm_cost_calc.utils import generic_cost_per_token
|
||||
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
|
||||
from litellm.litellm_core_utils.safe_json_loads import safe_json_loads
|
||||
from litellm.litellm_core_utils.served_output_texts import (
|
||||
|
|
@ -236,6 +239,7 @@ from litellm.repositories.verification_token_repository import (
|
|||
VerificationTokenRepository,
|
||||
)
|
||||
from litellm.router_utils.common_utils import resolve_model_group_alias
|
||||
from litellm.router_utils.deployment_pricing import pricing_models
|
||||
from litellm.secret_managers.main import str_to_bool
|
||||
from litellm.types.integrations.slack_alerting import DEFAULT_ALERT_TYPES
|
||||
from litellm.types.llms.openai import ResponsesAPIResponse
|
||||
|
|
@ -248,6 +252,7 @@ from litellm.types.proxy.policy_engine.pipeline_types import PipelineExecutionRe
|
|||
from litellm.types.utils import LLMResponseTypes, LoggedLiteLLMParams
|
||||
from litellm.utils import (
|
||||
_add_custom_logger_callback_to_specific_event, # pyright: ignore[reportPrivateUsage] # only string-to-logger helper
|
||||
_get_base_model_from_metadata, # pyright: ignore[reportPrivateUsage] # main exposes it as get_base_model_from_metadata
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
|
|
@ -939,20 +944,48 @@ def _failure_usage_to_lift(
|
|||
model_call_details: Mapping[str, object],
|
||||
request_body: Mapping[str, object],
|
||||
dispatched: bool,
|
||||
provider_billed: bool = False,
|
||||
deployment_id: str | None = None,
|
||||
base_model: str | None = None,
|
||||
) -> tuple[object, object] | None:
|
||||
"""A stream that broke mid-flight still billed the provider for the chunks
|
||||
already delivered; the streaming handler stashes that recovered usage and
|
||||
cost in model_call_details, so prefer it. Otherwise a request that was
|
||||
dispatched to a provider and failed without upstream usage gets an
|
||||
estimated input-side Usage with 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
|
||||
estimated input-side Usage. When the failure is one the provider bills
|
||||
anyway (``provider_billed``: a read timeout or a client disconnect, where
|
||||
the provider had the prompt and kept working after LiteLLM hung up) that
|
||||
usage is priced at the model's input rate, the way the success path prices it:
|
||||
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. Returns the
|
||||
(combined_usage_object, response_cost) pair to lift, or None."""
|
||||
recovered_usage: Final = model_call_details.get("combined_usage_object")
|
||||
if recovered_usage is not None:
|
||||
return recovered_usage, model_call_details.get("response_cost")
|
||||
if not dispatched or model_call_details.get(LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL):
|
||||
if not dispatched:
|
||||
return None
|
||||
if model_call_details.get(LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL) and not provider_billed:
|
||||
# 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")
|
||||
|
|
@ -962,14 +995,205 @@ 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:
|
||||
|
||||
|
||||
@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 _record_pricing(
|
||||
model_call_details: Mapping[str, object], record: Mapping[str, object], base_model: str | None
|
||||
) -> tuple[Mapping[str, object], str | None]:
|
||||
"""The ``(model_call_details, base_model)`` to price one earlier attempt with: the record's own
|
||||
deployment model and ``base_model`` when it served a different model than the final attempt
|
||||
(the provider is then inferred from that model), else the final attempt's."""
|
||||
record_model: Final = record.get("model")
|
||||
if not isinstance(record_model, str) or record_model == model_call_details.get("model"):
|
||||
return model_call_details, base_model
|
||||
record_base_model: Final = record.get("base_model")
|
||||
return (
|
||||
{**model_call_details, "model": record_model, "custom_llm_provider": None},
|
||||
record_base_model if isinstance(record_base_model, str) else None,
|
||||
)
|
||||
|
||||
|
||||
def _billed_earlier_attempt(
|
||||
model_call_details: Mapping[str, object],
|
||||
record: Mapping[str, object],
|
||||
index: int,
|
||||
estimated_usage: Usage,
|
||||
base_model: str | None,
|
||||
) -> BilledFailedAttempt:
|
||||
pricing_details, pricing_base_model = _record_pricing(model_call_details, record, base_model)
|
||||
return BilledFailedAttempt(
|
||||
attempted_retries=_attempt_number(record, index),
|
||||
deployment_id=_record_deployment_id(record),
|
||||
usage=estimated_usage,
|
||||
response_cost=_estimated_input_cost(
|
||||
pricing_details,
|
||||
estimated_usage.prompt_tokens,
|
||||
deployment_id=_record_deployment_id(record),
|
||||
base_model=pricing_base_model,
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
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 the record's own
|
||||
``base_model`` or model, which a fallback hop can make a different model than the final one).
|
||||
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(
|
||||
_billed_earlier_attempt(model_call_details, record, index, estimated_usage, base_model)
|
||||
for index, record in enumerate(_timed_out_earlier_records(litellm_params, final_identity))
|
||||
)
|
||||
|
||||
|
||||
def _deployment_priced_input_cost(
|
||||
model: str, custom_llm_provider: str | None, prompt_tokens: int, deployment_pricing: ModelInfo
|
||||
) -> float:
|
||||
"""Input price for the estimated prompt at the deployment's registered rates, read through
|
||||
``deployment_pricing_model_info`` and priced by ``generic_cost_per_token`` the way the success
|
||||
path prices them, so a long-context tier the deployment declares (``input_cost_per_token_above_200k_tokens``)
|
||||
is honoured instead of flattened to the standard rate. An entry it cannot price keeps the row at $0."""
|
||||
usage: Final = Usage(prompt_tokens=prompt_tokens, completion_tokens=0, total_tokens=prompt_tokens)
|
||||
try:
|
||||
prompt_cost, _ = generic_cost_per_token(
|
||||
model=model,
|
||||
usage=usage,
|
||||
custom_llm_provider=custom_llm_provider or str(deployment_pricing.get("litellm_provider") or ""),
|
||||
model_info=deployment_pricing,
|
||||
)
|
||||
except Exception as exc: # noqa: BLE001 # a malformed registered entry must not drop the failure row
|
||||
verbose_proxy_logger.debug("failure spend estimate: could not price deployment of %s: %s", model, exc)
|
||||
return 0.0
|
||||
return prompt_cost
|
||||
|
||||
|
||||
def _priced_input_cost(pricing_model: str, custom_llm_provider: str | None, prompt_tokens: int) -> float | None:
|
||||
try:
|
||||
estimated_input_cost, _ = litellm.cost_per_token(
|
||||
model=pricing_model,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
prompt_tokens=prompt_tokens,
|
||||
)
|
||||
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_usage, 0.0
|
||||
return estimated_input_cost
|
||||
|
||||
|
||||
def _estimated_input_cost(
|
||||
model_call_details: Mapping[str, object],
|
||||
prompt_tokens: int,
|
||||
deployment_id: str | None,
|
||||
base_model: str | None,
|
||||
) -> float:
|
||||
"""Input price for the estimated prompt: the deployment's registered rates when it has any (the id
|
||||
is ``Logging.get_router_model_id()``, never the flattened ``model_call_details`` copy a request body
|
||||
key can overwrite), else 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
|
||||
deployment_pricing: Final = deployment_pricing_model_info(model_id=deployment_id, deployment_model=model)
|
||||
if deployment_pricing is not None:
|
||||
return _deployment_priced_input_cost(model, custom_llm_provider, prompt_tokens, deployment_pricing)
|
||||
pricing_models_to_try: Final = pricing_models(base_model, model)
|
||||
candidate_costs: Final = (
|
||||
_priced_input_cost(pricing_model, custom_llm_provider, prompt_tokens) for pricing_model in pricing_models_to_try
|
||||
)
|
||||
return next((cost for cost in candidate_costs if cost is not None), 0.0)
|
||||
|
||||
|
||||
_EMPTY_LIFT: Final = MappingProxyType({})
|
||||
|
|
@ -1085,7 +1309,38 @@ def _call_type_for_route(route: str | None) -> str | None:
|
|||
_PROXY_ONLY_LLM_API_ERRORS: Final = (HTTPException, ProxyException, GuardrailRaisedException)
|
||||
|
||||
|
||||
def _failure_fields_to_lift(request_data: Mapping[str, object]) -> Mapping[str, object]:
|
||||
_EXCEPTION_CHAIN_MAX_DEPTH: Final = 16
|
||||
|
||||
|
||||
def _exception_chain(exception: BaseException) -> tuple[BaseException, ...]:
|
||||
chain: Final[list[BaseException]] = [exception] # mutable-ok: built once, returned frozen
|
||||
for _ in range(_EXCEPTION_CHAIN_MAX_DEPTH):
|
||||
following = chain[-1].__cause__ or chain[-1].__context__
|
||||
if following is None or following in chain:
|
||||
break
|
||||
chain.append(following)
|
||||
return tuple(chain)
|
||||
|
||||
|
||||
def _timed_out_reading_from_provider(exception: Exception) -> bool:
|
||||
"""A read timeout means the request was sent and the provider was working on it; a connect
|
||||
timeout (``httpx.ConnectTimeout``) never reached the provider, so it is not billed."""
|
||||
return isinstance(exception, litellm.Timeout) and any(
|
||||
isinstance(cause, httpx.ReadTimeout) for cause in _exception_chain(exception)
|
||||
)
|
||||
|
||||
|
||||
def _client_disconnected(exception: Exception) -> bool:
|
||||
"""The 499 the proxy itself raises when the caller hung up while the provider was still working on the
|
||||
request. A provider's own 499 arrives as a litellm exception and is a plain upstream failure."""
|
||||
return isinstance(exception, HTTPException) and exception.status_code == LITELLM_HTTP_STATUS_CLIENT_DISCONNECTED
|
||||
|
||||
|
||||
def _provider_billed_the_failure(exception: Exception) -> bool:
|
||||
return _timed_out_reading_from_provider(exception) or _client_disconnected(exception)
|
||||
|
||||
|
||||
def _failure_fields_to_lift(request_data: Mapping[str, object], provider_billed: bool = False) -> Mapping[str, object]:
|
||||
"""Failure-path callbacks run after ``litellm_logging_obj`` is popped from
|
||||
request_data (it is not serialisable), so the caller merges these fields
|
||||
onto request_data first: the request start and first-handoff instants for
|
||||
|
|
@ -1097,10 +1352,28 @@ def _failure_fields_to_lift(request_data: Mapping[str, object]) -> Mapping[str,
|
|||
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
|
||||
_success_already_logged: Final = isinstance(_logging_obj, Logging) and _logging_obj.async_success_log_enqueued
|
||||
_billed: Final = provider_billed and not _success_already_logged
|
||||
_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=_billed,
|
||||
deployment_id=_deployment_id,
|
||||
base_model=_base_model,
|
||||
)
|
||||
_prompt_estimate: Final = (
|
||||
_estimated_prompt_usage(_model_call_details, request_data)
|
||||
if _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")),
|
||||
|
|
@ -1109,6 +1382,7 @@ def _failure_fields_to_lift(request_data: Mapping[str, object]) -> Mapping[str,
|
|||
("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})
|
||||
|
||||
|
|
@ -3252,7 +3526,11 @@ class ProxyLogging:
|
|||
original_exception=original_exception,
|
||||
)
|
||||
|
||||
request_data.update(await offload_token_count(_failure_fields_to_lift)(request_data))
|
||||
request_data.update(
|
||||
await offload_token_count(_failure_fields_to_lift)(
|
||||
request_data, provider_billed=_provider_billed_the_failure(original_exception)
|
||||
)
|
||||
)
|
||||
|
||||
# Remove before callbacks iterate — not serialisable
|
||||
request_data.pop("litellm_logging_obj", None)
|
||||
|
|
@ -7790,7 +8068,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(
|
||||
|
|
@ -7824,7 +8106,7 @@ async def _create_spend_logs_with_poison_isolation(
|
|||
await repo.table.create_many(data=rows, skip_duplicates=True)
|
||||
return failure_budget
|
||||
except Exception as e:
|
||||
if not PrismaDBExceptionHandler.is_prisma_data_error(e):
|
||||
if not (PrismaDBExceptionHandler.is_prisma_data_error(e) or isinstance(e, UnicodeEncodeError)):
|
||||
raise
|
||||
if PrismaDBExceptionHandler.is_database_service_unavailable_error(e):
|
||||
raise
|
||||
|
|
|
|||
|
|
@ -14,7 +14,7 @@ from litellm.responses.litellm_completion_transformation.custom_tools import (
|
|||
from litellm.responses.litellm_completion_transformation.transformation import (
|
||||
LiteLLMCompletionResponsesConfig,
|
||||
)
|
||||
from litellm.responses.streaming_iterator import ResponsesAPIStreamingIterator
|
||||
from litellm.responses.streaming_iterator import ResponsesAPIStreamingIterator, stamp_responses_usage_cost
|
||||
from litellm.responses.utils import ResponsesAPIRequestUtils
|
||||
from litellm.types.llms.openai import (
|
||||
PART_UNION_TYPES,
|
||||
|
|
@ -34,6 +34,7 @@ from litellm.types.llms.openai import (
|
|||
ReasoningSummaryTextDoneEvent,
|
||||
ResponseCompletedEvent,
|
||||
ResponseCreatedEvent,
|
||||
ResponseIncompleteEvent,
|
||||
ResponseInProgressEvent,
|
||||
ResponseInputParam,
|
||||
ResponsesAPIOptionalRequestParams,
|
||||
|
|
@ -912,7 +913,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
|
|||
raise StopAsyncIteration
|
||||
|
||||
self.finished = self.is_stream_finished()
|
||||
response_completed_event: Final = self._emit_response_completed_event(self.litellm_model_response)
|
||||
response_completed_event: Final = self._emit_terminal_response_event(self.litellm_model_response)
|
||||
if response_completed_event:
|
||||
# Latch so wrappers (FallbackResponsesStreamWrapper) + proxy
|
||||
# container-ownership hook can read completed_response.
|
||||
|
|
@ -1258,7 +1259,46 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
|
|||
)
|
||||
return reasoning_aligned
|
||||
|
||||
def _emit_response_completed_event(self, litellm_model_response: ModelResponse) -> ResponseCompletedEvent | None:
|
||||
def partial_response_on_disconnect(self) -> ResponseIncompleteEvent | None:
|
||||
"""
|
||||
The response.incomplete event to bill when the client went away mid-stream.
|
||||
Built from the chat completion chunks relayed so far, so the usage is what
|
||||
stream_chunk_builder counts for them.
|
||||
"""
|
||||
if self.completed_response is not None or not self.collected_chat_completion_chunks:
|
||||
return None
|
||||
built: Final = stream_chunk_builder(
|
||||
chunks=self.collected_chat_completion_chunks,
|
||||
messages=LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( # pyright: ignore[reportUnknownMemberType] # the transformer's signature is partially untyped
|
||||
input=self.request_input,
|
||||
responses_api_request=dict(self.responses_api_request),
|
||||
),
|
||||
logging_obj=self.litellm_logging_obj,
|
||||
)
|
||||
if not isinstance(built, ModelResponse):
|
||||
return None
|
||||
responses_api_response: Final = (
|
||||
LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( # pyright: ignore[reportUnknownMemberType] # the transformer's signature is partially untyped
|
||||
request_input=self.request_input,
|
||||
chat_completion_response=built,
|
||||
responses_api_request=self.responses_api_request,
|
||||
)
|
||||
)
|
||||
partial: Final = self._with_encoded_response_id(
|
||||
responses_api_response.model_copy(
|
||||
update={
|
||||
"id": self._cached_response_id or responses_api_response.id,
|
||||
"status": "incomplete",
|
||||
"output": list(self._output_with_streamed_item_ids(responses_api_response)),
|
||||
}
|
||||
)
|
||||
)
|
||||
stamp_responses_usage_cost(partial, self.litellm_logging_obj)
|
||||
return ResponseIncompleteEvent(type=ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE, response=partial)
|
||||
|
||||
def _emit_terminal_response_event(
|
||||
self, litellm_model_response: ModelResponse
|
||||
) -> ResponseCompletedEvent | ResponseIncompleteEvent | None:
|
||||
if litellm_model_response:
|
||||
# Transform the response
|
||||
responses_api_response: Final = (
|
||||
|
|
@ -1278,6 +1318,11 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
|
|||
# Encode the response ID to match non-streaming behavior
|
||||
encoded_response: Final = self._with_encoded_response_id(responses_api_response)
|
||||
|
||||
if responses_api_response.status == "incomplete":
|
||||
return ResponseIncompleteEvent(
|
||||
type=ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE,
|
||||
response=encoded_response,
|
||||
)
|
||||
return ResponseCompletedEvent(
|
||||
type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED,
|
||||
response=encoded_response,
|
||||
|
|
|
|||
|
|
@ -117,6 +117,44 @@ _INCOMPLETE_REASON_BY_FINISH_REASON: Final[Mapping[str, Literal["max_output_toke
|
|||
)
|
||||
|
||||
|
||||
class _EchoedResponsesRequestParams(TypedDict, total=False):
|
||||
"""Responses API request fields that the response object echoes back."""
|
||||
|
||||
instructions: ReadOnly[str]
|
||||
metadata: ReadOnly[dict[str, object]]
|
||||
parallel_tool_calls: ReadOnly[bool]
|
||||
temperature: ReadOnly[float]
|
||||
tools: ReadOnly[list[dict[str, object]]]
|
||||
top_p: ReadOnly[float]
|
||||
max_output_tokens: ReadOnly[int]
|
||||
previous_response_id: ReadOnly[str]
|
||||
reasoning: ReadOnly[dict[str, object]]
|
||||
text: ReadOnly[dict[str, object]]
|
||||
truncation: ReadOnly[Literal["auto", "disabled"]]
|
||||
user: ReadOnly[str]
|
||||
store: ReadOnly[bool]
|
||||
|
||||
|
||||
_ECHOED_PARAMS_ADAPTER: Final = TypeAdapter(_EchoedResponsesRequestParams)
|
||||
|
||||
|
||||
def _is_echoable(name: str, value: object) -> bool:
|
||||
try:
|
||||
_ = _ECHOED_PARAMS_ADAPTER.validate_python({name: value})
|
||||
except ValidationError:
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
def _echoable_request_params(request: Mapping[str, object]) -> _EchoedResponsesRequestParams:
|
||||
echoable: Final = {
|
||||
name: request[name]
|
||||
for name in _EchoedResponsesRequestParams.__optional_keys__
|
||||
if name in request and _is_echoable(name, request[name])
|
||||
}
|
||||
return _ECHOED_PARAMS_ADAPTER.validate_python(echoable)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class ResponsesToolChatForm:
|
||||
chat_tools: tuple[ChatToolParam, ...]
|
||||
|
|
@ -2294,7 +2332,7 @@ class LiteLLMCompletionResponsesConfig:
|
|||
The corresponding responses API status value (one of ResponsesAPIStatus)
|
||||
"""
|
||||
if finish_reason is None:
|
||||
return "completed"
|
||||
return "incomplete"
|
||||
|
||||
# Map finish reasons to status
|
||||
if finish_reason in ["stop", "tool_calls", "function_call"]:
|
||||
|
|
@ -2309,11 +2347,18 @@ class LiteLLMCompletionResponsesConfig:
|
|||
def _incomplete_details_for_finish_reason(
|
||||
finish_reason: str | None,
|
||||
existing: IncompleteDetails | None,
|
||||
completion_tokens: int | None = None,
|
||||
max_output_tokens: int | None = None,
|
||||
) -> IncompleteDetails | None:
|
||||
if existing is not None:
|
||||
return existing
|
||||
if finish_reason is None:
|
||||
return None
|
||||
hit_cap: Final = (
|
||||
completion_tokens is not None
|
||||
and max_output_tokens is not None
|
||||
and completion_tokens >= max_output_tokens
|
||||
)
|
||||
return IncompleteDetails(reason="max_output_tokens") if hit_cap else None
|
||||
reason: Final = _INCOMPLETE_REASON_BY_FINISH_REASON.get(finish_reason)
|
||||
return IncompleteDetails(reason=reason) if reason is not None else None
|
||||
|
||||
|
|
@ -2433,9 +2478,13 @@ class LiteLLMCompletionResponsesConfig:
|
|||
if choices and len(choices) > 0:
|
||||
finish_reason = choices[0].finish_reason
|
||||
|
||||
echoed: Final = _echoable_request_params(responses_api_request)
|
||||
chat_usage: Final = getattr(chat_completion_response, "usage", None)
|
||||
incomplete_details: Final = LiteLLMCompletionResponsesConfig._incomplete_details_for_finish_reason(
|
||||
finish_reason=finish_reason,
|
||||
existing=getattr(chat_completion_response, "incomplete_details", None),
|
||||
completion_tokens=getattr(chat_usage, "completion_tokens", None),
|
||||
max_output_tokens=echoed.get("max_output_tokens"),
|
||||
)
|
||||
|
||||
responses_api_response: Final[ResponsesAPIResponse] = ResponsesAPIResponse(
|
||||
|
|
@ -2445,32 +2494,33 @@ class LiteLLMCompletionResponsesConfig:
|
|||
object="response",
|
||||
error=getattr(chat_completion_response, "error", None),
|
||||
incomplete_details=incomplete_details,
|
||||
instructions=getattr(chat_completion_response, "instructions", None),
|
||||
metadata=getattr(chat_completion_response, "metadata", {}),
|
||||
instructions=echoed.get("instructions"),
|
||||
metadata=echoed.get("metadata") or {},
|
||||
output=LiteLLMCompletionResponsesConfig._transform_chat_completion_choices_to_responses_output(
|
||||
chat_completion_response=chat_completion_response,
|
||||
choices=getattr(chat_completion_response, "choices", []),
|
||||
responses_api_request=responses_api_request,
|
||||
),
|
||||
parallel_tool_calls=getattr(chat_completion_response, "parallel_tool_calls", False),
|
||||
temperature=getattr(chat_completion_response, "temperature", 0),
|
||||
parallel_tool_calls=echoed.get("parallel_tool_calls", False),
|
||||
temperature=echoed.get("temperature"),
|
||||
tool_choice=LiteLLMCompletionResponsesConfig._transform_tool_choice_for_responses_api_response(
|
||||
responses_api_request.get("tool_choice")
|
||||
),
|
||||
tools=getattr(chat_completion_response, "tools", []),
|
||||
top_p=getattr(chat_completion_response, "top_p", None),
|
||||
max_output_tokens=getattr(chat_completion_response, "max_output_tokens", None),
|
||||
previous_response_id=getattr(chat_completion_response, "previous_response_id", None),
|
||||
reasoning=None,
|
||||
tools=echoed.get("tools") or [],
|
||||
top_p=echoed.get("top_p"),
|
||||
max_output_tokens=echoed.get("max_output_tokens"),
|
||||
previous_response_id=echoed.get("previous_response_id"),
|
||||
reasoning=echoed.get("reasoning"),
|
||||
status=LiteLLMCompletionResponsesConfig._map_chat_completion_finish_reason_to_responses_status(
|
||||
finish_reason
|
||||
),
|
||||
text={},
|
||||
truncation=getattr(chat_completion_response, "truncation", None),
|
||||
text=echoed.get("text") or {},
|
||||
truncation=echoed.get("truncation"),
|
||||
usage=LiteLLMCompletionResponsesConfig._transform_chat_completion_usage_to_responses_usage(
|
||||
chat_completion_response=chat_completion_response
|
||||
),
|
||||
user=getattr(chat_completion_response, "user", None),
|
||||
user=echoed.get("user"),
|
||||
store=echoed.get("store"),
|
||||
)
|
||||
responses_api_response._hidden_params = getattr(chat_completion_response, "_hidden_params", {})
|
||||
|
||||
|
|
|
|||
|
|
@ -547,6 +547,22 @@ def _will_bridge_to_chat_completions(
|
|||
)
|
||||
|
||||
|
||||
def provider_stores_the_conversation(model: str, custom_llm_provider: str | None, kwargs: Mapping[str, object]) -> bool:
|
||||
"""Whether the provider, not a LiteLLM spend log, holds this turn for a `previous_response_id` follow-up.
|
||||
|
||||
A bridged turn is replayed from its spend log, so the bridge needs that row in the DB, and a bridged
|
||||
first MCP round runs as an internal call that writes no spend log at all.
|
||||
"""
|
||||
api_base: Final = _api_base_kwarg(kwargs)
|
||||
return not _will_bridge_to_chat_completions(
|
||||
model,
|
||||
custom_llm_provider or litellm.get_llm_provider(model=model, api_base=api_base)[1],
|
||||
bool(kwargs.get("use_chat_completions_api")),
|
||||
kwargs.get("model_info"),
|
||||
api_base,
|
||||
)
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _prompt_management_sees_a_provisional_message_list(
|
||||
kwargs: dict[str, object], # mutable-ok: the signal is read and popped out of the caller's own kwargs
|
||||
|
|
|
|||
|
|
@ -41,10 +41,12 @@ from litellm.types.integrations.custom_logger import converted_stream_requested
|
|||
from litellm.types.llms.openai import (
|
||||
PART_UNION_TYPES,
|
||||
ResponseAPIUsage,
|
||||
ResponseIncompleteEvent,
|
||||
ResponsesAPIResponse,
|
||||
ResponsesAPIStreamEvents,
|
||||
ResponsesAPIStreamingResponse,
|
||||
)
|
||||
from litellm.types.responses.main import GenericResponseOutputItem, OutputText
|
||||
from litellm.types.utils import CallTypes
|
||||
from litellm.utils import async_post_call_success_deployment_hook
|
||||
|
||||
|
|
@ -289,6 +291,7 @@ class BaseResponsesAPIStreamingIterator:
|
|||
self._yielded_first_chunk = False
|
||||
self._generated_content = ""
|
||||
self._generated_tool_arguments = ""
|
||||
self._created_response: ResponsesAPIResponse | None = None
|
||||
self._completed_response_cached = False
|
||||
self._completed_response_logged = False
|
||||
self._completed_response_cache_hit: bool | None = None
|
||||
|
|
@ -450,6 +453,10 @@ class BaseResponsesAPIStreamingIterator:
|
|||
# Store the completed response (also for incomplete/failed so logging still fires)
|
||||
_chunk_type: Final = getattr(openai_responses_api_chunk, "type", None)
|
||||
openai_types = _get_openai_response_types()
|
||||
if _chunk_type == openai_types.ResponsesAPIStreamEvents.RESPONSE_CREATED:
|
||||
_created: Final[object] = getattr(openai_responses_api_chunk, "response", None)
|
||||
if isinstance(_created, ResponsesAPIResponse):
|
||||
self._created_response = _created
|
||||
if openai_responses_api_chunk and _chunk_type in (
|
||||
openai_types.ResponsesAPIStreamEvents.RESPONSE_COMPLETED,
|
||||
openai_types.ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE,
|
||||
|
|
@ -481,7 +488,7 @@ class BaseResponsesAPIStreamingIterator:
|
|||
else openai_responses_api_chunk.model_copy(update={"response": _billed_response})
|
||||
)
|
||||
self.completed_response = _terminal_chunk
|
||||
_stamp_responses_usage_cost(_billed_response, self.logging_obj)
|
||||
stamp_responses_usage_cost(_billed_response, self.logging_obj)
|
||||
|
||||
if _chunk_type == openai_types.ResponsesAPIStreamEvents.RESPONSE_FAILED:
|
||||
self._handle_logging_failed_response()
|
||||
|
|
@ -874,6 +881,34 @@ class BaseResponsesAPIStreamingIterator:
|
|||
except Exception:
|
||||
pass
|
||||
|
||||
def partial_response_on_disconnect(self) -> ResponseIncompleteEvent | None:
|
||||
"""
|
||||
The response.incomplete event to bill when the client went away mid-stream,
|
||||
the same shape the success handler assembles for a natural end of stream.
|
||||
Its response is the response.created payload (it carries the served
|
||||
service_tier) plus the text and tool arguments relayed so far, with counted
|
||||
usage since the provider never sent any. None when nothing was relayed yet
|
||||
or the stream already reached its terminal event and billed itself.
|
||||
"""
|
||||
created: Final = self._created_response
|
||||
generated: Final = self._generated_content + self._generated_tool_arguments
|
||||
if self.completed_response is not None or created is None or not generated:
|
||||
return None
|
||||
usage: Final = _estimate_usage_safely(
|
||||
self.model or "", self.request_data.get("input"), self.request_data, generated
|
||||
)
|
||||
if usage is None:
|
||||
return None
|
||||
partial: Final = created.model_copy(
|
||||
update={
|
||||
"status": "incomplete",
|
||||
"usage": usage,
|
||||
"output": _partial_output_items(self._generated_content),
|
||||
}
|
||||
)
|
||||
stamp_responses_usage_cost(partial, self.logging_obj)
|
||||
return ResponseIncompleteEvent(type=ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE, response=partial)
|
||||
|
||||
|
||||
async def call_post_streaming_hooks_for_testing(
|
||||
iterator: object, chunk: ResponsesAPIStreamingResponse
|
||||
|
|
@ -1347,6 +1382,37 @@ def _add_text_like_part_events(
|
|||
)
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class _BillsPartialStreamOnDisconnect(Protocol):
|
||||
def partial_response_on_disconnect(self) -> ResponseIncompleteEvent | None: ...
|
||||
|
||||
|
||||
def partial_responses_incomplete_event_on_disconnect(stream: object) -> ResponseIncompleteEvent | None:
|
||||
"""Ask the stream behind a proxy or router wrapper for the partial response to bill, if it can build one."""
|
||||
if not isinstance(stream, _BillsPartialStreamOnDisconnect):
|
||||
return None
|
||||
try:
|
||||
partial: Final = stream.partial_response_on_disconnect()
|
||||
except Exception as e: # noqa: BLE001 # disconnect billing is best-effort; never break stream teardown
|
||||
verbose_logger.debug("Failed to build partial Responses API response for disconnect billing: %s", e)
|
||||
return None
|
||||
return partial if isinstance(partial, ResponseIncompleteEvent) else None
|
||||
|
||||
|
||||
def _partial_output_items(generated_text: str) -> Sequence[GenericResponseOutputItem]:
|
||||
if not generated_text:
|
||||
return []
|
||||
return [
|
||||
GenericResponseOutputItem(
|
||||
type="message",
|
||||
id=f"msg_{uuid.uuid4()}",
|
||||
status="incomplete",
|
||||
role="assistant",
|
||||
content=[OutputText(type="output_text", text=generated_text, annotations=[])],
|
||||
)
|
||||
]
|
||||
|
||||
|
||||
def _billed_terminal_response(
|
||||
response_obj: object, estimate: Callable[[], ResponseAPIUsage | None] | None
|
||||
) -> ResponsesAPIResponse | None:
|
||||
|
|
@ -1453,7 +1519,7 @@ def _estimate_usage_safely(
|
|||
return None
|
||||
|
||||
|
||||
def _stamp_responses_usage_cost(
|
||||
def stamp_responses_usage_cost(
|
||||
response_obj: ResponsesAPIResponse | None, logging_obj: LiteLLMLoggingObj | None
|
||||
) -> None:
|
||||
if response_obj is None or logging_obj is None:
|
||||
|
|
@ -1479,7 +1545,7 @@ def build_synthetic_response_events(
|
|||
chunk_size: int,
|
||||
) -> list[ResponsesAPIStreamingResponse]:
|
||||
openai_types: Final = _get_openai_response_types()
|
||||
_stamp_responses_usage_cost(transformed, logging_obj)
|
||||
stamp_responses_usage_cost(transformed, logging_obj)
|
||||
|
||||
events: Final[list[ResponsesAPIStreamingResponse]] = [
|
||||
_build_response_status_event(openai_types.ResponsesAPIStreamEvents.RESPONSE_CREATED, transformed),
|
||||
|
|
@ -1696,6 +1762,8 @@ RESPONSES_WS_LOGGED_EVENT_TYPES: Final = [
|
|||
"error",
|
||||
]
|
||||
|
||||
_HISTORY_TERMINAL_EVENT_TYPES: Final = frozenset({"response.completed", "response.incomplete"})
|
||||
|
||||
RESPONSES_WS_MASKABLE_TEXT_BLOCK_TYPES: Final = frozenset({"input_text", "output_text", "text"})
|
||||
|
||||
_RESPONSES_WS_FAILURE_EVENT_TYPES: Final = frozenset({"error", "response.failed"})
|
||||
|
|
@ -2479,7 +2547,7 @@ class ManagedResponsesWebSocketHandler:
|
|||
@staticmethod
|
||||
def _extract_response_id(completed_event: _MutableJsonObject) -> str | None:
|
||||
"""
|
||||
Pull the raw (decoded) response ID out of a ``response.completed`` event.
|
||||
Pull the raw (decoded) response ID out of a terminal (``response.completed`` or ``response.incomplete``) event.
|
||||
Returns *None* if the event doesn't contain a usable ID.
|
||||
"""
|
||||
resp_obj: Final = completed_event.get("response", {})
|
||||
|
|
@ -2495,8 +2563,8 @@ class ManagedResponsesWebSocketHandler:
|
|||
completed_event: _MutableJsonObject,
|
||||
) -> list[dict[str, object]]:
|
||||
"""
|
||||
Convert the output items in a ``response.completed`` event into
|
||||
Responses API message dicts suitable for the next turn's ``input``.
|
||||
Convert output items in a terminal (``response.completed`` or ``response.incomplete``) event into Responses
|
||||
API message dicts suitable for the next turn's ``input``.
|
||||
"""
|
||||
resp_obj: Final = completed_event.get("response", {})
|
||||
if not isinstance(resp_obj, dict):
|
||||
|
|
@ -2737,11 +2805,11 @@ class ManagedResponsesWebSocketHandler:
|
|||
"""
|
||||
Stream ``litellm.aresponses`` and forward every chunk over the WebSocket.
|
||||
|
||||
Captures the ``response.completed`` event type from the chunk object
|
||||
directly (before serialization) to avoid a redundant JSON round-trip on
|
||||
every chunk. Returns the completed event dict, or ``None``.
|
||||
Captures the first terminal (``response.completed`` or ``response.incomplete``) event from the chunk object
|
||||
directly (before serialization) to avoid a redundant JSON round-trip on every chunk. Returns the terminal
|
||||
event dict, or ``None``.
|
||||
"""
|
||||
completed_event: _MutableJsonObject | None = None
|
||||
terminal_event: _MutableJsonObject | None = None
|
||||
stream_response: Final = await litellm.aresponses(model=model, **call_kwargs)
|
||||
async for chunk in stream_response:
|
||||
if chunk is None:
|
||||
|
|
@ -2751,31 +2819,31 @@ class ManagedResponsesWebSocketHandler:
|
|||
serialized = self._serialize_chunk(chunk)
|
||||
if serialized is None:
|
||||
continue
|
||||
if chunk_type == "response.completed" and completed_event is None:
|
||||
if chunk_type in _HISTORY_TERMINAL_EVENT_TYPES and terminal_event is None:
|
||||
try:
|
||||
completed_event = _load_json_object(serialized)
|
||||
terminal_event = _load_json_object(serialized)
|
||||
except Exception:
|
||||
pass
|
||||
try:
|
||||
await self.websocket.send_text(serialized)
|
||||
except Exception as send_exc:
|
||||
verbose_logger.debug("ManagedResponsesWS: error sending chunk to client: %s", send_exc)
|
||||
return completed_event # Client disconnected
|
||||
return completed_event
|
||||
return terminal_event # Client disconnected
|
||||
return terminal_event
|
||||
|
||||
def _save_turn_history(
|
||||
self,
|
||||
completed_event: _MutableJsonObject | None,
|
||||
terminal_event: _MutableJsonObject | None,
|
||||
prior_history: list[dict[str, object]],
|
||||
current_messages: list[dict[str, object]],
|
||||
) -> None:
|
||||
"""Store this turn in in-memory history for future previous_response_id lookups."""
|
||||
if completed_event is None:
|
||||
"""Store this terminal (``response.completed`` or ``response.incomplete``) event for future history lookups."""
|
||||
if terminal_event is None:
|
||||
return
|
||||
new_response_id: Final = self._extract_response_id(completed_event)
|
||||
new_response_id: Final = self._extract_response_id(terminal_event)
|
||||
if not new_response_id:
|
||||
return
|
||||
output_msgs: Final = self._extract_output_messages(completed_event)
|
||||
output_msgs: Final = self._extract_output_messages(terminal_event)
|
||||
all_messages: Final = prior_history + current_messages + output_msgs
|
||||
self._store_history(new_response_id, all_messages)
|
||||
verbose_logger.debug(
|
||||
|
|
@ -2801,7 +2869,8 @@ class ManagedResponsesWebSocketHandler:
|
|||
2. Prepend those messages to the current ``input`` so the model has full
|
||||
conversation context.
|
||||
3. After the stream completes, extract the new response ID and output
|
||||
messages from ``response.completed`` and store them in
|
||||
messages from the terminal (``response.completed`` or
|
||||
``response.incomplete``) event and store them in
|
||||
``self._session_history`` for the next turn.
|
||||
|
||||
This in-memory approach avoids the async DB-write race condition that
|
||||
|
|
@ -2855,13 +2924,13 @@ class ManagedResponsesWebSocketHandler:
|
|||
call_kwargs.update(self.extra_kwargs)
|
||||
|
||||
try:
|
||||
completed_event: Final = await self._stream_and_forward(model, call_kwargs)
|
||||
terminal_event: Final = await self._stream_and_forward(model, call_kwargs)
|
||||
except Exception as exc:
|
||||
verbose_logger.exception("ManagedResponsesWS: error processing response.create: %s", exc)
|
||||
await self._send_error(str(exc))
|
||||
return
|
||||
|
||||
self._save_turn_history(completed_event, prior_history, current_messages)
|
||||
self._save_turn_history(terminal_event, prior_history, current_messages)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Main entry point
|
||||
|
|
|
|||
|
|
@ -86,6 +86,7 @@ from litellm.litellm_core_utils.get_llm_provider_logic import (
|
|||
is_registered_custom_provider,
|
||||
)
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
|
||||
from litellm.litellm_core_utils.llm_cost_calc.utils import SERVICE_TIER_COST_KEY_SUFFIXES
|
||||
from litellm.litellm_core_utils.ptu_pricing import (
|
||||
PTU_COST_ATTRIBUTION_ENV_VAR,
|
||||
declares_ptu,
|
||||
|
|
@ -97,6 +98,7 @@ from litellm.litellm_core_utils.ptu_pricing import (
|
|||
)
|
||||
from litellm.litellm_core_utils.request_timeout_resolver import (
|
||||
get_configured_request_timeout,
|
||||
get_service_tier_timeout,
|
||||
)
|
||||
from litellm.litellm_core_utils.secret_redaction import redact_string
|
||||
from litellm.litellm_core_utils.sensitive_data_masker import (
|
||||
|
|
@ -198,6 +200,7 @@ from litellm.router_utils.cooldown_handlers import (
|
|||
is_background_response_cost_poll_not_found,
|
||||
is_caller_timeout_408,
|
||||
)
|
||||
from litellm.router_utils.deployment_pricing import deployment_pricing_models
|
||||
from litellm.router_utils.fallback_event_handlers import (
|
||||
MID_STREAM_FALLBACK_CONTROLS_KEY,
|
||||
AttemptedFallbackTargets,
|
||||
|
|
@ -317,12 +320,14 @@ from litellm.types.utils import (
|
|||
StandardLoggingRoutingDecision,
|
||||
Usage,
|
||||
all_litellm_params,
|
||||
is_server_derived_pricing_key,
|
||||
shared_backend_model_info,
|
||||
)
|
||||
from litellm.types.utils import ModelInfo as ModelMapInfo
|
||||
from litellm.utils import (
|
||||
CustomStreamWrapper,
|
||||
EmbeddingResponse,
|
||||
ModelInfoLookupMiss,
|
||||
ModelResponse,
|
||||
Rules,
|
||||
function_setup,
|
||||
|
|
@ -331,6 +336,8 @@ from litellm.utils import (
|
|||
get_secret,
|
||||
get_utc_datetime,
|
||||
is_region_allowed,
|
||||
lookup_model_info_for_registration,
|
||||
model_info_with_pricing_overrides,
|
||||
provider_rejectable_params,
|
||||
set_live_deployment_replay,
|
||||
)
|
||||
|
|
@ -363,6 +370,7 @@ if TYPE_CHECKING:
|
|||
from litellm.types.llms.base import BaseLiteLLMOpenAIResponseObject
|
||||
from litellm.types.llms.openai import (
|
||||
ResponseAPIUsage,
|
||||
ResponseIncompleteEvent,
|
||||
ResponseInputParam,
|
||||
ResponsesAPIResponse,
|
||||
)
|
||||
|
|
@ -470,6 +478,8 @@ def _stream_chunks_have_generated_content(chunks: Sequence[ModelResponseStream])
|
|||
_NO_SESSION_KWARGS: Final[Mapping[str, Mapping[str, object]]] = MappingProxyType({})
|
||||
_SESSION_ADAPTER: Final = TypeAdapter(Mapping[str, object])
|
||||
_SILENT_MODEL_ADAPTER: Final = TypeAdapter(str | list[str])
|
||||
_MODEL_INFO_ADAPTER: Final = TypeAdapter(Mapping[str, object])
|
||||
_LITELLM_PARAMS_ADAPTER: Final = TypeAdapter(Mapping[str, object])
|
||||
|
||||
|
||||
def _as_retry_skipped_deployment_ids(value: object) -> tuple[str, ...]:
|
||||
|
|
@ -629,6 +639,11 @@ class FallbackAwareAnthropicMessagesStream:
|
|||
self._source_iterator = fallback_response
|
||||
self.fallback_headers_adopted = True
|
||||
|
||||
def partial_response_on_disconnect(self) -> "ResponseIncompleteEvent | None":
|
||||
from litellm.responses.streaming_iterator import partial_responses_incomplete_event_on_disconnect
|
||||
|
||||
return partial_responses_incomplete_event_on_disconnect(self._source_iterator)
|
||||
|
||||
def __aiter__(self) -> "FallbackAwareAnthropicMessagesStream":
|
||||
return self
|
||||
|
||||
|
|
@ -3242,6 +3257,7 @@ class Router:
|
|||
from datetime import datetime
|
||||
|
||||
self._async_generator = async_generator
|
||||
self._billing_source: object = source_iterator
|
||||
# Mirror every attribute BaseResponsesAPIStreamingIterator.__init__
|
||||
# would have set. The wrapper bypasses super().__init__ (it has no
|
||||
# httpx.Response of its own and no provider config to drive), so
|
||||
|
|
@ -3290,10 +3306,29 @@ class Router:
|
|||
|
||||
def adopt_fallback_headers(self, fallback_response: object) -> tuple[dict[str, object], dict[str, object]]:
|
||||
prepared: Final = Router._prepare_fallback_hidden_params(fallback_response)
|
||||
self._billing_source = fallback_response
|
||||
self._hidden_params = {**prepared[0], "additional_headers": prepared[1]} # mutable-ok: stream metadata
|
||||
self.fallback_headers_adopted = True
|
||||
return prepared
|
||||
|
||||
def partial_response_on_disconnect(self) -> "ResponseIncompleteEvent | None":
|
||||
"""
|
||||
The event to bill on a client disconnect. After a mid-stream hop the primary
|
||||
stream's relayed tokens are folded into the fallback's event, the way
|
||||
_combine_responses_fallback_usage does for a natural end of stream.
|
||||
"""
|
||||
from litellm.responses.streaming_iterator import partial_responses_incomplete_event_on_disconnect
|
||||
|
||||
current: Final = partial_responses_incomplete_event_on_disconnect(self._billing_source)
|
||||
if self._billing_source is source_iterator:
|
||||
return current
|
||||
primary: Final = partial_responses_incomplete_event_on_disconnect(source_iterator)
|
||||
if current is None or primary is None:
|
||||
return current or primary
|
||||
if primary.response.usage is not None:
|
||||
Router._combine_responses_fallback_usage(current, primary.response.usage)
|
||||
return current
|
||||
|
||||
def __aiter__(self):
|
||||
return self
|
||||
|
||||
|
|
@ -4033,6 +4068,7 @@ class Router:
|
|||
"""Helper to get stream timeout from kwargs or deployment params"""
|
||||
return (
|
||||
kwargs.get("stream_timeout", None) # the params dynamically set by user
|
||||
or get_service_tier_timeout(kwargs.get("service_tier"), data, "stream_timeout_by_service_tier")
|
||||
or data.get("stream_timeout", None) # timeout set on litellm_params for this deployment
|
||||
or self.stream_timeout # timeout set on router
|
||||
or self.request_timeout # litellm_settings.request_timeout (per-attempt)
|
||||
|
|
@ -4044,6 +4080,7 @@ class Router:
|
|||
timeout: Final = (
|
||||
kwargs.get("timeout", None) # the params dynamically set by user
|
||||
or kwargs.get("request_timeout", None) # the params dynamically set by user
|
||||
or get_service_tier_timeout(kwargs.get("service_tier"), data, "timeout_by_service_tier")
|
||||
or data.get("timeout", None) # timeout set on litellm_params for this deployment
|
||||
or data.get("request_timeout", None) # timeout set on litellm_params for this deployment
|
||||
or self.request_timeout # litellm_settings.request_timeout (per-attempt)
|
||||
|
|
@ -8299,17 +8336,24 @@ class Router:
|
|||
When a retry or fallback happens, record which model group, deployment and attempt just failed and why,
|
||||
and count it toward the request-wide num_retries_per_request cap
|
||||
"""
|
||||
from litellm.litellm_core_utils.get_litellm_params import _get_base_model_from_litellm_call_metadata
|
||||
from litellm.types.router import RetryAttemptRecord
|
||||
|
||||
_metadata_var: Final = "litellm_metadata" if "litellm_metadata" in kwargs else "metadata"
|
||||
request_metadata: Final[Mapping[str, object]] = kwargs[_metadata_var]
|
||||
request_metadata: Final = cast("Mapping[str, object]", kwargs[_metadata_var]) # cast-ok: untyped kwargs
|
||||
model_group: Final = kwargs.get("model")
|
||||
model_info: Final = request_metadata.get("model_info")
|
||||
deployment_id: Final = model_info.get("id") if isinstance(model_info, Mapping) else None
|
||||
base_model: Final = _get_base_model_from_litellm_call_metadata(
|
||||
metadata=request_metadata
|
||||
) or self._deployment_params_base_model(deployment_id if isinstance(deployment_id, str) else None)
|
||||
deployment_model: Final = request_metadata.get("deployment")
|
||||
attempted_retries: Final = request_metadata.get("attempted_retries")
|
||||
attempt_record: Final[RetryAttemptRecord] = {
|
||||
"model_group": model_group if isinstance(model_group, str) else None,
|
||||
"deployment_id": deployment_id if isinstance(deployment_id, str) else None,
|
||||
"model": deployment_model if isinstance(deployment_model, str) else None,
|
||||
"base_model": base_model if isinstance(base_model, str) else None,
|
||||
"exception_type": type(e).__name__,
|
||||
"exception_string": str(e),
|
||||
"attempted_retries": attempted_retries if type(attempted_retries) is int else None,
|
||||
|
|
@ -8327,6 +8371,18 @@ class Router:
|
|||
kwargs[_metadata_var]["request_retry_count"] = request_retry_count # rebind-ok: same dict, read by the cap
|
||||
return kwargs
|
||||
|
||||
def _deployment_params_base_model(self, deployment_id: str | None) -> str | None:
|
||||
"""The ``base_model`` a deployment declares under ``litellm_params`` (the Azure named-deployment
|
||||
shape), the same key the success path's ``get_base_model_from_metadata`` reads."""
|
||||
deployment: Final = self.get_deployment(model_id=deployment_id) if deployment_id is not None else None
|
||||
if deployment is None:
|
||||
return None
|
||||
extra_params: Final = cast( # cast-ok: pydantic extras are untyped, base_model is isinstance-checked below
|
||||
"Mapping[str, object]", deployment.litellm_params.model_extra or {}
|
||||
)
|
||||
base_model: Final = extra_params.get("base_model")
|
||||
return base_model if isinstance(base_model, str) else None
|
||||
|
||||
def _update_usage(self, deployment_id: str, parent_otel_span: Span | None) -> int:
|
||||
"""
|
||||
Update deployment rpm for that minute
|
||||
|
|
@ -8866,34 +8922,7 @@ class Router:
|
|||
litellm_params=litellm_params,
|
||||
model_info=_model_info,
|
||||
)
|
||||
for field in CustomPricingLiteLLMParams.model_fields:
|
||||
if deployment.litellm_params.get(field) is not None:
|
||||
_model_info[field] = deployment.litellm_params[field]
|
||||
|
||||
Router._inherit_builtin_base_rates_for_off_peak(
|
||||
model_info=_model_info,
|
||||
backend_model=deployment.litellm_params.model,
|
||||
custom_llm_provider=deployment.litellm_params.custom_llm_provider,
|
||||
)
|
||||
if _model_info.get("input_cost_per_token") is not None:
|
||||
Router._inherit_builtin_cache_pricing(
|
||||
model_info=_model_info,
|
||||
backend_model=deployment.litellm_params.model,
|
||||
custom_llm_provider=deployment.litellm_params.custom_llm_provider,
|
||||
)
|
||||
Router._inherit_builtin_tiered_output_rate(
|
||||
model_info=_model_info,
|
||||
backend_model=deployment.litellm_params.model,
|
||||
custom_llm_provider=deployment.litellm_params.custom_llm_provider,
|
||||
)
|
||||
|
||||
## REGISTER MODEL INFO IN LITELLM MODEL COST MAP
|
||||
Router._register_deployment_in_model_cost(
|
||||
model_id=deployment.model_info.id,
|
||||
model_info=_model_info,
|
||||
model=deployment.litellm_params.model,
|
||||
custom_llm_provider=deployment.litellm_params.custom_llm_provider,
|
||||
)
|
||||
Router._register_deployment_pricing(deployment)
|
||||
|
||||
## Check if LLM Deployment is allowed for this deployment
|
||||
if self.deployment_is_active_for_environment(deployment=deployment) is not True:
|
||||
|
|
@ -9897,7 +9926,56 @@ class Router:
|
|||
return (backend_key,)
|
||||
|
||||
@staticmethod
|
||||
def _deployment_model_cost_payload(deployment: Deployment) -> dict: # mutable-ok: cost-map entry
|
||||
def _inherit_builtin_service_tier_pricing(
|
||||
model_info: dict, # mutable-ok: deployment cost-map entry filled in place
|
||||
backend_model: str,
|
||||
custom_llm_provider: str | None,
|
||||
) -> None:
|
||||
"""Inherit missing tier rates so a standalone entry does not fall back to custom standard rates."""
|
||||
if ptu_terms(model_info) is not None and is_ptu_cost_attribution_enabled():
|
||||
return
|
||||
if all(model_info.get(field) is None for field in ("input_cost_per_token", "output_cost_per_token")):
|
||||
return
|
||||
try:
|
||||
backend_info: Final = litellm.get_model_info(model=backend_model, custom_llm_provider=custom_llm_provider)
|
||||
except Exception: # noqa: BLE001 # get_model_info raises plain Exception for an unmapped backend model
|
||||
return
|
||||
backend_entry: Final = litellm.model_cost.get(backend_info.get("key") or "")
|
||||
if not isinstance(backend_entry, dict):
|
||||
return
|
||||
for field, backend_value in backend_entry.items():
|
||||
if not field.endswith(SERVICE_TIER_COST_KEY_SUFFIXES):
|
||||
continue
|
||||
if model_info.get(field) is not None or backend_value is None:
|
||||
continue
|
||||
model_info[field] = copy.deepcopy(backend_value)
|
||||
|
||||
@staticmethod
|
||||
def _cost_map_backend_model(deployment: Deployment) -> str:
|
||||
model_info_base_model: Final = deployment.model_info.base_model
|
||||
if isinstance(model_info_base_model, str) and model_info_base_model:
|
||||
return model_info_base_model
|
||||
params_base_model: Final = deployment.litellm_params.get("base_model")
|
||||
if isinstance(params_base_model, str) and params_base_model:
|
||||
return params_base_model
|
||||
params: Final = _LITELLM_PARAMS_ADAPTER.validate_python(deployment.litellm_params.model_dump(exclude_none=True))
|
||||
invoked_model: Final = params.get("model_id")
|
||||
if (
|
||||
isinstance(invoked_model, str)
|
||||
and invoked_model
|
||||
and (
|
||||
deployment.litellm_params.custom_llm_provider == "bedrock"
|
||||
or deployment.litellm_params.model.startswith("bedrock/")
|
||||
)
|
||||
):
|
||||
profile_model: Final = invoked_model if invoked_model.startswith("bedrock/") else f"bedrock/{invoked_model}"
|
||||
profile_info: Final = lookup_model_info_for_registration(model=profile_model, custom_llm_provider="bedrock")
|
||||
if not isinstance(profile_info, ModelInfoLookupMiss):
|
||||
return profile_model
|
||||
return deployment.litellm_params.model
|
||||
|
||||
@staticmethod
|
||||
def _deployment_model_cost_payload(deployment: Deployment) -> dict[str, object]: # mutable-ok: cost-map entry
|
||||
"""The ``model_info`` a deployment contributes to ``litellm.model_cost``.
|
||||
|
||||
Custom pricing lives on ``litellm_params`` rather than ``model_info``, and
|
||||
|
|
@ -9905,25 +9983,49 @@ class Router:
|
|||
both are folded back in here. That keeps this reproducible from a
|
||||
deployment alone, which is what lets a refresh rebuild the same entries.
|
||||
"""
|
||||
model_info: Final[dict] = deployment.model_info.model_dump(exclude_none=True) # mutable-ok: built in place
|
||||
model_info: Final = dict(
|
||||
_MODEL_INFO_ADAPTER.validate_python(deployment.model_info.model_dump(exclude_none=True))
|
||||
)
|
||||
for field in CustomPricingLiteLLMParams.model_fields:
|
||||
field_value = deployment.litellm_params.get(field)
|
||||
if field_value is not None:
|
||||
model_info[field] = field_value
|
||||
backend_model: Final = Router._cost_map_backend_model(deployment)
|
||||
params: Final = _LITELLM_PARAMS_ADAPTER.validate_python(deployment.litellm_params.model_dump(exclude_none=True))
|
||||
if (
|
||||
params.get("model_id")
|
||||
and backend_model != deployment.litellm_params.model
|
||||
and (
|
||||
deployment.litellm_params.custom_llm_provider == "bedrock"
|
||||
or deployment.litellm_params.model.startswith("bedrock/")
|
||||
)
|
||||
):
|
||||
profile_info: Final = lookup_model_info_for_registration(model=backend_model, custom_llm_provider="bedrock")
|
||||
if not isinstance(profile_info, ModelInfoLookupMiss):
|
||||
explicit_info: Final = _MODEL_INFO_ADAPTER.validate_python(model_info)
|
||||
priced_info: Final = model_info_with_pricing_overrides(profile_info, explicit_info)
|
||||
model_info.update(
|
||||
{key: value for key, value in priced_info.items() if is_server_derived_pricing_key(key)}
|
||||
)
|
||||
Router._inherit_builtin_base_rates_for_off_peak(
|
||||
model_info=model_info,
|
||||
backend_model=deployment.litellm_params.model,
|
||||
backend_model=backend_model,
|
||||
custom_llm_provider=deployment.litellm_params.custom_llm_provider,
|
||||
)
|
||||
if model_info.get("input_cost_per_token") is not None:
|
||||
Router._inherit_builtin_cache_pricing(
|
||||
model_info=model_info,
|
||||
backend_model=deployment.litellm_params.model,
|
||||
backend_model=backend_model,
|
||||
custom_llm_provider=deployment.litellm_params.custom_llm_provider,
|
||||
)
|
||||
Router._inherit_builtin_service_tier_pricing(
|
||||
model_info=model_info,
|
||||
backend_model=backend_model,
|
||||
custom_llm_provider=deployment.litellm_params.custom_llm_provider,
|
||||
)
|
||||
Router._inherit_builtin_tiered_output_rate(
|
||||
model_info=model_info,
|
||||
backend_model=deployment.litellm_params.model,
|
||||
backend_model=backend_model,
|
||||
custom_llm_provider=deployment.litellm_params.custom_llm_provider,
|
||||
)
|
||||
return model_info
|
||||
|
|
@ -10325,19 +10427,9 @@ class Router:
|
|||
)
|
||||
for deployment in deployments
|
||||
)
|
||||
params: Final = tuple(deployment.get("litellm_params") or MappingProxyType({}) for deployment in deployments)
|
||||
# base_model resolution mirrors get_router_model_info: unset or blank means the
|
||||
# deployment's own model name is the cost-map key.
|
||||
cost_map_keys: Final = tuple(
|
||||
dict.fromkeys( # deduplicates while preserving config order
|
||||
key
|
||||
for key in (
|
||||
model_info.get("base_model") or litellm_params.get("base_model") or litellm_params.get("model")
|
||||
for model_info, litellm_params in zip(model_infos, params)
|
||||
)
|
||||
if isinstance(key, str) and key
|
||||
)
|
||||
)
|
||||
raw_deployments: Final[tuple[object, ...]] = deployments
|
||||
pricing_keys: Final = (deployment_pricing_models(deployment) for deployment in raw_deployments)
|
||||
cost_map_keys: Final = tuple(dict.fromkeys(keys[0] for keys in pricing_keys if keys))
|
||||
return DeploymentModelListingInfo(
|
||||
cost_map_keys=cost_map_keys,
|
||||
max_input_tokens=self._widest_configured_limit(model_infos, "max_input_tokens"),
|
||||
|
|
|
|||
|
|
@ -27,6 +27,8 @@ from itertools import groupby
|
|||
from types import MappingProxyType
|
||||
from typing import Any, Final
|
||||
|
||||
from pydantic import TypeAdapter
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_router_logger
|
||||
from litellm.caching.caching import DualCache
|
||||
|
|
@ -37,6 +39,7 @@ from litellm.litellm_core_utils.core_helpers import (
|
|||
)
|
||||
from litellm.litellm_core_utils.duration_parser import duration_in_seconds
|
||||
from litellm.router_strategy.tag_based_routing import _get_tags_from_request_kwargs
|
||||
from litellm.router_utils.batch_utils import batch_charged_once_marker_key, batch_id_to_charge_once
|
||||
from litellm.router_utils.cooldown_callbacks import (
|
||||
_get_prometheus_logger_from_callbacks,
|
||||
)
|
||||
|
|
@ -46,6 +49,8 @@ from litellm.types.utils import BudgetConfig, GenericBudgetConfigType, StandardL
|
|||
from litellm.types.utils import BudgetConfig as GenericBudgetInfo
|
||||
|
||||
DEFAULT_REDIS_SYNC_INTERVAL: Final = 1
|
||||
_CALLBACK_FIELDS: Final = TypeAdapter(dict[str, object])
|
||||
_BATCH_CHARGE_VALUES: Final = TypeAdapter(tuple[object, float])
|
||||
|
||||
|
||||
class _LiteLLMParamsDictView:
|
||||
|
|
@ -500,14 +505,18 @@ class RouterBudgetLimiting(CustomLogger):
|
|||
if standard_logging_payload is None:
|
||||
raise ValueError("standard_logging_payload is required")
|
||||
|
||||
response_cost: Final[float] = standard_logging_payload.get("response_cost", 0)
|
||||
batch_response, response_cost = _BATCH_CHARGE_VALUES.validate_python(
|
||||
(response_obj, _CALLBACK_FIELDS.validate_python(standard_logging_payload).get("response_cost", 0))
|
||||
)
|
||||
model_id: Final[str] = str(standard_logging_payload.get("model_id", ""))
|
||||
custom_llm_provider: Final[str] = kwargs.get("litellm_params", {}).get("custom_llm_provider", None)
|
||||
if custom_llm_provider is None:
|
||||
raise ValueError("custom_llm_provider is required")
|
||||
custom_llm_provider: Final[str | None] = standard_logging_payload.get("custom_llm_provider")
|
||||
|
||||
budget_config: Final = self._get_budget_config_for_provider(custom_llm_provider)
|
||||
if budget_config:
|
||||
batch_id: Final = batch_id_to_charge_once(kwargs.get("call_type"), batch_response, response_cost)
|
||||
|
||||
budget_config: Final = (
|
||||
self._get_budget_config_for_provider(custom_llm_provider) if custom_llm_provider is not None else None
|
||||
)
|
||||
if custom_llm_provider is not None and budget_config is not None:
|
||||
# increment spend for provider
|
||||
spend_key: Final = f"provider_spend:{custom_llm_provider}:{budget_config.budget_duration}"
|
||||
start_time_key: Final = f"provider_budget_start_time:{custom_llm_provider}"
|
||||
|
|
@ -516,6 +525,7 @@ class RouterBudgetLimiting(CustomLogger):
|
|||
spend_key=spend_key,
|
||||
start_time_key=start_time_key,
|
||||
response_cost=response_cost,
|
||||
batch_id=batch_id,
|
||||
)
|
||||
|
||||
deployment_budget_config: Final = self._get_budget_config_for_deployment(model_id)
|
||||
|
|
@ -528,6 +538,7 @@ class RouterBudgetLimiting(CustomLogger):
|
|||
spend_key=deployment_spend_key,
|
||||
start_time_key=deployment_start_time_key,
|
||||
response_cost=response_cost,
|
||||
batch_id=batch_id,
|
||||
)
|
||||
|
||||
request_tags: Final = _get_tags_from_request_kwargs(
|
||||
|
|
@ -545,6 +556,7 @@ class RouterBudgetLimiting(CustomLogger):
|
|||
spend_key=_tag_spend_key,
|
||||
start_time_key=_tag_start_time_key,
|
||||
response_cost=response_cost,
|
||||
batch_id=batch_id,
|
||||
)
|
||||
|
||||
async def _increment_spend_for_key(
|
||||
|
|
@ -553,12 +565,15 @@ class RouterBudgetLimiting(CustomLogger):
|
|||
spend_key: str,
|
||||
start_time_key: str,
|
||||
response_cost: float,
|
||||
batch_id: str | None = None,
|
||||
):
|
||||
if budget_config.budget_duration is None:
|
||||
return
|
||||
|
||||
current_time: Final = datetime.now(timezone.utc).timestamp()
|
||||
ttl_seconds: Final = duration_in_seconds(budget_config.budget_duration)
|
||||
if batch_id is not None and not await self._claim_batch_charge(spend_key, batch_id, ttl_seconds):
|
||||
return
|
||||
|
||||
budget_start = await self._get_or_set_budget_start_time(
|
||||
start_time_key=start_time_key,
|
||||
|
|
@ -596,6 +611,17 @@ class RouterBudgetLimiting(CustomLogger):
|
|||
|
||||
verbose_router_logger.debug("Incremented spend for %s by %s", spend_key, response_cost)
|
||||
|
||||
async def _claim_batch_charge(self, spend_key: str, batch_id: str, ttl_seconds: int) -> bool:
|
||||
marker_key: Final = batch_charged_once_marker_key(spend_key=spend_key, batch_id=batch_id)
|
||||
redis_cache: Final = self.dual_cache.redis_cache
|
||||
if redis_cache is not None:
|
||||
polls: Final = await redis_cache.async_increment(key=marker_key, value=1, ttl=ttl_seconds, refresh_ttl=True)
|
||||
return polls == 1
|
||||
local_polls: Final = await self.dual_cache.in_memory_cache.async_increment(
|
||||
key=marker_key, value=1, ttl=ttl_seconds, refresh_ttl=True
|
||||
)
|
||||
return local_polls == 1
|
||||
|
||||
async def periodic_sync_in_memory_spend_with_redis(self):
|
||||
"""
|
||||
Handler that triggers sync_in_memory_spend_with_redis every DEFAULT_REDIS_SYNC_INTERVAL seconds
|
||||
|
|
|
|||
|
|
@ -1,5 +1,6 @@
|
|||
#### What this does ####
|
||||
# picks based on response time (for streaming, this is time to first token)
|
||||
from collections.abc import Mapping
|
||||
from datetime import datetime
|
||||
from typing import Final
|
||||
|
||||
|
|
@ -9,6 +10,30 @@ from litellm._logging import verbose_router_logger
|
|||
from litellm.caching.caching import DualCache
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.router_utils.batch_utils import is_batch_retrieve_call_type
|
||||
from litellm.router_utils.deployment_pricing import deployment_pricing_models, str_keyed
|
||||
|
||||
|
||||
def _deployment_cost_map_rates(deployment: object) -> tuple[float | None, float | None]:
|
||||
"""Input and output per-token prices from the first cost-map row among the deployment's
|
||||
pricing models that carries one; (None, None) when no row does, so the caller falls back to
|
||||
its unmapped default."""
|
||||
rows: Final = (str_keyed(_cost_map().get(model)) for model in deployment_pricing_models(deployment))
|
||||
row: Final = next((row for row in rows if _declares_token_price(row)), None)
|
||||
if row is None:
|
||||
return None, None
|
||||
return _as_rate(row.get("input_cost_per_token")), _as_rate(row.get("output_cost_per_token"))
|
||||
|
||||
|
||||
def _cost_map() -> Mapping[str, object]:
|
||||
return litellm.model_cost
|
||||
|
||||
|
||||
def _declares_token_price(row: Mapping[str, object]) -> bool:
|
||||
return row.get("input_cost_per_token") is not None or row.get("output_cost_per_token") is not None
|
||||
|
||||
|
||||
def _as_rate(value: object) -> float | None:
|
||||
return float(value) if isinstance(value, (int, float)) and not isinstance(value, bool) else None
|
||||
|
||||
|
||||
class LowestCostLoggingHandler(CustomLogger):
|
||||
|
|
@ -239,8 +264,8 @@ class LowestCostLoggingHandler(CustomLogger):
|
|||
or _deployment.get("model_info", {}).get("rpm", None)
|
||||
or float("inf")
|
||||
)
|
||||
item_litellm_model_name = _deployment.get("litellm_params", {}).get("model")
|
||||
item_litellm_model_cost_map = litellm.model_cost.get(item_litellm_model_name, {})
|
||||
deployment_view: object = _deployment
|
||||
map_input_cost, map_output_cost = _deployment_cost_map_rates(deployment_view)
|
||||
|
||||
# check if user provided input_cost_per_token and output_cost_per_token in litellm_params
|
||||
item_input_cost = None
|
||||
|
|
@ -252,10 +277,10 @@ class LowestCostLoggingHandler(CustomLogger):
|
|||
item_output_cost = _deployment.get("litellm_params", {}).get("output_cost_per_token")
|
||||
|
||||
if item_input_cost is None:
|
||||
item_input_cost = item_litellm_model_cost_map.get("input_cost_per_token", 5.0)
|
||||
item_input_cost = map_input_cost if map_input_cost is not None else 5.0
|
||||
|
||||
if item_output_cost is None:
|
||||
item_output_cost = item_litellm_model_cost_map.get("output_cost_per_token", 5.0)
|
||||
item_output_cost = map_output_cost if map_output_cost is not None else 5.0
|
||||
|
||||
# if litellm["model"] is not in model_cost map -> use item_cost = $10
|
||||
|
||||
|
|
|
|||
|
|
@ -3,6 +3,8 @@ import json
|
|||
from os import PathLike
|
||||
from typing import Final
|
||||
|
||||
from openai.types import Batch
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.types.llms.openai import FileTypes, OpenAIFilesPurpose
|
||||
from litellm.types.utils import CallTypes
|
||||
|
|
@ -189,3 +191,14 @@ def is_batch_retrieve_call_type(call_type: object) -> bool:
|
|||
routing strategies' own state, must not be fed from it.
|
||||
"""
|
||||
return isinstance(call_type, str) and call_type in BATCH_RETRIEVE_CALL_TYPES
|
||||
|
||||
|
||||
def batch_charged_once_marker_key(spend_key: str, batch_id: str) -> str:
|
||||
return f"{spend_key}:batch:{batch_id}"
|
||||
|
||||
|
||||
def batch_id_to_charge_once(call_type: object, response_obj: object, response_cost: float) -> str | None:
|
||||
"""A finished batch reports its whole cost on every poll, so its id is charged once per counter."""
|
||||
if response_cost <= 0 or not is_batch_retrieve_call_type(call_type):
|
||||
return None
|
||||
return response_obj.id if isinstance(response_obj, Batch) else None
|
||||
|
|
|
|||
37
litellm/router_utils/deployment_pricing.py
Normal file
37
litellm/router_utils/deployment_pricing.py
Normal file
|
|
@ -0,0 +1,37 @@
|
|||
from collections.abc import Iterator, Mapping
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
|
||||
from pydantic import TypeAdapter, ValidationError
|
||||
|
||||
_EMPTY: Final[Mapping[str, object]] = MappingProxyType({})
|
||||
_STR_KEYED: Final = TypeAdapter(dict[str, object])
|
||||
|
||||
|
||||
def pricing_models(base_model: object, model: object) -> tuple[str, ...]:
|
||||
"""Cost-map keys to price a deployment under, ``model_info.base_model`` first, then the
|
||||
deployment's own model, blanks and duplicates dropped. Every spend and routing surface that
|
||||
prices a deployment walks this order, so a ``base_model`` the cost map does not know falls
|
||||
through to the deployment model instead of billing or ranking at the unmapped default."""
|
||||
return tuple(dict.fromkeys(_named_models(base_model, model)))
|
||||
|
||||
|
||||
def deployment_pricing_models(deployment: object) -> tuple[str, ...]:
|
||||
"""``pricing_models`` read off a raw deployment dict (``model_info.base_model``, then
|
||||
``litellm_params.base_model``, then ``litellm_params.model``); malformed blocks count as unset."""
|
||||
view: Final = str_keyed(deployment)
|
||||
params: Final = str_keyed(view.get("litellm_params"))
|
||||
info: Final = str_keyed(view.get("model_info"))
|
||||
return pricing_models(info.get("base_model") or params.get("base_model"), params.get("model"))
|
||||
|
||||
|
||||
def _named_models(*candidates: object) -> Iterator[str]:
|
||||
return (candidate for candidate in candidates if isinstance(candidate, str) and candidate)
|
||||
|
||||
|
||||
def str_keyed(value: object) -> Mapping[str, object]:
|
||||
"""``value`` as a read-only str-keyed mapping, empty when it is not a dict."""
|
||||
try:
|
||||
return MappingProxyType(_STR_KEYED.validate_python(value))
|
||||
except ValidationError:
|
||||
return _EMPTY
|
||||
|
|
@ -1,20 +1,17 @@
|
|||
"""Resolve which reasoning_effort values a deployment, and by intersection a model group, accepts.
|
||||
|
||||
An entry that states its levels outright in reasoning_effort_levels is read first and wins
|
||||
whole, for a model whose set the per-level flags cannot express: Kimi K3 takes low, high and max,
|
||||
and no flag can drop medium because medium has none. Every other entry answers through the
|
||||
supports_*_reasoning_effort flags below, whose polarity mirrors how a request path reads that same
|
||||
flag. medium and high are unconditional for a reasoning model. minimal and low are opt-out:
|
||||
openai/chat/gpt_5_transformation.py refuses them only when the map says false. xhigh and max are
|
||||
opt-in. none is opt-out everywhere except the azure gpt-5 family, whose config raises
|
||||
UnsupportedParamsError without an explicit true.
|
||||
An entry that states its levels outright in reasoning_effort_levels wins whole, for a model whose set
|
||||
the per-level flags cannot express (Kimi K3 takes low, high and max, and no flag can drop medium).
|
||||
Every other entry answers through its supports_*_reasoning_effort flags, read with the polarity the
|
||||
chat gate uses: medium and high are unconditional, minimal and low are opt-out (refused only on an
|
||||
explicit false), xhigh is opt-in, and none is opt-out except on the azure gpt-5 family, whose chat
|
||||
config refuses it without an explicit true.
|
||||
|
||||
xhigh is gated on the request path by the openai and azure gpt-5 configs. max is not gated there at
|
||||
all: outside the gpt-6-astra rows every entry carrying supports_max_reasoning_effort is Claude-family,
|
||||
and anthropic/chat/transformation.py gates max on the output_config path while its reasoning_effort
|
||||
path maps any level to a thinking budget. Making max opt-in is a deliberate trade, then, since an
|
||||
explicit flag is the only signal that the tier is a real one rather than litellm rounding the level
|
||||
to a budget, and a missing flag costs advisory metadata rather than a rejected request.
|
||||
The Responses gate reads none, minimal, low and xhigh as opt-out. Neither gate judges max, which
|
||||
Chat Completions refuses and Responses honors on the gpt-6 family, so its flag is advisory and
|
||||
advertised opt-in: outside the gpt-6 rows every entry carrying supports_max_reasoning_effort is
|
||||
Claude-family, where anthropic/chat/transformation.py gates max on the output_config path, so an
|
||||
explicit flag is the only signal that the tier is real.
|
||||
|
||||
A deployment the map describes with no effort flags at all resolves to None rather than to the
|
||||
opt-out defaults. 689 of the map's 854 reasoning entries carry no flag, and the o-series, xai and
|
||||
|
|
|
|||
|
|
@ -8,6 +8,7 @@ from litellm.decisions.call import (
|
|||
SystemOneQuestions,
|
||||
SystemOneRequestFields,
|
||||
asend,
|
||||
log_caller_request,
|
||||
prepare_call,
|
||||
reject_other_format,
|
||||
send,
|
||||
|
|
@ -29,10 +30,12 @@ def _prepare_systemone_call(
|
|||
extra_headers: Mapping[str, str] | None,
|
||||
kwargs: Mapping[str, object],
|
||||
) -> DecisionsCall:
|
||||
fields: Final = SystemOneRequestFields(state=state, questions=questions)
|
||||
log_caller_request(kwargs, fields)
|
||||
reject_other_format(model=model, kwargs=kwargs, field_name="input", other_function="decisions")
|
||||
return prepare_call(
|
||||
model=model,
|
||||
fields=SystemOneRequestFields(state=state, questions=questions),
|
||||
fields=fields,
|
||||
safety_identifier=safety_identifier,
|
||||
api_key=api_key,
|
||||
api_base=api_base,
|
||||
|
|
|
|||
|
|
@ -22,7 +22,7 @@ class CacheControlMessageInjectionPoint(TypedDict):
|
|||
|
||||
|
||||
class CacheControlToolConfigInjectionPoint(TypedDict):
|
||||
"""Type for tool_config-level injection points (Bedrock)."""
|
||||
"""Type for tool_config-level injection points (Anthropic and Bedrock)."""
|
||||
|
||||
location: Literal["tool_config"]
|
||||
control: ChatCompletionCachedContent | None
|
||||
|
|
|
|||
|
|
@ -72,6 +72,8 @@ class ProviderConnection:
|
|||
request_timeout: float | None = None
|
||||
force_timeout: float | None = None
|
||||
stream_timeout: float | str | None = None
|
||||
timeout_by_service_tier: Mapping[str, float] | None = None
|
||||
stream_timeout_by_service_tier: Mapping[str, float] | None = None
|
||||
max_retries: int | None = None
|
||||
tenant_id: str | None = None
|
||||
client_id: str | None = None
|
||||
|
|
|
|||
|
|
@ -389,6 +389,9 @@ class GenericLiteLLMParams(CredentialLiteLLMParams, CustomPricingLiteLLMParams):
|
|||
otpm: int | None = None
|
||||
timeout: float | str | httpx.Timeout | None = None # if str, pass in as os.environ/
|
||||
stream_timeout: float | str | None = None # timeout when making stream=True calls, if str, pass in as os.environ/
|
||||
# per service_tier overrides of timeout / stream_timeout, used when the request's service_tier matches a key
|
||||
timeout_by_service_tier: Mapping[str, float] | None = None
|
||||
stream_timeout_by_service_tier: Mapping[str, float] | None = None
|
||||
max_retries: int | None = None
|
||||
drop_params: bool | str | None = None
|
||||
organization: str | None = None # for openai orgs
|
||||
|
|
@ -573,6 +576,8 @@ class LiteLLMParamsTypedDict(TypedDict, total=False):
|
|||
api_version: str | None
|
||||
timeout: float | str | httpx.Timeout | None
|
||||
stream_timeout: float | str | None
|
||||
timeout_by_service_tier: ReadOnly[Mapping[str, float] | None]
|
||||
stream_timeout_by_service_tier: ReadOnly[Mapping[str, float] | None]
|
||||
max_retries: int | None
|
||||
organization: list | str | None # for openai orgs
|
||||
provider_affinity_header: ReadOnly[str | None]
|
||||
|
|
@ -996,6 +1001,8 @@ class RouterModelGroupAliasItem(TypedDict):
|
|||
class RetryAttemptRecord(TypedDict):
|
||||
model_group: ReadOnly[str | None]
|
||||
deployment_id: ReadOnly[str | None]
|
||||
model: ReadOnly[str | None]
|
||||
base_model: ReadOnly[str | None]
|
||||
exception_type: ReadOnly[str]
|
||||
exception_string: ReadOnly[str]
|
||||
attempted_retries: ReadOnly[int | None]
|
||||
|
|
|
|||
|
|
@ -223,6 +223,12 @@ class SearchContextCostPerQuery(TypedDict, total=False):
|
|||
search_context_size_high: float
|
||||
|
||||
|
||||
class SearchFreeRequests(TypedDict):
|
||||
count: ReadOnly[int]
|
||||
period: ReadOnly[Literal["day", "month"]]
|
||||
shared_group: ReadOnly[str]
|
||||
|
||||
|
||||
class AgenticLoopParams(TypedDict, total=False):
|
||||
"""
|
||||
Parameters passed to agentic loop hooks (e.g., WebSearch interception).
|
||||
|
|
@ -284,6 +290,8 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False):
|
|||
input_cost_per_token_ultrafast: ReadOnly[float | None] # OpenAI ultrafast service tier pricing
|
||||
cache_creation_input_token_cost: float | None
|
||||
cache_creation_input_token_cost_above_200k_tokens: float | None
|
||||
cache_creation_input_token_cost_above_100k_tokens: ReadOnly[float | None]
|
||||
cache_creation_input_token_cost_above_1hr_above_100k_tokens: ReadOnly[float | None]
|
||||
cache_creation_input_token_cost_above_272k_tokens: float | None
|
||||
cache_creation_input_token_cost_above_272k_tokens_priority: float | None
|
||||
cache_creation_input_token_cost_above_272k_tokens_flex: float | None
|
||||
|
|
@ -298,15 +306,18 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False):
|
|||
cache_read_input_token_cost_priority: float | None # OpenAI priority service tier pricing
|
||||
cache_read_input_token_cost_ultrafast: ReadOnly[float | None] # OpenAI ultrafast service tier pricing
|
||||
cache_read_input_token_cost_above_200k_tokens: float | None
|
||||
cache_read_input_token_cost_above_100k_tokens: ReadOnly[float | None]
|
||||
cache_read_input_token_cost_above_200k_tokens_priority: float | None
|
||||
cache_read_input_token_cost_above_272k_tokens: float | None
|
||||
cache_read_input_token_cost_above_272k_tokens_priority: float | None
|
||||
cache_read_input_token_cost_above_272k_tokens_flex: float | None
|
||||
cache_read_input_token_cost_above_512k_tokens: float | None
|
||||
cache_read_input_token_cost_batches: ReadOnly[float | None]
|
||||
cache_read_input_token_cost_above_100k_tokens_batches: ReadOnly[float | None]
|
||||
cache_read_input_token_cost_above_272k_tokens_batches: ReadOnly[float | None]
|
||||
cache_creation_input_token_cost_batches: ReadOnly[float | None]
|
||||
cache_creation_input_token_cost_above_200k_tokens_batches: ReadOnly[float | None]
|
||||
cache_creation_input_token_cost_above_100k_tokens_batches: ReadOnly[float | None]
|
||||
cache_creation_input_token_cost_above_272k_tokens_batches: ReadOnly[float | None]
|
||||
# Smallest prefix this model will actually cache, whatever caching mechanism its provider uses.
|
||||
# Absent means the provider-agnostic default applies; see MINIMUM_PROMPT_CACHE_TOKEN_COUNT.
|
||||
|
|
@ -316,6 +327,7 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False):
|
|||
input_cost_per_audio_token: float | None
|
||||
input_cost_per_token_above_128k_tokens: float | None # only for vertex ai models
|
||||
input_cost_per_token_above_200k_tokens: float | None # only for vertex ai gemini-2.5-pro models
|
||||
input_cost_per_token_above_100k_tokens: ReadOnly[float | None]
|
||||
input_cost_per_token_above_200k_tokens_priority: float | None
|
||||
input_cost_per_token_above_272k_tokens: float | None # GPT-5.4/5.4-pro: prompts >272K priced at 2x input
|
||||
input_cost_per_token_above_272k_tokens_priority: float | None
|
||||
|
|
@ -333,8 +345,10 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False):
|
|||
input_cost_per_second: float | None # for OpenAI Speech models
|
||||
input_cost_per_token_batches: float | None
|
||||
input_cost_per_video_token_batches: ReadOnly[float | None]
|
||||
input_cost_per_token_above_100k_tokens_batches: ReadOnly[float | None]
|
||||
input_cost_per_token_above_272k_tokens_batches: ReadOnly[float | None]
|
||||
output_cost_per_token_batches: float | None
|
||||
output_cost_per_token_above_100k_tokens_batches: ReadOnly[float | None]
|
||||
output_cost_per_token_above_272k_tokens_batches: ReadOnly[float | None]
|
||||
output_cost_per_token: Required[float | None]
|
||||
output_cost_per_token_flex: float | None # OpenAI flex service tier pricing
|
||||
|
|
@ -353,6 +367,7 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False):
|
|||
output_cost_per_audio_token: float | None
|
||||
output_cost_per_token_above_128k_tokens: float | None # only for vertex ai models
|
||||
output_cost_per_token_above_200k_tokens: float | None # only for vertex ai gemini-2.5-pro models
|
||||
output_cost_per_token_above_100k_tokens: ReadOnly[float | None]
|
||||
output_cost_per_token_above_200k_tokens_priority: float | None
|
||||
output_cost_per_token_above_272k_tokens: float | None # GPT-5.4/5.4-pro: prompts >272K priced at 1.5x output
|
||||
output_cost_per_token_above_272k_tokens_priority: float | None
|
||||
|
|
@ -390,6 +405,7 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False):
|
|||
web_search_billing_unit: (
|
||||
Literal["per_query", "per_prompt"] | None
|
||||
) # "per_query" (Gemini 3.x) or "per_prompt" (Gemini 2.x)
|
||||
search_free_requests: ReadOnly[SearchFreeRequests | None]
|
||||
google_maps_grounding_cost_per_query: ReadOnly[float | None]
|
||||
citation_cost_per_token: float | None # Cost per citation token for Perplexity
|
||||
tiered_pricing: list[dict[str, Any]] | None # Tiered pricing structure for models like Dashscope
|
||||
|
|
@ -3731,6 +3747,8 @@ class CustomPricingLiteLLMParams(MirroredPricingParams):
|
|||
input_cost_per_token_priority: float | None = None
|
||||
input_cost_per_token_ultrafast: float | None = None
|
||||
cache_creation_input_token_cost_above_1hr: float | None = None
|
||||
cache_creation_input_token_cost_above_100k_tokens: float | None = None
|
||||
cache_creation_input_token_cost_above_1hr_above_100k_tokens: float | None = None
|
||||
cache_creation_input_token_cost_above_200k_tokens: float | None = None
|
||||
cache_creation_input_token_cost_above_272k_tokens: float | None = None
|
||||
cache_creation_input_token_cost_above_272k_tokens_priority: float | None = None
|
||||
|
|
@ -3742,13 +3760,16 @@ class CustomPricingLiteLLMParams(MirroredPricingParams):
|
|||
cache_read_input_token_cost_flex: float | None = None
|
||||
cache_read_input_token_cost_priority: float | None = None
|
||||
cache_read_input_token_cost_ultrafast: float | None = None
|
||||
cache_read_input_token_cost_above_100k_tokens: float | None = None
|
||||
cache_read_input_token_cost_above_200k_tokens: float | None = None
|
||||
cache_read_input_token_cost_above_200k_tokens_priority: float | None = None
|
||||
cache_read_input_token_cost_above_272k_tokens_priority: float | None = None
|
||||
cache_read_input_token_cost_above_272k_tokens_flex: float | None = None
|
||||
cache_read_input_token_cost_batches: float | None = None
|
||||
cache_read_input_token_cost_above_100k_tokens_batches: float | None = None
|
||||
cache_read_input_token_cost_above_272k_tokens_batches: float | None = None
|
||||
cache_creation_input_token_cost_batches: float | None = None
|
||||
cache_creation_input_token_cost_above_100k_tokens_batches: float | None = None
|
||||
cache_creation_input_token_cost_above_200k_tokens_batches: float | None = None
|
||||
cache_creation_input_token_cost_above_272k_tokens_batches: float | None = None
|
||||
cache_read_input_audio_token_cost: float | None = None
|
||||
|
|
@ -3757,10 +3778,12 @@ class CustomPricingLiteLLMParams(MirroredPricingParams):
|
|||
input_cost_per_audio_token: float | None = None
|
||||
input_cost_per_token_cache_hit: float | None = None
|
||||
input_cost_per_token_above_128k_tokens: float | None = None
|
||||
input_cost_per_token_above_100k_tokens: float | None = None
|
||||
input_cost_per_token_above_200k_tokens: float | None = None
|
||||
input_cost_per_token_above_200k_tokens_priority: float | None = None
|
||||
input_cost_per_token_above_272k_tokens_priority: float | None = None
|
||||
input_cost_per_token_above_272k_tokens_flex: float | None = None
|
||||
input_cost_per_token_above_100k_tokens_batches: float | None = None
|
||||
input_cost_per_token_above_272k_tokens_batches: float | None = None
|
||||
input_cost_per_query: float | None = None
|
||||
input_cost_per_image: float | None = None
|
||||
|
|
@ -3781,10 +3804,12 @@ class CustomPricingLiteLLMParams(MirroredPricingParams):
|
|||
output_cost_per_token_ultrafast: float | None = None
|
||||
output_cost_per_audio_token: float | None = None
|
||||
output_cost_per_token_above_128k_tokens: float | None = None
|
||||
output_cost_per_token_above_100k_tokens: float | None = None
|
||||
output_cost_per_token_above_200k_tokens: float | None = None
|
||||
output_cost_per_token_above_200k_tokens_priority: float | None = None
|
||||
output_cost_per_token_above_272k_tokens_priority: float | None = None
|
||||
output_cost_per_token_above_272k_tokens_flex: float | None = None
|
||||
output_cost_per_token_above_100k_tokens_batches: float | None = None
|
||||
output_cost_per_token_above_272k_tokens_batches: float | None = None
|
||||
output_cost_per_character_above_128k_tokens: float | None = None
|
||||
output_cost_per_image: float | None = None
|
||||
|
|
@ -3851,7 +3876,7 @@ def shared_backend_model_info(model_info: dict[str, Any]) -> dict[str, Any]:
|
|||
return {k: v for k, v in model_info.items() if k in SHARED_BACKEND_MODEL_INFO_FIELDS}
|
||||
|
||||
|
||||
ABOVE_THRESHOLD_COST_KEY_PATTERN: Final = re.compile(r"_above_\d+k?_tokens$")
|
||||
ABOVE_THRESHOLD_COST_KEY_PATTERN: Final = re.compile(r"_above_\d+k?_tokens(?:_batches)?$")
|
||||
|
||||
_PRICING_FIELD_EXEMPTIONS: Final[frozenset[str]] = frozenset({"output_vector_size"})
|
||||
|
||||
|
|
|
|||
150
litellm/utils.py
150
litellm/utils.py
|
|
@ -45,7 +45,7 @@ from httpx import Proxy
|
|||
from httpx._utils import get_environment_proxies
|
||||
from openai.lib import _parsing, _pydantic
|
||||
from openai.types.chat.completion_create_params import ResponseFormat
|
||||
from pydantic import BaseModel
|
||||
from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError, with_config
|
||||
|
||||
import litellm
|
||||
import litellm.litellm_core_utils
|
||||
|
|
@ -954,10 +954,6 @@ def _ocr_document_summary(document: object) -> str:
|
|||
|
||||
# Runs once per call to check if the user wants to send their data anywhere - PostHog/Sentry/Slack/etc.
|
||||
|
||||
_DECISION_CALL_TYPES: Final[frozenset[str]] = frozenset(
|
||||
(CallTypes.decisions.value, CallTypes.adecisions.value, CallTypes.systemone.value, CallTypes.asystemone.value)
|
||||
)
|
||||
|
||||
|
||||
def function_setup(
|
||||
original_function: str,
|
||||
|
|
@ -1207,9 +1203,6 @@ def function_setup(
|
|||
if isinstance(search_query, list)
|
||||
else search_query
|
||||
)
|
||||
elif call_type in _DECISION_CALL_TYPES:
|
||||
decisions_state: Final = args[1] if len(args) > 1 else kwargs.get("state") or kwargs.get("input") or ""
|
||||
messages = decisions_state if isinstance(decisions_state, str) else json.dumps(decisions_state, default=str)
|
||||
elif call_type in (CallTypes.image_edit.value, CallTypes.aimage_edit.value):
|
||||
messages = args[1] if len(args) > 1 else kwargs.get("prompt")
|
||||
elif call_type in (CallTypes.ocr.value, CallTypes.aocr.value):
|
||||
|
|
@ -1375,6 +1368,7 @@ def _schedule_async_success_logging(
|
|||
return
|
||||
|
||||
def _enqueue_async_logging() -> None:
|
||||
logging_obj.mark_async_success_log_enqueued()
|
||||
asyncio.create_task(
|
||||
_client_async_logging_helper(
|
||||
logging_obj=logging_obj,
|
||||
|
|
@ -3243,6 +3237,129 @@ def _get_builtin_model_info_for_registration(model: str) -> ModelInfo | None:
|
|||
|
||||
|
||||
_runtime_registered_model_cost: Final[dict[str, dict[str, object]]] = {} # mutable-ok: replayed on reload
|
||||
_runtime_registered_fields: Final[dict[str, frozenset[str]]] = {} # mutable-ok: fields each registration set
|
||||
_runtime_registered_fields_owner: list[int] = [0] # mutable-ok: id of the model_cost map the fields describe
|
||||
|
||||
|
||||
def _registered_fields_for_current_map() -> dict[str, frozenset[str]]:
|
||||
if _runtime_registered_fields_owner[0] != id(litellm.model_cost):
|
||||
_runtime_registered_fields.clear()
|
||||
_runtime_registered_fields_owner[0] = id(litellm.model_cost)
|
||||
return _runtime_registered_fields
|
||||
|
||||
|
||||
_RATE_TABLE: Final = TypeAdapter(Mapping[str, object])
|
||||
_RATE_TIERS: Final = TypeAdapter(tuple[Mapping[str, object], ...])
|
||||
_TOKEN_COST_FIELDS: Final = ("input_cost_per_token", "output_cost_per_token")
|
||||
_DIRECTION_RATE_FIELDS: Final = MappingProxyType(
|
||||
{
|
||||
"input_cost_per_token": ("input_cost_per_token",),
|
||||
"output_cost_per_token": ("output_cost_per_token", "output_cost_per_reasoning_token"),
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _is_overridden_rate(name: str, overridden_fields: tuple[str, ...]) -> bool:
|
||||
is_variant: Final = name.startswith(tuple(f"{rate_field}_" for rate_field in overridden_fields))
|
||||
return name in overridden_fields or (is_variant and "_cache_" not in name)
|
||||
|
||||
|
||||
def _without_overridden_rates(rates: Mapping[str, object], overridden_fields: tuple[str, ...]) -> Mapping[str, object]:
|
||||
return MappingProxyType(
|
||||
{name: rate for name, rate in rates.items() if not _is_overridden_rate(name, overridden_fields)}
|
||||
)
|
||||
|
||||
|
||||
def _rate_table_with_flat_prices(
|
||||
name: str, table: object, overridden_fields: tuple[str, ...], flat_prices: Mapping[str, object]
|
||||
) -> object:
|
||||
try:
|
||||
if name == "tiered_pricing":
|
||||
return [
|
||||
{**_without_overridden_rates(tier, overridden_fields), **flat_prices}
|
||||
for tier in _RATE_TIERS.validate_python(table)
|
||||
]
|
||||
if name == "off_peak_pricing":
|
||||
return dict(_without_overridden_rates(_RATE_TABLE.validate_python(table), overridden_fields))
|
||||
except ValidationError:
|
||||
return table
|
||||
return table
|
||||
|
||||
|
||||
def _without_overridden_price_variants(
|
||||
entry: Mapping[str, object], value: Mapping[str, object]
|
||||
) -> Mapping[str, object]:
|
||||
flat_prices: Final = MappingProxyType(
|
||||
{cost_field: value[cost_field] for cost_field in _TOKEN_COST_FIELDS if value.get(cost_field) is not None}
|
||||
)
|
||||
overridden_fields: Final = tuple(
|
||||
itertools.chain.from_iterable(_DIRECTION_RATE_FIELDS[cost_field] for cost_field in flat_prices)
|
||||
)
|
||||
if not overridden_fields:
|
||||
return entry
|
||||
return MappingProxyType(
|
||||
{
|
||||
name: _rate_table_with_flat_prices(name, field_value, overridden_fields, flat_prices)
|
||||
for name, field_value in entry.items()
|
||||
if not _is_overridden_rate(name, overridden_fields)
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
@with_config(ConfigDict(extra="allow"))
|
||||
class _PricingModelInfo(ModelInfo):
|
||||
pass
|
||||
|
||||
|
||||
_PRICING_MODEL_INFO: Final = TypeAdapter(_PricingModelInfo)
|
||||
|
||||
|
||||
def model_info_with_pricing_overrides(model_info: ModelInfo, registration: Mapping[str, object]) -> ModelInfo:
|
||||
from litellm.types.utils import is_server_derived_pricing_key
|
||||
|
||||
overrides: Final = {
|
||||
name: value for name, value in registration.items() if value is not None and is_server_derived_pricing_key(name)
|
||||
}
|
||||
return _PRICING_MODEL_INFO.validate_python(
|
||||
_update_dictionary(dict(_without_overridden_price_variants(model_info, overrides)), overrides)
|
||||
)
|
||||
|
||||
|
||||
def runtime_model_cost_registrations(
|
||||
model: str, custom_llm_provider: str
|
||||
) -> tuple[tuple[str, Mapping[str, object]], ...]:
|
||||
names: Final = _get_potential_model_names(model, custom_llm_provider)
|
||||
candidates: Final = (
|
||||
model,
|
||||
names["combined_model_name"],
|
||||
f"{custom_llm_provider}/{names['split_model']}",
|
||||
names["split_model"],
|
||||
names["stripped_model_name"],
|
||||
)
|
||||
registered: Final = (
|
||||
(key, litellm.model_cost.get(key), _registered_fields_for_current_map().get(key, frozenset()))
|
||||
for key in dict.fromkeys(candidates)
|
||||
)
|
||||
return tuple(
|
||||
(key, MappingProxyType({name: value for name, value in entry.items() if name in fields}))
|
||||
for key, entry, fields in registered
|
||||
if isinstance(entry, Mapping) and fields
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class ModelInfoLookupMiss:
|
||||
tripped: bool
|
||||
|
||||
|
||||
def lookup_model_info_for_registration(model: str, custom_llm_provider: str | None) -> ModelInfo | ModelInfoLookupMiss:
|
||||
try:
|
||||
info: Final = get_model_info(model=model, custom_llm_provider=custom_llm_provider)
|
||||
except ModelNotMappedError:
|
||||
return ModelInfoLookupMiss(tripped=False)
|
||||
except Exception:
|
||||
return ModelInfoLookupMiss(tripped=True)
|
||||
return ModelInfoLookupMiss(tripped=False) if is_generalized_model_info(info) else info
|
||||
|
||||
|
||||
class _LiveDeploymentReplay:
|
||||
|
|
@ -3284,6 +3401,7 @@ def reapply_runtime_model_cost_registrations() -> None:
|
|||
have no such owner are recorded and replayed, and a registration describing
|
||||
a single request opts out of even that.
|
||||
"""
|
||||
_runtime_registered_fields.clear()
|
||||
if _LiveDeploymentReplay.callback is not None:
|
||||
_LiveDeploymentReplay.callback()
|
||||
if _runtime_registered_model_cost:
|
||||
|
|
@ -3347,7 +3465,10 @@ def register_model(
|
|||
if persist_across_reloads:
|
||||
_registrations: Final[Mapping[str, Mapping[str, object]]] = loaded_model_cost
|
||||
for _registered_key, _registered_value in _registrations.items():
|
||||
_runtime_registered_model_cost[_registered_key] = dict(_registered_value) # mutable-ok: caller-owned
|
||||
_runtime_registered_model_cost[_registered_key] = _update_dictionary(
|
||||
dict(_runtime_registered_model_cost.get(_registered_key, {})),
|
||||
dict(_RATE_TABLE.validate_python(_registered_value)),
|
||||
)
|
||||
|
||||
_skip_get_model_info_providers: Final = PROVIDERS_THAT_AUTHENTICATE_ON_PROVIDER_INFO
|
||||
|
||||
|
|
@ -3413,6 +3534,10 @@ def register_model(
|
|||
## override / add new keys to the existing model cost dictionary
|
||||
updated_dictionary = _update_dictionary(existing_model, value)
|
||||
litellm.model_cost.setdefault(model_cost_key, {}).update(updated_dictionary)
|
||||
registered_fields = _registered_fields_for_current_map()
|
||||
registered_fields[str(model_cost_key)] = registered_fields.get(str(model_cost_key), frozenset()) | frozenset(
|
||||
name for name, field_value in value.items() if field_value is not None
|
||||
)
|
||||
|
||||
# Invalidate case-insensitive lookup map since model_cost was modified
|
||||
_invalidate_model_cost_lowercase_map()
|
||||
|
|
@ -5037,6 +5162,9 @@ def get_optional_params(
|
|||
EXTRA_BODY_ROUTING_KEYS: Final = frozenset({"model"})
|
||||
|
||||
|
||||
LITELLM_ONLY_PARAMS: Final = frozenset(("cache_control_injection_points",))
|
||||
|
||||
|
||||
def add_provider_specific_params_to_optional_params(
|
||||
optional_params: dict,
|
||||
passed_params: dict,
|
||||
|
|
@ -5052,8 +5180,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 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"] = {}
|
||||
|
|
@ -6318,6 +6447,7 @@ def _get_model_info_helper(
|
|||
supports_computer_use=_model_info.get("supports_computer_use", None),
|
||||
search_context_cost_per_query=_model_info.get("search_context_cost_per_query", None),
|
||||
web_search_billing_unit=_model_info.get("web_search_billing_unit", None),
|
||||
search_free_requests=_model_info.get("search_free_requests", None),
|
||||
google_maps_grounding_cost_per_query=_model_info.get("google_maps_grounding_cost_per_query", None),
|
||||
tpm=_model_info.get("tpm", None),
|
||||
rpm=_model_info.get("rpm", None),
|
||||
|
|
|
|||
File diff suppressed because it is too large
Load diff
|
|
@ -83,7 +83,12 @@
|
|||
"minimum": 0,
|
||||
"description": "USD per token written to the provider's prompt cache."
|
||||
},
|
||||
"cache_creation_input_token_cost_above_128k_tokens": {
|
||||
"cache_creation_input_token_cost_above_100k_tokens": {
|
||||
"type": "number",
|
||||
"minimum": 0,
|
||||
"description": "Rate applied once the prompt exceeds the token threshold in the field name."
|
||||
},
|
||||
"cache_creation_input_token_cost_above_100k_tokens_batches": {
|
||||
"type": "number",
|
||||
"minimum": 0,
|
||||
"description": "Rate applied once the prompt exceeds the token threshold in the field name."
|
||||
|
|
@ -93,6 +98,11 @@
|
|||
"minimum": 0,
|
||||
"description": "Rate applied once the prompt exceeds the token threshold in the field name."
|
||||
},
|
||||
"cache_creation_input_token_cost_above_1hr_above_100k_tokens": {
|
||||
"type": "number",
|
||||
"minimum": 0,
|
||||
"description": "Rate applied once the prompt exceeds the token threshold in the field name."
|
||||
},
|
||||
"cache_creation_input_token_cost_above_1hr_above_200k_tokens": {
|
||||
"type": "number",
|
||||
"minimum": 0,
|
||||
|
|
@ -103,11 +113,6 @@
|
|||
"minimum": 0,
|
||||
"description": "Rate applied once the prompt exceeds the token threshold in the field name."
|
||||
},
|
||||
"cache_creation_input_token_cost_above_256k_tokens": {
|
||||
"type": "number",
|
||||
"minimum": 0,
|
||||
"description": "Rate applied once the prompt exceeds the token threshold in the field name."
|
||||
},
|
||||
"cache_creation_input_token_cost_above_272k_tokens": {
|
||||
"type": "number",
|
||||
"minimum": 0,
|
||||
|
|
@ -118,6 +123,11 @@
|
|||
"minimum": 0,
|
||||
"description": "Rate applied once the prompt exceeds the token threshold in the field name."
|
||||
},
|
||||
"cache_creation_input_token_cost_above_272k_tokens_decisions": {
|
||||
"type": "number",
|
||||
"minimum": 0,
|
||||
"description": "Rate applied once the prompt exceeds the token threshold in the field name."
|
||||
},
|
||||
"cache_creation_input_token_cost_above_272k_tokens_flex": {
|
||||
"type": "number",
|
||||
"minimum": 0,
|
||||
|
|
@ -128,7 +138,7 @@
|
|||
"minimum": 0,
|
||||
"description": "Priority service-tier rate for the same-named base field."
|
||||
},
|
||||
"cache_creation_input_token_cost_above_32k_tokens": {
|
||||
"cache_creation_input_token_cost_above_272k_tokens_ultrafast": {
|
||||
"type": "number",
|
||||
"minimum": 0,
|
||||
"description": "Rate applied once the prompt exceeds the token threshold in the field name."
|
||||
|
|
@ -137,6 +147,10 @@
|
|||
"type": "number",
|
||||
"minimum": 0
|
||||
},
|
||||
"cache_creation_input_token_cost_decisions": {
|
||||
"type": "number",
|
||||
"minimum": 0
|
||||
},
|
||||
"cache_creation_input_token_cost_flex": {
|
||||
"type": "number",
|
||||
"minimum": 0,
|
||||
|
|
@ -147,6 +161,10 @@
|
|||
"minimum": 0,
|
||||
"description": "Priority service-tier rate for the same-named base field."
|
||||
},
|
||||
"cache_creation_input_token_cost_ultrafast": {
|
||||
"type": "number",
|
||||
"minimum": 0
|
||||
},
|
||||
"cache_read_input_audio_token_cost": {
|
||||
"type": "number",
|
||||
"minimum": 0
|
||||
|
|
@ -160,7 +178,12 @@
|
|||
"minimum": 0,
|
||||
"description": "USD per prompt token served from the provider's prompt cache."
|
||||
},
|
||||
"cache_read_input_token_cost_above_128k_tokens": {
|
||||
"cache_read_input_token_cost_above_100k_tokens": {
|
||||
"type": "number",
|
||||
"minimum": 0,
|
||||
"description": "Rate applied once the prompt exceeds the token threshold in the field name."
|
||||
},
|
||||
"cache_read_input_token_cost_above_100k_tokens_batches": {
|
||||
"type": "number",
|
||||
"minimum": 0,
|
||||
"description": "Rate applied once the prompt exceeds the token threshold in the field name."
|
||||
|
|
@ -170,16 +193,16 @@
|
|||
"minimum": 0,
|
||||
"description": "Rate applied once the prompt exceeds the token threshold in the field name."
|
||||
},
|
||||
"cache_read_input_token_cost_above_200k_tokens_batches": {
|
||||
"type": "number",
|
||||
"minimum": 0,
|
||||
"description": "Rate applied once the prompt exceeds the token threshold in the field name."
|
||||
},
|
||||
"cache_read_input_token_cost_above_200k_tokens_priority": {
|
||||
"type": "number",
|
||||
"minimum": 0,
|
||||
"description": "Priority service-tier rate for the same-named base field."
|
||||
},
|
||||
"cache_read_input_token_cost_above_256k_tokens": {
|
||||
"type": "number",
|
||||
"minimum": 0,
|
||||
"description": "Rate applied once the prompt exceeds the token threshold in the field name."
|
||||
},
|
||||
"cache_read_input_token_cost_above_272k_tokens": {
|
||||
"type": "number",
|
||||
"minimum": 0,
|
||||
|
|
@ -190,6 +213,11 @@
|
|||
"minimum": 0,
|
||||
"description": "Rate applied once the prompt exceeds the token threshold in the field name."
|
||||
},
|
||||
"cache_read_input_token_cost_above_272k_tokens_decisions": {
|
||||
"type": "number",
|
||||
"minimum": 0,
|
||||
"description": "Rate applied once the prompt exceeds the token threshold in the field name."
|
||||
},
|
||||
"cache_read_input_token_cost_above_272k_tokens_flex": {
|
||||
"type": "number",
|
||||
"minimum": 0,
|
||||
|
|
@ -200,7 +228,7 @@
|
|||
"minimum": 0,
|
||||
"description": "Priority service-tier rate for the same-named base field."
|
||||
},
|
||||
"cache_read_input_token_cost_above_32k_tokens": {
|
||||
"cache_read_input_token_cost_above_272k_tokens_ultrafast": {
|
||||
"type": "number",
|
||||
"minimum": 0,
|
||||
"description": "Rate applied once the prompt exceeds the token threshold in the field name."
|
||||
|
|
@ -210,10 +238,18 @@
|
|||
"minimum": 0,
|
||||
"description": "Rate applied once the prompt exceeds the token threshold in the field name."
|
||||
},
|
||||
"cache_read_input_token_cost_balanced": {
|
||||
"type": "number",
|
||||
"minimum": 0
|
||||
},
|
||||
"cache_read_input_token_cost_batches": {
|
||||
"type": "number",
|
||||
"minimum": 0
|
||||
},
|
||||
"cache_read_input_token_cost_decisions": {
|
||||
"type": "number",
|
||||
"minimum": 0
|
||||
},
|
||||
"cache_read_input_token_cost_flex": {
|
||||
"type": "number",
|
||||
"minimum": 0,
|
||||
|
|
@ -224,6 +260,10 @@
|
|||
"minimum": 0,
|
||||
"description": "Priority service-tier rate for the same-named base field."
|
||||
},
|
||||
"cache_read_input_token_cost_ultrafast": {
|
||||
"type": "number",
|
||||
"minimum": 0
|
||||
},
|
||||
"citation_cost_per_token": {
|
||||
"type": "number",
|
||||
"minimum": 0
|
||||
|
|
@ -235,6 +275,10 @@
|
|||
"comment": {
|
||||
"type": "string"
|
||||
},
|
||||
"cost_per_second": {
|
||||
"type": "number",
|
||||
"minimum": 0
|
||||
},
|
||||
"default_reasoning_effort": {
|
||||
"type": "string",
|
||||
"description": "Reasoning effort the provider applies when the request omits reasoning_effort. Gates whether a non-default temperature or the top_p/logprobs sampling params are accepted, which hold only when the effort resolves to 'none'.",
|
||||
|
|
@ -253,6 +297,10 @@
|
|||
"format": "date",
|
||||
"pattern": "^\\d{4}-(0[1-9]|1[0-2])-(0[1-9]|[12]\\d|3[01])$"
|
||||
},
|
||||
"explicit_cache_read_input_token_cost": {
|
||||
"type": "number",
|
||||
"minimum": 0
|
||||
},
|
||||
"gemini_audio_only_live": {
|
||||
"type": "boolean"
|
||||
},
|
||||
|
|
@ -336,11 +384,25 @@
|
|||
"type": "number",
|
||||
"minimum": 0
|
||||
},
|
||||
"input_cost_per_text_message": {
|
||||
"type": "number",
|
||||
"minimum": 0
|
||||
},
|
||||
"input_cost_per_token": {
|
||||
"type": "number",
|
||||
"minimum": 0,
|
||||
"description": "USD per prompt token."
|
||||
},
|
||||
"input_cost_per_token_above_100k_tokens": {
|
||||
"type": "number",
|
||||
"minimum": 0,
|
||||
"description": "Rate applied once the prompt exceeds the token threshold in the field name."
|
||||
},
|
||||
"input_cost_per_token_above_100k_tokens_batches": {
|
||||
"type": "number",
|
||||
"minimum": 0,
|
||||
"description": "Rate applied once the prompt exceeds the token threshold in the field name."
|
||||
},
|
||||
"input_cost_per_token_above_128k_tokens": {
|
||||
"type": "number",
|
||||
"minimum": 0,
|
||||
|
|
@ -351,16 +413,16 @@
|
|||
"minimum": 0,
|
||||
"description": "Rate applied once the prompt exceeds the token threshold in the field name."
|
||||
},
|
||||
"input_cost_per_token_above_200k_tokens_batches": {
|
||||
"type": "number",
|
||||
"minimum": 0,
|
||||
"description": "Rate applied once the prompt exceeds the token threshold in the field name."
|
||||
},
|
||||
"input_cost_per_token_above_200k_tokens_priority": {
|
||||
"type": "number",
|
||||
"minimum": 0,
|
||||
"description": "Priority service-tier rate for the same-named base field."
|
||||
},
|
||||
"input_cost_per_token_above_256k_tokens": {
|
||||
"type": "number",
|
||||
"minimum": 0,
|
||||
"description": "Rate applied once the prompt exceeds the token threshold in the field name."
|
||||
},
|
||||
"input_cost_per_token_above_272k_tokens": {
|
||||
"type": "number",
|
||||
"minimum": 0,
|
||||
|
|
@ -371,6 +433,11 @@
|
|||
"minimum": 0,
|
||||
"description": "Rate applied once the prompt exceeds the token threshold in the field name."
|
||||
},
|
||||
"input_cost_per_token_above_272k_tokens_decisions": {
|
||||
"type": "number",
|
||||
"minimum": 0,
|
||||
"description": "Rate applied once the prompt exceeds the token threshold in the field name."
|
||||
},
|
||||
"input_cost_per_token_above_272k_tokens_flex": {
|
||||
"type": "number",
|
||||
"minimum": 0,
|
||||
|
|
@ -381,7 +448,7 @@
|
|||
"minimum": 0,
|
||||
"description": "Priority service-tier rate for the same-named base field."
|
||||
},
|
||||
"input_cost_per_token_above_32k_tokens": {
|
||||
"input_cost_per_token_above_272k_tokens_ultrafast": {
|
||||
"type": "number",
|
||||
"minimum": 0,
|
||||
"description": "Rate applied once the prompt exceeds the token threshold in the field name."
|
||||
|
|
@ -391,6 +458,10 @@
|
|||
"minimum": 0,
|
||||
"description": "Rate applied once the prompt exceeds the token threshold in the field name."
|
||||
},
|
||||
"input_cost_per_token_balanced": {
|
||||
"type": "number",
|
||||
"minimum": 0
|
||||
},
|
||||
"input_cost_per_token_batches": {
|
||||
"type": "number",
|
||||
"minimum": 0,
|
||||
|
|
@ -400,6 +471,10 @@
|
|||
"type": "number",
|
||||
"minimum": 0
|
||||
},
|
||||
"input_cost_per_token_decisions": {
|
||||
"type": "number",
|
||||
"minimum": 0
|
||||
},
|
||||
"input_cost_per_token_flex": {
|
||||
"type": "number",
|
||||
"minimum": 0,
|
||||
|
|
@ -410,6 +485,10 @@
|
|||
"minimum": 0,
|
||||
"description": "Priority service-tier rate for the same-named base field."
|
||||
},
|
||||
"input_cost_per_token_ultrafast": {
|
||||
"type": "number",
|
||||
"minimum": 0
|
||||
},
|
||||
"input_cost_per_video_per_second": {
|
||||
"type": "number",
|
||||
"minimum": 0
|
||||
|
|
@ -474,6 +553,7 @@
|
|||
"completion",
|
||||
"embedding",
|
||||
"evaluation",
|
||||
"extraction",
|
||||
"guardrail",
|
||||
"image_edit",
|
||||
"image_generation",
|
||||
|
|
@ -656,6 +736,10 @@
|
|||
"type": "number",
|
||||
"minimum": 0
|
||||
},
|
||||
"output_cost_per_image_token_batches": {
|
||||
"type": "number",
|
||||
"minimum": 0
|
||||
},
|
||||
"output_cost_per_pixel": {
|
||||
"type": "number",
|
||||
"minimum": 0
|
||||
|
|
@ -698,6 +782,16 @@
|
|||
"minimum": 0,
|
||||
"description": "USD per generated token."
|
||||
},
|
||||
"output_cost_per_token_above_100k_tokens": {
|
||||
"type": "number",
|
||||
"minimum": 0,
|
||||
"description": "Rate applied once the prompt exceeds the token threshold in the field name."
|
||||
},
|
||||
"output_cost_per_token_above_100k_tokens_batches": {
|
||||
"type": "number",
|
||||
"minimum": 0,
|
||||
"description": "Rate applied once the prompt exceeds the token threshold in the field name."
|
||||
},
|
||||
"output_cost_per_token_above_128k_tokens": {
|
||||
"type": "number",
|
||||
"minimum": 0,
|
||||
|
|
@ -708,16 +802,16 @@
|
|||
"minimum": 0,
|
||||
"description": "Rate applied once the prompt exceeds the token threshold in the field name."
|
||||
},
|
||||
"output_cost_per_token_above_200k_tokens_batches": {
|
||||
"type": "number",
|
||||
"minimum": 0,
|
||||
"description": "Rate applied once the prompt exceeds the token threshold in the field name."
|
||||
},
|
||||
"output_cost_per_token_above_200k_tokens_priority": {
|
||||
"type": "number",
|
||||
"minimum": 0,
|
||||
"description": "Priority service-tier rate for the same-named base field."
|
||||
},
|
||||
"output_cost_per_token_above_256k_tokens": {
|
||||
"type": "number",
|
||||
"minimum": 0,
|
||||
"description": "Rate applied once the prompt exceeds the token threshold in the field name."
|
||||
},
|
||||
"output_cost_per_token_above_272k_tokens": {
|
||||
"type": "number",
|
||||
"minimum": 0,
|
||||
|
|
@ -728,6 +822,11 @@
|
|||
"minimum": 0,
|
||||
"description": "Rate applied once the prompt exceeds the token threshold in the field name."
|
||||
},
|
||||
"output_cost_per_token_above_272k_tokens_decisions": {
|
||||
"type": "number",
|
||||
"minimum": 0,
|
||||
"description": "Rate applied once the prompt exceeds the token threshold in the field name."
|
||||
},
|
||||
"output_cost_per_token_above_272k_tokens_flex": {
|
||||
"type": "number",
|
||||
"minimum": 0,
|
||||
|
|
@ -738,7 +837,7 @@
|
|||
"minimum": 0,
|
||||
"description": "Priority service-tier rate for the same-named base field."
|
||||
},
|
||||
"output_cost_per_token_above_32k_tokens": {
|
||||
"output_cost_per_token_above_272k_tokens_ultrafast": {
|
||||
"type": "number",
|
||||
"minimum": 0,
|
||||
"description": "Rate applied once the prompt exceeds the token threshold in the field name."
|
||||
|
|
@ -748,11 +847,19 @@
|
|||
"minimum": 0,
|
||||
"description": "Rate applied once the prompt exceeds the token threshold in the field name."
|
||||
},
|
||||
"output_cost_per_token_balanced": {
|
||||
"type": "number",
|
||||
"minimum": 0
|
||||
},
|
||||
"output_cost_per_token_batches": {
|
||||
"type": "number",
|
||||
"minimum": 0,
|
||||
"description": "USD per generated token via the provider's batch API."
|
||||
},
|
||||
"output_cost_per_token_decisions": {
|
||||
"type": "number",
|
||||
"minimum": 0
|
||||
},
|
||||
"output_cost_per_token_flex": {
|
||||
"type": "number",
|
||||
"minimum": 0,
|
||||
|
|
@ -763,6 +870,10 @@
|
|||
"minimum": 0,
|
||||
"description": "Priority service-tier rate for the same-named base field."
|
||||
},
|
||||
"output_cost_per_token_ultrafast": {
|
||||
"type": "number",
|
||||
"minimum": 0
|
||||
},
|
||||
"output_cost_per_video_per_second": {
|
||||
"type": "number",
|
||||
"minimum": 0
|
||||
|
|
@ -844,6 +955,33 @@
|
|||
},
|
||||
"additionalProperties": false
|
||||
},
|
||||
"search_free_requests": {
|
||||
"type": "object",
|
||||
"description": "Free Google Search grounding requests per period, shared across models in the same shared_group.",
|
||||
"required": [
|
||||
"count",
|
||||
"period",
|
||||
"shared_group"
|
||||
],
|
||||
"additionalProperties": false,
|
||||
"properties": {
|
||||
"count": {
|
||||
"type": "integer",
|
||||
"minimum": 0
|
||||
},
|
||||
"period": {
|
||||
"type": "string",
|
||||
"enum": [
|
||||
"day",
|
||||
"month"
|
||||
]
|
||||
},
|
||||
"shared_group": {
|
||||
"type": "string",
|
||||
"pattern": "\\S"
|
||||
}
|
||||
}
|
||||
},
|
||||
"source": {
|
||||
"type": "string",
|
||||
"description": "URL of the provider pricing/model page this entry was taken from."
|
||||
|
|
@ -918,6 +1056,24 @@
|
|||
"supports_audio_output": {
|
||||
"type": "boolean"
|
||||
},
|
||||
"supports_bedrock_converse_tool_result_images": {
|
||||
"type": "boolean"
|
||||
},
|
||||
"supports_bedrock_native_structured_output": {
|
||||
"type": "boolean"
|
||||
},
|
||||
"supports_bedrock_runtime_chat_completions_inline_reasoning": {
|
||||
"type": "boolean"
|
||||
},
|
||||
"supports_bedrock_runtime_chat_completions_response_format": {
|
||||
"type": "boolean"
|
||||
},
|
||||
"supports_bedrock_runtime_chat_completions_tools_with_reasoning": {
|
||||
"type": "boolean"
|
||||
},
|
||||
"supports_box_questions": {
|
||||
"type": "boolean"
|
||||
},
|
||||
"supports_computer_use": {
|
||||
"type": "boolean"
|
||||
},
|
||||
|
|
@ -990,6 +1146,9 @@
|
|||
"supports_reasoning": {
|
||||
"type": "boolean"
|
||||
},
|
||||
"supports_regex_lookaround": {
|
||||
"type": "boolean"
|
||||
},
|
||||
"supports_response_schema": {
|
||||
"type": "boolean"
|
||||
},
|
||||
|
|
@ -1088,6 +1247,31 @@
|
|||
"minimum": 0,
|
||||
"description": "Tokens added once per request to a local count estimate, for the provider's chat framing."
|
||||
},
|
||||
"token_count_model": {
|
||||
"type": "object",
|
||||
"required": [
|
||||
"model"
|
||||
],
|
||||
"additionalProperties": false,
|
||||
"properties": {
|
||||
"model": {
|
||||
"type": "string",
|
||||
"pattern": "\\S"
|
||||
},
|
||||
"region": {
|
||||
"type": "string",
|
||||
"pattern": "\\S"
|
||||
},
|
||||
"litellm_provider": {
|
||||
"type": "string",
|
||||
"enum": [
|
||||
"bedrock",
|
||||
"bedrock_mantle"
|
||||
]
|
||||
}
|
||||
},
|
||||
"description": "Model that counts tokens for this one, its counting provider, and the AWS region its count API runs in."
|
||||
},
|
||||
"token_count_tokenizer": {
|
||||
"type": "string",
|
||||
"description": "Tokenizer litellm counts this model's requests with when the provider cannot count them.",
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
[project]
|
||||
name = "litellm"
|
||||
version = "1.104.5"
|
||||
version = "1.104.6"
|
||||
description = "Library to easily interface with LLM API providers"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.10, <3.15"
|
||||
|
|
@ -76,7 +76,7 @@ proxy = [
|
|||
"httpx2>=2.5.0,<3",
|
||||
"pydantic>=2.12.0,<3",
|
||||
"litellm-proxy-extras==0.4.102.post1",
|
||||
"litellm-enterprise==0.1.71",
|
||||
"litellm-enterprise==0.1.71.post1",
|
||||
"RestrictedPython>=8.5,<9.0",
|
||||
"rich>=13.9.4,<14.0",
|
||||
"InquirerPy>=0.3.4,<1.0",
|
||||
|
|
@ -286,7 +286,7 @@ ci = [
|
|||
"langchain-openai==1.1.14",
|
||||
"langgraph>=1.2.4,<1.3.0",
|
||||
"langgraph-prebuilt>=1.1.0,<1.3.0",
|
||||
"claude-agent-sdk==0.1.44",
|
||||
"claude-agent-sdk==0.2.165",
|
||||
"google-cloud-speech==2.40.0",
|
||||
]
|
||||
healthcheck = [
|
||||
|
|
@ -355,7 +355,7 @@ litellm-enterprise = { workspace = true }
|
|||
members = ["enterprise", "litellm-proxy-extras"]
|
||||
|
||||
[tool.commitizen]
|
||||
version = "1.104.5"
|
||||
version = "1.104.6"
|
||||
version_files = [
|
||||
"pyproject.toml:^version",
|
||||
]
|
||||
|
|
|
|||
|
|
@ -87,6 +87,7 @@ ignored_function_names = [
|
|||
"_resolve_claude_code_session_router", # Tested through Claude Code session routing in test_router.py
|
||||
"_get_claude_code_session_router_binding", # Tested through the two-worker session routing test in test_router.py
|
||||
"_apply_updated_routing_strategy_args", # Tested via update_settings in test_lowest_latency.py (file lacks "router" in name)
|
||||
"_deployment_params_base_model", # Tested through log_retry in test_router_helper_utils.py
|
||||
]
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -0,0 +1,61 @@
|
|||
"""Live e2e: `/v1/messages/count_tokens` on a Bedrock Claude model that bedrock-runtime
|
||||
cannot count.
|
||||
|
||||
Claude Opus 4.8 is offered only through cross-region inference, and bedrock-runtime's
|
||||
CountTokens answers 400 for it. The proxy then has to count through bedrock-mantle's
|
||||
Anthropic count_tokens, and the answer must sit within a few percent of what `/v1/messages`
|
||||
bills as `usage.input_tokens`. The local tokenizer fallback undercounts these models by
|
||||
about 40%, so this is the line that proves the real count is served. Both calls go through
|
||||
the real Anthropic SDK, as a customer would use them.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
from anthropic.types import MessageParam
|
||||
from e2e_config import unique_marker
|
||||
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
|
||||
|
||||
CROSS_REGION_ONLY_CLAUDE_BACKEND: Final = "bedrock/global.anthropic.claude-opus-4-8"
|
||||
COUNT_TOLERANCE: Final = 0.05
|
||||
|
||||
|
||||
def _register(proxy: ProxyClient, resources: ResourceManager, backend: str) -> tuple[str, str]:
|
||||
model: Final = f"e2e-count-tokens-bedrock-{unique_marker()}"
|
||||
model_id: Final = proxy.create_model(
|
||||
model,
|
||||
LiteLLMParamsBody(
|
||||
model=backend,
|
||||
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",
|
||||
),
|
||||
)
|
||||
resources.defer(lambda: proxy.delete_model(model_id))
|
||||
return model, resources.key()
|
||||
|
||||
|
||||
class TestBedrockMessagesCountTokens:
|
||||
def test_count_matches_billed_input_tokens_for_a_model_bedrock_runtime_cannot_count(
|
||||
self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients
|
||||
) -> None:
|
||||
model, key = _register(proxy, resources, CROSS_REGION_ONLY_CLAUDE_BACKEND)
|
||||
client: Final = sdk.anthropic(key)
|
||||
prompt: Final = f"{unique_marker()} " + "The quick brown fox jumps over the lazy dog. " * 40
|
||||
message: Final[MessageParam] = {"role": "user", "content": prompt}
|
||||
|
||||
counted: Final = client.messages.count_tokens(model=model, messages=[message])
|
||||
answered: Final = client.messages.create(
|
||||
model=model, max_tokens=1, messages=[message], extra_body=NO_PROXY_CACHE
|
||||
)
|
||||
|
||||
billed: Final = answered.usage.input_tokens
|
||||
assert billed, answered.usage
|
||||
assert abs(counted.input_tokens - billed) <= billed * COUNT_TOLERANCE, (counted, answered.usage)
|
||||
|
|
@ -0,0 +1,349 @@
|
|||
import asyncio
|
||||
import json
|
||||
import re
|
||||
import threading
|
||||
import uuid
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from queue import SimpleQueue
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
from urllib.parse import urlsplit
|
||||
|
||||
import httpx
|
||||
import psutil
|
||||
import pytest
|
||||
import yaml
|
||||
from integration._support.client import Gateway, eventually
|
||||
from integration._support.database import read_rows
|
||||
from integration._support.process import owned_proxy_process
|
||||
from integration._support.wire import Reply, Request, Wire, wire_server
|
||||
from integration.providers.test_anthropic_safeguards_beta_wire import (
|
||||
ANTHROPIC_KEY,
|
||||
ANTHROPIC_MODEL,
|
||||
BETA,
|
||||
OMITTED,
|
||||
anthropic_peer,
|
||||
marker_of,
|
||||
messages_body,
|
||||
stream_frames,
|
||||
)
|
||||
from pydantic import JsonValue, TypeAdapter
|
||||
|
||||
_CONFIG_MODEL: Final = "anthropic-safeguards-chaos"
|
||||
_FOUNDRY_CONFIG_MODEL: Final = "azure-ai-safeguards-chaos"
|
||||
_YAML_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
|
||||
_MARKER: Final = re.compile(r"marker-([0-9a-f]{32})")
|
||||
_STARTED_WORKER: Final = re.compile(r"Started server process \[(\d+)\]")
|
||||
_MODEL_LIST: Final = json.dumps(
|
||||
{"object": "list", "data": [{"id": ANTHROPIC_MODEL, "object": "model", "owned_by": "anthropic"}]}
|
||||
).encode()
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _Call:
|
||||
stream: bool
|
||||
marker: str
|
||||
safeguards: bool = True
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _Served:
|
||||
call: _Call
|
||||
status: int
|
||||
text: str
|
||||
|
||||
|
||||
def _calls(count: int, stream: bool, safeguards: bool = True) -> tuple[_Call, ...]:
|
||||
return tuple(_Call(stream=stream, marker=uuid.uuid4().hex, safeguards=safeguards) for _ in range(count))
|
||||
|
||||
|
||||
def _streamed_reply(marker: str, abort_after: int | None = None, pause: float = 0) -> Reply:
|
||||
return Reply(
|
||||
content_type="text/event-stream",
|
||||
chunks=stream_frames(marker),
|
||||
abort_after=abort_after,
|
||||
pause_between_chunks=pause,
|
||||
)
|
||||
|
||||
|
||||
async def _send(client: httpx.AsyncClient, key: str, model: str, call: _Call) -> _Served:
|
||||
async with client.stream(
|
||||
"POST",
|
||||
"/v1/messages",
|
||||
json=(
|
||||
messages_body(model, call.marker, stream=call.stream)
|
||||
if call.safeguards
|
||||
else messages_body(model, call.marker, stream=call.stream, safeguards=OMITTED)
|
||||
),
|
||||
headers={"Authorization": f"Bearer {key}", "anthropic-version": "2023-06-01"},
|
||||
) as response:
|
||||
raw: Final = await response.aread()
|
||||
return _Served(call=call, status=response.status_code, text=raw.decode())
|
||||
|
||||
|
||||
async def _burst(
|
||||
base_url: str, key: str, model: str, calls: tuple[_Call, ...], *, tolerate_transport_errors: bool = False
|
||||
) -> tuple[_Served, ...]:
|
||||
async with httpx.AsyncClient(base_url=base_url, timeout=60, trust_env=False) as client:
|
||||
results: Final = await asyncio.gather(
|
||||
*(_send(client, key, model, call) for call in calls), return_exceptions=tolerate_transport_errors
|
||||
)
|
||||
for result in results:
|
||||
assert not isinstance(result, BaseException) or isinstance(result, httpx.TransportError), repr(result)
|
||||
return tuple(result for result in results if isinstance(result, _Served))
|
||||
|
||||
|
||||
def _assert_answered_with_its_own_marker(served: _Served) -> None:
|
||||
assert served.status == 200, served.text
|
||||
assert set(_MARKER.findall(served.text)) == {served.call.marker}, served.text
|
||||
assert f"msg_{served.call.marker}" in served.text, served.text
|
||||
if served.call.stream:
|
||||
assert "event: message_stop" in served.text, served.text
|
||||
|
||||
|
||||
def _assert_every_request_carried_the_beta(received: tuple[Request, ...], markers: frozenset[str]) -> None:
|
||||
assert sorted(marker_of(request) for request in received) == sorted(markers)
|
||||
missing: Final = [
|
||||
frozenset(request.headers) for request in received if request.headers.get("anthropic-beta") != BETA
|
||||
]
|
||||
assert missing == [], missing
|
||||
|
||||
|
||||
def _spend_success_ids(model: str, expected: frozenset[str]) -> None:
|
||||
rows: Final = eventually(
|
||||
lambda: read_rows(
|
||||
'SELECT request_id FROM "LiteLLM_SpendLogs" WHERE model_group=%s AND status=%s', (model, "success")
|
||||
),
|
||||
lambda found: expected <= {str(row["request_id"]) for row in found},
|
||||
seconds=70,
|
||||
)
|
||||
identities: Final = [row["request_id"] for row in rows]
|
||||
assert len(set(identities)) == len(identities), identities
|
||||
|
||||
|
||||
async def test_burst_with_aborted_streams_keeps_the_beta_on_every_forwarded_request(gateway: Gateway) -> None:
|
||||
calls: Final = tuple(_Call(stream=index % 2 == 0, marker=uuid.uuid4().hex) for index in range(30))
|
||||
aborted: Final = frozenset(call.marker for index, call in enumerate(calls) if call.stream and index % 3 == 0)
|
||||
|
||||
def respond(request: Request) -> Reply:
|
||||
marker: Final = marker_of(request)
|
||||
if marker in aborted:
|
||||
return _streamed_reply(marker, abort_after=0)
|
||||
return anthropic_peer(request)
|
||||
|
||||
with wire_server(respond) as wire, gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(model=f"anthropic/{ANTHROPIC_MODEL}", api_base=wire.url, api_key=ANTHROPIC_KEY)
|
||||
served: Final = await _burst(str(gateway.client.base_url), gateway.key, model, calls)
|
||||
assert len(served) == 30
|
||||
for item in served:
|
||||
if item.call.marker in aborted:
|
||||
assert item.status == 500, item.text
|
||||
assert '"type":"api_error"' in item.text and "marker-" not in item.text, item.text
|
||||
else:
|
||||
_assert_answered_with_its_own_marker(item)
|
||||
recovery: Final = _Call(stream=True, marker=uuid.uuid4().hex)
|
||||
(recovered,) = await _burst(str(gateway.client.base_url), gateway.key, model, (recovery,))
|
||||
_assert_answered_with_its_own_marker(recovered)
|
||||
_assert_every_request_carried_the_beta(wire.drain(), frozenset(call.marker for call in (*calls, recovery)))
|
||||
answered: Final = frozenset(f"msg_{call.marker}" for call in (*calls, recovery) if call.marker not in aborted)
|
||||
_spend_success_ids(model, answered)
|
||||
|
||||
|
||||
async def test_slow_streams_complete_with_the_beta_on_every_request(gateway: Gateway) -> None:
|
||||
calls: Final = _calls(10, stream=True)
|
||||
with (
|
||||
wire_server(lambda request: _streamed_reply(marker_of(request), pause=0.3)) as wire,
|
||||
gateway.scenario() as scenario,
|
||||
):
|
||||
model: Final = scenario.model(model=f"anthropic/{ANTHROPIC_MODEL}", api_base=wire.url, api_key=ANTHROPIC_KEY)
|
||||
served: Final = await _burst(str(gateway.client.base_url), gateway.key, model, calls)
|
||||
assert len(served) == 10
|
||||
for item in served:
|
||||
_assert_answered_with_its_own_marker(item)
|
||||
_assert_every_request_carried_the_beta(wire.drain(), frozenset(call.marker for call in calls))
|
||||
_spend_success_ids(model, frozenset(f"msg_{call.marker}" for call in calls))
|
||||
|
||||
|
||||
def _anthropic_deployment(wire: Wire) -> dict[str, JsonValue]:
|
||||
return {
|
||||
"model_name": _CONFIG_MODEL,
|
||||
"litellm_params": {"model": f"anthropic/{ANTHROPIC_MODEL}", "api_base": wire.url, "api_key": ANTHROPIC_KEY},
|
||||
}
|
||||
|
||||
|
||||
def _foundry_deployment(wire: Wire) -> dict[str, JsonValue]:
|
||||
return {
|
||||
"model_name": _FOUNDRY_CONFIG_MODEL,
|
||||
"litellm_params": {
|
||||
"model": "azure_ai/claude-sonnet-4-6",
|
||||
"api_base": wire.url,
|
||||
"api_key": "synthetic-foundry-key",
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def _chaos_config(tmp_path: Path, *deployments: dict[str, JsonValue]) -> Path:
|
||||
base: Final = _YAML_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()))
|
||||
config: Final = {**base, "model_list": list(deployments)}
|
||||
path: Final = tmp_path / "anthropic-safeguards-chaos.yaml"
|
||||
path.write_text(yaml.safe_dump(config))
|
||||
return path
|
||||
|
||||
|
||||
def _open_upstream_connections(pid: int, upstream: str) -> int:
|
||||
port: Final = urlsplit(upstream).port
|
||||
return sum(
|
||||
1
|
||||
for connection in psutil.Process(pid).net_connections(kind="tcp")
|
||||
if connection.status == psutil.CONN_ESTABLISHED and connection.raddr and connection.raddr.port == port
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.skip(reason="Needs OwnedProxy.replacing_worker, which this stable line's integration harness does not have")
|
||||
@pytest.mark.timeout(180)
|
||||
async def test_worker_sigkill_mid_burst_leaves_the_sibling_sending_the_beta(gateway: Gateway, tmp_path: Path) -> None:
|
||||
calls: Final = _calls(20, stream=False)
|
||||
release: Final = threading.Event()
|
||||
held_markers: Final[SimpleQueue[str]] = SimpleQueue()
|
||||
|
||||
def held(request: Request) -> Reply:
|
||||
if (request.method, request.target) == ("GET", "/v1/models"):
|
||||
return Reply(body=_MODEL_LIST)
|
||||
held_markers.put(marker_of(request))
|
||||
assert release.wait(timeout=60), "The burst was never released"
|
||||
return anthropic_peer(request)
|
||||
|
||||
with wire_server(held) as wire:
|
||||
path: Final = _chaos_config(tmp_path, _anthropic_deployment(wire))
|
||||
with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned:
|
||||
candidate: Final = owned.gateway
|
||||
workers: Final = eventually(
|
||||
lambda: tuple(int(found.group(1)) for found in _STARTED_WORKER.finditer(owned.log.read_text())),
|
||||
lambda pids: len(pids) == 2,
|
||||
seconds=30,
|
||||
)
|
||||
burst: Final = asyncio.create_task(
|
||||
_burst(
|
||||
str(candidate.client.base_url), candidate.key, _CONFIG_MODEL, calls, tolerate_transport_errors=True
|
||||
)
|
||||
)
|
||||
await asyncio.to_thread(eventually, held_markers.qsize, lambda size: size == 20, 60)
|
||||
held_by: Final = MappingProxyType({pid: _open_upstream_connections(pid, wire.url) for pid in workers})
|
||||
assert sum(held_by.values()) == 20, held_by
|
||||
victim_pid, survivor_pid = sorted(workers, key=held_by.__getitem__)
|
||||
with owned.replacing_worker(victim_pid) as kill:
|
||||
kill.freeze()
|
||||
await asyncio.to_thread(kill)
|
||||
release.set()
|
||||
served: Final = await burst
|
||||
assert held_by[survivor_pid] >= 10, held_by
|
||||
assert len(served) == held_by[survivor_pid], (held_by, len(served))
|
||||
for item in served:
|
||||
_assert_answered_with_its_own_marker(item)
|
||||
follow_up: Final = _Call(stream=False, marker=uuid.uuid4().hex)
|
||||
(answered,) = await _burst(str(candidate.client.base_url), candidate.key, _CONFIG_MODEL, (follow_up,))
|
||||
_assert_answered_with_its_own_marker(answered)
|
||||
received: Final = wire.drain()
|
||||
posts: Final = tuple(request for request in received if request.method == "POST")
|
||||
probes: Final[list[tuple[str, str]]] = [
|
||||
(request.method, request.target) for request in received if request.method != "POST"
|
||||
]
|
||||
assert set(probes) <= {("GET", "/v1/models")}, probes
|
||||
_assert_every_request_carried_the_beta(posts, frozenset(call.marker for call in (*calls, follow_up)))
|
||||
|
||||
|
||||
def _beta_map_gets(received: tuple[Request, ...]) -> int:
|
||||
return sum(1 for request in received if (request.method, request.target) == ("GET", "/beta.json"))
|
||||
|
||||
|
||||
async def _chat(base_url: str, key: str, marker: str) -> httpx.Response:
|
||||
async with httpx.AsyncClient(base_url=base_url, timeout=60, trust_env=False) as client:
|
||||
return await client.post(
|
||||
"/v1/chat/completions",
|
||||
json={
|
||||
"model": _CONFIG_MODEL,
|
||||
"max_tokens": 16,
|
||||
"messages": [{"role": "user", "content": f"Question marker-{marker}"}],
|
||||
},
|
||||
headers={"Authorization": f"Bearer {key}"},
|
||||
)
|
||||
|
||||
|
||||
async def _answered_while_the_beta_map_is_held(
|
||||
base_url: str, key: str, guarded: tuple[_Call, ...], chat_markers: tuple[str, ...]
|
||||
) -> None:
|
||||
served: Final = await _burst(base_url, key, _CONFIG_MODEL, guarded)
|
||||
chats: Final = await asyncio.gather(*(_chat(base_url, key, marker) for marker in chat_markers))
|
||||
for item in served:
|
||||
_assert_answered_with_its_own_marker(item)
|
||||
for marker, chat in zip(chat_markers, chats, strict=True):
|
||||
assert chat.status_code == 200, chat.text
|
||||
assert f"marker-{marker}" in chat.text, chat.text
|
||||
|
||||
|
||||
@pytest.mark.timeout(240)
|
||||
async def test_cold_workers_fetch_the_beta_map_only_for_filtered_safeguards_requests_off_the_shared_threads(
|
||||
gateway: Gateway, tmp_path: Path
|
||||
) -> None:
|
||||
beta_map: Final = Path("litellm/anthropic_beta_headers_config.json").read_bytes()
|
||||
release: Final = threading.Event()
|
||||
|
||||
def held_beta_map(_: Request) -> Reply:
|
||||
assert release.wait(timeout=180), "The beta map was never released"
|
||||
return Reply(body=beta_map)
|
||||
|
||||
def upstream(request: Request) -> Reply:
|
||||
if (request.method, request.target) == ("GET", "/v1/models"):
|
||||
return Reply(body=_MODEL_LIST)
|
||||
return anthropic_peer(request)
|
||||
|
||||
plain: Final = _calls(8, stream=False, safeguards=False)
|
||||
plain_foundry: Final = _calls(8, stream=False, safeguards=False)
|
||||
direct: Final = _calls(8, stream=False)
|
||||
foundry: Final = _calls(80, stream=False)
|
||||
direct_while_held: Final = _calls(8, stream=False)
|
||||
chat_markers: Final = tuple(uuid.uuid4().hex for _ in range(4))
|
||||
with wire_server(held_beta_map) as beta_map_wire, wire_server(upstream) as wire:
|
||||
path: Final = _chaos_config(tmp_path, _anthropic_deployment(wire), _foundry_deployment(wire))
|
||||
with owned_proxy_process(
|
||||
gateway,
|
||||
tmp_path,
|
||||
{"LITELLM_ANTHROPIC_BETA_HEADERS_URL": f"{beta_map_wire.url}/beta.json"},
|
||||
remove_environment=("LITELLM_LOCAL_ANTHROPIC_BETA_HEADERS",),
|
||||
config=path,
|
||||
workers=2,
|
||||
) as owned:
|
||||
base_url: Final = str(owned.gateway.client.base_url)
|
||||
key: Final = owned.gateway.key
|
||||
cold: Final = (
|
||||
*await _burst(base_url, key, _CONFIG_MODEL, plain),
|
||||
*await _burst(base_url, key, _FOUNDRY_CONFIG_MODEL, plain_foundry),
|
||||
*await _burst(base_url, key, _CONFIG_MODEL, direct),
|
||||
)
|
||||
for served_cold in cold:
|
||||
_assert_answered_with_its_own_marker(served_cold)
|
||||
assert beta_map_wire.received.qsize() == 0, "a request that never filters betas fetched the beta map"
|
||||
|
||||
held: Final = asyncio.create_task(_burst(base_url, key, _FOUNDRY_CONFIG_MODEL, foundry))
|
||||
try:
|
||||
await asyncio.to_thread(eventually, beta_map_wire.received.qsize, lambda size: size >= 1, 60)
|
||||
await asyncio.wait_for(
|
||||
_answered_while_the_beta_map_is_held(base_url, key, direct_while_held, chat_markers), timeout=45
|
||||
)
|
||||
assert not held.done(), "Foundry safeguards requests were answered before the beta map loaded"
|
||||
finally:
|
||||
release.set()
|
||||
for served_foundry in await held:
|
||||
_assert_answered_with_its_own_marker(served_foundry)
|
||||
|
||||
received: Final = tuple(request for request in wire.drain() if request.method == "POST")
|
||||
beta_by_marker: Final = MappingProxyType(
|
||||
{marker_of(request): request.headers.get("anthropic-beta") for request in received}
|
||||
)
|
||||
assert {beta_by_marker[call.marker] for call in (*direct, *direct_while_held, *foundry)} == {BETA}, (
|
||||
beta_by_marker
|
||||
)
|
||||
assert {beta_by_marker[call.marker] for call in (*plain, *plain_foundry)} == {None}, beta_by_marker
|
||||
fetches: Final = _beta_map_gets(beta_map_wire.drain())
|
||||
assert 1 <= fetches <= 2, fetches
|
||||
|
|
@ -0,0 +1,809 @@
|
|||
import itertools
|
||||
import json
|
||||
import re
|
||||
import time
|
||||
import uuid
|
||||
from collections.abc import Callable, Mapping
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
|
||||
import anthropic
|
||||
import httpx
|
||||
import openai
|
||||
import pytest
|
||||
from integration._support.client import Gateway, eventually
|
||||
from integration._support.database import read_rows
|
||||
from integration._support.process import owned_proxy
|
||||
from integration._support.vertex import service_account_json
|
||||
from integration._support.wire import Reply, Request, Wire, wire_server
|
||||
from pydantic import JsonValue, TypeAdapter
|
||||
|
||||
BETA: Final = "dangerous-tool-use-2026-09-03"
|
||||
SAFEGUARDS: Final[list[JsonValue]] = [
|
||||
{"type": "dangerous_tool_use", "classifier_context": {"v": 1, "permission_mode": "auto"}}
|
||||
]
|
||||
ANTHROPIC_MODEL: Final = "claude-sonnet-4-5-20250929"
|
||||
ANTHROPIC_KEY: Final = "synthetic-anthropic-key"
|
||||
_CLAUDE_CODE_BETAS: Final = (
|
||||
"claude-code-20250219",
|
||||
"interleaved-thinking-2025-05-14",
|
||||
"fine-grained-tool-streaming-2025-05-14",
|
||||
"context-management-2025-06-27",
|
||||
)
|
||||
_CLIENT_BETAS: Final = ",".join(_CLAUDE_CODE_BETAS)
|
||||
_CLIENT_BETAS_WITH_DANGEROUS: Final = ",".join((*_CLAUDE_CODE_BETAS, BETA))
|
||||
_MERGED_BETAS: Final = ",".join(sorted((*_CLAUDE_CODE_BETAS, BETA)))
|
||||
_FOUNDRY_MODEL: Final = "azure_ai/claude-sonnet-4-6"
|
||||
_FOUNDRY_KEY: Final = "synthetic-foundry-key"
|
||||
_FOUNDRY_TARGET: Final = "/anthropic/v1/messages"
|
||||
_FOUNDRY_ARRAY_ERROR: Final = "safeguards: Input should be a valid array"
|
||||
_FOUNDRY_EXTRA_INPUT_ERROR: Final = "safeguards: Extra inputs are not permitted"
|
||||
_ANTHROPIC_VERSION: Final = MappingProxyType({"anthropic-version": "2023-06-01"})
|
||||
_MESSAGES_KEYS: Final = frozenset({"max_tokens", "messages", "model", "stream"})
|
||||
_RESPONSES_KEYS: Final = frozenset({"include", "input", "max_output_tokens", "model"})
|
||||
_SDK_EXTRA: Final[dict[str, JsonValue]] = {"safeguards": SAFEGUARDS, "cache": {"no-cache": True}, "num_retries": 0}
|
||||
_MARKER: Final = re.compile(r"marker-([0-9a-f]{32})")
|
||||
_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _Omitted:
|
||||
pass
|
||||
|
||||
|
||||
OMITTED: Final = _Omitted()
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _Observed:
|
||||
method: str
|
||||
target: str
|
||||
marker: str
|
||||
headers: Mapping[str, str]
|
||||
body: Mapping[str, JsonValue]
|
||||
|
||||
@property
|
||||
def beta(self) -> str | None:
|
||||
return self.headers.get("anthropic-beta")
|
||||
|
||||
@property
|
||||
def keys(self) -> frozenset[str]:
|
||||
return frozenset(self.body)
|
||||
|
||||
@property
|
||||
def header_names(self) -> frozenset[str]:
|
||||
return frozenset(self.headers)
|
||||
|
||||
|
||||
def marker_of(request: Request) -> str:
|
||||
found: Final = _MARKER.search(request.body.decode())
|
||||
assert found is not None, request.body
|
||||
return found.group(1)
|
||||
|
||||
|
||||
def _message(marker: str) -> dict[str, JsonValue]:
|
||||
return {
|
||||
"id": f"msg_{marker}",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"model": ANTHROPIC_MODEL,
|
||||
"content": [{"type": "text", "text": f"answer marker-{marker}"}],
|
||||
"stop_reason": "end_turn",
|
||||
"stop_sequence": None,
|
||||
"usage": {"input_tokens": 12, "output_tokens": 5},
|
||||
}
|
||||
|
||||
|
||||
def stream_frames(marker: str) -> tuple[bytes, ...]:
|
||||
opening: Final[dict[str, JsonValue]] = {
|
||||
**_message(marker),
|
||||
"content": [],
|
||||
"stop_reason": None,
|
||||
"usage": {"input_tokens": 12, "output_tokens": 0},
|
||||
}
|
||||
events: Final[tuple[tuple[str, dict[str, JsonValue]], ...]] = (
|
||||
("message_start", {"message": opening}),
|
||||
("content_block_start", {"index": 0, "content_block": {"type": "text", "text": ""}}),
|
||||
("content_block_delta", {"index": 0, "delta": {"type": "text_delta", "text": f"answer marker-{marker}"}}),
|
||||
("content_block_stop", {"index": 0}),
|
||||
("message_delta", {"delta": {"stop_reason": "end_turn", "stop_sequence": None}, "usage": {"output_tokens": 5}}),
|
||||
("message_stop", {}),
|
||||
)
|
||||
return tuple(
|
||||
f"event: {kind}\ndata: {json.dumps({'type': kind, **payload})}\n\n".encode() for kind, payload in events
|
||||
)
|
||||
|
||||
|
||||
def anthropic_peer(request: Request) -> Reply:
|
||||
marker: Final = marker_of(request)
|
||||
if _JSON_OBJECT.validate_json(request.body).get("stream") is True:
|
||||
return Reply(content_type="text/event-stream", chunks=stream_frames(marker))
|
||||
return Reply(body=json.dumps(_message(marker)).encode())
|
||||
|
||||
|
||||
def _responses_peer(request: Request) -> Reply:
|
||||
marker: Final = marker_of(request)
|
||||
return Reply(
|
||||
body=json.dumps(
|
||||
{
|
||||
"id": f"resp_{marker}",
|
||||
"object": "response",
|
||||
"created_at": 1,
|
||||
"status": "completed",
|
||||
"model": "gpt-4o-mini",
|
||||
"output": [
|
||||
{
|
||||
"type": "message",
|
||||
"id": f"msg_{marker}",
|
||||
"status": "completed",
|
||||
"role": "assistant",
|
||||
"content": [{"type": "output_text", "text": f"answer marker-{marker}", "annotations": []}],
|
||||
}
|
||||
],
|
||||
"usage": {"input_tokens": 12, "output_tokens": 5, "total_tokens": 17},
|
||||
}
|
||||
).encode()
|
||||
)
|
||||
|
||||
|
||||
def _error_peer(status: int, kind: str, message: str) -> Callable[[Request], Reply]:
|
||||
body: Final = json.dumps({"type": "error", "error": {"type": kind, "message": message}}).encode()
|
||||
return lambda request: Reply(status=status, body=body)
|
||||
|
||||
|
||||
def messages_body(
|
||||
model: str,
|
||||
marker: str,
|
||||
*,
|
||||
safeguards: JsonValue | _Omitted = SAFEGUARDS,
|
||||
stream: bool = False,
|
||||
no_cache: bool = True,
|
||||
) -> dict[str, JsonValue]:
|
||||
base: Final[dict[str, JsonValue]] = {
|
||||
"model": model,
|
||||
"max_tokens": 16,
|
||||
"messages": [{"role": "user", "content": f"Question marker-{marker}"}],
|
||||
"num_retries": 0,
|
||||
**({"stream": True} if stream else {}),
|
||||
**({"cache": {"no-cache": True}} if no_cache else {}),
|
||||
}
|
||||
if isinstance(safeguards, _Omitted):
|
||||
return base
|
||||
return {**base, "safeguards": safeguards}
|
||||
|
||||
|
||||
def _observe(request: Request) -> _Observed:
|
||||
return _Observed(
|
||||
request.method, request.target, marker_of(request), request.headers, _JSON_OBJECT.validate_json(request.body)
|
||||
)
|
||||
|
||||
|
||||
def _only(wire: Wire) -> _Observed:
|
||||
(request,) = wire.drain()
|
||||
return _observe(request)
|
||||
|
||||
|
||||
def _post(gateway: Gateway, body: Mapping[str, JsonValue], *, beta: str | None = None) -> httpx.Response:
|
||||
headers: Final = {**_ANTHROPIC_VERSION, **({"anthropic-beta": beta} if beta else {})}
|
||||
return gateway.request("POST", "/v1/messages", body, headers=headers)
|
||||
|
||||
|
||||
def _assert_answered(response: httpx.Response, marker: str) -> None:
|
||||
assert response.status_code == 200, response.text
|
||||
assert response.json()["id"] == f"msg_{marker}", response.text
|
||||
|
||||
|
||||
def _assert_streamed(response: httpx.Response, marker: str) -> None:
|
||||
assert response.status_code == 200, response.text
|
||||
assert f"msg_{marker}" in response.text and f"answer marker-{marker}" in response.text, response.text
|
||||
assert "event: message_stop" in response.text, response.text
|
||||
|
||||
|
||||
def _spend_row(identity: str) -> None:
|
||||
rows: Final = eventually(
|
||||
lambda: read_rows('SELECT request_id, status FROM "LiteLLM_SpendLogs" WHERE request_id=%s', (identity,)),
|
||||
lambda found: len(found) == 1,
|
||||
seconds=70,
|
||||
)
|
||||
assert rows[0]["status"] == "success", rows
|
||||
|
||||
|
||||
def _cache_hit_spend_row(identity: str) -> None:
|
||||
rows: Final = eventually(
|
||||
lambda: read_rows(
|
||||
'SELECT request_id, status, cache_hit FROM "LiteLLM_SpendLogs" WHERE request_id LIKE %s',
|
||||
(f"{identity}\\_cache\\_hit%",),
|
||||
),
|
||||
lambda found: len(found) == 1,
|
||||
seconds=70,
|
||||
)
|
||||
assert (rows[0]["status"], str(rows[0]["cache_hit"])) == ("success", "True"), rows
|
||||
|
||||
|
||||
def _responses_spend_row(model: str) -> None:
|
||||
rows: Final = eventually(
|
||||
lambda: read_rows(
|
||||
'SELECT request_id, status FROM "LiteLLM_SpendLogs" WHERE model_group = %s AND call_type = %s',
|
||||
(model, "aresponses"),
|
||||
),
|
||||
lambda found: len(found) == 1,
|
||||
seconds=70,
|
||||
)
|
||||
assert rows[0]["status"] == "success", rows
|
||||
|
||||
|
||||
def _expect_messages(
|
||||
observed: _Observed,
|
||||
*,
|
||||
target: str,
|
||||
auth: tuple[str, str] | None,
|
||||
beta: str | None,
|
||||
safeguards: JsonValue | _Omitted = SAFEGUARDS,
|
||||
keys: frozenset[str] = _MESSAGES_KEYS,
|
||||
) -> None:
|
||||
assert (observed.method, observed.target) == ("POST", target)
|
||||
if auth is not None:
|
||||
credential_matches: Final = observed.headers.get(auth[0]) == auth[1]
|
||||
assert credential_matches, f"{auth[0]} differs from the deployment credential; headers {observed.header_names}"
|
||||
assert observed.beta == beta, observed.header_names
|
||||
if isinstance(safeguards, _Omitted):
|
||||
assert observed.keys == keys
|
||||
return
|
||||
assert observed.keys == keys | {"safeguards"}
|
||||
assert observed.body["safeguards"] == safeguards
|
||||
|
||||
|
||||
def _expect_anthropic_version(observed: _Observed) -> None:
|
||||
assert observed.headers.get("anthropic-version") == "2023-06-01", observed.header_names
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _ShapeCase:
|
||||
sent: JsonValue | _Omitted
|
||||
forwarded: JsonValue | _Omitted
|
||||
client_beta: str | None
|
||||
expected_beta: str | None
|
||||
stream: bool = False
|
||||
|
||||
|
||||
_ANTHROPIC_SHAPES: Final = MappingProxyType(
|
||||
{
|
||||
"D1_safeguards_without_client_header": _ShapeCase(SAFEGUARDS, SAFEGUARDS, None, BETA),
|
||||
"D2_safeguards_with_claude_code_header": _ShapeCase(
|
||||
SAFEGUARDS, SAFEGUARDS, _CLIENT_BETAS_WITH_DANGEROUS, _MERGED_BETAS
|
||||
),
|
||||
"D3_safeguards_with_claude_code_header_lacking_the_beta": _ShapeCase(
|
||||
SAFEGUARDS, SAFEGUARDS, _CLIENT_BETAS, _MERGED_BETAS
|
||||
),
|
||||
"D4_empty_safeguards": _ShapeCase([], [], None, BETA),
|
||||
"D5_null_safeguards": _ShapeCase(None, OMITTED, None, None),
|
||||
"D6_missing_safeguards": _ShapeCase(OMITTED, OMITTED, None, None),
|
||||
}
|
||||
)
|
||||
_FOUNDRY_SHAPES: Final = MappingProxyType(
|
||||
{
|
||||
"D9_streamed_safeguards_with_claude_code_header": _ShapeCase(
|
||||
SAFEGUARDS, SAFEGUARDS, _CLIENT_BETAS_WITH_DANGEROUS, BETA, stream=True
|
||||
),
|
||||
"D11_empty_safeguards": _ShapeCase([], [], None, BETA),
|
||||
"D12_null_safeguards": _ShapeCase(None, OMITTED, None, None),
|
||||
"D13_missing_safeguards_with_claude_code_header_lacking_the_beta": _ShapeCase(
|
||||
OMITTED, OMITTED, _CLIENT_BETAS, None
|
||||
),
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("row", tuple(_ANTHROPIC_SHAPES))
|
||||
def test_anthropic_messages_derive_the_dangerous_tool_use_beta_from_safeguards(gateway: Gateway, row: str) -> None:
|
||||
case: Final = _ANTHROPIC_SHAPES[row]
|
||||
marker: Final = uuid.uuid4().hex
|
||||
with wire_server(anthropic_peer) as wire, gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(model=f"anthropic/{ANTHROPIC_MODEL}", api_base=wire.url, api_key=ANTHROPIC_KEY)
|
||||
response: Final = _post(gateway, messages_body(model, marker, safeguards=case.sent), beta=case.client_beta)
|
||||
_assert_answered(response, marker)
|
||||
observed: Final = _only(wire)
|
||||
_expect_messages(
|
||||
observed,
|
||||
target="/v1/messages",
|
||||
auth=("x-api-key", ANTHROPIC_KEY),
|
||||
beta=case.expected_beta,
|
||||
safeguards=case.forwarded,
|
||||
)
|
||||
_expect_anthropic_version(observed)
|
||||
_spend_row(f"msg_{marker}")
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True, kw_only=True)
|
||||
class _DropCase:
|
||||
deployment: Callable[[Gateway], Mapping[str, JsonValue]]
|
||||
target: str
|
||||
auth: tuple[str, str] | None
|
||||
keys: frozenset[str] = _MESSAGES_KEYS
|
||||
|
||||
|
||||
_VERTEX_PROJECT: Final = "scripted-project"
|
||||
_VERTEX_LOCATION: Final = "us-east5"
|
||||
_VERTEX_MODEL: Final = "claude-opus-5-5"
|
||||
|
||||
_DROP_CASES: Final = MappingProxyType(
|
||||
{
|
||||
"anthropic": _DropCase(
|
||||
deployment=lambda _: {"model": f"anthropic/{ANTHROPIC_MODEL}", "api_key": ANTHROPIC_KEY},
|
||||
target="/v1/messages",
|
||||
auth=("x-api-key", ANTHROPIC_KEY),
|
||||
),
|
||||
"azure_ai": _DropCase(
|
||||
deployment=lambda _: {"model": _FOUNDRY_MODEL, "api_key": _FOUNDRY_KEY},
|
||||
target=_FOUNDRY_TARGET,
|
||||
auth=("x-api-key", _FOUNDRY_KEY),
|
||||
),
|
||||
"bedrock_claude_platform": _DropCase(
|
||||
deployment=lambda _: {
|
||||
"model": f"bedrock/claude_platform/{ANTHROPIC_MODEL}",
|
||||
"api_key": "synthetic-platform-key",
|
||||
"workspace_id": "ws-synthetic",
|
||||
"aws_region_name": "us-east-1",
|
||||
},
|
||||
target="/v1/messages",
|
||||
auth=("x-api-key", "synthetic-platform-key"),
|
||||
),
|
||||
"vertex_ai": _DropCase(
|
||||
deployment=lambda gateway: {
|
||||
"model": f"vertex_ai/{_VERTEX_MODEL}",
|
||||
"api_key": None,
|
||||
"vertex_project": _VERTEX_PROJECT,
|
||||
"vertex_location": _VERTEX_LOCATION,
|
||||
"vertex_credentials": service_account_json(_VERTEX_PROJECT, gateway.upstream_url.rstrip("/")),
|
||||
},
|
||||
target=(
|
||||
f"/v1/projects/{_VERTEX_PROJECT}/locations/{_VERTEX_LOCATION}"
|
||||
f"/publishers/anthropic/models/{_VERTEX_MODEL}:rawPredict"
|
||||
),
|
||||
auth=None,
|
||||
keys=frozenset({"anthropic_version", "max_tokens", "messages", "stream"}),
|
||||
),
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("row", tuple(_DROP_CASES))
|
||||
def test_a_deployment_dropping_safeguards_sends_neither_the_field_nor_the_beta(gateway: Gateway, row: str) -> None:
|
||||
case: Final = _DROP_CASES[row]
|
||||
marker: Final = uuid.uuid4().hex
|
||||
with wire_server(anthropic_peer) as wire, gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(
|
||||
api_base=wire.url, additional_drop_params=["safeguards"], **case.deployment(gateway)
|
||||
)
|
||||
response: Final = _post(gateway, messages_body(model, marker))
|
||||
_assert_answered(response, marker)
|
||||
_expect_messages(_only(wire), target=case.target, auth=case.auth, beta=None, safeguards=OMITTED, keys=case.keys)
|
||||
_spend_row(f"msg_{marker}")
|
||||
|
||||
|
||||
def test_a_deployment_dropping_context_management_sends_neither_the_field_nor_its_beta(gateway: Gateway) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
with wire_server(anthropic_peer) as wire, gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(
|
||||
model=f"anthropic/{ANTHROPIC_MODEL}",
|
||||
api_base=wire.url,
|
||||
api_key=ANTHROPIC_KEY,
|
||||
additional_drop_params=["context_management"],
|
||||
)
|
||||
body: Final = {
|
||||
**messages_body(model, marker, safeguards=OMITTED),
|
||||
"context_management": {"edits": [{"type": "clear_tool_uses_20250919"}]},
|
||||
}
|
||||
response: Final = _post(gateway, body)
|
||||
_assert_answered(response, marker)
|
||||
_expect_messages(
|
||||
_only(wire), target="/v1/messages", auth=("x-api-key", ANTHROPIC_KEY), beta=None, safeguards=OMITTED
|
||||
)
|
||||
_spend_row(f"msg_{marker}")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("row", tuple(_FOUNDRY_SHAPES))
|
||||
def test_azure_ai_messages_derive_the_beta_instead_of_forwarding_the_client_header(gateway: Gateway, row: str) -> None:
|
||||
case: Final = _FOUNDRY_SHAPES[row]
|
||||
marker: Final = uuid.uuid4().hex
|
||||
with wire_server(anthropic_peer) as wire, gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(model=_FOUNDRY_MODEL, api_base=wire.url, api_key=_FOUNDRY_KEY)
|
||||
response: Final = _post(
|
||||
gateway, messages_body(model, marker, safeguards=case.sent, stream=case.stream), beta=case.client_beta
|
||||
)
|
||||
if case.stream:
|
||||
_assert_streamed(response, marker)
|
||||
else:
|
||||
_assert_answered(response, marker)
|
||||
observed: Final = _only(wire)
|
||||
_expect_messages(
|
||||
observed,
|
||||
target=_FOUNDRY_TARGET,
|
||||
auth=("x-api-key", _FOUNDRY_KEY),
|
||||
beta=case.expected_beta,
|
||||
safeguards=case.forwarded,
|
||||
)
|
||||
_expect_anthropic_version(observed)
|
||||
_spend_row(f"msg_{marker}")
|
||||
|
||||
|
||||
def test_anthropic_sdk_sync_safeguards_derive_the_beta(gateway: Gateway) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
with wire_server(anthropic_peer) as wire, gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(model=f"anthropic/{ANTHROPIC_MODEL}", api_base=wire.url, api_key=ANTHROPIC_KEY)
|
||||
client: Final = anthropic.Anthropic(base_url=str(gateway.client.base_url), api_key=gateway.key, max_retries=0)
|
||||
reply: Final = client.messages.create(
|
||||
model=model,
|
||||
max_tokens=16,
|
||||
messages=[{"role": "user", "content": f"Question marker-{marker}"}],
|
||||
extra_body=_SDK_EXTRA,
|
||||
)
|
||||
assert reply.id == f"msg_{marker}", reply
|
||||
_expect_messages(_only(wire), target="/v1/messages", auth=("x-api-key", ANTHROPIC_KEY), beta=BETA)
|
||||
_spend_row(f"msg_{marker}")
|
||||
|
||||
|
||||
async def test_anthropic_sdk_async_stream_safeguards_derive_the_beta_and_log_spend(gateway: Gateway) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
with wire_server(anthropic_peer) as wire, gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(model=f"anthropic/{ANTHROPIC_MODEL}", api_base=wire.url, api_key=ANTHROPIC_KEY)
|
||||
client: Final = anthropic.AsyncAnthropic(
|
||||
base_url=str(gateway.client.base_url), api_key=gateway.key, max_retries=0
|
||||
)
|
||||
async with client.messages.with_streaming_response.create(
|
||||
model=model,
|
||||
max_tokens=16,
|
||||
messages=[{"role": "user", "content": f"Question marker-{marker}"}],
|
||||
stream=True,
|
||||
extra_body=_SDK_EXTRA,
|
||||
) as streamed:
|
||||
pieces: Final = [
|
||||
event.delta.text
|
||||
async for event in await streamed.parse()
|
||||
if event.type == "content_block_delta" and event.delta.type == "text_delta"
|
||||
]
|
||||
assert "".join(pieces) == f"answer marker-{marker}", pieces
|
||||
_expect_messages(_only(wire), target="/v1/messages", auth=("x-api-key", ANTHROPIC_KEY), beta=BETA)
|
||||
_spend_row(f"msg_{marker}")
|
||||
|
||||
|
||||
def test_azure_ai_sdk_sync_safeguards_derive_the_beta(gateway: Gateway) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
with wire_server(anthropic_peer) as wire, gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(model=_FOUNDRY_MODEL, api_base=wire.url, api_key=_FOUNDRY_KEY)
|
||||
client: Final = anthropic.Anthropic(base_url=str(gateway.client.base_url), api_key=gateway.key, max_retries=0)
|
||||
reply: Final = client.messages.create(
|
||||
model=model,
|
||||
max_tokens=16,
|
||||
messages=[{"role": "user", "content": f"Question marker-{marker}"}],
|
||||
extra_body=_SDK_EXTRA,
|
||||
)
|
||||
assert reply.id == f"msg_{marker}", reply
|
||||
_expect_messages(_only(wire), target=_FOUNDRY_TARGET, auth=("x-api-key", _FOUNDRY_KEY), beta=BETA)
|
||||
_spend_row(f"msg_{marker}")
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True, kw_only=True)
|
||||
class _ProviderCase:
|
||||
parameters: Mapping[str, JsonValue]
|
||||
auth: tuple[str, str] | None
|
||||
expected_beta: str | None
|
||||
target: str = "/v1/messages"
|
||||
api_base_suffix: str = ""
|
||||
model_info: Mapping[str, JsonValue] | None = None
|
||||
client_beta: str | None = None
|
||||
keys: frozenset[str] = _MESSAGES_KEYS
|
||||
extra_headers: Mapping[str, str] = MappingProxyType({})
|
||||
body_beta: list[str] | None = None
|
||||
|
||||
|
||||
_PROVIDERS: Final = MappingProxyType(
|
||||
{
|
||||
"D14_deepseek_with_claude_code_header": _ProviderCase(
|
||||
parameters={"model": "deepseek/deepseek-chat", "api_key": "synthetic-deepseek-key"},
|
||||
client_beta=_CLIENT_BETAS_WITH_DANGEROUS,
|
||||
target="/anthropic/v1/messages",
|
||||
auth=("x-api-key", "synthetic-deepseek-key"),
|
||||
expected_beta=None,
|
||||
),
|
||||
"D15_tencent_with_claude_code_header": _ProviderCase(
|
||||
parameters={"model": "tencent/deepseek-v4-flash", "api_key": "synthetic-tencent-key"},
|
||||
client_beta=_CLIENT_BETAS_WITH_DANGEROUS,
|
||||
auth=("x-api-key", "synthetic-tencent-key"),
|
||||
expected_beta=None,
|
||||
),
|
||||
"D16_minimax_with_claude_code_header": _ProviderCase(
|
||||
parameters={"model": "minimax/MiniMax-M2", "api_key": "synthetic-minimax-key"},
|
||||
client_beta=_CLIENT_BETAS_WITH_DANGEROUS,
|
||||
auth=("x-api-key", "synthetic-minimax-key"),
|
||||
expected_beta=None,
|
||||
),
|
||||
"D17_openai_compatible_native_passthrough": _ProviderCase(
|
||||
parameters={"model": "openai/claude-sonnet-4-5", "api_key": "synthetic-openai-like-key"},
|
||||
api_base_suffix="/v1",
|
||||
model_info={"supported_endpoints": ["/v1/messages"]},
|
||||
auth=("authorization", "Bearer synthetic-openai-like-key"),
|
||||
expected_beta=None,
|
||||
),
|
||||
"D18_sail_with_claude_code_header": _ProviderCase(
|
||||
parameters={"model": "sail/sail-claude", "api_key": "synthetic-sail-key"},
|
||||
client_beta=_CLIENT_BETAS_WITH_DANGEROUS,
|
||||
auth=("authorization", "Bearer synthetic-sail-key"),
|
||||
expected_beta=None,
|
||||
),
|
||||
"D19_edenai": _ProviderCase(
|
||||
parameters={"model": "edenai/claude-sonnet-4-5", "api_key": "synthetic-edenai-key"},
|
||||
auth=("authorization", "Bearer synthetic-edenai-key"),
|
||||
expected_beta=None,
|
||||
),
|
||||
"D21_bedrock_claude_platform": _ProviderCase(
|
||||
parameters={
|
||||
"model": f"bedrock/claude_platform/{ANTHROPIC_MODEL}",
|
||||
"api_key": "synthetic-platform-key",
|
||||
"workspace_id": "ws-synthetic",
|
||||
"aws_region_name": "us-east-1",
|
||||
},
|
||||
auth=("x-api-key", "synthetic-platform-key"),
|
||||
expected_beta=BETA,
|
||||
extra_headers={"anthropic-workspace-id": "ws-synthetic"},
|
||||
),
|
||||
"D22_bedrock_claude_platform_with_claude_code_header": _ProviderCase(
|
||||
parameters={
|
||||
"model": f"bedrock/claude_platform/{ANTHROPIC_MODEL}",
|
||||
"api_key": "synthetic-platform-key",
|
||||
"workspace_id": "ws-synthetic",
|
||||
"aws_region_name": "us-east-1",
|
||||
},
|
||||
client_beta=_CLIENT_BETAS_WITH_DANGEROUS,
|
||||
auth=("x-api-key", "synthetic-platform-key"),
|
||||
expected_beta=_MERGED_BETAS,
|
||||
extra_headers={"anthropic-workspace-id": "ws-synthetic"},
|
||||
),
|
||||
"D23_bedrock_invoke": _ProviderCase(
|
||||
parameters={
|
||||
"model": "bedrock/anthropic.claude-sonnet-4-5-20250929-v1:0",
|
||||
"api_key": "synthetic-bedrock-bearer",
|
||||
"aws_region_name": "us-east-1",
|
||||
},
|
||||
target="/model/anthropic.claude-sonnet-4-5-20250929-v1:0/invoke",
|
||||
auth=("authorization", "Bearer synthetic-bedrock-bearer"),
|
||||
expected_beta=None,
|
||||
keys=frozenset({"anthropic_beta", "anthropic_version", "max_tokens", "messages"}),
|
||||
body_beta=[BETA],
|
||||
),
|
||||
"D24_bedrock_mantle": _ProviderCase(
|
||||
parameters={
|
||||
"model": "bedrock_mantle/anthropic.claude-sonnet-5-v1:0",
|
||||
"api_key": "synthetic-mantle-bearer",
|
||||
"aws_region_name": "us-east-1",
|
||||
},
|
||||
target="/anthropic/v1/messages",
|
||||
auth=("authorization", "Bearer synthetic-mantle-bearer"),
|
||||
expected_beta=BETA,
|
||||
keys=frozenset({"max_tokens", "messages", "model"}),
|
||||
),
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("row", tuple(_PROVIDERS))
|
||||
def test_anthropic_compatible_providers_receive_the_beta_only_where_the_proxy_derives_it(
|
||||
gateway: Gateway, row: str
|
||||
) -> None:
|
||||
case: Final = _PROVIDERS[row]
|
||||
marker: Final = uuid.uuid4().hex
|
||||
with wire_server(anthropic_peer) as wire, gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(
|
||||
model_info=case.model_info, api_base=wire.url + case.api_base_suffix, **case.parameters
|
||||
)
|
||||
response: Final = _post(gateway, messages_body(model, marker), beta=case.client_beta)
|
||||
_assert_answered(response, marker)
|
||||
observed: Final = _only(wire)
|
||||
_expect_messages(observed, target=case.target, auth=case.auth, beta=case.expected_beta, keys=case.keys)
|
||||
for name, value in case.extra_headers.items():
|
||||
assert observed.headers.get(name) == value, observed.header_names
|
||||
assert observed.body.get("anthropic_beta") == case.body_beta, observed.body
|
||||
_spend_row(f"msg_{marker}")
|
||||
|
||||
|
||||
def test_github_copilot_messages_forward_safeguards_without_deriving_the_beta(
|
||||
gateway: Gateway, tmp_path: Path
|
||||
) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
token_dir: Final = tmp_path / "copilot"
|
||||
token_dir.mkdir()
|
||||
with wire_server(anthropic_peer) as wire:
|
||||
(token_dir / "api-key.json").write_text(
|
||||
json.dumps({"token": "synthetic-copilot-token", "expires_at": 4102444800, "endpoints": {"api": wire.url}})
|
||||
)
|
||||
with (
|
||||
owned_proxy(gateway, tmp_path, {"GITHUB_COPILOT_TOKEN_DIR": str(token_dir)}) as owned,
|
||||
owned.scenario() as scenario,
|
||||
):
|
||||
model: Final = scenario.model(model="github_copilot/claude-sonnet-4-5", api_base=wire.url)
|
||||
response: Final = _post(owned, messages_body(model, marker), beta=_CLIENT_BETAS_WITH_DANGEROUS)
|
||||
_assert_answered(response, marker)
|
||||
observed: Final = _only(wire)
|
||||
_expect_messages(
|
||||
observed, target="/v1/messages", auth=("authorization", "Bearer synthetic-copilot-token"), beta=None
|
||||
)
|
||||
assert observed.headers.get("openai-intent") == "messages-proxy", observed.header_names
|
||||
_spend_row(f"msg_{marker}")
|
||||
|
||||
|
||||
def test_non_native_deployments_bridge_to_the_responses_api_without_safeguards_or_the_beta(gateway: Gateway) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
with wire_server(_responses_peer) as wire, gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(
|
||||
model="openai/gpt-4o-mini", api_base=wire.url + "/v1", api_key="synthetic-openai-key"
|
||||
)
|
||||
response: Final = _post(gateway, messages_body(model, marker), beta=_CLIENT_BETAS_WITH_DANGEROUS)
|
||||
assert response.status_code == 200 and f"answer marker-{marker}" in response.text, response.text
|
||||
observed: Final = _only(wire)
|
||||
assert (observed.method, observed.target) == ("POST", "/v1/responses")
|
||||
assert observed.headers.get("authorization") == "Bearer synthetic-openai-key", observed.header_names
|
||||
assert observed.beta is None, observed.header_names
|
||||
assert observed.keys == _RESPONSES_KEYS, observed.keys
|
||||
_spend_row(str(_JSON_OBJECT.validate_json(response.content)["id"]))
|
||||
|
||||
|
||||
def test_chat_completions_forward_safeguards_verbatim_without_deriving_the_beta(gateway: Gateway) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
with wire_server(anthropic_peer) as wire, gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(model=f"anthropic/{ANTHROPIC_MODEL}", api_base=wire.url, api_key=ANTHROPIC_KEY)
|
||||
client: Final = openai.OpenAI(
|
||||
base_url=str(gateway.client.base_url).rstrip("/") + "/v1", api_key=gateway.key, max_retries=0
|
||||
)
|
||||
completion: Final = client.chat.completions.create(
|
||||
model=model, messages=[{"role": "user", "content": f"Question marker-{marker}"}], extra_body=_SDK_EXTRA
|
||||
)
|
||||
assert completion.choices[0].message.content == f"answer marker-{marker}", completion
|
||||
_expect_messages(
|
||||
_only(wire),
|
||||
target="/v1/messages",
|
||||
auth=("x-api-key", ANTHROPIC_KEY),
|
||||
beta=None,
|
||||
keys=frozenset({"max_tokens", "messages", "model"}),
|
||||
)
|
||||
_spend_row(completion.id)
|
||||
|
||||
|
||||
async def test_responses_stream_forwards_safeguards_verbatim_without_deriving_the_beta(gateway: Gateway) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
with wire_server(anthropic_peer) as wire, gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(model=f"anthropic/{ANTHROPIC_MODEL}", api_base=wire.url, api_key=ANTHROPIC_KEY)
|
||||
client: Final = openai.AsyncOpenAI(
|
||||
base_url=str(gateway.client.base_url).rstrip("/") + "/v1", api_key=gateway.key, max_retries=0
|
||||
)
|
||||
async with client.responses.with_streaming_response.create(
|
||||
model=model, input=f"Question marker-{marker}", stream=True, extra_body=_SDK_EXTRA
|
||||
) as streamed:
|
||||
events: Final = [event async for event in await streamed.parse()]
|
||||
deltas: Final = [event.delta for event in events if event.type == "response.output_text.delta"]
|
||||
assert "".join(deltas) == f"answer marker-{marker}", events
|
||||
assert any(event.type == "response.completed" for event in events), events
|
||||
_expect_messages(_only(wire), target="/v1/messages", auth=("x-api-key", ANTHROPIC_KEY), beta=None)
|
||||
_responses_spend_row(model)
|
||||
|
||||
|
||||
_FOUNDRY_REJECTED_SHAPES: Final[Mapping[str, JsonValue]] = MappingProxyType(
|
||||
{"S1_int": 1, "S2_empty_string": "", "S3_five_kb_string": "x" * 5000}
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("row", tuple(_FOUNDRY_REJECTED_SHAPES))
|
||||
def test_foundry_rejecting_a_non_array_safeguards_value_reaches_the_caller_after_one_attempt(
|
||||
gateway: Gateway, row: str
|
||||
) -> None:
|
||||
value: Final = _FOUNDRY_REJECTED_SHAPES[row]
|
||||
marker: Final = uuid.uuid4().hex
|
||||
with (
|
||||
wire_server(_error_peer(400, "invalid_request_error", _FOUNDRY_ARRAY_ERROR)) as wire,
|
||||
gateway.scenario() as scenario,
|
||||
):
|
||||
model: Final = scenario.model(model=_FOUNDRY_MODEL, api_base=wire.url, api_key=_FOUNDRY_KEY)
|
||||
response: Final = _post(gateway, messages_body(model, marker, safeguards=value))
|
||||
assert response.status_code == 400 and _FOUNDRY_ARRAY_ERROR in response.text, response.text
|
||||
_expect_messages(
|
||||
_only(wire), target=_FOUNDRY_TARGET, auth=("x-api-key", _FOUNDRY_KEY), beta=BETA, safeguards=value
|
||||
)
|
||||
|
||||
|
||||
def test_duplicate_safeguards_keys_resolve_to_the_last_value(gateway: Gateway) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
with wire_server(anthropic_peer) as wire, gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(model=_FOUNDRY_MODEL, api_base=wire.url, api_key=_FOUNDRY_KEY)
|
||||
raw: Final = (
|
||||
f'{{"model": "{model}", "max_tokens": 16, "num_retries": 0, "cache": {{"no-cache": true}}, '
|
||||
f'"messages": [{{"role": "user", "content": "Question marker-{marker}"}}], '
|
||||
f'"safeguards": {json.dumps(SAFEGUARDS)}, "safeguards": null}}'
|
||||
)
|
||||
response: Final = gateway.client.post(
|
||||
"/v1/messages",
|
||||
content=raw.encode(),
|
||||
headers={
|
||||
"Authorization": f"Bearer {gateway.key}",
|
||||
"content-type": "application/json",
|
||||
**_ANTHROPIC_VERSION,
|
||||
},
|
||||
)
|
||||
_assert_answered(response, marker)
|
||||
_expect_messages(
|
||||
_only(wire), target=_FOUNDRY_TARGET, auth=("x-api-key", _FOUNDRY_KEY), beta=None, safeguards=OMITTED
|
||||
)
|
||||
_spend_row(f"msg_{marker}")
|
||||
|
||||
|
||||
def test_the_same_safeguards_sent_twice_are_two_upstream_calls_with_two_spend_rows(gateway: Gateway) -> None:
|
||||
markers: Final = (uuid.uuid4().hex, uuid.uuid4().hex)
|
||||
with wire_server(anthropic_peer) as wire, gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(model=_FOUNDRY_MODEL, api_base=wire.url, api_key=_FOUNDRY_KEY)
|
||||
responses: Final = tuple(_post(gateway, messages_body(model, marker)) for marker in markers)
|
||||
for marker, response in zip(markers, responses, strict=True):
|
||||
_assert_answered(response, marker)
|
||||
received: Final = tuple(_observe(request) for request in wire.drain())
|
||||
assert len(received) == 2, received
|
||||
for observed in received:
|
||||
_expect_messages(observed, target=_FOUNDRY_TARGET, auth=("x-api-key", _FOUNDRY_KEY), beta=BETA)
|
||||
assert sorted(item.marker for item in received) == sorted(markers)
|
||||
for marker in markers:
|
||||
_spend_row(f"msg_{marker}")
|
||||
|
||||
|
||||
_UPSTREAM_ERRORS: Final = MappingProxyType(
|
||||
{
|
||||
"S6_unauthorized_deployment": (401, "authentication_error", "invalid x-api-key"),
|
||||
"S7_foundry_extra_inputs_rejection": (400, "invalid_request_error", _FOUNDRY_EXTRA_INPUT_ERROR),
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("row", tuple(_UPSTREAM_ERRORS))
|
||||
def test_upstream_errors_reach_the_caller_with_the_beta_sent_once(gateway: Gateway, row: str) -> None:
|
||||
status, kind, message = _UPSTREAM_ERRORS[row]
|
||||
marker: Final = uuid.uuid4().hex
|
||||
with wire_server(_error_peer(status, kind, message)) as wire, gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(model=_FOUNDRY_MODEL, api_base=wire.url, api_key=_FOUNDRY_KEY)
|
||||
response: Final = _post(gateway, messages_body(model, marker))
|
||||
assert response.status_code == status and message in response.text, response.text
|
||||
_expect_messages(_only(wire), target=_FOUNDRY_TARGET, auth=("x-api-key", _FOUNDRY_KEY), beta=BETA)
|
||||
|
||||
|
||||
def _post_noting_whether_upstream_was_called(
|
||||
gateway: Gateway, body: Mapping[str, JsonValue], wire: Wire
|
||||
) -> tuple[httpx.Response, bool]:
|
||||
before: Final = wire.received.qsize()
|
||||
response: Final = _post(gateway, body)
|
||||
assert response.status_code == 200, response.text
|
||||
return response, wire.received.qsize() > before
|
||||
|
||||
|
||||
def _until_served_from_cache(gateway: Gateway, body: Mapping[str, JsonValue], wire: Wire) -> tuple[httpx.Response, int]:
|
||||
deadline: Final = time.monotonic() + 20
|
||||
for attempt in itertools.count(1):
|
||||
response, called_upstream = _post_noting_whether_upstream_was_called(gateway, body, wire)
|
||||
if not called_upstream:
|
||||
return response, attempt
|
||||
assert time.monotonic() < deadline, "the response cache never served the repeated request"
|
||||
raise AssertionError("unreachable")
|
||||
|
||||
|
||||
def test_the_response_cache_serves_the_repeated_request_without_an_upstream_call(gateway: Gateway) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
with wire_server(anthropic_peer) as wire, gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(model=_FOUNDRY_MODEL, api_base=wire.url, api_key=_FOUNDRY_KEY)
|
||||
body: Final = messages_body(model, marker, no_cache=False)
|
||||
first: Final = _post(gateway, body)
|
||||
_assert_answered(first, marker)
|
||||
_expect_messages(_only(wire), target=_FOUNDRY_TARGET, auth=("x-api-key", _FOUNDRY_KEY), beta=BETA)
|
||||
_spend_row(f"msg_{marker}")
|
||||
hit, attempts = _until_served_from_cache(gateway, body, wire)
|
||||
assert hit.json()["id"] == first.json()["id"], hit.text
|
||||
assert hit.json()["content"] == first.json()["content"], hit.text
|
||||
_cache_hit_spend_row(f"msg_{marker}")
|
||||
misses: Final = tuple(_observe(request) for request in wire.drain())
|
||||
assert len(misses) == attempts - 1, (attempts, misses)
|
||||
for observed in misses:
|
||||
_expect_messages(observed, target=_FOUNDRY_TARGET, auth=("x-api-key", _FOUNDRY_KEY), beta=BETA)
|
||||
1206
tests/integration/providers/test_bedrock_mantle_count_tokens_wire.py
Normal file
1206
tests/integration/providers/test_bedrock_mantle_count_tokens_wire.py
Normal file
File diff suppressed because it is too large
Load diff
File diff suppressed because it is too large
Load diff
186
tests/integration/sdk/test_databricks_json_schema_refs_sdk.py
Normal file
186
tests/integration/sdk/test_databricks_json_schema_refs_sdk.py
Normal file
|
|
@ -0,0 +1,186 @@
|
|||
import json
|
||||
import uuid
|
||||
from collections.abc import Iterator, Mapping
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
from integration._support.wire import Reply, Request, wire_server
|
||||
from pydantic import BaseModel, ConfigDict, Field, JsonValue
|
||||
|
||||
import litellm
|
||||
|
||||
_QWEN_BACKEND: Final = "databricks-qwen35-122b-a10b"
|
||||
_CLAUDE_BACKEND: Final = "databricks-claude-haiku-4-5"
|
||||
_API_KEY: Final = "synthetic-databricks-key"
|
||||
_ORDER_JSON: Final = '{"buyer":{"name":"Ada"},"seller":{"name":"Bo"}}'
|
||||
_MESSAGES: Final[tuple[Mapping[str, str], ...]] = ({"role": "user", "content": "Return the order as JSON."},)
|
||||
|
||||
|
||||
class Person(BaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
name: str
|
||||
|
||||
|
||||
class Order(BaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
buyer: Person
|
||||
seller: Person
|
||||
|
||||
|
||||
class _StrictObject(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="ignore")
|
||||
additional_properties: bool | None = Field(default=None, alias="additionalProperties")
|
||||
|
||||
|
||||
class _OrderSchema(_StrictObject):
|
||||
defs: Mapping[str, _StrictObject] = Field(alias="$defs")
|
||||
|
||||
|
||||
class _JsonSchema(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="ignore")
|
||||
name: str
|
||||
strict: bool | None = None
|
||||
json_schema: dict[str, JsonValue] = Field(alias="schema")
|
||||
|
||||
|
||||
class _ResponseFormat(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="ignore")
|
||||
type: str
|
||||
json_schema: _JsonSchema
|
||||
|
||||
|
||||
class _Function(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="ignore")
|
||||
name: str
|
||||
parameters: dict[str, JsonValue]
|
||||
|
||||
|
||||
class _Tool(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="ignore")
|
||||
function: _Function
|
||||
|
||||
|
||||
class _Outbound(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="ignore")
|
||||
model: str
|
||||
response_format: _ResponseFormat | None = None
|
||||
tools: tuple[_Tool, ...] = ()
|
||||
|
||||
|
||||
class _CallerMessage(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="ignore")
|
||||
content: str
|
||||
|
||||
|
||||
class _CallerChoice(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="ignore")
|
||||
message: _CallerMessage
|
||||
|
||||
|
||||
class _CallerCompletion(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="ignore")
|
||||
choices: tuple[_CallerChoice, ...]
|
||||
|
||||
|
||||
def _completion(backend: str) -> Reply:
|
||||
claude: Final = "claude" in backend
|
||||
tool_call: Final[dict[str, JsonValue]] = {
|
||||
"id": "json-tool-call",
|
||||
"type": "function",
|
||||
"function": {"name": "json_tool_call", "arguments": _ORDER_JSON},
|
||||
}
|
||||
message: Final[dict[str, JsonValue]] = (
|
||||
{"role": "assistant", "content": None, "tool_calls": [tool_call]}
|
||||
if claude
|
||||
else {"role": "assistant", "content": _ORDER_JSON}
|
||||
)
|
||||
body: Final[dict[str, JsonValue]] = {
|
||||
"id": f"chatcmpl-{uuid.uuid4().hex}",
|
||||
"object": "chat.completion",
|
||||
"created": 1,
|
||||
"model": backend,
|
||||
"choices": [{"index": 0, "message": message, "finish_reason": "tool_calls" if claude else "stop"}],
|
||||
"usage": {"prompt_tokens": 9, "completion_tokens": 7, "total_tokens": 16},
|
||||
}
|
||||
return Reply(body=json.dumps(body).encode())
|
||||
|
||||
|
||||
def _refs(node: JsonValue) -> Iterator[str]:
|
||||
if isinstance(node, list):
|
||||
for item in node:
|
||||
yield from _refs(item)
|
||||
return
|
||||
if not isinstance(node, dict):
|
||||
return
|
||||
for key, value in node.items():
|
||||
if key == "$ref" and isinstance(value, str):
|
||||
yield value
|
||||
else:
|
||||
yield from _refs(value)
|
||||
|
||||
|
||||
def _sent_schema(backend: str, received: tuple[Request, ...]) -> dict[str, JsonValue]:
|
||||
assert [(request.method, request.target) for request in received] == [("POST", "/chat/completions")]
|
||||
outbound: Final = _Outbound.model_validate_json(received[0].body)
|
||||
assert outbound.model == backend
|
||||
if "claude" in backend:
|
||||
assert outbound.response_format is None
|
||||
assert [tool.function.name for tool in outbound.tools] == ["json_tool_call"]
|
||||
return outbound.tools[0].function.parameters
|
||||
assert outbound.response_format is not None
|
||||
assert (outbound.response_format.type, outbound.response_format.json_schema.name) == ("json_schema", "Order")
|
||||
assert outbound.response_format.json_schema.strict is True
|
||||
return outbound.response_format.json_schema.json_schema
|
||||
|
||||
|
||||
def _assert_strict_order_schema(schema: dict[str, JsonValue]) -> None:
|
||||
assert set(_refs(schema)) == {"#/$defs/Person"}, schema
|
||||
parsed: Final = _OrderSchema.model_validate(schema)
|
||||
assert (parsed.additional_properties, parsed.defs["Person"].additional_properties) == (False, False), schema
|
||||
|
||||
|
||||
def _caller_order(response: litellm.ModelResponse) -> Order:
|
||||
return Order.model_validate_json(
|
||||
_CallerCompletion.model_validate_json(response.model_dump_json()).choices[0].message.content
|
||||
)
|
||||
|
||||
|
||||
_BACKENDS: Final = (
|
||||
pytest.param(_QWEN_BACKEND, id="qwen"),
|
||||
pytest.param(_CLAUDE_BACKEND, id="claude"),
|
||||
)
|
||||
_EXPECTED_ORDER: Final = Order(buyer=Person(name="Ada"), seller=Person(name="Bo"))
|
||||
|
||||
|
||||
@pytest.mark.parametrize("backend", _BACKENDS)
|
||||
def test_databricks_pydantic_response_format_sends_strict_defs_refs(backend: str) -> None:
|
||||
with wire_server(lambda _request: _completion(backend)) as wire:
|
||||
response: Final = litellm.completion(
|
||||
model=f"databricks/{backend}",
|
||||
api_base=wire.url,
|
||||
api_key=_API_KEY,
|
||||
messages=[dict(message) for message in _MESSAGES],
|
||||
response_format=Order,
|
||||
num_retries=0,
|
||||
)
|
||||
received: Final = wire.drain()
|
||||
_assert_strict_order_schema(_sent_schema(backend, received))
|
||||
assert isinstance(response, litellm.ModelResponse)
|
||||
assert _caller_order(response) == _EXPECTED_ORDER
|
||||
|
||||
|
||||
@pytest.mark.parametrize("backend", _BACKENDS)
|
||||
async def test_databricks_pydantic_response_format_sends_strict_defs_refs_async(backend: str) -> None:
|
||||
with wire_server(lambda _request: _completion(backend)) as wire:
|
||||
response: Final = await litellm.acompletion(
|
||||
model=f"databricks/{backend}",
|
||||
api_base=wire.url,
|
||||
api_key=_API_KEY,
|
||||
messages=[dict(message) for message in _MESSAGES],
|
||||
response_format=Order,
|
||||
num_retries=0,
|
||||
)
|
||||
received: Final = wire.drain()
|
||||
_assert_strict_order_schema(_sent_schema(backend, received))
|
||||
assert isinstance(response, litellm.ModelResponse)
|
||||
assert _caller_order(response) == _EXPECTED_ORDER
|
||||
|
|
@ -1,14 +1,20 @@
|
|||
import json
|
||||
import os
|
||||
from datetime import datetime, timezone
|
||||
from functools import partial
|
||||
from hashlib import sha256
|
||||
from itertools import chain
|
||||
from typing import Final
|
||||
from uuid import uuid4
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from integration._support.client import JSON_OBJECT, Gateway, Scenario, eventually, object_value, string_value
|
||||
from integration._support.database import read_rows
|
||||
from integration._support.upstream import delete_scenario, register_scenario
|
||||
from integration._support.database import read_rows, write_rows
|
||||
from integration._support.prisma_clients import connected_proxy_prisma
|
||||
from integration._support.upstream import ScenarioHandle, delete_scenario, register_scenario
|
||||
from integration.cost_calculation.cost_tracking_case import JsonResponse, RoutedResponse, TextResponse
|
||||
from pydantic import JsonValue
|
||||
from pydantic import JsonValue, TypeAdapter
|
||||
|
||||
from litellm.constants import MAX_OBJECTS_PER_POLL_CYCLE
|
||||
|
||||
|
|
@ -59,17 +65,17 @@ def _accepting_routes() -> dict[str, JsonResponse | TextResponse]:
|
|||
}
|
||||
|
||||
|
||||
def _gone_at_provider_routes() -> RoutedResponse:
|
||||
def _gone_at_provider_routes(*, retryable: bool = False) -> RoutedResponse:
|
||||
return RoutedResponse(
|
||||
content_type="application/x-routed",
|
||||
routes={
|
||||
**_accepting_routes(),
|
||||
"GET /batches/batch-$REQUEST_ID": JsonResponse(
|
||||
content_type="application/json",
|
||||
status=404,
|
||||
status=503 if retryable else 404,
|
||||
body={
|
||||
"error": {
|
||||
"message": "No batch found with id 'batch-$REQUEST_ID'.",
|
||||
"message": "Temporary failure" if retryable else "No batch found with id 'batch-$REQUEST_ID'.",
|
||||
"type": "invalid_request_error",
|
||||
"param": "id",
|
||||
"code": "batch_not_found",
|
||||
|
|
@ -80,7 +86,7 @@ def _gone_at_provider_routes() -> RoutedResponse:
|
|||
)
|
||||
|
||||
|
||||
def _completed_routes() -> RoutedResponse:
|
||||
def _completed_routes(*, ready: bool = True) -> RoutedResponse:
|
||||
output_line: Final = {
|
||||
"id": "batch_req_1",
|
||||
"custom_id": "r1",
|
||||
|
|
@ -106,7 +112,8 @@ def _completed_routes() -> RoutedResponse:
|
|||
routes={
|
||||
**_accepting_routes(),
|
||||
"GET /batches/batch-$REQUEST_ID": JsonResponse(
|
||||
content_type="application/json", body=_batch("completed", "file-out-$REQUEST_ID")
|
||||
content_type="application/json",
|
||||
body=_batch("completed", "file-out-$REQUEST_ID") if ready else _batch("in_progress", None),
|
||||
),
|
||||
"GET /files/file-out-$REQUEST_ID/content": TextResponse(
|
||||
content_type="application/jsonl", body=json.dumps(output_line, separators=(",", ":")) + "\n"
|
||||
|
|
@ -115,7 +122,37 @@ def _completed_routes() -> RoutedResponse:
|
|||
)
|
||||
|
||||
|
||||
def _scripted_deployment(scenario: Scenario, marker: str, routes: RoutedResponse) -> str:
|
||||
def _batch_was_polled(upstream: ScenarioHandle) -> bool:
|
||||
response: Final = httpx.get(f"{upstream.control_url}/__observations", timeout=15, trust_env=False)
|
||||
response.raise_for_status()
|
||||
observations: Final = TypeAdapter(dict[str, list[dict[str, JsonValue]]]).validate_json(response.content)
|
||||
return any(
|
||||
item.get("method") == "GET" and string_value(item["path"]).startswith(f"/{upstream.scenario_id}/batches/")
|
||||
for item in observations["requests"]
|
||||
)
|
||||
|
||||
|
||||
def _missing_output_routes() -> RoutedResponse:
|
||||
return RoutedResponse(
|
||||
content_type="application/x-routed",
|
||||
routes={
|
||||
**_completed_routes().routes,
|
||||
"GET /files/file-out-$REQUEST_ID/content": JsonResponse(
|
||||
content_type="application/json",
|
||||
status=404,
|
||||
body={
|
||||
"error": {
|
||||
"message": "No file found with id 'file-out-$REQUEST_ID'.",
|
||||
"type": "invalid_request_error",
|
||||
"code": "file_not_found",
|
||||
}
|
||||
},
|
||||
),
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def _scripted_deployment(scenario: Scenario, marker: str, routes: RoutedResponse) -> tuple[str, ScenarioHandle]:
|
||||
scenario_id: Final = f"poll-{marker}-{sha256(os.urandom(16)).hexdigest()[:12]}"
|
||||
handle: Final = register_scenario(scenario_id, routes)
|
||||
scenario.cleanups.callback(delete_scenario, handle)
|
||||
|
|
@ -133,7 +170,7 @@ def _scripted_deployment(scenario: Scenario, marker: str, routes: RoutedResponse
|
|||
},
|
||||
)
|
||||
scenario.cleanups.callback(scenario.delete_model, string_value(object_value(created["model_info"])["id"]))
|
||||
return string_value(created["model_name"])
|
||||
return string_value(created["model_name"]), handle
|
||||
|
||||
|
||||
def _submitted_batch_id(gateway: Gateway, key: str, model_name: str) -> str:
|
||||
|
|
@ -173,20 +210,51 @@ def _managed_rows(batch_ids: tuple[str, ...]) -> list[dict[str, JsonValue]]:
|
|||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("retryable", "old_completed", "backlog"),
|
||||
[(False, False, 0), (False, True, 0), (True, False, 0), (True, True, 0), (True, False, 1000)],
|
||||
)
|
||||
@pytest.mark.timeout(180)
|
||||
@pytest.mark.covers("quota_management.spend_tracking.batch_costs.uncostable_rows_retire_so_newer_batches_are_costed")
|
||||
def test_batches_gone_at_provider_do_not_starve_a_newer_batch_out_of_cost_polling(gateway: Gateway) -> None:
|
||||
def test_uncostable_batches_do_not_starve_billable_batches(
|
||||
gateway: Gateway, retryable: bool, old_completed: bool, backlog: int
|
||||
) -> None:
|
||||
with gateway.scenario() as scenario:
|
||||
key: Final = scenario.key()
|
||||
owner: Final = f"batch-backlog-{uuid4().hex}"
|
||||
if backlog:
|
||||
_insert_scan_rows(owner, "old", backlog, datetime(2000, 1, 1, tzinfo=timezone.utc))
|
||||
scenario.cleanups.callback(
|
||||
write_rows, 'DELETE FROM "LiteLLM_ManagedObjectTable" WHERE created_by = %s', (owner,)
|
||||
)
|
||||
team_id: Final = scenario.team()
|
||||
key: Final = scenario.key(team_id=team_id)
|
||||
gone_batch_ids: Final = tuple(
|
||||
_submitted_batch_id(
|
||||
gateway, key, _scripted_deployment(scenario, f"gone{index}", _gone_at_provider_routes())
|
||||
gateway,
|
||||
key,
|
||||
_scripted_deployment(scenario, f"gone{index}", _gone_at_provider_routes(retryable=retryable))[0],
|
||||
)
|
||||
for index in range(MAX_OBJECTS_PER_POLL_CYCLE)
|
||||
)
|
||||
costable_batch_id: Final = _submitted_batch_id(
|
||||
gateway, key, _scripted_deployment(scenario, "costable", _completed_routes())
|
||||
missing_output_id: Final = _submitted_batch_id(
|
||||
gateway, key, _scripted_deployment(scenario, "missing-output", _missing_output_routes())[0]
|
||||
)
|
||||
model, upstream = _scripted_deployment(scenario, "costable", _completed_routes(ready=False))
|
||||
costable_batch_id: Final = _submitted_batch_id(gateway, key, model)
|
||||
for batch_id in (*gone_batch_ids, missing_output_id, costable_batch_id):
|
||||
scenario.cleanups.callback(
|
||||
write_rows,
|
||||
'DELETE FROM "LiteLLM_ManagedObjectTable" WHERE unified_object_id = %s',
|
||||
(batch_id,),
|
||||
)
|
||||
if old_completed:
|
||||
write_rows(
|
||||
'UPDATE "LiteLLM_ManagedObjectTable" SET status = %s, '
|
||||
"created_at = NOW() - INTERVAL '30 days' WHERE unified_object_id = %s",
|
||||
("completed", costable_batch_id),
|
||||
)
|
||||
assert eventually(lambda: _batch_was_polled(upstream), bool, seconds=120)
|
||||
register_scenario(upstream.scenario_id, _completed_routes(), control_url=upstream.control_url)
|
||||
spend_rows: Final = eventually(
|
||||
lambda: read_rows(
|
||||
'SELECT call_type, status, prompt_tokens, completion_tokens, spend FROM "LiteLLM_SpendLogs" '
|
||||
|
|
@ -209,4 +277,134 @@ def test_batches_gone_at_provider_do_not_starve_a_newer_batch_out_of_cost_pollin
|
|||
}
|
||||
]
|
||||
assert _managed_rows((costable_batch_id,)) == [{"batch_processed": True}]
|
||||
assert _managed_rows(gone_batch_ids) == [{"batch_processed": True}] * MAX_OBJECTS_PER_POLL_CYCLE
|
||||
assert eventually(
|
||||
lambda: _managed_rows((missing_output_id,)),
|
||||
lambda rows: rows == [{"batch_processed": True}],
|
||||
seconds=120,
|
||||
) == [{"batch_processed": True}]
|
||||
expected_blockers: Final = [{"batch_processed": not retryable}] * MAX_OBJECTS_PER_POLL_CYCLE
|
||||
blockers: Final = eventually(
|
||||
lambda: _managed_rows(gone_batch_ids), lambda rows: rows == expected_blockers, seconds=120
|
||||
)
|
||||
assert blockers == expected_blockers
|
||||
|
||||
expected_spend: Final = BATCH_COST_SHARE * (
|
||||
PROMPT_TOKENS * INPUT_COST_PER_TOKEN + COMPLETION_TOKENS * OUTPUT_COST_PER_TOKEN
|
||||
)
|
||||
for table, column, identity in (
|
||||
("LiteLLM_VerificationToken", "token", sha256(key.encode()).hexdigest()),
|
||||
("LiteLLM_TeamTable", "team_id", team_id),
|
||||
):
|
||||
balance: Final = eventually(
|
||||
partial(read_rows, f'SELECT spend FROM "{table}" WHERE {column} = %s', (identity,)),
|
||||
lambda rows: len(rows) == 1 and rows[0]["spend"] == pytest.approx(expected_spend),
|
||||
seconds=70,
|
||||
)
|
||||
assert balance == [{"spend": pytest.approx(expected_spend)}]
|
||||
|
||||
|
||||
def _insert_scan_rows(owner: str, group: str, size: int, created_at: datetime) -> tuple[str, ...]:
|
||||
prefix: Final = f"{owner}-{group}-"
|
||||
write_rows(
|
||||
'INSERT INTO "LiteLLM_ManagedObjectTable" '
|
||||
"(id, unified_object_id, model_object_id, file_object, file_purpose, status, created_by, created_at, updated_at) "
|
||||
"SELECT %s || lpad(n::text, 8, '0'), %s || n, %s || n, '{}'::jsonb, 'batch', 'completed', %s, "
|
||||
"%s::timestamptz, NOW() FROM generate_series(0, %s::int - 1) n",
|
||||
(prefix, prefix, prefix, owner, created_at.isoformat(), str(size)),
|
||||
)
|
||||
return tuple(f"{prefix}{index:08}" for index in range(size))
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("count", [2 * MAX_OBJECTS_PER_POLL_CYCLE, 2 * MAX_OBJECTS_PER_POLL_CYCLE + 1, 1000])
|
||||
async def test_arrivals_progress_while_retry_scan_advances_and_wraps(count: int) -> None:
|
||||
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
|
||||
from litellm.proxy.utils import ProxyLogging
|
||||
from litellm.router import Router
|
||||
from litellm_enterprise.proxy.common_utils.check_batch_cost import CheckBatchCost
|
||||
|
||||
owner: Final = f"batch-scan-{uuid4().hex}"
|
||||
|
||||
try:
|
||||
old_ids: Final = _insert_scan_rows(owner, "old", count, datetime(2000, 1, 1, tzinfo=timezone.utc))
|
||||
async with connected_proxy_prisma() as prisma:
|
||||
logging: Final = ProxyLogging(UserApiKeyCache())
|
||||
router: Final = Router()
|
||||
poller: Final = CheckBatchCost(logging, prisma, router)
|
||||
where: Final = {"created_by": owner, "batch_processed": False}
|
||||
first: Final = await poller._find_jobs(where)
|
||||
assert tuple(row.id for row in first) == old_ids[:MAX_OBJECTS_PER_POLL_CYCLE]
|
||||
|
||||
recent_ids: Final = _insert_scan_rows(
|
||||
owner, "recent", MAX_OBJECTS_PER_POLL_CYCLE + 1, datetime.now(timezone.utc)
|
||||
)
|
||||
|
||||
async def poll_with_arrival(turn: int) -> tuple[str, ...]:
|
||||
if turn % 4 == 0:
|
||||
_insert_scan_rows(owner, f"arrival-{turn}", 1, datetime.now(timezone.utc))
|
||||
page: Final = await poller._find_jobs(where)
|
||||
assert len(page) <= MAX_OBJECTS_PER_POLL_CYCLE
|
||||
write_rows(
|
||||
'UPDATE "LiteLLM_ManagedObjectTable" SET batch_processed = true '
|
||||
"WHERE created_by = %s AND id = ANY(%s) AND id LIKE %s",
|
||||
(owner, [row.id for row in page], f"{owner}-arrival-%"),
|
||||
)
|
||||
return tuple(row.id for row in page)
|
||||
|
||||
pages: Final = tuple(
|
||||
[await poll_with_arrival(turn) for turn in range(4 * (count // MAX_OBJECTS_PER_POLL_CYCLE + 2))]
|
||||
)
|
||||
selected: Final = tuple(chain.from_iterable(pages))
|
||||
assert all(selected.count(identity) >= 2 for identity in recent_ids)
|
||||
selected_old: Final = tuple(identity for identity in selected if identity in old_ids)
|
||||
assert selected_old[: count - MAX_OBJECTS_PER_POLL_CYCLE] == old_ids[MAX_OBJECTS_PER_POLL_CYCLE:]
|
||||
assert set(old_ids).issubset(selected_old)
|
||||
|
||||
latest_ids: Final = _insert_scan_rows(owner, "latest", 1, datetime.now(timezone.utc))
|
||||
followups: Final = tuple([await poller._find_jobs(where) for _ in range(12)])
|
||||
assert any(latest_ids[0] in tuple(row.id for row in page) for page in followups)
|
||||
|
||||
restarted: Final = CheckBatchCost(logging, prisma, router)
|
||||
restart_page: Final = await restarted._find_jobs(where)
|
||||
assert tuple(row.id for row in restart_page) == old_ids[:MAX_OBJECTS_PER_POLL_CYCLE]
|
||||
finally:
|
||||
write_rows('DELETE FROM "LiteLLM_ManagedObjectTable" WHERE created_by = %s', (owner,))
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("steady_arrivals", [False, True])
|
||||
@pytest.mark.parametrize("backlog", [0, 1000])
|
||||
async def test_recent_batches_are_revisited_without_scanning_old_backlog(steady_arrivals: bool, backlog: int) -> None:
|
||||
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
|
||||
from litellm.proxy.utils import ProxyLogging
|
||||
from litellm.router import Router
|
||||
from litellm_enterprise.proxy.common_utils.check_batch_cost import CheckBatchCost
|
||||
|
||||
owner: Final = f"batch-repeat-{uuid4().hex}"
|
||||
try:
|
||||
_insert_scan_rows(owner, "old", backlog, datetime(2000, 1, 1, tzinfo=timezone.utc))
|
||||
async with connected_proxy_prisma() as prisma:
|
||||
poller: Final = CheckBatchCost(ProxyLogging(UserApiKeyCache()), prisma, Router())
|
||||
where: Final = {"created_by": owner, "batch_processed": False}
|
||||
await poller._find_jobs(where)
|
||||
target: Final = _insert_scan_rows(owner, "target", 1, datetime.now(timezone.utc))[0]
|
||||
assert target in tuple(row.id for row in await poller._find_jobs(where))
|
||||
for cycle in range(3):
|
||||
if steady_arrivals:
|
||||
_insert_scan_rows(owner, f"arrival-{cycle}", 1, datetime.now(timezone.utc))
|
||||
pages: Final = tuple([await poller._find_jobs(where) for _ in range(4)])
|
||||
assert any(target in tuple(row.id for row in page) for page in pages)
|
||||
assert all(len(page) <= MAX_OBJECTS_PER_POLL_CYCLE for page in pages)
|
||||
write_rows(
|
||||
'UPDATE "LiteLLM_ManagedObjectTable" SET batch_processed = true WHERE created_by = %s AND id != %s',
|
||||
(owner, target),
|
||||
)
|
||||
write_rows(
|
||||
'UPDATE "LiteLLM_ManagedObjectTable" SET created_at = %s WHERE id = %s',
|
||||
(datetime(2000, 1, 1, tzinfo=timezone.utc), target),
|
||||
)
|
||||
recovery: Final = tuple([await poller._find_jobs(where) for _ in range(4)])
|
||||
assert any(target in tuple(row.id for row in page) for page in recovery)
|
||||
|
||||
finally:
|
||||
write_rows('DELETE FROM "LiteLLM_ManagedObjectTable" WHERE created_by = %s', (owner,))
|
||||
|
|
|
|||
|
|
@ -3,7 +3,7 @@ import pytest
|
|||
|
||||
import json
|
||||
from io import BytesIO
|
||||
from typing import Dict, List
|
||||
from typing import Dict, List, Final
|
||||
from litellm.router_utils.batch_utils import (
|
||||
replace_model_in_jsonl,
|
||||
_get_router_metadata_variable_name,
|
||||
|
|
@ -333,3 +333,31 @@ def test_is_batch_retrieve_call_type_matches_only_batch_retrieves():
|
|||
assert is_batch_retrieve_call_type(call_type.value) is False
|
||||
|
||||
assert is_batch_retrieve_call_type(None) is False
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("call_type", "cost", "eligible"),
|
||||
(("aretrieve_batch", 0.25, True), ("retrieve_batch", 0.25, True), ("acompletion", 0.25, False),
|
||||
("aretrieve_batch", 0, False), ("aretrieve_batch", -0.25, False)),
|
||||
)
|
||||
def test_batch_charge_identity_requires_a_cost_bearing_retrieval(call_type: str, cost: float, eligible: bool) -> None:
|
||||
from openai.types import Batch
|
||||
from litellm.router_utils.batch_utils import batch_id_to_charge_once
|
||||
|
||||
batch: Final = Batch(
|
||||
id="batch-budget", completion_window="24h", created_at=1, endpoint="/v1/chat/completions",
|
||||
input_file_id="file-budget", object="batch", status="completed",
|
||||
)
|
||||
assert batch_id_to_charge_once(call_type, batch, cost) == (batch.id if eligible else None)
|
||||
assert batch_id_to_charge_once(call_type, {"id": batch.id}, cost) is None
|
||||
|
||||
|
||||
def test_batch_charge_markers_keep_counters_and_batches_independent() -> None:
|
||||
from litellm.caching.in_memory_cache import InMemoryCache
|
||||
from litellm.router_utils.batch_utils import batch_charged_once_marker_key
|
||||
|
||||
cache: Final = InMemoryCache()
|
||||
cache.set_cache(batch_charged_once_marker_key("provider-spend", "first-batch"), 1)
|
||||
assert cache.get_cache(batch_charged_once_marker_key("provider-spend", "first-batch")) == 1
|
||||
assert cache.get_cache(batch_charged_once_marker_key("deployment-spend", "first-batch")) is None
|
||||
assert cache.get_cache(batch_charged_once_marker_key("provider-spend", "second-batch")) is None
|
||||
|
|
|
|||
|
|
@ -637,7 +637,12 @@ def test_log_retry(model_list: list[DeploymentTypedDict], metadata_key: str) ->
|
|||
"model": "gpt-3.5-turbo",
|
||||
"api_key": "sk-must-not-be-recorded",
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
metadata_key: {"model_info": {"id": "deployment-1"}, "attempted_retries": 2, "user_api_key": "sk-proxy"},
|
||||
metadata_key: {
|
||||
"model_info": {"id": "deployment-1", "base_model": "azure/gpt-3.5-turbo"},
|
||||
"deployment": "azure/my-35-deployment",
|
||||
"attempted_retries": 2,
|
||||
"user_api_key": "sk-proxy",
|
||||
},
|
||||
},
|
||||
e=rate_limit_error,
|
||||
)
|
||||
|
|
@ -645,6 +650,8 @@ def test_log_retry(model_list: list[DeploymentTypedDict], metadata_key: str) ->
|
|||
{
|
||||
"model_group": "gpt-3.5-turbo",
|
||||
"deployment_id": "deployment-1",
|
||||
"model": "azure/my-35-deployment",
|
||||
"base_model": "azure/gpt-3.5-turbo",
|
||||
"exception_type": "RateLimitError",
|
||||
"exception_string": "litellm.RateLimitError: slow down",
|
||||
"attempted_retries": 2,
|
||||
|
|
@ -656,6 +663,38 @@ def test_log_retry(model_list: list[DeploymentTypedDict], metadata_key: str) ->
|
|||
assert router.log_retry(kwargs=planted_kwargs, e=rate_limit_error)[metadata_key]["request_retry_count"] == 1
|
||||
|
||||
|
||||
def test_log_retry_reads_base_model_from_the_deployments_litellm_params() -> None:
|
||||
"""An Azure named deployment declares ``base_model`` under ``litellm_params``, not ``model_info``, and
|
||||
that is the model the failure row is priced at, so the retry record has to carry it too"""
|
||||
router = Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "azure-group",
|
||||
"litellm_params": {
|
||||
"model": "azure/my-named-deployment",
|
||||
"api_key": "sk-fake",
|
||||
"api_base": "https://example.openai.azure.com",
|
||||
"base_model": "azure/gpt-4o-mini",
|
||||
},
|
||||
"model_info": {"id": "dep-azure-named"},
|
||||
}
|
||||
]
|
||||
)
|
||||
new_kwargs = router.log_retry(
|
||||
kwargs={
|
||||
"model": "azure-group",
|
||||
"metadata": {
|
||||
"model_info": {"id": "dep-azure-named"},
|
||||
"deployment": "azure/my-named-deployment",
|
||||
"attempted_retries": 0,
|
||||
},
|
||||
},
|
||||
e=litellm.Timeout(message="slow", llm_provider="azure", model="azure/my-named-deployment"),
|
||||
)
|
||||
record = new_kwargs["metadata"]["previous_models"][0]
|
||||
assert (record["model"], record["base_model"]) == ("azure/my-named-deployment", "azure/gpt-4o-mini")
|
||||
|
||||
|
||||
def test_update_usage(model_list):
|
||||
"""Test if the '_update_usage' function is working correctly"""
|
||||
router = Router(model_list=model_list)
|
||||
|
|
|
|||
|
|
@ -1210,3 +1210,39 @@ class TestCoerceNumericFormFields:
|
|||
numeric_fields=self.numeric_fields,
|
||||
)
|
||||
assert result == {"n": 3, "temperature": None, "image": buffer}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_form_data_metadata_with_lone_surrogate_escape_is_rejected_like_a_json_body():
|
||||
"""``json.loads`` accepts ``"\\ud800"`` and yields a lone surrogate that no
|
||||
provider or database can utf-8 encode. The JSON body path already answers
|
||||
400, so the multipart ``metadata`` field must too, instead of letting the
|
||||
value through to the spend log writers."""
|
||||
mock_request = MagicMock()
|
||||
mock_request.form = AsyncMock(
|
||||
return_value=FormData({"model": "whisper-1", "metadata": '{"tags": ["poison-\\ud800-tag"]}'})
|
||||
)
|
||||
mock_request.headers = {"content-type": "multipart/form-data"}
|
||||
mock_request.scope = {}
|
||||
mock_request.state._cached_headers = None
|
||||
|
||||
with pytest.raises(ProxyException) as exc_info:
|
||||
await _read_request_body(mock_request)
|
||||
|
||||
assert exc_info.value.code == "400"
|
||||
assert "metadata" in exc_info.value.message
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_form_data_metadata_keeps_non_ascii_and_valid_surrogate_pairs():
|
||||
mock_request = MagicMock()
|
||||
mock_request.form = AsyncMock(
|
||||
return_value=FormData({"model": "whisper-1", "metadata": '{"tags": ["caf\\u00e9", "\\ud83d\\ude00"]}'})
|
||||
)
|
||||
mock_request.headers = {"content-type": "multipart/form-data"}
|
||||
mock_request.scope = {}
|
||||
mock_request.state._cached_headers = None
|
||||
|
||||
result = await _read_request_body(mock_request)
|
||||
|
||||
assert result["metadata"]["tags"] == ["café", "😀"]
|
||||
|
|
|
|||
|
|
@ -3843,28 +3843,57 @@ async def test_commit_spend_updates_retries_deadlock_on_every_entity_path(monkey
|
|||
proxy_logging.failure_handler.assert_not_called()
|
||||
|
||||
|
||||
def _responses_router() -> litellm.Router:
|
||||
return litellm.Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "native",
|
||||
"litellm_params": {"model": "openai/gpt-5.1", "api_key": "sk-test"},
|
||||
"model_info": {"id": "dep-native"},
|
||||
},
|
||||
{
|
||||
"model_name": "bridged",
|
||||
"litellm_params": {"model": "anthropic/claude-sonnet-4-5", "api_key": "sk-test"},
|
||||
"model_info": {"id": "dep-bridged"},
|
||||
},
|
||||
{
|
||||
"model_name": "forced-bridge",
|
||||
"litellm_params": {"model": "openai/gpt-5.1", "api_key": "sk-test", "use_chat_completions_api": True},
|
||||
"model_info": {"id": "dep-forced-bridge"},
|
||||
},
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"call_type, expects_flush",
|
||||
[("aresponses", True), ("responses", True), ("acompletion", False)],
|
||||
"call_type, model_id, expects_flush",
|
||||
[
|
||||
("aresponses", "dep-bridged", True),
|
||||
("responses", "dep-bridged", True),
|
||||
("aresponses", "dep-forced-bridge", True),
|
||||
("aresponses", None, True),
|
||||
("aresponses", "dep-native", False),
|
||||
("responses", "dep-native", False),
|
||||
("acompletion", "dep-bridged", False),
|
||||
],
|
||||
)
|
||||
async def test_insert_spend_log_asks_for_an_immediate_flush_on_rows_other_workers_read_back(
|
||||
call_type: str, expects_flush: bool
|
||||
async def test_insert_spend_log_asks_for_an_immediate_flush_only_on_rows_the_bridge_reads_back(
|
||||
call_type: str, model_id: str | None, expects_flush: bool
|
||||
):
|
||||
"""
|
||||
A `previous_response_id` chained straight off the previous turn reads the DB, so a
|
||||
Responses row cannot sit in this worker's queue until the monitor's next poll.
|
||||
A `previous_response_id` chained straight off a bridged turn reads that turn's row from the
|
||||
DB, so it cannot sit in this worker's queue until the monitor's next poll. A native Responses
|
||||
provider holds the conversation itself, so its row waits for the batch write like a chat row.
|
||||
"""
|
||||
db_writer = DBSpendUpdateWriter()
|
||||
prisma = _tool_usage_prisma()
|
||||
prisma.spend_log_flush_requested = asyncio.Event()
|
||||
payload = {"request_id": "req-1", "call_type": call_type, "model_id": model_id}
|
||||
|
||||
await db_writer._insert_spend_log_to_db(
|
||||
payload={"request_id": "req-1", "call_type": call_type},
|
||||
prisma_client=prisma,
|
||||
)
|
||||
await db_writer._insert_spend_log_to_db(payload=payload, prisma_client=prisma, llm_router=_responses_router())
|
||||
|
||||
assert prisma.spend_log_transactions == [{"request_id": "req-1", "call_type": call_type}]
|
||||
assert prisma.spend_log_transactions == [payload]
|
||||
assert prisma.spend_log_flush_requested.is_set() is expects_flush
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -2,12 +2,15 @@ import asyncio
|
|||
import json
|
||||
import logging
|
||||
from datetime import datetime
|
||||
from types import SimpleNamespace
|
||||
from typing import Final
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
import litellm
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.litellm_core_utils.internal_call_metadata import MODEL_ACCESS_GROUP_METADATA_KEY
|
||||
from litellm.proxy._types import SpendLogsPayload, UserAPIKeyAuth
|
||||
|
|
@ -27,7 +30,7 @@ from litellm.proxy.utils import ProxyUpdateSpend
|
|||
from litellm.proxy.spend_tracking.spend_event import SpendEventDecodeError, build_spend_event, decode_spend_event
|
||||
from litellm.proxy.spend_tracking.spend_event_producer import SpendEventProducer, UnixAddress
|
||||
from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload
|
||||
from litellm.types.utils import CallTypes, LiteLLMBatch, ModelResponse, Usage
|
||||
from litellm.types.utils import CallTypes, LiteLLMBatch, ModelResponse, PromptTokensDetailsWrapper, Usage
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -93,6 +96,65 @@ async def test_async_post_call_failure_hook():
|
|||
assert metadata["original_key"] == "original_value"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_each_earlier_billed_attempt_gets_its_own_spend_row():
|
||||
"""A timed-out request that the router retried writes one row per billed attempt: the final
|
||||
attempt under the request id, every earlier Timeout attempt under ``<id>-attempt-<n>`` with its
|
||||
own cost, attempt number and deployment."""
|
||||
from litellm.proxy.utils import BilledFailedAttempt
|
||||
|
||||
logger = _ProxyDBLogger()
|
||||
user_api_key_dict = UserAPIKeyAuth(api_key="test_api_key", user_id="u1", team_id="t1")
|
||||
usage = litellm.Usage(prompt_tokens=24, completion_tokens=0, total_tokens=24)
|
||||
request_data = {
|
||||
"model": "gpt-4.1",
|
||||
"litellm_call_id": "call-1",
|
||||
"standard_logging_object": {"id": "call-1", "model_id": "dep-final"},
|
||||
"metadata": {
|
||||
"attempted_retries": 2,
|
||||
"model_info": {"id": "dep-final", "input_cost_per_token": 0.5},
|
||||
"previous_models": ({"exception_type": "Timeout"}, {"exception_type": "Timeout"}),
|
||||
},
|
||||
"combined_usage_object": usage,
|
||||
"response_cost": 0.5,
|
||||
"billed_failed_attempts": (
|
||||
BilledFailedAttempt(attempted_retries=0, deployment_id="dep-a", usage=usage, response_cost=0.25),
|
||||
BilledFailedAttempt(attempted_retries=0, deployment_id="dep-final", usage=usage, response_cost=0.5),
|
||||
),
|
||||
}
|
||||
|
||||
with patch(
|
||||
"litellm.proxy.db.db_spend_update_writer.DBSpendUpdateWriter.update_database",
|
||||
new_callable=AsyncMock,
|
||||
) as mock_update_database:
|
||||
await logger.async_post_call_failure_hook(
|
||||
request_data=request_data,
|
||||
original_exception=litellm.Timeout("Request timed out.", model="gpt-4.1", llm_provider="openai"),
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
)
|
||||
|
||||
rows = [call.kwargs for call in mock_update_database.call_args_list]
|
||||
assert [row["response_cost"] for row in rows] == [0.5, 0.25, 0.5]
|
||||
final_row, attempt_0, attempt_1 = (row["kwargs"] for row in rows)
|
||||
assert final_row["litellm_call_id"] == "call-1"
|
||||
assert final_row["litellm_params"]["metadata"]["attempted_retries"] == 2
|
||||
assert attempt_0["litellm_call_id"] == "call-1-attempt-0"
|
||||
assert attempt_0["standard_logging_object"]["id"] == "call-1-attempt-0"
|
||||
assert attempt_0["litellm_params"]["metadata"]["attempted_retries"] == 0
|
||||
assert attempt_0["litellm_params"]["metadata"]["model_info"]["id"] == "dep-a"
|
||||
assert "previous_models" not in attempt_0["litellm_params"]["metadata"]
|
||||
assert "billed_failed_attempts" not in attempt_0
|
||||
assert attempt_0["combined_usage_object"] is usage
|
||||
assert attempt_1["litellm_call_id"] == "call-1-attempt-1"
|
||||
assert attempt_1["litellm_params"]["metadata"]["attempted_retries"] == 0
|
||||
assert attempt_1["litellm_params"]["metadata"]["model_info"]["id"] == "dep-final"
|
||||
assert final_row["litellm_params"]["metadata"]["model_info"]["id"] == "dep-final"
|
||||
for row in rows:
|
||||
assert row["token"] == "test_api_key"
|
||||
assert row["user_id"] == "u1"
|
||||
assert row["team_id"] == "t1"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_post_call_failure_hook_carries_guardrail_info_from_litellm_metadata():
|
||||
"""
|
||||
|
|
@ -2712,3 +2774,165 @@ async def test_track_cost_callback_failure_alert_never_carries_request_metadata_
|
|||
assert "headers" in failure_debug_lines[0]
|
||||
else:
|
||||
assert failure_debug_lines == []
|
||||
|
||||
|
||||
_BATCH_CALL_START: Final = datetime(2026, 1, 1)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_batch_accounting_failures_stay_with_their_originating_call() -> None:
|
||||
from litellm.proxy.hooks.proxy_track_cost_callback import track_batch_accounting
|
||||
|
||||
logger: Final = _ProxyDBLogger()
|
||||
dependency: Final = SimpleNamespace(failed_tracking_alert=AsyncMock())
|
||||
|
||||
async def fail_before_accounting(call_id: str) -> None:
|
||||
await logger.async_log_success_event(
|
||||
kwargs={
|
||||
"litellm_call_id": call_id,
|
||||
"call_type": "aretrieve_batch",
|
||||
"model": "batch-model",
|
||||
"stream": False,
|
||||
"response_cost": None,
|
||||
},
|
||||
response_obj=None,
|
||||
start_time=_BATCH_CALL_START,
|
||||
end_time=_BATCH_CALL_START,
|
||||
)
|
||||
|
||||
async def independent_batch() -> bool:
|
||||
with track_batch_accounting("independent") as outcome:
|
||||
await fail_before_accounting("independent")
|
||||
return outcome.retryable_failure
|
||||
|
||||
with patch("litellm.proxy.proxy_server.proxy_logging_obj", dependency):
|
||||
with track_batch_accounting("outer") as outer:
|
||||
await fail_before_accounting("unrelated")
|
||||
assert outer.retryable_failure is False
|
||||
with track_batch_accounting("inner") as inner:
|
||||
await fail_before_accounting("inner")
|
||||
assert inner.retryable_failure is True
|
||||
assert outer.retryable_failure is False
|
||||
independent, _ = await asyncio.gather(independent_batch(), fail_before_accounting("outer"))
|
||||
assert independent is True
|
||||
assert outer.retryable_failure is True
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
with track_batch_accounting("cancelled") as cancelled:
|
||||
raise asyncio.CancelledError
|
||||
await fail_before_accounting("cancelled")
|
||||
assert cancelled.retryable_failure is False
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_batch_accounting_does_not_retry_after_settling_a_budget_reservation() -> None:
|
||||
from litellm.proxy.hooks.proxy_track_cost_callback import track_batch_accounting
|
||||
|
||||
reservation: Final[dict[str, object]] = {"reserved_cost": 0.5, "entries": ()}
|
||||
dependency: Final = SimpleNamespace(failed_tracking_alert=AsyncMock())
|
||||
with patch("litellm.proxy.proxy_server.proxy_logging_obj", dependency):
|
||||
with track_batch_accounting("reserved") as outcome:
|
||||
await _ProxyDBLogger().async_log_success_event(
|
||||
kwargs={
|
||||
"litellm_call_id": "reserved",
|
||||
"call_type": "aretrieve_batch",
|
||||
"model": "batch-model",
|
||||
"stream": False,
|
||||
"response_cost": None,
|
||||
"litellm_params": {"metadata": {"user_api_key_budget_reservation": reservation}},
|
||||
},
|
||||
response_obj=None,
|
||||
start_time=_BATCH_CALL_START,
|
||||
end_time=_BATCH_CALL_START,
|
||||
)
|
||||
assert reservation["finalized"] is True
|
||||
assert outcome.retryable_failure is False
|
||||
|
||||
|
||||
class _SharedCounter:
|
||||
def __init__(self) -> None:
|
||||
self.totals: dict[str, float] = {}
|
||||
|
||||
async def async_increment(self, key: str, value: float, ttl: int | None = None) -> float:
|
||||
self.totals[key] = self.totals.get(key, 0) + value
|
||||
return self.totals[key]
|
||||
|
||||
|
||||
def _grounded_gemini_kwargs(response_cost: float) -> dict:
|
||||
return {
|
||||
"model": "vertex_ai/gemini-test",
|
||||
"call_type": "acompletion",
|
||||
"response_cost": response_cost,
|
||||
"litellm_params": {"metadata": {"user_api_key": "sk-test", "user_api_key_user_id": "user-1"}},
|
||||
"standard_logging_object": {
|
||||
"response_cost": response_cost,
|
||||
"request_tags": None,
|
||||
"cache_hit": None,
|
||||
"custom_llm_provider": "vertex_ai",
|
||||
"cost_breakdown": {
|
||||
"input_cost": 0.0,
|
||||
"output_cost": 0.0,
|
||||
"tool_usage_cost": 0.5,
|
||||
"total_cost": response_cost,
|
||||
},
|
||||
"model_map_information": {
|
||||
"model_map_key": "vertex_ai/gemini-test",
|
||||
"model_map_value": {
|
||||
"litellm_provider": "vertex_ai",
|
||||
"mode": "chat",
|
||||
"web_search_billing_unit": "per_query",
|
||||
"search_context_cost_per_query": {"search_context_size_medium": 0.1},
|
||||
"search_free_requests": {"count": 100, "period": "month", "shared_group": "test"},
|
||||
},
|
||||
},
|
||||
},
|
||||
"stream": False,
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("opted_in", [True, False])
|
||||
async def test_track_cost_callback_frees_grounding_under_the_allowance_only_when_opted_in(monkeypatch, opted_in):
|
||||
monkeypatch.setattr(litellm, "apply_search_free_tier", opted_in, raising=False)
|
||||
counter = _SharedCounter()
|
||||
kwargs = _grounded_gemini_kwargs(response_cost=1.0)
|
||||
response = ModelResponse(
|
||||
usage=Usage(
|
||||
prompt_tokens=10,
|
||||
completion_tokens=5,
|
||||
total_tokens=15,
|
||||
prompt_tokens_details=PromptTokensDetailsWrapper(web_search_requests=5),
|
||||
)
|
||||
)
|
||||
|
||||
with (
|
||||
patch( # test-quality-ok: increment_spend_counters is a proxy_server global the callback reads lazily, no seam
|
||||
"litellm.proxy.proxy_server.increment_spend_counters", new_callable=AsyncMock
|
||||
),
|
||||
patch( # test-quality-ok: update_cache is a proxy_server global the callback reads lazily, no seam
|
||||
"litellm.proxy.proxy_server.update_cache", new_callable=AsyncMock
|
||||
),
|
||||
patch( # test-quality-ok: callback imports proxy_logging_obj off proxy_server in its body, no seam
|
||||
"litellm.proxy.proxy_server.proxy_logging_obj"
|
||||
) as mock_proxy_logging,
|
||||
):
|
||||
mock_proxy_logging.failed_tracking_alert = AsyncMock()
|
||||
mock_proxy_logging.db_spend_update_writer.update_database = AsyncMock(return_value=True)
|
||||
mock_proxy_logging.slack_alerting_instance.customer_spend_alert = AsyncMock()
|
||||
mock_proxy_logging.internal_usage_cache.dual_cache.redis_cache = counter
|
||||
|
||||
await _ProxyDBLogger()._PROXY_track_cost_callback(
|
||||
kwargs=kwargs, completion_response=response, start_time=datetime.now(), end_time=datetime.now()
|
||||
)
|
||||
await asyncio.sleep(0)
|
||||
|
||||
mock_proxy_logging.failed_tracking_alert.assert_not_called()
|
||||
mock_proxy_logging.db_spend_update_writer.update_database.assert_awaited_once()
|
||||
|
||||
breakdown = kwargs["standard_logging_object"]["cost_breakdown"]
|
||||
if opted_in:
|
||||
assert kwargs["response_cost"] == pytest.approx(1.0 - 5 * 0.1)
|
||||
assert breakdown["original_cost"] == 1.0 and breakdown["discount_amount"] == pytest.approx(5 * 0.1)
|
||||
assert sum(counter.totals.values()) == 5
|
||||
else:
|
||||
assert kwargs["response_cost"] == 1.0 and "discount_amount" not in breakdown
|
||||
assert counter.totals == {}
|
||||
|
|
|
|||
|
|
@ -1249,6 +1249,32 @@ def test_resolve_llm_passthrough_timeout_validates_only_the_winning_value(
|
|||
assert resolve_llm_passthrough_timeout(kwargs=kwargs, litellm_params=litellm_params) == expected
|
||||
|
||||
|
||||
def test_resolve_llm_passthrough_timeout_uses_the_deployment_timeout_for_the_requests_service_tier():
|
||||
litellm_params = {
|
||||
"timeout": 610,
|
||||
"stream_timeout": 30,
|
||||
"timeout_by_service_tier": {"flex": 900},
|
||||
"stream_timeout_by_service_tier": {"flex": 950},
|
||||
}
|
||||
assert resolve_llm_passthrough_timeout(kwargs={"service_tier": "flex"}, litellm_params=litellm_params) == 900.0
|
||||
assert resolve_llm_passthrough_timeout(kwargs={"service_tier": "default"}, litellm_params=litellm_params) == 610.0
|
||||
assert (
|
||||
resolve_llm_passthrough_timeout(kwargs={"service_tier": "flex", "timeout": 20}, litellm_params=litellm_params)
|
||||
== 20.0
|
||||
)
|
||||
assert (
|
||||
resolve_llm_passthrough_timeout(kwargs={"service_tier": "flex", "stream": True}, litellm_params=litellm_params)
|
||||
== 950.0
|
||||
)
|
||||
assert (
|
||||
resolve_llm_passthrough_timeout(
|
||||
kwargs={"service_tier": "flex", "stream": True},
|
||||
litellm_params={"timeout": 610, "stream_timeout": 30, "timeout_by_service_tier": {"flex": 900}},
|
||||
)
|
||||
== 30.0
|
||||
)
|
||||
|
||||
|
||||
def test_resolve_llm_passthrough_timeout_rejects_a_non_numeric_winner():
|
||||
with pytest.raises(ValidationError):
|
||||
resolve_llm_passthrough_timeout(kwargs={"timeout": httpx.Timeout(30.0)})
|
||||
|
|
|
|||
|
|
@ -7637,6 +7637,118 @@ def _apply_stream_usage_tracking(
|
|||
)
|
||||
|
||||
|
||||
async def _start_partial_responses_stream(self, recorder: _RecordingSuccessLogger, *, relay_text: bool = True):
|
||||
from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig
|
||||
from litellm.responses.streaming_iterator import ResponsesAPIStreamingIterator
|
||||
|
||||
mocked: Final = await litellm.aresponses(
|
||||
model="gpt-4o-mini",
|
||||
input="tell me a story",
|
||||
mock_response="The codename is AZURE-FALCON-42 and the story is long.",
|
||||
stream=True,
|
||||
api_key="test-key",
|
||||
)
|
||||
logging_obj: Final = mocked.logging_obj
|
||||
logging_obj.dynamic_async_success_callbacks = [recorder]
|
||||
stream: Final = ResponsesAPIStreamingIterator(
|
||||
response=httpx.Response(200),
|
||||
model="gpt-4o-mini",
|
||||
responses_api_provider_config=OpenAIResponsesAPIConfig(),
|
||||
logging_obj=logging_obj,
|
||||
request_data={"model": "gpt-4o-mini", "input": "tell me a story"},
|
||||
custom_llm_provider="openai",
|
||||
)
|
||||
created: Final = {
|
||||
"type": "response.created",
|
||||
"sequence_number": 0,
|
||||
"response": {
|
||||
"id": "resp_partial",
|
||||
"object": "response",
|
||||
"created_at": 1700000000.0,
|
||||
"model": "gpt-4o-mini",
|
||||
"status": "in_progress",
|
||||
"output": [],
|
||||
"parallel_tool_calls": True,
|
||||
"tool_choice": "auto",
|
||||
"tools": [],
|
||||
"service_tier": "priority",
|
||||
},
|
||||
}
|
||||
delta: Final = {
|
||||
"type": "response.output_text.delta",
|
||||
"sequence_number": 1,
|
||||
"item_id": "msg_partial",
|
||||
"output_index": 0,
|
||||
"content_index": 0,
|
||||
"delta": "The codename is AZURE-FALCON-42",
|
||||
}
|
||||
for event in (created, delta) if relay_text else (created,):
|
||||
stream._process_chunk(json.dumps(event))
|
||||
return stream
|
||||
|
||||
async def _bill_and_wait(self, recorder: _RecordingSuccessLogger, streamed: object, logging_obj: object) -> bool:
|
||||
billed: Final = await _bill_partial_streamed_spend_on_disconnect({"litellm_logging_obj": logging_obj}, streamed)
|
||||
for _ in range(50):
|
||||
if recorder.success_events:
|
||||
break
|
||||
await asyncio.sleep(0.1)
|
||||
await asyncio.sleep(0.5)
|
||||
return billed
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_disconnect_bills_partial_native_responses_stream_with_counted_usage(self):
|
||||
from litellm.types.llms.openai import ResponsesAPIResponse
|
||||
|
||||
recorder: Final = _RecordingSuccessLogger()
|
||||
stream: Final = await self._start_partial_responses_stream(recorder)
|
||||
billed: Final = await self._bill_and_wait(recorder, stream, stream.logging_obj)
|
||||
|
||||
assert billed is True
|
||||
assert len(recorder.success_events) == 1
|
||||
partial_response: Final = recorder.success_events[0]["response_obj"]
|
||||
assert isinstance(partial_response, ResponsesAPIResponse)
|
||||
assert partial_response.status == "incomplete"
|
||||
assert getattr(partial_response, "service_tier") == "priority"
|
||||
assert partial_response.usage.output_tokens == litellm.token_counter(
|
||||
model="gpt-4o-mini", text="The codename is AZURE-FALCON-42"
|
||||
)
|
||||
assert partial_response.usage.input_tokens > 0
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_disconnect_bills_partial_default_route_messages_stream_through_the_responses_adapter(self):
|
||||
from litellm.llms.anthropic.pass_through.responses_adapters.streaming_iterator import (
|
||||
AnthropicResponsesStreamWrapper,
|
||||
)
|
||||
from litellm.router import FallbackAwareAnthropicMessagesStream
|
||||
from litellm.types.llms.openai import ResponsesAPIResponse
|
||||
|
||||
async def _sse_frames() -> AsyncGenerator[bytes, None]:
|
||||
yield b"event: message_start\n\n"
|
||||
|
||||
recorder: Final = _RecordingSuccessLogger()
|
||||
stream: Final = await self._start_partial_responses_stream(recorder)
|
||||
wrapper: Final = AnthropicResponsesStreamWrapper(responses_stream=stream, model="gpt-4o-mini")
|
||||
streamed: Final = FallbackAwareAnthropicMessagesStream(_sse_frames(), wrapper.async_anthropic_sse_wrapper())
|
||||
billed: Final = await self._bill_and_wait(recorder, streamed, stream.logging_obj)
|
||||
|
||||
assert billed is True
|
||||
assert len(recorder.success_events) == 1
|
||||
partial_response: Final = recorder.success_events[0]["response_obj"]
|
||||
assert isinstance(partial_response, ResponsesAPIResponse)
|
||||
assert partial_response.usage.output_tokens == litellm.token_counter(
|
||||
model="gpt-4o-mini", text="The codename is AZURE-FALCON-42"
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_disconnect_before_any_responses_output_bills_nothing(self):
|
||||
recorder: Final = _RecordingSuccessLogger()
|
||||
stream: Final = await self._start_partial_responses_stream(recorder, relay_text=False)
|
||||
billed: Final = await self._bill_and_wait(recorder, stream, stream.logging_obj)
|
||||
|
||||
assert billed is False
|
||||
assert recorder.success_events == []
|
||||
|
||||
|
||||
class TestApplyStreamUsageTracking:
|
||||
def test_default_injects_usage_and_marks_strip_for_chat_completions(self):
|
||||
data = {"stream": True, "model": "gpt-5.4-nano"}
|
||||
|
|
|
|||
|
|
@ -7764,6 +7764,7 @@ _PLANTED_STAMPS = {
|
|||
"attempted_fallbacks": 99,
|
||||
"original_model_group": "spoofed-group",
|
||||
"request_retry_count": -100,
|
||||
"previous_models": [{"exception_type": "Timeout", "deployment_id": "spoofed", "attempted_retries": 101}],
|
||||
"_client_output_ceiling": {"api_base": "https://attacker.example"},
|
||||
ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY: 10**9,
|
||||
"client_key": "client_value",
|
||||
|
|
@ -7798,6 +7799,7 @@ async def test_add_litellm_data_to_request_strips_router_reserved_stamps_from_bo
|
|||
assert "original_model_group" not in updated["metadata"]
|
||||
assert "_client_output_ceiling" not in updated["metadata"]
|
||||
assert "request_retry_count" not in updated["metadata"]
|
||||
assert "previous_models" not in updated["metadata"]
|
||||
assert ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY not in updated["metadata"]
|
||||
assert updated["metadata"]["client_key"] == "client_value"
|
||||
|
||||
|
|
|
|||
|
|
@ -58,9 +58,7 @@ class TestStripClientPricingOverrides:
|
|||
# The strip set is built from the model so additions are picked up
|
||||
# automatically — this test guards against the model and the strip
|
||||
# set drifting apart if someone replaces the auto-derivation later.
|
||||
assert _CLIENT_PRICING_CONTROL_FIELDS == frozenset(
|
||||
CustomPricingLiteLLMParams.model_fields.keys()
|
||||
)
|
||||
assert _CLIENT_PRICING_CONTROL_FIELDS >= frozenset(CustomPricingLiteLLMParams.model_fields.keys())
|
||||
# Sanity: the obvious top-level pricing fields are in the set.
|
||||
for field in (
|
||||
"input_cost_per_token",
|
||||
|
|
@ -84,6 +82,22 @@ class TestStripClientPricingOverrides:
|
|||
"messages": [{"role": "user", "content": "hi"}],
|
||||
}
|
||||
|
||||
def test_root_base_model_dropped(self):
|
||||
# ``base_model`` picks the cost-map entry a spend row is priced with
|
||||
# (success rows via the standard logging payload, billed failure rows
|
||||
# via the failure lift), so a caller naming a cheaper model lowers
|
||||
# their own recorded spend.
|
||||
data = {
|
||||
"model": "gpt-4o",
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"base_model": "gpt-4o-mini",
|
||||
}
|
||||
_strip_client_pricing_overrides(data)
|
||||
assert data == {
|
||||
"model": "gpt-4o",
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
}
|
||||
|
||||
def test_metadata_model_info_dropped(self):
|
||||
data = {
|
||||
"model": "gpt-4",
|
||||
|
|
@ -264,6 +278,7 @@ async def test_add_litellm_data_to_request_skips_strip_with_key_opt_in():
|
|||
"model": "gpt-4",
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"input_cost_per_token": 0.0001,
|
||||
"base_model": "gpt-4o-mini",
|
||||
"metadata": {"model_info": {"output_cost_per_token": 0.0002}},
|
||||
}
|
||||
|
||||
|
|
@ -279,6 +294,7 @@ async def test_add_litellm_data_to_request_skips_strip_with_key_opt_in():
|
|||
|
||||
assert updated["input_cost_per_token"] == 0.0001
|
||||
assert updated["metadata"]["model_info"] == {"output_cost_per_token": 0.0002}
|
||||
assert updated["base_model"] == "gpt-4o-mini"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
|
|||
|
|
@ -2,6 +2,7 @@ import datetime as real_datetime
|
|||
import smtplib
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from fastapi import HTTPException
|
||||
|
||||
|
|
@ -443,6 +444,47 @@ class TestPostCallFailureHookLiftsFirstApiCallStartTime:
|
|||
assert "litellm_logging_obj" not in request_data
|
||||
|
||||
|
||||
class TestPostCallFailureHookDoesNotInventAProviderHandoff:
|
||||
"""A failure the proxy raised before any provider call (here an auth error) is logged through
|
||||
_handle_logging_proxy_only_error, which fires pre_call for the callbacks. That must not stamp a
|
||||
handoff instant onto the request, or every downstream reader (preprocessing latency, the
|
||||
dispatched check behind the failure-row estimate) sees a provider call that never happened."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_proxy_only_failure_lifts_no_first_api_call_start_time(self):
|
||||
from datetime import datetime
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging
|
||||
from litellm.proxy._types import ProxyErrorTypes, UserAPIKeyAuth
|
||||
|
||||
logging_obj = Logging(
|
||||
model="gpt-4.1",
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
stream=False,
|
||||
call_type="acompletion",
|
||||
start_time=datetime.now(),
|
||||
litellm_call_id="pre-handoff-auth-failure",
|
||||
function_id="pre-handoff",
|
||||
)
|
||||
request_data: Final[dict[str, object]] = {"litellm_logging_obj": logging_obj, "model": "gpt-4.1", "metadata": {}}
|
||||
proxy_logging_obj = ProxyLogging(user_api_key_cache=DualCache())
|
||||
proxy_logging_obj.alert_types = []
|
||||
with (
|
||||
patch.object(proxy_logging_obj, "update_request_status", new=AsyncMock()),
|
||||
patch.object(proxy_logging_obj, "_dispatch_proxy_only_failure_handlers", new=AsyncMock()),
|
||||
):
|
||||
await proxy_logging_obj.post_call_failure_hook(
|
||||
request_data=request_data,
|
||||
original_exception=Exception("Authentication Error, invalid key"),
|
||||
user_api_key_dict=UserAPIKeyAuth(request_route="/v1/chat/completions"),
|
||||
error_type=ProxyErrorTypes.auth_error,
|
||||
)
|
||||
|
||||
assert request_data.get("first_api_call_start_time") is None
|
||||
assert "api_call_start_time" not in logging_obj.model_call_details
|
||||
|
||||
|
||||
class TestPostCallFailureHookLiftsRecoveredPartialSpend:
|
||||
"""A stream that broke mid-flight still billed the provider for the chunks
|
||||
already delivered. The streaming handler stashes that recovered usage and
|
||||
|
|
@ -687,11 +729,13 @@ class TestPostCallFailureHookEstimatesDispatchedInputTokens:
|
|||
"""A non-stream request that failed after dispatch (timeout, provider
|
||||
error) consumed provider-billed input tokens but recovered no usage.
|
||||
post_call_failure_hook must estimate the input side onto request_data so
|
||||
the spend log's failure row records what was sent instead of zero, while
|
||||
never charging spend for the failure (LIT-5690).
|
||||
the spend log's failure row records what was sent instead of zero
|
||||
(LIT-5690). On a timeout that estimate is priced at the model's input
|
||||
rate since the provider bills the prompt it accepted (LIT-145); other
|
||||
post-dispatch failures keep zero spend.
|
||||
"""
|
||||
|
||||
async def _run(self, request_data):
|
||||
async def _run(self, request_data, original_exception: Exception | None = None):
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
|
|
@ -701,17 +745,275 @@ class TestPostCallFailureHookEstimatesDispatchedInputTokens:
|
|||
with patch.object(proxy_logging_obj, "update_request_status", new=AsyncMock()):
|
||||
await proxy_logging_obj.post_call_failure_hook(
|
||||
request_data=request_data,
|
||||
original_exception=Exception("boom"),
|
||||
original_exception=original_exception or Exception("boom"),
|
||||
user_api_key_dict=UserAPIKeyAuth(),
|
||||
)
|
||||
|
||||
def _timeout(self, model: str, cause: Exception) -> Exception:
|
||||
try:
|
||||
try:
|
||||
raise cause
|
||||
except Exception as inner:
|
||||
raise litellm.Timeout("Request timed out.", model=model, llm_provider="openai") from inner
|
||||
except litellm.Timeout as timed_out:
|
||||
return timed_out
|
||||
|
||||
def _logging_obj(self, model_call_details):
|
||||
logging_obj = MagicMock()
|
||||
logging_obj.model_call_details = model_call_details
|
||||
return logging_obj
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_dispatched_failure_estimates_input_tokens_with_zero_cost(self):
|
||||
async def test_timed_out_request_estimates_input_tokens_priced_at_input_rate(self):
|
||||
from datetime import datetime
|
||||
|
||||
from litellm.types.utils import Usage
|
||||
|
||||
request_data = {
|
||||
"litellm_logging_obj": self._logging_obj(
|
||||
{
|
||||
"first_api_call_start_time": datetime.now(),
|
||||
"model": "gpt-3.5-turbo",
|
||||
"messages": [{"role": "user", "content": "count these input tokens please"}],
|
||||
"call_type": "acompletion",
|
||||
}
|
||||
),
|
||||
"metadata": {},
|
||||
"response_cost": 123.0,
|
||||
}
|
||||
await self._run(
|
||||
request_data, original_exception=self._timeout("gpt-3.5-turbo", httpx.ReadTimeout("no bytes in time"))
|
||||
)
|
||||
|
||||
estimated = request_data["combined_usage_object"]
|
||||
assert isinstance(estimated, Usage)
|
||||
assert estimated.prompt_tokens > 0
|
||||
assert estimated.completion_tokens == 0
|
||||
assert estimated.total_tokens == estimated.prompt_tokens
|
||||
expected_input_cost, _ = litellm.cost_per_token(model="gpt-3.5-turbo", prompt_tokens=estimated.prompt_tokens)
|
||||
assert expected_input_cost > 0
|
||||
assert request_data["response_cost"] == pytest.approx(expected_input_cost)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_client_disconnect_after_dispatch_is_priced_at_input_rate(self):
|
||||
from datetime import datetime
|
||||
|
||||
from fastapi import HTTPException
|
||||
|
||||
from litellm.constants import LITELLM_HTTP_STATUS_CLIENT_DISCONNECTED, LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL
|
||||
from litellm.types.utils import Usage
|
||||
|
||||
request_data = {
|
||||
"litellm_logging_obj": self._logging_obj(
|
||||
{
|
||||
"first_api_call_start_time": datetime.now(),
|
||||
"model": "gpt-3.5-turbo",
|
||||
"messages": [{"role": "user", "content": "count these input tokens please"}],
|
||||
"call_type": "acompletion",
|
||||
LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL: True,
|
||||
}
|
||||
),
|
||||
"metadata": {},
|
||||
}
|
||||
await self._run(
|
||||
request_data,
|
||||
original_exception=HTTPException(
|
||||
status_code=LITELLM_HTTP_STATUS_CLIENT_DISCONNECTED, detail="Client disconnected the request"
|
||||
),
|
||||
)
|
||||
|
||||
estimated = request_data["combined_usage_object"]
|
||||
assert isinstance(estimated, Usage)
|
||||
expected_input_cost, _ = litellm.cost_per_token(model="gpt-3.5-turbo", prompt_tokens=estimated.prompt_tokens)
|
||||
assert expected_input_cost > 0
|
||||
assert request_data["response_cost"] == pytest.approx(expected_input_cost)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_upstream_499_error_is_a_plain_failure_not_a_priced_disconnect(self):
|
||||
"""A provider answering HTTP 499 surfaces as litellm.APIError with status_code 499. Only the
|
||||
proxy's own disconnect HTTPException means the caller hung up mid-generation, so the upstream
|
||||
error keeps the estimated usage at zero cost like any other post-dispatch provider error."""
|
||||
from datetime import datetime
|
||||
|
||||
from litellm.constants import LITELLM_HTTP_STATUS_CLIENT_DISCONNECTED
|
||||
from litellm.types.utils import Usage
|
||||
|
||||
request_data = {
|
||||
"litellm_logging_obj": self._logging_obj(
|
||||
{
|
||||
"first_api_call_start_time": datetime(2026, 1, 1, 12, 0, 0),
|
||||
"model": "gpt-3.5-turbo",
|
||||
"messages": [{"role": "user", "content": "count these input tokens please"}],
|
||||
"call_type": "acompletion",
|
||||
}
|
||||
),
|
||||
"metadata": {},
|
||||
}
|
||||
await self._run(
|
||||
request_data,
|
||||
original_exception=litellm.APIError(
|
||||
status_code=LITELLM_HTTP_STATUS_CLIENT_DISCONNECTED,
|
||||
message="client closed request",
|
||||
llm_provider="openai",
|
||||
model="gpt-3.5-turbo",
|
||||
),
|
||||
)
|
||||
|
||||
estimated = request_data["combined_usage_object"]
|
||||
assert isinstance(estimated, Usage)
|
||||
assert estimated.prompt_tokens > 0
|
||||
assert request_data["response_cost"] == 0.0
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_client_disconnect_after_the_call_already_completed_is_not_priced_again(self):
|
||||
"""The LLM task finished and claimed its success log (which carries the real usage and cost)
|
||||
while a during_call guardrail task was still pending, then the caller hung up. The 499 raised
|
||||
for that disconnect must not price the prompt a second time."""
|
||||
from datetime import datetime
|
||||
|
||||
from fastapi import HTTPException
|
||||
|
||||
from litellm.constants import LITELLM_HTTP_STATUS_CLIENT_DISCONNECTED
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging
|
||||
|
||||
logging_obj = Logging(
|
||||
model="gpt-3.5-turbo",
|
||||
messages=[{"role": "user", "content": "count these input tokens please"}],
|
||||
stream=False,
|
||||
call_type="acompletion",
|
||||
start_time=datetime(2026, 1, 1, 12, 0, 0),
|
||||
litellm_call_id="completed-before-disconnect",
|
||||
function_id="fn",
|
||||
)
|
||||
logging_obj.model_call_details.update(
|
||||
{
|
||||
"first_api_call_start_time": datetime(2026, 1, 1, 12, 0, 1),
|
||||
"model": "gpt-3.5-turbo",
|
||||
"messages": [{"role": "user", "content": "count these input tokens please"}],
|
||||
"call_type": "acompletion",
|
||||
}
|
||||
)
|
||||
logging_obj.mark_async_success_log_enqueued()
|
||||
request_data = {"litellm_logging_obj": logging_obj, "metadata": {}}
|
||||
await self._run(
|
||||
request_data,
|
||||
original_exception=HTTPException(
|
||||
status_code=LITELLM_HTTP_STATUS_CLIENT_DISCONNECTED, detail="Client disconnected the request"
|
||||
),
|
||||
)
|
||||
|
||||
assert request_data.get("response_cost", 0.0) == 0.0
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_client_disconnect_with_a_claimed_but_deferred_success_log_is_still_priced(self):
|
||||
"""With a post_call guardrail registered the success log is claimed but parked in the deferred slot,
|
||||
and a 499 raised before the flush means it never runs. Nothing else bills that request, so the
|
||||
failure row keeps the estimated input cost."""
|
||||
from datetime import datetime
|
||||
|
||||
from fastapi import HTTPException
|
||||
|
||||
from litellm.constants import LITELLM_HTTP_STATUS_CLIENT_DISCONNECTED
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging
|
||||
|
||||
logging_obj = Logging(
|
||||
model="gpt-3.5-turbo",
|
||||
messages=[{"role": "user", "content": "count these input tokens please"}],
|
||||
stream=False,
|
||||
call_type="acompletion",
|
||||
start_time=datetime(2026, 1, 1, 12, 0, 0),
|
||||
litellm_call_id="deferred-success-never-flushed",
|
||||
function_id="fn",
|
||||
)
|
||||
logging_obj.model_call_details.update(
|
||||
{
|
||||
"first_api_call_start_time": datetime(2026, 1, 1, 12, 0, 1),
|
||||
"model": "gpt-3.5-turbo",
|
||||
"messages": [{"role": "user", "content": "count these input tokens please"}],
|
||||
"call_type": "acompletion",
|
||||
}
|
||||
)
|
||||
logging_obj._defer_async_logging = True
|
||||
logging_obj._enqueue_deferred_logging = lambda: None
|
||||
request_data = {"litellm_logging_obj": logging_obj, "metadata": {}}
|
||||
await self._run(
|
||||
request_data,
|
||||
original_exception=HTTPException(
|
||||
status_code=LITELLM_HTTP_STATUS_CLIENT_DISCONNECTED, detail="Client disconnected the request"
|
||||
),
|
||||
)
|
||||
|
||||
assert request_data["response_cost"] > 0.0
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_client_disconnect_before_provider_handoff_is_not_priced(self):
|
||||
"""A caller that hung up while the request was still in pre-call hooks never reached the
|
||||
provider, so there is no handoff instant on the logging object and the 499 row is not priced."""
|
||||
from datetime import datetime
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
from fastapi import HTTPException
|
||||
|
||||
from litellm.constants import LITELLM_HTTP_STATUS_CLIENT_DISCONNECTED
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
|
||||
logging_obj = Logging(
|
||||
model="gpt-3.5-turbo",
|
||||
messages=[{"role": "user", "content": "count these input tokens please"}],
|
||||
stream=False,
|
||||
call_type="acompletion",
|
||||
start_time=datetime.now(),
|
||||
litellm_call_id="lit145-pre-handoff",
|
||||
function_id="lit145",
|
||||
)
|
||||
assert logging_obj.model_call_details.get("first_api_call_start_time") is None
|
||||
request_data = {"litellm_logging_obj": logging_obj, "model": "gpt-3.5-turbo", "metadata": {}}
|
||||
|
||||
proxy_logging_obj = ProxyLogging(user_api_key_cache=DualCache())
|
||||
proxy_logging_obj.alert_types = []
|
||||
with (
|
||||
patch.object(proxy_logging_obj, "update_request_status", new=AsyncMock()),
|
||||
patch.object(proxy_logging_obj, "_dispatch_proxy_only_failure_handlers", new=AsyncMock()),
|
||||
):
|
||||
await proxy_logging_obj.post_call_failure_hook(
|
||||
request_data=request_data,
|
||||
original_exception=HTTPException(
|
||||
status_code=LITELLM_HTTP_STATUS_CLIENT_DISCONNECTED, detail="Client disconnected the request"
|
||||
),
|
||||
user_api_key_dict=UserAPIKeyAuth(request_route="/v1/chat/completions"),
|
||||
)
|
||||
|
||||
assert request_data.get("first_api_call_start_time") is None
|
||||
assert "combined_usage_object" not in request_data
|
||||
assert request_data.get("response_cost", 0.0) == 0.0
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_connect_timeout_never_reached_the_provider_stays_zero_cost(self):
|
||||
from datetime import datetime
|
||||
|
||||
from litellm.types.utils import Usage
|
||||
|
||||
request_data = {
|
||||
"litellm_logging_obj": self._logging_obj(
|
||||
{
|
||||
"first_api_call_start_time": datetime.now(),
|
||||
"model": "gpt-3.5-turbo",
|
||||
"messages": [{"role": "user", "content": "count these input tokens please"}],
|
||||
"call_type": "acompletion",
|
||||
}
|
||||
),
|
||||
"metadata": {},
|
||||
}
|
||||
await self._run(
|
||||
request_data, original_exception=self._timeout("gpt-3.5-turbo", httpx.ConnectTimeout("no route"))
|
||||
)
|
||||
|
||||
assert isinstance(request_data["combined_usage_object"], Usage)
|
||||
assert request_data["response_cost"] == 0.0
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_dispatched_non_timeout_failure_estimates_input_tokens_with_zero_cost(self):
|
||||
from datetime import datetime
|
||||
|
||||
from litellm.types.utils import Usage
|
||||
|
|
@ -733,8 +1035,6 @@ class TestPostCallFailureHookEstimatesDispatchedInputTokens:
|
|||
estimated = request_data["combined_usage_object"]
|
||||
assert isinstance(estimated, Usage)
|
||||
assert estimated.prompt_tokens > 0
|
||||
assert estimated.completion_tokens == 0
|
||||
assert estimated.total_tokens == estimated.prompt_tokens
|
||||
assert request_data["response_cost"] == 0.0
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -1901,7 +2201,355 @@ def test_usage_the_stream_already_recovered_beats_an_estimate():
|
|||
assert lifted == (recovered, 0.25)
|
||||
|
||||
|
||||
def test_a_request_that_reached_a_provider_bills_its_input_at_no_cost():
|
||||
def test_a_timed_out_request_on_an_unmapped_model_keeps_its_row_at_zero_cost():
|
||||
from litellm.proxy.utils import _failure_usage_to_lift
|
||||
|
||||
lifted = _failure_usage_to_lift(
|
||||
model_call_details={
|
||||
"call_type": "acompletion",
|
||||
"model": "my-unmapped-model",
|
||||
"custom_llm_provider": "openai",
|
||||
"messages": ONE_USER_MESSAGE,
|
||||
},
|
||||
request_body={},
|
||||
dispatched=True,
|
||||
provider_billed=True,
|
||||
)
|
||||
|
||||
assert lifted is not None
|
||||
usage, response_cost = lifted
|
||||
assert usage.prompt_tokens > 0
|
||||
assert response_cost == 0.0
|
||||
|
||||
|
||||
def test_a_timed_out_request_is_priced_at_the_rate_the_router_registered_for_the_deployment(monkeypatch):
|
||||
"""Router._create_deployment registers a deployment's own ``input_cost_per_token`` under its
|
||||
``model_info.id``; the failure row prices with that entry through the same
|
||||
``deployment_pricing_model_info`` the success path uses."""
|
||||
from litellm.proxy.utils import _failure_usage_to_lift
|
||||
|
||||
monkeypatch.setitem(
|
||||
litellm.model_cost,
|
||||
"dep-lit145",
|
||||
{"litellm_provider": "openai", "mode": "chat", "input_cost_per_token": 0.25, "output_cost_per_token": 1.0},
|
||||
)
|
||||
|
||||
lifted = _failure_usage_to_lift(
|
||||
model_call_details={
|
||||
"call_type": "acompletion",
|
||||
"model": FAILURE_USAGE_MODEL,
|
||||
"messages": ONE_USER_MESSAGE,
|
||||
},
|
||||
request_body={},
|
||||
dispatched=True,
|
||||
provider_billed=True,
|
||||
deployment_id="dep-lit145",
|
||||
)
|
||||
|
||||
assert lifted is not None
|
||||
usage, response_cost = lifted
|
||||
assert usage.prompt_tokens > 0
|
||||
assert response_cost == pytest.approx(usage.prompt_tokens * 0.25)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(("prompt_tokens", "expected_rate"), [(1_000, 1e-7), (250_000, 5e-7)])
|
||||
def test_a_timed_out_request_is_priced_at_the_deployments_long_context_tier(monkeypatch, prompt_tokens, expected_rate):
|
||||
"""A deployment that declares ``input_cost_per_token_above_200k_tokens`` next to its standard rate
|
||||
bills a 250k-token prompt at the tier on the success path; the failure row must not flatten it to
|
||||
the standard rate."""
|
||||
from litellm.proxy.utils import _estimated_input_cost
|
||||
|
||||
monkeypatch.setitem(
|
||||
litellm.model_cost,
|
||||
"dep-lit-f9-tiered",
|
||||
{
|
||||
"litellm_provider": "openai",
|
||||
"mode": "chat",
|
||||
"input_cost_per_token": 1e-7,
|
||||
"input_cost_per_token_above_200k_tokens": 5e-7,
|
||||
"output_cost_per_token": 4e-7,
|
||||
},
|
||||
)
|
||||
|
||||
response_cost = _estimated_input_cost(
|
||||
model_call_details={"model": FAILURE_USAGE_MODEL, "custom_llm_provider": "openai"},
|
||||
prompt_tokens=prompt_tokens,
|
||||
deployment_id="dep-lit-f9-tiered",
|
||||
base_model=None,
|
||||
)
|
||||
|
||||
assert response_cost == pytest.approx(prompt_tokens * expected_rate)
|
||||
|
||||
|
||||
def test_every_earlier_timed_out_attempt_is_billed_at_its_own_deployment_rate(monkeypatch):
|
||||
"""The router leaves one RetryAttemptRecord per failed attempt in metadata.previous_models, the
|
||||
final attempt included; the earlier Timeout ones are billed (each at the deployment that served
|
||||
it), a 429 attempt is not, and the final attempt is left to the failure row itself."""
|
||||
from litellm.proxy.utils import _billed_earlier_attempts
|
||||
|
||||
monkeypatch.setitem(
|
||||
litellm.model_cost,
|
||||
"dep-lit145-b",
|
||||
{"litellm_provider": "openai", "mode": "chat", "input_cost_per_token": 0.25, "output_cost_per_token": 1.0},
|
||||
)
|
||||
usage = litellm.Usage(prompt_tokens=24, completion_tokens=0, total_tokens=24)
|
||||
public_rate, _ = litellm.cost_per_token(model=FAILURE_USAGE_MODEL, prompt_tokens=24)
|
||||
|
||||
attempts = _billed_earlier_attempts(
|
||||
model_call_details={"model": FAILURE_USAGE_MODEL, "custom_llm_provider": "openai"},
|
||||
litellm_params={
|
||||
"metadata": {
|
||||
"attempted_retries": 3,
|
||||
"previous_models": (
|
||||
{"deployment_id": "dep-lit145-a", "exception_type": "Timeout", "attempted_retries": 0},
|
||||
{"deployment_id": "dep-lit145-b", "exception_type": "RateLimitError", "attempted_retries": 1},
|
||||
{"deployment_id": "dep-lit145-b", "exception_type": "Timeout", "attempted_retries": 2},
|
||||
{"deployment_id": "dep-lit145-a", "exception_type": "Timeout", "attempted_retries": 3},
|
||||
),
|
||||
}
|
||||
},
|
||||
estimated_usage=usage,
|
||||
base_model=None,
|
||||
deployment_id="dep-lit145-a",
|
||||
)
|
||||
|
||||
assert [(a.attempted_retries, a.deployment_id) for a in attempts] == [(0, "dep-lit145-a"), (2, "dep-lit145-b")]
|
||||
assert attempts[0].response_cost == pytest.approx(public_rate)
|
||||
assert attempts[0].response_cost > 0
|
||||
assert attempts[1].response_cost == pytest.approx(24 * 0.25)
|
||||
assert all(a.usage is usage for a in attempts)
|
||||
|
||||
|
||||
def test_an_earlier_attempt_on_another_model_is_priced_at_that_models_rate_not_the_final_models(monkeypatch):
|
||||
"""After a fallback hop the final attempt runs on another model group, so the records of the group
|
||||
that timed out first name their own deployment model (and its ``base_model`` when the deployment
|
||||
name is not in the cost map). Each is billed at that model's input rate, not at the final model's."""
|
||||
from litellm.proxy.utils import _billed_earlier_attempts
|
||||
|
||||
for pricing_model, provider, input_rate in (
|
||||
("lit145-final-model", "openai", 0.5),
|
||||
("lit145-first-model", "openai", 0.02),
|
||||
("azure/lit145-first-model", "azure", 0.02),
|
||||
):
|
||||
monkeypatch.setitem(
|
||||
litellm.model_cost,
|
||||
pricing_model,
|
||||
{"litellm_provider": provider, "mode": "chat", "input_cost_per_token": input_rate, "output_cost_per_token": 1.0},
|
||||
)
|
||||
usage = litellm.Usage(prompt_tokens=24, completion_tokens=0, total_tokens=24)
|
||||
|
||||
attempts = _billed_earlier_attempts(
|
||||
model_call_details={"model": "openai/lit145-final-model", "custom_llm_provider": "openai"},
|
||||
litellm_params={
|
||||
"metadata": {
|
||||
"attempted_retries": 1,
|
||||
"model_info": {"id": "dep-lit145-final"},
|
||||
"previous_models": (
|
||||
{
|
||||
"deployment_id": "dep-lit145-first",
|
||||
"model": "openai/lit145-first-model",
|
||||
"base_model": None,
|
||||
"exception_type": "Timeout",
|
||||
"attempted_retries": 0,
|
||||
},
|
||||
{
|
||||
"deployment_id": "dep-lit145-first-azure",
|
||||
"model": "azure/lit145-first-deployment",
|
||||
"base_model": "azure/lit145-first-model",
|
||||
"exception_type": "Timeout",
|
||||
"attempted_retries": 1,
|
||||
},
|
||||
{
|
||||
"deployment_id": "dep-lit145-final",
|
||||
"model": "openai/lit145-final-model",
|
||||
"base_model": None,
|
||||
"exception_type": "Timeout",
|
||||
"attempted_retries": 0,
|
||||
},
|
||||
{
|
||||
"deployment_id": "dep-lit145-final",
|
||||
"model": "openai/lit145-final-model",
|
||||
"base_model": None,
|
||||
"exception_type": "Timeout",
|
||||
"attempted_retries": 1,
|
||||
},
|
||||
),
|
||||
}
|
||||
},
|
||||
estimated_usage=usage,
|
||||
base_model=None,
|
||||
deployment_id="dep-lit145-final",
|
||||
)
|
||||
|
||||
assert [(a.attempted_retries, a.deployment_id, a.response_cost) for a in attempts] == [
|
||||
(0, "dep-lit145-first", pytest.approx(24 * 0.02)),
|
||||
(1, "dep-lit145-first-azure", pytest.approx(24 * 0.02)),
|
||||
(0, "dep-lit145-final", pytest.approx(24 * 0.5)),
|
||||
]
|
||||
|
||||
|
||||
def test_earlier_attempts_are_priced_from_the_prompt_estimate_not_the_final_streams_recovered_usage():
|
||||
"""A final attempt whose stream broke mid-flight lifts its recovered usage (prompt plus the chunks
|
||||
delivered); the earlier timed-out attempts never produced output, so they are priced from the
|
||||
input-only estimate and not from that recovered usage."""
|
||||
from litellm.litellm_core_utils.get_litellm_params import get_litellm_params
|
||||
from litellm.proxy.utils import _estimated_prompt_usage, _failure_fields_to_lift
|
||||
|
||||
logging_obj = _dispatched_failure_logging_obj(
|
||||
FAILURE_USAGE_MODEL,
|
||||
get_litellm_params(
|
||||
metadata={
|
||||
"attempted_retries": 1,
|
||||
"model_info": {"id": "dep-lit145-final"},
|
||||
"previous_models": (
|
||||
{"deployment_id": "dep-lit145-final", "exception_type": "Timeout", "attempted_retries": 0},
|
||||
),
|
||||
}
|
||||
),
|
||||
)
|
||||
recovered = litellm.Usage(prompt_tokens=500, completion_tokens=40, total_tokens=540)
|
||||
logging_obj.model_call_details["combined_usage_object"] = recovered
|
||||
logging_obj.model_call_details["response_cost"] = 0.01
|
||||
|
||||
lifted = _failure_fields_to_lift({"litellm_logging_obj": logging_obj}, provider_billed=True)
|
||||
|
||||
estimate = _estimated_prompt_usage(logging_obj.model_call_details, {})
|
||||
assert lifted["combined_usage_object"] is recovered
|
||||
assert [(a.usage.prompt_tokens, a.usage.completion_tokens) for a in lifted["billed_failed_attempts"]] == [
|
||||
(estimate.prompt_tokens, 0)
|
||||
]
|
||||
assert estimate.prompt_tokens < recovered.prompt_tokens
|
||||
|
||||
|
||||
def test_a_timed_out_attempt_from_the_group_a_fallback_left_is_still_billed():
|
||||
"""After a fallback hop the router restarts ``attempted_retries`` at 0 for the new model group and
|
||||
re-logs the attempt that exhausted the previous group, so the final attempt (dep-b, retry 1) must be
|
||||
told from the previous group's attempt with the same count (dep-a, retry 1), and the re-logged
|
||||
record must bill once."""
|
||||
from litellm.proxy.utils import _billed_earlier_attempts
|
||||
|
||||
usage = litellm.Usage(prompt_tokens=24, completion_tokens=0, total_tokens=24)
|
||||
|
||||
attempts = _billed_earlier_attempts(
|
||||
model_call_details={"model": FAILURE_USAGE_MODEL, "custom_llm_provider": "openai"},
|
||||
litellm_params={
|
||||
"metadata": {
|
||||
"attempted_retries": 1,
|
||||
"model_info": {"id": "dep-lit145-b"},
|
||||
"previous_models": (
|
||||
{"deployment_id": "dep-lit145-a", "exception_type": "Timeout", "attempted_retries": 0},
|
||||
{"deployment_id": "dep-lit145-a", "exception_type": "Timeout", "attempted_retries": 1},
|
||||
{"deployment_id": "dep-lit145-a", "exception_type": "Timeout", "attempted_retries": 1},
|
||||
{"deployment_id": "dep-lit145-b", "exception_type": "Timeout", "attempted_retries": 0},
|
||||
{"deployment_id": "dep-lit145-b", "exception_type": "Timeout", "attempted_retries": 1},
|
||||
),
|
||||
}
|
||||
},
|
||||
estimated_usage=usage,
|
||||
base_model=None,
|
||||
deployment_id="dep-lit145-b",
|
||||
)
|
||||
|
||||
assert [(a.attempted_retries, a.deployment_id) for a in attempts] == [
|
||||
(0, "dep-lit145-a"),
|
||||
(1, "dep-lit145-a"),
|
||||
(0, "dep-lit145-b"),
|
||||
]
|
||||
assert all(a.response_cost > 0 for a in attempts)
|
||||
|
||||
|
||||
def _dispatched_failure_logging_obj(model: str, litellm_params: dict) -> object:
|
||||
"""A Logging object as the failure hook sees it: litellm_params in the real get_litellm_params()
|
||||
shape (every key present, ``litellm_metadata`` None on chat completions) and a first-handoff
|
||||
instant, so the lift treats the request as dispatched."""
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging
|
||||
|
||||
logging_obj = Logging(
|
||||
model=model,
|
||||
messages=ONE_USER_MESSAGE,
|
||||
stream=False,
|
||||
call_type="acompletion",
|
||||
start_time=real_datetime.datetime.now(),
|
||||
litellm_call_id="call-lit145",
|
||||
function_id="fn-lit145",
|
||||
)
|
||||
logging_obj.update_environment_variables(model=model, user="", optional_params={}, litellm_params=litellm_params)
|
||||
logging_obj.model_call_details["first_api_call_start_time"] = real_datetime.datetime.now()
|
||||
logging_obj.model_call_details["messages"] = ONE_USER_MESSAGE
|
||||
logging_obj.model_call_details["call_type"] = "acompletion"
|
||||
return logging_obj
|
||||
|
||||
|
||||
def test_the_deployment_rate_is_found_on_the_logging_objects_real_litellm_params(monkeypatch):
|
||||
"""``get_litellm_params()`` always carries a ``litellm_metadata`` key (None on chat
|
||||
completions) next to the router-stamped ``metadata.model_info.id``; the failure row must
|
||||
still price at the deployment's own rate, like ``Logging.get_router_model_id`` on the
|
||||
success path."""
|
||||
from litellm.litellm_core_utils.get_litellm_params import get_litellm_params
|
||||
from litellm.proxy.utils import _failure_fields_to_lift
|
||||
|
||||
monkeypatch.setitem(
|
||||
litellm.model_cost,
|
||||
"dep-lit145-real-shape",
|
||||
{"litellm_provider": "openai", "mode": "chat", "input_cost_per_token": 0.25, "output_cost_per_token": 0.25},
|
||||
)
|
||||
litellm_params = get_litellm_params(metadata={"model_info": {"id": "dep-lit145-real-shape"}})
|
||||
assert "litellm_metadata" in litellm_params and litellm_params["litellm_metadata"] is None
|
||||
logging_obj = _dispatched_failure_logging_obj(FAILURE_USAGE_MODEL, litellm_params)
|
||||
|
||||
lifted = _failure_fields_to_lift({"litellm_logging_obj": logging_obj}, provider_billed=True)
|
||||
|
||||
assert lifted["response_cost"] == pytest.approx(lifted["combined_usage_object"].prompt_tokens * 0.25)
|
||||
|
||||
|
||||
def test_an_azure_deployment_is_priced_at_its_base_model_like_the_success_path():
|
||||
"""An Azure deployment name is not in the cost map; the success path prices it through
|
||||
``model_info.base_model`` (``get_base_model_from_metadata``) and so must the failure row."""
|
||||
from litellm.litellm_core_utils.get_litellm_params import get_litellm_params
|
||||
from litellm.proxy.utils import _failure_fields_to_lift
|
||||
|
||||
logging_obj = _dispatched_failure_logging_obj(
|
||||
"azure/lit145-not-in-cost-map",
|
||||
get_litellm_params(metadata={"model_info": {"base_model": "azure/gpt-4.1"}}),
|
||||
)
|
||||
logging_obj.model_call_details["custom_llm_provider"] = "azure"
|
||||
|
||||
lifted = _failure_fields_to_lift({"litellm_logging_obj": logging_obj}, provider_billed=True)
|
||||
|
||||
expected_cost, _ = litellm.cost_per_token(
|
||||
model="azure/gpt-4.1", prompt_tokens=lifted["combined_usage_object"].prompt_tokens
|
||||
)
|
||||
assert expected_cost > 0
|
||||
assert lifted["response_cost"] == pytest.approx(expected_cost)
|
||||
|
||||
|
||||
def test_a_litellm_params_key_in_the_request_body_cannot_zero_the_timeout_price():
|
||||
"""``model_call_details`` is flattened with the request's optional_params, so a body key
|
||||
named ``litellm_params`` can land there; the deployment rate must come from the logging
|
||||
object's own ``litellm_params`` instead."""
|
||||
from litellm.proxy.utils import _failure_usage_to_lift
|
||||
|
||||
lifted = _failure_usage_to_lift(
|
||||
model_call_details={
|
||||
"call_type": "acompletion",
|
||||
"model": FAILURE_USAGE_MODEL,
|
||||
"messages": ONE_USER_MESSAGE,
|
||||
"litellm_params": {"input_cost_per_token": 0},
|
||||
},
|
||||
request_body={},
|
||||
dispatched=True,
|
||||
provider_billed=True,
|
||||
deployment_id=None,
|
||||
)
|
||||
|
||||
assert lifted is not None
|
||||
usage, response_cost = lifted
|
||||
expected_cost, _ = litellm.cost_per_token(model=FAILURE_USAGE_MODEL, prompt_tokens=usage.prompt_tokens)
|
||||
assert response_cost == pytest.approx(expected_cost)
|
||||
assert response_cost > 0
|
||||
|
||||
|
||||
def test_a_request_that_timed_out_at_a_provider_bills_its_input_at_the_input_rate():
|
||||
from litellm.proxy.utils import _failure_usage_to_lift
|
||||
|
||||
lifted = _failure_usage_to_lift(
|
||||
|
|
@ -1912,13 +2560,16 @@ def test_a_request_that_reached_a_provider_bills_its_input_at_no_cost():
|
|||
},
|
||||
request_body={},
|
||||
dispatched=True,
|
||||
provider_billed=True,
|
||||
)
|
||||
|
||||
assert lifted is not None
|
||||
usage, response_cost = lifted
|
||||
assert usage.prompt_tokens > 0
|
||||
assert usage.completion_tokens == 0
|
||||
assert response_cost == 0.0
|
||||
expected_input_cost, _ = litellm.cost_per_token(model=FAILURE_USAGE_MODEL, prompt_tokens=usage.prompt_tokens)
|
||||
assert expected_input_cost > 0
|
||||
assert response_cost == pytest.approx(expected_input_cost)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
|
|
|
|||
|
|
@ -917,3 +917,53 @@ async def test_update_spend_logs_parks_failed_batch_in_redis_with_wire_safe_date
|
|||
parked = await buffer.get_spend_logs_from_redis_buffer(limit=10)
|
||||
assert mock_prisma_client.spend_log_transactions == []
|
||||
assert [(row["request_id"], row["startTime"]) for row in parked] == [("a", started.isoformat())]
|
||||
|
||||
|
||||
async def _create_many_while_postgres_is_restarting(**_: object) -> int:
|
||||
from prisma.engine.utils import handle_response_errors
|
||||
|
||||
handle_response_errors(
|
||||
MagicMock(),
|
||||
[
|
||||
{
|
||||
"error": "Error in connector: Can't reach database server at `127.0.0.1:5432`",
|
||||
"user_facing_error": {"error_code": "P1001", "message": "Can't reach database server", "meta": None},
|
||||
}
|
||||
],
|
||||
)
|
||||
return 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_spend_logs_retries_and_requeues_batch_when_postgres_restart_tears_down_prisma_engine(
|
||||
mock_prisma_client: Any, make_spend_log_row: Any, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
"""While Postgres restarts the prisma engine answers the write with a P1001
|
||||
payload whose ``meta`` is ``None``, which prisma's own response parser trips
|
||||
over before it can build a typed error. That outage must be retried like
|
||||
any other, and once retries exhaust the dequeued batch goes back to the head
|
||||
of the queue so the rows land when Postgres is back.
|
||||
"""
|
||||
|
||||
async def _fake_sleep(_: float) -> None:
|
||||
return None
|
||||
|
||||
monkeypatch.setattr(utils_mod.asyncio, "sleep", _fake_sleep)
|
||||
mock_prisma_client.db.litellm_spendlogs.create_many = AsyncMock(
|
||||
side_effect=_create_many_while_postgres_is_restarting
|
||||
)
|
||||
proxy_logging = MagicMock()
|
||||
proxy_logging.failure_handler = AsyncMock()
|
||||
mock_prisma_client.spend_log_transactions = [make_spend_log_row(request_id="c")]
|
||||
|
||||
with pytest.raises(AttributeError):
|
||||
await ProxyUpdateSpend.update_spend_logs(
|
||||
n_retry_times=2,
|
||||
prisma_client=mock_prisma_client,
|
||||
db_writer_client=None,
|
||||
proxy_logging_obj=proxy_logging,
|
||||
logs_to_process=[make_spend_log_row(request_id="a"), make_spend_log_row(request_id="b")],
|
||||
)
|
||||
|
||||
assert [row["request_id"] for row in mock_prisma_client.spend_log_transactions] == ["a", "b", "c"]
|
||||
assert mock_prisma_client.db.litellm_spendlogs.create_many.await_count == 3
|
||||
|
|
|
|||
|
|
@ -17,6 +17,7 @@ deterministic stand-ins so the arithmetic under test is the only variable.
|
|||
import json
|
||||
import logging
|
||||
from types import MappingProxyType
|
||||
from typing import Final, Literal
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
|
@ -640,7 +641,7 @@ async def test_calculate_vertex_disable_transform_path(monkeypatch):
|
|||
monkeypatch.setattr(
|
||||
bu,
|
||||
"calculate_vertex_ai_batch_cost_and_usage",
|
||||
lambda content, model, model_info=None: bu.BatchCostUsageResult(
|
||||
lambda content, model, model_info=None, vertex_location=None: bu.BatchCostUsageResult(
|
||||
cost=9.9,
|
||||
usage=Usage(prompt_tokens=1, completion_tokens=2, total_tokens=3),
|
||||
models=["gemini-2.0-flash-001"],
|
||||
|
|
@ -671,7 +672,7 @@ async def test_calculate_vertex_disable_transform_needs_model_name(monkeypatch):
|
|||
monkeypatch.setattr(
|
||||
bu,
|
||||
"calculate_vertex_ai_batch_cost_and_usage",
|
||||
lambda content, model, model_info=None: pytest.fail("raw vertex path should not run"),
|
||||
lambda content, model, model_info=None, vertex_location=None: pytest.fail("raw vertex path should not run"),
|
||||
)
|
||||
|
||||
result = await bu.calculate_batch_cost_and_usage(file_content_dictionary=[], custom_llm_provider="vertex_ai")
|
||||
|
|
@ -1112,6 +1113,36 @@ async def test_handle_completed_vertex_batch_computes_cost_usage_and_models(monk
|
|||
assert result.failed_requests == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_handle_completed_vertex_batch_regional_location_costs_uplift_over_global(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
import litellm.files.main as files_main
|
||||
|
||||
rows = [_vertex_openai_row("request-1", "gemini-3.6-flash", 10, 5)]
|
||||
|
||||
async def fake_afile_content(**kw):
|
||||
return type("R", (), {"content": _vertex_jsonl(rows)})()
|
||||
|
||||
monkeypatch.setattr(files_main, "afile_content", fake_afile_content)
|
||||
|
||||
global_result = await bu._handle_completed_batch(
|
||||
_batch("gs://litellm-bucket/output/predictions.jsonl"),
|
||||
custom_llm_provider="vertex_ai",
|
||||
litellm_params={"vertex_project": "proj-1", "vertex_location": "global"},
|
||||
vertex_location="global",
|
||||
)
|
||||
regional_result = await bu._handle_completed_batch(
|
||||
_batch("gs://litellm-bucket/output/predictions.jsonl"),
|
||||
custom_llm_provider="vertex_ai",
|
||||
litellm_params={"vertex_project": "proj-1", "vertex_location": "us-central1"},
|
||||
vertex_location="us-central1",
|
||||
)
|
||||
|
||||
uplift = litellm.model_cost["vertex_ai/gemini-3.6-flash"]["regional_endpoint_uplift_multiplier"]
|
||||
assert global_result.cost > 0
|
||||
assert regional_result.cost == pytest.approx(global_result.cost * uplift, rel=1e-9)
|
||||
assert regional_result.usage == global_result.usage
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_output_file_content_no_output_file_id_raises():
|
||||
with pytest.raises(ValueError, match="Output file id is None"):
|
||||
|
|
@ -1341,7 +1372,7 @@ async def test_handle_completed_batch_vertex_disable_transform_path(monkeypatch)
|
|||
monkeypatch.setattr(litellm, "disable_vertex_batch_output_transformation", True, raising=False)
|
||||
seen: dict = {}
|
||||
|
||||
def fake_vertex_calc(content, model, model_info=None):
|
||||
def fake_vertex_calc(content, model, model_info=None, vertex_location=None):
|
||||
seen["content"] = content
|
||||
seen["model"] = model
|
||||
return bu.BatchCostUsageResult(
|
||||
|
|
@ -2417,3 +2448,30 @@ async def test_flag_sends_every_vertex_row_down_the_native_path_when_a_model_is_
|
|||
|
||||
assert calls == []
|
||||
assert (result.successful_requests, result.failed_requests) == (0, 1)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("provider", ("openai", "anthropic"))
|
||||
@pytest.mark.parametrize("envelope", (1, "invalid", [1], None, {}))
|
||||
@pytest.mark.parametrize("include_valid", (True, False))
|
||||
def test_invalid_batch_envelopes_preserve_recoverable_cost(
|
||||
provider: Literal["openai", "anthropic"], envelope: object, include_valid: bool
|
||||
) -> None:
|
||||
invalid: Final[dict[str, object]] = {"response" if provider == "openai" else "result": envelope}
|
||||
valid: Final[dict[str, object]] = (
|
||||
_success_row(model="batch-accounting-control", usage=_usage(100, 20))
|
||||
if provider == "openai"
|
||||
else {
|
||||
"result": {
|
||||
"type": "succeeded",
|
||||
"message": {"model": "batch-accounting-control", "usage": {"input_tokens": 100, "output_tokens": 20}},
|
||||
}
|
||||
}
|
||||
)
|
||||
entries: Final = (valid, invalid, valid) if include_valid else (invalid,)
|
||||
result: Final = bu._aggregate_batch_cost_usage_models(
|
||||
entries, provider, model_info={"input_cost_per_token_batches": 0.001, "output_cost_per_token_batches": 0.002}
|
||||
)
|
||||
assert result.cost == pytest.approx(0.28 if include_valid else 0)
|
||||
assert result.usage.prompt_tokens == (200 if include_valid else 0)
|
||||
assert result.usage.completion_tokens == (40 if include_valid else 0)
|
||||
assert (result.successful_requests, result.failed_requests) == (2 if include_valid else 0, 1)
|
||||
|
|
|
|||
120
tests/unit/decisions/test_call.py
Normal file
120
tests/unit/decisions/test_call.py
Normal file
|
|
@ -0,0 +1,120 @@
|
|||
import base64
|
||||
import json
|
||||
from collections.abc import Mapping
|
||||
from datetime import datetime, timezone
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
|
||||
import litellm
|
||||
from litellm.decisions.main import _prepare_decisions_call
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging
|
||||
from litellm.systemone.main import _prepare_systemone_call
|
||||
|
||||
QUESTION_TEXT: Final = "Is this about a billing problem?"
|
||||
CONNECTION: Final = {
|
||||
"model": "hosted_vllm/jev-fake",
|
||||
"safety_identifier": None,
|
||||
"api_key": "fake-provider-key",
|
||||
"api_base": "http://127.0.0.1:4010",
|
||||
"timeout": None,
|
||||
"custom_llm_provider": None,
|
||||
"extra_headers": None,
|
||||
}
|
||||
|
||||
|
||||
def _logging_obj(call_type: str) -> Logging:
|
||||
return Logging(
|
||||
model="hosted_vllm/jev-fake",
|
||||
messages="placeholder",
|
||||
stream=False,
|
||||
call_type=call_type,
|
||||
start_time=datetime.now(timezone.utc),
|
||||
litellm_call_id="lit-9659",
|
||||
function_id="lit-9659",
|
||||
)
|
||||
|
||||
|
||||
def _logged_request(logging_obj: Logging) -> Mapping[str, object]:
|
||||
messages: Final = logging_obj.model_call_details["messages"]
|
||||
assert len(messages) == 1 and messages[0]["role"] == "user"
|
||||
return json.loads(messages[0]["content"])
|
||||
|
||||
|
||||
def test_systemone_logs_every_request_field_the_caller_sent():
|
||||
logging_obj: Final = _logging_obj("asystemone")
|
||||
|
||||
_prepare_systemone_call(
|
||||
**CONNECTION,
|
||||
state={"message": "My card was charged twice."},
|
||||
questions={"billing": {"type": "noul", "instructions": QUESTION_TEXT}},
|
||||
kwargs={"litellm_logging_obj": logging_obj},
|
||||
)
|
||||
|
||||
logged: Final = _logged_request(logging_obj)
|
||||
assert logged["state"] == {"message": "My card was charged twice."}
|
||||
assert logged["questions"]["billing"]["instructions"] == QUESTION_TEXT
|
||||
assert set(logged) == {"state", "questions"}
|
||||
|
||||
|
||||
def test_decisions_logs_every_request_field_the_caller_sent():
|
||||
logging_obj: Final = _logging_obj("adecisions")
|
||||
|
||||
_prepare_decisions_call(
|
||||
**CONNECTION,
|
||||
decision_input="My card was charged twice.",
|
||||
questions=[{"type": "predicate", "instructions": QUESTION_TEXT}],
|
||||
kwargs={"litellm_logging_obj": logging_obj},
|
||||
)
|
||||
|
||||
logged: Final = _logged_request(logging_obj)
|
||||
assert logged["input"] == "My card was charged twice."
|
||||
assert logged["questions"][0]["instructions"] == QUESTION_TEXT
|
||||
assert set(logged) == {"input", "questions"}
|
||||
|
||||
|
||||
def test_request_is_logged_even_when_it_is_rejected_before_the_provider_call():
|
||||
logging_obj: Final = _logging_obj("asystemone")
|
||||
|
||||
with pytest.raises(litellm.BadRequestError):
|
||||
_prepare_systemone_call(
|
||||
**CONNECTION,
|
||||
state={"message": "My card was charged twice."},
|
||||
questions={"billing": {"type": "noul", "instructions": QUESTION_TEXT}},
|
||||
kwargs={"litellm_logging_obj": logging_obj, "input": "sent to the wrong function"},
|
||||
)
|
||||
|
||||
logged: Final = _logged_request(logging_obj)
|
||||
assert logged["questions"]["billing"]["instructions"] == QUESTION_TEXT
|
||||
|
||||
|
||||
def test_decisions_request_is_logged_even_when_its_questions_fail_validation():
|
||||
logging_obj: Final = _logging_obj("adecisions")
|
||||
|
||||
with pytest.raises(litellm.BadRequestError):
|
||||
_prepare_decisions_call(
|
||||
**CONNECTION,
|
||||
decision_input="My card was charged twice.",
|
||||
questions=[{"type": "not-a-question-type", "instructions": QUESTION_TEXT}],
|
||||
kwargs={"litellm_logging_obj": logging_obj},
|
||||
)
|
||||
|
||||
logged: Final = _logged_request(logging_obj)
|
||||
assert logged["questions"][0]["instructions"] == QUESTION_TEXT
|
||||
|
||||
|
||||
def test_non_utf8_bytes_in_the_request_are_logged_and_rejected_with_a_400():
|
||||
logging_obj: Final = _logging_obj("adecisions")
|
||||
raw_input: Final = b"\xff\xfe not utf-8"
|
||||
|
||||
with pytest.raises(litellm.BadRequestError):
|
||||
_prepare_decisions_call(
|
||||
**CONNECTION,
|
||||
decision_input=raw_input,
|
||||
questions=[{"type": "predicate", "instructions": QUESTION_TEXT}],
|
||||
kwargs={"litellm_logging_obj": logging_obj},
|
||||
)
|
||||
|
||||
logged: Final = _logged_request(logging_obj)
|
||||
assert base64.urlsafe_b64decode(logged["input"]) == raw_input
|
||||
assert logged["questions"][0]["instructions"] == QUESTION_TEXT
|
||||
|
|
@ -9,6 +9,8 @@ with deployment credentials, bypassing the managed files access-control hooks.
|
|||
"""
|
||||
|
||||
import base64
|
||||
import json
|
||||
from typing import Final
|
||||
import pytest
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
|
@ -367,13 +369,14 @@ async def test_check_batch_cost_should_call_afile_content_directly_with_credenti
|
|||
# managed_files_obj.afile_content should NOT have been called
|
||||
mock_managed_files_hook.afile_content.assert_not_called()
|
||||
|
||||
# Verify the DB update writes batch_processed, status, and file_object
|
||||
mock_prisma.db.litellm_managedobjecttable.update.assert_called_once()
|
||||
update_call_kwargs = (
|
||||
mock_prisma.db.litellm_managedobjecttable.update.call_args.kwargs
|
||||
claim_data: Final = next(
|
||||
call.kwargs["data"]
|
||||
for call in mock_prisma.db.litellm_managedobjecttable.update_many.call_args_list
|
||||
if call.kwargs["where"].get("id") == mock_job.id
|
||||
)
|
||||
assert update_call_kwargs["data"]["batch_processed"] is True
|
||||
assert update_call_kwargs["data"]["status"] == "complete"
|
||||
assert (
|
||||
"file_object" in update_call_kwargs["data"]
|
||||
), "file_object must be written to DB so list_batches reads updated status"
|
||||
assert claim_data["batch_processed"] is True
|
||||
assert claim_data["status"] == "complete"
|
||||
persisted_batch: Final = json.loads(claim_data["file_object"])
|
||||
assert persisted_batch["status"] == "completed"
|
||||
assert persisted_batch["usage"]["prompt_tokens"] == 10
|
||||
assert persisted_batch["usage"]["completion_tokens"] == 5
|
||||
|
|
|
|||
|
|
@ -99,6 +99,31 @@ def test_generic_cost_per_token_bills_cache_reads_at_input_rate_when_no_cache_re
|
|||
assert completion_cost == pytest.approx(380 * 9.7e-7)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("model_info", "expected_prompt_cost"),
|
||||
[
|
||||
({"input_cost_per_token": 1e-7, "input_cost_per_character": 2.5e-8, "output_cost_per_token": 0.0}, 35 * 2.5e-8),
|
||||
({"input_cost_per_token": 1.5e-7, "output_cost_per_token": 0.0}, 11 * 1.5e-7),
|
||||
],
|
||||
ids=["character_priced_model_bills_characters_instead_of_tokens", "token_priced_model_ignores_the_character_count"],
|
||||
)
|
||||
def test_generic_cost_per_token_bills_a_character_counted_prompt_under_one_unit(
|
||||
model_info: ModelInfo, expected_prompt_cost: float
|
||||
) -> None:
|
||||
usage = Usage(
|
||||
prompt_tokens=11,
|
||||
completion_tokens=0,
|
||||
total_tokens=11,
|
||||
prompt_tokens_details=PromptTokensDetailsWrapper(character_count=35),
|
||||
)
|
||||
|
||||
prompt_cost, completion_cost = generic_cost_per_token(
|
||||
model="embedding-fixture", usage=usage, custom_llm_provider="vertex_ai", model_info=model_info
|
||||
)
|
||||
|
||||
assert (prompt_cost, completion_cost) == (pytest.approx(expected_prompt_cost), 0.0)
|
||||
|
||||
|
||||
def test_apply_provider_cache_read_default_only_derives_a_rate_for_fireworks() -> None:
|
||||
openai_info: ModelInfo = {"input_cost_per_token": 2e-6}
|
||||
fireworks_info: ModelInfo = {"input_cost_per_token": 2e-6}
|
||||
|
|
@ -450,6 +475,66 @@ def test_video_tokens_fallback_to_base_cost():
|
|||
assert round(completion_cost, 12) == round((600 + 1120) * 2e-6, 12)
|
||||
|
||||
|
||||
def test_prompt_audio_and_video_tokens_without_modality_rates_bill_at_tiered_input_rate():
|
||||
"""Prompt audio/video tokens with no modality rate are prompt tokens, so a >200k prompt bills them above-200k."""
|
||||
from unittest.mock import patch
|
||||
|
||||
mock_model_info = {
|
||||
"input_cost_per_token": 2e-6,
|
||||
"input_cost_per_token_above_200k_tokens": 4e-6,
|
||||
"output_cost_per_token": 1.2e-5,
|
||||
"output_cost_per_token_above_200k_tokens": 1.8e-5,
|
||||
}
|
||||
|
||||
usage = Usage(
|
||||
prompt_tokens=816032,
|
||||
completion_tokens=4559,
|
||||
total_tokens=820591,
|
||||
prompt_tokens_details=PromptTokensDetailsWrapper(text_tokens=9033, audio_tokens=14999, video_tokens=792000),
|
||||
completion_tokens_details=CompletionTokensDetailsWrapper(text_tokens=4559),
|
||||
)
|
||||
|
||||
with patch(
|
||||
"litellm.litellm_core_utils.llm_cost_calc.utils.get_model_info",
|
||||
return_value=mock_model_info,
|
||||
):
|
||||
prompt_cost, completion_cost = generic_cost_per_token(
|
||||
model="test-model", usage=usage, custom_llm_provider="vertex_ai"
|
||||
)
|
||||
|
||||
assert round(prompt_cost, 9) == round(816032 * 4e-6, 9)
|
||||
assert round(completion_cost, 9) == round(4559 * 1.8e-5, 9)
|
||||
|
||||
|
||||
def test_prompt_modality_rates_still_win_over_tiered_input_rate():
|
||||
"""An explicit input_cost_per_audio_token / input_cost_per_video_token beats the fallback, even above 200k."""
|
||||
from unittest.mock import patch
|
||||
|
||||
mock_model_info = {
|
||||
"input_cost_per_token": 2e-6,
|
||||
"input_cost_per_token_above_200k_tokens": 4e-6,
|
||||
"input_cost_per_audio_token": 3e-6,
|
||||
"input_cost_per_video_token": 1e-6,
|
||||
"output_cost_per_token": 1.2e-5,
|
||||
}
|
||||
|
||||
usage = Usage(
|
||||
prompt_tokens=300000,
|
||||
completion_tokens=10,
|
||||
total_tokens=300010,
|
||||
prompt_tokens_details=PromptTokensDetailsWrapper(text_tokens=100000, audio_tokens=100000, video_tokens=100000),
|
||||
completion_tokens_details=CompletionTokensDetailsWrapper(text_tokens=10),
|
||||
)
|
||||
|
||||
with patch(
|
||||
"litellm.litellm_core_utils.llm_cost_calc.utils.get_model_info",
|
||||
return_value=mock_model_info,
|
||||
):
|
||||
prompt_cost, _ = generic_cost_per_token(model="test-model", usage=usage, custom_llm_provider="vertex_ai")
|
||||
|
||||
assert round(prompt_cost, 9) == round(100000 * 4e-6 + 100000 * 3e-6 + 100000 * 1e-6, 9)
|
||||
|
||||
|
||||
def test_generic_cost_per_token_above_200k_tokens(_local_model_cost_map):
|
||||
# gemini-2.5-pro-exp-03-25 was removed; gemini-2.5-pro has same above-200k pricing
|
||||
model = "gemini-2.5-pro"
|
||||
|
|
@ -3801,3 +3886,220 @@ def test_azure_gpt_5_6_alias_matches_sol_pricing(_local_model_cost_map, region_p
|
|||
assert shared_cost_fields
|
||||
for field in shared_cost_fields:
|
||||
assert alias[field] == sol[field], field
|
||||
|
||||
|
||||
# Per-token rates read 2026-10-07 from https://platform.claude.com/docs/en/about-claude/pricing (direct and
|
||||
# azure_ai, which Microsoft bills at Anthropic's rates per
|
||||
# https://learn.microsoft.com/en-us/azure/foundry/foundry-models/concepts/claude-models-billing) and from the
|
||||
# AmazonBedrockFoundationModels price list at
|
||||
# https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json (Bedrock)
|
||||
@pytest.mark.parametrize(
|
||||
("model", "custom_llm_provider", "prompt_tokens", "input_rate", "cache_read_rate", "output_rate"),
|
||||
[
|
||||
("claude-haiku-5-5", "anthropic", 100_000, 1e-07, 1e-08, 5e-07),
|
||||
("claude-haiku-5-5", "anthropic", 100_001, 5e-07, 5e-08, 2.5e-06),
|
||||
("azure_ai/claude-haiku-5-5", "azure_ai", 100_000, 1e-07, 1e-08, 5e-07),
|
||||
("azure_ai/claude-haiku-5-5", "azure_ai", 100_001, 5e-07, 5e-08, 2.5e-06),
|
||||
("global.anthropic.claude-haiku-5-5", "bedrock", 100_000, 1e-07, 1e-08, 5e-07),
|
||||
("global.anthropic.claude-haiku-5-5", "bedrock", 100_001, 5e-07, 5e-08, 2.5e-06),
|
||||
("us.anthropic.claude-haiku-5-5", "bedrock", 100_000, 1.1e-07, 1.1e-08, 5.5e-07),
|
||||
("us.anthropic.claude-haiku-5-5", "bedrock", 100_001, 5.5e-07, 5.5e-08, 2.75e-06),
|
||||
("us-gov.anthropic.claude-haiku-5-5", "bedrock", 100_000, 1.2e-07, 1.2e-08, 6e-07),
|
||||
("us-gov.anthropic.claude-haiku-5-5", "bedrock", 100_001, 6e-07, 6e-08, 3e-06),
|
||||
("bedrock_mantle/anthropic.claude-haiku-5-5", "bedrock_mantle", 100_001, 5.5e-07, 5.5e-08, 2.75e-06),
|
||||
("bedrock/us-gov-west-1/anthropic.claude-haiku-5-5", "bedrock", 100_001, 6e-07, 6e-08, 3e-06),
|
||||
("vertex_ai/claude-haiku-5-5", "vertex_ai", 100_000, 1e-07, 1e-08, 5e-07),
|
||||
("vertex_ai/claude-haiku-5-5", "vertex_ai", 100_001, 5e-07, 5e-08, 2.5e-06),
|
||||
],
|
||||
)
|
||||
def test_generic_cost_per_token_claude_haiku_5_5_prompt_length_tiers(
|
||||
_local_model_cost_map: None,
|
||||
model: str,
|
||||
custom_llm_provider: str,
|
||||
prompt_tokens: int,
|
||||
input_rate: float,
|
||||
cache_read_rate: float,
|
||||
output_rate: float,
|
||||
) -> None:
|
||||
"""Claude Haiku 5.5 bills every token at 5x the base rates once the prompt is over 100,000 tokens."""
|
||||
cached_tokens: Final = 10_000
|
||||
completion_tokens: Final = 1_000
|
||||
usage: Final = Usage(
|
||||
prompt_tokens=prompt_tokens,
|
||||
completion_tokens=completion_tokens,
|
||||
total_tokens=prompt_tokens + completion_tokens,
|
||||
prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=cached_tokens),
|
||||
)
|
||||
|
||||
prompt_cost, completion_cost = generic_cost_per_token(
|
||||
model=model,
|
||||
usage=usage,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
)
|
||||
|
||||
assert prompt_cost == pytest.approx((prompt_tokens - cached_tokens) * input_rate + cached_tokens * cache_read_rate)
|
||||
assert completion_cost == pytest.approx(completion_tokens * output_rate)
|
||||
|
||||
|
||||
def test_vertex_regional_endpoint_uplift_scales_claude_haiku_5_5_over_100k_rates(
|
||||
_local_model_cost_map: None,
|
||||
) -> None:
|
||||
"""Vertex regional endpoints bill 1.1x the global rate on all token types
|
||||
(https://cloud.google.com/vertex-ai/generative-ai/pricing, 2026-10-07: regional
|
||||
over-100K input is $0.55/MTok), so the uplift scales the over-100k rates too."""
|
||||
cached_tokens: Final = 10_000
|
||||
prompt_tokens: Final = 100_001
|
||||
completion_tokens: Final = 1_000
|
||||
usage: Final = Usage(
|
||||
prompt_tokens=prompt_tokens,
|
||||
completion_tokens=completion_tokens,
|
||||
total_tokens=prompt_tokens + completion_tokens,
|
||||
prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=cached_tokens),
|
||||
)
|
||||
|
||||
prompt_cost, completion_cost = generic_cost_per_token(
|
||||
model="vertex_ai/claude-haiku-5-5",
|
||||
usage=usage,
|
||||
custom_llm_provider="vertex_ai",
|
||||
vertex_location="us-east5",
|
||||
)
|
||||
|
||||
assert prompt_cost == pytest.approx((prompt_tokens - cached_tokens) * 5.5e-07 + cached_tokens * 5.5e-08)
|
||||
assert completion_cost == pytest.approx(completion_tokens * 2.75e-06)
|
||||
|
||||
|
||||
# Batch rates read 2026-10-07 from the Batch processing table at
|
||||
# https://platform.claude.com/docs/en/about-claude/pricing: $0.05 / $0.25 per MTok input and $0.25 / $1.25 output,
|
||||
# up to and over 100,000 prompt tokens
|
||||
@pytest.mark.parametrize(
|
||||
("prompt_tokens", "input_rate", "output_rate"),
|
||||
[(100_000, 5e-08, 2.5e-07), (100_001, 2.5e-07, 1.25e-06)],
|
||||
)
|
||||
def test_batch_cost_calculator_claude_haiku_5_5_prompt_length_tiers(
|
||||
_local_model_cost_map: None,
|
||||
prompt_tokens: int,
|
||||
input_rate: float,
|
||||
output_rate: float,
|
||||
) -> None:
|
||||
from litellm.cost_calculator import batch_cost_calculator
|
||||
|
||||
completion_tokens: Final = 1_000
|
||||
|
||||
prompt_cost, completion_cost = batch_cost_calculator(
|
||||
usage=Usage(
|
||||
prompt_tokens=prompt_tokens,
|
||||
completion_tokens=completion_tokens,
|
||||
total_tokens=prompt_tokens + completion_tokens,
|
||||
),
|
||||
model="claude-haiku-5-5",
|
||||
custom_llm_provider="anthropic",
|
||||
)
|
||||
|
||||
assert prompt_cost == pytest.approx(prompt_tokens * input_rate)
|
||||
assert completion_cost == pytest.approx(completion_tokens * output_rate)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("prompt_tokens", "expected"),
|
||||
[
|
||||
(100_000, (5e-08, 2.5e-07, 5e-09, 6.25e-08)),
|
||||
(100_001, (2.5e-07, 1.25e-06, 2.5e-08, 3.125e-07)),
|
||||
],
|
||||
)
|
||||
def test_get_batch_cost_rates_claude_haiku_5_5_prompt_length_tiers(
|
||||
_local_model_cost_map: None,
|
||||
prompt_tokens: int,
|
||||
expected: tuple[float, float, float, float],
|
||||
) -> None:
|
||||
"""Cache write and cache read batch rates are 50% of the standard rates; Anthropic's batch table omits them."""
|
||||
from litellm.litellm_core_utils.llm_cost_calc.utils import get_batch_cost_rates
|
||||
|
||||
rates: Final = get_batch_cost_rates(
|
||||
litellm.get_model_info(model="claude-haiku-5-5", custom_llm_provider="anthropic"),
|
||||
Usage(prompt_tokens=prompt_tokens, completion_tokens=1, total_tokens=prompt_tokens + 1),
|
||||
"anthropic",
|
||||
)
|
||||
|
||||
assert (rates.input, rates.output, rates.cache_read, rates.cache_creation) == expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"field", ("output_cost_per_token", "cache_read_input_token_cost", "cache_creation_input_token_cost")
|
||||
)
|
||||
def test_flat_input_price_keeps_other_token_thresholds(field: str) -> None:
|
||||
from typing import Final, cast
|
||||
from litellm.litellm_core_utils.llm_cost_calc.utils import _get_token_base_cost
|
||||
|
||||
info: Final = cast(
|
||||
ModelInfo,
|
||||
{
|
||||
"input_cost_per_token": 0.001,
|
||||
"output_cost_per_token": 0.002,
|
||||
"cache_read_input_token_cost": 0.003,
|
||||
"cache_creation_input_token_cost": 0.004,
|
||||
f"{field}_above_300k_tokens": 0.005,
|
||||
},
|
||||
)
|
||||
rates: Final = _get_token_base_cost(
|
||||
model_info=info,
|
||||
usage=Usage(prompt_tokens=300_001, completion_tokens=3, total_tokens=300_004),
|
||||
)
|
||||
assert rates[0] == 0.001
|
||||
assert rates[1] == (0.005 if field == "output_cost_per_token" else 0.002)
|
||||
assert rates[2] == (0.005 if field == "cache_creation_input_token_cost" else 0.004)
|
||||
assert rates[4] == (0.005 if field == "cache_read_input_token_cost" else 0.003)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("prompt_tokens", (200_000, 200_001, 512_000, 512_001))
|
||||
@pytest.mark.parametrize("higher_rate", (0.0, 0.009))
|
||||
def test_mixed_token_thresholds_preserve_each_crossed_rate(prompt_tokens: int, higher_rate: float) -> None:
|
||||
from typing import Final, cast
|
||||
from litellm.litellm_core_utils.llm_cost_calc.utils import _get_token_base_cost
|
||||
|
||||
info: Final = cast(
|
||||
ModelInfo,
|
||||
{
|
||||
"input_cost_per_token": 0.001,
|
||||
"output_cost_per_token": 0.002,
|
||||
"cache_read_input_token_cost": 0.003,
|
||||
"cache_creation_input_token_cost": 0.004,
|
||||
"input_cost_per_token_above_200k_tokens": 0.005,
|
||||
"output_cost_per_token_above_200k_tokens": 0.006,
|
||||
"cache_creation_input_token_cost_above_200k_tokens": 0.007,
|
||||
"cache_read_input_token_cost_above_512k_tokens": higher_rate,
|
||||
"output_cost_per_token_above_512k_tokens": higher_rate,
|
||||
},
|
||||
)
|
||||
rates: Final = _get_token_base_cost(
|
||||
model_info=info,
|
||||
usage=Usage(prompt_tokens=prompt_tokens, completion_tokens=3, total_tokens=prompt_tokens + 3),
|
||||
)
|
||||
assert rates[0] == (0.005 if prompt_tokens > 200_000 else 0.001)
|
||||
assert rates[1] == (higher_rate if prompt_tokens > 512_000 else 0.006 if prompt_tokens > 200_000 else 0.002)
|
||||
assert rates[2] == (0.007 if prompt_tokens > 200_000 else 0.004)
|
||||
assert rates[4] == (higher_rate if prompt_tokens > 512_000 else 0.003)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("rate", (0.0, 0.001))
|
||||
def test_token_breakdown_uses_supplied_pricing_without_catalog_lookup(rate: float) -> None:
|
||||
model_info: Final[ModelInfo] = {
|
||||
"input_cost_per_token": 0.01,
|
||||
"output_cost_per_token": 0.02,
|
||||
"cache_read_input_token_cost": rate,
|
||||
"cache_creation_input_token_cost": rate,
|
||||
}
|
||||
usage: Final = Usage(
|
||||
prompt_tokens=20, completion_tokens=3, total_tokens=23,
|
||||
prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=5), cache_creation_input_tokens=5,
|
||||
)
|
||||
prompt_cost, output_cost = generic_cost_per_token(
|
||||
model="supplied-pricing-test", usage=usage, custom_llm_provider="bedrock", model_info=model_info,
|
||||
)
|
||||
breakdown: Final = get_token_type_cost_breakdown(
|
||||
model="supplied-pricing-test", usage=usage, custom_llm_provider="bedrock", model_info=model_info,
|
||||
)
|
||||
assert breakdown.cache_read_cost == pytest.approx(5 * rate)
|
||||
assert breakdown.cache_creation_cost == pytest.approx(5 * rate)
|
||||
assert breakdown.rates is not None
|
||||
assert prompt_cost == pytest.approx(10 * breakdown.rates.input_cost_per_token + breakdown.cache_read_cost + breakdown.cache_creation_cost)
|
||||
assert output_cost == pytest.approx(3 * breakdown.rates.output_cost_per_token)
|
||||
|
|
|
|||
|
|
@ -363,15 +363,15 @@ def test_web_search_provider_prefix_fallback_does_not_misprice_non_gemini_model(
|
|||
Regression for the provider-prefix fallback in _handle_web_search_cost. When the initial
|
||||
get_model_info lookup fails for a "/"-containing model, the retry re-resolves model_info from
|
||||
the prefix and must adopt that prefix's provider for routing. Otherwise an unrelated model
|
||||
(here OpenRouter, which carries no web search pricing) is re-resolved but still routed through
|
||||
(here Azure, which carries no web search pricing) is re-resolved but still routed through
|
||||
the request's vertex_ai Gemini calculator, which charges its $0.035 per_prompt default for a
|
||||
model that should cost nothing for web search.
|
||||
"""
|
||||
from litellm.types.utils import PromptTokensDetailsWrapper, Usage
|
||||
|
||||
model = "openrouter/google/gemini-3.1-flash-lite"
|
||||
model = "azure/gpt-4o"
|
||||
model_info = litellm.get_model_info(model)
|
||||
assert model_info["litellm_provider"] == "openrouter"
|
||||
assert model_info["litellm_provider"] == "azure"
|
||||
assert not model_info.get("search_context_cost_per_query")
|
||||
|
||||
usage = Usage(
|
||||
|
|
|
|||
|
|
@ -17,6 +17,7 @@ import pytest
|
|||
import litellm
|
||||
import litellm.litellm_core_utils.llm_response_utils.response_metadata as response_metadata_mod
|
||||
import litellm.proxy.common_request_processing as common_request_processing_mod
|
||||
from litellm.cost_calculator import _select_model_name_for_cost_calc
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging
|
||||
from litellm.litellm_core_utils.llm_response_utils.response_metadata import (
|
||||
ResponseMetadata,
|
||||
|
|
@ -651,3 +652,38 @@ def test_update_response_metadata_prices_per_second_deployment_from_its_stamped_
|
|||
|
||||
assert result._response_ms == pytest.approx(2000)
|
||||
assert result._hidden_params["response_cost"] == pytest.approx((0.02 + 0.04) * 2)
|
||||
|
||||
|
||||
def test_update_response_metadata_keeps_the_provider_model_after_the_proxy_rewrites_response_model():
|
||||
start_time: Final = datetime.datetime(2026, 9, 21, 12, 0, 0)
|
||||
logging_obj: Final = Logging(
|
||||
model="azure/ptu-deployment",
|
||||
messages=[{"role": "user", "content": "Hello"}],
|
||||
stream=False,
|
||||
call_type="completion",
|
||||
start_time=start_time,
|
||||
litellm_call_id="provider-model-response-metadata",
|
||||
function_id="f",
|
||||
)
|
||||
logging_obj.update_environment_variables(
|
||||
model="azure/ptu-deployment", litellm_params={}, optional_params={}, custom_llm_provider="azure"
|
||||
)
|
||||
result: Final = ModelResponse(
|
||||
model="gpt-5.4-nano-2026-03-17",
|
||||
usage=Usage(prompt_tokens=11, completion_tokens=7, total_tokens=18),
|
||||
)
|
||||
|
||||
update_response_metadata(
|
||||
result=result,
|
||||
logging_obj=logging_obj,
|
||||
model="azure/ptu-deployment",
|
||||
kwargs={},
|
||||
start_time=start_time,
|
||||
end_time=start_time + datetime.timedelta(seconds=2),
|
||||
)
|
||||
result.model = "model-group-alias"
|
||||
|
||||
assert (
|
||||
_select_model_name_for_cost_calc(model=None, completion_response=result, custom_llm_provider="azure")
|
||||
== "azure/gpt-5.4-nano-2026-03-17"
|
||||
)
|
||||
|
|
|
|||
|
|
@ -970,19 +970,16 @@ def test_shipped_tool_search_rule_version_boundaries(shipped_cost_map, model, pr
|
|||
|
||||
def test_shipped_tool_search_rule_fills_mapped_claude_entries_without_flag(shipped_cost_map):
|
||||
"""A mapped Claude 4.5+ entry with no supports_tool_search key gets it from the rule
|
||||
on Anthropic direct, Vertex and Bedrock, a mapped pre-4.5 entry stays without one,
|
||||
and Azure Foundry and reseller copies of the same model are not touched."""
|
||||
on Anthropic direct, Azure Foundry, Vertex and Bedrock, a mapped pre-4.5 entry stays
|
||||
without one, and reseller copies of the same model are not touched."""
|
||||
for key, model, provider in (
|
||||
("claude-opus-4-7", "claude-opus-4-7", "anthropic"),
|
||||
("azure_ai/claude-opus-5", "claude-opus-5", "azure_ai"),
|
||||
("vertex_ai/claude-opus-5", "claude-opus-5", "vertex_ai"),
|
||||
):
|
||||
assert "supports_tool_search" not in litellm.model_cost[key]
|
||||
assert litellm.get_model_info(model, custom_llm_provider=provider)["supports_tool_search"] is True
|
||||
|
||||
assert "supports_tool_search" not in litellm.model_cost["azure_ai/claude-opus-5"]
|
||||
azure_opus_5_info = litellm.get_model_info("claude-opus-5", custom_llm_provider="azure_ai")
|
||||
assert azure_opus_5_info.get("supports_tool_search") is None
|
||||
|
||||
assert match_fill_missing_generalizations("claude-opus-5", "bedrock")["supports_tool_search"] is True
|
||||
assert "supports_tool_search" not in match_fill_missing_generalizations("claude-opus-5", "azure_ai")
|
||||
assert match_fill_missing_generalizations("claude-opus-5", "azure_ai")["supports_tool_search"] is True
|
||||
assert match_fill_missing_generalizations("claude-opus-5", "perplexity") is None
|
||||
|
|
|
|||
|
|
@ -212,7 +212,6 @@ def test_shipped_backup_marks_claude_4_6_plus_adaptive_not_4_0():
|
|||
"anthropic.claude-opus-4-8",
|
||||
"vertex_ai/claude-opus-4-6@default",
|
||||
"us.anthropic.claude-sonnet-4-6",
|
||||
"openrouter/anthropic/claude-opus-4.7",
|
||||
"azure_ai/claude-opus-4-7",
|
||||
]:
|
||||
assert backup[adaptive]["supports_adaptive_thinking"] is True, adaptive
|
||||
|
|
|
|||
|
|
@ -20,13 +20,14 @@ from openai._legacy_response import HttpxBinaryResponseContent
|
|||
|
||||
import litellm
|
||||
from litellm._logging import session_id_var, trace_id_var
|
||||
from litellm.constants import SENTRY_DENYLIST, SENTRY_PII_DENYLIST
|
||||
from litellm.constants import LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL, SENTRY_DENYLIST, SENTRY_PII_DENYLIST
|
||||
from litellm.cost_calculator import ocr_batch_cost
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LitellmLogging
|
||||
from litellm.litellm_core_utils.litellm_logging import (
|
||||
_extract_response_obj_and_hidden_params,
|
||||
_get_status_fields,
|
||||
get_standard_logging_object_payload,
|
||||
set_callbacks,
|
||||
)
|
||||
from litellm.llms.base_llm.ocr.transformation import OCRUsageInfo
|
||||
|
|
@ -346,6 +347,20 @@ def test_get_masked_api_base(logging_obj):
|
|||
assert type(masked_api_base) == str
|
||||
|
||||
|
||||
@pytest.mark.parametrize("no_upstream_call", [False, True])
|
||||
def test_pre_call_stamps_the_provider_handoff_only_when_a_provider_call_is_made(logging_obj, no_upstream_call):
|
||||
"""The proxy fires pre_call for a failure it raised itself (auth, rate limit, client disconnect)
|
||||
so callbacks see it; that call never reached a provider, so neither handoff instant may be set."""
|
||||
logging_obj.update_environment_variables(litellm_params={}, optional_params={})
|
||||
if no_upstream_call:
|
||||
logging_obj.model_call_details[LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL] = True
|
||||
|
||||
logging_obj.pre_call(input="", api_key="")
|
||||
|
||||
assert ("api_call_start_time" in logging_obj.model_call_details) is not no_upstream_call
|
||||
assert (logging_obj.model_call_details.get("first_api_call_start_time") is not None) is not no_upstream_call
|
||||
|
||||
|
||||
def test_pre_call_tolerates_missing_api_base(logging_obj):
|
||||
"""Presigned batch retrieves (Mistral, Bedrock) build their own URL and pass api_base=None
|
||||
to pre_call; masking must not raise or the request's pre-call logging is silently lost."""
|
||||
|
|
@ -1329,6 +1344,56 @@ class TestRetrieveBatchCostPassesModelIdentity:
|
|||
assert captured["model_info"] is not None
|
||||
assert captured["model_info"]["input_cost_per_token"] == 0.0
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_forwards_vertex_location_of_the_deployment(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
from litellm.litellm_core_utils import litellm_logging as logging_module
|
||||
|
||||
captured: dict[str, object] = {}
|
||||
|
||||
from litellm.batches.batch_utils import BatchCostUsageResult
|
||||
from litellm.types.utils import LiteLLMBatch, Usage
|
||||
|
||||
async def fake_handle_completed_batch(**kwargs: object) -> BatchCostUsageResult:
|
||||
captured.update(kwargs)
|
||||
return BatchCostUsageResult(
|
||||
cost=1.25,
|
||||
usage=Usage(prompt_tokens=1800, completion_tokens=1000, total_tokens=2800),
|
||||
models=["gemini-3.5-flash"],
|
||||
successful_requests=1,
|
||||
failed_requests=0,
|
||||
)
|
||||
|
||||
monkeypatch.setattr(logging_module, "_handle_completed_batch", fake_handle_completed_batch)
|
||||
|
||||
obj = LitellmLogging(
|
||||
model="vertex_ai/gemini-3.5-flash",
|
||||
messages=[{"role": "user", "content": "Hey"}],
|
||||
stream=False,
|
||||
call_type="aretrieve_batch",
|
||||
start_time=time.time(),
|
||||
litellm_call_id="batch-call-2",
|
||||
function_id="f",
|
||||
)
|
||||
obj.custom_llm_provider = "vertex_ai"
|
||||
obj.litellm_params = {"vertex_project": "proj-1", "vertex_location": "us-central1"}
|
||||
|
||||
batch = LiteLLMBatch(
|
||||
id="batch_vertex",
|
||||
completion_window="24h",
|
||||
created_at=1,
|
||||
endpoint="/v1/chat/completions",
|
||||
input_file_id="file-in",
|
||||
object="batch",
|
||||
status="completed",
|
||||
output_file_id="file-out",
|
||||
)
|
||||
|
||||
with contextlib.suppress(Exception):
|
||||
await obj._async_success_handler_body(result=batch, start_time=None, end_time=None)
|
||||
|
||||
assert captured, "handle_completed_batch was never called"
|
||||
assert captured["vertex_location"] == "us-central1"
|
||||
|
||||
|
||||
class TestRetrieveBatchPricesOnlyFinalBatches:
|
||||
"""Regression (LIT-7048): retrieving a provider-id batch priced it on every poll.
|
||||
|
|
@ -6794,10 +6859,10 @@ async def test_streaming_success_callbacks_survive_guardrail_logging_hook_failur
|
|||
|
||||
def _resolve(custom_llm_provider, litellm_params, optional_params, model):
|
||||
from litellm.litellm_core_utils.litellm_logging import (
|
||||
_resolve_vertex_location_for_cost,
|
||||
resolve_vertex_location_for_cost,
|
||||
)
|
||||
|
||||
return _resolve_vertex_location_for_cost(
|
||||
return resolve_vertex_location_for_cost(
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
litellm_params=litellm_params,
|
||||
optional_params=optional_params,
|
||||
|
|
@ -8336,6 +8401,7 @@ class TestAzurePTUSpilloverCost:
|
|||
|
||||
ROUTER_MODEL_ID: Final = "ptu-spill-router-model-id"
|
||||
SERVED_MODEL: Final = "azure/spill-served-model-ptu"
|
||||
PTU_DEPLOYMENT: Final = "azure/spill-ptu-deployment"
|
||||
PTU_MODEL_INFO: Final = {
|
||||
"id": ROUTER_MODEL_ID,
|
||||
"team_id": "team-1",
|
||||
|
|
@ -8371,7 +8437,9 @@ class TestAzurePTUSpilloverCost:
|
|||
litellm.model_cost.pop(TestAzurePTUSpilloverCost.ROUTER_MODEL_ID, None)
|
||||
litellm.model_cost.pop(TestAzurePTUSpilloverCost.SERVED_MODEL, None)
|
||||
|
||||
def _logging_obj(self, model_info: dict, *, flag: str, litellm_rate: float, monkeypatch) -> LitellmLogging:
|
||||
def _logging_obj(
|
||||
self, model_info: dict, *, flag: str, litellm_rate: float, monkeypatch, deployment: str | None = None
|
||||
) -> LitellmLogging:
|
||||
monkeypatch.setenv("LITELLM_ENABLE_PTU_COST_ATTRIBUTION", flag)
|
||||
obj = LitellmLogging(
|
||||
model=self.SERVED_MODEL,
|
||||
|
|
@ -8388,7 +8456,7 @@ class TestAzurePTUSpilloverCost:
|
|||
optional_params={},
|
||||
litellm_params={
|
||||
"api_base": "",
|
||||
"metadata": {"model_info": model_info},
|
||||
"metadata": {"model_info": model_info, "deployment": deployment},
|
||||
"input_cost_per_token": litellm_rate,
|
||||
"output_cost_per_token": litellm_rate,
|
||||
},
|
||||
|
|
@ -8458,6 +8526,84 @@ class TestAzurePTUSpilloverCost:
|
|||
finally:
|
||||
self._unregister_models()
|
||||
|
||||
def test_spillover_payload_reports_the_pay_as_you_go_entry_billing_used(self, monkeypatch) -> None:
|
||||
self._register_models()
|
||||
try:
|
||||
obj = self._logging_obj(dict(self.PTU_MODEL_INFO), flag="True", litellm_rate=0.0, monkeypatch=monkeypatch)
|
||||
response = self._response()
|
||||
response._hidden_params["additional_headers"] = {"llm_provider-x-ms-is-spilled-over": "true"}
|
||||
response_cost: Final = obj._response_cost_calculator(result=response)
|
||||
|
||||
payload: Final = get_standard_logging_object_payload(
|
||||
kwargs={
|
||||
"model": self.SERVED_MODEL,
|
||||
"custom_llm_provider": "azure",
|
||||
"litellm_params": obj.litellm_params,
|
||||
"response_cost": response_cost,
|
||||
},
|
||||
init_response_obj=response,
|
||||
start_time=datetime.datetime.now(),
|
||||
end_time=datetime.datetime.now(),
|
||||
logging_obj=obj,
|
||||
status="success",
|
||||
)
|
||||
|
||||
assert response_cost == pytest.approx(self.EXPECTED_SPILL_COST)
|
||||
assert payload is not None
|
||||
model_map_value: Final = payload["model_map_information"]["model_map_value"]
|
||||
assert model_map_value is not None
|
||||
assert model_map_value["input_cost_per_token"] == litellm.model_cost[self.SERVED_MODEL]["input_cost_per_token"]
|
||||
assert model_map_value["output_cost_per_token"] == litellm.model_cost[self.SERVED_MODEL]["output_cost_per_token"]
|
||||
assert response_cost == pytest.approx(
|
||||
100 * model_map_value["input_cost_per_token"] + 50 * model_map_value["output_cost_per_token"]
|
||||
)
|
||||
finally:
|
||||
self._unregister_models()
|
||||
|
||||
def test_spillover_without_base_model_reports_the_pay_as_you_go_rates_billing_used(self, monkeypatch) -> None:
|
||||
self._register_models()
|
||||
litellm.model_cost[self.PTU_DEPLOYMENT] = {"litellm_provider": "azure", "mode": "chat"}
|
||||
try:
|
||||
obj = self._logging_obj(
|
||||
dict(self.PTU_MODEL_INFO),
|
||||
flag="True",
|
||||
litellm_rate=0.0,
|
||||
monkeypatch=monkeypatch,
|
||||
deployment=self.PTU_DEPLOYMENT,
|
||||
)
|
||||
response = self._response()
|
||||
response._hidden_params["additional_headers"] = {"llm_provider-x-ms-is-spilled-over": "true"}
|
||||
response._hidden_params["provider_response_model"] = response.model
|
||||
response_cost: Final = obj._response_cost_calculator(result=response)
|
||||
response.model = "spill-model-group-alias"
|
||||
|
||||
payload: Final = get_standard_logging_object_payload(
|
||||
kwargs={
|
||||
"model": self.SERVED_MODEL,
|
||||
"custom_llm_provider": "azure",
|
||||
"litellm_params": obj.litellm_params,
|
||||
"response_cost": response_cost,
|
||||
},
|
||||
init_response_obj=response,
|
||||
start_time=datetime.datetime.now(),
|
||||
end_time=datetime.datetime.now(),
|
||||
logging_obj=obj,
|
||||
status="success",
|
||||
)
|
||||
|
||||
assert response_cost == pytest.approx(self.EXPECTED_SPILL_COST)
|
||||
assert payload is not None
|
||||
model_map_information: Final = payload["model_map_information"]
|
||||
assert model_map_information["model_map_key"] != self.PTU_DEPLOYMENT
|
||||
model_map_value: Final = model_map_information["model_map_value"]
|
||||
assert model_map_value is not None
|
||||
assert response_cost == pytest.approx(
|
||||
100 * model_map_value["input_cost_per_token"] + 50 * model_map_value["output_cost_per_token"]
|
||||
)
|
||||
finally:
|
||||
litellm.model_cost.pop(self.PTU_DEPLOYMENT, None)
|
||||
self._unregister_models()
|
||||
|
||||
def test_spillover_header_does_not_touch_non_ptu_custom_pricing(self, monkeypatch) -> None:
|
||||
self._register_models()
|
||||
custom_model_id: Final = "non-ptu-custom-router-model-id"
|
||||
|
|
@ -8855,6 +9001,89 @@ async def test_prompt_management_with_unchanged_variables_replays_a_byte_identic
|
|||
assert len(messages_n_plus_one) == len(messages_n) + 2
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("input_cost_per_token", "output_cost_per_token"),
|
||||
[(1e-06, 2e-06), (0.0, 0.0)],
|
||||
ids=["custom_price", "explicit_zero_price"],
|
||||
)
|
||||
def test_standard_logging_payload_reports_the_deployment_entry_the_request_was_billed_from(
|
||||
input_cost_per_token: float, output_cost_per_token: float
|
||||
):
|
||||
"""A deployment with its own price keeps the catalog model as model_map_key but is logged at its own rates."""
|
||||
deployment_id: Final = "lit9660-priced-deployment"
|
||||
deployment_model_info: Final = {
|
||||
"id": deployment_id,
|
||||
"input_cost_per_token": input_cost_per_token,
|
||||
"output_cost_per_token": output_cost_per_token,
|
||||
}
|
||||
litellm.register_model(
|
||||
model_cost={
|
||||
deployment_id: {
|
||||
"litellm_provider": "openai",
|
||||
"mode": "chat",
|
||||
"input_cost_per_token": input_cost_per_token,
|
||||
"output_cost_per_token": output_cost_per_token,
|
||||
}
|
||||
},
|
||||
persist_across_reloads=False,
|
||||
)
|
||||
try:
|
||||
logging_obj: Final = LitellmLogging(
|
||||
model="gpt-5.4-nano",
|
||||
messages=[{"role": "user", "content": "Hi"}],
|
||||
stream=False,
|
||||
call_type="completion",
|
||||
start_time=time.time(),
|
||||
litellm_call_id="lit9660",
|
||||
function_id="fn",
|
||||
)
|
||||
litellm_params: Final = {
|
||||
"metadata": {
|
||||
"model_group": "lit9660-group",
|
||||
"deployment": "openai/gpt-5.4-nano",
|
||||
"model_info": deployment_model_info,
|
||||
}
|
||||
}
|
||||
logging_obj.update_environment_variables(
|
||||
model="gpt-5.4-nano", user="", optional_params={}, litellm_params=litellm_params, custom_llm_provider="openai"
|
||||
)
|
||||
response: Final = ModelResponse(
|
||||
model="gpt-5.4-nano",
|
||||
choices=[litellm.Choices(message=litellm.Message(role="assistant", content="hi"))],
|
||||
usage=litellm.Usage(prompt_tokens=10, completion_tokens=6, total_tokens=16),
|
||||
)
|
||||
|
||||
response_cost: Final = logging_obj._response_cost_calculator(result=response)
|
||||
payload: Final = get_standard_logging_object_payload(
|
||||
kwargs={
|
||||
"model": "gpt-5.4-nano",
|
||||
"custom_llm_provider": "openai",
|
||||
"litellm_params": litellm_params,
|
||||
"response_cost": response_cost,
|
||||
},
|
||||
init_response_obj=response,
|
||||
start_time=datetime.datetime.now(),
|
||||
end_time=datetime.datetime.now(),
|
||||
logging_obj=logging_obj,
|
||||
status="success",
|
||||
)
|
||||
|
||||
assert response_cost == 10 * input_cost_per_token + 6 * output_cost_per_token
|
||||
assert payload is not None
|
||||
assert payload["model_map_information"] == {
|
||||
"model_map_key": "openai/gpt-5.4-nano",
|
||||
"model_map_value": litellm.get_model_info(model=f"openai/{deployment_id}", custom_llm_provider="openai"),
|
||||
}
|
||||
model_map_value: Final = payload["model_map_information"]["model_map_value"]
|
||||
assert model_map_value is not None
|
||||
assert (model_map_value["input_cost_per_token"], model_map_value["output_cost_per_token"]) == (
|
||||
input_cost_per_token,
|
||||
output_cost_per_token,
|
||||
)
|
||||
finally:
|
||||
litellm.model_cost.pop(deployment_id, None)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_background_interaction_completion_logs_while_in_progress_handler_is_parked(monkeypatch):
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -29,7 +29,7 @@ from litellm.llms.bedrock.chat.invoke_transformations.anthropic_claude3_transfor
|
|||
from litellm.llms.vertex_ai.vertex_ai_partner_models.anthropic.transformation import (
|
||||
VertexAIAnthropicConfig,
|
||||
)
|
||||
from litellm.types.llms.anthropic import ANTHROPIC_BETA_HEADER_VALUES
|
||||
from litellm.types.llms.anthropic import ANTHROPIC_BETA_HEADER_VALUES, ANTHROPIC_TOOL_SEARCH_TOOL_TYPES
|
||||
from litellm.types.utils import ServerToolUse, Usage
|
||||
|
||||
|
||||
|
|
@ -6798,3 +6798,91 @@ def test_chat_dummy_tool_result_for_an_orphaned_tool_call_replays_a_byte_identic
|
|||
_assert_prefix_stable(requests)
|
||||
assert [m["role"] for m in requests[0]["messages"]] == ["user", "assistant", "user"]
|
||||
assert requests[0]["messages"][2]["content"][0]["type"] == "tool_result"
|
||||
|
||||
|
||||
def _make_transform_request(optional_params: dict, litellm_params: dict) -> dict:
|
||||
return AnthropicConfig().transform_request(
|
||||
model="claude-3-5-sonnet-20241022",
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
optional_params=optional_params,
|
||||
litellm_params=litellm_params,
|
||||
headers={},
|
||||
)
|
||||
|
||||
|
||||
def _two_function_tools() -> list[dict]:
|
||||
return [
|
||||
{
|
||||
"type": "function",
|
||||
"function": {"name": f"tool_{i}", "description": f"Tool {i}", "parameters": {"type": "object"}},
|
||||
}
|
||||
for i in range(2)
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"control, expected_control",
|
||||
[
|
||||
(None, {"type": "ephemeral"}),
|
||||
({"type": "ephemeral", "ttl": "1h"}, {"type": "ephemeral", "ttl": "1h"}),
|
||||
],
|
||||
)
|
||||
def test_transform_request_tool_config_point_marks_last_tool_and_drops_the_key(control, expected_control):
|
||||
"""LIT-9880: a tool_config injection point used to be spread into the Anthropic body, which
|
||||
Anthropic rejected with 400 Extra inputs are not permitted."""
|
||||
tools = _two_function_tools()
|
||||
tools_before = copy.deepcopy(tools)
|
||||
point = {"location": "tool_config", **({"control": control} if control is not None else {})}
|
||||
litellm_params = {"metadata": {}, "model_info": {"id": "deployment-1"}}
|
||||
|
||||
data = _make_transform_request(
|
||||
optional_params={"tools": tools, "cache_control_injection_points": [point]},
|
||||
litellm_params=litellm_params,
|
||||
)
|
||||
|
||||
assert "cache_control_injection_points" not in data
|
||||
assert data["tools"][-1]["cache_control"] == expected_control
|
||||
assert "cache_control" not in data["tools"][0]
|
||||
assert tools == tools_before
|
||||
assert litellm_params["metadata"]["litellm_gateway_injected_cache"] == "deployment-1"
|
||||
|
||||
|
||||
def test_transform_request_tool_config_point_keeps_the_clients_tool_marker():
|
||||
tools = _two_function_tools()
|
||||
tools[0]["cache_control"] = {"type": "ephemeral"}
|
||||
|
||||
data = _make_transform_request(
|
||||
optional_params={"tools": tools, "cache_control_injection_points": [{"location": "tool_config"}]},
|
||||
litellm_params={},
|
||||
)
|
||||
|
||||
assert "cache_control_injection_points" not in data
|
||||
assert data["tools"][0]["cache_control"] == {"type": "ephemeral"}
|
||||
assert "cache_control" not in data["tools"][1]
|
||||
|
||||
|
||||
def test_transform_request_tool_config_point_skips_tool_search_tools():
|
||||
tool_search_type = next(iter(ANTHROPIC_TOOL_SEARCH_TOOL_TYPES))
|
||||
tools = [*_two_function_tools(), {"type": tool_search_type, "name": "tool_search"}]
|
||||
|
||||
data = _make_transform_request(
|
||||
optional_params={"tools": tools, "cache_control_injection_points": [{"location": "tool_config"}]},
|
||||
litellm_params={},
|
||||
)
|
||||
|
||||
assert data["tools"][1]["cache_control"] == {"type": "ephemeral"}
|
||||
assert "cache_control" not in data["tools"][2]
|
||||
assert "cache_control" not in data["tools"][0]
|
||||
|
||||
|
||||
def test_transform_request_tool_config_point_without_tools_drops_the_key():
|
||||
litellm_params = {"metadata": {}}
|
||||
|
||||
data = _make_transform_request(
|
||||
optional_params={"cache_control_injection_points": [{"location": "tool_config"}]},
|
||||
litellm_params=litellm_params,
|
||||
)
|
||||
|
||||
assert "cache_control_injection_points" not in data
|
||||
assert "tools" not in data
|
||||
assert litellm_params["metadata"] == {}
|
||||
|
|
|
|||
Some files were not shown because too many files have changed in this diff Show more
Loading…
Add table
Reference in a new issue