fix!: backport the second batch of rc/1.107.0 CI triage fixes from main (#47090)

* fix(bedrock)!: block realtime replay after an input-first start and reject malformed Nova Sonic session settings (#46936)

* fix(bedrock): block realtime replay after input starts the session and reject non-string instructions and voice

* test(bedrock): name the malformed realtime setting test for its string cases

* fix(bedrock): replay the opening realtime message until Bedrock accepts the stream and refuse any fallback after commit

Also rejects a malformed modalities, input_sample_rate_hertz or tools setting with invalid_type, dropped under drop_params.

* fix(bedrock): replay a refused opening message on any fallback and reject OpenAI voices and non-string audio formats

* fix(realtime): read the client scope only when the websocket has one and reject any-case OpenAI and empty Bedrock voices

* test(bedrock): cover realtime replay, refusal and committed failure through a real proxy

Fix the scripted Bedrock peer's port import and the nested tool call cell's stale expected events, so the Nova Sonic integration file collects again

(cherry picked from commit 3ef0b35255)

* test(integration): write worker configs beside test-supplied configs (#47066)

Under worker isolation the harness rewrote every proxy config into the owned proxy's directory, so modules a test placed next to its own config (custom_auth) no longer resolved and the proxy died at boot. The copy now lands next to a test-supplied config, while configs inside the checkout still go to the worker directory so nothing is written into the repo tree

(cherry picked from commit e754857de9)

* test(e2e): keep the prompt-caching affinity test on the live provider (#47072)

With E2E_PROVIDER_CACHE=1 the test's Anthropic deployment went through the
shared response cache. Each call's recording expires on its own, so a run
could replay the cache-writing call while a follow-up went live against a
cold provider cache, recording creation instead of a read. That recording
then replayed as a failure for 24 hours. Mark the test provider_live, as
PROVIDER_CACHE.md requires for prompt-caching tests

(cherry picked from commit 6b45a196c4)

* test(integration): send post-kill requests on a fresh connection in two worker-kill tests (#47080)

kill() returns once the killed worker's main thread is a zombie, but its other threads can
still be exiting and holding its sockets open. A keep-alive connection pooled to that worker
still looks healthy, so the next request goes out on it and the closing socket answers with
a reset. The old reap wait used to hide this. A new connection can only be accepted by a
live worker, so the surviving-worker assertions now run on a fresh client

(cherry picked from commit c32eea634a)

* fix(caching): keep prompt cache usage on streamed cache hits (#47073)

* fix(caching): keep prompt cache usage on streamed cache hits

A stream=true request served from the response cache rebuilt its usage from
the prompt, completion and total counts only, so prompt_tokens_details and the
cache read and creation counts were dropped on the replayed stream. /v1/messages
clients routed to Bedrock Converse saw input_tokens equal to the whole prompt
and no cache tokens on a hit.

Both replays now build the usage through one helper that also restores
prompt_tokens_details and the integer cache read and creation counts, skipping
cached values that do not validate.

* fix(caching): coerce replayed token counts the way the live usage does

The replay accepted only exact ints for the token counts, so an integral float
or numeric string in a cached entry became 0 where the previous replay kept it.
Counts now go through pydantic int coercion and fall back only when they do not
validate.

(cherry picked from commit d346d769f0)

* fix(logging): keep the async request's logging payload when the sync success thread runs late (#47019)

* fix(logging): keep the async request's logging payload when sync callbacks run late

Treat the "cache" success callback as internal when deciding whether to start
the sync success thread for async requests. It has no branch in the sync
handler, so with `cache: true` that thread did nothing except rebuild and
overwrite model_call_details["standard_logging_object"].

When the sync handler does run for an async request (a real sync-only logger
is configured), it now builds its own payload and hands it only to its own
callbacks. The async handler keeps sole ownership of the shared payload, so
edits made by async logging hooks (for example the v3 rate limiter's mirrored
x-ratelimit headers that Prometheus reads) are no longer discarded. Sync
requests keep their existing behavior.

* fix(logging): keep sync cache hits logged and isolate the async-request payload build

The sync completion cache-hit path reused handle_sync_success_callbacks_for_async_calls,
so treating "cache" as internal silently dropped sync callbacks on sync cache hits. It now
submits the sync handler through its own method, as it effectively always did.

For async requests the sync handler builds its payload from a copy of the call details, so
message redaction cannot rewrite the shared standard_logging_object, and when it builds no
payload its callbacks keep the shared one instead of getting None.

* test(caching): lock sync cache-hit logging for completion, embedding and text_completion

* test(logging): keep sync hook edits off a caller-supplied payload for async requests

* fix(logging): give sync hooks a copy of the async payload when the thread builds none

* test(logging): check every legacy sync logger gets the thread's own payload

* test(logging): cover traceloop and supabase sync logging for async requests

(cherry picked from commit a4e0570bea)

* fix(cost): price exact cost-map names before price-free capability rules, and fix SCIM and MCP integration tests on CircleCI (#47037)

* fix(test-harness): preserve sibling auth modules in isolated worker configs

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(integration): expect servers in MCP connect-flow descriptions

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(integration): expect ptu_hours in daily activity metrics after #43043

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(integration): kill every serving worker once the SCIM push blocks on the row lock

* fix(cost): price exact cost-map names before price-free capability rule matches

* test(streaming): price a priced alias stream at its deployment model and keep the unpriceable case unmapped

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(cost): price an explicit base_model under its own provider after candidate reordering

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(cost): retain priced generalized candidates and skip explicit-pricing lookups

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(cost): assert custom pricing makes no remote model-info requests

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* refactor(cost): keep the per-candidate pricing provider final

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(router): pin a Bedrock application profile prompt for the one-hour ttl its base model sends

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(integration): assert the Nova Sonic tool-call argument delta carries the full arguments

* fix(cost): sort cost candidates from the static cost map and rules, without provider lookups

* test(integration): bill route-priced Bedrock chat streams and responses, and pin application profiles for an hour on messages and cache hits

---------

Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Co-authored-by: Yuneng Jiang <yuneng@berri.ai>
(cherry picked from commit 4a58784fa2)

* test(e2e): warm the memory test up with two batches before the measured one (#47087)

* test(e2e): warm the memory test up with two batches before the measured one

On a cold worker the second batch of failing requests can still raise the
allocator's high-water mark by about 50 MB while the live Python heap stays
flat, so the 48 MB budget failed about 1 in 10 Buildkite runs without a leak.
The test now runs two warmup batches, measures the third against the same
budget, and records every batch's settled growth per worker as an
rss_batch_growth JUnit property for later calibration

* test(e2e): record each memory batch's growth as soon as it settles

A failure in a later batch no longer drops the growth already measured, and
the worker keys are no longer sorted, which raised on a mix of named and
unnamed hosts

* test(e2e): silence the xunit2 record_property warning on the memory test

(cherry picked from commit 1fd7f81f97)

* test: import pytest and Final for the backported #47073 test on rc

On main these imports came from an earlier commit that is not on rc/1.107.0, so the
picked test failed collection with NameError

---------

Co-authored-by: Mateo Wang <277851410+mateo-berri@users.noreply.github.com>
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
yuneng-jiang 2026-10-17 00:20:23 -07:00 • committed by GitHub
parent a00d8a6ee1
commit 4bf1633452
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
36 changed files with 2099 additions and 304 deletions

View file

@ -512,7 +512,7 @@ class LLMCachingHandler:
)
if not _should_defer_streaming_cache_hit_callbacks(cached_result=cached_result):
logging_obj.handle_sync_success_callbacks_for_async_calls(
logging_obj.handle_sync_cache_hit_success_callbacks(
result=cached_result,
start_time=start_time,
end_time=end_time,

View file

@ -389,7 +389,7 @@ DEEPGRAM_DEFAULT_API_BASE: Final = "https://api.deepgram.com/v1"
NADIR_DEFAULT_API_BASE: Final = "https://api.getnadir.com/v1"
DEEPGRAM_LISTEN_DEFAULT_MODEL: Final = "nova-3"
BEDROCK_REALTIME_PENDING_SESSION_UPDATE_SCOPE_KEY: Final = "litellm.bedrock_realtime.pending_session_update"
BEDROCK_REALTIME_PENDING_OPENING_MESSAGE_SCOPE_KEY: Final = "litellm.bedrock_realtime.pending_opening_message"
BEDROCK_REALTIME_SESSION_COMMITTED_SCOPE_KEY: Final = "litellm.bedrock_realtime.session_committed"
BEDROCK_REALTIME_COMMITTED_FAILURE_SCOPE_KEY: Final = "litellm.bedrock_realtime.committed_failure"
BEDROCK_REALTIME_SDK_DISTRIBUTION: Final = "aws-sdk-bedrock-runtime"
@ -1613,6 +1613,7 @@ CLOUDZERO_EXPORT_INTERVAL_MINUTES: Final = get_env_int("CLOUDZERO_EXPORT_INTERVA
MCP_TOOL_NAME_PREFIX: Final = "mcp_tool"
MAXIMUM_TRACEBACK_LINES_TO_LOG: Final = get_env_int("MAXIMUM_TRACEBACK_LINES_TO_LOG", 100)
PASSTHROUGH_UPSTREAM_ERROR_BODY_MAX_LOG_CHARS: Final = 4096
INTERNAL_CALLBACK_NAMES: Final = frozenset({"cache", "vector_store_pre_call_hook"})
# Headers to control callbacks
X_LITELLM_DISABLE_CALLBACKS: Final = "x-litellm-disable-callbacks"

View file

@ -7,7 +7,7 @@ from collections.abc import Mapping, Sequence
from datetime import datetime
from functools import lru_cache
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, Literal, cast
from typing import TYPE_CHECKING, Any, Final, Literal, TypeAlias, cast
from httpx import Response
from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError
@ -147,6 +147,7 @@ from litellm.utils import (
TranscriptionResponse,
cached_get_model_info_helper,
token_counter,
unmapped_model_generalization,
)
if TYPE_CHECKING:
@ -863,6 +864,16 @@ def _base_model_pricing_provider(base_model: str, custom_llm_provider: str | Non
return own_model_info["litellm_provider"]
def _pricing_provider_for_cost_candidate(
candidate_provider: str | None,
request_provider: str | None,
current_provider: object,
) -> str | None:
if candidate_provider != request_provider:
return candidate_provider
return current_provider if isinstance(current_provider, str) else None
def select_model_name_for_cost_calc(
model: str | None,
completion_response: object | None,
@ -1549,14 +1560,20 @@ def completion_cost(
)
requested_model: Final = model
potential_model_names: Final = [
selected_model,
_get_response_model(completion_response),
]
if model is not None:
potential_model_names.append(model)
request_pricing_provider: Final = custom_llm_provider if isinstance(custom_llm_provider, str) else None
raw_candidates: Final = (
(
selected_model,
base_model_provider if base_model_provider is not None else request_pricing_provider,
),
(_get_response_model(completion_response), request_pricing_provider),
*(((model, request_pricing_provider),) if model is not None else ()),
)
potential_model_names: Final = (
raw_candidates if explicit_pricing else _priced_candidates_before_unpriced_rules(raw_candidates)
)
for idx, model in enumerate(potential_model_names):
for idx, (model, pricing_provider) in enumerate(potential_model_names):
try:
if verbose_logger.isEnabledFor(logging.DEBUG):
verbose_logger.debug("selected model name for cost calculation: %s", model)
@ -1895,13 +1912,6 @@ def completion_cost(
if litellm_logging_obj is not None:
request_model_for_cost = litellm_logging_obj.model
pricing_provider: str | None = (
base_model_provider
if base_model_provider is not None and idx == 0
else custom_llm_provider
if isinstance(custom_llm_provider, str)
else None
)
(
prompt_tokens_cost_usd_dollar,
completion_tokens_cost_usd_dollar,
@ -1909,7 +1919,11 @@ def completion_cost(
model=model,
prompt_tokens=prompt_tokens or 0,
completion_tokens=completion_tokens or 0,
custom_llm_provider=pricing_provider,
custom_llm_provider=_pricing_provider_for_cost_candidate(
candidate_provider=pricing_provider,
request_provider=request_pricing_provider,
current_provider=custom_llm_provider,
),
response_time_ms=response_time_ms,
region_name=None if explicit_pricing else region_name,
custom_cost_per_second=custom_cost_per_second,
@ -2008,7 +2022,11 @@ def completion_cost(
if cost_per_token_usage_object is not None and model:
_token_type_breakdown = get_token_type_cost_breakdown(
model=model,
custom_llm_provider=pricing_provider,
custom_llm_provider=_pricing_provider_for_cost_candidate(
candidate_provider=pricing_provider,
request_provider=request_pricing_provider,
current_provider=custom_llm_provider,
),
usage=cost_per_token_usage_object,
service_tier=service_tier,
data_residency=data_residency,
@ -2263,6 +2281,25 @@ def _cost_map_model_info(model: str, custom_llm_provider: str | None) -> ModelIn
return None
_CostCandidate: TypeAlias = tuple[str | None, str | None]
def _is_unpriced_generalized_candidate(candidate: _CostCandidate) -> bool:
model, pricing_provider = candidate
if model is None:
return False
rule_info: Final = unmapped_model_generalization(model, pricing_provider)
return rule_info is not None and not _cost_map_entry_prices_anything(rule_info)
def _priced_candidates_before_unpriced_rules(candidates: tuple[_CostCandidate, ...]) -> tuple[_CostCandidate, ...]:
unpriced: Final = tuple(_is_unpriced_generalized_candidate(candidate) for candidate in candidates)
return (
*(candidate for candidate, is_unpriced in zip(candidates, unpriced, strict=True) if not is_unpriced),
*(candidate for candidate, is_unpriced in zip(candidates, unpriced, strict=True) if is_unpriced),
)
def _raw_cost_map_entry(key: str) -> Mapping[str, object] | None:
raw_entry: Final = litellm.model_cost.get(key)
return raw_entry if isinstance(raw_entry, Mapping) else None

View file

@ -41,6 +41,7 @@ from litellm.constants import (
DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT,
DEFAULT_MOCK_RESPONSE_PROMPT_TOKEN_COUNT,
EMPTY_MAPPING,
INTERNAL_CALLBACK_NAMES,
PROVIDER_REQUEST_ID_HEADERS,
REDACTED_BY_LITELLM,
)
@ -2287,6 +2288,7 @@ class Logging(LiteLLMLoggingBaseClass):
start_time=start_time,
end_time=end_time,
cache_hit=cache_hit,
for_async_request=True,
**kwargs,
)
@ -2480,7 +2482,8 @@ class Logging(LiteLLMLoggingBaseClass):
start_time,
end_time,
build_logging_payload: bool = True,
):
publish_logging_payload: bool = True,
) -> StandardLoggingPayload | None:
"""Resolve hidden params, compute response cost, and emit the standard logging payload."""
self._surface_response_headers_from_result(logging_result)
hidden_params: Final = getattr(logging_result, HIDDEN_PARAMS_ATTR, {})
@ -2507,17 +2510,17 @@ class Logging(LiteLLMLoggingBaseClass):
self.model_call_details["response_cost"] = self.response_cost_calculator(result=logging_result)
if not build_logging_payload:
return
return None
self.model_call_details["standard_logging_object"] = self._build_standard_logging_payload(
logging_result, start_time, end_time
)
standard_logging_payload: Final[StandardLoggingPayload | None] = self.model_call_details.get(
"standard_logging_object"
standard_logging_payload: Final = self._build_standard_logging_payload(
logging_result, start_time, end_time, isolate_call_details=not publish_logging_payload
)
if not publish_logging_payload:
return standard_logging_payload
self.model_call_details["standard_logging_object"] = standard_logging_payload
if standard_logging_payload is not None:
emit_standard_logging_payload(standard_logging_payload)
return standard_logging_payload
async def _prepare_baseline_cache_estimate(self, response_obj: object) -> None:
if self.baseline_cache_context is None:
@ -2533,12 +2536,16 @@ class Logging(LiteLLMLoggingBaseClass):
await invalidate_baseline_cache(self, reason, completed=completed)
def _build_standard_logging_payload(
self, init_response_obj: object, start_time: dt_object, end_time: dt_object
self,
init_response_obj: object,
start_time: dt_object,
end_time: dt_object,
isolate_call_details: bool = False,
) -> StandardLoggingPayload | None:
"""Build StandardLoggingPayload and accumulate its construction time."""
_start: Final = time.time()
payload: Final = get_standard_logging_object_payload(
kwargs=self.model_call_details,
kwargs=dict(self.model_call_details) if isolate_call_details else self.model_call_details,
init_response_obj=init_response_obj,
start_time=start_time,
end_time=end_time,
@ -2555,11 +2562,16 @@ class Logging(LiteLLMLoggingBaseClass):
emit_standard_logging_payload(stored_payload)
def _transform_usage_objects(self, result: object) -> object:
return self._transform_usage_objects_for_payload(result, self.model_call_details.get("standard_logging_object"))
def _transform_usage_objects_for_payload(
self, result: object, standard_logging_payload: StandardLoggingPayload | None
) -> object:
if isinstance(result, ResponsesAPIResponse):
result = result.model_copy()
transformed_usage = ResponseAPILoggingUtils.transform_response_api_usage_to_chat_usage(result.usage)
setattr(result, "usage", transformed_usage)
if (standard_logging_payload := self.model_call_details.get("standard_logging_object")) is not None:
if standard_logging_payload is not None:
response_dict: Final = result.model_dump() if hasattr(result, "model_dump") else dict(result)
# Ensure usage is properly included with transformed chat format
if transformed_usage is not None:
@ -2568,7 +2580,7 @@ class Logging(LiteLLMLoggingBaseClass):
if hasattr(transformed_usage, "model_dump")
else dict(transformed_usage)
)
standard_logging_payload["response"] = response_dict
standard_logging_payload["response"] = response_dict # rebind-ok: mirrors the transformed usage
elif isinstance(result, TranscriptionResponse):
from litellm.litellm_core_utils.llm_cost_calc.usage_object_transformation import (
TranscriptionUsageObjectTransformation,
@ -2599,6 +2611,27 @@ class Logging(LiteLLMLoggingBaseClass):
standard_logging_object: StandardLoggingPayload | None = None,
build_logging_payload: bool = True,
) -> tuple[datetime.datetime, datetime.datetime, object]:
prepared: Final = self._prepare_success_logging(
result=result,
start_time=start_time,
end_time=end_time,
cache_hit=cache_hit,
standard_logging_object=standard_logging_object,
build_logging_payload=build_logging_payload,
publish_logging_payload=True,
)
return prepared[0], prepared[1], prepared[2]
def _prepare_success_logging(
self,
result: object,
start_time: datetime.datetime | None,
end_time: datetime.datetime | None,
cache_hit: bool | None,
standard_logging_object: StandardLoggingPayload | None,
build_logging_payload: bool,
publish_logging_payload: bool,
) -> tuple[datetime.datetime, datetime.datetime, object, StandardLoggingPayload | None]:
try:
if start_time is None:
start_time = self.start_time
@ -2629,29 +2662,21 @@ class Logging(LiteLLMLoggingBaseClass):
):
result = logging_result
if standard_logging_object is None and result is not None and self.stream is not True:
if self._is_recognized_call_type_for_logging(logging_result=logging_result) or isinstance(
logging_result, (dict, list)
):
self._process_hidden_params_and_response_cost(
logging_result=logging_result,
start_time=start_time,
end_time=end_time,
build_logging_payload=build_logging_payload,
)
elif standard_logging_object is not None:
self.model_call_details["standard_logging_object"] = standard_logging_object
else:
# Streaming reaches here before its cost is known, so the cost
# is seeded to None, but only when nothing has already
# established one. A stream that assembles into a response
# object recomputes the cost right after this; a pass-through
# stream cannot (its body is opaque) and carries the cost its
# upstream reported in the response headers, which an
# unconditional reset would discard.
self.model_call_details.setdefault("response_cost", None)
standard_logging_payload: Final = self._resolve_success_logging_payload(
result=result,
logging_result=logging_result,
start_time=start_time,
end_time=end_time,
standard_logging_object=standard_logging_object,
build_logging_payload=build_logging_payload,
publish_logging_payload=publish_logging_payload,
)
result = self._transform_usage_objects(result=result)
result = (
self._transform_usage_objects(result=result)
if publish_logging_payload
else self._transform_usage_objects_for_payload(result, standard_logging_payload)
)
if (
litellm.max_budget
@ -2670,10 +2695,47 @@ class Logging(LiteLLMLoggingBaseClass):
standard_built_in_tools_params=self.standard_built_in_tools_params,
)
return start_time, end_time, result
return start_time, end_time, result, standard_logging_payload
except Exception as e:
raise Exception(f"[Non-Blocking] LiteLLM.Success_Call Error: {e}")
def _resolve_success_logging_payload(
self,
result: object,
logging_result: object,
start_time: datetime.datetime,
end_time: datetime.datetime,
standard_logging_object: StandardLoggingPayload | None,
build_logging_payload: bool,
publish_logging_payload: bool,
) -> StandardLoggingPayload | None:
if standard_logging_object is not None:
if not publish_logging_payload:
return standard_logging_object.copy()
self.model_call_details["standard_logging_object"] = standard_logging_object
return standard_logging_object
if result is None or self.stream is True:
# Streaming reaches here before its cost is known, so the cost
# is seeded to None, but only when nothing has already
# established one. A stream that assembles into a response
# object recomputes the cost right after this; a pass-through
# stream cannot (its body is opaque) and carries the cost its
# upstream reported in the response headers, which an
# unconditional reset would discard.
self.model_call_details.setdefault("response_cost", None)
return None
if not self._is_recognized_call_type_for_logging(logging_result=logging_result) and not isinstance(
logging_result, (dict, list)
):
return None
return self._process_hidden_params_and_response_cost(
logging_result=logging_result,
start_time=start_time,
end_time=end_time,
build_logging_payload=build_logging_payload,
publish_logging_payload=publish_logging_payload,
)
def _is_recognized_call_type_for_logging(
self,
logging_result: object,
@ -2855,6 +2917,8 @@ class Logging(LiteLLMLoggingBaseClass):
start_time: datetime.datetime | None = None,
end_time: datetime.datetime | None = None,
cache_hit: bool | None = None,
*,
for_async_request: bool = False,
**kwargs: object, # kwargs-ok: forwarded to _success_handler_body
) -> None:
"""Restores trace_id/session_id contextvars once this attempt's own success
@ -2862,7 +2926,12 @@ class Logging(LiteLLMLoggingBaseClass):
try:
with post_response_phase():
return self._success_handler_body(
result=result, start_time=start_time, end_time=end_time, cache_hit=cache_hit, **kwargs
result=result,
start_time=start_time,
end_time=end_time,
cache_hit=cache_hit,
for_async_request=for_async_request,
**kwargs,
)
finally:
self._restore_correlation_context()
@ -2873,17 +2942,20 @@ class Logging(LiteLLMLoggingBaseClass):
start_time: datetime.datetime | None = None,
end_time: datetime.datetime | None = None,
cache_hit: bool | None = None,
for_async_request: bool = False,
**kwargs: Any, # kwargs-ok: forwarded from success_handler
) -> None:
verbose_logger.debug("Logging Details LiteLLM-Success Call: Cache_hit=%s", cache_hit)
if not self.should_run_logging(event_type="sync_success"): # prevent double logging
return
start_time, end_time, result = self._success_handler_helper_fn(
start_time, end_time, result, request_logging_payload = self._prepare_success_logging(
start_time=start_time,
end_time=end_time,
result=result,
cache_hit=cache_hit,
standard_logging_object=kwargs.get("standard_logging_object", None),
build_logging_payload=True,
publish_logging_payload=not for_async_request,
)
litellm_params = self.model_call_details.get("litellm_params", {})
is_sync_request: Final = self.is_sync_litellm_request(litellm_params)
@ -2912,16 +2984,45 @@ class Logging(LiteLLMLoggingBaseClass):
logged_streaming_response: Final = self._with_client_facing_stream_model(complete_streaming_response)
self.model_call_details["complete_streaming_response"] = logged_streaming_response
## STANDARDIZED LOGGING PAYLOAD
self.model_call_details["standard_logging_object"] = self._build_standard_logging_payload(
logged_streaming_response, start_time, end_time
if not for_async_request:
self.model_call_details["standard_logging_object"] = self._build_standard_logging_payload(
logged_streaming_response, start_time, end_time
)
standard_logging_payload: Final[StandardLoggingPayload | None] = self.model_call_details.get(
"standard_logging_object"
)
if standard_logging_payload is not None:
# Only emit for sync requests (async_success_handler handles async)
if is_sync_request:
emit_standard_logging_payload(standard_logging_payload)
callback_logging_payload: Final = (
self._build_standard_logging_payload(
self.model_call_details["complete_streaming_response"],
start_time,
end_time,
isolate_call_details=True,
)
standard_logging_payload: Final[StandardLoggingPayload | None] = self.model_call_details.get(
"standard_logging_object"
if for_async_request and complete_streaming_response is not None
else request_logging_payload
)
shared_logging_payload: Final[StandardLoggingPayload | None] = self.model_call_details.get(
"standard_logging_object"
)
isolated_logging_payload: Final = (
callback_logging_payload
if callback_logging_payload is not None or shared_logging_payload is None
else shared_logging_payload.copy()
)
call_details: dict[str, Any] # mutable-ok: legacy loggers take a plain dict
call_details = ( # rebind-ok: logging hooks may return a new dict
self.model_call_details
if not for_async_request
else (
dict(self.model_call_details)
if isolated_logging_payload is None
else {**self.model_call_details, "standard_logging_object": isolated_logging_payload}
)
if standard_logging_payload is not None:
# Only emit for sync requests (async_success_handler handles async)
if is_sync_request:
emit_standard_logging_payload(standard_logging_payload)
)
callbacks: Final = self.get_combined_callback_list(
dynamic_success_callbacks=self.dynamic_success_callbacks,
global_callbacks=litellm.success_callback,
@ -2929,7 +3030,7 @@ class Logging(LiteLLMLoggingBaseClass):
## REDACT MESSAGES ##
result = redact_message_input_output_from_logging(
model_call_details=(self.model_call_details if hasattr(self, "model_call_details") else {}),
model_call_details=call_details,
result=result,
)
## LOGGING HOOK ##
@ -2939,7 +3040,7 @@ class Logging(LiteLLMLoggingBaseClass):
if (
callback.should_run_guardrail(
data=self.model_call_details,
data=call_details,
event_type=GuardrailEventHooks.logging_only,
)
is not True
@ -2947,18 +3048,21 @@ class Logging(LiteLLMLoggingBaseClass):
continue
self.shadow_eval_request_snapshot = None
self.model_call_details, result = callback.logging_hook(
kwargs=self.model_call_details,
call_details, result = callback.logging_hook( # pyright: ignore[reportUnknownVariableType] # logging_hook is typed with a bare dict
kwargs=call_details,
result=result,
call_type=self.call_type,
)
elif isinstance(callback, CustomLogger):
self.model_call_details, result = callback.logging_hook(
kwargs=self.model_call_details,
call_details, result = callback.logging_hook( # pyright: ignore[reportUnknownVariableType] # logging_hook is typed with a bare dict
kwargs=call_details,
result=result,
call_type=self.call_type,
)
if not for_async_request:
self.model_call_details = call_details
call_details_view: Final = MappingProxyType[str, object](call_details)
self.mark_logging_complete(event_type="sync_success")
for callback in callbacks:
try:
@ -2972,7 +3076,7 @@ class Logging(LiteLLMLoggingBaseClass):
if callback == "promptlayer" and promptLayerLogger is not None:
print_verbose("reaches promptlayer for logging!")
promptLayerLogger.log_event(
kwargs=self.model_call_details,
kwargs=call_details,
response_obj=result,
start_time=start_time,
end_time=end_time,
@ -2980,7 +3084,7 @@ class Logging(LiteLLMLoggingBaseClass):
)
if callback == "supabase" and supabaseClient is not None:
print_verbose("reaches supabase for logging!")
kwargs = self.model_call_details
kwargs = call_details
# this only logs streaming once, complete_streaming_response exists i.e when stream ends
if self.stream:
@ -2988,10 +3092,10 @@ class Logging(LiteLLMLoggingBaseClass):
continue
else:
print_verbose("reaches supabase for streaming logging!")
result = self._model_call_details_view["complete_streaming_response"]
result = call_details_view["complete_streaming_response"]
model = self._model_call_details_view["model"]
messages = self._model_call_details_view["messages"]
model = call_details_view["model"]
messages = call_details_view["messages"]
optional_params = kwargs.get("optional_params", {})
litellm_params = kwargs.get("litellm_params", {})
supabaseClient.log_event(
@ -3011,7 +3115,7 @@ class Logging(LiteLLMLoggingBaseClass):
if callback == "wandb" and weightsBiasesLogger is not None:
print_verbose("reaches wandb for logging!")
weightsBiasesLogger.log_event(
kwargs=self.model_call_details,
kwargs=call_details,
response_obj=result,
start_time=start_time,
end_time=end_time,
@ -3020,7 +3124,7 @@ class Logging(LiteLLMLoggingBaseClass):
if callback == "logfire" and logfireLogger is not None:
verbose_logger.debug("reaches logfire for success logging!")
kwargs = {}
for k, v in self._model_call_details_view.items():
for k, v in call_details_view.items():
if k != "original_response": # copy.deepcopy raises errors as this could be a coroutine
kwargs[k] = v
@ -3030,10 +3134,10 @@ class Logging(LiteLLMLoggingBaseClass):
continue
else:
print_verbose("reaches logfire for streaming logging!")
result = self._model_call_details_view["complete_streaming_response"]
result = call_details_view["complete_streaming_response"]
logfireLogger.log_event(
kwargs=self.model_call_details,
kwargs=call_details,
response_obj=result,
start_time=start_time,
end_time=end_time,
@ -3044,7 +3148,7 @@ class Logging(LiteLLMLoggingBaseClass):
if callback == "lunary" and lunaryLogger is not None:
print_verbose("reaches lunary for logging!")
model = self.model
kwargs = self.model_call_details
kwargs = call_details
input = kwargs.get("messages", kwargs.get("input", None))
@ -3055,7 +3159,7 @@ class Logging(LiteLLMLoggingBaseClass):
if "complete_streaming_response" not in kwargs:
continue
else:
result = self._model_call_details_view["complete_streaming_response"]
result = call_details_view["complete_streaming_response"]
lunaryLogger.log_event(
type=type,
@ -3065,7 +3169,7 @@ class Logging(LiteLLMLoggingBaseClass):
input=input,
user_id=kwargs.get("user", None),
# user_props=self.model_call_details.get("user_props", None),
extra=self._model_call_details_view.get("optional_params", {}),
extra=call_details_view.get("optional_params", {}),
response_obj=result,
start_time=start_time,
end_time=end_time,
@ -3075,8 +3179,8 @@ class Logging(LiteLLMLoggingBaseClass):
if callback == "helicone" and heliconeLogger is not None:
print_verbose("reaches helicone for logging!")
model = self.model
messages = self._model_call_details_view["input"]
kwargs = self.model_call_details
messages = call_details_view["input"]
kwargs = call_details
# this only logs streaming once, complete_streaming_response exists i.e when stream ends
if self.stream:
@ -3084,7 +3188,7 @@ class Logging(LiteLLMLoggingBaseClass):
continue
else:
print_verbose("reaches helicone for streaming logging!")
result = self._model_call_details_view["complete_streaming_response"]
result = call_details_view["complete_streaming_response"]
heliconeLogger.log_success(
model=model,
@ -3099,7 +3203,7 @@ class Logging(LiteLLMLoggingBaseClass):
global langFuseLogger
print_verbose("reaches langfuse for success logging!")
kwargs = {}
for k, v in self._model_call_details_view.items():
for k, v in call_details_view.items():
if k != "original_response": # copy.deepcopy raises errors as this could be a coroutine
kwargs[k] = v
# this only logs streaming once, complete_streaming_response exists i.e when stream ends
@ -3112,7 +3216,7 @@ class Logging(LiteLLMLoggingBaseClass):
continue
else:
print_verbose("reaches langfuse for streaming logging!")
result = self._model_call_details_view["complete_streaming_response"]
result = call_details_view["complete_streaming_response"]
langfuse_logger_to_use = LangFuseHandler.get_langfuse_logger_for_request(
globalLangfuseLogger=langFuseLogger,
@ -3137,7 +3241,7 @@ class Logging(LiteLLMLoggingBaseClass):
)
if callback == "greenscale" and greenscaleLogger is not None:
kwargs = {}
for k, v in self._model_call_details_view.items():
for k, v in call_details_view.items():
if k != "original_response": # copy.deepcopy raises errors as this could be a coroutine
kwargs[k] = v
# this only logs streaming once, complete_streaming_response exists i.e when stream ends
@ -3150,7 +3254,7 @@ class Logging(LiteLLMLoggingBaseClass):
continue
else:
print_verbose("reaches greenscale for streaming logging!")
result = self._model_call_details_view["complete_streaming_response"]
result = call_details_view["complete_streaming_response"]
greenscaleLogger.log_event(
kwargs=kwargs,
@ -3161,7 +3265,7 @@ class Logging(LiteLLMLoggingBaseClass):
)
if callback == "athina" and athinaLogger is not None:
deep_copy = {}
for k, v in self._model_call_details_view.items():
for k, v in call_details_view.items():
deep_copy[k] = v
athinaLogger.log_event(
kwargs=deep_copy,
@ -3172,7 +3276,7 @@ class Logging(LiteLLMLoggingBaseClass):
)
if callback == "traceloop":
deep_copy = {}
for k, v in self._model_call_details_view.items():
for k, v in call_details_view.items():
if k != "original_response":
deep_copy[k] = v
traceloopLogger.log_event(
@ -3188,11 +3292,11 @@ class Logging(LiteLLMLoggingBaseClass):
if s3Logger is None:
s3Logger = S3Logger()
if self.stream:
if "complete_streaming_response" in self.model_call_details:
if "complete_streaming_response" in call_details:
print_verbose("S3Logger Logger: Got Stream Event - Completed Stream Response")
s3Logger.log_event(
kwargs=self.model_call_details,
response_obj=self.model_call_details["complete_streaming_response"],
kwargs=call_details,
response_obj=call_details["complete_streaming_response"],
start_time=start_time,
end_time=end_time,
print_verbose=print_verbose,
@ -3201,7 +3305,7 @@ class Logging(LiteLLMLoggingBaseClass):
print_verbose("S3Logger Logger: Got Stream Event - No complete stream response as yet")
else:
s3Logger.log_event(
kwargs=self.model_call_details,
kwargs=call_details,
response_obj=result,
start_time=start_time,
end_time=end_time,
@ -3215,19 +3319,17 @@ class Logging(LiteLLMLoggingBaseClass):
openMeterLogger = OpenMeterLogger()
if self.stream and complete_streaming_response is None:
openMeterLogger.log_stream_event(
kwargs=self.model_call_details,
kwargs=call_details,
response_obj=result,
start_time=start_time,
end_time=end_time,
)
else:
if self.stream and complete_streaming_response:
self.model_call_details["complete_response"] = self.model_call_details.get(
"complete_streaming_response", {}
)
result = self._model_call_details_view["complete_response"]
call_details["complete_response"] = call_details.get("complete_streaming_response", {})
result = call_details_view["complete_response"]
openMeterLogger.log_success_event(
kwargs=self.model_call_details,
kwargs=call_details,
response_obj=result,
start_time=start_time,
end_time=end_time,
@ -3240,20 +3342,18 @@ class Logging(LiteLLMLoggingBaseClass):
): # custom logger class
if self.stream and complete_streaming_response is None:
callback.log_stream_event(
kwargs=self.model_call_details,
kwargs=call_details,
response_obj=result,
start_time=start_time,
end_time=end_time,
)
else:
if self.stream and complete_streaming_response:
self.model_call_details["complete_response"] = self.model_call_details.get(
"complete_streaming_response", {}
)
result = self._model_call_details_view["complete_response"]
call_details["complete_response"] = call_details.get("complete_streaming_response", {})
result = call_details_view["complete_response"]
callback.log_success_event(
kwargs=self.model_call_details,
kwargs=call_details,
response_obj=result,
start_time=start_time,
end_time=end_time,
@ -3264,7 +3364,7 @@ class Logging(LiteLLMLoggingBaseClass):
print_verbose(f"success callbacks: Running Custom Callback Function - {callback}")
customLogger.log_event(
kwargs=self.model_call_details,
kwargs=call_details,
response_obj=result,
start_time=start_time,
end_time=end_time,
@ -4102,8 +4202,18 @@ class Logging(LiteLLMLoggingBaseClass):
start_time,
end_time,
cache_hit,
for_async_request=True,
)
def handle_sync_cache_hit_success_callbacks(
self,
result: object,
start_time: datetime.datetime,
end_time: datetime.datetime,
cache_hit: bool | None = None,
) -> None:
executor.submit(self.success_handler, result, start_time, end_time, cache_hit)
def _should_run_sync_callbacks_for_async_calls(self) -> bool:
"""
Returns:
@ -4115,7 +4225,10 @@ class Logging(LiteLLMLoggingBaseClass):
)
_filtered_success_callbacks = self._remove_internal_custom_logger_callbacks(_combined_sync_callbacks)
_filtered_success_callbacks = self._remove_internal_litellm_callbacks(_filtered_success_callbacks)
return len(_filtered_success_callbacks) > 0
return any(
not (isinstance(callback, str) and callback in INTERNAL_CALLBACK_NAMES)
for callback in _filtered_success_callbacks # pyright: ignore[reportUnknownVariableType] # callback lists are untyped
)
def _should_run_sync_failure_callbacks_for_async_calls(self) -> bool:
"""

View file

@ -251,14 +251,7 @@ async def convert_to_streaming_response_async(
setattr(
model_response_object,
"usage",
Usage(
completion_tokens=response_object["usage"].get("completion_tokens", 0),
prompt_tokens=response_object["usage"].get("prompt_tokens", 0),
total_tokens=response_object["usage"].get("total_tokens", 0),
completion_tokens_details=_cached_completion_tokens_details(
_CACHED_USAGE_DETAILS.validate_python(response_object["usage"])
),
),
_replayed_usage(_CACHED_USAGE_DETAILS.validate_python(response_object["usage"])),
)
if "id" in response_object:
@ -345,12 +338,10 @@ def convert_to_streaming_response(
model_response_object.choices = choice_list
if "usage" in response_object and response_object["usage"] is not None:
setattr(model_response_object, "usage", Usage())
model_response_object.usage.completion_tokens = response_object["usage"].get("completion_tokens", 0)
model_response_object.usage.prompt_tokens = response_object["usage"].get("prompt_tokens", 0)
model_response_object.usage.total_tokens = response_object["usage"].get("total_tokens", 0)
model_response_object.usage.completion_tokens_details = _cached_completion_tokens_details(
_CACHED_USAGE_DETAILS.validate_python(response_object["usage"])
setattr(
model_response_object,
"usage",
_replayed_usage(_CACHED_USAGE_DETAILS.validate_python(response_object["usage"])),
)
if "id" in response_object:
@ -414,13 +405,28 @@ class _CachedMessageProviderFields(TypedDict, total=False):
class _CachedUsageDetails(TypedDict, total=False):
completion_tokens: ReadOnly[object]
prompt_tokens: ReadOnly[object]
total_tokens: ReadOnly[object]
completion_tokens_details: ReadOnly[object]
prompt_tokens_details: ReadOnly[object]
cache_read_input_tokens: ReadOnly[object]
cache_creation_input_tokens: ReadOnly[object]
class _PromptCacheTokenCounts(TypedDict, total=False):
cache_read_input_tokens: ReadOnly[int]
cache_creation_input_tokens: ReadOnly[int]
_CACHED_MESSAGE_PROVIDER_FIELDS: Final = TypeAdapter(_CachedMessageProviderFields)
_CACHED_USAGE_DETAILS: Final = TypeAdapter(_CachedUsageDetails)
_FIELD_MAP: Final = TypeAdapter(dict[str, object])
_COMPLETION_TOKENS_DETAILS: Final = TypeAdapter(CompletionTokensDetailsWrapper)
_PROMPT_TOKENS_DETAILS: Final = TypeAdapter(PromptTokensDetailsWrapper)
_PROMPT_CACHE_TOKEN_COUNTS: Final = TypeAdapter(_PromptCacheTokenCounts)
_TOKEN_COUNT: Final = TypeAdapter(int)
_PROMPT_CACHE_TOKEN_FIELDS: Final = ("cache_read_input_tokens", "cache_creation_input_tokens")
def _cached_provider_specific_fields(message: _CachedMessageProviderFields) -> Mapping[str, object] | None:
@ -438,6 +444,44 @@ def _cached_completion_tokens_details(usage: _CachedUsageDetails) -> CompletionT
return None
def _cached_prompt_tokens_details(usage: _CachedUsageDetails) -> PromptTokensDetailsWrapper | None:
details: Final = usage.get("prompt_tokens_details")
if not isinstance(details, dict):
return None
try:
return _PROMPT_TOKENS_DETAILS.validate_python(details)
except ValidationError:
return None
def _cached_token_count(count: object) -> int | None:
try:
return _TOKEN_COUNT.validate_python(count)
except ValidationError:
return None
def _cached_prompt_cache_token_counts(usage: _CachedUsageDetails) -> _PromptCacheTokenCounts:
return _PROMPT_CACHE_TOKEN_COUNTS.validate_python(
{
field: count
for field in _PROMPT_CACHE_TOKEN_FIELDS
if (count := _cached_token_count(usage.get(field))) is not None
}
)
def _replayed_usage(usage: _CachedUsageDetails) -> Usage:
return Usage(
completion_tokens=_cached_token_count(usage.get("completion_tokens")) or 0,
prompt_tokens=_cached_token_count(usage.get("prompt_tokens")) or 0,
total_tokens=_cached_token_count(usage.get("total_tokens")) or 0,
completion_tokens_details=_cached_completion_tokens_details(usage),
prompt_tokens_details=_cached_prompt_tokens_details(usage),
**_cached_prompt_cache_token_counts(usage),
)
def handle_invalid_parallel_tool_calls(
tool_calls: list[ChatCompletionMessageToolCall | ChatCompletionMessageCustomToolCall],
) -> list[ChatCompletionMessageToolCall | ChatCompletionMessageCustomToolCall] | None:

View file

@ -1,10 +1,10 @@
import asyncio
import json
import traceback
from collections.abc import Coroutine, Mapping, Sequence
from collections.abc import Coroutine, Mapping, MutableMapping, Sequence
from dataclasses import dataclass
from enum import Enum, auto
from typing import TYPE_CHECKING, Any, Final, NoReturn, Protocol, TypedDict, cast
from typing import TYPE_CHECKING, Any, Final, NoReturn, Protocol, TypedDict, cast, runtime_checkable
from pydantic import TypeAdapter
from typing_extensions import ReadOnly
@ -12,6 +12,7 @@ from typing_extensions import ReadOnly
import litellm
from litellm._logging import redact_internal_details_from_client_message, verbose_logger
from litellm.constants import (
BEDROCK_REALTIME_PENDING_OPENING_MESSAGE_SCOPE_KEY,
REALTIME_METERED_USAGE_KEY,
REALTIME_SESSION_FAILURE_LOGGED_KEY,
REALTIME_SESSION_SUCCESS_LOGGED_KEY,
@ -98,6 +99,16 @@ class ScopedWebSocket(Protocol):
def scope(self) -> _ASGIScope: ...
@runtime_checkable
class _ReplayScopeWebSocket(Protocol):
@property
def scope(self) -> MutableMapping[str, object]: ...
def client_websocket_scope(websocket: object) -> MutableMapping[str, object] | None:
return websocket.scope if isinstance(websocket, _ReplayScopeWebSocket) else None
class _ClientWebSocket(ScopedWebSocket, Protocol):
async def send_text(self, data: str) -> None: ...
async def receive_text(self) -> str: ...
@ -154,6 +165,7 @@ class RealTimeStreaming:
logging_worker: _LoggingWorker = GLOBAL_LOGGING_WORKER,
):
self.websocket: _ClientWebSocket = websocket
self._replay_scope: Final = client_websocket_scope(websocket)
self.backend_ws = backend_ws
self.logging_obj = logging_obj
self._logging_worker = logging_worker
@ -1346,6 +1358,13 @@ class RealTimeStreaming:
return item
async def _receive_client_message(self) -> str | None:
refused_opening_message: Final = (
None
if self._replay_scope is None
else self._replay_scope.pop(BEDROCK_REALTIME_PENDING_OPENING_MESSAGE_SCOPE_KEY, None)
)
if isinstance(refused_opening_message, str):
return refused_opening_message
try:
return await self.websocket.receive_text()
except Exception as e: # noqa: BLE001 # whatever the client socket raises, the client is gone

View file

@ -8,7 +8,7 @@ import asyncio
import contextlib
import importlib.metadata
import json
from collections.abc import AsyncIterator, Awaitable, Callable, Mapping, MutableMapping
from collections.abc import AsyncIterator, Awaitable, Callable, Mapping, MutableMapping, Sequence
from dataclasses import dataclass
from types import MappingProxyType
from typing import Final, NoReturn, Protocol, runtime_checkable
@ -21,7 +21,7 @@ from litellm._uuid import uuid
from litellm.constants import (
BEDROCK_REALTIME_CLOSE_TIMEOUT_SECONDS,
BEDROCK_REALTIME_COMMITTED_FAILURE_SCOPE_KEY,
BEDROCK_REALTIME_PENDING_SESSION_UPDATE_SCOPE_KEY,
BEDROCK_REALTIME_PENDING_OPENING_MESSAGE_SCOPE_KEY,
BEDROCK_REALTIME_SDK_DISTRIBUTION,
BEDROCK_REALTIME_SDK_SUPPORTED_RANGE,
BEDROCK_REALTIME_SESSION_COMMITTED_SCOPE_KEY,
@ -37,9 +37,8 @@ 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
from .transformation import BedrockRealtimeConfig, BedrockSessionUpdateRejection, requested_modalities
_CLIENT_MODALITIES_ADAPTER: Final[TypeAdapter["list[str] | None"]] = TypeAdapter(list[str] | None)
_CLIENT_MESSAGE_ADAPTER: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue)
_EMPTY_JSON_OBJECT: Final[Mapping[str, JsonValue]] = MappingProxyType({})
@ -177,25 +176,35 @@ async def _client_messages(client_ws: RealtimeClientWebSocket, initial_message:
return
def _pending_session_update(scope: Mapping[str, object]) -> str | None:
"""A fallback attempt on the same websocket replays the session.update the failed attempt never acked."""
if scope.get(BEDROCK_REALTIME_SESSION_COMMITTED_SCOPE_KEY) is True:
committed_failure: Final = scope.get(BEDROCK_REALTIME_COMMITTED_FAILURE_SCOPE_KEY)
raise BedrockError(
status_code=400,
message=(
"Bedrock realtime session already committed to a provider stream; it cannot be replayed"
+ (f". The committed stream failed with: {committed_failure}" if committed_failure else "")
),
)
pending: Final = scope.get(BEDROCK_REALTIME_PENDING_SESSION_UPDATE_SCOPE_KEY)
def _is_session_committed(scope: Mapping[str, object]) -> bool:
return scope.get(BEDROCK_REALTIME_SESSION_COMMITTED_SCOPE_KEY) is True
def raise_if_session_committed(scope: Mapping[str, object] | None) -> None:
"""A fallback attempt on a websocket whose Bedrock stream already took input would silently start over."""
if scope is None or not _is_session_committed(scope):
return
committed_failure: Final = scope.get(BEDROCK_REALTIME_COMMITTED_FAILURE_SCOPE_KEY)
raise BedrockError(
status_code=400,
message=(
"Bedrock realtime session already committed to a provider stream; it cannot be replayed"
+ (f". The committed stream failed with: {committed_failure}" if committed_failure else "")
),
)
def _pending_opening_message(scope: Mapping[str, object]) -> str | None:
"""A fallback attempt on the same websocket replays the opening message Bedrock never accepted a stream for."""
raise_if_session_committed(scope)
pending: Final = scope.get(BEDROCK_REALTIME_PENDING_OPENING_MESSAGE_SCOPE_KEY)
return pending if isinstance(pending, str) else None
def _raise_provider_failure(scope: MutableMapping[str, object], failure: BaseException) -> NoReturn:
error: Final = _as_bedrock_error(failure)
verbose_proxy_logger.error("Bedrock Realtime: provider stream failed: %s", redact_string(str(error)))
if scope.get(BEDROCK_REALTIME_SESSION_COMMITTED_SCOPE_KEY) is True:
if _is_session_committed(scope):
scope[BEDROCK_REALTIME_COMMITTED_FAILURE_SCOPE_KEY] = redact_string(str(error))
raise error from failure
@ -207,26 +216,36 @@ def _parse_client_message(message: str) -> Mapping[str, JsonValue]:
return _EMPTY_JSON_OBJECT
async def _ack_session_update(
async def _send_opening_message(
client_ws: RealtimeClientWebSocket,
bedrock_stream: BedrockBidirectionalStream,
message: str,
bedrock_messages: Sequence[str],
send_to_bedrock: Callable[[str], Awaitable[None]],
) -> None:
"""Keep the opening message for a fallback to replay until Bedrock accepts the stream it opened."""
client_ws.scope[BEDROCK_REALTIME_PENDING_OPENING_MESSAGE_SCOPE_KEY] = message # rebind-ok: scope outlives attempts
for bedrock_message in bedrock_messages:
await send_to_bedrock(bedrock_message)
await bedrock_stream.await_output()
client_ws.scope.pop(BEDROCK_REALTIME_PENDING_OPENING_MESSAGE_SCOPE_KEY, None)
client_ws.scope[BEDROCK_REALTIME_SESSION_COMMITTED_SCOPE_KEY] = True # rebind-ok: scope outlives the attempt
async def _ack_session_update(
client_ws: RealtimeClientWebSocket,
transformation_config: BedrockRealtimeConfig,
model: str,
logging_obj: LiteLLMLogging | None,
parsed_client_message: Mapping[str, JsonValue],
) -> bool:
"""Ack the client's session.update once Bedrock accepted the stream; False means the client is gone."""
await bedrock_stream.await_output()
client_ws.scope.pop(BEDROCK_REALTIME_PENDING_SESSION_UPDATE_SCOPE_KEY, None)
client_ws.scope[BEDROCK_REALTIME_SESSION_COMMITTED_SCOPE_KEY] = True # rebind-ok: scope outlives the attempt
"""Ack the client's session.update; False means the client is gone."""
if logging_obj is None:
return True
requested_modalities: Final = _CLIENT_MODALITIES_ADAPTER.validate_python(
_json_dict(parsed_client_message.get("session")).get("modalities")
)
modalities: Final = requested_modalities(_json_dict(parsed_client_message.get("session")))
try:
await client_ws.send_text(
json.dumps(transformation_config.session_updated_event(model, logging_obj, requested_modalities))
json.dumps(transformation_config.session_updated_event(model, logging_obj, modalities))
)
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)
@ -314,7 +333,7 @@ class BedrockRealtime(BaseAWSLLM):
except ImportError as e:
raise _sdk_import_error(self._sdk_version_lookup(), e) from e
pending_session_update: Final = _pending_session_update(websocket.scope)
pending_opening_message: Final = _pending_opening_message(websocket.scope)
# Get AWS region
if aws_region_name is None:
@ -376,7 +395,7 @@ class BedrockRealtime(BaseAWSLLM):
try:
await self._run_session(
websocket, open_bidirectional_stream, model, logging_obj, pending_session_update, drop_params
websocket, open_bidirectional_stream, model, logging_obj, pending_opening_message, drop_params
)
finally:
await _close_bedrock_client(bedrock_client)
@ -387,7 +406,7 @@ class BedrockRealtime(BaseAWSLLM):
open_bidirectional_stream: Callable[[], Awaitable[BedrockBidirectionalStream]],
model: str,
logging_obj: LiteLLMLogging,
pending_session_update: str | None,
pending_opening_message: str | None,
drop_params: bool = False,
) -> None:
transformation_config: Final = BedrockRealtimeConfig(drop_params=drop_params)
@ -396,7 +415,7 @@ class BedrockRealtime(BaseAWSLLM):
verbose_proxy_logger.debug("Bedrock Realtime: Bidirectional stream established")
if pending_session_update is None:
if pending_opening_message is None:
await websocket.send_text(json.dumps(transformation_config.session_created_event(model, logging_obj)))
verbose_proxy_logger.debug("Bedrock Realtime: sent session.created to client on connect")
@ -418,7 +437,7 @@ class BedrockRealtime(BaseAWSLLM):
model,
session_state,
logging_obj,
initial_message=pending_session_update,
initial_message=pending_opening_message,
)
logged_events: Final = (
@ -531,18 +550,22 @@ class BedrockRealtime(BaseAWSLLM):
if not await _reject_session_update(client_ws, parsed_client_message, rejection):
break
continue
client_ws.scope[BEDROCK_REALTIME_PENDING_SESSION_UPDATE_SCOPE_KEY] = message
transformed_messages = transformation_config.transform_realtime_request(
message=message,
model=model,
session_configuration_request=session_state.get("session_configuration_request"),
)
for bedrock_message in transformed_messages:
await send_to_bedrock(bedrock_message)
if transformed_messages and not _is_session_committed(client_ws.scope):
await _send_opening_message(
client_ws, bedrock_stream, message, transformed_messages, send_to_bedrock
)
else:
for bedrock_message in transformed_messages:
await send_to_bedrock(bedrock_message)
if is_session_update and not await _ack_session_update(
client_ws, bedrock_stream, transformation_config, model, logging_obj, parsed_client_message
client_ws, transformation_config, model, logging_obj, parsed_client_message
):
break
finally:

View file

@ -123,26 +123,81 @@ def _integer_setting(value: JsonValue) -> int | None:
return None if number is None else int(number)
def _malformed_number_settings(session_config: Mapping[str, JsonValue]) -> tuple[str, ...]:
def requested_modalities(session_config: Mapping[str, JsonValue]) -> tuple[str, ...] | None:
modalities: Final = session_config.get("modalities")
if not isinstance(modalities, list):
return None
names: Final = tuple(modality for modality in modalities if isinstance(modality, str))
return names if len(names) == len(modalities) else None
_OPENAI_VOICES: Final = frozenset(
("alloy", "ash", "ballad", "coral", "echo", "sage", "shimmer", "verse", "marin", "cedar")
)
def _unplayable_voice(session_config: Mapping[str, JsonValue]) -> str | None:
voice: Final = session_config.get("voice")
if not isinstance(voice, str):
return None
return voice if not voice.strip() or voice.lower() in _OPENAI_VOICES else None
def _malformed_settings(session_config: Mapping[str, JsonValue]) -> tuple[tuple[str, str], ...]:
instructions: Final = session_config.get("instructions")
voice: Final = session_config.get("voice")
input_audio_format: Final = session_config.get("input_audio_format")
output_audio_format: Final = session_config.get("output_audio_format")
modalities: Final = session_config.get("modalities")
temperature: Final = session_config.get("temperature")
max_output_tokens: Final = session_config.get("max_response_output_tokens")
input_rate_hertz: Final = session_config.get("input_sample_rate_hertz")
output_rate_hertz: Final = session_config.get("output_sample_rate_hertz")
tools: Final = session_config.get("tools")
malformed: Final = (
("temperature", temperature is not None and _number_setting(temperature) is None),
("instructions", "a string", instructions is not None and not isinstance(instructions, str)),
("voice", "a string", voice is not None and not isinstance(voice, str)),
("tools", "a list", tools is not None and not isinstance(tools, list)),
(
"input_audio_format",
"a string",
input_audio_format is not None and not isinstance(input_audio_format, str),
),
(
"output_audio_format",
"a string",
output_audio_format is not None and not isinstance(output_audio_format, str),
),
(
"modalities",
"a list of strings",
modalities is not None and requested_modalities(session_config) is None,
),
("temperature", "a number", temperature is not None and _number_setting(temperature) is None),
(
"max_response_output_tokens",
"a number or 'inf'",
max_output_tokens is not None
and max_output_tokens != "inf"
and _integer_setting(max_output_tokens) is None,
),
("output_sample_rate_hertz", output_rate_hertz is not None and _integer_setting(output_rate_hertz) is None),
(
"input_sample_rate_hertz",
"a number",
input_rate_hertz is not None and _integer_setting(input_rate_hertz) is None,
),
(
"output_sample_rate_hertz",
"a number",
output_rate_hertz is not None and _integer_setting(output_rate_hertz) is None,
),
)
return tuple(name for name, is_malformed in malformed if is_malformed)
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", "session_update_not_applied"]
code: Literal["invalid_type", "invalid_value", "session_update_not_applied"]
message: str
param: str
@ -312,7 +367,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig):
return 8000 # G.711 typically uses 8kHz
return 24000 if is_output else 16000
def transform_session_update_event(self, json_message: dict) -> list[str]:
def transform_session_update_event(self, json_message: Mapping[str, JsonValue]) -> list[str]:
"""
Transform session.update event to Bedrock session configuration.
@ -325,7 +380,10 @@ class BedrockRealtimeConfig(BaseRealtimeConfig):
verbose_logger.debug("Handling session.update")
messages: Final[list[str]] = []
session_config: Final = json_message.get("session", {})
requested_session: Final = json_message.get("session")
session_config: Final[Mapping[str, JsonValue]] = (
requested_session if isinstance(requested_session, dict) else {}
)
locked_settings: Final = self._apply_locked_settings(session_config, self._started_settings)
# Settings Bedrock locks at prompt start may only change before it
@ -336,13 +394,14 @@ class BedrockRealtimeConfig(BaseRealtimeConfig):
self.output_sample_rate_hertz = locked_settings.output_sample_rate_hertz
# Update audio input configuration from session if provided
if "input_audio_format" in session_config:
input_format: Final = session_config["input_audio_format"]
input_format: Final = session_config.get("input_audio_format")
if isinstance(input_format, str):
self.input_sample_rate_hertz = self._map_audio_format_to_sample_rate(input_format, is_output=False)
# Allow direct override of the input sample rate if provided (custom extension)
if "input_sample_rate_hertz" in session_config:
self.input_sample_rate_hertz = session_config["input_sample_rate_hertz"]
requested_input_rate: Final = _integer_setting(session_config.get("input_sample_rate_hertz"))
if requested_input_rate is not None:
self.input_sample_rate_hertz = requested_input_rate
if self.prompt_started:
return messages
@ -352,7 +411,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig):
messages.append(json.dumps(session_start))
# Send prompt start
tools: Final[Sequence[JsonValue] | None] = session_config.get("tools")
tools: Final = session_config.get("tools")
prompt_start_config: Final[Mapping[str, object]] = {
"promptName": self.prompt_name,
"textOutputConfiguration": {"mediaType": self.text_media_type},
@ -370,7 +429,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig):
"toolUseOutputConfiguration": {"mediaType": "application/json"},
"toolConfiguration": {"tools": self._transform_tools_to_bedrock_format(tools)},
}
if tools
if isinstance(tools, list) and tools
else {}
),
}
@ -445,12 +504,14 @@ class BedrockRealtimeConfig(BaseRealtimeConfig):
if isinstance(requested_tools, list)
else base.tools_json
)
system_prompt: Final = base.system_prompt if instructions is None else str(instructions)
voice_id: Final = base.voice_id if voice is None else str(voice)
system_prompt: Final = instructions if isinstance(instructions, str) else base.system_prompt
voice_id: Final = (
voice if isinstance(voice, str) and _unplayable_voice(session_config) is None else base.voice_id
)
format_sample_rate: Final = (
base.output_sample_rate_hertz
if output_audio_format is None
else self._map_audio_format_to_sample_rate(str(output_audio_format), is_output=True)
self._map_audio_format_to_sample_rate(output_audio_format, is_output=True)
if isinstance(output_audio_format, str)
else base.output_sample_rate_hertz
)
requested_output_rate: Final = _integer_setting(session_config.get("output_sample_rate_hertz"))
output_sample_rate_hertz: Final = (
@ -476,16 +537,26 @@ class BedrockRealtimeConfig(BaseRealtimeConfig):
session_config: Final = json_message.get("session")
if not isinstance(session_config, Mapping):
return None
malformed: Final = () if self.drop_params else _malformed_number_settings(session_config)
malformed: Final = () if self.drop_params else _malformed_settings(session_config)
if malformed:
name, expected = malformed[0]
return BedrockSessionUpdateRejection(
code="invalid_type",
message=(
f"Invalid type for 'session.{malformed[0]}': Bedrock Nova Sonic needs a number for temperature "
"and output_sample_rate_hertz, and a number or 'inf' for max_response_output_tokens. "
f"Invalid type for 'session.{name}': Bedrock Nova Sonic needs {expected} there. "
"Set drop_params to drop it instead"
),
param=f"session.{malformed[0]}",
param=f"session.{name}",
)
unplayable_voice: Final = None if self.drop_params else _unplayable_voice(session_config)
if unplayable_voice is not None:
return BedrockSessionUpdateRejection(
code="invalid_value",
message=(
f"Invalid value for 'session.voice': '{unplayable_voice}' is not a Bedrock Nova Sonic voice. "
"Set drop_params to drop it instead"
),
param="session.voice",
)
if not self.prompt_started:
return None
@ -785,7 +856,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig):
List of Bedrock format messages (JSON strings)
"""
try:
json_message: Final[dict[str, object]] = json.loads(message)
json_message: Final[dict[str, JsonValue]] = json.loads(message)
except json.JSONDecodeError:
verbose_logger.warning("Invalid JSON message: %s", message[:200])
return []
@ -816,11 +887,11 @@ class BedrockRealtimeConfig(BaseRealtimeConfig):
self,
model: str,
logging_obj: LiteLLMLoggingObj,
modalities: list[str] | None = None,
modalities: Sequence[str] | None = None,
) -> OpenAIRealtimeStreamSession:
session: Final = OpenAIRealtimeStreamSession(
id=logging_obj.litellm_trace_id,
modalities=modalities if modalities is not None else ["text", "audio"],
modalities=list(modalities) if modalities is not None else ["text", "audio"],
)
if model is not None and isinstance(model, str):
session["model"] = model
@ -842,7 +913,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig):
self,
model: str,
logging_obj: LiteLLMLoggingObj,
modalities: list[str] | None = None,
modalities: Sequence[str] | None = None,
) -> OpenAIRealtimeStreamSessionEvents:
"""Build the OpenAI session.updated ack reflecting the client's requested modalities."""
return OpenAIRealtimeStreamSessionEvents(

View file

@ -75,6 +75,7 @@ from litellm.constants import (
DEFAULT_SHARED_HEALTH_CHECK_LOCK_TTL,
DEFAULT_SHARED_HEALTH_CHECK_TTL,
DEFAULT_SLACK_ALERTING_THRESHOLD,
INTERNAL_CALLBACK_NAMES,
LANGFUSE_SHUTDOWN_FLUSH_TIMEOUT_MILLIS,
LITELLM_EMBEDDING_PROVIDERS_SUPPORTING_INPUT_ARRAY_OF_TOKENS,
LITELLM_SETTINGS_SAFE_DB_OVERRIDES,
@ -19812,8 +19813,7 @@ def _hidden_runtime_callback_names(configured_callback_names: frozenset[str]) ->
if name in CustomLoggerRegistry.CALLBACK_CLASS_STR_TO_CLASS_TYPE
)
configured_modules: Final = frozenset(name.rsplit(".", 1)[0] for name in configured_callback_names if "." in name)
internal_callback_names: Final = frozenset({"cache", "vector_store_pre_call_hook"})
return internal_callback_names | frozenset(
return INTERNAL_CALLBACK_NAMES | frozenset(
callback_name
for callback_name, callback in litellm.logging_callback_manager.get_callback_objects()
if isinstance(callback, CustomGuardrail)

View file

@ -37,9 +37,10 @@ from ..litellm_core_utils.core_helpers import normalize_drop_params
from ..litellm_core_utils.credential_accessor import CredentialAccessor
from ..litellm_core_utils.get_litellm_params import get_litellm_params
from ..litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
from ..litellm_core_utils.realtime_streaming import client_websocket_scope
from ..llms.azure.common_utils import get_azure_ad_token
from ..llms.azure.realtime.handler import AzureOpenAIRealtime, azure_realtime_protocol_for_client
from ..llms.bedrock.realtime.handler import BedrockRealtime
from ..llms.bedrock.realtime.handler import BedrockRealtime, raise_if_session_committed
from ..llms.custom_httpx.http_handler import realtime_ssl_for_url
from ..llms.openai.live.handler import OpenAILiveSessions
from ..llms.openai.realtime.handler import OpenAIRealtime
@ -451,6 +452,7 @@ async def _arealtime(
For PROXY use only.
"""
raise_if_session_committed(client_websocket_scope(websocket))
headers = cast(dict | None, kwargs.get("headers"))
extra_headers: Final = cast(dict | None, kwargs.get("extra_headers"))
if headers is None:

View file

@ -4,8 +4,10 @@ Check if prompt caching is valid for a given deployment
Route to previously cached model id, if valid
"""
from collections.abc import Callable
from typing import Final, cast
from collections.abc import Callable, Mapping
from typing import Final, TypedDict, cast
from typing_extensions import ReadOnly
from litellm import verbose_logger
from litellm.caching.dual_cache import DualCache
@ -16,7 +18,7 @@ from litellm.integrations.anthropic_cache_control_hook import (
)
from litellm.integrations.custom_logger import CustomLogger, Span
from litellm.litellm_core_utils.token_counter import offload_token_count
from litellm.llms.bedrock.common_utils import supports_bedrock_prompt_cache_ttl
from litellm.llms.bedrock.common_utils import bedrock_cache_point_ttl_supported, bedrock_capability_model
from litellm.types.llms.openai import AllMessageValues
from litellm.types.utils import CallTypes, StandardLoggingPayload
from litellm.utils import get_prompt_cache_min_tokens, is_prompt_caching_valid_prompt
@ -24,6 +26,11 @@ from litellm.utils import get_prompt_cache_min_tokens, is_prompt_caching_valid_p
from ..prompt_caching_cache import PromptCachingCache
class _PromptCachingSuccessEventKwargs(TypedDict, total=False):
standard_logging_object: ReadOnly[StandardLoggingPayload | None]
litellm_params: ReadOnly[Mapping[str, object] | None]
def _get_min_token_count_for_deployments(healthy_deployments: list[dict]) -> int:
"""
Returns the lowest minimum cacheable prefix across a model group.
@ -49,9 +56,13 @@ def _get_min_token_count_for_deployments(healthy_deployments: list[dict]) -> int
)
def _honors_one_hour_cache_ttl(standard_logging_object: StandardLoggingPayload) -> bool:
def _honors_one_hour_cache_ttl(
standard_logging_object: StandardLoggingPayload, litellm_params: Mapping[str, object]
) -> bool:
if standard_logging_object["custom_llm_provider"] == "bedrock":
return supports_bedrock_prompt_cache_ttl(standard_logging_object["model"])
deployment_model: Final = standard_logging_object["model"]
capability_model: Final = bedrock_capability_model(deployment_model, litellm_params.get("base_model"))
return bedrock_cache_point_ttl_supported(capability_model, deployment_model)
model_map_value: Final = standard_logging_object["model_map_information"]["model_map_value"]
return model_map_value is not None and (model_map_value.get("cache_creation_input_token_cost_above_1hr") or 0) > 0
@ -121,7 +132,13 @@ class PromptCachingDeploymentCheck(CustomLogger):
return healthy_deployments
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
async def async_log_success_event(
self,
kwargs: _PromptCachingSuccessEventKwargs,
response_obj: object,
start_time: object,
end_time: object,
) -> None:
standard_logging_object: Final[StandardLoggingPayload | None] = kwargs.get("standard_logging_object", None)
if standard_logging_object is None:
@ -163,11 +180,15 @@ class PromptCachingDeploymentCheck(CustomLogger):
cache: Final = PromptCachingCache(
cache=self.cache,
)
raw_litellm_params: Final = kwargs.get("litellm_params")
litellm_params: Final[Mapping[str, object]] = (
raw_litellm_params if isinstance(raw_litellm_params, Mapping) else {}
)
await cache.async_add_model_id(
model_id=model_id,
messages=prompt_messages,
tools=None, # [TODO]: add tools once standard_logging_object supports it
honors_one_hour_ttl=_honors_one_hour_cache_ttl(standard_logging_object),
honors_one_hour_ttl=_honors_one_hour_cache_ttl(standard_logging_object, litellm_params),
)
return

View file

@ -151,6 +151,8 @@ class LoggingSurface(Protocol):
start_time: datetime.datetime | None = None,
end_time: datetime.datetime | None = None,
cache_hit: bool | None = None,
*,
for_async_request: bool = False,
**kwargs: object,
) -> None: ...

View file

@ -5936,6 +5936,15 @@ def _get_model_info_from_generalization(
return None
def unmapped_model_generalization(model: str, custom_llm_provider: str | None) -> Mapping[str, object] | None:
generalization: Final = _get_model_info_from_generalization(
model=model,
potential_model_names=get_potential_model_names(model, custom_llm_provider),
custom_llm_provider=custom_llm_provider,
)
return generalization[1] if generalization is not None else None
def _strip_mantle_region_prefix(model: str) -> str:
from litellm.llms.bedrock_mantle.common_utils import split_mantle_region_prefix

View file

@ -13,19 +13,26 @@ deployment that fails a handful of times in a row is benched (a bad-credential 4
included), the router answers "No deployments available" without retrying, and the
retry loop that leaks stops running; cooldown_time 0 keeps it running.
Two identical phases run back to back. The first is the warmup that grows the
proxy's caches and allocator arenas to their steady state, the second is the one
the budget applies to, so a healthy proxy shows the second phase adding roughly
nothing while a leaking one adds a fixed amount per request. RSS is read through
/debug/memory/summary on every configured replica; a burst of failing calls leaves
a transient bulge of garbage that gc reclaims within seconds, so each checkpoint
samples until no new worker has answered for a settle window and keeps the lowest
reading per worker. The growth is judged per worker (by replica address, hostname
and pid, since pods in their own pid namespaces report the same pids) so each
worker is compared with itself, and the two checkpoints must see the same workers:
a single load-balanced address reaches the workers behind it one answer at a time,
and a worker that answered only one checkpoint would otherwise drop out of the
comparison, which is where a leaking worker could hide.
Three identical batches run back to back. The first two are the warmup that grows
the proxy's caches and allocator arenas to their steady state, the third is the one
the budget applies to, so a healthy proxy shows the third batch adding roughly
nothing while a leaking one adds a fixed amount per request. One warmup batch is not
enough: each batch's in-flight retries peak well above the live heap, the allocator
keeps that high-water mark, and on a cold worker the second batch still set a new one
(about 50 MB of RSS with the live Python heap flat to within 1 MB), where the later
batches each settled under 25 MB. RSS is read through /debug/memory/summary on every
configured replica; a burst of failing calls leaves a transient bulge of garbage that
gc reclaims within seconds, so each checkpoint samples until no new worker has
answered for a settle window and keeps the lowest reading per worker. A checkpoint
runs before the warmup and after every batch, and each batch's settled growth per
worker rides out as a repeated `rss_batch_growth` JUnit property as soon as that
batch settles, so a run that fails keeps the curve up to the batch it failed in. The budget is judged per worker
(by replica address, hostname and pid, since pods in their own pid namespaces report
the same pids) so each worker is compared with itself, and the checkpoints around the
measured batch must see the same workers: a single load-balanced address reaches the
workers behind it one answer at a time, and a worker that answered only one
checkpoint would otherwise drop out of the comparison, which is where a leaking
worker could hide.
RSS alone is a coarse gauge: on the release stack (spend logs storing prompts,
json logs, prometheus and otel callbacks) the same v1.100.0 breadcrumbs grew RSS
@ -58,7 +65,7 @@ from __future__ import annotations
import json
import time
from collections.abc import Mapping, Sequence
from collections.abc import Callable, Mapping, Sequence
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass
from types import MappingProxyType
@ -91,6 +98,8 @@ pytestmark = [pytest.mark.e2e, pytest.mark.quiet_stack]
DEPLOYMENTS_PER_GROUP: Final = 2
RSS_SAMPLE_CAP: Final = 4 * MEMORY_RSS_SETTLE_SAMPLES
WARMUP_BATCHES: Final = 2
MEASURED_BATCH: Final = WARMUP_BATCHES + 1
@dataclass(frozen=True, slots=True)
@ -177,12 +186,52 @@ def _settled_rss_per_worker(proxy: ProxyClient) -> Mapping[WorkerKey, RssReading
)
def _record_batch_growth(
record_property: Callable[[str, object], None],
batch: int,
before: Mapping[WorkerKey, RssReading],
after: Mapping[WorkerKey, RssReading],
) -> None:
for growth in (WorkerGrowth(before[worker], after[worker]) for worker in before.keys() & after.keys()):
record_property(
"rss_batch_growth",
json.dumps(
{
"batch": batch,
"worker": growth.warm.where,
"before_mb": growth.warm.ram_usage_mb,
"after_mb": growth.after.ram_usage_mb,
"growth_mb": round(growth.growth_mb, 2),
}
),
)
def _settled_rss_around_measured_batch(
proxy: ProxyClient,
key: str,
model: str,
override: RouterSettingsOverride,
fallback: str,
record_property: Callable[[str, object], None],
before: Mapping[WorkerKey, RssReading],
batch: int,
) -> tuple[Mapping[WorkerKey, RssReading], Mapping[WorkerKey, RssReading]]:
calls: Final = _fail_many(proxy, key, model, override)
_assert_every_call_failed_through_fallback(calls, fallback)
after: Final = _settled_rss_per_worker(proxy)
_record_batch_growth(record_property, batch, before, after)
if batch == MEASURED_BATCH:
return before, after
return _settled_rss_around_measured_batch(proxy, key, model, override, fallback, record_property, after, batch + 1)
def _heaviest_worker_growth(
warm: Mapping[WorkerKey, RssReading], after: Mapping[WorkerKey, RssReading]
) -> WorkerGrowth:
assert warm.keys() == after.keys(), (
f"the workers answering /debug/memory/summary changed between the checkpoints, so not every worker can "
f"be compared with itself: gone after the measured batch {sorted(warm.keys() - after.keys())} (a worker "
f"the workers answering /debug/memory/summary changed around measured batch {MEASURED_BATCH}, so not every "
f"worker can be compared with itself: gone after it {sorted(warm.keys() - after.keys())} (a worker "
f"that died or was restarted under failing traffic, which is what an OOM kill looks like), first seen "
f"after it {sorted(after.keys() - warm.keys())} (the warm window never reached them, so they have no "
f"baseline; raise E2E_MEMORY_RSS_SETTLE_SAMPLES if the stack has more workers than the window covers)"
@ -247,6 +296,7 @@ class TestReliabilityMemory:
f"leaves the worker little room for traffic"
)
@pytest.mark.filterwarnings("ignore:record_property is incompatible with junit_family")
@pytest.mark.covers("reliability.perf.memory.under_slo")
@meta(
Subject(
@ -257,7 +307,11 @@ class TestReliabilityMemory:
)
)
def test_failing_requests_do_not_grow_rss_or_stored_request(
self, client: ComplexityRouterClient, resources: ResourceManager, scoped_key: str
self,
client: ComplexityRouterClient,
resources: ResourceManager,
scoped_key: str,
record_property: Callable[[str, object], None],
) -> None:
marker: Final = unique_marker()
primary: Final = f"reliability-memory-{marker}"
@ -278,19 +332,23 @@ class TestReliabilityMemory:
f"request into the stored snapshot the way the v1.100.0 ones did"
)
warmup: Final = _fail_many(client.proxy, scoped_key, primary, override)
_assert_every_call_failed_through_fallback(warmup, fallback)
warm: Final = _settled_rss_per_worker(client.proxy)
measured: Final = _fail_many(client.proxy, scoped_key, primary, override)
_assert_every_call_failed_through_fallback(measured, fallback)
after: Final = _settled_rss_per_worker(client.proxy)
warm, after = _settled_rss_around_measured_batch(
client.proxy,
scoped_key,
primary,
override,
fallback,
record_property,
_settled_rss_per_worker(client.proxy),
1,
)
heaviest: Final = _heaviest_worker_growth(warm, after)
assert heaviest.growth_mb <= MEMORY_RSS_BUDGET_MB, (
f"proxy RSS grew {heaviest.growth_mb:.1f} MB over a second batch of {MEMORY_REQUESTS_PER_PHASE} failing "
f"requests ({MEMORY_RETRIES_PER_REQUEST} retries each plus a fallback) after an identical warmup batch, "
f"past the {MEMORY_RSS_BUDGET_MB:.0f} MB budget: worker pid {heaviest.warm.worker_pid} on "
f"proxy RSS grew {heaviest.growth_mb:.1f} MB over measured batch {MEASURED_BATCH} of "
f"{MEMORY_REQUESTS_PER_PHASE} failing requests ({MEMORY_RETRIES_PER_REQUEST} retries each plus a fallback) "
f"after {WARMUP_BATCHES} identical warmup batches, past the {MEMORY_RSS_BUDGET_MB:.0f} MB budget: "
f"worker pid {heaviest.warm.worker_pid} on "
f"{heaviest.warm.hostname or 'an unnamed host'} behind {heaviest.warm.replica} settled at "
f"{heaviest.warm.ram_usage_mb:.1f} MB warm and "
f"{heaviest.after.ram_usage_mb:.1f} MB after; failing requests are leaking memory the way the "

View file

@ -43,6 +43,7 @@ FOLLOW_UPS = 3
class TestReliabilityPromptCachingAffinity:
@pytest.mark.provider_live
@pytest.mark.covers("reliability.cache.prompt_caching_model_select.returns_cached")
@meta(
Subject(

View file

@ -14,6 +14,12 @@ text message (the injected trigger-audio block is USER/AUDIO, so it never trips
``tool_result`` fires on the Nth inbound ``toolResult`` event, which the proxy sends inside
a TOOL contentStart/toolResult/contentEnd triple whose ``toolResultInputConfiguration``
``toolUseId`` names the call being answered.
A ``Refusal`` peer withholds each stream's initial response until that stream's opening arrived
(its SYSTEM prompt block or its first USER text block closed), then answers the HTTP error Bedrock
sends for an unknown model id. A ``StreamFailure`` peer accepts the stream, plays its turns, and
then sends an event-stream exception frame and ends the stream, as Bedrock does for a mid-call
model error. ``streams()`` groups the inbound events per opened stream.
"""
from __future__ import annotations
@ -26,22 +32,24 @@ import uuid
from collections.abc import Iterator, Mapping
from contextlib import contextmanager, suppress
from dataclasses import dataclass
from itertools import chain
from itertools import chain, count
from pathlib import Path
from queue import SimpleQueue
from typing import Final, Literal
from typing import Final, Literal, TypeVar
import h2.config
import h2.connection
import h2.events
import h2.exceptions
from integration._support.async_server import Handler, LoopbackServer
from integration._support.process import _free_port
from integration._support.ports import allocate_port
from integration._support.tls import server_context
from integration._support.upstream import aws_event_stream_frame
from pydantic import JsonValue, TypeAdapter
BedrockEvent = Mapping[str, JsonValue]
OpeningEnd = Literal["system_prompt", "user_input"]
Recorded = TypeVar("Recorded")
_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
@ -177,6 +185,20 @@ class Turn:
events: tuple[BedrockEvent, ...]
@dataclass(frozen=True, slots=True)
class Refusal:
after: OpeningEnd
status: int = 400
error_type: str = "ValidationException"
message: str = "The provided model identifier is invalid."
@dataclass(frozen=True, slots=True)
class StreamFailure:
exception_type: str = "modelStreamErrorException"
message: str = "The system encountered an unexpected error during processing. Try your request again."
def _event_body(event: BedrockEvent) -> tuple[str, Mapping[str, JsonValue]]:
body: Final = event["event"]
assert isinstance(body, Mapping) and len(body) == 1, event
@ -214,6 +236,17 @@ def _frame_payload(raw: bytes) -> bytes:
return raw[12 + header_length : -4]
def _exception_frame(failure: StreamFailure) -> bytes:
return aws_event_stream_frame(
{
":message-type": "exception",
":exception-type": failure.exception_type,
":content-type": "application/json",
},
json.dumps({"message": failure.message}).encode(),
)
def _decode_inbound(chunk_payload: bytes) -> BedrockEvent | None:
"""One signed chunk frame's payload -> the inner chunk's {"bytes": ...} event."""
inner_payload: Final = _frame_payload(chunk_payload)
@ -239,6 +272,7 @@ class _Inbound:
self.tool_use_ids: Final[dict[str, str]] = {}
self.tool_results_seen = 0
self.user_text_closed = False
self.system_text_closed = False
def record(self, event: BedrockEvent) -> None:
self.events.put(event)
@ -253,6 +287,8 @@ class _Inbound:
role_type: Final = self.open_blocks.pop(content_name, ("", ""))
if role_type == ("USER", "TEXT"):
self.user_text_closed = True
elif role_type == ("SYSTEM", "TEXT"):
self.system_text_closed = True
elif name == "toolResult":
self.tool_results_seen += 1
@ -283,10 +319,35 @@ class _Script:
return turn.events
def _opening_arrived(after: OpeningEnd, inbound: _Inbound) -> bool:
match after:
case "system_prompt":
return inbound.system_text_closed
case "user_input":
return inbound.user_text_closed
class _PeerState:
def __init__(self, turns: tuple[Turn, ...]) -> None:
def __init__(self, turns: tuple[Turn, ...], refusal: Refusal | None, failure: StreamFailure | None) -> None:
self.inbound: Final = _Inbound()
self.script: Final = _Script(turns)
self.refusal: Final = refusal
self.failure: Final = failure
self.ordinals: Final = count()
self.opened: Final = SimpleQueue[int]()
self.stream_events: Final = SimpleQueue[tuple[int, BedrockEvent]]()
def _drain(queue: SimpleQueue[Recorded], log: list[Recorded]) -> tuple[Recorded, ...]:
while not queue.empty():
log.append(queue.get()) # drained log, not a built value
return tuple(log)
@dataclass(frozen=True, slots=True)
class _Stream:
ordinal: int
inbound: _Inbound
@dataclass(frozen=True, slots=True)
@ -294,11 +355,16 @@ class BedrockRealtimePeer:
api_base: str
_state: _PeerState
_drained: list[BedrockEvent]
_opened: list[int]
_stream_events: list[tuple[int, BedrockEvent]]
def received(self) -> tuple[BedrockEvent, ...]:
while not self._state.inbound.events.empty():
self._drained.append(self._state.inbound.events.get()) # drained log, not a built value
return tuple(self._drained)
return _drain(self._state.inbound.events, self._drained)
def streams(self) -> tuple[tuple[BedrockEvent, ...], ...]:
opened: Final = _drain(self._state.opened, self._opened)
recorded: Final = _drain(self._state.stream_events, self._stream_events)
return tuple(tuple(event for ordinal, event in recorded if ordinal == stream) for stream in opened)
def tool_specs(self) -> tuple[Mapping[str, JsonValue], ...]:
return tuple(chain.from_iterable(_tool_specs(event) for event in self.received()))
@ -344,6 +410,7 @@ class _H2Session:
self._writer: Final = writer
self._h2: Final = h2.connection.H2Connection(config=h2.config.H2Configuration(client_side=False))
self._buffers: Final[dict[int, bytes]] = {}
self._streams: Final[dict[int, _Stream]] = {}
self._ended: set[int] = set()
self._closed = False
@ -367,10 +434,14 @@ class _H2Session:
def _apply(self, event: h2.events.Event) -> None:
if isinstance(event, h2.events.RequestReceived):
self._h2.send_headers(
event.stream_id,
[(":status", "200"), ("content-type", "application/vnd.amazon.eventstream")],
)
ordinal: Final = next(self._state.ordinals)
self._streams[event.stream_id] = _Stream(ordinal, _Inbound())
self._state.opened.put(ordinal)
if self._state.refusal is None:
self._h2.send_headers(
event.stream_id,
[(":status", "200"), ("content-type", "application/vnd.amazon.eventstream")],
)
elif isinstance(event, h2.events.DataReceived):
self._h2.acknowledge_received_data(event.flow_controlled_length, event.stream_id)
self._buffers[event.stream_id] = self._buffers.get(event.stream_id, b"") + event.data
@ -380,15 +451,44 @@ class _H2Session:
decoded: Final = _decode_inbound(payload)
if decoded is None:
continue
self._state.inbound.record(decoded)
if event.stream_id not in self._ended:
for outbound in self._state.script.ready_events(self._state.inbound, decoded):
self._h2.send_data(event.stream_id, _outbound_chunk(outbound))
self._record(event.stream_id, decoded)
elif isinstance(event, (h2.events.StreamEnded, h2.events.StreamReset)):
self._close_stream(event.stream_id)
elif isinstance(event, h2.events.ConnectionTerminated):
self._writer.close()
def _record(self, stream_id: int, decoded: BedrockEvent) -> None:
stream: Final = self._streams[stream_id]
self._state.inbound.record(decoded)
stream.inbound.record(decoded)
self._state.stream_events.put((stream.ordinal, decoded))
if stream_id in self._ended:
return
refusal: Final = self._state.refusal
if refusal is not None:
if _opening_arrived(refusal.after, stream.inbound):
self._refuse(stream_id, refusal)
return
outbound_events: Final = self._state.script.ready_events(self._state.inbound, decoded)
for outbound in outbound_events:
self._h2.send_data(stream_id, _outbound_chunk(outbound))
failure: Final = self._state.failure
if outbound_events and failure is not None and self._state.script.exhausted:
self._h2.send_data(stream_id, _exception_frame(failure))
self._end_stream(stream_id)
def _refuse(self, stream_id: int, refusal: Refusal) -> None:
self._h2.send_headers(
stream_id,
[
(":status", str(refusal.status)),
("content-type", "application/json"),
("x-amzn-errortype", refusal.error_type),
],
)
self._h2.send_data(stream_id, json.dumps({"message": refusal.message}).encode())
self._end_stream(stream_id)
def _close_stream(self, stream_id: int) -> None:
self._end_stream(stream_id)
self._closed = True
@ -411,19 +511,27 @@ def _session_handler(state: _PeerState) -> Handler:
@contextmanager
def bedrock_realtime_peer(turns: tuple[Turn, ...], cert: tuple[Path, Path]) -> Iterator[BedrockRealtimePeer]:
def bedrock_realtime_peer(
turns: tuple[Turn, ...],
cert: tuple[Path, Path],
*,
refusal: Refusal | None = None,
failure: StreamFailure | None = None,
) -> Iterator[BedrockRealtimePeer]:
"""Serve scripted `turns` on a loopback TLS h2 endpoint; yield the peer.
`cert` is the self-signed pair the owned proxy already trusts via SSL_CERT_FILE.
SigV4 headers are ignored; the stream closes cleanly when the client ends it.
"""
port: Final = _free_port()
port: Final = allocate_port()
context: Final = server_context(cert[0], cert[1])
context.set_alpn_protocols(["h2"])
state: Final = _PeerState(turns)
state: Final = _PeerState(turns, refusal, failure)
server: Final = LoopbackServer(port, _session_handler(state), ssl_context=context)
server.start()
try:
yield BedrockRealtimePeer(api_base=f"https://localhost:{port}", _state=state, _drained=[])
yield BedrockRealtimePeer(
api_base=f"https://localhost:{port}", _state=state, _drained=[], _opened=[], _stream_events=[]
)
finally:
server.stop()

View file

@ -628,8 +628,11 @@ def worker_configuration(config: Path, directory: Path, inherited: Mapping[str,
return f"http://127.0.0.1:{inherited['INTEGRATION_PROVIDER_PORT']}"
contents: Final = re.sub(r"http://127[.]0[.]0[.]1:(8190|8191)(?![0-9])", endpoint, config.read_text())
directory.mkdir(parents=True, exist_ok=True)
destination: Final = directory / f"worker-config-{uuid.uuid4().hex}.yaml"
checkouts: Final = (_harness_root(), _proxy_root().resolve())
inside_checkout: Final = any(config.resolve().is_relative_to(checkout) for checkout in checkouts)
target: Final = directory if inside_checkout else config.parent
target.mkdir(parents=True, exist_ok=True)
destination: Final = target / f"worker-config-{uuid.uuid4().hex}.yaml"
loaded: Final[object] = yaml.safe_load(contents)
if not isinstance(loaded, dict):
destination.write_text(contents)

View file

@ -1695,15 +1695,19 @@ def test_upstream_pause_and_worker_kill_preserve_required_body_status(
record_property("chaos_spend_row_counts", str(tuple(len(rows) for rows in spend_rows)))
with owned.replacing_worker(min(workers)) as kill:
kill()
post_kill: Final = tuple(
_post(candidate, path, body) for path, body, _expected, _identity in missing_calls[:5]
)
assert all(
response.status_code == 400 and response.json() == expected
for response, (_path, _body, expected, _identity) in zip(post_kill, missing_calls[:5])
), [response.text for response in post_kill]
recovered: Final = _post(candidate, "/v1/chat/completions", _chat_body(chat_model, "recovered", True))
assert recovered.status_code == 200, recovered.text
with httpx.Client(base_url=candidate.client.base_url, timeout=15, trust_env=False) as fresh:
survivor: Final = Gateway(fresh, candidate.key, candidate.upstream_url)
post_kill: Final = tuple(
_post(survivor, path, body) for path, body, _expected, _identity in missing_calls[:5]
)
assert all(
response.status_code == 400 and response.json() == expected
for response, (_path, _body, expected, _identity) in zip(post_kill, missing_calls[:5])
), [response.text for response in post_kill]
recovered: Final = _post(
survivor, "/v1/chat/completions", _chat_body(chat_model, "recovered", True)
)
assert recovered.status_code == 200, recovered.text
def test_eval_create_and_run_forward_client_metadata(gateway: Gateway) -> None:

View file

@ -702,6 +702,31 @@ BEDROCK_INVOKE_HAIKU_STREAM: Final = CostTrackingTestCase(
),
)
BEDROCK_INVOKE_HAIKU_CHAT_STREAM: Final = CostTrackingTestCase(
name="bedrock-invoke-haiku-chat-stream",
covers="quota_management.spend_tracking.scripted_wire.logs_cost",
model="bedrock/invoke/anthropic.claude-haiku-4-5-20251001-v1:0",
request={
"model": "$MODEL",
"messages": [{"role": "user", "content": "bedrock-invoke-haiku-chat-stream"}],
"stream": True,
"stream_options": {"include_usage": True},
"max_tokens": 412,
},
response=BEDROCK_INVOKE_HAIKU_STREAM.response,
expected=BEDROCK_INVOKE_HAIKU_STREAM.expected,
)
BEDROCK_INVOKE_HAIKU_RESPONSES: Final = CostTrackingTestCase(
name="bedrock-invoke-haiku-responses",
covers="quota_management.spend_tracking.cost_matrix.logs_cost",
model="bedrock/invoke/anthropic.claude-haiku-4-5-20251001-v1:0",
endpoint="/v1/responses",
request={"model": "$MODEL", "input": "bedrock-invoke-haiku-responses", "max_output_tokens": 412},
response=BEDROCK_INVOKE_HAIKU_JSON.response,
expected=BEDROCK_INVOKE_HAIKU_JSON.expected,
)
# azure_ai/gpt-5.4-mini-2026-03-17
AZURE_AI_GPT_5_4_MINI_LATEST: Final = CostTrackingTestCase(
@ -1741,6 +1766,8 @@ CASES: Final[tuple[CostTrackingTestCase, ...]] = (
XAI_GROK_5_PROVIDER_REPORTED_COST,
BEDROCK_INVOKE_HAIKU_JSON,
BEDROCK_INVOKE_HAIKU_STREAM,
BEDROCK_INVOKE_HAIKU_CHAT_STREAM,
BEDROCK_INVOKE_HAIKU_RESPONSES,
AZURE_AI_GPT_5_4_MINI_LATEST,
AZURE_AI_GPT_5_4_MINI_LATEST_STREAM,
AZURE_AI_CLAUDE_HAIKU_4_5_MESSAGES_WEB_SEARCH,

View file

@ -18,7 +18,6 @@ from pathlib import Path
from typing import Final, LiteralString
import httpx
import psutil
import psycopg
import pytest
from integration._support.budget import case_timeout_seconds
@ -55,9 +54,6 @@ LOCK_KEY_ROW_SQL: Final = 'SELECT token FROM "LiteLLM_VerificationToken" WHERE t
SESSIONS_BLOCKED_BY_SQL: Final = (
"SELECT count(*)::int AS waiting FROM pg_stat_activity WHERE %s::int = ANY(pg_blocking_pids(pid))"
)
CLIENT_PORTS_BLOCKED_BY_SQL: Final = (
"SELECT client_port FROM pg_stat_activity WHERE %s::int = ANY(pg_blocking_pids(pid)) AND client_port > 0"
)
TRANSACTIONS_OPEN_SINCE_SQL: Final = """
SELECT count(*)::int AS open FROM pg_stat_activity
WHERE datname = current_database() AND backend_type = 'client backend' AND pid <> pg_backend_pid()
@ -296,22 +292,6 @@ def _transactions_open_since(moment: str) -> int:
return _int_value(read_rows(TRANSACTIONS_OPEN_SINCE_SQL, (moment,)), "open")
def _client_ports_blocked_by(backend_pid: int) -> frozenset[int]:
rows: Final = read_rows(CLIENT_PORTS_BLOCKED_BY_SQL, (str(backend_pid),))
return frozenset(_int_value((row,), "client_port") for row in rows)
def _local_tcp_ports(pid: int) -> Iterator[int]:
worker: Final = psutil.Process(pid)
for process in (worker, *worker.children(recursive=True)):
yield from (connection.laddr.port for connection in process.net_connections(kind="tcp"))
def _worker_holding(owned: OwnedProxy, client_ports: frozenset[int]) -> int:
(worker,) = tuple(pid for pid in owned.worker_pids() if client_ports & frozenset(_local_tcp_ports(pid)))
return worker
def _roster_ids(team: str) -> tuple[str, ...]:
rows: Final = read_rows(TEAM_ROW_SQL, (team,))
assert len(rows) == 1, rows
@ -508,11 +488,16 @@ def test_a_worker_killed_mid_push_leaves_the_old_roster_and_a_retry_lands_the_ne
pinned_key: Final = _team_key_for(scenario, team, max(leaving))
with _holding_row_lock(LOCK_KEY_ROW_SQL, sha256(pinned_key.encode()).hexdigest()) as holder_pid:
push: Final = pool.submit(_push_outcome, candidate, team, arriving)
blocked_ports: Final = eventually(
lambda: _client_ports_blocked_by(holder_pid), bool, seconds=UNFIXED_PUSH_SECONDS
_ = eventually(
lambda: _sessions_blocked_by(holder_pid),
lambda waiting: waiting > 0,
seconds=UNFIXED_PUSH_SECONDS,
)
killed_at: Final = _database_now()
_ = owned.kill_worker(_worker_holding(owned, blocked_ports))
victims: Final = tuple(sorted(owned.worker_pids()))
assert victims, "expected at least one serving worker"
for pid in victims:
_ = owned.kill_worker(pid)
record_property("interrupted_push", push.result())
_ = eventually(lambda: _transactions_open_since(killed_at), lambda open: open == 0, seconds=ROLLBACK_SECONDS)
settled: Final = frozenset(_roster_ids(team))

View file

@ -860,6 +860,7 @@ def test_spec_client_discovers_the_aggregate_authorization_server_and_signs_in_t
"client_origin": "http://127.0.0.1:9",
"server_id": None,
"server_name": None,
"servers": None,
"connected": None,
}, described.text
completed: Final = gateway.client.post("/authorize/complete", data={"flow": connect_flow}, cookies=flow_cookies)
@ -969,6 +970,7 @@ def test_per_server_resource_scopes_the_session_bearer_to_the_named_server(gatew
"client_origin": "http://127.0.0.1:9",
"server_id": server_a,
"server_name": alias_a,
"servers": None,
"connected": True,
}, described.text
completed: Final = gateway.client.post("/authorize/complete", data={"flow": connect_flow}, cookies=flow_cookies)

View file

@ -519,10 +519,12 @@ def test_killing_one_worker_leaves_the_sibling_serving_the_fallback(fb_rig: Rig)
workers: Final = eventually(fb_rig.owned.worker_pids, lambda pids: len(pids) >= 2, seconds=30)
with fb_rig.owned.replacing_worker(min(workers)) as kill:
kill()
responses: Final = tuple(
fb_rig.gateway.request("POST", "/v1/chat/completions", _chat(fb_rig.models["fb"], _prompt("after kill")))
for _ in range(3)
)
with httpx.Client(base_url=fb_rig.gateway.client.base_url, timeout=15, trust_env=False) as fresh:
survivor: Final = Gateway(fresh, fb_rig.gateway.key, fb_rig.gateway.upstream_url)
responses: Final = tuple(
survivor.request("POST", "/v1/chat/completions", _chat(fb_rig.models["fb"], _prompt("after kill")))
for _ in range(3)
)
for response in responses:
assert response.status_code == 200, response.text
assert _FALLBACK_TEXT in response.text, response.text

View file

@ -8,6 +8,11 @@ dials: ``AWSCRTHTTPClient`` -> TLS ALPN h2 -> invoke-with-bidirectional-stream)
sides: the OpenAI realtime events the websocket client receives, and the toolSpec / toolResult
documents the proxy forwarded to the peer.
The fallback cells cover a websocket that outlives one provider attempt. A refused stream keeps the
client's opening message for the fallback, which gets it once whether it dials Bedrock again or
an OpenAI realtime deployment, and once Bedrock accepted a stream no later attempt may open a second
provider session on that websocket.
The file boots its own proxy instead of the shared gateway so it can pass ``SSL_CERT_FILE``
pointing at the peer's self-signed cert, which is the only way the CRT h2 client trusts it.
"""
@ -21,12 +26,19 @@ import uuid
from collections.abc import AsyncIterator, Callable, Iterator, Mapping
from contextlib import asynccontextmanager
from pathlib import Path
from types import MappingProxyType
from typing import Final
import httpx
import pytest
import websockets
import yaml
from integration._support.bedrock_realtime_peer import (
BedrockEvent,
BedrockRealtimePeer,
OpeningEnd,
Refusal,
StreamFailure,
Turn,
assistant_text,
bedrock_realtime_peer,
@ -43,12 +55,15 @@ from integration._support.client import (
object_value,
string_value,
)
from integration._support.process import owned_proxy
from integration._support.process import graceful_stop_seconds, owned_proxy
from integration._support.tls import write_self_signed_cert
from integration._support.upstream import delete_scenario, register_scenario
from integration.cost_calculation.cost_tracking_case import RealtimeResponse
from pydantic import JsonValue
from websockets.asyncio.client import ClientConnection
from websockets.exceptions import ConnectionClosed
pytestmark: Final = pytest.mark.timeout(180)
pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120)
DEPLOYMENT: Final = "bedrock/amazon.nova-2-sonic-v1:0"
WEATHER_PARAMETERS: Final = {"type": "object", "properties": {"location": {"type": "string"}}, "required": ["location"]}
@ -60,6 +75,30 @@ WEATHER_FUNCTION: Final = {
NESTED_TOOL: Final = {"type": "function", "function": WEATHER_FUNCTION}
FLAT_TOOL: Final = {"type": "function", **WEATHER_FUNCTION}
TOOL_RESULT_OUTPUT: Final = '{"temperature_f":72,"conditions":"sunny"}'
QUESTION: Final = "What is the weather in New York right now?"
ANSWER: Final = "It is 72 and sunny"
INSTRUCTIONS: Final = "Answer every weather question in one short sentence."
OPENAI_VOICE_MODEL: Final = "gpt-realtime-2.1"
BEDROCK_REPLAY_GROUPS: Final = MappingProxyType(
{
"system_prompt": ("sonic-refused-session-update-first", "sonic-replay-session-update-first"),
"user_input": ("sonic-refused-text-first", "sonic-replay-text-first"),
}
)
OPENAI_REPLAY_GROUPS: Final = ("sonic-refused-openai-fallback", "openai-replay-target")
OPENAI_COMMITTED_GROUPS: Final = ("sonic-committed-openai-fallback", "openai-committed-target")
FALLBACK_GROUPS: Final = (*BEDROCK_REPLAY_GROUPS.values(), OPENAI_REPLAY_GROUPS, OPENAI_COMMITTED_GROUPS)
SCRIPTED_DONE: Final[dict[str, JsonValue]] = {
"type": "response.done",
"event_id": "evt_$UNIQUE_ID",
"response": {
"id": "resp_$UNIQUE_ID",
"object": "realtime.response",
"status": "completed",
"output": [],
"usage": {"total_tokens": 0, "input_tokens": 0, "output_tokens": 0},
},
}
@pytest.fixture(scope="module")
@ -67,15 +106,35 @@ def cert(tmp_path_factory: pytest.TempPathFactory) -> tuple[Path, Path]:
return write_self_signed_cert(tmp_path_factory.mktemp("bedrock-sonic-cert"))
def _fallback_config(directory: Path) -> Path:
stock: Final = JSON_OBJECT.validate_python(
yaml.safe_load((Path(__file__).resolve().parents[1] / "proxy_config.yaml").read_text())
)
config: Final = {
**stock,
"router_settings": {
**object_value(stock["router_settings"]),
"fallbacks": [{source: [target]} for source, target in FALLBACK_GROUPS],
},
}
path: Final = directory / "bedrock-sonic-fallbacks.yaml"
path.write_text(yaml.safe_dump(config))
return path
@pytest.fixture(scope="module")
def candidate(tmp_path_factory: pytest.TempPathFactory, cert: tuple[Path, Path]) -> Iterator[Gateway]:
directory: Final = tmp_path_factory.mktemp("bedrock-sonic")
with gateway_from_environment() as base:
with owned_proxy(base, tmp_path_factory.mktemp("bedrock-sonic"), {"SSL_CERT_FILE": str(cert[0])}) as proxy:
with owned_proxy(base, directory, {"SSL_CERT_FILE": str(cert[0])}, config=_fallback_config(directory)) as proxy:
yield proxy
def _deployment(gateway: Gateway, scenario: Scenario, peer: BedrockRealtimePeer) -> str:
name: Final = f"bedrock-sonic-{uuid.uuid4().hex[:12]}"
return _bedrock_deployment(gateway, scenario, peer, f"bedrock-sonic-{uuid.uuid4().hex[:12]}")
def _bedrock_deployment(gateway: Gateway, scenario: Scenario, peer: BedrockRealtimePeer, name: str) -> str:
created: Final = gateway.post(
"/model/new",
{
@ -94,6 +153,44 @@ def _deployment(gateway: Gateway, scenario: Scenario, peer: BedrockRealtimePeer)
return name
def _openai_deployment(gateway: Gateway, scenario: Scenario, name: str) -> str:
scenario_id: Final = f"bedrock-sonic-fallback-{uuid.uuid4().hex[:12]}"
handle: Final = register_scenario(
scenario_id, RealtimeResponse(content_type="application/x-realtime", events=(SCRIPTED_DONE,))
)
scenario.cleanups.callback(delete_scenario, handle)
created: Final = gateway.post(
"/model/new",
{
"model_name": name,
"litellm_params": {
"model": f"openai/{OPENAI_VOICE_MODEL}",
"api_key": scenario_id,
"api_base": f"{gateway.upstream_url}/v1",
},
"model_info": {"mode": "realtime"},
},
)
scenario.cleanups.callback(scenario.delete_model, string_value(object_value(created["model_info"])["id"]))
return scenario_id
def _upstream_requests(gateway: Gateway, scenario_id: str) -> tuple[dict[str, JsonValue], ...]:
with httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream:
observations: Final = JSON_OBJECT.validate_json(upstream.get("/__observations").content)
requests: Final = observations["requests"]
assert isinstance(requests, list), observations
return tuple(
request
for request in map(object_value, requests)
if request["authorization"] == f"Bearer {scenario_id}" or request["api_key"] == scenario_id
)
def _upstream_frames(requests: tuple[dict[str, JsonValue], ...]) -> tuple[dict[str, JsonValue], ...]:
return tuple(object_value(request["body"]) for request in requests if request["method"] == "WEBSOCKET_FRAME")
def _session_update(tool: Mapping[str, JsonValue]) -> str:
return json.dumps(
{
@ -113,12 +210,16 @@ def _user_text() -> str:
"item": {
"type": "message",
"role": "user",
"content": [{"type": "input_text", "text": "What is the weather in New York right now?"}],
"content": [{"type": "input_text", "text": QUESTION}],
},
}
)
def _instructions_update() -> str:
return json.dumps({"type": "session.update", "session": {"instructions": INSTRUCTIONS}})
_RESPONSE_CREATE: Final = json.dumps({"type": "response.create"})
@ -207,10 +308,13 @@ def test_nested_tool_call_before_assistant_block(candidate: Gateway, cert: tuple
assert _tail_types(events) == (
"response.created",
"response.output_item.added",
"response.function_call_arguments.delta",
"response.function_call_arguments.done",
"response.output_item.done",
"response.done",
), events
(delta,) = tuple(event for event in events if event.get("type") == "response.function_call_arguments.delta")
assert string_value(delta["delta"]) == '{"location":"New York"}'
(call,) = _function_calls(events)
assert string_value(call["call_id"]) == "tooluse_A"
assert string_value(call["name"]) == "get_weather"
@ -376,3 +480,160 @@ def test_malformed_tool_content(candidate: Gateway, cert: tuple[Path, Path], con
events: Final = asyncio.run(_one_call_session(candidate, model, key, NESTED_TOOL))
(call,) = _function_calls(events)
assert string_value(call["arguments"]) == (content if content is not None else "{}"), events
def _of_type(events: tuple[Event, ...], event_type: str) -> tuple[Event, ...]:
return tuple(event for event in events if event.get("type") == event_type)
def _bedrock_fields(events: tuple[BedrockEvent, ...], name: str) -> tuple[dict[str, JsonValue], ...]:
bodies: Final = tuple(object_value(event["event"]) for event in events)
return tuple(object_value(body[name]) for body in bodies if name in body)
def _text_inputs(events: tuple[BedrockEvent, ...], role: str) -> tuple[str, ...]:
roles: Final = MappingProxyType(
{string_value(start["contentName"]): start.get("role") for start in _bedrock_fields(events, "contentStart")}
)
return tuple(
string_value(text["content"])
for text in _bedrock_fields(events, "textInput")
if roles.get(string_value(text["contentName"])) == role
)
def _answer_turn() -> Turn:
return Turn("user_input", assistant_text(ANSWER) + (usage(30, 12),))
def _answer_text(events: tuple[Event, ...]) -> str:
return "".join(string_value(event["delta"]) for event in _of_type(events, "response.text.delta"))
async def _opened_with(socket: ClientConnection, opening: tuple[str, ...], stop: frozenset[str]) -> tuple[Event, ...]:
created: Final = JSON_OBJECT.validate_json(await socket.recv())
for message in opening:
await socket.send(message)
return (created,) + await _events_until(socket, lambda event, _events: event.get("type") in stop)
async def _session(
gateway: Gateway, model: str, key: str, opening: tuple[str, ...], stop: frozenset[str]
) -> tuple[Event, ...]:
async with _realtime(gateway, model, key) as socket:
return await _opened_with(socket, opening, stop)
async def _session_until_closed(
gateway: Gateway, model: str, key: str, opening: tuple[str, ...]
) -> tuple[tuple[Event, ...], int | None]:
async with _realtime(gateway, model, key) as socket:
events: Final = await _opened_with(socket, opening, frozenset(("error", "session.created")))
if events[-1].get("type") != "error":
return events, None
try:
await asyncio.wait_for(socket.recv(), 30)
except ConnectionClosed:
return events, socket.close_code
return events, None
@pytest.mark.parametrize("after", ["system_prompt", "user_input"], ids=["session-update-first", "text-first"])
def test_refused_opening_is_replayed_on_the_bedrock_fallback(
candidate: Gateway, cert: tuple[Path, Path], after: OpeningEnd
) -> None:
source, target = BEDROCK_REPLAY_GROUPS[after]
update_first: Final = after == "system_prompt"
opening: Final = ((_instructions_update(),) if update_first else ()) + (_user_text(), _RESPONSE_CREATE)
opening_role, opening_text = ("SYSTEM", INSTRUCTIONS) if update_first else ("USER", QUESTION)
with (
bedrock_realtime_peer((), cert, refusal=Refusal(after)) as refusing,
bedrock_realtime_peer((_answer_turn(),), cert) as healthy,
candidate.scenario() as scenario,
):
_bedrock_deployment(candidate, scenario, refusing, source)
_bedrock_deployment(candidate, scenario, healthy, target)
events: Final = asyncio.run(
_session(
candidate, source, scenario.key(), opening, frozenset(("response.done", "error", "session.created"))
)
)
assert len(_of_type(events, "session.created")) == 1, events
assert len(_of_type(events, "session.updated")) == int(update_first), events
assert _of_type(events, "error") == (), events
assert len(_response_dones(events)) == 1, events
assert _answer_text(events) == ANSWER, events
(replayed,) = healthy.streams()
assert _text_inputs(replayed, "USER") == (QUESTION,), replayed
assert _text_inputs(replayed, "SYSTEM").count(INSTRUCTIONS) == int(update_first), replayed
refused: Final = refusing.streams()
assert refused and all(opening_text in _text_inputs(stream, opening_role) for stream in refused), refused
def test_refused_text_opening_is_replayed_on_the_openai_fallback(candidate: Gateway, cert: tuple[Path, Path]) -> None:
source, target = OPENAI_REPLAY_GROUPS
with (
bedrock_realtime_peer((), cert, refusal=Refusal("user_input")) as refusing,
candidate.scenario() as scenario,
):
_bedrock_deployment(candidate, scenario, refusing, source)
scenario_id: Final = _openai_deployment(candidate, scenario, target)
events: Final = asyncio.run(
_session(
candidate,
source,
scenario.key(),
(_user_text(), _RESPONSE_CREATE),
frozenset(("response.done", "error")),
)
)
assert _of_type(events, "error") == (), events
assert len(_response_dones(events)) == 1, events
requests: Final = _upstream_requests(candidate, scenario_id)
assert len(tuple(request for request in requests if request["method"] == "WEBSOCKET")) == 1, requests
assert _upstream_frames(requests) == (json.loads(_user_text()), json.loads(_RESPONSE_CREATE)), requests
def test_committed_session_failure_closes_the_socket_without_another_provider_stream(
candidate: Gateway, cert: tuple[Path, Path]
) -> None:
source, target = OPENAI_COMMITTED_GROUPS
with (
bedrock_realtime_peer((_answer_turn(),), cert, failure=StreamFailure()) as committed,
candidate.scenario() as scenario,
):
_bedrock_deployment(candidate, scenario, committed, source)
scenario_id: Final = _openai_deployment(candidate, scenario, target)
events, close_code = asyncio.run(
_session_until_closed(candidate, source, scenario.key(), (_user_text(), _RESPONSE_CREATE))
)
assert len(_of_type(events, "session.created")) == 1, events
assert len(_response_dones(events)) == 1, events
assert _answer_text(events) == ANSWER, events
assert string_value(events[-1]["type"]) == "error", events
assert close_code == 1011, events
assert _upstream_requests(candidate, scenario_id) == ()
assert len(committed.streams()) == 1
def test_voice_object_is_refused_and_never_reaches_bedrock(candidate: Gateway, cert: tuple[Path, Path]) -> None:
voice_update: Final = json.dumps({"type": "session.update", "session": {"voice": {"id": "voice_123"}}})
with bedrock_realtime_peer((_answer_turn(),), cert) as peer, candidate.scenario() as scenario:
model: Final = _deployment(candidate, scenario, peer)
events: Final = asyncio.run(
_session(
candidate,
model,
scenario.key(),
(voice_update, _user_text(), _RESPONSE_CREATE),
frozenset(("response.done", "session.created")),
)
)
errors: Final = tuple(object_value(event["error"]) for event in _of_type(events, "error"))
assert tuple((error.get("code"), error.get("param")) for error in errors) == (
("invalid_type", "session.voice"),
), events
assert _of_type(events, "session.updated") == (), events
assert len(_of_type(events, "session.created")) == 1, events
assert _answer_text(events) == ANSWER, events
assert "voice_123" not in json.dumps(peer.received())

View file

@ -84,7 +84,7 @@ _FAMILIES: Final = (
_Family("pin-bedrock-converse-1h", "pin-bedrock-converse-1h", one_hour=True, wire_one_hour=True),
_Family("pin-bedrock-converse-5m", "pin-bedrock-converse-5m", one_hour=False, wire_one_hour=False),
_Family("pin-bedrock-invoke-1h", "pin-bedrock-invoke-1h", one_hour=True, wire_one_hour=True),
_Family("pin-bedrock-arn-base", "pin-bedrock-arn-base", one_hour=False, wire_one_hour=False),
_Family("pin-bedrock-arn-base", "pin-bedrock-arn-base", one_hour=True, wire_one_hour=True),
_Family("pin-bedrock-arn", "pin-bedrock-arn", one_hour=False, wire_one_hour=False),
_Family("pin-azure-ai", "pin-azure-ai", one_hour=True, wire_one_hour=True),
_Family("pin-vertex", "pin-vertex", one_hour=True, wire_one_hour=True),
@ -449,19 +449,39 @@ def test_provider_family_pin_lifetime_agrees_with_the_ttl_the_wire_carries(rig:
spend_row_lands(response_id)
def test_messages_through_bedrock_invoke_pins_an_hour_and_sends_the_hour(rig: _Rig) -> None:
@pytest.mark.parametrize("group", ["pin-bedrock-invoke-1h", "pin-bedrock-arn-base"])
def test_messages_through_bedrock_pins_an_hour_and_sends_the_hour(rig: _Rig, group: str) -> None:
before: Final = _before(rig)
marker: Final = new_marker()
response: Final = rig.owned.gateway.request(
"POST", "/v1/messages", {**_messages_body(new_marker(), marker), "model": "pin-bedrock-invoke-1h"}
"POST", "/v1/messages", {**_messages_body(new_marker(), marker), "model": group}
)
assert response.status_code == 200, response.text
_, pin = one_new_pin(rig.redis, before)
assert pin.model_id == "pin-bedrock-invoke-1h", pin
assert pin.model_id == group, pin
assert lives_one_hour(pin), pin
assert [wire_carries_one_hour(request) for request in sent_with((rig.bedrock,), marker)] == [True]
def test_response_cache_hit_on_an_application_profile_rewrites_the_pin_for_its_base_model_hour(rig: _Rig) -> None:
before: Final = _before(rig)
marker: Final = new_marker()
body: Final = chat_body("pin-bedrock-arn-base", _first_turn(new_marker(), marker, ONE_HOUR_MARK))
first, served = post_chat(rig.owned.gateway, body)
key, pin = one_new_pin(rig.redis, before)
assert served == pin.model_id == "pin-bedrock-arn-base", (served, pin)
assert lives_one_hour(pin), pin
assert [wire_carries_one_hour(request) for request in sent_with(rig.wires(), marker)] == [True]
spend_row_lands(first)
rig.redis.delete(key)
second, _ = post_chat(rig.owned.gateway, body)
assert sent_with(rig.wires(), marker) == ()
assert second == first
rewritten: Final = eventually(lambda: pin_at(rig.redis, key), lambda found: found is not None, seconds=30)
assert rewritten is not None and rewritten.model_id == "pin-bedrock-arn-base", rewritten
assert lives_one_hour(rewritten), rewritten
@pytest.mark.parametrize("ttl", _MALFORMED_TTLS, ids=["int", "2h", "upper-1H", "empty", "null", "list", "5kb"])
def test_malformed_ttl_answers_and_pins_five_minutes_without_a_wire_ttl(rig: _Rig, ttl: JsonValue) -> None:
before: Final = _before(rig)

View file

@ -142,6 +142,7 @@ def test_failed_daily_user_rollup_commit_is_retried_so_spend_report_and_daily_ac
assert metrics == {
"spend": pytest.approx(0.06),
"flat_cost": pytest.approx(0.0),
"ptu_hours": pytest.approx(0.0),
"prompt_tokens": 20,
"completion_tokens": 20,
"cache_read_input_tokens": 0,

View file

@ -7,6 +7,7 @@ from dotenv import load_dotenv
load_dotenv()
import asyncio
import hashlib
import threading
import random
import pytest
@ -64,6 +65,46 @@ def setup_cache():
return cache
@pytest.mark.parametrize(
("call", "request_kwargs"),
[
(
completion,
{"model": "gpt-4o-mini", "messages": [{"role": "user", "content": "hi"}], "mock_response": "x"},
),
(embedding, {"model": "text-embedding-3-small", "input": ["hi"], "mock_response": [0.1, 0.2]}),
(litellm.text_completion, {"model": "gpt-3.5-turbo-instruct", "prompt": "hi", "mock_response": "x"}),
],
)
def test_sync_cache_hits_reach_sync_custom_logger(
monkeypatch: pytest.MonkeyPatch, call: Callable[..., object], request_kwargs: dict[str, object]
) -> None:
from litellm.integrations.custom_logger import CustomLogger
logged_cache_hits: Final[list[object]] = []
all_logged: Final = threading.Event()
class SyncRecorder(CustomLogger):
def log_success_event(
self, kwargs: dict[str, object], response_obj: object, start_time: object, end_time: object
) -> None:
logged_cache_hits.append(kwargs.get("cache_hit"))
if len(logged_cache_hits) == 3:
all_logged.set()
monkeypatch.setattr(litellm, "success_callback", [])
monkeypatch.setattr(litellm, "_async_success_callback", [])
monkeypatch.setattr(litellm, "input_callback", [])
monkeypatch.setattr(litellm, "callbacks", [SyncRecorder()])
monkeypatch.setattr(litellm, "cache", Cache(type=LiteLLMCacheType.LOCAL))
for _ in range(3):
call(**request_kwargs)
assert all_logged.wait(timeout=10)
assert logged_cache_hits == [None, True, True]
@pytest.fixture
def copilot_response_cache(monkeypatch: pytest.MonkeyPatch) -> Cache:
cache: Final = Cache(type=LiteLLMCacheType.LOCAL)
@ -1071,7 +1112,7 @@ def test_sync_get_cache_does_not_eagerly_log_streaming_responses_hits():
stream=True,
start_time=datetime.now(),
)
logging_obj.handle_sync_success_callbacks_for_async_calls = MagicMock()
logging_obj.handle_sync_cache_hit_success_callbacks = MagicMock()
kwargs = {
"model": original_model,
@ -1095,7 +1136,7 @@ def test_sync_get_cache_does_not_eagerly_log_streaming_responses_hits():
assert isinstance(
cached_response.cached_result, CachedResponsesAPIStreamingIterator
)
logging_obj.handle_sync_success_callbacks_for_async_calls.assert_not_called()
logging_obj.handle_sync_cache_hit_success_callbacks.assert_not_called()
def test_sync_get_cache_defers_streaming_completion_hit_callbacks():
@ -1115,7 +1156,7 @@ def test_sync_get_cache_defers_streaming_completion_hit_callbacks():
stream=True,
start_time=datetime.now(),
)
logging_obj.handle_sync_success_callbacks_for_async_calls = MagicMock()
logging_obj.handle_sync_cache_hit_success_callbacks = MagicMock()
kwargs = {
"model": original_model,
@ -1136,7 +1177,7 @@ def test_sync_get_cache_defers_streaming_completion_hit_callbacks():
)
assert cached_response.cached_result is not None
logging_obj.handle_sync_success_callbacks_for_async_calls.assert_not_called()
logging_obj.handle_sync_cache_hit_success_callbacks.assert_not_called()
def test_should_defer_streaming_cache_hit_callbacks_for_any_streaming_request():

View file

@ -20,6 +20,7 @@ import pytest
import yaml
from pydantic import TypeAdapter
from litellm.proxy.types_utils.utils import get_instance_fn
from tests.integration._support.client import Gateway, eventually
from tests.integration._support.process import (
OwnedProxy,
@ -36,6 +37,32 @@ from tests.integration._support.process import (
TESTS_DIR: Final = Path(__file__).resolve().parents[2]
@pytest.mark.parametrize("config_contents", ("[]\n", "general_settings: []\n", "general_settings: {}\n"))
def test_worker_configuration_preserves_sibling_auth_modules(
monkeypatch: pytest.MonkeyPatch,
tmp_path: Path,
config_contents: str,
) -> None:
monkeypatch.setenv("INTEGRATION_WORKER_ISOLATION", "1")
monkeypatch.setenv("INTEGRATION_PROVIDER_PORT", "8203")
monkeypatch.setenv("INTEGRATION_UPSTREAM_URL", "http://127.0.0.1:8202")
source: Final = tmp_path / "source.yaml"
source.write_text(config_contents)
auth_state: Final = tmp_path / "sibling_auth.json"
auth_state.write_text("authenticated")
(tmp_path / "sibling_auth.py").write_text(
"from pathlib import Path\n"
"def authenticate():\n"
" return Path(__file__).with_suffix('.json').read_text()\n"
)
configured: Final = worker_configuration(source, tmp_path / "worker")
assert get_instance_fn("sibling_auth.authenticate", config_file_path=str(configured))() == "authenticated"
auth_state.write_text("updated")
assert get_instance_fn("sibling_auth.authenticate", config_file_path=str(configured))() == "updated"
def test_worker_configuration_routes_stock_edges_to_its_own_servers_and_preserves_other_targets(
monkeypatch: pytest.MonkeyPatch, tmp_path: Path
) -> None:
@ -93,6 +120,42 @@ def test_worker_configuration_budgets_implicit_pools_without_changing_connection
assert yaml.safe_load(source.read_text()) == original
def test_worker_configuration_keeps_modules_beside_a_test_config_importable(
monkeypatch: pytest.MonkeyPatch, tmp_path: Path
) -> None:
monkeypatch.setenv("INTEGRATION_WORKER_ISOLATION", "1")
source: Final = tmp_path / "proxy.yaml"
source.write_text(yaml.safe_dump({"general_settings": {"custom_auth": "sibling_auth.authenticate"}}))
(tmp_path / "sibling_auth.py").write_text("def authenticate() -> str:\n return 'sibling'\n")
configured: Final = worker_configuration(source, tmp_path / "gateway")
assert configured != source
assert get_instance_fn("sibling_auth.authenticate", config_file_path=str(configured))() == "sibling"
def test_worker_configuration_writes_the_repository_config_copy_into_the_worker_directory(
monkeypatch: pytest.MonkeyPatch, tmp_path: Path
) -> None:
monkeypatch.setenv("INTEGRATION_WORKER_ISOLATION", "1")
monkeypatch.setenv("INTEGRATION_PROVIDER_PORT", "8203")
monkeypatch.setenv("INTEGRATION_UPSTREAM_URL", "http://127.0.0.1:8202")
monkeypatch.setenv("INTEGRATION_PROXY_ROOT", str(tmp_path / "product"))
configured: Final = worker_configuration(TESTS_DIR / "integration/proxy_config.yaml", tmp_path / "worker")
assert configured.parent == tmp_path / "worker"
def test_worker_configuration_writes_the_proxy_checkout_config_copy_into_the_worker_directory(
monkeypatch: pytest.MonkeyPatch, tmp_path: Path
) -> None:
monkeypatch.setenv("INTEGRATION_WORKER_ISOLATION", "1")
checkout: Final = tmp_path / "checkout"
monkeypatch.setenv("INTEGRATION_PROXY_ROOT", str(checkout))
source: Final = checkout / "tests/integration/proxy_config.yaml"
source.parent.mkdir(parents=True)
source.write_text("model_list: []\n")
configured: Final = worker_configuration(source, tmp_path / "worker")
assert configured.parent == tmp_path / "worker"
def test_serial_proxy_keeps_the_requested_configuration(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None:
monkeypatch.delenv("INTEGRATION_WORKER_ISOLATION", raising=False)
source: Final = tmp_path / "source.yaml"

View file

@ -155,6 +155,89 @@ async def test_a_replayed_stream_skips_cached_completion_tokens_details_that_do_
assert rebuilt.usage.completion_tokens == 7
def _prompt_cached_payload(content: str, cache_read: int, cache_write: int) -> dict:
cached = _async_payload(content=content)
cached["usage"] = litellm.Usage(
prompt_tokens=42145,
completion_tokens=6,
total_tokens=42151,
cache_read_input_tokens=cache_read,
cache_creation_input_tokens=cache_write,
).model_dump()
return cached
@pytest.mark.parametrize("content", ["", "Hello world! How are you?"], ids=["single_frame", "sliced"])
@pytest.mark.parametrize("cache_read, cache_write", [(42139, 0), (0, 42139)], ids=["read", "write"])
@pytest.mark.parametrize("replay", ["async", "sync"])
@pytest.mark.asyncio
async def test_a_replayed_stream_keeps_the_cached_prompt_cache_tokens(replay, cache_read, cache_write, content):
cached = _prompt_cached_payload(content, cache_read, cache_write)
chunks = await _collect_async(cached) if replay == "async" else list(convert_to_streaming_response(cached))
usage: Final = next(c.usage for c in chunks if getattr(c, "usage", None) is not None)
assert usage.prompt_tokens == 42145
assert usage.prompt_tokens_details.cached_tokens == cache_read
assert usage.prompt_tokens_details.cache_write_tokens == cache_write
assert usage.cache_read_input_tokens == cache_read
assert usage.cache_creation_input_tokens == cache_write
@pytest.mark.parametrize("replay", ["async", "sync"])
@pytest.mark.asyncio
async def test_a_replayed_stream_keeps_cached_tokens_reported_only_in_the_prompt_tokens_details(replay):
cached = _async_payload()
cached["usage"]["prompt_tokens_details"] = {"cached_tokens": 4, "audio_tokens": 1}
chunks = await _collect_async(cached) if replay == "async" else list(convert_to_streaming_response(cached))
rebuilt = litellm.stream_chunk_builder(chunks=chunks)
assert rebuilt.usage.prompt_tokens_details.cached_tokens == 4
assert rebuilt.usage.prompt_tokens_details.audio_tokens == 1
@pytest.mark.parametrize("replay", ["async", "sync"])
@pytest.mark.asyncio
async def test_a_replayed_stream_keeps_integral_counts_cached_as_floats_or_numeric_strings(replay):
cached = _async_payload()
cached["usage"] = {
"prompt_tokens": 100.0,
"completion_tokens": "6",
"total_tokens": 106.0,
"cache_read_input_tokens": 90.0,
}
chunks = await _collect_async(cached) if replay == "async" else list(convert_to_streaming_response(cached))
usage: Final = next(c.usage for c in chunks if getattr(c, "usage", None) is not None)
assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (100, 6, 106)
assert usage.cache_read_input_tokens == 90
@pytest.mark.parametrize(
"field, value",
[
("prompt_tokens_details", {"cached_tokens": "abc"}),
("prompt_tokens_details", [1, 2]),
("cache_read_input_tokens", "abc"),
("cache_creation_input_tokens", 1.5),
],
ids=["details_text", "details_list", "read_text", "write_fraction"],
)
@pytest.mark.parametrize("replay", ["async", "sync"])
@pytest.mark.asyncio
async def test_a_replayed_stream_skips_cached_prompt_cache_usage_that_does_not_validate(replay, field, value):
cached = _async_payload()
cached["usage"][field] = value
chunks = await _collect_async(cached) if replay == "async" else list(convert_to_streaming_response(cached))
rebuilt = litellm.stream_chunk_builder(chunks=chunks)
assert rebuilt.choices[0].message.content == "Hello world! How are you?"
assert rebuilt.usage.prompt_tokens == 5
@pytest.mark.asyncio
async def test_async_finish_reason_only_on_last_chunk():
chunks = await _collect_async(_async_payload())

View file

@ -7,6 +7,7 @@ import json
import logging
import os
import sys
import threading
import time
import traceback
from collections.abc import AsyncIterator, Callable, Iterator, Mapping, Sequence
@ -7540,6 +7541,252 @@ async def test_streaming_success_callbacks_survive_guardrail_logging_hook_failur
releasing.async_log_success_event.assert_awaited_once()
def _async_request_logging_obj(call_id: str, callbacks: list[CustomLogger], stream: bool = False) -> LitellmLogging:
logging_obj: Final = LitellmLogging(
model="gpt-4o-mini",
messages=[{"role": "user", "content": "hi"}],
stream=stream,
call_type="acompletion",
start_time=datetime.datetime.now(),
litellm_call_id=call_id,
function_id=call_id,
dynamic_async_success_callbacks=callbacks,
)
logging_obj.update_environment_variables(litellm_params={"acompletion": True, "metadata": {}}, optional_params={})
return logging_obj
@pytest.mark.parametrize(("success_callbacks", "expected"), [(["cache"], False), (["cache", "athina"], True)])
def test_cache_callback_alone_does_not_start_sync_callbacks_for_async_requests(
monkeypatch: pytest.MonkeyPatch, success_callbacks: list[str], expected: bool
) -> None:
monkeypatch.setattr(litellm, "success_callback", success_callbacks)
logging_obj: Final = _async_request_logging_obj("cache-gate", [])
assert logging_obj._should_run_sync_callbacks_for_async_calls() is expected
@pytest.mark.asyncio
@pytest.mark.parametrize("turn_off_message_logging", [False, True])
@pytest.mark.parametrize("make_response", [ModelResponse, _responses_api_response_with_text])
async def test_sync_success_handler_landing_after_logging_hook_keeps_hook_payload_edits(
monkeypatch: pytest.MonkeyPatch, make_response: Callable[[], object], turn_off_message_logging: bool
) -> None:
mirrored: Final = {"x-ratelimit-model_per_project-limit-requests": 600}
response: Final = make_response()
seen_edits: Final[list[tuple[object, object]]] = []
class MirroringHook(CustomLogger):
async def async_logging_hook(
self, kwargs: dict[str, object], result: object, call_type: str
) -> tuple[dict[str, object], object]:
payload = kwargs["standard_logging_object"]
assert isinstance(payload, dict)
payload["hidden_params"]["additional_headers"] = dict(mirrored)
payload["response"] = "masked by hook"
logging_obj.success_handler(result=response, for_async_request=True)
return kwargs, result
class Recorder(CustomLogger):
async def async_log_success_event(
self, kwargs: dict[str, object], response_obj: object, start_time: object, end_time: object
) -> None:
payload = kwargs["standard_logging_object"]
assert isinstance(payload, dict)
seen_edits.append((payload["hidden_params"]["additional_headers"], payload["response"]))
monkeypatch.setattr(litellm, "_async_success_callback", [])
monkeypatch.setattr(litellm, "success_callback", [])
monkeypatch.setattr(litellm, "turn_off_message_logging", turn_off_message_logging)
logging_obj: Final = _async_request_logging_obj("late-sync-rebuild", [MirroringHook(), Recorder()])
await logging_obj.async_success_handler(result=response)
assert seen_edits == [(mirrored, "masked by hook")]
class _ResponseEditingHook(CustomLogger):
def logging_hook(
self, kwargs: dict[str, object], result: object, call_type: str
) -> tuple[dict[str, object], object]:
payload = kwargs["standard_logging_object"]
assert isinstance(payload, dict)
payload["response"] = "edited by sync hook"
return kwargs, result
def test_sync_logging_hook_edits_stay_off_a_caller_supplied_payload_for_async_requests(
monkeypatch: pytest.MonkeyPatch,
) -> None:
sync_payloads: Final[list[object]] = []
class AthinaRecorder:
def log_event(
self,
kwargs: dict[str, object],
response_obj: object,
start_time: object,
end_time: object,
print_verbose: object,
) -> None:
sync_payloads.append(kwargs["standard_logging_object"])
monkeypatch.setattr(litellm, "_async_success_callback", [])
monkeypatch.setattr(litellm, "success_callback", [_ResponseEditingHook(), "athina"])
monkeypatch.setattr(litellm_logging, "athinaLogger", AthinaRecorder())
logging_obj: Final = _async_request_logging_obj("caller-supplied-payload", [])
supplied_payload: Final = {"id": "caller-supplied-payload", "response": "original"}
logging_obj.model_call_details["standard_logging_object"] = supplied_payload
logging_obj.success_handler(
result=ModelResponse(), for_async_request=True, standard_logging_object=supplied_payload
)
assert len(sync_payloads) == 1
sync_payload: Final = sync_payloads[0]
assert isinstance(sync_payload, dict)
assert sync_payload["response"] == "edited by sync hook"
assert supplied_payload == {"id": "caller-supplied-payload", "response": "original"}
assert logging_obj.model_call_details["standard_logging_object"] is supplied_payload
@pytest.mark.asyncio
async def test_sync_only_logger_keeps_the_async_payload_when_its_result_builds_none(
monkeypatch: pytest.MonkeyPatch,
) -> None:
sync_payloads: Final[list[object]] = []
class AthinaRecorder:
def log_event(
self,
kwargs: dict[str, object],
response_obj: object,
start_time: object,
end_time: object,
print_verbose: object,
) -> None:
sync_payloads.append(kwargs.get("standard_logging_object"))
monkeypatch.setattr(litellm, "_async_success_callback", [])
monkeypatch.setattr(litellm, "success_callback", ["athina"])
monkeypatch.setattr(litellm_logging, "athinaLogger", AthinaRecorder())
logging_obj: Final = LitellmLogging(
model="pass-through-model",
messages=[{"role": "user", "content": "hi"}],
stream=False,
call_type="pass_through_endpoint",
start_time=datetime.datetime.now(),
litellm_call_id="pass-through-sync-logger",
function_id="pass-through-sync-logger",
)
logging_obj.update_environment_variables(litellm_params={"metadata": {}}, optional_params={})
raw_body: Final = json.dumps({"answer": "plain text"})
await logging_obj.async_success_handler(result=raw_body)
async_payload: Final = logging_obj.model_call_details.get("standard_logging_object")
assert isinstance(async_payload, dict)
async_payload_before_sync: Final = dict(async_payload)
monkeypatch.setattr(litellm, "success_callback", [_ResponseEditingHook(), "athina"])
logging_obj.success_handler(result=raw_body, for_async_request=True)
assert sync_payloads == [{**async_payload_before_sync, "response": "edited by sync hook"}]
assert logging_obj.model_call_details["standard_logging_object"] is async_payload
assert async_payload == async_payload_before_sync
@pytest.mark.asyncio
@pytest.mark.parametrize("stream", [False, True])
@pytest.mark.parametrize("entrypoint", ["handle_sync_success_callbacks_for_async_calls", "dispatch_success_handlers"])
@pytest.mark.parametrize(
("callback_name", "logger_global"),
[
("athina", "athinaLogger"),
("promptlayer", "promptLayerLogger"),
("wandb", "weightsBiasesLogger"),
("lunary", "lunaryLogger"),
("helicone", "heliconeLogger"),
("greenscale", "greenscaleLogger"),
("s3", "s3Logger"),
("traceloop", "traceloopLogger"),
],
)
async def test_sync_only_logger_gets_its_own_payload_for_async_requests(
monkeypatch: pytest.MonkeyPatch, stream: bool, entrypoint: str, callback_name: str, logger_global: str
) -> None:
sync_logged: Final = threading.Event()
sync_payloads: Final[list[object]] = []
async_payloads: Final[list[object]] = []
class LegacyLoggerRecorder:
def _record(self, **call: object) -> None:
kwargs = call["kwargs"]
assert isinstance(kwargs, dict)
sync_payloads.append(kwargs["standard_logging_object"])
sync_logged.set()
log_event = log_success = _record
class AsyncRecorder(CustomLogger):
async def async_log_success_event(
self, kwargs: dict[str, object], response_obj: object, start_time: object, end_time: object
) -> None:
async_payloads.append(kwargs["standard_logging_object"])
monkeypatch.setattr(litellm, "_async_success_callback", [])
monkeypatch.setattr(litellm, "success_callback", [callback_name])
monkeypatch.setattr(litellm_logging, logger_global, LegacyLoggerRecorder(), raising=False)
response: Final = ModelResponse(id="chatcmpl-sync-only")
logging_obj: Final = _async_request_logging_obj("sync-only-logger", [AsyncRecorder()], stream=stream)
now: Final = datetime.datetime.now()
if entrypoint == "dispatch_success_handlers":
await logging_obj.dispatch_success_handlers(result=response, start_time=now, end_time=now)
else:
await logging_obj.async_success_handler(result=response, start_time=now, end_time=now)
logging_obj.handle_sync_success_callbacks_for_async_calls(result=response, start_time=now, end_time=now)
assert sync_logged.wait(timeout=10)
assert len(async_payloads) == 1
assert len(sync_payloads) == 1
async_payload: Final = async_payloads[0]
sync_payload: Final = sync_payloads[0]
assert isinstance(async_payload, dict)
assert isinstance(sync_payload, dict)
assert sync_payload["id"] == async_payload["id"]
assert sync_payload is not async_payload
assert logging_obj.model_call_details["standard_logging_object"] is async_payload
@pytest.mark.asyncio
@pytest.mark.parametrize("stream", [False, True])
async def test_supabase_logs_the_request_messages_and_response_for_async_requests(
monkeypatch: pytest.MonkeyPatch, stream: bool
) -> None:
supabase_logged: Final = threading.Event()
supabase_calls: Final[list[tuple[object, object]]] = []
class SupabaseRecorder:
def log_event(self, **call: object) -> None:
supabase_calls.append((call["messages"], call["response_obj"]))
supabase_logged.set()
monkeypatch.setattr(litellm, "_async_success_callback", [])
monkeypatch.setattr(litellm, "success_callback", ["supabase"])
monkeypatch.setattr(litellm_logging, "supabaseClient", SupabaseRecorder())
response: Final = ModelResponse(id="chatcmpl-supabase")
logging_obj: Final = _async_request_logging_obj("supabase-logger", [], stream=stream)
now: Final = datetime.datetime.now()
await logging_obj.dispatch_success_handlers(result=response, start_time=now, end_time=now)
assert supabase_logged.wait(timeout=10)
assert len(supabase_calls) == 1
messages, response_obj = supabase_calls[0]
assert messages == [{"role": "user", "content": "hi"}]
assert isinstance(response_obj, ModelResponse)
assert response_obj.id == "chatcmpl-supabase"
def _resolve(custom_llm_provider, litellm_params, optional_params, model):
from litellm.litellm_core_utils.litellm_logging import (
_resolve_vertex_location_for_cost,

View file

@ -13,6 +13,7 @@ from websockets.frames import Close
import litellm
from litellm.constants import (
BEDROCK_REALTIME_PENDING_OPENING_MESSAGE_SCOPE_KEY,
REALTIME_METERED_USAGE_KEY,
REALTIME_SESSION_FAILURE_LOGGED_KEY,
REALTIME_SESSION_SUCCESS_LOGGED_KEY,
@ -543,6 +544,32 @@ async def test_backend_to_client_stores_normalized_events_for_logging():
assert streaming.messages[0]["response"]["usage"]["total_tokens"] == 0
@pytest.mark.asyncio
async def test_client_ack_messages_replays_the_opening_message_a_refused_bedrock_stream_left():
opening_message: Final = json.dumps(
{
"type": "conversation.item.create",
"item": {"type": "message", "role": "user", "content": [{"type": "input_text", "text": "What is 17 * 3?"}]},
}
)
next_message: Final = json.dumps({"type": "response.create"})
client_ws = MagicMock()
client_ws.scope = {"headers": [], BEDROCK_REALTIME_PENDING_OPENING_MESSAGE_SCOPE_KEY: opening_message}
client_ws.receive_text = AsyncMock(side_effect=[next_message, Exception("connection closed")])
backend_ws = MagicMock()
backend_ws.send = AsyncMock()
logging_obj: Final = _mock_logging()
logging_obj.pre_call = MagicMock()
streaming = RealTimeStreaming(client_ws, backend_ws, logging_obj)
await streaming.client_ack_messages()
sent_types: Final = [json.loads(call.args[0])["type"] for call in backend_ws.send.call_args_list]
assert sent_types == ["conversation.item.create", "response.create"]
assert json.loads(backend_ws.send.call_args_list[0].args[0])["item"]["content"][0]["text"] == "What is 17 * 3?"
assert BEDROCK_REALTIME_PENDING_OPENING_MESSAGE_SCOPE_KEY not in client_ws.scope
@pytest.mark.asyncio
async def test_client_ack_messages_keeps_beta_session_shape_for_beta_clients():
client_ws = MagicMock()

View file

@ -12,8 +12,14 @@ content survives.
import asyncio
import json
from types import SimpleNamespace
from typing import AsyncIterator
from typing import AsyncIterator, Final
import pytest
from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import (
convert_to_streaming_response,
convert_to_streaming_response_async,
)
from litellm.llms.anthropic.pass_through.adapters.streaming_iterator import (
AnthropicStreamWrapper,
_CombinedChunkSplitter,
@ -139,6 +145,43 @@ def test_delayed_usage_chunk_preserves_cache_tokens():
assert message_delta["usage"]["cache_creation_input_tokens"] == 20
def _cached_bedrock_response(cache_read: int, cache_write: int) -> dict:
return {
"id": "chatcmpl-cached",
"object": "chat.completion",
"created": 1700000000,
"model": "us.anthropic.claude-opus-4-7",
"choices": [{"index": 0, "message": {"role": "assistant", "content": "pong"}, "finish_reason": "stop"}],
"usage": Usage(
prompt_tokens=42145,
completion_tokens=6,
total_tokens=42151,
cache_read_input_tokens=cache_read,
cache_creation_input_tokens=cache_write,
).model_dump(),
}
async def _replay_async(cached: dict) -> list[ModelResponseStream]:
return [chunk async for chunk in convert_to_streaming_response_async(cached)]
@pytest.mark.parametrize("cache_read, cache_write", [(42139, 0), (0, 42139)], ids=["read", "write"])
@pytest.mark.parametrize("replay", ["async", "sync"])
def test_a_cache_hit_replayed_as_a_messages_stream_keeps_the_prompt_cache_tokens(replay, cache_read, cache_write):
cached: Final = _cached_bedrock_response(cache_read, cache_write)
chunks: Final = (
asyncio.run(_replay_async(cached)) if replay == "async" else list(convert_to_streaming_response(cached))
)
events: Final = list(AnthropicStreamWrapper(completion_stream=iter(chunks), model="claude-opus-4-7"))
message_delta: Final = next(event for event in events if event.get("type") == "message_delta")
assert message_delta["usage"]["input_tokens"] == 6
assert message_delta["usage"].get("cache_read_input_tokens", 0) == cache_read
assert message_delta["usage"].get("cache_creation_input_tokens", 0) == cache_write
def test_splitter_passes_through_non_combined_chunks():
"""A chunk with content but no finish_reason is not split."""
chunk = ModelResponseStream(choices=[StreamingChoices(index=0, delta=Delta(content="partial"), finish_reason=None)])

View file

@ -17,6 +17,7 @@ from litellm.constants import (
from litellm.llms.bedrock.common_utils import BedrockError
from litellm.llms.bedrock.realtime.handler import BedrockRealtime
from litellm.llms.bedrock.realtime.transformation import BedrockRealtimeConfig
from litellm.realtime_api import main as realtime_main
@pytest.fixture(autouse=True)
@ -109,6 +110,14 @@ class UnavailableBedrockStream:
raise ServiceUnavailableException("fault injected: Bedrock realtime unavailable")
class RecordingRealtimeHandler:
def __init__(self):
self.calls = 0
async def async_realtime(self, **kwargs):
self.calls += 1
class FakeLogging:
def __init__(self, trace_id="trace-nova-sonic"):
self.litellm_trace_id = trace_id
@ -678,18 +687,23 @@ class TestBedrockRealtimeSessionLifecycle:
@pytest.mark.asyncio
@pytest.mark.parametrize(
"drop_params,expected_errors",
[(False, [("invalid_type", "session.temperature")]), (True, [])],
"session,drop_params,expected_errors",
[
({"temperature": "0.2", "modalities": ["text"]}, False, [("invalid_type", "session.temperature")]),
({"temperature": "0.2", "modalities": ["text"]}, True, []),
({"modalities": "text"}, False, [("invalid_type", "session.modalities")]),
({"modalities": "text"}, True, []),
],
)
async def test_malformed_number_setting_is_rejected_or_dropped_and_the_session_still_answers(
self, stub_aws_sdk_client, drop_params, expected_errors
async def test_malformed_setting_is_rejected_or_dropped_and_the_session_still_answers(
self, stub_aws_sdk_client, session, drop_params, expected_errors
):
stream: Final = ScriptedBedrockStream([])
stub_aws_sdk_client["streams"] = [stream]
user_text: Final = {"type": "message", "role": "user", "content": [{"type": "input_text", "text": "hi"}]}
websocket: Final = ConnectedClientWS(
[
json.dumps({"type": "session.update", "session": {"temperature": "0.2", "modalities": ["text"]}}),
json.dumps({"type": "session.update", "session": session}),
json.dumps({"type": "conversation.item.create", "item": user_text}),
]
)
@ -777,12 +791,26 @@ class TestBedrockRealtimeProviderFailurePropagation:
dispatched["logging_obj"] = RecordingLogging()
return dispatched
USER_TEXT: Final = json.dumps(
{
"type": "conversation.item.create",
"item": {"type": "message", "role": "user", "content": [{"type": "input_text", "text": "hi"}]},
}
)
REJECTED_SESSION_UPDATE: Final = json.dumps({"type": "session.update", "session": {"temperature": "0.2"}})
USER_AUDIO: Final = json.dumps({"type": "input_audio_buffer.append", "audio": "AAAA"})
@pytest.mark.asyncio
@pytest.mark.parametrize(
"client_messages",
[[SESSION_UPDATE], [USER_TEXT], [USER_AUDIO], [REJECTED_SESSION_UPDATE, USER_TEXT]],
ids=["session-update-first", "text-first", "audio-first", "input-after-rejected-update"],
)
async def test_mid_stream_failure_escapes_keeps_partial_spend_and_blocks_replay(
self, stub_aws_sdk_client, spend_dispatch
self, stub_aws_sdk_client, spend_dispatch, monkeypatch, client_messages: list[str]
):
handler = BedrockRealtime()
websocket = ConnectedClientWS([self.SESSION_UPDATE])
websocket = ConnectedClientWS(client_messages)
stream = ScriptedBedrockStream(self.TEXT_TURN, receiver_type=BreakingBedrockReceiver)
stub_aws_sdk_client["streams"] = [stream]
@ -814,6 +842,37 @@ class TestBedrockRealtimeProviderFailurePropagation:
assert not litellm.should_retry(replay.value.status_code), "the router must not retry the replay refusal"
assert "Nova Sonic stream broke" in replay.value.message, "the router surfaces the last attempt's error"
other_provider: Final = RecordingRealtimeHandler()
monkeypatch.setattr(realtime_main, "openai_realtime", other_provider)
with pytest.raises(BedrockError) as fallback:
await realtime_main._arealtime.__wrapped__(
model="openai/gpt-realtime", websocket=websocket, litellm_logging_obj=FakeLogging()
)
assert fallback.value.status_code == 400, "a fallback to another provider must not start a second session"
assert other_provider.calls == 0
@pytest.mark.asyncio
async def test_readiness_failure_on_an_input_first_start_replays_the_input(self, stub_aws_sdk_client):
handler = BedrockRealtime()
websocket = ConnectedClientWS([self.USER_TEXT])
healthy_stream = ScriptedBedrockStream([])
stub_aws_sdk_client["streams"] = [UnavailableBedrockStream(), healthy_stream]
with pytest.raises(BedrockError) as failure:
await handler.async_realtime(
model="amazon.nova-sonic-v1:0", websocket=websocket, logging_obj=FakeLogging(), **self.AWS_PARAMS
)
assert failure.value.status_code == 503
await handler.async_realtime(
model="amazon.nova-sonic-v1:0", websocket=websocket, logging_obj=FakeLogging(), **self.AWS_PARAMS
)
assert [json.loads(m)["type"] for m in websocket.sent_to_client] == ["session.created"]
replayed = [json.loads(chunk.value.bytes_.decode("utf-8")) for chunk in healthy_stream.input_stream.sent]
assert any(event["event"].get("textInput", {}).get("content") == "hi" for event in replayed)
@pytest.mark.asyncio
async def test_input_side_failure_keeps_spend_for_responses_already_delivered(
self, stub_aws_sdk_client, spend_dispatch

View file

@ -1655,9 +1655,17 @@ class TestBedrockRealtimeUserEventsAndUsage:
({"max_response_output_tokens": float("inf")}, "session.max_response_output_tokens"),
({"max_response_output_tokens": True}, "session.max_response_output_tokens"),
({"output_sample_rate_hertz": "16000"}, "session.output_sample_rate_hertz"),
({"instructions": {"role": "pirate"}}, "session.instructions"),
({"voice": {"id": "voice_123"}}, "session.voice"),
({"modalities": "text"}, "session.modalities"),
({"modalities": ["text", 1]}, "session.modalities"),
({"input_sample_rate_hertz": "16000"}, "session.input_sample_rate_hertz"),
({"tools": {"type": "function", "name": "lookup"}}, "session.tools"),
({"input_audio_format": {"type": "audio/pcm"}}, "session.input_audio_format"),
({"output_audio_format": {"type": "audio/pcmu"}}, "session.output_audio_format"),
],
)
def test_malformed_number_setting_is_rejected_without_drop_params(self, session, param):
def test_malformed_setting_is_rejected_without_drop_params(self, session, param):
rejection: Final = BedrockRealtimeConfig().session_update_rejection(
{"type": "session.update", "session": session}
)
@ -1677,6 +1685,59 @@ class TestBedrockRealtimeUserEventsAndUsage:
inference: Final = self._inference_configuration(config, session)
assert (inference["temperature"], inference["maxTokens"]) == (0.7, 1024)
def test_non_number_input_sample_rate_is_dropped_with_drop_params(self):
config: Final = BedrockRealtimeConfig(drop_params=True)
session: Final = {"input_sample_rate_hertz": "8000"}
assert config.session_update_rejection({"type": "session.update", "session": session}) is None
config.transform_session_update_event({"type": "session.update", "session": session})
assert config.input_sample_rate_hertz == BedrockRealtimeConfig().input_sample_rate_hertz
def test_non_string_instructions_and_voice_are_dropped_with_drop_params(self):
config: Final = BedrockRealtimeConfig(drop_params=True)
session: Final = {"instructions": {"role": "pirate"}, "voice": {"id": "voice_123"}}
assert config.session_update_rejection({"type": "session.update", "session": session}) is None
sent: Final = self._voice_and_system_prompt(config, session)
assert sent == self._voice_and_system_prompt(BedrockRealtimeConfig(), {})
@pytest.mark.parametrize("voice", ["alloy", "marin", "Alloy", "", " "])
def test_openai_or_empty_voice_is_rejected_without_drop_params(self, voice):
rejection: Final = BedrockRealtimeConfig().session_update_rejection(
{"type": "session.update", "session": {"voice": voice}}
)
assert rejection is not None
assert (rejection.code, rejection.param) == ("invalid_value", "session.voice")
assert f"'{voice}'" in rejection.message
def test_openai_voice_is_dropped_with_drop_params(self):
config: Final = BedrockRealtimeConfig(drop_params=True)
session: Final = {"voice": "alloy", "instructions": "Answer as a pirate."}
assert config.session_update_rejection({"type": "session.update", "session": session}) is None
voice, system_prompt = self._voice_and_system_prompt(config, session)
assert voice == self._voice_and_system_prompt(BedrockRealtimeConfig(), {})[0]
assert system_prompt == "Answer as a pirate."
def test_nova_voice_is_applied(self):
session: Final = {"voice": "tiffany"}
assert BedrockRealtimeConfig().session_update_rejection({"type": "session.update", "session": session}) is None
assert self._voice_and_system_prompt(BedrockRealtimeConfig(), session)[0] == "tiffany"
def test_g711_input_audio_format_sets_the_input_sample_rate(self):
config: Final = BedrockRealtimeConfig()
config.transform_session_update_event({"type": "session.update", "session": {"input_audio_format": "g711_ulaw"}})
assert config.input_sample_rate_hertz == 8000
@staticmethod
def _voice_and_system_prompt(config: BedrockRealtimeConfig, session: dict[str, object]) -> tuple[object, object]:
events: Final = [
json.loads(message)["event"]
for message in config.transform_session_update_event({"type": "session.update", "session": session})
]
voice: Final = next(
e["promptStart"]["audioOutputConfiguration"]["voiceId"] for e in events if "promptStart" in e
)
system_prompt: Final = next(e["textInput"]["content"] for e in events if "textInput" in e)
return voice, system_prompt
def test_null_number_settings_count_as_unset(self):
session: Final = {"temperature": None, "max_response_output_tokens": None}
config: Final = BedrockRealtimeConfig()

View file

@ -15,10 +15,12 @@ from litellm.integrations.anthropic_cache_control_hook import AnthropicCacheCont
from litellm.integrations.custom_logger import CustomLogger
from litellm.router_utils.pre_call_checks.prompt_caching_deployment_check import (
PromptCachingDeploymentCheck,
_honors_one_hour_cache_ttl,
_get_min_token_count_for_deployments,
)
from litellm.router_utils.prompt_caching_cache import PROMPT_CACHE_PIN_TTL_SECONDS, PromptCachingCache
from litellm.types.llms.openai import AllMessageValues
from litellm.types.utils import StandardLoggingPayload
from litellm.utils import get_prompt_cache_min_tokens, is_prompt_caching_valid_prompt, token_counter
MODEL_GROUP_ALIAS = "my-claude-group"
@ -126,6 +128,21 @@ def test_get_min_token_count_for_deployments_falls_back_to_default_for_empty_gro
assert _get_min_token_count_for_deployments([]) == DEFAULT_MINIMUM_PROMPT_CACHE_TOKEN_COUNT
def test_bedrock_application_profile_uses_its_base_model_for_one_hour_ttl() -> None:
standard_logging_object: Final = cast(
StandardLoggingPayload,
{
"custom_llm_provider": "bedrock",
"model": "bedrock/converse/arn:aws:bedrock:us-east-1:000000000000:application-inference-profile/pinprofile",
},
)
assert _honors_one_hour_cache_ttl(
standard_logging_object,
{"model": None, "base_model": "bedrock/us.anthropic.claude-sonnet-5-5"},
) is True
@pytest.mark.asyncio
async def test_async_filter_deployments_does_not_narrow_prompt_below_model_minimum():
"""

View file

@ -2,7 +2,7 @@ import json
import asyncio
import datetime
import time
from collections.abc import Mapping
from collections.abc import Iterator, Mapping
from types import MappingProxyType, SimpleNamespace
from typing import Final, Literal, cast
@ -17,6 +17,10 @@ from litellm import TranscriptionResponse, model_cost
from litellm.constants import REALTIME_METERED_USAGE_KEY
from litellm.cost_calculator import BaseTokenUsageProcessor, RealtimeAPITokenUsageProcessor, ResponsesWebSocketTokenUsageProcessor, batch_cost_calculator, completion_cost, cost_per_token, handle_realtime_stream_cost_calculation, response_cost_calculator
from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.fallback_generalizations import (
get_fallback_generalization_rules,
set_fallback_generalizations,
)
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
from litellm.litellm_core_utils.litellm_logging import Logging
from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import (
@ -75,6 +79,48 @@ def _local_model_cost_map(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url=""))
@pytest.fixture
def _claude_family_generalization() -> Iterator[None]:
previous_rules: Final = list(get_fallback_generalization_rules())
try:
set_fallback_generalizations(
[
{
"name": "claude-family-baseline",
"pattern": r"claude-[a-z]+-\d+(?:[-.]\d+)?(?:-\d{8})?",
"model_info": {"mode": "chat", "max_input_tokens": 200000, "max_output_tokens": 64000},
}
]
)
yield
finally:
set_fallback_generalizations(previous_rules)
@pytest.fixture
def _priced_claude_family_generalization() -> Iterator[None]:
previous_rules: Final = list(get_fallback_generalization_rules())
try:
set_fallback_generalizations(
[
{
"name": "claude-family-priced",
"pattern": r"claude-[a-z]+-\d+(?:[-.]\d+)?(?:-\d{8})?",
"model_info": {
"mode": "chat",
"max_input_tokens": 200000,
"max_output_tokens": 64000,
"input_cost_per_token": 0.0000321,
"output_cost_per_token": 0.0000765,
},
}
]
)
yield
finally:
set_fallback_generalizations(previous_rules)
def test_cost_per_token_duplicate_openai_prefix_matches_model_cost(monkeypatch):
"""
Router/proxy configs may use deployment ids like openai/openai/<model>. Cost lookup must
@ -149,6 +195,265 @@ def test_completion_cost_uses_response_model_for_dynamic_routing(_local_model_co
assert cost > 0, "Cost should be calculated using response model"
def test_completion_cost_prefers_exact_route_price_over_capability_generalization(
_local_model_cost_map: None,
_claude_family_generalization: None,
monkeypatch: pytest.MonkeyPatch,
) -> None:
priced_model: Final = "bedrock/invoke/anthropic.claude-haiku-4-5-20251001-v1:0"
priced_entry: Final = {
"litellm_provider": "bedrock",
"mode": "chat",
"input_cost_per_token": 0.0000123,
"output_cost_per_token": 0.0000456,
}
monkeypatch.setattr(litellm, "model_cost", {priced_model: priced_entry})
usage: Final = Usage(prompt_tokens=17, completion_tokens=9, total_tokens=26)
response: Final = ModelResponse(
model="anthropic.claude-haiku-4-5-20251001-v1:0",
choices=[],
usage=usage,
)
cost: Final = completion_cost(
completion_response=response,
model=priced_model,
custom_llm_provider="bedrock",
)
input_rate: Final = priced_entry["input_cost_per_token"]
output_rate: Final = priced_entry["output_cost_per_token"]
assert isinstance(input_rate, (int, float)), priced_entry
assert isinstance(output_rate, (int, float)), priced_entry
assert cost == usage.prompt_tokens * input_rate + usage.completion_tokens * output_rate
def test_completion_cost_keeps_priced_generalization_before_later_exact_candidate(
_local_model_cost_map: None,
_priced_claude_family_generalization: None,
monkeypatch: pytest.MonkeyPatch,
) -> None:
priced_model: Final = "bedrock/invoke/anthropic.claude-haiku-4-5-20251001-v1:0"
exact_entry: Final = {
"litellm_provider": "bedrock",
"mode": "chat",
"input_cost_per_token": 0.0000123,
"output_cost_per_token": 0.0000456,
}
monkeypatch.setattr(litellm, "model_cost", {priced_model: exact_entry})
usage: Final = Usage(prompt_tokens=17, completion_tokens=9, total_tokens=26)
response: Final = ModelResponse(
model="anthropic.claude-haiku-4-5-20251001-v1:0",
choices=[],
usage=usage,
)
cost: Final = completion_cost(
completion_response=response,
model=priced_model,
custom_llm_provider="bedrock",
)
assert cost == pytest.approx(usage.prompt_tokens * 0.0000321 + usage.completion_tokens * 0.0000765)
@pytest.fixture
def _per_second_priced_claude_family_generalization() -> Iterator[None]:
previous_rules: Final = list(get_fallback_generalization_rules())
try:
set_fallback_generalizations(
[
{
"name": "claude-family-per-second",
"pattern": r"claude-[a-z]+-\d+(?:[-.]\d+)?(?:-\d{8})?",
"model_info": {"cost_per_second": 0.25},
}
]
)
yield
finally:
set_fallback_generalizations(previous_rules)
def test_completion_cost_keeps_per_second_priced_generalization_before_later_exact_candidate(
_local_model_cost_map: None,
_per_second_priced_claude_family_generalization: None,
monkeypatch: pytest.MonkeyPatch,
) -> None:
priced_model: Final = "bedrock/invoke/anthropic.claude-haiku-4-5-20251001-v1:0"
monkeypatch.setattr(
litellm,
"model_cost",
{
priced_model: {
"litellm_provider": "bedrock",
"mode": "chat",
"input_cost_per_token": 0.0000123,
"output_cost_per_token": 0.0000456,
}
},
)
response: Final = ModelResponse(
model="anthropic.claude-haiku-4-5-20251001-v1:0",
choices=[],
usage=Usage(prompt_tokens=17, completion_tokens=9, total_tokens=26),
)
cost: Final = completion_cost(
completion_response=response,
model=priced_model,
custom_llm_provider="bedrock",
total_time=2000,
)
assert cost == pytest.approx(0.25 * 2)
@pytest.fixture
def _price_free_gpt_family_generalization() -> Iterator[None]:
previous_rules: Final = list(get_fallback_generalization_rules())
try:
set_fallback_generalizations(
[
{
"name": "gpt-family-baseline",
"pattern": r"gpt-\d+(?:\.\d+)?",
"model_info": {"mode": "chat", "supports_reasoning": True},
}
]
)
yield
finally:
set_fallback_generalizations(previous_rules)
def test_completion_cost_prices_rule_only_response_model_at_the_azure_deployment_rate(
_local_model_cost_map: None,
_price_free_gpt_family_generalization: None,
monkeypatch: pytest.MonkeyPatch,
) -> None:
deployment_model: Final = "azure/gpt-4o-2024-11-20"
monkeypatch.setattr(
litellm,
"model_cost",
{
deployment_model: {
"litellm_provider": "azure",
"mode": "chat",
"input_cost_per_token": 0.0000111,
"output_cost_per_token": 0.0000222,
}
},
)
usage: Final = Usage(prompt_tokens=17, completion_tokens=9, total_tokens=26)
response: Final = ModelResponse(model="gpt-9.9", choices=[], usage=usage)
cost: Final = completion_cost(
completion_response=response,
model=deployment_model,
custom_llm_provider="azure",
)
assert cost == pytest.approx(usage.prompt_tokens * 0.0000111 + usage.completion_tokens * 0.0000222)
def test_completion_cost_orders_candidates_without_extra_huggingface_config_requests(
_local_model_cost_map: None,
) -> None:
model: Final = "lazy-candidate-regression-org/lazy-candidate-model"
usage: Final = Usage(prompt_tokens=17, completion_tokens=9, total_tokens=26)
response: Final = ModelResponse(
model=model,
choices=[],
usage=usage,
)
with respx.mock(assert_all_called=False) as huggingface:
huggingface.get(host="huggingface.co").mock(
return_value=httpx.Response(200, json={"max_position_embeddings": 4096})
)
cost: Final = completion_cost(
completion_response=response,
model=f"huggingface/{model}",
custom_llm_provider="huggingface",
)
assert len(huggingface.calls) == 1
assert cost == 0.0
def test_completion_cost_skips_huggingface_config_request_with_custom_pricing(
_local_model_cost_map: None,
) -> None:
model: Final = "unmapped-cost-regression-org/unmapped-model"
usage: Final = Usage(prompt_tokens=17, completion_tokens=9, total_tokens=26)
response: Final = ModelResponse(
model=model,
choices=[],
usage=usage,
)
with respx.mock(assert_all_called=False) as huggingface:
huggingface.get(host="huggingface.co").mock(
return_value=httpx.Response(200, json={"max_position_embeddings": 4096})
)
cost: Final = completion_cost(
completion_response=response,
model=model,
custom_llm_provider="huggingface",
custom_pricing=True,
custom_cost_per_token={"input_cost_per_token": 0.001, "output_cost_per_token": 0.002},
)
assert len(huggingface.calls) == 0
assert cost == pytest.approx(usage.prompt_tokens * 0.001 + usage.completion_tokens * 0.002)
def test_completion_cost_returns_zero_for_capability_only_generalization(
_local_model_cost_map: None,
_claude_family_generalization: None,
monkeypatch: pytest.MonkeyPatch,
) -> None:
model: Final = "bedrock/invoke/anthropic.claude-haiku-9-9-20991231-v1:0"
monkeypatch.setattr(litellm, "model_cost", {})
response: Final = ModelResponse(
model="anthropic.claude-haiku-9-9-20991231-v1:0",
choices=[],
usage=Usage(prompt_tokens=17, completion_tokens=9, total_tokens=26),
)
cost: Final = completion_cost(
completion_response=response,
model=model,
custom_llm_provider="bedrock",
)
assert cost == 0.0
@pytest.mark.parametrize("custom_llm_provider", ["openai", "hosted_vllm"])
def test_completion_cost_prices_base_model_under_its_own_provider(
custom_llm_provider: str, _local_model_cost_map: None
) -> None:
response: Final = ModelResponse(
model="my-deploy",
choices=[],
usage=Usage(prompt_tokens=1000, completion_tokens=1000, total_tokens=2000),
)
cost: Final = completion_cost(
completion_response=response,
model=f"{custom_llm_provider}/my-deploy",
custom_llm_provider=custom_llm_provider,
base_model="claude-sonnet-4-5",
)
info: Final = litellm.get_model_info("claude-sonnet-4-5", custom_llm_provider="anthropic")
assert cost == pytest.approx(
1000 * info["input_cost_per_token"] + 1000 * info["output_cost_per_token"]
)
def test_completion_cost_strips_dated_azure_snapshot_model(_local_model_cost_map: None) -> None:
dated_response = ModelResponse(
model="gpt-5.6-luna-2099-01-01",

View file

@ -4556,7 +4556,7 @@ def test_stream_chunk_builder_skips_stamp_when_cost_is_unpriceable():
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
logging_obj: Final = LiteLLMLogging(
model="us.anthropic.claude-opus-5",
model="us.anthropic.claude-unpriced-999",
messages=[{"role": "user", "content": "hi"}],
stream=True,
call_type="completion",
@ -4582,6 +4582,41 @@ def test_stream_chunk_builder_skips_stamp_when_cost_is_unpriceable():
assert response._hidden_params.get("response_cost") is None
def test_stream_chunk_builder_prices_alias_response_at_its_deployment_model():
import time as time_module
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
deployment_model: Final = "us.anthropic.claude-opus-5"
logging_obj: Final = LiteLLMLogging(
model=deployment_model,
messages=[{"role": "user", "content": "hi"}],
stream=True,
call_type="completion",
start_time=time_module.time(),
litellm_call_id="stream-builder-alias-priced",
function_id="1",
)
logging_obj.model_call_details["custom_llm_provider"] = "bedrock"
logging_obj.optional_params = {}
usage_chunk: Final = _stream_builder_text_chunk("bedrock-claude-opus-5", "")
usage_chunk.usage = Usage(prompt_tokens=40, completion_tokens=5, total_tokens=45)
chunks: Final = [
_stream_builder_text_chunk("bedrock-claude-opus-5", "Hello ", finish_reason="stop"),
usage_chunk,
]
response: Final = litellm.stream_chunk_builder(
chunks=chunks, messages=[{"role": "user", "content": "hi"}], logging_obj=logging_obj
)
entry: Final = litellm.model_cost[deployment_model]
expected_cost: Final = 40 * entry["input_cost_per_token"] + 5 * entry["output_cost_per_token"]
assert response is not None
assert getattr(response.usage, "cost", None) == pytest.approx(expected_cost)
assert response._hidden_params.get("response_cost") == pytest.approx(expected_cost)
def test_stream_chunk_builder_keeps_provider_reported_usage_cost():
usage_chunk: Final = _stream_builder_text_chunk("gpt-4o", "")
usage_chunk.usage = Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15, cost=0.5)