diff --git a/enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py b/enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py index b2eda76f9ae..f40ced302ce 100644 --- a/enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py +++ b/enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py @@ -780,7 +780,10 @@ async def update_project( # Handle budget updates budget_fields = LiteLLM_BudgetTable.model_fields.keys() - budget_updates = {k: v for k, v in update_data.items() if k in budget_fields} + budget_updates = { + **{k: v for k, v in update_data.items() if k in budget_fields}, + **({"max_budget": None} if "max_budget" in data.model_fields_set and data.max_budget is None else {}), + } if budget_updates and existing_project.budget_id: # Update existing budget diff --git a/enterprise/pyproject.toml b/enterprise/pyproject.toml index 903c5155a12..c049bf68c46 100644 --- a/enterprise/pyproject.toml +++ b/enterprise/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "litellm-enterprise" -version = "0.1.66" +version = "0.1.67" 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.66" +version = "0.1.67" version_files = [ "pyproject.toml:^version", "../pyproject.toml:litellm-enterprise==", diff --git a/litellm-proxy-extras/litellm_proxy_extras/prisma_toolchain.py b/litellm-proxy-extras/litellm_proxy_extras/prisma_toolchain.py index b51de9609d3..9cd48fcf11a 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/prisma_toolchain.py +++ b/litellm-proxy-extras/litellm_proxy_extras/prisma_toolchain.py @@ -37,11 +37,13 @@ raised it above the deploy default keeps that larger budget for deploy unless the deploy override says otherwise. """ +import importlib.util import math import os import shutil import signal import subprocess +import sys from collections.abc import Mapping, Sequence from dataclasses import dataclass from pathlib import Path @@ -64,6 +66,7 @@ DEFAULT_PRISMA_BOOTSTRAP_TIMEOUT = 600.0 DEFAULT_PRISMA_MIGRATE_DEPLOY_TIMEOUT = 600.0 BOOTSTRAP_ARG = "--version" +PRISMA_CONSOLE_SCRIPT = "prisma" @dataclass(frozen=True) @@ -184,6 +187,28 @@ def _kill_process_group(process: "subprocess.Popen[str]") -> None: return +def prisma_cli_available() -> bool: + """Whether some way of running the Prisma CLI exists: the console script on PATH or the importable package.""" + if shutil.which(PRISMA_CONSOLE_SCRIPT) is not None: + return True + return importlib.util.find_spec(PRISMA_CONSOLE_SCRIPT) is not None + + +def resolve_prisma_argv(argv: Sequence[str]) -> tuple[str, ...]: + """Route a bare ``prisma`` command through ``python -m prisma`` when the console script is not on PATH. + + The console script and ``python -m prisma`` are the same entry point, but + only the module form survives an interpreter whose ``bin`` directory is + missing from PATH, which is how the proxy gets started under launchers and + init systems. Any other executable name is left untouched. + """ + if not argv or argv[0] != PRISMA_CONSOLE_SCRIPT: + return tuple(argv) + if shutil.which(PRISMA_CONSOLE_SCRIPT) is not None: + return tuple(argv) + return (sys.executable, "-m", PRISMA_CONSOLE_SCRIPT, *argv[1:]) + + def run_prisma( argv: Sequence[str], *, @@ -200,7 +225,7 @@ def run_prisma( text unless ``stdout``/``stderr`` say otherwise. """ with subprocess.Popen( - argv, + resolve_prisma_argv(argv), env=env, stdout=stdout, stderr=stderr, diff --git a/litellm-proxy-extras/pyproject.toml b/litellm-proxy-extras/pyproject.toml index 7d4c78088f1..f94591872a4 100644 --- a/litellm-proxy-extras/pyproject.toml +++ b/litellm-proxy-extras/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "litellm-proxy-extras" -version = "0.4.96" +version = "0.4.97" description = "Additional files for the LiteLLM Proxy. Reduces the size of the main litellm package." readme = "README.md" requires-python = ">=3.9" @@ -26,7 +26,7 @@ required-version = ">=0.10.9" module-root = "" [tool.commitizen] -version = "0.4.96" +version = "0.4.97" version_files = [ "pyproject.toml:^version", "../pyproject.toml:litellm-proxy-extras==", diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr/callbacks.rs b/litellm-rust/crates/python-bridge/src/routes/ocr/callbacks.rs index 1cbe8a179e3..c7e5f123c19 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr/callbacks.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/callbacks.rs @@ -58,6 +58,13 @@ impl PythonLogger { params.set_item(name, value)?; } } + for name in custom_pricing_fields(py)? { + if let Some(value) = kwargs.bind(py).get_item(&name)? + && !value.is_none() + { + params.set_item(name, value)?; + } + } update.set_item("litellm_params", params)?; update.set_item("custom_llm_provider", &pre_call.custom_llm_provider)?; self.object(py) @@ -120,6 +127,17 @@ impl PythonLogger { } } +fn custom_pricing_fields(py: Python<'_>) -> PyResult> { + py.import("litellm.types.utils")? + .getattr("CustomPricingLiteLLMParams")? + .getattr("model_fields")? + .cast_into::()? + .keys() + .iter() + .map(|name| name.extract::()) + .collect() +} + fn redact( py: Python<'_>, params: &Bound<'_, PyDict>, diff --git a/litellm/__init__.py b/litellm/__init__.py index 1dfd146a00e..ccfbf80369f 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -501,6 +501,7 @@ disable_copilot_system_to_assistant: bool = False # If false (default), convert public_mcp_servers: Optional[List[str]] = None public_mcp_hub_strict_whitelist: bool = True public_model_groups: Optional[List[str]] = None +public_skills_index: bool = False public_agent_groups: Optional[List[str]] = None agent_search_embedding_model: Optional[str] = None mcp_tool_search: Optional[Mapping[str, object]] = None diff --git a/litellm/caching/redis_cache.py b/litellm/caching/redis_cache.py index 1a024b6b2b0..d0cefcb6086 100644 --- a/litellm/caching/redis_cache.py +++ b/litellm/caching/redis_cache.py @@ -16,7 +16,7 @@ import inspect import json import logging import time -from collections.abc import Awaitable, Callable, Sequence +from collections.abc import Awaitable, Callable, Iterator, Sequence from contextvars import ContextVar from dataclasses import dataclass from datetime import timedelta @@ -317,19 +317,37 @@ def _is_redis_health_failure(exc: BaseException) -> bool: def _redis_timeout_error_types() -> tuple[type, ...]: """Health failures that are timeouts rather than unambiguous connectivity errors. - ``builtins.TimeoutError`` covers ``asyncio.TimeoutError`` and ``socket.timeout`` - (aliases since py3.11 / py3.10). ``redis.exceptions.TimeoutError`` does not subclass - either, so it is listed explicitly. + ``builtins.TimeoutError`` covers ``socket.timeout`` (an alias since py3.10) and, from + py3.11, ``asyncio.TimeoutError``; on py3.10 ``asyncio.TimeoutError`` is still its own + class, so it is listed explicitly. ``redis.exceptions.TimeoutError`` subclasses neither. """ try: from redis.exceptions import TimeoutError as RedisTimeoutError except ImportError: - return (TimeoutError,) - return (RedisTimeoutError, TimeoutError) + return (TimeoutError, asyncio.TimeoutError) + return (RedisTimeoutError, TimeoutError, asyncio.TimeoutError) + + +_MAX_EXCEPTION_CAUSE_DEPTH: Final = 20 + + +def _explicit_causes(exc: BaseException) -> Iterator[BaseException]: + current = exc # rebind-ok: advances one link per iteration of the bounded walk + for _ in range(_MAX_EXCEPTION_CAUSE_DEPTH): + yield current + if current.__cause__ is None: + return + current = current.__cause__ def _is_redis_timeout_failure(exc: BaseException) -> bool: - return isinstance(exc, _redis_timeout_error_types()) + """True when ``exc`` or any exception it was explicitly raised ``from`` is a timeout. + + redis-py's blocking pool reports a pool wait timeout as ``ConnectionError`` chained from + ``asyncio.TimeoutError``, which is a busy pool rather than an unreachable Redis. + """ + timeout_types: Final = _redis_timeout_error_types() + return any(isinstance(link, timeout_types) for link in _explicit_causes(exc)) class _BreakerMetrics: diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 50118af8a30..04cd971c16e 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -9,6 +9,7 @@ from typing import TYPE_CHECKING, Any, Final, Literal, cast from httpx import Response from pydantic import BaseModel +from typing_extensions import ReadOnly, TypedDict import litellm import litellm._logging @@ -311,6 +312,15 @@ def _transcription_usage_has_token_details( return (prompt_tokens_val > 0) or (completion_tokens_val > 0) +OCRPricingField = Literal["ocr_cost_per_page", "ocr_cost_per_credit", "annotation_cost_per_page"] + + +class OCRPricing(TypedDict, total=False): + ocr_cost_per_page: ReadOnly[float | None] + ocr_cost_per_credit: ReadOnly[float | None] + annotation_cost_per_page: ReadOnly[float | None] + + def cost_per_token( model: str = "", prompt_tokens: int = 0, @@ -345,6 +355,7 @@ def cost_per_token( response: Any | None = None, ### REQUEST MODEL ### request_model: str | None = None, # original request model for router detection + custom_model_info: OCRPricing | None = None, ) -> tuple[float, float]: """ Calculates the cost per token for a given model, prompt tokens, and completion tokens. @@ -559,6 +570,7 @@ def cost_per_token( model=model, custom_llm_provider=custom_llm_provider, response=response, + model_info=custom_model_info, ) elif ( call_type == "aretrieve_batch" @@ -1433,20 +1445,9 @@ def completion_cost( ) elif call_type in _VIDEO_CALL_TYPES: ### VIDEO GENERATION COST CALCULATION ### - # Extract custom model_info for deployment-specific pricing - _video_model_info: ModelInfo | None = None - if custom_pricing and litellm_logging_obj is not None: - _litellm_params = getattr(litellm_logging_obj, "litellm_params", None) - if _litellm_params is not None: - _video_model_info = next( - ( - model_info - for _metadata_key in ("metadata", "litellm_metadata") - if (model_info := (_litellm_params.get(_metadata_key) or {}).get("model_info")) - is not None - ), - None, - ) + _video_model_info: ModelInfo | None = _deployment_model_info( + litellm_logging_obj, custom_pricing, router_model_id + ) usage_obj = getattr(completion_response, "usage", None) duration_seconds: float | None = None @@ -1666,6 +1667,7 @@ def completion_cost( data_residency=data_residency, vertex_location=vertex_location, response=completion_response, + custom_model_info=_ocr_model_info(litellm_logging_obj, custom_pricing, router_model_id), ) # Get additional costs from provider (e.g., routing fees, infrastructure costs) @@ -1899,16 +1901,82 @@ def response_cost_calculator( raise e +def _deployment_model_info( + litellm_logging_obj: LitellmLoggingObject | None, + custom_pricing: bool | None, + router_model_id: str | None, +) -> ModelInfo | None: + if not custom_pricing: + return None + registered_deployment_info: Final = ( + _cost_map_model_info(router_model_id, None) + if router_model_id is not None and router_model_id in litellm.model_cost + else None + ) + if registered_deployment_info is not None: + return registered_deployment_info + if litellm_logging_obj is None: + return None + litellm_params: Final = getattr(litellm_logging_obj, "litellm_params", None) + if litellm_params is None: + return None + return next( + ( + model_info + for metadata_key in ("metadata", "litellm_metadata") + if (metadata := litellm_params.get(metadata_key)) and (model_info := metadata.get("model_info")) is not None + ), + None, + ) + + +def _ocr_model_info( + litellm_logging_obj: LitellmLoggingObject | None, + custom_pricing: bool | None, + router_model_id: str | None, +) -> OCRPricing | None: + deployment_info: Final = _deployment_model_info(litellm_logging_obj, custom_pricing, router_model_id) + litellm_params: Final = getattr(litellm_logging_obj, "litellm_params", None) if custom_pricing else None + if litellm_params is None: + return deployment_info + return _layered_ocr_pricing(litellm_params, deployment_info) + + +def _first_ocr_price(field: OCRPricingField, *sources: Mapping[str, object] | None) -> float | None: + return next( + (price for source in sources if source is not None and isinstance(price := source.get(field), int | float)), + None, + ) + + +def _layered_ocr_pricing(*sources: Mapping[str, object] | None) -> OCRPricing: + return OCRPricing( + ocr_cost_per_page=_first_ocr_price("ocr_cost_per_page", *sources), + ocr_cost_per_credit=_first_ocr_price("ocr_cost_per_credit", *sources), + annotation_cost_per_page=_first_ocr_price("annotation_cost_per_page", *sources), + ) + + +def _cost_map_model_info(model: str, custom_llm_provider: str | None) -> ModelInfo | None: + try: + return litellm.get_model_info(model=model, custom_llm_provider=custom_llm_provider) + except Exception: + return None + + def ocr_cost( model: str, custom_llm_provider: str | None, response: object | None = None, + model_info: OCRPricing | None = None, ) -> tuple[float, float]: """ Args: model: str - model name custom_llm_provider: Optional[str] - custom LLM provider response: Optional[Any] - response object + model_info: Optional[OCRPricing] - deployment-specific OCR pricing; each rate it sets + overrides the model cost map's, the rest fall back to the map Returns: Tuple[float, float]: cost of OCR processing @@ -1926,20 +1994,15 @@ def ocr_cost( if response.usage_info is None: raise ValueError("OCR response usage_info is None") - try: - model_info: ModelInfo | None = litellm.get_model_info(model=model, custom_llm_provider=custom_llm_provider) - except Exception: - model_info = None - credits: Final = getattr(response.usage_info, "credits", None) - cost_per_credit = None - if model_info is not None: - cost_per_credit = model_info.get("ocr_cost_per_credit") + pricing: Final = _layered_ocr_pricing(model_info, _cost_map_model_info(model, custom_llm_provider)) + + cost_per_credit: Final = pricing.get("ocr_cost_per_credit") if credits is not None and cost_per_credit is not None: return cost_per_credit * credits, 0.0 - ocr_cost_per_page: Final = model_info.get("ocr_cost_per_page") if model_info is not None else None - annotation_cost_per_page: Final = model_info.get("annotation_cost_per_page") if model_info is not None else None + ocr_cost_per_page: Final = pricing.get("ocr_cost_per_page") + annotation_cost_per_page: Final = pricing.get("annotation_cost_per_page") annotation_rate: Final = annotation_cost_per_page if annotation_cost_per_page is not None else ocr_cost_per_page pages_processed: Final = response.usage_info.pages_processed diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 8a976a966a6..39adea30828 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -5,6 +5,7 @@ import os import secrets from collections.abc import Mapping from datetime import datetime +from types import MappingProxyType from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, Optional, get_args from litellm._logging import verbose_logger @@ -1279,6 +1280,7 @@ class CustomGuardrail(CustomLogger): guardrail_response: Final = self._summarize_guardrail_response( response=response, original_inputs=original_inputs, + event_type=event_type, ) verbose_logger.debug("Guardrail response: %s", response) @@ -1298,6 +1300,7 @@ class CustomGuardrail(CustomLogger): self, response: object, original_inputs: Mapping[str, object] | None, + event_type: GuardrailEventHooks | None, ) -> object: """Reduce a hook's return value to what is safe to log as ``guardrail_response``. @@ -1305,15 +1308,21 @@ class CustomGuardrail(CustomLogger): returns the (possibly modified) request payload. Neither is a provider verdict, and logging them verbatim ships the user's prompt to every logging sink (OTEL spans, Datadog, spend logs), so both collapse to ``"allow"`` / ``"mask"`` by comparing - against ``original_inputs``, a copy taken before the hook ran. A string result is the - hook's own rejection message (the proxy turns it into a 400), not user input, so it is - logged as is. + against ``original_inputs``, a copy taken before the hook ran. A pre_call baseline only + holds the prompt-bearing keys, so the returned request is narrowed to those same keys + before the comparison. A string result is the hook's own rejection message (the proxy + turns it into a 400), not user input, so it is logged as is. """ if response is None: return {} if original_inputs is None or not isinstance(response, Mapping): return response - return "mask" if self._inputs_were_modified(original_inputs, response) else "allow" + compared_response: Final[Mapping[str, object]] = ( + MappingProxyType({key: value for key, value in response.items() if key in _PRE_CALL_CONTENT_KEYS}) + if event_type == GuardrailEventHooks.pre_call + else response + ) + return "mask" if self._inputs_were_modified(original_inputs, compared_response) else "allow" @staticmethod def _is_guardrail_intervention(e: Exception) -> bool: @@ -1355,8 +1364,8 @@ class CustomGuardrail(CustomLogger): raise e def _inputs_were_modified(self, original_inputs: Mapping[str, object], response: Mapping[str, object]) -> bool: - """True when any baseline key's value differs in ``response`` (mask), False otherwise (allow).""" - return any(response.get(key) != value for key, value in original_inputs.items()) + """True when any key of either mapping differs between them (mask), False otherwise (allow).""" + return any(original_inputs.get(key) != response.get(key) for key in original_inputs.keys() | response.keys()) def mask_content_in_string( self, @@ -1476,13 +1485,13 @@ def _original_inputs_for( ) -> dict | None: # mutable-ok: matches _process_response(original_inputs=) signature """Baseline the hook's return value is compared against to decide "allow" vs "mask". - ``apply_guardrail`` masks a fresh ``inputs`` dict, so that dict is the baseline. Pre-call - hooks edit the request in place and return it, so the baseline is a deep copy of the - prompt-bearing keys taken before the hook runs. + Hooks may edit their argument in place and return it, so the baseline is always a deep + copy taken before the hook runs: the whole ``inputs`` dict for ``apply_guardrail``, the + prompt-bearing request keys for pre-call hooks. """ if func_name == "apply_guardrail": inputs: Final = kwargs.get("inputs") - return inputs if isinstance(inputs, dict) else None + return copy.deepcopy(inputs) if isinstance(inputs, dict) else None if event_type != GuardrailEventHooks.pre_call: return None return {key: copy.deepcopy(value) for key, value in request_data.items() if key in _PRE_CALL_CONTENT_KEYS} diff --git a/litellm/integrations/gcs_bucket/gcs_bucket.py b/litellm/integrations/gcs_bucket/gcs_bucket.py index 31ceb338dcd..e338f490496 100644 --- a/litellm/integrations/gcs_bucket/gcs_bucket.py +++ b/litellm/integrations/gcs_bucket/gcs_bucket.py @@ -29,8 +29,6 @@ class GCSBucketLogger(GCSBucketBase, AdditionalLoggingUtils): def __init__(self, bucket_name: str | None = None) -> None: from litellm.proxy.proxy_server import premium_user - super().__init__(bucket_name=bucket_name) - self.batch_size = int(os.getenv("GCS_BATCH_SIZE", GCS_DEFAULT_BATCH_SIZE)) self.flush_interval = int(os.getenv("GCS_FLUSH_INTERVAL", GCS_DEFAULT_FLUSH_INTERVAL_SECONDS)) self.use_batched_logging = ( @@ -38,6 +36,7 @@ class GCSBucketLogger(GCSBucketBase, AdditionalLoggingUtils): ) self.flush_lock = asyncio.Lock() super().__init__( + bucket_name=bucket_name, flush_lock=self.flush_lock, batch_size=self.batch_size, flush_interval=self.flush_interval, diff --git a/litellm/litellm_core_utils/realtime_streaming.py b/litellm/litellm_core_utils/realtime_streaming.py index 4923bdda305..e3f8786a39a 100644 --- a/litellm/litellm_core_utils/realtime_streaming.py +++ b/litellm/litellm_core_utils/realtime_streaming.py @@ -445,6 +445,12 @@ class RealTimeStreaming: ) sent = False for msg in transformed: + if isinstance(msg, bytes): + await self.provider_config.pace_backend_send(msg) + await self.backend_ws.send(msg) + self._content_sent_after_setup = True + sent = True + continue try: msg_obj = _decode_json_object(msg) except (json.JSONDecodeError, TypeError): @@ -1013,7 +1019,7 @@ class RealTimeStreaming: cast(str, transcript), item_id=cast(str | None, event.get("item_id")), ) - if not blocked: + if not blocked and not self._is_transcription_session: await self._send_to_backend(json.dumps({"type": "response.create"})) continue ## LOGGING diff --git a/litellm/llms/base_llm/realtime/transformation.py b/litellm/llms/base_llm/realtime/transformation.py index cfcde7c6e9e..e44cccc1a62 100644 --- a/litellm/llms/base_llm/realtime/transformation.py +++ b/litellm/llms/base_llm/realtime/transformation.py @@ -1,4 +1,5 @@ from abc import ABC, abstractmethod +from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Any import httpx @@ -54,9 +55,12 @@ class BaseRealtimeConfig(ABC): message: str, model: str, session_configuration_request: str | None = None, - ) -> list[str]: + ) -> Sequence[str | bytes]: pass + async def pace_backend_send(self, message: bytes) -> None: + return None + def is_setup_message(self, msg_obj: dict) -> bool: return False @@ -79,7 +83,7 @@ class BaseRealtimeConfig(ABC): model: str, logging_session_id: str, session_configuration_request: str | None = None, - ) -> dict | OpenAIRealtimeStreamSessionEvents | None: + ) -> Mapping[str, object] | OpenAIRealtimeStreamSessionEvents | None: """ Optional hook for providers that defer session setup until client `session.update`. diff --git a/litellm/llms/meta/realtime/transformation.py b/litellm/llms/meta/realtime/transformation.py new file mode 100644 index 00000000000..1b8943f0cee --- /dev/null +++ b/litellm/llms/meta/realtime/transformation.py @@ -0,0 +1,719 @@ +import asyncio +import base64 +import binascii +import json +import math +import time +from collections.abc import Awaitable, Callable, Iterator, Mapping +from dataclasses import dataclass +from types import MappingProxyType +from typing import Final, Literal +from urllib.parse import urlparse, urlunparse + +from pydantic import JsonValue, TypeAdapter, ValidationError + +from litellm import verbose_logger +from litellm._uuid import uuid +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.llms.base_llm.realtime.transformation import BaseRealtimeConfig +from litellm.secret_managers.main import get_secret_str +from litellm.types.llms.meta import MuseAudioEncoding, MuseHandshake, MuseMode, MuseSampleRate +from litellm.types.llms.openai import ( + OpenAIRealtimeErrorEvent, + OpenAIRealtimeEvents, + OpenAIRealtimeInputAudioBufferSpeechEvent, + OpenAIRealtimeInputAudioTranscriptionCompleted, + OpenAIRealtimeInputAudioTranscriptionDelta, + OpenAIRealtimeServerVadTurnDetection, + OpenAIRealtimeTranscriptionSession, + OpenAIRealtimeTranscriptionSessionCreated, + OpenAIRealtimeTranscriptionSettings, +) +from litellm.types.realtime import ( + RealtimeInputAudioTranscriptionDurationUsage, + RealtimeInputAudioTranscriptionUsage, + RealtimeResponseTransformInput, + RealtimeResponseTypedDict, +) + +MUSE_MODEL: Final = "muse-voice-transcribe-1.0" +DEFAULT_MUSE_REALTIME_URL: Final = "wss://api.meta.ai/v1/asr/realtime" +SUPPORTED_SAMPLE_RATES: Final = frozenset((16_000, 24_000)) +SUPPORTED_LANGUAGES: Final = ( + "Arabic", + "Bengali", + "Dutch", + "English", + "French", + "German", + "Hebrew", + "Hindi", + "Indonesian", + "Italian", + "Japanese", + "Kannada", + "Korean", + "Malay", + "Mandarin Chinese", + "Marathi", + "Polish", + "Portuguese", + "Spanish", + "Tagalog", + "Tamil", + "Telugu", + "Thai", + "Turkish", + "Vietnamese", +) +_LANGUAGE_NAMES: Final = MappingProxyType({language.casefold(): language for language in SUPPORTED_LANGUAGES}) +_LANGUAGE_CODES: Final = MappingProxyType( + { + "ar": "Arabic", + "bn": "Bengali", + "de": "German", + "en": "English", + "es": "Spanish", + "fil": "Tagalog", + "fr": "French", + "he": "Hebrew", + "hi": "Hindi", + "id": "Indonesian", + "it": "Italian", + "iw": "Hebrew", + "ja": "Japanese", + "kn": "Kannada", + "ko": "Korean", + "ms": "Malay", + "mr": "Marathi", + "nl": "Dutch", + "pl": "Polish", + "pt": "Portuguese", + "ta": "Tamil", + "te": "Telugu", + "th": "Thai", + "tl": "Tagalog", + "tr": "Turkish", + "vi": "Vietnamese", + "zh": "Mandarin Chinese", + } +) +_SUPPORTED_TRANSCRIPTION_KEYS: Final = frozenset(("model", "language")) +_MAX_AUDIO_BACKLOG_SECONDS: Final = 4 +_PACKET_MS: Final = 80 +_END_STREAM: Final = '{"type":"endStream"}' +_PROVIDER_ERROR_MESSAGE: Final = "Meta Muse realtime transcription failed" +_JSON_ADAPTER: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue) +_EMPTY_OBJECT: Final[Mapping[str, JsonValue]] = MappingProxyType({}) +_SERVER_VAD: Final[OpenAIRealtimeServerVadTurnDetection] = {"type": "server_vad"} + + +class MuseProtocolError(ValueError): + pass + + +@dataclass(frozen=True, slots=True) +class MuseSessionConfig: + model: str + mode: MuseMode + sample_rate: MuseSampleRate + language_bias: tuple[str, ...] + + @property + def audio_encoding(self) -> MuseAudioEncoding: + return "PCM_16KHZ" if self.sample_rate == 16_000 else "PCM_24KHZ" + + @property + def bytes_per_second(self) -> int: + return self.sample_rate * 2 + + @property + def packet_bytes(self) -> int: + return self.bytes_per_second * _PACKET_MS // 1000 + + @property + def max_encoded_append_bytes(self) -> int: + return 4 * ((self.bytes_per_second * _MAX_AUDIO_BACKLOG_SECONDS + 2) // 3) + + def handshake(self, access_token: str) -> MuseHandshake: + base: Final[MuseHandshake] = { + "authorization": {"accessToken": access_token}, + "audioEncoding": self.audio_encoding, + "model": self.model, + "mode": self.mode, + "partialMode": "CUMULATIVE", + "emitAudioProgress": True, + } + if not self.language_bias: + return base + biased: Final[MuseHandshake] = {**base, "languageBias": self.language_bias} + return biased + + def openai_session(self, session_id: str) -> OpenAIRealtimeTranscriptionSession: + session: Final[OpenAIRealtimeTranscriptionSession] = { + "id": session_id, + "object": "realtime.transcription_session", + "type": "transcription", + "audio": { + "input": { + "format": {"type": "audio/pcm", "rate": self.sample_rate}, + "transcription": self._transcription_settings(), + "turn_detection": None if self.mode == "PUSH_TO_TALK" else _SERVER_VAD, + } + }, + } + return session + + def _transcription_settings(self) -> OpenAIRealtimeTranscriptionSettings: + base: Final[OpenAIRealtimeTranscriptionSettings] = {"model": self.model} + if not self.language_bias: + return base + localized: Final[OpenAIRealtimeTranscriptionSettings] = {**base, "language": self.language_bias[0]} + return localized + + +_DEFAULT_SESSION_CONFIG: Final = MuseSessionConfig( + model=MUSE_MODEL, mode="ENDPOINTING", sample_rate=24_000, language_bias=() +) + + +def _json_object(payload: str) -> Mapping[str, JsonValue]: + try: + value: Final = _JSON_ADAPTER.validate_json(payload) + except ValidationError: + raise MuseProtocolError("invalid JSON object") from None + if not isinstance(value, dict): + raise MuseProtocolError("message must be a JSON object") + return value + + +def _mapping(value: JsonValue | None, name: str) -> Mapping[str, JsonValue]: + if value is None: + return _EMPTY_OBJECT + if not isinstance(value, dict): + raise MuseProtocolError(f"{name} must be an object") + return value + + +def _string(value: JsonValue | None, name: str) -> str | None: + if value is None: + return None + if not isinstance(value, str): + raise MuseProtocolError(f"{name} must be a string") + return value + + +def _normalize_model(model: str) -> str: + return model.removeprefix("meta/").strip() + + +def _event_id() -> str: + return f"event_{uuid.uuid4().hex}" + + +def normalize_language(language: str) -> str: + value: Final = language.strip() + if not value: + raise MuseProtocolError("language must be non-empty") + documented_name: Final = _LANGUAGE_NAMES.get(value.casefold()) + if documented_name is not None: + return documented_name + primary: Final = value.replace("_", "-").split("-", 1)[0].casefold() + mapped_name: Final = _LANGUAGE_CODES.get(primary) + if mapped_name is None: + raise MuseProtocolError("unsupported Muse Voice language") + return mapped_name + + +def normalize_access_token(api_key: str) -> str: + stripped: Final = api_key.strip() + if not stripped: + raise ValueError("Meta API key is required") + parts: Final = stripped.split(None, 1) + if parts[0].casefold() != "bearer": + return f"Bearer {stripped}" + if len(parts) != 2 or not parts[1].strip(): + raise ValueError("Meta API key must include a token after Bearer") + return f"Bearer {parts[1].strip()}" + + +def build_muse_realtime_url(api_base: str | None) -> str: + if api_base is None: + return DEFAULT_MUSE_REALTIME_URL + parsed: Final = urlparse(api_base.strip()) + scheme: Final = "wss" if parsed.scheme == "https" else parsed.scheme + if ( + scheme != "wss" + or not parsed.hostname + or parsed.username is not None + or parsed.password is not None + or parsed.fragment + ): + raise ValueError("Meta api_base must be an absolute wss:// or https:// URL without credentials or a fragment") + netloc: Final = f"{parsed.hostname}:{parsed.port}" if parsed.port is not None else parsed.hostname + return urlunparse((scheme, netloc, "/v1/asr/realtime", "", "", "")) + + +def _parse_sample_rate(session: Mapping[str, JsonValue]) -> MuseSampleRate: + beta_format: Final = session.get("input_audio_format") + audio: Final = _mapping(session.get("audio"), "session.audio") + audio_input: Final = _mapping(audio.get("input"), "session.audio.input") + ga_format: Final = audio_input.get("format") + if beta_format is not None and ga_format is not None: + raise MuseProtocolError("input audio format must use either beta or GA layout") + if beta_format is not None: + if beta_format != "pcm16": + raise MuseProtocolError("Muse Voice requires pcm16 input audio") + return 24_000 + if ga_format is None: + return 24_000 + if isinstance(ga_format, str): + if ga_format != "pcm16": + raise MuseProtocolError("Muse Voice requires audio/pcm input audio") + return 24_000 + format_mapping: Final = _mapping(ga_format, "session.audio.input.format") + if format_mapping.get("type") != "audio/pcm": + raise MuseProtocolError("Muse Voice requires audio/pcm input audio") + channels: Final = format_mapping.get("channels", 1) + if isinstance(channels, bool) or channels != 1: + raise MuseProtocolError("Muse Voice requires mono input audio") + rate: Final = format_mapping.get("rate", 24_000) + if isinstance(rate, bool) or not isinstance(rate, int) or rate not in SUPPORTED_SAMPLE_RATES: + raise MuseProtocolError("Muse Voice supports PCM16 at 16000 Hz or 24000 Hz") + return 16_000 if rate == 16_000 else 24_000 + + +def _parse_mode(session: Mapping[str, JsonValue], audio_input: Mapping[str, JsonValue]) -> MuseMode: + turn_detection_present: Final = "turn_detection" in session or "turn_detection" in audio_input + turn_detection: Final = session.get("turn_detection", audio_input.get("turn_detection")) + if turn_detection_present and turn_detection is None: + return "PUSH_TO_TALK" + if turn_detection is None: + return "ENDPOINTING" + turn_detection_mapping: Final = _mapping(turn_detection, "turn_detection") + if turn_detection_mapping.get("type") not in (None, "server_vad"): + raise MuseProtocolError("Muse Voice supports server_vad turn detection or null") + return "ENDPOINTING" + + +def parse_session_update(payload: str, expected_model: str) -> MuseSessionConfig: + message: Final = _json_object(payload) + if message.get("type") not in ("session.update", "transcription_session.update"): + raise MuseProtocolError("expected session.update") + session: Final = _mapping(message.get("session"), "session") + if not session: + raise MuseProtocolError("session.update requires a session object") + if session.get("type") not in (None, "transcription", "realtime"): + raise MuseProtocolError("Muse Voice supports transcription sessions only") + audio: Final = _mapping(session.get("audio"), "session.audio") + audio_input: Final = _mapping(audio.get("input"), "session.audio.input") + beta_transcription: Final = session.get("input_audio_transcription") + ga_transcription: Final = audio_input.get("transcription") + if beta_transcription is not None and ga_transcription is not None: + raise MuseProtocolError("input transcription must use either beta or GA layout") + transcription: Final = _mapping( + beta_transcription if beta_transcription is not None else ga_transcription, + "input audio transcription", + ) + unsupported: Final = tuple(sorted(key for key in transcription if key not in _SUPPORTED_TRANSCRIPTION_KEYS)) + if unsupported: + verbose_logger.warning("Meta realtime: dropping unsupported transcription settings %s", unsupported) + requested_model: Final = _string(transcription.get("model"), "transcription model") + normalized_model: Final = _normalize_model(expected_model) + if normalized_model != MUSE_MODEL: + raise MuseProtocolError("unsupported Meta realtime model") + if requested_model is not None and _normalize_model(requested_model) != normalized_model: + raise MuseProtocolError("realtime session model cannot be changed") + language: Final = _string(transcription.get("language"), "language") + return MuseSessionConfig( + model=normalized_model, + mode=_parse_mode(session, audio_input), + sample_rate=_parse_sample_rate(session), + language_bias=() if language is None else (normalize_language(language),), + ) + + +def session_created_event(config: MuseSessionConfig, session_id: str) -> OpenAIRealtimeTranscriptionSessionCreated: + event: Final[OpenAIRealtimeTranscriptionSessionCreated] = { + "type": "session.created", + "event_id": _event_id(), + "session": config.openai_session(session_id), + } + return event + + +def error_event(message: str) -> OpenAIRealtimeErrorEvent: + event: Final[OpenAIRealtimeErrorEvent] = { + "type": "error", + "error": {"type": "server_error", "message": message}, + } + return event + + +def _speech_event( + event_type: Literal["input_audio_buffer.speech_started", "input_audio_buffer.speech_stopped"], item_id: str +) -> OpenAIRealtimeInputAudioBufferSpeechEvent: + event: Final[OpenAIRealtimeInputAudioBufferSpeechEvent] = { + "type": event_type, + "event_id": _event_id(), + "item_id": item_id, + } + return event + + +def _delta_event(item_id: str, delta: str) -> OpenAIRealtimeInputAudioTranscriptionDelta: + event: Final[OpenAIRealtimeInputAudioTranscriptionDelta] = { + "type": "conversation.item.input_audio_transcription.delta", + "event_id": _event_id(), + "item_id": item_id, + "content_index": 0, + "delta": delta, + } + return event + + +def _completed_event( + item_id: str, transcript: str, usage: RealtimeInputAudioTranscriptionUsage | None +) -> OpenAIRealtimeInputAudioTranscriptionCompleted: + event: Final[OpenAIRealtimeInputAudioTranscriptionCompleted] = { + "type": "conversation.item.input_audio_transcription.completed", + "event_id": _event_id(), + "item_id": item_id, + "content_index": 0, + "transcript": transcript, + } + if usage is None: + return event + billed: Final[OpenAIRealtimeInputAudioTranscriptionCompleted] = {**event, "usage": usage} + return billed + + +def _required_turn_id(message: Mapping[str, JsonValue], event: str) -> str: + value: Final = message.get("turnId") + if isinstance(value, bool) or not isinstance(value, (str, int)): + raise MuseProtocolError(f"{event} event has invalid turnId") + turn_id: Final = str(value).strip() + if not turn_id: + raise MuseProtocolError(f"{event} event has invalid turnId") + return turn_id + + +def _new_suffix(previous: str, current: str) -> str: + return current[len(previous) :] if current.startswith(previous) else "" + + +@dataclass(slots=True) +class _TurnState: + item_id: str + started: bool = False + start_emitted: bool = False + latest_partial: str | None = None + emitted_partial: str = "" + final_text: str | None = None + completed_emitted: bool = False + stopped: bool = False + stopped_emitted: bool = False + + def finish(self, transcript: str) -> None: + self.final_text = transcript + self.stopped = True + + def drain( + self, take_usage: Callable[[], RealtimeInputAudioTranscriptionUsage | None] + ) -> Iterator[OpenAIRealtimeEvents]: + has_content: Final = self.latest_partial is not None or self.final_text is not None + if (self.started or has_content) and not self.start_emitted: + self.start_emitted = True + yield _speech_event("input_audio_buffer.speech_started", self.item_id) + if self.latest_partial is not None and self.final_text is None: + delta: Final = _new_suffix(self.emitted_partial, self.latest_partial) + if delta: + self.emitted_partial = self.latest_partial + yield _delta_event(self.item_id, delta) + if self.stopped and not self.stopped_emitted: + self.stopped_emitted = True + yield _speech_event("input_audio_buffer.speech_stopped", self.item_id) + if self.final_text is not None and self.stopped_emitted and not self.completed_emitted: + self.completed_emitted = True + yield _completed_event(self.item_id, self.final_text, take_usage()) + + +class MuseEventTransformer: + def __init__(self, *, turn_limit: int = 128) -> None: + self._turns: dict[str, _TurnState] = {} # mutable-ok: bounded, insertion-ordered per-turn emit state + self._turn_limit: Final = turn_limit + self._active_turn_id: str | None = None + self._mode: MuseMode = "ENDPOINTING" + self._last_audio_processed_ms: float = 0.0 + self._unbilled_seconds: float = 0.0 + + def configure(self, config: MuseSessionConfig) -> None: + self._mode = config.mode + + def transform(self, message: Mapping[str, JsonValue]) -> tuple[OpenAIRealtimeEvents, ...]: + event_type: Final = message.get("type") + if event_type == "error": + return (error_event(_PROVIDER_ERROR_MESSAGE),) + if event_type == "audioProgress": + self._update_audio_progress(message) + return () + turn: Final = self._apply_turn_event(event_type, message) + if turn is None: + return () + return tuple(turn.drain(self.take_unbilled_usage)) + + def take_unbilled_usage(self) -> RealtimeInputAudioTranscriptionUsage | None: + seconds: Final = self._unbilled_seconds + if seconds <= 0: + return None + self._unbilled_seconds = 0.0 + usage: Final[RealtimeInputAudioTranscriptionDurationUsage] = {"type": "duration", "seconds": seconds} + return usage + + def _apply_turn_event(self, event_type: JsonValue | None, message: Mapping[str, JsonValue]) -> _TurnState | None: + match event_type: + case "speechStart": + return self._speech_start(message) + case "transcript": + return self._transcript(message) + case "speechEnd": + return self._speech_end(message) + case "speechComplete": + return self._speech_complete(message) + case _: + return None + + def _turn(self, turn_id: str) -> _TurnState: + existing: Final = self._turns.get(turn_id) + if existing is not None: + return existing + created: Final = _TurnState(item_id=turn_id) + self._turns[turn_id] = created + if len(self._turns) > self._turn_limit: + del self._turns[next(iter(self._turns))] + return created + + def _speech_start(self, message: Mapping[str, JsonValue]) -> _TurnState: + turn: Final = self._turn(_required_turn_id(message, "speechStart")) + if turn.stopped: + return turn + turn.started = True + self._active_turn_id = turn.item_id + return turn + + def _transcript(self, message: Mapping[str, JsonValue]) -> _TurnState | None: + transcript: Final = message.get("transcript") + if not isinstance(transcript, str): + raise MuseProtocolError("transcript event has invalid transcript") + if not transcript and message.get("turnId") is None and self._active_turn_id is None: + return None + turn: Final = self._turn(self._transcript_turn_id(message)) + if message.get("final") is True: + self._finish(turn, transcript) + elif turn.final_text is None: + turn.latest_partial = transcript + return turn + + def _speech_end(self, message: Mapping[str, JsonValue]) -> _TurnState: + turn: Final = self._turn(_required_turn_id(message, "speechEnd")) + turn.stopped = True + return turn + + def _speech_complete(self, message: Mapping[str, JsonValue]) -> _TurnState: + transcript: Final = message.get("transcript") + if not isinstance(transcript, str): + raise MuseProtocolError("speechComplete event has invalid transcript") + turn: Final = self._turn(_required_turn_id(message, "speechComplete")) + self._finish(turn, transcript) + return turn + + def _finish(self, turn: _TurnState, transcript: str) -> None: + turn.finish(transcript) + self._release_active(turn) + + def _release_active(self, turn: _TurnState) -> None: + if self._active_turn_id == turn.item_id: + self._active_turn_id = None + + def _update_audio_progress(self, message: Mapping[str, JsonValue]) -> None: + processed_ms: Final = message.get("audioProcessedMs") + if ( + isinstance(processed_ms, bool) + or not isinstance(processed_ms, (int, float)) + or not math.isfinite(processed_ms) + or processed_ms < 0 + ): + raise MuseProtocolError("audioProgress event has invalid audioProcessedMs") + if processed_ms <= self._last_audio_processed_ms: + return + self._unbilled_seconds += (float(processed_ms) - self._last_audio_processed_ms) / 1000 + self._last_audio_processed_ms = float(processed_ms) + + def _transcript_turn_id(self, message: Mapping[str, JsonValue]) -> str: + if message.get("turnId") is not None: + return _required_turn_id(message, "transcript") + if self._active_turn_id is not None: + return self._active_turn_id + if self._mode != "PUSH_TO_TALK": + raise MuseProtocolError("transcript event is missing turnId outside an active turn") + turn_id: Final = f"item_{uuid.uuid4().hex}" + self._active_turn_id = turn_id + return turn_id + + +class MetaRealtimeConfig(BaseRealtimeConfig): + def __init__( + self, + *, + monotonic: Callable[[], float] = time.monotonic, + sleep: Callable[[float], Awaitable[None]] = asyncio.sleep, + ) -> None: + self._monotonic: Final = monotonic + self._sleep: Final = sleep + self._transformer: Final = MuseEventTransformer() + self._access_token: str | None = None + self._config: MuseSessionConfig | None = None + self._pending_audio: bytes = b"" + self._end_stream_sent: bool = False + self._pacing_origin: float | None = None + self._sent_duration: float = 0.0 + + def validate_environment( + self, + headers: dict[str, str], # mutable-ok: BaseRealtimeConfig contract + model: str, + api_key: str | None = None, + ) -> dict[str, str]: # mutable-ok: BaseRealtimeConfig contract + token: Final = api_key or get_secret_str("META_API_KEY") + if token is None: + raise ValueError("api_key is required for Meta API calls") + self._access_token = normalize_access_token(token) + return headers + + def get_complete_url(self, api_base: str | None, model: str, api_key: str | None = None) -> str: + if _normalize_model(model) != MUSE_MODEL: + raise ValueError(f"Unsupported Meta realtime model: {model}") + return build_muse_realtime_url(api_base) + + def is_setup_message(self, msg_obj: Mapping[str, object]) -> bool: + return "authorization" in msg_obj + + def transform_session_created_event( + self, + model: str, + logging_session_id: str, + session_configuration_request: str | None = None, + ) -> OpenAIRealtimeTranscriptionSessionCreated: + return session_created_event(_DEFAULT_SESSION_CONFIG, logging_session_id) + + def transform_realtime_request( + self, + message: str, + model: str, + session_configuration_request: str | None = None, + ) -> tuple[str | bytes, ...]: + request: Final = _json_object(message) + event_type: Final = request.get("type") + if event_type in ("session.update", "transcription_session.update"): + return self._configure(message, model) + if event_type == "input_audio_buffer.append": + return self._append_audio(request) + if event_type == "input_audio_buffer.commit": + return self._flush_audio(end_stream=self._require_config().mode == "PUSH_TO_TALK") + if event_type == "input_audio_buffer.end": + return self._flush_audio(end_stream=True) + if event_type == "input_audio_buffer.clear": + self._pending_audio = b"" + return () + verbose_logger.debug("Meta realtime: dropping unsupported client event %s", event_type) + return () + + async def pace_backend_send(self, message: bytes) -> None: + now: Final = self._monotonic() + origin: Final = self._pacing_origin + effective_origin: Final = ( + now - self._sent_duration if origin is None or now > origin + self._sent_duration else origin + ) + delay: Final = effective_origin + self._sent_duration - now + if delay > 0: + await self._sleep(delay) + self._pacing_origin = effective_origin + self._sent_duration += len(message) / self._require_config().bytes_per_second + + def unbilled_usage_on_session_close(self, model: str) -> RealtimeInputAudioTranscriptionUsage | None: + return self._transformer.take_unbilled_usage() + + def transform_realtime_response( + self, + message: str | bytes, + model: str, + logging_obj: LiteLLMLoggingObj, + realtime_response_transform_input: RealtimeResponseTransformInput, + ) -> RealtimeResponseTypedDict: + payload: Final = message.decode("utf-8") if isinstance(message, bytes) else message + result: Final[RealtimeResponseTypedDict] = { + "response": list(self._backend_events(payload)), # mutable-ok: RealtimeResponseTypedDict.response is a list + "current_output_item_id": realtime_response_transform_input.get("current_output_item_id"), + "current_response_id": realtime_response_transform_input.get("current_response_id"), + "current_delta_chunks": realtime_response_transform_input.get("current_delta_chunks"), + "current_conversation_id": realtime_response_transform_input.get("current_conversation_id"), + "current_item_chunks": realtime_response_transform_input.get("current_item_chunks"), + "current_delta_type": realtime_response_transform_input.get("current_delta_type"), + "session_configuration_request": realtime_response_transform_input.get("session_configuration_request"), + } + return result + + def _backend_events(self, payload: str) -> tuple[OpenAIRealtimeEvents, ...]: + frame: Final = _json_object(payload) + session_id: Final = frame.get("sessionId") + if session_id is None: + return self._transformer.transform(frame) + if not isinstance(session_id, str) or not session_id.strip(): + raise MuseProtocolError("provider returned an invalid handshake response") + return (session_created_event(self._require_config(), session_id.strip()),) + + def _configure(self, message: str, model: str) -> tuple[str, ...]: + if self._config is not None: + verbose_logger.debug("Meta realtime: ignoring session.update after the Muse handshake was sent") + return () + access_token: Final = self._access_token + if access_token is None: + raise MuseProtocolError("Meta API key was not validated before the session was configured") + config: Final = parse_session_update(message, model) + self._config = config + self._transformer.configure(config) + return (json.dumps(config.handshake(access_token), separators=(",", ":")),) + + def _append_audio(self, request: Mapping[str, JsonValue]) -> tuple[bytes, ...]: + config: Final = self._require_config() + encoded: Final = request.get("audio") + if not isinstance(encoded, str): + raise MuseProtocolError("Audio must be a base64 string") + if len(encoded) > config.max_encoded_append_bytes: + raise MuseProtocolError("Audio append exceeds the four-second backlog limit") + try: + audio: Final = base64.b64decode(encoded, validate=True) + except (binascii.Error, ValueError): + raise MuseProtocolError("Audio must be valid base64") from None + if len(audio) % 2: + raise MuseProtocolError("PCM16 audio must contain complete samples") + buffered: Final = self._pending_audio + audio + packet_end: Final = len(buffered) - len(buffered) % config.packet_bytes + self._pending_audio = buffered[packet_end:] + return tuple( + buffered[start : start + config.packet_bytes] for start in range(0, packet_end, config.packet_bytes) + ) + + def _flush_audio(self, *, end_stream: bool) -> tuple[str | bytes, ...]: + remainder: Final = self._pending_audio + self._pending_audio = b"" + frames: Final[tuple[bytes, ...]] = (remainder,) if remainder else () + if not end_stream or self._end_stream_sent: + return frames + self._end_stream_sent = True + return (*frames, _END_STREAM) + + def _require_config(self) -> MuseSessionConfig: + if self._config is None: + raise MuseProtocolError("session.update must configure the Muse session before audio is sent") + return self._config diff --git a/litellm/llms/openai_like/providers.json b/litellm/llms/openai_like/providers.json index a458a209ea9..fe10293c420 100644 --- a/litellm/llms/openai_like/providers.json +++ b/litellm/llms/openai_like/providers.json @@ -173,7 +173,7 @@ "api_key_env": "META_API_KEY", "api_base_env": "META_API_BASE", "base_class": "openai_gpt", - "supported_endpoints": ["/v1/chat/completions", "/v1/responses", "/v1/messages"] + "supported_endpoints": ["/v1/chat/completions", "/v1/responses", "/v1/messages", "/v1/realtime"] }, "cognition": { "base_url": "https://api.cognition.ai/v1", diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 0d26f805565..de9867f9eee 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -4497,7 +4497,7 @@ }, "azure/eu/o1-2024-12-17": { "cache_read_input_token_cost": 8.25e-06, - "deprecation_date": "2026-10-21", + "deprecation_date": "2026-11-19", "input_cost_per_token": 1.65e-05, "litellm_provider": "azure", "max_input_tokens": 200000, @@ -4543,7 +4543,7 @@ }, "azure/eu/o3-mini-2025-01-31": { "cache_read_input_token_cost": 6.05e-07, - "deprecation_date": "2026-10-01", + "deprecation_date": "2026-11-19", "input_cost_per_token": 1.21e-06, "input_cost_per_token_batches": 6.05e-07, "litellm_provider": "azure", @@ -8730,7 +8730,7 @@ "supports_function_calling": true }, "azure/o1": { - "deprecation_date": "2026-10-21", + "deprecation_date": "2026-11-19", "cache_read_input_token_cost": 7.5e-06, "input_cost_per_token": 1.5e-05, "litellm_provider": "azure", @@ -8748,7 +8748,7 @@ }, "azure/o1-2024-12-17": { "cache_read_input_token_cost": 7.5e-06, - "deprecation_date": "2026-10-21", + "deprecation_date": "2026-11-19", "input_cost_per_token": 1.5e-05, "litellm_provider": "azure", "max_input_tokens": 200000, @@ -8825,7 +8825,7 @@ "supports_vision": false }, "azure/o3": { - "deprecation_date": "2026-10-21", + "deprecation_date": "2026-11-19", "cache_read_input_token_cost": 5e-07, "input_cost_per_token": 2e-06, "litellm_provider": "azure", @@ -8855,7 +8855,7 @@ "supports_vision": true }, "azure/o3-2025-04-16": { - "deprecation_date": "2026-10-21", + "deprecation_date": "2026-11-19", "cache_read_input_token_cost": 5e-07, "input_cost_per_token": 2e-06, "litellm_provider": "azure", @@ -8886,7 +8886,7 @@ }, "azure/o3-deep-research": { "cache_read_input_token_cost": 2.5e-06, - "deprecation_date": "2026-12-26", + "deprecation_date": "2026-11-19", "input_cost_per_token": 1e-05, "litellm_provider": "azure", "max_input_tokens": 200000, @@ -8923,7 +8923,7 @@ "supports_web_search": true }, "azure/o3-mini": { - "deprecation_date": "2026-10-01", + "deprecation_date": "2026-11-19", "cache_read_input_token_cost": 5.5e-07, "input_cost_per_token": 1.1e-06, "litellm_provider": "azure", @@ -8940,7 +8940,7 @@ }, "azure/o3-mini-2025-01-31": { "cache_read_input_token_cost": 5.5e-07, - "deprecation_date": "2026-10-01", + "deprecation_date": "2026-11-19", "input_cost_per_token": 1.1e-06, "litellm_provider": "azure", "max_input_tokens": 200000, @@ -8954,7 +8954,7 @@ "supports_vision": false }, "azure/o3-pro": { - "deprecation_date": "2026-12-17", + "deprecation_date": "2026-11-19", "input_cost_per_token": 2e-05, "input_cost_per_token_batches": 1e-05, "litellm_provider": "azure", @@ -8985,7 +8985,7 @@ "supports_vision": true }, "azure/o3-pro-2025-06-10": { - "deprecation_date": "2026-12-17", + "deprecation_date": "2026-11-19", "input_cost_per_token": 2e-05, "input_cost_per_token_batches": 1e-05, "litellm_provider": "azure", @@ -9016,7 +9016,7 @@ "supports_vision": true }, "azure/o4-mini": { - "deprecation_date": "2026-10-16", + "deprecation_date": "2026-11-19", "cache_read_input_token_cost": 2.75e-07, "input_cost_per_token": 1.1e-06, "litellm_provider": "azure", @@ -9047,7 +9047,7 @@ }, "azure/o4-mini-2025-04-16": { "cache_read_input_token_cost": 2.75e-07, - "deprecation_date": "2026-10-16", + "deprecation_date": "2026-11-19", "input_cost_per_token": 1.1e-06, "litellm_provider": "azure", "max_input_tokens": 200000, @@ -9600,7 +9600,7 @@ }, "azure/us/o1-2024-12-17": { "cache_read_input_token_cost": 8.25e-06, - "deprecation_date": "2026-10-21", + "deprecation_date": "2026-11-19", "input_cost_per_token": 1.65e-05, "litellm_provider": "azure", "max_input_tokens": 200000, @@ -9645,7 +9645,7 @@ "supports_vision": false }, "azure/us/o3-2025-04-16": { - "deprecation_date": "2026-10-21", + "deprecation_date": "2026-11-19", "cache_read_input_token_cost": 5.5e-07, "input_cost_per_token": 2.2e-06, "litellm_provider": "azure", @@ -9676,7 +9676,7 @@ }, "azure/us/o3-mini-2025-01-31": { "cache_read_input_token_cost": 6.05e-07, - "deprecation_date": "2026-10-01", + "deprecation_date": "2026-11-19", "input_cost_per_token": 1.21e-06, "input_cost_per_token_batches": 6.05e-07, "litellm_provider": "azure", @@ -9693,7 +9693,7 @@ }, "azure/us/o4-mini-2025-04-16": { "cache_read_input_token_cost": 3.1e-07, - "deprecation_date": "2026-10-16", + "deprecation_date": "2026-11-19", "input_cost_per_token": 1.21e-06, "litellm_provider": "azure", "max_input_tokens": 200000, @@ -14615,7 +14615,7 @@ }, "computer-use-preview": { "input_cost_per_token": 3e-06, - "litellm_provider": "azure", + "litellm_provider": "openai", "max_input_tokens": 8192, "max_output_tokens": 1024, "max_tokens": 1024, @@ -14633,12 +14633,14 @@ ], "supports_function_calling": true, "supports_parallel_function_calling": true, + "supports_pdf_input": true, "supports_prompt_caching": false, "supports_reasoning": true, - "supports_response_schema": true, + "supports_response_schema": false, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "source": "https://platform.openai.com/docs/models/computer-use-preview" }, "dall-e-2": { "deprecation_date": "2026-05-12", @@ -34720,6 +34722,22 @@ "supports_web_search": true, "supports_xhigh_reasoning_effort": true }, + "meta/muse-voice-transcribe-1.0": { + "input_cost_per_second": 0.00005, + "litellm_provider": "meta", + "mode": "audio_transcription", + "source": "https://dev.meta.ai/docs/speech-to-text", + "supported_endpoints": [ + "/v1/realtime" + ], + "supported_modalities": [ + "audio" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_input": true + }, "meta_llama/Llama-3.3-70B-Instruct": { "litellm_provider": "meta_llama", "max_input_tokens": 128000, @@ -39022,7 +39040,9 @@ "supports_max_reasoning_effort": true, "supports_tool_choice": true, "supports_vision": true, - "prompt_cache_min_tokens": 4096 + "prompt_cache_min_tokens": 4096, + "supports_response_schema": true, + "source": "https://openrouter.ai/api/v1/models" }, "openrouter/anthropic/claude-sonnet-4.5": { "input_cost_per_image": 0.0048, @@ -39170,18 +39190,20 @@ }, "openrouter/deepseek/deepseek-v3.2": { "input_cost_per_token": 2.69e-07, - "input_cost_per_token_cache_hit": 2.8e-08, + "input_cost_per_token_cache_hit": 1.345e-07, "litellm_provider": "openrouter", "max_input_tokens": 163840, - "max_output_tokens": 163840, - "max_tokens": 163840, + "max_output_tokens": 65536, + "max_tokens": 65536, "mode": "chat", "output_cost_per_token": 4e-07, "supports_assistant_prefill": true, "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_response_schema": true, + "source": "https://openrouter.ai/api/v1/models" }, "openrouter/deepseek/deepseek-v3.2-exp": { "input_cost_per_token": 2.7e-07, @@ -43491,6 +43513,7 @@ "supports_tool_choice": true }, "together_ai/openai/gpt-oss-20b": { + "deprecation_date": "2026-09-14", "input_cost_per_token": 5e-08, "litellm_provider": "together_ai", "max_input_tokens": 131072, @@ -43783,6 +43806,7 @@ "source": "https://docs.together.ai/docs/serverless-models" }, "together_ai/google/gemma-4-31B-it": { + "deprecation_date": "2026-09-14", "input_cost_per_token": 3.9e-07, "litellm_provider": "together_ai", "max_input_tokens": 262144, @@ -43797,6 +43821,7 @@ "supports_vision": true }, "together_ai/intfloat/multilingual-e5-large-instruct": { + "deprecation_date": "2026-09-14", "input_cost_per_token": 2e-08, "litellm_provider": "together_ai", "max_input_tokens": 514, @@ -43909,6 +43934,7 @@ "supports_tool_choice": true }, "together_ai/thinkingmachines/Inkling-Small": { + "deprecation_date": "2026-09-14", "cache_read_input_token_cost": 1e-07, "input_cost_per_token": 5e-07, "litellm_provider": "together_ai", @@ -59101,9 +59127,9 @@ "litellm_provider": "wandb", "mode": "chat", "supports_reasoning": true, - "input_cost_per_token": 0.00000131, - "output_cost_per_token": 0.00000396, - "cache_read_input_token_cost": 0.000000044, + "input_cost_per_token": 1.31e-06, + "output_cost_per_token": 3.96e-06, + "cache_read_input_token_cost": 4.4e-08, "supports_prompt_caching": true, "source": "https://wandb.ai/site/pricing/tokens/" }, @@ -59111,9 +59137,9 @@ "litellm_provider": "wandb", "mode": "chat", "supports_reasoning": true, - "input_cost_per_token": 0.0000001, - "output_cost_per_token": 0.00000015, - "cache_read_input_token_cost": 0.00000005, + "input_cost_per_token": 1e-07, + "output_cost_per_token": 1.5e-07, + "cache_read_input_token_cost": 5e-08, "supports_prompt_caching": true, "source": "https://wandb.ai/site/pricing/tokens/" }, @@ -60835,6 +60861,7 @@ "source": "https://docs.together.ai/docs/serverless-models" }, "together_ai/moonshotai/Kimi-K2.6": { + "deprecation_date": "2026-08-19", "input_cost_per_token": 1.2e-06, "output_cost_per_token": 4.5e-06, "cache_read_input_token_cost": 2e-07, @@ -60861,6 +60888,7 @@ "source": "https://api.together.xyz/v1/models" }, "together_ai/zai-org/GLM-5": { + "deprecation_date": "2026-06-22", "input_cost_per_token": 1e-06, "output_cost_per_token": 3.2e-06, "litellm_provider": "together_ai", @@ -60869,6 +60897,7 @@ "source": "https://api.together.xyz/v1/models" }, "together_ai/zai-org/GLM-5.1": { + "deprecation_date": "2026-07-10", "input_cost_per_token": 1.4e-06, "output_cost_per_token": 4.4e-06, "cache_read_input_token_cost": 2.6e-07, @@ -60886,6 +60915,7 @@ "source": "https://api.together.xyz/v1/models" }, "together_ai/Qwen/Qwen3-Coder-Next-FP8": { + "deprecation_date": "2026-05-14", "input_cost_per_token": 5e-07, "output_cost_per_token": 1.2e-06, "litellm_provider": "together_ai", @@ -60894,6 +60924,7 @@ "source": "https://api.together.xyz/v1/models" }, "together_ai/Qwen/Qwen3-VL-32B-Instruct": { + "deprecation_date": "2026-02-25", "input_cost_per_token": 5e-07, "output_cost_per_token": 1.5e-06, "litellm_provider": "together_ai", @@ -60902,6 +60933,7 @@ "source": "https://api.together.xyz/v1/models" }, "together_ai/Qwen/Qwen3-VL-8B-Instruct": { + "deprecation_date": "2026-04-16", "input_cost_per_token": 1.8e-07, "output_cost_per_token": 6.8e-07, "litellm_provider": "together_ai", @@ -60934,6 +60966,7 @@ "source": "https://api.together.xyz/v1/models" }, "together_ai/Qwen/QwQ-32B": { + "deprecation_date": "2025-11-13", "input_cost_per_token": 1.2e-06, "output_cost_per_token": 1.2e-06, "litellm_provider": "together_ai", diff --git a/litellm/ocr/legacy.py b/litellm/ocr/legacy.py index ddf6016dce3..a742be274b3 100644 --- a/litellm/ocr/legacy.py +++ b/litellm/ocr/legacy.py @@ -28,6 +28,7 @@ from litellm.llms.base_llm.ocr.transformation import ( from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler from litellm.ocr.input import FileReader from litellm.types.router import GenericLiteLLMParams +from litellm.types.utils import CustomPricingLiteLLMParams from litellm.utils import ProviderConfigManager, client base_llm_http_handler: Final = BaseLLMHTTPHandler() @@ -149,6 +150,7 @@ def _prepare_ocr_request( litellm_params={ "litellm_call_id": litellm_call_id, "api_base": resolved_api_base, + **litellm_params.model_dump(include=frozenset(CustomPricingLiteLLMParams.model_fields), exclude_none=True), }, custom_llm_provider=custom_llm_provider, ) diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index cb7a18cd107..7a110eff080 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -3234,6 +3234,17 @@ } ], "title": "User Email" + }, + "user_id": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "User Id" } }, "title": "KeyMetadata", @@ -5134,6 +5145,49 @@ "anthropic_skills" ] } + }, + "/v1/skills/{skill_id}/archive": { + "get": { + "description": "Stored skill upload, repacked so SKILL.md sits at the archive root.", + "operationId": "agent_skills_archive_v1_skills__skill_id__archive_get", + "parameters": [ + { + "in": "path", + "name": "skill_id", + "required": true, + "schema": { + "title": "Skill Id", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/zip": { + "schema": { + "type": "string" + } + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "summary": "Agent Skills Archive", + "tags": [ + "anthropic_skills" + ] + } } } }, diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 289fe086379..e6ed60ba177 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -3,7 +3,7 @@ import contextlib import json import logging import math -from collections.abc import AsyncGenerator, Awaitable, Callable, Coroutine, Mapping, Sequence +from collections.abc import AsyncGenerator, Awaitable, Callable, Mapping, Sequence from datetime import datetime from functools import lru_cache from types import MappingProxyType @@ -3210,20 +3210,24 @@ class ProxyBaseLLMRequestProcessing: end-of-stream blocks complete, so the spend log sees guardrail_information. - Three closure shapes, matching who owns logging for the stream: + Two closure shapes, matching who owns logging for the stream: - CustomStreamWrapper (chat completions) stores (assembled_response, cache_hit); the closure also runs non-apply_guardrail post-call hooks via _run_deferred_stream_guardrails. - - Bridged /v1/responses (LiteLLMCompletionStreamingIterator) shares - its inner CustomStreamWrapper's logging_obj, so it stores the same - (assembled_response, cache_hit) shape; the closure only dispatches - success logging, matching the route's pre-existing hook surface. - - Native anthropic_messages/aresponses iterators store a single - ready-made logging coroutine to enqueue. + - Every other anthropic_messages/aresponses stream gets a closure + that dispatches on the stored args shape, because the arming site + cannot tell the producers apart: native iterators store a single + ready-made logging coroutine to enqueue, while bridged streams + (LiteLLMCompletionStreamingIterator, and the plain SSE generator + AnthropicStreamWrapper returns for bridged /v1/messages) share + their inner CustomStreamWrapper's logging_obj and so store + (assembled_response, cache_hit); for those the closure only + dispatches success logging, matching the route's pre-existing + hook surface. - Raw async generators from passthrough routes bypass all three and - would orphan the closure, so they are not armed here. + Raw async generators from passthrough routes bypass both and would + orphan the closure, so they are not armed here. The router wraps iterators that cannot carry _hidden_params in HiddenParamsAsyncIteratorWrapper, so class sniffing runs on the @@ -3257,31 +3261,27 @@ class ProxyBaseLLMRequestProcessing: if route_type not in ("anthropic_messages", "aresponses") or not self._is_streaming_response(response): return - from litellm.responses.litellm_completion_transformation.streaming_iterator import ( - LiteLLMCompletionStreamingIterator, - ) - - if isinstance(unwrapped, LiteLLMCompletionStreamingIterator): - _captured_bridge_logging_obj: Final = logging_obj - - async def _on_deferred_bridged_stream_complete(assembled_response: object, cache_hit: object) -> None: - await _as_success_dispatcher(_captured_bridge_logging_obj).dispatch_success_handlers( - assembled_response, - cache_hit=cache_hit, - start_time=None, - end_time=None, - prefer_async_handlers=True, - ) - - logging_obj._on_deferred_stream_complete = _on_deferred_bridged_stream_complete - return - from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER - async def _on_deferred_native_stream_complete( - logging_coroutine: Coroutine[object, object, object], - ) -> None: - GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(async_coroutine=logging_coroutine) + _captured_native_logging_obj: Final = logging_obj + + async def _on_deferred_native_stream_complete(*args: object) -> None: + match args: + case (logging_coroutine,) if asyncio.iscoroutine(logging_coroutine): + GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(async_coroutine=logging_coroutine) + case (assembled_response, cache_hit): + await _as_success_dispatcher(_captured_native_logging_obj).dispatch_success_handlers( + assembled_response, + cache_hit=cache_hit, + start_time=None, + end_time=None, + prefer_async_handlers=True, + ) + case _: + verbose_proxy_logger.error( + "Deferred stream logging dropped: unexpected stored args shape %s", + tuple(type(arg).__name__ for arg in args), + ) logging_obj._on_deferred_stream_complete = _on_deferred_native_stream_complete diff --git a/litellm/proxy/common_utils/config_includes.py b/litellm/proxy/common_utils/config_includes.py new file mode 100644 index 00000000000..c1bb5ae952f --- /dev/null +++ b/litellm/proxy/common_utils/config_includes.py @@ -0,0 +1,132 @@ +import os +from collections.abc import Awaitable, Mapping +from types import MappingProxyType +from typing import Final, Protocol + +from litellm._logging import verbose_proxy_logger + +INCLUDE_KEY: Final = "include" + + +def resolve_include_file_path(include_file: str, declared_in: str, root_config_path: str) -> str: + """ + Resolve one `include` entry to the file it names, next to the config that declares it. + + A config written before nested entries resolved this way can name a file sitting next to the root + config instead, so that file is still read, with a warning naming where it was found. When both + files exist the one next to the declaring config wins and the other is named in a warning. + """ + declared_relative: Final = os.path.abspath(os.path.join(os.path.dirname(declared_in), include_file)) + root_relative: Final = os.path.abspath(os.path.join(os.path.dirname(root_config_path), include_file)) + if root_relative == declared_relative or not os.path.exists(root_relative): + return declared_relative + + if not os.path.exists(declared_relative): + verbose_proxy_logger.warning( + "Config include '%s' declared in %s was not found next to it, so %s was read instead. " + "Move the included file next to the config that declares it.", + include_file, + declared_in, + root_relative, + ) + return root_relative + + verbose_proxy_logger.warning( + "Config include '%s' declared in %s matches two files. %s sits next to that config and was read, " + "so %s was skipped. Rename one of the two to say which one you meant.", + include_file, + declared_in, + declared_relative, + root_relative, + ) + return declared_relative + + +class IncludeResolver(Protocol): + def __call__(self, include_entry: str, declared_in: str, /) -> str: ... + + +class ConfigReader(Protocol): + def __call__(self, location: str, /) -> Awaitable[Mapping[str, object]]: ... + + +def _merged_value(base_value: object, included_value: object) -> object: + if isinstance(included_value, list) and isinstance(base_value, list): + return [*base_value, *included_value] # mutable-ok: a merged config value stays the plain list the proxy loads + return included_value + + +def _merged_entry(base: Mapping[str, object], included: Mapping[str, object], key: str) -> object: + if key not in included: + return base[key] + return _merged_value(base.get(key), included[key]) + + +def _merged(base: Mapping[str, object], included: Mapping[str, object]) -> Mapping[str, object]: + return MappingProxyType({key: _merged_entry(base, included, key) for key in (*base, *included)}) + + +def _without_include(config: Mapping[str, object]) -> Mapping[str, object]: + return MappingProxyType({key: value for key, value in config.items() if key != INCLUDE_KEY}) + + +def include_entries(config: Mapping[str, object]) -> tuple[str, ...]: + if INCLUDE_KEY not in config: + return () + + entries: Final = config[INCLUDE_KEY] + if not isinstance(entries, list): + raise ValueError("'include' must be a list of file paths") + + paths: Final = tuple(entry for entry in entries if isinstance(entry, str)) + if len(paths) != len(entries): + raise ValueError("'include' must be a list of file paths") + + return paths + + +def _pending_from(config: Mapping[str, object], location: str) -> tuple[tuple[str, str], ...]: + return tuple((entry, location) for entry in include_entries(config)) + + +async def _resolve( + config: Mapping[str, object], + pending: tuple[tuple[str, str], ...], + loaded: frozenset[str], + resolve: IncludeResolver, + read: ConfigReader, +) -> Mapping[str, object]: + if not pending: + return _without_include(config) + + entry, declared_in = pending[0] + location: Final = resolve(entry, declared_in) + if location in loaded: + return await _resolve(config, pending[1:], loaded, resolve, read) + + included: Final = await read(location) + return await _resolve( + _merged(config, _without_include(included)), + (*pending[1:], *_pending_from(included, location)), + loaded | frozenset((location,)), + resolve, + read, + ) + + +async def resolve_includes( + config: Mapping[str, object], + location: str, + resolve: IncludeResolver, + read: ConfigReader, +) -> dict[str, object]: + """ + Merge every config named by the `include` directive into the config that declares it. + + List values are extended and every other value is overridden, `resolve` turns each entry into the + location it names relative to the config that declares it, a config already pulled in is neither + read nor merged a second time, and `read` decides where a location is read from, so the same merge + applies to configs on disk and to configs hosted in a bucket. + """ + merged: Final = await _resolve(config, _pending_from(config, location), frozenset((location,)), resolve, read) + return dict(merged) # mutable-ok: the proxy mutates the config it loads diff --git a/litellm/proxy/common_utils/load_config_utils.py b/litellm/proxy/common_utils/load_config_utils.py index 62649ad6ca1..4a082eb307b 100644 --- a/litellm/proxy/common_utils/load_config_utils.py +++ b/litellm/proxy/common_utils/load_config_utils.py @@ -1,12 +1,47 @@ +import asyncio import os -from typing import Final +import posixpath +from collections.abc import Awaitable, Mapping +from types import MappingProxyType +from typing import TYPE_CHECKING, Final, Protocol import yaml +from pydantic import TypeAdapter, ValidationError from litellm._logging import verbose_proxy_logger +from litellm.proxy.common_utils.config_includes import resolve_includes + +if TYPE_CHECKING: + from litellm.integrations.gcs_bucket.gcs_bucket_base import GCSBucketBase + +_BUCKET_CONFIG_ADAPTER: Final = TypeAdapter(dict[str, object]) -def get_file_contents_from_s3(bucket_name, object_key): +class BucketObjectFetcher(Protocol): + def __call__(self, object_key: str, /) -> Awaitable[Mapping[str, object] | None]: ... + + +class BucketObjectReader(Protocol): + def __call__(self, object_key: str, /) -> Awaitable[object | None]: ... + + +class SyncBucketObjectReader(Protocol): + def __call__(self, object_key: str, /) -> object | None: ... + + +def _parsed_config(object_key: str, file_contents: str) -> object | None: + try: + parsed: Final = yaml.safe_load(file_contents) + except yaml.YAMLError as e: + verbose_proxy_logger.error("Config object %s is not valid YAML: %s", object_key, e) + return None + return MappingProxyType({}) if parsed is None else parsed + + +def s3_object_reader(bucket_name: str) -> SyncBucketObjectReader: + """ + Build one reader for a whole config, so an `include` tree costs one S3 client rather than one per object. + """ try: # v0 rely on boto3 for authentication - allowing boto3 to handle IAM credentials etc import boto3 @@ -21,46 +56,147 @@ def get_file_contents_from_s3(bucket_name, object_key): aws_secret_access_key=credentials.secret_key, aws_session_token=credentials.token, # Optional, if using temporary credentials ) - verbose_proxy_logger.debug("Retrieving %s from S3 bucket: %s", object_key, bucket_name) - response: Final = s3_client.get_object(Bucket=bucket_name, Key=object_key) - verbose_proxy_logger.debug("Response: %s", response) - - # Read the file contents and directly parse YAML - file_contents: Final = response["Body"].read().decode("utf-8") - verbose_proxy_logger.debug("File contents retrieved from S3") - - # Parse YAML directly from string - config: Final = yaml.safe_load(file_contents) - return config - except ImportError as e: # this is most likely if a user is not using the litellm docker container verbose_proxy_logger.error("ImportError: %s", e) + return lambda object_key: None except Exception as e: - verbose_proxy_logger.error("Error retrieving file contents: %s", e) + verbose_proxy_logger.error("Error creating the S3 client for bucket %s: %s", bucket_name, e) + return lambda object_key: None + + def read(object_key: str) -> object | None: + try: + verbose_proxy_logger.debug("Retrieving %s from S3 bucket: %s", object_key, bucket_name) + response: Final = s3_client.get_object(Bucket=bucket_name, Key=object_key) + file_contents: Final = response["Body"].read().decode("utf-8") + except Exception as e: # noqa: BLE001 # any boto3 error must read as a missing object + verbose_proxy_logger.error("Error retrieving %s from S3 bucket %s: %s", object_key, bucket_name, e) + return None + + return _parsed_config(object_key, file_contents) + + return read + + +def get_file_contents_from_s3(bucket_name: str, object_key: str) -> object | None: + return s3_object_reader(bucket_name)(object_key) + + +def gcs_config_bucket(bucket_name: str) -> "GCSBucketBase | None": + """ + Build a plain GCS client for reading config objects. + + Reading a config out of a bucket is not GCS logging, so it neither needs the enterprise license + that gate covers nor the batching task the logger starts and never stops. + """ + try: + from litellm.integrations.gcs_bucket.gcs_bucket_base import GCSBucketBase + + return GCSBucketBase(bucket_name=bucket_name) + except Exception as e: # noqa: BLE001 # an unbuildable client must read as an unreadable bucket + verbose_proxy_logger.error("Error creating the GCS client for bucket %s: %s", bucket_name, e) return None -async def get_config_file_contents_from_gcs(bucket_name, object_key): +async def get_config_file_contents_from_gcs( + bucket_name: str, + object_key: str, + gcs_bucket: "GCSBucketBase | None" = None, +) -> object | None: try: - from litellm.integrations.gcs_bucket.gcs_bucket import GCSBucketLogger - - gcs_bucket: Final = GCSBucketLogger( - bucket_name=bucket_name, - ) - file_contents = await gcs_bucket.download_gcs_object(object_key) + bucket: Final = gcs_config_bucket(bucket_name) if gcs_bucket is None else gcs_bucket + if bucket is None: + return None + file_contents: Final = await bucket.download_gcs_object(object_key) if file_contents is None: raise Exception(f"File contents are None for {object_key}") - # file_contentis is a bytes object, so we need to convert it to yaml - file_contents = file_contents.decode("utf-8") - # convert to yaml - config: Final = yaml.safe_load(file_contents) - return config + decoded: Final = file_contents.decode("utf-8") except Exception as e: - verbose_proxy_logger.error("Error retrieving file contents: %s", e) + verbose_proxy_logger.error("Error retrieving %s from GCS bucket %s: %s", object_key, bucket_name, e) return None + return _parsed_config(object_key, decoded) + + +def resolve_include_object_key(config_object_key: str, include_entry: str) -> str: + """ + Resolve one `include` entry to the object key it names, relative to the config object's prefix. + + A leading "/" means the bucket root, mirroring how an absolute path on disk ignores the + directory the including config sits in. + """ + if include_entry.startswith("/"): + return posixpath.normpath(include_entry).lstrip("/") + return posixpath.normpath(posixpath.join(posixpath.dirname(config_object_key), include_entry)) + + +async def resolve_bucket_includes( + *, + config: Mapping[str, object], + object_key: str, + fetch: BucketObjectFetcher, +) -> dict[str, object]: + async def read(include_key: str) -> Mapping[str, object]: + included: Final = await fetch(include_key) + if included is None: + raise FileNotFoundError( + f"Included config could not be read from bucket: {include_key}. " + "The underlying bucket error is logged above." + ) + return included + + def resolve(include_entry: str, declared_in: str) -> str: + return resolve_include_object_key(declared_in, include_entry) + + return await resolve_includes(config=config, location=object_key, resolve=resolve, read=read) + + +async def bucket_object_reader(bucket_type: str | None, bucket_name: str) -> BucketObjectReader: + """ + Build one reader for a whole config, so an `include` tree costs one bucket client rather than one per object. + """ + if bucket_type != "gcs": + read_object: Final = await asyncio.to_thread(s3_object_reader, bucket_name) + + async def read_from_s3(object_key: str) -> object | None: + return await asyncio.to_thread(read_object, object_key) + + return read_from_s3 + + gcs_bucket: Final = gcs_config_bucket(bucket_name) + + async def read_from_gcs(object_key: str) -> object | None: + if gcs_bucket is None: + return None + return await get_config_file_contents_from_gcs(bucket_name, object_key, gcs_bucket) + + return read_from_gcs + + +async def get_config_from_bucket( + *, + bucket_type: str | None, + bucket_name: str, + object_key: str, +) -> dict[str, object] | None: + read: Final = await bucket_object_reader(bucket_type, bucket_name) + + async def fetch(key: str) -> Mapping[str, object] | None: + raw: Final = await read(key) + if raw is None: + return None + try: + return _BUCKET_CONFIG_ADAPTER.validate_python(raw) + except ValidationError as e: + raise ValueError(f"Config object in bucket is not a YAML mapping: {key}") from e + + config: Final = await fetch(object_key) + if not config: + return None + + return await resolve_bucket_includes(config=config, object_key=object_key, fetch=fetch) + def download_python_file_from_s3( bucket_name: str, @@ -136,11 +272,9 @@ async def download_python_file_from_gcs( bool: True if successful, False otherwise """ try: - from litellm.integrations.gcs_bucket.gcs_bucket import GCSBucketLogger + from litellm.integrations.gcs_bucket.gcs_bucket_base import GCSBucketBase - gcs_bucket: Final = GCSBucketLogger( - bucket_name=bucket_name, - ) + gcs_bucket: Final = GCSBucketBase(bucket_name=bucket_name) file_contents = await gcs_bucket.download_gcs_object(object_key) if file_contents is None: raise Exception(f"File contents are None for {object_key}") diff --git a/litellm/proxy/discovery_endpoints/__init__.py b/litellm/proxy/discovery_endpoints/__init__.py index a6401c2f1b4..52602f30b77 100644 --- a/litellm/proxy/discovery_endpoints/__init__.py +++ b/litellm/proxy/discovery_endpoints/__init__.py @@ -1,3 +1,4 @@ +from .agent_skills_endpoints import router as agent_skills_discovery_router from .ui_discovery_endpoints import router as ui_discovery_endpoints_router -__all__ = ["ui_discovery_endpoints_router"] +__all__ = ["agent_skills_discovery_router", "ui_discovery_endpoints_router"] diff --git a/litellm/proxy/discovery_endpoints/agent_skills_archive.py b/litellm/proxy/discovery_endpoints/agent_skills_archive.py new file mode 100644 index 00000000000..1f2fca3992e --- /dev/null +++ b/litellm/proxy/discovery_endpoints/agent_skills_archive.py @@ -0,0 +1,130 @@ +"""Repack a stored skill upload into the archive shape Agent Skills clients install from. + +Uploads follow the Anthropic Skills API layout, where every file sits under a single +top-level folder. Discovery clients read ``SKILL.md`` from the archive root, so that +folder is stripped and the zip is rebuilt with fixed entry timestamps, which keeps the +SHA-256 digest published in the index reproducible for identical uploads. +""" + +import hashlib +import io +import re +import zipfile +from collections.abc import Mapping +from dataclasses import dataclass +from types import MappingProxyType +from typing import Final + +import yaml + +MAX_ARCHIVE_UNPACKED_BYTES: Final = 50 * 1024 * 1024 +MAX_ARCHIVE_ENTRIES: Final = 1000 +SKILL_MANIFEST_FILENAME: Final = "SKILL.md" + +_ZIP_ENTRY_TIMESTAMP: Final = (1980, 1, 1, 0, 0, 0) +_ZIP_ENTRY_PERMISSIONS: Final = 0o644 << 16 +_FRONTMATTER_PATTERN: Final = re.compile(r"^---\s*\n(.*?)\n---\s*(?:\n|$)", re.DOTALL) +_WINDOWS_DRIVE_PATTERN: Final = re.compile(r"^[A-Za-z]:") +_EMPTY_FRONTMATTER: Final[Mapping[str, object]] = MappingProxyType({}) + + +@dataclass(frozen=True, slots=True) +class SkillArchive: + content: bytes + digest: str + declared_name: str | None + declared_description: str | None + + +def build_skill_archive(stored_content: bytes) -> SkillArchive | None: + """Return the installable archive for an upload, or None when it holds no root SKILL.md.""" + try: + with zipfile.ZipFile(io.BytesIO(stored_content)) as uploaded: + members: Final = _flattened_members(uploaded) + except (zipfile.BadZipFile, OSError, RuntimeError): + return None + + if members is None: + return None + + frontmatter: Final = _manifest_frontmatter(next(data for name, data in members if name == SKILL_MANIFEST_FILENAME)) + content: Final = _repack(members) + return SkillArchive( + content=content, + digest=f"sha256:{hashlib.sha256(content).hexdigest()}", + declared_name=_frontmatter_text(frontmatter, "name"), + declared_description=_frontmatter_text(frontmatter, "description"), + ) + + +def _flattened_members(uploaded: zipfile.ZipFile) -> tuple[tuple[str, bytes], ...] | None: + infos: Final = tuple(info for info in uploaded.infolist() if not info.is_dir()) + if not infos or len(infos) > MAX_ARCHIVE_ENTRIES: + return None + if sum(info.file_size for info in infos) > MAX_ARCHIVE_UNPACKED_BYTES: + return None + + normalized: Final = tuple((info, _normalized_path(info.filename)) for info in infos) + if any(path is None for _, path in normalized): + return None + + prefix: Final = _common_root_prefix(tuple(path for _, path in normalized if path is not None)) + flattened: Final = tuple((info, path[len(prefix) :]) for info, path in normalized if path is not None) + names: Final = frozenset(name for _, name in flattened) + if SKILL_MANIFEST_FILENAME not in names or len(names) != len(flattened): + return None + + return tuple((name, uploaded.read(info)) for info, name in sorted(flattened, key=lambda member: member[1])) + + +def _common_root_prefix(paths: tuple[str, ...]) -> str: + roots: Final = frozenset(path.split("/", 1)[0] for path in paths) + if len(roots) != 1 or not all("/" in path for path in paths): + return "" + return f"{next(iter(roots))}/" + + +def _normalized_path(raw_path: str) -> str | None: + if not raw_path or "\0" in raw_path or "\\" in raw_path: + return None + if raw_path.startswith("/") or _WINDOWS_DRIVE_PATTERN.match(raw_path): + return None + parts: Final = tuple(part for part in raw_path.split("/") if part) + if not parts or any(part in (".", "..") for part in parts): + return None + return "/".join(parts) + + +def _manifest_frontmatter(manifest: bytes) -> Mapping[str, object]: + match: Final = _FRONTMATTER_PATTERN.match(manifest.decode("utf-8", errors="replace")) + if match is None: + return _EMPTY_FRONTMATTER + try: + parsed: Final = yaml.safe_load(match.group(1)) + except yaml.YAMLError: + return _EMPTY_FRONTMATTER + if not isinstance(parsed, dict): + return _EMPTY_FRONTMATTER + return parsed + + +def _frontmatter_text(frontmatter: Mapping[str, object], key: str) -> str | None: + value: Final = frontmatter.get(key) + if not isinstance(value, str): + return None + return value.strip() or None + + +def _zip_entry(name: str) -> zipfile.ZipInfo: + entry: Final = zipfile.ZipInfo(filename=name, date_time=_ZIP_ENTRY_TIMESTAMP) + entry.compress_type = zipfile.ZIP_DEFLATED + entry.external_attr = _ZIP_ENTRY_PERMISSIONS + return entry + + +def _repack(members: tuple[tuple[str, bytes], ...]) -> bytes: + buffer: Final = io.BytesIO() + with zipfile.ZipFile(buffer, mode="w", compression=zipfile.ZIP_DEFLATED) as repacked: + for name, data in members: + repacked.writestr(_zip_entry(name), data) + return buffer.getvalue() diff --git a/litellm/proxy/discovery_endpoints/agent_skills_endpoints.py b/litellm/proxy/discovery_endpoints/agent_skills_endpoints.py new file mode 100644 index 00000000000..3084cbfd84f --- /dev/null +++ b/litellm/proxy/discovery_endpoints/agent_skills_endpoints.py @@ -0,0 +1,200 @@ +"""Serve skills stored on the proxy as an Agent Skills well-known discovery index. + +``npx skills add -a `` reads ``/.well-known/agent-skills/index.json`` +and downloads each entry's archive. Discovery clients send no credentials, so both +routes are unauthenticated and stay off until ``litellm_settings.public_skills_index`` +is enabled, which publishes every stored skill to anyone who can reach the proxy. +""" + +import asyncio +import re +from collections.abc import Sequence +from itertools import groupby +from operator import itemgetter +from types import MappingProxyType +from typing import Final + +from fastapi import APIRouter, Depends, HTTPException, Request, Response + +import litellm +from litellm._logging import verbose_proxy_logger +from litellm.caching.in_memory_cache import InMemoryCache +from litellm.models.skills import LiteLLM_SkillsTable +from litellm.proxy.discovery_endpoints.agent_skills_archive import SkillArchive, build_skill_archive +from litellm.types.proxy.discovery_endpoints.agent_skills_endpoints import ( + MAX_SKILL_DESCRIPTION_LENGTH, + MAX_SKILL_NAME_LENGTH, + AgentSkillsIndex, + AgentSkillsIndexEntry, +) + +MAX_INDEXED_SKILLS: Final = 1000 +MAX_CACHED_ARCHIVES: Final = 128 +MAX_CACHED_ARCHIVE_BYTES: Final = 512 * 1024 +ARCHIVE_CACHE_TTL_SECONDS: Final = 3600 + +_ARCHIVE_CACHE: Final = InMemoryCache( + max_size_in_memory=MAX_CACHED_ARCHIVES, + default_ttl=ARCHIVE_CACHE_TTL_SECONDS, + max_size_per_item=MAX_CACHED_ARCHIVE_BYTES // 1024, +) + +_NON_SLUG_PATTERN: Final = re.compile(r"[^a-z0-9]+") +_FALLBACK_SKILL_NAME: Final = "skill" + +router: Final = APIRouter(tags=["public", "skills"]) # mutable-ok: fastapi types tags as list[str | Enum] + + +class ZipArchiveResponse(Response): + """Response whose OpenAPI entry declares an application/zip download rather than JSON.""" + + media_type = "application/zip" + + +def ensure_index_enabled() -> None: + if litellm.public_skills_index is not True: + raise HTTPException(status_code=404, detail="Not Found") + + +async def stored_skills() -> Sequence[LiteLLM_SkillsTable]: + from litellm.llms.litellm_proxy.skills.handler import LiteLLMSkillsHandler + + return await LiteLLMSkillsHandler.list_skills(limit=MAX_INDEXED_SKILLS) + + +async def stored_skill(skill_id: str) -> LiteLLM_SkillsTable | None: + from litellm.llms.litellm_proxy.skills.handler import LiteLLMSkillsHandler + + try: + return await LiteLLMSkillsHandler.get_skill(skill_id) + except ValueError: + return None + + +@router.get( + "/.well-known/agent-skills/index.json", + response_model=AgentSkillsIndex, + dependencies=(Depends(ensure_index_enabled),), +) +@router.get( + "/.well-known/skills/index.json", + response_model=AgentSkillsIndex, + dependencies=(Depends(ensure_index_enabled),), + include_in_schema=False, +) +async def agent_skills_index( + request: Request, + skills: Sequence[LiteLLM_SkillsTable] = Depends(stored_skills), +) -> AgentSkillsIndex: + """Agent Skills v0.2.0 discovery index over every skill stored on this proxy.""" + from litellm.proxy.utils import get_custom_url + + installable: Final = await _installable(skills) + names: Final = _deduplicated(tuple(_base_name(skill, archive) for skill, archive in installable)) + + return AgentSkillsIndex( + skills=tuple( + AgentSkillsIndexEntry( + name=name, + type="archive", + description=_description(skill, archive, name), + url=get_custom_url( + request_base_url=str(request.base_url), + route=f"v1/skills/{skill.skill_id}/archive", + ), + digest=archive.digest, + ) + for (skill, archive), name in zip(installable, names, strict=True) + ) + ) + + +@router.get( + "/v1/skills/{skill_id}/archive", + dependencies=(Depends(ensure_index_enabled),), + response_class=ZipArchiveResponse, +) +async def agent_skills_archive( + skill_id: str, + skill: LiteLLM_SkillsTable | None = Depends(stored_skill), +) -> ZipArchiveResponse: + """Stored skill upload, repacked so SKILL.md sits at the archive root.""" + archive: Final = await _archive_for(skill) if skill is not None else None + if archive is None: + raise HTTPException(status_code=404, detail=f"No installable skill archive for: {skill_id}") + + return ZipArchiveResponse( + content=archive.content, + headers=MappingProxyType({"Content-Disposition": f'attachment; filename="{skill_id}.zip"'}), + ) + + +async def _installable( + skills: Sequence[LiteLLM_SkillsTable], +) -> tuple[tuple[LiteLLM_SkillsTable, SkillArchive], ...]: + built: Final = tuple([(skill, await _archive_for(skill)) for skill in reversed(skills)]) + return tuple((skill, archive) for skill, archive in built if archive is not None) + + +async def _archive_for(skill: LiteLLM_SkillsTable) -> SkillArchive | None: + if skill.file_content is None: + return None + + cache_key: Final = None if skill.updated_at is None else f"{skill.skill_id}:{skill.updated_at.isoformat()}" + cached: Final = None if cache_key is None else _ARCHIVE_CACHE.get_cache(cache_key) + if isinstance(cached, SkillArchive): + return cached + + archive: Final = await asyncio.to_thread(build_skill_archive, skill.file_content) + if archive is None: + verbose_proxy_logger.warning( + "Agent Skills index: skipping skill %s, its upload is not a zip holding SKILL.md at the root of a " + "single top-level folder", + skill.skill_id, + ) + return None + + if cache_key is not None and len(archive.content) <= MAX_CACHED_ARCHIVE_BYTES: + _ARCHIVE_CACHE.set_cache(cache_key, archive) + return archive + + +def _base_name(skill: LiteLLM_SkillsTable, archive: SkillArchive) -> str: + candidates: Final = (archive.declared_name, skill.display_title, skill.skill_id) + return next( + (slug for slug in (_slugify(candidate) for candidate in candidates) if slug is not None), + _FALLBACK_SKILL_NAME, + ) + + +def _slugify(raw: str | None) -> str | None: + if raw is None: + return None + return _NON_SLUG_PATTERN.sub("-", raw.lower()).strip("-")[:MAX_SKILL_NAME_LENGTH].rstrip("-") or None + + +def _deduplicated(names: Sequence[str]) -> tuple[str, ...]: + ordinals: Final = MappingProxyType( + { + position: ordinal + for _, duplicates in groupby(sorted(enumerate(names), key=itemgetter(1)), key=itemgetter(1)) + for ordinal, (position, _) in enumerate(duplicates) + } + ) + return tuple(_with_ordinal(name, ordinals[position]) for position, name in enumerate(names)) + + +def _with_ordinal(name: str, ordinal: int) -> str: + if ordinal == 0: + return name + suffix: Final = f"-{ordinal + 1}" + return f"{name[: MAX_SKILL_NAME_LENGTH - len(suffix)].rstrip('-')}{suffix}" + + +def _description(skill: LiteLLM_SkillsTable, archive: SkillArchive, name: str) -> str: + candidates: Final = (archive.declared_description, skill.description, skill.display_title) + chosen: Final = next( + (candidate.strip() for candidate in candidates if candidate is not None and candidate.strip()), + name, + ) + return chosen[:MAX_SKILL_DESCRIPTION_LENGTH] diff --git a/litellm/proxy/guardrails/guardrail_hooks/azure/prompt_shield.py b/litellm/proxy/guardrails/guardrail_hooks/azure/prompt_shield.py index de9618a44a1..a0724b75ec7 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/azure/prompt_shield.py +++ b/litellm/proxy/guardrails/guardrail_hooks/azure/prompt_shield.py @@ -341,6 +341,7 @@ class AzureContentSafetyPromptShieldGuardrail(AzureGuardrailBase, CustomGuardrai guardrail_response: Final = self._summarize_guardrail_response( response=response, original_inputs=original_inputs, + event_type=event_type, ) self.add_standard_logging_guardrail_information_to_request_data( guardrail_json_response=guardrail_response, diff --git a/litellm/proxy/management_endpoints/budget_management_endpoints.py b/litellm/proxy/management_endpoints/budget_management_endpoints.py index 62a24109dbb..81a607aaa43 100644 --- a/litellm/proxy/management_endpoints/budget_management_endpoints.py +++ b/litellm/proxy/management_endpoints/budget_management_endpoints.py @@ -14,6 +14,7 @@ All /budget management endpoints #### BUDGET TABLE MANAGEMENT #### import math from collections.abc import Mapping +from types import MappingProxyType from typing import Final from fastapi import APIRouter, Depends, HTTPException @@ -176,6 +177,10 @@ async def update_budget( recomputed_reset_at: Final = ( {"budget_reset_at": get_budget_reset_time(budget_duration=budget_obj.budget_duration)} if budget_obj.budget_duration is not None and "budget_reset_at" not in budget_obj.model_fields_set + else MappingProxyType({"budget_reset_at": None}) + if "budget_duration" in budget_obj.model_fields_set + and budget_obj.budget_duration is None + and "budget_reset_at" not in budget_obj.model_fields_set else {} ) diff --git a/litellm/proxy/management_endpoints/common_daily_activity.py b/litellm/proxy/management_endpoints/common_daily_activity.py index a8aef30107c..44ed0017e42 100644 --- a/litellm/proxy/management_endpoints/common_daily_activity.py +++ b/litellm/proxy/management_endpoints/common_daily_activity.py @@ -127,6 +127,7 @@ def _key_metadata(api_key_metadata: Mapping[str, _KeyMetadataDict], api_key: str return KeyMetadata( key_alias=meta.get("key_alias"), team_id=meta.get("team_id"), + user_id=meta.get("user_id"), user_email=meta.get("user_email"), ) diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index 10c11119006..000b7f874ee 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -1254,8 +1254,8 @@ def _update_internal_user_params(data_json: dict, data: UpdateUserRequest | Upda fields_set: Final = data.fields_set() if hasattr(data, "fields_set") else set() for k, v in data_json.items(): - if k == "max_budget": - if "max_budget" in fields_set: + if k in ("max_budget", "budget_duration"): + if k in fields_set: non_default_values[k] = v elif k == "model_max_budget": if k in fields_set: @@ -1283,8 +1283,10 @@ def _update_internal_user_params(data_json: dict, data: UpdateUserRequest | Upda from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time validate_budget_duration(non_default_values["budget_duration"]) - non_default_values["budget_reset_at"] = get_budget_reset_time( - budget_duration=non_default_values["budget_duration"] + non_default_values["budget_reset_at"] = ( + get_budget_reset_time(budget_duration=non_default_values["budget_duration"]) + if non_default_values["budget_duration"] is not None + else None ) if "max_budget" not in non_default_values: diff --git a/litellm/proxy/management_endpoints/tag_management_endpoints.py b/litellm/proxy/management_endpoints/tag_management_endpoints.py index b74aa1a4e16..ab33d4bd766 100644 --- a/litellm/proxy/management_endpoints/tag_management_endpoints.py +++ b/litellm/proxy/management_endpoints/tag_management_endpoints.py @@ -438,6 +438,7 @@ async def update_tag( user_api_key_dict=user_api_key_dict, prisma_client=prisma_client, litellm_proxy_admin_name=litellm_proxy_admin_name, + budget_duration_cleared="budget_duration" in tag.model_fields_set and tag.budget_duration is None, ) # Get model names for model_info diff --git a/litellm/proxy/management_helpers/utils.py b/litellm/proxy/management_helpers/utils.py index e2d7262fb69..f3bd4b0f6dd 100644 --- a/litellm/proxy/management_helpers/utils.py +++ b/litellm/proxy/management_helpers/utils.py @@ -3,6 +3,7 @@ from collections.abc import Callable, Mapping, MutableMapping, Sequence from datetime import datetime from functools import wraps +from types import MappingProxyType from typing import Any, Final, Protocol from fastapi import HTTPException, Request @@ -180,6 +181,7 @@ async def handle_budget_for_entity( user_api_key_dict: UserAPIKeyAuth, prisma_client: PrismaClient, litellm_proxy_admin_name: str, + budget_duration_cleared: bool = False, ) -> str | None: """ Common helper to handle budget creation/updates for entities (organizations, tags, etc). @@ -208,7 +210,14 @@ async def handle_budget_for_entity( # Extract budget fields from data _json_data: Final = data.model_dump(exclude_none=True) if hasattr(data, "model_dump") else data - _budget_data: Final = {k: v for k, v in _json_data.items() if k in budget_params} + _budget_data: Final = MappingProxyType( + { + k: _json_data.get(k) + for k in budget_params + if k in _json_data + or (k == "budget_duration" and existing_budget_id is not None and budget_duration_cleared) + } + ) # Check if budget_id is explicitly provided in the data data_budget_id: Final[str | None] = getattr(data, "budget_id", None) 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 30b75a7b482..7cec3bac207 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 @@ -216,11 +216,14 @@ class AnthropicPassthroughLoggingHandler: model=model, speed=AnthropicPassthroughLoggingHandler._cost_relevant_speed(request_body), ) - if response is None: - return None - AnthropicPassthroughLoggingHandler._recover_interrupted_stream_output_tokens( + if not isinstance(response, ModelResponse): + return response + recovered_usage: Final = AnthropicPassthroughLoggingHandler._recover_interrupted_stream_output_tokens( response=response, all_chunks=all_chunks, model=model ) + if recovered_usage is None: + return response + AnthropicPassthroughLoggingHandler._clear_placeholder_cost(response=response, usage=recovered_usage) return response @staticmethod @@ -259,7 +262,9 @@ class AnthropicPassthroughLoggingHandler: ) except Exception as e: # noqa: BLE001 # an uncostable partial stream still bills its tokens, at zero cost verbose_proxy_logger.warning( - "Anthropic passthrough: could not cost the partial usage of a failed stream (model=%s): %s", model, e + "Anthropic passthrough: could not cost the partial usage of an interrupted stream (model=%s): %s", + model, + e, ) return 0.0 @@ -359,7 +364,7 @@ class AnthropicPassthroughLoggingHandler: response: ModelResponse | TextCompletionResponse, all_chunks: Sequence[str | bytes], model: str, - ) -> None: + ) -> Usage | None: """ An Anthropic stream interrupted before its terminal ``message_delta`` (client disconnect) carries only the ``message_start`` ``output_tokens`` @@ -369,24 +374,24 @@ class AnthropicPassthroughLoggingHandler: untouched because their terminal ``message_delta`` short-circuits here. """ if not isinstance(response, ModelResponse): - return + return None if not AnthropicPassthroughLoggingHandler._stream_was_interrupted(all_chunks): - return + return None usage: Final = getattr(response, "usage", None) - if usage is None: - return + if not isinstance(usage, Usage): + return None output_text: Final = get_content_from_model_response(response) if not output_text: - return + return None try: recovered_output_tokens = litellm.token_counter(model=model, text=output_text, count_response_tokens=True) except Exception: verbose_proxy_logger.warning( "Could not re-tokenize interrupted stream output; keeping placeholder completion token count." ) - return + return None if recovered_output_tokens <= (usage.completion_tokens or 0): - return + return None usage.completion_tokens = recovered_output_tokens usage.total_tokens = (usage.prompt_tokens or 0) + recovered_output_tokens # Anthropic costing reads completion_tokens_details.text_tokens, so the @@ -395,6 +400,12 @@ class AnthropicPassthroughLoggingHandler: details: Final = getattr(usage, "completion_tokens_details", None) if details is not None and getattr(details, "text_tokens", None) is not None: details.text_tokens = recovered_output_tokens + return usage + + @staticmethod + def _clear_placeholder_cost(response: ModelResponse, usage: Usage) -> None: + usage.cost = None + response._hidden_params.pop("response_cost", None) # pyright: ignore[reportPrivateUsage] # no public accessor @staticmethod def _create_anthropic_response_logging_payload( diff --git a/litellm/proxy/policy_engine/attachment_registry.py b/litellm/proxy/policy_engine/attachment_registry.py index a8ead86ac36..76b2291774e 100644 --- a/litellm/proxy/policy_engine/attachment_registry.py +++ b/litellm/proxy/policy_engine/attachment_registry.py @@ -6,6 +6,7 @@ This allows the same policy to be attached to multiple scopes. """ from datetime import datetime, timezone +from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, TypedDict from litellm._logging import verbose_proxy_logger @@ -141,8 +142,11 @@ class AttachmentRegistry: ), key=_attachment_specificity, ) + broadest_attachment_by_policy: Final = MappingProxyType( + {attachment.policy: attachment for attachment in reversed(matching_attachments)} + ) unique_attachments: Final = tuple( - next(attachment for attachment in matching_attachments if attachment.policy == policy_name) + broadest_attachment_by_policy[policy_name] for policy_name in dict.fromkeys(attachment.policy for attachment in matching_attachments) ) diff --git a/litellm/proxy/prisma_migration.py b/litellm/proxy/prisma_migration.py index 1b95d24c011..7e3aff75cef 100644 --- a/litellm/proxy/prisma_migration.py +++ b/litellm/proxy/prisma_migration.py @@ -14,6 +14,8 @@ sys.path.insert(0, os.path.abspath("./")) from typing import Final +from litellm_proxy_extras.prisma_toolchain import resolve_prisma_argv + from litellm._logging import verbose_proxy_logger from litellm.proxy.proxy_cli import run_server from litellm.secret_managers.main import str_to_bool @@ -29,7 +31,7 @@ def main() -> int: run_server(run_server_args, standalone_mode=False) verbose_proxy_logger.info("Running 'prisma generate'...") - result: Final = subprocess.run(("prisma", "generate"), capture_output=True, text=True) + result: Final = subprocess.run(resolve_prisma_argv(("prisma", "generate")), capture_output=True, text=True) verbose_proxy_logger.info("'prisma generate' stdout: %s", result.stdout) if result.returncode != 0: diff --git a/litellm/proxy/proxy_cli.py b/litellm/proxy/proxy_cli.py index c2d60cd5488..01a3da08998 100644 --- a/litellm/proxy/proxy_cli.py +++ b/litellm/proxy/proxy_cli.py @@ -1267,73 +1267,69 @@ def run_server( flush=True, ) sys.exit(1) - try: - from litellm.secret_managers.main import get_secret + from litellm.secret_managers.main import get_secret - connection_url_params: Final = _build_db_connection_url_params( - connection_limit=db_connection_pool_limit, - pool_timeout=db_connection_timeout, - connect_timeout=db_connect_timeout, - socket_timeout=db_socket_timeout, - disable_prepared_statements=db_disable_prepared_statements, - extra_params=db_extra_connection_params, + connection_url_params: Final = _build_db_connection_url_params( + connection_limit=db_connection_pool_limit, + pool_timeout=db_connection_timeout, + connect_timeout=db_connect_timeout, + socket_timeout=db_socket_timeout, + disable_prepared_statements=db_disable_prepared_statements, + extra_params=db_extra_connection_params, + ) + lifetime_params: Final = idle_lifetime_params(general_settings.get("database_max_idle_connection_lifetime")) + if os.getenv("DATABASE_URL", None) is not None: + database_url = get_secret("DATABASE_URL", default_value=None) + resolved_url: Final[str | None] = str(database_url) if database_url else None + pg_options: Final[str] = _pg_options_with_timeouts( + _url_query_value(resolved_url, "options"), + db_statement_timeout, + db_lock_timeout, ) - lifetime_params: Final = idle_lifetime_params( - general_settings.get("database_max_idle_connection_lifetime") + writer_url: Final = ( + _with_query_value(resolved_url, "options", pg_options) + if resolved_url and pg_options + else resolved_url ) - if os.getenv("DATABASE_URL", None) is not None: - database_url = get_secret("DATABASE_URL", default_value=None) - resolved_url: Final[str | None] = str(database_url) if database_url else None - pg_options: Final[str] = _pg_options_with_timeouts( - _url_query_value(resolved_url, "options"), - db_statement_timeout, - db_lock_timeout, - ) - writer_url: Final = ( - _with_query_value(resolved_url, "options", pg_options) - if resolved_url and pg_options - else resolved_url - ) - modified_url = append_query_params( - writer_url, - connection_url_params, - ) - os.environ["DATABASE_URL"] = translate_libpq_ssl_params( - add_missing_query_params(modified_url, lifetime_params) - ) - if os.getenv("DIRECT_URL", None) is not None: - database_url = os.getenv("DIRECT_URL") - modified_url = append_query_params(database_url, connection_url_params) - os.environ["DIRECT_URL"] = translate_libpq_ssl_params( - add_missing_query_params(modified_url, lifetime_params) - ) - # The reader pool is a real pool against the same configured cap, so it - # gets the allowlisted pool params. Schema-affecting ones, including any - # the operator smuggled in through database_extra_connection_params, stay - # on the writer. Anything pinned on the replica URL wins, unlike the - # writer where the config is applied on top. - read_replica_url: Final[str | None] = os.getenv("DATABASE_URL_READ_REPLICA") - if read_replica_url: - reader_options: Final[str] = _pg_options_with_timeouts( - _url_query_value(read_replica_url, "options"), - db_statement_timeout, - db_lock_timeout, - ) - os.environ["DATABASE_URL_READ_REPLICA"] = translate_libpq_ssl_params( + modified_url = append_query_params( + writer_url, + connection_url_params, + ) + os.environ["DATABASE_URL"] = translate_libpq_ssl_params( + add_missing_query_params(modified_url, lifetime_params) + ) + if os.getenv("DIRECT_URL", None) is not None: + database_url = os.getenv("DIRECT_URL") + modified_url = append_query_params(database_url, connection_url_params) + os.environ["DIRECT_URL"] = translate_libpq_ssl_params( + add_missing_query_params(modified_url, lifetime_params) + ) + # The reader pool is a real pool against the same configured cap, so it + # gets the allowlisted pool params. Schema-affecting ones, including any + # the operator smuggled in through database_extra_connection_params, stay + # on the writer. Anything pinned on the replica URL wins, unlike the + # writer where the config is applied on top. + read_replica_url: Final[str | None] = os.getenv("DATABASE_URL_READ_REPLICA") + if read_replica_url: + reader_options: Final[str] = _pg_options_with_timeouts( + _url_query_value(read_replica_url, "options"), + db_statement_timeout, + db_lock_timeout, + ) + os.environ["DATABASE_URL_READ_REPLICA"] = translate_libpq_ssl_params( + add_missing_query_params( add_missing_query_params( - add_missing_query_params( - _with_query_value(read_replica_url, "options", reader_options) - if reader_options - else read_replica_url, - reader_shareable_params(connection_url_params), - ), - lifetime_params, - ) + _with_query_value(read_replica_url, "options", reader_options) + if reader_options + else read_replica_url, + reader_shareable_params(connection_url_params), + ), + lifetime_params, ) - subprocess.run(["prisma"], capture_output=True) - is_prisma_runnable = True - except FileNotFoundError: - is_prisma_runnable = False + ) + from litellm_proxy_extras.prisma_toolchain import prisma_cli_available + + is_prisma_runnable: Final = prisma_cli_available() if is_prisma_runnable: from litellm.proxy.db.check_migration import check_prisma_schema_diff @@ -1382,7 +1378,8 @@ def run_server( ) else: print( - f"Unable to connect to DB. DATABASE_URL found in environment, but prisma package not found." # noqa: F541 + "Unable to connect to DB. DATABASE_URL found in environment, but the prisma CLI is neither on " + "PATH nor importable as a package." ) pgbouncer_settings: Final = PgBouncerSettings() upstream_database_url: Final = os.getenv("DATABASE_URL") diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 606a590c24b..d9f8e04ebda 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -353,6 +353,7 @@ from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import ( AuthCacheInvalidationSubscriber, ) from litellm.proxy.common_utils.callback_utils import initialize_callbacks_on_proxy +from litellm.proxy.common_utils.config_includes import resolve_include_file_path, resolve_includes from litellm.proxy.common_utils.config_sync_pubsub import ConfigSyncSubscriber from litellm.proxy.common_utils.debug_utils import init_verbose_loggers from litellm.proxy.common_utils.debug_utils import router as debugging_endpoints_router @@ -371,10 +372,7 @@ from litellm.proxy.common_utils.http_parsing_utils import ( check_file_size_under_limit, get_form_data, ) -from litellm.proxy.common_utils.load_config_utils import ( - get_config_file_contents_from_gcs, - get_file_contents_from_s3, -) +from litellm.proxy.common_utils.load_config_utils import get_config_from_bucket from litellm.proxy.common_utils.model_deprecation import collect_model_deprecations from litellm.proxy.common_utils.model_listing_utils import ( ClaudeCodeRoutingNames, @@ -448,7 +446,10 @@ from litellm.proxy.db.proxy_worker_heartbeat import ( ProxyWorkerHeartbeat, ) from litellm.proxy.db.spend_counter_reseed import END_USER_COUNTER_PREFIX, SpendCounterReseed -from litellm.proxy.discovery_endpoints import ui_discovery_endpoints_router +from litellm.proxy.discovery_endpoints import ( + agent_skills_discovery_router, + ui_discovery_endpoints_router, +) from litellm.proxy.fine_tuning_endpoints.endpoints import router as fine_tuning_router from litellm.proxy.fine_tuning_endpoints.endpoints import set_fine_tuning_config from litellm.proxy.google_endpoints.endpoints import router as google_router @@ -4803,12 +4804,12 @@ class ProxyConfig: if config is None: raise Exception("Config cannot be None or Empty.") # Process includes - config = self._process_includes(config=config, base_dir=os.path.dirname(os.path.abspath(file_path or ""))) + config = await self._process_includes(config=config, config_file_path=os.path.abspath(file_path or "")) # verbose_proxy_logger.debug(f"loaded config={json.dumps(config, indent=4)}") return config - def _process_includes(self, config: dict, base_dir: str) -> dict: + async def _process_includes(self, config: dict, config_file_path: str) -> dict: """ Process includes by appending their contents to the main config @@ -4823,29 +4824,21 @@ class ProxyConfig: callbacks: ["prometheus"] ``` """ - if "include" not in config: - return config - if not isinstance(config["include"], list): - raise ValueError("'include' must be a list of file paths") + included_config_adapter: Final = TypeAdapter(dict[str, object]) - # Load and append all included files - for include_file in config["include"]: - file_path = os.path.join(base_dir, include_file) + def resolve(include_file: str, declared_in: str) -> str: + return resolve_include_file_path(include_file, declared_in, config_file_path) + + async def read_included(file_path: str) -> Mapping[str, object]: if not os.path.exists(file_path): raise FileNotFoundError(f"Included file not found: {file_path}") + try: + return included_config_adapter.validate_python(self._load_yaml_file(file_path)) + except ValidationError as e: + raise ValueError(f"Included config file is not a YAML mapping: {file_path}") from e - included_config = self._load_yaml_file(file_path) - # Simply update/extend the main config with included config - for key, value in included_config.items(): - if isinstance(value, list) and key in config: - config[key].extend(value) - else: - config[key] = value - - # Remove the include directive - del config["include"] - return config + return await resolve_includes(config=config, location=config_file_path, resolve=resolve, read=read_included) async def save_config(self, new_config: dict, include_env_vars: bool = False): global prisma_client, general_settings, user_config_file_path, store_model_in_db @@ -5203,15 +5196,19 @@ class ProxyConfig: global prisma_client, store_model_in_db # Load existing config - if os.environ.get("LITELLM_CONFIG_BUCKET_NAME") is not None: - bucket_name: Final = os.environ.get("LITELLM_CONFIG_BUCKET_NAME") + bucket_name: Final = os.environ.get("LITELLM_CONFIG_BUCKET_NAME") + if bucket_name is not None: object_key: Final = os.environ.get("LITELLM_CONFIG_BUCKET_OBJECT_KEY") bucket_type: Final = os.environ.get("LITELLM_CONFIG_BUCKET_TYPE") verbose_proxy_logger.debug("bucket_name: %s, object_key: %s", bucket_name, object_key) - if bucket_type == "gcs": - config = await get_config_file_contents_from_gcs(bucket_name=bucket_name, object_key=object_key) - else: - config = get_file_contents_from_s3(bucket_name=bucket_name, object_key=object_key) + if object_key is None: + raise Exception("LITELLM_CONFIG_BUCKET_OBJECT_KEY must be set to load the config from a bucket.") + + config = await get_config_from_bucket( + bucket_type=bucket_type, + bucket_name=bucket_name, + object_key=object_key, + ) if config is None: raise Exception("Unable to load config from given source.") @@ -18798,6 +18795,7 @@ app.include_router(user_agent_analytics_router) app.include_router(gateway_request_router) app.include_router(enterprise_router) app.include_router(ui_discovery_endpoints_router) +app.include_router(agent_skills_discovery_router) # Eager: /models/{name}:method overlaps with the OpenAI /models endpoint. app.include_router(google_router) diff --git a/litellm/types/llms/meta.py b/litellm/types/llms/meta.py new file mode 100644 index 00000000000..40ecd6f7b67 --- /dev/null +++ b/litellm/types/llms/meta.py @@ -0,0 +1,21 @@ +from typing import Literal, TypeAlias + +from typing_extensions import NotRequired, ReadOnly, TypedDict + +MuseMode: TypeAlias = Literal["PUSH_TO_TALK", "ENDPOINTING"] +MuseAudioEncoding: TypeAlias = Literal["PCM_16KHZ", "PCM_24KHZ"] +MuseSampleRate: TypeAlias = Literal[16000, 24000] + + +class MuseAuthorization(TypedDict): + accessToken: ReadOnly[str] + + +class MuseHandshake(TypedDict): + authorization: ReadOnly[MuseAuthorization] + audioEncoding: ReadOnly[MuseAudioEncoding] + model: ReadOnly[str] + mode: ReadOnly[MuseMode] + partialMode: ReadOnly[Literal["CUMULATIVE"]] + emitAudioProgress: ReadOnly[bool] + languageBias: NotRequired[ReadOnly[tuple[str, ...]]] diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py index c021a932c0e..7a86c27efae 100644 --- a/litellm/types/llms/openai.py +++ b/litellm/types/llms/openai.py @@ -2197,6 +2197,53 @@ class OpenAIRealtimeInputAudioBufferSpeechEvent(TypedDict): item_id: ReadOnly[str] +class OpenAIRealtimeErrorDetail(TypedDict): + type: ReadOnly[str] + message: ReadOnly[str] + + +class OpenAIRealtimeErrorEvent(TypedDict): + type: ReadOnly[Literal["error"]] + error: ReadOnly[OpenAIRealtimeErrorDetail] + + +class OpenAIRealtimeTranscriptionAudioFormat(TypedDict): + type: ReadOnly[Literal["audio/pcm"]] + rate: ReadOnly[int] + + +class OpenAIRealtimeTranscriptionSettings(TypedDict): + model: ReadOnly[str] + language: NotRequired[ReadOnly[str]] + + +class OpenAIRealtimeServerVadTurnDetection(TypedDict): + type: ReadOnly[Literal["server_vad"]] + + +class OpenAIRealtimeTranscriptionAudioInput(TypedDict): + format: ReadOnly[OpenAIRealtimeTranscriptionAudioFormat] + transcription: ReadOnly[OpenAIRealtimeTranscriptionSettings] + turn_detection: ReadOnly[OpenAIRealtimeServerVadTurnDetection | None] + + +class OpenAIRealtimeTranscriptionAudio(TypedDict): + input: ReadOnly[OpenAIRealtimeTranscriptionAudioInput] + + +class OpenAIRealtimeTranscriptionSession(TypedDict): + id: ReadOnly[str] + object: ReadOnly[Literal["realtime.transcription_session"]] + type: ReadOnly[Literal["transcription"]] + audio: ReadOnly[OpenAIRealtimeTranscriptionAudio] + + +class OpenAIRealtimeTranscriptionSessionCreated(TypedDict): + type: ReadOnly[Literal["session.created"]] + event_id: ReadOnly[str] + session: ReadOnly[OpenAIRealtimeTranscriptionSession] + + class OpenAIRealtimeInputAudioTranscriptionDelta(TypedDict): type: ReadOnly[Literal["conversation.item.input_audio_transcription.delta"]] event_id: ReadOnly[str] @@ -2211,6 +2258,7 @@ class OpenAIRealtimeInputAudioTranscriptionCompleted(TypedDict): item_id: ReadOnly[str] content_index: ReadOnly[int] transcript: ReadOnly[str] + usage: NotRequired[ReadOnly[Mapping[str, object]]] class OpenAIRealtimeCachedTokensDetails(TypedDict, total=False): @@ -2274,6 +2322,8 @@ OpenAIRealtimeEvents = ( | OpenAIRealtimeInputAudioBufferSpeechEvent | OpenAIRealtimeInputAudioTranscriptionDelta | OpenAIRealtimeInputAudioTranscriptionCompleted + | OpenAIRealtimeTranscriptionSessionCreated + | OpenAIRealtimeErrorEvent ) OpenAIRealtimeStreamList = list[OpenAIRealtimeEvents] diff --git a/litellm/types/proxy/discovery_endpoints/agent_skills_endpoints.py b/litellm/types/proxy/discovery_endpoints/agent_skills_endpoints.py new file mode 100644 index 00000000000..0d8bb29e172 --- /dev/null +++ b/litellm/types/proxy/discovery_endpoints/agent_skills_endpoints.py @@ -0,0 +1,25 @@ +"""Agent Skills discovery index, version 0.2.0. + +Schema: https://schemas.agentskills.io/discovery/0.2.0/schema.json +""" + +from typing import Final, Literal + +from pydantic import BaseModel, Field + +AGENT_SKILLS_DISCOVERY_SCHEMA_URL: Final = "https://schemas.agentskills.io/discovery/0.2.0/schema.json" +MAX_SKILL_NAME_LENGTH: Final = 64 +MAX_SKILL_DESCRIPTION_LENGTH: Final = 1024 + + +class AgentSkillsIndexEntry(BaseModel): + name: str + type: Literal["archive"] + description: str + url: str + digest: str + + +class AgentSkillsIndex(BaseModel): + discovery_schema: str = Field(default=AGENT_SKILLS_DISCOVERY_SCHEMA_URL, alias="$schema") + skills: tuple[AgentSkillsIndexEntry, ...] diff --git a/litellm/types/proxy/management_endpoints/common_daily_activity.py b/litellm/types/proxy/management_endpoints/common_daily_activity.py index 2b39c5dbb9b..090e5c42376 100644 --- a/litellm/types/proxy/management_endpoints/common_daily_activity.py +++ b/litellm/types/proxy/management_endpoints/common_daily_activity.py @@ -43,6 +43,7 @@ class KeyMetadata(BaseModel): key_alias: str | None = None team_id: str | None = None + user_id: str | None = None user_email: str | None = None diff --git a/litellm/types/realtime.py b/litellm/types/realtime.py index 17dc70126f3..30db794c96e 100644 --- a/litellm/types/realtime.py +++ b/litellm/types/realtime.py @@ -169,9 +169,19 @@ class RealtimeInputAudioTranscriptionUsageInputTokenDetails(TypedDict): audio_tokens: ReadOnly[int] -class RealtimeInputAudioTranscriptionUsage(TypedDict): +class RealtimeInputAudioTranscriptionTokenUsage(TypedDict): type: ReadOnly[Literal["tokens"]] input_tokens: ReadOnly[int] output_tokens: ReadOnly[int] total_tokens: ReadOnly[int] input_token_details: ReadOnly[RealtimeInputAudioTranscriptionUsageInputTokenDetails] + + +class RealtimeInputAudioTranscriptionDurationUsage(TypedDict): + type: ReadOnly[Literal["duration"]] + seconds: ReadOnly[float] + + +RealtimeInputAudioTranscriptionUsage = ( + RealtimeInputAudioTranscriptionTokenUsage | RealtimeInputAudioTranscriptionDurationUsage +) diff --git a/litellm/utils.py b/litellm/utils.py index d048265ecd3..9a6d2ea8642 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -9286,6 +9286,10 @@ class ProviderConfigManager: from litellm.llms.gemini.realtime.transformation import GeminiRealtimeConfig return GeminiRealtimeConfig() + if LlmProviders.META == provider: + from litellm.llms.meta.realtime.transformation import MetaRealtimeConfig + + return MetaRealtimeConfig() return None @staticmethod diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 0d26f805565..de9867f9eee 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -4497,7 +4497,7 @@ }, "azure/eu/o1-2024-12-17": { "cache_read_input_token_cost": 8.25e-06, - "deprecation_date": "2026-10-21", + "deprecation_date": "2026-11-19", "input_cost_per_token": 1.65e-05, "litellm_provider": "azure", "max_input_tokens": 200000, @@ -4543,7 +4543,7 @@ }, "azure/eu/o3-mini-2025-01-31": { "cache_read_input_token_cost": 6.05e-07, - "deprecation_date": "2026-10-01", + "deprecation_date": "2026-11-19", "input_cost_per_token": 1.21e-06, "input_cost_per_token_batches": 6.05e-07, "litellm_provider": "azure", @@ -8730,7 +8730,7 @@ "supports_function_calling": true }, "azure/o1": { - "deprecation_date": "2026-10-21", + "deprecation_date": "2026-11-19", "cache_read_input_token_cost": 7.5e-06, "input_cost_per_token": 1.5e-05, "litellm_provider": "azure", @@ -8748,7 +8748,7 @@ }, "azure/o1-2024-12-17": { "cache_read_input_token_cost": 7.5e-06, - "deprecation_date": "2026-10-21", + "deprecation_date": "2026-11-19", "input_cost_per_token": 1.5e-05, "litellm_provider": "azure", "max_input_tokens": 200000, @@ -8825,7 +8825,7 @@ "supports_vision": false }, "azure/o3": { - "deprecation_date": "2026-10-21", + "deprecation_date": "2026-11-19", "cache_read_input_token_cost": 5e-07, "input_cost_per_token": 2e-06, "litellm_provider": "azure", @@ -8855,7 +8855,7 @@ "supports_vision": true }, "azure/o3-2025-04-16": { - "deprecation_date": "2026-10-21", + "deprecation_date": "2026-11-19", "cache_read_input_token_cost": 5e-07, "input_cost_per_token": 2e-06, "litellm_provider": "azure", @@ -8886,7 +8886,7 @@ }, "azure/o3-deep-research": { "cache_read_input_token_cost": 2.5e-06, - "deprecation_date": "2026-12-26", + "deprecation_date": "2026-11-19", "input_cost_per_token": 1e-05, "litellm_provider": "azure", "max_input_tokens": 200000, @@ -8923,7 +8923,7 @@ "supports_web_search": true }, "azure/o3-mini": { - "deprecation_date": "2026-10-01", + "deprecation_date": "2026-11-19", "cache_read_input_token_cost": 5.5e-07, "input_cost_per_token": 1.1e-06, "litellm_provider": "azure", @@ -8940,7 +8940,7 @@ }, "azure/o3-mini-2025-01-31": { "cache_read_input_token_cost": 5.5e-07, - "deprecation_date": "2026-10-01", + "deprecation_date": "2026-11-19", "input_cost_per_token": 1.1e-06, "litellm_provider": "azure", "max_input_tokens": 200000, @@ -8954,7 +8954,7 @@ "supports_vision": false }, "azure/o3-pro": { - "deprecation_date": "2026-12-17", + "deprecation_date": "2026-11-19", "input_cost_per_token": 2e-05, "input_cost_per_token_batches": 1e-05, "litellm_provider": "azure", @@ -8985,7 +8985,7 @@ "supports_vision": true }, "azure/o3-pro-2025-06-10": { - "deprecation_date": "2026-12-17", + "deprecation_date": "2026-11-19", "input_cost_per_token": 2e-05, "input_cost_per_token_batches": 1e-05, "litellm_provider": "azure", @@ -9016,7 +9016,7 @@ "supports_vision": true }, "azure/o4-mini": { - "deprecation_date": "2026-10-16", + "deprecation_date": "2026-11-19", "cache_read_input_token_cost": 2.75e-07, "input_cost_per_token": 1.1e-06, "litellm_provider": "azure", @@ -9047,7 +9047,7 @@ }, "azure/o4-mini-2025-04-16": { "cache_read_input_token_cost": 2.75e-07, - "deprecation_date": "2026-10-16", + "deprecation_date": "2026-11-19", "input_cost_per_token": 1.1e-06, "litellm_provider": "azure", "max_input_tokens": 200000, @@ -9600,7 +9600,7 @@ }, "azure/us/o1-2024-12-17": { "cache_read_input_token_cost": 8.25e-06, - "deprecation_date": "2026-10-21", + "deprecation_date": "2026-11-19", "input_cost_per_token": 1.65e-05, "litellm_provider": "azure", "max_input_tokens": 200000, @@ -9645,7 +9645,7 @@ "supports_vision": false }, "azure/us/o3-2025-04-16": { - "deprecation_date": "2026-10-21", + "deprecation_date": "2026-11-19", "cache_read_input_token_cost": 5.5e-07, "input_cost_per_token": 2.2e-06, "litellm_provider": "azure", @@ -9676,7 +9676,7 @@ }, "azure/us/o3-mini-2025-01-31": { "cache_read_input_token_cost": 6.05e-07, - "deprecation_date": "2026-10-01", + "deprecation_date": "2026-11-19", "input_cost_per_token": 1.21e-06, "input_cost_per_token_batches": 6.05e-07, "litellm_provider": "azure", @@ -9693,7 +9693,7 @@ }, "azure/us/o4-mini-2025-04-16": { "cache_read_input_token_cost": 3.1e-07, - "deprecation_date": "2026-10-16", + "deprecation_date": "2026-11-19", "input_cost_per_token": 1.21e-06, "litellm_provider": "azure", "max_input_tokens": 200000, @@ -14615,7 +14615,7 @@ }, "computer-use-preview": { "input_cost_per_token": 3e-06, - "litellm_provider": "azure", + "litellm_provider": "openai", "max_input_tokens": 8192, "max_output_tokens": 1024, "max_tokens": 1024, @@ -14633,12 +14633,14 @@ ], "supports_function_calling": true, "supports_parallel_function_calling": true, + "supports_pdf_input": true, "supports_prompt_caching": false, "supports_reasoning": true, - "supports_response_schema": true, + "supports_response_schema": false, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "source": "https://platform.openai.com/docs/models/computer-use-preview" }, "dall-e-2": { "deprecation_date": "2026-05-12", @@ -34720,6 +34722,22 @@ "supports_web_search": true, "supports_xhigh_reasoning_effort": true }, + "meta/muse-voice-transcribe-1.0": { + "input_cost_per_second": 0.00005, + "litellm_provider": "meta", + "mode": "audio_transcription", + "source": "https://dev.meta.ai/docs/speech-to-text", + "supported_endpoints": [ + "/v1/realtime" + ], + "supported_modalities": [ + "audio" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_input": true + }, "meta_llama/Llama-3.3-70B-Instruct": { "litellm_provider": "meta_llama", "max_input_tokens": 128000, @@ -39022,7 +39040,9 @@ "supports_max_reasoning_effort": true, "supports_tool_choice": true, "supports_vision": true, - "prompt_cache_min_tokens": 4096 + "prompt_cache_min_tokens": 4096, + "supports_response_schema": true, + "source": "https://openrouter.ai/api/v1/models" }, "openrouter/anthropic/claude-sonnet-4.5": { "input_cost_per_image": 0.0048, @@ -39170,18 +39190,20 @@ }, "openrouter/deepseek/deepseek-v3.2": { "input_cost_per_token": 2.69e-07, - "input_cost_per_token_cache_hit": 2.8e-08, + "input_cost_per_token_cache_hit": 1.345e-07, "litellm_provider": "openrouter", "max_input_tokens": 163840, - "max_output_tokens": 163840, - "max_tokens": 163840, + "max_output_tokens": 65536, + "max_tokens": 65536, "mode": "chat", "output_cost_per_token": 4e-07, "supports_assistant_prefill": true, "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_response_schema": true, + "source": "https://openrouter.ai/api/v1/models" }, "openrouter/deepseek/deepseek-v3.2-exp": { "input_cost_per_token": 2.7e-07, @@ -43491,6 +43513,7 @@ "supports_tool_choice": true }, "together_ai/openai/gpt-oss-20b": { + "deprecation_date": "2026-09-14", "input_cost_per_token": 5e-08, "litellm_provider": "together_ai", "max_input_tokens": 131072, @@ -43783,6 +43806,7 @@ "source": "https://docs.together.ai/docs/serverless-models" }, "together_ai/google/gemma-4-31B-it": { + "deprecation_date": "2026-09-14", "input_cost_per_token": 3.9e-07, "litellm_provider": "together_ai", "max_input_tokens": 262144, @@ -43797,6 +43821,7 @@ "supports_vision": true }, "together_ai/intfloat/multilingual-e5-large-instruct": { + "deprecation_date": "2026-09-14", "input_cost_per_token": 2e-08, "litellm_provider": "together_ai", "max_input_tokens": 514, @@ -43909,6 +43934,7 @@ "supports_tool_choice": true }, "together_ai/thinkingmachines/Inkling-Small": { + "deprecation_date": "2026-09-14", "cache_read_input_token_cost": 1e-07, "input_cost_per_token": 5e-07, "litellm_provider": "together_ai", @@ -59101,9 +59127,9 @@ "litellm_provider": "wandb", "mode": "chat", "supports_reasoning": true, - "input_cost_per_token": 0.00000131, - "output_cost_per_token": 0.00000396, - "cache_read_input_token_cost": 0.000000044, + "input_cost_per_token": 1.31e-06, + "output_cost_per_token": 3.96e-06, + "cache_read_input_token_cost": 4.4e-08, "supports_prompt_caching": true, "source": "https://wandb.ai/site/pricing/tokens/" }, @@ -59111,9 +59137,9 @@ "litellm_provider": "wandb", "mode": "chat", "supports_reasoning": true, - "input_cost_per_token": 0.0000001, - "output_cost_per_token": 0.00000015, - "cache_read_input_token_cost": 0.00000005, + "input_cost_per_token": 1e-07, + "output_cost_per_token": 1.5e-07, + "cache_read_input_token_cost": 5e-08, "supports_prompt_caching": true, "source": "https://wandb.ai/site/pricing/tokens/" }, @@ -60835,6 +60861,7 @@ "source": "https://docs.together.ai/docs/serverless-models" }, "together_ai/moonshotai/Kimi-K2.6": { + "deprecation_date": "2026-08-19", "input_cost_per_token": 1.2e-06, "output_cost_per_token": 4.5e-06, "cache_read_input_token_cost": 2e-07, @@ -60861,6 +60888,7 @@ "source": "https://api.together.xyz/v1/models" }, "together_ai/zai-org/GLM-5": { + "deprecation_date": "2026-06-22", "input_cost_per_token": 1e-06, "output_cost_per_token": 3.2e-06, "litellm_provider": "together_ai", @@ -60869,6 +60897,7 @@ "source": "https://api.together.xyz/v1/models" }, "together_ai/zai-org/GLM-5.1": { + "deprecation_date": "2026-07-10", "input_cost_per_token": 1.4e-06, "output_cost_per_token": 4.4e-06, "cache_read_input_token_cost": 2.6e-07, @@ -60886,6 +60915,7 @@ "source": "https://api.together.xyz/v1/models" }, "together_ai/Qwen/Qwen3-Coder-Next-FP8": { + "deprecation_date": "2026-05-14", "input_cost_per_token": 5e-07, "output_cost_per_token": 1.2e-06, "litellm_provider": "together_ai", @@ -60894,6 +60924,7 @@ "source": "https://api.together.xyz/v1/models" }, "together_ai/Qwen/Qwen3-VL-32B-Instruct": { + "deprecation_date": "2026-02-25", "input_cost_per_token": 5e-07, "output_cost_per_token": 1.5e-06, "litellm_provider": "together_ai", @@ -60902,6 +60933,7 @@ "source": "https://api.together.xyz/v1/models" }, "together_ai/Qwen/Qwen3-VL-8B-Instruct": { + "deprecation_date": "2026-04-16", "input_cost_per_token": 1.8e-07, "output_cost_per_token": 6.8e-07, "litellm_provider": "together_ai", @@ -60934,6 +60966,7 @@ "source": "https://api.together.xyz/v1/models" }, "together_ai/Qwen/QwQ-32B": { + "deprecation_date": "2025-11-13", "input_cost_per_token": 1.2e-06, "output_cost_per_token": 1.2e-06, "litellm_provider": "together_ai", diff --git a/pyproject.toml b/pyproject.toml index 31c498ccbbb..62ce4b4fd61 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -67,8 +67,8 @@ proxy = [ "azure-identity>=1.25.2,<2.0", "azure-storage-blob>=12.28.0,<13.0", "mcp>=1.28.1,<2.0", - "litellm-proxy-extras==0.4.96", - "litellm-enterprise==0.1.66", + "litellm-proxy-extras==0.4.97", + "litellm-enterprise==0.1.67", "RestrictedPython>=8.5,<9.0", "rich>=13.9.4,<14.0", "InquirerPy>=0.3.4,<1.0", diff --git a/tests/enterprise/litellm_enterprise/proxy/management_endpoints/test_project_endpoints_prisma.py b/tests/enterprise/litellm_enterprise/proxy/management_endpoints/test_project_endpoints_prisma.py index c23b203feba..36878fa698c 100644 --- a/tests/enterprise/litellm_enterprise/proxy/management_endpoints/test_project_endpoints_prisma.py +++ b/tests/enterprise/litellm_enterprise/proxy/management_endpoints/test_project_endpoints_prisma.py @@ -1292,6 +1292,21 @@ async def test_update_project_leaves_metadata_untouched_when_no_limit_is_sent(mo assert "metadata" not in _written_project_data(mock_prisma) +@pytest.mark.asyncio +async def test_update_project_clears_only_the_explicit_budget_cap(monkeypatch): + mock_prisma = _project_update_mocks(monkeypatch, {}) + mock_prisma.db.litellm_projecttable.find_unique.return_value.budget_id = "budget-clear-test" + mock_prisma.db.litellm_budgettable.update = mock.AsyncMock() + + await _run_project_update("project-clear-test", max_budget=None) + + mock_prisma.db.litellm_budgettable.update.assert_awaited_once_with( + where={"budget_id": "budget-clear-test"}, + data={"max_budget": None, "updated_by": "1234"}, + ) + assert "max_budget" not in _written_project_data(mock_prisma) + + @pytest.mark.parametrize("entry", ["all-proxy-models", "*", "azure/*"]) def test_enforce_project_model_quota_rejects_entries_that_expand_at_request_time(entry): """A quota keyed on a wildcard entry is never applied by the limiter, so it fails loudly.""" diff --git a/tests/proxy_migration_tests/test_prisma_toolchain.py b/tests/proxy_migration_tests/test_prisma_toolchain.py index 0ed33193a9b..4e2274cc582 100644 --- a/tests/proxy_migration_tests/test_prisma_toolchain.py +++ b/tests/proxy_migration_tests/test_prisma_toolchain.py @@ -37,7 +37,10 @@ from litellm_proxy_extras.prisma_toolchain import ( node_binary_path, prisma_bootstrap_timeout, prisma_command_timeout, + prisma_cli_available, prisma_migrate_deploy_timeout, + resolve_prisma_argv, + run_prisma, ) from litellm_proxy_extras.utils import ProxyExtrasDBManager @@ -401,3 +404,95 @@ def test_every_prisma_command_timeout_is_overridable(module: str) -> None: f"{module} still hardcodes a Prisma timeout at lines {literals}; " "route it through prisma_command_timeout() so it can be raised without a release" ) + + +FAKE_PRISMA_MODULE_MAIN = """import json +import sys + +print(json.dumps({"module_argv": sys.argv[1:]})) +""" + + +def _write_fake_prisma_module(tmp_path: Path) -> Path: + package_dir = tmp_path / "fakemodule" / "prisma" + package_dir.mkdir(parents=True) + (package_dir / "__init__.py").write_text("") + (package_dir / "__main__.py").write_text(FAKE_PRISMA_MODULE_MAIN) + return package_dir.parent + + +def _empty_bin(tmp_path: Path) -> Path: + bin_dir = tmp_path / "emptybin" + bin_dir.mkdir() + return bin_dir + + +def test_run_prisma_uses_the_module_when_the_console_script_is_not_on_path( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + empty_bin = _empty_bin(tmp_path) + module_root = _write_fake_prisma_module(tmp_path) + monkeypatch.setenv("PATH", str(empty_bin)) + + result = run_prisma( + ["prisma", "migrate", "deploy"], + timeout=60, + env={"PATH": str(empty_bin), "PYTHONPATH": str(module_root)}, + ) + + assert json.loads(result.stdout) == {"module_argv": ["migrate", "deploy"]} + + +def test_run_prisma_prefers_the_console_script_on_path( + toolchain_env: tuple[Path, Path], tmp_path: Path +) -> None: + _, log_path = toolchain_env + module_root = _write_fake_prisma_module(tmp_path) + + result = run_prisma( + ["prisma", "--version"], + timeout=60, + env={**os.environ, "PYTHONPATH": str(module_root)}, + ) + + assert [call["args"] for call in _fake_prisma_calls(log_path)] == [["--version"]] + assert "module_argv" not in result.stdout + + +def test_resolve_prisma_argv_leaves_an_explicit_cli_path_alone( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setenv("PATH", str(_empty_bin(tmp_path))) + explicit = ("/app/.cache/prisma-python/prisma", "migrate", "deploy") + + assert resolve_prisma_argv(explicit) == explicit + + +def test_prisma_cli_is_unavailable_with_neither_script_nor_package( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setenv("PATH", str(_empty_bin(tmp_path))) + monkeypatch.delitem(sys.modules, "prisma", raising=False) + monkeypatch.setattr(sys, "path", []) + + assert prisma_cli_available() is False + + +def test_prisma_cli_is_available_through_the_package_alone( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setenv("PATH", str(_empty_bin(tmp_path))) + monkeypatch.delitem(sys.modules, "prisma", raising=False) + monkeypatch.setattr(sys, "path", [str(_write_fake_prisma_module(tmp_path))]) + + assert prisma_cli_available() is True + + +def test_prisma_cli_is_available_through_the_console_script_alone( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setenv("PATH", str(_write_fake_prisma(tmp_path))) + monkeypatch.delitem(sys.modules, "prisma", raising=False) + monkeypatch.setattr(sys, "path", []) + + assert prisma_cli_available() is True diff --git a/tests/test_litellm/caching/test_redis_cache.py b/tests/test_litellm/caching/test_redis_cache.py index 4bf7894829e..c1e3240adb7 100644 --- a/tests/test_litellm/caching/test_redis_cache.py +++ b/tests/test_litellm/caching/test_redis_cache.py @@ -1205,6 +1205,66 @@ async def test_a_probe_overtaken_by_a_later_outage_leaves_the_breaker_to_the_new assert breaker._state == breaker.CLOSED +@pytest.mark.asyncio +async def test_pool_wait_timeout_is_a_timeout_failure_not_hard_connectivity(): + """A saturated blocking pool must not open the breaker before the timeout minimum duration. + + redis-py's async BlockingConnectionPool gives up waiting for a free connection by raising + ConnectionError("No connection available.") chained from asyncio.TimeoutError. Redis itself + is healthy in that case, so the failure has to be classed as a timeout and stay behind the + duration gate instead of being counted as a hard connectivity failure. + """ + from fakeredis import FakeServer + from fakeredis.aioredis import FakeConnection + from redis.asyncio import BlockingConnectionPool, Redis + from redis.exceptions import ConnectionError as RedisConnectionError + + from litellm.caching.redis_cache import RedisCircuitBreaker, _run_under_circuit_breaker + + pool = BlockingConnectionPool(connection_class=FakeConnection, server=FakeServer(), max_connections=1, timeout=0.01) + client = Redis(connection_pool=pool) + breaker = RedisCircuitBreaker(failure_threshold=3, recovery_timeout=60, timeout_min_duration=5.0) + + busy_connection = await pool.get_connection() + try: + for _ in range(breaker.failure_threshold * 2): + with pytest.raises(RedisConnectionError, match="No connection available"): + await _run_under_circuit_breaker(breaker, "op", lambda: client.get("k")) + finally: + await pool.release(busy_connection) + + assert breaker.is_open() is False, "a busy pool is a timeout gated on duration, not a dead Redis" + assert await _run_under_circuit_breaker(breaker, "op", lambda: client.get("k")) is None + await client.aclose() + + +def test_timeout_classification_follows_the_explicit_cause_chain_only(): + from redis.exceptions import ConnectionError as RedisConnectionError + + from litellm.caching.redis_cache import _is_redis_timeout_failure + + def raise_chained_from_timeout() -> None: + try: + raise asyncio.TimeoutError() + except asyncio.TimeoutError as err: + raise RedisConnectionError("No connection available.") from err + + def raise_while_handling_timeout() -> None: + try: + raise asyncio.TimeoutError() + except asyncio.TimeoutError: + raise RedisConnectionError("refused") + + with pytest.raises(RedisConnectionError) as chained: + raise_chained_from_timeout() + with pytest.raises(RedisConnectionError) as contextual: + raise_while_handling_timeout() + + assert _is_redis_timeout_failure(chained.value) is True + assert _is_redis_timeout_failure(contextual.value) is False + assert _is_redis_timeout_failure(RedisConnectionError("refused")) is False + + class _RoundTripCountingRedis: """Fake redis.asyncio client: one round trip per awaited command or pipeline execute.""" diff --git a/tests/test_litellm/caching/test_redis_semantic_cache.py b/tests/test_litellm/caching/test_redis_semantic_cache.py index df990c43530..9884e9d9bc0 100644 --- a/tests/test_litellm/caching/test_redis_semantic_cache.py +++ b/tests/test_litellm/caching/test_redis_semantic_cache.py @@ -1,3 +1,5 @@ +from collections.abc import Iterator +from contextlib import contextmanager from importlib import import_module import sys from unittest.mock import AsyncMock, MagicMock, patch @@ -5,18 +7,19 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest +@contextmanager +def _fake_redisvl_modules(semantic_cache_mock: MagicMock, custom_vectorizer_mock: MagicMock) -> Iterator[None]: + with pytest.MonkeyPatch.context() as mp: + mp.setitem(sys.modules, "redisvl.extensions.llmcache", MagicMock(SemanticCache=semantic_cache_mock)) + mp.setitem(sys.modules, "redisvl.utils.vectorize", MagicMock(CustomTextVectorizer=custom_vectorizer_mock)) + yield + # Tests for RedisSemanticCache def test_redis_semantic_cache_initialization(monkeypatch): # Mock the redisvl import semantic_cache_mock = MagicMock() - with patch.dict( - "sys.modules", - { - "redisvl.extensions.llmcache": MagicMock(SemanticCache=semantic_cache_mock), - "redisvl.utils.vectorize": MagicMock(CustomTextVectorizer=MagicMock()), - }, - ): + with _fake_redisvl_modules(semantic_cache_mock, MagicMock()): from litellm.caching.redis_semantic_cache import RedisSemanticCache # Set environment variables @@ -44,15 +47,7 @@ def test_redis_semantic_cache_get_cache(monkeypatch): semantic_cache_mock = MagicMock() custom_vectorizer_mock = MagicMock() - with patch.dict( - "sys.modules", - { - "redisvl.extensions.llmcache": MagicMock(SemanticCache=semantic_cache_mock), - "redisvl.utils.vectorize": MagicMock( - CustomTextVectorizer=custom_vectorizer_mock - ), - }, - ): + with _fake_redisvl_modules(semantic_cache_mock, custom_vectorizer_mock): from litellm.caching.redis_semantic_cache import RedisSemanticCache # Set environment variables @@ -110,15 +105,7 @@ def test_redis_semantic_cache_rejects_unscoped_cache_hit(monkeypatch): semantic_cache_mock = MagicMock() custom_vectorizer_mock = MagicMock() - with patch.dict( - "sys.modules", - { - "redisvl.extensions.llmcache": MagicMock(SemanticCache=semantic_cache_mock), - "redisvl.utils.vectorize": MagicMock( - CustomTextVectorizer=custom_vectorizer_mock - ), - }, - ): + with _fake_redisvl_modules(semantic_cache_mock, custom_vectorizer_mock): from litellm.caching.redis_semantic_cache import RedisSemanticCache monkeypatch.setenv("REDIS_HOST", "localhost") @@ -162,15 +149,7 @@ def test_redis_semantic_cache_set_cache_stores_cache_key_filter(monkeypatch): semantic_cache_mock = MagicMock() custom_vectorizer_mock = MagicMock() - with patch.dict( - "sys.modules", - { - "redisvl.extensions.llmcache": MagicMock(SemanticCache=semantic_cache_mock), - "redisvl.utils.vectorize": MagicMock( - CustomTextVectorizer=custom_vectorizer_mock - ), - }, - ): + with _fake_redisvl_modules(semantic_cache_mock, custom_vectorizer_mock): from litellm.caching.redis_semantic_cache import RedisSemanticCache monkeypatch.setenv("REDIS_HOST", "localhost") @@ -210,15 +189,7 @@ def test_redis_semantic_cache_uses_isolated_index_for_old_schema(monkeypatch): ) custom_vectorizer_mock = MagicMock() - with patch.dict( - "sys.modules", - { - "redisvl.extensions.llmcache": MagicMock(SemanticCache=semantic_cache_mock), - "redisvl.utils.vectorize": MagicMock( - CustomTextVectorizer=custom_vectorizer_mock - ), - }, - ): + with _fake_redisvl_modules(semantic_cache_mock, custom_vectorizer_mock): from litellm.caching.redis_semantic_cache import RedisSemanticCache monkeypatch.setenv("REDIS_HOST", "localhost") @@ -252,15 +223,7 @@ def test_redis_semantic_cache_overwrites_stale_isolated_index(monkeypatch): ) custom_vectorizer_mock = MagicMock() - with patch.dict( - "sys.modules", - { - "redisvl.extensions.llmcache": MagicMock(SemanticCache=semantic_cache_mock), - "redisvl.utils.vectorize": MagicMock( - CustomTextVectorizer=custom_vectorizer_mock - ), - }, - ): + with _fake_redisvl_modules(semantic_cache_mock, custom_vectorizer_mock): from litellm.caching.redis_semantic_cache import RedisSemanticCache monkeypatch.setenv("REDIS_HOST", "localhost") @@ -292,15 +255,7 @@ def test_redis_semantic_cache_reraises_unexpected_isolated_index_error(monkeypat ) custom_vectorizer_mock = MagicMock() - with patch.dict( - "sys.modules", - { - "redisvl.extensions.llmcache": MagicMock(SemanticCache=semantic_cache_mock), - "redisvl.utils.vectorize": MagicMock( - CustomTextVectorizer=custom_vectorizer_mock - ), - }, - ): + with _fake_redisvl_modules(semantic_cache_mock, custom_vectorizer_mock): from litellm.caching.redis_semantic_cache import RedisSemanticCache monkeypatch.setenv("REDIS_HOST", "localhost") @@ -369,15 +324,15 @@ def test_redis_semantic_cache_builds_filter_expression(monkeypatch): def __eq__(self, value): return (self.field_name, value) - with patch.dict("sys.modules", {"redisvl.query.filter": MagicMock(Tag=FakeTag)}): - from litellm.caching.redis_semantic_cache import RedisSemanticCache + monkeypatch.setitem(sys.modules, "redisvl.query.filter", MagicMock(Tag=FakeTag)) + from litellm.caching.redis_semantic_cache import RedisSemanticCache - redis_semantic_cache = RedisSemanticCache.__new__(RedisSemanticCache) + redis_semantic_cache = RedisSemanticCache.__new__(RedisSemanticCache) - assert redis_semantic_cache._get_cache_key_filter_expression("test_key") == ( - RedisSemanticCache.CACHE_KEY_FIELD_NAME, - "test_key", - ) + assert redis_semantic_cache._get_cache_key_filter_expression("test_key") == ( + RedisSemanticCache.CACHE_KEY_FIELD_NAME, + "test_key", + ) @pytest.mark.asyncio @@ -386,15 +341,7 @@ async def test_redis_semantic_cache_async_get_cache(monkeypatch): semantic_cache_mock = MagicMock() custom_vectorizer_mock = MagicMock() - with patch.dict( - "sys.modules", - { - "redisvl.extensions.llmcache": MagicMock(SemanticCache=semantic_cache_mock), - "redisvl.utils.vectorize": MagicMock( - CustomTextVectorizer=custom_vectorizer_mock - ), - }, - ): + with _fake_redisvl_modules(semantic_cache_mock, custom_vectorizer_mock): from litellm.caching.redis_semantic_cache import RedisSemanticCache # Set environment variables @@ -449,15 +396,7 @@ async def test_redis_semantic_cache_async_get_cache_rejects_unscoped_hit(monkeyp semantic_cache_mock = MagicMock() custom_vectorizer_mock = MagicMock() - with patch.dict( - "sys.modules", - { - "redisvl.extensions.llmcache": MagicMock(SemanticCache=semantic_cache_mock), - "redisvl.utils.vectorize": MagicMock( - CustomTextVectorizer=custom_vectorizer_mock - ), - }, - ): + with _fake_redisvl_modules(semantic_cache_mock, custom_vectorizer_mock): from litellm.caching.redis_semantic_cache import RedisSemanticCache monkeypatch.setenv("REDIS_HOST", "localhost") @@ -499,15 +438,7 @@ async def test_redis_semantic_cache_async_set_cache_stores_cache_key_filter( semantic_cache_mock = MagicMock() custom_vectorizer_mock = MagicMock() - with patch.dict( - "sys.modules", - { - "redisvl.extensions.llmcache": MagicMock(SemanticCache=semantic_cache_mock), - "redisvl.utils.vectorize": MagicMock( - CustomTextVectorizer=custom_vectorizer_mock - ), - }, - ): + with _fake_redisvl_modules(semantic_cache_mock, custom_vectorizer_mock): from litellm.caching.redis_semantic_cache import RedisSemanticCache monkeypatch.setenv("REDIS_HOST", "localhost") @@ -1255,15 +1186,7 @@ def test_redis_init_defers_redisvl_construction(monkeypatch): semantic_cache_mock = MagicMock() custom_vectorizer_mock = MagicMock() - with patch.dict( - "sys.modules", - { - "redisvl.extensions.llmcache": MagicMock(SemanticCache=semantic_cache_mock), - "redisvl.utils.vectorize": MagicMock( - CustomTextVectorizer=custom_vectorizer_mock - ), - }, - ): + with _fake_redisvl_modules(semantic_cache_mock, custom_vectorizer_mock): from litellm.caching.redis_semantic_cache import RedisSemanticCache monkeypatch.setenv("REDIS_HOST", "localhost") @@ -1291,15 +1214,7 @@ def test_redis_failed_llmcache_build_is_not_memoized(monkeypatch): ) custom_vectorizer_mock = MagicMock() - with patch.dict( - "sys.modules", - { - "redisvl.extensions.llmcache": MagicMock(SemanticCache=semantic_cache_mock), - "redisvl.utils.vectorize": MagicMock( - CustomTextVectorizer=custom_vectorizer_mock - ), - }, - ): + with _fake_redisvl_modules(semantic_cache_mock, custom_vectorizer_mock): from litellm.caching.redis_semantic_cache import RedisSemanticCache monkeypatch.setenv("REDIS_HOST", "localhost") diff --git a/tests/test_litellm/integrations/gcs_bucket/test_gcs_bucket_base.py b/tests/test_litellm/integrations/gcs_bucket/test_gcs_bucket_base.py index 8d662311da1..a458752bed0 100644 --- a/tests/test_litellm/integrations/gcs_bucket/test_gcs_bucket_base.py +++ b/tests/test_litellm/integrations/gcs_bucket/test_gcs_bucket_base.py @@ -128,3 +128,20 @@ class TestGCSBucketBase: assert object_name.endswith("-target_uploadType_media") assert ".." not in object_name assert "?" not in object_name + + +class TestGCSBucketLoggerBucketName: + @pytest.mark.asyncio + async def test_the_bucket_name_it_is_constructed_with_survives(self, monkeypatch): + """Reading config.yaml out of a GCS bucket asks for that bucket, not the logging one (LIT-6982).""" + monkeypatch.setenv("GCS_BUCKET_NAME", "logging-bucket") + monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True) + + assert GCSBucketLogger(bucket_name="config-bucket").BUCKET_NAME == "config-bucket" + + @pytest.mark.asyncio + async def test_no_bucket_name_still_falls_back_to_the_environment(self, monkeypatch): + monkeypatch.setenv("GCS_BUCKET_NAME", "logging-bucket") + monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True) + + assert GCSBucketLogger().BUCKET_NAME == "logging-bucket" diff --git a/tests/test_litellm/integrations/test_custom_guardrail.py b/tests/test_litellm/integrations/test_custom_guardrail.py index bb4822eae57..2fd5fa76d8e 100644 --- a/tests/test_litellm/integrations/test_custom_guardrail.py +++ b/tests/test_litellm/integrations/test_custom_guardrail.py @@ -3016,3 +3016,62 @@ class TestPreCallHookResponseIsNotLoggedVerbatim: ) assert self._logged_response(data) == "mask" + + @pytest.mark.asyncio + async def test_pre_call_hook_adding_tools_logs_mask(self): + class ToolInjectingGuardrail(CustomGuardrail): + @log_guardrail_information + async def async_pre_call_hook( + self, + user_api_key_dict: UserAPIKeyAuth, + cache: object, + data: dict[str, object], + call_type: str, + ) -> dict[str, object]: + return {**data, "tools": [{"type": "function", "function": {"name": "guardrail_injected_tool"}}]} + + data = self._request() + await ToolInjectingGuardrail(guardrail_name="g").async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), cache=None, data=data, call_type="acompletion" + ) + + assert self._logged_response(data) == "mask" + + @pytest.mark.asyncio + async def test_apply_guardrail_adding_tools_logs_mask(self): + class ToolInjectingGuardrail(CustomGuardrail): + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict[str, object], + input_type: Literal["request", "response"], + logging_obj: Optional["LiteLLMLoggingObj"] = None, + ) -> GenericGuardrailAPIInputs: + return {**inputs, "tools": [{"type": "function", "function": {"name": "guardrail_injected_tool"}}]} + + data = self._request() + await ToolInjectingGuardrail(guardrail_name="g").apply_guardrail( + inputs={"texts": ["SECRET_PROMPT"]}, request_data=data, input_type="request" + ) + + assert self._logged_response(data) == "mask" + + @pytest.mark.asyncio + async def test_apply_guardrail_masking_inputs_in_place_logs_mask(self): + class InPlaceMaskingGuardrail(CustomGuardrail): + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict[str, object], + input_type: Literal["request", "response"], + logging_obj: Optional["LiteLLMLoggingObj"] = None, + ) -> GenericGuardrailAPIInputs: + inputs["texts"] = [""] + return inputs + + data = self._request() + await InPlaceMaskingGuardrail(guardrail_name="g").apply_guardrail( + inputs={"texts": ["SECRET_PROMPT"]}, request_data=data, input_type="request" + ) + + assert self._logged_response(data) == "mask" diff --git a/tests/test_litellm/integrations/test_langsmith_init.py b/tests/test_litellm/integrations/test_langsmith_init.py index 025aa86466c..0bc9e279fbf 100644 --- a/tests/test_litellm/integrations/test_langsmith_init.py +++ b/tests/test_litellm/integrations/test_langsmith_init.py @@ -1,5 +1,6 @@ +import asyncio import os -from unittest.mock import MagicMock, patch +from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -154,24 +155,22 @@ class TestLangsmithLoggerInit: assert logger._start_periodic_flush_task() is None mock_get_running_loop.assert_called_once() - @patch("asyncio.get_running_loop") - def test_langsmith_init_starts_periodic_flush_with_running_loop( - self, mock_get_running_loop - ): + @pytest.mark.asyncio + async def test_langsmith_init_starts_periodic_flush_with_running_loop(self): """Test that init schedules periodic flush when a running loop exists.""" - mock_loop = MagicMock() - mock_task = MagicMock() - mock_loop.create_task.return_value = mock_task - mock_get_running_loop.return_value = mock_loop - logger = LangsmithLogger( - langsmith_api_key="test-key", langsmith_project="test-project" + langsmith_api_key="test-key", langsmith_project="test-project", flush_interval=0.01 ) + batch_sent = asyncio.Event() + logger.async_send_batch = AsyncMock(side_effect=batch_sent.set) + logger.log_queue.append({"id": "run-id"}) - assert logger._flush_task == mock_task - mock_loop.create_task.assert_called_once() - scheduled_coro = mock_loop.create_task.call_args.args[0] - scheduled_coro.close() + flush_task = logger._flush_task + assert isinstance(flush_task, asyncio.Task) + await asyncio.wait_for(batch_sent.wait(), timeout=5) + flush_task.cancel() + with pytest.raises(asyncio.CancelledError): + await flush_task @pytest.mark.asyncio async def test_async_log_success_event_lazily_starts_periodic_flush(self): diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_openai_cache_write_cost.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_openai_cache_write_cost.py index 92cb417f96c..18acfeda07d 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_openai_cache_write_cost.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_openai_cache_write_cost.py @@ -39,7 +39,7 @@ def test_openai_cache_write_tokens_billed_at_the_cache_creation_rate(local_model input_rate = rates["input_cost_per_token"] cache_write_rate = rates["cache_creation_input_token_cost"] output_rate = rates["output_cost_per_token"] - assert cache_write_rate == pytest.approx(input_rate * 1.25) + assert cache_write_rate > input_rate prompt_tokens = 12317 cache_write_tokens = 12314 diff --git a/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py b/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py index 9c0f6f59463..2a33d84ec78 100644 --- a/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py +++ b/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py @@ -10,8 +10,6 @@ from websockets.exceptions import ConnectionClosed from websockets.frames import Close import litellm - - from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.litellm_core_utils.realtime_streaming import ( REALTIME_SESSION_SUCCESS_LOGGED_KEY, @@ -20,10 +18,6 @@ from litellm.litellm_core_utils.realtime_streaming import ( ) from litellm.llms.xai.realtime.transformation import XAIRealtimeNormalizer from litellm.types.guardrails import GuardrailEventHooks -from litellm.types.llms.openai import ( - OpenAIRealtimeStreamResponseBaseObject, - OpenAIRealtimeStreamSessionEvents, -) def _make_transcript_event(text: str, item_id: str = "item_x") -> bytes: @@ -161,6 +155,7 @@ async def test_backend_to_client_send_text_receives_str_not_bytes(): logging_obj = MagicMock() logging_obj.async_success_handler = AsyncMock() logging_obj.success_handler = MagicMock() + logging_obj.dispatch_success_handlers = AsyncMock() streaming = RealTimeStreaming(client_ws, backend_ws, logging_obj) await streaming.backend_to_client_send_messages() @@ -812,7 +807,6 @@ async def test_transcription_captured_in_backend_to_client(): Test that conversation.item.input_audio_transcription.completed events from the backend are captured as user input during the WebSocket session. """ - import litellm client_ws = MagicMock() client_ws.send_text = AsyncMock() @@ -838,6 +832,7 @@ async def test_transcription_captured_in_backend_to_client(): logging_obj.model_call_details = {"messages": "default-message-value"} logging_obj.async_success_handler = AsyncMock() logging_obj.success_handler = MagicMock() + logging_obj.dispatch_success_handlers = AsyncMock() streaming = RealTimeStreaming(client_ws, backend_ws, logging_obj) await streaming.backend_to_client_send_messages() @@ -883,6 +878,7 @@ async def test_transcription_session_captures_usage_and_skips_response_create(): logging_obj.model_call_details = {} logging_obj.async_success_handler = AsyncMock() logging_obj.success_handler = MagicMock() + logging_obj.dispatch_success_handlers = AsyncMock() streaming = RealTimeStreaming(client_ws, backend_ws, logging_obj) await streaming.backend_to_client_send_messages() @@ -1100,7 +1096,6 @@ def test_capture_transcription_usage_deduplicates_when_already_stored(): When the event is already in messages (logged via store_message), it must not be appended a second time by _capture_transcription_usage. """ - import litellm streaming = RealTimeStreaming(MagicMock(), MagicMock(), MagicMock()) # Add the event type to the default logged list so _should_store_message returns True. @@ -1409,7 +1404,6 @@ async def test_realtime_guardrail_blocks_prompt_injection(monkeypatch: pytest.Mo ) - @pytest.mark.asyncio async def test_realtime_guardrail_allows_clean_transcript(monkeypatch: pytest.MonkeyPatch): """ @@ -1466,7 +1460,6 @@ async def test_realtime_guardrail_allows_clean_transcript(monkeypatch: pytest.Mo assert len(response_creates) == 1, f"Clean transcript should trigger response.create, got: {sent_to_backend}" - @pytest.mark.asyncio async def test_realtime_text_input_guardrail_blocks_and_returns_error(monkeypatch: pytest.MonkeyPatch): """ @@ -1560,7 +1553,6 @@ async def test_realtime_text_input_guardrail_blocks_and_returns_error(monkeypatc assert len(original_items) == 0, f"Blocked item should not be forwarded to backend, got: {original_items}" - @pytest.mark.asyncio async def test_realtime_function_call_output_guardrail_blocks_and_returns_error(monkeypatch: pytest.MonkeyPatch): """ @@ -1649,7 +1641,6 @@ async def test_realtime_function_call_output_guardrail_blocks_and_returns_error( assert "test@example.com" not in sanitized_item["output"] - @pytest.mark.asyncio async def test_realtime_function_call_output_guardrail_allows_clean_output(monkeypatch: pytest.MonkeyPatch): """ @@ -1714,7 +1705,6 @@ async def test_realtime_function_call_output_guardrail_allows_clean_output(monke assert len(forwarded) == 1, f"Clean function_call_output should be forwarded, got: {forwarded}" - @pytest.mark.asyncio async def test_realtime_text_input_guardrail_uses_pre_call_mode(monkeypatch: pytest.MonkeyPatch): """ @@ -1750,7 +1740,6 @@ async def test_realtime_text_input_guardrail_uses_pre_call_mode(monkeypatch: pyt ) - @pytest.mark.asyncio async def test_realtime_session_created_injects_session_update_for_audio_guardrail(monkeypatch: pytest.MonkeyPatch): """ @@ -1807,7 +1796,6 @@ async def test_realtime_session_created_injects_session_update_for_audio_guardra ) - @pytest.mark.asyncio async def test_realtime_session_created_does_not_inject_session_update_for_pre_call_only( monkeypatch: pytest.MonkeyPatch, @@ -1852,7 +1840,6 @@ async def test_realtime_session_created_does_not_inject_session_update_for_pre_c assert len(session_updates) == 0, f"pre_call-only guardrail must not inject session.update, got: {sent_to_backend}" - @pytest.mark.asyncio async def test_pre_call_and_post_call_guardrails_do_not_disable_server_vad(monkeypatch: pytest.MonkeyPatch): """Model Armor-style pre_call + post_call must not gate audio VAD.""" @@ -1868,17 +1855,17 @@ async def test_pre_call_and_post_call_guardrails_do_not_disable_server_vad(monke litellm, "callbacks", [ - ModelArmorStyleGuardrail( - guardrail_name="model_armor_all_pre_call", - event_hook=GuardrailEventHooks.pre_call, - default_on=False, - ), - ModelArmorStyleGuardrail( - guardrail_name="model_armor_all_post_call", - event_hook=GuardrailEventHooks.post_call, - default_on=False, - ), - ], + ModelArmorStyleGuardrail( + guardrail_name="model_armor_all_pre_call", + event_hook=GuardrailEventHooks.pre_call, + default_on=False, + ), + ModelArmorStyleGuardrail( + guardrail_name="model_armor_all_post_call", + event_hook=GuardrailEventHooks.post_call, + default_on=False, + ), + ], ) client_ws = MagicMock() @@ -1902,7 +1889,6 @@ async def test_pre_call_and_post_call_guardrails_do_not_disable_server_vad(monke assert streaming._has_audio_transcription_guardrails() is False - @pytest.mark.asyncio async def test_end_session_after_n_fails_closes_connection(monkeypatch: pytest.MonkeyPatch): """ @@ -1949,7 +1935,6 @@ async def test_end_session_after_n_fails_closes_connection(monkeypatch: pytest.M assert streaming._violation_count == 2 - @pytest.mark.asyncio async def test_on_violation_end_session_closes_on_first_fail(monkeypatch: pytest.MonkeyPatch): """ @@ -1995,7 +1980,6 @@ async def test_on_violation_end_session_closes_on_first_fail(monkeypatch: pytest assert streaming._violation_count == 1 - @pytest.mark.asyncio async def test_provider_path_suppresses_duplicate_session_created_after_synthetic(): client_ws = MagicMock() @@ -2956,7 +2940,9 @@ async def test_log_messages_routes_async_logging_through_bounded_worker(): mock_worker.ensure_initialized_and_enqueue.assert_called_once() enqueued = mock_worker.ensure_initialized_and_enqueue.call_args - assert (enqueued.args or tuple(enqueued.kwargs.values()))[0] is logging_obj.dispatch_success_handlers.return_value + assert (enqueued.args or tuple(enqueued.kwargs.values()))[ + 0 + ] is logging_obj.dispatch_success_handlers.return_value logging_obj.dispatch_success_handlers.assert_called_once_with(streaming.messages, prefer_async_handlers=True) logging_obj.success_handler.assert_not_called() # the bare create_task path must no longer be used for success logging @@ -3041,6 +3027,7 @@ async def test_session_close_flushes_unbilled_transcription_usage(): logging_obj: Final = MagicMock() logging_obj.async_success_handler = AsyncMock() logging_obj.success_handler = MagicMock() + logging_obj.dispatch_success_handlers = AsyncMock() usage: Final[RealtimeInputAudioTranscriptionUsage] = { "type": "tokens", @@ -3116,6 +3103,7 @@ async def test_session_close_flush_noop_without_unbilled_usage(): logging_obj: Final = MagicMock() logging_obj.async_success_handler = AsyncMock() logging_obj.success_handler = MagicMock() + logging_obj.dispatch_success_handlers = AsyncMock() provider_config: Final = MagicMock() provider_config.unbilled_usage_on_session_close = MagicMock(return_value=None) @@ -3136,7 +3124,6 @@ async def test_session_close_flush_noop_without_unbilled_usage(): ) - _UPSTREAM_REFUSAL: Final = "Publisher model `publishers/google/models/gemini-live-2.5-flash` was not found" @@ -3204,9 +3191,7 @@ def _backend_ws_closing_with(*frames: bytes | Exception) -> MagicMock: def _relay_session(client_ws: MagicMock, backend_ws: MagicMock) -> _RelaySession: logging: Final = _RecordingLogging() worker: Final = _InlineLoggingWorker() - streaming: Final = RealTimeStreaming( - client_ws, backend_ws, logging, model="gpt-realtime", logging_worker=worker - ) + streaming: Final = RealTimeStreaming(client_ws, backend_ws, logging, model="gpt-realtime", logging_worker=worker) return _RelaySession(streaming=streaming, logging=logging, worker=worker) @@ -3412,3 +3397,136 @@ async def test_refused_session_does_not_stamp_the_reservation_ownership_marker() assert session.logging.logged_failures == (upstream_close,) assert REALTIME_SESSION_SUCCESS_LOGGED_KEY not in session.logging.model_call_details + + +@pytest.mark.asyncio +async def test_transformed_transcription_completion_never_sends_response_create(): + from typing import Final + + completed_event: Final = { + "type": "conversation.item.input_audio_transcription.completed", + "event_id": "event_1", + "item_id": "turn_1", + "content_index": 0, + "transcript": "private transcript", + "usage": {"type": "duration", "seconds": 0.5}, + } + provider_config: Final = MagicMock() + provider_config.requires_session_configuration.return_value = True + provider_config.transform_realtime_response.return_value = { + "response": completed_event, + "current_output_item_id": None, + "current_response_id": None, + "current_delta_chunks": None, + "current_conversation_id": None, + "current_item_chunks": None, + "current_delta_type": None, + "session_configuration_request": None, + } + provider_config.transform_realtime_request.return_value = (json.dumps({"type": "response.create"}),) + provider_config.is_setup_message.return_value = False + provider_config.is_content_message.return_value = False + client_ws: Final = MagicMock() + client_ws.send_text = AsyncMock() + backend_ws: Final = MagicMock() + backend_ws.send = AsyncMock() + + streaming: Final = RealTimeStreaming( + client_ws, + backend_ws, + MagicMock(), + provider_config=provider_config, + model="muse-voice-transcribe-1.0", + force_transcription_model="muse-voice-transcribe-1.0", + ) + + await streaming._handle_provider_config_message("{}") + + assert json.loads(client_ws.send_text.await_args.args[0]) == completed_event + backend_ws.send.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_transcription_session_still_runs_transcription_guardrail(monkeypatch: pytest.MonkeyPatch): + class BlockingGuardrail(CustomGuardrail): + async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): + raise ValueError("blocked transcript") + + guardrail: Final = BlockingGuardrail( + guardrail_name="transcription-blocker", + event_hook=GuardrailEventHooks.realtime_input_transcription, + default_on=True, + ) + monkeypatch.setattr(litellm, "callbacks", [guardrail]) + + completed_event: Final = { + "type": "conversation.item.input_audio_transcription.completed", + "event_id": "event_1", + "item_id": "turn_1", + "content_index": 0, + "transcript": "blocked transcript", + "usage": {"type": "duration", "seconds": 0.5}, + } + provider_config: Final = MagicMock() + provider_config.requires_session_configuration.return_value = True + provider_config.transform_realtime_response.return_value = { + "response": completed_event, + "current_output_item_id": None, + "current_response_id": None, + "current_delta_chunks": None, + "current_conversation_id": None, + "current_item_chunks": None, + "current_delta_type": None, + "session_configuration_request": None, + } + provider_config.transform_realtime_request.return_value = () + provider_config.is_setup_message.return_value = False + provider_config.is_content_message.return_value = False + client_ws: Final = MagicMock() + client_ws.send_text = AsyncMock() + backend_ws: Final = MagicMock() + backend_ws.send = AsyncMock() + + streaming: Final = RealTimeStreaming( + client_ws, + backend_ws, + MagicMock(), + provider_config=provider_config, + model="muse-voice-transcribe-1.0", + force_transcription_model="muse-voice-transcribe-1.0", + ) + + await streaming._handle_provider_config_message("{}") + + sent_to_client: Final = [json.loads(call.args[0]) for call in client_ws.send_text.await_args_list] + assert completed_event in sent_to_client + error_events: Final = [event for event in sent_to_client if event.get("type") == "error"] + assert len(error_events) == 1 + assert error_events[0]["error"]["type"] == "guardrail_violation" + backend_ws.send.assert_not_awaited() + assert streaming._violation_count == 1 + + +@pytest.mark.asyncio +async def test_provider_bytes_are_sent_raw_after_pacing(): + from typing import Final + + backend_ws: Final = MagicMock() + backend_ws.send = AsyncMock() + provider_config: Final = MagicMock() + provider_config.requires_session_configuration.return_value = True + provider_config.transform_realtime_request.return_value = (b"\x00\x01", '{"type":"endStream"}') + provider_config.pace_backend_send = AsyncMock() + provider_config.is_setup_message.return_value = False + streaming: Final = RealTimeStreaming( + MagicMock(), + backend_ws, + MagicMock(), + provider_config=provider_config, + model="muse-voice-transcribe-1.0", + ) + + assert await streaming._send_to_backend(json.dumps({"type": "input_audio_buffer.commit"})) is True + + assert [call.args[0] for call in backend_ws.send.await_args_list] == [b"\x00\x01", '{"type":"endStream"}'] + provider_config.pace_backend_send.assert_awaited_once_with(b"\x00\x01") diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py index 01f7a2fb7ab..997a97c6fd3 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py @@ -1246,7 +1246,7 @@ async def _flush_logging_worker(capture: "_SuccessPayloadCapture") -> None: await asyncio.sleep(0) try: await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), timeout=10.0) - except (asyncio.TimeoutError, RuntimeError): + except asyncio.TimeoutError: pass deadline = asyncio.get_running_loop().time() + 10.0 while not capture.payloads and asyncio.get_running_loop().time() < deadline: diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py index be33b2ee3b1..8043496f299 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py @@ -1,6 +1,7 @@ import asyncio import json from datetime import datetime +from unittest.mock import patch import pytest @@ -824,6 +825,174 @@ async def test_async_sse_wrapper_bills_partial_when_detached_drains_disabled(mon assert len(streaming_iterator_module._DETACHED_STREAM_DRAINS) == 0 +class _SuccessRecorder(CustomLogger): + def __init__(self): + super().__init__() + self.success_kwargs: list = [] + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + self.success_kwargs.append(kwargs) + + +def _make_priced_logging_obj(call_id: str, recorder: _SuccessRecorder, model: str) -> LiteLLMLoggingObj: + logging_obj = LiteLLMLoggingObj( + model=model, + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type="anthropic_messages", + start_time=datetime.now(), + litellm_call_id=call_id, + function_id=call_id, + dynamic_async_success_callbacks=[recorder], + ) + logging_obj.update_environment_variables( + model=model, + user="", + optional_params={}, + litellm_params={"custom_llm_provider": "anthropic"}, + custom_llm_provider="anthropic", + ) + return logging_obj + + +class _UpstreamClosedOnDetach: + """Upstream that yields its events and then, like a socket read, waits until it is closed.""" + + def __init__(self, events: tuple[dict, ...]): + self._events = iter(events) + self._closed = asyncio.Event() + + def __aiter__(self): + return self + + async def __anext__(self) -> dict: + if self._closed.is_set(): + raise StopAsyncIteration + try: + return next(self._events) + except StopIteration: + await self._closed.wait() + raise StopAsyncIteration + + async def aclose(self) -> None: + self._closed.set() + + +@pytest.mark.asyncio +async def test_client_disconnect_partial_billing_prices_recovered_tokens(monkeypatch): + """ + Regression (LIT-6872): a client disconnect that lands on partial billing + re-tokenizes the buffered text into completion_tokens, but the logged cost + stayed priced at the message_start placeholder (1 output token). The success + row's response_cost must match its recovered completion_tokens. + """ + import litellm + from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER + + monkeypatch.setattr(streaming_iterator_module, "ANTHROPIC_MESSAGES_MAX_DETACHED_STREAM_DRAINS", 0) + monkeypatch.setattr(streaming_iterator_module, "ANTHROPIC_MESSAGES_STREAM_RELAY_QUEUE_MAXSIZE", 4) + model = "claude-sonnet-5" + recorder = _SuccessRecorder() + iterator = BaseAnthropicMessagesStreamingIterator( + litellm_logging_obj=_make_priced_logging_obj("disconnect_partial_cost", recorder, model), + request_body={"model": model, "stream": True}, + ) + sentence = "The history of computing spans centuries of mechanical and electronic invention. " + + async def _stream(): + yield {"type": "message_start", "message": {"id": "msg_1", "usage": {"input_tokens": 29, "output_tokens": 1}}} + yield {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}} + for _ in range(100): + yield {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": sentence}} + yield {"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {"output_tokens": 1500}} + yield {"type": "message_stop"} + + enqueued: list = [] + + def _capture(async_coroutine): + enqueued.append(async_coroutine) + + with patch.object( # test-quality-ok: GLOBAL_LOGGING_WORKER is a process-global singleton with no injection seam + GLOBAL_LOGGING_WORKER, "ensure_initialized_and_enqueue", side_effect=_capture + ): + gen = iterator.async_sse_wrapper(_stream()) + for _ in range(4): + await gen.__anext__() + await gen.aclose() + for _ in range(500): + if enqueued: + break + await asyncio.sleep(0.01) + + assert len(enqueued) == 1, "client disconnect never reached partial billing" + await enqueued[0] + + assert len(recorder.success_kwargs) == 1 + logged = recorder.success_kwargs[0]["standard_logging_object"] + assert 1 < logged["completion_tokens"] < 1500 + prompt_cost, completion_cost = litellm.cost_per_token( + model=model, prompt_tokens=29, completion_tokens=logged["completion_tokens"] + ) + assert logged["response_cost"] == pytest.approx(prompt_cost + completion_cost) + + +@pytest.mark.asyncio +async def test_proxy_disconnect_closing_upstream_prices_recovered_tokens(): + """ + Regression (LIT-6872), proxy path: after a client disconnect the proxy's + shielded cleanup closes the upstream stream while the pump is still reading + it, so the pump bills the chunks collected so far without ever seeing + message_delta. That row's response_cost must be priced from its recovered + completion_tokens, not from the message_start placeholder. + """ + import litellm + from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER + + model = "claude-sonnet-5" + recorder = _SuccessRecorder() + iterator = BaseAnthropicMessagesStreamingIterator( + litellm_logging_obj=_make_priced_logging_obj("disconnect_upstream_closed", recorder, model), + request_body={"model": model, "stream": True}, + ) + sentence = "The history of computing spans centuries of mechanical and electronic invention. " + upstream = _UpstreamClosedOnDetach( + ( + {"type": "message_start", "message": {"id": "msg_1", "usage": {"input_tokens": 29, "output_tokens": 1}}}, + {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}, + *({"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": sentence}} for _ in range(6)), + ) + ) + enqueued: list = [] + + def _capture(async_coroutine): + enqueued.append(async_coroutine) + + with patch.object( # test-quality-ok: GLOBAL_LOGGING_WORKER is a process-global singleton with no injection seam + GLOBAL_LOGGING_WORKER, "ensure_initialized_and_enqueue", side_effect=_capture + ): + gen = iterator.async_sse_wrapper(upstream) + for _ in range(4): + await gen.__anext__() + await gen.aclose() + assert not enqueued, "billing must wait for the upstream read to end, not the client detach" + await upstream.aclose() + for _ in range(500): + if enqueued: + break + await asyncio.sleep(0.01) + + assert len(enqueued) == 1, "closing the upstream never reached partial billing" + await enqueued[0] + + assert len(recorder.success_kwargs) == 1 + logged = recorder.success_kwargs[0]["standard_logging_object"] + assert logged["completion_tokens"] > 1 + prompt_cost, completion_cost = litellm.cost_per_token( + model=model, prompt_tokens=29, completion_tokens=logged["completion_tokens"] + ) + assert logged["response_cost"] == pytest.approx(prompt_cost + completion_cost) + + @pytest.mark.asyncio async def test_async_sse_wrapper_aborts_upstream_when_detached_drain_cap_reached(monkeypatch): """ diff --git a/tests/test_litellm/llms/meta/realtime/test_meta_realtime_transformation.py b/tests/test_litellm/llms/meta/realtime/test_meta_realtime_transformation.py new file mode 100644 index 00000000000..a5d7e47fb65 --- /dev/null +++ b/tests/test_litellm/llms/meta/realtime/test_meta_realtime_transformation.py @@ -0,0 +1,683 @@ +import base64 +import itertools +import json +from typing import Final +from unittest.mock import MagicMock + +import pytest + +from litellm.llms.meta.realtime.transformation import ( + DEFAULT_MUSE_REALTIME_URL, + MUSE_MODEL, + MetaRealtimeConfig, + MuseEventTransformer, + MuseProtocolError, + MuseSessionConfig, + build_muse_realtime_url, + normalize_access_token, + normalize_language, + parse_session_update, + session_created_event, +) +from litellm.types.llms.meta import MuseMode +from litellm.types.realtime import RealtimeResponseTransformInput + +EMPTY_TRANSFORM_INPUT: Final[RealtimeResponseTransformInput] = { + "session_configuration_request": None, + "current_output_item_id": None, + "current_response_id": None, + "current_delta_chunks": None, + "current_item_chunks": None, + "current_conversation_id": None, + "current_delta_type": None, +} + + +def _event(event_type: str, **fields: object) -> str: + return json.dumps({"type": event_type, **fields}) + + +def _ga_session_update(rate: int = 24_000, turn_detection: object = "server_vad") -> str: + return _event( + "session.update", + session={ + "type": "transcription", + "audio": { + "input": { + "format": {"type": "audio/pcm", "rate": rate}, + "turn_detection": None if turn_detection is None else {"type": turn_detection}, + "transcription": {"model": f"meta/{MUSE_MODEL}"}, + } + }, + }, + ) + + +def _configured(rate: int = 24_000, turn_detection: object = "server_vad", **kwargs: object) -> MetaRealtimeConfig: + config = MetaRealtimeConfig(**kwargs) + config.validate_environment({}, MUSE_MODEL, api_key="secret-token") + config.transform_realtime_request(_ga_session_update(rate, turn_detection), MUSE_MODEL) + return config + + +def _backend_events(config: MetaRealtimeConfig, payload: str) -> list[dict[str, object]]: + response = config.transform_realtime_response(payload, MUSE_MODEL, MagicMock(), EMPTY_TRANSFORM_INPUT)["response"] + assert isinstance(response, list) + return response + + +def test_beta_session_translates_language_and_drops_non_openai_hints(): + config = parse_session_update( + _event( + "session.update", + session={ + "type": "transcription", + "input_audio_format": "pcm16", + "turn_detection": {"type": "server_vad"}, + "input_audio_transcription": { + "model": "meta/muse-voice-transcribe-1.0", + "language": "en-US", + "prompt": "must not become a keyword", + }, + }, + ), + "meta/muse-voice-transcribe-1.0", + ) + + assert config.sample_rate == 24_000 + assert config.packet_bytes == 3_840 + assert config.mode == "ENDPOINTING" + assert config.language_bias == ("English",) + assert config.handshake("Bearer token") == { + "mode": "ENDPOINTING", + "authorization": {"accessToken": "Bearer token"}, + "audioEncoding": "PCM_24KHZ", + "model": MUSE_MODEL, + "partialMode": "CUMULATIVE", + "emitAudioProgress": True, + "languageBias": ("English",), + } + assert "must not become a keyword" not in json.dumps(config.handshake("Bearer token")) + + +def test_ga_session_accepts_16khz_mono_push_to_talk(): + config = parse_session_update( + _event( + "session.update", + session={ + "type": "transcription", + "audio": { + "input": { + "format": {"type": "audio/pcm", "rate": 16000, "channels": 1}, + "turn_detection": None, + "transcription": {"model": MUSE_MODEL, "language": "zh-Hans"}, + } + }, + }, + ), + MUSE_MODEL, + ) + + assert config.sample_rate == 16_000 + assert config.packet_bytes == 2_560 + assert config.mode == "PUSH_TO_TALK" + assert config.language_bias == ("Mandarin Chinese",) + assert config.handshake("Bearer token")["audioEncoding"] == "PCM_16KHZ" + assert "languageBias" not in MuseSessionConfig(MUSE_MODEL, "ENDPOINTING", 24_000, ()).handshake("Bearer token") + + +@pytest.mark.parametrize( + ("source", "expected"), + [ + ("EN_us", "English"), + ("mandarin chinese", "Mandarin Chinese"), + ("fil-PH", "Tagalog"), + ("iw-IL", "Hebrew"), + ("pt-BR", "Portuguese"), + ], +) +def test_language_normalization_uses_official_muse_names(source: str, expected: str): + assert normalize_language(source) == expected + + +@pytest.mark.parametrize( + ("session", "message"), + [ + ({"input_audio_format": "g711_ulaw"}, "requires pcm16"), + ({"audio": {"input": {"format": {"type": "audio/pcm", "rate": 8000}}}}, "16000 Hz or 24000 Hz"), + ( + {"audio": {"input": {"format": {"type": "audio/pcm", "rate": 24000, "channels": 2}}}}, + "requires mono", + ), + ( + {"input_audio_format": "pcm16", "audio": {"input": {"format": {"type": "audio/pcm"}}}}, + "either beta or GA layout", + ), + ({"input_audio_transcription": {"model": "other-model"}}, "cannot be changed"), + ({"input_audio_transcription": {"language": "xx"}}, "unsupported Muse Voice language"), + ({"turn_detection": {"type": "semantic_vad"}}, "server_vad turn detection or null"), + ({"type": "realtime", "audio": {"input": {"turn_detection": {"type": "semantic_vad"}}}}, "server_vad"), + ], +) +def test_session_rejects_unsupported_audio_model_and_hints(session: dict[str, object], message: str): + with pytest.raises(MuseProtocolError, match=message): + parse_session_update(_event("session.update", session={"type": "transcription", **session}), MUSE_MODEL) + + +def test_session_created_event_exposes_openai_transcription_shape(): + config = parse_session_update( + _event( + "session.update", + session={ + "audio": { + "input": { + "format": {"type": "audio/pcm", "rate": 24000}, + "transcription": {"model": MUSE_MODEL, "language": "ja"}, + } + }, + }, + ), + MUSE_MODEL, + ) + + created = session_created_event(config, "provider-session") + + assert created["type"] == "session.created" + assert created["session"]["id"] == "provider-session" + assert created["session"]["type"] == "transcription" + assert created["session"]["audio"]["input"]["turn_detection"] == {"type": "server_vad"} + assert created["session"]["audio"]["input"]["transcription"] == {"model": MUSE_MODEL, "language": "Japanese"} + + +def test_turnless_empty_silence_transcript_is_ignored(): + transformer = MuseEventTransformer() + + assert transformer.transform(json.loads(_event("transcript", transcript="", final=True))) == () + + +def test_transcript_without_speech_start_synthesizes_start_before_delta(): + transformer = MuseEventTransformer() + + events = transformer.transform(json.loads(_event("transcript", turnId="turn-1", transcript="hello", final=False))) + + assert [event["type"] for event in events] == [ + "input_audio_buffer.speech_started", + "conversation.item.input_audio_transcription.delta", + ] + + +def test_cumulative_partials_emit_only_extensions_and_final_is_authoritative(): + transformer = MuseEventTransformer() + + def send(payload: str) -> tuple[dict[str, object], ...]: + return transformer.transform(json.loads(payload)) + + started = send(_event("speechStart", turnId="turn-1")) + first = send(_event("transcript", turnId="turn-1", transcript="hello", final=False)) + extension = send(_event("transcript", turnId="turn-1", transcript="hello world", final=False)) + rewrite = send(_event("transcript", turnId="turn-1", transcript="hullo world", final=False)) + completed = send(_event("speechComplete", turnId="turn-1", transcript="hullo world")) + + assert [event["type"] for event in started] == ["input_audio_buffer.speech_started"] + assert first[0]["delta"] == "hello" + assert extension[0]["delta"] == " world" + assert rewrite == () + assert completed[0]["type"] == "input_audio_buffer.speech_stopped" + assert completed[1]["type"] == "conversation.item.input_audio_transcription.completed" + assert completed[1]["item_id"] == "turn-1" + assert completed[1]["transcript"] == "hullo world" + assert send(_event("speechEnd", turnId="turn-1")) == () + + +def test_speech_end_then_speech_complete_emits_stopped_then_completed(): + transformer = MuseEventTransformer() + + transformer.transform(json.loads(_event("speechStart", turnId="turn-1"))) + stopped = transformer.transform(json.loads(_event("speechEnd", turnId="turn-1"))) + completed = transformer.transform(json.loads(_event("speechComplete", turnId="turn-1", transcript="done"))) + + assert [event["type"] for event in stopped] == ["input_audio_buffer.speech_stopped"] + assert [event["type"] for event in completed] == ["conversation.item.input_audio_transcription.completed"] + assert completed[0]["transcript"] == "done" + + +def test_turnless_partial_between_speech_end_and_speech_complete_stays_on_that_turn(): + transformer = MuseEventTransformer() + transformer.transform(json.loads(_event("speechStart", turnId="turn-1"))) + transformer.transform(json.loads(_event("transcript", transcript="what is", final=False))) + transformer.transform(json.loads(_event("speechEnd", turnId="turn-1"))) + + post_processed = transformer.transform( + json.loads(_event("transcript", transcript="what is the weather", final=False)) + ) + completed = transformer.transform( + json.loads(_event("speechComplete", turnId="turn-1", transcript="What is the weather?")) + ) + + assert _typed(post_processed) == [("conversation.item.input_audio_transcription.delta", "turn-1")] + assert post_processed[0]["delta"] == " the weather" + assert _typed(completed) == [("conversation.item.input_audio_transcription.completed", "turn-1")] + assert completed[0]["transcript"] == "What is the weather?" + + +def _typed(events: tuple[dict[str, object], ...]) -> list[tuple[object, object]]: + return [(event["type"], event["item_id"]) for event in events] + + +def test_overlapping_turns_emit_independently_and_correlate_by_item_id(): + transformer = MuseEventTransformer() + + def send(payload: str) -> list[tuple[object, object]]: + return _typed(transformer.transform(json.loads(payload))) + + assert send(_event("speechStart", turnId="turn-a")) == [("input_audio_buffer.speech_started", "turn-a")] + assert send(_event("speechStart", turnId="turn-b")) == [("input_audio_buffer.speech_started", "turn-b")] + assert send(_event("transcript", turnId="turn-b", transcript="second", final=False)) == [ + ("conversation.item.input_audio_transcription.delta", "turn-b") + ] + assert send(_event("speechComplete", turnId="turn-a", transcript="first")) == [ + ("input_audio_buffer.speech_stopped", "turn-a"), + ("conversation.item.input_audio_transcription.completed", "turn-a"), + ] + assert send(_event("speechEnd", turnId="turn-a")) == [] + assert send(_event("speechEnd", turnId="turn-b")) == [("input_audio_buffer.speech_stopped", "turn-b")] + assert send(_event("speechComplete", turnId="turn-b", transcript="second final")) == [ + ("conversation.item.input_audio_transcription.completed", "turn-b") + ] + + +def test_empty_vad_turn_is_closed_and_does_not_block_the_next_turn(): + transformer = MuseEventTransformer() + + def send(payload: str) -> list[tuple[object, object]]: + return _typed(transformer.transform(json.loads(payload))) + + assert send(_event("speechStart", turnId="noise")) == [("input_audio_buffer.speech_started", "noise")] + assert send(_event("speechEnd", turnId="noise")) == [("input_audio_buffer.speech_stopped", "noise")] + assert send(_event("speechStart", turnId="speech")) == [("input_audio_buffer.speech_started", "speech")] + assert send(_event("transcript", turnId="speech", transcript="hello", final=False)) == [ + ("conversation.item.input_audio_transcription.delta", "speech") + ] + assert send(_event("speechEnd", turnId="speech")) == [("input_audio_buffer.speech_stopped", "speech")] + assert send(_event("speechComplete", turnId="speech", transcript="hello world")) == [ + ("conversation.item.input_audio_transcription.completed", "speech") + ] + + +@pytest.mark.parametrize("transcript", ["", "late words"]) +def test_late_speech_complete_after_an_empty_speech_end_completes_that_item(transcript: str): + transformer = MuseEventTransformer() + transformer.transform(json.loads(_event("speechStart", turnId="turn-1"))) + transformer.transform(json.loads(_event("speechEnd", turnId="turn-1"))) + transformer.transform(json.loads(_event("speechStart", turnId="turn-2"))) + + (completed,) = transformer.transform(json.loads(_event("speechComplete", turnId="turn-1", transcript=transcript))) + + assert completed["type"] == "conversation.item.input_audio_transcription.completed" + assert completed["item_id"] == "turn-1" + assert completed["transcript"] == transcript + + +def test_push_to_talk_speech_complete_closes_the_turn_without_speech_end(): + transformer = MuseEventTransformer() + transformer.configure(MuseSessionConfig(MUSE_MODEL, "PUSH_TO_TALK", 24_000, ())) + + transformer.transform(json.loads(_event("speechStart", turnId="turn-1"))) + transformer.transform(json.loads(_event("transcript", turnId="turn-1", transcript="hel", final=False))) + events = transformer.transform(json.loads(_event("speechComplete", turnId="turn-1", transcript="hello"))) + + assert [event["type"] for event in events] == [ + "input_audio_buffer.speech_stopped", + "conversation.item.input_audio_transcription.completed", + ] + assert events[1]["transcript"] == "hello" + + +_TERMINAL_SIGNALS: Final = { + "speechEnd": _event("speechEnd", turnId="turn-1"), + "speechComplete": _event("speechComplete", turnId="turn-1", transcript="final words"), + "final": _event("transcript", turnId="turn-1", transcript="final words", final=True), +} +_TERMINAL_ORDERINGS: Final = tuple( + ordering for size in (1, 2, 3) for ordering in itertools.permutations(_TERMINAL_SIGNALS, size) +) + + +@pytest.mark.parametrize("mode", ["ENDPOINTING", "PUSH_TO_TALK"]) +@pytest.mark.parametrize("ordering", _TERMINAL_ORDERINGS, ids="-".join) +def test_every_terminal_signal_order_closes_the_turn_exactly_once(mode: MuseMode, ordering: tuple[str, ...]): + transformer = MuseEventTransformer() + transformer.configure(MuseSessionConfig(MUSE_MODEL, mode, 24_000, ())) + transformer.transform(json.loads(_event("speechStart", turnId="turn-1"))) + transformer.transform(json.loads(_event("transcript", turnId="turn-1", transcript="fin", final=False))) + + emitted = [ + event["type"] for signal in ordering for event in transformer.transform(json.loads(_TERMINAL_SIGNALS[signal])) + ] + replayed = [ + event["type"] for signal in ordering for event in transformer.transform(json.loads(_TERMINAL_SIGNALS[signal])) + ] + + has_text = bool(set(ordering) & {"speechComplete", "final"}) + assert emitted == [ + "input_audio_buffer.speech_stopped", + *(["conversation.item.input_audio_transcription.completed"] if has_text else []), + ] + assert replayed == [] + + +def test_push_to_talk_final_transcript_completes_without_speech_end(): + transformer = MuseEventTransformer() + transformer.configure(MuseSessionConfig(MUSE_MODEL, "PUSH_TO_TALK", 24_000, ())) + + events = transformer.transform(json.loads(_event("transcript", transcript="hello there", final=True))) + + assert [event["type"] for event in events] == [ + "input_audio_buffer.speech_started", + "input_audio_buffer.speech_stopped", + "conversation.item.input_audio_transcription.completed", + ] + assert events[2]["transcript"] == "hello there" + assert str(events[0]["item_id"]).startswith("item_") + + +def test_positive_audio_progress_deltas_attach_to_next_completion_and_speaker_is_ignored(): + transformer = MuseEventTransformer() + + def send(payload: str) -> tuple[dict[str, object], ...]: + return transformer.transform(json.loads(payload)) + + send(_event("audioProgress", audioProcessedMs=1000)) + send(_event("audioProgress", audioProcessedMs=750)) + send(_event("audioProgress", audioProcessedMs=1600)) + assert send(_event("speaker", turnId=42, label=" Speaker 2 ")) == () + completed = send(_event("speechComplete", turnId=42, transcript="hello")) + + assert "speaker" not in completed[-1] + assert completed[-1]["usage"] == {"type": "duration", "seconds": 1.6} + assert transformer.take_unbilled_usage() is None + assert send(_event("speechEnd", turnId=42)) == () + + +def test_trailing_audio_progress_is_returned_once(): + transformer = MuseEventTransformer() + + transformer.transform(json.loads(_event("audioProgress", audioProcessedMs=250))) + + assert transformer.take_unbilled_usage() == {"type": "duration", "seconds": 0.25} + assert transformer.take_unbilled_usage() is None + + +def test_finished_turn_ignores_late_duplicates(): + transformer = MuseEventTransformer() + + released = transformer.transform(json.loads(_event("speechComplete", turnId="turn-1", transcript="done"))) + + assert [event["type"] for event in released] == [ + "input_audio_buffer.speech_started", + "input_audio_buffer.speech_stopped", + "conversation.item.input_audio_transcription.completed", + ] + assert transformer.transform(json.loads(_event("speechComplete", turnId="turn-1", transcript="duplicate"))) == () + assert transformer.transform(json.loads(_event("speechEnd", turnId="turn-1"))) == () + assert transformer.transform(json.loads(_event("speechStart", turnId="turn-1"))) == () + assert ( + transformer.transform(json.loads(_event("transcript", turnId="turn-1", transcript="late", final=False))) == () + ) + + +def test_late_duplicate_speech_start_does_not_capture_the_next_turnless_transcript(): + transformer = MuseEventTransformer() + transformer.configure(MuseSessionConfig(MUSE_MODEL, "PUSH_TO_TALK", 24_000, ())) + transformer.transform(json.loads(_event("speechStart", turnId="turn-1"))) + transformer.transform(json.loads(_event("speechComplete", turnId="turn-1", transcript="first"))) + + assert transformer.transform(json.loads(_event("speechStart", turnId="turn-1"))) == () + events = transformer.transform(json.loads(_event("transcript", transcript="second", final=True))) + + assert [event["type"] for event in events] == [ + "input_audio_buffer.speech_started", + "input_audio_buffer.speech_stopped", + "conversation.item.input_audio_transcription.completed", + ] + assert events[2]["transcript"] == "second" + assert events[2]["item_id"] != "turn-1" + + +def test_turn_memory_is_bounded_by_turn_limit(): + transformer = MuseEventTransformer(turn_limit=2) + + transformer.transform(json.loads(_event("speechComplete", turnId="turn-1", transcript="one"))) + transformer.transform(json.loads(_event("speechComplete", turnId="turn-2", transcript="two"))) + assert transformer.transform(json.loads(_event("speechEnd", turnId="turn-1"))) == () + transformer.transform(json.loads(_event("speechComplete", turnId="turn-3", transcript="three"))) + + forgotten = transformer.transform(json.loads(_event("speechEnd", turnId="turn-1"))) + + assert [event["type"] for event in forgotten] == ["input_audio_buffer.speech_stopped"] + + +def test_provider_error_is_sanitized_and_encodable(): + token = "private-token" + provider_body = f"authorization failed for Bearer {token}" + transformed = MuseEventTransformer().transform( + json.loads(_event("error", code="AUTH", message=provider_body, request={"accessToken": token})) + ) + + encoded = json.dumps(transformed[0]) + assert json.loads(encoded)["error"] == { + "type": "server_error", + "message": "Meta Muse realtime transcription failed", + } + assert token not in encoded + assert provider_body not in encoded + + +@pytest.mark.parametrize( + ("raw", "expected"), + [("token", "Bearer token"), (" Bearer token ", "Bearer token"), ("bearer token", "Bearer token")], +) +def test_access_token_normalization_adds_single_bearer_prefix(raw: str, expected: str): + assert normalize_access_token(raw) == expected + + +@pytest.mark.parametrize("raw", ["", " ", "Bearer", " bearer "]) +def test_access_token_normalization_rejects_empty_tokens(raw: str): + with pytest.raises(ValueError, match=r"token|key is required"): + normalize_access_token(raw) + + +@pytest.mark.parametrize( + ("api_base", "expected"), + [ + (None, DEFAULT_MUSE_REALTIME_URL), + ("https://example.test/custom/path?ignored=yes", "wss://example.test/v1/asr/realtime"), + ("wss://example.test:8443/other", "wss://example.test:8443/v1/asr/realtime"), + ], +) +def test_realtime_url_pins_muse_path(api_base: str | None, expected: str): + assert build_muse_realtime_url(api_base) == expected + assert MetaRealtimeConfig().get_complete_url(api_base, f"meta/{MUSE_MODEL}") == expected + + +@pytest.mark.parametrize( + "api_base", + [ + "http://example.test", + "ws://example.test", + "wss://user:pass@example.test", + "wss://example.test/path#fragment", + "not-a-url", + ], +) +def test_realtime_url_rejects_insecure_or_ambiguous_bases(api_base: str): + with pytest.raises(ValueError, match="absolute wss:// or https://"): + build_muse_realtime_url(api_base) + + +def test_unsupported_model_is_rejected_before_connecting(): + with pytest.raises(ValueError, match="Unsupported Meta realtime model: meta/other-model"): + MetaRealtimeConfig().get_complete_url(None, "meta/other-model") + + +def test_missing_api_key_is_rejected(monkeypatch: pytest.MonkeyPatch): + monkeypatch.delenv("META_API_KEY", raising=False) + + with pytest.raises(ValueError, match="api_key is required for Meta API calls"): + MetaRealtimeConfig().validate_environment({}, MUSE_MODEL) + + +def test_bearer_token_travels_only_in_the_json_handshake(): + config = MetaRealtimeConfig() + headers = {"x-existing": "kept"} + + assert config.validate_environment(headers, MUSE_MODEL, api_key="secret-token") == {"x-existing": "kept"} + (handshake,) = config.transform_realtime_request(_ga_session_update(), MUSE_MODEL) + + assert isinstance(handshake, str) + assert json.loads(handshake)["authorization"] == {"accessToken": "Bearer secret-token"} + assert config.is_setup_message(json.loads(handshake)) is True + assert config.is_setup_message({"type": "input_audio_buffer.append"}) is False + assert config.transform_realtime_request(_ga_session_update(), MUSE_MODEL) == () + + +def test_synthetic_session_created_uses_default_transcription_shape(): + created = MetaRealtimeConfig().transform_session_created_event(f"meta/{MUSE_MODEL}", "trace-1") + + assert created["type"] == "session.created" + assert created["session"]["id"] == "trace-1" + assert created["session"]["audio"]["input"]["format"] == {"type": "audio/pcm", "rate": 24000} + assert created["session"]["audio"]["input"]["transcription"] == {"model": MUSE_MODEL} + + +def test_audio_before_session_update_is_rejected(): + config = MetaRealtimeConfig() + config.validate_environment({}, MUSE_MODEL, api_key="secret-token") + + with pytest.raises(MuseProtocolError, match=r"session\.update must configure"): + config.transform_realtime_request(_event("input_audio_buffer.append", audio="AAAA"), MUSE_MODEL) + + +@pytest.mark.parametrize(("rate", "packet_bytes"), [(16_000, 2_560), (24_000, 3_840)]) +def test_pcm_is_packetized_into_raw_binary_frames(rate: int, packet_bytes: int): + config = _configured(rate=rate) + pcm = b"\xff\xfe\x00\x80" * (packet_bytes // 2) + b"\x01\x02\x03\x04" + + frames = config.transform_realtime_request( + _event("input_audio_buffer.append", audio=base64.b64encode(pcm).decode()), MUSE_MODEL + ) + remainder = config.transform_realtime_request(_event("input_audio_buffer.commit"), MUSE_MODEL) + + assert frames == (pcm[:packet_bytes], pcm[packet_bytes : packet_bytes * 2]) + assert remainder == (pcm[packet_bytes * 2 :],) + + +@pytest.mark.parametrize( + ("audio", "message"), + [ + ("not base64!", "valid base64"), + (base64.b64encode(b"\x00").decode(), "complete samples"), + (12, "base64 string"), + ("A" * (4 * ((24_000 * 2 * 4 + 2) // 3) + 4), "four-second backlog"), + ], +) +def test_invalid_audio_appends_are_rejected(audio: object, message: str): + config = _configured() + + with pytest.raises(MuseProtocolError, match=message): + config.transform_realtime_request(_event("input_audio_buffer.append", audio=audio), MUSE_MODEL) + + +@pytest.mark.asyncio +async def test_backend_sends_are_paced_to_real_time(): + sleeps: list[float] = [] + + async def record_sleep(delay: float) -> None: + sleeps.append(delay) + + config = _configured(monotonic=lambda: 10.0, sleep=record_sleep) + packet = b"\x01\x02" * 1_920 + + await config.pace_backend_send(packet) + await config.pace_backend_send(packet) + await config.pace_backend_send(packet) + + assert sleeps == pytest.approx([0.08, 0.16]) + + +def test_endpointing_commit_flushes_without_end_stream_but_end_sends_it_once(): + config = _configured(turn_detection="server_vad") + + assert config.transform_realtime_request(_event("input_audio_buffer.commit"), MUSE_MODEL) == () + assert config.transform_realtime_request(_event("input_audio_buffer.end"), MUSE_MODEL) == ('{"type":"endStream"}',) + assert config.transform_realtime_request(_event("input_audio_buffer.end"), MUSE_MODEL) == () + + +def test_push_to_talk_commit_ends_the_stream_once(): + config = _configured(turn_detection=None) + config.transform_realtime_request( + _event("input_audio_buffer.append", audio=base64.b64encode(b"\x01\x02").decode()), MUSE_MODEL + ) + + assert config.transform_realtime_request(_event("input_audio_buffer.commit"), MUSE_MODEL) == ( + b"\x01\x02", + '{"type":"endStream"}', + ) + assert config.transform_realtime_request(_event("input_audio_buffer.end"), MUSE_MODEL) == () + + +def test_clear_drops_buffered_remainder_and_unknown_events_are_ignored(): + config = _configured() + config.transform_realtime_request( + _event("input_audio_buffer.append", audio=base64.b64encode(b"\x01\x02").decode()), MUSE_MODEL + ) + + assert config.transform_realtime_request(_event("input_audio_buffer.clear"), MUSE_MODEL) == () + assert config.transform_realtime_request(_event("response.create"), MUSE_MODEL) == () + assert config.transform_realtime_request(_event("input_audio_buffer.commit"), MUSE_MODEL) == () + + +def test_provider_ack_becomes_session_created_with_provider_id(): + config = _configured(rate=16_000, turn_detection=None) + + (created,) = _backend_events(config, json.dumps({"sessionId": " provider-session "})) + + assert created["type"] == "session.created" + assert created["session"]["id"] == "provider-session" + assert created["session"]["audio"]["input"]["format"]["rate"] == 16000 + assert created["session"]["audio"]["input"]["turn_detection"] is None + + +def test_provider_turn_events_and_close_usage_flow_through_config(): + config = _configured() + + assert _backend_events(config, json.dumps({"type": "audioProgress", "audioProcessedMs": 1349})) == [] + assert _backend_events(config, _event("speechStart", turnId="t1"))[0]["type"] == "input_audio_buffer.speech_started" + assert _backend_events(config, _event("speechEnd", turnId="t1"))[0]["type"] == "input_audio_buffer.speech_stopped" + completed = _backend_events(config, _event("speechComplete", turnId="t1", transcript="what is the weather")) + + assert [event["type"] for event in completed] == ["conversation.item.input_audio_transcription.completed"] + assert completed[0]["usage"] == {"type": "duration", "seconds": 1.349} + assert config.unbilled_usage_on_session_close(MUSE_MODEL) is None + + assert _backend_events(config, json.dumps({"type": "audioProgress", "audioProcessedMs": 2349})) == [] + assert config.unbilled_usage_on_session_close(MUSE_MODEL) == {"type": "duration", "seconds": 1.0} + + +def test_provider_error_frame_becomes_openai_error_without_leaking_token(): + config = _configured() + + (error,) = _backend_events(config, _event("error", message="bad token secret-token")) + + assert error == { + "type": "error", + "error": {"type": "server_error", "message": "Meta Muse realtime transcription failed"}, + } + assert "secret-token" not in json.dumps(error) + + +def test_invalid_provider_ack_is_rejected(): + config = _configured() + + with pytest.raises(MuseProtocolError, match="invalid handshake response"): + _backend_events(config, json.dumps({"sessionId": ""})) diff --git a/tests/test_litellm/ocr/test_legacy.py b/tests/test_litellm/ocr/test_legacy.py index a30976f89af..4b0b78f5a0f 100644 --- a/tests/test_litellm/ocr/test_legacy.py +++ b/tests/test_litellm/ocr/test_legacy.py @@ -11,7 +11,8 @@ import pytest import litellm from litellm.integrations.custom_logger import CustomLogger -from litellm.llms.base_llm.ocr.transformation import OCRResponse +from litellm.litellm_core_utils.litellm_logging import Logging, use_custom_pricing_for_model +from litellm.llms.base_llm.ocr.transformation import OCRPage, OCRResponse, OCRUsageInfo from litellm.llms.custom_httpx import llm_http_handler from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.ocr.legacy import _prepare_ocr_request @@ -198,3 +199,61 @@ def test_generic_azure_connection_still_applies_to_foundry_ocr(monkeypatch: pyte assert prepared.api_key == "generic-key" assert prepared.api_base == "https://generic.example.com" + + +PRICING_OCR_MODEL: Final = "mistral/some-unmapped-ocr-model-for-testing" +PRICING_DOCUMENT: Final = {"type": "document_url", "document_url": "https://example.com/doc.pdf"} + + +def _pricing_logging_obj() -> Logging: + return Logging( + model=PRICING_OCR_MODEL, + messages=[], + stream=False, + call_type="ocr", + start_time=None, + litellm_call_id="test-ocr-request-pricing", + function_id="1234", + ) + + +def _prepare_with_pricing(kwargs: dict[str, object]) -> Logging: + logging_obj: Final = _pricing_logging_obj() + _prepare_ocr_request( + model=PRICING_OCR_MODEL, + document=dict(PRICING_DOCUMENT), + api_key="test-key", + api_base=None, + timeout=None, + custom_llm_provider=None, + extra_headers=None, + kwargs={"litellm_logging_obj": logging_obj, **kwargs}, + ) + return logging_obj + + +def test_prepare_ocr_request_forwards_custom_pricing_to_logging_params() -> None: + logging_obj: Final = _prepare_with_pricing({"ocr_cost_per_page": 0.05, "ocr_cost_per_credit": 0.5}) + + assert logging_obj.litellm_params["ocr_cost_per_page"] == 0.05 + assert logging_obj.litellm_params["ocr_cost_per_credit"] == 0.5 + assert use_custom_pricing_for_model(logging_obj.litellm_params) is True + + +def test_prepare_ocr_request_without_custom_pricing_leaves_logging_params_unpriced() -> None: + logging_obj: Final = _prepare_with_pricing({}) + + assert "ocr_cost_per_page" not in logging_obj.litellm_params + assert use_custom_pricing_for_model(logging_obj.litellm_params) is False + + +def test_direct_ocr_call_bills_request_level_per_page_pricing() -> None: + assert PRICING_OCR_MODEL not in litellm.model_cost + logging_obj: Final = _prepare_with_pricing({"ocr_cost_per_page": 0.05}) + response: Final = OCRResponse( + pages=[OCRPage(index=index, markdown=f"page {index}") for index in range(3)], + model=PRICING_OCR_MODEL, + usage_info=OCRUsageInfo(pages_processed=3), + ) + + assert logging_obj._response_cost_calculator(result=response) == pytest.approx(0.05 * 3) diff --git a/tests/test_litellm/proxy/common_utils/test_load_config_utils.py b/tests/test_litellm/proxy/common_utils/test_load_config_utils.py index 524c260e94a..1042dbe9653 100644 --- a/tests/test_litellm/proxy/common_utils/test_load_config_utils.py +++ b/tests/test_litellm/proxy/common_utils/test_load_config_utils.py @@ -1,9 +1,18 @@ +import asyncio +import logging +import re +import threading from unittest.mock import MagicMock, mock_open, patch import pytest import yaml -from litellm.proxy.common_utils.load_config_utils import get_file_contents_from_s3 +from litellm.proxy.common_utils.load_config_utils import ( + gcs_config_bucket, + get_config_from_bucket, + get_file_contents_from_s3, + resolve_bucket_includes, +) class TestGetFileContentsFromS3: @@ -83,3 +92,385 @@ class TestGetFileContentsFromS3: # Verify yaml.safe_load was called with the decoded content mock_yaml_load.assert_called_once_with(yaml_content) + + +class TestBucketConfigIncludes: + + @staticmethod + def _bucket(objects): + async def fetch(object_key): + return objects.get(object_key) + + return fetch + + @pytest.mark.asyncio + async def test_include_resolves_against_the_config_objects_prefix(self): + merged = await resolve_bucket_includes( + config={"include": ["model_config.yaml"], "general_settings": {"master_key": "sk-1234"}}, + object_key="configs/prod/config.yaml", + fetch=self._bucket( + {"configs/prod/model_config.yaml": {"model_list": [{"model_name": "gpt-4o-mini"}]}} + ), + ) + + assert merged == { + "general_settings": {"master_key": "sk-1234"}, + "model_list": [{"model_name": "gpt-4o-mini"}], + } + + @pytest.mark.asyncio + async def test_include_with_a_leading_slash_reads_from_the_bucket_root(self): + merged = await resolve_bucket_includes( + config={"include": ["/shared/models.yaml"]}, + object_key="configs/prod/config.yaml", + fetch=self._bucket({"shared/models.yaml": {"model_list": [{"model_name": "shared"}]}}), + ) + + assert merged == {"model_list": [{"model_name": "shared"}]} + + @pytest.mark.asyncio + async def test_include_walks_out_of_the_prefix_with_dot_dot(self): + merged = await resolve_bucket_includes( + config={"include": ["../shared/models.yaml"]}, + object_key="configs/prod/config.yaml", + fetch=self._bucket({"configs/shared/models.yaml": {"model_list": [{"model_name": "shared"}]}}), + ) + + assert merged == {"model_list": [{"model_name": "shared"}]} + + @pytest.mark.asyncio + async def test_included_configs_may_declare_further_includes(self): + merged = await resolve_bucket_includes( + config={"include": ["models.yaml"]}, + object_key="configs/config.yaml", + fetch=self._bucket( + { + "configs/models.yaml": { + "include": ["extra/more_models.yaml"], + "model_list": [{"model_name": "first"}], + }, + "configs/extra/more_models.yaml": {"model_list": [{"model_name": "second"}]}, + } + ), + ) + + assert merged == {"model_list": [{"model_name": "first"}, {"model_name": "second"}]} + + @pytest.mark.asyncio + async def test_a_nested_include_resolves_against_the_object_that_declares_it(self): + merged = await resolve_bucket_includes( + config={"include": ["shared/models.yaml"]}, + object_key="configs/config.yaml", + fetch=self._bucket( + { + "configs/shared/models.yaml": { + "include": ["more_models.yaml"], + "model_list": [{"model_name": "first"}], + }, + "configs/shared/more_models.yaml": {"model_list": [{"model_name": "second"}]}, + "configs/more_models.yaml": {"model_list": [{"model_name": "wrong-prefix"}]}, + } + ), + ) + + assert merged == {"model_list": [{"model_name": "first"}, {"model_name": "second"}]} + + @pytest.mark.asyncio + async def test_an_object_pulled_in_twice_is_merged_once(self): + merged = await resolve_bucket_includes( + config={"include": ["a.yaml", "b.yaml"]}, + object_key="configs/config.yaml", + fetch=self._bucket( + { + "configs/a.yaml": {"include": ["shared.yaml"]}, + "configs/b.yaml": {"include": ["./shared.yaml"]}, + "configs/shared.yaml": {"model_list": [{"model_name": "shared"}]}, + } + ), + ) + + assert merged == {"model_list": [{"model_name": "shared"}]} + + @pytest.mark.asyncio + async def test_a_cycle_between_included_objects_terminates(self): + merged = await asyncio.wait_for( + resolve_bucket_includes( + config={"include": ["a.yaml"]}, + object_key="configs/config.yaml", + fetch=self._bucket( + { + "configs/a.yaml": {"include": ["b.yaml"], "model_list": [{"model_name": "from-a"}]}, + "configs/b.yaml": {"include": ["a.yaml"], "model_list": [{"model_name": "from-b"}]}, + } + ), + ), + timeout=10, + ) + + assert merged == {"model_list": [{"model_name": "from-a"}, {"model_name": "from-b"}]} + + @pytest.mark.asyncio + async def test_list_values_are_extended_and_other_values_are_overridden(self): + merged = await resolve_bucket_includes( + config={ + "include": ["models.yaml"], + "model_list": [{"model_name": "from-root"}], + "litellm_settings": {"drop_params": True}, + }, + object_key="config.yaml", + fetch=self._bucket( + { + "models.yaml": { + "model_list": [{"model_name": "from-include"}], + "litellm_settings": {"num_retries": 3}, + } + } + ), + ) + + assert merged == { + "model_list": [{"model_name": "from-root"}, {"model_name": "from-include"}], + "litellm_settings": {"num_retries": 3}, + } + + @pytest.mark.asyncio + async def test_a_missing_included_object_fails_loudly_with_its_key(self): + with pytest.raises(FileNotFoundError, match=re.escape("configs/prod/model_config.yaml")): + await resolve_bucket_includes( + config={"include": ["model_config.yaml"]}, + object_key="configs/prod/config.yaml", + fetch=self._bucket({}), + ) + + @pytest.mark.asyncio + async def test_a_non_list_include_fails_loudly(self): + with pytest.raises(ValueError, match="'include' must be a list of file paths"): + await resolve_bucket_includes( + config={"include": "model_config.yaml"}, + object_key="config.yaml", + fetch=self._bucket({}), + ) + + @pytest.mark.asyncio + async def test_get_config_from_bucket_merges_includes_over_s3(self, monkeypatch): + objects = { + "lit6982/config.yaml": { + "include": ["model_config.yaml"], + "general_settings": {"master_key": "sk-1234"}, + }, + "lit6982/model_config.yaml": {"model_list": [{"model_name": "included-model"}]}, + } + monkeypatch.setattr( + "litellm.proxy.common_utils.load_config_utils.s3_object_reader", + lambda bucket_name: objects.get, + ) + + config = await get_config_from_bucket( + bucket_type="s3", bucket_name="litellm-configs", object_key="lit6982/config.yaml" + ) + + assert config == { + "general_settings": {"master_key": "sk-1234"}, + "model_list": [{"model_name": "included-model"}], + } + + @pytest.mark.asyncio + async def test_the_blocking_s3_work_runs_off_the_event_loop_thread(self, monkeypatch): + loop_thread = threading.current_thread() + threads = [] + + def build_reader(bucket_name): + threads.append(threading.current_thread()) + + def read(object_key): + threads.append(threading.current_thread()) + return {"model_list": [{"model_name": "a-model"}]} + + return read + + monkeypatch.setattr("litellm.proxy.common_utils.load_config_utils.s3_object_reader", build_reader) + + await get_config_from_bucket(bucket_type="s3", bucket_name="litellm-configs", object_key="config.yaml") + + assert len(threads) == 2 and loop_thread not in threads + + @pytest.mark.asyncio + async def test_one_s3_client_serves_the_whole_include_tree(self, monkeypatch): + objects = { + "lit6982/config.yaml": {"include": ["model_config.yaml"]}, + "lit6982/model_config.yaml": {"model_list": [{"model_name": "included-model"}]}, + } + readers = [] + + def build_reader(bucket_name): + requested = [] + readers.append(requested) + + def read(object_key): + requested.append(object_key) + return objects.get(object_key) + + return read + + monkeypatch.setattr("litellm.proxy.common_utils.load_config_utils.s3_object_reader", build_reader) + + await get_config_from_bucket( + bucket_type="s3", bucket_name="litellm-configs", object_key="lit6982/config.yaml" + ) + + assert readers == [["lit6982/config.yaml", "lit6982/model_config.yaml"]] + + @pytest.mark.asyncio + async def test_an_empty_included_object_merges_as_an_empty_config(self, monkeypatch): + objects = { + "lit6982/config.yaml": "include:\n - empty.yaml\nmodel_list:\n - model_name: only-model\n", + "lit6982/empty.yaml": "", + } + + class FakeGCSBucket: + async def download_gcs_object(self, object_key): + return objects[object_key].encode("utf-8") + + monkeypatch.setattr( + "litellm.proxy.common_utils.load_config_utils.gcs_config_bucket", + lambda bucket_name: FakeGCSBucket(), + ) + + config = await get_config_from_bucket( + bucket_type="gcs", bucket_name="litellm-configs", object_key="lit6982/config.yaml" + ) + + assert config == {"model_list": [{"model_name": "only-model"}]} + + @pytest.mark.asyncio + async def test_get_config_from_bucket_merges_includes_over_gcs(self, monkeypatch): + objects = { + "lit6982/config.yaml": { + "include": ["model_config.yaml"], + "general_settings": {"master_key": "sk-1234"}, + }, + "lit6982/model_config.yaml": {"model_list": [{"model_name": "included-model"}]}, + } + + buckets = [] + + class FakeGCSBucket: + def __init__(self): + self.requested = [] + buckets.append(self) + + async def download_gcs_object(self, object_key): + self.requested.append(object_key) + return yaml.dump(objects[object_key]).encode("utf-8") + + monkeypatch.setattr( + "litellm.proxy.common_utils.load_config_utils.gcs_config_bucket", + lambda bucket_name: FakeGCSBucket(), + ) + + config = await get_config_from_bucket( + bucket_type="gcs", bucket_name="litellm-configs", object_key="lit6982/config.yaml" + ) + + assert config == { + "general_settings": {"master_key": "sk-1234"}, + "model_list": [{"model_name": "included-model"}], + } + assert [bucket.requested for bucket in buckets] == [ + ["lit6982/config.yaml", "lit6982/model_config.yaml"] + ] + + @pytest.mark.asyncio + async def test_get_config_from_bucket_returns_none_when_the_root_object_is_missing(self, monkeypatch): + monkeypatch.setattr( + "litellm.proxy.common_utils.load_config_utils.s3_object_reader", + lambda bucket_name: (lambda object_key: None), + ) + + assert ( + await get_config_from_bucket( + bucket_type="s3", bucket_name="litellm-configs", object_key="missing.yaml" + ) + is None + ) + + @pytest.mark.asyncio + async def test_an_object_pulled_in_twice_is_read_once(self): + objects = { + "configs/a.yaml": {"include": ["shared.yaml"]}, + "configs/b.yaml": {"include": ["./shared.yaml"]}, + "configs/shared.yaml": {"model_list": [{"model_name": "shared"}]}, + } + requested = [] + + async def fetch(object_key): + requested.append(object_key) + return objects.get(object_key) + + await resolve_bucket_includes( + config={"include": ["a.yaml", "b.yaml"]}, + object_key="configs/config.yaml", + fetch=fetch, + ) + + assert requested == ["configs/a.yaml", "configs/b.yaml", "configs/shared.yaml"] + + @pytest.mark.asyncio + async def test_an_empty_root_object_does_not_boot_an_empty_proxy(self, monkeypatch): + class FakeGCSBucket: + async def download_gcs_object(self, object_key): + return b"" + + monkeypatch.setattr( + "litellm.proxy.common_utils.load_config_utils.gcs_config_bucket", + lambda bucket_name: FakeGCSBucket(), + ) + + config = await get_config_from_bucket( + bucket_type="gcs", bucket_name="litellm-configs", object_key="lit6982/config.yaml" + ) + + assert config is None + + @pytest.mark.asyncio + async def test_an_object_that_is_not_valid_yaml_is_reported_as_a_yaml_error(self, monkeypatch, caplog): + class FakeGCSBucket: + async def download_gcs_object(self, object_key): + return b"model_list: [\n" + + monkeypatch.setattr( + "litellm.proxy.common_utils.load_config_utils.gcs_config_bucket", + lambda bucket_name: FakeGCSBucket(), + ) + + with caplog.at_level(logging.ERROR, logger="LiteLLM Proxy"): + config = await get_config_from_bucket( + bucket_type="gcs", bucket_name="litellm-configs", object_key="lit6982/config.yaml" + ) + + assert config is None + assert [ + record + for record in caplog.records + if "not valid YAML" in record.getMessage() and "lit6982/config.yaml" in record.getMessage() + ] + + +class TestGCSConfigBucketClient: + @pytest.mark.asyncio + async def test_reading_a_config_from_gcs_does_not_need_an_enterprise_license(self, monkeypatch): + monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", False) + + bucket = gcs_config_bucket("litellm-configs") + + assert bucket is not None + assert bucket.BUCKET_NAME == "litellm-configs" + + @pytest.mark.asyncio + async def test_reading_a_config_from_gcs_starts_no_background_task(self, monkeypatch): + monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True) + running_before = asyncio.all_tasks() + + gcs_config_bucket("litellm-configs") + + assert asyncio.all_tasks() - running_before == set() diff --git a/tests/test_litellm/proxy/common_utils/test_scheduled_job_stagger.py b/tests/test_litellm/proxy/common_utils/test_scheduled_job_stagger.py index 54e0aa74a25..3eabfc5c840 100644 --- a/tests/test_litellm/proxy/common_utils/test_scheduled_job_stagger.py +++ b/tests/test_litellm/proxy/common_utils/test_scheduled_job_stagger.py @@ -64,6 +64,10 @@ def _stagger(scheduler: AsyncIOScheduler, identity: str = "pod-a:1", **overrides return apply_scheduled_job_stagger(scheduler=scheduler, settings=_settings(**overrides), identity=identity) +def _trigger_of(scheduler: AsyncIOScheduler, job_id: str): + return next(job.trigger for job in scheduler.get_jobs() if job.id == job_id) + + def _fire_times(trigger, start: datetime, steps: int) -> tuple[datetime, ...]: """The fire times APScheduler would produce, each computed from the one before it""" return tuple( @@ -133,22 +137,24 @@ def test_default_cron_is_staggered_and_keeps_its_offset_on_every_later_fire(): applied = _stagger(scheduler) assert applied[PTU_ROLLUP_JOB_ID] > 0 - trigger = next(job.trigger for job in scheduler.get_jobs() if job.id == PTU_ROLLUP_JOB_ID) - fires = _fire_times(trigger, datetime(2026, 1, 1, 12, 0, tzinfo=timezone.utc), 3) + start = datetime(2026, 1, 1, 12, 0, tzinfo=timezone.utc) + fires = _fire_times(_trigger_of(scheduler, PTU_ROLLUP_JOB_ID), start, 3) expected = timedelta(minutes=15) + timedelta(seconds=applied[PTU_ROLLUP_JOB_ID]) assert [fire - fire.replace(hour=0, minute=0, second=0, microsecond=0) for fire in fires] == [expected] * 3 -async def test_explicit_offset_overrides_the_derived_one_and_zero_pins_a_job(): +def test_explicit_offset_overrides_the_derived_one_and_zero_pins_a_job(): scheduler = _with_jobs(_scheduler()) applied = _stagger(scheduler, offsets={"periodic_reload_job": 0, PTU_ROLLUP_JOB_ID: 7}) - unstaggered = _next_run_times(_with_jobs(_scheduler())) - staggered = _next_run_times(scheduler) assert applied["periodic_reload_job"] == 0 assert applied[PTU_ROLLUP_JOB_ID] == 7 - assert staggered[PTU_ROLLUP_JOB_ID] - unstaggered[PTU_ROLLUP_JOB_ID] == timedelta(seconds=7) + + start = datetime(2026, 1, 1, 12, 0, tzinfo=timezone.utc) + staggered = _trigger_of(scheduler, PTU_ROLLUP_JOB_ID) + unstaggered = _trigger_of(_with_jobs(_scheduler()), PTU_ROLLUP_JOB_ID) + assert _fire_times(staggered, start, 1)[0] - _fire_times(unstaggered, start, 1)[0] == timedelta(seconds=7) async def test_disabling_the_stagger_leaves_every_schedule_untouched(): diff --git a/tests/test_litellm/proxy/db/conftest.py b/tests/test_litellm/proxy/db/conftest.py index d3226b0ec50..bcb7794a20d 100644 --- a/tests/test_litellm/proxy/db/conftest.py +++ b/tests/test_litellm/proxy/db/conftest.py @@ -35,6 +35,14 @@ DB_ENV_KEYS = ( _db_env_snapshot_key = pytest.StashKey[dict[str, Optional[str]]]() +def _is_zombie(pid: int) -> bool: + try: + stat: Final = Path(f"/proc/{pid}/stat").read_text() + except OSError: + return False + return stat.rpartition(")")[2].split()[0] == "Z" + + def _db_env_snapshot() -> dict[str, Optional[str]]: return {key: os.environ.get(key) for key in DB_ENV_KEYS} @@ -136,6 +144,8 @@ class FakePrismaCli: os.kill(pid, 0) except ProcessLookupError: return True + if _is_zombie(pid): + return True time.sleep(0.05) return False diff --git a/tests/test_litellm/proxy/discovery_endpoints/test_agent_skills_archive.py b/tests/test_litellm/proxy/discovery_endpoints/test_agent_skills_archive.py new file mode 100644 index 00000000000..dd5ac230aac --- /dev/null +++ b/tests/test_litellm/proxy/discovery_endpoints/test_agent_skills_archive.py @@ -0,0 +1,106 @@ +import hashlib +import io +import zipfile + +from litellm.proxy.discovery_endpoints.agent_skills_archive import ( + MAX_ARCHIVE_ENTRIES, + build_skill_archive, +) + +MANIFEST = b"""--- +name: pdf-summarizer +description: Summarize a PDF into an executive brief. +--- + +Read the PDF, then write the brief. +""" + + +def zip_bytes(files: dict[str, bytes]) -> bytes: + buffer = io.BytesIO() + with zipfile.ZipFile(buffer, "w", zipfile.ZIP_DEFLATED) as archive: + for name, content in files.items(): + archive.writestr(name, content) + return buffer.getvalue() + + +def entries_of(content: bytes) -> dict[str, bytes]: + with zipfile.ZipFile(io.BytesIO(content)) as archive: + return {name: archive.read(name) for name in archive.namelist()} + + +def test_single_top_level_folder_is_stripped_so_skill_md_sits_at_the_root(): + archive = build_skill_archive( + zip_bytes( + { + "pdf-summarizer/SKILL.md": MANIFEST, + "pdf-summarizer/reference.md": b"page citations", + "pdf-summarizer/scripts/extract.py": b"print('hi')", + } + ) + ) + + assert archive is not None + assert entries_of(archive.content) == { + "SKILL.md": MANIFEST, + "reference.md": b"page citations", + "scripts/extract.py": b"print('hi')", + } + + +def test_digest_covers_the_repacked_bytes_and_is_stable_across_builds(): + upload = zip_bytes({"pdf-summarizer/SKILL.md": MANIFEST, "pdf-summarizer/reference.md": b"page citations"}) + + first = build_skill_archive(upload) + second = build_skill_archive(upload) + + assert first is not None and second is not None + assert first.digest == f"sha256:{hashlib.sha256(first.content).hexdigest()}" + assert first.content == second.content + + +def test_an_upload_that_is_already_flat_keeps_every_file_where_it_is(): + archive = build_skill_archive(zip_bytes({"SKILL.md": MANIFEST, "reference.md": b"page citations"})) + + assert archive is not None + assert sorted(entries_of(archive.content)) == ["SKILL.md", "reference.md"] + + +def test_manifest_frontmatter_supplies_the_declared_name_and_description(): + archive = build_skill_archive(zip_bytes({"pdf-summarizer/SKILL.md": MANIFEST})) + + assert archive is not None + assert archive.declared_name == "pdf-summarizer" + assert archive.declared_description == "Summarize a PDF into an executive brief." + + +def test_a_manifest_without_frontmatter_declares_nothing(): + archive = build_skill_archive(zip_bytes({"pdf-summarizer/SKILL.md": b"just prose, no frontmatter"})) + + assert archive is not None + assert archive.declared_name is None + assert archive.declared_description is None + + +def test_a_manifest_buried_below_the_stripped_folder_is_not_installable(): + assert build_skill_archive(zip_bytes({"pdf-summarizer/nested/SKILL.md": MANIFEST})) is None + + +def test_an_upload_with_no_manifest_is_not_installable(): + assert build_skill_archive(zip_bytes({"pdf-summarizer/reference.md": b"page citations"})) is None + + +def test_a_non_zip_upload_is_not_installable(): + assert build_skill_archive(MANIFEST) is None + + +def test_a_path_traversal_entry_is_not_installable(): + assert build_skill_archive(zip_bytes({"SKILL.md": MANIFEST, "../escape.md": b"nope"})) is None + + +def test_an_upload_over_the_entry_cap_is_not_installable(): + files = {"pdf-summarizer/SKILL.md": MANIFEST} | { + f"pdf-summarizer/file-{index}.md": b"x" for index in range(MAX_ARCHIVE_ENTRIES) + } + + assert build_skill_archive(zip_bytes(files)) is None diff --git a/tests/test_litellm/proxy/discovery_endpoints/test_agent_skills_endpoints.py b/tests/test_litellm/proxy/discovery_endpoints/test_agent_skills_endpoints.py new file mode 100644 index 00000000000..ea889e75ae8 --- /dev/null +++ b/tests/test_litellm/proxy/discovery_endpoints/test_agent_skills_endpoints.py @@ -0,0 +1,211 @@ +import hashlib +import io +import zipfile +from datetime import datetime, timezone + +import pytest +from fastapi import FastAPI +from fastapi.testclient import TestClient + +import litellm +from litellm.models.skills import LiteLLM_SkillsTable +from litellm.proxy.discovery_endpoints.agent_skills_endpoints import ( + router, + stored_skill, + stored_skills, +) +from litellm.types.proxy.discovery_endpoints.agent_skills_endpoints import ( + AGENT_SKILLS_DISCOVERY_SCHEMA_URL, +) + +WELL_KNOWN_PATHS = ("/.well-known/agent-skills/index.json", "/.well-known/skills/index.json") + +MANIFEST = b"""--- +name: pdf-summarizer +description: Summarize a PDF into an executive brief. +--- + +Read the PDF, then write the brief. +""" + + +def zip_bytes(files: dict[str, bytes]) -> bytes: + buffer = io.BytesIO() + with zipfile.ZipFile(buffer, "w", zipfile.ZIP_DEFLATED) as archive: + for name, content in files.items(): + archive.writestr(name, content) + return buffer.getvalue() + + +def skill( + skill_id: str, + *, + display_title: str | None = "PDF Summarizer", + description: str | None = None, + files: dict[str, bytes] | None = None, + updated_at: datetime | None = None, +) -> LiteLLM_SkillsTable: + return LiteLLM_SkillsTable( + skill_id=skill_id, + display_title=display_title, + description=description, + file_content=zip_bytes(files if files is not None else {"pdf-summarizer/SKILL.md": MANIFEST}), + updated_at=updated_at, + ) + + +def client_for(*skills: LiteLLM_SkillsTable) -> TestClient: + app = FastAPI() + app.include_router(router) + + def _skills() -> tuple[LiteLLM_SkillsTable, ...]: + return skills + + def _skill(skill_id: str) -> LiteLLM_SkillsTable | None: + return next((candidate for candidate in skills if candidate.skill_id == skill_id), None) + + app.dependency_overrides[stored_skills] = _skills + app.dependency_overrides[stored_skill] = _skill + return TestClient(app) + + +@pytest.fixture +def index_enabled(monkeypatch): + monkeypatch.setattr(litellm, "public_skills_index", True) + + +def test_discovery_is_absent_until_public_skills_index_is_enabled(monkeypatch): + monkeypatch.setattr(litellm, "public_skills_index", False) + client = client_for(skill("litellm_skill_1")) + + for path in WELL_KNOWN_PATHS: + assert client.get(path).status_code == 404 + assert client.get("/v1/skills/litellm_skill_1/archive").status_code == 404 + + +@pytest.mark.parametrize("path", WELL_KNOWN_PATHS) +def test_index_publishes_each_stored_skill_in_the_v0_2_0_shape(index_enabled, path): + client = client_for(skill("litellm_skill_1")) + + body = client.get(path).json() + + assert body["$schema"] == AGENT_SKILLS_DISCOVERY_SCHEMA_URL + assert len(body["skills"]) == 1 + entry = body["skills"][0] + assert entry["name"] == "pdf-summarizer" + assert entry["type"] == "archive" + assert entry["description"] == "Summarize a PDF into an executive brief." + assert entry["url"].endswith("/v1/skills/litellm_skill_1/archive") + assert entry["digest"].startswith("sha256:") + + +def test_index_digest_matches_the_bytes_the_archive_route_serves(index_enabled): + client = client_for(skill("litellm_skill_1")) + + entry = client.get(WELL_KNOWN_PATHS[0]).json()["skills"][0] + downloaded = client.get(entry["url"]) + + assert downloaded.status_code == 200 + assert downloaded.headers["content-type"] == "application/zip" + assert entry["digest"] == f"sha256:{hashlib.sha256(downloaded.content).hexdigest()}" + + +def test_install_name_falls_back_to_the_manifest_name_without_a_display_title(index_enabled): + client = client_for(skill("litellm_skill_1", display_title=None)) + + assert client.get(WELL_KNOWN_PATHS[0]).json()["skills"][0]["name"] == "pdf-summarizer" + + +@pytest.mark.parametrize( + "manifest, stored_description, expected", + [ + (MANIFEST, "registry copy", "Summarize a PDF into an executive brief."), + (b"no frontmatter here", "registry copy", "registry copy"), + (b"no frontmatter here", None, "PDF Summarizer"), + ], +) +def test_description_prefers_the_manifest_then_the_registry_then_the_title( + index_enabled, manifest, stored_description, expected +): + client = client_for( + skill( + "litellm_skill_1", + description=stored_description, + files={"pdf-summarizer/SKILL.md": manifest}, + ) + ) + + assert client.get(WELL_KNOWN_PATHS[0]).json()["skills"][0]["description"] == expected + + +def test_skills_sharing_a_title_get_distinct_install_names(index_enabled): + client = client_for( + skill("litellm_skill_2", files={"pdf-summarizer/SKILL.md": b"second"}), + skill("litellm_skill_1", files={"pdf-summarizer/SKILL.md": b"first"}), + ) + + names = [entry["name"] for entry in client.get(WELL_KNOWN_PATHS[0]).json()["skills"]] + + assert names == ["pdf-summarizer", "pdf-summarizer-2"] + + +def test_uploads_without_a_root_manifest_are_left_out_of_the_index(index_enabled): + client = client_for( + skill("litellm_skill_1"), + skill("litellm_skill_2", files={"pdf-summarizer/reference.md": b"no manifest"}), + ) + + body = client.get(WELL_KNOWN_PATHS[0]).json() + + assert [entry["url"].split("/")[-2] for entry in body["skills"]] == ["litellm_skill_1"] + assert client.get("/v1/skills/litellm_skill_2/archive").status_code == 404 + + +def test_archive_route_404s_for_a_skill_that_does_not_exist(index_enabled): + client = client_for(skill("litellm_skill_1")) + + assert client.get("/v1/skills/litellm_skill_missing/archive").status_code == 404 + + +def test_a_stored_skill_is_repacked_once_per_version(index_enabled): + stamp = datetime(2026, 9, 6, 9, 0, tzinfo=timezone.utc) + first = client_for(skill("litellm_skill_cached", files={"s/SKILL.md": MANIFEST}, updated_at=stamp)) + published = first.get(WELL_KNOWN_PATHS[0]).json()["skills"][0]["digest"] + + unchanged_row = client_for( + skill("litellm_skill_cached", files={"s/SKILL.md": MANIFEST, "s/extra.md": b"rewritten"}, updated_at=stamp) + ) + + assert unchanged_row.get(WELL_KNOWN_PATHS[0]).json()["skills"][0]["digest"] == published + assert hashlib.sha256(unchanged_row.get("/v1/skills/litellm_skill_cached/archive").content).hexdigest() == ( + published.removeprefix("sha256:") + ) + + +def test_a_skill_edited_since_the_last_read_is_republished(index_enabled): + stamp = datetime(2026, 9, 6, 9, 0, tzinfo=timezone.utc) + before = client_for(skill("litellm_skill_edited", files={"s/SKILL.md": MANIFEST}, updated_at=stamp)) + published = before.get(WELL_KNOWN_PATHS[0]).json()["skills"][0]["digest"] + + after = client_for( + skill( + "litellm_skill_edited", + files={"s/SKILL.md": MANIFEST, "s/extra.md": b"rewritten"}, + updated_at=datetime(2026, 9, 6, 10, 0, tzinfo=timezone.utc), + ) + ) + republished = after.get(WELL_KNOWN_PATHS[0]).json()["skills"][0]["digest"] + + assert republished != published + assert hashlib.sha256(after.get("/v1/skills/litellm_skill_edited/archive").content).hexdigest() == ( + republished.removeprefix("sha256:") + ) + + +def test_openapi_declares_the_archive_route_as_a_zip_download(index_enabled): + schema = client_for(skill("litellm_skill_1")).get("/openapi.json").json() + + content = schema["paths"]["/v1/skills/{skill_id}/archive"]["get"]["responses"]["200"]["content"] + + assert "application/zip" in content + assert "application/json" not in content diff --git a/tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py b/tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py index 11c3d2f8b20..c550a0a41d2 100644 --- a/tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py +++ b/tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py @@ -15,6 +15,7 @@ Streaming: CSW.__anext__ stores args on logging_obj at stream end. """ import asyncio +import logging from typing import Any, Final from unittest.mock import AsyncMock, MagicMock, patch @@ -1422,7 +1423,7 @@ class TestArmDeferredStreamDispatch: async def test_native_stream_closure_enqueues_single_coroutine(self): from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER - logging_obj, _ = self._dispatch_recording_logging_obj() + logging_obj, recorded = self._dispatch_recording_logging_obj() async def _agen(): yield b"x" @@ -1433,20 +1434,89 @@ class TestArmDeferredStreamDispatch: user_api_key_dict=MagicMock(), logging_obj=logging_obj, ) - closure = logging_obj._on_deferred_stream_complete - assert closure is not None + assert logging_obj._on_deferred_stream_complete is not None async def _logging_coroutine(): return None coro = _logging_coroutine() + logging_obj._deferred_stream_complete_args = (coro,) with patch.object( # test-quality-ok: GLOBAL_LOGGING_WORKER is a process-global singleton with no injection seam GLOBAL_LOGGING_WORKER, "ensure_initialized_and_enqueue" ) as mock_enqueue: - await closure(coro) + ProxyLogging._fire_deferred_stream_logging({"litellm_logging_obj": logging_obj}) + await asyncio.sleep(0) mock_enqueue.assert_called_once_with(async_coroutine=coro) + assert recorded == {} coro.close() + @pytest.mark.asyncio + @pytest.mark.parametrize("route_type", ["anthropic_messages", "aresponses"]) + async def test_raw_generator_stream_storing_csw_arg_shape_dispatches_success(self, route_type): + """Bridged /v1/messages returns AnthropicStreamWrapper's plain SSE + generator, which shares its inner CustomStreamWrapper's logging_obj and + so stores (assembled_response, cache_hit). The closure armed for a raw + generator must accept that shape too, or _fire_deferred_stream_logging + raises TypeError and the request loses its spend log and callbacks.""" + from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER + + logging_obj, recorded = self._dispatch_recording_logging_obj() + + async def _agen(): + yield b"x" + + self._processor()._arm_deferred_stream_dispatch( + response=_agen(), + route_type=route_type, + user_api_key_dict=MagicMock(), + logging_obj=logging_obj, + ) + + assembled = object() + logging_obj._deferred_stream_complete_args = (assembled, True) + with patch.object( # test-quality-ok: GLOBAL_LOGGING_WORKER is a process-global singleton with no injection seam + GLOBAL_LOGGING_WORKER, "ensure_initialized_and_enqueue" + ) as mock_enqueue: + ProxyLogging._fire_deferred_stream_logging({"litellm_logging_obj": logging_obj}) + await asyncio.sleep(0) + + mock_enqueue.assert_not_called() + assert recorded["result"] is assembled + assert recorded["cache_hit"] is True + assert recorded["prefer_async_handlers"] is True + + @pytest.mark.asyncio + @pytest.mark.parametrize("stored_args", [(object(),), (object(), object(), object())]) + async def test_raw_generator_stream_with_unknown_arg_shape_logs_and_drops(self, stored_args, caplog): + from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER + + logging_obj, recorded = self._dispatch_recording_logging_obj() + + async def _agen(): + yield b"x" + + self._processor()._arm_deferred_stream_dispatch( + response=_agen(), + route_type="anthropic_messages", + user_api_key_dict=MagicMock(), + logging_obj=logging_obj, + ) + + logging_obj._deferred_stream_complete_args = stored_args + with ( + patch.object( # test-quality-ok: GLOBAL_LOGGING_WORKER is a process-global singleton with no injection seam + GLOBAL_LOGGING_WORKER, "ensure_initialized_and_enqueue" + ) as mock_enqueue, + caplog.at_level(logging.ERROR, logger="LiteLLM Proxy"), + ): + ProxyLogging._fire_deferred_stream_logging({"litellm_logging_obj": logging_obj}) + await asyncio.sleep(0) + + mock_enqueue.assert_not_called() + assert recorded == {} + dropped = [r for r in caplog.records if r.getMessage().startswith("Deferred stream logging dropped")] + assert len(dropped) == 1 + @pytest.mark.asyncio async def test_csw_closure_routes_through_deferred_stream_guardrails(self, monkeypatch): from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper diff --git a/tests/test_litellm/proxy/management_endpoints/test_budget_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_budget_endpoints.py index 79d62f772bd..4b6815d7552 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_budget_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_budget_endpoints.py @@ -340,7 +340,8 @@ async def test_update_budget_recomputes_reset_at_when_duration_changes( @pytest.mark.asyncio -async def test_update_budget_preserves_explicit_reset_at(client_and_mocks): +@pytest.mark.parametrize("budget_duration", ["1d", None]) +async def test_update_budget_preserves_explicit_reset_at(client_and_mocks, budget_duration): """An explicit budget_reset_at from the caller always wins over recompute.""" client, _, mock_table = client_and_mocks captured = _capture_update_data(mock_table) @@ -350,7 +351,7 @@ async def test_update_budget_preserves_explicit_reset_at(client_and_mocks): "/budget/update", json={ "budget_id": "budget_explicit_reset", - "budget_duration": "1d", + "budget_duration": budget_duration, "budget_reset_at": explicit.isoformat(), }, ) @@ -377,8 +378,7 @@ async def test_update_budget_without_duration_leaves_reset_at_untouched( @pytest.mark.asyncio -async def test_update_budget_duration_none_does_not_recompute(client_and_mocks): - """Clearing budget_duration (explicit null) must not recompute against a None duration.""" +async def test_update_budget_duration_none_clears_obsolete_reset(client_and_mocks): client, _, mock_table = client_and_mocks captured = _capture_update_data(mock_table) @@ -389,7 +389,7 @@ async def test_update_budget_duration_none_does_not_recompute(client_and_mocks): assert resp.status_code == 200, resp.text assert "budget_duration" in captured and captured["budget_duration"] is None - assert "budget_reset_at" not in captured + assert captured["budget_reset_at"] is None @pytest.mark.asyncio diff --git a/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py b/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py index 6cd900cb041..71896a18f48 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py +++ b/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py @@ -613,6 +613,7 @@ def test_key_metadata_includes_recovered_user_email(): "dirty-key": { "key_alias": "batch-worker", "team_id": "team-1", + "user_id": "alice", "user_email": "alice@example.com", } }, @@ -620,6 +621,7 @@ def test_key_metadata_includes_recovered_user_email(): ) assert meta.key_alias == "batch-worker" + assert meta.user_id == "alice" assert meta.user_email == "alice@example.com" @@ -848,9 +850,11 @@ async def test_aggregated_activity_preserves_metadata_for_deleted_keys(): mock_deleted_key.token = "deleted-key-hash" mock_deleted_key.key_alias = "toto-test-2" mock_deleted_key.team_id = "69cd4b77-b095-4489-8c46-4f2f31d840a2" + mock_deleted_key.user_id = "deleted-key-owner" mock_prisma.db.litellm_deletedverificationtoken = MagicMock() mock_prisma.db.litellm_deletedverificationtoken.find_many = AsyncMock(return_value=[mock_deleted_key]) + mock_prisma.db.litellm_usertable.find_many = AsyncMock(return_value=[]) result = await get_daily_activity_aggregated( prisma_client=mock_prisma, @@ -871,6 +875,7 @@ async def test_aggregated_activity_preserves_metadata_for_deleted_keys(): key_data = chat_endpoint.api_key_breakdown["deleted-key-hash"] assert key_data.metadata.key_alias == "toto-test-2" assert key_data.metadata.team_id == "69cd4b77-b095-4489-8c46-4f2f31d840a2" + assert key_data.metadata.user_id == "deleted-key-owner" assert key_data.metrics.spend == 10.0 diff --git a/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py index 92b1ab1586d..0d8b19345f1 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py @@ -2097,6 +2097,22 @@ def test_update_internal_user_params_reset_max_budget_with_none(): assert non_default_values["user_id"] == "test_user" +def test_update_internal_user_params_explicit_duration_clear_overrides_role_default(monkeypatch): + import litellm + + monkeypatch.setattr(litellm, "internal_user_budget_duration", "30d") + data = UpdateUserRequest( + user_id="duration-clear-test", + user_role=LitellmUserRoles.INTERNAL_USER, + budget_duration=None, + ) + + updated = _update_internal_user_params(data_json=data.model_dump(exclude_unset=True), data=data) + + assert updated["budget_duration"] is None + assert updated["budget_reset_at"] is None + + def test_update_internal_user_params_ignores_other_nones(): """ Test that other fields are still filtered out if None diff --git a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py index d721be62efe..2d7397594aa 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py @@ -7,6 +7,7 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest +from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.proxy.pass_through_endpoints.llm_provider_handlers.anthropic_passthrough_logging_handler import ( AnthropicPassthroughLoggingHandler, @@ -1551,6 +1552,7 @@ class TestInterruptedStreamOutputTokenRecovery: return f"event: {event}\ndata: {json.dumps(data)}\n\n".encode() _MODEL = "claude-3-5-haiku-20241022" + _PRICED_MODEL = "claude-sonnet-5" _OUTPUT_TEXT = ( "The history of computing spans centuries, beginning with mechanical " "calculators and the abacus, advancing through Charles Babbage's " @@ -1559,7 +1561,7 @@ class TestInterruptedStreamOutputTokenRecovery: "century that gave rise to the modern information age." ) - def _interrupted_chunks(self, *, placeholder_output_tokens: int = 2): + def _interrupted_chunks(self, *, placeholder_output_tokens: int = 2, model: str | None = None): from litellm.proxy.pass_through_endpoints.streaming_handler import ( PassThroughStreamingHandler, ) @@ -1574,7 +1576,7 @@ class TestInterruptedStreamOutputTokenRecovery: "id": "msg_interrupted", "type": "message", "role": "assistant", - "model": self._MODEL, + "model": model or self._MODEL, "content": [], "stop_reason": None, "stop_sequence": None, @@ -1676,6 +1678,81 @@ class TestInterruptedStreamOutputTokenRecovery: # provider count is preserved verbatim. assert usage.completion_tokens == final + @pytest.mark.asyncio + async def test_interrupted_stream_logs_cost_of_recovered_tokens(self): + """ + Regression (LIT-6872): stream_chunk_builder stamps usage.cost and + _hidden_params["response_cost"] from the message_start placeholder before + the interrupted stream is re-tokenized, and the success handler prefers + that hidden cost over the recomputed one. The logged cost must price the + recovered completion tokens, not the placeholder. + """ + import litellm + + class _SuccessRecorder(CustomLogger): + def __init__(self): + super().__init__() + self.success_kwargs: list = [] + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + self.success_kwargs.append(kwargs) + + recorder = _SuccessRecorder() + logging_obj = LiteLLMLoggingObj( + model=self._PRICED_MODEL, + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type="pass_through_endpoint", + start_time=datetime.now(), + litellm_call_id="lit-6872", + function_id="lit-6872", + dynamic_async_success_callbacks=[recorder], + ) + logging_obj.update_environment_variables( + model=self._PRICED_MODEL, + user="", + optional_params={}, + litellm_params={"custom_llm_provider": "anthropic"}, + custom_llm_provider="anthropic", + ) + placeholder = 1 + handled = AnthropicPassthroughLoggingHandler._handle_logging_anthropic_collected_chunks( + litellm_logging_obj=logging_obj, + passthrough_success_handler_obj=MagicMock(), + url_route="/anthropic/v1/messages", + request_body={"model": self._PRICED_MODEL, "stream": True}, + endpoint_type="messages", + start_time=datetime.now(), + all_chunks=self._interrupted_chunks(placeholder_output_tokens=placeholder, model=self._PRICED_MODEL), + end_time=datetime.now(), + ) + await logging_obj.dispatch_success_handlers( + result=handled["result"], + start_time=datetime.now(), + end_time=datetime.now(), + cache_hit=False, + prefer_async_handlers=True, + **handled["kwargs"], + ) + for _ in range(300): + if recorder.success_kwargs: + break + await asyncio.sleep(0.01) + + assert len(recorder.success_kwargs) == 1 + logged = recorder.success_kwargs[0]["standard_logging_object"] + recovered_tokens = handled["result"].usage.completion_tokens + assert recovered_tokens > placeholder + assert logged["completion_tokens"] == recovered_tokens + prompt_cost, completion_cost = litellm.cost_per_token( + model=self._PRICED_MODEL, prompt_tokens=29, completion_tokens=recovered_tokens + ) + _, placeholder_completion_cost = litellm.cost_per_token( + model=self._PRICED_MODEL, prompt_tokens=29, completion_tokens=placeholder + ) + assert logged["response_cost"] == pytest.approx(prompt_cost + completion_cost) + assert logged["response_cost"] > prompt_cost + placeholder_completion_cost + class TestStreamFalseDeduplication: """ diff --git a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_passthrough_logging_handler.py b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_passthrough_logging_handler.py index a3d3ae32169..3d42301ed11 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_passthrough_logging_handler.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_passthrough_logging_handler.py @@ -1832,8 +1832,10 @@ class TestOpenAIPassthroughResponsesStreamingSpendLog: def setup_method(self): self.start_time = datetime.now() self.end_time = datetime.now() + + def _expected_spend(self) -> float: rates = litellm.model_cost[self.MODEL_MAP_KEY] - self.expected_spend = ( + return ( self.INPUT_TOKENS * rates["input_cost_per_token"] + self.OUTPUT_TOKENS * rates["output_cost_per_token"] ) @@ -1914,7 +1916,7 @@ class TestOpenAIPassthroughResponsesStreamingSpendLog: logging_obj.model_call_details["custom_llm_provider"] = "openai" return logging_obj - def test_streamed_responses_passthrough_spend_log_is_priced(self): + def test_streamed_responses_passthrough_spend_log_is_priced(self, local_model_cost_map): """The spend row books the same tokens, spend and `resp_` id as the buffered call.""" result = OpenAIPassthroughLoggingHandler._handle_logging_openai_collected_chunks( litellm_logging_obj=self._logging_obj(), @@ -1942,7 +1944,7 @@ class TestOpenAIPassthroughResponsesStreamingSpendLog: assert spend_log_row["prompt_tokens"] == self.INPUT_TOKENS assert spend_log_row["completion_tokens"] == self.OUTPUT_TOKENS assert spend_log_row["total_tokens"] == self.INPUT_TOKENS + self.OUTPUT_TOKENS - assert spend_log_row["spend"] == self.expected_spend + assert spend_log_row["spend"] == pytest.approx(self._expected_spend()) assert spend_log_row["request_id"] == self.RESPONSE_ID assert spend_log_row["model"] == "gpt-4o-mini" @@ -1967,7 +1969,6 @@ class TestOpenAIPassthroughEmbeddingsSpendLog: def setup_method(self): self.start_time = datetime.now() self.end_time = datetime.now() - self.expected_spend = self.PROMPT_TOKENS * litellm.model_cost[self.MODEL]["input_cost_per_token"] self.response_body = { "object": "list", "data": [{"object": "embedding", "index": 0, "embedding": [0.0, 1.0]}], @@ -1976,6 +1977,9 @@ class TestOpenAIPassthroughEmbeddingsSpendLog: } self.request_body = {"model": self.MODEL, "input": "hello"} + def _expected_spend(self) -> float: + return self.PROMPT_TOKENS * litellm.model_cost[self.MODEL]["input_cost_per_token"] + def _create_mock_httpx_response(self) -> httpx.Response: mock_response = MagicMock(spec=httpx.Response) mock_response.status_code = 200 @@ -2001,7 +2005,7 @@ class TestOpenAIPassthroughEmbeddingsSpendLog: ) return logging_obj - def test_embeddings_passthrough_spend_log_is_priced(self): + def test_embeddings_passthrough_spend_log_is_priced(self, local_model_cost_map): """The dispatched call books prompt tokens and cost onto the spend row.""" dispatched = PassThroughEndpointLogging().normalize_llm_passthrough_logging_payload( httpx_response=self._create_mock_httpx_response(), @@ -2020,7 +2024,7 @@ class TestOpenAIPassthroughEmbeddingsSpendLog: ) assert dispatched["standard_logging_response_object"] is not None - assert dispatched["kwargs"]["response_cost"] == self.expected_spend + assert dispatched["kwargs"]["response_cost"] == pytest.approx(self._expected_spend()) spend_log_row = get_logging_payload( kwargs=dispatched["kwargs"], @@ -2031,7 +2035,7 @@ class TestOpenAIPassthroughEmbeddingsSpendLog: assert spend_log_row["prompt_tokens"] == self.PROMPT_TOKENS assert spend_log_row["total_tokens"] == self.PROMPT_TOKENS - assert spend_log_row["spend"] == self.expected_spend + assert spend_log_row["spend"] == pytest.approx(self._expected_spend()) assert spend_log_row["model"] == self.MODEL assert spend_log_row["custom_llm_provider"] == "openai" assert spend_log_row["request_id"] == self.CALL_ID diff --git a/tests/test_litellm/proxy/policy_engine/test_attachment_registry.py b/tests/test_litellm/proxy/policy_engine/test_attachment_registry.py index 87b1bc56659..fa37a02a37c 100644 --- a/tests/test_litellm/proxy/policy_engine/test_attachment_registry.py +++ b/tests/test_litellm/proxy/policy_engine/test_attachment_registry.py @@ -4,6 +4,7 @@ Unit tests for AttachmentRegistry - tests policy attachment matching. Tests the main entry point: get_attached_policies() """ +import time from datetime import datetime, timezone from unittest.mock import AsyncMock, MagicMock @@ -222,6 +223,21 @@ class TestGetAttachedPolicies: # Should only appear once assert attached.count("multi-policy") == 1 + def test_many_distinct_policies_resolve_in_linear_time(self): + policy_count = 20_000 + registry = AttachmentRegistry() + registry.load_attachments( + [{"policy": f"policy-{index}", "scope": "*"} for index in range(policy_count)] + ) + context = PolicyMatchContext(team_alias="team", key_alias="key", model="gpt-4") + + started = time.perf_counter() + attached = registry.get_attached_policies(context) + elapsed = time.perf_counter() - started + + assert attached == [f"policy-{index}" for index in range(policy_count)] + assert elapsed < 1.0, f"{policy_count} attachments took {elapsed:.2f}s, dedup is no longer one pass" + def test_no_attachments_returns_empty(self): """Test empty attachments returns empty list.""" registry = AttachmentRegistry() diff --git a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py index e79448d0620..e109b650da7 100644 --- a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py +++ b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py @@ -8,6 +8,7 @@ Pins covered: from __future__ import annotations +import asyncio import json import logging import os @@ -712,22 +713,124 @@ async def test_ProxyConfig__get_config_from_file_missing_path_raises(): # --------------------------------------------------------------------------- -def test_ProxyConfig__process_includes_merges_files(tmp_path): +@pytest.mark.asyncio +async def test_ProxyConfig__process_includes_merges_files(tmp_path): inc = tmp_path / "models.yaml" inc.write_text("model_list:\n - model_name: gpt-4\n") pc = ProxyConfig() cfg = {"include": ["models.yaml"], "model_list": [], "litellm_settings": {}} - result = pc._process_includes(cfg, base_dir=str(tmp_path)) + result = await pc._process_includes(cfg, config_file_path=str(tmp_path / "config.yaml")) assert result == { "model_list": [{"model_name": "gpt-4"}], "litellm_settings": {}, } -def test_ProxyConfig__process_includes_missing_file_raises(tmp_path): +@pytest.mark.asyncio +async def test_ProxyConfig__process_includes_missing_file_raises(tmp_path): pc = ProxyConfig() with pytest.raises(FileNotFoundError): - pc._process_includes({"include": ["nope.yaml"]}, base_dir=str(tmp_path)) + await pc._process_includes({"include": ["nope.yaml"]}, config_file_path=str(tmp_path / "config.yaml")) + + +@pytest.mark.asyncio +async def test_ProxyConfig__process_includes_follows_nested_includes(tmp_path): + (tmp_path / "models.yaml").write_text("include:\n - more_models.yaml\nmodel_list:\n - model_name: first\n") + (tmp_path / "more_models.yaml").write_text("model_list:\n - model_name: second\n") + result = await ProxyConfig()._process_includes( + {"include": ["models.yaml"]}, config_file_path=str(tmp_path / "config.yaml") + ) + assert result == {"model_list": [{"model_name": "first"}, {"model_name": "second"}]} + + +@pytest.mark.asyncio +async def test_ProxyConfig__process_includes_resolves_a_nested_include_next_to_its_own_file(tmp_path): + (tmp_path / "shared").mkdir() + (tmp_path / "shared" / "models.yaml").write_text( + "include:\n - more_models.yaml\nmodel_list:\n - model_name: first\n" + ) + (tmp_path / "shared" / "more_models.yaml").write_text("model_list:\n - model_name: second\n") + (tmp_path / "more_models.yaml").write_text("model_list:\n - model_name: wrong-directory\n") + + result = await ProxyConfig()._process_includes( + {"include": ["shared/models.yaml"]}, config_file_path=str(tmp_path / "config.yaml") + ) + + assert result == {"model_list": [{"model_name": "first"}, {"model_name": "second"}]} + + +@pytest.mark.asyncio +async def test_ProxyConfig__process_includes_still_reads_a_nested_include_left_beside_the_root_config(tmp_path): + (tmp_path / "shared").mkdir() + (tmp_path / "shared" / "models.yaml").write_text( + "include:\n - more_models.yaml\nmodel_list:\n - model_name: first\n" + ) + (tmp_path / "more_models.yaml").write_text("model_list:\n - model_name: second\n") + + result = await ProxyConfig()._process_includes( + {"include": ["shared/models.yaml"]}, config_file_path=str(tmp_path / "config.yaml") + ) + + assert result == {"model_list": [{"model_name": "first"}, {"model_name": "second"}]} + + +@pytest.mark.asyncio +async def test_ProxyConfig__process_includes_names_both_files_when_a_nested_include_matches_two(tmp_path, caplog): + (tmp_path / "shared").mkdir() + (tmp_path / "shared" / "models.yaml").write_text( + "include:\n - more_models.yaml\nmodel_list:\n - model_name: first\n" + ) + (tmp_path / "shared" / "more_models.yaml").write_text("model_list:\n - model_name: next-to-the-declaring-file\n") + (tmp_path / "more_models.yaml").write_text("model_list:\n - model_name: next-to-the-root-config\n") + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + result = await ProxyConfig()._process_includes( + {"include": ["shared/models.yaml"]}, config_file_path=str(tmp_path / "config.yaml") + ) + + assert result == {"model_list": [{"model_name": "first"}, {"model_name": "next-to-the-declaring-file"}]} + assert [ + record + for record in caplog.records + if str(tmp_path / "shared" / "more_models.yaml") in record.getMessage() + and str(tmp_path / "more_models.yaml") in record.getMessage() + ] + + +@pytest.mark.asyncio +async def test_ProxyConfig__process_includes_merges_a_shared_file_once(tmp_path): + (tmp_path / "shared.yaml").write_text("model_list:\n - model_name: shared\n") + (tmp_path / "a.yaml").write_text("include:\n - shared.yaml\n") + (tmp_path / "b.yaml").write_text("include:\n - ./shared.yaml\n") + + result = await ProxyConfig()._process_includes( + {"include": ["a.yaml", "b.yaml"]}, config_file_path=str(tmp_path / "config.yaml") + ) + + assert result == {"model_list": [{"model_name": "shared"}]} + + +@pytest.mark.asyncio +async def test_ProxyConfig__process_includes_names_the_file_when_it_is_not_a_mapping(tmp_path): + (tmp_path / "models.yaml").write_text("- model_name: gpt-4\n") + + with pytest.raises(ValueError, match=re.escape(str(tmp_path / "models.yaml"))): + await ProxyConfig()._process_includes( + {"include": ["models.yaml"]}, config_file_path=str(tmp_path / "config.yaml") + ) + + +@pytest.mark.asyncio +async def test_ProxyConfig__process_includes_terminates_on_a_cycle(tmp_path): + (tmp_path / "a.yaml").write_text("include:\n - b.yaml\nmodel_list:\n - model_name: from-a\n") + (tmp_path / "b.yaml").write_text("include:\n - a.yaml\nmodel_list:\n - model_name: from-b\n") + + result = await asyncio.wait_for( + ProxyConfig()._process_includes({"include": ["a.yaml"]}, config_file_path=str(tmp_path / "config.yaml")), + timeout=10, + ) + + assert result == {"model_list": [{"model_name": "from-a"}, {"model_name": "from-b"}]} # --------------------------------------------------------------------------- @@ -1044,6 +1147,31 @@ async def test_ProxyConfig_get_config_loads_from_file(tmp_path, monkeypatch): } +@pytest.mark.asyncio +async def test_ProxyConfig_get_config_from_a_bucket_merges_includes(monkeypatch): + objects = { + "lit6982/config.yaml": { + "include": ["model_config.yaml"], + "general_settings": {"master_key": "sk-1234"}, + }, + "lit6982/model_config.yaml": {"model_list": [{"model_name": "included-model"}]}, + } + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) + monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", False) + monkeypatch.setattr( + "litellm.proxy.common_utils.load_config_utils.s3_object_reader", + lambda bucket_name: objects.get, + ) + monkeypatch.setenv("LITELLM_CONFIG_BUCKET_NAME", "litellm-configs") + monkeypatch.setenv("LITELLM_CONFIG_BUCKET_OBJECT_KEY", "lit6982/config.yaml") + monkeypatch.setenv("LITELLM_CONFIG_BUCKET_TYPE", "s3") + + cfg = await ProxyConfig().get_config() + + assert cfg["model_list"] == [{"model_name": "included-model"}] + assert "include" not in cfg + + @pytest.mark.asyncio async def test_ProxyConfig_get_config_missing_file_raises(monkeypatch): monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) diff --git a/tests/test_litellm/proxy/test_prisma_migration.py b/tests/test_litellm/proxy/test_prisma_migration.py index 729adcfb9e0..3fc69b34213 100644 --- a/tests/test_litellm/proxy/test_prisma_migration.py +++ b/tests/test_litellm/proxy/test_prisma_migration.py @@ -1,4 +1,6 @@ import os +import sys +from pathlib import Path from unittest.mock import MagicMock, patch import pytest @@ -64,3 +66,34 @@ class TestPrismaMigration: prisma_migration.main() mock_subprocess_run.assert_not_called() + + @patch("litellm.proxy.prisma_migration.subprocess.run") # test-quality-ok: the spawned argv is the behavior under test + @patch("litellm.proxy.prisma_migration.run_server") # test-quality-ok: run_server boots the whole proxy + def test_prisma_generate_runs_through_the_module_when_the_cli_is_not_on_path( + self, mock_run_server: MagicMock, mock_subprocess_run: MagicMock, tmp_path: Path + ) -> None: + mock_subprocess_run.return_value = MagicMock(returncode=0, stdout="", stderr="") + empty_bin: Path = tmp_path / "emptybin" + empty_bin.mkdir() + + with patch.dict(os.environ, {"PATH": str(empty_bin)}, clear=True): + assert prisma_migration.main() == 0 + + assert mock_subprocess_run.call_args.args[0] == (sys.executable, "-m", "prisma", "generate") + + @patch("litellm.proxy.prisma_migration.subprocess.run") # test-quality-ok: the spawned argv is the behavior under test + @patch("litellm.proxy.prisma_migration.run_server") # test-quality-ok: run_server boots the whole proxy + def test_prisma_generate_runs_the_console_script_when_it_is_on_path( + self, mock_run_server: MagicMock, mock_subprocess_run: MagicMock, tmp_path: Path + ) -> None: + mock_subprocess_run.return_value = MagicMock(returncode=0, stdout="", stderr="") + bin_dir: Path = tmp_path / "bin" + bin_dir.mkdir() + script: Path = bin_dir / "prisma" + script.write_text("#!/bin/sh\nexit 0\n") + script.chmod(0o755) + + with patch.dict(os.environ, {"PATH": str(bin_dir)}, clear=True): + assert prisma_migration.main() == 0 + + assert mock_subprocess_run.call_args.args[0] == ("prisma", "generate") diff --git a/tests/test_litellm/proxy/test_proxy_cli.py b/tests/test_litellm/proxy/test_proxy_cli.py index c76ff189a8a..e25e6a59884 100644 --- a/tests/test_litellm/proxy/test_proxy_cli.py +++ b/tests/test_litellm/proxy/test_proxy_cli.py @@ -1940,6 +1940,66 @@ class TestRunServerDbSetup: use_migrate=False, use_v2_resolver=False ) + @patch("atexit.register") + @patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database") # test-quality-ok: run_server always wires the DB; same isolation as the sibling CLI tests above + @patch("litellm.proxy.db.check_migration.check_prisma_schema_diff") # test-quality-ok: run_server always wires the DB; same isolation as the sibling CLI tests above + @patch("litellm.proxy.db.prisma_client.should_update_prisma_schema") # test-quality-ok: run_server always wires the DB; same isolation as the sibling CLI tests above + def test_migrations_run_when_the_prisma_cli_is_not_on_path( + self, + mock_should_update_schema, + mock_check_schema_diff, + mock_setup_database, + mock_atexit_register, + tmp_path, + capsys, + ): + from litellm.proxy.proxy_cli import run_server + + mock_should_update_schema.return_value = True + empty_bin = tmp_path / "emptybin" + empty_bin.mkdir() + + mock_proxy_module = MagicMock( + app=MagicMock(), + ProxyConfig=MagicMock(), + KeyManagementSettings=MagicMock(), + save_worker_config=MagicMock(), + ) + + clean_env = { + k: v + for k, v in os.environ.items() + if k not in ("DATABASE_URL", "DIRECT_URL") + } + clean_env["DATABASE_URL"] = "postgresql://test:test@localhost:5432/test" + clean_env["PATH"] = str(empty_bin) + + with ( + patch.dict(os.environ, clean_env, clear=True), + patch.dict( + "sys.modules", + { + "proxy_server": mock_proxy_module, + "litellm.proxy.proxy_server": mock_proxy_module, + }, + ), + patch( # test-quality-ok: same isolation as the sibling CLI tests above + "litellm.proxy.proxy_cli.ProxyInitializationHelpers._get_default_unvicorn_init_args" + ) as mock_get_args, + ): + mock_get_args.return_value = { + "app": "litellm.proxy.proxy_server:app", + "host": "localhost", + "port": 8000, + } + + run_server.main(["--local", "--skip_server_startup"], standalone_mode=False) + + assert "prisma CLI is neither on PATH" not in capsys.readouterr().out + mock_setup_database.assert_called_once_with( + use_migrate=True, use_v2_resolver=False + ) + @patch("subprocess.run") @patch("atexit.register") @patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database") diff --git a/tests/test_litellm/realtime_api/test_main.py b/tests/test_litellm/realtime_api/test_main.py index 0827bbcdc38..d3d41c5b54b 100644 --- a/tests/test_litellm/realtime_api/test_main.py +++ b/tests/test_litellm/realtime_api/test_main.py @@ -4,7 +4,6 @@ from types import TracebackType from typing import Final from unittest.mock import MagicMock, patch - import pytest import litellm @@ -152,6 +151,33 @@ async def test_vertex_credential_resolution_bounds_a_thread_offloaded_refresh(): assert time.monotonic() - start < 5 +@pytest.mark.asyncio +async def test_meta_realtime_dispatches_to_base_handler_with_meta_config(monkeypatch: pytest.MonkeyPatch): + from litellm.llms.meta.realtime.transformation import MetaRealtimeConfig + + captured: dict[str, object] = {} + + def mock_get_llm_provider(model, api_base, api_key): + return model.removeprefix("meta/"), "meta", None, api_base + + async def mock_async_realtime(**kwargs): + captured.update(kwargs) + + monkeypatch.setattr(realtime_main, "get_llm_provider", mock_get_llm_provider) + monkeypatch.setattr(realtime_main.base_llm_http_handler, "async_realtime", mock_async_realtime) + + await realtime_main._arealtime.__wrapped__( + model="meta/muse-voice-transcribe-1.0", + websocket=MagicMock(), + litellm_logging_obj=FakeLogging(), + query_params={"model": "meta/muse-voice-transcribe-1.0", "intent": "transcription"}, + ) + + assert isinstance(captured["provider_config"], MetaRealtimeConfig) + assert captured["model"] == "muse-voice-transcribe-1.0" + assert captured["query_params"] == {"model": "muse-voice-transcribe-1.0", "intent": "transcription"} + + @pytest.mark.asyncio async def test_arealtime_vertex_branch_resolves_credentials_under_a_bound(monkeypatch): """The wiring half of the regression: the vertex branch of _arealtime must diff --git a/tests/test_litellm/test_cost_calculator.py b/tests/test_litellm/test_cost_calculator.py index 0b339e3d525..f2ee2cdbd9a 100644 --- a/tests/test_litellm/test_cost_calculator.py +++ b/tests/test_litellm/test_cost_calculator.py @@ -1,4 +1,3 @@ - import json from pathlib import Path from typing import Final @@ -17,6 +16,8 @@ from litellm.cost_calculator import ( handle_realtime_stream_cost_calculation, response_cost_calculator, ) +from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.llms.base_llm.ocr.transformation import OCRPage, OCRResponse, OCRUsageInfo from litellm.types.llms.openai import OpenAIRealtimeStreamList from litellm.types.rerank import RerankResponse from litellm.types.utils import ( @@ -147,9 +148,7 @@ def test_jina_rerank_bills_total_tokens_at_input_rate_only(_local_model_cost_map def test_cost_calculator_with_response_cost_in_additional_headers(): class MockResponse(BaseModel): - _hidden_params = { - "additional_headers": {"llm_provider-x-litellm-response-cost": 1000} - } + _hidden_params = {"additional_headers": {"llm_provider-x-litellm-response-cost": 1000}} result = response_cost_calculator( response_object=MockResponse(), @@ -205,7 +204,9 @@ def test_vertex_lyria_speech_cost( call_type=call_type, ) - expected: Final = 0 if runtime_state == "custom_zero" else expected_cost * (2 if runtime_state == "custom_price" else 1) + expected: Final = ( + 0 if runtime_state == "custom_zero" else expected_cost * (2 if runtime_state == "custom_price" else 1) + ) assert cost == pytest.approx(expected) @@ -332,13 +333,12 @@ def test_cost_calculator_with_usage(_local_model_cost_map, monkeypatch): # Step 1: Test a model where input_cost_per_image_token is not set. # In this case the calculation should use input_cost_per_token as fallback. - assert ( - model_info.get("input_cost_per_image_token") is None - ), "Test case expects that input_cost_per_image_token is not set" + assert model_info.get("input_cost_per_image_token") is None, ( + "Test case expects that input_cost_per_image_token is not set" + ) expected_cost = ( - usage.prompt_tokens_details.audio_tokens - * model_info["input_cost_per_audio_token"] + usage.prompt_tokens_details.audio_tokens * model_info["input_cost_per_audio_token"] + usage.prompt_tokens_details.text_tokens * model_info["input_cost_per_token"] + usage.prompt_tokens_details.image_tokens * model_info["input_cost_per_token"] + usage.completion_tokens * model_info["output_cost_per_token"] @@ -373,12 +373,9 @@ def test_cost_calculator_with_usage(_local_model_cost_map, monkeypatch): ) expected_cost = ( - usage.prompt_tokens_details.audio_tokens - * temp_model_info_object["input_cost_per_audio_token"] - + usage.prompt_tokens_details.text_tokens - * temp_model_info_object["input_cost_per_token"] - + usage.prompt_tokens_details.image_tokens - * temp_model_info_object["input_cost_per_image_token"] + usage.prompt_tokens_details.audio_tokens * temp_model_info_object["input_cost_per_audio_token"] + + usage.prompt_tokens_details.text_tokens * temp_model_info_object["input_cost_per_token"] + + usage.prompt_tokens_details.image_tokens * temp_model_info_object["input_cost_per_image_token"] + usage.completion_tokens * temp_model_info_object["output_cost_per_token"] ) @@ -388,14 +385,11 @@ def test_cost_calculator_with_usage(_local_model_cost_map, monkeypatch): def test_transcription_cost_uses_token_pricing(_local_model_cost_map): from litellm import completion_cost - usage = Usage( prompt_tokens=14, completion_tokens=45, total_tokens=59, - prompt_tokens_details=PromptTokensDetailsWrapper( - text_tokens=0, audio_tokens=14 - ), + prompt_tokens_details=PromptTokensDetailsWrapper(text_tokens=0, audio_tokens=14), ) response = TranscriptionResponse(text="demo text") response.usage = usage @@ -439,7 +433,6 @@ def test_transcription_token_pricing_is_provider_aware(_local_model_cost_map): def test_transcription_cost_falls_back_to_duration(_local_model_cost_map): from litellm import completion_cost - response = TranscriptionResponse(text="demo text") response.duration = 10.0 @@ -460,7 +453,6 @@ def test_vertex_chirp_3_transcription_cost_from_duration(_local_model_cost_map): every transcription priced to $0.00 instead of using input_cost_per_second.""" from litellm import completion_cost - response = TranscriptionResponse(text="demo text") response.duration = 18.0 @@ -484,9 +476,7 @@ def test_handle_realtime_stream_cost_calculation(): {"type": "session.created", "session": {"model": "gpt-3.5-turbo"}}, { "type": "response.done", - "response": { - "usage": {"input_tokens": 100, "output_tokens": 50, "total_tokens": 150} - }, + "response": {"usage": {"input_tokens": 100, "output_tokens": 50, "total_tokens": 150}}, }, { "type": "response.done", @@ -517,9 +507,7 @@ def test_handle_realtime_stream_cost_calculation(): expected_cost = (300 * 0.0015 / 1000) + ( # input tokens (100 + 200) 150 * 0.002 / 1000 ) # output tokens (50 + 100) - assert ( - abs(cost - expected_cost) <= 0.00075 - ) # Allow small floating point differences + assert abs(cost - expected_cost) <= 0.00075 # Allow small floating point differences # Test with different model name in session results[0]["session"]["model"] = "gpt-4" @@ -599,14 +587,7 @@ def test_handle_realtime_stream_cost_calculation_stores_cost_breakdown(): assert logging_obj.cost_breakdown is not None assert logging_obj.cost_breakdown["input_cost"] > 0 assert logging_obj.cost_breakdown["output_cost"] > 0 - assert ( - abs( - logging_obj.cost_breakdown["input_cost"] - + logging_obj.cost_breakdown["output_cost"] - - total_cost - ) - < 1e-9 - ) + assert abs(logging_obj.cost_breakdown["input_cost"] + logging_obj.cost_breakdown["output_cost"] - total_cost) < 1e-9 assert abs(logging_obj.cost_breakdown["total_cost"] - total_cost) < 1e-9 @@ -680,9 +661,7 @@ def test_realtime_logging_object_allows_null_transcript_in_conversation_item_add }, ] - usage = RealtimeAPITokenUsageProcessor.collect_and_combine_usage_from_realtime_stream_results( - results=results - ) + usage = RealtimeAPITokenUsageProcessor.collect_and_combine_usage_from_realtime_stream_results(results=results) logging_result = RealtimeAPITokenUsageProcessor.create_logging_realtime_object( usage=usage, results=results, @@ -732,9 +711,7 @@ def test_realtime_logging_object_does_not_validate_unknown_event_types(): }, ] - usage = RealtimeAPITokenUsageProcessor.collect_and_combine_usage_from_realtime_stream_results( - results=results - ) + usage = RealtimeAPITokenUsageProcessor.collect_and_combine_usage_from_realtime_stream_results(results=results) # On unfixed code this raises pydantic ValidationError instead of returning. logging_result = RealtimeAPITokenUsageProcessor.create_logging_realtime_object( usage=usage, @@ -746,8 +723,7 @@ def test_realtime_logging_object_does_not_validate_unknown_event_types(): unknown_types = { r["type"] for r in logging_result.results - if r["type"] - in ("rate_limits.updated", "response.function_call_arguments.delta") + if r["type"] in ("rate_limits.updated", "response.function_call_arguments.delta") } assert unknown_types == { "rate_limits.updated", @@ -780,9 +756,7 @@ def test_realtime_transcription_duration_cost(monkeypatch): "type": "session.created", "session": { "type": "transcription", - "audio": { - "input": {"transcription": {"model": "gpt-realtime-whisper"}} - }, + "audio": {"input": {"transcription": {"model": "gpt-realtime-whisper"}}}, }, }, { @@ -797,9 +771,7 @@ def test_realtime_transcription_duration_cost(monkeypatch): }, ] - combined = RealtimeAPITokenUsageProcessor.collect_and_combine_usage_from_realtime_stream_results( - results=results - ) + combined = RealtimeAPITokenUsageProcessor.collect_and_combine_usage_from_realtime_stream_results(results=results) logging_obj = Logging( model="gpt-realtime-whisper", messages=[], @@ -892,9 +864,7 @@ def test_realtime_transcription_token_billed_fallback(monkeypatch): # gpt-4o-transcribe: input_cost_per_audio_token = 2.5e-06, input_cost_per_token = 2.5e-06, # output_cost_per_token = 1e-05 - model_info = litellm.get_model_info( - model="gpt-4o-transcribe", custom_llm_provider="openai" - ) + model_info = litellm.get_model_info(model="gpt-4o-transcribe", custom_llm_provider="openai") usage = { "type": "tokens", "input_tokens": 40, @@ -975,10 +945,7 @@ def test_get_transcription_model_falls_back_to_session_model(monkeypatch): mock_response=True, ) - assert ( - result._hidden_params["response_cost"] - > result_2._hidden_params["response_cost"] - ) + assert result._hidden_params["response_cost"] > result_2._hidden_params["response_cost"] model_info = router.get_deployment_model_info( model_id="my-unique-model-id", model_name="anthropic/claude-sonnet-4-5-20250929" @@ -1141,9 +1108,7 @@ def test_tiered_pricing_only_deployment_selects_router_model_id(): assert entry.get("input_cost_per_token") is None assert entry.get("tiered_pricing") is not None # The stripped shared alias must not carry tiered pricing. - assert ( - litellm.model_cost["dashscope/qwen-tier-only-test"].get("tiered_pricing") is None - ) + assert litellm.model_cost["dashscope/qwen-tier-only-test"].get("tiered_pricing") is None selected = _select_model_name_for_cost_calc( model="dashscope/qwen-tier-only-test", @@ -1223,9 +1188,7 @@ def test_azure_realtime_cost_calculator(_local_model_cost_map): combined_usage_object=Usage( prompt_tokens=100, completion_tokens=100, - prompt_tokens_details=PromptTokensDetailsWrapper( - text_tokens=10, audio_tokens=90 - ), + prompt_tokens_details=PromptTokensDetailsWrapper(text_tokens=10, audio_tokens=90), ), custom_llm_provider="azure", litellm_model_name="my-custom-azure-deployment", @@ -1244,7 +1207,6 @@ def test_azure_audio_output_cost_calculation(_local_model_cost_map): """ from litellm.types.utils import Choices, CompletionTokensDetailsWrapper, Message - # Scenario from issue #19764: # Input: 17 text tokens, 0 audio tokens # Output: 110 text tokens, 482 audio tokens @@ -1300,14 +1262,10 @@ def test_azure_audio_output_cost_calculation(_local_model_cost_map): wrong_total_cost = expected_input_cost + wrong_output_cost # Verify audio tokens are NOT charged at text rate (the bug) - assert ( - abs(cost - wrong_total_cost) > 0.001 - ), "Bug: Audio tokens are being charged at text token rate" + assert abs(cost - wrong_total_cost) > 0.001, "Bug: Audio tokens are being charged at text token rate" # Verify cost matches - assert ( - abs(cost - expected_total_cost) < 0.0000001 - ), f"Expected cost {expected_total_cost}, got {cost}" + assert abs(cost - expected_total_cost) < 0.0000001, f"Expected cost {expected_total_cost}, got {cost}" def test_default_image_cost_calculator(monkeypatch): @@ -1321,9 +1279,7 @@ def test_default_image_cost_calculator(monkeypatch): monkeypatch.setattr( litellm, "model_cost", - { - "azure/bf9001cd7209f5734ecb4ab937a5a0e2ba5f119708bd68f184db362930f9dc7b": temp_object - }, + {"azure/bf9001cd7209f5734ecb4ab937a5a0e2ba5f119708bd68f184db362930f9dc7b": temp_object}, ) args = { @@ -1539,9 +1495,7 @@ def test_gemini_25_implicit_caching_cost(): expected_cost = 0.00068708 # Allow for small floating point differences - assert ( - abs(result - expected_cost) < 1e-8 - ), f"Expected cost {expected_cost}, but got {result}" + assert abs(result - expected_cost) < 1e-8, f"Expected cost {expected_cost}, but got {result}" print(f"✓ Gemini 2.5 implicit caching cost calculation is correct: ${result:.8f}") @@ -1612,9 +1566,7 @@ def test_log_context_cost_calculation(): # Get model info to understand the pricing from litellm import get_model_info - model_info = get_model_info( - model="claude-4-sonnet-20250514", custom_llm_provider="anthropic" - ) + model_info = get_model_info(model="claude-4-sonnet-20250514", custom_llm_provider="anthropic") # Calculate expected cost based on actual model pricing input_cost_per_token = model_info.get("input_cost_per_token", 0) @@ -1622,12 +1574,8 @@ def test_log_context_cost_calculation(): cache_creation_cost_per_token = model_info.get("cache_creation_input_token_cost", 0) # Check if tiered pricing is applied - input_cost_above_200k = model_info.get( - "input_cost_per_token_above_200k_tokens", input_cost_per_token - ) - output_cost_above_200k = model_info.get( - "output_cost_per_token_above_200k_tokens", output_cost_per_token - ) + input_cost_above_200k = model_info.get("input_cost_per_token_above_200k_tokens", input_cost_per_token) + output_cost_above_200k = model_info.get("output_cost_per_token_above_200k_tokens", output_cost_per_token) cache_creation_above_200k = model_info.get( "cache_creation_input_token_cost_above_200k_tokens", cache_creation_cost_per_token, @@ -1635,31 +1583,23 @@ def test_log_context_cost_calculation(): print(f"DEBUG: Base input cost per token: ${input_cost_per_token:.2e}") print(f"DEBUG: Base output cost per token: ${output_cost_per_token:.2e}") - print( - f"DEBUG: Base cache creation cost per token: ${cache_creation_cost_per_token:.2e}" - ) + print(f"DEBUG: Base cache creation cost per token: ${cache_creation_cost_per_token:.2e}") # Handle tiered pricing - if not available, use base pricing if input_cost_above_200k is not None: - print( - f"DEBUG: Tiered input cost per token (>200k): ${input_cost_above_200k:.2e}" - ) + print(f"DEBUG: Tiered input cost per token (>200k): ${input_cost_above_200k:.2e}") else: print("DEBUG: No tiered input pricing available, using base pricing") input_cost_above_200k = input_cost_per_token if output_cost_above_200k is not None: - print( - f"DEBUG: Tiered output cost per token (>200k): ${output_cost_above_200k:.2e}" - ) + print(f"DEBUG: Tiered output cost per token (>200k): ${output_cost_above_200k:.2e}") else: print("DEBUG: No tiered output pricing available, using base pricing") output_cost_above_200k = output_cost_per_token if cache_creation_above_200k is not None: - print( - f"DEBUG: Tiered cache creation cost per token (>200k): ${cache_creation_above_200k:.2e}" - ) + print(f"DEBUG: Tiered cache creation cost per token (>200k): ${cache_creation_above_200k:.2e}") else: print("DEBUG: No tiered cache creation pricing available, using base pricing") cache_creation_above_200k = cache_creation_cost_per_token @@ -1673,13 +1613,9 @@ def test_log_context_cost_calculation(): print(f"DEBUG: Expected total: ${expected_total:.6f}") # Allow for small floating point differences - assert ( - abs(result - expected_total) < 1e-6 - ), f"Expected cost ${expected_total:.6f}, but got ${result:.6f}" + assert abs(result - expected_total) < 1e-6, f"Expected cost ${expected_total:.6f}, but got ${result:.6f}" - print( - f"✓ Log context cost calculation with tiered pricing is correct: ${result:.6f}" - ) + print(f"✓ Log context cost calculation with tiered pricing is correct: ${result:.6f}") print(f" - Input tokens (300k): ${expected_input_cost:.6f}") print(f" - Output tokens (50k): ${expected_output_cost:.6f}") print(f" - Cache creation (1k): ${expected_cache_cost:.6f}") @@ -1738,8 +1674,7 @@ def test_gemini_25_explicit_caching_cost_direct_usage(): expected_actual_cost = ( model_info["input_cost_per_token"] * usage.prompt_tokens_details.text_tokens - + model_info["cache_read_input_token_cost"] - * usage.prompt_tokens_details.cached_tokens + + model_info["cache_read_input_token_cost"] * usage.prompt_tokens_details.cached_tokens + model_info["output_cost_per_token"] * usage.completion_tokens ) @@ -1763,7 +1698,6 @@ def test_azure_ai_cache_cost_calculation(_local_model_cost_map): from litellm.litellm_core_utils.llm_cost_calc.utils import generic_cost_per_token from litellm.types.utils import PromptTokensDetailsWrapper, Usage - # Register a custom azure_ai model with cache pricing test_model_id = "test-azure-ai-claude-model" litellm.register_model( @@ -1812,13 +1746,12 @@ def test_azure_ai_cache_cost_calculation(_local_model_cost_map): print(f"Output cost: {output_cost}, Expected: {expected_output_cost}") print(f"Total cost: {total_cost}") - assert ( - abs(input_cost - expected_input_cost) < 1e-10 - ), f"Input cost mismatch: got {input_cost}, expected {expected_input_cost}" - assert ( - abs(output_cost - expected_output_cost) < 1e-10 - ), f"Output cost mismatch: got {output_cost}, expected {expected_output_cost}" - + assert abs(input_cost - expected_input_cost) < 1e-10, ( + f"Input cost mismatch: got {input_cost}, expected {expected_input_cost}" + ) + assert abs(output_cost - expected_output_cost) < 1e-10, ( + f"Output cost mismatch: got {output_cost}, expected {expected_output_cost}" + ) AZURE_GPT_5_6_MAP_KEYS = ( @@ -1887,6 +1820,7 @@ def test_azure_gpt_5_6_rates_match_azure_price_page(_local_model_cost_map, model for key in token_cost_keys: assert entry[key] == pytest.approx(global_entry[key] * 1.1), key + def test_vertex_regional_deployment_costs_uplift_over_global(monkeypatch): """ Regression for https://github.com/BerriAI/litellm/issues/34393: two Vertex @@ -1969,7 +1903,6 @@ def test_cost_discount_vertex_ai(monkeypatch): from litellm import completion_cost from litellm.types.utils import Usage - # Create mock response (use a model that exists in model_prices_and_context_window.json) response = ModelResponse( id="test-id", @@ -1998,7 +1931,6 @@ def test_cost_discount_vertex_ai(monkeypatch): custom_llm_provider="vertex_ai", ) - # Verify discount is applied (5% off means 95% of original cost) expected_cost = cost_without_discount * 0.95 assert cost_with_discount == pytest.approx(expected_cost, rel=1e-9) @@ -2016,7 +1948,6 @@ def test_cost_discount_not_applied_to_other_providers(monkeypatch): from litellm import completion_cost from litellm.types.utils import Usage - # Create mock response for OpenAI response = ModelResponse( id="test-id", @@ -2045,7 +1976,6 @@ def test_cost_discount_not_applied_to_other_providers(monkeypatch): custom_llm_provider="openai", ) - # Costs should be the same (no discount applied to OpenAI) assert cost_with_selective_discount == cost_without_discount @@ -2061,7 +1991,6 @@ def test_cost_margin_percentage(monkeypatch): from litellm import completion_cost from litellm.types.utils import Usage - # Create mock response response = ModelResponse( id="test-id", @@ -2090,7 +2019,6 @@ def test_cost_margin_percentage(monkeypatch): custom_llm_provider="openai", ) - # Verify margin is applied (10% margin means 110% of original cost) expected_cost = cost_without_margin * 1.10 assert cost_with_margin == pytest.approx(expected_cost, rel=1e-9) @@ -2108,7 +2036,6 @@ def test_cost_margin_fixed_amount(monkeypatch): from litellm import completion_cost from litellm.types.utils import Usage - # Create mock response response = ModelResponse( id="test-id", @@ -2137,7 +2064,6 @@ def test_cost_margin_fixed_amount(monkeypatch): custom_llm_provider="openai", ) - # Verify fixed margin is applied expected_cost = cost_without_margin + 0.001 assert cost_with_margin == pytest.approx(expected_cost, rel=1e-9) @@ -2155,7 +2081,6 @@ def test_cost_margin_combined(monkeypatch): from litellm import completion_cost from litellm.types.utils import Usage - # Create mock response response = ModelResponse( id="test-id", @@ -2175,9 +2100,7 @@ def test_cost_margin_combined(monkeypatch): ) # Set 8% margin + $0.0005 fixed for openai - monkeypatch.setattr(litellm, "cost_margin_config", { - "openai": {"percentage": 0.08, "fixed_amount": 0.0005} - }) + monkeypatch.setattr(litellm, "cost_margin_config", {"openai": {"percentage": 0.08, "fixed_amount": 0.0005}}) # Calculate cost with margin cost_with_margin = completion_cost( @@ -2186,7 +2109,6 @@ def test_cost_margin_combined(monkeypatch): custom_llm_provider="openai", ) - # Verify combined margin is applied expected_cost = cost_without_margin * 1.08 + 0.0005 assert cost_with_margin == pytest.approx(expected_cost, rel=1e-9) @@ -2204,7 +2126,6 @@ def test_cost_margin_global(monkeypatch): from litellm import completion_cost from litellm.types.utils import Usage - # Create mock response response = ModelResponse( id="test-id", @@ -2233,7 +2154,6 @@ def test_cost_margin_global(monkeypatch): custom_llm_provider="openai", ) - # Verify global margin is applied expected_cost = cost_without_margin * 1.05 assert cost_with_global_margin == pytest.approx(expected_cost, rel=1e-9) @@ -2251,7 +2171,6 @@ def test_cost_margin_provider_overrides_global(monkeypatch): from litellm import completion_cost from litellm.types.utils import Usage - # Create mock response response = ModelResponse( id="test-id", @@ -2280,16 +2199,13 @@ def test_cost_margin_provider_overrides_global(monkeypatch): custom_llm_provider="openai", ) - # Verify provider-specific margin is used (not global) expected_cost = cost_without_margin * 1.10 # 10% from provider, not 5% from global assert cost_with_provider_margin == pytest.approx(expected_cost, rel=1e-9) print("✓ Cost margin provider override test passed:") print(f" - Original cost: ${cost_without_margin:.6f}") - print( - f" - Cost with provider margin (10%, overrides 5% global): ${cost_with_provider_margin:.6f}" - ) + print(f" - Cost with provider margin (10%, overrides 5% global): ${cost_with_provider_margin:.6f}") print(f" - Margin added: ${cost_with_provider_margin - cost_without_margin:.6f}") @@ -2300,7 +2216,6 @@ def test_cost_margin_with_discount(monkeypatch): from litellm import completion_cost from litellm.types.utils import Usage - # Create mock response response = ModelResponse( id="test-id", @@ -2331,7 +2246,6 @@ def test_cost_margin_with_discount(monkeypatch): custom_llm_provider="openai", ) - # Verify: discount applied first, then margin # Base cost -> discount: base * 0.95 -> margin: (base * 0.95) * 1.10 expected_cost = base_cost * 0.95 * 1.10 @@ -2369,9 +2283,7 @@ def test_azure_image_generation_cost_calculator(): size=None, usage=ImageUsage( input_tokens=0, - input_tokens_details=ImageUsageInputTokensDetails( - image_tokens=0, text_tokens=0 - ), + input_tokens_details=ImageUsageInputTokensDetails(image_tokens=0, text_tokens=0), output_tokens=0, total_tokens=0, ), @@ -2401,7 +2313,6 @@ def test_completion_cost_extracts_service_tier_from_response(_local_model_cost_m """Test that completion_cost extracts service_tier from completion_response object.""" from litellm import completion_cost - # Test with gpt-5-nano which has flex pricing model = "gpt-5-nano" @@ -2442,23 +2353,18 @@ def test_completion_cost_extracts_service_tier_from_response(_local_model_cost_m assert flex_cost < standard_cost, "Flex cost should be less than standard cost" flex_ratio = flex_cost / standard_cost - assert ( - 0.45 <= flex_ratio <= 0.55 - ), f"Flex pricing should be ~50% of standard, got {flex_ratio:.2f}" + assert 0.45 <= flex_ratio <= 0.55, f"Flex pricing should be ~50% of standard, got {flex_ratio:.2f}" def test_completion_cost_extracts_service_tier_from_usage(_local_model_cost_map): """Test that completion_cost extracts service_tier from usage object.""" from litellm import completion_cost - # Test with gpt-5-nano which has flex pricing model = "gpt-5-nano" # Create usage object with service_tier - usage_with_service_tier = Usage( - prompt_tokens=1000, completion_tokens=500, total_tokens=1500 - ) + usage_with_service_tier = Usage(prompt_tokens=1000, completion_tokens=500, total_tokens=1500) # Set service_tier as an attribute on the usage object setattr(usage_with_service_tier, "service_tier", "flex") @@ -2476,9 +2382,7 @@ def test_completion_cost_extracts_service_tier_from_usage(_local_model_cost_map) ) # Create usage object without service_tier - usage_without_service_tier = Usage( - prompt_tokens=1000, completion_tokens=500, total_tokens=1500 - ) + usage_without_service_tier = Usage(prompt_tokens=1000, completion_tokens=500, total_tokens=1500) # Create ModelResponse with usage without service_tier response_standard = ModelResponse( @@ -2499,16 +2403,13 @@ def test_completion_cost_extracts_service_tier_from_usage(_local_model_cost_map) assert flex_cost < standard_cost, "Flex cost should be less than standard cost" flex_ratio = flex_cost / standard_cost - assert ( - 0.45 <= flex_ratio <= 0.55 - ), f"Flex pricing should be ~50% of standard, got {flex_ratio:.2f}" + assert 0.45 <= flex_ratio <= 0.55, f"Flex pricing should be ~50% of standard, got {flex_ratio:.2f}" def test_completion_cost_service_tier_priority(_local_model_cost_map): """Test that service_tier extraction follows priority: optional_params > completion_response > usage.""" from litellm import completion_cost - # Test with gpt-5-nano which has flex pricing model = "gpt-5-nano" @@ -2557,16 +2458,13 @@ def test_completion_cost_service_tier_priority(_local_model_cost_map): assert cost_from_usage > 0, "Cost from usage should be greater than 0" # Costs should be similar (all using flex) - assert ( - abs(cost_from_params - cost_from_usage) < 1e-6 - ), "Costs from params and usage should be similar (both flex)" + assert abs(cost_from_params - cost_from_usage) < 1e-6, "Costs from params and usage should be similar (both flex)" def test_completion_cost_service_tier_for_bedrock(_local_model_cost_map): """Test that Bedrock cost calculation applies service_tier-specific pricing.""" from litellm import completion_cost - model = "bedrock/us-east-1/test-bedrock-service-tier-cost-model" litellm.register_model( model_cost={ @@ -2622,7 +2520,6 @@ def test_completion_cost_service_tier_for_anthropic(_local_model_cost_map): from litellm import completion_cost from litellm.llms.anthropic.chat.transformation import AnthropicConfig - model = "claude-test-service-tier-cost-model" litellm.register_model( model_cost={ @@ -2675,7 +2572,6 @@ def test_completion_cost_anthropic_auto_tier_uses_served_priority_rate(_local_mo from litellm import completion_cost from litellm.llms.anthropic.chat.transformation import AnthropicConfig - model = "claude-test-auto-tier-cost-model" litellm.register_model( model_cost={ @@ -2769,7 +2665,6 @@ def test_completion_cost_non_string_service_tier_defers_to_served_tier(_local_mo from litellm import completion_cost from litellm.llms.anthropic.chat.transformation import AnthropicConfig - model = "claude-test-non-string-tier-cost-model" litellm.register_model( model_cost={ @@ -2819,7 +2714,6 @@ def test_completion_cost_non_string_response_service_tier_defers_to_served_tier( from litellm import completion_cost from litellm.llms.anthropic.chat.transformation import AnthropicConfig - model = "claude-test-response-non-string-tier-cost-model" litellm.register_model( model_cost={ @@ -2842,9 +2736,7 @@ def test_completion_cost_non_string_response_service_tier_defers_to_served_tier( }, reasoning_content=None, ) - response = ModelResponse( - usage=usage, model=model, service_tier={"name": "priority"} - ) + response = ModelResponse(usage=usage, model=model, service_tier={"name": "priority"}) cost = completion_cost( completion_response=response, @@ -2867,7 +2759,6 @@ def test_completion_cost_non_string_usage_service_tier_prices_standard(_local_mo """ from litellm import completion_cost - model = "claude-test-usage-non-string-tier-cost-model" litellm.register_model( model_cost={ @@ -2914,7 +2805,6 @@ def test_anthropic_cost_per_token_prices_cache_at_served_tier_with_multiplier(_l ) from litellm.types.utils import PromptTokensDetailsWrapper, Usage - model = "claude-test-priority-cache-fast-model" litellm.register_model( model_cost={ @@ -2940,9 +2830,7 @@ def test_anthropic_cost_per_token_prices_cache_at_served_tier_with_multiplier(_l ) usage.speed = "fast" - prompt_cost, completion_cost = anthropic_cost_per_token( - model=model, usage=usage, service_tier="priority" - ) + prompt_cost, completion_cost = anthropic_cost_per_token(model=model, usage=usage, service_tier="priority") expected_prompt = ((1000 - 200) * 6e-6 + 200 * 0.6e-6) * 2 expected_completion = 500 * 30e-6 * 2 @@ -3072,9 +2960,7 @@ def test_anthropic_fast_multiplier_only_on_models_with_fast_mode(_local_model_co "model", ["claude-sonnet-4-6", "claude-mythos-5", "claude-mythos-preview"], ) -def test_anthropic_us_data_residency_uplift_on_claude_4_6_and_later_models( - _local_model_cost_map, monkeypatch, model -): +def test_anthropic_us_data_residency_uplift_on_claude_4_6_and_later_models(_local_model_cost_map, monkeypatch, model): """ Anthropic bills every Claude 4.6+ model served with ``inference_geo="us"`` at 1.1x, and echoes that geo back in the response usage, so each of these real @@ -3139,29 +3025,27 @@ 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}" + assert usage.prompt_tokens_details.text_tokens == 9, ( + f"Expected text_tokens=9, got {usage.prompt_tokens_details.text_tokens}" + ) # 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}" + 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}" + 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 >= 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") + def test_gemini_without_cache_tokens_details(): """ @@ -3228,18 +3112,18 @@ def test_gemini_implicit_caching_cost_calculation(): usage = VertexGeminiConfig._calculate_usage(completion_response) # Verify parsing - assert ( - usage.cache_read_input_tokens == 8000 - ), f"cache_read_input_tokens should be 8000, got {usage.cache_read_input_tokens}" - assert ( - usage.prompt_tokens_details.cached_tokens == 8000 - ), f"cached_tokens should be 8000, got {usage.prompt_tokens_details.cached_tokens}" + assert usage.cache_read_input_tokens == 8000, ( + f"cache_read_input_tokens should be 8000, got {usage.cache_read_input_tokens}" + ) + assert usage.prompt_tokens_details.cached_tokens == 8000, ( + f"cached_tokens should be 8000, got {usage.prompt_tokens_details.cached_tokens}" + ) # CRITICAL: text_tokens should be (10000 - 8000) = 2000, NOT 10000 # This is the fix for issue #16341 - assert ( - usage.prompt_tokens_details.text_tokens == 2000 - ), f"text_tokens should be 2000 (10000 - 8000), got {usage.prompt_tokens_details.text_tokens}" + assert usage.prompt_tokens_details.text_tokens == 2000, ( + f"text_tokens should be 2000 (10000 - 8000), got {usage.prompt_tokens_details.text_tokens}" + ) # Verify cost calculation uses cached token pricing response = ModelResponse( @@ -3277,9 +3161,7 @@ def test_gemini_implicit_caching_cost_calculation(): f"Cached tokens may not be using reduced pricing." ) - print( - "✅ Issue #16341 fix verified: Gemini implicit caching cost calculated correctly" - ) + print("✅ Issue #16341 fix verified: Gemini implicit caching cost calculated correctly") def test_additional_costs_only_for_azure_ai(_local_model_cost_map): @@ -3293,7 +3175,6 @@ def test_additional_costs_only_for_azure_ai(_local_model_cost_map): """ from litellm.cost_calculator import _get_additional_costs - # Non-azure_ai providers should return None result = _get_additional_costs( model="gpt-4o", @@ -3436,12 +3317,7 @@ def test_custom_pricing_applies_cache_creation_input_cost_via_prompt_details(): }, ) - expected = ( - (4000 - 1000 - 500) * 0.0000025 - + 1000 * 0.00000025 - + 500 * 0.000003125 - + 100 * 0.000015 - ) + expected = (4000 - 1000 - 500) * 0.0000025 + 1000 * 0.00000025 + 500 * 0.000003125 + 100 * 0.000015 assert cost == pytest.approx(expected) @@ -3486,9 +3362,7 @@ def test_custom_pricing_applies_cache_creation_input_cost_via_cache_write_tokens }, ) - expected_prompt = ( - (4000 - 1000 - 500) * 0.0000025 + 1000 * 0.00000025 + 500 * 0.000003125 - ) + expected_prompt = (4000 - 1000 - 500) * 0.0000025 + 1000 * 0.00000025 + 500 * 0.000003125 expected_completion = 100 * 0.000015 assert prompt_cost == pytest.approx(expected_prompt) @@ -3528,10 +3402,7 @@ def test_extract_cache_read_tokens_zero_when_missing(): assert _extract_cache_read_tokens({}) == 0 assert _extract_cache_read_tokens({"cache_read_input_tokens": None}) == 0 - assert ( - _extract_cache_read_tokens({"prompt_tokens_details": {"cached_tokens": None}}) - == 0 - ) + assert _extract_cache_read_tokens({"prompt_tokens_details": {"cached_tokens": None}}) == 0 def test_extract_cache_creation_tokens_anthropic_top_level(): @@ -3573,12 +3444,7 @@ def test_extract_cache_creation_tokens_zero_when_missing(): assert _extract_cache_creation_tokens({}) == 0 assert _extract_cache_creation_tokens({"cache_creation_input_tokens": None}) == 0 - assert ( - _extract_cache_creation_tokens( - {"prompt_tokens_details": {"cache_write_tokens": None}} - ) - == 0 - ) + assert _extract_cache_creation_tokens({"prompt_tokens_details": {"cache_write_tokens": None}}) == 0 def test_custom_pricing_anthropic_style_cache_tokens_not_double_counted(): @@ -3705,7 +3571,6 @@ def test_completion_cost_logs_reasoning_and_cache_breakdown(_local_model_cost_ma from litellm.litellm_core_utils.litellm_logging import Logging from litellm.types.utils import Choices, CompletionTokensDetailsWrapper, Message - logging_obj = Logging( model="gemini-2.5-flash", messages=[{"role": "user", "content": "Hello"}], @@ -3732,12 +3597,8 @@ def test_completion_cost_logs_reasoning_and_cache_breakdown(_local_model_cost_ma prompt_tokens=209, completion_tokens=3996, total_tokens=4205, - completion_tokens_details=CompletionTokensDetailsWrapper( - reasoning_tokens=3114, text_tokens=882 - ), - prompt_tokens_details=PromptTokensDetailsWrapper( - cached_tokens=100, text_tokens=109 - ), + completion_tokens_details=CompletionTokensDetailsWrapper(reasoning_tokens=3114, text_tokens=882), + prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=100, text_tokens=109), ), ) @@ -3803,9 +3664,7 @@ def test_completion_cost_logs_the_rates_it_billed_at(monkeypatch): assert rates is not None assert rates.input_cost_per_token == pytest.approx(6e-6) assert rates.cache_read_input_token_cost == pytest.approx(6e-7) - assert logging_obj.cost_breakdown["cache_read_cost"] == pytest.approx( - 100_000 * rates.cache_read_input_token_cost - ) + assert logging_obj.cost_breakdown["cache_read_cost"] == pytest.approx(100_000 * rates.cache_read_input_token_cost) assert logging_obj.cost_breakdown["output_cost"] == pytest.approx(1_000 * rates.output_cost_per_token) @@ -3976,11 +3835,7 @@ def test_completion_cost_bills_interactions_api_response(): cost = completion_cost(completion_response=response, custom_llm_provider="gemini") reasoning_rate = model_info.get("output_cost_per_reasoning_token") or model_info["output_cost_per_token"] - expected = ( - 100 * model_info["input_cost_per_token"] - + 50 * model_info["output_cost_per_token"] - + 25 * reasoning_rate - ) + expected = 100 * model_info["input_cost_per_token"] + 50 * model_info["output_cost_per_token"] + 25 * reasoning_rate assert cost == pytest.approx(expected) assert cost > 0 @@ -4151,7 +4006,9 @@ def test_completion_cost_prices_anthropic_shaped_cache_read_tokens(_local_model_ assert cost == pytest.approx(3 * 4e-6 + 4014 * 4e-7 + 5 * 2e-5, rel=1e-9) -def _together_chat_response(model: str, prompt_tokens: int, completion_tokens: int, cached_tokens: int) -> ModelResponse: +def _together_chat_response( + model: str, prompt_tokens: int, completion_tokens: int, cached_tokens: int +) -> ModelResponse: return ModelResponse( id="chatcmpl-together-cache", choices=[{"finish_reason": "stop", "index": 0, "message": {"content": "acknowledged", "role": "assistant"}}], @@ -4219,6 +4076,8 @@ def test_completion_cost_together_metadata_only_model_still_uses_size_bucket(_lo ) assert cost == pytest.approx((23 + 15) * 8e-07, rel=1e-9) + + def test_select_model_name_strips_unregistered_alias_prefix(_local_model_cost_map): """A router-facing model_name alias containing "/" whose leading segment is NOT a registered provider must not be double-prefixed into a non-existent cost key. @@ -4459,9 +4318,7 @@ def test_every_one_hour_cache_write_rate_is_double_its_input_rate(): """Guard against pasting one model's 1h cache-write price onto another: every provider LiteLLM tracks (Anthropic, Bedrock, Vertex, Azure) publishes the 1h write at 2x input.""" - cost_map = json.loads( - (Path(__file__).parents[2] / "model_prices_and_context_window.json").read_text() - ) + cost_map = json.loads((Path(__file__).parents[2] / "model_prices_and_context_window.json").read_text()) one_hour_prefix = "cache_creation_input_token_cost_above_1hr" deviations = { (name, key): (entry["input_cost_per_token" + key[len(one_hour_prefix) :]], entry[key]) @@ -4667,9 +4524,7 @@ def test_batch_cost_calculator_gpt_6_astra_bills_half_the_standard_rate(_local_m usage = Usage(prompt_tokens=1000, completion_tokens=500, total_tokens=1500) - prompt_cost, completion_cost = batch_cost_calculator( - usage=usage, model="gpt-6-astra", custom_llm_provider="openai" - ) + prompt_cost, completion_cost = batch_cost_calculator(usage=usage, model="gpt-6-astra", custom_llm_provider="openai") assert prompt_cost == pytest.approx(1000 * 5e-6) assert completion_cost == pytest.approx(500 * 2.5e-5) @@ -4829,3 +4684,228 @@ def test_usage_without_cached_tokens_details_omits_key(): dumped = usage.prompt_tokens_details.model_dump() assert "cached_tokens_details" not in dumped assert "cached_tokens_details" not in usage.prompt_tokens_details.model_dump_json() + + +UNMAPPED_OCR_MODEL: Final = "azure_ai/some-unmapped-ocr-model-for-testing" +MAPPED_OCR_MODEL: Final = "mistral/mistral-ocr-4-0" + + +def _ocr_response(model: str, pages_processed: int, credits: float | None = None) -> OCRResponse: + return OCRResponse( + pages=[OCRPage(index=index, markdown=f"page {index}") for index in range(pages_processed)], + model=model, + usage_info=OCRUsageInfo(pages_processed=pages_processed, credits=credits), + ) + + +def _ocr_logging_obj(litellm_params: dict[str, object]) -> Logging: + logging_obj: Final = Logging( + model=UNMAPPED_OCR_MODEL, + messages=[], + stream=False, + call_type="ocr", + start_time=None, + litellm_call_id="test-ocr-custom-pricing", + function_id="1234", + ) + logging_obj.update_environment_variables(litellm_params=litellm_params, optional_params={}) + return logging_obj + + +@pytest.mark.parametrize("pages_processed", [1, 3, 10]) +def test_ocr_cost_uses_deployment_per_page_pricing_for_unmapped_model(pages_processed: int): + from litellm.cost_calculator import ocr_cost + + assert UNMAPPED_OCR_MODEL not in litellm.model_cost + cost, _ = ocr_cost( + model=UNMAPPED_OCR_MODEL, + custom_llm_provider="azure_ai", + response=_ocr_response(UNMAPPED_OCR_MODEL, pages_processed=pages_processed), + model_info={"ocr_cost_per_page": 0.004}, + ) + assert cost == pytest.approx(0.004 * pages_processed) + + +def test_ocr_cost_uses_deployment_annotation_only_pricing_for_unmapped_model(): + from litellm.cost_calculator import ocr_cost + + assert UNMAPPED_OCR_MODEL not in litellm.model_cost + response: Final = OCRResponse( + pages=[OCRPage(index=index, markdown=f"page {index}") for index in range(3)], + model=UNMAPPED_OCR_MODEL, + usage_info=OCRUsageInfo(pages_processed=3, pages_processed_annotation=2), + ) + cost, _ = ocr_cost( + model=UNMAPPED_OCR_MODEL, + custom_llm_provider="azure_ai", + response=response, + model_info={"annotation_cost_per_page": 0.01}, + ) + assert cost == pytest.approx(0.01 * 2) + + +def test_ocr_cost_annotation_only_override_keeps_mapped_per_page_rate(): + from litellm.cost_calculator import ocr_cost + + map_price: Final = litellm.model_cost[MAPPED_OCR_MODEL]["ocr_cost_per_page"] + response: Final = OCRResponse( + pages=[OCRPage(index=index, markdown=f"page {index}") for index in range(3)], + model=MAPPED_OCR_MODEL, + usage_info=OCRUsageInfo(pages_processed=3, pages_processed_annotation=2), + ) + cost, _ = ocr_cost( + model=MAPPED_OCR_MODEL, + custom_llm_provider="mistral", + response=response, + model_info={"annotation_cost_per_page": 0.01}, + ) + assert cost == pytest.approx(map_price * 3 + 0.01 * 2) + + +def test_ocr_cost_uses_deployment_per_credit_pricing_for_unmapped_model(): + from litellm.cost_calculator import ocr_cost + + cost, _ = ocr_cost( + model=UNMAPPED_OCR_MODEL, + custom_llm_provider="azure_ai", + response=_ocr_response(UNMAPPED_OCR_MODEL, pages_processed=2, credits=4), + model_info={"ocr_cost_per_credit": 0.25}, + ) + assert cost == pytest.approx(0.25 * 4) + + +def test_ocr_cost_unmapped_model_without_deployment_pricing_bills_zero(): + from litellm.cost_calculator import ocr_cost + + cost, _ = ocr_cost( + model=UNMAPPED_OCR_MODEL, + custom_llm_provider="azure_ai", + response=_ocr_response(UNMAPPED_OCR_MODEL, pages_processed=5), + model_info={"id": "some-deployment-id"}, + ) + assert cost == 0.0 + + +@pytest.mark.usefixtures("_local_model_cost_map") +def test_ocr_cost_deployment_pricing_overrides_cost_map_for_mapped_model(): + from litellm.cost_calculator import ocr_cost + + map_price: Final = litellm.get_model_info(MAPPED_OCR_MODEL)["ocr_cost_per_page"] + assert map_price is not None + override_price: Final = map_price * 10 + + cost, _ = ocr_cost( + model=MAPPED_OCR_MODEL, + custom_llm_provider="mistral", + response=_ocr_response(MAPPED_OCR_MODEL, pages_processed=2), + model_info={"ocr_cost_per_page": override_price}, + ) + assert cost == pytest.approx(override_price * 2) + + +@pytest.mark.usefixtures("_local_model_cost_map") +def test_ocr_cost_falls_through_to_cost_map_when_deployment_has_no_ocr_pricing(): + from litellm.cost_calculator import ocr_cost + + map_price: Final = litellm.get_model_info(MAPPED_OCR_MODEL)["ocr_cost_per_page"] + assert map_price is not None + + cost, _ = ocr_cost( + model=MAPPED_OCR_MODEL, + custom_llm_provider="mistral", + response=_ocr_response(MAPPED_OCR_MODEL, pages_processed=2), + model_info={"id": "some-deployment-id"}, + ) + assert cost == pytest.approx(map_price * 2) + + +@pytest.mark.usefixtures("_local_model_cost_map") +def test_ocr_cost_ignores_deployment_credit_pricing_when_response_reports_no_credits(): + from litellm.cost_calculator import ocr_cost + + map_price: Final = litellm.get_model_info(MAPPED_OCR_MODEL)["ocr_cost_per_page"] + assert map_price is not None + + cost, _ = ocr_cost( + model=MAPPED_OCR_MODEL, + custom_llm_provider="mistral", + response=_ocr_response(MAPPED_OCR_MODEL, pages_processed=2), + model_info={"ocr_cost_per_credit": 0.5}, + ) + assert cost == pytest.approx(map_price * 2) + + +@pytest.mark.parametrize("metadata_key", ["metadata", "litellm_metadata"]) +def test_completion_cost_ocr_reads_deployment_pricing_from_logging_metadata(metadata_key: str): + logging_obj = _ocr_logging_obj({metadata_key: {"model_info": {"ocr_cost_per_page": 0.004}}}) + + cost = completion_cost( + completion_response=_ocr_response(UNMAPPED_OCR_MODEL, pages_processed=3), + model=UNMAPPED_OCR_MODEL, + custom_llm_provider="azure_ai", + call_type="ocr", + custom_pricing=True, + litellm_logging_obj=logging_obj, + ) + assert cost == pytest.approx(0.004 * 3) + + +def test_completion_cost_ocr_prefers_pricing_registered_under_router_model_id(monkeypatch: pytest.MonkeyPatch): + deployment_id: Final = "ocr-deployment-priced-through-litellm-params" + monkeypatch.setitem( + litellm.model_cost, deployment_id, {"mode": "ocr", "litellm_provider": "azure_ai", "ocr_cost_per_page": 0.05} + ) + logging_obj = _ocr_logging_obj({"metadata": {"model_info": {"mode": "ocr"}}}) + + cost = completion_cost( + completion_response=_ocr_response(UNMAPPED_OCR_MODEL, pages_processed=3), + model=UNMAPPED_OCR_MODEL, + custom_llm_provider="azure_ai", + call_type="ocr", + custom_pricing=True, + router_model_id=deployment_id, + litellm_logging_obj=logging_obj, + ) + assert cost == pytest.approx(0.05 * 3) + + +def test_completion_cost_ocr_bills_request_level_pricing_for_direct_sdk_call(): + logging_obj = _ocr_logging_obj({"ocr_cost_per_page": 0.05}) + + cost = completion_cost( + completion_response=_ocr_response(UNMAPPED_OCR_MODEL, pages_processed=3), + model=UNMAPPED_OCR_MODEL, + custom_llm_provider="azure_ai", + call_type="ocr", + custom_pricing=True, + litellm_logging_obj=logging_obj, + ) + assert cost == pytest.approx(0.05 * 3) + + +def test_completion_cost_ocr_request_level_pricing_fills_in_deployment_model_info_without_ocr_pricing(): + logging_obj = _ocr_logging_obj({"ocr_cost_per_page": 0.05, "metadata": {"model_info": {"mode": "ocr"}}}) + + cost = completion_cost( + completion_response=_ocr_response(UNMAPPED_OCR_MODEL, pages_processed=3), + model=UNMAPPED_OCR_MODEL, + custom_llm_provider="azure_ai", + call_type="ocr", + custom_pricing=True, + litellm_logging_obj=logging_obj, + ) + assert cost == pytest.approx(0.05 * 3) + + +def test_completion_cost_ocr_ignores_deployment_pricing_without_custom_pricing_flag(): + logging_obj = _ocr_logging_obj({"metadata": {"model_info": {"ocr_cost_per_page": 0.004}}}) + + cost = completion_cost( + completion_response=_ocr_response(UNMAPPED_OCR_MODEL, pages_processed=3), + model=UNMAPPED_OCR_MODEL, + custom_llm_provider="azure_ai", + call_type="ocr", + custom_pricing=False, + litellm_logging_obj=logging_obj, + ) + assert cost == 0.0 diff --git a/tests/test_litellm/test_main.py b/tests/test_litellm/test_main.py index f71225c6fc5..4f7a51eb531 100644 --- a/tests/test_litellm/test_main.py +++ b/tests/test_litellm/test_main.py @@ -874,15 +874,25 @@ def test_responses_api_bridge_check_gpt_5_4_tools_with_default_reasoning_routes_ assert model_info.get("mode") == "responses" -@pytest.mark.parametrize("model_name", ["gpt-5.6-sol", "gpt-5.6-luna", "gpt-5.6-terra"]) +@pytest.mark.parametrize( + "model_name, expected_mode", + [ + pytest.param("gpt-5.6-sol", "responses", id="above-boundary-bridges"), + pytest.param("gpt-5.1", None, id="below-boundary-stays-chat"), + ], +) def test_responses_api_bridge_check_gpt_5_6_tools_with_default_reasoning_routes_to_responses( - monkeypatch, model_name + monkeypatch, model_name, expected_mode ): """ - The whole gpt-5.6 family must bridge on function tools alone. The bridge used to - require an explicit reasoning_effort, so a gpt-5.6 call carrying tools and no effort - was rejected with "Function tools with reasoning_effort are not supported for - gpt-5.6-sol in /v1/chat/completions". + gpt-5.6 must bridge on function tools alone. The bridge used to require an explicit + reasoning_effort, so a gpt-5.6 call carrying tools and no effort was rejected with + "Function tools with reasoning_effort are not supported for gpt-5.6-sol in + /v1/chat/completions". + + Paired with a model below the gpt-5.4 boundary, which must still stay on chat. The + gate parses the version and drops any suffix, so the family members bridge + identically and only the boundary distinguishes behaviour. """ import litellm from litellm.main import responses_api_bridge_check @@ -901,7 +911,7 @@ def test_responses_api_bridge_check_gpt_5_6_tools_with_default_reasoning_routes_ ) assert model == model_name - assert model_info.get("mode") == "responses" + assert model_info.get("mode") == expected_mode def test_responses_api_bridge_check_gpt_5_4_tools_with_reasoning_none_stays_chat(): @@ -3311,15 +3321,19 @@ def local_cost_map(monkeypatch): """The prices these tests assert are the checked-in ones. Setting the environment variable alone does not reload the map, so pin the map itself. - ``get_model_info`` is lru_cached, so pinning ``model_cost`` is not enough on its - own: a cached entry warmed against the network-fetched map keeps its old prices - and ``completion_cost`` bills at those while the assertions read the pinned map. - Clear on the way in and out so entries never leak across tests in either direction.""" + Prices are read through two separate lru_caches, so pinning ``model_cost`` is not + enough on its own: an entry warmed against the network-fetched map keeps its old + prices and billing reads those while the assertions read the pinned map. + ``_invalidate_model_cost_lowercase_map`` clears both caches, where + ``get_model_info.cache_clear`` reaches only one. Invalidate on the way in and out + so entries never leak across tests in either direction.""" + from litellm.utils import _invalidate_model_cost_lowercase_map + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) - litellm.get_model_info.cache_clear() + _invalidate_model_cost_lowercase_map() yield - litellm.get_model_info.cache_clear() + _invalidate_model_cost_lowercase_map() def test_a_streamed_response_bills_the_usage_the_provider_reported(local_cost_map): diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index d9ca7d9ddee..2cd5d528a69 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -4248,9 +4248,6 @@ def test_deepseek_flash_completion_cost(): _FIREWORKS_MODELS = [ ( "accounts/fireworks/models/glm-5p2", - 1.4e-06, - 4.4e-06, - 1.4e-07, 1048576, 131072, False, @@ -4258,9 +4255,6 @@ _FIREWORKS_MODELS = [ ), ( "accounts/fireworks/models/glm-5p1", - 1.4e-06, - 4.4e-06, - 2.6e-07, 202800, 131072, False, @@ -4268,9 +4262,6 @@ _FIREWORKS_MODELS = [ ), ( "accounts/fireworks/routers/glm-5p1-fast", - 2.8e-06, - 8.8e-06, - 5.2e-07, 202800, 131072, False, @@ -4278,9 +4269,6 @@ _FIREWORKS_MODELS = [ ), ( "accounts/fireworks/models/qwen3p7-plus", - 4e-07, - 1.6e-06, - 8e-08, 262144, 65536, True, @@ -4288,9 +4276,6 @@ _FIREWORKS_MODELS = [ ), ( "accounts/fireworks/models/minimax-m3", - 3e-07, - 1.2e-06, - 6e-08, 512000, 512000, True, @@ -4298,9 +4283,6 @@ _FIREWORKS_MODELS = [ ), ( "accounts/fireworks/models/minimax-m2p7", - 3e-07, - 1.2e-06, - 6e-08, 196608, 196608, False, @@ -4308,9 +4290,6 @@ _FIREWORKS_MODELS = [ ), ( "accounts/fireworks/models/kimi-k2p7-code", - 9.5e-07, - 4e-06, - 1.9e-07, 262144, 32768, True, @@ -4318,9 +4297,6 @@ _FIREWORKS_MODELS = [ ), ( "accounts/fireworks/routers/kimi-k2p7-code-fast", - 1.9e-06, - 8e-06, - 3.8e-07, 262144, 32768, True, @@ -4328,9 +4304,6 @@ _FIREWORKS_MODELS = [ ), ( "accounts/fireworks/models/kimi-k2p6", - 9.5e-07, - 4e-06, - 1.6e-07, 262144, 32768, True, @@ -4338,9 +4311,6 @@ _FIREWORKS_MODELS = [ ), ( "accounts/fireworks/routers/kimi-k2p6-fast", - 2e-06, - 8e-06, - 3e-07, 262144, 32768, True, @@ -4348,9 +4318,6 @@ _FIREWORKS_MODELS = [ ), ( "accounts/fireworks/models/gpt-oss-120b", - 1.5e-07, - 6e-07, - 1.5e-08, 131072, 32768, False, @@ -4358,9 +4325,6 @@ _FIREWORKS_MODELS = [ ), ( "accounts/fireworks/models/gpt-oss-20b", - 7e-08, - 3e-07, - 3.5e-08, 131072, 32768, False, @@ -4368,9 +4332,6 @@ _FIREWORKS_MODELS = [ ), ( "accounts/fireworks/models/deepseek-v4-pro", - 1.74e-06, - 3.48e-06, - 1.45e-07, 1048576, 384000, False, @@ -4378,9 +4339,6 @@ _FIREWORKS_MODELS = [ ), ( "accounts/fireworks/models/deepseek-v4-flash", - 1.4e-07, - 2.8e-07, - 2.8e-08, 1048576, 384000, False, @@ -4412,9 +4370,6 @@ _FIREWORKS_ROUTER_SHORT_FORMS = [ def _assert_fireworks_entry( model_cost, model_path, - expected_input, - expected_output, - expected_cache, expected_max_input, expected_max_output, expected_vision, @@ -4424,9 +4379,9 @@ def _assert_fireworks_entry( assert info is not None, f"fireworks_ai/{model_path} missing from model cost map" assert info["litellm_provider"] == "fireworks_ai" assert info["mode"] == "chat" - assert info["input_cost_per_token"] == expected_input - assert info["output_cost_per_token"] == expected_output - assert info["cache_read_input_token_cost"] == expected_cache + assert info["input_cost_per_token"] > 0 + assert info["output_cost_per_token"] > 0 + assert "cache_read_input_token_cost" in info assert info["max_input_tokens"] == expected_max_input assert info["max_output_tokens"] == expected_max_output assert info["max_tokens"] == expected_max_output diff --git a/tests/test_litellm_rust/ocr/test_lifecycle.py b/tests/test_litellm_rust/ocr/test_lifecycle.py index 1acad5527d8..e7ebc5b3018 100644 --- a/tests/test_litellm_rust/ocr/test_lifecycle.py +++ b/tests/test_litellm_rust/ocr/test_lifecycle.py @@ -79,6 +79,20 @@ async def test_proxy_metadata_remains_python_owned(ocr_server: RecordingServer) assert "metadata" not in ocr_server.requests[0].body +@pytest.mark.asyncio +async def test_request_level_custom_pricing_reaches_logging_params_and_bills_the_call( + ocr_server: RecordingServer, +) -> None: + recorder: Final = RecordingLogger() + response: Final = await call_aocr(ocr_server, callbacks=[recorder], ocr_cost_per_page=0.05) + events: Final = await recorder.wait_for_async("async_log_success_event") + + assert response.usage_info is not None and response.usage_info.pages_processed == 1 + assert events[0].kwargs["litellm_params"]["ocr_cost_per_page"] == 0.05 + assert response._hidden_params["response_cost"] == pytest.approx(0.05) + assert "ocr_cost_per_page" not in ocr_server.requests[0].body + + @pytest.mark.asyncio async def test_response_replacement_finalized_before_dispatch_in_caller_task(ocr_server: RecordingServer) -> None: caller: Final = asyncio.current_task() diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/common/useResourceList.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/hooks/common/useResourceList.test.tsx index 3ca67b082ec..b0c478920f3 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/common/useResourceList.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/common/useResourceList.test.tsx @@ -152,6 +152,28 @@ describe("useResourceList", () => { await waitFor(() => expect(lastCall().page_size).toBe(25)); }); + it("reports loading while a new search request is still pending", async () => { + let resolveSecond: ((value: ResourceListPage) => void) | undefined; + const fetchPage = vi.fn((query: ResourceListQuery) => { + calls.push(query); + if (calls.length === 1) return Promise.resolve(page([{ id: "a" }], 3)); + return new Promise>((resolve) => { + resolveSecond = resolve; + }); + }); + const { result } = renderList({ fetchPage }); + await waitFor(() => expect(result.current.rows).toEqual([{ id: "a" }])); + expect(result.current.isLoading).toBe(false); + + act(() => result.current.onSearchChange("zzz")); + await waitFor(() => expect(lastCall().q).toBe("zzz")); + expect(result.current.isLoading).toBe(true); + + act(() => resolveSecond?.(page([], 0))); + await waitFor(() => expect(result.current.isLoading).toBe(false)); + expect(result.current.rows).toEqual([]); + }); + it("surfaces a failed page as an error instead of empty rows", async () => { const fetchPage = vi.fn(() => Promise.reject(new Error("boom"))); const { result } = renderList({ fetchPage }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/common/useResourceList.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/common/useResourceList.ts index 8a6376b2248..983e537bd00 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/common/useResourceList.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/common/useResourceList.ts @@ -86,7 +86,7 @@ export function useResourceList(options: UseResourceListOptions): Re enabled, placeholderData: (previous) => previous, }; - const { data, isLoading, isFetching, error, refetch: refetchQuery } = useQuery(queryOptions); + const { data, isLoading, isPlaceholderData, isFetching, error, refetch: refetchQuery } = useQuery(queryOptions); const toFirstPage = useCallback(() => setPagination((previous) => ({ ...previous, pageIndex: 0 })), []); @@ -123,7 +123,7 @@ export function useResourceList(options: UseResourceListOptions): Re return { rows, rowCount: data?.meta.total_count ?? 0, - isLoading, + isLoading: isLoading || isPlaceholderData, isFetching, error, refetch, diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/projects/useUpdateProject.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/projects/useUpdateProject.ts index 8e6bad04a28..e6dec85128c 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/projects/useUpdateProject.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/projects/useUpdateProject.ts @@ -10,7 +10,7 @@ export interface ProjectUpdateParams { description?: string; team_id?: string; models?: string[]; - max_budget?: number; + max_budget?: number | null; blocked?: boolean; guardrails?: string[]; metadata?: Record; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AccessGroupBudgetModal.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AccessGroupBudgetModal.tsx index 8df2a6620f9..cf5cf78d08a 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AccessGroupBudgetModal.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AccessGroupBudgetModal.tsx @@ -98,7 +98,11 @@ const AccessGroupBudgetModal: React.FC = ({ )} > {({ id, value, onChange }) => ( - + onChange(next ?? undefined)} + /> )} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectModals/EditProjectModal.integration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectModals/EditProjectModal.integration.test.tsx index 5b84aa15dd1..2e39d45e660 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectModals/EditProjectModal.integration.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectModals/EditProjectModal.integration.test.tsx @@ -102,6 +102,20 @@ describe("EditProjectModal submit payload", () => { }); }); + it("should send an explicit clear after blanking a saved budget", async () => { + const user = setup(); + renderModal(); + + const budgetInput = screen.getByRole("spinbutton", { name: "Max Budget (USD)" }); + await user.clear(budgetInput); + await user.tab(); + expect(budgetInput).toHaveValue(null); + await save(user); + + await waitFor(() => expect(mutate).toHaveBeenCalled()); + expect(JSON.parse(JSON.stringify(variables().params))).toMatchObject({ max_budget: null }); + }); + it("includes the advanced fields once Advanced Settings has been opened, even after collapsing it again", async () => { const user = setup(); renderModal(); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectModals/EditProjectModal.tsx b/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectModals/EditProjectModal.tsx index 77f28b05ea5..31582da5b14 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectModals/EditProjectModal.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectModals/EditProjectModal.tsx @@ -86,7 +86,7 @@ function EditProjectForm({ project, onClose, onSuccess }: Omit onChange(toOptionalNumber(event.target.value))} + value={Number.isNaN(value) ? "" : value ?? ""} + onInput={(event) => { + if (event.currentTarget.validity.badInput || Number.isNaN(value)) { + onChange( + event.currentTarget.validity.badInput + ? Number.NaN + : toOptionalNumber(event.currentTarget.value) ?? null, + ); + } + }} + onChange={(event) => + onChange(event.target.validity.badInput ? Number.NaN : toOptionalNumber(event.target.value) ?? null) + } /> )} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectModals/projectFormSchema.ts b/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectModals/projectFormSchema.ts index d4c85d89616..6e6318a9609 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectModals/projectFormSchema.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectModals/projectFormSchema.ts @@ -22,7 +22,7 @@ export const projectFormSchema = z .pipe(z.string({ error: "Please select a team" }).min(1, "Please select a team")), description: z.string().optional(), models: z.array(z.string()), - max_budget: z.number().optional(), + max_budget: z.number().nullish(), isBlocked: z.boolean(), guardrails: z.array(z.string()).optional(), modelLimits: z.array(modelLimitSchema).optional(), diff --git a/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectModals/projectFormUtils.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectModals/projectFormUtils.test.ts index 28d791ddf54..e76555b3ca3 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectModals/projectFormUtils.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectModals/projectFormUtils.test.ts @@ -27,9 +27,9 @@ describe("buildProjectCreateParams", () => { expect(result.description).toBe("A description"); }); - it("should pass through max_budget when provided", () => { - const result = buildProjectCreateParams({ ...baseValues, max_budget: 50.0 }); - expect(result.max_budget).toBe(50.0); + it.each([50.0, 1e308])("should preserve a finite max_budget of %s", (maxBudget) => { + const result = buildProjectCreateParams({ ...baseValues, max_budget: maxBudget }); + expect(JSON.parse(JSON.stringify(result)).max_budget).toBe(maxBudget); }); it("should build model_rpm_limit from modelLimits entries", () => { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectModals/projectFormUtils.ts b/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectModals/projectFormUtils.ts index 71c1f5c79af..60e97939e69 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectModals/projectFormUtils.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectModals/projectFormUtils.ts @@ -16,6 +16,11 @@ const buildModelLimitMap = ( const buildMetadata = (entries: ProjectFormValues["metadata"]): Record | undefined => entries && Object.fromEntries(entries.flatMap((entry) => (entry.key ? [[entry.key, entry.value] as const] : []))); +const roundBudget = (value: number): number => { + const rounded = Math.round(value * 100) / 100; + return Number.isFinite(rounded) ? rounded : value; +}; + const buildProjectApiParams = (values: ProjectFormValues, sendEmpty: boolean) => { const limitEntries = values.modelLimits ?? []; const modelRpmLimit = buildModelLimitMap(limitEntries, (entry) => entry.rpm); @@ -35,7 +40,7 @@ const buildProjectApiParams = (values: ProjectFormValues, sendEmpty: boolean) => project_alias: values.project_alias, description: values.description, models: values.models ?? [], - max_budget: values.max_budget === undefined ? undefined : Math.round(values.max_budget * 100) / 100, + max_budget: values.max_budget == null ? undefined : roundBudget(values.max_budget), blocked: values.isBlocked ?? false, ...guardrailsParam, ...(keep(modelRpmLimit) && { model_rpm_limit: modelRpmLimit }), @@ -53,4 +58,7 @@ export const buildProjectCreateParams = (values: ProjectFormValues) => buildProj * /project/update leaves an omitted key untouched, so a limit the operator cleared has to go out as * an explicitly empty map. Omitting it is what silently kept a removed quota enforced. */ -export const buildProjectUpdateParams = (values: ProjectFormValues) => buildProjectApiParams(values, true); +export const buildProjectUpdateParams = (values: ProjectFormValues, savedMaxBudget?: number | null) => ({ + ...buildProjectApiParams(values, true), + ...(values.max_budget == null && savedMaxBudget != null ? { max_budget: null } : {}), +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/tag-management/_components/components/CreateTagModal.tsx b/ui/litellm-dashboard/src/app/(dashboard)/tag-management/_components/components/CreateTagModal.tsx index 6f6bdcb6fe4..adb5a167c35 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/tag-management/_components/components/CreateTagModal.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/tag-management/_components/components/CreateTagModal.tsx @@ -143,7 +143,11 @@ const CreateTagModal: React.FC = ({ visible, onCancel, onSu )} > {({ id, value, onChange }) => ( - + onChange(next ?? undefined)} + /> )} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/tag-management/_components/tag_info.tsx b/ui/litellm-dashboard/src/app/(dashboard)/tag-management/_components/tag_info.tsx index e8cb358c0cf..1648a99bb0e 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/tag-management/_components/tag_info.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/tag-management/_components/tag_info.tsx @@ -28,7 +28,7 @@ const tagEditShape = { description: z.string().optional(), models: z.array(z.string()).optional(), max_budget: z.union([z.string(), z.number()]).optional(), - budget_duration: z.string().optional(), + budget_duration: z.string().nullish(), }; const tagEditSchema = z.object(tagEditShape); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.tsx b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.tsx index 273e478528e..6c15b3c418d 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.tsx @@ -42,6 +42,7 @@ import { valueFormatterSpend } from "@/components/UsagePage/utils/value_formatte import EndpointUsage from "../EndpointUsage/EndpointUsage"; import ModelViewToggle, { ModelViewType } from "../ModelViewToggle"; import TopKeyView from "@/components/UsagePage/components/EntityUsage/TopKeyView"; +import KeyActivityPanel from "@/components/UsagePage/components/KeyActivityPanel"; import TopModelView from "./TopModelView"; import TeamUserSpendCard from "./TeamUserSpendCard"; @@ -654,7 +655,7 @@ const EntityUsage: React.FC = ({ { key: "keys", label: "Key Activity", - content: , + content: , }, { key: "endpoints", label: "Endpoint Activity", content: }, ]; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.tsx b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.tsx index a92d1209567..de353948db9 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.tsx @@ -30,6 +30,7 @@ import { ActivityMetrics, processActivityData } from "@/components/activity_metr import CloudZeroExportModal from "@/components/cloudzero_export_modal"; import UserDropdown from "@/components/common_components/UserDropdown"; import EntityUsageExportModal from "@/components/EntityUsageExport"; +import KeyActivityPanel from "@/components/UsagePage/components/KeyActivityPanel"; import { Team } from "@/components/key_team_helpers/key_list"; import { gatewayDailyActivityCall, @@ -886,7 +887,7 @@ const UsagePage: React.FC = ({ teams, organizations }) => { - + diff --git a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users.test.tsx index 6b0423067fd..10a13c1dbd7 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users.test.tsx @@ -344,5 +344,16 @@ describe("ViewUserDashboard", () => { expect(latest[4]).toBeNull(); expect(latest[2]).toBe(1); }); + + it("replaces the previous rows with the loading state while the search request is pending", async () => { + renderDashboard(); + expect(await screen.findByText("test@example.com")).toBeInTheDocument(); + + userListCall.mockReturnValue(new Promise(() => undefined)); + fireEvent.change(screen.getByPlaceholderText("Search by email or ID…"), { target: { value: "zzznomatch" } }); + + expect(await screen.findByText("Loading users…")).toBeInTheDocument(); + expect(screen.queryByText("test@example.com")).not.toBeInTheDocument(); + }); }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users.tsx b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users.tsx index 1ed23e7f523..20ce22b6444 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users.tsx @@ -295,7 +295,7 @@ const ViewUserDashboard: React.FC = ({ { expect(await openEditor(user)).toHaveValue(42); }); + + it("should keep Unlimited selected after saving and reopening the user", async () => { + const user = setup(); + render(); + + await openEditor(user); + await user.click(screen.getByRole("checkbox", { name: "Unlimited Budget" })); + await user.click(screen.getByRole("button", { name: /save changes/i })); + + await waitFor(() => expect(mockUserUpdateUserCall).toHaveBeenCalled()); + expect(mockUserUpdateUserCall.mock.calls[0][1]).toMatchObject({ max_budget: null }); + await openEditor(user); + expect(screen.getByRole("checkbox", { name: "Unlimited Budget" })).toBeChecked(); + }); + + it("should keep a cleared reset period after saving and reopening the user", async () => { + const user = setup(); + render(); + + await openEditor(user); + await user.click(screen.getByRole("combobox", { name: "Reset Budget" })); + await user.click(await screen.findByRole("option", { name: "n/a" })); + await user.click(screen.getByRole("button", { name: /save changes/i })); + + await waitFor(() => expect(mockUserUpdateUserCall).toHaveBeenCalled()); + expect(mockUserUpdateUserCall.mock.calls[0][1]).toMatchObject({ budget_duration: null }); + await openEditor(user); + expect(screen.getByRole("combobox", { name: "Reset Budget" })).toHaveTextContent("n/a"); + }); }); it("offers only the teams the user is not already a member of", async () => { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/user_info_view.tsx b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/user_info_view.tsx index eed39c8e585..e083e549552 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/user_info_view.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/user_info_view.tsx @@ -332,8 +332,9 @@ export default function UserInfoView({ user_email: formValues.user_email ?? userData.user_email, user_alias: formValues.user_alias ?? userData.user_alias, models: formValues.models ?? userData.models, - max_budget: formValues.max_budget ?? userData.max_budget, - budget_duration: formValues.budget_duration ?? userData.budget_duration, + max_budget: formValues.max_budget === undefined ? userData.max_budget : formValues.max_budget, + budget_duration: + formValues.budget_duration === undefined ? userData.budget_duration : formValues.budget_duration, metadata: formValues.metadata ?? userData.metadata, model_max_budget: formValues.model_max_budget ?? userData.model_max_budget, object_permission: mcpEntitlement diff --git a/ui/litellm-dashboard/src/components/Teams.tsx b/ui/litellm-dashboard/src/components/Teams.tsx index f0d21cc5350..dc531ea5dad 100644 --- a/ui/litellm-dashboard/src/components/Teams.tsx +++ b/ui/litellm-dashboard/src/components/Teams.tsx @@ -807,7 +807,7 @@ const Teams: React.FC = ({ accessToken, userID, userRole, premiumUser showNeverResets placeholder={budgetDurationPlaceholder} value={value} - onChange={onChange} + onChange={(next) => onChange(next ?? undefined)} /> )} diff --git a/ui/litellm-dashboard/src/components/TeamsPage/TeamsTable.test.tsx b/ui/litellm-dashboard/src/components/TeamsPage/TeamsTable.test.tsx index 95958994775..9699ea2b7d1 100644 --- a/ui/litellm-dashboard/src/components/TeamsPage/TeamsTable.test.tsx +++ b/ui/litellm-dashboard/src/components/TeamsPage/TeamsTable.test.tsx @@ -138,6 +138,14 @@ it("shows a loading state on initial load and hides the data", () => { expect(screen.queryByText("Acme Team")).not.toBeInTheDocument(); }); +it("replaces the previous rows with the loading state while a new search is pending", () => { + mockUseTeamsTable.mockReturnValue(teamsResult([mockTeam], {}, { isPlaceholderData: true, isFetching: true })); + renderTable(); + + expect(screen.getByText("Loading teams...")).toBeInTheDocument(); + expect(screen.queryByText("Acme Team")).not.toBeInTheDocument(); +}); + describe("sort contract – only backend-sortable columns are sortable", () => { it("requests the default created_at descending sort on first render", () => { renderTable(); diff --git a/ui/litellm-dashboard/src/components/TeamsPage/TeamsTable.tsx b/ui/litellm-dashboard/src/components/TeamsPage/TeamsTable.tsx index bde1c7724c3..07f4a0bcbc4 100644 --- a/ui/litellm-dashboard/src/components/TeamsPage/TeamsTable.tsx +++ b/ui/litellm-dashboard/src/components/TeamsPage/TeamsTable.tsx @@ -83,7 +83,8 @@ export function TeamsTable({ userRole, userID, onSelectTeam, onEditTeam, onDelet const { data: teamsResponse, - isPending: isLoading, + isPending, + isPlaceholderData, isFetching, refetch, } = useTeamsTable(tablePagination.pageIndex + 1, tablePagination.pageSize, teamListOptions); @@ -161,7 +162,7 @@ export function TeamsTable({ userRole, userID, onSelectTeam, onEditTeam, onDelet onColumnFiltersChange={handleColumnFiltersChange} enableColumnResizing columnResizeMode="onChange" - isLoading={isLoading} + isLoading={isPending || isPlaceholderData} loadingMessage="Loading teams..." noDataMessage="No teams found" fillHeight diff --git a/ui/litellm-dashboard/src/components/UsagePage/components/KeyActivityPanel.test.tsx b/ui/litellm-dashboard/src/components/UsagePage/components/KeyActivityPanel.test.tsx new file mode 100644 index 00000000000..693ac20a360 --- /dev/null +++ b/ui/litellm-dashboard/src/components/UsagePage/components/KeyActivityPanel.test.tsx @@ -0,0 +1,71 @@ +import { fireEvent, render, screen } from "@testing-library/react"; +import { describe, expect, it, vi } from "vitest"; + +import type { ModelActivityData } from "../types"; +import KeyActivityPanel from "./KeyActivityPanel"; + +vi.mock("@/components/activity_metrics", () => ({ + ActivityMetrics: ({ modelMetrics }: { modelMetrics: Record }) => ( +
    + {Object.keys(modelMetrics).map((hash) => ( +
  • {hash}
    • + ))} +
+ ), +})); + +function activity(label: string, user_email: string | null, user_id: string | null): ModelActivityData { + return { + label, + key_metadata: { key_alias: label, team_id: "team-1", user_id, user_email }, + total_requests: 1, + total_successful_requests: 1, + total_failed_requests: 0, + total_cache_read_input_tokens: 0, + total_cache_creation_input_tokens: 0, + total_tokens: 10, + prompt_tokens: 5, + completion_tokens: 5, + total_spend: 0.01, + top_api_keys: [], + top_models: [], + daily_data: [], + }; +} + +const keyMetrics: Record = { + "hash-alice": activity("alice-key", "alice@example.com", "user-alice"), + "hash-bob": activity("bob-key", "bob@example.com", "user-bob"), +}; + +describe("KeyActivityPanel", () => { + it("renders every key and the full count before searching", () => { + render(); + expect(screen.getByTestId("rendered-keys")).toHaveTextContent("hash-alicehash-bob"); + expect(screen.getByText("Showing 2 of 2 keys")).toBeInTheDocument(); + }); + + it("narrows the rendered keys to those matching the user email", () => { + render(); + fireEvent.change(screen.getByLabelText("Search keys"), { target: { value: "bob@example.com" } }); + expect(screen.getByTestId("rendered-keys")).toHaveTextContent("hash-bob"); + expect(screen.getByTestId("rendered-keys")).not.toHaveTextContent("hash-alice"); + expect(screen.getByText("Showing 1 of 2 keys")).toBeInTheDocument(); + }); + + it("shows an empty state instead of zeroed metrics when nothing matches", () => { + render(); + fireEvent.change(screen.getByLabelText("Search keys"), { target: { value: "carol" } }); + expect(screen.queryByTestId("rendered-keys")).not.toBeInTheDocument(); + expect(screen.getByText('No keys match "carol" in this date range')).toBeInTheDocument(); + }); + + it("clears the search and restores every key", () => { + render(); + fireEvent.change(screen.getByLabelText("Search keys"), { target: { value: "user-alice" } }); + expect(screen.getByTestId("rendered-keys")).toHaveTextContent("hash-alice"); + fireEvent.click(screen.getByLabelText("Clear key search")); + expect(screen.getByLabelText("Search keys")).toHaveValue(""); + expect(screen.getByTestId("rendered-keys")).toHaveTextContent("hash-alicehash-bob"); + }); +}); diff --git a/ui/litellm-dashboard/src/components/UsagePage/components/KeyActivityPanel.tsx b/ui/litellm-dashboard/src/components/UsagePage/components/KeyActivityPanel.tsx new file mode 100644 index 00000000000..8287a04d0c7 --- /dev/null +++ b/ui/litellm-dashboard/src/components/UsagePage/components/KeyActivityPanel.tsx @@ -0,0 +1,58 @@ +import { Search, X } from "lucide-react"; +import React, { useMemo, useState } from "react"; + +import { ActivityMetrics } from "@/components/activity_metrics"; +import { InputGroup, InputGroupAddon, InputGroupButton, InputGroupInput } from "@/components/ui/input-group"; + +import { filterKeyActivity } from "../keyActivityFilter"; +import type { ModelActivityData } from "../types"; + +interface KeyActivityPanelProps { + keyMetrics: Record; + hidePromptCachingMetrics?: boolean; +} + +const KeyActivityPanel: React.FC = ({ keyMetrics, hidePromptCachingMetrics = false }) => { + const [query, setQuery] = useState(""); + const filtered = useMemo(() => filterKeyActivity(keyMetrics, query), [keyMetrics, query]); + const totalKeys = Object.keys(keyMetrics).length; + const shownKeys = Object.keys(filtered).length; + const isFiltering = query.trim() !== ""; + + return ( +
+
+ + + + + setQuery(e.target.value)} + /> + {isFiltering && ( + + setQuery("")}> + + + + )} + + + Showing {shownKeys.toLocaleString()} of {totalKeys.toLocaleString()} keys + +
+ {isFiltering && totalKeys > 0 && shownKeys === 0 ? ( +

+ No keys match "{query.trim()}" in this date range +

+ ) : ( + + )} +
+ ); +}; + +export default KeyActivityPanel; diff --git a/ui/litellm-dashboard/src/components/UsagePage/keyActivityFilter.test.ts b/ui/litellm-dashboard/src/components/UsagePage/keyActivityFilter.test.ts new file mode 100644 index 00000000000..ce181f6b0c6 --- /dev/null +++ b/ui/litellm-dashboard/src/components/UsagePage/keyActivityFilter.test.ts @@ -0,0 +1,91 @@ +import { describe, expect, it } from "vitest"; + +import { filterKeyActivity, keyActivityMatches } from "./keyActivityFilter"; +import type { KeyMetadata, ModelActivityData } from "./types"; + +function activity(label: string, key_metadata?: KeyMetadata): ModelActivityData { + return { + label, + key_metadata, + total_requests: 1, + total_successful_requests: 1, + total_failed_requests: 0, + total_cache_read_input_tokens: 0, + total_cache_creation_input_tokens: 0, + total_tokens: 10, + prompt_tokens: 5, + completion_tokens: 5, + total_spend: 0.01, + top_api_keys: [], + top_models: [], + daily_data: [], + }; +} + +const aliceMeta: KeyMetadata = { + key_alias: "alice-batch", + team_id: "team-research", + user_id: "user-alice-1234", + user_email: "alice@example.com", +}; +const bobMeta: KeyMetadata = { + key_alias: null, + team_id: "team-research", + user_id: "user-bob-5678", + user_email: "bob@example.com", +}; +const alice = activity("alice-batch (team: research)", aliceMeta); +const bob = activity("bob@example.com (team: research)", bobMeta); +const orphan = activity("key-hash-deadbeef", { key_alias: null, team_id: null }); + +const keyMetrics: Record = { + "hash-alice": alice, + "hash-bob": bob, + deadbeef: orphan, +}; + +describe("keyActivityMatches", () => { + it("matches every key on an empty or whitespace query", () => { + expect(keyActivityMatches("deadbeef", orphan, "")).toBe(true); + expect(keyActivityMatches("deadbeef", orphan, " ")).toBe(true); + }); + + it("matches key alias case-insensitively", () => { + expect(keyActivityMatches("hash-alice", alice, "ALICE-batch")).toBe(true); + expect(keyActivityMatches("hash-bob", bob, "alice-batch")).toBe(false); + }); + + it("matches user email", () => { + expect(keyActivityMatches("hash-bob", bob, "bob@example")).toBe(true); + expect(keyActivityMatches("hash-alice", alice, "bob@example")).toBe(false); + }); + + it("matches user id", () => { + expect(keyActivityMatches("hash-alice", alice, "user-alice-1234")).toBe(true); + expect(keyActivityMatches("hash-bob", bob, "user-alice-1234")).toBe(false); + }); + + it("matches the key hash when the key has no alias or user metadata", () => { + expect(keyActivityMatches("deadbeef", orphan, "dead")).toBe(true); + expect(keyActivityMatches("deadbeef", orphan, "alice")).toBe(false); + }); + + it("trims surrounding whitespace from the query", () => { + expect(keyActivityMatches("hash-alice", alice, " alice@example.com ")).toBe(true); + }); +}); + +describe("filterKeyActivity", () => { + it("returns the same object when the query is blank", () => { + expect(filterKeyActivity(keyMetrics, "")).toBe(keyMetrics); + }); + + it("keeps only the keys matching the query, preserving their hashes", () => { + expect(Object.keys(filterKeyActivity(keyMetrics, "example.com"))).toEqual(["hash-alice", "hash-bob"]); + expect(filterKeyActivity(keyMetrics, "user-bob")).toEqual({ "hash-bob": bob }); + }); + + it("returns an empty record when nothing matches", () => { + expect(filterKeyActivity(keyMetrics, "nobody")).toEqual({}); + }); +}); diff --git a/ui/litellm-dashboard/src/components/UsagePage/keyActivityFilter.ts b/ui/litellm-dashboard/src/components/UsagePage/keyActivityFilter.ts new file mode 100644 index 00000000000..1e9a654fb2e --- /dev/null +++ b/ui/litellm-dashboard/src/components/UsagePage/keyActivityFilter.ts @@ -0,0 +1,20 @@ +import type { ModelActivityData } from "./types"; + +export function keyActivityMatches(apiKey: string, data: ModelActivityData, query: string): boolean { + const needle = query.trim().toLowerCase(); + if (needle === "") return true; + const meta = data.key_metadata; + return [apiKey, data.label, meta?.key_alias, meta?.user_id, meta?.user_email].some( + (field) => field?.toLowerCase().includes(needle) ?? false, + ); +} + +export function filterKeyActivity( + keyMetrics: Record, + query: string, +): Record { + if (query.trim() === "") return keyMetrics; + return Object.fromEntries( + Object.entries(keyMetrics).filter(([apiKey, data]) => keyActivityMatches(apiKey, data, query)), + ); +} diff --git a/ui/litellm-dashboard/src/components/UsagePage/types.ts b/ui/litellm-dashboard/src/components/UsagePage/types.ts index a10e9e68c4d..fd4f1350020 100644 --- a/ui/litellm-dashboard/src/components/UsagePage/types.ts +++ b/ui/litellm-dashboard/src/components/UsagePage/types.ts @@ -46,6 +46,7 @@ export interface KeyMetricWithMetadata { export interface KeyMetadata { key_alias: string | null; team_id: string | null; + user_id?: string | null; user_email?: string | null; tags?: { tag: string; usage: number }[]; } @@ -70,6 +71,7 @@ export interface TopModelData { export interface ModelActivityData { label: string; + key_metadata?: KeyMetadata; total_requests: number; total_successful_requests: number; total_failed_requests: number; diff --git a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx index c303dba697b..617b9209a41 100644 --- a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx +++ b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx @@ -289,6 +289,15 @@ it("should show a loading state on the initial load and hide the data", () => { expect(screen.queryByText("Test Key Alias")).not.toBeInTheDocument(); }); +it("replaces the previous rows with the loading state while a new search is pending", () => { + mockUseKeys.mockReturnValue(keysResult([mockKey], {}, { isPlaceholderData: true, isFetching: true })); + + renderWithProviders(); + + expect(screen.getByText("Loading keys...")).toBeInTheDocument(); + expect(screen.queryByText("Test Key Alias")).not.toBeInTheDocument(); +}); + it("should show 'No keys found' message when the key list is empty", () => { mockUseKeys.mockReturnValue(keysResult([])); diff --git a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx index d0367cb4d5b..907ee28bd05 100644 --- a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx +++ b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx @@ -128,7 +128,8 @@ export function VirtualKeysTable({ headerActions }: VirtualKeysTableProps) { const { data: keys, - isPending: isLoading, + isPending, + isPlaceholderData, isFetching, refetch, } = useKeys(tablePagination.pageIndex + 1, tablePagination.pageSize, keyListOptions); @@ -280,7 +281,7 @@ export function VirtualKeysTable({ headerActions }: VirtualKeysTableProps) { onColumnFiltersChange={handleColumnFiltersChange} enableColumnResizing columnResizeMode="onChange" - isLoading={isLoading} + isLoading={isPending || isPlaceholderData} loadingMessage="Loading keys..." noDataMessage="No keys found" fillHeight diff --git a/ui/litellm-dashboard/src/components/activity_metrics.test.tsx b/ui/litellm-dashboard/src/components/activity_metrics.test.tsx index b0fc8dc7866..914fe1872b6 100644 --- a/ui/litellm-dashboard/src/components/activity_metrics.test.tsx +++ b/ui/litellm-dashboard/src/components/activity_metrics.test.tsx @@ -655,6 +655,23 @@ describe("processActivityData", () => { expect(result["key1"].label).toBe("test-key-1 (team_id: team1)"); }); + it("retains the api key metadata so key activity can be searched by user", () => { + const metadata = { key_alias: "test-key-1", team_id: "team1", user_id: "user-1", user_email: "user1@example.com" }; + const withUser: { results: DailyData[] } = { + results: [ + createMockDailyData("2025-01-01", mockDailyActivity.results[0].metrics, { + ...EMPTY_BREAKDOWN, + api_keys: { key1: createMockKeyMetricWithMetadata(metadata, mockDailyActivity.results[0].metrics) }, + }), + ], + }; + + const result = processActivityData(withUser, "api_keys", MOCK_TEAMS); + + expect(result["key1"].key_metadata).toEqual(metadata); + expect(processActivityData(withUser, "models")["key1"]).toBeUndefined(); + }); + it("should process data for models key with data", () => { const dailyActivityWithModels: { results: DailyData[] } = { results: [ diff --git a/ui/litellm-dashboard/src/components/activity_metrics.tsx b/ui/litellm-dashboard/src/components/activity_metrics.tsx index f4348fb65ae..7c40a91be29 100644 --- a/ui/litellm-dashboard/src/components/activity_metrics.tsx +++ b/ui/litellm-dashboard/src/components/activity_metrics.tsx @@ -461,6 +461,7 @@ export const processActivityData = ( : key === "entities" ? (modelData as any).metadata?.agent_name || (modelData as any).metadata?.team_alias || model : model, + ...(key === "api_keys" ? { key_metadata: (modelData as KeyMetricWithMetadata).metadata } : {}), total_requests: 0, total_successful_requests: 0, total_failed_requests: 0, diff --git a/ui/litellm-dashboard/src/components/common_components/budget_duration_dropdown.tsx b/ui/litellm-dashboard/src/components/common_components/budget_duration_dropdown.tsx index 3c38907c597..40ee857634c 100644 --- a/ui/litellm-dashboard/src/components/common_components/budget_duration_dropdown.tsx +++ b/ui/litellm-dashboard/src/components/common_components/budget_duration_dropdown.tsx @@ -14,7 +14,7 @@ const DURATION_LABELS: Record = { interface BudgetDurationDropdownProps { id?: string; value?: string | null; - onChange?: (value: string | undefined) => void; + onChange?: (value: string | null) => void; className?: string; style?: React.CSSProperties; placeholder?: string; @@ -31,11 +31,7 @@ const BudgetDurationDropdown: React.FC = ({ showNeverResets = false, }) => { return ( - diff --git a/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx b/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx index 82580fd4667..b5789101f77 100644 --- a/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx +++ b/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx @@ -1021,7 +1021,7 @@ const CreateKey: React.FC = ({ team, teams, data, addKey, autoOp value={control.value as string | null | undefined} showNeverResets placeholder="Not set" - onChange={control.onChange} + onChange={(next) => control.onChange(next ?? undefined)} /> )} diff --git a/ui/litellm-dashboard/src/components/tag_management/types.tsx b/ui/litellm-dashboard/src/components/tag_management/types.tsx index 3cf17545fd6..88dfa28204d 100644 --- a/ui/litellm-dashboard/src/components/tag_management/types.tsx +++ b/ui/litellm-dashboard/src/components/tag_management/types.tsx @@ -41,7 +41,7 @@ export interface TagUpdateRequest { soft_budget?: number; tpm_limit?: number; rpm_limit?: number; - budget_duration?: string; + budget_duration?: string | null; } export interface TagDeleteRequest { diff --git a/ui/litellm-dashboard/src/components/team/EditMembership.tsx b/ui/litellm-dashboard/src/components/team/EditMembership.tsx index 036c4f1cc0b..909b5d56c97 100644 --- a/ui/litellm-dashboard/src/components/team/EditMembership.tsx +++ b/ui/litellm-dashboard/src/components/team/EditMembership.tsx @@ -157,7 +157,7 @@ const MemberModal = ({ onChange(next)} + onChange={(next) => onChange(mode === "add" ? next ?? undefined : next)} /> ); default: diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx index 8e3a17c2622..ffc83d0165e 100644 --- a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx @@ -1467,7 +1467,9 @@ const TeamInfoView: React.FC = ({ showNeverResets placeholder="Inherit team reset period" value={value === null ? NEVER_RESETS_BUDGET_DURATION : value} - onChange={(next) => onChange(next === NEVER_RESETS_BUDGET_DURATION ? null : next)} + onChange={(next) => + onChange(next === NEVER_RESETS_BUDGET_DURATION ? null : next ?? undefined) + } /> )} diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 9e844b992b2..c34975549e3 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -21,6 +21,26 @@ export interface paths { patch?: never; trace?: never; }; + "/.well-known/agent-skills/index.json": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** + * Agent Skills Index + * @description Agent Skills v0.2.0 discovery index over every skill stored on this proxy. + */ + get: operations["agent_skills_index__well_known_agent_skills_index_json_get"]; + put?: never; + post?: never; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; "/.well-known/jwks.json": { parameters: { query?: never; @@ -310,6 +330,26 @@ export interface paths { patch?: never; trace?: never; }; + "/.well-known/skills/index.json": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** + * Agent Skills Index + * @description Agent Skills v0.2.0 discovery index over every skill stored on this proxy. + */ + get: operations["agent_skills_index__well_known_skills_index_json_get"]; + put?: never; + post?: never; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; "/a2a/{agent_id}": { parameters: { query?: never; @@ -19982,6 +20022,26 @@ export interface paths { patch?: never; trace?: never; }; + "/v1/skills/{skill_id}/archive": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** + * Agent Skills Archive + * @description Stored skill upload, repacked so SKILL.md sits at the archive root. + */ + get: operations["agent_skills_archive_v1_skills__skill_id__archive_get"]; + put?: never; + post?: never; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; "/v1/threads": { parameters: { query?: never; @@ -23138,6 +23198,32 @@ export interface components { /** Tags */ tags?: string[]; }; + /** AgentSkillsIndex */ + AgentSkillsIndex: { + /** + * $Schema + * @default https://schemas.agentskills.io/discovery/0.2.0/schema.json + */ + $schema: string; + /** Skills */ + skills: components["schemas"]["AgentSkillsIndexEntry"][]; + }; + /** AgentSkillsIndexEntry */ + AgentSkillsIndexEntry: { + /** Description */ + description: string; + /** Digest */ + digest: string; + /** Name */ + name: string; + /** + * Type + * @constant + */ + type: "archive"; + /** Url */ + url: string; + }; /** * AlertType * @description Enum for alert types and management event types @@ -28374,6 +28460,8 @@ export interface components { team_id?: string | null; /** User Email */ user_email?: string | null; + /** User Id */ + user_id?: string | null; }; /** * KeyMetricWithMetadata @@ -40096,6 +40184,26 @@ export interface operations { }; }; }; + agent_skills_index__well_known_agent_skills_index_json_get: { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["AgentSkillsIndex"]; + }; + }; + }; + }; jwks_json__well_known_jwks_json_get: { parameters: { query?: never; @@ -40424,6 +40532,26 @@ export interface operations { }; }; }; + agent_skills_index__well_known_skills_index_json_get: { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["AgentSkillsIndex"]; + }; + }; + }; + }; invoke_agent_a2a_a2a__agent_id__post: { parameters: { query?: never; @@ -65154,6 +65282,37 @@ export interface operations { }; }; }; + agent_skills_archive_v1_skills__skill_id__archive_get: { + parameters: { + query?: never; + header?: never; + path: { + skill_id: string; + }; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/zip": string; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; create_threads_v1_threads_post: { parameters: { query?: never; diff --git a/uv.lock b/uv.lock index 9c659be658f..eb4cdef76f1 100644 --- a/uv.lock +++ b/uv.lock @@ -10,7 +10,7 @@ resolution-markers = [ ] [options] -exclude-newer = "2026-09-08T03:56:24.358378Z" +exclude-newer = "2026-09-09T21:39:49.468411Z" exclude-newer-span = "P3D" [manifest] @@ -4773,12 +4773,12 @@ proxy-dev = [ [[package]] name = "litellm-enterprise" -version = "0.1.66" +version = "0.1.67" source = { editable = "enterprise" } [[package]] name = "litellm-proxy-extras" -version = "0.4.96" +version = "0.4.97" source = { editable = "litellm-proxy-extras" } [[package]]