diff --git a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py index bd9440d4754..7290bd7d528 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py @@ -10,6 +10,7 @@ from dataclasses import replace as dataclasses_replace from datetime import datetime, timedelta, timezone from types import MappingProxyType from typing import TYPE_CHECKING, Final, List, Literal, Optional, Protocol, Tuple, cast +from pydantic import TypeAdapter, ValidationError from litellm_enterprise.proxy.common_utils.managed_object_attribution import ( ManagedObjectCreatorAttribution, @@ -46,6 +47,7 @@ def _error_names_resource(error: Exception, resource_id: str) -> bool: CHECK_BATCH_COST_USER_AGENT = "LiteLLM Proxy/CheckBatchCost" RECENT_BATCH_POLL_WINDOW: Final = timedelta(hours=48) +_FILE_OBJECT_ADAPTER: Final[TypeAdapter[dict[str, object]]] = TypeAdapter(dict[str, object]) PROVIDER_TERMINAL_BATCH_STATUSES: Final[Tuple[str, ...]] = ( "completed", @@ -500,6 +502,10 @@ class CheckBatchCost: llm_provider: str, bare_model_name: str, ) -> Optional[Tuple[str, str]]: + stored_model_id: Final = self._get_stored_model_id(job) + if stored_model_id is not None and self._batch_deployment_exists(stored_model_id): + return stored_model_id, job.unified_object_id + deployment_id = self._get_deployment_id_for_bare_model(bare_model_name, llm_provider) if deployment_id is None: verbose_proxy_logger.info( @@ -611,23 +617,36 @@ class CheckBatchCost: ) @staticmethod - def _get_input_file_id(job: "ManagedObjectRow") -> Optional[str]: - from litellm.types.utils import LiteLLMBatch - + def _parse_file_object(job: "ManagedObjectRow") -> Optional[dict[str, object]]: file_object: Final = job.file_object try: - parsed_file_object: Final[object] = ( + return _FILE_OBJECT_ADAPTER.validate_python( json.loads(file_object) if isinstance(file_object, str) else file_object ) - except (json.JSONDecodeError, ValueError): + except (json.JSONDecodeError, ValueError, ValidationError): return None - if not isinstance(parsed_file_object, dict): + + @classmethod + def _get_input_file_id(cls, job: "ManagedObjectRow") -> Optional[str]: + from litellm.types.utils import LiteLLMBatch + + parsed_file_object: Final = cls._parse_file_object(job) + if parsed_file_object is None: return None try: return LiteLLMBatch.model_validate(parsed_file_object).input_file_id except Exception: return None + @classmethod + def _get_stored_model_id(cls, job: "ManagedObjectRow") -> Optional[str]: + """The deployment that created the batch, written into file_object by the managed + files hook at create time. Older rows have no such key and fall back to matching + the bare model name against the router.""" + parsed_file_object: Final = cls._parse_file_object(job) + model_id: Final = parsed_file_object.get("model_id") if parsed_file_object else None + return model_id if isinstance(model_id, str) and model_id else None + async def _track_completed_batch_cost( self, job: "ManagedObjectRow", diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index 439770390f8..0819d3ded07 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py @@ -19,6 +19,7 @@ from typing import ( TypedDict, Union, cast, + runtime_checkable, ) from uuid import NAMESPACE_URL, uuid5 @@ -366,6 +367,12 @@ class _ManagedResourcePrismaClient(Protocol): def db(self) -> _ManagedResourceDatabase: ... +@runtime_checkable +class _WriterRoutedDatabase(Protocol): + @property + def writer(self) -> _ManagedResourceDatabase: ... + + class _SchedulerWithJobLookup(Protocol): def get_job(self, job_id: str) -> object: ... @@ -398,6 +405,15 @@ def _managed_object_table(prisma_client: _ManagedResourcePrismaClient) -> _Manag return prisma_client.db.litellm_managedobjecttable +def _managed_database(prisma_client: _ManagedResourcePrismaClient) -> _ManagedResourceDatabase: + return prisma_client.db + + +def _writer_database(db: _ManagedResourceDatabase) -> _ManagedResourceDatabase: + """Reads that feed a write must see the create that just happened; a lagging read replica may not""" + return db.writer if isinstance(db, _WriterRoutedDatabase) else db + + _JSON_VALUE_ADAPTER: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue) _POSTGRES_UNSTORABLE_CHARACTER: Final = "\x00" @@ -431,6 +447,25 @@ def _without_model_output( ) +def _with_model_id( + file_object: LiteLLMBatch | LiteLLMFineTuningJob | ResponsesAPIResponse, + model_id: str | None, +) -> LiteLLMBatch | LiteLLMFineTuningJob | ResponsesAPIResponse: + return file_object if model_id is None else file_object.model_copy(update={"model_id": model_id}) + + +async def _stored_model_id(managed_object_table: _ManagedObjectTableActions, unified_object_id: str) -> str | None: + row: Final = await managed_object_table.find_first(where={"unified_object_id": unified_object_id}) + if row is None: + return None + try: + parsed: Final = _JSON_VALUE_ADAPTER.validate_python(_decode_json_blob(row.file_object)) + except (json.JSONDecodeError, ValidationError): + return None + model_id: Final = parsed.get("model_id") if isinstance(parsed, dict) else None + return model_id if isinstance(model_id, str) and model_id else None + + 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 {} @@ -599,9 +634,17 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): persist_attribution: bool = False, create_if_missing: bool = True, batch_processed: bool = False, + creating_model_id: str | None = None, ) -> None: """Persist a managed object row, caching it and upserting it in the DB. + creating_model_id is the router deployment that created the object. It is kept + inside the stored file_object JSON as `model_id`, because the raw provider id an + unmanaged batch stores as unified_object_id says nothing about which deployment + made it, and the cost poller needs that deployment to read the result back. A + later observe of the row (create_if_missing=False) keeps the stored value rather + than overwriting it with whichever deployment served the retrieve. + persist_attribution is set only by the batch and background response creates, the callers that can speak for the creator; it gates the api_key, request_tags, and, for a response, end_user_id columns that the cost poll jobs bill against, so a later poll or @@ -622,7 +665,16 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): It is written only in the upsert create branch. """ verbose_logger.info(f"Storing LiteLLM Managed {file_purpose} object with id={unified_object_id} in cache") - stored_object: Final = _without_model_output(file_object) + managed_database: Final = _managed_database(self.prisma_client) + managed_object_table: Final = managed_database.litellm_managedobjecttable + preserved_model_id: Final = ( + creating_model_id + if creating_model_id is not None or create_if_missing + else await _stored_model_id( + _writer_database(managed_database).litellm_managedobjecttable, unified_object_id + ) + ) + stored_object: Final = _with_model_id(_without_model_output(file_object), preserved_model_id) litellm_managed_object = LiteLLM_ManagedObjectTable( unified_object_id=unified_object_id, model_object_id=model_object_id, @@ -654,7 +706,6 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): "status": file_object.status, "updated_by": user_api_key_dict.user_id, } - managed_object_table: Final = _managed_object_table(self.prisma_client) if not create_if_missing: await managed_object_table.update_many( where={"unified_object_id": unified_object_id}, @@ -2069,6 +2120,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): request_tags=request_tags_from_metadata(request_metadata if isinstance(request_metadata, dict) else {}), persist_attribution=is_batch_create, create_if_missing=is_batch_create, + creating_model_id=model_id if is_batch_create else None, ) if not is_batch_create: await self.store_batch_output_file_ownership( diff --git a/enterprise/pyproject.toml b/enterprise/pyproject.toml index 17209a29e15..ccf47150881 100644 --- a/enterprise/pyproject.toml +++ b/enterprise/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "litellm-enterprise" -version = "0.1.79" +version = "0.1.79.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.79" +version = "0.1.79.post1" version_files = [ "pyproject.toml:^version", "../pyproject.toml:litellm-enterprise==", diff --git a/litellm/caching/evicted_client_closer.py b/litellm/caching/evicted_client_closer.py index b22cd3aecd1..6f29e224249 100644 --- a/litellm/caching/evicted_client_closer.py +++ b/litellm/caching/evicted_client_closer.py @@ -29,6 +29,9 @@ it does not belong to. Queued clients are therefore bucketed by what it takes to close them, and each bucket is ordered by deadline, so a reap walks the entries that are due rather than the whole queue. +A bucket whose loop has been collected is dropped, and its clients are left to the +collector, as they are when the queue is full. + The queue holds its clients weakly, so waiting out a grace window never keeps alive anything the collector would have reclaimed first. """ @@ -36,6 +39,7 @@ alive anything the collector would have reclaimed first. import asyncio import contextlib import inspect +import itertools import threading import time import weakref @@ -53,6 +57,9 @@ _CLOSABLE_ANYWHERE: Final = "closable-anywhere" _CLOSABLE_ON_ANY_LOOP: Final = "closable-on-any-loop" _BucketKey = str | int +_LOOP_TOKENS: Final[weakref.WeakKeyDictionary[asyncio.AbstractEventLoop, int]] = weakref.WeakKeyDictionary() +_LOOP_TOKEN_COUNTER: Final = itertools.count() +_LOOP_TOKEN_LOCK: Final = threading.Lock() @dataclass(frozen=True, slots=True) @@ -63,12 +70,12 @@ class _PendingClose: collector would otherwise have reclaimed first. ``needs_loop`` is set for a client whose close is a coroutine; those can only - be closed from the event loop they were evicted on, recorded in ``loop_id``. + be closed from the event loop they were evicted on, recorded in ``loop_token``. A client that closes synchronously carries neither constraint. """ client_ref: "weakref.ref[object]" - loop_id: int | None + loop_token: int | None needs_loop: bool close_after: float @@ -77,16 +84,24 @@ def _bucket_key(pending: _PendingClose) -> _BucketKey: """Which reaps can close this entry: any at all, any running a loop, or one loop's.""" if not pending.needs_loop: return _CLOSABLE_ANYWHERE - if pending.loop_id is None: + if pending.loop_token is None: return _CLOSABLE_ON_ANY_LOOP - return pending.loop_id + return pending.loop_token -def _running_loop_id() -> int | None: +def running_loop_token() -> int | None: try: - return id(asyncio.get_running_loop()) + loop: Final = asyncio.get_running_loop() except RuntimeError: return None + try: + token: Final[int | None] = _LOOP_TOKENS.get(loop) + if token is not None: + return token + with _LOOP_TOKEN_LOCK: + return _LOOP_TOKENS.setdefault(loop, next(_LOOP_TOKEN_COUNTER)) + except TypeError: + return id(loop) def _close_function(client: object) -> Callable[[], object] | None: @@ -167,6 +182,8 @@ class EvictedClientCloser: self._clock = clock self._owned: weakref.WeakSet[object] = weakref.WeakSet() self._buckets: dict[_BucketKey, deque[_PendingClose]] = {} # mutable-ok: deadline-ordered queues + self._dead_loop_tokens: deque[int] = deque() # mutable-ok: loop finalizers append without taking the queue lock + self._watched_loop_tokens: set[int] = set() # mutable-ok: tracks tokens until their dead buckets are discarded self._pending_count = 0 self._queue_lock = threading.Lock() # the cache is reachable from every worker thread's loop self._close_tasks: set[asyncio.Task[None]] = set() # mutable-ok: strong refs to running closes @@ -197,13 +214,19 @@ class EvictedClientCloser: close_fn: Final = _close_function(client) if close_fn is None: return + loop_token: Final = running_loop_token() + needs_loop: Final = inspect.iscoroutinefunction(close_fn) if self._pending_count >= self._max_pending: - return + self._discard_dead_loops() + if self._pending_count >= self._max_pending: + return + if needs_loop and loop_token is not None: + self._watch_loop(loop_token) self._enqueue( _PendingClose( client_ref=weakref.ref(client), - loop_id=_running_loop_id(), - needs_loop=inspect.iscoroutinefunction(close_fn), + loop_token=loop_token, + needs_loop=needs_loop, close_after=self._clock() + self._grace_seconds, ) ) @@ -217,7 +240,7 @@ class EvictedClientCloser: if not self._pending_count: return now: Final = self._clock() - for pending in self._take_due(_running_loop_id(), now): + for pending in self._take_due(running_loop_token(), now): client = pending.client_ref() if client is None: continue @@ -230,6 +253,29 @@ class EvictedClientCloser: def pending_count(self) -> int: return self._pending_count + def _watch_loop(self, loop_token: int) -> None: + if loop_token in self._watched_loop_tokens: + return + try: + weakref.finalize(asyncio.get_running_loop(), self._dead_loop_tokens.append, loop_token) + except TypeError: + return + self._watched_loop_tokens.add(loop_token) + + def _discard_dead_loops(self) -> None: + with self._queue_lock: + self._discard_dead_loops_locked() + + def _discard_dead_loops_locked(self) -> None: + while self._dead_loop_tokens: + self._discard_dead_loop_token_locked(self._dead_loop_tokens.popleft()) + + def _discard_dead_loop_token_locked(self, loop_token: int) -> None: + self._watched_loop_tokens.discard(loop_token) + bucket: Final = self._buckets.pop(loop_token, None) + if bucket is not None: + self._pending_count -= len(bucket) + def _enqueue(self, pending: _PendingClose) -> None: """Append to the entry's bucket, dropping any dead entries it queues behind. @@ -238,6 +284,7 @@ class EvictedClientCloser: the front rather than having to be searched for. """ with self._queue_lock: + self._discard_dead_loops_locked() bucket: Final = self._buckets.setdefault(_bucket_key(pending), deque()) while bucket and bucket[0].client_ref() is None: bucket.popleft() @@ -245,9 +292,12 @@ class EvictedClientCloser: bucket.append(pending) self._pending_count += 1 - def _take_due(self, loop_id: int | None, now: float) -> tuple[_PendingClose, ...]: - buckets = (_CLOSABLE_ANYWHERE,) if loop_id is None else (_CLOSABLE_ANYWHERE, _CLOSABLE_ON_ANY_LOOP, loop_id) + def _take_due(self, loop_token: int | None, now: float) -> tuple[_PendingClose, ...]: + buckets = ( + (_CLOSABLE_ANYWHERE,) if loop_token is None else (_CLOSABLE_ANYWHERE, _CLOSABLE_ON_ANY_LOOP, loop_token) + ) with self._queue_lock: + self._discard_dead_loops_locked() return tuple(pending for key in buckets for pending in self._drain_locked(key, now)) def _drain_locked(self, key: _BucketKey, now: float) -> Iterator[_PendingClose]: diff --git a/litellm/caching/llm_caching_handler.py b/litellm/caching/llm_caching_handler.py index a89e43b78b4..ccc4b314e44 100644 --- a/litellm/caching/llm_caching_handler.py +++ b/litellm/caching/llm_caching_handler.py @@ -2,10 +2,9 @@ Add the event loop to the cache key, to prevent event loop closed errors. """ -import asyncio from typing import Final -from .evicted_client_closer import EvictedClientCloser, default_evicted_client_closer +from .evicted_client_closer import EvictedClientCloser, default_evicted_client_closer, running_loop_token from .in_memory_cache import InMemoryCache @@ -43,24 +42,22 @@ class LLMClientCache(InMemoryCache): self.evicted_client_closer.schedule(evicted) self.evicted_client_closer.reap() - def update_cache_key_with_event_loop(self, key): + def update_cache_key_with_event_loop(self, key: str) -> str: """ Add the event loop to the cache key, to prevent event loop closed errors. If none, use the key as is. """ - try: - event_loop: Final = asyncio.get_running_loop() - stringified_event_loop: Final = str(id(event_loop)) - return f"{key}-{stringified_event_loop}" - except RuntimeError: # handle no current running event loop - return key + token: Final = running_loop_token() + return f"{key}-{token}" if token is not None else key - def set_cache(self, key: str, value: object, litellm_owned_client: bool = False, **kwargs): + def _scoped_key(self, key: str, loop_bound: bool) -> str: + return self.update_cache_key_with_event_loop(key) if loop_bound else key + + def set_cache(self, key: str, value: object, litellm_owned_client: bool = False, loop_bound: bool = True, **kwargs): """``litellm_owned_client`` marks a client litellm built, so it may be closed once evicted.""" if litellm_owned_client: self.evicted_client_closer.mark_owned(value) - key = self.update_cache_key_with_event_loop(key) - return super().set_cache(key, value, **kwargs) + return super().set_cache(self._scoped_key(key, loop_bound), value, **kwargs) async def async_set_cache(self, key: str, value: object, litellm_owned_client: bool = False, **kwargs): if litellm_owned_client: @@ -68,13 +65,13 @@ class LLMClientCache(InMemoryCache): key = self.update_cache_key_with_event_loop(key) return await super().async_set_cache(key, value, **kwargs) - def get_cache(self, key, **kwargs): - key = self.update_cache_key_with_event_loop(key) + def get_cache(self, key: str, loop_bound: bool = True, **kwargs): + scoped_key: Final = self._scoped_key(key, loop_bound) self.evicted_client_closer.reap() - return super().get_cache(key, **kwargs) + return super().get_cache(scoped_key, **kwargs) - async def async_get_cache(self, key, **kwargs): + async def async_get_cache(self, key: str, **kwargs): key = self.update_cache_key_with_event_loop(key) return await super().async_get_cache(key, **kwargs) diff --git a/litellm/constants.py b/litellm/constants.py index ce002cc4ee8..9a994cd30b4 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -446,6 +446,7 @@ REDIS_DAILY_END_USER_SPEND_UPDATE_BUFFER_KEY: Final = "litellm_daily_end_user_sp REDIS_DAILY_AGENT_SPEND_UPDATE_BUFFER_KEY: Final = "litellm_daily_agent_spend_update_buffer" REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY: Final = "litellm_daily_tag_spend_update_buffer" REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY: Final = "litellm_window_spend_update_buffer" +REDIS_TOOL_DISCOVERY_BUFFER_KEY: Final = "litellm_tool_discovery_buffer" MAX_REDIS_BUFFER_DEQUEUE_COUNT: Final = get_env_int("MAX_REDIS_BUFFER_DEQUEUE_COUNT", 100) REDIS_SPEND_LOGS_BUFFER_KEY: Final = "litellm_spend_logs_buffer" REDIS_SPEND_LOGS_BUFFER_MAX_ROWS: Final = 100000 @@ -701,6 +702,7 @@ DD_TRACER_STREAMING_CHUNK_YIELD_RESOURCE: Final = os.getenv( ) LITELLM_HTTP_STATUS_CLIENT_DISCONNECTED: Final = 499 +LITELLM_CLIENT_DISCONNECT_DETAIL: Final = "Client disconnected the request" EMAIL_BUDGET_ALERT_TTL: Final = get_env_int("EMAIL_BUDGET_ALERT_TTL", 24 * 60 * 60) # 24 hours in seconds EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE: Final = get_env_float( @@ -1810,6 +1812,7 @@ LITELLM_UI_SESSION_DURATION: Final = os.getenv("LITELLM_UI_SESSION_DURATION", "2 ########################### DB CRON JOB NAMES ########################### DB_SPEND_UPDATE_JOB_NAME: Final = "db_spend_update_job" +TOOL_DISCOVERY_COMMIT_JOB_NAME: Final = "tool_discovery_commit_job" DB_DAILY_TAG_SPEND_UPDATE_JOB_NAME: Final = "db_daily_tag_spend_update_job" PROMETHEUS_EMIT_BUDGET_METRICS_JOB_NAME: Final = "prometheus_emit_budget_metrics" CLOUDZERO_EXPORT_USAGE_DATA_JOB_NAME: Final = "cloudzero_export_usage_data" @@ -1985,6 +1988,7 @@ SPEND_LOG_KEY_METADATA_CACHE_TTL: Final = 600 SPEND_LOG_KEY_METADATA_MISS_CACHE_TTL: Final = 30 SPEND_LOG_KEY_METADATA_CACHE_MAX_ITEMS: Final = 10000 SPEND_LOG_KEY_METADATA_QUERY_TIMEOUT_MS: Final = 5000 +SPEND_LOG_READ_STATEMENT_TIMEOUT_SECONDS: Final = 30 SPEND_LOG_KEY_METADATA_ROWS_PER_PROBE: Final = 100 # Short TTL for negative MCP access-group existence lookups. Keeps unauthenticated # callers from forcing a DB query per request for unknown names, while bounding diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index e05d3c4bfff..dcd9b905c9b 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -42,6 +42,8 @@ from litellm.litellm_core_utils.llm_cost_calc.utils import ( generic_cost_per_token, get_batch_cost_rates, get_billable_input_tokens, + get_cost_per_unit, + get_modality_token_cost_key, get_regional_uplift_multiplier, get_service_tier_cost_key, get_token_type_cost_breakdown, @@ -2832,12 +2834,17 @@ def _batch_rate( model_info: ModelInfo, key: Literal[ "input_cost_per_audio_token_batches", + "cache_read_input_audio_token_cost_batches", "input_cost_per_image_token_batches", "input_cost_per_video_token_batches", ], fallback: float, + *, + prompt_tokens: int = 0, ) -> float: - rate: Final = model_info.get(key) + context_key: Final = get_modality_token_cost_key(model_info, key, prompt_tokens) + context_rate: Final = get_cost_per_unit(model_info, context_key, None) + rate: Final = model_info.get(key) if context_rate is None else context_rate return fallback if rate is None else rate @@ -2965,6 +2972,7 @@ def _batch_prompt_cost( ) -> float: details: Final = parse_prompt_tokens_details(usage) cached_tokens: Final = details["cache_hit_tokens"] if cache_read_rate is not None else 0 + cached_audio_tokens: Final = details["cache_hit_audio_tokens"] if cache_read_rate is not None else 0 written_tokens: Final = details["cache_creation_tokens"] if cache_creation_rate is not None else 0 audio_tokens, image_tokens, video_tokens = ( details["audio_tokens"], @@ -2976,10 +2984,13 @@ def _batch_prompt_cost( ) return ( text_tokens * input_rate - + audio_tokens * _batch_rate(model_info, "input_cost_per_audio_token_batches", input_rate) + + audio_tokens + * _batch_rate(model_info, "input_cost_per_audio_token_batches", input_rate, prompt_tokens=usage.prompt_tokens) + image_tokens * _batch_rate(model_info, "input_cost_per_image_token_batches", input_rate) + video_tokens * _batch_rate(model_info, "input_cost_per_video_token_batches", input_rate) - + cached_tokens * (cache_read_rate or 0.0) + + (cached_tokens - cached_audio_tokens) * (cache_read_rate or 0.0) + + cached_audio_tokens + * _batch_rate(model_info, "cache_read_input_audio_token_cost_batches", cache_read_rate or 0.0) + written_tokens * (cache_creation_rate or 0.0) ) diff --git a/litellm/images/main.py b/litellm/images/main.py index be32ea4daed..5fd8e3e36bb 100644 --- a/litellm/images/main.py +++ b/litellm/images/main.py @@ -29,7 +29,6 @@ from litellm.litellm_core_utils.mock_functions import mock_image_generation from litellm.llms.base_llm import BaseImageEditConfig, BaseImageGenerationConfig from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler -from litellm.llms.custom_llm import CustomLLM from litellm.utils import exception_type, filter_out_litellm_params, get_litellm_params #################### Initialize provider clients #################### @@ -523,10 +522,7 @@ def image_generation( ) elif custom_llm_provider in litellm._custom_providers: # Assume custom LLM provider # Get the Custom Handler - custom_handler: CustomLLM | None = None - for item in litellm.custom_provider_map: - if item["provider"] == custom_llm_provider: - custom_handler = item["custom_handler"] + custom_handler: Final = ProviderConfigManager.get_custom_llm_handler(custom_llm_provider) if custom_handler is None: raise LiteLLMUnknownProvider(model=model, custom_llm_provider=custom_llm_provider) @@ -799,10 +795,7 @@ def image_edit( # Check for custom provider if custom_llm_provider in litellm._custom_providers: - custom_handler: CustomLLM | None = None - for item in litellm.custom_provider_map: - if item["provider"] == custom_llm_provider: - custom_handler = item["custom_handler"] + custom_handler: Final = ProviderConfigManager.get_custom_llm_handler(custom_llm_provider) if custom_handler is None: raise LiteLLMUnknownProvider(model=model, custom_llm_provider=custom_llm_provider) diff --git a/litellm/integrations/agentops/agentops.py b/litellm/integrations/agentops/agentops.py index 9043ab1853a..927b1b3c457 100644 --- a/litellm/integrations/agentops/agentops.py +++ b/litellm/integrations/agentops/agentops.py @@ -99,18 +99,14 @@ class AgentOps(OpenTelemetry): "Connection": "keep-alive", } - client: Final = get_httpx_client() - try: - response: Final = client.post( - url=auth_endpoint, - headers=headers, - json={"api_key": api_key}, - timeout=10, - ) + response: Final = get_httpx_client().post( + url=auth_endpoint, + headers=headers, + json={"api_key": api_key}, + timeout=10, + ) - if response.status_code != 200: - raise Exception(f"Failed to fetch auth token: {response.text}") + if response.status_code != 200: + raise Exception(f"Failed to fetch auth token: {response.text}") - return response.json() - finally: - client.close() + return response.json() diff --git a/litellm/litellm_core_utils/hidden_params.py b/litellm/litellm_core_utils/hidden_params.py index 714b258413e..4cb9dcd53f1 100644 --- a/litellm/litellm_core_utils/hidden_params.py +++ b/litellm/litellm_core_utils/hidden_params.py @@ -77,6 +77,11 @@ def get_hidden_params(obj: object) -> MutableMapping[str, object] | None: return _as_hidden_params_mapping(hidden_params) +def served_from_cache(obj: object) -> bool: + hidden_params: Final = get_hidden_params(obj) + return hidden_params is not None and hidden_params.get("cache_hit") is True + + def set_hidden_params(obj: object, hidden_params: dict[str, object] | HiddenParams) -> None: if isinstance(obj, dict): obj[_HIDDEN_PARAMS_ATTR] = hidden_params diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index eda1ed7e312..1d6250fc528 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -81,14 +81,19 @@ from litellm.litellm_core_utils.core_helpers import ( ) from litellm.litellm_core_utils.error_normalization import normalize_error from litellm.litellm_core_utils.get_litellm_params import get_litellm_params -from litellm.litellm_core_utils.hidden_params import HIDDEN_PARAMS_ATTR, set_hidden_param +from litellm.litellm_core_utils.hidden_params import HIDDEN_PARAMS_ATTR, served_from_cache, set_hidden_param from litellm.litellm_core_utils.internal_call_metadata import ( MODEL_ACCESS_GROUP_METADATA_KEY, is_unbilled_non_inference_call, ) from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import ( + bedrock_guardrail_trace, cost_breakdown_with_guardrail, guardrail_information_cost, + is_str_keyed_dict, + native_bedrock_guardrail_information, + object_mapping, + with_guardrail_entry, ) from litellm.litellm_core_utils.llm_cost_calc.tool_call_cost_tracking import ( StandardBuiltInToolCostTracking, @@ -2478,6 +2483,31 @@ class Logging(LiteLLMLoggingBaseClass): return self.model_call_details["response_headers"] = headers + def record_native_bedrock_guardrail(self, result: object) -> None: + """Bill the guardrail trace a Bedrock response carries when the deployment's ``guardrailConfig`` ran it, by + recording it as guardrail information on this call, the same entry the proxy-side Bedrock guardrail writes. + A cache hit carries the original response's trace but made no Bedrock call, so it bills nothing.""" + trace: Final = bedrock_guardrail_trace(result) + if trace is None or served_from_cache(result) or self.model_call_details.get("cache_hit") is True: + return + optional_params: Final = object_mapping(getattr(self, "optional_params", None)) or {} + guardrail_config: Final = object_mapping(optional_params.get("guardrailConfig")) or {} + identifier: Final = guardrail_config.get("guardrailIdentifier") + litellm_params_in: Final = object_mapping(getattr(self, "litellm_params", None)) or {} + region: Final = optional_params.get("aws_region_name") or litellm_params_in.get("aws_region_name") + entry: Final = native_bedrock_guardrail_information( + trace, + aws_region_name=region if isinstance(region, str) else None, + guardrail_identifier=identifier if isinstance(identifier, str) else None, + ) + litellm_params: Final[object] = self.model_call_details.get("litellm_params") + if entry is None or not is_str_keyed_dict(litellm_params): + return + metadata: Final[object] = litellm_params.get("metadata") + bucket: Final = {**metadata} if is_str_keyed_dict(metadata) else {} + litellm_params["metadata"] = bucket + with_guardrail_entry(bucket, entry) + def _merge_hidden_params_from_response_into_metadata(self, logging_result: object) -> None: """ Copy response._hidden_params into litellm_params.metadata['hidden_params']. @@ -2664,7 +2694,14 @@ class Logging(LiteLLMLoggingBaseClass): start_time = self.start_time if end_time is None: end_time = datetime.datetime.now() - if self.completion_start_time is None: + # A request re-dispatched to a fallback reuses this logging object but runs in + # its own window, so a first-token time from the previous attempt cannot belong + # to it. Keeping it makes the attempt report a negative TTFT. + first_token_seconds: Final = _timestamp_seconds(self.completion_start_time) + attempt_start_seconds: Final = _timestamp_seconds(start_time) + if first_token_seconds is None or ( + attempt_start_seconds is not None and first_token_seconds < attempt_start_seconds + ): self.completion_start_time = end_time self.model_call_details["completion_start_time"] = self.completion_start_time diff --git a/litellm/litellm_core_utils/llm_cost_calc/guardrail_cost.py b/litellm/litellm_core_utils/llm_cost_calc/guardrail_cost.py index 61ec2fc39c1..e4ad36f50bc 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/guardrail_cost.py +++ b/litellm/litellm_core_utils/llm_cost_calc/guardrail_cost.py @@ -1,15 +1,24 @@ import math -from collections.abc import Mapping +from collections.abc import Iterator, Mapping, Sequence +from functools import partial +from itertools import chain from typing import Annotated, Final -from pydantic import ConfigDict, Field, TypeAdapter, ValidationError +from pydantic import BaseModel, ConfigDict, Field, JsonValue, TypeAdapter, ValidationError +from typing_extensions import TypeIs # noqa: TID251 # TypeIs reaches typing only on 3.13 import litellm from litellm._logging import verbose_logger +from litellm.litellm_core_utils.core_helpers import ( + redact_nested_match_and_regex_keys, # pyright: ignore[reportUnknownVariableType] # signature uses bare dict +) +from litellm.types.guardrails import GuardrailEventHooks from litellm.types.llms.base import LiteLLMBaseModel -from litellm.types.utils import CostBreakdown +from litellm.types.utils import CostBreakdown, GuardrailStatus, StandardLoggingGuardrailInformation BEDROCK_GUARDRAIL_PRICING_KEY: Final = "bedrock/guardrails" +BEDROCK_TRACE_FIELDS: Final = ("trace", "amazon-bedrock-trace") +_INTERVENING_ACTIONS: Final = frozenset({"BLOCKED", "ANONYMIZED"}) class GuardrailPricing(LiteLLMBaseModel): @@ -96,6 +105,161 @@ def bedrock_guardrail_cost(usage_units: Mapping[str, int], aws_region_name: str return guardrail_cost_total(bedrock_guardrail_cost_by_unit(usage_units, aws_region_name)) +_OBJECT_MAPPING: Final = TypeAdapter(Mapping[str, object]) +_OBJECT_ITEMS: Final = TypeAdapter(tuple[object, ...]) + + +def object_mapping(value: object) -> Mapping[str, object] | None: + """``value`` as a str-keyed mapping, None for anything else. Vendor JSON arrives untyped, so this is the one place + that turns it into something the rest of the module can read without guessing.""" + try: + return _OBJECT_MAPPING.validate_python(value) + except ValidationError: + return None + + +def object_items(value: object) -> tuple[object, ...]: + if isinstance(value, (str, bytes, Mapping)): + return () + try: + return _OBJECT_ITEMS.validate_python(value) + except ValidationError: + return () + + +def is_str_keyed_dict( + value: object, +) -> TypeIs[dict[str, object]]: # guard-ok: JSON object keys are str # mutable-ok: the caller edits this dict in place + return isinstance(value, dict) + + +def _trace_field(response: object, field: str) -> Mapping[str, object] | None: + source: Final = object_mapping(response) + trace: Final = object_mapping(source.get(field) if source is not None else getattr(response, field, None)) + if trace is None or object_mapping(trace.get("guardrail")) is None: + return None + return trace + + +def bedrock_guardrail_trace(response: object) -> Mapping[str, object] | None: + """The guardrail trace a Bedrock response carries: Converse puts it under ``trace``, Invoke under + ``amazon-bedrock-trace``. Returned as is, or None when the response has no guardrail section.""" + traces: Final = map(partial(_trace_field, response), BEDROCK_TRACE_FIELDS) + return next((trace for trace in traces if trace is not None), None) + + +class _Assessment(BaseModel): + model_config = ConfigDict(frozen=True, extra="allow") + invocation_metrics: Mapping[str, JsonValue] = Field(default_factory=dict, alias="invocationMetrics") + + @property + def billable_usage(self) -> Mapping[str, int]: + usage: Final = self.invocation_metrics.get("usage") + if not isinstance(usage, dict): + return {} + return { + counter: units + for counter, units in usage.items() + if counter.endswith("Units") and isinstance(units, int) and not isinstance(units, bool) + } + + +class _GuardrailTrace(BaseModel): + """Both trace shapes: Converse ``inputAssessment`` {id: a} and ``outputAssessments`` {id: [a, ...]}, Invoke + ``input`` {id: a} and ``outputs`` [{id: a}, ...].""" + + model_config = ConfigDict(frozen=True, extra="allow") + input_assessment: Mapping[str, _Assessment] = Field(default_factory=dict, alias="inputAssessment") + input: Mapping[str, _Assessment] = Field(default_factory=dict) + output_assessments: Mapping[str, tuple[_Assessment, ...]] = Field(default_factory=dict, alias="outputAssessments") + outputs: tuple[Mapping[str, _Assessment], ...] = () + + @property + def assessments(self) -> tuple[tuple[str, _Assessment], ...]: + listed: Final = ( + (gid, a) for gid, items in self.output_assessments.items() for a in items + ) # comprehension-ok: id then its list + return ( + *self.input_assessment.items(), + *self.input.items(), + *listed, + *chain.from_iterable(by_id.items() for by_id in self.outputs), + ) + + +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) + + +def _summed_usage(usages: Sequence[Mapping[str, int]]) -> Mapping[str, int]: + counters: Final = tuple(dict.fromkeys(chain.from_iterable(usages))) + return {counter: sum(usage.get(counter, 0) for usage in usages) for counter in counters} + + +def _actions(node: JsonValue) -> Iterator[str]: + if isinstance(node, list): + yield from chain.from_iterable(map(_actions, node)) + return + if not isinstance(node, dict): + return + action: Final = node.get("action") + if isinstance(action, str): + yield action + yield from chain.from_iterable(map(_actions, node.values())) + + +def native_bedrock_guardrail_information( + trace: Mapping[str, object], aws_region_name: str | None, guardrail_identifier: str | None +) -> StandardLoggingGuardrailInformation | None: + """The guardrail entry for a trace Bedrock returned because the deployment's ``guardrailConfig`` ran with + ``trace: enabled``: every assessment's ``invocationMetrics.usage`` summed and priced like the proxy-side Bedrock + guardrail. None when the trace carries no usage, which is what ``trace: disabled`` leaves.""" + try: + guardrail: Final = _JSON_OBJECT.validate_python(trace.get("guardrail")) + pairs: Final = _GuardrailTrace.model_validate(guardrail).assessments + except ValidationError: + return None + usage: Final = _summed_usage(tuple(assessment.billable_usage for _, assessment in pairs)) + if not usage: + return None + cost_by_unit: Final = bedrock_guardrail_cost_by_unit(usage, aws_region_name) + name: Final = guardrail_identifier or pairs[0][0] + status: Final[GuardrailStatus] = ( + "guardrail_intervened" if any(action in _INTERVENING_ACTIONS for action in _actions(guardrail)) else "success" + ) + priced: Final[StandardLoggingGuardrailInformation] = ( + {"guardrail_cost_by_unit": cost_by_unit} if cost_by_unit is not None else {} + ) + return { + "guardrail_name": name, + "guardrail_id": name, + "guardrail_provider": "bedrock", + "guardrail_mode": GuardrailEventHooks.during_call, + "guardrail_response": redact_nested_match_and_regex_keys(guardrail), + "guardrail_status": status, + "guardrail_usage": usage, + "guardrail_cost": guardrail_cost_total(cost_by_unit), + **priced, + } + + +def guardrail_entries(raw: object) -> tuple[object, ...]: + """The entries under ``standard_logging_guardrail_information``, which older writers stored as one dict.""" + single: Final = object_mapping(raw) + return (single,) if single is not None else object_items(raw) + + +def with_guardrail_entry( + metadata: dict[str, object], # mutable-ok: the proxy's own request metadata, read later for headers and spend + entry: StandardLoggingGuardrailInformation, +) -> None: + """Append ``entry`` to the metadata's ``standard_logging_guardrail_information`` once, so the same trace recorded + from the response headers path and again from the success handler bills a single time.""" + existing: Final = guardrail_entries(metadata.get("standard_logging_guardrail_information")) + if entry in existing: + return + metadata["standard_logging_guardrail_information"] = [*existing, entry] # rebind-ok: appending here is the job + + AZURE_PROMPT_SHIELD_TEXT_RECORD_UNIT: Final = "text_records" diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py index 6664d54a744..2aa6142b3c5 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -1088,21 +1088,55 @@ def parse_completion_tokens_details(usage: Usage) -> CompletionTokensDetailsResu ) +def get_modality_token_cost_key( + model_info: ModelInfo, + cost_key: str, + prompt_tokens: int, + *, + threshold_is_inclusive: bool = False, +) -> str: + suffix: Final = next((suffix for suffix in _NON_STANDARD_THRESHOLD_SUFFIXES if cost_key.endswith(suffix)), "") + base_key: Final = cost_key.removesuffix(suffix) + threshold_key: Final = max( + ( + key + for key in model_info + if key.startswith(f"{base_key}_above_") + and key.endswith(("_tokens", f"_tokens{suffix}")) + and model_info.get(key) is not None + and _prompt_exceeds_threshold(prompt_tokens, _parse_above_token_threshold(key), threshold_is_inclusive) + ), + key=_parse_above_token_threshold, + default=None, + ) + if threshold_key is None: + return cost_key + return f"{threshold_key.removesuffix(suffix)}{suffix}" + + def _modality_token_cost( model_info: ModelInfo, modality_cost_key: str, tokens: int, prompt_base_cost: float, service_tier: str | None, + *, + prompt_tokens: int = 0, + threshold_is_inclusive: bool = False, ) -> 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: + cost_key: Final = get_modality_token_cost_key( + model_info, + get_service_tier_cost_key(modality_cost_key, service_tier), + prompt_tokens, + threshold_is_inclusive=threshold_is_inclusive, + ) + standard_key: Final = cost_key.removesuffix(get_service_tier_cost_key("", service_tier)) + if model_info.get(cost_key) is None and model_info.get(standard_key) is None: return float(tokens) * prompt_base_cost return calculate_cost_component(model_info, cost_key, tokens) @@ -1115,6 +1149,9 @@ def _calculate_input_cost( cache_creation_cost: float, cache_creation_cost_above_1hr: float, service_tier: str | None = None, + *, + prompt_tokens: int = 0, + threshold_is_inclusive: bool = False, ) -> float: """ Calculates the input cost for a given model, prompt tokens, and completion tokens. @@ -1149,6 +1186,8 @@ def _calculate_input_cost( prompt_tokens_details["audio_tokens"], prompt_base_cost, service_tier, + prompt_tokens=prompt_tokens, + threshold_is_inclusive=threshold_is_inclusive, ) ### IMAGE TOKEN COST @@ -1467,6 +1506,8 @@ def generic_cost_per_token( cache_creation_cost=cache_creation_cost, cache_creation_cost_above_1hr=cache_creation_cost_above_1hr, service_tier=service_tier, + prompt_tokens=usage.prompt_tokens, + threshold_is_inclusive=_uses_inclusive_token_thresholds(custom_llm_provider), ) ## CALCULATE OUTPUT COST @@ -1524,11 +1565,9 @@ def generic_cost_per_token( ## IMAGE COST if not is_text_tokens_total and image_tokens and image_tokens > 0: - _output_cost_per_image_token = get_cost_per_unit(resolved_model_info, "output_cost_per_image_token", None) - _output_cost_per_image_token = ( - _output_cost_per_image_token if _output_cost_per_image_token is not None else completion_base_cost - ) - completion_cost += float(image_tokens) * _output_cost_per_image_token + image_cost_key: Final = get_service_tier_cost_key("output_cost_per_image_token", service_tier) + image_rate: Final = get_cost_per_unit(resolved_model_info, image_cost_key, None) + completion_cost += float(image_tokens) * (completion_base_cost if image_rate is None else image_rate) ## VIDEO COST if not is_text_tokens_total and video_tokens and video_tokens > 0: diff --git a/litellm/litellm_core_utils/llm_cost_calc/zero_cost_diagnostic.py b/litellm/litellm_core_utils/llm_cost_calc/zero_cost_diagnostic.py index 9688b511ea8..3ed598b5e97 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/zero_cost_diagnostic.py +++ b/litellm/litellm_core_utils/llm_cost_calc/zero_cost_diagnostic.py @@ -5,6 +5,7 @@ from typing import Final from pydantic import TypeAdapter, ValidationError from typing_extensions import assert_never +from litellm.litellm_core_utils.llm_cost_calc.utils import parse_prompt_tokens_details from litellm.types.utils import ( CompletionTokensDetailsWrapper, PromptTokensDetailsWrapper, @@ -16,6 +17,8 @@ ZERO_COST_COUNTER_NAME: Final = "litellm_zero_cost_requests_total" _TEXT_INPUT_RATE: Final = "input_cost_per_token" _AUDIO_INPUT_RATE: Final = "input_cost_per_audio_token" +_CACHE_READ_INPUT_RATE: Final = "cache_read_input_token_cost" +_CACHE_READ_AUDIO_INPUT_RATE: Final = "cache_read_input_audio_token_cost" _TEXT_OUTPUT_RATE: Final = "output_cost_per_token" _AUDIO_OUTPUT_RATE: Final = "output_cost_per_audio_token" _RATE_KEY_MARKERS: Final = ("cost", "pricing") @@ -33,19 +36,38 @@ def _tokens(value: object) -> int: def used_pricing_keys(usage: Usage) -> tuple[str, ...]: - prompt_audio: Final = _audio_tokens(usage.prompt_tokens_details) + prompt_details: Final = parse_prompt_tokens_details(usage) + prompt_audio: Final = prompt_details["audio_tokens"] completion_audio: Final = _audio_tokens(usage.completion_tokens_details) - prompt_text: Final = _tokens(usage.prompt_tokens) - prompt_audio + prompt_text: Final = max( + _tokens(usage.prompt_tokens) + - prompt_details["cache_hit_tokens"] + - prompt_audio + - prompt_details["image_tokens"] + - prompt_details["video_tokens"], + 0, + ) + cached_text: Final = max(prompt_details["cache_hit_tokens"] - prompt_details["cache_hit_audio_tokens"], 0) completion_text: Final = _tokens(usage.completion_tokens) - completion_audio components: Final = ( (_TEXT_INPUT_RATE, prompt_text), (_AUDIO_INPUT_RATE, prompt_audio), + (_CACHE_READ_INPUT_RATE, cached_text), + (_CACHE_READ_AUDIO_INPUT_RATE, prompt_details["cache_hit_audio_tokens"]), (_TEXT_OUTPUT_RATE, completion_text), (_AUDIO_OUTPUT_RATE, completion_audio), ) return tuple(key for key, count in components if count > 0) +def _effective_pricing_key(key: str, pricing_entry: Mapping[str, object]) -> str: + if key == _CACHE_READ_AUDIO_INPUT_RATE and pricing_entry.get(key) is None: + return _CACHE_READ_INPUT_RATE if pricing_entry.get(_CACHE_READ_INPUT_RATE) is not None else _TEXT_INPUT_RATE + if key == _CACHE_READ_INPUT_RATE and pricing_entry.get(key) is None: + return _TEXT_INPUT_RATE + return key + + def _nested_pricing(value: object) -> Mapping[str, object] | tuple[object, ...] | None: try: return _NESTED_PRICING.validate_python(value) @@ -89,7 +111,8 @@ def diagnose_zero_cost( pricing_entry: Mapping[str, object], calculation_failed: bool, ) -> StandardLoggingZeroCostDiagnostic | None: - used_keys: Final = used_pricing_keys(usage) + effective_keys: Final = tuple(_effective_pricing_key(key, pricing_entry) for key in used_pricing_keys(usage)) + used_keys: Final = tuple(key for index, key in enumerate(effective_keys) if key not in effective_keys[:index]) if not used_keys: return None missing_keys: Final = tuple(key for key in used_keys if pricing_entry.get(key) is None) diff --git a/litellm/litellm_core_utils/redact_messages.py b/litellm/litellm_core_utils/redact_messages.py index c2cca5739cf..c8a0ffa956a 100644 --- a/litellm/litellm_core_utils/redact_messages.py +++ b/litellm/litellm_core_utils/redact_messages.py @@ -14,6 +14,8 @@ from collections.abc import Mapping from dataclasses import replace from typing import TYPE_CHECKING, Any, Final +from pydantic import JsonValue, TypeAdapter, ValidationError + import litellm from litellm.constants import REDACTED_BY_LITELLM from litellm.integrations.custom_logger import CustomLogger @@ -72,11 +74,55 @@ def _redacted_streaming_response_copy(streaming_response): return redacted_response +BEDROCK_GUARDRAIL_TRACE_FIELDS: Final = ("trace", "amazon-bedrock-trace") +_MATCHED_TEXT_KEYS: Final = frozenset({"match", "regex"}) +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) + + +def _without_matched_text(trace: Mapping[str, JsonValue]) -> Mapping[str, JsonValue]: + redacted: Final = dict(copy.deepcopy(trace)) + stack: Final[list[JsonValue]] = [redacted] # mutable-ok: iterative walk, recursive_detector bans recursion here + while stack: + node = stack.pop() + if isinstance(node, dict): + for key in _MATCHED_TEXT_KEYS & node.keys(): + node[key] = "[REDACTED]" + stack.extend(node.values()) + elif isinstance(node, list): + stack.extend(node) + return redacted + + +def _redacted_bedrock_trace(trace: object) -> object: + try: + return _without_matched_text(_JSON_OBJECT.validate_python(trace)) + except ValidationError: + return trace + + +def _redact_bedrock_guardrail_trace(response: object) -> None: + """A Bedrock guardrail trace echoes the matched text of its sensitive-information and word findings back.""" + mapping: Final = _as_mapping(response) + if isinstance(response, dict): + traces: Final = { + field: _redacted_bedrock_trace(value) + for field, value in (mapping or {}).items() + if field in BEDROCK_GUARDRAIL_TRACE_FIELDS + } + response.update(traces) + return + for field, value in ((name, getattr(response, name, None)) for name in BEDROCK_GUARDRAIL_TRACE_FIELDS): + if value is not None: + setattr(response, field, _redacted_bedrock_trace(value)) + + def _redact_streaming_response(streaming_response): if hasattr(streaming_response, "choices"): for choice in streaming_response.choices: _redact_choice_content(choice) redact_vertex_ai_metadata_from_logged_object(streaming_response) + if isinstance(streaming_response, litellm.ModelResponse): + _redact_bedrock_guardrail_trace(streaming_response) elif hasattr(streaming_response, "output"): _redact_responses_api_output(streaming_response.output) if hasattr(streaming_response, "reasoning") and streaming_response.reasoning is not None: @@ -234,6 +280,7 @@ def _redact_standard_logging_object(payload: Mapping[str, object]) -> dict[str, if isinstance(response.get("choices"), list): _redact_model_response_dict_choices(response["choices"], redacted_str) redact_vertex_ai_metadata_from_logged_object(response) + _redact_bedrock_guardrail_trace(response) elif isinstance(response, str): standard_logging_object["response"] = redacted_str else: @@ -393,11 +440,13 @@ def perform_redaction(model_call_details: dict, result, redact_streaming_respons for choice in _result.choices: _redact_choice_content(choice) redact_vertex_ai_metadata_from_logged_object(_result) + _redact_bedrock_guardrail_trace(_result) elif isinstance(_result, dict) and "choices" in _result: # Handle dict representation of ModelResponse (e.g., from model_dump()) if _result.get("choices") is not None: _redact_model_response_dict_choices(_result["choices"], REDACTED_BY_LITELLM) redact_vertex_ai_metadata_from_logged_object(_result) + _redact_bedrock_guardrail_trace(_result) elif isinstance(_result, dict) and "output" in _result: if isinstance(_result.get("output"), list): _redact_responses_api_output_dict(_result["output"], REDACTED_BY_LITELLM) diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index d505192184e..42418fda028 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -13,6 +13,7 @@ from typing_extensions import ReadOnly, Required import litellm from litellm._logging import verbose_logger from litellm.litellm_core_utils.hidden_params import HIDDEN_PARAMS_ATTR, set_hidden_params +from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import BEDROCK_TRACE_FIELDS, object_items, object_mapping from litellm.litellm_core_utils.prompt_templates.common_utils import ( get_content_from_model_response, # pyright: ignore[reportUnknownVariableType] # response param accepts an untyped dict ) @@ -1389,6 +1390,35 @@ def _simple_text_part(choices: Sequence[object]) -> str | None: return content if isinstance(content, str) else "" +_NO_TRACE_FIELDS: Final[Mapping[str, object]] = MappingProxyType({}) + + +def _carry_bedrock_guardrail_trace(chunks: Sequence[object], response: ModelResponse) -> None: + """Bedrock reports a guardrail trace once per stream, on the final metadata chunk, so the built response carries it + under the same field as the non-streaming response (``trace`` on Converse, ``amazon-bedrock-trace`` on Invoke).""" + traced: Final = map(_bedrock_trace_fields, reversed(chunks)) + carried: Final = next((fields for fields in traced if fields), _NO_TRACE_FIELDS) + for field, trace in carried.items(): + setattr(response, field, trace) + + +def _record_bedrock_guardrail_trace( + chunks: Sequence[object], response: ModelResponse, logging_obj: Optional["Logging"] +) -> None: + _carry_bedrock_guardrail_trace(chunks, response) + if logging_obj is not None: + logging_obj.record_native_bedrock_guardrail(response) + + +def _bedrock_trace_fields(chunk: object) -> Mapping[str, object]: + as_mapping: Final = object_mapping(chunk) + raw: Final = ( + as_mapping.get("provider_specific_fields") if as_mapping else getattr(chunk, "provider_specific_fields", None) + ) + fields: Final[Mapping[str, object]] = object_mapping(raw) or {} + return {field: fields[field] for field in BEDROCK_TRACE_FIELDS if field in fields} + + def stream_chunk_builder( chunks: list, messages: Sequence | None = None, @@ -1465,6 +1495,11 @@ def stream_chunk_builder( ) break + _record_bedrock_guardrail_trace( + object_items(chunks), # pyright: ignore[reportUnknownArgumentType] # chunks is an untyped list param + response, + logging_obj, + ) _stamp_streaming_usage_cost(usage, response, logging_obj) _set_stream_builder_response_cost(response, logging_obj) @@ -1652,6 +1687,11 @@ def stream_chunk_builder( _stamp_streaming_usage_cost(usage, response, logging_obj) _set_stream_builder_response_cost(response, logging_obj) + _record_bedrock_guardrail_trace( + object_items(chunks), # pyright: ignore[reportUnknownArgumentType] # chunks is an untyped list param + response, + logging_obj, + ) processor.apply_provider_assembled_streaming_metadata(response, chunks, logging_obj) return response except Exception as e: diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 46ad1b70abe..7135b9c1071 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -1431,11 +1431,7 @@ class CustomStreamWrapper: setattr( model_response, "usage", - litellm.Usage( - prompt_tokens=response_obj["usage"].get("prompt_tokens", None) or None, - completion_tokens=response_obj["usage"].get("completion_tokens", None) or None, - total_tokens=response_obj["usage"].get("total_tokens", None) or None, - ), + _usage_from_payload(response_obj["usage"]), ) elif isinstance(response_obj["usage"], Usage): setattr( @@ -1452,6 +1448,8 @@ class CustomStreamWrapper: return _ProviderChunkParsed(response_obj) def chunk_creator(self, chunk: Any): + if isinstance(chunk, ModelResponseStream) and chunk.provider_specific_fields: + self.logging_obj.record_native_bedrock_guardrail(chunk.provider_specific_fields) if self.custom_llm_provider != "cached_response" and isinstance( chunk, (ChatCompletionChunk, ModelResponseStream) ): @@ -1572,6 +1570,8 @@ class CustomStreamWrapper: model_response.usage = Usage(**chunk_usage.model_dump()) ## RETURN ARG + if isinstance(chunk, ModelResponseStream) and chunk.provider_specific_fields: + self.copy_model_response_level_provider_specific_fields(chunk, model_response) result: Final = self.return_processed_chunk_logic( completion_obj=completion_obj, model_response=model_response, @@ -2478,6 +2478,11 @@ class CustomStreamWrapper: _strip_sse_data_from_chunk = strip_sse_data_from_chunk +def _usage_from_payload(payload: Mapping[str, Any]) -> Usage: + """Build ``Usage`` from a wire-shaped usage dict, forwarding every key.""" + return litellm.Usage(**payload) + + def _cache_token_count(details: PromptTokensDetailsWrapper | None, keys: tuple[str, ...]) -> int: for key in keys: value = getattr(details, key, None) diff --git a/litellm/litellm_core_utils/token_counter.py b/litellm/litellm_core_utils/token_counter.py index 7ac3e9c80c5..0cf26b6075e 100644 --- a/litellm/litellm_core_utils/token_counter.py +++ b/litellm/litellm_core_utils/token_counter.py @@ -51,6 +51,7 @@ from litellm.types.llms.anthropic import ( AnthropicMessagesTextParam, AnthropicMessagesToolResultParam, AnthropicMessagesToolUseParam, + AnthropicSearchResultParam, ) from litellm.types.llms.openai import ( AllMessageValues, @@ -64,6 +65,10 @@ from litellm.types.llms.openai import ( from litellm.types.utils import Message, SelectTokenizerResponse +class UnsupportedContentTypeError(ValueError): + pass + + def get_modified_max_tokens( model: str, base_model: str, @@ -783,6 +788,7 @@ def _validate_anthropic_content(content: Mapping[str, object]) -> type: mapping: Final = { "tool_use": AnthropicMessagesToolUseParam, "tool_result": AnthropicMessagesToolResultParam, + "search_result": AnthropicSearchResultParam, } expected_cls: Final = mapping.get(content_type) if isinstance(content_type, str) else None @@ -938,7 +944,9 @@ def _count_anthropic_content( except Exception as e: if default_token_count is not None: return default_token_count - raise ValueError(f"Error counting field '{field_name}': {e}") + raise (UnsupportedContentTypeError if isinstance(e, UnsupportedContentTypeError) else ValueError)( + f"Error counting field '{field_name}': {e}" + ) return tokens @@ -952,6 +960,7 @@ def _count_content_list( | AnthropicMessagesTextParam | AnthropicMessagesImageParam | AnthropicMessagesDocumentParam + | AnthropicSearchResultParam ], use_default_image_token_count: bool, default_token_count: int | None, @@ -986,7 +995,7 @@ def _count_content_list( count_function, use_default_image_token_count, ) - elif c["type"] in ("tool_use", "tool_result"): + elif c["type"] in ("tool_use", "tool_result", "search_result"): num_tokens += _count_anthropic_content( c, count_function, @@ -1012,17 +1021,18 @@ def _count_content_list( num_tokens += count_function(tool_name) else: content_type = c.get("type", type(c).__name__) if isinstance(c, dict) else type(c).__name__ - raise ValueError( + raise UnsupportedContentTypeError( f"Invalid content item type: {content_type}. " f"Expected str or dict with 'type' field " f"(text, image_url, image, document, file, tool_use, tool_result, thinking, redacted_thinking, " - f"tool_reference)." + f"tool_reference, search_result)." ) return num_tokens except Exception as e: if default_token_count is not None: return default_token_count - raise ValueError( + error_type: Final = UnsupportedContentTypeError if isinstance(e, UnsupportedContentTypeError) else ValueError + raise error_type( f"Error getting number of tokens from content list: {e}, default_token_count={default_token_count}" ) diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index 1f09b3724f1..6984fc1c1cd 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/litellm/llms/bedrock/chat/converse_transformation.py @@ -2813,6 +2813,8 @@ class AmazonConverseConfig(BaseConfig): # Add "trace" from Bedrock guardrails - if user has opted in to returning it if "trace" in completion_response: setattr(model_response, "trace", completion_response["trace"]) + if logging_obj is not None: + logging_obj.record_native_bedrock_guardrail(model_response) # Add service_tier if present in Bedrock response # Map Bedrock serviceTier (object) to OpenAI service_tier (string) diff --git a/litellm/llms/bedrock/chat/invoke_handler.py b/litellm/llms/bedrock/chat/invoke_handler.py index 6ad400661be..cab7b165075 100644 --- a/litellm/llms/bedrock/chat/invoke_handler.py +++ b/litellm/llms/bedrock/chat/invoke_handler.py @@ -13,6 +13,7 @@ from litellm.caching.caching import InMemoryCache from litellm.constants import RESPONSE_FORMAT_TOOL_NAME from litellm.litellm_core_utils.core_helpers import map_finish_reason from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import object_mapping from litellm.litellm_core_utils.llm_response_utils.get_headers import provider_request_id_headers from litellm.litellm_core_utils.optional_imports import ensure_optional_import from litellm.litellm_core_utils.prompt_templates.common_utils import native_redacted_thinking_data @@ -21,6 +22,7 @@ from litellm.llms.anthropic.chat.handler import ( ) from litellm.llms.anthropic.chat.transformation import JsonModeRequest from litellm.llms.base_llm.base_utils import ResponseFormatSchema +from litellm.llms.bedrock.common_utils import BEDROCK_INVOKE_TRACE_FIELD from litellm.llms.custom_httpx.http_handler import ( AsyncHTTPHandler, HTTPHandler, @@ -951,8 +953,19 @@ class AmazonAnthropicClaudeStreamDecoder(AWSEventStreamDecoder): response_format_schema=response_format_schema, ) - def _chunk_parser(self, chunk_data: dict) -> ModelResponseStream: - return self.anthropic_model_response_iterator.chunk_parser(chunk=chunk_data) + def _chunk_parser(self, chunk_data: dict[str, object]) -> ModelResponseStream: + parsed: Final = self.anthropic_model_response_iterator.chunk_parser(chunk=chunk_data) + trace: Final = object_mapping(chunk_data.get(BEDROCK_INVOKE_TRACE_FIELD)) + if trace is None: + return parsed + return parsed.model_copy( + update={ + "provider_specific_fields": { + **(parsed.provider_specific_fields or {}), + BEDROCK_INVOKE_TRACE_FIELD: trace, + } + } + ) class AmazonDeepSeekR1StreamDecoder(AWSEventStreamDecoder): diff --git a/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py index 04747405375..1638f0c050c 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py @@ -5,6 +5,7 @@ import httpx from pydantic import TypeAdapter from litellm.anthropic_beta_headers_manager import filter_and_transform_beta_headers +from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import object_mapping from litellm.litellm_core_utils.prompt_templates.factory import ( convert_to_anthropic_image_obj, ) @@ -22,6 +23,7 @@ from litellm.llms.bedrock.chat.invoke_transformations.base_invoke_transformation AmazonInvokeConfig, ) from litellm.llms.bedrock.common_utils import ( + BEDROCK_INVOKE_TRACE_FIELD, apply_bedrock_invoke_structured_output, bedrock_supports_tool_search, get_anthropic_beta_from_headers, @@ -377,7 +379,7 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig): api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: - return AnthropicConfig.transform_response( + response: Final = AnthropicConfig.transform_response( self, model=model, raw_response=raw_response, @@ -391,3 +393,9 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig): api_key=api_key, json_mode=json_mode, ) + body: Final = object_mapping(raw_response.json()) or {} + trace: Final = object_mapping(body.get(BEDROCK_INVOKE_TRACE_FIELD)) + if trace is not None: + setattr(response, BEDROCK_INVOKE_TRACE_FIELD, trace) + logging_obj.record_native_bedrock_guardrail(response) + return response diff --git a/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py index 1b3b0c344df..2bd51702a81 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py @@ -60,7 +60,7 @@ def _response_header(headers: Mapping[str, str], name: str) -> str | None: return headers.get(name, None) -def _bedrock_invoke_guardrail_headers(raw_guardrail_config: object) -> "dict[str, str]": +def bedrock_invoke_guardrail_headers(raw_guardrail_config: object) -> "dict[str, str]": try: guardrail_config: Final = _GUARDRAIL_CONFIG_VALIDATOR.validate_python(raw_guardrail_config) except ValidationError as e: @@ -440,7 +440,7 @@ class AmazonInvokeConfig(BaseConfig, BaseAWSLLM): guardrail_headers: Final = ( () if raw_guardrail_config is None - else tuple(_bedrock_invoke_guardrail_headers(raw_guardrail_config).items()) + else tuple(bedrock_invoke_guardrail_headers(raw_guardrail_config).items()) ) owned_names, metadata_headers = bedrock_request_metadata_headers(litellm_params) return merge_bedrock_invoke_headers(headers, guardrail_headers, metadata_headers, owned_names) diff --git a/litellm/llms/bedrock/common_utils.py b/litellm/llms/bedrock/common_utils.py index 5862b05a619..999dd46b01a 100644 --- a/litellm/llms/bedrock/common_utils.py +++ b/litellm/llms/bedrock/common_utils.py @@ -42,6 +42,8 @@ if TYPE_CHECKING: from litellm.llms.custom_httpx.http_handler import HTTPHandler from litellm.types.llms.openai import AllMessageValues +BEDROCK_INVOKE_TRACE_FIELD: Final = "amazon-bedrock-trace" + _ERROR_REQUEST_URL: Final = "https://docs.litellm.ai/docs" _OPENAI_FAMILY_MODEL_RE: Final = re.compile(r"(^|[./])openai\.") diff --git a/litellm/llms/bedrock/cost_calculation.py b/litellm/llms/bedrock/cost_calculation.py index 0a2e39ab972..1b45531dda4 100644 --- a/litellm/llms/bedrock/cost_calculation.py +++ b/litellm/llms/bedrock/cost_calculation.py @@ -3,9 +3,13 @@ Helper util for handling bedrock-specific cost calculation - e.g.: prompt caching """ -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Final -from litellm.litellm_core_utils.llm_cost_calc.utils import generic_cost_per_token +import litellm +from litellm.litellm_core_utils.llm_cost_calc.utils import ( + generic_cost_per_token, + get_provider_specific_geo_multiplier, +) if TYPE_CHECKING: from litellm.types.utils import Usage @@ -17,9 +21,13 @@ def cost_per_token(model: str, usage: "Usage", service_tier: str | None = None) Follows the same logic as Anthropic's cost per token calculation. """ - return generic_cost_per_token( + model_info: Final = litellm.get_model_info(model=model, custom_llm_provider="bedrock") + prompt_cost, completion_cost = generic_cost_per_token( model=model, usage=usage, custom_llm_provider="bedrock", service_tier=service_tier, + model_info=model_info, ) + geo_multiplier: Final = get_provider_specific_geo_multiplier(model_info=model_info, usage=usage) + return prompt_cost * geo_multiplier, completion_cost * geo_multiplier diff --git a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py index 88ae0d4d514..00ef2673da1 100644 --- a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py +++ b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py @@ -32,6 +32,7 @@ from litellm.llms.base_llm.anthropic_messages.transformation import ( from litellm.llms.bedrock.chat.invoke_handler import AWSEventStreamDecoder from litellm.llms.bedrock.chat.invoke_transformations.base_invoke_transformation import ( AmazonInvokeConfig, + bedrock_invoke_guardrail_headers, ) from litellm.llms.bedrock.common_utils import ( BedrockError, @@ -114,8 +115,14 @@ class AmazonAnthropicClaudeMessagesConfig( api_key: str | None = None, api_base: str | None = None, ) -> tuple[dict, str | None]: + raw_guardrail_config: Final = litellm_params.get("guardrailConfig") # pyright: ignore[reportUnknownVariableType, reportUnknownMemberType] # litellm_params is an untyped dict param + guardrail_headers: Final = ( + () + if raw_guardrail_config is None + else tuple(bedrock_invoke_guardrail_headers(raw_guardrail_config).items()) # pyright: ignore[reportUnknownArgumentType] # read from the untyped dict param + ) owned_names, metadata_headers = bedrock_request_metadata_headers(litellm_params) - return merge_bedrock_invoke_headers(headers, (), metadata_headers, owned_names), api_base + return merge_bedrock_invoke_headers(headers, guardrail_headers, metadata_headers, owned_names), api_base # pyright: ignore[reportUnknownArgumentType] # headers is an untyped dict param def sign_request( self, diff --git a/litellm/llms/custom_httpx/aiohttp_transport.py b/litellm/llms/custom_httpx/aiohttp_transport.py index b8cb183c02f..9297e76624d 100644 --- a/litellm/llms/custom_httpx/aiohttp_transport.py +++ b/litellm/llms/custom_httpx/aiohttp_transport.py @@ -6,6 +6,7 @@ import ssl import sys import typing import urllib.request +import weakref from collections.abc import Callable, Generator from typing import ClassVar, Final @@ -109,6 +110,24 @@ def map_aiohttp_exceptions() -> Generator[None, None, None]: raise mapped_exc(message) from exc +def _release_transport_anchor(_transport: object) -> None: + """Release a transport anchored to a response stream. See ``_anchor_transport_to``.""" + + +def _anchor_transport_to(stream: "AiohttpResponseStream", transport: object) -> None: + """Keep the transport alive for as long as the stream can still read. + + The transport finalizer closes the session it owns, and nothing on the path from a + live response back to that session references the transport, so a client dropped + mid-stream would close the session under the body. ``weakref.finalize`` holds the + transport in its own registry rather than on the stream, which keeps the transport + out of the response's reference cycle (an exception carrying the response, for one) + so it is finalized by refcount once the anchor drops and can still schedule an async + close, instead of being collected in the same cyclic pass as the session it closes. + """ + weakref.finalize(stream, _release_transport_anchor, transport) + + class AiohttpResponseStream(httpx.AsyncByteStream): CHUNK_SIZE = 1024 * 16 @@ -289,6 +308,13 @@ class LiteLLMAiohttpTransport(AiohttpTransport): cls._background_close_tasks.add(task) task.add_done_callback(cls._on_close_task_done) + def __del__(self) -> None: + session: Final = getattr(self, "client", None) + if not getattr(self, "_owns_session", False) or not isinstance(session, ClientSession): + return + with contextlib.suppress(Exception): + self._close_recycled_session(session) + def get_valid_client_session(self) -> ClientSession: """ Helper to get a valid ClientSession for the current event loop. @@ -456,10 +482,12 @@ class LiteLLMAiohttpTransport(AiohttpTransport): # Re-raise if it's a different RuntimeError raise + stream: Final = AiohttpResponseStream(response) + _anchor_transport_to(stream, self) return httpx.Response( status_code=response.status, headers=response.headers, - stream=AiohttpResponseStream(response), + stream=stream, request=request, ) diff --git a/litellm/llms/custom_httpx/http_handler.py b/litellm/llms/custom_httpx/http_handler.py index cffeb44f12c..e7be1c9ce4d 100644 --- a/litellm/llms/custom_httpx/http_handler.py +++ b/litellm/llms/custom_httpx/http_handler.py @@ -1123,6 +1123,17 @@ class AsyncHTTPHandler: # Mirrors LiteLLMAiohttpTransport._background_close_tasks. _finalizer_close_tasks: ClassVar[set["asyncio.Task[None]"]] = set() # mutable-ok: strong refs for pending closes + @classmethod + async def _close_client_at_finalization(cls, client: httpx.AsyncClient) -> None: + try: + await client.aclose() + except (OSError, RuntimeError) as exc: + verbose_logger.debug("Error closing client at finalization: %s", exc) + finally: + task: Final = asyncio.current_task() + if task is not None: + cls._finalizer_close_tasks.discard(task) + @classmethod def _on_finalizer_close_done(cls, task: "asyncio.Task[None]") -> None: cls._finalizer_close_tasks.discard(task) @@ -1190,8 +1201,8 @@ class AsyncHTTPHandler: # here is the cross-loop path the transport refuses. self._dispose_wrapped_aiohttp_session() return - task: Final = loop.create_task(self._client.aclose()) cls: Final = type(self) + task: Final = loop.create_task(cls._close_client_at_finalization(self._client)) cls._finalizer_close_tasks.add(task) task.add_done_callback(cls._on_finalizer_close_done) except Exception: @@ -1823,7 +1834,7 @@ def get_httpx_client(params: dict | None = None) -> HTTPHandler: cache = LLMClientCache() setattr(litellm, "in_memory_llm_clients_cache", cache) - _cached_client: Final = cache.get_cache(_cache_key_name) + _cached_client: Final = cache.get_cache(_cache_key_name, loop_bound=False) if _cached_client: return _cached_client @@ -1839,6 +1850,7 @@ def get_httpx_client(params: dict | None = None) -> HTTPHandler: value=_new_client, ttl=_DEFAULT_TTL_FOR_HTTPX_CLIENTS, litellm_owned_client=True, + loop_bound=False, ) return _new_client diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 145e44f919e..3eff8bc3b18 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -2430,6 +2430,7 @@ class BaseLLMHTTPHandler: raw_response=response, logging_obj=logging_obj, ) + logging_obj.record_native_bedrock_guardrail(initial_response) return await self._finalize_anthropic_messages_response( initial_response=initial_response, @@ -7582,7 +7583,7 @@ class BaseLLMHTTPHandler: """ if client is None or not isinstance(client, AsyncHTTPHandler): async_httpx_client = get_async_httpx_client( - llm_provider=litellm.LlmProviders(custom_llm_provider), + llm_provider=custom_llm_provider, params={"ssl_verify": litellm_params.get("ssl_verify", None)}, ) else: @@ -7777,7 +7778,7 @@ class BaseLLMHTTPHandler: """ if client is None or not isinstance(client, AsyncHTTPHandler): async_httpx_client = get_async_httpx_client( - llm_provider=litellm.LlmProviders(custom_llm_provider), + llm_provider=custom_llm_provider, params={"ssl_verify": litellm_params.get("ssl_verify", None)}, ) else: @@ -7957,7 +7958,7 @@ class BaseLLMHTTPHandler: """ if client is None or not isinstance(client, AsyncHTTPHandler): async_httpx_client = get_async_httpx_client( - llm_provider=litellm.LlmProviders(custom_llm_provider), + llm_provider=custom_llm_provider, params={"ssl_verify": litellm_params.get("ssl_verify", None)}, ) else: @@ -8120,7 +8121,7 @@ class BaseLLMHTTPHandler: ): if client is None or not isinstance(client, AsyncHTTPHandler): async_httpx_client = get_async_httpx_client( - llm_provider=litellm.LlmProviders(custom_llm_provider), + llm_provider=custom_llm_provider, params={"ssl_verify": litellm_params.get("ssl_verify", None)}, ) else: @@ -8260,7 +8261,7 @@ class BaseLLMHTTPHandler: ): if client is None or not isinstance(client, AsyncHTTPHandler): async_httpx_client = get_async_httpx_client( - llm_provider=litellm.LlmProviders(custom_llm_provider), + llm_provider=custom_llm_provider, params={"ssl_verify": litellm_params.get("ssl_verify", None)}, ) else: @@ -8432,7 +8433,7 @@ class BaseLLMHTTPHandler: ): if client is None or not isinstance(client, AsyncHTTPHandler): async_httpx_client = get_async_httpx_client( - llm_provider=litellm.LlmProviders(custom_llm_provider), + llm_provider=custom_llm_provider, params={"ssl_verify": litellm_params.get("ssl_verify", None)}, ) else: @@ -8618,7 +8619,7 @@ class BaseLLMHTTPHandler: ): if client is None or not isinstance(client, AsyncHTTPHandler): async_httpx_client = get_async_httpx_client( - llm_provider=litellm.LlmProviders(custom_llm_provider), + llm_provider=custom_llm_provider, params={"ssl_verify": litellm_params.get("ssl_verify", None)}, ) else: @@ -8750,7 +8751,7 @@ class BaseLLMHTTPHandler: """ if client is None or not isinstance(client, AsyncHTTPHandler): async_httpx_client = get_async_httpx_client( - llm_provider=litellm.LlmProviders(custom_llm_provider), + llm_provider=custom_llm_provider, params={"ssl_verify": litellm_params.get("ssl_verify", None)}, ) else: @@ -8831,7 +8832,7 @@ class BaseLLMHTTPHandler: """ if client is None or not isinstance(client, AsyncHTTPHandler): async_httpx_client = get_async_httpx_client( - llm_provider=litellm.LlmProviders(custom_llm_provider), + llm_provider=custom_llm_provider, params={"ssl_verify": litellm_params.get("ssl_verify", None)}, ) else: @@ -9015,7 +9016,7 @@ class BaseLLMHTTPHandler: """ if client is None or not isinstance(client, AsyncHTTPHandler): async_httpx_client = get_async_httpx_client( - llm_provider=litellm.LlmProviders(custom_llm_provider), + llm_provider=custom_llm_provider, params={"ssl_verify": litellm_params.get("ssl_verify", None)}, ) else: diff --git a/litellm/llms/openai/common_utils.py b/litellm/llms/openai/common_utils.py index 4ab91392ad4..e8c4ffb54cc 100644 --- a/litellm/llms/openai/common_utils.py +++ b/litellm/llms/openai/common_utils.py @@ -253,6 +253,10 @@ def build_output_token_limit_response( return e.response.headers, (_aiter_once(chunk) if is_async else _iter_once(chunk)) +def _is_loop_bound(client_initialization_params: Mapping[str, object]) -> bool: + return client_initialization_params.get("is_async") is not False + + class BaseOpenAILLM: """ Base class for OpenAI LLMs for getting their httpx clients and SSL verification settings @@ -267,7 +271,9 @@ class BaseOpenAILLM: client_initialization_params=client_initialization_params, client_type=client_type, ) - _cached_client: Final = litellm.in_memory_llm_clients_cache.get_cache(_cache_key) + _cached_client: Final = litellm.in_memory_llm_clients_cache.get_cache( + _cache_key, loop_bound=_is_loop_bound(client_initialization_params) + ) return _cached_client @staticmethod @@ -308,6 +314,7 @@ class BaseOpenAILLM: value=openai_client, ttl=_DEFAULT_TTL_FOR_HTTPX_CLIENTS, litellm_owned_client=litellm_owned_client, + loop_bound=_is_loop_bound(client_initialization_params), ) @staticmethod diff --git a/litellm/llms/openai/openai.py b/litellm/llms/openai/openai.py index 27346ee936e..bf33070faa6 100644 --- a/litellm/llms/openai/openai.py +++ b/litellm/llms/openai/openai.py @@ -11,13 +11,14 @@ if TYPE_CHECKING: from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer import openai -from openai import AsyncOpenAI, OpenAI +from openai import AsyncOpenAI, AsyncStream, OpenAI, Stream from openai._base_client import make_request_options from openai._constants import RAW_RESPONSE_HEADER from openai._legacy_response import LegacyAPIResponse from openai._types import RequestOptions from openai.types import CreateEmbeddingResponse from openai.types.beta.assistant_deleted import AssistantDeleted +from openai.types.chat import ChatCompletion, ChatCompletionChunk from openai.types.file_deleted import FileDeleted from pydantic import BaseModel, TypeAdapter from typing_extensions import overload @@ -352,7 +353,14 @@ _NO_EXTRA_HEADERS: Final[Mapping[str, str]] = types.MappingProxyType({}) _SDK_OPTION_KEYS: Final = frozenset(("extra_headers", "extra_query", "extra_body")) -def _embedding_request_without_sdk_defaults( +def _raw_chat_response( + result: ChatCompletion | Stream[ChatCompletionChunk] | AsyncStream[ChatCompletionChunk], +) -> LegacyAPIResponse[ChatCompletion]: + assert isinstance(result, LegacyAPIResponse) + return result + + +def _request_without_sdk_transform( data: Mapping[str, object], timeout: float | httpx.Timeout ) -> tuple[Mapping[str, object], RequestOptions]: body: Final = {k: v for k, v in data.items() if k not in _SDK_OPTION_KEYS} @@ -515,7 +523,17 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): """ start_time: Final = time.time() try: - raw_response = await openai_aclient.chat.completions.with_raw_response.create(**data, timeout=timeout) + body, options = _request_without_sdk_transform(data, timeout) + raw_response = _raw_chat_response( + await openai_aclient.post( + "/chat/completions", + body=body, + cast_to=ChatCompletion, + options=options, + stream=bool(data.get("stream")), + stream_cls=AsyncStream[ChatCompletionChunk], + ) + ) end_time = time.time() if hasattr(raw_response, "headers"): @@ -556,7 +574,17 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): """ raw_response = None try: - raw_response = openai_client.chat.completions.with_raw_response.create(**data, timeout=timeout) + body, options = _request_without_sdk_transform(data, timeout) + raw_response = _raw_chat_response( + openai_client.post( + "/chat/completions", + body=body, + cast_to=ChatCompletion, + options=options, + stream=bool(data.get("stream")), + stream_cls=Stream[ChatCompletionChunk], + ) + ) if hasattr(raw_response, "headers"): headers = dict(raw_response.headers) @@ -1255,7 +1283,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): logging_obj: LiteLLMLoggingObj, ) -> LegacyAPIResponse[CreateEmbeddingResponse]: if "encoding_format" not in data: - body, options = _embedding_request_without_sdk_defaults(data, timeout) + body, options = _request_without_sdk_transform(data, timeout) bypass_response: Final = await openai_aclient.post( "/embeddings", body=body, options=options, cast_to=CreateEmbeddingResponse ) @@ -1272,7 +1300,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): logging_obj: LiteLLMLoggingObj, ) -> LegacyAPIResponse[CreateEmbeddingResponse]: if "encoding_format" not in data: - body, options = _embedding_request_without_sdk_defaults(data, timeout) + body, options = _request_without_sdk_transform(data, timeout) bypass_response: Final = openai_client.post( "/embeddings", body=body, options=options, cast_to=CreateEmbeddingResponse ) diff --git a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py index 257fafbb6b8..5faf24ac92b 100644 --- a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py +++ b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py @@ -40,6 +40,7 @@ from litellm.llms.custom_httpx.http_handler import ( get_httpx_client, ) from litellm.types.llms.anthropic import AnthropicThinkingParam +from litellm.types.llms.base import CachedTokensDetails from litellm.types.llms.gemini import BidiGenerateContentServerMessage from litellm.types.llms.openai import ( AllMessageValues, @@ -2004,26 +2005,13 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): elif modality == "DOCUMENT": cached_text_tokens = (cached_text_tokens or 0) + token_count - ## Calculate non-cached tokens by subtracting cached from total (per modality) - ## This is necessary because promptTokensDetails includes both cached and non-cached tokens - ## See: https://github.com/BerriAI/litellm/issues/18750 - if cached_text_tokens is not None and prompt_text_tokens is not None: - # Explicit caching: subtract cached tokens per modality from cacheTokensDetails - prompt_text_tokens = prompt_text_tokens - cached_text_tokens - elif ( + if ( cached_tokens is not None and prompt_text_tokens is not None and cached_text_tokens is None and "cacheTokensDetails" not in usage_metadata ): - # Implicit caching: only cachedContentTokenCount is provided (no cacheTokensDetails) - # Subtract from text tokens since implicit caching is primarily for text content - # See: https://github.com/BerriAI/litellm/issues/16341 prompt_text_tokens = prompt_text_tokens - cached_tokens - if cached_audio_tokens is not None and prompt_audio_tokens is not None: - prompt_audio_tokens = prompt_audio_tokens - cached_audio_tokens - if cached_image_tokens is not None and prompt_image_tokens is not None: - prompt_image_tokens = prompt_image_tokens - cached_image_tokens if cached_video_tokens is not None and prompt_video_tokens is not None: prompt_video_tokens = prompt_video_tokens - cached_video_tokens @@ -2044,9 +2032,24 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): billable_text_tokens: Final = ( prompt_text_tokens if prompt_text_tokens is None else prompt_text_tokens + billable_tool_use_prompt_tokens ) + cached_tokens_details_fields: Final = MappingProxyType( + { + key: token_count + for key, token_count in ( + ("text_tokens", cached_text_tokens), + ("audio_tokens", cached_audio_tokens), + ("image_tokens", cached_image_tokens), + ) + if token_count is not None + } + ) + cached_tokens_details: Final = ( + CachedTokensDetails.model_validate(cached_tokens_details_fields) if cached_tokens_details_fields else None + ) prompt_tokens_details = PromptTokensDetailsWrapper( cached_tokens=cached_tokens, + cached_tokens_details=cached_tokens_details, audio_tokens=prompt_audio_tokens, text_tokens=billable_text_tokens, image_tokens=prompt_image_tokens, diff --git a/litellm/main.py b/litellm/main.py index dcef613e314..83cc06c122f 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -232,7 +232,7 @@ from .llms.codestral.completion.handler import CodestralTextCompletion from .llms.cohere.embed import handler as cohere_embed from .llms.custom_httpx.aiohttp_handler import BaseLLMAIOHTTPHandler from .llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler -from .llms.custom_llm import CustomLLM, custom_chat_llm_router +from .llms.custom_llm import custom_chat_llm_router from .llms.databricks.embed.handler import DatabricksEmbeddingHandler from .llms.deprecated_providers import aleph_alpha from .llms.gdc.chat.transformation import GDCGeminiConfig @@ -4954,10 +4954,7 @@ def _complete_custom_providers( stream: Final = ctx.stream timeout: Final = ctx.timeout - custom_handler: CustomLLM | None = None - for item in litellm.custom_provider_map: - if item["provider"] == custom_llm_provider: - custom_handler = item["custom_handler"] + custom_handler: Final = ProviderConfigManager.get_custom_llm_handler(custom_llm_provider) if custom_handler is None: raise LiteLLMUnknownProvider(model=model, custom_llm_provider=custom_llm_provider) @@ -7366,10 +7363,7 @@ def embedding( litellm_params={}, ) elif custom_llm_provider in litellm._custom_providers: - custom_handler: CustomLLM | None = None - for item in litellm.custom_provider_map: - if item["provider"] == custom_llm_provider: - custom_handler = item["custom_handler"] + custom_handler: Final = ProviderConfigManager.get_custom_llm_handler(custom_llm_provider) if custom_handler is None: raise LiteLLMUnknownProvider(model=model, custom_llm_provider=custom_llm_provider) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 98fee42de3c..2c183527ce8 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -8899,7 +8899,11 @@ "output_cost_per_token_above_272k_tokens": 1.5e-05, "output_cost_per_token_above_272k_tokens_priority": 3e-05, "output_cost_per_token_priority": 2e-05, - "search_context_cost_per_query": {"search_context_size_high":0.01,"search_context_size_low":0.01,"search_context_size_medium":0.01}, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "source": "https://azure.microsoft.com/en-us/blog/gpt-6-astra-sol-and-luna-for-production-agents-in-microsoft-foundry/", "supported_endpoints": [ "/v1/chat/completions", @@ -8952,7 +8956,11 @@ "output_cost_per_token_above_272k_tokens": 1.5e-05, "output_cost_per_token_above_272k_tokens_priority": 3e-05, "output_cost_per_token_priority": 2e-05, - "search_context_cost_per_query": {"search_context_size_high":0.01,"search_context_size_low":0.01,"search_context_size_medium":0.01}, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "source": "https://azure.microsoft.com/en-us/blog/gpt-6-astra-sol-and-luna-for-production-agents-in-microsoft-foundry/", "supported_endpoints": [ "/v1/chat/completions", @@ -9644,7 +9652,11 @@ "output_cost_per_token_above_272k_tokens": 1.65e-05, "output_cost_per_token_above_272k_tokens_priority": 3.3e-05, "output_cost_per_token_priority": 2.2e-05, - "search_context_cost_per_query": {"search_context_size_high":0.01,"search_context_size_low":0.01,"search_context_size_medium":0.01}, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "source": "https://azure.microsoft.com/en-us/blog/gpt-6-astra-sol-and-luna-for-production-agents-in-microsoft-foundry/", "supported_endpoints": [ "/v1/chat/completions", @@ -16419,7 +16431,10 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_max_reasoning_effort": true, - "provider_specific_entry": {"us":1.1,"fast":2}, + "provider_specific_entry": { + "us": 1.1, + "fast": 2 + }, "supports_output_config": true, "supports_speed": true, "supports_fast_mode": true, @@ -28970,7 +28985,11 @@ "search_context_size_medium": 0.035, "search_context_size_high": 0.035 }, - "search_free_requests": {"count": 1500, "period": "day", "shared_group": "gemini-2-flash"}, + "search_free_requests": { + "count": 1500, + "period": "day", + "shared_group": "gemini-2-flash" + }, "google_maps_grounding_cost_per_query": 0.025, "input_cost_per_audio_token_batches": 5e-07, "input_cost_per_token_batches": 1.5e-07, @@ -28980,7 +28999,13 @@ "output_cost_per_token_batches": 1.25e-06, "output_cost_per_token_flex": 1.25e-06, "output_cost_per_token_priority": 4.5e-06, - "supports_image_size": false + "supports_image_size": false, + "input_cost_per_audio_token_above_200k_tokens_priority": 1.8e-06, + "input_cost_per_audio_token_above_200k_tokens_flex": 5e-07, + "cache_read_input_audio_token_cost_priority": 1.8e-07, + "input_cost_per_audio_token_flex": 5e-07, + "input_cost_per_audio_token_above_200k_tokens": 3e-07, + "input_cost_per_audio_token_above_200k_tokens_batches": 5e-07 }, "gemini-2.5-flash-image": { "deprecation_date": "2027-03-15", @@ -29085,8 +29110,13 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, - "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, - "web_search_billing_unit": "per_query" + "search_free_requests": { + "count": 5000, + "period": "month", + "shared_group": "gemini-3" + }, + "web_search_billing_unit": "per_query", + "output_cost_per_image_token_priority": 0.000216 }, "gemini-3-pro-image-preview": { "input_cost_per_image": 0.0011, @@ -29138,8 +29168,13 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, - "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, - "web_search_billing_unit": "per_query" + "search_free_requests": { + "count": 5000, + "period": "month", + "shared_group": "gemini-3" + }, + "web_search_billing_unit": "per_query", + "output_cost_per_image_token_priority": 0.000216 }, "gemini-3.1-flash-image": { "cache_read_input_token_cost": 5e-08, @@ -29189,8 +29224,13 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, - "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, - "web_search_billing_unit": "per_query" + "search_free_requests": { + "count": 5000, + "period": "month", + "shared_group": "gemini-3" + }, + "web_search_billing_unit": "per_query", + "output_cost_per_image_token_priority": 0.000108 }, "gemini-3.1-flash-image-preview": { "input_cost_per_image": 0.00056, @@ -29233,15 +29273,19 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, - "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, - "web_search_billing_unit": "per_query" + "search_free_requests": { + "count": 5000, + "period": "month", + "shared_group": "gemini-3" + }, + "web_search_billing_unit": "per_query", + "output_cost_per_image_token_priority": 0.000108 }, "gemini-3.1-flash-lite-image": { "cache_read_input_token_cost": 2.5e-08, "cache_read_input_token_cost_batches": 1.25e-08, "cache_read_input_token_cost_flex": 1.25e-08, "cache_read_input_token_cost_priority": 4.5e-08, - "deprecation_date": "2027-06-28", "input_cost_per_image": 0.00028, "input_cost_per_token": 2.5e-07, "input_cost_per_token_batches": 1.25e-07, @@ -29281,7 +29325,8 @@ "supports_response_schema": false, "supports_system_messages": true, "supports_video_input": true, - "supports_vision": true + "supports_vision": true, + "output_cost_per_image_token_priority": 5.4e-05 }, "gemini-3.1-flash-lite-preview": { "cache_read_input_token_cost": 2.5e-08, @@ -29329,7 +29374,11 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, - "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, + "search_free_requests": { + "count": 5000, + "period": "month", + "shared_group": "gemini-3" + }, "web_search_billing_unit": "per_query", "google_maps_grounding_cost_per_query": 0.014 }, @@ -29391,10 +29440,22 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, - "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, + "search_free_requests": { + "count": 5000, + "period": "month", + "shared_group": "gemini-3" + }, "web_search_billing_unit": "per_query", "google_maps_grounding_cost_per_query": 0.014, - "input_cost_per_audio_token_batches": 2.5e-07 + "input_cost_per_audio_token_batches": 2.5e-07, + "input_cost_per_audio_token_above_200k_tokens_priority": 9e-07, + "cache_read_input_audio_token_cost_flex": 2.5e-08, + "input_cost_per_audio_token_above_200k_tokens_flex": 2.5e-07, + "cache_read_input_audio_token_cost_priority": 9e-08, + "input_cost_per_audio_token_flex": 2.5e-07, + "input_cost_per_audio_token_above_200k_tokens": 5e-07, + "cache_read_input_audio_token_cost_batches": 2.5e-08, + "input_cost_per_audio_token_above_200k_tokens_batches": 2.5e-07 }, "gemini-3.5-flash-lite": { "deprecation_date": "2027-07-21", @@ -29452,7 +29513,11 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, - "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, + "search_free_requests": { + "count": 5000, + "period": "month", + "shared_group": "gemini-3" + }, "web_search_billing_unit": "per_query", "google_maps_grounding_cost_per_query": 0.014 }, @@ -29539,7 +29604,11 @@ "search_context_size_medium": 0.035, "search_context_size_high": 0.035 }, - "search_free_requests": {"count": 1500, "period": "day", "shared_group": "gemini-2-flash"}, + "search_free_requests": { + "count": 1500, + "period": "day", + "shared_group": "gemini-2-flash" + }, "google_maps_grounding_cost_per_query": 0.025, "input_cost_per_audio_token_batches": 5e-08, "input_cost_per_token_batches": 5e-08, @@ -29549,7 +29618,13 @@ "output_cost_per_token_batches": 2e-07, "output_cost_per_token_flex": 2e-07, "output_cost_per_token_priority": 7.2e-07, - "supports_image_size": false + "supports_image_size": false, + "input_cost_per_audio_token_above_200k_tokens_priority": 5.4e-07, + "input_cost_per_audio_token_above_200k_tokens_flex": 5e-08, + "cache_read_input_audio_token_cost_priority": 5.4e-08, + "input_cost_per_audio_token_flex": 5e-08, + "input_cost_per_audio_token_above_200k_tokens": 3e-07, + "input_cost_per_audio_token_above_200k_tokens_batches": 5e-08 }, "gemini-2.5-flash-lite-preview-09-2025": { "cache_read_input_token_cost": 1e-08, @@ -29594,7 +29669,11 @@ "search_context_size_medium": 0.035, "search_context_size_high": 0.035 }, - "search_free_requests": {"count": 1500, "period": "day", "shared_group": "gemini-2-flash"}, + "search_free_requests": { + "count": 1500, + "period": "day", + "shared_group": "gemini-2-flash" + }, "google_maps_grounding_cost_per_query": 0.025, "supports_image_size": false }, @@ -29641,7 +29720,11 @@ "search_context_size_medium": 0.035, "search_context_size_high": 0.035 }, - "search_free_requests": {"count": 1500, "period": "day", "shared_group": "gemini-2-flash"}, + "search_free_requests": { + "count": 1500, + "period": "day", + "shared_group": "gemini-2-flash" + }, "google_maps_grounding_cost_per_query": 0.025, "supports_image_size": false }, @@ -29829,7 +29912,11 @@ "search_context_size_medium": 0.035, "search_context_size_high": 0.035 }, - "search_free_requests": {"count": 10000, "period": "day", "shared_group": "gemini-2.5-pro"}, + "search_free_requests": { + "count": 10000, + "period": "day", + "shared_group": "gemini-2.5-pro" + }, "google_maps_grounding_cost_per_query": 0.025, "input_cost_per_token_above_200k_tokens_priority": 4.5e-06, "input_cost_per_token_batches": 6.25e-07, @@ -29838,7 +29925,9 @@ "output_cost_per_token_above_200k_tokens_priority": 2.7e-05, "output_cost_per_token_batches": 5e-06, "output_cost_per_token_flex": 5e-06, - "output_cost_per_token_priority": 1.8e-05 + "output_cost_per_token_priority": 1.8e-05, + "input_cost_per_token_above_200k_tokens_flex": 1.25e-06, + "output_cost_per_token_above_200k_tokens_flex": 7.5e-06 }, "gemini-3.1-pro-preview": { "cache_creation_input_token_cost_above_200k_tokens": 2.5e-07, @@ -29866,8 +29955,16 @@ "output_cost_per_token_flex": 6e-06, "output_cost_per_token_priority": 2.16e-05, "prompt_cache_min_tokens": 4096.0, - "search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014}, - "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, + "search_context_cost_per_query": { + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014, + "search_context_size_high": 0.014 + }, + "search_free_requests": { + "count": 5000, + "period": "month", + "shared_group": "gemini-3" + }, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", @@ -29896,7 +29993,9 @@ "supports_video_input": true, "supports_vision": true, "supports_web_search": true, - "web_search_billing_unit": "per_query" + "web_search_billing_unit": "per_query", + "input_cost_per_token_above_200k_tokens_flex": 2e-06, + "output_cost_per_token_above_200k_tokens_flex": 9e-06 }, "gemini-3.1-pro-preview-customtools": { "prompt_cache_min_tokens": 4096, @@ -29948,7 +30047,11 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, - "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, + "search_free_requests": { + "count": 5000, + "period": "month", + "shared_group": "gemini-3" + }, "web_search_billing_unit": "per_query", "google_maps_grounding_cost_per_query": 0.014 }, @@ -30005,7 +30108,11 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, - "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, + "search_free_requests": { + "count": 5000, + "period": "month", + "shared_group": "gemini-3" + }, "web_search_billing_unit": "per_query" }, "vertex_ai/gemini-3-flash-preview": { @@ -30029,8 +30136,16 @@ "output_cost_per_token_batches": 1.5e-06, "output_cost_per_token_flex": 1.5e-06, "output_cost_per_token_priority": 5.4e-06, - "search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014}, - "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, + "search_context_cost_per_query": { + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014, + "search_context_size_high": 0.014 + }, + "search_free_requests": { + "count": 5000, + "period": "month", + "shared_group": "gemini-3" + }, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", @@ -30058,7 +30173,13 @@ "supports_video_input": true, "supports_vision": true, "supports_web_search": true, - "web_search_billing_unit": "per_query" + "web_search_billing_unit": "per_query", + "input_cost_per_audio_token_above_200k_tokens_priority": 1.8e-06, + "input_cost_per_audio_token_above_200k_tokens_flex": 5e-07, + "cache_read_input_audio_token_cost_priority": 1.8e-07, + "input_cost_per_audio_token_flex": 5e-07, + "input_cost_per_audio_token_above_200k_tokens": 1e-06, + "input_cost_per_audio_token_above_200k_tokens_batches": 5e-07 }, "vertex_ai/gemini-3.5-flash": { "prompt_cache_min_tokens": 4096, @@ -30113,7 +30234,11 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, - "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, + "search_free_requests": { + "count": 5000, + "period": "month", + "shared_group": "gemini-3" + }, "web_search_billing_unit": "per_query", "google_maps_grounding_cost_per_query": 0.014, "input_cost_per_token_batches": 7.5e-07, @@ -30179,7 +30304,11 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, - "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, + "search_free_requests": { + "count": 5000, + "period": "month", + "shared_group": "gemini-3" + }, "web_search_billing_unit": "per_query", "google_maps_grounding_cost_per_query": 0.014 }, @@ -30241,7 +30370,11 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, - "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, + "search_free_requests": { + "count": 5000, + "period": "month", + "shared_group": "gemini-3" + }, "web_search_billing_unit": "per_query", "google_maps_grounding_cost_per_query": 0.014 }, @@ -30267,8 +30400,16 @@ "output_cost_per_token_priority": 6.75e-06, "prompt_cache_min_tokens": 4096.0, "regional_endpoint_uplift_multiplier": 1.1, - "search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014}, - "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, + "search_context_cost_per_query": { + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014, + "search_context_size_high": 0.014 + }, + "search_free_requests": { + "count": 5000, + "period": "month", + "shared_group": "gemini-3" + }, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", @@ -30415,8 +30556,16 @@ "output_cost_per_token_flex": 6e-06, "output_cost_per_token_priority": 2.16e-05, "prompt_cache_min_tokens": 4096.0, - "search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014}, - "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, + "search_context_cost_per_query": { + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014, + "search_context_size_high": 0.014 + }, + "search_free_requests": { + "count": 5000, + "period": "month", + "shared_group": "gemini-3" + }, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", @@ -30445,7 +30594,9 @@ "supports_video_input": true, "supports_vision": true, "supports_web_search": true, - "web_search_billing_unit": "per_query" + "web_search_billing_unit": "per_query", + "input_cost_per_token_above_200k_tokens_flex": 2e-06, + "output_cost_per_token_above_200k_tokens_flex": 9e-06 }, "vertex_ai/gemini-3.1-pro-preview-customtools": { "prompt_cache_min_tokens": 4096, @@ -30503,7 +30654,11 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, - "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, + "search_free_requests": { + "count": 5000, + "period": "month", + "shared_group": "gemini-3" + }, "web_search_billing_unit": "per_query", "google_maps_grounding_cost_per_query": 0.014 }, @@ -30831,7 +30986,11 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, - "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, + "search_free_requests": { + "count": 5000, + "period": "month", + "shared_group": "gemini-3" + }, "web_search_billing_unit": "per_query", "supports_reasoning": false }, @@ -30876,7 +31035,11 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, - "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, + "search_free_requests": { + "count": 5000, + "period": "month", + "shared_group": "gemini-3" + }, "web_search_billing_unit": "per_query" }, "gemini/gemini-3.1-flash-lite-preview": { @@ -30938,7 +31101,11 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, - "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, + "search_free_requests": { + "count": 5000, + "period": "month", + "shared_group": "gemini-3" + }, "web_search_billing_unit": "per_query", "google_maps_grounding_cost_per_query": 0.014, "input_cost_per_audio_token_batches": 2.5e-07 @@ -31041,11 +31208,15 @@ }, "gemini/gemini-2.5-flash": { "cache_read_input_audio_token_cost": 1e-07, + "cache_read_input_audio_token_cost_batches": 1e-07, + "cache_read_input_audio_token_cost_flex": 1e-07, + "cache_read_input_audio_token_cost_priority": 1.8e-07, "cache_read_input_token_cost": 3e-08, "cache_read_input_token_cost_batches": 3e-08, "cache_read_input_token_cost_flex": 3e-08, "cache_read_input_token_cost_priority": 5.4e-08, "input_cost_per_audio_token": 1e-06, + "input_cost_per_audio_token_flex": 5e-07, "input_cost_per_token": 3e-07, "litellm_provider": "gemini", "max_input_tokens": 1048576, @@ -31088,7 +31259,11 @@ "search_context_size_medium": 0.035, "search_context_size_high": 0.035 }, - "search_free_requests": {"count": 1500, "period": "day", "shared_group": "gemini-2-flash"}, + "search_free_requests": { + "count": 1500, + "period": "day", + "shared_group": "gemini-2-flash" + }, "google_maps_grounding_cost_per_query": 0.025, "input_cost_per_audio_token_batches": 5e-07, "input_cost_per_audio_token_priority": 1.8e-06, @@ -31201,7 +31376,11 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, - "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, + "search_free_requests": { + "count": 5000, + "period": "month", + "shared_group": "gemini-3" + }, "web_search_billing_unit": "per_query", "supports_reasoning": true }, @@ -31265,8 +31444,16 @@ "output_cost_per_token": 3e-06, "output_cost_per_token_batches": 1.5e-06, "rpm": 1000.0, - "search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014}, - "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, + "search_context_cost_per_query": { + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014, + "search_context_size_high": 0.014 + }, + "search_free_requests": { + "count": 5000, + "period": "month", + "shared_group": "gemini-3" + }, "source": "https://ai.google.dev/gemini-api/docs/pricing", "supported_endpoints": [ "/v1/chat/completions", @@ -31343,7 +31530,11 @@ "output_cost_per_token": 1.2e-05, "output_cost_per_token_batches": 6e-06, "rpm": 1000.0, - "search_context_cost_per_query": {"search_context_size_low":0.035,"search_context_size_medium":0.035,"search_context_size_high":0.035}, + "search_context_cost_per_query": { + "search_context_size_low": 0.035, + "search_context_size_medium": 0.035, + "search_context_size_high": 0.035 + }, "source": "https://ai.google.dev/gemini-api/docs/models/deep-research-pro-preview-12-2025", "supported_endpoints": [ "/v1/chat/completions", @@ -31370,11 +31561,15 @@ }, "gemini/gemini-2.5-flash-lite": { "cache_read_input_audio_token_cost": 3e-08, + "cache_read_input_audio_token_cost_batches": 3e-08, + "cache_read_input_audio_token_cost_flex": 3e-08, + "cache_read_input_audio_token_cost_priority": 5.4e-08, "cache_read_input_token_cost": 1e-08, "cache_read_input_token_cost_batches": 1e-08, "cache_read_input_token_cost_flex": 1e-08, "cache_read_input_token_cost_priority": 1.8e-08, "input_cost_per_audio_token": 3e-07, + "input_cost_per_audio_token_flex": 1.5e-07, "input_cost_per_token": 1e-07, "litellm_provider": "gemini", "max_input_tokens": 1048576, @@ -31417,7 +31612,11 @@ "search_context_size_medium": 0.035, "search_context_size_high": 0.035 }, - "search_free_requests": {"count": 1500, "period": "day", "shared_group": "gemini-2-flash"}, + "search_free_requests": { + "count": 1500, + "period": "day", + "shared_group": "gemini-2-flash" + }, "google_maps_grounding_cost_per_query": 0.025, "input_cost_per_audio_token_batches": 1.5e-07, "input_cost_per_audio_token_priority": 5.4e-07, @@ -31623,7 +31822,11 @@ "search_context_size_medium": 0.035, "search_context_size_high": 0.035 }, - "search_free_requests": {"count": 10000, "period": "day", "shared_group": "gemini-2.5-pro"}, + "search_free_requests": { + "count": 10000, + "period": "day", + "shared_group": "gemini-2.5-pro" + }, "google_maps_grounding_cost_per_query": 0.025, "input_cost_per_token_batches": 6.25e-07, "input_cost_per_token_flex": 6.25e-07, @@ -31664,12 +31867,16 @@ }, "gemini/gemini-3.1-flash-lite": { "cache_read_input_audio_token_cost": 5e-08, + "cache_read_input_audio_token_cost_batches": 2.5e-08, + "cache_read_input_audio_token_cost_flex": 2.5e-08, + "cache_read_input_audio_token_cost_priority": 9e-08, "cache_read_input_token_cost": 2.5e-08, "cache_read_input_token_cost_batches": 1.25e-08, "cache_read_input_token_cost_flex": 1.25e-08, "cache_read_input_token_cost_priority": 4.5e-08, "deprecation_date": "2027-05-07", "input_cost_per_audio_token": 5e-07, + "input_cost_per_audio_token_flex": 2.5e-07, "input_cost_per_token": 2.5e-07, "input_cost_per_token_batches": 1.25e-07, "input_cost_per_token_flex": 1.25e-07, @@ -31721,7 +31928,11 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, - "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, + "search_free_requests": { + "count": 5000, + "period": "month", + "shared_group": "gemini-3" + }, "web_search_billing_unit": "per_query", "google_maps_grounding_cost_per_query": 0.014, "input_cost_per_audio_token_batches": 2.5e-07, @@ -31784,16 +31995,24 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, - "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, + "search_free_requests": { + "count": 5000, + "period": "month", + "shared_group": "gemini-3" + }, "web_search_billing_unit": "per_query", "google_maps_grounding_cost_per_query": 0.014 }, "gemini/gemini-3-flash-preview": { "cache_read_input_audio_token_cost": 1e-07, + "cache_read_input_audio_token_cost_batches": 1e-07, + "cache_read_input_audio_token_cost_flex": 1e-07, + "cache_read_input_audio_token_cost_priority": 1.8e-07, "cache_read_input_token_cost": 5e-08, "cache_read_input_token_cost_batches": 5e-08, "cache_read_input_token_cost_flex": 5e-08, "input_cost_per_audio_token": 1e-06, + "input_cost_per_audio_token_flex": 5e-07, "input_cost_per_token": 5e-07, "litellm_provider": "gemini", "max_input_tokens": 1048576, @@ -31841,7 +32060,11 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, - "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, + "search_free_requests": { + "count": 5000, + "period": "month", + "shared_group": "gemini-3" + }, "web_search_billing_unit": "per_query", "google_maps_grounding_cost_per_query": 0.014, "input_cost_per_audio_token_batches": 5e-07, @@ -31905,7 +32128,11 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, - "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, + "search_free_requests": { + "count": 5000, + "period": "month", + "shared_group": "gemini-3" + }, "web_search_billing_unit": "per_query", "google_maps_grounding_cost_per_query": 0.014, "input_cost_per_token_batches": 7.5e-07, @@ -31972,7 +32199,11 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, - "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, + "search_free_requests": { + "count": 5000, + "period": "month", + "shared_group": "gemini-3" + }, "web_search_billing_unit": "per_query", "google_maps_grounding_cost_per_query": 0.014 }, @@ -32035,7 +32266,11 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, - "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, + "search_free_requests": { + "count": 5000, + "period": "month", + "shared_group": "gemini-3" + }, "web_search_billing_unit": "per_query", "google_maps_grounding_cost_per_query": 0.014 }, @@ -32098,7 +32333,11 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, - "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, + "search_free_requests": { + "count": 5000, + "period": "month", + "shared_group": "gemini-3" + }, "web_search_billing_unit": "per_query", "google_maps_grounding_cost_per_query": 0.014 }, @@ -32194,7 +32433,11 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, - "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, + "search_free_requests": { + "count": 5000, + "period": "month", + "shared_group": "gemini-3" + }, "web_search_billing_unit": "per_query", "google_maps_grounding_cost_per_query": 0.014, "input_cost_per_token_flex": 1e-06, @@ -32258,7 +32501,11 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, - "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, + "search_free_requests": { + "count": 5000, + "period": "month", + "shared_group": "gemini-3" + }, "web_search_billing_unit": "per_query", "google_maps_grounding_cost_per_query": 0.014, "input_cost_per_token_flex": 1e-06, @@ -32286,8 +32533,16 @@ "output_cost_per_token_batches": 1.5e-06, "output_cost_per_token_flex": 1.5e-06, "output_cost_per_token_priority": 5.4e-06, - "search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014}, - "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, + "search_context_cost_per_query": { + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014, + "search_context_size_high": 0.014 + }, + "search_free_requests": { + "count": 5000, + "period": "month", + "shared_group": "gemini-3" + }, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", @@ -32316,7 +32571,13 @@ "supports_url_context": true, "supports_vision": true, "supports_web_search": true, - "web_search_billing_unit": "per_query" + "web_search_billing_unit": "per_query", + "input_cost_per_audio_token_above_200k_tokens_priority": 1.8e-06, + "input_cost_per_audio_token_above_200k_tokens_flex": 5e-07, + "cache_read_input_audio_token_cost_priority": 1.8e-07, + "input_cost_per_audio_token_flex": 5e-07, + "input_cost_per_audio_token_above_200k_tokens": 1e-06, + "input_cost_per_audio_token_above_200k_tokens_batches": 5e-07 }, "gemini-omni-flash-preview": { "input_cost_per_audio_token": 1.5e-06, @@ -32402,7 +32663,11 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, - "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, + "search_free_requests": { + "count": 5000, + "period": "month", + "shared_group": "gemini-3" + }, "web_search_billing_unit": "per_query", "google_maps_grounding_cost_per_query": 0.014, "input_cost_per_token_batches": 7.5e-07, @@ -32468,7 +32733,11 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, - "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, + "search_free_requests": { + "count": 5000, + "period": "month", + "shared_group": "gemini-3" + }, "web_search_billing_unit": "per_query", "google_maps_grounding_cost_per_query": 0.014 }, @@ -32530,7 +32799,11 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, - "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, + "search_free_requests": { + "count": 5000, + "period": "month", + "shared_group": "gemini-3" + }, "web_search_billing_unit": "per_query", "google_maps_grounding_cost_per_query": 0.014 }, @@ -32555,8 +32828,16 @@ "output_cost_per_token_flex": 1.875e-06, "output_cost_per_token_priority": 6.75e-06, "prompt_cache_min_tokens": 4096.0, - "search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014}, - "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, + "search_context_cost_per_query": { + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014, + "search_context_size_high": 0.014 + }, + "search_free_requests": { + "count": 5000, + "period": "month", + "shared_group": "gemini-3" + }, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", @@ -41522,7 +41803,7 @@ "supports_xhigh_reasoning_effort": true }, "meta/muse-voice-transcribe-1.0": { - "input_cost_per_second": 0.00005, + "input_cost_per_second": 5e-05, "litellm_provider": "meta", "mode": "audio_transcription", "source": "https://dev.meta.ai/docs/speech-to-text", @@ -52342,8 +52623,8 @@ "/v1/images/generations", "/v1/images/edits", "/v1/batch" - ] - + ], + "output_cost_per_image_token_priority": 0.000216 }, "vertex_ai/gemini-3-pro-image-preview": { "input_cost_per_image": 0.0011, @@ -52376,8 +52657,8 @@ "/v1/images/generations", "/v1/images/edits", "/v1/batch" - ] - + ], + "output_cost_per_image_token_priority": 0.000216 }, "vertex_ai/gemini-3.1-flash-image": { "cache_read_input_token_cost": 5e-08, @@ -52409,8 +52690,8 @@ "/v1/images/generations", "/v1/images/edits", "/v1/batch" - ] - + ], + "output_cost_per_image_token_priority": 0.000108 }, "vertex_ai/gemini-3.1-flash-image-preview": { "input_cost_per_image": 0.00056, @@ -52434,15 +52715,14 @@ "/v1/images/generations", "/v1/images/edits", "/v1/batch" - ] - + ], + "output_cost_per_image_token_priority": 0.000108 }, "vertex_ai/gemini-3.1-flash-lite-image": { "cache_read_input_token_cost": 2.5e-08, "cache_read_input_token_cost_batches": 1.25e-08, "cache_read_input_token_cost_flex": 1.25e-08, "cache_read_input_token_cost_priority": 4.5e-08, - "deprecation_date": "2027-06-28", "input_cost_per_image": 0.00028, "input_cost_per_token": 2.5e-07, "input_cost_per_token_batches": 1.25e-07, @@ -52482,7 +52762,8 @@ "supports_response_schema": false, "supports_system_messages": true, "supports_video_input": true, - "supports_vision": true + "supports_vision": true, + "output_cost_per_image_token_priority": 5.4e-05 }, "vertex_ai/gemini-3.1-flash-lite-preview": { "cache_read_input_token_cost": 2.5e-08, @@ -52530,7 +52811,11 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, - "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, + "search_free_requests": { + "count": 5000, + "period": "month", + "shared_group": "gemini-3" + }, "web_search_billing_unit": "per_query", "google_maps_grounding_cost_per_query": 0.014 }, @@ -52593,10 +52878,22 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, - "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, + "search_free_requests": { + "count": 5000, + "period": "month", + "shared_group": "gemini-3" + }, "web_search_billing_unit": "per_query", "google_maps_grounding_cost_per_query": 0.014, - "input_cost_per_audio_token_batches": 2.5e-07 + "input_cost_per_audio_token_batches": 2.5e-07, + "input_cost_per_audio_token_above_200k_tokens_priority": 9e-07, + "cache_read_input_audio_token_cost_flex": 2.5e-08, + "input_cost_per_audio_token_above_200k_tokens_flex": 2.5e-07, + "cache_read_input_audio_token_cost_priority": 9e-08, + "input_cost_per_audio_token_flex": 2.5e-07, + "input_cost_per_audio_token_above_200k_tokens": 5e-07, + "cache_read_input_audio_token_cost_batches": 2.5e-08, + "input_cost_per_audio_token_above_200k_tokens_batches": 2.5e-07 }, "vertex_ai/gemini-3.5-flash-lite": { "deprecation_date": "2027-07-21", @@ -52655,7 +52952,11 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, - "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, + "search_free_requests": { + "count": 5000, + "period": "month", + "shared_group": "gemini-3" + }, "web_search_billing_unit": "per_query", "google_maps_grounding_cost_per_query": 0.014 }, @@ -61612,11 +61913,41 @@ "output_cost_per_second_480p": 0.096075, "output_cost_per_second_720p": 0.216, "output_cost_per_second_1080p": 0.5346, - "supported_modalities": ["text", "image", "video", "audio"], - "supported_output_modalities": ["video"], - "supported_endpoints": ["/v1/videos"], - "provider_specific_entry": {"volcengine_video_frames": {"480p": {"16:9": "854x480", "4:3": "752x560", "1:1": "640x640", "3:4": "560x752", "9:16": "480x854", "21:9": "992x432"}, "720p": {"16:9": "1280x720", "4:3": "1112x834", "1:1": "960x960", "3:4": "834x1112", "9:16": "720x1280", "21:9": "1470x630"}}}, - "metadata": {"comment": "Ark bills per output token, not per second: tokens = width * height * 24 * seconds / 1024, so a per-second rate per resolution tier is exact for the 16:9 geometry Ark publishes (1280x720 is 21600 tokens/s). List 70 CNY per million tokens at 480p and 720p and 77 at 1080p without a video input, matching Ark's per-resolution table and its 1080p 16:9 example of 3.74 CNY/s, converted at 7.0 CNY/USD like the other volcengine entries. Accepts up to 50 reference assets (30 images, 10 videos, 10 audios) and clips of 4-30s; per-asset billing beyond the free allowance is not encoded. Requests with a video input bill at a lower rate per token and are rejected, since their cost is not known at create time. provider_specific_entry.volcengine_video_frames is the frame Ark renders per resolution tier and ratio, verified with ffprobe at 480p for every ratio and at 720p for 16:9 and 4:3; the other 720p ratios are from the Seedance 2.0 model card."} + "supported_modalities": [ + "text", + "image", + "video", + "audio" + ], + "supported_output_modalities": [ + "video" + ], + "supported_endpoints": [ + "/v1/videos" + ], + "provider_specific_entry": { + "volcengine_video_frames": { + "480p": { + "16:9": "854x480", + "4:3": "752x560", + "1:1": "640x640", + "3:4": "560x752", + "9:16": "480x854", + "21:9": "992x432" + }, + "720p": { + "16:9": "1280x720", + "4:3": "1112x834", + "1:1": "960x960", + "3:4": "834x1112", + "9:16": "720x1280", + "21:9": "1470x630" + } + } + }, + "metadata": { + "comment": "Ark bills per output token, not per second: tokens = width * height * 24 * seconds / 1024, so a per-second rate per resolution tier is exact for the 16:9 geometry Ark publishes (1280x720 is 21600 tokens/s). List 70 CNY per million tokens at 480p and 720p and 77 at 1080p without a video input, matching Ark's per-resolution table and its 1080p 16:9 example of 3.74 CNY/s, converted at 7.0 CNY/USD like the other volcengine entries. Accepts up to 50 reference assets (30 images, 10 videos, 10 audios) and clips of 4-30s; per-asset billing beyond the free allowance is not encoded. Requests with a video input bill at a lower rate per token and are rejected, since their cost is not known at create time. provider_specific_entry.volcengine_video_frames is the frame Ark renders per resolution tier and ratio, verified with ffprobe at 480p for every ratio and at 720p for 16:9 and 4:3; the other 720p ratios are from the Seedance 2.0 model card." + } }, "volcengine/doubao-seedance-2-0-260128": { "litellm_provider": "volcengine", @@ -61627,11 +61958,41 @@ "output_cost_per_second_720p": 0.141942857, "output_cost_per_second_1080p": 0.354085714, "output_cost_per_second_4k": 0.722057143, - "supported_modalities": ["text", "image", "video", "audio"], - "supported_output_modalities": ["video"], - "supported_endpoints": ["/v1/videos"], - "provider_specific_entry": {"volcengine_video_frames": {"480p": {"16:9": "864x496", "4:3": "752x560", "1:1": "640x640", "3:4": "560x752", "9:16": "496x864", "21:9": "992x432"}, "720p": {"16:9": "1280x720", "4:3": "1112x834", "1:1": "960x960", "3:4": "834x1112", "9:16": "720x1280", "21:9": "1470x630"}}}, - "metadata": {"comment": "Ark bills per output token, not per second: tokens = width * height * 24 * seconds / 1024, so a per-second rate per resolution tier is exact for the 16:9 geometry Ark publishes (1280x720 is 21600 tokens/s; the 480p frame is 864x496, 10044 tokens/s, not the tier's nominal 854x480). List 46 CNY per million tokens at 480p and 720p, 51 at 1080p, and 26 at 4k without a video input, matching Ark's per-resolution table and its 1080p and 4k 16:9 examples of 2.48 and 5.05 CNY/s, converted at 7.0 CNY/USD like the other volcengine entries. Clips of 4-15s; 4k is exclusive to this tier. Requests with a video input bill at a lower rate per token and are rejected, since their cost is not known at create time. provider_specific_entry.volcengine_video_frames is the frame Ark renders per resolution tier and ratio, verified with ffprobe at 480p for every ratio and at 720p for 16:9 and 4:3; the other 720p ratios are from the Seedance 2.0 model card."} + "supported_modalities": [ + "text", + "image", + "video", + "audio" + ], + "supported_output_modalities": [ + "video" + ], + "supported_endpoints": [ + "/v1/videos" + ], + "provider_specific_entry": { + "volcengine_video_frames": { + "480p": { + "16:9": "864x496", + "4:3": "752x560", + "1:1": "640x640", + "3:4": "560x752", + "9:16": "496x864", + "21:9": "992x432" + }, + "720p": { + "16:9": "1280x720", + "4:3": "1112x834", + "1:1": "960x960", + "3:4": "834x1112", + "9:16": "720x1280", + "21:9": "1470x630" + } + } + }, + "metadata": { + "comment": "Ark bills per output token, not per second: tokens = width * height * 24 * seconds / 1024, so a per-second rate per resolution tier is exact for the 16:9 geometry Ark publishes (1280x720 is 21600 tokens/s; the 480p frame is 864x496, 10044 tokens/s, not the tier's nominal 854x480). List 46 CNY per million tokens at 480p and 720p, 51 at 1080p, and 26 at 4k without a video input, matching Ark's per-resolution table and its 1080p and 4k 16:9 examples of 2.48 and 5.05 CNY/s, converted at 7.0 CNY/USD like the other volcengine entries. Clips of 4-15s; 4k is exclusive to this tier. Requests with a video input bill at a lower rate per token and are rejected, since their cost is not known at create time. provider_specific_entry.volcengine_video_frames is the frame Ark renders per resolution tier and ratio, verified with ffprobe at 480p for every ratio and at 720p for 16:9 and 4:3; the other 720p ratios are from the Seedance 2.0 model card." + } }, "volcengine/doubao-seedance-2-0-fast-260128": { "litellm_provider": "volcengine", @@ -61640,11 +62001,41 @@ "output_cost_per_second": 0.114171429, "output_cost_per_second_480p": 0.053089714, "output_cost_per_second_720p": 0.114171429, - "supported_modalities": ["text", "image", "video", "audio"], - "supported_output_modalities": ["video"], - "supported_endpoints": ["/v1/videos"], - "provider_specific_entry": {"volcengine_video_frames": {"480p": {"16:9": "864x496", "4:3": "752x560", "1:1": "640x640", "3:4": "560x752", "9:16": "496x864", "21:9": "992x432"}, "720p": {"16:9": "1280x720", "4:3": "1112x834", "1:1": "960x960", "3:4": "834x1112", "9:16": "720x1280", "21:9": "1470x630"}}}, - "metadata": {"comment": "Ark bills per output token, not per second: tokens = width * height * 24 * seconds / 1024, so a per-second rate per resolution tier is exact for the 16:9 geometry Ark publishes (1280x720 is 21600 tokens/s; the 480p frame is 864x496, 10044 tokens/s, not the tier's nominal 854x480). List 37 CNY per million tokens without a video input, converted at 7.0 CNY/USD like the other volcengine entries. Clips of 4-15s at 480p or 720p. Requests with a video input bill at a lower rate per token and are rejected, since their cost is not known at create time. provider_specific_entry.volcengine_video_frames is the frame Ark renders per resolution tier and ratio, verified with ffprobe at 480p for every ratio and at 720p for 16:9 and 4:3; the other 720p ratios are from the Seedance 2.0 model card."} + "supported_modalities": [ + "text", + "image", + "video", + "audio" + ], + "supported_output_modalities": [ + "video" + ], + "supported_endpoints": [ + "/v1/videos" + ], + "provider_specific_entry": { + "volcengine_video_frames": { + "480p": { + "16:9": "864x496", + "4:3": "752x560", + "1:1": "640x640", + "3:4": "560x752", + "9:16": "496x864", + "21:9": "992x432" + }, + "720p": { + "16:9": "1280x720", + "4:3": "1112x834", + "1:1": "960x960", + "3:4": "834x1112", + "9:16": "720x1280", + "21:9": "1470x630" + } + } + }, + "metadata": { + "comment": "Ark bills per output token, not per second: tokens = width * height * 24 * seconds / 1024, so a per-second rate per resolution tier is exact for the 16:9 geometry Ark publishes (1280x720 is 21600 tokens/s; the 480p frame is 864x496, 10044 tokens/s, not the tier's nominal 854x480). List 37 CNY per million tokens without a video input, converted at 7.0 CNY/USD like the other volcengine entries. Clips of 4-15s at 480p or 720p. Requests with a video input bill at a lower rate per token and are rejected, since their cost is not known at create time. provider_specific_entry.volcengine_video_frames is the frame Ark renders per resolution tier and ratio, verified with ffprobe at 480p for every ratio and at 720p for 16:9 and 4:3; the other 720p ratios are from the Seedance 2.0 model card." + } }, "volcengine/doubao-seedance-2-0-mini-260615": { "litellm_provider": "volcengine", @@ -61653,11 +62044,39 @@ "output_cost_per_second": 0.070971429, "output_cost_per_second_480p": 0.033001714, "output_cost_per_second_720p": 0.070971429, - "supported_modalities": ["text", "image"], - "supported_output_modalities": ["video"], - "supported_endpoints": ["/v1/videos"], - "provider_specific_entry": {"volcengine_video_frames": {"480p": {"16:9": "864x496", "4:3": "752x560", "1:1": "640x640", "3:4": "560x752", "9:16": "496x864", "21:9": "992x432"}, "720p": {"16:9": "1280x720", "4:3": "1112x834", "1:1": "960x960", "3:4": "834x1112", "9:16": "720x1280", "21:9": "1470x630"}}}, - "metadata": {"comment": "Ark bills per output token, not per second: tokens = width * height * 24 * seconds / 1024, so a per-second rate per resolution tier is exact for the 16:9 geometry Ark publishes (1280x720 is 21600 tokens/s; the 480p frame is 864x496, 10044 tokens/s, not the tier's nominal 854x480). List 23 CNY per million tokens without a video input, converted at 7.0 CNY/USD like the other volcengine entries. Clips of 4-15s at 480p or 720p. Only text and image inputs are declared; the mini tier's reference video and audio caps were not verified. Requests with a video input bill at a lower rate per token and are rejected, since their cost is not known at create time. provider_specific_entry.volcengine_video_frames is the frame Ark renders per resolution tier and ratio, verified with ffprobe at 480p for every ratio and at 720p for 16:9 and 4:3; the other 720p ratios are from the Seedance 2.0 model card."} + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "video" + ], + "supported_endpoints": [ + "/v1/videos" + ], + "provider_specific_entry": { + "volcengine_video_frames": { + "480p": { + "16:9": "864x496", + "4:3": "752x560", + "1:1": "640x640", + "3:4": "560x752", + "9:16": "496x864", + "21:9": "992x432" + }, + "720p": { + "16:9": "1280x720", + "4:3": "1112x834", + "1:1": "960x960", + "3:4": "834x1112", + "9:16": "720x1280", + "21:9": "1470x630" + } + } + }, + "metadata": { + "comment": "Ark bills per output token, not per second: tokens = width * height * 24 * seconds / 1024, so a per-second rate per resolution tier is exact for the 16:9 geometry Ark publishes (1280x720 is 21600 tokens/s; the 480p frame is 864x496, 10044 tokens/s, not the tier's nominal 854x480). List 23 CNY per million tokens without a video input, converted at 7.0 CNY/USD like the other volcengine entries. Clips of 4-15s at 480p or 720p. Only text and image inputs are declared; the mini tier's reference video and audio caps were not verified. Requests with a video input bill at a lower rate per token and are rejected, since their cost is not known at create time. provider_specific_entry.volcengine_video_frames is the frame Ark renders per resolution tier and ratio, verified with ffprobe at 480p for every ratio and at 720p for 16:9 and 4:3; the other 720p ratios are from the Seedance 2.0 model card." + } }, "bedrock/us-east-1/zai.glm-5": { "input_cost_per_token": 1e-06, @@ -63518,7 +63937,13 @@ { "name": "claude-adaptive-thinking", "pattern": "claude-[a-z]+-(?:4[-._](?:[6-9]|[1-9]\\d)(?!\\d)|[5-9](?!\\d)(?:[-._]\\d{1,2}(?!\\d))?)", - "fill_missing_for_providers": ["anthropic", "azure_ai", "bedrock", "bedrock_converse", "vertex_ai-anthropic_models"], + "fill_missing_for_providers": [ + "anthropic", + "azure_ai", + "bedrock", + "bedrock_converse", + "vertex_ai-anthropic_models" + ], "description": "Claude at version 4.6 or higher, in any id shape that contains claude--: minors 4.6 through 4.99, any later major-minor, and bare 5+ majors so a new family shaped like claude-fable-5 matches. Requiring the claude- prefix keeps non-Claude names such as team-sonnet-5-1 out. The minor is capped at two digits so an 8-digit date suffix such as claude-opus-4-20250514 is never read as a >= 4.6 minor. Two-digit majors are deliberately not matched so ids like claude-opus-41 (4.1) are not read as major 41. Turns on adaptive thinking for new versions and new families with no code change.", "model_info": { "supports_adaptive_thinking": true @@ -63527,7 +63952,13 @@ { "name": "claude-legacy-thinking", "pattern": "claude-[a-z]+-4[-._]6(?!\\d)", - "fill_missing_for_providers": ["anthropic", "azure_ai", "bedrock", "bedrock_converse", "vertex_ai-anthropic_models"], + "fill_missing_for_providers": [ + "anthropic", + "azure_ai", + "bedrock", + "bedrock_converse", + "vertex_ai-anthropic_models" + ], "description": "Claude at version 4.6 exactly, in any id shape that contains claude--4-6 (dotted and underscored minors included, dated releases such as claude-sonnet-4-6-20260219 too). The 4.6 family is adaptive-thinking yet still accepts legacy thinking.type=enabled with budget_tokens, so the caller's hard budget cap is forwarded verbatim instead of being rewritten to an uncapped output_config.effort. The lookahead keeps two-digit minors such as 4-60 from matching. 4.7+ and 5+ majors reject the legacy shape and stay on the adaptive translation.", "model_info": { "supports_legacy_thinking": true @@ -63544,7 +63975,13 @@ { "name": "claude-mid-conversation-system", "pattern": "claude-[a-z]+-(?:4[-._](?:[89]|[1-9]\\d)(?!\\d)|[5-9](?!\\d)(?:[-._]\\d{1,2}(?!\\d))?)", - "fill_missing_for_providers": ["anthropic", "azure_ai", "bedrock", "bedrock_converse", "vertex_ai-anthropic_models"], + "fill_missing_for_providers": [ + "anthropic", + "azure_ai", + "bedrock", + "bedrock_converse", + "vertex_ai-anthropic_models" + ], "description": "Claude at version 4.8 or higher, in any id shape that contains claude--: minors 4.8 through 4.99, any later major-minor, and bare 5+ majors so a new family like claude-fable-5 matches. Two-digit majors are deliberately not matched so ids like claude-opus-41 (4.1) are not read as major 41. Anthropic introduced mid-conversation system messages with Opus 4.8 and every newer Claude keeps them; 4.7 and below reject the system role inside messages.", "model_info": { "supports_mid_conversation_system": true @@ -63553,7 +63990,13 @@ { "name": "claude-tool-search", "pattern": "claude-[a-z]+-(?:4[-._](?:[5-9]|[1-9]\\d)(?!\\d)|[5-9](?!\\d)(?:[-._]\\d{1,2}(?!\\d))?)", - "fill_missing_for_providers": ["anthropic", "azure_ai", "bedrock", "bedrock_converse", "vertex_ai-anthropic_models"], + "fill_missing_for_providers": [ + "anthropic", + "azure_ai", + "bedrock", + "bedrock_converse", + "vertex_ai-anthropic_models" + ], "description": "Claude at version 4.5 or higher, in any id shape that contains claude--: minors 4.5 through 4.99, any later major-minor, and bare 5+ majors so a new family like claude-fable-5 matches. Two-digit majors are deliberately not matched so ids like claude-opus-41 (4.1) are not read as major 41. Anthropic's tool search docs list every Claude 4.5 and newer model as supported and Opus 4.1 and earlier as unsupported, so the flag follows the version instead of a per-model list. azure_ai is included because Microsoft Foundry runs tool search for Claude on both Anthropic-hosted and Azure-hosted deployments.", "model_info": { "supports_tool_search": true @@ -63570,7 +64013,11 @@ { "name": "openai-reasoning-family-baseline", "pattern": "^(?!.*search-api)(?:[a-z0-9_.-]+/)*(?:ft:)?(?:o[1-9]\\d*(?![a-z0-9])|gpt-[5-9](?:\\.\\d+)?(?![0-9.])|(?:gpt-\\d+(?:\\.\\d+)?(?:-[a-z0-9]+)*-)?(?:codex|deep-research|chat-latest)(?![a-z0-9]))", - "fill_missing_for_providers": ["azure", "azure_ai", "openai"], + "fill_missing_for_providers": [ + "azure", + "azure_ai", + "openai" + ], "description": "OpenAI reasoning families by id shape, under any provider namespace and with an optional ft: prefix: the o-series (o1, o3-pro, o4-mini), gpt-5 through gpt-9 majors including dotted minors and suffixed variants (gpt-5.5-cyber, gpt-6-astra), and the codex, deep-research and chat-latest lines when standalone or on a gpt base. gpt-5-search-api is excluded because it is a search-only surface. Every model here is a reasoning model, and the Responses API drops the caller's reasoning param for any mapped OpenAI model whose info lacks supports_reasoning, so an id the registry has not named yet keeps its reasoning settings instead of silently losing them. Rules lose to exact entries. Carries no mode and no pricing, so cost stays on the standard unpriced behavior.", "model_info": { "supports_reasoning": true @@ -73263,7 +73710,9 @@ "mode": "chat", "output_cost_per_token": 4.4e-06, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", - "supported_regions": ["global"], + "supported_regions": [ + "global" + ], "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, @@ -74604,7 +75053,9 @@ }, "global.xai.grok-4.7": { "supports_regex_lookaround": false, - "supported_endpoints": ["/v1/chat/completions"], + "supported_endpoints": [ + "/v1/chat/completions" + ], "supports_bedrock_runtime_chat_completions_tools_with_reasoning": true, "supports_bedrock_runtime_chat_completions_response_format": true, "cache_read_input_token_cost": 5e-07, @@ -74630,7 +75081,9 @@ }, "us.xai.grok-4.7": { "supports_regex_lookaround": false, - "supported_endpoints": ["/v1/chat/completions"], + "supported_endpoints": [ + "/v1/chat/completions" + ], "supports_bedrock_runtime_chat_completions_tools_with_reasoning": true, "supports_bedrock_runtime_chat_completions_response_format": true, "cache_read_input_token_cost": 5.5e-07, @@ -75079,7 +75532,11 @@ "output_cost_per_token": 7.5e-06, "output_cost_per_token_batches": 3.75e-06, "rpm": 1000.0, - "search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014}, + "search_context_cost_per_query": { + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014, + "search_context_size_high": 0.014 + }, "source": "https://ai.google.dev/gemini-api/docs/pricing", "supported_endpoints": [ "/v1/chat/completions", @@ -75124,7 +75581,11 @@ "output_cost_per_token": 7.5e-06, "output_cost_per_token_batches": 3.75e-06, "output_cost_per_token_priority": 1.35e-05, - "search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014}, + "search_context_cost_per_query": { + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014, + "search_context_size_high": 0.014 + }, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", @@ -75150,7 +75611,8 @@ "supports_video_input": true, "supports_vision": true, "supports_web_search": true, - "web_search_billing_unit": "per_query" + "web_search_billing_unit": "per_query", + "output_cost_per_image_token_priority": 5.4e-05 }, "vertex_ai/gemini-nano-banana-2.1": { "cache_read_input_token_cost": 1.5e-07, @@ -75170,7 +75632,11 @@ "output_cost_per_token": 7.5e-06, "output_cost_per_token_batches": 3.75e-06, "output_cost_per_token_priority": 1.35e-05, - "search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014}, + "search_context_cost_per_query": { + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014, + "search_context_size_high": 0.014 + }, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", @@ -75196,7 +75662,8 @@ "supports_video_input": true, "supports_vision": true, "supports_web_search": true, - "web_search_billing_unit": "per_query" + "web_search_billing_unit": "per_query", + "output_cost_per_image_token_priority": 5.4e-05 }, "azure_ai/grok-4.7": { "cache_read_input_token_cost": 5e-07, @@ -75252,7 +75719,7 @@ "supports_tool_choice": true, "supports_vision": true }, -"claude-haiku-5-5": { + "claude-haiku-5-5": { "input_cost_per_token_above_100k_tokens_batches": 2.5e-07, "output_cost_per_token_above_100k_tokens_batches": 1.25e-06, "cache_creation_input_token_cost_above_100k_tokens_batches": 3.125e-07, @@ -76179,7 +76646,11 @@ "output_cost_per_reasoning_token": 5e-06, "output_cost_per_token": 5e-06, "output_cost_per_token_batches": 2.5e-06, - "search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014}, + "search_context_cost_per_query": { + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014, + "search_context_size_high": 0.014 + }, "source": "https://ai.google.dev/gemini-api/docs/pricing", "supported_endpoints": [ "/v1/chat/completions", @@ -76225,7 +76696,11 @@ "output_cost_per_reasoning_token": 5e-06, "output_cost_per_token": 5e-06, "output_cost_per_token_batches": 2.5e-06, - "search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014}, + "search_context_cost_per_query": { + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014, + "search_context_size_high": 0.014 + }, "source": "https://ai.google.dev/gemini-api/docs/pricing", "supported_endpoints": [ "/v1/chat/completions", @@ -76279,7 +76754,11 @@ "output_cost_per_token_priority": 1.8e-05, "prompt_cache_min_tokens": 4096.0, "rpm": 2000.0, - "search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014}, + "search_context_cost_per_query": { + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014, + "search_context_size_high": 0.014 + }, "source": "https://ai.google.dev/gemini-api/docs/pricing", "supported_endpoints": [ "/v1/chat/completions", @@ -76338,7 +76817,11 @@ "output_cost_per_token_priority": 1.8e-05, "prompt_cache_min_tokens": 6144.0, "regional_endpoint_uplift_multiplier": 1.1, - "search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014}, + "search_context_cost_per_query": { + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014, + "search_context_size_high": 0.014 + }, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", @@ -76396,7 +76879,11 @@ "output_cost_per_token_priority": 3.6e-05, "prompt_cache_min_tokens": 6144.0, "regional_endpoint_uplift_multiplier": 1.1, - "search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014}, + "search_context_cost_per_query": { + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014, + "search_context_size_high": 0.014 + }, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", @@ -76624,7 +77111,11 @@ "mode": "chat", "output_cost_per_token": 4e-05, "prompt_cache_min_tokens": 512.0, - "search_context_cost_per_query": {"search_context_size_high":0.02,"search_context_size_low":0.02,"search_context_size_medium":0.02}, + "search_context_cost_per_query": { + "search_context_size_high": 0.02, + "search_context_size_low": 0.02, + "search_context_size_medium": 0.02 + }, "source": "https://learn.microsoft.com/en-us/azure/foundry/foundry-models/concepts/claude-models", "supports_adaptive_thinking": true, "supports_assistant_prefill": false, @@ -76988,5 +77479,605 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true + }, + "bedrock/claude_platform/claude-fable-5": { + "supports_web_search": true, + "supports_anthropic_compaction": true, + "cache_creation_input_token_cost": 1.25e-05, + "cache_creation_input_token_cost_above_1hr": 2e-05, + "cache_creation_input_token_cost_batches": 6.25e-06, + "cache_read_input_token_cost": 1e-06, + "cache_read_input_token_cost_batches": 5e-07, + "input_cost_per_token": 1e-05, + "input_cost_per_token_batches": 5e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5e-05, + "output_cost_per_token_batches": 2.5e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "thinking_always_on": true, + "supports_mid_conversation_system": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true, + "provider_specific_entry": { + "us": 1.1 + }, + "supports_output_config": true, + "prompt_cache_min_tokens": 512, + "supports_native_structured_output": true, + "source": "https://platform.claude.com/docs/en/about-claude/pricing#claude-platform-on-aws-pricing" + }, + "bedrock/claude_platform/claude-fable-5-1": { + "supports_web_search": true, + "supports_anthropic_compaction": true, + "cache_creation_input_token_cost": 1.25e-05, + "cache_creation_input_token_cost_above_1hr": 2e-05, + "cache_creation_input_token_cost_batches": 6.25e-06, + "cache_read_input_token_cost": 2.5e-07, + "cache_read_input_token_cost_batches": 1.25e-07, + "input_cost_per_token": 1e-05, + "input_cost_per_token_batches": 5e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5e-05, + "output_cost_per_token_batches": 2.5e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "thinking_always_on": true, + "supports_mid_conversation_system": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_forced_tool_use": false, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true, + "provider_specific_entry": { + "us": 1.1 + }, + "supports_output_config": true, + "prompt_cache_min_tokens": 512, + "supports_native_structured_output": true, + "source": "https://platform.claude.com/docs/en/about-claude/pricing#claude-platform-on-aws-pricing" + }, + "bedrock/claude_platform/claude-haiku-4-5": { + "cache_creation_input_token_cost": 1.25e-06, + "cache_creation_input_token_cost_above_1hr": 2e-06, + "cache_creation_input_token_cost_batches": 6.25e-07, + "cache_read_input_token_cost": 1e-07, + "cache_read_input_token_cost_batches": 5e-08, + "input_cost_per_token": 1e-06, + "input_cost_per_token_batches": 5e-07, + "litellm_provider": "bedrock", + "max_input_tokens": 200000, + "max_output_tokens": 64000, + "max_tokens": 64000, + "mode": "chat", + "output_cost_per_token": 5e-06, + "output_cost_per_token_batches": 2.5e-06, + "supports_assistant_prefill": true, + "supports_function_calling": true, + "supports_native_structured_output": true, + "supports_computer_use": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true, + "prompt_cache_min_tokens": 4096, + "source": "https://platform.claude.com/docs/en/about-claude/pricing#claude-platform-on-aws-pricing" + }, + "bedrock/claude_platform/claude-haiku-5-5": { + "input_cost_per_token_above_100k_tokens_batches": 2.5e-07, + "output_cost_per_token_above_100k_tokens_batches": 1.25e-06, + "cache_creation_input_token_cost_above_100k_tokens_batches": 3.125e-07, + "cache_read_input_token_cost_above_100k_tokens_batches": 2.5e-08, + "supports_anthropic_compaction": true, + "cache_creation_input_token_cost": 1.25e-07, + "cache_creation_input_token_cost_above_1hr": 2e-07, + "cache_creation_input_token_cost_batches": 6.25e-08, + "cache_read_input_token_cost": 1e-08, + "cache_read_input_token_cost_batches": 5e-09, + "input_cost_per_token": 1e-07, + "input_cost_per_token_batches": 5e-08, + "litellm_provider": "bedrock", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5e-07, + "output_cost_per_token_batches": 2.5e-07, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "supports_mid_conversation_system": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_native_structured_output": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true, + "provider_specific_entry": { + "us": 1.1 + }, + "supports_output_config": true, + "supports_forced_tool_use": true, + "thinking_always_on": false, + "prompt_cache_min_tokens": 512, + "source": "https://platform.claude.com/docs/en/about-claude/pricing#claude-platform-on-aws-pricing", + "supports_web_search": true, + "input_cost_per_token_above_100k_tokens": 5e-07, + "output_cost_per_token_above_100k_tokens": 2.5e-06, + "cache_creation_input_token_cost_above_100k_tokens": 6.25e-07, + "cache_creation_input_token_cost_above_1hr_above_100k_tokens": 1e-06, + "cache_read_input_token_cost_above_100k_tokens": 5e-08 + }, + "bedrock/claude_platform/claude-opus-4-5": { + "cache_creation_input_token_cost": 6.25e-06, + "cache_creation_input_token_cost_above_1hr": 1e-05, + "cache_creation_input_token_cost_batches": 3.125e-06, + "cache_read_input_token_cost": 5e-07, + "cache_read_input_token_cost_batches": 2.5e-07, + "input_cost_per_token": 5e-06, + "input_cost_per_token_batches": 2.5e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 200000, + "max_output_tokens": 64000, + "max_tokens": 64000, + "mode": "chat", + "output_cost_per_token": 2.5e-05, + "output_cost_per_token_batches": 1.25e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_assistant_prefill": true, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_thinking_cache_preservation": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_native_structured_output": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_output_config": true, + "prompt_cache_min_tokens": 4096, + "source": "https://platform.claude.com/docs/en/about-claude/pricing#claude-platform-on-aws-pricing" + }, + "bedrock/claude_platform/claude-opus-4-6": { + "supports_web_search": true, + "supports_anthropic_compaction": true, + "cache_creation_input_token_cost": 6.25e-06, + "cache_creation_input_token_cost_above_1hr": 1e-05, + "cache_creation_input_token_cost_batches": 3.125e-06, + "cache_read_input_token_cost": 5e-07, + "cache_read_input_token_cost_batches": 2.5e-07, + "input_cost_per_token": 5e-06, + "input_cost_per_token_batches": 2.5e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 2.5e-05, + "output_cost_per_token_batches": 1.25e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "supports_legacy_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_thinking_cache_preservation": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_native_structured_output": true, + "supports_tool_choice": true, + "supports_vision": true, + "provider_specific_entry": { + "us": 1.1 + }, + "supports_output_config": true, + "supports_max_reasoning_effort": true, + "prompt_cache_min_tokens": 4096, + "source": "https://platform.claude.com/docs/en/about-claude/pricing#claude-platform-on-aws-pricing" + }, + "bedrock/claude_platform/claude-opus-4-7": { + "supports_web_search": true, + "supports_anthropic_compaction": true, + "cache_creation_input_token_cost": 6.25e-06, + "cache_creation_input_token_cost_above_1hr": 1e-05, + "cache_creation_input_token_cost_batches": 3.125e-06, + "cache_read_input_token_cost": 5e-07, + "cache_read_input_token_cost_batches": 2.5e-07, + "input_cost_per_token": 5e-06, + "input_cost_per_token_batches": 2.5e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 2.5e-05, + "output_cost_per_token_batches": 1.25e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_thinking_cache_preservation": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_native_structured_output": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true, + "provider_specific_entry": { + "us": 1.1 + }, + "supports_output_config": true, + "prompt_cache_min_tokens": 2048, + "source": "https://platform.claude.com/docs/en/about-claude/pricing#claude-platform-on-aws-pricing" + }, + "bedrock/claude_platform/claude-opus-4-8": { + "supports_web_search": true, + "supports_anthropic_compaction": true, + "cache_creation_input_token_cost": 6.25e-06, + "cache_creation_input_token_cost_above_1hr": 1e-05, + "cache_creation_input_token_cost_batches": 3.125e-06, + "cache_read_input_token_cost": 5e-07, + "cache_read_input_token_cost_batches": 2.5e-07, + "input_cost_per_token": 5e-06, + "input_cost_per_token_batches": 2.5e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 2.5e-05, + "output_cost_per_token_batches": 1.25e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "supports_mid_conversation_system": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_thinking_cache_preservation": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_native_structured_output": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true, + "provider_specific_entry": { + "us": 1.1 + }, + "supports_output_config": true, + "prompt_cache_min_tokens": 1024, + "source": "https://platform.claude.com/docs/en/about-claude/pricing#claude-platform-on-aws-pricing" + }, + "bedrock/claude_platform/claude-opus-5": { + "supports_web_search": true, + "supports_anthropic_compaction": true, + "cache_creation_input_token_cost": 6.25e-06, + "cache_creation_input_token_cost_above_1hr": 1e-05, + "cache_creation_input_token_cost_batches": 3.125e-06, + "cache_read_input_token_cost": 5e-07, + "cache_read_input_token_cost_batches": 2.5e-07, + "input_cost_per_token": 5e-06, + "input_cost_per_token_batches": 2.5e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 2.5e-05, + "output_cost_per_token_batches": 1.25e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "supports_mid_conversation_system": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_thinking_cache_preservation": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_native_structured_output": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true, + "provider_specific_entry": { + "us": 1.1 + }, + "supports_output_config": true, + "prompt_cache_min_tokens": 512, + "source": "https://platform.claude.com/docs/en/about-claude/pricing#claude-platform-on-aws-pricing" + }, + "bedrock/claude_platform/claude-opus-5-5": { + "supports_web_search": true, + "supports_anthropic_compaction": true, + "cache_creation_input_token_cost": 5e-06, + "cache_creation_input_token_cost_above_1hr": 8e-06, + "cache_creation_input_token_cost_batches": 2.5e-06, + "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_batches": 1e-07, + "input_cost_per_token": 4e-06, + "input_cost_per_token_batches": 2e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 2e-05, + "output_cost_per_token_batches": 1e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "thinking_always_on": true, + "supports_mid_conversation_system": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_forced_tool_use": false, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_native_structured_output": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true, + "provider_specific_entry": { + "us": 1.1 + }, + "supports_output_config": true, + "prompt_cache_min_tokens": 512, + "source": "https://platform.claude.com/docs/en/about-claude/pricing#claude-platform-on-aws-pricing" + }, + "bedrock/claude_platform/claude-sonnet-4-5": { + "cache_creation_input_token_cost": 3.75e-06, + "cache_creation_input_token_cost_above_1hr": 6e-06, + "cache_read_input_token_cost": 3e-07, + "input_cost_per_token": 3e-06, + "cache_creation_input_token_cost_batches": 1.875e-06, + "cache_read_input_token_cost_batches": 1.5e-07, + "deprecation_date": "2026-11-30", + "input_cost_per_token_batches": 1.5e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 200000, + "max_output_tokens": 64000, + "max_tokens": 64000, + "mode": "chat", + "output_cost_per_token": 1.5e-05, + "output_cost_per_token_batches": 7.5e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_assistant_prefill": true, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_native_structured_output": true, + "supports_tool_choice": true, + "supports_vision": true, + "prompt_cache_min_tokens": 1024, + "source": "https://platform.claude.com/docs/en/about-claude/pricing#claude-platform-on-aws-pricing" + }, + "bedrock/claude_platform/claude-sonnet-4-6": { + "supports_web_search": true, + "supports_anthropic_compaction": true, + "cache_creation_input_token_cost": 3.75e-06, + "cache_creation_input_token_cost_above_1hr": 6e-06, + "cache_creation_input_token_cost_batches": 1.875e-06, + "cache_read_input_token_cost": 3e-07, + "cache_read_input_token_cost_batches": 1.5e-07, + "input_cost_per_token": 3e-06, + "input_cost_per_token_batches": 1.5e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 1.5e-05, + "output_cost_per_token_batches": 7.5e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "supports_legacy_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_thinking_cache_preservation": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_native_structured_output": true, + "supports_max_reasoning_effort": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_output_config": true, + "prompt_cache_min_tokens": 1024, + "provider_specific_entry": { + "us": 1.1 + }, + "source": "https://platform.claude.com/docs/en/about-claude/pricing#claude-platform-on-aws-pricing" + }, + "bedrock/claude_platform/claude-sonnet-5": { + "supports_web_search": true, + "supports_anthropic_compaction": true, + "cache_creation_input_token_cost": 2.5e-06, + "cache_creation_input_token_cost_above_1hr": 4e-06, + "cache_creation_input_token_cost_batches": 1.25e-06, + "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_batches": 1e-07, + "input_cost_per_token": 2e-06, + "input_cost_per_token_batches": 1e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 1e-05, + "output_cost_per_token_batches": 5e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "supports_mid_conversation_system": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_native_structured_output": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_thinking_cache_preservation": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true, + "provider_specific_entry": { + "us": 1.1 + }, + "supports_output_config": true, + "prompt_cache_min_tokens": 1024, + "source": "https://platform.claude.com/docs/en/about-claude/pricing#claude-platform-on-aws-pricing" + }, + "bedrock/claude_platform/claude-sonnet-5-5": { + "supports_anthropic_compaction": true, + "cache_creation_input_token_cost": 2.5e-06, + "cache_creation_input_token_cost_above_1hr": 4e-06, + "cache_creation_input_token_cost_batches": 1.25e-06, + "cache_read_input_token_cost": 1e-07, + "cache_read_input_token_cost_batches": 5e-08, + "input_cost_per_token": 2e-06, + "input_cost_per_token_batches": 1e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 1e-05, + "output_cost_per_token_batches": 5e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "supports_mid_conversation_system": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_native_structured_output": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true, + "provider_specific_entry": { + "us": 1.1 + }, + "supports_output_config": true, + "supports_forced_tool_use": false, + "thinking_always_on": true, + "prompt_cache_min_tokens": 512, + "source": "https://platform.claude.com/docs/en/about-claude/pricing#claude-platform-on-aws-pricing", + "supports_web_search": true } } diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 4272bfde3f5..5e812eec477 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -5198,6 +5198,7 @@ class ToolDiscoveryQueueItem(TypedDict, total=False): team_id: str | None # team that triggered discovery key_alias: str | None # human-readable key alias user_agent: str | None # HTTP User-Agent of the caller + call_count: ReadOnly[int] # invocations folded into this row since the previous flush from litellm.models.managed_files import ( # noqa: E402 diff --git a/litellm/proxy/analytics_endpoints/cache_activity.py b/litellm/proxy/analytics_endpoints/cache_activity.py index 8a3f71e656e..17e250a2dbe 100644 --- a/litellm/proxy/analytics_endpoints/cache_activity.py +++ b/litellm/proxy/analytics_endpoints/cache_activity.py @@ -2,31 +2,22 @@ import asyncio import json from collections.abc import Sequence from datetime import datetime -from typing import Final, Protocol +from typing import TYPE_CHECKING, Final from pydantic import TypeAdapter from litellm.proxy._types import LiteLLMRoutes from litellm.proxy.auth.authorization import ReadScope, read_scope_bind_params +from litellm.proxy.utils import query_raw_within_statement_timeout from litellm.types.llms.base import LiteLLMBaseModel +if TYPE_CHECKING: + from litellm.proxy.utils import PrismaClient + UNKNOWN_CALL_TYPE: Final = "Unknown" INFO_ROUTES_JSON: Final = json.dumps(LiteLLMRoutes.info_routes.value) -class _SupportsQueryRaw(Protocol): - """The single database operation the cache-activity queries issue.""" - - async def query_raw(self, query: str, *args: object) -> Sequence[object]: ... - - -class _SupportsRawQueryDb(Protocol): - """A prisma client handle, narrowed to the raw-query surface used here.""" - - @property - def db(self) -> _SupportsQueryRaw: ... - - class CacheActivityGroup(LiteLLMBaseModel): call_type: str api_requests: int @@ -174,7 +165,7 @@ def compute_totals(groups: Sequence[CacheActivityGroup]) -> CacheActivityTotals: async def get_cache_activity( - prisma_client: _SupportsRawQueryDb, + prisma_client: "PrismaClient", read_scope: ReadScope, start_date: datetime, end_date: datetime, @@ -185,14 +176,46 @@ async def get_cache_activity( models_json: Final = json.dumps(list(models)) scope_params: Final = read_scope_bind_params(read_scope) group_rows, error_rows, key_alias_rows, model_rows = await asyncio.gather( - prisma_client.db.query_raw( - GROUPS_SQL, start_date, end_date, key_aliases_json, models_json, INFO_ROUTES_JSON, *scope_params + query_raw_within_statement_timeout( + prisma_client, + GROUPS_SQL, + start_date, + end_date, + key_aliases_json, + models_json, + INFO_ROUTES_JSON, + *scope_params, + what="Cache activity", ), - prisma_client.db.query_raw( - ERROR_BREAKDOWN_SQL, start_date, end_date, key_aliases_json, models_json, INFO_ROUTES_JSON, *scope_params + query_raw_within_statement_timeout( + prisma_client, + ERROR_BREAKDOWN_SQL, + start_date, + end_date, + key_aliases_json, + models_json, + INFO_ROUTES_JSON, + *scope_params, + what="Cache activity", + ), + query_raw_within_statement_timeout( + prisma_client, + KEY_ALIAS_OPTIONS_SQL, + start_date, + end_date, + INFO_ROUTES_JSON, + *scope_params, + what="Cache activity", + ), + query_raw_within_statement_timeout( + prisma_client, + MODEL_OPTIONS_SQL, + start_date, + end_date, + INFO_ROUTES_JSON, + *scope_params, + what="Cache activity", ), - prisma_client.db.query_raw(KEY_ALIAS_OPTIONS_SQL, start_date, end_date, INFO_ROUTES_JSON, *scope_params), - prisma_client.db.query_raw(MODEL_OPTIONS_SQL, start_date, end_date, INFO_ROUTES_JSON, *scope_params), ) groups: Final = _groups_adapter.validate_python(group_rows or []) return CacheActivityResponse( diff --git a/litellm/proxy/anthropic_endpoints/endpoints.py b/litellm/proxy/anthropic_endpoints/endpoints.py index b12676e08f3..a763bb3a341 100644 --- a/litellm/proxy/anthropic_endpoints/endpoints.py +++ b/litellm/proxy/anthropic_endpoints/endpoints.py @@ -18,6 +18,7 @@ from litellm.anthropic_interface.exceptions import ( AnthropicExceptionMapping, ) from litellm.integrations.custom_guardrail import ModifyResponseException +from litellm.litellm_core_utils.token_counter import UnsupportedContentTypeError from litellm.litellm_core_utils.tool_fields import first_validation_error from litellm.llms.anthropic.pass_through.context_management import ( AnthropicContextManagementError, @@ -413,6 +414,10 @@ async def count_tokens( messages_format="anthropic", **_count_tokens_body_options(data, auto_drop_params), ) + except UnsupportedContentTypeError as e: + return _anthropic_error_json_response( + ProxyException(message=str(e), type="invalid_request_error", param="messages", code=400), request + ) except HTTPException as e: return _anthropic_error_json_response(proxy_exception_from_http_exception(e, {}), request) except ProxyException as e: diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index d04a84e73a4..c6562dc2273 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -6,6 +6,7 @@ import math from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Callable, Coroutine, Mapping, Sequence from datetime import datetime from functools import lru_cache +from itertools import chain from types import MappingProxyType from typing import ( TYPE_CHECKING, @@ -38,6 +39,7 @@ from litellm.anthropic_interface.exceptions import AnthropicErrorSseFrame, anthr from litellm.constants import ( DD_TRACER_STREAMING_CHUNK_YIELD_RESOURCE, DEFAULT_MAX_RECURSE_DEPTH, + LITELLM_CLIENT_DISCONNECT_DETAIL, LITELLM_DETAILED_TIMING, LITELLM_HTTP_STATUS_CLIENT_DISCONNECTED, MAX_LITELLM_CALL_ID_LENGTH, @@ -68,7 +70,11 @@ from litellm.litellm_core_utils.get_supported_openai_params import ( ) from litellm.litellm_core_utils.internal_call_metadata import is_unbilled_non_inference_call_from_params from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj -from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import guardrail_information_cost +from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import ( + guardrail_entries, + guardrail_information_cost, + object_mapping, +) from litellm.litellm_core_utils.llm_response_utils.get_headers import ( get_response_headers, ) @@ -298,6 +304,23 @@ async def close_guarded_stream(stream: object) -> None: ) +def _guardrail_entries_for_headers( + request_metadata_bucket: object, logging_obj: LiteLLMLoggingObj +) -> tuple[object, ...]: + """Guardrail entries the proxy hooks wrote on the request plus the ones the call recorded on its own logging + object (a Bedrock deployment's native guardrail trace), each once: the hooks copy theirs into the logging object.""" + litellm_params: Final = object_mapping(logging_obj.model_call_details.get("litellm_params")) or {} + buckets: Final = (object_mapping(request_metadata_bucket), object_mapping(litellm_params.get("metadata"))) + entries: Final = tuple( + chain.from_iterable( + guardrail_entries(metadata.get("standard_logging_guardrail_information")) + for metadata in buckets + if metadata is not None + ) + ) + return tuple(entry for index, entry in enumerate(entries) if entry not in entries[:index]) + + def resolve_litellm_call_id(client_call_id: str | None) -> str: if client_call_id is not None and 0 < len(client_call_id) <= MAX_LITELLM_CALL_ID_LENGTH: return client_call_id @@ -1784,7 +1807,7 @@ def _has_attribute_error_in_chain(exc: Exception) -> bool: return False -_CLIENT_DISCONNECT_DETAIL: Final = "Client disconnected the request" +_CLIENT_DISCONNECT_DETAIL: Final = LITELLM_CLIENT_DISCONNECT_DETAIL @runtime_checkable @@ -3186,9 +3209,9 @@ class ProxyBaseLLMRequestProcessing: if is_unbilled_non_inference_call_from_params(logging_obj.call_type, logging_obj.litellm_params, response) else computed_cost_for_headers ) - _, request_metadata_bucket = get_or_create_metadata_bucket(self.data) + request_metadata_bucket: Final[object] = get_or_create_metadata_bucket(self.data)[1] # pyright: ignore[reportUnknownVariableType, reportUnknownMemberType] # self.data is an untyped dict guardrail_cost_for_headers: Final = guardrail_information_cost( - request_metadata_bucket.get("standard_logging_guardrail_information") + _guardrail_entries_for_headers(request_metadata_bucket, logging_obj) # pyright: ignore[reportUnknownArgumentType] # bucket comes from the untyped request dict ) response_cost_for_headers: Final = ( (llm_cost_for_headers if isinstance(llm_cost_for_headers, (int, float)) else 0.0) diff --git a/litellm/proxy/common_utils/http_parsing_utils.py b/litellm/proxy/common_utils/http_parsing_utils.py index f1e409781ef..01795affed2 100644 --- a/litellm/proxy/common_utils/http_parsing_utils.py +++ b/litellm/proxy/common_utils/http_parsing_utils.py @@ -5,16 +5,19 @@ from types import MappingProxyType, UnionType from typing import Annotated, Any, Final, Literal, Union, get_args, get_origin import orjson -from fastapi import Request, UploadFile, status +from fastapi import HTTPException, Request, UploadFile, status from pydantic import TypeAdapter from starlette._utils import get_route_path -from starlette.requests import HTTPConnection +from starlette.datastructures import FormData +from starlette.requests import ClientDisconnect, HTTPConnection from typing_extensions import NotRequired, ReadOnly, Required, assert_never from litellm._logging import verbose_proxy_logger from litellm.constants import ( AZURE_SPEECH_PASS_THROUGH_ROUTE_PREFIX, CLIENT_REQUESTED_MODEL_SCOPE_KEY, + LITELLM_CLIENT_DISCONNECT_DETAIL, + LITELLM_HTTP_STATUS_CLIENT_DISCONNECTED, MAX_REQUEST_BODY_SIZE_TO_REPAIR_MB, ) from litellm.integrations.otel.runtime import phase_event @@ -208,6 +211,29 @@ def _parse_form_metadata(raw_metadata: str) -> object: return parsed +async def _read_form_data(request: Request) -> FormData: + try: + form_data: Final = await request.form() + _mark_body_received(_declared_content_length(request.headers)) + except ClientDisconnect: + raise + except Exception as e: + # ``request.form()`` raises on malformed multipart (missing + # boundary, malformed chunk encoding, …). Surface as 400 so + # the auth-time pre-read does not silently cache ``{}`` while + # a later raw-body re-read sees the original payload — + # banned-param checks must see the same body the handler + # acts on. + verbose_proxy_logger.error("Invalid form payload: %s", e) + raise ProxyException( + message=f"Invalid form payload: {e}", + type="invalid_request_error", + param="request_body", + code=status.HTTP_400_BAD_REQUEST, + ) + return form_data + + async def read_request_body(request: Request | None) -> dict: """ Safely read the request body and parse it as JSON. @@ -218,10 +244,9 @@ async def read_request_body(request: Request | None) -> dict: Returns: - dict: Parsed request data as a dictionary or an empty dictionary if parsing fails """ + if request is None: + return {} try: - if request is None: - return {} - if is_otlp_trace_request(request): return {} @@ -238,24 +263,7 @@ async def read_request_body(request: Request | None) -> dict: _mark_body_received(len(binary_body)) parsed_body = _parse_binary_body(binary_body) elif _is_form_content_type(content_type): - try: - form_data: Final = await request.form() - _mark_body_received(_declared_content_length(request.headers)) - except Exception as e: - # ``request.form()`` raises on malformed multipart (missing - # boundary, malformed chunk encoding, …). Surface as 400 so - # the auth-time pre-read does not silently cache ``{}`` while - # a later raw-body re-read sees the original payload — - # banned-param checks must see the same body the handler - # acts on. - verbose_proxy_logger.error("Invalid form payload: %s", e) - raise ProxyException( - message=f"Invalid form payload: {e}", - type="invalid_request_error", - param="request_body", - code=status.HTTP_400_BAD_REQUEST, - ) - form_body: Final = dict(form_data) + form_body: Final = dict(await _read_form_data(request)) raw_metadata: Final = form_body.get("metadata") parsed_body = ( {**form_body, "metadata": _parse_form_metadata(raw_metadata)} @@ -318,6 +326,17 @@ async def read_request_body(request: Request | None) -> dict: # Re-raise ProxyException as-is verbose_proxy_logger.error("Invalid JSON payload received: %s", e) raise + except ClientDisconnect as e: + from litellm.proxy.proxy_server import ( + _close_dangling_otel_server_span, # pyright: ignore[reportPrivateUsage] # proxy_server keeps the span-close helper private; the Starlette HTTPException handler does not stamp the OTel server span opened at auth time, so the body reader must before it aborts + ) + + verbose_proxy_logger.info("Client disconnected before the request body was fully received") + disconnected: Final = HTTPException( + status_code=LITELLM_HTTP_STATUS_CLIENT_DISCONNECTED, detail=LITELLM_CLIENT_DISCONNECT_DETAIL + ) + _close_dangling_otel_server_span(request, LITELLM_HTTP_STATUS_CLIENT_DISCONNECTED, exc=disconnected) + raise disconnected from e except Exception as e: # Catch unexpected errors to avoid crashes verbose_proxy_logger.exception("Unexpected error reading request body - %s", e) diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 3d8678d70fb..127376b98c4 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -37,7 +37,6 @@ from litellm.constants import ( from litellm.litellm_core_utils.litellm_logging import coerce_model_access_groups from litellm.litellm_core_utils.safe_json_loads import safe_json_loads from litellm.proxy._types import ( - DB_CONNECTION_ERROR_TYPES, DB_RETRY_SAFE_ERROR_TYPES, BaseDailySpendTransaction, DailyAgentSpendTransaction, @@ -74,7 +73,7 @@ from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import ( WindowSpendTransaction, WindowSpendUpdateQueue, ) -from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler +from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler, db_write_failure_is_requeue_safe from litellm.proxy.db.model_usage_rollup import build_model_usage_transaction from litellm.proxy.db.request_error_tracking import request_error_accumulator from litellm.proxy.db.rollup_lock_timeout import apply_rollup_lock_timeout @@ -231,16 +230,6 @@ class _DailySpendCommit(Protocol[_DailySpendTransactionT]): ) -> None: ... -_DATA_REJECTED_SQLSTATE_CLASSES: Final = frozenset({"22", "23"}) - - -def _spend_commit_failure_is_requeue_safe(e: Exception) -> bool: - if isinstance(e, DB_CONNECTION_ERROR_TYPES): - return isinstance(e, DB_RETRY_SAFE_ERROR_TYPES) - sqlstate: Final = PrismaDBExceptionHandler.postgres_sqlstate(e) - return sqlstate is None or sqlstate[:2] not in _DATA_REJECTED_SQLSTATE_CLASSES - - _SpendTableName = Literal[ "user_list_transactions", "end_user_list_transactions", @@ -278,7 +267,7 @@ def _spend_tables_left_to_send( _SPEND_TABLE_COMMIT_ORDER[len(committed)] if len(committed) < len(_SPEND_TABLE_COMMIT_ORDER) else None ) dropped: Final[frozenset[_SpendTableName]] = ( - frozenset() if in_flight is None or _spend_commit_failure_is_requeue_safe(failure) else frozenset({in_flight}) + frozenset() if in_flight is None or db_write_failure_is_requeue_safe(failure) else frozenset({in_flight}) ) if dropped and in_flight is not None: spend_log_error( @@ -1757,7 +1746,7 @@ class DBSpendUpdateWriter: window_spend_transactions=window_spend_update_transactions, ) except Exception as e: - if not _spend_commit_failure_is_requeue_safe(e): + if not db_write_failure_is_requeue_safe(e): uncommitted.pop("window_spend_update_transactions", None) spend_log_error( "Spend tracking - dropped %d budget window increments: the failed statement may have " @@ -1923,7 +1912,7 @@ class DBSpendUpdateWriter: window_spend_transactions=window_spend_update_transactions, ) except Exception as e: # noqa: BLE001 # the increments go back on the queue; the rest of the flush must run - if _spend_commit_failure_is_requeue_safe(e): + if db_write_failure_is_requeue_safe(e): spend_log_error( "Spend tracking - failed to commit budget window spend updates. " "Re-queued %d window increments for retry on next tick. Error: %s", @@ -1941,9 +1930,6 @@ class DBSpendUpdateWriter: exc=e, ) - ################## Tool Registry Upserts ################## - await self._flush_tool_discovery_queue(prisma_client=prisma_client) - async def commit_daily_tag_spend_to_db( self, prisma_client: PrismaClient, @@ -2075,20 +2061,6 @@ class DBSpendUpdateWriter: ) raise - async def _flush_tool_discovery_queue( - self, - prisma_client: PrismaClient, - ) -> None: - """Flush ToolDiscoveryQueue and batch-upsert new tools into LiteLLM_ToolTable.""" - from litellm.proxy.db.tool_registry_writer import batch_upsert_tools - - try: - items: Final = self.tool_discovery_queue.flush() - if items: - await batch_upsert_tools(prisma_client=prisma_client, items=items) - except Exception as e: - verbose_proxy_logger.debug("_flush_tool_discovery_queue error (non-blocking): %s", e) - @staticmethod async def handle_spend_update_failure( e: Exception, @@ -2100,7 +2072,6 @@ class DBSpendUpdateWriter: """Retry a failed spend-update transaction on connection errors, deadlocks or a rollup ``lock_timeout`` (55P03), else re-raise. All three roll the transaction back before any increment applied, so re-sending the same batch cannot double-count.""" - from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler from litellm.proxy.utils import raise_failed_update_spend_exception is_retryable = ( @@ -2609,7 +2580,7 @@ class DBSpendUpdateWriter: _mark_daily_spend_commit_started() _mark_daily_spend_commit_finished() except Exception as batch_error: - if _spend_commit_failure_is_requeue_safe(batch_error): + if db_write_failure_is_requeue_safe(batch_error): spend_log_error( "Daily %s spend batch upsert failed. Table: %s, Rows: %d, Error: %s", entity_type, diff --git a/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py b/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py index 898aee493df..22f55ea9153 100644 --- a/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py +++ b/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py @@ -9,8 +9,11 @@ import json from collections.abc import Mapping, Sequence from datetime import datetime from functools import reduce +from itertools import chain from typing import TYPE_CHECKING, Any, Final, Literal, TypeAlias, TypeVar, cast +from pydantic import TypeAdapter + from redis.exceptions import RedisError from litellm._internal_context import with_service_target @@ -26,6 +29,7 @@ from litellm.constants import ( REDIS_DAILY_TEAM_SPEND_UPDATE_BUFFER_KEY, REDIS_SPEND_LOGS_BUFFER_KEY, REDIS_SPEND_LOGS_BUFFER_MAX_ROWS, + REDIS_TOOL_DISCOVERY_BUFFER_KEY, REDIS_UPDATE_BUFFER_KEY, REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY, ) @@ -41,12 +45,17 @@ from litellm.proxy._types import ( DBSpendUpdateTransactions, Litellm_EntityType, SpendUpdateQueueItem, + ToolDiscoveryQueueItem, ) from litellm.proxy.db.db_transaction_queue.base_update_queue import service_logger_obj from litellm.proxy.db.db_transaction_queue.daily_spend_update_queue import ( DailySpendUpdateQueue, ) from litellm.proxy.db.db_transaction_queue.spend_update_queue import SpendUpdateQueue +from litellm.proxy.db.db_transaction_queue.tool_discovery_queue import ( + ToolDiscoveryQueue, + aggregate_tool_discoveries, +) from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import ( WindowSpendTransaction, WindowSpendUpdateQueue, @@ -133,6 +142,9 @@ def _merged_entity_transactions( ) +_TOOL_DISCOVERY_ITEMS: Final = TypeAdapter(list[ToolDiscoveryQueueItem]) + + class RedisUpdateBuffer: """ Handles buffering database `UPDATE` transactions in Redis before committing them to the database @@ -551,6 +563,56 @@ class RedisUpdateBuffer: str(e), ) + @with_service_target(SPEND_QUEUE_TARGET) + async def store_tool_discoveries_in_redis(self, tool_discovery_queue: ToolDiscoveryQueue) -> None: + """Move this worker's pending tool rows to Redis so only the leader writes LiteLLM_ToolTable.""" + if self.redis_cache is None: + return + items: Final = tool_discovery_queue.flush() + if not items: + return + try: + _ = await self.redis_cache.async_rpush_pipeline( + rpush_list=[ + RedisPipelineRpushOperation(key=REDIS_TOOL_DISCOVERY_BUFFER_KEY, values=[safe_dumps(items)]) + ] + ) + except (RedisError, OSError, asyncio.TimeoutError) as e: + verbose_proxy_logger.error( + "Tool discovery - failed to push %d tool rows to Redis, restoring them for the next tick. Error: %s", + len(items), + str(e), + ) + for item in items: + tool_discovery_queue.add_update(item) + + @with_service_target(SPEND_QUEUE_TARGET) + async def restore_tool_discoveries_to_redis(self, items: Sequence[ToolDiscoveryQueueItem]) -> None: + """Put drained tool rows back for the next drain after a write that provably did not apply.""" + if self.redis_cache is None or not items: + return + _ = await self.redis_cache.async_rpush_pipeline( + rpush_list=[ + RedisPipelineRpushOperation(key=REDIS_TOOL_DISCOVERY_BUFFER_KEY, values=[safe_dumps(list(items))]) + ] + ) + + @with_service_target(SPEND_QUEUE_TARGET) + async def get_all_tool_discoveries_from_redis_buffer(self) -> tuple[ToolDiscoveryQueueItem, ...]: + """Drain every worker's tool rows and fold them to one row per tool_name.""" + if self.redis_cache is None: + return () + popped: Final = await self.redis_cache.async_lpop_pipeline( + lpop_list=[ + RedisPipelineLpopOperation(key=REDIS_TOOL_DISCOVERY_BUFFER_KEY, count=MAX_REDIS_BUFFER_DEQUEUE_COUNT) + ] + ) + entries: Final = popped[0] if popped else None + if not entries: + return () + batches: Final = tuple(_TOOL_DISCOVERY_ITEMS.validate_json(entry) for entry in entries) + return aggregate_tool_discoveries(chain.from_iterable(batches)) + @with_service_target(SPEND_QUEUE_TARGET) async def store_spend_logs_in_redis( self, diff --git a/litellm/proxy/db/db_transaction_queue/tool_discovery_queue.py b/litellm/proxy/db/db_transaction_queue/tool_discovery_queue.py index d5f6a447131..e9ec040023c 100644 --- a/litellm/proxy/db/db_transaction_queue/tool_discovery_queue.py +++ b/litellm/proxy/db/db_transaction_queue/tool_discovery_queue.py @@ -1,52 +1,57 @@ """ -In-memory buffer for tool registry upserts. - -Unlike SpendUpdateQueue (which aggregates increments), ToolDiscoveryQueue -uses set-deduplication: each unique tool_name is only queued once per flush -cycle (~30s). The seen-set is cleared on every flush so that call_count -increments in subsequent cycles rather than stopping after the first flush. +In-memory buffer for tool registry upserts, aggregated by tool_name so one flush carries one row per +tool with the number of times it was invoked since the previous flush. """ +from collections.abc import Iterable from typing import Final +from pydantic import TypeAdapter + from litellm._logging import verbose_proxy_logger from litellm.proxy._types import ToolDiscoveryQueueItem +_ITEM: Final = TypeAdapter(ToolDiscoveryQueueItem) + + +def _merged(existing: ToolDiscoveryQueueItem | None, item: ToolDiscoveryQueueItem) -> ToolDiscoveryQueueItem: + """The first sighting keeps its origin and caller fields, later sightings only add to call_count.""" + increment: Final = item.get("call_count", 1) + if existing is None: + return _ITEM.validate_python({**item, "call_count": increment}) + return _ITEM.validate_python({**existing, "call_count": existing.get("call_count", 1) + increment}) + + +def aggregate_tool_discoveries(items: Iterable[ToolDiscoveryQueueItem]) -> tuple[ToolDiscoveryQueueItem, ...]: + """Collapse items to one per tool_name, summing call_count, in first-seen order.""" + by_name: Final[dict[str, ToolDiscoveryQueueItem]] = {} # mutable-ok: single pass fold keyed by tool_name + for item in items: + if not (tool_name := item.get("tool_name", "")): + continue + by_name[tool_name] = _merged(by_name.get(tool_name), item) + return tuple(by_name.values()) + class ToolDiscoveryQueue: - """ - In-memory buffer for tool registry upserts. - - Deduplicates by tool_name within each flush cycle: a tool is only queued - once per ~30s batch, so call_count increments once per flush cycle the - tool appears in (not once per invocation, but not once per pod lifetime - either). The seen-set is cleared on flush so subsequent batches can - re-count the same tool. - """ + """In-memory buffer for tool registry upserts, one pending row per tool_name with a running call_count.""" def __init__(self) -> None: - self._seen_tool_names: set[str] = set() - self._pending: list[ToolDiscoveryQueueItem] = [] + self._pending: dict[str, ToolDiscoveryQueueItem] = {} # mutable-ok: folded per add_update, cleared on flush def add_update(self, item: ToolDiscoveryQueueItem) -> None: - """Enqueue a tool discovery item if tool_name has not been seen before.""" + """Count one more invocation of item's tool_name, keeping the first sighting's origin and caller fields.""" tool_name: Final = item.get("tool_name", "") if not tool_name: return - if tool_name in self._seen_tool_names: - verbose_proxy_logger.debug("ToolDiscoveryQueue: skipping already-seen tool %s", tool_name) - return - self._seen_tool_names.add(tool_name) - self._pending.append(item) + self._pending[tool_name] = _merged(self._pending.get(tool_name), item) verbose_proxy_logger.debug( - "ToolDiscoveryQueue: queued new tool %s (origin=%s)", + "ToolDiscoveryQueue: tool %s call_count=%s (origin=%s)", tool_name, + self._pending[tool_name].get("call_count"), item.get("origin"), ) def flush(self) -> list[ToolDiscoveryQueueItem]: - """Return and clear all pending items. Resets seen-set so the next - flush cycle can re-count the same tools.""" - items, self._pending = self._pending, [] - self._seen_tool_names.clear() + """Return and clear all pending rows.""" + items, self._pending = list(self._pending.values()), {} return items diff --git a/litellm/proxy/db/exception_handler.py b/litellm/proxy/db/exception_handler.py index 379e0012ea8..669df7454e3 100644 --- a/litellm/proxy/db/exception_handler.py +++ b/litellm/proxy/db/exception_handler.py @@ -8,6 +8,7 @@ from pydantic import TypeAdapter, ValidationError from litellm._logging import verbose_proxy_logger from litellm.proxy._types import ( DB_CONNECTION_ERROR_TYPES, + DB_RETRY_SAFE_ERROR_TYPES, ProxyErrorTypes, ProxyException, ) @@ -33,6 +34,8 @@ _CONNECTION_CAPACITY_PHRASES: Final = ( _PRISMA_POOL_TIMEOUT_CODE: Final = "P2024" _LOCK_TIMEOUT_SQLSTATE: Final = "55P03" _LOCK_TIMEOUT_PHRASE: Final = "canceling statement due to lock timeout" +_STATEMENT_TIMEOUT_SQLSTATE: Final = "57014" +_STATEMENT_TIMEOUT_PHRASE: Final = "canceling statement due to statement timeout" def _exception_chain(e: BaseException) -> Iterator[BaseException]: @@ -283,6 +286,17 @@ class PrismaDBExceptionHandler: return True return _LOCK_TIMEOUT_PHRASE in str(e).lower() + @staticmethod + def is_statement_timeout_error(e: Exception) -> bool: + """True iff Postgres cancelled a statement for exceeding ``statement_timeout`` (SQLSTATE 57014).""" + import prisma + + if not isinstance(e, _exception_types(prisma.errors.PrismaError)): + return False + if PrismaDBExceptionHandler.postgres_sqlstate(e) == _STATEMENT_TIMEOUT_SQLSTATE: + return True + return _STATEMENT_TIMEOUT_PHRASE in str(e).lower() + @staticmethod def postgres_sqlstate(e: Exception) -> str | None: """The SQLSTATE Postgres attached to a failed statement, as prisma surfaces it, or None.""" @@ -637,3 +651,14 @@ async def call_with_db_reconnect_retry( # propagate — repeated reconnect-loops are the watchdog's job, not # this helper's. return await coro_factory() + + +_DATA_REJECTED_SQLSTATE_CLASSES: Final = frozenset({"22", "23"}) + + +def db_write_failure_is_requeue_safe(e: Exception) -> bool: + """Whether a failed write provably never applied, so its rows can be sent again without double counting.""" + if isinstance(e, DB_CONNECTION_ERROR_TYPES): + return isinstance(e, DB_RETRY_SAFE_ERROR_TYPES) + sqlstate: Final = PrismaDBExceptionHandler.postgres_sqlstate(e) + return sqlstate is None or sqlstate[:2] not in _DATA_REJECTED_SQLSTATE_CLASSES diff --git a/litellm/proxy/db/tool_registry_writer.py b/litellm/proxy/db/tool_registry_writer.py index 2e3530d7613..f7ec88945a1 100644 --- a/litellm/proxy/db/tool_registry_writer.py +++ b/litellm/proxy/db/tool_registry_writer.py @@ -14,8 +14,17 @@ from typing import TYPE_CHECKING, Final, Protocol from pydantic import TypeAdapter from litellm._logging import verbose_proxy_logger +from litellm.constants import TOOL_DISCOVERY_COMMIT_JOB_NAME from litellm.proxy._types import ToolDiscoveryQueueItem -from litellm.proxy.db.exception_handler import call_with_db_reconnect_retry +from litellm.proxy.db.db_span import db_span +from litellm.proxy.db.db_transaction_queue.pod_lock_manager import PodLockManager +from litellm.proxy.db.db_transaction_queue.redis_update_buffer import RedisUpdateBuffer +from litellm.proxy.db.db_transaction_queue.tool_discovery_queue import ( + ToolDiscoveryQueue, + aggregate_tool_discoveries, +) +from litellm.proxy.db.exception_handler import call_with_db_reconnect_retry, db_write_failure_is_requeue_safe +from litellm.proxy.db.tool_table_bulk_upsert import build_tool_bulk_upsert from litellm.repositories.object_permission_repository import ObjectPermissionRepository from litellm.repositories.prisma_protocols import TableActions from litellm.repositories.table_repositories import ToolRepository @@ -105,58 +114,75 @@ def _row_to_model(row: object) -> LiteLLM_ToolTableRow: async def batch_upsert_tools( prisma_client: "PrismaClient", - items: list[ToolDiscoveryQueueItem], + items: Sequence[ToolDiscoveryQueueItem], ) -> None: - """ - Batch-upsert tool registry rows via Prisma. + """Write a batch of discovered tools to LiteLLM_ToolTable in a single statement. Raises on failure so the + caller decides whether the rows can be retried.""" + rows: Final = aggregate_tool_discoveries(items) + if not rows: + return + sql, params = build_tool_bulk_upsert(rows) + async with db_span("commit_tool_discoveries", "LiteLLM_ToolTable"): + await prisma_client.db.execute_raw(sql, *params) # pyright: ignore[reportAny] # untyped prisma client + verbose_proxy_logger.debug("tool_registry_writer: upserted %d tool(s)", len(rows)) - On first insert: sets input_policy/output_policy = "untrusted" (default), call_count = 1. - On conflict: increments call_count; preserves existing policies. - """ + +async def flush_tool_discovery_queue(prisma_client: "PrismaClient", tool_discovery_queue: ToolDiscoveryQueue) -> None: + """One worker's own rows straight to LiteLLM_ToolTable, the path without the Redis buffer.""" + items: Final = tool_discovery_queue.flush() if not items: return try: - data: Final = [item for item in items if item.get("tool_name")] - if not data: - return - now: Final = datetime.now(timezone.utc) - table: Final = _tool_table_actions(prisma_client) - for item in data: - tool_name = item.get("tool_name", "") - origin = item.get("origin") or "user_defined" - created_by = item.get("created_by") or "system" - key_hash = item.get("key_hash") - team_id = item.get("team_id") - key_alias = item.get("key_alias") - user_agent = item.get("user_agent") - await table.upsert( - where={"tool_name": tool_name}, - data={ - "create": { - "tool_id": str(uuid.uuid4()), - "tool_name": tool_name, - "origin": origin, - "input_policy": "untrusted", - "output_policy": "untrusted", - "call_count": 1, - "created_by": created_by, - "updated_by": created_by, - "key_hash": key_hash, - "team_id": team_id, - "key_alias": key_alias, - "user_agent": user_agent, - "last_used_at": now, - }, - "update": { - "call_count": {"increment": 1}, - "updated_at": now, - "last_used_at": now, - }, - }, + await batch_upsert_tools(prisma_client=prisma_client, items=items) + except Exception as e: # noqa: BLE001 # a lost tool row only delays its discovery, the scheduler job must not fail + verbose_proxy_logger.debug("flush_tool_discovery_queue error (non-blocking): %s", e) + + +async def commit_tool_discoveries( + prisma_client: "PrismaClient", + tool_discovery_queue: ToolDiscoveryQueue, + redis_update_buffer: RedisUpdateBuffer, + pod_lock_manager: PodLockManager, + use_redis_buffer: bool, +) -> None: + """Scheduled alongside update_spend on its own lock, so tool table writes never ride the spend commit. + + With the Redis buffer every worker hands its rows to Redis and the holder of the tool lock writes them all in + one statement. Without it each worker writes its own rows.""" + if not use_redis_buffer: + await flush_tool_discovery_queue(prisma_client=prisma_client, tool_discovery_queue=tool_discovery_queue) + return + await redis_update_buffer.store_tool_discoveries_in_redis(tool_discovery_queue=tool_discovery_queue) + if not await pod_lock_manager.acquire_lock(cronjob_id=TOOL_DISCOVERY_COMMIT_JOB_NAME): + return + try: + await _commit_tool_discoveries_from_redis(prisma_client=prisma_client, redis_update_buffer=redis_update_buffer) + finally: + await pod_lock_manager.release_lock(cronjob_id=TOOL_DISCOVERY_COMMIT_JOB_NAME) + + +async def _commit_tool_discoveries_from_redis( + prisma_client: "PrismaClient", redis_update_buffer: RedisUpdateBuffer +) -> None: + """Tool lock holder only: drain every worker's tool rows from Redis and write them in one statement.""" + items: Final = await redis_update_buffer.get_all_tool_discoveries_from_redis_buffer() + if not items: + return + try: + await batch_upsert_tools(prisma_client=prisma_client, items=items) + except Exception as e: # noqa: BLE001 # the rows go back to redis or are dropped, the scheduler job must not fail + if db_write_failure_is_requeue_safe(e): + await redis_update_buffer.restore_tool_discoveries_to_redis(items) + verbose_proxy_logger.error( + "Tool discovery - %d tool rows went back to Redis for the next tick. Error: %s", len(items), str(e) ) - verbose_proxy_logger.debug("tool_registry_writer: upserted %d tool(s)", len(data)) - except Exception as e: - verbose_proxy_logger.error("tool_registry_writer batch_upsert_tools error: %s", e) + return + verbose_proxy_logger.error( + "Tool discovery - dropped %d tool rows: the failed statement may have applied or the database " + "refused the data, so re-sending it is not safe. Error: %s", + len(items), + str(e), + ) _NO_OWNERS: Final[Mapping[str, ToolDiscoveryUser]] = MappingProxyType({}) diff --git a/litellm/proxy/db/tool_table_bulk_upsert.py b/litellm/proxy/db/tool_table_bulk_upsert.py new file mode 100644 index 00000000000..e586042ae7c --- /dev/null +++ b/litellm/proxy/db/tool_table_bulk_upsert.py @@ -0,0 +1,71 @@ +"""One multi-row ``INSERT ... ON CONFLICT ("tool_name") DO UPDATE`` per batch of discovered tools. + +LiteLLM_ToolTable holds one row per tool name the proxy has seen. A batch arrives already folded to one +entry per tool name with a summed call_count, and this statement inserts the names the table has not seen +with the untrusted policies the columns default to, and for the names it has adds the batch's call_count +while leaving the policies and the discovering caller's fields alone. Same shape as daily_spend_bulk_upsert. +""" + +import uuid +from collections.abc import Mapping, Sequence +from itertools import chain +from types import MappingProxyType +from typing import Final, TypeAlias + +from litellm.proxy._types import ToolDiscoveryQueueItem + +SqlValue: TypeAlias = str | int | None + +_TOOL_ROW_COLUMNS: Final[tuple[str, ...]] = ( + "tool_id", + "tool_name", + "origin", + "call_count", + "created_by", + "updated_by", + "key_hash", + "team_id", + "key_alias", + "user_agent", +) +_CASTS: Final[Mapping[str, str]] = MappingProxyType({"call_count": "integer"}) + + +def _row_params(item: ToolDiscoveryQueueItem) -> tuple[SqlValue, ...]: + created_by: Final = item.get("created_by") or "system" + return ( + str(uuid.uuid4()), + item.get("tool_name", ""), + item.get("origin") or "user_defined", + item.get("call_count", 1), + created_by, + created_by, + item.get("key_hash"), + item.get("team_id"), + item.get("key_alias"), + item.get("user_agent"), + ) + + +def build_tool_bulk_upsert(items: Sequence[ToolDiscoveryQueueItem]) -> tuple[str, tuple[SqlValue, ...]]: + """The single statement writing one batch of distinct tool_names, plus its positional arguments.""" + width: Final = len(_TOOL_ROW_COLUMNS) + rows: Final = ", ".join( + "(" + + ", ".join( + f"${row_index * width + offset + 1}::{_CASTS.get(column, 'text')}" + for offset, column in enumerate(_TOOL_ROW_COLUMNS) + ) + + ", (NOW() AT TIME ZONE 'UTC'), (NOW() AT TIME ZONE 'UTC'))" + for row_index in range(len(items)) + ) + columns: Final = ", ".join(f'"{column}"' for column in _TOOL_ROW_COLUMNS) + sql: Final = ( + f'INSERT INTO "LiteLLM_ToolTable" ({columns}, "last_used_at", "updated_at")\n' + f"VALUES {rows}\n" + 'ON CONFLICT ("tool_name") DO UPDATE SET\n' + ' "call_count" = "LiteLLM_ToolTable"."call_count" + EXCLUDED."call_count",\n' + ' "last_used_at" = EXCLUDED."last_used_at",\n' + ' "updated_at" = EXCLUDED."updated_at"' + ) + return sql, tuple(chain.from_iterable(map(_row_params, items))) diff --git a/litellm/proxy/management_endpoints/prompt_caching_requests.py b/litellm/proxy/management_endpoints/prompt_caching_requests.py index 2f00d263d84..d52c0255b08 100644 --- a/litellm/proxy/management_endpoints/prompt_caching_requests.py +++ b/litellm/proxy/management_endpoints/prompt_caching_requests.py @@ -14,9 +14,7 @@ from litellm.proxy.spend_tracking.savings import ( marks_gateway_injection, prompt_caching_savings_for_request, ) -from litellm.proxy.spend_tracking.spend_tracking_utils import ( - _query_raw_rows, # pyright: ignore[reportPrivateUsage] # existing typed spend-query adapter; rows validated below -) +from litellm.proxy.utils import query_raw_within_statement_timeout from litellm.types.integrations.anthropic_cache_control_hook import GATEWAY_INJECTED_CACHE_METADATA_KEY from litellm.types.llms.base import LiteLLMBaseModel from litellm.types.management_endpoints.prompt_caching_requests import ( @@ -158,7 +156,7 @@ async def get_prompt_caching_requests( else cursor_start_time ) rows: Final = _REQUEST_ROWS.validate_python( - await _query_raw_rows( + await query_raw_within_statement_timeout( prisma_client, prompt_caching_requests_sql(filter), start.isoformat(), @@ -166,8 +164,8 @@ async def get_prompt_caching_requests( page_size + 1, cursor_time.isoformat() if cursor_time is not None else None, cursor_request_id, + what="Prompt caching requests", ) - or () ) def current_router() -> "Router | None": diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 0f40662ad83..6f7f60f36d4 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -205,7 +205,12 @@ from litellm.proxy.management_helpers.utils import ( add_new_member, management_endpoint_wrapper, ) -from litellm.proxy.utils import PrismaClient, ProxyLogging, handle_exception_on_proxy +from litellm.proxy.utils import ( + PrismaClient, + ProxyLogging, + handle_exception_on_proxy, + query_raw_within_statement_timeout, +) from litellm.repositories.budget_repository import BudgetRepository from litellm.repositories.chunked_in import find_many_in from litellm.repositories.organization_repository import OrganizationRepository @@ -6969,6 +6974,9 @@ class _TeamUserSpendDbRow(TypedDict): failed_requests: ReadOnly[int] +_TEAM_USER_SPEND_ROWS: Final = TypeAdapter(tuple[_TeamUserSpendDbRow, ...]) + + @router.get( "/team/spend/by_user", response_model=TeamUserSpendResponse, @@ -7019,12 +7027,16 @@ async def get_team_spend_by_user( own_user_only: Final = scope.api_key_filter is not None user_param: Final = (user_api_key_dict.user_id or "",) if own_user_only else () async with db_span("team_user_spend", "LiteLLM_SpendLogs"): - rows: Final[Sequence[_TeamUserSpendDbRow]] = await prisma_client.db.query_raw( - _team_user_spend_sql(team_count=len(scoped_team_ids), restrict_to_user=own_user_only), - start_date, - end_date, - *scoped_team_ids, - *user_param, + rows: Final = _TEAM_USER_SPEND_ROWS.validate_python( + await query_raw_within_statement_timeout( + prisma_client, + _team_user_spend_sql(team_count=len(scoped_team_ids), restrict_to_user=own_user_only), + start_date, + end_date, + *scoped_team_ids, + *user_param, + what="Spend per user", + ) ) results: Final = tuple( TeamUserSpendRow( diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 464cfee270a..978882ef3f1 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -4259,10 +4259,16 @@ async def watsonx_proxy_route( if provider_config is None: raise HTTPException(status_code=404, detail="Watsonx passthrough config not found") + watsonx_credentials: Final = passthrough_endpoint_router.get_proxy_wide_deployment_credentials( + custom_llm_provider="watsonx" + ) + api_key: Final = watsonx_credentials.api_key if watsonx_credentials is not None else None + api_base: Final = watsonx_credentials.api_base if watsonx_credentials is not None else None + # Get complete URL with version parameter complete_url, _ = provider_config.get_complete_url( - api_base=None, - api_key=None, + api_base=api_base, + api_key=api_key, model="", endpoint=endpoint, request_query_params=None, @@ -4276,8 +4282,8 @@ async def watsonx_proxy_route( messages=[], optional_params={}, litellm_params={}, - api_key=None, - api_base=None, + api_key=api_key, + api_base=api_base, ) # Check for streaming diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py index a7f3a4f84a7..02c079e433d 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py @@ -5,7 +5,7 @@ from datetime import datetime from typing import TYPE_CHECKING, Any, Final, cast import httpx -from pydantic import ConfigDict, TypeAdapter +from pydantic import ConfigDict, TypeAdapter, ValidationError import litellm from litellm._logging import verbose_proxy_logger @@ -25,6 +25,7 @@ from litellm.llms.anthropic.pass_through.stream_assembly import ( collapse_pure_text_chunks, split_sse_chunk_into_events, ) +from litellm.llms.bedrock.common_utils import BEDROCK_INVOKE_TRACE_FIELD from litellm.proxy._types import PassThroughEndpointLoggingTypedDict from litellm.proxy.auth.auth_utils import get_end_user_id_from_request_body from litellm.proxy.pass_through_endpoints.llm_provider_handlers.batch_attribution import ( @@ -55,6 +56,29 @@ else: _JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) +def _event_payload(event: str) -> Mapping[str, object] | None: + data_lines: Final = tuple( + line.strip()[len("data:") :].strip() for line in event.splitlines() if line.strip().startswith("data:") + ) + if not data_lines or data_lines[0] in ("", "[DONE]"): + return None + try: + return _JSON_OBJECT.validate_json(data_lines[0]) + except ValidationError: + return None + + +def _bedrock_trace_payload(all_chunks: Sequence[str | bytes]) -> Mapping[str, object] | None: + payloads: Final = ( + _event_payload(event) + for chunk in reversed(all_chunks) + for event in reversed(split_sse_chunk_into_events(chunk)) + ) # comprehension-ok: chunk then event, read newest first + return next( + (payload for payload in payloads if payload is not None and BEDROCK_INVOKE_TRACE_FIELD in payload), None + ) + + class AnthropicPassthroughLoggingHandler: @staticmethod def anthropic_passthrough_handler( @@ -106,6 +130,7 @@ class AnthropicPassthroughLoggingHandler: json_mode=False, litellm_params={}, ) + logging_obj.record_native_bedrock_guardrail(response_body) # pyright: ignore[reportUnknownArgumentType] # response_body is an untyped dict param kwargs = AnthropicPassthroughLoggingHandler._create_anthropic_response_logging_payload( litellm_model_response=litellm_model_response, @@ -522,6 +547,7 @@ class AnthropicPassthroughLoggingHandler: "result": None, "kwargs": {}, } + litellm_logging_obj.record_native_bedrock_guardrail(_bedrock_trace_payload(all_chunks)) kwargs: Final = AnthropicPassthroughLoggingHandler._create_anthropic_response_logging_payload( litellm_model_response=complete_streaming_response, model=model, diff --git a/litellm/proxy/pass_through_endpoints/passthrough_endpoint_router.py b/litellm/proxy/pass_through_endpoints/passthrough_endpoint_router.py index 3d1be882779..018c1ce524c 100644 --- a/litellm/proxy/pass_through_endpoints/passthrough_endpoint_router.py +++ b/litellm/proxy/pass_through_endpoints/passthrough_endpoint_router.py @@ -1,5 +1,6 @@ import json from collections.abc import Callable +from dataclasses import dataclass from typing import TYPE_CHECKING, Final from fastapi import HTTPException @@ -57,6 +58,12 @@ def _resolve_oidc_reference(credential: str | None) -> str | None: ) from exc +@dataclass(frozen=True, slots=True) +class _PassThroughCredentials: + api_key: str + api_base: str | None + + class PassthroughEndpointRouter: """ Use this class to Get credentials for pass-through endpoints @@ -89,6 +96,35 @@ class PassthroughEndpointRouter: ) return _resolve_oidc_reference(get_secret_str(_env_variable_name)) + def get_proxy_wide_deployment_credentials(self, custom_llm_provider: str) -> _PassThroughCredentials | None: + llm_router: Final = self.llm_router_getter() + if llm_router is None: + return None + for deployment in llm_router.get_model_list() or (): + litellm_params: LiteLLMParamsTypedDict = deployment["litellm_params"] + api_key: str | None = self._resolve_matching_deployment_api_key( + litellm_params=litellm_params, + custom_llm_provider=custom_llm_provider, + region_name=None, + ) + if api_key is None: + continue + resolved_api_key: str | None = _resolve_oidc_reference(api_key) + if not resolved_api_key: + return None + return _PassThroughCredentials( + api_key=resolved_api_key, + api_base=self._resolve_matching_deployment_api_base(litellm_params), + ) + return None + + def _resolve_matching_deployment_api_base(self, litellm_params: LiteLLMParamsTypedDict) -> str | None: + credential_name: Final = litellm_params.get("litellm_credential_name") + credential_values: Final = ( + CredentialAccessor.get_credential_values(credential_name) if credential_name is not None else None + ) + return _get_str_value(credential_values, "api_base") or litellm_params.get("api_base") + def _get_deployment_api_key( self, custom_llm_provider: str, diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 6dcde1aacae..80970cba90b 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -10888,6 +10888,19 @@ class ProxyStartupEvent: misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME, ) + ### TOOL DISCOVERY (own lock, same interval, never rides the spend commit) ### + from litellm.proxy.utils import commit_tool_discoveries + + scheduler.add_job( # pyright: ignore[reportUnknownMemberType] # apscheduler is untyped + commit_tool_discoveries, + "interval", + seconds=batch_writing_interval, + args=[prisma_client, proxy_logging_obj], + id="commit_tool_discoveries_job", + replace_existing=True, + misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME, + ) + ### UPDATE DAILY TAG SPEND (separate scheduler job with longer interval) ### ## Reduces QPS as there are more tags for a single request tag_spend_update_interval: Final = int(batch_writing_interval * DAILY_TAG_SPEND_BATCH_MULTIPLIER) diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index 31732f2dbc9..2d68cd698e8 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py @@ -65,7 +65,7 @@ from litellm.proxy.spend_tracking.spend_tracking_utils import ( get_spend_by_team, get_spend_by_team_and_customer, ) -from litellm.proxy.utils import handle_exception_on_proxy +from litellm.proxy.utils import handle_exception_on_proxy, query_raw_within_statement_timeout from litellm.repositories.prisma_protocols import TableActions from litellm.repositories.table_repositories import SpendLogsRepository from litellm.repositories.team_repository import TeamRepository @@ -268,6 +268,20 @@ async def _query_raw_or_none(prisma_client: PrismaClient, sql_query: str, *args: return await _query_raw(prisma_client, sql_query, *args) +async def _query_spend_logs(prisma_client: PrismaClient, sql_query: str, *args: object) -> Sequence[_RowT]: + """``_query_raw`` for raw ``LiteLLM_SpendLogs`` reads, which Postgres cancels at the statement timeout.""" + return cast( # cast-ok: callers declare the row shape their SELECT list produces, as with _query_raw + "Sequence[_RowT]", await query_raw_within_statement_timeout(prisma_client, sql_query, *args) + ) + + +async def _query_spend_logs_or_none( + prisma_client: PrismaClient, sql_query: str, *args: object +) -> Sequence[_RowT] | None: + """``_query_spend_logs`` for the call sites that guard the result against ``None``.""" + return await _query_spend_logs(prisma_client, sql_query, *args) + + class _TeamTable(Protocol): """The subset of the Prisma team table API this module uses.""" @@ -1567,7 +1581,9 @@ async def get_global_spend_report( ORDER BY total_cost DESC; """ - db_response = await _query_raw_or_none(prisma_client, sql_query, start_date_obj, end_date_obj, api_key) + db_response = await _query_spend_logs_or_none( + prisma_client, sql_query, start_date_obj, end_date_obj, api_key + ) if db_response is None: return [] @@ -1610,7 +1626,7 @@ async def get_global_spend_report( ORDER BY total_cost DESC; """ - db_response = await _query_raw_or_none( + db_response = await _query_spend_logs_or_none( prisma_client, sql_query, start_date_obj, end_date_obj, internal_user_id ) if db_response is None: @@ -1677,7 +1693,7 @@ async def get_global_spend_report( group_by_day; """ - db_response = await _query_raw_or_none(prisma_client, sql_query, start_date_obj, end_date_obj) + db_response = await _query_spend_logs_or_none(prisma_client, sql_query, start_date_obj, end_date_obj) if db_response is None: return [] @@ -1718,12 +1734,14 @@ async def get_global_spend_report( ORDER BY total_cost DESC; """ - db_response = await _query_raw_or_none(prisma_client, sql_query, start_date_obj, end_date_obj) + db_response = await _query_spend_logs_or_none(prisma_client, sql_query, start_date_obj, end_date_obj) if db_response is None: return [] return db_response + except HTTPException: + raise except Exception as e: raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, @@ -1967,7 +1985,7 @@ async def get_key_spend_report( caller_value=LiteLLMProxyRequestSetup.get_logged_api_key(user_api_key_dict), scope_name="api_key", ) - db_response: Sequence[Mapping[str, object]] | None = await _query_raw_or_none( + db_response: Sequence[Mapping[str, object]] | None = await _query_spend_logs_or_none( prisma_client, _scoped_spend_report_sql(scope_column="api_key"), start_date_obj, @@ -2009,7 +2027,7 @@ async def get_user_spend_report( caller_value=user_api_key_dict.user_id, scope_name="internal_user_id", ) - db_response: Sequence[Mapping[str, object]] | None = await _query_raw_or_none( + db_response: Sequence[Mapping[str, object]] | None = await _query_spend_logs_or_none( prisma_client, _scoped_spend_report_sql(scope_column="user"), start_date_obj, @@ -2051,7 +2069,7 @@ async def get_team_spend_report( caller_value=user_api_key_dict.team_id, scope_name="team_id", ) - db_response: Sequence[Mapping[str, object]] | None = await _query_raw_or_none( + db_response: Sequence[Mapping[str, object]] | None = await _query_spend_logs_or_none( prisma_client, _scoped_spend_report_sql(scope_column="team_id"), start_date_obj, @@ -2093,7 +2111,7 @@ async def get_organization_spend_report( organization_id=organization_id, prisma_client=prisma_client, ) - db_response: Sequence[Mapping[str, object]] | None = await _query_raw_or_none( + db_response: Sequence[Mapping[str, object]] | None = await _query_spend_logs_or_none( prisma_client, _ORG_SPEND_REPORT_SQL, start_date_obj, @@ -3019,7 +3037,7 @@ async def ui_view_spend_logs( LIMIT ${p} ) AS bounded_matches """ - count_rows: Final[Sequence[_SpendLogsCountRow] | None] = await _query_raw_or_none( + count_rows: Final[Sequence[_SpendLogsCountRow] | None] = await _query_spend_logs_or_none( prisma_client, count_query, *sql_params, SPEND_LOGS_PAGINATION_COUNT_CAP + 1 ) raw_total: Final = int(count_rows[0]["total_count"]) if count_rows else 0 @@ -3050,7 +3068,7 @@ async def ui_view_spend_logs( ) sql_params.extend([page_size, skip]) - data: Final = await prisma_client.db.query_raw(sql_query, *sql_params) + data: Final[Sequence[Mapping[str, object]]] = await _query_spend_logs(prisma_client, sql_query, *sql_params) if request_id is not None and not is_v2 and not is_admin_view: await _assert_user_owns_fetched_spend_rows( @@ -3118,7 +3136,7 @@ async def _fetch_session_representatives( ORDER BY {_SESSION_GROUP_KEY_SQL}, {_SESSION_REPRESENTATIVE_ORDER_SQL} ) AS session_representatives """ - rep_rows: Final[Sequence[dict[str, object]]] = await _query_raw( # mutable-ok: rows are enriched in place + rep_rows: Final[Sequence[dict[str, object]]] = await _query_spend_logs( # mutable-ok: rows are enriched in place prisma_client, rep_query, *sql_params, @@ -3148,7 +3166,7 @@ async def _count_grouped_sessions( LIMIT ${next_param_index} ) AS bounded_sessions """ - count_rows: Final[Sequence[_SpendLogsCountRow]] = await _query_raw( + count_rows: Final[Sequence[_SpendLogsCountRow]] = await _query_spend_logs( prisma_client, count_query, *sql_params, SPEND_LOGS_PAGINATION_COUNT_CAP + 1 ) raw_total: Final = int(count_rows[0]["total_count"]) if count_rows else 0 @@ -3222,7 +3240,9 @@ async def _ui_session_grouped_spend_logs( page_rows: Final[Sequence[_SessionPageRow]] = ( () if page_limit <= 0 - else await _query_raw(prisma_client, page_query, *sql_params, *cursor_params, page_limit + 1, *offset_params) + else await _query_spend_logs( + prisma_client, page_query, *sql_params, *cursor_params, page_limit + 1, *offset_params + ) ) has_more: Final = len(page_rows) > page_limit @@ -4707,7 +4727,7 @@ async def _build_ui_spend_logs_response( if (row.get("api_key") if isinstance(row, dict) else getattr(row, "api_key", None)) is not None } ) - rows: Final[Sequence[_SessionSpendRow]] = await _query_raw( + rows: Final[Sequence[_SessionSpendRow]] = await _query_spend_logs( prisma_client, f""" SELECT s.*, COALESCE(m.session_models, ARRAY[]::text[]) AS session_models @@ -4778,7 +4798,7 @@ async def _build_ui_spend_logs_response( if row.get("session_id") and row.get("api_key") is not None for models in (TypeAdapter(list[str]).validate_python(row.get("session_models") or ()),) } - except PrismaError: + except (PrismaError, HTTPException): verbose_proxy_logger.debug( "Failed to enrich session spend aggregates for spend logs UI", exc_info=True, diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index 8149026a00a..97c7a3cab57 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -63,7 +63,7 @@ from litellm.proxy._types import ( ) from litellm.proxy.route_llm_request import ProxyModelNotFoundError from litellm.proxy.spend_tracking.spend_log_error_logger import spend_log_error -from litellm.proxy.utils import PrismaClient, hash_token +from litellm.proxy.utils import PrismaClient, hash_token, query_raw_within_statement_timeout from litellm.types.router import DeploymentTypedDict, LiteLLM_Params from litellm.types.utils import ( BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN, @@ -985,7 +985,8 @@ async def _query_raw_rows( sql_query: str, *args: object, ) -> Sequence[Mapping[str, object]] | None: - return await prisma_client.db.query_raw(sql_query, *args) + """A raw ``LiteLLM_SpendLogs`` read that Postgres cancels at the statement timeout.""" + return await query_raw_within_statement_timeout(prisma_client, sql_query, *args) async def get_spend_by_team( diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index c06bf54fbe7..486028297b3 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -64,6 +64,7 @@ from litellm.constants import ( PROXY_REJECTED_BEFORE_ROUTING_KEY, REDIS_SPEND_LOGS_BUFFER_DEQUEUE_COUNT, SPEND_LOG_QUEUE_MAX_BYTES, + SPEND_LOG_READ_STATEMENT_TIMEOUT_SECONDS, SPEND_LOG_WRITE_BATCH_MAX_BYTES, SPEND_LOG_WRITE_BATCH_MAX_ROWS, ) @@ -180,6 +181,7 @@ from litellm.proxy.auth.route_checks import RouteChecks from litellm.proxy.common_utils.callback_utils import add_guardrail_to_applied_guardrails_header from litellm.proxy.common_utils.config_sync_pubsub import publish_config_param_change from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from litellm.proxy.db import tool_registry_writer from litellm.proxy.db.create_views import ( create_missing_views, create_view_tolerating_race, @@ -187,6 +189,7 @@ from litellm.proxy.db.create_views import ( ) from litellm.proxy.db.db_span import db_span from litellm.proxy.db.db_spend_update_writer import DBSpendUpdateWriter +from litellm.proxy.db.db_transaction_queue.redis_update_buffer import RedisUpdateBuffer from litellm.proxy.db.db_url_settings import ( DatabaseURLSettings, add_missing_query_params, @@ -5035,6 +5038,39 @@ def jsonify_object(data: dict) -> dict: _deprecated_key_cache: Final[LimitedSizeOrderedDict] = LimitedSizeOrderedDict(max_size=1000) _DEPRECATED_KEY_CACHE_TTL_SECONDS: Final = 60 _PRISMA_DEFAULT_TX_TIMEOUT: Final = timedelta(seconds=5) +_SPEND_LOG_READ_STATEMENT_TIMEOUT: Final = timedelta(seconds=SPEND_LOG_READ_STATEMENT_TIMEOUT_SECONDS) +_STATEMENT_TIMEOUT_TRANSACTION_GRACE: Final = timedelta(seconds=5) + + +class _ReadTransactionClient(Protocol): + @property + def read_db(self) -> PrismaWrapper: ... + + +async def query_raw_within_statement_timeout( + prisma_client: _ReadTransactionClient, + sql: str, + *params: object, + what: str = "Spend logs", + timeout: timedelta = _SPEND_LOG_READ_STATEMENT_TIMEOUT, +) -> Sequence[Mapping[str, object]]: + """One raw read on the read target that Postgres cancels at ``timeout``. + + Prisma's HTTP client gives up on a statement after 30s anyway, so without + this the Admin UI's error is the same but Postgres keeps running the scan. + A cancelled read surfaces as a 503 asking for a narrower date range. + """ + try: + async with prisma_client.read_db.tx(timeout=timeout + _STATEMENT_TIMEOUT_TRANSACTION_GRACE) as transaction: + await transaction.execute_raw(f"SET LOCAL statement_timeout = {int(timeout / timedelta(milliseconds=1))}") + return await transaction.query_raw(sql, *params) + except Exception as e: + if not PrismaDBExceptionHandler.is_statement_timeout_error(e): + raise + raise HTTPException( + status_code=status.HTTP_503_SERVICE_UNAVAILABLE, + detail={"error": f"{what} took longer than {int(timeout.total_seconds())}s. Pick a shorter date range"}, + ) from e async def _lookup_deprecated_key( @@ -8289,6 +8325,18 @@ class ProxyUpdateSpend: return False +async def commit_tool_discoveries(prisma_client: PrismaClient, proxy_logging_obj: ProxyLogging) -> None: + """Scheduler job writing discovered tools to LiteLLM_ToolTable, on its own lock, next to update_spend.""" + writer: Final = proxy_logging_obj.db_spend_update_writer + await tool_registry_writer.commit_tool_discoveries( + prisma_client=prisma_client, + tool_discovery_queue=writer.tool_discovery_queue, + redis_update_buffer=writer.redis_update_buffer, + pod_lock_manager=writer.pod_lock_manager, + use_redis_buffer=RedisUpdateBuffer.should_commit_spend_updates_to_redis(), + ) + + async def update_spend( prisma_client: PrismaClient, db_writer_client: AsyncHTTPHandler | None, diff --git a/litellm/responses/litellm_completion_transformation/handler.py b/litellm/responses/litellm_completion_transformation/handler.py index 0ab2052abb0..ff26fc4f509 100644 --- a/litellm/responses/litellm_completion_transformation/handler.py +++ b/litellm/responses/litellm_completion_transformation/handler.py @@ -122,7 +122,6 @@ class LiteLLMCompletionTransformationHandler: litellm_completion_response: Final[ModelResponse | litellm.CustomStreamWrapper] = await litellm.acompletion( **acompletion_args, ) - if isinstance(litellm_completion_response, ModelResponse): responses_api_response: Final[ResponsesAPIResponse] = ( LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( diff --git a/litellm/router.py b/litellm/router.py index 8b705f69e5b..41706bdf539 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -1224,7 +1224,7 @@ class Router: }, ] - router = Router(model_list=model_list, fallbacks=[{"azure-gpt-3.5-turbo": "openai-gpt-3.5-turbo"}]) + router = Router(model_list=model_list, fallbacks=[{"azure-gpt-3.5-turbo": ["openai-gpt-3.5-turbo"]}]) ``` """ @@ -1431,6 +1431,10 @@ class Router: if default_fallbacks is not None or litellm.default_fallbacks is not None: _fallbacks = default_fallbacks or litellm.default_fallbacks + if isinstance(_fallbacks, str): + raise TypeError( + f"default_fallbacks must be a list of model names, got '{_fallbacks}'. Use ['{_fallbacks}']" + ) if self.fallbacks is not None: self.fallbacks.append({"*": _fallbacks}) else: @@ -2564,7 +2568,7 @@ class Router: self._initialize_core_endpoints() self._initialize_specialized_endpoints() - def validate_fallbacks(self, fallback_param: list | None): + def validate_fallbacks(self, fallback_param: Sequence[Mapping[str, object]] | None): """ Validate the fallbacks parameter. """ @@ -2577,6 +2581,13 @@ class Router: raise ValueError( f"Dictionary '{fallback_dict}' must have exactly one key, but has {len(fallback_dict)} keys." ) + ((model_group, targets),) = fallback_dict.items() + if isinstance(targets, list) or check_non_standard_fallback_format(fallbacks=[fallback_dict]): + continue + raise TypeError( + f"Fallbacks for '{model_group}' must be a list of model names, got '{targets}'. " + f"Use {{'{model_group}': ['{targets}']}}" + ) def _add_encrypted_content_affinity_check(self, enable_global_affinity: bool) -> None: def _move_before_deployment_affinity( diff --git a/litellm/router_utils/cooldown_handlers.py b/litellm/router_utils/cooldown_handlers.py index bb64d912c4c..cdca5009b0e 100644 --- a/litellm/router_utils/cooldown_handlers.py +++ b/litellm/router_utils/cooldown_handlers.py @@ -8,12 +8,12 @@ Router cooldown handlers import asyncio import math -from collections.abc import Mapping +from collections.abc import Mapping, Sequence from datetime import datetime from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final -from pydantic import TypeAdapter +from pydantic import TypeAdapter, ValidationError import litellm from litellm._internal_context import service_target @@ -176,7 +176,14 @@ def _should_cooldown_based_on_deployment_policy( override that safety net, matching `_has_explicit_allowed_fails_policy_for_exception`. """ allowed_fails_from_policy: Final = _resolve_allowed_fails_from_policy(dep_policy, original_exception) - if allowed_fails_from_policy is None and dep_allowed_fails is not None and is_single_deployment_model_group: + if ( + allowed_fails_from_policy is None + and is_single_deployment_model_group + and ( + dep_allowed_fails is not None + or litellm_router_instance.get_allowed_fails_from_policy(original_exception) is None + ) + ): return False allowed_fails_override: Final[int | None] = ( @@ -390,6 +397,50 @@ def _should_run_cooldown_logic( return True +_FALLBACK_ENTRIES_ADAPTER: Final = TypeAdapter(tuple[object, ...]) +_FALLBACK_ENTRY_ADAPTER: Final = TypeAdapter(Mapping[str, list[str | Mapping[str, object]] | str]) + + +def _fallback_target_name(target: str | Mapping[str, object]) -> str: + return target if isinstance(target, str) else str(target.get("model", "")) + + +def _fallback_target_names( + targets: Sequence[str | Mapping[str, object]] | str, +) -> list[str]: # mutable-ok: mirrors get_fallback_model_group's fallbacks shape + return [targets] if isinstance(targets, str) else [_fallback_target_name(target) for target in targets] + + +def _fallback_chain_names( + entry: object, +) -> Mapping[str, list[str]] | None: # mutable-ok: mirrors get_fallback_model_group's fallbacks shape + try: + validated: Final = _FALLBACK_ENTRY_ADAPTER.validate_python(entry) + except ValidationError: + return None + return {key: _fallback_target_names(targets) for key, targets in validated.items()} + + +def _fallback_entries(configured: object) -> tuple[object, ...]: + try: + return _FALLBACK_ENTRIES_ADAPTER.validate_python(configured) + except ValidationError: + return () + + +def _has_fallback_group(litellm_router_instance: LitellmRouter, requested_model_group: str | None) -> bool: + from litellm.router_utils.fallback_event_handlers import get_fallback_model_group + + configured: Final[object] = getattr(litellm_router_instance, "fallbacks", None) + if requested_model_group is None or not configured: + return False + chains: Final = tuple( + chain for chain in map(_fallback_chain_names, _fallback_entries(configured)) if chain is not None + ) + fallback_group, _ = get_fallback_model_group(fallbacks=chains, model_group=requested_model_group) + return bool(fallback_group) + + def _should_cooldown_deployment( litellm_router_instance: LitellmRouter, deployment: str, @@ -422,6 +473,9 @@ def _should_cooldown_deployment( litellm_router_instance.routing_group_has_alternatives(requested_model_group) or litellm_router_instance.team_model_has_alternatives(deployment) ) + is_lone_deployment: Final = is_single_deployment_model_group and not _has_fallback_group( + litellm_router_instance, requested_model_group + ) exception_status_int: Final = cast_exception_status_to_int(exception_status) dep_policy, dep_allowed_fails = _get_deployment_cooldown_policy(litellm_router_instance, deployment) @@ -441,7 +495,7 @@ def _should_cooldown_deployment( original_exception, dep_policy, dep_allowed_fails, - is_single_deployment_model_group, + is_lone_deployment, ) ## BASE CASE - single deployment @@ -486,6 +540,8 @@ def _should_cooldown_deployment( return False else: + if is_lone_deployment and litellm_router_instance.get_allowed_fails_from_policy(original_exception) is None: + return False return should_cooldown_based_on_allowed_fails_policy( litellm_router_instance=litellm_router_instance, deployment=deployment, diff --git a/litellm/types/llms/anthropic.py b/litellm/types/llms/anthropic.py index c3d257117ed..f113f833a46 100644 --- a/litellm/types/llms/anthropic.py +++ b/litellm/types/llms/anthropic.py @@ -209,6 +209,13 @@ class AnthropicMessagesTextParam(TypedDict, total=False): prompt_cache_breakpoint: ReadOnly[PromptCacheBreakpoint] +class AnthropicSearchResultParam(TypedDict, total=False): + type: Required[ReadOnly[Literal["search_result"]]] + source: Required[ReadOnly[str]] + title: Required[ReadOnly[str]] + content: Required[ReadOnly[Iterable[AnthropicMessagesTextParam]]] + + class AnthropicMessagesToolUseParam(TypedDict, total=False): type: Required[Literal["tool_use"]] id: str @@ -357,6 +364,7 @@ class AnthropicMessagesToolResultParam(TypedDict, total=False): AnthropicMessagesToolResultContent | AnthropicMessagesImageParam | AnthropicMessagesDocumentParam + | AnthropicSearchResultParam | ToolReference ] ) @@ -368,6 +376,7 @@ AnthropicMessagesUserMessageValues = ( | AnthropicMessagesImageParam | AnthropicMessagesToolResultParam | AnthropicMessagesDocumentParam + | AnthropicSearchResultParam | AnthropicMessagesContainerUploadParam ) diff --git a/litellm/types/llms/custom_llm.py b/litellm/types/llms/custom_llm.py index 6b583b0a83d..bb12449617d 100644 --- a/litellm/types/llms/custom_llm.py +++ b/litellm/types/llms/custom_llm.py @@ -1,11 +1,12 @@ from typing import TYPE_CHECKING -from typing_extensions import TypedDict +from typing_extensions import ReadOnly, TypedDict if TYPE_CHECKING: + from litellm.llms.base_llm.videos.transformation import BaseVideoConfig from litellm.llms.custom_llm import CustomLLM class CustomLLMItem(TypedDict): provider: str - custom_handler: "CustomLLM" + custom_handler: ReadOnly["CustomLLM | BaseVideoConfig"] diff --git a/litellm/types/utils.py b/litellm/types/utils.py index cce1c0cfd70..104ab9d928d 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -321,6 +321,8 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False): cache_creation_input_token_cost_ultrafast: ReadOnly[float | None] # OpenAI ultrafast service tier pricing cache_read_input_token_cost: float | None cache_read_input_audio_token_cost: ReadOnly[float | None] + cache_read_input_audio_token_cost_flex: ReadOnly[float | None] + cache_read_input_audio_token_cost_priority: ReadOnly[float | None] cache_read_input_image_token_cost: ReadOnly[float | None] cache_read_input_token_cost_flex: float | None # OpenAI flex service tier pricing cache_read_input_token_cost_priority: float | None # OpenAI priority service tier pricing @@ -335,6 +337,7 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False): cache_read_input_token_cost_above_272k_tokens_ultrafast: ReadOnly[float | None] cache_read_input_token_cost_above_512k_tokens: float | None cache_read_input_token_cost_batches: ReadOnly[float | None] + cache_read_input_audio_token_cost_batches: ReadOnly[float | None] cache_read_input_token_cost_above_200k_tokens_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] @@ -348,6 +351,12 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False): off_peak_pricing: ReadOnly[OffPeakPricing | None] # time-windowed off-peak rates input_cost_per_character: float | None # only for vertex ai models input_cost_per_audio_token: float | None + input_cost_per_audio_token_flex: ReadOnly[float | None] + input_cost_per_audio_token_above_200k_tokens: ReadOnly[float | None] + input_cost_per_audio_token_above_200k_tokens_flex: ReadOnly[float | None] + input_cost_per_audio_token_above_200k_tokens_priority: ReadOnly[float | None] + input_cost_per_audio_token_above_200k_tokens_batches: ReadOnly[float | None] + input_cost_per_audio_token_priority: ReadOnly[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] @@ -372,10 +381,12 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False): input_cost_per_token_batches: float | None input_cost_per_video_token_batches: ReadOnly[float | None] input_cost_per_token_above_200k_tokens_batches: ReadOnly[float | None] + input_cost_per_token_above_200k_tokens_flex: 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_200k_tokens_batches: ReadOnly[float | None] + output_cost_per_token_above_200k_tokens_flex: ReadOnly[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] @@ -408,6 +419,7 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False): output_cost_per_pixel: ReadOnly[float | None] output_cost_per_image_token: float | None output_cost_per_image_token_batches: ReadOnly[float | None] + output_cost_per_image_token_priority: ReadOnly[float | None] output_cost_per_video_token: float | None # for gemini omni models with video output output_vector_size: int | None output_cost_per_reasoning_token: float | None @@ -3964,14 +3976,23 @@ class CustomPricingLiteLLMParams(MirroredPricingParams): cache_read_input_token_cost_above_100k_tokens_batches: float | None = None cache_read_input_token_cost_above_200k_tokens_batches: float | None = None cache_read_input_token_cost_above_272k_tokens_batches: float | None = None + cache_read_input_audio_token_cost_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 + cache_read_input_audio_token_cost_flex: float | None = None + cache_read_input_audio_token_cost_priority: float | None = None cache_read_input_image_token_cost: float | None = None input_cost_per_character_above_128k_tokens: float | None = None input_cost_per_audio_token: float | None = None + input_cost_per_audio_token_flex: float | None = None + input_cost_per_audio_token_above_200k_tokens: float | None = None + input_cost_per_audio_token_above_200k_tokens_flex: float | None = None + input_cost_per_audio_token_above_200k_tokens_priority: float | None = None + input_cost_per_audio_token_above_200k_tokens_batches: float | None = None + input_cost_per_audio_token_priority: 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 @@ -3981,6 +4002,7 @@ class CustomPricingLiteLLMParams(MirroredPricingParams): input_cost_per_token_above_272k_tokens_flex: float | None = None input_cost_per_token_above_272k_tokens_ultrafast: float | None = None input_cost_per_token_above_200k_tokens_batches: float | None = None + input_cost_per_token_above_200k_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 @@ -4010,12 +4032,14 @@ class CustomPricingLiteLLMParams(MirroredPricingParams): output_cost_per_token_above_272k_tokens_flex: float | None = None output_cost_per_token_above_272k_tokens_ultrafast: float | None = None output_cost_per_token_above_200k_tokens_batches: float | None = None + output_cost_per_token_above_200k_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 output_cost_per_image_token: float | None = None output_cost_per_image_token_batches: float | None = None + output_cost_per_image_token_priority: float | None = None output_cost_per_video_token: float | None = None output_cost_per_reasoning_token: float | None = None output_cost_per_reasoning_token_flex: float | None = None diff --git a/litellm/utils.py b/litellm/utils.py index 5e3c2c886ec..51dc6db456a 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -438,6 +438,7 @@ if TYPE_CHECKING: get_num_retries_from_retry_policy, reset_retry_policy, ) + from litellm.llms.custom_llm import CustomLLM from litellm.types.llms.anthropic import ( ANTHROPIC_API_ONLY_HEADERS, AnthropicThinkingParam, @@ -6542,6 +6543,13 @@ def get_model_info_helper( ), cache_read_input_token_cost=_model_info.get("cache_read_input_token_cost", None), cache_read_input_audio_token_cost=_model_info.get("cache_read_input_audio_token_cost", None), + cache_read_input_audio_token_cost_flex=_model_info.get("cache_read_input_audio_token_cost_flex", None), + cache_read_input_audio_token_cost_priority=_model_info.get( + "cache_read_input_audio_token_cost_priority", None + ), + cache_read_input_audio_token_cost_batches=_model_info.get( + "cache_read_input_audio_token_cost_batches", None + ), prompt_cache_min_tokens=_model_info.get("prompt_cache_min_tokens", None), cache_read_input_token_cost_above_200k_tokens=_model_info.get( "cache_read_input_token_cost_above_200k_tokens", None @@ -6608,6 +6616,20 @@ def get_model_info_helper( input_cost_per_second=_model_info.get("input_cost_per_second", None), input_cost_per_text_message=_model_info.get("input_cost_per_text_message", None), input_cost_per_audio_token=_model_info.get("input_cost_per_audio_token", None), + input_cost_per_audio_token_flex=_model_info.get("input_cost_per_audio_token_flex", None), + input_cost_per_audio_token_above_200k_tokens=_model_info.get( + "input_cost_per_audio_token_above_200k_tokens", None + ), + input_cost_per_audio_token_above_200k_tokens_flex=_model_info.get( + "input_cost_per_audio_token_above_200k_tokens_flex", None + ), + input_cost_per_audio_token_above_200k_tokens_priority=_model_info.get( + "input_cost_per_audio_token_above_200k_tokens_priority", None + ), + input_cost_per_audio_token_above_200k_tokens_batches=_model_info.get( + "input_cost_per_audio_token_above_200k_tokens_batches", None + ), + input_cost_per_audio_token_priority=_model_info.get("input_cost_per_audio_token_priority", None), input_cost_per_image_token=_model_info.get("input_cost_per_image_token", None), input_cost_per_video_token=_model_info.get("input_cost_per_video_token", None), input_cost_per_audio_token_batches=_model_info.get("input_cost_per_audio_token_batches", None), @@ -6620,6 +6642,9 @@ def get_model_info_helper( input_cost_per_token_above_200k_tokens_batches=_model_info.get( "input_cost_per_token_above_200k_tokens_batches" ), + input_cost_per_token_above_200k_tokens_flex=_model_info.get( + "input_cost_per_token_above_200k_tokens_flex", None + ), input_cost_per_token_above_272k_tokens_batches=_model_info.get( "input_cost_per_token_above_272k_tokens_batches" ), @@ -6627,6 +6652,9 @@ def get_model_info_helper( output_cost_per_token_above_200k_tokens_batches=_model_info.get( "output_cost_per_token_above_200k_tokens_batches" ), + output_cost_per_token_above_200k_tokens_flex=_model_info.get( + "output_cost_per_token_above_200k_tokens_flex", None + ), output_cost_per_token_above_272k_tokens_batches=_model_info.get( "output_cost_per_token_above_272k_tokens_batches" ), @@ -6688,6 +6716,7 @@ def get_model_info_helper( output_cost_per_pixel=_model_info.get("output_cost_per_pixel", None), output_cost_per_image_token=_model_info.get("output_cost_per_image_token", None), output_cost_per_image_token_batches=_model_info.get("output_cost_per_image_token_batches", None), + output_cost_per_image_token_priority=_model_info.get("output_cost_per_image_token_priority", None), output_cost_per_video_token=_model_info.get("output_cost_per_video_token", None), output_vector_size=_model_info.get("output_vector_size", None), citation_cost_per_token=_model_info.get("citation_cost_per_token", None), @@ -10140,8 +10169,13 @@ class ProviderConfigManager: @staticmethod def get_provider_video_config( model: str | None, - provider: LlmProviders, + provider: LlmProviders | str, ) -> BaseVideoConfig | None: + if not isinstance(provider, LlmProviders): + custom_video_config: Final = ProviderConfigManager.get_custom_provider_video_config(provider) + if custom_video_config is not None or provider not in LlmProvidersSet: + return custom_video_config + return ProviderConfigManager.get_provider_video_config(model=model, provider=LlmProviders(provider)) if LlmProviders.OPENAI == provider: from litellm.llms.openai.videos.transformation import OpenAIVideoConfig @@ -10188,6 +10222,34 @@ class ProviderConfigManager: return litellm.EdenAIVideoConfig() return None + @staticmethod + def get_custom_llm_handler(provider: str | None) -> CustomLLM | None: + """The `CustomLLM` registered for `provider` in `litellm.custom_provider_map`, if its handler is one.""" + from litellm.llms.custom_llm import CustomLLM + + return next( + ( + item["custom_handler"] + for item in litellm.custom_provider_map + if item["provider"] == provider and isinstance(item["custom_handler"], CustomLLM) + ), + None, + ) + + @staticmethod + def get_custom_provider_video_config(provider: str) -> BaseVideoConfig | None: + """The `BaseVideoConfig` registered for `provider` in `litellm.custom_provider_map`, if its handler is one.""" + from litellm.llms.base_llm.videos.transformation import BaseVideoConfig + + return next( + ( + item["custom_handler"] + for item in litellm.custom_provider_map + if item["provider"] == provider and isinstance(item["custom_handler"], BaseVideoConfig) + ), + None, + ) + @staticmethod def get_provider_container_config( provider: LlmProviders, diff --git a/litellm/videos/main.py b/litellm/videos/main.py index 85e651b687c..66eb10e30db 100644 --- a/litellm/videos/main.py +++ b/litellm/videos/main.py @@ -26,6 +26,17 @@ from litellm.types.videos.utils import decode_video_id_with_provider from litellm.utils import ProviderConfigManager, client from litellm.videos.utils import VideoGenerationRequestUtils + +def _unsupported_video_call( + call_name: str, custom_llm_provider: str, model: str | None = None +) -> litellm.BadRequestError: + return litellm.BadRequestError( + message=f"{call_name} is not supported for {custom_llm_provider}", + model=model or "", + llm_provider=custom_llm_provider, + ) + + #################### Initialize provider clients #################### llm_http_handler: BaseLLMHTTPHandler = BaseLLMHTTPHandler() @@ -40,7 +51,7 @@ async def avideo_generation( size: str | None = None, user: str | None = None, timeout=600, # default to 10 minutes - custom_llm_provider=None, + custom_llm_provider: str | None = None, # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs. # The extra values given here take precedence over values defined on the client or passed to this method. extra_headers: dict[str, object] | None = None, @@ -170,7 +181,7 @@ def video_generation( size: str | None = None, user: str | None = None, timeout=600, # default to 10 minutes - custom_llm_provider=None, + custom_llm_provider: str | None = None, # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs. # The extra values given here take precedence over values defined on the client or passed to this method. extra_headers: dict[str, object] | None = None, @@ -208,11 +219,11 @@ def video_generation( # get provider config video_generation_provider_config: BaseVideoConfig | None = ProviderConfigManager.get_provider_video_config( model=model, - provider=litellm.LlmProviders(custom_llm_provider), + provider=custom_llm_provider, ) if video_generation_provider_config is None: - raise ValueError(f"video generation is not supported for {custom_llm_provider}") + raise _unsupported_video_call("video generation", custom_llm_provider, model) local_vars.update(kwargs) # Get VideoGenerationOptionalRequestParams with only valid parameters @@ -329,11 +340,11 @@ def video_content( # get provider config video_provider_config: Final[BaseVideoConfig | None] = ProviderConfigManager.get_provider_video_config( model=None, - provider=litellm.LlmProviders(custom_llm_provider), + provider=custom_llm_provider, ) if video_provider_config is None: - raise ValueError(f"video support download is not supported for {custom_llm_provider}") + raise _unsupported_video_call("video support download", custom_llm_provider) local_vars.update(kwargs) # For video content download, we don't need complex optional parameter handling @@ -456,7 +467,7 @@ async def avideo_remix( video_id: str, prompt: str, timeout=600, # default to 10 minutes - custom_llm_provider=None, + custom_llm_provider: str | None = None, # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs. # The extra values given here take precedence over values defined on the client or passed to this method. extra_headers: dict[str, object] | None = None, @@ -559,7 +570,7 @@ def video_remix( video_id: str, prompt: str, timeout=600, # default to 10 minutes - custom_llm_provider=None, + custom_llm_provider: str | None = None, # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs. # The extra values given here take precedence over values defined on the client or passed to this method. extra_headers: dict[str, object] | None = None, @@ -598,11 +609,11 @@ def video_remix( # get provider config video_remix_provider_config: Final[BaseVideoConfig | None] = ProviderConfigManager.get_provider_video_config( model=None, - provider=litellm.LlmProviders(custom_llm_provider), + provider=custom_llm_provider, ) if video_remix_provider_config is None: - raise ValueError(f"video remix is not supported for {custom_llm_provider}") + raise _unsupported_video_call("video remix", custom_llm_provider) local_vars.update(kwargs) # For video remix, we need the video_id and prompt @@ -660,7 +671,7 @@ async def avideo_list( order: str | None = None, api_key: str | None = None, timeout=600, # default to 10 minutes - custom_llm_provider=None, + custom_llm_provider: str | None = None, # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs. # The extra values given here take precedence over values defined on the client or passed to this method. extra_headers: dict[str, object] | None = None, @@ -774,7 +785,7 @@ def video_list( limit: int | None = None, order: str | None = None, timeout=600, # default to 10 minutes - custom_llm_provider=None, + custom_llm_provider: str | None = None, # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs. # The extra values given here take precedence over values defined on the client or passed to this method. extra_headers: dict[str, object] | None = None, @@ -810,11 +821,11 @@ def video_list( # get provider config video_list_provider_config: Final[BaseVideoConfig | None] = ProviderConfigManager.get_provider_video_config( model=None, - provider=litellm.LlmProviders(custom_llm_provider), + provider=custom_llm_provider, ) if video_list_provider_config is None: - raise ValueError(f"video list is not supported for {custom_llm_provider}") + raise _unsupported_video_call("video list", custom_llm_provider) local_vars.update(kwargs) # For video list, we need the query parameters @@ -871,7 +882,7 @@ def video_list( async def avideo_status( video_id: str, timeout=600, # default to 10 minutes - custom_llm_provider=None, + custom_llm_provider: str | None = None, # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs. # The extra values given here take precedence over values defined on the client or passed to this method. extra_headers: dict[str, object] | None = None, @@ -970,7 +981,7 @@ def video_status( def video_status( video_id: str, timeout=600, # default to 10 minutes - custom_llm_provider=None, + custom_llm_provider: str | None = None, # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs. # The extra values given here take precedence over values defined on the client or passed to this method. extra_headers: dict[str, object] | None = None, @@ -1031,11 +1042,11 @@ def video_status( # get provider config video_status_provider_config: Final[BaseVideoConfig | None] = ProviderConfigManager.get_provider_video_config( model=None, - provider=litellm.LlmProviders(custom_llm_provider), + provider=custom_llm_provider, ) if video_status_provider_config is None: - raise ValueError(f"video status is not supported for {custom_llm_provider}") + raise _unsupported_video_call("video status", custom_llm_provider) local_vars.update(kwargs) # For video status, we need the video_id @@ -1088,7 +1099,7 @@ async def avideo_create_character( name: str, video: FileTypes, timeout=600, - custom_llm_provider=None, + custom_llm_provider: str | None = None, extra_headers: dict[str, object] | None = None, extra_query: dict[str, object] | None = None, extra_body: dict[str, object] | None = None, @@ -1143,7 +1154,7 @@ def video_create_character( name: str, video: FileTypes, timeout=600, - custom_llm_provider=None, + custom_llm_provider: str | None = None, extra_headers: dict[str, object] | None = None, extra_query: dict[str, object] | None = None, extra_body: dict[str, object] | None = None, @@ -1172,11 +1183,11 @@ def video_create_character( provider_config: Final[BaseVideoConfig | None] = ProviderConfigManager.get_provider_video_config( model=None, - provider=litellm.LlmProviders(custom_llm_provider), + provider=custom_llm_provider, ) if provider_config is None: - raise ValueError(f"video create character is not supported for {custom_llm_provider}") + raise _unsupported_video_call("video create character", custom_llm_provider) local_vars.update(kwargs) request_params: Final[dict] = {"name": name} @@ -1218,7 +1229,7 @@ def video_create_character( async def avideo_get_character( character_id: str, timeout=600, - custom_llm_provider=None, + custom_llm_provider: str | None = None, extra_headers: dict[str, object] | None = None, extra_query: dict[str, object] | None = None, extra_body: dict[str, object] | None = None, @@ -1268,7 +1279,7 @@ async def avideo_get_character( def video_get_character( character_id: str, timeout=600, - custom_llm_provider=None, + custom_llm_provider: str | None = None, extra_headers: dict[str, object] | None = None, extra_query: dict[str, object] | None = None, extra_body: dict[str, object] | None = None, @@ -1297,11 +1308,11 @@ def video_get_character( provider_config: Final[BaseVideoConfig | None] = ProviderConfigManager.get_provider_video_config( model=None, - provider=litellm.LlmProviders(custom_llm_provider), + provider=custom_llm_provider, ) if provider_config is None: - raise ValueError(f"video get character is not supported for {custom_llm_provider}") + raise _unsupported_video_call("video get character", custom_llm_provider) local_vars.update(kwargs) request_params: Final[dict] = {"character_id": character_id} @@ -1343,7 +1354,7 @@ async def avideo_edit( video_id: str, prompt: str, timeout=600, - custom_llm_provider=None, + custom_llm_provider: str | None = None, extra_headers: dict[str, object] | None = None, extra_query: dict[str, object] | None = None, extra_body: dict[str, object] | None = None, @@ -1398,7 +1409,7 @@ def video_edit( video_id: str, prompt: str, timeout=600, - custom_llm_provider=None, + custom_llm_provider: str | None = None, extra_headers: dict[str, object] | None = None, extra_query: dict[str, object] | None = None, extra_body: dict[str, object] | None = None, @@ -1430,11 +1441,11 @@ def video_edit( provider_config: Final[BaseVideoConfig | None] = ProviderConfigManager.get_provider_video_config( model=None, - provider=litellm.LlmProviders(custom_llm_provider), + provider=custom_llm_provider, ) if provider_config is None: - raise ValueError(f"video edit is not supported for {custom_llm_provider}") + raise _unsupported_video_call("video edit", custom_llm_provider) local_vars.update(kwargs) request_params: Final[dict] = {"video_id": video_id, "prompt": prompt} @@ -1480,7 +1491,7 @@ async def avideo_extension( prompt: str, seconds: str, timeout=600, - custom_llm_provider=None, + custom_llm_provider: str | None = None, extra_headers: dict[str, object] | None = None, extra_query: dict[str, object] | None = None, extra_body: dict[str, object] | None = None, @@ -1534,7 +1545,7 @@ def video_extension( prompt: str, seconds: str, timeout=600, - custom_llm_provider=None, + custom_llm_provider: str | None = None, extra_headers: dict[str, object] | None = None, extra_query: dict[str, object] | None = None, extra_body: dict[str, object] | None = None, @@ -1564,11 +1575,11 @@ def video_extension( provider_config: Final[BaseVideoConfig | None] = ProviderConfigManager.get_provider_video_config( model=None, - provider=litellm.LlmProviders(custom_llm_provider), + provider=custom_llm_provider, ) if provider_config is None: - raise ValueError(f"video extension is not supported for {custom_llm_provider}") + raise _unsupported_video_call("video extension", custom_llm_provider) local_vars.update(kwargs) request_params: Final[dict] = { diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 98fee42de3c..2c183527ce8 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -8899,7 +8899,11 @@ "output_cost_per_token_above_272k_tokens": 1.5e-05, "output_cost_per_token_above_272k_tokens_priority": 3e-05, "output_cost_per_token_priority": 2e-05, - "search_context_cost_per_query": {"search_context_size_high":0.01,"search_context_size_low":0.01,"search_context_size_medium":0.01}, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "source": "https://azure.microsoft.com/en-us/blog/gpt-6-astra-sol-and-luna-for-production-agents-in-microsoft-foundry/", "supported_endpoints": [ "/v1/chat/completions", @@ -8952,7 +8956,11 @@ "output_cost_per_token_above_272k_tokens": 1.5e-05, "output_cost_per_token_above_272k_tokens_priority": 3e-05, "output_cost_per_token_priority": 2e-05, - "search_context_cost_per_query": {"search_context_size_high":0.01,"search_context_size_low":0.01,"search_context_size_medium":0.01}, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "source": "https://azure.microsoft.com/en-us/blog/gpt-6-astra-sol-and-luna-for-production-agents-in-microsoft-foundry/", "supported_endpoints": [ "/v1/chat/completions", @@ -9644,7 +9652,11 @@ "output_cost_per_token_above_272k_tokens": 1.65e-05, "output_cost_per_token_above_272k_tokens_priority": 3.3e-05, "output_cost_per_token_priority": 2.2e-05, - "search_context_cost_per_query": {"search_context_size_high":0.01,"search_context_size_low":0.01,"search_context_size_medium":0.01}, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, "source": "https://azure.microsoft.com/en-us/blog/gpt-6-astra-sol-and-luna-for-production-agents-in-microsoft-foundry/", "supported_endpoints": [ "/v1/chat/completions", @@ -16419,7 +16431,10 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "supports_max_reasoning_effort": true, - "provider_specific_entry": {"us":1.1,"fast":2}, + "provider_specific_entry": { + "us": 1.1, + "fast": 2 + }, "supports_output_config": true, "supports_speed": true, "supports_fast_mode": true, @@ -28970,7 +28985,11 @@ "search_context_size_medium": 0.035, "search_context_size_high": 0.035 }, - "search_free_requests": {"count": 1500, "period": "day", "shared_group": "gemini-2-flash"}, + "search_free_requests": { + "count": 1500, + "period": "day", + "shared_group": "gemini-2-flash" + }, "google_maps_grounding_cost_per_query": 0.025, "input_cost_per_audio_token_batches": 5e-07, "input_cost_per_token_batches": 1.5e-07, @@ -28980,7 +28999,13 @@ "output_cost_per_token_batches": 1.25e-06, "output_cost_per_token_flex": 1.25e-06, "output_cost_per_token_priority": 4.5e-06, - "supports_image_size": false + "supports_image_size": false, + "input_cost_per_audio_token_above_200k_tokens_priority": 1.8e-06, + "input_cost_per_audio_token_above_200k_tokens_flex": 5e-07, + "cache_read_input_audio_token_cost_priority": 1.8e-07, + "input_cost_per_audio_token_flex": 5e-07, + "input_cost_per_audio_token_above_200k_tokens": 3e-07, + "input_cost_per_audio_token_above_200k_tokens_batches": 5e-07 }, "gemini-2.5-flash-image": { "deprecation_date": "2027-03-15", @@ -29085,8 +29110,13 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, - "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, - "web_search_billing_unit": "per_query" + "search_free_requests": { + "count": 5000, + "period": "month", + "shared_group": "gemini-3" + }, + "web_search_billing_unit": "per_query", + "output_cost_per_image_token_priority": 0.000216 }, "gemini-3-pro-image-preview": { "input_cost_per_image": 0.0011, @@ -29138,8 +29168,13 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, - "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, - "web_search_billing_unit": "per_query" + "search_free_requests": { + "count": 5000, + "period": "month", + "shared_group": "gemini-3" + }, + "web_search_billing_unit": "per_query", + "output_cost_per_image_token_priority": 0.000216 }, "gemini-3.1-flash-image": { "cache_read_input_token_cost": 5e-08, @@ -29189,8 +29224,13 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, - "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, - "web_search_billing_unit": "per_query" + "search_free_requests": { + "count": 5000, + "period": "month", + "shared_group": "gemini-3" + }, + "web_search_billing_unit": "per_query", + "output_cost_per_image_token_priority": 0.000108 }, "gemini-3.1-flash-image-preview": { "input_cost_per_image": 0.00056, @@ -29233,15 +29273,19 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, - "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, - "web_search_billing_unit": "per_query" + "search_free_requests": { + "count": 5000, + "period": "month", + "shared_group": "gemini-3" + }, + "web_search_billing_unit": "per_query", + "output_cost_per_image_token_priority": 0.000108 }, "gemini-3.1-flash-lite-image": { "cache_read_input_token_cost": 2.5e-08, "cache_read_input_token_cost_batches": 1.25e-08, "cache_read_input_token_cost_flex": 1.25e-08, "cache_read_input_token_cost_priority": 4.5e-08, - "deprecation_date": "2027-06-28", "input_cost_per_image": 0.00028, "input_cost_per_token": 2.5e-07, "input_cost_per_token_batches": 1.25e-07, @@ -29281,7 +29325,8 @@ "supports_response_schema": false, "supports_system_messages": true, "supports_video_input": true, - "supports_vision": true + "supports_vision": true, + "output_cost_per_image_token_priority": 5.4e-05 }, "gemini-3.1-flash-lite-preview": { "cache_read_input_token_cost": 2.5e-08, @@ -29329,7 +29374,11 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, - "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, + "search_free_requests": { + "count": 5000, + "period": "month", + "shared_group": "gemini-3" + }, "web_search_billing_unit": "per_query", "google_maps_grounding_cost_per_query": 0.014 }, @@ -29391,10 +29440,22 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, - "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, + "search_free_requests": { + "count": 5000, + "period": "month", + "shared_group": "gemini-3" + }, "web_search_billing_unit": "per_query", "google_maps_grounding_cost_per_query": 0.014, - "input_cost_per_audio_token_batches": 2.5e-07 + "input_cost_per_audio_token_batches": 2.5e-07, + "input_cost_per_audio_token_above_200k_tokens_priority": 9e-07, + "cache_read_input_audio_token_cost_flex": 2.5e-08, + "input_cost_per_audio_token_above_200k_tokens_flex": 2.5e-07, + "cache_read_input_audio_token_cost_priority": 9e-08, + "input_cost_per_audio_token_flex": 2.5e-07, + "input_cost_per_audio_token_above_200k_tokens": 5e-07, + "cache_read_input_audio_token_cost_batches": 2.5e-08, + "input_cost_per_audio_token_above_200k_tokens_batches": 2.5e-07 }, "gemini-3.5-flash-lite": { "deprecation_date": "2027-07-21", @@ -29452,7 +29513,11 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, - "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, + "search_free_requests": { + "count": 5000, + "period": "month", + "shared_group": "gemini-3" + }, "web_search_billing_unit": "per_query", "google_maps_grounding_cost_per_query": 0.014 }, @@ -29539,7 +29604,11 @@ "search_context_size_medium": 0.035, "search_context_size_high": 0.035 }, - "search_free_requests": {"count": 1500, "period": "day", "shared_group": "gemini-2-flash"}, + "search_free_requests": { + "count": 1500, + "period": "day", + "shared_group": "gemini-2-flash" + }, "google_maps_grounding_cost_per_query": 0.025, "input_cost_per_audio_token_batches": 5e-08, "input_cost_per_token_batches": 5e-08, @@ -29549,7 +29618,13 @@ "output_cost_per_token_batches": 2e-07, "output_cost_per_token_flex": 2e-07, "output_cost_per_token_priority": 7.2e-07, - "supports_image_size": false + "supports_image_size": false, + "input_cost_per_audio_token_above_200k_tokens_priority": 5.4e-07, + "input_cost_per_audio_token_above_200k_tokens_flex": 5e-08, + "cache_read_input_audio_token_cost_priority": 5.4e-08, + "input_cost_per_audio_token_flex": 5e-08, + "input_cost_per_audio_token_above_200k_tokens": 3e-07, + "input_cost_per_audio_token_above_200k_tokens_batches": 5e-08 }, "gemini-2.5-flash-lite-preview-09-2025": { "cache_read_input_token_cost": 1e-08, @@ -29594,7 +29669,11 @@ "search_context_size_medium": 0.035, "search_context_size_high": 0.035 }, - "search_free_requests": {"count": 1500, "period": "day", "shared_group": "gemini-2-flash"}, + "search_free_requests": { + "count": 1500, + "period": "day", + "shared_group": "gemini-2-flash" + }, "google_maps_grounding_cost_per_query": 0.025, "supports_image_size": false }, @@ -29641,7 +29720,11 @@ "search_context_size_medium": 0.035, "search_context_size_high": 0.035 }, - "search_free_requests": {"count": 1500, "period": "day", "shared_group": "gemini-2-flash"}, + "search_free_requests": { + "count": 1500, + "period": "day", + "shared_group": "gemini-2-flash" + }, "google_maps_grounding_cost_per_query": 0.025, "supports_image_size": false }, @@ -29829,7 +29912,11 @@ "search_context_size_medium": 0.035, "search_context_size_high": 0.035 }, - "search_free_requests": {"count": 10000, "period": "day", "shared_group": "gemini-2.5-pro"}, + "search_free_requests": { + "count": 10000, + "period": "day", + "shared_group": "gemini-2.5-pro" + }, "google_maps_grounding_cost_per_query": 0.025, "input_cost_per_token_above_200k_tokens_priority": 4.5e-06, "input_cost_per_token_batches": 6.25e-07, @@ -29838,7 +29925,9 @@ "output_cost_per_token_above_200k_tokens_priority": 2.7e-05, "output_cost_per_token_batches": 5e-06, "output_cost_per_token_flex": 5e-06, - "output_cost_per_token_priority": 1.8e-05 + "output_cost_per_token_priority": 1.8e-05, + "input_cost_per_token_above_200k_tokens_flex": 1.25e-06, + "output_cost_per_token_above_200k_tokens_flex": 7.5e-06 }, "gemini-3.1-pro-preview": { "cache_creation_input_token_cost_above_200k_tokens": 2.5e-07, @@ -29866,8 +29955,16 @@ "output_cost_per_token_flex": 6e-06, "output_cost_per_token_priority": 2.16e-05, "prompt_cache_min_tokens": 4096.0, - "search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014}, - "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, + "search_context_cost_per_query": { + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014, + "search_context_size_high": 0.014 + }, + "search_free_requests": { + "count": 5000, + "period": "month", + "shared_group": "gemini-3" + }, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", @@ -29896,7 +29993,9 @@ "supports_video_input": true, "supports_vision": true, "supports_web_search": true, - "web_search_billing_unit": "per_query" + "web_search_billing_unit": "per_query", + "input_cost_per_token_above_200k_tokens_flex": 2e-06, + "output_cost_per_token_above_200k_tokens_flex": 9e-06 }, "gemini-3.1-pro-preview-customtools": { "prompt_cache_min_tokens": 4096, @@ -29948,7 +30047,11 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, - "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, + "search_free_requests": { + "count": 5000, + "period": "month", + "shared_group": "gemini-3" + }, "web_search_billing_unit": "per_query", "google_maps_grounding_cost_per_query": 0.014 }, @@ -30005,7 +30108,11 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, - "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, + "search_free_requests": { + "count": 5000, + "period": "month", + "shared_group": "gemini-3" + }, "web_search_billing_unit": "per_query" }, "vertex_ai/gemini-3-flash-preview": { @@ -30029,8 +30136,16 @@ "output_cost_per_token_batches": 1.5e-06, "output_cost_per_token_flex": 1.5e-06, "output_cost_per_token_priority": 5.4e-06, - "search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014}, - "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, + "search_context_cost_per_query": { + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014, + "search_context_size_high": 0.014 + }, + "search_free_requests": { + "count": 5000, + "period": "month", + "shared_group": "gemini-3" + }, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", @@ -30058,7 +30173,13 @@ "supports_video_input": true, "supports_vision": true, "supports_web_search": true, - "web_search_billing_unit": "per_query" + "web_search_billing_unit": "per_query", + "input_cost_per_audio_token_above_200k_tokens_priority": 1.8e-06, + "input_cost_per_audio_token_above_200k_tokens_flex": 5e-07, + "cache_read_input_audio_token_cost_priority": 1.8e-07, + "input_cost_per_audio_token_flex": 5e-07, + "input_cost_per_audio_token_above_200k_tokens": 1e-06, + "input_cost_per_audio_token_above_200k_tokens_batches": 5e-07 }, "vertex_ai/gemini-3.5-flash": { "prompt_cache_min_tokens": 4096, @@ -30113,7 +30234,11 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, - "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, + "search_free_requests": { + "count": 5000, + "period": "month", + "shared_group": "gemini-3" + }, "web_search_billing_unit": "per_query", "google_maps_grounding_cost_per_query": 0.014, "input_cost_per_token_batches": 7.5e-07, @@ -30179,7 +30304,11 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, - "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, + "search_free_requests": { + "count": 5000, + "period": "month", + "shared_group": "gemini-3" + }, "web_search_billing_unit": "per_query", "google_maps_grounding_cost_per_query": 0.014 }, @@ -30241,7 +30370,11 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, - "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, + "search_free_requests": { + "count": 5000, + "period": "month", + "shared_group": "gemini-3" + }, "web_search_billing_unit": "per_query", "google_maps_grounding_cost_per_query": 0.014 }, @@ -30267,8 +30400,16 @@ "output_cost_per_token_priority": 6.75e-06, "prompt_cache_min_tokens": 4096.0, "regional_endpoint_uplift_multiplier": 1.1, - "search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014}, - "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, + "search_context_cost_per_query": { + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014, + "search_context_size_high": 0.014 + }, + "search_free_requests": { + "count": 5000, + "period": "month", + "shared_group": "gemini-3" + }, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", @@ -30415,8 +30556,16 @@ "output_cost_per_token_flex": 6e-06, "output_cost_per_token_priority": 2.16e-05, "prompt_cache_min_tokens": 4096.0, - "search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014}, - "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, + "search_context_cost_per_query": { + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014, + "search_context_size_high": 0.014 + }, + "search_free_requests": { + "count": 5000, + "period": "month", + "shared_group": "gemini-3" + }, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", @@ -30445,7 +30594,9 @@ "supports_video_input": true, "supports_vision": true, "supports_web_search": true, - "web_search_billing_unit": "per_query" + "web_search_billing_unit": "per_query", + "input_cost_per_token_above_200k_tokens_flex": 2e-06, + "output_cost_per_token_above_200k_tokens_flex": 9e-06 }, "vertex_ai/gemini-3.1-pro-preview-customtools": { "prompt_cache_min_tokens": 4096, @@ -30503,7 +30654,11 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, - "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, + "search_free_requests": { + "count": 5000, + "period": "month", + "shared_group": "gemini-3" + }, "web_search_billing_unit": "per_query", "google_maps_grounding_cost_per_query": 0.014 }, @@ -30831,7 +30986,11 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, - "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, + "search_free_requests": { + "count": 5000, + "period": "month", + "shared_group": "gemini-3" + }, "web_search_billing_unit": "per_query", "supports_reasoning": false }, @@ -30876,7 +31035,11 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, - "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, + "search_free_requests": { + "count": 5000, + "period": "month", + "shared_group": "gemini-3" + }, "web_search_billing_unit": "per_query" }, "gemini/gemini-3.1-flash-lite-preview": { @@ -30938,7 +31101,11 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, - "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, + "search_free_requests": { + "count": 5000, + "period": "month", + "shared_group": "gemini-3" + }, "web_search_billing_unit": "per_query", "google_maps_grounding_cost_per_query": 0.014, "input_cost_per_audio_token_batches": 2.5e-07 @@ -31041,11 +31208,15 @@ }, "gemini/gemini-2.5-flash": { "cache_read_input_audio_token_cost": 1e-07, + "cache_read_input_audio_token_cost_batches": 1e-07, + "cache_read_input_audio_token_cost_flex": 1e-07, + "cache_read_input_audio_token_cost_priority": 1.8e-07, "cache_read_input_token_cost": 3e-08, "cache_read_input_token_cost_batches": 3e-08, "cache_read_input_token_cost_flex": 3e-08, "cache_read_input_token_cost_priority": 5.4e-08, "input_cost_per_audio_token": 1e-06, + "input_cost_per_audio_token_flex": 5e-07, "input_cost_per_token": 3e-07, "litellm_provider": "gemini", "max_input_tokens": 1048576, @@ -31088,7 +31259,11 @@ "search_context_size_medium": 0.035, "search_context_size_high": 0.035 }, - "search_free_requests": {"count": 1500, "period": "day", "shared_group": "gemini-2-flash"}, + "search_free_requests": { + "count": 1500, + "period": "day", + "shared_group": "gemini-2-flash" + }, "google_maps_grounding_cost_per_query": 0.025, "input_cost_per_audio_token_batches": 5e-07, "input_cost_per_audio_token_priority": 1.8e-06, @@ -31201,7 +31376,11 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, - "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, + "search_free_requests": { + "count": 5000, + "period": "month", + "shared_group": "gemini-3" + }, "web_search_billing_unit": "per_query", "supports_reasoning": true }, @@ -31265,8 +31444,16 @@ "output_cost_per_token": 3e-06, "output_cost_per_token_batches": 1.5e-06, "rpm": 1000.0, - "search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014}, - "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, + "search_context_cost_per_query": { + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014, + "search_context_size_high": 0.014 + }, + "search_free_requests": { + "count": 5000, + "period": "month", + "shared_group": "gemini-3" + }, "source": "https://ai.google.dev/gemini-api/docs/pricing", "supported_endpoints": [ "/v1/chat/completions", @@ -31343,7 +31530,11 @@ "output_cost_per_token": 1.2e-05, "output_cost_per_token_batches": 6e-06, "rpm": 1000.0, - "search_context_cost_per_query": {"search_context_size_low":0.035,"search_context_size_medium":0.035,"search_context_size_high":0.035}, + "search_context_cost_per_query": { + "search_context_size_low": 0.035, + "search_context_size_medium": 0.035, + "search_context_size_high": 0.035 + }, "source": "https://ai.google.dev/gemini-api/docs/models/deep-research-pro-preview-12-2025", "supported_endpoints": [ "/v1/chat/completions", @@ -31370,11 +31561,15 @@ }, "gemini/gemini-2.5-flash-lite": { "cache_read_input_audio_token_cost": 3e-08, + "cache_read_input_audio_token_cost_batches": 3e-08, + "cache_read_input_audio_token_cost_flex": 3e-08, + "cache_read_input_audio_token_cost_priority": 5.4e-08, "cache_read_input_token_cost": 1e-08, "cache_read_input_token_cost_batches": 1e-08, "cache_read_input_token_cost_flex": 1e-08, "cache_read_input_token_cost_priority": 1.8e-08, "input_cost_per_audio_token": 3e-07, + "input_cost_per_audio_token_flex": 1.5e-07, "input_cost_per_token": 1e-07, "litellm_provider": "gemini", "max_input_tokens": 1048576, @@ -31417,7 +31612,11 @@ "search_context_size_medium": 0.035, "search_context_size_high": 0.035 }, - "search_free_requests": {"count": 1500, "period": "day", "shared_group": "gemini-2-flash"}, + "search_free_requests": { + "count": 1500, + "period": "day", + "shared_group": "gemini-2-flash" + }, "google_maps_grounding_cost_per_query": 0.025, "input_cost_per_audio_token_batches": 1.5e-07, "input_cost_per_audio_token_priority": 5.4e-07, @@ -31623,7 +31822,11 @@ "search_context_size_medium": 0.035, "search_context_size_high": 0.035 }, - "search_free_requests": {"count": 10000, "period": "day", "shared_group": "gemini-2.5-pro"}, + "search_free_requests": { + "count": 10000, + "period": "day", + "shared_group": "gemini-2.5-pro" + }, "google_maps_grounding_cost_per_query": 0.025, "input_cost_per_token_batches": 6.25e-07, "input_cost_per_token_flex": 6.25e-07, @@ -31664,12 +31867,16 @@ }, "gemini/gemini-3.1-flash-lite": { "cache_read_input_audio_token_cost": 5e-08, + "cache_read_input_audio_token_cost_batches": 2.5e-08, + "cache_read_input_audio_token_cost_flex": 2.5e-08, + "cache_read_input_audio_token_cost_priority": 9e-08, "cache_read_input_token_cost": 2.5e-08, "cache_read_input_token_cost_batches": 1.25e-08, "cache_read_input_token_cost_flex": 1.25e-08, "cache_read_input_token_cost_priority": 4.5e-08, "deprecation_date": "2027-05-07", "input_cost_per_audio_token": 5e-07, + "input_cost_per_audio_token_flex": 2.5e-07, "input_cost_per_token": 2.5e-07, "input_cost_per_token_batches": 1.25e-07, "input_cost_per_token_flex": 1.25e-07, @@ -31721,7 +31928,11 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, - "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, + "search_free_requests": { + "count": 5000, + "period": "month", + "shared_group": "gemini-3" + }, "web_search_billing_unit": "per_query", "google_maps_grounding_cost_per_query": 0.014, "input_cost_per_audio_token_batches": 2.5e-07, @@ -31784,16 +31995,24 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, - "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, + "search_free_requests": { + "count": 5000, + "period": "month", + "shared_group": "gemini-3" + }, "web_search_billing_unit": "per_query", "google_maps_grounding_cost_per_query": 0.014 }, "gemini/gemini-3-flash-preview": { "cache_read_input_audio_token_cost": 1e-07, + "cache_read_input_audio_token_cost_batches": 1e-07, + "cache_read_input_audio_token_cost_flex": 1e-07, + "cache_read_input_audio_token_cost_priority": 1.8e-07, "cache_read_input_token_cost": 5e-08, "cache_read_input_token_cost_batches": 5e-08, "cache_read_input_token_cost_flex": 5e-08, "input_cost_per_audio_token": 1e-06, + "input_cost_per_audio_token_flex": 5e-07, "input_cost_per_token": 5e-07, "litellm_provider": "gemini", "max_input_tokens": 1048576, @@ -31841,7 +32060,11 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, - "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, + "search_free_requests": { + "count": 5000, + "period": "month", + "shared_group": "gemini-3" + }, "web_search_billing_unit": "per_query", "google_maps_grounding_cost_per_query": 0.014, "input_cost_per_audio_token_batches": 5e-07, @@ -31905,7 +32128,11 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, - "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, + "search_free_requests": { + "count": 5000, + "period": "month", + "shared_group": "gemini-3" + }, "web_search_billing_unit": "per_query", "google_maps_grounding_cost_per_query": 0.014, "input_cost_per_token_batches": 7.5e-07, @@ -31972,7 +32199,11 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, - "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, + "search_free_requests": { + "count": 5000, + "period": "month", + "shared_group": "gemini-3" + }, "web_search_billing_unit": "per_query", "google_maps_grounding_cost_per_query": 0.014 }, @@ -32035,7 +32266,11 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, - "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, + "search_free_requests": { + "count": 5000, + "period": "month", + "shared_group": "gemini-3" + }, "web_search_billing_unit": "per_query", "google_maps_grounding_cost_per_query": 0.014 }, @@ -32098,7 +32333,11 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, - "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, + "search_free_requests": { + "count": 5000, + "period": "month", + "shared_group": "gemini-3" + }, "web_search_billing_unit": "per_query", "google_maps_grounding_cost_per_query": 0.014 }, @@ -32194,7 +32433,11 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, - "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, + "search_free_requests": { + "count": 5000, + "period": "month", + "shared_group": "gemini-3" + }, "web_search_billing_unit": "per_query", "google_maps_grounding_cost_per_query": 0.014, "input_cost_per_token_flex": 1e-06, @@ -32258,7 +32501,11 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, - "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, + "search_free_requests": { + "count": 5000, + "period": "month", + "shared_group": "gemini-3" + }, "web_search_billing_unit": "per_query", "google_maps_grounding_cost_per_query": 0.014, "input_cost_per_token_flex": 1e-06, @@ -32286,8 +32533,16 @@ "output_cost_per_token_batches": 1.5e-06, "output_cost_per_token_flex": 1.5e-06, "output_cost_per_token_priority": 5.4e-06, - "search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014}, - "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, + "search_context_cost_per_query": { + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014, + "search_context_size_high": 0.014 + }, + "search_free_requests": { + "count": 5000, + "period": "month", + "shared_group": "gemini-3" + }, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", @@ -32316,7 +32571,13 @@ "supports_url_context": true, "supports_vision": true, "supports_web_search": true, - "web_search_billing_unit": "per_query" + "web_search_billing_unit": "per_query", + "input_cost_per_audio_token_above_200k_tokens_priority": 1.8e-06, + "input_cost_per_audio_token_above_200k_tokens_flex": 5e-07, + "cache_read_input_audio_token_cost_priority": 1.8e-07, + "input_cost_per_audio_token_flex": 5e-07, + "input_cost_per_audio_token_above_200k_tokens": 1e-06, + "input_cost_per_audio_token_above_200k_tokens_batches": 5e-07 }, "gemini-omni-flash-preview": { "input_cost_per_audio_token": 1.5e-06, @@ -32402,7 +32663,11 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, - "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, + "search_free_requests": { + "count": 5000, + "period": "month", + "shared_group": "gemini-3" + }, "web_search_billing_unit": "per_query", "google_maps_grounding_cost_per_query": 0.014, "input_cost_per_token_batches": 7.5e-07, @@ -32468,7 +32733,11 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, - "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, + "search_free_requests": { + "count": 5000, + "period": "month", + "shared_group": "gemini-3" + }, "web_search_billing_unit": "per_query", "google_maps_grounding_cost_per_query": 0.014 }, @@ -32530,7 +32799,11 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, - "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, + "search_free_requests": { + "count": 5000, + "period": "month", + "shared_group": "gemini-3" + }, "web_search_billing_unit": "per_query", "google_maps_grounding_cost_per_query": 0.014 }, @@ -32555,8 +32828,16 @@ "output_cost_per_token_flex": 1.875e-06, "output_cost_per_token_priority": 6.75e-06, "prompt_cache_min_tokens": 4096.0, - "search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014}, - "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, + "search_context_cost_per_query": { + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014, + "search_context_size_high": 0.014 + }, + "search_free_requests": { + "count": 5000, + "period": "month", + "shared_group": "gemini-3" + }, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", @@ -41522,7 +41803,7 @@ "supports_xhigh_reasoning_effort": true }, "meta/muse-voice-transcribe-1.0": { - "input_cost_per_second": 0.00005, + "input_cost_per_second": 5e-05, "litellm_provider": "meta", "mode": "audio_transcription", "source": "https://dev.meta.ai/docs/speech-to-text", @@ -52342,8 +52623,8 @@ "/v1/images/generations", "/v1/images/edits", "/v1/batch" - ] - + ], + "output_cost_per_image_token_priority": 0.000216 }, "vertex_ai/gemini-3-pro-image-preview": { "input_cost_per_image": 0.0011, @@ -52376,8 +52657,8 @@ "/v1/images/generations", "/v1/images/edits", "/v1/batch" - ] - + ], + "output_cost_per_image_token_priority": 0.000216 }, "vertex_ai/gemini-3.1-flash-image": { "cache_read_input_token_cost": 5e-08, @@ -52409,8 +52690,8 @@ "/v1/images/generations", "/v1/images/edits", "/v1/batch" - ] - + ], + "output_cost_per_image_token_priority": 0.000108 }, "vertex_ai/gemini-3.1-flash-image-preview": { "input_cost_per_image": 0.00056, @@ -52434,15 +52715,14 @@ "/v1/images/generations", "/v1/images/edits", "/v1/batch" - ] - + ], + "output_cost_per_image_token_priority": 0.000108 }, "vertex_ai/gemini-3.1-flash-lite-image": { "cache_read_input_token_cost": 2.5e-08, "cache_read_input_token_cost_batches": 1.25e-08, "cache_read_input_token_cost_flex": 1.25e-08, "cache_read_input_token_cost_priority": 4.5e-08, - "deprecation_date": "2027-06-28", "input_cost_per_image": 0.00028, "input_cost_per_token": 2.5e-07, "input_cost_per_token_batches": 1.25e-07, @@ -52482,7 +52762,8 @@ "supports_response_schema": false, "supports_system_messages": true, "supports_video_input": true, - "supports_vision": true + "supports_vision": true, + "output_cost_per_image_token_priority": 5.4e-05 }, "vertex_ai/gemini-3.1-flash-lite-preview": { "cache_read_input_token_cost": 2.5e-08, @@ -52530,7 +52811,11 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, - "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, + "search_free_requests": { + "count": 5000, + "period": "month", + "shared_group": "gemini-3" + }, "web_search_billing_unit": "per_query", "google_maps_grounding_cost_per_query": 0.014 }, @@ -52593,10 +52878,22 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, - "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, + "search_free_requests": { + "count": 5000, + "period": "month", + "shared_group": "gemini-3" + }, "web_search_billing_unit": "per_query", "google_maps_grounding_cost_per_query": 0.014, - "input_cost_per_audio_token_batches": 2.5e-07 + "input_cost_per_audio_token_batches": 2.5e-07, + "input_cost_per_audio_token_above_200k_tokens_priority": 9e-07, + "cache_read_input_audio_token_cost_flex": 2.5e-08, + "input_cost_per_audio_token_above_200k_tokens_flex": 2.5e-07, + "cache_read_input_audio_token_cost_priority": 9e-08, + "input_cost_per_audio_token_flex": 2.5e-07, + "input_cost_per_audio_token_above_200k_tokens": 5e-07, + "cache_read_input_audio_token_cost_batches": 2.5e-08, + "input_cost_per_audio_token_above_200k_tokens_batches": 2.5e-07 }, "vertex_ai/gemini-3.5-flash-lite": { "deprecation_date": "2027-07-21", @@ -52655,7 +52952,11 @@ "search_context_size_medium": 0.014, "search_context_size_high": 0.014 }, - "search_free_requests": {"count": 5000, "period": "month", "shared_group": "gemini-3"}, + "search_free_requests": { + "count": 5000, + "period": "month", + "shared_group": "gemini-3" + }, "web_search_billing_unit": "per_query", "google_maps_grounding_cost_per_query": 0.014 }, @@ -61612,11 +61913,41 @@ "output_cost_per_second_480p": 0.096075, "output_cost_per_second_720p": 0.216, "output_cost_per_second_1080p": 0.5346, - "supported_modalities": ["text", "image", "video", "audio"], - "supported_output_modalities": ["video"], - "supported_endpoints": ["/v1/videos"], - "provider_specific_entry": {"volcengine_video_frames": {"480p": {"16:9": "854x480", "4:3": "752x560", "1:1": "640x640", "3:4": "560x752", "9:16": "480x854", "21:9": "992x432"}, "720p": {"16:9": "1280x720", "4:3": "1112x834", "1:1": "960x960", "3:4": "834x1112", "9:16": "720x1280", "21:9": "1470x630"}}}, - "metadata": {"comment": "Ark bills per output token, not per second: tokens = width * height * 24 * seconds / 1024, so a per-second rate per resolution tier is exact for the 16:9 geometry Ark publishes (1280x720 is 21600 tokens/s). List 70 CNY per million tokens at 480p and 720p and 77 at 1080p without a video input, matching Ark's per-resolution table and its 1080p 16:9 example of 3.74 CNY/s, converted at 7.0 CNY/USD like the other volcengine entries. Accepts up to 50 reference assets (30 images, 10 videos, 10 audios) and clips of 4-30s; per-asset billing beyond the free allowance is not encoded. Requests with a video input bill at a lower rate per token and are rejected, since their cost is not known at create time. provider_specific_entry.volcengine_video_frames is the frame Ark renders per resolution tier and ratio, verified with ffprobe at 480p for every ratio and at 720p for 16:9 and 4:3; the other 720p ratios are from the Seedance 2.0 model card."} + "supported_modalities": [ + "text", + "image", + "video", + "audio" + ], + "supported_output_modalities": [ + "video" + ], + "supported_endpoints": [ + "/v1/videos" + ], + "provider_specific_entry": { + "volcengine_video_frames": { + "480p": { + "16:9": "854x480", + "4:3": "752x560", + "1:1": "640x640", + "3:4": "560x752", + "9:16": "480x854", + "21:9": "992x432" + }, + "720p": { + "16:9": "1280x720", + "4:3": "1112x834", + "1:1": "960x960", + "3:4": "834x1112", + "9:16": "720x1280", + "21:9": "1470x630" + } + } + }, + "metadata": { + "comment": "Ark bills per output token, not per second: tokens = width * height * 24 * seconds / 1024, so a per-second rate per resolution tier is exact for the 16:9 geometry Ark publishes (1280x720 is 21600 tokens/s). List 70 CNY per million tokens at 480p and 720p and 77 at 1080p without a video input, matching Ark's per-resolution table and its 1080p 16:9 example of 3.74 CNY/s, converted at 7.0 CNY/USD like the other volcengine entries. Accepts up to 50 reference assets (30 images, 10 videos, 10 audios) and clips of 4-30s; per-asset billing beyond the free allowance is not encoded. Requests with a video input bill at a lower rate per token and are rejected, since their cost is not known at create time. provider_specific_entry.volcengine_video_frames is the frame Ark renders per resolution tier and ratio, verified with ffprobe at 480p for every ratio and at 720p for 16:9 and 4:3; the other 720p ratios are from the Seedance 2.0 model card." + } }, "volcengine/doubao-seedance-2-0-260128": { "litellm_provider": "volcengine", @@ -61627,11 +61958,41 @@ "output_cost_per_second_720p": 0.141942857, "output_cost_per_second_1080p": 0.354085714, "output_cost_per_second_4k": 0.722057143, - "supported_modalities": ["text", "image", "video", "audio"], - "supported_output_modalities": ["video"], - "supported_endpoints": ["/v1/videos"], - "provider_specific_entry": {"volcengine_video_frames": {"480p": {"16:9": "864x496", "4:3": "752x560", "1:1": "640x640", "3:4": "560x752", "9:16": "496x864", "21:9": "992x432"}, "720p": {"16:9": "1280x720", "4:3": "1112x834", "1:1": "960x960", "3:4": "834x1112", "9:16": "720x1280", "21:9": "1470x630"}}}, - "metadata": {"comment": "Ark bills per output token, not per second: tokens = width * height * 24 * seconds / 1024, so a per-second rate per resolution tier is exact for the 16:9 geometry Ark publishes (1280x720 is 21600 tokens/s; the 480p frame is 864x496, 10044 tokens/s, not the tier's nominal 854x480). List 46 CNY per million tokens at 480p and 720p, 51 at 1080p, and 26 at 4k without a video input, matching Ark's per-resolution table and its 1080p and 4k 16:9 examples of 2.48 and 5.05 CNY/s, converted at 7.0 CNY/USD like the other volcengine entries. Clips of 4-15s; 4k is exclusive to this tier. Requests with a video input bill at a lower rate per token and are rejected, since their cost is not known at create time. provider_specific_entry.volcengine_video_frames is the frame Ark renders per resolution tier and ratio, verified with ffprobe at 480p for every ratio and at 720p for 16:9 and 4:3; the other 720p ratios are from the Seedance 2.0 model card."} + "supported_modalities": [ + "text", + "image", + "video", + "audio" + ], + "supported_output_modalities": [ + "video" + ], + "supported_endpoints": [ + "/v1/videos" + ], + "provider_specific_entry": { + "volcengine_video_frames": { + "480p": { + "16:9": "864x496", + "4:3": "752x560", + "1:1": "640x640", + "3:4": "560x752", + "9:16": "496x864", + "21:9": "992x432" + }, + "720p": { + "16:9": "1280x720", + "4:3": "1112x834", + "1:1": "960x960", + "3:4": "834x1112", + "9:16": "720x1280", + "21:9": "1470x630" + } + } + }, + "metadata": { + "comment": "Ark bills per output token, not per second: tokens = width * height * 24 * seconds / 1024, so a per-second rate per resolution tier is exact for the 16:9 geometry Ark publishes (1280x720 is 21600 tokens/s; the 480p frame is 864x496, 10044 tokens/s, not the tier's nominal 854x480). List 46 CNY per million tokens at 480p and 720p, 51 at 1080p, and 26 at 4k without a video input, matching Ark's per-resolution table and its 1080p and 4k 16:9 examples of 2.48 and 5.05 CNY/s, converted at 7.0 CNY/USD like the other volcengine entries. Clips of 4-15s; 4k is exclusive to this tier. Requests with a video input bill at a lower rate per token and are rejected, since their cost is not known at create time. provider_specific_entry.volcengine_video_frames is the frame Ark renders per resolution tier and ratio, verified with ffprobe at 480p for every ratio and at 720p for 16:9 and 4:3; the other 720p ratios are from the Seedance 2.0 model card." + } }, "volcengine/doubao-seedance-2-0-fast-260128": { "litellm_provider": "volcengine", @@ -61640,11 +62001,41 @@ "output_cost_per_second": 0.114171429, "output_cost_per_second_480p": 0.053089714, "output_cost_per_second_720p": 0.114171429, - "supported_modalities": ["text", "image", "video", "audio"], - "supported_output_modalities": ["video"], - "supported_endpoints": ["/v1/videos"], - "provider_specific_entry": {"volcengine_video_frames": {"480p": {"16:9": "864x496", "4:3": "752x560", "1:1": "640x640", "3:4": "560x752", "9:16": "496x864", "21:9": "992x432"}, "720p": {"16:9": "1280x720", "4:3": "1112x834", "1:1": "960x960", "3:4": "834x1112", "9:16": "720x1280", "21:9": "1470x630"}}}, - "metadata": {"comment": "Ark bills per output token, not per second: tokens = width * height * 24 * seconds / 1024, so a per-second rate per resolution tier is exact for the 16:9 geometry Ark publishes (1280x720 is 21600 tokens/s; the 480p frame is 864x496, 10044 tokens/s, not the tier's nominal 854x480). List 37 CNY per million tokens without a video input, converted at 7.0 CNY/USD like the other volcengine entries. Clips of 4-15s at 480p or 720p. Requests with a video input bill at a lower rate per token and are rejected, since their cost is not known at create time. provider_specific_entry.volcengine_video_frames is the frame Ark renders per resolution tier and ratio, verified with ffprobe at 480p for every ratio and at 720p for 16:9 and 4:3; the other 720p ratios are from the Seedance 2.0 model card."} + "supported_modalities": [ + "text", + "image", + "video", + "audio" + ], + "supported_output_modalities": [ + "video" + ], + "supported_endpoints": [ + "/v1/videos" + ], + "provider_specific_entry": { + "volcengine_video_frames": { + "480p": { + "16:9": "864x496", + "4:3": "752x560", + "1:1": "640x640", + "3:4": "560x752", + "9:16": "496x864", + "21:9": "992x432" + }, + "720p": { + "16:9": "1280x720", + "4:3": "1112x834", + "1:1": "960x960", + "3:4": "834x1112", + "9:16": "720x1280", + "21:9": "1470x630" + } + } + }, + "metadata": { + "comment": "Ark bills per output token, not per second: tokens = width * height * 24 * seconds / 1024, so a per-second rate per resolution tier is exact for the 16:9 geometry Ark publishes (1280x720 is 21600 tokens/s; the 480p frame is 864x496, 10044 tokens/s, not the tier's nominal 854x480). List 37 CNY per million tokens without a video input, converted at 7.0 CNY/USD like the other volcengine entries. Clips of 4-15s at 480p or 720p. Requests with a video input bill at a lower rate per token and are rejected, since their cost is not known at create time. provider_specific_entry.volcengine_video_frames is the frame Ark renders per resolution tier and ratio, verified with ffprobe at 480p for every ratio and at 720p for 16:9 and 4:3; the other 720p ratios are from the Seedance 2.0 model card." + } }, "volcengine/doubao-seedance-2-0-mini-260615": { "litellm_provider": "volcengine", @@ -61653,11 +62044,39 @@ "output_cost_per_second": 0.070971429, "output_cost_per_second_480p": 0.033001714, "output_cost_per_second_720p": 0.070971429, - "supported_modalities": ["text", "image"], - "supported_output_modalities": ["video"], - "supported_endpoints": ["/v1/videos"], - "provider_specific_entry": {"volcengine_video_frames": {"480p": {"16:9": "864x496", "4:3": "752x560", "1:1": "640x640", "3:4": "560x752", "9:16": "496x864", "21:9": "992x432"}, "720p": {"16:9": "1280x720", "4:3": "1112x834", "1:1": "960x960", "3:4": "834x1112", "9:16": "720x1280", "21:9": "1470x630"}}}, - "metadata": {"comment": "Ark bills per output token, not per second: tokens = width * height * 24 * seconds / 1024, so a per-second rate per resolution tier is exact for the 16:9 geometry Ark publishes (1280x720 is 21600 tokens/s; the 480p frame is 864x496, 10044 tokens/s, not the tier's nominal 854x480). List 23 CNY per million tokens without a video input, converted at 7.0 CNY/USD like the other volcengine entries. Clips of 4-15s at 480p or 720p. Only text and image inputs are declared; the mini tier's reference video and audio caps were not verified. Requests with a video input bill at a lower rate per token and are rejected, since their cost is not known at create time. provider_specific_entry.volcengine_video_frames is the frame Ark renders per resolution tier and ratio, verified with ffprobe at 480p for every ratio and at 720p for 16:9 and 4:3; the other 720p ratios are from the Seedance 2.0 model card."} + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "video" + ], + "supported_endpoints": [ + "/v1/videos" + ], + "provider_specific_entry": { + "volcengine_video_frames": { + "480p": { + "16:9": "864x496", + "4:3": "752x560", + "1:1": "640x640", + "3:4": "560x752", + "9:16": "496x864", + "21:9": "992x432" + }, + "720p": { + "16:9": "1280x720", + "4:3": "1112x834", + "1:1": "960x960", + "3:4": "834x1112", + "9:16": "720x1280", + "21:9": "1470x630" + } + } + }, + "metadata": { + "comment": "Ark bills per output token, not per second: tokens = width * height * 24 * seconds / 1024, so a per-second rate per resolution tier is exact for the 16:9 geometry Ark publishes (1280x720 is 21600 tokens/s; the 480p frame is 864x496, 10044 tokens/s, not the tier's nominal 854x480). List 23 CNY per million tokens without a video input, converted at 7.0 CNY/USD like the other volcengine entries. Clips of 4-15s at 480p or 720p. Only text and image inputs are declared; the mini tier's reference video and audio caps were not verified. Requests with a video input bill at a lower rate per token and are rejected, since their cost is not known at create time. provider_specific_entry.volcengine_video_frames is the frame Ark renders per resolution tier and ratio, verified with ffprobe at 480p for every ratio and at 720p for 16:9 and 4:3; the other 720p ratios are from the Seedance 2.0 model card." + } }, "bedrock/us-east-1/zai.glm-5": { "input_cost_per_token": 1e-06, @@ -63518,7 +63937,13 @@ { "name": "claude-adaptive-thinking", "pattern": "claude-[a-z]+-(?:4[-._](?:[6-9]|[1-9]\\d)(?!\\d)|[5-9](?!\\d)(?:[-._]\\d{1,2}(?!\\d))?)", - "fill_missing_for_providers": ["anthropic", "azure_ai", "bedrock", "bedrock_converse", "vertex_ai-anthropic_models"], + "fill_missing_for_providers": [ + "anthropic", + "azure_ai", + "bedrock", + "bedrock_converse", + "vertex_ai-anthropic_models" + ], "description": "Claude at version 4.6 or higher, in any id shape that contains claude--: minors 4.6 through 4.99, any later major-minor, and bare 5+ majors so a new family shaped like claude-fable-5 matches. Requiring the claude- prefix keeps non-Claude names such as team-sonnet-5-1 out. The minor is capped at two digits so an 8-digit date suffix such as claude-opus-4-20250514 is never read as a >= 4.6 minor. Two-digit majors are deliberately not matched so ids like claude-opus-41 (4.1) are not read as major 41. Turns on adaptive thinking for new versions and new families with no code change.", "model_info": { "supports_adaptive_thinking": true @@ -63527,7 +63952,13 @@ { "name": "claude-legacy-thinking", "pattern": "claude-[a-z]+-4[-._]6(?!\\d)", - "fill_missing_for_providers": ["anthropic", "azure_ai", "bedrock", "bedrock_converse", "vertex_ai-anthropic_models"], + "fill_missing_for_providers": [ + "anthropic", + "azure_ai", + "bedrock", + "bedrock_converse", + "vertex_ai-anthropic_models" + ], "description": "Claude at version 4.6 exactly, in any id shape that contains claude--4-6 (dotted and underscored minors included, dated releases such as claude-sonnet-4-6-20260219 too). The 4.6 family is adaptive-thinking yet still accepts legacy thinking.type=enabled with budget_tokens, so the caller's hard budget cap is forwarded verbatim instead of being rewritten to an uncapped output_config.effort. The lookahead keeps two-digit minors such as 4-60 from matching. 4.7+ and 5+ majors reject the legacy shape and stay on the adaptive translation.", "model_info": { "supports_legacy_thinking": true @@ -63544,7 +63975,13 @@ { "name": "claude-mid-conversation-system", "pattern": "claude-[a-z]+-(?:4[-._](?:[89]|[1-9]\\d)(?!\\d)|[5-9](?!\\d)(?:[-._]\\d{1,2}(?!\\d))?)", - "fill_missing_for_providers": ["anthropic", "azure_ai", "bedrock", "bedrock_converse", "vertex_ai-anthropic_models"], + "fill_missing_for_providers": [ + "anthropic", + "azure_ai", + "bedrock", + "bedrock_converse", + "vertex_ai-anthropic_models" + ], "description": "Claude at version 4.8 or higher, in any id shape that contains claude--: minors 4.8 through 4.99, any later major-minor, and bare 5+ majors so a new family like claude-fable-5 matches. Two-digit majors are deliberately not matched so ids like claude-opus-41 (4.1) are not read as major 41. Anthropic introduced mid-conversation system messages with Opus 4.8 and every newer Claude keeps them; 4.7 and below reject the system role inside messages.", "model_info": { "supports_mid_conversation_system": true @@ -63553,7 +63990,13 @@ { "name": "claude-tool-search", "pattern": "claude-[a-z]+-(?:4[-._](?:[5-9]|[1-9]\\d)(?!\\d)|[5-9](?!\\d)(?:[-._]\\d{1,2}(?!\\d))?)", - "fill_missing_for_providers": ["anthropic", "azure_ai", "bedrock", "bedrock_converse", "vertex_ai-anthropic_models"], + "fill_missing_for_providers": [ + "anthropic", + "azure_ai", + "bedrock", + "bedrock_converse", + "vertex_ai-anthropic_models" + ], "description": "Claude at version 4.5 or higher, in any id shape that contains claude--: minors 4.5 through 4.99, any later major-minor, and bare 5+ majors so a new family like claude-fable-5 matches. Two-digit majors are deliberately not matched so ids like claude-opus-41 (4.1) are not read as major 41. Anthropic's tool search docs list every Claude 4.5 and newer model as supported and Opus 4.1 and earlier as unsupported, so the flag follows the version instead of a per-model list. azure_ai is included because Microsoft Foundry runs tool search for Claude on both Anthropic-hosted and Azure-hosted deployments.", "model_info": { "supports_tool_search": true @@ -63570,7 +64013,11 @@ { "name": "openai-reasoning-family-baseline", "pattern": "^(?!.*search-api)(?:[a-z0-9_.-]+/)*(?:ft:)?(?:o[1-9]\\d*(?![a-z0-9])|gpt-[5-9](?:\\.\\d+)?(?![0-9.])|(?:gpt-\\d+(?:\\.\\d+)?(?:-[a-z0-9]+)*-)?(?:codex|deep-research|chat-latest)(?![a-z0-9]))", - "fill_missing_for_providers": ["azure", "azure_ai", "openai"], + "fill_missing_for_providers": [ + "azure", + "azure_ai", + "openai" + ], "description": "OpenAI reasoning families by id shape, under any provider namespace and with an optional ft: prefix: the o-series (o1, o3-pro, o4-mini), gpt-5 through gpt-9 majors including dotted minors and suffixed variants (gpt-5.5-cyber, gpt-6-astra), and the codex, deep-research and chat-latest lines when standalone or on a gpt base. gpt-5-search-api is excluded because it is a search-only surface. Every model here is a reasoning model, and the Responses API drops the caller's reasoning param for any mapped OpenAI model whose info lacks supports_reasoning, so an id the registry has not named yet keeps its reasoning settings instead of silently losing them. Rules lose to exact entries. Carries no mode and no pricing, so cost stays on the standard unpriced behavior.", "model_info": { "supports_reasoning": true @@ -73263,7 +73710,9 @@ "mode": "chat", "output_cost_per_token": 4.4e-06, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", - "supported_regions": ["global"], + "supported_regions": [ + "global" + ], "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, @@ -74604,7 +75053,9 @@ }, "global.xai.grok-4.7": { "supports_regex_lookaround": false, - "supported_endpoints": ["/v1/chat/completions"], + "supported_endpoints": [ + "/v1/chat/completions" + ], "supports_bedrock_runtime_chat_completions_tools_with_reasoning": true, "supports_bedrock_runtime_chat_completions_response_format": true, "cache_read_input_token_cost": 5e-07, @@ -74630,7 +75081,9 @@ }, "us.xai.grok-4.7": { "supports_regex_lookaround": false, - "supported_endpoints": ["/v1/chat/completions"], + "supported_endpoints": [ + "/v1/chat/completions" + ], "supports_bedrock_runtime_chat_completions_tools_with_reasoning": true, "supports_bedrock_runtime_chat_completions_response_format": true, "cache_read_input_token_cost": 5.5e-07, @@ -75079,7 +75532,11 @@ "output_cost_per_token": 7.5e-06, "output_cost_per_token_batches": 3.75e-06, "rpm": 1000.0, - "search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014}, + "search_context_cost_per_query": { + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014, + "search_context_size_high": 0.014 + }, "source": "https://ai.google.dev/gemini-api/docs/pricing", "supported_endpoints": [ "/v1/chat/completions", @@ -75124,7 +75581,11 @@ "output_cost_per_token": 7.5e-06, "output_cost_per_token_batches": 3.75e-06, "output_cost_per_token_priority": 1.35e-05, - "search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014}, + "search_context_cost_per_query": { + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014, + "search_context_size_high": 0.014 + }, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", @@ -75150,7 +75611,8 @@ "supports_video_input": true, "supports_vision": true, "supports_web_search": true, - "web_search_billing_unit": "per_query" + "web_search_billing_unit": "per_query", + "output_cost_per_image_token_priority": 5.4e-05 }, "vertex_ai/gemini-nano-banana-2.1": { "cache_read_input_token_cost": 1.5e-07, @@ -75170,7 +75632,11 @@ "output_cost_per_token": 7.5e-06, "output_cost_per_token_batches": 3.75e-06, "output_cost_per_token_priority": 1.35e-05, - "search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014}, + "search_context_cost_per_query": { + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014, + "search_context_size_high": 0.014 + }, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", @@ -75196,7 +75662,8 @@ "supports_video_input": true, "supports_vision": true, "supports_web_search": true, - "web_search_billing_unit": "per_query" + "web_search_billing_unit": "per_query", + "output_cost_per_image_token_priority": 5.4e-05 }, "azure_ai/grok-4.7": { "cache_read_input_token_cost": 5e-07, @@ -75252,7 +75719,7 @@ "supports_tool_choice": true, "supports_vision": true }, -"claude-haiku-5-5": { + "claude-haiku-5-5": { "input_cost_per_token_above_100k_tokens_batches": 2.5e-07, "output_cost_per_token_above_100k_tokens_batches": 1.25e-06, "cache_creation_input_token_cost_above_100k_tokens_batches": 3.125e-07, @@ -76179,7 +76646,11 @@ "output_cost_per_reasoning_token": 5e-06, "output_cost_per_token": 5e-06, "output_cost_per_token_batches": 2.5e-06, - "search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014}, + "search_context_cost_per_query": { + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014, + "search_context_size_high": 0.014 + }, "source": "https://ai.google.dev/gemini-api/docs/pricing", "supported_endpoints": [ "/v1/chat/completions", @@ -76225,7 +76696,11 @@ "output_cost_per_reasoning_token": 5e-06, "output_cost_per_token": 5e-06, "output_cost_per_token_batches": 2.5e-06, - "search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014}, + "search_context_cost_per_query": { + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014, + "search_context_size_high": 0.014 + }, "source": "https://ai.google.dev/gemini-api/docs/pricing", "supported_endpoints": [ "/v1/chat/completions", @@ -76279,7 +76754,11 @@ "output_cost_per_token_priority": 1.8e-05, "prompt_cache_min_tokens": 4096.0, "rpm": 2000.0, - "search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014}, + "search_context_cost_per_query": { + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014, + "search_context_size_high": 0.014 + }, "source": "https://ai.google.dev/gemini-api/docs/pricing", "supported_endpoints": [ "/v1/chat/completions", @@ -76338,7 +76817,11 @@ "output_cost_per_token_priority": 1.8e-05, "prompt_cache_min_tokens": 6144.0, "regional_endpoint_uplift_multiplier": 1.1, - "search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014}, + "search_context_cost_per_query": { + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014, + "search_context_size_high": 0.014 + }, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", @@ -76396,7 +76879,11 @@ "output_cost_per_token_priority": 3.6e-05, "prompt_cache_min_tokens": 6144.0, "regional_endpoint_uplift_multiplier": 1.1, - "search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014}, + "search_context_cost_per_query": { + "search_context_size_low": 0.014, + "search_context_size_medium": 0.014, + "search_context_size_high": 0.014 + }, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", @@ -76624,7 +77111,11 @@ "mode": "chat", "output_cost_per_token": 4e-05, "prompt_cache_min_tokens": 512.0, - "search_context_cost_per_query": {"search_context_size_high":0.02,"search_context_size_low":0.02,"search_context_size_medium":0.02}, + "search_context_cost_per_query": { + "search_context_size_high": 0.02, + "search_context_size_low": 0.02, + "search_context_size_medium": 0.02 + }, "source": "https://learn.microsoft.com/en-us/azure/foundry/foundry-models/concepts/claude-models", "supports_adaptive_thinking": true, "supports_assistant_prefill": false, @@ -76988,5 +77479,605 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true + }, + "bedrock/claude_platform/claude-fable-5": { + "supports_web_search": true, + "supports_anthropic_compaction": true, + "cache_creation_input_token_cost": 1.25e-05, + "cache_creation_input_token_cost_above_1hr": 2e-05, + "cache_creation_input_token_cost_batches": 6.25e-06, + "cache_read_input_token_cost": 1e-06, + "cache_read_input_token_cost_batches": 5e-07, + "input_cost_per_token": 1e-05, + "input_cost_per_token_batches": 5e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5e-05, + "output_cost_per_token_batches": 2.5e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "thinking_always_on": true, + "supports_mid_conversation_system": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true, + "provider_specific_entry": { + "us": 1.1 + }, + "supports_output_config": true, + "prompt_cache_min_tokens": 512, + "supports_native_structured_output": true, + "source": "https://platform.claude.com/docs/en/about-claude/pricing#claude-platform-on-aws-pricing" + }, + "bedrock/claude_platform/claude-fable-5-1": { + "supports_web_search": true, + "supports_anthropic_compaction": true, + "cache_creation_input_token_cost": 1.25e-05, + "cache_creation_input_token_cost_above_1hr": 2e-05, + "cache_creation_input_token_cost_batches": 6.25e-06, + "cache_read_input_token_cost": 2.5e-07, + "cache_read_input_token_cost_batches": 1.25e-07, + "input_cost_per_token": 1e-05, + "input_cost_per_token_batches": 5e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5e-05, + "output_cost_per_token_batches": 2.5e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "thinking_always_on": true, + "supports_mid_conversation_system": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_forced_tool_use": false, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true, + "provider_specific_entry": { + "us": 1.1 + }, + "supports_output_config": true, + "prompt_cache_min_tokens": 512, + "supports_native_structured_output": true, + "source": "https://platform.claude.com/docs/en/about-claude/pricing#claude-platform-on-aws-pricing" + }, + "bedrock/claude_platform/claude-haiku-4-5": { + "cache_creation_input_token_cost": 1.25e-06, + "cache_creation_input_token_cost_above_1hr": 2e-06, + "cache_creation_input_token_cost_batches": 6.25e-07, + "cache_read_input_token_cost": 1e-07, + "cache_read_input_token_cost_batches": 5e-08, + "input_cost_per_token": 1e-06, + "input_cost_per_token_batches": 5e-07, + "litellm_provider": "bedrock", + "max_input_tokens": 200000, + "max_output_tokens": 64000, + "max_tokens": 64000, + "mode": "chat", + "output_cost_per_token": 5e-06, + "output_cost_per_token_batches": 2.5e-06, + "supports_assistant_prefill": true, + "supports_function_calling": true, + "supports_native_structured_output": true, + "supports_computer_use": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true, + "prompt_cache_min_tokens": 4096, + "source": "https://platform.claude.com/docs/en/about-claude/pricing#claude-platform-on-aws-pricing" + }, + "bedrock/claude_platform/claude-haiku-5-5": { + "input_cost_per_token_above_100k_tokens_batches": 2.5e-07, + "output_cost_per_token_above_100k_tokens_batches": 1.25e-06, + "cache_creation_input_token_cost_above_100k_tokens_batches": 3.125e-07, + "cache_read_input_token_cost_above_100k_tokens_batches": 2.5e-08, + "supports_anthropic_compaction": true, + "cache_creation_input_token_cost": 1.25e-07, + "cache_creation_input_token_cost_above_1hr": 2e-07, + "cache_creation_input_token_cost_batches": 6.25e-08, + "cache_read_input_token_cost": 1e-08, + "cache_read_input_token_cost_batches": 5e-09, + "input_cost_per_token": 1e-07, + "input_cost_per_token_batches": 5e-08, + "litellm_provider": "bedrock", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 5e-07, + "output_cost_per_token_batches": 2.5e-07, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "supports_mid_conversation_system": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_native_structured_output": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true, + "provider_specific_entry": { + "us": 1.1 + }, + "supports_output_config": true, + "supports_forced_tool_use": true, + "thinking_always_on": false, + "prompt_cache_min_tokens": 512, + "source": "https://platform.claude.com/docs/en/about-claude/pricing#claude-platform-on-aws-pricing", + "supports_web_search": true, + "input_cost_per_token_above_100k_tokens": 5e-07, + "output_cost_per_token_above_100k_tokens": 2.5e-06, + "cache_creation_input_token_cost_above_100k_tokens": 6.25e-07, + "cache_creation_input_token_cost_above_1hr_above_100k_tokens": 1e-06, + "cache_read_input_token_cost_above_100k_tokens": 5e-08 + }, + "bedrock/claude_platform/claude-opus-4-5": { + "cache_creation_input_token_cost": 6.25e-06, + "cache_creation_input_token_cost_above_1hr": 1e-05, + "cache_creation_input_token_cost_batches": 3.125e-06, + "cache_read_input_token_cost": 5e-07, + "cache_read_input_token_cost_batches": 2.5e-07, + "input_cost_per_token": 5e-06, + "input_cost_per_token_batches": 2.5e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 200000, + "max_output_tokens": 64000, + "max_tokens": 64000, + "mode": "chat", + "output_cost_per_token": 2.5e-05, + "output_cost_per_token_batches": 1.25e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_assistant_prefill": true, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_thinking_cache_preservation": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_native_structured_output": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_output_config": true, + "prompt_cache_min_tokens": 4096, + "source": "https://platform.claude.com/docs/en/about-claude/pricing#claude-platform-on-aws-pricing" + }, + "bedrock/claude_platform/claude-opus-4-6": { + "supports_web_search": true, + "supports_anthropic_compaction": true, + "cache_creation_input_token_cost": 6.25e-06, + "cache_creation_input_token_cost_above_1hr": 1e-05, + "cache_creation_input_token_cost_batches": 3.125e-06, + "cache_read_input_token_cost": 5e-07, + "cache_read_input_token_cost_batches": 2.5e-07, + "input_cost_per_token": 5e-06, + "input_cost_per_token_batches": 2.5e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 2.5e-05, + "output_cost_per_token_batches": 1.25e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "supports_legacy_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_thinking_cache_preservation": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_native_structured_output": true, + "supports_tool_choice": true, + "supports_vision": true, + "provider_specific_entry": { + "us": 1.1 + }, + "supports_output_config": true, + "supports_max_reasoning_effort": true, + "prompt_cache_min_tokens": 4096, + "source": "https://platform.claude.com/docs/en/about-claude/pricing#claude-platform-on-aws-pricing" + }, + "bedrock/claude_platform/claude-opus-4-7": { + "supports_web_search": true, + "supports_anthropic_compaction": true, + "cache_creation_input_token_cost": 6.25e-06, + "cache_creation_input_token_cost_above_1hr": 1e-05, + "cache_creation_input_token_cost_batches": 3.125e-06, + "cache_read_input_token_cost": 5e-07, + "cache_read_input_token_cost_batches": 2.5e-07, + "input_cost_per_token": 5e-06, + "input_cost_per_token_batches": 2.5e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 2.5e-05, + "output_cost_per_token_batches": 1.25e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_thinking_cache_preservation": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_native_structured_output": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true, + "provider_specific_entry": { + "us": 1.1 + }, + "supports_output_config": true, + "prompt_cache_min_tokens": 2048, + "source": "https://platform.claude.com/docs/en/about-claude/pricing#claude-platform-on-aws-pricing" + }, + "bedrock/claude_platform/claude-opus-4-8": { + "supports_web_search": true, + "supports_anthropic_compaction": true, + "cache_creation_input_token_cost": 6.25e-06, + "cache_creation_input_token_cost_above_1hr": 1e-05, + "cache_creation_input_token_cost_batches": 3.125e-06, + "cache_read_input_token_cost": 5e-07, + "cache_read_input_token_cost_batches": 2.5e-07, + "input_cost_per_token": 5e-06, + "input_cost_per_token_batches": 2.5e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 2.5e-05, + "output_cost_per_token_batches": 1.25e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "supports_mid_conversation_system": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_thinking_cache_preservation": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_native_structured_output": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true, + "provider_specific_entry": { + "us": 1.1 + }, + "supports_output_config": true, + "prompt_cache_min_tokens": 1024, + "source": "https://platform.claude.com/docs/en/about-claude/pricing#claude-platform-on-aws-pricing" + }, + "bedrock/claude_platform/claude-opus-5": { + "supports_web_search": true, + "supports_anthropic_compaction": true, + "cache_creation_input_token_cost": 6.25e-06, + "cache_creation_input_token_cost_above_1hr": 1e-05, + "cache_creation_input_token_cost_batches": 3.125e-06, + "cache_read_input_token_cost": 5e-07, + "cache_read_input_token_cost_batches": 2.5e-07, + "input_cost_per_token": 5e-06, + "input_cost_per_token_batches": 2.5e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 2.5e-05, + "output_cost_per_token_batches": 1.25e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "supports_mid_conversation_system": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_thinking_cache_preservation": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_native_structured_output": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true, + "provider_specific_entry": { + "us": 1.1 + }, + "supports_output_config": true, + "prompt_cache_min_tokens": 512, + "source": "https://platform.claude.com/docs/en/about-claude/pricing#claude-platform-on-aws-pricing" + }, + "bedrock/claude_platform/claude-opus-5-5": { + "supports_web_search": true, + "supports_anthropic_compaction": true, + "cache_creation_input_token_cost": 5e-06, + "cache_creation_input_token_cost_above_1hr": 8e-06, + "cache_creation_input_token_cost_batches": 2.5e-06, + "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_batches": 1e-07, + "input_cost_per_token": 4e-06, + "input_cost_per_token_batches": 2e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 2e-05, + "output_cost_per_token_batches": 1e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "thinking_always_on": true, + "supports_mid_conversation_system": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_forced_tool_use": false, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_native_structured_output": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true, + "provider_specific_entry": { + "us": 1.1 + }, + "supports_output_config": true, + "prompt_cache_min_tokens": 512, + "source": "https://platform.claude.com/docs/en/about-claude/pricing#claude-platform-on-aws-pricing" + }, + "bedrock/claude_platform/claude-sonnet-4-5": { + "cache_creation_input_token_cost": 3.75e-06, + "cache_creation_input_token_cost_above_1hr": 6e-06, + "cache_read_input_token_cost": 3e-07, + "input_cost_per_token": 3e-06, + "cache_creation_input_token_cost_batches": 1.875e-06, + "cache_read_input_token_cost_batches": 1.5e-07, + "deprecation_date": "2026-11-30", + "input_cost_per_token_batches": 1.5e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 200000, + "max_output_tokens": 64000, + "max_tokens": 64000, + "mode": "chat", + "output_cost_per_token": 1.5e-05, + "output_cost_per_token_batches": 7.5e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_assistant_prefill": true, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_native_structured_output": true, + "supports_tool_choice": true, + "supports_vision": true, + "prompt_cache_min_tokens": 1024, + "source": "https://platform.claude.com/docs/en/about-claude/pricing#claude-platform-on-aws-pricing" + }, + "bedrock/claude_platform/claude-sonnet-4-6": { + "supports_web_search": true, + "supports_anthropic_compaction": true, + "cache_creation_input_token_cost": 3.75e-06, + "cache_creation_input_token_cost_above_1hr": 6e-06, + "cache_creation_input_token_cost_batches": 1.875e-06, + "cache_read_input_token_cost": 3e-07, + "cache_read_input_token_cost_batches": 1.5e-07, + "input_cost_per_token": 3e-06, + "input_cost_per_token_batches": 1.5e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 1.5e-05, + "output_cost_per_token_batches": 7.5e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "supports_legacy_thinking": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_thinking_cache_preservation": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_native_structured_output": true, + "supports_max_reasoning_effort": true, + "supports_tool_choice": true, + "supports_vision": true, + "supports_output_config": true, + "prompt_cache_min_tokens": 1024, + "provider_specific_entry": { + "us": 1.1 + }, + "source": "https://platform.claude.com/docs/en/about-claude/pricing#claude-platform-on-aws-pricing" + }, + "bedrock/claude_platform/claude-sonnet-5": { + "supports_web_search": true, + "supports_anthropic_compaction": true, + "cache_creation_input_token_cost": 2.5e-06, + "cache_creation_input_token_cost_above_1hr": 4e-06, + "cache_creation_input_token_cost_batches": 1.25e-06, + "cache_read_input_token_cost": 2e-07, + "cache_read_input_token_cost_batches": 1e-07, + "input_cost_per_token": 2e-06, + "input_cost_per_token_batches": 1e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 1e-05, + "output_cost_per_token_batches": 5e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "supports_mid_conversation_system": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_native_structured_output": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_thinking_cache_preservation": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true, + "provider_specific_entry": { + "us": 1.1 + }, + "supports_output_config": true, + "prompt_cache_min_tokens": 1024, + "source": "https://platform.claude.com/docs/en/about-claude/pricing#claude-platform-on-aws-pricing" + }, + "bedrock/claude_platform/claude-sonnet-5-5": { + "supports_anthropic_compaction": true, + "cache_creation_input_token_cost": 2.5e-06, + "cache_creation_input_token_cost_above_1hr": 4e-06, + "cache_creation_input_token_cost_batches": 1.25e-06, + "cache_read_input_token_cost": 1e-07, + "cache_read_input_token_cost_batches": 5e-08, + "input_cost_per_token": 2e-06, + "input_cost_per_token_batches": 1e-06, + "litellm_provider": "bedrock", + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 1e-05, + "output_cost_per_token_batches": 5e-06, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_adaptive_thinking": true, + "supports_mid_conversation_system": true, + "supports_assistant_prefill": false, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_native_structured_output": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_sampling_params": false, + "supports_tool_choice": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true, + "supports_max_reasoning_effort": true, + "provider_specific_entry": { + "us": 1.1 + }, + "supports_output_config": true, + "supports_forced_tool_use": false, + "thinking_always_on": true, + "prompt_cache_min_tokens": 512, + "source": "https://platform.claude.com/docs/en/about-claude/pricing#claude-platform-on-aws-pricing", + "supports_web_search": true } } diff --git a/pyproject.toml b/pyproject.toml index c26d9a7c6ec..d7deb25237b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "litellm" -version = "1.108.0" +version = "1.108.1" description = "Library to easily interface with LLM API providers" readme = "README.md" requires-python = ">=3.10, <3.15" @@ -79,7 +79,7 @@ proxy = [ "httpx2>=2.5.0,<3", "pydantic>=2.12.0,<3", "litellm-proxy-extras==0.4.110", - "litellm-enterprise==0.1.79", + "litellm-enterprise==0.1.79.post1", "RestrictedPython>=8.5,<9.0", "rich>=13.9.4,<14.0", "InquirerPy>=0.3.4,<1.0", @@ -393,7 +393,7 @@ litellm-enterprise = { workspace = true } members = ["enterprise", "litellm-proxy-extras"] [tool.commitizen] -version = "1.108.0" +version = "1.108.1" version_files = [ "pyproject.toml:^version", ] diff --git a/tests/integration/_support/custom_video_provider.py b/tests/integration/_support/custom_video_provider.py new file mode 100644 index 00000000000..3513d1e3186 --- /dev/null +++ b/tests/integration/_support/custom_video_provider.py @@ -0,0 +1,181 @@ +"""A customer-owned video provider for the proxy's `custom_provider_map`: a job queue with its own wire shape.""" + +from collections.abc import Mapping +from types import MappingProxyType +from typing import Final + +import httpx +from httpx._types import RequestFiles + +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.llms.base_llm.videos.transformation import BaseVideoConfig +from litellm.llms.custom_llm import CustomLLM +from litellm.types.router import GenericLiteLLMParams +from litellm.types.videos.main import VideoCreateOptionalRequestParams, VideoObject +from litellm.types.videos.utils import encode_video_id_with_provider, extract_original_video_id + +_STATUS_BY_STATE: Final = MappingProxyType({"queued": "queued", "running": "in_progress", "done": "completed"}) + + +def _video_object(job: Mapping[str, object], model: str | None, custom_llm_provider: str | None) -> VideoObject: + job_id: Final = str(job["job_id"]) + seconds: Final = job.get("seconds") + video: Final = VideoObject( + id=encode_video_id_with_provider(job_id, custom_llm_provider, model) if custom_llm_provider else job_id, + object="video", + status=_STATUS_BY_STATE[str(job["state"])], + created_at=0, + model=model, + seconds=str(seconds) if seconds is not None else None, + ) + return video + + +class QueueVideoConfig(BaseVideoConfig): + def get_supported_openai_params(self, model: str) -> list: # mutable-ok: BaseVideoConfig contract + return ["seconds"] + + def map_openai_params( + self, + video_create_optional_params: VideoCreateOptionalRequestParams, + model: str, + drop_params: bool, + ) -> dict: # mutable-ok: BaseVideoConfig contract + return {key: value for key, value in video_create_optional_params.items() if key == "seconds"} + + def validate_environment( + self, + headers: dict, # mutable-ok: BaseVideoConfig contract + model: str, + api_key: str | None = None, + litellm_params: GenericLiteLLMParams | None = None, + ) -> dict: # mutable-ok: BaseVideoConfig contract + key: Final = (litellm_params.api_key if litellm_params is not None else None) or api_key or "" + return {**headers, "X-Api-Key": key} + + def get_complete_url( + self, + model: str, + api_base: str | None, + litellm_params: dict, # mutable-ok: BaseVideoConfig contract + ) -> str: + if api_base is None: + raise ValueError("api_base is required") + return f"{api_base.rstrip('/')}/jobs" + + def transform_video_create_request( + self, + model: str, + prompt: str, + api_base: str, + video_create_optional_request_params: dict, # mutable-ok: BaseVideoConfig contract + litellm_params: GenericLiteLLMParams, + headers: dict, # mutable-ok: BaseVideoConfig contract + ) -> tuple[dict, RequestFiles, str]: # mutable-ok: BaseVideoConfig contract + seconds: Final = video_create_optional_request_params.get("seconds") + return ({"model": model, "prompt": prompt, "seconds": seconds}, [], api_base) + + def transform_video_create_response( + self, + model: str, + raw_response: httpx.Response, + logging_obj: LiteLLMLoggingObj, + custom_llm_provider: str | None = None, + request_data: dict | None = None, + ) -> VideoObject: + return _video_object(raw_response.json(), model, custom_llm_provider) + + def transform_video_status_retrieve_request( + self, + video_id: str, + api_base: str, + litellm_params: GenericLiteLLMParams, + headers: dict, # mutable-ok: BaseVideoConfig contract + ) -> tuple[str, dict]: # mutable-ok: BaseVideoConfig contract + return (f"{api_base}/{extract_original_video_id(video_id)}", {}) + + def transform_video_status_retrieve_response( + self, + raw_response: httpx.Response, + logging_obj: LiteLLMLoggingObj, + custom_llm_provider: str | None = None, + client: object = None, + ) -> VideoObject: + return _video_object(raw_response.json(), None, custom_llm_provider) + + def transform_video_content_request( + self, + video_id: str, + api_base: str, + litellm_params: GenericLiteLLMParams, + headers: dict, # mutable-ok: BaseVideoConfig contract + variant: str | None = None, + ) -> tuple[str, dict]: # mutable-ok: BaseVideoConfig contract + return (f"{api_base}/{extract_original_video_id(video_id)}/file", {}) + + def transform_video_content_response( + self, + raw_response: httpx.Response, + logging_obj: LiteLLMLoggingObj, + client: object = None, + ) -> bytes: + return raw_response.content + + def transform_video_list_request( + self, + api_base: str, + litellm_params: GenericLiteLLMParams, + headers: dict, # mutable-ok: BaseVideoConfig contract + after: str | None = None, + limit: int | None = None, + order: str | None = None, + extra_query: dict[str, object] | None = None, + ) -> tuple[str, dict]: # mutable-ok: BaseVideoConfig contract + return (api_base, {}) + + def transform_video_list_response( + self, + raw_response: httpx.Response, + logging_obj: LiteLLMLoggingObj, + custom_llm_provider: str | None = None, + ) -> dict[str, str]: # mutable-ok: BaseVideoConfig contract + return raw_response.json() + + def transform_video_remix_request( + self, + video_id: str, + prompt: str, + api_base: str, + litellm_params: GenericLiteLLMParams, + headers: dict, # mutable-ok: BaseVideoConfig contract + extra_body: dict[str, object] | None = None, + ) -> tuple[str, dict]: # mutable-ok: BaseVideoConfig contract + raise NotImplementedError("the queue backend has no remix") + + def transform_video_remix_response( + self, + raw_response: httpx.Response, + logging_obj: LiteLLMLoggingObj, + custom_llm_provider: str | None = None, + ) -> VideoObject: + raise NotImplementedError("the queue backend has no remix") + + def transform_video_delete_request( + self, + video_id: str, + api_base: str, + litellm_params: GenericLiteLLMParams, + headers: dict, # mutable-ok: BaseVideoConfig contract + ) -> tuple[str, dict]: # mutable-ok: BaseVideoConfig contract + raise NotImplementedError("the queue backend has no delete") + + def transform_video_delete_response( + self, + raw_response: httpx.Response, + logging_obj: LiteLLMLoggingObj, + ) -> VideoObject: + raise NotImplementedError("the queue backend has no delete") + + +VIDEO_CONFIG: Final = QueueVideoConfig() +IMAGE_ONLY_HANDLER: Final = CustomLLM() diff --git a/tests/integration/observability/test_guardrail_effects.py b/tests/integration/observability/test_guardrail_effects.py index 462480592ef..8409648bec8 100644 --- a/tests/integration/observability/test_guardrail_effects.py +++ b/tests/integration/observability/test_guardrail_effects.py @@ -1,7 +1,9 @@ import asyncio import json import re +import socket import threading +import time import uuid from collections.abc import Callable, Iterator, Mapping from concurrent.futures import ThreadPoolExecutor @@ -11,6 +13,7 @@ from dataclasses import dataclass from pathlib import Path from types import MappingProxyType from typing import Final +from urllib.parse import urlsplit import anthropic import httpx @@ -2928,3 +2931,60 @@ def test_disconnect_scans_keep_recording_after_a_worker_is_killed(gateway: Gatew assert all(rig.secret(index) in line for index, line in enumerate(received)), received assert _post_call_statuses(rig, rig.model, rows=8) == (("success",),) * 8 assert rig.owned.process.poll() is None + + +def _openai_upstream(request: Request) -> Reply: + if request.target.endswith("/vector_stores"): + return Reply(body=json.dumps({"id": "vs_" + uuid.uuid4().hex, "object": "vector_store", "name": None}).encode()) + return Reply(body=_chat_completion_body("control-reply")) + + +def _post_half_of_the_body_then_hang_up(gateway: Gateway, path: str, body: bytes) -> None: + address: Final = urlsplit(str(gateway.client.base_url)) + assert address.hostname is not None and address.port is not None, address + head: Final = ( + f"POST {path} HTTP/1.1\r\nHost: {address.netloc}\r\nAuthorization: Bearer {gateway.key}\r\n" + f"Content-Type: application/json\r\nContent-Length: {len(body)}\r\n\r\n" + ).encode() + with socket.create_connection((address.hostname, address.port)) as connection: + connection.sendall(head + body[: len(body) // 2]) + time.sleep(0.5) + + +@pytest.mark.parametrize( + ("path", "body"), + ( + pytest.param( + "/v1/chat/completions", + b'{"model": "MODEL", "messages": [{"role": "user", "content": "Say OK"}], "max_tokens": 16}', + id="chat-completions-where-an-empty-body-is-invalid", + ), + pytest.param( + "/v1/vector_stores", + b'{"name": "integration-store", "metadata": {"owner": "integration", "purpose": "half upload"}}', + id="vector-stores-where-an-empty-body-is-valid", + ), + ), +) +def test_client_disconnect_during_the_body_upload_runs_no_endpoint_and_writes_no_spend_row( + gateway: Gateway, tmp_path: Path, path: str, body: bytes +) -> None: + with ( + wire_server(_openai_upstream) as upstream, + owned_proxy_process( + gateway, tmp_path, {"OPENAI_BASE_URL": upstream.url + "/v1", "OPENAI_API_KEY": "synthetic-openai-key"} + ) as owned, + owned.gateway.scenario() as scenario, + ): + model: Final = scenario.model(api_base=upstream.url + "/v1", api_key="synthetic-openai-key") + _post_half_of_the_body_then_hang_up(owned.gateway, path, body.replace(b"MODEL", model.encode())) + control: Final = owned.gateway.chat(model) + rows: Final = eventually( + lambda: owned.gateway.rows('SELECT request_id, call_type, status FROM "LiteLLM_SpendLogs"', ()), + lambda values: any(row["request_id"] == control["id"] for row in values), + seconds=70, + ) + assert [(row["call_type"], row["status"]) for row in rows] == [("acompletion", "success")], rows + assert [request.target for request in upstream.drain()] == ["/v1/chat/completions"], ( + "the half-uploaded request must never reach the provider" + ) diff --git a/tests/integration/providers/test_bedrock_native_guardrail_trace_spend.py b/tests/integration/providers/test_bedrock_native_guardrail_trace_spend.py new file mode 100644 index 00000000000..b1f2d5a9181 --- /dev/null +++ b/tests/integration/providers/test_bedrock_native_guardrail_trace_spend.py @@ -0,0 +1,423 @@ +"""A Bedrock deployment whose ``guardrailConfig`` has ``trace: enabled`` bills the policy units Bedrock reports in the +trace as guardrail cost, on every surface the deployment serves: ``/v1/chat/completions`` and ``/v1/responses`` over +Converse and Invoke, streamed and not, and the four ``/bedrock/model/{model}`` pass-through actions. The cost rides +``x-litellm-response-cost`` on non-streaming responses and the spend row's ``spend`` and ``guardrail_information`` +everywhere. ``trace: disabled`` leaves no usage, so nothing is billed beyond the tokens (LIT-10101)""" + +import base64 +import json +import uuid +from collections.abc import Iterator, Mapping +from dataclasses import dataclass +from hashlib import sha256 +from pathlib import Path +from typing import Final + +import boto3 +import pytest +import yaml +from botocore.config import Config +from integration._support.budget import case_timeout_seconds +from integration._support.client import Gateway, Scenario, eventually +from integration._support.process import graceful_stop_seconds, owned_proxy +from integration._support.upstream import _aws_event_frame +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import JsonValue, TypeAdapter + +import litellm + +pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 240) + +_MODEL_ID: Final = "anthropic.claude-sonnet-4-5-20250929-v1:0" +_EVENT_STREAM: Final = "application/vnd.amazon.eventstream" +_GUARDRAIL_ID: Final = "native-guardrail" +_INPUT_COST_PER_TOKEN: Final = 0.001 +_OUTPUT_COST_PER_TOKEN: Final = 0.002 +_PROMPT_TOKENS: Final = 5 +_COMPLETION_TOKENS: Final = 1 +_INPUT_USAGE: Final[Mapping[str, int]] = {"topicPolicyUnits": 1, "contentPolicyUnits": 1, "wordPolicyUnits": 1} +_OUTPUT_USAGE: Final[Mapping[str, int]] = {"topicPolicyUnits": 2, "sensitiveInformationPolicyUnits": 2} +_SUMMED_USAGE: Final[Mapping[str, int]] = { + "topicPolicyUnits": 3, + "contentPolicyUnits": 1, + "wordPolicyUnits": 1, + "sensitiveInformationPolicyUnits": 2, +} +_ACTIONS: Final = ("converse", "converse-stream", "invoke", "invoke-with-response-stream") +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_SPEND_SQL: Final = 'SELECT spend, metadata FROM "LiteLLM_SpendLogs" WHERE api_key=%s ORDER BY "startTime"' + + +def _llm_cost() -> float: + return _PROMPT_TOKENS * _INPUT_COST_PER_TOKEN + _COMPLETION_TOKENS * _OUTPUT_COST_PER_TOKEN + + +def _guardrail_cost() -> float: + """What the proxy's own cost map prices the traced units at, so the test owns the relation, not AWS's rates.""" + rates: Final = litellm.model_cost["bedrock/guardrails"]["guardrail_cost_per_unit"] + return sum(units * rates[counter] for counter, units in _SUMMED_USAGE.items()) + + +def _action(request: Request) -> str: + return request.target.partition("?")[0].rsplit("/", 1)[-1] + + +def _traced(request: Request) -> bool: + if _action(request).startswith("converse"): + config: Final = json.loads(request.body).get("guardrailConfig") or {} + return config.get("trace") == "enabled" + return request.headers.get("x-amzn-bedrock-trace", "").upper() == "ENABLED" + + +def _assessment(usage: Mapping[str, int]) -> dict[str, JsonValue]: + return {"invocationMetrics": {"usage": dict(usage), "guardrailProcessingLatency": 7}} + + +def _converse_trace() -> dict[str, JsonValue]: + return { + "guardrail": { + "inputAssessment": {_GUARDRAIL_ID: _assessment(_INPUT_USAGE)}, + "outputAssessments": {_GUARDRAIL_ID: [_assessment(_OUTPUT_USAGE)]}, + } + } + + +def _invoke_trace() -> dict[str, JsonValue]: + return { + "guardrail": { + "input": {_GUARDRAIL_ID: _assessment(_INPUT_USAGE)}, + "outputs": [{_GUARDRAIL_ID: _assessment(_OUTPUT_USAGE)}], + } + } + + +def _bedrock(request: Request) -> Reply: + """Bedrock's runtime as the caller sees it: a one-token answer, and the guardrail trace with its policy units only + when the request enabled tracing, in the shape of the action that served it.""" + traced: Final = _traced(request) + action: Final = _action(request) + usage: Final = {"inputTokens": _PROMPT_TOKENS, "outputTokens": _COMPLETION_TOKENS, "totalTokens": 6} + converse_trace: Final = {"trace": _converse_trace()} if traced else {} + if action == "converse": + body: Final = { + "output": {"message": {"role": "assistant", "content": [{"text": "ok"}]}}, + "stopReason": "end_turn", + "usage": usage, + "metrics": {"latencyMs": 3}, + **converse_trace, + } + return Reply(body=json.dumps(body).encode()) + if action == "converse-stream" and "toolConfig" in json.loads(request.body): + return Reply( + chunks=( + _aws_event_frame("messageStart", {"role": "assistant"}, "sc", "u"), + _aws_event_frame( + "contentBlockStart", + {"start": {"toolUse": {"toolUseId": "tool-1", "name": "lookup"}}, "contentBlockIndex": 0}, + "sc", + "u", + ), + _aws_event_frame( + "contentBlockDelta", + {"delta": {"toolUse": {"input": '{"q": "x"}'}}, "contentBlockIndex": 0}, + "sc", + "u", + ), + _aws_event_frame("contentBlockStop", {"contentBlockIndex": 0}, "sc", "u"), + _aws_event_frame("messageStop", {"stopReason": "tool_use"}, "sc", "u"), + _aws_event_frame( + "metadata", {"usage": usage, "metrics": {"latencyMs": 3}, **converse_trace}, "sc", "u" + ), + ), + content_type=_EVENT_STREAM, + ) + if action == "converse-stream": + return Reply( + chunks=( + _aws_event_frame("messageStart", {"role": "assistant"}, "sc", "u"), + _aws_event_frame("contentBlockDelta", {"delta": {"text": "ok"}, "contentBlockIndex": 0}, "sc", "u"), + _aws_event_frame("contentBlockStop", {"contentBlockIndex": 0}, "sc", "u"), + _aws_event_frame("messageStop", {"stopReason": "end_turn"}, "sc", "u"), + _aws_event_frame( + "metadata", {"usage": usage, "metrics": {"latencyMs": 3}, **converse_trace}, "sc", "u" + ), + ), + content_type=_EVENT_STREAM, + ) + invoke_trace: Final = {"amazon-bedrock-trace": _invoke_trace()} if traced else {} + if action == "invoke": + message: Final = { + "id": f"msg_{uuid.uuid4().hex[:12]}", + "type": "message", + "role": "assistant", + "model": "claude", + "content": [{"type": "text", "text": "ok"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": _PROMPT_TOKENS, "output_tokens": _COMPLETION_TOKENS}, + **invoke_trace, + } + return Reply(body=json.dumps(message).encode()) + events: Final = ( + { + "type": "message_start", + "message": { + "id": f"msg_{uuid.uuid4().hex[:12]}", + "type": "message", + "role": "assistant", + "model": "claude", + "content": [], + "stop_reason": None, + "stop_sequence": None, + "usage": {"input_tokens": _PROMPT_TOKENS, "output_tokens": 0}, + }, + }, + {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}, + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "ok"}}, + {"type": "content_block_stop", "index": 0}, + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn", "stop_sequence": None}, + "usage": {"output_tokens": _COMPLETION_TOKENS}, + }, + { + "type": "message_stop", + "amazon-bedrock-invocationMetrics": { + "inputTokenCount": _PROMPT_TOKENS, + "outputTokenCount": _COMPLETION_TOKENS, + "invocationLatency": 3, + "firstByteLatency": 1, + }, + **invoke_trace, + }, + ) + return Reply( + chunks=tuple( + _aws_event_frame("chunk", {"bytes": base64.b64encode(json.dumps(event).encode()).decode()}, "sc", "u") + for event in events + ), + content_type=_EVENT_STREAM, + ) + + +def _deployment(wire: Wire, model: str, trace: str) -> Mapping[str, JsonValue]: + return { + "model": model, + "api_base": wire.url, + "aws_access_key_id": "AKIASCRIPTEDPROVIDER", + "aws_secret_access_key": "scripted-secret", + "aws_region_name": "us-east-1", + "use_in_pass_through": True, + "input_cost_per_token": _INPUT_COST_PER_TOKEN, + "output_cost_per_token": _OUTPUT_COST_PER_TOKEN, + "guardrailConfig": {"guardrailIdentifier": _GUARDRAIL_ID, "guardrailVersion": "1", "trace": trace}, + } + + +@dataclass(frozen=True, slots=True) +class Rig: + proxy: Gateway + scenario: Scenario + converse_model: str + invoke_model: str + untraced_model: str + + def key(self) -> str: + return self.scenario.key(key_alias=f"native-trace-{uuid.uuid4().hex[:8]}") + + def runtime(self, monkeypatch: pytest.MonkeyPatch, key: str): # noqa: ANN201 # boto3 clients are untyped + monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", key) + return boto3.client( + "bedrock-runtime", + endpoint_url=f"{str(self.proxy.client.base_url).rstrip('/')}/bedrock", + region_name="us-east-1", + config=Config(read_timeout=case_timeout_seconds(), retries={"total_max_attempts": 1}), + ) + + +@pytest.fixture(scope="module") +def rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]: + suffix: Final = uuid.uuid4().hex[:12] + converse_model: Final = f"bedrock-traced-converse-{suffix}" + invoke_model: Final = f"bedrock-traced-invoke-{suffix}" + untraced_model: Final = f"bedrock-untraced-{suffix}" + directory: Final = tmp_path_factory.mktemp("bedrock-native-guardrail-trace") + base_config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + with wire_server(_bedrock) as wire: + config: Final = { + **base_config, + "model_list": [ + *base_config.get("model_list", []), + {"model_name": converse_model, "litellm_params": _deployment(wire, f"bedrock/{_MODEL_ID}", "enabled")}, + { + "model_name": invoke_model, + "litellm_params": _deployment(wire, f"bedrock/invoke/{_MODEL_ID}", "enabled"), + }, + {"model_name": untraced_model, "litellm_params": _deployment(wire, f"bedrock/{_MODEL_ID}", "disabled")}, + ], + } + config_path: Final = directory / "bedrock-native-guardrail-trace.yaml" + config_path.write_text(yaml.safe_dump(config)) + with owned_proxy(session_gateway, directory, {}, config=config_path) as proxy, proxy.scenario() as scenario: + yield Rig( + proxy=proxy, + scenario=scenario, + converse_model=converse_model, + invoke_model=invoke_model, + untraced_model=untraced_model, + ) + + +@dataclass(frozen=True, slots=True) +class SpendRow: + spend: float + guardrail_information: tuple[dict[str, JsonValue], ...] + + +def _spend_row(proxy: Gateway, key: str) -> SpendRow: + digest: Final = sha256(key.encode()).hexdigest() + (row,) = eventually(lambda: proxy.rows(_SPEND_SQL, (digest,)), lambda rows: len(rows) == 1, seconds=70) + metadata: Final = _JSON_OBJECT.validate_python( + json.loads(row["metadata"]) if isinstance(row["metadata"], str) else row["metadata"] + ) + raw_information: Final = metadata.get("guardrail_information") + entries: Final = raw_information if isinstance(raw_information, list) else [raw_information] + return SpendRow( + spend=float(str(row["spend"])), + guardrail_information=tuple(_JSON_OBJECT.validate_python(entry) for entry in entries if entry is not None), + ) + + +def _assert_billed_as_guardrail_cost(row: SpendRow) -> None: + assert row.spend == pytest.approx(_llm_cost() + _guardrail_cost()) + (entry,) = row.guardrail_information + assert entry["guardrail_provider"] == "bedrock" + assert entry["guardrail_name"] == _GUARDRAIL_ID + assert entry["guardrail_usage"] == _SUMMED_USAGE + assert isinstance(entry["guardrail_cost"], float) and entry["guardrail_cost"] == pytest.approx(_guardrail_cost()) + + +def _chat_body(model: str, stream: bool) -> dict[str, JsonValue]: + return {"model": model, "messages": [{"role": "user", "content": "hello"}], "max_tokens": 16, "stream": stream} + + +def _responses_body(model: str, stream: bool) -> dict[str, JsonValue]: + return {"model": model, "input": "hello", "stream": stream} + + +@pytest.mark.parametrize("api", ["converse", "invoke"]) +@pytest.mark.parametrize("endpoint", ["/v1/chat/completions", "/v1/responses"]) +def test_non_streaming_response_prices_the_trace_into_the_cost_header_and_the_spend_row( + rig: Rig, endpoint: str, api: str +) -> None: + key: Final = rig.key() + model: Final = rig.converse_model if api == "converse" else rig.invoke_model + body: Final = _chat_body(model, False) if endpoint == "/v1/chat/completions" else _responses_body(model, False) + response: Final = rig.proxy.request("POST", endpoint, body, key=key) + assert response.is_success, response.text[:400] + assert float(response.headers["x-litellm-response-cost"]) == pytest.approx(_llm_cost() + _guardrail_cost()) + _assert_billed_as_guardrail_cost(_spend_row(rig.proxy, key)) + + +@pytest.mark.parametrize("api", ["converse", "invoke"]) +@pytest.mark.parametrize("endpoint", ["/v1/chat/completions", "/v1/responses"]) +def test_streamed_response_prices_the_trace_of_the_final_event_into_the_spend_row( + rig: Rig, endpoint: str, api: str +) -> None: + key: Final = rig.key() + model: Final = rig.converse_model if api == "converse" else rig.invoke_model + body: Final = _chat_body(model, True) if endpoint == "/v1/chat/completions" else _responses_body(model, True) + with rig.proxy.client.stream( + "POST", endpoint, json=body, headers={"Authorization": f"Bearer {key}"}, timeout=case_timeout_seconds() + ) as response: + frames: Final = tuple(line for line in response.iter_lines() if line.startswith("data:")) + assert response.is_success and frames + _assert_billed_as_guardrail_cost(_spend_row(rig.proxy, key)) + + +@pytest.mark.parametrize("action", _ACTIONS) +def test_passthrough_action_prices_the_trace_into_the_spend_row( + rig: Rig, monkeypatch: pytest.MonkeyPatch, action: str +) -> None: + key: Final = rig.key() + runtime: Final = rig.runtime(monkeypatch, key) + messages: Final = [{"role": "user", "content": [{"text": "hello"}]}] + invoke_body: Final = json.dumps( + { + "anthropic_version": "bedrock-2023-05-31", + "max_tokens": 16, + "messages": [{"role": "user", "content": "hello"}], + } + ) + if action == "converse": + getattr(runtime, "converse")(modelId=rig.converse_model, messages=messages) + elif action == "converse-stream": + list(getattr(runtime, "converse_stream")(modelId=rig.converse_model, messages=messages)["stream"]) + elif action == "invoke": + getattr(runtime, "invoke_model")(modelId=rig.converse_model, body=invoke_body)["body"].read() + else: + list( + getattr(runtime, "invoke_model_with_response_stream")(modelId=rig.converse_model, body=invoke_body)["body"] + ) + _assert_billed_as_guardrail_cost(_spend_row(rig.proxy, key)) + + +def test_disabled_trace_bills_only_the_tokens(rig: Rig) -> None: + key: Final = rig.key() + response: Final = rig.proxy.request("POST", "/v1/chat/completions", _chat_body(rig.untraced_model, False), key=key) + assert response.is_success, response.text[:400] + assert "trace" not in response.json() + assert float(response.headers["x-litellm-response-cost"]) == pytest.approx(_llm_cost()) + row: Final = _spend_row(rig.proxy, key) + assert row.spend == pytest.approx(_llm_cost()) + assert row.guardrail_information == () + + +def _messages_body(model: str, stream: bool) -> dict[str, JsonValue]: + return { + "model": model, + "messages": [{"role": "user", "content": f"hello from the messages api, stream={stream}"}], + "max_tokens": 16, + "stream": stream, + } + + +def test_messages_non_streaming_response_applies_the_guardrail_and_prices_its_trace(rig: Rig) -> None: + key: Final = rig.key() + response: Final = rig.proxy.request("POST", "/v1/messages", _messages_body(rig.invoke_model, False), key=key) + assert response.is_success, response.text[:400] + assert "amazon-bedrock-trace" in response.json() + _assert_billed_as_guardrail_cost(_spend_row(rig.proxy, key)) + + +def test_messages_streamed_response_prices_the_trace_of_the_message_stop_event(rig: Rig) -> None: + key: Final = rig.key() + with rig.proxy.client.stream( + "POST", + "/v1/messages", + json=_messages_body(rig.invoke_model, True), + headers={"Authorization": f"Bearer {key}"}, + timeout=case_timeout_seconds(), + ) as response: + body: Final = b"".join(response.iter_bytes()) + assert response.is_success and b"message_stop" in body, body[:400] + _assert_billed_as_guardrail_cost(_spend_row(rig.proxy, key)) + + +def test_passthrough_streamed_tool_call_prices_the_trace_of_the_final_event( + rig: Rig, monkeypatch: pytest.MonkeyPatch +) -> None: + key: Final = rig.key() + runtime: Final = rig.runtime(monkeypatch, key) + tool_config: Final = { + "tools": [{"toolSpec": {"name": "lookup", "inputSchema": {"json": {"type": "object", "properties": {}}}}}] + } + events: Final = list( + getattr(runtime, "converse_stream")( + modelId=rig.converse_model, + messages=[{"role": "user", "content": [{"text": "hello"}]}], + toolConfig=tool_config, + )["stream"] + ) + assert any("contentBlockStart" in event for event in events) + _assert_billed_as_guardrail_cost(_spend_row(rig.proxy, key)) diff --git a/tests/integration/providers/test_count_tokens_error_shapes_wire.py b/tests/integration/providers/test_count_tokens_error_shapes_wire.py index 200f5e81673..077eac385dc 100644 --- a/tests/integration/providers/test_count_tokens_error_shapes_wire.py +++ b/tests/integration/providers/test_count_tokens_error_shapes_wire.py @@ -68,6 +68,65 @@ def _never_called(request: Request) -> Reply: raise AssertionError(f"count reached the provider without a key: {request.target}") +@pytest.mark.parametrize("nested", (False, True), ids=("user-content", "tool-result")) +def test_local_messages_count_accepts_search_result_text(gateway: Gateway, nested: bool) -> None: + search: Final[JsonValue] = { + "type": "search_result", + "source": "https://example.invalid/notes", + "title": "Project notes", + "content": [{"type": "text", "text": "The code word is 7f3a."}], + "citations": {"enabled": True}, + } + content: Final = [{"type": "tool_result", "tool_use_id": "toolu_1", "content": [search]}] if nested else [search] + text: Final[list[JsonValue]] = [ + {"type": "text", "text": "https://example.invalid/notes"}, + {"type": "text", "text": "Project notes"}, + {"type": "text", "text": "The code word is 7f3a."}, + ] + equivalent: Final = [{"type": "tool_result", "tool_use_id": "toolu_1", "content": text}] if nested else text + with gateway.scenario() as scenario: + model: Final = scenario.model(model="hosted_vllm/zai-org/GLM-5.3", api_base="http://127.0.0.1:9/v1") + counted: Final = gateway.request( + "POST", "/v1/messages/count_tokens", {"model": model, "messages": [{"role": "user", "content": content}]} + ) + control: Final = gateway.request( + "POST", "/v1/messages/count_tokens", {"model": model, "messages": [{"role": "user", "content": equivalent}]} + ) + assert counted.status_code == control.status_code == 200, (counted.text, control.text) + tokens: Final = _JSON_OBJECT.validate_json(counted.content)["input_tokens"] + assert tokens == _JSON_OBJECT.validate_json(control.content)["input_tokens"], (counted.text, control.text) + + +@pytest.mark.parametrize("nested", (False, True), ids=("user-content", "tool-result")) +@pytest.mark.parametrize( + ("block", "status", "error_type", "message"), + ( + ( + {"type": "future_block", "future_block": {"value": "ignore me"}}, + 400, + "invalid_request_error", + "future_block", + ), + ({"type": "image"}, 500, "api_error", "source"), + ), + ids=("unknown-block", "unrelated-counting-error"), +) +def test_local_messages_count_classifies_only_unsupported_content_as_an_invalid_request( + gateway: Gateway, nested: bool, block: JsonValue, status: int, error_type: str, message: str +) -> None: + content: Final = [{"type": "tool_result", "tool_use_id": "toolu_1", "content": [block]}] if nested else [block] + with gateway.scenario() as scenario: + model: Final = scenario.model(model="hosted_vllm/zai-org/GLM-5.3", api_base="http://127.0.0.1:9/v1") + response: Final = gateway.request( + "POST", "/v1/messages/count_tokens", {"model": model, "messages": [{"role": "user", "content": content}]} + ) + assert response.status_code == status, response.text + payload: Final = _JSON_OBJECT.validate_json(response.content) + error: Final = payload["error"] + assert payload["type"] == "error" and isinstance(error, dict), payload + assert error["type"] == error_type and message in str(error["message"]), payload + + @pytest.mark.parametrize( ("deployment", "missing"), ( diff --git a/tests/integration/providers/test_custom_video_provider_wire.py b/tests/integration/providers/test_custom_video_provider_wire.py new file mode 100644 index 00000000000..901d5a077a1 --- /dev/null +++ b/tests/integration/providers/test_custom_video_provider_wire.py @@ -0,0 +1,112 @@ +"""A provider registered in `custom_provider_map` with a `BaseVideoConfig` serves /v1/videos create, status and content.""" + +import json +import shutil +from contextlib import ExitStack +from pathlib import Path +from typing import Final + +from integration._support.client import Gateway +from integration._support.process import owned_proxy_process +from integration._support.wire import Reply, Request, Wire, wire_server + +_SUPPORT_MODULE: Final = Path(__file__).resolve().parents[1] / "_support" / "custom_video_provider.py" +_PROMPT: Final = "a cat walking across a bridge" +_MP4: Final = b"\x00\x00\x00\x18ftypmp42" + b"queue-video-bytes" + + +def _queue_backend(request: Request) -> Reply: + if (request.method, request.target) == ("POST", "/jobs"): + return Reply(body=json.dumps({"job_id": "job-7", "state": "queued"}).encode()) + if (request.method, request.target) == ("GET", "/jobs/job-7"): + return Reply(body=json.dumps({"job_id": "job-7", "state": "done", "seconds": 8}).encode()) + if (request.method, request.target) == ("GET", "/jobs/job-7/file"): + return Reply(body=_MP4, content_type="video/mp4") + return Reply(status=404, body=json.dumps({"error": f"no route {request.method} {request.target}"}).encode()) + + +def _write_config(path: Path, wire: Wire) -> Path: + shutil.copy(_SUPPORT_MODULE, path.parent / "custom_video_provider.py") + path.write_text( + json.dumps( + { + "model_list": [ + { + "model_name": "queue-video", + "litellm_params": { + "model": "queue-video/seedance-1", + "api_base": wire.url, + "api_key": "queue-secret", + }, + }, + { + "model_name": "image-only-video", + "litellm_params": { + "model": "image-only/seedance-1", + "api_base": wire.url, + "api_key": "unused", + }, + }, + ], + "litellm_settings": { + "custom_provider_map": [ + {"provider": "queue-video", "custom_handler": "custom_video_provider.VIDEO_CONFIG"}, + {"provider": "image-only", "custom_handler": "custom_video_provider.IMAGE_ONLY_HANDLER"}, + ] + }, + } + ), + encoding="utf-8", + ) + return path + + +def _lower_headers(request: Request) -> dict[str, str]: + return {name.lower(): value for name, value in request.headers.items()} + + +def test_custom_provider_map_base_video_config_serves_create_status_and_content( + gateway: Gateway, tmp_path: Path +) -> None: + with ExitStack() as stack: + wire: Final = stack.enter_context(wire_server(_queue_backend)) + config: Final = _write_config(tmp_path / "custom-video.yaml", wire) + owned: Final = stack.enter_context(owned_proxy_process(gateway, tmp_path, {}, config=config)) + + created: Final = owned.gateway.request( + "POST", "/v1/videos", {"model": "queue-video", "prompt": _PROMPT, "seconds": "8"} + ) + assert created.status_code == 200, created.text + created_body: Final = created.json() + assert created_body["status"] == "queued", created_body + video_id: Final = created_body["id"] + assert isinstance(video_id, str) and video_id.startswith("video_"), created_body + + submits: Final = wire.drain() + assert [(r.method, r.target) for r in submits] == [("POST", "/jobs")], submits + assert json.loads(submits[0].body) == {"model": "seedance-1", "prompt": _PROMPT, "seconds": "8"} + assert _lower_headers(submits[0])["x-api-key"] == "queue-secret" + + status: Final = owned.gateway.request("GET", f"/v1/videos/{video_id}") + assert status.status_code == 200, status.text + assert status.json()["status"] == "completed", status.text + + content: Final = owned.gateway.request("GET", f"/v1/videos/{video_id}/content") + assert content.status_code == 200, content.text + assert content.content == _MP4 + + polls: Final = wire.drain() + assert [(r.method, r.target) for r in polls] == [("GET", "/jobs/job-7"), ("GET", "/jobs/job-7/file")], polls + assert all(_lower_headers(r)["x-api-key"] == "queue-secret" for r in polls) + + +def test_custom_provider_without_a_video_config_is_refused_with_400(gateway: Gateway, tmp_path: Path) -> None: + with ExitStack() as stack: + wire: Final = stack.enter_context(wire_server(_queue_backend)) + config: Final = _write_config(tmp_path / "custom-video.yaml", wire) + owned: Final = stack.enter_context(owned_proxy_process(gateway, tmp_path, {}, config=config)) + + refused: Final = owned.gateway.request("POST", "/v1/videos", {"model": "image-only-video", "prompt": _PROMPT}) + assert refused.status_code == 400, refused.text + assert "video generation is not supported for image-only" in refused.text + assert wire.drain() == () diff --git a/tests/integration/providers/test_watsonx_passthrough_router_credentials_wire.py b/tests/integration/providers/test_watsonx_passthrough_router_credentials_wire.py new file mode 100644 index 00000000000..f20e6d35cf9 --- /dev/null +++ b/tests/integration/providers/test_watsonx_passthrough_router_credentials_wire.py @@ -0,0 +1,81 @@ +import json +import uuid +from pathlib import Path +from typing import Final + +import yaml +from integration._support.client import Gateway, JsonValue, object_value +from integration._support.process import owned_proxy +from integration._support.wire import Reply, Request, wire_server + +_DEPLOYMENT_KEY: Final = "synthetic-watsonx-deployment-key" +_IAM_TOKEN: Final = "synthetic-watsonx-iam-token" +_WATSONX_ENV: Final = ( + "WATSONX_API_KEY", + "WATSONX_APIKEY", + "WX_API_KEY", + "WATSONX_API_BASE", + "WATSONX_URL", + "WX_URL", + "WML_URL", + "WATSONX_TOKEN", + "WATSONX_ZENAPIKEY", +) + + +def _config(directory: Path, api_base: str) -> Path: + stock: Final = object_value(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())) + config: Final[dict[str, JsonValue]] = { + **stock, + "model_list": [ + { + "model_name": "granite", + "litellm_params": { + "model": "watsonx/ibm/granite-3-8b-instruct", + "api_key": _DEPLOYMENT_KEY, + "api_base": api_base, + "project_id": "synthetic-project", + "use_in_pass_through": True, + }, + } + ], + } + path: Final = directory / "watsonx-passthrough.yaml" + path.write_text(yaml.safe_dump(config)) + return path + + +def test_watsonx_passthrough_uses_the_flagged_deployment_key_and_base(gateway: Gateway, tmp_path: Path) -> None: + marker: Final = "watsonx-" + uuid.uuid4().hex + + def respond(request: Request) -> Reply: + if request.method == "POST" and request.target == "/identity/token": + assert f"apikey={_DEPLOYMENT_KEY}".encode() in request.body, request.body + return Reply(body=json.dumps({"access_token": _IAM_TOKEN, "expires_in": 3600}).encode()) + if request.method == "POST" and request.target.startswith("/ml/v1/text/generation?"): + assert request.headers["authorization"] == f"Bearer {_IAM_TOKEN}", request.headers + assert marker.encode() in request.body, request.body[:400] + return Reply( + body=json.dumps( + {"model_id": "ibm/granite-3-8b-instruct", "results": [{"generated_text": marker}]} + ).encode() + ) + return Reply(status=404) + + with wire_server(respond) as wire: + with owned_proxy( + gateway, + tmp_path, + {"WATSONX_IAM_URL": f"{wire.url}/identity/token"}, + config=_config(tmp_path, wire.url), + remove_environment=_WATSONX_ENV, + ) as candidate: + response: Final = candidate.request( + "POST", + "/watsonx/ml/v1/text/generation", + {"model_id": "ibm/granite-3-8b-instruct", "input": marker, "project_id": "synthetic-project"}, + ) + assert response.status_code == 200, response.text + assert response.json()["results"][0]["generated_text"] == marker + forwarded: Final = tuple((request.method, request.target.split("?")[0]) for request in wire.drain()) + assert forwarded == (("POST", "/identity/token"), ("POST", "/ml/v1/text/generation")), forwarded diff --git a/tests/integration/routing/test_payment_required_cooldown.py b/tests/integration/routing/test_payment_required_cooldown.py index 12190efee26..ca8242682a9 100644 --- a/tests/integration/routing/test_payment_required_cooldown.py +++ b/tests/integration/routing/test_payment_required_cooldown.py @@ -1,5 +1,6 @@ from __future__ import annotations +import itertools import json import re import uuid @@ -390,6 +391,79 @@ def test_a_router_allowed_fails_policy_cools_a_lone_402_deployment_past_its_coun assert markers_received(refusing) == (first.marker, second.marker) +def failing_reply(request: Request) -> Reply: + return model_list_reply() if is_model_info_probe(request) else anthropic_error(500, "upstream down") + + +def assert_failed(outcome: Outcome) -> None: + assert outcome.status == 500, (outcome.status, outcome.detail) + assert "upstream down" in outcome.detail, outcome.detail + + +@pytest.mark.timeout(900) +def test_a_router_allowed_fails_count_leaves_a_lone_deployment_warm_but_benches_one_with_a_sibling_or_fallback( + tmp_path: Path, +) -> None: + lone_group: Final = new_group() + pair_group: Final = new_group() + primary_group: Final = new_group() + backup_group: Final = new_group() + with ExitStack() as stack: + upstream: Final = stack.enter_context(gateway_from_environment()) + cache: Final = stack.enter_context(owned_redis(tmp_path)) + settings: Final = router_settings(cache, allowed_fails=1, fallbacks=[{primary_group: [backup_group]}]) + owned: Final = stack.enter_context( + owned_proxy_process( + upstream, + tmp_path, + {}, + config=cooldown_config(tmp_path, settings), + remove_environment=("PROMETHEUS_MULTIPROC_DIR",), + ) + ) + failing: Final = stack.enter_context(wire_server(failing_reply)) + healthy: Final = stack.enter_context(wire_server(healthy_reply)) + scenario: Final = stack.enter_context(owned.gateway.scenario()) + lone_id: Final = deployment(scenario, lone_group, anthropic_params(failing)) + failed: Final = tuple(raw_chat(owned.gateway, lone_group, new_marker()) for _ in range(4)) + for outcome in failed: + assert_failed(outcome) + assert not cooled(cache, lone_id) + assert markers_received(failing) == markers(failed) + + paired_id: Final = two_deployments(scenario, pair_group, failing, healthy) + history: Final = tuple( + itertools.takewhile( + lambda outcome: not cooled(cache, paired_id), + (raw_chat(owned.gateway, pair_group, new_marker()) for _ in range(40)), + ) + ) + assert cooled(cache, paired_id), [(outcome.status, outcome.detail[:80]) for outcome in history] + assert sum(outcome.status == 500 for outcome in history) == 1 + served: Final = tuple(raw_chat(owned.gateway, pair_group, new_marker()) for _ in range(10)) + assert_served(served, chat_identity) + assert markers_received(healthy) == (*markers(o for o in history if o.status == 200), *markers(served)) + + primary_id: Final = deployment(scenario, primary_group, anthropic_params(failing)) + deployment(scenario, backup_group, openai_params(healthy)) + failing.drain() + healthy.drain() + fell_back: Final = tuple( + itertools.takewhile( + lambda outcome: not cooled(cache, primary_id), + (raw_chat(owned.gateway, primary_group, new_marker()) for _ in range(40)), + ) + ) + assert cooled(cache, primary_id) and fell_back, [(o.status, o.detail[:80]) for o in fell_back] + assert_served(fell_back, chat_identity) + backed: Final = tuple(raw_chat(owned.gateway, primary_group, new_marker()) for _ in range(4)) + assert_served(backed, chat_identity) + received: Final = markers_received(failing) + assert received[: len(fell_back)] == markers(fell_back) + assert len(received) <= len(fell_back) + 1, received + assert {*markers(fell_back), *markers(backed)} <= set(markers_received(healthy)) + + def test_a_deployment_allowed_fails_policy_cools_a_lone_402_deployment(rig: CooldownRig) -> None: group: Final = new_group() policy: Final = {"allowed_fails_policy": {"BadRequestErrorAllowedFails": 0}} diff --git a/tests/integration/sdk/test_aiohttp_session_rebuild_wire.py b/tests/integration/sdk/test_aiohttp_session_rebuild_wire.py index 1ec1f486bed..361f98cf943 100644 --- a/tests/integration/sdk/test_aiohttp_session_rebuild_wire.py +++ b/tests/integration/sdk/test_aiohttp_session_rebuild_wire.py @@ -11,11 +11,30 @@ from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from typing import Final import pytest -from pydantic import JsonValue, TypeAdapter +from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter CONFIGURED_KEEPALIVE_SECONDS: Final = 1 IDLE_SECONDS: Final = 2 +EVICTED_CLIENTS: Final = 30 RESPONSES: Final = TypeAdapter(list[dict[str, JsonValue]]) +CHAT_COMPLETION: Final = { + "id": "chatcmpl-scripted", + "object": "chat.completion", + "created": 1700000000, + "model": "gpt-5.4-nano", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "hi"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, +} + + +class EvictionReport(BaseModel): + model_config = ConfigDict(frozen=True) + + cache_size: int + cached_clients: int + established_connections: int + unclosed_messages: tuple[str, ...] + REBUILT_SESSION_EXCHANGE: Final = textwrap.dedent( """ @@ -38,6 +57,47 @@ REBUILT_SESSION_EXCHANGE: Final = textwrap.dedent( ) +EVICTED_CLIENTS_EXCHANGE: Final = textwrap.dedent( + """ + import asyncio, gc, json, sys + import psutil + import litellm + from litellm.llms.custom_httpx.aiohttp_transport import LiteLLMAiohttpTransport + + unclosed = [] + + def established(port: int) -> int: + connections = psutil.Process().net_connections("tcp") + return sum(1 for c in connections if c.status == "ESTABLISHED" and c.raddr and c.raddr.port == port) + + async def one_call(base_url: str, index: int) -> None: + await litellm.acompletion( + model="openai/gpt-5.4-nano", + api_key=f"sk-distinct-{index}", + api_base=base_url, + messages=[{"role": "user", "content": "hi"}], + ) + + async def main(base_url: str, port: int, evicted: int) -> None: + asyncio.get_running_loop().set_exception_handler(lambda loop, ctx: unclosed.append(ctx["message"])) + cache = litellm.in_memory_llm_clients_cache + for index in range(cache.max_size_in_memory + evicted): + await one_call(base_url, index) + gc.collect() + await asyncio.gather(*LiteLLMAiohttpTransport._background_close_tasks) + await asyncio.sleep(0) + print(json.dumps({ + "cache_size": cache.max_size_in_memory, + "cached_clients": len(cache.cache_dict), + "established_connections": established(port), + "unclosed_messages": unclosed, + })) + + asyncio.run(main(sys.argv[1], int(sys.argv[2]), int(sys.argv[3]))) + """ +) + + class _ConnectionCountingPeer(ThreadingHTTPServer): daemon_threads = True @@ -62,7 +122,10 @@ class _ConnectionHandler(BaseHTTPRequestHandler): def do_POST(self) -> None: self.rfile.read(int(self.headers["Content-Length"])) - body: Final = json.dumps({"connection": self.connection_number}).encode() + payload: Final = ( + CHAT_COMPLETION if self.path.endswith("/chat/completions") else {"connection": self.connection_number} + ) + body: Final = json.dumps(payload).encode() self.send_response(200) self.send_header("Content-Type", "application/json") self.send_header("Content-Length", str(len(body))) @@ -107,3 +170,32 @@ def test_rebuilt_shared_session_drops_idle_connection_after_configured_keepalive ) -> None: observed: Final = _rebuilt_session_exchange(connection_counting_peer) assert observed == [{"connection": 1}, {"connection": 2}], observed + + +def _evicted_clients_exchange(base_url: str) -> EvictionReport: + completed: Final = subprocess.run( + [ + sys.executable, + "-P", + "-c", + EVICTED_CLIENTS_EXCHANGE, + base_url, + base_url.rsplit(":", 1)[1], + str(EVICTED_CLIENTS), + ], + capture_output=True, + text=True, + timeout=300, + check=False, + ) + assert completed.returncode == 0, completed.stderr + return EvictionReport.model_validate_json(completed.stdout) + + +def test_sdk_clients_evicted_from_a_full_client_cache_leave_no_open_session_or_connection( + connection_counting_peer: str, +) -> None: + report: Final = _evicted_clients_exchange(connection_counting_peer) + assert report.cached_clients == report.cache_size, report + assert report.unclosed_messages == (), report + assert report.established_connections <= report.cache_size, report diff --git a/tests/integration/sdk/test_bedrock_native_guardrail_shared_metadata_sdk.py b/tests/integration/sdk/test_bedrock_native_guardrail_shared_metadata_sdk.py new file mode 100644 index 00000000000..d0235310aaa --- /dev/null +++ b/tests/integration/sdk/test_bedrock_native_guardrail_shared_metadata_sdk.py @@ -0,0 +1,67 @@ +"""An SDK caller that reuses one ``metadata=`` dict across ``litellm.completion`` calls to a Bedrock deployment whose +``guardrailConfig`` has ``trace: enabled`` is billed, on each call, for that call's guardrail trace only, and the +caller's dict is left as the caller wrote it (LIT-10239)""" + +import json +from collections.abc import Mapping +from typing import Final + +import pytest +from integration._support.client import eventually +from integration._support.wire import Reply, Request, wire_server +from pydantic import JsonValue + +import litellm + +_MODEL: Final = "bedrock/anthropic.claude-sonnet-4-5-20250929-v1:0" +_PROMPT_TOKENS: Final = 5 +_COMPLETION_TOKENS: Final = 1 +_USAGE: Final[Mapping[str, int]] = {"contentPolicyUnits": 1, "topicPolicyUnits": 1} +_GUARDRAIL_IDS: Final = ("guardrail-one", "guardrail-two") + + +def _bedrock(request: Request) -> Reply: + """Bedrock's Converse runtime: a one-token answer and the trace of the guardrail the request named.""" + guardrail_id: Final = json.loads(request.body)["guardrailConfig"]["guardrailIdentifier"] + body: Final[dict[str, JsonValue]] = { + "output": {"message": {"role": "assistant", "content": [{"text": "ok"}]}}, + "stopReason": "end_turn", + "usage": {"inputTokens": _PROMPT_TOKENS, "outputTokens": _COMPLETION_TOKENS, "totalTokens": 6}, + "metrics": {"latencyMs": 3}, + "trace": {"guardrail": {"inputAssessment": {guardrail_id: {"invocationMetrics": {"usage": dict(_USAGE)}}}}}, + } + return Reply(body=json.dumps(body).encode()) + + +def _guardrail_cost() -> float: + rates: Final = litellm.model_cost["bedrock/guardrails"]["guardrail_cost_per_unit"] + return sum(units * rates[counter] for counter, units in _USAGE.items()) + + +def _llm_cost() -> float: + prices: Final = litellm.model_cost[_MODEL.removeprefix("bedrock/")] + return _PROMPT_TOKENS * prices["input_cost_per_token"] + _COMPLETION_TOKENS * prices["output_cost_per_token"] + + +def test_each_call_sharing_one_metadata_dict_is_billed_its_own_trace_only() -> None: + billed: list[float] = [] # mutable-ok: the success callback appends what each call was billed + litellm.success_callback.append( + lambda kwargs, response, start, end: billed.append(kwargs["standard_logging_object"]["response_cost"]) + ) + shared_metadata: Final = {"tag": "x"} + with wire_server(_bedrock) as wire: + for guardrail_id in _GUARDRAIL_IDS: + litellm.completion( + model=_MODEL, + messages=[{"role": "user", "content": "hello"}], + api_base=wire.url, + aws_access_key_id="AKIASCRIPTED", + aws_secret_access_key="scripted", + aws_region_name="us-east-1", + guardrailConfig={"guardrailIdentifier": guardrail_id, "guardrailVersion": "1", "trace": "enabled"}, + metadata=shared_metadata, + ) + costs: Final = eventually(lambda: tuple(billed), lambda seen: len(seen) == len(_GUARDRAIL_IDS), seconds=30) + + assert costs == pytest.approx((_llm_cost() + _guardrail_cost(),) * len(_GUARDRAIL_IDS)) + assert shared_metadata == {"tag": "x"} diff --git a/tests/integration/sdk/test_router_cooldown_fallback_sdk.py b/tests/integration/sdk/test_router_cooldown_fallback_sdk.py index 747adeee41d..0a5bfd339fc 100644 --- a/tests/integration/sdk/test_router_cooldown_fallback_sdk.py +++ b/tests/integration/sdk/test_router_cooldown_fallback_sdk.py @@ -12,6 +12,7 @@ from pydantic import JsonValue, TypeAdapter import litellm from litellm import Router +from litellm.router_utils.cooldown_handlers import async_get_cooldown_deployments from litellm.types.router import DeploymentTypedDict _MODEL: Final = "gpt-5.4" @@ -118,3 +119,30 @@ async def test_cooldown_skips_only_the_failed_duplicate_model_deployment() -> No "/primary/chat/completions", "/backup/chat/completions", ) + + +@pytest.mark.parametrize( + "fallbacks", + [[{"primary": "backup"}], [{"*": "backup"}]], + ids=["string_fallback_target", "string_catch_all_target"], +) +@pytest.mark.asyncio +async def test_string_valued_fallbacks_still_cool_down_the_failing_primary(fallbacks: list[dict[str, object]]) -> None: + with wire_server(_peer) as wire: + router: Final = Router( + model_list=[ + _deployment("primary", f"{wire.url}/primary", "primary-deployment"), + _deployment("backup", f"{wire.url}/backup", "backup-deployment"), + ], + allowed_fails=0, + num_retries=0, + ) + router.fallbacks = fallbacks + with pytest.raises(litellm.AuthenticationError): + await router.acompletion( + model="primary", + messages=[{"role": "user", "content": f"cooldown {uuid.uuid4().hex}"}], + ) + cooled: Final = await async_get_cooldown_deployments(litellm_router_instance=router, parent_otel_span=None) + + assert tuple(cooled) == ("primary-deployment",) diff --git a/tests/integration/spend/test_unmanaged_vertex_batch_deployment_routing_wire.py b/tests/integration/spend/test_unmanaged_vertex_batch_deployment_routing_wire.py new file mode 100644 index 00000000000..2d2d279e18e --- /dev/null +++ b/tests/integration/spend/test_unmanaged_vertex_batch_deployment_routing_wire.py @@ -0,0 +1,189 @@ +"""Two Vertex deployments share a base model but live in different projects with different +buckets. An unmanaged gs:// batch created on the second one must be polled through that +deployment (LIT-10326). Before the fix the cost poller matched the base model name in the +gs:// path to the first deployment in the list and asked the wrong project for the job.""" + +import json +import re +from collections.abc import Callable +from contextlib import ExitStack +from pathlib import Path +from typing import Final +from uuid import uuid4 + +import pytest +import yaml +from integration._support.client import JSON_OBJECT, Gateway, eventually, object_value, string_value +from integration._support.process import owned_proxy_process +from integration._support.vertex import service_account_json +from integration._support.wire import Reply, Request, wire_server +from pydantic import JsonValue + +LOCATION: Final = "us-central1" +BASE_MODEL: Final = "gemini-2.5-flash" +TEAMS: Final = ("a", "b") +_JOB_PATH: Final = re.compile( + r"^/v1/projects/(?P[^/]+)/locations/(?P[^/]+)/batchPredictionJobs(?:/(?P\d+))?$" +) + + +def _project(team: str) -> str: + return f"scripted-project-{team}" + + +def _deployment(team: str) -> str: + return f"{BASE_MODEL}-team-{team}" + + +def _input_file(team: str) -> str: + return f"gs://scripted-bucket-{team}/litellm-vertex-files/publishers/google/models/{BASE_MODEL}/{uuid4().hex}/input.jsonl" + + +def _job_number(team: str) -> str: + return f"{TEAMS.index(team) + 1}{uuid4().int % 10**17:017d}" + + +def _vertex_peer(jobs: dict[str, tuple[str, str]]) -> Callable[[Request], Reply]: + """A Vertex double that remembers which project each job was created in. A retrieve + from any other project answers the 404 real Vertex gives (the job does not exist + there), a retrieve from the right one reports the job still running.""" + + def respond(request: Request) -> Reply: + path: Final = _JOB_PATH.match(request.target.split("?", 1)[0]) + assert path is not None, request.target + assert request.headers.get("authorization") == "Bearer scripted-token", request.headers + project, location, job = path.group("project"), path.group("location"), path.group("job") + if request.method == "POST": + team: Final = project.rsplit("-", 1)[-1] + uri: Final = re.search(r'"(gs://[^"]+)"', request.body.decode()) + assert uri is not None and uri.group(1).startswith(f"gs://scripted-bucket-{team}/"), (project, request.body) + created, created_uri = jobs.setdefault(project, (_job_number(team), uri.group(1))) + return Reply(body=json.dumps(_job(project, location, created, created_uri, "JOB_STATE_PENDING")).encode()) + assert request.method == "GET" and job is not None, (request.method, request.target) + known: Final = jobs.get(project) + if known is None or known[0] != job: + return Reply( + status=404, + body=json.dumps( + { + "error": { + "code": 404, + "message": f"The BatchPredictionJob does not exist: {job}", + "status": "NOT_FOUND", + } + } + ).encode(), + ) + return Reply(body=json.dumps(_job(project, location, job, known[1], "JOB_STATE_RUNNING")).encode()) + + return respond + + +def _job(project: str, location: str, job: str, input_uri: str, state: str) -> dict[str, JsonValue]: + return { + "name": f"projects/{project}/locations/{location}/batchPredictionJobs/{job}", + "displayName": "litellm-vertex-batch-scripted", + "model": f"publishers/google/models/{BASE_MODEL}", + "inputConfig": {"gcsSource": {"uris": [input_uri]}, "instancesFormat": "jsonl"}, + "outputConfig": { + "predictionsFormat": "jsonl", + "gcsDestination": {"outputUriPrefix": input_uri.rsplit("/", 1)[0]}, + }, + "state": state, + "createTime": "2026-10-30T20:00:00.000000Z", + "updateTime": "2026-10-30T20:00:00.000000Z", + } + + +def _owned_config(path: Path, gateway: Gateway, wire_url: str) -> Path: + config: Final = JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())) + path.write_text( + yaml.safe_dump( + { + **config, + "model_list": [ + { + "model_name": _deployment(team), + "litellm_params": { + "model": f"vertex_ai/{BASE_MODEL}", + "api_base": wire_url, + "vertex_project": _project(team), + "vertex_location": LOCATION, + "vertex_credentials": service_account_json( + _project(team), gateway.upstream_url.rstrip("/") + ), + "gcs_bucket_name": f"scripted-bucket-{team}", + }, + } + for team in TEAMS + ], + "general_settings": { + **object_value(config.get("general_settings") or {}), + "track_unmanaged_batch_cost": True, + }, + "litellm_settings": { + **object_value(config.get("litellm_settings") or {}), + "enable_loadbalancing_on_batch_endpoints": True, + "user_url_allowed_hosts": ["127.0.0.1"], + }, + } + ) + ) + return path + + +def _create_batch(gateway: Gateway, key: str, team: str) -> str: + response: Final = gateway.request( + "POST", + "/v1/batches", + { + "input_file_id": _input_file(team), + "endpoint": "/v1/chat/completions", + "completion_window": "24h", + "model": _deployment(team), + }, + key=key, + ) + assert response.status_code == 200, response.text + return string_value(JSON_OBJECT.validate_json(response.content)["id"]) + + +def _retrieve_targets(requests: tuple[Request, ...], job: str) -> tuple[str, ...]: + return tuple( + request.target for request in requests if request.method == "GET" and request.target.endswith(f"/{job}") + ) + + +@pytest.mark.timeout(240) +def test_unmanaged_batch_on_the_second_deployment_sharing_a_base_model_is_polled_through_its_own_project( + gateway: Gateway, tmp_path: Path +) -> None: + jobs: Final[dict[str, tuple[str, str]]] = {} # mutable-ok: the wire double records which project created each job + with ExitStack() as stack: + wire: Final = stack.enter_context(wire_server(_vertex_peer(jobs))) + config: Final = _owned_config(tmp_path / "two-vertex-teams.yaml", gateway, wire.url) + owned: Final = stack.enter_context(owned_proxy_process(gateway, tmp_path, {}, config=config)) + scenario: Final = stack.enter_context(owned.gateway.scenario()) + key: Final = scenario.key(models=[_deployment(team) for team in TEAMS]) + + batch_a: Final = _create_batch(owned.gateway, key, "a") + batch_b: Final = _create_batch(owned.gateway, key, "b") + assert (batch_a, batch_b) == (jobs[_project("a")][0], jobs[_project("b")][0]), (batch_a, batch_b, jobs) + + seen: Final[ + list[Request] + ] = [] # mutable-ok: drain empties the wire queue, so samples accumulate across retries + + def polled_so_far() -> tuple[Request, ...]: + seen.extend(wire.drain()) + return tuple(seen) + + polled: Final = eventually( + polled_so_far, + lambda requests: bool(_retrieve_targets(requests, batch_b)) and bool(_retrieve_targets(requests, batch_a)), + seconds=90, + ) + for team, batch in (("a", batch_a), ("b", batch_b)): + assert all( + target.startswith(f"/v1/projects/{_project(team)}/") for target in _retrieve_targets(polled, batch) + ), (team, _retrieve_targets(polled, batch)) diff --git a/tests/unit/caching/_event_loop_reuse.py b/tests/unit/caching/_event_loop_reuse.py new file mode 100644 index 00000000000..97cf8fe31f0 --- /dev/null +++ b/tests/unit/caching/_event_loop_reuse.py @@ -0,0 +1,20 @@ +from collections.abc import Callable +from typing import Final, TypeVar + +import pytest + +LoopResult = TypeVar("LoopResult") + + +def runs_until_a_loop_address_repeats( + run_on_new_loop: Callable[[], tuple[int, LoopResult]], + earlier: tuple[tuple[int, LoopResult], ...] = (), + limit: int = 50, +) -> tuple[tuple[int, LoopResult], ...]: + if len(earlier) == limit: + pytest.skip(f"none of {limit} collected loops had its address handed to a later loop") + latest: Final = run_on_new_loop() + runs: Final = (*earlier, latest) + if latest[0] in {loop_id for loop_id, _ in earlier}: + return runs + return runs_until_a_loop_address_repeats(run_on_new_loop, runs, limit) diff --git a/tests/unit/caching/test_evicted_client_closer.py b/tests/unit/caching/test_evicted_client_closer.py index 7679e276621..936eb37a733 100644 --- a/tests/unit/caching/test_evicted_client_closer.py +++ b/tests/unit/caching/test_evicted_client_closer.py @@ -10,6 +10,7 @@ never closed, because litellm does not own its lifecycle. import asyncio import gc import weakref +from typing import Final from unittest.mock import AsyncMock import httpx @@ -18,6 +19,7 @@ from redis.asyncio import ConnectionPool, Redis from litellm.caching.evicted_client_closer import EvictedClientCloser from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler +from tests.unit.caching._event_loop_reuse import runs_until_a_loop_address_repeats class FakeClock: @@ -41,6 +43,14 @@ class AsyncClient: self.closed = True +class RecordingAsyncClient: + def __init__(self) -> None: + self.close_calls = 0 + + async def aclose(self) -> None: + self.close_calls += 1 + + class SyncClient: def __init__(self) -> None: self.closed = False @@ -266,6 +276,100 @@ def test_sync_client_evicted_outside_an_event_loop_is_still_closed(): assert closer.pending_count == 0 +def test_async_client_from_collected_loop_is_not_closed_by_fresh_loops() -> None: + clock: Final = FakeClock() + closer: Final = EvictedClientCloser(grace_seconds=0, clock=clock) + + async def reap_then_schedule(client: RecordingAsyncClient) -> None: + closer.reap() + await asyncio.sleep(0) + closer.schedule(client) + + def run_on_new_loop() -> tuple[int, RecordingAsyncClient]: + client: Final = RecordingAsyncClient() + closer.mark_owned(client) + loop: Final = asyncio.new_event_loop() + loop_id: Final = id(loop) + try: + loop.run_until_complete(reap_then_schedule(client)) + finally: + loop.close() + del loop + gc.collect() + return loop_id, client + + runs: Final = runs_until_a_loop_address_repeats(run_on_new_loop) + + assert all(client.close_calls == 0 for _, client in runs) + + +def test_async_client_queued_on_a_live_loop_is_closed_only_by_that_loop() -> None: + clock: Final = FakeClock() + closer: Final = EvictedClientCloser(grace_seconds=0, clock=clock) + client: Final = RecordingAsyncClient() + closer.mark_owned(client) + owner_loop: Final = asyncio.new_event_loop() + other_loop: Final = asyncio.new_event_loop() + + async def schedule() -> None: + closer.schedule(client) + + async def reap() -> None: + closer.reap() + await asyncio.sleep(0) + + try: + owner_loop.run_until_complete(schedule()) + other_loop.run_until_complete(reap()) + closes_after_other_loop_reaped: Final = client.close_calls + pending_after_other_loop_reaped: Final = closer.pending_count + owner_loop.run_until_complete(reap()) + finally: + other_loop.close() + owner_loop.close() + + assert closes_after_other_loop_reaped == 0 + assert pending_after_other_loop_reaped == 1 + assert client.close_calls == 1 + assert closer.pending_count == 0 + + +def test_closes_queued_on_a_collected_loop_free_capacity() -> None: + clock: Final = FakeClock() + closer: Final = EvictedClientCloser(grace_seconds=0, max_pending=1, clock=clock) + first: Final = RecordingAsyncClient() + closer.mark_owned(first) + loop_a: Final = asyncio.new_event_loop() + + async def schedule_first() -> None: + closer.schedule(first) + + try: + loop_a.run_until_complete(schedule_first()) + finally: + loop_a.close() + del loop_a + gc.collect() + + second: Final = RecordingAsyncClient() + closer.mark_owned(second) + loop_b: Final = asyncio.new_event_loop() + + async def schedule_and_reap_second() -> None: + closer.schedule(second) + closer.reap() + await asyncio.sleep(0) + + try: + loop_b.run_until_complete(schedule_and_reap_second()) + finally: + loop_b.close() + + assert first.close_calls == 0 + assert second.close_calls == 1 + assert closer.pending_count == 0 + + @pytest.mark.asyncio async def test_an_async_client_waits_for_a_loop_rather_than_being_dropped(): clock = FakeClock() @@ -289,8 +393,7 @@ async def test_an_async_client_waits_for_a_loop_rather_than_being_dropped(): @pytest.mark.asyncio -async def test_a_client_evicted_on_another_event_loop_is_left_alone(): - """Closing a client bound to a different loop would schedule work on that loop.""" +async def test_a_client_evicted_on_a_collected_loop_is_dropped(): clock = FakeClock() closer = make_closer(clock) client = AsyncClient() @@ -310,7 +413,7 @@ async def test_a_client_evicted_on_another_event_loop_is_left_alone(): await asyncio.sleep(0.05) assert client.closed is False - assert closer.pending_count == 1 + assert closer.pending_count == 0 @pytest.mark.asyncio diff --git a/tests/unit/caching/test_llm_caching_handler.py b/tests/unit/caching/test_llm_caching_handler.py index 0c0fe8ca3c5..465ba2f9150 100644 --- a/tests/unit/caching/test_llm_caching_handler.py +++ b/tests/unit/caching/test_llm_caching_handler.py @@ -10,12 +10,16 @@ See: https://github.com/BerriAI/litellm/pull/22247 import asyncio, gc, httpx, importlib, litellm, os, threading, time, weakref import warnings +from typing import Final +import aiohttp import pytest +from openai import AsyncOpenAI from litellm.caching.evicted_client_closer import EvictedClientCloser from litellm.caching.llm_caching_handler import LLMClientCache +from litellm.llms.custom_httpx.aiohttp_transport import LiteLLMAiohttpTransport from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from litellm.llms.custom_httpx.http_handler import( @@ -26,6 +30,41 @@ from litellm.llms.custom_httpx.http_handler import( from litellm.types.utils import LlmProviders from litellm.utils import _invalidate_model_cost_lowercase_map from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome +from tests.unit.caching._event_loop_reuse import runs_until_a_loop_address_repeats + + +async def _get_and_set_cache(cache: LLMClientCache) -> object | None: + cached: Final = await cache.async_get_cache("client") + await cache.async_set_cache("client", object()) + return cached + + +def _cache_result_on_new_loop(cache: LLMClientCache) -> tuple[int, object | None]: + loop: Final = asyncio.new_event_loop() + loop_id: Final = id(loop) + try: + result: Final = loop.run_until_complete(_get_and_set_cache(cache)) + finally: + loop.close() + del loop + gc.collect() + return loop_id, result + + +async def _get_cache_key(cache: LLMClientCache) -> str: + return cache.update_cache_key_with_event_loop("client") + + +def _cache_key_on_new_loop(cache: LLMClientCache) -> tuple[int, str]: + loop: Final = asyncio.new_event_loop() + loop_id: Final = id(loop) + try: + key: Final = loop.run_until_complete(_get_cache_key(cache)) + finally: + loop.close() + del loop + gc.collect() + return loop_id, key class MockAsyncClient: @@ -48,6 +87,44 @@ class MockSyncClient: self.closed = True +def test_async_client_cache_does_not_reuse_clients_between_collected_loops() -> None: + cache: Final = LLMClientCache() + runs: Final = runs_until_a_loop_address_repeats(lambda: _cache_result_on_new_loop(cache)) + + assert all(result is None for _, result in runs) + + +def test_async_client_cache_keys_differ_for_every_collected_loop() -> None: + cache: Final = LLMClientCache() + runs: Final = runs_until_a_loop_address_repeats(lambda: _cache_key_on_new_loop(cache)) + keys: Final = tuple(key for _, key in runs) + + assert len(set(keys)) == len(keys) + + +class _LoopWithoutWeakReferences: + __slots__ = () + + +def _cache_key_while_running(cache: LLMClientCache, loop: _LoopWithoutWeakReferences) -> str: + asyncio._set_running_loop(loop) + try: + return cache.update_cache_key_with_event_loop("client") + finally: + asyncio._set_running_loop(None) + + +def test_cache_keys_stay_per_loop_on_loops_without_weak_references() -> None: + cache: Final = LLMClientCache() + loop_a: Final = _LoopWithoutWeakReferences() + loop_b: Final = _LoopWithoutWeakReferences() + + first_key_on_a: Final = _cache_key_while_running(cache, loop_a) + + assert _cache_key_while_running(cache, loop_a) == first_key_on_a + assert _cache_key_while_running(cache, loop_b) != first_key_on_a + + @pytest.mark.asyncio async def test_remove_key_does_not_close_async_client(): """ @@ -204,6 +281,36 @@ async def test_evicted_litellm_owned_client_is_closed_once_the_grace_window_elap assert client.closed is True +@pytest.mark.asyncio +async def test_an_evicted_owned_client_collected_inside_its_grace_window_still_closes_its_aiohttp_session(): + """ + The closer holds an evicted client weakly for the whole grace window. An SDK + process that exits, or a proxy that evicts and never calls the cache again, + collects the client first, so only its own finalizer can still close the + aiohttp session (LIT-9767). + """ + clock = _FakeClock() + cache = LLMClientCache( + max_size_in_memory=2, + evicted_client_closer=EvictedClientCloser(grace_seconds=900.0, clock=clock), + ) + session = aiohttp.ClientSession() + client = AsyncOpenAI( + api_key="sk-evicted", + http_client=httpx.AsyncClient(transport=LiteLLMAiohttpTransport(client=session)), + ) + cache.set_cache("client-key", client, litellm_owned_client=True, ttl=600) + + cache.ttl_dict = {key: 0 for key in cache.ttl_dict} + cache.expiration_heap = [(0, key) for _, key in cache.expiration_heap] + cache.evict_cache() + del client + gc.collect() + await asyncio.gather(*LiteLLMAiohttpTransport._background_close_tasks) + + assert session.closed, "the evicted client was collected inside the grace window and left its session open" + + @pytest.mark.asyncio async def test_evicted_caller_supplied_client_is_never_closed(): """litellm does not own a client the caller passed in, so it must stay open.""" diff --git a/tests/unit/enterprise/proxy/test_managed_files_hook.py b/tests/unit/enterprise/proxy/test_managed_files_hook.py index 01f20689e89..5384e34f59b 100644 --- a/tests/unit/enterprise/proxy/test_managed_files_hook.py +++ b/tests/unit/enterprise/proxy/test_managed_files_hook.py @@ -1420,6 +1420,8 @@ def _make_object_store_instance(): mock_prisma = MagicMock() mock_prisma.db.litellm_managedobjecttable.upsert = AsyncMock() mock_prisma.db.litellm_managedobjecttable.update_many = AsyncMock() + mock_prisma.db.litellm_managedobjecttable.find_first = AsyncMock(return_value=None) + mock_prisma.db.writer = mock_prisma.db return ( PROXY_LiteLLMManagedFiles( @@ -2144,3 +2146,117 @@ async def test_store_unified_file_id_caches_the_storage_location_the_db_row_gets assert (cached.storage_backend, cached.storage_url) == ("litellm_db", "litellm_db://content-row-1") create_data = file_table.upsert.await_args.kwargs["data"]["create"] assert (create_data["storage_backend"], create_data["storage_url"]) == ("litellm_db", "litellm_db://content-row-1") + + +@pytest.mark.asyncio +async def test_create_stores_the_creating_deployment_in_file_object(): + """Regression (LIT-10326): an unmanaged gs:// batch stores the raw Vertex job id, which + says nothing about which of two deployments sharing a base model created it. The create + must keep the router's deployment id inside the stored file_object so the cost poller + reads the result through that deployment instead of the first base-model match.""" + managed_files, mock_prisma = _make_object_store_instance() + creator: Final = UserAPIKeyAuth(api_key="sk-team-b", user_id="bob", team_id="team-b", parent_otel_span=None) + + await managed_files.store_unified_object_id( + unified_object_id="projects/123/locations/us/batchPredictionJobs/456", + file_object=_make_batch_response(status="validating"), + litellm_parent_otel_span=None, + model_object_id="projects/123/locations/us/batchPredictionJobs/456", + file_purpose="batch", + user_api_key_dict=creator, + persist_attribution=True, + create_if_missing=True, + creating_model_id="deploy-team-b", + ) + + created: Final = mock_prisma.db.litellm_managedobjecttable.upsert.await_args.kwargs["data"]["create"] + assert json.loads(created["file_object"])["model_id"] == "deploy-team-b" + + +@pytest.mark.asyncio +async def test_poll_refresh_keeps_the_stored_creating_deployment(): + """A later observe of the row (a GET through whichever deployment the client named) + rewrites file_object; it must carry the create-time model_id forward, not drop it or + replace it with the deployment that served the retrieve.""" + managed_files, mock_prisma = _make_object_store_instance() + stored_row: Final = MagicMock() + stored_row.file_object = json.dumps( + {**json.loads(_make_batch_response(status="validating").model_dump_json()), "model_id": "deploy-team-b"} + ) + mock_prisma.db.litellm_managedobjecttable.find_first = AsyncMock(return_value=stored_row) + observer: Final = UserAPIKeyAuth(api_key="sk-team-b", user_id="bob", team_id="team-b", parent_otel_span=None) + + await managed_files.store_unified_object_id( + unified_object_id="projects/123/locations/us/batchPredictionJobs/456", + file_object=_make_batch_response(status="completed"), + litellm_parent_otel_span=None, + model_object_id="projects/123/locations/us/batchPredictionJobs/456", + file_purpose="batch", + user_api_key_dict=observer, + persist_attribution=False, + create_if_missing=False, + ) + + written: Final = json.loads( + mock_prisma.db.litellm_managedobjecttable.update_many.await_args.kwargs["data"]["file_object"] + ) + assert written["model_id"] == "deploy-team-b" + assert written["status"] == "completed" + + +@pytest.mark.asyncio +async def test_poll_refresh_reads_the_stored_creating_deployment_from_the_writer(): + """With a read replica configured, a retrieve that lands before the replica has the create + must still find the create-time model_id, so the lookup goes to the writer. Prisma may hand + the Json column back already decoded, so the row here is a dict rather than JSON text.""" + managed_files, mock_prisma = _make_object_store_instance() + stored_row: Final = MagicMock() + stored_row.file_object = { + **json.loads(_make_batch_response(status="validating").model_dump_json()), + "model_id": "deploy-team-b", + } + mock_prisma.db.litellm_managedobjecttable.find_first = AsyncMock(return_value=None) + mock_prisma.db.writer = MagicMock() + mock_prisma.db.writer.litellm_managedobjecttable.find_first = AsyncMock(return_value=stored_row) + observer: Final = UserAPIKeyAuth(api_key="sk-team-b", user_id="bob", team_id="team-b", parent_otel_span=None) + + await managed_files.store_unified_object_id( + unified_object_id="projects/123/locations/us/batchPredictionJobs/456", + file_object=_make_batch_response(status="completed"), + litellm_parent_otel_span=None, + model_object_id="projects/123/locations/us/batchPredictionJobs/456", + file_purpose="batch", + user_api_key_dict=observer, + persist_attribution=False, + create_if_missing=False, + ) + + written: Final = json.loads( + mock_prisma.db.litellm_managedobjecttable.update_many.await_args.kwargs["data"]["file_object"] + ) + assert written["model_id"] == "deploy-team-b" + mock_prisma.db.litellm_managedobjecttable.find_first.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_batch_create_hook_passes_the_router_deployment_as_creating_model_id(): + """The post-call hook on a batch create hands the router's model_id to the object store; + a retrieve of the same batch hands None so it cannot overwrite the creator.""" + from litellm_enterprise.proxy.hooks.managed_files import BATCH_CREATE_HIDDEN_PARAM + + managed_files: Final = _make_managed_files_instance() + created: Final = _make_batch_response(model_id="deploy-team-b", output_file_id=None, status="validating") + created._hidden_params[BATCH_CREATE_HIDDEN_PARAM] = True + retrieved: Final = _make_batch_response(model_id="deploy-team-a", output_file_id=None, status="in_progress") + + with patch("litellm.proxy.proxy_server.llm_router", MagicMock()): + await managed_files.async_post_call_success_hook( + data={}, user_api_key_dict=_make_user_api_key_dict(), response=created + ) + await managed_files.async_post_call_success_hook( + data={}, user_api_key_dict=_make_user_api_key_dict(), response=retrieved + ) + + first, second = managed_files.store_unified_object_id.await_args_list + assert first.kwargs["creating_model_id"] == "deploy-team-b" + assert second.kwargs["creating_model_id"] is None diff --git a/tests/unit/integrations/test_agentops.py b/tests/unit/integrations/test_agentops.py index 5d4055ac75f..0c966534e6b 100644 --- a/tests/unit/integrations/test_agentops.py +++ b/tests/unit/integrations/test_agentops.py @@ -1,10 +1,14 @@ import os -from unittest.mock import MagicMock, patch +from typing import Final +from unittest.mock import patch import pytest +import respx - +import litellm +from litellm.caching.llm_caching_handler import LLMClientCache from litellm.integrations.agentops.agentops import AgentOps, AgentOpsConfig +from litellm.llms.custom_httpx.http_handler import get_httpx_client @pytest.fixture @@ -111,3 +115,17 @@ def test_agentops_initialization_no_auth(): assert agentops.resource_attributes["deployment.environment"] == "test_env" assert agentops.resource_attributes["telemetry.sdk.name"] == "agentops" assert "project.id" not in agentops.resource_attributes + + +def test_auth_fetch_leaves_the_shared_sync_client_open_for_other_callers( + monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter +) -> None: + monkeypatch.setattr(litellm, "in_memory_llm_clients_cache", LLMClientCache()) + auth_endpoint: Final = "https://agentops.example/v3/auth/token" + respx_mock.post(auth_endpoint).respond(json={"token": "jwt", "project_id": "project-1"}) + client_another_caller_is_using: Final = get_httpx_client().client + + agentops: Final = AgentOps(config=AgentOpsConfig(api_key="key", auth_endpoint=auth_endpoint)) + + assert agentops.resource_attributes["project.id"] == "project-1" + assert not client_another_caller_is_using.is_closed diff --git a/tests/unit/litellm_core_utils/llm_cost_calc/test_guardrail_cost.py b/tests/unit/litellm_core_utils/llm_cost_calc/test_guardrail_cost.py index aa2fc0b9a45..949f55b8a82 100644 --- a/tests/unit/litellm_core_utils/llm_cost_calc/test_guardrail_cost.py +++ b/tests/unit/litellm_core_utils/llm_cost_calc/test_guardrail_cost.py @@ -1,7 +1,11 @@ +import json +from typing import Final import pytest +from pydantic import JsonValue import litellm +from litellm.litellm_core_utils.llm_cost_calc import guardrail_cost from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import ( bedrock_guardrail_cost, bedrock_guardrail_cost_by_unit, @@ -230,3 +234,118 @@ def test_guardrail_information_cost_skips_malformed_entry_keeps_siblings(): ] assert guardrail_information_cost(entries) == pytest.approx(0.0003) assert guardrail_information_cost({"guardrail_cost": 0.5, "guardrail_cost_in_spend": "maybe"}) == 0.0 + + +_TRACE_USAGE_IN: Final = { + "topicPolicyUnits": 1, + "contentPolicyUnits": 1, + "wordPolicyUnits": 1, + "sensitiveInformationPolicyUnits": 0, + "contextualGroundingPolicyUnits": 0, + "automatedReasoningPolicies": 0, +} +_TRACE_USAGE_OUT: Final = { + **_TRACE_USAGE_IN, + "topicPolicyUnits": 2, + "contentPolicyUnits": 2, + "sensitiveInformationPolicyUnits": 2, +} +_TRACE_PRICING: Final = { + "bedrock/guardrails": { + "litellm_provider": "bedrock", + "mode": "guardrails", + "guardrail_cost_per_unit": { + "topicPolicyUnits": 0.00015, + "contentPolicyUnits": 0.00015, + "wordPolicyUnits": 0.0, + "sensitiveInformationPolicyUnits": 0.0001, + }, + } +} + + +def _converse_trace(action: str = "NONE") -> dict[str, JsonValue]: + return { + "guardrail": { + "inputAssessment": {"gid": {"invocationMetrics": {"usage": _TRACE_USAGE_IN}}}, + "outputAssessments": { + "gid": [ + { + "topicPolicy": {"topics": [{"name": "t", "type": "DENY", "action": action}]}, + "invocationMetrics": {"usage": _TRACE_USAGE_OUT}, + } + ] + }, + } + } + + +def _invoke_trace() -> dict[str, JsonValue]: + return { + "guardrail": { + "input": {"gid": {"invocationMetrics": {"usage": _TRACE_USAGE_IN}}}, + "outputs": [{"gid": {"invocationMetrics": {"usage": _TRACE_USAGE_OUT}}}], + } + } + + +@pytest.mark.parametrize("trace", [_converse_trace(), _invoke_trace()], ids=["converse", "invoke"]) +def test_native_trace_sums_every_assessment_and_prices_it_like_the_proxy_guardrail(monkeypatch, trace) -> None: + monkeypatch.setattr(litellm, "model_cost", _TRACE_PRICING) + entry: Final = guardrail_cost.native_bedrock_guardrail_information( + trace, aws_region_name=None, guardrail_identifier=None + ) + assert entry is not None + assert entry["guardrail_usage"] == { + "topicPolicyUnits": 3, + "contentPolicyUnits": 3, + "wordPolicyUnits": 2, + "sensitiveInformationPolicyUnits": 2, + "contextualGroundingPolicyUnits": 0, + } + assert entry["guardrail_cost"] == pytest.approx(3 * 0.00015 + 3 * 0.00015 + 2 * 0.0001) + assert entry["guardrail_cost_by_unit"]["contextualGroundingPolicyUnits"] is None + assert entry["guardrail_name"] == "gid" + assert entry["guardrail_provider"] == "bedrock" + assert entry["guardrail_status"] == "success" + assert guardrail_cost.guardrail_information_cost([entry]) == pytest.approx(entry["guardrail_cost"]) + + +def test_native_trace_marks_a_blocked_assessment_intervened_and_prefers_the_configured_identifier(monkeypatch) -> None: + monkeypatch.setattr(litellm, "model_cost", _TRACE_PRICING) + entry: Final = guardrail_cost.native_bedrock_guardrail_information( + _converse_trace("BLOCKED"), aws_region_name="us-east-1", guardrail_identifier="configured-id" + ) + assert entry is not None + assert entry["guardrail_status"] == "guardrail_intervened" + assert entry["guardrail_name"] == "configured-id" + + +def test_native_trace_without_usage_is_not_billed(monkeypatch) -> None: + monkeypatch.setattr(litellm, "model_cost", _TRACE_PRICING) + unpriced: Final = {"guardrail": {"inputAssessment": {"gid": {"topicPolicy": {"topics": []}}}}} + assert guardrail_cost.native_bedrock_guardrail_information(unpriced, None, None) is None + assert guardrail_cost.bedrock_guardrail_trace({"usage": {}}) is None + + +def test_with_guardrail_entry_records_the_same_trace_once() -> None: + entry: Final = {"guardrail_name": "gid", "guardrail_cost": 0.001} + metadata: dict = {"standard_logging_guardrail_information": {"guardrail_name": "hook"}} + guardrail_cost.with_guardrail_entry(metadata, entry) # pyright: ignore[reportArgumentType] # partial entry + guardrail_cost.with_guardrail_entry(metadata, entry) # pyright: ignore[reportArgumentType] # partial entry + assert metadata["standard_logging_guardrail_information"] == [{"guardrail_name": "hook"}, entry] + + +def test_native_trace_entry_redacts_the_pii_matches_before_they_reach_spend_logs(monkeypatch) -> None: + monkeypatch.setattr(litellm, "model_cost", _TRACE_PRICING) + trace: Final = _converse_trace() + trace["guardrail"]["inputAssessment"]["gid"]["sensitiveInformationPolicy"] = { + "piiEntities": [{"type": "US_SSN", "match": "123-45-6789", "action": "ANONYMIZED"}] + } + entry: Final = guardrail_cost.native_bedrock_guardrail_information( + trace, aws_region_name=None, guardrail_identifier=None + ) + assert entry is not None + assert entry["guardrail_status"] == "guardrail_intervened" + assert "123-45-6789" not in json.dumps(entry["guardrail_response"]) + assert entry["guardrail_usage"]["sensitiveInformationPolicyUnits"] == 2 diff --git a/tests/unit/litellm_core_utils/llm_cost_calc/test_utils.py b/tests/unit/litellm_core_utils/llm_cost_calc/test_utils.py index 142d49bfff1..1d2a5e4cc9b 100644 --- a/tests/unit/litellm_core_utils/llm_cost_calc/test_utils.py +++ b/tests/unit/litellm_core_utils/llm_cost_calc/test_utils.py @@ -11,7 +11,8 @@ import respx import litellm from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER -from litellm.types.utils import ModelResponse, PromptTokensDetailsWrapper, Usage +from litellm.types.llms.base import CachedTokensDetails +from litellm.types.utils import CompletionTokensDetailsWrapper, ModelResponse, PromptTokensDetailsWrapper, Usage TIER_MODEL: Final = "tier-priced-test-model" TIER_ROW: Final[Mapping[str, float]] = MappingProxyType( @@ -31,6 +32,204 @@ PROMPT_TOKENS: Final = 1000 CACHED_TOKENS: Final = 200 COMPLETION_TOKENS: Final = 500 TIER_API_BASE: Final = "https://tier-pricing.invalid/v1" +AUDIO_TIER_MODEL: Final = "tier-priced-gemini-audio-test-model" +AUDIO_TIER_ROW: Final[Mapping[str, float]] = MappingProxyType( + { + "input_cost_per_token": 4e-06, + "output_cost_per_token": 8e-06, + "cache_read_input_token_cost": 1e-06, + "input_cost_per_audio_token": 3e-06, + "cache_read_input_audio_token_cost": 2e-06, + "input_cost_per_audio_token_flex": 1e-06, + "input_cost_per_audio_token_priority": 2e-06, + "cache_read_input_audio_token_cost_flex": 4e-07, + "cache_read_input_audio_token_cost_priority": 8e-07, + } +) + + +@pytest.mark.parametrize("image_tokens", (60, 100)) +@pytest.mark.parametrize( + ("service_tier", "priority_rate", "image_rate"), + ((None, 11e-06, 7e-06), ("priority", 11e-06, 11e-06), ("flex", 11e-06, 7e-06), ("priority", 0.0, 0.0)), +) +def test_completion_cost_uses_priority_image_output_rate( + local_model_cost_map: None, service_tier: str | None, priority_rate: float, image_rate: float, image_tokens: int +) -> None: + model: Final = "gemini-synthetic-image-priority" + litellm.register_model( + { + model: { + "litellm_provider": "vertex_ai", + "mode": "chat", + "input_cost_per_token": 3e-06, + "output_cost_per_token": 5e-06, + "output_cost_per_image_token": 7e-06, + "output_cost_per_image_token_priority": priority_rate, + } + } + ) + response: Final = ModelResponse( + model=model, + usage=Usage( + prompt_tokens=20, + completion_tokens=100, + total_tokens=120, + completion_tokens_details=CompletionTokensDetailsWrapper(image_tokens=image_tokens), + ), + ) + cost: Final = litellm.completion_cost( + completion_response=response, model=model, custom_llm_provider="vertex_ai", service_tier=service_tier + ) + assert cost == pytest.approx(20 * 3e-06 + (100 - image_tokens) * 5e-06 + image_tokens * image_rate) + + +@pytest.mark.parametrize("prompt_tokens", (199_999, 200_000, 200_001)) +@pytest.mark.parametrize( + ("service_tier", "flat_rate", "context_rate"), + ((None, 7e-06, 11e-06), ("priority", 13e-06, 17e-06), ("flex", 19e-06, 23e-06)), +) +def test_completion_cost_uses_total_prompt_for_audio_context_tiers( + local_model_cost_map: None, prompt_tokens: int, service_tier: str | None, flat_rate: float, context_rate: float +) -> None: + model: Final = "gemini-synthetic-audio-context" + litellm.register_model( + { + model: { + "litellm_provider": "vertex_ai", + "mode": "chat", + "input_cost_per_token": 3e-06, + "output_cost_per_token": 5e-06, + "input_cost_per_audio_token": 7e-06, + "input_cost_per_audio_token_above_200k_tokens": 11e-06, + "input_cost_per_audio_token_priority": 13e-06, + "input_cost_per_audio_token_above_200k_tokens_priority": 17e-06, + "input_cost_per_audio_token_flex": 19e-06, + "input_cost_per_audio_token_above_200k_tokens_flex": 23e-06, + } + } + ) + response: Final = ModelResponse( + model=model, + usage=Usage( + prompt_tokens=prompt_tokens, + completion_tokens=10, + total_tokens=prompt_tokens + 10, + prompt_tokens_details=PromptTokensDetailsWrapper(audio_tokens=100), + ), + ) + cost: Final = litellm.completion_cost( + completion_response=response, model=model, custom_llm_provider="vertex_ai", service_tier=service_tier + ) + audio_rate: Final = context_rate if prompt_tokens > 200_000 else flat_rate + assert cost == pytest.approx((prompt_tokens - 100) * 3e-06 + 100 * audio_rate + 10 * 5e-06) + + +@pytest.mark.parametrize("prompt_tokens", (200_000, 200_001)) +def test_batch_cost_uses_audio_context_and_image_output_tiers(local_model_cost_map: None, prompt_tokens: int) -> None: + from litellm.cost_calculator import batch_cost_calculator + + model: Final = "gemini-synthetic-batch-audio-context" + litellm.register_model( + { + model: { + "litellm_provider": "vertex_ai", + "mode": "chat", + "input_cost_per_token": 6e-06, + "output_cost_per_token": 10e-06, + "input_cost_per_token_batches": 3e-06, + "output_cost_per_token_batches": 5e-06, + "input_cost_per_token_above_200k_tokens_batches": 13e-06, + "output_cost_per_token_above_200k_tokens_batches": 17e-06, + "output_cost_per_image_token_batches": 19e-06, + "input_cost_per_audio_token_batches": 7e-06, + "input_cost_per_audio_token_above_200k_tokens_batches": 11e-06, + } + } + ) + usage: Final = Usage( + prompt_tokens=prompt_tokens, + completion_tokens=20, + total_tokens=prompt_tokens + 20, + prompt_tokens_details=PromptTokensDetailsWrapper(audio_tokens=100), + completion_tokens_details=CompletionTokensDetailsWrapper(image_tokens=10), + ) + cost: Final = batch_cost_calculator(model=model, custom_llm_provider="vertex_ai", usage=usage) + audio_rate: Final = 11e-06 if prompt_tokens > 200_000 else 7e-06 + input_rate: Final = 13e-06 if prompt_tokens > 200_000 else 3e-06 + output_rate: Final = 17e-06 if prompt_tokens > 200_000 else 5e-06 + assert sum(cost) == pytest.approx( + (prompt_tokens - 100) * input_rate + 100 * audio_rate + 10 * output_rate + 10 * 19e-06 + ) + + +@pytest.mark.parametrize("prompt_tokens", (200_000, 200_001)) +def test_completion_cost_uses_flex_context_input_and_output_rates( + local_model_cost_map: None, prompt_tokens: int +) -> None: + model: Final = "gemini-synthetic-flex-context" + litellm.register_model( + { + model: { + "litellm_provider": "vertex_ai", + "mode": "chat", + "input_cost_per_token": 6e-06, + "output_cost_per_token": 10e-06, + "input_cost_per_token_flex": 3e-06, + "output_cost_per_token_flex": 5e-06, + "input_cost_per_token_above_200k_tokens_flex": 7e-06, + "output_cost_per_token_above_200k_tokens_flex": 11e-06, + } + } + ) + response: Final = ModelResponse( + model=model, usage=Usage(prompt_tokens=prompt_tokens, completion_tokens=10, total_tokens=prompt_tokens + 10) + ) + cost: Final = litellm.completion_cost( + completion_response=response, model=model, custom_llm_provider="vertex_ai", service_tier="flex" + ) + input_rate: Final = 7e-06 if prompt_tokens > 200_000 else 3e-06 + output_rate: Final = 11e-06 if prompt_tokens > 200_000 else 5e-06 + assert cost == pytest.approx(prompt_tokens * input_rate + 10 * output_rate) + + +@pytest.mark.parametrize( + ("audio_fields", "audio_rate"), + ( + ({}, 3e-06), + ({"input_cost_per_audio_token": 7e-06}, 7e-06), + ({"input_cost_per_audio_token_above_200k_tokens": 11e-06}, 11e-06), + ({"input_cost_per_audio_token_above_200k_tokens_priority": 0.0}, 0.0), + ), +) +def test_audio_context_rates_preserve_fallbacks_and_explicit_zero( + local_model_cost_map: None, audio_fields: Mapping[str, float], audio_rate: float +) -> None: + model: Final = "gemini-synthetic-audio-fallback" + litellm.register_model( + { + model: { + "litellm_provider": "vertex_ai", + "mode": "chat", + "input_cost_per_token": 3e-06, + "output_cost_per_token": 5e-06, + **audio_fields, + } + } + ) + response: Final = ModelResponse( + model=model, + usage=Usage( + prompt_tokens=200_001, + completion_tokens=10, + total_tokens=200_011, + prompt_tokens_details=PromptTokensDetailsWrapper(audio_tokens=100), + ), + ) + cost: Final = litellm.completion_cost( + completion_response=response, model=model, custom_llm_provider="vertex_ai", service_tier="priority" + ) + assert cost == pytest.approx(199_901 * 3e-06 + 100 * audio_rate + 10 * 5e-06) def _cost_at(prices: Mapping[str, float], column_suffix: str) -> float: @@ -78,6 +277,53 @@ def test_completion_cost_bills_the_price_columns_of_the_service_tier( assert cost == pytest.approx(_cost_at(TIER_ROW, column_suffix)) +@pytest.mark.parametrize( + ("service_tier", "audio_rate_key", "cache_audio_rate_key"), + ( + pytest.param("flex", "input_cost_per_audio_token_flex", "cache_read_input_audio_token_cost_flex"), + pytest.param("priority", "input_cost_per_audio_token_priority", "cache_read_input_audio_token_cost_priority"), + ), +) +def test_gemini_audio_and_cached_audio_use_service_tier_rates( + local_model_cost_map: None, service_tier: str, audio_rate_key: str, cache_audio_rate_key: str +) -> None: + litellm.register_model( + { + AUDIO_TIER_MODEL: { + "litellm_provider": "gemini", + "mode": "chat", + **AUDIO_TIER_ROW, + } + } + ) + usage: Final = Usage( + prompt_tokens=40, + completion_tokens=0, + total_tokens=40, + prompt_tokens_details=PromptTokensDetailsWrapper( + text_tokens=20, + audio_tokens=20, + cached_tokens=10, + cached_tokens_details=CachedTokensDetails(audio_tokens=10), + ), + ) + response: Final = ModelResponse(model=AUDIO_TIER_MODEL, usage=usage) + + cost: Final = litellm.completion_cost( + completion_response=response, + model=AUDIO_TIER_MODEL, + custom_llm_provider="gemini", + service_tier=service_tier, + ) + + expected: Final = ( + 20 * AUDIO_TIER_ROW["input_cost_per_token"] + + 10 * AUDIO_TIER_ROW[audio_rate_key] + + 10 * AUDIO_TIER_ROW[cache_audio_rate_key] + ) + assert cost == pytest.approx(expected) + + LONG_CONTEXT_TIER_MODEL: Final = "long-context-tier-priced-test-model" LONG_CONTEXT_TIER_ROW: Final[Mapping[str, float]] = MappingProxyType( { diff --git a/tests/unit/litellm_core_utils/llm_cost_calc/test_zero_cost_diagnostic.py b/tests/unit/litellm_core_utils/llm_cost_calc/test_zero_cost_diagnostic.py index 921dde8bf7b..bdc14c4934b 100644 --- a/tests/unit/litellm_core_utils/llm_cost_calc/test_zero_cost_diagnostic.py +++ b/tests/unit/litellm_core_utils/llm_cost_calc/test_zero_cost_diagnostic.py @@ -9,6 +9,7 @@ from litellm.litellm_core_utils.llm_cost_calc.zero_cost_diagnostic import ( used_pricing_keys, zero_cost_warning, ) +from litellm.types.llms.base import CachedTokensDetails from litellm.types.utils import CompletionTokensDetailsWrapper, PromptTokensDetailsWrapper, Usage PER_SECOND_ENTRY: Final = {"cost_per_second": 0.00042} @@ -135,6 +136,36 @@ def test_audio_tokens_need_the_audio_rates() -> None: assert diagnostic["missing_pricing_keys"] == ("input_cost_per_audio_token", "output_cost_per_audio_token") +def test_cached_audio_uses_its_cache_rate_in_zero_cost_diagnostic() -> None: + usage: Final = Usage( + prompt_tokens=20, + completion_tokens=0, + total_tokens=20, + prompt_tokens_details=PromptTokensDetailsWrapper( + text_tokens=0, + audio_tokens=20, + cached_tokens=20, + cached_tokens_details=CachedTokensDetails(audio_tokens=20), + ), + ) + pricing_entry: Final = { + "input_cost_per_token": 1e-6, + "input_cost_per_audio_token": 2e-6, + "cache_read_input_audio_token_cost": 0, + } + + assert used_pricing_keys(usage) == ("cache_read_input_audio_token_cost",) + assert ( + diagnose_zero_cost( + usage=usage, + pricing_model="gemini-audio", + pricing_entry=pricing_entry, + calculation_failed=False, + ) + is None + ) + + def test_warning_names_the_request_the_entry_the_missing_keys_and_the_counter() -> None: diagnostic = diagnose_zero_cost( usage=TEXT_USAGE, pricing_model="dep-1", pricing_entry=PER_SECOND_ENTRY, calculation_failed=False diff --git a/tests/unit/litellm_core_utils/test_chat_completion_agentic_loop.py b/tests/unit/litellm_core_utils/test_chat_completion_agentic_loop.py index e891706df13..3cbc3591009 100644 --- a/tests/unit/litellm_core_utils/test_chat_completion_agentic_loop.py +++ b/tests/unit/litellm_core_utils/test_chat_completion_agentic_loop.py @@ -25,13 +25,13 @@ from typing import Any, Dict, Final, List, Optional, Tuple from unittest.mock import AsyncMock, MagicMock, patch import pytest - +from openai._legacy_response import LegacyAPIResponse import litellm -from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.code_interpreter_interception.handler import ( CodeInterpreterInterceptionLogger, ) +from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.chat_completion_agentic_loop import ( maybe_run_chat_completion_agentic_loop, ) @@ -45,9 +45,9 @@ from litellm.types.integrations.custom_logger import ( AgenticLoopRequestPatch, ) from litellm.types.utils import ( + ChatCompletionMessageToolCall, Choices, Function, - ChatCompletionMessageToolCall, Message, ModelResponse, ) @@ -142,7 +142,7 @@ def _raw_response_for(model_response: ModelResponse) -> MagicMock: (an object exposing `.headers` and `.parse()` -> something with model_dump).""" parsed = MagicMock() parsed.model_dump.return_value = model_response.model_dump() - raw = MagicMock() + raw = MagicMock(spec=LegacyAPIResponse) raw.headers = {} raw.parse.return_value = parsed return raw @@ -171,7 +171,7 @@ async def test_internal_control_fields_never_leak_into_provider_body(restore_cal ] ) mock_client = MagicMock() - mock_client.chat.completions.with_raw_response.create = create + mock_client.post = create response = await litellm.acompletion( model="openai/gpt-4o-mini", diff --git a/tests/unit/litellm_core_utils/test_exception_mapping_utils.py b/tests/unit/litellm_core_utils/test_exception_mapping_utils.py index 02ebae6dcf6..b24348cfaa8 100644 --- a/tests/unit/litellm_core_utils/test_exception_mapping_utils.py +++ b/tests/unit/litellm_core_utils/test_exception_mapping_utils.py @@ -2462,7 +2462,6 @@ def test_openai_gateway_timeout_error(): Test that the OpenAI gateway timeout error is raised """ openai_client = OpenAI() - mapped_target = openai_client.chat.completions.with_raw_response # type: ignore def _return_exception(*args, **kwargs): @@ -2498,8 +2497,8 @@ def test_openai_gateway_timeout_error(): with pytest.raises(litellm.Timeout) as exc_info: with patch.object( - mapped_target, - "create", + openai_client, + "post", side_effect=_return_exception, ): litellm.completion( diff --git a/tests/unit/litellm_core_utils/test_litellm_logging.py b/tests/unit/litellm_core_utils/test_litellm_logging.py index f4a39f17bed..48cea556c1a 100644 --- a/tests/unit/litellm_core_utils/test_litellm_logging.py +++ b/tests/unit/litellm_core_utils/test_litellm_logging.py @@ -2444,6 +2444,44 @@ async def test_async_success_handler_runs_async_callbacks_in_the_post_response_p assert in_post_response_phase() is False, "the phase must not leak into the request task" +@pytest.mark.parametrize("stream", [True, False], ids=["streaming", "non_streaming"]) +def test_reused_logging_obj_does_not_report_the_previous_attempts_first_token(logging_obj, stream: bool): + """A fallback reuses the request's logging object but runs in its own window. + + When the first attempt already recorded a first-token time, the fallback must not + report it. Otherwise ``completion_start_time`` lands before the attempt's + ``start_time``, the callbacks see a negative TTFT, and latency-based routing caches + a negative sample that empties its buffer window. + """ + logging_obj.stream = stream + logging_obj.model_call_details["litellm_params"] = {} + logging_obj.litellm_params = {} + + first_attempt_start: Final = datetime.datetime(2026, 1, 1, 12, 0, 0) + logging_obj.update_completion_start_time(first_attempt_start + datetime.timedelta(seconds=1)) + fallback_start: Final = first_attempt_start + datetime.timedelta(seconds=600) + fallback_end: Final = fallback_start + datetime.timedelta(seconds=5) + + observed: dict[str, object] = {} + + class Recorder(CustomLogger): + def log_success_event(self, kwargs, response_obj, start_time, end_time) -> None: + observed["completion_start_time"] = kwargs.get("completion_start_time") + + with patch.object(logging_obj, "get_combined_callback_list", return_value=[Recorder()]): + logging_obj.success_handler( + result=_success_response(), + start_time=fallback_start, + end_time=fallback_end, + ) + + completion_start_time: Final = observed["completion_start_time"] + assert isinstance(completion_start_time, datetime.datetime) + assert completion_start_time >= fallback_start, ( + "the fallback attempt reported a first-token time that predates its own start" + ) + + def test_is_sync_litellm_request(): assert LitellmLogging.is_sync_litellm_request({}) is True assert LitellmLogging.is_sync_litellm_request({"acompletion": True}) is False @@ -12470,3 +12508,88 @@ def test_standard_logging_payload_reports_the_deployment_entry_the_request_was_b ) finally: litellm.model_cost.pop(deployment_id, None) + + +def test_record_native_bedrock_guardrail_bills_when_the_metadata_is_none() -> None: + import datetime + + from litellm.litellm_core_utils.litellm_logging import Logging + + logging_obj = Logging( + model="bedrock/anthropic.claude-sonnet-4-5-20250929-v1:0", + messages=[{"role": "user", "content": "hello"}], + stream=False, + call_type="completion", + start_time=datetime.datetime.now(), + litellm_call_id="call-none-metadata", + function_id="fn", + ) + logging_obj.model_call_details["litellm_params"] = {"metadata": None} + trace = {"guardrail": {"inputAssessment": {"gid": {"invocationMetrics": {"usage": {"contentPolicyUnits": 2}}}}}} + + logging_obj.record_native_bedrock_guardrail({"trace": trace}) + + (entry,) = logging_obj.model_call_details["litellm_params"]["metadata"]["standard_logging_guardrail_information"] + assert entry["guardrail_usage"] == {"contentPolicyUnits": 2} + assert entry["guardrail_cost"] > 0 + + +def test_record_native_bedrock_guardrail_bills_nothing_for_a_response_served_from_cache() -> None: + import datetime + + from litellm.litellm_core_utils.litellm_logging import Logging + + logging_obj = Logging( + model="bedrock/anthropic.claude-sonnet-4-5-20250929-v1:0", + messages=[{"role": "user", "content": "hello"}], + stream=False, + call_type="completion", + start_time=datetime.datetime.now(), + litellm_call_id="call-cache-hit", + function_id="fn", + ) + logging_obj.model_call_details["litellm_params"] = {"metadata": {}} + trace = {"guardrail": {"inputAssessment": {"gid": {"invocationMetrics": {"usage": {"contentPolicyUnits": 2}}}}}} + cached = litellm.ModelResponse(choices=[{"message": {"role": "assistant", "content": "hi"}}], model="x") + setattr(cached, "trace", trace) + cached._hidden_params = {"cache_hit": True} + + logging_obj.record_native_bedrock_guardrail(cached) + logging_obj.model_call_details["cache_hit"] = True + logging_obj.record_native_bedrock_guardrail({"trace": trace}) + + assert "standard_logging_guardrail_information" not in logging_obj.model_call_details["litellm_params"]["metadata"] + + +def test_record_native_bedrock_guardrail_bills_each_call_its_own_trace_when_calls_share_one_metadata_dict() -> None: + import datetime + + from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import guardrail_information_cost + + shared_metadata = {"tag": "x"} + traces = tuple( + {"guardrail": {"inputAssessment": {"gid": {"invocationMetrics": {"usage": {"contentPolicyUnits": units}}}}}} + for units in (1, 2) + ) + costs = [] + for call_index, trace in enumerate(traces): + logging_obj = Logging( + model="bedrock/anthropic.claude-sonnet-4-5-20250929-v1:0", + messages=[{"role": "user", "content": "hello"}], + stream=False, + call_type="completion", + start_time=datetime.datetime.now(), + litellm_call_id=f"call-shared-metadata-{call_index}", + function_id="fn", + ) + logging_obj.model_call_details["litellm_params"] = {"metadata": shared_metadata} + logging_obj.record_native_bedrock_guardrail({"trace": trace}) + recorded = logging_obj.model_call_details["litellm_params"]["metadata"][ + "standard_logging_guardrail_information" + ] + costs.append(guardrail_information_cost(recorded)) + + (first_cost, second_cost) = costs + assert second_cost == pytest.approx(2 * first_cost) + assert shared_metadata == {"tag": "x"} diff --git a/tests/unit/litellm_core_utils/test_redact_messages.py b/tests/unit/litellm_core_utils/test_redact_messages.py index d14f236f765..3c77a334b58 100644 --- a/tests/unit/litellm_core_utils/test_redact_messages.py +++ b/tests/unit/litellm_core_utils/test_redact_messages.py @@ -1798,3 +1798,4 @@ async def test_redaction_with_metadata_completion_api(): response = standard_logging_payload["response"] assert response["choices"][0]["message"]["content"] == "redacted-by-litellm" assert standard_logging_payload["messages"][0]["content"] == "redacted-by-litellm" + diff --git a/tests/unit/litellm_core_utils/test_streaming_chunk_builder_utils.py b/tests/unit/litellm_core_utils/test_streaming_chunk_builder_utils.py index f8cd0416b1c..a46c8bed4ad 100644 --- a/tests/unit/litellm_core_utils/test_streaming_chunk_builder_utils.py +++ b/tests/unit/litellm_core_utils/test_streaming_chunk_builder_utils.py @@ -8,7 +8,6 @@ import httpx import pytest import respx - import litellm from litellm import ChatCompletionUsageBlock, stream_chunk_builder from litellm.types.utils import GenericStreamingChunk @@ -19,6 +18,7 @@ from litellm.types.utils import ( ChatCompletionMessageToolCall, Delta, Function, + GenericStreamingChunk, ModelResponseStream, PromptTokensDetails, ServerToolUse, @@ -194,11 +194,7 @@ def test_get_combined_thinking_content_preserves_interleaved_blocks(): } ] ), - make_chunk( - thinking_blocks=[ - {"type": "thinking", "thinking": None, "signature": "sig_block1"} - ] - ), + make_chunk(thinking_blocks=[{"type": "thinking", "thinking": None, "signature": "sig_block1"}]), make_chunk( thinking_blocks=[ { @@ -216,16 +212,10 @@ def test_get_combined_thinking_content_preserves_interleaved_blocks(): } ] ), - make_chunk( - thinking_blocks=[ - {"type": "thinking", "thinking": None, "signature": "sig_block2"} - ] - ), + make_chunk(thinking_blocks=[{"type": "thinking", "thinking": None, "signature": "sig_block2"}]), ] - thinking_chunks = [ - chunk for chunk in chunks if chunk["choices"][0]["delta"].get("thinking_blocks") - ] + thinking_chunks = [chunk for chunk in chunks if chunk["choices"][0]["delta"].get("thinking_blocks")] processor = ChunkProcessor(chunks=chunks) result = processor.get_combined_thinking_content(thinking_chunks) @@ -295,9 +285,7 @@ def test_cache_read_input_tokens_retained(): prompt_tokens=11779, total_tokens=11784, completion_tokens_details=None, - prompt_tokens_details=PromptTokensDetails( - audio_tokens=None, cached_tokens=11775 - ), + prompt_tokens_details=PromptTokensDetails(audio_tokens=None, cached_tokens=11775), cache_creation_input_tokens=4, cache_read_input_tokens=11775, ), @@ -331,9 +319,7 @@ def test_cache_read_input_tokens_retained(): prompt_tokens=0, total_tokens=214, completion_tokens_details=None, - prompt_tokens_details=PromptTokensDetails( - audio_tokens=None, cached_tokens=0 - ), + prompt_tokens_details=PromptTokensDetails(audio_tokens=None, cached_tokens=0), cache_creation_input_tokens=0, cache_read_input_tokens=0, ), @@ -393,10 +379,7 @@ def test_streaming_preserves_anthropic_1hr_cache_creation_breakdown(): ) # Sanity: the delta event genuinely lacks the breakdown - this is the input # condition that used to defeat cost calc. - assert ( - getattr(message_delta_usage.prompt_tokens_details, "cache_creation_token_details", None) - is None - ) + assert getattr(message_delta_usage.prompt_tokens_details, "cache_creation_token_details", None) is None def _usage_chunk(usage, finish_reason): return ModelResponseStream( @@ -515,9 +498,7 @@ def test_cache_read_input_tokens_retained_genericstreamingchunk(): prompt_tokens=1234, total_tokens=1239, completion_tokens_details=None, - prompt_tokens_details=PromptTokensDetails( - audio_tokens=None, cached_tokens=543 - ).model_dump(), + prompt_tokens_details=PromptTokensDetails(audio_tokens=None, cached_tokens=543).model_dump(), ), index=2, ) @@ -534,6 +515,7 @@ def test_cache_read_input_tokens_retained_genericstreamingchunk(): assert usage.prompt_tokens_details.cached_tokens == 543 + def test_stream_chunk_builder_litellm_usage_chunks(): """ Validate ChunkProcessor.calculate_usage uses provided usage fields from streaming chunks @@ -607,9 +589,7 @@ def test_stream_chunk_builder_litellm_usage_chunks(): chunks = [chunk1, chunk2] processor = ChunkProcessor(chunks=chunks) - usage = processor.calculate_usage( - chunks=chunks, model="gemini/gemini-2.5-flash-lite", completion_output="" - ) + usage = processor.calculate_usage(chunks=chunks, model="gemini/gemini-2.5-flash-lite", completion_output="") assert usage.prompt_tokens == 50 assert usage.completion_tokens == 27 @@ -653,15 +633,11 @@ def test_calculate_usage_honors_openai_sdk_completion_usage_chunks(): provider_specific_fields=None, stream_options={"include_usage": True}, ) - usage_chunk.usage = CompletionUsage( - prompt_tokens=20, completion_tokens=60, total_tokens=80, cost=0.000704 - ) + usage_chunk.usage = CompletionUsage(prompt_tokens=20, completion_tokens=60, total_tokens=80, cost=0.000704) assert type(usage_chunk.usage) is CompletionUsage chunks = [content_chunk, usage_chunk] - usage = ChunkProcessor(chunks=chunks).calculate_usage( - chunks=chunks, model="mantle-claude", completion_output="" - ) + usage = ChunkProcessor(chunks=chunks).calculate_usage(chunks=chunks, model="mantle-claude", completion_output="") assert usage.prompt_tokens == 20 assert usage.completion_tokens == 60 @@ -684,9 +660,7 @@ def test_get_model_from_chunks_azure_model_router(): {"model": "gpt-4.1-nano-2025-04-14", "id": "chatcmpl-123", "choices": []}, ] - result = ChunkProcessor._get_model_from_chunks( - chunks=chunks, first_chunk_model="azure-model-router" - ) + result = ChunkProcessor._get_model_from_chunks(chunks=chunks, first_chunk_model="azure-model-router") # Should return the actual model, not the request model assert result == "gpt-4.1-nano-2025-04-14" @@ -697,9 +671,7 @@ def test_get_model_from_chunks_azure_model_router(): {"model": "gpt-4", "id": "chatcmpl-456", "choices": []}, ] - result_same = ChunkProcessor._get_model_from_chunks( - chunks=chunks_same_model, first_chunk_model="gpt-4" - ) + result_same = ChunkProcessor._get_model_from_chunks(chunks=chunks_same_model, first_chunk_model="gpt-4") # Should return the first chunk's model when all are the same assert result_same == "gpt-4" @@ -775,9 +747,7 @@ def test_stream_chunk_builder_anthropic_web_search(): chunks = [chunk1, chunk2] processor = ChunkProcessor(chunks=chunks) - usage = processor.calculate_usage( - chunks=chunks, model="claude-sonnet-4-5-20250929", completion_output="" - ) + usage = processor.calculate_usage(chunks=chunks, model="claude-sonnet-4-5-20250929", completion_output="") assert usage.prompt_tokens == 50 assert usage.completion_tokens == 27 @@ -929,15 +899,11 @@ def test_stream_chunk_builder_dict_snapshot_preserves_hidden_provider_fields(): ], ) chunk_dict = chunk.model_dump() - chunk_dict["_hidden_params"] = { - "provider_specific_fields": {"traffic_type": "default"} - } + chunk_dict["_hidden_params"] = {"provider_specific_fields": {"traffic_type": "default"}} response = stream_chunk_builder(chunks=[chunk_dict]) assert response is not None - assert ( - response._hidden_params["provider_specific_fields"]["traffic_type"] == "default" - ) + assert response._hidden_params["provider_specific_fields"]["traffic_type"] == "default" def test_stream_chunk_builder_propagates_vertex_ai_metadata_from_chunks(): @@ -982,10 +948,7 @@ def test_stream_chunk_builder_propagates_vertex_ai_metadata_from_chunks(): assert getattr(response, "vertex_ai_grounding_metadata") == grounding_metadata assert getattr(response, "vertex_ai_url_context_metadata") == url_context_metadata assert response._hidden_params["vertex_ai_grounding_metadata"] == grounding_metadata - assert ( - response._hidden_params["vertex_ai_url_context_metadata"] - == url_context_metadata - ) + assert response._hidden_params["vertex_ai_url_context_metadata"] == url_context_metadata dumped = response.model_dump() assert dumped["vertex_ai_grounding_metadata"] == grounding_metadata @@ -1032,9 +995,7 @@ def test_stream_chunk_builder_uses_assembled_model_for_provider_metadata(): def test_stream_chunk_builder_propagates_vertex_ai_safety_results(): """Assembled response must expose safety data under the non-streaming field name.""" - safety_ratings = [ - [{"category": "HARM_CATEGORY_HATE_SPEECH", "probability": "NEGLIGIBLE"}] - ] + safety_ratings = [[{"category": "HARM_CATEGORY_HATE_SPEECH", "probability": "NEGLIGIBLE"}]] chunk = ModelResponseStream( id="chatcmpl-vertex-safety", @@ -1076,18 +1037,12 @@ def test_stream_chunk_builder_propagates_vertex_ai_metadata_from_dict_chunks(): ) ], ).model_dump() - chunk_dict["_hidden_params"] = { - "vertex_ai_grounding_metadata": [{"webSearchQueries": ["test query"]}] - } + chunk_dict["_hidden_params"] = {"vertex_ai_grounding_metadata": [{"webSearchQueries": ["test query"]}]} response = stream_chunk_builder(chunks=[chunk_dict]) assert response is not None - assert getattr(response, "vertex_ai_grounding_metadata") == [ - {"webSearchQueries": ["test query"]} - ] - assert response.model_dump()["vertex_ai_grounding_metadata"] == [ - {"webSearchQueries": ["test query"]} - ] + assert getattr(response, "vertex_ai_grounding_metadata") == [{"webSearchQueries": ["test query"]}] + assert response.model_dump()["vertex_ai_grounding_metadata"] == [{"webSearchQueries": ["test query"]}] def test_cost_field_in_usage_chunks(): @@ -1096,29 +1051,21 @@ def test_cost_field_in_usage_chunks(): id="chatcmpl-1", created=1745513206, model="openrouter/claude", - choices=[ - StreamingChoices(finish_reason=None, index=0, delta=Delta(content="Hi")) - ], + choices=[StreamingChoices(finish_reason=None, index=0, delta=Delta(content="Hi"))], usage=chunk1_usage, ) - chunk2_usage = Usage( - completion_tokens=5, prompt_tokens=10, total_tokens=15, cost=0.00025 - ) + chunk2_usage = Usage(completion_tokens=5, prompt_tokens=10, total_tokens=15, cost=0.00025) chunk2 = ModelResponseStream( id="chatcmpl-1", created=1745513207, model="openrouter/claude", - choices=[ - StreamingChoices(finish_reason="stop", index=0, delta=Delta(content="")) - ], + choices=[StreamingChoices(finish_reason="stop", index=0, delta=Delta(content=""))], usage=chunk2_usage, ) processor = ChunkProcessor(chunks=[chunk1, chunk2]) - usage = processor.calculate_usage( - chunks=[chunk1, chunk2], model="openrouter/claude", completion_output="Hi" - ) + usage = processor.calculate_usage(chunks=[chunk1, chunk2], model="openrouter/claude", completion_output="Hi") assert hasattr(usage, "cost") assert usage.cost == 0.00025 @@ -1201,25 +1148,19 @@ def test_prompt_tokens_details_survive_later_usage_chunk_without_details(): id="chatcmpl-1", created=1745513206, model="openai/gpt-5.6-sol", - choices=[ - StreamingChoices(finish_reason=None, index=0, delta=Delta(content="Hi")) - ], + choices=[StreamingChoices(finish_reason=None, index=0, delta=Delta(content="Hi"))], usage=Usage( prompt_tokens=6017, completion_tokens=4, total_tokens=6021, - prompt_tokens_details=PromptTokensDetailsWrapper( - cached_tokens=6004, cache_write_tokens=10 - ), + prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=6004, cache_write_tokens=10), ), ) chunk_without_details = ModelResponseStream( id="chatcmpl-1", created=1745513207, model="openai/gpt-5.6-sol", - choices=[ - StreamingChoices(finish_reason="stop", index=0, delta=Delta(content="")) - ], + choices=[StreamingChoices(finish_reason="stop", index=0, delta=Delta(content=""))], usage=Usage(prompt_tokens=6017, completion_tokens=4, total_tokens=6021), ) @@ -1557,9 +1498,7 @@ def test_calculate_usage_fills_unknown_split_from_reasoning_estimate( assert usage.completion_tokens_details.text_tokens == expected_text_tokens -def _openai_chunk( - choices: Sequence[Mapping[str, object]], usage: Mapping[str, int] | None = None -) -> dict[str, object]: +def _openai_chunk(choices: Sequence[Mapping[str, object]], usage: Mapping[str, int] | None = None) -> dict[str, object]: base: Final = { "id": "chatcmpl-lit6552", "object": "chat.completion.chunk", @@ -1994,7 +1933,11 @@ def test_stream_chunk_builder_litellm_tool_call(respx_mock: respx.MockRouter) -> ) assert [tool_call.to_dict() for tool_call in response.choices[0].message.tool_calls] == [ - {"function": {"arguments": '{"city":"Boston"}', "name": "get_weather"}, "id": "call_weather", "type": "function"} + { + "function": {"arguments": '{"city":"Boston"}', "name": "get_weather"}, + "id": "call_weather", + "type": "function", + } ] assert (response.usage.prompt_tokens, response.usage.completion_tokens, response.usage.total_tokens) == (5, 3, 8) @@ -2016,7 +1959,9 @@ def test_stream_chunk_builder_litellm_tool_call_regular_message(respx_mock: resp def test_stream_chunk_builder_multiple_tool_calls() -> None: - def tool_call_chunk(*tool_calls: ChatCompletionDeltaToolCall, finish_reason: str | None = None) -> ModelResponseStream: + def tool_call_chunk( + *tool_calls: ChatCompletionDeltaToolCall, finish_reason: str | None = None + ) -> ModelResponseStream: return _legacy_stream_chunk("openai/gpt-4o", tool_calls=tool_calls, finish_reason=finish_reason) chunks: Final = ( @@ -2027,7 +1972,9 @@ def test_stream_chunk_builder_multiple_tool_calls() -> None: ), tool_call_chunk(ChatCompletionDeltaToolCall(index=0, function=Function(arguments='{"base": 3, '))), tool_call_chunk( - ChatCompletionDeltaToolCall(id="call_add", index=1, type="function", function=Function(name="add", arguments="")) + ChatCompletionDeltaToolCall( + id="call_add", index=1, type="function", function=Function(name="add", arguments="") + ) ), tool_call_chunk(ChatCompletionDeltaToolCall(index=0, function=Function(arguments='"exponent": 5}'))), tool_call_chunk( @@ -2041,7 +1988,11 @@ def test_stream_chunk_builder_multiple_tool_calls() -> None: assert response is not None assert response.choices[0].finish_reason == "tool_calls" assert [tool_call.to_dict() for tool_call in response.choices[0].message.tool_calls] == [ - {"function": {"arguments": '{"base": 3, "exponent": 5}', "name": "exponentiate"}, "id": "call_exp", "type": "function"}, + { + "function": {"arguments": '{"base": 3, "exponent": 5}', "name": "exponentiate"}, + "id": "call_exp", + "type": "function", + }, {"function": {"arguments": '{"x": 1, "y": 2}', "name": "add"}, "id": "call_add", "type": "function"}, ] @@ -2075,9 +2026,7 @@ def test_grok_bug() -> None: response: Final = stream_chunk_builder(chunks=chunks) assert response is not None - assert json.loads(response.choices[0].message.tool_calls[0].function.arguments) == { - "location": "Boston" - } + assert json.loads(response.choices[0].message.tool_calls[0].function.arguments) == {"location": "Boston"} def test_stream_chunk_builder_openai_audio_output_usage() -> None: @@ -2180,3 +2129,50 @@ def test_stream_chunk_builder_keeps_reasoning_items_on_text_only_replies(): assert response is not None assert response.choices[0].message.content == "Hello world" assert response.choices[0].message.reasoning_items == reasoning_items + + +def test_stream_chunk_builder_carries_the_bedrock_guardrail_trace_and_records_it_on_the_logging_object( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A Bedrock stream reports its guardrail trace on the last chunk only; the built response carries it like the + non-streaming response and the call bills it as guardrail cost.""" + from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices + + monkeypatch.setattr( + litellm, + "model_cost", + { + "bedrock/guardrails": { + "litellm_provider": "bedrock", + "mode": "guardrails", + "guardrail_cost_per_unit": {"topicPolicyUnits": 0.00015}, + } + }, + ) + trace = {"guardrail": {"inputAssessment": {"gid": {"invocationMetrics": {"usage": {"topicPolicyUnits": 2}}}}}} + chunks = [ + ModelResponseStream(id="s", model="m", choices=[StreamingChoices(index=0, delta=Delta(content="hi"))]), + ModelResponseStream( + id="s", + model="m", + choices=[StreamingChoices(index=0, delta=Delta(), finish_reason="stop")], + provider_specific_fields={"trace": trace}, + ), + ] + logging_obj = Logging( + model="bedrock/m", + messages=[{"role": "user", "content": "q"}], + stream=True, + call_type="acompletion", + start_time=None, + litellm_call_id="call", + function_id="fn", + ) + logging_obj.model_call_details["litellm_params"] = {"metadata": {}} + built = stream_chunk_builder(chunks=chunks, messages=[{"role": "user", "content": "q"}], logging_obj=logging_obj) + assert built is not None + assert getattr(built, "trace") == trace + (entry,) = logging_obj.model_call_details["litellm_params"]["metadata"]["standard_logging_guardrail_information"] + assert entry["guardrail_usage"] == {"topicPolicyUnits": 2} + assert entry["guardrail_cost"] == pytest.approx(2 * 0.00015) diff --git a/tests/unit/litellm_core_utils/test_streaming_handler.py b/tests/unit/litellm_core_utils/test_streaming_handler.py index 4347a6639bd..637811f6387 100644 --- a/tests/unit/litellm_core_utils/test_streaming_handler.py +++ b/tests/unit/litellm_core_utils/test_streaming_handler.py @@ -6621,3 +6621,102 @@ async def test_leaving_an_async_with_block_closes_the_provider_stream(): async with _wrapper_over(stream): pass assert stream.close_calls == 1 + + +def _final_chunk_with_dict_usage(choices) -> ModelResponseStream: + """A stream chunk whose ``usage`` is a plain dict. + + ``ModelResponseStream.__init__`` coerces a dict ``usage`` kwarg into + ``litellm.Usage``, so the dict is attached after construction — the same + shape a custom ``streaming_decoder`` (or any third-party iterator that + bypasses litellm chunk construction) delivers to ``chunk_creator``. + """ + chunk: Final = ModelResponseStream( + id="chatcmpl-dict-usage-test", + created=1754900000, + model="gpt-4o", + object="chat.completion.chunk", + choices=choices, + ) + chunk.usage = { # mutable-ok: wire-shaped usage payload under test + "prompt_tokens": 100, + "completion_tokens": 10, + "total_tokens": 110, + "prompt_tokens_details": {"cached_tokens": 80}, # mutable-ok: wire payload + "completion_tokens_details": {"reasoning_tokens": 5}, # mutable-ok: wire payload + "cache_creation_input_tokens": 42, + } + return chunk + + +def test_dispatch_provider_chunk_dict_usage_preserves_token_details( + logging_obj: Logging, +): + """The openai else-branch of ``_dispatch_provider_chunk`` handles three + usage shapes: ``Usage`` (passed through), ``BaseModel`` + (``Usage(**model_dump())``), and plain ``dict``. The dict arm used to + rebuild ``litellm.Usage`` from only prompt/completion/total tokens, + silently dropping ``prompt_tokens_details`` (cached_tokens), + ``completion_tokens_details`` (reasoning_tokens), and cache token counts + that both sibling arms preserve. + """ + wrapper: Final = CustomStreamWrapper( + completion_stream=None, + model="gpt-4o", + logging_obj=logging_obj, + custom_llm_provider="openai", + stream_options={"include_usage": True}, # mutable-ok: constructor payload + ) + chunk: Final = _final_chunk_with_dict_usage( + choices=[ + StreamingChoices(finish_reason="stop", index=0, delta=Delta(content=None)) + ] # mutable-ok: chunk payload + ) + model_response: Final = wrapper.model_response_creator() + wrapper._dispatch_provider_chunk( + chunk=chunk, + model_response=model_response, + completion_obj={"content": ""}, # mutable-ok: dispatch scratch dict + ) + + usage: Final = model_response.usage + assert isinstance(usage, Usage) + assert usage.prompt_tokens == 100 + assert usage.completion_tokens == 10 + assert usage.total_tokens == 110 + assert usage.prompt_tokens_details is not None, ( + "dict-usage arm dropped prompt_tokens_details (its Usage/BaseModel sibling arms preserve it)" + ) + assert usage.prompt_tokens_details.cached_tokens == 80 + assert usage.completion_tokens_details is not None + assert usage.completion_tokens_details.reasoning_tokens == 5 + assert getattr(usage, "cache_creation_input_tokens", None) == 42 + + +def test_chunk_creator_usage_only_chunk_with_dict_usage_keeps_details( + logging_obj: Logging, +): + """End-to-end through ``chunk_creator``: a usage-only final chunk (empty + ``choices``, ``stream_options.include_usage`` set) takes the early-return + path, so whatever the dict-usage arm wrote is exactly what the consumer + receives — there is no later aggregation to repair it. + """ + wrapper: Final = CustomStreamWrapper( + completion_stream=None, + model="gpt-4o", + logging_obj=logging_obj, + custom_llm_provider="openai", + stream_options={"include_usage": True}, # mutable-ok: constructor payload + ) + chunk: Final = _final_chunk_with_dict_usage(choices=[]) # mutable-ok: empty choices list + result: Final = wrapper.chunk_creator(chunk=chunk) + + assert result is not None + usage: Final = result.usage + assert isinstance(usage, Usage) + assert usage.prompt_tokens_details is not None, ( + "usage-only early return delivered a Usage stripped of prompt_tokens_details to the stream consumer" + ) + assert usage.prompt_tokens_details.cached_tokens == 80 + assert usage.completion_tokens_details is not None + assert usage.completion_tokens_details.reasoning_tokens == 5 diff --git a/tests/unit/litellm_core_utils/test_token_counter.py b/tests/unit/litellm_core_utils/test_token_counter.py index 84c139e867b..bba233e8cb9 100644 --- a/tests/unit/litellm_core_utils/test_token_counter.py +++ b/tests/unit/litellm_core_utils/test_token_counter.py @@ -1401,6 +1401,34 @@ def test_anthropic_image_block_with_empty_base64_data(): assert tokens > 0 +def test_token_counter_counts_anthropic_search_result_text() -> None: + from litellm.litellm_core_utils.token_counter import _count_content_list + + search_result: Final = { + "type": "search_result", + "source": "https://example.invalid/notes", + "title": "Project notes", + "content": [{"type": "text", "text": "The code word is 7f3a."}], + "citations": {"enabled": True}, + } + + assert _count_content_list(len, [search_result], False, None) == len( + "https://example.invalid/notesProject notesThe code word is 7f3a." + ) + + +@pytest.mark.parametrize("nested", (False, True)) +def test_unknown_content_type_remains_a_value_error_with_an_optional_fallback(nested: bool) -> None: + from litellm.litellm_core_utils.token_counter import _count_content_list + + unknown: Final = {"type": "future_block"} + content: Final = [{"type": "tool_result", "tool_use_id": "toolu_1", "content": [unknown]}] if nested else [unknown] + + with pytest.raises(ValueError, match="Invalid content item type: future_block"): + _count_content_list(len, content, False, None) + assert _count_content_list(len, content, False, 7) == 7 + + def test_anthropic_image_block_without_source_raises(): """An `image` block with no `source` raises, matching the OpenAI `image_url`-without-`url` behavior.""" from litellm.litellm_core_utils.token_counter import _count_content_list diff --git a/tests/unit/llms/azure/chat/test_azure_chat_o_series_transformation.py b/tests/unit/llms/azure/chat/test_azure_chat_o_series_transformation.py index 3630ef1e354..3a463ad4eee 100644 --- a/tests/unit/llms/azure/chat/test_azure_chat_o_series_transformation.py +++ b/tests/unit/llms/azure/chat/test_azure_chat_o_series_transformation.py @@ -1,9 +1,9 @@ import json -import traceback -from typing import Callable, Optional +from typing import Final from unittest.mock import MagicMock, patch import pytest +from openai._legacy_response import LegacyAPIResponse import litellm from litellm import ModelResponse @@ -130,7 +130,7 @@ def test_azure_o3_streaming(): api_version="2024-02-15-preview", ) - with patch.object(client.chat.completions.with_raw_response, "create") as mock_create: + with patch.object(client, "post") as mock_create: try: completion( model="azure/o3-mini", @@ -141,7 +141,7 @@ def test_azure_o3_streaming(): except Exception as e: print(e) assert mock_create.call_count == 1 - assert "stream" in mock_create.call_args.kwargs + assert "stream" in mock_create.call_args.kwargs["body"] def test_azure_o_series_routing(): @@ -157,7 +157,7 @@ def test_azure_o_series_routing(): api_version="2024-02-15-preview", ) - with patch.object(client.chat.completions.with_raw_response, "create") as mock_create: + with patch.object(client, "post") as mock_create: try: completion( model="azure/o_series/my-random-deployment-name", @@ -168,17 +168,19 @@ def test_azure_o_series_routing(): except Exception as e: print(e) assert mock_create.call_count == 1 - assert "stream" not in mock_create.call_args.kwargs + assert "stream" not in mock_create.call_args.kwargs["body"] @patch("litellm.main.azure_o1_chat_completions._get_openai_client") def test_openai_o_series_max_retries_0(mock_get_openai_client): import litellm - mock_get_openai_client.return_value.chat.completions.with_raw_response.create.return_value.headers = {} - mock_get_openai_client.return_value.chat.completions.with_raw_response.create.return_value.parse.return_value = ( - ModelResponse(choices=[{"message": {"role": "assistant", "content": "Hello"}}]) - ) + raw_response: Final = MagicMock(spec=LegacyAPIResponse) + raw_response.headers = {} + raw_response.parse.return_value = ModelResponse(choices=[{"message": {"role": "assistant", "content": "Hello"}}]) + client_stub: Final = MagicMock() + client_stub.post.return_value = raw_response + mock_get_openai_client.return_value = client_stub litellm.set_verbose = True response = litellm.completion( model="azure/o1-preview", @@ -233,7 +235,7 @@ async def test_azure_o1_series_response_format_extra_params(): response_format = {"type": "json_object"} tool_choice = "auto" with patch.object( - client.chat.completions.with_raw_response, "create" + client, "post" ) as mock_client: try: await litellm.acompletion( @@ -251,7 +253,7 @@ async def test_azure_o1_series_response_format_extra_params(): print(f"Error: {e}") mock_client.assert_called_once() - request_body = mock_client.call_args.kwargs + request_body = mock_client.call_args.kwargs["body"] print("request_body: ", json.dumps(request_body, indent=4)) assert request_body["tools"] == tools diff --git a/tests/unit/llms/azure/test_azure.py b/tests/unit/llms/azure/test_azure.py index 7907344ae64..57d16bec4db 100644 --- a/tests/unit/llms/azure/test_azure.py +++ b/tests/unit/llms/azure/test_azure.py @@ -313,7 +313,7 @@ def test_azure_openai_gpt_4o_naming(monkeypatch): number: str = Field(description="total number of days in a week") days: list[str] = Field(description="name of days in a week") - with patch.object(client.chat.completions.with_raw_response, "create") as mock_post: + with patch.object(client, "post") as mock_post: try: completion( model="azure/gpt4o", @@ -326,9 +326,9 @@ def test_azure_openai_gpt_4o_naming(monkeypatch): mock_post.assert_called_once() - print(mock_post.call_args.kwargs) + print(mock_post.call_args.kwargs["body"]) - assert "tool_calls" not in mock_post.call_args.kwargs + assert "tool_calls" not in mock_post.call_args.kwargs["body"] @pytest.mark.parametrize( @@ -339,7 +339,6 @@ def test_azure_openai_gpt_4o_naming(monkeypatch): ], ) def test_azure_gpt_4o_with_tool_call_and_response_format(api_version): - from litellm import completion from typing import Optional from pydantic import BaseModel import litellm @@ -382,7 +381,7 @@ def test_azure_gpt_4o_with_tool_call_and_response_format(api_version): } ] - with patch.object(client.chat.completions.with_raw_response, "create") as mock_post: + with patch.object(client, "post") as mock_post: mock_post.return_value.headers = {} mock_post.return_value.parse.return_value = litellm.ModelResponse( choices=[{"message": {"role": "assistant", "content": InvestigationOutput().model_dump_json()}}] @@ -411,9 +410,9 @@ def test_azure_gpt_4o_with_tool_call_and_response_format(api_version): mock_post.assert_called_once() if api_version == "2024-10-21": - assert "response_format" in mock_post.call_args.kwargs + assert "response_format" in mock_post.call_args.kwargs["body"] else: - assert "response_format" not in mock_post.call_args.kwargs + assert "response_format" not in mock_post.call_args.kwargs["body"] assert response.choices[0].message.content == InvestigationOutput().model_dump_json() @@ -609,7 +608,6 @@ def test_azure_with_content_safety_error(): """ Verify user can access innererror from the Azure OpenAI exception """ - from litellm import completion from litellm.exceptions import ContentPolicyViolationError from litellm.litellm_core_utils.exception_mapping_utils import exception_type from unittest.mock import MagicMock diff --git a/tests/unit/llms/bedrock/chat/test_invoke_handler.py b/tests/unit/llms/bedrock/chat/test_invoke_handler.py index 069438e25c0..cf14190ff8a 100644 --- a/tests/unit/llms/bedrock/chat/test_invoke_handler.py +++ b/tests/unit/llms/bedrock/chat/test_invoke_handler.py @@ -19,6 +19,7 @@ from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging from litellm.litellm_core_utils.litellm_logging import _get_provider_request_id from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper from litellm.llms.bedrock.chat.invoke_handler import ( + AmazonAnthropicClaudeStreamDecoder, AmazonOpenAICompatibleStreamDecoder, AWSEventStreamDecoder, _stream_decoder, @@ -1120,6 +1121,7 @@ async def test_converse_streaming_completion_yields_text_and_finish_reason( ), ) ) + def handle_request(request: httpx.Request) -> httpx.Response: assert request.url.path.endswith("/converse-stream") return httpx.Response(200, content=event_stream_body, headers=_event_stream_headers()) @@ -1179,9 +1181,7 @@ def test_mistral_invoke_stream_yields_text_and_finish_reason( assert "".join(chunk.choices[0].delta.content or "" for chunk in chunks) == ( "A court case moves through the courts." ) - assert [chunk.choices[0].finish_reason for chunk in chunks if chunk.choices[0].finish_reason] == [ - "stop" - ] + assert [chunk.choices[0].finish_reason for chunk in chunks if chunk.choices[0].finish_reason] == ["stop"] def test_converse_streaming_tool_call_arguments_are_valid_json( @@ -1721,6 +1721,7 @@ def test_filter_headers_for_aws_signature(): @pytest.mark.parametrize("missing", ["botocore", "unrelated_dependency"]) def test_invoke_decoder_reports_only_missing_aws_dependency(missing): from unittest.mock import patch + from litellm.llms.bedrock.chat.invoke_handler import AWSEventStreamDecoder failure = ModuleNotFoundError("missing dependency", name=missing) @@ -1835,9 +1836,7 @@ async def test_converse_stream_thinking_merges_into_content_with_think_tags() -> "contentBlockDelta", {"contentBlockIndex": 0, "delta": {"reasoningContent": {"signature": "c2lnbmF0dXJl"}}} ), _converse_event_frame("contentBlockStop", {"contentBlockIndex": 0}), - _converse_event_frame( - "contentBlockDelta", {"contentBlockIndex": 1, "delta": {"text": "Hi I am Anthropic, "}} - ), + _converse_event_frame("contentBlockDelta", {"contentBlockIndex": 1, "delta": {"text": "Hi I am Anthropic, "}}), _converse_event_frame( "contentBlockDelta", {"contentBlockIndex": 1, "delta": {"text": "I am a helpful assistant"}} ), diff --git a/tests/unit/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py b/tests/unit/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py index dffe2c36e20..30bbc7e5179 100644 --- a/tests/unit/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py +++ b/tests/unit/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py @@ -3835,3 +3835,32 @@ def test_bedrock_removed_tool_change_does_not_add_beta(explicit_beta: bool) -> N assert result["messages"] == [{"role": "user", "content": "Reply with OK"}] assert result.get("anthropic_beta", []).count(beta) == int(explicit_beta) + + +@pytest.mark.parametrize("trace", ["enabled", "disabled"]) +def test_bedrock_messages_applies_the_deployment_guardrail_config_as_invoke_headers(trace: str) -> None: + from litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation import ( + AmazonAnthropicClaudeMessagesConfig, + ) + + headers, _ = AmazonAnthropicClaudeMessagesConfig().validate_anthropic_messages_environment( + headers={}, + model="anthropic.claude-sonnet-5-v1:0", + messages=[], + optional_params={}, + litellm_params={"guardrailConfig": {"guardrailIdentifier": "gid", "guardrailVersion": "2", "trace": trace}}, + ) + assert headers["X-Amzn-Bedrock-GuardrailIdentifier"] == "gid" + assert headers["X-Amzn-Bedrock-GuardrailVersion"] == "2" + assert headers["X-Amzn-Bedrock-Trace"] == trace.upper() + + +def test_bedrock_messages_without_guardrail_config_sends_no_guardrail_headers() -> None: + from litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation import ( + AmazonAnthropicClaudeMessagesConfig, + ) + + headers, _ = AmazonAnthropicClaudeMessagesConfig().validate_anthropic_messages_environment( + headers={}, model="anthropic.claude-sonnet-5-v1:0", messages=[], optional_params={}, litellm_params={} + ) + assert not any(name.startswith("X-Amzn-Bedrock-Guardrail") for name in headers) diff --git a/tests/unit/llms/bedrock/test_cost_calculation.py b/tests/unit/llms/bedrock/test_cost_calculation.py new file mode 100644 index 00000000000..80465cab5fe --- /dev/null +++ b/tests/unit/llms/bedrock/test_cost_calculation.py @@ -0,0 +1,41 @@ +import json +from importlib.resources import files +from typing import Final + +import pytest + +import litellm +from litellm.types.utils import Usage + +CLAUDE_PLATFORM_PREFIX: Final = "bedrock/claude_platform/" +CLAUDE_PLATFORM_KEYS: Final = tuple( + key + for key in json.loads(files("litellm").joinpath("model_prices_and_context_window_backup.json").read_text()) + if key.startswith(CLAUDE_PLATFORM_PREFIX) +) + + +def test_claude_platform_rows_exist() -> None: + assert CLAUDE_PLATFORM_KEYS + + +@pytest.mark.parametrize("inference_geo", ["global", "us"]) +@pytest.mark.parametrize("key", CLAUDE_PLATFORM_KEYS) +def test_claude_platform_cost_matches_claude_api_cost(local_model_cost_map: None, key: str, inference_geo: str) -> None: + usage: Final = Usage( + prompt_tokens=1_000, + completion_tokens=200, + cache_read_input_tokens=300, + cache_creation_input_tokens=250, + inference_geo=inference_geo, + ) + + claude_platform_cost: Final = litellm.cost_per_token( + model=key.removeprefix("bedrock/"), custom_llm_provider="bedrock", usage_object=usage + ) + claude_api_cost: Final = litellm.cost_per_token( + model=key.removeprefix(CLAUDE_PLATFORM_PREFIX), custom_llm_provider="anthropic", usage_object=usage + ) + + assert claude_platform_cost == pytest.approx(claude_api_cost) + assert sum(claude_platform_cost) > 0 diff --git a/tests/unit/llms/custom_httpx/test_aiohttp_transport.py b/tests/unit/llms/custom_httpx/test_aiohttp_transport.py index 0765371dc4a..aed808223bd 100644 --- a/tests/unit/llms/custom_httpx/test_aiohttp_transport.py +++ b/tests/unit/llms/custom_httpx/test_aiohttp_transport.py @@ -1,8 +1,11 @@ import aiohttp as aiohttp_aiohttp_handler, asyncio, importlib, litellm import concurrent.futures +import gc import socket import sys import struct +import threading +import weakref from typing import Final import aiohttp @@ -1310,3 +1313,128 @@ async def test_reset_after_request_is_not_replayed(monkeypatch: pytest.MonkeyPat assert len(received_requests) == 1 finally: await handler.close() + + +async def _new_session() -> aiohttp.ClientSession: + return aiohttp.ClientSession() + + +async def _until_closed(session: aiohttp.ClientSession) -> None: + while not session.closed: + await asyncio.sleep(0) + + +@pytest.mark.asyncio +async def test_a_collected_transport_closes_the_session_it_owns(): + """An SDK client evicted from the cache is collected long before the closer's grace window ends.""" + session = aiohttp.ClientSession() + transport = LiteLLMAiohttpTransport(client=session) + + del transport + gc.collect() + await asyncio.gather(*LiteLLMAiohttpTransport._background_close_tasks) + + assert session.closed, "the owned session outlived the transport that was collected before closing it" + + +@pytest.mark.asyncio +async def test_a_collected_transport_leaves_a_shared_session_open(): + session = aiohttp.ClientSession() + try: + transport = LiteLLMAiohttpTransport(client=session, owns_session=False) + + del transport + gc.collect() + await asyncio.gather(*LiteLLMAiohttpTransport._background_close_tasks) + + assert not session.closed, "a transport closed a session it does not own" + finally: + await session.close() + + +def test_a_transport_collected_after_its_loop_closed_still_closes_its_session(caplog: pytest.LogCaptureFixture): + """asyncio.run() closes the loop the session lives on before the evicted client is collected.""" + loop = asyncio.new_event_loop() + session = loop.run_until_complete(_new_session()) + loop.close() + transport = LiteLLMAiohttpTransport(client=session) + + del transport + gc.collect() + + assert session.closed, "the owned session stayed open after its transport was collected on a closed loop" + del session + gc.collect() + assert "Unclosed client session" not in caplog.text, caplog.text + + +def test_a_transport_collected_while_its_loop_runs_elsewhere_closes_its_session_on_that_loop(): + loop = asyncio.new_event_loop() + thread = threading.Thread(target=loop.run_forever, daemon=True) + thread.start() + session = asyncio.run_coroutine_threadsafe(_new_session(), loop).result(timeout=10) + transport = LiteLLMAiohttpTransport(client=session) + try: + del transport + gc.collect() + asyncio.run_coroutine_threadsafe(_until_closed(session), loop).result(timeout=10) + + assert session.closed, "the owned session stayed open after its transport was collected from another thread" + finally: + loop.call_soon_threadsafe(loop.stop) + thread.join(timeout=10) + loop.close() + + +@pytest.mark.asyncio +async def test_a_stream_keeps_reading_after_the_client_that_opened_it_is_collected(): + from aiohttp import web + + release: Final = asyncio.Event() + body: Final = tuple(f"chunk{i}\n".encode() for i in range(4)) + + async def handler(request: web.Request) -> web.StreamResponse: + response: Final = web.StreamResponse() + await response.prepare(request) + await response.write(body[0]) + await release.wait() + for chunk in body[1:]: + await response.write(chunk) + await response.write_eof() + return response + + app: Final = web.Application() + app.router.add_get("/stream", handler) + runner: Final = web.AppRunner(app) + await runner.setup() + sock: Final = socket.socket() + sock.bind(("127.0.0.1", 0)) + port: Final[int] = sock.getsockname()[1] + await web.SockSite(runner, sock).start() + + transport: Final = LiteLLMAiohttpTransport(client=aiohttp.ClientSession) + client: Final = httpx.AsyncClient(transport=transport) + transport_ref: Final = weakref.ref(transport) + try: + response: Final = await client.send(client.build_request("GET", f"http://127.0.0.1:{port}/stream"), stream=True) + session: Final = transport.client + assert isinstance(session, aiohttp.ClientSession) + chunks: Final = response.aiter_bytes() + first: Final = await chunks.__anext__() + + del client, transport + gc.collect() + await asyncio.gather(*LiteLLMAiohttpTransport._background_close_tasks) + release.set() + + rest: Final = [chunk async for chunk in chunks] + assert first + b"".join(rest) == b"".join(body), (first, rest) + assert not session.closed, "the session was closed under a stream that was still being read" + + del response, chunks + gc.collect() + await asyncio.gather(*LiteLLMAiohttpTransport._background_close_tasks) + assert transport_ref() is None, "the finished stream kept the transport alive" + assert session.closed, "the owned session stayed open after its last stream finished" + finally: + await runner.cleanup() diff --git a/tests/unit/llms/custom_httpx/test_http_handler.py b/tests/unit/llms/custom_httpx/test_http_handler.py index bd4ed4c99cc..4c8342bc32d 100644 --- a/tests/unit/llms/custom_httpx/test_http_handler.py +++ b/tests/unit/llms/custom_httpx/test_http_handler.py @@ -1,6 +1,5 @@ import asyncio, importlib, json, time import gc -import io import os import pathlib import ssl @@ -1388,6 +1387,55 @@ async def test_finalizer_with_running_loop_schedules_close_and_holds_task_ref(): assert not (AsyncHTTPHandler._finalizer_close_tasks & scheduled) +async def _drop_unused_handler() -> None: + AsyncHTTPHandler(timeout=61.0) + gc.collect() + + +def test_finalizer_close_does_not_keep_a_loop_closed_right_after_it_alive(): + loop: Final = asyncio.new_event_loop() + loop.run_until_complete(_drop_unused_handler()) + loop.close() + closed_loop: Final = weakref.ref(loop) + del loop + gc.collect() + + assert closed_loop() is None + + +class _TransportFailingToClose(httpx.AsyncBaseTransport): + async def aclose(self) -> None: + raise RuntimeError("close failed") + + +async def _drop_handler_whose_close_fails() -> None: + AsyncHTTPHandler(timeout=61.0, transport=_TransportFailingToClose()) + gc.collect() + + +def test_failed_finalizer_close_in_a_loops_last_step_reports_no_unretrieved_exception(): + loop: Final = asyncio.new_event_loop() + exception_handler: Final = MagicMock() + loop.set_exception_handler(exception_handler) + loop.run_until_complete(_drop_handler_whose_close_fails()) + loop.close() + del loop + gc.collect() + + exception_handler.assert_not_called() + + +async def _sync_client_with_timeout_62() -> HTTPHandler: + return get_httpx_client(params={"timeout": 62.0}) + + +def test_sync_client_is_shared_by_short_lived_loops_and_callers_outside_a_loop(): + built_in_first_loop: Final = asyncio.run(_sync_client_with_timeout_62()) + + assert asyncio.run(_sync_client_with_timeout_62()) is built_in_first_loop + assert get_httpx_client(params={"timeout": 62.0}) is built_in_first_loop + + @pytest.mark.asyncio async def test_sync_close_helper_respects_session_ownership(): """The loop-less fallback closes only sessions the transport owns; a diff --git a/tests/unit/llms/litellm_proxy/chat/test_litellm_proxy_chat_transformation.py b/tests/unit/llms/litellm_proxy/chat/test_litellm_proxy_chat_transformation.py index 86fb6fa15d8..9b8788d3769 100644 --- a/tests/unit/llms/litellm_proxy/chat/test_litellm_proxy_chat_transformation.py +++ b/tests/unit/llms/litellm_proxy/chat/test_litellm_proxy_chat_transformation.py @@ -1,15 +1,16 @@ import json +import re from io import BytesIO from pathlib import Path from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import httpx +import pytest from openai import AsyncOpenAI, OpenAI +from openai._legacy_response import LegacyAPIResponse from openai.types import CreateEmbeddingResponse, Embedding from openai.types.create_embedding_response import Usage -import pytest -import re import litellm from litellm import completion @@ -18,7 +19,6 @@ from litellm.llms.litellm_proxy.chat.transformation import LiteLLMProxyChatConfi from tests._vcr_conftest_common import rewound_new_episodes_cassette from tests.capturing_transport import CapturingTransport - _GATEWAY_EMBEDDING_RESPONSE: Final = CreateEmbeddingResponse( object="list", data=(Embedding(object="embedding", index=0, embedding=(0.1, 0.2, 0.3)),), @@ -98,7 +98,7 @@ async def test_litellm_gateway_from_sdk(): openai_client = OpenAI(api_key="fake-key") with patch.object( - openai_client.chat.completions.with_raw_response, "create", new=MagicMock() + openai_client, "post", new=MagicMock() ) as mock_call: try: completion( @@ -114,9 +114,9 @@ async def test_litellm_gateway_from_sdk(): mock_call.assert_called_once() - print("Call KWARGS - {}".format(mock_call.call_args.kwargs)) + print("Call KWARGS - {}".format(mock_call.call_args.kwargs["body"])) - assert "hello" in mock_call.call_args.kwargs["extra_body"] + assert "hello" in mock_call.call_args.kwargs["options"]["extra_json"] @pytest.mark.parametrize("is_async", (False, True)) @@ -264,7 +264,7 @@ async def test_litellm_gateway_from_sdk_rerank(is_async): patch_target = client.post with patch.object(client, "post", new=mock_method): - mock_response = MagicMock() + mock_response = MagicMock(spec=LegacyAPIResponse) # Create a mock response similar to OpenAI's rerank response mock_response.text = json.dumps( @@ -408,7 +408,7 @@ async def test_litellm_gateway_from_sdk_structured_output(): openai_client = OpenAI(api_key="fake-key") with patch.object( - openai_client.chat.completions, "create", new=MagicMock() + openai_client, "post", new=MagicMock() ) as mock_call: try: litellm.completion( @@ -427,8 +427,8 @@ async def test_litellm_gateway_from_sdk_structured_output(): mock_call.assert_called_once() - print("Call KWARGS - {}".format(mock_call.call_args.kwargs)) - json_schema = mock_call.call_args.kwargs["response_format"] + print("Call KWARGS - {}".format(mock_call.call_args.kwargs["body"])) + json_schema = mock_call.call_args.kwargs["body"]["response_format"] assert "json_schema" in json_schema @@ -486,7 +486,7 @@ def test_litellm_gateway_from_sdk_with_response_cost_in_additional_headers(): openai_client = OpenAI(api_key="fake-key") # Create mock response object - mock_response = MagicMock() + mock_response = MagicMock(spec=LegacyAPIResponse) mock_response.headers = {"x-litellm-response-cost": "120"} mock_response.parse.return_value = litellm.ModelResponse( **{ @@ -528,8 +528,7 @@ def test_litellm_gateway_from_sdk_with_response_cost_in_additional_headers(): ) with patch.object( - openai_client.chat.completions.with_raw_response, - "create", + openai_client, "post", return_value=mock_response, ) as mock_call: response = litellm.completion( @@ -564,7 +563,7 @@ async def test_litellm_gateway_image_generation_direct(is_async): "created": 1, "data": [{"url": "https://example.com/image.png"}], } - mock_raw_response = MagicMock() + mock_raw_response = MagicMock(spec=LegacyAPIResponse) mock_raw_response.parse.return_value = mock_openai_response mock_raw_response.headers = {} diff --git a/tests/unit/llms/nvidia_nim/test_nvidia_nim.py b/tests/unit/llms/nvidia_nim/test_nvidia_nim.py index e5298c7ba39..4c62dc81bb3 100644 --- a/tests/unit/llms/nvidia_nim/test_nvidia_nim.py +++ b/tests/unit/llms/nvidia_nim/test_nvidia_nim.py @@ -1,6 +1,6 @@ import json from typing import Final -from unittest.mock import AsyncMock, MagicMock, patch +from unittest.mock import AsyncMock, patch import httpx import pytest @@ -8,7 +8,7 @@ from openai.types import CreateEmbeddingResponse, Embedding from openai.types.create_embedding_response import Usage as EmbeddingUsage import litellm -from litellm import ModelResponse, completion +from litellm import completion from tests.capturing_transport import CapturingTransport @@ -94,7 +94,7 @@ def test_chat_completion_nvidia_nim_with_tools(): ] with patch.object( - client.chat.completions.with_raw_response, "create" + client, "post" ) as mock_client: try: completion( @@ -116,7 +116,7 @@ def test_chat_completion_nvidia_nim_with_tools(): # Add assertions to check the request mock_client.assert_called_once() - request_body = mock_client.call_args.kwargs + request_body = mock_client.call_args.kwargs["body"] print("request_body: ", request_body) @@ -208,7 +208,7 @@ def test_completion_nvidia_nim(): api_key="fake-api-key", ) - with patch.object(client.chat.completions.with_raw_response, "create") as mock_client: + with patch.object(client, "post") as mock_client: try: completion( model=model_name, @@ -226,7 +226,7 @@ def test_completion_nvidia_nim(): print(e) mock_client.assert_called_once() - request_body = mock_client.call_args.kwargs + request_body = mock_client.call_args.kwargs["body"] assert request_body["messages"] == [ { diff --git a/tests/unit/llms/openai/test_o_series_transformation.py b/tests/unit/llms/openai/test_o_series_transformation.py index 6ee882b7d32..c62523cf756 100644 --- a/tests/unit/llms/openai/test_o_series_transformation.py +++ b/tests/unit/llms/openai/test_o_series_transformation.py @@ -1,5 +1,4 @@ -from typing import Final -from unittest.mock import AsyncMock, MagicMock, patch +from unittest.mock import patch import pytest @@ -70,7 +69,7 @@ async def test_o1_handle_system_role(model): client = AsyncOpenAI(api_key="fake-api-key") with patch.object( - client.chat.completions.with_raw_response, "create" + client, "post" ) as mock_client: try: await litellm.acompletion( @@ -83,7 +82,7 @@ async def test_o1_handle_system_role(model): print(f"Error: {e}") mock_client.assert_called_once() - request_body = mock_client.call_args.kwargs + request_body = mock_client.call_args.kwargs["body"] print("request_body: ", request_body) @@ -141,7 +140,7 @@ async def test_o1_max_completion_tokens(model: str): client = AsyncOpenAI(api_key="fake-api-key") with patch.object( - client.chat.completions.with_raw_response, "create" + client, "post" ) as mock_client: try: await litellm.acompletion( @@ -154,7 +153,7 @@ async def test_o1_max_completion_tokens(model: str): print(f"Error: {e}") mock_client.assert_called_once() - request_body = mock_client.call_args.kwargs + request_body = mock_client.call_args.kwargs["body"] print("request_body: ", request_body) diff --git a/tests/unit/llms/openai/test_openai.py b/tests/unit/llms/openai/test_openai.py index e37fd4b6ae5..749dc823f20 100644 --- a/tests/unit/llms/openai/test_openai.py +++ b/tests/unit/llms/openai/test_openai.py @@ -5,12 +5,13 @@ import os from collections.abc import Mapping from types import ModuleType from typing import Final, Literal, Protocol -from unittest.mock import Mock, patch +from unittest.mock import MagicMock, Mock, patch import httpx import pytest import respx from openai import AsyncOpenAI, OpenAI +from openai._legacy_response import LegacyAPIResponse from openai._types import Response as SDKResponse from pydantic import BaseModel, ConfigDict, JsonValue @@ -1596,7 +1597,7 @@ async def test_openai_prediction_param_mock(): client = AsyncOpenAI(api_key="fake-api-key") with patch.object( - client.chat.completions.with_raw_response, "create" + client, "post" ) as mock_client: try: await litellm.acompletion( @@ -1615,7 +1616,7 @@ async def test_openai_prediction_param_mock(): print(f"Error: {e}") mock_client.assert_called_once() - request_body = mock_client.call_args.kwargs + request_body = mock_client.call_args.kwargs["body"] # Verify the request contains the prediction parameter assert "prediction" in request_body @@ -1627,10 +1628,12 @@ async def test_openai_prediction_param_mock(): def test_openai_max_retries_0(mock_get_openai_client): import litellm - mock_get_openai_client.return_value.chat.completions.with_raw_response.create.return_value.headers = {} - mock_get_openai_client.return_value.chat.completions.with_raw_response.create.return_value.parse.return_value = ( - ModelResponse(choices=[{"message": {"role": "assistant", "content": "Hello"}}]) - ) + raw_response: Final = MagicMock(spec=LegacyAPIResponse) + raw_response.headers = {} + raw_response.parse.return_value = ModelResponse(choices=[{"message": {"role": "assistant", "content": "Hello"}}]) + client_stub: Final = MagicMock() + client_stub.post.return_value = raw_response + mock_get_openai_client.return_value = client_stub litellm.set_verbose = True response = litellm.completion( model="gpt-4o-mini", @@ -1739,7 +1742,7 @@ async def test_openai_safety_identifier_parameter(): client = AsyncOpenAI(api_key="fake-api-key") with patch.object( - client.chat.completions.with_raw_response, "create" + client, "post" ) as mock_client: try: await litellm.acompletion( @@ -1752,7 +1755,7 @@ async def test_openai_safety_identifier_parameter(): print(f"Error: {e}") mock_client.assert_called_once() - request_body = mock_client.call_args.kwargs + request_body = mock_client.call_args.kwargs["body"] # Verify the request contains the safety_identifier parameter assert "safety_identifier" in request_body @@ -1768,7 +1771,7 @@ def test_openai_safety_identifier_parameter_sync(): client = OpenAI(api_key="fake-api-key") with patch.object( - client.chat.completions.with_raw_response, "create" + client, "post" ) as mock_client: try: litellm.completion( @@ -1781,7 +1784,7 @@ def test_openai_safety_identifier_parameter_sync(): print(f"Error: {e}") mock_client.assert_called_once() - request_body = mock_client.call_args.kwargs + request_body = mock_client.call_args.kwargs["body"] # Verify the request contains the safety_identifier parameter assert "safety_identifier" in request_body @@ -1798,7 +1801,7 @@ async def test_openai_service_tier_parameter(): client = AsyncOpenAI(api_key="fake-api-key") with patch.object( - client.chat.completions.with_raw_response, "create" + client, "post" ) as mock_client: try: await litellm.acompletion( @@ -1811,7 +1814,7 @@ async def test_openai_service_tier_parameter(): print(f"Error: {e}") mock_client.assert_called_once() - request_body = mock_client.call_args.kwargs + request_body = mock_client.call_args.kwargs["body"] # Verify the request contains the service_tier parameter assert "service_tier" in request_body, "service_tier should be in request body" @@ -1829,7 +1832,7 @@ def test_openai_service_tier_parameter_sync(): client = OpenAI(api_key="fake-api-key") with patch.object( - client.chat.completions.with_raw_response, "create" + client, "post" ) as mock_client: try: litellm.completion( @@ -1842,7 +1845,7 @@ def test_openai_service_tier_parameter_sync(): print(f"Error: {e}") mock_client.assert_called_once() - request_body = mock_client.call_args.kwargs + request_body = mock_client.call_args.kwargs["body"] # Verify the request contains the service_tier parameter assert "service_tier" in request_body, "service_tier should be in request body" @@ -1874,3 +1877,63 @@ def test_responses_gpt54_with_xhigh_reasoning(): assert request_body["model"] == "openai/gpt-5.4" assert request_body["reasoning"] == {"effort": "xhigh"} + + +_CHAT_RESPONSE_JSON: Final = { + "id": "chatcmpl-wire", + "object": "chat.completion", + "created": 1, + "model": "gpt-5.4", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, +} +_WIRE_REQUEST: Final = { + "model": "gpt-5.4", + "messages": [ + {"role": "system", "content": "sys"}, + { + "role": "user", + "content": [ + {"type": "text", "text": "what is in the image"}, + {"type": "image_url", "image_url": {"url": "https://example.com/a.png", "detail": "low"}}, + ], + }, + ], + "tools": [{"type": "function", "function": {"name": "f", "parameters": {"type": "object", "properties": {}}}}], + "tool_choice": "auto", + "response_format": {"type": "json_schema", "json_schema": {"name": "r", "schema": {"type": "object"}}}, + "max_completion_tokens": 7, + "reasoning_effort": "none", + "n": 2, + "seed": 3, +} + + +@pytest.mark.parametrize("is_async", (False, True)) +@respx.mock +async def test_chat_completion_sends_the_caller_params_extra_body_and_extra_headers_unchanged( + is_async: bool, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + route: Final = respx.post("https://api.openai.com/v1/chat/completions").mock( + return_value=httpx.Response(200, json=_CHAT_RESPONSE_JSON) + ) + kwargs: Final = { + **_WIRE_REQUEST, + "model": "openai/gpt-5.4", + "api_base": "https://api.openai.com/v1", + "api_key": "sk-test", + "num_retries": 0, + "extra_body": {"vendor_flag": True}, + "extra_headers": {"x-caller": "unit"}, + } + + if is_async: + await litellm.acompletion(**kwargs) + else: + litellm.completion(**kwargs) + + sent: Final = route.calls.last.request + assert json.loads(sent.content) == {**_WIRE_REQUEST, "vendor_flag": True} + assert sent.headers["x-caller"] == "unit" + assert sent.headers["authorization"] == "Bearer sk-test" diff --git a/tests/unit/llms/openai/test_openai_common_utils.py b/tests/unit/llms/openai/test_openai_common_utils.py index 3c7dff34296..5004f5236bd 100644 --- a/tests/unit/llms/openai/test_openai_common_utils.py +++ b/tests/unit/llms/openai/test_openai_common_utils.py @@ -1,3 +1,4 @@ +import asyncio from typing import Final from unittest.mock import MagicMock, patch @@ -164,6 +165,28 @@ def test_get_openai_client_initialization_param_fields(client_type): assert "self" not in result +def _loop_test_client_params(is_async: bool) -> dict[str, object]: + return {"api_key": "sk-loop-scope-test", "is_async": is_async} + + +async def _cache_openai_client(client: openai.OpenAI | openai.AsyncOpenAI, is_async: bool) -> None: + BaseOpenAILLM.set_cached_openai_client(client, "openai", _loop_test_client_params(is_async)) + + +async def _cached_openai_client(is_async: bool) -> object: + return BaseOpenAILLM.get_cached_openai_client(_loop_test_client_params(is_async), "openai") + + +def test_sync_openai_client_cached_in_one_loop_is_reused_by_the_next_while_async_is_not(): + sync_client: Final = openai.OpenAI(api_key="sk-loop-scope-test") + async_client: Final = openai.AsyncOpenAI(api_key="sk-loop-scope-test") + asyncio.run(_cache_openai_client(sync_client, is_async=False)) + asyncio.run(_cache_openai_client(async_client, is_async=True)) + + assert asyncio.run(_cached_openai_client(is_async=False)) is sync_client + assert asyncio.run(_cached_openai_client(is_async=True)) is None + + @pytest.mark.parametrize("client_type", ["openai", "azure"]) def test_get_openai_client_cache_key(client_type): """Verify get_openai_client_cache_key doesn't raise on tuple + tuple concatenation.""" diff --git a/tests/unit/llms/openai/test_openai_empty_response.py b/tests/unit/llms/openai/test_openai_empty_response.py index 26b28f967db..1505d9ef049 100644 --- a/tests/unit/llms/openai/test_openai_empty_response.py +++ b/tests/unit/llms/openai/test_openai_empty_response.py @@ -2,13 +2,13 @@ Test for issue #17209: Clearer error when LLM endpoint returns empty response """ -from unittest.mock import MagicMock, patch +from unittest.mock import MagicMock import pytest +from openai._legacy_response import LegacyAPIResponse - -from litellm.llms.openai.openai import OpenAIChatCompletion from litellm.llms.openai.common_utils import OpenAIError +from litellm.llms.openai.openai import OpenAIChatCompletion class TestEmptyResponseHandling: @@ -22,12 +22,12 @@ class TestEmptyResponseHandling: openai_chat = OpenAIChatCompletion() # Mock the raw response to return an empty string from parse() - mock_raw_response = MagicMock() + mock_raw_response = MagicMock(spec=LegacyAPIResponse) mock_raw_response.headers = {} mock_raw_response.parse.return_value = "" # Empty string response mock_client = MagicMock() - mock_client.chat.completions.with_raw_response.create.return_value = ( + mock_client.post.return_value = ( mock_raw_response ) @@ -48,12 +48,12 @@ class TestEmptyResponseHandling: """Test that None response also produces a clear error""" openai_chat = OpenAIChatCompletion() - mock_raw_response = MagicMock() + mock_raw_response = MagicMock(spec=LegacyAPIResponse) mock_raw_response.headers = {} mock_raw_response.parse.return_value = None mock_client = MagicMock() - mock_client.chat.completions.with_raw_response.create.return_value = ( + mock_client.post.return_value = ( mock_raw_response ) @@ -75,12 +75,12 @@ class TestEmptyResponseHandling: mock_response = MagicMock() mock_response.model_dump.return_value = {"choices": []} - mock_raw_response = MagicMock() + mock_raw_response = MagicMock(spec=LegacyAPIResponse) mock_raw_response.headers = {"x-request-id": "123"} mock_raw_response.parse.return_value = mock_response mock_client = MagicMock() - mock_client.chat.completions.with_raw_response.create.return_value = ( + mock_client.post.return_value = ( mock_raw_response ) @@ -104,12 +104,12 @@ class TestEmptyResponseHandling: # Create a mock response WITHOUT model_dump (like an AsyncStream/Iterator) mock_stream = MagicMock(spec=[]) # spec=[] means no attributes - mock_raw_response = MagicMock() + mock_raw_response = MagicMock(spec=LegacyAPIResponse) mock_raw_response.headers = {"x-request-id": "123"} mock_raw_response.parse.return_value = mock_stream mock_client = MagicMock() - mock_client.chat.completions.with_raw_response.create.return_value = ( + mock_client.post.return_value = ( mock_raw_response ) diff --git a/tests/unit/llms/openai/test_use_chat_completions_api_no_leak.py b/tests/unit/llms/openai/test_use_chat_completions_api_no_leak.py index 7013afc7a5f..58f1ff60ba5 100644 --- a/tests/unit/llms/openai/test_use_chat_completions_api_no_leak.py +++ b/tests/unit/llms/openai/test_use_chat_completions_api_no_leak.py @@ -9,6 +9,7 @@ body. OpenAI/Anthropic reject unknown body params with HTTP 400. from unittest.mock import MagicMock +from openai._legacy_response import LegacyAPIResponse import litellm from litellm.types.utils import all_litellm_params @@ -47,12 +48,12 @@ def test_completion_does_not_leak_flag_into_provider_request_body(): }, } - mock_raw_response = MagicMock() + mock_raw_response = MagicMock(spec=LegacyAPIResponse) mock_raw_response.headers = {} mock_raw_response.parse.return_value = mock_response mock_client = MagicMock() - mock_client.chat.completions.with_raw_response.create.return_value = ( + mock_client.post.return_value = ( mock_raw_response ) @@ -65,7 +66,7 @@ def test_completion_does_not_leak_flag_into_provider_request_body(): ) create_kwargs = ( - mock_client.chat.completions.with_raw_response.create.call_args.kwargs + mock_client.post.call_args.kwargs ) assert "use_chat_completions_api" not in create_kwargs assert "use_chat_completions_api" not in (create_kwargs.get("extra_body") or {}) diff --git a/tests/unit/llms/perplexity/chat/test_transformation.py b/tests/unit/llms/perplexity/chat/test_transformation.py index 6a70fd1232d..0cb6335f5a8 100644 --- a/tests/unit/llms/perplexity/chat/test_transformation.py +++ b/tests/unit/llms/perplexity/chat/test_transformation.py @@ -4,6 +4,7 @@ import os from unittest.mock import MagicMock, patch import pytest +from openai._legacy_response import LegacyAPIResponse import litellm from litellm import completion @@ -92,7 +93,7 @@ class TestPerplexityReasoning: pydantic_obj = ChatCompletion(**response_object) def _return_pydantic_obj(*args, **kwargs): - new_response = MagicMock() + new_response = MagicMock(spec=LegacyAPIResponse) new_response.headers = {"content-type": "application/json"} new_response.parse.return_value = pydantic_obj return new_response @@ -100,8 +101,7 @@ class TestPerplexityReasoning: openai_client = OpenAI(api_key="fake-api-key") with patch.object( - openai_client.chat.completions.with_raw_response, - "create", + openai_client, "post", side_effect=_return_pydantic_obj, ) as mock_client: response = completion( @@ -121,7 +121,7 @@ class TestPerplexityReasoning: # Get the request data from the mock call call_args = mock_client.call_args - request_data = call_args.kwargs + request_data = call_args.kwargs["body"] # Verify reasoning_effort was included in the request assert "reasoning_effort" in request_data diff --git a/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py b/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py index 5c0e9f1a305..3d7296ffa7b 100644 --- a/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py +++ b/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py @@ -13,7 +13,9 @@ from pydantic import BaseModel import litellm from litellm import ModelResponse, completion, embedding, image_generation +from litellm.cost_calculator import batch_cost_calculator from litellm.litellm_core_utils.exception_mapping_utils import exception_type +from litellm.litellm_core_utils.llm_cost_calc.utils import generic_cost_per_token, parse_prompt_tokens_details from litellm.llms.anthropic.pass_through.messages import handler as anthropic_messages_handler from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.llms.gemini.chat.transformation import GoogleAIStudioGeminiConfig @@ -22,8 +24,9 @@ from litellm.llms.vertex_ai.vertex_llm_base import VertexBase from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( VertexGeminiConfig, ) +from litellm.types.llms.base import CachedTokensDetails from litellm.types.llms.vertex_ai import GeminiFinishReason, UsageMetadata -from litellm.types.utils import ChoiceLogprobs, Usage +from litellm.types.utils import ChoiceLogprobs, ModelInfo, Usage from litellm.utils import _invalidate_model_cost_lowercase_map, CustomStreamWrapper from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from litellm.integrations.custom_logger import CustomLogger @@ -1271,6 +1274,7 @@ def test_vertex_ai_usage_metadata_accumulates_duplicate_modalities(): v = VertexGeminiConfig() usage_metadata = { "promptTokenCount": 210, + "cachedContentTokenCount": 54, "candidatesTokenCount": 50, "totalTokenCount": 260, "promptTokensDetails": [ @@ -1292,9 +1296,10 @@ def test_vertex_ai_usage_metadata_accumulates_duplicate_modalities(): usage_metadata = UsageMetadata(**usage_metadata) result = v.calculate_usage(completion_response={"usageMetadata": usage_metadata}) - # prompt details are total - cached per modality - assert result.prompt_tokens_details.text_tokens == 16 # 20 - 4 - assert result.prompt_tokens_details.image_tokens == 140 # (90 + 100) - (40 + 10) + assert result.prompt_tokens_details.text_tokens == 20 + assert result.prompt_tokens_details.image_tokens == 190 + assert result.prompt_tokens_details.cached_tokens == 54 + assert result.prompt_tokens_details.cached_tokens_details == CachedTokensDetails(text_tokens=4, image_tokens=50) # candidates details accumulate duplicate modalities assert result.completion_tokens_details.text_tokens == 20 # 15 + 5 @@ -1302,6 +1307,118 @@ def test_vertex_ai_usage_metadata_accumulates_duplicate_modalities(): assert result.completion_tokens == 50 +def _gemini_usage_with_cached_audio(cached_audio_tokens: int) -> Usage: + cache_tokens_details: Final = ( + ( + {"modality": "TEXT", "tokenCount": 5}, + {"modality": "AUDIO", "tokenCount": cached_audio_tokens}, + ) + if cached_audio_tokens > 0 + else ({"modality": "TEXT", "tokenCount": 5},) + ) + completion_response: Final = { + "usageMetadata": { + "promptTokenCount": 40, + "cachedContentTokenCount": 5 + cached_audio_tokens, + "promptTokensDetails": ( + {"modality": "TEXT", "tokenCount": 20}, + {"modality": "AUDIO", "tokenCount": 20}, + ), + "cacheTokensDetails": cache_tokens_details, + "candidatesTokenCount": 7, + "totalTokenCount": 47, + } + } + return VertexGeminiConfig.calculate_usage(completion_response=completion_response) + + +def test_gemini_raw_usage_metadata_prices_cached_audio_at_its_batch_rate() -> None: + model_info: Final[ModelInfo] = { + "input_cost_per_token": 2e-6, + "input_cost_per_token_batches": 1e-6, + "input_cost_per_audio_token": 6e-6, + "input_cost_per_audio_token_batches": 3e-6, + "cache_read_input_token_cost": 8e-7, + "cache_read_input_token_cost_batches": 4e-7, + "cache_read_input_audio_token_cost": 5e-7, + "cache_read_input_audio_token_cost_batches": 2e-7, + } + usage: Final = _gemini_usage_with_cached_audio(5) + expected_prompt_cost: Final = ( + 15 * model_info["input_cost_per_token_batches"] + + 15 * model_info["input_cost_per_audio_token_batches"] + + 5 * model_info["cache_read_input_token_cost_batches"] + + 5 * model_info["cache_read_input_audio_token_cost_batches"] + ) + + prompt_cost, _ = batch_cost_calculator( + usage=usage, + model="gemini-audio-batch-test-model", + custom_llm_provider="gemini", + model_info=model_info, + ) + + assert prompt_cost == pytest.approx(expected_prompt_cost) + + +def test_gemini_raw_usage_metadata_prices_cached_audio_at_its_standard_rate() -> None: + model_info: Final[ModelInfo] = { + "input_cost_per_token": 2e-6, + "input_cost_per_audio_token": 6e-6, + "cache_read_input_token_cost": 8e-7, + "cache_read_input_audio_token_cost": 5e-7, + } + usage: Final = _gemini_usage_with_cached_audio(5) + expected_prompt_cost: Final = ( + 15 * model_info["input_cost_per_token"] + + 15 * model_info["input_cost_per_audio_token"] + + 5 * model_info["cache_read_input_token_cost"] + + 5 * model_info["cache_read_input_audio_token_cost"] + ) + + prompt_cost, _ = generic_cost_per_token( + model="gemini-audio-standard-test-model", + usage=usage, + custom_llm_provider="gemini", + model_info=model_info, + ) + + assert prompt_cost == pytest.approx(expected_prompt_cost) + + +def test_gemini_usage_totals_are_subtracted_once_and_audio_cache_rate_matches_audio_rate() -> None: + model_info: Final[ModelInfo] = { + "input_cost_per_token": 2e-6, + "input_cost_per_audio_token": 6e-6, + "cache_read_input_token_cost": 8e-7, + "cache_read_input_audio_token_cost": 6e-6, + } + cached_usage: Final = _gemini_usage_with_cached_audio(5) + uncached_usage: Final = _gemini_usage_with_cached_audio(0) + cached_details: Final = parse_prompt_tokens_details(cached_usage) + cached_prompt_cost, _ = generic_cost_per_token( + model="gemini-audio-standard-test-model", + usage=cached_usage, + custom_llm_provider="gemini", + model_info=model_info, + ) + uncached_prompt_cost, _ = generic_cost_per_token( + model="gemini-audio-standard-test-model", + usage=uncached_usage, + custom_llm_provider="gemini", + model_info=model_info, + ) + + assert cached_usage.prompt_tokens_details.text_tokens == 20 + assert cached_usage.prompt_tokens_details.audio_tokens == 20 + assert cached_usage.prompt_tokens_details.cached_tokens_details == CachedTokensDetails( + text_tokens=5, audio_tokens=5 + ) + assert cached_details["text_tokens"] == 15 + assert cached_details["audio_tokens"] == 15 + assert cached_prompt_cost == pytest.approx(uncached_prompt_cost) + + def test_vertex_ai_map_thinking_param_with_budget_tokens_0(): """ If budget_tokens is 0, do not set includeThoughts to True @@ -5532,7 +5649,7 @@ def test_vertex_ai_usage_metadata_with_document_tokens_in_prompt(): def test_vertex_ai_usage_metadata_with_document_tokens_cached(): - """Test that cached DOCUMENT tokens are correctly subtracted from prompt text tokens.""" + """Test cached DOCUMENT token details and unchanged prompt cost.""" v = VertexGeminiConfig() usage_metadata_dict = { @@ -5556,10 +5673,24 @@ def test_vertex_ai_usage_metadata_with_document_tokens_cached(): completion_response = {"usageMetadata": usage_metadata_dict} result = v.calculate_usage(completion_response=completion_response) - # DOCUMENT cached tokens map to cached_text_tokens, so: - # text_tokens = (8 TEXT + 774 DOCUMENT) - 400 cached = 382 - assert result.prompt_tokens_details.text_tokens == 382, "text_tokens should be (8 + 774) - 400 cached = 382" + assert result.prompt_tokens_details.text_tokens == 782 assert result.prompt_tokens_details.cached_tokens == 400 + assert result.prompt_tokens_details.cached_tokens_details.text_tokens == 400 + + model_info: Final[ModelInfo] = { + "input_cost_per_token": 2e-6, + "cache_read_input_token_cost": 5e-7, + } + prompt_cost, _ = generic_cost_per_token( + model="gemini-document-cache-test-model", + usage=result, + custom_llm_provider="gemini", + model_info=model_info, + ) + + assert prompt_cost == pytest.approx( + 382 * model_info["input_cost_per_token"] + 400 * model_info["cache_read_input_token_cost"] + ) def test_async_streaming_uses_custom_client(): diff --git a/tests/unit/llms/volcengine/test_volcengine.py b/tests/unit/llms/volcengine/test_volcengine.py index d472d52d2db..aecd3879afa 100644 --- a/tests/unit/llms/volcengine/test_volcengine.py +++ b/tests/unit/llms/volcengine/test_volcengine.py @@ -1,8 +1,6 @@ -import os -import sys from unittest.mock import MagicMock, patch -from pydantic import BaseModel +from openai._legacy_response import LegacyAPIResponse from litellm.llms.volcengine.chat.transformation import ( VolcEngineChatConfig as VolcEngineConfig, @@ -117,7 +115,7 @@ class TestVolcEngineConfig: client = OpenAI(api_key="test_api_key") - mock_raw_response = MagicMock() + mock_raw_response = MagicMock(spec=LegacyAPIResponse) mock_raw_response.headers = { "x-request-id": "123", "openai-organization": "org-123", @@ -126,9 +124,7 @@ class TestVolcEngineConfig: } mock_raw_response.parse.return_value = ModelResponse() - with patch.object( - client.chat.completions.with_raw_response, "create", mock_raw_response - ) as mock_create: + with patch.object(client, "post", MagicMock(return_value=mock_raw_response)) as mock_create: completion( model="volcengine/doubao-seed-1.6", messages=[ @@ -144,11 +140,6 @@ class TestVolcEngineConfig: ) mock_create.assert_called_once() - print(mock_create.call_args.kwargs) + print(mock_create.call_args.kwargs["body"]) # Fixed: thinking disabled should appear in extra_body with original structure - assert ( - "extra_body" in mock_create.call_args.kwargs - and "thinking" in mock_create.call_args.kwargs.get("extra_body", {}) - and mock_create.call_args.kwargs.get("extra_body", {})["thinking"] - == {"type": "disabled"} - ) + assert mock_create.call_args.kwargs["options"]["extra_json"]["thinking"] == {"type": "disabled"} diff --git a/tests/unit/proxy/analytics_endpoints/test_analytics_endpoints.py b/tests/unit/proxy/analytics_endpoints/test_analytics_endpoints.py index 1c706978fed..2d28292a487 100644 --- a/tests/unit/proxy/analytics_endpoints/test_analytics_endpoints.py +++ b/tests/unit/proxy/analytics_endpoints/test_analytics_endpoints.py @@ -26,6 +26,7 @@ from litellm.proxy.analytics_endpoints.cache_activity import ( ADMIN: Final = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin-user") INTERNAL_USER: Final = UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, user_id="user-1") +from tests.unit.proxy.db.fake_prisma_engine import statement_timeout_error, statement_timeout_reads async def _no_teams(auth: UserAPIKeyAuth) -> tuple[str, ...]: @@ -70,10 +71,15 @@ MODEL_ROWS = [{"model": "gpt-5.1"}] def build_prisma(query_raw: AsyncMock) -> MagicMock: prisma = MagicMock() - prisma.db.query_raw = query_raw + statement_timeout_reads(prisma, query_raw) return prisma +def _spend_log_reads(prisma: MagicMock) -> AsyncMock: + """The raw reads the cache dashboard ran under the statement timeout.""" + return prisma.read_db.tx.return_value.query_raw + + def dispatching_query_raw() -> AsyncMock: async def dispatch(sql: str, *params: object) -> list[dict[str, object]]: if "error_code" in sql: @@ -120,7 +126,9 @@ async def test_filters_are_passed_to_sql_as_json_arrays(mock_prisma: MagicMock): await _activity(ADMIN, key_aliases=["my-key"], models=["gpt-5.1", "claude-opus-4-8"]) filtered_calls = [ - call for call in mock_prisma.db.query_raw.call_args_list if call.args[0] in (GROUPS_SQL, ERROR_BREAKDOWN_SQL) + call + for call in _spend_log_reads(mock_prisma).call_args_list + if call.args[0] in (GROUPS_SQL, ERROR_BREAKDOWN_SQL) ] assert len(filtered_calls) == 2 for call in filtered_calls: @@ -135,7 +143,9 @@ async def test_every_query_excludes_the_same_info_routes(mock_prisma: MagicMock) the SQL does with it is covered against Postgres in tests/integration/spend/test_cache_activity.py.""" await _activity(ADMIN) - exclusions_by_query = {call.args[0]: json.loads(call.args[-4]) for call in mock_prisma.db.query_raw.call_args_list} + exclusions_by_query = { + call.args[0]: json.loads(call.args[-4]) for call in _spend_log_reads(mock_prisma).call_args_list + } assert set(exclusions_by_query) == {GROUPS_SQL, ERROR_BREAKDOWN_SQL, KEY_ALIAS_OPTIONS_SQL, MODEL_OPTIONS_SQL} for excluded_call_types in exclusions_by_query.values(): assert excluded_call_types == LiteLLMRoutes.info_routes.value @@ -147,7 +157,7 @@ async def test_every_query_excludes_the_same_info_routes(mock_prisma: MagicMock) async def test_admin_queries_grant_all_rows(mock_prisma: MagicMock): await _activity(ADMIN) - scope_by_query = {call.args[0]: call.args[-3:] for call in mock_prisma.db.query_raw.call_args_list} + scope_by_query = {call.args[0]: call.args[-3:] for call in _spend_log_reads(mock_prisma).call_args_list} assert set(scope_by_query) == {GROUPS_SQL, ERROR_BREAKDOWN_SQL, KEY_ALIAS_OPTIONS_SQL, MODEL_OPTIONS_SQL} for scope_args in scope_by_query.values(): assert scope_args == (True, None, "[]") @@ -157,7 +167,7 @@ async def test_admin_queries_grant_all_rows(mock_prisma: MagicMock): async def test_internal_user_queries_are_scoped_to_own_and_team_rows(mock_prisma: MagicMock): await _activity(INTERNAL_USER, team_lookup=_team_nine) - scope_by_query = {call.args[0]: call.args[-3:] for call in mock_prisma.db.query_raw.call_args_list} + scope_by_query = {call.args[0]: call.args[-3:] for call in _spend_log_reads(mock_prisma).call_args_list} assert set(scope_by_query) == {GROUPS_SQL, ERROR_BREAKDOWN_SQL, KEY_ALIAS_OPTIONS_SQL, MODEL_OPTIONS_SQL} for scope_args in scope_by_query.values(): assert scope_args == (False, "user-1", json.dumps(["team-9"])) @@ -168,7 +178,7 @@ async def test_key_without_user_id_grants_no_rows(mock_prisma: MagicMock): no_user = UserAPIKeyAuth(user_role=LitellmUserRoles.TEAM, user_id=None) await _activity(no_user) - for call in mock_prisma.db.query_raw.call_args_list: + for call in _spend_log_reads(mock_prisma).call_args_list: assert call.args[-3:] == (False, None, "[]") @@ -178,7 +188,27 @@ async def test_rejects_malformed_dates_with_400(mock_prisma: MagicMock): await _activity(ADMIN, start_date="07/01/2026") assert exc_info.value.status_code == 400 - mock_prisma.db.query_raw.assert_not_called() + _spend_log_reads(mock_prisma).assert_not_called() + + +@pytest.mark.asyncio +async def test_every_cache_activity_read_is_cancelled_by_postgres_after_30s(mock_prisma: MagicMock): + await _activity(ADMIN) + + transaction = mock_prisma.read_db.tx.return_value + assert transaction.statements == ("SET LOCAL statement_timeout = 30000",) * 4 + + +@pytest.mark.asyncio +async def test_cancelled_cache_activity_read_is_reported_as_503_with_a_date_range_hint(monkeypatch: pytest.MonkeyPatch): + prisma = build_prisma(AsyncMock(side_effect=statement_timeout_error())) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", prisma) + + with pytest.raises(HTTPException) as exc_info: + await _activity(ADMIN) + + assert exc_info.value.status_code == 503 + assert exc_info.value.detail == {"error": "Cache activity took longer than 30s. Pick a shorter date range"} def test_totals_ratio_is_zero_without_requests(): diff --git a/tests/unit/proxy/common_utils/test_check_batch_cost.py b/tests/unit/proxy/common_utils/test_check_batch_cost.py index e5f289f2944..aff145f1e27 100644 --- a/tests/unit/proxy/common_utils/test_check_batch_cost.py +++ b/tests/unit/proxy/common_utils/test_check_batch_cost.py @@ -1774,6 +1774,53 @@ class TestUnmanagedVertexRouting: assert result == ("deploy-ft", "8823717160934178816") + def _two_team_deployments(self): + """Two Vertex deployments sharing the base model, team A listed first.""" + return [ + { + "model_name": f"gemini-2.5-flash-team-{team}", + "litellm_params": {"model": "vertex_ai/gemini-2.5-flash", "custom_llm_provider": "vertex_ai"}, + "model_info": {"id": f"deploy-team-{team}"}, + } + for team in ("a", "b") + ] + + def test_flag_on_prefers_the_deployment_stored_on_the_row_over_the_first_base_model_match(self): + """Regression (LIT-10326): two deployments share gemini-2.5-flash with different + buckets. A batch created on team B's deployment stores deploy-team-b in file_object, + so the poller must read it back through team B, not the first match (team A).""" + router: Final = MagicMock() + router.resolve_model_name_from_model_id.return_value = None + router.get_model_list.return_value = self._two_team_deployments() + router.get_deployment = MagicMock(return_value=self._vertex_deployment()) + instance: Final = self._instance(track_unmanaged=True, router=router) + job: Final = self._job( + file_object=json.dumps({**json.loads(_unmanaged_vertex_file_object()), "model_id": "deploy-team-b"}) + ) + + with patch(_IS_B64, return_value=False): + result: Final = instance._resolve_job_routing(job, MagicMock()) + + assert result == ("deploy-team-b", "8823717160934178816") + router.get_deployment.assert_called_once_with(model_id="deploy-team-b") + + def test_flag_on_falls_back_to_base_model_match_when_the_stored_deployment_is_gone(self): + """The stored deployment was removed from the config after the batch was created: + the row still routes through today's base-model lookup instead of being skipped.""" + router: Final = MagicMock() + router.resolve_model_name_from_model_id.return_value = None + router.get_model_list.return_value = self._two_team_deployments() + router.get_deployment = MagicMock(return_value=None) + instance: Final = self._instance(track_unmanaged=True, router=router) + job: Final = self._job( + file_object=json.dumps({**json.loads(_unmanaged_vertex_file_object()), "model_id": "deploy-removed"}) + ) + + with patch(_IS_B64, return_value=False): + result: Final = instance._resolve_job_routing(job, MagicMock()) + + assert result == ("deploy-team-a", "8823717160934178816") + def test_flag_on_skips_non_vertex_deployment_sharing_model_group(self): """Flag on, but the only deployment for the model group is a non-vertex_ai provider: must not be selected, even though the model group name matches.""" diff --git a/tests/unit/proxy/common_utils/test_http_parsing_utils.py b/tests/unit/proxy/common_utils/test_http_parsing_utils.py index cd5820fa2fd..beb2e35f823 100644 --- a/tests/unit/proxy/common_utils/test_http_parsing_utils.py +++ b/tests/unit/proxy/common_utils/test_http_parsing_utils.py @@ -7,6 +7,7 @@ from unittest.mock import AsyncMock, MagicMock, patch import orjson import pytest +from fastapi import HTTPException from fastapi.testclient import TestClient from starlette.datastructures import FormData from starlette.requests import Request @@ -1592,3 +1593,44 @@ async def test_form_data_metadata_keeps_non_ascii_and_valid_surrogate_pairs(): result = await read_request_body(mock_request) assert result["metadata"]["tags"] == ["café", "😀"] + + +def _request_whose_client_disconnects_mid_upload(content_type: str) -> Request: + scope: Final = { + "type": "http", + "method": "POST", + "path": "/v1/chat/completions", + "headers": [(b"content-type", content_type.encode()), (b"content-length", b"96")], + "query_string": b"", + } + + async def receive() -> Message: + return {"type": "http.disconnect"} + + return Request(scope, receive) + + +@pytest.mark.parametrize( + "content_type", + ("application/json", "multipart/form-data; boundary=synthetic-boundary", "application/x-protobuf"), +) +async def test_read_request_body_aborts_with_499_when_the_client_disconnects_mid_upload(content_type: str) -> None: + request: Final = _request_whose_client_disconnects_mid_upload(content_type) + with pytest.raises(HTTPException) as raised: + await read_request_body(request) + assert (raised.value.status_code, raised.value.detail) == (499, "Client disconnected the request") + assert _safe_get_request_parsed_body(request) is None, "a half-received body must never be cached as {}" + + +async def test_read_request_body_closes_the_otel_server_span_when_the_client_disconnects_mid_upload( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from litellm.proxy import proxy_server + + span: Final = MagicMock() + monkeypatch.setattr(proxy_server, "open_telemetry_logger", MagicMock(), raising=False) + request: Final = _request_whose_client_disconnects_mid_upload("application/json") + request.state.parent_otel_span = span + with pytest.raises(HTTPException): + await read_request_body(request) + assert (span.end.called, request.state.parent_otel_span) == (True, None), "the auth-time span must not dangle" diff --git a/tests/unit/proxy/db/db_transaction_queue/test_redis_update_buffer.py b/tests/unit/proxy/db/db_transaction_queue/test_redis_update_buffer.py index 5ecc8784384..76a2dbf73f1 100644 --- a/tests/unit/proxy/db/db_transaction_queue/test_redis_update_buffer.py +++ b/tests/unit/proxy/db/db_transaction_queue/test_redis_update_buffer.py @@ -5,7 +5,6 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest - from litellm.proxy.db.db_transaction_queue.redis_update_buffer import RedisUpdateBuffer from litellm.proxy.proxy_server import ProxyStartupEvent @@ -701,3 +700,18 @@ async def test_store_spend_logs_in_redis_is_off_unless_transaction_buffering_is_ assert await buffer.store_spend_logs_in_redis([{"request_id": "a"}]) is False assert redis.rows == [] + + +@pytest.mark.asyncio +async def test_tool_discoveries_go_back_on_the_queue_when_the_redis_push_fails(): + from litellm.proxy.db.db_transaction_queue.tool_discovery_queue import ToolDiscoveryQueue + + redis_cache = MagicMock() + redis_cache.async_rpush_pipeline = AsyncMock(side_effect=ConnectionError("redis down")) + queue = ToolDiscoveryQueue() + queue.add_update({"tool_name": "get_weather", "origin": "user_defined"}) + queue.add_update({"tool_name": "get_weather", "origin": "user_defined"}) + + await RedisUpdateBuffer(redis_cache=redis_cache).store_tool_discoveries_in_redis(queue) + + assert queue.flush() == [{"tool_name": "get_weather", "origin": "user_defined", "call_count": 2}] diff --git a/tests/unit/proxy/db/db_transaction_queue/test_tool_discovery_queue.py b/tests/unit/proxy/db/db_transaction_queue/test_tool_discovery_queue.py index e400ad16e84..0a015255f2e 100644 --- a/tests/unit/proxy/db/db_transaction_queue/test_tool_discovery_queue.py +++ b/tests/unit/proxy/db/db_transaction_queue/test_tool_discovery_queue.py @@ -5,9 +5,9 @@ Unit tests for ToolDiscoveryQueue. import pytest - from litellm.proxy.db.db_transaction_queue.tool_discovery_queue import ( ToolDiscoveryQueue, + aggregate_tool_discoveries, ) @@ -16,57 +16,50 @@ def queue(): return ToolDiscoveryQueue() -def test_add_single_tool(queue): +def test_a_tool_seen_once_flushes_with_call_count_one(queue): queue.add_update({"tool_name": "my_tool", "origin": "user_defined"}) - items = queue.flush() - assert len(items) == 1 - assert items[0]["tool_name"] == "my_tool" - assert items[0]["origin"] == "user_defined" + assert queue.flush() == [{"tool_name": "my_tool", "origin": "user_defined", "call_count": 1}] -def test_deduplication_same_name(queue): - """Adding the same tool_name twice should only keep the first.""" - queue.add_update({"tool_name": "tool_a", "origin": "mcp_server"}) - queue.add_update({"tool_name": "tool_a", "origin": "user_defined"}) - items = queue.flush() - assert len(items) == 1 - assert items[0]["origin"] == "mcp_server" # first wins +def test_repeated_sightings_fold_into_one_row_counting_every_call(queue): + queue.add_update({"tool_name": "tool_a", "origin": "mcp_server", "created_by": "alice"}) + queue.add_update({"tool_name": "tool_a", "origin": "user_defined", "created_by": "bob"}) + queue.add_update({"tool_name": "tool_a"}) + assert queue.flush() == [{"tool_name": "tool_a", "origin": "mcp_server", "created_by": "alice", "call_count": 3}] -def test_deduplication_different_names(queue): +def test_distinct_tools_each_get_their_own_row(queue): queue.add_update({"tool_name": "tool_a"}) queue.add_update({"tool_name": "tool_b"}) - items = queue.flush() - assert len(items) == 2 - names = {i["tool_name"] for i in items} - assert names == {"tool_a", "tool_b"} - - -def test_flush_clears_pending(queue): - queue.add_update({"tool_name": "tool_x"}) - items1 = queue.flush() - assert len(items1) == 1 - items2 = queue.flush() - assert len(items2) == 0 - - -def test_seen_names_reset_after_flush(queue): - """Seen-set is cleared on flush so the same tool can re-enter the next cycle.""" queue.add_update({"tool_name": "tool_a"}) - queue.flush() - queue.add_update({"tool_name": "tool_a"}) # same tool, new cycle - items = queue.flush() - assert len(items) == 1 - assert items[0]["tool_name"] == "tool_a" + assert queue.flush() == [ + {"tool_name": "tool_a", "call_count": 2}, + {"tool_name": "tool_b", "call_count": 1}, + ] + + +def test_flush_clears_pending_and_the_next_cycle_counts_from_zero(queue): + queue.add_update({"tool_name": "tool_x"}) + queue.add_update({"tool_name": "tool_x"}) + assert queue.flush() == [{"tool_name": "tool_x", "call_count": 2}] + assert queue.flush() == [] + queue.add_update({"tool_name": "tool_x"}) + assert queue.flush() == [{"tool_name": "tool_x", "call_count": 1}] def test_empty_tool_name_ignored(queue): queue.add_update({"tool_name": ""}) queue.add_update({"tool_name": None}) # type: ignore[arg-type] - items = queue.flush() - assert len(items) == 0 + assert queue.flush() == [] -def test_flush_returns_list(queue): - result = queue.flush() - assert isinstance(result, list) +def test_aggregate_sums_call_counts_across_workers_flushes(): + worker_a = [{"tool_name": "get_weather", "origin": "user_defined", "call_count": 3}] + worker_b = [ + {"tool_name": "get_weather", "origin": "mcp_server", "call_count": 2}, + {"tool_name": "search", "call_count": 1}, + ] + assert aggregate_tool_discoveries([*worker_a, *worker_b]) == ( + {"tool_name": "get_weather", "origin": "user_defined", "call_count": 5}, + {"tool_name": "search", "call_count": 1}, + ) diff --git a/tests/unit/proxy/db/fake_prisma_engine.py b/tests/unit/proxy/db/fake_prisma_engine.py index 4221eeced5a..082659a3adc 100644 --- a/tests/unit/proxy/db/fake_prisma_engine.py +++ b/tests/unit/proxy/db/fake_prisma_engine.py @@ -2,8 +2,10 @@ marking the DB I/O witness the way ``_TrackedPrismaEngine`` does, so the producer under test emits its service event.""" -from typing import TypeVar -from unittest.mock import AsyncMock +from typing import Final, TypeVar +from unittest.mock import AsyncMock, MagicMock + +from prisma.errors import DataError from litellm.proxy.db.log_db_metrics import record_db_io @@ -16,3 +18,73 @@ def engine_call(return_value: _T | None = None) -> AsyncMock: return return_value return AsyncMock(side_effect=run) + + +class RecordingReadTransaction: + """The transaction ``query_raw_within_statement_timeout`` opens on ``read_db``, keeping its ``SET LOCAL`` statements.""" + + def __init__(self, query_raw: AsyncMock) -> None: + self.statements: tuple[str, ...] = () + self.query_raw: AsyncMock = query_raw + + async def execute_raw(self, sql: str) -> int: + self.statements = (*self.statements, sql) + return 0 + + async def __aenter__(self) -> "RecordingReadTransaction": + return self + + async def __aexit__(self, *exc_info: object) -> None: + pass + + +def statement_timeout_reads(prisma: MagicMock, query_raw: AsyncMock) -> RecordingReadTransaction: + """Make ``prisma.read_db.tx()`` hand out one recording transaction whose raw reads answer with ``query_raw``.""" + transaction: Final = RecordingReadTransaction(query_raw) + prisma.read_db.tx = MagicMock(return_value=transaction) + return transaction + + +def statement_timeout_error() -> DataError: + """The error prisma raises when Postgres cancels a statement under ``statement_timeout`` (SQLSTATE 57014).""" + message: Final = "canceling statement due to statement timeout" + return DataError( + data={ + "user_facing_error": { + "is_panic": False, + "message": f"Error querying the database: {message}", + "meta": {"code": "57014", "message": message}, + } + } + ) + + +class _DelegatingReadTransaction: + """``read_db.tx()`` of a hand-written prisma fake: raw reads go to the fake's ``db.query_raw``.""" + + def __init__(self, db: object) -> None: + self.db: Final = db + self.statements: tuple[str, ...] = () + + async def execute_raw(self, sql: str) -> int: + self.statements = (*self.statements, sql) + return 0 + + async def query_raw(self, sql: str, *args: object) -> object: + return await getattr(self.db, "query_raw")(sql, *args) + + async def __aenter__(self) -> "_DelegatingReadTransaction": + return self + + async def __aexit__(self, *exc_info: object) -> None: + pass + + +class ReadTargetOf: + """What a prisma fake returns from ``read_db``: ``tx()`` opens a transaction over the fake's ``db``.""" + + def __init__(self, db: object) -> None: + self.db: Final = db + + def tx(self, **kwargs: object) -> _DelegatingReadTransaction: + return _DelegatingReadTransaction(self.db) diff --git a/tests/unit/proxy/db/test_db_spend_update_writer.py b/tests/unit/proxy/db/test_db_spend_update_writer.py index d440b99a632..b2f006aa52f 100644 --- a/tests/unit/proxy/db/test_db_spend_update_writer.py +++ b/tests/unit/proxy/db/test_db_spend_update_writer.py @@ -12,6 +12,8 @@ from unittest.mock import AsyncMock, MagicMock, call, patch import httpx import pytest + +from litellm.types.caching import RedisPipelineLpopOperation, RedisPipelineRpushOperation from prisma.errors import DataError as PrismaDataError from prisma.errors import RawQueryError from redis.exceptions import DataError @@ -2365,6 +2367,33 @@ async def test_commit_spend_updates_uses_pipeline(): mock_redis_update_buffer.get_all_daily_tag_spend_update_transactions_from_redis_buffer.assert_not_called() +class _InMemoryRedis: + def __init__(self) -> None: + self.store: dict[str, str] = {} + self.lists: dict[str, list[object]] = {} + + async def async_set_cache(self, key: str, value: str, nx: bool = False, **kwargs: object) -> bool | None: + if nx and key in self.store: + return None + self.store[key] = json.dumps(value) + return True + + async def async_get_cache(self, key: str, **kwargs: object) -> str | None: + raw = self.store.get(key) + return json.loads(raw) if raw is not None else None + + async def async_delete_cache(self, key: str) -> int: + return 1 if self.store.pop(key, None) is not None else 0 + + async def async_rpush_pipeline(self, rpush_list: list[RedisPipelineRpushOperation]) -> list[int]: + for op in rpush_list: + self.lists.setdefault(op["key"], []).extend(op["values"]) + return [len(self.lists[op["key"]]) for op in rpush_list] + + async def async_lpop_pipeline(self, lpop_list: list[RedisPipelineLpopOperation]) -> list[list[object] | None]: + return [self.lists.pop(op["key"], None) for op in lpop_list] + + @pytest.mark.asyncio async def test_commit_with_redis_requeues_all_on_db_failure(): """ @@ -3213,7 +3242,6 @@ async def test_failed_window_spend_commit_requeues_the_increments_and_continues_ await db_writer.window_spend_update_queue.add_update(transaction) db = _WindowSpendFakeDB() db.query_raw = AsyncMock(side_effect=Exception("connection reset")) - db_writer._flush_tool_discovery_queue = AsyncMock() await db_writer._commit_spend_updates_to_db_without_redis_buffer( prisma_client=_WindowSpendFakePrisma(db), @@ -3221,7 +3249,6 @@ async def test_failed_window_spend_commit_requeues_the_increments_and_continues_ proxy_logging_obj=MagicMock(), ) - db_writer._flush_tool_discovery_queue.assert_called_once() requeued = await db_writer.window_spend_update_queue.flush_and_get_aggregated_window_spend_transactions() assert requeued == (transaction,) @@ -3297,7 +3324,6 @@ async def test_failed_daily_spend_commit_is_requeued_only_when_the_rows_are_prov db_writer = DBSpendUpdateWriter() await db_writer.daily_spend_update_queue.add_update({"user-key": _daily_txn(user_id="user-1")}) db = _DailySpendFakeDB(failing_table="LiteLLM_DailyUserSpend", failure=failure) - db_writer._flush_tool_discovery_queue = AsyncMock() proxy_logging_obj = MagicMock() proxy_logging_obj.failure_handler = AsyncMock() @@ -3323,7 +3349,6 @@ async def test_failed_daily_spend_commit_drops_only_the_batch_that_was_sent(): {f"user-{i:03d}": _daily_txn(user_id=f"user-{i:03d}") for i in range(150)} ) db = _DailySpendFakeDB(failing_table="LiteLLM_DailyUserSpend", failure=httpx.ReadTimeout("no reply")) - db_writer._flush_tool_discovery_queue = AsyncMock() proxy_logging_obj = MagicMock() proxy_logging_obj.failure_handler = AsyncMock() @@ -3351,7 +3376,6 @@ async def test_failed_daily_spend_commit_requeues_the_rows_and_flushes_the_other team_txn = {key: value for key, value in _daily_txn().items() if key != "user_id"} | {"team_id": "team-1"} await db_writer.daily_team_spend_update_queue.add_update({"team-key": team_txn}) db = _DailySpendFakeDB(failing_table="LiteLLM_DailyUserSpend") - db_writer._flush_tool_discovery_queue = AsyncMock() proxy_logging_obj = MagicMock() proxy_logging_obj.failure_handler = AsyncMock() @@ -3362,7 +3386,6 @@ async def test_failed_daily_spend_commit_requeues_the_rows_and_flushes_the_other assert _daily_upserts(db, "LiteLLM_DailyUserSpend") == [] (team_upsert,) = _daily_upserts(db, "LiteLLM_DailyTeamSpend") assert _row_values(team_upsert, "team_id") == ["team-1"] - db_writer._flush_tool_discovery_queue.assert_called_once() db.failing_table = None await db_writer._commit_spend_updates_to_db_without_redis_buffer( @@ -3589,7 +3612,6 @@ async def test_failed_window_spend_commit_is_requeued_only_when_the_rows_are_pro ) await db_writer.window_spend_update_queue.add_update(transaction) db = _WindowSpendFailureDB(failure) - db_writer._flush_tool_discovery_queue = AsyncMock() await db_writer._commit_spend_updates_to_db_without_redis_buffer( prisma_client=_WindowSpendFakePrisma(db), diff --git a/tests/unit/proxy/db/test_exception_handler.py b/tests/unit/proxy/db/test_exception_handler.py index 75346ca08a6..a4c16c0949b 100644 --- a/tests/unit/proxy/db/test_exception_handler.py +++ b/tests/unit/proxy/db/test_exception_handler.py @@ -951,3 +951,18 @@ def test_is_lock_timeout_error_recognises_a_statement_cancelled_under_lock_timeo ) def test_is_lock_timeout_error_excludes_other_failures(error: Exception) -> None: assert PrismaDBExceptionHandler.is_lock_timeout_error(error) is False + + +@pytest.mark.parametrize( + ("error", "expected"), + [ + (_lock_timeout_error(message="canceling statement due to statement timeout", code="57014"), True), + (_lock_timeout_error(message="canceling statement due to statement timeout", code=None), True), + (_lock_timeout_error(), False), + (_capacity_error("sorry, too many clients already"), False), + (RuntimeError("canceling statement due to statement timeout"), False), + ], + ids=["meta_57014", "message_only", "lock_timeout_55P03", "capacity_53300", "not_prisma"], +) +def test_is_statement_timeout_error(error: Exception, expected: bool) -> None: + assert PrismaDBExceptionHandler.is_statement_timeout_error(error) is expected diff --git a/tests/unit/proxy/db/test_prisma_query_span.py b/tests/unit/proxy/db/test_prisma_query_span.py index 9f2ba28ea7e..55ba03d9f2c 100644 --- a/tests/unit/proxy/db/test_prisma_query_span.py +++ b/tests/unit/proxy/db/test_prisma_query_span.py @@ -19,6 +19,7 @@ from litellm.proxy.db.prisma_query_span import UNKNOWN_PRISMA_QUERY, parse_prism _REPO: Final = Path(__file__).resolve().parents[4] _SOURCE_ROOTS: Final = ("litellm", "enterprise", "litellm-proxy-extras") _RAW_METHODS: Final = frozenset({"query_first", "query_raw", "execute_raw"}) +_CROSS_MODULE_FORWARDERS: Final = frozenset({"query_raw_within_statement_timeout"}) _MODEL_METHODS: Final = frozenset( { "find_unique", @@ -297,7 +298,7 @@ def _parameter_site( """A statement that arrives as a parameter: a ``query_raw`` forwarder adds no round trip of its own, any other helper is named by what its callers in the module hand it.""" function: Final = next(p for p in parents if isinstance(p, (ast.AsyncFunctionDef, ast.FunctionDef))) - if function.name in _RAW_METHODS: + if function.name in _RAW_METHODS or function.name in _CROSS_MODULE_FORWARDERS: return _PrismaCallSite(location, method, "forwarder", f"(callers of {function.name})") callers: Final = tuple( node diff --git a/tests/unit/proxy/db/test_tool_registry_writer.py b/tests/unit/proxy/db/test_tool_registry_writer.py index c9df665741d..76898b0eeab 100644 --- a/tests/unit/proxy/db/test_tool_registry_writer.py +++ b/tests/unit/proxy/db/test_tool_registry_writer.py @@ -3,22 +3,28 @@ Unit tests for tool_registry_writer.py — uses a mock prisma client that exposes litellm_tooltable.upsert / find_many / find_unique. """ +import json from datetime import datetime, timezone from unittest.mock import AsyncMock, MagicMock import pytest -from prisma.errors import PrismaError - +from prisma.errors import PrismaError, RawQueryError +from litellm.constants import DB_SPEND_UPDATE_JOB_NAME, REDIS_TOOL_DISCOVERY_BUFFER_KEY, TOOL_DISCOVERY_COMMIT_JOB_NAME +from litellm.proxy.db.db_transaction_queue.pod_lock_manager import PodLockManager +from litellm.proxy.db.db_transaction_queue.redis_update_buffer import RedisUpdateBuffer +from litellm.proxy.db.db_transaction_queue.tool_discovery_queue import ToolDiscoveryQueue from litellm.proxy.db.tool_registry_writer import ( ToolPolicyRegistry, batch_upsert_tools, + commit_tool_discoveries, get_tool, get_tool_policy_registry, get_tools_by_names, list_tools, update_tool_policy, ) +from litellm.types.caching import RedisPipelineLpopOperation, RedisPipelineRpushOperation def _mock_row(**kwargs): @@ -62,6 +68,7 @@ def _make_prisma( prisma = MagicMock() prisma.db.litellm_tooltable = MagicMock() prisma.db.litellm_tooltable.upsert = AsyncMock(return_value=upsert_return) + prisma.db.execute_raw = AsyncMock(return_value=1) prisma.db.litellm_tooltable.find_many = AsyncMock( return_value=find_many_rows if find_many_rows is not None else [] ) @@ -73,28 +80,32 @@ def _make_prisma( return prisma +def _statement(prisma) -> tuple[str, tuple[object, ...]]: + prisma.db.execute_raw.assert_awaited_once() + sql, *params = prisma.db.execute_raw.call_args.args + return sql, tuple(params) + + @pytest.mark.asyncio -async def test_batch_upsert_tools_calls_upsert(): +async def test_batch_upsert_tools_inserts_a_new_tool_with_the_batch_call_count_and_adds_it_on_conflict(): prisma = _make_prisma() - items = [{"tool_name": "tool_a", "origin": "mcp_server", "created_by": None}] - await batch_upsert_tools(prisma, items) - prisma.db.litellm_tooltable.upsert.assert_awaited_once() - call_kw = prisma.db.litellm_tooltable.upsert.call_args.kwargs - assert call_kw["where"] == {"tool_name": "tool_a"} - assert call_kw["data"]["create"]["tool_name"] == "tool_a" - assert call_kw["data"]["create"]["origin"] == "mcp_server" - assert call_kw["data"]["create"]["input_policy"] == "untrusted" - assert call_kw["data"]["create"]["output_policy"] == "untrusted" - assert call_kw["data"]["create"]["call_count"] == 1 - assert call_kw["data"]["update"]["call_count"] == {"increment": 1} - assert "updated_at" in call_kw["data"]["update"] + await batch_upsert_tools( + prisma, [{"tool_name": "tool_a", "origin": "mcp_server", "created_by": None, "call_count": 4}] + ) + sql, params = _statement(prisma) + assert sql.startswith('INSERT INTO "LiteLLM_ToolTable"') + assert 'ON CONFLICT ("tool_name") DO UPDATE SET' in sql + assert '"call_count" = "LiteLLM_ToolTable"."call_count" + EXCLUDED."call_count"' in sql + assert '"input_policy"' not in sql, "policies come from the table defaults and are never overwritten" + assert params[1:] == ("tool_a", "mcp_server", 4, "system", "system", None, None, None, None) + prisma.db.litellm_tooltable.upsert.assert_not_awaited() @pytest.mark.asyncio async def test_batch_upsert_tools_empty_list(): prisma = _make_prisma() await batch_upsert_tools(prisma, []) - prisma.db.litellm_tooltable.upsert.assert_not_awaited() + prisma.db.execute_raw.assert_not_awaited() @pytest.mark.asyncio @@ -102,21 +113,22 @@ async def test_batch_upsert_tools_skips_empty_names(): prisma = _make_prisma() items = [{"tool_name": "", "origin": None}, {"tool_name": None}] # type: ignore[list-item] await batch_upsert_tools(prisma, items) - prisma.db.litellm_tooltable.upsert.assert_not_awaited() + prisma.db.execute_raw.assert_not_awaited() @pytest.mark.asyncio -async def test_batch_upsert_multiple_tools_calls_upsert_per_tool(): +async def test_batch_upsert_many_tools_is_one_statement_with_one_row_per_distinct_tool(): prisma = _make_prisma() items = [ - {"tool_name": "tool_a", "origin": "mcp_server", "created_by": None}, - {"tool_name": "tool_b", "origin": "user_defined", "created_by": "alice"}, + {"tool_name": "tool_a", "origin": "mcp_server", "created_by": None, "call_count": 2}, + {"tool_name": "tool_b", "origin": "user_defined", "created_by": "alice", "call_count": 1}, + {"tool_name": "tool_a", "origin": "user_defined", "created_by": "bob", "call_count": 5}, ] await batch_upsert_tools(prisma, items) - assert prisma.db.litellm_tooltable.upsert.await_count == 2 - calls = prisma.db.litellm_tooltable.upsert.call_args_list - assert calls[0].kwargs["where"]["tool_name"] == "tool_a" - assert calls[1].kwargs["where"]["tool_name"] == "tool_b" + sql, params = _statement(prisma) + assert sql.count("$1::text") == 1 and "$11::text" in sql and "$21::text" not in sql + assert params[1:10] == ("tool_a", "mcp_server", 7, "system", "system", None, None, None, None) + assert params[11:20] == ("tool_b", "user_defined", 1, "alice", "alice", None, None, None, None) @pytest.mark.asyncio @@ -264,20 +276,14 @@ async def test_update_tool_policy_calls_upsert_then_get_tool(): assert call_kw["where"] == {"tool_name": "my_tool"} assert call_kw["data"]["update"]["input_policy"] == "blocked" assert call_kw["data"]["update"]["updated_by"] == "admin" - prisma.db.litellm_tooltable.find_unique.assert_awaited_with( - where={"tool_name": "my_tool"} - ) + prisma.db.litellm_tooltable.find_unique.assert_awaited_with(where={"tool_name": "my_tool"}) @pytest.mark.asyncio async def test_get_tools_by_names_returns_policy_map(): rows = [ - _mock_row( - tool_name="tool_a", input_policy="trusted", output_policy="untrusted" - ), - _mock_row( - tool_name="tool_b", input_policy="blocked", output_policy="untrusted" - ), + _mock_row(tool_name="tool_a", input_policy="trusted", output_policy="untrusted"), + _mock_row(tool_name="tool_b", input_policy="blocked", output_policy="untrusted"), ] prisma = _make_prisma(find_many_rows=rows) result = await get_tools_by_names(prisma, ["tool_a", "tool_b"]) @@ -285,9 +291,7 @@ async def test_get_tools_by_names_returns_policy_map(): "tool_a": ("trusted", "untrusted"), "tool_b": ("blocked", "untrusted"), } - prisma.db.litellm_tooltable.find_many.assert_awaited_once_with( - where={"tool_name": {"in": ["tool_a", "tool_b"]}} - ) + prisma.db.litellm_tooltable.find_many.assert_awaited_once_with(where={"tool_name": {"in": ["tool_a", "tool_b"]}}) @pytest.mark.asyncio @@ -382,12 +386,8 @@ async def test_sync_tool_policy_from_db_retries_on_transport_error_first_read(): return [] mock_prisma_client = MagicMock() - mock_prisma_client.db.litellm_tooltable.find_many = AsyncMock( - side_effect=_flaky_find_many - ) - mock_prisma_client.db.litellm_objectpermissiontable.find_many = AsyncMock( - return_value=[] - ) + mock_prisma_client.db.litellm_tooltable.find_many = AsyncMock(side_effect=_flaky_find_many) + mock_prisma_client.db.litellm_objectpermissiontable.find_many = AsyncMock(return_value=[]) mock_prisma_client.attempt_db_reconnect = AsyncMock(return_value=True) mock_prisma_client._db_auth_reconnect_timeout_seconds = 2.0 mock_prisma_client._db_auth_reconnect_lock_timeout_seconds = 0.1 @@ -397,10 +397,7 @@ async def test_sync_tool_policy_from_db_retries_on_transport_error_first_read(): assert len(invocations) == 2 mock_prisma_client.attempt_db_reconnect.assert_awaited_once() reconnect_kwargs = mock_prisma_client.attempt_db_reconnect.await_args.kwargs - assert ( - reconnect_kwargs["reason"] - == "sync_tool_policy_from_db_tools_lookup_failure" - ) + assert reconnect_kwargs["reason"] == "sync_tool_policy_from_db_tools_lookup_failure" assert registry.is_initialized() @@ -421,9 +418,7 @@ async def test_sync_tool_policy_from_db_retries_on_transport_error_second_read() mock_prisma_client = MagicMock() mock_prisma_client.db.litellm_tooltable.find_many = AsyncMock(return_value=[]) - mock_prisma_client.db.litellm_objectpermissiontable.find_many = AsyncMock( - side_effect=_flaky_perms_find_many - ) + mock_prisma_client.db.litellm_objectpermissiontable.find_many = AsyncMock(side_effect=_flaky_perms_find_many) mock_prisma_client.attempt_db_reconnect = AsyncMock(return_value=True) mock_prisma_client._db_auth_reconnect_timeout_seconds = 2.0 mock_prisma_client._db_auth_reconnect_lock_timeout_seconds = 0.1 @@ -433,7 +428,142 @@ async def test_sync_tool_policy_from_db_retries_on_transport_error_second_read() assert len(perms_invocations) == 2 mock_prisma_client.attempt_db_reconnect.assert_awaited_once() reconnect_kwargs = mock_prisma_client.attempt_db_reconnect.await_args.kwargs - assert ( - reconnect_kwargs["reason"] - == "sync_tool_policy_from_db_perms_lookup_failure" + assert reconnect_kwargs["reason"] == "sync_tool_policy_from_db_perms_lookup_failure" + + +class _InMemoryRedis: + def __init__(self) -> None: + self.store: dict[str, str] = {} + self.lists: dict[str, list[object]] = {} + + async def async_set_cache(self, key: str, value: str, nx: bool = False, **kwargs: object) -> bool | None: + if nx and key in self.store: + return None + self.store[key] = json.dumps(value) + return True + + async def async_get_cache(self, key: str, **kwargs: object) -> str | None: + raw = self.store.get(key) + return json.loads(raw) if raw is not None else None + + async def async_delete_cache(self, key: str) -> int: + return 1 if self.store.pop(key, None) is not None else 0 + + async def async_rpush_pipeline(self, rpush_list: list[RedisPipelineRpushOperation]) -> list[int]: + for op in rpush_list: + self.lists.setdefault(op["key"], []).extend(op["values"]) + return [len(self.lists[op["key"]]) for op in rpush_list] + + async def async_lpop_pipeline(self, lpop_list: list[RedisPipelineLpopOperation]) -> list[list[object] | None]: + return [self.lists.pop(op["key"], None) for op in lpop_list] + + +def _postgres_rejection(sqlstate: str) -> RawQueryError: + return RawQueryError( + data={"user_facing_error": {"error_code": "P2010", "meta": {"code": sqlstate, "message": "db error"}}} ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "lock_held_by_another_pod", + [None, DB_SPEND_UPDATE_JOB_NAME, TOOL_DISCOVERY_COMMIT_JOB_NAME], + ids=["free", "spend_lock_held", "tool_lock_held"], +) +async def test_commit_with_redis_writes_every_workers_tools_in_one_statement_from_the_tool_lock_holder_only( + lock_held_by_another_pod: str | None, +): + redis = _InMemoryRedis() + if lock_held_by_another_pod is not None: + assert await PodLockManager(redis_cache=redis).acquire_lock(cronjob_id=lock_held_by_another_pod) + other_worker = ToolDiscoveryQueue() + other_worker.add_update({"tool_name": "get_weather", "origin": "mcp_server"}) + other_worker.add_update({"tool_name": "search", "origin": "user_defined"}) + await RedisUpdateBuffer(redis_cache=redis).store_tool_discoveries_in_redis(other_worker) + + queue = ToolDiscoveryQueue() + queue.add_update({"tool_name": "get_weather", "origin": "user_defined"}) + queue.add_update({"tool_name": "get_weather", "origin": "user_defined"}) + prisma_client = MagicMock() + prisma_client.db.execute_raw = AsyncMock(return_value=2) + prisma_client.db.litellm_tooltable.upsert = AsyncMock() + + await commit_tool_discoveries( + prisma_client=prisma_client, + tool_discovery_queue=queue, + redis_update_buffer=RedisUpdateBuffer(redis_cache=redis), + pod_lock_manager=PodLockManager(redis_cache=redis), + use_redis_buffer=True, + ) + + assert queue.flush() == [], "every worker hands its tools to redis each tick" + prisma_client.db.litellm_tooltable.upsert.assert_not_awaited() + if lock_held_by_another_pod == TOOL_DISCOVERY_COMMIT_JOB_NAME: + prisma_client.db.execute_raw.assert_not_awaited() + assert len(redis.lists[REDIS_TOOL_DISCOVERY_BUFFER_KEY]) == 2, "non-holders leave their tools for the holder" + return + assert REDIS_TOOL_DISCOVERY_BUFFER_KEY not in redis.lists, "the spend lock must not gate the tool drain" + prisma_client.db.execute_raw.assert_awaited_once() + sql, *params = prisma_client.db.execute_raw.call_args.args + assert sql.startswith('INSERT INTO "LiteLLM_ToolTable"') and 'ON CONFLICT ("tool_name")' in sql + rows = {params[i + 1]: params[i + 3] for i in range(0, len(params), 10)} + assert rows == {"get_weather": 3, "search": 1}, f"one row per tool summed across workers, got {params}" + + +@pytest.mark.asyncio +async def test_commit_without_redis_writes_only_this_workers_tools_in_one_statement_and_never_touches_redis(): + redis = _InMemoryRedis() + queue = ToolDiscoveryQueue() + queue.add_update({"tool_name": "get_weather", "origin": "user_defined"}) + queue.add_update({"tool_name": "get_weather", "origin": "user_defined"}) + queue.add_update({"tool_name": "search", "origin": "user_defined"}) + prisma_client = MagicMock() + prisma_client.db.execute_raw = AsyncMock(return_value=2) + + await commit_tool_discoveries( + prisma_client=prisma_client, + tool_discovery_queue=queue, + redis_update_buffer=RedisUpdateBuffer(redis_cache=redis), + pod_lock_manager=PodLockManager(redis_cache=redis), + use_redis_buffer=False, + ) + + assert redis.lists == {}, "without the buffer flag nothing goes through redis" + assert queue.flush() == [] + prisma_client.db.execute_raw.assert_awaited_once() + sql, *params = prisma_client.db.execute_raw.call_args.args + assert 'ON CONFLICT ("tool_name")' in sql + rows = {params[i + 1]: params[i + 3] for i in range(0, len(params), 10)} + assert rows == {"get_weather": 2, "search": 1}, f"got {params}" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("failure", "requeued"), + [ + (Exception("connection reset"), True), + (_postgres_rejection("23502"), False), + ], + ids=["connection_error_goes_back_to_redis", "rejected_row_is_dropped"], +) +async def test_failed_tool_table_write_requeues_the_rows_only_when_they_provably_did_not_apply( + failure: Exception, requeued: bool +): + redis = _InMemoryRedis() + queue = ToolDiscoveryQueue() + queue.add_update({"tool_name": "get_weather", "origin": "user_defined"}) + prisma_client = MagicMock() + prisma_client.db.execute_raw = AsyncMock(side_effect=failure) + + await commit_tool_discoveries( + prisma_client=prisma_client, + tool_discovery_queue=queue, + redis_update_buffer=RedisUpdateBuffer(redis_cache=redis), + pod_lock_manager=PodLockManager(redis_cache=redis), + use_redis_buffer=True, + ) + + prisma_client.db.execute_raw.assert_awaited_once() + left_in_redis = [json.loads(entry) for entry in redis.lists.get(REDIS_TOOL_DISCOVERY_BUFFER_KEY, [])] + expected = [[{"tool_name": "get_weather", "origin": "user_defined", "call_count": 1}]] if requeued else [] + assert left_in_redis == expected, f"requeued={requeued} but redis holds {left_in_redis}" diff --git a/tests/unit/proxy/db/test_tool_table_bulk_upsert.py b/tests/unit/proxy/db/test_tool_table_bulk_upsert.py new file mode 100644 index 00000000000..d1e79030f9d --- /dev/null +++ b/tests/unit/proxy/db/test_tool_table_bulk_upsert.py @@ -0,0 +1,38 @@ +"""Tests for the single-statement LiteLLM_ToolTable upsert (LIT-9997).""" + +from litellm.proxy.db.tool_table_bulk_upsert import build_tool_bulk_upsert + + +def test_one_row_per_item_and_the_conflict_update_only_moves_the_counters(): + sql, params = build_tool_bulk_upsert( + [ + {"tool_name": "tool_a", "origin": "mcp_server", "created_by": None, "call_count": 4}, + { + "tool_name": "tool_b", + "origin": "user_defined", + "created_by": "alice", + "call_count": 1, + "key_hash": "hash", + "team_id": "team", + "key_alias": "alias", + "user_agent": "agent", + }, + ] + ) + assert sql.startswith('INSERT INTO "LiteLLM_ToolTable"') + assert sql.count("$1::text") == 1 and "$11::text" in sql and "$21::text" not in sql + assert "$4::integer" in sql and "$14::integer" in sql + assert 'ON CONFLICT ("tool_name") DO UPDATE SET' in sql + assert '"call_count" = "LiteLLM_ToolTable"."call_count" + EXCLUDED."call_count"' in sql + assert '"last_used_at" = EXCLUDED."last_used_at"' in sql + assert '"input_policy"' not in sql and '"output_policy"' not in sql, "policies stay on the table defaults" + assert '"origin" =' not in sql and '"key_hash" =' not in sql, "an existing row keeps who discovered it" + assert len(params) == 20 + assert params[1:10] == ("tool_a", "mcp_server", 4, "system", "system", None, None, None, None) + assert params[11:20] == ("tool_b", "user_defined", 1, "alice", "alice", "hash", "team", "alias", "agent") + assert params[0] != params[10], "each row gets its own tool_id" + + +def test_defaults_fill_origin_created_by_and_call_count(): + _, params = build_tool_bulk_upsert([{"tool_name": "tool_a"}]) + assert params[1:6] == ("tool_a", "user_defined", 1, "system", "system") diff --git a/tests/unit/proxy/management_endpoints/test_prompt_caching_requests.py b/tests/unit/proxy/management_endpoints/test_prompt_caching_requests.py index 39dcc9630df..ea61177f773 100644 --- a/tests/unit/proxy/management_endpoints/test_prompt_caching_requests.py +++ b/tests/unit/proxy/management_endpoints/test_prompt_caching_requests.py @@ -1,5 +1,6 @@ from collections.abc import Mapping from typing import Final +from unittest.mock import AsyncMock, MagicMock import httpx import pytest @@ -8,6 +9,7 @@ from fastapi import FastAPI from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.management_endpoints.prompt_caching_requests import router +from tests.unit.proxy.db.fake_prisma_engine import statement_timeout_error, statement_timeout_reads pytestmark = pytest.mark.usefixtures("local_model_cost_map") @@ -67,3 +69,37 @@ async def test_incomplete_cursor_is_rejected( ) as client: response: Final = await client.get(_URL, params={"start_date": _START, "end_date": _END, **params}) assert response.status_code == 400 + + +@pytest.mark.asyncio +async def test_the_requests_read_is_cancelled_by_postgres_after_30s(monkeypatch: pytest.MonkeyPatch) -> None: + from litellm.proxy import proxy_server + + prisma: Final = MagicMock() + transaction: Final = statement_timeout_reads(prisma, AsyncMock(return_value=[])) + monkeypatch.setattr(proxy_server, "prisma_client", prisma) + async with httpx.AsyncClient( + transport=httpx.ASGITransport(app=_app(LitellmUserRoles.PROXY_ADMIN)), base_url="http://test" + ) as client: + response: Final = await client.get(_URL, params={"start_date": _START, "end_date": _END}) + + assert response.status_code == 200, response.text + assert transaction.statements == ("SET LOCAL statement_timeout = 30000",) + + +@pytest.mark.asyncio +async def test_a_cancelled_requests_read_is_a_503_with_a_date_range_hint(monkeypatch: pytest.MonkeyPatch) -> None: + from litellm.proxy import proxy_server + + prisma: Final = MagicMock() + statement_timeout_reads(prisma, AsyncMock(side_effect=statement_timeout_error())) + monkeypatch.setattr(proxy_server, "prisma_client", prisma) + async with httpx.AsyncClient( + transport=httpx.ASGITransport(app=_app(LitellmUserRoles.PROXY_ADMIN)), base_url="http://test" + ) as client: + response: Final = await client.get(_URL, params={"start_date": _START, "end_date": _END}) + + assert response.status_code == 503, response.text + assert response.json() == { + "detail": {"error": "Prompt caching requests took longer than 30s. Pick a shorter date range"} + } diff --git a/tests/unit/proxy/management_endpoints/test_team_endpoints.py b/tests/unit/proxy/management_endpoints/test_team_endpoints.py index a379b424aeb..0a85ba30aa1 100644 --- a/tests/unit/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_team_endpoints.py @@ -10,6 +10,7 @@ from unittest.mock import AsyncMock, MagicMock, PropertyMock, call, patch import httpx import litellm +from tests.unit.proxy.db.fake_prisma_engine import statement_timeout_error, statement_timeout_reads import pytest from fastapi import HTTPException, Request from fastapi.testclient import TestClient @@ -16396,12 +16397,15 @@ async def test_get_team_spend_by_user_admin_groups_spend_logs_by_team_and_user(m alpha = _team_spend_by_user_team("team-alpha", "Team Alpha", Member(user_id="alice", role="admin"), []) beta = _team_spend_by_user_team("team-beta", "Team Beta", Member(user_id="alice", role="user"), []) mock_db_client.db.litellm_teamtable.find_many = AsyncMock(return_value=[alpha, beta]) - mock_db_client.db.query_raw = AsyncMock( - return_value=[ - _team_spend_by_user_db_row("team-alpha", "alice", 0.5, 3), - _team_spend_by_user_db_row("team-alpha", "bob", 0.25, 2), - _team_spend_by_user_db_row("team-beta", "alice", 0.1, 1), - ] + transaction = statement_timeout_reads( + mock_db_client, + AsyncMock( + return_value=[ + _team_spend_by_user_db_row("team-alpha", "alice", 0.5, 3), + _team_spend_by_user_db_row("team-alpha", "bob", 0.25, 2), + _team_spend_by_user_db_row("team-beta", "alice", 0.1, 1), + ] + ), ) response = await get_team_spend_by_user( @@ -16411,7 +16415,9 @@ async def test_get_team_spend_by_user_admin_groups_spend_logs_by_team_and_user(m end_date="2026-09-04", ) - sql, *params = mock_db_client.db.query_raw.call_args.args + assert any("statement_timeout = 30000" in statement for statement in transaction.statements) + mock_db_client.db.tx.assert_not_called() + sql, *params = transaction.query_raw.call_args.args assert params == ["2026-09-01", "2026-09-04", "team-alpha", "team-beta"] assert 'FROM "LiteLLM_SpendLogs" sl' in sql assert 'sl."startTime" >= $1::timestamp' in sql @@ -16438,7 +16444,7 @@ async def test_get_team_spend_by_user_team_admin_sees_every_member(mock_db_clien caller = UserAPIKeyAuth(user_id="alice", user_role=LitellmUserRoles.INTERNAL_USER) alpha = _team_spend_by_user_team("team-alpha", "Team Alpha", Member(user_id="alice", role="admin"), []) mock_db_client.db.litellm_teamtable.find_many = AsyncMock(return_value=[alpha]) - mock_db_client.db.query_raw = AsyncMock(return_value=[]) + transaction = statement_timeout_reads(mock_db_client, AsyncMock(return_value=[])) mock_db_client.writer_db.litellm_usertable.find_unique = AsyncMock( return_value=_team_spend_by_user_caller("alice", ["team-alpha"]) ) @@ -16447,7 +16453,7 @@ async def test_get_team_spend_by_user_team_admin_sees_every_member(mock_db_clien user_api_key_dict=caller, team_ids="team-alpha", start_date="2026-09-01", end_date="2026-09-04" ) - sql, *params = mock_db_client.db.query_raw.call_args.args + sql, *params = transaction.query_raw.call_args.args assert params == ["2026-09-01", "2026-09-04", "team-alpha"] assert 'sl."user" = $' not in sql @@ -16460,7 +16466,9 @@ async def test_get_team_spend_by_user_plain_member_only_sees_own_row(mock_db_cli alpha = _team_spend_by_user_team("team-alpha", "Team Alpha", Member(user_id="bob", role="user"), ["/key/info"]) mock_db_client.db.litellm_teamtable.find_many = AsyncMock(return_value=[alpha]) mock_db_client.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[]) - mock_db_client.db.query_raw = AsyncMock(return_value=[_team_spend_by_user_db_row("team-alpha", "bob", 0.25, 2)]) + transaction = statement_timeout_reads( + mock_db_client, AsyncMock(return_value=[_team_spend_by_user_db_row("team-alpha", "bob", 0.25, 2)]) + ) mock_db_client.writer_db.litellm_usertable.find_unique = AsyncMock( return_value=_team_spend_by_user_caller("bob", ["team-alpha"]) ) @@ -16469,7 +16477,7 @@ async def test_get_team_spend_by_user_plain_member_only_sees_own_row(mock_db_cli user_api_key_dict=caller, team_ids="team-alpha", start_date="2026-09-01", end_date="2026-09-04" ) - sql, *params = mock_db_client.db.query_raw.call_args.args + sql, *params = transaction.query_raw.call_args.args assert params == ["2026-09-01", "2026-09-04", "team-alpha", "bob"] assert "sl.team_id IN ($3)" in sql assert 'AND sl."user" = $4' in sql @@ -16481,7 +16489,7 @@ async def test_get_team_spend_by_user_member_of_other_team_gets_404(mock_db_clie from litellm.proxy.management_endpoints.team_endpoints import get_team_spend_by_user caller = UserAPIKeyAuth(user_id="bob", user_role=LitellmUserRoles.INTERNAL_USER) - mock_db_client.db.query_raw = AsyncMock(return_value=[]) + transaction = statement_timeout_reads(mock_db_client, AsyncMock(return_value=[])) mock_db_client.writer_db.litellm_usertable.find_unique = AsyncMock( return_value=_team_spend_by_user_caller("bob", ["team-alpha"]) ) @@ -16492,7 +16500,7 @@ async def test_get_team_spend_by_user_member_of_other_team_gets_404(mock_db_clie ) assert exc_info.value.status_code == 404 - mock_db_client.db.query_raw.assert_not_called() + transaction.query_raw.assert_not_called() @pytest.mark.asyncio @@ -16510,7 +16518,7 @@ async def test_get_team_spend_by_user_member_of_other_team_gets_404(mock_db_clie async def test_get_team_spend_by_user_rejects_bad_input(mock_db_client, team_ids, start_date, end_date, expected_error): from litellm.proxy.management_endpoints.team_endpoints import get_team_spend_by_user - mock_db_client.db.query_raw = AsyncMock(return_value=[]) + transaction = statement_timeout_reads(mock_db_client, AsyncMock(return_value=[])) admin = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN) with pytest.raises(HTTPException) as exc_info: @@ -16520,7 +16528,28 @@ async def test_get_team_spend_by_user_rejects_bad_input(mock_db_client, team_ids assert exc_info.value.status_code == 400 assert expected_error in str(exc_info.value.detail) - mock_db_client.db.query_raw.assert_not_called() + transaction.query_raw.assert_not_called() + + +@pytest.mark.asyncio +async def test_get_team_spend_by_user_statement_timeout_returns_503(mock_db_client): + from litellm.proxy.management_endpoints.team_endpoints import get_team_spend_by_user + + admin = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN) + alpha = _team_spend_by_user_team("team-alpha", "Team Alpha", Member(user_id="alice", role="admin"), []) + mock_db_client.db.litellm_teamtable.find_many = AsyncMock(return_value=[alpha]) + statement_timeout_reads(mock_db_client, AsyncMock(side_effect=statement_timeout_error())) + + with pytest.raises(HTTPException) as exc_info: + await get_team_spend_by_user( + user_api_key_dict=admin, + team_ids="team-alpha", + start_date="2026-09-01", + end_date="2026-09-04", + ) + + assert exc_info.value.status_code == 503 + assert "longer than 30s" in str(exc_info.value.detail) class _TeamRowWithOrganization(LiteLLM_TeamTable): diff --git a/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py b/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py index 4770ecd5485..66bafcbc423 100644 --- a/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py +++ b/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py @@ -1,7 +1,7 @@ import asyncio import json from datetime import datetime -from typing import Any, Dict, Final, List +from typing import Final from unittest.mock import AsyncMock, MagicMock, Mock, patch import httpx @@ -330,9 +330,10 @@ class TestAzureAnthropicCostCalculation: self, mock_completion_cost ): """Test that custom_llm_provider is passed to completion_cost for Azure AI Anthropic""" - from litellm.types.utils import ModelResponse from datetime import datetime + from litellm.types.utils import ModelResponse + mock_completion_cost.return_value = 0.001 logging_obj = self._create_mock_logging_obj( @@ -365,9 +366,10 @@ class TestAzureAnthropicCostCalculation: @patch("litellm.completion_cost") def test_cost_calculation_without_custom_llm_provider(self, mock_completion_cost): """Test that cost calculation works without custom_llm_provider (standard Anthropic)""" - from litellm.types.utils import ModelResponse from datetime import datetime + from litellm.types.utils import ModelResponse + mock_completion_cost.return_value = 0.001 # No custom_llm_provider in model_call_details @@ -401,9 +403,10 @@ class TestAzureAnthropicCostCalculation: self, mock_completion_cost ): """Test that provider prefix is not duplicated if already present in model name""" - from litellm.types.utils import ModelResponse from datetime import datetime + from litellm.types.utils import ModelResponse + mock_completion_cost.return_value = 0.001 logging_obj = self._create_mock_logging_obj( @@ -906,9 +909,10 @@ class TestAnthropicBatchPassthroughCostTracking: mock_request_body, ): """Test that model gets 'anthropic/' prefix when not found in router""" + import base64 + from litellm.llms.anthropic.batches.transformation import AnthropicBatchesConfig from litellm.types.utils import LiteLLMBatch - import base64 # Model not in router - returns same model name mock_get_model_id.return_value = "claude-sonnet-4-5-20250929" @@ -1186,7 +1190,7 @@ class TestAnthropicBatchPassthroughCostTracking: class TestBuildCompleteStreamingResponseRobustness: """_build_complete_streaming_response must tolerate non-standard SSE frames.""" - def _build(self, chunks: List[str]): + def _build(self, chunks: list[str]): return AnthropicPassthroughLoggingHandler._build_complete_streaming_response( all_chunks=chunks, litellm_logging_obj=MagicMock(), @@ -1969,9 +1973,6 @@ class TestStreamFalseDeduplication: This is the _is_assembled_stream_success gate: with stream=False it always returned False and the guard was permanently disabled. """ - from litellm.types.passthrough_endpoints.pass_through_endpoints import ( - EndpointType, - ) from litellm.types.utils import ModelResponse # Simulate what pass_through_endpoints.py now does after stream detection @@ -2856,8 +2857,6 @@ def test_create_anthropic_response_logging_payload(mock_logging_obj, metadata_pa def test_handle_logging_anthropic_collected_chunks(all_chunks): from litellm.proxy.pass_through_endpoints.llm_provider_handlers.anthropic_passthrough_logging_handler import ( AnthropicPassthroughLoggingHandler, - PassthroughStandardLoggingPayload, - EndpointType, ) from litellm.types.utils import ModelResponse @@ -3058,3 +3057,65 @@ def test_streamed_messages_end_turn_logs_no_stop_details(mock_logging_obj): assert logged.choices[0].finish_reason == "stop" provider_fields: Final = logged.choices[0].message.provider_specific_fields or {} assert "stop_details" not in provider_fields + + +_BEDROCK_TRACE_PRICING: Final = { + "bedrock/guardrails": { + "litellm_provider": "bedrock", + "mode": "guardrail", + "guardrail_cost_per_unit": {"topicPolicyUnits": 0.00015, "contentPolicyUnits": 0.00015}, + } +} +_BEDROCK_INVOKE_TRACE: Final = { + "guardrail": { + "input": {"gid": {"invocationMetrics": {"usage": {"topicPolicyUnits": 1, "contentPolicyUnits": 1}}}}, + "outputs": [{"gid": {"invocationMetrics": {"usage": {"topicPolicyUnits": 2, "contentPolicyUnits": 2}}}}], + } +} + + +def _recorded_guardrail_entries(logging_obj: LiteLLMLoggingObj) -> list: + return logging_obj.model_call_details["litellm_params"]["metadata"]["standard_logging_guardrail_information"] + + +def test_anthropic_passthrough_handler_bills_the_bedrock_guardrail_trace_of_a_messages_response( + monkeypatch, mock_httpx_response, mock_response, mock_logging_obj +) -> None: + monkeypatch.setattr(litellm, "model_cost", _BEDROCK_TRACE_PRICING) + mock_logging_obj.model_call_details["litellm_params"] = {"metadata": {}} + AnthropicPassthroughLoggingHandler.anthropic_passthrough_handler( + httpx_response=mock_httpx_response, + response_body={**mock_response, "amazon-bedrock-trace": _BEDROCK_INVOKE_TRACE}, + logging_obj=mock_logging_obj, + url_route="/v1/messages", + result="success", + start_time=datetime(2025, 1, 1), + end_time=datetime(2025, 1, 1), + cache_hit=False, + ) + (entry,) = _recorded_guardrail_entries(mock_logging_obj) + assert entry["guardrail_provider"] == "bedrock" + assert entry["guardrail_usage"] == {"topicPolicyUnits": 3, "contentPolicyUnits": 3} + assert entry["guardrail_cost"] == pytest.approx(6 * 0.00015) + + +def test_anthropic_passthrough_stream_bills_the_bedrock_guardrail_trace_of_the_message_stop_event( + monkeypatch, mock_logging_obj, all_chunks +) -> None: + monkeypatch.setattr(litellm, "model_cost", _BEDROCK_TRACE_PRICING) + mock_logging_obj.model_call_details["litellm_params"] = {"metadata": {}} + stop_event: Final = json.dumps({"type": "message_stop", "amazon-bedrock-trace": _BEDROCK_INVOKE_TRACE}) + chunks: Final = [*all_chunks, f"event: message_stop\ndata: {stop_event}\n\n"] + AnthropicPassthroughLoggingHandler.handle_logging_anthropic_collected_chunks( + litellm_logging_obj=mock_logging_obj, + passthrough_success_handler_obj=MagicMock(), + url_route="/v1/messages", + request_body={"model": "claude", "messages": []}, + endpoint_type=MagicMock(), + start_time=datetime(2025, 1, 1), + all_chunks=chunks, + end_time=datetime(2025, 1, 1), + ) + (entry,) = _recorded_guardrail_entries(mock_logging_obj) + assert entry["guardrail_usage"] == {"topicPolicyUnits": 3, "contentPolicyUnits": 3} + assert entry["guardrail_cost"] == pytest.approx(6 * 0.00015) diff --git a/tests/unit/proxy/pass_through_endpoints/test_watsonx_proxy_route.py b/tests/unit/proxy/pass_through_endpoints/test_watsonx_proxy_route.py index 5500bb0aad9..8efd8a760f5 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_watsonx_proxy_route.py +++ b/tests/unit/proxy/pass_through_endpoints/test_watsonx_proxy_route.py @@ -437,3 +437,120 @@ class TestWatsonxProxyRoute: call_args["target"] == f"https://us-south.ml.cloud.ibm.com/{endpoint_path}" ) + + +_WATSONX_ENV = ( + "WATSONX_API_KEY", + "WATSONX_APIKEY", + "WX_API_KEY", + "WATSONX_API_BASE", + "WATSONX_URL", + "WX_URL", + "WML_URL", + "WATSONX_TOKEN", + "WATSONX_ZENAPIKEY", +) + + +def _watsonx_deployment(api_key: str, api_base: str, use_in_pass_through: bool) -> dict: + return { + "model_name": "granite", + "litellm_params": { + "model": "watsonx/ibm/granite-3-8b-instruct", + "api_key": api_key, + "api_base": api_base, + "project_id": "proj", + "use_in_pass_through": use_in_pass_through, + }, + } + + +def _post_request(body: dict) -> MagicMock: + request = MagicMock(spec=Request) + request.method = "POST" + request.query_params = {} + request.headers = {"content-type": "application/json"} + request.json = AsyncMock(return_value=body) + return request + + +def _iam_response(token: str) -> MagicMock: + response = MagicMock() + response.json.return_value = {"access_token": token, "expires_in": 3600} + return response + + +async def _route_target_and_iam_apikey( + monkeypatch, llm_router: litellm.Router, env: dict[str, str] +) -> tuple[str, str, str]: + from litellm.llms.watsonx.common_utils import iam_token_cache + from litellm.proxy.pass_through_endpoints.passthrough_endpoint_router import PassthroughEndpointRouter + + for name in _WATSONX_ENV: + monkeypatch.delenv(name, raising=False) + for name, value in env.items(): + monkeypatch.setenv(name, value) + monkeypatch.setattr(litellm, "credential_list", []) + iam_token_cache.flush_cache() + monkeypatch.setattr( + "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.passthrough_endpoint_router", + PassthroughEndpointRouter(llm_router_getter=lambda: llm_router), + ) + endpoint_func = AsyncMock(return_value={"results": []}) + with ( + patch.object(litellm.module_level_client, "post", return_value=_iam_response("iam-token-1")) as iam_post, + patch( + "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.create_pass_through_route", + return_value=endpoint_func, + ) as create_route, + ): + await watsonx_proxy_route( + endpoint="ml/v1/text/generation", + request=_post_request({"input": "hi", "project_id": "proj"}), + fastapi_response=MagicMock(spec=Response), + user_api_key_dict=MagicMock(spec=UserAPIKeyAuth), + ) + route_kwargs = create_route.call_args.kwargs + return ( + route_kwargs["target"], + route_kwargs["custom_headers"]["Authorization"], + iam_post.call_args.kwargs["data"]["apikey"], + ) + + +@pytest.mark.asyncio +async def test_watsonx_route_uses_the_use_in_pass_through_deployment_key_and_base(monkeypatch): + llm_router = litellm.Router( + model_list=[ + _watsonx_deployment("wx-deployment-key", "https://eu-de.ml.cloud.ibm.com", use_in_pass_through=True) + ] + ) + + target, authorization, iam_apikey = await _route_target_and_iam_apikey(monkeypatch, llm_router, env={}) + + assert target == "https://eu-de.ml.cloud.ibm.com/ml/v1/text/generation" + assert authorization == "Bearer iam-token-1" + assert iam_apikey == "wx-deployment-key" + + +@pytest.mark.asyncio +async def test_watsonx_route_falls_back_to_env_when_no_deployment_is_flagged(monkeypatch): + llm_router = litellm.Router( + model_list=[ + _watsonx_deployment("wx-deployment-key", "https://eu-de.ml.cloud.ibm.com", use_in_pass_through=False) + ] + ) + + target, authorization, iam_apikey = await _route_target_and_iam_apikey( + monkeypatch, + llm_router, + env={ + "WX_API_KEY": "wx-env-key", + "WATSONX_API_KEY": "watsonx-env-key", + "WATSONX_API_BASE": "https://us-south.ml.cloud.ibm.com", + }, + ) + + assert target == "https://us-south.ml.cloud.ibm.com/ml/v1/text/generation" + assert authorization == "Bearer iam-token-1" + assert iam_apikey == "wx-env-key" diff --git a/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py b/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py index f01d70ab46f..2deca6bc0a9 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py +++ b/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py @@ -18,6 +18,11 @@ import litellm import litellm.proxy.proxy_server as ps from litellm.proxy.auth.authorization import OwnedRows from litellm.proxy.auth.authorization_dependencies import get_log_team_lookup, load_permitted_log_team_ids +from tests.unit.proxy.db.fake_prisma_engine import ( + ReadTargetOf, + statement_timeout_error, + statement_timeout_reads, +) def _default_date_range(): @@ -256,6 +261,8 @@ def make_ui_spend_logs_mock_prisma(mock_spend_logs, filter_fn, team_lookup_fn=No return [row for row in ordered[skip : skip + page_size]] class MockPrismaClient: + read_db = property(lambda self: ReadTargetOf(self.db)) + def __init__(self): self.db = MockDB() self.db.litellm_spendlogs = self.db @@ -313,6 +320,7 @@ async def test_is_admin_view_safe_exception(): async def test_can_team_member_view_log_none_team_id(): # team_id=None should immediately return False class MockPrisma: + read_db = property(lambda self: ReadTargetOf(self.db)) class DB: class TeamTable: async def find_unique(self, where: dict): @@ -334,6 +342,7 @@ async def test_can_team_member_view_log_none_team_id(): async def test_can_team_member_view_log_team_not_found(monkeypatch): # Non-existent team should return False class MockPrisma: + read_db = property(lambda self: ReadTargetOf(self.db)) class DB: class TeamTable: async def find_unique(self, where: dict): @@ -373,6 +382,7 @@ async def test_can_team_member_view_log_not_admin(monkeypatch): } class MockPrisma: + read_db = property(lambda self: ReadTargetOf(self.db)) class DB: class TeamTable: async def find_unique(self, where: dict): @@ -411,6 +421,7 @@ async def test_can_team_member_view_log_admin(monkeypatch): } class MockPrisma: + read_db = property(lambda self: ReadTargetOf(self.db)) class DB: class TeamTable: async def find_unique(self, where: dict): @@ -476,6 +487,7 @@ def _make_owner_lookup_prisma(rows): return _emulate_spend_log_owner_lookup(rows, sql_query, params) class MockPrisma: + read_db = property(lambda self: ReadTargetOf(self.db)) def __init__(self): self.db = MockDB() @@ -981,6 +993,7 @@ async def test_ui_view_spend_logs_sort_by_and_sort_order( return [row for row in sorted_logs[skip : skip + page_size]] class MockPrismaClient: + read_db = property(lambda self: ReadTargetOf(self.db)) def __init__(self): self.db = MagicMock() self.db.litellm_spendlogs = MagicMock() @@ -1042,6 +1055,7 @@ async def test_ui_view_spend_logs_sort_validation_errors(client, monkeypatch, so return 0 class MockPrismaClient: + read_db = property(lambda self: ReadTargetOf(self.db)) def __init__(self): self.db = MagicMock() self.db.litellm_spendlogs = MagicMock() @@ -1118,6 +1132,7 @@ async def test_ui_view_spend_logs_sort_by_request_duration_ms(client, monkeypatc return [row for row in sorted_logs[skip : skip + page_size]] class MockPrismaClient: + read_db = property(lambda self: ReadTargetOf(self.db)) def __init__(self): self.db = MagicMock() self.db.litellm_spendlogs = MagicMock() @@ -1214,6 +1229,7 @@ async def test_ui_view_spend_logs_sort_by_model(client, monkeypatch, sort_order, return [row for row in sorted_logs[skip : skip + page_size]] class MockPrismaClient: + read_db = property(lambda self: ReadTargetOf(self.db)) def __init__(self): self.db = MagicMock() self.db.litellm_spendlogs = MagicMock() @@ -1326,6 +1342,7 @@ async def test_ui_view_spend_logs_sort_by_ttft_ms(client, monkeypatch): return [{k: v for k, v in row.items() if k != "_ttft_ms"} for row in sorted_logs[skip : skip + page_size]] class MockPrismaClient: + read_db = property(lambda self: ReadTargetOf(self.db)) def __init__(self): self.db = MagicMock() self.db.litellm_spendlogs = MagicMock() @@ -2046,6 +2063,7 @@ async def test_ui_view_session_spend_logs_pagination(client, monkeypatch): return [mock_spend_logs[0]] class MockPrismaClient: + read_db = property(lambda self: ReadTargetOf(self.db)) def __init__(self): self.db = MockDB() self.db.litellm_spendlogs = self.db @@ -2098,6 +2116,7 @@ async def test_ui_view_session_spend_logs_rehydrates_metadata_jsonb_text(client, return [dict(raw_row)] class MockPrismaClient: + read_db = property(lambda self: ReadTargetOf(self.db)) def __init__(self): self.db = MockDB() self.db.litellm_spendlogs = self.db @@ -2146,6 +2165,7 @@ async def test_ui_view_session_spend_logs_includes_permitted_team_logs(client, m ] class MockPrismaClient: + read_db = property(lambda self: ReadTargetOf(self.db)) def __init__(self): self.db = MockDB() self.db.litellm_spendlogs = self.db @@ -2623,6 +2643,7 @@ async def test_ui_view_request_response_rejects_foreign_row_inserted_after_owner ] class MockPrisma: + read_db = property(lambda self: ReadTargetOf(self.db)) def __init__(self): self.db = MockDB() @@ -2654,6 +2675,7 @@ async def test_ui_view_request_response_custom_logger_denies_foreign_payload_own return [] class MockPrisma: + read_db = property(lambda self: ReadTargetOf(self.db)) def __init__(self): self.db = MockDB() @@ -3956,6 +3978,7 @@ async def test_view_spend_logs_summarize_parameter(client, monkeypatch): ] class MockPrismaClient: + read_db = property(lambda self: ReadTargetOf(self.db)) def __init__(self): self.db = MockDB() @@ -4051,6 +4074,7 @@ async def test_view_spend_logs_bounds_row_count(client, monkeypatch): return [{}] * min(kwargs.get("take", 0), self.available_rows) class MockPrismaClient: + read_db = property(lambda self: ReadTargetOf(self.db)) def __init__(self): self.db = MockDB() @@ -4174,6 +4198,7 @@ def _tag_query_prisma_client(assert_query_raw): return assert_query_raw(sql_query, params) class MockPrismaClient: + read_db = property(lambda self: ReadTargetOf(self.db)) def __init__(self): self.db = MockDB() @@ -4393,6 +4418,7 @@ async def test_view_spend_logs_with_date_range_summarized(client, monkeypatch): return mock_summarized_response class MockPrismaClient: + read_db = property(lambda self: ReadTargetOf(self.db)) def __init__(self): self.db = MockDB() @@ -4455,6 +4481,7 @@ async def test_view_spend_logs_summarize_groups_by_day_in_sql(client, monkeypatc return mock_rows class MockPrismaClient: + read_db = property(lambda self: ReadTargetOf(self.db)) def __init__(self): self.db = MockDB() @@ -4520,6 +4547,7 @@ async def test_view_spend_logs_summarize_empty_rows(client, monkeypatch): return [] class MockPrismaClient: + read_db = property(lambda self: ReadTargetOf(self.db)) def __init__(self): self.db = MockDB() @@ -4559,6 +4587,7 @@ async def test_view_spend_logs_summarize_unhashed_api_key_without_padding(client return mock_rows class MockPrismaClient: + read_db = property(lambda self: ReadTargetOf(self.db)) def __init__(self): self.db = MockDB() @@ -4873,6 +4902,7 @@ async def test_build_ui_spend_logs_response_dict_rows_session_counts(): } ] ) + statement_timeout_reads(mock_prisma, mock_prisma.db.query_raw) result = await _build_ui_spend_logs_response( prisma_client=mock_prisma, @@ -4938,6 +4968,7 @@ async def test_build_ui_spend_logs_response_caps_session_models(): } ] ) + statement_timeout_reads(mock_prisma, mock_prisma.db.query_raw) result = await _build_ui_spend_logs_response( prisma_client=mock_prisma, @@ -5002,6 +5033,7 @@ async def test_build_ui_spend_logs_response_key_split_session_gets_per_key_aggre }, ] ) + statement_timeout_reads(mock_prisma, mock_prisma.db.query_raw) result = await _build_ui_spend_logs_response( prisma_client=mock_prisma, @@ -5055,6 +5087,7 @@ async def test_build_ui_spend_logs_response_empty_api_key_keeps_session_aggregat } ] ) + statement_timeout_reads(mock_prisma, mock_prisma.db.query_raw) result = await _build_ui_spend_logs_response( prisma_client=mock_prisma, @@ -5110,6 +5143,7 @@ async def test_build_ui_spend_logs_response_sums_multi_round_session_spend(): } ] ) + statement_timeout_reads(mock_prisma, mock_prisma.db.query_raw) result = await _build_ui_spend_logs_response( prisma_client=mock_prisma, @@ -5192,6 +5226,7 @@ async def test_build_ui_spend_logs_response_sums_multi_round_session_tokens(): } ] ) + statement_timeout_reads(mock_prisma, mock_prisma.db.query_raw) result = await _build_ui_spend_logs_response( prisma_client=mock_prisma, @@ -5272,6 +5307,7 @@ async def test_build_ui_spend_logs_response_sums_multi_round_session_duration(): } ] ) + statement_timeout_reads(mock_prisma, mock_prisma.db.query_raw) result = await _build_ui_spend_logs_response( prisma_client=mock_prisma, @@ -5297,6 +5333,35 @@ async def test_build_ui_spend_logs_response_sums_multi_round_session_duration(): ) +@pytest.mark.asyncio +async def test_build_ui_spend_logs_response_keeps_the_page_when_postgres_cancels_the_session_totals_read(): + """The page already fetched is returned without session totals, like any other failed enrichment read.""" + from litellm.proxy.spend_tracking.spend_management_endpoints import ( + _build_ui_spend_logs_response, + ) + + dict_rows = [ + {"request_id": "req-1", "session_id": "sess-slow", "call_type": "completion", "api_key": "k"}, + {"request_id": "req-2", "session_id": "sess-slow", "call_type": "completion", "api_key": "k"}, + ] + mock_prisma = MagicMock() + mock_prisma.db.query_raw = AsyncMock(side_effect=statement_timeout_error()) + statement_timeout_reads(mock_prisma, mock_prisma.db.query_raw) + + result = await _build_ui_spend_logs_response( + prisma_client=mock_prisma, + data=dict_rows, + total_records=2, + page=1, + page_size=50, + total_pages=1, + enrich_session_counts=True, + ) + + assert [row["request_id"] for row in result["data"]] == ["req-1", "req-2"] + assert all(row["session_total_count"] == 1 for row in result["data"]) + + @pytest.mark.asyncio async def test_build_ui_spend_logs_response_session_cache_hit_count(): """ @@ -5330,6 +5395,7 @@ async def test_build_ui_spend_logs_response_session_cache_hit_count(): } ] ) + statement_timeout_reads(mock_prisma, mock_prisma.db.query_raw) result = await _build_ui_spend_logs_response( prisma_client=mock_prisma, @@ -5376,6 +5442,7 @@ async def test_can_team_member_view_log_with_spend_logs_permission(monkeypatch): } class MockPrisma: + read_db = property(lambda self: ReadTargetOf(self.db)) class DB: class TeamTable: async def find_unique(self, where: dict): @@ -5412,6 +5479,7 @@ async def test_can_team_member_view_log_without_spend_logs_permission(monkeypatc } class MockPrisma: + read_db = property(lambda self: ReadTargetOf(self.db)) class DB: class TeamTable: async def find_unique(self, where: dict): @@ -6101,6 +6169,7 @@ async def test_ui_view_spend_logs_rehydrates_metadata_jsonb_text(client, monkeyp return [{**raw_row, "total_count": 1}] class MockPrismaClient: + read_db = property(lambda self: ReadTargetOf(self.db)) def __init__(self): self.db = MagicMock() self.db.litellm_spendlogs = MagicMock() @@ -6185,6 +6254,7 @@ async def test_ui_view_spend_logs_metadata_invalid_json_falls_back_to_empty_dict return [{**raw_row, "total_count": 1}] class MockPrismaClient: + read_db = property(lambda self: ReadTargetOf(self.db)) def __init__(self): self.db = MagicMock() self.db.litellm_spendlogs = MagicMock() @@ -6529,6 +6599,7 @@ _SCOPED_SPEND_REPORT_PATHS = [ def _spend_report_mock_prisma(query_raw_returns=None, team_rows=None, user_row=None): pc = MagicMock() pc.db.query_raw = AsyncMock(return_value=query_raw_returns if query_raw_returns is not None else []) + statement_timeout_reads(pc, pc.db.query_raw) pc.db.litellm_teamtable.find_many = AsyncMock(return_value=team_rows if team_rows is not None else []) pc.db.litellm_usertable.find_unique = AsyncMock(return_value=user_row) return pc @@ -7105,6 +7176,7 @@ def _session_grouped_mock_prisma(session_page_rows, session_total, representativ mock_prisma = MagicMock() mock_prisma.db = MagicMock() mock_prisma.db.query_raw = AsyncMock(side_effect=mock_query_raw) + statement_timeout_reads(mock_prisma, mock_prisma.db.query_raw) return mock_prisma @@ -7150,6 +7222,7 @@ def _session_grouped_paginating_prisma(sessions, counted_total=None): mock_prisma = MagicMock() mock_prisma.db = MagicMock() mock_prisma.db.query_raw = AsyncMock(side_effect=mock_query_raw) + statement_timeout_reads(mock_prisma, mock_prisma.db.query_raw) return mock_prisma @@ -7473,6 +7546,7 @@ async def test_ui_view_spend_logs_group_by_session_offset_for_non_starttime_sort return [] mock_prisma.db.query_raw = AsyncMock(side_effect=mock_query_raw) + statement_timeout_reads(mock_prisma, mock_prisma.db.query_raw) monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) monkeypatch.setattr( "litellm.proxy.spend_tracking.spend_management_endpoints._is_admin_view_safe", @@ -7523,6 +7597,7 @@ async def test_ui_view_spend_logs_search_returns_flat_rows_when_grouping_by_sess mock_prisma = MagicMock() mock_prisma.db = MagicMock() mock_prisma.db.query_raw = AsyncMock(side_effect=mock_query_raw) + statement_timeout_reads(mock_prisma, mock_prisma.db.query_raw) monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin_user" diff --git a/tests/unit/proxy/spend_tracking/test_spend_query_optimization.py b/tests/unit/proxy/spend_tracking/test_spend_query_optimization.py index 93fae093340..68e009d1815 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_query_optimization.py +++ b/tests/unit/proxy/spend_tracking/test_spend_query_optimization.py @@ -9,6 +9,8 @@ import datetime from datetime import timezone from unittest.mock import AsyncMock, MagicMock +from collections.abc import Mapping + import pytest from litellm.proxy.auth.authorization_dependencies import get_log_team_lookup @@ -16,6 +18,7 @@ from litellm.proxy.spend_tracking.spend_tracking_utils import ( get_spend_by_team, get_spend_by_team_and_customer, ) +from tests.unit.proxy.db.fake_prisma_engine import statement_timeout_error, statement_timeout_reads @pytest.mark.asyncio @@ -30,10 +33,8 @@ async def test_spend_query_uses_timestamp_filtering(): """ # Mock prisma client mock_prisma = MagicMock() - mock_db = MagicMock() mock_query_raw = AsyncMock(return_value=[]) - mock_db.query_raw = mock_query_raw - mock_prisma.db = mock_db + statement_timeout_reads(mock_prisma, mock_query_raw) # Use timezone-aware datetime objects start_date = datetime.datetime(2024, 1, 1, tzinfo=timezone.utc) @@ -168,7 +169,7 @@ async def test_spend_logs_ui_wraps_params_in_at_time_zone_utc(monkeypatch): mock_prisma = MagicMock() mock_prisma.db = MagicMock() - mock_prisma.db.query_raw = AsyncMock(return_value=[]) + statement_timeout_reads(mock_prisma, AsyncMock(return_value=[])) mock_prisma.db.litellm_spendlogs = MagicMock() mock_prisma.db.litellm_spendlogs.count = AsyncMock(return_value=0) @@ -195,22 +196,27 @@ async def test_spend_logs_ui_wraps_params_in_at_time_zone_utc(monkeypatch): user_api_key_dict=auth, ) - assert mock_prisma.db.query_raw.called, "query_raw should have been called" - sql = mock_prisma.db.query_raw.call_args[0][0] + assert _spend_log_reads(mock_prisma).called, "query_raw should have been called" + sql = _spend_log_reads(mock_prisma).call_args[0][0] assert sql.count("AT TIME ZONE 'UTC'") >= 2, ( f"/spend/logs/ui must wrap both `startTime` bounds with `AT TIME ZONE 'UTC'`. SQL was:\n{sql}" ) +def _spend_log_reads(mock_prisma: MagicMock) -> AsyncMock: + """The raw reads ``/spend/logs/ui`` ran under the statement timeout.""" + return mock_prisma.read_db.tx.return_value.query_raw + + def _make_ui_spend_logs_mock(count_total, page_rows): """ - Build a prisma mock whose first `query_raw` (the bounded count) returns - `count_total` and whose second `query_raw` (the page data) returns + Build a prisma mock whose first statement-timeout read (the bounded count) returns + `count_total` and whose second read (the page data) returns `page_rows`. """ mock_prisma = MagicMock() mock_prisma.db = MagicMock() - mock_prisma.db.query_raw = AsyncMock(side_effect=[[{"total_count": count_total}], page_rows]) + statement_timeout_reads(mock_prisma, AsyncMock(side_effect=[[{"total_count": count_total}], page_rows])) mock_prisma.db.litellm_spendlogs = MagicMock() mock_prisma.db.litellm_spendlogs.count = AsyncMock(return_value=0) return mock_prisma @@ -260,7 +266,7 @@ async def test_spend_logs_ui_uses_bounded_count_not_full_scan(monkeypatch): mock_prisma.db.litellm_spendlogs.count.assert_not_called() - count_call = mock_prisma.db.query_raw.call_args_list[0] + count_call = _spend_log_reads(mock_prisma).call_args_list[0] count_sql = count_call[0][0] assert "COUNT(*) OVER ()" not in count_sql assert "LIMIT" in count_sql and "FROM (" in count_sql, ( @@ -268,7 +274,7 @@ async def test_spend_logs_ui_uses_bounded_count_not_full_scan(monkeypatch): ) assert count_call[0][-1] == SPEND_LOGS_PAGINATION_COUNT_CAP + 1, "the bounded count must probe at most cap+1 rows" - page_sql = mock_prisma.db.query_raw.call_args_list[1][0][0] + page_sql = _spend_log_reads(mock_prisma).call_args_list[1][0][0] assert "COUNT(*) OVER ()" not in page_sql, ( f"the page query must not carry a window count that forces a full-window scan. SQL was:\n{page_sql}" ) @@ -342,7 +348,7 @@ async def test_spend_logs_ui_empty_page_reports_zero_total(monkeypatch): # page. mock_prisma = MagicMock() mock_prisma.db = MagicMock() - mock_prisma.db.query_raw = AsyncMock(side_effect=[[{"total_count": 0}], []]) + statement_timeout_reads(mock_prisma, AsyncMock(side_effect=[[{"total_count": 0}], []])) mock_prisma.db.litellm_spendlogs = MagicMock() mock_prisma.db.litellm_spendlogs.count = AsyncMock(return_value=0) @@ -391,7 +397,7 @@ async def test_spend_logs_ui_out_of_range_page_keeps_total(monkeypatch): # out-of-range page (empty). mock_prisma = MagicMock() mock_prisma.db = MagicMock() - mock_prisma.db.query_raw = AsyncMock(side_effect=[[{"total_count": 7}], []]) + statement_timeout_reads(mock_prisma, AsyncMock(side_effect=[[{"total_count": 7}], []])) mock_prisma.db.litellm_spendlogs = MagicMock() mock_prisma.db.litellm_spendlogs.count = AsyncMock(return_value=0) @@ -432,9 +438,8 @@ async def test_get_spend_by_team_binds_optional_team_filter(): guard for LIT-4125 (the team query previously carried no team_id predicate). """ mock_prisma = MagicMock() - mock_prisma.db = MagicMock() mock_query_raw = AsyncMock(return_value=[]) - mock_prisma.db.query_raw = mock_query_raw + statement_timeout_reads(mock_prisma, mock_query_raw) start_date = datetime.datetime(2024, 1, 1, tzinfo=timezone.utc) end_date = datetime.datetime(2024, 1, 31, tzinfo=timezone.utc) @@ -484,7 +489,7 @@ async def test_global_spend_report_team_group_forwards_team_id(monkeypatch): mock_prisma = MagicMock() mock_prisma.db = MagicMock() - mock_prisma.db.query_raw = AsyncMock(return_value=[]) + statement_timeout_reads(mock_prisma, AsyncMock(return_value=[])) monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True) @@ -499,9 +504,9 @@ async def test_global_spend_report_team_group_forwards_team_id(monkeypatch): customer_id=None, ) - assert mock_prisma.db.query_raw.called, "query_raw should have been called" - sql = mock_prisma.db.query_raw.call_args[0][0] - params = mock_prisma.db.query_raw.call_args[0][1:] + assert _spend_log_reads(mock_prisma).called, "query_raw should have been called" + sql = _spend_log_reads(mock_prisma).call_args[0][0] + params = _spend_log_reads(mock_prisma).call_args[0][1:] assert "team_x" in params, "team_id must be forwarded into the DB query params" assert "sl.team_id = $3" in sql, f"team query must filter on team_id. SQL was:\n{sql}" @@ -542,7 +547,7 @@ async def test_spend_logs_ui_group_by_session_paginates_sessions(monkeypatch): mock_prisma = MagicMock() mock_prisma.db = MagicMock() - mock_prisma.db.query_raw = AsyncMock(side_effect=mock_query_raw) + statement_timeout_reads(mock_prisma, AsyncMock(side_effect=mock_query_raw)) monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) auth = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin") @@ -567,7 +572,7 @@ async def test_spend_logs_ui_group_by_session_paginates_sessions(monkeypatch): session_cursor=None, ) - emitted = [call[0] for call in mock_prisma.db.query_raw.call_args_list] + emitted = [call[0] for call in _spend_log_reads(mock_prisma).call_args_list] page_sql = emitted[0][0] assert f"GROUP BY {group_key}" in page_sql, f"page must select sessions, not calls. SQL was:\n{page_sql}" assert "OFFSET" not in page_sql, "the startTime page must be keyset-selected, not offset-selected" @@ -635,10 +640,10 @@ async def test_spend_logs_ui_group_by_session_offset_pages_for_other_sorts(monke ) group_key = "COALESCE(NULLIF(session_id, ''), request_id), api_key" - count_sql = mock_prisma.db.query_raw.call_args_list[0][0][0] + count_sql = _spend_log_reads(mock_prisma).call_args_list[0][0][0] assert f"GROUP BY {group_key}" in count_sql - page_call = mock_prisma.db.query_raw.call_args_list[1][0] + page_call = _spend_log_reads(mock_prisma).call_args_list[1][0] page_sql = page_call[0] assert f"DISTINCT ON ({group_key})" in page_sql assert "ORDER BY spend DESC" in page_sql @@ -684,8 +689,108 @@ async def test_spend_logs_ui_request_id_lookup_with_grouping_returns_exact_row(m group_by_session=True, ) - page_call = mock_prisma.db.query_raw.call_args_list[1] + page_call = _spend_log_reads(mock_prisma).call_args_list[1] assert "request_id = $" in page_call[0][0], "the request_id equality filter must survive grouping" assert "req-deep-link" in page_call[0] assert [row["request_id"] for row in response["data"]] == ["req-deep-link"] assert response["total"] == 1 + + +async def _spend_logs_ui_page(mock_prisma: MagicMock, monkeypatch: pytest.MonkeyPatch) -> Mapping[str, object]: + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + from litellm.proxy.spend_tracking.spend_management_endpoints import ui_view_spend_logs + + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) + mock_request = MagicMock() + mock_request.url.path = "/spend/logs/ui" + return await ui_view_spend_logs( + log_team_lookup=await get_log_team_lookup(), + request=mock_request, + api_key=None, + user_id=None, + request_id=None, + search=None, + start_date="2026-02-16 00:00:00", + end_date="2026-02-16 23:59:59", + page=1, + page_size=50, + sort_by="startTime", + sort_order="desc", + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin"), + ) + + +@pytest.mark.asyncio +async def test_spend_logs_ui_reads_are_cancelled_by_postgres_after_30s(monkeypatch: pytest.MonkeyPatch): + """Prisma's HTTP client gives up at 30s either way; without SET LOCAL statement_timeout the scan kept running.""" + mock_prisma = MagicMock() + transaction = statement_timeout_reads(mock_prisma, AsyncMock(side_effect=[[{"total_count": 0}], []])) + + await _spend_logs_ui_page(mock_prisma, monkeypatch) + + assert transaction.statements == ("SET LOCAL statement_timeout = 30000",) * 2 + assert mock_prisma.read_db.tx.call_args.kwargs["timeout"] > datetime.timedelta(seconds=30) + + +@pytest.mark.asyncio +async def test_spend_logs_ui_cancelled_read_is_reported_as_503_with_a_date_range_hint(monkeypatch: pytest.MonkeyPatch): + from litellm.proxy._types import ProxyException + + mock_prisma = MagicMock() + statement_timeout_reads(mock_prisma, AsyncMock(side_effect=statement_timeout_error())) + + with pytest.raises(ProxyException) as exc_info: + await _spend_logs_ui_page(mock_prisma, monkeypatch) + + assert str(exc_info.value.code) == "503" + assert "Spend logs took longer than 30s. Pick a shorter date range" in exc_info.value.message + + +@pytest.mark.asyncio +async def test_global_spend_report_returns_the_503_hint_not_its_catch_all_400_when_postgres_cancels( + monkeypatch: pytest.MonkeyPatch, +): + from fastapi import HTTPException + + from litellm.proxy.spend_tracking.spend_management_endpoints import get_global_spend_report + + mock_prisma = MagicMock() + statement_timeout_reads(mock_prisma, AsyncMock(side_effect=statement_timeout_error())) + monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) + + with pytest.raises(HTTPException) as exc_info: + await get_global_spend_report( + start_date="2026-02-01", + end_date="2026-02-16", + group_by="api_key", + api_key=None, + internal_user_id=None, + team_id=None, + customer_id=None, + ) + + assert exc_info.value.status_code == 503 + assert "Spend logs took longer than 30s. Pick a shorter date range" in str(exc_info.value.detail) + + +@pytest.mark.asyncio +async def test_global_spend_report_reads_are_cancelled_by_postgres_after_30s(monkeypatch: pytest.MonkeyPatch): + from litellm.proxy.spend_tracking.spend_management_endpoints import get_global_spend_report + + mock_prisma = MagicMock() + transaction = statement_timeout_reads(mock_prisma, AsyncMock(return_value=[])) + monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) + + await get_global_spend_report( + start_date="2026-02-01", + end_date="2026-02-16", + group_by="api_key", + api_key=None, + internal_user_id=None, + team_id=None, + customer_id=None, + ) + + assert transaction.statements == ("SET LOCAL statement_timeout = 30000",) diff --git a/tests/unit/proxy/test_proxy_utils.py b/tests/unit/proxy/test_proxy_utils.py index 20a13bdc566..b5c67de6eb7 100644 --- a/tests/unit/proxy/test_proxy_utils.py +++ b/tests/unit/proxy/test_proxy_utils.py @@ -3426,3 +3426,57 @@ async def test_evict_config_param_clears_the_local_layer_and_survives_a_redis_ou await evict_config_param("router_settings", cache=cache) assert await cache.in_memory_cache.async_get_cache(_config_cache_key("router_settings")) is None + + +@pytest.mark.asyncio +async def test_query_raw_within_statement_timeout_has_postgres_cancel_the_read_at_30s() -> None: + from datetime import timedelta + + from litellm.proxy.utils import query_raw_within_statement_timeout + from tests.unit.proxy.db.fake_prisma_engine import statement_timeout_reads + + prisma: Final = MagicMock() + transaction: Final = statement_timeout_reads(prisma, AsyncMock(return_value=[{"spend": 1.5}])) + + rows: Final = await query_raw_within_statement_timeout( + prisma, 'SELECT spend FROM "LiteLLM_SpendLogs" WHERE "user" = $1', "u1" + ) + + assert rows == [{"spend": 1.5}] + assert transaction.statements == ("SET LOCAL statement_timeout = 30000",) + transaction.query_raw.assert_awaited_once_with('SELECT spend FROM "LiteLLM_SpendLogs" WHERE "user" = $1', "u1") + assert prisma.read_db.tx.call_args.kwargs["timeout"] > timedelta(seconds=30), ( + "Prisma must outlive the Postgres cancel" + ) + + +@pytest.mark.asyncio +async def test_query_raw_within_statement_timeout_reports_a_cancelled_read_as_503() -> None: + from litellm.proxy.utils import query_raw_within_statement_timeout + from tests.unit.proxy.db.fake_prisma_engine import statement_timeout_error, statement_timeout_reads + + prisma: Final = MagicMock() + statement_timeout_reads(prisma, AsyncMock(side_effect=statement_timeout_error())) + + with pytest.raises(HTTPException) as exc_info: + await query_raw_within_statement_timeout(prisma, "SELECT 1", what="Spend per user") + + assert exc_info.value.status_code == 503 + assert exc_info.value.detail == {"error": "Spend per user took longer than 30s. Pick a shorter date range"} + + +@pytest.mark.asyncio +async def test_query_raw_within_statement_timeout_leaves_other_database_errors_alone() -> None: + from prisma.errors import DataError + + from litellm.proxy.utils import query_raw_within_statement_timeout + from tests.unit.proxy.db.fake_prisma_engine import statement_timeout_reads + + prisma: Final = MagicMock() + failure: Final = DataError(data={"user_facing_error": {"is_panic": False, "message": "relation does not exist"}}) + statement_timeout_reads(prisma, AsyncMock(side_effect=failure)) + + with pytest.raises(DataError) as exc_info: + await query_raw_within_statement_timeout(prisma, "SELECT 1") + + assert exc_info.value is failure diff --git a/tests/unit/router_utils/test_cooldown_handlers.py b/tests/unit/router_utils/test_cooldown_handlers.py index 2376e0f4d0c..b9a043d7d2b 100644 --- a/tests/unit/router_utils/test_cooldown_handlers.py +++ b/tests/unit/router_utils/test_cooldown_handlers.py @@ -2,7 +2,7 @@ import asyncio import importlib import random import time -from collections.abc import Mapping +from collections.abc import Mapping, Sequence from typing import Final from unittest.mock import MagicMock, patch @@ -20,6 +20,7 @@ from litellm.router_utils.cooldown_callbacks import router_cooldown_event_callba from litellm.router_utils.cooldown_handlers import ( _get_deployment_cooldown_policy, _has_explicit_allowed_fails_policy_for_exception, + _has_fallback_group, _increment_allowed_fails, _is_cooldown_required, _resolve_allowed_fails_from_policy, @@ -41,6 +42,7 @@ from litellm.router_utils.router_callbacks.track_deployment_metrics import ( increment_deployment_successes_for_current_minute, ) from litellm.types.router import ( + RouterRateLimitError, AllowedFailsPolicy, DeploymentTypedDict, LiteLLMParamsTypedDict, @@ -2927,10 +2929,65 @@ async def test_bad_request_error_does_not_cool_down_a_deployment( assert tuple(cooldown_deployments) == () -@pytest.mark.parametrize("use_policy", [False, True], ids=["allowed_fails", "allowed_fails_policy"]) +TIMED_OUT_DEPLOYMENT: Final = { + "model_name": "test-model", + "litellm_params": {"model": "openai/test-model", "api_key": "test-key"}, + "model_info": {"id": "test-deployment"}, +} +BACKUP_DEPLOYMENT: Final = { + "model_name": "backup-model", + "litellm_params": {"model": "openai/backup-model", "api_key": "test-key", "api_base": "https://backup.test/v1"}, + "model_info": {"id": "backup-deployment"}, +} +BACKUP_REPLY: Final = { + "id": "chatcmpl-backup", + "object": "chat.completion", + "created": 1700000000, + "model": "backup-model", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, +} + + +async def _lone_call_outcome(router: Router) -> type[Exception]: + try: + await router.acompletion(model="test-model", messages=[{"role": "user", "content": "timeout test"}]) + except (litellm.Timeout, RouterRateLimitError) as error: + return type(error) + pytest.fail("the timed-out deployment answered") + + +BENCHED: Final = (litellm.Timeout, litellm.Timeout, RouterRateLimitError, RouterRateLimitError, RouterRateLimitError) + + +@pytest.mark.parametrize( + ("router_kwargs", "model_info", "outcomes"), + [ + ({"allowed_fails": 1}, {}, (litellm.Timeout,) * 5), + ({"allowed_fails": 1}, {"allowed_fails_policy": {}}, (litellm.Timeout,) * 5), + ({"allowed_fails_policy": AllowedFailsPolicy(RateLimitErrorAllowedFails=1)}, {}, (litellm.Timeout,) * 5), + ({"allowed_fails_policy": AllowedFailsPolicy(TimeoutErrorAllowedFails=1)}, {}, BENCHED), + ({"allowed_fails": 1, "fallbacks": [{"test-model": ["backup-model"]}]}, {}, BENCHED), + ( + {"allowed_fails": 1, "fallbacks": [{"test-model": [{"model": "backup-model", "temperature": 0}]}]}, + {}, + BENCHED, + ), + ], + ids=[ + "allowed_fails", + "allowed_fails_with_an_empty_deployment_policy", + "policy_names_another_exception", + "policy_names_timeout", + "allowed_fails_with_a_fallback_group", + "allowed_fails_with_a_dict_fallback_target", + ], +) @pytest.mark.asyncio -async def test_allowed_fails_cools_a_single_deployment( - use_policy: bool, +async def test_router_allowed_fails_benches_a_lone_deployment_only_when_the_policy_names_the_exception( + router_kwargs: Mapping[str, object], + model_info: Mapping[str, object], + outcomes: tuple[type[Exception], ...], respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, ) -> None: @@ -2938,38 +2995,80 @@ async def test_allowed_fails_cools_a_single_deployment( route: Final = respx_mock.post("https://api.openai.com/v1/chat/completions").mock( side_effect=httpx.ReadTimeout("upstream timeout") ) - policy_kwargs: Final = ( - {"allowed_fails_policy": AllowedFailsPolicy(TimeoutErrorAllowedFails=1)} - if use_policy - else {"allowed_fails": 1} + deployment: Final = {**TIMED_OUT_DEPLOYMENT, "model_info": {**TIMED_OUT_DEPLOYMENT["model_info"], **model_info}} + router: Final = Router(model_list=[deployment], cooldown_time=300, num_retries=0, **router_kwargs) + + seen: Final = tuple([await _lone_call_outcome(router) for _ in range(5)]) + + assert seen == outcomes + assert len(route.calls) == outcomes.count(litellm.Timeout) + + +@pytest.mark.parametrize( + ("fallbacks", "expected"), + [ + ([{"test-model": "backup-model"}], True), + ([{"*": "backup-model"}], True), + ([{"test-model": ["backup-model"]}], True), + ([{"test-model": 5}, {"*": "backup-model"}], True), + ([{"test-model": 5}], False), + ([{"other-model": ["backup-model"]}], False), + ], + ids=[ + "string_target", + "string_catch_all_from_default_fallbacks", + "list_target", + "unparseable_entry_does_not_hide_the_catch_all", + "unparseable_entry_alone", + "another_groups_chain", + ], +) +def test_has_fallback_group_reads_string_targets_and_skips_unparseable_entries( + fallbacks: Sequence[Mapping[str, object]], expected: bool +) -> None: + router: Final = Router(model_list=[TIMED_OUT_DEPLOYMENT]) + router.fallbacks = list(fallbacks) + + assert _has_fallback_group(router, "test-model") is expected + + +async def _group_call_outcome(router: Router) -> str: + try: + response: Final = await router.acompletion( + model="test-model", messages=[{"role": "user", "content": "timeout test"}] + ) + except litellm.Timeout: + return "timeout" + return response.id + + +@pytest.mark.asyncio +async def test_router_allowed_fails_still_benches_a_deployment_that_has_a_sibling_in_its_group( + respx_mock: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + timed_out: Final = respx_mock.post("https://api.openai.com/v1/chat/completions").mock( + side_effect=httpx.ReadTimeout("upstream timeout") ) + respx_mock.post("https://backup.test/v1/chat/completions").mock(return_value=httpx.Response(200, json=BACKUP_REPLY)) router: Final = Router( - model_list=[ - { - "model_name": "test-model", - "litellm_params": {"model": "openai/test-model", "api_key": "test-key"}, - "model_info": {"id": "test-deployment"}, - } - ], + model_list=[TIMED_OUT_DEPLOYMENT, {**BACKUP_DEPLOYMENT, "model_name": "test-model"}], cooldown_time=300, num_retries=0, - **policy_kwargs, + allowed_fails=1, ) - for _ in range(2): - with pytest.raises(litellm.Timeout): - await router.acompletion( - model="test-model", - messages=[{"role": "user", "content": "timeout test"}], - ) - + while len(timed_out.calls) < 2: + await _group_call_outcome(router) cooldown_deployments: Final = await async_get_cooldown_deployments( - litellm_router_instance=router, - parent_otel_span=None, + litellm_router_instance=router, parent_otel_span=None ) + served: Final = tuple([await _group_call_outcome(router) for _ in range(5)]) - assert len(route.calls) == 2 assert tuple(cooldown_deployments) == ("test-deployment",) + assert served == ("chatcmpl-backup",) * 5 + assert len(timed_out.calls) == 2 def test_per_user_session_upstream_errors_never_cool_down_the_shared_deployment(single_deployment_router): diff --git a/tests/unit/router_utils/test_fallback_event_handlers.py b/tests/unit/router_utils/test_fallback_event_handlers.py index 56184fa299a..3d4c1f2d4b3 100644 --- a/tests/unit/router_utils/test_fallback_event_handlers.py +++ b/tests/unit/router_utils/test_fallback_event_handlers.py @@ -1034,7 +1034,12 @@ class TestTriggerCooldownForFailedDeployment: "model_name": "fallback-model", "litellm_params": {"model": "openai/gpt-5.6", "api_key": "sk-fake"}, "model_info": {"id": "fallback-deployment"}, - } + }, + { + "model_name": "fallback-model", + "litellm_params": {"model": "openai/gpt-5.6", "api_key": "sk-fake"}, + "model_info": {"id": "fallback-sibling"}, + }, ], allowed_fails=0, cooldown_time=60, diff --git a/tests/unit/test_cost_calculator.py b/tests/unit/test_cost_calculator.py index 6cbe4084186..71baf6067a1 100644 --- a/tests/unit/test_cost_calculator.py +++ b/tests/unit/test_cost_calculator.py @@ -16,6 +16,7 @@ import litellm from litellm import TranscriptionResponse, model_cost from litellm.constants import REALTIME_METERED_USAGE_KEY from litellm.cost_calculator import BaseTokenUsageProcessor, RealtimeAPITokenUsageProcessor, ResponsesWebSocketTokenUsageProcessor, batch_cost_calculator, completion_cost, cost_per_token, handle_realtime_stream_cost_calculation, response_cost_calculator +from litellm.litellm_core_utils.llm_cost_calc.utils import parse_prompt_tokens_details from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.fallback_generalizations import ( get_fallback_generalization_rules, @@ -3215,14 +3216,7 @@ def test_anthropic_us_data_residency_uplift_on_claude_4_6_and_later_models(_loca def test_gemini_cache_tokens_details_no_negative_values(): - """ - Test for Issue #18750: Negative text_tokens with Gemini caching - - When using Gemini with explicit caching, the response includes cacheTokensDetails - which breaks down cached tokens by modality. This test ensures that: - 1. text_tokens is never negative - 2. We correctly subtract cached tokens per modality (not total) - """ + """Ensure issue #18750 cache details remain non-negative after cost parsing.""" from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( VertexGeminiConfig, ) @@ -3250,27 +3244,15 @@ def test_gemini_cache_tokens_details_no_negative_values(): usage = VertexGeminiConfig.calculate_usage(completion_response) - # Text tokens should be non-cached text only: 9402 - 9393 = 9 - assert usage.prompt_tokens_details.text_tokens == 9, ( - f"Expected text_tokens=9, got {usage.prompt_tokens_details.text_tokens}" - ) + prompt_details: Final = parse_prompt_tokens_details(usage) - # Image tokens should be non-cached image only: 258 - 258 = 0 - assert usage.prompt_tokens_details.image_tokens == 0, ( - f"Expected image_tokens=0, got {usage.prompt_tokens_details.image_tokens}" - ) - - # Total cached should match - assert usage.prompt_tokens_details.cached_tokens == 9651, ( - f"Expected cached_tokens=9651, got {usage.prompt_tokens_details.cached_tokens}" - ) - - # MOST IMPORTANT: text_tokens should NEVER be negative - assert usage.prompt_tokens_details.text_tokens >= 0, ( - f"BUG: text_tokens is negative ({usage.prompt_tokens_details.text_tokens})! This was the issue in #18750" - ) - - print("✅ Issue #18750 fix verified: text_tokens is correctly calculated and non-negative") + assert usage.prompt_tokens_details.text_tokens == 9402 + assert usage.prompt_tokens_details.image_tokens == 258 + assert usage.prompt_tokens_details.cached_tokens == 9651 + assert usage.prompt_tokens_details.cached_tokens_details == CachedTokensDetails(text_tokens=9393, image_tokens=258) + assert prompt_details["text_tokens"] == 9 + assert prompt_details["image_tokens"] == 0 + assert prompt_details["text_tokens"] >= 0 def test_gemini_without_cache_tokens_details(): @@ -5868,6 +5850,45 @@ def test_batch_cost_calculator_prices_modalities_and_cached_tokens_together_in_t assert prompt_cost == pytest.approx(298_926 * 3e-6 + 64 * 5e-6 + 10 * 3e-6 + 1_000 * 3e-7) +def test_batch_cost_calculator_prices_cached_audio_at_its_batch_rate() -> None: + model_info: Final[ModelInfo] = { + "input_cost_per_token": 2e-6, + "input_cost_per_token_batches": 1e-6, + "input_cost_per_audio_token": 6e-6, + "input_cost_per_audio_token_batches": 3e-6, + "cache_read_input_token_cost": 8e-7, + "cache_read_input_token_cost_batches": 4e-7, + "cache_read_input_audio_token_cost": 5e-7, + "cache_read_input_audio_token_cost_batches": 2e-7, + } + usage: Final = Usage( + prompt_tokens=40, + completion_tokens=0, + total_tokens=40, + prompt_tokens_details=PromptTokensDetailsWrapper( + text_tokens=20, + audio_tokens=20, + cached_tokens=10, + cached_tokens_details=CachedTokensDetails(text_tokens=5, audio_tokens=5), + ), + ) + + prompt_cost, _ = batch_cost_calculator( + usage=usage, + model="gemini-audio-batch-test-model", + custom_llm_provider="gemini", + model_info=model_info, + ) + + expected: Final = ( + 15 * model_info["input_cost_per_token_batches"] + + 15 * model_info["input_cost_per_audio_token_batches"] + + 5 * model_info["cache_read_input_token_cost_batches"] + + 5 * model_info["cache_read_input_audio_token_cost_batches"] + ) + assert prompt_cost == pytest.approx(expected) + + QWEN3_NEXT_REGIONS: Final = ("ap-northeast-1", "ap-south-1", "ap-southeast-2", "eu-west-1", "eu-west-2", "sa-east-1") diff --git a/tests/unit/test_main.py b/tests/unit/test_main.py index 21f8c994b67..6bb69def307 100644 --- a/tests/unit/test_main.py +++ b/tests/unit/test_main.py @@ -20,16 +20,23 @@ from unittest.mock import AsyncMock, MagicMock, patch import httpx import pytest import respx -from openai import APITimeoutError, AsyncOpenAI +from openai import AsyncOpenAI +from openai._legacy_response import LegacyAPIResponse from openai.types.chat.chat_completion import ChatCompletion import litellm -from litellm import acompletion, completion -from litellm import acompletion_with_retries, aresponses_with_retries, completion_with_retries, responses_with_retries +from litellm import ( + acompletion, + acompletion_with_retries, + aresponses_with_retries, + completion, + completion_with_retries, + responses_with_retries, +) from litellm import main as litellm_main -from litellm.constants import CONTROL_OPTIONS_KEY from litellm.caching.base_cache import BaseCache from litellm.caching.caching import Cache +from litellm.constants import CONTROL_OPTIONS_KEY from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.custom_prompt_management import CustomPromptManagement from litellm.litellm_core_utils.core_helpers import get_litellm_metadata_from_kwargs @@ -130,7 +137,7 @@ def test_completion_missing_role(openai_api_response): client = OpenAI(api_key="test_api_key") - mock_raw_response = MagicMock() + mock_raw_response = MagicMock(spec=LegacyAPIResponse) mock_raw_response.headers = { "x-request-id": "123", "openai-organization": "org-123", @@ -142,7 +149,7 @@ def test_completion_missing_role(openai_api_response): print(f"openai_api_response: {openai_api_response}") with patch.object( - client.chat.completions.with_raw_response, "create", MagicMock(return_value=mock_raw_response) + client, "post", MagicMock(return_value=mock_raw_response) ) as mock_create: litellm.completion( model="gpt-4o-mini", @@ -218,7 +225,7 @@ async def test_url_with_format_param_openai(model, sync_mode): } ], } - with patch.object(client.chat.completions.with_raw_response, "create") as mock_client: + with patch.object(client, "post") as mock_client: try: if sync_mode: response = completion(**args, client=client) @@ -232,7 +239,7 @@ async def test_url_with_format_param_openai(model, sync_mode): print(mock_client.call_args.kwargs) - json_str = json.dumps(mock_client.call_args.kwargs) + json_str = json.dumps(mock_client.call_args.kwargs["body"]) assert "format" not in json_str @@ -3052,9 +3059,10 @@ def test_stream_chunk_builder_thinking_blocks(): assert response.choices[0].message.thinking_blocks is not None -from litellm.llms.openai.openai import OpenAIChatCompletion import traceback +from litellm.llms.openai.openai import OpenAIChatCompletion + user_message = "Write a short poem about the sky" messages = [{"content": user_message, "role": "user"}] @@ -3824,7 +3832,7 @@ def test_completion_forwards_store_and_prompt_cache_key_to_openai(): client = OpenAI(api_key="fake-api-key") - with patch.object(client.chat.completions.with_raw_response, "create") as mock_client: + with patch.object(client, "post") as mock_client: try: litellm.completion( model="openai/gpt-4o", @@ -3837,7 +3845,7 @@ def test_completion_forwards_store_and_prompt_cache_key_to_openai(): print(e) mock_client.assert_called_once() - request_body = mock_client.call_args.kwargs + request_body = mock_client.call_args.kwargs["body"] assert request_body["store"] is False assert request_body["prompt_cache_key"] == "test-cache-key" @@ -3852,7 +3860,7 @@ async def test_acompletion_forwards_store_and_prompt_cache_key_to_openai(): client = AsyncOpenAI(api_key="fake-api-key") - with patch.object(client.chat.completions.with_raw_response, "create") as mock_client: + with patch.object(client, "post") as mock_client: try: await litellm.acompletion( model="openai/gpt-4o", @@ -3865,7 +3873,7 @@ async def test_acompletion_forwards_store_and_prompt_cache_key_to_openai(): print(e) mock_client.assert_called_once() - request_body = mock_client.call_args.kwargs + request_body = mock_client.call_args.kwargs["body"] assert request_body["store"] is False assert request_body["prompt_cache_key"] == "test-cache-key" @@ -3879,7 +3887,7 @@ def test_completion_omits_store_and_prompt_cache_key_when_not_passed(): client = OpenAI(api_key="fake-api-key") - with patch.object(client.chat.completions.with_raw_response, "create") as mock_client: + with patch.object(client, "post") as mock_client: try: litellm.completion( model="openai/gpt-4o", @@ -3890,7 +3898,7 @@ def test_completion_omits_store_and_prompt_cache_key_when_not_passed(): print(e) mock_client.assert_called_once() - request_body = mock_client.call_args.kwargs + request_body = mock_client.call_args.kwargs["body"] assert "store" not in request_body assert "prompt_cache_key" not in request_body @@ -4085,8 +4093,7 @@ def test_openai_model_without_a_provider_still_routes_to_openai(): from openai import OpenAI client = OpenAI(api_key="fake-key") - raw_response = client.chat.completions.with_raw_response - with patch.object(raw_response, "create") as mock_create, contextlib.suppress(Exception): + with patch.object(client, "post") as mock_create, contextlib.suppress(Exception): litellm.completion( model="gpt-4o", messages=[{"role": "user", "content": "hello"}], @@ -4097,7 +4104,7 @@ def test_openai_model_without_a_provider_still_routes_to_openai(): def _openai_chat_create_kwargs(client, **completion_kwargs): - with patch.object(client.chat.completions.with_raw_response, "create") as mock_client: + with patch.object(client, "post") as mock_client: with contextlib.suppress(Exception): litellm.completion( messages=[{"role": "system", "content": "sys"}, {"role": "user", "content": "hi"}], @@ -4122,7 +4129,7 @@ def test_completion_custom_api_base_sends_no_prompt_cache_breakpoint_for_gpt_5_6 from openai import OpenAI client = OpenAI(api_key="fake-api-key", base_url="http://127.0.0.1:9/v1") - request_body = _openai_chat_create_kwargs(client, model="gpt-5.6", api_base="http://127.0.0.1:9/v1") + request_body = _openai_chat_create_kwargs(client, model="gpt-5.6", api_base="http://127.0.0.1:9/v1")["body"] assert request_body["messages"][0] == {"role": "system", "content": "sys", "cache_control": {"type": "ephemeral"}} assert "prompt_cache_breakpoint" not in json.dumps(request_body["messages"]) @@ -4134,7 +4141,7 @@ def test_completion_custom_base_url_sends_no_prompt_cache_breakpoint_for_gpt_5_6 from openai import OpenAI client = OpenAI(api_key="fake-api-key", base_url="http://127.0.0.1:9/v1") - request_body = _openai_chat_create_kwargs(client, model="gpt-5.6", base_url="http://127.0.0.1:9/v1") + request_body = _openai_chat_create_kwargs(client, model="gpt-5.6", base_url="http://127.0.0.1:9/v1")["body"] assert request_body["messages"][0] == {"role": "system", "content": "sys", "cache_control": {"type": "ephemeral"}} assert "prompt_cache_breakpoint" not in json.dumps(request_body["messages"]) @@ -4147,7 +4154,7 @@ async def test_acompletion_custom_base_url_sends_no_prompt_cache_breakpoint_for_ from openai import AsyncOpenAI client = AsyncOpenAI(api_key="fake-api-key", base_url="http://127.0.0.1:9/v1") - with patch.object(client.chat.completions.with_raw_response, "create") as mock_create: + with patch.object(client, "post") as mock_create: with contextlib.suppress(Exception): await litellm.acompletion( model="gpt-5.6", @@ -4158,7 +4165,7 @@ async def test_acompletion_custom_base_url_sends_no_prompt_cache_breakpoint_for_ ) mock_create.assert_called_once() - request_body = mock_create.call_args.kwargs + request_body = mock_create.call_args.kwargs["body"] assert request_body["messages"][0] == {"role": "system", "content": "sys", "cache_control": {"type": "ephemeral"}} assert "prompt_cache_breakpoint" not in json.dumps(request_body["messages"]) @@ -4170,12 +4177,12 @@ def test_completion_default_api_base_sends_prompt_cache_breakpoint_for_gpt_5_6() from openai import OpenAI client = OpenAI(api_key="fake-api-key") - request_body = _openai_chat_create_kwargs(client, model="gpt-5.6") + request = _openai_chat_create_kwargs(client, model="gpt-5.6") - assert request_body["messages"][0]["content"] == [ + assert request["body"]["messages"][0]["content"] == [ {"type": "text", "text": "sys", "prompt_cache_breakpoint": {"mode": "explicit"}} ] - assert request_body["extra_body"]["prompt_cache_options"] == {"mode": "implicit"} + assert request["options"]["extra_json"]["prompt_cache_options"] == {"mode": "implicit"} _SUBSCRIPTION_OAUTH_CREDENTIAL = "Bearer sk-ant-oat01-fake-subscription-token-for-testing-0123456789" @@ -5223,7 +5230,7 @@ def test_drop_params_false_still_rejects_an_invalid_stream_chunk_size() -> None: def test_acompletion_params(): import inspect - from litellm.types.completion import CompletionRequest + acompletion_params_odict = inspect.signature(acompletion).parameters completion_params_dict = inspect.signature(completion).parameters @@ -5253,7 +5260,7 @@ def test_acompletion_params(): def _openai_mock_response(*args: object, **kwargs: object) -> MagicMock: - new_response: Final = MagicMock() + new_response: Final = MagicMock(spec=LegacyAPIResponse) new_response.headers = {"hello": "world"} response_object: Final = { "id": "chatcmpl-123", @@ -5288,7 +5295,7 @@ def test_null_role_response(): openai_client = openai.OpenAI() with patch.object( - openai_client.chat.completions, "create", side_effect=_openai_mock_response + openai_client, "post", side_effect=_openai_mock_response ) as mock_response: response = litellm.completion( model="gpt-3.5-turbo", @@ -5378,7 +5385,7 @@ def test_completion_perplexity_api(): pydantic_obj = ChatCompletion(**response_object) def _return_pydantic_obj(*args, **kwargs): - new_response = MagicMock() + new_response = MagicMock(spec=LegacyAPIResponse) new_response.headers = {"hello": "world"} new_response.parse.return_value = pydantic_obj @@ -5387,8 +5394,7 @@ def test_completion_perplexity_api(): openai_client = OpenAI() with patch.object( - openai_client.chat.completions.with_raw_response, - "create", + openai_client, "post", side_effect=_return_pydantic_obj, ) as mock_client: # litellm.set_verbose= True @@ -5432,7 +5438,7 @@ async def test_openai_compatible_custom_api_base(provider): openai_client = OpenAI(api_key="fake-key") with patch.object( - openai_client.chat.completions, "create", new=MagicMock() + openai_client, "post", new=MagicMock() ) as mock_call: try: completion( @@ -5450,7 +5456,7 @@ async def test_openai_compatible_custom_api_base(provider): print("Call KWARGS - {}".format(mock_call.call_args.kwargs)) - assert "hello" in mock_call.call_args.kwargs["extra_body"] + assert "hello" in mock_call.call_args.kwargs["options"]["extra_json"] @pytest.mark.parametrize( @@ -5483,7 +5489,7 @@ async def test_openai_compatible_custom_api_video(provider): openai_client = OpenAI(api_key="fake-key") with patch.object( - openai_client.chat.completions, "create", new=MagicMock() + openai_client, "post", new=MagicMock() ) as mock_call: try: completion( @@ -7346,7 +7352,7 @@ def _openai_hallucinated_tool_call_mock_response( *args: object, **kwargs: object, ) -> MagicMock: - new_response: Final = MagicMock() + new_response: Final = MagicMock(spec=LegacyAPIResponse) new_response.headers = {"hello": "world"} response_object: Final = { "id": "chatcmpl-123", @@ -7404,8 +7410,8 @@ def test_openai_hallucinated_tool_call(): openai_client = openai.OpenAI() with patch.object( - openai_client.chat.completions, - "create", + openai_client, + "post", side_effect=_openai_hallucinated_tool_call_mock_response, ) as mock_response: response = litellm.completion( @@ -7527,7 +7533,7 @@ def test_completion_novita_ai(): openai_client = OpenAI(api_key="fake-key") with patch.object( - openai_client.chat.completions.with_raw_response, "create" + openai_client, "post", return_value=MagicMock(spec=LegacyAPIResponse) ) as mock_call: mock_call.return_value.headers = {} mock_call.return_value.parse.return_value = litellm.ModelResponse( @@ -7546,11 +7552,11 @@ def test_completion_novita_ai(): # Verify model is passed correctly assert ( - mock_call.call_args.kwargs["model"] + mock_call.call_args.kwargs["body"]["model"] == "meta-llama/llama-3.3-70b-instruct" ) # Verify messages are passed correctly - assert mock_call.call_args.kwargs["messages"] == messages + assert mock_call.call_args.kwargs["body"]["messages"] == messages except Exception as e: pytest.fail(f"Error occurred: {e}") @@ -7573,8 +7579,8 @@ def test_completion_novita_ai_dynamic_params(api_key): openai_client = OpenAI(api_key="fake-key") with patch.object( - openai_client.chat.completions, - "create", + openai_client, + "post", side_effect=Exception("Invalid API key"), ) as mock_call: with pytest.raises(Exception, match="Invalid API key") as exc_info: @@ -7606,7 +7612,7 @@ def test_completion_openai_metadata(monkeypatch, enable_preview_features): monkeypatch.setattr(litellm, "enable_preview_features", enable_preview_features) with patch.object( - client.chat.completions.with_raw_response, "create", return_value=MagicMock() + client, "post", return_value=MagicMock() ) as mock_completion: try: resp = litellm.completion( @@ -7620,11 +7626,11 @@ def test_completion_openai_metadata(monkeypatch, enable_preview_features): mock_completion.assert_called_once() if enable_preview_features: - assert mock_completion.call_args.kwargs["metadata"] == { + assert mock_completion.call_args.kwargs["body"]["metadata"] == { "my-test-key": "my-test-value" } else: - assert "metadata" not in mock_completion.call_args.kwargs + assert "metadata" not in mock_completion.call_args.kwargs["body"] AZURE_TTS_BASE: Final = "https://tts.example.azure.com" diff --git a/tests/unit/test_read_rc_version.py b/tests/unit/test_read_rc_version.py index 7b8848916ff..b31e870ad40 100644 --- a/tests/unit/test_read_rc_version.py +++ b/tests/unit/test_read_rc_version.py @@ -45,5 +45,7 @@ def test_the_repo_pyproject_is_read_when_no_path_is_given( monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str] ) -> None: monkeypatch.chdir(_REPO_ROOT) + if read_rc_version.RELEASE_VERSION.fullmatch(read_rc_version.read_version(_REPO_ROOT / "pyproject.toml")) is None: + pytest.skip("release line pyproject carries a patch version, the rc workflow only runs on X.Y.0") assert read_rc_version.main(["read_rc_version.py"]) == 0 assert capsys.readouterr().out == f"version={read_rc_version.read_version(_REPO_ROOT / 'pyproject.toml')}\n" diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index 3a8fc49da7a..8cde45e7603 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -10548,7 +10548,12 @@ class TestBackgroundResponseCostPollCooldown: "model_name": "gpt-4.1", "litellm_params": {"model": "openai/gpt-4.1"}, "model_info": {"id": "dep-1"}, - } + }, + { + "model_name": "gpt-4.1", + "litellm_params": {"model": "openai/gpt-4.1"}, + "model_info": {"id": "dep-2"}, + }, ], allowed_fails=0, ) @@ -10676,7 +10681,12 @@ class TestCallerTimeoutCooldown: "model_name": "slow-model", "litellm_params": {"model": "openai/gpt-5.6", "api_key": "sk-fake"}, "model_info": {"id": "dep-1"}, - } + }, + { + "model_name": "slow-model", + "litellm_params": {"model": "openai/gpt-5.6", "api_key": "sk-fake"}, + "model_info": {"id": "dep-2"}, + }, ], allowed_fails=0, cooldown_time=120, diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index 381883eb284..0a408e7c8ab 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -969,6 +969,16 @@ def test_aaamodel_prices_and_context_window_json_is_valid(): "cache_creation_input_token_cost_above_1hr_above_100k_tokens": {"type": "number"}, "cache_creation_input_token_cost_above_1hr_above_200k_tokens": {"type": "number"}, "cache_read_input_audio_token_cost": {"type": "number"}, + "cache_read_input_audio_token_cost_batches": {"type": "number"}, + "cache_read_input_audio_token_cost_flex": {"type": "number"}, + "cache_read_input_audio_token_cost_priority": {"type": "number"}, + "input_cost_per_audio_token_above_200k_tokens": {"type": "number"}, + "input_cost_per_audio_token_above_200k_tokens_batches": {"type": "number"}, + "input_cost_per_audio_token_above_200k_tokens_flex": {"type": "number"}, + "input_cost_per_audio_token_above_200k_tokens_priority": {"type": "number"}, + "input_cost_per_token_above_200k_tokens_flex": {"type": "number"}, + "output_cost_per_image_token_priority": {"type": "number"}, + "output_cost_per_token_above_200k_tokens_flex": {"type": "number"}, "cache_read_input_image_token_cost": {"type": "number"}, "audio_transcription_config": {"type": "string"}, "deprecation_date": {"type": "string"}, @@ -976,6 +986,7 @@ def test_aaamodel_prices_and_context_window_json_is_valid(): "input_cost_per_audio_per_second_above_128k_tokens": {"type": "number"}, "google_maps_grounding_cost_per_query": {"type": "number"}, "input_cost_per_audio_token": {"type": "number"}, + "input_cost_per_audio_token_flex": {"type": "number"}, "input_cost_per_image_token": {"type": "number"}, "input_cost_per_character": {"type": "number"}, "input_cost_per_character_above_128k_tokens": {"type": "number"}, @@ -1576,6 +1587,38 @@ def test_get_provider_rerank_config(): assert isinstance(config, HostedVLLMRerankConfig) +def test_get_provider_video_config_resolves_a_custom_provider_map_base_video_config( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A `BaseVideoConfig` registered in `custom_provider_map` answers for its provider name, a `CustomLLM` + registered there does not, and a built-in name still resolves to its built-in config even when a `CustomLLM` + shadows it in the map.""" + from litellm.llms.custom_llm import CustomLLM + from litellm.llms.openai.videos.transformation import OpenAIVideoConfig + + class QueueVideoConfig(OpenAIVideoConfig): + pass + + video_config: Final = QueueVideoConfig() + monkeypatch.setattr( + litellm, + "custom_provider_map", + [ + {"provider": "queue-video", "custom_handler": video_config}, + {"provider": "image-only", "custom_handler": CustomLLM()}, + {"provider": "openai", "custom_handler": CustomLLM()}, + ], + ) + monkeypatch.setattr(litellm, "_custom_providers", ["queue-video", "image-only", "openai"]) + + assert ProviderConfigManager.get_provider_video_config(model="seedance-1", provider="queue-video") is video_config + assert ProviderConfigManager.get_provider_video_config(model="seedance-1", provider="image-only") is None + assert isinstance( + ProviderConfigManager.get_provider_video_config(model="sora-2", provider="openai"), OpenAIVideoConfig + ) + assert ProviderConfigManager.get_provider_video_config(model="seedance-1", provider="not-registered") is None + + def test_get_provider_text_to_speech_config_vertex_gemini_skips_cloud_tts(): """Regression for LIT-6501: mapping vertex Gemini TTS params through Google Cloud TTS dropped response_format before the speech_to_completion bridge could honor it.""" @@ -7664,7 +7707,6 @@ def test_openai_non_text_embedding_3_drop_params_per_call(): litellm.drop_params = prev_drop_params - @pytest.mark.usefixtures( "_vcr_outcome_gate_local_testing", "isolate_litellm_state", diff --git a/tests/unit/test_utils_get_model_info.py b/tests/unit/test_utils_get_model_info.py index bfaac44a483..302a7d60cdd 100644 --- a/tests/unit/test_utils_get_model_info.py +++ b/tests/unit/test_utils_get_model_info.py @@ -61,6 +61,30 @@ def test_get_model_info_custom_llm_with_same_name_vllm(monkeypatch): assert model_info["input_cost_per_token"] == 0.0 +def test_get_model_info_passes_through_gemini_tiered_audio_rates() -> None: + model: Final = "gemini-tiered-audio-rates-test" + prices: Final = { + "input_cost_per_audio_token_flex": 1e-6, + "input_cost_per_audio_token_priority": 2e-6, + "cache_read_input_audio_token_cost_flex": 3e-7, + "cache_read_input_audio_token_cost_priority": 4e-7, + "cache_read_input_audio_token_cost_batches": 5e-7, + } + litellm.register_model( + { + model: { + "litellm_provider": "gemini", + "mode": "chat", + **prices, + } + } + ) + + model_info: Final = litellm.get_model_info(model=model, custom_llm_provider="gemini") + + assert {key: model_info[key] for key in prices} == prices + + def test_get_model_info_ollama_chat(): from litellm.llms.ollama.completion.transformation import OllamaConfig @@ -136,6 +160,7 @@ def _enforce_bedrock_converse_models( info.get("litellm_provider") == "bedrock" and info.get("mode") == "chat" and model not in whitelist_models + and BedrockModelInfo.get_bedrock_route(model) != "claude_platform" and not ( (base_model := BedrockModelInfo.get_base_model(model)) != model and model_cost.get(base_model, {}).get("litellm_provider") == "bedrock_converse" diff --git a/tests/unit/videos/test_main.py b/tests/unit/videos/test_main.py index 22e1e5c05eb..76e9b062476 100644 --- a/tests/unit/videos/test_main.py +++ b/tests/unit/videos/test_main.py @@ -32,7 +32,7 @@ reflect production. from contextlib import ExitStack from dataclasses import dataclass -from typing import Any, Dict +from typing import Any, Dict, Final from unittest.mock import MagicMock, patch import pytest @@ -205,9 +205,7 @@ def test_video_edit__dispatch_and_provider_from_id(seams): def test_video_extension__dispatch_and_provider_from_id(seams): - result = videos_main.video_extension( - video_id=AZURE_VIDEO_ID, prompt="continue", seconds="5" - ) + result = videos_main.video_extension(video_id=AZURE_VIDEO_ID, prompt="continue", seconds="5") seams.assert_only("video_extension_handler") assert result is seams.handler.video_extension_handler.return_value @@ -285,9 +283,7 @@ def test_generation__mock_response_short_circuits(seams): def test_list__mock_response_short_circuits(seams): - resp = videos_main.video_list( - mock_response=[{"id": "v1", "object": "video", "status": "completed"}] - ) + resp = videos_main.video_list(mock_response=[{"id": "v1", "object": "video", "status": "completed"}]) assert isinstance(resp, list) assert resp[0].id == "v1" @@ -318,7 +314,7 @@ def test_get_character__mock_response_short_circuits(seams): def test_unsupported_provider_raises_without_dispatch(seams): seams.get_config.return_value = None - with pytest.raises(litellm.APIConnectionError): + with pytest.raises(litellm.BadRequestError, match="video status is not supported for azure"): videos_main.video_status(video_id=AZURE_VIDEO_ID) seams.handler.video_status_handler.assert_not_called() @@ -333,9 +329,7 @@ def test_unsupported_provider_raises_without_dispatch(seams): async def test_avideo_generation__delegates_with_async_flag(): sentinel = VideoObject(id="v-async", object="video", status="queued") with ( - patch.object( - videos_main, "video_generation", MagicMock(return_value=sentinel) - ) as sync, + patch.object(videos_main, "video_generation", MagicMock(return_value=sentinel)) as sync, patch.object( litellm, "get_llm_provider", @@ -352,9 +346,7 @@ async def test_avideo_generation__delegates_with_async_flag(): @pytest.mark.asyncio async def test_avideo_status__delegates_untouched(): sentinel = VideoObject(id="v-async", object="video", status="queued") - with patch.object( - videos_main, "video_status", MagicMock(return_value=sentinel) - ) as sync: + with patch.object(videos_main, "video_status", MagicMock(return_value=sentinel)) as sync: result = await videos_main.avideo_status(video_id="video_plain") assert result is sentinel @@ -367,9 +359,7 @@ async def test_avideo_content__pre_decodes_provider_before_delegating(): """avideo_content resolves the provider from the encoded id itself before handing off, so the sync worker receives the decoded provider, not None.""" sentinel = b"mp4-bytes" - with patch.object( - videos_main, "video_content", MagicMock(return_value=sentinel) - ) as sync: + with patch.object(videos_main, "video_content", MagicMock(return_value=sentinel)) as sync: result = await videos_main.avideo_content(video_id=AZURE_VIDEO_ID) assert result is sentinel @@ -393,9 +383,7 @@ DB_YAML_CREDS = { CREDENTIAL_OPERATIONS = [ ( "video_generation_handler", - lambda: videos_main.video_generation( - prompt="p", model="sora-2", **DB_YAML_CREDS - ), + lambda: videos_main.video_generation(prompt="p", model="sora-2", **DB_YAML_CREDS), ), ( "video_status_handler", @@ -407,21 +395,15 @@ CREDENTIAL_OPERATIONS = [ ), ( "video_remix_handler", - lambda: videos_main.video_remix( - video_id=AZURE_VIDEO_ID, prompt="p", **DB_YAML_CREDS - ), + lambda: videos_main.video_remix(video_id=AZURE_VIDEO_ID, prompt="p", **DB_YAML_CREDS), ), ( "video_edit_handler", - lambda: videos_main.video_edit( - video_id=AZURE_VIDEO_ID, prompt="p", **DB_YAML_CREDS - ), + lambda: videos_main.video_edit(video_id=AZURE_VIDEO_ID, prompt="p", **DB_YAML_CREDS), ), ( "video_extension_handler", - lambda: videos_main.video_extension( - video_id=AZURE_VIDEO_ID, prompt="p", seconds="5", **DB_YAML_CREDS - ), + lambda: videos_main.video_extension(video_id=AZURE_VIDEO_ID, prompt="p", seconds="5", **DB_YAML_CREDS), ), ( "video_list_handler", @@ -429,9 +411,7 @@ CREDENTIAL_OPERATIONS = [ ), ( "video_create_character_handler", - lambda: videos_main.video_create_character( - name="hero", video=MagicMock(name="vid"), **DB_YAML_CREDS - ), + lambda: videos_main.video_create_character(name="hero", video=MagicMock(name="vid"), **DB_YAML_CREDS), ), ( "video_get_character_handler", @@ -453,3 +433,33 @@ def test_db_yaml_credentials_reach_every_handler(seams, handler_name, invoke): assert litellm_params.get("api_base") == DB_YAML_CREDS["api_base"] assert litellm_params.get("api_version") == DB_YAML_CREDS["api_version"] assert litellm_params.get("vertex_project") == DB_YAML_CREDS["vertex_project"] + + +def test_video_generation__custom_provider_map_base_video_config_is_dispatched(monkeypatch: pytest.MonkeyPatch) -> None: + from litellm.llms.openai.videos.transformation import OpenAIVideoConfig + + class QueueVideoConfig(OpenAIVideoConfig): + pass + + video_config: Final = QueueVideoConfig() + monkeypatch.setattr(litellm, "custom_provider_map", [{"provider": "queue-video", "custom_handler": video_config}]) + monkeypatch.setattr(litellm, "_custom_providers", ["queue-video"]) + handler: Final = MagicMock(spec=BaseLLMHTTPHandler) + with patch.object(videos_main, "base_llm_http_handler", handler): + videos_main.video_generation(model="queue-video/seedance-1", prompt="a cat", api_base="http://backend.test") + kw: Final = handler.video_generation_handler.call_args.kwargs + assert kw["video_generation_provider_config"] is video_config + assert kw["custom_llm_provider"] == "queue-video" + assert kw["model"] == "seedance-1" + + +def test_video_generation__custom_provider_without_a_video_config_is_a_400(monkeypatch: pytest.MonkeyPatch) -> None: + from litellm.llms.custom_llm import CustomLLM + + monkeypatch.setattr(litellm, "custom_provider_map", [{"provider": "image-only", "custom_handler": CustomLLM()}]) + monkeypatch.setattr(litellm, "_custom_providers", ["image-only"]) + handler: Final = MagicMock(spec=BaseLLMHTTPHandler) + with patch.object(videos_main, "base_llm_http_handler", handler): + with pytest.raises(litellm.BadRequestError, match="video generation is not supported for image-only"): + videos_main.video_generation(model="image-only/seedance-1", prompt="a cat") + handler.video_generation_handler.assert_not_called() diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/TeamUserSpendCard.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/TeamUserSpendCard.test.tsx new file mode 100644 index 00000000000..c8d14ab4461 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/TeamUserSpendCard.test.tsx @@ -0,0 +1,39 @@ +import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; +import { render, screen } from "@testing-library/react"; +import React from "react"; +import { describe, expect, it, vi } from "vitest"; + +import { teamSpendByUserCall } from "@/components/networking"; + +import TeamUserSpendCard from "./TeamUserSpendCard"; + +vi.mock("@/components/networking", async (importOriginal) => ({ + ...(await importOriginal()), + teamSpendByUserCall: vi.fn(), +})); + +const renderCard = () => { + const client = new QueryClient({ defaultOptions: { queries: { retry: false, gcTime: 0 } } }); + return render( + + + , + ); +}; + +describe("TeamUserSpendCard", () => { + it("shows the server's error instead of an empty table when the request fails", async () => { + const timeoutMessage = "Spend per user took longer than 30s. Pick a shorter date range or fewer teams"; + vi.mocked(teamSpendByUserCall).mockRejectedValueOnce(new Error(timeoutMessage)); + + renderCard(); + + expect(await screen.findByText(timeoutMessage)).toBeInTheDocument(); + expect(screen.queryByText("No user spend in this range")).not.toBeInTheDocument(); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/TeamUserSpendCard.tsx b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/TeamUserSpendCard.tsx index cd488cbd4ea..eb48088f919 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/TeamUserSpendCard.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/TeamUserSpendCard.tsx @@ -73,7 +73,7 @@ const columns: ColumnDef[] = [ const TeamUserSpendCard: React.FC = ({ accessToken, startTime, endTime, teamIds }) => { const hasTeams = teamIds.length > 0; - const { data, isLoading } = useQuery({ + const { data, isLoading, error } = useQuery({ queryKey: ["teamSpendByUser", startTime?.toISOString(), endTime?.toISOString(), teamIds], queryFn: () => accessToken && startTime && endTime ? teamSpendByUserCall(accessToken, startTime, endTime, teamIds) : null, @@ -100,15 +100,21 @@ const TeamUserSpendCard: React.FC = ({ accessToken, star

Attributed per request from spend logs, so it includes JWT/SSO traffic that does not use a virtual key

- + {error !== null ? ( +

+ {error instanceof Error ? error.message : "Spend per user failed to load"} +

+ ) : ( + + )} ); }; diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 867ed435961..89cc9157359 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -35048,6 +35048,12 @@ export interface components { cache_creation_input_token_cost_ultrafast?: number | null; /** Cache Read Input Audio Token Cost */ cache_read_input_audio_token_cost?: number | null; + /** Cache Read Input Audio Token Cost Batches */ + cache_read_input_audio_token_cost_batches?: number | null; + /** Cache Read Input Audio Token Cost Flex */ + cache_read_input_audio_token_cost_flex?: number | null; + /** Cache Read Input Audio Token Cost Priority */ + cache_read_input_audio_token_cost_priority?: number | null; /** Cache Read Input Image Token Cost */ cache_read_input_image_token_cost?: number | null; /** Cache Read Input Token Cost */ @@ -35123,8 +35129,20 @@ export interface components { input_cost_per_audio_per_second_above_128k_tokens?: number | null; /** Input Cost Per Audio Token */ input_cost_per_audio_token?: number | null; + /** Input Cost Per Audio Token Above 200K Tokens */ + input_cost_per_audio_token_above_200k_tokens?: number | null; + /** Input Cost Per Audio Token Above 200K Tokens Batches */ + input_cost_per_audio_token_above_200k_tokens_batches?: number | null; + /** Input Cost Per Audio Token Above 200K Tokens Flex */ + input_cost_per_audio_token_above_200k_tokens_flex?: number | null; + /** Input Cost Per Audio Token Above 200K Tokens Priority */ + input_cost_per_audio_token_above_200k_tokens_priority?: number | null; /** Input Cost Per Audio Token Batches */ input_cost_per_audio_token_batches?: number | null; + /** Input Cost Per Audio Token Flex */ + input_cost_per_audio_token_flex?: number | null; + /** Input Cost Per Audio Token Priority */ + input_cost_per_audio_token_priority?: number | null; /** Input Cost Per Character */ input_cost_per_character?: number | null; /** Input Cost Per Character Above 128K Tokens */ @@ -35157,6 +35175,8 @@ export interface components { input_cost_per_token_above_200k_tokens?: number | null; /** Input Cost Per Token Above 200K Tokens Batches */ input_cost_per_token_above_200k_tokens_batches?: number | null; + /** Input Cost Per Token Above 200K Tokens Flex */ + input_cost_per_token_above_200k_tokens_flex?: number | null; /** Input Cost Per Token Above 200K Tokens Priority */ input_cost_per_token_above_200k_tokens_priority?: number | null; /** Input Cost Per Token Above 272K Tokens */ @@ -35264,6 +35284,8 @@ export interface components { output_cost_per_image_token?: number | null; /** Output Cost Per Image Token Batches */ output_cost_per_image_token_batches?: number | null; + /** Output Cost Per Image Token Priority */ + output_cost_per_image_token_priority?: number | null; /** Output Cost Per Pixel */ output_cost_per_pixel?: number | null; /** Output Cost Per Reasoning Token */ @@ -35298,6 +35320,8 @@ export interface components { output_cost_per_token_above_200k_tokens?: number | null; /** Output Cost Per Token Above 200K Tokens Batches */ output_cost_per_token_above_200k_tokens_batches?: number | null; + /** Output Cost Per Token Above 200K Tokens Flex */ + output_cost_per_token_above_200k_tokens_flex?: number | null; /** Output Cost Per Token Above 200K Tokens Priority */ output_cost_per_token_above_200k_tokens_priority?: number | null; /** Output Cost Per Token Above 272K Tokens */ @@ -50283,6 +50307,12 @@ export interface components { cache_creation_input_token_cost_ultrafast?: number | null; /** Cache Read Input Audio Token Cost */ cache_read_input_audio_token_cost?: number | null; + /** Cache Read Input Audio Token Cost Batches */ + cache_read_input_audio_token_cost_batches?: number | null; + /** Cache Read Input Audio Token Cost Flex */ + cache_read_input_audio_token_cost_flex?: number | null; + /** Cache Read Input Audio Token Cost Priority */ + cache_read_input_audio_token_cost_priority?: number | null; /** Cache Read Input Image Token Cost */ cache_read_input_image_token_cost?: number | null; /** Cache Read Input Token Cost */ @@ -50358,8 +50388,20 @@ export interface components { input_cost_per_audio_per_second_above_128k_tokens?: number | null; /** Input Cost Per Audio Token */ input_cost_per_audio_token?: number | null; + /** Input Cost Per Audio Token Above 200K Tokens */ + input_cost_per_audio_token_above_200k_tokens?: number | null; + /** Input Cost Per Audio Token Above 200K Tokens Batches */ + input_cost_per_audio_token_above_200k_tokens_batches?: number | null; + /** Input Cost Per Audio Token Above 200K Tokens Flex */ + input_cost_per_audio_token_above_200k_tokens_flex?: number | null; + /** Input Cost Per Audio Token Above 200K Tokens Priority */ + input_cost_per_audio_token_above_200k_tokens_priority?: number | null; /** Input Cost Per Audio Token Batches */ input_cost_per_audio_token_batches?: number | null; + /** Input Cost Per Audio Token Flex */ + input_cost_per_audio_token_flex?: number | null; + /** Input Cost Per Audio Token Priority */ + input_cost_per_audio_token_priority?: number | null; /** Input Cost Per Character */ input_cost_per_character?: number | null; /** Input Cost Per Character Above 128K Tokens */ @@ -50392,6 +50434,8 @@ export interface components { input_cost_per_token_above_200k_tokens?: number | null; /** Input Cost Per Token Above 200K Tokens Batches */ input_cost_per_token_above_200k_tokens_batches?: number | null; + /** Input Cost Per Token Above 200K Tokens Flex */ + input_cost_per_token_above_200k_tokens_flex?: number | null; /** Input Cost Per Token Above 200K Tokens Priority */ input_cost_per_token_above_200k_tokens_priority?: number | null; /** Input Cost Per Token Above 272K Tokens */ @@ -50499,6 +50543,8 @@ export interface components { output_cost_per_image_token?: number | null; /** Output Cost Per Image Token Batches */ output_cost_per_image_token_batches?: number | null; + /** Output Cost Per Image Token Priority */ + output_cost_per_image_token_priority?: number | null; /** Output Cost Per Pixel */ output_cost_per_pixel?: number | null; /** Output Cost Per Reasoning Token */ @@ -50533,6 +50579,8 @@ export interface components { output_cost_per_token_above_200k_tokens?: number | null; /** Output Cost Per Token Above 200K Tokens Batches */ output_cost_per_token_above_200k_tokens_batches?: number | null; + /** Output Cost Per Token Above 200K Tokens Flex */ + output_cost_per_token_above_200k_tokens_flex?: number | null; /** Output Cost Per Token Above 200K Tokens Priority */ output_cost_per_token_above_200k_tokens_priority?: number | null; /** Output Cost Per Token Above 272K Tokens */ diff --git a/uv.lock b/uv.lock index 81ed8c0f427..674ffa049ef 100644 --- a/uv.lock +++ b/uv.lock @@ -4572,7 +4572,7 @@ wheels = [ [[package]] name = "litellm" -version = "1.108.0" +version = "1.108.1" source = { editable = "." } dependencies = [ { name = "aiohttp" }, @@ -5072,7 +5072,7 @@ provides-extras = ["test"] [[package]] name = "litellm-enterprise" -version = "0.1.79" +version = "0.1.79.post1" source = { editable = "enterprise" } [[package]]