mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-12 08:42:18 +00:00
Merge main into litellm_openai_sdk3_docker
Keeps both sides in the OpenAI common utils, handler, and tests. Main's upstream error body and credential stripping only accepted an httpx response, while SDK 3's errors carry an httpx2 one, so the body was dropped and a query secret stayed on e.response. Both now accept httpx2 and copy it into an httpx response, with tests that fail without the change
This commit is contained in:
commit
943a96019e
127 changed files with 9224 additions and 1237 deletions
|
|
@ -872,6 +872,11 @@ def _populate_provider_model_sets(model_cost_map: Mapping[str, Mapping[str, obje
|
|||
# ignore the 'up-to', '-to-' model names -> not real models. just for cost tracking based on model params.
|
||||
if "-to-" not in key and "fireworks-ai-default" not in key:
|
||||
fireworks_ai_models.add(key)
|
||||
elif value.get("litellm_provider") == "together_ai":
|
||||
if _is_past_deprecation_date(value):
|
||||
together_ai_models.discard(key)
|
||||
elif not key.startswith("together-ai-"):
|
||||
together_ai_models.add(key)
|
||||
elif value.get("litellm_provider") == "baseten":
|
||||
if _is_past_deprecation_date(value):
|
||||
baseten_models.discard(key)
|
||||
|
|
|
|||
|
|
@ -7,6 +7,7 @@ from typing import Final, Literal
|
|||
|
||||
from litellm.litellm_core_utils.env_utils import (
|
||||
get_env_float,
|
||||
get_env_float_in_range,
|
||||
get_env_float_or_none,
|
||||
get_env_int,
|
||||
get_env_int_in_range,
|
||||
|
|
@ -149,6 +150,9 @@ DEFAULT_FAILURE_THRESHOLD_PERCENT: Final = get_env_float(
|
|||
"DEFAULT_FAILURE_THRESHOLD_PERCENT", 0.5
|
||||
) # default cooldown a deployment if 50% of requests fail in a given minute
|
||||
DEFAULT_MAX_TOKENS: Final = get_env_int("DEFAULT_MAX_TOKENS", 4096)
|
||||
NO_MAX_TOKENS_TPM_RESERVATION_FRACTION: Final = get_env_float_in_range(
|
||||
"NO_MAX_TOKENS_TPM_RESERVATION_FRACTION", 0.25, 0.01, 1.0
|
||||
)
|
||||
DEFAULT_ALLOWED_FAILS: Final = get_env_int("DEFAULT_ALLOWED_FAILS", 3)
|
||||
DEFAULT_REDIS_SYNC_INTERVAL: Final = get_env_int("DEFAULT_REDIS_SYNC_INTERVAL", 1)
|
||||
DEFAULT_COOLDOWN_TIME_SECONDS: Final = get_env_int("DEFAULT_COOLDOWN_TIME_SECONDS", 5)
|
||||
|
|
@ -416,6 +420,8 @@ REALTIME_CREDENTIAL_RESOLUTION_TIMEOUT_SECONDS: Final = get_env_float(
|
|||
"REALTIME_CREDENTIAL_RESOLUTION_TIMEOUT_SECONDS", 20.0
|
||||
)
|
||||
BEDROCK_REALTIME_CLOSE_TIMEOUT_SECONDS: Final = get_env_float("BEDROCK_REALTIME_CLOSE_TIMEOUT_SECONDS", 5.0)
|
||||
REALTIME_LAZY_SETUP_DEADLINE_SECONDS: Final = 240.0
|
||||
REALTIME_LAZY_SETUP_CONTENT_GRACE_SECONDS: Final = 1.0
|
||||
|
||||
# RFC 6455 caps the close frame payload at 125 bytes, 2 of which carry the status code
|
||||
WEBSOCKET_CLOSE_REASON_MAX_BYTES: Final = 123
|
||||
|
|
@ -1270,7 +1276,7 @@ empower_models: Final = set(
|
|||
]
|
||||
)
|
||||
|
||||
together_ai_models: Final[set] = set(
|
||||
together_ai_models: Final[set[str]] = set(
|
||||
[
|
||||
# llama llms - chat
|
||||
"togethercomputer/llama-2-70b-chat",
|
||||
|
|
|
|||
|
|
@ -1544,6 +1544,12 @@ def _response_time_ms_for_cost(
|
|||
return litellm_logging_obj.get_response_ms()
|
||||
|
||||
|
||||
def _breakdown_total(litellm_logging_obj: LitellmLoggingObject | None) -> float | None:
|
||||
breakdown: Final = None if litellm_logging_obj is None else litellm_logging_obj.cost_breakdown
|
||||
total: Final[object] = None if breakdown is None else breakdown.get("total_cost")
|
||||
return float(total) if isinstance(total, (int, float)) and not isinstance(total, bool) else None
|
||||
|
||||
|
||||
def _adjusted_provider_billed_cost(
|
||||
reported_cost: float,
|
||||
custom_llm_provider: str | None,
|
||||
|
|
@ -1561,9 +1567,9 @@ def _adjusted_provider_billed_cost(
|
|||
)
|
||||
_store_cost_breakdown_in_logging_obj(
|
||||
litellm_logging_obj=litellm_logging_obj,
|
||||
prompt_tokens_cost_usd_dollar=None,
|
||||
completion_tokens_cost_usd_dollar=None,
|
||||
cost_for_built_in_tools_cost_usd_dollar=None,
|
||||
prompt_tokens_cost_usd_dollar=0.0,
|
||||
completion_tokens_cost_usd_dollar=adjusted_cost,
|
||||
cost_for_built_in_tools_cost_usd_dollar=0.0,
|
||||
total_cost_usd_dollar=adjusted_cost,
|
||||
original_cost=reported_cost,
|
||||
discount_percent=discount_percent,
|
||||
|
|
@ -2376,6 +2382,15 @@ def response_cost_calculator(
|
|||
)
|
||||
provider_response_cost: Final = get_response_cost_from_hidden_params(hidden_params)
|
||||
if provider_response_cost is not None and not custom_pricing:
|
||||
if _breakdown_total(litellm_logging_obj) != provider_response_cost:
|
||||
_store_cost_breakdown_in_logging_obj(
|
||||
litellm_logging_obj=litellm_logging_obj,
|
||||
prompt_tokens_cost_usd_dollar=0.0,
|
||||
completion_tokens_cost_usd_dollar=provider_response_cost,
|
||||
cost_for_built_in_tools_cost_usd_dollar=0.0,
|
||||
total_cost_usd_dollar=provider_response_cost,
|
||||
original_cost=provider_response_cost,
|
||||
)
|
||||
return provider_response_cost
|
||||
|
||||
response_cost = completion_cost(
|
||||
|
|
|
|||
|
|
@ -9,9 +9,11 @@
|
|||
|
||||
## LiteLLM versions of the OpenAI Exception Types
|
||||
|
||||
import datetime
|
||||
import enum
|
||||
import sys
|
||||
from collections.abc import Mapping, Sequence
|
||||
from typing import Any, Final
|
||||
from typing import TYPE_CHECKING, Any, Final
|
||||
|
||||
import httpx
|
||||
import openai
|
||||
|
|
@ -20,6 +22,9 @@ import litellm
|
|||
from litellm.types.utils import LiteLLMCommonStrings
|
||||
from litellm.types.vector_stores import VectorStoreSearchFailure
|
||||
|
||||
if TYPE_CHECKING:
|
||||
import httpx2
|
||||
|
||||
|
||||
class RateLimitErrorCategory(str, enum.Enum):
|
||||
"""
|
||||
|
|
@ -115,6 +120,77 @@ def validate_rate_limit_type(value: object) -> str | None:
|
|||
return None
|
||||
|
||||
|
||||
def upstream_http_response(upstream: object) -> "httpx.Response | httpx2.Response | None":
|
||||
if isinstance(upstream, httpx.Response):
|
||||
return upstream
|
||||
if "httpx2" not in sys.modules:
|
||||
return None
|
||||
import httpx2
|
||||
|
||||
return upstream if isinstance(upstream, httpx2.Response) else None
|
||||
|
||||
|
||||
def _upstream_body(response: "httpx.Response | httpx2.Response | None") -> bytes:
|
||||
if response is None:
|
||||
return b""
|
||||
try:
|
||||
return response.content
|
||||
except RuntimeError:
|
||||
return b""
|
||||
|
||||
|
||||
def _request_without_credentials(
|
||||
upstream: "httpx.Response | httpx2.Response | None", fallback_url: str
|
||||
) -> httpx.Request:
|
||||
if upstream is None:
|
||||
return httpx.Request(method="POST", url=fallback_url)
|
||||
try:
|
||||
upstream_request: Final = upstream.request
|
||||
except RuntimeError:
|
||||
return httpx.Request(method="POST", url=fallback_url)
|
||||
return httpx.Request(
|
||||
method=upstream_request.method,
|
||||
url=httpx.URL(str(upstream_request.url)).copy_with(userinfo=b"", query=None, fragment=None),
|
||||
)
|
||||
|
||||
|
||||
def _upstream_elapsed(upstream: "httpx.Response | httpx2.Response | None") -> datetime.timedelta | None:
|
||||
if upstream is None:
|
||||
return None
|
||||
try:
|
||||
return upstream.elapsed
|
||||
except RuntimeError:
|
||||
return None
|
||||
|
||||
|
||||
def _upstream_headers(upstream: "httpx.Response | httpx2.Response | None") -> httpx.Headers:
|
||||
match upstream:
|
||||
case None:
|
||||
return httpx.Headers()
|
||||
case httpx.Response(headers=headers):
|
||||
return httpx.Headers(headers)
|
||||
case _:
|
||||
return httpx.Headers(upstream.headers.multi_items())
|
||||
|
||||
|
||||
def error_response_without_credentials(status_code: int, upstream: object, fallback_url: str) -> httpx.Response:
|
||||
upstream_response: Final = upstream_http_response(upstream)
|
||||
response: Final = httpx.Response(
|
||||
status_code=status_code,
|
||||
stream=httpx.ByteStream(_upstream_body(upstream_response)),
|
||||
request=_request_without_credentials(upstream_response, fallback_url),
|
||||
)
|
||||
_ = response.read()
|
||||
response.headers = _upstream_headers(upstream_response)
|
||||
elapsed: Final = _upstream_elapsed(upstream_response)
|
||||
if elapsed is not None:
|
||||
response.elapsed = elapsed
|
||||
return response
|
||||
|
||||
|
||||
_NO_UPSTREAM_REQUEST_URL: Final = "https://cloud.google.com/vertex-ai/"
|
||||
|
||||
|
||||
_MINIMAL_ERROR_RESPONSE: httpx.Response | None = None
|
||||
|
||||
|
||||
|
|
@ -519,20 +595,11 @@ class RateLimitError(openai.RateLimitError):
|
|||
# headers stay reachable on `e.response.headers` for callers that
|
||||
# explicitly want them; only the proxy-supplied `headers=` kwarg
|
||||
# makes it onto `self.headers`.
|
||||
_response_headers: Final = getattr(response, "headers", None) if response is not None else None
|
||||
self.headers: dict[str, str] | None = {k: str(v) for k, v in headers.items()} if headers else None
|
||||
# Mirrors FastAPI HTTPException.detail so the same instance can be
|
||||
# serialized through both the ProxyException and HTTPException paths.
|
||||
self.detail = detail if detail is not None else self.message
|
||||
self.response = httpx.Response(
|
||||
status_code=429,
|
||||
headers=_response_headers,
|
||||
content=response.content if response is not None else None,
|
||||
request=httpx.Request(
|
||||
method="POST",
|
||||
url=" https://cloud.google.com/vertex-ai/",
|
||||
),
|
||||
)
|
||||
self.response = error_response_without_credentials(429, response, _NO_UPSTREAM_REQUEST_URL)
|
||||
super().__init__(
|
||||
self.message,
|
||||
response=self.response, # pyright: ignore[reportArgumentType] # SDK accepts HTTPX at runtime
|
||||
|
|
@ -757,15 +824,7 @@ class ServiceUnavailableError(openai.APIStatusError):
|
|||
self.litellm_debug_info = litellm_debug_info
|
||||
self.max_retries = max_retries
|
||||
self.num_retries = num_retries
|
||||
_response_headers: Final = getattr(response, "headers", None) if response is not None else None
|
||||
self.response = httpx.Response(
|
||||
status_code=self.status_code,
|
||||
headers=_response_headers,
|
||||
request=httpx.Request(
|
||||
method="POST",
|
||||
url=" https://cloud.google.com/vertex-ai/",
|
||||
),
|
||||
)
|
||||
self.response = error_response_without_credentials(self.status_code, response, _NO_UPSTREAM_REQUEST_URL)
|
||||
super().__init__(
|
||||
self.message, response=self.response, body=None
|
||||
) # Call the base class constructor with the parameters it needs
|
||||
|
|
@ -808,15 +867,7 @@ class BadGatewayError(openai.APIStatusError):
|
|||
self.litellm_debug_info = litellm_debug_info
|
||||
self.max_retries = max_retries
|
||||
self.num_retries = num_retries
|
||||
_response_headers: Final = getattr(response, "headers", None) if response is not None else None
|
||||
self.response = httpx.Response(
|
||||
status_code=self.status_code,
|
||||
headers=_response_headers,
|
||||
request=httpx.Request(
|
||||
method="POST",
|
||||
url=" https://cloud.google.com/vertex-ai/",
|
||||
),
|
||||
)
|
||||
self.response = error_response_without_credentials(self.status_code, response, _NO_UPSTREAM_REQUEST_URL)
|
||||
super().__init__(
|
||||
self.message,
|
||||
response=self.response, # pyright: ignore[reportArgumentType] # SDK accepts HTTPX at runtime
|
||||
|
|
@ -862,15 +913,7 @@ class InternalServerError(openai.InternalServerError):
|
|||
self.litellm_debug_info = litellm_debug_info
|
||||
self.max_retries = max_retries
|
||||
self.num_retries = num_retries
|
||||
_response_headers: Final = getattr(response, "headers", None) if response is not None else None
|
||||
self.response = httpx.Response(
|
||||
status_code=self.status_code,
|
||||
headers=_response_headers,
|
||||
request=httpx.Request(
|
||||
method="POST",
|
||||
url=" https://cloud.google.com/vertex-ai/",
|
||||
),
|
||||
)
|
||||
self.response = error_response_without_credentials(self.status_code, response, _NO_UPSTREAM_REQUEST_URL)
|
||||
super().__init__(
|
||||
self.message,
|
||||
response=self.response, # pyright: ignore[reportArgumentType] # SDK accepts HTTPX at runtime
|
||||
|
|
@ -908,6 +951,7 @@ class APIError(openai.APIError):
|
|||
max_retries: int | None = None,
|
||||
num_retries: int | None = None,
|
||||
body: object | None = None,
|
||||
response: httpx.Response | None = None,
|
||||
):
|
||||
self.status_code = status_code
|
||||
self.message = f"litellm.APIError: {message}"
|
||||
|
|
@ -916,6 +960,8 @@ class APIError(openai.APIError):
|
|||
self.litellm_debug_info = litellm_debug_info
|
||||
self.max_retries = max_retries
|
||||
self.num_retries = num_retries
|
||||
if response is not None:
|
||||
self.response = response
|
||||
if request is None:
|
||||
request = httpx.Request(method="POST", url="https://api.openai.com/v1")
|
||||
super().__init__(self.message, request=request, body=body)
|
||||
|
|
|
|||
|
|
@ -23,6 +23,7 @@ from litellm._logging import verbose_logger
|
|||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.integrations.custom_prompt_management import CustomPromptManagement
|
||||
from litellm.integrations.prompt_management_base import PromptManagementClient
|
||||
from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap
|
||||
from litellm.litellm_core_utils.prompt_templates.common_utils import (
|
||||
cache_control_as_prompt_cache_breakpoint,
|
||||
with_prompt_cache_breakpoint,
|
||||
|
|
@ -175,7 +176,10 @@ def _hosted_entry_flag(entry: Mapping[str, object], resolve_provider: Callable[[
|
|||
|
||||
|
||||
def _hosted_openai_dialect_flag(
|
||||
model: str, custom_llm_provider: str | None, resolve_provider: Callable[[str], str | None]
|
||||
model: str,
|
||||
custom_llm_provider: str | None,
|
||||
resolve_provider: Callable[[str], str | None],
|
||||
model_cost: Mapping[str, Mapping[str, object]] | None = None,
|
||||
) -> bool | None:
|
||||
"""
|
||||
Explicit opt-in for an OpenAI-shaped deployment served by another provider.
|
||||
|
|
@ -194,14 +198,19 @@ def _hosted_openai_dialect_flag(
|
|||
"""
|
||||
import litellm
|
||||
|
||||
exact_entry: Final = litellm.model_cost.get(model)
|
||||
exact_provider: Final = exact_entry.get("litellm_provider") if isinstance(exact_entry, dict) else None
|
||||
cost_map: Final = (
|
||||
cast(Mapping[str, Mapping[str, object]], litellm.model_cost) # cast-ok: the live map holds catalog entries
|
||||
if model_cost is None
|
||||
else model_cost
|
||||
)
|
||||
exact_entry: Final = cost_map.get(model)
|
||||
exact_provider: Final = exact_entry.get("litellm_provider") if isinstance(exact_entry, Mapping) else None
|
||||
if custom_llm_provider is None and (exact_provider == "openai" or ("/" not in model and exact_provider is None)):
|
||||
return None
|
||||
provider: Final = custom_llm_provider or resolve_provider(model)
|
||||
if provider is None or provider == "openai":
|
||||
return None
|
||||
if isinstance(exact_entry, dict) and exact_provider == provider:
|
||||
if isinstance(exact_entry, Mapping) and exact_provider == provider:
|
||||
return _hosted_entry_flag(exact_entry, lambda: provider)
|
||||
from litellm.utils import get_potential_model_names
|
||||
|
||||
|
|
@ -214,8 +223,8 @@ def _hosted_openai_dialect_flag(
|
|||
names["stripped_model_name"],
|
||||
names["provider_prefixed_model_name"],
|
||||
)
|
||||
entries: Final = (litellm.model_cost.get(candidate) for candidate in candidates)
|
||||
flags: Final = (_hosted_entry_flag(entry, lambda: provider) for entry in entries if isinstance(entry, dict))
|
||||
entries: Final = (cost_map.get(candidate) for candidate in candidates)
|
||||
flags: Final = (_hosted_entry_flag(entry, lambda: provider) for entry in entries if isinstance(entry, Mapping))
|
||||
return next((flag for flag in flags if flag is not None), None)
|
||||
|
||||
|
||||
|
|
@ -537,6 +546,18 @@ class AnthropicCacheControlHook(CustomPromptManagement):
|
|||
return False
|
||||
return prompt_cache_options is not None or targets_openai_api(api_base)
|
||||
|
||||
@staticmethod
|
||||
def catalog_targets_openai_prompt_cache_breakpoint(model: str, custom_llm_provider: str) -> bool:
|
||||
return (
|
||||
_hosted_openai_dialect_flag(
|
||||
model,
|
||||
custom_llm_provider,
|
||||
AnthropicCacheControlHook._resolve_provider,
|
||||
GetModelCostMap.loaded_model_cost_map(),
|
||||
)
|
||||
is True
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _resolve_provider(model: str) -> str | None:
|
||||
return _served_provider_for_model(model)
|
||||
|
|
|
|||
|
|
@ -9,6 +9,7 @@ from types import MappingProxyType
|
|||
from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, Optional, cast, get_args
|
||||
|
||||
import httpx
|
||||
from pydantic import TypeAdapter
|
||||
|
||||
from litellm._internal_context import with_service_target
|
||||
from litellm._logging import verbose_logger
|
||||
|
|
@ -218,6 +219,15 @@ def _request_is_streaming(data: object, event_type: GuardrailEventHooks | None =
|
|||
)
|
||||
|
||||
|
||||
_REQUEST_FIELDS: Final = TypeAdapter(Mapping[str, object])
|
||||
|
||||
|
||||
def _guardrails_listed_in(metadata: object) -> list[str | dict[str, DynamicGuardrailParams]]:
|
||||
if not isinstance(metadata, dict):
|
||||
return []
|
||||
return metadata.get("guardrails") or []
|
||||
|
||||
|
||||
class CustomGuardrail(CustomLogger):
|
||||
# If True, during_call runs async_moderation_hook instead of the unified apply_guardrail path.
|
||||
use_native_during_call_hook: ClassVar[bool] = False
|
||||
|
|
@ -791,28 +801,22 @@ class CustomGuardrail(CustomLogger):
|
|||
return True
|
||||
raise
|
||||
|
||||
def get_guardrail_from_metadata(self, data: dict) -> list[str] | list[dict[str, DynamicGuardrailParams]]:
|
||||
def get_guardrail_from_metadata(self, data: dict) -> list[str | dict[str, DynamicGuardrailParams]]:
|
||||
"""
|
||||
Returns the guardrail(s) to be run from the metadata or root
|
||||
Returns the guardrail(s) to be run from the root, else from both metadata fields together
|
||||
"""
|
||||
|
||||
if "guardrails" in data:
|
||||
return data["guardrails"]
|
||||
# Check both metadata locations. For regular endpoints move_guardrails_to_metadata
|
||||
# writes to "metadata"; for thread/assistant endpoints it writes to
|
||||
# "litellm_metadata". We check the one that actually contains the "guardrails"
|
||||
# key so that a non-empty litellm_metadata without guardrails does not shadow
|
||||
# the merged list stored in metadata (which would cause team guardrails to be
|
||||
# silently skipped while default_on=True policy guardrails still fire).
|
||||
for meta_key in ("metadata", "litellm_metadata"):
|
||||
meta = data.get(meta_key) or {}
|
||||
if isinstance(meta, dict) and "guardrails" in meta:
|
||||
return meta.get("guardrails") or []
|
||||
return []
|
||||
request_fields: Final = _REQUEST_FIELDS.validate_python(data)
|
||||
return [
|
||||
*_guardrails_listed_in(request_fields.get("metadata")),
|
||||
*_guardrails_listed_in(request_fields.get("litellm_metadata")),
|
||||
]
|
||||
|
||||
def _guardrail_is_in_requested_guardrails(
|
||||
self,
|
||||
requested_guardrails: list[str] | list[dict[str, DynamicGuardrailParams]],
|
||||
requested_guardrails: Sequence[str | dict[str, DynamicGuardrailParams]],
|
||||
) -> bool:
|
||||
for _guardrail in requested_guardrails:
|
||||
if isinstance(_guardrail, dict):
|
||||
|
|
|
|||
|
|
@ -398,14 +398,17 @@ class WebSearchInterceptionLogger(CustomLogger):
|
|||
(t for t in tools if is_anthropic_native_web_search_tool(t)),
|
||||
None,
|
||||
)
|
||||
emit_native_blocks: Final = native_tool is not None or (
|
||||
kwargs is not None and bool(kwargs.get(WEBSEARCH_EMIT_NATIVE_BLOCKS_KEY))
|
||||
)
|
||||
|
||||
outcome: Final = await self._short_circuit_search_outcome(query, kwargs=kwargs)
|
||||
search_result_text: Final = WebSearchTransformation.search_outcome_text(outcome)
|
||||
|
||||
content: Final[list[dict[str, object]]] = []
|
||||
if native_tool is not None:
|
||||
if emit_native_blocks:
|
||||
tool_use_id: Final = f"srvtoolu_{uuid.uuid4().hex}"
|
||||
tool_name: Final = native_tool.get("name") or "web_search"
|
||||
tool_name: Final = (native_tool.get("name") if native_tool is not None else None) or "web_search"
|
||||
content.append(
|
||||
{
|
||||
"type": "server_tool_use",
|
||||
|
|
@ -415,10 +418,10 @@ class WebSearchInterceptionLogger(CustomLogger):
|
|||
}
|
||||
)
|
||||
content.append(
|
||||
WebSearchTransformation.build_web_search_outcome_block(tool_use_id=tool_use_id, outcome=outcome)
|
||||
WebSearchTransformation.build_web_search_outcome_block(
|
||||
tool_use_id=tool_use_id, outcome=outcome, with_snippets=False
|
||||
)
|
||||
)
|
||||
# Keep the text block so non-native short-circuit callers (Claude Code,
|
||||
# github_copilot, etc.) see the same payload they always have.
|
||||
content.append({"type": "text", "text": search_result_text})
|
||||
|
||||
response: Final[dict[str, object]] = {
|
||||
|
|
@ -435,7 +438,7 @@ class WebSearchInterceptionLogger(CustomLogger):
|
|||
verbose_logger.debug(
|
||||
"WebSearchInterception: Short-circuit search completed, returning synthetic response (%s chars, native_blocks=%s)",
|
||||
len(search_result_text),
|
||||
native_tool is not None,
|
||||
emit_native_blocks,
|
||||
)
|
||||
return response
|
||||
|
||||
|
|
|
|||
|
|
@ -412,6 +412,7 @@ class WebSearchTransformation:
|
|||
def build_web_search_tool_result_block(
|
||||
tool_use_id: str,
|
||||
search_response: SearchResponse | None,
|
||||
with_snippets: bool = True,
|
||||
) -> dict[str, object]:
|
||||
"""
|
||||
Build an Anthropic-native ``web_search_tool_result`` content block.
|
||||
|
|
@ -445,6 +446,8 @@ class WebSearchTransformation:
|
|||
``litellm.asearch()``. If None or empty, the block is still
|
||||
emitted with an empty result list (signals "search ran, no
|
||||
results" rather than "search did not run").
|
||||
with_snippets: False when a text block in the same reply already
|
||||
carries the snippets.
|
||||
"""
|
||||
items: Final[list[dict[str, object]]] = []
|
||||
if search_response is not None:
|
||||
|
|
@ -460,7 +463,7 @@ class WebSearchTransformation:
|
|||
"title": title,
|
||||
"page_age": page_age,
|
||||
"encrypted_content": "",
|
||||
"snippet": getattr(r, "snippet", "") or "",
|
||||
**({"snippet": getattr(r, "snippet", "") or ""} if with_snippets else {}),
|
||||
}
|
||||
)
|
||||
return {
|
||||
|
|
@ -481,12 +484,15 @@ class WebSearchTransformation:
|
|||
}
|
||||
|
||||
@staticmethod
|
||||
def build_web_search_outcome_block(tool_use_id: str, outcome: SearchOutcome) -> dict[str, object]:
|
||||
def build_web_search_outcome_block(
|
||||
tool_use_id: str, outcome: SearchOutcome, with_snippets: bool = True
|
||||
) -> dict[str, object]:
|
||||
match outcome:
|
||||
case SearchSucceeded(response=response):
|
||||
return WebSearchTransformation.build_web_search_tool_result_block(
|
||||
tool_use_id=tool_use_id,
|
||||
search_response=response,
|
||||
with_snippets=with_snippets,
|
||||
)
|
||||
case SearchFailed(error_code=error_code):
|
||||
return WebSearchTransformation.build_web_search_tool_result_error_block(
|
||||
|
|
|
|||
|
|
@ -1017,13 +1017,44 @@ def chunk_usage(chunk: object) -> object:
|
|||
return getattr(chunk, "usage", None)
|
||||
|
||||
|
||||
USAGE_COST_BILLED_PROVIDERS: Final = frozenset({"kyma"})
|
||||
PROVIDERS_BILLED_FROM_THEIR_OWN_USAGE_FIELDS: Final = frozenset({"xai", "perplexity", "opencode_go"})
|
||||
|
||||
|
||||
UPSTREAM_USAGE_COST: Final = "upstream_usage_cost"
|
||||
|
||||
|
||||
def _numeric(value: object) -> float | None:
|
||||
return float(value) if isinstance(value, (int, float)) and not isinstance(value, bool) else None
|
||||
|
||||
|
||||
def stamp_litellm_bill_on_usage(response: object, usage: object, bill: float) -> None:
|
||||
"""Writes litellm's bill into ``usage.cost`` and the hidden ``response_cost`` cache, keeping the upstream's own
|
||||
``usage.cost`` under ``upstream_usage_cost`` so a later pricing pass still bills the upstream's charge instead of
|
||||
reading litellm's bill back as if the upstream had reported it."""
|
||||
hidden: Final[object] = getattr(response, "_hidden_params", None)
|
||||
if isinstance(hidden, dict):
|
||||
hidden_params: Final = cast(dict[str, object], hidden) # cast-ok: hidden params ride on the response as a dict
|
||||
upstream: Final = (
|
||||
hidden_params[UPSTREAM_USAGE_COST]
|
||||
if UPSTREAM_USAGE_COST in hidden_params
|
||||
else _numeric(getattr(usage, "cost", None))
|
||||
)
|
||||
hidden_params.update({"response_cost": bill, UPSTREAM_USAGE_COST: upstream})
|
||||
setattr(usage, "cost", bill)
|
||||
|
||||
|
||||
def billed_usage_cost(response: object, custom_llm_provider: str | None) -> float | None:
|
||||
if custom_llm_provider not in USAGE_COST_BILLED_PROVIDERS:
|
||||
"""The charge an OpenAI-compatible upstream reports in ``usage.cost``, skipped for providers billed from their own
|
||||
usage fields. Once litellm wrote its own bill into ``usage.cost`` (the stream builder and the Responses terminal
|
||||
event do), the upstream's figure is read from ``upstream_usage_cost`` instead."""
|
||||
if custom_llm_provider in PROVIDERS_BILLED_FROM_THEIR_OWN_USAGE_FIELDS:
|
||||
return None
|
||||
cost: Final = provider_reported_usage_cost(_usage_field(response, "usage"), custom_llm_provider)
|
||||
hidden: Final = as_str_mapping(getattr(response, "_hidden_params", None))
|
||||
cost: Final = (
|
||||
_numeric(hidden[UPSTREAM_USAGE_COST])
|
||||
if hidden is not None and UPSTREAM_USAGE_COST in hidden
|
||||
else provider_reported_usage_cost(_usage_field(response, "usage"), custom_llm_provider)
|
||||
)
|
||||
return cost if cost is not None and math.isfinite(cost) and cost >= 0 else None
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -71,7 +71,15 @@ def get_env_int_in_range(env_var: str, default: int, minimum: int, maximum: int)
|
|||
Values outside the range fall back to the default and warn, so a misconfigured knob can
|
||||
neither crash the caller nor silently change the meaning of what it computes.
|
||||
"""
|
||||
value: Final = get_env_int(env_var, default)
|
||||
return _within_range(env_var, get_env_int(env_var, default), default, minimum, maximum)
|
||||
|
||||
|
||||
def get_env_float_in_range(env_var: str, default: float, minimum: float, maximum: float) -> float:
|
||||
"""Float counterpart of `get_env_int_in_range`."""
|
||||
return _within_range(env_var, get_env_float(env_var, default), default, minimum, maximum)
|
||||
|
||||
|
||||
def _within_range(env_var: str, value: _NumberT, default: _NumberT, minimum: _NumberT, maximum: _NumberT) -> _NumberT:
|
||||
if minimum <= value <= maximum:
|
||||
return value
|
||||
logging.getLogger("LiteLLM").warning(
|
||||
|
|
|
|||
|
|
@ -38,6 +38,8 @@ from ..exceptions import (
|
|||
ServiceUnavailableError,
|
||||
Timeout,
|
||||
UnprocessableEntityError,
|
||||
error_response_without_credentials,
|
||||
upstream_http_response,
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -310,6 +312,17 @@ def _litellm_proxy_response(
|
|||
)
|
||||
|
||||
|
||||
def _without_credentials(response: httpx.Response | None) -> httpx.Response | None:
|
||||
upstream: Final = upstream_http_response(response)
|
||||
if upstream is None:
|
||||
return response
|
||||
return error_response_without_credentials(upstream.status_code, upstream, "https://docs.litellm.ai/docs")
|
||||
|
||||
|
||||
def _upstream_reply(response: httpx.Response | None) -> httpx.Response | None:
|
||||
return response if response is not None and (response.headers or response.content) else None
|
||||
|
||||
|
||||
def _map_openai_exception(
|
||||
*,
|
||||
model: str,
|
||||
|
|
@ -320,7 +333,7 @@ def _map_openai_exception(
|
|||
exception_provider: str,
|
||||
extra_information: str,
|
||||
) -> None:
|
||||
response: Final = _litellm_proxy_response(original_exception, custom_llm_provider)
|
||||
response: Final = _without_credentials(_litellm_proxy_response(original_exception, custom_llm_provider))
|
||||
if (cyber_policy_body := _openai_cyber_policy_body(original_exception)) is not None:
|
||||
raise ContentPolicyViolationError(
|
||||
message=f"ContentPolicyViolationError: {exception_provider} - {original_exception}",
|
||||
|
|
@ -383,6 +396,7 @@ def _map_openai_exception(
|
|||
model=model,
|
||||
llm_provider=custom_llm_provider,
|
||||
litellm_debug_info=extra_information,
|
||||
response=_upstream_reply(response),
|
||||
)
|
||||
elif (
|
||||
("invalid_request_error" in error_str and "content_policy_violation" in error_str)
|
||||
|
|
@ -454,6 +468,7 @@ def _map_openai_exception(
|
|||
model=model,
|
||||
llm_provider=custom_llm_provider,
|
||||
body=getattr(original_exception, "body", None),
|
||||
response=_upstream_reply(response),
|
||||
)
|
||||
elif "Request too large" in error_str:
|
||||
raise RateLimitError(
|
||||
|
|
@ -481,6 +496,7 @@ def _map_openai_exception(
|
|||
model=model,
|
||||
request=_request,
|
||||
litellm_debug_info=extra_information,
|
||||
response=_upstream_reply(response),
|
||||
)
|
||||
elif hasattr(original_exception, "status_code"):
|
||||
if original_exception.status_code == 400:
|
||||
|
|
@ -514,6 +530,7 @@ def _map_openai_exception(
|
|||
model=model,
|
||||
llm_provider=custom_llm_provider,
|
||||
litellm_debug_info=extra_information,
|
||||
response=_upstream_reply(response),
|
||||
)
|
||||
elif original_exception.status_code == 422:
|
||||
raise BadRequestError(
|
||||
|
|
@ -565,6 +582,7 @@ def _map_openai_exception(
|
|||
llm_provider=custom_llm_provider,
|
||||
litellm_debug_info=extra_information,
|
||||
exception_status_code=original_exception.status_code,
|
||||
response=_upstream_reply(response),
|
||||
)
|
||||
else:
|
||||
raise APIError(
|
||||
|
|
@ -574,6 +592,7 @@ def _map_openai_exception(
|
|||
model=model,
|
||||
request=getattr(original_exception, "request", None),
|
||||
litellm_debug_info=extra_information,
|
||||
response=_upstream_reply(response),
|
||||
)
|
||||
else:
|
||||
# if no status code then it is an APIConnectionError: https://github.com/openai/openai-python#handling-errors
|
||||
|
|
@ -2021,6 +2040,7 @@ def _map_azure_exception(
|
|||
exception_provider: str,
|
||||
extra_information: str,
|
||||
) -> None:
|
||||
response: Final = _without_credentials(getattr(original_exception, "response", None))
|
||||
message = get_error_message(error_obj=original_exception)
|
||||
if message is None:
|
||||
if hasattr(original_exception, "message"):
|
||||
|
|
@ -2056,7 +2076,7 @@ def _map_azure_exception(
|
|||
llm_provider="azure",
|
||||
model=model,
|
||||
litellm_debug_info=extra_information,
|
||||
response=getattr(original_exception, "response", None),
|
||||
response=response,
|
||||
)
|
||||
elif "DeploymentNotFound" in error_str:
|
||||
raise NotFoundError(
|
||||
|
|
@ -2064,7 +2084,7 @@ def _map_azure_exception(
|
|||
llm_provider="azure",
|
||||
model=model,
|
||||
litellm_debug_info=extra_information,
|
||||
response=getattr(original_exception, "response", None),
|
||||
response=response,
|
||||
)
|
||||
elif azure_error_code == "content_policy_violation" or ExceptionCheckers.is_azure_content_policy_violation_error(
|
||||
error_str
|
||||
|
|
@ -2078,6 +2098,7 @@ def _map_azure_exception(
|
|||
model=model,
|
||||
extra_information=extra_information,
|
||||
original_exception=original_exception,
|
||||
response=response,
|
||||
)
|
||||
elif azure_error_code == "invalid_encrypted_content" or "could not be verified" in error_str:
|
||||
helpful_message: Final = (
|
||||
|
|
@ -2096,7 +2117,7 @@ def _map_azure_exception(
|
|||
llm_provider="azure",
|
||||
model=model,
|
||||
litellm_debug_info=extra_information,
|
||||
response=getattr(original_exception, "response", None),
|
||||
response=response,
|
||||
body=getattr(original_exception, "body", None),
|
||||
)
|
||||
elif getattr(original_exception, "status_code", None) in (None, 400) and (
|
||||
|
|
@ -2108,7 +2129,7 @@ def _map_azure_exception(
|
|||
llm_provider="azure",
|
||||
model=model,
|
||||
litellm_debug_info=extra_information,
|
||||
response=getattr(original_exception, "response", None),
|
||||
response=response,
|
||||
body=getattr(original_exception, "body", None),
|
||||
)
|
||||
elif "invalid_request_error" in error_str and getattr(original_exception, "status_code", None) in (None, 400):
|
||||
|
|
@ -2117,7 +2138,7 @@ def _map_azure_exception(
|
|||
llm_provider="azure",
|
||||
model=model,
|
||||
litellm_debug_info=extra_information,
|
||||
response=getattr(original_exception, "response", None),
|
||||
response=response,
|
||||
body=getattr(original_exception, "body", None),
|
||||
)
|
||||
elif any(marker in error_str for marker in _OPENAI_CLIENT_MISSING_CREDENTIALS_MESSAGES):
|
||||
|
|
@ -2126,7 +2147,7 @@ def _map_azure_exception(
|
|||
llm_provider=custom_llm_provider,
|
||||
model=model,
|
||||
litellm_debug_info=extra_information,
|
||||
response=getattr(original_exception, "response", None),
|
||||
response=response,
|
||||
)
|
||||
elif "Connection error" in error_str:
|
||||
raise APIConnectionError(
|
||||
|
|
@ -2142,7 +2163,7 @@ def _map_azure_exception(
|
|||
llm_provider="azure",
|
||||
model=model,
|
||||
litellm_debug_info=extra_information,
|
||||
response=getattr(original_exception, "response", None),
|
||||
response=response,
|
||||
body=getattr(original_exception, "body", None),
|
||||
)
|
||||
elif original_exception.status_code == 401:
|
||||
|
|
@ -2151,7 +2172,7 @@ def _map_azure_exception(
|
|||
llm_provider="azure",
|
||||
model=model,
|
||||
litellm_debug_info=extra_information,
|
||||
response=getattr(original_exception, "response", None),
|
||||
response=response,
|
||||
)
|
||||
elif original_exception.status_code == 404:
|
||||
raise NotFoundError(
|
||||
|
|
@ -2159,7 +2180,7 @@ def _map_azure_exception(
|
|||
llm_provider="azure",
|
||||
model=model,
|
||||
litellm_debug_info=extra_information,
|
||||
response=getattr(original_exception, "response", None),
|
||||
response=response,
|
||||
)
|
||||
elif original_exception.status_code == 408:
|
||||
raise Timeout(
|
||||
|
|
@ -2167,6 +2188,7 @@ def _map_azure_exception(
|
|||
model=model,
|
||||
litellm_debug_info=extra_information,
|
||||
llm_provider="azure",
|
||||
response=_upstream_reply(response),
|
||||
)
|
||||
elif original_exception.status_code == 422:
|
||||
raise BadRequestError(
|
||||
|
|
@ -2174,7 +2196,7 @@ def _map_azure_exception(
|
|||
model=model,
|
||||
llm_provider="azure",
|
||||
litellm_debug_info=extra_information,
|
||||
response=getattr(original_exception, "response", None),
|
||||
response=response,
|
||||
)
|
||||
elif original_exception.status_code == 429:
|
||||
raise RateLimitError(
|
||||
|
|
@ -2182,7 +2204,7 @@ def _map_azure_exception(
|
|||
model=model,
|
||||
llm_provider="azure",
|
||||
litellm_debug_info=extra_information,
|
||||
response=getattr(original_exception, "response", None),
|
||||
response=response,
|
||||
)
|
||||
elif original_exception.status_code == 502:
|
||||
raise BadGatewayError(
|
||||
|
|
@ -2190,7 +2212,7 @@ def _map_azure_exception(
|
|||
model=model,
|
||||
llm_provider="azure",
|
||||
litellm_debug_info=extra_information,
|
||||
response=getattr(original_exception, "response", None),
|
||||
response=response,
|
||||
)
|
||||
elif original_exception.status_code == 503:
|
||||
raise ServiceUnavailableError(
|
||||
|
|
@ -2198,7 +2220,7 @@ def _map_azure_exception(
|
|||
model=model,
|
||||
llm_provider="azure",
|
||||
litellm_debug_info=extra_information,
|
||||
response=getattr(original_exception, "response", None),
|
||||
response=response,
|
||||
)
|
||||
elif original_exception.status_code == 504: # gateway timeout error
|
||||
raise Timeout(
|
||||
|
|
@ -2207,6 +2229,7 @@ def _map_azure_exception(
|
|||
litellm_debug_info=extra_information,
|
||||
llm_provider="azure",
|
||||
exception_status_code=original_exception.status_code,
|
||||
response=_upstream_reply(response),
|
||||
)
|
||||
else:
|
||||
raise APIError(
|
||||
|
|
@ -2216,6 +2239,7 @@ def _map_azure_exception(
|
|||
litellm_debug_info=extra_information,
|
||||
model=model,
|
||||
request=httpx.Request(method="POST", url="https://openai.com/"),
|
||||
response=_upstream_reply(response),
|
||||
)
|
||||
else:
|
||||
# if no status code then it is an APIConnectionError: https://github.com/openai/openai-python#handling-errors
|
||||
|
|
@ -2591,17 +2615,20 @@ def exception_type(
|
|||
model=model,
|
||||
llm_provider=custom_llm_provider,
|
||||
litellm_debug_info=extra_information,
|
||||
response=_upstream_reply(
|
||||
_without_credentials(_litellm_proxy_response(mappable_exception, custom_llm_provider))
|
||||
),
|
||||
)
|
||||
|
||||
if (
|
||||
custom_llm_provider == "litellm_proxy"
|
||||
): # handle special case where calling litellm proxy + exception str contains error message
|
||||
extract_and_raise_litellm_exception(
|
||||
response=_litellm_proxy_response(mappable_exception, custom_llm_provider),
|
||||
response=_without_credentials(_litellm_proxy_response(mappable_exception, custom_llm_provider)),
|
||||
error_str=error_str,
|
||||
model=model,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
body=getattr(original_exception, "body", None),
|
||||
body=getattr(mappable_exception, "body", None),
|
||||
)
|
||||
if (
|
||||
custom_llm_provider == "openai"
|
||||
|
|
|
|||
|
|
@ -85,6 +85,7 @@ OPTIONAL_KWARGS_KEYS: Final = (
|
|||
"github_copilot_user_session",
|
||||
"fireworks_forward_user_id",
|
||||
"responses_api_version",
|
||||
"allowed_openai_params",
|
||||
PROVIDER_AFFINITY_HEADER_KWARG_KEY,
|
||||
}
|
||||
)
|
||||
|
|
|
|||
|
|
@ -3,6 +3,8 @@ Helper functions for health check calls.
|
|||
"""
|
||||
|
||||
import base64
|
||||
import itertools
|
||||
import re
|
||||
from collections.abc import Awaitable, Callable, Mapping
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Final, Literal
|
||||
|
|
@ -26,6 +28,9 @@ IMAGE_EDIT_HEALTH_CHECK_PROMPT: Final = (
|
|||
)
|
||||
|
||||
ANTHROPIC_MESSAGES_HEALTH_CHECK_MAX_TOKENS: Final = 16
|
||||
WILDCARD_HEALTH_CHECK_SERVED_MODEL_COUNT: Final = 2
|
||||
WILDCARD_HEALTH_CHECK_PROVIDER_MODEL_COUNT: Final = 3
|
||||
BEDROCK_ANY_MODEL_ROUTES: Final = frozenset({"converse", "invoke"})
|
||||
EXTRACTION_HEALTH_CHECK_FIELDS: Final[Mapping[str, Mapping[str, str]]] = MappingProxyType(
|
||||
{"reachable": MappingProxyType({"type": "text", "instructions": "Any word from the text"})}
|
||||
)
|
||||
|
|
@ -90,6 +95,62 @@ def _ocr_health_check_document(model: str, custom_llm_provider: str) -> Document
|
|||
return native(model, custom_llm_provider)
|
||||
|
||||
|
||||
def _provider_base_model(custom_llm_provider: str, model: str) -> str | None:
|
||||
from litellm.types.utils import LlmProviders
|
||||
from litellm.utils import ProviderConfigManager
|
||||
|
||||
try:
|
||||
model_info: Final = ProviderConfigManager.get_provider_model_info(
|
||||
model=None, provider=LlmProviders(custom_llm_provider)
|
||||
)
|
||||
return None if model_info is None else model_info.get_base_model(model)
|
||||
except (ValueError, NotImplementedError):
|
||||
return None
|
||||
|
||||
|
||||
def _is_route_segment(custom_llm_provider: str, segment: str, model_after_segment: str) -> bool:
|
||||
if custom_llm_provider == "bedrock":
|
||||
from litellm.llms.bedrock.common_utils import get_all_bedrock_regions
|
||||
|
||||
return segment in BEDROCK_ANY_MODEL_ROUTES or segment in get_all_bedrock_regions()
|
||||
base_model: Final = _provider_base_model(custom_llm_provider, f"{segment}/{model_after_segment}")
|
||||
return base_model is not None and not base_model.startswith(f"{segment}/")
|
||||
|
||||
|
||||
def wildcard_health_check_models(wildcard: str, custom_llm_provider: str) -> tuple[str, ...]:
|
||||
from litellm.litellm_core_utils.llm_request_utils import rank_chat_models_from_llm_provider
|
||||
from litellm.router_utils.pattern_match_deployments import PatternMatchRouter
|
||||
|
||||
provider_prefix: Final = f"{custom_llm_provider}/"
|
||||
segments: Final = wildcard.removeprefix(provider_prefix).split("/")
|
||||
route_segments: Final = tuple(
|
||||
itertools.takewhile(
|
||||
lambda segment_index: _is_route_segment(
|
||||
custom_llm_provider, segments[segment_index], "/".join(segments[segment_index + 1 :])
|
||||
),
|
||||
range(len(segments) - 1),
|
||||
)
|
||||
)
|
||||
route_prefix: Final = "".join(f"{segments[segment_index]}/" for segment_index in route_segments)
|
||||
model_regex: Final = re.compile(PatternMatchRouter().pattern_to_regex("/".join(segments[len(route_segments) :])))
|
||||
ranked_model_ids: Final = tuple(
|
||||
(ranked_model.model.removeprefix(provider_prefix), ranked_model)
|
||||
for ranked_model in rank_chat_models_from_llm_provider(custom_llm_provider)
|
||||
)
|
||||
served_model_ids: Final = sorted(
|
||||
(ranked for ranked in ranked_model_ids if model_regex.fullmatch(ranked[0])),
|
||||
key=lambda ranked: (ranked[1].is_past_deprecation_date, ranked[1].route_rank, ranked[1].cost, ranked[0]),
|
||||
)
|
||||
served_probes: Final = itertools.islice(
|
||||
dict.fromkeys(f"{provider_prefix}{route_prefix}{model_id}" for model_id, _ in served_model_ids),
|
||||
WILDCARD_HEALTH_CHECK_SERVED_MODEL_COUNT,
|
||||
)
|
||||
provider_probes: Final = (
|
||||
f"{provider_prefix}{model_id}" for model_id, _ in ranked_model_ids[:WILDCARD_HEALTH_CHECK_PROVIDER_MODEL_COUNT]
|
||||
)
|
||||
return tuple(dict.fromkeys((*served_probes, *provider_probes)))
|
||||
|
||||
|
||||
class HealthCheckHelpers:
|
||||
@staticmethod
|
||||
async def ahealth_check_wildcard_models(
|
||||
|
|
@ -99,23 +160,16 @@ class HealthCheckHelpers:
|
|||
litellm_logging_obj: "Logging",
|
||||
) -> dict:
|
||||
from litellm import acompletion
|
||||
from litellm.litellm_core_utils.llm_request_utils import (
|
||||
pick_cheapest_chat_models_from_llm_provider,
|
||||
)
|
||||
|
||||
# this is a wildcard model, we need to pick a random model from the provider
|
||||
cheapest_models = pick_cheapest_chat_models_from_llm_provider(custom_llm_provider=custom_llm_provider, n=3)
|
||||
if len(cheapest_models) == 0:
|
||||
probe_models: Final = wildcard_health_check_models(wildcard=model, custom_llm_provider=custom_llm_provider)
|
||||
if len(probe_models) == 0:
|
||||
raise Exception(
|
||||
f"Unable to health check wildcard model for provider {custom_llm_provider}. Add a model on your config.yaml or contribute here - https://github.com/BerriAI/litellm/blob/main/model_prices_and_context_window.json"
|
||||
f"Unable to health check wildcard model {model} for provider {custom_llm_provider}: the cost map has "
|
||||
"no chat model for this provider. Set model_info.health_check_model to a model this wildcard serves"
|
||||
)
|
||||
if len(cheapest_models) > 1:
|
||||
fallback_models = cheapest_models[1:] # Pick the last 2 models from the shuffled list
|
||||
else:
|
||||
fallback_models = None
|
||||
model_params["model"] = cheapest_models[0]
|
||||
model_params["model"] = probe_models[0]
|
||||
model_params["litellm_logging_obj"] = litellm_logging_obj
|
||||
model_params["fallbacks"] = fallback_models
|
||||
model_params["fallbacks"] = list(probe_models[1:]) or None
|
||||
model_params["max_tokens"] = model_params.get("max_tokens", 16) # GPT-5 models require max_output_tokens >= 16
|
||||
await acompletion(**model_params)
|
||||
return {}
|
||||
|
|
|
|||
|
|
@ -1,6 +1,11 @@
|
|||
import datetime
|
||||
from collections.abc import Iterable, Iterator, Mapping
|
||||
from dataclasses import dataclass
|
||||
from typing import Final, cast
|
||||
|
||||
from pydantic import TypeAdapter
|
||||
from typing_extensions import NotRequired, ReadOnly, TypedDict
|
||||
|
||||
import litellm
|
||||
from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH
|
||||
from litellm.litellm_core_utils.get_llm_provider_logic import inferred_provider
|
||||
|
|
@ -151,19 +156,52 @@ def pick_cheapest_chat_models_from_llm_provider(custom_llm_provider: str, n=1):
|
|||
Returns:
|
||||
list[str]: A list of the n cheapest chat models.
|
||||
"""
|
||||
if custom_llm_provider not in litellm.models_by_provider:
|
||||
return []
|
||||
return [ranked_model.model for ranked_model in rank_chat_models_from_llm_provider(custom_llm_provider)[:n]]
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class RankedChatModel:
|
||||
model: str
|
||||
route_rank: int
|
||||
cost: float
|
||||
is_past_deprecation_date: bool
|
||||
|
||||
|
||||
def rank_chat_models_from_llm_provider(custom_llm_provider: str) -> tuple[RankedChatModel, ...]:
|
||||
known_models: Final[Iterable[str]] = litellm.models_by_provider.get(custom_llm_provider, ())
|
||||
ranked_rows: Final = (
|
||||
(route_rank, cost, model)
|
||||
for model, cost in _chat_row_costs(known_models, custom_llm_provider)
|
||||
ranked_models: Final = (
|
||||
RankedChatModel(
|
||||
model=model, route_rank=route_rank, cost=cost, is_past_deprecation_date=is_past_deprecation_date
|
||||
)
|
||||
for model, cost, is_past_deprecation_date in _chat_row_costs(known_models, custom_llm_provider)
|
||||
if (route_rank := _route_rank(model, custom_llm_provider)) is not None
|
||||
)
|
||||
return [model for _, _, model in sorted(ranked_rows, key=lambda row: row[:2])[:n]]
|
||||
return tuple(sorted(ranked_models, key=lambda ranked_model: (ranked_model.route_rank, ranked_model.cost)))
|
||||
|
||||
|
||||
def _chat_row_costs(models: Iterable[str], custom_llm_provider: str) -> Iterator[tuple[str, float]]:
|
||||
class _CostMapDeprecation(TypedDict):
|
||||
deprecation_date: NotRequired[ReadOnly[object]]
|
||||
|
||||
|
||||
_COST_MAP_DEPRECATION: Final = TypeAdapter(_CostMapDeprecation)
|
||||
|
||||
|
||||
def _is_past_deprecation_date(cost_map_key: str) -> bool:
|
||||
deprecation_date: Final = _COST_MAP_DEPRECATION.validate_python(
|
||||
litellm.model_cost.get(cost_map_key) or {} # pyright: ignore[reportUnknownMemberType] # untyped cost map JSON
|
||||
).get("deprecation_date")
|
||||
if not isinstance(deprecation_date, str):
|
||||
return False
|
||||
try:
|
||||
return (
|
||||
datetime.date.fromisoformat(deprecation_date)
|
||||
<= datetime.datetime.now(datetime.timezone.utc).astimezone().date()
|
||||
)
|
||||
except ValueError:
|
||||
return False
|
||||
|
||||
|
||||
def _chat_row_costs(models: Iterable[str], custom_llm_provider: str) -> Iterator[tuple[str, float, bool]]:
|
||||
for model in models:
|
||||
try:
|
||||
model_info = litellm.get_model_info(model=model, custom_llm_provider=custom_llm_provider)
|
||||
|
|
@ -171,7 +209,11 @@ def _chat_row_costs(models: Iterable[str], custom_llm_provider: str) -> Iterator
|
|||
continue
|
||||
if model_info.get("mode") != "chat":
|
||||
continue
|
||||
yield model, (model_info.get("input_cost_per_token") or 0.0) + (model_info.get("output_cost_per_token") or 0.0)
|
||||
yield (
|
||||
model,
|
||||
(model_info.get("input_cost_per_token") or 0.0) + (model_info.get("output_cost_per_token") or 0.0),
|
||||
_is_past_deprecation_date(model_info["key"]),
|
||||
)
|
||||
|
||||
|
||||
def _route_rank(model: str, custom_llm_provider: str) -> int | None:
|
||||
|
|
|
|||
|
|
@ -5878,38 +5878,36 @@ def function_call_prompt(
|
|||
]
|
||||
|
||||
|
||||
def response_schema_prompt(model: str, response_schema: dict, llm_provider: str) -> str:
|
||||
def response_schema_prompt(model: str, response_schema: Mapping[str, object], llm_provider: str) -> str:
|
||||
"""
|
||||
Decides if a user-defined custom prompt or default needs to be used
|
||||
|
||||
Returns the prompt str that's passed to the model as a user message
|
||||
"""
|
||||
custom_prompt_details: dict | None = None
|
||||
response_schema_as_message: Final = [{"role": "user", "content": f"{response_schema}"}]
|
||||
if f"{model}/response_schema_prompt" in litellm.custom_prompt_dict:
|
||||
custom_prompt_details = litellm.custom_prompt_dict[
|
||||
f"{model}/response_schema_prompt"
|
||||
] # allow user to define custom response schema prompt by model
|
||||
elif "response_schema_prompt" in litellm.custom_prompt_dict:
|
||||
custom_prompt_details = litellm.custom_prompt_dict["response_schema_prompt"]
|
||||
|
||||
if custom_prompt_details is not None:
|
||||
return render_custom_prompt_template(
|
||||
template=custom_prompt_details, messages=response_schema_as_message, model=model, llm_provider=llm_provider
|
||||
)
|
||||
else:
|
||||
prompts: Final[Mapping[str, Mapping[str, object]]] = litellm.custom_prompt_dict
|
||||
model_prompt_key: Final = f"{model}/response_schema_prompt"
|
||||
custom_prompt_details: Final = prompts.get(
|
||||
model_prompt_key if model_prompt_key in prompts else "response_schema_prompt"
|
||||
)
|
||||
if custom_prompt_details is None:
|
||||
return default_response_schema_prompt(response_schema=response_schema)
|
||||
return render_custom_prompt_template(
|
||||
template=custom_prompt_details,
|
||||
messages=[{"role": "user", "content": f"{response_schema}"}],
|
||||
model=model,
|
||||
llm_provider=llm_provider,
|
||||
)
|
||||
|
||||
|
||||
def default_response_schema_prompt(response_schema: dict) -> str:
|
||||
def default_response_schema_prompt(response_schema: Mapping[str, object]) -> str:
|
||||
"""
|
||||
Used if provider/model doesn't support 'response_schema' param.
|
||||
|
||||
This is the default prompt. Allow user to override this with a custom_prompt.
|
||||
"""
|
||||
prompt_str: Final = f"""Use this JSON schema:
|
||||
```json
|
||||
{response_schema}
|
||||
```json
|
||||
{json.dumps(response_schema, default=str)}
|
||||
```"""
|
||||
return prompt_str
|
||||
|
||||
|
|
|
|||
|
|
@ -1,5 +1,6 @@
|
|||
import asyncio
|
||||
import json
|
||||
import time
|
||||
import traceback
|
||||
from collections.abc import Callable, Coroutine, Mapping, MutableMapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
|
|
@ -15,13 +16,19 @@ from litellm._logging import redact_internal_details_from_client_message, verbos
|
|||
from litellm._uuid import uuid
|
||||
from litellm.constants import (
|
||||
BEDROCK_REALTIME_PENDING_OPENING_MESSAGE_SCOPE_KEY,
|
||||
REALTIME_LAZY_SETUP_CONTENT_GRACE_SECONDS,
|
||||
REALTIME_LAZY_SETUP_DEADLINE_SECONDS,
|
||||
REALTIME_METERED_USAGE_KEY,
|
||||
REALTIME_SESSION_FAILURE_LOGGED_KEY,
|
||||
REALTIME_SESSION_SUCCESS_LOGGED_KEY,
|
||||
)
|
||||
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
|
||||
from litellm.llms.base_llm.realtime.transcription_protocol import RealtimeTranscriptionProtocolError
|
||||
from litellm.llms.base_llm.realtime.transformation import BaseRealtimeConfig, RealtimeBackend
|
||||
from litellm.llms.base_llm.realtime.transformation import (
|
||||
BaseRealtimeConfig,
|
||||
RealtimeBackend,
|
||||
RealtimeSessionUpdateRejection,
|
||||
)
|
||||
from litellm.types.llms.openai import (
|
||||
OpenAIRealtimeEvents,
|
||||
OpenAIRealtimeOutputItemDone,
|
||||
|
|
@ -125,6 +132,65 @@ class _LoggingWorker(Protocol):
|
|||
_DECODED_JSON: Final = TypeAdapter(object)
|
||||
|
||||
|
||||
_SESSION_FIELDS: Final = TypeAdapter(dict[str, object])
|
||||
_NO_SESSION_FIELDS: Final[Mapping[str, object]] = MappingProxyType({})
|
||||
_SESSION_MERGE_MAX_DEPTH: Final = 8
|
||||
|
||||
|
||||
class _SessionUpdateFrame(TypedDict):
|
||||
type: ReadOnly[str]
|
||||
session: ReadOnly[Mapping[str, object]]
|
||||
|
||||
|
||||
def _merge_session_fields(
|
||||
base: Mapping[str, object], update: Mapping[str, object], depth: int = 0
|
||||
) -> Mapping[str, object]:
|
||||
def merged_value(key: str) -> object:
|
||||
if key not in update:
|
||||
return base[key]
|
||||
value: Final = update[key]
|
||||
base_value: Final = base.get(key)
|
||||
if depth < _SESSION_MERGE_MAX_DEPTH and isinstance(value, Mapping) and isinstance(base_value, Mapping):
|
||||
return _merge_session_fields(
|
||||
_SESSION_FIELDS.validate_python(base_value), _SESSION_FIELDS.validate_python(value), depth + 1
|
||||
)
|
||||
return update[key]
|
||||
|
||||
return MappingProxyType({key: merged_value(key) for key in dict.fromkeys((*base, *update))})
|
||||
|
||||
|
||||
def _session_update_frame(session: Mapping[str, object]) -> str:
|
||||
frame: Final[_SessionUpdateFrame] = {"type": "session.update", "session": session}
|
||||
return json.dumps(frame, default=_SESSION_FIELDS.validate_python)
|
||||
|
||||
|
||||
_HELD_SESSION_UPDATE_TOO_LARGE: Final = RealtimeSessionUpdateRejection(
|
||||
code="session_update_too_large",
|
||||
message=(
|
||||
"The session.update frames sent before the first conversation content add up to more than 10 MB, "
|
||||
"so this session.update was not applied"
|
||||
),
|
||||
param="session",
|
||||
)
|
||||
|
||||
|
||||
def _unscreened_audio_error_event(client_event_id: object) -> Mapping[str, object]:
|
||||
return {
|
||||
"type": "error",
|
||||
"event_id": f"event_{uuid.uuid4()}",
|
||||
"error": {
|
||||
"type": "invalid_request_error",
|
||||
"code": "audio_input_not_screenable",
|
||||
"message": (
|
||||
"A realtime_input_transcription guardrail is on, and Gemini Live only transcribes a turn after it "
|
||||
"starts answering it, so audio input was not sent. Send text with conversation.item.create instead"
|
||||
),
|
||||
"param": None,
|
||||
"event_id": client_event_id if isinstance(client_event_id, str) else None,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def _decode_json_object(payload: str) -> Mapping[str, object]:
|
||||
"""Decode a realtime frame into its top-level field mapping."""
|
||||
return json.loads(payload)
|
||||
|
|
@ -243,8 +309,17 @@ class RealTimeStreaming:
|
|||
event_normalizer: RealtimeEventNormalizer | None = None,
|
||||
usage_meter: RealtimeUsageMeter | None = None,
|
||||
logging_worker: _LoggingWorker = GLOBAL_LOGGING_WORKER,
|
||||
lazy_setup_deadline_seconds: float = REALTIME_LAZY_SETUP_DEADLINE_SECONDS,
|
||||
lazy_setup_content_grace_seconds: float = REALTIME_LAZY_SETUP_CONTENT_GRACE_SECONDS,
|
||||
clock: Callable[[], float] = time.monotonic,
|
||||
):
|
||||
self.websocket: _ClientWebSocket = websocket
|
||||
self._clock: Final = clock
|
||||
self._lazy_setup_deadline: Final = clock() + lazy_setup_deadline_seconds
|
||||
self._lazy_setup_content_grace_seconds: Final = lazy_setup_content_grace_seconds
|
||||
self._held_client_content: tuple[str, ...] = ()
|
||||
self._held_client_content_bytes: int = 0
|
||||
self._held_client_content_deadline: float | None = None
|
||||
self._replay_scope: Final = client_websocket_scope(websocket)
|
||||
self.backend_ws = backend_ws
|
||||
self.logging_obj = logging_obj
|
||||
|
|
@ -300,6 +375,10 @@ class RealTimeStreaming:
|
|||
# Gemini Live rejects a follow-up BidiGenerateContentSetup once any
|
||||
# content (realtimeInput / clientContent / toolResponse) has been sent.
|
||||
self._content_sent_after_setup: bool = False
|
||||
self._held_session_update: Mapping[str, object] | None = None
|
||||
self._applied_session_fields: Mapping[str, object] | None = None
|
||||
self._held_session_update_bytes: int = 0
|
||||
self._unscreened_audio_refused: bool = False
|
||||
# Whether this is a transcription-only session (session.type == "transcription",
|
||||
# e.g. gpt-realtime-whisper). Such sessions must not be sent response.create and
|
||||
# their input_audio_transcription.completed usage drives duration-based cost.
|
||||
|
|
@ -337,6 +416,17 @@ class RealTimeStreaming:
|
|||
]
|
||||
)
|
||||
_CLIENT_AUDIO_BUFFER_COMMIT_TYPES = frozenset(["input_audio_buffer.commit", "input_audio_buffer.end"])
|
||||
_CLIENT_CONTENT_TYPES = frozenset(
|
||||
[
|
||||
"conversation.item.create",
|
||||
"response.create",
|
||||
"input_audio_buffer.append",
|
||||
"input_audio_buffer.commit",
|
||||
"input_audio_buffer.end",
|
||||
]
|
||||
)
|
||||
_HELD_CONTENT_TYPES: Final = frozenset(["conversation.item.create", "input_audio_buffer.append"])
|
||||
_DEFERRED_SETUP_HELD_CONTENT_TYPES: Final = frozenset(["conversation.item.create"])
|
||||
_AUDIO_FORMAT_MAP: dict[str, dict[str, str | int]] = {
|
||||
"pcm16": {"type": "audio/pcm", "rate": 24000},
|
||||
"g711_ulaw": {"type": "audio/G711-ulaw", "rate": 8000},
|
||||
|
|
@ -377,7 +467,7 @@ class RealTimeStreaming:
|
|||
message_obj: dict[str, Any] = cast(dict[str, Any], message)
|
||||
else:
|
||||
message_obj = cast(dict[str, Any], json.loads(message))
|
||||
self._collect_tool_calls_from_response_done(cast(dict, message_obj))
|
||||
self._collect_tool_calls_from_response_done(message_obj)
|
||||
if not self._should_store_message(message_obj):
|
||||
return
|
||||
try:
|
||||
|
|
@ -502,7 +592,7 @@ class RealTimeStreaming:
|
|||
self.store_message(flush_event)
|
||||
self._capture_transcription_usage(flush_event)
|
||||
|
||||
def _collect_tool_calls_from_response_done(self, event_obj: dict | OpenAIRealtimeEvents) -> None:
|
||||
def _collect_tool_calls_from_response_done(self, event_obj: Mapping[str, object] | OpenAIRealtimeEvents) -> None:
|
||||
"""Extract function_call items from response.done events for spend logging."""
|
||||
try:
|
||||
if event_obj.get("type") != "response.done":
|
||||
|
|
@ -607,6 +697,167 @@ class RealTimeStreaming:
|
|||
self._usage_meter.observe_client_event(sent_event)
|
||||
return True
|
||||
|
||||
def _builds_setup_lazily(self) -> bool:
|
||||
if self.provider_config is None:
|
||||
return False
|
||||
return self.provider_config.builds_setup_from_session_update()
|
||||
|
||||
def _holds_session_updates_until_first_content(self) -> bool:
|
||||
return self.session_configuration_request is None and self._builds_setup_lazily()
|
||||
|
||||
async def _rejected_setup_session_update(self, message: str) -> bool:
|
||||
provider_config: Final = self.provider_config
|
||||
if provider_config is None or not provider_config.builds_setup_from_session_update():
|
||||
return False
|
||||
client_event: Final = _decode_json_object(message)
|
||||
session: Final = client_event.get("session")
|
||||
update: Final = _SESSION_FIELDS.validate_python(session) if isinstance(session, Mapping) else _NO_SESSION_FIELDS
|
||||
if self._leaves_applied_setup_unchanged(update):
|
||||
await self._acknowledge_session_setup(self._applied_session_fields or _NO_SESSION_FIELDS)
|
||||
return True
|
||||
rejection: Final = provider_config.setup_session_update_rejection(
|
||||
self.model, update, setup_sent=self.session_configuration_request is not None
|
||||
)
|
||||
if rejection is None:
|
||||
return False
|
||||
await self.websocket.send_text(json.dumps(rejection.error_event(client_event.get("event_id"))))
|
||||
return True
|
||||
|
||||
def _held_session_update_message(self) -> str:
|
||||
return _session_update_frame(self._held_session_update or _NO_SESSION_FIELDS)
|
||||
|
||||
async def _hold_session_update(self, message: str) -> None:
|
||||
client_event: Final = _decode_json_object(message)
|
||||
if self._held_session_update_bytes + len(message.encode("utf-8")) > RealTimeStreaming._MAX_BUFFERED_BYTES:
|
||||
await self.websocket.send_text(
|
||||
json.dumps(_HELD_SESSION_UPDATE_TOO_LARGE.error_event(client_event.get("event_id")))
|
||||
)
|
||||
return
|
||||
session: Final = client_event.get("session")
|
||||
update: Final = _SESSION_FIELDS.validate_python(session) if isinstance(session, Mapping) else _NO_SESSION_FIELDS
|
||||
held: Final = _merge_session_fields(self._held_session_update or _NO_SESSION_FIELDS, update)
|
||||
self._held_session_update = held
|
||||
self._held_session_update_bytes = len(_session_update_frame(held).encode("utf-8"))
|
||||
await self._acknowledge_session_setup(held)
|
||||
|
||||
def _setup_preview(self, session: Mapping[str, object]) -> str | None:
|
||||
if self.provider_config is None:
|
||||
return None
|
||||
preview: Final = self.provider_config.transform_realtime_request(
|
||||
_session_update_frame(session), self.model, None
|
||||
)
|
||||
return next(
|
||||
(frame for frame in preview if isinstance(frame, str) and "setup" in _decode_json_object(frame)), None
|
||||
)
|
||||
|
||||
def _leaves_applied_setup_unchanged(self, update: Mapping[str, object]) -> bool:
|
||||
applied: Final = self._applied_session_fields
|
||||
if self.session_configuration_request is None or applied is None:
|
||||
return False
|
||||
applied_setup: Final = self._setup_preview(applied)
|
||||
merged_setup: Final = self._setup_preview(_merge_session_fields(applied, update))
|
||||
return (
|
||||
applied_setup is not None
|
||||
and merged_setup is not None
|
||||
and _decode_json_object(applied_setup) == _decode_json_object(merged_setup)
|
||||
)
|
||||
|
||||
async def _acknowledge_session_setup(self, session: Mapping[str, object]) -> None:
|
||||
if self.provider_config is None:
|
||||
return
|
||||
setup: Final = self._setup_preview(session)
|
||||
if setup is None:
|
||||
return
|
||||
acknowledged: Final = self.provider_config.transform_session_created_event(
|
||||
self.model, str(self.logging_obj.litellm_trace_id), setup
|
||||
)
|
||||
if acknowledged is None:
|
||||
return
|
||||
event: Final[Mapping[str, object]] = {**acknowledged, "type": "session.updated"}
|
||||
event_str: Final = json.dumps(event)
|
||||
self.store_message(event_str)
|
||||
await self._send_event_to_client(event, event_str)
|
||||
|
||||
async def _send_held_session_update_as_setup(self) -> None:
|
||||
held: Final = self._held_session_update
|
||||
if held is None:
|
||||
return
|
||||
message: Final = self._held_session_update_message()
|
||||
self._held_session_update = None
|
||||
self._held_session_update_bytes = 0
|
||||
if await self._send_to_backend(message):
|
||||
self._applied_session_fields = held
|
||||
await self._send_held_client_content()
|
||||
|
||||
def _holds_client_content(self, msg_type: str | None, message: str) -> bool:
|
||||
held_types: Final = (
|
||||
RealTimeStreaming._DEFERRED_SETUP_HELD_CONTENT_TYPES
|
||||
if self._uses_deferred_backend_setup()
|
||||
else RealTimeStreaming._HELD_CONTENT_TYPES
|
||||
)
|
||||
return (
|
||||
msg_type in held_types
|
||||
and self._held_session_update is not None
|
||||
and self._holds_session_updates_until_first_content()
|
||||
and not self._pending_messages_until_setup
|
||||
and len(self._held_client_content) < RealTimeStreaming._MAX_BUFFERED_MESSAGES
|
||||
and self._held_client_content_bytes + len(message.encode("utf-8")) <= RealTimeStreaming._MAX_BUFFERED_BYTES
|
||||
)
|
||||
|
||||
def _hold_client_content(self, message: str) -> None:
|
||||
if self._held_client_content_deadline is None:
|
||||
self._held_client_content_deadline = self._clock() + self._lazy_setup_content_grace_seconds
|
||||
self._held_client_content = (*self._held_client_content, message)
|
||||
self._held_client_content_bytes += len(message.encode("utf-8"))
|
||||
|
||||
async def _send_held_client_content(self) -> None:
|
||||
held: Final = self._held_client_content
|
||||
self._held_client_content = ()
|
||||
self._held_client_content_bytes = 0
|
||||
self._held_client_content_deadline = None
|
||||
for message in held:
|
||||
await self._send_to_backend(message)
|
||||
|
||||
def _setup_deadline(self) -> float:
|
||||
content_deadline: Final = self._held_client_content_deadline
|
||||
return (
|
||||
self._lazy_setup_deadline if content_deadline is None else min(self._lazy_setup_deadline, content_deadline)
|
||||
)
|
||||
|
||||
async def _send_setup_before_content(self, msg_type: str | None) -> None:
|
||||
if self._holds_session_updates_until_first_content():
|
||||
await self._send_held_session_update_as_setup()
|
||||
await self._send_default_setup_before_first_client_message(msg_type)
|
||||
|
||||
async def _receive_client_message_before_setup_deadline(self) -> str | None:
|
||||
if not self._holds_session_updates_until_first_content():
|
||||
return await self._receive_client_message()
|
||||
receive: Final = asyncio.ensure_future(self._receive_client_message())
|
||||
try:
|
||||
done, _ = await asyncio.wait((receive,), timeout=self._setup_deadline() - self._clock())
|
||||
if not done:
|
||||
await self._send_setup_before_content(None)
|
||||
return await receive
|
||||
finally:
|
||||
receive.cancel()
|
||||
|
||||
async def _send_default_setup_before_first_client_message(self, msg_type: str | None) -> None:
|
||||
if self.session_configuration_request is not None or self.provider_config is None:
|
||||
return
|
||||
if msg_type == "session.update" or (
|
||||
self._uses_deferred_backend_setup()
|
||||
and (msg_type in RealTimeStreaming._CLIENT_AUDIO_BUFFER_TYPES or self._pending_messages_until_setup)
|
||||
):
|
||||
return
|
||||
default_setup: Final = self.provider_config.session_configuration_request(self.model)
|
||||
if not default_setup:
|
||||
return
|
||||
setup: Final = self.maybe_inject_guardrail_auto_response_disable(default_setup)
|
||||
await self.backend_ws.send(setup)
|
||||
self._cache_session_configuration_request(setup)
|
||||
self._applied_session_fields = _NO_SESSION_FIELDS
|
||||
await self._send_held_client_content()
|
||||
|
||||
def _enforce_transcription_session_model(self, message: str) -> str:
|
||||
"""Force client transcription session updates to the authorized model.
|
||||
|
||||
|
|
@ -851,6 +1102,20 @@ class RealTimeStreaming:
|
|||
self._guardrail_cancel_event_ids = self._guardrail_cancel_event_ids | {event_id}
|
||||
await self._send_to_backend(json.dumps({"type": "response.cancel", "event_id": event_id}))
|
||||
|
||||
async def _send_blocked_turn_to_backend(self, pre_block_backend_message: str | None, item_id: str | None) -> None:
|
||||
if pre_block_backend_message is None and self._is_transcription_session:
|
||||
return
|
||||
await self._send_setup_before_content("conversation.item.create")
|
||||
# Deliver any caller-supplied backend message FIRST so that
|
||||
# protocol contracts requiring a specific ordering (e.g.
|
||||
# Gemini Live's mandatory ``toolResponse`` after a
|
||||
# ``toolCall``) are honored before the guardrail's own
|
||||
# clientContent / cancel messages are sent.
|
||||
if pre_block_backend_message is not None:
|
||||
await self._send_to_backend(pre_block_backend_message)
|
||||
if not self._is_transcription_session:
|
||||
await self._withdraw_blocked_turn(item_id)
|
||||
|
||||
async def _withdraw_blocked_turn(self, item_id: str | None) -> None:
|
||||
await self._send_guardrail_cancel()
|
||||
if item_id is not None and self._deletes_conversation_items:
|
||||
|
|
@ -958,6 +1223,21 @@ class RealTimeStreaming:
|
|||
"""Transcription-only sessions have no assistant turn to gate and reject realtime session updates."""
|
||||
return not self._is_transcription_session and self._has_audio_transcription_guardrails()
|
||||
|
||||
def _refuses_unscreened_audio(self, msg_type: str | None) -> bool:
|
||||
return (
|
||||
msg_type in RealTimeStreaming._CLIENT_AUDIO_BUFFER_TYPES
|
||||
and self._builds_setup_lazily()
|
||||
and self._should_disable_vad_auto_response()
|
||||
)
|
||||
|
||||
async def _refuse_unscreened_audio(self, msg_type: str | None, message: str) -> None:
|
||||
if self._unscreened_audio_refused and msg_type not in RealTimeStreaming._CLIENT_AUDIO_BUFFER_COMMIT_TYPES:
|
||||
return
|
||||
self._unscreened_audio_refused = True
|
||||
await self.websocket.send_text(
|
||||
json.dumps(_unscreened_audio_error_event(_decode_json_object(message).get("event_id")))
|
||||
)
|
||||
|
||||
async def run_realtime_guardrails(
|
||||
self,
|
||||
transcript: str,
|
||||
|
|
@ -1030,15 +1310,7 @@ class RealTimeStreaming:
|
|||
# Use realtime_violation_message if configured; fall back to guardrail error text.
|
||||
error_msg = getattr(callback, "realtime_violation_message", None) or safe_msg
|
||||
|
||||
# Deliver any caller-supplied backend message FIRST so that
|
||||
# protocol contracts requiring a specific ordering (e.g.
|
||||
# Gemini Live's mandatory ``toolResponse`` after a
|
||||
# ``toolCall``) are honored before the guardrail's own
|
||||
# clientContent / cancel messages are sent.
|
||||
if pre_block_backend_message is not None:
|
||||
await self._send_to_backend(pre_block_backend_message)
|
||||
if not self._is_transcription_session:
|
||||
await self._withdraw_blocked_turn(item_id)
|
||||
await self._send_blocked_turn_to_backend(pre_block_backend_message, item_id)
|
||||
# Send the policy violation hint (shows as small gray status text in UI).
|
||||
await self.websocket.send_text(
|
||||
json.dumps(
|
||||
|
|
@ -1136,6 +1408,9 @@ class RealTimeStreaming:
|
|||
break
|
||||
finally:
|
||||
self._flushing_pending_messages_until_setup = False
|
||||
if self.session_created_sent_to_client and self._builds_setup_lazily():
|
||||
await self._maybe_send_guardrail_turn_detection_update()
|
||||
continue
|
||||
if self.session_created_sent_to_client:
|
||||
# A synthetic session.created (with placeholder defaults) was
|
||||
# already forwarded to the client when we connected. The
|
||||
|
|
@ -1295,7 +1570,11 @@ class RealTimeStreaming:
|
|||
return BackendClose(code=1011, reason="proxy failed while relaying the upstream websocket")
|
||||
|
||||
def _backend_refused_session(self, close: BackendClose) -> bool:
|
||||
return close.code != 1000 and not self._backend_sent_frames
|
||||
return (
|
||||
close.code != 1000
|
||||
and not self._backend_sent_frames
|
||||
and not self._holds_session_updates_until_first_content()
|
||||
)
|
||||
|
||||
async def _log_session_end(self, transport_refusal: Exception | None) -> None:
|
||||
provider_refusal: Final = None if self.provider_config is None else self.provider_config.session_refusal()
|
||||
|
|
@ -1488,7 +1767,7 @@ class RealTimeStreaming:
|
|||
client_event: _ClientEventFrame
|
||||
try:
|
||||
while True:
|
||||
message = await self._receive_client_message()
|
||||
message = await self._receive_client_message_before_setup_deadline()
|
||||
if message is None:
|
||||
return ClientLoopExit.CLIENT_DISCONNECTED
|
||||
|
||||
|
|
@ -1600,6 +1879,7 @@ class RealTimeStreaming:
|
|||
if (
|
||||
msg_type == "session.update"
|
||||
and self.session_configuration_request is None
|
||||
and self._held_session_update is None
|
||||
and not self._guardrail_turn_detection_update_sent
|
||||
and self._should_disable_vad_auto_response()
|
||||
and (
|
||||
|
|
@ -1663,12 +1943,25 @@ class RealTimeStreaming:
|
|||
except (json.JSONDecodeError, AttributeError):
|
||||
pass
|
||||
|
||||
if self._refuses_unscreened_audio(msg_type):
|
||||
await self._refuse_unscreened_audio(msg_type, message)
|
||||
continue
|
||||
|
||||
## LOGGING
|
||||
# Log after any in-place modifications (GA remap, guardrail
|
||||
# turn_detection injection) so audit logs reflect what we
|
||||
# actually forward to the backend.
|
||||
self.store_input(message=message)
|
||||
|
||||
if self._holds_client_content(msg_type, message):
|
||||
self._hold_client_content(message)
|
||||
continue
|
||||
if self._held_client_content and msg_type != "session.update":
|
||||
await self._send_setup_before_content(msg_type)
|
||||
|
||||
if msg_type in RealTimeStreaming._CLIENT_CONTENT_TYPES and self._uses_deferred_backend_setup():
|
||||
await self._send_setup_before_content(msg_type)
|
||||
|
||||
if self._should_buffer_client_message_until_setup(message):
|
||||
self._buffer_pending_message_until_setup(message)
|
||||
continue
|
||||
|
|
@ -1697,6 +1990,15 @@ class RealTimeStreaming:
|
|||
# ``_maybe_send_guardrail_turn_detection_update`` would retry.
|
||||
if msg_type == "session.update" and self.provider_config is None:
|
||||
await self._send_session_gate_once()
|
||||
if msg_type == "session.update" and await self._rejected_setup_session_update(message):
|
||||
continue
|
||||
if msg_type == "session.update" and self._holds_session_updates_until_first_content():
|
||||
await self._hold_session_update(message)
|
||||
if self._pending_messages_until_setup:
|
||||
await self._send_held_session_update_as_setup()
|
||||
continue
|
||||
if msg_type in RealTimeStreaming._CLIENT_CONTENT_TYPES:
|
||||
await self._send_setup_before_content(msg_type)
|
||||
try:
|
||||
sent = await self._send_to_backend(message)
|
||||
except RealtimeTranscriptionProtocolError as e:
|
||||
|
|
|
|||
|
|
@ -13,9 +13,10 @@ from typing_extensions import ReadOnly, Required
|
|||
import litellm
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.litellm_core_utils.core_helpers import (
|
||||
USAGE_COST_BILLED_PROVIDERS,
|
||||
PROVIDERS_BILLED_FROM_THEIR_OWN_USAGE_FIELDS,
|
||||
chunk_usage,
|
||||
provider_reported_usage_cost,
|
||||
stamp_litellm_bill_on_usage,
|
||||
)
|
||||
from litellm.litellm_core_utils.hidden_params import HIDDEN_PARAMS_ATTR, set_hidden_params
|
||||
from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import BEDROCK_TRACE_FIELDS, object_items, object_mapping
|
||||
|
|
@ -1380,9 +1381,6 @@ def stream_chunk_builder_text_completion(
|
|||
return TextCompletionResponse(**response)
|
||||
|
||||
|
||||
_CALCULATOR_PRICED_REPORTED_COST_PROVIDERS: Final = frozenset({LlmProviders.XAI.value}) | USAGE_COST_BILLED_PROVIDERS
|
||||
|
||||
|
||||
def _logged_provider(logging_obj: Optional["Logging"]) -> object:
|
||||
if logging_obj is None:
|
||||
return None
|
||||
|
|
@ -1390,10 +1388,6 @@ def _logged_provider(logging_obj: Optional["Logging"]) -> object:
|
|||
return provider
|
||||
|
||||
|
||||
def _reported_cost_is_priced_by_calculator(logging_obj: Optional["Logging"]) -> bool:
|
||||
return _logged_provider(logging_obj) in _CALCULATOR_PRICED_REPORTED_COST_PROVIDERS
|
||||
|
||||
|
||||
def _stream_builder_response_cost(
|
||||
response: ModelResponse, logging_obj: Optional["Logging"], chunks: Sequence[object]
|
||||
) -> float | None:
|
||||
|
|
@ -1404,9 +1398,6 @@ def _stream_builder_response_cost(
|
|||
for chunk in reversed(chunks)
|
||||
)
|
||||
return next((cost for cost in reported if cost is not None), None)
|
||||
usage_cost: Final = getattr(getattr(response, "usage", None), "cost", None)
|
||||
if isinstance(usage_cost, (int, float)) and not _reported_cost_is_priced_by_calculator(logging_obj):
|
||||
return float(usage_cost)
|
||||
if logging_obj is not None:
|
||||
return None
|
||||
provider_hint: Final = response._hidden_params.get( # pyright: ignore[reportPrivateUsage] # no public accessor
|
||||
|
|
@ -1453,13 +1444,15 @@ def _set_stream_builder_response_cost(
|
|||
def _stamp_streaming_usage_cost(usage: Usage, response: ModelResponse, logging_obj: Optional["Logging"]) -> None:
|
||||
if logging_obj is None:
|
||||
return
|
||||
if isinstance(getattr(usage, "cost", None), (int, float)):
|
||||
if _logged_provider(logging_obj) in PROVIDERS_BILLED_FROM_THEIR_OWN_USAGE_FIELDS and isinstance(
|
||||
getattr(usage, "cost", None), (int, float)
|
||||
):
|
||||
return
|
||||
computed_cost: Final = logging_obj.response_cost_calculator(result=response)
|
||||
if not isinstance(computed_cost, (int, float)) or computed_cost <= 0:
|
||||
if isinstance(computed_cost, bool) or not isinstance(computed_cost, (int, float)) or computed_cost < 0:
|
||||
return
|
||||
setattr(usage, "cost", computed_cost)
|
||||
_set_hidden_response_cost(response, float(computed_cost))
|
||||
if computed_cost > 0 or isinstance(getattr(usage, "cost", None), (int, float)):
|
||||
stamp_litellm_bill_on_usage(response, usage, float(computed_cost))
|
||||
|
||||
|
||||
def _field_of(value: object, name: str) -> object:
|
||||
|
|
|
|||
|
|
@ -17,6 +17,7 @@ import anyio
|
|||
import anyio.lowlevel
|
||||
import httpx
|
||||
import tiktoken
|
||||
from pydantic import TypeAdapter
|
||||
from tiktoken.model import MODEL_TO_ENCODING
|
||||
from typing_extensions import ParamSpec, TypeVar
|
||||
|
||||
|
|
@ -69,6 +70,7 @@ from litellm.types.llms.anthropic import (
|
|||
AnthropicMessagesToolResultParam,
|
||||
AnthropicMessagesToolUseParam,
|
||||
AnthropicSearchResultParam,
|
||||
CompactionBlock,
|
||||
)
|
||||
from litellm.types.llms.openai import (
|
||||
AllMessageValues,
|
||||
|
|
@ -937,6 +939,8 @@ def _count_messages(
|
|||
default_token_count,
|
||||
params.pdf_pricing,
|
||||
)
|
||||
elif key == "provider_specific_fields":
|
||||
num_tokens += sum(map(params.count_function, _compaction_summaries(value)))
|
||||
elif key == "search_results" and isinstance(value, list):
|
||||
from litellm.litellm_core_utils.prompt_templates.common_utils import (
|
||||
extract_search_results_text,
|
||||
|
|
@ -951,6 +955,25 @@ def _count_messages(
|
|||
return num_tokens
|
||||
|
||||
|
||||
_OBJECT_LIST: Final = TypeAdapter(list[object])
|
||||
|
||||
|
||||
def _compaction_summaries(provider_specific_fields: object) -> tuple[str, ...]:
|
||||
fields: Final = as_str_mapping(provider_specific_fields)
|
||||
blocks: Final = fields.get("compaction_blocks") if fields is not None else None
|
||||
block_fields: Final = (
|
||||
tuple(map(as_str_mapping, _OBJECT_LIST.validate_python(blocks))) if isinstance(blocks, list) else ()
|
||||
)
|
||||
return tuple(
|
||||
content
|
||||
for block in block_fields
|
||||
if block is not None
|
||||
and block.get("type") == "compaction"
|
||||
and isinstance(content := block.get("content"), str)
|
||||
and content
|
||||
)
|
||||
|
||||
|
||||
def _count_extra(
|
||||
count_function: TokenCounterFunction,
|
||||
tools: list[ChatCompletionToolParam] | None,
|
||||
|
|
@ -1648,6 +1671,11 @@ def _count_anthropic_content(
|
|||
return tokens
|
||||
|
||||
|
||||
_TEXT_FIELD_BY_BLOCK_TYPE: Final = MappingProxyType(
|
||||
{"text": "text", "thinking": "thinking", "redacted_thinking": "thinking", "compaction": "content"}
|
||||
)
|
||||
|
||||
|
||||
def _count_content_list(
|
||||
count_function: TokenCounterFunction,
|
||||
content_list: str
|
||||
|
|
@ -1659,6 +1687,7 @@ def _count_content_list(
|
|||
| AnthropicMessagesImageParam
|
||||
| AnthropicMessagesDocumentParam
|
||||
| AnthropicSearchResultParam
|
||||
| CompactionBlock
|
||||
],
|
||||
use_default_image_token_count: bool,
|
||||
default_token_count: int | None,
|
||||
|
|
@ -1670,8 +1699,10 @@ def _count_content_list(
|
|||
for c in content_list:
|
||||
if isinstance(c, str):
|
||||
num_tokens += count_function(c)
|
||||
elif c["type"] == "text":
|
||||
num_tokens += count_function(str(c.get("text", "")))
|
||||
elif c["type"] in _TEXT_FIELD_BY_BLOCK_TYPE:
|
||||
block_text = str(c.get(_TEXT_FIELD_BY_BLOCK_TYPE[c["type"]]) or "")
|
||||
if block_text:
|
||||
num_tokens += count_function(block_text)
|
||||
elif c["type"] == "image_url":
|
||||
image_url = c.get("image_url")
|
||||
num_tokens += _count_image_tokens(image_url, use_default_image_token_count)
|
||||
|
|
@ -1704,12 +1735,6 @@ def _count_content_list(
|
|||
default_token_count,
|
||||
pdf_pricing,
|
||||
)
|
||||
elif c["type"] in ("thinking", "redacted_thinking"):
|
||||
# Claude extended thinking content block
|
||||
# Count the thinking text and skip the opaque blobs (signature, redacted data)
|
||||
thinking_text = str(c.get("thinking", ""))
|
||||
if thinking_text:
|
||||
num_tokens += count_function(thinking_text)
|
||||
elif c["type"] == "tool_reference":
|
||||
# Anthropic tool-search reference block: a lightweight pointer to
|
||||
# a deferred tool, e.g. {"type": "tool_reference", "tool_name": ...}.
|
||||
|
|
@ -1727,7 +1752,7 @@ def _count_content_list(
|
|||
f"Invalid content item type: {content_type}. "
|
||||
f"Expected str or dict with 'type' field "
|
||||
f"(text, image_url, image, document, file, tool_use, tool_result, thinking, redacted_thinking, "
|
||||
f"tool_reference, search_result)."
|
||||
f"compaction, tool_reference, search_result)."
|
||||
)
|
||||
return num_tokens
|
||||
except Exception as e:
|
||||
|
|
|
|||
|
|
@ -24,7 +24,7 @@ from pydantic import ConfigDict, TypeAdapter, ValidationError, with_config
|
|||
from typing_extensions import ReadOnly, TypedDict, assert_never
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.litellm_core_utils.core_helpers import as_str_mapping
|
||||
from litellm.litellm_core_utils.core_helpers import as_str_mapping, iter_items
|
||||
from litellm.llms.anthropic.chat.transformation import AnthropicConfig
|
||||
from litellm.llms.anthropic.messages.adapters.transformation import (
|
||||
OPENAI_MAX_TOOL_NAME_LENGTH,
|
||||
|
|
@ -61,6 +61,7 @@ from litellm.llms.base_llm.guardrail_translation.utils import (
|
|||
unappliable_request_rewrite,
|
||||
)
|
||||
from litellm.types.llms.anthropic import (
|
||||
ANTHROPIC_SERVER_TOOL_USE_ID_PREFIX,
|
||||
AllAnthropicToolsValues,
|
||||
AnthropicMessagesRequest,
|
||||
)
|
||||
|
|
@ -153,6 +154,7 @@ class _TextAttrBlock(Protocol):
|
|||
text: str
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class _WritableMessage(Protocol):
|
||||
@overload
|
||||
def get(self, key: str, /) -> object | None: ...
|
||||
|
|
@ -195,6 +197,140 @@ class _ToolCallShape:
|
|||
arguments: str
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _WebSearchString:
|
||||
block_idx: int
|
||||
result_idx: int | None
|
||||
field: str
|
||||
value: str
|
||||
|
||||
|
||||
def _streamed_web_search_result_blocks(response: ModelResponse) -> tuple[Mapping[str, object], ...]:
|
||||
choice: Final = response.choices[0] if response.choices else None
|
||||
fields: Final = as_str_mapping(choice.message.provider_specific_fields) if isinstance(choice, Choices) else None
|
||||
blocks: Final = (
|
||||
(as_str_mapping(item) for item in iter_items(fields.get("web_search_results")))
|
||||
if fields is not None and isinstance(fields.get("web_search_results"), list)
|
||||
else ()
|
||||
)
|
||||
return tuple(
|
||||
deepcopy(block) for block in blocks if block is not None and block.get("type") == "web_search_tool_result"
|
||||
)
|
||||
|
||||
|
||||
def _web_search_results(block: Mapping[str, object]) -> Iterable[object]:
|
||||
return iter_items(block.get("content")) if isinstance(block.get("content"), list) else ()
|
||||
|
||||
|
||||
def _web_search_result_strings(block_idx: int, result_idx: int, result: object) -> tuple[_WebSearchString, ...]:
|
||||
fields: Final = as_str_mapping(result)
|
||||
if fields is None or fields.get("type") != "web_search_result":
|
||||
return ()
|
||||
values: Final = ((field, fields.get(field)) for field in ("title", "url"))
|
||||
return tuple(
|
||||
_WebSearchString(block_idx, result_idx, field, value)
|
||||
for field, value in values
|
||||
if isinstance(value, str) and value
|
||||
)
|
||||
|
||||
|
||||
def _is_server_tool_use(block: Mapping[str, object]) -> bool:
|
||||
block_id: Final = block.get("id")
|
||||
match block.get("type"):
|
||||
case "server_tool_use":
|
||||
return True
|
||||
case "tool_use":
|
||||
return isinstance(block_id, str) and block_id.startswith(ANTHROPIC_SERVER_TOOL_USE_ID_PREFIX)
|
||||
case _:
|
||||
return False
|
||||
|
||||
|
||||
def _web_search_block_strings(
|
||||
block_idx: int, block: Mapping[str, object], *, with_query: bool
|
||||
) -> tuple[_WebSearchString, ...]:
|
||||
"""The strings a client reads off a native web search block: the ``server_tool_use`` query and each
|
||||
``web_search_tool_result`` result's title and url."""
|
||||
match block.get("type"):
|
||||
case "server_tool_use" | "tool_use" if with_query and _is_server_tool_use(block):
|
||||
tool_input: Final = as_str_mapping(block.get("input"))
|
||||
query: Final = None if tool_input is None else tool_input.get("query")
|
||||
return (_WebSearchString(block_idx, None, "query", query),) if isinstance(query, str) and query else ()
|
||||
case "web_search_tool_result":
|
||||
return tuple(
|
||||
chain.from_iterable(
|
||||
_web_search_result_strings(block_idx, result_idx, result)
|
||||
for result_idx, result in enumerate(_web_search_results(block))
|
||||
)
|
||||
)
|
||||
case _:
|
||||
return ()
|
||||
|
||||
|
||||
def _rewritten_web_search_result(
|
||||
result_idx: int, result: object, rewrites: Mapping[tuple[int | None, str], str]
|
||||
) -> object:
|
||||
fields: Final = as_str_mapping(result)
|
||||
if fields is None:
|
||||
return result
|
||||
return {
|
||||
**fields,
|
||||
**{field: rewrites[(result_idx, field)] for field in ("title", "url") if (result_idx, field) in rewrites},
|
||||
}
|
||||
|
||||
|
||||
def _rewritten_web_search_block(
|
||||
block: Mapping[str, object], rewrites: Sequence[tuple[_WebSearchString, str]]
|
||||
) -> Mapping[str, object]:
|
||||
by_location: Final = MappingProxyType({(target.result_idx, target.field): value for target, value in rewrites})
|
||||
if _is_server_tool_use(block):
|
||||
query: Final = by_location.get((None, "query"))
|
||||
tool_input: Final = as_str_mapping(block.get("input")) or MappingProxyType({})
|
||||
return block if query is None else {**block, "input": {**tool_input, "query": query}}
|
||||
if not isinstance(block.get("content"), list):
|
||||
return block
|
||||
return {
|
||||
**block,
|
||||
"content": [
|
||||
_rewritten_web_search_result(result_idx, result, by_location)
|
||||
for result_idx, result in enumerate(_web_search_results(block))
|
||||
],
|
||||
}
|
||||
|
||||
|
||||
def _write_back_web_search_block(content_block: object, rewrites: Sequence[tuple[_WebSearchString, str]]) -> None:
|
||||
block: Final = as_str_mapping(content_block)
|
||||
if block is None or not isinstance(content_block, _WritableMessage):
|
||||
return
|
||||
rewritten: Final = _rewritten_web_search_block(block, rewrites)
|
||||
writable: Final = _as_writable(content_block)
|
||||
for key in ("input", "content"):
|
||||
if key in rewritten:
|
||||
writable[key] = rewritten[key]
|
||||
|
||||
|
||||
def _split_web_search_texts(
|
||||
guardrailed: Sequence[str], web_search_count: int, reply_count: int
|
||||
) -> tuple[Sequence[str], Sequence[str]]:
|
||||
split_at: Final = web_search_count if len(guardrailed) == web_search_count + reply_count else 0
|
||||
return guardrailed[:split_at], guardrailed[split_at:]
|
||||
|
||||
|
||||
def _web_search_rewrites_by_block(
|
||||
strings: Sequence[_WebSearchString], guardrailed: Sequence[object]
|
||||
) -> Mapping[int, tuple[tuple[_WebSearchString, str], ...]]:
|
||||
rewrites: Final = tuple(
|
||||
(target, value)
|
||||
for target, value in zip(strings, guardrailed)
|
||||
if isinstance(value, str) and value != target.value
|
||||
)
|
||||
return MappingProxyType(
|
||||
{
|
||||
block_idx: tuple(rewrite for rewrite in rewrites if rewrite[0].block_idx == block_idx)
|
||||
for block_idx in {target.block_idx for target, _ in rewrites}
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _is_client_tool_use(block: Mapping[str, object]) -> bool:
|
||||
return (
|
||||
block.get("type") == "tool_use"
|
||||
|
|
@ -289,6 +425,7 @@ def _write_back_message_text(message: _WritableMessage, target: MessageTextTarge
|
|||
|
||||
_TOOL_USE_INPUT_ADAPTER: Final = TypeAdapter(dict[str, object])
|
||||
_DECODED_SSE_DATA: Final = TypeAdapter(object)
|
||||
_REBUILT_TOOL_CALL_BLOCK_TYPES: Final = frozenset({"tool_use", "server_tool_use"})
|
||||
_RELEASED_TOOL_USE_STOP: Final = (
|
||||
b"event: message_delta\n"
|
||||
b'data: {"type": "message_delta", "delta": {"stop_reason": "tool_use", "stop_sequence": null}, '
|
||||
|
|
@ -315,6 +452,51 @@ def _write_back_tool_use(
|
|||
block["name"] = shape.name
|
||||
|
||||
|
||||
def _streamed_query_strings(
|
||||
tool_call_blocks: Sequence[tuple[int, Mapping[str, object]]], tool_calls: Sequence[_ToolCallShape]
|
||||
) -> tuple[_WebSearchString, ...]:
|
||||
if len(tool_call_blocks) != len(tool_calls):
|
||||
return ()
|
||||
return tuple(
|
||||
chain.from_iterable(
|
||||
_web_search_block_strings(
|
||||
index, {**block, "input": _rewritten_tool_use_input(tool_call.arguments)}, with_query=True
|
||||
)
|
||||
for (index, block), tool_call in zip(tool_call_blocks, tool_calls)
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def _query_rewrites(strings: Sequence[_WebSearchString], guardrailed: Sequence[object]) -> Mapping[int, str]:
|
||||
return MappingProxyType(
|
||||
{
|
||||
target.block_idx: value
|
||||
for target, value in zip(strings, guardrailed)
|
||||
if isinstance(value, str) and value != target.value
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _with_query(tool_call: _ToolCallShape, query: str) -> _ToolCallShape:
|
||||
arguments: Final = _rewritten_tool_use_input(tool_call.arguments)
|
||||
if arguments is None:
|
||||
return tool_call
|
||||
return _ToolCallShape(name=tool_call.name, arguments=json.dumps({**arguments, "query": query}))
|
||||
|
||||
|
||||
def _with_query_rewrites(
|
||||
tool_calls: tuple[_ToolCallShape, ...],
|
||||
tool_call_blocks: Sequence[tuple[int, Mapping[str, object]]],
|
||||
query_rewrites: Mapping[int, str],
|
||||
) -> tuple[_ToolCallShape, ...]:
|
||||
if not query_rewrites or len(tool_call_blocks) != len(tool_calls):
|
||||
return tool_calls
|
||||
return tuple(
|
||||
_with_query(tool_call, query_rewrites[index]) if index in query_rewrites else tool_call
|
||||
for (index, _), tool_call in zip(tool_call_blocks, tool_calls)
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _SSEFieldRewrite:
|
||||
"""One field of one nested section of a buffered SSE event, rewritten."""
|
||||
|
|
@ -402,7 +584,12 @@ class AnthropicMessagesHandler(BaseTranslation):
|
|||
def post_call_hook_response(self, response: object) -> object:
|
||||
if not isinstance(response, ModelResponse):
|
||||
return response
|
||||
return self.adapter.translate_openai_response_to_anthropic(response)
|
||||
translated: Final = self.adapter.translate_openai_response_to_anthropic(response)
|
||||
result_blocks: Final = _streamed_web_search_result_blocks(response)
|
||||
if not result_blocks:
|
||||
return translated
|
||||
content: Final = translated.get("content") or ()
|
||||
return {**translated, "content": [*result_blocks, *content]}
|
||||
|
||||
@staticmethod
|
||||
def _build_streaming_usage_response(
|
||||
|
|
@ -1273,9 +1460,17 @@ class AnthropicMessagesHandler(BaseTranslation):
|
|||
task_mappings,
|
||||
tool_calls_to_check,
|
||||
)
|
||||
block_fields: Final = tuple(map(self._output_block_fields, response_content))
|
||||
web_search_strings: Final = tuple(
|
||||
chain.from_iterable(
|
||||
_web_search_block_strings(content_idx, fields[1], with_query=True)
|
||||
for content_idx, fields in enumerate(block_fields)
|
||||
if fields is not None
|
||||
)
|
||||
)
|
||||
|
||||
# Step 2: Apply guardrail to all texts in batch
|
||||
if texts_to_check or tool_calls_to_check:
|
||||
if texts_to_check or tool_calls_to_check or web_search_strings:
|
||||
request_data = self._prepare_request_data(
|
||||
request_data,
|
||||
response,
|
||||
|
|
@ -1284,7 +1479,7 @@ class AnthropicMessagesHandler(BaseTranslation):
|
|||
)
|
||||
|
||||
inputs: Final = self._build_guardrail_inputs(
|
||||
texts_to_check,
|
||||
[*(target.value for target in web_search_strings), *texts_to_check],
|
||||
images_to_check,
|
||||
tool_calls_to_check,
|
||||
response,
|
||||
|
|
@ -1297,14 +1492,19 @@ class AnthropicMessagesHandler(BaseTranslation):
|
|||
logging_obj=litellm_logging_obj,
|
||||
)
|
||||
|
||||
guardrailed_texts: Final = guardrailed_inputs.get("texts", [])
|
||||
web_search_rewrites, reply_rewrites = _split_web_search_texts(
|
||||
guardrailed_inputs.get("texts", []), len(web_search_strings), len(texts_to_check)
|
||||
)
|
||||
|
||||
# Step 3: Map guardrail responses back to original response structure
|
||||
await self._apply_guardrail_responses_to_output(
|
||||
response=response,
|
||||
responses=guardrailed_texts,
|
||||
responses=list(reply_rewrites),
|
||||
task_mappings=task_mappings,
|
||||
)
|
||||
self._write_back_web_search_strings(
|
||||
response_content, _web_search_rewrites_by_block(web_search_strings, web_search_rewrites)
|
||||
)
|
||||
|
||||
verbose_proxy_logger.debug("Anthropic Messages: Processed output response: %s", response)
|
||||
|
||||
|
|
@ -1350,8 +1550,14 @@ class AnthropicMessagesHandler(BaseTranslation):
|
|||
string_so_far = first_choice.message.content
|
||||
pre_guardrail_tool_calls: Final = _tool_call_shapes(tool_calls_list or ())
|
||||
guardrail_inputs: Final = GenericGuardrailAPIInputs()
|
||||
if string_so_far:
|
||||
guardrail_inputs["texts"] = [string_so_far]
|
||||
leading_texts: Final = (string_so_far,) if isinstance(string_so_far, str) and string_so_far else ()
|
||||
tool_call_blocks: Final = self._streamed_tool_call_blocks(responses_so_far)
|
||||
streamed_query_strings: Final = _streamed_query_strings(tool_call_blocks, pre_guardrail_tool_calls)
|
||||
streamed_result_strings: Final = self._streamed_web_search_result_strings(responses_so_far)
|
||||
streamed_web_search_strings: Final = (*streamed_result_strings, *streamed_query_strings)
|
||||
texts_to_check: Final = [*(target.value for target in streamed_web_search_strings), *leading_texts]
|
||||
if texts_to_check:
|
||||
guardrail_inputs["texts"] = texts_to_check
|
||||
if tool_calls_list:
|
||||
guardrail_inputs["tool_calls"] = tool_calls_list
|
||||
|
||||
|
|
@ -1374,29 +1580,42 @@ class AnthropicMessagesHandler(BaseTranslation):
|
|||
responses_so_far, request_data
|
||||
)
|
||||
raise
|
||||
guardrailed_texts: Final = _guardrailed_inputs.get("texts")
|
||||
web_search_rewrites, reply_rewrites = _split_web_search_texts(
|
||||
_guardrailed_inputs.get("texts") or [], len(streamed_web_search_strings), len(leading_texts)
|
||||
)
|
||||
if (
|
||||
deliver_ended_stream_rewrites
|
||||
and isinstance(string_so_far, str)
|
||||
and string_so_far
|
||||
and guardrailed_texts
|
||||
and guardrailed_texts[0] != string_so_far
|
||||
and reply_rewrites
|
||||
and reply_rewrites[0] != string_so_far
|
||||
):
|
||||
self._write_ended_stream_text_rewrite(
|
||||
responses_so_far,
|
||||
guardrailed_texts[0],
|
||||
reply_rewrites[0],
|
||||
guardrail_name=guardrail_to_apply.guardrail_name or "unknown",
|
||||
)
|
||||
if deliver_ended_stream_rewrites:
|
||||
self._write_ended_stream_web_search_rewrites(
|
||||
responses_so_far,
|
||||
_web_search_rewrites_by_block(streamed_result_strings, web_search_rewrites),
|
||||
)
|
||||
if deliver_ended_stream_rewrites:
|
||||
returned_tool_calls: Final = _guardrailed_inputs.get("tool_calls")
|
||||
self._write_ended_stream_tool_call_rewrites(
|
||||
responses_so_far,
|
||||
pre_guardrail_tool_calls=pre_guardrail_tool_calls,
|
||||
post_guardrail_tool_calls=_tool_call_shapes(
|
||||
returned_tool_calls
|
||||
if isinstance(returned_tool_calls, list)
|
||||
and len(returned_tool_calls) == len(pre_guardrail_tool_calls)
|
||||
else tool_calls_list or ()
|
||||
post_guardrail_tool_calls=_with_query_rewrites(
|
||||
_tool_call_shapes(
|
||||
returned_tool_calls
|
||||
if isinstance(returned_tool_calls, list)
|
||||
and len(returned_tool_calls) == len(pre_guardrail_tool_calls)
|
||||
else tool_calls_list or ()
|
||||
),
|
||||
tool_call_blocks,
|
||||
_query_rewrites(
|
||||
streamed_query_strings, web_search_rewrites[len(streamed_result_strings) :]
|
||||
),
|
||||
),
|
||||
guardrail_name=guardrail_to_apply.guardrail_name or "unknown",
|
||||
)
|
||||
|
|
@ -1451,6 +1670,14 @@ class AnthropicMessagesHandler(BaseTranslation):
|
|||
request_data["litellm_metadata"] = user_metadata
|
||||
return request_data
|
||||
|
||||
@staticmethod
|
||||
def _write_back_web_search_strings(
|
||||
response_content: Sequence[object],
|
||||
rewrites_by_block: Mapping[int, tuple[tuple[_WebSearchString, str], ...]],
|
||||
) -> None:
|
||||
for block_idx, rewrites in rewrites_by_block.items():
|
||||
_write_back_web_search_block(_content_block_at(response_content, block_idx), rewrites)
|
||||
|
||||
@staticmethod
|
||||
def _get_response_content(response: object) -> Sequence[object]:
|
||||
"""Extract content list from a dict or object response."""
|
||||
|
|
@ -1573,30 +1800,22 @@ class AnthropicMessagesHandler(BaseTranslation):
|
|||
) -> None:
|
||||
"""Deliver ended-stream guardrail tool-call rewrites by rewriting the
|
||||
buffered chunks in place: the rebuilt response lists tool calls in the
|
||||
order of the stream's ``tool_use`` blocks, so the nth rewritten call lands
|
||||
on the nth block, its first ``input_json_delta`` carrying the full rewritten
|
||||
order of the stream's ``tool_use`` and ``server_tool_use`` blocks, so the nth
|
||||
rewritten call lands on the nth block, its first ``input_json_delta`` carrying the full rewritten
|
||||
arguments, every later one blanked, and ``content_block_start`` carrying the
|
||||
rewritten name. Blocks that do not line up with the rebuilt tool calls make
|
||||
the rewrite undeliverable, so the pipeline executor discards it and releases
|
||||
the original chunks."""
|
||||
if post_guardrail_tool_calls == pre_guardrail_tool_calls:
|
||||
return
|
||||
block_indices: Final = tuple(
|
||||
index
|
||||
for item in responses_so_far
|
||||
for event in cls._iter_sse_events(item)
|
||||
if event.get("type") == "content_block_start"
|
||||
and isinstance(block := event.get("content_block"), Mapping)
|
||||
and block.get("type") == "tool_use"
|
||||
and isinstance(index := event.get("index"), int)
|
||||
)
|
||||
block_indices: Final = tuple(index for index, _ in cls._streamed_tool_call_blocks(responses_so_far))
|
||||
if len(block_indices) != len(post_guardrail_tool_calls):
|
||||
from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite
|
||||
|
||||
raise UndeliverableStreamRewrite(
|
||||
guardrail_name,
|
||||
f"the guardrail returned {len(post_guardrail_tool_calls)} tool calls for a stream that carried "
|
||||
f"{len(block_indices)} tool_use blocks",
|
||||
f"{len(block_indices)} tool_use or server_tool_use blocks",
|
||||
)
|
||||
rewrites_by_block: Final = MappingProxyType(
|
||||
{
|
||||
|
|
@ -1636,6 +1855,60 @@ class AnthropicMessagesHandler(BaseTranslation):
|
|||
),
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def _streamed_web_search_result_strings(cls, responses_so_far: Sequence[object]) -> tuple[_WebSearchString, ...]:
|
||||
events: Final = chain.from_iterable(cls._iter_sse_events(item) for item in responses_so_far)
|
||||
result_blocks: Final = (
|
||||
(index, block)
|
||||
for event in events
|
||||
if event.get("type") == "content_block_start"
|
||||
and isinstance(index := event.get("index"), int)
|
||||
and (block := as_str_mapping(event.get("content_block"))) is not None
|
||||
)
|
||||
return tuple(
|
||||
chain.from_iterable(
|
||||
_web_search_block_strings(index, block, with_query=False) for index, block in result_blocks
|
||||
)
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def _write_ended_stream_web_search_rewrites(
|
||||
cls,
|
||||
responses_so_far: MutableSequence[object], # mutable-ok: rewrites the caller's buffered chunks in place
|
||||
rewrites_by_block: Mapping[int, tuple[tuple[_WebSearchString, str], ...]],
|
||||
) -> None:
|
||||
"""Deliver ended-stream guardrail rewrites of web search result titles and urls into the
|
||||
``content_block_start`` that carries each whole ``web_search_tool_result`` block."""
|
||||
if not rewrites_by_block:
|
||||
return
|
||||
|
||||
def rewrite_result_block(event: Mapping[str, object]) -> _SSEFieldRewrite | None:
|
||||
index: Final = event.get("index")
|
||||
block: Final = as_str_mapping(event.get("content_block"))
|
||||
rewrites: Final = rewrites_by_block.get(index) if isinstance(index, int) else None
|
||||
if event.get("type") != "content_block_start" or block is None or rewrites is None:
|
||||
return None
|
||||
return _SSEFieldRewrite(
|
||||
"content_block", "content", _rewritten_web_search_block(block, rewrites).get("content")
|
||||
)
|
||||
|
||||
cls._rewrite_ended_stream_events(responses_so_far, rewrite_result_block)
|
||||
|
||||
@classmethod
|
||||
def _streamed_tool_call_blocks(
|
||||
cls, responses_so_far: Sequence[object]
|
||||
) -> tuple[tuple[int, Mapping[str, object]], ...]:
|
||||
events: Final = chain.from_iterable(cls._iter_sse_events(item) for item in responses_so_far)
|
||||
return tuple(
|
||||
(index, block)
|
||||
for event in events
|
||||
if event.get("type") == "content_block_start"
|
||||
and (block := as_str_mapping(event.get("content_block"))) is not None
|
||||
and isinstance(block_type := block.get("type"), str)
|
||||
and block_type in _REBUILT_TOOL_CALL_BLOCK_TYPES
|
||||
and isinstance(index := event.get("index"), int)
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def _streamed_tool_use_ids(cls, responses_so_far: Sequence[object]) -> Mapping[int, str]:
|
||||
"""The id of each buffered ``tool_use`` block, by its content block index."""
|
||||
|
|
@ -1738,8 +2011,10 @@ class AnthropicMessagesHandler(BaseTranslation):
|
|||
def get_streaming_scan_key(self, responses_so_far: Sequence[object]) -> StreamingScanKey | None:
|
||||
stream_ended: Final = self._check_streaming_has_ended(responses_so_far)
|
||||
tool_use_fingerprints: Final = self._streamed_tool_use_fingerprints(responses_so_far)
|
||||
streamed_text: Final = self.get_streaming_string_so_far(responses_so_far)
|
||||
result_strings: Final = self._streamed_web_search_result_strings(responses_so_far) if stream_ended else ()
|
||||
return StreamingScanKey(
|
||||
texts=(self.get_streaming_string_so_far(responses_so_far),),
|
||||
texts=(*(target.value for target in result_strings), streamed_text),
|
||||
tool_calls=tool_use_fingerprints if stream_ended else (),
|
||||
stream_ended=stream_ended,
|
||||
tool_calls_in_flight=bool(tool_use_fingerprints) and not stream_ended,
|
||||
|
|
|
|||
|
|
@ -8,12 +8,13 @@ import re
|
|||
from collections.abc import Container, Mapping, MutableMapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime, timezone
|
||||
from itertools import chain
|
||||
from types import MappingProxyType
|
||||
from typing import Annotated, Any, ClassVar, Final, Literal, TypeVar
|
||||
from urllib.parse import quote, urlparse
|
||||
|
||||
import httpx
|
||||
from pydantic import BeforeValidator, ConfigDict, Field, StrictBool, TypeAdapter, ValidationError
|
||||
from pydantic import BeforeValidator, ConfigDict, Field, JsonValue, StrictBool, TypeAdapter, ValidationError
|
||||
from typing_extensions import ReadOnly, TypedDict
|
||||
|
||||
import litellm
|
||||
|
|
@ -24,6 +25,7 @@ from litellm.constants import (
|
|||
DEFAULT_REASONING_EFFORT_XHIGH_THINKING_BUDGET,
|
||||
)
|
||||
from litellm.exceptions import UnsupportedParamsError
|
||||
from litellm.litellm_core_utils.core_helpers import as_str_mapping
|
||||
from litellm.litellm_core_utils.prompt_templates.common_utils import (
|
||||
get_file_ids_from_messages,
|
||||
is_encrypted_reasoning_block,
|
||||
|
|
@ -46,6 +48,7 @@ from litellm.llms.anthropic.wif import (
|
|||
from litellm.llms.base_llm.base_utils import BaseLLMModelInfo, BaseTokenCounter
|
||||
from litellm.llms.base_llm.chat.transformation import BaseLLMException
|
||||
from litellm.types.llms.anthropic import (
|
||||
ANTHROPIC_BETA_HEADER_VALUES,
|
||||
ANTHROPIC_HOSTED_TOOLS,
|
||||
ANTHROPIC_MID_CONVERSATION_OUTPUT_CONFIG_BETA_HEADER,
|
||||
ANTHROPIC_MID_CONVERSATION_TOOL_CHANGES_BETA_HEADER,
|
||||
|
|
@ -118,6 +121,7 @@ _CLAUDE_CODE_BILLING_HEADER_PREFIX: Final = "x-anthropic-billing-header:"
|
|||
OBJECT_MAPPING_ADAPTER: Final = TypeAdapter(dict[object, object])
|
||||
OBJECT_LIST_ADAPTER: Final = TypeAdapter(list[object])
|
||||
_DECODED_JSON: Final = TypeAdapter(object)
|
||||
_JSON_VALUE: Final = TypeAdapter(JsonValue)
|
||||
|
||||
|
||||
_CLAUDE_CODE_USER_AGENT_PREFIXES: Final = ("claude-cli/", "claude-code/")
|
||||
|
|
@ -177,24 +181,68 @@ def is_anthropic_messages_url(url: str) -> bool:
|
|||
return parsed_url.hostname == "api.anthropic.com" or parsed_url.path.removesuffix("/").endswith("/v1/messages")
|
||||
|
||||
|
||||
def _message_blocks(message: object) -> Sequence[object]:
|
||||
fields: Final = as_str_mapping(message)
|
||||
blocks: Final = validated_object_list(fields.get("content")) if fields is not None else None
|
||||
return blocks or ()
|
||||
|
||||
|
||||
def _is_signed_compaction_block(block: object) -> bool:
|
||||
fields: Final = as_str_mapping(block)
|
||||
if fields is None or fields.get("type") != "compaction":
|
||||
return False
|
||||
signature: Final = fields.get("signature")
|
||||
return isinstance(signature, str) and bool(signature)
|
||||
|
||||
|
||||
def carries_signed_compaction_block(messages: Sequence[object]) -> bool:
|
||||
return any(map(_is_signed_compaction_block, chain.from_iterable(map(_message_blocks, messages))))
|
||||
|
||||
|
||||
def requires_native_compaction_beta(
|
||||
custom_llm_provider: str,
|
||||
optional_params: Mapping[str, object],
|
||||
messages: Sequence[object],
|
||||
) -> bool:
|
||||
return custom_llm_provider == "anthropic" and (
|
||||
optional_params.get("compaction") is not None
|
||||
or any(
|
||||
isinstance(block, Mapping)
|
||||
and block.get("type") == "compaction"
|
||||
and isinstance(block.get("signature"), str)
|
||||
and bool(block.get("signature"))
|
||||
for message in messages
|
||||
if isinstance(message, Mapping)
|
||||
for content in (message.get("content"),)
|
||||
if isinstance(content, (list, tuple))
|
||||
for block in content
|
||||
optional_params.get("compaction") is not None or carries_signed_compaction_block(messages)
|
||||
)
|
||||
|
||||
|
||||
def anthropic_context_management(context_management: JsonValue) -> JsonValue:
|
||||
if not isinstance(context_management, list):
|
||||
return context_management
|
||||
from litellm.llms.anthropic.chat.transformation import AnthropicConfig
|
||||
|
||||
edits: Final = AnthropicConfig.map_openai_context_management_to_anthropic(
|
||||
[entry for entry in context_management if isinstance(entry, dict)]
|
||||
)
|
||||
return context_management if edits is None else _JSON_VALUE.validate_python(edits)
|
||||
|
||||
|
||||
def _context_management_edits(spec: JsonValue) -> Sequence[object]:
|
||||
mapping: Final = as_str_mapping(spec)
|
||||
return (validated_object_list(mapping.get("edits")) if mapping is not None else None) or ()
|
||||
|
||||
|
||||
def context_management_with_edits(context_management: JsonValue) -> JsonValue:
|
||||
spec: Final = anthropic_context_management(context_management)
|
||||
return spec if _context_management_edits(spec) else None
|
||||
|
||||
|
||||
def context_management_betas(context_management: JsonValue) -> tuple[str, ...]:
|
||||
edits: Final = _context_management_edits(anthropic_context_management(context_management))
|
||||
edit_types: Final = tuple(edit.get("type") for edit in map(as_str_mapping, edits) if edit is not None)
|
||||
return tuple(
|
||||
beta
|
||||
for beta, used in (
|
||||
(ANTHROPIC_BETA_HEADER_VALUES.COMPACT_2026_01_12.value, "compact_20260112" in edit_types),
|
||||
(
|
||||
ANTHROPIC_BETA_HEADER_VALUES.CONTEXT_MANAGEMENT_2025_06_27.value,
|
||||
any(edit_type != "compact_20260112" for edit_type in edit_types),
|
||||
),
|
||||
)
|
||||
if used
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -1367,8 +1415,6 @@ class AnthropicModelInfo(BaseLLMModelInfo):
|
|||
if is_vertex_request is True:
|
||||
# Vertex AI requires web search beta header for web search to work
|
||||
if web_search_tool_used:
|
||||
from litellm.types.llms.anthropic import ANTHROPIC_BETA_HEADER_VALUES
|
||||
|
||||
headers["anthropic-beta"] = ANTHROPIC_BETA_HEADER_VALUES.WEB_SEARCH_2025_03_05.value
|
||||
elif len(all_betas) > 0:
|
||||
headers["anthropic-beta"] = ",".join(all_betas)
|
||||
|
|
|
|||
|
|
@ -39,6 +39,7 @@ class AnthropicCountTokensHandler(AnthropicCountTokensConfig):
|
|||
tools: list[dict[str, JsonValue]] | None = None,
|
||||
system: JsonValue = None,
|
||||
optional_params: Mapping[str, JsonValue] | None = None,
|
||||
context_management: JsonValue = None,
|
||||
drop_params: bool = False,
|
||||
) -> dict[str, JsonValue]:
|
||||
"""
|
||||
|
|
@ -70,6 +71,7 @@ class AnthropicCountTokensHandler(AnthropicCountTokensConfig):
|
|||
tools=tools,
|
||||
system=system,
|
||||
optional_params=optional_params,
|
||||
context_management=context_management,
|
||||
drop_params=drop_params,
|
||||
)
|
||||
|
||||
|
|
@ -81,7 +83,7 @@ class AnthropicCountTokensHandler(AnthropicCountTokensConfig):
|
|||
verbose_logger.debug("Making request to: %s", endpoint_url)
|
||||
|
||||
# Get required headers
|
||||
headers: Final = self.get_count_tokens_headers(auth_header)
|
||||
headers: Final = self.get_count_tokens_headers(auth_header, messages, context_management)
|
||||
|
||||
# Use LiteLLM's async httpx client
|
||||
async_client: Final = get_async_httpx_client(llm_provider=litellm.LlmProviders.ANTHROPIC)
|
||||
|
|
|
|||
|
|
@ -132,6 +132,7 @@ class AnthropicTokenCounter(BaseTokenCounter):
|
|||
tools=tools,
|
||||
system=system,
|
||||
optional_params=count_options,
|
||||
context_management=context_management.spec if context_management is not None else None,
|
||||
drop_params=drops_unsupported_params(litellm_params.get("drop_params")),
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -27,11 +27,15 @@ from litellm.litellm_core_utils.prompt_templates.mid_conversation_system import
|
|||
)
|
||||
from litellm.litellm_core_utils.url_utils import join_api_base
|
||||
from litellm.llms.anthropic.common_utils import (
|
||||
carries_signed_compaction_block,
|
||||
context_management_betas,
|
||||
context_management_with_edits,
|
||||
is_claude_model_target,
|
||||
merge_anthropic_beta_headers,
|
||||
tools_without_custom_field,
|
||||
)
|
||||
from litellm.llms.anthropic.wif import resolve_anthropic_base
|
||||
from litellm.types.llms.anthropic import ANTHROPIC_BETA_HEADER_VALUES
|
||||
from litellm.types.llms.openai import AllMessageValues, ChatCompletionImageObject, ValidUserMessageContentTypes
|
||||
from litellm.types.utils import TokenCountMessagesFormat, TokenCountResponse
|
||||
|
||||
|
|
@ -376,6 +380,7 @@ class AnthropicCountTokensConfig:
|
|||
tools: list[dict[str, JsonValue]] | None = None,
|
||||
system: JsonValue = None,
|
||||
optional_params: Mapping[str, JsonValue] | None = None,
|
||||
context_management: JsonValue = None,
|
||||
drop_params: bool = False,
|
||||
) -> dict[str, JsonValue]: # mutable-ok: provider transport requires JSON dictionaries
|
||||
"""
|
||||
|
|
@ -408,6 +413,7 @@ class AnthropicCountTokensConfig:
|
|||
drop_custom=drop_params,
|
||||
),
|
||||
),
|
||||
("context_management", context_management_with_edits(context_management)),
|
||||
)
|
||||
if value is not None
|
||||
}
|
||||
|
|
@ -419,18 +425,33 @@ class AnthropicCountTokensConfig:
|
|||
)
|
||||
)
|
||||
|
||||
def get_count_tokens_headers(self, auth_header: Mapping[str, str]) -> dict[str, str]:
|
||||
def get_count_tokens_headers(
|
||||
self,
|
||||
auth_header: Mapping[str, str],
|
||||
messages: Sequence[Mapping[str, JsonValue]],
|
||||
context_management: JsonValue = None,
|
||||
) -> dict[str, str]:
|
||||
"""The count-tokens headers around a resolved Anthropic auth header
|
||||
(``AnthropicModelInfo.get_auth_header``): x-api-key for a static key, an Authorization
|
||||
bearer for ``ANTHROPIC_AUTH_TOKEN`` and for sk-ant-oat tokens, whose mandatory oauth beta
|
||||
merges with the token-counting beta instead of replacing it."""
|
||||
merges with the token-counting beta instead of replacing it. A history that carries a
|
||||
signed compaction block also needs the compaction beta, or Anthropic rejects the block
|
||||
as an unknown content type, and ``context_management`` edits need the betas inference sends."""
|
||||
compaction_betas: Final = (
|
||||
(ANTHROPIC_BETA_HEADER_VALUES.COMPACT_2026_09_04.value,)
|
||||
if carries_signed_compaction_block(messages)
|
||||
else ()
|
||||
)
|
||||
request_betas: Final = (
|
||||
ANTHROPIC_TOKEN_COUNTING_BETA_VERSION,
|
||||
*compaction_betas,
|
||||
*context_management_betas(context_management),
|
||||
)
|
||||
return {
|
||||
"Content-Type": "application/json",
|
||||
"anthropic-version": "2023-06-01",
|
||||
**auth_header,
|
||||
"anthropic-beta": merge_anthropic_beta_headers(
|
||||
auth_header.get("anthropic-beta"), ANTHROPIC_TOKEN_COUNTING_BETA_VERSION
|
||||
),
|
||||
"anthropic-beta": merge_anthropic_beta_headers(auth_header.get("anthropic-beta"), request_betas),
|
||||
}
|
||||
|
||||
def validate_request(
|
||||
|
|
|
|||
|
|
@ -647,12 +647,14 @@ class LiteLLMMessagesToCompletionTransformationHandler:
|
|||
request_data["output_config"] = extra_kwargs["output_config"]
|
||||
|
||||
custom_llm_provider: Final = extra_kwargs.get("custom_llm_provider")
|
||||
api_base: Final = extra_kwargs.get("api_base")
|
||||
(
|
||||
openai_request,
|
||||
tool_name_mapping,
|
||||
) = ANTHROPIC_ADAPTER.translate_completion_input_params_with_tool_mapping(
|
||||
request_data,
|
||||
custom_llm_provider=custom_llm_provider if isinstance(custom_llm_provider, str) else None,
|
||||
api_base=api_base if isinstance(api_base, str) else None,
|
||||
drop_params=drops_unsupported_params(extra_kwargs.get("drop_params")),
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -24,6 +24,7 @@ OPENAI_MAX_TOOL_NAME_LENGTH: Final = 64
|
|||
TOOL_NAME_HASH_LENGTH: Final = 8
|
||||
TOOL_NAME_PREFIX_LENGTH: Final = OPENAI_MAX_TOOL_NAME_LENGTH - TOOL_NAME_HASH_LENGTH - 1 # 55
|
||||
PROVIDERS_PROXYING_AN_UNKNOWN_BACKEND: Final = frozenset({"litellm_proxy"})
|
||||
PROVIDERS_SERVED_BY_THE_MISTRAL_CHAT_CONFIG: Final = frozenset({"mistral", "codestral"})
|
||||
_COMPACTION_BLOCK: Final = TypeAdapter(dict[str, JsonValue])
|
||||
|
||||
|
||||
|
|
@ -255,7 +256,12 @@ class AnthropicAdapter:
|
|||
return result
|
||||
|
||||
def translate_completion_input_params_with_tool_mapping(
|
||||
self, kwargs, *, custom_llm_provider: str | None = None, drop_params: bool = True
|
||||
self,
|
||||
kwargs,
|
||||
*,
|
||||
custom_llm_provider: str | None = None,
|
||||
api_base: str | None = None,
|
||||
drop_params: bool = True,
|
||||
) -> tuple[ChatCompletionRequest | None, dict[str, str]]:
|
||||
"""
|
||||
Translate Anthropic request params to OpenAI format, returning tool name mapping.
|
||||
|
|
@ -289,6 +295,7 @@ class AnthropicAdapter:
|
|||
) = LiteLLMAnthropicMessagesAdapter().translate_anthropic_to_openai(
|
||||
anthropic_message_request=request_body,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
api_base=api_base,
|
||||
drop_params=drop_params,
|
||||
)
|
||||
|
||||
|
|
@ -1001,7 +1008,9 @@ class LiteLLMAnthropicMessagesAdapter:
|
|||
)
|
||||
|
||||
@staticmethod
|
||||
def _supports_prompt_cache_key(model: str | None, custom_llm_provider: str | None) -> bool:
|
||||
def _supports_prompt_cache_key(
|
||||
prompt_cache_key: str, model: str | None, custom_llm_provider: str | None, api_base: str | None
|
||||
) -> bool:
|
||||
if not model or not custom_llm_provider:
|
||||
return False
|
||||
if custom_llm_provider in PROVIDERS_PROXYING_AN_UNKNOWN_BACKEND:
|
||||
|
|
@ -1009,7 +1018,17 @@ class LiteLLMAnthropicMessagesAdapter:
|
|||
supported_params: Final = litellm.get_supported_openai_params(
|
||||
model=model, custom_llm_provider=custom_llm_provider
|
||||
)
|
||||
return "prompt_cache_key" in (supported_params or ())
|
||||
if "prompt_cache_key" not in (supported_params or ()):
|
||||
return False
|
||||
if custom_llm_provider not in PROVIDERS_SERVED_BY_THE_MISTRAL_CHAT_CONFIG:
|
||||
return True
|
||||
_, _, _, resolved_api_base = litellm.get_llm_provider(
|
||||
model=model, custom_llm_provider=custom_llm_provider, api_base=api_base
|
||||
)
|
||||
mistral_config: Final = litellm.MistralConfig()
|
||||
return mistral_config.takes_prompt_cache_key_value(
|
||||
prompt_cache_key
|
||||
) and mistral_config.accepts_prompt_cache_key_at(resolved_api_base)
|
||||
|
||||
@staticmethod
|
||||
def _target_declares_reasoning_effort(model: str, custom_llm_provider: str | None) -> bool:
|
||||
|
|
@ -1042,6 +1061,7 @@ class LiteLLMAnthropicMessagesAdapter:
|
|||
new_kwargs: ChatCompletionRequest,
|
||||
*,
|
||||
custom_llm_provider: str | None = None,
|
||||
api_base: str | None = None,
|
||||
) -> None:
|
||||
"""Translate metadata fields from Anthropic request to OpenAI request."""
|
||||
if "metadata" in anthropic_message_request:
|
||||
|
|
@ -1050,7 +1070,7 @@ class LiteLLMAnthropicMessagesAdapter:
|
|||
new_kwargs["user"] = metadata["user_id"]
|
||||
prompt_cache_key: Final = prompt_cache_key_from_user_id(metadata["user_id"])
|
||||
if prompt_cache_key is not None and self._supports_prompt_cache_key(
|
||||
anthropic_message_request.get("model"), custom_llm_provider
|
||||
prompt_cache_key, anthropic_message_request.get("model"), custom_llm_provider, api_base
|
||||
):
|
||||
new_kwargs["prompt_cache_key"] = prompt_cache_key
|
||||
|
||||
|
|
@ -1244,6 +1264,7 @@ class LiteLLMAnthropicMessagesAdapter:
|
|||
anthropic_message_request: AnthropicMessagesRequest,
|
||||
*,
|
||||
custom_llm_provider: str | None = None,
|
||||
api_base: str | None = None,
|
||||
preserve_midturn_system: bool = False,
|
||||
drop_params: bool = True,
|
||||
) -> tuple[ChatCompletionRequest, dict[str, str]]:
|
||||
|
|
@ -1282,6 +1303,7 @@ class LiteLLMAnthropicMessagesAdapter:
|
|||
anthropic_message_request=anthropic_message_request,
|
||||
new_kwargs=new_kwargs,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
api_base=api_base,
|
||||
)
|
||||
## CONVERT TOOL CHOICE
|
||||
self._translate_tool_choice_to_openai(
|
||||
|
|
|
|||
|
|
@ -12,11 +12,20 @@ import json
|
|||
from collections.abc import Mapping, Sequence
|
||||
from typing import Final, cast
|
||||
|
||||
from pydantic import TypeAdapter
|
||||
|
||||
from litellm.llms.anthropic.messages.native.utils import anthropic_message_delta
|
||||
from litellm.types.llms.anthropic_messages.anthropic_response import (
|
||||
AnthropicMessagesResponse,
|
||||
)
|
||||
|
||||
_CITATIONS: Final = TypeAdapter(tuple[object, ...])
|
||||
|
||||
|
||||
def _citations_of(block: Mapping[str, object]) -> tuple[object, ...]:
|
||||
citations: Final = block.get("citations")
|
||||
return _CITATIONS.validate_python(citations) if isinstance(citations, list) else ()
|
||||
|
||||
|
||||
class FakeAnthropicMessagesStreamIterator:
|
||||
"""
|
||||
|
|
@ -47,16 +56,30 @@ class FakeAnthropicMessagesStreamIterator:
|
|||
|
||||
def _create_content_block_chunks(self, block_dict: Mapping[str, object], index: int) -> list[bytes]:
|
||||
"""Build SSE chunks for a single content block."""
|
||||
chunks: Final = []
|
||||
chunks: Final = list[bytes]()
|
||||
block_type: Final = block_dict.get("type")
|
||||
|
||||
if block_type == "text":
|
||||
citations: Final = _citations_of(block_dict)
|
||||
content_block_start = {
|
||||
"type": "content_block_start",
|
||||
"index": index,
|
||||
"content_block": {"type": "text", "text": ""},
|
||||
"content_block": {"type": "text", "text": "", "citations": ()}
|
||||
if citations
|
||||
else {"type": "text", "text": ""},
|
||||
}
|
||||
chunks.append(f"event: content_block_start\ndata: {json.dumps(content_block_start)}\n\n".encode())
|
||||
citation_deltas: Final = (
|
||||
{
|
||||
"type": "content_block_delta",
|
||||
"index": index,
|
||||
"delta": {"type": "citations_delta", "citation": citation},
|
||||
}
|
||||
for citation in citations
|
||||
)
|
||||
chunks.extend(
|
||||
f"event: content_block_delta\ndata: {json.dumps(delta)}\n\n".encode() for delta in citation_deltas
|
||||
)
|
||||
text: Final = block_dict.get("text", "")
|
||||
content_block_delta = {
|
||||
"type": "content_block_delta",
|
||||
|
|
@ -89,36 +112,23 @@ class FakeAnthropicMessagesStreamIterator:
|
|||
}
|
||||
chunks.append(f"event: content_block_delta\ndata: {json.dumps(signature_delta)}\n\n".encode())
|
||||
|
||||
elif block_type == "redacted_thinking":
|
||||
content_block_start = {
|
||||
elif block_type in ("tool_use", "server_tool_use"):
|
||||
empty_input: Final[Mapping[str, object]] = {}
|
||||
tool_use_start: Final = {
|
||||
"type": "content_block_start",
|
||||
"index": index,
|
||||
"content_block": {"type": "redacted_thinking"},
|
||||
"content_block": {**block_dict, "input": empty_input},
|
||||
}
|
||||
chunks.append(f"event: content_block_start\ndata: {json.dumps(content_block_start)}\n\n".encode())
|
||||
|
||||
elif block_type == "tool_use":
|
||||
content_block_start = {
|
||||
"type": "content_block_start",
|
||||
"index": index,
|
||||
"content_block": {
|
||||
"type": "tool_use",
|
||||
"id": block_dict.get("id"),
|
||||
"name": block_dict.get("name"),
|
||||
"input": {},
|
||||
},
|
||||
}
|
||||
chunks.append(f"event: content_block_start\ndata: {json.dumps(content_block_start)}\n\n".encode())
|
||||
input_data: Final = block_dict.get("input", {})
|
||||
content_block_delta = {
|
||||
chunks.append(f"event: content_block_start\ndata: {json.dumps(tool_use_start)}\n\n".encode())
|
||||
tool_use_delta: Final = {
|
||||
"type": "content_block_delta",
|
||||
"index": index,
|
||||
"delta": {
|
||||
"type": "input_json_delta",
|
||||
"partial_json": json.dumps(input_data),
|
||||
"partial_json": json.dumps(block_dict.get("input", {})),
|
||||
},
|
||||
}
|
||||
chunks.append(f"event: content_block_delta\ndata: {json.dumps(content_block_delta)}\n\n".encode())
|
||||
chunks.append(f"event: content_block_delta\ndata: {json.dumps(tool_use_delta)}\n\n".encode())
|
||||
|
||||
else:
|
||||
passthrough_start: Final = {
|
||||
|
|
|
|||
|
|
@ -43,6 +43,7 @@ from ...common_utils import (
|
|||
OBJECT_LIST_ADAPTER,
|
||||
AnthropicError,
|
||||
AnthropicModelInfo,
|
||||
context_management_betas,
|
||||
is_claude_model_target,
|
||||
merge_anthropic_beta_headers,
|
||||
optionally_handle_anthropic_oauth,
|
||||
|
|
@ -745,6 +746,8 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig):
|
|||
)
|
||||
if transformed_context_management is not None:
|
||||
anthropic_messages_optional_request_params["context_management"] = transformed_context_management
|
||||
elif drops_unsupported_params(litellm_params.drop_params):
|
||||
anthropic_messages_optional_request_params.pop("context_management")
|
||||
|
||||
####### get required params for all anthropic messages requests ######
|
||||
# Lazy %s: the f-string previously stringified the entire messages
|
||||
|
|
@ -873,28 +876,7 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig):
|
|||
if requires_native_compaction_beta(custom_llm_provider, optional_params, messages):
|
||||
beta_values.add(ANTHROPIC_BETA_HEADER_VALUES.COMPACT_2026_09_04.value)
|
||||
|
||||
# Check for context management
|
||||
context_management_param: Final = optional_params.get("context_management")
|
||||
if context_management_param is not None:
|
||||
# Check edits array for compact_20260112 type
|
||||
edits: Final = context_management_param.get("edits", ())
|
||||
has_compact = False
|
||||
has_other = False
|
||||
|
||||
for edit in edits:
|
||||
edit_type = edit.get("type", "")
|
||||
if edit_type == "compact_20260112":
|
||||
has_compact = True
|
||||
else:
|
||||
has_other = True
|
||||
|
||||
# Add compact header if any compact edits exist
|
||||
if has_compact:
|
||||
beta_values.add(ANTHROPIC_BETA_HEADER_VALUES.COMPACT_2026_01_12.value)
|
||||
|
||||
# Add context management header if any other edits exist
|
||||
if has_other:
|
||||
beta_values.add(ANTHROPIC_BETA_HEADER_VALUES.CONTEXT_MANAGEMENT_2025_06_27.value)
|
||||
beta_values.update(context_management_betas(optional_params.get("context_management")))
|
||||
|
||||
# Check for structured outputs. Anthropic's newer request shape nests
|
||||
# the schema under output_config.format; the older top-level
|
||||
|
|
|
|||
|
|
@ -42,6 +42,7 @@ from ..openai.common_utils import (
|
|||
build_output_token_limit_response,
|
||||
is_output_token_limit_error,
|
||||
model_dump_provider_response,
|
||||
upstream_error_response,
|
||||
)
|
||||
from .common_utils import (
|
||||
AzureOpenAIError,
|
||||
|
|
@ -410,6 +411,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM):
|
|||
raise AzureOpenAIError(
|
||||
status_code=status_code,
|
||||
message=str(e),
|
||||
response=upstream_error_response(e),
|
||||
headers=error_headers,
|
||||
body=error_body,
|
||||
)
|
||||
|
|
@ -673,6 +675,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM):
|
|||
raise AzureOpenAIError(
|
||||
status_code=status_code,
|
||||
message=message,
|
||||
response=upstream_error_response(e),
|
||||
headers=error_headers,
|
||||
body=error_body,
|
||||
)
|
||||
|
|
@ -866,7 +869,12 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM):
|
|||
if error_headers is None and error_response:
|
||||
error_headers = getattr(error_response, "headers", None)
|
||||
error_text = error_response.text
|
||||
raise AzureOpenAIError(status_code=status_code, message=error_text, headers=error_headers)
|
||||
raise AzureOpenAIError(
|
||||
status_code=status_code,
|
||||
message=error_text,
|
||||
response=upstream_error_response(e),
|
||||
headers=error_headers,
|
||||
)
|
||||
|
||||
async def make_async_azure_httpx_request(
|
||||
self,
|
||||
|
|
@ -1197,7 +1205,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM):
|
|||
except Exception as e:
|
||||
error_code: Final[object] = getattr(e, "status_code", None)
|
||||
if error_code is not None:
|
||||
raise AzureOpenAIError(status_code=error_code, message=str(e))
|
||||
raise AzureOpenAIError(status_code=error_code, message=str(e), response=upstream_error_response(e))
|
||||
else:
|
||||
raise AzureOpenAIError(status_code=500, message=str(e))
|
||||
|
||||
|
|
|
|||
|
|
@ -8,6 +8,7 @@ from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
|
|||
from litellm.litellm_core_utils.prompt_templates.factory import prompt_factory
|
||||
from litellm.utils import CustomStreamWrapper, ModelResponse, TextCompletionResponse
|
||||
|
||||
from ...openai.common_utils import upstream_error_response
|
||||
from ...openai.completion.transformation import OpenAITextCompletionConfig
|
||||
from ..common_utils import AzureOpenAIError, BaseAzureLLM
|
||||
|
||||
|
|
@ -183,7 +184,12 @@ class AzureTextCompletion(BaseAzureLLM):
|
|||
error_response: Final[object] = getattr(e, "response", None)
|
||||
if error_headers is None and error_response:
|
||||
error_headers = getattr(error_response, "headers", None)
|
||||
raise AzureOpenAIError(status_code=status_code, message=str(e), headers=error_headers)
|
||||
raise AzureOpenAIError(
|
||||
status_code=status_code,
|
||||
message=str(e),
|
||||
headers=error_headers,
|
||||
response=upstream_error_response(e),
|
||||
)
|
||||
|
||||
async def acompletion(
|
||||
self,
|
||||
|
|
@ -244,7 +250,12 @@ class AzureTextCompletion(BaseAzureLLM):
|
|||
error_response: Final[object] = getattr(e, "response", None)
|
||||
if error_headers is None and error_response:
|
||||
error_headers = getattr(error_response, "headers", None)
|
||||
raise AzureOpenAIError(status_code=status_code, message=str(e), headers=error_headers)
|
||||
raise AzureOpenAIError(
|
||||
status_code=status_code,
|
||||
message=str(e),
|
||||
headers=error_headers,
|
||||
response=upstream_error_response(e),
|
||||
)
|
||||
|
||||
def streaming(
|
||||
self,
|
||||
|
|
@ -355,4 +366,9 @@ class AzureTextCompletion(BaseAzureLLM):
|
|||
error_response: Final[object] = getattr(e, "response", None)
|
||||
if error_headers is None and error_response:
|
||||
error_headers = getattr(error_response, "headers", None)
|
||||
raise AzureOpenAIError(status_code=status_code, message=str(e), headers=error_headers)
|
||||
raise AzureOpenAIError(
|
||||
status_code=status_code,
|
||||
message=str(e),
|
||||
headers=error_headers,
|
||||
response=upstream_error_response(e),
|
||||
)
|
||||
|
|
|
|||
|
|
@ -1,5 +1,7 @@
|
|||
from typing import Final
|
||||
|
||||
import httpx
|
||||
|
||||
from litellm.exceptions import ContentPolicyViolationError
|
||||
|
||||
|
||||
|
|
@ -14,6 +16,7 @@ class AzureOpenAIExceptionMapping:
|
|||
model: str,
|
||||
extra_information: str,
|
||||
original_exception: Exception,
|
||||
response: httpx.Response | None,
|
||||
) -> ContentPolicyViolationError:
|
||||
"""
|
||||
Create a content policy violation error
|
||||
|
|
@ -39,7 +42,7 @@ class AzureOpenAIExceptionMapping:
|
|||
llm_provider="azure",
|
||||
model=model,
|
||||
litellm_debug_info=extra_information,
|
||||
response=getattr(original_exception, "response", None),
|
||||
response=response,
|
||||
provider_specific_fields={
|
||||
# Preserve legacy key for backward compatibility.
|
||||
"innererror": inner_error,
|
||||
|
|
|
|||
|
|
@ -6,15 +6,19 @@ from typing import TYPE_CHECKING, Final, cast
|
|||
|
||||
import httpx
|
||||
from httpx import Response
|
||||
from pydantic import TypeAdapter, ValidationError
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.integrations.anthropic_cache_control_hook import AnthropicCacheControlHook
|
||||
from litellm.litellm_core_utils.prompt_templates.common_utils import (
|
||||
content_parts,
|
||||
convert_content_list_to_str,
|
||||
drop_tool_reference_parts_from_tool_messages,
|
||||
filter_value_from_dict,
|
||||
is_pdf_file_part,
|
||||
prompt_cache_breakpoint_for_wire,
|
||||
with_prompt_cache_breakpoint,
|
||||
)
|
||||
from litellm.llms.azure.common_utils import BaseAzureLLM
|
||||
from litellm.llms.azure_ai.common_utils import (
|
||||
|
|
@ -27,7 +31,7 @@ from litellm.llms.openai.common_utils import drop_params_from_unprocessable_enti
|
|||
from litellm.llms.openai.openai import OpenAIChatCompletionResponseIterator, OpenAIConfig
|
||||
from litellm.llms.xai.chat.transformation import XAIChatConfig
|
||||
from litellm.secret_managers.main import get_secret_str
|
||||
from litellm.types.llms.openai import AllMessageValues
|
||||
from litellm.types.llms.openai import AllMessageValues, PromptCacheBreakpoint
|
||||
from litellm.types.router import GenericLiteLLMParams
|
||||
from litellm.types.utils import ModelResponse, ModelResponseStream, ProviderField
|
||||
from litellm.utils import add_path_to_api_base, supports_tool_choice
|
||||
|
|
@ -62,6 +66,72 @@ def _has_list_only_content_part(message: Mapping[str, object]) -> bool:
|
|||
return any(_is_list_only_content_part(part) for part in content_parts(message))
|
||||
|
||||
|
||||
_CONTENT_PARTS_ADAPTER: Final = TypeAdapter(list[dict[str, object]])
|
||||
|
||||
|
||||
def _text_parts(message: Mapping[str, object]) -> tuple[Mapping[str, object], ...]:
|
||||
content: Final = message.get("content")
|
||||
if not isinstance(content, list):
|
||||
return ()
|
||||
try:
|
||||
parts: Final = _CONTENT_PARTS_ADAPTER.validate_python(content)
|
||||
except ValidationError:
|
||||
return ()
|
||||
return tuple(part for part in parts if part.get("type") == "text")
|
||||
|
||||
|
||||
def _with_only_text_parts(message: AllMessageValues) -> AllMessageValues:
|
||||
if _has_list_only_content_part(message) or not (text_parts := _text_parts(message)):
|
||||
return message
|
||||
text_only: Final = {**message, "content": list(text_parts)}
|
||||
return cast(AllMessageValues, text_only) # cast-ok: the same message, text parts only
|
||||
|
||||
|
||||
_FIELDS: Final = TypeAdapter(Mapping[str, object])
|
||||
|
||||
|
||||
def _foundry_breakpoint(marker: object) -> PromptCacheBreakpoint | None:
|
||||
if prompt_cache_breakpoint_for_wire(marker, drop_params=True) is None:
|
||||
return None
|
||||
return PromptCacheBreakpoint(mode="explicit")
|
||||
|
||||
|
||||
def _part_with_foundry_breakpoint(part: object) -> object:
|
||||
if not isinstance(part, Mapping):
|
||||
return part
|
||||
fields: Final = _FIELDS.validate_python(part)
|
||||
if "prompt_cache_breakpoint" not in fields:
|
||||
return fields
|
||||
unmarked: Final = {key: value for key, value in fields.items() if key != "prompt_cache_breakpoint"}
|
||||
return with_prompt_cache_breakpoint(unmarked, _foundry_breakpoint(fields["prompt_cache_breakpoint"]))
|
||||
|
||||
|
||||
def _with_breakpoints_for_foundry(message: AllMessageValues, drop_params: bool) -> AllMessageValues:
|
||||
renamed: Final = AnthropicCacheControlHook.cache_control_as_prompt_cache_breakpoints(_with_only_text_parts(message))
|
||||
parts: Final = content_parts(renamed) if drop_params else ()
|
||||
if not parts:
|
||||
return renamed
|
||||
foundry_message: Final = {**renamed, "content": [_part_with_foundry_breakpoint(part) for part in parts]}
|
||||
return cast(AllMessageValues, foundry_message) # cast-ok: the same message, its breakpoints in Foundry's shape
|
||||
|
||||
|
||||
def _deployment_reads_breakpoints(model: str, litellm_params: Mapping[str, object]) -> bool:
|
||||
base_model: Final = litellm_params.get("base_model")
|
||||
named_models: Final = (base_model, model) if isinstance(base_model, str) and base_model else (model,)
|
||||
model_info: Final = litellm_params.get("model_info")
|
||||
if not isinstance(model_info, Mapping):
|
||||
return any(
|
||||
AnthropicCacheControlHook.targets_openai_prompt_cache_breakpoint(name, "azure_ai") for name in named_models
|
||||
)
|
||||
deployment_flag: Final = _FIELDS.validate_python(model_info).get("supports_prompt_cache_breakpoint")
|
||||
if isinstance(deployment_flag, bool):
|
||||
return deployment_flag
|
||||
return any(
|
||||
AnthropicCacheControlHook.catalog_targets_openai_prompt_cache_breakpoint(name, "azure_ai")
|
||||
for name in named_models
|
||||
)
|
||||
|
||||
|
||||
class AzureAIChatCompletionResponseIterator(OpenAIChatCompletionResponseIterator):
|
||||
def chunk_parser(self, chunk: dict[str, object]) -> ModelResponseStream: # mutable-ok: mirrors override contract
|
||||
usage: Final = chunk.get("usage")
|
||||
|
|
@ -251,6 +321,8 @@ class AzureAIStudioConfig(OpenAIConfig):
|
|||
self,
|
||||
messages: list[AllMessageValues],
|
||||
model: str,
|
||||
reads_breakpoints: bool | None = None,
|
||||
drop_params: bool = False,
|
||||
) -> list:
|
||||
"""
|
||||
- Azure AI Studio doesn't support content as a list. This handles:
|
||||
|
|
@ -260,13 +332,31 @@ class AzureAIStudioConfig(OpenAIConfig):
|
|||
Azure AI Foundry backends set additionalProperties=false and reject
|
||||
these with "Extra inputs are not permitted", which breaks multi-turn
|
||||
Anthropic-format clients that echo thinking blocks back as history.
|
||||
2. Transforms list content to a string.
|
||||
2. Transforms list content to a string, except on deployments that read OpenAI explicit
|
||||
prompt cache breakpoints (the deployment's own model_info opt-in, else the cost-map entry
|
||||
of its base_model or name, as loaded and without a sibling deployment's opt-in): only a
|
||||
content part can carry the breakpoint (a message-level one is silently ignored), so those
|
||||
keep their text parts, with each Anthropic `cache_control` mark renamed to a breakpoint the
|
||||
way the openai provider renames it. Every other part is dropped, exactly as the string
|
||||
conversion drops it, before the marks are renamed, so a message-level mark lands on the
|
||||
last text part. Under drop_params each breakpoint goes out as `{"mode": "explicit"}`, the only
|
||||
shape Foundry takes, and a malformed one is dropped.
|
||||
3. If message contains an image, audio, or PDF file part, send as is (user-intended)
|
||||
|
||||
Operates on a deep copy so the caller's messages keep their thinking blocks
|
||||
and provider metadata, which a fallback to another provider still needs.
|
||||
"""
|
||||
stripped_messages: Final = copy.deepcopy(drop_tool_reference_parts_from_tool_messages(messages))
|
||||
deployment_reads_breakpoints: Final = (
|
||||
AnthropicCacheControlHook.targets_openai_prompt_cache_breakpoint(model, "azure_ai")
|
||||
if reads_breakpoints is None
|
||||
else reads_breakpoints
|
||||
)
|
||||
copied_messages: Final = copy.deepcopy(drop_tool_reference_parts_from_tool_messages(messages))
|
||||
stripped_messages: Final = (
|
||||
[_with_breakpoints_for_foundry(message, drop_params) for message in copied_messages]
|
||||
if deployment_reads_breakpoints
|
||||
else copied_messages
|
||||
)
|
||||
for message in stripped_messages:
|
||||
message_dict = cast(dict, message) # cast-ok: TypedDict is a runtime dict stripped on our copy
|
||||
for field in NON_OPENAI_SPEC_MESSAGE_FIELDS:
|
||||
|
|
@ -275,6 +365,10 @@ class AzureAIStudioConfig(OpenAIConfig):
|
|||
if _has_list_only_content_part(message):
|
||||
continue
|
||||
|
||||
if deployment_reads_breakpoints and (text_parts := _text_parts(message)):
|
||||
message_dict["content"] = list(text_parts)
|
||||
continue
|
||||
|
||||
texts = convert_content_list_to_str(message=message)
|
||||
if texts:
|
||||
message["content"] = texts
|
||||
|
|
@ -329,7 +423,16 @@ class AzureAIStudioConfig(OpenAIConfig):
|
|||
if extra_body and isinstance(extra_body, dict):
|
||||
optional_params.update(extra_body)
|
||||
optional_params.pop("max_retries", None)
|
||||
return super().transform_request(model, messages, optional_params, litellm_params, headers)
|
||||
return {
|
||||
"model": model,
|
||||
"messages": self._transform_messages(
|
||||
messages,
|
||||
model,
|
||||
_deployment_reads_breakpoints(model, _FIELDS.validate_python(litellm_params)),
|
||||
bool(litellm_params.get("drop_params") or litellm.drop_params),
|
||||
),
|
||||
**optional_params,
|
||||
}
|
||||
|
||||
def transform_response(
|
||||
self,
|
||||
|
|
|
|||
|
|
@ -1,10 +1,12 @@
|
|||
import dataclasses
|
||||
from abc import ABC, abstractmethod
|
||||
from collections.abc import Mapping, Sequence
|
||||
from types import TracebackType
|
||||
from typing import TYPE_CHECKING, Any, Protocol
|
||||
from typing import TYPE_CHECKING, Any, Literal, Protocol
|
||||
|
||||
import httpx
|
||||
|
||||
from litellm._uuid import uuid
|
||||
from litellm.types.llms.openai import OpenAIRealtimeStreamSessionEvents
|
||||
from litellm.types.realtime import (
|
||||
RealtimeInputAudioTranscriptionUsage,
|
||||
|
|
@ -22,6 +24,26 @@ else:
|
|||
LiteLLMLoggingObj = Any
|
||||
|
||||
|
||||
@dataclasses.dataclass(frozen=True, slots=True)
|
||||
class RealtimeSessionUpdateRejection:
|
||||
code: Literal["invalid_type", "invalid_value", "session_update_not_applied", "session_update_too_large"]
|
||||
message: str
|
||||
param: str
|
||||
|
||||
def error_event(self, client_event_id: object) -> dict[str, object]:
|
||||
return {
|
||||
"type": "error",
|
||||
"event_id": f"event_{uuid.uuid4()}",
|
||||
"error": {
|
||||
"type": "invalid_request_error",
|
||||
"code": self.code,
|
||||
"message": self.message,
|
||||
"param": self.param,
|
||||
"event_id": client_event_id if isinstance(client_event_id, str) else None,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
class RealtimeBackend(Protocol):
|
||||
async def __aenter__(self) -> object: ...
|
||||
|
||||
|
|
@ -93,6 +115,14 @@ class BaseRealtimeConfig(ABC):
|
|||
def session_configuration_request(self, model: str) -> str | None: # message sent to setup the realtime session
|
||||
return None
|
||||
|
||||
def builds_setup_from_session_update(self) -> bool:
|
||||
return False
|
||||
|
||||
def setup_session_update_rejection(
|
||||
self, model: str, session: Mapping[str, object], *, setup_sent: bool
|
||||
) -> RealtimeSessionUpdateRejection | None:
|
||||
return None
|
||||
|
||||
def unbilled_usage_on_session_close(self, model: str) -> RealtimeInputAudioTranscriptionUsage | None:
|
||||
return None
|
||||
|
||||
|
|
|
|||
|
|
@ -17,7 +17,6 @@ from pydantic import JsonValue, TypeAdapter
|
|||
|
||||
import litellm
|
||||
from litellm._logging import redact_string, verbose_proxy_logger
|
||||
from litellm._uuid import uuid
|
||||
from litellm.constants import (
|
||||
BEDROCK_REALTIME_CLOSE_TIMEOUT_SECONDS,
|
||||
BEDROCK_REALTIME_COMMITTED_FAILURE_SCOPE_KEY,
|
||||
|
|
@ -31,13 +30,14 @@ from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix
|
|||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
|
||||
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
|
||||
from litellm.litellm_core_utils.realtime_streaming import DefaultLoggedRealTimeEventTypes
|
||||
from litellm.llms.base_llm.realtime.transformation import RealtimeSessionUpdateRejection
|
||||
from litellm.types.llms.bedrock import AwsAuthParams
|
||||
from litellm.types.llms.openai import OpenAIRealtimeEvents
|
||||
from litellm.types.realtime import RealtimeResponseTransformInput
|
||||
|
||||
from ..base_aws_llm import BaseAWSLLM, run_aws_signing
|
||||
from ..common_utils import BedrockError
|
||||
from .transformation import BedrockRealtimeConfig, BedrockSessionUpdateRejection, requested_modalities
|
||||
from .transformation import BedrockRealtimeConfig, requested_modalities
|
||||
|
||||
_CLIENT_MESSAGE_ADAPTER: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue)
|
||||
|
||||
|
|
@ -256,23 +256,11 @@ async def _ack_session_update(
|
|||
async def _reject_session_update(
|
||||
client_ws: RealtimeClientWebSocket,
|
||||
parsed_client_message: Mapping[str, JsonValue],
|
||||
rejection: BedrockSessionUpdateRejection,
|
||||
rejection: RealtimeSessionUpdateRejection,
|
||||
) -> bool:
|
||||
"""Send an error event for a session.update Bedrock cannot apply; False means the client is gone."""
|
||||
client_event_id: Final = parsed_client_message.get("event_id")
|
||||
error_event: Final = {
|
||||
"type": "error",
|
||||
"event_id": f"event_{uuid.uuid4()}",
|
||||
"error": {
|
||||
"type": "invalid_request_error",
|
||||
"code": rejection.code,
|
||||
"message": rejection.message,
|
||||
"param": rejection.param,
|
||||
"event_id": client_event_id if isinstance(client_event_id, str) else None,
|
||||
},
|
||||
}
|
||||
try:
|
||||
await client_ws.send_text(json.dumps(error_event))
|
||||
await client_ws.send_text(json.dumps(rejection.error_event(parsed_client_message.get("event_id"))))
|
||||
except Exception as e: # noqa: BLE001 # any send failure means the client is gone
|
||||
verbose_proxy_logger.debug("Client to Bedrock forwarding ended: %s", e, exc_info=True)
|
||||
return False
|
||||
|
|
|
|||
|
|
@ -11,7 +11,7 @@ import math
|
|||
import uuid as uuid_lib
|
||||
from collections.abc import Mapping, Sequence
|
||||
from types import MappingProxyType
|
||||
from typing import Final, Literal
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
from pydantic import ConfigDict, JsonValue, TypeAdapter
|
||||
|
|
@ -19,7 +19,7 @@ from pydantic import ConfigDict, JsonValue, TypeAdapter
|
|||
from litellm._logging 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.llms.base_llm.realtime.transformation import BaseRealtimeConfig, RealtimeSessionUpdateRejection
|
||||
from litellm.llms.bedrock.common_utils import BedrockError
|
||||
from litellm.llms.bedrock.realtime.trigger_audio import ready_trigger_pcm
|
||||
from litellm.types.llms.base import LiteLLMBaseModel
|
||||
|
|
@ -196,13 +196,6 @@ def _malformed_settings(session_config: Mapping[str, JsonValue]) -> tuple[tuple[
|
|||
return tuple((name, expected) for name, expected, is_malformed in malformed if is_malformed)
|
||||
|
||||
|
||||
@dataclasses.dataclass(frozen=True, slots=True)
|
||||
class BedrockSessionUpdateRejection:
|
||||
code: Literal["invalid_type", "invalid_value", "session_update_not_applied"]
|
||||
message: str
|
||||
param: str
|
||||
|
||||
|
||||
@dataclasses.dataclass(frozen=True, slots=True)
|
||||
class BedrockLockedSessionSettings:
|
||||
"""The session settings Bedrock Nova Sonic locks in once the prompt starts."""
|
||||
|
|
@ -533,7 +526,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig):
|
|||
max_tokens=max_tokens,
|
||||
)
|
||||
|
||||
def session_update_rejection(self, json_message: Mapping[str, JsonValue]) -> BedrockSessionUpdateRejection | None:
|
||||
def session_update_rejection(self, json_message: Mapping[str, JsonValue]) -> RealtimeSessionUpdateRejection | None:
|
||||
"""The error for a session.update this session cannot apply, or None when it applies."""
|
||||
session_config: Final = json_message.get("session")
|
||||
if not isinstance(session_config, Mapping):
|
||||
|
|
@ -541,7 +534,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig):
|
|||
malformed: Final = () if self.drop_params else _malformed_settings(session_config)
|
||||
if malformed:
|
||||
name, expected = malformed[0]
|
||||
return BedrockSessionUpdateRejection(
|
||||
return RealtimeSessionUpdateRejection(
|
||||
code="invalid_type",
|
||||
message=(
|
||||
f"Invalid type for 'session.{name}': Bedrock Nova Sonic needs {expected} there. "
|
||||
|
|
@ -551,7 +544,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig):
|
|||
)
|
||||
unplayable_voice: Final = None if self.drop_params else _unplayable_voice(session_config)
|
||||
if unplayable_voice is not None:
|
||||
return BedrockSessionUpdateRejection(
|
||||
return RealtimeSessionUpdateRejection(
|
||||
code="invalid_value",
|
||||
message=(
|
||||
f"Invalid value for 'session.voice': '{unplayable_voice}' is not a Bedrock Nova Sonic voice. "
|
||||
|
|
@ -562,7 +555,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig):
|
|||
if not self.prompt_started:
|
||||
return None
|
||||
if self._apply_locked_settings(session_config, self._started_settings) != self._started_settings:
|
||||
return BedrockSessionUpdateRejection(
|
||||
return RealtimeSessionUpdateRejection(
|
||||
code="session_update_not_applied",
|
||||
message="Bedrock Nova Sonic cannot change tools, instructions, voice, output audio format, temperature or max output tokens after the session has started, so this session.update was not applied",
|
||||
param="session",
|
||||
|
|
|
|||
|
|
@ -26,6 +26,10 @@ class BigModelChatConfig(ZAIChatConfig):
|
|||
def custom_llm_provider(self) -> str:
|
||||
return "bigmodel"
|
||||
|
||||
@property
|
||||
def prompts_unsupported_json_schema(self) -> bool:
|
||||
return False
|
||||
|
||||
def get_openai_compatible_provider_info(self, api_base: str | None, api_key: str | None) -> tuple[str, str | None]:
|
||||
return get_bigmodel_url(api_base, "chat"), get_bigmodel_api_key(api_key)
|
||||
|
||||
|
|
|
|||
|
|
@ -6973,14 +6973,19 @@ class BaseLLMHTTPHandler:
|
|||
),
|
||||
)
|
||||
|
||||
# Auto-send session setup if the provider requires it (e.g.
|
||||
# Gemini/Vertex AI Live needs a `setup` before any realtime_input).
|
||||
# Auto-send session setup at connect unless the provider builds it from the
|
||||
# client's first session.update (Gemini/Vertex AI Live accept exactly one
|
||||
# setup, so RealTimeStreaming sends theirs on the first client message).
|
||||
# Build the streaming handler first so a transcription guardrail's
|
||||
# auto-response disable can be folded into this one setup: Gemini
|
||||
# rejects a second setup, so a follow-up disable would be dropped
|
||||
# and the guardrail bypassed.
|
||||
_session_config: str | None = None
|
||||
if provider_config.requires_session_configuration():
|
||||
sends_setup_at_connect: Final = (
|
||||
provider_config.requires_session_configuration()
|
||||
and not provider_config.builds_setup_from_session_update()
|
||||
)
|
||||
if sends_setup_at_connect:
|
||||
_session_config = provider_config.session_configuration_request(model)
|
||||
if _session_config:
|
||||
_session_config = realtime_streaming.maybe_inject_guardrail_auto_response_disable(
|
||||
|
|
|
|||
|
|
@ -7,6 +7,7 @@ from collections import OrderedDict
|
|||
from collections.abc import Mapping, Sequence
|
||||
from typing import Final, cast
|
||||
|
||||
from pydantic import TypeAdapter
|
||||
from typing_extensions import ReadOnly, Required, TypedDict
|
||||
|
||||
import litellm
|
||||
|
|
@ -15,8 +16,9 @@ from litellm._uuid import uuid
|
|||
from litellm.litellm_core_utils.core_helpers import as_str_mapping
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.litellm_core_utils.url_utils import join_api_base
|
||||
from litellm.llms.base_llm.realtime.transformation import BaseRealtimeConfig
|
||||
from litellm.llms.base_llm.realtime.transformation import BaseRealtimeConfig, RealtimeSessionUpdateRejection
|
||||
from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
|
||||
OPENAI_TO_GEMINI_TTS_VOICES,
|
||||
VertexGeminiConfig,
|
||||
)
|
||||
from litellm.responses.litellm_completion_transformation.transformation import (
|
||||
|
|
@ -76,27 +78,82 @@ MAP_GEMINI_FIELD_TO_OPENAI_EVENT: Final[dict[str, OpenAIRealtimeEventTypes | Res
|
|||
|
||||
# Keys the main transform loop handles; siblings like ``usageMetadata`` are skipped.
|
||||
_KNOWN_GEMINI_TOP_LEVEL_KEYS: Final[set] = {map_key.split(".", 1)[0] for map_key in MAP_GEMINI_FIELD_TO_OPENAI_EVENT}
|
||||
_GEMINI_LIVE_RESPONSE_MODALITIES: Final = frozenset({"TEXT", "AUDIO"})
|
||||
|
||||
|
||||
OPENAI_STOCK_REALTIME_VOICES: Final[frozenset[str]] = frozenset(
|
||||
{"alloy", "ash", "ballad", "cedar", "coral", "echo", "marin", "sage", "shimmer", "verse"}
|
||||
)
|
||||
|
||||
# Per https://docs.cloud.google.com/vertex-ai/generative-ai/docs/live-api/configure-language-voice (2026-10-28)
|
||||
GEMINI_LIVE_DOCUMENTED_VOICES: Final[frozenset[str]] = frozenset(
|
||||
{
|
||||
"achernar",
|
||||
"achird",
|
||||
"algenib",
|
||||
"algieba",
|
||||
"alnilam",
|
||||
"aoede",
|
||||
"autonoe",
|
||||
"callirrhoe",
|
||||
"charon",
|
||||
"despina",
|
||||
"enceladus",
|
||||
"erinome",
|
||||
"fenrir",
|
||||
"gacrux",
|
||||
"iapetus",
|
||||
"kore",
|
||||
"laomedeia",
|
||||
"leda",
|
||||
"orus",
|
||||
"puck",
|
||||
"pulcherrima",
|
||||
"rasalgethi",
|
||||
"sadachbia",
|
||||
"sadaltager",
|
||||
"schedar",
|
||||
"sulafat",
|
||||
"umbriel",
|
||||
"vindemiatrix",
|
||||
"zephyr",
|
||||
"zubenelgenubi",
|
||||
}
|
||||
)
|
||||
|
||||
def _gemini_live_speech_config(voice: object) -> Mapping[str, object] | None:
|
||||
|
||||
def _is_undocumented_live_voice(voice: object) -> bool:
|
||||
return (
|
||||
isinstance(voice, str)
|
||||
and voice.lower() not in OPENAI_STOCK_REALTIME_VOICES
|
||||
and OPENAI_TO_GEMINI_TTS_VOICES.get(voice.lower(), voice).lower() not in GEMINI_LIVE_DOCUMENTED_VOICES
|
||||
)
|
||||
|
||||
|
||||
def _gemini_live_speech_config(voice: object, drop_undocumented: bool) -> Mapping[str, object] | None:
|
||||
"""Build the Gemini Live speechConfig for a client-requested voice.
|
||||
|
||||
OpenAI stock voice names have no Gemini equivalent and Gemini Live closes
|
||||
the session on an unknown voice, so they are dropped with a warning and
|
||||
the model keeps its default voice. Every other name is forwarded verbatim.
|
||||
OpenAI stock voice names have no Gemini equivalent, so they are always
|
||||
dropped, and so is any name outside Google's documented Live voices when
|
||||
``drop_undocumented`` is set, keeping the model's default voice. Every
|
||||
other name is forwarded verbatim.
|
||||
"""
|
||||
if isinstance(voice, str) and voice.lower() in OPENAI_STOCK_REALTIME_VOICES:
|
||||
if not isinstance(voice, str):
|
||||
return VertexGeminiConfig().map_audio_params({"voice": voice})
|
||||
if voice.lower() in OPENAI_STOCK_REALTIME_VOICES:
|
||||
verbose_logger.warning(
|
||||
"Gemini Realtime: voice %s is an OpenAI voice with no Gemini equivalent; "
|
||||
"dropping it so the session keeps the model's default voice.",
|
||||
voice,
|
||||
)
|
||||
return None
|
||||
if drop_undocumented and _is_undocumented_live_voice(voice):
|
||||
verbose_logger.warning(
|
||||
"Gemini Realtime: voice %s is not a documented Gemini Live voice; "
|
||||
"drop_params is set, so it is dropped and the session keeps the model's default voice.",
|
||||
voice,
|
||||
)
|
||||
return None
|
||||
return VertexGeminiConfig().map_audio_params({"voice": voice})
|
||||
|
||||
|
||||
|
|
@ -142,16 +199,34 @@ def _base64_decoded_byte_count(data: str) -> int:
|
|||
return max(len(data) * 3 // 4 - padding, 0)
|
||||
|
||||
|
||||
_JSON_ARRAY: Final = TypeAdapter(tuple[object, ...])
|
||||
|
||||
|
||||
def _chat_completion_tool(tool: object) -> dict[str, object]:
|
||||
fields: Final = as_str_mapping(tool) or {}
|
||||
if fields.get("type") != "function" or "function" in fields:
|
||||
return {**fields}
|
||||
return {"type": "function", "function": {key: value for key, value in fields.items() if key != "type"}}
|
||||
|
||||
|
||||
def _system_instruction_text(system_instruction: HttpxContentType | None) -> str | None:
|
||||
parts: Final = system_instruction.get("parts", []) if system_instruction else []
|
||||
texts: Final = tuple(text for part in parts if isinstance(text := part.get("text"), str))
|
||||
return "\n".join(texts) if texts else None
|
||||
|
||||
|
||||
class GeminiRealtimeConfig(BaseRealtimeConfig):
|
||||
_TOOL_CALL_ID_TO_NAME_MAX = 256 # LRU cap for call_id→name mapping
|
||||
|
||||
def __init__(self):
|
||||
def __init__(self, drop_params: bool = False):
|
||||
super().__init__()
|
||||
self.drop_params = drop_params
|
||||
self._tool_call_id_to_name: OrderedDict[str, str] = OrderedDict()
|
||||
# Gemini Live sometimes emits usageMetadata in a standalone frame between
|
||||
# turns; buffer it here so the next response.done carries the token counts.
|
||||
self._pending_usage_metadata: dict | None = None
|
||||
self._unbilled_input_audio_bytes: int = 0
|
||||
self._manual_activity_open: bool = False
|
||||
|
||||
def is_setup_message(self, msg_obj: dict) -> bool:
|
||||
return "setup" in msg_obj
|
||||
|
|
@ -167,6 +242,10 @@ class GeminiRealtimeConfig(BaseRealtimeConfig):
|
|||
"""Google AI Studio closes the socket (1007) on padding or silence tuning next to ``disabled: true``."""
|
||||
return True
|
||||
|
||||
def _rejects_undocumented_voices(self, model: str) -> bool:
|
||||
"""Google AI Studio closes the socket (1007) on an unknown voice only for its native-audio Live models."""
|
||||
return bool(GeminiRealtimeConfig._model_cost_entry(model).get("gemini_native_audio"))
|
||||
|
||||
@staticmethod
|
||||
def _usage_detail_alias(details: object, defaults: dict[str, int]) -> dict[str, object]:
|
||||
detail_mapping: Final = as_str_mapping(details)
|
||||
|
|
@ -272,17 +351,26 @@ class GeminiRealtimeConfig(BaseRealtimeConfig):
|
|||
except (json.JSONDecodeError, TypeError, AttributeError):
|
||||
return False
|
||||
|
||||
def _start_manual_activity(self, session_configuration_request: str | None) -> tuple[str, ...]:
|
||||
if self._manual_activity_open or not self._manual_turn_detection_enabled(session_configuration_request):
|
||||
return ()
|
||||
self._manual_activity_open = True
|
||||
activity_start: Final[BidiGenerateContentRealtimeInput] = {"activityStart": {}}
|
||||
return (json.dumps({"realtimeInput": activity_start}),)
|
||||
|
||||
def _end_manual_activity(self) -> tuple[str, ...]:
|
||||
if not self._manual_activity_open:
|
||||
return ()
|
||||
self._manual_activity_open = False
|
||||
activity_end: Final[BidiGenerateContentRealtimeInput] = {"activityEnd": {}}
|
||||
return (json.dumps({"realtimeInput": activity_end}),)
|
||||
|
||||
def _handle_input_audio_buffer_commit_or_end(self, session_configuration_request: str | None) -> list[str]:
|
||||
"""Map OpenAI buffer commit/end to Gemini Live turn-boundary signals."""
|
||||
if self._manual_turn_detection_enabled(session_configuration_request):
|
||||
realtime_input_dict: BidiGenerateContentRealtimeInput = {
|
||||
"activityEnd": True,
|
||||
}
|
||||
verbose_logger.debug("Gemini Realtime: Sending activityEnd realtimeInput to backend")
|
||||
else:
|
||||
realtime_input_dict = {"audioStreamEnd": True}
|
||||
verbose_logger.debug("Gemini Realtime: Sending audioStreamEnd realtimeInput to backend")
|
||||
return [json.dumps({"realtimeInput": realtime_input_dict})]
|
||||
return list(self._end_manual_activity())
|
||||
audio_stream_end: Final[BidiGenerateContentRealtimeInput] = {"audioStreamEnd": True}
|
||||
return [json.dumps({"realtimeInput": audio_stream_end})]
|
||||
|
||||
def map_automatic_turn_detection(self, value: OpenAIRealtimeTurnDetection) -> AutomaticActivityDetection:
|
||||
"""Map OpenAI ``server_vad`` to Gemini ``automaticActivityDetection``.
|
||||
|
|
@ -325,7 +413,7 @@ class GeminiRealtimeConfig(BaseRealtimeConfig):
|
|||
"voice",
|
||||
]
|
||||
|
||||
def map_openai_params(self, optional_params: dict, non_default_params: dict) -> dict:
|
||||
def map_openai_params(self, optional_params: dict, non_default_params: dict, *, model: str) -> dict:
|
||||
if "generationConfig" not in optional_params:
|
||||
optional_params["generationConfig"] = {}
|
||||
for key, value in non_default_params.items():
|
||||
|
|
@ -333,12 +421,12 @@ class GeminiRealtimeConfig(BaseRealtimeConfig):
|
|||
optional_params["systemInstruction"] = HttpxContentType(role="user", parts=[{"text": value}])
|
||||
elif key == "temperature":
|
||||
optional_params["generationConfig"]["temperature"] = value
|
||||
elif key == "max_response_output_tokens":
|
||||
elif key == "max_response_output_tokens" and isinstance(value, int):
|
||||
optional_params["generationConfig"]["maxOutputTokens"] = value
|
||||
elif key == "modalities":
|
||||
optional_params["generationConfig"]["responseModalities"] = [
|
||||
modality.upper() for modality in cast(list[str], value)
|
||||
]
|
||||
optional_params["generationConfig"]["responseModalities"] = list(
|
||||
GeminiRealtimeConfig._requested_response_modalities(cast(object, value))
|
||||
)
|
||||
elif key == "tools":
|
||||
from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
|
||||
VertexGeminiConfig,
|
||||
|
|
@ -347,34 +435,43 @@ class GeminiRealtimeConfig(BaseRealtimeConfig):
|
|||
vertex_gemini_config = VertexGeminiConfig()
|
||||
# Tools should be at the top level of setup, not inside generationConfig
|
||||
optional_params["tools"] = vertex_gemini_config.map_function(
|
||||
value=value, optional_params=optional_params
|
||||
value=[_chat_completion_tool(tool) for tool in _JSON_ARRAY.validate_python(value)],
|
||||
optional_params=optional_params,
|
||||
)
|
||||
elif key == "input_audio_transcription" and value is not None:
|
||||
optional_params["inputAudioTranscription"] = {}
|
||||
elif key == "turn_detection" and value is not None:
|
||||
value_typed = cast(OpenAIRealtimeTurnDetection, value)
|
||||
if (
|
||||
isinstance(value_typed, dict)
|
||||
and value_typed.get("type") == "semantic_vad"
|
||||
and "create_response" not in value_typed
|
||||
):
|
||||
# Pipecat/OpenAI GA semantic VAD — skip; Gemini uses its own VAD.
|
||||
# Only skip when there is no create_response override so that
|
||||
# a guardrail-injected create_response:false is not dropped.
|
||||
continue
|
||||
transformed_audio_activity_config = self.map_automatic_turn_detection(value_typed)
|
||||
if transformed_audio_activity_config:
|
||||
optional_params["realtimeInputConfig"] = BidiGenerateContentRealtimeInputConfig(
|
||||
automaticActivityDetection=transformed_audio_activity_config
|
||||
)
|
||||
elif key == "turn_detection":
|
||||
realtime_input_config = self._realtime_input_config(cast(object, value))
|
||||
if realtime_input_config is not None:
|
||||
optional_params["realtimeInputConfig"] = realtime_input_config
|
||||
elif key == "voice":
|
||||
speech_config = _gemini_live_speech_config(value)
|
||||
speech_config = _gemini_live_speech_config(
|
||||
value, self.drop_params and self._rejects_undocumented_voices(model)
|
||||
)
|
||||
if speech_config:
|
||||
optional_params["generationConfig"]["speechConfig"] = speech_config
|
||||
if len(optional_params["generationConfig"]) == 0:
|
||||
optional_params.pop("generationConfig")
|
||||
return optional_params
|
||||
|
||||
def _realtime_input_config(self, turn_detection: object) -> BidiGenerateContentRealtimeInputConfig | None:
|
||||
if turn_detection is None:
|
||||
return BidiGenerateContentRealtimeInputConfig(automaticActivityDetection={"disabled": True})
|
||||
value_typed: Final = cast(OpenAIRealtimeTurnDetection, turn_detection)
|
||||
if (
|
||||
isinstance(value_typed, dict)
|
||||
and value_typed.get("type") == "semantic_vad"
|
||||
and "create_response" not in value_typed
|
||||
):
|
||||
# Pipecat/OpenAI GA semantic VAD — skip; Gemini uses its own VAD.
|
||||
# Only skip when there is no create_response override so that
|
||||
# a guardrail-injected create_response:false is not dropped.
|
||||
return None
|
||||
transformed_audio_activity_config: Final = self.map_automatic_turn_detection(value_typed)
|
||||
if not transformed_audio_activity_config:
|
||||
return None
|
||||
return BidiGenerateContentRealtimeInputConfig(automaticActivityDetection=transformed_audio_activity_config)
|
||||
|
||||
@staticmethod
|
||||
def _extract_turn_detection(session: dict) -> dict | None:
|
||||
"""Extract turn_detection from a session.update payload.
|
||||
|
|
@ -423,6 +520,8 @@ class GeminiRealtimeConfig(BaseRealtimeConfig):
|
|||
if isinstance(input_cfg, dict):
|
||||
if "input_audio_transcription" not in normalized and "transcription" in input_cfg:
|
||||
normalized["input_audio_transcription"] = input_cfg["transcription"]
|
||||
if "turn_detection" not in normalized and "turn_detection" in input_cfg:
|
||||
normalized["turn_detection"] = input_cfg["turn_detection"]
|
||||
output_cfg: Final = audio.get("output")
|
||||
if isinstance(output_cfg, dict) and output_cfg.get("voice"):
|
||||
normalized["voice"] = output_cfg["voice"]
|
||||
|
|
@ -434,7 +533,7 @@ class GeminiRealtimeConfig(BaseRealtimeConfig):
|
|||
return normalized
|
||||
|
||||
@staticmethod
|
||||
def _model_cost_entry(model: str) -> dict:
|
||||
def _model_cost_entry(model: str) -> Mapping[str, object]:
|
||||
entry = litellm.model_cost.get(model)
|
||||
if entry is None:
|
||||
stripped: Final = model.split("/", 1)[-1]
|
||||
|
|
@ -455,18 +554,122 @@ class GeminiRealtimeConfig(BaseRealtimeConfig):
|
|||
return "TEXT" if GeminiRealtimeConfig._is_text_only_live_model(model) else "AUDIO"
|
||||
|
||||
@staticmethod
|
||||
def _coerce_response_modalities(model: str, modalities: Sequence[object]) -> tuple[str, ...]:
|
||||
"""Swap responseModalities a Live model cannot produce: TEXT to AUDIO for
|
||||
audio-only models, AUDIO to TEXT for text-only ones (e.g. transcribe-live)."""
|
||||
normalized: Final = tuple(
|
||||
modality.upper() if isinstance(modality, str) else str(modality).upper() for modality in modalities
|
||||
def _requested_response_modalities(value: object) -> tuple[str, ...]:
|
||||
if isinstance(value, str):
|
||||
return (value.upper(),)
|
||||
if not isinstance(value, (list, tuple)):
|
||||
return ()
|
||||
return tuple(str(modality).upper() for modality in _JSON_ARRAY.validate_python(value))
|
||||
|
||||
@staticmethod
|
||||
def _unsupported_response_modalities(model: str, requested: tuple[str, ...]) -> tuple[str, ...]:
|
||||
unknown: Final = tuple(modality for modality in requested if modality not in _GEMINI_LIVE_RESPONSE_MODALITIES)
|
||||
if unknown:
|
||||
return unknown
|
||||
if GeminiRealtimeConfig._is_audio_only_live_model(model) and frozenset(requested) == frozenset(("TEXT",)):
|
||||
return ("TEXT",)
|
||||
if GeminiRealtimeConfig._is_text_only_live_model(model) and frozenset(requested) == frozenset(("AUDIO",)):
|
||||
return ("AUDIO",)
|
||||
return ()
|
||||
|
||||
def setup_session_update_rejection(
|
||||
self, model: str, session: Mapping[str, object], *, setup_sent: bool
|
||||
) -> RealtimeSessionUpdateRejection | None:
|
||||
if self.drop_params:
|
||||
return None
|
||||
if setup_sent:
|
||||
return RealtimeSessionUpdateRejection(
|
||||
code="session_update_not_applied",
|
||||
message=(
|
||||
"Gemini Live accepts one session configuration per connection, and this session already sent "
|
||||
"its configuration, so this session.update was not applied. Put every setting in session.update "
|
||||
"before the first conversation content, or set drop_params to drop it instead"
|
||||
),
|
||||
param="session",
|
||||
)
|
||||
return self._modality_rejection(model, session) or self._voice_rejection(model, session)
|
||||
|
||||
@staticmethod
|
||||
def _modality_rejection(model: str, session: Mapping[str, object]) -> RealtimeSessionUpdateRejection | None:
|
||||
param: Final = "output_modalities" if "output_modalities" in session else "modalities"
|
||||
unsupported: Final = GeminiRealtimeConfig._unsupported_response_modalities(
|
||||
model, GeminiRealtimeConfig._requested_response_modalities(session.get(param))
|
||||
)
|
||||
if not unsupported:
|
||||
return None
|
||||
return RealtimeSessionUpdateRejection(
|
||||
code="invalid_value",
|
||||
message=(
|
||||
f"Invalid value for 'session.{param}': {model} cannot produce "
|
||||
f"{', '.join(modality.lower() for modality in unsupported)} output. Set drop_params to drop it instead"
|
||||
),
|
||||
param=f"session.{param}",
|
||||
)
|
||||
|
||||
def _voice_rejection(self, model: str, session: Mapping[str, object]) -> RealtimeSessionUpdateRejection | None:
|
||||
audio_output: Final = as_str_mapping((as_str_mapping(session.get("audio")) or {}).get("output")) or {}
|
||||
output_voice: Final = audio_output.get("voice")
|
||||
voice: Final = output_voice or session.get("voice")
|
||||
if not (self._rejects_undocumented_voices(model) and _is_undocumented_live_voice(voice)):
|
||||
return None
|
||||
param: Final = "session.audio.output.voice" if output_voice else "session.voice"
|
||||
return RealtimeSessionUpdateRejection(
|
||||
code="invalid_value",
|
||||
message=(
|
||||
f"Invalid value for '{param}': {voice} is not a Gemini Live voice {model} offers. "
|
||||
"Set drop_params to drop it instead"
|
||||
),
|
||||
param=param,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _coerce_response_modalities(model: str, modalities: Sequence[object]) -> tuple[str, ...]:
|
||||
"""Drop modalities Gemini Live cannot produce at all, swap the ones this
|
||||
model cannot (TEXT to AUDIO for audio-only models, AUDIO to TEXT for
|
||||
text-only ones such as transcribe-live), then keep a single modality
|
||||
because Gemini Live rejects a setup naming two."""
|
||||
return GeminiRealtimeConfig._single_response_modality(
|
||||
GeminiRealtimeConfig._swap_unsupported_response_modalities(
|
||||
model, GeminiRealtimeConfig._known_response_modalities(model, modalities)
|
||||
)
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _known_response_modalities(model: str, modalities: Sequence[object]) -> tuple[str, ...]:
|
||||
normalized: Final = tuple(str(modality).upper() for modality in modalities)
|
||||
known: Final = tuple(modality for modality in normalized if modality in _GEMINI_LIVE_RESPONSE_MODALITIES)
|
||||
if len(known) < len(normalized):
|
||||
verbose_logger.warning(
|
||||
"Gemini Live: %s only produces TEXT or AUDIO; ignoring the requested %s",
|
||||
model,
|
||||
", ".join(modality for modality in normalized if modality not in _GEMINI_LIVE_RESPONSE_MODALITIES),
|
||||
)
|
||||
return known or (GeminiRealtimeConfig._default_response_modality(model),)
|
||||
|
||||
@staticmethod
|
||||
def _swap_unsupported_response_modalities(model: str, normalized: tuple[str, ...]) -> tuple[str, ...]:
|
||||
if GeminiRealtimeConfig._is_audio_only_live_model(model) and "TEXT" in normalized:
|
||||
verbose_logger.warning(
|
||||
"Gemini Live: %s only produces audio; downgrading the requested TEXT response modality to AUDIO",
|
||||
model,
|
||||
)
|
||||
return tuple(modality for modality in normalized if modality != "TEXT") or ("AUDIO",)
|
||||
if GeminiRealtimeConfig._is_text_only_live_model(model) and "AUDIO" in normalized:
|
||||
return tuple(modality for modality in normalized if modality != "AUDIO") or ("TEXT",)
|
||||
return normalized
|
||||
|
||||
@staticmethod
|
||||
def _single_response_modality(modalities: tuple[str, ...]) -> tuple[str, ...]:
|
||||
if len(modalities) <= 1:
|
||||
return modalities
|
||||
chosen: Final = "AUDIO" if "AUDIO" in modalities else modalities[0]
|
||||
verbose_logger.warning(
|
||||
"Gemini Live: at most one response modality is allowed per session; using %s for the requested %s",
|
||||
chosen,
|
||||
", ".join(modalities),
|
||||
)
|
||||
return (chosen,)
|
||||
|
||||
@staticmethod
|
||||
def _finalize_gemini_live_setup(model: str, setup: dict[str, object]) -> dict[str, object]:
|
||||
generation_config: Final = setup.get("generationConfig")
|
||||
|
|
@ -504,7 +707,9 @@ class GeminiRealtimeConfig(BaseRealtimeConfig):
|
|||
# would be silently dropped because ``map_openai_params`` only
|
||||
# recognises the flat OpenAI-beta key names.
|
||||
session_payload = self._normalize_session_payload_for_mapping(session_payload)
|
||||
new_overrides: Final = self.map_openai_params(optional_params={}, non_default_params=session_payload)
|
||||
new_overrides: Final = self.map_openai_params(
|
||||
optional_params={}, non_default_params=session_payload, model=model
|
||||
)
|
||||
|
||||
if session_configuration_request is None:
|
||||
generation_config: Final = new_overrides.setdefault("generationConfig", {})
|
||||
|
|
@ -618,14 +823,13 @@ class GeminiRealtimeConfig(BaseRealtimeConfig):
|
|||
message_str = str(message)
|
||||
raise ValueError(f"Invalid JSON message: {message_str}")
|
||||
|
||||
messages: Final[list[str]] = []
|
||||
msg_type: Final = json_message.get("type")
|
||||
|
||||
if msg_type == "session.update":
|
||||
return self._handle_session_update(json_message, model, session_configuration_request)
|
||||
|
||||
if msg_type == "response.create":
|
||||
return [] # Gemini responds automatically; nothing to forward
|
||||
return list(self._end_manual_activity())
|
||||
|
||||
if msg_type == "conversation.item.create":
|
||||
return self._handle_conversation_item(json_message)
|
||||
|
|
@ -643,8 +847,7 @@ class GeminiRealtimeConfig(BaseRealtimeConfig):
|
|||
|
||||
gemini_msg: Final = json.dumps({"realtimeInput": realtime_input_dict})
|
||||
verbose_logger.debug("Gemini Realtime: Sending audio realtimeInput to backend")
|
||||
messages.append(gemini_msg)
|
||||
return messages
|
||||
return [*self._start_manual_activity(session_configuration_request), gemini_msg]
|
||||
|
||||
if msg_type in ("input_audio_buffer.commit", "input_audio_buffer.end"):
|
||||
return self._handle_input_audio_buffer_commit_or_end(session_configuration_request)
|
||||
|
|
@ -660,22 +863,21 @@ class GeminiRealtimeConfig(BaseRealtimeConfig):
|
|||
logging_session_id: str,
|
||||
session_configuration_request: str | None = None,
|
||||
) -> OpenAIRealtimeStreamSessionEvents:
|
||||
if session_configuration_request:
|
||||
session_configuration_request_dict: BidiGenerateContentSetup = _parse_setup(session_configuration_request)
|
||||
else:
|
||||
session_configuration_request_dict = {}
|
||||
session_configuration_request_dict: Final = _parse_setup(
|
||||
session_configuration_request or self.session_configuration_request(model)
|
||||
)
|
||||
|
||||
_model: Final = session_configuration_request_dict.get("model") or model
|
||||
generation_config: Final = session_configuration_request_dict.get("generationConfig", {}) or {}
|
||||
gemini_modalities: Final = generation_config.get("responseModalities", ["AUDIO"])
|
||||
_modalities: Final = [modality.lower() for modality in cast(list[str], gemini_modalities)]
|
||||
_system_instruction: Final = session_configuration_request_dict.get("systemInstruction")
|
||||
instructions: Final = _system_instruction_text(session_configuration_request_dict.get("systemInstruction"))
|
||||
session: Final = OpenAIRealtimeStreamSession(
|
||||
id=logging_session_id,
|
||||
modalities=_modalities,
|
||||
)
|
||||
if _system_instruction is not None and isinstance(_system_instruction, str):
|
||||
session["instructions"] = _system_instruction
|
||||
if instructions is not None:
|
||||
session["instructions"] = instructions
|
||||
if _model is not None and isinstance(_model, str):
|
||||
# Strip Vertex/AI Studio path prefixes to expose the bare model name.
|
||||
if "/models/" in _model:
|
||||
|
|
@ -1662,6 +1864,9 @@ class GeminiRealtimeConfig(BaseRealtimeConfig):
|
|||
# Deferred setup opt-in: litellm.gemini_live_defer_setup = True
|
||||
return not litellm.gemini_live_defer_setup
|
||||
|
||||
def builds_setup_from_session_update(self) -> bool:
|
||||
return True
|
||||
|
||||
def session_configuration_request(self, model: str) -> str:
|
||||
"""
|
||||
|
||||
|
|
|
|||
|
|
@ -9,11 +9,16 @@ Docs - https://docs.mistral.ai/api/
|
|||
from collections.abc import AsyncIterator, Coroutine, Iterator, Mapping, Sequence
|
||||
from itertools import accumulate, count
|
||||
from typing import TYPE_CHECKING, Final, Literal, cast, get_type_hints, overload
|
||||
from urllib.parse import urlparse
|
||||
|
||||
import httpx
|
||||
from pydantic import TypeAdapter
|
||||
from typing_extensions import TypeIs # noqa: TID251 # TypeIs reaches typing only on 3.13
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.exceptions import UnsupportedParamsError
|
||||
from litellm.litellm_core_utils.core_helpers import as_str_mapping
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.litellm_core_utils.prompt_templates.common_utils import (
|
||||
convert_content_list_to_str,
|
||||
|
|
@ -30,7 +35,15 @@ from litellm.router_utils.reasoning_effort_capability import (
|
|||
nearest_declared_reasoning_effort,
|
||||
)
|
||||
from litellm.secret_managers.main import get_secret_str
|
||||
from litellm.types.llms.mistral import MistralTextBlock, MistralThinkingBlock, MistralToolCallMessage
|
||||
from litellm.types.llms.mistral import (
|
||||
MistralCacheKeyParams,
|
||||
MistralFilePart,
|
||||
MistralReasoningEffortParams,
|
||||
MistralRequestRouting,
|
||||
MistralTextBlock,
|
||||
MistralThinkingBlock,
|
||||
MistralToolCallMessage,
|
||||
)
|
||||
from litellm.types.llms.openai import AllMessageValues, ChatCompletionAssistantMessage, ChatCompletionTextObject
|
||||
from litellm.types.utils import ModelResponse, ModelResponseStream
|
||||
from litellm.utils import convert_to_model_response_object, supports_reasoning
|
||||
|
|
@ -57,6 +70,17 @@ def _accepted_reasoning_effort(model: str, requested: str, custom_llm_provider:
|
|||
|
||||
|
||||
_NO_TOOL_CALL_TURN: Final[tuple[int, frozenset[str]]] = (-1, frozenset())
|
||||
_MISTRAL_REQUEST_ROUTING: Final = TypeAdapter(MistralRequestRouting)
|
||||
_MISTRAL_CACHE_KEY_PARAMS: Final = TypeAdapter(MistralCacheKeyParams)
|
||||
_MISTRAL_REASONING_EFFORT_PARAMS: Final = TypeAdapter(MistralReasoningEffortParams)
|
||||
_MISTRAL_FILE_PART: Final = TypeAdapter(MistralFilePart)
|
||||
|
||||
|
||||
def _reasoning_level_and_summary(reasoning_effort: object) -> tuple[object, object]:
|
||||
requested: Final = as_str_mapping(reasoning_effort)
|
||||
if requested is None:
|
||||
return reasoning_effort, None
|
||||
return requested.get("effort"), requested.get("summary")
|
||||
|
||||
|
||||
def _is_json_array(value: object) -> TypeIs[Sequence[object]]: # guard-ok: a decoded JSON array is a list
|
||||
|
|
@ -88,6 +112,8 @@ class MistralConfig(OpenAIGPTConfig):
|
|||
- `response_format` (object or null): An object specifying the format that the model must output. Setting to { "type": "json_object" } enables JSON mode, which guarantees the message the model generates is in JSON. When using JSON mode you MUST also instruct the model to produce JSON yourself with a system or a user message.
|
||||
"""
|
||||
|
||||
_is_base_class = False
|
||||
|
||||
temperature: int | None = None
|
||||
top_p: int | None = None
|
||||
max_tokens: int | None = None
|
||||
|
|
@ -131,6 +157,7 @@ class MistralConfig(OpenAIGPTConfig):
|
|||
def get_supported_openai_params(self, model: str) -> list[str]:
|
||||
is_magistral: Final = "magistral" in model.lower()
|
||||
accepts_reasoning_effort: Final = self._accepts_reasoning(model)
|
||||
accepts_prompt_cache_key: Final = self.custom_llm_provider == "mistral"
|
||||
return [
|
||||
"stream",
|
||||
"temperature",
|
||||
|
|
@ -145,6 +172,7 @@ class MistralConfig(OpenAIGPTConfig):
|
|||
"parallel_tool_calls",
|
||||
*(("thinking",) if is_magistral else ()),
|
||||
*(("reasoning_effort",) if accepts_reasoning_effort else ()),
|
||||
*(("prompt_cache_key",) if accepts_prompt_cache_key else ()),
|
||||
]
|
||||
|
||||
def _map_tool_choice(self, tool_choice: str) -> str:
|
||||
|
|
@ -211,15 +239,84 @@ class MistralConfig(OpenAIGPTConfig):
|
|||
optional_params["extra_body"] = {"random_seed": value}
|
||||
if param == "response_format":
|
||||
optional_params["response_format"] = value
|
||||
if param == "reasoning_effort" and "magistral" not in model.lower():
|
||||
optional_params["reasoning_effort"] = _accepted_reasoning_effort(model, value, self.custom_llm_provider)
|
||||
if (
|
||||
param == "reasoning_effort"
|
||||
and "magistral" not in model.lower()
|
||||
and (
|
||||
reasoning_level := self._reasoning_level_to_send(
|
||||
model, _MISTRAL_REASONING_EFFORT_PARAMS.validate_python(non_default_params), drop_params
|
||||
)
|
||||
)
|
||||
is not None
|
||||
):
|
||||
optional_params["reasoning_effort"] = reasoning_level
|
||||
if param in ("reasoning_effort", "thinking") and "magistral" in model.lower():
|
||||
# Flag that we need to add reasoning system prompt
|
||||
optional_params["_add_reasoning_prompt"] = True
|
||||
if param == "parallel_tool_calls":
|
||||
optional_params["parallel_tool_calls"] = value
|
||||
if param in ("parallel_tool_calls", "prompt_cache_key"):
|
||||
optional_params[param] = value
|
||||
return optional_params
|
||||
|
||||
def _reasoning_level_to_send(
|
||||
self, model: str, requested: MistralReasoningEffortParams, drop_params: bool
|
||||
) -> str | None:
|
||||
level, summary = _reasoning_level_and_summary(requested.get("reasoning_effort"))
|
||||
sendable: Final = (
|
||||
_accepted_reasoning_effort(model, level, self.custom_llm_provider) if isinstance(level, str) else None
|
||||
)
|
||||
if sendable is not None and summary is None:
|
||||
return sendable
|
||||
if drop_params or litellm.drop_params is True:
|
||||
return sendable
|
||||
raise UnsupportedParamsError(
|
||||
message=(
|
||||
f"{self.custom_llm_provider} does not support parameters: ['reasoning_effort'], for model={model}, "
|
||||
"since the Mistral API takes it only as a level such as 'high', with no reasoning summary. "
|
||||
"To drop it, set `litellm.drop_params=True` or for proxy:\n\n`litellm_settings:\n drop_params: true`\n"
|
||||
),
|
||||
llm_provider=self.custom_llm_provider,
|
||||
model=model,
|
||||
)
|
||||
|
||||
def accepts_prompt_cache_key_at(self, api_base: str | None) -> bool:
|
||||
resolved_api_base, _ = self._get_openai_compatible_provider_info(api_base, None)
|
||||
hostname: Final = urlparse(resolved_api_base).hostname or ""
|
||||
return hostname == "mistral.ai" or hostname.endswith(".mistral.ai")
|
||||
|
||||
@staticmethod
|
||||
def takes_prompt_cache_key_value(prompt_cache_key: object) -> bool:
|
||||
return isinstance(prompt_cache_key, str) and prompt_cache_key.isascii()
|
||||
|
||||
def _prompt_cache_key_rejection(self, prompt_cache_key: object, api_base: str | None) -> str | None:
|
||||
if not self.accepts_prompt_cache_key_at(api_base):
|
||||
return "since only the Mistral API (*.mistral.ai) takes it"
|
||||
if not self.takes_prompt_cache_key_value(prompt_cache_key):
|
||||
return "since the Mistral API takes it only as a string of ASCII characters"
|
||||
return None
|
||||
|
||||
def _prompt_cache_key_withheld(self, model: str, prompt_cache_key: object, routing: MistralRequestRouting) -> bool:
|
||||
requested_api_base: Final = routing.get("api_base")
|
||||
api_base: Final = requested_api_base if isinstance(requested_api_base, str) else None
|
||||
allowed_openai_params: Final = routing.get("allowed_openai_params")
|
||||
if isinstance(allowed_openai_params, (list, tuple)) and "prompt_cache_key" in allowed_openai_params:
|
||||
return False
|
||||
rejection: Final = self._prompt_cache_key_rejection(prompt_cache_key, api_base)
|
||||
if rejection is None:
|
||||
return False
|
||||
if litellm.drop_params is True or routing.get("drop_params") is True:
|
||||
return True
|
||||
requested_provider: Final = routing.get("custom_llm_provider")
|
||||
provider: Final = requested_provider if isinstance(requested_provider, str) else self.custom_llm_provider
|
||||
raise UnsupportedParamsError(
|
||||
message=(
|
||||
f"{provider} does not support parameters: ['prompt_cache_key'], for model={model}, {rejection}. "
|
||||
"To send it anyway, pass `allowed_openai_params=['prompt_cache_key']`. To drop it, set "
|
||||
"`litellm.drop_params=True` or for proxy:\n\n`litellm_settings:\n drop_params: true`\n"
|
||||
),
|
||||
llm_provider=provider,
|
||||
model=model,
|
||||
)
|
||||
|
||||
def _get_openai_compatible_provider_info(self, api_base: str | None, api_key: str | None) -> tuple[str, str | None]:
|
||||
# mistral is openai compatible, we just need to set this to custom_openai and have the api_base be https://api.mistral.ai
|
||||
api_base = (
|
||||
|
|
@ -359,24 +456,31 @@ class MistralConfig(OpenAIGPTConfig):
|
|||
messages = self._handle_message_with_file(messages)
|
||||
return messages
|
||||
|
||||
@staticmethod
|
||||
def _with_mistral_file_id(block: Mapping[str, object]) -> Mapping[str, object]:
|
||||
file_part: Final = block.get("file") if block.get("type") == "file" else None
|
||||
file_id: Final = (
|
||||
_MISTRAL_FILE_PART.validate_python(file_part).get("file_id") if isinstance(file_part, Mapping) else None
|
||||
)
|
||||
if not file_id:
|
||||
return block
|
||||
return {**{key: value for key, value in block.items() if key != "file"}, "file_id": file_id}
|
||||
|
||||
@classmethod
|
||||
def _with_mistral_file_ids(cls, message: AllMessageValues) -> AllMessageValues:
|
||||
content: Final = message.get("content")
|
||||
if not isinstance(content, list):
|
||||
return message
|
||||
blocks: Final[tuple[Mapping[str, object], ...]] = tuple(content)
|
||||
return cast( # cast-ok: only file blocks change, to the file_id shape the Mistral API takes
|
||||
AllMessageValues, {**message, "content": [cls._with_mistral_file_id(block) for block in blocks]}
|
||||
)
|
||||
|
||||
def _handle_message_with_file(self, messages: list[AllMessageValues]) -> list[AllMessageValues]:
|
||||
"""
|
||||
Mistral API supports only 'file_id' in message content with type 'file'.
|
||||
"""
|
||||
for m in messages:
|
||||
_content_block = m.get("content")
|
||||
if _content_block and isinstance(_content_block, list):
|
||||
if any(c.get("type") == "file" for c in _content_block):
|
||||
# If file content is present, we get file_id from 'file' attribute of content block
|
||||
# then replace 'file' with 'file_id' and assign the value of 'file_id' attribute to it.
|
||||
file_contents = [c for c in _content_block if c.get("type") == "file"]
|
||||
for file_content in file_contents:
|
||||
file_id = file_content.get("file", {}).get("file_id")
|
||||
if file_id:
|
||||
# Replace 'file' with 'file_id'
|
||||
file_content["file_id"] = file_id # pyright: ignore[reportGeneralTypeIssues] # legacy in-place rewrite of the block shape
|
||||
file_content.pop("file", None)
|
||||
return messages
|
||||
return [self._with_mistral_file_ids(m) for m in messages]
|
||||
|
||||
def _add_reasoning_system_prompt_if_needed(
|
||||
self, messages: list[AllMessageValues], optional_params: dict
|
||||
|
|
@ -794,7 +898,16 @@ class MistralConfig(OpenAIGPTConfig):
|
|||
if "magistral" in model.lower() and optional_params.get("_add_reasoning_prompt", False):
|
||||
messages = self._add_reasoning_system_prompt_if_needed(messages, optional_params)
|
||||
|
||||
upstream_params: Final = {key: value for key, value in optional_params.items() if key != "client_metadata"}
|
||||
cache_key_params: Final = _MISTRAL_CACHE_KEY_PARAMS.validate_python(optional_params)
|
||||
withheld_params: Final = (
|
||||
("client_metadata", "prompt_cache_key")
|
||||
if "prompt_cache_key" in cache_key_params
|
||||
and self._prompt_cache_key_withheld(
|
||||
model, cache_key_params["prompt_cache_key"], _MISTRAL_REQUEST_ROUTING.validate_python(litellm_params)
|
||||
)
|
||||
else ("client_metadata",)
|
||||
)
|
||||
upstream_params: Final = {key: value for key, value in optional_params.items() if key not in withheld_params}
|
||||
|
||||
# Call parent transform_request which handles _transform_messages
|
||||
return super().transform_request(
|
||||
|
|
|
|||
|
|
@ -30,6 +30,7 @@ if TYPE_CHECKING:
|
|||
from aiohttp import ClientSession
|
||||
|
||||
import litellm
|
||||
from litellm.exceptions import error_response_without_credentials
|
||||
from litellm.litellm_core_utils.completion_timeout import CompletionTimeout
|
||||
from litellm.litellm_core_utils.core_helpers import as_str_mapping, decoded_json_object
|
||||
from litellm.litellm_core_utils.token_counter import token_counter, usage_tokenizer_model
|
||||
|
|
@ -169,6 +170,14 @@ def openai_sdk_key_params(api_key: str | None) -> OpenAISdkKeyParams:
|
|||
return {"api_key": sdk_api_key, "admin_api_key": sdk_api_key}
|
||||
|
||||
|
||||
def upstream_error_response(error: Exception) -> httpx.Response | None:
|
||||
if not isinstance(error, (openai.APIStatusError, httpx.HTTPStatusError)):
|
||||
return None
|
||||
return error_response_without_credentials(
|
||||
error.response.status_code, error.response, "https://docs.litellm.ai/docs"
|
||||
)
|
||||
|
||||
|
||||
_DECODED_JSON: Final = TypeAdapter(object)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -10,7 +10,7 @@ from litellm.types.llms.openai import AllMessageValues, OpenAITextCompletionUser
|
|||
from litellm.types.utils import LlmProviders, ModelResponse, TextCompletionResponse
|
||||
from litellm.utils import ProviderConfigManager
|
||||
|
||||
from ..common_utils import BaseOpenAILLM, OpenAIError
|
||||
from ..common_utils import BaseOpenAILLM, OpenAIError, upstream_error_response
|
||||
from ..workload_identity import build_async_openai_client, build_openai_client
|
||||
from .transformation import OpenAITextCompletionConfig
|
||||
|
||||
|
|
@ -164,7 +164,12 @@ class OpenAITextCompletion(BaseLLM):
|
|||
error_response: Final[object] = getattr(e, "response", None)
|
||||
if error_headers is None and error_response:
|
||||
error_headers = getattr(error_response, "headers", None)
|
||||
raise OpenAIError(status_code=status_code, message=error_text, headers=error_headers)
|
||||
raise OpenAIError(
|
||||
status_code=status_code,
|
||||
message=error_text,
|
||||
response=upstream_error_response(e),
|
||||
headers=error_headers,
|
||||
)
|
||||
|
||||
async def acompletion(
|
||||
self,
|
||||
|
|
@ -220,7 +225,12 @@ class OpenAITextCompletion(BaseLLM):
|
|||
error_response: Final[object] = getattr(e, "response", None)
|
||||
if error_headers is None and error_response:
|
||||
error_headers = getattr(error_response, "headers", None)
|
||||
raise OpenAIError(status_code=status_code, message=error_text, headers=error_headers)
|
||||
raise OpenAIError(
|
||||
status_code=status_code,
|
||||
message=error_text,
|
||||
response=upstream_error_response(e),
|
||||
headers=error_headers,
|
||||
)
|
||||
|
||||
def streaming(
|
||||
self,
|
||||
|
|
@ -261,7 +271,12 @@ class OpenAITextCompletion(BaseLLM):
|
|||
error_response: object = getattr(e, "response", None)
|
||||
if error_headers is None and error_response:
|
||||
error_headers = getattr(error_response, "headers", None)
|
||||
raise OpenAIError(status_code=status_code, message=error_text, headers=error_headers)
|
||||
raise OpenAIError(
|
||||
status_code=status_code,
|
||||
message=error_text,
|
||||
response=upstream_error_response(e),
|
||||
headers=error_headers,
|
||||
)
|
||||
streamwrapper: Final = CustomStreamWrapper(
|
||||
completion_stream=response,
|
||||
model=model,
|
||||
|
|
|
|||
|
|
@ -14,7 +14,7 @@ from litellm.utils import ProviderConfigManager
|
|||
|
||||
from ...base_llm.image_variations.transformation import BaseImageVariationConfig
|
||||
from ...custom_httpx.llm_http_handler import LiteLLMLoggingObj
|
||||
from ..common_utils import OpenAIError
|
||||
from ..common_utils import OpenAIError, upstream_error_response
|
||||
from ..workload_identity import build_async_openai_client, build_openai_client
|
||||
|
||||
|
||||
|
|
@ -129,7 +129,12 @@ class OpenAIImageVariationsHandler:
|
|||
error_response: Final[object] = getattr(e, "response", None)
|
||||
if error_headers is None and error_response:
|
||||
error_headers = getattr(error_response, "headers", None)
|
||||
raise OpenAIError(status_code=status_code, message=error_text, headers=error_headers)
|
||||
raise OpenAIError(
|
||||
status_code=status_code,
|
||||
message=error_text,
|
||||
response=upstream_error_response(e),
|
||||
headers=error_headers,
|
||||
)
|
||||
|
||||
def image_variations(
|
||||
self,
|
||||
|
|
@ -245,4 +250,9 @@ class OpenAIImageVariationsHandler:
|
|||
error_response: Final[object] = getattr(e, "response", None)
|
||||
if error_headers is None and error_response:
|
||||
error_headers = getattr(error_response, "headers", None)
|
||||
raise OpenAIError(status_code=status_code, message=error_text, headers=error_headers)
|
||||
raise OpenAIError(
|
||||
status_code=status_code,
|
||||
message=error_text,
|
||||
response=upstream_error_response(e),
|
||||
headers=error_headers,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -27,6 +27,7 @@ import litellm
|
|||
from litellm import LlmProviders
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.constants import DEFAULT_MAX_RETRIES
|
||||
from litellm.exceptions import UnsupportedParamsError
|
||||
from litellm.files.types import FileContentStreamingResult
|
||||
from litellm.litellm_core_utils.core_helpers import set_provider_response_headers_in_hidden_params
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
|
|
@ -65,6 +66,7 @@ from .common_utils import (
|
|||
is_output_token_limit_error,
|
||||
model_dump_provider_response,
|
||||
openai_sdk_key_params,
|
||||
upstream_error_response,
|
||||
)
|
||||
from .workload_identity import (
|
||||
build_async_openai_client,
|
||||
|
|
@ -971,8 +973,8 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM):
|
|||
litellm.remove_index_from_tool_calls(messages=messages)
|
||||
else:
|
||||
raise e
|
||||
except OpenAIError as e:
|
||||
raise e
|
||||
except (OpenAIError, UnsupportedParamsError):
|
||||
raise
|
||||
except Exception as e:
|
||||
status_code: Final = getattr(e, "status_code", 500)
|
||||
error_headers = getattr(e, "headers", None)
|
||||
|
|
@ -984,6 +986,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM):
|
|||
raise OpenAIError(
|
||||
status_code=status_code,
|
||||
message=error_text,
|
||||
response=upstream_error_response(e),
|
||||
headers=error_headers,
|
||||
body=error_body,
|
||||
)
|
||||
|
|
@ -1111,6 +1114,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM):
|
|||
raise OpenAIError(
|
||||
status_code=status_code,
|
||||
message=message,
|
||||
response=upstream_error_response(e),
|
||||
headers=error_headers,
|
||||
body=exception_body,
|
||||
)
|
||||
|
|
@ -1283,6 +1287,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM):
|
|||
raise OpenAIError(
|
||||
status_code=getattr(e, "status_code", 500),
|
||||
message=str(e),
|
||||
response=upstream_error_response(e),
|
||||
headers=error_headers,
|
||||
body=exception_body,
|
||||
)
|
||||
|
|
@ -1410,7 +1415,12 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM):
|
|||
error_response: Final[object] = getattr(e, "response", None)
|
||||
if error_headers is None and error_response:
|
||||
error_headers = getattr(error_response, "headers", None)
|
||||
raise OpenAIError(status_code=status_code, message=error_text, headers=error_headers)
|
||||
raise OpenAIError(
|
||||
status_code=status_code,
|
||||
message=error_text,
|
||||
response=upstream_error_response(e),
|
||||
headers=error_headers,
|
||||
)
|
||||
|
||||
def embedding(
|
||||
self,
|
||||
|
|
@ -1500,7 +1510,12 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM):
|
|||
error_response: Final[object] = getattr(e, "response", None)
|
||||
if error_headers is None and error_response:
|
||||
error_headers = getattr(error_response, "headers", None)
|
||||
raise OpenAIError(status_code=status_code, message=error_text, headers=error_headers)
|
||||
raise OpenAIError(
|
||||
status_code=status_code,
|
||||
message=error_text,
|
||||
response=upstream_error_response(e),
|
||||
headers=error_headers,
|
||||
)
|
||||
|
||||
async def aimage_generation(
|
||||
self,
|
||||
|
|
@ -1667,7 +1682,11 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM):
|
|||
original_response=str(e),
|
||||
)
|
||||
if hasattr(e, "status_code"):
|
||||
raise OpenAIError(status_code=getattr(e, "status_code", 500), message=str(e))
|
||||
raise OpenAIError(
|
||||
status_code=getattr(e, "status_code", 500),
|
||||
message=str(e),
|
||||
response=upstream_error_response(e),
|
||||
)
|
||||
else:
|
||||
raise OpenAIError(status_code=500, message=str(e))
|
||||
|
||||
|
|
|
|||
|
|
@ -39,8 +39,8 @@ class VertexAIRealtimeConfig(GeminiRealtimeConfig):
|
|||
(they require async I/O) and injected at construction time.
|
||||
"""
|
||||
|
||||
def __init__(self, access_token: str, project: str, location: str) -> None:
|
||||
super().__init__()
|
||||
def __init__(self, access_token: str, project: str, location: str, drop_params: bool = False) -> None:
|
||||
super().__init__(drop_params=drop_params)
|
||||
self._access_token = access_token
|
||||
self._project = project
|
||||
self._location = location
|
||||
|
|
@ -51,6 +51,9 @@ class VertexAIRealtimeConfig(GeminiRealtimeConfig):
|
|||
def _rejects_vad_tuning_with_detection_off(self) -> bool:
|
||||
return False
|
||||
|
||||
def _rejects_undocumented_voices(self, model: str) -> bool:
|
||||
return True
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# URL
|
||||
# ------------------------------------------------------------------
|
||||
|
|
@ -163,7 +166,7 @@ class VertexAIRealtimeConfig(GeminiRealtimeConfig):
|
|||
# settings would be silently dropped because ``map_openai_params`` only
|
||||
# recognises the flat OpenAI-beta key names.
|
||||
session_params = self._normalize_session_payload_for_mapping(session_params)
|
||||
setup_config: Final = self.map_openai_params(optional_params={}, non_default_params=session_params)
|
||||
setup_config: Final = self.map_openai_params(optional_params={}, non_default_params=session_params, model=model)
|
||||
|
||||
# Use full Vertex AI model path
|
||||
setup_config["model"] = self._vertex_model_path(model)
|
||||
|
|
@ -177,15 +180,16 @@ class VertexAIRealtimeConfig(GeminiRealtimeConfig):
|
|||
# ``silence_duration_ms``). ``map_automatic_turn_detection`` sets
|
||||
# ``disabled=True`` whenever ``create_response`` is absent or
|
||||
# ``False``. Force ``disabled=False`` only when the client did
|
||||
# not explicitly request ``create_response: False`` — that path
|
||||
# not explicitly request ``create_response: False`` or
|
||||
# ``turn_detection: null`` — the first of those paths
|
||||
# is how transcription guardrails suppress automatic responses,
|
||||
# and overriding it here would silently bypass the guardrail.
|
||||
# Vertex Live has no "VAD on, no auto-response" mode, so callers
|
||||
# that need that behaviour must accept that VAD is off.
|
||||
client_turn_detection: Final = session_params.get("turn_detection")
|
||||
client_disabled_auto_response: Final = (
|
||||
isinstance(client_turn_detection, dict) and client_turn_detection.get("create_response") is False
|
||||
)
|
||||
"turn_detection" in session_params and client_turn_detection is None
|
||||
) or (isinstance(client_turn_detection, dict) and client_turn_detection.get("create_response") is False)
|
||||
realtime_input_config: Final = setup_config.setdefault("realtimeInputConfig", {})
|
||||
automatic_detection: Final = realtime_input_config.setdefault("automaticActivityDetection", {})
|
||||
if not client_disabled_auto_response:
|
||||
|
|
@ -261,4 +265,5 @@ def vertex_realtime_config(
|
|||
access_token=access_token,
|
||||
project=project,
|
||||
location=VertexBase.get_vertex_region(vertex_region=location, model=model),
|
||||
drop_params=drop_params,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -3,7 +3,7 @@ from typing import Final
|
|||
|
||||
from litellm.exceptions import UnsupportedParamsError
|
||||
from litellm.litellm_core_utils.core_helpers import drops_unsupported_params
|
||||
from litellm.llms.anthropic.common_utils import AnthropicModelInfo
|
||||
from litellm.llms.anthropic.common_utils import AnthropicModelInfo, context_management_betas
|
||||
from litellm.llms.anthropic.messages.native.transformation import (
|
||||
AnthropicMessagesConfig,
|
||||
messages_carry_output_config,
|
||||
|
|
@ -171,28 +171,7 @@ class VertexAIPartnerModelsAnthropicMessagesConfig(AnthropicMessagesConfig, Vert
|
|||
if existing_beta:
|
||||
beta_values.update(b.strip() for b in existing_beta.split(","))
|
||||
|
||||
# Check for context management
|
||||
context_management_param: Final = optional_params.get("context_management")
|
||||
if context_management_param is not None:
|
||||
# Check edits array for compact_20260112 type
|
||||
edits: Final = context_management_param.get("edits", [])
|
||||
has_compact = False
|
||||
has_other = False
|
||||
|
||||
for edit in edits:
|
||||
edit_type = edit.get("type", "")
|
||||
if edit_type == "compact_20260112":
|
||||
has_compact = True
|
||||
else:
|
||||
has_other = True
|
||||
|
||||
# Add compact header if any compact edits exist
|
||||
if has_compact:
|
||||
beta_values.add(ANTHROPIC_BETA_HEADER_VALUES.COMPACT_2026_01_12.value)
|
||||
|
||||
# Add context management header if any other edits exist
|
||||
if has_other:
|
||||
beta_values.add(ANTHROPIC_BETA_HEADER_VALUES.CONTEXT_MANAGEMENT_2025_06_27.value)
|
||||
beta_values.update(context_management_betas(optional_params.get("context_management")))
|
||||
|
||||
# Check for web search tool
|
||||
for tool in tools:
|
||||
|
|
|
|||
|
|
@ -1,11 +1,43 @@
|
|||
from collections.abc import Mapping
|
||||
from typing import Final
|
||||
|
||||
from typing_extensions import TypeIs # noqa: TID251 # TypeIs reaches typing only on 3.13
|
||||
|
||||
import litellm
|
||||
from litellm.litellm_core_utils.prompt_templates.factory import response_schema_prompt
|
||||
from litellm.secret_managers.main import get_secret_str
|
||||
from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolParam
|
||||
from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolParam, ChatCompletionUserMessage
|
||||
from litellm.utils import ModelInfoLookupMiss, lookup_model_info_for_registration
|
||||
|
||||
from ...openai.chat.gpt_transformation import OpenAIGPTConfig
|
||||
|
||||
ZAI_API_BASE: Final = "https://api.z.ai/api/paas/v4"
|
||||
TOKEN_LIMIT_PARAMS: Final = frozenset({"max_tokens", "max_completion_tokens"})
|
||||
|
||||
|
||||
def _is_json_object(
|
||||
value: object,
|
||||
) -> TypeIs[Mapping[str, object]]: # guard-ok: cost-map rows and request params are JSON
|
||||
return isinstance(value, Mapping)
|
||||
|
||||
|
||||
def _json_field(value: object, key: str) -> object:
|
||||
return value.get(key) if _is_json_object(value) else None
|
||||
|
||||
|
||||
def _exceeds_output_limit(model: str, max_tokens: object) -> bool:
|
||||
info: Final = lookup_model_info_for_registration(model=model, custom_llm_provider="zai")
|
||||
limit: Final = None if isinstance(info, ModelInfoLookupMiss) else info.get("max_output_tokens")
|
||||
return isinstance(max_tokens, int) and isinstance(limit, int) and max_tokens > limit
|
||||
|
||||
|
||||
def _schema_without_native_support(model: str, response_format: object) -> Mapping[str, object] | None:
|
||||
if _json_field(response_format, "type") != "json_schema":
|
||||
return None
|
||||
if litellm.supports_response_schema(model=model, custom_llm_provider="zai"):
|
||||
return None
|
||||
schema: Final = _json_field(_json_field(response_format, "json_schema"), "schema")
|
||||
return schema if _is_json_object(schema) else None
|
||||
|
||||
|
||||
class ZAIChatConfig(OpenAIGPTConfig):
|
||||
|
|
@ -13,6 +45,10 @@ class ZAIChatConfig(OpenAIGPTConfig):
|
|||
def custom_llm_provider(self) -> str | None:
|
||||
return "zai"
|
||||
|
||||
@property
|
||||
def prompts_unsupported_json_schema(self) -> bool:
|
||||
return True
|
||||
|
||||
def _get_openai_compatible_provider_info(
|
||||
self, api_base: str | None, api_key: str | None
|
||||
) -> tuple[str | None, str | None]:
|
||||
|
|
@ -40,9 +76,11 @@ class ZAIChatConfig(OpenAIGPTConfig):
|
|||
# GLM/ZAI supports cache_control, so return messages and tools unchanged
|
||||
return messages, tools
|
||||
|
||||
def get_supported_openai_params(self, model: str) -> list:
|
||||
def get_supported_openai_params(self, model: str) -> list[str]:
|
||||
base_params: Final = [
|
||||
"max_tokens",
|
||||
"max_completion_tokens",
|
||||
"response_format",
|
||||
"stream",
|
||||
"stream_options",
|
||||
"temperature",
|
||||
|
|
@ -52,8 +90,6 @@ class ZAIChatConfig(OpenAIGPTConfig):
|
|||
"tool_choice",
|
||||
]
|
||||
|
||||
import litellm
|
||||
|
||||
try:
|
||||
if litellm.supports_reasoning(model=model, custom_llm_provider=self.custom_llm_provider):
|
||||
base_params.append("thinking")
|
||||
|
|
@ -61,3 +97,53 @@ class ZAIChatConfig(OpenAIGPTConfig):
|
|||
pass
|
||||
|
||||
return base_params
|
||||
|
||||
def map_openai_params(
|
||||
self,
|
||||
non_default_params: Mapping[str, object],
|
||||
optional_params: Mapping[str, object],
|
||||
model: str,
|
||||
drop_params: bool,
|
||||
) -> dict[str, object]: # mutable-ok: request pipeline adds parameters to the returned body
|
||||
supported_params: Final = self.get_supported_openai_params(model)
|
||||
sent: Final = {param: value for param, value in non_default_params.items() if param in supported_params}
|
||||
max_completion_tokens: Final = sent.get("max_completion_tokens")
|
||||
falls_back_to_max_tokens: Final = max_completion_tokens is None or (
|
||||
(litellm.drop_params or drop_params) and _exceeds_output_limit(model, max_completion_tokens)
|
||||
)
|
||||
max_tokens: Final = sent.get("max_tokens") if falls_back_to_max_tokens else max_completion_tokens
|
||||
return {
|
||||
**optional_params,
|
||||
**{param: value for param, value in sent.items() if param not in TOKEN_LIMIT_PARAMS},
|
||||
**({} if max_tokens is None else {"max_tokens": max_tokens}),
|
||||
}
|
||||
|
||||
def transform_request(
|
||||
self,
|
||||
model: str,
|
||||
messages: list[AllMessageValues], # mutable-ok: OpenAIGPTConfig.transform_request signature
|
||||
optional_params: dict[str, object], # mutable-ok: OpenAIGPTConfig.transform_request signature
|
||||
litellm_params: dict[str, object], # mutable-ok: OpenAIGPTConfig.transform_request signature
|
||||
headers: dict[str, str], # mutable-ok: OpenAIGPTConfig.transform_request signature
|
||||
) -> dict[str, object]: # mutable-ok: inherited contract returns the request body
|
||||
schema: Final = (
|
||||
_schema_without_native_support(model, optional_params.get("response_format"))
|
||||
if self.prompts_unsupported_json_schema
|
||||
else None
|
||||
)
|
||||
schema_message: Final = (
|
||||
()
|
||||
if schema is None
|
||||
else (
|
||||
ChatCompletionUserMessage(
|
||||
role="user", content=response_schema_prompt(model=model, response_schema=schema, llm_provider="zai")
|
||||
),
|
||||
)
|
||||
)
|
||||
return super().transform_request(
|
||||
model,
|
||||
[*messages, *schema_message],
|
||||
optional_params if schema is None else {**optional_params, "response_format": {"type": "json_object"}},
|
||||
litellm_params,
|
||||
headers,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -6199,6 +6199,7 @@ def completion(
|
|||
PROVIDER_AFFINITY_HEADER_KWARG_KEY,
|
||||
"fireworks_forward_user_id",
|
||||
"responses_api_version",
|
||||
"allowed_openai_params",
|
||||
)
|
||||
if key in kwargs
|
||||
},
|
||||
|
|
|
|||
File diff suppressed because it is too large
Load diff
|
|
@ -8,7 +8,7 @@ from typing import Final
|
|||
|
||||
from fastapi import APIRouter, Depends, HTTPException, Request, Response, status
|
||||
from fastapi.responses import JSONResponse
|
||||
from pydantic import JsonValue, ValidationError
|
||||
from pydantic import JsonValue, TypeAdapter, ValidationError
|
||||
from typing_extensions import ReadOnly, TypedDict
|
||||
|
||||
import litellm
|
||||
|
|
@ -54,6 +54,7 @@ from litellm.proxy.litellm_pre_call_utils import (
|
|||
should_auto_drop_params_for_agentic_cli,
|
||||
)
|
||||
from litellm.router import Router
|
||||
from litellm.types.utils import TokenCountResponse
|
||||
|
||||
router: Final = APIRouter()
|
||||
|
||||
|
|
@ -334,6 +335,15 @@ def _count_tokens_body_options(data: Mapping[str, JsonValue], auto_drop_params:
|
|||
}
|
||||
|
||||
|
||||
_PROVIDER_COUNT_REPLY: Final[TypeAdapter[Mapping[str, object] | None]] = TypeAdapter(Mapping[str, object] | None)
|
||||
|
||||
|
||||
def _provider_context_management(response: TokenCountResponse) -> Mapping[str, object]:
|
||||
original: Final = _PROVIDER_COUNT_REPLY.validate_python(response.original_response)
|
||||
context_management: Final = original.get("context_management") if original is not None else None
|
||||
return {} if context_management is None else {"context_management": context_management}
|
||||
|
||||
|
||||
@router.post(
|
||||
"/v1/messages/count_tokens",
|
||||
tags=["[beta] Anthropic Messages Token Counting"],
|
||||
|
|
@ -413,7 +423,11 @@ async def count_tokens(
|
|||
log_llm_api_exception(e, litellm_call_id)
|
||||
return _anthropic_error_json_response(proxy_exception_from_route_error(e), request)
|
||||
else:
|
||||
return {"input_tokens": token_response.total_tokens, **count_source_fields(token_response)}
|
||||
return {
|
||||
"input_tokens": token_response.total_tokens,
|
||||
**_provider_context_management(token_response),
|
||||
**count_source_fields(token_response),
|
||||
}
|
||||
|
||||
|
||||
@router.post(
|
||||
|
|
|
|||
|
|
@ -446,7 +446,7 @@ def _request_has_trace_id(request_body: Mapping[str, object], request: Request,
|
|||
from litellm.proxy.litellm_pre_call_utils import (
|
||||
_trace_id_from_traceparent, # pyright: ignore[reportPrivateUsage] # reuse W3C traceparent parser
|
||||
get_chain_id_from_headers,
|
||||
metadata_variable_name_for_route,
|
||||
metadata_variable_name_for_request,
|
||||
)
|
||||
|
||||
if get_chain_id_from_headers(dict(request.headers)):
|
||||
|
|
@ -460,7 +460,7 @@ def _request_has_trace_id(request_body: Mapping[str, object], request: Request,
|
|||
traceparent: Final[str | None] = request.headers.get("traceparent")
|
||||
if isinstance(traceparent, str) and _trace_id_from_traceparent(traceparent):
|
||||
return True
|
||||
metadata_variable_name: Final = metadata_variable_name_for_route(get_request_route(request))
|
||||
metadata_variable_name: Final = metadata_variable_name_for_request(get_request_route(request), request_body)
|
||||
selected_metadata: Final[object] = request_body.get(metadata_variable_name)
|
||||
if isinstance(selected_metadata, dict):
|
||||
selected_metadata_dict: Final = _TRACE_ID_METADATA_ADAPTER.validate_python(selected_metadata)
|
||||
|
|
@ -1367,11 +1367,15 @@ async def common_checks(
|
|||
if not skip_all_budget_checks:
|
||||
# Key metadata.tags are injected into request_body here so the tag budget
|
||||
# check can read them; this mutation must run before the gathered checks.
|
||||
from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup
|
||||
from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup, metadata_variable_name_for_request
|
||||
|
||||
if valid_token is not None:
|
||||
for tag_target in LiteLLMProxyRequestSetup.pre_auth_tag_merge_targets(request, request_body):
|
||||
LiteLLMProxyRequestSetup.pre_seed_litellm_metadata_for_route(request_data=tag_target, route=route)
|
||||
LiteLLMProxyRequestSetup.merge_caller_body_tags(
|
||||
request_data=tag_target,
|
||||
metadata_variable_name=metadata_variable_name_for_request(route, tag_target),
|
||||
)
|
||||
LiteLLMProxyRequestSetup.apply_key_tags_pre_auth(request_data=tag_target, user_api_key_dict=valid_token)
|
||||
LiteLLMProxyRequestSetup.apply_authenticated_baggage_tags(request, request_body)
|
||||
|
||||
|
|
|
|||
|
|
@ -14,14 +14,16 @@ from litellm.litellm_core_utils.credential_accessor import CredentialAccessor
|
|||
from litellm.llms.azure.ptu_capacity import is_str_keyed_mapping
|
||||
from litellm.proxy._types import SpecialModelNames, UserAPIKeyAuth
|
||||
from litellm.repositories.object_permission_repository import ObjectPermissionRepository
|
||||
from litellm.responses.main import strip_responses_routing_prefix
|
||||
from litellm.router import Router
|
||||
from litellm.router_utils.fallback_event_handlers import get_fallback_model_group
|
||||
from litellm.router_utils.pattern_match_deployments import PatternMatchRouter
|
||||
from litellm.types.router import CredentialLiteLLMParams, LiteLLM_Params
|
||||
from litellm.types.utils import LlmProviders
|
||||
from litellm.utils import ProviderConfigManager, get_valid_models
|
||||
from litellm.utils import ProviderConfigManager, get_valid_models, potential_model_names_in_lookup_order
|
||||
|
||||
LLM_PROVIDER_PREFIXES: Final = frozenset(llm_provider.value for llm_provider in LlmProviders)
|
||||
_PROVIDERS_SENDING_THE_DEPLOYMENT_PATH_AS_THE_MODEL_ID: Final = frozenset({LlmProviders.TOGETHER_AI.value})
|
||||
|
||||
_CREDENTIAL_LITELLM_PARAM_FIELDS = set(CredentialLiteLLMParams.model_fields)
|
||||
|
||||
|
|
@ -331,6 +333,25 @@ def _routes_back_to(
|
|||
)
|
||||
|
||||
|
||||
def _resolves_to_catalog_model(model: str, provider: str, provider_model_ids: frozenset[str]) -> bool:
|
||||
provider_id_prefix: Final = f"{provider}/"
|
||||
resolved_model: Final = (
|
||||
f"{provider_id_prefix}{strip_responses_routing_prefix(model.removeprefix(provider_id_prefix))}"
|
||||
)
|
||||
return any(
|
||||
name.removeprefix(provider_id_prefix) in provider_model_ids
|
||||
for name in potential_model_names_in_lookup_order(resolved_model, provider)
|
||||
)
|
||||
|
||||
|
||||
def _routes_to_catalog_model(
|
||||
routes_to_model: Callable[[str, str], bool], provider: str, provider_model_ids: frozenset[str]
|
||||
) -> Callable[[str, str], bool]:
|
||||
return lambda listed_id, model: (
|
||||
routes_to_model(listed_id, model) and _resolves_to_catalog_model(model, provider, provider_model_ids)
|
||||
)
|
||||
|
||||
|
||||
def _ids_routing_to_catalog_models(
|
||||
name_pattern: str,
|
||||
deployment_pattern: str,
|
||||
|
|
@ -344,8 +365,15 @@ def _ids_routing_to_catalog_models(
|
|||
return ()
|
||||
provider_id_prefix: Final = f"{provider}/"
|
||||
provider_model_ids: Final = tuple(catalog_id.removeprefix(provider_id_prefix) for catalog_id in catalog_ids)
|
||||
deployment_model_pattern_without_provider: Final = deployment_pattern.removeprefix(provider_id_prefix)
|
||||
catalog_pattern: Final = _without_routing_segments(
|
||||
deployment_pattern.removeprefix(provider_id_prefix), provider, name_pattern, provider_model_ids
|
||||
deployment_model_pattern_without_provider, provider, name_pattern, provider_model_ids
|
||||
)
|
||||
routing_prefix: Final = deployment_model_pattern_without_provider.removesuffix(catalog_pattern)
|
||||
reaches_catalog_model: Final = (
|
||||
_routes_to_catalog_model(routes_to_model, provider, frozenset(provider_model_ids))
|
||||
if routing_prefix and provider in _PROVIDERS_SENDING_THE_DEPLOYMENT_PATH_AS_THE_MODEL_ID
|
||||
else routes_to_model
|
||||
)
|
||||
deployment_model_pattern: Final = f"{provider_id_prefix}{catalog_pattern}"
|
||||
deployment_regex: Final = re.compile(PatternMatchRouter().pattern_to_regex(deployment_model_pattern))
|
||||
|
|
@ -369,7 +397,7 @@ def _ids_routing_to_catalog_models(
|
|||
for listed_id, catalog_id in listed_ids
|
||||
if listed_id
|
||||
and _routes_back_to(
|
||||
listed_id, name_regex, (deployment_model_pattern, catalog_id), deployment_pattern, routes_to_model
|
||||
listed_id, name_regex, (deployment_model_pattern, catalog_id), deployment_pattern, reaches_catalog_model
|
||||
)
|
||||
)
|
||||
)
|
||||
|
|
|
|||
|
|
@ -2401,9 +2401,10 @@ class ProxyBaseLLMRequestProcessing:
|
|||
version=version,
|
||||
proxy_config=proxy_config,
|
||||
)
|
||||
request_fields: Final = _REQUEST_FIELDS.validate_python(self.data)
|
||||
metadata_variable_name: Final = get_metadata_variable_name(request, request_fields)
|
||||
if logging_auth is not None:
|
||||
metadata_variable_name: Final = get_metadata_variable_name(request)
|
||||
request_metadata: Final = self.data.get(metadata_variable_name)
|
||||
request_metadata: Final = request_fields.get(metadata_variable_name)
|
||||
if isinstance(request_metadata, dict):
|
||||
self.data[metadata_variable_name] = {
|
||||
**TypeAdapter(dict[str, object]).validate_python(request_metadata),
|
||||
|
|
@ -2431,12 +2432,11 @@ class ProxyBaseLLMRequestProcessing:
|
|||
|
||||
# Store queue time in metadata after add_litellm_data_to_request to ensure it's preserved
|
||||
if queue_time_seconds is not None:
|
||||
_metadata_variable_name: Final = get_metadata_variable_name(request)
|
||||
if _metadata_variable_name not in self.data:
|
||||
self.data[_metadata_variable_name] = {}
|
||||
if not isinstance(self.data[_metadata_variable_name], dict):
|
||||
self.data[_metadata_variable_name] = {}
|
||||
self.data[_metadata_variable_name]["queue_time_seconds"] = queue_time_seconds
|
||||
if metadata_variable_name not in self.data:
|
||||
self.data[metadata_variable_name] = {}
|
||||
if not isinstance(self.data[metadata_variable_name], dict):
|
||||
self.data[metadata_variable_name] = {}
|
||||
self.data[metadata_variable_name]["queue_time_seconds"] = queue_time_seconds
|
||||
|
||||
if isinstance(model, str):
|
||||
reject_url_valued_destination("model", model)
|
||||
|
|
|
|||
|
|
@ -42,7 +42,11 @@ from litellm.caching.redis_batch import (
|
|||
active_request_redis_batch,
|
||||
)
|
||||
from litellm.caching.redis_cache import log_redis_failure
|
||||
from litellm.constants import DYNAMIC_RATE_LIMIT_ERROR_THRESHOLD_PER_MINUTE, INTERNAL_CALL_ORIGIN_METADATA_KEY
|
||||
from litellm.constants import (
|
||||
DYNAMIC_RATE_LIMIT_ERROR_THRESHOLD_PER_MINUTE,
|
||||
INTERNAL_CALL_ORIGIN_METADATA_KEY,
|
||||
NO_MAX_TOKENS_TPM_RESERVATION_FRACTION,
|
||||
)
|
||||
from litellm.cost_calculator import get_usage_object
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.litellm_core_utils.audio_utils.utils import estimate_audio_tokens
|
||||
|
|
@ -508,11 +512,6 @@ REDIS_NODE_HASHTAG_NAME: Final = "all_keys"
|
|||
# *some* output budget; these define that fallback estimate.
|
||||
DEFAULT_MAX_TOKENS_ESTIMATE: Final = 4096
|
||||
DEFAULT_CHARS_PER_TOKEN: Final = 4
|
||||
# Fraction of the available output budget reserved as the upfront floor when
|
||||
# the request omits max_tokens. Applied to both DEFAULT_MAX_TOKENS_ESTIMATE
|
||||
# (baseline floor) and to the smallest configured TPM limit (capped floor for
|
||||
# small per-tenant TPM caps).
|
||||
_TPM_FLOOR_FRACTION: Final = 4
|
||||
# Both embeddings and the Responses API put their prompt in data["input"],
|
||||
# but only embeddings have no output tokens. Every "is this an embedding"
|
||||
# check on data["input"] must exclude these call types, or a Responses call
|
||||
|
|
@ -1003,10 +1002,10 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
per-tenant cap can't be tripped by the floor alone. Returns the
|
||||
baseline floor when no limit is provided.
|
||||
"""
|
||||
baseline: Final = DEFAULT_MAX_TOKENS_ESTIMATE // _TPM_FLOOR_FRACTION
|
||||
baseline: Final = int(DEFAULT_MAX_TOKENS_ESTIMATE * NO_MAX_TOKENS_TPM_RESERVATION_FRACTION)
|
||||
if min_configured_tpm_limit is None:
|
||||
return baseline
|
||||
return min(baseline, max(1, min_configured_tpm_limit // _TPM_FLOOR_FRACTION))
|
||||
return min(baseline, max(1, int(min_configured_tpm_limit * NO_MAX_TOKENS_TPM_RESERVATION_FRACTION)))
|
||||
|
||||
@staticmethod
|
||||
def _is_embedding_request(data: object, call_type: str | None) -> bool:
|
||||
|
|
@ -1138,7 +1137,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
if call_type in RESPONSES_API_CALL_TYPES
|
||||
else base_capped_floor
|
||||
)
|
||||
baseline_floor: Final = DEFAULT_MAX_TOKENS_ESTIMATE // _TPM_FLOOR_FRACTION
|
||||
baseline_floor: Final = int(DEFAULT_MAX_TOKENS_ESTIMATE * NO_MAX_TOKENS_TPM_RESERVATION_FRACTION)
|
||||
is_embedding: Final = PROXY_MaxParallelRequestsHandler_v3._is_embedding_request(data, call_type)
|
||||
if (
|
||||
capped_floor >= baseline_floor
|
||||
|
|
@ -3313,7 +3312,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
if key_agent_id:
|
||||
return key_agent_id
|
||||
metadata: Final = data.get("metadata") or {}
|
||||
return metadata.get("agent_id")
|
||||
litellm_metadata: Final = data.get("litellm_metadata") or {}
|
||||
return metadata.get("agent_id") or litellm_metadata.get("agent_id")
|
||||
|
||||
def _get_session_id_from_data(self, data: dict) -> str | None:
|
||||
"""Extract session_id from request metadata or litellm_session_id."""
|
||||
|
|
|
|||
|
|
@ -37,7 +37,7 @@ from litellm.proxy.db.db_spend_update_writer import (
|
|||
get_llm_router,
|
||||
)
|
||||
from litellm.proxy.hooks.search_free_tier import apply_search_free_tier
|
||||
from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup, metadata_variable_name_for_route
|
||||
from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup, metadata_variable_name_for_request
|
||||
from litellm.proxy.route_llm_request import ProxyModelNotFoundError
|
||||
from litellm.proxy.spend_tracking.spend_counter_batch import post_call_counter_keys, spend_counter_batch_scope
|
||||
from litellm.proxy.spend_tracking.spend_event import (
|
||||
|
|
@ -137,19 +137,34 @@ _CAPTURED_IDENTITY_CALL_TYPES: Final[frozenset[str]] = frozenset(
|
|||
_OBJECT_MAPPING: Final = TypeAdapter(ObjectMapping)
|
||||
|
||||
|
||||
def _proxy_stamped_used_client_oauth_token(
|
||||
request_data: Mapping[str, object], request_route: str | None
|
||||
) -> bool | None:
|
||||
def _proxy_stamped_metadata(request_data: Mapping[str, object], request_route: str | None) -> Mapping[str, object]:
|
||||
proxy_bucket: Final = (
|
||||
get_metadata_variable_name_from_kwargs(request_data)
|
||||
if request_route is None
|
||||
else metadata_variable_name_for_route(request_route)
|
||||
else metadata_variable_name_for_request(request_route, request_data)
|
||||
)
|
||||
proxy_metadata: Final = request_data.get(proxy_bucket)
|
||||
stamped: Final = proxy_metadata.get("used_client_oauth_token") if isinstance(proxy_metadata, dict) else None
|
||||
return _str_keyed_fields(request_data.get(proxy_bucket))
|
||||
|
||||
|
||||
def _proxy_stamped_used_client_oauth_token(
|
||||
request_data: Mapping[str, object], request_route: str | None
|
||||
) -> bool | None:
|
||||
stamped: Final = _proxy_stamped_metadata(request_data, request_route).get("used_client_oauth_token")
|
||||
return stamped if isinstance(stamped, bool) else None
|
||||
|
||||
|
||||
_PROXY_STAMPED_REQUEST_FIELDS: Final = ("requester_ip_address", "user_agent")
|
||||
|
||||
|
||||
def _proxy_stamped_request_fields(
|
||||
request_data: Mapping[str, object], request_route: str | None
|
||||
) -> Mapping[str, object]:
|
||||
if request_route is None:
|
||||
return {}
|
||||
proxy_metadata: Final = _proxy_stamped_metadata(request_data, request_route)
|
||||
return {field: proxy_metadata[field] for field in _PROXY_STAMPED_REQUEST_FIELDS if field in proxy_metadata}
|
||||
|
||||
|
||||
def _str_keyed_fields(value: object) -> Mapping[str, object]:
|
||||
return value if is_str_keyed_mapping(value) else {}
|
||||
|
||||
|
|
@ -455,7 +470,13 @@ class _ProxyDBLogger(CustomLogger):
|
|||
_metadata["used_client_oauth_token"] = _proxy_stamped_used_client_oauth_token(request_fields, request_route)
|
||||
|
||||
existing_metadata: Final[dict] = request_data.get("metadata", None) or {}
|
||||
existing_metadata.update({**_routed_metadata_fallback(request_fields), **_metadata})
|
||||
existing_metadata.update(
|
||||
{
|
||||
**_routed_metadata_fallback(request_fields),
|
||||
**_proxy_stamped_request_fields(request_fields, request_route),
|
||||
**_metadata,
|
||||
}
|
||||
)
|
||||
|
||||
litellm_metadata_bucket: Final = request_data.get("litellm_metadata")
|
||||
litellm_metadata: Final[ObjectMapping] = (
|
||||
|
|
|
|||
|
|
@ -379,6 +379,8 @@ LITELLM_METADATA_ROUTES: Final = (
|
|||
"interactions",
|
||||
)
|
||||
|
||||
CURSOR_CHAT_COMPLETIONS_ROUTE: Final = "/cursor/chat/completions"
|
||||
|
||||
LITELLM_TRACE_CONTROL_METADATA_FIELDS: Final = frozenset(
|
||||
{
|
||||
"mask_input",
|
||||
|
|
@ -636,7 +638,7 @@ def _has_tag_budget(tag: "LiteLLM_TagTable | None") -> bool:
|
|||
return tag.litellm_budget_table.max_budget is not None
|
||||
|
||||
|
||||
def _normalized_metadata_slot(
|
||||
def normalized_metadata_slot(
|
||||
request_data: MutableMapping[str, object], metadata_variable_name: str
|
||||
) -> dict[str, object]:
|
||||
"""Return the request's metadata slot as a dict, normalising it in place first.
|
||||
|
|
@ -653,6 +655,12 @@ def _normalized_metadata_slot(
|
|||
return normalized
|
||||
|
||||
|
||||
def _holds_metadata_object(value: object) -> bool:
|
||||
if value is None or isinstance(value, dict):
|
||||
return True
|
||||
return isinstance(value, str) and isinstance(_DECODED_JSON.validate_python(safe_json_loads(value)), dict)
|
||||
|
||||
|
||||
def _metadata_object_or_error(field: str, value: object) -> Mapping[str, object] | ProxyException:
|
||||
if value is None:
|
||||
return {}
|
||||
|
|
@ -863,7 +871,7 @@ def _strip_router_reserved_metadata(
|
|||
)
|
||||
|
||||
|
||||
def get_metadata_variable_name(request: Request) -> str:
|
||||
def get_metadata_variable_name(request: Request, data: Mapping[str, object] | None = None) -> str:
|
||||
"""
|
||||
Helper to return what the "metadata" field should be called in the request data
|
||||
|
||||
|
|
@ -874,7 +882,10 @@ def get_metadata_variable_name(request: Request) -> str:
|
|||
# Inline imports — auth_utils/route_checks participate in a proxy import cycle.
|
||||
from litellm.proxy.auth.auth_utils import get_request_route # noqa: PLC0415
|
||||
|
||||
return metadata_variable_name_for_route(get_request_route(request))
|
||||
route: Final = get_request_route(request)
|
||||
if data is None:
|
||||
return metadata_variable_name_for_route(route)
|
||||
return metadata_variable_name_for_request(route, data)
|
||||
|
||||
|
||||
_get_metadata_variable_name: Final = get_metadata_variable_name
|
||||
|
|
@ -890,6 +901,25 @@ def metadata_variable_name_for_route(route: str) -> Literal["metadata", "litellm
|
|||
return "metadata"
|
||||
|
||||
|
||||
def is_chat_completions_body(data: Mapping[str, object]) -> bool:
|
||||
messages: Final = data.get("messages")
|
||||
if isinstance(messages, list) and messages:
|
||||
return True
|
||||
return "messages" in data and "input" not in data
|
||||
|
||||
|
||||
def is_cursor_responses_request(route: str, body: Mapping[str, object]) -> bool:
|
||||
return CURSOR_CHAT_COMPLETIONS_ROUTE in route and not is_chat_completions_body(body)
|
||||
|
||||
|
||||
def metadata_variable_name_for_request(
|
||||
route: str, body: Mapping[str, object]
|
||||
) -> Literal["metadata", "litellm_metadata"]:
|
||||
if is_cursor_responses_request(route, body):
|
||||
return "litellm_metadata"
|
||||
return metadata_variable_name_for_route(route)
|
||||
|
||||
|
||||
def _promoted_trace_control_fields(
|
||||
requester_metadata: Mapping[str, object],
|
||||
litellm_metadata: Mapping[str, object],
|
||||
|
|
@ -2309,7 +2339,9 @@ class LiteLLMProxyRequestSetup:
|
|||
``_tag_max_budget_check`` resolve to ``litellm_metadata``, leaving
|
||||
header tags invisible to per-tag budget enforcement.
|
||||
"""
|
||||
if any(metadata_route in route for metadata_route in LITELLM_METADATA_ROUTES):
|
||||
if any(metadata_route in route for metadata_route in LITELLM_METADATA_ROUTES) or is_cursor_responses_request(
|
||||
route, request_data
|
||||
):
|
||||
request_data.setdefault("litellm_metadata", {})
|
||||
|
||||
@staticmethod
|
||||
|
|
@ -2380,7 +2412,7 @@ class LiteLLMProxyRequestSetup:
|
|||
return
|
||||
tags: Final = _STRING_TAGS_ADAPTER.validate_python(raw_tags)
|
||||
slot: Final = metadata_key if metadata_key is not None else get_metadata_variable_name_from_kwargs(request_data)
|
||||
metadata: Final = _normalized_metadata_slot(request_data, slot)
|
||||
metadata: Final = normalized_metadata_slot(request_data, slot)
|
||||
existing: Final = metadata.get("tags")
|
||||
if existing and not isinstance(existing, list):
|
||||
return
|
||||
|
|
@ -2393,6 +2425,34 @@ class LiteLLMProxyRequestSetup:
|
|||
tags=tuple(tag for tag in tags if tag not in sent_tags)
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def merge_caller_body_tags(
|
||||
request_data: MutableMapping[str, object], # mutable-ok: the request body the tag merges write into
|
||||
metadata_variable_name: str | None = None,
|
||||
) -> None:
|
||||
slot: Final = (
|
||||
metadata_variable_name
|
||||
if metadata_variable_name is not None
|
||||
else get_metadata_variable_name_from_kwargs(request_data)
|
||||
)
|
||||
if slot == "metadata" or not _holds_metadata_object(request_data.get(slot)):
|
||||
return
|
||||
raw_caller_metadata: Final = request_data.get("metadata")
|
||||
caller_metadata: Final = (
|
||||
_DECODED_JSON.validate_python(safe_json_loads(raw_caller_metadata))
|
||||
if isinstance(raw_caller_metadata, str)
|
||||
else raw_caller_metadata
|
||||
)
|
||||
caller_tags: Final = caller_metadata.get("tags") if isinstance(caller_metadata, dict) else None
|
||||
if not isinstance(caller_tags, list) or not caller_tags:
|
||||
return
|
||||
metadata: Final = normalized_metadata_slot(request_data, slot)
|
||||
existing_tags: Final = metadata.get("tags")
|
||||
metadata["tags"] = LiteLLMProxyRequestSetup.merge_tags(
|
||||
request_tags=existing_tags if isinstance(existing_tags, list) else None,
|
||||
tags_to_add=caller_tags,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def apply_key_tags_pre_auth(
|
||||
request_data: dict,
|
||||
|
|
@ -2408,7 +2468,7 @@ class LiteLLMProxyRequestSetup:
|
|||
return
|
||||
|
||||
_metadata_variable_name: Final = get_metadata_variable_name_from_kwargs(request_data)
|
||||
metadata: Final = _normalized_metadata_slot(request_data, _metadata_variable_name)
|
||||
metadata: Final = normalized_metadata_slot(request_data, _metadata_variable_name)
|
||||
|
||||
existing_tags: Final = metadata.get("tags")
|
||||
metadata["tags"] = LiteLLMProxyRequestSetup.merge_tags(
|
||||
|
|
@ -2460,7 +2520,7 @@ class LiteLLMProxyRequestSetup:
|
|||
# from (litellm_metadata vs metadata) so the merged tags are visible
|
||||
# to _tag_max_budget_check.
|
||||
_metadata_variable_name: Final = get_metadata_variable_name_from_kwargs(request_data)
|
||||
metadata: Final = _normalized_metadata_slot(request_data, _metadata_variable_name)
|
||||
metadata: Final = normalized_metadata_slot(request_data, _metadata_variable_name)
|
||||
|
||||
existing_tags: Final = metadata.get("tags")
|
||||
metadata["tags"] = LiteLLMProxyRequestSetup.merge_tags(
|
||||
|
|
@ -2637,7 +2697,7 @@ async def add_litellm_data_to_request(
|
|||
}
|
||||
|
||||
safe_add_api_version_from_query_params(data, request)
|
||||
_metadata_variable_name: Final = get_metadata_variable_name(request)
|
||||
_metadata_variable_name: Final = get_metadata_variable_name(request, _REQUEST_FIELDS.validate_python(data))
|
||||
if data.get(_metadata_variable_name, None) is None:
|
||||
data[_metadata_variable_name] = {}
|
||||
|
||||
|
|
@ -2974,17 +3034,7 @@ async def add_litellm_data_to_request(
|
|||
tags_to_add=tags,
|
||||
)
|
||||
|
||||
_caller_body_metadata: Final = data.get("metadata") if _metadata_variable_name != "metadata" else None
|
||||
_caller_body_tags: Final = (
|
||||
_caller_body_metadata.get("tags")
|
||||
if isinstance(_caller_body_metadata, dict) and isinstance(_caller_body_metadata.get("tags"), list)
|
||||
else None
|
||||
)
|
||||
if _caller_body_tags:
|
||||
data[_metadata_variable_name]["tags"] = LiteLLMProxyRequestSetup.merge_tags( # rebind-ok: matches file idiom
|
||||
request_tags=data[_metadata_variable_name].get("tags"),
|
||||
tags_to_add=_caller_body_tags,
|
||||
)
|
||||
LiteLLMProxyRequestSetup.merge_caller_body_tags(request_data=data, metadata_variable_name=_metadata_variable_name)
|
||||
|
||||
# caller_tags: the tags the request's headers and body carry (x-litellm-tags, body "tags", and body
|
||||
# "metadata.tags" / "litellm_metadata.tags" read before the mapped tags are added below),
|
||||
|
|
|
|||
|
|
@ -98,15 +98,11 @@ def _count_change(sent: tuple[object, ...] | None, returned: tuple[object, ...]
|
|||
return (len(sent), len(returned))
|
||||
|
||||
|
||||
def _tool_call_mismatch_reason(
|
||||
sent: tuple[tuple[object, object], ...] | None, returned: tuple[tuple[object, object], ...] | None
|
||||
) -> str | None:
|
||||
if sent == returned:
|
||||
return None
|
||||
def _tool_call_count_mismatch_reason(sent: Sequence[object] | None, returned: Sequence[object] | None) -> str | None:
|
||||
sent_count: Final = len(sent or ())
|
||||
returned_count: Final = len(returned or ())
|
||||
if sent_count == returned_count:
|
||||
return "the legacy hook changed a tool call's name or arguments, which this path cannot write back"
|
||||
return None
|
||||
return f"the legacy hook returned {returned_count} tool calls for a stream that carried {sent_count}"
|
||||
|
||||
|
||||
|
|
@ -205,11 +201,10 @@ class _LegacyHookStreamAdapter(CustomGuardrail):
|
|||
non-streaming hooks, an exception it raises ends the stream through the executor's
|
||||
fail/error classification, and the response it hands back, or the one it changed in place
|
||||
and returned ``None`` for, is re-scanned by the same translation so its texts reach the
|
||||
client through the translation's ended-stream write-back. A
|
||||
replacement whose scanned texts do not line up with the originals, or whose tool calls
|
||||
differ from them, is undeliverable, so the executor releases the original chunks. A stream
|
||||
that carried no text to scan, such as a tool-only Anthropic message, stays deliverable as
|
||||
long as the hook left the tool calls alone."""
|
||||
client through the translation's ended-stream write-back. Tool calls the hook rewrote go to the translation's
|
||||
tool-call write-back like a unified guardrail's. A replacement whose scanned texts or tool
|
||||
calls do not line up with the originals by count is undeliverable, so the executor releases
|
||||
the original chunks."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
|
|
@ -258,14 +253,19 @@ class _LegacyHookStreamAdapter(CustomGuardrail):
|
|||
guardrail_name,
|
||||
f"the legacy hook returned {returned_text_count} texts for a stream that carried {sent_text_count}",
|
||||
)
|
||||
tool_call_mismatch: Final = _tool_call_mismatch_reason(
|
||||
_tool_call_shapes(inputs.get("tool_calls")), _tool_call_shapes(rescanned.get("tool_calls"))
|
||||
sent_tool_calls: Final = inputs.get("tool_calls")
|
||||
rewritten_tool_calls: Final = rescanned.get("tool_calls")
|
||||
tool_call_count_mismatch: Final = _tool_call_count_mismatch_reason(sent_tool_calls, rewritten_tool_calls)
|
||||
if tool_call_count_mismatch is not None:
|
||||
raise UndeliverableStreamRewrite(guardrail_name, tool_call_count_mismatch)
|
||||
texts_update: Final[GenericGuardrailAPIInputs] = {"texts": rewritten} if rewritten else {}
|
||||
tool_calls_update: Final[GenericGuardrailAPIInputs] = (
|
||||
{"tool_calls": rewritten_tool_calls}
|
||||
if rewritten_tool_calls is not None
|
||||
and _rewrote(_tool_call_shapes(sent_tool_calls), _tool_call_shapes(rewritten_tool_calls))
|
||||
else {}
|
||||
)
|
||||
if tool_call_mismatch is not None:
|
||||
raise UndeliverableStreamRewrite(guardrail_name, tool_call_mismatch)
|
||||
if not rewritten:
|
||||
return inputs
|
||||
rewritten_inputs: Final[GenericGuardrailAPIInputs] = {**inputs, "texts": rewritten}
|
||||
rewritten_inputs: Final[GenericGuardrailAPIInputs] = {**inputs, **texts_update, **tool_calls_update}
|
||||
return rewritten_inputs
|
||||
|
||||
async def _rescan(
|
||||
|
|
|
|||
|
|
@ -79,6 +79,7 @@ from litellm.proxy.hooks.responses_id_security import ResponsesIDSecurity, decry
|
|||
from litellm.proxy.litellm_pre_call_utils import (
|
||||
_CLIENT_PRICING_METADATA_FIELDS, # pyright: ignore[reportPrivateUsage] # the HTTP path's own strip list
|
||||
_UNTRUSTED_METADATA_CONTROL_FIELDS, # pyright: ignore[reportPrivateUsage] # the HTTP path's own strip list
|
||||
is_chat_completions_body,
|
||||
)
|
||||
from litellm.proxy.pass_through_endpoints.llm_provider_handlers.batch_attribution import (
|
||||
request_tags_from_request_data,
|
||||
|
|
@ -213,13 +214,6 @@ def _normalize_tool_dialect(
|
|||
return {**data, **{key: value for key, value in replaceable if key in data}}
|
||||
|
||||
|
||||
def _is_chat_completions_body(data: Mapping[str, object]) -> bool:
|
||||
messages: Final = data.get("messages")
|
||||
if isinstance(messages, list) and messages:
|
||||
return True
|
||||
return "messages" in data and "input" not in data
|
||||
|
||||
|
||||
_CURSOR_THINKING_SEPARATOR: Final = "-thinking-"
|
||||
_CURSOR_FAST_SUFFIX: Final = "-fast"
|
||||
_CURSOR_THINKING_LEVELS: Final[frozenset[str]] = frozenset(get_args(REASONING_EFFORT))
|
||||
|
|
@ -261,7 +255,7 @@ def _resolve_cursor_model_variant(
|
|||
resolved: Final = {**data, "model": variant.base_model}
|
||||
if variant.reasoning_effort is None:
|
||||
return resolved
|
||||
if _is_chat_completions_body(data):
|
||||
if is_chat_completions_body(data):
|
||||
if "reasoning_effort" in data:
|
||||
return resolved
|
||||
return {**resolved, "reasoning_effort": variant.reasoning_effort}
|
||||
|
|
@ -646,7 +640,7 @@ async def cursor_chat_completions(
|
|||
|
||||
raw_body: Final = await read_request_body(request=request)
|
||||
|
||||
if _is_chat_completions_body(raw_body):
|
||||
if is_chat_completions_body(raw_body):
|
||||
# Genuine chat completions body (Cursor sends these for models whose BYOK it
|
||||
# already fixed); delegate so behavior matches /chat/completions exactly.
|
||||
# Keyed on messages CONTENT, not key presence: Cursor can send a null or
|
||||
|
|
|
|||
|
|
@ -250,7 +250,11 @@ from litellm.proxy.hooks.sensitive_data_routing import ( # noqa: F401, RUF100
|
|||
PROXY_SensitiveDataRoutingHandler,
|
||||
_PROXY_SensitiveDataRoutingHandler, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export
|
||||
)
|
||||
from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup, add_guardrails_from_auth_metadata
|
||||
from litellm.proxy.litellm_pre_call_utils import (
|
||||
LiteLLMProxyRequestSetup,
|
||||
add_guardrails_from_auth_metadata,
|
||||
normalized_metadata_slot,
|
||||
)
|
||||
from litellm.proxy.management_helpers.key_settings_audit import with_settings_updated_at
|
||||
from litellm.proxy.policy_engine.pipeline_executor import (
|
||||
PipelineExecutor,
|
||||
|
|
@ -3973,9 +3977,7 @@ class ProxyLogging:
|
|||
)
|
||||
request_data["litellm_logging_obj"] = litellm_logging_obj # rebind-ok: lifted then popped by the caller
|
||||
request_data.update({REJECTED_BEFORE_REQUEST_SETUP_KEY: REJECTED_BEFORE_REQUEST_SETUP})
|
||||
if "metadata" not in request_data:
|
||||
request_data["metadata"] = {}
|
||||
request_data["metadata"].update(user_api_key_logged_metadata)
|
||||
normalized_metadata_slot(request_data, "metadata").update(user_api_key_logged_metadata)
|
||||
|
||||
if litellm_logging_obj is not None:
|
||||
## UPDATE LOGGING INPUT
|
||||
|
|
|
|||
|
|
@ -759,7 +759,7 @@ def _dispatched_model_name(
|
|||
from litellm.litellm_core_utils.get_llm_provider_logic import resolve_model_and_provider_for_metadata
|
||||
|
||||
provider_model, _ = resolve_model_and_provider_for_metadata(model=model, custom_llm_provider=custom_llm_provider)
|
||||
return _strip_responses_routing_prefix(provider_model)
|
||||
return strip_responses_routing_prefix(provider_model)
|
||||
|
||||
|
||||
def _will_bridge_to_chat_completions(
|
||||
|
|
@ -1245,7 +1245,7 @@ def _pop_use_chat_completions_api_kw(kwargs: dict[str, object]) -> bool:
|
|||
_RESPONSES_ROUTING_PREFIX: Final = "responses/"
|
||||
|
||||
|
||||
def _strip_responses_routing_prefix(model: str) -> str:
|
||||
def strip_responses_routing_prefix(model: str) -> str:
|
||||
if not model.startswith(_RESPONSES_ROUTING_PREFIX):
|
||||
return model
|
||||
return model[len(_RESPONSES_ROUTING_PREFIX) :]
|
||||
|
|
@ -1273,7 +1273,7 @@ def _resolve_model_provider_for_responses(
|
|||
litellm_params.api_key = dynamic_api_key
|
||||
if dynamic_api_base is not None:
|
||||
litellm_params.api_base = dynamic_api_base
|
||||
return _strip_responses_routing_prefix(provider_model), resolved_provider
|
||||
return strip_responses_routing_prefix(provider_model), resolved_provider
|
||||
|
||||
|
||||
def _apply_managed_file_id_mapping(
|
||||
|
|
@ -2893,7 +2893,7 @@ async def _aresponses_websocket(
|
|||
api_key=api_key,
|
||||
litellm_params=litellm_params,
|
||||
)
|
||||
resolved_model: Final = _strip_responses_routing_prefix(provider_model)
|
||||
resolved_model: Final = strip_responses_routing_prefix(provider_model)
|
||||
|
||||
litellm_params_dict["data_residency"] = infer_openai_data_residency(
|
||||
_custom_llm_provider,
|
||||
|
|
|
|||
|
|
@ -43,6 +43,7 @@ from litellm.litellm_core_utils.core_helpers import (
|
|||
process_response_headers,
|
||||
provider_reported_usage_cost,
|
||||
set_response_cost_in_hidden_params,
|
||||
stamp_litellm_bill_on_usage,
|
||||
)
|
||||
from litellm.litellm_core_utils.exception_mapping_utils import with_deployment_api_base
|
||||
from litellm.litellm_core_utils.internal_call_metadata import (
|
||||
|
|
@ -1872,7 +1873,7 @@ def stamp_responses_usage_cost(
|
|||
if isinstance(bill, bool) or not isinstance(bill, (int, float)):
|
||||
return
|
||||
if bill > 0 or isinstance(getattr(usage_obj, "cost", None), (int, float)):
|
||||
setattr(usage_obj, "cost", float(bill))
|
||||
stamp_litellm_bill_on_usage(response_obj, usage_obj, float(bill))
|
||||
|
||||
|
||||
def opening_item_payload(payload: Mapping[str, object]) -> Mapping[str, object]:
|
||||
|
|
|
|||
|
|
@ -377,6 +377,7 @@ from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check impor
|
|||
EncryptedContentAffinityCheck,
|
||||
)
|
||||
from litellm.router_utils.pre_call_checks.io_token_rate_limit_check import (
|
||||
async_io_token_refund_failure,
|
||||
build_io_token_rate_limit_headers,
|
||||
deployment_has_io_token_limits,
|
||||
refund_stale_reservation_before_retry,
|
||||
|
|
@ -9500,7 +9501,7 @@ class Router(DeprecatedRouterMembers):
|
|||
|
||||
raise original_exception
|
||||
|
||||
async def make_call(self, original_function: Any, *args, **kwargs):
|
||||
async def make_call(self, original_function: Any, *args: object, **kwargs: object):
|
||||
"""
|
||||
Handler for making a call to the .completion()/.embeddings()/etc. functions.
|
||||
"""
|
||||
|
|
@ -9509,17 +9510,26 @@ class Router(DeprecatedRouterMembers):
|
|||
endpoint_function: Final = (
|
||||
generic_function if callable(generic_function) else _OPAQUE_VALUE_ADAPTER.validate_python(original_function)
|
||||
)
|
||||
with endpoint_selection(
|
||||
getattr(endpoint_function, "__name__", None),
|
||||
getattr(endpoint_function, "__module__", None),
|
||||
model=_OPAQUE_VALUE_ADAPTER.validate_python(model_group),
|
||||
):
|
||||
response = original_function(*args, **kwargs)
|
||||
if coroutine_checker.is_async_callable(response) or inspect.isawaitable(response):
|
||||
response = await response
|
||||
try:
|
||||
with endpoint_selection(
|
||||
getattr(endpoint_function, "__name__", None),
|
||||
getattr(endpoint_function, "__module__", None),
|
||||
model=_OPAQUE_VALUE_ADAPTER.validate_python(model_group),
|
||||
):
|
||||
response = original_function(*args, **kwargs)
|
||||
if coroutine_checker.is_async_callable(response) or inspect.isawaitable(response):
|
||||
response = await response
|
||||
except Exception:
|
||||
with contextlib.suppress(Exception):
|
||||
await async_io_token_refund_failure(
|
||||
self.cache, kwargs, parent_otel_span=get_parent_otel_span_from_kwargs(kwargs)
|
||||
)
|
||||
raise
|
||||
await self.increment_deployment_usage_for_response(response=response, request_kwargs=kwargs)
|
||||
## PROCESS RESPONSE HEADERS
|
||||
response = await self.set_response_headers(response=response, model_group=model_group, request_kwargs=kwargs)
|
||||
response = await self.set_response_headers(
|
||||
response=response, model_group=model_group if isinstance(model_group, str) else None, request_kwargs=kwargs
|
||||
)
|
||||
|
||||
return response
|
||||
|
||||
|
|
|
|||
|
|
@ -1,7 +1,8 @@
|
|||
"""
|
||||
Separate ITPM/OTPM (input/output tokens per minute) deployment rate limits.
|
||||
|
||||
- Pre-call: atomically reserve estimated_input against ITPM and max_tokens against OTPM
|
||||
- Pre-call: atomically reserve estimated_input against ITPM and max_tokens against OTPM; a request
|
||||
without max_tokens reserves the model's max output, capped at a fraction of the OTPM limit
|
||||
- Post-call: reconcile ITPM to actual input tokens and OTPM to actual output tokens
|
||||
- Cached prompt-read tokens are excluded from ITPM post-call accounting
|
||||
|
||||
|
|
@ -22,6 +23,7 @@ from litellm import token_counter
|
|||
from litellm._internal_context import with_service_target
|
||||
from litellm._logging import verbose_router_logger
|
||||
from litellm.caching.dual_cache import DualCache
|
||||
from litellm.constants import NO_MAX_TOKENS_TPM_RESERVATION_FRACTION
|
||||
from litellm.litellm_core_utils.token_counter import deployment_openai_model, offload_token_count
|
||||
from litellm.llms.azure.ptu_capacity import is_str_keyed_mapping
|
||||
from litellm.router_utils.routing_read_batch import ROUTER_USAGE_TARGET
|
||||
|
|
@ -127,10 +129,15 @@ def _get_cache_keys(deployment: dict, current_minute: str) -> tuple[str, str] |
|
|||
return itpm_key, otpm_key
|
||||
|
||||
|
||||
def _tokenizer_model(deployment: object) -> str:
|
||||
def _served_model(deployment: object) -> str:
|
||||
params: Final = deployment.get("litellm_params") if is_str_keyed_mapping(deployment) else None
|
||||
served: Final = params.get("model") if is_str_keyed_mapping(params) else None
|
||||
return deployment_openai_model(deployment) or (served if isinstance(served, str) else "")
|
||||
return served if isinstance(served, str) else ""
|
||||
|
||||
|
||||
def _tokenizer_and_served_models(deployment: object) -> tuple[str, str]:
|
||||
served: Final = _served_model(deployment)
|
||||
return deployment_openai_model(deployment) or served, served
|
||||
|
||||
|
||||
def _estimate_input_tokens(request_kwargs: dict[str, Any] | None, model: str = "") -> int:
|
||||
|
|
@ -160,24 +167,61 @@ def _model_max_output_tokens(model_name: str) -> int | None:
|
|||
return None
|
||||
|
||||
|
||||
def _resolve_max_tokens(request_kwargs: dict[str, Any] | None, deployment: dict) -> int:
|
||||
if request_kwargs:
|
||||
# An explicit max_tokens=0 must be honored, not treated as absent and
|
||||
# replaced by the model default.
|
||||
explicit = request_kwargs.get("max_tokens")
|
||||
if explicit is None:
|
||||
explicit = request_kwargs.get("max_completion_tokens")
|
||||
if explicit is None:
|
||||
explicit = request_kwargs.get("max_output_tokens")
|
||||
if explicit is not None:
|
||||
return max(0, int(explicit))
|
||||
_CHAT_OUTPUT_CAP_KEYS: Final = ("max_tokens", "max_completion_tokens")
|
||||
_RESPONSES_OUTPUT_CAP_KEYS: Final = ("max_output_tokens",)
|
||||
_OUTPUT_CAP_KEYS: Final = _CHAT_OUTPUT_CAP_KEYS + _RESPONSES_OUTPUT_CAP_KEYS
|
||||
|
||||
model_name: Final = (deployment.get("litellm_params") or {}).get("model")
|
||||
if model_name:
|
||||
model_max: Final = _model_max_output_tokens(model_name)
|
||||
if model_max is not None:
|
||||
return model_max
|
||||
return 4096
|
||||
|
||||
def _explicit_max_tokens(
|
||||
request_kwargs: Mapping[str, object] | None, cap_keys: tuple[str, ...] = _OUTPUT_CAP_KEYS
|
||||
) -> int | None:
|
||||
if not request_kwargs:
|
||||
return None
|
||||
explicit: Final = next(
|
||||
(request_kwargs[key] for key in cap_keys if request_kwargs.get(key) is not None),
|
||||
None,
|
||||
)
|
||||
match explicit:
|
||||
case int() | float() | str():
|
||||
return max(0, int(explicit))
|
||||
case _:
|
||||
return None
|
||||
|
||||
|
||||
def _sent_deployment_cap_keys(request_kwargs: Mapping[str, object] | None) -> tuple[str, ...]:
|
||||
is_responses_request: Final = (
|
||||
request_kwargs is not None and "input" in request_kwargs and "messages" not in request_kwargs
|
||||
)
|
||||
return _RESPONSES_OUTPUT_CAP_KEYS if is_responses_request else _CHAT_OUTPUT_CAP_KEYS
|
||||
|
||||
|
||||
def _deployment_output_cap(request_kwargs: Mapping[str, object] | None, deployment: Mapping[str, object]) -> int | None:
|
||||
params: Final = deployment.get("litellm_params")
|
||||
return (
|
||||
_explicit_max_tokens(params, _sent_deployment_cap_keys(request_kwargs))
|
||||
if is_str_keyed_mapping(params)
|
||||
else None
|
||||
)
|
||||
|
||||
|
||||
def _default_otpm_reservation(model_name: str, otpm_limit: int) -> int:
|
||||
model_max: Final = _model_max_output_tokens(model_name) if model_name else None
|
||||
share_of_limit: Final = max(1, int(otpm_limit * NO_MAX_TOKENS_TPM_RESERVATION_FRACTION))
|
||||
return min(4096 if model_max is None else model_max, share_of_limit)
|
||||
|
||||
|
||||
def _otpm_reservation(
|
||||
request_kwargs: Mapping[str, object] | None,
|
||||
deployment: Mapping[str, object],
|
||||
model_name: str,
|
||||
otpm_limit: int,
|
||||
) -> int:
|
||||
from_request: Final = _explicit_max_tokens(request_kwargs)
|
||||
if from_request is not None:
|
||||
return from_request
|
||||
default: Final = _default_otpm_reservation(model_name, otpm_limit)
|
||||
from_deployment: Final = _deployment_output_cap(request_kwargs, deployment)
|
||||
return default if from_deployment is None else min(from_deployment, default)
|
||||
|
||||
|
||||
def _get_usage_tokens(usage: object) -> tuple[int, int, int]:
|
||||
|
|
@ -210,6 +254,12 @@ def _extract_response_usage(response_obj: object) -> object:
|
|||
return getattr(response_obj, "usage", None)
|
||||
|
||||
|
||||
def _openai_shaped_usage(usage: object) -> object:
|
||||
if isinstance(usage, dict) and litellm.AnthropicConfig.is_anthropic_usage_object(usage):
|
||||
return litellm.AnthropicConfig().calculate_usage(usage_object=usage, reasoning_content=None)
|
||||
return usage
|
||||
|
||||
|
||||
def _usage_is_present(usage: object) -> bool:
|
||||
"""
|
||||
True only if usage carries an actual input/output breakdown.
|
||||
|
|
@ -238,7 +288,7 @@ def _resolve_reconcile_usage_tokens(
|
|||
fields. When usage cannot be resolved, return ``usage_resolved=False`` so
|
||||
callers keep the pre-call reservation instead of refunding it as zero usage.
|
||||
"""
|
||||
usage: Final = _extract_response_usage(response_obj)
|
||||
usage: Final = _openai_shaped_usage(_extract_response_usage(response_obj))
|
||||
if _usage_is_present(usage):
|
||||
prompt_tokens, completion_tokens, cached_tokens = _get_usage_tokens(usage)
|
||||
return max(0, prompt_tokens - cached_tokens), completion_tokens, True
|
||||
|
|
@ -261,6 +311,21 @@ def _resolve_reconcile_usage_tokens(
|
|||
return 0, 0, False
|
||||
|
||||
|
||||
def _key_for_this_minute(reserved_key: str) -> str:
|
||||
counter, _, _ = reserved_key.rpartition(":")
|
||||
return f"{counter}:{get_utc_datetime().strftime('%H-%M')}"
|
||||
|
||||
|
||||
def _reconcile_increments(
|
||||
reservation: tuple[int, int, str | None, str | None], billable_input: int, completion_tokens: int
|
||||
) -> tuple[tuple[str, int], ...]:
|
||||
itpm_reserved, otpm_reserved, itpm_key, otpm_key = reservation
|
||||
deltas: Final = ((itpm_key, billable_input - itpm_reserved), (otpm_key, completion_tokens - otpm_reserved))
|
||||
return tuple(
|
||||
(key if delta < 0 else _key_for_this_minute(key), delta) for key, delta in deltas if key is not None and delta
|
||||
)
|
||||
|
||||
|
||||
def _stash_reservation_in_metadata(
|
||||
request_kwargs: dict[str, object] | None,
|
||||
*,
|
||||
|
|
@ -311,7 +376,9 @@ def _reservation_channels(kwargs: Mapping[str, object] | None) -> tuple[object,
|
|||
return ()
|
||||
top_level: Final = (kwargs.get("metadata"), kwargs.get("litellm_metadata"))
|
||||
litellm_params: Final = _as_mutable_mapping(kwargs.get("litellm_params"))
|
||||
from_params: Final = () if litellm_params is None else (litellm_params.get("metadata"),)
|
||||
from_params: Final = (
|
||||
() if litellm_params is None else (litellm_params.get("metadata"), litellm_params.get("litellm_metadata"))
|
||||
)
|
||||
standard_logging_object: Final = _as_mutable_mapping(kwargs.get("standard_logging_object"))
|
||||
from_logging_object: Final = () if standard_logging_object is None else (standard_logging_object.get("metadata"),)
|
||||
return top_level + from_params + from_logging_object
|
||||
|
|
@ -433,9 +500,11 @@ def io_token_pre_call_check(
|
|||
return deployment
|
||||
|
||||
request_kwargs: Final = get_io_token_rate_limit_request_kwargs()
|
||||
_model: Final = _tokenizer_model(deployment)
|
||||
_model, served_model = _tokenizer_and_served_models(deployment)
|
||||
estimated_input: Final = _estimate_input_tokens(request_kwargs, model=_model)
|
||||
max_tokens: Final = _resolve_max_tokens(request_kwargs, deployment)
|
||||
otpm_to_reserve: Final = (
|
||||
_otpm_reservation(request_kwargs, deployment, served_model, otpm_limit) if otpm_limit is not None else 0
|
||||
)
|
||||
|
||||
dt: Final = get_utc_datetime()
|
||||
current_minute: Final = dt.strftime("%H-%M")
|
||||
|
|
@ -458,7 +527,7 @@ def io_token_pre_call_check(
|
|||
)
|
||||
|
||||
if otpm_limit is not None:
|
||||
otpm_reserved = 0 if max_tokens == 0 else _reservation_value(max_tokens, otpm_limit)
|
||||
otpm_reserved = otpm_to_reserve
|
||||
try:
|
||||
_sync_increment_with_rollback(
|
||||
dual_cache,
|
||||
|
|
@ -497,9 +566,11 @@ async def async_io_token_pre_call_check(
|
|||
return deployment
|
||||
|
||||
request_kwargs: Final = get_io_token_rate_limit_request_kwargs()
|
||||
_model: Final = _tokenizer_model(deployment)
|
||||
_model, served_model = _tokenizer_and_served_models(deployment)
|
||||
estimated_input: Final = await offload_token_count(_estimate_input_tokens)(request_kwargs, model=_model)
|
||||
max_tokens: Final = _resolve_max_tokens(request_kwargs, deployment)
|
||||
otpm_to_reserve: Final = (
|
||||
_otpm_reservation(request_kwargs, deployment, served_model, otpm_limit) if otpm_limit is not None else 0
|
||||
)
|
||||
|
||||
dt: Final = get_utc_datetime()
|
||||
current_minute: Final = dt.strftime("%H-%M")
|
||||
|
|
@ -523,7 +594,7 @@ async def async_io_token_pre_call_check(
|
|||
)
|
||||
|
||||
if otpm_limit is not None:
|
||||
otpm_reserved = 0 if max_tokens == 0 else _reservation_value(max_tokens, otpm_limit)
|
||||
otpm_reserved = otpm_to_reserve
|
||||
try:
|
||||
await _increment_with_rollback(
|
||||
dual_cache,
|
||||
|
|
@ -564,7 +635,8 @@ def io_token_reconcile_success(
|
|||
) -> None:
|
||||
request_kwargs: Final[Mapping[str, object] | None] = kwargs
|
||||
response: Final[object] = response_obj
|
||||
itpm_reserved, otpm_reserved, itpm_key, otpm_key = _read_reservation_from_kwargs(request_kwargs)
|
||||
reservation: Final = _read_reservation_from_kwargs(request_kwargs)
|
||||
itpm_reserved, otpm_reserved, itpm_key, otpm_key = reservation
|
||||
if itpm_key is None and otpm_key is None:
|
||||
return
|
||||
|
||||
|
|
@ -572,23 +644,8 @@ def io_token_reconcile_success(
|
|||
|
||||
try:
|
||||
if usage_resolved:
|
||||
if itpm_key is not None:
|
||||
itpm_delta: Final = billable_input - itpm_reserved
|
||||
if itpm_delta != 0:
|
||||
dual_cache.increment_cache(
|
||||
key=itpm_key,
|
||||
value=itpm_delta,
|
||||
ttl=RoutingArgsTTL,
|
||||
)
|
||||
|
||||
if otpm_key is not None:
|
||||
otpm_delta: Final = completion_tokens - otpm_reserved
|
||||
if otpm_delta != 0:
|
||||
dual_cache.increment_cache(
|
||||
key=otpm_key,
|
||||
value=otpm_delta,
|
||||
ttl=RoutingArgsTTL,
|
||||
)
|
||||
for key, delta in _reconcile_increments(reservation, billable_input, completion_tokens):
|
||||
dual_cache.increment_cache(key=key, value=delta, ttl=RoutingArgsTTL)
|
||||
else:
|
||||
verbose_router_logger.debug(
|
||||
"[IO TOKEN LIMIT] usage missing; keeping reservation (itpm_reserved=%s, otpm_reserved=%s)",
|
||||
|
|
@ -618,39 +675,23 @@ async def async_io_token_reconcile_success(
|
|||
) -> None:
|
||||
request_kwargs: Final[Mapping[str, object] | None] = kwargs
|
||||
response: Final[object] = response_obj
|
||||
itpm_reserved, otpm_reserved, itpm_key, otpm_key = _read_reservation_from_kwargs(request_kwargs)
|
||||
reservation: Final = _read_reservation_from_kwargs(request_kwargs)
|
||||
itpm_reserved, otpm_reserved, itpm_key, otpm_key = reservation
|
||||
if itpm_key is None and otpm_key is None:
|
||||
return
|
||||
|
||||
billable_input, completion_tokens, usage_resolved = _resolve_reconcile_usage_tokens(request_kwargs, response)
|
||||
|
||||
# Reconcile against the exact key that held the reservation (which encodes
|
||||
# the reservation's minute), not a key recomputed at response time. This
|
||||
# tracks actual usage even when the pre-call estimate was 0, and avoids a
|
||||
# minute-boundary mismatch for calls that span into the next minute. Always
|
||||
# clear the stash afterwards (even if an increment throws) so a retry or a
|
||||
# duplicate success event can't re-process it.
|
||||
# A refund goes back to the minute that held the reservation, while usage above the reservation is charged
|
||||
# to the minute the response finished in, so a call that runs past its minute still counts against a live
|
||||
# window. Always clear the stash afterwards (even if an increment throws) so a retry or a duplicate success
|
||||
# event can't re-process it.
|
||||
try:
|
||||
if usage_resolved:
|
||||
if itpm_key is not None:
|
||||
itpm_delta: Final = billable_input - itpm_reserved
|
||||
if itpm_delta != 0:
|
||||
await dual_cache.async_increment_cache(
|
||||
key=itpm_key,
|
||||
value=itpm_delta,
|
||||
ttl=RoutingArgsTTL,
|
||||
parent_otel_span=parent_otel_span,
|
||||
)
|
||||
|
||||
if otpm_key is not None:
|
||||
otpm_delta: Final = completion_tokens - otpm_reserved
|
||||
if otpm_delta != 0:
|
||||
await dual_cache.async_increment_cache(
|
||||
key=otpm_key,
|
||||
value=otpm_delta,
|
||||
ttl=RoutingArgsTTL,
|
||||
parent_otel_span=parent_otel_span,
|
||||
)
|
||||
for key, delta in _reconcile_increments(reservation, billable_input, completion_tokens):
|
||||
await dual_cache.async_increment_cache(
|
||||
key=key, value=delta, ttl=RoutingArgsTTL, parent_otel_span=parent_otel_span
|
||||
)
|
||||
else:
|
||||
verbose_router_logger.debug(
|
||||
"[IO TOKEN LIMIT] usage missing; keeping reservation (itpm_reserved=%s, otpm_reserved=%s)",
|
||||
|
|
|
|||
|
|
@ -26,7 +26,6 @@ from litellm.caching.dual_cache import DualCache
|
|||
from litellm.caching.redis_cache import RedisCircuitBreakerOpenError
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.router_utils.pre_call_checks.io_token_rate_limit_check import (
|
||||
ITPM_RESERVED_KEY,
|
||||
async_io_token_pre_call_check,
|
||||
async_io_token_reconcile_success,
|
||||
async_io_token_refund_failure,
|
||||
|
|
@ -485,24 +484,12 @@ class ModelRateLimitingCheck(CustomLogger):
|
|||
|
||||
try:
|
||||
standard_logging_object: Final[StandardLoggingPayload | None] = kwargs.get("standard_logging_object")
|
||||
|
||||
# IO token reconciliation works purely from the cache keys stashed in
|
||||
# kwargs/metadata, so it must run before the model_id guard below
|
||||
# (which only the TPM-tracking path needs). Otherwise a request whose
|
||||
# standard_logging_object lacks model_id would never return its
|
||||
# reservation, leaving the counter elevated until the TTL expires.
|
||||
slo_metadata: Final = (standard_logging_object.get("metadata") or {}) if standard_logging_object else {}
|
||||
kwargs_metadata: Final = kwargs.get("metadata") or {}
|
||||
if ITPM_RESERVED_KEY in slo_metadata or ITPM_RESERVED_KEY in kwargs_metadata:
|
||||
await async_io_token_reconcile_success(
|
||||
self.dual_cache,
|
||||
kwargs,
|
||||
response_obj,
|
||||
parent_otel_span=get_parent_otel_span_from_kwargs(kwargs),
|
||||
)
|
||||
# Fall through: a deployment can also configure tpm/rpm alongside
|
||||
# itpm/otpm, and that path's pre-call check reads the tpm_key
|
||||
# counter tracked below, so it must still be incremented here.
|
||||
await async_io_token_reconcile_success(
|
||||
self.dual_cache,
|
||||
kwargs,
|
||||
response_obj,
|
||||
parent_otel_span=get_parent_otel_span_from_kwargs(kwargs),
|
||||
)
|
||||
|
||||
if standard_logging_object is None:
|
||||
return
|
||||
|
|
@ -557,17 +544,7 @@ class ModelRateLimitingCheck(CustomLogger):
|
|||
"""
|
||||
try:
|
||||
standard_logging_object: Final[StandardLoggingPayload | None] = kwargs.get("standard_logging_object")
|
||||
slo_metadata: Final = (standard_logging_object.get("metadata") or {}) if standard_logging_object else {}
|
||||
kwargs_metadata: Final = kwargs.get("metadata") or {}
|
||||
if ITPM_RESERVED_KEY in slo_metadata or ITPM_RESERVED_KEY in kwargs_metadata:
|
||||
io_token_reconcile_success(
|
||||
self.dual_cache,
|
||||
kwargs,
|
||||
response_obj,
|
||||
)
|
||||
# Fall through: a deployment can also configure tpm/rpm alongside
|
||||
# itpm/otpm, and that path's pre-call check reads the tpm_key
|
||||
# counter tracked below, so it must still be incremented here.
|
||||
io_token_reconcile_success(self.dual_cache, kwargs, response_obj)
|
||||
|
||||
if standard_logging_object is None:
|
||||
return
|
||||
|
|
|
|||
|
|
@ -713,6 +713,23 @@ class AnthropicResponseContentBlockRedactedThinking(LiteLLMBaseModel):
|
|||
data: str
|
||||
|
||||
|
||||
class AnthropicResponseContentBlockServerToolUse(LiteLLMBaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
type: Literal["server_tool_use"]
|
||||
id: str
|
||||
name: str
|
||||
input: Mapping[str, object]
|
||||
|
||||
|
||||
class AnthropicResponseContentBlockWebSearchToolResult(LiteLLMBaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
type: Literal["web_search_tool_result"]
|
||||
tool_use_id: str
|
||||
content: Sequence[Mapping[str, object]] | Mapping[str, object]
|
||||
|
||||
|
||||
class AnthropicResponseUsageBlock(LiteLLMBaseModel):
|
||||
model_config = ConfigDict(extra="allow")
|
||||
|
||||
|
|
@ -744,6 +761,8 @@ class AnthropicResponse(LiteLLMBaseModel):
|
|||
| AnthropicResponseContentBlockToolUse
|
||||
| AnthropicResponseContentBlockThinking
|
||||
| AnthropicResponseContentBlockRedactedThinking
|
||||
| AnthropicResponseContentBlockServerToolUse
|
||||
| AnthropicResponseContentBlockWebSearchToolResult
|
||||
]
|
||||
"""Content generated by the model."""
|
||||
|
||||
|
|
|
|||
|
|
@ -81,6 +81,14 @@ class BidiGenerateContentServerMessage(TypedDict, total=False):
|
|||
"""Output only. The setup complete message."""
|
||||
|
||||
|
||||
class BidiGenerateContentActivityStart(TypedDict):
|
||||
pass
|
||||
|
||||
|
||||
class BidiGenerateContentActivityEnd(TypedDict):
|
||||
pass
|
||||
|
||||
|
||||
class BidiGenerateContentRealtimeInput(TypedDict, total=False):
|
||||
text: str
|
||||
"""The text to be sent to the model."""
|
||||
|
|
@ -94,11 +102,11 @@ class BidiGenerateContentRealtimeInput(TypedDict, total=False):
|
|||
audioStreamEnd: bool
|
||||
"""Output only. If true, indicates that the audio stream has ended."""
|
||||
|
||||
activityStart: bool
|
||||
"""Output only. If true, indicates that the activity has started."""
|
||||
activityStart: BidiGenerateContentActivityStart
|
||||
"""Marks the start of user activity when automatic activity detection is disabled."""
|
||||
|
||||
activityEnd: bool
|
||||
"""Output only. If true, indicates that the activity has ended."""
|
||||
activityEnd: BidiGenerateContentActivityEnd
|
||||
"""Marks the end of user activity when automatic activity detection is disabled."""
|
||||
|
||||
|
||||
StartOfSpeechSensitivityEnum = Literal[
|
||||
|
|
|
|||
|
|
@ -23,3 +23,22 @@ class MistralThinkingBlock(TypedDict):
|
|||
type: Literal["thinking"]
|
||||
thinking: list[MistralTextBlock]
|
||||
closed: NotRequired[ReadOnly[bool]]
|
||||
|
||||
|
||||
class MistralRequestRouting(TypedDict, total=False):
|
||||
api_base: ReadOnly[object]
|
||||
drop_params: ReadOnly[object]
|
||||
allowed_openai_params: ReadOnly[object]
|
||||
custom_llm_provider: ReadOnly[object]
|
||||
|
||||
|
||||
class MistralCacheKeyParams(TypedDict, total=False):
|
||||
prompt_cache_key: ReadOnly[object]
|
||||
|
||||
|
||||
class MistralReasoningEffortParams(TypedDict, total=False):
|
||||
reasoning_effort: ReadOnly[object]
|
||||
|
||||
|
||||
class MistralFilePart(TypedDict, total=False):
|
||||
file_id: ReadOnly[object]
|
||||
|
|
|
|||
|
|
@ -6708,6 +6708,26 @@ def _first_registered_match(
|
|||
return next(matches, (None, None))
|
||||
|
||||
|
||||
def _model_cost_lookup_order(
|
||||
model: str, potential_model_names: PotentialModelNamesAndCustomLLMProvider
|
||||
) -> tuple[str, ...]:
|
||||
return (
|
||||
potential_model_names["combined_model_name"],
|
||||
model,
|
||||
potential_model_names["region_free_combined_model_name"],
|
||||
potential_model_names["split_model"],
|
||||
potential_model_names["combined_stripped_model_name"],
|
||||
potential_model_names["stripped_model_name"],
|
||||
potential_model_names["provider_prefixed_model_name"],
|
||||
)
|
||||
|
||||
|
||||
def potential_model_names_in_lookup_order(model: str, custom_llm_provider: str) -> tuple[str, ...]:
|
||||
return _model_cost_lookup_order(
|
||||
model, get_potential_model_names(model=model, custom_llm_provider=custom_llm_provider)
|
||||
)
|
||||
|
||||
|
||||
def _get_model_info_from_generalization(
|
||||
model: str,
|
||||
potential_model_names: PotentialModelNamesAndCustomLLMProvider,
|
||||
|
|
@ -6726,15 +6746,7 @@ def _get_model_info_from_generalization(
|
|||
still had a priced exact name to try. O(number of rules); only call after the
|
||||
exact lookups have missed.
|
||||
"""
|
||||
candidates: Final = (
|
||||
potential_model_names["combined_model_name"],
|
||||
model,
|
||||
potential_model_names["region_free_combined_model_name"],
|
||||
potential_model_names["split_model"],
|
||||
potential_model_names["combined_stripped_model_name"],
|
||||
potential_model_names["stripped_model_name"],
|
||||
potential_model_names["provider_prefixed_model_name"],
|
||||
)
|
||||
candidates: Final = _model_cost_lookup_order(model, potential_model_names)
|
||||
if any(get_model_cost_key(candidate) is not None for candidate in candidates):
|
||||
return None
|
||||
for candidate in candidates:
|
||||
|
|
@ -6973,11 +6985,6 @@ def get_model_info_helper(
|
|||
model=model, custom_llm_provider=custom_llm_provider or declared_authenticating_provider(model)
|
||||
)
|
||||
|
||||
combined_model_name: Final = potential_model_names["combined_model_name"]
|
||||
stripped_model_name: Final = potential_model_names["stripped_model_name"]
|
||||
combined_stripped_model_name: Final = potential_model_names["combined_stripped_model_name"]
|
||||
provider_prefixed_model_name: Final = potential_model_names["provider_prefixed_model_name"]
|
||||
split_model: Final = potential_model_names["split_model"]
|
||||
custom_llm_provider = potential_model_names["custom_llm_provider"]
|
||||
model_cost_custom_llm_provider: Final = custom_llm_provider
|
||||
#########################
|
||||
|
|
@ -7044,16 +7051,9 @@ def get_model_info_helper(
|
|||
and strip it. Tried last so no model that already resolves through 1-6 can change.
|
||||
"""
|
||||
|
||||
lookup_order: Final = (
|
||||
combined_model_name,
|
||||
model,
|
||||
potential_model_names["region_free_combined_model_name"],
|
||||
split_model,
|
||||
combined_stripped_model_name,
|
||||
stripped_model_name,
|
||||
provider_prefixed_model_name,
|
||||
lookup: Final = _first_registered_match(
|
||||
_model_cost_lookup_order(model, potential_model_names), model_cost_custom_llm_provider
|
||||
)
|
||||
lookup: Final = _first_registered_match(lookup_order, model_cost_custom_llm_provider)
|
||||
key: str | None = lookup[0]
|
||||
_model_info: dict[str, Any] | None = lookup[1]
|
||||
|
||||
|
|
@ -10911,7 +10911,7 @@ class ProviderConfigManager:
|
|||
if LlmProviders.GEMINI == provider:
|
||||
from litellm.llms.gemini.realtime.transformation import GeminiRealtimeConfig
|
||||
|
||||
return GeminiRealtimeConfig()
|
||||
return GeminiRealtimeConfig(drop_params=drop_params)
|
||||
if LlmProviders.META == provider:
|
||||
from litellm.llms.meta.realtime.transformation import MetaRealtimeConfig
|
||||
|
||||
|
|
|
|||
File diff suppressed because it is too large
Load diff
21
tests/benchmarks/test_realtime_session_merge_benchmarks.py
Normal file
21
tests/benchmarks/test_realtime_session_merge_benchmarks.py
Normal file
|
|
@ -0,0 +1,21 @@
|
|||
from typing import Final
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm.litellm_core_utils.realtime_streaming import _merge_session_fields
|
||||
|
||||
NESTED_LEVELS: Final = 8
|
||||
LEAVES: Final = 2000
|
||||
|
||||
|
||||
def _nested_session(depth: int) -> dict[str, object]:
|
||||
return {"child": _nested_session(depth - 1)} if depth else {f"k{index}": index for index in range(LEAVES)}
|
||||
|
||||
|
||||
@pytest.mark.benchmark
|
||||
def test_merging_a_session_update_that_repeats_every_nested_level() -> None:
|
||||
session: Final = _nested_session(NESTED_LEVELS)
|
||||
|
||||
merged: Final = _merge_session_fields(session, session)
|
||||
|
||||
assert merged == session
|
||||
|
|
@ -85,6 +85,7 @@ IGNORE_FUNCTIONS = [
|
|||
"_sort_processed_sets", # bounded by the nesting depth of the log-record extra it walks (a finite JSON tree, no cycles possible).
|
||||
"scrub_json_strings", # max depth set (MAX_SCRUB_DEPTH); fails closed by returning "[Filtered]" for anything nested past the cap.
|
||||
"_without_requester_ip_address", # walks only the fresh tree _sanitize_request_body_for_spend_logs_payload returns (mapping cycles already cut), so it never goes deeper than that sanitizer did.
|
||||
"_merge_session_fields", # max depth set (_SESSION_MERGE_MAX_DEPTH); deeper mappings are replaced wholesale by the later session.update.
|
||||
]
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -47,6 +47,7 @@ class RealtimeProvider:
|
|||
id: str
|
||||
alias: str
|
||||
litellm_params: LiteLLMParamsBody
|
||||
modalities: tuple[str, ...] = ("text",)
|
||||
|
||||
|
||||
PROVIDERS = (
|
||||
|
|
@ -75,6 +76,7 @@ PROVIDERS = (
|
|||
model="gemini/gemini-3.1-flash-live-preview",
|
||||
api_key="os.environ/GEMINI_API_KEY",
|
||||
),
|
||||
modalities=("audio",),
|
||||
),
|
||||
RealtimeProvider(
|
||||
"vertex_ai",
|
||||
|
|
@ -84,6 +86,7 @@ PROVIDERS = (
|
|||
vertex_location="us-central1",
|
||||
vertex_credentials="os.environ/VERTEXAI_CREDENTIALS",
|
||||
),
|
||||
modalities=("audio",),
|
||||
),
|
||||
RealtimeProvider(
|
||||
"bedrock",
|
||||
|
|
|
|||
|
|
@ -124,7 +124,8 @@ def test_text_conversation(
|
|||
session.send(
|
||||
SessionUpdate(
|
||||
session=SessionConfig(
|
||||
instructions="You are a terse assistant. Reply in one short sentence."
|
||||
modalities=list(provider.modalities),
|
||||
instructions="You are a terse assistant. Reply in one short sentence.",
|
||||
)
|
||||
)
|
||||
)
|
||||
|
|
@ -159,6 +160,7 @@ def test_tool_call_round_trip(
|
|||
session.send(
|
||||
SessionUpdate(
|
||||
session=SessionConfig(
|
||||
modalities=list(provider.modalities),
|
||||
tools=[WEATHER_TOOL],
|
||||
tool_choice="auto",
|
||||
instructions=(
|
||||
|
|
|
|||
|
|
@ -2,6 +2,7 @@
|
|||
|
||||
PARITY pairs the streamed case with its plain twin where both bill the same row."""
|
||||
|
||||
from collections.abc import Mapping
|
||||
from typing import Final, cast
|
||||
|
||||
from integration.cost_calculation.cost_tracking_case import (
|
||||
|
|
@ -11878,7 +11879,215 @@ GPT_5_6_TERRA_MESSAGES_BRIDGED_STREAM: Final = GPT_5_6_TERRA_MESSAGES_BRIDGED_JS
|
|||
}
|
||||
)
|
||||
|
||||
|
||||
# An OpenAI-compatible gateway that reports what it charged under usage.cost (#47533). The reported total is the bill,
|
||||
# split across input and output by the ratio of the cost map rates (100 * 1e-6 : 20 * 3e-6), on plain and streamed
|
||||
# replies alike. A deployment the operator priced keeps the operator's rates on both.
|
||||
_GATEWAY_REPORTED_USAGE: Final = {"prompt_tokens": 100, "completion_tokens": 20, "total_tokens": 120, "cost": 0.037}
|
||||
_GATEWAY_REPORTED_EXPECTED: Final = ExactExpected(
|
||||
spend=0.037, input_cost=0.0, output_cost=0.037, prompt_tokens=100, completion_tokens=20
|
||||
)
|
||||
_GATEWAY_OPERATOR_PRICED: Final = Deployment(input_cost_per_token=2e-06, output_cost_per_token=6e-06)
|
||||
_GATEWAY_OPERATOR_PRICED_EXPECTED: Final = ExactExpected(
|
||||
spend=0.00032, input_cost=0.0002, output_cost=0.00012, prompt_tokens=100, completion_tokens=20
|
||||
)
|
||||
_GATEWAY_CHUNK: Final = {
|
||||
"id": "chatcmpl-$REQUEST_ID",
|
||||
"object": "chat.completion.chunk",
|
||||
"created": 1,
|
||||
"model": "gateway-reported-cost-fixture",
|
||||
}
|
||||
|
||||
GATEWAY_REPORTED_COST: Final = CostTrackingTestCase(
|
||||
name="gateway-reported-cost",
|
||||
covers="quota_management.spend_tracking.cost_matrix.logs_cost",
|
||||
model="gateway-reported-cost-fixture",
|
||||
endpoint="/v1/chat/completions",
|
||||
request={"model": "$MODEL", "messages": [{"role": "user", "content": "reported cost"}], "stream": False},
|
||||
response=JsonResponse(
|
||||
content_type="application/json",
|
||||
body={
|
||||
"id": "chatcmpl-$REQUEST_ID",
|
||||
"object": "chat.completion",
|
||||
"created": 1,
|
||||
"model": "gateway-reported-cost-fixture",
|
||||
"choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}],
|
||||
"usage": _GATEWAY_REPORTED_USAGE,
|
||||
},
|
||||
),
|
||||
expected=_GATEWAY_REPORTED_EXPECTED,
|
||||
)
|
||||
|
||||
GATEWAY_REPORTED_COST_STREAM: Final = CostTrackingTestCase(
|
||||
name="gateway-reported-cost-stream",
|
||||
covers="quota_management.spend_tracking.cost_matrix.logs_cost",
|
||||
model="gateway-reported-cost-fixture",
|
||||
endpoint="/v1/chat/completions",
|
||||
request={
|
||||
"model": "$MODEL",
|
||||
"messages": [{"role": "user", "content": "reported cost"}],
|
||||
"stream": True,
|
||||
"stream_options": {"include_usage": True},
|
||||
},
|
||||
response=SseResponse(
|
||||
content_type="text/event-stream",
|
||||
frames=sse_frames(
|
||||
{**_GATEWAY_CHUNK, "choices": [{"index": 0, "delta": {"role": "assistant", "content": "ok"}}]},
|
||||
{**_GATEWAY_CHUNK, "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}]},
|
||||
{**_GATEWAY_CHUNK, "choices": [], "usage": _GATEWAY_REPORTED_USAGE},
|
||||
done=True,
|
||||
),
|
||||
),
|
||||
expected=_GATEWAY_REPORTED_EXPECTED,
|
||||
)
|
||||
|
||||
GATEWAY_REPORTED_COST_OPERATOR_PRICED: Final = GATEWAY_REPORTED_COST.model_copy(
|
||||
update={
|
||||
"name": "gateway-reported-cost-operator_priced",
|
||||
"deployment": _GATEWAY_OPERATOR_PRICED,
|
||||
"expected": _GATEWAY_OPERATOR_PRICED_EXPECTED,
|
||||
}
|
||||
)
|
||||
|
||||
GATEWAY_REPORTED_COST_OPERATOR_PRICED_STREAM: Final = GATEWAY_REPORTED_COST_STREAM.model_copy(
|
||||
update={
|
||||
"name": "gateway-reported-cost-operator_priced-stream",
|
||||
"deployment": _GATEWAY_OPERATOR_PRICED,
|
||||
"expected": _GATEWAY_OPERATOR_PRICED_EXPECTED,
|
||||
}
|
||||
)
|
||||
|
||||
_GATEWAY_RESPONSES_USAGE: Final[Mapping[str, JsonValue]] = {"input_tokens": 100, "output_tokens": 20, "total_tokens": 120, "cost": 0.037}
|
||||
_GATEWAY_RESPONSES_ITEM: Final[Mapping[str, JsonValue]] = {
|
||||
"id": "msg_gateway",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"status": "completed",
|
||||
"content": [{"type": "output_text", "text": "ok", "annotations": []}],
|
||||
}
|
||||
_GATEWAY_RESPONSES_BODY: Final[Mapping[str, JsonValue]] = {
|
||||
"id": "resp_$REQUEST_ID",
|
||||
"object": "response",
|
||||
"created_at": 1,
|
||||
"status": "completed",
|
||||
"model": "gateway-reported-cost-fixture",
|
||||
"output": [_GATEWAY_RESPONSES_ITEM],
|
||||
"usage": _GATEWAY_RESPONSES_USAGE,
|
||||
}
|
||||
|
||||
GATEWAY_REPORTED_COST_RESPONSES: Final = CostTrackingTestCase(
|
||||
name="gateway-reported-cost-responses",
|
||||
covers="quota_management.spend_tracking.cost_matrix.logs_cost",
|
||||
model="gateway-reported-cost-fixture",
|
||||
endpoint="/v1/responses",
|
||||
request={"model": "$MODEL", "input": "reported cost", "cache": {"no-cache": True}},
|
||||
response=JsonResponse(content_type="application/json", body=_GATEWAY_RESPONSES_BODY),
|
||||
expected=_GATEWAY_REPORTED_EXPECTED,
|
||||
)
|
||||
|
||||
GATEWAY_REPORTED_COST_RESPONSES_STREAM: Final = GATEWAY_REPORTED_COST_RESPONSES.model_copy(
|
||||
update={
|
||||
"name": "gateway-reported-cost-responses-stream",
|
||||
"request": {**GATEWAY_REPORTED_COST_RESPONSES.request, "stream": True},
|
||||
"response": SseResponse(
|
||||
content_type="text/event-stream",
|
||||
frames=sse_frames(
|
||||
{
|
||||
"type": "response.created",
|
||||
"sequence_number": 0,
|
||||
"response": {**_GATEWAY_RESPONSES_BODY, "status": "in_progress", "output": [], "usage": None},
|
||||
},
|
||||
{
|
||||
"type": "response.output_text.delta",
|
||||
"sequence_number": 1,
|
||||
"item_id": "msg_gateway",
|
||||
"output_index": 0,
|
||||
"content_index": 0,
|
||||
"delta": "ok",
|
||||
},
|
||||
{"type": "response.completed", "sequence_number": 2, "response": _GATEWAY_RESPONSES_BODY},
|
||||
),
|
||||
),
|
||||
}
|
||||
)
|
||||
|
||||
GATEWAY_REPORTED_COST_RESPONSES_OPERATOR_PRICED: Final = GATEWAY_REPORTED_COST_RESPONSES.model_copy(
|
||||
update={
|
||||
"name": "gateway-reported-cost-responses-operator_priced",
|
||||
"deployment": _GATEWAY_OPERATOR_PRICED,
|
||||
"expected": _GATEWAY_OPERATOR_PRICED_EXPECTED,
|
||||
}
|
||||
)
|
||||
|
||||
GATEWAY_REPORTED_COST_RESPONSES_OPERATOR_PRICED_STREAM: Final = GATEWAY_REPORTED_COST_RESPONSES_STREAM.model_copy(
|
||||
update={
|
||||
"name": "gateway-reported-cost-responses-operator_priced-stream",
|
||||
"deployment": _GATEWAY_OPERATOR_PRICED,
|
||||
"expected": _GATEWAY_OPERATOR_PRICED_EXPECTED,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
_GATEWAY_ZERO_USAGE: Final = {**_GATEWAY_REPORTED_USAGE, "cost": 0}
|
||||
|
||||
GATEWAY_REPORTED_ZERO_OPERATOR_PRICED: Final = GATEWAY_REPORTED_COST_OPERATOR_PRICED.model_copy(
|
||||
update={
|
||||
"name": "gateway-reported-zero-operator_priced",
|
||||
"response": JsonResponse(
|
||||
content_type="application/json",
|
||||
body={**GATEWAY_REPORTED_COST.response.body, "usage": _GATEWAY_ZERO_USAGE},
|
||||
),
|
||||
}
|
||||
)
|
||||
|
||||
GATEWAY_REPORTED_ZERO_OPERATOR_PRICED_STREAM: Final = GATEWAY_REPORTED_COST_OPERATOR_PRICED_STREAM.model_copy(
|
||||
update={
|
||||
"name": "gateway-reported-zero-operator_priced-stream",
|
||||
"response": SseResponse(
|
||||
content_type="text/event-stream",
|
||||
frames=sse_frames(
|
||||
{**_GATEWAY_CHUNK, "choices": [{"index": 0, "delta": {"role": "assistant", "content": "ok"}}]},
|
||||
{**_GATEWAY_CHUNK, "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}]},
|
||||
{**_GATEWAY_CHUNK, "choices": [], "usage": _GATEWAY_ZERO_USAGE},
|
||||
done=True,
|
||||
),
|
||||
),
|
||||
}
|
||||
)
|
||||
|
||||
_GATEWAY_OPENROUTER_OPERATOR_PRICED: Final = _GATEWAY_OPERATOR_PRICED.model_copy(
|
||||
update={"model": "openrouter/vendor/gateway-reported-cost-fixture"}
|
||||
)
|
||||
|
||||
GATEWAY_REPORTED_COST_OPENROUTER_OPERATOR_PRICED: Final = GATEWAY_REPORTED_COST_OPERATOR_PRICED.model_copy(
|
||||
update={
|
||||
"name": "gateway-reported-cost-openrouter-operator_priced",
|
||||
"deployment": _GATEWAY_OPENROUTER_OPERATOR_PRICED,
|
||||
}
|
||||
)
|
||||
|
||||
GATEWAY_REPORTED_COST_OPENROUTER_OPERATOR_PRICED_STREAM: Final = (
|
||||
GATEWAY_REPORTED_COST_OPERATOR_PRICED_STREAM.model_copy(
|
||||
update={
|
||||
"name": "gateway-reported-cost-openrouter-operator_priced-stream",
|
||||
"deployment": _GATEWAY_OPENROUTER_OPERATOR_PRICED,
|
||||
}
|
||||
)
|
||||
)
|
||||
|
||||
CASES: Final[tuple[CostTrackingTestCase, ...]] = (
|
||||
GATEWAY_REPORTED_COST,
|
||||
GATEWAY_REPORTED_COST_STREAM,
|
||||
GATEWAY_REPORTED_COST_OPERATOR_PRICED,
|
||||
GATEWAY_REPORTED_COST_OPERATOR_PRICED_STREAM,
|
||||
GATEWAY_REPORTED_COST_RESPONSES,
|
||||
GATEWAY_REPORTED_COST_RESPONSES_STREAM,
|
||||
GATEWAY_REPORTED_COST_RESPONSES_OPERATOR_PRICED,
|
||||
GATEWAY_REPORTED_COST_RESPONSES_OPERATOR_PRICED_STREAM,
|
||||
GATEWAY_REPORTED_ZERO_OPERATOR_PRICED,
|
||||
GATEWAY_REPORTED_ZERO_OPERATOR_PRICED_STREAM,
|
||||
GATEWAY_REPORTED_COST_OPENROUTER_OPERATOR_PRICED,
|
||||
GATEWAY_REPORTED_COST_OPENROUTER_OPERATOR_PRICED_STREAM,
|
||||
GPT_5_6_SLASH_ALIAS,
|
||||
GPT_5_6_SLASH_ALIAS_STREAM,
|
||||
GPT_6_SOL_RESPONSES_DEPLOYMENT_JSON,
|
||||
|
|
@ -12120,6 +12329,22 @@ CASES: Final[tuple[CostTrackingTestCase, ...]] = (
|
|||
)
|
||||
|
||||
PARITY: Final[tuple[StreamParityTestCase, ...]] = (
|
||||
StreamParityTestCase(plain=GATEWAY_REPORTED_COST, streamed=GATEWAY_REPORTED_COST_STREAM),
|
||||
StreamParityTestCase(
|
||||
plain=GATEWAY_REPORTED_COST_OPERATOR_PRICED, streamed=GATEWAY_REPORTED_COST_OPERATOR_PRICED_STREAM
|
||||
),
|
||||
StreamParityTestCase(
|
||||
plain=GATEWAY_REPORTED_ZERO_OPERATOR_PRICED, streamed=GATEWAY_REPORTED_ZERO_OPERATOR_PRICED_STREAM
|
||||
),
|
||||
StreamParityTestCase(
|
||||
plain=GATEWAY_REPORTED_COST_OPENROUTER_OPERATOR_PRICED,
|
||||
streamed=GATEWAY_REPORTED_COST_OPENROUTER_OPERATOR_PRICED_STREAM,
|
||||
),
|
||||
StreamParityTestCase(plain=GATEWAY_REPORTED_COST_RESPONSES, streamed=GATEWAY_REPORTED_COST_RESPONSES_STREAM),
|
||||
StreamParityTestCase(
|
||||
plain=GATEWAY_REPORTED_COST_RESPONSES_OPERATOR_PRICED,
|
||||
streamed=GATEWAY_REPORTED_COST_RESPONSES_OPERATOR_PRICED_STREAM,
|
||||
),
|
||||
StreamParityTestCase(plain=GPT_5_3_CODEX_INPUT_TEXT, streamed=GPT_5_3_CODEX_STREAM),
|
||||
StreamParityTestCase(plain=GPT_5_3_CODEX_INPUT_TEXT, streamed=GPT_5_3_CODEX_STREAM_INCOMPLETE),
|
||||
StreamParityTestCase(plain=GPT_5_3_CODEX_INPUT_TEXT, streamed=GPT_5_3_CODEX_STREAM_UNVALIDATED),
|
||||
|
|
|
|||
|
|
@ -26,6 +26,12 @@ COST_MAP_ENTRIES: Final[dict[str, dict[str, JsonValue]]] = {
|
|||
"input_cost_per_token_above_128k_tokens": 2.2e-07,
|
||||
"output_cost_per_token_above_128k_tokens": 8.8e-07,
|
||||
},
|
||||
"gateway-reported-cost-fixture": {
|
||||
"litellm_provider": "openai",
|
||||
"mode": "chat",
|
||||
"input_cost_per_token": 1e-06,
|
||||
"output_cost_per_token": 3e-06,
|
||||
},
|
||||
"openrouter/anthropic/claude-sonnet-5": {
|
||||
"litellm_provider": "openrouter",
|
||||
"mode": "chat",
|
||||
|
|
|
|||
|
|
@ -29,6 +29,24 @@ def sse_events(text: str) -> tuple[tuple[str, dict[str, object]], ...]:
|
|||
)
|
||||
|
||||
|
||||
def streamed_blocks(events: tuple[tuple[str, dict[str, object]], ...]) -> tuple[dict[str, object], ...]:
|
||||
input_json: Final = tuple(
|
||||
(event["index"], event["delta"]["partial_json"])
|
||||
for name, event in events
|
||||
if name == "content_block_delta" and event["delta"]["type"] == "input_json_delta"
|
||||
)
|
||||
return tuple(
|
||||
{
|
||||
**event["content_block"],
|
||||
"input": json.loads("".join(part for index, part in input_json if index == event["index"])),
|
||||
}
|
||||
if any(index == event["index"] for index, _ in input_json)
|
||||
else event["content_block"]
|
||||
for name, event in events
|
||||
if name == "content_block_start"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.covers("other.provider_wire.bedrock.websearch_interception_streamed_capped_turn_ends_with_native_results")
|
||||
def test_streamed_web_search_turn_capped_by_max_agentic_loops_ends_turn_with_snippets_and_ordered_blocks(
|
||||
gateway: Gateway, tmp_path: Path
|
||||
|
|
@ -136,7 +154,7 @@ def test_streamed_web_search_turn_capped_by_max_agentic_loops_ends_turn_with_sni
|
|||
for earlier, earlier_event in events[:position]
|
||||
if earlier == "content_block_start"
|
||||
}, response.text
|
||||
started: Final = tuple(event["content_block"] for name, event in events if name == "content_block_start")
|
||||
started: Final = streamed_blocks(events)
|
||||
search_ids: Final = tuple(block["id"] for block in started if block["type"] == "server_tool_use")
|
||||
assert search_ids and all(search_id.startswith("srvtoolu_") for search_id in search_ids), response.text
|
||||
assert started[-1] == {"type": "text", "text": ""}, response.text
|
||||
|
|
@ -162,7 +180,12 @@ def test_streamed_web_search_turn_capped_by_max_agentic_loops_ends_turn_with_sni
|
|||
)
|
||||
), response.text
|
||||
assert (
|
||||
"".join(event["delta"]["text"] for name, event in events if name == "content_block_delta") == "turn-1"
|
||||
"".join(
|
||||
event["delta"]["text"]
|
||||
for name, event in events
|
||||
if name == "content_block_delta" and event["delta"]["type"] == "text_delta"
|
||||
)
|
||||
== "turn-1"
|
||||
), response.text
|
||||
assert [event["delta"]["stop_reason"] for name, event in events if name == "message_delta"] == [
|
||||
"end_turn"
|
||||
|
|
|
|||
|
|
@ -3,11 +3,14 @@ from __future__ import annotations
|
|||
import json
|
||||
import os
|
||||
import uuid
|
||||
from collections.abc import Mapping
|
||||
from hashlib import sha256
|
||||
from importlib.resources import files
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from integration._support.client import Gateway, Scenario
|
||||
from integration._support.client import Gateway, Scenario, eventually, list_value, object_value
|
||||
from integration._support.responses_vendor import same_response
|
||||
from integration._support.wire import Reply, Request, Wire, wire_server
|
||||
from integration._support.worker_readiness import wait_for_worker_readiness
|
||||
|
|
@ -38,6 +41,7 @@ _CHAT_SUFFIXES: Final[tuple[tuple[str, str], ...]] = (("", ""),) + tuple(
|
|||
(f"-thinking-{level}", level) for level in _CHAT_THINKING_LEVELS
|
||||
)
|
||||
_PROMPT: Final = "Inspect the requested change."
|
||||
_REQUEST_TAG: Final = "cursor-integration"
|
||||
_INPUT: Final = [{"role": "user", "content": [{"type": "input_text", "text": _PROMPT}]}]
|
||||
_CUSTOM_TOOL: Final = {
|
||||
"type": "custom",
|
||||
|
|
@ -107,16 +111,30 @@ def _stream(identity: str) -> tuple[bytes, ...]:
|
|||
) + (b"data: [DONE]\n\n",)
|
||||
|
||||
|
||||
def _chat_completion(stream: bool) -> bytes | tuple[bytes, ...]:
|
||||
def _chat_completion(stream: bool, identity: str) -> bytes | tuple[bytes, ...]:
|
||||
if stream:
|
||||
return (
|
||||
b'data: {"id":"chat_cursor_1","object":"chat.completion.chunk","created":1,"model":"gpt-5","choices":[{"index":0,"delta":{"content":"Patch inspected."},"finish_reason":null}]}\n\n',
|
||||
b'data: {"id":"chat_cursor_1","object":"chat.completion.chunk","created":1,"model":"gpt-5","choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}\n\n',
|
||||
b"data: [DONE]\n\n",
|
||||
)
|
||||
return tuple(
|
||||
f"data: {json.dumps(chunk)}\n\n".encode()
|
||||
for chunk in (
|
||||
{
|
||||
"id": identity,
|
||||
"object": "chat.completion.chunk",
|
||||
"created": 1,
|
||||
"model": _MODEL,
|
||||
"choices": [{"index": 0, "delta": {"content": "Patch inspected."}, "finish_reason": None}],
|
||||
},
|
||||
{
|
||||
"id": identity,
|
||||
"object": "chat.completion.chunk",
|
||||
"created": 1,
|
||||
"model": _MODEL,
|
||||
"choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}],
|
||||
},
|
||||
)
|
||||
) + (b"data: [DONE]\n\n",)
|
||||
return json.dumps(
|
||||
{
|
||||
"id": "chat_cursor_1",
|
||||
"id": identity,
|
||||
"object": "chat.completion",
|
||||
"created": 1,
|
||||
"model": _MODEL,
|
||||
|
|
@ -133,11 +151,12 @@ def _chat_completion(stream: bool) -> bytes | tuple[bytes, ...]:
|
|||
).encode()
|
||||
|
||||
|
||||
def _deployment(scenario: Scenario, wire: Wire) -> str:
|
||||
def _deployment(scenario: Scenario, wire: Wire, **pricing: JsonValue) -> str:
|
||||
model: Final = scenario.model(
|
||||
model=f"openai/{_MODEL}",
|
||||
api_key=_API_KEY,
|
||||
api_base=f"{wire.url}/v1",
|
||||
**pricing,
|
||||
)
|
||||
wait_for_worker_readiness(scenario.gateway, model, 1, _PROXY_WORKERS)
|
||||
return model
|
||||
|
|
@ -167,7 +186,7 @@ def test_cursor_responses_input_uses_responses_wire_and_chat_output(
|
|||
assert request.method == "POST"
|
||||
assert request.headers["authorization"] == f"Bearer {_API_KEY}", dict(request.headers)
|
||||
if request.target == "/v1/chat/completions":
|
||||
completion: Final = _chat_completion(stream)
|
||||
completion: Final = _chat_completion(stream, raw_id)
|
||||
return Reply(
|
||||
content_type="text/event-stream" if stream else "application/json",
|
||||
chunks=completion if isinstance(completion, tuple) else None,
|
||||
|
|
@ -175,7 +194,7 @@ def test_cursor_responses_input_uses_responses_wire_and_chat_output(
|
|||
)
|
||||
assert request.target == "/v1/responses", request.target
|
||||
body: Final = _JSON_OBJECT.validate_json(request.body)
|
||||
assert {key: value for key, value in body.items() if key != "metadata"} == {
|
||||
assert body == {
|
||||
"model": _MODEL,
|
||||
"input": _INPUT,
|
||||
**reasoning,
|
||||
|
|
@ -259,41 +278,69 @@ def test_cursor_responses_input_uses_responses_wire_and_chat_output(
|
|||
assert [request.target for request in drain_contract_requests(wire)] == ["/v1/responses"]
|
||||
|
||||
|
||||
def test_cursor_responses_body_carries_no_proxy_metadata(gateway: Gateway) -> None:
|
||||
pytest.skip(
|
||||
"BUG: /cursor/chat/completions forwards proxy-internal metadata (key hash, team, user, client ip, user agent) "
|
||||
"to the provider in the Responses body"
|
||||
)
|
||||
@pytest.mark.parametrize(
|
||||
"client_metadata", ({}, {"cursor_conversation": "integration"}), ids=("no_client_metadata", "client_metadata")
|
||||
)
|
||||
@pytest.mark.parametrize("stream", (False, True))
|
||||
def test_cursor_responses_body_sends_only_client_metadata_and_bills_the_key(
|
||||
gateway: Gateway,
|
||||
stream: bool,
|
||||
client_metadata: Mapping[str, JsonValue],
|
||||
) -> None:
|
||||
raw_id: Final = f"resp_cursor_metadata_{uuid.uuid4().hex}"
|
||||
alias: Final = f"integration-cursor-{uuid.uuid4().hex}"
|
||||
sent_metadata: Final = {"metadata": client_metadata} if client_metadata else {}
|
||||
|
||||
def respond(request: Request) -> Reply:
|
||||
if request.method == "GET" and request.target == "/v1/models":
|
||||
return model_discovery_reply(_MODEL)
|
||||
assert request.method == "POST" and request.target == "/v1/responses", request.target
|
||||
if stream:
|
||||
return Reply(content_type="text/event-stream", chunks=_stream(raw_id))
|
||||
return Reply(body=json.dumps(_response(raw_id)).encode())
|
||||
|
||||
with wire_server(respond) as wire, gateway.scenario() as scenario:
|
||||
model: Final = _deployment(scenario, wire)
|
||||
key: Final = scenario.key(key_alias=alias)
|
||||
hashed_key: Final = sha256(key.encode()).hexdigest()
|
||||
response: Final = gateway.request(
|
||||
"POST",
|
||||
"/cursor/chat/completions",
|
||||
{
|
||||
"model": f"{model}-thinking-high",
|
||||
"model": model,
|
||||
"input": _INPUT,
|
||||
"messages": [],
|
||||
"tools": [_CUSTOM_TOOL],
|
||||
"stream": stream,
|
||||
**sent_metadata,
|
||||
},
|
||||
key=key,
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
requests: Final = drain_contract_requests(wire)
|
||||
assert len(requests) == 1, requests
|
||||
body: Final = _JSON_OBJECT.validate_json(requests[0].body)
|
||||
assert body == {
|
||||
assert _JSON_OBJECT.validate_json(requests[0].body) == {
|
||||
"model": _MODEL,
|
||||
"input": _INPUT,
|
||||
"reasoning": {"effort": "high"},
|
||||
"tools": [_CUSTOM_TOOL],
|
||||
"stream": stream,
|
||||
**sent_metadata,
|
||||
}, requests[0].body
|
||||
(row,) = eventually(
|
||||
lambda: gateway.rows(
|
||||
'SELECT call_type, spend, metadata FROM "LiteLLM_SpendLogs" WHERE api_key = %s', (hashed_key,)
|
||||
),
|
||||
lambda rows: len(rows) == 1,
|
||||
seconds=70,
|
||||
)
|
||||
assert row["call_type"] == "aresponses", row
|
||||
assert object_value(row["metadata"])["user_api_key_alias"] == alias, row
|
||||
assert isinstance(row["spend"], float) and row["spend"] > 0, row
|
||||
eventually(
|
||||
lambda: gateway.rows('SELECT spend FROM "LiteLLM_VerificationToken" WHERE token = %s', (hashed_key,)),
|
||||
lambda tokens: len(tokens) == 1 and tokens[0]["spend"] == pytest.approx(row["spend"]),
|
||||
seconds=70,
|
||||
)
|
||||
|
||||
|
||||
def test_cursor_stream_reports_usage_when_the_client_asks_for_it(gateway: Gateway) -> None:
|
||||
|
|
@ -341,10 +388,12 @@ def test_cursor_messages_body_uses_chat_completions_wire(
|
|||
suffix: str,
|
||||
reasoning_effort: str,
|
||||
) -> None:
|
||||
chat_id: Final = f"chat_cursor_{uuid.uuid4().hex}"
|
||||
|
||||
def respond(request: Request) -> Reply:
|
||||
if request.method == "GET" and request.target == "/v1/models":
|
||||
return model_discovery_reply(_MODEL)
|
||||
completion: Final = _chat_completion(stream)
|
||||
completion: Final = _chat_completion(stream, chat_id)
|
||||
return Reply(
|
||||
content_type="text/event-stream" if stream else "application/json",
|
||||
chunks=completion if isinstance(completion, tuple) else None,
|
||||
|
|
@ -367,6 +416,7 @@ def test_cursor_messages_body_uses_chat_completions_wire(
|
|||
"stream": stream,
|
||||
},
|
||||
key=key,
|
||||
headers={"x-litellm-tags": _REQUEST_TAG},
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
if stream:
|
||||
|
|
@ -393,7 +443,7 @@ def test_cursor_messages_body_uses_chat_completions_wire(
|
|||
payload: Final = _JSON_OBJECT.validate_json(response.content)
|
||||
identity: Final = payload.get("id")
|
||||
created: Final = payload.get("created")
|
||||
assert isinstance(identity, str) and same_response(identity, "chat_cursor_1"), response.text
|
||||
assert isinstance(identity, str) and same_response(identity, chat_id), response.text
|
||||
assert isinstance(created, int) and created > 0, response.text
|
||||
assert payload == {
|
||||
"id": identity,
|
||||
|
|
@ -456,3 +506,72 @@ def test_cursor_messages_body_uses_chat_completions_wire(
|
|||
**({"stream_options": {"include_usage": True}} if stream else {}),
|
||||
**({"reasoning_effort": reasoning_effort} if reasoning_effort else {}),
|
||||
}, requests[0].body
|
||||
(row,) = eventually(
|
||||
lambda: gateway.rows(
|
||||
'SELECT call_type, team_id, request_tags FROM "LiteLLM_SpendLogs" WHERE api_key = %s',
|
||||
(sha256(key.encode()).hexdigest(),),
|
||||
),
|
||||
lambda rows: len(rows) == 1,
|
||||
seconds=70,
|
||||
)
|
||||
assert (row["call_type"], row["team_id"]) == ("acompletion", team), row
|
||||
assert _REQUEST_TAG in list_value(row["request_tags"]), row
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("body_shape", "metadata_as_json_string"), (("messages", False), ("input", False), ("input", True))
|
||||
)
|
||||
def test_cursor_body_keeps_client_and_key_tags_and_enforces_the_client_tag_budget(
|
||||
gateway: Gateway, body_shape: str, metadata_as_json_string: bool
|
||||
) -> None:
|
||||
chat_id: Final = f"chat_cursor_{uuid.uuid4().hex}"
|
||||
response_id: Final = f"resp_cursor_tags_{uuid.uuid4().hex}"
|
||||
budget_tag: Final = f"cursor-tag-budget-{uuid.uuid4().hex}"
|
||||
key_tag: Final = f"cursor-key-tag-{uuid.uuid4().hex}"
|
||||
client_metadata: Final = {"tags": [budget_tag]}
|
||||
|
||||
def respond(request: Request) -> Reply:
|
||||
if request.method == "GET" and request.target == "/v1/models":
|
||||
return model_discovery_reply(_MODEL)
|
||||
if request.target == "/v1/responses":
|
||||
return Reply(body=json.dumps(_response(response_id)).encode())
|
||||
completion: Final = _chat_completion(False, chat_id)
|
||||
return Reply(body=completion if isinstance(completion, bytes) else b"")
|
||||
|
||||
def delete_tag() -> None:
|
||||
gateway.post("/tag/delete", {"name": budget_tag})
|
||||
|
||||
with wire_server(respond) as wire, gateway.scenario() as scenario:
|
||||
model: Final = _deployment(scenario, wire, input_cost_per_token=0.01, output_cost_per_token=0.01)
|
||||
key: Final = scenario.key(models=[model], metadata={"tags": [key_tag]})
|
||||
gateway.post("/tag/new", {"name": budget_tag, "max_budget": 0.0001})
|
||||
scenario.cleanups.callback(delete_tag)
|
||||
|
||||
def tagged_chat(text: str) -> httpx.Response:
|
||||
return gateway.request(
|
||||
"POST",
|
||||
"/cursor/chat/completions",
|
||||
{
|
||||
"model": model,
|
||||
body_shape: [{"role": "user", "content": text}],
|
||||
"metadata": json.dumps(client_metadata) if metadata_as_json_string else client_metadata,
|
||||
},
|
||||
key=key,
|
||||
)
|
||||
|
||||
first: Final = tagged_chat("first tagged chat")
|
||||
assert first.status_code == 200, first.text
|
||||
(row,) = eventually(
|
||||
lambda: gateway.rows(
|
||||
'SELECT request_tags FROM "LiteLLM_SpendLogs" WHERE api_key = %s',
|
||||
(sha256(key.encode()).hexdigest(),),
|
||||
),
|
||||
lambda rows: len(rows) == 1,
|
||||
seconds=70,
|
||||
)
|
||||
assert {budget_tag, key_tag} <= set(list_value(row["request_tags"])), row
|
||||
blocked: Final = eventually(
|
||||
lambda: tagged_chat("tag budget probe"), lambda response: response.status_code != 200, seconds=70
|
||||
)
|
||||
assert blocked.status_code == 422, blocked.text
|
||||
assert "budget" in blocked.text.lower(), blocked.text
|
||||
|
|
|
|||
|
|
@ -153,15 +153,16 @@ def _foundry_deployment(rig: _Rig, scenario: Scenario, *, bare: bool) -> str:
|
|||
return name
|
||||
|
||||
|
||||
def _assert_foundry_chat_wire(received: Request, prompt: str, *, options: JsonValue | None) -> None:
|
||||
def _assert_foundry_chat_wire(received: Request, user_content: JsonValue, *, options: JsonValue | None) -> None:
|
||||
body: Final = body_of(received)
|
||||
marked_system: Final = [{"type": "text", "text": SYSTEM, "prompt_cache_breakpoint": EXPLICIT}]
|
||||
assert urlsplit(received.target).netloc == _FOUNDRY_HOST, received.target
|
||||
assert urlsplit(received.target).path == _FOUNDRY_CHAT_PATH, received.target
|
||||
assert "prompt_cache_breakpoint" not in received.body.decode(), received.body
|
||||
assert body["model"] == _FOUNDRY_MODEL, received.body
|
||||
assert body["messages"] == [{"role": "system", "content": SYSTEM}, {"role": "user", "content": prompt}], (
|
||||
received.body
|
||||
)
|
||||
assert body["messages"] == [
|
||||
{"role": "system", "content": marked_system},
|
||||
{"role": "user", "content": user_content},
|
||||
], received.body
|
||||
assert body.get("prompt_cache_options") == options, received.body
|
||||
|
||||
|
||||
|
|
@ -185,7 +186,7 @@ def test_b9_b12_foundry_gpt6_responses_reach_the_openai_dialect_on_the_foundry_h
|
|||
|
||||
|
||||
@pytest.mark.timeout(_CELL_TIMEOUT)
|
||||
def test_b10_foundry_gpt6_messages_bridge_to_chat_without_a_breakpoint(rig: _Rig) -> None:
|
||||
def test_b10_foundry_gpt6_messages_bridge_to_chat_keeps_the_injected_breakpoint(rig: _Rig) -> None:
|
||||
marker: Final = fresh_marker()
|
||||
with rig.gateway.scenario() as scenario:
|
||||
name: Final = _foundry_deployment(rig, scenario, bare=False)
|
||||
|
|
@ -193,12 +194,12 @@ def test_b10_foundry_gpt6_messages_bridge_to_chat_without_a_breakpoint(rig: _Rig
|
|||
rig.gateway, rig.wire, "messages", request_body("messages", name, prompt_text(marker))
|
||||
)
|
||||
assert_answered(outcome, marker)
|
||||
_assert_foundry_chat_wire(received, prompt_text(marker), options=IMPLICIT)
|
||||
_assert_foundry_chat_wire(received, [{"type": "text", "text": prompt_text(marker)}], options=IMPLICIT)
|
||||
success_row(name, marker)
|
||||
|
||||
|
||||
@pytest.mark.timeout(_CELL_TIMEOUT)
|
||||
def test_b11_foundry_gpt6_chat_stays_on_the_plain_chat_wire(rig: _Rig) -> None:
|
||||
def test_b11_foundry_gpt6_chat_stays_on_the_chat_wire_with_the_injected_breakpoint(rig: _Rig) -> None:
|
||||
marker: Final = fresh_marker()
|
||||
with rig.gateway.scenario() as scenario:
|
||||
name: Final = _foundry_deployment(rig, scenario, bare=False)
|
||||
|
|
|
|||
|
|
@ -355,7 +355,13 @@ def _voice_scenario_created_after_the_first_frame(transcript: str) -> RealtimeRe
|
|||
|
||||
|
||||
def _live_scenario() -> RealtimeResponse:
|
||||
return RealtimeResponse(content_type="application/x-realtime", events=())
|
||||
return RealtimeResponse(
|
||||
content_type="application/x-realtime",
|
||||
events=(
|
||||
{"serverContent": {"modelTurn": {"parts": [{"text": "scripted $REQUEST_ID"}]}}},
|
||||
{"serverContent": {"turnComplete": True}},
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def _realtime_update(turn_detection: JsonValue) -> dict[str, JsonValue]:
|
||||
|
|
@ -1348,10 +1354,14 @@ def test_gemini_ai_studio_voice_session_with_vad_tuning_connects_under_the_guard
|
|||
_ws_base(_owned_url(guardrail_proxy)),
|
||||
f"model={model}",
|
||||
key,
|
||||
(_step(_update_frame(update), until="session.updated"),),
|
||||
(
|
||||
_step(_update_frame(update), until="session.updated"),
|
||||
_step(_user_text_item(CLEAN_TRANSCRIPT), RESPONSE_CREATE, until=RESPONSE_DONE),
|
||||
),
|
||||
)
|
||||
observed: Final = _observed(upstream_url, handle.scenario_id)
|
||||
assert session.types == ("session.created", "session.updated"), session
|
||||
assert session.types[:2] == ("session.created", "session.updated"), session
|
||||
assert session.types[-1] == RESPONSE_DONE, session
|
||||
assert _setup_detection(observed) == {"disabled": True}, _sent(observed)
|
||||
|
||||
|
||||
|
|
@ -1365,10 +1375,14 @@ def test_vertex_live_voice_session_keeps_vad_tuning_under_the_guardrail(guardrai
|
|||
_ws_base(_owned_url(guardrail_proxy)),
|
||||
f"model={model}",
|
||||
key,
|
||||
(_step(_update_frame(_realtime_update(TUNED_SERVER_VAD)), until="session.updated"),),
|
||||
(
|
||||
_step(_update_frame(_realtime_update(TUNED_SERVER_VAD)), until="session.updated"),
|
||||
_step(_user_text_item(CLEAN_TRANSCRIPT), RESPONSE_CREATE, until=RESPONSE_DONE),
|
||||
),
|
||||
)
|
||||
observed: Final = _observed(upstream_url, handle.scenario_id)
|
||||
assert session.types == ("session.created", "session.updated"), session
|
||||
assert session.types[:2] == ("session.created", "session.updated"), session
|
||||
assert session.types[-1] == RESPONSE_DONE, session
|
||||
assert _setup_detection(observed) == GEMINI_TUNED_DETECTION, _sent(observed)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -0,0 +1,137 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import dataclasses
|
||||
import threading
|
||||
import uuid
|
||||
from collections.abc import Callable, Iterator
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from dataclasses import dataclass
|
||||
from datetime import UTC, datetime
|
||||
from pathlib import Path
|
||||
from queue import SimpleQueue
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from integration._support.client import JSON_OBJECT, gateway_from_environment
|
||||
from integration._support.openai_wire import chat_reply
|
||||
from integration._support.process import OwnedProxy, graceful_stop_seconds, owned_proxy_process
|
||||
from integration._support.wire import Reply, Request, Wire, wire_server
|
||||
from integration.providers._cache_control_marks_support import owned_config
|
||||
from integration.routing.test_payment_required_mapping import OPENAI_MODEL, is_model_info_probe, model_list_reply
|
||||
|
||||
pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120)
|
||||
|
||||
GROUPS: Final = ("otpm-fraction-first", "otpm-fraction-retry")
|
||||
OTPM_LIMIT: Final = 1000
|
||||
FRACTION: Final = 0.5
|
||||
DEFAULT_RESERVATION: Final = int(OTPM_LIMIT * FRACTION)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Rig:
|
||||
owned: OwnedProxy
|
||||
upstream: Wire
|
||||
gates: SimpleQueue[threading.Event]
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Admissions:
|
||||
held_forwarded_max_tokens: bool
|
||||
statuses: tuple[int, int, int]
|
||||
one_minute: bool
|
||||
|
||||
|
||||
def _reply_for(gates: SimpleQueue[threading.Event]) -> Callable[[Request], Reply]:
|
||||
def reply(request: Request) -> Reply:
|
||||
if is_model_info_probe(request):
|
||||
return model_list_reply()
|
||||
if JSON_OBJECT.validate_json(request.body).get("stream") is True:
|
||||
return dataclasses.replace(
|
||||
chat_reply("chatcmpl-held", OPENAI_MODEL, "held", stream=True), gate_before_first=gates.get(timeout=60)
|
||||
)
|
||||
return chat_reply("chatcmpl-capped", OPENAI_MODEL, "capped", stream=False)
|
||||
|
||||
return reply
|
||||
|
||||
|
||||
def _config(directory: Path, upstream: Wire) -> Path:
|
||||
return owned_config(
|
||||
directory,
|
||||
[
|
||||
{
|
||||
"model_name": group,
|
||||
"litellm_params": {
|
||||
"model": OPENAI_MODEL,
|
||||
"api_base": upstream.url + "/v1",
|
||||
"api_key": "synthetic-openai-key",
|
||||
"otpm": OTPM_LIMIT,
|
||||
},
|
||||
"model_info": {"id": group},
|
||||
}
|
||||
for group in GROUPS
|
||||
],
|
||||
litellm_settings={"cache": False},
|
||||
router_settings={"optional_pre_call_checks": ["enforce_model_rate_limits"], "num_retries": 0},
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]:
|
||||
directory: Final = tmp_path_factory.mktemp("otpm-fraction")
|
||||
gates: Final[SimpleQueue[threading.Event]] = SimpleQueue()
|
||||
with (
|
||||
gateway_from_environment() as gateway,
|
||||
wire_server(_reply_for(gates)) as upstream,
|
||||
owned_proxy_process(
|
||||
gateway,
|
||||
directory,
|
||||
{"NO_MAX_TOKENS_TPM_RESERVATION_FRACTION": str(FRACTION)},
|
||||
config=_config(directory, upstream),
|
||||
) as owned,
|
||||
):
|
||||
yield Rig(owned, upstream, gates)
|
||||
|
||||
|
||||
def _chat(rig: Rig, group: str, **params: int | bool) -> httpx.Response:
|
||||
return rig.owned.gateway.request(
|
||||
"POST",
|
||||
"/v1/chat/completions",
|
||||
{"model": group, "messages": [{"role": "user", "content": f"fraction {uuid.uuid4().hex}"}], **params},
|
||||
)
|
||||
|
||||
|
||||
def _await_forwarded_chat(upstream: Wire) -> Request:
|
||||
request: Final = upstream.received.get(timeout=60)
|
||||
return _await_forwarded_chat(upstream) if is_model_info_probe(request) else request
|
||||
|
||||
|
||||
def _utc_minute() -> str:
|
||||
return datetime.now(UTC).strftime("%Y-%m-%d %H:%M")
|
||||
|
||||
|
||||
def _admissions_beside_a_held_stream(rig: Rig, group: str) -> Admissions:
|
||||
release: Final = threading.Event()
|
||||
rig.gates.put(release)
|
||||
started: Final = _utc_minute()
|
||||
with ThreadPoolExecutor(max_workers=1) as pool:
|
||||
held: Final = pool.submit(_chat, rig, group, stream=True)
|
||||
try:
|
||||
forwarded: Final = JSON_OBJECT.validate_json(_await_forwarded_chat(rig.upstream).body)
|
||||
over: Final = _chat(rig, group, max_tokens=OTPM_LIMIT - DEFAULT_RESERVATION + 1)
|
||||
fits: Final = _chat(rig, group, max_tokens=OTPM_LIMIT - DEFAULT_RESERVATION)
|
||||
ended: Final = _utc_minute()
|
||||
finally:
|
||||
release.set()
|
||||
return Admissions(
|
||||
held_forwarded_max_tokens="max_tokens" in forwarded,
|
||||
statuses=(over.status_code, fits.status_code, held.result(timeout=60).status_code),
|
||||
one_minute=started == ended,
|
||||
)
|
||||
|
||||
|
||||
def test_reservation_fraction_env_var_sets_what_a_request_without_max_tokens_holds(rig: Rig) -> None:
|
||||
first: Final = _admissions_beside_a_held_stream(rig, GROUPS[0])
|
||||
admissions: Final = first if first.one_minute else _admissions_beside_a_held_stream(rig, GROUPS[1])
|
||||
|
||||
assert admissions == Admissions(held_forwarded_max_tokens=False, statuses=(429, 200, 200), one_minute=True)
|
||||
|
|
@ -3699,6 +3699,20 @@ def test_get_guardrail_from_metadata() -> None:
|
|||
) == [{"test-guardrail": {"extra_body": {"key": "value"}}}]
|
||||
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"caller_metadata",
|
||||
[{"guardrails": []}, {"guardrails": ["not-a-guardrail"]}, {"guardrails": None}],
|
||||
)
|
||||
def test_caller_metadata_guardrails_never_hide_the_proxy_guardrails_in_litellm_metadata(caller_metadata: dict) -> None:
|
||||
guardrail: Final = CustomGuardrail(
|
||||
guardrail_name="key-guardrail", event_hook=GuardrailEventHooks.pre_call, default_on=False
|
||||
)
|
||||
request_data: Final = {"metadata": caller_metadata, "litellm_metadata": {"guardrails": ["key-guardrail"]}}
|
||||
|
||||
assert guardrail.should_run_guardrail(request_data, GuardrailEventHooks.pre_call) is True
|
||||
|
||||
|
||||
def test_guardrail_is_in_requested_guardrails() -> None:
|
||||
guardrail: Final = CustomGuardrail(guardrail_name="test-guardrail")
|
||||
|
||||
|
|
|
|||
|
|
@ -5,7 +5,9 @@ Tests the short-circuit path that detects web-search-only /v1/messages requests
|
|||
and executes the search directly without routing through the backend LLM.
|
||||
"""
|
||||
|
||||
import datetime
|
||||
import json
|
||||
from collections.abc import Iterator, Mapping, Sequence
|
||||
from typing import Final
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
|
|
@ -16,6 +18,7 @@ import litellm
|
|||
from litellm.integrations.websearch_interception.handler import (
|
||||
WebSearchInterceptionLogger,
|
||||
)
|
||||
from litellm.messages.dispatch import MessagesResult
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Detection tests
|
||||
|
|
@ -73,6 +76,42 @@ class TestTryShortCircuitSearch:
|
|||
assert "Result" in text_block["text"]
|
||||
mock_search.assert_called_once_with("Search for Claude Code releases")
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_emits_native_blocks_when_tools_already_converted(self, perplexity_search: respx.Route):
|
||||
"""Pre-request/deployment hooks convert the native tool to the LiteLLM
|
||||
standard shape before the short-circuit runs and record that in
|
||||
WEBSEARCH_EMIT_NATIVE_BLOCKS_KEY; the flag alone must be enough for
|
||||
native blocks to be emitted."""
|
||||
from litellm.integrations.websearch_interception.handler import (
|
||||
WEBSEARCH_EMIT_NATIVE_BLOCKS_KEY,
|
||||
)
|
||||
|
||||
logger: Final = WebSearchInterceptionLogger(enabled_providers=["hosted_vllm"])
|
||||
|
||||
result: Final = await logger.try_short_circuit_search(
|
||||
model="hosted_vllm/some-model",
|
||||
messages=[{"role": "user", "content": "Search for something"}],
|
||||
tools=[
|
||||
{
|
||||
"name": "litellm_web_search",
|
||||
"description": "Search the web for information.",
|
||||
"input_schema": {
|
||||
"type": "object",
|
||||
"properties": {"query": {"type": "string"}},
|
||||
"required": ["query"],
|
||||
},
|
||||
}
|
||||
],
|
||||
custom_llm_provider="hosted_vllm",
|
||||
kwargs={WEBSEARCH_EMIT_NATIVE_BLOCKS_KEY: True},
|
||||
)
|
||||
|
||||
assert perplexity_search.call_count == 1
|
||||
assert result is not None
|
||||
assert [b["type"] for b in result["content"]] == ["server_tool_use", "web_search_tool_result", "text"]
|
||||
assert result["content"][0]["name"] == "web_search"
|
||||
assert [r["url"] for r in result["content"][1]["content"]] == ["https://example.com/releases"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_does_not_short_circuit_mixed_tools(self):
|
||||
"""Mix of web_search and other tools → NOT short-circuited"""
|
||||
|
|
@ -468,6 +507,89 @@ class TestShortCircuitEntryPoint:
|
|||
assert text_block["text"] == "results"
|
||||
|
||||
|
||||
class TestShortCircuitThroughMessagesEntryPoint:
|
||||
"""The pre-request hook renames the native tool before the short-circuit
|
||||
runs, so only a request through anthropic_messages shows whether native
|
||||
clients still get their server_tool_use and web_search_tool_result blocks."""
|
||||
|
||||
NATIVE_TOOL: Final = {"type": "web_search_20250305", "name": "web_search"}
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _intercept_web_search(self, perplexity_search: respx.Route) -> Iterator[respx.Route]:
|
||||
with patch("litellm.callbacks", [WebSearchInterceptionLogger(enabled_providers=["bedrock"])]):
|
||||
yield perplexity_search
|
||||
|
||||
@staticmethod
|
||||
async def _send(*, tools: Sequence[Mapping[str, object]], stream: bool, **kwargs: object) -> MessagesResult:
|
||||
return await litellm.anthropic_messages(
|
||||
model="bedrock/converse/us.anthropic.claude-sonnet-5-5",
|
||||
messages=[{"role": "user", "content": "latest python release"}],
|
||||
tools=tools,
|
||||
max_tokens=64,
|
||||
stream=stream,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_native_tool_gets_native_blocks(self, perplexity_search: respx.Route):
|
||||
response: Final = await self._send(tools=[self.NATIVE_TOOL], stream=False)
|
||||
|
||||
assert perplexity_search.call_count == 1
|
||||
assert [block["type"] for block in response["content"]] == [
|
||||
"server_tool_use",
|
||||
"web_search_tool_result",
|
||||
"text",
|
||||
]
|
||||
server_tool_use, tool_result, text = response["content"]
|
||||
assert server_tool_use["name"] == "web_search"
|
||||
assert server_tool_use["input"] == {"query": "latest python release"}
|
||||
assert tool_result["tool_use_id"] == server_tool_use["id"]
|
||||
assert [result["url"] for result in tool_result["content"]] == ["https://example.com/releases"]
|
||||
assert "Release notes" in text["text"]
|
||||
assert "Release notes" not in json.dumps(tool_result)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_native_blocks_reply_can_be_costed(self):
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging
|
||||
|
||||
logging_obj: Final = Logging(
|
||||
model="bedrock/converse/us.anthropic.claude-sonnet-5-5",
|
||||
messages=[{"role": "user", "content": "latest python release"}],
|
||||
stream=False,
|
||||
call_type="anthropic_messages",
|
||||
start_time=datetime.datetime(2026, 1, 1),
|
||||
litellm_call_id="short-circuit-cost",
|
||||
function_id="short-circuit-cost",
|
||||
)
|
||||
logging_obj.update_environment_variables(litellm_params={}, optional_params={})
|
||||
response: Final = await self._send(tools=[self.NATIVE_TOOL], stream=False, litellm_logging_obj=logging_obj)
|
||||
|
||||
assert logging_obj.response_cost_calculator(result=response) == 0.0
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_native_tool_streams_native_blocks(self):
|
||||
stream: Final = await self._send(tools=[self.NATIVE_TOOL], stream=True)
|
||||
body: Final = b"".join([chunk async for chunk in stream])
|
||||
|
||||
starts: Final = [
|
||||
json.loads(line.removeprefix(b"data: "))["content_block"]
|
||||
for line in body.splitlines()
|
||||
if line.startswith(b"data: ") and b'"content_block_start"' in line
|
||||
]
|
||||
assert [block["type"] for block in starts] == ["server_tool_use", "web_search_tool_result", "text"]
|
||||
assert starts[1]["tool_use_id"] == starts[0]["id"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_litellm_standard_tool_stays_text_only(self):
|
||||
from litellm.integrations.websearch_interception.tools import (
|
||||
get_litellm_web_search_tool,
|
||||
)
|
||||
|
||||
response: Final = await self._send(tools=[get_litellm_web_search_tool()], stream=False)
|
||||
|
||||
assert [block["type"] for block in response["content"]] == ["text"]
|
||||
|
||||
|
||||
class TestShortCircuitSafeguardResults:
|
||||
"""A short-circuited reply never reaches the backend, so the proxy answers Claude Code's safeguards itself."""
|
||||
|
||||
|
|
|
|||
|
|
@ -69,9 +69,7 @@ def test_preserve_upstream_non_openai_attributes_keeps_scalar_and_structured_val
|
|||
scalar_attribute="scalar",
|
||||
structured_attribute={"items": ["value"]},
|
||||
)
|
||||
model_response: Final = ModelResponseStream(
|
||||
choices=[StreamingChoices(index=0, delta=Delta(content="hello"))]
|
||||
)
|
||||
model_response: Final = ModelResponseStream(choices=[StreamingChoices(index=0, delta=Delta(content="hello"))])
|
||||
|
||||
preserve_upstream_non_openai_attributes(model_response, original_chunk)
|
||||
|
||||
|
|
@ -262,9 +260,7 @@ class TestMapFinishReasonAnthropic:
|
|||
("content_filtered", "content_filter"),
|
||||
],
|
||||
)
|
||||
def test_anthropic_finish_reasons(
|
||||
self, provider_reason: str, expected: str
|
||||
) -> None:
|
||||
def test_anthropic_finish_reasons(self, provider_reason: str, expected: str) -> None:
|
||||
assert map_finish_reason(provider_reason) == expected
|
||||
|
||||
def test_refusal(self):
|
||||
|
|
@ -336,9 +332,7 @@ class TestMapFinishReasonZhipu:
|
|||
|
||||
|
||||
class TestMapFinishReasonOpenAIPassthrough:
|
||||
@pytest.mark.parametrize(
|
||||
"reason", ["stop", "length", "tool_calls", "function_call", "content_filter"]
|
||||
)
|
||||
@pytest.mark.parametrize("reason", ["stop", "length", "tool_calls", "function_call", "content_filter"])
|
||||
def test_openai_values_pass_through(self, reason):
|
||||
assert map_finish_reason(reason) == reason
|
||||
|
||||
|
|
@ -349,9 +343,7 @@ class TestMapFinishReasonGenericError:
|
|||
assert map_finish_reason("error") == "stop"
|
||||
|
||||
def test_lowercase_error_does_not_warn(self, mocker):
|
||||
warn = mocker.patch(
|
||||
"litellm.litellm_core_utils.core_helpers.verbose_logger.warning"
|
||||
)
|
||||
warn = mocker.patch("litellm.litellm_core_utils.core_helpers.verbose_logger.warning")
|
||||
assert map_finish_reason("error") == "stop"
|
||||
warn.assert_not_called()
|
||||
|
||||
|
|
@ -369,8 +361,7 @@ class TestFinishReasonMapOutputsAreValid:
|
|||
"""Every value in _FINISH_REASON_MAP must be a valid OpenAI finish reason."""
|
||||
for provider_reason, openai_reason in _FINISH_REASON_MAP.items():
|
||||
assert openai_reason in VALID_OPENAI_FINISH_REASONS, (
|
||||
f"Mapped value '{openai_reason}' (from '{provider_reason}') "
|
||||
f"is not a valid OpenAI finish reason"
|
||||
f"Mapped value '{openai_reason}' (from '{provider_reason}') is not a valid OpenAI finish reason"
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -380,28 +371,18 @@ class TestRedactNestedMatchAndRegexKeys:
|
|||
"assessments": [
|
||||
{
|
||||
"sensitiveInformationPolicy": {
|
||||
"piiEntities": [
|
||||
{"type": "NAME", "match": "secret-name", "action": "BLOCKED"}
|
||||
]
|
||||
},
|
||||
"wordPolicy": {
|
||||
"customWords": [{"match": "badword", "action": "BLOCKED"}]
|
||||
"piiEntities": [{"type": "NAME", "match": "secret-name", "action": "BLOCKED"}]
|
||||
},
|
||||
"wordPolicy": {"customWords": [{"match": "badword", "action": "BLOCKED"}]},
|
||||
}
|
||||
],
|
||||
"regex": "should-redact-key-named-regex",
|
||||
}
|
||||
out = redact_nested_match_and_regex_keys(payload)
|
||||
assert out["assessments"][0]["sensitiveInformationPolicy"]["piiEntities"][0][
|
||||
"match"
|
||||
] == "[REDACTED]"
|
||||
assert out["assessments"][0]["wordPolicy"]["customWords"][0]["match"] == (
|
||||
"[REDACTED]"
|
||||
)
|
||||
assert out["assessments"][0]["sensitiveInformationPolicy"]["piiEntities"][0]["match"] == "[REDACTED]"
|
||||
assert out["assessments"][0]["wordPolicy"]["customWords"][0]["match"] == ("[REDACTED]")
|
||||
assert out["regex"] == "[REDACTED]"
|
||||
assert payload["assessments"][0]["sensitiveInformationPolicy"]["piiEntities"][
|
||||
0
|
||||
]["match"] == "secret-name"
|
||||
assert payload["assessments"][0]["sensitiveInformationPolicy"]["piiEntities"][0]["match"] == "secret-name"
|
||||
|
||||
def test_passes_through_none_and_str(self):
|
||||
assert redact_nested_match_and_regex_keys(None) is None
|
||||
|
|
@ -564,13 +545,17 @@ class TestIsExpectedClientError:
|
|||
assert is_expected_client_error(over_budget) is True
|
||||
|
||||
litellm_limit = RateLimitError(
|
||||
message="key over rpm", llm_provider="anthropic", model="claude-haiku-4-5",
|
||||
message="key over rpm",
|
||||
llm_provider="anthropic",
|
||||
model="claude-haiku-4-5",
|
||||
category=RateLimitErrorCategory.LITELLM_RATE_LIMIT,
|
||||
)
|
||||
assert is_expected_client_error(litellm_limit) is True
|
||||
|
||||
vendor_limit = RateLimitError(
|
||||
message="rate limited upstream", llm_provider="anthropic", model="claude-haiku-4-5",
|
||||
message="rate limited upstream",
|
||||
llm_provider="anthropic",
|
||||
model="claude-haiku-4-5",
|
||||
category=RateLimitErrorCategory.VENDOR_RATE_LIMIT,
|
||||
)
|
||||
assert is_expected_client_error(vendor_limit) is False
|
||||
|
|
@ -803,3 +788,34 @@ class TestReceiveStampSurvivesASecondMetadataBucket:
|
|||
}
|
||||
|
||||
assert get_litellm_metadata_from_kwargs(kwargs)["litellm_received_at"] is own
|
||||
|
||||
|
||||
def test_billed_usage_cost_bills_any_chat_upstream_without_its_own_reader() -> None:
|
||||
from types import SimpleNamespace
|
||||
|
||||
from litellm.litellm_core_utils.core_helpers import billed_usage_cost
|
||||
from litellm.types.utils import ModelResponse, Usage
|
||||
|
||||
usage = Usage(prompt_tokens=100, completion_tokens=20, total_tokens=120, cost=0.037)
|
||||
chat = ModelResponse(usage=usage)
|
||||
assert billed_usage_cost(chat, "openai") == 0.037
|
||||
assert billed_usage_cost(chat, "openrouter") == 0.037
|
||||
assert billed_usage_cost(chat, "xai") is None
|
||||
assert billed_usage_cost(SimpleNamespace(usage=usage), "openai") == 0.037
|
||||
|
||||
|
||||
def test_billed_usage_cost_keeps_the_upstream_charge_after_litellm_wrote_its_bill_into_usage_cost() -> None:
|
||||
from litellm.litellm_core_utils.core_helpers import billed_usage_cost, stamp_litellm_bill_on_usage
|
||||
from litellm.types.utils import ModelResponse, Usage
|
||||
|
||||
reported = ModelResponse(usage=Usage(prompt_tokens=100, completion_tokens=20, total_tokens=120, cost=0.037))
|
||||
stamp_litellm_bill_on_usage(reported, reported.usage, 0.0407)
|
||||
stamp_litellm_bill_on_usage(reported, reported.usage, 0.0407)
|
||||
assert reported.usage.cost == 0.0407
|
||||
assert reported._hidden_params["response_cost"] == 0.0407
|
||||
assert billed_usage_cost(reported, "openai") == 0.037
|
||||
|
||||
unreported = ModelResponse(usage=Usage(prompt_tokens=100, completion_tokens=20, total_tokens=120))
|
||||
stamp_litellm_bill_on_usage(unreported, unreported.usage, 0.0015)
|
||||
assert unreported.usage.cost == 0.0015
|
||||
assert billed_usage_cost(unreported, "openai") is None
|
||||
|
|
|
|||
|
|
@ -6,6 +6,7 @@ import pytest
|
|||
from litellm._logging import verbose_logger
|
||||
from litellm.litellm_core_utils.env_utils import (
|
||||
get_env_float,
|
||||
get_env_float_in_range,
|
||||
get_env_float_or_none,
|
||||
get_env_int,
|
||||
get_env_int_or_none,
|
||||
|
|
@ -94,3 +95,12 @@ def test_unparsable_stdout_log_limit_does_not_break_logging(
|
|||
assert "first line" in messages
|
||||
assert "second line" in messages
|
||||
assert len([message for message in messages if raw in message]) == 1
|
||||
|
||||
|
||||
@pytest.mark.parametrize(("raw", "expected"), [("0.5", 0.5), ("1", 1.0), ("4", 0.25), ("0", 0.25), ("half", 0.25)])
|
||||
def test_float_in_range_keeps_values_inside_the_range_and_falls_back_otherwise(
|
||||
monkeypatch: pytest.MonkeyPatch, raw: str, expected: float
|
||||
) -> None:
|
||||
monkeypatch.setenv("ENV_UTILS_TEST_FRACTION", raw)
|
||||
|
||||
assert get_env_float_in_range("ENV_UTILS_TEST_FRACTION", 0.25, 0.01, 1.0) == expected
|
||||
|
|
|
|||
|
|
@ -5,6 +5,7 @@ from typing import Any, Final, Literal, Optional, Union
|
|||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import httpx
|
||||
import httpx2
|
||||
import openai
|
||||
import pytest
|
||||
from fastapi import HTTPException
|
||||
|
|
@ -24,6 +25,7 @@ from litellm.llms.base_llm.chat.transformation import BaseLLMException
|
|||
from litellm.llms.bedrock.common_utils import BedrockError
|
||||
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
|
||||
from litellm.llms.openai.common_utils import OpenAIError
|
||||
from litellm.llms.vertex_ai.common_utils import VertexAIError
|
||||
from litellm.proxy.common_utils.openai_error_payload import openai_error_type
|
||||
from litellm.types.utils import LlmProviders
|
||||
|
||||
|
|
@ -54,8 +56,8 @@ def test_openai_cyber_policy_uses_the_structured_error_code(shape: str, code: st
|
|||
assert caught.value.code == code
|
||||
assert caught.value.body == body
|
||||
assert caught.value.status_code == 400
|
||||
assert caught.value.response is response
|
||||
assert caught.value.request is request
|
||||
assert caught.value.response.status_code == 400
|
||||
assert caught.value.request.url == request.url
|
||||
assert caught.value.response.headers["x-request-id"] == "cyber-policy-request"
|
||||
|
||||
|
||||
|
|
@ -2938,3 +2940,110 @@ def test_transport_timeout_maps_to_timeout(
|
|||
|
||||
assert exc_info.value.status_code == 408
|
||||
assert f"Request timed out. Error_str: {reason}" in str(exc_info.value)
|
||||
|
||||
|
||||
def test_bedrock_server_error_keeps_the_aws_body_without_the_signed_request() -> None:
|
||||
body: Final = '{"message":"Bedrock is unable to process your request."}'
|
||||
signed_request: Final = httpx.Request(
|
||||
"POST",
|
||||
"https://bedrock-runtime.us-east-1.amazonaws.com/model/anthropic.claude-opus-5-5/converse",
|
||||
headers={"authorization": "AWS4-HMAC-SHA256 Credential=AKIDSIGNED/20261029/us-east-1/bedrock/aws4_request"},
|
||||
content=b'{"messages":[{"role":"user","content":[{"text":"hi"}]}]}',
|
||||
)
|
||||
upstream: Final = httpx.Response(
|
||||
500, content=body, headers={"x-amzn-requestid": "req-aws"}, request=signed_request
|
||||
)
|
||||
|
||||
with pytest.raises(litellm.ServiceUnavailableError) as raised:
|
||||
exception_type(
|
||||
model="bedrock/anthropic.claude-opus-5-5",
|
||||
original_exception=BedrockError(status_code=500, message=body, response=upstream),
|
||||
custom_llm_provider="bedrock",
|
||||
)
|
||||
|
||||
assert raised.value.response.text == body
|
||||
assert raised.value.response.headers["x-amzn-requestid"] == "req-aws"
|
||||
assert "AKIDSIGNED" not in str(raised.value.request.headers)
|
||||
assert raised.value.request.content == b""
|
||||
|
||||
|
||||
@pytest.mark.parametrize("provider", ("gemini", "vertex_ai"))
|
||||
def test_gemini_server_error_response_carries_the_upstream_error_body(provider: str) -> None:
|
||||
body: Final = '{"error":{"code":500,"message":"Internal error encountered.","status":"INTERNAL"}}'
|
||||
|
||||
with pytest.raises(litellm.InternalServerError) as raised:
|
||||
exception_type(
|
||||
model="gemini-3.1-pro",
|
||||
original_exception=VertexAIError(status_code=500, message=body, headers=None),
|
||||
custom_llm_provider=provider,
|
||||
)
|
||||
|
||||
assert raised.value.response.status_code == 500
|
||||
assert raised.value.response.json() == json.loads(body)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("status_code", "message", "mapped_error"),
|
||||
(
|
||||
(500, "The server had an error processing your request. Sorry about that!", litellm.InternalServerError),
|
||||
(520, "Web server is returning an unknown error", litellm.InternalServerError),
|
||||
(504, "A timeout occurred", litellm.Timeout),
|
||||
(500, "Mistral API raised a streaming error", litellm.APIError),
|
||||
),
|
||||
)
|
||||
def test_openai_error_matched_by_its_message_keeps_the_upstream_body(
|
||||
status_code: int, message: str, mapped_error: type[Exception]
|
||||
) -> None:
|
||||
body: Final = json.dumps({"error": {"message": message, "type": "server_error", "param": None, "code": None}})
|
||||
sent_request: Final = httpx.Request(
|
||||
"POST",
|
||||
"https://api.openai.com/v1/chat/completions",
|
||||
headers={"authorization": "Bearer sk-openai-secret"},
|
||||
content=b'{"messages":[{"role":"user","content":"hi"}]}',
|
||||
)
|
||||
upstream: Final = httpx.Response(
|
||||
status_code, content=body, headers={"x-request-id": "req-openai"}, request=sent_request
|
||||
)
|
||||
|
||||
with pytest.raises(mapped_error) as raised:
|
||||
exception_type(
|
||||
model="gpt-5.6",
|
||||
original_exception=openai.APIStatusError(message, response=upstream, body=json.loads(body)),
|
||||
custom_llm_provider="openai",
|
||||
)
|
||||
|
||||
error_response: Final = getattr(raised.value, "response")
|
||||
assert error_response.text == body
|
||||
assert error_response.headers["x-request-id"] == "req-openai"
|
||||
assert "sk-openai-secret" not in str(error_response.request.headers)
|
||||
assert error_response.request.content == b""
|
||||
|
||||
|
||||
@pytest.mark.parametrize("custom_llm_provider", ("openai", "azure"))
|
||||
def test_sdk_error_on_httpx2_keeps_the_upstream_body_without_the_provider_credentials(
|
||||
custom_llm_provider: str,
|
||||
) -> None:
|
||||
body: Final = json.dumps(
|
||||
{"error": {"message": "Unsupported value: 'temperature'", "type": "invalid_request_error", "code": None}}
|
||||
)
|
||||
sent_request: Final = httpx2.Request(
|
||||
"POST",
|
||||
"https://resource.openai.azure.com/openai/v1/chat/completions?api-key=sdk-query-secret",
|
||||
headers={"authorization": "Bearer sk-openai-secret"},
|
||||
content=b'{"messages":[{"role":"user","content":"hi"}]}',
|
||||
)
|
||||
upstream: Final = httpx2.Response(400, content=body, headers={"x-request-id": "req-sdk3"}, request=sent_request)
|
||||
|
||||
with pytest.raises(litellm.BadRequestError) as raised:
|
||||
exception_type(
|
||||
model="gpt-5.6",
|
||||
original_exception=openai.BadRequestError(
|
||||
"Unsupported value: 'temperature'", response=upstream, body=json.loads(body)
|
||||
),
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
)
|
||||
|
||||
assert raised.value.response.text == body
|
||||
assert raised.value.response.headers["x-request-id"] == "req-sdk3"
|
||||
assert str(raised.value.response.request.url) == "https://resource.openai.azure.com/openai/v1/chat/completions"
|
||||
assert "sk-openai-secret" not in str(raised.value.response.request.headers)
|
||||
|
|
|
|||
|
|
@ -4,7 +4,7 @@ from litellm.litellm_core_utils.get_supported_openai_params import (
|
|||
get_supported_openai_params,
|
||||
)
|
||||
|
||||
BEDROCK_REAL_MODEL = "eu.anthropic.claude-haiku-4-5-20251001-v1:0"
|
||||
BEDROCK_REAL_MODEL = "us.anthropic.claude-haiku-4-5-20251001-v1:0"
|
||||
BEDROCK_LABEL = "claude-haiku-4-5"
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -8,6 +8,7 @@ from types import MappingProxyType
|
|||
from typing import Final, Literal
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
import responses
|
||||
import respx
|
||||
|
|
@ -19,6 +20,7 @@ from litellm.litellm_core_utils.health_check_helpers import (
|
|||
HealthCheckHelpers,
|
||||
default_health_check_mode,
|
||||
native_health_check_mode,
|
||||
wildcard_health_check_models,
|
||||
)
|
||||
from litellm.main import ahealth_check
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
|
|
@ -1449,6 +1451,339 @@ async def test_groq_wildcard_health_check_uses_a_concrete_model(
|
|||
assert response == {}
|
||||
|
||||
|
||||
_WILDCARD_CATALOG: Final = MappingProxyType(
|
||||
{
|
||||
"databricks": (
|
||||
"databricks/databricks-claude-fixture-d",
|
||||
"databricks/databricks-claude-fixture-c",
|
||||
"databricks/databricks-claude-fixture-b",
|
||||
"databricks/databricks-claude-fixture-a",
|
||||
"databricks/databricks-cheap-chat",
|
||||
"databricks/databricks-claude-fixture-retired",
|
||||
),
|
||||
"bedrock": ("meta.llama-fixture", "anthropic.claude-fixture"),
|
||||
"cloudflare": ("cloudflare/@cf/meta/fixture-chat",),
|
||||
"vertex_ai": ("claude-fixture@default",),
|
||||
}
|
||||
)
|
||||
_WILDCARD_CATALOG_COST: Final = MappingProxyType(
|
||||
{
|
||||
"databricks/databricks-claude-fixture-retired": 1e-10,
|
||||
"databricks/databricks-cheap-chat": 1e-9,
|
||||
"databricks/databricks-claude-fixture-a": 2e-9,
|
||||
"databricks/databricks-claude-fixture-b": 3e-9,
|
||||
"databricks/databricks-claude-fixture-c": 4e-9,
|
||||
"databricks/databricks-claude-fixture-d": 5e-9,
|
||||
"meta.llama-fixture": 1e-9,
|
||||
"anthropic.claude-fixture": 2e-9,
|
||||
"cloudflare/@cf/meta/fixture-chat": 1e-9,
|
||||
"claude-fixture@default": 1e-9,
|
||||
}
|
||||
)
|
||||
_DATABRICKS_FIXTURE_API_BASE: Final = "https://fixture.cloud.databricks.com/serving-endpoints"
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def wildcard_catalog(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
providers: Final = {model: provider for provider, models in _WILDCARD_CATALOG.items() for model in models}
|
||||
litellm.register_model(
|
||||
{
|
||||
model: {
|
||||
"litellm_provider": providers[model],
|
||||
"mode": "chat",
|
||||
"max_tokens": 1024,
|
||||
"input_cost_per_token": cost,
|
||||
"output_cost_per_token": cost,
|
||||
**({"deprecation_date": "2000-01-01"} if model.endswith("-retired") else {}),
|
||||
}
|
||||
for model, cost in _WILDCARD_CATALOG_COST.items()
|
||||
}
|
||||
)
|
||||
for provider, models in _WILDCARD_CATALOG.items():
|
||||
monkeypatch.setitem(litellm.models_by_provider, provider, list(models))
|
||||
|
||||
|
||||
_DATABRICKS_PROVIDER_PROBES: Final = (
|
||||
"databricks/databricks-claude-fixture-retired",
|
||||
"databricks/databricks-cheap-chat",
|
||||
"databricks/databricks-claude-fixture-a",
|
||||
)
|
||||
_BEDROCK_PROVIDER_PROBES: Final = ("bedrock/meta.llama-fixture", "bedrock/anthropic.claude-fixture")
|
||||
_DATABRICKS_SERVED_CLAUDE_PROBES: Final = (
|
||||
"databricks/databricks-claude-fixture-a",
|
||||
"databricks/databricks-claude-fixture-b",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("custom_llm_provider", "wildcard", "expected"),
|
||||
[
|
||||
("databricks", "databricks-claude-*", (*_DATABRICKS_SERVED_CLAUDE_PROBES, *_DATABRICKS_PROVIDER_PROBES[:2])),
|
||||
(
|
||||
"databricks",
|
||||
"databricks/databricks-claude-*",
|
||||
(*_DATABRICKS_SERVED_CLAUDE_PROBES, *_DATABRICKS_PROVIDER_PROBES[:2]),
|
||||
),
|
||||
("databricks", "*", (*_DATABRICKS_PROVIDER_PROBES[1:], _DATABRICKS_PROVIDER_PROBES[0])),
|
||||
("databricks", "databricks-claude-fixture-r*", _DATABRICKS_PROVIDER_PROBES),
|
||||
("databricks", "databricks-llama-*", _DATABRICKS_PROVIDER_PROBES),
|
||||
("databricks", "*-fixture", _DATABRICKS_PROVIDER_PROBES),
|
||||
("databricks", "team-endpoints/*", _DATABRICKS_PROVIDER_PROBES),
|
||||
("bedrock", "converse/anthropic.*", ("bedrock/converse/anthropic.claude-fixture", *_BEDROCK_PROVIDER_PROBES)),
|
||||
(
|
||||
"bedrock",
|
||||
"converse/*",
|
||||
(
|
||||
"bedrock/converse/meta.llama-fixture",
|
||||
"bedrock/converse/anthropic.claude-fixture",
|
||||
*_BEDROCK_PROVIDER_PROBES,
|
||||
),
|
||||
),
|
||||
("bedrock", "invoke/anthropic.*", ("bedrock/invoke/anthropic.claude-fixture", *_BEDROCK_PROVIDER_PROBES)),
|
||||
("bedrock", "us-east-1/anthropic.*", ("bedrock/us-east-1/anthropic.claude-fixture", *_BEDROCK_PROVIDER_PROBES)),
|
||||
("bedrock", "agent/anthropic.*", _BEDROCK_PROVIDER_PROBES),
|
||||
("bedrock", "nova/*", _BEDROCK_PROVIDER_PROBES),
|
||||
("bedrock", "nova-2/*", _BEDROCK_PROVIDER_PROBES),
|
||||
("bedrock", "mantle/*", _BEDROCK_PROVIDER_PROBES),
|
||||
("bedrock", "openai/*", _BEDROCK_PROVIDER_PROBES),
|
||||
("bedrock", "chat_completions/*", _BEDROCK_PROVIDER_PROBES),
|
||||
("bedrock", "arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/*", _BEDROCK_PROVIDER_PROBES),
|
||||
(
|
||||
"bedrock",
|
||||
"converse/arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/*",
|
||||
_BEDROCK_PROVIDER_PROBES,
|
||||
),
|
||||
("cloudflare", "@cf/*", ("cloudflare/@cf/meta/fixture-chat",)),
|
||||
("vertex_ai", "claude-*", ("vertex_ai/claude-fixture@default",)),
|
||||
("vertex_ai", "meta/*", ("vertex_ai/claude-fixture@default",)),
|
||||
("fixture-gateway", "team/*", ()),
|
||||
],
|
||||
)
|
||||
@pytest.mark.usefixtures("wildcard_catalog")
|
||||
def test_wildcard_health_check_probes_the_cheapest_served_chat_models_before_the_provider_cheapest(
|
||||
custom_llm_provider: str, wildcard: str, expected: tuple[str, ...]
|
||||
) -> None:
|
||||
assert wildcard_health_check_models(wildcard=wildcard, custom_llm_provider=custom_llm_provider) == expected
|
||||
|
||||
|
||||
def _fail_first_databricks_claude_probe(request: httpx.Request) -> httpx.Response:
|
||||
model: Final = json.loads(request.content)["model"]
|
||||
if model == "databricks-claude-fixture-a":
|
||||
return httpx.Response(500, json={"error_code": "INTERNAL_ERROR", "message": "fixture outage"})
|
||||
return httpx.Response(
|
||||
200,
|
||||
json={
|
||||
"id": "chatcmpl-health",
|
||||
"object": "chat.completion",
|
||||
"created": 1,
|
||||
"model": model,
|
||||
"choices": [{"index": 0, "message": {"role": "assistant", "content": "2"}, "finish_reason": "stop"}],
|
||||
"usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2},
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.usefixtures("wildcard_catalog")
|
||||
async def test_partial_wildcard_health_check_falls_back_within_the_wildcard_before_any_other_model(
|
||||
monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter
|
||||
) -> None:
|
||||
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
|
||||
litellm.in_memory_llm_clients_cache.flush_cache()
|
||||
upstream: Final = respx_mock.post(f"{_DATABRICKS_FIXTURE_API_BASE}/chat/completions").mock(
|
||||
side_effect=_fail_first_databricks_claude_probe
|
||||
)
|
||||
|
||||
response: Final = await ahealth_check(
|
||||
{
|
||||
"model": "databricks/databricks-claude-*",
|
||||
"api_key": "fake-key",
|
||||
"api_base": _DATABRICKS_FIXTURE_API_BASE,
|
||||
"messages": [{"role": "user", "content": "What is 1 + 1?"}],
|
||||
}
|
||||
)
|
||||
|
||||
probed_models: Final = [json.loads(call.request.content)["model"] for call in upstream.calls]
|
||||
assert probed_models[0] == "databricks-claude-fixture-a"
|
||||
assert probed_models[-1] == "databricks-claude-fixture-b"
|
||||
assert all(model.startswith("databricks-claude-") for model in probed_models)
|
||||
assert "databricks-claude-fixture-a" not in probed_models[1:]
|
||||
assert response == {}
|
||||
|
||||
|
||||
def _answer_only_the_retired_databricks_claude_model(request: httpx.Request) -> httpx.Response:
|
||||
model: Final = json.loads(request.content)["model"]
|
||||
if model != "databricks-claude-fixture-retired":
|
||||
return httpx.Response(404, json={"error_code": "ENDPOINT_NOT_FOUND", "message": f"{model} does not exist"})
|
||||
return _fail_first_databricks_claude_probe(request)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.usefixtures("wildcard_catalog")
|
||||
async def test_partial_wildcard_health_check_stays_healthy_while_any_model_the_provider_probes_answers(
|
||||
monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter
|
||||
) -> None:
|
||||
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
|
||||
litellm.in_memory_llm_clients_cache.flush_cache()
|
||||
upstream: Final = respx_mock.post(f"{_DATABRICKS_FIXTURE_API_BASE}/chat/completions").mock(
|
||||
side_effect=_answer_only_the_retired_databricks_claude_model
|
||||
)
|
||||
|
||||
response: Final = await ahealth_check(
|
||||
{
|
||||
"model": "databricks/databricks-claude-*",
|
||||
"api_key": "fake-key",
|
||||
"api_base": _DATABRICKS_FIXTURE_API_BASE,
|
||||
"messages": [{"role": "user", "content": "What is 1 + 1?"}],
|
||||
}
|
||||
)
|
||||
|
||||
assert [json.loads(call.request.content)["model"] for call in upstream.calls] == [
|
||||
"databricks-claude-fixture-a",
|
||||
"databricks-claude-fixture-b",
|
||||
"databricks-claude-fixture-retired",
|
||||
]
|
||||
assert response == {}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.usefixtures("wildcard_catalog")
|
||||
async def test_partial_wildcard_health_check_that_fails_on_every_probe_tries_its_own_models_then_the_provider_probes(
|
||||
monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter
|
||||
) -> None:
|
||||
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
|
||||
litellm.in_memory_llm_clients_cache.flush_cache()
|
||||
upstream: Final = respx_mock.post(f"{_DATABRICKS_FIXTURE_API_BASE}/chat/completions").respond(
|
||||
401, json={"error_code": "UNAUTHENTICATED", "message": "fixture bad token"}
|
||||
)
|
||||
|
||||
response: Final = await ahealth_check(
|
||||
{
|
||||
"model": "databricks/databricks-claude-*",
|
||||
"api_key": "fake-key",
|
||||
"api_base": _DATABRICKS_FIXTURE_API_BASE,
|
||||
"messages": [{"role": "user", "content": "What is 1 + 1?"}],
|
||||
}
|
||||
)
|
||||
|
||||
assert [json.loads(call.request.content)["model"] for call in upstream.calls] == [
|
||||
"databricks-claude-fixture-a",
|
||||
"databricks-claude-fixture-b",
|
||||
"databricks-claude-fixture-retired",
|
||||
"databricks-cheap-chat",
|
||||
]
|
||||
assert "error" in response
|
||||
|
||||
|
||||
_DATABRICKS_EQUAL_PRICE_MODELS: Final = (
|
||||
"databricks/databricks-free-d",
|
||||
"databricks/databricks-free-c",
|
||||
"databricks/databricks-free-b",
|
||||
"databricks/databricks-free-a",
|
||||
)
|
||||
|
||||
|
||||
def _register_equal_price_databricks_models(monkeypatch: pytest.MonkeyPatch, listed_order: tuple[str, ...]) -> None:
|
||||
litellm.register_model(
|
||||
{
|
||||
model: {
|
||||
"litellm_provider": "databricks",
|
||||
"mode": "chat",
|
||||
"max_tokens": 1024,
|
||||
"input_cost_per_token": 0.0,
|
||||
"output_cost_per_token": 0.0,
|
||||
}
|
||||
for model in listed_order
|
||||
}
|
||||
)
|
||||
monkeypatch.setitem(litellm.models_by_provider, "databricks", list(listed_order))
|
||||
|
||||
|
||||
@pytest.mark.parametrize("listed_order", [_DATABRICKS_EQUAL_PRICE_MODELS, _DATABRICKS_EQUAL_PRICE_MODELS[::-1]])
|
||||
def test_partial_wildcard_health_check_orders_its_own_equal_price_models_by_id_whatever_the_listing_order(
|
||||
monkeypatch: pytest.MonkeyPatch, listed_order: tuple[str, ...]
|
||||
) -> None:
|
||||
_register_equal_price_databricks_models(monkeypatch, listed_order)
|
||||
|
||||
assert wildcard_health_check_models(wildcard="databricks/databricks-free-*", custom_llm_provider="databricks")[
|
||||
:2
|
||||
] == ("databricks/databricks-free-a", "databricks/databricks-free-b")
|
||||
|
||||
|
||||
def _answer_only_the_first_listed_databricks_free_model(request: httpx.Request) -> httpx.Response:
|
||||
model: Final = json.loads(request.content)["model"]
|
||||
if model != "databricks-free-d":
|
||||
return httpx.Response(404, json={"error_code": "ENDPOINT_NOT_FOUND", "message": f"{model} does not exist"})
|
||||
return _fail_first_databricks_claude_probe(request)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_wildcard_health_check_falls_back_to_equal_price_provider_models_in_the_provider_listing_order(
|
||||
monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter
|
||||
) -> None:
|
||||
_register_equal_price_databricks_models(monkeypatch, _DATABRICKS_EQUAL_PRICE_MODELS)
|
||||
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
|
||||
litellm.in_memory_llm_clients_cache.flush_cache()
|
||||
upstream: Final = respx_mock.post(f"{_DATABRICKS_FIXTURE_API_BASE}/chat/completions").mock(
|
||||
side_effect=_answer_only_the_first_listed_databricks_free_model
|
||||
)
|
||||
|
||||
response: Final = await ahealth_check(
|
||||
{
|
||||
"model": "databricks/databricks-claude-*",
|
||||
"api_key": "fake-key",
|
||||
"api_base": _DATABRICKS_FIXTURE_API_BASE,
|
||||
"messages": [{"role": "user", "content": "What is 1 + 1?"}],
|
||||
}
|
||||
)
|
||||
|
||||
assert [json.loads(call.request.content)["model"] for call in upstream.calls] == ["databricks-free-d"]
|
||||
assert response == {}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.usefixtures("wildcard_catalog")
|
||||
async def test_partial_wildcard_health_check_without_a_served_cost_map_model_probes_the_provider_cheapest_chat_model(
|
||||
monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter
|
||||
) -> None:
|
||||
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
|
||||
litellm.in_memory_llm_clients_cache.flush_cache()
|
||||
upstream: Final = respx_mock.post(f"{_DATABRICKS_FIXTURE_API_BASE}/chat/completions").mock(
|
||||
side_effect=_fail_first_databricks_claude_probe
|
||||
)
|
||||
|
||||
response: Final = await ahealth_check(
|
||||
{
|
||||
"model": "databricks/databricks-llama-*",
|
||||
"api_key": "fake-key",
|
||||
"api_base": _DATABRICKS_FIXTURE_API_BASE,
|
||||
"messages": [{"role": "user", "content": "What is 1 + 1?"}],
|
||||
}
|
||||
)
|
||||
|
||||
assert json.loads(upstream.calls[0].request.content)["model"] == "databricks-claude-fixture-retired"
|
||||
assert response == {}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_partial_wildcard_health_check_on_a_provider_without_cost_map_chat_models_names_health_check_model(
|
||||
monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter
|
||||
) -> None:
|
||||
monkeypatch.setitem(litellm.models_by_provider, "databricks", [])
|
||||
|
||||
response: Final = await ahealth_check(
|
||||
{
|
||||
"model": "databricks/databricks-llama-*",
|
||||
"api_key": "fake-key",
|
||||
"api_base": _DATABRICKS_FIXTURE_API_BASE,
|
||||
"messages": [{"role": "user", "content": "What is 1 + 1?"}],
|
||||
}
|
||||
)
|
||||
|
||||
assert "model_info.health_check_model" in response["error"], response["error"]
|
||||
assert respx_mock.calls.call_count == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cohere_rerank_health_check_returns_a_successful_response(
|
||||
monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter
|
||||
|
|
@ -1632,7 +1967,12 @@ async def test_wildcard_health_check_with_pinned_provider_calls_a_model_without_
|
|||
monkeypatch.setitem(
|
||||
litellm.model_cost,
|
||||
"fake-together-chat",
|
||||
{"input_cost_per_token": 1e-9, "output_cost_per_token": 1e-9, "mode": "chat", "litellm_provider": "together_ai"},
|
||||
{
|
||||
"input_cost_per_token": 1e-9,
|
||||
"output_cost_per_token": 1e-9,
|
||||
"mode": "chat",
|
||||
"litellm_provider": "together_ai",
|
||||
},
|
||||
)
|
||||
monkeypatch.setitem(litellm.models_by_provider, "together_ai", ("fake-together-chat",))
|
||||
litellm.in_memory_llm_clients_cache.flush_cache()
|
||||
|
|
@ -1647,4 +1987,3 @@ async def test_wildcard_health_check_with_pinned_provider_calls_a_model_without_
|
|||
assert "error" not in result, result
|
||||
assert json.loads(together_route.calls.last.request.content)["model"] == "fake-together-chat"
|
||||
assert not other_hosts.called
|
||||
|
||||
|
|
|
|||
|
|
@ -8,6 +8,7 @@ from litellm.litellm_core_utils.llm_request_utils import (
|
|||
ensure_extra_body_is_safe,
|
||||
flatten_form_field_values,
|
||||
pick_cheapest_chat_models_from_llm_provider,
|
||||
rank_chat_models_from_llm_provider,
|
||||
serialize_multipart_form_fields,
|
||||
)
|
||||
|
||||
|
|
@ -149,6 +150,24 @@ def deepseek_rows_whose_cheapest_ids_route_elsewhere(monkeypatch: pytest.MonkeyP
|
|||
"deepseek/fake-cheapest-embedding": {"input_cost_per_token": 1e-13, "mode": "embedding"},
|
||||
"openai/fake-free-chat": {"input_cost_per_token": 0.0, "output_cost_per_token": 0.0, "mode": "chat"},
|
||||
"deepseek/fake-cheap-chat": {"input_cost_per_token": 1e-12, "output_cost_per_token": 1e-12, "mode": "chat"},
|
||||
"deepseek/fake-retired-chat": {
|
||||
"input_cost_per_token": 1e-14,
|
||||
"output_cost_per_token": 1e-14,
|
||||
"mode": "chat",
|
||||
"deprecation_date": "2000-01-01",
|
||||
},
|
||||
"deepseek/fake-retiring-chat": {
|
||||
"input_cost_per_token": 2e-12,
|
||||
"output_cost_per_token": 2e-12,
|
||||
"mode": "chat",
|
||||
"deprecation_date": "2999-01-01",
|
||||
},
|
||||
"deepseek/fake-misdated-chat": {
|
||||
"input_cost_per_token": 2.5e-12,
|
||||
"output_cost_per_token": 2.5e-12,
|
||||
"mode": "chat",
|
||||
"deprecation_date": "next spring",
|
||||
},
|
||||
"fake-free-chat-without-provider": {"input_cost_per_token": 0.0, "output_cost_per_token": 0.0, "mode": "chat"},
|
||||
}
|
||||
for name, row in rows.items():
|
||||
|
|
@ -156,12 +175,73 @@ def deepseek_rows_whose_cheapest_ids_route_elsewhere(monkeypatch: pytest.MonkeyP
|
|||
monkeypatch.setitem(litellm.models_by_provider, PICKER_PROVIDER, tuple(rows))
|
||||
|
||||
|
||||
def test_pick_cheapest_chat_models_prefers_provider_routed_ids_and_skips_ids_routed_elsewhere(
|
||||
def test_pick_cheapest_chat_models_prefers_provider_routed_ids_skips_ids_routed_elsewhere_and_keeps_retired_ones(
|
||||
deepseek_rows_whose_cheapest_ids_route_elsewhere: None,
|
||||
) -> None:
|
||||
assert pick_cheapest_chat_models_from_llm_provider(PICKER_PROVIDER, n=10) == [
|
||||
"deepseek/fake-retired-chat",
|
||||
"deepseek/fake-cheap-chat",
|
||||
"deepseek/fake-retiring-chat",
|
||||
"deepseek/fake-misdated-chat",
|
||||
"deepseek/fake-pricey-chat",
|
||||
"fake-free-chat-without-provider",
|
||||
]
|
||||
assert pick_cheapest_chat_models_from_llm_provider(PICKER_PROVIDER, n=1) == ["deepseek/fake-cheap-chat"]
|
||||
assert pick_cheapest_chat_models_from_llm_provider(PICKER_PROVIDER, n=1) == ["deepseek/fake-retired-chat"]
|
||||
|
||||
|
||||
def test_rank_chat_models_flags_only_the_rows_whose_cost_map_deprecation_date_has_passed(
|
||||
deepseek_rows_whose_cheapest_ids_route_elsewhere: None,
|
||||
) -> None:
|
||||
assert {
|
||||
ranked_model.model: ranked_model.is_past_deprecation_date
|
||||
for ranked_model in rank_chat_models_from_llm_provider(PICKER_PROVIDER)
|
||||
} == {
|
||||
"deepseek/fake-retired-chat": True,
|
||||
"deepseek/fake-cheap-chat": False,
|
||||
"deepseek/fake-retiring-chat": False,
|
||||
"deepseek/fake-misdated-chat": False,
|
||||
"deepseek/fake-pricey-chat": False,
|
||||
"fake-free-chat-without-provider": False,
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"listed_order",
|
||||
[("deepseek/fake-tie-b", "deepseek/fake-tie-a"), ("deepseek/fake-tie-a", "deepseek/fake-tie-b")],
|
||||
)
|
||||
def test_pick_cheapest_chat_models_keeps_the_provider_listing_order_for_equal_price_models(
|
||||
monkeypatch: pytest.MonkeyPatch, listed_order: tuple[str, str]
|
||||
) -> None:
|
||||
for name in listed_order:
|
||||
monkeypatch.setitem(
|
||||
litellm.model_cost,
|
||||
name,
|
||||
{
|
||||
"input_cost_per_token": 1e-12,
|
||||
"output_cost_per_token": 1e-12,
|
||||
"mode": "chat",
|
||||
"litellm_provider": PICKER_PROVIDER,
|
||||
},
|
||||
)
|
||||
monkeypatch.setitem(litellm.models_by_provider, PICKER_PROVIDER, listed_order)
|
||||
|
||||
assert pick_cheapest_chat_models_from_llm_provider(PICKER_PROVIDER, n=2) == list(listed_order)
|
||||
|
||||
|
||||
def test_rank_chat_models_flags_a_retired_model_the_provider_lists_without_its_cost_map_prefix(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
rows: Final = {
|
||||
"vertex_ai/fake-retired-vertex-chat": {"input_cost_per_token": 1e-14, "deprecation_date": "2000-01-01"},
|
||||
"vertex_ai/fake-live-vertex-chat": {"input_cost_per_token": 1e-12, "deprecation_date": "2999-01-01"},
|
||||
}
|
||||
for name, row in rows.items():
|
||||
monkeypatch.setitem(litellm.model_cost, name, {**row, "mode": "chat", "litellm_provider": "vertex_ai"})
|
||||
monkeypatch.setitem(
|
||||
litellm.models_by_provider, "vertex_ai", tuple(name.removeprefix("vertex_ai/") for name in rows)
|
||||
)
|
||||
|
||||
assert [
|
||||
(ranked_model.model, ranked_model.is_past_deprecation_date)
|
||||
for ranked_model in rank_chat_models_from_llm_provider("vertex_ai")
|
||||
] == [("fake-retired-vertex-chat", True), ("fake-live-vertex-chat", False)]
|
||||
|
|
|
|||
|
|
@ -1,7 +1,7 @@
|
|||
import asyncio
|
||||
import base64
|
||||
import json
|
||||
from collections.abc import Coroutine, Mapping, Sequence
|
||||
from collections.abc import Callable, Coroutine, Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from typing import Final, Literal
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
|
@ -14,6 +14,7 @@ from websockets.frames import Close
|
|||
import litellm
|
||||
from litellm.constants import (
|
||||
BEDROCK_REALTIME_PENDING_OPENING_MESSAGE_SCOPE_KEY,
|
||||
REALTIME_LAZY_SETUP_DEADLINE_SECONDS,
|
||||
REALTIME_METERED_USAGE_KEY,
|
||||
REALTIME_SESSION_FAILURE_LOGGED_KEY,
|
||||
REALTIME_SESSION_SUCCESS_LOGGED_KEY,
|
||||
|
|
@ -21,6 +22,7 @@ from litellm.constants import (
|
|||
from litellm.integrations.custom_guardrail import CustomGuardrail
|
||||
from litellm.litellm_core_utils.realtime_streaming import (
|
||||
RealTimeStreaming,
|
||||
_merge_session_fields,
|
||||
client_sent_openai_beta_realtime_header,
|
||||
)
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
|
|
@ -1255,11 +1257,15 @@ async def test_client_ack_caches_setup_to_prevent_duplicate_session_update_setup
|
|||
side_effect=[
|
||||
json.dumps({"type": "session.update", "session": {"tools": []}}),
|
||||
json.dumps({"type": "session.update", "session": {"tools": []}}),
|
||||
_TEXT_ITEM_CREATE,
|
||||
_RESPONSE_CREATE,
|
||||
Exception("client done"),
|
||||
]
|
||||
)
|
||||
websocket.send_text = AsyncMock()
|
||||
|
||||
provider_config = MagicMock()
|
||||
provider_config.setup_session_update_rejection = MagicMock(return_value=None)
|
||||
|
||||
def _transform(message: str, model: str, session_configuration_request=None):
|
||||
if session_configuration_request is None:
|
||||
|
|
@ -2285,6 +2291,7 @@ async def test_guardrail_turn_detection_injected_into_first_session_update_defer
|
|||
logging_obj.success_handler = MagicMock()
|
||||
|
||||
provider_config = MagicMock()
|
||||
provider_config.builds_setup_from_session_update = MagicMock(return_value=False)
|
||||
transformed_messages = []
|
||||
|
||||
def mock_transform(msg, model, session_config):
|
||||
|
|
@ -2351,6 +2358,7 @@ async def test_guardrail_turn_detection_injection_tolerates_non_dict_value(
|
|||
logging_obj.success_handler = MagicMock()
|
||||
|
||||
provider_config = MagicMock()
|
||||
provider_config.builds_setup_from_session_update = MagicMock(return_value=False)
|
||||
transformed_messages = []
|
||||
|
||||
def mock_transform(msg, model, session_config):
|
||||
|
|
@ -2413,6 +2421,7 @@ async def test_subsequent_session_update_cannot_reenable_vad_when_guardrails_act
|
|||
logging_obj.success_handler = MagicMock()
|
||||
|
||||
provider_config = MagicMock()
|
||||
provider_config.setup_session_update_rejection = MagicMock(return_value=None)
|
||||
transformed_messages = []
|
||||
|
||||
def mock_transform(msg, model, session_config):
|
||||
|
|
@ -2465,6 +2474,7 @@ async def test_follow_up_setup_updates_cached_session_configuration_request():
|
|||
logging_obj.success_handler = MagicMock()
|
||||
|
||||
provider_config = MagicMock()
|
||||
provider_config.setup_session_update_rejection = MagicMock(return_value=None)
|
||||
follow_up_setup = json.dumps(
|
||||
{
|
||||
"setup": {
|
||||
|
|
@ -2498,6 +2508,636 @@ async def test_follow_up_setup_updates_cached_session_configuration_request():
|
|||
assert streaming.session_configuration_request == follow_up_setup
|
||||
|
||||
|
||||
def _gemini_default_mode_streaming(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
client_messages: Sequence[str],
|
||||
model: str = "gemini-live-2.5-flash-native-audio",
|
||||
drop_params: bool = False,
|
||||
defer_setup: bool = False,
|
||||
) -> tuple[RealTimeStreaming, MagicMock]:
|
||||
monkeypatch.setattr(litellm, "gemini_live_defer_setup", defer_setup, raising=False)
|
||||
from litellm.llms.gemini.realtime.transformation import GeminiRealtimeConfig
|
||||
|
||||
client_ws = MagicMock()
|
||||
client_ws.receive_text = AsyncMock(side_effect=[*client_messages, ConnectionClosed(None, None)])
|
||||
client_ws.send_text = AsyncMock()
|
||||
backend_ws = MagicMock()
|
||||
backend_ws.send = AsyncMock()
|
||||
streaming = RealTimeStreaming(
|
||||
client_ws,
|
||||
backend_ws,
|
||||
MagicMock(),
|
||||
provider_config=GeminiRealtimeConfig(drop_params=drop_params),
|
||||
model=model,
|
||||
)
|
||||
return streaming, backend_ws
|
||||
|
||||
|
||||
def _events_sent_to_client(streaming: RealTimeStreaming) -> list[dict]:
|
||||
return [json.loads(call.args[0]) for call in streaming.websocket.send_text.await_args_list]
|
||||
|
||||
|
||||
_TEXT_ITEM_CREATE = json.dumps(
|
||||
{
|
||||
"type": "conversation.item.create",
|
||||
"item": {"type": "message", "role": "user", "content": [{"type": "input_text", "text": "hi"}]},
|
||||
}
|
||||
)
|
||||
|
||||
_RESPONSE_CREATE = json.dumps({"type": "response.create"})
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_default_mode_sends_default_setup_before_first_non_session_client_message(monkeypatch):
|
||||
"""Gemini Live setup is no longer pre-sent at connect, so a client that never
|
||||
sends session.update still gets the provider default setup ahead of its
|
||||
first content frame."""
|
||||
streaming, backend_ws = _gemini_default_mode_streaming(monkeypatch, [_TEXT_ITEM_CREATE])
|
||||
|
||||
await streaming.client_ack_messages()
|
||||
|
||||
frames = [json.loads(call.args[0]) for call in backend_ws.send.await_args_list]
|
||||
assert [next(iter(frame)) for frame in frames] == ["setup", "clientContent"]
|
||||
assert frames[0]["setup"]["generationConfig"]["responseModalities"] == ["AUDIO"]
|
||||
assert streaming.session_configuration_request == backend_ws.send.await_args_list[0].args[0]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_default_mode_session_update_first_builds_the_only_setup(monkeypatch):
|
||||
session_update = json.dumps({"type": "session.update", "session": {"modalities": ["text"]}})
|
||||
streaming, backend_ws = _gemini_default_mode_streaming(
|
||||
monkeypatch, [session_update, _TEXT_ITEM_CREATE, _RESPONSE_CREATE], model="gemini-live-2.5-flash"
|
||||
)
|
||||
|
||||
await streaming.client_ack_messages()
|
||||
|
||||
frames = [json.loads(call.args[0]) for call in backend_ws.send.await_args_list]
|
||||
assert [next(iter(frame)) for frame in frames] == ["setup", "clientContent"]
|
||||
assert frames[0]["setup"]["generationConfig"]["responseModalities"] == ["TEXT"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_default_mode_merges_every_session_update_sent_before_the_first_content_frame(monkeypatch):
|
||||
"""Clients configure a session over several session.update messages. Gemini
|
||||
Live takes one setup, so every update sent before the first content frame is
|
||||
merged into that setup instead of only the first one applying."""
|
||||
instructions_update = json.dumps({"type": "session.update", "session": {"instructions": "Be brief."}})
|
||||
modalities_update = json.dumps({"type": "session.update", "session": {"modalities": ["text"]}})
|
||||
streaming, backend_ws = _gemini_default_mode_streaming(
|
||||
monkeypatch, [instructions_update, modalities_update, _TEXT_ITEM_CREATE, _RESPONSE_CREATE], model="gemini-live-2.5-flash"
|
||||
)
|
||||
|
||||
await streaming.client_ack_messages()
|
||||
|
||||
frames = [json.loads(call.args[0]) for call in backend_ws.send.await_args_list]
|
||||
assert [next(iter(frame)) for frame in frames] == ["setup", "clientContent"]
|
||||
assert frames[0]["setup"]["generationConfig"]["responseModalities"] == ["TEXT"]
|
||||
assert frames[0]["setup"]["systemInstruction"]["parts"][0]["text"] == "Be brief."
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_default_mode_merges_nested_audio_settings_across_session_updates(monkeypatch):
|
||||
vad_update = json.dumps(
|
||||
{"type": "session.update", "session": {"audio": {"input": {"turn_detection": {"silence_duration_ms": 1500}}}}}
|
||||
)
|
||||
voice_update = json.dumps({"type": "session.update", "session": {"audio": {"output": {"voice": "Kore"}}}})
|
||||
streaming, backend_ws = _gemini_default_mode_streaming(
|
||||
monkeypatch, [vad_update, voice_update, _TEXT_ITEM_CREATE, _RESPONSE_CREATE], model="gemini-live-2.5-flash"
|
||||
)
|
||||
|
||||
await streaming.client_ack_messages()
|
||||
|
||||
setup = json.loads(backend_ws.send.await_args_list[0].args[0])["setup"]
|
||||
assert setup["realtimeInputConfig"]["automaticActivityDetection"]["silenceDurationMs"] == 1500
|
||||
assert setup["generationConfig"]["speechConfig"]["voiceConfig"]["prebuiltVoiceConfig"]["voiceName"] == "Kore"
|
||||
|
||||
|
||||
def test_merge_session_fields_replaces_nested_mappings_wholesale_past_the_depth_cap():
|
||||
def nested(levels: int, key: str) -> dict:
|
||||
return {key: levels, **({"child": nested(levels - 1, key)} if levels else {})}
|
||||
|
||||
def level(node, depth: int):
|
||||
return node if depth == 0 else level(node["child"], depth - 1)
|
||||
|
||||
merged = _merge_session_fields(nested(10, "keep"), nested(10, "new"))
|
||||
|
||||
assert all({"keep", "new"} <= set(level(merged, depth)) for depth in range(9))
|
||||
assert "keep" not in level(merged, 9)
|
||||
assert level(merged, 9)["new"] == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_default_mode_acknowledges_each_held_session_update_before_any_backend_frame(monkeypatch):
|
||||
"""A client that waits for session.updated before sending content must not
|
||||
hang while its updates are held, so each one is acknowledged right away
|
||||
with the modalities the merged setup will carry."""
|
||||
instructions_update = json.dumps({"type": "session.update", "session": {"instructions": "Be brief."}})
|
||||
modalities_update = json.dumps({"type": "session.update", "session": {"modalities": ["text"]}})
|
||||
streaming, backend_ws = _gemini_default_mode_streaming(
|
||||
monkeypatch, [instructions_update, modalities_update], model="gemini-live-2.5-flash"
|
||||
)
|
||||
|
||||
await streaming.client_ack_messages()
|
||||
|
||||
acks = [json.loads(call.args[0]) for call in streaming.websocket.send_text.await_args_list]
|
||||
assert [ack["type"] for ack in acks] == ["session.updated", "session.updated"]
|
||||
assert [ack["session"].get("output_modalities") or ack["session"].get("modalities") for ack in acks] == [
|
||||
["audio"],
|
||||
["text"],
|
||||
]
|
||||
assert acks[1]["session"]["instructions"] == "Be brief."
|
||||
backend_ws.send.assert_not_awaited()
|
||||
|
||||
|
||||
_LATE_INSTRUCTIONS_UPDATE = json.dumps(
|
||||
{"type": "session.update", "event_id": "evt_late", "session": {"instructions": "Be brief."}}
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_default_mode_answers_a_session_update_after_the_first_content_with_an_error(monkeypatch):
|
||||
streaming, backend_ws = _gemini_default_mode_streaming(monkeypatch, [_TEXT_ITEM_CREATE, _LATE_INSTRUCTIONS_UPDATE])
|
||||
|
||||
await streaming.client_ack_messages()
|
||||
|
||||
errors: Final = [event["error"] for event in _events_sent_to_client(streaming) if event["type"] == "error"]
|
||||
assert [(error["code"], error["event_id"]) for error in errors] == [("session_update_not_applied", "evt_late")]
|
||||
frames: Final = [json.loads(call.args[0]) for call in backend_ws.send.await_args_list]
|
||||
assert [next(iter(frame)) for frame in frames] == ["setup", "clientContent"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_default_mode_drops_a_late_session_update_quietly_under_drop_params(monkeypatch):
|
||||
streaming, backend_ws = _gemini_default_mode_streaming(
|
||||
monkeypatch, [_TEXT_ITEM_CREATE, _LATE_INSTRUCTIONS_UPDATE], drop_params=True
|
||||
)
|
||||
|
||||
await streaming.client_ack_messages()
|
||||
|
||||
assert _events_sent_to_client(streaming) == []
|
||||
assert backend_ws.send.await_count == 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("model", "requested_output"),
|
||||
[
|
||||
("gemini-live-2.5-flash-native-audio", {"modalities": ["text"]}),
|
||||
("gemini-live-2.5-flash", {"output_modalities": ["image"]}),
|
||||
],
|
||||
)
|
||||
async def test_default_mode_rejects_an_output_modality_the_model_cannot_produce(monkeypatch, model, requested_output):
|
||||
update: Final = json.dumps(
|
||||
{"type": "session.update", "session": {**requested_output, "instructions": "Be brief."}}
|
||||
)
|
||||
streaming, backend_ws = _gemini_default_mode_streaming(monkeypatch, [update, _TEXT_ITEM_CREATE], model=model)
|
||||
|
||||
await streaming.client_ack_messages()
|
||||
|
||||
events: Final = _events_sent_to_client(streaming)
|
||||
assert [(event["type"], event["error"]["code"]) for event in events] == [("error", "invalid_value")]
|
||||
setup: Final = json.loads(backend_ws.send.await_args_list[0].args[0])["setup"]
|
||||
assert "systemInstruction" not in setup
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("modalities", "drop_params"),
|
||||
[(["text", "audio"], False), (["text"], True)],
|
||||
ids=["openai-default-pair", "text-dropped-under-drop-params"],
|
||||
)
|
||||
async def test_default_mode_applies_the_update_with_audio_output_on_an_audio_only_model(
|
||||
monkeypatch, modalities, drop_params
|
||||
):
|
||||
update: Final = json.dumps(
|
||||
{"type": "session.update", "session": {"modalities": modalities, "instructions": "Be brief."}}
|
||||
)
|
||||
streaming, backend_ws = _gemini_default_mode_streaming(
|
||||
monkeypatch, [update, _TEXT_ITEM_CREATE, _RESPONSE_CREATE], drop_params=drop_params
|
||||
)
|
||||
|
||||
await streaming.client_ack_messages()
|
||||
|
||||
assert [event["type"] for event in _events_sent_to_client(streaming)] == ["session.updated"]
|
||||
setup: Final = json.loads(backend_ws.send.await_args_list[0].args[0])["setup"]
|
||||
assert setup["generationConfig"]["responseModalities"] == ["AUDIO"]
|
||||
assert setup["systemInstruction"]["parts"][0]["text"] == "Be brief."
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_default_mode_sends_no_extra_session_updated_when_gemini_confirms_the_setup(monkeypatch):
|
||||
"""Each held session.update was already acknowledged, so Gemini's setupComplete
|
||||
must not reach the client as one more session.updated nobody asked for."""
|
||||
monkeypatch.setattr(litellm, "gemini_live_defer_setup", False, raising=False)
|
||||
from litellm.llms.gemini.realtime.transformation import GeminiRealtimeConfig
|
||||
|
||||
model: Final = "gemini-live-2.5-flash-native-audio"
|
||||
config: Final = GeminiRealtimeConfig()
|
||||
client_ws: Final = MagicMock()
|
||||
client_ws.send_text = AsyncMock()
|
||||
backend_ws: Final = MagicMock()
|
||||
backend_ws.send = AsyncMock()
|
||||
backend_ws.recv = AsyncMock(side_effect=[json.dumps({"setupComplete": {}}).encode(), ConnectionClosed(None, None)])
|
||||
logging_obj: Final = _mock_logging()
|
||||
logging_obj.litellm_trace_id = "trace_lazy"
|
||||
logging_obj.async_success_handler = AsyncMock()
|
||||
logging_obj.success_handler = MagicMock()
|
||||
streaming: Final = RealTimeStreaming(client_ws, backend_ws, logging_obj, provider_config=config, model=model)
|
||||
streaming.session_created_sent_to_client = True
|
||||
streaming.session_configuration_request = config.session_configuration_request(model)
|
||||
|
||||
await streaming.backend_to_client_send_messages()
|
||||
|
||||
assert [json.loads(call.args[0])["type"] for call in client_ws.send_text.await_args_list] == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_default_mode_folds_transcription_guardrail_vad_disable_into_the_merged_setup(monkeypatch):
|
||||
"""A transcription guardrail gates every turn, so its auto-response disable
|
||||
has to land in the one setup built from the held session.updates, injected
|
||||
once rather than into every update that gets merged."""
|
||||
monkeypatch.setattr(litellm, "callbacks", [_transcription_guardrail()])
|
||||
debug = MagicMock()
|
||||
monkeypatch.setattr("litellm.litellm_core_utils.realtime_streaming.verbose_logger.debug", debug)
|
||||
instructions_update = json.dumps({"type": "session.update", "session": {"instructions": "Be brief."}})
|
||||
vad_update = json.dumps(
|
||||
{"type": "session.update", "session": {"turn_detection": {"type": "server_vad", "silence_duration_ms": 900}}}
|
||||
)
|
||||
streaming, backend_ws = _gemini_default_mode_streaming(
|
||||
monkeypatch, [instructions_update, vad_update, _TEXT_ITEM_CREATE, _RESPONSE_CREATE], model="gemini-live-2.5-flash"
|
||||
)
|
||||
|
||||
await streaming.client_ack_messages()
|
||||
|
||||
frames = [json.loads(call.args[0]) for call in backend_ws.send.await_args_list]
|
||||
assert [next(iter(frame)) for frame in frames] == ["setup", "clientContent"]
|
||||
assert frames[0]["setup"]["realtimeInputConfig"]["automaticActivityDetection"] == {"disabled": True}
|
||||
assert frames[0]["setup"]["systemInstruction"]["parts"][0]["text"] == "Be brief."
|
||||
assert streaming._guardrail_turn_detection_update_sent is True
|
||||
injected = [call for call in debug.call_args_list if str(call.args[0]).startswith("Injected turn_detection")]
|
||||
assert len(injected) == 1
|
||||
|
||||
|
||||
def _blocking_text_guardrail() -> CustomGuardrail:
|
||||
class _BlockingTextGuardrail(CustomGuardrail):
|
||||
async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None):
|
||||
if any("zebrafruit" in text for text in inputs.get("texts", [])):
|
||||
raise ValueError("Content blocked")
|
||||
return inputs
|
||||
|
||||
return _BlockingTextGuardrail(
|
||||
guardrail_name="test_text_guard", event_hook=GuardrailEventHooks.pre_call, default_on=True
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_default_mode_sends_the_setup_before_a_guardrail_blocks_the_first_message(monkeypatch):
|
||||
"""Gemini Live closes the socket when content arrives before the setup, so a
|
||||
guardrail that blocks the first typed message must not reach the backend first."""
|
||||
monkeypatch.setattr(litellm, "callbacks", [_blocking_text_guardrail()])
|
||||
instructions_update = json.dumps({"type": "session.update", "session": {"instructions": "Be brief."}})
|
||||
blocked_item = json.dumps(
|
||||
{
|
||||
"type": "conversation.item.create",
|
||||
"item": {"type": "message", "role": "user", "content": [{"type": "input_text", "text": "zebrafruit"}]},
|
||||
}
|
||||
)
|
||||
streaming, backend_ws = _gemini_default_mode_streaming(
|
||||
monkeypatch, [instructions_update, blocked_item], model="gemini-live-2.5-flash"
|
||||
)
|
||||
|
||||
await streaming.client_ack_messages()
|
||||
|
||||
frames: Final = [json.loads(call.args[0]) for call in backend_ws.send.await_args_list]
|
||||
assert [next(iter(frame)) for frame in frames] == ["setup", "clientContent"]
|
||||
assert frames[0]["setup"]["systemInstruction"]["parts"][0]["text"] == "Be brief."
|
||||
assert "zebrafruit" not in json.dumps(frames[1])
|
||||
errors: Final = [event["error"]["type"] for event in _events_sent_to_client(streaming) if event["type"] == "error"]
|
||||
assert errors == ["guardrail_violation"]
|
||||
|
||||
|
||||
class _SteppedClock:
|
||||
def __init__(self, epoch: float) -> None:
|
||||
self.now = epoch
|
||||
|
||||
def __call__(self) -> float:
|
||||
return self.now
|
||||
|
||||
|
||||
def _client_ws_receiving(*steps: Callable[[], Coroutine[object, object, str]]) -> MagicMock:
|
||||
remaining: Final = iter(steps)
|
||||
|
||||
async def receive_text() -> str:
|
||||
return await next(remaining)()
|
||||
|
||||
client_ws: Final = MagicMock()
|
||||
client_ws.receive_text = receive_text
|
||||
client_ws.send_text = AsyncMock()
|
||||
return client_ws
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("epoch", [0.0, 1e9])
|
||||
async def test_default_mode_sends_the_held_setup_when_no_content_arrives_before_the_deadline(monkeypatch, epoch):
|
||||
"""Vertex AI refuses a setup sent minutes after connecting, so an idle session
|
||||
sends what it holds at the deadline and answers any later session.update as late."""
|
||||
monkeypatch.setattr(litellm, "gemini_live_defer_setup", False, raising=False)
|
||||
from litellm.llms.gemini.realtime.transformation import GeminiRealtimeConfig
|
||||
|
||||
clock: Final = _SteppedClock(epoch)
|
||||
setup_sent: Final = asyncio.Event()
|
||||
|
||||
async def instructions_update_then_the_deadline_passes() -> str:
|
||||
clock.now += REALTIME_LAZY_SETUP_DEADLINE_SECONDS
|
||||
return json.dumps({"type": "session.update", "session": {"instructions": "Be brief."}})
|
||||
|
||||
async def late_update_once_the_setup_is_out() -> str:
|
||||
await setup_sent.wait()
|
||||
return json.dumps(
|
||||
{"type": "session.update", "event_id": "evt_late", "session": {"instructions": "Answer in French."}}
|
||||
)
|
||||
|
||||
async def disconnect() -> str:
|
||||
raise ConnectionClosed(None, None)
|
||||
|
||||
backend_ws: Final = MagicMock()
|
||||
backend_ws.send = AsyncMock(side_effect=lambda frame: setup_sent.set())
|
||||
streaming: Final = RealTimeStreaming(
|
||||
_client_ws_receiving(instructions_update_then_the_deadline_passes, late_update_once_the_setup_is_out, disconnect),
|
||||
backend_ws,
|
||||
MagicMock(),
|
||||
provider_config=GeminiRealtimeConfig(),
|
||||
model="gemini-live-2.5-flash",
|
||||
clock=clock,
|
||||
)
|
||||
|
||||
await asyncio.wait_for(streaming.client_ack_messages(), timeout=5)
|
||||
|
||||
frames: Final = [json.loads(call.args[0]) for call in backend_ws.send.await_args_list]
|
||||
assert [next(iter(frame)) for frame in frames] == ["setup"]
|
||||
assert frames[0]["setup"]["systemInstruction"]["parts"][0]["text"] == "Be brief."
|
||||
events: Final = _events_sent_to_client(streaming)
|
||||
assert [(event["type"], event.get("error", {}).get("code")) for event in events] == [
|
||||
("session.updated", None),
|
||||
("error", "session_update_not_applied"),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_default_mode_refuses_held_session_updates_past_the_size_cap(monkeypatch):
|
||||
def instructions_update(text: str) -> str:
|
||||
return json.dumps({"type": "session.update", "event_id": f"evt_{text[0]}", "session": {"instructions": text}})
|
||||
|
||||
half_the_cap: Final = RealTimeStreaming._MAX_BUFFERED_BYTES // 2
|
||||
streaming, backend_ws = _gemini_default_mode_streaming(
|
||||
monkeypatch,
|
||||
[instructions_update("a" * half_the_cap), instructions_update("b" * half_the_cap), _TEXT_ITEM_CREATE, _RESPONSE_CREATE],
|
||||
model="gemini-live-2.5-flash",
|
||||
)
|
||||
|
||||
await streaming.client_ack_messages()
|
||||
|
||||
events: Final = _events_sent_to_client(streaming)
|
||||
assert [(event["type"], event.get("error", {}).get("event_id")) for event in events] == [
|
||||
("session.updated", None),
|
||||
("error", "evt_b"),
|
||||
]
|
||||
assert events[1]["error"]["code"] == "session_update_too_large"
|
||||
setup: Final = json.loads(backend_ws.send.await_args_list[0].args[0])["setup"]
|
||||
assert setup["systemInstruction"]["parts"][0]["text"].startswith("a")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("defer_setup", [False, True])
|
||||
async def test_a_session_update_between_the_first_items_and_response_create_joins_the_setup(monkeypatch, defer_setup):
|
||||
modalities_update: Final = json.dumps({"type": "session.update", "session": {"modalities": ["text"]}})
|
||||
instructions_update: Final = json.dumps({"type": "session.update", "session": {"instructions": "Be brief."}})
|
||||
streaming, backend_ws = _gemini_default_mode_streaming(
|
||||
monkeypatch,
|
||||
[modalities_update, _TEXT_ITEM_CREATE, instructions_update, _RESPONSE_CREATE],
|
||||
model="gemini-live-2.5-flash",
|
||||
defer_setup=defer_setup,
|
||||
)
|
||||
|
||||
await streaming.client_ack_messages()
|
||||
|
||||
frames: Final = [json.loads(call.args[0]) for call in backend_ws.send.await_args_list]
|
||||
assert [next(iter(frame)) for frame in frames] == ["setup", "clientContent"]
|
||||
assert frames[0]["setup"]["generationConfig"]["responseModalities"] == ["TEXT"]
|
||||
assert frames[0]["setup"]["systemInstruction"]["parts"][0]["text"] == "Be brief."
|
||||
assert [event["type"] for event in _events_sent_to_client(streaming)] == ["session.updated", "session.updated"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_default_mode_applies_a_session_update_that_arrives_while_audio_streams_in(monkeypatch):
|
||||
modalities_update: Final = json.dumps({"type": "session.update", "session": {"modalities": ["text"]}})
|
||||
append: Final = json.dumps({"type": "input_audio_buffer.append", "audio": "AAAA"})
|
||||
instructions_update: Final = json.dumps({"type": "session.update", "session": {"instructions": "Be brief."}})
|
||||
commit: Final = json.dumps({"type": "input_audio_buffer.commit"})
|
||||
streaming, backend_ws = _gemini_default_mode_streaming(
|
||||
monkeypatch, [modalities_update, append, instructions_update, commit], model="gemini-live-2.5-flash"
|
||||
)
|
||||
|
||||
await streaming.client_ack_messages()
|
||||
|
||||
frames: Final = [json.loads(call.args[0]) for call in backend_ws.send.await_args_list]
|
||||
assert [next(iter(frame)) for frame in frames][:2] == ["setup", "realtimeInput"]
|
||||
assert frames[0]["setup"]["systemInstruction"]["parts"][0]["text"] == "Be brief."
|
||||
assert frames[1]["realtimeInput"]["audio"]["data"] == "AAAA"
|
||||
assert "error" not in [event["type"] for event in _events_sent_to_client(streaming)]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("epoch", [0.0, 1e9])
|
||||
async def test_held_items_go_out_with_the_setup_when_no_response_create_follows(monkeypatch, epoch):
|
||||
monkeypatch.setattr(litellm, "gemini_live_defer_setup", False, raising=False)
|
||||
from litellm.llms.gemini.realtime.transformation import GeminiRealtimeConfig
|
||||
|
||||
item_sent: Final = asyncio.Event()
|
||||
|
||||
async def instructions_update() -> str:
|
||||
return json.dumps({"type": "session.update", "session": {"instructions": "Be brief."}})
|
||||
|
||||
async def text_item() -> str:
|
||||
return _TEXT_ITEM_CREATE
|
||||
|
||||
async def disconnect_once_the_item_is_out() -> str:
|
||||
await item_sent.wait()
|
||||
raise ConnectionClosed(None, None)
|
||||
|
||||
backend_ws: Final = MagicMock()
|
||||
backend_ws.send = AsyncMock(
|
||||
side_effect=lambda frame: item_sent.set() if "clientContent" in json.loads(frame) else None
|
||||
)
|
||||
streaming: Final = RealTimeStreaming(
|
||||
_client_ws_receiving(instructions_update, text_item, disconnect_once_the_item_is_out),
|
||||
backend_ws,
|
||||
MagicMock(),
|
||||
provider_config=GeminiRealtimeConfig(),
|
||||
model="gemini-live-2.5-flash",
|
||||
lazy_setup_content_grace_seconds=0,
|
||||
clock=_SteppedClock(epoch),
|
||||
)
|
||||
|
||||
await asyncio.wait_for(streaming.client_ack_messages(), timeout=5)
|
||||
|
||||
frames: Final = [json.loads(call.args[0]) for call in backend_ws.send.await_args_list]
|
||||
assert [next(iter(frame)) for frame in frames] == ["setup", "clientContent"]
|
||||
assert frames[0]["setup"]["systemInstruction"]["parts"][0]["text"] == "Be brief."
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_default_mode_acknowledges_an_unchanged_session_update_after_content(monkeypatch):
|
||||
"""Pipecat resends its unchanged session.update right after its first
|
||||
conversation items and treats any error event as fatal."""
|
||||
session_update: Final = json.dumps(
|
||||
{"type": "session.update", "session": {"modalities": ["text"], "instructions": "Be brief."}}
|
||||
)
|
||||
streaming, backend_ws = _gemini_default_mode_streaming(
|
||||
monkeypatch, [session_update, _TEXT_ITEM_CREATE, _RESPONSE_CREATE, session_update], model="gemini-live-2.5-flash"
|
||||
)
|
||||
|
||||
await streaming.client_ack_messages()
|
||||
|
||||
frames: Final = [json.loads(call.args[0]) for call in backend_ws.send.await_args_list]
|
||||
assert [next(iter(frame)) for frame in frames] == ["setup", "clientContent"]
|
||||
events: Final = _events_sent_to_client(streaming)
|
||||
assert [event["type"] for event in events] == ["session.updated", "session.updated"]
|
||||
assert events[0]["session"] == events[1]["session"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_deferred_setup_merges_session_updates_until_the_first_content(monkeypatch):
|
||||
"""Deferred setup merges every session.update sent before the first content,
|
||||
the same as the default mode, and answers a changed one after it as late."""
|
||||
monkeypatch.setattr(litellm, "gemini_live_defer_setup", True, raising=False)
|
||||
from litellm.llms.gemini.realtime.transformation import GeminiRealtimeConfig
|
||||
|
||||
client_ws: Final = MagicMock()
|
||||
client_ws.receive_text = AsyncMock(
|
||||
side_effect=[
|
||||
json.dumps({"type": "session.update", "session": {"modalities": ["text"]}}),
|
||||
json.dumps({"type": "session.update", "session": {"instructions": "Be brief."}}),
|
||||
_TEXT_ITEM_CREATE,
|
||||
_RESPONSE_CREATE,
|
||||
json.dumps(
|
||||
{"type": "session.update", "event_id": "evt_late", "session": {"instructions": "Answer in French."}}
|
||||
),
|
||||
ConnectionClosed(None, None),
|
||||
]
|
||||
)
|
||||
client_ws.send_text = AsyncMock()
|
||||
backend_ws: Final = MagicMock()
|
||||
backend_ws.send = AsyncMock()
|
||||
streaming: Final = RealTimeStreaming(
|
||||
client_ws, backend_ws, MagicMock(), provider_config=GeminiRealtimeConfig(), model="gemini-live-2.5-flash"
|
||||
)
|
||||
|
||||
await streaming.client_ack_messages()
|
||||
|
||||
frames: Final = [json.loads(call.args[0]) for call in backend_ws.send.await_args_list]
|
||||
assert [next(iter(frame)) for frame in frames] == ["setup", "clientContent"]
|
||||
assert frames[0]["setup"]["generationConfig"]["responseModalities"] == ["TEXT"]
|
||||
assert frames[0]["setup"]["systemInstruction"]["parts"][0]["text"] == "Be brief."
|
||||
errors: Final = [event["error"] for event in _events_sent_to_client(streaming) if event["type"] == "error"]
|
||||
assert [(error["code"], error["event_id"]) for error in errors] == [("session_update_not_applied", "evt_late")]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_deferred_setup_sends_the_default_setup_before_text_content_that_comes_first(monkeypatch):
|
||||
"""Gemini Live closes the socket on content sent before any setup, so text that arrives
|
||||
before a session.update in deferred mode gets the model's default setup first."""
|
||||
monkeypatch.setattr(litellm, "gemini_live_defer_setup", True, raising=False)
|
||||
from litellm.llms.gemini.realtime.transformation import GeminiRealtimeConfig
|
||||
|
||||
client_ws: Final = MagicMock()
|
||||
client_ws.receive_text = AsyncMock(
|
||||
side_effect=[
|
||||
_TEXT_ITEM_CREATE,
|
||||
json.dumps({"type": "session.update", "event_id": "evt_late", "session": {"modalities": ["text"]}}),
|
||||
ConnectionClosed(None, None),
|
||||
]
|
||||
)
|
||||
client_ws.send_text = AsyncMock()
|
||||
backend_ws: Final = MagicMock()
|
||||
backend_ws.send = AsyncMock()
|
||||
streaming: Final = RealTimeStreaming(
|
||||
client_ws, backend_ws, MagicMock(), provider_config=GeminiRealtimeConfig(), model="gemini-live-2.5-flash"
|
||||
)
|
||||
|
||||
await streaming.client_ack_messages()
|
||||
|
||||
frames: Final = [json.loads(call.args[0]) for call in backend_ws.send.await_args_list]
|
||||
assert [next(iter(frame)) for frame in frames] == ["setup", "clientContent"]
|
||||
errors: Final = [event["error"] for event in _events_sent_to_client(streaming) if event["type"] == "error"]
|
||||
assert [(error["code"], error["event_id"]) for error in errors] == [("session_update_not_applied", "evt_late")]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_deferred_setup_sends_a_session_update_that_follows_buffered_audio_at_once(monkeypatch):
|
||||
append: Final = json.dumps({"type": "input_audio_buffer.append", "audio": "AAAA"})
|
||||
commit: Final = json.dumps({"type": "input_audio_buffer.commit"})
|
||||
session_update: Final = json.dumps({"type": "session.update", "session": {"instructions": "Be a pirate."}})
|
||||
streaming, backend_ws = _gemini_default_mode_streaming(
|
||||
monkeypatch, [append, commit, session_update], model="gemini-live-2.5-flash", defer_setup=True
|
||||
)
|
||||
|
||||
await streaming.client_ack_messages()
|
||||
|
||||
frames: Final = [json.loads(call.args[0]) for call in backend_ws.send.await_args_list]
|
||||
assert [next(iter(frame)) for frame in frames] == ["setup"]
|
||||
assert "Be a pirate." in json.dumps(frames[0]["setup"]["systemInstruction"])
|
||||
|
||||
|
||||
@pytest.mark.parametrize("defer_setup", [False, True])
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_transcription_guardrail_refuses_gemini_audio_it_cannot_screen(monkeypatch, defer_setup):
|
||||
monkeypatch.setattr(litellm, "callbacks", [_transcription_guardrail()])
|
||||
session_update: Final = json.dumps(
|
||||
{"type": "session.update", "session": {"turn_detection": {"type": "server_vad", "create_response": False}}}
|
||||
)
|
||||
append: Final = json.dumps({"type": "input_audio_buffer.append", "audio": "AAAA"})
|
||||
commit: Final = json.dumps({"type": "input_audio_buffer.commit", "event_id": "commit_1"})
|
||||
streaming, backend_ws = _gemini_default_mode_streaming(
|
||||
monkeypatch,
|
||||
[session_update, append, append, commit, json.dumps({"type": "response.create"}), _TEXT_ITEM_CREATE],
|
||||
model="gemini-live-2.5-flash",
|
||||
defer_setup=defer_setup,
|
||||
)
|
||||
|
||||
await streaming.client_ack_messages()
|
||||
|
||||
frames: Final = [json.loads(call.args[0]) for call in backend_ws.send.await_args_list]
|
||||
assert [next(iter(frame)) for frame in frames] == ["setup", "clientContent"]
|
||||
refusals: Final = [
|
||||
event["error"]["event_id"]
|
||||
for event in _events_sent_to_client(streaming)
|
||||
if event["type"] == "error" and event["error"]["code"] == "audio_input_not_screenable"
|
||||
]
|
||||
assert refusals == [None, "commit_1"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("defer_setup", [False, True])
|
||||
@pytest.mark.parametrize("control_type", ["response.cancel", "conversation.item.delete", "input_audio_buffer.clear"])
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_control_event_before_session_update_leaves_the_setup_open(monkeypatch, defer_setup, control_type):
|
||||
control_event: Final = json.dumps({"type": control_type, "item_id": "item_x"})
|
||||
session_update: Final = json.dumps({"type": "session.update", "session": {"modalities": ["text"]}})
|
||||
streaming, backend_ws = _gemini_default_mode_streaming(
|
||||
monkeypatch,
|
||||
[control_event, session_update, _TEXT_ITEM_CREATE, _RESPONSE_CREATE],
|
||||
model="gemini-live-2.5-flash",
|
||||
defer_setup=defer_setup,
|
||||
)
|
||||
|
||||
await streaming.client_ack_messages()
|
||||
|
||||
frames: Final = [json.loads(call.args[0]) for call in backend_ws.send.await_args_list]
|
||||
assert [next(iter(frame)) for frame in frames] == ["setup", "clientContent"]
|
||||
assert frames[0]["setup"]["generationConfig"]["responseModalities"] == ["TEXT"]
|
||||
assert [event["type"] for event in _events_sent_to_client(streaming) if event["type"] == "error"] == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_deferred_setup_buffers_audio_until_backend_setup_complete(monkeypatch):
|
||||
"""Pipecat may send audio before session.update when setup is deferred."""
|
||||
|
|
@ -2533,14 +3173,18 @@ async def test_deferred_setup_buffers_audio_until_backend_setup_complete(monkeyp
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_deferred_setup_sends_session_update_before_buffered_audio(monkeypatch):
|
||||
async def test_deferred_setup_sends_the_held_session_update_when_more_audio_arrives(monkeypatch):
|
||||
monkeypatch.setattr(litellm, "gemini_live_defer_setup", True, raising=False)
|
||||
from litellm.llms.gemini.realtime.transformation import GeminiRealtimeConfig
|
||||
|
||||
client_ws = MagicMock()
|
||||
client_ws.send_text = AsyncMock()
|
||||
audio_msg = json.dumps({"type": "input_audio_buffer.append", "audio": "AA=="})
|
||||
session_update = json.dumps({"type": "session.update", "session": {"modalities": ["audio"]}})
|
||||
client_ws.receive_text = AsyncMock(side_effect=[audio_msg, session_update, ConnectionClosed(None, None)])
|
||||
later_audio_msg = json.dumps({"type": "input_audio_buffer.append", "audio": "AQ=="})
|
||||
client_ws.receive_text = AsyncMock(
|
||||
side_effect=[audio_msg, session_update, later_audio_msg, ConnectionClosed(None, None)]
|
||||
)
|
||||
backend_ws = MagicMock()
|
||||
backend_ws.send = AsyncMock()
|
||||
logging_obj: Final = _mock_logging()
|
||||
|
|
@ -2560,7 +3204,7 @@ async def test_deferred_setup_sends_session_update_before_buffered_audio(monkeyp
|
|||
sent_payload = json.loads(backend_ws.send.await_args_list[0].args[0])
|
||||
assert "setup" in sent_payload
|
||||
assert "realtimeInput" not in sent_payload
|
||||
assert streaming._pending_messages_until_setup == [audio_msg]
|
||||
assert streaming._pending_messages_until_setup == [audio_msg, later_audio_msg]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -3523,6 +4167,24 @@ async def test_upstream_refusal_after_a_synthetic_session_created_still_logs_a_f
|
|||
assert session.logging.logged_sessions == ()
|
||||
|
||||
|
||||
class _LazySetupRealtimeConfig(_PassthroughRealtimeConfig):
|
||||
def builds_setup_from_session_update(self) -> bool:
|
||||
return True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_upstream_outage_before_a_held_setup_was_sent_logs_the_session_not_a_refusal():
|
||||
upstream_close: Final = ConnectionClosed(Close(1011, "upstream outage"), None)
|
||||
session: Final = _relay_session(
|
||||
_client_ws_that_never_sends(), _backend_ws_closing_with(upstream_close), _LazySetupRealtimeConfig()
|
||||
)
|
||||
|
||||
await session.run()
|
||||
|
||||
assert session.logging.logged_failures == ()
|
||||
assert len(session.logging.logged_sessions) == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_upstream_close_after_relayed_events_still_logs_the_session_as_success():
|
||||
client_ws: Final = _client_ws_that_never_sends()
|
||||
|
|
|
|||
|
|
@ -2664,7 +2664,7 @@ def test_perplexity_streaming_dict_cost_bills_through_its_own_calculator():
|
|||
) == pytest.approx(0.00503)
|
||||
|
||||
|
||||
def test_openai_compatible_streaming_cost_is_priced_from_the_cost_map():
|
||||
def test_openai_compatible_streaming_reported_cost_is_billed_instead_of_the_cost_map():
|
||||
import litellm
|
||||
from litellm.cost_calculator import (
|
||||
get_response_cost_from_hidden_params,
|
||||
|
|
@ -2686,7 +2686,7 @@ def test_openai_compatible_streaming_cost_is_priced_from_the_cost_map():
|
|||
id="chatcmpl-openai-compatible",
|
||||
model=model,
|
||||
choices=[],
|
||||
usage=Usage(completion_tokens=5, prompt_tokens=10, total_tokens=15, cost=3_144_000),
|
||||
usage=Usage(completion_tokens=5, prompt_tokens=10, total_tokens=15, cost=0.037),
|
||||
)
|
||||
|
||||
CustomStreamWrapper._propagate_usage_cost_to_hidden_params(complete_response, "openai")
|
||||
|
|
@ -2698,7 +2698,7 @@ def test_openai_compatible_streaming_cost_is_priced_from_the_cost_map():
|
|||
custom_llm_provider="openai",
|
||||
call_type="completion",
|
||||
optional_params={},
|
||||
) == pytest.approx(2e-5)
|
||||
) == pytest.approx(0.037)
|
||||
|
||||
|
||||
def test_xai_streaming_reported_cost_still_takes_the_margin(monkeypatch):
|
||||
|
|
|
|||
|
|
@ -1962,6 +1962,97 @@ def test_token_counter_with_redacted_thinking_content():
|
|||
|
||||
assert token_counter(model=model, messages=with_block) == token_counter(model=model, messages=without_block)
|
||||
|
||||
|
||||
def test_token_counter_with_compaction_block():
|
||||
"""
|
||||
A replayed on-demand compaction block (Anthropic's signed summary of the earlier turns) counts its
|
||||
summary text like a text block and nothing for its signature. It used to raise, which turned every
|
||||
/v1/messages/count_tokens call on a compacted history into a 500 once the provider count fell back here.
|
||||
"""
|
||||
model = "anthropic/claude-sonnet-5-5"
|
||||
summary = "The user is building a recipe app in Python and asked for one-sentence class descriptions."
|
||||
follow_up = {"role": "user", "content": [{"type": "text", "text": "Now do the same for Ingredient."}]}
|
||||
signed_block = {"type": "compaction", "content": summary, "signature": "EqQBCkYIBRgCKkBjZ2xhc3M" * 40}
|
||||
|
||||
as_compaction = [{"role": "assistant", "content": [signed_block]}, follow_up]
|
||||
unsigned = [{"role": "assistant", "content": [{"type": "compaction", "content": summary}]}, follow_up]
|
||||
as_text = [{"role": "assistant", "content": [{"type": "text", "text": summary}]}, follow_up]
|
||||
|
||||
assert token_counter(model=model, messages=as_compaction) == token_counter(model=model, messages=as_text)
|
||||
assert token_counter(model=model, messages=unsigned) == token_counter(model=model, messages=as_text)
|
||||
|
||||
|
||||
def test_token_counter_counts_chat_compaction_blocks_like_their_summary_text():
|
||||
"""
|
||||
Chat clients replay a compaction block in the assistant message's provider_specific_fields, and inference
|
||||
sends it to Anthropic as a content block, so the local count includes its summary text and skips its signature.
|
||||
"""
|
||||
model = "anthropic/claude-sonnet-5-5"
|
||||
summary = "The user is building a recipe app in Python and asked for one-sentence class descriptions."
|
||||
follow_up = {"role": "user", "content": "Now do the same for Ingredient."}
|
||||
signed_block = {"type": "compaction", "content": summary, "signature": "EqQBCkYIBRgCKkBjZ2xhc3M" * 40}
|
||||
|
||||
replayed = [
|
||||
{"role": "assistant", "content": None, "provider_specific_fields": {"compaction_blocks": [signed_block]}},
|
||||
follow_up,
|
||||
]
|
||||
as_text = [{"role": "assistant", "content": summary}, follow_up]
|
||||
|
||||
assert token_counter(model=model, messages=replayed) == token_counter(model=model, messages=as_text)
|
||||
|
||||
|
||||
class _BosEncoder:
|
||||
def encode_batch_fast(self, texts: list[str]) -> list[list[int]]:
|
||||
return [[0] * (len(text.split()) + 1) for text in texts]
|
||||
|
||||
|
||||
def test_token_counter_charges_nothing_for_a_chat_compaction_block_with_no_summary():
|
||||
"""
|
||||
A tokenizer that prepends a BOS token charges one token for the empty string, so a replayed
|
||||
compaction block with an empty summary must not reach it, like an empty thinking block.
|
||||
"""
|
||||
model = "anthropic/claude-sonnet-5-5"
|
||||
tokenizer = {"type": "huggingface_tokenizer", "tokenizer": _BosEncoder()}
|
||||
follow_up = {"role": "user", "content": "Now do the same for Ingredient."}
|
||||
replayed = [
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": None,
|
||||
"provider_specific_fields": {"compaction_blocks": [{"type": "compaction", "content": ""}]},
|
||||
},
|
||||
follow_up,
|
||||
]
|
||||
bare = [{"role": "assistant", "content": None}, follow_up]
|
||||
|
||||
assert token_counter(model=model, messages=replayed, custom_tokenizer=tokenizer) == token_counter(
|
||||
model=model, messages=bare, custom_tokenizer=tokenizer
|
||||
)
|
||||
|
||||
|
||||
def test_count_content_list_skips_the_tokenizer_for_empty_block_text():
|
||||
"""
|
||||
An empty thinking block, a redacted_thinking block (no thinking text at all) and a compaction
|
||||
block with no summary carry no tokens, even under a tokenizer that charges for the empty string
|
||||
(the llama-2 tokenizer prepends its BOS token to every encode call). Counting them through the
|
||||
tokenizer added one token per block on such models.
|
||||
"""
|
||||
from litellm.litellm_core_utils.token_counter import _count_content_list
|
||||
|
||||
def charges_for_empty_text(text: str) -> int:
|
||||
return len(text.split()) + 1
|
||||
|
||||
empty_blocks = [
|
||||
{"type": "thinking", "thinking": "", "signature": "EqcLCkYICxgCKkCrqu6lP..."},
|
||||
{"type": "thinking", "thinking": None, "signature": "EqcLCkYICxgCKkCrqu6lP..."},
|
||||
{"type": "redacted_thinking", "data": "EmwKAhgBEgy3va3pzix/LafPsn4aDFIT2Xlxh0L5L8rLVyIw"},
|
||||
{"type": "compaction", "content": ""},
|
||||
{"type": "text", "text": ""},
|
||||
]
|
||||
|
||||
assert _count_content_list(charges_for_empty_text, empty_blocks, False, None) == 0
|
||||
assert _count_content_list(charges_for_empty_text, [{"type": "thinking", "thinking": "two words"}], False, None) == 3
|
||||
|
||||
|
||||
def test_token_counter_with_tool_reference_block():
|
||||
"""
|
||||
Regression test: a message containing an Anthropic tool-search
|
||||
|
|
|
|||
|
|
@ -6,6 +6,7 @@ with guardrail transformations, specifically testing edge cases with empty choic
|
|||
"""
|
||||
|
||||
import json
|
||||
import re
|
||||
from typing import Any, Literal, Optional, Final
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
|
|
@ -552,13 +553,25 @@ class TestAnthropicMessagesHandlerStreamingOutputProcessing:
|
|||
assert chunks == original
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_deliver_ended_stream_tool_use_rewrite_with_server_tool_use_block_fails_closed(self):
|
||||
from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite
|
||||
|
||||
async def test_deliver_ended_stream_tool_use_rewrite_lands_on_server_tool_use_and_tool_use_blocks(self):
|
||||
handler = AnthropicMessagesHandler()
|
||||
server_tool_use = [
|
||||
("content_block_start", {"type": "content_block_start", "index": 0, "content_block": {"type": "server_tool_use", "id": "srvtoolu_1", "name": "web_search", "input": {}}}),
|
||||
("content_block_delta", {"type": "content_block_delta", "index": 0, "delta": {"type": "input_json_delta", "partial_json": '{"query": "fruit"}'}}),
|
||||
(
|
||||
"content_block_start",
|
||||
{
|
||||
"type": "content_block_start",
|
||||
"index": 0,
|
||||
"content_block": {"type": "server_tool_use", "id": "srvtoolu_1", "name": "web_search", "input": {}},
|
||||
},
|
||||
),
|
||||
(
|
||||
"content_block_delta",
|
||||
{
|
||||
"type": "content_block_delta",
|
||||
"index": 0,
|
||||
"delta": {"type": "input_json_delta", "partial_json": '{"query": "fruit"}'},
|
||||
},
|
||||
),
|
||||
("content_block_stop", {"type": "content_block_stop", "index": 0}),
|
||||
]
|
||||
tool_use = self._ended_tool_use_sse_chunks()
|
||||
|
|
@ -568,6 +581,26 @@ class TestAnthropicMessagesHandlerStreamingOutputProcessing:
|
|||
+ [chunk.replace(b'"index": 0', b'"index": 1') for chunk in tool_use[1:]]
|
||||
)
|
||||
|
||||
await handler.process_output_streaming_response(
|
||||
responses_so_far=chunks,
|
||||
guardrail_to_apply=self._argument_masking_guardrail(),
|
||||
litellm_logging_obj=MagicMock(),
|
||||
deliver_ended_stream_rewrites=True,
|
||||
)
|
||||
|
||||
assert self._partial_jsons(chunks) == ['{"fruit": "[MASKED]"}', '{"fruit": "[MASKED]"}', "", ""]
|
||||
raw = b"".join(chunks).decode()
|
||||
assert '"id": "srvtoolu_1"' in raw and '"id": "toolu_1"' in raw
|
||||
assert "persim" not in raw and '"query"' not in raw
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_deliver_ended_stream_tool_use_rewrite_on_an_unindexed_tool_block_fails_closed(self):
|
||||
from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite
|
||||
|
||||
handler = AnthropicMessagesHandler()
|
||||
tool_use = self._ended_tool_use_sse_chunks()
|
||||
chunks = tool_use[:1] + [tool_use[1].replace(b'"index": 0', b'"index": "0"')] + tool_use[2:]
|
||||
|
||||
with pytest.raises(UndeliverableStreamRewrite):
|
||||
await handler.process_output_streaming_response(
|
||||
responses_so_far=chunks,
|
||||
|
|
@ -576,6 +609,270 @@ class TestAnthropicMessagesHandlerStreamingOutputProcessing:
|
|||
deliver_ended_stream_rewrites=True,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _python_masking_guardrail() -> CustomGuardrail:
|
||||
class MaskPython(CustomGuardrail):
|
||||
async def apply_guardrail(
|
||||
self,
|
||||
inputs: GenericGuardrailAPIInputs,
|
||||
request_data: dict[str, object],
|
||||
input_type: Literal["request", "response"],
|
||||
logging_obj: Optional[LiteLLMLoggingObj] = None,
|
||||
) -> GenericGuardrailAPIInputs:
|
||||
for tool_call in inputs.get("tool_calls", []):
|
||||
if isinstance(tool_call, ChatCompletionMessageToolCall):
|
||||
tool_call.function.arguments = re.sub(
|
||||
"python", "[MASKED]", tool_call.function.arguments, flags=re.I
|
||||
)
|
||||
return {
|
||||
**inputs,
|
||||
"texts": [re.sub("python", "[MASKED]", text, flags=re.I) for text in inputs.get("texts", [])],
|
||||
}
|
||||
|
||||
return MaskPython(guardrail_name="test")
|
||||
|
||||
@staticmethod
|
||||
def _python_texts_masking_guardrail() -> CustomGuardrail:
|
||||
class MaskPythonTexts(CustomGuardrail):
|
||||
async def apply_guardrail(
|
||||
self,
|
||||
inputs: GenericGuardrailAPIInputs,
|
||||
request_data: dict[str, object],
|
||||
input_type: Literal["request", "response"],
|
||||
logging_obj: Optional[LiteLLMLoggingObj] = None,
|
||||
) -> GenericGuardrailAPIInputs:
|
||||
return {
|
||||
**inputs,
|
||||
"texts": [re.sub("python", "[MASKED]", text, flags=re.I) for text in inputs.get("texts", [])],
|
||||
}
|
||||
|
||||
return MaskPythonTexts(guardrail_name="test")
|
||||
|
||||
@staticmethod
|
||||
def _web_search_result_block() -> dict[str, object]:
|
||||
return {
|
||||
"type": "web_search_tool_result",
|
||||
"tool_use_id": "srvtoolu_1",
|
||||
"content": [
|
||||
{
|
||||
"type": "web_search_result",
|
||||
"url": "https://docs.python.org/3.14/",
|
||||
"title": "Python 3.14 docs",
|
||||
"encrypted_content": "",
|
||||
"page_age": None,
|
||||
},
|
||||
{
|
||||
"type": "web_search_result",
|
||||
"url": "https://peps.example/745",
|
||||
"title": "PEP 745",
|
||||
"encrypted_content": "",
|
||||
"page_age": None,
|
||||
},
|
||||
],
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def _web_search_stream_chunks(cls, server_tool_block_type: str = "server_tool_use") -> list[bytes]:
|
||||
events = [
|
||||
(
|
||||
"message_start",
|
||||
{
|
||||
"type": "message_start",
|
||||
"message": {
|
||||
"id": "msg_1",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"model": "claude-sonnet-5-5",
|
||||
"content": [],
|
||||
"stop_reason": None,
|
||||
"usage": {"input_tokens": 1, "output_tokens": 0},
|
||||
},
|
||||
},
|
||||
),
|
||||
(
|
||||
"content_block_start",
|
||||
{
|
||||
"type": "content_block_start",
|
||||
"index": 0,
|
||||
"content_block": {
|
||||
"type": server_tool_block_type,
|
||||
"id": "srvtoolu_1",
|
||||
"name": "web_search",
|
||||
"input": {},
|
||||
},
|
||||
},
|
||||
),
|
||||
(
|
||||
"content_block_delta",
|
||||
{
|
||||
"type": "content_block_delta",
|
||||
"index": 0,
|
||||
"delta": {"type": "input_json_delta", "partial_json": '{"query": "python 3.14 release date"}'},
|
||||
},
|
||||
),
|
||||
("content_block_stop", {"type": "content_block_stop", "index": 0}),
|
||||
(
|
||||
"content_block_start",
|
||||
{"type": "content_block_start", "index": 1, "content_block": cls._web_search_result_block()},
|
||||
),
|
||||
("content_block_stop", {"type": "content_block_stop", "index": 1}),
|
||||
(
|
||||
"content_block_start",
|
||||
{"type": "content_block_start", "index": 2, "content_block": {"type": "text", "text": ""}},
|
||||
),
|
||||
(
|
||||
"content_block_delta",
|
||||
{
|
||||
"type": "content_block_delta",
|
||||
"index": 2,
|
||||
"delta": {"type": "text_delta", "text": "Python 3.14 shipped"},
|
||||
},
|
||||
),
|
||||
("content_block_stop", {"type": "content_block_stop", "index": 2}),
|
||||
(
|
||||
"message_delta",
|
||||
{
|
||||
"type": "message_delta",
|
||||
"delta": {"stop_reason": "end_turn", "stop_sequence": None},
|
||||
"usage": {"output_tokens": 2},
|
||||
},
|
||||
),
|
||||
("message_stop", {"type": "message_stop"}),
|
||||
]
|
||||
return [f"event: {name}\ndata: {json.dumps(payload)}\n\n".encode() for name, payload in events]
|
||||
|
||||
@pytest.mark.parametrize("server_tool_block_type", ["server_tool_use", "tool_use"])
|
||||
@pytest.mark.parametrize(
|
||||
"guardrail_factory",
|
||||
[_python_masking_guardrail, _python_texts_masking_guardrail],
|
||||
ids=["texts_and_tool_calls", "texts_only"],
|
||||
)
|
||||
@pytest.mark.asyncio
|
||||
async def test_deliver_ended_stream_rewrites_mask_the_web_search_query_results_and_text(
|
||||
self, guardrail_factory, server_tool_block_type
|
||||
):
|
||||
chunks = self._web_search_stream_chunks(server_tool_block_type)
|
||||
|
||||
await AnthropicMessagesHandler().process_output_streaming_response(
|
||||
responses_so_far=chunks,
|
||||
guardrail_to_apply=guardrail_factory(),
|
||||
litellm_logging_obj=MagicMock(),
|
||||
deliver_ended_stream_rewrites=True,
|
||||
)
|
||||
|
||||
raw = b"".join(chunks).decode()
|
||||
assert "python" not in raw.lower()
|
||||
result_block = next(
|
||||
json.loads(line[len("data:") :])["content_block"]
|
||||
for line in raw.split("\n")
|
||||
if line.startswith("data:") and '"web_search_tool_result"' in line
|
||||
)
|
||||
assert [(result["title"], result["url"]) for result in result_block["content"]] == [
|
||||
("[MASKED] 3.14 docs", "https://docs.[MASKED].org/3.14/"),
|
||||
("PEP 745", "https://peps.example/745"),
|
||||
]
|
||||
assert '"id": "srvtoolu_1"' in raw and '"tool_use_id": "srvtoolu_1"' in raw
|
||||
|
||||
@staticmethod
|
||||
def _last_text_redacting_guardrail() -> CustomGuardrail:
|
||||
class RedactLastText(CustomGuardrail):
|
||||
async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None):
|
||||
return {**inputs, "texts": ["[REDACTED] " + inputs["texts"][-1]]}
|
||||
|
||||
return RedactLastText(guardrail_name="test")
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_output_response_keeps_the_reply_as_the_last_text_a_guardrail_scans(self):
|
||||
response = {
|
||||
"id": "msg_1",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"model": "claude-sonnet-5-5",
|
||||
"content": [
|
||||
{
|
||||
"type": "server_tool_use",
|
||||
"id": "srvtoolu_1",
|
||||
"name": "web_search",
|
||||
"input": {"query": "python 3.14 release date"},
|
||||
},
|
||||
self._web_search_result_block(),
|
||||
{"type": "text", "text": "Python 3.14 shipped"},
|
||||
],
|
||||
"stop_reason": "end_turn",
|
||||
}
|
||||
|
||||
await AnthropicMessagesHandler().process_output_response(
|
||||
response=response,
|
||||
guardrail_to_apply=self._last_text_redacting_guardrail(),
|
||||
litellm_logging_obj=MagicMock(),
|
||||
)
|
||||
|
||||
assert response["content"][0]["input"] == {"query": "python 3.14 release date"}
|
||||
assert response["content"][1] == self._web_search_result_block()
|
||||
assert response["content"][2] == {"type": "text", "text": "[REDACTED] Python 3.14 shipped"}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_deliver_ended_stream_keeps_the_reply_as_the_last_text_a_guardrail_scans(self):
|
||||
chunks = self._web_search_stream_chunks()
|
||||
|
||||
await AnthropicMessagesHandler().process_output_streaming_response(
|
||||
responses_so_far=chunks,
|
||||
guardrail_to_apply=self._last_text_redacting_guardrail(),
|
||||
litellm_logging_obj=MagicMock(),
|
||||
deliver_ended_stream_rewrites=True,
|
||||
)
|
||||
|
||||
raw = b"".join(chunks).decode()
|
||||
assert "[REDACTED] Python 3.14 shipped" in raw
|
||||
assert raw.count("[REDACTED]") == 1
|
||||
assert json.dumps(self._web_search_result_block()) in raw
|
||||
|
||||
def test_ended_stream_scan_key_differs_from_a_mid_stream_key_by_the_web_search_results(self):
|
||||
handler = AnthropicMessagesHandler()
|
||||
chunks = self._web_search_stream_chunks()
|
||||
|
||||
mid_stream_key = handler.get_streaming_scan_key(chunks[:-2])
|
||||
ended_key = handler.get_streaming_scan_key(chunks)
|
||||
|
||||
assert mid_stream_key.texts == ("Python 3.14 shipped",)
|
||||
assert ended_key != mid_stream_key
|
||||
assert ended_key.texts[-1] == "Python 3.14 shipped"
|
||||
assert {"Python 3.14 docs", "https://docs.python.org/3.14/"} <= set(ended_key.texts)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_output_response_masks_the_web_search_query_results_and_text(self):
|
||||
response = {
|
||||
"id": "msg_1",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"model": "claude-sonnet-5-5",
|
||||
"content": [
|
||||
{
|
||||
"type": "server_tool_use",
|
||||
"id": "srvtoolu_1",
|
||||
"name": "web_search",
|
||||
"input": {"query": "python 3.14 release date"},
|
||||
},
|
||||
self._web_search_result_block(),
|
||||
{"type": "text", "text": "Python 3.14 shipped"},
|
||||
],
|
||||
"stop_reason": "end_turn",
|
||||
}
|
||||
|
||||
await AnthropicMessagesHandler().process_output_response(
|
||||
response=response,
|
||||
guardrail_to_apply=self._python_masking_guardrail(),
|
||||
litellm_logging_obj=MagicMock(),
|
||||
)
|
||||
|
||||
assert "python" not in json.dumps(response).lower()
|
||||
assert response["content"][0]["input"] == {"query": "[MASKED] 3.14 release date"}
|
||||
assert [(result["title"], result["url"]) for result in response["content"][1]["content"]] == [
|
||||
("[MASKED] 3.14 docs", "https://docs.[MASKED].org/3.14/"),
|
||||
("PEP 745", "https://peps.example/745"),
|
||||
]
|
||||
assert response["content"][2] == {"type": "text", "text": "[MASKED] 3.14 shipped"}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ended_stream_rewrite_leaves_chunks_untouched_by_default(self):
|
||||
handler = AnthropicMessagesHandler()
|
||||
|
|
|
|||
|
|
@ -1,6 +1,7 @@
|
|||
import json
|
||||
import os
|
||||
import sys
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
|
||||
|
|
@ -70,3 +71,70 @@ def test_prepare_completion_kwargs_keeps_prompt_cache_key_through_responses_rero
|
|||
)
|
||||
assert completion_kwargs["model"] == "openai/responses/gpt-5.6-luna"
|
||||
assert completion_kwargs["prompt_cache_key"] == "session-abc"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("api_base", "azure_api_base_env", "derives_key"),
|
||||
[
|
||||
(None, None, True),
|
||||
("https://api.mistral.ai/v1", None, True),
|
||||
("https://my-resource.openai.azure.com/openai/v1", None, False),
|
||||
(None, "https://my-resource.services.ai.azure.com/openai", False),
|
||||
],
|
||||
)
|
||||
def test_prepare_completion_kwargs_derives_prompt_cache_key_only_for_a_mistral_hosted_endpoint(
|
||||
monkeypatch: pytest.MonkeyPatch, api_base: str | None, azure_api_base_env: str | None, derives_key: bool
|
||||
) -> None:
|
||||
if azure_api_base_env is None:
|
||||
monkeypatch.delenv("MISTRAL_AZURE_API_BASE", raising=False)
|
||||
else:
|
||||
monkeypatch.setenv("MISTRAL_AZURE_API_BASE", azure_api_base_env)
|
||||
|
||||
completion_kwargs: Final = _prepare(
|
||||
"mistral/mistral-small-latest", {"custom_llm_provider": "mistral", "api_base": api_base}
|
||||
)
|
||||
|
||||
assert completion_kwargs["user"] == CLAUDE_CODE_USER_ID
|
||||
assert ("prompt_cache_key" in completion_kwargs) is derives_key
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("codestral_api_base_env", "derives_key"),
|
||||
[(None, True), ("https://my-resource.services.ai.azure.com/openai/v1", False)],
|
||||
)
|
||||
def test_prepare_completion_kwargs_derives_prompt_cache_key_for_codestral_only_on_its_mistral_host(
|
||||
monkeypatch: pytest.MonkeyPatch, codestral_api_base_env: str | None, derives_key: bool
|
||||
) -> None:
|
||||
monkeypatch.delenv("MISTRAL_AZURE_API_BASE", raising=False)
|
||||
if codestral_api_base_env is None:
|
||||
monkeypatch.delenv("CODESTRAL_API_BASE", raising=False)
|
||||
else:
|
||||
monkeypatch.setenv("CODESTRAL_API_BASE", codestral_api_base_env)
|
||||
|
||||
completion_kwargs: Final = _prepare("codestral/codestral-latest", {"custom_llm_provider": "codestral"})
|
||||
|
||||
assert ("prompt_cache_key" in completion_kwargs) is derives_key
|
||||
|
||||
|
||||
def test_prepare_completion_kwargs_derives_no_prompt_cache_key_mistral_cannot_take_from_a_non_ascii_session() -> None:
|
||||
non_ascii_user_id: Final = json.dumps({"device_id": "d" * 64, "account_uuid": "", "session_id": "sitzung-ü"})
|
||||
|
||||
mistral_kwargs: Final = LiteLLMMessagesToCompletionTransformationHandler._prepare_completion_kwargs(
|
||||
max_tokens=1024,
|
||||
messages=MESSAGES,
|
||||
model="mistral/mistral-small-latest",
|
||||
metadata={"user_id": non_ascii_user_id},
|
||||
thinking=None,
|
||||
extra_kwargs={"custom_llm_provider": "mistral"},
|
||||
)[0]
|
||||
openai_kwargs: Final = LiteLLMMessagesToCompletionTransformationHandler._prepare_completion_kwargs(
|
||||
max_tokens=1024,
|
||||
messages=MESSAGES,
|
||||
model="openai/gpt-5.6-luna",
|
||||
metadata={"user_id": non_ascii_user_id},
|
||||
thinking=None,
|
||||
extra_kwargs={"custom_llm_provider": "openai"},
|
||||
)[0]
|
||||
|
||||
assert "prompt_cache_key" not in mistral_kwargs
|
||||
assert openai_kwargs["prompt_cache_key"] == "sitzung-ü"
|
||||
|
|
|
|||
|
|
@ -24,7 +24,7 @@ from litellm.llms.tencent.messages.transformation import TencentAnthropicMessage
|
|||
from litellm.llms.vertex_ai.vertex_ai_partner_models.anthropic.experimental_pass_through.transformation import (
|
||||
VertexAIPartnerModelsAnthropicMessagesConfig,
|
||||
)
|
||||
from litellm.types.llms.anthropic import AnthropicFinishReason
|
||||
from litellm.types.llms.anthropic import ANTHROPIC_BETA_HEADER_VALUES, AnthropicFinishReason
|
||||
from litellm.types.router import GenericLiteLLMParams
|
||||
from tests.unit.llms.anthropic.test_anthropic_wif import (
|
||||
ScriptedPoster,
|
||||
|
|
@ -321,6 +321,66 @@ def _derived_beta_header(config: AnthropicMessagesConfig) -> dict:
|
|||
)
|
||||
|
||||
|
||||
_COMPACT_EDIT_BETA: Final = ANTHROPIC_BETA_HEADER_VALUES.COMPACT_2026_01_12.value
|
||||
_CLEARING_EDIT_BETA: Final = ANTHROPIC_BETA_HEADER_VALUES.CONTEXT_MANAGEMENT_2025_06_27.value
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("context_management", "betas"),
|
||||
(
|
||||
pytest.param({"edits": [{"type": "compact_20260112"}]}, {_COMPACT_EDIT_BETA}, id="compact"),
|
||||
pytest.param({"edits": [{"type": "clear_tool_uses_20250919"}]}, {_CLEARING_EDIT_BETA}, id="clearing"),
|
||||
pytest.param(
|
||||
{"edits": [{"type": "compact_20260112"}, {"type": "clear_thinking_20251015"}]},
|
||||
{_COMPACT_EDIT_BETA, _CLEARING_EDIT_BETA},
|
||||
id="both",
|
||||
),
|
||||
pytest.param(
|
||||
[{"type": "compaction", "compact_threshold": 150000}], {_COMPACT_EDIT_BETA}, id="openai_list_form"
|
||||
),
|
||||
),
|
||||
)
|
||||
def test_context_management_edits_add_the_beta_each_edit_needs(
|
||||
context_management: JsonValue, betas: set[str], local_beta_headers_config: None
|
||||
) -> None:
|
||||
headers: Final = AnthropicMessagesConfig()._update_headers_with_anthropic_beta(
|
||||
headers={}, optional_params={"context_management": context_management}, messages=[]
|
||||
)
|
||||
|
||||
assert set(headers["anthropic-beta"].split(",")) == betas
|
||||
|
||||
|
||||
_UNMAPPABLE_LIST_FORM: Final = [{"type": "clear_tool_uses_20250919"}]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("config_class", [AnthropicMessagesConfig, VertexAIPartnerModelsAnthropicMessagesConfig])
|
||||
@pytest.mark.parametrize(
|
||||
("context_management", "drop_params", "forwarded"),
|
||||
(
|
||||
pytest.param(_UNMAPPABLE_LIST_FORM, True, None, id="unmappable_dropped_under_drop_params"),
|
||||
pytest.param(_UNMAPPABLE_LIST_FORM, False, _UNMAPPABLE_LIST_FORM, id="unmappable_kept_for_the_provider_400"),
|
||||
pytest.param(
|
||||
[{"type": "compaction"}],
|
||||
True,
|
||||
{"edits": [{"type": "compact_20260112"}]},
|
||||
id="mappable_still_mapped_under_drop_params",
|
||||
),
|
||||
),
|
||||
)
|
||||
def test_context_management_the_list_form_cannot_map_is_dropped_only_under_drop_params(
|
||||
config_class: type[AnthropicMessagesConfig], context_management: JsonValue, drop_params: bool, forwarded: JsonValue
|
||||
) -> None:
|
||||
request: Final = config_class().transform_anthropic_messages_request(
|
||||
model="claude-haiku-5-5",
|
||||
messages=[{"role": "user", "content": "Hi"}],
|
||||
anthropic_messages_optional_request_params={"max_tokens": 50, "context_management": context_management},
|
||||
litellm_params=GenericLiteLLMParams(drop_params=drop_params),
|
||||
headers={},
|
||||
)
|
||||
|
||||
assert request.get("context_management") == forwarded
|
||||
|
||||
|
||||
@pytest.mark.parametrize("config", _classifier_configs(), ids=type)
|
||||
def test_routes_whose_platform_runs_the_classifier_keep_safeguards_and_the_beta(
|
||||
config: AnthropicMessagesConfig, local_beta_headers_config: None
|
||||
|
|
|
|||
|
|
@ -1,6 +1,10 @@
|
|||
import json
|
||||
from collections.abc import Iterable, Mapping
|
||||
from typing import Final
|
||||
from collections.abc import Iterable, Mapping, Sequence
|
||||
from typing import Final, cast
|
||||
|
||||
import anthropic
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
from litellm.llms.anthropic.messages.native.fake_stream_iterator import (
|
||||
FakeAnthropicMessagesStreamIterator,
|
||||
|
|
@ -91,3 +95,140 @@ def test_message_delta_omits_absent_optional_fields():
|
|||
"usage": {"input_tokens": 10, "output_tokens": 2},
|
||||
}
|
||||
assert _message_delta(response) == {"stop_reason": "end_turn", "stop_sequence": None}
|
||||
|
||||
|
||||
_WEB_SEARCH_RESULT_BLOCK: Final = {
|
||||
"type": "web_search_tool_result",
|
||||
"tool_use_id": "srvtoolu_01",
|
||||
"content": [
|
||||
{
|
||||
"type": "web_search_result",
|
||||
"url": "https://example.com/a",
|
||||
"title": "Example A",
|
||||
"encrypted_content": "abc123",
|
||||
}
|
||||
],
|
||||
}
|
||||
|
||||
|
||||
def _content_events(content: Sequence[Mapping[str, object]]) -> tuple[tuple[str, dict[str, object]], ...]:
|
||||
response: Final = {
|
||||
"id": "msg_ws",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"model": "claude-sonnet-4-6",
|
||||
"content": list(content),
|
||||
"stop_reason": "end_turn",
|
||||
"usage": {"input_tokens": 10, "output_tokens": 5},
|
||||
}
|
||||
return tuple(
|
||||
event for event in _decoded_events(FakeAnthropicMessagesStreamIterator(response=response))
|
||||
if event[0].startswith("content_block_")
|
||||
)
|
||||
|
||||
|
||||
def test_server_tool_use_streams_empty_start_then_input_json_delta():
|
||||
"""A real Anthropic stream opens server_tool_use with an empty input and
|
||||
sends the query as an input_json_delta, which SDK accumulators read."""
|
||||
caller: Final = {"type": "code_execution_20260120", "tool_id": "srvtoolu_00"}
|
||||
events: Final = _content_events(
|
||||
[
|
||||
{
|
||||
"type": "server_tool_use",
|
||||
"id": "srvtoolu_01",
|
||||
"name": "web_search",
|
||||
"input": {"query": "example search"},
|
||||
"caller": caller,
|
||||
}
|
||||
]
|
||||
)
|
||||
assert [event_type for event_type, _ in events] == [
|
||||
"content_block_start",
|
||||
"content_block_delta",
|
||||
"content_block_stop",
|
||||
]
|
||||
assert events[0][1]["content_block"] == {
|
||||
"type": "server_tool_use",
|
||||
"id": "srvtoolu_01",
|
||||
"name": "web_search",
|
||||
"input": {},
|
||||
"caller": caller,
|
||||
}
|
||||
delta: Final = events[1][1]["delta"]
|
||||
assert isinstance(delta, dict)
|
||||
assert delta["type"] == "input_json_delta"
|
||||
assert json.loads(delta["partial_json"]) == {"query": "example search"}
|
||||
|
||||
|
||||
def test_web_search_tool_result_rides_whole_in_content_block_start():
|
||||
events: Final = _content_events([_WEB_SEARCH_RESULT_BLOCK])
|
||||
assert [event_type for event_type, _ in events] == ["content_block_start", "content_block_stop"]
|
||||
assert events[0][1]["content_block"] == _WEB_SEARCH_RESULT_BLOCK
|
||||
|
||||
|
||||
def test_native_search_blocks_keep_indices_aligned_with_text():
|
||||
events: Final = _content_events(
|
||||
[
|
||||
{"type": "server_tool_use", "id": "srvtoolu_01", "name": "web_search", "input": {"query": "q"}},
|
||||
_WEB_SEARCH_RESULT_BLOCK,
|
||||
{"type": "text", "text": "Here is what I found."},
|
||||
]
|
||||
)
|
||||
starts: Final = tuple(payload for event_type, payload in events if event_type == "content_block_start")
|
||||
stops: Final = tuple(payload["index"] for event_type, payload in events if event_type == "content_block_stop")
|
||||
assert [(start["index"], cast(dict[str, object], start["content_block"])["type"]) for start in starts] == [
|
||||
(0, "server_tool_use"),
|
||||
(1, "web_search_tool_result"),
|
||||
(2, "text"),
|
||||
]
|
||||
assert stops == (0, 1, 2)
|
||||
|
||||
|
||||
_CALLER: Final = {"type": "code_execution_20260120", "tool_id": "srvtoolu_00"}
|
||||
|
||||
_REPLAYED_BLOCKS: Final = (
|
||||
{"type": "thinking", "thinking": "Look up the release page.", "signature": "sig_01"},
|
||||
{"type": "redacted_thinking", "data": "EncryptedThinking=="},
|
||||
{"type": "tool_use", "id": "toolu_01", "name": "get_weather", "input": {"city": "Paris"}, "caller": _CALLER},
|
||||
{"type": "server_tool_use", "id": "srvtoolu_01", "name": "web_fetch", "input": {"url": "https://example.com/a"}, "caller": _CALLER},
|
||||
_WEB_SEARCH_RESULT_BLOCK,
|
||||
{
|
||||
"type": "text",
|
||||
"text": "Example A is the newest release.",
|
||||
"citations": [
|
||||
{
|
||||
"type": "web_search_result_location",
|
||||
"cited_text": "Example A was released today.",
|
||||
"url": "https://example.com/a",
|
||||
"title": "Example A",
|
||||
"encrypted_index": "idx_01",
|
||||
}
|
||||
],
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def _sdk_rebuilt_content(block: Mapping[str, object]) -> list[dict[str, object]]:
|
||||
response: Final = {
|
||||
"id": "msg_replay",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"model": "claude-sonnet-5-5",
|
||||
"content": [block],
|
||||
"stop_reason": "end_turn",
|
||||
"usage": {"input_tokens": 10, "output_tokens": 5},
|
||||
}
|
||||
body: Final = b"".join(FakeAnthropicMessagesStreamIterator(response=response))
|
||||
transport: Final = httpx.MockTransport(
|
||||
lambda _request: httpx.Response(200, content=body, headers={"content-type": "text/event-stream"})
|
||||
)
|
||||
client: Final = anthropic.Anthropic(api_key="sk-ant-test", http_client=httpx.Client(transport=transport))
|
||||
with client.messages.stream(
|
||||
model="claude-sonnet-5-5", max_tokens=64, messages=[{"role": "user", "content": "hi"}]
|
||||
) as stream:
|
||||
return [content_block.model_dump(exclude_none=True) for content_block in stream.get_final_message().content]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("block", _REPLAYED_BLOCKS, ids=lambda block: str(block["type"]))
|
||||
def test_anthropic_sdk_rebuilds_every_block_field_from_the_stream(block: Mapping[str, object]):
|
||||
assert _sdk_rebuilt_content(block) == [block]
|
||||
|
|
|
|||
|
|
@ -26,7 +26,24 @@ from litellm.llms.azure_ai.anthropic.count_tokens.token_counter import AzureAIAn
|
|||
from litellm.llms.azure_ai.anthropic.count_tokens.transformation import (
|
||||
AzureAIAnthropicCountTokensConfig,
|
||||
)
|
||||
from litellm.types.utils import TokenCountMessagesFormat, TokenCountResponse
|
||||
from litellm.types.utils import CountContextManagement, TokenCountMessagesFormat, TokenCountResponse
|
||||
from litellm.constants import ANTHROPIC_TOKEN_COUNTING_BETA_VERSION
|
||||
from litellm.types.llms.anthropic import ANTHROPIC_BETA_HEADER_VALUES
|
||||
|
||||
COMPACTION_BETA = ANTHROPIC_BETA_HEADER_VALUES.COMPACT_2026_09_04.value
|
||||
COMPACTED_HISTORY = [
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{
|
||||
"type": "compaction",
|
||||
"content": "The user is building a recipe app and asked for one-sentence class descriptions.",
|
||||
"signature": "EqQBCkYIBRgCKkBjZ2xhc3M" * 40,
|
||||
}
|
||||
],
|
||||
},
|
||||
{"role": "user", "content": "Now do the same for Ingredient."},
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
|
|
@ -890,3 +907,115 @@ async def test_counter_sends_the_structured_output_tool_choice_and_thinking_chat
|
|||
inference["tool_choice"],
|
||||
inference["tools"],
|
||||
)
|
||||
|
||||
|
||||
def test_headers_add_the_compaction_beta_only_for_a_compacted_history():
|
||||
config = AnthropicCountTokensConfig()
|
||||
auth_header = {"x-api-key": "sk-ant-api03-test-key"}
|
||||
|
||||
with_block = config.get_count_tokens_headers(auth_header, COMPACTED_HISTORY)["anthropic-beta"].split(",")
|
||||
plain = config.get_count_tokens_headers(auth_header, [{"role": "user", "content": "hi"}])["anthropic-beta"]
|
||||
|
||||
assert COMPACTION_BETA in with_block
|
||||
assert ANTHROPIC_TOKEN_COUNTING_BETA_VERSION in with_block
|
||||
assert plain == ANTHROPIC_TOKEN_COUNTING_BETA_VERSION
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_handler_sends_the_compaction_beta_with_a_compacted_history(httpx_transport_clients):
|
||||
with respx.mock:
|
||||
route = respx.post("https://gateway.example/v1/messages/count_tokens").mock(
|
||||
return_value=httpx.Response(200, json={"input_tokens": 404})
|
||||
)
|
||||
result = await AnthropicCountTokensHandler().handle_count_tokens_request(
|
||||
model="claude-sonnet-5-5",
|
||||
messages=COMPACTED_HISTORY,
|
||||
auth_header={"x-api-key": "sk-ant-api03-test-key"},
|
||||
api_base="https://gateway.example",
|
||||
)
|
||||
|
||||
assert result == {"input_tokens": 404}
|
||||
betas = route.calls.last.request.headers["anthropic-beta"].split(",")
|
||||
assert COMPACTION_BETA in betas
|
||||
assert ANTHROPIC_TOKEN_COUNTING_BETA_VERSION in betas
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.usefixtures("httpx_transport_clients")
|
||||
@pytest.mark.parametrize(
|
||||
("context_management", "forwarded", "sent_beta", "unsent_beta"),
|
||||
(
|
||||
pytest.param(
|
||||
{"edits": [{"type": "compact_20260112"}]},
|
||||
{"edits": [{"type": "compact_20260112"}]},
|
||||
ANTHROPIC_BETA_HEADER_VALUES.COMPACT_2026_01_12.value,
|
||||
ANTHROPIC_BETA_HEADER_VALUES.CONTEXT_MANAGEMENT_2025_06_27.value,
|
||||
id="compact",
|
||||
),
|
||||
pytest.param(
|
||||
{"edits": [{"type": "clear_tool_uses_20250919"}]},
|
||||
{"edits": [{"type": "clear_tool_uses_20250919"}]},
|
||||
ANTHROPIC_BETA_HEADER_VALUES.CONTEXT_MANAGEMENT_2025_06_27.value,
|
||||
ANTHROPIC_BETA_HEADER_VALUES.COMPACT_2026_01_12.value,
|
||||
id="clearing",
|
||||
),
|
||||
pytest.param(
|
||||
[{"type": "compaction", "compact_threshold": 150000}],
|
||||
{"edits": [{"type": "compact_20260112", "trigger": {"type": "input_tokens", "value": 150000}}]},
|
||||
ANTHROPIC_BETA_HEADER_VALUES.COMPACT_2026_01_12.value,
|
||||
ANTHROPIC_BETA_HEADER_VALUES.CONTEXT_MANAGEMENT_2025_06_27.value,
|
||||
id="openai_list_form",
|
||||
),
|
||||
),
|
||||
)
|
||||
async def test_counter_forwards_context_management_with_the_beta_inference_sends(
|
||||
context_management: JsonValue, forwarded: JsonValue, sent_beta: str, unsent_beta: str
|
||||
) -> None:
|
||||
"""Anthropic counts an unsigned compaction block only under the compact_20260112 edit that inference sends,
|
||||
and a clearing edit changes the count, so the count carries the request's context_management as inference
|
||||
sends it, the OpenAI list form mapped to edits, and its beta."""
|
||||
with respx.mock:
|
||||
route: Final = respx.post("https://gateway.example/v1/messages/count_tokens").mock(
|
||||
return_value=httpx.Response(200, json={"input_tokens": 259})
|
||||
)
|
||||
result: Final = await AnthropicTokenCounter().count_tokens(
|
||||
model_to_use="claude-opus-5-5",
|
||||
messages=[{"role": "assistant", "content": [{"type": "compaction", "content": "Recipe app summary."}]}],
|
||||
contents=None,
|
||||
deployment={"litellm_params": {"api_key": "sk-ant-api03-test-key", "api_base": "https://gateway.example"}},
|
||||
context_management=CountContextManagement(spec=context_management, drop_invalid_edits=False),
|
||||
)
|
||||
|
||||
sent: Final = route.calls.last.request
|
||||
betas: Final = sent.headers["anthropic-beta"].split(",")
|
||||
assert result is not None and result.total_tokens == 259
|
||||
assert TypeAdapter(dict[str, JsonValue]).validate_json(sent.content)["context_management"] == forwarded
|
||||
assert (sent_beta in betas, unsent_beta in betas) == (True, False)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.usefixtures("httpx_transport_clients")
|
||||
@pytest.mark.parametrize("context_management", ({}, {"edits": []}, []), ids=("empty_object", "empty_edits", "empty_list"))
|
||||
async def test_counter_leaves_out_context_management_without_edits(context_management: JsonValue) -> None:
|
||||
"""Anthropic's count endpoint rejects a context_management with no edits, which inference accepts, so the count
|
||||
leaves it out and adds no context-management beta."""
|
||||
with respx.mock:
|
||||
route: Final = respx.post("https://gateway.example/v1/messages/count_tokens").mock(
|
||||
return_value=httpx.Response(200, json={"input_tokens": 11})
|
||||
)
|
||||
result: Final = await AnthropicTokenCounter().count_tokens(
|
||||
model_to_use="claude-opus-5-5",
|
||||
messages=[{"role": "user", "content": "Hello"}],
|
||||
contents=None,
|
||||
deployment={"litellm_params": {"api_key": "sk-ant-api03-test-key", "api_base": "https://gateway.example"}},
|
||||
context_management=CountContextManagement(spec=context_management, drop_invalid_edits=False),
|
||||
)
|
||||
|
||||
sent: Final = route.calls.last.request
|
||||
betas: Final = frozenset(sent.headers.get("anthropic-beta", "").split(","))
|
||||
assert result is not None and result.total_tokens == 11
|
||||
assert "context_management" not in TypeAdapter(dict[str, JsonValue]).validate_json(sent.content)
|
||||
assert not betas & {
|
||||
ANTHROPIC_BETA_HEADER_VALUES.COMPACT_2026_01_12.value,
|
||||
ANTHROPIC_BETA_HEADER_VALUES.CONTEXT_MANAGEMENT_2025_06_27.value,
|
||||
}
|
||||
|
|
|
|||
|
|
@ -48,7 +48,7 @@ FEDERATED_DEPLOYMENT = {
|
|||
def count_tokens_headers_for(api_key: str) -> dict[str, str]:
|
||||
auth_header = AnthropicModelInfo.get_auth_header(api_key=api_key)
|
||||
assert auth_header is not None
|
||||
return AnthropicCountTokensConfig().get_count_tokens_headers(auth_header)
|
||||
return AnthropicCountTokensConfig().get_count_tokens_headers(auth_header, [{"role": "user", "content": "hi"}])
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
|
|
|
|||
|
|
@ -1,3 +1,4 @@
|
|||
import datetime
|
||||
import json
|
||||
from collections.abc import Awaitable, Callable, Mapping
|
||||
from typing import Final
|
||||
|
|
@ -12,6 +13,7 @@ import respx
|
|||
import litellm
|
||||
from litellm.exceptions import BadRequestError, ContentPolicyViolationError, ContextWindowExceededError
|
||||
from litellm.litellm_core_utils.exception_mapping_utils import exception_type
|
||||
from litellm.llms.azure.common_utils import AzureOpenAIError
|
||||
|
||||
|
||||
@respx.mock
|
||||
|
|
@ -632,3 +634,202 @@ def test_azure_content_policy_wins_over_context_length_code(monkeypatch: pytest.
|
|||
error: Final = _azure_chat_error(monkeypatch, _azure_error_response(400, policy_error))
|
||||
|
||||
assert type(error) is ContentPolicyViolationError
|
||||
|
||||
|
||||
_AZURE_REFUSAL_BODY: Final = '{"error":{"message":"azure refused","type":"upstream_refusal","code":"refused"}}'
|
||||
_AZURE_DEPLOYMENT_NOT_FOUND_BODY: Final = '{"error":{"code":"DeploymentNotFound","message":"no such deployment"}}'
|
||||
_AZURE_INTERNAL_ERROR_BODY: Final = '{"error":{"code":"InternalServerError","message":"Internal server error"}}'
|
||||
_AZURE_CONTENT_FILTER_BODY: Final = json.dumps(
|
||||
{
|
||||
"error": {
|
||||
"message": "The response was filtered due to the prompt triggering Azure OpenAI's content management policy.",
|
||||
"type": None,
|
||||
"param": "prompt",
|
||||
"code": "content_filter",
|
||||
"status": 400,
|
||||
"innererror": {
|
||||
"code": "ResponsibleAIPolicyViolation",
|
||||
"content_filter_result": {"violence": {"filtered": True, "severity": "high"}},
|
||||
},
|
||||
}
|
||||
}
|
||||
)
|
||||
_AZURE_KEY: Final = "azure-provider-secret"
|
||||
_AZURE_BASE: Final = "https://azure-upstream.openai.azure.com"
|
||||
_AZURE_ARGS: Final = {"api_base": _AZURE_BASE, "api_version": "2024-10-21", "api_key": _AZURE_KEY, "max_retries": 0}
|
||||
|
||||
|
||||
async def _azure_completion() -> None:
|
||||
litellm.completion(model="azure/gpt-4.1-mini", messages=[{"role": "user", "content": "hi"}], **_AZURE_ARGS)
|
||||
|
||||
|
||||
async def _azure_completion_stream() -> None:
|
||||
list(
|
||||
litellm.completion(
|
||||
model="azure/gpt-4.1-mini", messages=[{"role": "user", "content": "hi"}], stream=True, **_AZURE_ARGS
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
async def _azure_acompletion() -> None:
|
||||
await litellm.acompletion(model="azure/gpt-4.1-mini", messages=[{"role": "user", "content": "hi"}], **_AZURE_ARGS)
|
||||
|
||||
|
||||
async def _azure_acompletion_stream() -> None:
|
||||
stream: Final = await litellm.acompletion(
|
||||
model="azure/gpt-4.1-mini", messages=[{"role": "user", "content": "hi"}], stream=True, **_AZURE_ARGS
|
||||
)
|
||||
async for _ in stream:
|
||||
pass
|
||||
|
||||
|
||||
async def _azure_embedding() -> None:
|
||||
litellm.embedding(model="azure/text-embedding-3-large", input=["hi"], **_AZURE_ARGS)
|
||||
|
||||
|
||||
async def _azure_aembedding() -> None:
|
||||
await litellm.aembedding(model="azure/text-embedding-3-large", input=["hi"], **_AZURE_ARGS)
|
||||
|
||||
|
||||
async def _azure_image_generation() -> None:
|
||||
litellm.image_generation(model="azure/gpt-image-2", prompt="a cat", **_AZURE_ARGS)
|
||||
|
||||
|
||||
async def _azure_text_completion() -> None:
|
||||
litellm.text_completion(model="azure_text/gpt-35-turbo-instruct", prompt="hi", **_AZURE_ARGS)
|
||||
|
||||
|
||||
async def _azure_text_completion_stream() -> None:
|
||||
list(litellm.text_completion(model="azure_text/gpt-35-turbo-instruct", prompt="hi", stream=True, **_AZURE_ARGS))
|
||||
|
||||
|
||||
async def _azure_atext_completion() -> None:
|
||||
await litellm.atext_completion(model="azure_text/gpt-35-turbo-instruct", prompt="hi", **_AZURE_ARGS)
|
||||
|
||||
|
||||
async def _azure_atext_completion_stream() -> None:
|
||||
stream: Final = await litellm.atext_completion(
|
||||
model="azure_text/gpt-35-turbo-instruct", prompt="hi", stream=True, **_AZURE_ARGS
|
||||
)
|
||||
async for _ in stream:
|
||||
pass
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"route",
|
||||
(
|
||||
_azure_completion,
|
||||
_azure_completion_stream,
|
||||
_azure_acompletion,
|
||||
_azure_acompletion_stream,
|
||||
_azure_embedding,
|
||||
_azure_aembedding,
|
||||
_azure_image_generation,
|
||||
_azure_text_completion,
|
||||
_azure_text_completion_stream,
|
||||
_azure_atext_completion,
|
||||
_azure_atext_completion_stream,
|
||||
),
|
||||
ids=lambda route: route.__name__.removeprefix("_azure_"),
|
||||
)
|
||||
@pytest.mark.parametrize(
|
||||
("status_code", "body", "mapped_error"),
|
||||
(
|
||||
(400, _AZURE_REFUSAL_BODY, BadRequestError),
|
||||
(400, _AZURE_CONTENT_FILTER_BODY, BadRequestError),
|
||||
(403, _AZURE_REFUSAL_BODY, litellm.APIError),
|
||||
(404, _AZURE_DEPLOYMENT_NOT_FOUND_BODY, litellm.NotFoundError),
|
||||
(500, _AZURE_INTERNAL_ERROR_BODY, openai.APIError),
|
||||
(502, _AZURE_REFUSAL_BODY, litellm.BadGatewayError),
|
||||
(504, _AZURE_REFUSAL_BODY, litellm.Timeout),
|
||||
),
|
||||
)
|
||||
@respx.mock
|
||||
async def test_azure_error_keeps_the_upstream_body_without_the_api_key(
|
||||
route: Callable[[], Awaitable[None]],
|
||||
status_code: int,
|
||||
body: str,
|
||||
mapped_error: type[Exception],
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
|
||||
respx.post(url__regex=r"https://azure-upstream\.openai\.azure\.com/.*").mock(
|
||||
return_value=httpx.Response(status_code, content=body, headers={"content-type": "application/json"})
|
||||
)
|
||||
|
||||
with pytest.raises(mapped_error) as raised:
|
||||
await route()
|
||||
|
||||
error_response: Final = getattr(raised.value, "response")
|
||||
assert error_response.status_code == status_code
|
||||
assert error_response.text == body
|
||||
assert error_response.json() == json.loads(body)
|
||||
for exposed in (error_response.request, getattr(raised.value, "request")):
|
||||
assert _AZURE_KEY not in str(exposed.headers)
|
||||
assert exposed.content == b""
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"route",
|
||||
(
|
||||
_azure_completion,
|
||||
_azure_completion_stream,
|
||||
_azure_acompletion,
|
||||
_azure_acompletion_stream,
|
||||
_azure_embedding,
|
||||
_azure_aembedding,
|
||||
_azure_text_completion,
|
||||
_azure_atext_completion,
|
||||
_azure_atext_completion_stream,
|
||||
),
|
||||
ids=lambda route: route.__name__.removeprefix("_azure_"),
|
||||
)
|
||||
@respx.mock
|
||||
async def test_azure_error_response_keeps_the_upstream_elapsed_time(
|
||||
route: Callable[[], Awaitable[None]], monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
|
||||
respx.post(url__regex=r"https://azure-upstream\.openai\.azure\.com/.*").mock(
|
||||
return_value=httpx.Response(400, content=_AZURE_REFUSAL_BODY, headers={"content-type": "application/json"})
|
||||
)
|
||||
|
||||
with pytest.raises(BadRequestError) as raised:
|
||||
await route()
|
||||
|
||||
assert raised.value.response.elapsed > datetime.timedelta(0)
|
||||
|
||||
|
||||
def test_azure_client_side_missing_key_error_keeps_the_empty_placeholder_response() -> None:
|
||||
missing_key: Final = openai.OpenAIError(
|
||||
"The api_key client option must be set either by passing api_key to the client or by setting the "
|
||||
"OPENAI_API_KEY environment variable"
|
||||
)
|
||||
|
||||
with pytest.raises(litellm.AuthenticationError) as raised:
|
||||
exception_type(model="azure/gpt-4.1-mini", original_exception=missing_key, custom_llm_provider="azure")
|
||||
|
||||
assert raised.value.response.status_code == 401
|
||||
assert raised.value.response.text == ""
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("status_code", "message", "mapped_error"),
|
||||
(
|
||||
(500, "azure_client is not an instance of AzureOpenAI or OpenAI", litellm.APIError),
|
||||
(408, "the deployment did not answer in time", litellm.Timeout),
|
||||
(504, "the deployment did not answer in time", litellm.Timeout),
|
||||
),
|
||||
)
|
||||
def test_azure_error_litellm_raises_itself_has_no_upstream_response(
|
||||
status_code: int, message: str, mapped_error: type[Exception]
|
||||
) -> None:
|
||||
with pytest.raises(mapped_error) as raised:
|
||||
exception_type(
|
||||
model="azure/gpt-4.1-mini",
|
||||
original_exception=AzureOpenAIError(status_code=status_code, message=message),
|
||||
custom_llm_provider="azure",
|
||||
)
|
||||
|
||||
assert getattr(raised.value, "response", None) is None
|
||||
|
|
|
|||
|
|
@ -1,6 +1,7 @@
|
|||
import json
|
||||
import traceback
|
||||
from collections.abc import Callable
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
|
|
@ -9,6 +10,7 @@ import respx
|
|||
from httpx import Response
|
||||
|
||||
import litellm
|
||||
from litellm import Router
|
||||
from litellm.llms.custom_httpx.http_handler import HTTPHandler
|
||||
from litellm.litellm_core_utils.get_model_cost_map import get_model_cost_map
|
||||
from litellm.llms.azure_ai.azure_model_router.transformation import (
|
||||
|
|
@ -885,6 +887,360 @@ def test_azure_ai_stripping_does_not_mutate_caller_messages():
|
|||
}
|
||||
|
||||
|
||||
def test_azure_ai_keeps_prompt_cache_breakpoint_for_a_model_that_supports_it(_local_model_cost_map):
|
||||
"""
|
||||
Foundry GPT-6 deployments take OpenAI explicit prompt caching, and the breakpoint only
|
||||
counts on a content part, so a system message carrying one must reach the request body as
|
||||
the content-part list it arrived in while the Anthropic-only fields are still stripped.
|
||||
"""
|
||||
request: Final = AzureAIStudioConfig().transform_request(
|
||||
model="gpt-6-astra",
|
||||
messages=[
|
||||
{
|
||||
"role": "system",
|
||||
"content": [
|
||||
{
|
||||
"type": "text",
|
||||
"text": "stable instructions",
|
||||
"prompt_cache_breakpoint": {"mode": "explicit"},
|
||||
"cache_control": {"type": "ephemeral"},
|
||||
}
|
||||
],
|
||||
},
|
||||
{"role": "user", "content": [{"type": "text", "text": "hi"}]},
|
||||
],
|
||||
optional_params={"prompt_cache_options": {"mode": "explicit"}},
|
||||
litellm_params={},
|
||||
headers={},
|
||||
)
|
||||
|
||||
assert request["messages"][0]["content"] == [
|
||||
{"type": "text", "text": "stable instructions", "prompt_cache_breakpoint": {"mode": "explicit"}}
|
||||
]
|
||||
assert request["messages"][1]["content"] == [{"type": "text", "text": "hi"}]
|
||||
assert request["prompt_cache_options"] == {"mode": "explicit"}
|
||||
|
||||
|
||||
def test_azure_ai_still_flattens_list_content_for_a_model_without_the_breakpoint_flag(_local_model_cost_map):
|
||||
request: Final = AzureAIStudioConfig().transform_request(
|
||||
model="gpt-4o",
|
||||
messages=[
|
||||
{
|
||||
"role": "system",
|
||||
"content": [
|
||||
{"type": "text", "text": "stable ", "prompt_cache_breakpoint": {"mode": "explicit"}},
|
||||
{"type": "text", "text": "instructions"},
|
||||
],
|
||||
},
|
||||
{"role": "user", "content": [{"type": "text", "text": "hi"}]},
|
||||
],
|
||||
optional_params={},
|
||||
litellm_params={},
|
||||
headers={},
|
||||
)
|
||||
|
||||
assert request["messages"][0]["content"] == "stable instructions"
|
||||
assert request["messages"][1]["content"] == "hi"
|
||||
assert not _find_key_anywhere(request["messages"], "prompt_cache_breakpoint")
|
||||
|
||||
|
||||
def test_azure_ai_drops_thinking_parts_from_history_while_keeping_the_breakpoint(_local_model_cost_map):
|
||||
"""
|
||||
A multi-turn client can echo an assistant turn back with thinking parts inside its content
|
||||
list. Foundry rejects those parts, and the string conversion always dropped them, so a model
|
||||
that keeps content parts must still drop them while the text parts and the breakpoint stay.
|
||||
"""
|
||||
request: Final = AzureAIStudioConfig().transform_request(
|
||||
model="gpt-6-astra",
|
||||
messages=[
|
||||
{
|
||||
"role": "system",
|
||||
"content": [
|
||||
{"type": "text", "text": "stable instructions", "prompt_cache_breakpoint": {"mode": "explicit"}}
|
||||
],
|
||||
},
|
||||
{"role": "user", "content": "read the file"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{"type": "thinking", "thinking": "The user wants me to read a file.", "signature": "sig"},
|
||||
{"type": "redacted_thinking", "data": "opaque"},
|
||||
{"type": "text", "text": "Reading it now."},
|
||||
],
|
||||
},
|
||||
{"role": "user", "content": [{"type": "text", "text": "thanks"}]},
|
||||
],
|
||||
optional_params={"prompt_cache_options": {"mode": "explicit"}},
|
||||
litellm_params={},
|
||||
headers={},
|
||||
)
|
||||
|
||||
assert request["messages"][0]["content"] == [
|
||||
{"type": "text", "text": "stable instructions", "prompt_cache_breakpoint": {"mode": "explicit"}}
|
||||
]
|
||||
assert request["messages"][2]["content"] == [{"type": "text", "text": "Reading it now."}]
|
||||
assert request["messages"][3]["content"] == [{"type": "text", "text": "thanks"}]
|
||||
assert not _find_key_anywhere(request["messages"], "thinking")
|
||||
assert not _find_key_anywhere(request["messages"], "data")
|
||||
|
||||
|
||||
|
||||
def test_azure_ai_renames_a_cache_control_mark_to_a_breakpoint_for_a_model_that_reads_them(_local_model_cost_map):
|
||||
request: Final = AzureAIStudioConfig().transform_request(
|
||||
model="gpt-6-astra",
|
||||
messages=[
|
||||
{
|
||||
"role": "system",
|
||||
"content": [{"type": "text", "text": "stable instructions", "cache_control": {"type": "ephemeral"}}],
|
||||
},
|
||||
{"role": "user", "content": "hi", "cache_control": {"type": "ephemeral"}},
|
||||
],
|
||||
optional_params={"prompt_cache_options": {"mode": "explicit"}},
|
||||
litellm_params={},
|
||||
headers={},
|
||||
)
|
||||
|
||||
assert request["messages"][0]["content"] == [
|
||||
{"type": "text", "text": "stable instructions", "prompt_cache_breakpoint": {"mode": "explicit"}}
|
||||
]
|
||||
assert request["messages"][1]["content"] == [
|
||||
{"type": "text", "text": "hi", "prompt_cache_breakpoint": {"mode": "explicit"}}
|
||||
]
|
||||
assert not _find_key_anywhere(request["messages"], "cache_control")
|
||||
|
||||
|
||||
def test_azure_ai_drops_a_file_part_next_to_the_breakpoint_like_the_string_conversion(_local_model_cost_map):
|
||||
messages: Final = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "summarize", "prompt_cache_breakpoint": {"mode": "explicit"}},
|
||||
{"type": "file", "file": {"file_id": "file-held-by-another-provider"}},
|
||||
],
|
||||
}
|
||||
]
|
||||
|
||||
kept: Final = AzureAIStudioConfig().transform_request(
|
||||
model="gpt-6-astra", messages=messages, optional_params={}, litellm_params={}, headers={}
|
||||
)
|
||||
flattened: Final = AzureAIStudioConfig().transform_request(
|
||||
model="gpt-4o", messages=messages, optional_params={}, litellm_params={}, headers={}
|
||||
)
|
||||
|
||||
assert kept["messages"][0]["content"] == [
|
||||
{"type": "text", "text": "summarize", "prompt_cache_breakpoint": {"mode": "explicit"}}
|
||||
]
|
||||
assert flattened["messages"][0]["content"] == "summarize"
|
||||
|
||||
|
||||
def test_azure_ai_sends_a_pdf_part_next_to_the_breakpoint_as_is(_local_model_cost_map):
|
||||
content: Final = [
|
||||
{"type": "file", "file": {"filename": "document.pdf", "file_data": "data:application/pdf;base64,JVBERi0xLjQK"}},
|
||||
{"type": "text", "text": "What is the codeword?", "prompt_cache_breakpoint": {"mode": "explicit"}},
|
||||
]
|
||||
|
||||
request: Final = AzureAIStudioConfig().transform_request(
|
||||
model="gpt-6-astra",
|
||||
messages=[{"role": "user", "content": content}],
|
||||
optional_params={},
|
||||
litellm_params={},
|
||||
headers={},
|
||||
)
|
||||
|
||||
assert request["messages"] == [{"role": "user", "content": content}]
|
||||
|
||||
|
||||
def test_azure_ai_puts_a_message_level_mark_on_the_last_text_part_when_a_file_part_follows(_local_model_cost_map):
|
||||
request: Final = AzureAIStudioConfig().transform_request(
|
||||
model="gpt-6-astra",
|
||||
messages=[
|
||||
{
|
||||
"role": "user",
|
||||
"cache_control": {"type": "ephemeral"},
|
||||
"content": [
|
||||
{"type": "text", "text": "summarize"},
|
||||
{"type": "file", "file": {"file_id": "file-held-by-another-provider"}},
|
||||
],
|
||||
}
|
||||
],
|
||||
optional_params={"prompt_cache_options": {"mode": "explicit"}},
|
||||
litellm_params={},
|
||||
headers={},
|
||||
)
|
||||
|
||||
assert request["messages"][0]["content"] == [
|
||||
{"type": "text", "text": "summarize", "prompt_cache_breakpoint": {"mode": "explicit"}}
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("registered", "keeps_parts"),
|
||||
[
|
||||
({"azure_ai/qa-foundry-alias": {"litellm_provider": "azure_ai", "supports_prompt_cache_breakpoint": True}}, True),
|
||||
({"qa-foundry-alias": {"litellm_provider": "openai", "supports_prompt_cache_breakpoint": True}}, False),
|
||||
],
|
||||
ids=["azure_ai-entry-opts-in", "openai-entry-does-not-speak-for-azure_ai"],
|
||||
)
|
||||
def test_azure_ai_keeps_parts_only_when_the_azure_ai_entry_reads_breakpoints(
|
||||
_local_model_cost_map, registered: dict[str, dict[str, object]], keeps_parts: bool
|
||||
):
|
||||
litellm.register_model(model_cost=registered, persist_across_reloads=False, add_to_provider_model_sets=False)
|
||||
breakpoint_part: Final = {"type": "text", "text": "stable", "prompt_cache_breakpoint": {"mode": "explicit"}}
|
||||
|
||||
request: Final = AzureAIStudioConfig().transform_request(
|
||||
model="qa-foundry-alias",
|
||||
messages=[{"role": "system", "content": [breakpoint_part]}],
|
||||
optional_params={},
|
||||
litellm_params={"custom_llm_provider": "azure_ai", "drop_params": True},
|
||||
headers={},
|
||||
)
|
||||
|
||||
assert request["messages"][0]["content"] == ([breakpoint_part] if keeps_parts else "stable")
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"registration_order", [("opted-in", "opted-out", "unflagged"), ("unflagged", "opted-out", "opted-in")]
|
||||
)
|
||||
def test_azure_ai_deployment_reads_only_its_own_breakpoint_flag_not_a_siblings_sharing_its_model_string(
|
||||
_local_model_cost_map, respx_mock: respx.MockRouter, registration_order: tuple[str, ...]
|
||||
) -> None:
|
||||
breakpoint_flags: Final = MappingProxyType(
|
||||
{
|
||||
"opted-in": {"supports_prompt_cache_breakpoint": True},
|
||||
"opted-out": {"supports_prompt_cache_breakpoint": False},
|
||||
}
|
||||
)
|
||||
def deployment(host: str, model_info: dict[str, object]) -> dict[str, object]:
|
||||
return {
|
||||
"model_name": host,
|
||||
"litellm_params": {
|
||||
"model": "azure_ai/qa-shared-deployment",
|
||||
"api_base": f"https://{host}.services.ai.azure.com/models",
|
||||
"api_key": "azure-ai-test-key",
|
||||
},
|
||||
"model_info": model_info,
|
||||
}
|
||||
|
||||
routes: Final = {
|
||||
host: respx_mock.post(host=f"{host}.services.ai.azure.com", path="/models/chat/completions").mock(
|
||||
return_value=Response(
|
||||
200,
|
||||
json={
|
||||
"id": "chatcmpl-azure-ai",
|
||||
"object": "chat.completion",
|
||||
"created": 1700000000,
|
||||
"model": "qa-shared-deployment",
|
||||
"choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}],
|
||||
"usage": {"prompt_tokens": 5, "completion_tokens": 1, "total_tokens": 6},
|
||||
},
|
||||
)
|
||||
)
|
||||
for host in registration_order
|
||||
}
|
||||
router: Final = Router(
|
||||
model_list=[
|
||||
deployment(host, {"id": f"{host}-id", **breakpoint_flags.get(host, {})}) for host in registration_order
|
||||
]
|
||||
)
|
||||
marked_system: Final = {
|
||||
"role": "system",
|
||||
"content": [{"type": "text", "text": "stable", "cache_control": {"type": "ephemeral"}}],
|
||||
}
|
||||
|
||||
for host in registration_order:
|
||||
router.completion(model=host, messages=[marked_system, {"role": "user", "content": "hi"}])
|
||||
|
||||
def sent_system_content(host: str) -> object:
|
||||
return json.loads(routes[host].calls.last.request.content)["messages"][0]["content"]
|
||||
|
||||
assert sent_system_content("opted-in") == [
|
||||
{"type": "text", "text": "stable", "prompt_cache_breakpoint": {"mode": "explicit"}}
|
||||
]
|
||||
assert sent_system_content("opted-out") == "stable"
|
||||
assert sent_system_content("unflagged") == "stable"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"litellm_params",
|
||||
[
|
||||
{"base_model": "azure_ai/gpt-6-astra"},
|
||||
{"base_model": "azure_ai/gpt-6-astra", "model_info": {"id": "my-astra-deployment-id"}},
|
||||
],
|
||||
)
|
||||
def test_azure_ai_keeps_parts_for_a_custom_deployment_name_whose_base_model_reads_breakpoints(
|
||||
_local_model_cost_map, litellm_params: dict[str, object]
|
||||
):
|
||||
breakpoint_part: Final = {"type": "text", "text": "stable", "prompt_cache_breakpoint": {"mode": "explicit"}}
|
||||
|
||||
request: Final = AzureAIStudioConfig().transform_request(
|
||||
model="my-astra-deployment",
|
||||
messages=[{"role": "system", "content": [breakpoint_part]}],
|
||||
optional_params={},
|
||||
litellm_params=litellm_params,
|
||||
headers={},
|
||||
)
|
||||
|
||||
assert request["messages"][0]["content"] == [breakpoint_part]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"litellm_params",
|
||||
[
|
||||
{"base_model": "azure/cost-tracking-row"},
|
||||
{"base_model": "azure/cost-tracking-row", "model_info": {"id": "astra-deployment-id"}},
|
||||
],
|
||||
)
|
||||
def test_azure_ai_falls_back_to_the_model_name_when_the_base_model_row_does_not_read_breakpoints(
|
||||
_local_model_cost_map, litellm_params: dict[str, object]
|
||||
):
|
||||
breakpoint_part: Final = {"type": "text", "text": "stable", "prompt_cache_breakpoint": {"mode": "explicit"}}
|
||||
|
||||
request: Final = AzureAIStudioConfig().transform_request(
|
||||
model="gpt-6-astra",
|
||||
messages=[{"role": "system", "content": [breakpoint_part]}],
|
||||
optional_params={},
|
||||
litellm_params=litellm_params,
|
||||
headers={},
|
||||
)
|
||||
|
||||
assert request["messages"][0]["content"] == [breakpoint_part]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"sent_marker, litellm_params, wire_part",
|
||||
[
|
||||
(
|
||||
{"mode": "explicit", "ttl": "30m"},
|
||||
{"drop_params": True},
|
||||
{"type": "text", "text": "stable", "prompt_cache_breakpoint": {"mode": "explicit"}},
|
||||
),
|
||||
({"type": "ephemeral"}, {"drop_params": True}, {"type": "text", "text": "stable"}),
|
||||
(
|
||||
{"mode": "explicit", "ttl": "30m"},
|
||||
{},
|
||||
{"type": "text", "text": "stable", "prompt_cache_breakpoint": {"mode": "explicit", "ttl": "30m"}},
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_azure_ai_reshapes_breakpoints_for_foundry_only_under_drop_params(
|
||||
_local_model_cost_map,
|
||||
sent_marker: dict[str, object],
|
||||
litellm_params: dict[str, object],
|
||||
wire_part: dict[str, object],
|
||||
):
|
||||
request: Final = AzureAIStudioConfig().transform_request(
|
||||
model="gpt-6-astra",
|
||||
messages=[
|
||||
{"role": "system", "content": [{"type": "text", "text": "stable", "prompt_cache_breakpoint": sent_marker}]}
|
||||
],
|
||||
optional_params={},
|
||||
litellm_params=litellm_params,
|
||||
headers={},
|
||||
)
|
||||
|
||||
assert request["messages"][0]["content"] == [wire_part]
|
||||
|
||||
|
||||
def test_azure_ai_reference_only_tool_result_is_sent_as_empty_text():
|
||||
messages: Final = [
|
||||
{"role": "user", "content": "load the WebFetch tool"},
|
||||
|
|
|
|||
|
|
@ -197,6 +197,29 @@ def test_chat_does_not_send_the_openai_organization_to_bigmodel(monkeypatch: pyt
|
|||
assert "openai-organization" not in route.calls.last.request.headers
|
||||
|
||||
|
||||
@respx.mock
|
||||
def test_chat_sends_json_schema_response_format_as_given() -> None:
|
||||
route: Final = respx.post(f"{BASE}/api/paas/v4/chat/completions").respond(
|
||||
json={
|
||||
"id": "chat_test",
|
||||
"object": "chat.completion",
|
||||
"created": 1,
|
||||
"model": "glm-4.6",
|
||||
"choices": [{"index": 0, "message": {"role": "assistant", "content": "{}"}, "finish_reason": "stop"}],
|
||||
"usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2},
|
||||
}
|
||||
)
|
||||
messages: Final = [{"role": "user", "content": "Ann likes 7"}]
|
||||
response_format: Final = {
|
||||
"type": "json_schema",
|
||||
"json_schema": {"name": "profile", "schema": {"type": "object", "properties": {"n": {"type": "integer"}}}},
|
||||
}
|
||||
_ = litellm.completion(model="bigmodel/glm-4.6", messages=messages, response_format=response_format)
|
||||
body: Final = json.loads(route.calls.last.request.content)
|
||||
assert body["response_format"] == response_format
|
||||
assert body["messages"] == messages
|
||||
|
||||
|
||||
@pytest.mark.parametrize("is_async", (False, True))
|
||||
@pytest.mark.asyncio
|
||||
@respx.mock
|
||||
|
|
|
|||
|
|
@ -14,6 +14,7 @@ import httpcore
|
|||
import httpx
|
||||
import pytest
|
||||
import respx
|
||||
import websockets.exceptions
|
||||
from botocore.credentials import RefreshableCredentials
|
||||
from h2.events import ConnectionTerminated
|
||||
|
||||
|
|
@ -2787,6 +2788,194 @@ async def test_async_realtime_error_event_send_failure_still_closes():
|
|||
assert client_ws.events == [("close", (1011, "Internal server error: vertex token refresh exploded"))]
|
||||
|
||||
|
||||
_GEMINI_LIVE_SETUP_COMPLETE = json.dumps({"setupComplete": {}})
|
||||
_GEMINI_LIVE_TEXT_TURN = json.dumps({"serverContent": {"modelTurn": {"parts": [{"text": "Blue."}]}}})
|
||||
_GEMINI_LIVE_TURN_COMPLETE = json.dumps({"serverContent": {"turnComplete": True}})
|
||||
|
||||
|
||||
class _ScriptedRealtimeBackend:
|
||||
"""Backend socket that answers the first setup frame with the scripted replies
|
||||
and closes once the client is done, so both bridge loops run to completion."""
|
||||
|
||||
def __init__(self, replies_after_setup):
|
||||
self.sent = []
|
||||
self.closed = asyncio.Event()
|
||||
self._replies_after_setup = list(replies_after_setup)
|
||||
self._inbox = asyncio.Queue()
|
||||
|
||||
async def __aenter__(self):
|
||||
return self
|
||||
|
||||
async def __aexit__(self, exc_type, exc, tb):
|
||||
return None
|
||||
|
||||
async def send(self, payload):
|
||||
self.sent.append(payload)
|
||||
if "setup" in json.loads(payload):
|
||||
for reply in self._replies_after_setup:
|
||||
self._inbox.put_nowait(reply)
|
||||
|
||||
def close_after_replies(self):
|
||||
self._inbox.put_nowait(None)
|
||||
|
||||
async def recv(self, decode=True):
|
||||
reply = await self._inbox.get()
|
||||
if reply is None:
|
||||
self.closed.set()
|
||||
raise websockets.exceptions.ConnectionClosedOK(None, None)
|
||||
return reply
|
||||
|
||||
|
||||
class _ScriptedRealtimeClient:
|
||||
def __init__(self, messages, backend):
|
||||
self.scope = {"headers": []}
|
||||
self.exceptions = websockets.exceptions
|
||||
self.sent = []
|
||||
self._messages = list(messages)
|
||||
self._backend = backend
|
||||
|
||||
async def send_text(self, payload):
|
||||
self.sent.append(payload)
|
||||
|
||||
async def receive_text(self):
|
||||
if self._messages:
|
||||
return self._messages.pop(0)
|
||||
self._backend.close_after_replies()
|
||||
await self._backend.closed.wait()
|
||||
raise websockets.exceptions.ConnectionClosedOK(None, None)
|
||||
|
||||
async def close(self, code=None, reason=None):
|
||||
return None
|
||||
|
||||
|
||||
def _realtime_logging_obj():
|
||||
logging_obj = Mock()
|
||||
logging_obj.litellm_trace_id = "trace-lit-6974"
|
||||
logging_obj.async_success_handler = AsyncMock()
|
||||
logging_obj.dispatch_failure_handlers = AsyncMock()
|
||||
return logging_obj
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"model, session, backend_modality, client_modality, backend_turn, drop_params",
|
||||
[
|
||||
(
|
||||
"gemini-live-2.5-flash",
|
||||
{"type": "realtime", "output_modalities": ["text"]},
|
||||
"TEXT",
|
||||
"text",
|
||||
[_GEMINI_LIVE_TEXT_TURN],
|
||||
False,
|
||||
),
|
||||
("gemini-live-2.5-flash", {"modalities": ["text"]}, "TEXT", "text", [_GEMINI_LIVE_TEXT_TURN], False),
|
||||
(
|
||||
"gemini-live-2.5-flash-native-audio",
|
||||
{"type": "realtime", "output_modalities": ["text"]},
|
||||
"AUDIO",
|
||||
"audio",
|
||||
[],
|
||||
True,
|
||||
),
|
||||
],
|
||||
)
|
||||
async def test_async_realtime_vertex_first_session_update_builds_the_only_setup(
|
||||
model, session, backend_modality, client_modality, backend_turn, drop_params: bool, monkeypatch
|
||||
):
|
||||
"""Regression for text-only Vertex Live sessions answered in audio: the
|
||||
connect-time default setup (AUDIO) used to go out before the client's
|
||||
session.update, and Gemini accepts only one setup per connection, so the
|
||||
requested modality was dropped and billed at the audio rate. The setup now
|
||||
goes out on the client's first message, built from session.update when that
|
||||
comes first. Under drop_params an audio-only model still gets AUDIO, and the
|
||||
client learns the downgrade from session.updated."""
|
||||
from litellm.llms.vertex_ai.realtime.transformation import VertexAIRealtimeConfig
|
||||
|
||||
monkeypatch.setattr(litellm, "gemini_live_defer_setup", False, raising=False)
|
||||
backend = _ScriptedRealtimeBackend([_GEMINI_LIVE_SETUP_COMPLETE, *backend_turn, _GEMINI_LIVE_TURN_COMPLETE])
|
||||
client_ws = _ScriptedRealtimeClient(
|
||||
[
|
||||
json.dumps({"type": "session.update", "session": session}),
|
||||
json.dumps(
|
||||
{
|
||||
"type": "conversation.item.create",
|
||||
"item": {
|
||||
"type": "message",
|
||||
"role": "user",
|
||||
"content": [{"type": "input_text", "text": "What color is the sky? One word."}],
|
||||
},
|
||||
}
|
||||
),
|
||||
json.dumps({"type": "response.create"}),
|
||||
],
|
||||
backend,
|
||||
)
|
||||
handler = BaseLLMHTTPHandler()
|
||||
|
||||
with patch.object(handler, "_open_realtime_backend_ws", AsyncMock(return_value=backend)):
|
||||
await asyncio.wait_for(
|
||||
handler.async_realtime(
|
||||
model=model,
|
||||
websocket=client_ws,
|
||||
logging_obj=_realtime_logging_obj(),
|
||||
provider_config=VertexAIRealtimeConfig(
|
||||
access_token="fake-token", project="fake-project", location="us-central1", drop_params=drop_params
|
||||
),
|
||||
headers={},
|
||||
),
|
||||
timeout=10,
|
||||
)
|
||||
|
||||
backend_frames = [json.loads(frame) for frame in backend.sent]
|
||||
setups = [frame["setup"] for frame in backend_frames if "setup" in frame]
|
||||
assert len(setups) == 1
|
||||
assert "setup" in backend_frames[0]
|
||||
assert setups[0]["generationConfig"]["responseModalities"] == [backend_modality]
|
||||
assert any("clientContent" in frame for frame in backend_frames[1:])
|
||||
|
||||
client_events = [json.loads(frame) for frame in client_ws.sent]
|
||||
event_types = [event["type"] for event in client_events]
|
||||
assert event_types[:2] == ["session.created", "session.updated"]
|
||||
assert client_events[1]["session"]["modalities"] == [client_modality]
|
||||
if backend_turn:
|
||||
assert "response.output_text.delta" in event_types
|
||||
assert "response.output_audio_transcript.delta" not in event_types
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_realtime_sends_setup_at_connect_when_provider_cannot_build_it_from_session_update():
|
||||
"""Bedrock Nova Sonic style providers keep the connect-time setup and get no
|
||||
synthetic session.created."""
|
||||
setup = json.dumps({"event": {"sessionStart": {}}})
|
||||
backend = _ScriptedRealtimeBackend([])
|
||||
client_ws = _ScriptedRealtimeClient([], backend)
|
||||
provider_config = Mock()
|
||||
provider_config.get_complete_url.return_value = "wss://backend.example/live"
|
||||
provider_config.validate_environment.return_value = {}
|
||||
provider_config.requires_session_configuration.return_value = True
|
||||
provider_config.builds_setup_from_session_update.return_value = False
|
||||
provider_config.session_configuration_request.return_value = setup
|
||||
provider_config.unbilled_usage_on_session_close.return_value = None
|
||||
provider_config.open_backend = AsyncMock(return_value=None)
|
||||
provider_config.wrap_backend.side_effect = lambda opened: opened
|
||||
handler = BaseLLMHTTPHandler()
|
||||
|
||||
with patch.object(handler, "_open_realtime_backend_ws", AsyncMock(return_value=backend)):
|
||||
await asyncio.wait_for(
|
||||
handler.async_realtime(
|
||||
model="amazon.nova-sonic-v1:0",
|
||||
websocket=client_ws,
|
||||
logging_obj=_realtime_logging_obj(),
|
||||
provider_config=provider_config,
|
||||
headers={},
|
||||
),
|
||||
timeout=10,
|
||||
)
|
||||
|
||||
assert backend.sent == [setup]
|
||||
assert client_ws.sent == []
|
||||
|
||||
|
||||
class _JSONBodyAudioTranscriptionConfig(BaseAudioTranscriptionConfig):
|
||||
def get_supported_openai_params(self, model):
|
||||
return []
|
||||
|
|
|
|||
|
|
@ -1,3 +1,4 @@
|
|||
import base64
|
||||
import json
|
||||
from collections.abc import Mapping
|
||||
from typing import Final, cast
|
||||
|
|
@ -7,7 +8,7 @@ import pytest
|
|||
|
||||
|
||||
import litellm
|
||||
from litellm.llms.gemini.realtime.transformation import GeminiRealtimeConfig
|
||||
from litellm.llms.gemini.realtime.transformation import GEMINI_LIVE_DOCUMENTED_VOICES, GeminiRealtimeConfig
|
||||
from litellm.types.llms.gemini import BidiGenerateContentServerMessage
|
||||
|
||||
|
||||
|
|
@ -439,6 +440,51 @@ def test_gemini_realtime_session_update_with_tools():
|
|||
assert "parameters" in function_decl
|
||||
|
||||
|
||||
def _schema_nodes(schema: object):
|
||||
if isinstance(schema, dict):
|
||||
yield schema
|
||||
for value in schema.values():
|
||||
yield from _schema_nodes(value)
|
||||
elif isinstance(schema, list):
|
||||
for value in schema:
|
||||
yield from _schema_nodes(value)
|
||||
|
||||
|
||||
def test_gemini_session_update_converts_flat_realtime_tool_schemas_to_gemini_schemas():
|
||||
session_update: Final = {
|
||||
"type": "session.update",
|
||||
"session": {
|
||||
"tools": [
|
||||
{
|
||||
"type": "function",
|
||||
"name": "lookup_weather",
|
||||
"description": "Look up the weather in a city.",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"$defs": {
|
||||
"City": {"type": "object", "properties": {"name": {"type": "string"}}, "required": ["name"]}
|
||||
},
|
||||
"properties": {"city": {"$ref": "#/$defs/City"}, "note": {"type": ["string", "null"]}},
|
||||
"required": ["city"],
|
||||
},
|
||||
}
|
||||
]
|
||||
},
|
||||
}
|
||||
|
||||
messages: Final = GeminiRealtimeConfig().transform_realtime_request(
|
||||
json.dumps(session_update), "gemini-2.5-flash", session_configuration_request=None
|
||||
)
|
||||
|
||||
declaration: Final = json.loads(messages[0])["setup"]["tools"][0]["function_declarations"][0]
|
||||
assert declaration["name"] == "lookup_weather"
|
||||
parameters: Final = declaration["parameters"]
|
||||
assert parameters["properties"]["city"]["properties"]["name"]["type"] == "string"
|
||||
nodes: Final = tuple(_schema_nodes(parameters))
|
||||
assert not any("$defs" in node or "$ref" in node for node in nodes)
|
||||
assert all(isinstance(node["type"], str) for node in nodes if "type" in node)
|
||||
|
||||
|
||||
def test_gemini_session_update_defaults_to_audio_modality():
|
||||
config = GeminiRealtimeConfig()
|
||||
|
||||
|
|
@ -613,6 +659,35 @@ def test_gemini_audio_only_live_models_drop_text_from_text_audio_combo(patch_gem
|
|||
assert setup["generationConfig"]["responseModalities"] == ["AUDIO"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("output_modalities", ["text", ["text"]])
|
||||
def test_gemini_session_update_accepts_a_single_string_output_modality(output_modalities):
|
||||
"""Regression: a bare string was iterated per character, so output_modalities
|
||||
"text" reached Vertex as responseModalities ["T"], which it rejects with 1007."""
|
||||
config = GeminiRealtimeConfig()
|
||||
session_update = {"type": "session.update", "session": {"type": "realtime", "output_modalities": output_modalities}}
|
||||
|
||||
messages = config.transform_realtime_request(
|
||||
json.dumps(session_update), "gemini-live-2.5-flash", session_configuration_request=None
|
||||
)
|
||||
|
||||
assert json.loads(messages[0])["setup"]["generationConfig"]["responseModalities"] == ["TEXT"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("output_modalities", "expected"),
|
||||
[(["video", "text"], ["TEXT"]), (["video"], ["AUDIO"]), (None, ["AUDIO"])],
|
||||
)
|
||||
def test_gemini_session_update_drops_modalities_live_cannot_produce(output_modalities, expected):
|
||||
config = GeminiRealtimeConfig()
|
||||
session_update = {"type": "session.update", "session": {"type": "realtime", "output_modalities": output_modalities}}
|
||||
|
||||
messages = config.transform_realtime_request(
|
||||
json.dumps(session_update), "gemini-live-2.5-flash", session_configuration_request=None
|
||||
)
|
||||
|
||||
assert json.loads(messages[0])["setup"]["generationConfig"]["responseModalities"] == expected
|
||||
|
||||
|
||||
def test_gemini_non_live_model_preserves_text_modality():
|
||||
config = GeminiRealtimeConfig()
|
||||
session_update = {
|
||||
|
|
@ -1396,22 +1471,71 @@ def test_gemini_input_audio_buffer_clear_is_local_noop():
|
|||
assert messages == []
|
||||
|
||||
|
||||
def test_gemini_input_audio_buffer_commit_maps_to_activity_end_when_manual_vad():
|
||||
_MANUAL_VAD_SETUP = json.dumps({"setup": {"realtimeInputConfig": {"automaticActivityDetection": {"disabled": True}}}})
|
||||
_APPEND = json.dumps({"type": "input_audio_buffer.append", "audio": base64.b64encode(b"\x00\x01" * 80).decode()})
|
||||
|
||||
|
||||
def _realtime_input_keys(config: GeminiRealtimeConfig, client_event: str, setup: str | None) -> list[str]:
|
||||
frames = config.transform_realtime_request(client_event, "gemini-live-2.5-flash-native-audio", setup)
|
||||
return [key for frame in frames for key in json.loads(frame)["realtimeInput"]]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("turn_end", ["input_audio_buffer.commit", "response.create"])
|
||||
def test_gemini_manual_vad_wraps_each_audio_turn_in_activity_start_and_end(turn_end):
|
||||
"""With activity detection off, Gemini Live answers only after activityEnd closes an
|
||||
activityStart, and it closes the socket on any other shape of those markers."""
|
||||
config = GeminiRealtimeConfig()
|
||||
setup = {
|
||||
"setup": {
|
||||
"realtimeInputConfig": {
|
||||
"automaticActivityDetection": {"disabled": True},
|
||||
}
|
||||
}
|
||||
}
|
||||
messages = config.transform_realtime_request(
|
||||
json.dumps({"type": "input_audio_buffer.commit"}),
|
||||
"gemini-live-2.5-flash-native-audio",
|
||||
session_configuration_request=json.dumps(setup),
|
||||
end = json.dumps({"type": turn_end})
|
||||
|
||||
turns = [
|
||||
[_realtime_input_keys(config, event, _MANUAL_VAD_SETUP) for event in (_APPEND, _APPEND, end, end)]
|
||||
for _ in range(2)
|
||||
]
|
||||
|
||||
assert turns == [[["activityStart", "audio"], ["audio"], ["activityEnd"], []]] * 2
|
||||
|
||||
|
||||
def test_gemini_manual_vad_activity_markers_are_empty_messages():
|
||||
config = GeminiRealtimeConfig()
|
||||
commit = json.dumps({"type": "input_audio_buffer.commit"})
|
||||
|
||||
first_append = config.transform_realtime_request(_APPEND, "gemini-live-2.5-flash", _MANUAL_VAD_SETUP)
|
||||
turn_end = config.transform_realtime_request(commit, "gemini-live-2.5-flash", _MANUAL_VAD_SETUP)
|
||||
|
||||
assert json.loads(first_append[0]) == {"realtimeInput": {"activityStart": {}}}
|
||||
assert [json.loads(frame) for frame in turn_end] == [{"realtimeInput": {"activityEnd": {}}}]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("session", "expected_keys"),
|
||||
[
|
||||
({"turn_detection": None}, [["activityStart", "audio"], ["activityEnd"]]),
|
||||
({"type": "realtime", "audio": {"input": {"turn_detection": None}}}, [["activityStart", "audio"], ["activityEnd"]]),
|
||||
({"type": "realtime", "audio": {"input": {}}}, [["audio"], ["audioStreamEnd"]]),
|
||||
],
|
||||
)
|
||||
def test_gemini_null_turn_detection_makes_the_client_mark_each_turn(session, expected_keys):
|
||||
"""OpenAI's turn_detection null turns server VAD off, so the client commits each turn itself."""
|
||||
config = GeminiRealtimeConfig()
|
||||
setup = config.transform_realtime_request(
|
||||
json.dumps({"type": "session.update", "session": session}), "gemini-live-2.5-flash-native-audio", None
|
||||
)
|
||||
assert len(messages) == 1
|
||||
assert json.loads(messages[0]) == {"realtimeInput": {"activityEnd": True}}
|
||||
commit = json.dumps({"type": "input_audio_buffer.commit"})
|
||||
|
||||
keys = [_realtime_input_keys(config, event, setup[0]) for event in (_APPEND, commit)]
|
||||
|
||||
assert keys == expected_keys
|
||||
|
||||
|
||||
def test_gemini_automatic_vad_sends_no_activity_markers():
|
||||
config = GeminiRealtimeConfig()
|
||||
commit = json.dumps({"type": "input_audio_buffer.commit"})
|
||||
response_create = json.dumps({"type": "response.create"})
|
||||
|
||||
keys = [_realtime_input_keys(config, event, None) for event in (_APPEND, commit)]
|
||||
|
||||
assert keys == [["audio"], ["audioStreamEnd"]]
|
||||
assert config.transform_realtime_request(response_create, "gemini-live-2.5-flash", None) == []
|
||||
|
||||
|
||||
def test_gemini_tool_call_id_to_name_evicts_oldest_when_capped():
|
||||
|
|
@ -1883,13 +2007,75 @@ def test_map_openai_params_drops_stock_voice_case_insensitively():
|
|||
"""
|
||||
cfg = GeminiRealtimeConfig()
|
||||
|
||||
dropped = cfg.map_openai_params(optional_params={}, non_default_params={"voice": "Alloy"})
|
||||
dropped = cfg.map_openai_params(optional_params={}, non_default_params={"voice": "Alloy"}, model="gemini-3.8-live")
|
||||
assert "speechConfig" not in dropped.get("generationConfig", {})
|
||||
|
||||
passthrough = cfg.map_openai_params(optional_params={}, non_default_params={"voice": "Kore"})
|
||||
passthrough = cfg.map_openai_params(
|
||||
optional_params={}, non_default_params={"voice": "Kore"}, model="gemini-3.8-live"
|
||||
)
|
||||
assert passthrough["generationConfig"]["speechConfig"]["voiceConfig"]["prebuiltVoiceConfig"]["voiceName"] == "Kore"
|
||||
|
||||
|
||||
_DOCUMENTED_VOICE: Final = sorted(GEMINI_LIVE_DOCUMENTED_VOICES)[0].capitalize()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("drop_params", [False, True])
|
||||
def test_google_ai_studio_forwards_a_documented_voice(drop_params):
|
||||
config: Final = GeminiRealtimeConfig(drop_params=drop_params)
|
||||
mapped: Final = config.map_openai_params(
|
||||
optional_params={}, non_default_params={"voice": _DOCUMENTED_VOICE}, model="gemini-3.8-live"
|
||||
)
|
||||
|
||||
speech_config: Final = mapped["generationConfig"]["speechConfig"]
|
||||
assert speech_config["voiceConfig"]["prebuiltVoiceConfig"]["voiceName"] == _DOCUMENTED_VOICE
|
||||
assert config.setup_session_update_rejection("gemini-3.8-live", {"voice": _DOCUMENTED_VOICE}, setup_sent=False) is None
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("session", "param"),
|
||||
[
|
||||
({"voice": "Nova"}, "session.voice"),
|
||||
({"type": "realtime", "audio": {"output": {"voice": "Nova"}}}, "session.audio.output.voice"),
|
||||
],
|
||||
)
|
||||
def test_google_ai_studio_rejects_an_undocumented_voice_on_the_session_update(session, param):
|
||||
rejection: Final = GeminiRealtimeConfig().setup_session_update_rejection(
|
||||
"gemini-2.5-flash-native-audio-preview-12-2025", session, setup_sent=False
|
||||
)
|
||||
|
||||
assert rejection is not None
|
||||
assert (rejection.code, rejection.param) == ("invalid_value", param)
|
||||
|
||||
|
||||
def _setup_speech_config(config: GeminiRealtimeConfig, model: str, session: dict[str, object]) -> object:
|
||||
messages: Final = config.transform_realtime_request(
|
||||
json.dumps({"type": "session.update", "session": session}), model, session_configuration_request=None
|
||||
)
|
||||
return json.loads(messages[0])["setup"].get("generationConfig", {}).get("speechConfig")
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"session",
|
||||
[{"voice": "Nova"}, {"type": "realtime", "audio": {"output": {"voice": "Nova"}}}],
|
||||
)
|
||||
def test_google_ai_studio_drops_an_undocumented_voice_with_drop_params(session):
|
||||
config: Final = GeminiRealtimeConfig(drop_params=True)
|
||||
model: Final = "gemini-2.5-flash-native-audio-preview-12-2025"
|
||||
|
||||
assert config.setup_session_update_rejection(model, session, setup_sent=False) is None
|
||||
assert _setup_speech_config(config, model, session) is None
|
||||
|
||||
|
||||
@pytest.mark.parametrize("drop_params", [False, True])
|
||||
@pytest.mark.parametrize("model", ["gemini-3.8-live", "gemini-3.1-flash-live-preview"])
|
||||
def test_google_ai_studio_forwards_any_voice_to_models_that_accept_it(model, drop_params):
|
||||
config: Final = GeminiRealtimeConfig(drop_params=drop_params)
|
||||
|
||||
assert config.setup_session_update_rejection(model, {"voice": "Nova"}, setup_sent=False) is None
|
||||
speech_config: Final = _setup_speech_config(config, model, {"voice": "Nova"})
|
||||
assert speech_config == {"voiceConfig": {"prebuiltVoiceConfig": {"voiceName": "Nova"}}}
|
||||
|
||||
|
||||
@pytest.fixture(autouse=False)
|
||||
def patch_gemini_transcribe_live_cost_map_entry(monkeypatch):
|
||||
"""Inject the gemini-3.5-transcribe-live registry entry locally.
|
||||
|
|
@ -1954,6 +2140,55 @@ def test_gemini_transcribe_live_coerces_audio_modality_to_text(modalities, patch
|
|||
assert setup["generationConfig"]["responseModalities"] == ["TEXT"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model", ["gemini-3.5-transcribe-live", "gemini/gemini-3.5-transcribe-live"])
|
||||
def test_gemini_transcribe_live_synthetic_session_created_reports_text_output(
|
||||
model, patch_gemini_transcribe_live_cost_map_entry
|
||||
):
|
||||
session = GeminiRealtimeConfig().transform_session_created_event(model, "sess_transcribe")["session"]
|
||||
|
||||
assert session["modalities"] == ["text"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("model", "modalities", "producible"),
|
||||
[
|
||||
("gemini-3.5-transcribe-live", ["audio"], False),
|
||||
("gemini-3.5-transcribe-live", ["audio", "text"], True),
|
||||
("gemini-live-2.5-flash-native-audio", ["text"], False),
|
||||
("gemini-live-2.5-flash-native-audio", ["text", "audio"], True),
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize("drop_params", [False, True])
|
||||
def test_gemini_live_rejects_only_output_modalities_the_model_cannot_produce_at_all(
|
||||
model,
|
||||
modalities,
|
||||
producible,
|
||||
drop_params,
|
||||
patch_gemini_transcribe_live_cost_map_entry,
|
||||
patch_gemini_audio_cost_map_entries,
|
||||
):
|
||||
rejection = GeminiRealtimeConfig(drop_params=drop_params).setup_session_update_rejection(
|
||||
model, {"output_modalities": modalities}, setup_sent=False
|
||||
)
|
||||
|
||||
assert (rejection is None) is (drop_params or producible)
|
||||
assert rejection is None or rejection.param == "session.output_modalities"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("drop_params", [False, True])
|
||||
def test_gemini_provider_realtime_config_carries_drop_params(drop_params):
|
||||
from litellm.types.utils import LlmProviders
|
||||
from litellm.utils import ProviderConfigManager
|
||||
|
||||
model = "gemini-live-2.5-flash-native-audio"
|
||||
config = ProviderConfigManager.get_provider_realtime_config(
|
||||
model=model, provider=LlmProviders.GEMINI, drop_params=drop_params
|
||||
)
|
||||
|
||||
rejection = config.setup_session_update_rejection(model, {"modalities": ["text"]}, setup_sent=False)
|
||||
assert (rejection is None) is drop_params
|
||||
|
||||
|
||||
def test_gemini_chat_model_with_text_output_modalities_keeps_audio_eager_setup(
|
||||
patch_gemini_transcribe_live_cost_map_entry,
|
||||
):
|
||||
|
|
@ -2169,6 +2404,47 @@ def test_unbilled_usage_on_session_close_flushes_trailing_audio(patch_gemini_tra
|
|||
assert config.unbilled_usage_on_session_close("gemini-3.5-transcribe-live") is None
|
||||
|
||||
|
||||
def test_gemini_realtime_two_output_modalities_collapse_to_audio():
|
||||
"""Gemini Live closes the socket with 1007 when the setup names two response
|
||||
modalities, so a GA session.update asking for ["text", "audio"] must reach the
|
||||
backend as AUDIO only, the same choice the beta shim makes."""
|
||||
config = GeminiRealtimeConfig()
|
||||
session_update = {
|
||||
"type": "session.update",
|
||||
"session": {"output_modalities": ["text", "audio"], "instructions": "Be brief."},
|
||||
}
|
||||
|
||||
messages = config.transform_realtime_request(
|
||||
json.dumps(session_update),
|
||||
"gemini-live-2.5-flash",
|
||||
session_configuration_request=None,
|
||||
)
|
||||
|
||||
assert len(messages) == 1
|
||||
assert json.loads(messages[0])["setup"]["generationConfig"]["responseModalities"] == ["AUDIO"]
|
||||
|
||||
|
||||
def test_gemini_realtime_inf_max_response_output_tokens_is_not_forwarded():
|
||||
"""OpenAI's default max_response_output_tokens is the string "inf", which Gemini
|
||||
rejects with 1007 (TYPE_INT32), so only integer limits map to maxOutputTokens."""
|
||||
config = GeminiRealtimeConfig()
|
||||
|
||||
def generation_config_for(limit):
|
||||
session_update = {
|
||||
"type": "session.update",
|
||||
"session": {"modalities": ["text"], "max_response_output_tokens": limit},
|
||||
}
|
||||
messages = config.transform_realtime_request(
|
||||
json.dumps(session_update),
|
||||
"gemini-live-2.5-flash",
|
||||
session_configuration_request=None,
|
||||
)
|
||||
return json.loads(messages[0])["setup"]["generationConfig"]
|
||||
|
||||
assert "maxOutputTokens" not in generation_config_for("inf")
|
||||
assert generation_config_for(200)["maxOutputTokens"] == 200
|
||||
|
||||
|
||||
def _grounded_live_frame(grounding_metadata: Mapping[str, object] | None) -> Mapping[str, object]:
|
||||
"""One Live server frame. Grounding metadata and usageMetadata arrive together, as Vertex sends them."""
|
||||
from typing import Final
|
||||
|
|
|
|||
Some files were not shown because too many files have changed in this diff Show more
Loading…
Add table
Reference in a new issue