diff --git a/litellm/litellm_core_utils/get_model_cost_map.py b/litellm/litellm_core_utils/get_model_cost_map.py index 6e3769260cc..b68db22183a 100644 --- a/litellm/litellm_core_utils/get_model_cost_map.py +++ b/litellm/litellm_core_utils/get_model_cost_map.py @@ -324,10 +324,16 @@ def _classify_fetch_content(content: bytes, url: str, etag: str | None) -> _Fetc return ModelCostMapReloaded(model_cost_map=model_cost_map, revision=git_blob_id(content), etag=etag) -def _attempt_fetch_s3(url: str, timeout: int) -> _FetchAttemptOutcome: +def _split_url(url: str) -> urllib.parse.SplitResult | ModelCostMapReloadUnavailable: + try: + return urllib.parse.urlsplit(url) + except ValueError as error: + return ModelCostMapReloadUnavailable(reason=f"invalid URL {url}: {error}") + + +def _attempt_fetch_s3(parts: urllib.parse.SplitResult, url: str, timeout: int) -> _FetchAttemptOutcome: from botocore.exceptions import BotoCoreError, ClientError - parts: Final = urllib.parse.urlsplit(url) bucket: Final = parts.netloc key: Final = parts.path.lstrip("/") if not bucket or not key: @@ -343,9 +349,11 @@ def _attempt_fetch_s3(url: str, timeout: int) -> _FetchAttemptOutcome: async def _attempt_fetch(client: _AsyncGetClient, url: str, timeout: int) -> _FetchAttemptOutcome: - parts: Final = urllib.parse.urlsplit(url) + parts: Final = _split_url(url) + if isinstance(parts, ModelCostMapReloadUnavailable): + return parts if parts.scheme == "s3": - return await asyncio.to_thread(_attempt_fetch_s3, url, timeout) + return await asyncio.to_thread(_attempt_fetch_s3, parts, url, timeout) try: response: Final = await client.get(url, timeout=timeout) except (httpx.HTTPError, httpx.InvalidURL) as e: @@ -354,9 +362,11 @@ async def _attempt_fetch(client: _AsyncGetClient, url: str, timeout: int) -> _Fe def _attempt_fetch_sync(client: _SyncGetClient, url: str, timeout: int) -> _FetchAttemptOutcome: - parts: Final = urllib.parse.urlsplit(url) + parts: Final = _split_url(url) + if isinstance(parts, ModelCostMapReloadUnavailable): + return parts if parts.scheme == "s3": - return _attempt_fetch_s3(url, timeout) + return _attempt_fetch_s3(parts, url, timeout) try: response: Final = client.get(url, timeout=timeout) except (httpx.HTTPError, httpx.InvalidURL) as e: diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 4ff4c43428f..1bccef7dd9e 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -216,6 +216,7 @@ if TYPE_CHECKING: FakeAnthropicMessagesStreamIterator, ) from litellm.llms.base_llm.passthrough.transformation import BasePassthroughConfig + from litellm.responses.streaming_iterator import ResponsesWsFrameGuard from litellm.types.google_genai.main import GenerateContentResponse from litellm.types.llms.openai_evals import ( CancelEvalResponse, @@ -6867,6 +6868,7 @@ class BaseLLMHTTPHandler: custom_llm_provider: str | None = None, first_message: str | None = None, request_defaults: ResponsesWebSocketRequestDefaults | None = None, + frame_guard: "ResponsesWsFrameGuard | None" = None, **kwargs: object, ) -> Exception | None: """ @@ -6882,6 +6884,35 @@ class BaseLLMHTTPHandler: """ _ws_quota_callbacks: Final = _collect_ws_project_quota_callbacks() + _ws_guardrail_callbacks: list = [] + _ws_output_guardrail_callbacks: list = [] + try: + import litellm as _litellm + + # Use duck-typing so any guardrail that exposes the PII + # masking interface works, not just _OPTIONAL_PresidioPIIMasking. + # This avoids a layering violation (SDK importing from proxy). + _ws_guardrail_callbacks = [ + cb + for cb in _litellm.callbacks + if callable(getattr(cb, "check_pii", None)) + and callable(getattr(cb, "get_presidio_settings_from_request_data", None)) + and callable(getattr(cb, "_unmask_pii_text", None)) + and getattr(cb, "output_parse_pii", False) + ] + _ws_output_guardrail_callbacks = [ + cb + for cb in _litellm.callbacks + if callable(getattr(cb, "check_pii", None)) + and callable(getattr(cb, "get_presidio_settings_from_request_data", None)) + and getattr(cb, "apply_to_output", False) + ] + except Exception as _guardrail_exc: + verbose_logger.warning( + "Responses WebSocket: failed to collect guardrail callbacks — PII masking will be skipped. Error: %s", + _guardrail_exc, + ) + if responses_api_provider_config is None or not responses_api_provider_config.supports_native_websocket(): from litellm.responses.streaming_iterator import ( ManagedResponsesWebSocketHandler, @@ -6899,6 +6930,8 @@ class BaseLLMHTTPHandler: custom_llm_provider=custom_llm_provider, first_message=first_message, quota_callbacks=_ws_quota_callbacks, + guardrail_callbacks=_ws_guardrail_callbacks, # pyright: ignore[reportUnknownArgumentType] # duck-typed callbacks from litellm.callbacks + frame_guard=frame_guard, **kwargs, ) await handler.run() @@ -6980,36 +7013,6 @@ class BaseLLMHTTPHandler: if litellm_metadata: _request_data["litellm_metadata"] = litellm_metadata - _ws_guardrail_callbacks: list = [] - _ws_output_guardrail_callbacks: list = [] - try: - import litellm as _litellm - - # Use duck-typing so any guardrail that exposes the PII - # masking interface works, not just _OPTIONAL_PresidioPIIMasking. - # This avoids a layering violation (SDK importing from proxy). - _ws_guardrail_callbacks = [ - cb - for cb in _litellm.callbacks - if callable(getattr(cb, "check_pii", None)) - and callable(getattr(cb, "get_presidio_settings_from_request_data", None)) - and callable(getattr(cb, "_unmask_pii_text", None)) - and getattr(cb, "output_parse_pii", False) - ] - _ws_output_guardrail_callbacks = [ - cb - for cb in _litellm.callbacks - if callable(getattr(cb, "check_pii", None)) - and callable(getattr(cb, "get_presidio_settings_from_request_data", None)) - and getattr(cb, "apply_to_output", False) - ] - except Exception as _guardrail_exc: - verbose_logger.warning( - "Responses WebSocket: failed to collect guardrail " - "callbacks — PII masking will be skipped. Error: %s", - _guardrail_exc, - ) - streaming: Final = ResponsesWebSocketStreaming( websocket=websocket, backend_ws=cast(ClientConnection, backend_ws), @@ -7023,6 +7026,7 @@ class BaseLLMHTTPHandler: authorized_model=model, custom_llm_provider=custom_llm_provider, request_defaults=request_defaults, + frame_guard=frame_guard, ) return await streaming.bidirectional_forward() diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 58338f185c7..8545a5b568e 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -65561,6 +65561,40 @@ "supports_vision": false, "source": "https://deepinfra.com/pricing" }, + "deepinfra/deepseek-ai/DeepSeek-V4.1-Flash": { + "max_tokens": 131072, + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "input_cost_per_token": 1.4e-07, + "output_cost_per_token": 4.2e-07, + "cache_read_input_token_cost": 4.2e-09, + "supports_prompt_caching": true, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true, + "supports_reasoning": true, + "supports_vision": true, + "source": "https://deepinfra.com/deepseek-ai/DeepSeek-V4.1-Flash" + }, + "deepinfra/deepseek-ai/DeepSeek-V4.1-Flash-1017": { + "max_tokens": 131072, + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "input_cost_per_token": 1.4e-07, + "output_cost_per_token": 4.2e-07, + "cache_read_input_token_cost": 4.2e-09, + "supports_prompt_caching": true, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true, + "supports_reasoning": true, + "supports_vision": true, + "source": "https://deepinfra.com/deepseek-ai/DeepSeek-V4.1-Flash-1017" + }, "deepinfra/openai/gpt-oss-120b-Ultra": { "max_tokens": 131072, "max_input_tokens": 131072, diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 27bfd03d7f8..1e16eb626fe 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -1,6 +1,7 @@ import asyncio import copy import enum +import gc import importlib import inspect import io @@ -1746,6 +1747,8 @@ async def proxy_startup_event(app: FastAPI) -> AsyncGenerator[ProxyLifespanState async with AsyncExitStack() as admin_mcp_stack: try: await admin_mcp_stack.enter_async_context(admin_mcp_lifespan(app)) + gc.freeze() + verbose_proxy_logger.warning("gc.freeze pinned %d startup objects", gc.get_freeze_count()) yield state finally: if model_info_scheduler is not None and model_info_scheduler.running: diff --git a/litellm/proxy/response_api_endpoints/endpoints.py b/litellm/proxy/response_api_endpoints/endpoints.py index 207e966ef28..aebcab75601 100644 --- a/litellm/proxy/response_api_endpoints/endpoints.py +++ b/litellm/proxy/response_api_endpoints/endpoints.py @@ -1,6 +1,7 @@ import asyncio import contextlib import json +import uuid from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Callable, Mapping, Sequence from dataclasses import dataclass from enum import Enum @@ -28,10 +29,11 @@ from pydantic import ConfigDict, TypeAdapter, ValidationError from starlette.websockets import WebSocket, WebSocketDisconnect from typing_extensions import ReadOnly, TypedDict +import litellm from litellm._logging import verbose_proxy_logger from litellm.constants import EMPTY_MAPPING, MANAGED_BACKGROUND_RESPONSE_METADATA_KEY, TERMINAL_RESPONSE_STATUSES from litellm.integrations.custom_guardrail import ModifyResponseException -from litellm.litellm_core_utils.core_helpers import as_str_mapping +from litellm.litellm_core_utils.core_helpers import as_str_mapping, get_metadata_variable_name_from_kwargs from litellm.litellm_core_utils.internal_call_metadata import is_background_response, requests_background from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup from litellm.litellm_core_utils.tool_fields import first_validation_error @@ -58,11 +60,17 @@ from litellm.proxy.common_utils.http_parsing_utils import ( # noqa: F401 # leg read_request_body, safe_set_request_parsed_body, ) +from litellm.proxy.hooks.proxy_track_cost_callback import ProxyDBLogger from litellm.proxy.hooks.responses_id_security import ResponsesIDSecurity, decrypted_response_id +from litellm.proxy.litellm_pre_call_utils import ( + _CLIENT_PRICING_METADATA_FIELDS, # pyright: ignore[reportPrivateUsage] # the HTTP path's own strip list + _UNTRUSTED_METADATA_CONTROL_FIELDS, # pyright: ignore[reportPrivateUsage] # the HTTP path's own strip list +) from litellm.proxy.pass_through_endpoints.llm_provider_handlers.batch_attribution import ( request_tags_from_request_data, ) from litellm.proxy.route_llm_request import raise_if_required_body_param_missing +from litellm.responses.streaming_iterator import GuardedResponsesWsFrame, ResponsesWsFrameGuard from litellm.router import Router from litellm.types.llms.base import LiteLLMBaseModel from litellm.types.llms.openai import ( @@ -1810,6 +1818,143 @@ async def _enforce_responses_ws_first_frame_model_auth( ) +_WS_FRAME: Final = TypeAdapter(dict[str, object]) +_GUARDRAIL_INFORMATION: Final = TypeAdapter(list[dict[str, object]]) +_PII_TOKENS: Final = TypeAdapter(dict[str, str]) +_GUARDRAIL_INFORMATION_KEY: Final = "standard_logging_guardrail_information" +_GUARDED_FRAME_FIELDS: Final = ("input", "instructions", "tools") +_CLIENT_METADATA_KEYS: Final = ("metadata", "litellm_metadata") +# the guardrail list is merged from key, team and the admission frame; a later frame cannot shrink it +_CLIENT_GUARDRAIL_SELECTORS: Final = ("guardrails", "guardrail_config", "policies") + + +class _PreCallGuardrailWalk(Protocol): + async def pre_call_hook( + self, + *, + user_api_key_dict: UserAPIKeyAuth, + data: dict[str, object], # mutable-ok: ProxyLogging.pre_call_hook hands the guardrails a dict they edit + call_type: Literal["aresponses"], + guardrails_only: bool, + ) -> dict[str, object]: ... # mutable-ok: ProxyLogging.pre_call_hook returns the guardrails' own dict + + +def _responses_ws_frame_guard( + admitted: Mapping[str, object], user_api_key_dict: UserAPIKeyAuth, proxy_logging_obj: _PreCallGuardrailWalk +) -> ResponsesWsFrameGuard: + """Socket admission skips guardrails (the first frame's ``input`` is not part of the admitted request), so every + ``response.create`` frame gets the pre_call guardrail walk the HTTP ``/v1/responses`` path runs on its body.""" + metadata_key: Final = get_metadata_variable_name_from_kwargs(admitted) + admitted_bucket: Final = admitted.get(metadata_key) + bucket: Final = { + key: value + for key, value in ( + _WS_FRAME.validate_python(admitted_bucket) if isinstance(admitted_bucket, dict) else {} + ).items() + if key != _GUARDRAIL_INFORMATION_KEY + } + model: Final = admitted.get("model") + + # a blocked frame is billed as its own failed request; only the first turn settles the admission reservation + rejected_frame_auth: Final = user_api_key_dict.model_copy(update={"budget_reservation": None}) + + async def guard(frame: str) -> GuardedResponsesWsFrame: + frame_obj: Final = _WS_FRAME.validate_json(frame) + nested: Final = frame_obj.get("response") + if isinstance(nested, dict) and any(key in frame_obj for key in _GUARDED_FRAME_FIELDS): + raise ValueError( + "response.create carries input, instructions or tools both at the top level and under response" + ) + body: Final = _WS_FRAME.validate_python(nested) if isinstance(nested, dict) else frame_obj + frame_data: Final = { + **{ + key: _without_untrusted_metadata(key, value) + for key, value in body.items() + if key != "type" and key not in _CLIENT_GUARDRAIL_SELECTORS + }, + "model": model, + "litellm_call_id": str(uuid.uuid4()), + # shallow: the bucket carries proxy-owned objects (the OTEL parent span) that cannot be deep-copied + metadata_key: {**bucket}, + } + try: + guarded_data: Final = await proxy_logging_obj.pre_call_hook( + user_api_key_dict=user_api_key_dict, data=frame_data, call_type="aresponses", guardrails_only=True + ) + except Exception as e: + await _bill_rejected_frame(rejected_frame_auth, e, frame_data) + raise + if guarded_data.get("model") != model: + rerouted: Final = ValueError( + f"a guardrail rerouted this turn to {guarded_data.get('model')!r}; an open Responses WebSocket" + f" is bound to {model!r}, send this turn over HTTP /v1/responses" + ) + await _bill_rejected_frame(rejected_frame_auth, rerouted, frame_data) + raise rerouted + # a guardrail may delete a field (the Responses translation pops ``instructions``), so start without them + guarded_fields: Final = {key: guarded_data[key] for key in _GUARDED_FRAME_FIELDS if key in guarded_data} + kept: Final = {key: value for key, value in body.items() if key not in _GUARDED_FRAME_FIELDS} + rewritten: Final = ( + {**frame_obj, "response": {**kept, **guarded_fields}} + if isinstance(nested, dict) + else {**kept, **guarded_fields} + ) + return GuardedResponsesWsFrame( + frame=json.dumps(rewritten), + guardrail_information=_recorded_guardrail_information(guarded_data.get(metadata_key)), + pii_tokens=_recorded_pii_tokens(guarded_data.get("metadata")), + ) + + return guard + + +def _without_untrusted_metadata(key: str, value: object) -> object: + """Like the HTTP path, drop the proxy-owned ``user_api_key_*`` and guardrail control keys a client may forge.""" + if key not in _CLIENT_METADATA_KEYS or not isinstance(value, dict): + return value + return { + name: item + for name, item in _WS_FRAME.validate_python(value).items() + if not name.startswith("user_api_key_") + and name not in _UNTRUSTED_METADATA_CONTROL_FIELDS + and name not in _CLIENT_PRICING_METADATA_FIELDS + and name not in ("headers", "pii_tokens", *_CLIENT_GUARDRAIL_SELECTORS) + } + + +async def _bill_rejected_frame( + user_api_key_dict: UserAPIKeyAuth, + exception: Exception, + frame_data: dict[str, object], # mutable-ok: the spend logger stamps the failure row on the request dict +) -> None: + """Only the spend logger sees a rejected frame; the other failure hooks (the rate limiter's slot release, + alerting) belong to the live connection, which stays open. Billing never replaces the guardrail's own error.""" + for callback in litellm.callbacks: + if isinstance(callback, ProxyDBLogger): + try: + await callback.async_post_call_failure_hook( # pyright: ignore[reportUnknownMemberType] # bare dict + request_data=frame_data, original_exception=exception, user_api_key_dict=user_api_key_dict + ) + except Exception: # noqa: BLE001 # billing never replaces the guardrail's own error + verbose_proxy_logger.exception("[Non-Blocking] Responses WebSocket: billing a rejected frame failed") + + +def _recorded_pii_tokens(metadata: object) -> Mapping[str, str]: + """Presidio stores the tokens it masked a frame with in ``metadata.pii_tokens``; the socket unmasks with them.""" + recorded: Final = _WS_FRAME.validate_python(metadata).get("pii_tokens") if isinstance(metadata, dict) else None + return MappingProxyType(_PII_TOKENS.validate_python(recorded)) if isinstance(recorded, dict) else _NO_PII_TOKENS + + +_NO_PII_TOKENS: Final[Mapping[str, str]] = MappingProxyType({}) + + +def _recorded_guardrail_information(bucket: object) -> tuple[Mapping[str, object], ...]: + recorded: Final = ( + _WS_FRAME.validate_python(bucket).get(_GUARDRAIL_INFORMATION_KEY) if isinstance(bucket, dict) else None + ) + return tuple(_GUARDRAIL_INFORMATION.validate_python(recorded)) if recorded else () + + async def _responses_websocket_session( websocket: WebSocket, model: str | None, @@ -1885,6 +2030,7 @@ async def _responses_websocket_session( user_api_base=user_api_base, model=resolved_model, route_type="_aresponses_websocket", + skip_guardrails=True, ) except Exception as e: verbose_proxy_logger.exception("Responses WebSocket pre-call error") @@ -1906,7 +2052,10 @@ async def _responses_websocket_session( return routed_data: Final = dict( - data, user_api_key_dict=user_api_key_dict, **_routing_hints_from_first_ws_frame(first_message) + data, + user_api_key_dict=user_api_key_dict, + responses_ws_frame_guard=_responses_ws_frame_guard(data, user_api_key_dict, proxy_logging_obj), + **_routing_hints_from_first_ws_frame(first_message), ) # Phase 2: route to upstream provider try: diff --git a/litellm/responses/main.py b/litellm/responses/main.py index b7385c3b3d5..d32d48c51ca 100644 --- a/litellm/responses/main.py +++ b/litellm/responses/main.py @@ -62,6 +62,7 @@ from litellm.types.llms.openai import ( # Handle ResponseText import with fallback if TYPE_CHECKING: + from litellm.responses.streaming_iterator import ResponsesWsFrameGuard from litellm.types.llms.openai import ResponseText else: ResponseText = str # Fallback for ResponseText import @@ -2598,6 +2599,7 @@ async def _aresponses_websocket( api_base: str | None = None, api_key: str | None = None, timeout: float | None = None, + responses_ws_frame_guard: "ResponsesWsFrameGuard | None" = None, **kwargs, ) -> Exception | None: """ @@ -2690,5 +2692,6 @@ async def _aresponses_websocket( litellm_metadata=_build_litellm_metadata_for_ws(kwargs), custom_llm_provider=_custom_llm_provider, request_defaults=_build_responses_websocket_request_defaults(deployment_kwargs), + frame_guard=responses_ws_frame_guard, **remaining_kwargs, ) diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 5897b0ac954..145189616af 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -7,12 +7,23 @@ import time import traceback import uuid from collections.abc import AsyncIterable, Awaitable, Callable, Coroutine, Iterable, Mapping, Sequence -from dataclasses import dataclass +from dataclasses import dataclass, field from datetime import datetime from functools import lru_cache from itertools import chain from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, Literal, NoReturn, Protocol, cast, overload, runtime_checkable +from typing import ( + TYPE_CHECKING, + Any, + Final, + Literal, + NoReturn, + Protocol, + TypeAlias, + cast, + overload, + runtime_checkable, +) import httpx from openai._streaming import SSEDecoder @@ -1932,6 +1943,73 @@ RESPONSES_WS_MASKABLE_TEXT_BLOCK_TYPES: Final = frozenset({"input_text", "output _RESPONSES_WS_FAILURE_EVENT_TYPES: Final = frozenset({"error", "response.failed"}) _PER_TURN_METADATA_KEYS: Final = frozenset({"metadata", "litellm_metadata"}) +_GUARDRAIL_INFORMATION_KEY: Final = "standard_logging_guardrail_information" + + +_PII_DELTA_EVENT_TYPES: Final = frozenset( + { + "response.output_text.delta", + "response.reasoning_summary_text.delta", + "response.refusal.delta", + "response.function_call_arguments.delta", + } +) + + +def _unmask_responses_event(response_str: str, pii_tokens: Mapping[str, str], unmask_pii_text: _UnmasksPiiText) -> str: + """Put the original PII back into one client-bound event: ``response.completed`` output text, or a delta.""" + try: + evt_obj: Final = _load_json_object(response_str) + except (json.JSONDecodeError, TypeError): + return response_str + event_type: Final = evt_obj.get("type") + + if event_type == "response.completed": + modified = False + response_obj: Final = evt_obj.get("response") + if not _is_json_object(response_obj): + return response_str + output_items: Final = response_obj.get("output") + for output_item in output_items if _is_json_array(output_items) else []: + if not _is_json_object(output_item): + continue + content = output_item.get("content") + if not _is_json_array(content): + continue + for content_block in content: + if not _is_json_object(content_block): + continue + text = content_block.get("text") + if isinstance(text, str): + unmasked = unmask_pii_text(text, pii_tokens) + if unmasked != text: + content_block["text"] = unmasked + modified = True + return json.dumps(evt_obj) if modified else response_str + + if event_type in _PII_DELTA_EVENT_TYPES: + delta: Final = evt_obj.get("delta") + if isinstance(delta, str): + unmasked = unmask_pii_text(delta, pii_tokens) + if unmasked != delta: + evt_obj["delta"] = unmasked + return json.dumps(evt_obj) + + return response_str + + +@dataclass(frozen=True, slots=True) +class GuardedResponsesWsFrame: + """A ``response.create`` frame after the proxy's pre_call guardrails ran on it: the frame to forward (guardrails + may rewrite ``input`` and ``instructions``), the guardrail records that bill against this turn only, and the + Presidio ``pii_tokens`` the walk masked with, which the handler needs to unmask the provider's reply.""" + + frame: str + guardrail_information: tuple[Mapping[str, object], ...] = () + pii_tokens: Mapping[str, str] = field(default_factory=lambda: MappingProxyType({})) + + +ResponsesWsFrameGuard: TypeAlias = Callable[[str], Awaitable[GuardedResponsesWsFrame]] _RESPONSES_WS_TERMINAL_EVENT_TYPES: Final = _HISTORY_TERMINAL_EVENT_TYPES | _RESPONSES_WS_FAILURE_EVENT_TYPES @@ -1964,6 +2042,30 @@ def _user_texts_in_input_item(item: object) -> tuple[str, ...]: ) +def _set_turn_guardrail_information( + turn_logging_obj: LiteLLMLoggingObj, guardrail_information: tuple[Mapping[str, object], ...] +) -> None: + """A turn bills the guardrails that ran on its own ``response.create`` frame, never the records the + connection-level logging object carries from an earlier frame.""" + litellm_params: Final = _json_object_or_none( + turn_logging_obj.litellm_params # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # Logging.litellm_params is untyped + ) + if litellm_params is None: + return + for key in _PER_TURN_METADATA_KEYS: + _replace_guardrail_information(litellm_params.get(key), guardrail_information) + + +def _replace_guardrail_information(bucket: object, guardrail_information: tuple[Mapping[str, object], ...]) -> None: + if not _is_json_object(bucket): + return + bucket.pop(_GUARDRAIL_INFORMATION_KEY, None) + if guardrail_information: + bucket[_GUARDRAIL_INFORMATION_KEY] = [ # rebind-ok: the turn copy's own metadata bucket + dict(entry) for entry in guardrail_information + ] + + def _logged_user_messages(client_frame: object) -> tuple[Mapping[str, object], ...]: frame: Final = _json_object_or_none(client_frame) if frame is None or frame.get("type") != "response.create": @@ -2092,11 +2194,13 @@ class ResponsesWebSocketStreaming: authorized_model: str | None = None, custom_llm_provider: str | None = None, request_defaults: ResponsesWebSocketRequestDefaults | None = None, + frame_guard: ResponsesWsFrameGuard | None = None, ): self.websocket = websocket self.backend_ws = backend_ws self.logging_obj = logging_obj self.user_api_key_dict = user_api_key_dict + self.frame_guard: ResponsesWsFrameGuard | None = frame_guard self.request_data: dict[str, object] = request_data or {} litellm_metadata: Final = self.request_data.get("litellm_metadata") self.litellm_metadata: dict[str, object] = litellm_metadata if _is_json_object(litellm_metadata) else {} @@ -2112,12 +2216,15 @@ class ResponsesWebSocketStreaming: self.authorized_model: str | None = authorized_model self.request_defaults: ResponsesWebSocketRequestDefaults | None = request_defaults - def _start_turn(self, client_frame: str | None) -> None: + def _start_turn( + self, client_frame: str | None, guardrail_information: tuple[Mapping[str, object], ...] = () + ) -> None: start_time: Final = datetime.now() # noqa: DTZ005 # naive to match the logging pipeline start_time turn_logging_obj: Final = self.logging_obj.copy_for_new_call( litellm_call_id=str(uuid.uuid4()), start_time=start_time ) self._detach_budget_reservation() + _set_turn_guardrail_information(turn_logging_obj, guardrail_information) turn_logging_obj.model_call_details["messages"] = list(_logged_user_messages(client_frame)) turn_logging_obj.pre_call(input=client_frame, api_key="") self._open_turns = (*self._open_turns, _ResponsesWsTurn(logging_obj=turn_logging_obj)) @@ -2136,10 +2243,46 @@ class ResponsesWebSocketStreaming: for key, value in litellm_params.items() } - def _record_client_frame(self, client_frame: str) -> None: + def _record_client_frame( + self, client_frame: str, guardrail_information: tuple[Mapping[str, object], ...] = () + ) -> None: frame: Final = _json_object_or_none(client_frame) if frame is not None and frame.get("type") == "response.create": - self._start_turn(client_frame) + self._start_turn(client_frame, guardrail_information) + + async def _guard_frame(self, message: str) -> GuardedResponsesWsFrame | None: + """Run the proxy's pre_call guardrails on one ``response.create`` frame. + + Returns None when a guardrail rejects the frame: the client gets an ``error`` event and the + frame is dropped while the connection stays open for the next turn. + """ + frame: Final = _json_object_or_none(message) + if self.frame_guard is None or frame is None or frame.get("type") != "response.create": + return GuardedResponsesWsFrame(frame=message) + try: + return await self.frame_guard(message) + except Exception as e: # noqa: BLE001 # any guardrail rejection maps to one client-facing error event + verbose_logger.info("Responses WebSocket: pre_call guardrail rejected a response.create frame: %s", e) + await self._send_error_event("invalid_request_error", str(e)) + return None + + async def _forward_guarded_frame(self, message: str) -> None: + guarded: Final = await self._guard_frame(message) + if guarded is None: + return + if guarded.pii_tokens: + self._remember_pii_tokens(guarded.pii_tokens) + masked: Final = await self._mask_response_create(guarded.frame) + self._record_client_frame(masked, guarded.guardrail_information) + await self.backend_ws.send(masked) + + def _remember_pii_tokens(self, pii_tokens: Mapping[str, str]) -> None: + metadata: Final = self.request_data.get("metadata") + known: Final = metadata.get("pii_tokens") if _is_json_object(metadata) else None + self.request_data["metadata"] = { + **(metadata if _is_json_object(metadata) else {}), + "pii_tokens": {**(known if _is_str_mapping(known) else {}), **pii_tokens}, + } def _record_backend_event(self, event: str) -> None: event_obj: Final = _json_object_or_none(event) @@ -2383,14 +2526,7 @@ class ResponsesWebSocketStreaming: return json.dumps(msg_obj) if modified else message # Delta event types whose ``delta`` field may contain PII tokens. - _DELTA_EVENT_TYPES = frozenset( - { - "response.output_text.delta", - "response.reasoning_summary_text.delta", - "response.refusal.delta", - "response.function_call_arguments.delta", - } - ) + _DELTA_EVENT_TYPES = _PII_DELTA_EVENT_TYPES # Terminal events that carry the full output text or tool-call arguments # already delivered by ``response.completed``. Suppressed when output masking @@ -2430,47 +2566,9 @@ class ResponsesWebSocketStreaming: if not pii_tokens: return response_str - try: - evt_obj: Final = _load_json_object(response_str) - except (json.JSONDecodeError, TypeError): - return response_str - cb: Final = self.guardrail_callbacks[0] unmask_pii_text: Final[_UnmasksPiiText] = getattr(cb, _UNMASK_PII_TEXT_ATTR) - event_type: Final = evt_obj.get("type") - - if event_type == "response.completed": - modified = False - response_obj: Final = evt_obj.get("response") - if not _is_json_object(response_obj): - return response_str - output_items: Final = response_obj.get("output") - for output_item in output_items if _is_json_array(output_items) else []: - if not _is_json_object(output_item): - continue - content = output_item.get("content") - if not _is_json_array(content): - continue - for content_block in content: - if not _is_json_object(content_block): - continue - text = content_block.get("text") - if isinstance(text, str): - unmasked = unmask_pii_text(text, pii_tokens) - if unmasked != text: - content_block["text"] = unmasked - modified = True - return json.dumps(evt_obj) if modified else response_str - - if event_type in self._DELTA_EVENT_TYPES: - delta: Final = evt_obj.get("delta") - if isinstance(delta, str): - unmasked = unmask_pii_text(delta, pii_tokens) - if unmasked != delta: - evt_obj["delta"] = unmasked - return json.dumps(evt_obj) - - return response_str + return _unmask_responses_event(response_str, pii_tokens, unmask_pii_text) async def _mask_response_completed(self, response_str: str) -> str: """ @@ -2564,38 +2662,29 @@ class ResponsesWebSocketStreaming: self.quota_callbacks, self.user_api_key_dict, self.authorized_model, message ) except RateLimitError as e: - try: - await self.websocket.send_text( - json.dumps( - { - "type": "error", - "error": { - "type": "rate_limit_exceeded", - "message": str(e), - }, - } - ) - ) - except Exception: # noqa: BLE001, S110 # client may already be gone - pass + await self._send_error_event("rate_limit_exceeded", str(e)) return False return True + async def _send_error_event(self, error_type: str, message: str) -> None: + try: + await self.websocket.send_text( + json.dumps({"type": "error", "error": {"type": error_type, "message": message}}) + ) + except Exception: # noqa: BLE001, S110 # client may already be gone + pass + async def client_to_backend(self) -> None: """Forward response.create events from client to backend.""" try: if self.first_message is not None and await self._enforce_or_reject_frame(self.first_message): - masked_first: Final = await self._mask_response_create(self.first_message) - self._record_client_frame(masked_first) - await self.backend_ws.send(masked_first) + await self._forward_guarded_frame(self.first_message) while True: message = await self.websocket.receive_text() if not await self._enforce_or_reject_frame(message): continue - masked = await self._mask_response_create(message) - self._record_client_frame(masked) - await self.backend_ws.send(masked) + await self._forward_guarded_frame(message) except Exception as e: verbose_logger.debug("Responses WS client_to_backend ended: %s", e) @@ -2672,10 +2761,15 @@ class ManagedResponsesWebSocketHandler: custom_llm_provider: str | None = None, first_message: str | None = None, quota_callbacks: Sequence[ProjectQuotaCallback] | None = None, + frame_guard: ResponsesWsFrameGuard | None = None, + guardrail_callbacks: Sequence[PresidioGuardrailCallback] | None = None, **kwargs: object, ) -> None: self.websocket = websocket self.model = model + self.frame_guard = frame_guard + self.guardrail_callbacks: tuple[PresidioGuardrailCallback, ...] = tuple(guardrail_callbacks or ()) + self._pii_tokens: Mapping[str, str] = {} self.logging_obj = logging_obj self.user_api_key_dict = user_api_key_dict self.litellm_metadata: Mapping[str, object] = litellm_metadata or {} @@ -2983,6 +3077,27 @@ class ManagedResponsesWebSocketHandler: if self.litellm_metadata: call_kwargs["litellm_metadata"] = dict(self.litellm_metadata) + async def _guard_frame(self, raw_message: str) -> GuardedResponsesWsFrame | None: + """Same per-frame pre_call guardrail walk as the native socket: a rejected frame errors and is not sent.""" + if self.frame_guard is None: + return GuardedResponsesWsFrame(frame=raw_message) + try: + return await self.frame_guard(raw_message) + except Exception as e: # noqa: BLE001 # any guardrail rejection maps to one client-facing error event + verbose_logger.info("ManagedResponsesWS: pre_call guardrail rejected a response.create frame: %s", e) + await self._send_error(str(e), error_type="invalid_request_error") + return None + + @staticmethod + def _metadata_with_turn_guardrails( + existing: object, guardrail_information: tuple[Mapping[str, object], ...] + ) -> Mapping[str, object]: + """Each turn is its own ``litellm.aresponses`` call, so its guardrail records ride on that call's metadata.""" + return { + **(existing if _is_json_object(existing) else {}), + _GUARDRAIL_INFORMATION_KEY: [dict(record) for record in guardrail_information], + } + @staticmethod def _update_proxy_request(call_kwargs: dict[str, Any], model: str) -> None: """Update proxy_server_request body so spend logs record the full request.""" @@ -3026,12 +3141,19 @@ class ManagedResponsesWebSocketHandler: except Exception: pass try: - await self.websocket.send_text(serialized) + await self.websocket.send_text(self._unmask_for_client(serialized)) except Exception as send_exc: verbose_logger.debug("ManagedResponsesWS: error sending chunk to client: %s", send_exc) return terminal_event # Client disconnected return terminal_event + def _unmask_for_client(self, serialized: str) -> str: + """The provider and the turn history keep Presidio's tokens, the client gets the original text back.""" + if not self._pii_tokens or not self.guardrail_callbacks: + return serialized + unmask_pii_text: Final[_UnmasksPiiText] = getattr(self.guardrail_callbacks[0], _UNMASK_PII_TEXT_ATTR) + return _unmask_responses_event(serialized, self._pii_tokens, unmask_pii_text) + def _save_turn_history( self, terminal_event: _MutableJsonObject | None, @@ -3101,7 +3223,13 @@ class ManagedResponsesWebSocketHandler: await self._send_error(str(e), error_type="rate_limit_exceeded") return - call_kwargs: Final = self._build_base_call_kwargs(msg_obj) + guarded: Final = await self._guard_frame(raw_message) + if guarded is None: + return + if guarded.pii_tokens: + self._pii_tokens = {**self._pii_tokens, **guarded.pii_tokens} + + call_kwargs: Final = self._build_base_call_kwargs(_load_json_object(guarded.frame)) call_kwargs["stream"] = True # A frame that repeats the connection's public alias (model_group) must @@ -3121,6 +3249,10 @@ class ManagedResponsesWebSocketHandler: self._apply_history(call_kwargs, previous_response_id, current_messages, prior_history) self._inject_credentials(call_kwargs, model=model) + if guarded.guardrail_information: + call_kwargs["litellm_metadata"] = self._metadata_with_turn_guardrails( + call_kwargs.get("litellm_metadata"), guarded.guardrail_information + ) self._update_proxy_request(call_kwargs, requested_model or self.model_group or model) call_kwargs.update(self.extra_kwargs) diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 58338f185c7..8545a5b568e 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -65561,6 +65561,40 @@ "supports_vision": false, "source": "https://deepinfra.com/pricing" }, + "deepinfra/deepseek-ai/DeepSeek-V4.1-Flash": { + "max_tokens": 131072, + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "input_cost_per_token": 1.4e-07, + "output_cost_per_token": 4.2e-07, + "cache_read_input_token_cost": 4.2e-09, + "supports_prompt_caching": true, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true, + "supports_reasoning": true, + "supports_vision": true, + "source": "https://deepinfra.com/deepseek-ai/DeepSeek-V4.1-Flash" + }, + "deepinfra/deepseek-ai/DeepSeek-V4.1-Flash-1017": { + "max_tokens": 131072, + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "input_cost_per_token": 1.4e-07, + "output_cost_per_token": 4.2e-07, + "cache_read_input_token_cost": 4.2e-09, + "supports_prompt_caching": true, + "litellm_provider": "deepinfra", + "mode": "chat", + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_response_schema": true, + "supports_reasoning": true, + "supports_vision": true, + "source": "https://deepinfra.com/deepseek-ai/DeepSeek-V4.1-Flash-1017" + }, "deepinfra/openai/gpt-oss-120b-Ultra": { "max_tokens": 131072, "max_input_tokens": 131072, diff --git a/tests/integration/cost_calculation/cases/anthropic.py b/tests/integration/cost_calculation/cases/anthropic.py index a0c81d8f06d..5ad4afa9530 100644 --- a/tests/integration/cost_calculation/cases/anthropic.py +++ b/tests/integration/cost_calculation/cases/anthropic.py @@ -6899,6 +6899,35 @@ CLAUDE_OPUS_5_5_RESPONSES_BRIDGED_JSON: Final = CLAUDE_OPUS_5_5_MESSAGES_INPUT_T } ) +def _responses_bridged_stream_case( + sibling: CostTrackingTestCase, stream_sibling: CostTrackingTestCase, name: str +) -> CostTrackingTestCase: + return sibling.model_copy( + update={ + "name": name, + "request": {**sibling.request, "stream": True}, + "response": stream_sibling.response, + } + ) + + +CLAUDE_FABLE_5_1_RESPONSES_BRIDGED_STREAM: Final = _responses_bridged_stream_case( + CLAUDE_FABLE_5_1_RESPONSES_BRIDGED_JSON, CLAUDE_FABLE_5_1_MESSAGES_STREAM, + "claude-fable-5-1-responses_bridged_stream", +) +CLAUDE_HAIKU_5_5_RESPONSES_BRIDGED_STREAM: Final = _responses_bridged_stream_case( + CLAUDE_HAIKU_5_5_RESPONSES_BRIDGED_JSON, CLAUDE_HAIKU_5_5_MESSAGES_STREAM, + "claude-haiku-5-5-responses_bridged_stream", +) +CLAUDE_OPUS_5_5_RESPONSES_BRIDGED_STREAM: Final = _responses_bridged_stream_case( + CLAUDE_OPUS_5_5_RESPONSES_BRIDGED_JSON, CLAUDE_OPUS_5_5_MESSAGES_STREAM, + "claude-opus-5-5-responses_bridged_stream", +) +CLAUDE_SONNET_5_5_RESPONSES_BRIDGED_STREAM: Final = _responses_bridged_stream_case( + CLAUDE_SONNET_5_5_RESPONSES_BRIDGED_JSON, CLAUDE_SONNET_5_5_MESSAGES_STREAM, + "claude-sonnet-5-5-responses_bridged_stream", +) + CLAUDE_HAIKU_5_5_TIERED_INPUT_ABOVE_100K: Final = with_usage( CLAUDE_HAIKU_5_5_INPUT_TEXT, "claude-haiku-5-5-tiered_input_above_100k", @@ -6994,6 +7023,10 @@ CASES: Final[tuple[CostTrackingTestCase, ...]] = ( CLAUDE_SONNET_5_RESPONSES_BRIDGED_JSON, CLAUDE_SONNET_5_5_RESPONSES_BRIDGED_JSON, CLAUDE_HAIKU_5_5_RESPONSES_BRIDGED_JSON, + CLAUDE_FABLE_5_1_RESPONSES_BRIDGED_STREAM, + CLAUDE_HAIKU_5_5_RESPONSES_BRIDGED_STREAM, + CLAUDE_OPUS_5_5_RESPONSES_BRIDGED_STREAM, + CLAUDE_SONNET_5_5_RESPONSES_BRIDGED_STREAM, CLAUDE_FABLE_5_RESPONSES_BRIDGED_JSON, CLAUDE_FABLE_5_1_RESPONSES_BRIDGED_JSON, CLAUDE_OPUS_5_5_RESPONSES_BRIDGED_JSON, @@ -7119,6 +7152,18 @@ CASES: Final[tuple[CostTrackingTestCase, ...]] = ( ) PARITY: Final[tuple[StreamParityTestCase, ...]] = ( + StreamParityTestCase( + plain=CLAUDE_FABLE_5_1_RESPONSES_BRIDGED_JSON, streamed=CLAUDE_FABLE_5_1_RESPONSES_BRIDGED_STREAM + ), + StreamParityTestCase( + plain=CLAUDE_HAIKU_5_5_RESPONSES_BRIDGED_JSON, streamed=CLAUDE_HAIKU_5_5_RESPONSES_BRIDGED_STREAM + ), + StreamParityTestCase( + plain=CLAUDE_OPUS_5_5_RESPONSES_BRIDGED_JSON, streamed=CLAUDE_OPUS_5_5_RESPONSES_BRIDGED_STREAM + ), + StreamParityTestCase( + plain=CLAUDE_SONNET_5_5_RESPONSES_BRIDGED_JSON, streamed=CLAUDE_SONNET_5_5_RESPONSES_BRIDGED_STREAM + ), StreamParityTestCase(plain=CLAUDE_HAIKU_4_5_INPUT_TEXT, streamed=CLAUDE_HAIKU_4_5_STREAM), StreamParityTestCase( plain=CLAUDE_HAIKU_4_5_RESPONSE_MODEL_OVERRIDE, streamed=CLAUDE_HAIKU_4_5_STREAM_RESPONSE_MODEL_OVERRIDE diff --git a/tests/integration/cost_calculation/cases/azure.py b/tests/integration/cost_calculation/cases/azure.py index 1360f65b802..d9f17c12bda 100644 --- a/tests/integration/cost_calculation/cases/azure.py +++ b/tests/integration/cost_calculation/cases/azure.py @@ -21,6 +21,8 @@ from integration.cost_calculation.cost_tracking_case import ( ) from integration.cost_calculation.bridge_cost_parity.case import BridgeCostParityTestCase, bridge_pairs from integration.cost_calculation.stream_parity.case import StreamParityTestCase, sse_frames +from integration.cost_calculation.cases.openai import _responses_json_body, _responses_stream_case +from pydantic import JsonValue def _response_for_model( @@ -47,6 +49,37 @@ def _stream_for_model(response: SseResponse, model: str) -> SseResponse: ) +def _azure_responses_json_case( + base: CostTrackingTestCase, name: str, response_model: str, content_filters: JsonValue +) -> CostTrackingTestCase: + response_case: Final = bridge_pairs((base,), ("/v1/responses",))[0].bridged + return response_case.model_copy( + update={ + "name": name, + "response": JsonResponse( + content_type="application/json", + body=_responses_json_body( + response_model, + include_content_filters=True, + content_filters=content_filters, + ), + ), + } + ) + + +def _azure_messages_stream_case( + messages_json: CostTrackingTestCase, stream_case: CostTrackingTestCase, name: str +) -> CostTrackingTestCase: + return messages_json.model_copy( + update={ + "name": name, + "request": {**messages_json.request, "stream": True}, + "response": stream_case.response, + } + ) + + # azure/gpt-5.4-mini AZURE_GPT_5_4_MINI_INPUT_TEXT: Final = CostTrackingTestCase( name="azure-gpt-5.4-mini-input_text", @@ -5812,6 +5845,85 @@ AZURE_GPT_6_ASTRA_CACHE_WRITE_5M: Final = with_response_usage( ExactExpected(spend=0.22375, input_cost=0.2, output_cost=0.02375, prompt_tokens=12928, completion_tokens=380), ) +AZURE_GPT_5_6_TERRA_TIERED_INPUT_ABOVE_272K: Final = with_usage( + AZURE_GPT_5_6_TERRA_INPUT_TEXT, + "azure/gpt-5.6-terra-tiered_input_above_272k", + { + "prompt_tokens": 300000, + "total_tokens": 300412, + "prompt_tokens_details": {"audio_tokens": 0, "cache_write_tokens": 0, "cached_tokens": 0}, + }, + ExactExpected( + spend=1.3281576, input_cost=1.32, output_cost=0.0081576, prompt_tokens=300000, completion_tokens=412 + ), +) +AZURE_GPT_5_6_TERRA_TIERED_CACHE_READ_ABOVE_272K: Final = with_usage( + AZURE_GPT_5_6_TERRA_INPUT_TEXT, + "azure/gpt-5.6-terra-tiered_cache_read_above_272k", + { + "prompt_tokens": 300000, + "total_tokens": 300412, + "prompt_tokens_details": {"audio_tokens": 0, "cache_write_tokens": 0, "cached_tokens": 200000}, + }, + ExactExpected( + spend=0.5361576, input_cost=0.528, output_cost=0.0081576, prompt_tokens=300000, completion_tokens=412 + ), +) + + +AZURE_GPT_6_1_SOL_RESPONSES_JSON: Final = _azure_responses_json_case( + AZURE_GPT_6_1_SOL_INPUT_TEXT, "azure-gpt-6.1-sol-responses_json", "gpt-6.1-sol", None +) +AZURE_GPT_6_LUNA_RESPONSES_JSON: Final = _azure_responses_json_case( + AZURE_GPT_6_LUNA_INPUT_TEXT, "azure-gpt-6-luna-responses_json", "gpt-6-luna", None +) +AZURE_GPT_6_ASTRA_RESPONSES_JSON: Final = _azure_responses_json_case( + AZURE_GPT_6_ASTRA_INPUT_TEXT, "azure-gpt-6-astra-responses_json", "gpt-6-astra", [] +) +AZURE_GPT_5_6_TERRA_RESPONSES_JSON: Final = _azure_responses_json_case( + AZURE_GPT_5_6_TERRA_INPUT_TEXT, "azure-gpt-5.6-terra-responses_json", "gpt-5.6-terra", None +) +AZURE_GPT_6_1_SOL_RESPONSES_STREAM: Final = _responses_stream_case( + AZURE_GPT_6_1_SOL_RESPONSES_JSON, + "azure-gpt-6.1-sol-responses_stream", + "gpt-6.1-sol", + 14, + include_content_filters=True, +) +AZURE_GPT_6_LUNA_RESPONSES_STREAM: Final = _responses_stream_case( + AZURE_GPT_6_LUNA_RESPONSES_JSON, + "azure-gpt-6-luna-responses_stream", + "gpt-6-luna", + 25, + include_content_filters=True, +) +AZURE_GPT_6_ASTRA_RESPONSES_STREAM: Final = _responses_stream_case( + AZURE_GPT_6_ASTRA_RESPONSES_JSON, + "azure-gpt-6-astra-responses_stream", + "gpt-6-astra", + 16, + include_content_filters=True, + content_filters=[], +) +AZURE_GPT_5_6_TERRA_RESPONSES_STREAM: Final = _responses_stream_case( + AZURE_GPT_5_6_TERRA_RESPONSES_JSON, + "azure-gpt-5.6-terra-responses_stream", + "gpt-5.6-terra", + 15, + include_content_filters=True, +) +AZURE_GPT_6_1_SOL_MESSAGES_BRIDGED_STREAM: Final = _azure_messages_stream_case( + AZURE_GPT_6_1_SOL_MESSAGES_BRIDGED_JSON, AZURE_GPT_6_1_SOL_STREAM, "azure-gpt-6.1-sol-messages_bridged_stream" +) +AZURE_GPT_6_LUNA_MESSAGES_BRIDGED_STREAM: Final = _azure_messages_stream_case( + AZURE_GPT_6_LUNA_MESSAGES_BRIDGED_JSON, AZURE_GPT_6_LUNA_STREAM, "azure-gpt-6-luna-messages_bridged_stream" +) +AZURE_GPT_6_ASTRA_MESSAGES_BRIDGED_STREAM: Final = _azure_messages_stream_case( + AZURE_GPT_6_ASTRA_MESSAGES_BRIDGED_JSON, AZURE_GPT_6_ASTRA_STREAM, "azure-gpt-6-astra-messages_bridged_stream" +) +AZURE_GPT_5_6_TERRA_MESSAGES_BRIDGED_STREAM: Final = _azure_messages_stream_case( + AZURE_GPT_5_6_TERRA_MESSAGES_BRIDGED_JSON, AZURE_GPT_5_6_TERRA_STREAM, "azure-gpt-5.6-terra-messages_bridged_stream" +) CASES: Final[tuple[CostTrackingTestCase, ...]] = ( AZURE_GPT_6_SOL_MESSAGES_BRIDGED_JSON, @@ -5903,6 +6015,8 @@ CASES: Final[tuple[CostTrackingTestCase, ...]] = ( AZURE_GPT_6_LUNA_CACHE_WRITE_5M, AZURE_GPT_6_ASTRA_CACHE_WRITE_5M, AZURE_GPT_5_6_TERRA_INPUT_TEXT, + AZURE_GPT_5_6_TERRA_TIERED_INPUT_ABOVE_272K, + AZURE_GPT_5_6_TERRA_TIERED_CACHE_READ_ABOVE_272K, AZURE_GPT_5_6_TERRA_STREAM, AZURE_GPT_5_6_TERRA_SERVICE_TIER_FLEX, AZURE_GPT_5_6_TERRA_SERVICE_TIER_PRIORITY, @@ -5911,6 +6025,18 @@ CASES: Final[tuple[CostTrackingTestCase, ...]] = ( AZURE_GPT_5_6_TERRA_CACHE_WRITE_5M, AZURE_GPT_5_6_TERRA_REASONING, AZURE_GPT_5_6_TERRA_TOOL_CALL, + AZURE_GPT_6_1_SOL_RESPONSES_JSON, + AZURE_GPT_6_1_SOL_RESPONSES_STREAM, + AZURE_GPT_6_LUNA_RESPONSES_JSON, + AZURE_GPT_6_LUNA_RESPONSES_STREAM, + AZURE_GPT_6_ASTRA_RESPONSES_JSON, + AZURE_GPT_6_ASTRA_RESPONSES_STREAM, + AZURE_GPT_5_6_TERRA_RESPONSES_JSON, + AZURE_GPT_5_6_TERRA_RESPONSES_STREAM, + AZURE_GPT_6_1_SOL_MESSAGES_BRIDGED_STREAM, + AZURE_GPT_6_LUNA_MESSAGES_BRIDGED_STREAM, + AZURE_GPT_6_ASTRA_MESSAGES_BRIDGED_STREAM, + AZURE_GPT_5_6_TERRA_MESSAGES_BRIDGED_STREAM, ) PARITY: Final[tuple[StreamParityTestCase, ...]] = ( @@ -5924,6 +6050,22 @@ PARITY: Final[tuple[StreamParityTestCase, ...]] = ( StreamParityTestCase(plain=AZURE_GPT_6_ASTRA_INPUT_TEXT, streamed=AZURE_GPT_6_ASTRA_STREAM), StreamParityTestCase(plain=AZURE_GPT_6_1_SOL_INPUT_TEXT, streamed=AZURE_GPT_6_1_SOL_STREAM), StreamParityTestCase(plain=AZURE_GPT_5_6_TERRA_INPUT_TEXT, streamed=AZURE_GPT_5_6_TERRA_STREAM), + StreamParityTestCase(plain=AZURE_GPT_6_1_SOL_RESPONSES_JSON, streamed=AZURE_GPT_6_1_SOL_RESPONSES_STREAM), + StreamParityTestCase(plain=AZURE_GPT_6_LUNA_RESPONSES_JSON, streamed=AZURE_GPT_6_LUNA_RESPONSES_STREAM), + StreamParityTestCase(plain=AZURE_GPT_6_ASTRA_RESPONSES_JSON, streamed=AZURE_GPT_6_ASTRA_RESPONSES_STREAM), + StreamParityTestCase(plain=AZURE_GPT_5_6_TERRA_RESPONSES_JSON, streamed=AZURE_GPT_5_6_TERRA_RESPONSES_STREAM), + StreamParityTestCase( + plain=AZURE_GPT_6_1_SOL_MESSAGES_BRIDGED_JSON, streamed=AZURE_GPT_6_1_SOL_MESSAGES_BRIDGED_STREAM + ), + StreamParityTestCase( + plain=AZURE_GPT_6_LUNA_MESSAGES_BRIDGED_JSON, streamed=AZURE_GPT_6_LUNA_MESSAGES_BRIDGED_STREAM + ), + StreamParityTestCase( + plain=AZURE_GPT_6_ASTRA_MESSAGES_BRIDGED_JSON, streamed=AZURE_GPT_6_ASTRA_MESSAGES_BRIDGED_STREAM + ), + StreamParityTestCase( + plain=AZURE_GPT_5_6_TERRA_MESSAGES_BRIDGED_JSON, streamed=AZURE_GPT_5_6_TERRA_MESSAGES_BRIDGED_STREAM + ), ) BRIDGE_COST_PARITY: Final[tuple[BridgeCostParityTestCase, ...]] = ( @@ -5938,4 +6080,10 @@ BRIDGE_COST_PARITY: Final[tuple[BridgeCostParityTestCase, ...]] = ( BridgeCostParityTestCase( native=AZURE_GPT_5_6_TERRA_INPUT_TEXT, bridged=AZURE_GPT_5_6_TERRA_MESSAGES_BRIDGED_JSON ), + BridgeCostParityTestCase(native=AZURE_GPT_6_1_SOL_STREAM, bridged=AZURE_GPT_6_1_SOL_MESSAGES_BRIDGED_STREAM), + BridgeCostParityTestCase(native=AZURE_GPT_6_LUNA_STREAM, bridged=AZURE_GPT_6_LUNA_MESSAGES_BRIDGED_STREAM), + BridgeCostParityTestCase(native=AZURE_GPT_6_ASTRA_STREAM, bridged=AZURE_GPT_6_ASTRA_MESSAGES_BRIDGED_STREAM), + BridgeCostParityTestCase( + native=AZURE_GPT_5_6_TERRA_STREAM, bridged=AZURE_GPT_5_6_TERRA_MESSAGES_BRIDGED_STREAM + ), ) diff --git a/tests/integration/cost_calculation/cases/bedrock_converse.py b/tests/integration/cost_calculation/cases/bedrock_converse.py index f78bdfecfa3..858db06cafc 100644 --- a/tests/integration/cost_calculation/cases/bedrock_converse.py +++ b/tests/integration/cost_calculation/cases/bedrock_converse.py @@ -6088,6 +6088,71 @@ US_ANTHROPIC_CLAUDE_HAIKU_5_5_TIERED_CACHE_WRITE_1H_ABOVE_100K: Final = with_usa ), ) +def _us_anthropic_messages_bridge_pairs( + input_case: CostTrackingTestCase, stream_case: CostTrackingTestCase +) -> tuple[BridgeCostParityTestCase, ...]: + return tuple( + BridgeCostParityTestCase( + native=pair.native, + bridged=pair.bridged.model_copy( + update={ + "name": ( + f"us.anthropic." + f"{pair.native.name.removeprefix('us.anthropic.').removesuffix('-input_text').removesuffix('-stream')}" + f"-messages_bridged_{'stream' if pair.native.request.get('stream') is True else 'json'}" + ) + } + ), + ) + for pair in bridge_pairs((input_case, stream_case), ("/v1/messages",)) + ) + + +def _us_anthropic_responses_stream_case( + json_case: CostTrackingTestCase, stream_case: CostTrackingTestCase, name: str +) -> CostTrackingTestCase: + return json_case.model_copy( + update={ + "name": name, + "request": {**json_case.request, "stream": True}, + "response": stream_case.response, + } + ) + + +US_ANTHROPIC_CLAUDE_FABLE_5_1_RESPONSES_BRIDGED_STREAM: Final = _us_anthropic_responses_stream_case( + US_ANTHROPIC_CLAUDE_FABLE_5_1_RESPONSES_BRIDGED_JSON, + US_ANTHROPIC_CLAUDE_FABLE_5_1_STREAM, + "us.anthropic.claude-fable-5-1-responses_bridged_stream", +) +US_ANTHROPIC_CLAUDE_HAIKU_5_5_RESPONSES_BRIDGED_STREAM: Final = _us_anthropic_responses_stream_case( + US_ANTHROPIC_CLAUDE_HAIKU_5_5_RESPONSES_BRIDGED_JSON, + US_ANTHROPIC_CLAUDE_HAIKU_5_5_STREAM, + "us.anthropic.claude-haiku-5-5-responses_bridged_stream", +) +US_ANTHROPIC_CLAUDE_OPUS_5_5_RESPONSES_BRIDGED_STREAM: Final = _us_anthropic_responses_stream_case( + US_ANTHROPIC_CLAUDE_OPUS_5_5_RESPONSES_BRIDGED_JSON, + US_ANTHROPIC_CLAUDE_OPUS_5_5_STREAM, + "us.anthropic.claude-opus-5-5-responses_bridged_stream", +) +US_ANTHROPIC_CLAUDE_SONNET_5_5_RESPONSES_BRIDGED_STREAM: Final = _us_anthropic_responses_stream_case( + US_ANTHROPIC_CLAUDE_SONNET_5_5_RESPONSES_BRIDGED_JSON, + US_ANTHROPIC_CLAUDE_SONNET_5_5_STREAM, + "us.anthropic.claude-sonnet-5-5-responses_bridged_stream", +) +_US_ANTHROPIC_FABLE_5_1_MESSAGES_PAIRS: Final = _us_anthropic_messages_bridge_pairs( + US_ANTHROPIC_CLAUDE_FABLE_5_1_INPUT_TEXT, US_ANTHROPIC_CLAUDE_FABLE_5_1_STREAM +) +_US_ANTHROPIC_HAIKU_5_5_MESSAGES_PAIRS: Final = _us_anthropic_messages_bridge_pairs( + US_ANTHROPIC_CLAUDE_HAIKU_5_5_INPUT_TEXT, US_ANTHROPIC_CLAUDE_HAIKU_5_5_STREAM +) +_US_ANTHROPIC_OPUS_5_5_MESSAGES_PAIRS: Final = _us_anthropic_messages_bridge_pairs( + US_ANTHROPIC_CLAUDE_OPUS_5_5_INPUT_TEXT, US_ANTHROPIC_CLAUDE_OPUS_5_5_STREAM +) +_US_ANTHROPIC_SONNET_5_5_MESSAGES_PAIRS: Final = _us_anthropic_messages_bridge_pairs( + US_ANTHROPIC_CLAUDE_SONNET_5_5_INPUT_TEXT, US_ANTHROPIC_CLAUDE_SONNET_5_5_STREAM +) + CASES: Final[tuple[CostTrackingTestCase, ...]] = ( US_ANTHROPIC_CLAUDE_SONNET_5_RESPONSES_BRIDGED_JSON, US_ANTHROPIC_CLAUDE_SONNET_5_5_RESPONSES_BRIDGED_JSON, @@ -6210,6 +6275,14 @@ CASES: Final[tuple[CostTrackingTestCase, ...]] = ( US_ANTHROPIC_CLAUDE_FABLE_5_1_THINKING, US_ANTHROPIC_CLAUDE_OPUS_5_5_UPSTREAM_400_ZERO_SPEND, US_ANTHROPIC_CLAUDE_OPUS_5_5_INPUT_TEXT_ROLLUPS, + US_ANTHROPIC_CLAUDE_FABLE_5_1_RESPONSES_BRIDGED_STREAM, + US_ANTHROPIC_CLAUDE_HAIKU_5_5_RESPONSES_BRIDGED_STREAM, + US_ANTHROPIC_CLAUDE_OPUS_5_5_RESPONSES_BRIDGED_STREAM, + US_ANTHROPIC_CLAUDE_SONNET_5_5_RESPONSES_BRIDGED_STREAM, + *tuple(pair.bridged for pair in _US_ANTHROPIC_FABLE_5_1_MESSAGES_PAIRS), + *tuple(pair.bridged for pair in _US_ANTHROPIC_HAIKU_5_5_MESSAGES_PAIRS), + *tuple(pair.bridged for pair in _US_ANTHROPIC_OPUS_5_5_MESSAGES_PAIRS), + *tuple(pair.bridged for pair in _US_ANTHROPIC_SONNET_5_5_MESSAGES_PAIRS), ) PARITY: Final[tuple[StreamParityTestCase, ...]] = ( @@ -6235,6 +6308,38 @@ PARITY: Final[tuple[StreamParityTestCase, ...]] = ( StreamParityTestCase(plain=US_ANTHROPIC_CLAUDE_HAIKU_5_5_INPUT_TEXT, streamed=US_ANTHROPIC_CLAUDE_HAIKU_5_5_STREAM), StreamParityTestCase(plain=US_ANTHROPIC_CLAUDE_FABLE_5_INPUT_TEXT, streamed=US_ANTHROPIC_CLAUDE_FABLE_5_STREAM), StreamParityTestCase(plain=US_ANTHROPIC_CLAUDE_FABLE_5_1_INPUT_TEXT, streamed=US_ANTHROPIC_CLAUDE_FABLE_5_1_STREAM), + StreamParityTestCase( + plain=US_ANTHROPIC_CLAUDE_FABLE_5_1_RESPONSES_BRIDGED_JSON, + streamed=US_ANTHROPIC_CLAUDE_FABLE_5_1_RESPONSES_BRIDGED_STREAM, + ), + StreamParityTestCase( + plain=US_ANTHROPIC_CLAUDE_HAIKU_5_5_RESPONSES_BRIDGED_JSON, + streamed=US_ANTHROPIC_CLAUDE_HAIKU_5_5_RESPONSES_BRIDGED_STREAM, + ), + StreamParityTestCase( + plain=US_ANTHROPIC_CLAUDE_OPUS_5_5_RESPONSES_BRIDGED_JSON, + streamed=US_ANTHROPIC_CLAUDE_OPUS_5_5_RESPONSES_BRIDGED_STREAM, + ), + StreamParityTestCase( + plain=US_ANTHROPIC_CLAUDE_SONNET_5_5_RESPONSES_BRIDGED_JSON, + streamed=US_ANTHROPIC_CLAUDE_SONNET_5_5_RESPONSES_BRIDGED_STREAM, + ), + StreamParityTestCase( + plain=_US_ANTHROPIC_FABLE_5_1_MESSAGES_PAIRS[0].bridged, + streamed=_US_ANTHROPIC_FABLE_5_1_MESSAGES_PAIRS[1].bridged, + ), + StreamParityTestCase( + plain=_US_ANTHROPIC_HAIKU_5_5_MESSAGES_PAIRS[0].bridged, + streamed=_US_ANTHROPIC_HAIKU_5_5_MESSAGES_PAIRS[1].bridged, + ), + StreamParityTestCase( + plain=_US_ANTHROPIC_OPUS_5_5_MESSAGES_PAIRS[0].bridged, + streamed=_US_ANTHROPIC_OPUS_5_5_MESSAGES_PAIRS[1].bridged, + ), + StreamParityTestCase( + plain=_US_ANTHROPIC_SONNET_5_5_MESSAGES_PAIRS[0].bridged, + streamed=_US_ANTHROPIC_SONNET_5_5_MESSAGES_PAIRS[1].bridged, + ), ) BRIDGE_COST_PARITY: Final[tuple[BridgeCostParityTestCase, ...]] = ( @@ -6267,4 +6372,24 @@ BRIDGE_COST_PARITY: Final[tuple[BridgeCostParityTestCase, ...]] = ( BridgeCostParityTestCase( native=US_ANTHROPIC_CLAUDE_FABLE_5_1_INPUT_TEXT, bridged=US_ANTHROPIC_CLAUDE_FABLE_5_1_RESPONSES_BRIDGED_JSON ), + BridgeCostParityTestCase( + native=US_ANTHROPIC_CLAUDE_FABLE_5_1_STREAM, + bridged=US_ANTHROPIC_CLAUDE_FABLE_5_1_RESPONSES_BRIDGED_STREAM, + ), + BridgeCostParityTestCase( + native=US_ANTHROPIC_CLAUDE_HAIKU_5_5_STREAM, + bridged=US_ANTHROPIC_CLAUDE_HAIKU_5_5_RESPONSES_BRIDGED_STREAM, + ), + BridgeCostParityTestCase( + native=US_ANTHROPIC_CLAUDE_OPUS_5_5_STREAM, + bridged=US_ANTHROPIC_CLAUDE_OPUS_5_5_RESPONSES_BRIDGED_STREAM, + ), + BridgeCostParityTestCase( + native=US_ANTHROPIC_CLAUDE_SONNET_5_5_STREAM, + bridged=US_ANTHROPIC_CLAUDE_SONNET_5_5_RESPONSES_BRIDGED_STREAM, + ), + *_US_ANTHROPIC_FABLE_5_1_MESSAGES_PAIRS, + *_US_ANTHROPIC_HAIKU_5_5_MESSAGES_PAIRS, + *_US_ANTHROPIC_OPUS_5_5_MESSAGES_PAIRS, + *_US_ANTHROPIC_SONNET_5_5_MESSAGES_PAIRS, ) diff --git a/tests/integration/cost_calculation/cases/bedrock_mantle.py b/tests/integration/cost_calculation/cases/bedrock_mantle.py index 2fa940d66f4..3cd9150033a 100644 --- a/tests/integration/cost_calculation/cases/bedrock_mantle.py +++ b/tests/integration/cost_calculation/cases/bedrock_mantle.py @@ -18,7 +18,7 @@ from integration.cost_calculation.cases.anthropic import ( CLAUDE_OPUS_5_5_THINKING, CLAUDE_OPUS_5_5_TOOL_CALL, ) -from integration.cost_calculation.bridge_cost_parity.case import BridgeCostParityTestCase +from integration.cost_calculation.bridge_cost_parity.case import BridgeCostParityTestCase, bridge_pairs from integration.cost_calculation.cost_tracking_case import ( CostTrackingTestCase, Deployment, @@ -535,14 +535,40 @@ _MANTLE_GPT_CHAT_REQUEST: Final = { "messages": [{"role": "user", "content": "cost tracking fixture"}], "max_completion_tokens": 64, } +_MANTLE_GPT_TOOL_REQUEST: Final = { + **_MANTLE_GPT_RESPONSES_REQUEST, + "tools": [ + { + "type": "function", + "name": "get_weather", + "description": "Get the weather for a city", + "parameters": { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + }, + } + ], + "tool_choice": "required", +} +_MANTLE_GPT_CACHE_WRITE_REQUEST: Final = { + **_MANTLE_GPT_RESPONSES_REQUEST, + "input": "cacheable pricing-harness prefix " * 1000, + "prompt_cache_key": "cost-cases-20261017", + "max_output_tokens": 64, +} def _mantle_gpt_usage( - input_tokens: int, output_tokens: int, reasoning_tokens: int, cached_tokens: int = 0 + input_tokens: int, + output_tokens: int, + reasoning_tokens: int, + cached_tokens: int = 0, + cache_write_tokens: int = 0, ) -> dict[str, JsonValue]: return { "input_tokens": input_tokens, - "input_tokens_details": {"cache_write_tokens": 0, "cached_tokens": cached_tokens}, + "input_tokens_details": {"cache_write_tokens": cache_write_tokens, "cached_tokens": cached_tokens}, "output_tokens": output_tokens, "output_tokens_details": {"reasoning_tokens": reasoning_tokens}, "total_tokens": input_tokens + output_tokens, @@ -628,6 +654,20 @@ def _mantle_gpt_cases( response_expected, _MANTLE_GPT_CHAT_REQUEST, ) + chat_stream_case: Final = chat_case.model_copy( + update={ + "name": f"bedrock_mantle-{provider_model}-chat_stream", + "request": {**chat_case.request, "stream": True}, + "response": stream_case.response, + } + ) + messages_pairs: Final = bridge_pairs((response_case, stream_case), ("/v1/messages",)) + messages_json: Final = messages_pairs[0].bridged.model_copy( + update={"name": f"bedrock_mantle-{provider_model}-messages_bridged_json"} + ) + messages_stream: Final = messages_pairs[1].bridged.model_copy( + update={"name": f"bedrock_mantle-{provider_model}-messages_bridged_stream"} + ) cache_case: Final = _mantle_gpt_json_case( BEDROCK_MANTLE_OPENAI_GPT_6_SOL_RESPONSES_JSON, provider_model, @@ -660,7 +700,73 @@ def _mantle_gpt_cases( tiered_boundary_expected, {**_MANTLE_GPT_RESPONSES_REQUEST, "max_output_tokens": 412}, ) - return (response_case, stream_case, reasoning_case, chat_case, cache_case, tiered_input, tiered_cache, tiered_boundary) + return ( + response_case, + stream_case, + reasoning_case, + chat_case, + cache_case, + tiered_input, + tiered_cache, + tiered_boundary, + chat_stream_case, + messages_json, + messages_stream, + ) + + +def _mantle_gpt_audit3_cases( + provider_model: str, + tool_input_tokens: int, + tool_output_tokens: int, + tool_reasoning_tokens: int, + tool_expected: ExactExpected, + cache_output_tokens: int, + cache_reasoning_tokens: int, + cache_expected: ExactExpected, +) -> tuple[CostTrackingTestCase, ...]: + tool_case: Final = _mantle_gpt_json_case( + BEDROCK_MANTLE_OPENAI_GPT_6_SOL_RESPONSES_JSON, + provider_model, + "tool_call", + _mantle_gpt_usage(tool_input_tokens, tool_output_tokens, tool_reasoning_tokens), + tool_expected, + _MANTLE_GPT_TOOL_REQUEST, + ) + if not isinstance(tool_case.response, JsonResponse): + raise TypeError(f"{tool_case.name}: expected a JSON response fixture") + tool_response: Final = tool_case.response.model_copy( + update={ + "body": { + **tool_case.response.body, + "output": [ + { + "arguments": '{"city":"Paris"}', + "call_id": "call_$REQUEST_ID", + "id": "fc_$REQUEST_ID", + "name": "get_weather", + "status": "completed", + "type": "function_call", + } + ], + } + } + ) + tool_case_with_response: Final = tool_case.model_copy(update={"response": tool_response}) + cache_case: Final = _mantle_gpt_json_case( + BEDROCK_MANTLE_OPENAI_GPT_6_SOL_RESPONSES_JSON, + provider_model, + "cache_write_5m", + _mantle_gpt_usage( + 1667, + cache_output_tokens, + cache_reasoning_tokens, + cache_write_tokens=1665, + ), + cache_expected, + _MANTLE_GPT_CACHE_WRITE_REQUEST, + ) + return (tool_case_with_response, cache_case) def _mantle_gpt_bridge(case: CostTrackingTestCase) -> BridgeCostParityTestCase: @@ -723,6 +829,83 @@ _MANTLE_GPT_56_TERRA_CASES: Final = _mantle_gpt_cases( ExactExpected(spend=0.84537376, input_cost=0.83776, output_cost=0.00761376, prompt_tokens=272000, completion_tokens=412), ) +_MANTLE_GPT_AUDIT3_CASES: Final = ( + *_mantle_gpt_audit3_cases( + "openai.gpt-6.1-sol", + 50, + 20, + 0, + ExactExpected(spend=0.000363, input_cost=0.000121, output_cost=0.000242, prompt_tokens=50, completion_tokens=20), + 28, + 0, + ExactExpected( + spend=0.005381765, + input_cost=0.005042965, + output_cost=0.0003388, + prompt_tokens=1667, + completion_tokens=28, + cache_creation_cost=0.005038125, + ), + ), + *_mantle_gpt_audit3_cases( + "openai.gpt-6-luna", + 51, + 18, + 0, + ExactExpected( + spend=0.000018612, input_cost=0.000006732, output_cost=0.00001188, prompt_tokens=51, completion_tokens=18 + ), + 64, + 40, + ExactExpected( + spend=0.000317229, + input_cost=0.000274989, + output_cost=0.00004224, + prompt_tokens=1667, + completion_tokens=64, + cache_creation_cost=0.000274725, + ), + ), + *_mantle_gpt_audit3_cases( + "openai.gpt-6-astra", + 51, + 18, + 0, + ExactExpected( + spend=0.0020163, input_cost=0.0007293, output_cost=0.001287, prompt_tokens=51, completion_tokens=18 + ), + 27, + 0, + ExactExpected( + spend=0.031720975, + input_cost=0.029790475, + output_cost=0.0019305, + prompt_tokens=1667, + completion_tokens=27, + cache_creation_cost=0.029761875, + ), + ), + *_mantle_gpt_audit3_cases( + "openai.gpt-5.6-terra", + 51, + 28, + 8, + ExactExpected( + spend=0.00067452, input_cost=0.00015708, output_cost=0.00051744, prompt_tokens=51, completion_tokens=28 + ), + 28, + 0, + ExactExpected( + spend=0.00693385, + input_cost=0.00641641, + output_cost=0.00051744, + prompt_tokens=1667, + completion_tokens=28, + cache_creation_cost=0.00641025, + ), + ), +) + _MANTLE_CLAUDE_OPUS_MODEL: Final = "bedrock_mantle/anthropic.claude-opus-5-5" _ADAPTIVE_USAGE: Final = { "input_tokens": 24, @@ -737,7 +920,25 @@ _ADAPTIVE_EXPECTED: Final = ExactExpected( spend=0.0007524, input_cost=0.0001584, output_cost=0.000594, prompt_tokens=24, completion_tokens=18 ) -BEDROCK_MANTLE_CLAUDE_OPUS_5_5_INPUT_TEXT: Final = with_usage( + +def _mantle_without_usage_fields(case: CostTrackingTestCase) -> CostTrackingTestCase: + if not isinstance(case.response, JsonResponse): + raise TypeError(f"{case.name}: expected a JSON response fixture") + usage: Final = case.response.body.get("usage") + if not isinstance(usage, dict): + raise TypeError(f"{case.name}: expected a usage object") + response: Final = case.response.model_copy( + update={ + "body": { + **case.response.body, + "usage": {key: value for key, value in usage.items() if key not in {"inference_geo", "speed"}}, + } + } + ) + return case.model_copy(update={"response": response}) + + +BEDROCK_MANTLE_CLAUDE_OPUS_5_5_INPUT_TEXT: Final = _mantle_without_usage_fields(with_usage( CLAUDE_OPUS_5_5_INPUT_TEXT, "bedrock_mantle-anthropic.claude-opus-5-5-input_text", _ADAPTIVE_USAGE, @@ -751,7 +952,7 @@ BEDROCK_MANTLE_CLAUDE_OPUS_5_5_INPUT_TEXT: Final = with_usage( "allowed_openai_params": ["reasoning_effort"], }, } -) +)) _MANTLE_CLAUDE_STREAM_RESPONSE: Final = SseResponse( content_type="text/event-stream", @@ -841,7 +1042,7 @@ BEDROCK_MANTLE_CLAUDE_OPUS_5_5_MESSAGES_STREAM: Final = CLAUDE_OPUS_5_5_MESSAGES } ) -BEDROCK_MANTLE_CLAUDE_OPUS_5_5_CACHE_READ: Final = with_usage( +BEDROCK_MANTLE_CLAUDE_OPUS_5_5_CACHE_READ: Final = _mantle_without_usage_fields(with_usage( CLAUDE_OPUS_5_5_CACHE_READ, "bedrock_mantle-anthropic.claude-opus-5-5-cache_read", { @@ -856,7 +1057,7 @@ BEDROCK_MANTLE_CLAUDE_OPUS_5_5_CACHE_READ: Final = with_usage( ExactExpected( spend=0.00308319, input_cost=0.00153219, output_cost=0.001551, prompt_tokens=4225, completion_tokens=47 ), -).model_copy(update={"model": _MANTLE_CLAUDE_OPUS_MODEL}) +).model_copy(update={"model": _MANTLE_CLAUDE_OPUS_MODEL})) BEDROCK_MANTLE_CLAUDE_OPUS_5_5_CACHE_WRITE_5M: Final = with_usage( CLAUDE_OPUS_5_5_CACHE_WRITE_5M, @@ -892,12 +1093,12 @@ BEDROCK_MANTLE_CLAUDE_OPUS_5_5_CACHE_WRITE_1H: Final = with_usage( ), ).model_copy(update={"model": _MANTLE_CLAUDE_OPUS_MODEL}) -BEDROCK_MANTLE_CLAUDE_OPUS_5_5_THINKING: Final = with_usage( +BEDROCK_MANTLE_CLAUDE_OPUS_5_5_THINKING: Final = _mantle_without_usage_fields(with_usage( CLAUDE_OPUS_5_5_THINKING, "bedrock_mantle-anthropic.claude-opus-5-5-thinking", _ADAPTIVE_USAGE, _ADAPTIVE_EXPECTED, -).model_copy(update={"model": _MANTLE_CLAUDE_OPUS_MODEL}) +).model_copy(update={"model": _MANTLE_CLAUDE_OPUS_MODEL})) BEDROCK_MANTLE_CLAUDE_OPUS_5_5_TOOL_CALL: Final = with_usage( CLAUDE_OPUS_5_5_TOOL_CALL, @@ -921,6 +1122,21 @@ BEDROCK_MANTLE_CLAUDE_OPUS_5_5_CHAT_BRIDGED_JSON: Final = BEDROCK_MANTLE_CLAUDE_ } ) +BEDROCK_MANTLE_CLAUDE_OPUS_5_5_RESPONSES_JSON: Final = BEDROCK_MANTLE_CLAUDE_OPUS_5_5_INPUT_TEXT.model_copy( + update={ + "name": "bedrock_mantle-anthropic.claude-opus-5-5-responses_json", + "endpoint": "/v1/responses", + "request": {"model": "$MODEL", "input": "summarize the attached material in one line"}, + } +) +BEDROCK_MANTLE_CLAUDE_OPUS_5_5_RESPONSES_STREAM: Final = BEDROCK_MANTLE_CLAUDE_OPUS_5_5_STREAM.model_copy( + update={ + "name": "bedrock_mantle-anthropic.claude-opus-5-5-responses_stream", + "endpoint": "/v1/responses", + "request": {"model": "$MODEL", "input": "summarize the attached material in one line", "stream": True}, + } +) + BRIDGE_COST_PARITY: Final[tuple[BridgeCostParityTestCase, ...]] = ( _mantle_gpt_bridge(_MANTLE_GPT_61_SOL_CASES[0]), _mantle_gpt_bridge(_MANTLE_GPT_6_LUNA_CASES[0]), @@ -930,6 +1146,14 @@ BRIDGE_COST_PARITY: Final[tuple[BridgeCostParityTestCase, ...]] = ( native=BEDROCK_MANTLE_CLAUDE_OPUS_5_5_MESSAGES_INPUT_TEXT, bridged=BEDROCK_MANTLE_CLAUDE_OPUS_5_5_CHAT_BRIDGED_JSON, ), + *tuple( + BridgeCostParityTestCase(native=cases[0], bridged=cases[9]) + for cases in (_MANTLE_GPT_61_SOL_CASES, _MANTLE_GPT_6_LUNA_CASES, _MANTLE_GPT_6_ASTRA_CASES, _MANTLE_GPT_56_TERRA_CASES) + ), + *tuple( + BridgeCostParityTestCase(native=cases[1], bridged=cases[10]) + for cases in (_MANTLE_GPT_61_SOL_CASES, _MANTLE_GPT_6_LUNA_CASES, _MANTLE_GPT_6_ASTRA_CASES, _MANTLE_GPT_56_TERRA_CASES) + ), ) CASES: Final[tuple[CostTrackingTestCase, ...]] = ( @@ -945,11 +1169,14 @@ CASES: Final[tuple[CostTrackingTestCase, ...]] = ( *_MANTLE_GPT_6_LUNA_CASES, *_MANTLE_GPT_6_ASTRA_CASES, *_MANTLE_GPT_56_TERRA_CASES, + *_MANTLE_GPT_AUDIT3_CASES, BEDROCK_MANTLE_CLAUDE_OPUS_5_5_INPUT_TEXT, BEDROCK_MANTLE_CLAUDE_OPUS_5_5_STREAM, BEDROCK_MANTLE_CLAUDE_OPUS_5_5_MESSAGES_INPUT_TEXT, BEDROCK_MANTLE_CLAUDE_OPUS_5_5_MESSAGES_STREAM, BEDROCK_MANTLE_CLAUDE_OPUS_5_5_CHAT_BRIDGED_JSON, + BEDROCK_MANTLE_CLAUDE_OPUS_5_5_RESPONSES_JSON, + BEDROCK_MANTLE_CLAUDE_OPUS_5_5_RESPONSES_STREAM, BEDROCK_MANTLE_CLAUDE_OPUS_5_5_CACHE_READ, BEDROCK_MANTLE_CLAUDE_OPUS_5_5_CACHE_WRITE_5M, BEDROCK_MANTLE_CLAUDE_OPUS_5_5_CACHE_WRITE_1H, @@ -972,4 +1199,16 @@ PARITY: Final[tuple[StreamParityTestCase, ...]] = ( plain=BEDROCK_MANTLE_CLAUDE_OPUS_5_5_MESSAGES_INPUT_TEXT, streamed=BEDROCK_MANTLE_CLAUDE_OPUS_5_5_MESSAGES_STREAM, ), + StreamParityTestCase( + plain=BEDROCK_MANTLE_CLAUDE_OPUS_5_5_RESPONSES_JSON, + streamed=BEDROCK_MANTLE_CLAUDE_OPUS_5_5_RESPONSES_STREAM, + ), + *tuple( + StreamParityTestCase(plain=cases[3], streamed=cases[8]) + for cases in (_MANTLE_GPT_61_SOL_CASES, _MANTLE_GPT_6_LUNA_CASES, _MANTLE_GPT_6_ASTRA_CASES, _MANTLE_GPT_56_TERRA_CASES) + ), + *tuple( + StreamParityTestCase(plain=cases[9], streamed=cases[10]) + for cases in (_MANTLE_GPT_61_SOL_CASES, _MANTLE_GPT_6_LUNA_CASES, _MANTLE_GPT_6_ASTRA_CASES, _MANTLE_GPT_56_TERRA_CASES) + ), ) diff --git a/tests/integration/cost_calculation/cases/gemini.py b/tests/integration/cost_calculation/cases/gemini.py index bd64f4aeb11..ceca816c300 100644 --- a/tests/integration/cost_calculation/cases/gemini.py +++ b/tests/integration/cost_calculation/cases/gemini.py @@ -2,8 +2,10 @@ PARITY pairs the streamed case with its plain twin where both bill the same row.""" +import json from typing import Final, cast +from pydantic import JsonValue from integration.cost_calculation.bridge_cost_parity.case import BridgeCostParityTestCase, bridge_pairs from integration.cost_calculation.cost_tracking_case import ( CostTrackingTestCase, @@ -17,26 +19,74 @@ from integration.cost_calculation.cost_tracking_case import ( from integration.cost_calculation.stream_parity.case import StreamParityTestCase, sse_frames -def _gemini_case_response(response: JsonResponse | SseResponse, model_version: str) -> JsonResponse | SseResponse: +def _gemini_case_response( + response: JsonResponse | SseResponse, + model_version: str, + usage_fields: dict[str, JsonValue] | None = None, +) -> JsonResponse | SseResponse: match response: case JsonResponse(): - return response.model_copy(update={"body": {**response.body, "modelVersion": model_version}}) + usage: Final = response.body.get("usageMetadata") + body: Final = {**response.body, "modelVersion": model_version} + return response.model_copy( + update={ + "body": { + **body, + **( + { + "usageMetadata": { + **cast(dict[str, JsonValue], usage), + **usage_fields, + } + } + if usage_fields is not None and isinstance(usage, dict) + else {} + ), + } + } + ) case SseResponse(): return response.model_copy( update={ - "frames": tuple(frame.replace("gemini-3.8-flash", model_version) for frame in response.frames) + "frames": tuple( + _gemini_case_frame(frame, model_version, usage_fields) for frame in response.frames + ) } ) +def _gemini_case_frame( + frame: str, model_version: str, usage_fields: dict[str, JsonValue] | None +) -> str: + replaced: Final = frame.replace("gemini-3.8-flash", model_version) + if usage_fields is None: + return replaced + prefix, separator, payload = replaced.partition("data: ") + if not separator or payload == "[DONE]": + return replaced + event: Final = cast(dict[str, JsonValue], json.loads(payload)) + usage: Final = event.get("usageMetadata") + if not isinstance(usage, dict): + return replaced + return f"{prefix}data: {json.dumps({**event, 'usageMetadata': {**usage, **usage_fields}})}" + + def _gemini_case( - sibling: CostTrackingTestCase, name: str, model: str, model_version: str, expected: ExactExpected + sibling: CostTrackingTestCase, + name: str, + model: str, + model_version: str, + expected: ExactExpected, + *, + usage_fields: dict[str, JsonValue] | None = None, ) -> CostTrackingTestCase: return sibling.model_copy( update={ "name": name, "model": model, - "response": _gemini_case_response(cast(JsonResponse | SseResponse, sibling.response), model_version), + "response": _gemini_case_response( + cast(JsonResponse | SseResponse, sibling.response), model_version, usage_fields + ), "expected": expected, } ) @@ -1540,6 +1590,7 @@ GEMINI_GEMINI_3_8_FLASH_INPUT_TEXT: Final = CostTrackingTestCase( "candidatesTokenCount": 412, "totalTokenCount": 2252, "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 1840}], + "serviceTier": "standard", }, "modelVersion": "gemini-3.8-flash", }, @@ -1594,6 +1645,7 @@ GEMINI_GEMINI_3_8_FLASH_CACHE_READ: Final = CostTrackingTestCase( "totalTokenCount": 13308, "cachedContentTokenCount": 12288, "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 12928}], + "serviceTier": "standard", }, "modelVersion": "gemini-3.8-flash", }, @@ -1653,6 +1705,7 @@ GEMINI_GEMINI_3_8_FLASH_REASONING: Final = CostTrackingTestCase( "thoughtsTokenCount": 3480, "totalTokenCount": 5280, "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 1240}], + "serviceTier": "standard", }, "modelVersion": "gemini-3.8-flash", }, @@ -1934,6 +1987,7 @@ GEMINI_GEMINI_3_8_FLASH_STREAM: Final = CostTrackingTestCase( "candidatesTokenCount": 412, "totalTokenCount": 2252, "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 1840}], + "serviceTier": "standard", }, "modelVersion": "gemini-3.8-flash", }, @@ -3275,19 +3329,46 @@ GEMINI_GEMINI_3_1_PRO_CACHE_CREATION: Final = CostTrackingTestCase( GEMINI_GEMINI_4_ARGON_INPUT_TEXT: Final = _gemini_case( GEMINI_GEMINI_3_8_FLASH_INPUT_TEXT, "gemini-gemini-4-argon-input_text", "gemini/gemini-4-argon", - "gemini-4-argon", ExactExpected(spend=0.00858, input_cost=0.004048, output_cost=0.004532, prompt_tokens=1840, completion_tokens=412) + "gemini-4-argon", ExactExpected(spend=0.00858, input_cost=0.004048, output_cost=0.004532, prompt_tokens=1840, completion_tokens=412), + usage_fields={"candidatesTokensDetails": [{"modality": "TEXT", "tokenCount": 412}]}, ) GEMINI_GEMINI_4_ARGON_STREAM: Final = _gemini_case( GEMINI_GEMINI_3_8_FLASH_STREAM, "gemini-gemini-4-argon-stream", "gemini/gemini-4-argon", - "gemini-4-argon", ExactExpected(spend=0.00858, input_cost=0.004048, output_cost=0.004532, prompt_tokens=1840, completion_tokens=412) + "gemini-4-argon", ExactExpected(spend=0.00858, input_cost=0.004048, output_cost=0.004532, prompt_tokens=1840, completion_tokens=412), + usage_fields={"candidatesTokensDetails": [{"modality": "TEXT", "tokenCount": 412}]}, ) GEMINI_GEMINI_4_ARGON_CACHE_READ: Final = _gemini_case( GEMINI_GEMINI_3_8_FLASH_CACHE_READ, "gemini-gemini-4-argon-cache_read", "gemini/gemini-4-argon", - "gemini-4-argon", ExactExpected(spend=0.00693968, input_cost=0.00275968, output_cost=0.00418, prompt_tokens=12928, completion_tokens=380) + "gemini-4-argon", ExactExpected(spend=0.00693968, input_cost=0.00275968, output_cost=0.00418, prompt_tokens=12928, completion_tokens=380), + usage_fields={"candidatesTokensDetails": [{"modality": "TEXT", "tokenCount": 380}]}, ) GEMINI_GEMINI_4_ARGON_REASONING: Final = _gemini_case( GEMINI_GEMINI_3_8_FLASH_REASONING, "gemini-gemini-4-argon-reasoning", "gemini/gemini-4-argon", - "gemini-4-argon", ExactExpected(spend=0.047168, input_cost=0.002728, output_cost=0.04444, prompt_tokens=1240, completion_tokens=4040) + "gemini-4-argon", ExactExpected(spend=0.047168, input_cost=0.002728, output_cost=0.04444, prompt_tokens=1240, completion_tokens=4040), + usage_fields={"candidatesTokensDetails": [{"modality": "TEXT", "tokenCount": 560}]}, +) +GEMINI_GEMINI_4_ARGON_SERVICE_TIER_FLEX: Final = _gemini_case( + GEMINI_GEMINI_3_8_FLASH_SERVICE_TIER_FLEX, + "gemini-gemini-4-argon-service_tier_flex", + "gemini/gemini-4-argon", + "gemini-4-argon", + ExactExpected(spend=0.00429, input_cost=0.002024, output_cost=0.002266, prompt_tokens=1840, completion_tokens=412), +) +GEMINI_GEMINI_4_ARGON_SERVICE_TIER_PRIORITY: Final = _gemini_case( + GEMINI_GEMINI_3_8_FLASH_SERVICE_TIER_PRIORITY, + "gemini-gemini-4-argon-service_tier_priority", + "gemini/gemini-4-argon", + "gemini-4-argon", + ExactExpected(spend=0.015444, input_cost=0.0072864, output_cost=0.0081576, prompt_tokens=1840, completion_tokens=412), +) +GEMINI_GEMINI_3_1_PRO_PREVIEW_SERVICE_TIER_PRIORITY: Final = _gemini_case( + GEMINI_GEMINI_3_8_FLASH_SERVICE_TIER_PRIORITY, + "gemini-gemini-3.1-pro-preview-service_tier_priority", + "gemini/gemini-3.1-pro-preview", + "gemini-3.1-pro-preview", + ExactExpected( + spend=0.02018016, input_cost=0.0086112, output_cost=0.01156896, prompt_tokens=1840, completion_tokens=412 + ), ) GEMINI_GEMINI_4_ARGON_TOOL_CALL: Final = _gemini_case( GEMINI_GEMINI_3_8_FLASH_TOOL_CALL, "gemini-gemini-4-argon-tool_call", "gemini/gemini-4-argon", @@ -3296,28 +3377,60 @@ GEMINI_GEMINI_4_ARGON_TOOL_CALL: Final = _gemini_case( GEMINI_GEMINI_3_1_PRO_PREVIEW_INPUT_TEXT: Final = _gemini_case( GEMINI_GEMINI_3_8_FLASH_INPUT_TEXT, "gemini-gemini-3.1-pro-preview-input_text", "gemini/gemini-3.1-pro-preview", "gemini-3.1-pro-preview", - ExactExpected(spend=0.0112112, input_cost=0.004784, output_cost=0.0064272, prompt_tokens=1840, completion_tokens=412) + ExactExpected(spend=0.0112112, input_cost=0.004784, output_cost=0.0064272, prompt_tokens=1840, completion_tokens=412), + usage_fields={"serviceTier": "standard"}, ) GEMINI_GEMINI_3_1_PRO_PREVIEW_STREAM: Final = _gemini_case( GEMINI_GEMINI_3_8_FLASH_STREAM, "gemini-gemini-3.1-pro-preview-stream", "gemini/gemini-3.1-pro-preview", "gemini-3.1-pro-preview", - ExactExpected(spend=0.0112112, input_cost=0.004784, output_cost=0.0064272, prompt_tokens=1840, completion_tokens=412) + ExactExpected(spend=0.0112112, input_cost=0.004784, output_cost=0.0064272, prompt_tokens=1840, completion_tokens=412), + usage_fields={"serviceTier": "standard"}, ) GEMINI_GEMINI_3_1_PRO_PREVIEW_CACHE_READ: Final = _gemini_case( GEMINI_GEMINI_3_8_FLASH_CACHE_READ, "gemini-gemini-3.1-pro-preview-cache_read", "gemini/gemini-3.1-pro-preview", "gemini-3.1-pro-preview", - ExactExpected(spend=0.01078688, input_cost=0.00485888, output_cost=0.005928, prompt_tokens=12928, completion_tokens=380) + ExactExpected(spend=0.01078688, input_cost=0.00485888, output_cost=0.005928, prompt_tokens=12928, completion_tokens=380), + usage_fields={"serviceTier": "standard"}, ) GEMINI_GEMINI_3_1_PRO_PREVIEW_REASONING: Final = _gemini_case( GEMINI_GEMINI_3_8_FLASH_REASONING, "gemini-gemini-3.1-pro-preview-reasoning", "gemini/gemini-3.1-pro-preview", "gemini-3.1-pro-preview", - ExactExpected(spend=0.066248, input_cost=0.003224, output_cost=0.063024, prompt_tokens=1240, completion_tokens=4040) + ExactExpected(spend=0.066248, input_cost=0.003224, output_cost=0.063024, prompt_tokens=1240, completion_tokens=4040), + usage_fields={"serviceTier": "standard"}, ) GEMINI_GEMINI_3_1_PRO_PREVIEW_TOOL_CALL: Final = _gemini_case( GEMINI_GEMINI_3_8_FLASH_TOOL_CALL, "gemini-gemini-3.1-pro-preview-tool_call", "gemini/gemini-3.1-pro-preview", "gemini-3.1-pro-preview", ExactExpected(spend=0.0112112, input_cost=0.004784, output_cost=0.0064272, prompt_tokens=1840, completion_tokens=412) ) +GEMINI_GEMINI_3_1_PRO_PREVIEW_TIERED_INPUT_ABOVE_200K: Final = _gemini_case( + GEMINI_GEMINI_3_1_PRO_TIERED_INPUT_ABOVE_200K, + "gemini-gemini-3.1-pro-preview-tiered_input_above_200k", + "gemini/gemini-3.1-pro-preview", + "gemini-3.1-pro-preview", + ExactExpected( + spend=1.079468, input_cost=1.06496, output_cost=0.014508, prompt_tokens=204800, completion_tokens=620 + ), + usage_fields={"serviceTier": "standard"}, +) +GEMINI_GEMINI_3_1_PRO_PREVIEW_TIERED_CACHE_READ_ABOVE_200K: Final = _gemini_case( + GEMINI_GEMINI_3_1_PRO_TIERED_CACHE_READ_ABOVE_200K, + "gemini-gemini-3.1-pro-preview-tiered_cache_read_above_200k", + "gemini/gemini-3.1-pro-preview", + "gemini-3.1-pro-preview", + ExactExpected( + spend=0.13742976, input_cost=0.12619776, output_cost=0.011232, prompt_tokens=205824, completion_tokens=480 + ), + usage_fields={"serviceTier": "standard"}, +) +GEMINI_GEMINI_4_ARGON_ROUTE_BRIDGES: Final = bridge_pairs( + (GEMINI_GEMINI_4_ARGON_INPUT_TEXT, GEMINI_GEMINI_4_ARGON_STREAM), + ("/v1/responses", "/v1/messages"), +) +GEMINI_GEMINI_3_1_PRO_PREVIEW_ROUTE_BRIDGES: Final = bridge_pairs( + (GEMINI_GEMINI_3_1_PRO_PREVIEW_INPUT_TEXT, GEMINI_GEMINI_3_1_PRO_PREVIEW_STREAM), + ("/v1/responses", "/v1/messages"), +) CASES: Final[tuple[CostTrackingTestCase, ...]] = ( GEMINI_GEMINI_3_1_PRO_INPUT_TEXT, @@ -3379,12 +3492,19 @@ CASES: Final[tuple[CostTrackingTestCase, ...]] = ( GEMINI_GEMINI_4_ARGON_STREAM, GEMINI_GEMINI_4_ARGON_CACHE_READ, GEMINI_GEMINI_4_ARGON_REASONING, + GEMINI_GEMINI_4_ARGON_SERVICE_TIER_FLEX, + GEMINI_GEMINI_4_ARGON_SERVICE_TIER_PRIORITY, GEMINI_GEMINI_4_ARGON_TOOL_CALL, GEMINI_GEMINI_3_1_PRO_PREVIEW_INPUT_TEXT, GEMINI_GEMINI_3_1_PRO_PREVIEW_STREAM, GEMINI_GEMINI_3_1_PRO_PREVIEW_CACHE_READ, + GEMINI_GEMINI_3_1_PRO_PREVIEW_TIERED_INPUT_ABOVE_200K, + GEMINI_GEMINI_3_1_PRO_PREVIEW_TIERED_CACHE_READ_ABOVE_200K, GEMINI_GEMINI_3_1_PRO_PREVIEW_REASONING, GEMINI_GEMINI_3_1_PRO_PREVIEW_TOOL_CALL, + GEMINI_GEMINI_3_1_PRO_PREVIEW_SERVICE_TIER_PRIORITY, + *(pair.bridged for pair in GEMINI_GEMINI_4_ARGON_ROUTE_BRIDGES), + *(pair.bridged for pair in GEMINI_GEMINI_3_1_PRO_PREVIEW_ROUTE_BRIDGES), ) PARITY: Final[tuple[StreamParityTestCase, ...]] = ( @@ -3410,6 +3530,22 @@ PARITY: Final[tuple[StreamParityTestCase, ...]] = ( StreamParityTestCase( plain=GEMINI_GEMINI_3_1_PRO_PREVIEW_INPUT_TEXT, streamed=GEMINI_GEMINI_3_1_PRO_PREVIEW_STREAM ), + StreamParityTestCase( + plain=GEMINI_GEMINI_4_ARGON_ROUTE_BRIDGES[0].bridged, + streamed=GEMINI_GEMINI_4_ARGON_ROUTE_BRIDGES[2].bridged, + ), + StreamParityTestCase( + plain=GEMINI_GEMINI_4_ARGON_ROUTE_BRIDGES[1].bridged, + streamed=GEMINI_GEMINI_4_ARGON_ROUTE_BRIDGES[3].bridged, + ), + StreamParityTestCase( + plain=GEMINI_GEMINI_3_1_PRO_PREVIEW_ROUTE_BRIDGES[0].bridged, + streamed=GEMINI_GEMINI_3_1_PRO_PREVIEW_ROUTE_BRIDGES[2].bridged, + ), + StreamParityTestCase( + plain=GEMINI_GEMINI_3_1_PRO_PREVIEW_ROUTE_BRIDGES[1].bridged, + streamed=GEMINI_GEMINI_3_1_PRO_PREVIEW_ROUTE_BRIDGES[3].bridged, + ), ) BRIDGE_COST_PARITY: Final[tuple[BridgeCostParityTestCase, ...]] = ( @@ -3422,6 +3558,6 @@ BRIDGE_COST_PARITY: Final[tuple[BridgeCostParityTestCase, ...]] = ( ), ("/v1/responses", "/v1/messages"), ), - *bridge_pairs((GEMINI_GEMINI_4_ARGON_INPUT_TEXT,), ("/v1/responses",)), - *bridge_pairs((GEMINI_GEMINI_3_1_PRO_PREVIEW_INPUT_TEXT,), ("/v1/responses",)), + *GEMINI_GEMINI_4_ARGON_ROUTE_BRIDGES, + *GEMINI_GEMINI_3_1_PRO_PREVIEW_ROUTE_BRIDGES, ) diff --git a/tests/integration/cost_calculation/cases/misc.py b/tests/integration/cost_calculation/cases/misc.py index 1a8b2fcb66b..760f77f91cc 100644 --- a/tests/integration/cost_calculation/cases/misc.py +++ b/tests/integration/cost_calculation/cases/misc.py @@ -19,13 +19,17 @@ from integration.cost_calculation.cost_tracking_case import ( RecountRates, SseResponse, WavUpload, + with_usage, ) from integration.cost_calculation.stream_parity.case import StreamParityTestCase, sse_frames +from integration.cost_calculation.bridge_cost_parity.case import BridgeCostParityTestCase, bridge_pairs from integration.cost_calculation.cases.anthropic import ( CLAUDE_FABLE_5_1_CACHE_READ, CLAUDE_FABLE_5_1_CACHE_WRITE_1H, CLAUDE_FABLE_5_1_CACHE_WRITE_5M, CLAUDE_FABLE_5_1_INPUT_TEXT, + CLAUDE_FABLE_5_1_MESSAGES_INPUT_TEXT, + CLAUDE_FABLE_5_1_MESSAGES_STREAM, CLAUDE_FABLE_5_1_STREAM, CLAUDE_FABLE_5_1_THINKING, CLAUDE_FABLE_5_1_TOOL_CALL, @@ -33,12 +37,30 @@ from integration.cost_calculation.cases.anthropic import ( CLAUDE_HAIKU_5_5_CACHE_WRITE_1H, CLAUDE_HAIKU_5_5_CACHE_WRITE_5M, CLAUDE_HAIKU_5_5_INPUT_TEXT, + CLAUDE_HAIKU_5_5_MESSAGES_INPUT_TEXT, + CLAUDE_HAIKU_5_5_MESSAGES_STREAM, CLAUDE_HAIKU_5_5_STREAM, CLAUDE_HAIKU_5_5_THINKING, CLAUDE_HAIKU_5_5_TOOL_CALL, ) +def _azure_ai_claude_messages_case( + response_case: CostTrackingTestCase, + request_case: CostTrackingTestCase, + name: str, + model: str, +) -> CostTrackingTestCase: + return response_case.model_copy( + update={ + "name": name, + "model": model, + "endpoint": "/v1/messages", + "request": request_case.request, + } + ) + + def _azure_ai_claude_usage( input_tokens: int, output_tokens: int, @@ -2027,7 +2049,66 @@ AZURE_AI_CLAUDE_HAIKU_5_5_TOOL_CALL: Final = _azure_ai_claude_shape_case( ), ) - +AZURE_AI_CLAUDE_FABLE_5_1_MESSAGES_INPUT_TEXT: Final = _azure_ai_claude_messages_case( + AZURE_AI_CLAUDE_FABLE_5_1_INPUT_TEXT, + CLAUDE_FABLE_5_1_MESSAGES_INPUT_TEXT, + "azure_ai/claude-fable-5-1-messages_input_text", + "azure_ai/claude-fable-5-1", +) +AZURE_AI_CLAUDE_FABLE_5_1_MESSAGES_STREAM: Final = _azure_ai_claude_messages_case( + AZURE_AI_CLAUDE_FABLE_5_1_STREAM, + CLAUDE_FABLE_5_1_MESSAGES_STREAM, + "azure_ai/claude-fable-5-1-messages_stream", + "azure_ai/claude-fable-5-1", +) +AZURE_AI_CLAUDE_FABLE_5_1_ROUTE_BRIDGES: Final = bridge_pairs( + (AZURE_AI_CLAUDE_FABLE_5_1_MESSAGES_INPUT_TEXT, AZURE_AI_CLAUDE_FABLE_5_1_MESSAGES_STREAM), + ("/v1/responses",), +) +AZURE_AI_CLAUDE_HAIKU_5_5_MESSAGES_INPUT_TEXT: Final = _azure_ai_claude_messages_case( + AZURE_AI_CLAUDE_HAIKU_5_5_INPUT_TEXT, + CLAUDE_HAIKU_5_5_MESSAGES_INPUT_TEXT, + "azure_ai/claude-haiku-5-5-messages_input_text", + "azure_ai/claude-haiku-5-5", +) +AZURE_AI_CLAUDE_HAIKU_5_5_MESSAGES_STREAM: Final = _azure_ai_claude_messages_case( + AZURE_AI_CLAUDE_HAIKU_5_5_STREAM, + CLAUDE_HAIKU_5_5_MESSAGES_STREAM, + "azure_ai/claude-haiku-5-5-messages_stream", + "azure_ai/claude-haiku-5-5", +) +AZURE_AI_CLAUDE_HAIKU_5_5_ROUTE_BRIDGES: Final = bridge_pairs( + (AZURE_AI_CLAUDE_HAIKU_5_5_MESSAGES_INPUT_TEXT, AZURE_AI_CLAUDE_HAIKU_5_5_MESSAGES_STREAM), + ("/v1/responses",), +) +AZURE_AI_CLAUDE_HAIKU_5_5_TIERED_INPUT_ABOVE_100K: Final = with_usage( + AZURE_AI_CLAUDE_HAIKU_5_5_INPUT_TEXT, + "azure_ai/claude-haiku-5-5-tiered_input_above_100k", + { + "input_tokens": 150000, + "output_tokens": 480, + "cache_read_input_tokens": 0, + "cache_creation_input_tokens": 0, + "cache_creation": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 0}, + }, + ExactExpected( + spend=0.14478, input_cost=0.1425, output_cost=0.00228, prompt_tokens=150000, completion_tokens=480 + ), +) +AZURE_AI_CLAUDE_HAIKU_5_5_TIERED_CACHE_READ_ABOVE_100K: Final = with_usage( + AZURE_AI_CLAUDE_HAIKU_5_5_CACHE_READ, + "azure_ai/claude-haiku-5-5-tiered_cache_read_above_100k", + { + "input_tokens": 4096, + "output_tokens": 480, + "cache_read_input_tokens": 110592, + "cache_creation_input_tokens": 0, + "cache_creation": {"ephemeral_5m_input_tokens": 0, "ephemeral_1h_input_tokens": 0}, + }, + ExactExpected( + spend=0.01667744, input_cost=0.01439744, output_cost=0.00228, prompt_tokens=114688, completion_tokens=480 + ), +) CASES: Final[tuple[CostTrackingTestCase, ...]] = ( PERPLEXITY_SONAR_NEXT_OFF_PEAK, NOVA_NEXT_TRANSCRIPTIONS_PER_SECOND, @@ -2088,13 +2169,21 @@ CASES: Final[tuple[CostTrackingTestCase, ...]] = ( AZURE_AI_CLAUDE_FABLE_5_1_CACHE_WRITE_1H, AZURE_AI_CLAUDE_FABLE_5_1_THINKING, AZURE_AI_CLAUDE_FABLE_5_1_TOOL_CALL, + AZURE_AI_CLAUDE_FABLE_5_1_MESSAGES_INPUT_TEXT, + AZURE_AI_CLAUDE_FABLE_5_1_MESSAGES_STREAM, + *(pair.bridged for pair in AZURE_AI_CLAUDE_FABLE_5_1_ROUTE_BRIDGES), AZURE_AI_CLAUDE_HAIKU_5_5_INPUT_TEXT, AZURE_AI_CLAUDE_HAIKU_5_5_STREAM, AZURE_AI_CLAUDE_HAIKU_5_5_CACHE_READ, + AZURE_AI_CLAUDE_HAIKU_5_5_TIERED_INPUT_ABOVE_100K, + AZURE_AI_CLAUDE_HAIKU_5_5_TIERED_CACHE_READ_ABOVE_100K, AZURE_AI_CLAUDE_HAIKU_5_5_CACHE_WRITE_5M, AZURE_AI_CLAUDE_HAIKU_5_5_CACHE_WRITE_1H, AZURE_AI_CLAUDE_HAIKU_5_5_THINKING, AZURE_AI_CLAUDE_HAIKU_5_5_TOOL_CALL, + AZURE_AI_CLAUDE_HAIKU_5_5_MESSAGES_INPUT_TEXT, + AZURE_AI_CLAUDE_HAIKU_5_5_MESSAGES_STREAM, + *(pair.bridged for pair in AZURE_AI_CLAUDE_HAIKU_5_5_ROUTE_BRIDGES), ) PARITY: Final[tuple[StreamParityTestCase, ...]] = ( @@ -2105,4 +2194,25 @@ PARITY: Final[tuple[StreamParityTestCase, ...]] = ( StreamParityTestCase( plain=AZURE_AI_CLAUDE_HAIKU_5_5_INPUT_TEXT, streamed=AZURE_AI_CLAUDE_HAIKU_5_5_STREAM ), + StreamParityTestCase( + plain=AZURE_AI_CLAUDE_FABLE_5_1_MESSAGES_INPUT_TEXT, + streamed=AZURE_AI_CLAUDE_FABLE_5_1_MESSAGES_STREAM, + ), + StreamParityTestCase( + plain=AZURE_AI_CLAUDE_FABLE_5_1_ROUTE_BRIDGES[0].bridged, + streamed=AZURE_AI_CLAUDE_FABLE_5_1_ROUTE_BRIDGES[1].bridged, + ), + StreamParityTestCase( + plain=AZURE_AI_CLAUDE_HAIKU_5_5_MESSAGES_INPUT_TEXT, + streamed=AZURE_AI_CLAUDE_HAIKU_5_5_MESSAGES_STREAM, + ), + StreamParityTestCase( + plain=AZURE_AI_CLAUDE_HAIKU_5_5_ROUTE_BRIDGES[0].bridged, + streamed=AZURE_AI_CLAUDE_HAIKU_5_5_ROUTE_BRIDGES[1].bridged, + ), +) + +BRIDGE_COST_PARITY: Final[tuple[BridgeCostParityTestCase, ...]] = ( + *AZURE_AI_CLAUDE_FABLE_5_1_ROUTE_BRIDGES, + *AZURE_AI_CLAUDE_HAIKU_5_5_ROUTE_BRIDGES, ) diff --git a/tests/integration/cost_calculation/cases/openai.py b/tests/integration/cost_calculation/cases/openai.py index 85ac05dd69d..f610fd9591a 100644 --- a/tests/integration/cost_calculation/cases/openai.py +++ b/tests/integration/cost_calculation/cases/openai.py @@ -34,6 +34,256 @@ def _stream_for_model(response: SseResponse, model: str) -> SseResponse: return response.model_copy(update={"frames": tuple(frame.replace("gpt-6-luna", model) for frame in response.frames)}) +def _responses_json_body( + model: str, + *, + include_content_filters: bool = False, + content_filters: JsonValue = None, + include_reasoning: bool = True, +) -> dict[str, JsonValue]: + reasoning_item: Final[dict[str, JsonValue]] = { + "content": [], + "encrypted_content": "$ENCRYPTED_CONTENT", + "id": "rs_$REQUEST_ID", + "summary": [], + "type": "reasoning", + } + message_item: Final[dict[str, JsonValue]] = { + "content": [ + { + "annotations": [], + "logprobs": [], + "text": "scripted response", + "type": "output_text", + } + ], + "id": "msg_$REQUEST_ID", + "phase": "final_answer", + "role": "assistant", + "status": "completed", + "type": "message", + } + output: Final = [reasoning_item, message_item] if include_reasoning else [message_item] + body: Final[dict[str, JsonValue]] = { + "access_programs": [], + "background": False, + "completed_at": 1791855453, + "created_at": 1791855452, + "error": None, + "frequency_penalty": None, + "id": "resp_$REQUEST_ID", + "incomplete_details": None, + "instructions": None, + "max_output_tokens": None, + "max_tool_calls": None, + "metadata": {}, + "moderation": None, + "model": model, + "object": "response", + "output": output, + "parallel_tool_calls": True, + "presence_penalty": None, + "previous_response_id": None, + "prompt_cache_key": None, + "prompt_cache_retention": None, + "reasoning": {"effort": "medium"}, + "safety_identifier": None, + "service_tier": "default", + "status": "completed", + "store": True, + "temperature": 1.0, + "text": None, + "tool_choice": "auto", + "tool_usage": None, + "tools": [], + "top_logprobs": None, + "top_p": 1.0, + "truncation": "disabled", + "usage": { + "input_tokens": 1840, + "input_tokens_details": {"cache_write_tokens": 0, "cached_tokens": 0}, + "output_tokens": 412, + "output_tokens_details": {"reasoning_tokens": 0}, + "total_tokens": 2252, + }, + "user": None, + } + return {**body, "content_filters": content_filters} if include_content_filters else body + + +def _responses_stream_response( + model: str, + delta_count: int, + *, + include_content_filters: bool = False, + content_filters: JsonValue = None, + include_reasoning: bool = True, +) -> SseResponse: + response_body: Final = _responses_json_body( + model, + include_content_filters=include_content_filters, + content_filters=content_filters, + include_reasoning=include_reasoning, + ) + in_progress_body: Final = { + **response_body, + "output": [], + "status": "in_progress", + "usage": None, + } + reasoning_item: Final = { + "content": [], + "encrypted_content": "$ENCRYPTED_CONTENT", + "id": "rs_$REQUEST_ID", + "summary": [], + "type": "reasoning", + } + message_item: Final = { + "content": [], + "id": "msg_$REQUEST_ID", + "role": "assistant", + "status": "in_progress", + "type": "message", + } + delta_events: Final = tuple( + { + "content_index": 0, + "delta": "scripted response" if index == 0 else "", + "item_id": "msg_$REQUEST_ID", + "logprobs": [], + "obfuscation": None, + "output_index": 1 if include_reasoning else 0, + "sequence_number": (6 if include_reasoning else 4) + index, + "type": "response.output_text.delta", + } + for index in range(delta_count) + ) + delta_start: Final = 6 if include_reasoning else 4 + message_output_index: Final = 1 if include_reasoning else 0 + reasoning_events: Final = ( + ( + { + "item": reasoning_item, + "output_index": 0, + "sequence_number": 2, + "type": "response.output_item.added", + }, + { + "item": reasoning_item, + "output_index": 0, + "sequence_number": 3, + "type": "response.output_item.done", + }, + ) + if include_reasoning + else () + ) + completed_message_item: Final = { + "content": [ + { + "annotations": [], + "logprobs": [], + "text": "scripted response", + "type": "output_text", + } + ], + "id": "msg_$REQUEST_ID", + "phase": "final_answer", + "role": "assistant", + "status": "completed", + "type": "message", + } + return SseResponse( + content_type="text/event-stream", + frames=sse_frames( + { + "response": in_progress_body, + "sequence_number": 0, + "type": "response.created", + }, + { + "response": in_progress_body, + "sequence_number": 1, + "type": "response.in_progress", + }, + *reasoning_events, + { + "item": message_item, + "output_index": message_output_index, + "sequence_number": 4 if include_reasoning else 2, + "type": "response.output_item.added", + }, + { + "content_index": 0, + "item_id": "msg_$REQUEST_ID", + "output_index": message_output_index, + "part": {"annotations": [], "logprobs": [], "text": "", "type": "output_text"}, + "sequence_number": 5 if include_reasoning else 3, + "type": "response.content_part.added", + }, + *delta_events, + { + "content_index": 0, + "item_id": "msg_$REQUEST_ID", + "logprobs": [], + "output_index": message_output_index, + "sequence_number": delta_start + delta_count, + "text": "scripted response", + "type": "response.output_text.done", + }, + { + "content_index": 0, + "item_id": "msg_$REQUEST_ID", + "output_index": message_output_index, + "part": { + "annotations": [], + "logprobs": [], + "text": "scripted response", + "type": "output_text", + }, + "sequence_number": delta_start + delta_count + 1, + "type": "response.content_part.done", + }, + { + "item": completed_message_item, + "output_index": message_output_index, + "sequence_number": delta_start + delta_count + 2, + "type": "response.output_item.done", + }, + { + "response": response_body, + "sequence_number": delta_start + delta_count + 3, + "type": "response.completed", + }, + ), + ) + + +def _responses_stream_case( + base: CostTrackingTestCase, + name: str, + response_model: str, + delta_count: int, + *, + include_content_filters: bool = False, + content_filters: JsonValue = None, + include_reasoning: bool = True, +) -> CostTrackingTestCase: + return base.model_copy( + update={ + "name": name, + "request": {**base.request, "stream": True}, + "response": _responses_stream_response( + response_model, + delta_count, + include_content_filters=include_content_filters, + content_filters=content_filters, + include_reasoning=include_reasoning, + ), + } + ) + + # gpt-5.3-codex GPT_5_3_CODEX_INPUT_TEXT: Final = CostTrackingTestCase( name="gpt-5.3-codex-input_text", @@ -11356,6 +11606,31 @@ GPT_5_6_TERRA_TOOL_CALL: Final = GPT_6_LUNA_TOOL_CALL.model_copy( } ) +GPT_5_6_TERRA_TIERED_INPUT_ABOVE_272K: Final = with_usage( + GPT_5_6_TERRA_INPUT_TEXT, + "gpt-5.6-terra-tiered_input_above_272k", + { + "prompt_tokens": 300000, + "total_tokens": 300412, + "prompt_tokens_details": {"audio_tokens": 0, "cache_write_tokens": 0, "cached_tokens": 0}, + }, + ExactExpected( + spend=1.5696408, input_cost=1.56, output_cost=0.0096408, prompt_tokens=300000, completion_tokens=412 + ), +) +GPT_5_6_TERRA_TIERED_CACHE_READ_ABOVE_272K: Final = with_usage( + GPT_5_6_TERRA_INPUT_TEXT, + "gpt-5.6-terra-tiered_cache_read_above_272k", + { + "prompt_tokens": 300000, + "total_tokens": 300412, + "prompt_tokens_details": {"audio_tokens": 0, "cache_write_tokens": 0, "cached_tokens": 200000}, + }, + ExactExpected( + spend=0.6336408, input_cost=0.624, output_cost=0.0096408, prompt_tokens=300000, completion_tokens=412 + ), +) + GPT_5_6_SLASH_ALIAS: Final = GPT_5_6_INPUT_TEXT.model_copy( update={"name": "gpt-5.6-slash_alias", "alias": "openai/custom/gpt-5.6-$MARKER"} @@ -11446,6 +11721,63 @@ GPT_5_6_RESPONSES_IMAGE_GENERATION_TOOL_TOKENS: Final = CostTrackingTestCase( ), ) +GPT_6_1_SOL_RESPONSES_NATIVE_STREAM: Final = _responses_stream_case( + GPT_6_1_SOL_RESPONSES_NATIVE_JSON, + "gpt-6.1-sol-responses_native_stream", + "gpt-6.1-sol", + 1, + include_reasoning=False, +) +GPT_6_LUNA_RESPONSES_NATIVE_STREAM: Final = _responses_stream_case( + GPT_6_LUNA_RESPONSES_NATIVE_JSON, + "gpt-6-luna-responses_native_stream", + "gpt-6-luna", + 45, +) +GPT_6_ASTRA_RESPONSES_NATIVE_STREAM: Final = _responses_stream_case( + GPT_6_ASTRA_RESPONSES_NATIVE_JSON, + "gpt-6-astra-responses_native_stream", + "gpt-6-astra", + 36, +) +GPT_5_6_TERRA_RESPONSES_NATIVE_STREAM: Final = _responses_stream_case( + GPT_5_6_TERRA_RESPONSES_NATIVE_JSON, + "gpt-5.6-terra-responses_native_stream", + "gpt-5.6-terra", + 95, +) +GPT_5_6_TERRA_MESSAGES_BRIDGED_JSON: Final = bridge_pairs( + (GPT_5_6_TERRA_RESPONSES_NATIVE_JSON,), ("/v1/messages",) +)[0].bridged.model_copy(update={"name": "gpt-5.6-terra-messages_bridged_json"}) +GPT_6_1_SOL_MESSAGES_BRIDGED_STREAM: Final = GPT_6_1_SOL_MESSAGES_BRIDGED_JSON.model_copy( + update={ + "name": "gpt-6.1-sol-messages_bridged_stream", + "request": {**GPT_6_1_SOL_MESSAGES_BRIDGED_JSON.request, "stream": True}, + "response": GPT_6_1_SOL_RESPONSES_NATIVE_STREAM.response, + } +) +GPT_6_LUNA_MESSAGES_BRIDGED_STREAM: Final = GPT_6_LUNA_MESSAGES_BRIDGED_JSON.model_copy( + update={ + "name": "gpt-6-luna-messages_bridged_stream", + "request": {**GPT_6_LUNA_MESSAGES_BRIDGED_JSON.request, "stream": True}, + "response": GPT_6_LUNA_RESPONSES_NATIVE_STREAM.response, + } +) +GPT_6_ASTRA_MESSAGES_BRIDGED_STREAM: Final = GPT_6_ASTRA_MESSAGES_BRIDGED_JSON.model_copy( + update={ + "name": "gpt-6-astra-messages_bridged_stream", + "request": {**GPT_6_ASTRA_MESSAGES_BRIDGED_JSON.request, "stream": True}, + "response": GPT_6_ASTRA_RESPONSES_NATIVE_STREAM.response, + } +) +GPT_5_6_TERRA_MESSAGES_BRIDGED_STREAM: Final = GPT_5_6_TERRA_MESSAGES_BRIDGED_JSON.model_copy( + update={ + "name": "gpt-5.6-terra-messages_bridged_stream", + "request": {**GPT_5_6_TERRA_MESSAGES_BRIDGED_JSON.request, "stream": True}, + "response": GPT_5_6_TERRA_RESPONSES_NATIVE_STREAM.response, + } +) + CASES: Final[tuple[CostTrackingTestCase, ...]] = ( GPT_5_6_SLASH_ALIAS, GPT_5_6_SLASH_ALIAS_STREAM, @@ -11659,6 +11991,8 @@ CASES: Final[tuple[CostTrackingTestCase, ...]] = ( GPT_6_LUNA_CACHE_WRITE_5M, GPT_6_ASTRA_CACHE_WRITE_5M, GPT_5_6_TERRA_INPUT_TEXT, + GPT_5_6_TERRA_TIERED_INPUT_ABOVE_272K, + GPT_5_6_TERRA_TIERED_CACHE_READ_ABOVE_272K, GPT_5_6_TERRA_STREAM, GPT_5_6_TERRA_RESPONSES_NATIVE_JSON, GPT_5_6_TERRA_RESPONSES_DEPLOYMENT_JSON, @@ -11669,6 +12003,15 @@ CASES: Final[tuple[CostTrackingTestCase, ...]] = ( GPT_5_6_TERRA_SERVICE_TIER_FLEX, GPT_5_6_TERRA_SERVICE_TIER_PRIORITY, GPT_5_6_TERRA_TOOL_CALL, + GPT_6_1_SOL_RESPONSES_NATIVE_STREAM, + GPT_6_LUNA_RESPONSES_NATIVE_STREAM, + GPT_6_ASTRA_RESPONSES_NATIVE_STREAM, + GPT_5_6_TERRA_RESPONSES_NATIVE_STREAM, + GPT_5_6_TERRA_MESSAGES_BRIDGED_JSON, + GPT_6_1_SOL_MESSAGES_BRIDGED_STREAM, + GPT_6_LUNA_MESSAGES_BRIDGED_STREAM, + GPT_6_ASTRA_MESSAGES_BRIDGED_STREAM, + GPT_5_6_TERRA_MESSAGES_BRIDGED_STREAM, ) PARITY: Final[tuple[StreamParityTestCase, ...]] = ( @@ -11698,6 +12041,16 @@ PARITY: Final[tuple[StreamParityTestCase, ...]] = ( StreamParityTestCase(plain=GPT_6_ASTRA_INPUT_TEXT, streamed=GPT_6_ASTRA_STREAM), StreamParityTestCase(plain=GPT_6_1_SOL_INPUT_TEXT, streamed=GPT_6_1_SOL_STREAM), StreamParityTestCase(plain=GPT_5_6_TERRA_INPUT_TEXT, streamed=GPT_5_6_TERRA_STREAM), + StreamParityTestCase(plain=GPT_6_1_SOL_RESPONSES_NATIVE_JSON, streamed=GPT_6_1_SOL_RESPONSES_NATIVE_STREAM), + StreamParityTestCase(plain=GPT_6_LUNA_RESPONSES_NATIVE_JSON, streamed=GPT_6_LUNA_RESPONSES_NATIVE_STREAM), + StreamParityTestCase(plain=GPT_6_ASTRA_RESPONSES_NATIVE_JSON, streamed=GPT_6_ASTRA_RESPONSES_NATIVE_STREAM), + StreamParityTestCase(plain=GPT_5_6_TERRA_RESPONSES_NATIVE_JSON, streamed=GPT_5_6_TERRA_RESPONSES_NATIVE_STREAM), + StreamParityTestCase(plain=GPT_6_1_SOL_MESSAGES_BRIDGED_JSON, streamed=GPT_6_1_SOL_MESSAGES_BRIDGED_STREAM), + StreamParityTestCase(plain=GPT_6_LUNA_MESSAGES_BRIDGED_JSON, streamed=GPT_6_LUNA_MESSAGES_BRIDGED_STREAM), + StreamParityTestCase(plain=GPT_6_ASTRA_MESSAGES_BRIDGED_JSON, streamed=GPT_6_ASTRA_MESSAGES_BRIDGED_STREAM), + StreamParityTestCase( + plain=GPT_5_6_TERRA_MESSAGES_BRIDGED_JSON, streamed=GPT_5_6_TERRA_MESSAGES_BRIDGED_STREAM + ), ) BRIDGE_COST_PARITY: Final[tuple[BridgeCostParityTestCase, ...]] = ( @@ -11726,6 +12079,13 @@ BRIDGE_COST_PARITY: Final[tuple[BridgeCostParityTestCase, ...]] = ( BridgeCostParityTestCase(native=GPT_6_LUNA_RESPONSES_NATIVE_JSON, bridged=GPT_6_LUNA_MESSAGES_BRIDGED_JSON), BridgeCostParityTestCase(native=GPT_6_ASTRA_RESPONSES_NATIVE_JSON, bridged=GPT_6_ASTRA_MESSAGES_BRIDGED_JSON), BridgeCostParityTestCase(native=GPT_6_1_SOL_RESPONSES_NATIVE_JSON, bridged=GPT_6_1_SOL_MESSAGES_BRIDGED_JSON), + BridgeCostParityTestCase(native=GPT_5_6_TERRA_RESPONSES_NATIVE_JSON, bridged=GPT_5_6_TERRA_MESSAGES_BRIDGED_JSON), + BridgeCostParityTestCase(native=GPT_6_1_SOL_RESPONSES_NATIVE_STREAM, bridged=GPT_6_1_SOL_MESSAGES_BRIDGED_STREAM), + BridgeCostParityTestCase(native=GPT_6_LUNA_RESPONSES_NATIVE_STREAM, bridged=GPT_6_LUNA_MESSAGES_BRIDGED_STREAM), + BridgeCostParityTestCase(native=GPT_6_ASTRA_RESPONSES_NATIVE_STREAM, bridged=GPT_6_ASTRA_MESSAGES_BRIDGED_STREAM), + BridgeCostParityTestCase( + native=GPT_5_6_TERRA_RESPONSES_NATIVE_STREAM, bridged=GPT_5_6_TERRA_MESSAGES_BRIDGED_STREAM + ), BridgeCostParityTestCase(native=GPT_6_SOL_RESPONSES_DEPLOYMENT_JSON, bridged=GPT_6_SOL_CHAT_BRIDGED_JSON), BridgeCostParityTestCase(native=GPT_6_LUNA_RESPONSES_DEPLOYMENT_JSON, bridged=GPT_6_LUNA_CHAT_BRIDGED_JSON), BridgeCostParityTestCase(native=GPT_6_ASTRA_RESPONSES_DEPLOYMENT_JSON, bridged=GPT_6_ASTRA_CHAT_BRIDGED_JSON), diff --git a/tests/integration/cost_calculation/cases/vertex_ai.py b/tests/integration/cost_calculation/cases/vertex_ai.py index 761078ef593..0e95eb865bf 100644 --- a/tests/integration/cost_calculation/cases/vertex_ai.py +++ b/tests/integration/cost_calculation/cases/vertex_ai.py @@ -2,6 +2,7 @@ PARITY pairs the streamed case with its plain twin where both bill the same row.""" +import json from typing import Final, cast from pydantic import JsonValue @@ -11,13 +12,19 @@ from integration.cost_calculation.cases.anthropic import ( CLAUDE_FABLE_5_1_CACHE_WRITE_1H, CLAUDE_FABLE_5_1_CACHE_WRITE_5M, CLAUDE_FABLE_5_1_INPUT_TEXT, + CLAUDE_FABLE_5_1_MESSAGES_INPUT_TEXT, + CLAUDE_FABLE_5_1_MESSAGES_STREAM, CLAUDE_FABLE_5_1_STREAM, CLAUDE_FABLE_5_1_THINKING, CLAUDE_FABLE_5_1_TOOL_CALL, + CLAUDE_HAIKU_5_5_MESSAGES_INPUT_TEXT, + CLAUDE_HAIKU_5_5_MESSAGES_STREAM, CLAUDE_OPUS_5_5_CACHE_READ, CLAUDE_OPUS_5_5_CACHE_WRITE_1H, CLAUDE_OPUS_5_5_CACHE_WRITE_5M, CLAUDE_OPUS_5_5_INPUT_TEXT, + CLAUDE_OPUS_5_5_MESSAGES_INPUT_TEXT, + CLAUDE_OPUS_5_5_MESSAGES_STREAM, CLAUDE_OPUS_5_5_STREAM, CLAUDE_OPUS_5_5_THINKING, CLAUDE_OPUS_5_5_TOOL_CALL, @@ -25,6 +32,8 @@ from integration.cost_calculation.cases.anthropic import ( CLAUDE_SONNET_5_5_CACHE_WRITE_1H, CLAUDE_SONNET_5_5_CACHE_WRITE_5M, CLAUDE_SONNET_5_5_INPUT_TEXT, + CLAUDE_SONNET_5_5_MESSAGES_INPUT_TEXT, + CLAUDE_SONNET_5_5_MESSAGES_STREAM, CLAUDE_SONNET_5_5_STREAM, CLAUDE_SONNET_5_5_THINKING, CLAUDE_SONNET_5_5_TOOL_CALL, @@ -41,26 +50,72 @@ from integration.cost_calculation.cost_tracking_case import ( from integration.cost_calculation.stream_parity.case import StreamParityTestCase, sse_frames -def _gemini_case_response(response: JsonResponse | SseResponse, model_version: str) -> JsonResponse | SseResponse: +def _gemini_case_response( + response: JsonResponse | SseResponse, + model_version: str, + usage_fields: dict[str, JsonValue] | None = None, +) -> JsonResponse | SseResponse: match response: case JsonResponse(): - return response.model_copy(update={"body": {**response.body, "modelVersion": model_version}}) + usage: Final = response.body.get("usageMetadata") + body: Final = {**response.body, "modelVersion": model_version} + return response.model_copy( + update={ + "body": { + **body, + **( + { + "usageMetadata": { + **cast(dict[str, JsonValue], usage), + **usage_fields, + } + } + if usage_fields is not None and isinstance(usage, dict) + else {} + ), + } + } + ) case SseResponse(): return response.model_copy( update={ - "frames": tuple(frame.replace("gemini-3.8-flash", model_version) for frame in response.frames) + "frames": tuple( + _gemini_case_frame(frame, model_version, usage_fields) for frame in response.frames + ) } ) +def _gemini_case_frame(frame: str, model_version: str, usage_fields: dict[str, JsonValue] | None) -> str: + replaced: Final = frame.replace("gemini-3.8-flash", model_version) + if usage_fields is None: + return replaced + prefix, separator, payload = replaced.partition("data: ") + if not separator or payload == "[DONE]": + return replaced + event: Final = cast(dict[str, JsonValue], json.loads(payload)) + usage: Final = event.get("usageMetadata") + if not isinstance(usage, dict): + return replaced + return f"{prefix}data: {json.dumps({**event, 'usageMetadata': {**usage, **usage_fields}})}" + + def _gemini_case( - sibling: CostTrackingTestCase, name: str, model: str, expected: ExactExpected + sibling: CostTrackingTestCase, + name: str, + model: str, + expected: ExactExpected, + *, + usage_fields: dict[str, JsonValue] | None = None, + model_version: str = "gemini-3.1-pro-preview", ) -> CostTrackingTestCase: return sibling.model_copy( update={ "name": name, "model": model, - "response": _gemini_case_response(cast(JsonResponse | SseResponse, sibling.response), "gemini-3.1-pro-preview"), + "response": _gemini_case_response( + cast(JsonResponse | SseResponse, sibling.response), model_version, usage_fields + ), "expected": expected, } ) @@ -97,6 +152,7 @@ def _vertex_claude_shape_case( expected: ExactExpected, response_content: list[dict[str, JsonValue]] | None = None, adaptive_thinking: bool = False, + thinking_effort: str = "medium", ) -> CostTrackingTestCase: response: Final = cast(JsonResponse, sibling.response) request: Final = ( @@ -107,7 +163,7 @@ def _vertex_claude_shape_case( if key not in {"reasoning_effort", "allowed_openai_params"} }, "thinking": {"type": "adaptive"}, - "output_config": {"effort": "medium"}, + "output_config": {"effort": thinking_effort}, "allowed_openai_params": ["thinking", "output_config"], } if adaptive_thinking @@ -1716,6 +1772,8 @@ GEMINI_3_8_FLASH_INPUT_TEXT: Final = CostTrackingTestCase( "candidatesTokenCount": 412, "totalTokenCount": 2252, "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 1840}], + "candidatesTokensDetails": [{"modality": "TEXT", "tokenCount": 412}], + "trafficType": "ON_DEMAND", }, "modelVersion": "gemini-3.8-flash", }, @@ -1770,6 +1828,8 @@ GEMINI_3_8_FLASH_CACHE_READ: Final = CostTrackingTestCase( "totalTokenCount": 13308, "cachedContentTokenCount": 12288, "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 12928}], + "candidatesTokensDetails": [{"modality": "TEXT", "tokenCount": 380}], + "trafficType": "ON_DEMAND", }, "modelVersion": "gemini-3.8-flash", }, @@ -2117,6 +2177,8 @@ GEMINI_3_8_FLASH_STREAM: Final = CostTrackingTestCase( "candidatesTokenCount": 412, "totalTokenCount": 2252, "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 1840}], + "candidatesTokensDetails": [{"modality": "TEXT", "tokenCount": 412}], + "trafficType": "ON_DEMAND", }, "modelVersion": "gemini-3.8-flash", }, @@ -3122,28 +3184,258 @@ VERTEX_AI_CLAUDE_FABLE_5_1_TOOL_CALL: Final = _vertex_claude_shape_case( ), ) +VERTEX_AI_CLAUDE_HAIKU_5_5_INPUT_TEXT: Final = _vertex_claude_shape_case( + VERTEX_AI_CLAUDE_SONNET_5_5_INPUT_TEXT, + "vertex_ai/claude-haiku-5-5-input_text", + "vertex_ai/claude-haiku-5-5", + _vertex_claude_usage(input_tokens=1840, output_tokens=412), + ExactExpected(spend=0.00507, input_cost=0.002392, output_cost=0.002678, prompt_tokens=1840, completion_tokens=412), +) +VERTEX_AI_CLAUDE_HAIKU_5_5_STREAM: Final = VERTEX_AI_CLAUDE_SONNET_5_5_STREAM.model_copy( + update={ + "name": "vertex_ai/claude-haiku-5-5-stream", + "model": "vertex_ai/claude-haiku-5-5", + "response": _vertex_claude_stream_response("claude-haiku-5-5", "end_turn", ("scripted answer",)), + "expected": VERTEX_AI_CLAUDE_HAIKU_5_5_INPUT_TEXT.expected, + } +) +VERTEX_AI_CLAUDE_HAIKU_5_5_CACHE_READ: Final = _vertex_claude_shape_case( + VERTEX_AI_CLAUDE_SONNET_5_5_CACHE_READ, + "vertex_ai/claude-haiku-5-5-cache_read", + "vertex_ai/claude-haiku-5-5", + _vertex_claude_usage(input_tokens=14, output_tokens=6, cache_read_input_tokens=27003), + ExactExpected( + spend=0.00356759, + input_cost=0.00352859, + output_cost=0.000039, + prompt_tokens=27017, + completion_tokens=6, + cache_read_cost=0.00351039, + ), +) +VERTEX_AI_CLAUDE_HAIKU_5_5_CACHE_WRITE_5M: Final = _vertex_claude_shape_case( + VERTEX_AI_CLAUDE_SONNET_5_5_CACHE_WRITE_5M, + "vertex_ai/claude-haiku-5-5-cache_write_5m", + "vertex_ai/claude-haiku-5-5", + _vertex_claude_usage( + input_tokens=14, output_tokens=6, cache_creation_input_tokens=27011, ephemeral_5m_input_tokens=27011 + ), + ExactExpected( + spend=0.043950075, + input_cost=0.043911075, + output_cost=0.000039, + prompt_tokens=27025, + completion_tokens=6, + cache_creation_cost=0.043892875, + ), +) +VERTEX_AI_CLAUDE_HAIKU_5_5_CACHE_WRITE_1H: Final = _vertex_claude_shape_case( + VERTEX_AI_CLAUDE_SONNET_5_5_CACHE_WRITE_1H, + "vertex_ai/claude-haiku-5-5-cache_write_1h", + "vertex_ai/claude-haiku-5-5", + _vertex_claude_usage( + input_tokens=14, output_tokens=6, cache_creation_input_tokens=27011, ephemeral_1h_input_tokens=27011 + ), + ExactExpected( + spend=0.0702858, + input_cost=0.0702468, + output_cost=0.000039, + prompt_tokens=27025, + completion_tokens=6, + cache_creation_cost=0.0702286, + ), +) +VERTEX_AI_CLAUDE_HAIKU_5_5_THINKING: Final = _vertex_claude_shape_case( + VERTEX_AI_CLAUDE_SONNET_5_5_THINKING, + "vertex_ai/claude-haiku-5-5-thinking", + "vertex_ai/claude-haiku-5-5", + _vertex_claude_usage(input_tokens=18, output_tokens=144, thinking_tokens=37), + ExactExpected( + spend=0.0009594, input_cost=0.0000234, output_cost=0.000936, prompt_tokens=18, completion_tokens=144 + ), + adaptive_thinking=True, + thinking_effort="low", +) +VERTEX_AI_CLAUDE_HAIKU_5_5_TOOL_CALL: Final = _vertex_claude_shape_case( + VERTEX_AI_CLAUDE_SONNET_5_5_TOOL_CALL, + "vertex_ai/claude-haiku-5-5-tool_call", + "vertex_ai/claude-haiku-5-5", + _vertex_claude_usage(input_tokens=385, output_tokens=50), + ExactExpected( + spend=0.0008255, input_cost=0.0005005, output_cost=0.000325, prompt_tokens=385, completion_tokens=50 + ), +) +VERTEX_AI_CLAUDE_HAIKU_5_5_MESSAGES_INPUT_TEXT: Final = CLAUDE_HAIKU_5_5_MESSAGES_INPUT_TEXT.model_copy( + update={ + "name": "vertex_ai/claude-haiku-5-5-messages_input_text", + "model": "vertex_ai/claude-haiku-5-5", + "endpoint": "/v1/messages", + "response": VERTEX_AI_CLAUDE_HAIKU_5_5_INPUT_TEXT.response, + "expected": VERTEX_AI_CLAUDE_HAIKU_5_5_INPUT_TEXT.expected, + } +) +VERTEX_AI_CLAUDE_HAIKU_5_5_MESSAGES_STREAM: Final = CLAUDE_HAIKU_5_5_MESSAGES_STREAM.model_copy( + update={ + "name": "vertex_ai/claude-haiku-5-5-messages_stream", + "model": "vertex_ai/claude-haiku-5-5", + "endpoint": "/v1/messages", + "response": VERTEX_AI_CLAUDE_HAIKU_5_5_STREAM.response, + "expected": VERTEX_AI_CLAUDE_HAIKU_5_5_INPUT_TEXT.expected, + } +) + +VERTEX_AI_CLAUDE_FABLE_5_1_MESSAGES_INPUT_TEXT: Final = CLAUDE_FABLE_5_1_MESSAGES_INPUT_TEXT.model_copy( + update={ + "name": "vertex_ai/claude-fable-5-1-messages_input_text", + "model": "vertex_ai/claude-fable-5-1", + "endpoint": "/v1/messages", + "response": VERTEX_AI_CLAUDE_FABLE_5_1_INPUT_TEXT.response, + "expected": VERTEX_AI_CLAUDE_FABLE_5_1_INPUT_TEXT.expected, + } +) +VERTEX_AI_CLAUDE_FABLE_5_1_MESSAGES_STREAM: Final = CLAUDE_FABLE_5_1_MESSAGES_STREAM.model_copy( + update={ + "name": "vertex_ai/claude-fable-5-1-messages_stream", + "model": "vertex_ai/claude-fable-5-1", + "endpoint": "/v1/messages", + "response": VERTEX_AI_CLAUDE_FABLE_5_1_STREAM.response, + "expected": VERTEX_AI_CLAUDE_FABLE_5_1_INPUT_TEXT.expected, + } +) +VERTEX_AI_CLAUDE_FABLE_5_1_BRIDGE_PAIRS: Final = bridge_pairs( + (VERTEX_AI_CLAUDE_FABLE_5_1_MESSAGES_INPUT_TEXT, VERTEX_AI_CLAUDE_FABLE_5_1_MESSAGES_STREAM), + ("/v1/responses",), +) +VERTEX_AI_CLAUDE_OPUS_5_5_MESSAGES_INPUT_TEXT: Final = CLAUDE_OPUS_5_5_MESSAGES_INPUT_TEXT.model_copy( + update={ + "name": "vertex_ai/claude-opus-5-5-messages_input_text", + "model": "vertex_ai/claude-opus-5-5", + "endpoint": "/v1/messages", + "response": VERTEX_AI_CLAUDE_OPUS_5_5_INPUT_TEXT.response, + "expected": VERTEX_AI_CLAUDE_OPUS_5_5_INPUT_TEXT.expected, + } +) +VERTEX_AI_CLAUDE_OPUS_5_5_MESSAGES_STREAM: Final = CLAUDE_OPUS_5_5_MESSAGES_STREAM.model_copy( + update={ + "name": "vertex_ai/claude-opus-5-5-messages_stream", + "model": "vertex_ai/claude-opus-5-5", + "endpoint": "/v1/messages", + "response": VERTEX_AI_CLAUDE_OPUS_5_5_STREAM.response, + "expected": VERTEX_AI_CLAUDE_OPUS_5_5_INPUT_TEXT.expected, + } +) +VERTEX_AI_CLAUDE_OPUS_5_5_BRIDGE_PAIRS: Final = bridge_pairs( + (VERTEX_AI_CLAUDE_OPUS_5_5_MESSAGES_INPUT_TEXT, VERTEX_AI_CLAUDE_OPUS_5_5_MESSAGES_STREAM), + ("/v1/responses",), +) +VERTEX_AI_CLAUDE_SONNET_5_5_MESSAGES_INPUT_TEXT: Final = CLAUDE_SONNET_5_5_MESSAGES_INPUT_TEXT.model_copy( + update={ + "name": "vertex_ai/claude-sonnet-5-5-messages_input_text", + "model": "vertex_ai/claude-sonnet-5-5", + "endpoint": "/v1/messages", + "response": VERTEX_AI_CLAUDE_SONNET_5_5_INPUT_TEXT.response, + "expected": VERTEX_AI_CLAUDE_SONNET_5_5_INPUT_TEXT.expected, + } +) +VERTEX_AI_CLAUDE_SONNET_5_5_MESSAGES_STREAM: Final = CLAUDE_SONNET_5_5_MESSAGES_STREAM.model_copy( + update={ + "name": "vertex_ai/claude-sonnet-5-5-messages_stream", + "model": "vertex_ai/claude-sonnet-5-5", + "endpoint": "/v1/messages", + "response": VERTEX_AI_CLAUDE_SONNET_5_5_STREAM.response, + "expected": VERTEX_AI_CLAUDE_SONNET_5_5_INPUT_TEXT.expected, + } +) +VERTEX_AI_CLAUDE_SONNET_5_5_BRIDGE_PAIRS: Final = bridge_pairs( + (VERTEX_AI_CLAUDE_SONNET_5_5_MESSAGES_INPUT_TEXT, VERTEX_AI_CLAUDE_SONNET_5_5_MESSAGES_STREAM), + ("/v1/responses",), +) +VERTEX_AI_CLAUDE_HAIKU_5_5_BRIDGE_PAIRS: Final = bridge_pairs( + (VERTEX_AI_CLAUDE_HAIKU_5_5_MESSAGES_INPUT_TEXT, VERTEX_AI_CLAUDE_HAIKU_5_5_MESSAGES_STREAM), + ("/v1/responses",), +) + GEMINI_3_1_PRO_PREVIEW_INPUT_TEXT: Final = _gemini_case( GEMINI_3_8_FLASH_INPUT_TEXT, "vertex_ai/gemini-3.1-pro-preview-input_text", "vertex_ai/gemini-3.1-pro-preview", - ExactExpected(spend=0.0094864, input_cost=0.004048, output_cost=0.0054384, prompt_tokens=1840, completion_tokens=412) + ExactExpected(spend=0.0094864, input_cost=0.004048, output_cost=0.0054384, prompt_tokens=1840, completion_tokens=412), + usage_fields={ + "candidatesTokensDetails": [{"modality": "TEXT", "tokenCount": 412}], + "trafficType": "ON_DEMAND", + }, ) GEMINI_3_1_PRO_PREVIEW_STREAM: Final = _gemini_case( GEMINI_3_8_FLASH_STREAM, "vertex_ai/gemini-3.1-pro-preview-stream", "vertex_ai/gemini-3.1-pro-preview", - ExactExpected(spend=0.0094864, input_cost=0.004048, output_cost=0.0054384, prompt_tokens=1840, completion_tokens=412) + ExactExpected(spend=0.0094864, input_cost=0.004048, output_cost=0.0054384, prompt_tokens=1840, completion_tokens=412), + usage_fields={ + "candidatesTokensDetails": [{"modality": "TEXT", "tokenCount": 412}], + "trafficType": "ON_DEMAND", + }, ) GEMINI_3_1_PRO_PREVIEW_CACHE_READ: Final = _gemini_case( GEMINI_3_8_FLASH_CACHE_READ, "vertex_ai/gemini-3.1-pro-preview-cache_read", "vertex_ai/gemini-3.1-pro-preview", - ExactExpected(spend=0.00912736, input_cost=0.00411136, output_cost=0.005016, prompt_tokens=12928, completion_tokens=380) + ExactExpected(spend=0.00912736, input_cost=0.00411136, output_cost=0.005016, prompt_tokens=12928, completion_tokens=380), + usage_fields={ + "candidatesTokensDetails": [{"modality": "TEXT", "tokenCount": 380}], + "trafficType": "ON_DEMAND", + }, ) GEMINI_3_1_PRO_PREVIEW_REASONING: Final = _gemini_case( GEMINI_3_1_PRO_REASONING, "vertex_ai/gemini-3.1-pro-preview-reasoning", "vertex_ai/gemini-3.1-pro-preview", - ExactExpected(spend=0.056056, input_cost=0.002728, output_cost=0.053328, prompt_tokens=1240, completion_tokens=4040) + ExactExpected(spend=0.056056, input_cost=0.002728, output_cost=0.053328, prompt_tokens=1240, completion_tokens=4040), + usage_fields={ + "candidatesTokensDetails": [{"modality": "TEXT", "tokenCount": 560}], + "trafficType": "ON_DEMAND", + }, +) +GEMINI_3_8_FLASH_REASONING: Final = _gemini_case( + GEMINI_3_1_PRO_REASONING, + "vertex_ai/gemini-3.8-flash-reasoning", + "vertex_ai/gemini-3.8-flash", + ExactExpected( + spend=0.01144, + input_cost=0.0006448, + output_cost=0.0107952, + prompt_tokens=1240, + completion_tokens=4040, + ), + usage_fields={ + "candidatesTokensDetails": [{"modality": "TEXT", "tokenCount": 560}], + "trafficType": "ON_DEMAND", + }, + model_version="gemini-3.8-flash", ) GEMINI_3_1_PRO_PREVIEW_TOOL_CALL: Final = _gemini_case( GEMINI_3_8_FLASH_TOOL_CALL, "vertex_ai/gemini-3.1-pro-preview-tool_call", "vertex_ai/gemini-3.1-pro-preview", ExactExpected(spend=0.0094864, input_cost=0.004048, output_cost=0.0054384, prompt_tokens=1840, completion_tokens=412) ) +GEMINI_3_1_PRO_PREVIEW_TIERED_INPUT_ABOVE_200K: Final = _gemini_case( + GEMINI_3_1_PRO_TIERED_INPUT_ABOVE_200K, + "vertex_ai/gemini-3.1-pro-preview-tiered_input_above_200k", + "vertex_ai/gemini-3.1-pro-preview", + ExactExpected( + spend=0.913396, input_cost=0.90112, output_cost=0.012276, prompt_tokens=204800, completion_tokens=620 + ), + usage_fields={ + "candidatesTokensDetails": [{"modality": "TEXT", "tokenCount": 620}], + "trafficType": "ON_DEMAND", + }, +) +GEMINI_3_1_PRO_PREVIEW_TIERED_CACHE_READ_ABOVE_200K: Final = _gemini_case( + GEMINI_3_1_PRO_TIERED_CACHE_READ_ABOVE_200K, + "vertex_ai/gemini-3.1-pro-preview-tiered_cache_read_above_200k", + "vertex_ai/gemini-3.1-pro-preview", + ExactExpected( + spend=0.11628672, input_cost=0.10678272, output_cost=0.009504, prompt_tokens=205824, completion_tokens=480 + ), + usage_fields={ + "candidatesTokensDetails": [{"modality": "TEXT", "tokenCount": 480}], + "trafficType": "ON_DEMAND", + "cacheTokensDetails": [{"modality": "TEXT", "tokenCount": 201728}], + }, +) + GEMINI_4_ARGON_INPUT_TEXT: Final = GEMINI_3_8_FLASH_INPUT_TEXT.model_copy( update={ "name": "vertex_ai/gemini-4-argon-input_text", @@ -3301,6 +3593,15 @@ GEMINI_4_ARGON_TOOL_CALL: Final = GEMINI_3_8_FLASH_TOOL_CALL.model_copy( } ) +GEMINI_4_ARGON_ROUTE_BRIDGES: Final = bridge_pairs( + (GEMINI_4_ARGON_INPUT_TEXT, GEMINI_4_ARGON_STREAM), + ("/v1/responses", "/v1/messages"), +) +GEMINI_3_1_PRO_PREVIEW_ROUTE_BRIDGES: Final = bridge_pairs( + (GEMINI_3_1_PRO_PREVIEW_INPUT_TEXT, GEMINI_3_1_PRO_PREVIEW_STREAM), + ("/v1/responses", "/v1/messages"), +) + CASES: Final[tuple[CostTrackingTestCase, ...]] = ( GEMINI_3_1_PRO_INPUT_TEXT, GEMINI_3_1_PRO_CACHE_READ, @@ -3326,6 +3627,7 @@ CASES: Final[tuple[CostTrackingTestCase, ...]] = ( GEMINI_3_1_PRO_STREAM_TOOL_CALL, GEMINI_3_1_PRO_STREAM_FULL_USAGE, GEMINI_3_8_FLASH_INPUT_TEXT, + GEMINI_3_8_FLASH_REASONING, GEMINI_3_8_FLASH_CACHE_READ, GEMINI_3_8_FLASH_SERVICE_TIER_FLEX, GEMINI_3_8_FLASH_SERVICE_TIER_PRIORITY, @@ -3367,9 +3669,30 @@ CASES: Final[tuple[CostTrackingTestCase, ...]] = ( VERTEX_AI_CLAUDE_FABLE_5_1_CACHE_WRITE_1H, VERTEX_AI_CLAUDE_FABLE_5_1_THINKING, VERTEX_AI_CLAUDE_FABLE_5_1_TOOL_CALL, + VERTEX_AI_CLAUDE_HAIKU_5_5_INPUT_TEXT, + VERTEX_AI_CLAUDE_HAIKU_5_5_STREAM, + VERTEX_AI_CLAUDE_HAIKU_5_5_CACHE_READ, + VERTEX_AI_CLAUDE_HAIKU_5_5_CACHE_WRITE_5M, + VERTEX_AI_CLAUDE_HAIKU_5_5_CACHE_WRITE_1H, + VERTEX_AI_CLAUDE_HAIKU_5_5_THINKING, + VERTEX_AI_CLAUDE_HAIKU_5_5_TOOL_CALL, + VERTEX_AI_CLAUDE_HAIKU_5_5_MESSAGES_INPUT_TEXT, + VERTEX_AI_CLAUDE_HAIKU_5_5_MESSAGES_STREAM, + *(pair.bridged for pair in VERTEX_AI_CLAUDE_HAIKU_5_5_BRIDGE_PAIRS), + VERTEX_AI_CLAUDE_FABLE_5_1_MESSAGES_INPUT_TEXT, + VERTEX_AI_CLAUDE_FABLE_5_1_MESSAGES_STREAM, + *(pair.bridged for pair in VERTEX_AI_CLAUDE_FABLE_5_1_BRIDGE_PAIRS), + VERTEX_AI_CLAUDE_OPUS_5_5_MESSAGES_INPUT_TEXT, + VERTEX_AI_CLAUDE_OPUS_5_5_MESSAGES_STREAM, + *(pair.bridged for pair in VERTEX_AI_CLAUDE_OPUS_5_5_BRIDGE_PAIRS), + VERTEX_AI_CLAUDE_SONNET_5_5_MESSAGES_INPUT_TEXT, + VERTEX_AI_CLAUDE_SONNET_5_5_MESSAGES_STREAM, + *(pair.bridged for pair in VERTEX_AI_CLAUDE_SONNET_5_5_BRIDGE_PAIRS), GEMINI_3_1_PRO_PREVIEW_INPUT_TEXT, GEMINI_3_1_PRO_PREVIEW_STREAM, GEMINI_3_1_PRO_PREVIEW_CACHE_READ, + GEMINI_3_1_PRO_PREVIEW_TIERED_INPUT_ABOVE_200K, + GEMINI_3_1_PRO_PREVIEW_TIERED_CACHE_READ_ABOVE_200K, GEMINI_3_1_PRO_PREVIEW_REASONING, GEMINI_3_1_PRO_PREVIEW_TOOL_CALL, GEMINI_4_ARGON_INPUT_TEXT, @@ -3377,6 +3700,8 @@ CASES: Final[tuple[CostTrackingTestCase, ...]] = ( GEMINI_4_ARGON_CACHE_READ, GEMINI_4_ARGON_REASONING, GEMINI_4_ARGON_TOOL_CALL, + *(pair.bridged for pair in GEMINI_4_ARGON_ROUTE_BRIDGES), + *(pair.bridged for pair in GEMINI_3_1_PRO_PREVIEW_ROUTE_BRIDGES), ) PARITY: Final[tuple[StreamParityTestCase, ...]] = ( @@ -3397,8 +3722,59 @@ PARITY: Final[tuple[StreamParityTestCase, ...]] = ( ), StreamParityTestCase(plain=VERTEX_AI_CLAUDE_OPUS_5_5_INPUT_TEXT, streamed=VERTEX_AI_CLAUDE_OPUS_5_5_STREAM), StreamParityTestCase(plain=VERTEX_AI_CLAUDE_FABLE_5_1_INPUT_TEXT, streamed=VERTEX_AI_CLAUDE_FABLE_5_1_STREAM), + StreamParityTestCase( + plain=VERTEX_AI_CLAUDE_HAIKU_5_5_INPUT_TEXT, streamed=VERTEX_AI_CLAUDE_HAIKU_5_5_STREAM + ), + StreamParityTestCase( + plain=VERTEX_AI_CLAUDE_HAIKU_5_5_MESSAGES_INPUT_TEXT, + streamed=VERTEX_AI_CLAUDE_HAIKU_5_5_MESSAGES_STREAM, + ), + StreamParityTestCase( + plain=VERTEX_AI_CLAUDE_HAIKU_5_5_BRIDGE_PAIRS[0].bridged, + streamed=VERTEX_AI_CLAUDE_HAIKU_5_5_BRIDGE_PAIRS[1].bridged, + ), + StreamParityTestCase( + plain=VERTEX_AI_CLAUDE_FABLE_5_1_MESSAGES_INPUT_TEXT, + streamed=VERTEX_AI_CLAUDE_FABLE_5_1_MESSAGES_STREAM, + ), + StreamParityTestCase( + plain=VERTEX_AI_CLAUDE_FABLE_5_1_BRIDGE_PAIRS[0].bridged, + streamed=VERTEX_AI_CLAUDE_FABLE_5_1_BRIDGE_PAIRS[1].bridged, + ), + StreamParityTestCase( + plain=VERTEX_AI_CLAUDE_OPUS_5_5_MESSAGES_INPUT_TEXT, + streamed=VERTEX_AI_CLAUDE_OPUS_5_5_MESSAGES_STREAM, + ), + StreamParityTestCase( + plain=VERTEX_AI_CLAUDE_OPUS_5_5_BRIDGE_PAIRS[0].bridged, + streamed=VERTEX_AI_CLAUDE_OPUS_5_5_BRIDGE_PAIRS[1].bridged, + ), + StreamParityTestCase( + plain=VERTEX_AI_CLAUDE_SONNET_5_5_MESSAGES_INPUT_TEXT, + streamed=VERTEX_AI_CLAUDE_SONNET_5_5_MESSAGES_STREAM, + ), + StreamParityTestCase( + plain=VERTEX_AI_CLAUDE_SONNET_5_5_BRIDGE_PAIRS[0].bridged, + streamed=VERTEX_AI_CLAUDE_SONNET_5_5_BRIDGE_PAIRS[1].bridged, + ), StreamParityTestCase(plain=GEMINI_3_1_PRO_PREVIEW_INPUT_TEXT, streamed=GEMINI_3_1_PRO_PREVIEW_STREAM), StreamParityTestCase(plain=GEMINI_4_ARGON_INPUT_TEXT, streamed=GEMINI_4_ARGON_STREAM), + StreamParityTestCase( + plain=GEMINI_4_ARGON_ROUTE_BRIDGES[0].bridged, + streamed=GEMINI_4_ARGON_ROUTE_BRIDGES[2].bridged, + ), + StreamParityTestCase( + plain=GEMINI_4_ARGON_ROUTE_BRIDGES[1].bridged, + streamed=GEMINI_4_ARGON_ROUTE_BRIDGES[3].bridged, + ), + StreamParityTestCase( + plain=GEMINI_3_1_PRO_PREVIEW_ROUTE_BRIDGES[0].bridged, + streamed=GEMINI_3_1_PRO_PREVIEW_ROUTE_BRIDGES[2].bridged, + ), + StreamParityTestCase( + plain=GEMINI_3_1_PRO_PREVIEW_ROUTE_BRIDGES[1].bridged, + streamed=GEMINI_3_1_PRO_PREVIEW_ROUTE_BRIDGES[3].bridged, + ), ) BRIDGE_COST_PARITY: Final[tuple[BridgeCostParityTestCase, ...]] = ( @@ -3406,6 +3782,10 @@ BRIDGE_COST_PARITY: Final[tuple[BridgeCostParityTestCase, ...]] = ( (GEMINI_3_8_FLASH_INPUT_TEXT, GEMINI_3_8_FLASH_CACHE_READ, GEMINI_3_8_FLASH_STREAM), ("/v1/responses", "/v1/messages"), ), - *bridge_pairs((GEMINI_3_1_PRO_PREVIEW_INPUT_TEXT,), ("/v1/responses",)), - *bridge_pairs((GEMINI_4_ARGON_INPUT_TEXT,), ("/v1/responses",)), + *GEMINI_3_1_PRO_PREVIEW_ROUTE_BRIDGES, + *GEMINI_4_ARGON_ROUTE_BRIDGES, + *VERTEX_AI_CLAUDE_HAIKU_5_5_BRIDGE_PAIRS, + *VERTEX_AI_CLAUDE_FABLE_5_1_BRIDGE_PAIRS, + *VERTEX_AI_CLAUDE_OPUS_5_5_BRIDGE_PAIRS, + *VERTEX_AI_CLAUDE_SONNET_5_5_BRIDGE_PAIRS, ) diff --git a/tests/integration/cost_calculation/cost_map.py b/tests/integration/cost_calculation/cost_map.py index 2ef42249a7e..ba7be7647da 100644 --- a/tests/integration/cost_calculation/cost_map.py +++ b/tests/integration/cost_calculation/cost_map.py @@ -889,6 +889,26 @@ COST_MAP_ENTRIES: Final[dict[str, dict[str, JsonValue]]] = { "max_output_tokens": 128000, "max_tokens": 128000, }, + "vertex_ai/claude-haiku-5-5": { + "litellm_provider": "vertex_ai-anthropic_models", + "mode": "chat", + "input_cost_per_token": 1.3e-06, + "input_cost_per_token_above_100k_tokens": 6.5e-06, + "output_cost_per_token": 6.5e-06, + "output_cost_per_token_above_100k_tokens": 3.25e-05, + "cache_creation_input_token_cost": 1.625e-06, + "cache_creation_input_token_cost_above_1hr": 2.6e-06, + "cache_creation_input_token_cost_above_100k_tokens": 8.125e-06, + "cache_creation_input_token_cost_above_1hr_above_100k_tokens": 1.3e-05, + "cache_read_input_token_cost": 1.3e-07, + "cache_read_input_token_cost_above_100k_tokens": 6.5e-07, + "supports_function_calling": True, + "supports_prompt_caching": True, + "supports_reasoning": True, + "max_input_tokens": 1000000, + "max_output_tokens": 128000, + "max_tokens": 128000, + }, "groq/qwen/qwen3.8-27b": { "litellm_provider": "groq", "mode": "chat", diff --git a/tests/integration/cost_calculation/test_responses_websocket_guardrail_spend.py b/tests/integration/cost_calculation/test_responses_websocket_guardrail_spend.py new file mode 100644 index 00000000000..8af70d39186 --- /dev/null +++ b/tests/integration/cost_calculation/test_responses_websocket_guardrail_spend.py @@ -0,0 +1,168 @@ +"""A native Responses WebSocket runs the proxy's pre_call guardrails on every ``response.create`` frame, bills each +turn's spend row with the guardrail cost of that turn alone, and bills a blocked frame's guardrail scan as a failed +turn (LIT-9744)""" + +import asyncio +import hashlib +import json +import uuid +from collections.abc import Iterator +from dataclasses import dataclass +from pathlib import Path +from typing import Final + +import pytest +import websockets +import yaml +from integration._support.client import Gateway, eventually, gateway_from_environment +from integration._support.database import read_rows +from integration._support.process import graceful_stop_seconds, owned_proxy +from integration._support.responses_ws_peer import PROVIDER_MODEL, USAGE, create, responses_peer, turn +from integration._support.tls import write_self_signed_cert +from pydantic import JsonValue + +pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120) + +_INPUT_COST_PER_TOKEN: Final = 0.002 +_OUTPUT_COST_PER_TOKEN: Final = 0.004 +_COST_PER_WORD: Final = 0.1 +_MARKER: Final = " [guarded]" +_BLOCKED_WORD: Final = "forbidden" +_GUARDRAIL_MODULE: Final = f""" +from litellm.integrations.custom_guardrail import CustomGuardrail + + +class TurnGuardrail(CustomGuardrail): + async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): + blocks = [block for item in data.get("input") or [] for block in item.get("content") or []] + words = sum(len(block["text"].split()) for block in blocks) + blocked = any({_BLOCKED_WORD!r} in block["text"] for block in blocks) + self.add_standard_logging_guardrail_information_to_request_data( + guardrail_json_response={{"words": words}}, + request_data=data, + guardrail_status="guardrail_intervened" if blocked else "success", + tracing_detail={{"guardrail_cost": round({_COST_PER_WORD} * words, 10)}}, + ) + if blocked: + raise ValueError("turn guardrail blocked the request") + for block in blocks: + block["text"] = block["text"] + {_MARKER!r} + return data +""" + + +@dataclass(frozen=True, slots=True) +class _Rig: + gateway: Gateway + cert: tuple[Path, Path] + + +def _config(directory: Path) -> Path: + (directory / "turn_guardrail.py").write_text(_GUARDRAIL_MODULE) + configuration: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + configuration["guardrails"] = [ + { + "guardrail_name": "turn-guardrail", + "litellm_params": {"guardrail": "turn_guardrail.TurnGuardrail", "mode": "pre_call", "default_on": True}, + } + ] + path: Final = directory / "responses-ws-guardrail.yaml" + path.write_text(yaml.safe_dump(configuration)) + return path + + +@pytest.fixture(scope="module") +def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Rig]: + directory: Final = tmp_path_factory.mktemp("responses-ws-guardrail") + cert: Final = write_self_signed_cert(directory) + with ( + gateway_from_environment() as shared, + owned_proxy(shared, directory, {"SSL_CERT_FILE": str(cert[0])}, config=_config(directory)) as owned, + ): + yield _Rig(owned, cert) + + +def _nested(frame: str) -> str: + """The ``{"type": "response.create", "response": {...}}`` wire shape the WS layer accepts next to the flat one.""" + flat: Final = json.loads(frame) + return json.dumps({"type": flat["type"], "response": {key: value for key, value in flat.items() if key != "type"}}) + + +async def _session( + proxy_url: str, key: str, model: str, frames: tuple[str, ...] +) -> tuple[tuple[dict[str, JsonValue], ...], ...]: + proxy: Final = proxy_url.rstrip("/").replace("http://", "ws://") + async with websockets.connect( + f"{proxy}/v1/responses?model={model}", + additional_headers={"Authorization": f"Bearer {key}"}, + open_timeout=10, + ) as connection: + return tuple([await turn(connection, frame) for frame in frames]) + + +def _spend_rows(key: str) -> list[dict[str, JsonValue]]: + return read_rows( + 'SELECT spend, status FROM "LiteLLM_SpendLogs" WHERE api_key = %s ORDER BY "startTime"', + (hashlib.sha256(key.encode()).hexdigest(),), + ) + + +def _key_spend(key: str) -> float: + (row,) = read_rows( + 'SELECT spend FROM "LiteLLM_VerificationToken" WHERE token = %s', (hashlib.sha256(key.encode()).hexdigest(),) + ) + return float(row["spend"]) + + +def _llm_cost() -> float: + return USAGE["input_tokens"] * _INPUT_COST_PER_TOKEN + USAGE["output_tokens"] * _OUTPUT_COST_PER_TOKEN + + +def test_every_turn_runs_the_pre_call_guardrail_and_bills_only_its_own_guardrail_cost(rig: _Rig) -> None: + with responses_peer(rig.cert) as peer, rig.gateway.scenario() as scenario: + model: Final = scenario.model( + model=f"openai/{PROVIDER_MODEL}", + api_base=peer.url, + input_cost_per_token=_INPUT_COST_PER_TOKEN, + output_cost_per_token=_OUTPUT_COST_PER_TOKEN, + ) + key: Final = scenario.key(models=[model]) + tag: Final = uuid.uuid4().hex + texts: Final = ( + f"one {tag}", + f"one two {tag}", + f"one two three {tag}", + f"{_BLOCKED_WORD} {tag}", + f"one {tag}", + f"{_BLOCKED_WORD} again {tag}", + ) + frames: Final = tuple(create(model, text) for text in texts[:-1]) + (_nested(create(model, texts[-1])),) + turns: Final = asyncio.run(_session(str(rig.gateway.client.base_url), key, model, frames)) + + assert [events[-1]["type"] for events in turns] == [ + "response.completed", + "response.completed", + "response.completed", + "error", + "response.completed", + "error", + ], [[event["type"] for event in events] for events in turns] + assert "blocked" in str(turns[3][-1]["error"]), turns[3] + + forwarded: Final = tuple(peer.frames.get_nowait() for _ in range(peer.frames.qsize())) + assert [frame["input"][0]["content"][0]["text"] for frame in forwarded] == [ + f"{text}{_MARKER}" for text in texts if _BLOCKED_WORD not in text + ], forwarded + + rows: Final = eventually(lambda: _spend_rows(key), lambda values: len(values) == 6, seconds=70) + assert [row["status"] for row in rows] == ["success", "success", "success", "failure", "success", "failure"], ( + rows + ) + expected: Final = [ + (0.0 if _BLOCKED_WORD in text else _llm_cost()) + _COST_PER_WORD * len(text.split()) for text in texts + ] + assert [float(row["spend"]) for row in rows] == pytest.approx(expected), rows + total: Final = sum(expected) + assert eventually( + lambda: _key_spend(key), lambda spend: abs(spend - total) < 1e-6, seconds=70 + ) == pytest.approx(total) diff --git a/tests/integration/management/test_scim_admin_group.py b/tests/integration/management/test_scim_admin_group.py index 1c0cf42ad13..b7bf3cfe968 100644 --- a/tests/integration/management/test_scim_admin_group.py +++ b/tests/integration/management/test_scim_admin_group.py @@ -661,10 +661,6 @@ def test_scim_admin_group_put_members_promotes_and_demotes(gateway: Gateway, tmp @pytest.mark.timeout(300) def test_scim_group_put_returns_the_updated_members(gateway: Gateway, tmp_path: Path) -> None: - pytest.skip( - "BUG: PUT /scim/v2/Groups/{id} returns the members list from before the update, so the 200 body omits " - "the members it just added" - ) group_name: Final = f"scim-admin-put-response-{uuid.uuid4().hex}" with _scim_admin_group_fixture(gateway, tmp_path, group_name) as fixture: expected_member: Final = ScimGroupMember( diff --git a/tests/integration/providers/test_rust_anthropic_safeguards_beta_wire.py b/tests/integration/providers/test_rust_anthropic_safeguards_beta_wire.py index 6ed7e34a583..75c5efc3aa0 100644 --- a/tests/integration/providers/test_rust_anthropic_safeguards_beta_wire.py +++ b/tests/integration/providers/test_rust_anthropic_safeguards_beta_wire.py @@ -334,7 +334,6 @@ def test_deepseek_messages_forward_safeguards_without_adding_a_beta(rig: _Rig, b def test_python_route_adds_the_beta_for_safeguards(gateway: Gateway, scenario: Scenario) -> None: - pytest.skip("BUG: Python /v1/messages sends safeguards without the dangerous-tool-use beta, Anthropic answers 400") with wire_server(anthropic_peer) as wire: deployment: Final = scenario.model(model=f"anthropic/{BACKEND}", api_base=wire.url, api_key=ANTHROPIC_KEY) marker: Final = new_marker() diff --git a/tests/integration/proxy_config.yaml b/tests/integration/proxy_config.yaml index 37bcc4683dd..4e2cab56d27 100644 --- a/tests/integration/proxy_config.yaml +++ b/tests/integration/proxy_config.yaml @@ -473,6 +473,13 @@ model_list: vertex_project: synthetic-vertex-project vertex_location: global vertex_credentials: '{"type": "external_account", "audience": "synthetic-vertex-audience", "subject_token_type": "urn:ietf:params:oauth:token-type:jwt", "token_url": "http://127.0.0.1:8190/_oauth/token", "credential_source": {"url": "http://127.0.0.1:8190/health"}}' + - model_name: vertex_ai/claude-haiku-5-5@default + litellm_params: + model: vertex_ai/claude-haiku-5-5@default + api_base: http://127.0.0.1:8191 + vertex_project: synthetic-vertex-project + vertex_location: global + vertex_credentials: '{"type": "external_account", "audience": "synthetic-vertex-audience", "subject_token_type": "urn:ietf:params:oauth:token-type:jwt", "token_url": "http://127.0.0.1:8190/_oauth/token", "credential_source": {"url": "http://127.0.0.1:8190/health"}}' - model_name: vertex_ai/claude-opus-5-5@default litellm_params: model: vertex_ai/claude-opus-5-5@default diff --git a/tests/integration/translation/chat_completions/bases/vertex_ai.py b/tests/integration/translation/chat_completions/bases/vertex_ai.py index d1022da78c7..447566b0ee3 100644 --- a/tests/integration/translation/chat_completions/bases/vertex_ai.py +++ b/tests/integration/translation/chat_completions/bases/vertex_ai.py @@ -577,6 +577,154 @@ CLAUDE_SONNET_5_5_TEST_CASE: Final = TranslationTestCase( ) +CLAUDE_HAIKU_5_5_THINKING_SIGNATURE: Final = ( + "CAQS8gcKEggSEAIYAjgBQgh0aGlua2luZxIMKxM+8GdyGfxkJrU6Ggw3AMSc5G/SaulUVJkiMJgPNRytmwmfF4bj7WugkKlRxBr+SRwcL9iQ" + "+S5uvmIi/D8thAXX+oVg67diUo3rxiqNB0KrfSB/DXD50OomDcy90ai3a+RMcwIy9wEMzOdOIFqQ24tUMm1YorcQLQpresEnXA9ItiMAzP+P" + "8NMJ1br7O3JHYLIi/12UzVr36AKPrMJQk86VBn5fVxo0xcphyl071Y30m3wO0INGtU+d37IXSrYCJzQdnUSKxekBK28fxsA4tmxx1BE/1Hfl" + "UTsq36+CE2MatE0ymL2pQA+povoJ1iVlB/PBZvkOvpRcQubSwm8RqJrBfmse8baU2UJot0VNNC7OprQ3wYx1j3BGmXOQCGkvCn6jd2lPLjLS" + "F1CjbWQs7X18gB69y2sUosm3aI+MnQNHkUqe8JtIaii42uu3Xe/JheeRTRrcsuAdMZJ76ID9wRZnf3/QCHaMyrYwHXf5SsoUjX3Ofs5MuQT0" + "ISf4T7q+QGf51Bxb4jdfF879h061ZSejzknCHI3fV3SwuldjIIV+FD50d63iLqM/D8GOjaGfDrYugkkdBjCNEQkTpn9NQmqalsrJWOdQQrhc" + "L8TDZG4HBasd/kGr7EnlEbLKrGaWbAWYo4vssoUE7Q57wPumuMrYh4OKInksQTnB0OV9OIZEiTgb9fyaxJBCUiypuQDGMlAP8nb0TQF1vxjR" + "vsXV/qnN0ovR6/7raePN4NM4nMEXESRQyC22Ml88z+WZhe2VG5CLH8tFNapisH+msyqtGyDa+SEie/e6xcDfI1fLAZwEHM34/YjW+C1E1PPd" + "RsGBE9Ms6JUaPrcITwImVKMLGH1pj+TsS+qK7JTGjH8FKwR1Eds/dJTFDEQaAYcmjXWxDR5HBknx7xv90K2uoKt4KmAgulWCQv/LaEAt27lm" + "aGP+4PQ9tpb/Otsub8spmcXKfNqdjWTDH3aEF5MsJ5eSddn92pFV9JAo4oFcj1fUTujfwumGCYDADjLBmyvQao9S42zpO5idmYrHWHt7h7m6" + "BoU5b4uKVS4w4ELEl8y+Oq1uiQLSB8e0cxLarWfCiiTL6v6aJUDEvUTQU+O3XiOxNMCkc00J7e1TnNnnoFqs0sqhV0mMT346vLHocBMVn2iE" + "ZSwFb9OWlHO1e9Y7ccd8XL7KjXtoZeuBTrdYlxrf3g6QHVZg3OOM3OVSOWawFU0tBZUUgH6rYH4oHh8JqwGmbTI1mtAtUX89KrHzGWgUs619" + "SPbFptVh8Vfb3zr0D0OzDqjzv1KCNW/tZcJnCaL9Zl6tAt2v5FiBTOjIbBgB" +) + +CLAUDE_HAIKU_5_5_TEST_CASE: Final = TranslationTestCase( + scenario="basic", + litellm_endpoint="/v1/chat/completions", + litellm_request={**CHAT_COMPLETIONS_REQUEST, "model": "vertex_ai/claude-haiku-5-5@default"}, + expected_provider_endpoint="/v1/projects/synthetic-vertex-project/locations/global/publishers/anthropic/models/claude-haiku-5-5@default:rawPredict", + expected_provider_headers={ + "authorization": "Bearer scripted-token", + "anthropic-version": "2023-06-01", + "content-type": "application/json", + "x-api-key": "scripted-token", + }, + expected_provider_request=VERTEX_AI_ANTHROPIC_MESSAGES_REQUEST, + mock_provider_response={ + **VERTEX_AI_ANTHROPIC_MESSAGES_RESPONSE, + "model": "claude-haiku-5-5", + "id": "msg_vrtx_011Cg8AfAmbYAQCTSseeQSZJ", + "content": [ + { + "type": "thinking", + "thinking": ( + "The system prompt mentions a chat_completions client for this review, but that's not actually in my " + "available toolset — I only have search_files and read_file. I'll note that discrepancy and proceed by " + "explaining which file I need, then calling read_file on shop/validation.py, probably starting around the " + "first 60 lines to get context near line 18.\n\n" + ), + "signature": CLAUDE_HAIKU_5_5_THINKING_SIGNATURE, + }, + { + "type": "text", + "text": ( + "I'll read `shop/validation.py` next, starting at the `validate_quantity` definition on line 18 and " + "including the surrounding context. The project instructions mention a chat_completions client, but none " + "is available in my tools, so I'm using the read-only file tools." + ), + }, + { + "type": "tool_use", + "id": "toolu_vrtx_017jjdVAoN8bus5ng5NkzafQ", + "name": "read_file", + "input": {"path": "shop/validation.py", "start_line": 1, "end_line": 60}, + }, + ], + "usage": { + "input_tokens": 4, + "cache_creation_input_tokens": 876, + "cache_read_input_tokens": 0, + "cache_creation": {"ephemeral_5m_input_tokens": 876, "ephemeral_1h_input_tokens": 0}, + "output_tokens": 348, + "output_tokens_details": {"thinking_tokens": 168}, + }, + }, + expected_litellm_response={ + **CHAT_COMPLETIONS_RESPONSE, + "model": "vertex_ai/claude-haiku-5-5@default", + "choices": [ + { + "finish_reason": "tool_calls", + "index": 0, + "message": { + "content": ( + "I'll read `shop/validation.py` next, starting at the `validate_quantity` definition on line 18 and " + "including the surrounding context. The project instructions mention a chat_completions client, but " + "none is available in my tools, so I'm using the read-only file tools." + ), + "role": "assistant", + "tool_calls": [ + { + "index": 2, + "function": { + "arguments": '{"path": "shop/validation.py", "start_line": 1, "end_line": 60}', + "name": "read_file", + }, + "id": "toolu_vrtx_017jjdVAoN8bus5ng5NkzafQ", + "type": "function", + } + ], + "reasoning_content": ( + "The system prompt mentions a chat_completions client for this review, but that's not actually in my " + "available toolset — I only have search_files and read_file. I'll note that discrepancy and proceed " + "by explaining which file I need, then calling read_file on shop/validation.py, probably starting " + "around the first 60 lines to get context near line 18.\n\n" + ), + "thinking_blocks": [ + { + "type": "thinking", + "thinking": ( + "The system prompt mentions a chat_completions client for this review, but that's not actually " + "in my available toolset — I only have search_files and read_file. I'll note that discrepancy " + "and proceed by explaining which file I need, then calling read_file on shop/validation.py, " + "probably starting around the first 60 lines to get context near line 18.\n\n" + ), + "signature": CLAUDE_HAIKU_5_5_THINKING_SIGNATURE, + } + ], + "provider_specific_fields": { + "citations": None, + "thinking_blocks": [ + { + "type": "thinking", + "thinking": ( + "The system prompt mentions a chat_completions client for this review, but that's not " + "actually in my available toolset — I only have search_files and read_file. I'll note " + "that discrepancy and proceed by explaining which file I need, then calling read_file on " + "shop/validation.py, probably starting around the first 60 lines to get context near " + "line 18.\n\n" + ), + "signature": CLAUDE_HAIKU_5_5_THINKING_SIGNATURE, + } + ], + }, + }, + } + ], + "usage": { + "completion_tokens": 348, + "prompt_tokens": 880, + "total_tokens": 1228, + "completion_tokens_details": {"reasoning_tokens": 168, "text_tokens": 180}, + "prompt_tokens_details": { + "cached_tokens": 0, + "text_tokens": 4, + "cache_write_tokens": 876, + "cache_creation_tokens": 876, + "cache_creation_token_details": {"ephemeral_5m_input_tokens": 876, "ephemeral_1h_input_tokens": 0}, + }, + "cache_creation_input_tokens": 876, + "cache_read_input_tokens": 0, + }, + }, +) + + CLAUDE_OPUS_5_5_TEST_CASE: Final = TranslationTestCase( scenario="basic", litellm_endpoint="/v1/chat/completions", diff --git a/tests/integration/translation/chat_completions/basic/test_chat_completions_basic_vertex_ai.py b/tests/integration/translation/chat_completions/basic/test_chat_completions_basic_vertex_ai.py index 71cbd84543e..7b7f5a7876c 100644 --- a/tests/integration/translation/chat_completions/basic/test_chat_completions_basic_vertex_ai.py +++ b/tests/integration/translation/chat_completions/basic/test_chat_completions_basic_vertex_ai.py @@ -5,6 +5,7 @@ from integration.translation.case import TranslationTestCase from integration.translation.chat_completions.bases.vertex_ai import ( CLAUDE_FABLE_5_1_TEST_CASE, CLAUDE_HAIKU_4_5_TEST_CASE, + CLAUDE_HAIKU_5_5_TEST_CASE, CLAUDE_OPUS_4_8_TEST_CASE, CLAUDE_OPUS_5_5_TEST_CASE, CLAUDE_SONNET_4_6_TEST_CASE, @@ -26,6 +27,7 @@ from integration.translation.runner import assert_translation CLAUDE_SONNET_5_TEST_CASE, CLAUDE_OPUS_4_8_TEST_CASE, CLAUDE_SONNET_5_5_TEST_CASE, + CLAUDE_HAIKU_5_5_TEST_CASE, CLAUDE_OPUS_5_5_TEST_CASE, CLAUDE_FABLE_5_1_TEST_CASE, GEMINI_3_5_FLASH_TEST_CASE, diff --git a/tests/integration/translation/messages/bases/vertex_ai.py b/tests/integration/translation/messages/bases/vertex_ai.py index 6f506b39d93..120b764a63c 100644 --- a/tests/integration/translation/messages/bases/vertex_ai.py +++ b/tests/integration/translation/messages/bases/vertex_ai.py @@ -10,6 +10,7 @@ from integration.translation.templates.vertex_ai import ( ) GEMINI_3_5_FLASH_THOUGHT_SIGNATURE: Final = "AY89a1+kriNs9YsPegJQxgjTWsMlxsqPU/UkHDyz4EK+3FRrE52ygqPIHTvkblgKKiO5Fn8R9D3iu//i+YnuuZ+e7W0yxUoFScas89b5rn3vZHsihnUA7VTe3b7beth1IftteoTGZ3F8OSLrSIP8HKn4xTIx35eeNEt5oiD3K5rTQrdklqO277HvTDhu/LuZhZ0pzaWciGQEc61epIBBpC2ygu8V4L8TekUVZnz6qzSzxipdiQStQ8dZou43334hMnxOR6JjJCGxqt/1K+C0vXX6YP6f+Bph5VhPSnvy2Z1uApi3TUxAcsDDyVewIr+ha2G6DPhN1Oi5uZZNhmbvH01Uq5OOquF1uEICq2m/35LhJIbDNM9N7oI1Z6L94P9xdugeFXM2S7wdJfc8sdg09vEzNeA++GK25zRuB4hU7XyaG4XnWjmu2kswi4sf+EyrO/K0P2tb/ZCJu7kJTtWTiD1wgSIl1eqS8ORKKsXDN2hIbFrtsNR50WIvVzd6" +CLAUDE_HAIKU_5_5_MESSAGES_THINKING_SIGNATURE: Final = "CAQS4wYKEggSEAIYAjgBQgh0aGlua2luZxIMgzBnOw4aYVVqPKRPGgyphQhlJ0pt7tgasWIiMOWM4ATud/TWad7qo/Eh9ZSXhK3IcgnwdgqVeXog9D1q/y69qnyrHdN3a8RYUrAgtyr+BYkQvRW/cF0vH0l86vIM95N77usMXdZKRX3I8XW68KpQedYwFoR12sS9q7A24CcN5UMyQljcF76GhmAFnvpk4yEs7pk8b7ZRSXzBsU6wJfMYCoGXC9r3hZN8m4sBHLJGJRv5ekzEVGVvGCYooE1fBZ8ik/j2DCNOyVdDs4/7xzkqRo7dgFsNJI9Nm7+HgwEdvb1rORObJWooeYb8tfofTFVO43qvNdvM2p2d1Dcf6lsQmvE5fgDFc0t75AuvSZUrgDNe2bMXd2USxxL9YtJ1I9ae96hZAnyOqYZlyiuZf5lZtQQSRouVIdZvFUNonOHLy/Uyw/BGF2bO1/5rIFJBDJC8C6BpbyfDQPpal4YkALulDdgeJWI6dmJYUiK+A+3zcVgQh7at5JPbG5n0tQgO2IFkayqLjCf67AaGS6Xq8J97F+FSwppKvS0TvRbTEvIK1iO3HZZRwcRIcFBn5Llvv6R9JFknXEtohUqvfTnZfuPQ104vxqm91G2rFK1R6C1md6Ze3pJWD0wgkgZyQ7PIKsAmBF1qNoniFi5GVRJx6XHSzr4ZAPDDO1aYABIRl5OWilrkQP3xedj5pcAjJqFuto3cXlhQSNUrocOpEV5IIXxSJc5g5VwYAybpMU4FVImzkCJGJaHLWQKRo0k4IcvXLxBPNXZhop2CkEeazZskF4de0jn1dJa3sQ9rT+7fNVEWPvP+lkBqu/lfObBPVh5jGeouDrRcJvyeYKNZ0NQBdznM5rv4rN8HczeYg1WaYuNOMBvkqT5Zg/UhLHhHV7NWBSa0zt9DPevU/S3OBR1E0YIzCtPFTa0Ia4gGTX40hEc+supCA3h4HpehqMqETGe++Ah+U10HRhlb6lVfRhkU0ypdPqtpodgw/OdUOy7f3/sc31xknudQMoEQqv48F4gR34ee18Op+w1MGn0FAzAFE95D2BuJObRqUhZ6N6nWtlWOqJjCPR+0OrwR3JlxWh19Vxps/ASL6HaJKjClVG2ffojvmVnY8aHjEfI8ifbHE6EYAQ==" CLAUDE_HAIKU_4_5_TEST_CASE: Final = TranslationTestCase( scenario="basic", @@ -486,6 +487,106 @@ CLAUDE_SONNET_5_5_TEST_CASE: Final = TranslationTestCase( ) +CLAUDE_HAIKU_5_5_TEST_CASE: Final = TranslationTestCase( + scenario="basic", + litellm_endpoint="/v1/messages", + litellm_request={**MESSAGES_REQUEST, "model": "vertex_ai/claude-haiku-5-5@default"}, + expected_provider_endpoint="/v1/projects/synthetic-vertex-project/locations/global/publishers/anthropic/models/claude-haiku-5-5@default:rawPredict", + expected_provider_headers={"authorization": "Bearer scripted-token", "content-type": "application/json"}, + expected_provider_request={ + **VERTEX_AI_ANTHROPIC_MESSAGES_REQUEST, + "messages": MESSAGES_REQUEST["messages"], + "stream": False, + "system": [ + { + "type": "text", + "text": ( + "You help maintain a synthetic Python project. Read the relevant files before suggesting a fix. " + "Explain which file you need next, then call a read-only tool. Never change files. Use the messages " + "client for this project review. Consider zero, negative and non-integer quantities before choosing " + "the next file to inspect." + ), + }, + ], + "tools": MESSAGES_REQUEST["tools"], + }, + mock_provider_response={ + **VERTEX_AI_ANTHROPIC_MESSAGES_RESPONSE, + "model": "claude-haiku-5-5", + "id": "msg_vrtx_011Cg8BfjWDUi4ZYQxpy6aVM", + "content": [ + { + "type": "thinking", + "thinking": ( + "I found validate_quantity defined around line 18, so I'll read the validation.py file from the start " + "through roughly line 60 to capture the full context, keeping in mind I'll need to check how it handles " + "zero, negative, and non-integer cases.\n\n" + ), + "signature": CLAUDE_HAIKU_5_5_MESSAGES_THINKING_SIGNATURE, + }, + { + "type": "text", + "text": ( + "I'll read `shop/validation.py` first. The search put `validate_quantity` at line 18, so I'll start " + "there and include the lines around it to see how zero, negative, and non-integer values are handled." + ), + }, + { + "type": "tool_use", + "id": "toolu_vrtx_01YSY1kAV4ayCLVdBJkSv4dt", + "name": "read_file", + "input": {"path": "shop/validation.py", "start_line": 1, "end_line": 60}, + }, + ], + "usage": { + "input_tokens": 4, + "cache_creation_input_tokens": 873, + "cache_read_input_tokens": 0, + "cache_creation": {"ephemeral_5m_input_tokens": 873, "ephemeral_1h_input_tokens": 0}, + "output_tokens": 292, + "output_tokens_details": {"thinking_tokens": 129}, + }, + }, + expected_litellm_response={ + **VERTEX_AI_ANTHROPIC_MESSAGES_RESPONSE, + "model": "vertex_ai/claude-haiku-5-5@default", + "id": "msg_vrtx_011Cg8BfjWDUi4ZYQxpy6aVM", + "content": [ + { + "type": "thinking", + "thinking": ( + "I found validate_quantity defined around line 18, so I'll read the validation.py file from the start " + "through roughly line 60 to capture the full context, keeping in mind I'll need to check how it handles " + "zero, negative, and non-integer cases.\n\n" + ), + "signature": CLAUDE_HAIKU_5_5_MESSAGES_THINKING_SIGNATURE, + }, + { + "type": "text", + "text": ( + "I'll read `shop/validation.py` first. The search put `validate_quantity` at line 18, so I'll start " + "there and include the lines around it to see how zero, negative, and non-integer values are handled." + ), + }, + { + "type": "tool_use", + "id": "toolu_vrtx_01YSY1kAV4ayCLVdBJkSv4dt", + "name": "read_file", + "input": {"path": "shop/validation.py", "start_line": 1, "end_line": 60}, + }, + ], + "usage": { + "input_tokens": 4, + "cache_creation_input_tokens": 873, + "cache_read_input_tokens": 0, + "cache_creation": {"ephemeral_5m_input_tokens": 873, "ephemeral_1h_input_tokens": 0}, + "output_tokens": 292, + "output_tokens_details": {"thinking_tokens": 129}, + }, + }, +) + + CLAUDE_OPUS_5_5_TEST_CASE: Final = TranslationTestCase( scenario="basic", litellm_endpoint="/v1/messages", diff --git a/tests/integration/translation/messages/basic/test_messages_basic_vertex_ai.py b/tests/integration/translation/messages/basic/test_messages_basic_vertex_ai.py index a5b01dfd2e4..4f74a648d8e 100644 --- a/tests/integration/translation/messages/basic/test_messages_basic_vertex_ai.py +++ b/tests/integration/translation/messages/basic/test_messages_basic_vertex_ai.py @@ -5,6 +5,7 @@ from integration.translation.case import TranslationTestCase from integration.translation.messages.bases.vertex_ai import ( CLAUDE_FABLE_5_1_TEST_CASE, CLAUDE_HAIKU_4_5_TEST_CASE, + CLAUDE_HAIKU_5_5_TEST_CASE, CLAUDE_OPUS_4_8_TEST_CASE, CLAUDE_OPUS_5_5_TEST_CASE, CLAUDE_SONNET_4_6_TEST_CASE, @@ -26,6 +27,7 @@ from integration.translation.runner import assert_translation CLAUDE_SONNET_5_TEST_CASE, CLAUDE_OPUS_4_8_TEST_CASE, CLAUDE_SONNET_5_5_TEST_CASE, + CLAUDE_HAIKU_5_5_TEST_CASE, CLAUDE_OPUS_5_5_TEST_CASE, CLAUDE_FABLE_5_1_TEST_CASE, GEMINI_3_5_FLASH_TEST_CASE, diff --git a/tests/integration/translation/responses/bases/vertex_ai.py b/tests/integration/translation/responses/bases/vertex_ai.py index 734ebf13613..59637e380ce 100644 --- a/tests/integration/translation/responses/bases/vertex_ai.py +++ b/tests/integration/translation/responses/bases/vertex_ai.py @@ -12,6 +12,7 @@ from integration.translation.templates.vertex_ai import ( ) GEMINI_3_5_FLASH_THOUGHT_SIGNATURE: Final = "AY89a1+kriNs9YsPegJQxgjTWsMlxsqPU/UkHDyz4EK+3FRrE52ygqPIHTvkblgKKiO5Fn8R9D3iu//i+YnuuZ+e7W0yxUoFScas89b5rn3vZHsihnUA7VTe3b7beth1IftteoTGZ3F8OSLrSIP8HKn4xTIx35eeNEt5oiD3K5rTQrdklqO277HvTDhu/LuZhZ0pzaWciGQEc61epIBBpC2ygu8V4L8TekUVZnz6qzSzxipdiQStQ8dZou43334hMnxOR6JjJCGxqt/1K+C0vXX6YP6f+Bph5VhPSnvy2Z1uApi3TUxAcsDDyVewIr+ha2G6DPhN1Oi5uZZNhmbvH01Uq5OOquF1uEICq2m/35LhJIbDNM9N7oI1Z6L94P9xdugeFXM2S7wdJfc8sdg09vEzNeA++GK25zRuB4hU7XyaG4XnWjmu2kswi4sf+EyrO/K0P2tb/ZCJu7kJTtWTiD1wgSIl1eqS8ORKKsXDN2hIbFrtsNR50WIvVzd6" +CLAUDE_HAIKU_5_5_RESPONSES_THINKING_SIGNATURE: Final = "CAQS0gYKEggSEAIYAjgBQgh0aGlua2luZxIMcbnLkNvj2D6NoBl7Ggw+Qgax2/68TAR1maEiMH6XEeKEqR1ewAiuj01fMNuD1KI7pCTXUsWWr6H7nSzP5BiRGXRrDCIxPYu2Xbvx7yrtBbbyId4ReZzx21wZ2+d0vkkpZM82PXY72wQRJOHLR4JNLRjiR2Jn+FrZ8mBo2mofFrlTSuioOVHc1hheiH4InRK57w2X9xaTXa+mwoWW5epPieHD26+NNJirF9vB832fvqDVOd+f8N0ZDUXUut8zl/Jxx+0BMOktMqyfehIOGumQmyv49E3JHVQGrx+Dd3KZjaQXrsRFf1K3AJ7p5XElR+cYis40W7EL3DE9EAtriZuDEt0fTct79fzZIZgHd6Dnua/rGRnjNjaaZKxGjnhfkVHwsZW4HqTNgwWQiD3lt2ncoLfsfd2IWwMEI5uWpN3GtvRkQvfzidmgGmx3dI8g0nBvLPuq0jzVIuuoKIKqzT7A5ae3Gs9eTlSDr2wOGfdyNf6Q2TLw0w7EegI4wYSExKhFkx0EytCGHJcbN0WClXXmHsRtrXnP7sxCMMY7B8TIpAoN7SXyNjJizOGgzoYL7PNUgqp5nVz/C9vTvudhSWbjLITdxFJ5hearL2rMhryfHBKvRGYC9RMy/EjjtINNwYpxdVX29mUE+8KT+41HymtJ1Ewjtoo8V/7XQnQk3cB7X2RFHcR89hAM/pSsB14ofMFoDk+jJitS+sCDZI0Xh7/ySU0o/OI4l+oD/yctZTHx/mIOiWRtnAOrveNT1ms0/2+4tcO0wi3pRrOB2s5ugYc1I3+GcRwILEIM3a4ABd5Dw2nW2fgUhfSTTGURm4Z91ciih8pKUS1ZPatFTOJ1ehKeUCXflkqY2PHnftiNB/8GzlomsQ2NoRw8UXD00oMmBvSvZpPN7khFLxWHtxLjleKkeKCzN7QZkOqvjAnUwy9TacT2ZvSXIxrvaep7dOyrlMEta/1KsYEcWuBxLr87Va4txXnmoV4RqGgaauNcR2wOGyRFJppxKAgYbmMnxams7cv/7IY2RVSgSQf2k4hFQYzRwrHngjapuZzgNS7UCMjvJJs4VdIypVOATsZdPzJToMgWaiEoRZ3zApQDgYhhGAE=" CLAUDE_HAIKU_4_5_TEST_CASE: Final = TranslationTestCase( scenario="basic", @@ -790,6 +791,160 @@ CLAUDE_SONNET_5_5_TEST_CASE: Final = TranslationTestCase( ) +CLAUDE_HAIKU_5_5_TEST_CASE: Final = TranslationTestCase( + scenario="basic", + litellm_endpoint="/v1/responses", + litellm_request={**RESPONSES_REQUEST, "model": "vertex_ai/claude-haiku-5-5@default"}, + expected_provider_endpoint="/v1/projects/synthetic-vertex-project/locations/global/publishers/anthropic/models/claude-haiku-5-5@default:rawPredict", + expected_provider_headers={ + "authorization": "Bearer scripted-token", + "anthropic-version": "2023-06-01", + "content-type": "application/json", + "x-api-key": "scripted-token", + }, + expected_provider_request={ + **VERTEX_AI_ANTHROPIC_MESSAGES_REQUEST, + "system": [ + { + "type": "text", + "text": ( + "You help maintain a synthetic Python project. Read the relevant files before suggesting a fix. " + "Explain which file you need next, then call a read-only tool. Never change files. Use the responses " + "client for this project review. Consider zero, negative and non-integer quantities before choosing " + "the next file to inspect." + ), + }, + ], + }, + mock_provider_response={ + **VERTEX_AI_ANTHROPIC_MESSAGES_RESPONSE, + "model": "claude-haiku-5-5", + "id": "msg_vrtx_011Cg8BhBrWmiX5ZtMxpM9zA", + "content": [ + { + "type": "thinking", + "thinking": ( + "I should read shop/validation.py to check how it handles zero, negative, and non-integer values " + "before moving to the next file—I'll grab the first 60 lines since I don't know the full length yet.\n\n" + ), + "signature": CLAUDE_HAIKU_5_5_RESPONSES_THINKING_SIGNATURE, + }, + { + "type": "text", + "text": ( + "I'll read the validation function in `shop/validation.py`, starting around the match on line 18." + ), + }, + { + "type": "tool_use", + "id": "toolu_vrtx_01Jxxsojf34b9skuGH47XkJd", + "name": "read_file", + "input": {"path": "shop/validation.py", "start_line": 1, "end_line": 60}, + }, + ], + "usage": { + "input_tokens": 4, + "cache_creation_input_tokens": 873, + "cache_read_input_tokens": 0, + "cache_creation": {"ephemeral_5m_input_tokens": 873, "ephemeral_1h_input_tokens": 0}, + "output_tokens": 255, + "output_tokens_details": {"thinking_tokens": 126}, + }, + }, + expected_litellm_response={ + **RESPONSES_RESPONSE, + "model": "vertex_ai/claude-haiku-5-5@default", + "output": [ + { + "type": "reasoning", + "id": ANY, + "status": "completed", + "role": "assistant", + "content": [ + { + "type": "output_text", + "text": ( + "I should read shop/validation.py to check how it handles zero, negative, and non-integer " + "values before moving to the next file—I'll grab the first 60 lines since I don't know the " + "full length yet.\n\n" + ), + "annotations": [], + }, + ], + "phase": None, + "encrypted_content": json.dumps( + [ + { + "type": "thinking", + "thinking": ( + "I should read shop/validation.py to check how it handles zero, negative, and non-integer " + "values before moving to the next file—I'll grab the first 60 lines since I don't know the " + "full length yet.\n\n" + ), + "signature": CLAUDE_HAIKU_5_5_RESPONSES_THINKING_SIGNATURE, + }, + ], + separators=(",", ":"), + ), + }, + { + "type": "message", + "id": ANY, + "status": "completed", + "role": "assistant", + "content": [ + { + "type": "output_text", + "text": "I'll read the validation function in `shop/validation.py`, starting around the match on line 18.", + "annotations": [], + }, + ], + "phase": None, + }, + { + "arguments": '{"path": "shop/validation.py", "start_line": 1, "end_line": 60}', + "call_id": "toolu_vrtx_01Jxxsojf34b9skuGH47XkJd", + "name": "read_file", + "type": "function_call", + "id": "fc_toolu_vrtx_01Jxxsojf34b9skuGH47XkJd", + "namespace": None, + "status": "completed", + }, + ], + "usage": { + "input_tokens": 877, + "input_tokens_details": { + "audio_tokens": None, + "cached_tokens": 0, + "cached_tokens_details": None, + "image_tokens": None, + "text_tokens": 4, + "video_tokens": None, + "cache_write_tokens": 873, + }, + "output_tokens": 255, + "output_tokens_details": {"audio_tokens": None, "reasoning_tokens": 126, "text_tokens": 129}, + "total_tokens": 1132, + "cost": None, + }, + "provider_specific_fields": { + "citations": None, + "thinking_blocks": [ + { + "type": "thinking", + "thinking": ( + "I should read shop/validation.py to check how it handles zero, negative, and non-integer " + "values before moving to the next file—I'll grab the first 60 lines since I don't know the " + "full length yet.\n\n" + ), + "signature": CLAUDE_HAIKU_5_5_RESPONSES_THINKING_SIGNATURE, + }, + ], + }, + }, +) + + CLAUDE_OPUS_5_5_TEST_CASE: Final = TranslationTestCase( scenario="basic", litellm_endpoint="/v1/responses", diff --git a/tests/integration/translation/responses/basic/test_responses_basic_vertex_ai.py b/tests/integration/translation/responses/basic/test_responses_basic_vertex_ai.py index 67c3cf4f7a9..bd9118482d0 100644 --- a/tests/integration/translation/responses/basic/test_responses_basic_vertex_ai.py +++ b/tests/integration/translation/responses/basic/test_responses_basic_vertex_ai.py @@ -5,6 +5,7 @@ from integration.translation.case import TranslationTestCase from integration.translation.responses.bases.vertex_ai import ( CLAUDE_FABLE_5_1_TEST_CASE, CLAUDE_HAIKU_4_5_TEST_CASE, + CLAUDE_HAIKU_5_5_TEST_CASE, CLAUDE_OPUS_4_8_TEST_CASE, CLAUDE_OPUS_5_5_TEST_CASE, CLAUDE_SONNET_4_6_TEST_CASE, @@ -26,6 +27,7 @@ from integration.translation.runner import assert_translation CLAUDE_SONNET_5_TEST_CASE, CLAUDE_OPUS_4_8_TEST_CASE, CLAUDE_SONNET_5_5_TEST_CASE, + CLAUDE_HAIKU_5_5_TEST_CASE, CLAUDE_OPUS_5_5_TEST_CASE, CLAUDE_FABLE_5_1_TEST_CASE, GEMINI_3_5_FLASH_TEST_CASE, diff --git a/tests/unit/litellm_core_utils/test_get_model_cost_map.py b/tests/unit/litellm_core_utils/test_get_model_cost_map.py index 0256ecf24ed..eca9bbdc170 100644 --- a/tests/unit/litellm_core_utils/test_get_model_cost_map.py +++ b/tests/unit/litellm_core_utils/test_get_model_cost_map.py @@ -686,6 +686,25 @@ def test_boot_load_does_not_retry_non_retryable_failure(): assert source["fallback_reason"] is not None +def test_boot_load_falls_back_when_urlsplit_rejects_the_url(): + client, calls = _mock_client([httpx.Response(200, content=_real_map_bytes())], client_cls=httpx.Client) + sleeper = _SyncSleepRecorder() + + cost_map = get_model_cost_map( + url="http://[bad/model_prices.json", + sleep=sleeper, + rng=random.Random(0), + client=client, + ) + assert calls["count"] == 0 + assert sleeper.waits == [] + assert _retry_threads() == [] + assert cost_map.keys() >= _load_root_cost_map().keys() - {"sample_spec", FALLBACK_GENERALIZATIONS_KEY} + source = get_model_cost_map_source_info() + assert source["source"] == "local" + assert "invalid URL" in (source["fallback_reason"] or "") + + def test_boot_load_respects_local_env_override(monkeypatch): """LITELLM_LOCAL_MODEL_COST_MAP=True still short-circuits to the backup with zero HTTP.""" monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") diff --git a/tests/unit/proxy/response_api_endpoints/test_endpoints.py b/tests/unit/proxy/response_api_endpoints/test_endpoints.py index a7b955f5b84..52e95901a16 100644 --- a/tests/unit/proxy/response_api_endpoints/test_endpoints.py +++ b/tests/unit/proxy/response_api_endpoints/test_endpoints.py @@ -480,7 +480,12 @@ def _advertised_cost_split(headers: Mapping[str, str]) -> list[tuple[float, floa @pytest.mark.asyncio @pytest.mark.parametrize( "background, stored, status, expected_billed_events", - [(True, True, "completed", 1), (True, True, "in_progress", 0), (True, False, "completed", 0), (False, True, "completed", 0)], + [ + (True, True, "completed", 1), + (True, True, "in_progress", 0), + (True, False, "completed", 0), + (False, True, "completed", 0), + ], ) async def test_reading_a_finished_background_response_bills_it( monkeypatch: pytest.MonkeyPatch, background: bool, stored: bool, status: str, expected_billed_events: int @@ -531,8 +536,12 @@ async def test_reading_a_finished_background_response_bills_it( prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=1) with ( - patch("litellm.proxy.proxy_server.prisma_client", prisma_client), # test-quality-ok: the endpoint reads the proxy singletons at call time - patch("litellm.proxy.proxy_server.llm_router", MagicMock()), # test-quality-ok: the endpoint reads the proxy singletons at call time + patch( + "litellm.proxy.proxy_server.prisma_client", prisma_client + ), # test-quality-ok: the endpoint reads the proxy singletons at call time + patch( + "litellm.proxy.proxy_server.llm_router", MagicMock() + ), # test-quality-ok: the endpoint reads the proxy singletons at call time patch("litellm.aget_responses", side_effect=AssertionError("billing must not read the provider again")), ): await _bill_finished_background_response( @@ -662,9 +671,15 @@ async def _read_twice_on_two_pods( router.aget_responses = AsyncMock(side_effect=AssertionError("billing must not read the provider again")) with ( - patch("litellm.proxy.proxy_server.prisma_client", prisma_client), # test-quality-ok: the endpoint reads the proxy singletons at call time - patch("litellm.proxy.proxy_server.llm_router", router), # test-quality-ok: the endpoint reads the proxy singletons at call time - patch("litellm.proxy.proxy_server.proxy_logging_obj", proxy_logging), # test-quality-ok: the endpoint reads the proxy singletons at call time + patch( + "litellm.proxy.proxy_server.prisma_client", prisma_client + ), # test-quality-ok: the endpoint reads the proxy singletons at call time + patch( + "litellm.proxy.proxy_server.llm_router", router + ), # test-quality-ok: the endpoint reads the proxy singletons at call time + patch( + "litellm.proxy.proxy_server.proxy_logging_obj", proxy_logging + ), # test-quality-ok: the endpoint reads the proxy singletons at call time patch("litellm.aget_responses", side_effect=AssertionError("billing must not read the provider again")), ): await asyncio.gather( @@ -702,7 +717,9 @@ async def test_a_finished_response_the_create_could_not_store_is_billed_once_und monkeypatch, table, response_id, - UserAPIKeyAuth(api_key="sk-reader", user_id="bob", team_id="team-beta", org_id="org-gamma", end_user_id="customer-42"), + UserAPIKeyAuth( + api_key="sk-reader", user_id="bob", team_id="team-beta", org_id="org-gamma", end_user_id="customer-42" + ), ) (billed,) = payloads @@ -746,7 +763,9 @@ async def test_a_finished_read_never_bills_a_response_already_billed_or_billed_b ) table: Final = _UniqueManagedObjectTable(stored) - payloads: Final = await _read_twice_on_two_pods(monkeypatch, table, response_id, UserAPIKeyAuth(api_key="sk-reader")) + payloads: Final = await _read_twice_on_two_pods( + monkeypatch, table, response_id, UserAPIKeyAuth(api_key="sk-reader") + ) assert payloads == () assert table.rows == stored @@ -878,9 +897,15 @@ async def test_a_finished_read_or_cancel_bills_usage_only_the_configured_deploym "litellm.proxy.response_api_endpoints.endpoints.ProxyBaseLLMRequestProcessing", return_value=processor, ), - patch("litellm.proxy.proxy_server.prisma_client", stored.prisma_client), # test-quality-ok: the endpoint reads the proxy singletons at call time - patch("litellm.proxy.proxy_server.llm_router", router), # test-quality-ok: the endpoint reads the proxy singletons at call time - patch("litellm.proxy.proxy_server.proxy_logging_obj", stored.proxy_logging), # test-quality-ok: the endpoint reads the proxy singletons at call time + patch( + "litellm.proxy.proxy_server.prisma_client", stored.prisma_client + ), # test-quality-ok: the endpoint reads the proxy singletons at call time + patch( + "litellm.proxy.proxy_server.llm_router", router + ), # test-quality-ok: the endpoint reads the proxy singletons at call time + patch( + "litellm.proxy.proxy_server.proxy_logging_obj", stored.proxy_logging + ), # test-quality-ok: the endpoint reads the proxy singletons at call time patch("litellm.aget_responses", side_effect=AssertionError("the configured deployment answers billing reads")), ): await getattr(endpoints, endpoint_name)( @@ -941,9 +966,15 @@ async def test_deleting_a_finished_unbilled_background_response_bills_it_before_ "litellm.proxy.response_api_endpoints.endpoints.ProxyBaseLLMRequestProcessing", return_value=processor, ), - patch("litellm.proxy.proxy_server.prisma_client", stored.prisma_client), # test-quality-ok: the endpoint reads the proxy singletons at call time - patch("litellm.proxy.proxy_server.llm_router", router), # test-quality-ok: the endpoint reads the proxy singletons at call time - patch("litellm.proxy.proxy_server.proxy_logging_obj", stored.proxy_logging), # test-quality-ok: the endpoint reads the proxy singletons at call time + patch( + "litellm.proxy.proxy_server.prisma_client", stored.prisma_client + ), # test-quality-ok: the endpoint reads the proxy singletons at call time + patch( + "litellm.proxy.proxy_server.llm_router", router + ), # test-quality-ok: the endpoint reads the proxy singletons at call time + patch( + "litellm.proxy.proxy_server.proxy_logging_obj", stored.proxy_logging + ), # test-quality-ok: the endpoint reads the proxy singletons at call time ): result: Final = await delete_response( response_id=stored.response_id, @@ -989,7 +1020,9 @@ async def test_a_delete_the_proxy_refuses_never_fetches_or_bills_the_response(ca caller: Final = UserAPIKeyAuth(api_key="sk-caller", user_role=caller_role) async def pipeline(**_: object) -> DeleteResponseResult: - await security.async_pre_call_hook(caller, MagicMock(), {"response_id": stored.response_id}, "adelete_responses") + await security.async_pre_call_hook( + caller, MagicMock(), {"response_id": stored.response_id}, "adelete_responses" + ) return DeleteResponseResult(id=stored.response_id, object="response", deleted=True) processor: Final = MagicMock() @@ -1002,9 +1035,15 @@ async def test_a_delete_the_proxy_refuses_never_fetches_or_bills_the_response(ca "litellm.proxy.response_api_endpoints.endpoints.ProxyBaseLLMRequestProcessing", return_value=processor, ), - patch("litellm.proxy.proxy_server.prisma_client", stored.prisma_client), # test-quality-ok: the endpoint reads the proxy singletons at call time - patch("litellm.proxy.proxy_server.llm_router", router), # test-quality-ok: the endpoint reads the proxy singletons at call time - patch("litellm.proxy.proxy_server.proxy_logging_obj", stored.proxy_logging), # test-quality-ok: the endpoint reads the proxy singletons at call time + patch( + "litellm.proxy.proxy_server.prisma_client", stored.prisma_client + ), # test-quality-ok: the endpoint reads the proxy singletons at call time + patch( + "litellm.proxy.proxy_server.llm_router", router + ), # test-quality-ok: the endpoint reads the proxy singletons at call time + patch( + "litellm.proxy.proxy_server.proxy_logging_obj", stored.proxy_logging + ), # test-quality-ok: the endpoint reads the proxy singletons at call time ): with contextlib.suppress(HTTPException): await delete_response( @@ -1058,9 +1097,15 @@ def _capture_spend_events(monkeypatch: pytest.MonkeyPatch) -> _SpendEvents: @contextlib.contextmanager def _proxy_over(prisma_client: MagicMock, proxy_logging: MagicMock, router: MagicMock): with ( - patch("litellm.proxy.proxy_server.prisma_client", prisma_client), # test-quality-ok: the endpoint reads the proxy singletons at call time - patch("litellm.proxy.proxy_server.llm_router", router), # test-quality-ok: the endpoint reads the proxy singletons at call time - patch("litellm.proxy.proxy_server.proxy_logging_obj", proxy_logging), # test-quality-ok: the endpoint reads the proxy singletons at call time + patch( + "litellm.proxy.proxy_server.prisma_client", prisma_client + ), # test-quality-ok: the endpoint reads the proxy singletons at call time + patch( + "litellm.proxy.proxy_server.llm_router", router + ), # test-quality-ok: the endpoint reads the proxy singletons at call time + patch( + "litellm.proxy.proxy_server.proxy_logging_obj", proxy_logging + ), # test-quality-ok: the endpoint reads the proxy singletons at call time ): yield @@ -1211,7 +1256,9 @@ async def test_a_streamed_background_create_is_billed_wherever_it_sets_backgroun "litellm.proxy.response_polling.polling_handler.should_use_polling_for_request", return_value=False, ), - patch("litellm.proxy.proxy_server.select_data_generator", provider_events), # test-quality-ok: the endpoint reads the proxy singletons at call time + patch( + "litellm.proxy.proxy_server.select_data_generator", provider_events + ), # test-quality-ok: the endpoint reads the proxy singletons at call time ): await responses_api( request=_request_for("POST", "/v1/responses", body), @@ -1329,7 +1376,12 @@ async def test_background_create_the_provider_confirms_is_stored_under_its_provi } request = Request( - {"type": "http", "method": "POST", "path": "/v1/responses", "headers": [(b"content-type", b"application/json")]}, + { + "type": "http", + "method": "POST", + "path": "/v1/responses", + "headers": [(b"content-type", b"application/json")], + }, receive, ) @@ -1918,9 +1970,327 @@ class TestResponsesWSFirstFrameModelAuth: assert routed["input"] == replayed_input assert routed["previous_response_id"] == "resp_prev" assert processor.common_processing_pre_call_logic.await_args.kwargs["model"] == "gpt-4o-mini" + assert processor.common_processing_pre_call_logic.await_args.kwargs["skip_guardrails"] is True + assert callable(routed["responses_ws_frame_guard"]) assert mock_route_request.await_args.kwargs["route_type"] == "_aresponses_websocket" ws.close.assert_not_awaited() + @pytest.mark.asyncio + async def test_frame_guard_runs_the_guardrail_walk_on_the_frame_body_and_returns_its_records(self): + from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard + + async def pre_call_hook(user_api_key_dict, data, call_type, guardrails_only): + assert call_type == "aresponses" and guardrails_only is True + assert "type" not in data and data["model"] == "gpt-4o-mini" + assert data["litellm_metadata"]["user_api_key_hash"] == "hash" + assert "standard_logging_guardrail_information" not in data["litellm_metadata"] + data["litellm_metadata"]["standard_logging_guardrail_information"] = [{"guardrail_name": "turn"}] + return {**data, "input": data["input"] + " [guarded]"} + + proxy_logging_obj = MagicMock() + proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=pre_call_hook) + admitted = { + "model": "gpt-4o-mini", + "litellm_metadata": { + "user_api_key_hash": "hash", + "standard_logging_guardrail_information": [{"guardrail_name": "admission"}], + }, + } + guard = _responses_ws_frame_guard(admitted, MagicMock(), proxy_logging_obj) + + first = await guard(json.dumps({"type": "response.create", "model": "alias", "input": "hi"})) + second = await guard(json.dumps({"type": "response.create", "input": "again"})) + + assert json.loads(first.frame) == {"type": "response.create", "model": "alias", "input": "hi [guarded]"} + assert first.guardrail_information == ({"guardrail_name": "turn"},) + assert json.loads(second.frame)["input"] == "again [guarded]" + assert admitted["litellm_metadata"]["standard_logging_guardrail_information"] == [ + {"guardrail_name": "admission"} + ] + + @pytest.mark.asyncio + async def test_frame_guard_bills_a_rejected_frame_as_a_failed_request(self): + import litellm + from litellm.proxy.hooks.proxy_track_cost_callback import ProxyDBLogger + from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard + + async def pre_call_hook(user_api_key_dict, data, call_type, guardrails_only): + data["litellm_metadata"]["standard_logging_guardrail_information"] = [{"guardrail_cost": 0.5}] + raise ValueError("blocked") + + proxy_logging_obj = MagicMock() + proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=pre_call_hook) + proxy_logging_obj.post_call_failure_hook = AsyncMock() + spend_logger = MagicMock(spec=ProxyDBLogger) + spend_logger.async_post_call_failure_hook = AsyncMock() + rate_limiter = MagicMock() + rate_limiter.async_post_call_failure_hook = AsyncMock() + user_api_key_dict = MagicMock() + guard = _responses_ws_frame_guard( + {"model": "gpt-4o-mini", "litellm_metadata": {}}, user_api_key_dict, proxy_logging_obj + ) + + with ( + patch.object(litellm, "callbacks", [rate_limiter, spend_logger]), + pytest.raises(ValueError, match="blocked"), + ): + await guard(json.dumps({"type": "response.create", "input": "hi"})) + + proxy_logging_obj.post_call_failure_hook.assert_not_awaited() + rate_limiter.async_post_call_failure_hook.assert_not_awaited() + spend_logger.async_post_call_failure_hook.assert_awaited_once() + failure = spend_logger.async_post_call_failure_hook.await_args.kwargs + user_api_key_dict.model_copy.assert_called_once_with(update={"budget_reservation": None}) + assert failure["user_api_key_dict"] is user_api_key_dict.model_copy.return_value + assert str(failure["original_exception"]) == "blocked" + assert failure["request_data"]["litellm_metadata"]["standard_logging_guardrail_information"] == [ + {"guardrail_cost": 0.5} + ] + assert failure["request_data"]["input"] == "hi" + assert len(failure["request_data"]["litellm_call_id"]) == 36 + + @pytest.mark.asyncio + async def test_frame_guard_guards_the_nested_response_body_and_copies_back_tools(self): + from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard + + async def pre_call_hook(user_api_key_dict, data, call_type, guardrails_only): + return {**data, "input": data["input"] + " [guarded]", "tools": data["tools"][:1]} + + proxy_logging_obj = MagicMock() + proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=pre_call_hook) + guard = _responses_ws_frame_guard( + {"model": "gpt-4o-mini", "litellm_metadata": {}}, MagicMock(), proxy_logging_obj + ) + tools = [{"type": "function", "name": "a"}, {"type": "function", "name": "b"}] + + guarded = await guard(json.dumps({"type": "response.create", "response": {"input": "hi", "tools": tools}})) + + assert proxy_logging_obj.pre_call_hook.await_args.kwargs["data"]["input"] == "hi" + assert json.loads(guarded.frame) == { + "type": "response.create", + "response": {"input": "hi [guarded]", "tools": tools[:1]}, + } + + @pytest.mark.asyncio + async def test_frame_guard_strips_the_proxy_owned_metadata_keys_a_client_sends(self): + from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard + + proxy_logging_obj = MagicMock() + proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=lambda **kwargs: kwargs["data"]) + guard = _responses_ws_frame_guard( + {"model": "gpt-4o-mini", "litellm_metadata": {}}, MagicMock(), proxy_logging_obj + ) + metadata = { + "user_api_key_metadata": {"disable_global_guardrails": True}, + "disable_global_guardrails": True, + "headers": {"x-litellm-tags": "forged"}, + "tag": "kept", + } + + await guard(json.dumps({"type": "response.create", "input": "hi", "metadata": metadata})) + + assert proxy_logging_obj.pre_call_hook.await_args.kwargs["data"]["metadata"] == {"tag": "kept"} + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "frame", + [ + {"type": "response.create", "input": "forbidden", "guardrails": []}, + {"type": "response.create", "input": "forbidden", "metadata": {"guardrails": []}}, + {"type": "response.create", "response": {"input": "forbidden", "guardrails": [], "policies": []}}, + ], + ) + async def test_frame_guard_keeps_the_admitted_guardrail_list_when_a_frame_tries_to_replace_it(self, frame): + from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard + + proxy_logging_obj = MagicMock() + proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=lambda **kwargs: kwargs["data"]) + guard = _responses_ws_frame_guard( + {"model": "gpt-4o-mini", "litellm_metadata": {"guardrails": ["key-required"]}}, + MagicMock(), + proxy_logging_obj, + ) + + await guard(json.dumps(frame)) + + scanned = proxy_logging_obj.pre_call_hook.await_args.kwargs["data"] + assert "guardrails" not in scanned and "policies" not in scanned + assert "guardrails" not in (scanned.get("metadata") or {}) + assert scanned["litellm_metadata"]["guardrails"] == ["key-required"] + + @pytest.mark.asyncio + async def test_frame_guard_copies_an_admitted_bucket_that_holds_the_otel_parent_span(self): + import threading + + from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard + + proxy_logging_obj = MagicMock() + proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=lambda **kwargs: kwargs["data"]) + span = threading.RLock() + guard = _responses_ws_frame_guard( + {"model": "gpt-4o-mini", "litellm_metadata": {"litellm_parent_otel_span": span}}, + MagicMock(), + proxy_logging_obj, + ) + + guarded = await guard(json.dumps({"type": "response.create", "input": "hi"})) + + assert json.loads(guarded.frame) == {"type": "response.create", "input": "hi"} + assert ( + proxy_logging_obj.pre_call_hook.await_args.kwargs["data"]["litellm_metadata"]["litellm_parent_otel_span"] + is span + ) + + @pytest.mark.asyncio + async def test_frame_guard_rejects_a_turn_a_guardrail_rerouted_to_another_model(self): + from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard + + from litellm.proxy.hooks.proxy_track_cost_callback import ProxyDBLogger + + async def pre_call_hook(user_api_key_dict, data, call_type, guardrails_only): + data["litellm_metadata"]["standard_logging_guardrail_information"] = [{"guardrail_cost": 0.5}] + return {**data, "model": "onprem-llama"} + + proxy_logging_obj = MagicMock() + proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=pre_call_hook) + spend_logger = MagicMock(spec=ProxyDBLogger) + spend_logger.async_post_call_failure_hook = AsyncMock() + guard = _responses_ws_frame_guard( + {"model": "gpt-4o-mini", "litellm_metadata": {}}, MagicMock(), proxy_logging_obj + ) + + with ( + patch.object(litellm, "callbacks", [spend_logger]), + pytest.raises(ValueError, match="rerouted this turn to 'onprem-llama'"), + ): + await guard(json.dumps({"type": "response.create", "input": "my SSN is 123-45-6789"})) + + failure = spend_logger.async_post_call_failure_hook.await_args.kwargs + assert "rerouted this turn to 'onprem-llama'" in str(failure["original_exception"]) + assert failure["request_data"]["litellm_metadata"]["standard_logging_guardrail_information"] == [ + {"guardrail_cost": 0.5} + ] + + @pytest.mark.asyncio + @pytest.mark.parametrize("nested", [False, True]) + async def test_frame_guard_drops_a_client_s_guardrail_cost_records_so_a_blocked_frame_bills_the_real_scan( + self, nested: bool + ): + import litellm + from litellm.proxy.hooks.proxy_track_cost_callback import ProxyDBLogger + from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard + + async def pre_call_hook(user_api_key_dict, data, call_type, guardrails_only): + data["litellm_metadata"]["standard_logging_guardrail_information"] = [{"guardrail_cost": 0.5}] + raise ValueError("blocked") + + proxy_logging_obj = MagicMock() + proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=pre_call_hook) + spend_logger = MagicMock(spec=ProxyDBLogger) + spend_logger.async_post_call_failure_hook = AsyncMock() + guard = _responses_ws_frame_guard( + {"model": "gpt-4o-mini", "litellm_metadata": {}}, MagicMock(), proxy_logging_obj + ) + body = { + "input": "forbidden", + "metadata": {"trace": "t1", "standard_logging_guardrail_information": [], "model_info": {"x": 1}}, + } + frame = {"type": "response.create", "response": body} if nested else {"type": "response.create", **body} + + with patch.object(litellm, "callbacks", [spend_logger]), pytest.raises(ValueError, match="blocked"): + await guard(json.dumps(frame)) + + billed = spend_logger.async_post_call_failure_hook.await_args.kwargs["request_data"] + assert billed["metadata"] == {"trace": "t1"} + assert billed["litellm_metadata"]["standard_logging_guardrail_information"] == [{"guardrail_cost": 0.5}] + + @pytest.mark.asyncio + @pytest.mark.parametrize("nested", [False, True]) + async def test_frame_guard_forwards_a_frame_without_the_instructions_a_guardrail_deleted(self, nested: bool): + from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard + + async def pre_call_hook(user_api_key_dict, data, call_type, guardrails_only): + return {key: value for key, value in data.items() if key != "instructions"} + + proxy_logging_obj = MagicMock() + proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=pre_call_hook) + guard = _responses_ws_frame_guard( + {"model": "gpt-4o-mini", "litellm_metadata": {}}, MagicMock(), proxy_logging_obj + ) + body = {"instructions": "SECRET system prompt", "input": "hi", "store": False} + frame = {"type": "response.create", "response": body} if nested else {"type": "response.create", **body} + + with patch.object(litellm, "callbacks", []): + guarded = await guard(json.dumps(frame)) + + forwarded = json.loads(guarded.frame) + assert forwarded["type"] == "response.create" + sent_body = forwarded["response"] if nested else {k: v for k, v in forwarded.items() if k != "type"} + assert sent_body == {"input": "hi", "store": False} + + @pytest.mark.asyncio + async def test_frame_guard_hands_back_the_pii_tokens_presidio_masked_with_and_drops_forged_ones(self): + from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard + + async def pre_call_hook(user_api_key_dict, data, call_type, guardrails_only): + assert data["metadata"] == {"trace": "t1"}, "a client cannot seed pii_tokens" + data["metadata"]["pii_tokens"] = {"": "alice@example.com"} + return {**data, "input": "email please"} + + proxy_logging_obj = MagicMock() + proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=pre_call_hook) + guard = _responses_ws_frame_guard( + {"model": "gpt-4o-mini", "litellm_metadata": {}}, MagicMock(), proxy_logging_obj + ) + + with patch.object(litellm, "callbacks", []): + guarded = await guard( + json.dumps( + { + "type": "response.create", + "input": "email alice@example.com please", + "metadata": {"trace": "t1", "pii_tokens": {"": "mallory@example.com"}}, + } + ) + ) + + assert json.loads(guarded.frame)["input"] == "email please" + assert dict(guarded.pii_tokens) == {"": "alice@example.com"} + + @pytest.mark.asyncio + async def test_frame_guard_rejects_a_frame_that_carries_input_at_both_levels(self): + from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard + + proxy_logging_obj = MagicMock() + proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=lambda **kwargs: kwargs["data"]) + guard = _responses_ws_frame_guard( + {"model": "gpt-4o-mini", "litellm_metadata": {}}, MagicMock(), proxy_logging_obj + ) + + with patch.object(litellm, "callbacks", []), pytest.raises(ValueError, match="both at the top level"): + await guard(json.dumps({"type": "response.create", "input": "forbidden", "response": {}})) + + proxy_logging_obj.pre_call_hook.assert_not_awaited() + + @pytest.mark.asyncio + async def test_frame_guard_keeps_the_guardrail_error_when_billing_the_rejected_frame_fails(self): + import litellm + from litellm.proxy.hooks.proxy_track_cost_callback import ProxyDBLogger + from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard + + proxy_logging_obj = MagicMock() + proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=ValueError("blocked")) + spend_logger = MagicMock(spec=ProxyDBLogger) + spend_logger.async_post_call_failure_hook = AsyncMock(side_effect=RuntimeError("db down")) + guard = _responses_ws_frame_guard( + {"model": "gpt-4o-mini", "litellm_metadata": {}}, MagicMock(), proxy_logging_obj + ) + + with patch.object(litellm, "callbacks", [spend_logger]), pytest.raises(ValueError, match="blocked"): + await guard(json.dumps({"type": "response.create", "input": "hi"})) + + spend_logger.async_post_call_failure_hook.assert_awaited_once() + @pytest.mark.asyncio @pytest.mark.parametrize("provider_rejected", [True, False]) async def test_endpoint_books_a_provider_rejected_connection_as_a_failed_request(self, provider_rejected: bool): @@ -3614,16 +3984,18 @@ class TestResponsesInputTokens: model_list=[ { "model_name": "gpt-5.2", - "litellm_params": {"model": "openai/gpt-5.2", "api_key": "sk-test", "additional_drop_params": ["tools"]}, + "litellm_params": { + "model": "openai/gpt-5.2", + "api_key": "sk-test", + "additional_drop_params": ["tools"], + }, }, {"model_name": "backup", "litellm_params": {"model": "openai/gpt-5.4-mini", "api_key": "sk-test"}}, ], fallbacks=[{"gpt-5.2": ["backup"]}], ) - response, counter = self._post_input_tokens( - {"model": "gpt-5.2", "input": "hi", "tools": 5}, llm_router=router - ) + response, counter = self._post_input_tokens({"model": "gpt-5.2", "input": "hi", "tools": 5}, llm_router=router) assert response.status_code == 200, response.text assert counter.await_args.kwargs["request"].tools is None diff --git a/tests/unit/responses/test_responses_websocket_all_providers.py b/tests/unit/responses/test_responses_websocket_all_providers.py index 2b78aa98b04..b38808fb57c 100644 --- a/tests/unit/responses/test_responses_websocket_all_providers.py +++ b/tests/unit/responses/test_responses_websocket_all_providers.py @@ -285,6 +285,151 @@ class TestManagedWebSocketHandlerIntegration: assert captured["model"] == "bedrock_mantle/openai.gpt-5.5" + @pytest.mark.asyncio + async def test_managed_socket_runs_the_frame_guard_on_every_turn(self, monkeypatch): + import json + from unittest.mock import AsyncMock, MagicMock + + import litellm + from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.responses.streaming_iterator import GuardedResponsesWsFrame, ManagedResponsesWebSocketHandler + + async def guard(frame: str) -> GuardedResponsesWsFrame: + body = json.loads(frame) + if "forbidden" in body["input"]: + raise ValueError("blocked by guardrail") + return GuardedResponsesWsFrame( + frame=json.dumps({**body, "input": body["input"].upper()}), + guardrail_information=({"guardrail_name": "g", "guardrail_cost": 0.5},), + ) + + async def empty_stream(): + return + yield {} + + aresponses: Final = AsyncMock(side_effect=lambda **kwargs: empty_stream()) + monkeypatch.setattr(litellm, "aresponses", aresponses) + websocket: Final = MagicMock() + websocket.send_text = AsyncMock() + handler: Final = ManagedResponsesWebSocketHandler( + websocket=websocket, + model="anthropic/claude-sonnet-4-5", + logging_obj=Logging( + model="anthropic/claude-sonnet-4-5", + messages=[], + stream=True, + call_type="aresponses", + start_time=0, + litellm_call_id="test-id", + function_id="test-func", + ), + litellm_metadata={"user_api_key_hash": "k1"}, + frame_guard=guard, + ) + + await handler._process_response_create(json.dumps({"type": "response.create", "input": "say forbidden"})) + await handler._process_response_create(json.dumps({"type": "response.create", "input": "hello"})) + + aresponses.assert_awaited_once() + error_event: Final = json.loads(websocket.send_text.await_args_list[0].args[0]) + assert error_event["error"] == {"type": "invalid_request_error", "message": "blocked by guardrail"} + sent: Final = aresponses.await_args.kwargs + assert sent["input"] == "HELLO" + assert sent["litellm_metadata"]["user_api_key_hash"] == "k1" + assert sent["litellm_metadata"]["standard_logging_guardrail_information"] == [ + {"guardrail_name": "g", "guardrail_cost": 0.5} + ] + + @pytest.mark.asyncio + async def test_managed_socket_unmasks_the_reply_with_the_frame_guard_s_pii_tokens(self, monkeypatch): + import json + from unittest.mock import AsyncMock, MagicMock + + import litellm + from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.responses.streaming_iterator import GuardedResponsesWsFrame, ManagedResponsesWebSocketHandler + + presidio = _FakeWSGuardrail() + + async def guard(frame: str) -> GuardedResponsesWsFrame: + body = json.loads(frame) + masked = await presidio.check_pii(body["input"], True, None, walk := {}) + return GuardedResponsesWsFrame( + frame=json.dumps({**body, "input": masked}), pii_tokens=walk["metadata"]["pii_tokens"] + ) + + provider_saw: dict = {} + + async def fake_aresponses(**kwargs): + provider_saw.update(kwargs) + + async def stream(): + yield {"type": "response.output_text.delta", "delta": "Sent to "} + yield { + "type": "response.completed", + "response": { + "id": "resp_1", + "output": [ + {"type": "message", "content": [{"type": "output_text", "text": "Sent to "}]} + ], + }, + } + + return stream() + + monkeypatch.setattr(litellm, "aresponses", fake_aresponses) + websocket: Final = MagicMock() + websocket.send_text = AsyncMock() + handler: Final = ManagedResponsesWebSocketHandler( + websocket=websocket, + model="anthropic/claude-sonnet-4-5", + logging_obj=Logging( + model="anthropic/claude-sonnet-4-5", + messages=[], + stream=True, + call_type="aresponses", + start_time=0, + litellm_call_id="test-id", + function_id="test-func", + ), + frame_guard=guard, + guardrail_callbacks=[presidio], + ) + + await handler._process_response_create( + json.dumps({"type": "response.create", "input": "email alice@example.com please"}) + ) + + assert provider_saw["input"] == "email please" + sent: Final = [json.loads(call.args[0]) for call in websocket.send_text.await_args_list] + assert sent[0]["delta"] == "Sent to alice@example.com" + assert sent[1]["response"]["output"][0]["content"][0]["text"] == "Sent to alice@example.com" + + @pytest.mark.asyncio + async def test_http_handler_hands_the_frame_guard_to_a_managed_socket(self): + import json + from unittest.mock import AsyncMock, MagicMock + + from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler + + guard: Final = AsyncMock(side_effect=ValueError("blocked by guardrail")) + websocket: Final = MagicMock() + websocket.send_text = AsyncMock() + websocket.receive_text = AsyncMock(side_effect=ConnectionError("client left")) + + await BaseLLMHTTPHandler().async_responses_websocket( + model="anthropic/claude-sonnet-4-5", + websocket=websocket, + logging_obj=MagicMock(), + responses_api_provider_config=None, + first_message=json.dumps({"type": "response.create", "input": "say forbidden"}), + frame_guard=guard, + ) + + guard.assert_awaited_once() + error_event: Final = json.loads(websocket.send_text.await_args.args[0]) + assert error_event["error"]["message"] == "blocked by guardrail" + @pytest.mark.parametrize("terminal_type", ["response.completed", "response.incomplete"]) @pytest.mark.asyncio async def test_truncated_turn_is_kept_as_history_for_the_next_turn(self, monkeypatch, terminal_type: str): @@ -1319,6 +1464,97 @@ class TestWebSocketProjectQuotaEnforcement: mock_websocket.send_text.assert_called_once() assert "rate_limit_exceeded" in mock_websocket.send_text.call_args[0][0] + @pytest.mark.asyncio + async def test_native_handler_runs_the_frame_guard_on_every_turn_and_bills_each_turn_its_own_guardrail_records( + self, + ): + from unittest.mock import AsyncMock, MagicMock + + from litellm.responses.streaming_iterator import GuardedResponsesWsFrame, ResponsesWebSocketStreaming + + def connection_logging(): + logging_obj = MagicMock() + logging_obj.litellm_params = { + "metadata": { + "standard_logging_guardrail_information": [{"guardrail_name": "admission", "guardrail_cost": 9.0}] + } + } + logging_obj.model_call_details = {} + logging_obj.copy_for_new_call.side_effect = lambda **_: connection_logging() + return logging_obj + + async def guard(frame: str) -> GuardedResponsesWsFrame: + body = json.loads(frame) + return GuardedResponsesWsFrame( + frame=json.dumps({**body, "input": body["input"] + " [guarded]"}), + guardrail_information=({"guardrail_name": "turn", "guardrail_cost": 0.1 * len(body["input"])},), + ) + + frames = [json.dumps({"type": "response.create", "input": text}) for text in ("hi", "hello")] + mock_backend_ws = MagicMock() + mock_backend_ws.send = AsyncMock() + mock_websocket = MagicMock() + mock_websocket.send_text = AsyncMock() + mock_websocket.receive_text = AsyncMock(side_effect=[frames[1], RuntimeError("client closed")]) + + handler = ResponsesWebSocketStreaming( + websocket=mock_websocket, + backend_ws=mock_backend_ws, + logging_obj=connection_logging(), + first_message=frames[0], + frame_guard=guard, + ) + await handler.client_to_backend() + + sent = [json.loads(call.args[0]) for call in mock_backend_ws.send.call_args_list] + assert [frame["input"] for frame in sent] == ["hi [guarded]", "hello [guarded]"] + assert [ + turn.logging_obj.litellm_params["metadata"]["standard_logging_guardrail_information"] + for turn in handler._open_turns + ] == [ + [{"guardrail_name": "turn", "guardrail_cost": pytest.approx(0.2)}], + [{"guardrail_name": "turn", "guardrail_cost": pytest.approx(0.5)}], + ] + + @pytest.mark.asyncio + async def test_native_handler_drops_a_frame_the_guard_rejects_and_keeps_the_socket_open(self): + from unittest.mock import AsyncMock, MagicMock + + from litellm.responses.streaming_iterator import GuardedResponsesWsFrame, ResponsesWebSocketStreaming + + async def guard(frame: str) -> GuardedResponsesWsFrame: + if "forbidden" in frame: + raise ValueError("turn guardrail blocked the request") + return GuardedResponsesWsFrame(frame=frame) + + blocked = json.dumps({"type": "response.create", "input": "forbidden"}) + allowed = json.dumps({"type": "response.create", "input": "fine"}) + mock_backend_ws = MagicMock() + mock_backend_ws.send = AsyncMock() + mock_websocket = MagicMock() + mock_websocket.send_text = AsyncMock() + mock_websocket.receive_text = AsyncMock(side_effect=[allowed, RuntimeError("client closed")]) + logging_obj = MagicMock() + logging_obj.copy_for_new_call.return_value = logging_obj + + handler = ResponsesWebSocketStreaming( + websocket=mock_websocket, + backend_ws=mock_backend_ws, + logging_obj=logging_obj, + first_message=blocked, + frame_guard=guard, + ) + await handler.client_to_backend() + + error_event = json.loads(mock_websocket.send_text.call_args[0][0]) + assert error_event["type"] == "error" + assert error_event["error"] == { + "type": "invalid_request_error", + "message": "turn guardrail blocked the request", + } + assert [json.loads(call.args[0])["input"] for call in mock_backend_ws.send.call_args_list] == ["fine"] + assert len(handler._open_turns) == 1 + @pytest.mark.asyncio async def test_native_handler_forwards_frame_allowed_by_quota_callback(self): from unittest.mock import AsyncMock, MagicMock @@ -1905,6 +2141,34 @@ def _make_streaming(**kwargs): class TestNativeWebSocketGuardrailMasking: """Exercises the input/output PII masking hooks on ResponsesWebSocketStreaming.""" + @pytest.mark.asyncio + async def test_a_reply_to_a_frame_the_guard_masked_is_unmasked_with_the_guard_s_tokens(self): + from unittest.mock import AsyncMock, MagicMock + + from litellm.responses.streaming_iterator import GuardedResponsesWsFrame + + guardrail = _FakeWSGuardrail() + + async def guard(frame: str) -> GuardedResponsesWsFrame: + masked = await guardrail.check_pii(json.loads(frame)["input"], True, None, walk := {}) + return GuardedResponsesWsFrame( + frame=json.dumps({"type": "response.create", "input": masked}), + pii_tokens=walk["metadata"]["pii_tokens"], + ) + + backend_ws = MagicMock() + backend_ws.send = AsyncMock() + handler = _make_streaming( + backend_ws=backend_ws, request_data={}, guardrail_callbacks=[guardrail], frame_guard=guard + ) + handler.logging_obj.copy_for_new_call.return_value = handler.logging_obj + + await handler._forward_guarded_frame(json.dumps({"type": "response.create", "input": "email alice@example.com"})) + + assert json.loads(backend_ws.send.await_args.args[0])["input"] == "email " + reply = json.dumps({"type": "response.output_text.delta", "delta": "Sent to "}) + assert json.loads(handler._unmask_response_event(reply))["delta"] == "Sent to alice@example.com" + @pytest.mark.asyncio async def test_mask_response_create_flat_string_input(self): guardrail = _FakeWSGuardrail()