mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-19 00:01:29 +00:00
Merge remote-tracking branch 'origin/litellm_internal_staging' into litellm_team_callback_resolution
This commit is contained in:
commit
8536c5cdf6
76 changed files with 5426 additions and 492 deletions
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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<Vec<String>> {
|
||||
py.import("litellm.types.utils")?
|
||||
.getattr("CustomPricingLiteLLMParams")?
|
||||
.getattr("model_fields")?
|
||||
.cast_into::<PyDict>()?
|
||||
.keys()
|
||||
.iter()
|
||||
.map(|name| name.extract::<String>())
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn redact(
|
||||
py: Python<'_>,
|
||||
params: &Bound<'_, PyDict>,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
@ -310,6 +311,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,
|
||||
|
|
@ -344,6 +354,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.
|
||||
|
|
@ -558,6 +569,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"
|
||||
|
|
@ -1432,20 +1444,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
|
||||
|
|
@ -1665,6 +1666,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)
|
||||
|
|
@ -1898,16 +1900,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
|
||||
|
|
@ -1925,20 +1993,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
|
||||
|
|
|
|||
|
|
@ -30,8 +30,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 = (
|
||||
|
|
@ -39,6 +37,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,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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`.
|
||||
|
||||
|
|
|
|||
719
litellm/llms/meta/realtime/transformation.py
Normal file
719
litellm/llms/meta/realtime/transformation.py
Normal file
|
|
@ -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
|
||||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
@ -34717,6 +34719,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,
|
||||
|
|
@ -39019,7 +39037,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,
|
||||
|
|
@ -39167,18 +39187,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,
|
||||
|
|
@ -43488,6 +43510,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,
|
||||
|
|
@ -43780,6 +43803,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,
|
||||
|
|
@ -43794,6 +43818,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,
|
||||
|
|
@ -43906,6 +43931,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",
|
||||
|
|
@ -57470,6 +57496,14 @@
|
|||
"model_info": {
|
||||
"supports_reasoning": true
|
||||
}
|
||||
},
|
||||
{
|
||||
"name": "openai-reasoning-family-baseline",
|
||||
"pattern": "^(?!.*search-api)(?:[a-z0-9_.-]+/)*(?:ft:)?(?:o[1-9]\\d*(?![a-z0-9])|gpt-[5-9](?:\\.\\d+)?(?![0-9.])|(?:gpt-\\d+(?:\\.\\d+)?(?:-[a-z0-9]+)*-)?(?:codex|deep-research|chat-latest)(?![a-z0-9]))",
|
||||
"description": "OpenAI reasoning families by id shape, under any provider namespace and with an optional ft: prefix: the o-series (o1, o3-pro, o4-mini), gpt-5 through gpt-9 majors including dotted minors and suffixed variants (gpt-5.5-cyber, gpt-6-astra), and the codex, deep-research and chat-latest lines when standalone or on a gpt base. gpt-5-search-api is excluded because it is a search-only surface. Every model here is a reasoning model, and the Responses API drops the caller's reasoning param for any mapped OpenAI model whose info lacks supports_reasoning, so an id the registry has not named yet keeps its reasoning settings instead of silently losing them. Rules lose to exact entries. Carries no mode and no pricing, so cost stays on the standard unpriced behavior.",
|
||||
"model_info": {
|
||||
"supports_reasoning": true
|
||||
}
|
||||
}
|
||||
]
|
||||
},
|
||||
|
|
@ -59098,9 +59132,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/"
|
||||
},
|
||||
|
|
@ -59108,9 +59142,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/"
|
||||
},
|
||||
|
|
@ -60832,6 +60866,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,
|
||||
|
|
@ -60858,6 +60893,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",
|
||||
|
|
@ -60866,6 +60902,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,
|
||||
|
|
@ -60883,6 +60920,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",
|
||||
|
|
@ -60891,6 +60929,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",
|
||||
|
|
@ -60899,6 +60938,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",
|
||||
|
|
@ -60931,6 +60971,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",
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
]
|
||||
}
|
||||
}
|
||||
}
|
||||
},
|
||||
|
|
|
|||
|
|
@ -6,10 +6,13 @@ from typing import TYPE_CHECKING, Final
|
|||
|
||||
from pydantic import BaseModel, TypeAdapter
|
||||
|
||||
from litellm.proxy._types import LiteLLMRoutes
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.proxy.utils import PrismaClient
|
||||
|
||||
UNKNOWN_CALL_TYPE: Final = "Unknown"
|
||||
INFO_ROUTES_JSON: Final = json.dumps(LiteLLMRoutes.info_routes.value)
|
||||
|
||||
|
||||
class CacheActivityGroup(BaseModel):
|
||||
|
|
@ -69,6 +72,7 @@ GROUPS_SQL: Final = """
|
|||
OR COALESCE(vt."key_alias", 'Unnamed Key') IN (SELECT jsonb_array_elements_text($3::jsonb)))
|
||||
AND ($4::jsonb = '[]'::jsonb
|
||||
OR sl."model" IN (SELECT jsonb_array_elements_text($4::jsonb)))
|
||||
AND sl."call_type" NOT IN (SELECT jsonb_array_elements_text($5::jsonb))
|
||||
GROUP BY 1
|
||||
ORDER BY (COUNT(*)) DESC
|
||||
"""
|
||||
|
|
@ -89,6 +93,7 @@ ERROR_BREAKDOWN_SQL: Final = """
|
|||
OR COALESCE(vt."key_alias", 'Unnamed Key') IN (SELECT jsonb_array_elements_text($3::jsonb)))
|
||||
AND ($4::jsonb = '[]'::jsonb
|
||||
OR sl."model" IN (SELECT jsonb_array_elements_text($4::jsonb)))
|
||||
AND sl."call_type" NOT IN (SELECT jsonb_array_elements_text($5::jsonb))
|
||||
GROUP BY 1, 2, 3
|
||||
ORDER BY (COUNT(*)) DESC
|
||||
"""
|
||||
|
|
@ -100,6 +105,7 @@ KEY_ALIAS_OPTIONS_SQL: Final = """
|
|||
WHERE
|
||||
sl."startTime" >= ($1::timestamptz AT TIME ZONE 'UTC')
|
||||
AND sl."startTime" < (($2::timestamptz + INTERVAL '1 day') AT TIME ZONE 'UTC')
|
||||
AND sl."call_type" NOT IN (SELECT jsonb_array_elements_text($3::jsonb))
|
||||
ORDER BY 1
|
||||
"""
|
||||
|
||||
|
|
@ -110,6 +116,7 @@ MODEL_OPTIONS_SQL: Final = """
|
|||
sl."startTime" >= ($1::timestamptz AT TIME ZONE 'UTC')
|
||||
AND sl."startTime" < (($2::timestamptz + INTERVAL '1 day') AT TIME ZONE 'UTC')
|
||||
AND sl."model" != ''
|
||||
AND sl."call_type" NOT IN (SELECT jsonb_array_elements_text($3::jsonb))
|
||||
ORDER BY 1
|
||||
"""
|
||||
|
||||
|
|
@ -152,10 +159,12 @@ async def get_cache_activity(
|
|||
key_aliases_json: Final = json.dumps(list(key_aliases))
|
||||
models_json: Final = json.dumps(list(models))
|
||||
group_rows, error_rows, key_alias_rows, model_rows = await asyncio.gather(
|
||||
prisma_client.db.query_raw(GROUPS_SQL, start_date, end_date, key_aliases_json, models_json),
|
||||
prisma_client.db.query_raw(ERROR_BREAKDOWN_SQL, start_date, end_date, key_aliases_json, models_json),
|
||||
prisma_client.db.query_raw(KEY_ALIAS_OPTIONS_SQL, start_date, end_date),
|
||||
prisma_client.db.query_raw(MODEL_OPTIONS_SQL, start_date, end_date),
|
||||
prisma_client.db.query_raw(GROUPS_SQL, start_date, end_date, key_aliases_json, models_json, INFO_ROUTES_JSON),
|
||||
prisma_client.db.query_raw(
|
||||
ERROR_BREAKDOWN_SQL, start_date, end_date, key_aliases_json, models_json, INFO_ROUTES_JSON
|
||||
),
|
||||
prisma_client.db.query_raw(KEY_ALIAS_OPTIONS_SQL, start_date, end_date, INFO_ROUTES_JSON),
|
||||
prisma_client.db.query_raw(MODEL_OPTIONS_SQL, start_date, end_date, INFO_ROUTES_JSON),
|
||||
)
|
||||
groups: Final = _groups_adapter.validate_python(group_rows or [])
|
||||
return CacheActivityResponse(
|
||||
|
|
|
|||
132
litellm/proxy/common_utils/config_includes.py
Normal file
132
litellm/proxy/common_utils/config_includes.py
Normal file
|
|
@ -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
|
||||
|
|
@ -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}")
|
||||
|
|
|
|||
|
|
@ -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"]
|
||||
|
|
|
|||
130
litellm/proxy/discovery_endpoints/agent_skills_archive.py
Normal file
130
litellm/proxy/discovery_endpoints/agent_skills_archive.py
Normal file
|
|
@ -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()
|
||||
200
litellm/proxy/discovery_endpoints/agent_skills_endpoints.py
Normal file
200
litellm/proxy/discovery_endpoints/agent_skills_endpoints.py
Normal file
|
|
@ -0,0 +1,200 @@
|
|||
"""Serve skills stored on the proxy as an Agent Skills well-known discovery index.
|
||||
|
||||
``npx skills add <proxy url> -a <agent>`` 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]
|
||||
|
|
@ -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"),
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -5973,7 +5973,7 @@ async def list_keys(
|
|||
key_hash: str | None = Query(None, description="Filter keys by key hash"),
|
||||
key_alias: str | None = Query(
|
||||
None,
|
||||
description="Filter keys by key alias. Exact match by default; set substring_matching=true (admin only) for case-insensitive substring matching.",
|
||||
description="Filter keys by key alias. Exact match by default; set substring_matching=true for case-insensitive substring matching.",
|
||||
),
|
||||
search: str | None = Query(
|
||||
None,
|
||||
|
|
@ -5994,7 +5994,7 @@ async def list_keys(
|
|||
agent_id: str | None = Query(None, description="Filter keys by agent ID"),
|
||||
substring_matching: bool = Query(
|
||||
False,
|
||||
description="If true (proxy admins only), match user_id/key_alias as case-insensitive substrings instead of exact values. Defaults to false: /key/list matched these exactly before substring search was added, and an exact user_id/key_alias filter must never return another user's keys.",
|
||||
description="If true, match key_alias (any caller) and user_id (proxy admins only) as case-insensitive substrings instead of exact values. Defaults to false: /key/list matched these exactly before substring search was added, and an exact user_id filter must never return another user's keys.",
|
||||
),
|
||||
expires: str | None = Query(
|
||||
None,
|
||||
|
|
@ -6088,13 +6088,14 @@ async def list_keys(
|
|||
LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY.value,
|
||||
]
|
||||
|
||||
# Substring matching is opt-in (admin-only). /key/list matched user_id and
|
||||
# key_alias exactly before substring search was added; auto-applying a
|
||||
# substring match to every admin call broke that contract and let a caller
|
||||
# passing an exact user_id (e.g. an integration scoping to one user with an
|
||||
# admin key) receive other users' keys (user_id="alice" -> "alice2"). Exact
|
||||
# by default restores the prior behavior; the dashboard opts in explicitly.
|
||||
# Substring matching is opt-in. /key/list matched user_id and key_alias
|
||||
# exactly before substring search was added; auto-applying a substring
|
||||
# match to every admin call broke that contract and let a caller passing
|
||||
# an exact user_id (e.g. an integration scoping to one user with an admin
|
||||
# key) receive other users' keys (user_id="alice" -> "alice2"). Exact by
|
||||
# default restores the prior behavior; the dashboard opts in explicitly.
|
||||
use_substring_matching: Final = substring_matching and is_proxy_admin
|
||||
use_key_alias_substring_matching: Final = substring_matching
|
||||
|
||||
# Admins may omit user_id to list all keys; non-admins are scoped to self.
|
||||
if not user_id and not is_proxy_admin:
|
||||
|
|
@ -6121,6 +6122,7 @@ async def list_keys(
|
|||
access_group_id=access_group_id,
|
||||
agent_id=agent_id,
|
||||
use_substring_matching=use_substring_matching,
|
||||
use_key_alias_substring_matching=use_key_alias_substring_matching,
|
||||
expires_filter=expires if isinstance(expires, str) else None,
|
||||
search=search,
|
||||
)
|
||||
|
|
@ -6366,6 +6368,7 @@ def _build_key_filter_conditions(
|
|||
access_group_id: str | None = None,
|
||||
agent_id: str | None = None,
|
||||
use_substring_matching: bool = False,
|
||||
use_key_alias_substring_matching: bool = False,
|
||||
expires_filter: str | None = None,
|
||||
search: str | None = None,
|
||||
) -> Mapping[str, object]:
|
||||
|
|
@ -6461,7 +6464,7 @@ def _build_key_filter_conditions(
|
|||
*(
|
||||
(
|
||||
{"key_alias": {"contains": key_alias, "mode": "insensitive"}}
|
||||
if use_substring_matching
|
||||
if use_key_alias_substring_matching
|
||||
else {"key_alias": key_alias},
|
||||
)
|
||||
if key_alias and isinstance(key_alias, str)
|
||||
|
|
@ -6507,6 +6510,7 @@ async def _list_key_helper(
|
|||
access_group_id: str | None = None,
|
||||
agent_id: str | None = None,
|
||||
use_substring_matching: bool = False,
|
||||
use_key_alias_substring_matching: bool = False,
|
||||
expires_filter: str | None = None,
|
||||
search: str | None = None,
|
||||
) -> KeyListResponseObject:
|
||||
|
|
@ -6546,6 +6550,7 @@ async def _list_key_helper(
|
|||
access_group_id=access_group_id,
|
||||
agent_id=agent_id,
|
||||
use_substring_matching=use_substring_matching,
|
||||
use_key_alias_substring_matching=use_key_alias_substring_matching,
|
||||
expires_filter=expires_filter,
|
||||
search=search,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
||||
|
|
|
|||
|
|
@ -93,6 +93,7 @@ from litellm.integrations.custom_logger import CustomLogger
|
|||
from litellm.integrations.prometheus import PrometheusLogger
|
||||
from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting
|
||||
from litellm.integrations.SlackAlerting.utils import _add_langfuse_trace_id_to_alert
|
||||
from litellm.litellm_core_utils.api_route_to_call_types import get_call_types_for_route
|
||||
from litellm.litellm_core_utils.core_helpers import (
|
||||
coerce_token_limit,
|
||||
get_or_create_metadata_bucket,
|
||||
|
|
@ -879,6 +880,18 @@ def _failure_usage_to_lift(
|
|||
_EMPTY_LIFT: Final = MappingProxyType({})
|
||||
|
||||
|
||||
def _call_type_for_route(route: str | None) -> str | None:
|
||||
"""The route's call type when it maps to a single operation (its async and sync variants);
|
||||
None for routes shared by several operations, since the method is not known here."""
|
||||
if route is None:
|
||||
return None
|
||||
call_types: Final = get_call_types_for_route(route)
|
||||
if not call_types:
|
||||
return None
|
||||
operations: Final = frozenset(call_type.value.removeprefix("a") for call_type in call_types)
|
||||
return call_types[0].value if len(operations) == 1 else None
|
||||
|
||||
|
||||
def _failure_fields_to_lift(request_data: Mapping[str, object]) -> Mapping[str, object]:
|
||||
"""Failure-path callbacks run after ``litellm_logging_obj`` is popped from
|
||||
request_data (it is not serialisable), so the caller merges these fields
|
||||
|
|
@ -2549,10 +2562,6 @@ class ProxyLogging:
|
|||
|
||||
@staticmethod
|
||||
def _stream_requires_guardrail_translation(user_api_key_dict: UserAPIKeyAuth) -> bool:
|
||||
from litellm.litellm_core_utils.api_route_to_call_types import (
|
||||
get_call_types_for_route,
|
||||
)
|
||||
|
||||
route: Final = user_api_key_dict.request_route
|
||||
if not route:
|
||||
return False
|
||||
|
|
@ -3020,6 +3029,7 @@ class ProxyLogging:
|
|||
start_time=datetime.now(),
|
||||
**request_data,
|
||||
)
|
||||
request_data["litellm_logging_obj"] = litellm_logging_obj # rebind-ok: lifted then popped by the caller
|
||||
if "metadata" not in request_data:
|
||||
request_data["metadata"] = {}
|
||||
request_data["metadata"].update(user_api_key_logged_metadata)
|
||||
|
|
@ -3044,25 +3054,23 @@ class ProxyLogging:
|
|||
)
|
||||
|
||||
input: list | str | dict = ""
|
||||
normalized_call_type: str | None = None
|
||||
body_shape_call_type: str | None = None
|
||||
if "messages" in request_data and isinstance(request_data["messages"], list):
|
||||
input = request_data["messages"]
|
||||
litellm_logging_obj.model_call_details["messages"] = input
|
||||
if litellm_logging_obj.call_type != CallTypes.pass_through.value:
|
||||
normalized_call_type = CallTypes.acompletion.value
|
||||
body_shape_call_type = CallTypes.acompletion.value
|
||||
elif "prompt" in request_data and isinstance(request_data["prompt"], str):
|
||||
input = request_data["prompt"]
|
||||
litellm_logging_obj.model_call_details["prompt"] = input
|
||||
if litellm_logging_obj.call_type != CallTypes.pass_through.value:
|
||||
normalized_call_type = CallTypes.atext_completion.value
|
||||
body_shape_call_type = CallTypes.atext_completion.value
|
||||
elif "input" in request_data and isinstance(request_data["input"], list):
|
||||
input = request_data["input"]
|
||||
litellm_logging_obj.model_call_details["input"] = input
|
||||
if litellm_logging_obj.call_type != CallTypes.pass_through.value:
|
||||
normalized_call_type = CallTypes.aembedding.value
|
||||
if normalized_call_type is not None:
|
||||
litellm_logging_obj.call_type = normalized_call_type
|
||||
litellm_logging_obj.model_call_details["call_type"] = normalized_call_type
|
||||
body_shape_call_type = CallTypes.aembedding.value
|
||||
resolved_call_type: Final = _call_type_for_route(route) or body_shape_call_type
|
||||
if resolved_call_type is not None and litellm_logging_obj.call_type != CallTypes.pass_through.value:
|
||||
litellm_logging_obj.call_type = resolved_call_type
|
||||
litellm_logging_obj.model_call_details["call_type"] = resolved_call_type
|
||||
# Pass-through endpoints are logged via the callback loop's
|
||||
# async_post_call_failure_hook — skip pre_call and failure handlers.
|
||||
if litellm_logging_obj.call_type == CallTypes.pass_through.value:
|
||||
|
|
|
|||
21
litellm/types/llms/meta.py
Normal file
21
litellm/types/llms/meta.py
Normal file
|
|
@ -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, ...]]]
|
||||
|
|
@ -2190,6 +2190,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]
|
||||
|
|
@ -2204,6 +2251,7 @@ class OpenAIRealtimeInputAudioTranscriptionCompleted(TypedDict):
|
|||
item_id: ReadOnly[str]
|
||||
content_index: ReadOnly[int]
|
||||
transcript: ReadOnly[str]
|
||||
usage: NotRequired[ReadOnly[Mapping[str, object]]]
|
||||
|
||||
|
||||
class OpenAIRealtimeUsageTokenDetails(TypedDict):
|
||||
|
|
@ -2260,6 +2308,8 @@ OpenAIRealtimeEvents = (
|
|||
| OpenAIRealtimeInputAudioBufferSpeechEvent
|
||||
| OpenAIRealtimeInputAudioTranscriptionDelta
|
||||
| OpenAIRealtimeInputAudioTranscriptionCompleted
|
||||
| OpenAIRealtimeTranscriptionSessionCreated
|
||||
| OpenAIRealtimeErrorEvent
|
||||
)
|
||||
|
||||
OpenAIRealtimeStreamList = list[OpenAIRealtimeEvents]
|
||||
|
|
|
|||
|
|
@ -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, ...]
|
||||
|
|
@ -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
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
)
|
||||
|
|
|
|||
|
|
@ -9285,6 +9285,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
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
@ -34717,6 +34719,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,
|
||||
|
|
@ -39019,7 +39037,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,
|
||||
|
|
@ -39167,18 +39187,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,
|
||||
|
|
@ -43488,6 +43510,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,
|
||||
|
|
@ -43780,6 +43803,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,
|
||||
|
|
@ -43794,6 +43818,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,
|
||||
|
|
@ -43906,6 +43931,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",
|
||||
|
|
@ -57470,6 +57496,14 @@
|
|||
"model_info": {
|
||||
"supports_reasoning": true
|
||||
}
|
||||
},
|
||||
{
|
||||
"name": "openai-reasoning-family-baseline",
|
||||
"pattern": "^(?!.*search-api)(?:[a-z0-9_.-]+/)*(?:ft:)?(?:o[1-9]\\d*(?![a-z0-9])|gpt-[5-9](?:\\.\\d+)?(?![0-9.])|(?:gpt-\\d+(?:\\.\\d+)?(?:-[a-z0-9]+)*-)?(?:codex|deep-research|chat-latest)(?![a-z0-9]))",
|
||||
"description": "OpenAI reasoning families by id shape, under any provider namespace and with an optional ft: prefix: the o-series (o1, o3-pro, o4-mini), gpt-5 through gpt-9 majors including dotted minors and suffixed variants (gpt-5.5-cyber, gpt-6-astra), and the codex, deep-research and chat-latest lines when standalone or on a gpt base. gpt-5-search-api is excluded because it is a search-only surface. Every model here is a reasoning model, and the Responses API drops the caller's reasoning param for any mapped OpenAI model whose info lacks supports_reasoning, so an id the registry has not named yet keeps its reasoning settings instead of silently losing them. Rules lose to exact entries. Carries no mode and no pricing, so cost stays on the standard unpriced behavior.",
|
||||
"model_info": {
|
||||
"supports_reasoning": true
|
||||
}
|
||||
}
|
||||
]
|
||||
},
|
||||
|
|
@ -59098,9 +59132,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/"
|
||||
},
|
||||
|
|
@ -59108,9 +59142,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/"
|
||||
},
|
||||
|
|
@ -60832,6 +60866,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,
|
||||
|
|
@ -60858,6 +60893,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",
|
||||
|
|
@ -60866,6 +60902,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,
|
||||
|
|
@ -60883,6 +60920,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",
|
||||
|
|
@ -60891,6 +60929,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",
|
||||
|
|
@ -60899,6 +60938,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",
|
||||
|
|
@ -60931,6 +60971,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",
|
||||
|
|
|
|||
102
tests/proxy_behavior/spend/test_cache_activity.py
Normal file
102
tests/proxy_behavior/spend/test_cache_activity.py
Normal file
|
|
@ -0,0 +1,102 @@
|
|||
"""
|
||||
Behavior tests for the cache analytics queries against a real Postgres. The info-route
|
||||
exclusion and the Unknown grouping live in SQL, so these tests are the ones that exercise
|
||||
them; the endpoint wiring is unit-tested in
|
||||
tests/test_litellm/proxy/analytics_endpoints/test_analytics_endpoints.py.
|
||||
"""
|
||||
|
||||
import json
|
||||
import uuid
|
||||
from datetime import datetime
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm.proxy.analytics_endpoints.cache_activity import (
|
||||
ERROR_BREAKDOWN_SQL,
|
||||
GROUPS_SQL,
|
||||
INFO_ROUTES_JSON,
|
||||
KEY_ALIAS_OPTIONS_SQL,
|
||||
MODEL_OPTIONS_SQL,
|
||||
)
|
||||
|
||||
pytestmark = pytest.mark.asyncio(loop_scope="session")
|
||||
|
||||
DAY: Final = datetime(2001, 3, 7)
|
||||
AT_NOON: Final = DAY.replace(hour=12)
|
||||
RUN: Final = uuid.uuid4()
|
||||
INFERENCE_KEY: Final = f"ca-inference-{RUN}"
|
||||
INFO_ONLY_KEY: Final = f"ca-info-only-{RUN}"
|
||||
INFERENCE_ALIAS: Final = f"alias-inference-{RUN}"
|
||||
INFO_ONLY_ALIAS: Final = f"alias-info-only-{RUN}"
|
||||
INFERENCE_MODEL: Final = f"gpt-5.4-mini-{RUN}"
|
||||
INFO_ONLY_MODEL: Final = f"ghost-model-{RUN}"
|
||||
NO_FILTER: Final = "[]"
|
||||
|
||||
|
||||
async def _spend_log(db, api_key: str, call_type: str, status: str, model: str = "", error_code: str = "") -> None:
|
||||
metadata: Final = {"error_information": {"error_code": error_code, "error_class": "ProxyException"}}
|
||||
await db.execute_raw(
|
||||
'INSERT INTO "LiteLLM_SpendLogs" ("request_id", "call_type", "api_key", "startTime", "endTime", "model", '
|
||||
'"status", "metadata") VALUES ($1, $2, $3, $4::timestamp, $4::timestamp, $5, $6, $7::jsonb)',
|
||||
str(uuid.uuid4()),
|
||||
call_type,
|
||||
api_key,
|
||||
AT_NOON,
|
||||
model,
|
||||
status,
|
||||
json.dumps(metadata if status == "failure" else {}),
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture(scope="module", autouse=True)
|
||||
async def seeded(db):
|
||||
for token, alias in ((INFERENCE_KEY, INFERENCE_ALIAS), (INFO_ONLY_KEY, INFO_ONLY_ALIAS)):
|
||||
await db.execute_raw(
|
||||
'INSERT INTO "LiteLLM_VerificationToken" ("token", "key_alias") VALUES ($1, $2)', token, alias
|
||||
)
|
||||
await _spend_log(db, INFERENCE_KEY, "acompletion", "success", model=INFERENCE_MODEL)
|
||||
await _spend_log(db, INFERENCE_KEY, "acompletion", "failure", model=INFERENCE_MODEL, error_code="429")
|
||||
await _spend_log(db, INFERENCE_KEY, "", "failure", error_code="401")
|
||||
await _spend_log(db, INFERENCE_KEY, "/model/info", "failure", error_code="401")
|
||||
await _spend_log(db, INFO_ONLY_KEY, "/v1/models", "failure", model=INFO_ONLY_MODEL, error_code="401")
|
||||
await _spend_log(db, INFO_ONLY_KEY, "/key/info", "success")
|
||||
yield
|
||||
keys: Final = [INFERENCE_KEY, INFO_ONLY_KEY]
|
||||
await db.execute_raw('DELETE FROM "LiteLLM_SpendLogs" WHERE "api_key" = ANY($1::text[])', keys)
|
||||
await db.execute_raw('DELETE FROM "LiteLLM_VerificationToken" WHERE "token" = ANY($1::text[])', keys)
|
||||
|
||||
|
||||
async def _groups(db, key_aliases: list[str]) -> dict[str, dict]:
|
||||
rows: Final = await db.query_raw(GROUPS_SQL, DAY, DAY, json.dumps(key_aliases), NO_FILTER, INFO_ROUTES_JSON)
|
||||
return {row["call_type"]: row for row in rows}
|
||||
|
||||
|
||||
async def test_groups_drop_info_routes_and_keep_unknown_for_rows_without_an_endpoint(db):
|
||||
groups: Final = await _groups(db, [INFERENCE_ALIAS])
|
||||
assert set(groups) == {"acompletion", "Unknown"}
|
||||
assert (groups["acompletion"]["api_requests"], groups["acompletion"]["failed_requests"]) == (1, 1)
|
||||
assert (groups["Unknown"]["api_requests"], groups["Unknown"]["failed_requests"]) == (0, 1)
|
||||
|
||||
|
||||
async def test_key_with_only_info_route_traffic_has_no_groups(db):
|
||||
assert await _groups(db, [INFO_ONLY_ALIAS]) == {}
|
||||
|
||||
|
||||
async def test_error_breakdown_drops_info_routes(db):
|
||||
rows: Final = await db.query_raw(
|
||||
ERROR_BREAKDOWN_SQL, DAY, DAY, json.dumps([INFERENCE_ALIAS, INFO_ONLY_ALIAS]), NO_FILTER, INFO_ROUTES_JSON
|
||||
)
|
||||
assert {(row["call_type"], row["error_code"], row["count"]) for row in rows} == {
|
||||
("acompletion", "429", 1),
|
||||
("Unknown", "401", 1),
|
||||
}
|
||||
|
||||
|
||||
async def test_filter_options_only_offer_values_that_return_analytics(db):
|
||||
key_alias_rows: Final = await db.query_raw(KEY_ALIAS_OPTIONS_SQL, DAY, DAY, INFO_ROUTES_JSON)
|
||||
model_rows: Final = await db.query_raw(MODEL_OPTIONS_SQL, DAY, DAY, INFO_ROUTES_JSON)
|
||||
key_aliases: Final = {row["key_alias"] for row in key_alias_rows}
|
||||
models: Final = {row["model"] for row in model_rows}
|
||||
assert INFERENCE_ALIAS in key_aliases and INFO_ONLY_ALIAS not in key_aliases
|
||||
assert INFERENCE_MODEL in models and INFO_ONLY_MODEL not in models
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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."""
|
||||
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -677,3 +677,48 @@ def test_deployment_model_info_beats_the_seeded_rule_defaults(shipped_cost_map):
|
|||
|
||||
assert litellm.model_cost[model]["supports_reasoning"] is False
|
||||
assert litellm.supports_reasoning(model="some-org/NoThink-1", custom_llm_provider="wandb") is False
|
||||
|
||||
|
||||
def test_shipped_rules_flag_unmapped_openai_reasoning_families(shipped_cost_map):
|
||||
for model in (
|
||||
"gpt-5.7-nova",
|
||||
"openai/gpt-6",
|
||||
"ft:gpt-5.1-2025-11-13:org::abc",
|
||||
"o5-mini",
|
||||
"gpt-5.6-codex-max",
|
||||
"o4-mini-deep-research-2027-01-01",
|
||||
"gpt-5.7-chat-latest",
|
||||
"azure/gpt-5.7-cyber",
|
||||
"openai/codex-mini-latest-2027",
|
||||
):
|
||||
assert model not in litellm.model_cost, model
|
||||
assert match_capability_generalizations(model) == {"supports_reasoning": True}, model
|
||||
info = litellm.get_model_info("gpt-5.7-nova", custom_llm_provider="openai")
|
||||
assert info["litellm_provider"] == "openai"
|
||||
assert info["supports_reasoning"] is True
|
||||
assert info.get("mode") is None
|
||||
assert not info.get("input_cost_per_token")
|
||||
assert litellm.supports_reasoning(model="gpt-5.7-nova", custom_llm_provider="openai") is True
|
||||
|
||||
|
||||
def test_shipped_openai_reasoning_rule_skips_non_reasoning_gpt_ids(shipped_cost_map):
|
||||
for model in (
|
||||
"gpt-4o",
|
||||
"gpt-4.1-nano-new",
|
||||
"gpt-oss-120b",
|
||||
"gpt-realtime-2027",
|
||||
"gpt-image-2",
|
||||
"gpt-5-search-api-2027-01-01",
|
||||
"omni-moderation-new",
|
||||
"text-embedding-4",
|
||||
"vendor/my-codex-embedding",
|
||||
"some-codex-model",
|
||||
"azure/gpt-35-turbo-0125-custom",
|
||||
"github_copilot/gpt-41-copilot-new",
|
||||
):
|
||||
assert match_capability_generalizations(model) is None, model
|
||||
|
||||
|
||||
def test_shipped_openai_reasoning_rule_loses_to_mapped_entries(shipped_cost_map):
|
||||
assert "gpt-5-search-api" in litellm.model_cost
|
||||
assert litellm.supports_reasoning(model="gpt-5-search-api", custom_llm_provider="openai") is False
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -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": ""}))
|
||||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -12,10 +12,13 @@ from unittest.mock import AsyncMock, MagicMock
|
|||
import pytest
|
||||
from fastapi import HTTPException
|
||||
|
||||
from litellm.proxy._types import LiteLLMRoutes
|
||||
from litellm.proxy.analytics_endpoints.analytics_endpoints import get_global_activity
|
||||
from litellm.proxy.analytics_endpoints.cache_activity import (
|
||||
ERROR_BREAKDOWN_SQL,
|
||||
GROUPS_SQL,
|
||||
KEY_ALIAS_OPTIONS_SQL,
|
||||
MODEL_OPTIONS_SQL,
|
||||
CacheActivityGroup,
|
||||
compute_totals,
|
||||
)
|
||||
|
|
@ -112,6 +115,21 @@ async def test_filters_are_passed_to_sql_as_json_arrays(mock_prisma: MagicMock):
|
|||
assert call.args[4] == json.dumps(["gpt-5.1", "claude-opus-4-8"])
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_every_query_excludes_the_same_info_routes(mock_prisma: MagicMock):
|
||||
"""Regression for LIT-5884: failed info-route calls are spend-logged but are not inference traffic, so
|
||||
the groups, error breakdown and both filter-option queries all receive the same exclusion list. What
|
||||
the SQL does with it is covered against Postgres in tests/proxy_behavior/spend/test_cache_activity.py."""
|
||||
await get_global_activity(start_date="2026-07-01", end_date="2026-07-27", key_aliases=[], models=[])
|
||||
|
||||
exclusions_by_query = {call.args[0]: json.loads(call.args[-1]) for call in mock_prisma.db.query_raw.call_args_list}
|
||||
assert set(exclusions_by_query) == {GROUPS_SQL, ERROR_BREAKDOWN_SQL, KEY_ALIAS_OPTIONS_SQL, MODEL_OPTIONS_SQL}
|
||||
for excluded_call_types in exclusions_by_query.values():
|
||||
assert excluded_call_types == LiteLLMRoutes.info_routes.value
|
||||
assert {"/model/info", "/v1/models", "/key/info"} <= set(excluded_call_types)
|
||||
assert "" not in excluded_call_types
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_rejects_malformed_dates_with_400(mock_prisma: MagicMock):
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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():
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -21,6 +21,7 @@ from litellm.proxy._types import (
|
|||
LiteLLM_BudgetTable,
|
||||
LiteLLM_OrganizationTable,
|
||||
LiteLLM_ProjectTableCachedObj,
|
||||
LiteLLM_TeamTable,
|
||||
LiteLLM_TeamTableCachedObj,
|
||||
LiteLLM_UserTable,
|
||||
LiteLLM_VerificationToken,
|
||||
|
|
@ -6416,7 +6417,7 @@ def test_build_key_filter_conditions_key_alias_narrows_team_admin_visibility():
|
|||
admin_team_ids=["team-a"],
|
||||
member_team_ids=["team-a"],
|
||||
include_created_by_keys=False,
|
||||
use_substring_matching=True,
|
||||
use_key_alias_substring_matching=True,
|
||||
)
|
||||
assert {"key_alias": {"contains": "member-key", "mode": "insensitive"}} in where_substring["AND"], (
|
||||
f"substring key_alias not ANDed: {where_substring}"
|
||||
|
|
@ -9361,7 +9362,7 @@ async def test_build_key_filter_team_id_scoped():
|
|||
async def test_build_key_filter_admin_substring_matching():
|
||||
"""
|
||||
Admin callers get substring (contains + insensitive) matching for user_id
|
||||
and key_alias when use_substring_matching=True.
|
||||
and key_alias when both substring flags are set.
|
||||
"""
|
||||
from litellm.proxy.management_endpoints.key_management_endpoints import (
|
||||
_build_key_filter_conditions,
|
||||
|
|
@ -9381,6 +9382,7 @@ async def test_build_key_filter_admin_substring_matching():
|
|||
member_team_ids=None,
|
||||
include_created_by_keys=False,
|
||||
use_substring_matching=True,
|
||||
use_key_alias_substring_matching=True,
|
||||
)
|
||||
|
||||
assert where["AND"][0]["user_id"] == {"contains": user_id, "mode": "insensitive"}
|
||||
|
|
@ -15152,8 +15154,8 @@ async def test_list_keys_admin_substring_opt_in():
|
|||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_keys_non_admin_cannot_opt_into_substring():
|
||||
"""substring_matching is admin-only: a non-admin requesting it still gets
|
||||
exact matching, scoped to their own user_id."""
|
||||
"""user_id substring matching is admin-only: a non-admin requesting it still
|
||||
gets exact matching, scoped to their own user_id."""
|
||||
user = UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, user_id="alice")
|
||||
kwargs = await _list_keys_capture_helper_kwargs(
|
||||
user, user_id=None, substring_matching=True
|
||||
|
|
@ -15162,6 +15164,108 @@ async def test_list_keys_non_admin_cannot_opt_into_substring():
|
|||
assert kwargs["user_id"] == "alice"
|
||||
|
||||
|
||||
def _prisma_where_matches(row, where):
|
||||
for field, expected in where.items():
|
||||
if field == "AND":
|
||||
if not all(_prisma_where_matches(row, child) for child in expected):
|
||||
return False
|
||||
elif field == "OR":
|
||||
if not any(_prisma_where_matches(row, child) for child in expected):
|
||||
return False
|
||||
elif isinstance(expected, dict):
|
||||
value = getattr(row, field)
|
||||
if "in" in expected and value not in expected["in"]:
|
||||
return False
|
||||
if "not" in expected and value == expected["not"]:
|
||||
return False
|
||||
if "contains" in expected:
|
||||
haystack, needle = value or "", expected["contains"]
|
||||
if expected.get("mode") == "insensitive":
|
||||
haystack, needle = haystack.lower(), needle.lower()
|
||||
if needle not in haystack:
|
||||
return False
|
||||
elif getattr(row, field) != expected:
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
class _InMemoryVerificationTokenTable:
|
||||
def __init__(self, rows):
|
||||
self.rows = rows
|
||||
|
||||
async def find_many(self, where, **kwargs):
|
||||
return [row for row in self.rows if _prisma_where_matches(row, where)]
|
||||
|
||||
async def count(self, where):
|
||||
return len(await self.find_many(where))
|
||||
|
||||
|
||||
def _team_key(token, key_alias, user_id):
|
||||
return LiteLLM_VerificationToken(token=token, key_alias=key_alias, user_id=user_id, team_id="team-a")
|
||||
|
||||
|
||||
_TEAM_A_KEYS = (
|
||||
_team_key("tok-alice-first", "app_llmhub_first.last", "alice"),
|
||||
_team_key("tok-alice-other", "alice_other_key", "alice"),
|
||||
_team_key("tok-bob-first", "bob_First_key", "bob"),
|
||||
_team_key("tok-svc-first", "service_first_key", None),
|
||||
)
|
||||
|
||||
|
||||
def _list_team_a_keys_as(user_role, members_with_roles, query):
|
||||
from fastapi import FastAPI
|
||||
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
from litellm.proxy.management_endpoints.key_management_endpoints import router
|
||||
|
||||
mock_prisma_client = AsyncMock()
|
||||
mock_prisma_client.db.litellm_verificationtoken = _InMemoryVerificationTokenTable(_TEAM_A_KEYS)
|
||||
mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(
|
||||
return_value=LiteLLM_UserTable(user_id="alice", teams=["team-a"], organization_memberships=[])
|
||||
)
|
||||
mock_prisma_client.db.litellm_teamtable.find_many = AsyncMock(
|
||||
return_value=[LiteLLM_TeamTable(team_id="team-a", members_with_roles=members_with_roles)]
|
||||
)
|
||||
test_app = FastAPI()
|
||||
test_app.include_router(router)
|
||||
test_app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth(user_role=user_role, user_id="alice")
|
||||
with patch( # test-quality-ok: /key/list reads the prisma client from the proxy_server module global, no injection point
|
||||
"litellm.proxy.proxy_server.prisma_client", mock_prisma_client
|
||||
):
|
||||
response = TestClient(test_app).get(
|
||||
f"/key/list?team_id=team-a&include_team_keys=true&include_created_by_keys=true&{query}"
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
return sorted(response.json()["keys"])
|
||||
|
||||
|
||||
_ALICE_TEAM_ADMIN = [Member(user_id="alice", role="admin"), Member(user_id="bob", role="user")]
|
||||
_ALICE_TEAM_MEMBER = [Member(user_id="alice", role="user"), Member(user_id="bob", role="user")]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"user_role",
|
||||
[LitellmUserRoles.INTERNAL_USER, LitellmUserRoles.INTERNAL_USER_VIEW_ONLY, LitellmUserRoles.TEAM],
|
||||
)
|
||||
def test_list_keys_team_admin_key_alias_substring_returns_every_matching_team_key(user_role):
|
||||
keys = _list_team_a_keys_as(user_role, _ALICE_TEAM_ADMIN, "key_alias=first&substring_matching=true")
|
||||
assert keys == ["tok-alice-first", "tok-bob-first", "tok-svc-first"]
|
||||
|
||||
|
||||
def test_list_keys_team_member_key_alias_substring_stays_within_own_visibility():
|
||||
keys = _list_team_a_keys_as(
|
||||
LitellmUserRoles.INTERNAL_USER, _ALICE_TEAM_MEMBER, "key_alias=first&substring_matching=true"
|
||||
)
|
||||
assert keys == ["tok-alice-first", "tok-svc-first"]
|
||||
|
||||
|
||||
def test_list_keys_key_alias_stays_exact_without_substring_matching():
|
||||
assert _list_team_a_keys_as(LitellmUserRoles.INTERNAL_USER, _ALICE_TEAM_ADMIN, "key_alias=first") == []
|
||||
assert _list_team_a_keys_as(
|
||||
LitellmUserRoles.INTERNAL_USER, _ALICE_TEAM_ADMIN, "key_alias=app_llmhub_first.last"
|
||||
) == ["tok-alice-first"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_keys_search_is_honored_for_non_admin():
|
||||
"""LIT-4741: unlike substring_matching, `search` is not admin-gated. A non-admin's
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
|
|
|||
|
|
@ -33,9 +33,7 @@ def test_is_proxy_only_llm_api_truth_table(proxy_logging):
|
|||
snapshot. Covers no-route, non-LLM route, HTTPException on LLM route,
|
||||
and auth-error short-circuit."""
|
||||
snapshot = {
|
||||
"no_route": proxy_logging._is_proxy_only_llm_api_error(
|
||||
original_exception=Exception(), route=None
|
||||
),
|
||||
"no_route": proxy_logging._is_proxy_only_llm_api_error(original_exception=Exception(), route=None),
|
||||
"non_llm_route": proxy_logging._is_proxy_only_llm_api_error(
|
||||
original_exception=HTTPException(status_code=429, detail="rate"),
|
||||
route="/random/path",
|
||||
|
|
@ -158,9 +156,7 @@ async def test_post_call_failure_hook_non_http_exception_in_callback_swallowed(
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_handle_logging_proxy_only_path_uses_existing_logging_obj(
|
||||
proxy_logging, make_user_api_key_auth
|
||||
):
|
||||
async def test_handle_logging_proxy_only_path_uses_existing_logging_obj(proxy_logging, make_user_api_key_auth):
|
||||
logging_obj = MagicMock()
|
||||
logging_obj.call_type = "acompletion"
|
||||
logging_obj.model_call_details = {}
|
||||
|
|
@ -183,10 +179,7 @@ async def test_handle_logging_proxy_only_path_uses_existing_logging_obj(
|
|||
snapshot = {
|
||||
"input_logged": "messages" in logging_obj.model_call_details,
|
||||
"call_type_normalized": logging_obj.call_type,
|
||||
"marker_present": logging_obj.model_call_details.get(
|
||||
LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL
|
||||
)
|
||||
is True,
|
||||
"marker_present": logging_obj.model_call_details.get(LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL) is True,
|
||||
"async_failure_called": logging_obj.async_failure_handler.called,
|
||||
}
|
||||
assert snapshot == {
|
||||
|
|
@ -198,9 +191,7 @@ async def test_handle_logging_proxy_only_path_uses_existing_logging_obj(
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_handle_logging_proxy_only_path_skips_for_pass_through(
|
||||
proxy_logging, make_user_api_key_auth
|
||||
):
|
||||
async def test_handle_logging_proxy_only_path_skips_for_pass_through(proxy_logging, make_user_api_key_auth):
|
||||
from litellm.types.utils import CallTypes
|
||||
|
||||
logging_obj = MagicMock()
|
||||
|
|
@ -248,9 +239,7 @@ async def test_handle_logging_proxy_only_path_no_logging_obj_creates_one(
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_handle_logging_proxy_only_path_propagates_async_failure_raises(
|
||||
proxy_logging, make_user_api_key_auth
|
||||
):
|
||||
async def test_handle_logging_proxy_only_path_propagates_async_failure_raises(proxy_logging, make_user_api_key_auth):
|
||||
logging_obj = MagicMock()
|
||||
logging_obj.call_type = "acompletion"
|
||||
logging_obj.model_call_details = {}
|
||||
|
|
@ -267,3 +256,65 @@ async def test_handle_logging_proxy_only_path_propagates_async_failure_raises(
|
|||
route="/chat/completions",
|
||||
original_exception=Exception("x"),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"route, request_data, expected_call_type",
|
||||
[
|
||||
("/v1/chat/completions", {}, "acompletion"),
|
||||
("/chat/completions", {"model": "m", "messages": [{"role": "user", "content": "hi"}]}, "acompletion"),
|
||||
("/v1/messages", {"model": "m", "messages": [{"role": "user", "content": "hi"}]}, "anthropic_messages"),
|
||||
("/v1/responses", {"model": "m", "input": "hi"}, "aresponses"),
|
||||
("/v1/embeddings", {"model": "m", "input": ["hi"]}, "aembedding"),
|
||||
("/model/info", {}, "/model/info"),
|
||||
],
|
||||
)
|
||||
async def test_post_call_failure_hook_lifts_route_call_type_for_gate_rejections(
|
||||
proxy_logging, make_user_api_key_auth, route, request_data, expected_call_type
|
||||
):
|
||||
"""Regression for LIT-5884: the matched route, not the body shape, sets the
|
||||
spend-log call_type for requests rejected before dispatch."""
|
||||
proxy_logging.alert_types = []
|
||||
await proxy_logging.post_call_failure_hook(
|
||||
request_data=request_data,
|
||||
original_exception=Exception("Authentication Error, No api key passed in."),
|
||||
user_api_key_dict=make_user_api_key_auth(request_route=route),
|
||||
error_type=ProxyErrorTypes.auth_error,
|
||||
route=route,
|
||||
)
|
||||
assert request_data["call_type"] == expected_call_type
|
||||
assert "start_time" in request_data
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_post_call_failure_hook_falls_back_to_body_shape_without_a_route(proxy_logging, make_user_api_key_auth):
|
||||
proxy_logging.alert_types = []
|
||||
request_data = {"model": "m", "messages": [{"role": "user", "content": "hi"}]}
|
||||
await proxy_logging.post_call_failure_hook(
|
||||
request_data=request_data,
|
||||
original_exception=Exception("Authentication Error, No api key passed in."),
|
||||
user_api_key_dict=make_user_api_key_auth(request_route="/chat/completions"),
|
||||
error_type=ProxyErrorTypes.auth_error,
|
||||
)
|
||||
assert request_data["call_type"] == "acompletion"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("route", ["/v1/files", "/files/file-abc", "/v1/containers"])
|
||||
async def test_post_call_failure_hook_keeps_the_route_for_multi_operation_routes(
|
||||
proxy_logging, make_user_api_key_auth, route
|
||||
):
|
||||
"""Routes shared by several operations (POST create vs GET list) cannot be attributed without the
|
||||
method, so a rejected request there is filed under its route, not under whichever operation the
|
||||
mapping lists first."""
|
||||
proxy_logging.alert_types = []
|
||||
request_data: dict = {}
|
||||
await proxy_logging.post_call_failure_hook(
|
||||
request_data=request_data,
|
||||
original_exception=Exception("Authentication Error, No api key passed in."),
|
||||
user_api_key_dict=make_user_api_key_auth(request_route=route),
|
||||
error_type=ProxyErrorTypes.auth_error,
|
||||
route=route,
|
||||
)
|
||||
assert request_data["call_type"] == route
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -17,6 +17,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 (
|
||||
|
|
@ -4768,3 +4770,228 @@ def test_collect_and_combine_realtime_usage_stores_partitioned_text_tokens() ->
|
|||
assert combined.completion_tokens_details.reasoning_tokens == 95
|
||||
assert combined.completion_tokens_details.text_tokens == 38
|
||||
assert combined.completion_tokens_details.audio_tokens == 0
|
||||
|
||||
|
||||
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
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -236,6 +236,35 @@ describe("CacheDashboard cache analytics charts", () => {
|
|||
expect(screen.queryByText(/Failed requests by error code/)).not.toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("explains the Unknown bucket only when a group has no recorded endpoint", async () => {
|
||||
const { rerender } = renderDashboard();
|
||||
await screen.findByText(REQUESTS_CHART_TITLE);
|
||||
expect(screen.queryByText(/recorded no endpoint/)).not.toBeInTheDocument();
|
||||
|
||||
useCacheActivity.mockReturnValue({
|
||||
data: {
|
||||
...cacheActivity,
|
||||
groups: [
|
||||
...cacheActivity.groups,
|
||||
{
|
||||
call_type: "Unknown",
|
||||
api_requests: 0,
|
||||
cache_hits: 0,
|
||||
failed_requests: 121000,
|
||||
cached_completion_tokens: 0,
|
||||
generated_completion_tokens: 0,
|
||||
},
|
||||
],
|
||||
},
|
||||
refetch: vi.fn(),
|
||||
});
|
||||
rerender(<CacheDashboard accessToken="sk-test" token="tok" userRole="Admin" userID="u1" premiumUser={false} />);
|
||||
|
||||
expect(
|
||||
within(cardTitled(REQUESTS_CHART_TITLE)).getByText(/Unknown groups spend logs that recorded no endpoint/),
|
||||
).toHaveTextContent("not necessarily LLM API requests");
|
||||
});
|
||||
|
||||
it("formats y-axis ticks with compact notation", async () => {
|
||||
renderDashboard();
|
||||
const { requestsCard, tokensCard } = await findChartCards();
|
||||
|
|
|
|||
|
|
@ -35,6 +35,11 @@ const REQUEST_SERIES = {
|
|||
failed: "Failed requests",
|
||||
} as const;
|
||||
|
||||
const UNKNOWN_CALL_TYPE = "Unknown";
|
||||
|
||||
const UNKNOWN_CALL_TYPE_NOTE =
|
||||
"Unknown groups spend logs that recorded no endpoint. Older proxy versions wrote those for requests rejected before routing, so they are not necessarily LLM API requests.";
|
||||
|
||||
const toChartDatum = (group: CacheActivityGroup) => ({
|
||||
name: group.call_type,
|
||||
[REQUEST_SERIES.apiRequests]: group.api_requests,
|
||||
|
|
@ -103,6 +108,7 @@ const CacheDashboard: React.FC<CachePageProps> = ({ accessToken, token, userRole
|
|||
const uniqueApiKeys = activity?.filter_options.key_aliases ?? [];
|
||||
const uniqueModels = activity?.filter_options.models ?? [];
|
||||
const chartData = (activity?.groups ?? []).map(toChartDatum);
|
||||
const hasUnknownGroup = (activity?.groups ?? []).some((group) => group.call_type === UNKNOWN_CALL_TYPE);
|
||||
const activeDrilldownCallType = resolveDrilldownCallType(errorDrilldownCallType, activity?.groups ?? []);
|
||||
|
||||
const handleRefreshClick = () => {
|
||||
|
|
@ -288,6 +294,7 @@ const CacheDashboard: React.FC<CachePageProps> = ({ accessToken, token, userRole
|
|||
<p className="text-sm text-muted-foreground">
|
||||
Click a red failed-requests segment to see which error codes caused those failures.
|
||||
</p>
|
||||
{hasUnknownGroup && <p className="mt-1 text-sm text-muted-foreground">{UNKNOWN_CALL_TYPE_NOTE}</p>}
|
||||
<BarChart
|
||||
data={chartData}
|
||||
stack={true}
|
||||
|
|
|
|||
|
|
@ -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<EntityUsageProps> = ({
|
|||
{
|
||||
key: "keys",
|
||||
label: "Key Activity",
|
||||
content: <ActivityMetrics modelMetrics={keyMetrics} hidePromptCachingMetrics={entityType === "agent"} />,
|
||||
content: <KeyActivityPanel keyMetrics={keyMetrics} hidePromptCachingMetrics={entityType === "agent"} />,
|
||||
},
|
||||
{ key: "endpoints", label: "Endpoint Activity", content: <EndpointUsage userSpendData={spendData} /> },
|
||||
];
|
||||
|
|
|
|||
|
|
@ -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<UsagePageProps> = ({ teams, organizations }) => {
|
|||
<ActivityMetrics modelMetrics={modelMetrics} />
|
||||
</TabsContent>
|
||||
<TabsContent value="keys" keepMounted>
|
||||
<ActivityMetrics modelMetrics={keyMetrics} />
|
||||
<KeyActivityPanel keyMetrics={keyMetrics} />
|
||||
</TabsContent>
|
||||
<TabsContent value="mcp" keepMounted>
|
||||
<ActivityMetrics modelMetrics={mcpServerMetrics} />
|
||||
|
|
|
|||
|
|
@ -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<string, ModelActivityData> }) => (
|
||||
<ul data-testid="rendered-keys">
|
||||
{Object.keys(modelMetrics).map((hash) => (
|
||||
<li key={hash}>{hash}</li>
|
||||
))}
|
||||
</ul>
|
||||
),
|
||||
}));
|
||||
|
||||
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<string, ModelActivityData> = {
|
||||
"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(<KeyActivityPanel keyMetrics={keyMetrics} />);
|
||||
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(<KeyActivityPanel keyMetrics={keyMetrics} />);
|
||||
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(<KeyActivityPanel keyMetrics={keyMetrics} />);
|
||||
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(<KeyActivityPanel keyMetrics={keyMetrics} />);
|
||||
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");
|
||||
});
|
||||
});
|
||||
|
|
@ -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<string, ModelActivityData>;
|
||||
hidePromptCachingMetrics?: boolean;
|
||||
}
|
||||
|
||||
const KeyActivityPanel: React.FC<KeyActivityPanelProps> = ({ 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 (
|
||||
<div className="space-y-4">
|
||||
<div className="mt-2 flex items-center gap-3">
|
||||
<InputGroup className="max-w-md">
|
||||
<InputGroupAddon>
|
||||
<Search className="size-4 text-muted-foreground" />
|
||||
</InputGroupAddon>
|
||||
<InputGroupInput
|
||||
aria-label="Search keys"
|
||||
placeholder="Search by key alias, key hash, user ID, or email"
|
||||
value={query}
|
||||
onChange={(e) => setQuery(e.target.value)}
|
||||
/>
|
||||
{isFiltering && (
|
||||
<InputGroupAddon align="inline-end">
|
||||
<InputGroupButton size="icon-xs" aria-label="Clear key search" onClick={() => setQuery("")}>
|
||||
<X />
|
||||
</InputGroupButton>
|
||||
</InputGroupAddon>
|
||||
)}
|
||||
</InputGroup>
|
||||
<span className="text-sm text-muted-foreground">
|
||||
Showing {shownKeys.toLocaleString()} of {totalKeys.toLocaleString()} keys
|
||||
</span>
|
||||
</div>
|
||||
{isFiltering && totalKeys > 0 && shownKeys === 0 ? (
|
||||
<p className="rounded-lg border p-6 text-center text-sm text-muted-foreground">
|
||||
No keys match "{query.trim()}" in this date range
|
||||
</p>
|
||||
) : (
|
||||
<ActivityMetrics modelMetrics={filtered} hidePromptCachingMetrics={hidePromptCachingMetrics} />
|
||||
)}
|
||||
</div>
|
||||
);
|
||||
};
|
||||
|
||||
export default KeyActivityPanel;
|
||||
|
|
@ -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<string, ModelActivityData> = {
|
||||
"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({});
|
||||
});
|
||||
});
|
||||
|
|
@ -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<string, ModelActivityData>,
|
||||
query: string,
|
||||
): Record<string, ModelActivityData> {
|
||||
if (query.trim() === "") return keyMetrics;
|
||||
return Object.fromEntries(
|
||||
Object.entries(keyMetrics).filter(([apiKey, data]) => keyActivityMatches(apiKey, data, query)),
|
||||
);
|
||||
}
|
||||
|
|
@ -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;
|
||||
|
|
|
|||
|
|
@ -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: [
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
163
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
163
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
|
|
@ -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;
|
||||
|
|
@ -19984,6 +20024,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;
|
||||
|
|
@ -23140,6 +23200,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
|
||||
|
|
@ -28376,6 +28462,8 @@ export interface components {
|
|||
team_id?: string | null;
|
||||
/** User Email */
|
||||
user_email?: string | null;
|
||||
/** User Id */
|
||||
user_id?: string | null;
|
||||
};
|
||||
/**
|
||||
* KeyMetricWithMetadata
|
||||
|
|
@ -40098,6 +40186,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;
|
||||
|
|
@ -40426,6 +40534,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;
|
||||
|
|
@ -50324,7 +50452,7 @@ export interface operations {
|
|||
organization_id?: string | null;
|
||||
/** @description Filter keys by key hash */
|
||||
key_hash?: string | null;
|
||||
/** @description Filter keys by key alias. Exact match by default; set substring_matching=true (admin only) for case-insensitive substring matching. */
|
||||
/** @description Filter keys by key alias. Exact match by default; set substring_matching=true for case-insensitive substring matching. */
|
||||
key_alias?: string | null;
|
||||
/** @description Combined search: matches keys whose token (key hash) equals the value OR whose key_alias contains it (case-insensitive). */
|
||||
search?: string | null;
|
||||
|
|
@ -50348,7 +50476,7 @@ export interface operations {
|
|||
access_group_id?: string | null;
|
||||
/** @description Filter keys by agent ID */
|
||||
agent_id?: string | null;
|
||||
/** @description If true (proxy admins only), match user_id/key_alias as case-insensitive substrings instead of exact values. Defaults to false: /key/list matched these exactly before substring search was added, and an exact user_id/key_alias filter must never return another user's keys. */
|
||||
/** @description If true, match key_alias (any caller) and user_id (proxy admins only) as case-insensitive substrings instead of exact values. Defaults to false: /key/list matched these exactly before substring search was added, and an exact user_id filter must never return another user's keys. */
|
||||
substring_matching?: boolean;
|
||||
/** @description Filter keys by expiration. 'expired' returns keys whose expires is in the past; 'active' returns keys that never expire or expire in the future. Omit to return keys regardless of expiration. */
|
||||
expires?: string | null;
|
||||
|
|
@ -65156,6 +65284,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;
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue