mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-27 06:13:07 +00:00
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 commit3ef0b35255) * 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 commite754857de9) * 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 commit6b45a196c4) * 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 commitc32eea634a) * 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 commitd346d769f0) * 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 commita4e0570bea) * 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 commit4a58784fa2) * 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 commit1fd7f81f97) * 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:
parent
a00d8a6ee1
commit
4bf1633452
36 changed files with 2099 additions and 304 deletions
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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: ...
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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 "
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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))
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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())
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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():
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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())
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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)])
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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():
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue